"""Feedback can select different compositions with the same activity signature.

Run: python example.py --output outputs
The reaction model and numerical exploration are reusable. Rational root isolation
and the fixed fold certificate are distinguished from floating-point trajectories.
"""
from __future__ import annotations

# EDITABLE INPUTS: nondimensional maintained activities and rates from Table 1.
from fractions import Fraction as Q
A_FEED = Q(6)
B_FEED = Q(27)
ZH_FORWARD = Q(16)
ZH_REVERSE = Q(2)
REVERSE_CORE_RATE = Q(1, 100000)
OUTPUT_LOSS = Q(1, 10000)
FEEDBACK_FRACTION = Q(1)
REVERSE_OUTPUT_INPUT = Q(0)  # positive values explore the reversible-sink extension
FEEDBACK_SWEEP = tuple(Q(i, 100) for i in range(10, 101))
TRAJECTORY_END = 1200.0
INITIAL_PERTURBATION = 0.02
SAMPLE_COUNT = 241
CUSTOM_INITIAL_STATES = ()  # optional additional (A, B, z, H) positive states
MANUSCRIPT_SHA256 = '83d9a0ecb93ca62c09b473fd55615b34712b44e5d6077598706b156d4d01cb44'

import argparse
import csv
import hashlib
import json
import math
import platform
from dataclasses import dataclass, replace
from functools import cached_property
from pathlib import Path

import numpy as np
import sympy as sp
from scipy.integrate import quad, solve_ivp

SPECIES = ('A', 'B', 'z', 'H')


@dataclass(frozen=True)
class Parameters:
    a: Q = A_FEED
    b: Q = B_FEED
    u: Q = ZH_FORWARD
    v: Q = ZH_REVERSE
    e: Q = REVERSE_CORE_RATE
    d: Q = OUTPUT_LOSS
    g: Q = FEEDBACK_FRACTION
    eta: Q = REVERSE_OUTPUT_INPUT

    def __post_init__(self):
        for name in self.__dataclass_fields__:
            object.__setattr__(self, name, Q(str(getattr(self, name))))
        if min(self.a, self.b, self.u, self.v, self.e, self.d) <= 0:
            raise ValueError('All six reaction rates must be positive.')
        if not 0 <= self.g <= 1 or self.eta < 0:
            raise ValueError('Routing is in [0,1]; reverse output input is nonnegative.')

    @property
    def flagship(self):
        return ((self.a, self.b, self.u, self.v, self.d, self.g, self.eta)
                == (6, 27, 16, 2, Q(1,10000), 1, 0)
                and Q(1,200000) <= self.e <= Q(1,50000))


@dataclass(frozen=True)
class ReversiblePair:
    """Reduced mass action after maintained external activities are substituted."""
    name: str
    reactants: tuple[int, ...]
    products: tuple[int, ...]
    forward: Q
    reverse: Q

    @property
    def change(self):
        return np.subtract(self.products, self.reactants)

    def current(self, state):
        def monomial(powers):
            return math.prod(value**power for value, power in zip(state, powers))
        return self.forward*monomial(self.reactants)-self.reverse*monomial(self.products)


class CoupledReactor:
    """Complete four-species field, with routing and reversible output as parameters."""
    def __init__(self, parameters=Parameters()):
        self.p = parameters
        p = parameters
        self.pairs = (
            ReversiblePair('dynamic fork', (1,0,0,0), (0,1,1,0), p.g, p.g),
            ReversiblePair('buffered fork', (1,0,0,0), (0,1,0,0), 1-p.g, 1-p.g),
            ReversiblePair('z+F / H', (0,0,1,0), (0,0,0,1), p.u, Q(1)),
            ReversiblePair('H / 2z', (0,0,0,1), (0,0,2,0), Q(1), p.v),
            ReversiblePair('A reservoir', (0,0,0,0), (1,0,0,0), p.a, Q(1)),
            ReversiblePair('B reservoir', (0,0,0,0), (0,1,0,0), p.b, Q(1)),
            ReversiblePair('B+G / 2A', (0,1,0,0), (2,0,0,0), p.e, p.e),
            ReversiblePair('H / waste', (0,0,0,1), (0,0,0,0), p.d, p.eta),
        )

    def field(self, time, state):
        # Expanded form is the same sum of reaction changes and currents; tests
        # check that identity exactly. This form makes repeated integration cheap.
        A, B, z, H = state
        a,b,u,v,e,d,g,eta = (float(getattr(self.p,k)) for k in ('a','b','u','v','e','d','g','eta'))
        c = 1-g+g*z
        return np.array([a-2*A+c*B+2*e*(B-A*A),
                         b+A-(1+c)*B-e*(B-A*A),
                         g*(A-B*z)-u*z-2*v*z*z+3*H,
                         u*z+v*z*z-(2+d)*H+eta])

    def jacobian(self, state):
        A,B,z,H = state
        u,v,e,d,g = (float(getattr(self.p,k)) for k in ('u','v','e','d','g'))
        return np.array([[-2-4*e*A,1-g+g*z+2*e,g*B,0],
                         [1+2*e*A,-2+g-g*z-e,-g*B,0],
                         [g,-g*z,-g*B-u-4*v*z,3],
                         [0,0,u+2*v*z,-2-d]])

    def activity(self, state):
        # The paper's two distinguished cores use the dynamic fork. At g != 1
        # no flagship eventual-signature theorem is asserted.
        j = self.pairs[0].current(state)
        k = self.pairs[6].current(state)
        ell = self.pairs[2].current(state)
        m = self.pairs[3].current(state)
        return np.array([-j+2*k, j-k, -ell+2*m, ell-m], dtype=float)

    def integrate(self, initial, end=TRAJECTORY_END, method='Radau', rtol=2e-10):
        initial = np.asarray(initial, dtype=float)
        if initial.shape != (4,) or not np.isfinite(initial).all() or np.min(initial) <= 0 or end <= 0:
            raise ValueError('Use a finite positive four-species state and horizon.')
        sol = solve_ivp(self.field, (0,end), initial, method=method,
                        jac=lambda t,x:self.jacobian(x), rtol=rtol, atol=rtol/100,
                        dense_output=True)
        if not sol.success or np.min(sol.y) <= 0:
            raise RuntimeError('Numerical trajectory failed or lost positivity.')
        return sol


class StationaryResponse:
    """Stationary elimination, never used as a scalar time-evolution equation."""
    def __init__(self, reactor):
        self.reactor = reactor
        self.p = reactor.p

    @cached_property
    def symbolic(self):
        p = self.p
        z = sp.Symbol('z')
        if p.g == 0:
            raise ValueError('Use isolated_state() at zero feedback.')
        B = (p.a+2*p.b)/(3-p.g+p.g*z)
        K = (p.v*(1+2*p.d)*z*z-p.u*(1-p.d)*z-3*p.eta)/(2+p.d)
        A = z*B+K/p.g
        H = (p.u*z+p.v*z*z+p.eta)/(2+p.d)
        residual = p.e*A*A+A+(1-p.e)*B-p.a-p.b
        poly = sp.Poly(sp.cancel(residual).as_numer_denom()[0], z)
        return z, (A,B,z,H), residual, poly

    def reconstruct(self, z):
        p = self.p
        if not p.g:
            raise ValueError('Zero routing uses the independent-module solution.')
        B = (p.a+2*p.b)/(3-p.g+p.g*z)
        K = (p.v*(1+2*p.d)*z*z-p.u*(1-p.d)*z-3*p.eta)/(2+p.d)
        return (z*B+K/p.g, B, z, (p.u*z+p.v*z*z+p.eta)/(2+p.d))

    def isolated_state(self):
        p = self.p
        if p.g != 0:
            raise ValueError('This solution is for g=0 only.')
        B = float((p.a+2*p.b)/3)
        rhs = float(p.a+p.b-(1-p.e)*Q(str(B)))
        A = 2*rhs/(1+math.sqrt(1+4*float(p.e)*rhs))
        c = float(p.u*(1-p.d)); lead=float(p.v*(1+2*p.d))
        z = (c+math.sqrt(c*c+12*lead*float(p.eta)))/(2*lead)
        state=np.array([A,B,z,float((p.u*z+p.v*z*z+p.eta)/(2+p.d))])
        return state if np.min(state)>0 else None

    def exact_roots(self, digits=14):
        """Exact polynomial real-root isolation plus rational positivity enclosure."""
        z, expressions, _, poly = self.symbolic
        result = []
        for (lo,hi), multiplicity in sp.polys.polytools.intervals(poly, eps=sp.Rational(1,10**digits)):
            if hi <= 0:
                continue
            box = Interval(Q(lo),Q(hi))
            state_boxes = [interval_evaluate(f, {z:box}) for f in expressions]
            if any(x.hi <= 0 for x in state_boxes):
                continue
            if any(x.lo <= 0 for x in state_boxes):
                raise ArithmeticError('Positivity unresolved; refine root intervals.')
            result.append({'z_interval':[str(lo),str(hi)], 'multiplicity':multiplicity,
                'state_intervals':[x.as_json() for x in state_boxes],
                'state':list(map(float,self.reconstruct((Q(lo)+Q(hi))/2)))})
        return result

    def numerical_states(self):
        if self.p.g == 0:
            state=self.isolated_state()
            return [] if state is None else [state]
        roots=np.roots([float(c) for c in self.symbolic[3].all_coeffs()])
        states=[]
        for root in sorted(roots,key=lambda v:v.real):
            if abs(root.imag)<1e-8 and root.real>0:
                x=np.array(self.reconstruct(root.real),dtype=float)
                if min(x)>0 and max(abs(self.reactor.field(0,x)))<1e-7:
                    states.append(x)
        return states

    def clamped_ab(self, z):
        if z <= 0:raise ValueError('Clamp must be positive.')
        p=self.p;c=1-p.g+p.g*z
        rhs=float((p.a*(1+c+p.e)+(c+2*p.e)*p.b)/(2+c))
        A=2*rhs/(1+math.sqrt(1+4*float(p.e)*rhs))
        B=float((p.b+A+p.e*A*A)/(1+c+p.e))
        return np.array([A,B])

    def clamped_zh(self, A, B):
        if min(A,B)<=0:raise ValueError('Clamps must be positive.')
        p=self.p
        a=float(p.v*(1+2*p.d));b=float((2+p.d)*p.g*B-p.u*(1-p.d))
        c=float((2+p.d)*p.g*A+3*p.eta)
        disc=math.sqrt(b*b+4*a*c)
        z=2*c/(b+disc) if b>=0 and c else (-b+disc)/(2*a)
        return np.array([z,float((p.u*z+p.v*z*z+p.eta)/(2+p.d))])


class ResponsePotential:
    """Paper equations (29)-(33); valid selection scope is checked explicitly."""
    def __init__(self, reactor):
        if not reactor.p.flagship:
            raise ValueError('Response-potential theorem requires the flagship family.')
        self.reactor=reactor;self.e=float(reactor.p.e);self.d=float(reactor.p.d)

    def response(self,B):
        q=33-B+self.e*B
        a=2*q/(1+math.sqrt(1+4*self.e*q))
        c=self.e*(1+2*a)/(1+2*self.e*a)
        return a+B,c

    @staticmethod
    def inside(state):
        A,B,z,H=state
        return min(state)>=0 and A+B<=34 and z+7*H/4<=384 and z<=12 and B>=2

    @staticmethod
    def L(z):return math.log1p(-1/(z+2))

    def phi(self,H):
        return 2*(2+self.d)*H/(16+math.sqrt(256+8*(2+self.d)*H))

    def residuals(self,state):
        A,B,z,H=state;h,c=self.response(B)
        return np.array([A+B-h,60-(2+z)*B,h-(1+z)*B-16*z-4*z*z+3*H,
                         16*z+2*z*z-(2+self.d)*H])

    def value(self,state):
        if not self.inside(state):raise ValueError('Evaluate this potential inside Omega.')
        A,B,z,H=state;h,c=self.response(B);r=A+B-h
        # U is integrated analytically. R and P use deterministic quadrature.
        U=4*z-12*math.log1p(z)+16*math.log1p(z/2)
        R=quad(lambda v:3*self.L(self.phi(v)),0,H,epsabs=2e-12,epsrel=2e-12)[0]
        P=quad(lambda v:math.log(v/60)+self.response(v)[1]*math.log1p(-v/60),
               20,B,epsabs=2e-12,epsrel=2e-12)[0]
        return B*math.log(z+2)-h*self.L(z)+U-3*H*self.L(z)+P+R+r*r/2

    def dissipation(self,state):
        r,F,Z,K=self.residuals(state)
        return r*r/4+F*F/1000+Z*Z/364+K*K/4000

    def derivative(self,state):
        A,B,z,H=state;h,c=self.response(B);r=A+B-h
        beta=60/B-2;Z=self.residuals(state)[2]
        vB=math.log(B*(z+2)/60)-c*(self.L(z)-self.L(beta))
        vz=-Z/((z+1)*(z+2));vH=3*(self.L(self.phi(H))-self.L(z))
        f=self.reactor.field(0,state)
        return float(vB*f[1]+vz*f[2]+vH*f[3]+r*(f[0]+(1-c)*f[1]))


class NumericalSelector:
    """A numerical energy-and-side test; NOT a validated flow enclosure."""
    def __init__(self,potential,saddle,tolerance=1e-8):
        self.potential=potential;self.saddle=np.array(saddle)
        self.barrier=potential.value(saddle);self.tolerance=tolerance

    def inspect(self,state):
        if not self.potential.inside(state):return 'outside_absorbing_region'
        if self.barrier-self.potential.value(state)<=self.tolerance:return 'unresolved'
        side=state[1]-self.saddle[1]
        if abs(side)<=self.tolerance:return 'unresolved'
        return 'low_z_sink' if side>0 else 'high_z_sink'


@dataclass(frozen=True)
class Interval:
    """Closed rational intervals: no floating-point acceptance tests."""
    lo: Q
    hi: Q | None = None
    def __post_init__(self):
        object.__setattr__(self,'lo',Q(self.lo))
        object.__setattr__(self,'hi',Q(self.lo if self.hi is None else self.hi))
        if self.lo>self.hi:raise ValueError('Inverted interval')
    @staticmethod
    def cast(v):return v if isinstance(v,Interval) else Interval(v)
    def __add__(self,v):
        v=self.cast(v);return Interval(self.lo+v.lo,self.hi+v.hi)
    __radd__=__add__
    def __neg__(self):return Interval(-self.hi,-self.lo)
    def __sub__(self,v):return self+-self.cast(v)
    def __rsub__(self,v):return self.cast(v)+-self
    def __mul__(self,v):
        v=self.cast(v);s=[a*b for a in (self.lo,self.hi) for b in (v.lo,v.hi)]
        return Interval(min(s),max(s))
    __rmul__=__mul__
    def __pow__(self,n):
        if n<0:
            if self.lo<=0<=self.hi:raise ZeroDivisionError('Interval contains zero')
            return Interval(1/self.hi,1/self.lo)**(-n)
        if n==0:return Interval(1)
        if n%2:return Interval(self.lo**n,self.hi**n)
        s=[self.lo**n,self.hi**n]
        return Interval(0 if self.lo<=0<=self.hi else min(s),max(s))
    def size(self):return max(abs(self.lo),abs(self.hi))
    def as_json(self):return [str(self.lo),str(self.hi)]


def interval_evaluate(expr,values):
    expr=sp.sympify(expr)
    if expr.is_Rational:return Interval(Q(expr))
    if expr.is_Symbol:return values[expr]
    if expr.is_Add:return sum((interval_evaluate(v,values) for v in expr.args),Interval(0))
    if expr.is_Mul:
        out=Interval(1)
        for v in expr.args:out=out*interval_evaluate(v,values)
        return out
    if expr.is_Pow and expr.exp.is_Integer:return interval_evaluate(expr.base,values)**int(expr.exp)
    raise ValueError('Only rational expressions are accepted.')


class FoldCertificate:
    """Recompute Appendix B's fixed rational contraction and Hurwitz tests.

    Adapted from the manuscript's certify_fold.py, with no cached acceptance data.
    This certificate is independent of the editable exploration parameters.
    """
    @staticmethod
    def check():
        z,g=sp.symbols('z g');e=sp.Rational(1,100000);h=sp.Rational(20001,10000)
        D=3-g+g*z;B=60/D;K=(20004*z*z-159984*z)/20001
        A=z*B+K/g;H=(16*z+2*z*z)/h
        residual=e*A*A+A+(1-e)*B-33
        f=sp.Matrix([residual,sp.diff(residual,z)])
        derivative=f.jacobian([z,g])
        center=[Q('1.617456101591'),Q('0.970187992295')]
        radius=[Q(1,10**9),Q(1,10**10)]
        point={v:Interval(c) for v,c in zip((z,g),center)}
        box={v:Interval(c-r,c+r) for v,c,r in zip((z,g),center,radius)}
        C=sp.Matrix(2,2,lambda i,j:sp.Rational(interval_evaluate(derivative[i,j],point).lo)).inv()
        M=[[Interval(int(i==j))-sum((Q(C[i,k])*interval_evaluate(derivative[k,j],box)
             for k in range(2)),Interval(0)) for j in range(2)] for i in range(2)]
        error=[-sum((Q(C[i,k])*interval_evaluate(f[k],point) for k in range(2)),Interval(0))
               +sum((M[i][j]*Interval(-radius[j],radius[j]) for j in range(2)),Interval(0)) for i in range(2)]
        norm=max(sum(M[i][j].size()*radius[j]/radius[i] for j in range(2)) for i in range(2))
        assert norm<1 and all(error[i].size()<radius[i] for i in range(2))
        qg=interval_evaluate(sp.diff(residual,g),box)
        qzz=interval_evaluate(sp.diff(residual,z,2),box)
        assert qg.lo>0 and qzz.hi<0
        aa,bb=sp.symbols('aa bb')
        J=sp.Matrix([[-2-4*e*aa,1-g+g*z+2*e,g*bb,0],
                     [1+2*e*aa,-2+g-g*z-e,-g*bb,0],
                     [g,-g*z,-g*bb-16-8*z,3],[0,0,16+4*z,-h]])
        coefficients=J.charpoly().all_coeffs()
        assert sp.cancel(coefficients[-1].subs({aa:A,bb:B})-g*h*D*sp.diff(residual,z))==0
        values={**box,aa:interval_evaluate(A,box),bb:interval_evaluate(B,box)}
        cubic=[interval_evaluate(v,values) for v in coefficients[1:4]]
        gap=cubic[0]*cubic[1]-cubic[2]
        assert min(v.lo for v in cubic)>0 and gap.lo>0
        state=[interval_evaluate(v,box) for v in (A,B,z,H)]
        assert min(v.lo for v in state)>0
        return {'scope':'fixed e=1/100000; local fold only; Lean proofs not rerun',
                'box':{str(k):v.as_json() for k,v in box.items()},
                'contraction_norm':str(norm),'newton_errors':[v.as_json() for v in error],
                'Qg':qg.as_json(),'Qzz':qzz.as_json(),'state':[v.as_json() for v in state],
                'cubic_coefficients':[v.as_json() for v in cubic],'hurwitz_gap':gap.as_json()}


def write_csv(path,headers,rows):
    with path.open('w',newline='') as stream:
        writer=csv.writer(stream);writer.writerow(headers);writer.writerows(rows)


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output',type=Path,default=Path('outputs'))
    args=parser.parse_args();out=args.output;out.mkdir(parents=True,exist_ok=True)
    model=CoupledReactor();response=StationaryResponse(model)
    roots=response.exact_roots() if model.p.g else []
    states=[np.array(r['state']) for r in roots] if model.p.g else response.numerical_states()
    equilibria=[]
    for x in states:
        eig=np.linalg.eigvals(model.jacobian(x));activity=model.activity(x)
        equilibria.append({'state':x.tolist(),'eigenvalues_real':eig.real.tolist(),
            'eigenvalues_imag':eig.imag.tolist(),'numerical_sink':bool(max(eig.real)<0),
            'activity':activity.tolist(),'signature':[bool(min(activity[:2])>0),bool(min(activity[2:])>0)],
            'field_residual':float(max(abs(model.field(0,x))))})
    sweep=[]
    for g in FEEDBACK_SWEEP:
        trial=CoupledReactor(replace(model.p,g=g))
        for x in StationaryResponse(trial).numerical_states():
            spectral=float(max(np.linalg.eigvals(trial.jacobian(x)).real))
            sweep.append([float(g),*x,spectral])
    write_csv(out/'feedback_sweep.csv',['g',*SPECIES,'largest_eigenvalue_real'],sweep)
    write_csv(out/'equilibria.csv',[*SPECIES,'AB_A','AB_B','ZH_z','ZH_H'],
              [d['state']+d['activity'] for d in equilibria])
    trajectories=[];diagnostics=[]
    initials=list(CUSTOM_INITIAL_STATES)
    potential=selector=None
    if model.p.flagship and len(states)==3:
        potential=ResponsePotential(model);selector=NumericalSelector(potential,states[1])
        eigenvalues,vectors=np.linalg.eig(model.jacobian(states[1]))
        direction=vectors[:,np.argmax(eigenvalues.real)].real
        direction=direction/np.linalg.norm(direction)*np.sign(direction[1])
        initials=[states[1]+sign*INITIAL_PERTURBATION*direction for sign in (1,-1)]+initials
    elif not initials:
        initials=[(1,1,1,1)]
    times=np.linspace(0,TRAJECTORY_END,SAMPLE_COUNT)
    for i,initial in enumerate(initials):
        sol=model.integrate(initial);fine=model.integrate(initial,method='BDF',rtol=2e-12)
        values=sol.sol(times).T;acts=np.array([model.activity(x) for x in values])
        energy=np.array([potential.value(x)-selector.barrier if potential and potential.inside(x) else np.nan for x in values])
        verdicts=[selector.inspect(x) if selector else 'outside_theorem_scope' for x in values]
        selected=next(((t,v) for t,v in zip(times,verdicts) if v.endswith('_sink')),None)
        diagnostics.append({'initial':list(map(float,initial)),'endpoint':values[-1].tolist(),
            'first_sampled_numerical_selection':selected,'last_verdict':verdicts[-1],
            'independent_solver_max_difference':float(np.max(abs(values-fine.sol(times).T))),
            'largest_sampled_energy_increase':float(max(np.diff(energy))) if np.isfinite(energy).all() else None,
            'maximum_dV_plus_D':float(max(potential.derivative(x)+potential.dissipation(x) for x in values)) if potential and all(potential.inside(x) for x in values) else None})
        write_csv(out/f'trajectory_{i}.csv',['time',*SPECIES,'AB_A','AB_B','ZH_z','ZH_H','V_minus_saddle','numerical_verdict'],
                  zip(times,*values.T,*acts.T,energy,verdicts))
        trajectories.append((values,acts,energy))
    fold=FoldCertificate.check()
    result={'parameters':{k:str(getattr(model.p,k)) for k in model.p.__dataclass_fields__},
            'flagship_theorem_scope':model.p.flagship,'exact_stationary_roots':roots,
            'equilibria':equilibria,'trajectories':diagnostics,'fixed_fold_certificate':fold,
            'scope':'exact rational stationary isolation and local fold; numerical flows and selection only'}
    (out/'results.json').write_text(json.dumps(result,indent=2)+'\n')
    lines=[f'Positive stationary states: {len(states)}; flagship theorem scope: {model.p.flagship}.']
    lines += [f'z={d["state"][2]:.10f}; numerical sink={d["numerical_sink"]}; core activity signature={d["signature"]}.' for d in equilibria]
    lines += [f'Trajectory {i}: {d["last_verdict"]}; first sampled test={d["first_sampled_numerical_selection"]}.' for i,d in enumerate(diagnostics)]
    lines += ['Fixed local fold: exact rational contraction and cubic Hurwitz checks pass.',
              'Trajectory verdicts use floating point; they are not validated basin certificates.']
    (out/'console.txt').write_text('\n'.join(lines)+'\n');print('\n'.join(lines))
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,ax=plt.subplots(figsize=(7,4),layout='constrained')
    data=np.array(sweep)
    for stable,color,label in [(True,'#287a99','Numerical sink'),(False,'#bd642f','Numerical unstable state')]:
        keep=(data[:,-1]<0)==stable
        ax.scatter(data[keep,0],data[keep,3],s=9,color=color,label=label)
    if (model.p.a,model.p.b,model.p.u,model.p.v,model.p.e,model.p.d,model.p.eta)==(6,27,16,2,Q(1,100000),Q(1,10000),0):
        ax.plot(.970187992295,1.617456101591,'kD',ms=5,label='Rationally isolated local fold')
    ax.set(xlabel='Dynamic feedback fraction g (total fork capacity = 1)',ylabel='Stationary z',title='Stationary compositions versus feedback\nfraction')
    ax.legend(fontsize=8);ax.grid(alpha=.2)
    fig.savefig(out/'feedback.png',dpi=180);fig.savefig(out/'feedback.svg');plt.close(fig)
    fig,axes=plt.subplots(3,1,figsize=(8,8),layout='constrained',sharex=True)
    for i,(values,acts,energy) in enumerate(trajectories):
        color=['#287a99','#bd642f'][i%2]
        axes[0].plot(times,values[:,2],color=color,label=f'Initial state {i}')
        axes[1].plot(times,energy,color=color)
        for col,style,label in [(1,'--','AB: B production'),(2,'-','ZH: z production'),(3,':','ZH: H production')]:
            axes[2].plot(times,acts[:,col],style,color=color,label=label if i==0 else None)
    axes[0].set(ylabel='z',title='Concentration trajectories from two initial\nstates');axes[0].legend(fontsize=8)
    axes[1].axhline(0,color='gray',lw=.7);axes[1].set(ylabel='V - V(saddle)')
    if potential is None:axes[1].text(.5,.5,'Potential test outside theorem scope',ha='center',transform=axes[1].transAxes)
    axes[2].set(xlabel='Time (reference-rate units)',ylabel='Core production',ylim=(-30,30))
    axes[2].set_yscale('symlog',linthresh=1e-4)
    axes[2].legend(fontsize=8,ncol=3)
    for ax in axes:ax.grid(alpha=.2)
    fig.savefig(out/'selection.png',dpi=180);fig.savefig(out/'selection.svg');plt.close(fig)
    digest=lambda p:hashlib.sha256(p.read_bytes()).hexdigest()
    (out/'run_metadata.json').write_text(json.dumps({'manuscript_sha256':MANUSCRIPT_SHA256,
        'source_sha256':digest(Path(__file__)),'python':platform.python_version(),
        'output_sha256':{p.name:digest(p) for p in sorted(out.iterdir()) if p.is_file() and p.name!='run_metadata.json'}},indent=2)+'\n')


if __name__=='__main__':main()
