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"}, \
               "NOAH":{"LAT":"lon","LAT":"lat"}}

ModelDimNames={"NOAH":{"west_east":"lon","south_north":"lat"}}

FillValue=-1.0e+10

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"
noahfile="etzh/qrain_1989.nc"

c=xr.open_dataset(cnrmfile).rename(name_dict=ModelVarNames["CNRM"]).sel(time=slice('1989-02-01','1989-03-31'))
j=xr.open_dataset(julesfile).sel(time=slice('1989-02-01','1989-03-31'))
n=xr.open_dataset(noahfile).rename_dims(dims_dict=ModelDimNames["NOAH"]).rename(name_dict=ModelVarNames["NOAH"]).sel(time=slice('1989-02-01','1989-03-31'))

xr.where(n["Rainf"] < FillValue, np.nan, n["Rainf"])
xr.where(j["Rainf"] < FillValue, np.nan, j["Rainf"])

pdfhdl = PdfPages("ComparisonRainf.pdf")

c["Rainf"].isel(lon=[110],lat=[140]).plot(label="CNRM")
j["Rainf"].isel(x=[110],y=[140]).plot(label="JULES")
n["Rainf"].isel(lon=[110],lat=[140]).plot(label="NOAH")
plt.legend()

pdfhdl.savefig()
plt.close()


d=(c["Rainf"].mean(dim="time",skipna=True)-n["Rainf"].mean(dim="time",skipna=True))*86400


d.plot()
plt.title("Rainfall difference CNRM-NOAH [mm/d]")
pdfhdl.savefig()
plt.close()

pdfhdl.close()