"""Unit tests for the stuck-at fault simulators (the SAF* classes).

Each fault is identified by ``line.index * 2 + polarity`` where polarity 0 is
stuck-at-0 and polarity 1 is stuck-at-1.  ``classify_faults`` returns a dict
with the detected-by-simulation set under ``'DS'`` and the not-observed set
under ``'NO'``.

The three simulators are exercised side-by-side:
  * ``SAFSimSimple``      - brute force, exact.
  * ``SAFSimIncremental`` - incremental cone re-simulation, exact.
  * ``SAFSimPPSFP``       - FFR-based, deliberately *optimistic*: its stem
    observability is hard-wired to "always observable" (see the FIXME in
    ``ppsfp.py``), so its DS set is a super-set of the exact DS set.
"""

import itertools

import numpy as np
import pytest

from kyupy import bench, logic
from kyupy.techlib import KYUPY
from fsim.static import FaultSet
from fsim.simple import SAFSimSimple
from fsim.incremental import SAFSimIncremental
from fsim.ppsfp import SAFSimPPSFP


# The two exact simulators must always agree; PPSFP only over-approximates.
ALL_ALGS = [SAFSimSimple, SAFSimIncremental, SAFSimPPSFP]


def make_patterns(*specs):
    """Build a 2D pattern array from per-pattern strings via ``logic.mvarray``.

    Each string holds one character per s-node, ordered as the inputs and
    outputs are defined in the netlist (output positions are don't-cares here,
    overwritten by simulation).  The result has shape ``(n_s_nodes, n_patterns)``
    as expected by the simulators; a lone pattern is reshaped to a single column.
    """
    pat = logic.mvarray(*specs)
    return pat if pat.ndim == 2 else pat.reshape(-1, 1)


def classify(alg, circuit, faults, patterns):
    # PPSFP requires the batch size to match the pattern count (as main.py does).
    sim = alg(circuit, patterns.shape[1])
    return sim.classify_faults(list(faults), patterns)


def saf(line, polarity):
    return line.index * 2 + polarity


# --------------------------------------------------------------------------- #
# Inline single-gate circuits with hand-computed expectations.
# --------------------------------------------------------------------------- #

@pytest.mark.parametrize('alg', ALL_ALGS)
def test_inv_single_pattern(alg):
    # s-nodes are 'i','o'.  i=0 => golden o=1.  Detect faults that flip o off 1.
    c = bench.parse('input(i) output(o) o=INV(i)')
    oline = c.forks['o'].ins[0]
    iline = oline.driver.ins[0]
    pat = make_patterns('0-')  # i=0 (o position is a don't-care)
    faults = [saf(oline, 0), saf(oline, 1), saf(iline, 0), saf(iline, 1)]
    r = classify(alg, c, faults, pat)
    # o s-a-0 forces 0 (!=1) -> detected; i s-a-1 forces i=1 -> o=0 -> detected.
    assert r['DS'] == {saf(oline, 0), saf(iline, 1)}
    # o s-a-1 already matches golden; i s-a-0 keeps i=0 -> nothing changes.
    assert r['NO'] == {saf(oline, 1), saf(iline, 0)}


@pytest.mark.parametrize('alg', ALL_ALGS)
def test_inv_both_polarities_detected(alg):
    # Applying both i=0 and i=1 sensitises every fault of the inverter.
    c = bench.parse('input(i) output(o) o=INV(i)')
    oline = c.forks['o'].ins[0]
    iline = oline.driver.ins[0]
    pat = make_patterns('00', '10')  # i=0 and i=1
    faults = [saf(oline, 0), saf(oline, 1), saf(iline, 0), saf(iline, 1)]
    r = classify(alg, c, faults, pat)
    assert r['DS'] == set(faults)
    assert r['NO'] == set()


@pytest.mark.parametrize('alg', ALL_ALGS)
def test_and2_output(alg):
    # s-nodes are 'i0','i1','o'.
    c = bench.parse('input(i0,i1) output(o) o=AND2(i0,i1)')
    oline = c.forks['o'].ins[0]
    o_sa0, o_sa1 = saf(oline, 0), saf(oline, 1)

    # i0=i1=1 => golden o=1: only o s-a-0 is observable.
    r = classify(alg, c, [o_sa0, o_sa1], make_patterns('110'))
    assert r['DS'] == {o_sa0}
    assert r['NO'] == {o_sa1}

    # i0=0,i1=1 => golden o=0: only o s-a-1 is observable.
    r = classify(alg, c, [o_sa0, o_sa1], make_patterns('010'))
    assert r['DS'] == {o_sa1}
    assert r['NO'] == {o_sa0}


@pytest.mark.parametrize('alg', ALL_ALGS)
def test_two_level_nand_propagation(alg):
    # o = NAND2(NAND2(a,b), c); s-nodes 'a','b','c','o'.  a=b=c=1: n=0, o=1.
    c = bench.parse('input(a,b,c) output(o) n=NAND2(a,b) o=NAND2(n,c)')
    nline = c.forks['n'].ins[0]
    oline = c.forks['o'].ins[0]
    faults = [saf(nline, 0), saf(nline, 1), saf(oline, 0), saf(oline, 1)]
    r = classify(alg, c, faults, make_patterns('1110'))
    # n already 0 (s-a-0 invisible); n s-a-1 -> o=0; o s-a-0 visible; o already 1.
    assert r['DS'] == {saf(nline, 1), saf(oline, 0)}
    assert r['NO'] == {saf(nline, 0), saf(oline, 1)}


# --------------------------------------------------------------------------- #
# c17 fixture: small, fully testable combinational benchmark.
# --------------------------------------------------------------------------- #

def _c17_exhaustive_patterns():
    """All 2**5 input combinations of c17.

    s-nodes are inputs 1,2,3,6,7 then outputs 22,23 (netlist order); the two
    output positions are don't-cares appended after the five input bits.
    """
    return make_patterns(*[''.join(combo) + '--'
                           for combo in itertools.product('01', repeat=5)])


@pytest.mark.parametrize('alg', ALL_ALGS)
def test_c17_fully_detected(alg, c17_bench, c17_resolved):
    # c17 has no redundant faults, so exhaustive patterns detect every fault.
    fs = FaultSet(c17_bench, KYUPY, c17_resolved)
    faults = list(fs.saf_equiv_classes.keys())
    pat = _c17_exhaustive_patterns()
    r = classify(alg, c17_resolved, faults, pat)
    assert r['DS'] == set(faults)
    assert r['NO'] == set()


@pytest.mark.parametrize('alg', ALL_ALGS)
def test_c17_output_fault_class(alg, c17_resolved):
    # All-zero inputs: output 22 settles to 0, so only its s-a-1 is observable.
    o22 = c17_resolved.forks['22'].ins[0]
    pat = make_patterns('00000--')  # all inputs 0
    r = classify(alg, c17_resolved, [saf(o22, 0), saf(o22, 1)], pat)
    assert r['DS'] == {saf(o22, 1)}
    assert r['NO'] == {saf(o22, 0)}


# --------------------------------------------------------------------------- #
# s27 fixture: exposes the exact-vs-optimistic difference between simulators.
# --------------------------------------------------------------------------- #

def _s27_setup(s27_bench, s27_resolved):
    fs = FaultSet(s27_bench, KYUPY, s27_resolved)
    faults = list(fs.saf_equiv_classes.keys())
    rng = np.random.default_rng(1)
    pat = rng.choice([logic.ZERO, logic.ONE],
                     size=(len(s27_resolved.s_nodes(KYUPY)), 64)).astype(np.uint8)
    return faults, pat


def test_s27_exact_simulators_agree(s27_bench, s27_resolved):
    faults, pat = _s27_setup(s27_bench, s27_resolved)
    assert len(faults) == 32  # collapsed fault count, see test_fault_set.test_s27

    simple = classify(SAFSimSimple, s27_resolved, faults, pat)
    incr = classify(SAFSimIncremental, s27_resolved, faults, pat)

    # The two exact simulators must classify every fault identically.
    assert simple['DS'] == incr['DS']
    assert simple['NO'] == incr['NO']
    # This deterministic pattern set leaves exactly one fault unobserved.
    assert len(simple['DS']) == 31
    assert len(simple['NO']) == 1


def test_s27_ppsfp_over_approximates(s27_bench, s27_resolved):
    faults, pat = _s27_setup(s27_bench, s27_resolved)

    exact = classify(SAFSimSimple, s27_resolved, faults, pat)
    ppsfp = classify(SAFSimPPSFP, s27_resolved, faults, pat)

    # PPSFP's optimistic stem observability can only add detections, never drop.
    assert exact['DS'] <= ppsfp['DS']
    # Here it reports the lone exact-NO fault as detected too.
    assert ppsfp['DS'] == exact['DS'] | exact['NO']


@pytest.mark.parametrize('alg', ALL_ALGS)
def test_ppsfp_multibranch_stem_observability(alg):
    # Stem `s` fans out to two AND gates; with sel1=0,sel2=1 the stem is
    # observable only through the o2 branch (the stem's *second* fanout line).
    c = bench.parse('input(a,b,sel1,sel2) output(o1,o2) '
                    's=NAND2(a,b) o1=AND2(s,sel1) o2=AND2(s,sel2)')
    a_sa0 = saf(c.forks['s'].ins[0].driver.ins[0], 0)
    # a=1,b=1,sel1=0,sel2=1: a s-a-0 flips s, propagating through o2 only.
    pat = make_patterns('1101--')
    r = classify(alg, c, [a_sa0], pat)
    assert r['DS'] == {a_sa0}