Pamhyr2/src/Model/WeatherParameters/WeatherParametersList.py

# WeatherParametersList.py -- Pamhyr
# Copyright (C) 2023-2025  INRAE
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

# -*- coding: utf-8 -*-

from copy import copy
from tools import trace, timer

from Model.Tools.PamhyrListExt import PamhyrModelList
from Model.Except import NotImplementedMethodeError

from Model.WeatherParameters.WeatherParameters import (
    WeatherParameters,
    WeatherParametersDefault,
)
from Model.WeatherParameters.WeatherParametersTypes import (
    NotDefined, WeatherParam, AirTemperature, SpecificHumidity,
    GlobalRadiation, ReferenceWindSpeed, GroundwaterTemperature,
    GroundwaterFlowRate, Albedo, ShadingCoefficient, CloudCoverFraction
)


class WeatherParametersList(PamhyrModelList):
    _sub_classes = [
        WeatherParameters,
    ]
    _default_types = [
        "AT", "SH", "GR", "RWS", "GT", "GFR", "ALB", "SC", "CCF"
    ]
    _default_values = {
        "AT": 20.0,
        "SH": 0.001,
        "GR": 10.0,
        "RWS": 1.0,
        "GT": 0.0,
        "GFR": 0.0,
        "ALB": 0.06,
        "SC": 0.15,
        "CCF": 0.5,
    }

    def __init__(self, status=None):
        super(WeatherParametersList, self).__init__(status=status)
        self._default = []
        self._init_missing_default_values()

    @classmethod
    def _db_load(cls, execute, data=None):
        new = cls(status=data['status'])

        if data is None:
            data = {}

        new._lst = WeatherParameters._db_load(
            execute, data
        )

        data['loaded_pid'] = set()
        new._default = WeatherParametersDefault._db_load(
            execute, data
        )
        new._init_missing_default_values()
        new._sort_default_values()

        return new

    def _db_save(self, execute, data=None):
        execute(
            "DELETE FROM weather_parameters " +
            f"WHERE scenario = {self._status.scenario_id}"
        )
        execute(
            "DELETE FROM weather_parameters_data " +
            f"WHERE scenario = {self._status.scenario_id}"
        )
        execute(
            "DELETE FROM weather_parameters_default " +
            f"WHERE scenario = {self._status.scenario_id}"
        )

        if data is None:
            data = {}

        for wp in self._lst:
            wp._db_save(execute, data=data)

        for wp_default in self._default:
            wp_default._db_save(execute, data=data)

        return True

    def _data_traversal(self,
                        predicate=lambda obj, data: True,
                        modifier=lambda obj, data: None,
                        data={}):
        if predicate(self, data):
            modifier(self, data)

        for el in self._lst:
            el._data_traversal(predicate, modifier, data)

        for el in self._default:
            el._data_traversal(predicate, modifier, data)

    def _init_missing_default_values(self):
        existing = [wp_default.type for wp_default in self.default]
        for type in self._default_types:
            if type not in existing:
                self._default.append(
                    WeatherParametersDefault(
                        type=type,
                        value=self._default_values[type],
                        status=self._status,
                    )
                )

        self._sort_default_values()

    def _sort_default_values(self):
        self._default.sort(
            key=lambda wp_default: self._default_types.index(wp_default.type)
            if wp_default.type in self._default_types
            else len(self._default_types)
        )

    def new(self, index, type="ND"):
        if type == "AT":
            n = AirTemperature(status=self._status)
        elif type == "SH":
            n = SpecificHumidity(status=self._status)
        elif type == "GR":
            n = GlobalRadiation(status=self._status)
        elif type == "RWS":
            n = ReferenceWindSpeed(status=self._status)
        elif type == "GT":
            n = GroundwaterTemperature(status=self._status)
        elif type == "GFR":
            n = GroundwaterFlowRate(status=self._status)
        elif type == "ALB":
            n = Albedo(status=self._status)
        elif type == "SC":
            n = ShadingCoefficient(status=self._status)
        elif type == "CCF":
            n = CloudCoverFraction(status=self._status)
        else:
            n = NotDefined(status=self._status)
        self._lst.insert(index, n)
        self._status.modified()
        return n

    def split_reach(self, reach, profile, reach1, reach2):
        profiles = reach.reach.profiles
        split_index = profiles.index(profile)
        parts = [
            (
                reach1,
                0,
                split_index,
                dict(zip(
                    profiles[:split_index + 1],
                    reach1.reach.profiles
                ))
            ),
            (
                reach2,
                split_index,
                len(profiles) - 1,
                dict(zip(
                    profiles[split_index:],
                    reach2.reach.profiles
                ))
            ),
        ]
        weather_parameters = [
            weather_parameter
            for weather_parameter in self.lst
            if weather_parameter.reach is reach
        ]
        clones = []

        for weather_parameter in weather_parameters:
            if (weather_parameter.begin_section not in profiles or
                    weather_parameter.end_section not in profiles):
                continue

            begin = profiles.index(weather_parameter.begin_section)
            end = profiles.index(weather_parameter.end_section)
            lower, upper = sorted((begin, end))
            reverse = begin > end

            for new_reach, part_lower, part_upper, sections in parts:
                clipped_lower = max(lower, part_lower)
                clipped_upper = min(upper, part_upper)
                if clipped_lower >= clipped_upper:
                    continue

                begin_index = clipped_upper if reverse else clipped_lower
                end_index = clipped_lower if reverse else clipped_upper
                clones.append(weather_parameter.cloned_for(
                    new_reach,
                    sections[profiles[begin_index]],
                    sections[profiles[end_index]]
                ))

        if clones:
            self._lst.extend(clones)
            self._status.modified()

    @property
    def Weather_Parameters_List(self):
        return self.lst

    @property
    def default(self):
        return self._default.copy()

    def default_for_type(self, type):
        for wp_default in self.default:
            if wp_default.type == type:
                return wp_default
        return None