import datetime, numpy as np, netCDF4, sys
from grid import Grid
from cwt import cwt_kg31to3
from basics import printStat, writePFTMAP
debug = False

path_orc = "/home/surface5/vbastri/BDD/ESACCI-LC/ORCHIDEE_NEW/0.25/PFTmap_2010.nc"
path_luh = "/home/surface5/vbastri/BDD/LUH2/GCB2022/states.nc"
path_kg  = "/home/surface5/vbastri/BDD/KG2017/KG_720x1440.nc"
path_c4  = "/home/surface5/vbastri/BDD/C4FRAC/C4frac_720x1440.nc"
path_wb  = "/home/surface5/vbastri/BDD/ESACCI-LC/WATER/WB_720x1440.nc"
path_out = "netcdf/PFTmap_%(year)s.nc"
luh_year0 = 850
merge_year = 2010
build_years = np.arange(850, 2022+1)

atts = dict(description = "ORCHIDEE PFT map derived from ESACCI Land Cover map v2.0.7b for 2010, aggregated using user-tool v4.3, cross-walking table v2.4, C3/C4 for grasses from C.Still (May 2018), C3/C4 for crops from LUH2-GCB2022, followed by historical reconstrution using G. Hurtt data LUH2-GCB2022 with merging rules v2.0",
            contact = "V.Bastrikov (vladislav.bastrikov@lsce.ipsl.fr), P.Peylin (peylin@lsce.ipsl.fr)",
            web = "https://orchidas.lsce.ipsl.fr/dev/lccci/",
            date = datetime.datetime.now().strftime("%d/%m/%Y %H:%M"))

# ORCHIDEE PFT indices
BS    = 0
TREBF = 1
TRDBF = 2
TEENF = 3
TEEBF = 4
TEDBF = 5
BOENF = 6
BODBF = 7
BODNF = 8
TEGR  = 9
C4GR  = 10
C3CR  = 11
C4CR  = 12
TRGR  = 13
BOGR  = 14

# LUH class indices
PRIMF = 0
PRIMN = 1
SECDF = 2
SECDN = 3
URBAN = 4
C3ANN = 5
C4ANN = 6
C3PER = 7
C4PER = 8
C3NFX = 9
PASTR = 10
RANGE = 11
ICWTR = 12
LUH_LIST = ["primf", "primn", "secdf", "secdn", "urban", "c3ann", "c4ann", "c3per", "c4per", "c3nfx", "pastr", "range"]

# Function to safely divide two streams: if denominator equals zero, division result is set to zero
# During division denominator zeros are changed to units, so that python does not warn on errors
def divide(num, den): return np.where(den == 0, 0, num / np.where(den == 0, 1, den))

# Function to make plots for data streams
def plot(**kwargs):
    for name, data in kwargs.items():
        for idx, item in enumerate(data if len(data.shape)>2 else [data]):
            grid.plotmap(np.ma.masked_where(item==0, item), dpi=150, vmin=0, vmax=1,
                         filename="%s%s.png" % (name.lower(), idx if len(data.shape)>2 else ""), title="%s%s" % (name.upper(), idx if len(data.shape)>2 else ""),
                         info="min=%.3g\nmax=%.3g\nmean=%.3g\nnzpts=%s" % (item.min(), item.max(), grid.meanmask(np.ma.masked_where(item==0, item)), (item>0).sum()))

# Function to extend land by 1 pixel outside land-sea mask
def extend(veg):
    newveg = np.ma.masked_all((len(veg), len(lat), len(lon)))
    loc_area = np.where(veg[0].mask == False, grid.area, 0)
    wveg = veg * loc_area
    for ilat in range(len(lat)):
        for ilon in range(len(lon)):
            if veg[0,ilat,ilon] is not np.ma.masked:
                newveg[:,ilat,ilon] = veg[:,ilat,ilon]
            else:
                if wveg[:, max(0, ilat-1) : ilat+2, max(0, ilon-1) : ilon+2].count() == 0: continue
                newveg[:,ilat,ilon] = wveg[:, max(0, ilat-1) : ilat+2, max(0, ilon-1) : ilon+2].sum(axis=-1).sum(axis=-1) / loc_area[max(0, ilat-1) : ilat+2, max(0, ilon-1) : ilon+2].sum()
    return newveg

# Read present-day ORCHIDEE map to be merged with LUH
print("\nREAD ORCHIDEE")
nc = netCDF4.Dataset(path_orc)
lat = nc.variables["lat"][:]
lon = nc.variables["lon"][:]
orc = nc.variables["maxvegetfrac"][0].astype("f8").filled(0.)
nc.close()
printStat(orc=orc)
orc = divide(orc, orc.sum(axis=0))
printStat(orcn=orc)
grid = Grid(lat=lat, lon=lon)
if debug: plot(orc=orc, orcsum=orc.sum(axis=0))

# Read Koeppen-Geiger map with 31 classes
print("\nREAD KG")
nc = netCDF4.Dataset(path_kg)
if not np.all(lat == nc.variables["lat"][:]) or not np.all(lon == nc.variables["lon"][:]): print("Different grid"); sys.exit()
kg31 = nc.variables["kg"][:].astype("f8").filled(0)
nc.close()
printStat(kg31=kg31)
kg31 = divide(kg31, kg31.sum(axis=0))
printStat(kg31n=kg31)

# Aggregate KG into 3 groups
kg = np.zeros((len(cwt_kg31to3[0])-1, len(lat), len(lon)))
for idx, rule in enumerate(cwt_kg31to3[1:]): kg = kg + kg31[idx] * np.array(rule[1:])[:,None,None]
printStat(kg3=kg)
if debug: plot(kg=kg)
# Fill KG where undefined with the dominant class across each latitude (or last class if no data in a latitude)
domclass = np.where(kg.sum(axis=(0,2)) == 0, len(kg)-1, kg.sum(axis=2).argmax(axis=0))
for idx in range(len(lat)): kg[domclass[idx], idx, kg[:,idx,:].sum(axis=0) == 0] = 1
printStat(kg3f=kg)
if debug: plot(kgf=kg)

print("\nREAD C4")
nc = netCDF4.Dataset(path_c4)
if not np.all(lat == nc.variables["lat"][:]) or not np.all(lon == nc.variables["lon"][:]): print("Different grid"); sys.exit()
c4 = nc.variables["C4frac"][:].astype("f8").filled(0)
nc.close()
printStat(c4=c4)
if debug: plot(c4=c4)

print("\nREAD WB")
nc = netCDF4.Dataset(path_wb)
if not np.all(lat == nc.variables["lat"][:]) or not np.all(lon == nc.variables["lon"][:]): print("Different grid"); sys.exit()
wb = nc.variables["wb"][:].astype("f8").filled(0)
nc.close()
printStat(wb=wb)
lsm = np.where(wb[1] > 0.001, 1, 0)
print("land pts=%s" % lsm.sum())
if debug: plot(wb=wb, lsm=lsm)

# Read LUH for the merging year
print("\nREAD LUH")
nc = netCDF4.Dataset(path_luh)
if not np.all(lat == nc.variables["lat"][:]) or not np.all(lon == nc.variables["lon"][:]): print("Different grid"); sys.exit()
luh = np.zeros((len(LUH_LIST)+1, len(lat), len(lon))) # array for all LUH maps plus ICWTR
for idx, item in enumerate(LUH_LIST): luh[idx] = nc.variables[item][merge_year-luh_year0].astype("f8") # do not fill with 0 to calculate icwtr over land (ice and water)
luh[ICWTR] = np.clip(1 - luh.sum(axis=0), 0, 1)
luh = np.ma.masked_invalid(luh).filled(0)
luh = divide(luh, luh.sum(axis=0))
nc.close()
printStat(luh=luh)
if debug: plot(luh=luh, luhsum=luh.sum(axis=0))

onlyorc = np.where((orc.sum(axis=0) > 0) & (luh.sum(axis=0) == 0), 1, 0)
onlyluh = np.where((orc.sum(axis=0) == 0) & (luh.sum(axis=0) > 0), 1, 0)
print("ORC>0 LUH=0", onlyorc.sum())
print("ORC=0 LUH>0", onlyluh.sum())
if debug: plot(onlyorc=onlyorc, onlyluh=onlyluh)

print("\nBUILD FRACTIONS")
luh_C3  = luh[C3ANN] + luh[C3PER] + luh[C3NFX]
luh_C4  = luh[C4ANN] + luh[C4PER]
luh_ant = luh_C3 + luh_C4 + luh[PASTR]

orc_gr = orc[TEGR] + orc[TRGR] + orc[BOGR] + orc[C4GR]
orc_cr = orc[C3CR] + orc[C4CR]
gr_ant = np.clip(luh_ant - orc_cr, 0, np.inf)
gr_nat = np.clip(orc_gr  - gr_ant, 0, np.inf)
gr_tr = gr_nat * divide(orc[TRGR], orc_gr)
gr_te = gr_nat * divide(orc[TEGR], orc_gr)
gr_bo = gr_nat * divide(orc[BOGR], orc_gr)
gr_C4 = gr_nat * divide(orc[C4GR], orc_gr)
total_nat = orc[BS] + orc[TREBF] + orc[TRDBF] + orc[TEENF] + orc[TEEBF] + orc[TEDBF] + orc[BOENF] + orc[BODBF] + orc[BODNF] + gr_tr + gr_te + gr_bo + gr_C4

frac = divide(orc, total_nat)
frac[TRGR] = divide(gr_tr, total_nat)
frac[TEGR] = divide(gr_te, total_nat)
frac[BOGR] = divide(gr_bo, total_nat)
frac[C4GR] = divide(gr_C4, total_nat)
frac[C3CR] = 0
frac[C4CR] = 0
printStat(frac=frac)
print("ORC>0 FRAC=0", ((orc.sum(axis=0) > 0) & (frac.sum(axis=0) == 0)).sum()) # these are pixels with crops in ORCHIDEE (161)
print("LUH>0 FRAC=0", ((luh.sum(axis=0) > 0) & (frac.sum(axis=0) == 0)).sum()) # these are mostly ICWTR pixels in LUH, few as PRIMN (361)
print("ORC>0 LUH>0 FRAC=0", ((orc.sum(axis=0) > 0) & (luh.sum(axis=0) > 0) & (frac.sum(axis=0) == 0)).sum()) # 40
if debug: plot(frac=frac, fracsum=frac.sum(axis=0))

nc = netCDF4.Dataset(path_luh)
for year in build_years:
    print("\nREAD LUH FOR YEAR %s" % year)
    yluh = np.zeros((len(LUH_LIST)+1, len(lat), len(lon)))
    for idx, item in enumerate(LUH_LIST): yluh[idx] = nc.variables[item][year-luh_year0].astype("f8")
    yluh[ICWTR] = np.clip(1 - yluh.sum(axis=0), 0, 1)
    yluh = np.ma.masked_invalid(yluh).filled(0)
    yluh = divide(yluh, yluh.sum(axis=0))
    printStat(yluh=yluh)

    yluh_nat = yluh[PRIMF] + yluh[PRIMN] + yluh[SECDF] + yluh[SECDN] + yluh[RANGE] + yluh[URBAN] + yluh[ICWTR]
    yluh_C3  = yluh[C3ANN] + yluh[C3PER] + yluh[C3NFX]
    yluh_C4  = yluh[C4ANN] + yluh[C4PER]
    
    print("\nBUILD ORC FOR YEAR %s" % year)    
    yorc = frac * yluh_nat
    yorc[TRGR] = yorc[TRGR] + kg[0] * (1 - c4) * yluh[PASTR]
    yorc[TEGR] = yorc[TEGR] + kg[1] * (1 - c4) * yluh[PASTR]
    yorc[BOGR] = yorc[BOGR] + kg[2] * (1 - c4) * yluh[PASTR]
    yorc[C4GR] = yorc[C4GR] + c4 * yluh[PASTR]
    yorc[C3CR] = yluh_C3
    yorc[C4CR] = yluh_C4
    printStat(yorc=yorc)
    print("0<YORC<0.999", ((yorc.sum(axis=0)>0) & (yorc.sum(axis=0)<0.999)).sum())
    print("ORC>0 LUH>0 YORC<0.999", ((orc.sum(axis=0)>0) & (luh.sum(axis=0)>0) & (yorc.sum(axis=0)<0.999)).sum())
    print("ORC>0 LUH=0 YORC=0", ((orc.sum(axis=0)>0) & (luh.sum(axis=0)==0) & (yorc.sum(axis=0)==0)).sum())
    print("ORC>0 YORC<0.999", ((orc.sum(axis=0)>0) & (yorc.sum(axis=0)<0.999)).sum())
    if debug: plot(yorc=yorc)
    
    # POST PROCESS
    # - drop pixels where ORC has no info
    # - use data from ORC where LUH has no info or problematic balance
    # - drop fractions below 0.001
    # - normalize to 1
    # - if OCEAN>0.999 apply mask
    # - filter out Antarctica (lat < S60)
    yorc = np.where((orc.sum(axis=0)==0), 0, yorc)
    yorc = np.where((orc.sum(axis=0)>0) & (yorc.sum(axis=0)<0.999), orc, yorc)
    yorc = np.where(yorc < 0.001, 0, yorc)
    yorc = divide(yorc, yorc.sum(axis=0))
    yorc = np.ma.where(wb[0] < 0.999, yorc, np.ma.masked)
    yorc[:, lat < -60, :] = np.ma.masked
    printStat(orcn=yorc)
    if debug: plot(yorcn=yorc, yorcnsum=yorc.sum(axis=0))

    # Check that resulting LSM is identical to ORC original
    m1 = np.where(orc.sum(axis=0)>0, 1, 0)
    m2 = np.where(yorc.sum(axis=0)>0, 1, 0)
    print(np.all(m1==m2))
    
    # Extend by 1 pixel
    yorc = extend(yorc)
    printStat(orcext=yorc)
    
    writePFTMAP(yorc, lat=lat, lon=lon, year=year, filename=path_out % dict(year=year), atts=atts)

nc.close()