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

import datetime
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages

import argparse

ModelVarNames={"ORCH":{"time_counter":"time"}, \
               "JULES":{"time":"time","longitude":"lon","latitude":"lat"}, \
               "CNRM":{"time_counter":"time"}}


cnrmfile="../CNRM/ETHZ_Avg/TS_DA/CNRM_ETHZ_Avg_19890101_20131231_1D_Rainf.nc"
julesfile="../UKMO/ETHZ_Avg/TS_DA/JULES_ETHZ_Avg_19890101_20131231_1D_Rainf.nc"
ecmwffile="osm_o_wat_19900101_1_Rainf_daily_latinverted.nc"

##x=xr.open_dataset(julesfile)
##print(x)

dates=['1990-01-01','1990-01-31']
li=[110, 140]
pdf=True
FillValue=1.0e+20

c=xr.open_dataset(cnrmfile).rename(name_dict=ModelVarNames["CNRM"]).sel(time=slice(dates[0],dates[1]))
j=xr.open_dataset(julesfile).rename(name_dict=ModelVarNames["JULES"]).sel(time=slice(dates[0],dates[1]))
n=xr.open_dataset(ecmwffile).sel(time=slice(dates[0],dates[1]))


LOC=[c["lon"].data[li[0]], c["lat"].data[li[1]]]

if pdf :
    plt.switch_backend('agg')
    pdfhdl = PdfPages("ComparisonRainf.pdf")

c["Rainf"].sel(lon=LOC[0],lat=LOC[1], method="nearest").plot(label="CNRM")
j["Rainf"].isel(x=li[0],y=li[1]).plot(label="JULES")
n["Rainf"].sel(lon=LOC[0],lat=LOC[1], method="nearest").plot(label="ECMWF")
plt.legend()

if pdf :
    pdfhdl.savefig()
    plt.close()
else :
    plt.show()

print("XXXXX TS plot done ")

cm=c["Rainf"].mean(dim="time",skipna=True)
nm=n["Rainf"].mean(dim="time",skipna=True)
jm=j["Rainf"].mean(dim="time",skipna=True)

jm.data[jm.data > FillValue]=np.nan

cn=(cm.data-nm.data)*86400
print("Difference range (CNRM - ECMWF) : ", np.nanmin(cn), np.nanmax(cn))
jn=(jm.data-nm.data)*86400
print("Difference range (JULES - ECMWF) : ", np.nanmin(jn), np.nanmax(jn))

fig, ax = plt.subplots()
im=ax.pcolormesh(cm["lon"].data, cm["lat"].data, cn)
fig.colorbar(im, ax=ax)
ax.set_title("Rainfall difference (CNRM - ECMWF) [mm/d]")

if pdf :
    pdfhdl.savefig()
    plt.close()
else :
    plt.show()

fig, ax = plt.subplots()
im=ax.pcolormesh(cm["lon"].data, cm["lat"].data, jn)
fig.colorbar(im, ax=ax)
ax.set_title("Rainfall difference (JULES - ECMWF) [mm/d]")

if pdf :
    pdfhdl.savefig()
    plt.close()
else :
    plt.show()

    
if pdf :
    pdfhdl.close()