#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Dec 10 2019

@author: Martin Renoult
correspondence: martin.renoult@misu.su.se
"""


## Library

from pymc3 import *
import numpy as np
import matplotlib.pyplot as plt
import scipy.stats as stat
from scipy.optimize import curve_fit
from adjustText import adjust_text
import math

#------------------------------------------------------------------------------
## Lists to save the data while computing

list_predict_t = list()
list_predict_t_stats_66 = list()
list_predict_t_stats_90 = list()

lb_90 = list()
ub_90 = list()
lb_66 = list()
ub_66 = list()

#------------------------------------------------------------------------------
## Model data
# x = ECS
# y = Tropical temperature

# Pliocene from PlioMIP1
# In the following order:
# [CCSM4, IPSLCM5A, MIROC4m, GISS-E2-R, COSMOS, MRI-CGCM2.3, HadCM3,NorESM-L,
# FGOALS-g2, GISS-E2-1-G, IPSL-CM6A-LR, NorESM1-F,CESM2,EC-EARTH3.3]

# FGOALS has not been recomputed: Taken from Hargreaves et Annan, 2016

# ECS of FGOALS and HadCM3 are taken from LGM ECS for consistency
#x = [3.2, 3.4, 4.05, 2.8, 4.1, 3.2, 3.3, 3.1,
#     3.37, 2.6,4.50, 2.29,5.3,4.3]
#
#y = [1.0256042, 1.3337059, 1.9900227, 1.1576538, 2.1806774, 1.151104, 1.9331722, 1.445343,
#     2.14, 0.9211941,2.1174774, 1.3736095,3.4950447,2.94250999999997]

# Latest model versions approach

x = [3.2, 4.05, 4.1, 3.2, 3.3,
     3.37, 2.6,4.50, 2.29,5.3,4.3]

y = [1.0256042, 1.9900227, 2.1806774, 1.151104, 1.9331722,
     2.14, 0.9211941,2.1174774, 1.3736095,3.4950447,2.94250999999997]

# PlioMIP1 only
#x = [3.2, 3.4, 4.05, 2.8, 4.1, 3.2, 3.1, 3.1,3.7]

# Modified ECS, new HadCM3 and FGOALS-g2 to be consistent with LGM
#x = [3.2, 3.4, 4.05, 2.8, 4.1, 3.2, 3.3, 3.1,3.37]
#y = [1.0256042, 1.3337059, 1.9900227, 1.1576538, 2.1806774, 1.151104, 1.9331722, 1.445343,2.14]

# Full PlioMIP1 plot
#pliomip1x = [3.2, 3.4, 4.05, 2.8, 4.1, 3.2, 3.3, 3.1,3.37]
#pliomip1y = [1.0256042, 1.3337059, 1.9900227, 1.1576538, 2.1806774, 1.151104, 1.9331722, 1.445343,2.14]

# Latest version plot
pliomip1x = [3.2, 4.05, 4.1, 3.2, 3.3, 3.37]
pliomip1y = [1.0256042, 1.9900227, 2.1806774, 1.151104, 1.9331722 ,2.14]

pliomip2x = [2.6, 4.50, 2.29]
pliomip2y = [0.9211941,2.1174774, 1.3736095]
#------------------------------------------------------------------------------
### Uncomment to see the data distribution
#
#fig, ax = plt.subplots(figsize=(7, 7))
#
# #plt.plot(pliomip1x, pliomip1y, '.', label='PlioMIP1', ms=17, color='#009999', mec='#006666')
# #plt.plot(pliomip2x, pliomip2y, '.', label='PlioMIP2', ms=17, color='#cc0066',mec='#800040')
#
#plt.xlim(-1, 6)
#plt.ylim(-0.5, 3)
#plt.legend(loc='upper left', bbox_to_anchor=(0.2, 0.35), fancybox=True)
#plt.xlabel('Climate Sensitivity (K)', labelpad=-40, weight='bold')
#plt.ylabel('Mid-Pliocene Tropical Ocean SST anomaly (K)', position=(0,0.8), weight='bold')
#ax.spines['top'].set_visible(False)
#ax.spines['right'].set_visible(False)
#ax.spines['bottom'].set_position(('data', 0))
#ax.spines['left'].set_position(('data', 0))
#ax.spines['left'].set_linewidth(2)
#ax.spines['bottom'].set_linewidth(2)
#ax.tick_params('x', direction='out', pad=20, width=2)
#ax.tick_params('y', width=2)
#plt.yticks(ticks=[0.5, 1, 1.5, 2, 2.5, 3,3.5], labels=['0.5', '1','1.5', '2','2.5', '3','3.5'],
#           weight='bold')
#plt.xticks(ticks=[1, 2, 3, 4, 5, 6], labels=['1', '2','3',  '4','5', '6'],
#           weight='bold')

#------------------------------------------------------------------------------
## MCMC model

with Model() as model: # model specifications in PyMC3 are wrapped in a with-statement
    # Define prior means, sd or precision (see conjugate for precision)
    
    mn_intercept = 0.0
    mn_slope = 1.0
    sd_reg = 1.0
    
    precis_intercept = np.sqrt(0.5)
    precis_slope = np.sqrt(0.5)
        
    # Define priors: Non-conjugate approach
    
    sigma = HalfCauchy('Sigma', beta=5, testval=1.)
    intercept = Normal('Intercept', mn_intercept, sd=sd_reg)
    x_coeff = Normal('x', mn_slope, sd=sd_reg)

    # Priors for a Normal-Inverse Gamma conjugate approach. 
    # In conjugate approach, priors on intercept and slope depends on a scaled sigma
    # and need to be defined that way to match mathematical equations.
    # This part should be mainly used for comparison / check with the conjugate approach code
    
    # sigma = InverseGamma('Sigma',alpha=0.5,beta=0.5)
    # intercept = Normal('Intercept', mn_intercept, sd=sigma/precis_intercept)
    # x_coeff = Normal('x', mn_slope, sd=sigma/precis_slope)

    # Define likelihood
    
    likelihood = Normal('y', mu=intercept + x_coeff * x, 
                        sd=sigma, observed=y)
    
    # Inference! 4 jobs in parallel (convergence check)
    # By default, the sampling method is NUTS
    
    trace = sample(progressbar=True, draws=100000, cores=4, tune=5000)

# Extract the data of the trace
values_x = trace['x']
values_intercept = trace['Intercept']
values_sigma = trace['Sigma']

# Gelman-Rubin test for convergence of the model
# If BFMI = Gelman-Rubin, then you have convergence
# It compares the variance between the chains to the variance inside a chain 
# and both variances should be equal if all the chains (the model) converged

#bfmi = bfmi(trace)
#max_gr = max(np.max(gr_stats) for gr_stats in gelman_rubin(trace).values())
#
#(energyplot(trace, legend=True, figsize=(6, 4))
#   .set_title("BFMI = {}\nGelman-Rubin = {}".format(bfmi, max_gr)));

#------------------------------------------------------------------------------
## Predicted temperature calculation
## Create predicted ensemble for 5-95% estimate

# Discrete sample of sensitivity
ran = np.linspace(0, 10, 500)

# Loop for predicted temperature based on trace and line above
for j in ran:
    predicted_t = values_x * j + values_intercept + np.random.normal(loc=0, scale=values_sigma)
    
    # Calculate and save the 5-95% interval of the prediction
    stats_predict_t_90 = np.percentile(predicted_t, q=(5,95))
    # Calculate and save the 17-83% interval of the prediction
    stats_predict_t_66 = np.percentile(predicted_t, q=(17,83))

    # Save in a list the intervals for every sample of sensitivity "ran"
    list_predict_t_stats_66.append(stats_predict_t_66)
    list_predict_t_stats_90.append(stats_predict_t_90)
    
#------------------------------------------------------------------------------
## Bayesian framework

# Priors on sensitivity
#prior_S = np.random.uniform(0, 20, size=400000)
#prior_S = stat.gamma.rvs(a=1, loc=0, scale=5, size=400000)

def truncated_cauchy_rvs(loc=0, scale=1, a=-1, b=1, size=None):
    """
    Generate random samples from a truncated Cauchy distribution.

    `loc` and `scale` are the location and scale parameters of the distribution.
    `a` and `b` define the interval [a, b] to which the distribution is to be
    limited.

    With the default values of the parameters, the samples are generated
    from the standard Cauchy distribution limited to the interval [-1, 1].
    """
    ua = np.arctan((a - loc)/scale)/np.pi + 0.5
    ub = np.arctan((b - loc)/scale)/np.pi + 0.5
    U = np.random.uniform(ua, ub, size=size)
    rvs =  loc + scale * np.tan(np.pi*(U - 0.5))
    return rvs

# Truncated-at-zero Cauchy distribution
prior_S = truncated_cauchy_rvs(loc=2.5, scale=3, a=1/math.inf, b=math.inf, size=400000)

# Compute 5-95% and 17-83% prior intervals
prior_stats_90 = np.percentile(prior_S, q=(5, 95))
prior_stats_66 = np.percentile(prior_S, q=(17, 83))

# Model to generate a single point based on the prior on S
def gen_mod(alpha, s, beta, error):
    return alpha * s + beta + np.random.normal(loc=0, scale=error)

# Likelihood estimate
def likelihood(sim, obs, std):
    return stat.norm.pdf(x=sim, loc=obs, scale=std)

# Generate temperatures
model_T = gen_mod(values_x, prior_S, values_intercept, values_sigma)

# "Real" observed data
# Tropical Pliocene T
T = 0.8
stdT = 1
gauss_obs = np.random.normal(loc=T, scale=stdT, size=800000)

obs_stats_90 = np.percentile(gauss_obs, q=(5, 95))

# Create weights through importance sampling
weight = likelihood(model_T, T, stdT)
weight = weight/weight.sum()

# Bayesian updating of the prior with importance sampling
posterior = np.random.choice(prior_S, size=100000, p=weight)

post_median = np.median(posterior)

# Compute 5-95% and 17-83% posterior intervals 
post_stats_90 = np.percentile(posterior, q=(5, 95))
post_stats_66 = np.percentile(posterior, q=(17, 83))

#------------------------------------------------------------------------------
## Plot part

## 1st plot: Trace plot
plt.figure(figsize=(7, 7))
traceplot(trace)
plt.tight_layout()
#plt.savefig('Trace_PlioMIP.pdf')

                       #-------------------------------

## 2nd plot: BLR
# Plot the data
fig, ax = plt.subplots(figsize=(7, 7))

# Range of the plotted MCMC lines and plot of the lines
range_eval = np.linspace(-1, 10, 100)
plots.plot_posterior_predictive_glm(trace, samples=100, eval=range_eval, 
                                    label='Predictive regression lines', alpha=0.35)

# Plot the 5-95% interval
for h in list_predict_t_stats_90:
    lb_90.append(h[0])
    ub_90.append(h[1])

for g in list_predict_t_stats_66:
    lb_66.append(g[0])
    ub_66.append(g[1])
    
# Compute running mean to smooth the confidence interval
rand_rm = np.convolve(ran, np.ones((50,))/50, mode='valid')

low_rm_90 = np.convolve(lb_90, np.ones((50,))/50, mode='valid')
up_rm_90 = np.convolve(ub_90, np.ones((50,))/50, mode='valid')

low_rm_66 = np.convolve(lb_66, np.ones((50,))/50, mode='valid')
up_rm_66 = np.convolve(ub_66, np.ones((50,))/50, mode='valid')

plt.plot(rand_rm, low_rm_90, linestyle='-', color='red', label='5-95% interval', alpha=0.75, linewidth=2)
plt.plot(rand_rm, up_rm_90, linestyle='-', color='red', alpha=0.75, linewidth=2)
plt.plot(rand_rm, low_rm_66, linestyle='--', color='red', label='17-83% interval', alpha=0.75, linewidth=2)
plt.plot(rand_rm, up_rm_66, linestyle='--', color='red', alpha=0.75, linewidth=2)

ylim = plt.ylim(-0.5, 3)
xlim = plt.xlim(-1,8)

# Plot 2 std on the figure instead of one (esthetic change)
stdT = 1.6

# Line for observed value, 2 standard deviations
plt.axvline(x=0.2, ymin=(-0.5/(ylim[0]-ylim[1]))-((T-stdT)/(ylim[0]-ylim[1])), ymax=(-0.5/(ylim[0]-ylim[1]))-((T+stdT)/(ylim[0]-ylim[1])), 
            color='#009900', label='5-95% observed value', linewidth=2)
plt.axvline(x=0.2, ymin=(-0.5/(ylim[0]-ylim[1]))-((T-stdT)/(ylim[0]-ylim[1])), ymax=(-0.5/(ylim[0]-ylim[1]))-((T-stdT)/(ylim[0]-ylim[1])), 
            color='#009900', marker='v')
plt.axvline(x=0.2, ymin=(-0.5/(ylim[0]-ylim[1]))-((T+stdT)/(ylim[0]-ylim[1])), ymax=(-0.5/(ylim[0]-ylim[1]))-((T+stdT)/(ylim[0]-ylim[1])), 
            color='#009900', marker='^')
            
plt.axvline(x=0.2, ymin=(-0.5/(ylim[0]-ylim[1]))-((T)/(ylim[0]-ylim[1])), ymax=(-0.5/(ylim[0]-ylim[1]))-((T)/(ylim[0]-ylim[1])), 
            color='#009900', marker='.', ms=12)
            
plt.axhline(y=0.08, xmin=(-1/(xlim[0]-xlim[1]))-(post_stats_90[0]/(xlim[0]-xlim[1])), 
            xmax=(-1/(xlim[0]-xlim[1]))-(post_stats_90[1]/(xlim[0]-xlim[1])), c='#9933ff', label='5-95% posterior', linewidth=2)
plt.axhline(y=0.08, xmin=(-1/(xlim[0]-xlim[1]))-(post_stats_90[0]/(xlim[0]-xlim[1])), 
            xmax=(-1/(xlim[0]-xlim[1]))-(post_stats_90[0]/(xlim[0]-xlim[1])), marker='<', c='#9933ff')
plt.axhline(y=0.08, xmin=(-1/(xlim[0]-xlim[1]))-(post_stats_90[1]/(xlim[0]-xlim[1])), 
            xmax=(-1/(xlim[0]-xlim[1]))-(post_stats_90[1]/(xlim[0]-xlim[1])), marker='>', c='#9933ff')

plt.axhline(y=0.08, xmin=(-1/(xlim[0]-xlim[1]))-(post_median/(xlim[0]-xlim[1])),
            xmax=(-1/(xlim[0]-xlim[1]))-(post_median/(xlim[0]-xlim[1])), c='#9933ff', marker='.', ms=12)            

plt.plot(pliomip1x, pliomip1y, '.', label='PlioMIP1',markersize=17, color='#009999', mec='#006666')
plt.plot(pliomip2x, pliomip2y, '.', label='PlioMIP2',markersize=17, color='#cc0066',mec='#800040')

# Adjust text function. Computes location of model numbers based on all points (esthetic)
texts = [plt.text(x[i], y[i], '%s' %(i+15), ha='center', va='center', fontsize=15) for i in range(0, 1, 1)]
adjust_text(texts)
#texts2 = [plt.text(x[i], y[i], '%s' %(i+16), ha='center', va='center', fontsize=15) for i in range(1, 2, 1)]
#adjust_text(texts2)
#texts3 = [plt.text(x[i], y[i], '%s' %(i+17), ha='center', va='center', fontsize=15) for i in range(2, 4, 1)]
#adjust_text(texts3)
#texts4 = [plt.text(x[i], y[i], '%s' %(i+17), ha='center', va='center', fontsize=15) for i in range(4, 5, 1)]
#adjust_text(texts4)
#texts6 = [plt.text(x[i], y[i], '%s' %(i+18), ha='center', va='center', fontsize=15) for i in range(5, 6, 1)]
#adjust_text(texts6)
#texts5 = [plt.text(x[i], y[i], '%s' %(i+20), ha='center', va='center', fontsize=15) for i in range(6, 9, 1)]
#adjust_text(texts5)
         
# Make it pretty
plt.legend(loc='best', bbox_to_anchor=(0.6, 0.45), fancybox=True, edgecolor='k')
plt.xlabel('Climate sensitivity (K)', labelpad=10, position=(0.6,0),fontsize=16)
plt.ylabel('Mid-Pliocene tropical (30° S - 30° N) \nocean SST anomaly (K)',position=(0,0.6),fontsize=16)
ax.spines['top'].set_alpha(0)
ax.spines['right'].set_visible(False)
ax.spines['bottom'].set_position(('data', 0))
ax.spines['left'].set_position(('data', 0))
ax.spines['top'].set_position(('data', 0))
ax.spines['left'].set_linewidth(2)
ax.spines['bottom'].set_linewidth(2)
ax.tick_params('x', direction='out', pad=5, width=2)
ax.tick_params('y', width=2)
plt.yticks(ticks=[0.5, 1, 1.5, 2, 2.5, 3], labels=['0.5', '1','1.5', '2','2.5', '3'], fontsize=14)
plt.xticks(ticks=np.arange(1,9,1), fontsize=14)
plt.tight_layout()
#plt.savefig('Bayes_PlioMIP.pdf', dpi=300)

                        #-------------------------------              

## 3rd plot: Posterior
fig, ax = plt.subplots(figsize=(7,7))

## Prior distribution line plot
scale = 5
loc = 2.5
a = 0
b = 10
k = 3
cauchy_scale = 3

x = np.linspace(0, 15, 1000)

# /!\ *1.3 account for the truncation at zero for the Cauchy distribution (esthetic approximation for the Figure)
cauchy = (1/((np.pi*cauchy_scale)*(1+((x-loc)/cauchy_scale)**2)))*1.3
#uniform = np.linspace(1/(b-a), 1/(b-a), 1000)
#gamma = (x**(k-1)*np.exp(-x/scale))/(scale**k)

plt.plot(x, cauchy, '-', color='darkorange', label='Prior',linewidth=4, linestyle='--')
#plt.plot(x, uniform, '-', c='darkorange', label='Prior',linewidth=4, linestyle='--')
#plt.plot(x, gamma, '-', c='darkorange', label='Prior',linewidth=4, linestyle='--')

# Fit function. Doesn't work well sometimes...
# Changing the value of alpha plot the true histogram behind
def fit_function(x, A, beta, B, mu, sigma):
    return (A * np.exp(-x/beta) + B * np.exp(-1.0 * (x - mu)**2 / (2 * sigma**2)))
     
n, bins, patches = plt.hist(posterior, density=True, bins=500, alpha=0)
xspace = np.linspace(0, 8, 1000000)
binscenters = np.array([0.5 * (bins[ijk] + bins[ijk+1]) for ijk in range(len(bins)-1)])
popt, pcov = curve_fit(fit_function, xdata=binscenters, ydata=n)
plt.plot(xspace, fit_function(xspace, *popt), color='darkorange', linewidth=4, label='Posterior')

ylim = plt.ylim(-0.02, 0.5)
xlim = plt.xlim(-1,8)

plt.xlabel('Climate sensitivity (K)',fontsize=16)
plt.ylabel('Probability',fontsize=16)

plt.axhline(y=0.008, xmin=(-1/(xlim[0]-xlim[1]))-(post_stats_90[0]/(xlim[0]-xlim[1])), 
            xmax=(-1/(xlim[0]-xlim[1]))-(post_stats_90[1]/(xlim[0]-xlim[1])), c='#9933ff', label='5-95% estimate', linewidth=2)
plt.axhline(y=0.008, xmin=(-1/(xlim[0]-xlim[1]))-(post_stats_90[0]/(xlim[0]-xlim[1])), 
            xmax=(-1/(xlim[0]-xlim[1]))-(post_stats_90[0]/(xlim[0]-xlim[1])), marker='<', c='#9933ff')
plt.axhline(y=0.008, xmin=(-1/(xlim[0]-xlim[1]))-(post_stats_90[1]/(xlim[0]-xlim[1])), 
            xmax=(-1/(xlim[0]-xlim[1]))-(post_stats_90[1]/(xlim[0]-xlim[1])), marker='>', c='#9933ff')

plt.axhline(y=0.008, xmin=(-1/(xlim[0]-xlim[1]))-(post_median/(xlim[0]-xlim[1])),
            xmax=(-1/(xlim[0]-xlim[1]))-(post_median/(xlim[0]-xlim[1])), c='#9933ff', marker='.', ms=12)

plt.axhline(y=0.02, xmin=(-1/(xlim[0]-xlim[1]))-(post_stats_66[0]/(xlim[0]-xlim[1])), 
            xmax=(-1/(xlim[0]-xlim[1]))-(post_stats_66[1]/(xlim[0]-xlim[1])), linestyle=':', c='#9933ff', label='17-83% estimate', linewidth=2)
plt.axhline(y=0.02, xmin=(-1/(xlim[0]-xlim[1]))-(post_stats_66[0]/(xlim[0]-xlim[1])), 
            xmax=(-1/(xlim[0]-xlim[1]))-(post_stats_66[0]/(xlim[0]-xlim[1])), marker='<', c='#9933ff')
plt.axhline(y=0.02, xmin=(-1/(xlim[0]-xlim[1]))-(post_stats_66[1]/(xlim[0]-xlim[1])), 
            xmax=(-1/(xlim[0]-xlim[1]))-(post_stats_66[1]/(xlim[0]-xlim[1])), marker='>', c='#9933ff')

plt.axhline(y=0.02, xmin=(-1/(xlim[0]-xlim[1]))-(post_median/(xlim[0]-xlim[1])),
            xmax=(-1/(xlim[0]-xlim[1]))-(post_median/(xlim[0]-xlim[1])), c='#9933ff', marker='.', ms=12)

ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['bottom'].set_position(('data', 0))
ax.spines['left'].set_position(('data', 0))
ax.spines['left'].set_linewidth(2)
ax.spines['bottom'].set_linewidth(2)
ax.tick_params('x', width=2)
ax.tick_params('y', width=2)
plt.xticks(ticks=np.arange(1,9,1),fontsize=14)
plt.yticks([0.1, 0.2, 0.3, 0.4, 0.5],fontsize=14)
plt.legend(loc='upper right', edgecolor='k')
plt.tight_layout()
#plt.savefig('Posterior_PlioMIP.pdf', dpi=300)