#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Apr 26 15:15:34 2024

@author: bwj
"""


# Code used to prepare, invert, and plot oxygen isotope data as described in
# Mineart et al., 2024. Additional files required to run this code fully are
# o_isotope_invert.py, nan_helper.py, and matrix_average.py as described in
#  Johnson and Wing (2020), Nature Geoscience

# %% Load Packages

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Apr 15 09:42:36 2022

@author: bwj
"""


# %% Read in data


import scipy as sp
from scipy import stats
from scipy.interpolate import interp1d
import matplotlib
import matplotlib.patches as patches
import random
from nan_helper import nan_helper
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from o_isotope_invert import o_isotope_invert
from matrix_average import matrix_average

font = {'family' : 'normal',
        'weight' : 'bold',
        'size'   : 22}

matplotlib.rc('font', **font)



sturg = pd.ExcelFile('Edited sturgeon lake isotopes.xlsx')
sturgeon_data = pd.read_excel(sturg, 'Edited')

pecos_data = pd.read_excel('pecos_data.xlsx')

age_data = pd.read_excel('zircon_ages.xlsx')

# %% Zircon weights
# filter out data older than 4.5 Ga
age_filtered = age_data[age_data['Best Age (Ma)'] < 4500]
zirc_bin_width = 30
zirc_n, zirc_bins, zirc_patches = plt.hist(np.subtract(4.5, age_filtered['Best Age (Ma)']/1000),
                                           bins=np.arange(min(age_filtered['Best Age (Ma)'])/1000, max(
                                               age_filtered['Best Age (Ma)'])/1000 + zirc_bin_width/1000, zirc_bin_width/1000),
                                           color='k', alpha=0.25)

# If 98% of zircons are preserved (Puetz et al., 2017), we can multiply each bin by 1.02 to get the "actual" abundance


perc_preserved = 0.98
detrend = np.linspace(0, (1-perc_preserved)*len(zirc_n), num=len(zirc_n))
detrend = np.flip(detrend)
zirc_n_detrend = detrend*zirc_n+zirc_n

plt.hist(np.subtract(4.5, age_filtered['Best Age (Ma)']/1000),
         bins=np.arange(min(age_filtered['Best Age (Ma)'])/1000, max(
             age_filtered['Best Age (Ma)'])/1000 + zirc_bin_width/1000, zirc_bin_width/1000),
         color='k', alpha=0.25)

# %% prep data and define uncertainties
# ~7 estimate from King et al., 2000 and based on alteration of felsic material
d18Osturgeon = 7
d18Osturgeon_high = 8
d18Osturgeon_low = 6

error_pecos = 0.6
error_sturgeon = 0.2
error_chlorite = 20

# %%

d18Onans = nan_helper(sturgeon_data['d18Owr'])
loc_nana = nan_helper(sturgeon_data['UTM E'])

all_nans_idx = []
for irow in range(0, len(d18Onans[0][:])):
    if d18Onans[0][irow] | loc_nana[0][irow] == True:
        all_nans_idx.append(irow)

sturgeon_trimmed = sturgeon_data.drop(all_nans_idx)

sturgeon_temp = sturgeon_trimmed['low_T']
sturgeon_temp_hiT = sturgeon_trimmed['high_T']*0.75


# %%Sturgeon Lake
# Run the inversion over a variety of conditions spanning
# range of estimates for starting rock and high vs low temperature
# of alteration 
# ---- best guess starting rock d18O = 7
W_R_St = np.ones(len(sturgeon_trimmed))

sturgeon_initial = []
for ientry in range(0, len(sturgeon_trimmed['d18Owr'])):
    leave_out = random.randrange(0, len(sturgeon_trimmed['d18Owr']))
    d18Oin = sturgeon_trimmed['d18Owr'].drop(
        sturgeon_trimmed.loc[sturgeon_trimmed.index == leave_out].index)
    temp_in = sturgeon_temp.drop(
        sturgeon_temp.loc[sturgeon_temp.index == leave_out].index)
    x_in = sturgeon_trimmed['UTM E'].drop(
        sturgeon_trimmed.loc[sturgeon_trimmed.index == leave_out].index)
    z_in = sturgeon_trimmed['UTM N'].drop(
        sturgeon_trimmed.loc[sturgeon_trimmed.index == leave_out].index)
    [iso_grid_St, iso_grid_trim_St, x_grid_St, z_grid_St, conc_matrix_St, change_matrixtL,
     qsolved_taperSt, moles_fluidSt, W_R_taperSt, water_initial_d18OSt, water_outgoing_d18OSt,
     temp_testSt, temp_gridSt, moles_O_rockSt, epsilon_r_wSt] = o_isotope_invert(d18Oin, temp_in, x_in, z_in, d18Osturgeon)

    sturgeon_initial.append(water_initial_d18OSt)
    W_R_St[ientry] = W_R_taperSt


W_R_St_HT = np.ones(len(sturgeon_trimmed))

sturgeon_initial_highT = []
for ientry in range(0, len(sturgeon_trimmed['d18Owr'])):
    leave_out = random.randrange(0, len(sturgeon_trimmed['d18Owr']))
    d18Oin = sturgeon_trimmed['d18Owr'].drop(
        sturgeon_trimmed.loc[sturgeon_trimmed.index == leave_out].index)
    temp_in = sturgeon_temp_hiT.drop(
        sturgeon_temp_hiT.loc[sturgeon_temp_hiT.index == leave_out].index)
    x_in = sturgeon_trimmed['UTM E'].drop(
        sturgeon_trimmed.loc[sturgeon_trimmed.index == leave_out].index)
    z_in = sturgeon_trimmed['UTM N'].drop(
        sturgeon_trimmed.loc[sturgeon_trimmed.index == leave_out].index)
    [iso_grid_St_HT, iso_grid_trim_St_HT, x_grid_St_HT, z_grid_St_HT, conc_matrix_St_HT, change_matrixtL_HT,
     qsolved_taperSt_HT, moles_fluidSt_HT, W_R_taperSt_HT, water_initial_d18OSt_HT, water_outgoing_d18OSt_HT,
     temp_testSt_HT, temp_gridSt_HT, moles_O_rockSt_HT, epsilon_r_wSt_HT] = o_isotope_invert(d18Oin, temp_in, x_in, z_in, d18Osturgeon)

    sturgeon_initial_highT.append(water_initial_d18OSt_HT)
    W_R_St_HT[ientry] = W_R_taperSt_HT

# ---- low starting rock d18O = 6
W_R_St_low = np.ones(len(sturgeon_trimmed))

sturgeon_initial_low = []
for ientry in range(0, len(sturgeon_trimmed['d18Owr'])):
    leave_out = random.randrange(0, len(sturgeon_trimmed['d18Owr']))
    d18Oin = sturgeon_trimmed['d18Owr'].drop(
        sturgeon_trimmed.loc[sturgeon_trimmed.index == leave_out].index)
    temp_in = sturgeon_temp.drop(
        sturgeon_temp.loc[sturgeon_temp.index == leave_out].index)
    x_in = sturgeon_trimmed['UTM E'].drop(
        sturgeon_trimmed.loc[sturgeon_trimmed.index == leave_out].index)
    z_in = sturgeon_trimmed['UTM N'].drop(
        sturgeon_trimmed.loc[sturgeon_trimmed.index == leave_out].index)
    [iso_grid_St, iso_grid_trim_St, x_grid_St, z_grid_St, conc_matrix_St, change_matrixtL,
     qsolved_taperSt, moles_fluidSt, W_R_taperSt, water_initial_d18OSt, water_outgoing_d18OSt,
     temp_testSt, temp_gridSt, moles_O_rockSt, epsilon_r_wSt] = o_isotope_invert(d18Oin, temp_in, x_in, z_in, d18Osturgeon_low)

    sturgeon_initial_low.append(water_initial_d18OSt)
    W_R_St_low[ientry] = W_R_taperSt


W_R_St_HT_low = np.ones(len(sturgeon_trimmed))

sturgeon_initial_highT_low = []
for ientry in range(0, len(sturgeon_trimmed['d18Owr'])):
    leave_out = random.randrange(0, len(sturgeon_trimmed['d18Owr']))
    d18Oin = sturgeon_trimmed['d18Owr'].drop(
        sturgeon_trimmed.loc[sturgeon_trimmed.index == leave_out].index)
    temp_in = sturgeon_temp_hiT.drop(
        sturgeon_temp_hiT.loc[sturgeon_temp_hiT.index == leave_out].index)
    x_in = sturgeon_trimmed['UTM E'].drop(
        sturgeon_trimmed.loc[sturgeon_trimmed.index == leave_out].index)
    z_in = sturgeon_trimmed['UTM N'].drop(
        sturgeon_trimmed.loc[sturgeon_trimmed.index == leave_out].index)
    [iso_grid_St_HT, iso_grid_trim_St_HT, x_grid_St_HT, z_grid_St_HT, conc_matrix_St_HT, change_matrixtL_HT,
     qsolved_taperSt_HT, moles_fluidSt_HT, W_R_taperSt_HT, water_initial_d18OSt_HT, water_outgoing_d18OSt_HT,
     temp_testSt_HT, temp_gridSt_HT, moles_O_rockSt_HT, epsilon_r_wSt_HT] = o_isotope_invert(d18Oin, temp_in, x_in, z_in, d18Osturgeon_low)

    sturgeon_initial_highT_low.append(water_initial_d18OSt_HT)
    W_R_St_HT_low[ientry] = W_R_taperSt_HT
    
# ---- high starting rock d18O = 8
W_R_St_high = np.ones(len(sturgeon_trimmed))

sturgeon_initial_high = []
for ientry in range(0, len(sturgeon_trimmed['d18Owr'])):
    leave_out = random.randrange(0, len(sturgeon_trimmed['d18Owr']))
    d18Oin = sturgeon_trimmed['d18Owr'].drop(
        sturgeon_trimmed.loc[sturgeon_trimmed.index == leave_out].index)
    temp_in = sturgeon_temp.drop(
        sturgeon_temp.loc[sturgeon_temp.index == leave_out].index)
    x_in = sturgeon_trimmed['UTM E'].drop(
        sturgeon_trimmed.loc[sturgeon_trimmed.index == leave_out].index)
    z_in = sturgeon_trimmed['UTM N'].drop(
        sturgeon_trimmed.loc[sturgeon_trimmed.index == leave_out].index)
    [iso_grid_St, iso_grid_trim_St, x_grid_St, z_grid_St, conc_matrix_St, change_matrixtL,
     qsolved_taperSt, moles_fluidSt, W_R_taperSt, water_initial_d18OSt, water_outgoing_d18OSt,
     temp_testSt, temp_gridSt, moles_O_rockSt, epsilon_r_wSt] = o_isotope_invert(d18Oin, temp_in, x_in, z_in, d18Osturgeon_high)

    sturgeon_initial_high.append(water_initial_d18OSt)
    W_R_St_high[ientry] = W_R_taperSt


W_R_St_HT_high = np.ones(len(sturgeon_trimmed))

sturgeon_initial_highT_high = []
for ientry in range(0, len(sturgeon_trimmed['d18Owr'])):
    leave_out = random.randrange(0, len(sturgeon_trimmed['d18Owr']))
    d18Oin = sturgeon_trimmed['d18Owr'].drop(
        sturgeon_trimmed.loc[sturgeon_trimmed.index == leave_out].index)
    temp_in = sturgeon_temp_hiT.drop(
        sturgeon_temp_hiT.loc[sturgeon_temp_hiT.index == leave_out].index)
    x_in = sturgeon_trimmed['UTM E'].drop(
        sturgeon_trimmed.loc[sturgeon_trimmed.index == leave_out].index)
    z_in = sturgeon_trimmed['UTM N'].drop(
        sturgeon_trimmed.loc[sturgeon_trimmed.index == leave_out].index)
    [iso_grid_St_HT, iso_grid_trim_St_HT, x_grid_St_HT, z_grid_St_HT, conc_matrix_St_HT, change_matrixtL_HT,
     qsolved_taperSt_HT, moles_fluidSt_HT, W_R_taperSt_HT, water_initial_d18OSt_HT, water_outgoing_d18OSt_HT,
     temp_testSt_HT, temp_gridSt_HT, moles_O_rockSt_HT, epsilon_r_wSt_HT] = o_isotope_invert(d18Oin, temp_in, x_in, z_in, d18Osturgeon_high)

    sturgeon_initial_highT_high.append(water_initial_d18OSt_HT)
    W_R_St_HT_high[ientry] = W_R_taperSt_HT
    
#%% Then, take the average and std of the different inversion iterations
sturgeon_inits = list(sturgeon_initial)+list(sturgeon_initial_highT)+list(sturgeon_initial_low)+list(sturgeon_initial_highT_low)+list(sturgeon_initial_high)+list(sturgeon_initial_highT_high)
sturgeon_avg = np.mean(sturgeon_inits)
sturgeon_std = np.std(sturgeon_inits)

sturgeon_wr_all = list(W_R_St_HT_high)+list(W_R_St_high)+list(W_R_St_HT_low)+list(W_R_St_HT_low)+list(W_R_St_low)+list(W_R_St)
# %% Pecos
#--Similarly for pecos, we have some uncertainty in d18O and chlorite 
# measurements. So, we'll run a few iterations for upper and lower 
#uncertainties for each along with mean values 
# 7.1 is least altered?.2*5.5 + 0.8*7.2 #7.2 granitic 20% , 5.5 basaltic 8o%

d18Oinit_pecos = 7.1
pecos_del18O = pd.Series.tolist(pecos_data['d18O'])
pecos_temp = np.multiply(0.75, pd.Series.tolist(pecos_data['calc temp']))
# pecos_temp = (np.log(pecos_data['d18O']/log_fitA)/log_fitexp)-x_shift
x = pd.Series.tolist(pecos_data['Distance from hinge'])
y = pd.Series.tolist(pecos_data['Elevation (m)'])

numsamples = len(x)
pecos_initial = np.zeros(numsamples)
water_outgoing = np.zeros(numsamples)
moles_fluid = np.zeros(numsamples)
W_R = np.zeros(numsamples)
final_rock = np.zeros(numsamples)
moles_Orock = np.zeros(numsamples)

# Nominal setup
for irun in range(0, numsamples):
    del18O_in = np.add(pecos_del18O, 0)
    Temp_in = np.add(pecos_temp, 0)
    x_in = x
    y_in = y

    num_takeout = np.random.randint(0, len(del18O_in))

    del18O_in = np.delete(del18O_in, num_takeout)
    Temp_in = np.delete(Temp_in, num_takeout)
    x_in = np.delete(x, num_takeout)
    y_in = np.delete(y, num_takeout)

    [iso_grid, iso_grid_trim, x_grid_trim, z_grid_trim,
     conc_matrix_test, change_matrix, qsolved_taper, moles_fluid_pecos, W_R_taper,
     water_initial_d18O, water_outgoing_d18O, temp_test, temp_grid, moles_O_rock_pecos, epsilon_r_w
     ] = \
        o_isotope_invert(del18O_in, Temp_in, x_in, y_in,
                         d18Oinit_pecos)  # mx_center_points,y_center_points,summed_horz_vectors,summed_vert_vectors

    pecos_initial[irun] = water_initial_d18O
    water_outgoing[irun] = water_outgoing_d18O
    moles_fluid[irun] = moles_fluid_pecos
    W_R[irun] = W_R_taper
    final_rock[irun] = del18O_in.mean()

#High d18O, high temp
pecos_initial_HH = np.zeros(numsamples)
W_R_HH = np.zeros(numsamples)
for irun in range(0, numsamples):
    del18O_in = np.add(pecos_del18O, error_pecos)
    Temp_in = np.add(pecos_temp, error_chlorite)
    x_in = x
    y_in = y

    num_takeout = np.random.randint(0, len(del18O_in))

    del18O_in = np.delete(del18O_in, num_takeout)
    Temp_in = np.delete(Temp_in, num_takeout)
    x_in = np.delete(x, num_takeout)
    y_in = np.delete(y, num_takeout)

    [iso_grid, iso_grid_trim, x_grid_trim, z_grid_trim,
     conc_matrix_test, change_matrix, qsolved_taper, moles_fluid_pecos, W_R_taper,
     water_initial_d18O, water_outgoing_d18O, temp_test, temp_grid, moles_O_rock_pecos, epsilon_r_w
     ] = \
        o_isotope_invert(del18O_in, Temp_in, x_in, y_in,
                         d18Oinit_pecos)  # mx_center_points,y_center_points,summed_horz_vectors,summed_vert_vectors

    pecos_initial_HH[irun] = water_initial_d18O
    W_R_HH[irun] = W_R_taper
    
#High d18O, low temp
pecos_initial_HL = np.zeros(numsamples)
W_R_HL = np.zeros(numsamples)
for irun in range(0, numsamples):
    del18O_in = np.add(pecos_del18O, error_pecos)
    Temp_in = np.add(pecos_temp, -error_chlorite)
    x_in = x
    y_in = y

    num_takeout = np.random.randint(0, len(del18O_in))

    del18O_in = np.delete(del18O_in, num_takeout)
    Temp_in = np.delete(Temp_in, num_takeout)
    x_in = np.delete(x, num_takeout)
    y_in = np.delete(y, num_takeout)

    [iso_grid, iso_grid_trim, x_grid_trim, z_grid_trim,
     conc_matrix_test, change_matrix, qsolved_taper, moles_fluid_pecos, W_R_taper,
     water_initial_d18O, water_outgoing_d18O, temp_test, temp_grid, moles_O_rock_pecos, epsilon_r_w
     ] = \
        o_isotope_invert(del18O_in, Temp_in, x_in, y_in,
                         d18Oinit_pecos)  # mx_center_points,y_center_points,summed_horz_vectors,summed_vert_vectors

    pecos_initial_HL[irun] = water_initial_d18O
    W_R_HL[irun] = W_R_taper
#Low d18O, high temp
pecos_initial_LH = np.zeros(numsamples)
W_R_LH = np.zeros(numsamples)
for irun in range(0, numsamples):
    del18O_in = np.add(pecos_del18O, -error_pecos)
    Temp_in = np.add(pecos_temp, error_chlorite)
    x_in = x
    y_in = y

    num_takeout = np.random.randint(0, len(del18O_in))

    del18O_in = np.delete(del18O_in, num_takeout)
    Temp_in = np.delete(Temp_in, num_takeout)
    x_in = np.delete(x, num_takeout)
    y_in = np.delete(y, num_takeout)

    [iso_grid, iso_grid_trim, x_grid_trim, z_grid_trim,
     conc_matrix_test, change_matrix, qsolved_taper, moles_fluid_pecos, W_R_taper,
     water_initial_d18O, water_outgoing_d18O, temp_test, temp_grid, moles_O_rock_pecos, epsilon_r_w
     ] = \
        o_isotope_invert(del18O_in, Temp_in, x_in, y_in,
                         d18Oinit_pecos)  # mx_center_points,y_center_points,summed_horz_vectors,summed_vert_vectors

    pecos_initial_LH[irun] = water_initial_d18O
    W_R_LH[irun] = W_R_taper
#Low d18O, low temp
pecos_initial_LL = np.zeros(numsamples)
W_R_LL = np.zeros(numsamples)
for irun in range(0, numsamples):
    del18O_in = np.add(pecos_del18O, -error_pecos)
    Temp_in = np.add(pecos_temp, -error_chlorite)
    x_in = x
    y_in = y

    num_takeout = np.random.randint(0, len(del18O_in))

    del18O_in = np.delete(del18O_in, num_takeout)
    Temp_in = np.delete(Temp_in, num_takeout)
    x_in = np.delete(x, num_takeout)
    y_in = np.delete(y, num_takeout)

    [iso_grid, iso_grid_trim, x_grid_trim, z_grid_trim,
     conc_matrix_test, change_matrix, qsolved_taper, moles_fluid_pecos, W_R_taper,
     water_initial_d18O, water_outgoing_d18O, temp_test, temp_grid, moles_O_rock_pecos, epsilon_r_w
     ] = \
        o_isotope_invert(del18O_in, Temp_in, x_in, y_in,
                         d18Oinit_pecos)  # mx_center_points,y_center_points,summed_horz_vectors,summed_vert_vectors

    pecos_initial_LL[irun] = water_initial_d18O
    W_R_LL[irun] = W_R_taper
    
#%% Take the average and stdev 
pecos_inits = list(pecos_initial)+list(pecos_initial_HH)+list(pecos_initial_HL)+list(pecos_initial_LH)+list(pecos_initial_LL)
pecos_avg = np.mean(pecos_inits)
pecos_std = np.std(pecos_inits)

pecos_wr_all = list(W_R_LL)+list(W_R_LH)+list(W_R_HL)+list(W_R_HH)+list(W_R)
# %% Seawater oxygen isotope exchange model


t = 4.4  # time in Gyr

Wo = 7  # original seawater d18O
W_SS = -1  # steady state
num_steps = 100  # num of initial model steps
time = np.linspace(0, 4.5, num=num_steps)  # sample every 250 myr
weath_time_on_twostep = 4.5-2.4  # in Ga
weath_time_on = 4.5-1.9
weath_time_early = 4.5-4.43
weath_time_late = 4.5-0.9

# rate constants in Gyr-1, from Muehlenbachs, 1998
k_weath = 8  # nominal 8continental weathering
k_growth = 1.2  # nominal 1.2continental growth
k_hiT = 14.6  # nominal 14.6high temperature seafloor
k_loT = 1.7  # nominal 1.7low temp seafloor/seafloor weathering
k_W_recycling = 0.6  # nominal 0.6water recycling at subduction zones

# fractionations (permil) btwn rock and water, from Muehlenbachs, 1998 except weathering, which we tuned to reproduce -1permil ocean

Delt_weath = 12  # mueh = 9.6 nominal 13, newer 17
Delt_growth = 9.8  # mueh = 9.8
Delt_hiT_mid = 1.5  # meuh = 4.1, nominal = 1.5
Delt_hiT = 4.1
# Delt_hiT_mid = np.zeros(len(time))

Delt_lowT = 9.3  # mueh = 9.3
Delt_water_recycling = 2.5  # mueh = 2.5


# calculate steady state in 250 myr increments
del_graniteo = np.linspace(7.8, 7.8, num=num_steps)

del_basalto = 5.5
del_WR = 7
bb = 0.4
bb2 = 0.1


Delt_hiT_change = (Delt_hiT-Delt_hiT_mid)+(Delt_hiT-Delt_hiT_mid) * \
    0.5*(1+np.tanh((np.subtract(time, weath_time_on)/bb))) - 1
Delt_hiT_change_late = (Delt_hiT-Delt_hiT_mid)+(Delt_hiT-Delt_hiT_mid) * \
    0.5*(1+np.tanh((np.subtract(time, weath_time_late)/bb))) - 1
Delt_hiT_twostep = (Delt_hiT-Delt_hiT_mid)+(Delt_hiT-Delt_hiT_mid) * \
    0.5*(1+np.tanh((np.subtract(time, weath_time_on_twostep)/bb2))) - 1

k_growth_change = 0.5*k_growth * \
    (1+np.tanh((np.subtract(time, weath_time_on)/bb)))
k_growth_late = 0.5*k_growth * \
    (1+np.tanh((np.subtract(time, weath_time_late)/bb)))
k_growth_early = 0.5*k_growth * \
    (1+np.tanh((np.subtract(time, weath_time_early)/bb)))

# Calculate a value for weathering and growth related to zircon peak abundance.
# As a first pass, just compare it directly to the max 50 myr peak of zircon abundance
# doesn't yet take into account preservation bias etc. Could change that at some point
per_zirc = zirc_n_detrend/zirc_n[-1]

k_growth_zirc = per_zirc*k_growth
k_growth_zirc_interp = np.interp(time, zirc_bins[1:], k_growth_zirc)

k_weathering_change = 0.5*k_weath * \
    (1+np.tanh((np.subtract(time, weath_time_on)/bb)))
k_weathering_late = 0.5*k_weath * \
    (1+np.tanh((np.subtract(time, weath_time_late)/bb)))
k_weathering_early = 0.5*k_weath * \
    (1+np.tanh((np.subtract(time, weath_time_early)/bb)))

k_weath_zirc = per_zirc*k_weath
k_weath_zirc_interp = np.interp(time, zirc_bins[1:], k_weath_zirc)

two_step_time = 4.5-2
k_growth_twostep = 0.5*k_growth * \
    (1+np.tanh((np.subtract(time, weath_time_on_twostep)/(bb2))))


weath_time_mid = 4.5 - 1.65
k_weathering_mid = 0.5*k_weath * \
    (1+np.tanh((np.subtract(time, weath_time_mid)/bb)))
k_weathering_two_step = 0.5*k_weath * \
    (1+np.tanh((np.subtract(time, weath_time_mid)/bb2)))

k_loT_change = k_loT*np.ones(time.size)  # keep it the same
k_hiT_change = k_hiT*np.ones(time.size)

k_water_change = k_W_recycling*np.ones(time.size)


del_steady_change = np.zeros(time.size)
del_steady_early = np.zeros(time.size)
del_steady_late = np.zeros(time.size)
del_steady_two_step = np.zeros(time.size)
del_steady_zirc = np.zeros(time.size)

k_sum = np.zeros(time.size)
k_sum_early = np.zeros(time.size)
k_sum_late = np.zeros(time.size)
k_sum_twostep = np.zeros(time.size)
k_sum_zirc = np.zeros(time.size)

for istep in range(0, time.size):
    top = np.sum([k_weathering_change[istep]*(del_graniteo[istep]-Delt_weath),
                  k_growth_change[istep]*(del_graniteo[istep]-Delt_growth),
                  k_hiT_change[istep]*(del_basalto-Delt_hiT_change[istep]),
                  k_loT_change[istep]*(del_basalto-Delt_lowT),
                  k_water_change[istep]*(del_WR-Delt_water_recycling)])
    top_two_step = np.sum([k_weathering_two_step[istep]*(del_graniteo[istep]-Delt_weath),
                           k_growth_twostep[istep] *
                           (del_graniteo[istep]-Delt_growth),
                           k_hiT_change[istep] *
                           (del_basalto-Delt_hiT_twostep[istep]),
                           k_loT_change[istep]*(del_basalto-Delt_lowT),
                           k_water_change[istep]*(del_WR-Delt_water_recycling)])
    top_early = np.sum([k_weathering_early[istep]*(del_graniteo[istep]-Delt_weath),
                        k_growth_early[istep] *
                        (del_graniteo[istep]-Delt_growth),
                        k_hiT_change[istep]*(del_basalto-Delt_hiT),
                        k_loT_change[istep]*(del_basalto-Delt_lowT),
                        k_water_change[istep]*(del_WR-Delt_water_recycling)])
    top_late = np.sum([k_weathering_late[istep]*(del_graniteo[istep]-Delt_weath),
                       k_growth_late[istep]*(del_graniteo[istep]-Delt_growth),
                       k_hiT_change[istep] *
                       (del_basalto-Delt_hiT_change_late[istep]),
                       k_loT_change[istep]*(del_basalto-Delt_lowT),
                       k_water_change[istep]*(del_WR-Delt_water_recycling)])
    top_zirc = np.sum([k_weath_zirc_interp[istep]*(del_graniteo[istep]-Delt_weath),
                       k_growth_zirc_interp[istep] *
                       (del_graniteo[istep]-Delt_growth),
                       k_hiT_change[istep] *
                       (del_basalto-Delt_hiT_change_late[istep]),
                       k_loT_change[istep]*(del_basalto-Delt_lowT),
                       k_water_change[istep]*(del_WR-Delt_water_recycling)])

    k_sum[istep] = np.sum([k_weathering_change[istep], k_growth_change[istep],
                          k_hiT_change[istep], k_loT_change[istep], k_water_change[istep]])
    k_sum_early[istep] = np.sum([k_weathering_early[istep], k_growth_early[istep],
                                k_hiT_change[istep], k_loT_change[istep], k_water_change[istep]])
    k_sum_late[istep] = np.sum([k_weathering_late[istep], k_growth_late[istep],
                               k_hiT_change[istep], k_loT_change[istep], k_water_change[istep]])
    k_sum_twostep[istep] = np.sum([k_weathering_two_step[istep], k_growth_twostep[istep],
                                  k_hiT_change[istep], k_loT_change[istep], k_water_change[istep]])
    k_sum_zirc[istep] = np.sum([k_weath_zirc_interp[istep], k_growth_zirc_interp[istep],
                               k_hiT_change[istep], k_loT_change[istep], k_water_change[istep]])

    del_steady_change[istep] = top/k_sum[istep]
    del_steady_early[istep] = top_early/k_sum_early[istep]
    del_steady_late[istep] = top_late/k_sum_late[istep]
    del_steady_two_step[istep] = top_two_step/k_sum_twostep[istep]
    del_steady_zirc[istep] = top_zirc/k_sum_zirc[istep]

# calculate dW at for each steady state
time_new = np.linspace(0.01, 4.5, num=1000)
f1 = sp.interpolate.interp1d(time, del_steady_change)
f2 = sp.interpolate.interp1d(time, k_sum)
steady_interp = f1(time_new)
k_sum_interp = f2(time_new)
f1_late = sp.interpolate.interp1d(time, del_steady_late)
f2_late = sp.interpolate.interp1d(time, k_sum_late)
steady_interp_late = f1_late(time_new)
k_sum_interp_late = f2_late(time_new)

steady_interp_late = f1_late(time_new)
k_sum_interp_late = f2_late(time_new)

f1_early = sp.interpolate.interp1d(time, del_steady_early)
f2_early = sp.interpolate.interp1d(time, k_sum_early)
steady_interp_early = f1_early(time_new)
k_sum_interp_early = f2_early(time_new)

f1_twostep = sp.interpolate.interp1d(time, del_steady_two_step)
f2_twostep = sp.interpolate.interp1d(time, k_sum_twostep)
steady_interp_twostep = f1_twostep(time_new)
k_sum_interp_twostep = f2_twostep(time_new)

f1_zirc = sp.interpolate.interp1d(time, del_steady_zirc)
f2_zirc = sp.interpolate.interp1d(time, k_sum_zirc)
steady_interp_zirc = f1_zirc(time_new)
k_sum_interp_zirc = f2_zirc(time_new)

dW_middle = np.add(np.subtract(Wo, steady_interp) *
                   np.exp(-np.multiply(time_new, k_sum_interp)), steady_interp)
dW_early = np.add(np.subtract(Wo, steady_interp_early) *
                  np.exp(-np.multiply(time_new, k_sum_interp_early)), steady_interp_early)
dW_late = np.add(np.subtract(Wo, steady_interp_late) *
                 np.exp(-np.multiply(time_new, k_sum_interp_late)), steady_interp_late)
dW_twostep = np.add(np.subtract(Wo, steady_interp_twostep) *
                    np.exp(-np.multiply(time_new, k_sum_interp_twostep)), steady_interp_twostep)

dW_zirc = np.add(np.subtract(Wo, steady_interp_zirc) *
                 np.exp(-np.multiply(time_new, k_sum_interp_zirc)), steady_interp_zirc)

decay_const_low = 0.02
decay_const_high = 0.04
dW_decay_low = []  # np.zeros(len(time_new))
dW_decay_high = []  # np.zeros(len(time_new))
whatstep = []


for istep in range(0, len(time_new)):
    if time_new[istep] <= 1.5:
        dW_decay_low.append(np.add(np.subtract(Wo, steady_interp[-1])*np.exp(-np.multiply(
            time_new[istep], decay_const_low*k_sum[-1])), steady_interp[-1]))
        dW_decay_high.append(np.add(np.subtract(Wo, steady_interp[-1])*np.exp(-np.multiply(
            time_new[istep], decay_const_high*k_sum[-1])), steady_interp[-1]))
        temp_Wo_low = dW_decay_low[istep]
        temp_Wo_high = dW_decay_high[istep]
        whatstep.append(istep)
knickpoint = whatstep[-1]
time_late = np.flip(time_new[-1] - time_new[knickpoint+1:], 0)
low_test = []
high_test = []
for istep in range(0, len(time_late)):
    dW_decay_low.append(np.add(np.subtract(
        temp_Wo_low, steady_interp[-1])*np.exp(-np.multiply(time_late[istep], 0.4*k_sum[-1])), steady_interp[-1]))
    dW_decay_high.append(np.add(np.subtract(
        temp_Wo_high, steady_interp[-1])*np.exp(-np.multiply(time_late[istep], 0.5*k_sum[-1])), steady_interp[-1]))
    low_test.append(np.add(np.subtract(temp_Wo_low, steady_interp[-1])*np.exp(-np.multiply(
        time_late[istep], 0.03*k_sum[-1])), steady_interp[-1]))

# %% Oxygen isotope contour Plots

plt.close('all')
O_contours = [2, 4, 6, 8, 10]
size = 15
iso_colormap = plt.cm.cividis

markersize = 35
ecolor = 'k'
fcolor = 'w'
marker = 'o'
shrink = 0.65
cbarfont = 14

contours = plt.figure(figsize=(15, 6))
plt.clf()

plt.subplot(2, 2, 1)
plt.contourf(x_grid_trim, z_grid_trim, iso_grid_trim,
             O_contours, cmap=iso_colormap)
cbar = plt.colorbar(orientation='horizontal', shrink=shrink)
cbar.ax.set_xlabel('$\delta^{18}$O (‰ V-SMOW)', fontsize=cbarfont)
plt.scatter(x, y, s=markersize, marker='o', c=fcolor, ec='k')
plt.scatter(x_grid_trim, z_grid_trim, s=markersize /
            2, marker="*", color='black', alpha=0.5)

plt.ylim([z_grid_trim.min(), z_grid_trim.max()])
plt.title('Pecos')
plt.xlabel('')
plt.ylabel('')
ax = plt.gca()
ax.set_xticklabels('')


temp_contours = [280, 320, 360, 400, 440]
ax = plt.subplot(2, 2, 3)


cplot = plt.contourf(x_grid_trim, z_grid_trim, temp_grid,
                     extend='both', cmap='Spectral_r')
# fig.add_axes(cax)
# fig.colorbar(im, cax = cax, orientation = 'horizontal')

cbar = plt.colorbar(orientation='horizontal', shrink=shrink)
cbar.ax.set_xlabel('Temperature ($^\circ$C)', fontsize=cbarfont)
cbar.ax.get_yaxis().labelpad = 23
plt.scatter(x_grid_trim, z_grid_trim, s=markersize /
            2, marker="*", color='black', alpha=0.5)
plt.scatter(x, y, s=markersize, marker='o', c=fcolor, ec='k')


plt.ylim([z_grid_trim.min(), z_grid_trim.max()])

plt.xlabel('Horizontal (m)')
plt.ylabel('Vertical (m)')


O_contours = [-1, 2, 5, 8, 11, 15]
plt.subplot(2, 2, 2)
cplot = plt.contourf(x_grid_St, z_grid_St, iso_grid_St,
                     O_contours, extend='both', cmap=iso_colormap)
cbar = plt.colorbar(orientation='horizontal', shrink=shrink)
cbar.ax.set_xlabel('$\delta^{18}$O (‰ V-SMOW)', fontsize=cbarfont)
plt.scatter(sturgeon_data['UTM E'], sturgeon_data['UTM N'],
            s=markersize, marker='o', c=fcolor, ec='k')
plt.scatter(x_grid_St, z_grid_St, s=markersize/2,
            marker="*", color='black', alpha=0.5)
plt.ylim([z_grid_St.min(), z_grid_St.max()])
plt.xlim([x_grid_St.min(), x_grid_St.max()])
plt.title('Sturgeon Lake')
plt.xlabel('')
plt.ylabel('Northing')
ax = plt.gca()
ax.set_xticklabels('')
# plt.gca().set_aspect('auto')

plt.subplot(2, 2, 4)
temp_contours_st = [100, 150, 200, 250, 300, 350, 400]
cplot = plt.contourf(x_grid_St, z_grid_St, temp_gridSt,
                     temp_contours_st, extend='both', cmap='Spectral_r')
cbar = plt.colorbar(orientation='horizontal', shrink=shrink)
cbar.ax.set_xlabel('Temperature ($^\circ$C)', fontsize=cbarfont)
plt.scatter(sturgeon_data['UTM E'], sturgeon_data['UTM N'],
            s=markersize, marker='o', c=fcolor, ec='k')
plt.scatter(x_grid_St, z_grid_St, s=markersize/2,
            marker="*", color='black', alpha=0.5)
plt.ylim([z_grid_St.min(), z_grid_St.max()])
plt.xlim([x_grid_St.min(), x_grid_St.max()])

plt.xlabel('Easting (m)')
plt.ylabel('Northing (m)')


contours.savefig('contours.svg', format='svg', dpi=1200, bbox_inches='tight')
# %% initial d18O histograms
binwidth = .1
c = '#7eb54e'


inversion_hists = plt.figure(figsize=(8, 6))
plt.clf()

ax2 = plt.gca()

sturgeon_weights = np.ones_like(sturgeon_inits)/float(len(sturgeon_inits))
ax2.hist(sturgeon_inits, bins=np.arange(min(sturgeon_inits), max(sturgeon_inits) +
         binwidth, binwidth), weights=sturgeon_weights, alpha=0.6, color='xkcd:blue', ec='k')

# sturgeon_weights_hiT = np.ones_like(
#     sturgeon_initial_highT)/float(len(sturgeon_initial_highT))
# ax2.hist(sturgeon_initial_highT, bins=np.arange(min(sturgeon_initial_highT), max(sturgeon_initial_highT) +
#          binwidth, binwidth), weights=sturgeon_weights, alpha=0.6, color='xkcd:salmon', ec='k')


pecos_weights = np.ones_like(pecos_inits)/float(len(pecos_inits))
ax2.hist(pecos_inits, bins=np.arange(min(pecos_inits), max(pecos_inits) +
          binwidth, binwidth), weights=pecos_weights, alpha=0.6, color='xkcd:plum', ec='k')


# 'Noranda','Snow Lake',
plt.legend(['Sturgeon', 'Pecos'])
plt.xlabel('$\delta^{18}$O (‰ V-SMOW)')
plt.ylabel('Fracton of runs')

inversion_hists.savefig('inversion_hists.svg',
                        format='svg', dpi=1200, bbox_inches='tight')

# %% F/R ratio histograms
binwidth = 0.01
F_R_hist = plt.figure(figsize=(8, 6))
plt.clf()
ax1 = plt.gca()

sturgeon_weights_FR = np.ones_like(sturgeon_wr_all)/float(len(sturgeon_wr_all))
ax1.hist(sturgeon_wr_all, bins=np.arange(min(sturgeon_wr_all), max(sturgeon_wr_all) + binwidth, binwidth),
         weights=sturgeon_weights_FR, alpha=0.6, color='xkcd:blue', ec='k')

# sturgeon_weights_hiT_FR = np.ones_like(W_R_St_HT)/float(len(W_R_St_HT))
# ax1.hist(W_R_St_HT, bins=np.arange(min(W_R_St_HT), max(W_R_St_HT) + binwidth, binwidth),
#          weights=sturgeon_weights_hiT_FR, alpha=0.6, color='xkcd:salmon', ec='k')


pecos_weights = np.ones_like(pecos_wr_all)/float(len(pecos_wr_all))
ax1.hist(pecos_wr_all, bins=np.arange(min(pecos_wr_all), max(pecos_wr_all) + binwidth, binwidth),
         weights=pecos_weights, alpha=0.6, color='xkcd:plum', ec='k')


plt.legend(['Sturgeon', 'Pecos'],
           loc=9)  # 'Noranda','Snow Lake',
plt.xlabel('F/R ratio')
plt.ylabel('Fracton of runs')

F_R_hist.savefig('f_r_hists.svg', format='svg', dpi=1200, bbox_inches='tight')

# %% Ocean steady state plot
font = {'family': 'normal',
        'weight': 'bold',
        'size': 18}

matplotlib.rc('font', **font)

time_labels = ['4.5', '4', '3.5', '3', '2.5', '2', '1.5', '1', '0.5', '0']
time_ticks = np.linspace(0, 4.5, 10)
binwidth = 30

plt.clf()
seawater_evo = plt.figure(figsize=(12, 6))
ax2 = plt.subplot(1, 1, 1)  # subplot(2,1,2) is now active
ax3 = ax2.twinx()
ax3.hist(np.subtract(4.5, age_filtered['Best Age (Ma)']/1000),
         bins=np.arange(min(age_filtered['Best Age (Ma)'])/1000, max(
             age_filtered['Best Age (Ma)'])/1000 + binwidth/1000, binwidth/1000),
         color='k', alpha=0.15)
# ax3.plot(zirc_n_detrend,'o')

plt.gca().invert_xaxis()
# ax3.set_yticklabels('')
plt.xlabel('Age (Ga)')
ax3.set_ylabel('Number of zircons')


plt.xlabel('Age (Ga)')
ax2.set_ylabel('Seawater $\delta^{18}$O (‰- VSMOW)')
# ax3.set_yticklabels('')


modern_sw = [0.2, -0.5, -1]
modern_std = [0.2, 0.5, 0.3]

old_inverse = [0.2, -0.5, -1,3.3]
old_error = [0.2, 0.5, 0.3, 0.1]
new_data = [pecos_avg, sturgeon_avg]
new_error = [pecos_std, sturgeon_std]
EPR_age = 0.002
furoko_age = 0.014
solea_age = 0.09
pecos_age = 1.72
sturgeon_age = 2.732
pano_age = 3.24

old_ages = np.subtract(4.5,[EPR_age, furoko_age, solea_age,pano_age])
new_ages = np.subtract(4.5, [ pecos_age, sturgeon_age])  # Hydrothermal cell inversion results snow_age,noranda_age,
other_color = 'xkcd:dull green'  # np.divide([219,168,133],255)
our_color = 'xkcd:pumpkin'  # np.divide([144,110,110],255)
mag_color = 'xkcd:dusky rose'
pope_color = 'xkcd:dark teal'

# # #'Hadean  emergence',

mag_oc = patches.Rectangle((0, 6), 0.1, 2, linewidth=1,
                           edgecolor='k', facecolor=mag_color)
ax2.add_patch(mag_oc)

Pope = patches.Rectangle((4.5-3.75, 0.8), 0.1, 3, linewidth=1,
                         edgecolor='k', facecolor=pope_color, label='Ophiolite')
ax2.add_patch(Pope)

Hodel = patches.Rectangle((4.5-0.71, -2.28), 0.1, 1.95,
                          linewidth=1, edgecolor='k', facecolor=other_color)
ax2.add_patch(Hodel)

old_inverse = ax2.errorbar(old_ages,old_inverse,yerr=old_error, fmt='o', markerfacecolor='k',
                       mec='k', markersize=8, elinewidth=4, capsize=5, ecolor='k')
inverse = ax2.errorbar(new_ages, new_data, yerr=new_error, fmt='o', markerfacecolor=our_color,
                       mec='k', markersize=8, elinewidth=4, capsize=5, ecolor='k')

# emerge = ax2.errorbar(time_new,dW_twostep,yerr=0.5,fmt='k-',linewidth=2,alpha=0.1)

zirc_color = 'xkcd:dull orange'
zirc_fmt = zirc_color+':'

zirc = ax2.errorbar(time_new, dW_zirc, yerr=1,
                    color=zirc_color, fmt='--', linewidth=2, alpha=.25)


ax2.legend([mag_oc, Pope, Hodel, old_inverse,inverse, zirc], ['Water equilibrated with magma ocean', 'Isua serpentinites', 'Ophiolites', 'Previous inversions','This study', 'Emergence scenarios'],
           # ,emerge ,'Emergence scenarios', 'Tracer inversion estimate'
           loc='upper center', bbox_to_anchor=(.4, 1.02), facecolor='w', framealpha=1)


ax2.errorbar(time_new, dW_middle, yerr=1, fmt='k--', linewidth=2, alpha=.1)
# ax3.plot(time_new,dW_late,'k-',linewidth=2)


Ordo = patches.Rectangle((4.5-0.5, -2), 0.1, 4, linewidth=1,
                         edgecolor='k', facecolor=other_color)
ax2.add_patch(Ordo)
Samail = patches.Rectangle((4.5-0.085, -1.4), 0.1, 2,
                           linewidth=1, edgecolor='k', facecolor=other_color)
ax2.add_patch(Samail)


holmden = patches.Rectangle(
    (2.5, -2), 0.1, 4, linewidth=1, edgecolor='k', facecolor=other_color)
ax2.add_patch(holmden)
ax2.set_ylim([-3.5, 8.5])

ax3.set_xlim([0, 4.5])

plt.xlim([0, 4.5])  # plt.set_xticks(time_ticks, time_labels)
ax3.set_xticklabels(time_labels)
ax2.set_xlabel('Age (Ga)')

seawater_evo.savefig('seawater_evolution.svg', format='svg',
                     dpi=1200, bbox_inches='tight')

# %% Plots for steady state water model fluxes through time
model_forcings = plt.subplots(figsize=(15, 6))
plt.clf()

# x_labels = ['4.5','3.5','2.5','1.5','0.5']
ax1 = plt.subplot(1, 2, 1)

plt.plot(time, k_weath_zirc_interp, 'k:')
plt.plot(time, k_growth_zirc_interp, 'k')
plt.text(0, 25, '(a)')
plt.ylabel('[(O-process/yr)/O-ocean]*10$^9$')
plt.xlabel('Age (Ga)')
ax1.set_xticks(time_ticks)
ax1.set_xticklabels(time_labels)
ax2 = plt.subplot(122, sharey=ax1)

plt.plot(time, k_weathering_mid, 'k:')
plt.plot(time, k_growth_change, 'k')
plt.text(0, 25, '(b)')
plt.legend(['Weathering', 'Continental Growth'])
ax2.tick_params('y', labelleft=False)
ax2.set_xticks(time_ticks)
ax2.set_xticklabels(time_labels)
plt.xlabel('Age (Ga)')


# %% Zircon abundance and detrended alone

zircon_only = plt.figure(figsize=(15, 9))
plt.hist(np.subtract(4.5, age_filtered['Best Age (Ma)']/1000),
         bins=np.arange(min(age_filtered['Best Age (Ma)'])/1000, max(
             age_filtered['Best Age (Ma)'])/1000 + binwidth/1000, binwidth/1000),
         color='k', alpha=0.15)

zirc_interp = np.interp(time, zirc_bins[1:], zirc_n_detrend)

plt.plot(time, zirc_interp, 'k')
ax1 = plt.gca()
ax1.set_xticks(time_ticks)
ax1.set_xticklabels(time_labels)
plt.legend(['Detrended'])
plt.xlim([0, 4.5])
# ax3.plot(zirc_n_detrend,'o')

# plt.gca().invert_xaxis()
# ax3.set_yticklabels('')
plt.ylabel('Number of zircons')
plt.xlabel('Age (Ga)')
ax3.set_ylabel('Number of zircons')
