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

"""
HBV Model Configuration.

Pydantic schema and configuration adapter for HBV-96 model settings.
"""

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

from pydantic import BaseModel, Field, field_validator

from symfluence.models.base.base_config import AutoGeneratedConfigAdapter


class HBVConfig(BaseModel):
    """
    HBV-96 model configuration schema.

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

    # =========================================================================
    # Runtime Configuration
    # =========================================================================
    spatial_mode: Literal['lumped', 'distributed', 'auto'] = Field(
        default='auto',
        description="Spatial modeling mode. 'auto' infers from domain definition."
    )

    routing_integration: str = Field(
        default='none',
        description="Routing model integration. 'mizuRoute' for distributed routing, 'none' for standalone."
    )

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

    use_gpu: bool = Field(
        default=False,
        description="Use GPU acceleration (requires JAX with CUDA support)."
    )

    jit_compile: bool = Field(
        default=True,
        description="JIT compile model for faster execution (JAX backend only)."
    )

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

    timestep_hours: int = Field(
        default=24,
        ge=1,
        le=24,
        description=(
            "Model timestep in hours. Default is 24 (daily). Set to 1 for hourly simulation. "
            "Sub-daily timesteps require parameter scaling (see HBV-96 documentation). "
            "Flux rates (cfmax, perc) use linear scaling; recession coefficients (k0, k1, k2) "
            "use exact exponential scaling: k_sub = 1 - (1-k_daily)^(dt/24)."
        )
    )

    # =========================================================================
    # Calibration Configuration
    # =========================================================================
    params_to_calibrate: str = Field(
        default='tt,cfmax,fc,lp,beta,k0,k1,k2,uzl,perc,maxbas',
        description="Comma-separated list of parameters to calibrate."
    )

    use_gradient_calibration: bool = Field(
        default=True,
        description="Use gradient-based optimization (requires JAX). Falls back to evolutionary if False."
    )

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

    # =========================================================================
    # Initial State Values
    # =========================================================================
    initial_snow: float = Field(
        default=0.0,
        ge=0.0,
        description="Initial snow storage (mm SWE)."
    )

    initial_sm: float = Field(
        default=150.0,
        ge=0.0,
        description="Initial soil moisture (mm)."
    )

    initial_suz: float = Field(
        default=10.0,
        ge=0.0,
        description="Initial upper zone storage (mm)."
    )

    initial_slz: float = Field(
        default=10.0,
        ge=0.0,
        description="Initial lower zone storage (mm)."
    )

    # =========================================================================
    # Default Parameter Values (used when not calibrating)
    # =========================================================================
    default_tt: float = Field(default=0.0, description="Default threshold temperature (°C)")
    default_cfmax: float = Field(default=3.5, description="Default degree-day factor (mm/°C/day)")
    default_sfcf: float = Field(default=0.9, description="Default snowfall correction factor")
    default_cfr: float = Field(default=0.05, description="Default refreezing coefficient")
    default_cwh: float = Field(default=0.1, description="Default snow water holding capacity")
    default_fc: float = Field(default=250.0, description="Default field capacity (mm)")
    default_lp: float = Field(default=0.7, description="Default LP threshold (fraction of FC)")
    default_beta: float = Field(default=2.5, description="Default soil shape coefficient")
    default_k0: float = Field(default=0.3, description="Default surface runoff coefficient (1/day)")
    default_k1: float = Field(default=0.1, description="Default interflow coefficient (1/day)")
    default_k2: float = Field(default=0.01, description="Default baseflow coefficient (1/day)")
    default_uzl: float = Field(default=30.0, description="Default upper zone threshold (mm)")
    default_perc: float = Field(default=2.5, description="Default percolation rate (mm/day)")
    default_maxbas: float = Field(default=2.5, description="Default routing length (days)")

    # =========================================================================
    # PET Configuration
    # =========================================================================
    pet_method: Literal['input', 'hamon', 'thornthwaite'] = 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 Hamon/Thornthwaite)."
    )

    # =========================================================================
    # Distributed Routing Configuration
    # =========================================================================
    distributed_routing: Literal['none', 'muskingum_cunge', 'mizuroute'] = Field(
        default='none',
        description=(
            "River network routing method for distributed mode. "
            "'none' = no internal routing (use with mizuRoute externally), "
            "'muskingum_cunge' = internal differentiable Muskingum-Cunge routing, "
            "'mizuroute' = use external mizuRoute routing."
        )
    )

    distributed_param_mode: Literal['uniform', 'per_gru', 'regionalized'] = Field(
        default='uniform',
        description=(
            "HBV parameter mode for distributed simulation. "
            "'uniform' = same parameters for all GRUs (fewer parameters, simpler calibration), "
            "'per_gru' = individual parameters per GRU (more flexible, more parameters), "
            "'regionalized' = parameters derived from GRU attributes via transfer functions."
        )
    )

    routing_substep_method: Literal['fixed', 'adaptive'] = Field(
        default='adaptive',
        description=(
            "Sub-stepping method for Muskingum-Cunge routing stability. "
            "'fixed' = use fixed number of substeps, "
            "'adaptive' = automatically determine substeps based on Courant number."
        )
    )

    routing_max_substeps: int = Field(
        default=10,
        ge=1,
        le=100,
        description="Maximum number of sub-timesteps for Muskingum-Cunge routing stability."
    )

    default_mannings_n: float = Field(
        default=0.035,
        ge=0.01,
        le=0.2,
        description="Default Manning's n roughness coefficient for channel routing."
    )

    default_channel_slope: float = Field(
        default=0.001,
        ge=0.00001,
        le=0.1,
        description="Default channel slope (m/m) when not available from data."
    )

    # =========================================================================
    # Output Configuration
    # =========================================================================
    save_states: bool = Field(
        default=False,
        description="Save internal state variables (snow, soil moisture, etc.) to output."
    )

    output_frequency: Literal['daily', 'timestep'] = Field(
        default='daily',
        description="Output frequency. 'timestep' preserves sub-daily resolution if available."
    )

    @field_validator('params_to_calibrate')
    @classmethod
    def validate_params(cls, v: str) -> str:
        """Validate that parameter names are recognized."""
        from .model import PARAM_BOUNDS
        param_list = [p.strip() for p in v.split(',')]
        invalid = [p for p in param_list if p not in PARAM_BOUNDS]
        if invalid:
            raise ValueError(f"Unknown HBV parameters: {invalid}. Valid: {list(PARAM_BOUNDS.keys())}")
        return v

    def get_default_params(self) -> Dict[str, float]:
        """Get dictionary of default parameter values."""
        return {
            'tt': self.default_tt,
            'cfmax': self.default_cfmax,
            'sfcf': self.default_sfcf,
            'cfr': self.default_cfr,
            'cwh': self.default_cwh,
            'fc': self.default_fc,
            'lp': self.default_lp,
            'beta': self.default_beta,
            'k0': self.default_k0,
            'k1': self.default_k1,
            'k2': self.default_k2,
            'uzl': self.default_uzl,
            'perc': self.default_perc,
            'maxbas': self.default_maxbas,
        }

    def get_calibration_params(self) -> List[str]:
        """Get list of parameters to calibrate."""
        return [p.strip() for p in self.params_to_calibrate.split(',')]

    model_config = {'extra': 'forbid'}


class HBVConfigAdapter(AutoGeneratedConfigAdapter):
    """
    Configuration adapter for HBV model.

    Maps between config dictionary keys and HBVConfig schema.
    Uses AutoGeneratedConfigAdapter pattern for automatic field mapping.
    """

    MODEL_NAME = "HBV"
    CONFIG_PREFIX = "HBV_"

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

    @classmethod
    def get_defaults(cls) -> Dict[str, Any]:
        """Get default values for HBV configuration."""
        return {
            'HBV_SPATIAL_MODE': 'auto',
            'HBV_ROUTING_INTEGRATION': 'none',
            'HBV_BACKEND': 'jax',
            'HBV_USE_GPU': False,
            'HBV_JIT_COMPILE': True,
            'HBV_WARMUP_DAYS': 365,
            'HBV_TIMESTEP_HOURS': 24,
            'HBV_PARAMS_TO_CALIBRATE': 'tt,cfmax,fc,lp,beta,k0,k1,k2,uzl,perc,maxbas',
            'HBV_USE_GRADIENT_CALIBRATION': True,
            'HBV_CALIBRATION_METRIC': 'KGE',
            'HBV_INITIAL_SNOW': 0.0,
            'HBV_INITIAL_SM': 150.0,
            'HBV_INITIAL_SUZ': 10.0,
            'HBV_INITIAL_SLZ': 10.0,
            'HBV_PET_METHOD': 'input',
            'HBV_DISTRIBUTED_ROUTING': 'none',
            'HBV_DISTRIBUTED_PARAM_MODE': 'uniform',
            'HBV_ROUTING_SUBSTEP_METHOD': 'adaptive',
            'HBV_ROUTING_MAX_SUBSTEPS': 10,
            'HBV_DEFAULT_MANNINGS_N': 0.035,
            'HBV_DEFAULT_CHANNEL_SLOPE': 0.001,
            'HBV_SAVE_STATES': False,
            'HBV_OUTPUT_FREQUENCY': 'daily',
            # Default parameter values
            'HBV_DEFAULT_TT': 0.0,
            'HBV_DEFAULT_CFMAX': 3.5,
            'HBV_DEFAULT_SFCF': 0.9,
            'HBV_DEFAULT_CFR': 0.05,
            'HBV_DEFAULT_CWH': 0.1,
            'HBV_DEFAULT_FC': 250.0,
            'HBV_DEFAULT_LP': 0.7,
            'HBV_DEFAULT_BETA': 2.5,
            'HBV_DEFAULT_K0': 0.3,
            'HBV_DEFAULT_K1': 0.1,
            'HBV_DEFAULT_K2': 0.01,
            'HBV_DEFAULT_UZL': 30.0,
            'HBV_DEFAULT_PERC': 2.5,
            'HBV_DEFAULT_MAXBAS': 2.5,
        }

    @classmethod
    def _get_legacy_transformers(cls) -> Dict[str, tuple[str, type]]:
        """Get field transformation functions (Legacy format for from_dict)."""
        return {
            # Map config dict keys to HBVConfig field names
            'HBV_SPATIAL_MODE': ('spatial_mode', str),
            'HBV_ROUTING_INTEGRATION': ('routing_integration', str),
            'HBV_BACKEND': ('backend', str),
            'HBV_USE_GPU': ('use_gpu', bool),
            'HBV_JIT_COMPILE': ('jit_compile', bool),
            'HBV_WARMUP_DAYS': ('warmup_days', int),
            'HBV_TIMESTEP_HOURS': ('timestep_hours', int),
            'HBV_PARAMS_TO_CALIBRATE': ('params_to_calibrate', str),
            'HBV_USE_GRADIENT_CALIBRATION': ('use_gradient_calibration', bool),
            'HBV_CALIBRATION_METRIC': ('calibration_metric', str),
            'HBV_INITIAL_SNOW': ('initial_snow', float),
            'HBV_INITIAL_SM': ('initial_sm', float),
            'HBV_INITIAL_SUZ': ('initial_suz', float),
            'HBV_INITIAL_SLZ': ('initial_slz', float),
            'HBV_PET_METHOD': ('pet_method', str),
            'HBV_LATITUDE': ('latitude', float),
            'HBV_DISTRIBUTED_ROUTING': ('distributed_routing', str),
            'HBV_DISTRIBUTED_PARAM_MODE': ('distributed_param_mode', str),
            'HBV_ROUTING_SUBSTEP_METHOD': ('routing_substep_method', str),
            'HBV_ROUTING_MAX_SUBSTEPS': ('routing_max_substeps', int),
            'HBV_DEFAULT_MANNINGS_N': ('default_mannings_n', float),
            'HBV_DEFAULT_CHANNEL_SLOPE': ('default_channel_slope', float),
            'HBV_SAVE_STATES': ('save_states', bool),
            'HBV_OUTPUT_FREQUENCY': ('output_frequency', str),
            # Default parameters
            'HBV_DEFAULT_TT': ('default_tt', float),
            'HBV_DEFAULT_CFMAX': ('default_cfmax', float),
            'HBV_DEFAULT_SFCF': ('default_sfcf', float),
            'HBV_DEFAULT_CFR': ('default_cfr', float),
            'HBV_DEFAULT_CWH': ('default_cwh', float),
            'HBV_DEFAULT_FC': ('default_fc', float),
            'HBV_DEFAULT_LP': ('default_lp', float),
            'HBV_DEFAULT_BETA': ('default_beta', float),
            'HBV_DEFAULT_K0': ('default_k0', float),
            'HBV_DEFAULT_K1': ('default_k1', float),
            'HBV_DEFAULT_K2': ('default_k2', float),
            'HBV_DEFAULT_UZL': ('default_uzl', float),
            'HBV_DEFAULT_PERC': ('default_perc', float),
            'HBV_DEFAULT_MAXBAS': ('default_maxbas', float),
        }

    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any]) -> HBVConfig:
        """
        Create HBVConfig from a configuration dictionary.

        Args:
            config_dict: Dictionary with HBV_* keys

        Returns:
            HBVConfig instance
        """
        transformers = cls._get_legacy_transformers()
        defaults = cls.get_defaults()

        # Build kwargs for HBVConfig
        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 HBVConfig(**kwargs)

    @classmethod
    def to_dict(cls, config: HBVConfig) -> Dict[str, Any]:
        """
        Convert HBVConfig to configuration dictionary.

        Args:
            config: HBVConfig instance

        Returns:
            Dictionary with HBV_* keys
        """
        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
