"""
Armory Sweeps - Hyperparameter optimization

This module provides hyperparameter sweep functionality similar to W&B Sweeps.
"""

import os
import random
import itertools
from typing import Any, Dict, List, Optional, Callable, Union
from dataclasses import dataclass
from enum import Enum

from armory.logger import get_logger

logger = get_logger("sweep")


class SweepMethod(Enum):
    """Sweep search methods"""
    GRID = "grid"
    RANDOM = "random"
    BAYES = "bayes"


class ParameterDistribution(Enum):
    """Parameter distribution types"""
    UNIFORM = "uniform"
    LOG_UNIFORM = "log_uniform"
    NORMAL = "normal"
    LOG_NORMAL = "log_normal"
    CATEGORICAL = "categorical"
    INT_UNIFORM = "int_uniform"
    CONSTANT = "constant"


@dataclass
class SweepParameter:
    """Definition of a sweep parameter"""
    name: str
    distribution: ParameterDistribution
    values: Optional[List[Any]] = None  # For categorical
    min: Optional[float] = None
    max: Optional[float] = None
    mean: Optional[float] = None
    std: Optional[float] = None
    value: Optional[Any] = None  # For constant

    def sample(self) -> Any:
        """Sample a value from this parameter's distribution"""
        import math

        if self.distribution == ParameterDistribution.CATEGORICAL:
            return random.choice(self.values)
        elif self.distribution == ParameterDistribution.CONSTANT:
            return self.value
        elif self.distribution == ParameterDistribution.UNIFORM:
            return random.uniform(self.min, self.max)
        elif self.distribution == ParameterDistribution.LOG_UNIFORM:
            log_min = math.log(self.min)
            log_max = math.log(self.max)
            return math.exp(random.uniform(log_min, log_max))
        elif self.distribution == ParameterDistribution.INT_UNIFORM:
            return random.randint(int(self.min), int(self.max))
        elif self.distribution == ParameterDistribution.NORMAL:
            return random.gauss(self.mean or 0, self.std or 1)
        elif self.distribution == ParameterDistribution.LOG_NORMAL:
            return math.exp(random.gauss(self.mean or 0, self.std or 1))
        else:
            raise ValueError(f"Unknown distribution: {self.distribution}")


class SweepConfig:
    """
    Sweep configuration.

    Example:
        sweep_config = {
            "method": "random",
            "metric": {"name": "val/loss", "goal": "minimize"},
            "parameters": {
                "learning_rate": {"min": 1e-5, "max": 1e-2, "distribution": "log_uniform"},
                "batch_size": {"values": [16, 32, 64, 128]},
                "epochs": {"value": 10},
            }
        }
    """

    def __init__(self, config: Dict[str, Any]):
        self.method = SweepMethod(config.get("method", "random"))
        self.metric = config.get("metric", {})
        self.metric_name = self.metric.get("name", "loss")
        self.metric_goal = self.metric.get("goal", "minimize")
        self.parameters: Dict[str, SweepParameter] = {}
        self.run_cap = config.get("run_cap")
        self.early_terminate = config.get("early_terminate")

        # Parse parameters
        for name, param_config in config.get("parameters", {}).items():
            self.parameters[name] = self._parse_parameter(name, param_config)

    def _parse_parameter(self, name: str, config: Dict[str, Any]) -> SweepParameter:
        """Parse a parameter configuration"""
        if "values" in config:
            return SweepParameter(
                name=name,
                distribution=ParameterDistribution.CATEGORICAL,
                values=config["values"],
            )
        elif "value" in config:
            return SweepParameter(
                name=name,
                distribution=ParameterDistribution.CONSTANT,
                value=config["value"],
            )
        else:
            dist_str = config.get("distribution", "uniform")
            dist = ParameterDistribution(dist_str)
            return SweepParameter(
                name=name,
                distribution=dist,
                min=config.get("min"),
                max=config.get("max"),
                mean=config.get("mean"),
                std=config.get("std"),
            )

    def sample_config(self) -> Dict[str, Any]:
        """Sample a configuration for random/bayes sweeps"""
        return {name: param.sample() for name, param in self.parameters.items()}

    def grid_configs(self) -> List[Dict[str, Any]]:
        """Generate all configurations for grid sweep"""
        # Get all values for each parameter
        param_values = {}
        for name, param in self.parameters.items():
            if param.distribution == ParameterDistribution.CATEGORICAL:
                param_values[name] = param.values
            elif param.distribution == ParameterDistribution.CONSTANT:
                param_values[name] = [param.value]
            else:
                # For continuous distributions, sample some values
                param_values[name] = [param.sample() for _ in range(5)]

        # Generate all combinations
        keys = list(param_values.keys())
        values = [param_values[k] for k in keys]
        configs = []
        for combo in itertools.product(*values):
            configs.append(dict(zip(keys, combo)))
        return configs


class Sweep:
    """
    Armory Sweep controller.

    Example:
        import armory
        from armory import Sweep

        sweep_config = {
            "method": "random",
            "metric": {"name": "val/loss", "goal": "minimize"},
            "parameters": {
                "learning_rate": {"min": 1e-5, "max": 1e-2, "distribution": "log_uniform"},
                "batch_size": {"values": [16, 32, 64]},
            }
        }

        sweep = Sweep(sweep_config, project="my-project")

        def train():
            run = armory.init()
            lr = run.config.learning_rate
            batch_size = run.config.batch_size
            # ... training code ...
            run.log({"val/loss": final_loss})
            run.finish()

        sweep.run(train, count=10)
    """

    def __init__(
        self,
        config: Union[Dict[str, Any], SweepConfig],
        project: str = "default",
        entity: Optional[str] = None,
    ):
        if isinstance(config, dict):
            self.config = SweepConfig(config)
        else:
            self.config = config
        self.project = project
        self.entity = entity
        self._runs: List[Dict[str, Any]] = []
        self._best_run: Optional[Dict[str, Any]] = None
        self._sweep_id = f"sweep-{random.randint(1000, 9999)}"

    @property
    def sweep_id(self) -> str:
        return self._sweep_id

    @property
    def best_run(self) -> Optional[Dict[str, Any]]:
        return self._best_run

    def _get_next_config(self) -> Dict[str, Any]:
        """Get the next configuration to try"""
        if self.config.method == SweepMethod.GRID:
            configs = self.config.grid_configs()
            run_index = len(self._runs) % len(configs)
            return configs[run_index]
        elif self.config.method == SweepMethod.RANDOM:
            return self.config.sample_config()
        elif self.config.method == SweepMethod.BAYES:
            # Simple Bayesian optimization: explore early, exploit later
            if len(self._runs) < 5:
                return self.config.sample_config()
            else:
                # Use best config with small perturbations
                if self._best_run:
                    best_config = self._best_run.get("config", {})
                    new_config = {}
                    for name, param in self.config.parameters.items():
                        if param.distribution == ParameterDistribution.CATEGORICAL:
                            # 70% chance to keep best, 30% to explore
                            if random.random() < 0.7:
                                new_config[name] = best_config.get(name, param.sample())
                            else:
                                new_config[name] = param.sample()
                        elif param.distribution == ParameterDistribution.CONSTANT:
                            new_config[name] = param.value
                        else:
                            # Perturb continuous parameters
                            best_val = best_config.get(name, param.sample())
                            perturbation = random.gauss(0, 0.1) * best_val
                            new_val = best_val + perturbation
                            if param.min is not None:
                                new_val = max(param.min, new_val)
                            if param.max is not None:
                                new_val = min(param.max, new_val)
                            new_config[name] = new_val
                    return new_config
                else:
                    return self.config.sample_config()
        else:
            return self.config.sample_config()

    def _update_best(self, run_result: Dict[str, Any]):
        """Update best run based on metric"""
        metric_value = run_result.get("metric")
        if metric_value is None:
            return

        if self._best_run is None:
            self._best_run = run_result
        else:
            best_value = self._best_run.get("metric")
            if self.config.metric_goal == "minimize":
                if metric_value < best_value:
                    self._best_run = run_result
            else:
                if metric_value > best_value:
                    self._best_run = run_result

    def run(
        self,
        function: Callable,
        count: int = 10,
    ):
        """
        Run the sweep.

        Args:
            function: Training function to run for each config
            count: Number of runs to perform
        """
        # Import run module directly to avoid circular imports
        from armory.run import init as armory_init, Run

        logger.info(f"Starting sweep {self._sweep_id}")
        logger.info(f"Method: {self.config.method.value}")
        logger.info(f"Metric: {self.config.metric_name} ({self.config.metric_goal})")
        logger.info(f"Parameters: {list(self.config.parameters.keys())}")

        for i in range(count):
            config = self._get_next_config()
            logger.info(f"Run {i+1}/{count}: {config}")

            # Initialize run with sweep config
            run = armory_init(
                project=self.project,
                name=f"{self._sweep_id}-run-{i+1}",
                config=config,
                tags=["sweep", self._sweep_id],
            )

            try:
                # Run the training function
                function()

                # Get the metric from summary
                metric_value = run.summary.get(self.config.metric_name)

                run_result = {
                    "run_id": run.id,
                    "config": config,
                    "metric": metric_value,
                    "index": i,
                }
                self._runs.append(run_result)
                self._update_best(run_result)

                logger.info(f"Run {i+1} complete: {self.config.metric_name}={metric_value}")

            except Exception as e:
                logger.error(f"Run {i+1} failed: {e}")
                run.finish(exit_code=1)
            finally:
                if not run._finished:
                    run.finish()

        logger.info("Sweep complete!")
        if self._best_run:
            logger.info(f"Best run: {self._best_run}")

    def agent(self, function: Callable):
        """
        Create a sweep agent (for distributed sweeps).

        Args:
            function: Training function
        """
        # For now, just run once with a sampled config
        self.run(function, count=1)


def sweep(
    sweep_config: Dict[str, Any],
    function: Optional[Callable] = None,
    project: str = "default",
    count: int = 10,
) -> Union[str, None]:
    """
    Initialize and optionally run a sweep.

    Args:
        sweep_config: Sweep configuration dictionary
        function: Training function (if provided, runs immediately)
        project: Project name
        count: Number of runs

    Returns:
        Sweep ID if no function provided, None otherwise

    Example:
        # Create sweep and run later
        sweep_id = armory.sweep(config, project="my-project")

        # Or create and run immediately
        armory.sweep(config, function=train, count=10)
    """
    s = Sweep(sweep_config, project=project)

    if function is not None:
        s.run(function, count=count)
        return None
    else:
        return s.sweep_id


def agent(
    sweep_id: str,
    function: Callable,
    count: int = 1,
):
    """
    Run a sweep agent.

    Args:
        sweep_id: Sweep ID to join
        function: Training function
        count: Number of runs for this agent
    """
    # For now, create a new sweep with default config
    # In a real implementation, this would fetch config from server
    logger.info(f"Agent joining sweep {sweep_id}")
    logger.info(f"Running {count} trial(s)")

    for i in range(count):
        function()
