# Python function to compute BIO1-19
# using monthly tas, tasmax, tasmin, and prec
# Input: monthly tas, tasmax, tasmin, prec in a year
#             temperature is in K*10
#              precipitation is in mm*100
# Please see Tabl 7.2 in the CHELSA_tech_specification_V2.pdf
# Output: B1,B2,...,B19
# All input variables are 3-D in the shape of [time, lat, lon]
# All output variables are 2-D in the shape of [lat,lon]
# All variables are numpy arrays

import numpy as np
from osgeo import gdal
from BIOVAR import biovar
import os

DEBUG=0

# define a function to write the BIO
# def wrtbio(fname,ds,rows,cols,layer,arr_out):
#     driver = gdal.GetDriverByName("GTiff")
#     outdata = driver.Create(fname, cols, rows, 1, gdal.GDT_UInt16)
#     outdata.SetGeoTransform(ds.GetGeoTransform())##sets same geotransform as input
#     outdata.SetProjection(ds.GetProjection())##sets same projection as input
#     outdata.GetRasterBand(1).WriteArray(arr_out)
#     outdata.FlushCache() ##save to disk!!
#     outdata = None

# New writer (Float32 + explicit NoData + compression)
def wrtbio(fname, ds, rows, cols, layer, arr_out):
    driver = gdal.GetDriverByName("GTiff")

    # Create float32 GeoTIFF, compressed
    outdata = driver.Create(
        fname, cols, rows, 1, gdal.GDT_Float32,
        options=["TILED=YES", "COMPRESS=DEFLATE", "PREDICTOR=2"]
    )
    outdata.SetGeoTransform(ds.GetGeoTransform())  # same geotransform as input
    outdata.SetProjection(ds.GetProjection())      # same projection as input

    band = outdata.GetRasterBand(1)

    # Explicit NoData for downstream tools (e.g., GEE)
    nodata = -9999.0
    band.SetNoDataValue(nodata)

    # Write array safely
    a = np.asarray(arr_out, dtype=np.float32)
    a[~np.isfinite(a)] = nodata
    band.WriteArray(a)

    band.FlushCache()
    outdata.FlushCache()
    outdata = None
    
#determine the shape of the tif
ds = gdal.Open(r"/D2/CHELSA/Cote/monthly/pr/Cote_CHELSA_pr_01_1980_V.2.1.tif")
#print(ds.RasterCount )
band = ds.GetRasterBand(1)
tas = band.ReadAsArray()
ix = ds.RasterXSize
iy = ds.RasterYSize
it=14
#print(it,iy,ix)
#b=list
for j in range(1980,2018):
#for j in range(1980,1981):
    print(j)
    rain=np.zeros((it,iy,ix))
    t2m=np.zeros((it,iy,ix))
    tmax=np.zeros((it,iy,ix))
    tmin=np.zeros((it,iy,ix))
    yy=str(j)
    dest='/D2/CHELSA/Cote/yearly/bio/'+yy
    os.system('mkdir -p '+dest)
    for i in range(it):
        ii=str(i+1).zfill(2)
        if i>11:
            ii=str(i-11).zfill(2)
            yy=str(j+1)
        ds = gdal.Open(r'/D2/CHELSA/Cote/monthly/pr/Cote_CHELSA_pr_'+ii+'_'+yy+'_V.2.1.tif')
        band = ds.GetRasterBand(1)
        tas = band.ReadAsArray()
        rain[i,:,:]=tas[:,:]/100 #the monthly pr is in mm*100, see Table 7.2 in the Techreport
        del ds
        ds = gdal.Open(r'/D2/CHELSA/Cote/monthly/tas/Cote_CHELSA_tas_'+ii+'_'+yy+'_V.2.1.tif')
        band = ds.GetRasterBand(1)
        tas = band.ReadAsArray()
        t2m[i,:,:]=tas[:,:]/10 -273.15 #the monthly tas, tasmax, tasmin all in K*10, Table 7.2
        del ds
        ds = gdal.Open(r'/D2/CHELSA/Cote/monthly/tasmax/Cote_CHELSA_tasmax_'+ii+'_'+yy+'_V.2.1.tif')
        band = ds.GetRasterBand(1)
        tas = band.ReadAsArray()
        tmax[i,:,:]=tas[:,:]/10 -273.15
        del ds
        ds = gdal.Open(r'/D2/CHELSA/Cote/monthly/tasmin/Cote_CHELSA_tasmin_'+ii+'_'+yy+'_V.2.1.tif')
        band = ds.GetRasterBand(1)
        tas = band.ReadAsArray()
        tmin[i,:,:]=tas[:,:]/10 -273.15
    if DEBUG:
        #some diagnostics about the rain, t2m, tmax, and tmin
        print("=== Input diagnostics for year", j, "===")

        for name, arr in [("rain", rain),
                      ("t2m", t2m),
                      ("tmax", tmax),
                      ("tmin", tmin)]:
            arr = np.asarray(arr)
            finite = np.isfinite(arr)
            print(name,
                  "dtype", arr.dtype,
                  "min", np.nanmin(arr),
                  "max", np.nanmax(arr),
                  "mean", np.nanmean(arr),
                  "nanfrac", np.mean(~finite))
            
            # also print percentiles to understand scale
            vals = arr[finite]
            print(name, "percentiles (50,90,99,99.9):",
                  np.percentile(vals, [50, 90, 99, 99.9]))
    #Now we comput BIOs
    (b1,b2,b3,b4,b4a,b5,b6,b7,b8,b9,
    b10,b11,b12,b13,b14,b15,b16,b17,
    b18,b19) = biovar(t2m,tmax,tmin,rain)
    if DEBUG:
        #Confirming whether biovar() produced NaNs/invalids
        for name, arr in [("b12", b12), ("b13", b13), ("b16", b16)]:
            arr = np.asarray(arr)
            print(name,
                  "dtype", arr.dtype,
                  "min", np.nanmin(arr),
                  "max", np.nanmax(arr),
                  "nanfrac", np.mean(~np.isfinite(arr)))
        for name, arr in [("b1", b1), ("b2", b2), ("b4", b4), ("b5", b5), ("b6", b6), ("b7", b7), ("b4a", b4a)]:
            arr = np.asarray(arr)
            print(name, "min", np.nanmin(arr), "max", np.nanmax(arr))
        
    #break

    #write to tif
    ll=[b1,b2,b3,b4,b4a,b5,b6,b7,b8,b9,
    b10,b11,b12,b13,b14,b15,b16,b17,
    b18,b19]
    for nn in range(len(ll)):
        if nn < 4:
            snn = str(nn+1)
        elif nn == 4:
            snn = '4a'
        else:
            snn=str(nn)
        fname=dest+'/Cote.bio'+snn+'.'+str(j)+'.tif'
        wrtbio(fname,ds,iy,ix,1,ll[nn])

'''
def bio1(tas,tasmax,tasmin,prec):
    #inputs are 12-month variables
    B1 = np.mean(tas,axis=0)
    B2 = np.mean(tasmax-tasmin, axis=0)
    B4 = np.std(tas,axis=0)
    B4a = B4/(B1+273.15)*100
    B5 = tasmax.max(axis=0)
    B6 = tasmin.min(axis=0)
    B7 = B5 - B6
    B3 = B2/B7*100
    B12 = np.sum(prec,axis=0)
    B13 = prec.max(axis=0)
    B14 = prec.min(axis=0)
    B15 = np.std(prec,axis=0)/(1+B12/12)*100
    return B1,B2,B3,B4,B4a,B5,B6,B7, \
                B12,B13,B14,B15

def bio2(tas,tasmax,tasmin,prec):
    #inputs are 14-month variables starting from January
    #and ending in February next year
    B8 = 
    return B8, B9, B10, B11, B16,B17,B18,B19
'''