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

class ModelEvalConfig(BaseModel):
    """Configuration for evaluating a single forecasting model or an ensemble."""
    name: str = Field(..., description="Name of the forecasting model or ensemble.")
    type: Literal['model', 'ensemble'] = Field(..., description="Type of evaluation artifact: 'model' for a single checkpoint, 'ensemble' for an ensemble definition JSON.")
    model_path: str = Field(..., description="Path to the saved PyTorch model file (.ckpt for type='model') or the ensemble definition JSON file (.json for type='ensemble').")
    model_config_path: str = Field(..., description="Path to the forecasting config (YAML) used for this model training (or for the best trial in an ensemble).")
    target_scaler_path: Optional[str] = Field(None, description="Path to the target scaler file for the single model (or will be loaded per fold for ensemble).")

class OptimizationRunConfig(BaseModel):
    """Main configuration for running battery optimization with multiple models/ensembles."""
    initial_b: float = Field(..., description="Initial state of charge of the battery (MWh).")
    max_capacity: float = Field(..., description="Maximum energy capacity of the battery (MWh).")
    max_rate: float = Field(..., description="Maximum charge/discharge power rate of the battery (MW).")
    optimization_horizon_hours: int = Field(24, gt=0, description="The length of the time window (in hours) for optimization.")
    models: List[ModelEvalConfig] = Field(..., description="List of forecasting models or ensembles to evaluate.")