"""Aggregate actual trials into Figure 1 style plots and the flow-rate table."""
import csv
import json
from pathlib import Path

from common import case_guard, load_config, load_data, parse_args, save_json, write_csv
import numpy as np
import torch
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt


def read_trajectory(path, budget):
    with open(path) as f:
        rows = list(csv.DictReader(f))
    if len(rows) != budget or [int(r['iteration']) for r in rows] != list(range(1, budget + 1)):
        raise RuntimeError(f'Incomplete or nonsequential trial: {path}')
    regrets = np.array([float(r['regret']) for r in rows])
    if not np.isfinite(regrets).all() or (regrets < -1e-7).any():
        raise ValueError(f'Invalid regret: {path}')
    cumulative = regrets.cumsum()
    if not np.allclose(cumulative, [float(r['cumulative_regret']) for r in rows]):
        raise ValueError('Saved cumulative regret does not match sum of instant regrets')
    return regrets


def save_fig(fig, out, name):
    for ext in ('png', 'pdf'):
        fig.savefig(out / f'{name}.{ext}', dpi=300, bbox_inches='tight')
    plt.close(fig)


def main():
    args = parse_args(__doc__).parse_args()
    cfg = load_config(args.config)
    case_guard(cfg)
    out = Path(cfg['paths']['output_dir']) / 'results'
    n = cfg['experiment']['iterations']
    seeds = cfg['experiment']['seeds']
    times = np.arange(1, n + 1)
    actual = {'NOTS': np.stack([read_trajectory(out / f'trial_{s}/trajectory.csv', n) for s in seeds])}
    if cfg['experiment']['run_random_baseline']:
        actual['RS'] = np.stack([read_trajectory(out / f'rs_{s}.csv', n) for s in seeds])
    statistics = {}
    for method, regrets in actual.items():
        cum = regrets.cumsum(axis=1)
        avg = cum / times
        finals = avg[:, -1]
        statistics[method] = {'final_average_regret_mean': float(finals.mean()),
                              'final_average_regret_std': float(finals.std(ddof=0)),
                              'final_average_regret_sample_std': float(finals.std(ddof=1)),
                              'per_trial_final_average_regret': finals.tolist(),
                              'trials': len(seeds), 'iterations': n, 'provenance': 'measured'}
        write_csv(out / f'{method.lower()}_curves.csv', [
            {'iteration': int(t), 'cumulative_mean': float(cum[:, i].mean()),
             'cumulative_std': float(cum[:, i].std()), 'average_mean': float(avg[:, i].mean()),
             'average_std': float(avg[:, i].std())} for i, t in enumerate(times)])
    save_json(out / 'metrics.json', {'metric_definition': 'mean_trials(sum_t(true_max - true_selected_t)/300)',
                                    'std_ddof': 0, 'seeds': seeds, 'methods': statistics,
                                    'published_references': cfg['reference'],
                                    'unrerun_baselines': ['BO', 'BFO', 'GP-TS', 'STO-NTS']})
    rows = []
    for method, ref in cfg['reference'].items():
        value = statistics.get(method)
        rows.append({'method': method, 'measured_mean': value['final_average_regret_mean'] if value else '',
                     'measured_std': value['final_average_regret_std'] if value else '',
                     'paper_mean': ref['mean'], 'paper_std': ref['std'],
                     'status': 'measured' if value else 'published_reference_only'})
    write_csv(out / 'table1_flow_rates.csv', rows)
    plt.rcParams.update({'font.size': 10, 'axes.spines.top': False, 'axes.spines.right': False})
    fig = plt.figure(figsize=(12, 7.2), constrained_layout=True)
    grid = fig.add_gridspec(2, 4, height_ratios=[1.1, 1])
    ax0, ax1 = fig.add_subplot(grid[0, :2]), fig.add_subplot(grid[0, 2:])
    for method, regrets in actual.items():
        cumulative = regrets.cumsum(axis=1)
        for ax, matrix in [(ax0, cumulative), (ax1, cumulative / times)]:
            mean, std = matrix.mean(0), matrix.std(0)
            line, = ax.plot(times, mean, label=f'{method} (measured)', lw=2)
            ax.fill_between(times, mean - std, mean + std, color=line.get_color(), alpha=0.18)
    ax0.set_ylabel('Cumulative regret $R_t$')
    ax1.set_ylabel('Average regret $R_t/t$')
    for ax in (ax0, ax1):
        ax.set_xlabel('Oracle query')
        ax.legend()
        ax.grid(alpha=0.2)
    x, y, meta = load_data(cfg)
    oracle = torch.load(out / 'oracle.pt', map_location='cpu', weights_only=True)
    best, worst = oracle['best_index'], oracle['worst_index']
    # Common scales for best/worst pressure, original permeability for display.
    panels = [(x[best, 0], 'Best permeability a', 'gray', x.min().item(), x.max().item()),
              (y[best, 0], 'Best pressure u', 'viridis', y.min().item(), y.max().item()),
              (x[worst, 0], 'Worst permeability a', 'gray', x.min().item(), x.max().item()),
              (y[worst, 0], 'Worst pressure u', 'viridis', y.min().item(), y.max().item())]
    for i, (a, title, cmap, vmin, vmax) in enumerate(panels):
        ax = fig.add_subplot(grid[1, i])
        im = ax.imshow(a.numpy(), origin='lower', extent=[0, 1, 0, 1], cmap=cmap,
                       interpolation='nearest', vmin=vmin, vmax=vmax)
        ax.set_title(title)
        ax.set_xlabel('x')
        ax.set_ylabel('y')
        fig.colorbar(im, ax=ax, shrink=0.7)
    fig.suptitle('Darcy negative total flow rates: 10 independent trials, 300 queries\nShading: ±1 trial standard deviation; reference-only methods are omitted')
    save_fig(fig, out, 'figure1')
    inf = torch.load(out / 'inference_test.pt', map_location='cpu', weights_only=True)
    target, pred = inf['target'][0, 0].numpy(), inf['prediction'][0, 0].numpy()
    fig, axes = plt.subplots(1, 4, figsize=(12, 3), constrained_layout=True)
    for ax, field, title, cmap in zip(axes,
        [inf['x'][0, 0].numpy(), target, pred, np.abs(pred - target)],
        ['Test permeability', 'Oracle pressure', 'Loaded checkpoint pressure', 'Absolute error'],
        ['gray', 'viridis', 'viridis', 'magma']):
        kwargs = {'vmin': float(min(target.min(), pred.min())), 'vmax': float(max(target.max(), pred.max()))} if 'pressure' in title else {}
        im = ax.imshow(field, origin='lower', extent=[0, 1, 0, 1], cmap=cmap, **kwargs)
        ax.set_title(title)
        fig.colorbar(im, ax=ax, shrink=0.7)
    save_fig(fig, out, 'checkpoint_inference')
    fig, ax = plt.subplots(figsize=(8, 4), constrained_layout=True)
    for seed in seeds:
        with (out / f'trial_{seed}/training_loss.jsonl').open() as f:
            losses = [json.loads(line) for line in f]
        final_epochs = {r['iteration']: r['mse'] for r in losses if r['epoch'] == cfg['training']['epochs']}
        ax.plot(list(final_epochs), list(final_epochs.values()), alpha=0.6, label=f'seed {seed}')
    ax.set(xlabel='Oracle query', ylabel='Final epoch training pressure MSE', yscale='log')
    ax.legend(ncol=2, fontsize=8)
    save_fig(fig, out, 'training_loss')
    measured = statistics['NOTS']
    report = ('# Darcy flow rates reproduction results\n\n'
              f'NOTS measured R300/300 = {measured["final_average_regret_mean"]:.8g} ± '
              f'{measured["final_average_regret_std"]:.8g} (population std, ten trials).\n\n'
              'Paper Table 1 reference: 0.012 ± 0.001. This is a literature value, not measured data.\n\n'
              'All measured trajectories contain 300 queries including one random initial query. '
              'Training and acquisition use the same physical pressure scale. '
              'The paper main-text flow sign conflicts with Appendix D.1; the latter is used. '
              'Published BO/BFO/GP-TS/STO-NTS entries are reference-only, not reruns.\n')
    (out / 'report.md').write_text(report)
    print('RESULT ' + json.dumps(statistics), flush=True)


if __name__ == '__main__':
    main()