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

"""
HEC-HMS Model Configuration.

Pydantic schema and configuration adapter for HEC-HMS model settings.
"""

from typing import Any, Dict, List, Literal, Optional

from pydantic import BaseModel, Field

from symfluence.models.base.base_config import AutoGeneratedConfigAdapter


class HECHMSConfig(BaseModel):
    """
    HEC-HMS model configuration schema.

    Defines all configurable aspects of the HEC-HMS model including runtime options,
    calibration settings, and initial state values.
    """

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

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

    # =========================================================================
    # Calibration Configuration
    # =========================================================================
    params_to_calibrate: str = Field(
        default='default',
        description=(
            "Comma-separated list of parameters to calibrate. "
            "'default' or empty string triggers use of all available parameters."
        )
    )

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

    # =========================================================================
    # PET Configuration
    # =========================================================================
    pet_method: Literal['input', 'oudin', 'hamon'] = Field(
        default='input',
        description="PET calculation method. 'input' uses provided PET, others compute from temperature."
    )

    latitude: Optional[float] = Field(
        default=None,
        ge=-90.0,
        le=90.0,
        description="Catchment latitude for PET calculation (required for Oudin/Hamon)."
    )

    # =========================================================================
    # Default Parameter Values
    # =========================================================================
    default_px_temp: float = Field(default=1.0, description="Rain/snow partition temperature (deg C)")
    default_base_temp: float = Field(default=0.0, description="Base temperature for snowmelt (deg C)")
    default_ati_meltrate_coeff: float = Field(default=0.98, description="ATI meltrate coefficient")
    default_meltrate_max: float = Field(default=5.0, description="Maximum melt rate (mm/deg C/day)")
    default_meltrate_min: float = Field(default=1.0, description="Minimum melt rate (mm/deg C/day)")
    default_cold_limit: float = Field(default=10.0, description="Cold limit (mm)")
    default_ati_cold_rate_coeff: float = Field(default=0.1, description="ATI cold rate coefficient")
    default_water_capacity: float = Field(default=0.05, description="Snowpack liquid water capacity (fraction)")
    default_cn: float = Field(default=65.0, description="SCS Curve Number")
    default_initial_abstraction_ratio: float = Field(default=0.2, description="Ia/S ratio")
    default_tc: float = Field(default=3.0, description="Time of concentration (days)")
    default_r_coeff: float = Field(default=5.0, description="Clark storage coefficient (days)")
    default_gw_storage_coeff: float = Field(default=30.0, description="GW storage coefficient (days)")
    default_deep_perc_fraction: float = Field(default=0.1, description="Deep percolation fraction")

    # =========================================================================
    # Initial State Values
    # =========================================================================
    initial_snow_swe: float = Field(default=0.0, ge=0.0, description="Initial SWE (mm)")
    initial_gw_storage: float = Field(default=10.0, ge=0.0, description="Initial GW storage (mm)")

    def get_default_params(self) -> Dict[str, float]:
        """Get dictionary of default parameter values."""
        return {
            'px_temp': self.default_px_temp,
            'base_temp': self.default_base_temp,
            'ati_meltrate_coeff': self.default_ati_meltrate_coeff,
            'meltrate_max': self.default_meltrate_max,
            'meltrate_min': self.default_meltrate_min,
            'cold_limit': self.default_cold_limit,
            'ati_cold_rate_coeff': self.default_ati_cold_rate_coeff,
            'water_capacity': self.default_water_capacity,
            'cn': self.default_cn,
            'initial_abstraction_ratio': self.default_initial_abstraction_ratio,
            'tc': self.default_tc,
            'r_coeff': self.default_r_coeff,
            'gw_storage_coeff': self.default_gw_storage_coeff,
            'deep_perc_fraction': self.default_deep_perc_fraction,
        }

    def get_calibration_params(self) -> List[str]:
        """Get list of parameters to calibrate."""
        if not self.params_to_calibrate or self.params_to_calibrate == 'default':
            from .parameters import PARAM_BOUNDS
            return list(PARAM_BOUNDS.keys())
        return [p.strip() for p in self.params_to_calibrate.split(',')]

    model_config = {'extra': 'forbid'}


class HecHmsConfigAdapter(AutoGeneratedConfigAdapter):
    """
    Configuration adapter for HEC-HMS model.

    Maps between config dictionary keys and HECHMSConfig schema.
    """

    MODEL_NAME = "HECHMS"
    CONFIG_PREFIX = "HECHMS_"

    @classmethod
    def get_config_schema(cls) -> type:
        """Return the Pydantic config schema class."""
        return HECHMSConfig

    @classmethod
    def get_defaults(cls) -> Dict[str, Any]:
        """Get default values for HEC-HMS configuration."""
        return {
            'HECHMS_BACKEND': 'jax',
            'HECHMS_WARMUP_DAYS': 365,
            'HECHMS_PARAMS_TO_CALIBRATE': 'default',
            'HECHMS_CALIBRATION_METRIC': 'KGE',
            'HECHMS_PET_METHOD': 'input',
            'HECHMS_INITIAL_SNOW_SWE': 0.0,
            'HECHMS_INITIAL_GW_STORAGE': 10.0,
            # Default parameters
            'HECHMS_DEFAULT_PX_TEMP': 1.0,
            'HECHMS_DEFAULT_BASE_TEMP': 0.0,
            'HECHMS_DEFAULT_ATI_MELTRATE_COEFF': 0.98,
            'HECHMS_DEFAULT_MELTRATE_MAX': 5.0,
            'HECHMS_DEFAULT_MELTRATE_MIN': 1.0,
            'HECHMS_DEFAULT_COLD_LIMIT': 10.0,
            'HECHMS_DEFAULT_ATI_COLD_RATE_COEFF': 0.1,
            'HECHMS_DEFAULT_WATER_CAPACITY': 0.05,
            'HECHMS_DEFAULT_CN': 65.0,
            'HECHMS_DEFAULT_INITIAL_ABSTRACTION_RATIO': 0.2,
            'HECHMS_DEFAULT_TC': 3.0,
            'HECHMS_DEFAULT_R_COEFF': 5.0,
            'HECHMS_DEFAULT_GW_STORAGE_COEFF': 30.0,
            'HECHMS_DEFAULT_DEEP_PERC_FRACTION': 0.1,
        }

    @classmethod
    def _get_legacy_transformers(cls) -> Dict[str, tuple]:
        """Get field transformation functions."""
        return {
            'HECHMS_BACKEND': ('backend', str),
            'HECHMS_WARMUP_DAYS': ('warmup_days', int),
            'HECHMS_PARAMS_TO_CALIBRATE': ('params_to_calibrate', str),
            'HECHMS_CALIBRATION_METRIC': ('calibration_metric', str),
            'HECHMS_PET_METHOD': ('pet_method', str),
            'HECHMS_LATITUDE': ('latitude', float),
            'HECHMS_INITIAL_SNOW_SWE': ('initial_snow_swe', float),
            'HECHMS_INITIAL_GW_STORAGE': ('initial_gw_storage', float),
            # Default parameters
            'HECHMS_DEFAULT_PX_TEMP': ('default_px_temp', float),
            'HECHMS_DEFAULT_BASE_TEMP': ('default_base_temp', float),
            'HECHMS_DEFAULT_ATI_MELTRATE_COEFF': ('default_ati_meltrate_coeff', float),
            'HECHMS_DEFAULT_MELTRATE_MAX': ('default_meltrate_max', float),
            'HECHMS_DEFAULT_MELTRATE_MIN': ('default_meltrate_min', float),
            'HECHMS_DEFAULT_COLD_LIMIT': ('default_cold_limit', float),
            'HECHMS_DEFAULT_ATI_COLD_RATE_COEFF': ('default_ati_cold_rate_coeff', float),
            'HECHMS_DEFAULT_WATER_CAPACITY': ('default_water_capacity', float),
            'HECHMS_DEFAULT_CN': ('default_cn', float),
            'HECHMS_DEFAULT_INITIAL_ABSTRACTION_RATIO': ('default_initial_abstraction_ratio', float),
            'HECHMS_DEFAULT_TC': ('default_tc', float),
            'HECHMS_DEFAULT_R_COEFF': ('default_r_coeff', float),
            'HECHMS_DEFAULT_GW_STORAGE_COEFF': ('default_gw_storage_coeff', float),
            'HECHMS_DEFAULT_DEEP_PERC_FRACTION': ('default_deep_perc_fraction', float),
        }

    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any]) -> HECHMSConfig:
        """Create HECHMSConfig from a configuration dictionary."""
        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 HECHMSConfig(**kwargs)

    @classmethod
    def to_dict(cls, config: HECHMSConfig) -> Dict[str, Any]:
        """Convert HECHMSConfig to configuration dictionary."""
        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
