"""Confounder identification following Cinelli, Forney & Pearl (2022).

Ports inferConfounders.R — Models 1-6, 8, 13-15.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from ..variables import Measure, UnobservedVariable
from ..graph import parents, children, ancestors

if TYPE_CHECKING:
    from ..conceptual_model import ConceptualModel
    from ..variables import AbstractVariable


def is_observed(cm: "ConceptualModel", variable_name: str) -> bool:
    """Return True if *variable_name* refers to a Measure in *cm*."""
    var = cm.variables.get(variable_name)
    if var is None:
        return False
    if isinstance(var, UnobservedVariable):
        return False
    return isinstance(var, Measure)


def get_mediators(cm: "ConceptualModel", x: str, y: str) -> list[str]:
    """Find observed mediators on directed paths from *x* to *y*.

    A mediator is a node Z on a path x -> Z -> y (exactly two directed edges).
    """
    from ..graph import find_paths

    mediators: list[str] = []
    for path in find_paths(cm.graph, x, y):
        # A path of length 3 means x -> Z -> y
        if len(path) == 3:
            z = path[1]
            if is_observed(cm, z):
                mediators.append(z)
    return mediators


def infer_confounders(cm: "ConceptualModel", iv: "AbstractVariable", dv: "AbstractVariable") -> list[str]:
    """Identify confounders to include as controls.

    Implements Models 1-6, 8, 13-15 from Cinelli, Forney & Pearl.
    """
    confounders: list[str] = []
    gr = cm.graph
    x = iv.name
    y = dv.name

    iv_parents = parents(gr, x)
    dv_parents = parents(gr, y)

    # Model 1: Common observed cause of IV and DV
    common_parents = set(iv_parents) & set(dv_parents)
    for p in common_parents:
        if is_observed(cm, p):
            confounders.append(p)

    # Model 2: Unobserved common ancestor of IV and DV; Z mediates U -> Z -> X
    for ip in iv_parents:
        iv_grandparents_unobs = [
            v for v in parents(gr, ip) if not is_observed(cm, v)
        ]
        for gp in iv_grandparents_unobs:
            if gp in dv_parents and is_observed(cm, ip):
                confounders.append(ip)

    # Model 3: Unobserved common ancestor of IV and DV; Z mediates U -> Z -> Y
    for dp in dv_parents:
        dv_grandparents_unobs = [
            v for v in parents(gr, dp) if not is_observed(cm, v)
        ]
        for gp in dv_grandparents_unobs:
            if gp in iv_parents and is_observed(cm, dp) and dp != x:
                confounders.append(dp)

    # Model 4: Common cause of X and any mediator between X and Y
    mediators = get_mediators(cm, x, y)
    for m in mediators:
        m_parents = parents(gr, m)
        shared = set(iv_parents) & set(m_parents)
        for sp in shared:
            if is_observed(cm, sp):
                confounders.append(sp)

    # Model 5: Unobserved common ancestor of IV and mediator; Z mediates U -> Z -> M
    iv_parents_unobs = [ip for ip in iv_parents if not is_observed(cm, ip)]
    for m in mediators:
        m_parents = parents(gr, m)
        for mp in m_parents:
            if mp != x:
                m_grandparents = parents(gr, mp)
                shared_anc = set(iv_parents_unobs) & set(m_grandparents)
                if shared_anc:
                    confounders.append(mp)

    # Model 6: Unobserved parent of mediator; Z mediates U -> Z -> X
    for m in mediators:
        m_parents = parents(gr, m)
        m_parents_unobs = [mp for mp in m_parents if not is_observed(cm, mp)]
        for ip in iv_parents:
            iv_grandparents = parents(gr, ip)
            shared_anc = set(iv_grandparents) & set(m_parents_unobs)
            if shared_anc:
                confounders.append(ip)

    # Neutral Controls ---------------------------------------------------

    # Model 8: Parent of Y unrelated to X (precision gain)
    for dp in dv_parents:
        if is_observed(cm, dp) and dp != x:
            d_anc = ancestors(gr, dp)
            if d_anc == {dp}:
                a_children = children(gr, dp)
                if len(a_children) == 1 and a_children[0] == y:
                    confounders.append(dp)

    # Model 13: Parent of mediator (precision gain)
    for m in mediators:
        m_parents = parents(gr, m)
        for mp in m_parents:
            if mp != x and is_observed(cm, mp):
                mp_anc = ancestors(gr, mp)
                if mp_anc == {mp}:
                    mp_ch = children(gr, mp)
                    if len(mp_ch) == 1 and mp_ch[0] == m:
                        confounders.append(mp)

    # Model 14: Child of X (post-treatment, precision for selection bias)
    iv_children = children(gr, x)
    for ic in iv_children:
        if ic != x and ic != y and is_observed(cm, ic):
            ic_anc = ancestors(gr, ic)
            if len(ic_anc) == 2 and ic in ic_anc and x in ic_anc:
                ic_ch = children(gr, ic)
                if len(ic_ch) == 0:
                    confounders.append(ic)

    # Model 15: Child of X whose child is also child of Unobserved parent of Y
    iv_children_obs = [
        ic for ic in children(gr, x)
        if ic != x and is_observed(cm, ic)
    ]
    dv_parents_unobs = [dp for dp in dv_parents if not is_observed(cm, dp)]
    for ico in iv_children_obs:
        iv_grandchildren = children(gr, ico)
        for ig in iv_grandchildren:
            if is_observed(cm, ig):
                for dp in dv_parents_unobs:
                    dp_children = children(gr, dp)
                    if ig in dp_children:
                        confounders.append(ico)

    # Deduplicate while preserving order
    seen: set[str] = set()
    unique: list[str] = []
    for c in confounders:
        if c not in seen:
            seen.add(c)
            unique.append(c)
    return unique
