"""Self-limited consumers coupled to the paper's loaded resident network.

Run: python example.py --output outputs
Concentrations/time are nondimensional; numerical trajectories are not proofs.
"""
from __future__ import annotations

# EDITABLE STUDY INPUTS: paper illustrations, not measured kinetic parameters.
RESIDENT_E = 1e-5
SUPPLY = 0.05
INITIAL_RESIDENT = (3., 10., 0.05, 3.)  # A, B, z, H
INITIAL_RESERVOIR = 0.001
INITIAL_CONSUMERS = (1e-8, 1.)
TARGET_SHARES = (0.25, 0.75)
TARGET_KAPPA = 1.
HORIZON_A = 600.
HORIZON_B = 1200.
SUPPLY_SWEEP = tuple(float(x) for x in (0.05, 0.02, 0.01, 0.005, 0.002, 0.001, 0.0005, 0.0002, 0.0001, 0.00005))
CONCENTRATION_SCALE_MICROMOLAR = 1.  # illustrative dimensionalization only
TIME_SCALE_HOURS = 1.

import argparse
import csv
from dataclasses import dataclass, replace
from fractions import Fraction as F
import hashlib
import json
from pathlib import Path
import platform
import numpy as np
from scipy.integrate import solve_ivp
from scipy.optimize import least_squares

MANUSCRIPT_SHA256 = '8f00f7290e3e9a3f9a9d437d54ae386a3755ec295b37a4ef171d3ea3d449a835'


@dataclass(frozen=True)
class Reaction:
    """Directed monomial mass action: 2X uses x*x, no factorial divisor."""
    name: str
    inputs: tuple[int, ...]
    outputs: tuple[int, ...]
    rate: object

    def flux(self, x):
        value = self.rate
        for i, power in enumerate(self.inputs):
            if power: value = value * x[i]**power
        return value

    def derivative(self, x, j):
        if not self.inputs[j]: return 0
        value = self.rate * self.inputs[j]
        for i, power in enumerate(self.inputs):
            exponent = power - (i == j)
            if exponent > 0: value = value * x[i]**exponent
        return value


@dataclass(frozen=True)
class Resident:
    # Each directed rate can vary independently, including both reverse pairs.
    rates: tuple = (6, 27, 1, 1, 1, 1, F(1,100000), F(1,100000), 16, 2, 1, 1, F(1,10000))

    @classmethod
    def reference(cls, e=RESIDENT_E):
        rates=list(cls().rates); rates[6:8]=[e,e]; return cls(tuple(rates))


@dataclass(frozen=True)
class Consumers:
    copying: tuple
    loss: tuple
    limitation: tuple

    def __post_init__(self):
        n=len(self.copying)
        if n<2 or len(self.loss)!=n or len(self.limitation)!=n: raise ValueError('Need matching vectors for at least two consumers.')
        if any(float(v)<0 or not np.isfinite(float(v)) for vs in (self.copying,self.loss,self.limitation) for v in vs): raise ValueError('Rates must be finite and nonnegative.')

    @classmethod
    def equal(cls, n): return cls((1,)*n,(F(1,2),)*n,(n,)*n)

    @classmethod
    def target(cls, shares, kappa=1):
        q=tuple(F(str(v)) for v in shares); k=F(str(kappa))
        if k<=0 or any(v<=0 for v in q) or sum(q)!=1: raise ValueError('Positive shares must sum exactly to one, with positive kappa.')
        n=len(q); return cls((1,)*n,(F(1,2),)*n,tuple(k/v for v in q))

    def target_data(self):
        if len(set(self.copying))!=1 or len(set(self.loss))!=1 or min(self.limitation)<=0:
            raise ValueError('Composition theorem requires common growth and positive self-limitation.')
        k=1/sum(1/F(v) for v in self.limitation)
        return k, tuple(k/F(v) for v in self.limitation)


@dataclass(frozen=True)
class Reservoir:
    feed: object
    washout: object

    @classmethod
    def reference(cls, d): return cls(d,d)


class Reactor:
    """Composable literal reaction list; reservoir=None selects maintained A.

    Arbitrary rates define a model but do not inherit the paper's theorems.
    """
    def __init__(self, resident, consumers, reservoir=None):
        self.resident,self.consumers,self.reservoir=resident,consumers,reservoir
        self.offset=4+(reservoir is not None)
        self.names=('A','B','z','H')+(('R',) if reservoir else ())+tuple(f'X{i+1}' for i in range(len(consumers.copying)))
        self.reactions=self._build()
        if any(not np.isfinite(float(r.rate)) or float(r.rate)<0 for r in self.reactions): raise ValueError('Invalid directed rate.')
        self.matrix=np.array([[b-a for a,b in zip(r.inputs,r.outputs)] for r in self.reactions],dtype=float).T
        self.powers=np.array([r.inputs for r in self.reactions]);self.rates=np.array([r.rate for r in self.reactions],dtype=float)

    def _build(self):
        n=len(self.names); reactions=[]
        def add(name, inputs, outputs, rate):
            def vector(d): return tuple(d.get(i,0) for i in range(n))
            reactions.append(Reaction(name,vector(inputs),vector(outputs),rate))
        specs=[({}, {0:1}),({}, {1:1}),({0:1},{1:1,2:1}),({1:1,2:1},{0:1}),({0:1},{}),({1:1},{}),({1:1},{0:2}),({0:2},{1:1}),({2:1},{3:1}),({2:2},{3:1}),({3:1},{2:1}),({3:1},{2:2}),({3:1},{})]
        if len(self.resident.rates)!=13: raise ValueError('Exactly thirteen resident rates required.')
        for j,((a,b),k) in enumerate(zip(specs,self.resident.rates)): add(f'resident_{j+1}',a,b,k)
        for j,(k,mu,rho) in enumerate(zip(self.consumers.copying,self.consumers.loss,self.consumers.limitation)):
            i=self.offset+j; inputs={2:1,i:1}
            if self.reservoir: inputs[4]=1
            add(f'copy_{j+1}',inputs,{i:2},k);add(f'loss_{j+1}',{i:1},{},mu);add(f'limit_{j+1}',{i:2},{i:1},rho)
        if self.reservoir:
            add('reservoir_feed',{}, {4:1},self.reservoir.feed);add('reservoir_washout',{4:1},{},self.reservoir.washout)
        return tuple(reactions)

    def field(self, x):
        """Supports exact Fractions and SymPy as well as ordinary numbers."""
        values=[0]*len(x)
        for r in self.reactions:
            v=r.flux(x)
            for i,(a,b) in enumerate(zip(r.inputs,r.outputs)):
                if b!=a: values[i]=values[i]+(b-a)*v
        return values

    def jacobian(self,x):
        n=len(x); J=[[0]*n for _ in range(n)]
        for r in self.reactions:
            for j in range(n):
                v=r.derivative(x,j)
                for i,(a,b) in enumerate(zip(r.inputs,r.outputs)):
                    if b!=a:J[i][j]=J[i][j]+(b-a)*v
        return J

    def numerical_field(self,x): return self.matrix @ (self.rates*np.prod(np.asarray(x)[None,:]**self.powers,axis=1))

    def with_rates(self,rates):
        """Order: 13 resident, (copy,loss,limit) per consumer, then feed/washout."""
        if len(rates)!=len(self.reactions):raise ValueError('Incorrect rate vector length.')
        n=len(self.consumers.copying); c=rates[13:13+3*n]
        return Reactor(Resident(tuple(rates[:13])),Consumers(tuple(c[::3]),tuple(c[1::3]),tuple(c[2::3])),Reservoir(*rates[-2:]) if self.reservoir else None)

    def initial(self, resident=INITIAL_RESIDENT, consumers=INITIAL_CONSUMERS, reservoir=INITIAL_RESERVOIR):
        x=np.array([*resident,*([reservoir] if self.reservoir else []),*consumers],float)
        if len(x)!=len(self.names) or not np.all(np.isfinite(x)) or min(x)<=0:raise ValueError('Strictly positive initial state of matching size required.')
        return x


@dataclass
class Trajectory:
    time: np.ndarray
    state: np.ndarray
    ledgers: np.ndarray
    model: Reactor
    ledger_names=('tau','uptake','reservoir_washout','linear_loss','quadratic_loss','variance_loss')

    @property
    def consumers(self):return self.state[:,self.model.offset:]
    @property
    def total(self):return self.consumers.sum(axis=1)

    def balance_error(self):
        if not self.model.reservoir: return None
        inventory=self.state[:,4]+self.total
        return float(np.max(np.abs(inventory-inventory[0]-float(self.model.reservoir.feed)*self.time+self.ledgers[:,2:5].sum(axis=1))))

    def save(self,path):
        np.savetxt(path,np.column_stack([self.time,self.state,self.ledgers]),delimiter=',',header=','.join(('time',*self.model.names,*self.ledger_names)),comments='')


class Simulator:
    """Log consumers preserve rare positive species without artificial immigration.

    Integrated rewards make reservoir balances independent of sampling quadrature.
    Negative resident outputs or log underflow fail visibly; nothing is clipped.
    """
    def __init__(self,model):self.model=model

    def run(self, initial, times, method='Radau'):
        m=self.model; o=m.offset; initial=np.asarray(initial,float); times=np.asarray(times,float)
        if len(initial)!=len(m.names) or min(initial)<=0 or times[0]!=0 or np.any(np.diff(times)<=0):raise ValueError('Positive state and strictly increasing times starting at zero required.')
        try: kappa,q=m.consumers.target_data();q=np.array(q,float)
        except ValueError:kappa,q=None,None
        mu=np.array(m.consumers.loss,float);rho=np.array(m.consumers.limitation,float);copy=np.array(m.consumers.copying,float)
        def rhs(t,y):
            x=np.exp(y[o:len(initial)]); state=np.r_[y[:o],x];R=state[4] if m.reservoir else 1.;z=state[2];S=sum(x)
            field=m.numerical_field(state);growth=copy*R*z-mu-rho*x
            uptake=R*z*(copy@x);wash=float(m.reservoir.washout)*R if m.reservoir else 0.
            variance=float(rho@((x-q*S)**2)) if q is not None else 0.
            return np.r_[field[:o],growth,S,uptake,wash,mu@x,rho@(x*x),variance]
        y0=np.r_[initial[:o],np.log(initial[o:]),np.zeros(6)]
        sol=solve_ivp(rhs,(0,times[-1]),y0,t_eval=times,method=method,rtol=2e-10,atol=2e-12)
        if not sol.success:raise RuntimeError(sol.message)
        state=sol.y[:len(initial)].T.copy();state[:,o:]=np.exp(state[:,o:])
        if not np.all(np.isfinite(state)) or np.any(state<=0):raise RuntimeError('Nonpositive output or log underflow; reduce horizon or rescale.')
        return Trajectory(times,state,sol.y[len(initial):].T,m)


class OperatingAnalysis:
    @staticmethod
    def ceiling(d,kappa=1):
        if d<0 or kappa<=0:raise ValueError('Nonnegative supply and positive self-limitation required.')
        return 4*d/(np.sqrt(1+16*kappa*d)+1) # stable equivalent of (sqrt-1)/(4 kappa)

    @staticmethod
    def equilibrium(model,guess):
        """Positive numerical root; per-capita residual excludes zero-consumer roots."""
        o=model.offset
        def residual(logx):
            x=np.exp(logx);f=model.numerical_field(x);f[o:]/=x[o:];return f
        fit=least_squares(residual,np.log(guess),bounds=(-40,12),xtol=1e-13,ftol=1e-13,gtol=1e-13,max_nfev=2000)
        error=float(np.max(np.abs(residual(fit.x))))
        if not fit.success or error>2e-9:raise RuntimeError(f'Unresolved positive equilibrium: {error:g}')
        return np.exp(fit.x),error

    @staticmethod
    def spectral_split(model,state):
        k,q=model.consumers.target_data();q=np.array(q,float);o=model.offset;n=len(q);S=sum(state[o:])
        if not np.allclose(state[o:]/S,q,atol=1e-8):raise ValueError('State is not on the target manifold.')
        # Columns map (environment, total, n-1 composition perturbations) to x.
        T=np.zeros((o+n,o+n));T[:o,:o]=np.eye(o);T[o:,o]=q
        for j in range(n-1):T[o+j,o+1+j]=S;T[-1,o+1+j]=-S
        J=np.array(model.jacobian(state),float);transformed=np.linalg.solve(T,J@T)
        expected=-float(k)*S*np.eye(n-1)
        return {'composition_eigenvalue':-float(k)*S,'multiplicity':n-1,
            'composition_block_error':float(np.max(np.abs(transformed[o+1:,o+1:]-expected))),
            'coupling_error':float(max(np.max(np.abs(transformed[:o+1,o+1:])),np.max(np.abs(transformed[o+1:,:o+1])))),
            'reduced_eigenvalues':[[float(v.real),float(v.imag)] for v in np.linalg.eigvals(transformed[:o+1,:o+1])]}


@dataclass(frozen=True)
class UnitScales:
    concentration:float=CONCENTRATION_SCALE_MICROMOLAR
    time:float=TIME_SCALE_HOURS

    def __post_init__(self):
        if self.concentration<=0 or self.time<=0:raise ValueError('Positive unit scales required.')
    def concentration_value(self,x):return self.concentration*x
    def time_value(self,t):return self.time*t
    def rate_constant(self,k,order):return k/(self.concentration**(order-1)*self.time)
    def flux(self,j):return self.concentration*j/self.time


@dataclass(frozen=True)
class ExponentialFloor:
    """Exact positive prefactor * exp(-exponent); never coerce floor to float."""
    prefactor:F
    exponent:F

    def divided(self,n):return ExponentialFloor(self.prefactor/F(n),self.exponent)
    def record(self):return {'prefactor':str(self.prefactor),'negative_exponent':str(self.exponent),'expression':f'({self.prefactor}) * exp(-({self.exponent}))','scope':'Paper theorem evaluated symbolically; not a numerical abundance estimate or independently rerun global proof.'}


class PermanenceBounds:
    @staticmethod
    def aggregate(n,system='B',dmin=F(1,20),robust=False):
        if n<2 or system not in ('A','B') or (system=='B' and dmin<=0):raise ValueError('Outside theorem parameter domain.')
        b=F(288)/F(dmin) if system=='B' else F(0)
        gamma=F(1,16 if robust else 8) if system=='B' else F(1,8 if robust else 4)
        beta=n+(120000002+6*b if robust else 120000000+3*b) if system=='B' else n+(60000001 if robust else 60000000)
        return ExponentialFloor(gamma/(2*beta),F(300000000000)+b/2)

    @staticmethod
    def report(n,dmin=F(1,20),dmax=F(1,20)):
        if dmax<dmin:raise ValueError('Invalid supply interval.')
        result={}
        for system in ('A','B'):
            for robust in (False,True):
                floor=PermanenceBounds.aggregate(n,system,dmin,robust)
                label=system+('2' if robust else '1')
                recovery=floor.divided(8*(n+1) if robust else 2*n)
                result[label]={'aggregate':floor.record(),'consumer_recovery':recovery.record()}
                if robust:
                    result[label]['allowed_rate_radius']=f'min(1/10^18, {"dmin/4, 1/[1000(1+288/dmin)], " if system=="B" else ""}aggregate_floor/1000)'
        result['warning']='The full robust radius includes the exponentially tiny floor/1000. Using 1e-18 alone is NOT sufficient. Floors require strictly positive initial states; zero consumer faces stay invariant.'
        return result


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)
    resident=Resident.reference();n=len(INITIAL_CONSUMERS);equal=Consumers.equal(n)
    models={'maintained':Reactor(resident,equal),'reservoir':Reactor(resident,equal,Reservoir.reference(SUPPLY)),
        'target':Reactor(resident,Consumers.target(TARGET_SHARES,TARGET_KAPPA),Reservoir.reference(SUPPLY)),
        'zero_supply':Reactor(resident,equal,Reservoir.reference(0)),
        'no_limitation':Reactor(resident,Consumers((1,1),(F(1,2),F(3,4)),(0,0)))}
    trajectories={}
    for name,model in models.items():
        end=HORIZON_A if name in ('maintained','no_limitation') else HORIZON_B
        times=np.unique(np.r_[np.linspace(0,80,321),np.linspace(80,end,901)])
        initial=model.initial(consumers=(1e-8,1.)) if name=='no_limitation' else model.initial()
        trajectories[name]=Simulator(model).run(initial,times);trajectories[name].save(out/(name+'.csv'))
    main_run=trajectories['reservoir'];independent=Simulator(models['reservoir']).run(models['reservoir'].initial(),main_run.time,method='BDF')
    equilibrium,error=OperatingAnalysis.equilibrium(models['reservoir'],main_run.state[-1]);guess=equilibrium
    branch=[]
    for d in SUPPLY_SWEEP:
        model=Reactor(resident,equal,Reservoir.reference(d));guess,err=OperatingAnalysis.equilibrium(model,guess)
        branch.append([d,*guess,sum(guess[model.offset:]),OperatingAnalysis.ceiling(d),err])
    np.savetxt(out/'stationary_branch.csv',branch,delimiter=',',header=','.join(['supply',*models['reservoir'].names,'total','average_ceiling','per_capita_residual']),comments='')
    from certificates import local_certificate, donor_certificate
    local=local_certificate();donor=donor_certificate()
    (out/'local_certificate.json').write_text(json.dumps(local,indent=2)+'\n');(out/'donor_certificate.json').write_text(json.dumps(donor,indent=2)+'\n')
    result={'trajectories':{name:{'final_consumers':run.consumers[-1].tolist(),'final_total':float(run.total[-1]),'balance_max_error':run.balance_error(),
        'minimum_consumer':float(np.min(run.consumers)),'integrated_uptake':float(run.ledgers[-1,1]),'integrated_abundance':float(run.ledgers[-1,0]),'variance_loss':float(run.ledgers[-1,5])} for name,run in trajectories.items()},
        'equilibrium':equilibrium.tolist(),'per_capita_residual':error,'spectral_split':OperatingAnalysis.spectral_split(models['reservoir'],equilibrium),
        'independent_solver_max_difference':float(np.max(np.abs(main_run.state-independent.state))),
        'reservoir_copying_fraction':float(1-equilibrium[4]),'copying_flux':float(SUPPLY*(1-equilibrium[4])),
        'illustrative_units':{'concentration_scale_micromolar':CONCENTRATION_SCALE_MICROMOLAR,'time_scale_hours':TIME_SCALE_HOURS,
            'copying_flux_micromolar_per_hour':float(UnitScales().flux(SUPPLY*(1-equilibrium[4]))),'calibrated':False},
        'asymptotic_average_ceiling':OperatingAnalysis.ceiling(SUPPLY),'floors':PermanenceBounds.report(n,F(str(SUPPLY)),F(str(SUPPLY))),
        'scope':'Numerical trajectories do not prove permanence or global convergence. Exact local rational certificate applies only to its fixed reference center/rates and ellipsoid. Lean not rerun.'}
    (out/'results.json').write_text(json.dumps(result,indent=2)+'\n')
    lines=[f'Reference reservoir equilibrium total: {sum(equilibrium[5:]):.10g}; average ceiling: {OperatingAnalysis.ceiling(SUPPLY):.10g}.',
        f'Copying uses {100*(1-equilibrium[4]):.6g}% of reservoir inflow; unused reservoir washout uses the remainder.',
        f'Target consumers: {trajectories["target"].consumers[-1]}; proportions are prescribed by inverse self-limitation.',
        f'Reservoir balance error: {main_run.balance_error():.3g}; independent BDF comparison: {result["independent_solver_max_difference"]:.3g}.',
        f'Exact local certificate passes: {local["accepted"]}; consumer floor {local["consumer_floor"]} within its ellipsoid.',
        'Global floors are retained as exact prefactor times exp(-exponent), not rounded to zero.']
    (out/'console.txt').write_text('\n'.join(lines)+'\n');print('\n'.join(lines))
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axs=plt.subplots(1,2,figsize=(10,4.2),layout='constrained')
    for ax,key,title in zip(axs,['maintained','reservoir'],['Maintained precursor','Consumed dynamic reservoir']):
        run=trajectories[key]
        for i in range(n):ax.semilogy(run.time,run.consumers[:,i],label=f'X{i+1}')
        ax.set(xlabel='Time (nondimensional)',ylabel='Concentration (nondimensional)',title=title);ax.legend();ax.grid(alpha=.2)
    fig.savefig(out/'consumer_recovery.png',dpi=180);fig.savefig(out/'consumer_recovery.svg');plt.close(fig)
    fig,axs=plt.subplots(1,2,figsize=(10,4.2),layout='constrained')
    for key,label in [('reservoir','Equal target'),('target','Prescribed 1:3 target')]:
        run=trajectories[key];axs[0].plot(run.ledgers[:,0],run.consumers[:,0]/run.total,label=label)
    axs[0].set(xlabel='Accumulated abundance tau = integral S',ylabel='Consumer fraction X1 / S',title='Consumer proportions versus accumulated\nabundance');axs[0].legend(fontsize=8)
    ar=np.array(branch);ds=ar[:,0];z0=float(F(donor['z_interval'][0])+F(donor['z_interval'][1]))/2
    axs[1].loglog(ds,ar[:,-3],'.-',label='Numerical stationary total');axs[1].loglog(ds,ar[:,-2],'--',label='Proved time-average ceiling');axs[1].loglog(ds,2*(1-1/(2*z0))*ds,':',label='Small-supply tangent')
    axs[1].set(xlabel='Supply d (nondimensional)',ylabel='Total concentration',title='Stationary abundance and supply-dependent\nceiling');axs[1].legend(fontsize=8)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'composition_supply.png',dpi=180);fig.savefig(out/'composition_supply.svg');plt.close(fig)
    digest=lambda p:hashlib.sha256(p.read_bytes()).hexdigest();base=Path(__file__).resolve().parent
    (out/'run_metadata.json').write_text(json.dumps({'manuscript_sha256':MANUSCRIPT_SHA256,'source_sha256':digest(Path(__file__)),
        'input_sha256':{f:digest(base/f) for f in ('certificates.py','local_certificate_inputs.json')},'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()