import sympy as sp
import numpy as np
import pandas as pd
from functools import cache
from .effects import get_simulations
from ..core.helper import get_nodes, _arrows, _NodeSign
import networkx as nx
from typing import List, Optional


def _parse_observations(s: str) -> tuple[tuple[str, int], ...]:
    if not s:
        return tuple()
    return tuple(_NodeSign.from_str(obs.strip()).to_tuple() 
                for obs in s.split(","))

@cache
def marginal_likelihood(G: nx.DiGraph, perturb: str, observe: str, n_sim: int = 10000, distribution: str = "uniform", seed: int = 42) -> float:
    """Calculate proportion of simulations matching qualitative observations.

    Args:
        G: NetworkX DiGraph representing signed digraph model
        perturb: Node and sign to perturb
        observe: String of observations
        n_sim: Number of simulations
        distribution: Distribution for sampling 
        seed: Random seed

    Returns:
        float: Marginal likelihood
    """
    sims = get_simulations(G, n_sim, distribution, seed,
                          _NodeSign.from_str(perturb).to_tuple(),
                          _parse_observations(observe))
    return sum(sims["valid_sims"]) / n_sim

@cache
def model_validation(G: nx.DiGraph, perturb: str, observe: str, n_sim: int = 10000, distribution: str = "uniform", seed: int = 42) -> pd.DataFrame:
    """Compare marginal likelihoods from alternative model structures.

    Args:
        G: NetworkX DiGraph representing signed digraph model
        perturb: Node and sign to perturb
        observe: String of observations
        
    Returns:
        pd.DataFrame: Marginal likelihood comparison for model variants
    """
    dashed_edges = [(u, v) for u, v, d in G.edges(data=True) if d.get("dashes", False)]
    variants = []
    edge_presence = []
    for i in range(2 ** len(dashed_edges)):
        G_variant = G.copy()
        presence = []
        for j, (u, v) in enumerate(dashed_edges):
            if not (i & (1 << j)):
                G_variant.remove_edge(u, v)
            presence.append(bool(i & (1 << j)))
        variants.append(G_variant)
        edge_presence.append(presence)
    
    likelihoods = [marginal_likelihood(G, perturb, observe, n_sim, distribution, seed) for G in variants]
    columns = ["Marginal likelihood"] + [_arrows(G, [u, v]) for u, v in dashed_edges]
    rows = []
    for i in range(len(variants)):
        row = {
            "Marginal likelihood": likelihoods[i]
        }
        for j, (u, v) in enumerate(dashed_edges):
            row[_arrows(G, [u, v])] = "\u2713" if edge_presence[i][j] else ""
        rows.append(row)
    df = pd.DataFrame(rows, columns=columns).sort_values("Marginal likelihood", ascending=False).reset_index(drop=True)    
    df["Marginal likelihood"] = df["Marginal likelihood"].apply(lambda x: f"{x:.3f}")
    return df

@cache
def posterior_predictions(G: nx.DiGraph, perturb: str, observe: str = "", n_sim: int = 10000, dist: str = "uniform", seed: int = 42) -> sp.Matrix:
    """Calculate model predictions conditioned on observations.

    Args:
        G: NetworkX DiGraph representing signed digraph model
        perturb: Node and sign to perturb
        observe: String of observations
        n_sim: Number of simulations
        dist: Distribution for sampling
        seed: Random seed
        
    Returns:
        sp.Matrix: Predictions conditioned on observations
    """
    sims = get_simulations(G, n_sim, dist, seed,
                          _NodeSign.from_str(perturb).to_tuple(),
                          _parse_observations(observe))
    state_nodes, output_nodes = get_nodes(G, "state"), get_nodes(G, "output")
    n, m = len(state_nodes), len(output_nodes)
    valid_count = sum(sims["valid_sims"])
    tmat = sims["tmat"]
    if valid_count == 0:
        return sp.Matrix([np.nan] * (n + m))
    effects = np.array(
        [
            e[: n + m] if len(e) >= n + m else np.pad(e, (0, n + m - len(e)))
            for e, v in zip(sims["effects"], sims["valid_sims"])
            if v
        ]
    )
    positive = np.sum(effects > 0, axis=0)
    negative = np.sum(effects < 0, axis=0)
    smat = positive / valid_count
    tmat_np = np.array(tmat.tolist(), dtype=bool)
    perturb_node = _NodeSign.from_str(perturb).node
    perturb_index = sims["all_nodes"].index(perturb_node)
    smat = [np.nan if not tmat_np[i, perturb_index] else smat[i] for i in range(n + m)]
    
    if observe:
        observations = _parse_observations(observe)
        for node, value in observations:
            index = (
                state_nodes.index(node)
                if node in state_nodes
                else (n + output_nodes.index(node) if node in output_nodes else None)
            )
            if index is not None:
                smat[index] = 1 if value > 0 else (0 if value < 0 else np.nan)

    smat = np.where(negative > positive, -negative / valid_count, smat)
    return sp.Matrix(smat)

@cache
def diagnose_observations(G: nx.DiGraph, observe: str, n_sim: int = 10000, distribution: str = "uniform", seed: int = 42) -> pd.DataFrame:
    """Identify possible perturbations from marginal likelihoods.

    Args:
        G: NetworkX DiGraph representing signed digraph model
        observe: String of observations
        
    Returns:
        pd.DataFrame: Ranked perturbations matching observations
    """
    perturb_nodes = get_nodes(G, "state") + get_nodes(G, "input")
    results = []
    for node in perturb_nodes:
        for sign in ["+", "-"]:
            try:
                likelihood = marginal_likelihood(G, f"{node}:{sign}", observe, n_sim, distribution, seed)
                results.append({"Input": node, "Sign": sign, "Marginal likelihood": likelihood})
            except Exception as e:
                print(f"Error for node {node} with sign {sign}: {str(e)}")
    return pd.DataFrame(results).sort_values("Marginal likelihood", ascending=False).reset_index(drop=True)


def bayes_factors(G_list: List[nx.DiGraph], perturb: str, observe: str,
                 n_sim: int = 10000, distribution: str = "uniform", 
                 seed: int = 42, names: Optional[List[str]] = None) -> pd.DataFrame:
    """Calculate Bayes factors from the ratio of marginal likelihoods of alternative models.

    Args:
        G_list (List[nx.DiGraph]): List of NetworkX DiGraphs representing alternative models
        perturb (str): Node and sign to perturb
        observe (str): String of observations
        n_sim (int): Number of simulations
        distribution (str): Distribution for sampling
        seed (int): Random seed
        names (Optional[List[str]]): Optional list of model names
        
    Returns:
        pd.DataFrame: DataFrame containing Bayes factors
    """
    likelihoods = [marginal_likelihood(G, perturb, observe, n_sim, distribution, seed) for G in G_list]
    model_names = names if names and len(names) == len(G_list) else [f"Model {chr(65+i)}" for i in range(len(G_list))]
    bayes_factors = {f"{model_names[i]}/{model_names[j]}": (
        float("inf") if likelihoods[j] == 0 and likelihoods[i] > 0 else
        0 if likelihoods[j] == 0 else likelihoods[i] / likelihoods[j]
    ) for i in range(len(G_list)) for j in range(i + 1, len(G_list))}

    return pd.DataFrame({
        "Model comparison": list(bayes_factors.keys()),
        "Likelihood 1": [likelihoods[i] for i in range(len(G_list)) for j in range(i + 1, len(G_list))],
        "Likelihood 2": [likelihoods[j] for i in range(len(G_list)) for j in range(i + 1, len(G_list))],
        "Bayes factor": list(bayes_factors.values()),
    })