import climt
import sympl
import numpy as np
from datetime import timedelta

import os
os.system("rm 1_yr_spinup_nodamping_T42_30vlvl.nc")

### -----------------------
# 9th Dec 2019
# Script to initialise a simple dry gcm from rest and spin up for one year, 
# saving all fields as go along.
### -----------------------

# Set timestep
timestep = timedelta(minutes=20)

# Convection
convection = climt.EmanuelConvection()

# Radiation: only call every 1 hr
radiation_timestep = timedelta(minutes=60)
radiation = sympl.UpdateFrequencyWrapper(climt.GrayLongwaveRadiation(), 
                                         radiation_timestep)
# Simple physics (boundary layer): use tendencies to work better in spectral space
simple_physics = sympl.TimeDifferencingWrapper(climt.SimplePhysics())

# Set up grid (3degx3deg, 15 v levels)
grid = climt.get_grid(nx=128, ny=64, nz=30, x_name='lon', y_name='lat')

dycore = climt.GFSDynamicalCore(
        [simple_physics, radiation,
         convection], number_of_damped_levels=0
    )

state = climt.get_default_state([dycore], grid_state=grid)

# Initialise saving object
fields_to_store = ['air_temperature', 'air_pressure', 'eastward_wind',
                   'northward_wind', 'surface_pressure',
                   'specific_humidity', 'surface_temperature',
                   'latitude', 'longitude',
                   'convective_heating_rate']

netcdf_monitor = sympl.NetCDFMonitor('1_yr_spinup_nodamping_T42_30vlvl.nc',
                                    store_names=fields_to_store,
                                    write_on_store=True)

# Step forward
diagnostics, new_state = dycore(state, timestep)
state.update(diagnostics)
state.update(new_state)
state['time'] += timestep

#%%time
for step in range(1*24*3*365):
    diagnostics, new_state = dycore(state, timestep)
    state.update(new_state)
    state['time'] += timestep
    if step%(24*3) == 0:
        print(state['time'])
        netcdf_monitor.store(state)
print("Finished at ", state['time'])