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

class AutoGluonConfig(BaseModel):
    eval_metric: str = Field(
        ..., 
        description="Primary metric to optimize. Allowed: 'accuracy', 'log_loss', 'f1', 'roc_auc', 'precision', 'recall'."
    )
    
    preset: Optional[Literal[
        'best_quality', 'high_quality', 'good_quality', 'medium_quality'
    ]] = Field(
        default='good_quality',
        description=(
            "Optional preset configuration. "
            "'best_quality' 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_quality'. 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]."
    )


class RecommendationConfig(BaseModel):
    algorithm : Literal[
        "user_based",
        "item_based"
    ] = Field(
        ...,
        description=(
            "Type of collaborative filtering to apply. "
            "'user_based' computes similarity between users, "
            "'item_based' computes similarity between items."
        )
    )

    top_k_similar : int = Field(
        ...,
        ge=5,
        le=50,
        description=(
            "Number of nearest neighbors used for recommendation. "
            "Minimum: 5, Maximum: 50. "
            "Higher values produce smoother recommendations but increase computation cost."
        )
    )

    time_limit: int = Field(
        ...,
        ge=60,
        le=3600,
        description=(
            "Total computation time in seconds allocated for similarity matrix construction "
            "and neighborhood computation. "
            "Minimum: 60 seconds, Maximum: 3600 seconds. "
            "Value must be calculated using the recommendation system time estimation formula."
        )
    )

    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]."
    )

    similarity_threshold: float = Field(
        ...,
        ge=0.0,
        le=1.0,
        description=(
            "Minimum similarity score required for a user or item to be "
            "considered a valid neighbor. "
            "Neighbors with similarity below this threshold are ignored. "
            "Typical values: "
            "0.2–0.4 for cosine similarity, "
            "0.3–0.6 for pearson correlation."
        )
    )
