# Copyright 2021 Prediktor AS
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import List, Union
import numpy as np
import pandas as pd
from scipy.sparse import coo_matrix, eye
from .value_parser import parse_nodeid


def supertypes_of_nodes(
    types_to_supertype: List[str],
    type_nodes: pd.DataFrame,
    type_references: pd.DataFrame,
) -> pd.DataFrame:
    # Computes the closure
    closure = typing_transitive_reflexive(type_nodes, type_references)
    closure = closure.set_index("Trg")
    typeindex = pd.Index(types_to_supertype)
    closure = closure[closure.index.isin(typeindex)].reset_index()
    closure = closure.rename(columns={"Src": "supertype", "Trg": "type"})

    return closure


def hierarchical_references_trg_has_modelling_rule(
    references: pd.DataFrame, type_references: pd.DataFrame, type_nodes: pd.DataFrame
):
    hierarchical_refs = hierarchical_references(
        inst_references=references,
        type_references=type_references,
        type_nodes=type_nodes,
    )
    has_modelling_rule = has_modelling_rule_references(
        inst_references=references,
        type_references=type_references,
        type_nodes=type_nodes,
    )[["Src"]].rename(columns={"Src": "id"})
    hierarchical_refs_has_modelling_rule = hierarchical_refs.set_index("Trg").join(
        has_modelling_rule.set_index("id"), how="inner"
    )
    hierarchical_refs_has_modelling_rule = (
        hierarchical_refs_has_modelling_rule.reset_index().rename(
            columns={"index": "Trg"}
        )
    )
    return hierarchical_refs_has_modelling_rule[["Src", "Trg", "ReferenceType"]].copy()


# Useful for keeping forward refs for methods when instantiating
def hierarchical_references_trg_has_no_modelling_rule(
    references: pd.DataFrame, type_references: pd.DataFrame, type_nodes: pd.DataFrame
):
    hierarchical_refs = hierarchical_references(
        inst_references=references,
        type_references=type_references,
        type_nodes=type_nodes,
    )
    has_modelling_rule = has_modelling_rule_references(
        inst_references=references,
        type_references=type_references,
        type_nodes=type_nodes,
    )[["Src"]].rename(columns={"Src": "id"})
    has_modelling_rule = has_modelling_rule.set_index("id")
    hierarchical_refs_has_no_modelling_rule = hierarchical_refs.set_index(
        "Trg", drop=False
    )
    trg_no_modelling_rule = ~hierarchical_refs_has_no_modelling_rule.index.isin(
        has_modelling_rule.index
    )
    hierarchical_refs_has_no_modelling_rule = hierarchical_refs_has_no_modelling_rule[
        trg_no_modelling_rule
    ].copy()
    return hierarchical_refs_has_no_modelling_rule[
        ["Src", "Trg", "ReferenceType"]
    ].copy()


def find_relatives(
    nodes: pd.DataFrame,
    nodes_key_col: str,
    edges: pd.DataFrame,
    relative_type: str,
    cutoff: int = None,
    keep_paths: bool = False,
) -> pd.DataFrame:
    """
    Find the relatives to the provided nodes

    The

    Parameters:
    nodes (pd.DataFrame): The dataframe with nodes to find relatives from.
    nodes_key_col (str): The key column to use for indexing in the nodes DataFrame
    edges(pd.DataFrame): The edges to follow in the traverse
    cutoff (int): Max number of steps to go for the traverse
    keep_paths(bool):

    Returns:


    """
    new_nodes = nodes.copy()
    orig_cols = nodes.columns.values.tolist()
    new_nodes[0] = nodes[nodes_key_col]
    new_nodes["len_path"] = 0
    new_nodes = new_nodes.set_index(0, drop=False)
    path_cols = [0]
    if relative_type.lower().startswith("a"):
        col_src = "Trg"
        col_trg = "Src"
    else:
        col_src = "Src"
        col_trg = "Trg"
    # Set the index on reference column based on relativetype(direction)
    edge_join = edges.set_index(col_src)
    result = [new_nodes]
    i = 1
    while (new_nodes.shape[0] > 0) and ((cutoff is None) or (cutoff >= i)):
        joined_nodes = new_nodes.join(edge_join, how="inner")
        #
        new_nodes = joined_nodes[orig_cols + [col_trg] + path_cols].copy()
        if len(new_nodes) > 0:
            if keep_paths:
                new_nodes[i] = new_nodes[col_trg]
                path_cols.append(i)
            else:
                new_nodes[0] = new_nodes[col_trg]
            new_nodes["len_path"] = i
            new_nodes = new_nodes.set_index(col_trg)
            result.append(new_nodes)
        i = i + 1
    if len(result) > 1:
        resconc = pd.concat(result).reset_index()
    else:
        resconc = result[0]

    resconc["end"] = resconc[0]
    if keep_paths:
        for i in range(1, i - 1):
            resconc["end"] = resconc[i].combine_first(resconc["end"])

    return resconc


def resolve_id_from_nodeid(nodes: pd.DataFrame, nodeid: str):
    nodeid_instance = parse_nodeid(nodeid)
    nodes = nodes.set_index("NodeId")
    theid = nodes.loc[nodeid_instance, "id"]
    return theid


def resolve_ids_from_browsenames(nodes: pd.DataFrame, browsenames: List[str]):
    nodes = nodes.set_index("BrowseName")
    hst_id = nodes.loc[browsenames, "id"].reset_index(drop=True)
    return hst_id


def has_subtype_references(
    references: pd.DataFrame, type_nodes: pd.DataFrame, id_col: str
) -> pd.DataFrame:
    hst_id = type_nodes.loc[
        (type_nodes["NodeClass"] == "UAReferenceType")
        & (type_nodes["BrowseName"] == "HasSubtype"),
        id_col,
    ].iloc[0]
    subtype = references[references["ReferenceType"] == hst_id].copy()
    return subtype


def has_type_definition_references(
    references: pd.DataFrame, type_nodes: pd.DataFrame
) -> pd.DataFrame:
    htd_id = type_nodes.loc[
        (type_nodes["NodeClass"] == "UAReferenceType")
        & (type_nodes["BrowseName"] == "HasTypeDefinition"),
        "id",
    ].iloc[0]
    htd = references[references["ReferenceType"] == htd_id].copy()
    return htd


def has_property_references(
    inst_references: pd.DataFrame,
    type_references: pd.DataFrame,
    type_nodes: pd.DataFrame,
) -> pd.DataFrame:
    propref_id = type_nodes.loc[
        (type_nodes["NodeClass"] == "UAReferenceType")
        & (type_nodes["BrowseName"] == "HasProperty"),
        "id",
    ].iloc[0]
    prop_ref = constrain_to_reference_type(
        inst_references, type_nodes, type_references, [propref_id]
    )
    return prop_ref


def has_modelling_rule_references(
    inst_references: pd.DataFrame,
    type_references: pd.DataFrame,
    type_nodes: pd.DataFrame,
) -> pd.DataFrame:
    propref_id = type_nodes.loc[
        (type_nodes["NodeClass"] == "UAReferenceType")
        & (type_nodes["BrowseName"] == "HasModellingRule"),
        "id",
    ].iloc[0]
    prop_ref = constrain_to_reference_type(
        inst_references, type_nodes, type_references, [propref_id]
    )
    return prop_ref


def hierarchical_references(
    inst_references: pd.DataFrame,
    type_references: pd.DataFrame,
    type_nodes: pd.DataFrame,
) -> pd.DataFrame:
    hierref_id = type_nodes.loc[
        (type_nodes["NodeClass"] == "UAReferenceType")
        & (type_nodes["BrowseName"] == "HierarchicalReferences"),
        "id",
    ].iloc[0]
    prop_ref = constrain_to_reference_type(
        inst_references, type_nodes, type_references, [hierref_id]
    )
    return prop_ref


def fast_transitive_closure(references: pd.DataFrame) -> pd.DataFrame:
    # Quickly and memory-efficiently compute the transitive closure of a reference
    # Transitive closure: https://en.wikipedia.org/wiki/Transitive_closure
    # We do not consider ReferenceType or store underlying paths in this closure
    # The references are converted to a matrix where each unique Src/Trg is assigned a column and row
    # We have row(Src), col(Trg) == 1 if there is a reference from Src to Trg, 0 otherwise
    # A sparse matrix representation is used in order to reduce memory consumption.
    # The transitive closure is given by iterative dot products of this matrix with itself until we reach a fixed point
    # At the end of each iteration, a cell is converted to 1 if it is above 0 or 0 otherwise.
    # This is in order to ensure that the algorithm terminates
    assert (
        references["Src"] == references["Trg"]
    ).sum() == 0, "There should be no references r such that r(n,n)"
    codes, uniques = pd.factorize(pd.concat([references["Src"], references["Trg"]]))
    src_codes = codes[0 : references.shape[0]]
    trg_codes = codes[references.shape[0] :]
    data = np.ones(references.shape[0])
    sparmat = coo_matrix(
        (data, (src_codes, trg_codes)),
        shape=(uniques.shape[0], uniques.shape[0]),
        dtype=bool,
    )
    sparmat = sparmat + eye(uniques.shape[0])
    fixedp = False
    while not fixedp:
        before = sparmat
        sparmat = sparmat.dot(sparmat) > 0
        if before.sum(axis=None) == sparmat.sum(axis=None):
            fixedp = True

    sparmat = sparmat - eye(uniques.shape[0])
    sparmat = sparmat.tocoo()
    src = sparmat.row
    trg = sparmat.col
    src_nids = uniques.take(src)
    trg_nids = uniques.take(trg)
    hierarchy = pd.DataFrame({"Src": src_nids, "Trg": trg_nids})
    return hierarchy


def fast_hierarchy_transitive_closure(
    inst_references: pd.DataFrame,
    constrained_references: pd.DataFrame,
    type_nodes: pd.DataFrame,
    type_references: pd.DataFrame,
    interesting_types: pd.Series,
) -> pd.DataFrame:
    # Computes a navigation hierarchy between instances of interesting_types
    # The navigation hierarchy is first computed between all types (given references in inst_references)
    # Then, it is restricted to only concern instances of interesting_types
    constrained_references = constrained_references[["Src", "Trg", "ReferenceType"]]
    hierarchy = fast_transitive_closure(constrained_references)

    types = type_nodes[type_nodes["NodeClass"].str.endswith("Type")].set_index(
        "BrowseName"
    )
    interesting_ids_types = types.loc[pd.Index(interesting_types), "id"].to_list()

    sub_inter_types = subtypes_of_nodes(
        interesting_ids_types, type_nodes, type_references
    )
    inter_subtypes_index = pd.Index(sub_inter_types["subtype"])
    htd = has_type_definition_references(inst_references, type_nodes).set_index("Trg")
    htd = htd[htd.index.isin(inter_subtypes_index)]
    inst_inter_type_index = pd.Index(htd["Src"])

    hierarchy = hierarchy.set_index("Src")
    hierarchy = hierarchy[hierarchy.index.isin(inst_inter_type_index)].reset_index()
    hierarchy = hierarchy.set_index("Trg")
    hierarchy = hierarchy[hierarchy.index.isin(inst_inter_type_index)].reset_index()
    return hierarchy


def signal_variables(
    instance_references: pd.DataFrame,
    type_nodes: pd.DataFrame,
    type_references: pd.DataFrame,
) -> pd.DataFrame:
    reftypes = type_nodes[type_nodes["NodeClass"] == "UAVariableType"]
    signalvar_type = reftypes.loc[
        reftypes["BrowseName"] == "BaseDataVariableType", "id"
    ].iloc[0]
    signalvar_types = subtypes_of_nodes([signalvar_type], type_nodes, type_references)
    signalvar_types_index = pd.Index(signalvar_types["subtype"])
    inst_htd = has_type_definition_references(instance_references, type_nodes)
    inst_htd = inst_htd.set_index("Trg")
    signal_vars = inst_htd[inst_htd.index.isin(signalvar_types_index)]["Src"]

    return signal_vars


def property_variables(
    instance_references: pd.DataFrame,
    type_nodes: pd.DataFrame,
    type_references: pd.DataFrame,
) -> pd.DataFrame:
    reftypes = type_nodes[type_nodes["NodeClass"] == "UAVariableType"]
    propvar_type = reftypes.loc[reftypes["BrowseName"] == "PropertyType", "id"].iloc[0]
    propvar_types = subtypes_of_nodes([propvar_type], type_nodes, type_references)
    propvar_types_index = pd.Index(propvar_types["subtype"])
    inst_htd = has_type_definition_references(instance_references, type_nodes)
    inst_htd = inst_htd.set_index("Trg")
    prop_vars = inst_htd[inst_htd.index.isin(propvar_types_index)]["Src"]

    return prop_vars


def typing_transitive_reflexive(
    type_nodes: pd.DataFrame, type_references: pd.DataFrame
) -> pd.DataFrame:
    subtype = has_subtype_references(type_references, type_nodes, "id")[["Src", "Trg"]]
    # Add transitive closure
    subtype_closure = fast_transitive_closure(subtype)
    unique_types = pd.concat([type_references["Src"], type_references["Trg"]]).unique()
    # Add reflexive closure
    subtype_closure = pd.concat(
        [subtype_closure, pd.DataFrame({"Src": unique_types, "Trg": unique_types})]
    )
    return subtype_closure


def subtypes_of_nodes(
    types_to_subtype: Union[pd.Series, List[str]],
    type_nodes: pd.DataFrame,
    type_references: pd.DataFrame,
) -> pd.DataFrame:
    # Computes the closure
    closure = typing_transitive_reflexive(type_nodes, type_references)
    closure = closure.set_index("Src")
    typeindex = pd.Index(types_to_subtype)
    closure = closure[closure.index.isin(typeindex)].reset_index()
    closure = closure.rename(columns={"Trg": "subtype", "Src": "type"})

    return closure


def closest_parent(
    inst_nodes: pd.DataFrame,
    inst_references: pd.DataFrame,
    type_nodes: pd.DataFrame,
    type_references: pd.DataFrame,
    interesting_types: List[str],
    uavar_type: str,
) -> pd.DataFrame:
    sub_inter_types = subtypes_of_nodes(interesting_types, type_nodes, type_references)
    inter_subtypes_index = pd.Index(sub_inter_types["subtype"])
    htd = has_type_definition_references(inst_references, type_nodes).set_index("Trg")
    htd = htd[htd.index.isin(inter_subtypes_index)]
    inst_inter_type_index = pd.Index(htd["Src"])

    # Remove incoming edges to objects of desired type
    inst_references = inst_references.set_index("Trg")
    inst_references = inst_references[
        ~inst_references.index.isin(inst_inter_type_index)
    ].reset_index()

    signals = signal_variables(inst_references, type_nodes, type_references)
    signals_index = pd.Index(signals)

    if uavar_type == "Property":
        inst_references = inst_references.set_index("Trg")
        inst_references = inst_references[
            ~inst_references.index.isin(signals_index)
        ].reset_index()

    inst_references_trans = fast_transitive_closure(inst_references[["Src", "Trg"]])
    inst_nodes_ns_index = inst_nodes[["id", "ns"]].set_index("id")
    inst_references_trans = (
        inst_references_trans.set_index("Src")
        .join(inst_nodes_ns_index)
        .reset_index()
        .rename(columns={"index": "Src", "ns": "Src_ns"})
    )
    inst_references_trans = (
        inst_references_trans.set_index("Trg")
        .join(inst_nodes_ns_index)
        .reset_index()
        .rename(columns={"index": "Trg", "ns": "Trg_ns"})
    )

    if uavar_type == "Property":
        properties = property_variables(inst_references, type_nodes, type_references)
        inst_references_trans = inst_references_trans.set_index("Trg")
        inst_references_trans = inst_references_trans[
            inst_references_trans.index.isin(pd.Index(properties))
        ]
        inst_references_trans = inst_references_trans.reset_index().set_index("Src")
        inst_references_trans = inst_references_trans[
            (
                inst_references_trans.index.isin(signals_index)
                | inst_references_trans.index.isin(inst_inter_type_index)
            )
        ]
        inst_references_trans["is_signal"] = inst_references_trans.index.isin(
            signals_index
        )
        inst_references_trans = inst_references_trans.reset_index()
        inst_references_trans = inst_references_trans.sort_values(
            by=["Trg", "is_signal"], ascending=[True, False]
        )
        inst_references_trans = inst_references_trans.drop_duplicates(keep="first")
        inst_references_trans = inst_references_trans.drop(columns="is_signal")
        return inst_references_trans.rename(
            columns={
                "Src": "instance",
                "Trg": "property",
                "Src_ns": "instance_ns",
                "Trg_ns": "property_ns",
            }
        )

    if uavar_type == "Signal":
        inst_references_trans = inst_references_trans.set_index("Trg")
        inst_references_trans = inst_references_trans[
            inst_references_trans.index.isin(signals_index)
        ]
        inst_references_trans = inst_references_trans.reset_index().set_index("Src")
        inst_references_trans = inst_references_trans[
            inst_references_trans.index.isin(inst_inter_type_index)
        ].reset_index()
        inst_references_trans = inst_references_trans.drop_duplicates(keep="first")
        return inst_references_trans.rename(
            columns={
                "Src": "instance",
                "Trg": "variable",
                "Src_ns": "instance_ns",
                "Trg_ns": "variable_ns",
            }
        )


def constrain_to_reference_type(
    inst_references: pd.DataFrame,
    type_nodes: pd.DataFrame,
    type_references: pd.DataFrame,
    reference_types: List[str],
) -> pd.DataFrame:
    subtypes_of_reference = subtypes_of_nodes(
        reference_types, type_nodes, type_references
    )
    inst_references = inst_references.set_index("ReferenceType")
    inst_references = inst_references[
        inst_references.index.isin(pd.Index(subtypes_of_reference["subtype"]))
    ]
    inst_references = inst_references.reset_index()
    return inst_references
