from datetime import datetime

import netCDF4
import numpy as np

from barotropic.model import BarotropicModel
from barotropic.io import NetCDFWriter

"""
Explanation: An object-oriented atmospheric model
To give a domain-specific example of how object-oriented programming techniques may be useful, we've put together a simple atmospheric model in Python, using object-oriented programming.
A spectral barotropic vorticity model
We're going to present a simple barotropic vorticity equation model. This is a 2-dimensional model of a homegeneous,  non-divergent, incompressible flow solved on the surface of a sphere. The governing equation is simply:
$\dfrac{\mathrm{D} \omega}{\mathrm{D} t} = 0$
where $\omega$ is the absolute vorticity. Given the spherical geometry of the problem, it is convenient to use a spectral method, representing the prognostic variables as finite summations of spherical harmonics.
Why an object-oriented model
There are several advantages to implementing such a model using object-oriented programming. Firstly we can make clear our separation of concerns:

A model object knows how to take an atmospheric state and evolve it according to some dynamical rules.
A spectral transforms object knows how to transform fields bewteen grid and spectral space, and how to evaluate derivatives in spectral space.
An output writer object knows how to take a model state and write it to an output file.

Why an object-oriented model
A second advantage of the object-oriented design is that we can encapsulate the model state within an object. This means we can construct multiple model instances with different inputs and configurations, and they can never interfere with one another.
These design concepts make running experiments with the model simple too (think of running experiments as another separated concern).
Tour of the model
Running barotropic models
In this notebook we'll run a couple of different barotropic model configurations and write their output to NetCDF files.
We won't have to change any of the code in the barotropic package. We can create multiple independent model instances, because all model state is stored within the model instance (i.e., no global state in the barotropic module).
The first job is to import the modules/objects we'll need:
End of explanation
"""


hour = 3600
run_time = 5 * 24 * hour
output_interval = 6 * hour

"""
Explanation: Before we start creating models we'll just create a few variables to control how long we run each model for and how often we write output to file. Both models will be run for 5 days and write output every 6 hours:
End of explanation
"""


vrt = np.load('input_data/idealized.npy')
nlat, nlon = vrt.shape
truncation = nlon // 3

"""
Explanation: An idealized initial condition
The first model we'll run uses an idealized initial condition as in Held (1985).
The initial vorticity is derived from a zonal wind:
$u(\theta) = 25 \cos(\theta)\, - 30 \cos^3(\theta) + 300 \sin^2(\theta)\,\cos^6(\theta)$
where $\theta$ is latitude. The associated vorticity (assuming zero meridional flow) is computed and the following perturbation is added:
$\zeta^\prime = \dfrac{A}{2}\, \cos (\theta) \, \exp{\left[-\left((\theta - \theta_0)/\theta_w\right)^2\right]} \, \cos (m \lambda)$
with $m = 4$, $\theta_0 = 45^\circ\textrm{N}$, $\theta_w = 15^\circ$, and $A = 8\times 10^{-5}$, and where $\lambda$ is longitude.
To save time and effort, this initial condition has been pre-generated on an N64 Gaussian grid, and is contained in the NumPy save file input_data/idealized.npy.
We can load this file using np.load, determine the grid size and work out a good trunction:
End of explanation
"""


dt = 1800
start_time = datetime(2000, 1, 1)
model_ideal = BarotropicModel(vrt, truncation, dt, start_time)

"""
Explanation: Next we can create a barotropic model using this initial condition.
For this model resolution a time-step of 1800 seconds is a good enough choice.
We'll pick an arbitrary start date, it doesn't mean anything in this case.
End of explanation
"""


saver_ideal = NetCDFWriter(model_ideal, 'model_idealized.nc')

"""
Explanation: We want to run this model, saving output every 3 hours.
For this we're going to need a way to write model data to file.
Within the barotropic.io module is a class called NetCDFWriter that we can use for this.
It is already imported so we can go ahead and use it, it takes a model and a filename as input.
End of explanation
"""


# The run_with_snapshots method is an iterator that yields the
# current model time at the given interval. This means we can
# loop over the iterator saving the current model state.
for t in model_ideal.run_with_snapshots(
        run_time, snapshot_interval=output_interval):
    print('Saving output at time: {}'.format(model_ideal.valid_time))
    saver_ideal.save()
    
# Make sure the output file is closed when we are done:
saver_ideal.close()

"""
Explanation: Later we can call the save() method of the NetCDFWriter to save model state to file:
Now we have both a model and a saver we can run the model.
The model object has a method called run_with_snapshots() which allows us to run a model for a specified amount of time, pausing at specified intervals.
We can tell the model to pause every 3 hours and during this pause we can write output using our saver:
End of explanation
"""


!ncdump -h model_idealized.nc

"""
Explanation: We now have a file 'model_idealized.nc' that contains our model output:
End of explanation
"""


# Load initial condition from file:
ds = netCDF4.Dataset('input_data/ecmwf.201611010000.nc')
vrt = ds.variables['atmosphere_relative_vorticity'][:]
nlat, nlon = vrt.shape
truncation = nlon // 3

# Use a shorter 900 second time-step and the correct start time:
dt = 900
start_time = datetime(2016, 11, 1, 0)

# Create the model and a saver:
model_ecmwf = BarotropicModel(vrt, truncation, dt, start_time)
saver_ecmwf = NetCDFWriter(model_ecmwf, 'model_ecmwf.nc')

# Run the model (same method as before)
for t in model_ecmwf.run_with_snapshots(
        run_time, snapshot_interval=output_interval):
    print('Saving output at time: {}'.format(model_ecmwf.valid_time))
    saver_ecmwf.save()
saver_ecmwf.close()

"""
Explanation: A realistic initial condition
We can also run the barotropic model with a more realistic initial condition.
We'll set up a second model that is initialized using vorticity from ECMWF analysis valid at 2016-11-01 0000z.
This initial condition is stored in the NetCDF file input_data/ecmwf.201611010000.nc.
This file contains data on an N80 Gaussian grid, corresponding to a spectral resolution of T106.
This is higher resolution than the previous model, so we'll need to make the time-step shorter:
End of explanation
"""


import cartopy.crs as ccrs
from cube_browser import Browser, Pcolormesh
import iris
import matplotlib.pyplot as plt


iris.FUTURE.netcdf_promote = True

"""
Explanation: Visualizing the output : Cube Browser
We are going to visualize the output of the two models using a package called cube_browser.
This package is built on-top of Iris/cartopy/matplotlib and integrates with the Jupyter notebook.
End of explanation
"""


vrt_ideal = iris.load_cube('model_idealized.nc',
                           'atmosphere_relative_vorticity')

"""
Explanation: Loading data
We can load the data as normal using iris:
End of explanation
"""


def make_nps_axes():
    ax = plt.axes(projection=ccrs.NorthPolarStereo())
    ax.coastlines()
    ax.set_extent([-180, 180, 30, 90], crs=ccrs.PlateCarree())
    ax.set_title('Relative Vorticity / s-1')
    return ax

ax = make_nps_axes()
plot = Pcolormesh(vrt_ideal, ax, cmap='viridis')
Browser([plot]).display()

"""
Explanation: Drawing the plot
We'll draw a plot by first creating an axes to draw into, then creating a Pcolormesh object using cube_browser. We can then use the Browser object to see the plot.
End of explanation
"""


vrt_ecmwf = iris.load_cube('model_ecmwf.nc',
                           'atmosphere_relative_vorticity')
ax = make_nps_axes()
plot = Pcolormesh(vrt_ecmwf, ax, cmap='viridis')
Browser([plot]).display()

"""
Explanation: A second plot
We can do the same thing for the other model run:
End of explanation
"""


