"""
Optimus Research Center - Laboratory 1: Algorithm Benchmarking Engine
Coordinates rigorous, reproducible execution and comparison of algorithms,
measuring wall time, CPU time, peak memory, and accuracy against baselines.
"""

import time
import psutil
import numpy as np
from typing import Dict, Any, List, Optional, Union

from app.lab_algorithms.algorithms.optimization import (
    GeneticAlgorithmOptimizer,
    BayesianSurrogateOptimizer,
    SimulatedAnnealingOptimizer,
    DifferentialEvolutionOptimizer
)
from app.lab_algorithms.algorithms.predictive import (
    RidgePredictor,
    GradientBoostingPredictor,
    KNNPredictor
)
from app.lab_algorithms.algorithms.search_planning import (
    AStarSearch,
    DijkstraSearch,
    CriticalPathPlanner
)
from app.lab_algorithms.algorithms.anomaly import (
    IsolationForestAnomalyDetector,
    RollingZScoreAnomalyDetector
)


class AlgorithmBenchmarkingEngine:
    """Executes algorithms under uniform telemetry and produces comparative metrics."""

    @classmethod
    def benchmark_optimization(
        cls,
        algorithms: List[Union[str, Dict[str, Any]]],
        problem_type: str = "rastrigin",
        dim: int = 5,
        random_seed: int = 42,
        parameter_configs: Optional[Dict[str, Dict[str, Any]]] = None,
        parameter_grid: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """
        Benchmarks optimization algorithms on standard mathematical benchmark functions.
        (e.g., Sphere: f(x)=sum(x^2), Rastrigin: multimodal challenging function).
        Applies physically specified candidate parameters to authentic executors.
        """
        def objective(x: np.ndarray) -> float:
            # Maximization objective (negated error so global optimum is 0.0)
            p_norm = problem_type.lower()
            if p_norm == "sphere":
                return -float(np.sum(x ** 2))
            elif p_norm == "rosenbrock":
                if len(x) < 2:
                    return -float((1.0 - x[0]) ** 2)
                return -float(np.sum(100.0 * (x[1:] - x[:-1] ** 2.0) ** 2.0 + (1.0 - x[:-1]) ** 2.0))
            elif p_norm == "ackley":
                d = max(1, len(x))
                term1 = -20.0 * np.exp(-0.2 * np.sqrt(np.sum(x ** 2) / d))
                term2 = -np.exp(np.sum(np.cos(2.0 * np.pi * x)) / d)
                return -float(term1 + term2 + 20.0 + np.e)
            # Rastrigin function
            return -float(10 * len(x) + np.sum(x**2 - 10 * np.cos(2 * np.pi * x)))

        results = []
        proc = psutil.Process()
        baseline_fitness = None

        # Expand parameter_grid into algorithms if provided
        expanded_algorithms = []
        if parameter_grid:
            for item in algorithms:
                alg_name = item.get("algorithm", str(item)) if isinstance(item, dict) else str(item)
                if alg_name in parameter_grid and isinstance(parameter_grid[alg_name], list):
                    for i, p in enumerate(parameter_grid[alg_name]):
                        expanded_algorithms.append({
                            "algorithm": alg_name,
                            "candidate_name": f"{alg_name}_cfg_{i+1}",
                            "params": p
                        })
                elif alg_name in parameter_grid and isinstance(parameter_grid[alg_name], dict):
                    expanded_algorithms.append({
                        "algorithm": alg_name,
                        "candidate_name": f"{alg_name}_cfg_1",
                        "params": parameter_grid[alg_name]
                    })
                else:
                    expanded_algorithms.append(item)
            algorithms = expanded_algorithms

        for idx, item in enumerate(algorithms):
            if isinstance(item, dict):
                alg_name = item.get("algorithm", "UnknownAlgorithm")
                candidate_name = item.get("candidate_name", alg_name)
                cand_params = item.get("params", item.get("parameters", {}))
            else:
                alg_name = str(item)
                candidate_name = alg_name
                cand_params = (parameter_configs or {}).get(alg_name, {})

            mem_before = proc.memory_info().rss / (1024 * 1024)
            start_wall = time.perf_counter()
            start_cpu = time.process_time()

            output = {}
            applied_params = {}
            clean_name = alg_name.lower().replace(" ", "_").replace("-", "")

            if clean_name in {"genetic_algorithm", "geneticalgorithm", "ga"}:
                ga_kwargs = {
                    "population_size": int(cand_params.get("population_size", 30)),
                    "generations": int(cand_params.get("generations", 25)),
                    "crossover_rate": float(cand_params.get("crossover_rate", 0.8)),
                    "mutation_rate": float(cand_params.get("mutation_rate", 0.15)),
                    "random_seed": random_seed
                }
                ga = GeneticAlgorithmOptimizer(**ga_kwargs)
                output = ga.optimize(objective, dim=dim)
                applied_params = ga_kwargs

            elif clean_name in {"bayesian_surrogate", "bayesian", "bayesiansurrogate", "bayesianoptimizer", "bo"}:
                bo_kwargs = {
                    "n_init_samples": int(cand_params.get("n_init_samples", 5)),
                    "n_iterations": int(cand_params.get("n_iterations", 15)),
                    "random_seed": random_seed
                }
                bo = BayesianSurrogateOptimizer(**bo_kwargs)
                output = bo.optimize(objective, dim=dim)
                applied_params = bo_kwargs

            elif clean_name in {"simulated_annealing", "simulatedannealing", "sa"}:
                sa_kwargs = {
                    "initial_temp": float(cand_params.get("initial_temp", 50.0)),
                    "cooling_rate": float(cand_params.get("cooling_rate", 0.92)),
                    "iterations_per_temp": int(cand_params.get("iterations_per_temp", 20)),
                    "random_seed": random_seed
                }
                sa = SimulatedAnnealingOptimizer(**sa_kwargs)
                output = sa.optimize(objective, dim=dim)
                applied_params = sa_kwargs

            elif clean_name in {"differential_evolution", "differentialevolution", "de"}:
                de_kwargs = {
                    "population_size": int(cand_params.get("population_size", 30)),
                    "generations": int(cand_params.get("generations", 25)),
                    "mutation_factor": float(cand_params.get("mutation_factor", 0.7)),
                    "crossover_rate": float(cand_params.get("crossover_rate", 0.8)),
                    "random_seed": random_seed
                }
                de = DifferentialEvolutionOptimizer(**de_kwargs)
                output = de.optimize(objective, dim=dim)
                applied_params = de_kwargs

            elif clean_name in {"hill_climbing", "hillclimbing", "random_search"}:
                raise ValueError(
                    f"UNSUPPORTED_CANDIDATE: Algorithm '{alg_name}' has no authentic executor implemented. "
                    f"Masquerading under SimulatedAnnealing is strictly prohibited."
                )
            else:
                raise ValueError(f"Unknown optimization algorithm: {alg_name}")

            wall_time = time.perf_counter() - start_wall
            cpu_time = time.process_time() - start_cpu
            mem_after = proc.memory_info().rss / (1024 * 1024)
            peak_mem = max(0.01, round(mem_after - mem_before, 2))

            best_fit = round(float(output.get("best_fitness", 0.0)), 6)

            if idx == 0:
                outcome = "BASELINE"
                baseline_fitness = best_fit
            else:
                outcome = "POSITIVE" if best_fit > baseline_fitness else "NEGATIVE"

            results.append({
                "algorithm": alg_name,
                "candidate_name": candidate_name,
                "parameters": applied_params,
                "wall_time_seconds": round(wall_time, 4),
                "cpu_time_seconds": round(cpu_time, 4),
                "peak_memory_mb": peak_mem,
                "metrics": {
                    "best_fitness": round(best_fit, 6),
                    "problem_type": problem_type,
                    "dimension": dim
                },
                "outcome": outcome,
                "random_seed": random_seed,
                "details": output
            })

        # Rank candidates by fitness descending
        ranked = sorted(results, key=lambda r: r["metrics"]["best_fitness"], reverse=True)

        baseline_name = algorithms[0]["candidate_name"] if isinstance(algorithms[0], dict) else str(algorithms[0])
        return {
            "experiment_type": "ALGORITHM_BENCHMARK_OPTIMIZATION",
            "problem": problem_type,
            "dimension": dim,
            "runs": results,
            "top_candidate": ranked[0]["candidate_name"],
            "top_algorithm": ranked[0]["algorithm"],
            "top_parameters": ranked[0]["parameters"],
            "top_fitness": ranked[0]["metrics"]["best_fitness"],
            "baseline_algorithm": baseline_name,
            "baseline_fitness": baseline_fitness
        }

    @classmethod
    def benchmark_predictive(
        cls,
        algorithms: List[str],
        n_samples: int = 200,
        n_features: int = 10,
        random_seed: int = 42
    ) -> Dict[str, Any]:
        """Benchmarks regression/predictive algorithms on standardized synthetic data."""
        rng = np.random.RandomState(random_seed)
        X = rng.randn(n_samples, n_features)
        true_w = rng.randn(n_features)
        y = np.dot(X, true_w) + rng.normal(0, 0.2, size=n_samples)

        split = int(n_samples * 0.75)
        X_train, X_test = X[:split], X[split:]
        y_train, y_test = y[:split], y[split:]

        proc = psutil.Process()
        results = []
        baseline_mse = None

        for idx, alg_name in enumerate(algorithms):
            mem_before = proc.memory_info().rss / (1024 * 1024)
            start_wall = time.perf_counter()
            start_cpu = time.process_time()

            if alg_name == "ridge_regression":
                pred = RidgePredictor(alpha=1.0, random_seed=random_seed)
                eval_metrics = pred.train_and_eval(X_train, y_train, X_test, y_test)
            elif alg_name == "gradient_boosting":
                pred = GradientBoostingPredictor(n_estimators=30, learning_rate=0.1, random_seed=random_seed)
                eval_metrics = pred.train_and_eval(X_train, y_train, X_test, y_test)
            elif alg_name == "knn_regressor":
                pred = KNNPredictor(n_neighbors=5)
                eval_metrics = pred.train_and_eval(X_train, y_train, X_test, y_test)
            else:
                raise ValueError(f"Unknown predictive algorithm: {alg_name}")

            wall_time = time.perf_counter() - start_wall
            cpu_time = time.process_time() - start_cpu
            mem_after = proc.memory_info().rss / (1024 * 1024)
            peak_mem = max(0.01, round(mem_after - mem_before, 2))

            mse = eval_metrics["mse"]
            if idx == 0:
                outcome = "BASELINE"
                baseline_mse = mse
            else:
                outcome = "POSITIVE" if mse < baseline_mse else "NEGATIVE"

            results.append({
                "algorithm": alg_name,
                "wall_time_seconds": round(wall_time, 4),
                "cpu_time_seconds": round(cpu_time, 4),
                "peak_memory_mb": peak_mem,
                "metrics": eval_metrics,
                "outcome": outcome,
                "random_seed": random_seed
            })

        # Rank candidates by MSE ascending
        ranked = sorted(results, key=lambda r: r["metrics"]["mse"])

        return {
            "experiment_type": "ALGORITHM_BENCHMARK_PREDICTIVE",
            "samples": n_samples,
            "features": n_features,
            "runs": results,
            "top_candidate": ranked[0]["algorithm"],
            "top_mse": ranked[0]["metrics"]["mse"],
            "baseline_algorithm": algorithms[0]
        }
