#!/usr/bin/env python3
"""
Trajectory Analysis Script
Analyzes training trajectories for signs of model collapse/degradation.

Usage:
    python analyze_trajectories.py [--trajectories-dir PATH] [--output-dir PATH]

Requirements:
    pip install matplotlib
"""

import json
import os
import glob
import argparse
from collections import defaultdict

def parse_args():
    parser = argparse.ArgumentParser(description='Analyze training trajectories for issues')
    parser.add_argument('--trajectories-dir', type=str, default='trajectories',
                        help='Path to trajectories directory (default: trajectories)')
    parser.add_argument('--output-dir', type=str, default='.',
                        help='Path to output directory for graphs (default: current directory)')
    parser.add_argument('--no-graph', action='store_true',
                        help='Skip generating graphs (text analysis only)')
    return parser.parse_args()


def collect_trajectory_data(trajectories_dir):
    """Collect and analyze all trajectory data from the given directory."""
    steps_data = []
    
    # Get all step directories
    step_dirs = [d for d in os.listdir(trajectories_dir) if d.startswith('step_')]
    steps = sorted([int(d.replace('step_', '')) for d in step_dirs])
    
    print(f"Found {len(steps)} training steps (range: {min(steps)} to {max(steps)})")
    print(f"Halfway point: step {steps[len(steps)//2]}")
    print()
    
    for step in steps:
        step_dir = os.path.join(trajectories_dir, f"step_{step}", "train")
        files = glob.glob(os.path.join(step_dir, "*.json"))
        
        total = len(files)
        if total == 0:
            continue
            
        # Metrics to collect
        tool_call_count = 0
        truncated_think_count = 0
        endoftext_count = 0
        im_end_count = 0
        empty_parsed_count = 0
        very_short_count = 0
        
        rewards = []
        tokens = []
        num_steps = []
        
        for f in files:
            try:
                with open(f) as fp:
                    data = json.load(fp)
                
                raw = str(data.get('raw_final_message', ''))
                parsed = str(data.get('parsed_final_message', ''))
                
                # Check for issues
                if 'tool_call' in raw.lower():
                    tool_call_count += 1
                if '<think>' in raw and '</think>' not in raw:
                    truncated_think_count += 1
                if '<|endoftext|>' in raw:
                    endoftext_count += 1
                if '<|im_end|>' in raw:
                    im_end_count += 1
                if not parsed or len(parsed.strip()) < 10:
                    empty_parsed_count += 1
                if len(raw) < 50:
                    very_short_count += 1
                
                # Collect metrics
                rewards.append(data.get('total_reward', 0))
                if 'metrics_dict' in data:
                    tokens.append(data['metrics_dict'].get('tokens', 0))
                    num_steps.append(data['metrics_dict'].get('steps', 0))
                    
            except Exception as e:
                continue
        
        steps_data.append({
            'step': step,
            'total_files': total,
            'tool_call_pct': tool_call_count / total * 100,
            'truncated_think_pct': truncated_think_count / total * 100,
            'endoftext_pct': endoftext_count / total * 100,
            'im_end_pct': im_end_count / total * 100,
            'empty_parsed_pct': empty_parsed_count / total * 100,
            'very_short_pct': very_short_count / total * 100,
            'avg_reward': sum(rewards) / len(rewards) if rewards else 0,
            'min_reward': min(rewards) if rewards else 0,
            'max_reward': max(rewards) if rewards else 0,
            'avg_tokens': sum(tokens) / len(tokens) if tokens else 0,
            'avg_steps': sum(num_steps) / len(num_steps) if num_steps else 0,
        })
    
    return steps_data


def print_statistics_table(steps_data):
    """Print detailed statistics table."""
    print("=" * 100)
    print("STATISTICS BY TRAINING STEP")
    print("=" * 100)
    print(f"{'Step':<6} {'Files':<6} {'AvgReward':<10} {'MinReward':<10} {'MaxReward':<10} {'AvgTokens':<10} {'AvgSteps':<8}")
    print("-" * 100)
    
    for d in steps_data:
        print(f"{d['step']:<6} {d['total_files']:<6} {d['avg_reward']:<10.3f} {d['min_reward']:<10.3f} "
              f"{d['max_reward']:<10.3f} {d['avg_tokens']:<10.1f} {d['avg_steps']:<8.1f}")


def print_issues_table(steps_data):
    """Print issues detected by step."""
    print("\n" + "=" * 100)
    print("ISSUES DETECTED BY TRAINING STEP")
    print("=" * 100)
    print(f"{'Step':<6} {'ToolCall%':<10} {'TruncThink%':<12} {'EndOfText%':<11} {'EmptyParsed%':<13} {'VeryShort%':<11}")
    print("-" * 100)
    
    for d in steps_data:
        # Only print steps with issues
        if (d['tool_call_pct'] > 5 or d['truncated_think_pct'] > 1 or 
            d['endoftext_pct'] > 1 or d['empty_parsed_pct'] > 5):
            print(f"{d['step']:<6} {d['tool_call_pct']:<10.1f} {d['truncated_think_pct']:<12.1f} "
                  f"{d['endoftext_pct']:<11.1f} {d['empty_parsed_pct']:<13.1f} {d['very_short_pct']:<11.1f}")


def compute_aggregate_comparison(steps_data, halfway_step):
    """Compute aggregate statistics before and after halfway point."""
    metrics = ['tool_call_pct', 'truncated_think_pct', 'endoftext_pct', 'empty_parsed_pct', 'very_short_pct']
    
    before = {m: [] for m in metrics}
    before['rewards'] = []
    before['tokens'] = []
    before['total_files'] = 0
    
    after = {m: [] for m in metrics}
    after['rewards'] = []
    after['tokens'] = []
    after['total_files'] = 0
    
    for d in steps_data:
        target = before if d['step'] < halfway_step else after
        target['total_files'] += d['total_files']
        target['rewards'].append(d['avg_reward'])
        target['tokens'].append(d['avg_tokens'])
        for m in metrics:
            target[m].append(d[m])
    
    return before, after


def print_aggregate_comparison(before, after, halfway_step):
    """Print aggregate comparison before/after halfway point."""
    print("\n" + "=" * 100)
    print(f"AGGREGATE COMPARISON: Before vs After Step {halfway_step}")
    print("=" * 100)
    
    metrics = [
        ('tool_call_pct', 'Tool calls in final message'),
        ('truncated_think_pct', 'Truncated <think> tags'),
        ('endoftext_pct', '<|endoftext|> in output'),
        ('empty_parsed_pct', 'Empty parsed messages'),
        ('very_short_pct', 'Very short responses'),
    ]
    
    print(f"\n{'Metric':<35} {'Before':<15} {'After':<15} {'Change':<15}")
    print("-" * 80)
    
    for key, name in metrics:
        before_avg = sum(before[key]) / len(before[key]) if before[key] else 0
        after_avg = sum(after[key]) / len(after[key]) if after[key] else 0
        change = after_avg - before_avg
        print(f"{name:<35} {before_avg:<15.1f}% {after_avg:<15.1f}% {change:+.1f}%")
    
    # Rewards and tokens
    before_reward = sum(before['rewards']) / len(before['rewards']) if before['rewards'] else 0
    after_reward = sum(after['rewards']) / len(after['rewards']) if after['rewards'] else 0
    before_tokens = sum(before['tokens']) / len(before['tokens']) if before['tokens'] else 0
    after_tokens = sum(after['tokens']) / len(after['tokens']) if after['tokens'] else 0
    
    print(f"\n{'Average Reward':<35} {before_reward:<15.3f} {after_reward:<15.3f} {after_reward - before_reward:+.3f}")
    print(f"{'Average Tokens':<35} {before_tokens:<15.1f} {after_tokens:<15.1f} {after_tokens - before_tokens:+.1f}")
    
    return {
        'before_tool_call': sum(before['tool_call_pct']) / len(before['tool_call_pct']),
        'after_tool_call': sum(after['tool_call_pct']) / len(after['tool_call_pct']),
        'before_truncated': sum(before['truncated_think_pct']) / len(before['truncated_think_pct']),
        'after_truncated': sum(after['truncated_think_pct']) / len(after['truncated_think_pct']),
        'before_endoftext': sum(before['endoftext_pct']) / len(before['endoftext_pct']),
        'after_endoftext': sum(after['endoftext_pct']) / len(after['endoftext_pct']),
    }


def print_key_findings(comparison_stats):
    """Print key findings summary."""
    print("\n" + "=" * 100)
    print("KEY FINDINGS")
    print("=" * 100)
    
    findings = []
    
    tool_change = comparison_stats['after_tool_call'] - comparison_stats['before_tool_call']
    if tool_change > 10:
        findings.append(f"1. CRITICAL: Tool calls in final message increased by {tool_change:+.1f}% "
                       f"({comparison_stats['before_tool_call']:.1f}% → {comparison_stats['after_tool_call']:.1f}%)")
        findings.append("   → Model is confused about when to use tools vs. produce final answers")
    
    trunc_change = comparison_stats['after_truncated'] - comparison_stats['before_truncated']
    if trunc_change > 5:
        findings.append(f"\n2. CRITICAL: Truncated <think> tags increased by {trunc_change:+.1f}% "
                       f"({comparison_stats['before_truncated']:.1f}% → {comparison_stats['after_truncated']:.1f}%)")
        findings.append("   → Model failing to complete reasoning process")
    
    eof_change = comparison_stats['after_endoftext'] - comparison_stats['before_endoftext']
    if eof_change > 2:
        findings.append(f"\n3. WARNING: EndOfText tokens in output increased by {eof_change:+.1f}% "
                       f"({comparison_stats['before_endoftext']:.1f}% → {comparison_stats['after_endoftext']:.1f}%)")
        findings.append("   → Special tokens leaking into generated text (model degeneration)")
    
    if findings:
        for f in findings:
            print(f)
        
        print("\n" + "-" * 80)
        print("DIAGNOSIS: Training collapse/mode collapse detected")
        print("-" * 80)
        print("""
Recommendations:
1. Reduce learning rate or implement learning rate decay
2. Add KL divergence penalty to prevent drift from base model
3. Check for reward hacking around the collapse point
4. Investigate batch composition changes
5. Consider early stopping or checkpointing before collapse
""")
    else:
        print("No significant issues detected in the trajectory data.")


def generate_graphs(steps_data, output_dir):
    """Generate visualization graphs."""
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        print("\nWARNING: matplotlib not installed. Skipping graph generation.")
        print("Install with: pip install matplotlib")
        return
    
    # Extract data for plotting
    x = [d['step'] for d in steps_data]
    tool_call_pct = [d['tool_call_pct'] for d in steps_data]
    truncated_think_pct = [d['truncated_think_pct'] for d in steps_data]
    endoftext_pct = [d['endoftext_pct'] for d in steps_data]
    avg_reward = [d['avg_reward'] for d in steps_data]
    avg_tokens = [d['avg_tokens'] for d in steps_data]
    
    halfway = x[len(x)//2]
    
    # Create figure with 4 subplots
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    fig.suptitle(f'Training Trajectory Analysis - Issues Around Step {halfway}', 
                 fontsize=14, fontweight='bold')
    
    # 1. Tool calls in final message
    ax1 = axes[0, 0]
    ax1.plot(x, tool_call_pct, 'b-', linewidth=1.5, label='Tool calls in final msg')
    ax1.axvline(x=halfway, color='red', linestyle='--', alpha=0.7, label=f'Halfway (step {halfway})')
    ax1.fill_between(x, tool_call_pct, alpha=0.3)
    ax1.set_xlabel('Training Step')
    ax1.set_ylabel('Percentage (%)')
    ax1.set_title('1. Tool Calls in Final Message\n(Model confused about when to use tools)', fontsize=11)
    ax1.legend(loc='upper left')
    ax1.grid(True, alpha=0.3)
    ax1.set_ylim(0, 100)
    
    # 2. Truncated think tags
    ax2 = axes[0, 1]
    ax2.plot(x, truncated_think_pct, 'orange', linewidth=1.5, label='Truncated <think> tags')
    ax2.axvline(x=halfway, color='red', linestyle='--', alpha=0.7, label=f'Halfway (step {halfway})')
    ax2.fill_between(x, truncated_think_pct, alpha=0.3, color='orange')
    ax2.set_xlabel('Training Step')
    ax2.set_ylabel('Percentage (%)')
    ax2.set_title('2. Truncated <think> Tags\n(Incomplete reasoning blocks)', fontsize=11)
    ax2.legend(loc='upper left')
    ax2.grid(True, alpha=0.3)
    ax2.set_ylim(0, 100)
    
    # 3. EndOfText tokens
    ax3 = axes[1, 0]
    ax3.plot(x, endoftext_pct, 'green', linewidth=1.5, label='<|endoftext|> in output')
    ax3.axvline(x=halfway, color='red', linestyle='--', alpha=0.7, label=f'Halfway (step {halfway})')
    ax3.fill_between(x, endoftext_pct, alpha=0.3, color='green')
    ax3.set_xlabel('Training Step')
    ax3.set_ylabel('Percentage (%)')
    ax3.set_title('3. EndOfText Tokens in Output\n(Special tokens leaking into text)', fontsize=11)
    ax3.legend(loc='upper left')
    ax3.grid(True, alpha=0.3)
    if max(endoftext_pct) > 0:
        ax3.set_ylim(0, max(endoftext_pct) * 1.2)
    
    # 4. Average reward with token count
    ax4 = axes[1, 1]
    ax4.plot(x, avg_reward, 'purple', linewidth=1.5, label='Avg Reward')
    ax4.axvline(x=halfway, color='red', linestyle='--', alpha=0.7, label=f'Halfway (step {halfway})')
    ax4.axhline(y=0, color='gray', linestyle='-', alpha=0.5)
    ax4.set_xlabel('Training Step')
    ax4.set_ylabel('Average Reward', color='purple')
    ax4.tick_params(axis='y', labelcolor='purple')
    ax4.set_title('4. Average Reward & Token Count\n(Performance collapse)', fontsize=11)
    ax4.grid(True, alpha=0.3)
    
    # Secondary y-axis for tokens
    ax4b = ax4.twinx()
    ax4b.plot(x, [t/1000 for t in avg_tokens], 'cyan', linewidth=1.5, linestyle='--', label='Avg Tokens (k)')
    ax4b.set_ylabel('Avg Tokens (thousands)', color='cyan')
    ax4b.tick_params(axis='y', labelcolor='cyan')
    
    # Combined legend
    lines1, labels1 = ax4.get_legend_handles_labels()
    lines2, labels2 = ax4b.get_legend_handles_labels()
    ax4.legend(lines1 + lines2, labels1 + labels2, loc='lower left')
    
    plt.tight_layout()
    
    output_path = os.path.join(output_dir, 'trajectory_analysis.png')
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    print(f"\nGraph saved to: {output_path}")
    
    # Also save a detailed CSV
    csv_path = os.path.join(output_dir, 'trajectory_analysis.csv')
    with open(csv_path, 'w') as f:
        headers = ['step', 'total_files', 'avg_reward', 'min_reward', 'max_reward', 
                   'avg_tokens', 'avg_steps', 'tool_call_pct', 'truncated_think_pct',
                   'endoftext_pct', 'empty_parsed_pct', 'very_short_pct']
        f.write(','.join(headers) + '\n')
        for d in steps_data:
            row = [str(d[h]) for h in headers]
            f.write(','.join(row) + '\n')
    print(f"CSV data saved to: {csv_path}")


def main():
    args = parse_args()
    
    print("=" * 100)
    print("TRAJECTORY ANALYSIS REPORT")
    print("=" * 100)
    print(f"\nAnalyzing trajectories from: {args.trajectories_dir}")
    print()
    
    # Collect data
    steps_data = collect_trajectory_data(args.trajectories_dir)
    
    if not steps_data:
        print("ERROR: No trajectory data found!")
        return
    
    # Determine halfway point
    steps = [d['step'] for d in steps_data]
    halfway_step = steps[len(steps)//2]
    
    # Print statistics
    print_statistics_table(steps_data)
    print_issues_table(steps_data)
    
    # Compute and print aggregate comparison
    before, after = compute_aggregate_comparison(steps_data, halfway_step)
    comparison_stats = print_aggregate_comparison(before, after, halfway_step)
    
    # Print key findings
    print_key_findings(comparison_stats)
    
    # Generate graphs
    if not args.no_graph:
        generate_graphs(steps_data, args.output_dir)


if __name__ == '__main__':
    main()