#!/usr/bin/env python
# -*- coding: utf-8 -*-
import string
import numpy as np
import matplotlib.pyplot as plt
from pylab import *
from mpl_toolkits.basemap import Basemap
from mpl_toolkits.basemap import shiftgrid, cm
from netCDF4 import Dataset
import arctic_map # function to regrid coast limits
import cartesian_grid_test # function to convert grid from polar to cartesian
import scipy.special
import ffgrid2
import map_ffgrid
from matplotlib import colors
from matplotlib.font_manager import FontProperties
import map_cartesian_grid


###############################
# time period characteristics #
###############################
MONTH = np.array(['01', '02', '03', '04', '05', '06', '07', '08', '09', '10', '11', '12'])
month = np.array(['JANUARY', 'FEBRUARY', 'MARCH', 'APRIL', 'MAY', 'JUNE', 'JULY', 'AUGUST', 'SEPTEMBER', 'OCTOBER', 'NOVEMBER', 'DECEMBER'])
month_day = np.array([31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31])
M = len(month)


########################
# grid characteristics #
########################
x0 = -3000. # min limit of grid
x1 = 2500. # max limit of grid
dx = 40.
xvec = np.arange(x0, x1+dx, dx)
nx = len(xvec) 
y0 = -3000. # min limit of grid
y1 = 3000. # max limit of grid
dy = 40.
yvec = np.arange(y0, y1+dy, dy)
ny = len(yvec)




############################################
# time evolution (monthly) in a given zone #
############################################
#zone 1 (seasonal ice) : yi = 960. // yf = 1360. // xi = -680. // xf = -320. (North Beaufort Sea)
#zone 2 (multiyear ice) : yi = 320. // yf = 720. // xi = -1080. // xf = -720. (North Canadian Archipelago)
#zone 3 (young ice) : yi = 1880. // yf = 2280. // xi = -480. // xf = -120. (Chukchi Sea)

# select borders of zone
yi = 320. 
yf = 720.
xi = -1080.
xf = -720.

#find corresponding index in xvec and yvec
xxi = np.where(xvec == xi)[0][0]
xxf = np.where(xvec == xf)[0][0]
yyi = np.where(yvec == yi)[0][0]
yyf = np.where(yvec == yf)[0][0]

len(xvec[xxi:xxf+1])
len(yvec[yyi:yyf+1])


mean_grad_ratio_zone = np.zeros([M, 31], float) # 2D-array of zonal mean gradient ratio for each day in each month 
std_grad_ratio_zone = np.zeros([M, 31], float) # 2D-array of zonal std of gradient ratio for each day in each month 

mean_spec_anom_23_zone = np.zeros([M, 31], float) # 2D-array of zonal mean emis_spec spatial anomaly for each day in each month (at 23GHz)
std_spec_anom_23_zone = np.zeros([M, 31], float) # 2D-array of zonal std of emis_spec spatial anomaly for each day in each month (at 23GHz)

mean_spec_anom_89_zone = np.zeros([M, 31], float ) # 2D-array of zonal mean emis_spec spatial anomaly for each day in each month (at 89GHz)
std_spec_anom_89_zone = np.zeros([M, 31], float) # 2D-array of zonal std of emis_spec spatial anomaly for each day in each month (at 89GHz)

mean_ratio_anom_89_zone = np.zeros([M, 31], float) # 2D-array of zonal mean emis_ratio spatial anomaly for each day in each month (at 89GHz)
std_ratio_anom_89_zone = np.zeros([M, 31], float) # 2D-array of zonal std of emis_ratio spatial anomaly for each day in each month (at 89GHz)

S = np.zeros([M, 31], float) # number of data pixels in selected zone, per day and per month

for imo in range (0, M): 
    # daily read gradient ratio file
    print 'read daily gradient ratio for month ' + month[imo]
    fichier_grad_ratio = Dataset('/net/argos/data/parvati/lahlod/ARCTIC/monthly_GLACE/gridded_data/cartesian_grid/res_40/cartesian_grid_daily_grad_ratio_spec_23-89_near_nadir_AMSUA_' + month[imo] + '2009.nc', 'r', format = 'NETCDF3_CLASSIC')
    gr = fichier_grad_ratio.variables['grad_ratio'][:]
    fichier_grad_ratio.close()
    # read daily emis anomaly file for 23GHz
    print 'read daily emis anomaly 23GHz for month ' + month[imo]
    fichier_anom23 = Dataset('/net/argos/data/parvati/lahlod/ARCTIC/monthly_GLACE/gridded_data/cartesian_grid/res_40/cartesian_grid_daily_data_lamb_spec_ratio_anomaly_near_nadir_AMSUA23_' + month[imo] + '2009.nc', 'r', format = 'NETCDF3_CLASSIC')
    sa_23 = fichier_anom23.variables['spec_anomaly'][:]
    fichier_anom23.close()
    # read daily emis anomaly file for 89GHz
    print 'read daily emis anomaly 89GHz for month ' + month[imo]
    fichier_anom89 = Dataset('/net/argos/data/parvati/lahlod/ARCTIC/monthly_GLACE/gridded_data/cartesian_grid/res_40/cartesian_grid_daily_data_lamb_spec_ratio_anomaly_near_nadir_AMSUA89_' + month[imo] + '2009.nc', 'r', format = 'NETCDF3_CLASSIC')
    sa_89 = fichier_anom89.variables['spec_anomaly'][:]
    ra_89 = fichier_anom89.variables['ratio_anomaly'][:]
    fichier_anom89.close()
    grad_ratio_vec = np.zeros([month_day[imo], len(xvec[xxi : xxf+1]) * len(yvec[yyi : yyf+1])], float)
    spec_anom_23_vec = np.zeros([month_day[imo], len(xvec[xxi : xxf+1]) * len(yvec[yyi : yyf+1])], float)
    spec_anom_89_vec = np.zeros([month_day[imo], len(xvec[xxi : xxf+1]) * len(yvec[yyi : yyf+1])], float)
    ratio_anom_89_vec = np.zeros([month_day[imo], len(xvec[xxi : xxf+1]) * len(yvec[yyi : yyf+1])], float)
    print 'calculate daily mean and std zonal gradient ratio'
    for ijr in range (0, month_day[imo]):
        print 'date ' + str(ijr+1) + ' ' + month[imo] + ' 2009'
        S[imo, ijr] = shape(gr[ijr, yyi : yyf + 1, xxi : xxf + 1][nonzero(isnan(gr[ijr, yyi : yyf + 1, xxi : xxf + 1]) == False)])[0]
        # gradient ratio
        grad_ratio_vec[ijr, :] = reshape(gr[ijr, yyi : yyf + 1, xxi : xxf + 1], size(gr[ijr, yyi : yyf + 1, xxi : xxf + 1]))
        mean_grad_ratio_zone[imo, ijr] = mean(grad_ratio_vec[ijr, :][nonzero(isnan(grad_ratio_vec[ijr, :]) == False)])
        std_grad_ratio_zone[imo, ijr] = sqrt((1. / (size(gr[ijr, yyi : yyf+1, xxi : xxf+1]) - 1.)) * sum((grad_ratio_vec[ijr, :][nonzero(isnan(grad_ratio_vec[ijr, :]) == False)] - mean_grad_ratio_zone[imo, ijr])**2))
        # spec anomaly 23GHz
        spec_anom_23_vec[ijr, :] = reshape(sa_23[ijr, yyi : yyf + 1, xxi : xxf + 1], size(sa_23[ijr, yyi : yyf + 1, xxi : xxf + 1]))
        mean_spec_anom_23_zone[imo, ijr] = mean(spec_anom_23_vec[ijr, :][nonzero(isnan(spec_anom_23_vec[ijr, :]) == False)])
        std_spec_anom_23_zone[imo, ijr] = sqrt((1. / (size(sa_23[ijr, yyi : yyf+1, xxi : xxf+1]) - 1.)) * sum((spec_anom_23_vec[ijr, :][nonzero(isnan(spec_anom_23_vec[ijr, :]) == False)] - mean_spec_anom_23_zone[imo, ijr])**2))
        # spec anomaly 89GHz
        spec_anom_89_vec[ijr, :] = reshape(sa_89[ijr, yyi : yyf + 1, xxi : xxf + 1], size(sa_89[ijr, yyi : yyf + 1, xxi : xxf + 1]))
        mean_spec_anom_89_zone[imo, ijr] = mean(spec_anom_89_vec[ijr, :][nonzero(isnan(spec_anom_89_vec[ijr, :]) == False)])
        std_spec_anom_89_zone[imo, ijr] = sqrt((1. / (size(sa_89[ijr, yyi : yyf+1, xxi : xxf+1]) - 1.)) * sum((spec_anom_89_vec[ijr, :][nonzero(isnan(spec_anom_89_vec[ijr, :]) == False)] - mean_spec_anom_89_zone[imo, ijr])**2))
        # ratio anomaly 89GHz
        ratio_anom_89_vec[ijr, :] = reshape(ra_89[ijr, yyi : yyf + 1, xxi : xxf + 1], size(ra_89[ijr, yyi : yyf + 1, xxi : xxf + 1]))
        mean_ratio_anom_89_zone[imo, ijr] = mean(ratio_anom_89_vec[ijr, :][nonzero(isnan(ratio_anom_89_vec[ijr, :]) == False)])
        std_ratio_anom_89_zone[imo, ijr] = sqrt((1. / (size(ra_89[ijr, yyi : yyf+1, xxi : xxf+1]) - 1.)) * sum((ratio_anom_89_vec[ijr, :][nonzero(isnan(ratio_anom_89_vec[ijr, :]) == False)] - mean_ratio_anom_89_zone[imo, ijr])**2))
        if (isnan(mean_grad_ratio_zone[imo, ijr]) == True):
            std_grad_ratio_zone[imo, ijr] = NaN
        if (isnan(mean_spec_anom_23_zone[imo, ijr]) == True):
            std_spec_anom_23_zone[imo, ijr] = NaN
        if (isnan(mean_spec_anom_89_zone[imo, ijr]) == True):
            std_spec_anom_89_zone[imo, ijr] = NaN
        if (isnan(mean_ratio_anom_89_zone[imo, ijr]) == True):
            std_ratio_anom_89_zone[imo, ijr] = NaN



# append daily zonal gradient ratio for study of the whole year 2009
mean_year_grad_ratio = np.append(mean_grad_ratio_zone[0, 0 : month_day[0]], mean_grad_ratio_zone[1, 0 : month_day[1]])
std_year_grad_ratio = np.append(std_grad_ratio_zone[0, 0 : month_day[0]], std_grad_ratio_zone[1, 0 : month_day[1]])

mean_year_spec_anom_23 = np.append(mean_spec_anom_23_zone[0, 0 : month_day[0]], mean_spec_anom_23_zone[1, 0 : month_day[1]])
std_year_spec_anom_23 = np.append(std_spec_anom_23_zone[0, 0 : month_day[0]], std_spec_anom_23_zone[1, 0 : month_day[1]])

mean_year_spec_anom_89 = np.append(mean_spec_anom_89_zone[0, 0 : month_day[0]], mean_spec_anom_89_zone[1, 0 : month_day[1]])
std_year_spec_anom_89 = np.append(std_spec_anom_89_zone[0, 0 : month_day[0]], std_spec_anom_89_zone[1, 0 : month_day[1]])

mean_year_ratio_anom_89 = np.append(mean_ratio_anom_89_zone[0, 0 : month_day[0]], mean_ratio_anom_89_zone[1, 0 : month_day[1]])
std_year_ratio_anom_89 = np.append(std_ratio_anom_89_zone[0, 0 : month_day[0]], std_ratio_anom_89_zone[1, 0 : month_day[1]])

year_S = np.append(S[0, 0 : month_day[0]], S[1, 0 : month_day[1]])

for imo in range (2, M):
    # gradient ratio
    mean_year_grad_ratio = np.append(mean_year_grad_ratio, mean_grad_ratio_zone[imo, 0 : month_day[imo]])
    std_year_grad_ratio = np.append(std_year_grad_ratio, std_grad_ratio_zone[imo, 0 : month_day[imo]])
    # spec anomaly 23GHz
    mean_year_spec_anom_23 = np.append(mean_year_spec_anom_23, mean_spec_anom_23_zone[imo, 0 : month_day[imo]])
    std_year_spec_anom_23 = np.append(std_year_spec_anom_23, std_spec_anom_23_zone[imo, 0 : month_day[imo]])
    # spec anomaly 89GHz
    mean_year_spec_anom_89 = np.append(mean_year_spec_anom_89, mean_spec_anom_89_zone[imo, 0 : month_day[imo]])
    std_year_spec_anom_89 = np.append(std_year_spec_anom_89, std_spec_anom_89_zone[imo, 0 : month_day[imo]])
    # ratio anomaly 89GHz
    mean_year_ratio_anom_89 = np.append(mean_year_ratio_anom_89, mean_ratio_anom_89_zone[imo, 0 : month_day[imo]])
    std_year_ratio_anom_89 = np.append(std_year_ratio_anom_89, std_ratio_anom_89_zone[imo, 0 : month_day[imo]])
    # number of data points in area of study
    year_S = np.append(year_S, S[imo, 0 : month_day[imo]])



######################################################################################
# calculate standard deviation of emissivity parameters over the year 2009 (364 days)#
######################################################################################
print 'calculate standard deviation of emissivity parameters over the year 2009'
time_std1 = sqrt((1./364.)*sum((mean_year_grad_ratio[nonzero(isnan(mean_year_grad_ratio) == False)] - mean(mean_year_grad_ratio[nonzero(isnan(mean_year_grad_ratio) == False)]))**2))

time_std2 = sqrt((1./364.)*sum((mean_year_spec_anom_23[nonzero(isnan(mean_year_spec_anom_23) == False)] - mean(mean_year_spec_anom_23[nonzero(isnan(mean_year_spec_anom_23) == False)]))**2))

time_std3 = sqrt((1./364.)*sum((mean_year_spec_anom_89[nonzero(isnan(mean_year_spec_anom_89) == False)] - mean(mean_year_spec_anom_89[nonzero(isnan(mean_year_spec_anom_89) == False)]))**2))

time_std4 = sqrt((1./364.)*sum((mean_year_ratio_anom_89[nonzero(isnan(mean_year_ratio_anom_89) == False)] - mean(mean_year_ratio_anom_89[nonzero(isnan(mean_year_ratio_anom_89) == False)]))**2))




#########################
# plot daily parameters #
#########################
vec_months = np.array([0, 31, 59, 90, 120, 151, 181, 212, 243, 273, 304, 334])
ion()
figure()
subplot(2, 1, 1)
plot(mean_year_grad_ratio, 'b', label = 'grad ratio')
plot(mean_year_spec_anom_23, 'r', label = 'spec anom 23GHz')
plot(mean_year_spec_anom_89, 'g', label = 'spec anom 89GHz')
plot(np.zeros([365], float), '--k')
xticks(vec_months, month, rotation = 25)
xlim(0, 365)
fontP = FontProperties()
fontP.set_size('small')
legend(loc = 3, prop = fontP)
grid()
subplot(2, 1, 2)
plot(mean_year_ratio_anom_89, 'y', label = 'ratio anom 89GHz')
plot(np.zeros([365], float), '--k')
xlim(0, 365)
xticks(vec_months, month, rotation = 25)
legend(loc = 1, prop = fontP)
grid()
plt.savefig('/usr/home/lahlod/twice_d/fig_output_ARCTIC/fig_output_sea_ice_study/ice_class_AMSUA/sub_classification/study_by_zones/monthly_evolution_grad_ratio_emis_anomaly_params_zone_Chukchi_Sea.png')
############################################
# plot daily number of data points in area #
############################################
vec_months = np.array([0, 31, 59, 90, 120, 151, 181, 212, 243, 273, 304, 334])
figure()
plot(year_S, label = 'mean=' + str(mean(year_S))[0:5] + ' ; std=' + str(std(year_S))[0:5])
xticks(vec_months, month, rotation = 25)
xlim(0, 365)
ylabel('Number of data points in area')
fontP = FontProperties()
fontP.set_size('small')
legend(loc = 1, prop = fontP)
title('North Canadian Archipelago')
plt.savefig('/usr/home/lahlod/twice_d/fig_output_ARCTIC/fig_output_sea_ice_study/ice_class_AMSUA/sub_classification/study_by_zones/monthly_evolution_nber_data_points_zone_North_Canadian_Archipelago.png')



subplot(2, 1, 2)
plot(std_year_grad_ratio, '--b', label = 'std grad ratio')
plot(std_year_grad_ratio, '--r', label = 'std grad ratio')
plot(std_year_grad_ratio, '--g', label = 'std grad ratio')
plot(std_year_grad_ratio, '--y', label = 'std grad ratio')
xticks(vec_months, month, rotation = 25)
legend(loc = 2, prop = fontP)
xlim(0, 365)



####################
# map studied zone #
####################
ion()
x_ind, y_ind, z_ind, volume = arctic_map.arctic_map_lat50()
x_coast = x_ind
y_coast = y_ind
z_coast = z_ind
map_cartesian_grid.draw_map_cartes_l(x_coast, y_coast, z_coast, volume, xvec[xxi : xxf + 1], yvec[yyi : yyf + 1], gr[0, yyi : yyf + 1, xxi : xxf + 1], gr[0, :, :][nonzero(isnan(gr[0, :, :]) == False)].min(), gr[0, :, :][nonzero(isnan(gr[0, :, :]) == False)].max(), 0.001, cm.s3pcpn_l_r, 'daily gradient ratio (01-01-2009)')
title('area of study - Chukchi Sea')
plt.savefig('/usr/home/lahlod/twice_d/fig_output_ARCTIC/fig_output_sea_ice_study/ice_class_AMSUA/sub_classification/study_by_zones/map_grad_ratio_zone_Chukchi_Sea.png')



############################################
# read emissivity parameters in all Arctic #
############################################
cumul_params = np.zeros([M, ny, nx], float)
grad_ratio = np.zeros([M, ny, nx], float)
spec_anom_23 = np.zeros([M, ny, nx], float)
spec_anom_89 = np.zeros([M, ny, nx], float)
ratio_anom_89 = np.zeros([M, ny, nx], float)
CP = np.zeros([M, ny, nx], float)
for imo in range (0, M): 
    # daily read gradient ratio file
    print 'read daily gradient ratio for month ' + month[imo]
    fichier_grad_ratio = Dataset('/net/argos/data/parvati/lahlod/ARCTIC/monthly_GLACE/gridded_data/cartesian_grid/res_40/cartesian_grid_daily_grad_ratio_spec_23-89_near_nadir_AMSUA_' + month[imo] + '2009.nc', 'r', format = 'NETCDF3_CLASSIC')
    gr = fichier_grad_ratio.variables['grad_ratio'][:]
    fichier_grad_ratio.close()
    # read daily emis anomaly file for 23GHz
    print 'read daily emis anomaly 23GHz for month ' + month[imo]
    fichier_anom23 = Dataset('/net/argos/data/parvati/lahlod/ARCTIC/monthly_GLACE/gridded_data/cartesian_grid/res_40/cartesian_grid_daily_data_lamb_spec_ratio_anomaly_near_nadir_AMSUA23_' + month[imo] + '2009.nc', 'r', format = 'NETCDF3_CLASSIC')
    sa_23 = fichier_anom23.variables['spec_anomaly'][:]
    fichier_anom23.close()
    # read daily emis anomaly file for 89GHz
    print 'read daily emis anomaly 89GHz for month ' + month[imo]
    fichier_anom89 = Dataset('/net/argos/data/parvati/lahlod/ARCTIC/monthly_GLACE/gridded_data/cartesian_grid/res_40/cartesian_grid_daily_data_lamb_spec_ratio_anomaly_near_nadir_AMSUA89_' + month[imo] + '2009.nc', 'r', format = 'NETCDF3_CLASSIC')
    sa_89 = fichier_anom89.variables['spec_anomaly'][:]
    ra_89 = fichier_anom89.variables['ratio_anomaly'][:]
    fichier_anom89.close()
    for ilon in range (0, nx):
        for ilat in range (0, ny):
            # calculate monthly mean of emissivity parameters
            grad_ratio[imo, ilat, ilon] = mean(gr[:, ilat, ilon][nonzero(isnan(gr[:, ilat, ilon]) == False)])
            spec_anom_23[imo, ilat, ilon] = mean(sa_23[:, ilat, ilon][nonzero(isnan(sa_23[:, ilat, ilon]) == False)])
            spec_anom_89[imo, ilat, ilon] = mean(sa_89[:, ilat, ilon][nonzero(isnan(sa_89[:, ilat, ilon]) == False)])
            ratio_anom_89[imo, ilat, ilon] = mean(ra_89[:, ilat, ilon][nonzero(isnan(ra_89[:, ilat, ilon]) == False)])
            # calculate monthly cumulation index = (sum(abs(all emissivity parameters))/maximum of this sum)*100 (to have a percentage)
            cumul_params[imo, ilat, ilon] = abs(grad_ratio[imo, ilat, ilon]) + abs(spec_anom_23[imo, ilat, ilon]) + abs(spec_anom_89[imo, ilat, ilon]) + abs(ratio_anom_89[imo, ilat, ilon])
    CP[imo, :, :] = (cumul_params[imo, :, :]/max(cumul_params[imo, :, :][nonzero(isnan(cumul_params[imo, :, :]) == False)])) * 100.


#######################
# map cuulation index #
#######################
#ion()
x_ind, y_ind, z_ind, volume = arctic_map.arctic_map_lat50()
x_coast = x_ind
y_coast = y_ind
z_coast = z_ind
for imo in range (0, M):
    print 'map month ' + month[imo]
    map_cartesian_grid.draw_map_cartes_l(x_coast, y_coast, z_coast, volume, xvec, yvec, CP[imo, :, :], 0., 100., 1., cm.s3pcpn_l_r, 'Cumulation index (%)')
    title('AMSUA - ' + month[imo] + ' 2009')
    savefig('/usr/home/lahlod/twice_d/fig_output_ARCTIC/fig_output_sea_ice_study/ice_class_AMSUA/sub_classification/maps/cumul_params/cumul_all_parameters/map_cumulation_index_'+ str(MONTH[imo]) + month[imo] + '2009.png')