import numpy as np
import matplotlib.pyplot as plt
from scipy import stats
import statistics
from math import sqrt
import netCDF4 as nc 

def area_calculate_nonuniform(lon,lat):
    Re = 6.3781e6
    area = np.zeros((len(lat),len(lon)))*np.nan
    for II in range(len(lon)-1):
        for JJ in range(len(lat)-1):
            area[JJ+1,II+1] = 2*np.pi*Re**2*(abs(np.sin(lat[JJ+1]*np.pi/180.)-np.sin(lat[JJ]*np.pi/180.)))*abs((lon[II+1]-lon[II]))/360.
    dlat = abs(lat[1] - lat[0])
    for II in range(len(lon)-1):
        area[0,II+1] = 2*np.pi*Re**2*(abs(np.sin(lat[0]*np.pi/180.)-np.sin((lat[0]-dlat)*np.pi/180.)))*abs((lon[II+1]-lon[II]))/360.
    area[:,0] = area[:,-1].copy()
    return area
al = 0.5
C = 30
s = 1970
n = ['1', '0125', '025', '05', '2', '3', '4', '5', '6', '7', '8']
filename = ['1xCO2.B1850LENS.n21.f09_g16.cam.h0.', '0.125xCO2.B1850LENS.n21.f09_g16.cam.h0.', '0.25xCO2.B1850LENS.n21.f09_g16.cam.h0.', '0.5xCO2.B1850LENS.n21.f09_g16.cam.h0.', '2xCO2.B1850LENS.n21.f09_g16.02.cam.h0.', '3xCO2.B1850LENS.n21.f09_g16.cam.h0.', '4xCO2.B1850LENS.n21.f09_g16.cam.h0.', '5xCO2.B1850LENS.n21.f09_g16.02.cam.h0.','6xCO2.B1850LENS.n21.f09_g16.cam.h0.', '7xCO2.B1850LENS.n21.f09_g16.cam.h0.', '8xCO2.B1850LENS.n21.f09_g16.cam.h0.']
filedataset = ['1xCO2.B1850LENS.n21.f09_g16', '0.125xCO2.B1850LENS.n21.f09_g16', '0.25xCO2.B1850LENS.n21.f09_g16', '0.5xCO2.B1850LENS.n21.f09_g16', '2xCO2.B1850LENS.n21.f09_g16.02', '3xCO2.B1850LENS.n21.f09_g16', '4xCO2.B1850LENS.n21.f09_g16', '5xCO2.B1850LENS.n21.f09_g16.02', '6xCO2.B1850LENS.n21.f09_g16', '7xCO2.B1850LENS.n21.f09_g16', '8xCO2.B1850LENS.n21.f09_g16']
cc = ['#990099', '#483D8B', '#6495ED', '#00AA00', '#77FF00', '#CCFF33', '#FFFF00', '#FF8800', '#FF0000', '#AA0000']
labels = ['0.125xCO2', '0.25xCO2', '0.5xCO2', '2xCO2', '3xCO2', '4xCO2', '5xCO2', '6xCO2', '7xCO2', '8xCO2']
plt.figure(figsize=(14, 12), dpi=300)
f = nc.Dataset(r'/work/yuchiaol/from_ivan/'+str(filedataset[0])+'/atm/hist/'+str(filename[0])+'1970-01.nc')
lon = f.variables['lon'][:].data.copy()
lat = f.variables['lat'][:].data.copy()
area = area_calculate_nonuniform(lon,lat)

##########################################   (a) SAT   ############################################
plt.subplot(221)
for N in range(len(n)):
    M = 12
    Y = 30
    m = np.array(("01", "02", "03", "04", "05", "06", "07", "08", "09", "10", "11", "12"))
    y = np.arange(1970, 2000, 1)
    arctic_temp = np.zeros((C, 12)) * np.nan
    for YY in range(Y):
        for MM in range(M):
            f = nc.Dataset(r'/work/yuchiaol/from_ivan/'+str(filedataset[N])+'/atm/hist/'+str(filename[N])+str(y[YY])+"-"+str(m[MM])+'.nc')
            a = f.variables['TREFHT'][:,:,:].data.copy()
            arctic_temp[YY, MM] = np.nansum(a[0,-32:,:]*area[-32:,:])/np.nansum(area[-32:,:])
    arctic_temp_year = np.zeros(12) * np.nan
    sat = np.zeros((30, 12)) * np.nan
    var_sat = np.zeros((30, 12)) * np.nan
    for j in range(12):
        arctic_temp_year[j] = np.mean(arctic_temp[:, j])
        sat[:,j] = arctic_temp[:, j].copy()
    if N==0:
        arctic_1_temp = arctic_temp_year.copy()
        sat_1 = sat.copy()
    else:
        month = np.array((1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12))
        arctic_y = np.zeros(12) * np.nan
        for a in range(6):
            arctic_y[a] = (arctic_temp_year[6+a]-arctic_1_temp[6+a])
            var_sat[:,a] = sat[:,6+a]-sat_1[:,6+a]
        for a in range(6):
            arctic_y[6+a] = (arctic_temp_year[a]-arctic_1_temp[a])
            var_sat[:,6+a] = sat[:,a]-sat_1[:,a]
        for e in range(12):
            bottom, top = stats.t.interval(0.95, 30-1, loc=var_sat[:,e].mean()-var_sat[:,0].mean(), scale=np.std(var_sat[:,e])/np.sqrt(30))
            plt.errorbar(month[e], np.nanmean(var_sat[:,e])-np.nanmean(var_sat[:,0]), yerr=(top-bottom)/2, color=cc[N-1], capsize=7, capthick=1.5, alpha=al)
        if (N-1)<5:
            plt.plot(month, arctic_y-arctic_y[0], color=cc[N-1], linestyle='-', label=labels[N-1], marker='o')
        else:
            plt.plot(month, arctic_y-arctic_y[0], color=cc[N-1], linestyle='-', marker='o')

plt.xticks([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], ['Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec', 'Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun'])
plt.yticks([-12, -11, -10, -9, -8, -7, -6, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13], ["", "", "-10", "", "", "", "", "-5", "", "", "", "", "0", "", "", "", "", "5", "", "", "", "", "10", "", "", ""])
plt.ylim([-12, 13])
plt.ylabel("K", fontsize=14)
plt.title("(a) Arctic (60$^\circ$N) SAT response", fontsize=16)
plt.legend(bbox_to_anchor=(1.2, 1.15), frameon=False, ncol=5, fontsize=11,labelspacing=0.02)

##########################################   (b) SIE   ############################################
plt.subplot(222)
for N in range(len(n)):
    M = 12
    Y = 30
    m = np.array(("01", "02", "03", "04", "05", "06", "07", "08", "09", "10", "11", "12"))
    y = np.arange(1970, 2000, 1)
    SIE = np.zeros((C, 12)) * np.nan
    for YY in range(Y):
        for MM in range(M):
            f = nc.Dataset(r'/work/yuchiaol/from_ivan/'+str(filedataset[N])+'/atm/hist/'+str(filename[N])+str(y[YY])+"-"+str(m[MM])+'.nc')
            a = f.variables['ICEFRAC'][:,:,:].data.copy()
            SIE[YY, MM] = np.nansum(a[0,-32:,:]*area[-32:, :])
    SIE_year = np.zeros(12) * np.nan
    ice = np.zeros((30, 12)) * np.nan
    var_ice = np.zeros((30, 12)) * np.nan
    for j in range(12):
        SIE_year[j] = np.mean(SIE[-30:, j])
        ice[:,j] = SIE[:, j].copy()
    if N==0:
        SIE_1_year = SIE_year
        ice_1 = ice
    else:
        month = np.array((1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12))
        tt = np.zeros(12) * np.nan
        for a in range(6):
            tt[a] = SIE_year[6+a]-SIE_1_year[6+a]
            var_ice[:,a] = ice[:,6+a]-ice_1[:,6+a]
        for a in range(6):
            tt[6+a] = SIE_year[a]-SIE_1_year[a]
            var_ice[:,6+a] = ice[:,a]-ice_1[:,a]
        for e in range(12):
            bottom, top = stats.t.interval(0.95, 30-1, loc=var_ice[:,e].mean()/10**12, scale=np.std(var_ice[:,e]/10**12)/np.sqrt(30))
            plt.errorbar(month[e], np.nanmean(var_ice[:,e])/10**12, yerr=(top-bottom)/2, color=cc[N-1], capsize=7, capthick=1.5, alpha=al)
        if (N-1)>4:
            plt.plot(month, tt/10**12, color=cc[N-1], label=labels[N-1], marker='o')
        else:
            plt.plot(month, tt/10**12, color=cc[N-1], marker='o')

plt.xticks([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], ['Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec', 'Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun'])
plt.yticks([-16, -15, -14, -13, -12, -11, -10, -9, -8, -7, -6, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10], ["", "-15", "", "", "", "", "-10", "", "", "", "", "-5", "", "", "", "", "0", "", "", "", "", "5", "", "", "", "", "10"])
plt.ylim([-16, 10])
plt.title("(b) Arctic (60$^\circ$N) SIE response", fontsize=16)
plt.ylabel("x${10^6}$ ${km^2}$", fontsize=14)
plt.legend(bbox_to_anchor=(1, 1.15), frameon=False, ncol=5, fontsize=11,labelspacing=0.02)
##########################################   (c) tur   ############################################
plt.subplot(223)
for N in range(len(n)):
    M = 12
    Y = 30
    m = np.array(("01", "02", "03", "04", "05", "06", "07", "08", "09", "10", "11", "12"))
    y = np.arange(1970, 2000, 1)
    SE = np.zeros((C, 12)) * np.nan
    LE = np.zeros((C, 12)) * np.nan
    for YY in range(Y):
        for MM in range(M):
            f = nc.Dataset(r'/work/yuchiaol/from_ivan/'+str(filedataset[N])+'/atm/hist/'+str(filename[N])+str(y[YY])+"-"+str(m[MM])+'.nc')
            a = f.variables['SHFLX'][:,:,:].data.copy()
            SE[YY, MM] = np.nansum(a[0,-32:,:]*area[-32:,:])/np.nansum(area[-32:,:])
            a = f.variables['LHFLX'][:,:,:].data.copy()
            LE[YY, MM] = np.nansum(a[0,-32:,:]*area[-32:,:])/np.nansum(area[-32:,:])
    SE_year = np.zeros(12) * np.nan
    LE_year = np.zeros(12) * np.nan
    var = np.zeros((30, 12)) * np.nan
    var_turbulent = np.zeros((30, 12)) * np.nan
    for j in range(12):
        SE_year[j] = np.mean(SE[:, j])
        LE_year[j] = np.mean(LE[:, j])
        var[:, j] = SE[:, j]+LE[:, j]
    turbulent = SE_year + LE_year
    if N==0:
        turbulent_1 = turbulent
        var_1 = var
        print(np.nanmean(turbulent_1))
    else:
        month = np.array((1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12))
        tt = np.zeros(12) * np.nan
        for a in range(6):
            tt[a] = turbulent[6+a]-turbulent_1[6+a]
            var_turbulent[:,a] = var[:,6+a]-var_1[:,6+a]
        for a in range(6):
            tt[6+a] = turbulent[a]-turbulent_1[a]
            var_turbulent[:,6+a] = var[:,a]-var_1[:,a]
        for e in range(12):
            bottom, top = stats.t.interval(0.95, 30-1, loc=var_turbulent[:,e].mean(), scale=np.std(var_turbulent[:,e])/np.sqrt(30))
            plt.errorbar(month[e], np.nanmean(var_turbulent[:,e]), yerr=(top-bottom)/2, color=cc[N-1], capsize=7, capthick=1.5, alpha=al)
        plt.plot(month, tt, color=cc[N-1], label=labels[N-1], marker='o')
plt.xticks([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], ['Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec', 'Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun'])
plt.yticks([-20, -18, -16, -14, -12, -10, -8, -6, -4, -2, 0, 2, 4, 6, 8, 10, 12, 14, 16, 18], ['-20', '-18', '-16', '-14', '-12', '-10', '-8', '-6', '-4', '-2', '0', '2', '4', '6', '8', '10', '12', '14', '16', '18'])
plt.ylim([-20, 18])
plt.title("(c) Turbulent Heat Flux Response (60$^\circ$N)", fontsize=16)
plt.ylabel("W/${m^2}$", fontsize=14)

##########################################   (d) AAF   ############################################
plt.subplot(224)
a_y = np.zeros((len(n)-1, 12))*np.nan
for N in range(len(n)):
    M = 12
    Y = 30
    m = np.array(("01", "02", "03", "04", "05", "06", "07", "08", "09", "10", "11", "12"))
    y = np.arange(1970, 2000, 1)
    temp = np.zeros((C, 12)) * np.nan
    temp_year = np.zeros(12) * np.nan
    for YY in range(Y):
        for MM in range(M):
            f = nc.Dataset(r'/work/yuchiaol/from_ivan/'+str(filedataset[N])+'/atm/hist/'+str(filename[N])+str(y[YY])+"-"+str(m[MM])+'.nc')
            a = f.variables['TREFHT'][:,:,:].data.copy()
            temp[YY, MM] = np.nansum(a[0,:,:]*area[:,:])/np.nansum(area[:,:])
    arctic_temp = np.zeros((C, 12)) * np.nan
    for YY in range(Y):
        for MM in range(M):
            f = nc.Dataset(r'/work/yuchiaol/from_ivan/'+str(filedataset[N])+'/atm/hist/'+str(filename[N])+str(y[YY])+"-"+str(m[MM])+'.nc')
            a = f.variables['TREFHT'][:,:,:].data.copy()
            arctic_temp[YY, MM] = np.nansum(a[0,-32:,:]*area[-32:,:])/np.nansum(area[-32:,:])
    for j in range(12):
        temp_year[j] = np.mean(temp[:, j])   
    if N==0:
        temp_1_year = temp_year.copy()
        temp_1 = temp.copy()
        arctic_temp_1 = arctic_temp.copy()
    else:
        month = np.array((1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12))
        arctic_y = np.zeros(12) * np.nan
        arctic_temp_year = np.nanmean(arctic_temp, axis=0)
        for a in range(6):
            arctic_y[a] = (arctic_temp_year[6+a]-arctic_1_temp[6+a])/(temp_year[6+a]-temp_1_year[6+a])
        for a in range(6):
            arctic_y[6+a] = (arctic_temp_year[a]-arctic_1_temp[a])/(temp_year[a]-temp_1_year[a])
        AAF = (arctic_temp-arctic_temp_1)/(temp-temp_1)
        a_y[N-1, :] = arctic_y
        for i in range(12):
            bottom, top = stats.t.interval(0.95, 30-1, loc=AAF[:, i].mean(), scale=np.std(AAF[:, i])/np.sqrt(30))
            plt.errorbar(month[i], arctic_y[i], yerr=(top-bottom)/2, color=cc[N-1], capsize=7, capthick=1.5, alpha=al)

for e in range(len(n)-1):
    plt.plot(month, a_y[e, :], color=cc[e], linestyle='-', label=labels[e], marker='o')
plt.xticks([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], ['Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec', 'Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun'])
# plt.yticks([-12, -11, -10, -9, -8, -7, -6, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], ["", "", "-10", "", "", "", "", "-5", "", "", "", "", "0", "", "", "", "", "5", "", "", "", "", "10", "", ""])
# plt.ylim([0, 28])
plt.ylabel("", fontsize=14)
plt.title("(d) AAF (60$^\circ$N)", fontsize=16)




plt.tight_layout()
plt.savefig(r'fig4.jpg')