import math
import random
from dataclasses import dataclass
from typing import Any, Dict, List, Mapping, Sequence, Tuple

from core import clamp, require


@dataclass
class EdgeCounts:
    cc: float
    cd: float
    dd: float

    @property
    def total(self) -> float:
        return float(self.cc + self.cd + self.dd)


@dataclass
class EdgeMetrics:
    mu: float
    eta: float
    h: float


def edge_counts(states: Sequence[int], edges: Sequence[Tuple[int, int]]) -> EdgeCounts:
    cc = 0.0
    cd = 0.0
    dd = 0.0
    for (u, v) in edges:
        su = 1 if states[u] else 0
        sv = 1 if states[v] else 0
        if su == 1 and sv == 1:
            cc += 1.0
        elif su == 0 and sv == 0:
            dd += 1.0
        else:
            cd += 1.0
    return EdgeCounts(cc=cc, cd=cd, dd=dd)


def edge_metrics_from_counts(c: EdgeCounts) -> EdgeMetrics:
    m = c.total
    if m <= 0.0:
        return EdgeMetrics(mu=0.0, eta=0.0, h=0.0)

    p_cc = c.cc / m
    p_cd = c.cd / m
    p_dd = c.dd / m

    mu = p_cc + 0.5 * p_cd
    eta = p_cd

    h = 0.0
    for p in (p_cc, p_cd, p_dd):
        if p > 0.0:
            h -= p * math.log(p)
    h /= math.log(3.0)

    return EdgeMetrics(mu=float(mu), eta=float(eta), h=float(h))


def sample_edge_counts(states: Sequence[int], edges: Sequence[Tuple[int, int]], q: float, rng: random.Random) -> EdgeCounts:
    cc = 0.0
    cd = 0.0
    dd = 0.0
    qq = clamp(float(q), 0.0, 1.0)
    if qq <= 0.0:
        return EdgeCounts(cc=0.0, cd=0.0, dd=0.0)

    for (u, v) in edges:
        if rng.random() >= qq:
            continue
        su = 1 if states[u] else 0
        sv = 1 if states[v] else 0
        if su == 1 and sv == 1:
            cc += 1.0
        elif su == 0 and sv == 0:
            dd += 1.0
        else:
            cd += 1.0
    return EdgeCounts(cc=cc, cd=cd, dd=dd)


def apply_count_noise(c: EdgeCounts, noise_cfg: Mapping[str, Any], rng: random.Random) -> EdgeCounts:
    typ = str(require(noise_cfg, "type", "observation.count_noise")).lower()
    if typ == "none":
        return c

    total = c.total
    if total <= 0.0:
        return c

    if typ == "gaussian":
        sigma = float(require(noise_cfg, "sigma", "observation.count_noise"))
        cc = max(0.0, c.cc + rng.gauss(0.0, sigma * total))
        cd = max(0.0, c.cd + rng.gauss(0.0, sigma * total))
        dd = max(0.0, c.dd + rng.gauss(0.0, sigma * total))
        s = cc + cd + dd
        if s <= 0.0:
            return EdgeCounts(cc=total / 3.0, cd=total / 3.0, dd=total / 3.0)
        scale = total / s
        return EdgeCounts(cc=cc * scale, cd=cd * scale, dd=dd * scale)

    raise ValueError(f"Unsupported observation.count_noise.type: {typ}")


def observed_counts(states: Sequence[int], edges: Sequence[Tuple[int, int]], obs_cfg: Mapping[str, Any], rng: random.Random) -> EdgeCounts:
    q = float(require(obs_cfg, "edge_sample_prob", "observation"))
    raw = sample_edge_counts(states, edges, q=q, rng=rng)
    noise_cfg = obs_cfg.get("count_noise", {"type": "none"})
    return apply_count_noise(raw, noise_cfg=noise_cfg, rng=rng)


def community_edge_counts(
    states: Sequence[int],
    edges: Sequence[Tuple[int, int]],
    community_by_node: Sequence[int],
    k: int,
) -> List[EdgeCounts]:
    kk = int(k)
    if kk <= 0:
        raise ValueError("community_edge_counts: k must be > 0")
    out = [EdgeCounts(cc=0.0, cd=0.0, dd=0.0) for _ in range(kk)]
    for (u, v) in edges:
        cu = int(community_by_node[int(u)])
        cv = int(community_by_node[int(v)])
        if cu != cv:
            continue
        if cu < 0 or cu >= kk:
            continue
        su = 1 if states[u] else 0
        sv = 1 if states[v] else 0
        if su == 1 and sv == 1:
            out[cu].cc += 1.0
        elif su == 0 and sv == 0:
            out[cu].dd += 1.0
        else:
            out[cu].cd += 1.0
    return out


def sample_community_edge_counts(
    states: Sequence[int],
    edges: Sequence[Tuple[int, int]],
    community_by_node: Sequence[int],
    k: int,
    q: float,
    rng: random.Random,
) -> List[EdgeCounts]:
    kk = int(k)
    if kk <= 0:
        raise ValueError("sample_community_edge_counts: k must be > 0")
    out = [EdgeCounts(cc=0.0, cd=0.0, dd=0.0) for _ in range(kk)]

    qq = clamp(float(q), 0.0, 1.0)
    if qq <= 0.0:
        return out

    for (u, v) in edges:
        if rng.random() >= qq:
            continue
        cu = int(community_by_node[int(u)])
        cv = int(community_by_node[int(v)])
        if cu != cv:
            continue
        if cu < 0 or cu >= kk:
            continue
        su = 1 if states[u] else 0
        sv = 1 if states[v] else 0
        if su == 1 and sv == 1:
            out[cu].cc += 1.0
        elif su == 0 and sv == 0:
            out[cu].dd += 1.0
        else:
            out[cu].cd += 1.0

    return out


def observed_community_counts(
    states: Sequence[int],
    edges: Sequence[Tuple[int, int]],
    community_by_node: Sequence[int],
    k: int,
    obs_cfg: Mapping[str, Any],
    rng: random.Random,
) -> List[EdgeCounts]:
    q = float(require(obs_cfg, "edge_sample_prob", "observation"))
    raw = sample_community_edge_counts(states, edges, community_by_node=community_by_node, k=k, q=q, rng=rng)
    noise_cfg = obs_cfg.get("count_noise", {"type": "none"})
    return [apply_count_noise(c, noise_cfg=noise_cfg, rng=rng) for c in raw]


def mu_dispersion(mu_by_comm: Sequence[float]) -> float:
    xs = [float(x) for x in mu_by_comm]
    if not xs:
        return 0.0
    m = sum(xs) / float(len(xs))
    v = 0.0
    for x in xs:
        d = float(x) - float(m)
        v += d * d
    v /= float(len(xs))
    return float(math.sqrt(v))


def mu_by_comm_and_dispersion_from_counts(
    comm_counts: Sequence[EdgeCounts],
    *,
    min_total_edges: float = 1.0,
) -> Tuple[List[float], float]:
    mus: List[float] = []
    for c in comm_counts:
        if float(c.total) < float(min_total_edges):
            continue
        m = edge_metrics_from_counts(c)
        mus.append(float(m.mu))
    return mus, mu_dispersion(mus)


def mu_hat_by_comm_and_dispersion_from_counts(
    comm_counts: Sequence[EdgeCounts],
    *,
    meso_alpha0: Sequence[float],
    min_total_edges: float = 1.0,
) -> Tuple[List[Tuple[int, float]], float]:
    if len(meso_alpha0) != 3:
        raise ValueError("meso_alpha0 must be a list of length 3")
    a0 = [float(x) for x in meso_alpha0]
    if any(x <= 0.0 for x in a0):
        raise ValueError("meso_alpha0 must be positive")

    pairs: List[Tuple[int, float]] = []
    for i, c in enumerate(comm_counts):
        if float(c.total) < float(min_total_edges):
            continue

        a_cc = float(a0[0]) + float(c.cc)
        a_cd = float(a0[1]) + float(c.cd)
        a_dd = float(a0[2]) + float(c.dd)
        s = float(a_cc + a_cd + a_dd)
        mu_hat = 0.0
        if s > 0.0:
            p_cc = float(a_cc) / s
            p_cd = float(a_cd) / s
            mu_hat = float(p_cc + 0.5 * p_cd)
        pairs.append((int(i), float(mu_hat)))

    mu_disp = mu_dispersion([float(x[1]) for x in pairs])
    return pairs, float(mu_disp)


def _dirichlet_sample_p3(alpha: Sequence[float], rng: random.Random) -> Tuple[float, float, float]:
    xs: List[float] = []
    s = 0.0
    for a in alpha:
        aa = float(a)
        if aa <= 0.0:
            aa = 1e-12
        x = rng.gammavariate(aa, 1.0)
        xs.append(x)
        s += x
    if s <= 0.0:
        return (1.0 / 3.0, 1.0 / 3.0, 1.0 / 3.0)
    return (float(xs[0]) / float(s), float(xs[1]) / float(s), float(xs[2]) / float(s))


def mu_risk_by_comm_and_dispersion_from_counts(
    comm_counts: Sequence[EdgeCounts],
    *,
    meso_alpha0: Sequence[float],
    mu_min: float,
    min_total_edges: float = 1.0,
    k_mc: int = 64,
    rng: random.Random,
) -> Tuple[List[Tuple[int, float]], float]:
    if len(meso_alpha0) != 3:
        raise ValueError("meso_alpha0 must be a list of length 3")
    a0 = [float(x) for x in meso_alpha0]
    if any(x <= 0.0 for x in a0):
        raise ValueError("meso_alpha0 must be positive")

    kk = int(k_mc)
    if kk <= 0:
        raise ValueError("k_mc must be > 0")
    mu_min_f = float(mu_min)

    pairs: List[Tuple[int, float]] = []
    for i, c in enumerate(comm_counts):
        if float(c.total) < float(min_total_edges):
            continue

        alpha = [float(a0[0]) + float(c.cc), float(a0[1]) + float(c.cd), float(a0[2]) + float(c.dd)]
        below = 0
        for _ in range(kk):
            p_cc, p_cd, _p_dd = _dirichlet_sample_p3(alpha, rng=rng)
            mu_s = float(p_cc + 0.5 * float(p_cd))
            if float(mu_s) < float(mu_min_f):
                below += 1
        risk = float(below) / float(kk)
        pairs.append((int(i), float(risk)))

    disp = mu_dispersion([float(x[1]) for x in pairs])
    return pairs, float(disp)