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

class AutoGluonConfig(BaseModel):
    eval_metric: str = Field(
        ..., 
        description="Primary metric to optimize. Binary: 'f1', 'roc_auc', 'accuracy', 'balanced_accuracy', 'precision', 'recall', 'log_loss', 'mcc'. Multiclass: 'accuracy', 'balanced_accuracy', 'log_loss', 'f1_macro', 'f1_micro', 'f1_weighted'. Regression: 'rmse', 'mae', 'r2', 'mse', 'mape'."
    )
    
    preset: Literal["best_quality", "high_quality", "good_quality", "medium_quality", "optimize_for_deployment", "interpretable"] = Field(
        ..., 
        description="Preset configurations. 'best_quality' enables bagging/stacking. 'high_quality' balances accuracy/speed. 'good_quality' is the recommended default. 'medium_quality' is for fast prototyping. 'optimize_for_deployment' minimizes footprint."
    )
    
    additional_metrics: List[str] = Field(
        ..., 
        description="List of additional metrics to track. Options: 'f1', 'precision', 'recall', 'accuracy', 'roc_auc', 'balanced_accuracy', 'log_loss', 'rmse', 'mae', 'r2', 'mse', 'mcc', 'mape'."
    )
    
    time_limit: int = Field(
        ..., 
        description="Total training time in seconds. AutoGluon will distribute this across models. Small datasets: 300, Medium: 3600, Large: 7200+."
    )
    
    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 and their hyperparameters. Can be a string alias or a dict of model keys: 
        'GBM' (LightGBM), 'CAT' (CatBoost), 'XGB' (XGBoost), 'NN' (Neural Net), 'FASTAI' (FastAI NN), 
        'RF' (Random Forest), 'XT' (Extra Trees), 'KNN' (K-Neighbors), 'LR' (Linear Regression)."""
    )

    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 if bagging is not used. Defaults to 0.1-0.2 depending on dataset size."
    )