"""
EXACT, EXHAUSTIVE verification of Lemma 1, Theorem 1, Theorem 2 and Theorem 4
over the COMPLETE space of 4-trait evolutionary selection models:

    all 543 labelled DAGs on X_1..X_4  x  all 2^4 = 16 choices of pa(S)
    = 8,688 static models G,  each unrolled to T = 1,2,3,
    x all 55 disjoint triples (A,B,C) with A,B non-empty.

Nothing here is sampled or simulated: every number is an exact combinatorial
count at machine precision.  A random-model sweep at d = 5..10 extends the same
checks beyond the exhaustive range.

Run:  python3 exhaustive.py
"""
import itertools, json, os, time
import numpy as np
from evosel import (DG, SEL, dsep, cpdag, cpdag_key, clique_augmented,
                    multidomain_augmented, evolutionary_graph, evo_counts,
                    random_static_dag)

OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'outputs')
os.makedirs(OUT, exist_ok=True)


# --------------------------------------------------------------------------
def all_static_models(d):
    """Every labelled DAG on d traits x every subset of X as pa(S)."""
    pairs = list(itertools.combinations(range(d), 2))
    dags = []
    for code in itertools.product((0, 1, 2), repeat=len(pairs)):
        g = DG(range(d))
        for (a, b), c in zip(pairs, code):
            if c == 1:
                g.add(a, b)
            elif c == 2:
                g.add(b, a)
        if g.is_acyclic():
            dags.append([(u, v) for (u, v) in g.edges()])
    models = []
    for es in dags:
        for r in range(d + 1):
            for ps in itertools.combinations(range(d), r):
                G = DG(list(range(d)) + [SEL], es)
                for p in ps:
                    G.add(p, SEL)
                models.append(G)
    return dags, models


def triples(d):
    """All (A,B,C) disjoint subsets of X with A,B non-empty, de-duplicated
    under the A<->B symmetry."""
    out = []
    for code in itertools.product((0, 1, 2, 3), repeat=d):
        A = [i for i in range(d) if code[i] == 0]
        B = [i for i in range(d) if code[i] == 1]
        C = [i for i in range(d) if code[i] == 2]
        if A and B and A[0] < B[0]:
            out.append((A, B, C))
    return out


def naive_dag(G, d):
    """The selection-blind baseline: G with S simply deleted."""
    g = DG(range(d))
    for j in range(d):
        for i in G.pa[j]:
            if i != SEL:
                g.add(i, j)
    return g


# --------------------------------------------------------------------------
def run_dsep_claims(d=4, Ts=(1, 2, 3)):
    """Lemma 1 (implication + converse failure) and Theorem 1 (biconditional),
    plus the two destructive controls."""
    _, models = all_static_models(d)
    trs = triples(d)
    R = {'d': d, 'n_models': len(models), 'n_triples_per_model': len(trs),
         'Ts': list(Ts)}
    lem_tested = lem_viol = lem_conv_fail = 0
    thm_tested = thm_viol = 0
    naive_tested = naive_mism = 0
    noinherit_tested = noinherit_conv_fail = 0
    degenerate_ok = degenerate_n = 0
    per_T = {}
    for T in Ts:
        t_l = t_v = t_c = t_t = t_tv = t_nm = 0
        for G in models:
            gT = evolutionary_graph(G, d, T)
            gp = clique_augmented(G, d)
            gn = naive_dag(G, d)
            Sall = [('S', t) for t in range(T)]
            for (A, B, C) in trs:
                AT = [('X', i, T) for i in A]
                BT = [('X', i, T) for i in B]
                CT = [('X', i, T) for i in C] + Sall
                ev = dsep(gT, AT, BT, CT)
                st = dsep(G, A, B, C + [SEL])       # static model G
                pl = dsep(gp, A, B, C)              # clique-augmented G^+
                nv = dsep(gn, A, B, C)              # S-deleted baseline
                t_l += 1
                t_v += (ev and not st)              # Lemma 1 violation
                t_c += (st and not ev)              # converse failure witness
                t_t += 1
                t_tv += (ev != pl)                  # Theorem 1 violation
                t_nm += (ev != nv)                  # control: naive mismatch
        per_T[T] = {'lemma1_tested': t_l, 'lemma1_violations': t_v,
                    'lemma1_converse_failures': t_c,
                    'theorem1_tested': t_t, 'theorem1_violations': t_tv,
                    'control_naive_Sdeleted_mismatches': t_nm}
        lem_tested += t_l; lem_viol += t_v; lem_conv_fail += t_c
        thm_tested += t_t; thm_viol += t_tv
        naive_tested += t_t; naive_mism += t_nm
    R['per_T'] = per_T
    R['lemma1'] = {'tested': lem_tested, 'violations': lem_viol,
                   'converse_failure_witnesses': lem_conv_fail}
    R['theorem1'] = {'tested': thm_tested, 'violations': thm_viol,
                     'agreement': 1.0 - thm_viol / max(thm_tested, 1)}
    R['control_naive'] = {'tested': naive_tested, 'mismatches': naive_mism,
                          'agreement': 1.0 - naive_mism / max(naive_tested, 1)}

    # ---- Theorem 1, implication 1: the d-separations do not depend on T
    diff_T = 0
    for G in models:
        sigs = []
        for T in (1, 2, 3, 4):
            gT = evolutionary_graph(G, d, T)
            Sall = [('S', t) for t in range(T)]
            sigs.append(tuple(dsep(gT, [('X', i, T) for i in A],
                                   [('X', i, T) for i in B],
                                   [('X', i, T) for i in C] + Sall)
                              for (A, B, C) in trs))
        diff_T += (len(set(sigs)) != 1)
    R['theorem1_implication1_T_invariance'] = {
        'models': len(models), 'models_whose_dsep_set_changes_with_T': diff_T,
        'Ts_compared': [1, 2, 3, 4]}

    # ---- Theorem 1, implication 2: pa(S) = {}  =>  G^+ == G minus S
    for G in models:
        if not G.pa[SEL]:
            degenerate_n += 1
            degenerate_ok += (sorted(clique_augmented(G, d).edges())
                              == sorted(naive_dag(G, d).edges()))
    R['theorem1_implication2_degenerate'] = {
        'models_with_no_selection_parents': degenerate_n,
        'models_where_Gplus_equals_G_minus_S': degenerate_ok}

    # ---- destructive control: delete the inheritance edges eps^t -> eps^t+1.
    # Without inheritance the repeated selection cannot propagate, so the
    # Lemma-1 converse failures must disappear.
    T = 3
    for G in models:
        gT = evolutionary_graph(G, d, T)
        g2 = DG(gT.nodes)
        for (u, v) in gT.edges():
            if not (u[0] == 'e' and v[0] == 'e'):
                g2.add(u, v)
        gn = naive_dag(G, d)
        Sall = [('S', t) for t in range(T)]
        for (A, B, C) in trs:
            ev = dsep(g2, [('X', i, T) for i in A], [('X', i, T) for i in B],
                      [('X', i, T) for i in C] + Sall)
            noinherit_tested += 1
            noinherit_conv_fail += (dsep(G, A, B, C + [SEL]) and not ev)
    R['control_no_inheritance'] = {
        'tested': noinherit_tested,
        'lemma1_converse_failures_without_inheritance': noinherit_conv_fail}
    return R


# --------------------------------------------------------------------------
def run_theorem2(d=4):
    """Theorem 2: adjacency soundness+completeness, orientation soundness, and
    orientation COMPLETENESS (checked by exhaustive search for the alternative
    DAG G' the theorem asserts must exist)."""
    _, models = all_static_models(d)
    keys, info = [], []
    for G in models:
        gp = clique_augmented(G, d)
        dd, uu = cpdag(gp)
        keys.append(cpdag_key(dd, uu, d))
        rel = {}
        for i, j in itertools.combinations(range(d), 2):
            rel[(i, j)] = 1 if G.has(i, j) else (2 if G.has(j, i) else 0)
        info.append((G, dd, uu, rel))
    # group models by the CPDAG that Algorithm 1 would output
    groups = {}
    for k, (G, dd, uu, rel) in zip(keys, info):
        groups.setdefault(k, []).append(rel)

    n_pairs = adj_bad = 0
    n_or = or_bad = 0
    n_un = un_incomplete = 0
    n_or_into_anS = 0
    for k, (G, dd, uu, rel) in zip(keys, info):
        anS = G.ancestors([SEL]) - {SEL}
        adj = set()
        for (u, v) in dd:
            adj.add(frozenset((u, v)))
        adj |= set(uu)
        for i, j in itertools.combinations(range(d), 2):
            n_pairs += 1
            truth = (G.has(i, j) or G.has(j, i) or ({i, j} <= anS))
            adj_bad += (truth != (frozenset((i, j)) in adj))
        for (u, v) in dd:
            n_or += 1
            or_bad += not (G.has(u, v) and v not in anS)
            n_or_into_anS += (v in anS)
        alts = groups[k]
        for e in uu:
            i, j = sorted(tuple(e))
            n_un += 1
            if not any(r[(i, j)] != rel[(i, j)] for r in alts):
                un_incomplete += 1
    return {'d': d, 'n_models': len(models),
            'adjacency_pairs_tested': n_pairs, 'adjacency_violations': adj_bad,
            'oriented_edges_tested': n_or, 'orientation_soundness_violations': or_bad,
            'oriented_edges_whose_head_is_in_an(S)': n_or_into_anS,
            'unoriented_edges_tested': n_un,
            'orientation_completeness_failures': un_incomplete,
            'n_distinct_cpdags': len(groups)}


def run_theorem2_control(d=4, n=3000, seed=3):
    """Destructive control for Theorem 2: corrupt one oriented edge of the
    CPDAG (reverse it) and confirm the soundness checker fires."""
    rng = np.random.default_rng(seed)
    _, models = all_static_models(d)
    idx = rng.choice(len(models), size=n, replace=False)
    fired = tested = 0
    for k in idx:
        G = models[k]
        gp = clique_augmented(G, d)
        dd, uu = cpdag(gp)
        anS = G.ancestors([SEL]) - {SEL}
        if not dd:
            continue
        u, v = list(dd)[int(rng.integers(len(dd)))]
        tested += 1
        fired += not (G.has(v, u) and u not in anS)   # reversed edge must fail
    return {'corrupted_models_tested': tested, 'checker_fired': fired}


# --------------------------------------------------------------------------
def run_theorem4(d=4, max_models=None, seed=5):
    """Theorem 4: P_X keeps every orientation of C (monotonicity), stays sound,
    and strictly improves on some models.  Control: drop the an_G(S) expansion
    that Theorem 3 prescribes and confirm soundness breaks."""
    _, models = all_static_models(d)
    rng = np.random.default_rng(seed)
    if max_models and max_models < len(models):
        models = [models[i] for i in rng.choice(len(models), max_models, replace=False)]
    Isets = [list(s) for r in range(d + 2)
             for s in itertools.combinations(list(range(d)) + [SEL], r)]
    n = mono_bad = sound_bad = strict = tot = 0
    extra_edges = 0
    ctrl_n = ctrl_bad = 0
    adj_bad = 0
    for G in models:
        gp = clique_augmented(G, d)
        C, Cu = cpdag(gp)
        anS = G.ancestors([SEL]) - {SEL}
        for I in Isets:
            if not I:
                continue
            gpi = multidomain_augmented(G, d, I)
            forced = [('zeta', v) for v in gpi.ch['zeta']]
            P, Pu = cpdag(gpi, forced=forced)
            PX = set((u, v) for (u, v) in P if u != 'zeta' and v != 'zeta')
            tot += 1
            n += len(PX)
            mono_bad += not (C <= PX)
            bad = sum(1 for (u, v) in PX if not (G.has(u, v) and v not in anS))
            sound_bad += bad
            strict += (len(PX) > len(C))
            extra_edges += len(PX - C)
            # adjacency characterisation must still hold for P_X
            adjP = set(frozenset((u, v)) for (u, v) in PX) | \
                   set(e for e in Pu if 'zeta' not in e)
            for i, j in itertools.combinations(range(d), 2):
                truth = (G.has(i, j) or G.has(j, i) or ({i, j} <= anS))
                adj_bad += (truth != (frozenset((i, j)) in adjP))
            # ---- destructive control: omit Theorem 3's an_G(S) expansion
            g2 = DG(list(range(d)) + ['zeta'])
            for (u, v) in gp.edges():
                g2.add(u, v)
            for x in [v for v in I if v != SEL]:
                g2.add('zeta', x)
            f2 = [('zeta', v) for v in g2.ch['zeta']]
            P2, _ = cpdag(g2, forced=f2)
            P2X = set((u, v) for (u, v) in P2 if u != 'zeta' and v != 'zeta')
            ctrl_n += 1
            ctrl_bad += sum(1 for (u, v) in P2X
                            if not (G.has(u, v) and v not in anS))
    return {'d': d, 'n_models': len(models), 'n_I_sets': len(Isets) - 1,
            'model_x_Iset_configurations': tot,
            'monotonicity_violations': mono_bad,
            'orientation_soundness_violations': sound_bad,
            'adjacency_violations': adj_bad,
            'total_oriented_edges_multi_domain': n,
            'configurations_with_strictly_more_orientations': strict,
            'frac_strict_improvement': strict / max(tot, 1),
            'additional_oriented_edges_vs_single_domain': extra_edges,
            'control_no_anS_expansion': {
                'configurations': ctrl_n,
                'orientation_soundness_violations': ctrl_bad}}


# --------------------------------------------------------------------------
def run_random_sweep(ds=(5, 6, 7, 8, 9, 10), n_models=40, n_triples=200,
                     Ts=(1, 2, 3, 5), seed=11):
    """Extend the exact checks beyond the exhaustive range with random models."""
    rng = np.random.default_rng(seed)
    tot = viol_l = conv = viol_t = naive_mis = 0
    or_tested = or_bad = adj_tested = adj_bad = 0
    for d in ds:
        for _ in range(n_models):
            G = random_static_dag(d, rng, avg_deg=2.0)
            gp = clique_augmented(G, d)
            gn = naive_dag(G, d)
            anS = G.ancestors([SEL]) - {SEL}
            dd, uu = cpdag(gp)
            for i, j in itertools.combinations(range(d), 2):
                adj_tested += 1
                truth = (G.has(i, j) or G.has(j, i) or ({i, j} <= anS))
                has = frozenset((i, j)) in (set(frozenset(e) for e in dd) | uu)
                adj_bad += (truth != has)
            for (u, v) in dd:
                or_tested += 1
                or_bad += not (G.has(u, v) and v not in anS)
            for T in Ts:
                gT = evolutionary_graph(G, d, T)
                Sall = [('S', t) for t in range(T)]
                for _ in range(n_triples):
                    code = rng.integers(0, 4, size=d)
                    A = [i for i in range(d) if code[i] == 0]
                    B = [i for i in range(d) if code[i] == 1]
                    C = [i for i in range(d) if code[i] == 2]
                    if not A or not B:
                        continue
                    ev = dsep(gT, [('X', i, T) for i in A],
                              [('X', i, T) for i in B],
                              [('X', i, T) for i in C] + Sall)
                    tot += 1
                    viol_l += (ev and not dsep(G, A, B, C + [SEL]))
                    conv += (dsep(G, A, B, C + [SEL]) and not ev)
                    viol_t += (ev != dsep(gp, A, B, C))
                    naive_mis += (ev != dsep(gn, A, B, C))
    return {'ds': list(ds), 'models_per_d': n_models, 'Ts': list(Ts),
            'dsep_relations_tested': tot,
            'lemma1_violations': viol_l, 'lemma1_converse_failures': conv,
            'theorem1_violations': viol_t,
            'control_naive_Sdeleted_mismatches': naive_mis,
            'theorem2_adjacency_pairs': adj_tested,
            'theorem2_adjacency_violations': adj_bad,
            'theorem2_oriented_edges': or_tested,
            'theorem2_orientation_violations': or_bad}


# --------------------------------------------------------------------------
def run_definition1(seed=17):
    """Definition 1: acyclicity + the closed-form node/edge counts + all four
    edge families, over an exhaustive d=4 sweep and a random d=6..20 sweep."""
    rng = np.random.default_rng(seed)
    _, models = all_static_models(4)
    n = acyc = vok = eok = fam_ok = 0
    for G in models:
        for T in (1, 2, 3):
            g = evolutionary_graph(G, 4, T)
            V, E = evo_counts(G, 4, T)
            n += 1
            acyc += g.is_acyclic()
            vok += (len(g.nodes) == V)
            eok += (g.n_edges() == E)
            fams = set()
            for (u, v) in g.edges():
                if u[0] == 'X' and v[0] == 'X':
                    fams.add('trait->trait')
                elif u[0] == 'X' and v[0] == 'S':
                    fams.add('trait->S')
                elif u[0] == 'e' and v[0] == 'X':
                    fams.add('eps->trait')
                elif u[0] == 'e' and v[0] == 'e':
                    fams.add('eps->eps')
            need = {'eps->trait', 'eps->eps'}
            if any(u != SEL for j in range(4) for u in G.pa[j]):
                need.add('trait->trait')
            if G.pa[SEL]:
                need.add('trait->S')
            fam_ok += (fams == need)
    big = []
    for d in (6, 8, 10, 15, 20):
        for T in (1, 2, 3, 5):
            for r in range(3):
                G = random_static_dag(d, rng, avg_deg=2.0)
                g = evolutionary_graph(G, d, T)
                V, E = evo_counts(G, d, T)
                big.append((g.is_acyclic(), len(g.nodes) == V, g.n_edges() == E,
                            len(g.nodes), g.n_edges(), d, T))
    return {'exhaustive_d4': {'constructions': n, 'acyclic': acyc,
                              'node_count_formula_matches': vok,
                              'edge_count_formula_matches': eok,
                              'edge_families_exactly_as_defined': fam_ok},
            'random_large': {'constructions': len(big),
                             'acyclic': sum(b[0] for b in big),
                             'node_count_matches': sum(b[1] for b in big),
                             'edge_count_matches': sum(b[2] for b in big),
                             'max_nodes': max(b[3] for b in big),
                             'max_edges': max(b[4] for b in big)},
            'node_formula': '|V| = 2d(T+1) + T',
            'edge_formula': '|E| = |E_G^X|(T+1) + |pa_G(S)|T + d(T+1) + dT'}


if __name__ == '__main__':
    R = {}
    t0 = time.time()
    R['definition1'] = run_definition1(); print('def1', time.time() - t0)
    R['theorem2'] = run_theorem2(); print('thm2', time.time() - t0)
    R['theorem2_control'] = run_theorem2_control(); print('thm2c', time.time() - t0)
    R['theorem4'] = run_theorem4(); print('thm4', time.time() - t0)
    R['random_sweep'] = run_random_sweep(); print('sweep', time.time() - t0)
    R['dsep_claims'] = run_dsep_claims(); print('dsep', time.time() - t0)
    R['runtime_sec'] = time.time() - t0
    json.dump(R, open(os.path.join(OUT, 'exhaustive.json'), 'w'), indent=1)
    print(json.dumps({k: v for k, v in R.items() if k != 'dsep_claims'}, indent=1)[:4000])