import scipy.stats
import iris
import iris.plot as iplt
import matplotlib.pyplot as plt
import numpy as np
import iris.util
import os, sys
import pandas as pd
import cartopy.crs as ccrs

### 
### The PAMIP data used is freely available to download from the Earth System Grid Federation https://esgf.github.io/index.html. 
###

def read_full_ens(inmod,exp,invar,mon_anal_str,mon_anal,f_base):
  indir_base = "directory holding monthly mean PAMIP data" 
  inmod_dir = inmod
  cubelist = iris.load(indir_base+'/'+invar+'/'+exp+'/'+inmod_dir+'/'+invar+'*'+inmod+'*.nc')    
  nrunid =  len(cubelist)
  j=0  
  cube_mon = cubelist[j]  
  cube_sea = season_agg(cube_mon)
  cube_sea_tot = cube_sea.extract(iris.Constraint(clim_season=mon_anal_str))  
  for j in range(1, nrunid):   
    print('j = '+str(j))
    cube_mon = cubelist[j]  
    cube_sea_all = season_agg(cube_mon)
    cube_sea = cube_sea_all.extract(iris.Constraint(clim_season=mon_anal_str))
    cube_sea_tot.data = cube_sea_tot.data + cube_sea.data
    cube_sea_avg = cube_sea_tot/nrunid

  iris.save(cube_sea_avg, 'nc_files/'+pamip_exp_arr[i]+'_'+f_base+'.nc')
   

def season_agg(cube, seasons=None):
  ### Create seasonal means (unique and specified 3-month seasons)
  if (seasons == None): 
    seasons=['mam', 'jja', 'son', 'djf']

  if len(cube.coords('clim_season')) == 0:
    iris.coord_categorisation.add_season(cube, 'time', name='clim_season', seasons=seasons)

  if len(cube.coords('season_year')) == 0:
    iris.coord_categorisation.add_season_year(cube, 'time', name='season_year', seasons=seasons)

  # Keep only those times where we can produce seasonal means using exactly 3 months
  # (i.e., remove times from cubelist where only 1 or 2 times exist for that season)
  clim_seasons = cube.coords('clim_season')[0].points
  season_years = cube.coords('season_year')[0].points
  ntimes = len(cube.coords('time')[0].points)    
  keep_ind = np.zeros((0), dtype=int)
  for i in range(0,ntimes):
    ind = np.where( (clim_seasons == clim_seasons[i]) & (season_years == season_years[i]) )[0]
    n_months_in_season = len(clim_seasons[i]) # length of string, usually 3 (e.g., 'djfm' = 4)
    if (len(ind) == n_months_in_season): keep_ind = np.append(keep_ind,i)

  cube = cube[keep_ind]
  cube_out = cube.aggregated_by(['clim_season', 'season_year'], iris.analysis.MEAN)
  return cube_out


var_anal = sys.argv[1]
mon_anal_str = sys.argv[2]
read_data = int(sys.argv[3])
#modnum = int(sys.argv[3])
if mon_anal_str == 'djf': mon_anal = [1,2,12]
if mon_anal_str == 'mam': mon_anal = [3,4,5]
if mon_anal_str == 'jja': mon_anal = [6,7,8]
if mon_anal_str == 'son': mon_anal = [9,10,11]
pamip_exp_arr = ['pdSST-pdSIC', 'pdSST-futArcSIC']
modname_arr=['AWI-CM-1-1-MR', 'CanESM5', 'CESM2', 'CNRM-CM6-1', 'E3SM-1-0', 'EC-Earth3', 'FGOALS-f3-L', 'IPSL-CM6A-LR', 'NorESM2-LM','TaiESM1','HadGEM3-GC31-MM','MIROC6']

nmods = len(modname_arr)
if (nmods > 1): 
  modlabel = 'multi-model-mean'
else: 
  modlabel = modname_arr[0]    

if read_data == 1: 
  for k in range(0,nmods):     
    fname_base = var_anal+'_'+modname_arr[k]+'_'+mon_anal_str
    for i in range(0,2):
      print('i = '+str(i))      
      read_full_ens(modname_arr[k],pamip_exp_arr[i],var_anal,mon_anal_str,mon_anal,fname_base)  

####### Load fut and pd climatologies 
ref_grd = iris.load_cube('nc_files/pdSST-pdSIC_'+var_anal+'_CanESM5_'+mon_anal_str+'.nc')
cube_pdsic_var_tot = iris.cube.copy.deepcopy(ref_grd)
cube_futArcsic_var_tot = iris.cube.copy.deepcopy(ref_grd)

k = 0
fname_base = var_anal+'_'+modname_arr[k]+'_'+mon_anal_str
cube_pdsic_var = iris.load_cube('nc_files/pdSST-pdSIC_'+fname_base+'.nc')
cube_pdsic_var_regrid = cube_pdsic_var.regrid(ref_grd, iris.analysis.Linear()) 
cube_pdsic_var_tot.data = cube_pdsic_var_regrid.data
cube_futArcsic_var = iris.load_cube('nc_files/pdSST-futArcSIC_'+fname_base+'.nc')
cube_futArcsic_var_regrid = cube_futArcsic_var.regrid(ref_grd, iris.analysis.Linear())
cube_futArcsic_var_tot.data = cube_futArcsic_var_regrid.data
for k in range(1,nmods):
  fname_base = var_anal+'_'+modname_arr[k]+'_'+mon_anal_str
  cube_pdsic_var = iris.load_cube('nc_files/pdSST-pdSIC_'+fname_base+'.nc')
  cube_pdsic_var_regrid = cube_pdsic_var.regrid(ref_grd, iris.analysis.Linear()) 
  cube_pdsic_var_tot.data = cube_pdsic_var_tot.data + cube_pdsic_var_regrid.data
  cube_futArcsic_var = iris.load_cube('nc_files/pdSST-futArcSIC_'+fname_base+'.nc')
  cube_futArcsic_var_regrid = cube_futArcsic_var.regrid(ref_grd, iris.analysis.Linear())
  cube_futArcsic_var_tot.data = cube_futArcsic_var_tot.data + cube_futArcsic_var_regrid.data

fname_save = var_anal+'_'+modlabel+'_'+mon_anal_str
cube_pdsic_var_mmm = cube_pdsic_var_tot/nmods
cube_futArcsic_var_mmm = cube_futArcsic_var_tot/nmods
iris.save(cube_pdsic_var_mmm, 'nc_files/'+pamip_exp_arr[0]+'_'+fname_save+'.nc')
iris.save(cube_futArcsic_var_mmm, 'nc_files/'+pamip_exp_arr[1]+'_'+fname_save+'.nc')

if var_anal == 'pr': 
### Change units to mm/day ###
  cube_pdsic_var_mmm = cube_pdsic_var_mmm*86400
  cube_futArcsic_var_mmm = cube_futArcsic_var_mmm*86400

###### Take differences
cube_diffsic_var_mmm = cube_futArcsic_var_mmm-cube_pdsic_var_mmm
fname_diff = pamip_exp_arr[1]+'_minus_'+pamip_exp_arr[0]+'_'+fname_save
iris.save(cube_diffsic_var_mmm, 'nc_files/'+fname_diff+'.nc')

###### Make plots
if var_anal == 'pr': 
  var_unit='mm / day'
  levels = (np.arange(13) - 6)*.05
  var_tit = 'Precipitation rate' 
  title_txt = var_anal+'_'+modlabel+'_'+mon_anal_str

if var_anal == 'ua':
  var_unit='m / s'
  var_tit = 'Westerly wind at 850 hPa'
  levels = (np.arange(13) - 6)*.15
  title_txt = var_anal+'_'+modlabel+'_'+mon_anal_str
  
if var_anal == 'tas':
  var_unit='K'
  var_tit = 'Surface air temperature' 
  levels = (np.arange(13) - 6)*.05
  title_txt = var_anal+'_'+modlabel+'_'+mon_anal_str
  
brewer_cmap = plt.get_cmap('brewer_RdBu_11')
ax=plt.figure(figsize=(6, 6), dpi=300)
ax=plt.subplot(1,1,1, projection=ccrs.PlateCarree())
ax.set_extent([-50, 30, 20, 85], ccrs.PlateCarree())
ax.coastlines()
plt.title(var_tit)
cf=iplt.contourf(cube_diffsic_var_mmm, cmap=brewer_cmap, levels=levels, extend = "both")
colorbar = plt.colorbar(cf, orientation='horizontal')
colorbar.set_label(var_unit, fontsize=8)
plt.show()
#plt.savefig('plots/'+fname_diff+'.png')
#plt.close()