#!/usr/bin/env python3
"""
Extract DQN Tuning Results from Checkpoints
Since Optuna study wasn't preserved, we need to backtest checkpoints to determine best hyperparameters.
"""

import os
import json
from pathlib import Path
from datetime import datetime

def analyze_checkpoints():
    """Analyze checkpoint directory structure"""
    base_dir = Path("ml/tuning_checkpoints")
    
    if not base_dir.exists():
        print(f"Error: {base_dir} does not exist")
        return
    
    trials = []
    for trial_dir in sorted(base_dir.iterdir()):
        if not trial_dir.is_dir() or not trial_dir.name.startswith("trial_"):
            continue
        
        trial_num = trial_dir.name.replace("trial_", "")
        checkpoints = list(trial_dir.glob("*.safetensors"))
        
        if not checkpoints:
            print(f"⚠️  {trial_dir.name}: No checkpoints found")
            continue
        
        # Get file metadata
        checkpoint = checkpoints[-1]  # Use final checkpoint
        stats = checkpoint.stat()
        
        trial_info = {
            "trial_num": int(trial_num) if trial_num.isdigit() else trial_num,
            "num_checkpoints": len(checkpoints),
            "final_checkpoint": checkpoint.name,
            "file_size": stats.st_size,
            "created": datetime.fromtimestamp(stats.st_ctime).isoformat(),
            "modified": datetime.fromtimestamp(stats.st_mtime).isoformat()
        }
        
        trials.append(trial_info)
        print(f"✓ Trial {trial_num:2s}: {len(checkpoints)} checkpoints, {stats.st_size/1024:.1f} KB")
    
    return trials

def create_backtest_script(trials):
    """Create a script to backtest each checkpoint"""
    script_path = Path("backtest_dqn_trials.sh")
    
    with open(script_path, 'w') as f:
        f.write("#!/bin/bash\n")
        f.write("# Backtest all DQN trial checkpoints to determine best hyperparameters\n\n")
        f.write("set -e\n\n")
        f.write('RESULTS_FILE="dqn_backtest_results.json"\n')
        f.write('echo "[" > $RESULTS_FILE\n\n')
        
        for i, trial in enumerate(trials):
            if isinstance(trial["trial_num"], int):
                trial_num = trial["trial_num"]
                checkpoint_path = f"ml/tuning_checkpoints/trial_{trial_num}/{trial['final_checkpoint']}"
                
                f.write(f"echo 'Backtesting trial {trial_num}...'\n")
                f.write(f"# TODO: Add actual backtest command here\n")
                f.write(f"# cargo run --example backtest_dqn -- --checkpoint {checkpoint_path}\n\n")
        
        f.write('echo "]" >> $RESULTS_FILE\n')
        f.write('echo "Results saved to $RESULTS_FILE"\n')
    
    os.chmod(script_path, 0o755)
    print(f"\n✅ Created backtest script: {script_path}")

def create_search_space_reference():
    """Create a reference document for the hyperparameter search space"""
    content = """# DQN Hyperparameter Search Space (36 Trials)

Based on tuning_config.yaml:

## Search Space

### learning_rate
- Type: loguniform
- Range: [0.0001, 0.01]
- Distribution: Logarithmic between 1e-4 and 1e-2

### batch_size
- Type: categorical
- Choices: [64, 128, 256]

### gamma (discount factor)
- Type: uniform
- Range: [0.95, 0.99]

## Objective
- Metric: sharpe_ratio
- Direction: maximize

## Pruning Strategy
- Enabled: true
- Strategy: median
- Warmup trials: 2

## Trial Summary

Total Trials: 36 completed (out of 50 requested)
Stopped early: User interrupted or median pruning

## Next Steps

1. **Option A: Backtest All Checkpoints** (Recommended)
   - Test each of the 36 checkpoint files with real market data
   - Measure Sharpe ratio for each trial
   - Extract hyperparameters from top 5 performers
   - Estimated time: 3-4 hours

2. **Option B: Use Default Best-Practice Hyperparameters**
   - learning_rate: 0.001 (middle of loguniform range)
   - batch_size: 128 (balanced memory/performance)
   - gamma: 0.97 (standard DQN discount factor)
   - Trade-off: Faster but suboptimal

3. **Option C: Resume Tuning**
   - Continue from trial 36 to complete 50 trials
   - Requires original tuning job ID and Optuna study
   - Estimated time: 2-3 hours additional

## Recommendation

**Use Option A** if PPO tuning is blocked on DQN results.
**Use Option B** if immediate PPO tuning is priority and can iterate later.

"""
    
    path = Path("DQN_TUNING_SEARCH_SPACE.md")
    with open(path, 'w') as f:
        f.write(content)
    print(f"✅ Created search space reference: {path}")

def main():
    print("=" * 70)
    print("DQN TUNING CHECKPOINT ANALYSIS")
    print("=" * 70)
    print()
    
    trials = analyze_checkpoints()
    
    if trials:
        print(f"\n📊 Summary:")
        print(f"   Total trials with checkpoints: {len(trials)}")
        
        # Calculate statistics
        avg_size = sum(t["file_size"] for t in trials) / len(trials)
        print(f"   Average checkpoint size: {avg_size/1024:.1f} KB")
        
        # Save trial info
        with open("dqn_trial_metadata.json", 'w') as f:
            json.dump(trials, f, indent=2)
        print(f"\n✅ Saved trial metadata to: dqn_trial_metadata.json")
        
        # Create helper scripts
        create_backtest_script(trials)
        create_search_space_reference()
    else:
        print("\n❌ No valid trials found")
    
    print("\n" + "=" * 70)
    print("Next Steps:")
    print("1. Review DQN_TUNING_SEARCH_SPACE.md for options")
    print("2. Choose backtesting strategy (A, B, or C)")
    print("3. For Option A: Implement backtest logic in backtest_dqn_trials.sh")
    print("=" * 70)

if __name__ == "__main__":
    main()