import os, sys
import iris
import numpy
import iris.plot as iplt
import matplotlib.pyplot as plt
import seaborn
seaborn.set_context('talk')

cwd = os.getcwd()
repo_dir = '/'
for directory in cwd.split('/')[1:]:
    repo_dir = os.path.join(repo_dir, directory)
    if directory == 'ocean-analysis':
        break

modules_dir = os.path.join(repo_dir, 'modules')
sys.path.append(modules_dir)

import general_io as gio
import timeseries
import convenient_universal as uconv


%matplotlib inline

"""
Explanation: Analysing energy changes at sub-global scales
In CMIP5 climate model simulations, drift can lead to substantial non-closure of the global energy budget (i.e. discrepancies between changes in netTOA flux, ocean surface heat flux and ocean heat storage). Foruntately, it has been shown that in most cases this issue can be remedied (i.e. budget closure can be acheived) by dedriting the netTOA, ocean surface heat flux and OHC data (Hobbs et al 2016; My Anaylsis).   
Of course, many analyses are interested in changes to the energy budget at sub-global scales. In general, there are a few different approaches to energy budget analysis at these scales, which I compare and contrast below:
1. Flux integral
2. Linear trend
3. OHC tendency
(I also need to look at the issue of regridding of ocean data.)
End of explanation
"""


def calc_anomaly(cube):
    """Calculate the anomaly."""
    
    anomaly = cube.copy()
    anomaly.data = anomaly.data - anomaly.data[0]
    
    return anomaly


def get_data(model):
    """Get the hfds, ohc and hfbasin data"""

    rndt_file = '/g/data/r87/dbi599/DRSv2/CMIP5/%s/historical/yr/atmos/r1i1p1/rndt/latest/dedrifted/rndt-zonal-sum_Ayr_%s_historical_r1i1p1_cumsum-all.nc' %(model, model)
    hfds_file = '/g/data/r87/dbi599/DRSv2/CMIP5/%s/historical/yr/ocean/r1i1p1/hfds/latest/dedrifted/hfds-zonal-sum_Oyr_%s_historical_r1i1p1_cumsum-all.nc' %(model, model)
    ohc_file = '/g/data/r87/dbi599/DRSv2/CMIP5/%s/historical/yr/ocean/r1i1p1/ohc/latest/dedrifted/ohc-zonal-sum_Oyr_%s_historical_r1i1p1_all.nc' %(model, model)
    hfbasin_file = '/g/data/r87/dbi599/DRSv2/CMIP5/%s/historical/yr/ocean/r1i1p1/hfbasin/latest/dedrifted/hfbasin-global_Oyr_%s_historical_r1i1p1_cumsum-all.nc' %(model, model)
    
    time_constraint = gio.get_time_constraint(['1861-01-01', '2005-12-31'])
    
    rndt_cube = iris.load_cube(rndt_file, 'TOA Incoming Net Radiation' & time_constraint)
    hfds_cube = iris.load_cube(hfds_file, 'surface_downward_heat_flux_in_sea_water' & time_constraint)
    ohc_cube = iris.load_cube(ohc_file, 'ocean heat content' & time_constraint)
    hfbasin_cube = iris.load_cube(hfbasin_file, 'northward_ocean_heat_transport' & time_constraint)
    
    rndt_anomaly = calc_anomaly(rndt_cube)
    hfds_anomaly = calc_anomaly(hfds_cube)
    ohc_anomaly = calc_anomaly(ohc_cube)
    hfbasin_anomaly = calc_anomaly(hfbasin_cube)
    
    ocean_convergence = ohc_anomaly[-1, ::] - hfds_anomaly[-1, ::]
    hfbasin_inferred = ocean_convergence.copy()
    hfbasin_inferred.data = numpy.ma.cumsum(-1 * ocean_convergence.data)
    
    atmos_convergence = hfds_anomaly[-1, ::] - rndt_anomaly[-1, ::]
    hfatmos_inferred = atmos_convergence.copy()
    hfatmos_inferred.data = numpy.ma.cumsum(-1 * atmos_convergence.data)
    
    return rndt_anomaly, hfds_anomaly, ohc_anomaly, hfbasin_anomaly, hfbasin_inferred, hfatmos_inferred

def plot_uptake_storage(ohc_anomaly, hfds_anomaly, rndt_anomaly):
    """Plot the ocean heat uptake and storage"""

    fig, ax = plt.subplots(figsize=[12, 6])

    iplt.plot(ohc_anomaly[-1, ::], color='blue', label='OHC')
    iplt.plot(hfds_anomaly[-1, ::], color='orange', label='hfds')
    iplt.plot(rndt_anomaly[-1, ::], color='red', label='rndt')

    plt.xlabel('latitude')
    plt.ylabel('J')
    plt.xlim(-90, 90)

    plt.axhline(y=0, color='0.5', linestyle='--')

    plt.legend()
    plt.show()

def plot_transport(hfbasin_anomaly_data, hfbasin_anomaly_inferred, hfatmos_anomaly_inferred):
    """Plot the northward ocean heat transport"""

    fig, ax = plt.subplots(figsize=[12, 6])

    iplt.plot(hfbasin_anomaly_data[-1, ::], color='purple', label='northward OHT')
    iplt.plot(hfbasin_anomaly_inferred, color='purple', linestyle='--', label='inferred northward OHT')
    iplt.plot(hfatmos_anomaly_inferred, color='green', linestyle='--', label='inferred northward AHT')

    plt.xlabel('latitude')
    plt.ylabel('J')
    plt.xlim(-90, 90)

    plt.axhline(y=0, color='0.5', linestyle='--')

    plt.legend()
    plt.show()

"""
Explanation: 1. Flux integral approach
On a global scale, the most common approach to analysing changes in the oceanic energy budget is to plot the change in ocean heat content against the cumulative sum of the ocean surface heat flux (i.e. take the temporal integral of the surface flux to get the corresponding change in heat) (e.g. Hobbs et al 2016; my anaylsis). This approach works well, because any non-zero accumulation of ocean surface heat flux (after dedrifting) will be reflected in the (dedrifted) change in OHC, because in a stationary climate the globally integrated ocean surface heat flux will be zero.
(I think this is essentially the approach that Frolicher et al (2015) take.)
End of explanation
"""


hfds_anomaly, ohc_anomaly, hfbasin_anomaly, hfbasin_inferred = get_data('GISS-E2-R')

plot_uptake_storage(ohc_anomaly, hfds_anomaly)

ohc_anomaly[-1, ::].data.sum()

hfds_anomaly[-1, ::].data.sum()

"""
Explanation: GISS-E2-R model
End of explanation
"""


plot_transport(hfbasin_anomaly, hfbasin_inferred)

hfbasin_anomaly[-1, ::].data.sum()

hfbasin_inferred.data.sum()

"""
Explanation: The OHC and HFDS totals match the corresponding global timeseries.
End of explanation
"""


hfds_anomaly, ohc_anomaly, hfbasin_anomaly, hfbasin_inferred = get_data('CSIRO-Mk3-6-0')

plot_uptake_storage(ohc_anomaly, hfds_anomaly)

ohc_anomaly[-1, ::].data.sum()

hfds_anomaly[-1, ::].data.sum()

hfds_anomaly.coord('latitude')

data_lat_bounds = hfds_anomaly.coord('latitude').bounds
data_lat_diffs = numpy.apply_along_axis(lambda x: x[1] - x[0], 1, data_lat_bounds)
data_lat_diffs

hfds_scaled = hfds_anomaly / data_lat_diffs

new_points = [('latitude', numpy.arange(-89.25, 90, 1.5))]
hfds_new_grid = hfds_scaled.interpolate(new_points, iris.analysis.Linear())

hfds_new = hfds_new_grid * 1.5

iplt.plot(hfds_new[-1, ::], color='orange', label='hfds')
plt.xlabel('latitude')
plt.ylabel('J')
plt.xlim(-90, 90)

plt.show()

hfds_new[-1, ::].data.sum()

plot_transport(hfbasin_anomaly, hfbasin_inferred)

"""
Explanation: CSIRO-Mk3-6-0 model
End of explanation
"""


rndt_anomaly, hfds_anomaly, ohc_anomaly, hfbasin_anomaly, hfbasin_inferred, hfatmos_inferred = get_data('NorESM1-M')

test = False == hfds_anomaly.coord('latitude')
test

plot_uptake_storage(ohc_anomaly, hfds_anomaly, rndt_anomaly)

ohc_anomaly[-1, ::].data.sum()

hfds_anomaly[-1, ::].data.sum()

rndt_anomaly[-1, ::].data.sum()

"""
Explanation: NorESM1-M model
End of explanation
"""


plot_transport(hfbasin_anomaly, hfbasin_inferred, hfatmos_inferred)

"""
Explanation: The HFDS totals match the corresponding global timeseries, but it seems regridding has wrecked the OHC values (the amplitude has dropped appreciably).
End of explanation
"""


