import numpy as np
import xarray as xr
import sys
import os
import glob
import argparse

def GetLandMask(filetemp) :
    gr=xr.open_dataset(filetemp.replace("XXXX","RadT"))
    mask=gr["RadT"].data[0,:,:]
    mask[~np.isnan(mask)]=1.0
    gr.close()
    return mask

def UpdateMaskInFile(var, infile, outfile, mask) :
    grin=xr.open_dataset(infile)
    grin[var].data=grin[var].data*mask
    grin.to_netcdf(outfile, mode='w')
    grin.close()
    return mask

    return
#
####################################################################"
#

FinalDir="/bdd/MEDI/workspaces/polcher/LIAISEoutputs/UKMO/"

parser = argparse.ArgumentParser(description='Corrects the land/sea mask in Rainf and Snowf and also borrows potnetial evaporation.')
parser.add_argument('--experiment', "-e", nargs="*", type=str, default=["WFDE5_CRU_GPCC"], \
                    help="Experiment file to be processed.")
args = parser.parse_args()

#
# Clean-up arguments received.
#
if args.experiment[0].find(".def") >= 0 :
    exp=args.experiment[0].replace(".def","")
else :
    exp=args.experiment[0]

    
DataLoc=glob.glob(FinalDir+"/"+exp+"/TS_DA/*RadT.nc")[0].replace("RadT","XXXX")
landmask=GetLandMask(DataLoc)

if exp.find("WFDE") >= 0 :
    ORCH="/bdd/MEDI/workspaces/polcher/IGCM_OUT/OL2/PROD/LIAISE/O2-LIAISE-WFDE5/SRF/Analyse/TS_DA/O2-LIAISE-WFDE5_19890101_20131231_1D_PotEvap.nc"
elif exp.find("IPSL_Alt") >= 0 :
    ORCH="/bdd/MEDI/workspaces/polcher/IGCM_OUT/OL2/PROD/LIAISE/O2-LIAISE-IPSLalt2/SRF/Analyse/TS_DA/O2-LIAISE-IPSLalt2_19890101_20131231_1D_PotEvap.nc"
elif exp.find("IPSL_Avg") >= 0 :
    ORCH="/bdd/MEDI/workspaces/polcher/IGCM_OUT/OL2/PROD/LIAISE/O2-LIAISE-IPSLavg2/SRF/Analyse/TS_DA/O2-LIAISE-IPSLavg2_19890101_20131231_1D_PotEvap.nc"
elif exp.find("ETHZ_Avg") >= 0 :
    ORCH="/bdd/MEDI/workspaces/polcher/IGCM_OUT/OL2/PROD/LIAISE/O2-LIAISE-ETHZavg2/SRF/Analyse/TS_DA/O2-LIAISE-ETHZavg2_19890101_20131231_1D_PotEvap.nc"
else :
    sys.exit("Unknown experiment")

pet=DataLoc.replace("XXXX.nc","PotEvap.nc")
print("Using as substitut for PostEvap :", ORCH)
print("to produce file ", pet)
UpdateMaskInFile("PotEvap", ORCH, pet, landmask)