River.py

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

from Model.DB import SQLSubModel

from Model.Network.Node import Node
from Model.Network.Edge import Edge
from Model.Network.Graph import Graph

from Model.Geometry.Profile import Profile
from Model.Geometry.Reach import Reach

from Model.BoundaryCondition.BoundaryConditionList import BoundaryConditionList
from Model.LateralContribution.LateralContributionList import LateralContributionList
from Model.InitialConditions.InitialConditionsDict import InitialConditionsDict
from Model.Stricklers.StricklersList import StricklersList
from Model.Section.SectionList import SectionList
from Model.SolverParameters.SolverParametersList import SolverParametersList

from Solver.Solvers import solver_type_list

class RiverNode(Node, SQLSubModel):
    _sub_classes = []

    def __init__(self, id:str, name:str,
                 x:float, y:float,
                 status = None):
        super(RiverNode, self).__init__(
            id, name,
            x, y,
            status = status
        )

        self._locker = None

    @classmethod
    def _sql_create(cls, execute):
        execute("""
          CREATE TABLE river_node(
            id INTEGER NOT NULL PRIMARY KEY,
            name TEXT NOT NULL,
            x REAL NOT NULL,
            y REAL NOT NULL
          )
        """)

        cls._create_submodel(execute)
        return True

    @classmethod
    def _sql_update(cls, execute, version):
        return True

    @classmethod
    def _sql_load(cls, execute, data = None):
        nodes = []

        table = execute("SELECT id, name, x, y FROM river_node")
        for row in table:
            # Update id counter
            cls._id_cnt = max(cls._id_cnt, row[0])
            # Create new node
            nodes.append(cls(*row, **data))

        return nodes

    def _sql_save(self, execute, data = None):
        sql = (
            "INSERT OR REPLACE INTO river_node(id, name, x, y) VALUES (" +
            f"{self.id}, '{self._sql_format(self.name)}', " +
            f"{self.x}, {self.y}"+
            ")"
        )
        execute(sql)

        return True

    @property
    def locker(self):
        return self._locker

    @locker.setter
    def locker(self, locker):
        self._locker = locker


class RiverReach(Edge, SQLSubModel):
    _sub_classes = [
        Reach,
        # SectionList,
    ]

    def __init__(self, id:str, name:str,
                 node1:RiverNode = None,
                 node2:RiverNode = None,
                 status = None):
        super(RiverReach, self).__init__(
            id, name,
            node1, node2,
            status = status
        )

        self._reach = Reach(status=self._status, parent=self)
        self._sections = SectionList(status=self._status)

    @classmethod
    def _sql_create(cls, execute):
        execute("""
          CREATE TABLE river_reach(
            id INTEGER NOT NULL PRIMARY KEY,
            name TEXT NOT NULL,
            enable BOOLEAN NOT NULL,
            node1 INTEGER,
            node2 INTEGER,
            FOREIGN KEY(node1) REFERENCES river_node(id),
            FOREIGN KEY(node2) REFERENCES river_node(id)
          )
        """)

        cls._create_submodel(execute)
        return True

    @classmethod
    def _sql_update(cls, execute, version):
        return True

    @classmethod
    def _sql_load(cls, execute, data = None):
        reachs = []

        if data is None:
            data = {}

        table = execute("SELECT id, name, enable, node1, node2 FROM river_reach")
        for row in table:
            # Update id counter
            cls._id_cnt = max(cls._id_cnt, row[0])
            # Create new reach
            id = row[0]
            name = row[1]
            enable = (row[2] == 1)
            # Get nodes corresponding to db foreign key id
            node1 = next(filter(lambda n: n.id == row[3], data["nodes"]))
            node2 = next(filter(lambda n: n.id == row[4], data["nodes"]))

            new = cls(id, name, node1, node2, status = data["status"])
            new.enable(enable = enable)

            data["reach"] = id
            data["parent"] = new
            new._reach = Reach._sql_load(execute, data)

            reachs.append(new)

        return reachs

    def _sql_save(self, execute, data = None):
        sql = (
            "INSERT OR REPLACE INTO " +
            "river_reach(id, name, enable, node1, node2) "+
            "VALUES (" +
            f"{self.id}, '{self._sql_format(self._name)}', " +
            f"{self._sql_format(self.is_enable())},"
            f"{self.node1.id}, {self.node2.id}"+
            ")"
        )
        execute(sql)

        objs = [self._reach]
        return self._save_submodel(execute, objs, data)

    @property
    def reach(self):
        return self._reach

    @property
    def sections(self):
        return self._sections

class River(Graph, SQLSubModel):
    _sub_classes = [
        RiverNode,
        RiverReach,
        BoundaryConditionList,
        LateralContributionList,
        InitialConditionsDict,
        StricklersList,
        SolverParametersList,
    ]

    def __init__(self, status=None):
        super(River, self).__init__(status=status)

        # Replace Node and Edge ctor by custom ctor
        self._node_ctor = RiverNode
        self._edge_ctor = RiverReach

        self._current_reach = None
        self._boundary_condition = BoundaryConditionList(status=self._status)
        self._lateral_contribution = LateralContributionList(status=self._status)
        self._initial_conditions = InitialConditionsDict(status=self._status)
        self._stricklers = StricklersList(status=self._status)
        self._parameters = {}

    @classmethod
    def _sql_create(cls, execute):
        cls._create_submodel(execute)
        return True

    @classmethod
    def _sql_update(cls, execute, version):
        cls._update_submodel(execute, version)
        return True

    @classmethod
    def _sql_load(cls, execute, data = None):
        # Network
        new = cls(data["status"])
        new._nodes = RiverNode._sql_load(
            execute,
            data
        )
        data["nodes"] = new.nodes()

        new._edges = RiverReach._sql_load(
            execute,
            data
        )
        data["edges"] = new.edges()

        # Boundary Condition
        new._boundary_condition = BoundaryConditionList._sql_load(
            execute,
            data
        )
        # Lateral Contribution
        new._lateral_contribution = LateralContributionList._sql_load(
            execute,
            data
        )

        # Initial conditions
        new._initial_conditions = InitialConditionsDict._sql_load(
            execute,
            data
        )

        # Stricklers
        new._stricklers = StricklersList._sql_load(
            execute,
            data
        )

        # Parameters
        for solver in new._parameters:
            data["solver"] = solver
            new._parameters[solver] = SolverParametersList(
                execute,
                data
            )

        return new

    def _sql_save(self, execute, data = None):
        objs = (self._nodes + self._edges)
        objs.append(self._boundary_condition)
        objs.append(self._initial_conditions)
        objs.append(self._lateral_contribution)
        objs.append(self._stricklers)

        for solver in self._parameters:
            objs.append(self._parameters[solver])

        self._save_submodel(execute, objs, data)
        return True

    @property
    def reach(self):
        return self._reach

    @property
    def sections(self):
        return self._sections

    @property
    def boundary_condition(self):
        return self._boundary_condition

    @property
    def lateral_contribution(self):
        return self._lateral_contribution

    @property
    def initial_conditions(self):
        return self._initial_conditions

    @property
    def stricklers(self):
        return self._stricklers

    def strickler(self, name):
        ret = list(
            filter(
                lambda s: s.name == name or str(s) == name,
                self._stricklers.stricklers
            )
        )

        if len(ret) == 0:
            return None

        return ret[0]


    @property
    def parameters(self):
        return self._parameters

    def get_params(self, solver):
        if solver in self._parameters:
            return self._parameters[solver]

        new = SolverParametersList(
            solver_type = solver_type_list[solver],
            status = self._status
        )
        self._parameters[solver] = new
        self._status.modified()
        return self._parameters[solver]

    def has_current_reach(self):
        return self._current_reach is not None

    def current_reach(self):
        return self._current_reach

    def set_current_reach(self, reach):
        self._current_reach = reach