from pydantic import BaseModel, Field
from typing import List, Literal, Optional, Union, Dict

# ============================================================================
# CLASSIFICATION MODELS (Binary, Multi, Timeseries Binary, Timeseries Multi)
# ============================================================================

class ClassificationUserExperimentConfig(BaseModel):
    """User experiment configuration for classification training"""

    eval_metric: str = Field(
        ...,
        description="Primary evaluation metric (e.g. 'f1' for binary, 'f1_macro' for multi-class)"
    )

    preset: Literal['extreme', 'best', 'high', 'good', 'medium'] = Field(
        default='good',
        description="AutoGluon preset configuration"
    )

    additional_metrics: List[str] = Field(
        default=["accuracy", "precision", "recall", "f1"],
        description="Extra metrics to compute"
    )

    time_limit: int = Field(
        default=60,
        gt=0,
        description="Training time budget in seconds (any integer > 0)"
    )

    models: List[str] = Field(
        default=["GBM", "CAT", "XGB", "XT", "RF", "KNN"],
        description="Models to include: GBM, CAT, XGB, XT, RF, KNN"
    )

    split_test_size: float = Field(
        default=0.2,
        gt=0.0,
        le=0.5,
        description="Fraction of data held out for validation. Choose from: [0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5]."
    )

    num_bag_sets: int = Field(
        default=2,
        gt=0,
        description="Number of bagging sets (any integer > 0)"
    )

    num_bag_folds: int = Field(
        default=5,
        gt=0,
        description="CV folds per bagging set (any integer > 0)"
    )

    num_stack_levels: int = Field(
        default=1,
        ge=0,
        description="Stacking depth (any integer >= 0)"
    )



# ============================================================================
# REGRESSION MODELS (Regression, Timeseries Regression)
# ============================================================================

class RegressionUserExperimentConfig(BaseModel):
    """User experiment configuration for regression training"""

    eval_metric: str = Field(
        default="root_mean_squared_error",
        description="Primary metric (any AutoGluon-supported regression metric)"
    )

    preset: Literal['extreme', 'best', 'high', 'good', 'medium'] = Field(
        default='good',
        description="AutoGluon preset"
    )

    additional_metrics: List[str] = Field(
        default=["mean_absolute_error", "r2", "mean_squared_error", "median_absolute_error"],
        description="Extra metrics to track"
    )

    time_limit: int = Field(
        default=60,
        gt=0,
        description="Training time budget in seconds (any integer > 0)"
    )

    models: List[str] = Field(
        default=["GBM", "CAT", "XGB", "XT", "RF", "KNN"],
        description="Models to include: GBM, CAT, XGB, XT, RF, KNN"
    )

    split_test_size: float = Field(
        default=0.2,
        gt=0.0,
        le=0.5,
        description="Fraction of data held out for validation. Choose from: [0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5]."
    )

    num_bag_sets: int = Field(
        default=2,
        gt=0,
        description="Number of bagging sets (any integer > 0)"
    )

    num_bag_folds: int = Field(
        default=5,
        gt=0,
        description="CV folds per bagging set (any integer > 0)"
    )

    num_stack_levels: int = Field(
        default=1,
        ge=0,
        description="Stacking depth (any integer >= 0)"
    )



# ============================================================================
# FORECASTING MODEL
# ============================================================================

class ForecastingUserExperimentConfig(BaseModel):
    """User experiment configuration for forecasting"""

    eval_metric: str = Field(
        default="MASE",
        description="Primary evaluation metric (e.g. MASE, RMSE, MAPE)"
    )

    preset: Literal['extreme', 'best', 'high', 'good', 'medium'] = Field(
        default="medium",
        description="Forecasting preset"
    )

    additional_metrics: List[str] = Field(
        default=["RMSE", "MAPE", "sMAPE"],
        description="Extra metrics to track"
    )

    time_limit: int = Field(
        default=60,
        gt=0,
        description="Training time budget in seconds (any integer > 0)"
    )

    models: List[str] = Field(
        default=["ETS", "ARIMA", "Theta", "DeepAR", "AutoETS"],
        description="Forecasting models to include"
    )

    split_test_size: float = Field(
        default=0.2,
        gt=0.0,
        le=0.5,
        description="Fraction of data held out for validation. Choose from: [0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5]."
    )



# ============================================================================
# ANOMALY DETECTION MODEL
# ============================================================================

class AnomalyDetectionUserExperimentConfig(BaseModel):
    """User experiment configuration for anomaly detection"""

    eval_metric: str = Field(
        default="f1",
        description="Primary evaluation metric (e.g. 'f1', 'precision', 'recall', 'roc_auc')"
    )

    preset: Literal['extreme', 'best', 'high', 'good', 'medium'] = Field(
        default='good',
        description="Quality preset"
    )

    additional_metrics: List[str] = Field(
        default=["precision", "recall", "roc_auc"],
        description="Extra metrics to track"
    )

    time_limit: int = Field(
        default=60,
        gt=0,
        description="Time budget in seconds (any integer > 0)"
    )

    models: List[str] = Field(
        default=["IForest", "LOF", "COPOD", "ECOD", "KNN", "OCSVM", "HBOS", "LODA"],
        description="Detectors to use"
    )

    split_test_size: float = Field(
        default=0.2,
        gt=0.0,
        le=0.5,
        description="Fraction of data held out for validation. Choose from: [0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5]."
    )

    threshold_method: Literal["std", "contamination", "elbow", "otsu", "auto"] = Field(
        default="auto",
        description="Threshold determination method"
    )

    threshold_stds: float = Field(
        default=3.0,
        gt=0,
        description="Std devs for STD threshold method (used only when threshold_method='std')"
    )

    contamination: Optional[float] = Field(
        default=None,
        description="Expected contamination fraction 0 < x < 1 (used only when threshold_method='contamination')"
    )

    top_n_models: int = Field(
        default=10,
        gt=0,
        description="Top models to keep (any integer > 0)"
    )

    max_combinations: int = Field(
        default=50,
        gt=0,
        description="Max ensemble combinations to evaluate (any integer > 0)"
    )

    lof_n_neighbors: List[int] = Field(
        default=[15, 20, 25, 35],
        description="LOF neighbor sweep values (e.g. [15, 20, 25, 35])"
    )

    iforest_estimators: List[int] = Field(
        default=[100, 150],
        description="IForest estimator sweep values (e.g. [100, 150])"
    )



# ============================================================================
# RECOMMENDATION MODEL
# ============================================================================

class RecommendationUserExperimentConfig(BaseModel):
    """User experiment configuration for recommendation systems"""

    eval_metric: str = Field(
        default="precision_at_k",
        description="Primary evaluation metric (e.g. 'precision_at_k', 'recall_at_k', 'ndcg', 'map')"
    )

    preset: Literal['extreme', 'best', 'high', 'good', 'medium'] = Field(
        default='good',
        description="Quality preset"
    )

    additional_metrics: List[str] = Field(
        default=["recall_at_k", "ndcg", "map"],
        description="Extra metrics to track"
    )

    time_limit: int = Field(
        default=60,
        gt=0,
        description="Time budget in seconds (any integer > 0)"
    )

    models: List[str] = Field(
        default=["user_based_cf"],
        description="Algorithms to use: user_based_cf, item_based_cf, svd, als (max 4)"
    )

    split_test_size: float = Field(
        default=0.2,
        gt=0.0,
        le=0.5,
        description="Fraction of interactions held out for evaluation. Choose from: [0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5]."
    )

    customer_id_column: str = Field(
        ...,
        description="User/customer ID column (required)"
    )

    product_id: str = Field(
        ...,
        description="Item/product ID column (required)"
    )

    max_pca_components: int = Field(
        default=9,
        gt=0,
        description="Max PCA components (any integer > 0)"
    )

    pca_variance_threshold: float = Field(
        default=0.95,
        gt=0.0,
        le=1.0,
        description="Variance threshold for PCA (0 < x <= 1)"
    )

    scoring_method: str = Field(
        default="weighted_mean",
        description="Score aggregation method"
    )

    similarity_threshold: float = Field(
        default=0.7,
        ge=0.0,
        le=1.0,
        description="Min similarity for user-based CF (0 <= x <= 1)"
    )

    top_k_similar: int = Field(
        default=10,
        gt=0,
        description="Similar users/items to consider (any integer > 0)"
    )

    min_user_interactions: int = Field(
        default=3,
        gt=0,
        description="Min interactions required per user (any integer > 0)"
    )

    item_similarity_threshold: float = Field(
        default=0.05,
        ge=0.0,
        le=1.0,
        description="Min similarity for item-based CF (0 <= x <= 1)"
    )

    svd_factors: int = Field(
        default=50,
        gt=0,
        description="SVD latent factors (any integer > 0)"
    )

    svd_iterations: int = Field(
        default=20,
        gt=0,
        description="SVD iterations (any integer > 0)"
    )

    als_factors: int = Field(
        default=50,
        gt=0,
        description="ALS latent factors (any integer > 0)"
    )

    als_iterations: int = Field(
        default=15,
        gt=0,
        description="ALS iterations (any integer > 0)"
    )

    als_regularization: float = Field(
        default=0.1,
        gt=0,
        description="ALS regularization strength (any float > 0)"
    )

    als_alpha: float = Field(
        default=40.0,
        gt=0,
        description="ALS confidence scaling (any float > 0)"
    )



# ============================================================================
# CLUSTERING MODEL
# ============================================================================

class ClusteringUserExperimentConfig(BaseModel):
    """User experiment configuration for clustering"""

    eval_metric: str = Field(
        default="silhouette",
        description="Primary evaluation metric (e.g. 'silhouette', 'davies_bouldin', 'calinski_harabasz')"
    )

    preset: Literal['extreme', 'best', 'high', 'good', 'medium'] = Field(
        default='good',
        description="Quality preset"
    )

    additional_metrics: List[str] = Field(
        default=["davies_bouldin", "calinski_harabasz"],
        description="Extra metrics to track"
    )

    time_limit: int = Field(
        default=120,
        gt=0,
        description="Time budget in seconds (any integer > 0)"
    )

    models: List[str] = Field(
        default=["kmeans", "minibatch_kmeans", "agglomerative", "gmm", "hdbscan", "dbscan", "spectral"],
        description="Clustering algorithms to use (all 7 by default)"
    )

    split_test_size: float = Field(
        default=0.2,
        gt=0.0,
        le=0.5,
        description="Fraction of data held out for validation. Choose from: [0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5]."
    )

    min_clusters: int = Field(
        default=2,
        ge=2,
        description="Min cluster count to search (any integer >= 2)"
    )

    max_clusters: int = Field(
        default=15,
        ge=2,
        description="Max cluster count to search (any integer >= min_clusters)"
    )

    top_n_models: int = Field(
        default=15,
        gt=0,
        description="Top models to save (any integer > 0)"
    )

    clustering_metric_weights: Dict[str, float] = Field(
        default={
            "silhouette": 0.4,
            "davies_bouldin": 0.3,
            "calinski_harabasz": 0.2,
            "noise_penalty": 0.1
        },
        description="Composite score weights (must sum to 1.0)"
    )

    dbscan_eps: Optional[float] = Field(
        default=None,
        description="DBSCAN epsilon radius - if omitted, sweeps 0.1 to 2.0 (any float > 0 or None for auto-sweep)"
    )

    dbscan_min_samples: int = Field(
        default=5,
        gt=0,
        description="DBSCAN min samples in neighborhood (any integer > 0)"
    )

    agglomerative_linkage: Literal["ward", "complete", "average", "single"] = Field(
        default="ward",
        description="Linkage method for agglomerative clustering"
    )



# ============================================================================
# LEGACY MODELS (for backward compatibility)
# ============================================================================

class AutoGluonConfig(BaseModel):
    """Legacy model - kept for backward compatibility"""
    eval_metric: str = Field(
        ...,
        description="Primary metric to optimize. Allowed: 'accuracy', 'log_loss', 'f1', 'roc_auc', 'precision', 'recall'."
    )

    preset: Optional[Literal[
        'extreme', 'best', 'high', 'good', 'medium'
    ]] = Field(
        default='good',
        description=(
            "Optional preset configuration. "
            "'best' enables bagging/stacking for maximum accuracy. "
            "If not provided, AutoGluon default behavior will be used."
        )
    )

    additional_metrics: List[str] = Field(
        ...,
        description="List of additional metrics to track. Allowed: 'accuracy', 'log_loss', 'f1', 'roc_auc', 'precision', 'recall'."
    )

    time_limit: int = Field(
        ...,
        description="Total training time in seconds. AutoGluon will distribute this across models. Small datasets (upto 500000 rows): Value between 60 seconds to 600 seconds, Medium (upto 7000000 rows): Value between 601 seconds to 1800 seconds, Large (7000000+ rows): Value between 1800 seconds to 3600 seconds."
    )

    num_bag_folds: int = Field(
        ...,
        description="Number of folds for k-fold bagging. 0 = no bagging. 5-10 is standard for 'best'. Bagging reduces variance and allows the model to be trained on all data (if refit_full=True)."
    )

    num_bag_sets: int = Field(
        ...,
        description="Number of bagging sets. Each set repeats k-fold bagging to reduce variance further. Only used if num_bag_folds > 0. Usually 1-3."
    )

    num_stack_levels: int = Field(
        ...,
        description="Levels of stacking. 0 = no stacking, 1 = one level (models trained on base model predictions). Higher values increase accuracy but exponentially increase training time."
    )

    models: list[str] = Field(
        ...,
        description="""Models to train.
        'GBM',
        'CAT',
        'XGB',
        'RF',
        'XT',
        'KNN'. """
    )

    fit_weighted_ensemble: bool = Field(
        ...,
        description="Whether to fit an ensemble that weights predictions of base models to improve accuracy. Usually recommended to keep True."
    )

    split_test_size: float = Field(
        ...,
        description="Fraction of data held out for validation. You MUST choose exactly one value from: [0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5]."
    )
