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

# 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 + NoData + compression); keep old above unchanged ---
def wrtbio_float32(fname, ds, rows, cols, layer, arr_out):
    driver = gdal.GetDriverByName("GTiff")
    outdata = driver.Create(
        fname, cols, rows, 1, gdal.GDT_Float32,
        options=["TILED=YES", "COMPRESS=DEFLATE", "PREDICTOR=2"]
    )
    outdata.SetGeoTransform(ds.GetGeoTransform())
    outdata.SetProjection(ds.GetProjection())
    band = outdata.GetRasterBand(1)
    nodata = -9999.0
    band.SetNoDataValue(nodata)

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

    band.FlushCache()
    outdata.FlushCache()
    outdata = None

for j in range(1981,2021,10):
    j1=j+9
    if j==2011:
        j1=j+6
    decd=str(j)+'-'+str(j1)
    dest='/D2/CHELSA/Cote/decade-mean/'+decd
    os.system('mkdir -p '+dest)
    for nn in range(20):
        if nn < 4:
            snn = str(nn+1)
        if nn == 4:
            snn = str(nn)+'a'
        if nn>4:
            snn=str(nn)

        fname='Cote.bio'+snn+'.'+str(j)+'.tif'
        ds0 = gdal.Open(r"/D2/CHELSA/Cote/yearly/bio/"+str(j)+'/'+fname)

        # Optional safety: handle missing files
        if ds0 is None:
            print("Missing:", "/D2/CHELSA/Cote/yearly/bio/"+str(j)+'/'+fname)
            continue

        band = ds0.GetRasterBand(1)
        tas0 = band.ReadAsArray()
        tas0 = tas0.astype(float)

        ix = ds0.RasterXSize   # cols
        iy = ds0.RasterYSize   # rows

        for jj in range(j+1,j1+1):
            fname='Cote.bio'+snn+'.'+str(jj)+'.tif'
            ds = gdal.Open(r"/D2/CHELSA/Cote/yearly/bio/"+str(jj)+'/'+fname)
            if ds is None:
                print("Missing:", "/D2/CHELSA/Cote/yearly/bio/"+str(jj)+'/'+fname)
                continue
            band = ds.GetRasterBand(1)
            tas = band.ReadAsArray()
            tas0 = tas0+tas
            del ds

        tas0=tas0/(j1-j+1)

        # OLD (breaks fractional BIOs)
        # tas0=np.rint(tas0).astype(int)

        fname=dest+'/Cote.bio'+snn+'.'+decd+'.tif'

        # BUG: your original call swapped rows/cols
        # wrtbio(fname,ds0,ix,iy,1,tas0)

        # Corrected: rows=iy, cols=ix
        # If you insist on keeping UInt16, use this line (not recommended):
        # wrtbio(fname, ds0, iy, ix, 1, np.rint(tas0).astype(np.uint16))

        # Recommended: Float32 output preserves decimals and avoids overflow
        wrtbio_float32(fname, ds0, iy, ix, 1, tas0)
