"""
AdamOps ModelOps Module

Provides model training for regression, classification, and clustering.
"""

from typing import Any, Dict, List, Optional, Tuple, Union
import numpy as np
import pandas as pd
import joblib
from pathlib import Path

from sklearn.linear_model import Ridge, Lasso, ElasticNet, LogisticRegression
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
from sklearn.ensemble import (
    GradientBoostingClassifier, GradientBoostingRegressor,
    RandomForestClassifier, RandomForestRegressor
)
from sklearn.naive_bayes import GaussianNB
from sklearn.neighbors import KNeighborsClassifier, KNeighborsRegressor
from sklearn.cluster import KMeans, DBSCAN, AgglomerativeClustering
from sklearn.mixture import GaussianMixture

try:
    import xgboost as xgb
    XGB_AVAILABLE = True
except ImportError:
    XGB_AVAILABLE = False

try:
    import lightgbm as lgb
    LGB_AVAILABLE = True
except ImportError:
    LGB_AVAILABLE = False

from adamops.utils.logging import get_logger
from adamops.utils.helpers import infer_task_type

logger = get_logger(__name__)


# Model Registry
REGRESSION_MODELS = {
    "ridge": Ridge,
    "lasso": Lasso,
    "elasticnet": ElasticNet,
    "decision_tree": DecisionTreeRegressor,
    "random_forest": RandomForestRegressor,
    "gradient_boosting": GradientBoostingRegressor,
    "knn": KNeighborsRegressor,
}

CLASSIFICATION_MODELS = {
    "logistic": LogisticRegression,
    "decision_tree": DecisionTreeClassifier,
    "random_forest": RandomForestClassifier,
    "gradient_boosting": GradientBoostingClassifier,
    "naive_bayes": GaussianNB,
    "knn": KNeighborsClassifier,
}

CLUSTERING_MODELS = {
    "kmeans": KMeans,
    "dbscan": DBSCAN,
    "hierarchical": AgglomerativeClustering,
    "gmm": GaussianMixture,
}

# Add XGBoost if available
if XGB_AVAILABLE:
    REGRESSION_MODELS["xgboost"] = xgb.XGBRegressor
    CLASSIFICATION_MODELS["xgboost"] = xgb.XGBClassifier

# Add LightGBM if available
if LGB_AVAILABLE:
    REGRESSION_MODELS["lightgbm"] = lgb.LGBMRegressor
    CLASSIFICATION_MODELS["lightgbm"] = lgb.LGBMClassifier


class TrainedModel:
    """Wrapper for trained models with metadata."""
    
    def __init__(self, model: Any, task: str, algorithm: str, params: Dict,
                 feature_names: Optional[List[str]] = None):
        self.model = model
        self.task = task
        self.algorithm = algorithm
        self.params = params
        self.feature_names = feature_names
        self.is_fitted = True
    
    def predict(self, X: Union[pd.DataFrame, np.ndarray]) -> np.ndarray:
        """Make predictions."""
        return self.model.predict(X)
    
    def predict_proba(self, X: Union[pd.DataFrame, np.ndarray]) -> np.ndarray:
        """Predict probabilities (classification only)."""
        if hasattr(self.model, 'predict_proba'):
            return self.model.predict_proba(X)
        raise ValueError("Model does not support probability predictions")
    
    def save(self, filepath: Union[str, Path]) -> None:
        """Save model to file."""
        filepath = Path(filepath)
        filepath.parent.mkdir(parents=True, exist_ok=True)
        joblib.dump(self, filepath)
        logger.info(f"Model saved to {filepath}")
    
    @classmethod
    def load(cls, filepath: Union[str, Path]) -> "TrainedModel":
        """Load model from file."""
        return joblib.load(filepath)


def get_available_models(task: str = "classification") -> List[str]:
    """Get list of available models for a task."""
    if task == "classification":
        return list(CLASSIFICATION_MODELS.keys())
    elif task == "regression":
        return list(REGRESSION_MODELS.keys())
    elif task == "clustering":
        return list(CLUSTERING_MODELS.keys())
    return []


def train(
    X: Union[pd.DataFrame, np.ndarray], y: Union[pd.Series, np.ndarray],
    task: str = "auto", algorithm: str = "random_forest",
    params: Optional[Dict] = None, random_state: int = 42
) -> TrainedModel:
    """
    Train a model.
    
    Args:
        X: Features.
        y: Target.
        task: 'classification', 'regression', or 'auto'.
        algorithm: Model algorithm name.
        params: Model hyperparameters.
        random_state: Random seed.
    
    Returns:
        TrainedModel: Trained model wrapper.
    """
    # Auto-detect task type
    if task == "auto":
        task = infer_task_type(y)
        logger.info(f"Auto-detected task: {task}")
    
    # Get model class
    if task in ["classification", "multiclass"]:
        if algorithm not in CLASSIFICATION_MODELS:
            raise ValueError(f"Unknown classification algorithm: {algorithm}")
        model_class = CLASSIFICATION_MODELS[algorithm]
    elif task == "regression":
        if algorithm not in REGRESSION_MODELS:
            raise ValueError(f"Unknown regression algorithm: {algorithm}")
        model_class = REGRESSION_MODELS[algorithm]
    else:
        raise ValueError(f"Unknown task: {task}")
    
    # Set default params
    default_params = {"random_state": random_state}
    if params:
        default_params.update(params)
    
    # Filter params for model
    import inspect
    sig = inspect.signature(model_class)
    valid_params = {k: v for k, v in default_params.items() if k in sig.parameters}
    
    # Create and train model
    logger.info(f"Training {algorithm} for {task}")
    model = model_class(**valid_params)
    model.fit(X, y)
    
    feature_names = X.columns.tolist() if isinstance(X, pd.DataFrame) else None
    
    return TrainedModel(
        model=model, task=task, algorithm=algorithm,
        params=valid_params, feature_names=feature_names
    )


def train_regression(
    X: Union[pd.DataFrame, np.ndarray], y: Union[pd.Series, np.ndarray],
    algorithm: str = "ridge", params: Optional[Dict] = None
) -> TrainedModel:
    """Train a regression model."""
    return train(X, y, task="regression", algorithm=algorithm, params=params)


def train_classification(
    X: Union[pd.DataFrame, np.ndarray], y: Union[pd.Series, np.ndarray],
    algorithm: str = "random_forest", params: Optional[Dict] = None
) -> TrainedModel:
    """Train a classification model."""
    return train(X, y, task="classification", algorithm=algorithm, params=params)


def train_clustering(
    X: Union[pd.DataFrame, np.ndarray], algorithm: str = "kmeans",
    n_clusters: int = 3, params: Optional[Dict] = None
) -> Tuple[Any, np.ndarray]:
    """
    Train a clustering model.
    
    Returns:
        (model, labels): Fitted model and cluster labels.
    """
    if algorithm not in CLUSTERING_MODELS:
        raise ValueError(f"Unknown clustering algorithm: {algorithm}")
    
    model_class = CLUSTERING_MODELS[algorithm]
    default_params = {}
    
    # Set n_clusters for algorithms that support it
    if algorithm in ["kmeans", "hierarchical"]:
        default_params["n_clusters"] = n_clusters
    elif algorithm == "gmm":
        default_params["n_components"] = n_clusters
    
    if params:
        default_params.update(params)
    
    logger.info(f"Training {algorithm} clustering")
    model = model_class(**default_params)
    
    if algorithm == "gmm":
        model.fit(X)
        labels = model.predict(X)
    else:
        labels = model.fit_predict(X)
    
    return model, labels


def cross_validate(
    X: Union[pd.DataFrame, np.ndarray], y: Union[pd.Series, np.ndarray],
    task: str = "classification", algorithm: str = "random_forest",
    cv: int = 5, scoring: Optional[str] = None, params: Optional[Dict] = None
) -> Dict[str, Any]:
    """
    Cross-validate a model.
    
    Returns:
        Dict with train_scores, test_scores, and mean/std values.
    """
    from sklearn.model_selection import cross_validate as sklearn_cv
    
    # Get model
    if task == "classification":
        model = CLASSIFICATION_MODELS[algorithm]()
    else:
        model = REGRESSION_MODELS[algorithm]()
    
    if params:
        model.set_params(**params)
    
    if scoring is None:
        scoring = "accuracy" if task == "classification" else "r2"
    
    logger.info(f"Cross-validating {algorithm} with {cv} folds")
    
    results = sklearn_cv(model, X, y, cv=cv, scoring=scoring, return_train_score=True)
    
    return {
        "train_scores": results["train_score"].tolist(),
        "test_scores": results["test_score"].tolist(),
        "train_mean": float(results["train_score"].mean()),
        "train_std": float(results["train_score"].std()),
        "test_mean": float(results["test_score"].mean()),
        "test_std": float(results["test_score"].std()),
        "fit_time": float(results["fit_time"].mean()),
    }


def compare_models(
    X: Union[pd.DataFrame, np.ndarray], y: Union[pd.Series, np.ndarray],
    task: str = "classification", algorithms: Optional[List[str]] = None,
    cv: int = 5, scoring: Optional[str] = None
) -> pd.DataFrame:
    """
    Compare multiple models.
    
    Returns:
        DataFrame with model comparison results.
    """
    if algorithms is None:
        algorithms = list(CLASSIFICATION_MODELS.keys()) if task == "classification" \
            else list(REGRESSION_MODELS.keys())
    
    results = []
    for algo in algorithms:
        try:
            cv_results = cross_validate(X, y, task, algo, cv, scoring)
            results.append({
                "algorithm": algo,
                "cv_mean": cv_results["test_mean"],
                "cv_std": cv_results["test_std"],
                "train_mean": cv_results["train_mean"],
                "fit_time": cv_results["fit_time"],
            })
        except Exception as e:
            logger.warning(f"Failed to train {algo}: {e}")
    
    df = pd.DataFrame(results).sort_values("cv_mean", ascending=False)
    return df
