#!/usr/bin/env python3
"""Independent surface-cell controls; no FEM state, solver, or limiter is changed.

For inward distance s in [0, width], assume h=kappa*s and rho=A*s**n.
Compare exact, uniformly weighted Q2 L2 projection with interpolation at MFEM's
default L2 Q2 Gauss-Legendre nodes. In one dimension Q2 means degree <= 2.
Extrema, zeros, and negative integrals use quadratic algebra, not plot sampling.

The normalized model is t**n, t=s/width. Physical density is
A*width**n*p(t). Curved metrics, tangential variation, higher terms in h,
and changes in kappa under refinement are deliberately outside this control.

Example:
  python negative_density_reference_model.py --self-test
  python negative_density_reference_model.py --output /tmp/surface-control --png

Only the optional PNG writer requires matplotlib; all arithmetic is stdlib.
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
from decimal import Decimal, localcontext
from fractions import Fraction
from pathlib import Path
import unittest


INDICES = (Fraction(1), Fraction(3, 2), Fraction(3))
INVERSE_MOMENTS = ((9, -36, 30), (-36, 192, -180), (30, -180, 180))
METHODS = ("exact_l2_projection", "gauss_legendre_interpolation")


def decimal(value):
    if isinstance(value, Fraction):
        return Decimal(value.numerator) / Decimal(value.denominator)
    return Decimal(str(value))


def nodes():
    r = (Decimal(3) / 5).sqrt()
    return ((1 - r) / 2, Decimal("0.5"), (1 + r) / 2)


def exact_value(t, n):
    return t * t.sqrt() if n == Fraction(3, 2) else t ** int(n)


def projection(n):
    """Invert the degree-two Hilbert moment matrix with rational arithmetic."""
    moments = [1 / (n + j + 1) for j in range(3)]
    return tuple(sum(a * b for a, b in zip(row, moments))
                 for row in INVERSE_MOMENTS)


def interpolation(n):
    x = nodes()
    values = [exact_value(t, n) for t in x]
    slope = (values[2] - values[0]) / (x[2] - x[0])
    curvature = (values[2] + values[0] - 2 * values[1]) * 10 / 3
    return (values[1] - slope / 2 + curvature / 4,
            slope - curvature, curvature)


def coefficients(n, method):
    if method == METHODS[0] or n in (1, 3):
        # For n=3 the cubic nodal polynomial is the shifted Legendre P3;
        # hence its remainder is also the exact L2 projection. Using that
        # identity avoids classifying a rounded zero curvature as an n=1 root.
        return tuple(decimal(c) for c in projection(n))
    return interpolation(n)


def value(c, t):
    return c[0] + t * (c[1] + t * c[2])


def integral(c, left=Decimal(0), right=Decimal(1)):
    return sum(a * (right ** (j + 1) - left ** (j + 1)) / (j + 1)
               for j, a in enumerate(c))


def analyze(c):
    a, b, q = c
    if q == 0:
        roots = [] if b == 0 else [-a / b]
    else:
        discriminant = b * b - 4 * a * q
        roots = ([] if discriminant < 0 else
                 sorted({(-b - discriminant.sqrt()) / (2 * q),
                         (-b + discriminant.sqrt()) / (2 * q)}))
    unit_roots = [t for t in roots if 0 <= t <= 1]
    extrema = [Decimal(0), Decimal(1)]
    if q != 0 and 0 < -b / (2 * q) < 1:
        extrema.append(-b / (2 * q))
    minimum_at = min(extrema, key=lambda t: value(c, t))
    bounds = sorted({Decimal(0), Decimal(1), *unit_roots})
    negative = [(left, right) for left, right in zip(bounds, bounds[1:])
                if value(c, (left + right) / 2) < 0]
    return {
        "roots_in_cell": unit_roots,
        "minimum_at": minimum_at,
        "minimum": value(c, minimum_at),
        "negative_intervals": negative,
        "negative_width_fraction": sum((b - a for a, b in negative), Decimal(0)),
        "negative_mass": -sum((integral(c, a, b) for a, b in negative), Decimal(0)),
        "cell_mean": integral(c),
        "boundary_value": a,
        "interior_endpoint_value": value(c, Decimal(1)),
        "nonnegative_on_closed_cell": value(c, minimum_at) >= 0,
    }


def as_numbers(item):
    if isinstance(item, (Decimal, Fraction)):
        result = float(item)
        if not math.isfinite(result):
            raise ValueError("Nonfinite derived output")
        return result
    if isinstance(item, dict):
        return {k: as_numbers(v) for k, v in item.items()}
    if isinstance(item, (tuple, list)):
        return [as_numbers(v) for v in item]
    return item


def case(n, method):
    c = coefficients(n, method)
    result = analyze(c)
    exact_mean = 1 / (decimal(n) + 1)
    if n == 1:
        exact_roots, exact_minimum = "0", "0"
    elif n == 3:
        exact_roots, exact_minimum = "1/5-sqrt(6)/30, 1/5+sqrt(6)/30", "-1/100"
    elif method == METHODS[0]:
        exact_roots, exact_minimum = "(sqrt(174)-12)/30", "-2/105"
    else:
        exact_roots = "(-c1+sqrt(c1^2-4*c0*c2))/(2*c2); c_j from the algebraic interpolation formula"
        exact_minimum = "c0 from the algebraic interpolation formula"
    result.update({
        "n": n, "method": method, "coefficients_c0_c1_c2": c,
        "coefficients_decimal": [str(v) for v in c],
        "exact_projection_coefficients": [str(v) for v in projection(n)],
        "roots_in_cell_exact": exact_roots, "minimum_exact": exact_minimum,
        "exact_cell_mean": exact_mean,
        "mean_error": result["cell_mean"] - exact_mean,
        "negative_mass_over_exact_cell_mass": result["negative_mass"] / exact_mean,
        "values_at_gl_nodes": [value(c, t) for t in nodes()],
        "exact_values_at_gl_nodes": [exact_value(t, n) for t in nodes()],
    })
    return result


def refinement_row(result, width, amplitude):
    n = result["n"]
    scale = amplitude * exact_value(width, n)
    return {
        "n": n, "method": result["method"], "width": width,
        "amplitude_A": amplitude, "minimum_density": scale * result["minimum"],
        "minimum_distance": width * result["minimum_at"],
        "negative_width": width * result["negative_width_fraction"],
        "negative_mass_per_unit_area": scale * width * result["negative_mass"],
        "exact_mass_per_unit_area": scale * width * result["exact_cell_mean"],
        "negative_mass_over_exact_cell_mass": result["negative_mass_over_exact_cell_mass"],
    }


class ReferenceChecks(unittest.TestCase):
    def test_exact_moments(self):
        for n in INDICES:
            c = projection(n)
            for j in range(3):
                self.assertEqual(sum(a / (i + j + 1) for i, a in enumerate(c)),
                                 1 / (n + j + 1))

    def test_known_coefficients(self):
        self.assertEqual(projection(Fraction(1)), (0, 1, 0))
        self.assertEqual(projection(Fraction(3, 2)),
                         (Fraction(-2, 105), Fraction(16, 35), Fraction(4, 7)))
        self.assertEqual(projection(Fraction(3)),
                         (Fraction(1, 20), Fraction(-3, 5), Fraction(3, 2)))

    def test_interpolation_and_integer_identity(self):
        for n in INDICES:
            c = interpolation(n)
            for t in nodes():
                self.assertLess(abs(value(c, t) - exact_value(t, n)), Decimal("1e-55"))
            if n in (1, 3):
                for actual, expected in zip(c, coefficients(n, METHODS[1])):
                    self.assertLess(abs(actual - expected), Decimal("1e-55"))

    def test_positive_dofs_do_not_certify_polynomial(self):
        for n in (Fraction(3, 2), Fraction(3)):
            result = case(n, METHODS[1])
            self.assertTrue(all(v > 0 for v in result["values_at_gl_nodes"]))
            self.assertLess(result["minimum"], 0)

    def test_roots_and_extrema(self):
        for method in METHODS:
            self.assertTrue(case(Fraction(1), method)["nonnegative_on_closed_cell"])
            for n in (Fraction(3, 2), Fraction(3)):
                result = case(n, method)
                c = coefficients(n, method)
                for t in result["roots_in_cell"]:
                    self.assertLess(abs(value(c, t)), Decimal("1e-55"))
                for left, right in result["negative_intervals"]:
                    self.assertLess(value(c, (left + right) / 2), 0)
                self.assertGreater(result["negative_mass"], 0)
        cubic = case(Fraction(3), METHODS[0])
        self.assertEqual(cubic["minimum_at"], Decimal("0.2"))
        self.assertEqual(cubic["minimum"], Decimal("-0.01"))
        self.assertGreater(cubic["boundary_value"], 0)
        self.assertEqual(cubic["roots_in_cell"],
                         [Decimal(".2") - Decimal(6).sqrt() / 30,
                          Decimal(".2") + Decimal(6).sqrt() / 30])

    def test_gauss_mean_and_nonconservation_of_interpolation(self):
        weights = (Decimal(5) / 18, Decimal(4) / 9, Decimal(5) / 18)
        for n in INDICES:
            mean = sum(w * exact_value(t, n) for w, t in zip(weights, nodes()))
            self.assertLess(abs(mean - integral(interpolation(n))), Decimal("1e-55"))
        self.assertLess(case(Fraction(3, 2), METHODS[1])["mean_error"], Decimal("-1e-4"))

    def test_refinement_scales(self):
        for n in (Fraction(3, 2), Fraction(3)):
            for method in METHODS:
                result = case(n, method)
                coarse = refinement_row(result, Decimal(1), Decimal(7))
                fine = refinement_row(result, Decimal(".5"), Decimal(7))
                factor = exact_value(Decimal(".5"), n)
                for key, expected in (("minimum_density", factor),
                                      ("negative_width", Decimal(".5")),
                                      ("negative_mass_per_unit_area", factor / 2)):
                    self.assertLess(abs(fine[key] / coarse[key] - expected), Decimal("1e-55"))


def write_csv(path, rows):
    rows = [as_numbers(row) for row in rows]
    with path.open("w", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=list(rows[0]))
        writer.writeheader()
        for row in rows:
            writer.writerow({k: json.dumps(v) if isinstance(v, (list, dict)) else v
                             for k, v in row.items()})


def plot(output, results, samples):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    fig, axes = plt.subplots(2, 3, figsize=(12, 6), constrained_layout=True)
    ts = [i / (samples - 1) for i in range(samples)]
    for column, n in enumerate(INDICES):
        for row in range(2):
            ax = axes[row, column]
            ax.plot(ts, [t ** float(n) for t in ts], "k", label="Exact density")
            for result in results:
                if result["n"] != n:
                    continue
                c = [float(v) for v in result["coefficients_c0_c1_c2"]]
                ax.plot(ts, [value(c, t) for t in ts],
                        "--" if result["method"] == METHODS[1] else "-",
                        label="GL interpolation" if result["method"] == METHODS[1] else "L2 projection")
            ax.axhline(0, color="0.7", linewidth=0.8)
            ax.scatter([float(t) for t in nodes()],
                       [float(exact_value(t, n)) for t in nodes()], s=18, color="k")
            ax.set_title(f"n = {float(n):g}" + ("; near-surface detail" if row else ""))
            ax.set_xlabel("Inward distance / cell width")
            if row:
                ax.set_xlim(0, 0.35)
                ax.set_ylim(-0.027, 0.075 if n == 3 else 0.23)
            if column == 0:
                ax.set_ylabel("Density / (A width^n)")
    axes[0, 0].legend(fontsize=8)
    fig.suptitle("Uniform 1D surface-cell model: positive nodal data can yield a negative polynomial")
    fig.savefig(output / "profiles.png", dpi=170)
    plt.close(fig)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, help="Fresh directory for JSON and CSV artifacts")
    parser.add_argument("--self-test", action="store_true")
    parser.add_argument("--png", action="store_true", help="Also plot profiles (requires matplotlib)")
    parser.add_argument("--widths", nargs="+", default=["1", ".5", ".25", ".125"])
    parser.add_argument("--amplitude", default="1", help="A in rho=A*s^n")
    parser.add_argument("--samples", type=int, default=501, help="Plot/CSV samples, not extrema checks")
    args = parser.parse_args()
    widths, amplitude = [Decimal(x) for x in args.widths], Decimal(args.amplitude)
    if any(not x.is_finite() or x <= 0 for x in (*widths, amplitude)):
        parser.error("Widths and amplitude must be finite and positive")
    if args.samples < 2 or (args.png and args.output is None):
        parser.error("At least two samples are required; --png requires --output")
    with localcontext() as context:
        context.prec = 70
        if args.self_test:
            test_result = unittest.TextTestRunner(verbosity=2).run(
                unittest.defaultTestLoader.loadTestsFromTestCase(ReferenceChecks))
            if not test_result.wasSuccessful():
                raise SystemExit(1)
            if args.output is None:
                return
        results = [case(n, method) for n in INDICES for method in METHODS]
        report = {
            "schema": "negative_density_reference_model_v1",
            "producer_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
            "model": "h=kappa*s; rho=A*s^n; t=s/width in [0,1]; uniform physical weight",
            "scope": "Analytic mechanism control, not a fit or accuracy certificate for the 3D star",
            "arithmetic": "Exact rational projection; 70-digit decimal algebraic interpolation/roots",
            "gl_nodes": nodes(), "gl_nodes_exact": "(1-sqrt(3/5))/2, 1/2, (1+sqrt(3/5))/2",
            "node_source": "MFEM L2_FECollection defaults to BasisType::GaussLegendre; Q2 tensor element",
            "projection_derivation": "c=H^-1 b; H_ij=1/(i+j+1); b_i=1/(n+i+1), i,j=0,1,2",
            "interpolation_derivation": "v_j=t_j^n; b=(v_+-v_-)/sqrt(3/5); c=(10/3)*(v_++v_--2v_0); p=v_0+b*(t-.5)+c*(t-.5)^2",
            "negative_interval_convention": "Interior points are strictly negative; endpoints are roots except a negative physical-cell endpoint",
            "refinement": "rho=A*width^n*p(t); negative mass/area scales width^(n+1); relative first-cell negativity need not disappear",
            "cases": results,
        }
        if args.output is None:
            print(json.dumps(as_numbers(report), indent=2, allow_nan=False))
            return
        args.output.mkdir(parents=True, exist_ok=False)
        write_csv(args.output / "summary.csv", results)
        write_csv(args.output / "refinement.csv",
                  [refinement_row(result, width, amplitude) for result in results for width in widths])
        profiles = []
        for result in results:
            for i in range(args.samples):
                t = Decimal(i) / (args.samples - 1)
                profiles.append({"n": result["n"], "method": result["method"], "t": t,
                                 "exact": exact_value(t, result["n"]),
                                 "approximation": value(result["coefficients_c0_c1_c2"], t)})
        write_csv(args.output / "profiles.csv", profiles)
        if args.png:
            plot(args.output, results, args.samples)
        report["artifacts_sha256"] = {p.name: hashlib.sha256(p.read_bytes()).hexdigest()
                                      for p in sorted(args.output.iterdir())}
        (args.output / "summary.json").write_text(json.dumps(as_numbers(report), indent=2, allow_nan=False) + "\n")
        print(f"Wrote six analytic controls to {args.output}")


if __name__ == "__main__":
    main()