"""Equilibrium geometry, independent kinetics and full mass-action dynamics."""
from dataclasses import dataclass
from fractions import Fraction as Q
import numpy as np
from scipy.integrate import solve_ivp
import phos_sharp as ps
import phos_capacity as pc

@dataclass(frozen=True)
class Geometry:
    roots:tuple
    enzyme_ratio:Q
    def __post_init__(self):
        object.__setattr__(self,'roots',tuple(Q(str(x)) for x in self.roots));object.__setattr__(self,'enzyme_ratio',Q(str(self.enzyme_ratio)))
        if not self.roots or len(self.roots)%2!=1 or min(self.roots)<=1 or self.enzyme_ratio<=0:raise ValueError('An odd nonempty list of roots >1 and positive enzyme ratio is required.')
    @property
    def n(self):return (len(self.roots)+1)//2
    def construct(self):
        rec=ps.build(list(self.roots),self.enzyme_ratio)
        if not rec['positive'] or max(ps.ratio(x) for x in self.roots)>=self.enzyme_ratio:raise ValueError('This geometry fails coefficient positivity or the admitted root domain.')
        return rec

@dataclass(frozen=True)
class Kinetics:
    catalytic_scales:tuple
    kinase_dissociation:tuple
    phosphatase_dissociation:tuple
    binding_scale:Q=Q(1)
    def rates(self,geometry):
        rec=geometry.construct();n=geometry.n
        arrays=[tuple(Q(str(x)) for x in a) for a in (self.catalytic_scales,self.kinase_dissociation,self.phosphatase_dissociation)];eps=Q(str(self.binding_scale))
        if any(len(a)!=n or min(a)<=0 for a in arrays) or not 0<eps<=1:raise ValueError('Positive site arrays of length n and binding scale in (0,1] required.')
        return pc.retune(rec,*arrays,eps)
    @staticmethod
    def from_coalesced_currents(n,r,currents=None,binding_scale=Q(1)):
        if type(n)is not int or n<2:raise ValueError('The coalesced-current construction needs n>=2.')
        rec=pc.coalesced(n,Q(r));Y=rec['states'][0]['z'][2*n+3:];currents=tuple(Q(str(x)) for x in (currents or [1]*n))
        if len(currents)!=n or min(currents)<=0:raise ValueError('Positive current for every site required.')
        return Kinetics(tuple(w/y for w,y in zip(currents,Y)),(Q(1),)*n,(Q(1),)*n,Q(binding_scale))

def stability(rates,state):
    matrix=ps.reduced_jacobian(rates,state);coef=ps.charpoly(matrix)
    try:count,signs=ps.routh_rhp(coef)
    except AssertionError:return dict(status='unresolved_regular_routh',characteristic=list(map(str,coef)))
    return dict(status='certified_regular_routh',unstable=count,signs=signs,characteristic=list(map(str,coef)),determinant=str(ps.det(matrix)))

def ordered_design(n,max_halvings=24):
    """Bounded implementation of the paper's ordered existence construction."""
    if type(n)is not int or not 2<=n<=12 or type(max_halvings)is not int or max_halvings<0:raise ValueError('Use 2<=n<=12 and a nonnegative stage budget.')
    r=Q(2)
    for _ in range(max_halvings+1):
        rec=ps.build([Q(3)]*(2*n-1),r)
        if rec['positive']:
            state=rec['states'][0]['z'];_,_,K=pc.loaded(n,state,2*r,Q(2))
            if pc.hurwitz([row[:n-1] for row in K[:n-1]]):break
        r*=2
    else:return dict(status='unknown_budget',stage='prefix')
    def simple_zero(matrix):
        cp=ps.charpoly(matrix)
        if cp[-1]!=0 or cp[-2]==0:return False
        try:return ps.routh_rhp(cp[:-1])[0]==0
        except AssertionError:return False
    last=Q(1)
    for _ in range(max_halvings+1):
        if simple_zero([[K[i][j]*(last if j==n-1 else 1) for j in range(n)] for i in range(n)]):break
        last/=2
    else:return dict(status='unknown_budget',stage='last_current')
    currents=(Q(1),)*(n-1)+(last,);kinetics=Kinetics.from_coalesced_currents(n,r,currents);eps=Q(1)
    for _ in range(max_halvings+1):
        candidate=Kinetics(kinetics.catalytic_scales,kinetics.kinase_dissociation,kinetics.phosphatase_dissociation,eps)
        if simple_zero(ps.reduced_jacobian(candidate.rates(Geometry((Q(3),)*(2*n-1),r)),state)):break
        eps/=2
    else:return dict(status='unknown_budget',stage='binding')
    delta=Q(1,2)
    for _ in range(max_halvings+1):
        xs=tuple(Q(3)+delta*(j-n+1) for j in range(2*n-1))
        if min(xs)>1:
            try:
                g=Geometry(xs,r);rec=g.construct();rates=candidate.rates(g);rows=[stability(rates,s['z']) for s in rec['states']]
                if all(row.get('unstable')==j%2 for j,row in enumerate(rows)):return dict(status='certified_design',n=n,r=r,delta=delta,last_current=last,binding_scale=eps,geometry=g,kinetics=candidate,unstable=[row['unstable'] for row in rows])
            except ValueError:pass
        delta/=2
    return dict(status='unknown_budget',stage='splitting')

class Reactor:
    """n-site physical kinetics in class coordinates q=(S1..Sn,C1..Cn,Y1..Yn).
    Independent totals and rates can be changed; no equilibrium is hard-coded.
    """
    def __init__(self,rates,totals):
        self.rates=np.array(rates,float);self.n=len(rates);n=self.n;self.totals=np.asarray(totals,float)
        if n<1 or self.rates.shape!=(n,6) or self.totals.shape!=(3,) or min(self.rates.ravel())<=0 or min(self.totals)<=0 or not np.all(np.isfinite(self.rates)) or not np.all(np.isfinite(self.totals)):raise ValueError('Positive finite six-rate site rows and three totals required.')
        self.indices=np.array(list(range(1,n+1))+list(range(n+3,3*n+3)));self.P=np.zeros((3*n+3,3*n));self.P[0]=-1;self.P[self.indices,np.arange(3*n)]=1;self.P[n+1,n:2*n]=-1;self.P[n+2,2*n:]=-1
        self.offset=np.zeros(3*n+3);self.offset[0]=self.totals[2];self.offset[n+1:n+3]=self.totals[:2]
        self.N=np.zeros((3*n+3,6*n));self.reactants=[]
        for i in range(n):
            E,F,C,Y=n+1,n+2,n+3+i,2*n+3+i
            for j,(reac,prod) in enumerate([([i,E],[C]),([C],[i,E]),([C],[i+1,E]),([i+1,F],[Y]),([Y],[i+1,F]),([Y],[i,F])]):
                for k in reac:self.N[k,6*i+j]-=1
                for k in prod:self.N[k,6*i+j]+=1
                self.reactants.append(reac)
        self.labels=[f'S{i}' for i in range(n+1)]+['E','F']+[f'C{i+1}' for i in range(n)]+[f'Y{i+1}' for i in range(n)]
    def species(self,q):return self.offset+self.P@np.asarray(q)
    def flux(self,z):return np.array([k*np.prod(z[reac]) for k,reac in zip(self.rates.ravel(),self.reactants)])
    def field(self,t,q):return (self.N@self.flux(self.species(q)))[self.indices]
    def jacobian(self,t,q):
        z=self.species(q);D=np.zeros((6*self.n,3*self.n+3))
        for j,(k,react) in enumerate(zip(self.rates.ravel(),self.reactants)):
            for l,s in enumerate(react):D[j,s]+=k*np.prod([z[v] for m,v in enumerate(react) if m!=l])
        return (self.N@D)[self.indices]@self.P
    def integrate(self,state,duration,samples=1001,method='LSODA'):
        state=np.asarray(state,float)
        if state.shape!=(3*self.n+3,) or np.min(state)<=0 or duration<=0:raise ValueError('Physical full initial state and positive duration required.')
        q=state[self.indices]
        if max(abs(self.species(q)-state))>1e-9:raise ValueError('Initial state does not have the selected conserved totals.')
        sol=solve_ivp(self.field,(0,duration),q,method=method,jac=self.jacobian,rtol=2e-10,atol=1e-13,dense_output=True)
        if not sol.success:raise RuntimeError(sol.message)
        times=np.r_[0,np.geomspace(max(duration*1e-7,1e-7),duration,samples-1)];Y=sol.sol(times);X=self.offset[:,None]+self.P@Y
        if np.min(X)<=0:raise ArithmeticError('Numerical positivity failed.')
        ledger=max(max(abs(np.array(ps.totals(list(z)),float)-self.totals)) for z in X.T)
        return dict(time=times,species=X,endpoint=X[:,-1],conservation_drift=float(ledger),minimum_concentration=float(np.min(X)))