"""Figures and summary table for the season replay (dissertation Chapter 4)."""
from pathlib import Path
import sys
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

OUT = Path(sys.argv[1])
d = pd.read_csv(OUT / 'season_replay_metrics.csv')
COLORS = {'learning_loop': '#2a78d6', 'synthetic_only': '#eb6834', 'real_only': '#1baf7a'}
LABELS = {'learning_loop': 'Learning loop (prior + real outcomes)', 'synthetic_only': 'Synthetic only (never updated)',
          'real_only': 'Real outcomes only (no prior)'}
MARKERS = {'learning_loop': 'o', 'synthetic_only': 's', 'real_only': '^'}
plt.rcParams.update({'font.family': 'DejaVu Sans', 'font.size': 10, 'axes.edgecolor': '#b5b4ad',
                     'axes.labelcolor': '#52514e', 'xtick.color': '#52514e', 'ytick.color': '#52514e',
                     'axes.spines.top': False, 'axes.spines.right': False})


def panel(ax, frame, metric, title, ylabel):
    for method in ['synthetic_only', 'real_only', 'learning_loop']:
        g = frame[frame.method.eq(method)].groupby('season')[metric].agg(['mean', 'std']).reset_index()
        if not len(g):
            continue
        ax.plot(g.season, g['mean'], color=COLORS[method], lw=2, marker=MARKERS[method], ms=6,
                label=LABELS[method], zorder=3)
        ax.fill_between(g.season, g['mean'] - g['std'].fillna(0), g['mean'] + g['std'].fillna(0),
                        color=COLORS[method], alpha=.12, lw=0)
    ax.set_title(title, loc='left', fontsize=11, color='#0b0b0b')
    ax.set_xlabel('Seasons of recorded outcomes')
    ax.set_ylabel(ylabel)
    ax.grid(axis='y', color='#e6e5e0', lw=.8)
    ax.set_xticks(sorted(frame.season.unique()))


crops = [('maize', 'Maize'), ('potato', 'Potato'), ('tomato', 'Tomato')]
fig, axes = plt.subplots(1, 3, figsize=(13, 4.2))
for ax, (crop, name) in zip(axes, crops):
    panel(ax, d[d.target.eq('yield_' + crop)], 'mae_t_ha', f'{name} yield error', 'Mean absolute error (t/ha)')
    ax.set_ylim(bottom=0)
handles, labels = axes[0].get_legend_handles_labels()
fig.legend(handles, labels, loc='lower center', ncol=3, frameon=False, bbox_to_anchor=(.5, -.02))
fig.suptitle('Yield forecast error on 1,500 unseen farms per crop (mean of 3 runs, band = 1 s.d.)', x=.01, ha='left',
             fontsize=12, color='#0b0b0b')
fig.tight_layout(rect=(0, .08, 1, .95))
fig.savefig(OUT / 'fig_yield_learning_curve.png', dpi=200)

c = d[d.target.eq('credit')]
fig, axes = plt.subplots(1, 3, figsize=(13, 4.2))
panel(axes[0], c, 'log_loss', 'Repayment log loss (lower is better)', 'Log loss')
panel(axes[1], c, 'ece', 'Calibration error (lower is better)', 'Expected calibration error')
panel(axes[2], c, 'roc_auc', 'Ranking power (higher is better)', 'ROC AUC')
axes[1].set_ylim(bottom=0)
handles, labels = axes[0].get_legend_handles_labels()
fig.legend(handles, labels, loc='lower center', ncol=3, frameon=False, bbox_to_anchor=(.5, -.02))
fig.suptitle('Repayment model on 4,500 unseen loans (mean of 3 runs, band = 1 s.d.)', x=.01, ha='left',
             fontsize=12, color='#0b0b0b')
fig.tight_layout(rect=(0, .08, 1, .95))
fig.savefig(OUT / 'fig_credit_learning_curve.png', dpi=200)

# Summary table: season 0, 1 and final season.
last = int(d.season.max())
rows = []
for target in ['yield_maize', 'yield_potato', 'yield_tomato', 'credit']:
    metric = 'mae_t_ha' if target != 'credit' else 'log_loss'
    for method in ['synthetic_only', 'real_only', 'learning_loop']:
        g = d[d.target.eq(target) & d.method.eq(method)].groupby('season')[metric].mean()
        row = dict(target=target, metric=metric, method=method)
        for s in [0, 1, 3, last]:
            row[f'season_{s}'] = round(float(g.get(s, float('nan'))), 3)
        rows.append(row)
extra = []
for method in ['synthetic_only', 'real_only', 'learning_loop']:
    g = c[c.method.eq(method)].groupby('season')[['roc_auc', 'ece', 'brier']].mean()
    for metric in ['roc_auc', 'ece', 'brier']:
        row = dict(target='credit', metric=metric, method=method)
        for s in [0, 1, 3, last]:
            row[f'season_{s}'] = round(float(g[metric].get(s, float('nan'))), 3)
        extra.append(row)
table = pd.DataFrame(rows + extra)
table.to_csv(OUT / 'season_replay_summary_table.csv', index=False)
print(table.to_string(index=False))