"""Render a noise comparison from four native Robustness exports.

Run with Python, pandas, and matplotlib beside the saved result JSON files.
No simulation is run by this script. Lines connect measured conditions only.
"""

import json
from pathlib import Path

import matplotlib

matplotlib.use('Agg')
import matplotlib.pyplot as plt
import pandas as pd


def main():
    folder = Path(__file__).resolve().parent
    experiments = [(0, 'all-strategies-result.json'), (1, 'noise-01-result.json'),
                   (10, 'noise-10-result.json'), (25, 'noise-25-result.json')]
    frames, conditions = [], []
    reference = None
    for percent, filename in experiments:
        result = json.loads((folder / filename).read_text(encoding='utf-8'))
        config = result['config'].copy()
        assert config.pop('execution_error_rate') == percent / 100
        if reference is None:
            reference = config
        assert config == reference, 'Only the execution error probability may change'
        rows = pd.DataFrame(result['seed_results'])
        assert len(rows) == 600 and not rows.duplicated(['seed', 'strategy']).any()
        assert set(rows['seed']) == set(range(10))
        assert set(rows['strategy']) == set(config['strategy_names'])
        assert (rows['rounds'] == 5900).all()
        ranks = rows.groupby('seed')['points_per_round'].rank(method='min', ascending=False)
        assert (ranks == rows['rank']).all()
        summary = rows.groupby('strategy').agg(
            mean_payoff=('points_per_round', 'mean'),
            mean_cooperation=('cooperation_rate', 'mean'),
        )
        summary['rank'] = summary['mean_payoff'].rank(method='min', ascending=False).astype(int)
        summary['error_percent'] = percent
        frames.append(summary.reset_index())
        winner = summary['mean_payoff'].idxmax()
        conditions.append({'error_percent': percent, 'leader': winner,
                           'leader_payoff': summary.loc[winner, 'mean_payoff'],
                           'field_payoff': rows['points_per_round'].mean(),
                           'field_cooperation': rows['cooperation_rate'].mean()})
    combined = pd.concat(frames, ignore_index=True)
    combined.to_csv(folder / 'noise-comparison.csv', index=False)
    conditions = pd.DataFrame(conditions)
    conditions.to_csv(folder / 'noise-conditions.csv', index=False)
    ranking = combined.pivot(index='strategy', columns='error_percent', values='rank')
    ranking = ranking.sort_values(0)
    print(conditions.to_string(index=False))
    selected = ['PastTrauma', 'HoldingAGrudge', 'TitForTat', 'GenerousTitForTat', 'WinStayLoseShift', 'BadCop']
    print(ranking.loc[selected].to_string())
    print('Strategies changing rank at 1%:', int((ranking[0] != ranking[1]).sum()))
    print('Largest rank gains at 25%:', (ranking[0] - ranking[25]).sort_values(ascending=False).head().to_dict())

    paper, ink, muted = '#101414', '#edf4ef', '#b2c2ba'
    plt.rcParams.update({'font.family': 'DejaVu Sans', 'text.color': ink,
                         'axes.labelcolor': muted, 'svg.fonttype': 'path'})
    fig, ax = plt.subplots(figsize=(14, 10), facecolor=paper)
    fig.subplots_adjust(left=0.09, right=0.95, top=0.80, bottom=0.25)
    ax.set_facecolor(paper)
    cmap = plt.get_cmap('turbo')
    highlights = {'PastTrauma': '#4ade80', 'HoldingAGrudge': '#facc15',
                  'TitForTat': '#38bdf8', 'GenerousTitForTat': '#c4b5fd',
                  'WinStayLoseShift': '#f472b6', 'BadCop': '#fb923c'}
    for name, row in ranking.iterrows():
        color = highlights.get(name, cmap((row[0] - 1) / 59))
        chosen = name in selected
        ax.plot(range(4), row[[0, 1, 10, 25]], color=color,
                linewidth=2.7 if chosen else 1.25, alpha=1 if chosen else 0.42,
                marker='o', markersize=5 if chosen else 2.5,
                zorder=4 if chosen else 2, label=name if chosen else None)
    ax.set_ylim(61, 0)
    ax.set_xlim(-0.10, 3.10)
    ax.set_xticks(range(4), ['0% errors', '1% errors', '10% errors', '25% errors'], fontsize=12)
    ax.set_yticks([1, 10, 20, 30, 40, 50, 60])
    ax.set_ylabel('Rank by mean points per round across ten seeds', fontsize=12, labelpad=15)
    ax.tick_params(colors=muted)
    ax.grid(axis='y', color='#33413a', alpha=0.8)
    ax.grid(axis='x', color='#33413a', alpha=0.5)
    for spine in ax.spines.values():
        spine.set_visible(False)
    fig.text(0.09, 0.94, 'When mistakes reshuffle the field', fontsize=25, weight='bold')
    fig.text(0.09, 0.895, '60 strategies. The same ten seeds. Increasing chances of an intended move flipping.', fontsize=12, color=muted)
    fig.text(0.09, 0.855, 'Each line follows one strategy; crossings mark reversals in rank. Higher is better.', fontsize=12, color=muted)
    legend = ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.12), ncol=3,
                       frameon=False, fontsize=11, labelcolor=ink, handlelength=2.6)
    for line in legend.get_lines():
        line.set_linewidth(3)
    fig.text(0.09, 0.075, '100 rounds per match | 1 repetition | seeds 0 through 9 | hidden horizon | no self-play', fontsize=10, color=muted)
    fig.text(0.09, 0.045, 'Measured conditions are equally spaced. Connecting lines do not represent unmeasured runs or play over time.', fontsize=10, color=muted)
    fig.savefig(folder / 'noise-comparison.svg', facecolor=paper)
    svg = folder / 'noise-comparison.svg'
    svg.write_text('\n'.join(line.rstrip() for line in svg.read_text(encoding='utf-8').splitlines()) + '\n', encoding='utf-8')
    fig.savefig(folder / 'noise-comparison.png', dpi=160, facecolor=paper)
    plt.close(fig)


if __name__ == '__main__':
    main()