Source code for improver.fire_weather

# (C) Crown Copyright, Met Office. All rights reserved.
#
# This file is part of 'IMPROVER' and is released under the BSD 3-Clause license.
# See LICENSE in the root of the repository for full licensing details.
"""Canadian Forest Fire Weather Index System components."""

import warnings
from abc import abstractmethod
from copy import deepcopy
from typing import Union, cast

import numpy as np
from iris.cube import Cube, CubeList
from iris.exceptions import ConstraintMismatchError

from improver import BasePlugin
from improver.utilities.common_input_handle import as_cubelist
from improver.utilities.copy_metadata import CopyMetadata
from improver.utilities.load import load_baseline_cube




class FireWeatherBase(BasePlugin):
    """
    Abstract base class for the Canadian Forest Fire Weather Index (CFFWI) System
    calculations.

    This class provides common functionality for all fire weather
    components, including:

    - Standardised cube loading and validation
    - Fixed unit conversions for all cube types (non-configurable)
    - Output cube creation
    - Process orchestration

    The CFFWI system requires specific units
    for all calculations. These are fixed and cannot be overridden:

    - Temperature: degrees Celsius (Celsius)
    - Precipitation: millimeters (mm)
    - Relative humidity: dimensionless fraction (1)
    - Wind speed: kilometers per hour (km/h)
    - All fire weather indices: dimensionless (1)

    Subclasses must define class attributes:

    - METADATA_SOURCE_CUBE: The name of the input cube from which the
        metadata will be sourced for the output_cube. For downstream
        datasets that take iterative datasets as input this should
        be the iterative dataset.
    - INPUT_CUBE_NAMES: List of standard names for required input cubes
    - OUTPUT_CUBE_NAME: Standard name for the output cube
    - REQUIRES_MONTH: Boolean indicating if month parameter is required

    Subclasses must implement:

    - _calculate(): Method that performs the actual calculation
    """

    # Fixed unit conversions for all cube types used in fire weather calculations
    # These units are required by the CFFWI system and cannot be changed
    _REQUIRED_UNITS: dict[str, str] = {
        "temperature": "Celsius",
        "precipitation": "mm",
        "relative_humidity": "1",
        "wind_speed": "km/h",
        # Fire weather indices are dimensionless
        "fine_fuel_moisture_code": "1",
        "duff_moisture_code": "1",
        "drought_code": "1",
        "initial_spread_index": "1",
        "build_up_index": "1",
        "fire_weather_index": "1",
        "fire_severity_index": "1",
        # Disambiguated input indices (used by FFMC, DMC, DC, and ISI calculations)
        "input_ffmc": "1",
        "input_dmc": "1",
        "input_dc": "1",
    }

    # Class attributes to be overridden by subclasses
    METADATA_SOURCE_CUBE: str = ""
    INPUT_CUBE_NAMES: list[str] = []
    OUTPUT_CUBE_NAME: str = ""
    REQUIRES_MONTH: bool = False
    # Optional: mapping of input cube names to attribute names (for disambiguation)
    INPUT_ATTRIBUTE_MAPPINGS: dict[str, str] = {}
    # Optional: valid output range for the specific index (min, max)
    # None means no validation for that bound
    VALID_OUTPUT_RANGE: tuple[float | None, float | None] | None = None

    # Valid ranges for input validation (attribute_name: (min, max))
    # None means no validation for that bound
    _VALID_RANGES: dict[str, tuple[float | None, float | None]] = {
        "temperature": (-100.0, 100.0),  # Reasonable temperature range in Celsius
        "precipitation": (0.0, None),  # Must be non-negative
        "relative_humidity": (0.0, 101.0),  # Percentage
        "wind_speed": (0.0, None),  # Must be non-negative
        "input_ffmc": (0.0, 101.0),  # Valid FFMC range
        "input_dmc": (0.0, None),  # DMC is non-negative
        "input_dc": (0.0, None),  # DC is non-negative
        "initial_spread_index": (0.0, 100.0),  # ISI valid range
        "build_up_index": (0.0, 500.0),  # BUI valid range
        "fire_weather_index": (0.0, 100.0),  # FWI valid range
    }



    def load_input_cubes(
        self,
        cubes: tuple[Cube, ...] | CubeList,
        month: int | None = None,
        input_cube_names: list[str] = None,
    ):
        """Loads the required input cubes for the calculation. These are stored
        internally as Cube objects.

        Args:
            cubes:
                Input cubes containing the necessary data.
            month:
                Month of the year (1-12), required only if REQUIRES_MONTH is True.
                Defaults to None.
            input_cube_names:
                A list of input_cube_names if different from self.INPUT_CUBE_NAMES
                Defaults to None.

        Raises:
            ValueError:
                If the number of cubes does not match the expected number, if
                month is required but not provided, or if month is out of range.
        """
        if input_cube_names is None:
            input_cube_names = self.INPUT_CUBE_NAMES

        if len(cubes) != len(input_cube_names):
            raise ValueError(
                f"Expected {len(input_cube_names)} cubes, found {len(cubes)}"
            )

        if self.REQUIRES_MONTH:
            if month is None:
                raise ValueError(
                    f"{self.__class__.__name__} requires a month parameter"
                )
            if not (1 <= month <= 12):
                raise ValueError(f"Month must be between 1 and 12, got {month}")
            self.month = month

        # Load cubes by extracting them using their standard names
        loaded_cubes = tuple(
            cast(Cube, CubeList(cubes).extract_cube(n)) for n in input_cube_names
        )

        # Assign cubes to instance attributes and convert units in a single loop
        for loaded_cube, cube_name in zip(loaded_cubes, input_cube_names):
            attr_name = self._get_attribute_name(cube_name)
            cube = deepcopy(loaded_cube)  # Avoid modifying input cubes in-place
            setattr(self, attr_name, cube)

            # Convert to required units if defined
            if attr_name in self._REQUIRED_UNITS:
                cube.convert_units(self._REQUIRED_UNITS[attr_name])

            # Validate input ranges
            self._validate_input_range(cube, attr_name)




    def _get_attribute_name(self, standard_name: str) -> str:
        """Convert a cube standard name to an attribute name.

        Args:
            standard_name:
                The cube's standard name

        Returns:
            The attribute name to use for storing the cube

        Examples:
            "air_temperature" -> "temperature"
            "lwe_thickness_of_precipitation_amount" -> "precipitation"
            "fine_fuel_moisture_code" -> "input_ffmc" (if INPUT_ATTRIBUTE_MAPPINGS is set)
        """
        # Check class-specific mappings first (for disambiguation)
        if standard_name in self.INPUT_ATTRIBUTE_MAPPINGS:
            return self.INPUT_ATTRIBUTE_MAPPINGS[standard_name]

        # Strip common prefixes and suffixes to create cleaner attribute names
        name = standard_name.removeprefix("lwe_thickness_of_").removeprefix("air_")
        name = name.removesuffix("_amount")

        return name




    def _validate_input_range(self, cube: Cube, attr_name: str) -> None:
        """Validate that input data falls within expected physical ranges.
        If values fall outside the valid range for this input type, then emit a warning.

        Args:
            cube:
                The input cube to validate
            attr_name:
                The attribute name for the cube

        Warns:
            UserWarning:
                If any values fall outside the valid range for this input type

        Raises:
            ValueError:
                If data contains NaN or Inf values
        """
        if attr_name not in self._VALID_RANGES:
            return  # No validation defined for this input type

        min_val, max_val = self._VALID_RANGES[attr_name]
        data = cube.data

        # Check for NaN values
        if np.any(np.isnan(data)):
            raise ValueError(f"{attr_name} contains NaN (Not a Number) values")

        # Check for infinite values
        if np.any(np.isinf(data)):
            raise ValueError(f"{attr_name} contains infinite values")

        # Check minimum bound if defined
        if min_val is not None and np.any(data < min_val):
            actual_min = float(np.min(data))
            warnings.warn(
                f"{attr_name} contains values below valid minimum: "
                f"found {actual_min}, expected >= {min_val}. "
                f"This may indicate unusual conditions or invalid input data.",
                UserWarning,
                stacklevel=3,
            )

        # Check maximum bound if defined
        if max_val is not None and np.any(data > max_val):
            actual_max = float(np.max(data))
            warnings.warn(
                f"{attr_name} contains values above valid maximum: "
                f"found {actual_max}, expected <= {max_val}. "
                f"This may indicate unusual conditions or invalid input data.",
                UserWarning,
                stacklevel=3,
            )




    def _make_output_cube(
        self, data: np.ndarray, template_cube: Cube | None = None
    ) -> Cube:
        """Creates an output cube with specified data and metadata.

        For classes that use precipitation data (FFMC, DMC, DC), automatically
        updates the 'forecast_reference_time' and 'time' coordinates from the
        precipitation cube to reflect the 24-hour accumulation period.

        Args:
            data:
                The output data array
            template_cube:
                The cube to use as a template for metadata.
                If None, uses the first input cube. Defaults to None

        Returns:
            The output cube containing the output data with proper metadata
            and coordinates.
        """
        if template_cube is None:
            # Use first input cube as template
            first_attr = self._get_attribute_name(self.INPUT_CUBE_NAMES[0])
            template_cube = getattr(self, first_attr)

        output_cube = template_cube.copy(data=data.astype(np.float32))
        output_cube.rename(self.OUTPUT_CUBE_NAME)
        output_cube.units = "1"

        # If this class uses precipitation, update time coordinates from precipitation cube
        if hasattr(self, "precipitation"):
            copy = CopyMetadata(aux_coord=["forecast_reference_time", "time"])
            output_cube = copy.process(output_cube, self.precipitation)
            output_cube.coord("forecast_reference_time").bounds = None
            output_cube.coord("time").bounds = None

        return self._set_metadata(output_cube)




    @abstractmethod
    def _calculate(self) -> np.ndarray:
        """Perform the fire weather calculation.

        This method must be implemented by subclasses to perform
        the specific calculation logic for that component.

        Raises:
            NotImplementedError:
                This method must be implemented by subclasses.
        """
        raise NotImplementedError("Subclasses must implement the _calculate method.")




    def process(self, *cubes: Union[Cube, CubeList], month: int | None = None) -> Cube:
        """Calculate the fire weather component.

        Args:
            cubes:
                Input cubes as specified by INPUT_CUBE_NAMES
            month:
                Month parameter (1-12), required only if REQUIRES_MONTH is True
                Defaults to None.

        Returns:
            The calculated output cube.

        Warns:
            UserWarning:
                If output values fall outside typical expected ranges
        """
        cubes = as_cubelist(*cubes)
        self.load_input_cubes(cubes, month)
        output_data = self._calculate()
        output_cube = self._make_output_cube(output_data)

        # Check if output values are within expected ranges
        self._validate_output_range(output_cube)

        return output_cube




    def _validate_output_range(self, output_cube: Cube) -> None:
        """Check if output values fall within expected ranges and issue warnings if not.

        Uses the VALID_OUTPUT_RANGE class attribute if defined.

        Args:
            output_cube:
                The output cube to validate

        Warns:
            UserWarning:
                If output contains NaN, Inf, or values outside expected ranges
        """
        # Check if class-specific VALID_OUTPUT_RANGE is defined
        if self.VALID_OUTPUT_RANGE is None:
            return  # No validation defined for this output type

        min_val, max_val = self.VALID_OUTPUT_RANGE
        data = output_cube.data
        output_name = output_cube.name()

        # Check for NaN values
        if np.any(np.isnan(data)):
            warnings.warn(
                f"{output_name} contains NaN (Not a Number) values. "
                f"This indicates a calculation error or invalid input data.",
                UserWarning,
                stacklevel=3,
            )
            return

        # Check for infinite values
        if np.any(np.isinf(data)):
            warnings.warn(
                f"{output_name} contains infinite values. "
                f"This indicates a calculation error or invalid input data.",
                UserWarning,
                stacklevel=3,
            )
            return

        # Check for values outside expected range
        if min_val is not None and max_val is not None:
            if np.any(data < min_val) or np.any(data > max_val):
                actual_min = float(np.min(data))
                actual_max = float(np.max(data))
                warnings.warn(
                    f"{output_name} contains values outside feasible range "
                    f"[{min_val}, {max_val}]: found range [{actual_min:.2f}, {actual_max:.2f}]. "
                    f"This may indicate unusual conditions or invalid input data.",
                    UserWarning,
                    stacklevel=3,
                )
        elif min_val is not None:
            if np.any(data < min_val):
                actual_min = float(np.min(data))
                warnings.warn(
                    f"{output_name} contains values below feasible minimum "
                    f"{min_val}: found minimum {actual_min:.2f}. "
                    f"This may indicate unusual conditions or invalid input data.",
                    UserWarning,
                    stacklevel=3,
                )
        elif max_val is not None:
            if np.any(data > max_val):
                actual_max = float(np.max(data))
                warnings.warn(
                    f"{output_name} contains values above feasible maximum "
                    f"{max_val}: found maximum {actual_max:.2f}. "
                    f"This may indicate unusual conditions or invalid input data.",
                    UserWarning,
                    stacklevel=3,
                )




    def _set_metadata(self, output_cube: Cube) -> None:
        """
        Add metadata to the output_cube, sourced from either the
        INPUT_ATTRIBUTE_MAPPING of the METADATA_SOURCE_CUBE, the
        METADATA_SOURCE_CUBE attribute itself, or a standard name
        stripped of common prefixes and suffixes.

        Args:
            output_cube:
                The output cube

        Raise:
            NotImplementedError: If METADATA_SOURCE_CUBE is not defined
              or the named cube does not contain the necessary metadata

        Note:
            Metadata is required on all fire weather datasets to monitor the
            build up period of the iterative datasets (Fine Fuel Moisture Code,
            Duff Moisture Code and Drought Code).

            These values are also added to datasets which are not iterative but which
            take iterative datasets as inputs, because the context of when the build up
            period was begun and whether the data is ready is important for stakeholders.
            These downstream datasets include the Initial Spread Index, the Build Up Index,
            the Fire Weather Index and the Fire Severity Index.

        """
        if not self.METADATA_SOURCE_CUBE:
            raise NotImplementedError(
                "A METADATA_SOURCE_CUBE is required for fire weather metadata handling."
            )

        attr_name = self._get_attribute_name(self.METADATA_SOURCE_CUBE)
        try:
            start_date_cube = getattr(self, attr_name)
            start_date = start_date_cube.attributes["iteration_start_date"]
            iteration_count = start_date_cube.attributes["iteration_count"]
            analysis_ready = str(start_date_cube.attributes["analysis_ready"])
        except (AttributeError, KeyError):
            raise NotImplementedError(
                "METADATA_SOURCE_CUBE must match an available input cube with all the "
                "required attributes: `iteration_start_date`, `iteration_count` and `analysis_ready`."
            )

        output_cube.attributes["iteration_start_date"] = start_date
        output_cube.attributes["iteration_count"] = iteration_count
        output_cube.attributes["analysis_ready"] = analysis_ready
        return output_cube






class IterativeFireWeatherBase(FireWeatherBase):
    """
    Iterative abstract base class for Iterative Fire Weather calculations.

    Extends common functionality provided by FireWeatherBase to provide
    iterative functionality for Fire Weather values that take the previous
    days outputs as an input to the current days calculation.

    Subclasses must define class attributes:

    - INPUT_CUBE_NAMES: List of standard names for required input cubes
    - OUTPUT_CUBE_NAME: Standard name for the output cube
    - REQUIRES_MONTH: Boolean indicating if month parameter is required
    - STARTING_VALUE: Integer providing starting value for iterative
        calculations in the absense of input data for the preceding day.
    - LAG_TIME: Integer representing the number of days needed after
        starting calculations from the STARTING_VALUE, before outputs
        should be considered scientifically valid.
    - METADATA_SOURCE_CUBE: The name of the input cube from which the
        metadata will be sourced for the output_cube. For iterative
        fire weather classes this must match the OUTPUT_CUBE_NAME.

    Subclasses must implement:

    - _calculate(): Method that performs the actual calculation

    """

    STARTING_VALUE: int
    LAG_TIME: int
    REFERENCE_CUBE_NAME = "air_temperature"



    def process(
        self,
        *cubes: Union[Cube, CubeList],
        month: int | None = None,
        initialise: bool = False,
    ) -> Cube:
        """
        Args:
            cubes:
                One or more input cubes as specified by INPUT_CUBE_NAMES. When initialise is True cubes should
                exclude the OUTPUT_CUBE_NAME, which should otherwise be given as the iterative input.
            month:
                Month parameter (1-12), required only if REQUIRES_MONTH is True
            initialise:
                True when starting the iterative process else False

        Returns:
            The calculated output cube.

        Warns:
            UserWarning:
                If output values fall outside typical expected ranges

        Raises:
            ValueError: If an output cube is given with initialise=True

        """
        cubes = as_cubelist(*cubes)
        try:
            output_cube = cast(
                Cube, CubeList(cubes).extract_cube(self.OUTPUT_CUBE_NAME)
            )
        except (IndexError, ConstraintMismatchError):
            output_cube = None

        if initialise:
            if output_cube is not None:
                raise ValueError(
                    f"Unexpected output cube '{self.OUTPUT_CUBE_NAME}' supplied when attempting initialisation"
                )
            input_only_cube_names = [
                c for c in self.INPUT_CUBE_NAMES if c not in self.OUTPUT_CUBE_NAME
            ]
            self.load_input_cubes(cubes, month, input_only_cube_names)
            output_cube = self._initialise_baseline_cube(cubes)
            cubes = as_cubelist(*cubes, output_cube)

        output_cube = super().process(cubes, month=month)
        return self._record_lag_time_state(output_cube)




    def _initialise_baseline_cube(self, cubes: tuple[Cube, ...] | CubeList) -> Cube:
        """Create a baseline cube from the reference cube and set iteration_start_date.

        Args:
            cubes:
                Input cubes as specified by INPUT_CUBE_NAMES except OUTPUT_CUBE_NAME

        Raises:
            ValueError: If the REFERENCE_CUBE is not present in cubes
            ValueError: If the REFERENCE_CUBE has a iteration_start_date attribute

        """
        try:
            reference_cube = cast(
                Cube, CubeList(cubes).extract_cube(self.REFERENCE_CUBE_NAME)
            )
        except (IndexError, ConstraintMismatchError) as exc:
            raise ValueError(
                f"Reference cube '{self.REFERENCE_CUBE_NAME}' not found during initialisation process"
            ) from exc

        cube = load_baseline_cube(
            reference_cube,
            self.STARTING_VALUE,
            self.OUTPUT_CUBE_NAME,
            self._REQUIRED_UNITS[self.OUTPUT_CUBE_NAME],
        )
        if "iteration_start_date" in cube.attributes:
            raise ValueError("Unexpected metadata on reference_cube")

        time_coord = self.precipitation.coord("time").copy()
        cube.attributes["iteration_start_date"] = str(
            time_coord.units.num2pydate(time_coord.points)[0]
        )
        cube.attributes["iteration_count"] = 0
        cube.attributes["analysis_ready"] = "False"

        return cube




    def _record_lag_time_state(self, cube: Cube) -> None:
        """Check metadata attributes and warn if LAG_TIME has not been exceeded.

        Args:
            cube:
                The output cube with metadata values to compare to LAG_TIME

        Raises:
            ValueError: If cube has missing metadata attributes

        Warns:
            UserWarning:
                If runtime is less than LAG_TIME

        """
        metadata_in_cube_attributes = [
            "iteration_start_date" in cube.attributes,
            "iteration_count" in cube.attributes,
            "analysis_ready" in cube.attributes,
        ]
        if not all(metadata_in_cube_attributes):
            raise ValueError(
                f"{cube.name()} has missing metadata attributes. To start "
                f"initialise process set `initialise=True`."
            )

        iteration_count = cube.attributes["iteration_count"]
        if iteration_count < self.LAG_TIME:
            cube.attributes["analysis_ready"] = "False"
            warnings.warn(
                f"{cube.name()} is {iteration_count} iterations in "
                f"to its spin-up period of {self.LAG_TIME}.",
                UserWarning,
                stacklevel=3,
            )
        else:
            cube.attributes["analysis_ready"] = "True"

        cube.attributes["iteration_count"] = iteration_count + 1
        return cube