## Module to run the minimisation
from WM_Methods import OTM
## The linear optimisation module
import cvxpy as cp 
## Module to track runtime of cells and loops
import time
from tqdm.notebook import tqdm
## Suppress warnings related to division by zero
import warnings
warnings.filterwarnings('ignore')
## Module to load files and handle array computations
import xarray as xr
import numpy as np
## Modules to plot outputs
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import pandas as pd
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from matplotlib import gridspec

from mpl_toolkits.axes_grid1.inset_locator import inset_axes

# Flag to set Q_prior (False = 0; True = actual fluxes)
SF = True
# Setting the value of the biases (only works if SF = True)
HF_offset = -5 #W/m2
FW_offset = 0.5 #mm/day
# Specifying the number of BSP bins, 2^n
tree_depth = 4

volnorming = 10**15 #normalising coeffcients
areanorming = 10**12 #normalising coeffcients

# Establish basic constants 
yr2sec = 365.25*24*60*60
# delta(time) in the ACCESS-CM2 model
dt_OG = xr.open_mfdataset('Data/dtime*.nc')
dt_OG['time'] = np.arange(dt_OG.average_DT.size).astype(float)
dt = dt_OG.rename({'time':'Time'})
# Reference density, specific heat capacity and reference salinity
rho = 1035 #kg/m^3
Cp = 3992.10322329649
S0=35

# Range of years over which 'early' and 'late' are defined
dyrs = 9
# Start and end year
init_early = 1979
init_late = 2006
Early_period = (np.array([init_early,init_early+dyrs]) - init_early)*12+((-1850+1979)*12)
Late_period = (np.array([init_late,init_late+dyrs]) - init_early)*12+((-1850+1979)*12)
range_yrs = init_late-init_early+1

## Convert Offset from W/m2 to W, and from mm/day to kg/s
HF_offset_W = HF_offset*ACCESS_BSP_data.A_sum #W
FW_offset_kg_s = FW_offset*ACCESS_BSP_data.A_sum*rho/(1000*24*3600) #kg/s

## Load Q_prior data here
if SF:
    ## Calculate the cumulative time integrated surface fluxes 
    HFDS_cumsum = (((ACCESS_BSP_data.hfds_sum+HF_offset_W)*dt.average_DT).cumsum('Time').isel(Time=-1))\
        -(((ACCESS_BSP_data.hfds_sum+HF_offset_W)*dt.average_DT).cumsum('Time').isel(Time=Early_period[0])) # units: J
    WFO_cumsum = ((ACCESS_BSP_data.wfo_sum+FW_offset_kg_s)*dt.average_DT).cumsum('Time').isel(Time=-1)\
         -(((ACCESS_BSP_data.wfo_sum+FW_offset_kg_s)*dt.average_DT).cumsum('Time').isel(Time=Early_period[0]))# units: kg

    ## The final dflux value is then scaled by the time between the middle of t1 and the middle of t2
    dhfds = HFDS_cumsum*((np.mean(Late_period)-np.mean(Early_period))/(Late_period[-1]-Early_period[0]))
    dwfo = WFO_cumsum*((np.mean(Late_period)-np.mean(Early_period))/(Late_period[-1]-Early_period[0]))

    ## Convert dflux to equivalent T or S change
    dT_hfds = dhfds/(Cp*rho*V_early) # units: C
    dS_wfo = -dwfo*S0/(rho*V_early) # units: g/kg

    time = ACCESS_BSP_data.Time.values

# Flatten the early and late variables to a 1D array
Vol_1 = V_early.values.flatten()
Vol_2 = V_late.values.flatten()
S_1 = SA_early.values.flatten()-S0 # Remove reference salinity S0
S_2 = SA_late.values.flatten()-S0 # Remove reference salinity S0
T_1 = CT_early.values.flatten()
T_2 = CT_late.values.flatten()
A_1 = A_early.values.flatten()
A_2 = A_late.values.flatten()

# Do the same for basin index
Basin_1 = np.zeros_like(V_early)
Basin_2 = np.zeros_like(V_early)
Basin_names = []
for i in range(np.array(Basins).size):
    Basin_1[i,:] = i
    Basin_2[i,:] = i
    for j in range(V_early.shape[-1]):
        #... and for basin name
        Basin_names.append(Basins[i])



Basin_1_inds = Basin_1.flatten()
Basin_2_inds = Basin_2.flatten()

#... and for the edges of the BSP bins
## Here we calculate the mean TS edges averaged over both early and late times
S_start = (0.5*(Part_early.values[:,:,0]+Part_late.values[:,:,0])).flatten()-S0 # Remove reference salinity S0
S_end = (0.5*(Part_early.values[:,:,1]+Part_late.values[:,:,1])).flatten()-S0 # Remove reference salinity S0
T_start = (0.5*(Part_early.values[:,:,2]+Part_late.values[:,:,2])).flatten()
T_end = (0.5*(Part_early.values[:,:,3]+Part_late.values[:,:,3])).flatten()

# Any NaNs are zeroed out
S_1[np.isnan(S_1)] = 0
S_2[np.isnan(S_2)] = 0
T_1[np.isnan(T_1)] = 0
T_2[np.isnan(T_2)] = 0

## Add the Q_prior to S_0; T_0 to get the 'initial' state ##
if SF:
    S_pre = SA_early.values.flatten()-S0
    S_1 = SA_early.values.flatten()-S0+dS_wfo.values.flatten()
    T_pre = CT_early.values.flatten()
    T_1 = CT_early.values.flatten()+dT_hfds.values.flatten()

    S_1[np.isnan(S_1)] = 0
    T_1[np.isnan(T_1)] = 0
    S_1[~np.isfinite(S_1)] = 0
    T_1[~np.isfinite(T_1)] = 0

## Here, we create the tracers and volumes matrices, which will be fed into the MTM function

volumes = np.stack((Vol_1, Vol_2), axis=0)/volnorming # Shape: [2 x N]

salinities = np.stack((S_1, S_2), axis=0)
temps = np.stack((T_1, T_2), axis=0)

tracers = np.stack((salinities, temps),axis=1) # Shape: [2 x M x N], where M = 2 for just T and S, and M>2 for T,S+other tracers

#### We have loaded the standard deviations into an intermediate netCDF file as the std is calculated on the full, time-varying T/S distribution, which is quite large in size

ST_scale = xr.open_mfdataset('Data/ST_standard_devs.nc')

# Weights are 1/std
T_scale = 1/ST_scale.T_std.values
S_scale = 1/ST_scale.S_std.values

# Now create the weight matrix, which is a combination of 1/std(tracer) and 1/Area
# Load A_late for the weights
A_2_modified = A_2.copy()
# Set the minimum area (as zero will lead to inf values)
# This is only relevant for the Q_prior = 0 case which does not work with a 'hard' area constraint
A_2_modified[A_2_modified==0] = np.nanmin(A_2_modified[A_2_modified>0]) 

area_weight = np.sqrt(areanorming/A_2_modified)

## Define the 'hard' area constraint that will be fed into the MTM function and enforce that non-outcropping 
## water masses have no net air-sea flux. Only works with Q_prior = known. 
area_cons = A_2/A_2-1
area_cons[np.isnan(area_cons)] = 1

weights = np.stack((S_scale*area_weight, T_scale*area_weight), axis=0) # Shape: An [M x N] matrix

hard_area = np.stack((area_cons,area_cons), axis=0)

## Break down the Mixing and Adjustment matrices into their constituent tracers
dT_mixing = Mixing[1,:]
dS_mixing = Mixing[0,:]
dS_adj = Adjustment[0,:]
dT_adj = Adjustment[1,:]

### Calculate the total ocean heat content (OHC) and net Q_prior heat flux (hfds) change between t=0 and t=1 in each Basin
OHC_early = rho*Cp*((ACCESS_BSP_data.T_mean-273.15)*ACCESS_BSP_data.V_sum).isel(Time=slice(Early_period[0],Early_period[1])).mean('Time').sum('Depth') # J
OHC_late = rho*Cp*((ACCESS_BSP_data.T_mean-273.15)*ACCESS_BSP_data.V_sum).isel(Time=slice(Late_period[0],Late_period[1])).mean('Time').sum('Depth') # J

hfds_early = ((ACCESS_BSP_data.hfds_sum+HF_offset_W)*dt.average_DT).isel(Time=slice(Early_period[0],Late_period[1])).cumsum('Time').sel(Time=slice(Early_period[0],Early_period[1])).mean('Time').sum('Depth') # J
hfds_late = ((ACCESS_BSP_data.hfds_sum+HF_offset_W)*dt.average_DT).isel(Time=slice(Early_period[0],Late_period[1])).cumsum('Time').sel(Time=slice(Late_period[0],Late_period[1])).mean('Time').sum('Depth') # J

hfds_change = ((hfds_late-hfds_early)/(range_yrs*yr2sec*10**12)).values # units: TW
OHC_change = ((OHC_late-OHC_early)/(range_yrs*yr2sec*10**12)).values # units: TW
### Calculate the change in net Q_adjustment heat flux between t=0 and t=1 in each Basin
dhfds_adj = dT_adj.reshape(V_late.shape)*(Cp*rho*V_late.values) # units: J
hfds_adj_change = np.nansum(dhfds_adj, axis=-1)/(range_yrs*yr2sec*10**12) # units: TW

### First we define an array of connections between each basin
# 1 = connected (sign convention: North = +ve, East = +ve), 0 = disconnected
d_ij = np.zeros((Basins.size,Basins.size))

d_ij[0,:] = [1, -1, 0, 0, 0, 0, 0, 0, 0]
d_ij[1,:] = [1, 1, -1, 0, 0, 0, 0, 0, 0]
d_ij[2,:] = [0, 1, 1, -1, 0, 0, 0, 0, 0]
d_ij[3,:] = [0, 0, 1, 1, 0, 0, 0, 0, -1]
d_ij[4,:] = [0, 0, 0, 0, 1, 0, 1, 0, -1]
d_ij[5,:] = [0, 0, 0, 0, 0, 1, 1, 0, -1]
d_ij[6,:] = [0, 0, 0, 0, -1, -1, 1, 1, 0]
d_ij[7,:] = [0, 0, 0, 0, 0, 0, -1, 1, 0]
d_ij[8,:] = [0, 0, 0, 1, 1, 1, 0, 0, 1]

### The section transport is then g_ij between Basins times this sign convention, in Sverdrups
section_trans = np.zeros((Vol_1.size,Vol_1.size))
for i in (range(Vol_1.size)):
    for j in range(Vol_2.size):
        if d_ij[int(Basin_1_inds[i]), int(Basin_2_inds[j])]!=0:
            if Basin_names[i] != Basin_names[j]:
                section_trans[i,j] = g_ij[i,j]/(range_yrs*yr2sec*10**6)*d_ij[int(Basin_1_inds[i]), int(Basin_2_inds[j])] # units: Sv

### Sum the section transport across all initial water masses so you get a total transport from one basin to another

section_trans_temp = np.zeros((Vol_1.size, Basins.size))
for j in range(Basins.size):
    section_trans_temp[:,j] = np.nansum(section_trans[:,V_early.shape[1]*j:(V_early.shape[1]*(j+1))],axis=-1) # units: Sv

### basin_connex is an array of numbers which will be used to pinpoint the indices in the 144x144 array that belong to Basins i and j when calculating connectivity
basin_connex = np.array([[0,1],
                [1,2],
                [2,3],
                [3,8],
                [4,8],
                [4,6],
                [5,8],
                [6,5],
                [7,6]]).astype(int)
### T and section_joined are the temps and corresponding transports between a given connected Basin pair (as defined by basin_connex)
# So, sum(T*section) should give you the total net heat transp[ort between the Basin pair

T_joined = np.zeros((int(np.size(basin_connex)/2), int(V_early.shape[1]*2)))
section_joined = np.zeros((int(np.size(basin_connex)/2), int(V_early.shape[1]*2)))

for i in range(int(np.size(basin_connex)/2)):
    T_joined[i,:] = np.concatenate([T_1[V_early.shape[1]*basin_connex[i][0]:(V_early.shape[1]*(basin_connex[i][0]+1))],T_1[V_early.shape[1]*basin_connex[i][1]:(V_early.shape[1]*(basin_connex[i][1]+1))]])
    section_joined[i,:] = np.concatenate([section_trans_temp[V_early.shape[1]*basin_connex[i][0]:(V_early.shape[1]*(basin_connex[i][0]+1)),basin_connex[i][1]],\
        section_trans_temp[V_early.shape[1]*basin_connex[i][1]:(V_early.shape[1]*(basin_connex[i][1]+1)),basin_connex[i][0]]])

heat_trans = T_joined*section_joined*rho*Cp*10**6 # units: W
### Sum along the entire section to get the total heat flux across the section
HF_section_tot = np.nansum(heat_trans, axis=-1)/10**12 # units: TW

### The total heat flux into a basin will be the sum of northward and southward section transport, so we need to define the total transport out as this sum:
net_heat_trans_Basins = [HF_section_tot[0],\
    HF_section_tot[1]-HF_section_tot[0],\
        HF_section_tot[2]-HF_section_tot[1],\
            HF_section_tot[3]-HF_section_tot[2],\
                HF_section_tot[4]-HF_section_tot[5],\
                    HF_section_tot[6]-HF_section_tot[7],\
                        HF_section_tot[7]-HF_section_tot[8]+HF_section_tot[5],\
                            HF_section_tot[8],\
                            -HF_section_tot[4]-HF_section_tot[6]-HF_section_tot[3]]

### where net_heat_trans_Basins[j] corresponds to the net transport into (or out of) Basin[j]

### The same routine is applied to salt as above. For more details on the code, see comments in above heat transport calculations.

OSC_early = ((-ACCESS_BSP_data.S_mean/S0*ACCESS_BSP_data.V_sum)).isel(Time=slice(Early_period[0],Early_period[1])).mean('Time').sum('Depth') # m3
OSC_late = ((-ACCESS_BSP_data.S_mean/S0*ACCESS_BSP_data.V_sum)).isel(Time=slice(Late_period[0],Late_period[1])).mean('Time').sum('Depth') # m3

wfo_early = ((ACCESS_BSP_data.wfo_sum+FW_offset_kg_s)*dt.average_DT).isel(Time=slice(Early_period[0],Late_period[1])).cumsum('Time').sel(Time=slice(Early_period[0],Early_period[1])).mean('Time').sum('Depth')/rho # m3
wfo_late = ((ACCESS_BSP_data.wfo_sum+FW_offset_kg_s)*dt.average_DT).isel(Time=slice(Early_period[0],Late_period[1])).cumsum('Time').sel(Time=slice(Late_period[0],Late_period[1])).mean('Time').sum('Depth')/rho # m3

wfo_change = ((wfo_late-wfo_early)/(range_yrs*yr2sec*10**6)).values # units: Sv
OSC_change = ((OSC_late-OSC_early)/(range_yrs*yr2sec*10**6)).values # units: Sv

dwfo_adj = -dS_adj.reshape(V_late.shape)*(rho*V_late.values)/S0 # units: kg
wfo_adj_change = np.nansum(dwfo_adj, axis=-1)/(rho*range_yrs*yr2sec*10**6) # units: Sv


S_joined = np.zeros((int(np.size(basin_connex)/2), int(V_early.shape[1]*2)))
for i in range(int(np.size(basin_connex)/2)):
    S_joined[i,:] = np.concatenate([S_1[V_early.shape[1]*basin_connex[i][0]:(V_early.shape[1]*(basin_connex[i][0]+1))],S_1[V_early.shape[1]*basin_connex[i][1]:(V_early.shape[1]*(basin_connex[i][1]+1))]])

salt_trans = S_joined/S0*section_joined # units: Sv
FWF_section_tot = np.nansum(salt_trans, axis=-1) # units: Sv

net_salt_trans_Basins = [FWF_section_tot[0],\
    FWF_section_tot[1]-FWF_section_tot[0],\
        FWF_section_tot[2]-FWF_section_tot[1],\
            FWF_section_tot[3]-FWF_section_tot[2],\
                FWF_section_tot[4]-FWF_section_tot[5],\
                    FWF_section_tot[6]-FWF_section_tot[7],\
                        FWF_section_tot[7]-FWF_section_tot[8]+FWF_section_tot[5],\
                            FWF_section_tot[8],\
                            -FWF_section_tot[4]-FWF_section_tot[6]-FWF_section_tot[3]]

### Load the Basin and water mass masks for ACCESS-CM2
ACCESS_masks = xr.open_mfdataset('Data/ACCESS_masks.nc')
basin_mask = ACCESS_masks.basin_mask.values
WM_mask = ACCESS_masks.water_mass_mask.values

## Load the latitudes and longitudes for plotting
lats = ACCESS_masks.geolat_t.values
lons = ACCESS_masks.geolon_t.values

if SF:
    ## Get Q_prior in W/m2 and mm/day
    dhfds_flux = (dhfds/(range_yrs*yr2sec*A_early)).values
    dwfo_flux = (dwfo/(range_yrs*yr2sec*A_early*rho)*1000*24*3600).values

## Get Q_adj in W/m2 and mm/day
dhfds_flux_adj = dhfds_adj/(range_yrs*yr2sec*A_late.values)
dwfo_flux_adj = dwfo_adj/(range_yrs*yr2sec*A_late.values*rho)*1000*24*3600

### Map the Q_prior onto the relevant basin and water mass based on the defined masks
if SF:
        hfds_tot = 0
        wfo_tot = 0
        for i in (range(SA_early.shape[-1])):
                tmp_dhfds = 0
                tmp_dwfo = 0
                for j in range(Basins.size):
                        tmp = basin_mask[j,:,:]*WM_mask[i,:,:]*dhfds_flux[j,i]
                        tmp[np.isnan(tmp)] = 0
                        tmp_dhfds = tmp+tmp_dhfds
                        tmp = basin_mask[j,:,:]*WM_mask[i,:,:]*dwfo_flux[j,i]
                        tmp[np.isnan(tmp)] = 0
                        tmp_dwfo = tmp+tmp_dwfo
                hfds_tot = tmp_dhfds+hfds_tot
                wfo_tot = tmp_dwfo+wfo_tot

### Now do the same for the Q_adj
dS_adj_tot = 0
dT_adj_tot = 0
for i in (range(SA_early.shape[-1])):
    tmp_dS_adj = 0
    tmp_dT_adj = 0
    for j in range(Basins.size):
            tmp = basin_mask[j,:,:]*WM_mask[i,:,:]*dwfo_flux_adj[j,i]
            tmp[np.isnan(tmp)] = 0
            tmp_dS_adj = tmp+tmp_dS_adj
            tmp = basin_mask[j,:,:]*WM_mask[i,:,:]*dhfds_flux_adj[j,i]
            tmp[np.isnan(tmp)] = 0
            tmp_dT_adj = tmp+tmp_dT_adj
    dS_adj_tot = tmp_dS_adj+dS_adj_tot
    dT_adj_tot = tmp_dT_adj+dT_adj_tot