import os
import re
import numpy as np
import pandas as pd
from scipy import stats
import matplotlib.pyplot as plt

PLOT = False
START_YEAR = 2020
END_YEAR = 2100

n_runs = 100000

df = pd.read_csv('cauchy-values.csv', index_col=0)
years = np.arange(START_YEAR, END_YEAR + 1)

# Squeeze distribution to zero in 2020
df.loc[2020, 'scale'] = 0.0001
df.loc[2020, 'min'] = df.loc[2020, 'loc'] - 0.0001
df.loc[2020, 'max'] = df.loc[2020, 'loc'] + 0.0001

# Interpolate intermediate values
df = df.reindex(years).interpolate(method='cubic')

# Prepare array for SLR values
slr = np.zeros([len(years), n_runs], dtype=float)

for i, (year, row) in enumerate(df.iterrows()):
    # Get probability distribution
    dist = stats.cauchy(loc=row['loc'], scale=row['scale'])

    # Generate random samples
    for factor in range(2, 10):
        s_raw = dist.rvs(n_runs * factor)

        # Take first samples within valid range
        s = s_raw[(s_raw > row['min']) & (s_raw < row['max'])]

        if len(s) > n_runs:
            break  # Success
        else:
            continue  # We need more samples, so try larger factor

    # Add the requried number of samples
    slr[i] = s[:n_runs]

# Sort each row to make SLR trajectories smooth
slr = np.sort(slr, axis=1)

# Randomise run order (column-wise)
slr = np.random.permutation(slr.T).T

# Set first year to zero
slr[0, :] = df.loc[2020, 'loc']

# Plot first few trajectories
if PLOT:
    fig, ax = plt.subplots(1,
                           3,
                           figsize=(12, 5),
                           sharey=True,
                           gridspec_kw={
                               'wspace': 0.05,
                               'width_ratios': [3, 1, 1]
                           })

    ax[0].plot(years, slr[:, :100], c='#444444', lw=0.2)
    ax[1].hist(slr[-1, :],
               bins=100,
               fc='#cccccc',
               ec='#aaaaaa',
               orientation='horizontal')

    i = len(years) - 1
    dff = df.T.loc['5':'95', years[i]]
    ax[2].hist(
        slr[i, :],
        bins=100,
        fc='#cccccc',
        ec='#aaaaaa',
        orientation='horizontal',
        cumulative=True,
    )
    ax[2].plot(dff.index.astype(int) / 100 * n_runs,
               dff.values,
               'o',
               c='C3',
               label='IPCC AR6 data')

    ax[0].set_xlim(right=years[i])

    ax[0].set_title(f'SLR trajectories\n(first 100 out of {n_runs:,} runs)')
    ax[1].set_title(f'Probability\ndistribution\nin year {years[i]}')
    ax[2].set_title(f'Cumulative\ndistribution\nin year {years[i]}')
    ax[0].set_ylabel('SLR (m)', labelpad=10)

    ax[2].legend()

    ax[0].spines['top'].set_visible(False)
    ax[0].spines['right'].set_visible(False)

    for a in ax[1:]:
        a.spines['top'].set_visible(False)
        a.spines['right'].set_visible(False)
        a.spines['bottom'].set_visible(False)
        a.xaxis.set_visible(False)