# quantum_plane_timed_stress_examples.sage
#
# Timed examples for the quantum-plane skew-polynomial implementation
# over K(t)[x; sigma], K=Q(omega), sigma(t)=omega*t.
#
# Uses implementation of quantum_plane_factorization.sage
#
# Run:
#     sage quantum_plane_timed_stress_examples.sage
#
# Optional controls:
#     QP_LEVEL=core|stress|all|xl     default: all
#     QP_ONLY=C1,I1,S3                comma-separated example ids
#     QP_VERBOSE=1                    pass verbose=True into factorizations
#     QP_PRINT_POLYS=1                print full polynomials/factors
#     QP_SUMMARY_CSV=path.csv         write timing summary CSV
#     QP_IMPL=/path/to/impl.sage      choose implementation file
#
# The examples are organized around the paper's stages:
#   - Algorithm 4.1: central bound factorization and right-gcd splitting.
#   - Corollary 3.2 / norm tests: exact irreducibility certificates.
#   - Algorithms 6.3, 7.2, 8.1: certifying modular search for
#     irreducible-bound pieces, followed here by exact sparse certification
#     for constructed toy cases when the modular local factor does not lift.
#
# Every returned factorization is checked by exact multiplication in
# K(t)[x; sigma].  The irreducible-bound stage remains certifying, not
# completeness-proved, exactly as in the paper.

import os
import sys
import time
import csv
import traceback
from itertools import product

# -----------------------------------------------------------------------------
# System configuration cruft
# -----------------------------------------------------------------------------

try:
    _HERE = os.path.dirname(os.path.abspath(__file__))
except NameError:
    _HERE = os.getcwd()

_IMPL_NAMES = []
if os.environ.get("QP_IMPL", None):
    _IMPL_NAMES.append(os.environ["QP_IMPL"])
_IMPL_NAMES += [
    "quantum_plane_factorization.sage",
]

_IMPL_PATH = None
_TRIED_IMPL_PATHS = []
for _name in _IMPL_NAMES:
    _candidates = [_name]
    if not os.path.isabs(_name):
        _candidates.insert(0, os.path.join(_HERE, _name))
    for _path in _candidates:
        if _path in _TRIED_IMPL_PATHS:
            continue
        _TRIED_IMPL_PATHS.append(_path)
        if os.path.exists(_path):
            _IMPL_PATH = _path
            break
    if _IMPL_PATH is not None:
        break

if _IMPL_PATH is None:
    raise RuntimeError("Could not find the implementation file. Put "
                       "quantum_plane_full_stack_solution_patch.sage in this directory, "
                       "or set QP_IMPL=/path/to/implementation.sage")

print("Loading implementation:", _IMPL_PATH)
load(_IMPL_PATH)


# -----------------------------------------------------------------------------
# Runtime controls
# -----------------------------------------------------------------------------

_LEVEL = os.environ.get("QP_LEVEL", "all").strip().lower()
if _LEVEL not in ["core", "stress", "all", "xl"]:
    raise ValueError("QP_LEVEL must be core, stress, all, or xl")

_ONLY_ENV = os.environ.get("QP_ONLY", "").strip()
_ONLY = None
if _ONLY_ENV:
    _ONLY = set([z.strip() for z in _ONLY_ENV.split(",") if z.strip()])

_VERBOSE = os.environ.get("QP_VERBOSE", "0") not in ["", "0", "false", "False", "no", "No"]
_PRINT_POLYS = os.environ.get("QP_PRINT_POLYS", "0") not in ["", "0", "false", "False", "no", "No"]
_SUMMARY_CSV = os.environ.get("QP_SUMMARY_CSV", os.path.join(_HERE, "quantum_plane_timing_summary.csv"))

_RESULTS = []


# -----------------------------------------------------------------------------
# Helper functions
# -----------------------------------------------------------------------------

def banner(title):
    print("\n" + "=" * 96)
    print(title)
    print("=" * 96)


def subbanner(title):
    print("\n" + "-" * 96)
    print(title)
    print("-" * 96)


def skew_product(Q, factors):
    out = Q.one()
    for h in factors:
        out = out * h
    return out


def maybe_print(label, obj):
    if _PRINT_POLYS:
        print(label, obj)


def degree_multiset(factors):
    return sorted([ZZ(h.degree()) for h in factors])


def degrees_string(factors_or_degrees):
    if factors_or_degrees is None:
        return ""
    vals = []
    for z in factors_or_degrees:
        try:
            vals.append(ZZ(z.degree()))
        except Exception:
            vals.append(ZZ(z))
    return "+".join([str(v) for v in sorted(vals)])


def bound_shape(Q, P):
    try:
        fac = P.factor()
        pieces = []
        for item in fac:
            pieces.append("deg%s^%s" % (ZZ(item[0].degree()), ZZ(item[1])))
        return " * ".join(pieces)
    except Exception:
        try:
            return "deg%s" % ZZ(P.degree())
        except Exception:
            return "unknown"


def compact_bound(P, max_len=90):
    s = str(P)
    if len(s) <= max_len:
        return s
    return s[:max_len-3] + "..."


def print_short_report(Q, f, P, factors=None, extra=None):
    print("deg_x(f) =", f.degree())
    print("bound degree in X =", P.degree())
    print("bound shape =", bound_shape(Q, P))
    print("bound =", compact_bound(P))
    print("bound irreducible?", Q.central_is_irreducible(P))
    if factors is not None:
        print("factor degrees =", degrees_string(factors))
        print("verified product equals input:", Q.verify_product(factors, f))
    if extra:
        for k in sorted(extra.keys()):
            print("%s = %s" % (k, extra[k]))


def assert_degrees(factors, expected):
    if expected is None:
        return
    got = degree_multiset(factors)
    exp = sorted([ZZ(d) for d in expected])
    if got != exp:
        raise AssertionError("degree multiset mismatch: got %s, expected %s" % (got, exp))


def m3_fast_modular_options(Q3, interpolation_degree_bound=0, pbound=5):
    # Deliberately small: enough to exercise good T-specialization / inert-prime
    # construction without wandering into huge F_{p^6} coefficient fields.
    return {
        "interpolation_degree_bound": interpolation_degree_bound,
        "specializations": [Q3.K(2) + Q3.omega],
        "min_specializations": 1,
        "primes_per_specialization": 1,
        "prime_search_bound": pbound,
        "min_crt_primes": 1,
        "height_bound": 5,
        "max_lambda_cardinality": 20000,
        "max_modular_candidates_per_prime": 120,
        "max_exact_specialized_candidates": 80,
        "max_crt_states": 500,
        "max_uniform_factor_trials": 0,
        "max_factorization_iterator": 0,
        "max_right_irreducible_divisors": 0,
        "verbose_modular": _VERBOSE,
    }


def sparse_recursive_factor(Q, f, e, coeff_bound=1, poly_degree=1,
                            include_ratios=False, verbose=False):
    # Exact recursive sparse right-factor search of a fixed expected degree e.
    f = f.monic()
    e = ZZ(e)
    if f.degree() == e:
        return [f]
    h, q = Q.find_sparse_right_factor_degree(f, e,
                                             coeff_bound=coeff_bound,
                                             poly_degree=poly_degree,
                                             include_ratios=include_ratios,
                                             verbose=verbose)
    if h is None:
        raise RuntimeError("sparse recursive factorization found no degree-%s right factor" % e)
    return (sparse_recursive_factor(Q, q.monic(), e,
                                    coeff_bound=coeff_bound,
                                    poly_degree=poly_degree,
                                    include_ratios=include_ratios,
                                    verbose=verbose)
            + [h.monic()])


def central_solve(Q, factors, expected_degrees=None):
    f = skew_product(Q, factors).monic()
    P = Q.bound(f)
    rough = Q.central_decomposition(f, verbose=_VERBOSE)
    if not Q.verify_product(rough, f):
        raise AssertionError("central decomposition product check failed")
    final = Q.factor(f, use_modular=False, verbose=_VERBOSE)
    if not Q.verify_product(final, f):
        raise AssertionError("final central-only product check failed")
    assert_degrees(final, expected_degrees)
    return f, P, rough, final


def full_stack_solve(Q, f, modular_options, coeff_bound=1, poly_degree=1,
                     include_ratios=False, expected_degrees=None):
    P = Q.bound(f)
    rough = Q.central_decomposition(f, verbose=_VERBOSE)
    if not Q.verify_product(rough, f):
        raise AssertionError("rough central product check failed")
    final = Q.factor(f,
                     use_modular=True,
                     modular_options=modular_options,
                     coeff_bound=coeff_bound,
                     poly_degree=poly_degree,
                     include_ratios=include_ratios,
                     verbose=_VERBOSE)
    if not Q.verify_product(final, f):
        raise AssertionError("final product check failed")
    assert_degrees(final, expected_degrees)
    return P, rough, final


def sparse_bound_poly_solve(Q, pi, h, e, expected_count=3,
                            coeff_bound=1, poly_degree=1):
    F = Q.central_to_skew(pi)
    P = Q.bound(F)
    if P != pi.monic():
        raise AssertionError("bound mismatch: got %s expected %s" % (P, pi.monic()))
    if h**expected_count != F:
        raise AssertionError("constructed power identity h^%s == pi failed" % expected_count)
    rough = Q.central_decomposition(F, verbose=_VERBOSE)
    if not Q.verify_product(rough, F):
        raise AssertionError("central decomposition product check failed")
    factors = sparse_recursive_factor(Q, F, e=e,
                                      coeff_bound=coeff_bound,
                                      poly_degree=poly_degree,
                                      include_ratios=False,
                                      verbose=_VERBOSE)
    if not Q.verify_product(factors, F):
        raise AssertionError("sparse factorization product check failed")
    assert_degrees(factors, [e] * expected_count)
    return F, P, rough, factors


def run_example(example):
    ex_id = example["id"]
    if _ONLY is not None and ex_id not in _ONLY:
        return
    level = example.get("level", "core")
    if _LEVEL == "core" and level != "core":
        return
    if _LEVEL == "stress" and level not in ["stress"]:
        return
    if _LEVEL == "all" and level == "xl":
        return
    # _LEVEL == xl runs all examples, including xl.

    banner("%s: %s" % (ex_id, example["title"]))
    print("stage:", example.get("stage", ""))
    print("shape:", example.get("shape", ""))
    if example.get("notes", ""):
        print("notes:", example.get("notes", ""))

    start = time.time()
    row = {
        "id": ex_id,
        "level": level,
        "title": example["title"],
        "m": example.get("m", ""),
        "deg_x": example.get("deg_x", ""),
        "bound_deg_X": example.get("bound_deg_X", ""),
        "bound_shape": example.get("bound_shape", ""),
        "stage": example.get("stage", ""),
        "expected_factor_degrees": example.get("expected_factor_degrees", ""),
        "status": "STARTED",
        "seconds": "",
        "message": "",
    }
    try:
        Q, f, P, factors, extra = example["runner"]()
        elapsed = time.time() - start
        row["status"] = "OK"
        row["seconds"] = "%.3f" % elapsed
        row["m"] = str(getattr(Q, "m", row["m"]))
        row["deg_x"] = str(ZZ(f.degree()))
        row["bound_deg_X"] = str(ZZ(P.degree()))
        row["bound_shape"] = bound_shape(Q, P)
        row["expected_factor_degrees"] = degrees_string(factors)
        row["message"] = "verified"
        _RESULTS.append(row)
        print_short_report(Q, f, P, factors=factors, extra=extra)
        maybe_print("f =", f)
        if _PRINT_POLYS:
            for i, h in enumerate(factors, 1):
                print("  factor[%s] = %s" % (i, h))
                print("      bound =", Q.bound(h))
        print("elapsed seconds: %.3f" % elapsed)
    except Exception as err:
        elapsed = time.time() - start
        row["status"] = "FAIL"
        row["seconds"] = "%.3f" % elapsed
        row["message"] = str(err).replace("\n", " | ")
        _RESULTS.append(row)
        print("FAILED after %.3f seconds" % elapsed)
        print("error:", err)
        traceback.print_exc()


# -----------------------------------------------------------------------------
# Example definitions
# -----------------------------------------------------------------------------

def make_examples():
    examples = []

    def add(ex_id, level, title, m, deg_x, bound_deg_X, shape, stage,
            expected_factor_degrees, runner, notes=""):
        examples.append({
            "id": ex_id,
            "level": level,
            "title": title,
            "m": str(m),
            "deg_x": str(deg_x),
            "bound_deg_X": str(bound_deg_X),
            "bound_shape": shape,
            "stage": stage,
            "expected_factor_degrees": expected_factor_degrees,
            "runner": runner,
            "shape": shape,
            "notes": notes,
        })

    # ------------------------------------------------------------------
    # CORE examples: similar in size to the previous suite.
    # ------------------------------------------------------------------

    def C1():
        Q = QuantumPlane(2)
        x, t = Q.x(), Q.t
        f, P, rough, final = central_solve(Q, [x - t, x - (t + 1)], [1, 1])
        return Q, f, P, final, {"rough_degrees": degrees_string(rough)}
    add("C1", "core", "m=2 degree-2 central split into two linear factors",
        2, 2, 2, "central bound product of two linear factors", "central only",
        "1+1", C1)

    def C2():
        Q = QuantumPlane(2)
        x, t = Q.x(), Q.t
        factors = [x - (t + Q.F(c)) for c in [0, 1, 2, 3, 4]]
        f, P, rough, final = central_solve(Q, factors, [1, 1, 1, 1, 1])
        return Q, f, P, final, {"rough_degrees": degrees_string(rough)}
    add("C2", "core", "m=2 degree-5 central split into five linear factors",
        2, 5, 5, "five distinct linear central factors", "central only",
        "1+1+1+1+1", C2)

    def C3():
        Q = QuantumPlane(3)
        x, t = Q.x(), Q.t
        factors = [x - (t + Q.F(c)) for c in [0, 1, 2, 3, 4]]
        f, P, rough, final = central_solve(Q, factors, [1, 1, 1, 1, 1])
        return Q, f, P, final, {"rough_degrees": degrees_string(rough)}
    add("C3", "core", "m=3 degree-5 central split into five linear factors",
        3, 5, 5, "five factors X-(T+c^3)", "central only",
        "1+1+1+1+1", C3)

    def I1():
        Q = QuantumPlane(3)
        x, t = Q.x(), Q.t
        h = x**2 - t
        g = h**2
        opts = m3_fast_modular_options(Q, interpolation_degree_bound=0, pbound=5)
        P, rough, final = full_stack_solve(Q, g, opts,
                                           coeff_bound=1, poly_degree=1,
                                           include_ratios=False,
                                           expected_degrees=[2, 2])
        return Q, g, P, final, {"rough_degrees": degrees_string(rough),
                                "constructed_factor": "x^2 - t"}
    add("I1", "core", "m=3 irreducible-bound example g=(x^2-t)^2",
        3, 4, 2, "irreducible bound X^2-T", "modular attempt + exact sparse certification",
        "2+2", I1)

    def I2():
        Q = QuantumPlane(3)
        x, t = Q.x(), Q.t
        h = x**2 - t * (t**3 + 1)
        g = h**2
        opts = m3_fast_modular_options(Q, interpolation_degree_bound=1, pbound=5)
        P, rough, final = full_stack_solve(Q, g, opts,
                                           coeff_bound=1, poly_degree=4,
                                           include_ratios=False,
                                           expected_degrees=[2, 2])
        return Q, g, P, final, {"rough_degrees": degrees_string(rough),
                                "constructed_factor": "x^2 - t*(T+1)"}
    add("I2", "core", "m=3 T-dependent irreducible-bound square",
        3, 4, 2, "irreducible bound X^2 - T*(T+1)^3",
        "modular attempt + exact sparse certification", "2+2", I2,
        notes="The factor x^2-t*(T+1) has nonconstant dependence on the central parameter T.")

    def B1():
        Q = QuantumPlane(3)
        x, t = Q.x(), Q.t
        X, T = Q.X, Q.T
        pi = Q.RX(X**2 - T)
        h = x**2 - t
        F, P, rough, final = sparse_bound_poly_solve(Q, pi, h, e=2, expected_count=3,
                                                     coeff_bound=1, poly_degree=1)
        return Q, F, P, final, {"rough_degrees": degrees_string(rough),
                                "identity": "X^2-T=(x^2-t)^3"}
    add("B1", "core", "factor irreducible central bound polynomial X^2-T",
        3, 6, 2, "irreducible central polynomial, skew degree 6", "central irreducible + sparse exact",
        "2+2+2", B1)

    # ------------------------------------------------------------------
    # STRESS examples: larger but designed to remain practical.
    # ------------------------------------------------------------------

    def S1():
        Q = QuantumPlane(3)
        x, t = Q.x(), Q.t
        factors = [x - (t + Q.F(c)) for c in range(8)]
        f, P, rough, final = central_solve(Q, factors, [1] * 8)
        return Q, f, P, final, {"rough_degrees": degrees_string(rough),
                                "linear_factors": "x-(t+c), c=0..7"}
    add("S1", "stress", "m=3 degree-8 central split into eight linear factors",
        3, 8, 8, "eight distinct central linear factors", "central only",
        "+".join(["1"] * 8), S1,
        notes="This mainly stresses bound computation and repeated right-gcd splitting.")

    def S2():
        Q = QuantumPlane(3)
        x, t = Q.x(), Q.t
        # Mixed central split: degree-4 irreducible-bound piece, degree-2
        # T-dependent irreducible-bound piece, and six linear pieces.
        h4 = x**4 - t
        h2T = x**2 - t * (t**3 + 1)
        linears = [x - (t + Q.F(c)) for c in [2, 3, 4, 5, 6, 7]]
        factors = [h4, h2T] + linears
        f, P, rough, final = central_solve(Q, factors, [4, 2, 1, 1, 1, 1, 1, 1])
        return Q, f, P, final, {"rough_degrees": degrees_string(rough),
                                "constructed_degrees": "4+2+1+1+1+1+1+1"}
    add("S2", "stress", "m=3 degree-12 mixed central decomposition",
        3, 12, 12, "irreducible central factors of degrees 4,2 plus six linear factors",
        "central only + degree-bound tests", "4+2+1+1+1+1+1+1", S2,
        notes="This is central-stage-complete but not just linear factors.")

    def S3():
        Q = QuantumPlane(5)
        x, t = Q.x(), Q.t
        factors = [x - (t + Q.F(c)) for c in range(6)]
        f, P, rough, final = central_solve(Q, factors, [1] * 6)
        return Q, f, P, final, {"rough_degrees": degrees_string(rough),
                                "constant_field": "Q(zeta_5)"}
    add("S3", "stress", "m=5 degree-6 central split over Q(zeta_5)",
        5, 6, 6, "six distinct central linear factors over K(T)", "central only",
        "+".join(["1"] * 6), S3,
        notes="This increases the constant-field degree to phi(5)=4.")

    def S4():
        Q = QuantumPlane(3)
        x, t = Q.x(), Q.t
        h = x**4 - t
        g = h**2
        opts = m3_fast_modular_options(Q, interpolation_degree_bound=0, pbound=5)
        P, rough, final = full_stack_solve(Q, g, opts,
                                           coeff_bound=1, poly_degree=1,
                                           include_ratios=False,
                                           expected_degrees=[4, 4])
        return Q, g, P, final, {"rough_degrees": degrees_string(rough),
                                "constructed_factor": "x^4 - t"}
    add("S4", "stress", "m=3 degree-8 irreducible-bound square (x^4-t)^2",
        3, 8, 4, "irreducible bound X^4-T", "modular attempt + exact sparse certification",
        "4+4", S4,
        notes="Larger version of I1; finite-field budget is capped at p=5.")

    def S5():
        Q = QuantumPlane(3)
        x, t = Q.x(), Q.t
        X, T = Q.X, Q.T
        pi = Q.RX(X**5 - T)
        h = x**5 - t
        F, P, rough, final = sparse_bound_poly_solve(Q, pi, h, e=5, expected_count=3,
                                                     coeff_bound=1, poly_degree=1)
        return Q, F, P, final, {"rough_degrees": degrees_string(rough),
                                "identity": "X^5-T=(x^5-t)^3"}
    add("S5", "stress", "m=3 skew degree-15: factor irreducible bound X^5-T",
        3, 15, 5, "irreducible central polynomial X^5-T", "central irreducible + sparse exact",
        "5+5+5", S5,
        notes="This pushes the irreducible-bound polynomial example to degree 15.")

    # ------------------------------------------------------------------
    # XL optional examples.  These are not run under QP_LEVEL=all.
    # ------------------------------------------------------------------

    def X1():
        Q = QuantumPlane(3)
        x, t = Q.x(), Q.t
        factors = [x - (t + Q.F(c)) for c in range(12)]
        f, P, rough, final = central_solve(Q, factors, [1] * 12)
        return Q, f, P, final, {"rough_degrees": degrees_string(rough),
                                "linear_factors": "x-(t+c), c=0..11"}
    add("X1", "xl", "XL m=3 degree-12 central split into twelve linear factors",
        3, 12, 12, "twelve distinct central linear factors", "central only",
        "+".join(["1"] * 12), X1)

    def X2():
        Q = QuantumPlane(3)
        x, t = Q.x(), Q.t
        X, T = Q.X, Q.T
        pi = Q.RX(X**7 - T)
        h = x**7 - t
        F, P, rough, final = sparse_bound_poly_solve(Q, pi, h, e=7, expected_count=3,
                                                     coeff_bound=1, poly_degree=1)
        return Q, F, P, final, {"rough_degrees": degrees_string(rough),
                                "identity": "X^7-T=(x^7-t)^3"}
    add("X2", "xl", "XL m=3 skew degree-21: factor irreducible bound X^7-T",
        3, 21, 7, "irreducible central polynomial X^7-T", "central irreducible + sparse exact",
        "7+7+7", X2,
        notes="Optional heavier example; run with QP_LEVEL=xl or QP_ONLY=X2.")

    return examples


# -----------------------------------------------------------------------------
# Summary output
# -----------------------------------------------------------------------------

def print_summary():
    banner("TIMING SUMMARY")
    if not _RESULTS:
        print("No examples were selected. Check QP_LEVEL or QP_ONLY.")
        return

    cols = ["id", "level", "m", "deg_x", "bound_deg_X", "expected_factor_degrees",
            "seconds", "status", "stage"]
    widths = {}
    for c in cols:
        widths[c] = max(len(c), max([len(str(r.get(c, ""))) for r in _RESULTS]))
    header = "  ".join([c.ljust(widths[c]) for c in cols])
    print(header)
    print("  ".join(["-" * widths[c] for c in cols]))
    for r in _RESULTS:
        print("  ".join([str(r.get(c, "")).ljust(widths[c]) for c in cols]))

    total = sum([float(r["seconds"]) for r in _RESULTS if r.get("seconds", "")])
    print("\nTotal wall time for selected examples: %.3f seconds" % total)

    print("\nDetailed shapes:")
    for r in _RESULTS:
        print("  {id}: {title}".format(**r))
        print("      shape: m={m}, deg_x={deg_x}, bound_deg_X={bound_deg_X}, factors={expected_factor_degrees}".format(**r))
        print("      bound shape: %s" % r.get("bound_shape", ""))
        print("      time/status: %s sec, %s" % (r.get("seconds", ""), r.get("status", "")))
        if r.get("message", "") not in ["", "verified"]:
            print("      message: %s" % r.get("message", ""))

    try:
        with open(_SUMMARY_CSV, "w", newline="") as fh:
            fieldnames = ["id", "level", "title", "m", "deg_x", "bound_deg_X",
                          "bound_shape", "stage", "expected_factor_degrees",
                          "status", "seconds", "message"]
            writer = csv.DictWriter(fh, fieldnames=fieldnames)
            writer.writeheader()
            for r in _RESULTS:
                writer.writerow({k: r.get(k, "") for k in fieldnames})
        print("\nWrote CSV summary to:", _SUMMARY_CSV)
    except Exception as err:
        print("\nCould not write CSV summary:", err)


def run_all_examples():
    print("Quantum-plane timed example suite")
    print("Implementation:", _IMPL_PATH)
    print("QP_LEVEL =", _LEVEL)
    if _ONLY is not None:
        print("QP_ONLY =", sorted(_ONLY))
    print("QP_VERBOSE =", _VERBOSE)
    print("QP_PRINT_POLYS =", _PRINT_POLYS)

    examples = make_examples()
    for ex in examples:
        run_example(ex)

    print_summary()

    failed = [r for r in _RESULTS if r.get("status") != "OK"]
    if failed:
        raise RuntimeError("%s selected example(s) failed; see timing summary above" % len(failed))


if __name__ == "__main__":
    run_all_examples()