####################################3
#
## seasonal climlab energy balance model
##
# python3/climblab code
#
# 23.10.2023 0000.0003b
#

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
import climlab
from climlab import constants as const
from climlab.process.diagnostic import DiagnosticProcess
from climlab.domain.field import Field, global_mean

numyears=10
plotvar=0 ## 1,2,3 lot temp, ice, mean albedo

#S1=1365.2*1
S1=1361*1

co2=1

orbit={'ecc': 0.0167643, 'long_peri': 280.32687, 'obliquity': 23.459277, 'S0':S1}
#orbit={'ecc': 0.0, 'long_peri': 0.0, 'obliquity': 0.0, 'S0':S1}

class tanalbedo(DiagnosticProcess):

    def __init__(self, **kwargs):
        super(tanalbedo, self).__init__(**kwargs)

        self.add_diagnostic('albedo')
        Ts = self.state['Ts']
        self._compute_fixed()

    def _compute_fixed(self):
        Ts = self.state['Ts']
        try:
            lon, lat = np.meshgrid(self.lon, self.lat)
        except:
            lat = self.lat
        phi = lat

        try:
            albedo=np.zeros(len(phi));
            albedo=0.42-0.20*np.tanh(0.052*(Ts-3))

        except:
            albedo = np.zeros_like(phi)

        dom = next(iter(self.domains.values()))
        self.albedo = Field(albedo, domain=dom)

    def _compute(self):
        self._compute_fixed()

        return {}


# creating EBM model

#ebm= climlab.EBM(CO2=co2,orbit={'ecc': 0.0167643, 'long_peri': 280.32687, 'obliquity': 23.459277, 'S0':S1})
ebm0= climlab.EBM_seasonal(water_depth=10.0, a0=0.3, num_lat=90, lum_lon=None, num_lev=10,num_lon=None
, orb=orbit)


ebm=climlab.process_like(ebm0)


#ebm.step_forward()
#print(ebm.diagnostics)


#quit(-1)

surface = ebm.domains['Ts']
# define new insolation and SW process



ebm.remove_subprocess('insolation')
insolation = climlab.radiation.DailyInsolation(domains=surface, orb = orbit, **ebm.param)
insolation.S0=S1

##sun = climlab.radiation.DailyInsolation(domains=model.Ts.domain)
ebm.add_subprocess('insolation', insolation)

#ebm.step_forward()


#print(insolation.diagnostics)

#print (insolation.insolation)
#print (np.max(insolation.insolation))


##print(insolation.S0)

#quit(-1)

ebm.remove_subprocess('albedo')
alb = climlab.surface.albedo.StepFunctionAlbedo(state=ebm.state, Tf=-10, **ebm.param)
#alb = climlab.surface.albedo.ConstantAlbedo(domains=surface, **ebm.param)
#alb = tanalbedo(state=ebm.state, **ebm.param)
ebm.add_subprocess('albedo', alb)
ebm.remove_subprocess('SW')
SW = climlab.radiation.absorbed_shorwave.SimpleAbsorbedShortwave(insolation=insolation.insolation, state = ebm.state, albedo = alb.albedo, **ebm.param)
ebm.add_subprocess('SW', SW)
ebm.remove_subprocess('Lw')
LW = climlab.radiation.aplusbt.AplusBT_CO2(CO2=co2,state=ebm.state, **ebm.param)
ebm.add_subprocess('LW', LW)


#print(SW.diagnostics)




#quit(-1)

#ebm.CO2=co2
ebm.remove_subprocess('diffusion')
diff = climlab.dynamics.MeridionalMoistDiffusion(state=ebm.state, timestep=ebm.timestep)
ebm.add_subprocess('diffusion', diff)


#print (ebm)


ebm.step_forward()

#ebm.diagnostics

#ebm.integrate_years(numyears)
ebm.integrate_converge()

#print(ebm.Ts)

#plt.plot(ebm.Ts)

#plt.show()

num_steps_per_year = int(ebm.time['num_steps_per_year'])
mean_year = np.empty(num_steps_per_year)
for m in range(num_steps_per_year):
	ebm.step_forward()
	mean_year[m] = ebm.global_mean_temperature()
Tmean_year = np.mean(mean_year)

print(round(Tmean_year,2))



if(plotvar==0):
        num_steps_per_year = int(ebm.time['num_steps_per_year'])
        Tyear = np.empty((ebm.lat.size, num_steps_per_year))
        for m in range(num_steps_per_year):
            ebm.step_forward()
            Tyear[:,m] = np.squeeze(ebm.Ts)
        Tmin=np.min(Tyear)
        Tmax=np.max(Tyear)
        
        fig = plt.figure(figsize=(5,5))
        ax = fig.add_subplot(111)
        
        factor = 365. / num_steps_per_year
        cmap1=plt.cm.seismic
        cmap1=plt.cm.winter
        #cmap1=plt.cm.cool_r
        cmap1=plt.cm.cool
        cmap1=cmap1.reversed()      
        #levels1=[-80,-70,-60,-50,-40,-30]
        levels2=[-100,-70,-60,-50,-40,-30,-20,-10,0,10,20,30,40,50,60,80,100]
        cax = ax.contourf(factor * np.arange(num_steps_per_year),
                      ebm.lat, Tyear[:,:], 
                      cmap=cmap1, vmin=Tmin, vmax=Tmax, antialiased=False, levels=levels2)
        cs1 = ax.contour(factor * np.arange(num_steps_per_year),
                      ebm.lat, Tyear[:,:], 
                      colors='#00005f', vmin=Tmin, vmax=Tmax, levels=levels2)
        ax.clabel(cs1, cs1.levels, inline=True, fontsize=14)                     
        #cbar1 = plt.colorbar(cax)
        title1='Snowball Earth - temperatures throughout the year deg C \n if S0= '+ str(S1) +' W m-2 , pressure of CO2= '+str(co2)+' ppm volume'
        ax.set_title(title1, fontsize=12)
        ax.tick_params(axis='x', labelsize=12)
        ax.tick_params(axis='y', labelsize=12)
        ax.set_xlabel('Days of year', fontsize=13)
        ax.set_ylabel('Latitude', fontsize=13)
        plt.savefig('1000dpi.svg', dpi=1000)

if(plotvar==1):
        if 'Tf' in ebm.subprocess['albedo'].param.keys():
            Tf = ebm.subprocess['albedo'].param['Tf']
        else:
            print('No ice considered in this model. Can not plot.')

        num_steps_per_year = int(ebm.time['num_steps_per_year'])
        ice_year = np.empty((ebm.lat.size, num_steps_per_year))
        for m in range(num_steps_per_year):
            ebm.step_forward()
            ice_year[:,m] = np.where(np.squeeze(ebm.Ts) <= Tf, 0, 1)
        
        fig = plt.figure(figsize=(5,5))
        ax = fig.add_subplot(111)
        
        factor = 365. / num_steps_per_year
        cax = ax.contourf(factor * np.arange(num_steps_per_year),
                      ebm.lat, ice_year[:,:], 
                      cmap=plt.cm.seismic, vmin=0, vmax=1, levels=2)
        cbar1 = plt.colorbar(cax)
        
        ax.set_title('Ice throughout the year', fontsize=14)
        ax.set_xlabel('Days of year', fontsize=14)
        ax.set_ylabel('Latitude', fontsize=14)


if(plotvar==2):
        fig = plt.figure(figsize=(7.5,5))

        # Temperature plot
        ax2 = fig.add_subplot(111)
        ax2.plot(ebm.lat,ebm.albedo)

        ax2.set_title('Albedo', fontsize=14)
        ax2.set_xlabel('latitude', fontsize=10)
        ax2.set_ylabel('', fontsize=12)

        ax2.set_xticks([-90,-60,-30,0,30,60,90])
        ax2.set_xlim([-90,90])
        ax2.set_ylim([0,1])
        ax2.grid()
    
        plt.show()



plt.show()