#!/usr/bin/env python3
"""Reproduce the explicit counterexample in Sections 3–4 of the supplied paper.

Run: python unstable_cores.py --output results
Dependencies: sympy, numpy, matplotlib. No network access or input data needed.
Exact algebra establishes the claims for this example; plots use floating point.
This does not run Lean or prove the general results discussed in Section 5.
"""
import argparse
import csv
import hashlib
import json
import platform
from collections import Counter
from itertools import combinations, product
from pathlib import Path

import sympy as sp

Q = sp.Rational
S = sp.Matrix([[-1, 0, 0, 1, 0], [0, -4, 0, 6, 0],
               [4, -2, -4, 1, 2], [-2, 0, 2, -2, 2]])
Y = sp.Matrix([[1, 0, 0, 0, 0], [0, 4, 0, 0, 0],
               [0, 2, 4, 0, 0], [2, 0, 0, 2, 0]])
FLUX = sp.Matrix([2, 3, 2, 2, 2])
E0 = Q(5187522222797, 46658000000)
F0 = Q(28200219703, 46658000000)
LAMBDA = Q(1, 500) + sp.I * Q(51, 500)
VECTOR = sp.Matrix([-152886431705 + 15563289455*sp.I,
                    23331004034 - 30699360512*sp.I,
                    7890610882 + 34396287629*sp.I, 1399740000])


def reactivity(e, f):
    """Rows are reactions; columns are species. Entries are rate derivatives."""
    return sp.Matrix([[1, 0, 0, e], [0, Q(1, 50), Q(1, 1000), 0],
                      [0, 0, Q(5, 11), 0], [0, 0, 0, f], [0, 0, 0, 0]])


def require(condition, explanation):
    # Unlike assert, these checks still run with python -O.
    if not condition:
        raise RuntimeError(explanation)


def subsets(n):
    for size in range(1, n + 1):
        yield from combinations(range(n), size)


def children():
    """Choose one distinct reactant reaction for each selected species."""
    choices = [[j for j in range(5) if Y[i, j] > 0] for i in range(4)]
    found = []
    for species in subsets(4):
        for reactions in product(*(choices[i] for i in species)):
            if len(set(reactions)) == len(reactions):
                found.append((species, reactions))
    return found


def is_restriction(child, parent):
    assignments = dict(zip(*parent))
    return all(assignments.get(i) == j for i, j in zip(*child))


def child_matrix(child):
    return S.extract(*child)


def matrix_strings(a):
    return [[str(x) for x in row] for row in a.tolist()]


def write_csv(path, fields, rows):
    with path.open('w', newline='', encoding='utf-8') as stream:
        writer = csv.DictWriter(stream, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)


def exact_checks():
    log = []
    def say(message):
        log.append(message)

    r = reactivity(E0, F0)
    g = S*r
    require(all(x >= 0 and x.is_integer for x in Y) and
            all(x >= 0 and x.is_integer for x in S+Y), 'Invalid reaction data')
    require(all(Y[i, j]*(S+Y)[i, j] == 0 for i in range(4) for j in range(5)),
            'A catalyst was introduced')
    require(all((r[j, i] > 0) if Y[i, j] > 0 else (r[j, i] == 0)
                for i in range(4) for j in range(5)), 'Incorrect derivative support')
    require(all(x > 0 for x in FLUX) and S*FLUX == sp.zeros(4, 1), 'Not an equilibrium')
    require(S[:, :4].det() == 48 and S[:, :4]*r[:4, :] == g,
            'Inflow check failed')
    say('PASS: nonnegative integer reaction data; no catalysts; correct derivative support.')
    say('PASS: all equilibrium reaction rates are positive; S*(2,3,2,2,2) = (0,0,0,0).')
    say('PASS: the constant inflow supplies balance and contributes zero to the Jacobian.')

    # Differentiate the actual nonlinear rates, not a finite-difference surrogate.
    x = sp.symbols('x0:4', positive=True)
    rates = sp.Matrix([FLUX[j]*sp.prod(x[i]**(r[j, i]/FLUX[j])
                                     for i in range(4)) for j in range(5)])
    at_one = dict.fromkeys(x, 1)
    require(rates.subs(at_one) == FLUX, 'Wrong rates at the equilibrium')
    require((S*rates).jacobian(x).subs(at_one) == g, 'Wrong nonlinear Jacobian')
    say('PASS: exact differentiation of the generalized power-law rates gives G = S*R.')
    residual = (g*VECTOR-LAMBDA*VECTOR).applyfunc(sp.simplify)
    require(VECTOR != sp.zeros(4, 1) and residual == sp.zeros(4, 1), 'Eigenpair failed')
    say('PASS: G*v - (1/500 + 51*i/500)*v = (0,0,0,0), exactly.')
    z = sp.Symbol('z')
    chi = g.charpoly(z).as_expr()
    target = ((125000*z**2-500*z+1301)*
              (513238000*z**2+116236430203*z+13005481800)/64154750000000)
    require(sp.expand(chi-target) == 0, 'Characteristic factorization failed')
    coefficients = sp.Poly(chi, z).all_coeffs()[1:]
    a, b, c, d = coefficients
    delta2, delta3 = a*b-c, a*b*c-a*a*d-c*c
    require(all(v > 0 for v in coefficients) and delta2 > 0 and delta3 < 0,
            'Quartic sign checks failed')
    say('PASS: all four polynomial coefficients are positive, but the third Hurwitz determinant is negative.')

    all_children = children()
    counts = dict(sorted(Counter(len(k[0]) for k in all_children).items()))
    require(counts == {1: 6, 2: 11, 3: 6, 4: 1}, 'Child enumeration failed')
    maximal = [k for k in all_children if not any(k != p and is_restriction(k, p)
                                                  for p in all_children)]
    identity = ((0, 1, 2, 3), (0, 1, 2, 3))
    child_a = ((0, 2, 3), (0, 1, 3))
    child_b = ((1, 2, 3), (1, 2, 0))
    child_2 = ((2, 3), (1, 0))
    require(set(maximal) == {identity, child_a, child_b, child_2}, 'Wrong maximal selections')
    say('PASS: 24 selections: 6 of size 1, 11 of size 2, 6 of size 3, 1 of size 4.')
    say('PASS: exactly four maximal selections; every selection belongs to at least one.')

    # For each weighted symmetric matrix, check ALL principal minors exactly.
    # Their nonnegativity proves positive semidefiniteness, including zero cases.
    covers = [
        ('four-species', identity, [100, 35, 40, 140]),
        ('three-species A', child_a, [10, 2, 7]),
        ('two-species', child_2, [1, 2]),
        ('face 1,3', ((1, 3), (1, 0)), [1, 2]),
        ('face 2,3', ((2, 3), (2, 0)), [1, 2]),
    ]
    certificates = []
    for name, child, weights in covers:
        p = sp.diag(*weights)
        a_child = child_matrix(child)
        m = -(p*a_child+a_child.T*p)
        minors = [m.extract(i, i).det() for i in subsets(m.rows)]
        require(all(v >= 0 for v in minors), 'Invalid weighted certificate: '+name)
        certificates.append(dict(name=name, matrix=matrix_strings(m), weights=weights,
                                 all_principal_minors=list(map(str, minors))))
    d0, d1, d2 = sp.symbols('d0 d1 d2', positive=True)
    scaled_b = child_matrix(child_b)*sp.diag(d0, d1, d2)
    factor_b = z*(z+4*d0)*(z+4*d1+2*d2)
    require(sp.expand(scaled_b.charpoly(z).as_expr()-factor_b) == 0,
            'Exceptional factorization failed')
    say('PASS: all five weighted certificates have nonnegative principal minors.')
    say('PASS: exceptional polynomial = z*(z+4*d0)*(z+4*d1+2*d2) for symbolic positive d0,d1,d2.')
    p0, p1, p2 = sp.symbols('p0 p1 p2', positive=True)
    pb = sp.diag(p0, p1, p2)
    mb = -(pb*child_matrix(child_b)+child_matrix(child_b).T*pb)
    require(sp.expand(mb[1:3, 1:3].det()+(4*p1-2*p2)**2) == 0,
            'Exceptional weight obstruction failed')
    require(sp.expand(mb.det().subs(p2, 2*p1)+32*p1**3) == 0,
            'Exceptional determinant obstruction failed')
    say('PASS: the exceptional child has no positive diagonal weight certificate; its exact polynomial is needed.')

    records = []
    for number, child in enumerate(all_children, 1):
        if child == child_b:
            method, label, weights_text = 'symbolic cubic factorization', 'one zero eigenvalue', ''
        else:
            matches = [(name, parent, weights) for name, parent, weights in covers
                       if is_restriction(child, parent)]
            require(bool(matches), 'Uncovered selection: '+str(child))
            name, parent, weights = matches[0]
            local_weights = [weights[parent[0].index(i)] for i in child[0]]
            a_child = child_matrix(child)
            p = sp.diag(*local_weights)
            m = -(p*a_child+a_child.T*p)
            require(all(m.extract(i, i).det() >= 0 for i in subsets(m.rows)),
                    'Restricted certificate failed')
            # The only singular weighted case has roots 0 and -(4*d0+2*d1).
            if a_child.det() == 0:
                require(sp.expand((a_child*sp.diag(d0, d1)).charpoly(z).as_expr()
                                  -z*(z+4*d0+2*d1)) == 0, 'Singular face failed')
                label = 'one zero eigenvalue'
            else:
                require(all(m[:i, :i].det() > 0 for i in range(1, m.rows+1)),
                        'Strict negativity certificate failed')
                label = 'all real parts strictly negative'
            method, weights_text = 'weighted certificate: '+name, str(local_weights)
        records.append(dict(number=number, size=len(child[0]),
                            assignment='; '.join(f'X{i}->r{j}' for i, j in zip(*child)),
                            matrix=str(child_matrix(child).tolist()), certificate=method,
                            weights=weights_text, every_positive_scaling=label))
    require(sum(r['every_positive_scaling'] == 'one zero eigenvalue' for r in records) == 2,
            'Unexpected boundary count')
    say('PASS: all 24 selections covered for EVERY positive column scaling: 22 strictly negative; 2 with a zero eigenvalue.')

    # The same network with two varying derivatives gives a useful parameter map.
    e, f = sp.symbols('E F', positive=True)
    symbolic_g = S*reactivity(e, f)
    a, b, c, d = symbolic_g.charpoly(z).all_coeffs()[1:]
    h2, h3 = sp.expand(a*b-c), sp.expand(a*b*c-a*a*d-c*c)
    require(all(v > 0 for expr in [a, b, c, d, h2]
                for v in sp.Poly(expr, e, f).coeffs()), 'Positive coefficient check failed')
    # Thus in this slice: h3>0 means strictly stable; h3<0 means unstable.
    # At h3=0 a purely imaginary pair occurs; no nonlinear cycle is asserted.
    presets = []
    for name, ev, fv in [('decaying oscillations', sp.Integer(10), Q(3, 5)),
                        ('paper example', E0, F0)]:
        h = h3.subs({e: ev, f: fv})
        require(h > 0 if name.startswith('decaying') else h < 0, 'Preset failed')
        presets.append(dict(name=name, E=str(ev), F=str(fv), delta3=str(h),
                            behavior='decay' if h > 0 else 'growth'))
    boundary = sp.Poly(h3.subs(f, Q(3, 5)), e)
    positive_roots = [v for v in sp.nroots(boundary) if abs(sp.im(v)) < 1e-12 and sp.re(v) > 0]
    require(len(positive_roots) == 1, 'Unexpected boundary at F=0.6')
    say(f'PARAMETER MAP: at F=0.6 the boundary is E approximately {float(sp.re(positive_roots[0])):.9f}.')
    say('SCOPE: exact counterexample checked in SymPy; figures are numerical; Lean was not run.')
    report = dict(S=matrix_strings(S), Y=matrix_strings(Y), R=matrix_strings(r),
                  G=matrix_strings(g), rates=[str(v) for v in rates],
                  E=str(E0), F=str(F0), eigenvalue=str(LAMBDA),
                  eigenvector=[str(v) for v in VECTOR], residual=[str(v) for v in residual],
                  characteristic_polynomial=str(sp.factor(chi)),
                  coefficients=list(map(str, coefficients)), delta2=str(delta2), delta3=str(delta3),
                  child_counts=counts, certificates=certificates, children=records,
                  parameter_coefficients=list(map(str, [a, b, c, d])),
                  parameter_delta2=str(h2), parameter_delta3=str(h3), presets=presets,
                  boundary_at_F_0_6=str(positive_roots[0]),
                  proof_scope='Sections 3–4: this explicit counterexample; no Lean execution')
    return report, log, (e, f, symbolic_g, h3)


def make_figures(out, symbolic):
    import numpy as np
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    from matplotlib.colors import TwoSlopeNorm, LinearSegmentedColormap
    plt.rcParams.update({'figure.facecolor': '#f7f3e9', 'axes.facecolor': '#f7f3e9',
                         'text.color': '#24282b', 'axes.labelcolor': '#24282b',
                         'axes.spines.top': False, 'axes.spines.right': False,
                         'font.size': 11, 'axes.titlesize': 14, 'svg.fonttype': 'none'})
    colors = ['#ad4c35', '#187b83', '#76558c', '#526984']
    def save(fig, name):
        fig.savefig(out/(name+'.png'), dpi=180, bbox_inches='tight')
        fig.savefig(out/(name+'.svg'), bbox_inches='tight')
        plt.close(fig)

    # This is an exact linear solution evaluated numerically, not an ODE solver.
    t = np.linspace(0, 600, 2401)
    vector = np.array(VECTOR, dtype=complex).ravel()
    vector /= np.max(np.abs(vector))
    mode = np.real(np.exp(complex(LAMBDA)*t[:, None])*vector[None, :])
    fig, axes = plt.subplots(1, 2, figsize=(12, 4.4), layout='constrained')
    for i in range(4):
        axes[0].plot(t, mode[:, i], color=colors[i], label=f'X{i}', lw=1.3)
    axes[0].plot(t, np.exp(t/500), '--', color='#555555', lw=1, label='growth envelope')
    axes[0].plot(t, -np.exp(t/500), '--', color='#555555', lw=1)
    axes[0].set(title='Small deviations oscillate and grow', xlabel='Time (model units)',
                ylabel='Deviation / chosen initial scale')
    axes[0].legend(ncol=3, fontsize=9, loc='upper left')
    axes[1].plot(mode[:, 0], mode[:, 1], color=colors[0], lw=1.2)
    axes[1].scatter(mode[[0, -1], 0], mode[[0, -1], 1], c=['#187b83', '#24282b'], s=40)
    axes[1].annotate('start', (mode[0, 0], mode[0, 1]), xytext=(6, 8), textcoords='offset points')
    axes[1].annotate('end', (mode[-1, 0], mode[-1, 1]), xytext=(6, 8), textcoords='offset points')
    axes[1].set(title='The same motion, viewed in two species', xlabel='X0 deviation / initial scale',
                ylabel='X1 deviation / initial scale')
    fig.suptitle('Linear motion near the equilibrium: amplitude doubles in 346.6 time units', fontsize=15)
    save(fig, 'growing_oscillations')
    write_csv(out/'linear_mode.csv', ['time', 'X0', 'X1', 'X2', 'X3'],
              [dict(zip(['time', 'X0', 'X1', 'X2', 'X3'], [tt, *row])) for tt, row in zip(t, mode)])

    e, f, symbolic_g, h3 = symbolic
    evalues, fvalues = np.linspace(1, 250, 250), np.linspace(.02, 2, 199)
    ee, ff = np.meshgrid(evalues, fvalues)
    # Build a stacked matrix array to obtain actual numerical eigenvalues at every point.
    matrices = np.zeros(ee.shape+(4, 4))
    for i in range(4):
        for j in range(4):
            matrices[..., i, j] = sp.lambdify((e, f), symbolic_g[i, j], 'numpy')(ee, ff)
    eigenvalues = np.linalg.eigvals(matrices)
    growth = eigenvalues.real.max(axis=-1)
    h = sp.lambdify((e, f), h3, 'numpy')(ee, ff)
    away = (np.abs(growth) > 1e-8) & (np.abs(h) > 1e-6)
    require(np.all((growth[away] > 0) == (h[away] < 0)), 'Map disagrees with polynomial test')
    require(np.all(np.max(np.abs(eigenvalues.imag), axis=-1)[growth > 1e-8] > 0),
            'Unexpected nonoscillatory instability')
    fig, ax = plt.subplots(figsize=(10, 5.4), layout='constrained')
    cmap = LinearSegmentedColormap.from_list('decay_growth', ['#187b83', '#f7f3e9', '#b34632'])
    mesh = ax.pcolormesh(ee, ff, growth, shading='auto', cmap=cmap,
                         norm=TwoSlopeNorm(vmin=-.015, vcenter=0, vmax=.015), rasterized=True)
    ax.contour(ee, ff, h, levels=[0], colors='#24282b', linewidths=1.4)
    ax.scatter([float(E0)], [float(F0)], marker='*', s=160, color='#24282b', zorder=4)
    ax.annotate('Paper example', (float(E0), float(F0)), xytext=(12, 12), textcoords='offset points')
    ax.scatter([10], [.6], marker='s', s=40, color='#24282b')
    ax.annotate('Decay preset', (10, .6), xytext=(12, -20), textcoords='offset points')
    ax.text(45, 1.7, 'DECAY', weight='bold')
    ax.text(190, .45, 'GROWTH', weight='bold')
    ax.set(xlabel='E: sensitivity of reaction r0 to species X3',
           ylabel='F: sensitivity of reaction r3 to species X3',
           title='All 24 matrices in the proposed test pass throughout this map')
    fig.colorbar(mesh, ax=ax, extend='both', label='Largest eigenvalue real part (1 / model time)')
    save(fig, 'parameter_map')
    write_csv(out/'parameter_map.csv', ['E', 'F', 'largest_real_part', 'delta3_numerical'],
              [dict(E=ev, F=fv, largest_real_part=gv, delta3_numerical=hv)
               for ev, fv, gv, hv in zip(ee.ravel(), ff.ravel(), growth.ravel(), h.ravel())])

    g = np.array(S*reactivity(E0, F0), dtype=float)
    eig = np.linalg.eigvals(g)
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.2), layout='constrained')
    for ax in axes:
        ax.axvline(0, color='#555555', lw=1)
        ax.axhline(0, color='#bbbbbb', lw=.7)
        ax.scatter(eig.real, eig.imag, c=[colors[0] if z.real > 0 else colors[1] for z in eig], s=55)
        ax.set(xlabel='Real part: negative = decay; positive = growth', ylabel='Imaginary part')
    axes[0].set(title='All four eigenvalues', xlim=(-240, 10), ylim=(-.14, .14))
    axes[1].set(title='Close view of the growth boundary', xlim=(-.14, .022), ylim=(-.14, .14))
    axes[1].annotate('0.002 + 0.102 i', (.002, .102), xytext=(-100, 8), textcoords='offset points')
    axes[1].annotate('0.002 - 0.102 i', (.002, -.102), xytext=(-100, -15), textcoords='offset points')
    save(fig, 'eigenvalues')
    return dict(numpy=np.__version__, matplotlib=matplotlib.__version__,
                map_points=int(ee.size), map_sign_check='passed away from numerical boundary',
                eigenvalues=[[float(z.real), float(z.imag)] for z in eig],
                period=float(2*np.pi/(51/500)), doubling_time=float(500*np.log(2)))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', type=Path, default=Path('results'))
    parser.add_argument('--exact-only', action='store_true', help='Skip figures; only SymPy is needed')
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)
    report, log, symbolic = exact_checks()
    if not args.exact_only:
        numeric = make_figures(args.output, symbolic)
        report['numerical_outputs'] = numeric
        log.append(f"FIGURES: {numeric['map_points']} parameter points; polynomial/eigenvalue sign check passed.")
        log.append(f"LINEAR MODE: period {numeric['period']:.6f}; amplitude doubling time {numeric['doubling_time']:.6f}.")
        log.append('NUMERICAL EIGENVALUES: '+str(numeric['eigenvalues']))
    report['environment'] = dict(python=platform.python_version(), sympy=sp.__version__)
    report['source_sha256'] = hashlib.sha256(Path(__file__).read_bytes()).hexdigest()
    (args.output/'verification.json').write_text(json.dumps(report, indent=2)+'\n', encoding='utf-8')
    write_csv(args.output/'children.csv', list(report['children'][0]), report['children'])
    transcript = '\n'.join(log)+'\n'
    (args.output/'console.txt').write_text(transcript, encoding='utf-8')
    print(transcript, end='')


if __name__ == '__main__':
    main()