# River.py -- Pamhyr river model
# Copyright (C) 2023  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 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.Friction.FrictionList import FrictionList
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, status=data["status"]))

        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,
        FrictionList,
    ]

    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._frictions = FrictionList(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 cls._update_submodel(execute, version)

    @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)

            new._frictions = FrictionList._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)

        if data is None:
            data = {}

        data["reach"] = self

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

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

    @property
    def frictions(self):
        return self._frictions

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):
        new = cls(data["status"])

        # Stricklers (Stricklers is load in first because it's needed
        # for reachs)
        new._stricklers = StricklersList._sql_load(
            execute,
            data
        )
        data["stricklers"] = new._stricklers

        # Network
        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
        )

        # Parameters
        new._parameters = SolverParametersList._sql_load(
            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 frictions(self):
        return self._frictions

    @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):
        if self.enable_edges_counts() == 1:
            return True

        return self._current_reach is not None

    def current_reach(self):
        ee = self.enable_edges()
        if len(ee) == 1:
            return ee[0]

        return self._current_reach

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