#!/usr/bin/env python3
"""Full training analysis: plot curves, verify reward accuracy, analyze degradation.

Loads TensorBoard data from both Bias and NoBias 500ep trainings,
plots training curves with rolling averages, and analyzes the degradation pattern.

Usage:
    conda run -n pycuda_3_10 python -u analyze_final.py
"""
from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

_REPO = Path(__file__).resolve().parents[3]
if str(_REPO) not in sys.path:
    sys.path.insert(0, str(_REPO))

from tensorboard.backend.event_processing.event_accumulator import EventAccumulator

OUT_DIR = Path(__file__).resolve().parent / "output" / "degradation_diag"
BIAS_DIR = Path(__file__).resolve().parent / "output" / "bias_seed42_s512_lr0.0003_e3_v3long"
NOBIAS_DIR = Path(__file__).resolve().parent / "output" / "nobias_seed42_s512_lr0.0003_e3_v3long"


def load_tb(tb_dir, tag):
    ea = EventAccumulator(str(tb_dir))
    ea.Reload()
    tags = ea.Tags()['scalars']
    if tag not in tags:
        return None
    events = ea.Scalars(tag)
    by_step = {}
    for e in events:
        by_step[e.step] = e.value
    steps = sorted(by_step.keys())
    return np.array(steps), np.array([by_step[s] for s in steps])


def rolling_avg(vals, window=20):
    """Rolling average with edge handling."""
    if len(vals) < window:
        return vals
    result = np.convolve(vals, np.ones(window)/window, mode='valid')
    # Pad edges
    pad = (len(vals) - len(result)) // 2
    result = np.pad(result, (pad, len(vals) - len(result) - pad), mode='edge')
    return result


def main():
    print("=== Final Training Analysis ===")

    all_data = {}
    for name, run_dir in [("Bias", BIAS_DIR), ("NoBias", NOBIAS_DIR)]:
        data = {}
        for tag in ['eval/avg_reward', 'eval/r_cd', 'eval/r_cl', 'eval/r_sim']:
            short = tag.split('/')[-1]
            result = load_tb(run_dir / "tb", tag)
            if result is not None:
                data[short] = result
        all_data[name] = data
        if 'avg_reward' in data:
            steps, rewards = data['avg_reward']
            print(f"  {name}: {len(steps)} episodes, best={np.max(rewards):.4f} at Ep{steps[np.argmax(rewards)]}")

    # --- Plot 1: Training curves with rolling average ---
    fig, axes = plt.subplots(2, 2, figsize=(16, 12))
    fig.suptitle('Training Curves: Bias vs NoBias (V3, 500ep, double_buffer)', fontsize=14)

    titles = {
        'avg_reward': 'Total Reward',
        'r_cd': 'r_cd (drag, Gaussian+EMA)',
        'r_cl': 'r_cl (lift, Gaussian+EMA)',
        'r_sim': 'r_sim (normalized DTW)',
    }
    colors = {"Bias": "blue", "NoBias": "red"}

    for idx, (key, title) in enumerate(titles.items()):
        ax = axes[idx // 2, idx % 2]
        for name, data in all_data.items():
            if key not in data:
                continue
            steps, vals = data[key]
            # Raw
            ax.plot(steps, vals, alpha=0.2, color=colors[name], linewidth=0.5)
            # Rolling average
            roll = rolling_avg(vals, window=20)
            ax.plot(steps, roll, color=colors[name], linewidth=2, label=f"{name} (20-ep avg)")
            # Best line
            best_idx = np.argmax(vals)
            ax.axhline(vals[best_idx], color=colors[name], linestyle=':', alpha=0.3)

        ax.set_title(title)
        ax.set_xlabel("Episode")
        ax.legend()
        ax.grid(True, alpha=0.3)

    plt.tight_layout()
    out1 = OUT_DIR / "final_training_curves.png"
    plt.savefig(str(out1), dpi=120)
    plt.close()
    print(f"  Saved: {out1}")

    # --- Plot 2: Degradation analysis ---
    fig, axes = plt.subplots(2, 1, figsize=(14, 10))
    fig.suptitle('Degradation Analysis: Pre-peak vs Post-peak', fontsize=14)

    for idx, (name, data) in enumerate(all_data.items()):
        if 'avg_reward' not in data:
            continue
        steps, rewards = data['avg_reward']
        best_idx = np.argmax(rewards)
        best_ep = steps[best_idx]

        ax = axes[idx]
        ax.plot(steps, rewards, alpha=0.3, color='gray', linewidth=0.5)
        ax.plot(steps, rolling_avg(rewards, 20), color='blue', linewidth=2, label='20-ep avg')
        ax.axvline(best_ep, color='red', linestyle='--', alpha=0.7, label=f'Peak at Ep{best_ep}')
        ax.axhline(rewards[best_idx], color='red', linestyle=':', alpha=0.3, label=f'Best={rewards[best_idx]:.3f}')

        # Mark regions
        pre_peak = rewards[:best_idx]
        post_peak = rewards[best_idx:]
        ax.axhspan(np.mean(post_peak) - np.std(post_peak),
                   np.mean(post_peak) + np.std(post_peak),
                   alpha=0.1, color='orange', label=f'Post-peak: {np.mean(post_peak):.3f}±{np.std(post_peak):.3f}')

        ax.set_title(f"{name}: Peak={rewards[best_idx]:.3f} (Ep{best_ep}), Post-peak mean={np.mean(post_peak):.3f}")
        ax.set_xlabel("Episode")
        ax.legend(fontsize=8)
        ax.grid(True, alpha=0.3)

    plt.tight_layout()
    out2 = OUT_DIR / "degradation_analysis.png"
    plt.savefig(str(out2), dpi=120)
    plt.close()
    print(f"  Saved: {out2}")

    # --- Print detailed analysis ---
    print("\n" + "=" * 90)
    print("DETAILED ANALYSIS")
    print("=" * 90)

    for name, data in all_data.items():
        if 'avg_reward' not in data:
            continue
        steps, rewards = data['avg_reward']
        best_idx = np.argmax(rewards)
        best_ep = steps[best_idx]

        print(f"\n--- {name} ---")
        print(f"  Total episodes: {len(steps)}")
        print(f"  Best reward: {rewards[best_idx]:.4f} at Ep{best_ep}")

        # Phase analysis
        for phase_name, sl in [("Early (1-50)", slice(0, 50)),
                                ("Mid (50-best)", slice(50, best_idx)),
                                (f"Peak (Ep{best_ep})", slice(best_idx, best_idx+1)),
                                ("Post-peak+50", slice(best_idx, best_idx+50)),
                                ("Post-peak+100", slice(best_idx, best_idx+100)),
                                ("Last 50", slice(max(0, len(rewards)-50), len(rewards)))]:
            seg_start = sl.start if sl.start is not None else 0
            seg_end = sl.stop if sl.stop is not None else len(rewards)
            if seg_start >= len(rewards) or seg_end > len(rewards):
                continue
            segment = rewards[sl]
            if len(segment) == 0:
                continue
            print(f"  {phase_name:>20}: mean={np.mean(segment):.4f}, std={np.std(segment):.4f}, "
                  f"min={np.min(segment):.4f}, max={np.max(segment):.4f}")

        # Component analysis at peak vs post-peak
        for comp in ['r_cd', 'r_cl', 'r_sim']:
            if comp not in data:
                continue
            _, vals = data[comp]
            peak_val = vals[best_idx]
            post50 = vals[best_idx:best_idx+50]
            last50 = vals[-50:]
            print(f"  {comp}: peak={peak_val:.4f}, post50_mean={np.mean(post50):.4f}, "
                  f"last50_mean={np.mean(last50):.4f}")

    # --- Reward accuracy check ---
    print("\n" + "=" * 90)
    print("REWARD ACCURACY CHECK")
    print("=" * 90)
    print("  Symmetry-off during eval: YES (code verified)")
    print("  NoBias Ep1 should be ~0.08 (zero action, no control):")
    if 'NoBias' in all_data and 'avg_reward' in all_data['NoBias']:
        _, r = all_data['NoBias']['avg_reward']
        print(f"    NoBias Ep1 = {r[0]:.4f} -> {'CORRECT' if 0.05 < r[0] < 0.15 else 'SUSPICIOUS'}")
    print("  Bias Ep1 should be ~0.55 (bias action, partial control):")
    if 'Bias' in all_data and 'avg_reward' in all_data['Bias']:
        _, r = all_data['Bias']['avg_reward']
        print(f"    Bias Ep1 = {r[0]:.4f} -> {'CORRECT' if 0.50 < r[0] < 0.65 else 'SUSPICIOUS'}")

    print("\n=== Done ===")


if __name__ == "__main__":
    main()
