Pamhyr2/src/Model/Scenario.py

# Scenario.py -- Pamhyr
# Copyright (C) 2024-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 -*-

import os
import logging

from tools import logger_exception, timer, flatten

from Model.Tools.PamhyrDB import SQLSubModel

logger = logging.getLogger()


class Scenario(SQLSubModel):
    _id_cnt = 0
    _sub_classes = []

    tables_with_deleted_column = [
        # Adists
        "output_rk_adists",
        "boundary_condition_adists", "boundary_condition_data_adists",
        "lateral_contribution_adists", "lateral_contribution_data_adists",
        "initial_conditions_adists", "initial_conditions_adists_spec",
        "d90_adists", "d90_adists_spec",
        "dif_adists_spec",
        "pollutants", "pollutants_characteristics",
        # Hydraulic
        "additional_files",
        "boundary_condition", "boundary_condition_data",
        "lateral_contribution", "lateral_contribution_data",
        "friction", "stricklers",
        "hydraulic_structures",
        "hydraulic_structures_basic",
        "initial_conditions",
        "sedimentary_layer", "sedimentary_layer_layer",
        "reservoir", "reservoir_data",
        "rep_lines",
        "geometry_pointXYZ", "geometry_profileXYZ",
        "river_reach", "river_node",
        "geotiff", "reservoir", "reservoir_data",
        # Temperature
        "boundary_condition_temperature", "boundary_condition_data_temperature",
        "initial_conditions_temperature", "initial_conditions_temperature_spec",
        "weather_parameters", "weather_parameters_data",
    ]

    related_tables = tables_with_deleted_column + [
        "dif_adists",
        "solver_parameter",
        "hydraulic_structures_basic_value",
        "results", "results_data", "results_add_data",
    ]

    def __init__(self,
                 id: int = -1,
                 name: str = "",
                 description: str = "",
                 x: int = 1000.0, y: int = 1000.0,
                 revision: int = 0,
                 parent=None):
        super(Scenario, self).__init__()

        self._set_id(id)

        self._x, self._y = x, y
        self._name = name
        self._description = description
        self._revision = revision
        self._parent = parent

    def _set_id(self, id):
        if id == -1:
            self._id = Scenario._id_cnt
        else:
            self._id = id

        Scenario._id_cnt = max(
            self._id + 1, Scenario._id_cnt + 1
        )

    @classmethod
    def _db_create(cls, execute):
        execute("""
          CREATE TABLE scenario(
            id INTEGER PRIMARY KEY,
            x REAL NOT NULL DEFAULT 1000,
            y REAL NOT NULL DEFAULT 1000,
            name TEXT NOT NULL,
            description TEXT NOT NULL,
            revision INTEGER NOT NULL,
            parent_id INTEGER REFERENCES scenario(id)
          )
        """)

        cls._create_submodel(execute)
        return True

    @classmethod
    def _db_add_default(cls, execute):
        execute(
            "INSERT OR REPLACE INTO " +
            "scenario(id, x, y, name, description, revision, parent_id) " +
            "VALUES (\n" +
            "  0, 1000, 1000, 'default', 'Default scenario',\n" +
            "  0, NULL\n" +
            ")"
        )

    @classmethod
    def create_db_add_scenario(cls):
        return "scenario INTEGER NOT NULL DEFAULT 0"

    @classmethod
    def create_db_add_scenario_fk(cls):
        return "FOREIGN KEY(scenario) REFERENCES scenario(id)"

    @classmethod
    def _db_update(cls, execute, version, data=None):
        major, minor, release = version.strip().split(".")

        if major == "0" and int(minor) < 2:
            cls._db_create(execute)
            cls._db_add_default(execute)

        return True

    @classmethod
    def update_db_add_scenario(cls, execute, table):
        execute(
            f"ALTER TABLE {table} " +
            "ADD COLUMN scenario INTEGER NOT NULL DEFAULT 0"
        )

        # execute(
        #     f"ALTER TABLE {table} " +
        #     "ADD CONSTRAINT fk_scenario FOREIGN KEY (scenario) " +
        #     "REFERENCES scenario(id)"
        # )

    @classmethod
    def _db_load(cls, execute, data=None):
        scenarios = {}

        table = execute(
            "SELECT id, x, y, name, description, revision, parent_id " +
            "FROM scenario " +
            "ORDER BY id ASC"
        )

        for row in table:
            it = iter(row)

            id = next(it)
            x, y = next(it), next(it)
            name = next(it)
            desc = next(it)
            revi = next(it)
            parent = next(it)

            if parent is not None:
                parent = scenarios[parent]

            new = cls(
                id=id, x=x, y=y,
                name=name, description=desc,
                revision=revi, parent=parent
            )
            scenarios[id] = new

        return scenarios

    def _db_save(self, execute, data=None):
        if self.is_deleted():
            return self.drop_all(execute)

        parent = 'NULL'
        if self.parent is not None:
            parent = self.parent._id

        execute(
            "INSERT OR REPLACE INTO " +
            "scenario(id, x, y, name, description, revision, parent_id) " +
            "VALUES (" +
            f"{self._id}, " +
            f"{self.x}, {self.y}, " +
            f"'{self._db_format(self.name)}', " +
            f"'{self._db_format(self.description)}', " +
            f"{self._revision}, " +
            f"{parent}" +
            ")"
        )

        return True

    def drop_all(self, execute):
        execute(f"DELETE FROM scenario WHERE id = {self.id}")

        tables = self.related_tables

        for table in tables:
            execute(
                f"DELETE FROM {table} " +
                f"WHERE scenario = {self.id}"
            )

        return True

    def get_parent_branch(self):
        def aux(scenario, acc):
            if scenario is None:
                return acc

            return aux(scenario.parent, [scenario.id] + acc)

        return aux(self, [])

    @timer
    def clean_deleted_data(self, execute):
        tables = self.tables_with_deleted_column
        branch = self.get_parent_branch()

        all_ids = []

        for table in tables:
            if self.parent is None:
                # This scenario is the default scenario, so we can
                # delete all data marked as deleted
                execute(
                    f"DELETE FROM {table} " +
                    f"WHERE scenario = {self.id} " +
                    "AND deleted = TRUE"
                )
                continue

            # Select pamhyr_id for each deleted data in this scenario
            # who do not exists into parents scenarios
            ids = execute(
                f"""
                  SELECT pamhyr_id FROM {table}
                  WHERE deleted = TRUE
                  AND scenario = {self.id}
                  AND pamhyr_id NOT IN (
                    SELECT pamhyr_id FROM {table}
                    WHERE scenario IN
                      ({', '.join(map(str, branch[:-1]))})
                  )
                """, fetch_one=False
            )
            if ids is None or len(ids) == 0:
                continue

            ids = flatten(ids)

            logger.debug(
                f"({self.name}) Drop deleted data into '{table}' : {ids}"
            )

            execute(
                f"DELETE FROM {table} " +
                f"WHERE scenario = {self.id} " +
                "AND pamhyr_id IN " +
                f"({', '.join(map(str, ids))})"
            )

            all_ids += ids

        return all_ids

    @property
    def id(self):
        return self._id

    @property
    def x(self):
        return self._x

    @property
    def y(self):
        return self._y

    @property
    def name(self):
        if self._name == "":
            return f"Child of '{self._parent.name}' ({self.id})"

        return self._name

    @name.setter
    def name(self, name):
        self._name = name

    @property
    def description(self):
        if self._description == "":
            return f"Child of '{self._parent.name}' ({self.id})"

        return self._description

    @description.setter
    def description(self, description):
        self._description = description

    @property
    def revision(self):
        return self._revision

    @revision.setter
    def revision(self, revision):
        self._revision = revision

    @property
    def parent(self):
        return self._parent

    def set_pos(self, x, y):
        self._x = x
        self._y = y

    def __setitem__(self, key, value):
        if key == "name":
            self.name = value
        elif key == "description":
            self.description = value

    def __getitem__(self, key):
        if key == "id":
            return self.id
        if key == "name":
            return self.name
        if key == "description":
            return self.description
        if key == "parent":
            return self.parent

        return None

    def workdir(self):
        srep = ""
        if self.id != 0:
            srep = os.path.join(
                self.parent.workdir(),
                "scenario_" + str(self.id)
            )

        return srep