"""
Visualization for Representation Learning Dynamics experiment.
================================================================
Generates publication-quality figures from experiment results.
"""

import json
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from pathlib import Path
from typing import Dict, List, Optional
import argparse


def load_results(results_path: str) -> Dict:
    with open(results_path) as f:
        return json.load(f)


def extract_metric_series(history: List[Dict], metric_name: str) -> tuple:
    """Extract (steps, values) for a metric from history."""
    steps = [h['step'] for h in history if metric_name in h]
    values = [h[metric_name] for h in history if metric_name in h]
    return np.array(steps), np.array(values)


def plot_training_curves(results: Dict, output_dir: str):
    """Plot training loss and task accuracies across all phases."""
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))

    # Phase 1
    p1 = results['phase1_history']
    steps_p1 = [h['step'] for h in p1]
    loss_p1 = [h['train_loss'] for h in p1]
    acc_add_p1 = [h.get('eval/add_test_acc', 0) for h in p1]
    acc_sub_p1 = [h.get('eval/subtract_test_acc', 0) for h in p1]

    # Phase 2 A→A
    p2aa = results['phase2_aa_history']
    steps_aa = [h['step'] + steps_p1[-1] for h in p2aa] if p2aa else []
    loss_aa = [h['train_loss'] for h in p2aa]
    acc_add_aa = [h.get('eval/add_test_acc', 0) for h in p2aa]
    acc_sub_aa = [h.get('eval/subtract_test_acc', 0) for h in p2aa]

    # Phase 2 A→B
    p2ab = results['phase2_ab_history']
    steps_ab = [h['step'] + steps_p1[-1] for h in p2ab] if p2ab else []
    loss_ab = [h['train_loss'] for h in p2ab]
    acc_add_ab = [h.get('eval/add_test_acc', 0) for h in p2ab]
    acc_sub_ab = [h.get('eval/subtract_test_acc', 0) for h in p2ab]

    # Training loss
    ax = axes[0, 0]
    ax.plot(steps_p1, loss_p1, 'k-', label='Phase 1 (Add)', linewidth=2)
    if steps_aa:
        ax.plot(steps_aa, loss_aa, 'b-', label='A→A (Continue Add)', linewidth=2)
    if steps_ab:
        ax.plot(steps_ab, loss_ab, 'r-', label='A→B (Switch to Sub)', linewidth=2)
    ax.axvline(x=steps_p1[-1] if steps_p1 else 0, color='gray', linestyle='--',
               alpha=0.5, label='Phase transition')
    ax.set_xlabel('Training Step')
    ax.set_ylabel('Loss')
    ax.set_title('Training Loss')
    ax.legend()
    ax.set_yscale('log')

    # Addition accuracy
    ax = axes[0, 1]
    ax.plot(steps_p1, acc_add_p1, 'k-', label='Phase 1', linewidth=2)
    if steps_aa:
        ax.plot(steps_aa, acc_add_aa, 'b-', label='A→A', linewidth=2)
    if steps_ab:
        ax.plot(steps_ab, acc_add_ab, 'r-', label='A→B', linewidth=2)
    ax.axvline(x=steps_p1[-1] if steps_p1 else 0, color='gray',
               linestyle='--', alpha=0.5)
    ax.set_xlabel('Training Step')
    ax.set_ylabel('Accuracy')
    ax.set_title('Task A (Addition) Accuracy')
    ax.legend()
    ax.set_ylim(-0.05, 1.05)

    # Subtraction accuracy
    ax = axes[1, 0]
    ax.plot(steps_p1, acc_sub_p1, 'k-', label='Phase 1', linewidth=2)
    if steps_aa:
        ax.plot(steps_aa, acc_sub_aa, 'b-', label='A→A', linewidth=2)
    if steps_ab:
        ax.plot(steps_ab, acc_sub_ab, 'r-', label='A→B', linewidth=2)
    ax.axvline(x=steps_p1[-1] if steps_p1 else 0, color='gray',
               linestyle='--', alpha=0.5)
    ax.set_xlabel('Training Step')
    ax.set_ylabel('Accuracy')
    ax.set_title('Task B (Subtraction) Accuracy')
    ax.legend()
    ax.set_ylim(-0.05, 1.05)

    # Gradient alignment
    ax = axes[1, 1]
    ga_p1 = [h.get('gradient_alignment_a_vs_b', 0) for h in p1]
    ga_aa = [h.get('gradient_alignment_a_vs_b', 0) for h in p2aa]
    ga_ab = [h.get('gradient_alignment_a_vs_b', 0) for h in p2ab]
    ax.plot(steps_p1, ga_p1, 'k-', label='Phase 1', linewidth=2)
    if steps_aa:
        ax.plot(steps_aa, ga_aa, 'b-', label='A→A', linewidth=2)
    if steps_ab:
        ax.plot(steps_ab, ga_ab, 'r-', label='A→B', linewidth=2)
    ax.axvline(x=steps_p1[-1] if steps_p1 else 0, color='gray',
               linestyle='--', alpha=0.5)
    ax.axhline(y=0, color='gray', linestyle=':', alpha=0.3)
    ax.set_xlabel('Training Step')
    ax.set_ylabel('Cosine Similarity')
    ax.set_title('Gradient Alignment (Task A vs Task B)')
    ax.legend()

    plt.tight_layout()
    plt.savefig(f'{output_dir}/training_curves.png', dpi=150, bbox_inches='tight')
    plt.close()
    print(f"Saved: {output_dir}/training_curves.png")


def plot_cka_dynamics(results: Dict, output_dir: str):
    """Plot CKA drift from Phase 1 end across all layers."""
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    n_layers = results['config']['n_layers'] + 1

    for layer_idx in range(n_layers):
        metric = f'layer_{layer_idx}/cka_vs_phase1'

        # A→A branch
        p2aa = results['phase2_aa_history']
        steps_aa = [h['step'] for h in p2aa if metric in h]
        vals_aa = [h[metric] for h in p2aa if metric in h]

        # A→B branch
        p2ab = results['phase2_ab_history']
        steps_ab = [h['step'] for h in p2ab if metric in h]
        vals_ab = [h[metric] for h in p2ab if metric in h]

        label = f'Layer {layer_idx}' if layer_idx > 0 else 'Embedding'
        axes[0].plot(steps_aa, vals_aa, '-', label=label, linewidth=1.5)
        axes[1].plot(steps_ab, vals_ab, '-', label=label, linewidth=1.5)

    axes[0].set_title('Branch A→A: CKA vs Phase 1 End')
    axes[0].set_xlabel('Training Step')
    axes[0].set_ylabel('CKA Similarity')
    axes[0].legend()
    axes[0].set_ylim(0, 1.05)

    axes[1].set_title('Branch A→B: CKA vs Phase 1 End')
    axes[1].set_xlabel('Training Step')
    axes[1].set_ylabel('CKA Similarity')
    axes[1].legend()
    axes[1].set_ylim(0, 1.05)

    plt.tight_layout()
    plt.savefig(f'{output_dir}/cka_dynamics.png', dpi=150, bbox_inches='tight')
    plt.close()
    print(f"Saved: {output_dir}/cka_dynamics.png")


def plot_attention_entropy(results: Dict, output_dir: str):
    """Plot attention entropy per head over training."""
    n_layers = results['config']['n_layers']
    n_heads = results['config']['n_heads']

    fig, axes = plt.subplots(n_layers, 2, figsize=(14, 4 * n_layers))
    if n_layers == 1:
        axes = axes.reshape(1, 2)

    for layer_idx in range(n_layers):
        for head_idx in range(n_heads):
            metric = f'layer_{layer_idx+1}/head_{head_idx}_entropy'

            # A→A
            p2aa = results['phase2_aa_history']
            steps_aa = [h['step'] for h in p2aa if metric in h]
            vals_aa = [h[metric] for h in p2aa if metric in h]
            axes[layer_idx, 0].plot(steps_aa, vals_aa, label=f'Head {head_idx}')

            # A→B
            p2ab = results['phase2_ab_history']
            steps_ab = [h['step'] for h in p2ab if metric in h]
            vals_ab = [h[metric] for h in p2ab if metric in h]
            axes[layer_idx, 1].plot(steps_ab, vals_ab, label=f'Head {head_idx}')

        axes[layer_idx, 0].set_title(f'Layer {layer_idx+1} — A→A')
        axes[layer_idx, 0].set_ylabel('Entropy (bits)')
        axes[layer_idx, 0].legend()
        axes[layer_idx, 1].set_title(f'Layer {layer_idx+1} — A→B')
        axes[layer_idx, 1].legend()

    axes[-1, 0].set_xlabel('Training Step')
    axes[-1, 1].set_xlabel('Training Step')

    plt.tight_layout()
    plt.savefig(f'{output_dir}/attention_entropy.png', dpi=150, bbox_inches='tight')
    plt.close()
    print(f"Saved: {output_dir}/attention_entropy.png")


def plot_cka_heatmaps(results: Dict, output_dir: str):
    """Plot CKA cross-layer heatmaps for final model comparisons."""
    heatmaps = results['cka_heatmaps']

    fig, axes = plt.subplots(1, 3, figsize=(18, 5))

    titles = ['A→A vs A→B', 'A→A vs Phase 1 End', 'A→B vs Phase 1 End']
    keys = ['aa_vs_ab', 'aa_vs_p1', 'ab_vs_p1']

    for ax, title, key in zip(axes, titles, keys):
        hm = np.array(heatmaps[key])
        im = ax.imshow(hm, cmap='viridis', vmin=0, vmax=1, aspect='auto')
        ax.set_title(title)
        ax.set_xlabel('Layer (model 2)')
        ax.set_ylabel('Layer (model 1)')
        # Add text annotations
        for i in range(hm.shape[0]):
            for j in range(hm.shape[1]):
                color = 'white' if hm[i, j] < 0.5 else 'black'
                ax.text(j, i, f'{hm[i,j]:.2f}', ha='center', va='center',
                        fontsize=8, color=color)
        plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)

    plt.tight_layout()
    plt.savefig(f'{output_dir}/cka_heatmaps.png', dpi=150, bbox_inches='tight')
    plt.close()
    print(f"Saved: {output_dir}/cka_heatmaps.png")


def plot_subspace_angles(results: Dict, output_dir: str):
    """Plot subspace angle divergence between branches."""
    n_layers = results['config']['n_layers'] + 1

    fig, ax = plt.subplots(figsize=(10, 5))

    for layer_idx in range(n_layers):
        metric = f'layer_{layer_idx}/subspace_angle_vs_phase1'

        p2aa = results['phase2_aa_history']
        steps_aa = [h['step'] for h in p2aa if metric in h]
        vals_aa = [h[metric] for h in p2aa if metric in h]

        p2ab = results['phase2_ab_history']
        steps_ab = [h['step'] for h in p2ab if metric in h]
        vals_ab = [h[metric] for h in p2ab if metric in h]

        label = f'Layer {layer_idx}' if layer_idx > 0 else 'Embedding'
        if steps_aa:
            ax.plot(steps_aa, vals_aa, '--', label=f'{label} (A→A)',
                    alpha=0.7, linewidth=1.5)
        if steps_ab:
            ax.plot(steps_ab, vals_ab, '-', label=f'{label} (A→B)',
                    linewidth=2)

    ax.set_xlabel('Training Step')
    ax.set_ylabel('Mean Subspace Angle (degrees)')
    ax.set_title('Subspace Angle Drift from Phase 1 End')
    ax.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.tight_layout()
    plt.savefig(f'{output_dir}/subspace_angles.png', dpi=150, bbox_inches='tight')
    plt.close()
    print(f"Saved: {output_dir}/subspace_angles.png")


def plot_weight_changes(results: Dict, output_dir: str):
    """Plot weight change magnitude per block."""
    n_blocks = results['config']['n_layers']

    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    for block_idx in range(n_blocks):
        metric_init = f'block_{block_idx}/weight_change_from_init'
        metric_p1 = f'block_{block_idx}/weight_change_from_phase1'

        # A→A
        p2aa = results['phase2_aa_history']
        steps = [h['step'] for h in p2aa if metric_p1 in h]
        vals = [h[metric_p1] for h in p2aa if metric_p1 in h]
        axes[0].plot(steps, vals, label=f'Block {block_idx}', linewidth=2)

        # A→B
        p2ab = results['phase2_ab_history']
        steps = [h['step'] for h in p2ab if metric_p1 in h]
        vals = [h[metric_p1] for h in p2ab if metric_p1 in h]
        axes[1].plot(steps, vals, label=f'Block {block_idx}', linewidth=2)

    axes[0].set_title('A→A: Weight Change from Phase 1')
    axes[0].set_xlabel('Training Step')
    axes[0].set_ylabel('L2 Norm of Weight Delta')
    axes[0].legend()

    axes[1].set_title('A→B: Weight Change from Phase 1')
    axes[1].set_xlabel('Training Step')
    axes[1].set_ylabel('L2 Norm of Weight Delta')
    axes[1].legend()

    plt.tight_layout()
    plt.savefig(f'{output_dir}/weight_changes.png', dpi=150, bbox_inches='tight')
    plt.close()
    print(f"Saved: {output_dir}/weight_changes.png")


def generate_all_plots(results_path: str, output_dir: str = None):
    """Generate all visualization plots from experiment results."""
    results = load_results(results_path)
    if output_dir is None:
        output_dir = str(Path(results_path).parent)

    Path(output_dir).mkdir(parents=True, exist_ok=True)

    plot_training_curves(results, output_dir)
    plot_cka_dynamics(results, output_dir)
    plot_attention_entropy(results, output_dir)
    plot_cka_heatmaps(results, output_dir)
    plot_subspace_angles(results, output_dir)
    plot_weight_changes(results, output_dir)

    print(f"\nAll plots saved to {output_dir}/")


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--results', type=str, default='results/experiment_results.json')
    parser.add_argument('--output-dir', type=str, default=None)
    args = parser.parse_args()
    generate_all_plots(args.results, args.output_dir)