"""
Hyperparameter tuning script for gradient ascent optimization.

This script performs a systematic search over hyperparameter combinations
to find the optimal configuration for maximum evaluation scores.
"""

import subprocess
import json
import argparse
from pathlib import Path
from datetime import datetime
import itertools
import numpy as np
from typing import Dict, List, Any
import re


class HyperparameterTuner:
    """Hyperparameter tuner for gradient ascent."""
    
    def __init__(
        self,
        output_dir: str = "tuning_results",
        max_samples: int = 30,
        num_steps: int = 20,
        dataset_type: str = "pickapic",
        model_variant: str = "lpo",
        cuda_id: int = 0,
        metrics: List[str] = None
    ):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        self.max_samples = max_samples
        self.num_steps = num_steps
        self.dataset_type = dataset_type
        self.model_variant = model_variant
        self.cuda_id = cuda_id
        self.metrics = metrics or ["clip", "aesthetic", "pickscore", "hpsv2", "imagereward"]
        
        # Store results
        self.results = []
        self.baseline_results = None
        
    def define_search_space(self) -> List[Dict[str, Any]]:
        """Define the hyperparameter search space - FULL GRID SEARCH.
        
        Tests all combinations of parameters including momentum overrides for configs that support it.
        """
        
        # Define all parameter values
        cfg_scales = [3.0, 5.0, 7.5] #  
        
        # All available gradient configs from grad_ascent_configs.py
        grad_configs = [
            # "constant",
            # "linear",
            "cosine_nesterov",
            # "low_to_high_nesterov",
            # "high_to_low_nesterov",
            "low_to_high_momentum",
            "high_to_low_momentum",
        ]
        
        num_grad_steps_list = [1, 2] # 5, 7, 10
        grad_step_sizes = [0.001, 0.005, 0.01, 0.05] # 
        momentums = [0.5, 0.8, 0.9]  # 
        
        # Generate ALL combinations using itertools.product
        configs = []
        for cfg, grad_cfg, num_steps, step_size, momentum in itertools.product(
            cfg_scales, grad_configs, num_grad_steps_list, grad_step_sizes, momentums
        ):
            configs.append({
                "cfg_scale": cfg,
                "grad_config": grad_cfg,
                "num_grad_steps": num_steps,
                "grad_step_size": step_size,
                "momentum": momentum,
            })
        
        print(f"\nGenerated {len(configs)} total configurations")
        print(f"  cfg_scales: {len(cfg_scales)}")
        print(f"  grad_configs: {len(grad_configs)}")
        print(f"  num_grad_steps: {len(num_grad_steps_list)}")
        print(f"  grad_step_sizes: {len(grad_step_sizes)}")
        print(f"  momentums: {len(momentums)}")
        print(f"  Total: {len(cfg_scales)} × {len(grad_configs)} × {len(num_grad_steps_list)} × {len(grad_step_sizes)} × {len(momentums)} = {len(configs)}")
        
        return configs
    
    def run_baseline(self) -> Dict[str, float]:
        """Run baseline evaluation once."""
        print("\n" + "="*80)
        print("RUNNING BASELINE EVALUATION")
        print("="*80)
        
        # Use median cfg_scale for baseline
        cfg_scale = 5.0
        
        output_dir = self.output_dir / "baseline"
        
        cmd = [
            "python", "eval.py",
            "--model_variant", self.model_variant,
            "--dataset_type", self.dataset_type,
            "--max_samples", str(self.max_samples),
            "--num_steps", str(self.num_steps),
            "--cfg_scale", str(cfg_scale),
            "--output_dir", str(output_dir),
            "--cuda", str(self.cuda_id),
            "--mode", "baseline",
            "--metrics", *self.metrics,
        ]
        
        print(f"Command: {' '.join(cmd)}")
        
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
            
            # Parse results from output
            metrics = self._parse_metrics(result.stdout, "baseline")
            
            print(f"\nBaseline Results:")
            for metric, value in metrics.items():
                print(f"  {metric}: {value:.4f}")
            
            self.baseline_results = {
                "cfg_scale": cfg_scale,
                "metrics": metrics,
            }
            
            return metrics
            
        except subprocess.CalledProcessError as e:
            print(f"Error running baseline: {e}")
            print(f"Stdout: {e.stdout}")
            print(f"Stderr: {e.stderr}")
            return {}
    
    def run_experiment(self, config: Dict[str, Any]) -> Dict[str, Any]:
        """Run a single experiment with given hyperparameters."""
        
        # Create output directory for this config
        config_name = f"cfg{config['cfg_scale']}_" \
                     f"{config['grad_config']}_" \
                     f"steps{config['num_grad_steps']}_" \
                     f"lr{config['grad_step_size']}_" \
                     f"mom{config['momentum']}"
        
        output_dir = self.output_dir / config_name
        
        # Build command
        cmd = [
            "python", "eval.py",
            "--model_variant", self.model_variant,
            "--dataset_type", self.dataset_type,
            "--grad_config", config["grad_config"],
            "--max_samples", str(self.max_samples),
            "--num_steps", str(self.num_steps),
            "--cfg_scale", str(config["cfg_scale"]),
            "--output_dir", str(output_dir),
            "--cuda", str(self.cuda_id),
            "--mode", "gradient_ascent",
            "--metrics", *self.metrics,
            # Override config parameters
            "--override_num_grad_steps", str(config["num_grad_steps"]),
            "--override_grad_step_size", str(config["grad_step_size"]),
            "--override_momentum", str(config["momentum"]),
        ]
        
        print(f"\nRunning experiment: {config_name}")
        print(f"Config: {config}")
        
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
            
            # Parse metrics from output
            metrics = self._parse_metrics(result.stdout, "gradient_ascent")
            
            # Compute improvement over baseline
            improvements = {}
            if self.baseline_results:
                baseline_metrics = self.baseline_results["metrics"]
                for metric, value in metrics.items():
                    if metric in baseline_metrics:
                        baseline_val = baseline_metrics[metric]
                        if baseline_val != 0:
                            improvement = ((value - baseline_val) / abs(baseline_val)) * 100
                            improvements[f"{metric}_improvement"] = improvement
            
            result_dict = {
                "config": config,
                "metrics": metrics,
                "improvements": improvements,
                "output_dir": str(output_dir),
                "timestamp": datetime.now().isoformat(),
            }
            
            print(f"Results:")
            for metric, value in metrics.items():
                print(f"  {metric}: {value:.4f}")
            if improvements:
                print(f"Improvements over baseline:")
                for metric, value in improvements.items():
                    print(f"  {metric}: {value:+.2f}%")
            
            return result_dict
            
        except subprocess.CalledProcessError as e:
            print(f"Error running experiment: {e}")
            print(f"Stderr: {e.stderr}")
            return {
                "config": config,
                "error": str(e),
                "timestamp": datetime.now().isoformat(),
            }
    
    def _parse_metrics(self, output: str, mode: str) -> Dict[str, float]:
        """Parse metrics from eval.py output."""
        metrics = {}
        
        # Look for the summary section
        lines = output.split('\n')
        
        # Pattern to match metric lines like "  Reward: 0.1234"
        metric_patterns = {
            "reward": r"Reward:\s+([-+]?\d*\.?\d+)",
            "clip": r"CLIP Score:\s+([-+]?\d*\.?\d+)",
            "aesthetic": r"Aesthetic Score:\s+([-+]?\d*\.?\d+)",
            "pickscore": r"PickScore:\s+([-+]?\d*\.?\d+)",
            "hpsv2": r"HPSv2 Score:\s+([-+]?\d*\.?\d+)",
            "hpsv21": r"HPSv2\.1 Score:\s+([-+]?\d*\.?\d+)",
            "imagereward": r"ImageReward:\s+([-+]?\d*\.?\d+)",
            "fid": r"FID:\s+([-+]?\d*\.?\d+)",
        }
        
        for line in lines:
            for metric_name, pattern in metric_patterns.items():
                match = re.search(pattern, line)
                if match:
                    metrics[metric_name] = float(match.group(1))
        
        return metrics
    
    def compute_aggregate_score(self, metrics: Dict[str, float]) -> float:
        """
        Compute aggregate score for ranking configurations.
        
        Uses weighted combination of metrics (higher is better for most,
        except FID which is lower is better).
        """
        weights = {
            "reward": 1.0,
            "clip": 0.8,
            "aesthetic": 0.8,
            "pickscore": 1.0,
            "hpsv2": 1.0,
            "hpsv21": 1.0,
            "imagereward": 1.0,
            "fid": -0.5,  # Negative weight (lower FID is better)
        }
        
        score = 0.0
        total_weight = 0.0
        
        for metric, value in metrics.items():
            if metric in weights:
                score += weights[metric] * value
                total_weight += abs(weights[metric])
        
        # Normalize by total weight
        if total_weight > 0:
            score /= total_weight
        
        return score
    
    def run_search(
        self, 
        search_type: str = "grid",
        start_idx: int = 0,
        end_idx: int = None
    ) -> List[Dict[str, Any]]:
        """
        Run hyperparameter search.
        
        Args:
            search_type: Type of search ("grid" or "random")
            start_idx: Starting index for experiments (for GPU distribution)
            end_idx: Ending index for experiments (for GPU distribution)
        """
        all_configs = self.define_search_space()
        
        print("\n" + "="*80)
        print("HYPERPARAMETER SEARCH CONFIGURATION")
        print("="*80)
        print(f"Dataset: {self.dataset_type}")
        print(f"Model: {self.model_variant}")
        print(f"Samples: {self.max_samples}")
        print(f"Inference steps: {self.num_steps}")
        print(f"Metrics: {', '.join(self.metrics)}")
        
        # Select subset of configs if indices provided
        if search_type == "grid":
            configs = all_configs
        elif search_type == "random":
            # Random sample from all configs
            n_samples = min(50, len(all_configs))
            indices = np.random.choice(len(all_configs), n_samples, replace=False)
            configs = [all_configs[i] for i in indices]
        else:
            raise ValueError(f"Unknown search type: {search_type}")
        
        # Apply index slicing for GPU distribution
        if end_idx is None:
            end_idx = len(configs)
        configs = configs[start_idx:end_idx]
        
        print(f"\nTotal configurations: {len(all_configs)}")
        print(f"Assigned to this worker: {len(configs)} (indices {start_idx} to {end_idx})")
        
        # Run baseline first
        if self.baseline_results is None:
            self.run_baseline()
        
        # Run experiments
        print("\n" + "="*80)
        print("RUNNING EXPERIMENTS")
        print("="*80)
        
        for i, config in enumerate(configs, 1):
            print(f"\n{'='*80}")
            print(f"Experiment {i}/{len(configs)}")
            print(f"{'='*80}")
            
            result = self.run_experiment(config)
            self.results.append(result)
            
            # Save intermediate results
            self._save_results()
        
        return self.results
    
    def _generate_grid_configs(self, search_space: Dict[str, List[Any]]) -> List[Dict[str, Any]]:
        """Generate all combinations for grid search."""
        keys = list(search_space.keys())
        values = list(search_space.values())
        
        configs = []
        for combination in itertools.product(*values):
            config = dict(zip(keys, combination))
            configs.append(config)
        
        return configs
    
    def _generate_random_configs(
        self, 
        search_space: Dict[str, List[Any]], 
        n_samples: int = 20
    ) -> List[Dict[str, Any]]:
        """Generate random configurations for random search."""
        configs = []
        
        for _ in range(n_samples):
            config = {}
            for param, values in search_space.items():
                config[param] = np.random.choice(values)
            configs.append(config)
        
        return configs
    
    def _save_results(self):
        """Save results to JSON file."""
        results_file = self.output_dir / "tuning_results.json"
        
        data = {
            "baseline": self.baseline_results,
            "experiments": self.results,
            "timestamp": datetime.now().isoformat(),
            "config": {
                "max_samples": self.max_samples,
                "num_steps": self.num_steps,
                "dataset_type": self.dataset_type,
                "model_variant": self.model_variant,
            }
        }
        
        with open(results_file, 'w') as f:
            json.dump(data, f, indent=2)
        
        print(f"\nResults saved to: {results_file}")
    
    def analyze_results(self) -> Dict[str, Any]:
        """Analyze results and find best configuration."""
        if not self.results:
            print("No results to analyze!")
            return {}
        
        print("\n" + "="*80)
        print("ANALYSIS: FINDING BEST CONFIGURATION")
        print("="*80)
        
        # Filter out failed experiments
        successful_results = [r for r in self.results if "metrics" in r]
        
        if not successful_results:
            print("No successful experiments!")
            return {}
        
        # Compute aggregate scores
        for result in successful_results:
            metrics = result["metrics"]
            result["aggregate_score"] = self.compute_aggregate_score(metrics)
        
        # Sort by aggregate score
        successful_results.sort(key=lambda x: x["aggregate_score"], reverse=True)
        
        # Print top 5 configurations
        print("\nTop 5 Configurations:")
        print("="*80)
        
        for i, result in enumerate(successful_results[:5], 1):
            print(f"\n#{i} - Aggregate Score: {result['aggregate_score']:.4f}")
            print(f"Config: {result['config']}")
            print(f"Metrics:")
            for metric, value in result['metrics'].items():
                print(f"  {metric}: {value:.4f}")
            if result.get('improvements'):
                print(f"Improvements over baseline:")
                for metric, value in result['improvements'].items():
                    print(f"  {metric}: {value:+.2f}%")
        
        # Save best config
        best_result = successful_results[0]
        best_config_file = self.output_dir / "best_config.json"
        
        with open(best_config_file, 'w') as f:
            json.dump({
                "config": best_result["config"],
                "metrics": best_result["metrics"],
                "aggregate_score": best_result["aggregate_score"],
                "improvements": best_result.get("improvements", {}),
            }, f, indent=2)
        
        print(f"\n✓ Best configuration saved to: {best_config_file}")
        
        return best_result


def main():
    parser = argparse.ArgumentParser(description="Hyperparameter tuning for gradient ascent")
    parser.add_argument("--output_dir", type=str, default="tuning_results",
                        help="Directory to save tuning results")
    parser.add_argument("--max_samples", type=int, default=30,
                        help="Number of samples to use for tuning")
    parser.add_argument("--num_steps", type=int, default=20,
                        help="Number of inference steps (fixed)")
    parser.add_argument("--dataset_type", type=str, default="pickapic",
                        choices=["coco", "pickapic"],
                        help="Dataset to use")
    parser.add_argument("--model_variant", type=str, default="lpo",
                        choices=["origin", "spo", "diffusion_dpo", "lpo"],
                        help="Model variant to use")
    parser.add_argument("--cuda", type=int, default=0,
                        help="CUDA device ID")
    parser.add_argument("--search_type", type=str, default="grid",
                        choices=["grid", "random"],
                        help="Type of hyperparameter search")
    parser.add_argument("--metrics", type=str, nargs="+",
                        default=["clip", "aesthetic", "pickscore", "hpsv2", "imagereward"],
                        help="Metrics to evaluate")
    parser.add_argument("--start_idx", type=int, default=0,
                        help="Starting index for experiments (for GPU distribution)")
    parser.add_argument("--end_idx", type=int, default=None,
                        help="Ending index for experiments (for GPU distribution)")
    
    args = parser.parse_args()
    
    # Create tuner
    tuner = HyperparameterTuner(
        output_dir=args.output_dir,
        max_samples=args.max_samples,
        num_steps=args.num_steps,
        dataset_type=args.dataset_type,
        model_variant=args.model_variant,
        cuda_id=args.cuda,
        metrics=args.metrics,
    )
    
    # Run search
    results = tuner.run_search(
        search_type=args.search_type,
        start_idx=args.start_idx,
        end_idx=args.end_idx
    )
    
    # Analyze results
    best_result = tuner.analyze_results()
    
    print("\n" + "="*80)
    print("TUNING COMPLETE!")
    print("="*80)
    print(f"Total experiments: {len(results)}")
    print(f"Results directory: {args.output_dir}")
    
    if best_result:
        print(f"\nBest configuration:")
        print(json.dumps(best_result["config"], indent=2))
        print(f"\nAggregate score: {best_result['aggregate_score']:.4f}")


if __name__ == "__main__":
    main()