"""
View training results and benchmark comparisons.
"""

import json
import os
from datetime import datetime

checkpoint_dir = './checkpoints/benchmark_run'

print("="*70)
print("BENCHMARK TRAINING RESULTS")
print("="*70)

if not os.path.exists(checkpoint_dir):
    print("Training directory not found. Training may still be in progress.")
    exit(1)

# Check for checkpoints
checkpoints = [f for f in os.listdir(checkpoint_dir) if f.endswith('.pt')]
if checkpoints:
    print(f"\n✓ Found {len(checkpoints)} checkpoint(s)")
    for cp in sorted(checkpoints):
        path = os.path.join(checkpoint_dir, cp)
        size = os.path.getsize(path) / (1024*1024)
        mtime = datetime.fromtimestamp(os.path.getmtime(path))
        print(f"  - {cp}: {size:.1f} MB, {mtime.strftime('%Y-%m-%d %H:%M:%S')}")

# Check for evaluation results
eval_files = [f for f in os.listdir(checkpoint_dir) if f.startswith('eval_results')]
if eval_files:
    print(f"\n✓ Found {len(eval_files)} evaluation result(s)")
    for eval_file in sorted(eval_files):
        path = os.path.join(checkpoint_dir, eval_file)
        try:
            with open(path, 'r') as f:
                results = json.load(f)
            
            print(f"\n{eval_file}:")
            for dataset, metrics in results.items():
                if dataset != 'model_info':
                    ppl = metrics.get('perplexity', 'N/A')
                    acc = metrics.get('accuracy', 'N/A')
                    if isinstance(acc, float):
                        acc = acc * 100
                    print(f"  {dataset}:")
                    print(f"    Perplexity: {ppl:.2f if isinstance(ppl, float) else ppl}")
                    print(f"    Accuracy: {acc:.2f if isinstance(acc, float) else acc}%")
                    if metrics.get('top_k_accuracy'):
                        top5 = metrics['top_k_accuracy'] * 100
                        print(f"    Top-5 Accuracy: {top5:.2f}%")
        except Exception as e:
            print(f"  Error reading {eval_file}: {e}")

# Check for final results
final_results = os.path.join(checkpoint_dir, 'final_benchmark_results.json')
if os.path.exists(final_results):
    print("\n" + "="*70)
    print("FINAL BENCHMARK RESULTS")
    print("="*70)
    
    with open(final_results, 'r') as f:
        results = json.load(f)
    
    print("\nPerplexity (lower is better):")
    for dataset, metrics in results.items():
        if dataset != 'model_info':
            ppl = metrics.get('perplexity', 'N/A')
            print(f"  {dataset}: {ppl:.2f if isinstance(ppl, float) else ppl}")
    
    print("\nAccuracy (higher is better):")
    for dataset, metrics in results.items():
        if dataset != 'model_info':
            acc = metrics.get('accuracy', 'N/A')
            if isinstance(acc, float):
                acc = acc * 100
            print(f"  {dataset}: {acc:.2f if isinstance(acc, float) else acc}%")

else:
    print("\nFinal results not yet available. Training may still be in progress.")

print("\n" + "="*70)