"""
AdamOps AutoML Module

Provides automated model selection and hyperparameter tuning.
"""

from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import time
import numpy as np
import pandas as pd
from sklearn.model_selection import cross_val_score

from adamops.utils.logging import get_logger
from adamops.utils.helpers import infer_task_type
from adamops.models.modelops import (
    CLASSIFICATION_MODELS, REGRESSION_MODELS, TrainedModel, train
)

logger = get_logger(__name__)

try:
    import optuna
    optuna.logging.set_verbosity(optuna.logging.WARNING)
    OPTUNA_AVAILABLE = True
except ImportError:
    OPTUNA_AVAILABLE = False


# Default hyperparameter search spaces
PARAM_SPACES = {
    "random_forest": {
        "n_estimators": ("int", 50, 300),
        "max_depth": ("int", 3, 20),
        "min_samples_split": ("int", 2, 20),
        "min_samples_leaf": ("int", 1, 10),
    },
    "gradient_boosting": {
        "n_estimators": ("int", 50, 300),
        "max_depth": ("int", 3, 10),
        "learning_rate": ("float", 0.01, 0.3),
        "min_samples_split": ("int", 2, 20),
    },
    "xgboost": {
        "n_estimators": ("int", 50, 300),
        "max_depth": ("int", 3, 12),
        "learning_rate": ("float", 0.01, 0.3),
        "subsample": ("float", 0.6, 1.0),
        "colsample_bytree": ("float", 0.6, 1.0),
    },
    "lightgbm": {
        "n_estimators": ("int", 50, 300),
        "max_depth": ("int", 3, 12),
        "learning_rate": ("float", 0.01, 0.3),
        "num_leaves": ("int", 20, 100),
        "subsample": ("float", 0.6, 1.0),
    },
    "ridge": {
        "alpha": ("float", 0.001, 100.0, "log"),
    },
    "lasso": {
        "alpha": ("float", 0.001, 100.0, "log"),
    },
    "knn": {
        "n_neighbors": ("int", 1, 30),
        "weights": ("categorical", ["uniform", "distance"]),
    },
}


class AutoMLResult:
    """Results from AutoML run."""
    
    def __init__(self):
        self.best_model: Optional[TrainedModel] = None
        self.best_score: float = 0.0
        self.best_algorithm: str = ""
        self.best_params: Dict = {}
        self.leaderboard: List[Dict] = []
        self.time_elapsed: float = 0.0
    
    def summary(self) -> str:
        lines = [
            "=" * 50, "AutoML Results", "=" * 50,
            f"Best Algorithm: {self.best_algorithm}",
            f"Best Score: {self.best_score:.4f}",
            f"Time Elapsed: {self.time_elapsed:.1f}s",
            "", "Leaderboard:",
        ]
        for i, entry in enumerate(self.leaderboard[:10], 1):
            lines.append(f"  {i}. {entry['algorithm']}: {entry['score']:.4f}")
        return "\n".join(lines)


def grid_search(
    X, y, algorithm: str, param_grid: Dict[str, List],
    task: str = "classification", cv: int = 5, scoring: Optional[str] = None
) -> Tuple[Dict, float]:
    """Grid search hyperparameter tuning."""
    from sklearn.model_selection import GridSearchCV
    
    models = CLASSIFICATION_MODELS if task == "classification" else REGRESSION_MODELS
    if algorithm not in models:
        raise ValueError(f"Unknown algorithm: {algorithm}")
    
    model = models[algorithm]()
    scoring = scoring or ("accuracy" if task == "classification" else "r2")
    
    grid = GridSearchCV(model, param_grid, cv=cv, scoring=scoring, n_jobs=-1)
    grid.fit(X, y)
    
    return grid.best_params_, grid.best_score_


def random_search(
    X, y, algorithm: str, param_distributions: Dict,
    task: str = "classification", cv: int = 5, n_iter: int = 50,
    scoring: Optional[str] = None
) -> Tuple[Dict, float]:
    """Random search hyperparameter tuning."""
    from sklearn.model_selection import RandomizedSearchCV
    
    models = CLASSIFICATION_MODELS if task == "classification" else REGRESSION_MODELS
    model = models[algorithm]()
    scoring = scoring or ("accuracy" if task == "classification" else "r2")
    
    search = RandomizedSearchCV(
        model, param_distributions, n_iter=n_iter, cv=cv,
        scoring=scoring, n_jobs=-1, random_state=42
    )
    search.fit(X, y)
    
    return search.best_params_, search.best_score_


def bayesian_search(
    X, y, algorithm: str, task: str = "classification",
    cv: int = 5, n_trials: int = 50, scoring: Optional[str] = None,
    param_space: Optional[Dict] = None
) -> Tuple[Dict, float]:
    """Bayesian optimization with Optuna."""
    if not OPTUNA_AVAILABLE:
        raise ImportError("Optuna required. Install with: pip install optuna")
    
    models = CLASSIFICATION_MODELS if task == "classification" else REGRESSION_MODELS
    scoring = scoring or ("accuracy" if task == "classification" else "r2")
    space = param_space or PARAM_SPACES.get(algorithm, {})
    
    def objective(trial):
        params = {}
        for name, spec in space.items():
            if spec[0] == "int":
                params[name] = trial.suggest_int(name, spec[1], spec[2])
            elif spec[0] == "float":
                log = len(spec) > 3 and spec[3] == "log"
                params[name] = trial.suggest_float(name, spec[1], spec[2], log=log)
            elif spec[0] == "categorical":
                params[name] = trial.suggest_categorical(name, spec[1])
        
        model = models[algorithm](**params)
        scores = cross_val_score(model, X, y, cv=cv, scoring=scoring)
        return scores.mean()
    
    study = optuna.create_study(direction="maximize")
    study.optimize(objective, n_trials=n_trials, show_progress_bar=False)
    
    return study.best_params, study.best_value


def run(
    X: Union[pd.DataFrame, np.ndarray], y: Union[pd.Series, np.ndarray],
    task: str = "auto", algorithms: Optional[List[str]] = None,
    tuning: str = "bayesian", cv: int = 5, time_limit: int = 3600,
    n_trials: int = 50, scoring: Optional[str] = None
) -> AutoMLResult:
    """
    Run AutoML.
    
    Args:
        X: Features.
        y: Target.
        task: 'classification', 'regression', or 'auto'.
        algorithms: Algorithms to try (None for all).
        tuning: 'grid', 'random', 'bayesian', or 'none'.
        cv: Cross-validation folds.
        time_limit: Max time in seconds.
        n_trials: Trials per algorithm for tuning.
        scoring: Scoring metric.
    
    Returns:
        AutoMLResult: Results with best model and leaderboard.
    """
    start_time = time.time()
    result = AutoMLResult()
    
    # Auto-detect task
    if task == "auto":
        task = infer_task_type(y)
        logger.info(f"Auto-detected task: {task}")
    
    # Get algorithms
    models = CLASSIFICATION_MODELS if task == "classification" else REGRESSION_MODELS
    if algorithms is None:
        algorithms = list(models.keys())
    
    scoring = scoring or ("accuracy" if task == "classification" else "r2")
    logger.info(f"Running AutoML with {len(algorithms)} algorithms")
    
    for algo in algorithms:
        if time.time() - start_time > time_limit:
            logger.warning("Time limit reached")
            break
        
        try:
            logger.info(f"Tuning {algo}...")
            
            if tuning == "bayesian" and algo in PARAM_SPACES and OPTUNA_AVAILABLE:
                best_params, score = bayesian_search(
                    X, y, algo, task, cv, min(n_trials, 30), scoring
                )
            elif tuning == "none":
                model = models[algo]()
                scores = cross_val_score(model, X, y, cv=cv, scoring=scoring)
                best_params, score = {}, scores.mean()
            else:
                # Default to random search
                from scipy.stats import randint, uniform
                param_dist = {}
                if algo in PARAM_SPACES:
                    for name, spec in PARAM_SPACES[algo].items():
                        if spec[0] == "int":
                            param_dist[name] = randint(spec[1], spec[2])
                        elif spec[0] == "float":
                            param_dist[name] = uniform(spec[1], spec[2] - spec[1])
                
                if param_dist:
                    best_params, score = random_search(
                        X, y, algo, param_dist, task, cv, min(n_trials, 20), scoring
                    )
                else:
                    model = models[algo]()
                    scores = cross_val_score(model, X, y, cv=cv, scoring=scoring)
                    best_params, score = {}, scores.mean()
            
            result.leaderboard.append({
                "algorithm": algo, "score": score, "params": best_params
            })
            
            if score > result.best_score:
                result.best_score = score
                result.best_algorithm = algo
                result.best_params = best_params
                
        except Exception as e:
            logger.warning(f"Failed {algo}: {e}")
    
    # Sort leaderboard
    result.leaderboard.sort(key=lambda x: x["score"], reverse=True)
    
    # Train best model
    if result.best_algorithm:
        result.best_model = train(
            X, y, task, result.best_algorithm, result.best_params
        )
    
    result.time_elapsed = time.time() - start_time
    logger.info(f"AutoML complete. Best: {result.best_algorithm} ({result.best_score:.4f})")
    
    return result


def quick_run(
    X: Union[pd.DataFrame, np.ndarray], y: Union[pd.Series, np.ndarray],
    task: str = "auto"
) -> TrainedModel:
    """Quick AutoML run with defaults."""
    result = run(X, y, task, tuning="none", n_trials=10)
    return result.best_model
