"""Worked examples for Lessons 20–23. Run from any directory; no files/network.

These are the video's worked rules, not answers to the independent checkpoints.
"""
from itertools import product
from pathlib import Path
import sys

# Locate this checkout rather than requiring a package installation or PYTHONPATH.
candidates = [Path.cwd(), *Path(__file__).resolve().parents]
REPOSITORY = next((path for path in candidates if (path / 'cm_build.py').is_file()), None)
if REPOSITORY is None:
    raise SystemExit('Run this script from the CM_Computation repository root (containing cm_build.py).')
sys.path.insert(0, str(REPOSITORY))

from bitset_backend import bitset_to_bool_array, eval_cm_node_bitset
from cm_build import compile_expr_to_cm, eval_cm_boolean
from cm_build_lazy import compile_expr_to_cm_lazy
from cm_exprlib import And, Not, Or, Var, Xor
from cm_ir import compile_expr, materialize_cm


def check_layouts(expression, scalar_rule):
    """Compare every lookup with a rule written directly from the input bits."""
    names = ["x0", "x1", "x2", "x3"]
    layouts = [
        (["x0", "x1"], ["x2", "x3"]),
        (["x0", "x2"], ["x1", "x3"]),
        (["x0", "x3"], ["x1", "x2"]),
    ]
    compiled = compile_expr(expression)
    first_matrix = None
    checks = 0
    for rows, columns in layouts:
        matrix = compile_expr_to_cm(expression, rows, columns, {})
        reused = materialize_cm(compiled.node, rows, columns, {})
        lazy_result = compile_expr_to_cm_lazy(expression, rows, columns, {})
        assert matrix.tolist() == reused.tolist() == lazy_result.tolist()
        if first_matrix is None:
            first_matrix = matrix
        for values in product((0, 1), repeat=4):
            assignment = dict(zip(names, values))
            expected = scalar_rule(*values)
            actual = eval_cm_boolean(matrix, rows, columns, assignment, {})
            assert actual == expected, (rows, columns, values, actual, expected)
            checks += 1
    return compiled, first_matrix, checks


def run_lab():
    # Var constructs an expression input; it does not set that input to 0 or 1.
    A, B, C, D = [Var(i) for i in range(4)]
    rule = Xor(And(A, B), Or(C, D))

    def rule_oracle(a_bit, b_bit, c_bit, d_bit):
        # Parentheses declare the rule independently of any matrix lookup.
        return (a_bit & b_bit) ^ (c_bit | d_bit)

    compiled, matrix, rule_checks = check_layouts(rule, rule_oracle)
    rows, columns = ["x0", "x1"], ["x2", "x3"]
    assignment = {"x0": 1, "x1": 1, "x2": 0, "x3": 0}
    lookup = eval_cm_boolean(matrix, rows, columns, assignment, {})
    assert lookup == 1  # 1100 has row index 3 and column index 0.
    packed = eval_cm_node_bitset(compiled.node, ["x0", "x1", "x2", "x3"])
    vector = bitset_to_bool_array(packed, n_vars=4).astype(int).tolist()
    assert vector == [rule_oracle(*values) for values in product((0, 1), repeat=4)]
    assert packed == 7918 and sum(vector) == 10

    # Lesson 23: authenticated AND (owner OR administrator) AND NOT suspended.
    A, O, D, S = [Var(i) for i in range(4)]
    policy = And(And(A, Or(O, D)), Not(S))

    def policy_oracle(a_bit, o_bit, d_bit, s_bit):
        return int(a_bit and (o_bit or d_bit) and not s_bit)

    _, policy_matrix, policy_checks = check_layouts(policy, policy_oracle)
    allowed = sum(policy_oracle(*values) for values in product((0, 1), repeat=4))
    assert allowed == 3
    return {
        "rule_matrix": matrix.astype(int).tolist(),
        "lookup_1100": lookup,
        "packed": packed,
        "accepted_rule_cases": sum(vector),
        "policy_matrix": policy_matrix.astype(int).tolist(),
        "policy_allowed": allowed,
        "rule_checks": rule_checks,
        "policy_checks": policy_checks,
    }


if __name__ == "__main__":
    result = run_lab()
    print("Rule matrix (AB rows / CD columns, false-first):")
    for row in result["rule_matrix"]:
        print(row)
    print(f"1100 -> row 3, column 0 -> {result['lookup_1100']}")
    print(f"Packed output: {result['packed']}; accepted cases: {result['accepted_rule_cases']}/16")
    print("Policy matrix (AO rows / DS columns, false-first):")
    for row in result["policy_matrix"]:
        print(row)
    print(f"Policy allowed: {result['policy_allowed']}/16")
    print(f"PASS: {result['rule_checks']} rule + {result['policy_checks']} policy = 96 scalar lookup checks")
    print("Direct, reused-compiled and current lazy-wrapper arrays agree in all three layouts.")