# SPDX-License-Identifier: Apache-2.0
# Copyright (C) 2024-2026 SYMFLUENCE Team <dev@symfluence.org>

"""
SAC-SMA Model Configuration.

Pydantic schema and configuration adapter for SAC-SMA + Snow-17 model settings.
"""

from typing import Any, Dict, Literal, Optional

from pydantic import BaseModel, Field

from symfluence.models.base.base_config import AutoGeneratedConfigAdapter


class SacSmaConfig(BaseModel):
    """SAC-SMA + Snow-17 model configuration schema."""

    # Runtime configuration
    warmup_days: int = Field(
        default=365, ge=0,
        description="Number of warmup days for model spinup."
    )

    pet_method: Literal['input', 'hamon'] = Field(
        default='input',
        description="PET method. 'input' uses provided PET, 'hamon' computes from temperature."
    )

    latitude: Optional[float] = Field(
        default=None, ge=-90.0, le=90.0,
        description="Catchment latitude for PET and melt factor calculations."
    )

    si: float = Field(
        default=100.0, ge=0.0,
        description="SWE threshold for 100% areal snow coverage (mm)."
    )

    save_states: bool = Field(
        default=False,
        description="Save internal state variables to output."
    )

    # Calibration configuration
    params_to_calibrate: str = Field(
        default='all',
        description="Comma-separated list of parameters to calibrate, or 'all' for all 26."
    )

    calibration_metric: Literal['KGE', 'NSE'] = Field(
        default='KGE',
        description="Objective function for calibration."
    )

    snow_module: Literal['snow17', 'none'] = Field(
        default='snow17',
        description="Snow module: 'snow17' for coupled mode, 'none' for standalone SAC-SMA."
    )

    backend: Literal['jax', 'numpy'] = Field(
        default='numpy',
        description="Computation backend: 'jax' for autodiff/JIT, 'numpy' for standard."
    )

    model_config = {'extra': 'forbid'}


class SacSmaConfigAdapter(AutoGeneratedConfigAdapter):
    """Configuration adapter for SAC-SMA model."""

    MODEL_NAME = "SACSMA"
    CONFIG_PREFIX = "SACSMA_"

    @classmethod
    def get_config_schema(cls) -> type:
        return SacSmaConfig

    @classmethod
    def get_defaults(cls) -> Dict[str, Any]:
        return {
            'SACSMA_WARMUP_DAYS': 365,
            'SACSMA_PET_METHOD': 'input',
            'SACSMA_SI': 100.0,
            'SACSMA_SAVE_STATES': False,
            'SACSMA_PARAMS_TO_CALIBRATE': 'all',
            'SACSMA_CALIBRATION_METRIC': 'KGE',
            'SACSMA_SNOW_MODULE': 'snow17',
            'SACSMA_BACKEND': 'numpy',
        }

    @classmethod
    def _get_legacy_transformers(cls) -> Dict[str, tuple]:
        return {
            'SACSMA_WARMUP_DAYS': ('warmup_days', int),
            'SACSMA_PET_METHOD': ('pet_method', str),
            'SACSMA_LATITUDE': ('latitude', float),
            'SACSMA_SI': ('si', float),
            'SACSMA_SAVE_STATES': ('save_states', bool),
            'SACSMA_PARAMS_TO_CALIBRATE': ('params_to_calibrate', str),
            'SACSMA_CALIBRATION_METRIC': ('calibration_metric', str),
            'SACSMA_SNOW_MODULE': ('snow_module', str),
            'SACSMA_BACKEND': ('backend', str),
        }

    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any]) -> SacSmaConfig:
        transformers = cls._get_legacy_transformers()
        defaults = cls.get_defaults()
        kwargs = {}
        for config_key, (field_name, field_type) in transformers.items():
            value = config_dict.get(config_key, defaults.get(config_key))
            if value is not None:
                try:
                    kwargs[field_name] = field_type(value)
                except (ValueError, TypeError):
                    kwargs[field_name] = value
        return SacSmaConfig(**kwargs)

    @classmethod
    def to_dict(cls, config: SacSmaConfig) -> Dict[str, Any]:
        transformers = cls._get_legacy_transformers()
        result = {}
        for config_key, (field_name, _) in transformers.items():
            if hasattr(config, field_name):
                result[config_key] = getattr(config, field_name)
        return result
