"""Reusable full mass-action reactor, chart, protocols and numerical experiments.

All numerical trajectories and shooting results are numerical evidence, not
interval enclosures. The independent rational kernel lives in phos.py.
"""
from dataclasses import dataclass, replace
from typing import Protocol
import numpy as np
from scipy.integrate import solve_ivp
from scipy.optimize import root
import phos


@dataclass(frozen=True)
class EquilibriumDesign:
    substrate: tuple
    free_kinase: str
    free_phosphatase: str
    kinase_complexes: tuple
    phosphatase_complexes: tuple
    flux: tuple
    kinase_reverse: tuple
    phosphatase_reverse: tuple

    def exact(self, clamp_kinase=False, clamp_phosphatase=False):
        n=len(self.flux)
        if n<1 or len(self.substrate)!=n+1 or any(len(v)!=n for v in [self.kinase_complexes,self.phosphatase_complexes,self.kinase_reverse,self.phosphatase_reverse]):
            raise ValueError('One substrate level more than sites; all site arrays must match.')
        values=(*self.substrate,self.free_kinase,self.free_phosphatase,*self.kinase_complexes,*self.phosphatase_complexes,*self.flux,*self.kinase_reverse,*self.phosphatase_reverse)
        if any(phos.R_(v)<=0 for v in values):raise ValueError('Physical designs require strictly positive values.')
        return phos.Model(n,list(values[:3*n+3]),self.flux,self.kinase_reverse,self.phosphatase_reverse,clampE=clamp_kinase,clampF=clamp_phosphatase)

    def with_reverse(self,r):
        return replace(self,phosphatase_reverse=(str(r),*self.phosphatase_reverse[1:]))

    def append_site(self,load='1/100'):
        e=phos.R_(load)
        if e<=0:raise ValueError('A physical new site needs positive load.')
        return replace(self,substrate=(*self.substrate,str(2*e)),kinase_complexes=(*self.kinase_complexes,str(e)),phosphatase_complexes=(*self.phosphatase_complexes,str(e)),flux=(*self.flux,str(e)),kinase_reverse=(*self.kinase_reverse,'1'),phosphatase_reverse=(*self.phosphatase_reverse,'1'))


class BindingProtocol(Protocol):
    def factor(self,time:float)->float:...


@dataclass(frozen=True)
class ConstantBinding:
    multiplier:float=1.
    def factor(self,time):return self.multiplier


@dataclass(frozen=True)
class SinusoidalBinding:
    amplitude:float
    frequency:float
    def __post_init__(self):
        if not 0<=self.amplitude<1 or self.frequency<=0:raise ValueError('Positive periodically modulated binding required.')
    def factor(self,time):return 1+self.amplitude*np.cos(self.frequency*time)


class MassActionReactor:
    """Fixed rates on one compatibility class; deviation coordinates lift to species.

    Supply arbitrary positive equilibrium/flux data to the design builder, or
    replace rates on this instance through with_rates(). Reversible drive adds
    reverse catalysis and does NOT keep the designed irreversible equilibrium.
    """
    def __init__(self,model,protocol=None,drive=None,relaxation=1.):
        self.model=model;self.n=model.n;self.dim=3*self.n
        self.xstar=np.array(model.x,dtype=float)
        self.P=np.array(model.P,dtype=float);self.R=np.array(model.Rm,dtype=float);self.N=np.array(model.N,dtype=float)
        self.kon=np.array(model.kon,dtype=float);self.koff=np.array(model.koff,dtype=float);self.kcat=np.array(model.kcat,dtype=float)
        if relaxation<=0:raise ValueError('Relaxation scale must be positive.')
        self.kon/=relaxation;self.koff=(self.koff+self.kcat)/relaxation-self.kcat
        if np.any(self.koff<=0):raise ValueError('Relaxation path would create nonpositive dissociation.')
        self.lev=np.array(model.lev);self.enz=np.array(model.enz);self.cpx=np.array(model.cpx);self.out=np.array(model.out)
        self.protocol=protocol or ConstantBinding()
        if self.protocol.factor(0)<=0:raise ValueError('Binding multiplier must be positive.')
        self.reverse=np.zeros(2*self.n) if drive is None else self.kcat*self.kon/self.koff*np.exp(-float(drive))
        self.drive=drive
        self.labels=[f'S{i}' for i in range(self.n+1)]+['E','F']+[f'C{i+1}' for i in range(self.n)]+[f'D{i+1}' for i in range(self.n)]
        self.conservation=np.zeros((3,len(self.xstar)))
        self.conservation[0,model.iE]=1;self.conservation[0,self.cpx[:self.n]]=1
        self.conservation[1,model.iF]=1;self.conservation[1,self.cpx[self.n:]]=1
        self.conservation[2,:self.n+1]=1;self.conservation[2,self.cpx]=1

    def with_rates(self,kon,koff,kcat):
        if self.drive is not None:raise ValueError('Rebuild the driven reactor so detailed balance is recalculated.')
        other=MassActionReactor(self.model,self.protocol)
        arrays=[np.asarray(v,dtype=float) for v in (kon,koff,kcat)]
        if any(v.shape!=(2*self.n,) or np.any(v<=0) or not np.all(np.isfinite(v)) for v in arrays):raise ValueError('Each rate array must contain 2n positive finite values.')
        other.kon,other.koff,other.kcat=[v.copy() for v in arrays];return other

    def species(self,y):return self.xstar+self.P@np.asarray(y)

    def flux(self,time,x):
        on=self.kon.copy();on[self.n]*=self.protocol.factor(time)
        v=np.empty(6*self.n);v[0::3]=on*x[self.lev]*x[self.enz]
        v[1::3]=self.koff*x[self.cpx]
        v[2::3]=self.kcat*x[self.cpx]-self.reverse*x[self.out]*x[self.enz]
        return v

    def full_field(self,time,x):
        f=self.N@self.flux(time,x)
        if self.model.clampE:f[self.model.iE]=0
        if self.model.clampF:f[self.model.iF]=0
        return f

    def field(self,time,y):return self.R@self.full_field(time,self.species(y))

    def jacobian(self,time,y):
        x=self.species(y);on=self.kon.copy();on[self.n]*=self.protocol.factor(time)
        D=np.zeros((6*self.n,len(x)));j=np.arange(2*self.n)
        D[3*j,self.lev]=on*x[self.enz];D[3*j,self.enz]=on*x[self.lev]
        D[3*j+1,self.cpx]=self.koff;D[3*j+2,self.cpx]=self.kcat
        D[3*j+2,self.out]-=self.reverse*x[self.enz];D[3*j+2,self.enz]-=self.reverse*x[self.out]
        full=self.N@D
        if self.model.clampE:full[self.model.iE]=0
        if self.model.clampF:full[self.model.iF]=0
        return self.R@full@self.P

    def integrate(self,y0,duration,samples=1201,ledger=False):
        y0=np.asarray(y0,dtype=float)
        if duration<=0 or y0.shape!=(self.dim,) or np.min(self.species(y0))<=0:raise ValueError('Require positive time and a strictly positive initial species state.')
        def rhs(t,w):
            y=w[:self.dim];dy=self.field(t,y)
            return np.r_[dy,self.flux(t,self.species(y))[2::3]] if ledger else dy
        init=np.r_[y0,np.zeros(2*self.n)] if ledger else y0
        sol=solve_ivp(rhs,(0,duration),init,method='LSODA',rtol=2e-10,atol=2e-12,t_eval=np.linspace(0,duration,samples))
        if not sol.success:raise RuntimeError(sol.message)
        X=self.xstar[:,None]+self.P@sol.y[:self.dim]
        if np.min(X)<=0:raise RuntimeError('Numerical trajectory lost positivity; refine the solve.')
        return sol,X

    def static_jacobian(self):
        J=self.jacobian(0,np.zeros(self.dim));n=self.n
        return J[:n,:n]-J[:n,n:]@np.linalg.solve(J[n:,n:],J[n:,:n])

    def static_complexes(self,u_deviation):
        """Local numerical implicit solve; reject failed or unphysical branches."""
        u=np.asarray(u_deviation)
        initial=np.r_[u,np.zeros(2*self.n)]
        if np.min(self.species(initial))>0 and np.max(np.abs(self.field(0,initial)[self.n:]))<1e-12:
            return initial[self.n:],self.field(0,initial)[:self.n]
        result=root(lambda z:self.field(0,np.r_[u,z])[self.n:],np.zeros(2*self.n),jac=lambda z:self.jacobian(0,np.r_[u,z])[self.n:,self.n:])
        y=np.r_[u,result.x]
        if not result.success or np.max(np.abs(self.field(0,y)[self.n:]))>1e-9 or np.min(self.species(y))<=0:raise RuntimeError('No converged physical local algebraic solution found.')
        return result.x,self.field(0,y)[:self.n]

    def periodic_orbit(self,y0,period):
        """Numerical Newton shooting; never accepted as an interval certificate."""
        d=self.dim;anchor=np.asarray(y0).copy();direction=self.field(0,anchor)
        def flow(y,T):
            def rhs(t,w):return np.r_[self.field(t,w[:d]),(self.jacobian(t,w[:d])@w[d:].reshape(d,d)).ravel()]
            sol=solve_ivp(rhs,(0,T),np.r_[y,np.eye(d).ravel()],method='LSODA',rtol=2e-11,atol=2e-13)
            if not sol.success:raise RuntimeError(sol.message)
            return sol.y[:d,-1],sol.y[d:,-1].reshape(d,d)
        y=anchor.copy();T=float(period)
        for _ in range(15):
            end,V=flow(y,T);res=np.r_[end-y,direction@(y-anchor)]
            if np.max(abs(res))<3e-10:break
            bordered=np.block([[V-np.eye(d),self.field(0,end)[:,None]],[direction[None,:],np.zeros((1,1))]])
            change=np.linalg.solve(bordered,-res);y+=change[:d];T+=change[-1]
            if T<=0 or np.min(self.species(y))<=0:raise RuntimeError('Shooting left the physical branch.')
        end,V=flow(y,T);residual=float(np.max(abs(end-y)))
        if residual>1e-8 or np.linalg.norm(self.field(0,y))<1e-5:raise RuntimeError('Shooting did not produce a nonstationary periodic orbit.')
        return y,T,np.linalg.eigvals(V),residual