"""Analyse causal pathways, cycles and complementary feedback."""

import pandas as pd
import networkx as nx
import sympy as sp
from functools import cache
from typing import Optional, Literal
from ..core.structure import create_matrix
from ..core.stability import system_feedback, net_feedback, absolute_feedback, weighted_feedback
from ..core.helper import get_nodes, get_positive, get_negative, get_weight, _check_direct_io_edges, _sign_string, _arrows

@cache
def get_cycles(G: nx.DiGraph) -> sp.Matrix:
    """Find all feedback cycles in the signed digraph.

    Args:
        G: NetworkX DiGraph representing signed digraph model

    Returns:
        sp.Matrix: Products of interactions along each cycle

    References:
        - Mason, S.J. (1953). Feedback Theory-Some Properties of Signal Flow Graphs. Proceedings of the IRE 41, 1144–1156.
        - Levins, R. (1974). The qualitative analysis of partially specified systems. Annals of the New York Academy of Sciences 231, 123–138.
        - Puccia, C.J., Levins, R. (1985). Qualitative Modeling of Complex Systems: An Introduction to Loop Analysis and Time Averaging. Harvard University Press.

    Examples:
        ```python
        from qmm import get_cycles, load_digraph
        get_cycles(load_digraph("snowshoe_rp"))
        # Matrix([
        # [           -a_P,P],
        # [           -a_R,R],
        # [     -a_C,P*a_P,C],
        # [     -a_C,R*a_R,C],
        # [a_C,P*a_P,R*a_R,C]])
        ```
    """
    A = create_matrix(G, form="symbolic")
    nodes = get_nodes(G, "state")
    node_id = {n: i for i, n in enumerate(nodes)}
    cycle_nodes = sorted(
        [c[c.index(min(c)):] + c[:c.index(min(c))] for c in nx.simple_cycles(G)],
        key=lambda x: (len(x), x)
    )
    C = [c + [c[0]] for c in cycle_nodes]
    cycles = sp.Matrix([sp.prod([A[node_id[c[i + 1]], node_id[c[i]]] for i in range(len(c) - 1)]) for c in C])
    return cycles

@cache
def cycles_table(G: nx.DiGraph) -> pd.DataFrame:
    """Find all feedback cycles in the signed digraph.

    Args:
        G: NetworkX DiGraph representing signed digraph model

    Returns:
        pd.DataFrame: Table with cycle length, path representation, and sign

    References:
        - Mason, S.J. (1953). Feedback Theory-Some Properties of Signal Flow Graphs. Proceedings of the IRE 41, 1144–1156.
        - Levins, R. (1974). The qualitative analysis of partially specified systems. Annals of the New York Academy of Sciences 231, 123–138.
        - Puccia, C.J., Levins, R. (1985). Qualitative Modeling of Complex Systems: An Introduction to Loop Analysis and Time Averaging. Harvard University Press.

    Examples:
        ```python
        from qmm import cycles_table, load_digraph
        cycles_table(load_digraph("snowshoe_rp"))
        #    Length                                          Cycle Sign
        # 0       1                                P $\\multimap$ P    −
        # 1       1                                R $\\multimap$ R    −
        # 2       2                C $\\rightarrow$ P $\\multimap$ C    −
        # 3       2                C $\\multimap$ R $\\rightarrow$ C    −
        # 4       3  C $\\multimap$ R $\\rightarrow$ P $\\multimap$ C    +
        ```
    """
    cycle_nodes = sorted(
        [c[c.index(min(c)):] + c[:c.index(min(c))] for c in nx.simple_cycles(G)],
        key=lambda x: (len(x), x)
    )
    all_cycles = [cycle + [cycle[0]] for cycle in cycle_nodes]
    cycle_signs = [_sign_string(G, path) for path in all_cycles]
    cycles_df = pd.DataFrame(
        {
            "Length": [len(nodes) for nodes in cycle_nodes],
            "Cycle": [_arrows(G, path) for path in all_cycles],
            "Sign": cycle_signs,
        }
    )
    return cycles_df

@cache
def get_paths(
    G: nx.DiGraph,
    source: str,
    target: str,
    form: Literal["symbolic", "signed", "binary"] = "symbolic",
) -> sp.Matrix:
    """Find all causal pathways between two nodes.

    Args:
        G: NetworkX DiGraph representing signed digraph model
        source: Source node
        target: Target node (can be same as source for self-effect)
        form: Type of path products ('symbolic', 'signed', or 'binary')

    Returns:
        sp.Matrix: Products of interactions along each path

    References:
        - Mason, S.J. (1953). Feedback Theory-Some Properties of Signal Flow Graphs. Proceedings of the IRE 41, 1144–1156.
        - Levins, R. (1974). The qualitative analysis of partially specified systems. Annals of the New York Academy of Sciences 231, 123–138.
        - Puccia, C.J., Levins, R. (1985). Qualitative Modeling of Complex Systems: An Introduction to Loop Analysis and Time Averaging. Harvard University Press.

    Examples:
        ```python
        from qmm import get_paths, load_digraph
        get_paths(load_digraph("snowshoe_io"), 'Inp1', 'Out1', form='symbolic')
        # Matrix([
        # [a_C,R*a_P,C*b_R,Inp1*c_Out1,P],
        # [     -a_C,R*b_R,Inp1*c_Out1,C],
        # [     -a_P,C*b_C,Inp1*c_Out1,P],
        # [            b_C,Inp1*c_Out1,C]])

        get_paths(load_digraph("snowshoe_io"), 'Inp1', 'Out1', form='signed')
        # Matrix([
        # [ 1],
        # [-1],
        # [-1],
        # [ 1]])

        get_paths(load_digraph("snowshoe_io"), 'Inp1', 'Out1', form='binary')
        # Matrix([
        # [1],
        # [1],
        # [1],
        # [1]])
        ```
    """
    all_nodes = get_nodes(G, "all")
    if source not in all_nodes or target not in all_nodes:
        raise ValueError("Invalid source or target node")
    if source == target:
        return sp.Matrix([sp.Integer(1)])
    if not nx.has_path(G, source, target):
        return sp.Matrix([sp.Integer(0)])
    path_nodes = list(nx.all_simple_paths(G, source, target))
    paths = []
    for p in path_nodes:
        effect = sp.Integer(1)
        for i in range(len(p) - 1):
            u, v = p[i], p[i + 1]
            s = G[u][v].get("sign", 1)
            if form == "binary":
                effect *= sp.Integer(1)
            elif form == "signed":
                effect *= sp.Integer(s)
            else:
                u_cat = G.nodes[u].get("category", "state")
                v_cat = G.nodes[v].get("category", "state")
                if u_cat == "input" and v_cat == "output":
                    prefix = "d"
                elif u_cat == "input":
                    prefix = "b"
                elif v_cat == "output":
                    prefix = "c"
                else:
                    prefix = "a"
                effect *= sp.Symbol(f"{prefix}_{v},{u}") * s
        paths.append(effect)
    return sp.Matrix(paths)

@cache
def paths_table(G: nx.DiGraph, source: str, target: str) -> Optional[pd.DataFrame]:
    """Create table of paths between nodes.

    Args:
        G (nx.DiGraph): NetworkX DiGraph representing signed digraph model
        source (str): Source node
        target (str): Target node

    Returns:
        Optional[pd.DataFrame]: DataFrame containing path information or None if no paths exist

    References:
        - Mason, S.J. (1953). Feedback Theory-Some Properties of Signal Flow Graphs. Proceedings of the IRE 41, 1144–1156.
        - Levins, R. (1974). The qualitative analysis of partially specified systems. Annals of the New York Academy of Sciences 231, 123–138.
        - Puccia, C.J., Levins, R. (1985). Qualitative Modeling of Complex Systems: An Introduction to Loop Analysis and Time Averaging. Harvard University Press.

    Examples:
        ```python
        from qmm import paths_table, load_digraph
        paths_table(load_digraph("snowshoe_io"), 'Inp1', 'Out1')
        #    Length                                                                     Path Sign
        # 0       4  Inp1 $\\rightarrow$ R $\\rightarrow$ C $\\rightarrow$ P $\\rightarrow$ Out1    +
        # 1       3                    Inp1 $\\rightarrow$ R $\\rightarrow$ C $\\multimap$ Out1    −
        # 2       3                    Inp1 $\\multimap$ C $\\rightarrow$ P $\\rightarrow$ Out1    −
        # 3       2                                      Inp1 $\\multimap$ C $\\multimap$ Out1    +
        ```
    """
    nodes = get_nodes(G, "all")
    if source not in nodes or target not in nodes or source == target:
        raise ValueError("Invalid source or target node")
    if not nx.has_path(G, source, target):
        return None
    paths = list(nx.all_simple_paths(G, source, target))
    paths_df = pd.DataFrame(
        {
            "Length": [len(path) - 1 for path in paths],
            "Path": [_arrows(G, path) for path in paths],
            "Sign": [_sign_string(G, path) for path in paths],
        }
    )
    return paths_df

@cache
def complementary_feedback(
    G: nx.DiGraph,
    source: str,
    target: str,
    form: Literal["symbolic", "signed", "binary"] = "symbolic",
) -> sp.Matrix:
    """Calculate feedback from state nodes not on paths between source and target.

    Args:
        G: NetworkX DiGraph representing signed digraph model
        source: Source node
        target: Target node (can be same as source for self-effect)
        form: Type of feedback ('symbolic', 'signed', or 'binary')

    Returns:
        sp.Matrix: Feedback cycles in complementary subsystem

    References:
        - Mason, S.J. (1953). Feedback Theory-Some Properties of Signal Flow Graphs. Proceedings of the IRE 41, 1144–1156.
        - Levins, R. (1974). The qualitative analysis of partially specified systems. Annals of the New York Academy of Sciences 231, 123–138.
        - Puccia, C.J., Levins, R. (1985). Qualitative Modeling of Complex Systems: An Introduction to Loop Analysis and Time Averaging. Harvard University Press.

    Examples:
        ```python
        from qmm import complementary_feedback, load_digraph
        complementary_feedback(load_digraph("snowshoe_io"), 'Inp1', 'Out1', form='symbolic')
        # Matrix([
        # [          -1],
        # [      -a_P,P],
        # [      -a_R,R],
        # [-a_P,P*a_R,R]])

        complementary_feedback(load_digraph("snowshoe_io"), 'Inp1', 'Out1', form='signed')
        # Matrix([
        # [-1],
        # [-1],
        # [-1],
        # [-1]])

        complementary_feedback(load_digraph("snowshoe_io"), 'Inp1', 'Out1', form='binary')
        # Matrix([
        # [1],
        # [1],
        # [1],
        # [1]])
        ```
    """
    state_nodes = get_nodes(G, "state")
    all_nodes = get_nodes(G, "all")
    if source not in all_nodes or target not in all_nodes:
        raise ValueError("Invalid source or target node")
    if source == target:
        paths = [[source]]
    elif not nx.has_path(G, source, target):
        return sp.Matrix([sp.Integer(0)])
    else:
        paths = list(nx.all_simple_paths(G, source, target))
    
    feedback_funcs = {"symbolic": system_feedback, "signed": net_feedback, "binary": absolute_feedback}
    if form not in feedback_funcs:
        raise ValueError("Invalid form. Choose 'symbolic', 'signed', or 'binary'.")
    
    feedback = []
    for path in paths:
        subsystem_nodes = [n for n in state_nodes if n not in path]
        if not subsystem_nodes:
            feedback.append(sp.Integer(1) if form == "binary" else sp.Integer(-1))
        else:
            subsystem = G.subgraph(subsystem_nodes).copy()
            feedback.append(feedback_funcs[form](subsystem, level=len(subsystem_nodes))[0])
    return sp.Matrix([sp.expand_mul(f) for f in feedback])

@cache
def system_paths(
    G: nx.DiGraph,
    source: str,
    target: str,
    form: Literal["symbolic", "signed", "binary"] = "symbolic",
) -> sp.Matrix:
    """Calculate combined effect of paths and complementary feedback.

    Args:
        G: NetworkX DiGraph representing signed digraph model
        source: Source node
        target: Target node (can be same as source for self-effect)
        form: Type of computation ('symbolic', 'signed', or 'binary')

    Returns:
        sp.Matrix: Total effects including paths and feedback

    References:
        - Mason, S.J. (1953). Feedback Theory-Some Properties of Signal Flow Graphs. Proceedings of the IRE 41, 1144–1156.
        - Levins, R. (1974). The qualitative analysis of partially specified systems. Annals of the New York Academy of Sciences 231, 123–138.
        - Puccia, C.J., Levins, R. (1985). Qualitative Modeling of Complex Systems: An Introduction to Loop Analysis and Time Averaging. Harvard University Press.

    Examples:
        ```python
        from qmm import system_paths, load_digraph
        system_paths(load_digraph("snowshoe_io"), 'Inp1', 'Out1', form='symbolic')
        # Matrix([
        # [ a_C,R*a_P,C*b_R,Inp1*c_Out1,P],
        # [-a_C,R*a_P,P*b_R,Inp1*c_Out1,C],
        # [-a_P,C*a_R,R*b_C,Inp1*c_Out1,P],
        # [ a_P,P*a_R,R*b_C,Inp1*c_Out1,C]])

        system_paths(load_digraph("snowshoe_io"), 'Inp1', 'Out1', form='signed')
        # Matrix([
        # [ 1],
        # [-1],
        # [-1],
        # [ 1]])

        system_paths(load_digraph("snowshoe_io"), 'Inp1', 'Out1', form='binary')
        # Matrix([
        # [1],
        # [1],
        # [1],
        # [1]])
        ```
    """
    _check_direct_io_edges(G)
    path = get_paths(G, source, target, form=form)
    feedback = complementary_feedback(G, source, target, form=form)
    if form == "binary":
        effect = path.multiply_elementwise(feedback)
    else:
        effect = path.multiply_elementwise(feedback) / sp.Integer(-1)
    return sp.Matrix([sp.expand_mul(e) for e in effect])

@cache
def weighted_paths(G: nx.DiGraph, source: str, target: str) -> sp.Matrix:
    """Calculate ratio of net to total path effects.

    Args:
        G: NetworkX DiGraph representing signed digraph model
        source: Source node
        target: Target node

    Returns:
        sp.Matrix: Net-to-total ratios for path predictions

    Examples:
        ```python
        from qmm import weighted_paths, load_digraph
        weighted_paths(load_digraph("snowshoe_io"), 'Inp1', 'Out1')
        # Matrix([
        # [ 1],
        # [-1],
        # [-1],
        # [ 1]])
        ```
    """
    _check_direct_io_edges(G)
    state_nodes = get_nodes(G, "state")
    all_nodes = get_nodes(G, "all")
    if source not in all_nodes or target not in all_nodes or source == target:
        raise ValueError("Invalid source or target node")
    path_nodes = list(nx.all_simple_paths(G, source, target))
    wgt_effects = []
    for path in path_nodes:
        subsystem_nodes = [n for n in state_nodes if n not in path]
        if not subsystem_nodes:
            feedback = sp.Integer(-1)
        else:
            subsystem = G.subgraph(subsystem_nodes).copy()
            feedback = weighted_feedback(subsystem, level=len(subsystem_nodes))
            if feedback[0] == sp.nan:
                feedback = sp.Integer(0)
        sign = 1
        for i in range(len(path) - 1):
            sign *= G[path[i]][path[i + 1]].get('sign', 1)
        wgt_effect = sp.Integer(-1) * sign * feedback
        wgt_effects.append(wgt_effect)
    return sp.Matrix(wgt_effects)

@cache
def path_metrics(G: nx.DiGraph, source: str, target: str) -> pd.DataFrame:
    """Calculate comprehensive metrics for paths between nodes.

    Args:
        G: NetworkX DiGraph representing signed digraph model
        source: Source node
        target: Target node

    Returns:
        pd.DataFrame: Metrics including path length, sign, and feedback

    Examples:
        ```python
        from qmm import path_metrics, load_digraph
        path_metrics(load_digraph("snowshoe_io"), 'Inp1', 'Out1')
        #    Length                 Path Path sign Complementary subsystem Net feedback Absolute feedback Positive feedback Negative feedback Weighted feedback Weighted path
        # 0       4  Inp1, R, C, P, Out1         +                    None           -1                 1                 0                 1                -1             1
        # 1       3     Inp1, R, C, Out1         −                       P           -1                 1                 0                 1                -1            -1
        # 2       3     Inp1, C, P, Out1         −                       R           -1                 1                 0                 1                -1            -1
        # 3       2        Inp1, C, Out1         +                    R, P           -1                 1                 0                 1                -1             1
        ```
    """
    _check_direct_io_edges(G)
    state_nodes = get_nodes(G, "state")
    all_nodes = get_nodes(G, "all")
    if source not in all_nodes or target not in all_nodes or source == target:
        raise ValueError("Invalid source or target node")
    if not nx.has_path(G, source, target):
        return pd.DataFrame()
    paths = list(nx.all_simple_paths(G, source, target))
    subsystem_nodes = [[n for n in state_nodes if n not in set(path)] for path in paths]
    net_fb = complementary_feedback(G, source=source, target=target, form="signed")
    absolute_fb = complementary_feedback(G, source=source, target=target, form="binary")
    path_signs = get_paths(G, source=source, target=target, form="signed")
    weighted_fb = get_weight(net_fb, absolute_fb, sp.Integer(0))
    positive_fb = get_positive(net_fb, absolute_fb)
    negative_fb = get_negative(net_fb, absolute_fb)
    weighted_path = weighted_paths(G, source, target)
    n = len(paths)
    paths_df = pd.DataFrame(
        {
            "Length": [len(path) - 1 for path in paths],
            "Path": [", ".join(str(x) for x in path) for path in paths],
            "Path sign": ["+" if sign == 1 else "−" for sign in path_signs[:n]],
            "Complementary subsystem": [
                ", ".join(str(x) for x in nodes) if nodes else None for nodes in subsystem_nodes
            ],
            "Net feedback": [net_fb[i] for i in range(n)],
            "Absolute feedback": [absolute_fb[i] for i in range(n)],
            "Positive feedback": [positive_fb[i] for i in range(n)],
            "Negative feedback": [negative_fb[i] for i in range(n)],
            "Weighted feedback": [weighted_fb[i] for i in range(n)],
            "Weighted path": [weighted_path[i] for i in range(n)],
        }
    )
    return paths_df
