#!/usr/bin/env python3
"""Small independent log-density controls for rho(t)=t**n on [0,1].

No stellar solve or FEM library is used. A three-coordinate exponential moment
fit illustrates a changed density approximation, not a production proposal.
Polynomial extrema are analytic; integrals use independently generated Gauss
rules, with 64/128-point comparison. Only stdlib is required.
"""
import argparse
import csv
from fractions import Fraction
import hashlib
import json
import math
from pathlib import Path
import unittest


def dot(a, b):
    return math.fsum(x * y for x, y in zip(a, b))


def norm(a):
    return math.sqrt(dot(a, a))


def gauss(count):
    rule = []
    for j in range(1, count + 1):
        z = math.cos(math.pi * (j - .25) / (count + .5))
        for _ in range(40):
            p0, p1 = 1., z
            for degree in range(2, count + 1):
                p0, p1 = p1, ((2 * degree - 1) * z * p1 - (degree - 1) * p0) / degree
            derivative = count * (z * p1 - p0) / (z * z - 1)
            change = p1 / derivative
            z -= change
            if abs(change) < 2e-16:
                break
        rule.append(((1 + z) / 2, 1 / ((1 - z * z) * derivative * derivative)))
    return sorted(rule)


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


def interpolate(values):
    separation = math.sqrt(.6)
    slope = (values[2] - values[0]) / separation
    curvature = (10 / 3) * (values[2] + values[0] - 2 * values[1])
    return [values[1] - slope / 2 + curvature / 4, slope - curvature, curvature]


def project(n):
    n = Fraction(str(n))
    moments = [1 / (n + j + 1) for j in range(3)]
    return [float(sum(a * b for a, b in zip(row, moments)))
            for row in ((9, -36, 30), (-36, 192, -180), (30, -180, 180))]


def extrema(c):
    points = [0., 1.]
    if c[2] and 0 < -c[1] / (2 * c[2]) < 1:
        points.append(-c[1] / (2 * c[2]))
    return min(polynomial(c, t) for t in points), max(polynomial(c, t) for t in points)


def basis(t):
    return [1., math.sqrt(3) * (2 * t - 1), math.sqrt(5) * (6 * t * t - 6 * t + 1)]


def orthogonal_coefficients(c):
    return [c[0] + c[1] / 2 + c[2] / 3, (c[1] + c[2]) / (2 * math.sqrt(3)), c[2] / (6 * math.sqrt(5))]


def power_coefficients(a):
    return [a[0] - math.sqrt(3) * a[1] + math.sqrt(5) * a[2],
            2 * math.sqrt(3) * a[1] - 6 * math.sqrt(5) * a[2], 6 * math.sqrt(5) * a[2]]


def target_moments(n):
    return [1 / (n + 1), math.sqrt(3) * (2 / (n + 2) - 1 / (n + 1)),
            math.sqrt(5) * (6 / (n + 3) - 6 / (n + 2) + 1 / (n + 1))]


def moments_and_jacobian(density, rule):
    samples = [(weight * density(t), basis(t)) for t, weight in rule]
    moments = [math.fsum(weight * phi[i] for weight, phi in samples) for i in range(3)]
    jacobian = [[math.fsum(weight * phi[i] * phi[j] for weight, phi in samples) for j in range(3)] for i in range(3)]
    return moments, jacobian


def solve(matrix, rhs):
    a = [row[:] + [value] for row, value in zip(matrix, rhs)]
    for column in range(3):
        pivot = max(range(column, 3), key=lambda i: abs(a[i][column]))
        a[column], a[pivot] = a[pivot], a[column]
        if not abs(a[column][column]) > 1e-16:
            raise ArithmeticError("Unresolved tiny moment Jacobian")
        for row in range(column + 1, 3):
            factor = a[row][column] / a[column][column]
            for j in range(column, 4):
                a[row][j] -= factor * a[column][j]
    result = [0.] * 3
    for row in (2, 1, 0):
        result[row] = (a[row][3] - dot(a[row][row + 1:3], result[row + 1:])) / a[row][row]
    if norm([dot(row, result) - value for row, value in zip(matrix, rhs)]) > 1e-12 * norm(rhs):
        raise ArithmeticError("Tiny solve failed independent defect check")
    return result


def determinant(a):
    return (a[0][0] * (a[1][1] * a[2][2] - a[1][2] * a[2][1])
            - a[0][1] * (a[1][0] * a[2][2] - a[1][2] * a[2][0])
            + a[0][2] * (a[1][0] * a[2][1] - a[1][1] * a[2][0]))


def eigenvalues(a):
    # Eigenvalues of the three-by-three symmetric weighted mass matrix itself;
    # no Gram-squared surrogate or full stellar conditioning claim is involved.
    mean = sum(a[i][i] for i in range(3)) / 3
    centered = [[a[i][j] - (mean if i == j else 0) for j in range(3)] for i in range(3)]
    scale = math.sqrt(sum(value * value for row in centered for value in row) / 6)
    if scale == 0:
        return [mean] * 3
    angle = math.acos(max(-1., min(1., determinant([[x / scale for x in row] for row in centered]) / 2))) / 3
    values = sorted(mean + 2 * scale * math.cos(angle + j * 2 * math.pi / 3) for j in range(3))
    if abs(sum(values) - 3 * mean) > 1e-12 or abs(math.prod(values) - determinant(a)) > 1e-12:
        raise ArithmeticError("Tiny spectral invariants failed")
    return values


def fit_log_moments(n, initial, rule):
    target = target_moments(n)
    a = orthogonal_coefficients(initial)
    for iteration in range(40):
        moments, jacobian = moments_and_jacobian(lambda t: math.exp(dot(a, basis(t))), rule)
        residual = [value - wanted for value, wanted in zip(moments, target)]
        if norm(residual) < 2e-13 * norm(target):
            return power_coefficients(a), iteration, norm(residual) / norm(target)
        direction = solve(jacobian, [-value for value in residual])
        old_norm = norm(residual)
        for backtrack in range(20):
            step = 2. ** -backtrack
            trial = [value + step * change for value, change in zip(a, direction)]
            try:
                trial_moments, _ = moments_and_jacobian(lambda t: math.exp(dot(trial, basis(t))), rule)
                new_norm = norm([value - wanted for value, wanted in zip(trial_moments, target)])
            except OverflowError:
                continue
            if new_norm < (1 - 1e-4 * step) * old_norm:
                a = trial
                break
        else:
            raise ArithmeticError("Tiny moment fit failed line search")
    raise ArithmeticError("Tiny moment fit did not converge")


def controls(n):
    nodes = [t for t, _ in gauss(3)]
    direct = interpolate([t ** n for t in nodes])
    logs = interpolate([n * math.log(t) for t in nodes])
    log_mass = moments_and_jacobian(lambda t: math.exp(polynomial(logs, t)), gauss(64))[0][0]
    shifted = logs[:]
    shifted[0] += math.log((1 / (n + 1)) / log_mass)
    fitted, iterations, relative = fit_log_moments(n, logs, gauss(64))
    return [("density_q2_interpolation_or_exp_coefficients", "polynomial", direct, None),
            ("density_q2_l2_projection", "polynomial", project(n), None),
            ("pointwise_exp_q2_log_interpolation", "exp", logs, None),
            ("pointwise_exp_q2_log_mass_shift", "exp", shifted, None),
            ("pointwise_exp_q2_log_moment_fit", "exp", fitted,
             {"iterations": iterations, "relative_moment_defect_64": relative}),
            ("matched_factor_t_power_n_exp_zero", "factor", [0., 0., 0.], None)]


def evaluate_case(n, name, kind, coefficients, fit):
    if kind == "polynomial":
        density = lambda t: polynomial(coefficients, t)
        minimum, maximum = extrema(coefficients)
    elif kind == "exp":
        density = lambda t: math.exp(polynomial(coefficients, t))
        lo, hi = extrema(coefficients)
        minimum, maximum = math.exp(lo), math.exp(hi)
    else:
        density = lambda t: t ** n
        minimum, maximum = 0., 1.
    moments, jacobian = moments_and_jacobian(density, gauss(128))
    coarse_moments, _ = moments_and_jacobian(density, gauss(64))
    residual = [a - b for a, b in zip(moments, target_moments(n))]
    error = math.sqrt(math.fsum(w * (density(t) - t ** n) ** 2 for t, w in gauss(128)))
    eig = eigenvalues(jacobian) if kind in ("exp", "factor") else None
    return dict(n=n, name=name, kind=kind, coefficients=coefficients,
                boundary_density=density(0.), minimum_density=minimum, maximum_density=maximum,
                cell_mass=moments[0], relative_mass_error=moments[0] * (n + 1) - 1,
                relative_l2_error=error * math.sqrt(2 * n + 1),
                moment_defect=residual, relative_moment_defect=norm(residual) / norm(target_moments(n)),
                moment_quadrature_64_vs_128_maximum=max(abs(a - b) for a, b in zip(moments, coarse_moments)),
                log_coordinate_jacobian_eigenvalues=eig,
                log_coordinate_jacobian_condition=eig[-1] / eig[0] if eig else None,
                fit=fit)


class Tests(unittest.TestCase):
    def test_gauss(self):
        for degree in range(12):
            self.assertAlmostEqual(math.fsum(w * t ** degree for t, w in gauss(8)), 1 / (degree + 1), places=14)

    def test_log_coefficients_do_not_fix_polynomial(self):
        for n in (1.5, 3.):
            nodes = [t for t, _ in gauss(3)]
            a = interpolate([math.exp(n * math.log(t)) for t in nodes])
            b = interpolate([t ** n for t in nodes])
            self.assertLess(max(abs(x - y) for x, y in zip(a, b)), 2e-15)
            self.assertLess(extrema(a)[0], 0)

    def test_fit_and_vacuum(self):
        for n in (1., 1.5, 3.):
            results = [evaluate_case(n, *case) for case in controls(n)]
            fitted = next(row for row in results if row["name"].endswith("moment_fit"))
            self.assertLess(fitted["relative_moment_defect"], 3e-13)
            self.assertGreater(fitted["boundary_density"], 0)
            self.assertGreater(fitted["relative_l2_error"], 1e-3)
            exact = results[-1]
            self.assertEqual(exact["boundary_density"], 0)
            self.assertEqual(exact["relative_l2_error"], 0)

    def test_jacobian(self):
        a, direction, step = [-1., .2, -.1], [.3, -.1, .2], 1e-5
        rule = gauss(64)
        _, jacobian = moments_and_jacobian(lambda t: math.exp(dot(a, basis(t))), rule)
        plus, _ = moments_and_jacobian(lambda t: math.exp(dot([x + step * y for x, y in zip(a, direction)], basis(t))), rule)
        minus, _ = moments_and_jacobian(lambda t: math.exp(dot([x - step * y for x, y in zip(a, direction)], basis(t))), rule)
        error = [(x - y) / (2 * step) - dot(row, direction) for x, y, row in zip(plus, minus, jacobian)]
        self.assertLess(norm(error), 2e-10)
        self.assertEqual(eigenvalues([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]), [1., 1., 1.])
        self.assertLess(max(abs(x - y) for x, y in zip(eigenvalues([[2., 1., 0.], [1., 2., 0.], [0., 0., 4.]]), [1., 3., 4.])), 2e-14)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test:
        result = unittest.TextTestRunner(verbosity=2).run(unittest.defaultTestLoader.loadTestsFromTestCase(Tests))
        if not result.wasSuccessful():
            raise SystemExit(1)
        if args.output is None:
            return
    cases = [evaluate_case(n, *case) for n in (1., 1.5, 3.) for case in controls(n)]
    data = dict(schema="log_density_reference_v1", producer_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
                scope="Uniform 1D frozen t^n profile; no FEM, coupled star solve, virial test, or general-EOS validation.",
                eigenvalue_scope="Three-dimensional log-coordinate weighted mass matrix in L2-orthonormal polynomial basis on [0,1].",
                scaling="For rho=A*Delta^n*t^n, local reference-basis Jacobian=A*Delta^(n+1)*H_n; relative to ordinary mass the scale is A*Delta^n.",
                cases=cases)
    if args.output is None:
        print(json.dumps(data, indent=2, allow_nan=False))
        return
    args.output.mkdir(parents=True, exist_ok=False)
    with (args.output / "summary.csv").open("w", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=list(cases[0]))
        writer.writeheader()
        for row in cases:
            writer.writerow({key: json.dumps(value) if isinstance(value, (list, dict)) else value for key, value in row.items()})
    with (args.output / "profiles.csv").open("w", newline="") as stream:
        writer = csv.writer(stream)
        writer.writerow(["n", "method", "t", "target", "density"])
        for row in cases:
            for j in range(401):
                t = j / 400
                value = (t ** row["n"] if row["kind"] == "factor" else
                         math.exp(polynomial(row["coefficients"], t)) if row["kind"] == "exp" else
                         polynomial(row["coefficients"], t))
                writer.writerow([row["n"], row["name"], t, t ** row["n"], value])
    data["artifacts_sha256"] = {path.name: hashlib.sha256(path.read_bytes()).hexdigest() for path in sorted(args.output.iterdir())}
    (args.output / "summary.json").write_text(json.dumps(data, indent=2, allow_nan=False) + "\n")
    print("Wrote", args.output)


if __name__ == "__main__":
    main()