"""Classical mass action with identical stoichiometry/support, opposite stability.

Exact certificates plus reusable equilibrium-centred rate and Jacobian models.
"""
from __future__ import annotations
import argparse
import csv
from dataclasses import dataclass
from fractions import Fraction as F
from itertools import combinations,permutations,product
import hashlib
import json
from math import lcm,log
from pathlib import Path
import platform
import time
import numpy as np

# USER INPUTS ---------------------------------------------------------------
# Rows are species, columns are reactions. Abstract nondimensional model.
BASE_REACTANTS = ((1,0,0,0,0),(0,4,0,0,0),(0,2,4,0,0),(2,0,0,2,0))
BASE_PRODUCTS = ((0,0,0,1,0),(0,0,0,6,0),(4,0,0,1,2),(0,0,2,0,2))
EQUILIBRIUM_FLUX = (2,3,2,2,2)
COMPACT_EXPONENT = 200
COMPACT_STATE = ('1','235','585','1/257')
EXACT_STATE = ('2','600','33000','93316000000')
EXACT_PADDED_ORDERS = (11,7500,5187522222797,28200219703)
EXPONENT_SWEEP = tuple(range(4,241,2)) # integer exponents; each point reconstructs k
PAPER_SHA256 = '46399af83f2c5ce5870601ab84cb955d314a12a57dfb4b01bb6ca4fa8ff37f62'
# These high-order reactions are mathematical witnesses, not calibrated chemistry.
# --------------------------------------------------------------------------


def rational_matrix(values): return np.array([[F(v) for v in row] for row in values],dtype=object)
def strings(a): return np.vectorize(str)(a).tolist()


def determinant(a):
    a=np.asarray(a,dtype=object);n=len(a)
    if a.shape!=(n,n): raise ValueError('Determinant needs a square matrix.')
    result=F(0)
    for order in permutations(range(n)):
        term=F((-1)**sum(order[i]>order[j] for i in range(n) for j in range(i+1,n)))
        for i,j in enumerate(order): term*=a[i,j]
        result+=term
    return result


def characteristic(a):
    """Descending coefficients of det(zI-A), by exact principal minors."""
    a=np.asarray(a,dtype=object);n=len(a)
    return [F(1)]+[(-1)**k*sum((determinant(a[np.ix_(s,s)]) for s in combinations(range(n),k)),F(0)) for k in range(1,n+1)]


def quartic_certificate(a):
    coefficients=characteristic(a)
    if len(coefficients)!=5: raise ValueError('Quartic certificate needs four species.')
    _,a1,a2,a3,a4=coefficients;d2=a1*a2-a3;d3=d2*a3-a1*a1*a4
    # Strict signs suffice here; boundary or other sign patterns remain unclassified.
    stable=all(v>0 for v in (a1,a2,a3,a4,d2,d3))
    pair=all(v>0 for v in (a1,a2,a3,a4,d2)) and d3<0
    return dict(coefficients=list(map(str,coefficients)),delta2=str(d2),delta3=str(d3),
                classification='strictly stable' if stable else 'two right-half-plane roots' if pair else 'unclassified boundary/sign pattern')


@dataclass(frozen=True)
class Network:
    reactants: np.ndarray
    products: np.ndarray

    def __post_init__(self):
        y=rational_matrix(self.reactants);p=rational_matrix(self.products)
        if y.shape!=p.shape or y.ndim!=2 or min(y.shape)<1: raise ValueError('Nonempty equally shaped complex matrices required.')
        if any(v<0 or v.denominator!=1 for v in (*y.flat,*p.flat)): raise ValueError('Complexes require nonnegative integers.')
        y.setflags(write=False);p.setflags(write=False)
        object.__setattr__(self,'reactants',y);object.__setattr__(self,'products',p)

    @property
    def stoichiometry(self): return self.products-self.reactants

    def pad_to(self,orders):
        orders=rational_matrix(orders)
        if orders.shape!=self.reactants.shape: raise ValueError('Padding shape mismatch.')
        pad=orders-self.reactants
        if any(v<0 for v in pad.flat) or np.any((orders>0)!=(self.reactants>0)):
            raise ValueError('Padding must increase existing incidences without adding new ones.')
        return Network(orders,self.products+pad)

    def same_skeleton(self,other):
        return np.array_equal(self.stoichiometry,other.stoichiometry) and np.array_equal(self.reactants>0,other.reactants>0)

    def children(self):
        n,m=self.reactants.shape;children=[]
        for size in range(1,n+1):
            for species in combinations(range(n),size):
                choices=[tuple(j for j in range(m) if self.reactants[i,j]>0) for i in species]
                for reactions in product(*choices):
                    if len(set(reactions))==size:
                        children.append((species,reactions,self.stoichiometry[np.ix_(species,reactions)]))
        return children


class EquilibriumModel:
    """Mass action represented by a network, positive stationary flux, and state.

    r(x) = v * exp(Y.T @ log(x/xbar)); avoids expanding enormous k monomials.
    """
    def __init__(self,network:Network,state,flux):
        self.network=network;self.state=np.array([F(x) for x in state],dtype=object);self.flux=np.array([F(v) for v in flux],dtype=object)
        n,m=network.reactants.shape
        if self.state.shape!=(n,) or self.flux.shape!=(m,) or any(x<=0 for x in (*self.state,*self.flux)):
            raise ValueError('Positive state and flux of the correct sizes required.')
        if any(v!=0 for v in network.stoichiometry@self.flux): raise ValueError('Flux must lie exactly in the stoichiometric kernel.')
        self.y=np.array(network.reactants,dtype=float);self.s=np.array(network.stoichiometry,dtype=float)
        self.x=np.array(self.state,dtype=float);self.v=np.array(self.flux,dtype=float)

    def reactivity(self): return self.flux[:,None]*self.network.reactants.T/self.state[None,:]
    def jacobian(self): return self.network.stoichiometry@self.reactivity()

    def log_rate_constants(self):
        # Logarithms of exact rational factors; no attempt to expand powers.
        return np.array([log(v.numerator)-log(v.denominator) for v in self.flux])-self.y.T@np.array([log(x.numerator)-log(x.denominator) for x in self.state])

    def rates(self,state):
        state=np.asarray(state)
        if state.shape!=self.x.shape or np.any(np.real(state)<=0): raise ValueError('Rate evaluation requires positive concentrations.')
        with np.errstate(over='raise',invalid='raise'):
            return self.v*np.exp(self.y.T@np.log(state/self.x))

    def rhs(self,t,state): return self.s@self.rates(state)

    def numerical_jacobian(self,state):
        state=np.asarray(state)
        return self.s@(self.rates(state)[:,None]*self.y.T/state[None,:])

    def simulate(self,initial,times,rtol=1e-8,atol=1e-10):
        """Optional nonlinear positive-state simulation in log coordinates.

        Highly stiff witnesses need care; this is not a periodic-orbit detector.
        """
        from scipy.integrate import solve_ivp
        initial=np.asarray(initial,dtype=float);times=np.asarray(times,dtype=float)
        if initial.shape!=self.x.shape or np.any(initial<=0) or len(times)<2 or np.any(np.diff(times)<=0):
            raise ValueError('Positive initial state and increasing output times required.')
        def rhs(t,q):
            with np.errstate(over='raise',invalid='raise'):
                state=self.x*np.exp(q);return self.s@(self.v*np.exp(self.y.T@q))/state
        def jac(t,q):
            state=self.x*np.exp(q);drift=rhs(t,q)
            return (self.numerical_jacobian(state)*state[None,:])/state[:,None]-np.diag(drift)
        result=solve_ivp(rhs,(times[0],times[-1]),np.log(initial/self.x),t_eval=times,method='Radau',jac=jac,rtol=rtol,atol=atol)
        if not result.success: raise RuntimeError(result.message)
        return self.x[:,None]*np.exp(result.y)


class RationalLift:
    """Construct the theorem's uniform-state integer-exponent padding exactly."""
    @staticmethod
    def build(network,flux,target):
        flux=np.array([F(v) for v in flux],dtype=object);target=rational_matrix(target)
        n,m=network.reactants.shape
        if flux.shape!=(m,) or any(v<=0 for v in flux) or target.shape!=(m,n): raise ValueError('Invalid flux or reactivity shape.')
        if any(v<0 for v in target.flat) or np.any((target>0)!=(network.reactants.T>0)):
            raise ValueError('Target must be positive exactly on the existing reactant incidences.')
        q=target.T/flux[None,:];denominator=lcm(*(v.denominator for v in q.flat))
        scale=1+sum(network.reactants.flat);state=int(scale*denominator)
        model=EquilibriumModel(network.pad_to(q*state),[state]*n,flux)
        if not np.array_equal(model.reactivity(),target): raise ArithmeticError('Lift failed to reconstruct target reactivity.')
        return model


def poly_product(a,b):
    result={}
    for u,x in a.items():
        for v,y in b.items():
            w=tuple(i+j for i,j in zip(u,v));result[w]=result.get(w,F(0))+x*y
    return {w:x for w,x in result.items() if x}


def scaled_coefficients(a):
    """Coefficient polynomials of det(zI-A diag(d)), as sparse exact monomials."""
    n=len(a);result=[]
    for k in range(1,n+1):
        terms={}
        for s in combinations(range(n),k):
            value=(-1)**k*determinant(a[np.ix_(s,s)])
            if value: terms[tuple(int(i in s) for i in range(n))]=value
        result.append(terms)
    return result


def base_stability_certificate(base,flux):
    b=base.stoichiometry@np.diag([F(v) for v in flux])@base.reactants.T
    if any(b[0,j]!=0 for j in range(1,4)) or b[0,0]>=0: raise ValueError('Base block certificate does not apply.')
    a1,a2,a3=scaled_coefficients(b[1:,1:]);delta=poly_product(a1,a2)
    for powers,value in a3.items(): delta[powers]=delta.get(powers,F(0))-value
    delta={p:c for p,c in delta.items() if c}
    if not all(terms and all(c>0 for c in terms.values()) for terms in (a1,a2,a3,delta)):
        raise ArithmeticError('Strict all-positive-scaling certificate failed.')
    return dict(template=strings(b),isolated_eigenvalue_coefficient=str(b[0,0]),
        cubic_coefficients=[{','.join(map(str,p)):str(c) for p,c in terms.items()} for terms in (a1,a2,a3)],
        cubic_hurwitz_gap={','.join(map(str,p)):str(c) for p,c in delta.items()})


def child_certificates(network):
    children=network.children();lookup={(s,r):a for s,r,a in children}
    specifications=[((0,1,2,3),(0,1,2,3),(100,35,40,140)),
                    ((0,2,3),(0,1,3),(10,2,7)),((2,3),(1,0),(1,2)),
                    ((1,3),(1,0),(1,2)),((2,3),(2,0),(1,2))]
    certificates=[];covered=set()
    for species,reactions,weight in specifications:
        a=lookup[(species,reactions)];p=np.diag(weight);m=-(p@a+a.T@p)
        minors=[determinant(m[np.ix_(s,s)]) for k in range(1,len(m)+1) for s in combinations(range(len(m)),k)]
        if any(v<0 for v in minors): raise ArithmeticError('Weighted symmetrization is not positive semidefinite.')
        certificates.append(dict(species=species,reactions=reactions,weights=weight,symmetrization=strings(m),principal_minors=list(map(str,minors))))
        mapping=dict(zip(species,reactions))
        covered.update((s,r) for s,r,a in children if all(mapping.get(i)==j for i,j in zip(s,r)))
    # Remaining maximal child: exact scaled characteristic z(z+4d0)(z+4d1+2d2).
    singular=((1,2,3),(1,2,0));a=lookup[singular];actual=scaled_coefficients(a)
    expected=[{(1,0,0):F(4),(0,1,0):F(4),(0,0,1):F(2)},
              {(1,1,0):F(16),(1,0,1):F(8)},{}]
    if actual!=expected: raise ArithmeticError('Singular child factorization failed.')
    covered.add(singular)
    if covered!=set(lookup): raise ArithmeticError('Certificate coverage misses a child selection.')
    return dict(child_count=len(children),covered_count=len(covered),dissipativity=certificates,
        singular_child=dict(species=singular[0],reactions=singular[1],factorization='z (z+4d0) (z+4d1+2d2)'),
        conclusion='Every child is D-nonunstable; the singular child has a permanent zero eigenvalue.')


def exact_eigenpair(jacobian):
    real=np.array([-152886431705,23331004034,7890610882,1399740000],dtype=object)
    imag=np.array([15563289455,-30699360512,34396287629,0],dtype=object)
    a,b=F(1,500),F(51,500)
    residual_real=jacobian@real-a*real+b*imag;residual_imag=jacobian@imag-b*real-a*imag
    if any(v!=0 for v in (*residual_real,*residual_imag)): raise ArithmeticError('Exact unstable eigenpair failed.')
    return dict(real_part=str(a),imaginary_part=str(b),real_vector=list(map(str,real)),imaginary_vector=list(map(str,imag)),residual='exactly zero')


def compact_network(base,exponent):
    if not isinstance(exponent,int) or exponent<4: raise ValueError('Compact family exponent must be an integer at least four.')
    orders=base.reactants.copy();orders[2,2]=exponent;orders[3,0]=exponent
    return base.pad_to(orders)


def write_csv(path,rows):
    with path.open('w',newline='',encoding='utf-8') as f:
        writer=csv.DictWriter(f,fieldnames=list(rows[0]));writer.writeheader();writer.writerows(rows)


def plot(base,compact,sweep,output):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    with plt.rc_context({'font.size':11,'figure.facecolor':'white','savefig.facecolor':'white','svg.fonttype':'none'}):
        incidences=list(zip(*np.nonzero(base.reactants>0)));x=np.arange(len(incidences))
        fig,ax=plt.subplots(figsize=(7.2,4.5),layout='constrained')
        ax.bar(x-.18,[float(base.reactants[i,j]) for i,j in incidences],width=.36,color='#0072B2',label='Original orders')
        ax.bar(x+.18,[float(compact.reactants[i,j]) for i,j in incidences],width=.36,color='#D55E00',label='Compact padded orders')
        ax.set_xticks(x,[f'X{i} in r{j}' for i,j in incidences]);ax.tick_params(axis='x',labelsize=9)
        ax.set(yscale='log',ylabel='Reactant multiplicity (log scale)');ax.legend(fontsize=9)
        ax.spines[['top','right']].set_visible(False)
        fig.savefig(output/'orders.png',dpi=220);fig.savefig(output/'orders.svg');plt.close(fig)
        fig,ax=plt.subplots(figsize=(7.2,4.5),layout='constrained')
        ax.axhline(0,color='#444444',linewidth=1)
        ax.plot([r['exponent'] for r in sweep],[r['max_real_eigenvalue'] for r in sweep],color='#0072B2',linewidth=1.5)
        for classification,color,marker in [('strictly stable','#0072B2','o'),('two right-half-plane roots','#D55E00','s')]:
            rows=[r for r in sweep if r['classification']==classification]
            ax.scatter([r['exponent'] for r in rows],[r['max_real_eigenvalue'] for r in rows],s=12,c=color,marker=marker,label=classification)
        ax.set(xlabel='Common integer order in X2 of r2 and X3 of r0',ylabel='Largest real part of Jacobian eigenvalues')
        ax.legend(fontsize=9);ax.spines[['top','right']].set_visible(False);ax.grid(axis='y',color='#e3e6e8')
        fig.savefig(output/'stability.png',dpi=220);fig.savefig(output/'stability.svg');plt.close(fig)


def main():
    parser=argparse.ArgumentParser(description=__doc__);parser.add_argument('--output',type=Path,default=Path('outputs'))
    args=parser.parse_args();args.output.mkdir(parents=True,exist_ok=True);start=time.perf_counter()
    inputs=dict(base_reactants=BASE_REACTANTS,base_products=BASE_PRODUCTS,flux=EQUILIBRIUM_FLUX,compact_exponent=COMPACT_EXPONENT,
        compact_state=COMPACT_STATE,exact_state=EXACT_STATE,exact_padded_orders=EXACT_PADDED_ORDERS,exponent_sweep=EXPONENT_SWEEP)
    intro='Resolved inputs: '+json.dumps(inputs);print(intro,flush=True)
    base=Network(BASE_REACTANTS,BASE_PRODUCTS);orders=base.reactants.copy()
    for (i,j),value in zip(((2,1),(2,2),(3,0),(3,3)),EXACT_PADDED_ORDERS): orders[i,j]=value
    exact=EquilibriumModel(base.pad_to(orders),EXACT_STATE,EQUILIBRIUM_FLUX)
    compact=EquilibriumModel(compact_network(base,COMPACT_EXPONENT),COMPACT_STATE,EQUILIBRIUM_FLUX)
    stable=EquilibriumModel(base,COMPACT_STATE,EQUILIBRIUM_FLUX)
    lift=RationalLift.build(base,EQUILIBRIUM_FLUX,exact.reactivity())
    eigenpair=exact_eigenpair(exact.jacobian());children=child_certificates(base)
    if any(not base.same_skeleton(m.network) for m in (exact,compact,lift)): raise ArithmeticError('Changed skeleton.')
    base_certificate=base_stability_certificate(base,EQUILIBRIUM_FLUX)
    models={'base':stable,'compact':compact,'exact':exact,'uniform_lift':lift};records={};spectra=[]
    for name,model in models.items():
        eigenvalues=np.linalg.eigvals(np.array(model.jacobian(),dtype=float))
        spectra.extend(dict(model=name,real=float(v.real),imaginary=float(v.imag)) for v in eigenvalues)
        records[name]=dict(reactants=strings(model.network.reactants),products=strings(model.network.products),
            state=list(map(str,model.state)),flux=list(map(str,model.flux)),reactivity=strings(model.reactivity()),
            jacobian=strings(model.jacobian()),log_rate_constants=model.log_rate_constants().tolist(),
            rate_constant_definition='k[j] = flux[j] / product(state[i] ** reactants[i,j])',certificate=quartic_certificate(model.jacobian()))
    sweep=[]
    for exponent in EXPONENT_SWEEP:
        model=EquilibriumModel(compact_network(base,exponent),COMPACT_STATE,EQUILIBRIUM_FLUX);jac=model.jacobian()
        certificate=quartic_certificate(jac)
        sweep.append(dict(exponent=exponent,max_real_eigenvalue=float(max(np.linalg.eigvals(np.array(jac,dtype=float)).real)),
            delta2=certificate['delta2'],delta3=certificate['delta3'],classification=certificate['classification']))
    write_csv(args.output/'sweep.csv',sweep);write_csv(args.output/'spectra.csv',spectra)
    write_csv(args.output/'children.csv',[dict(species=' '.join(map(str,s)),reactions=' '.join(map(str,r)),matrix=json.dumps(strings(a))) for s,r,a in base.children()])
    (args.output/'models.json').write_text(json.dumps(records,indent=2)+'\n',encoding='utf-8')
    certificates=dict(base_stability=base_certificate,children=children,eigenpair=eigenpair)
    (args.output/'certificates.json').write_text(json.dumps(certificates,indent=2)+'\n',encoding='utf-8')
    plot(base,compact.network,sweep,args.output)
    summary=dict(inputs=inputs,same_skeleton=True,child_count=children['child_count'],exact_eigenpair=eigenpair,
        compact_certificate=records['compact']['certificate'],base_certificate=records['base']['certificate'],
        uniform_lift_state=str(lift.state[0]),uniform_lift_max_exponent=str(max(lift.network.reactants.flat)),
        first_sampled_unstable_exponent=next((r['exponent'] for r in sweep if r['classification']=='two right-half-plane roots'),None),
        evidence='Rational certificates for flux, lifting, eigenpair, child coverage and Hurwitz signs. Spectral coordinates are floating-point visualizations. No Hopf or periodic-orbit claim; no Lean compilation.')
    transcript=json.dumps(summary,indent=2)
    (args.output/'summary.json').write_text(transcript+'\n',encoding='utf-8')
    (args.output/'console.txt').write_text(intro+'\n'+transcript+'\n',encoding='utf-8')
    import scipy,matplotlib
    metadata=dict(paper_sha256=PAPER_SHA256,source_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
        python=platform.python_version(),numpy=np.__version__,scipy=scipy.__version__,matplotlib=matplotlib.__version__,platform=platform.platform(),
        elapsed_seconds=time.perf_counter()-start,command='python example.py --output outputs',seed_policy='No random sampling.',
        output_sha256={p.name:hashlib.sha256(p.read_bytes()).hexdigest() for p in sorted(args.output.iterdir()) if p.is_file() and p.name!='run_metadata.json'})
    (args.output/'run_metadata.json').write_text(json.dumps(metadata,indent=2)+'\n',encoding='utf-8');print(transcript)


if __name__=='__main__': main()