#!/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)




########################################################################
# open NETCDF files of anomaly data and ice spec, lamb, diff and ratio #
########################################################################
frequ = np.array(['23', '30', '50', '89'])

anom_spec = np.zeros([4, M, ny, nx], float)
anom_lamb = np.zeros([4, M, ny, nx], float)
anom_diff = np.zeros([4, M, ny, nx], float)
anom_ratio = np.zeros([4, M, ny, nx], float)

ice_spec = np.zeros([4, M, ny, nx], float)
ice_lamb = np.zeros([4, M, ny, nx], float)
ice_diff = np.zeros([4, M, ny, nx], float)
ice_ratio = np.zeros([4, M, ny, nx], float)

for ifr in range (0, 4):
   for imo in range (0, M):
        print 'stack in file month ' + str(month[imo])
        print 'open anomaly file'
        fichier_anom = Dataset('/net/argos/data/parvati/lahlod/ARCTIC/AMSUA_ice_class/sub_classification/cartesian_grid_map_sea_ice_extent_with-AMSUA23-and-30_param_anomaly_' + month[imo] + '2009_AMSUA' + str(frequ[ifr]) + '.nc', 'r', format='NETCDF3_CLASSIC')
        anom_spec[ifr, imo, :, :] = fichier_anom.variables['spec_anomaly'][:]
        anom_lamb[ifr, imo, :, :] = fichier_anom.variables['lamb_anomay'][:]
        anom_diff[ifr, imo, :, :] = fichier_anom.variables['diff_anomaly'][:]
        anom_ratio[ifr, imo, :, :] = fichier_anom.variables['ratio_anomaly'][:]
        fichier_anom.close()
        print 'open ice file'
        fichier_ice = Dataset('/net/argos/data/parvati/lahlod/ARCTIC/AMSUA_ice_class/sub_classification/cartesian_grid_map_sea_ice_extent_with-AMSUA23-and-30_' + month[imo] + '2009_AMSUA' + str(frequ[ifr]) + '_spec_lamb_thresholds.nc', 'r', format='NETCDF3_CLASSIC')
        ice_spec[ifr, imo, :, :] = fichier_ice.variables['spec_ice_area'][:]
        ice_lamb[ifr, imo, :, :] = fichier_ice.variables['lamb_ice_area'][:]
        ice_diff[ifr, imo, :, :] = fichier_ice.variables['diff_ice_area'][:]
        ice_ratio[ifr, imo, :, :] = fichier_ice.variables['ratio_ice_area'][:]
        fichier_ice.close()


'''
###############
# map anomaly #
###############
print 'map monthly anomaly'
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
ifr = 0 # anomaly of emissivity ratio from AMSUA89GHz database
for imo in range (0, M):
    print 'month ' + str(month[imo])
    map_cartesian_grid.draw_map_cartes_l(x_coast, y_coast, z_coast, volume, xvec, yvec, anom_spec[ifr, imo, :, :], anom_spec[ifr, imo,:,:][nonzero(isnan(anom_spec[ifr, imo,:,:]) == False)].min(), anom_spec[ifr, imo, :, :][nonzero(isnan(anom_spec[ifr, imo , :, :]) == False)].max(), 0.001, cm.s3pcpn_l_r, 'Anomaly of emissivity spec')
    title('AMSUA ' + str(frequ[ifr]) + 'GHz - ' + str(month[imo]) + ' 2009')
    print 'save figure in .png file'
    plt.savefig('/usr/home/lahlod/twice_d/fig_output_ARCTIC/fig_output_sea_ice_study/ice_class_AMSUA/sub_classification/maps_sea_ice_extent/emiss_anomaly/spec/anomaly_spec_map_AMSUA' + str(frequ[ifr]) + '_with_AMSUA23-and-30-spec_'+str(month[imo])+'_2009.png')
#cm.s3pcpn_l_r
'''


################################################
# calculate gradient ratio for two frequencies #
################################################
print 'calculate gradient ratio'
grad_ratio = np.zeros([M, ny, nx], float)
hist_ratio = np.zeros([M, 50], float)
corresp_ratio = np.zeros([M, 50], float)
for imo in range (0, M):
    print 'month ' + month[imo]
    for ilat in range (0, ny):
        for ilon in range (0, nx):
            grad_ratio[imo, ilat, ilon] = (ice_spec[0, imo, ilat, ilon] - ice_spec[3, imo, ilat, ilon]) / (ice_spec[0, imo, ilat, ilon] + ice_spec[3, imo, ilat, ilon])
    '''
    # compute histogram distribution of gradient ratio to find characteristic threshold
    grad_ratio_vec = reshape(grad_ratio[imo, :, :], size(grad_ratio[imo, :, :]))[nonzero(isnan(reshape(grad_ratio[imo, :, :], size(grad_ratio[imo, :, :]))) == False)]
    hist_ratio[imo, :] = hist(grad_ratio_vec, bins = 50, normed = True, histtype='step')[0]
    for ibin in range (0, 50):
        corresp_ratio[imo, ibin] = mean(hist(grad_ratio_vec, bins = 50, normed = True, histtype='step')[1][ibin : ibin + 2])

# plot histogram
c = np.array(['r', 'b', 'c', 'm', 'y', 'g'])
figure()
for imo in range (0, 6):
    plot(corresp_ratio[imo], hist_ratio[imo], c = str(c[imo]), label = str(month[imo]))

grid()
xlim(corresp_ratio.min(), corresp_ratio.max())
xlabel('emissivity ratio')
ylabel('frequency of occurence')
fontP = FontProperties()
fontP.set_size('small')
legend(loc = 2, prop = fontP)
#plt.savefig('/usr/home/lahlod/twice_d/fig_output_ARCTIC/fig_output_sea_ice_study/ice_class_AMSUA/emiss_spec_AMSUA' + str(frequ) + '_JANUARY-JUNE_2009.png')
## plot six following months of spec emissivity histograms ##
figure()
for imo in range (6, M):
    plot(corresp_ratio[imo], hist_ratio[imo], c = str(c[imo - 6]), label = str(month[imo]))

grid()
xlim(corresp_ratio.min(), corresp_ratio.max())
xlabel('emissivity ratio')
ylabel('frequency of occurence')
legend(loc = 1, prop = fontP)
#plt.savefig('/usr/home/lahlod/twice_d/fig_output_ARCTIC/fig_output_sea_ice_study/ice_class_AMSUA/emiss_spec_AMSUA' + str(frequ) + '_JULY-DECEMBER_2009.png')
'''

################################
# map cumulation of parameters #
################################
cumul1 = anom_spec[0, :, :, :] - grad_ratio # correlation
cumul2 = anom_spec[3, :, :, :] + grad_ratio # anti correlation
cumul3 = anom_spec[3, :, :, :] + anom_ratio[3, :, :, :] # anti correlation
cumul4 = abs(anom_spec[0, :, :, :]) + abs(anom_spec[3, :, :, :]) + abs(grad_ratio) + abs(anom_ratio[3, :, :, :])




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):
    '''
    # map cumul1
    map_cartesian_grid.draw_map_cartes_l(x_coast, y_coast, z_coast, volume, xvec, yvec, cumul1[imo, :, :], -0.25, 0.14, 0.01, cm.s3pcpn_l_r, 'Cumulation [emis spec anomaly 23 - grad ratio 23-89]')
    title('AMSUA - ' + month[imo] + ' 2009')
    plt.savefig('/usr/home/lahlod/twice_d/fig_output_ARCTIC/fig_output_sea_ice_study/ice_class_AMSUA/sub_classification/maps_sea_ice_extent/cumul_params/map_cumulation_spec-anom-23_grad-ratio_' + month[imo] +'2009.png')
    # map cumul2
    map_cartesian_grid.draw_map_cartes_l(x_coast, y_coast, z_coast, volume, xvec, yvec, cumul2[imo, :, :], -0.17, 0.17, 0.01, cm.s3pcpn_l_r, 'Cumulation [emis spec anomaly 89 + grad ratio 23-89]')
    title('AMSUA - ' + month[imo] + ' 2009')
    plt.savefig('/usr/home/lahlod/twice_d/fig_output_ARCTIC/fig_output_sea_ice_study/ice_class_AMSUA/sub_classification/maps_sea_ice_extent/cumul_params/map_cumulation_spec-anom-89_grad-ratio_' + month[imo] +'2009.png')
    # map cumul3
    map_cartesian_grid.draw_map_cartes_l(x_coast, y_coast, z_coast, volume, xvec, yvec, cumul3[imo, :, :], -5.46, 2.14, 0.01, cm.s3pcpn_l_r, 'Cumulation [emis spec anomaly 89 + emis ratio 89]')
    title('AMSUA - ' + month[imo] + ' 2009')
    plt.savefig('/usr/home/lahlod/twice_d/fig_output_ARCTIC/fig_output_sea_ice_study/ice_class_AMSUA/sub_classification/maps_sea_ice_extent/cumul_params/map_cumulation_spec-anom-89_ratio-anom-89_' + month[imo] +'2009.png')
    '''
    # map cumul4 
    map_cartesian_grid.draw_map_cartes_l(x_coast, y_coast, z_coast, volume, xvec, yvec, cumul4[imo, :, :], cumul4[imo, :, :][nonzero(isnan(cumul4[imo, :, :]) == False)].min(), cumul4[imo, :, :][nonzero(isnan(cumul4[imo, :, :]) == False)].max(), 0.1, cm.s3pcpn_l_r, 'Cumulation all parameters')
    title('AMSUA - ' + month[imo] + ' 2009')
    plt.savefig('/usr/home/lahlod/twice_d/fig_output_ARCTIC/fig_output_sea_ice_study/ice_class_AMSUA/sub_classification/maps_sea_ice_extent/cumul_params/map_cumulation_all_parameters_' + month[imo] +'2009.png')





'''
figure()
scatter(x_coast, y_coast, c = z_coast, s = volume)
#cmap = cm.s3pcpn_l_r
levels = np.arange(-1, 11, 1)
cs = contour(xvec, yvec, cumul[imo, :, :]) 
cs = contour(xvec, yvec, new_anom_spec0[imo, :, :], levels) 
cs = contour(xvec, yvec, new_anom_ratio[imo, :, :], levels) 
cs = contour(xvec, yvec, new_grad_ratio[imo, :, :], levels) 
cbar = colorbar(cs, ticks = [-1, 0, 1, 2, 3, 4, 5])
cbar.ax.set_yticklabels(['emis_spec_89', 'emis_spec_23', 'emis_ratio_89', 'grad_ratio_23-89', '-'])
#cbar.set_label(text)
xlim(-3500, 2700.)
ylim(-4000, 2800.)
'''







'''
############################################
# time evolution (monthly) in a given zone #
############################################
# read monthly std of emis 
ifr = 3
fichier_temporal_std = Dataset('/net/argos/data/parvati/lahlod/ARCTIC/monthly_GLACE/gridded_data/cartesian_grid/res_40/cartesian_grid_monthly_mean-std_data_lamb_spec_near_nadir_AMSUB' + str(frequ[ifr]) + '_' + month[imo] + '2009.nc', 'r', format='NETCDF3_CLASSIC')



#zone 1 : yi = 400. // yf = 720. // xi = -1080. // xf = -920. (Canadian Archipelago)
#zone 2 : yi = 920. // yf = 1280. // xi = -720. // xf = -560. (Chukchi Sea, Bering Detroit)
# select borders of zone
yi = 920. 
yf = 1280.
xi = -720.
xf = -560.
#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]

# define vectors
anom_spec0_zone = np.zeros([M], float)
anom_spec3_zone = np.zeros([M], float)
anom_ratio_zone = np.zeros([M], float)
grad_ratio_zone = np.zeros([M], float)
std_as0 = np.zeros([M], float)
std_as3 = np.zeros([M], float)
std_ar = np.zeros([M], float)
std_gr = np.zeros([M], float)
for imo in range (0, M):
    # anom spec emis from 23GHz
    anom_spec0_vec = reshape(anom_spec[0, imo, yyi : yyf + 1, xxi : xxf + 1], size(anom_spec[0, imo, yyi : yyf + 1, xxi : xxf + 1]))
    if (len(anom_spec0_vec[nonzero(isnan(anom_spec0_vec) == True)]) != 0):
        anom_spec0_zone[imo] = NaN
        std_as0[imo] = NaN
        anom_spec3_zone[imo] = NaN
        std_as3[imo] = NaN
        anom_ratio_zone[imo] = NaN
        std_ar[imo] = NaN
        grad_ratio_zone[imo] = NaN
        std_gr[imo] = NaN
        continue
    else:
        anom_spec0_zone[imo] = anom_spec0_vec.mean()
        std_as0[imo] = sqrt((1. / (size(anom_spec[0, imo, yyi : yyf + 1, xxi : xxf + 1]) - 1.)) * sum((anom_spec0_vec[:] - anom_spec0_zone[imo])**2))
        # anom spec emis from 89GHz
        anom_spec3_vec = reshape(anom_spec[3, imo, yyi : yyf + 1, xxi : xxf + 1], size(anom_spec[3, imo, yyi : yyf + 1, xxi : xxf + 1]))
        anom_spec3_zone[imo] = anom_spec3_vec.mean()
        std_as3[imo] = sqrt((1. / (size(anom_spec[3, imo, yyi : yyf + 1, xxi : xxf + 1]) - 1.)) * sum((anom_spec3_vec[:] - anom_spec3_zone[imo])**2))
        # anom emis ratio from 89GHz
        anom_ratio_vec = reshape(anom_ratio[3, imo, yyi : yyf + 1, xxi : xxf + 1], size(anom_ratio[3, imo, yyi : yyf + 1, xxi : xxf + 1]))
        anom_ratio_zone[imo] = anom_ratio_vec.mean()
        std_ar[imo] = sqrt((1. / (size(anom_ratio[3, imo, yyi : yyf + 1, xxi : xxf + 1]) - 1.)) * sum((anom_ratio_vec[:] - anom_ratio_zone[imo])**2))
        # gradient ratio from 23 and 89GHz
        grad_ratio_vec = reshape(grad_ratio[imo, yyi : yyf + 1, xxi : xxf + 1], size(grad_ratio[imo, yyi : yyf + 1, xxi : xxf + 1]))
        grad_ratio_zone[imo] = grad_ratio_vec.mean()
        std_gr[imo] = sqrt((1. / (size(grad_ratio[imo, yyi : yyf + 1, xxi : xxf + 1]) - 1.)) * sum((grad_ratio_vec[:] - grad_ratio_zone[imo])**2))


figure()
fontP = FontProperties()
fontP.set_size('small')
subplot(2,1,1)
errorbar(arange(0, M, 1), anom_spec0_zone, yerr = std_as0, label = 'emis spec anomaly 23GHz')
errorbar(arange(0, M, 1), anom_spec3_zone, yerr = std_as3, label = 'emis spec anomaly 89GHz')
errorbar(arange(0, M, 1), grad_ratio_zone, yerr = std_gr, label = 'gradient ratio 23-89GHz')
plot(arange(0, M, 1), np.zeros([M], float), '--k')
legend(loc = 3, prop = fontP)
grid()
xticks(arange(0, M, 1), month, rotation = 15)
xlim(-1, M)
subplot(2,1,2)
errorbar(arange(0, M, 1), anom_ratio_zone, yerr = std_ar, label = 'emis ratio anomaly 89GHz')
plot(arange(0, M, 1), np.zeros([M], float), '--k')
legend(loc = 2, prop = fontP)
grid()
xticks(arange(0, M, 1), month, rotation = 15)
xlim(-1, M)
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_emis_params_zone_North_Beaufort_Sea.png')
### map the selected zone ###
map_cartesian_grid.draw_map_cartes_l(x_coast, y_coast, z_coast, volume, xvec[xxi : xxf], yvec[yyi : yyf], grad_ratio[imo, yyi : yyf, xxi : xxf], grad_ratio[imo, yyi : yyf, xxi : xxf][nonzero(isnan(grad_ratio[imo, yyi : yyf, xxi : xxf]) == False)].min(), grad_ratio[imo, yyi : yyf, xxi : xxf][nonzero(isnan(grad_ratio[imo, yyi : yyf, xxi : xxf]) == False)].max(), 0.001, cm.s3pcpn_l_r, 'selected zone')
plt.savefig('/usr/home/lahlod/twice_d/fig_output_ARCTIC/fig_output_sea_ice_study/ice_class_AMSUA/sub_classification/study_by_zones/map_zone_North_Beaufort_Sea.png')
'''



'''
############################################################################
# calculate correlation between anom ratio and gradient ratio (spec emiss) #
############################################################################
corr_mat = np.zeros([M, 4, 4], float)
for imo in range (0, M):
    a = reshape(anom_spec[0, imo, :, :], size(anom_spec[0, imo, :, :]))[nonzero(isnan(reshape(anom_spec[0, imo, :, :], size(anom_spec[0, imo, :, :]))) == False)]
    b = reshape(anom_spec[3, imo, :, :], size(anom_spec[3, imo, :, :]))[nonzero(isnan(reshape(anom_spec[3, imo, :, :], size(anom_spec[3, imo, :, :]))) == False)]
    c = reshape(anom_ratio[3, imo, :, :], size(anom_ratio[3, imo, :, :]))[nonzero(isnan(reshape(anom_ratio[3, imo, :, :], size(anom_ratio[3, imo, :, :]))) == False)]
    d = reshape(grad_ratio[imo, :, :], size(grad_ratio[imo, :, :]))[nonzero(isnan(reshape(grad_ratio[imo, :, :], size(grad_ratio[imo, :, :]))) == False)]
    params = np.array([a, b, c, d])
    for ii in range (0, 4):
        for jj in range (0, 4):
            corr_mat[imo, ii, jj] = corrcoef(params[ii], params[jj])[0][1]

figure()
fontP = FontProperties()
fontP.set_size('small')
plot(arange(0, M, 1), corr_mat[:, 0, 1], 'r', label = 'spec anomaly 23GHz - spec anomaly 89GHz')
plot(arange(0, M, 1), corr_mat[:, 0, 2], 'g', label = 'spec anomaly 23GHz - ratio anomaly 89GHz')
plot(arange(0, M, 1), corr_mat[:, 0, 3], 'b', label = 'spec anomaly 23GHz - gradient ratio')
plot(arange(0, M, 1), corr_mat[:, 1, 2], 'c', label = 'spec anomaly 89GHz - ratio anomaly 89GHz')
plot(arange(0, M, 1), corr_mat[:, 1, 3], 'm', label = 'spec anomaly 89GHz - gradient ratio')
plot(arange(0, M, 1), corr_mat[:, 2, 3], 'y', label = 'ratio anomaly 89GHz - gradient ratio')
plot(arange(0, M, 1), zeros([M], float), '--k')
xlim(-1, M)
ylim(-2, 1.)
legend(loc = 8, prop = fontP)
xticks(arange(0, M, 1), month, rotation = 25)
yticks(np.arange(-2., 1., 0.1))
grid()
plt.savefig('/usr/home/lahlod/twice_d/fig_output_ARCTIC/fig_output_sea_ice_study/ice_class_AMSUA/sub_classification/correlation_params_2009.png')
'''

'''
corr_mat = np.zeros([4, ny, nx], float)
for ifr in range (0, 4):
    for ilat in range (0, ny):
        for ilon in range (0, nx):
            corr_mat[ifr, ilat, ilon] = corrcoef(grad_ratio[:, ilat, ilon], anom_spec[ifr, :, ilat, ilon])[0][1]

for ifr in range (0, 4):
    map_cartesian_grid.draw_map_cartes_l(x_coast, y_coast, z_coast, volume, xvec, yvec, corr_mat[ifr, :, :], -1., 1., 0.01, cm.s3pcpn_l_r, 'correlation [emis ratio anomaly - gradient ratio]')
    title('AMSUA ' + str(frequ[ifr]) + 'GHz - year 2009')
    plt.savefig('/usr/home/lahlod/twice_d/fig_output_ARCTIC/fig_output_sea_ice_study/ice_class_AMSUA/sub_classification/maps_sea_ice_extent/correlation_map_grad-ratio_emis-spec-anom_AMSUA' + str(frequ[ifr]) + '_with_AMSUA23-and-30-spec_year2009.png')
'''