"""Full consumer dynamics and readouts; no inherited-population theorem is imported."""
from dataclasses import dataclass
from fractions import Fraction as F
import math
import numpy as np
from scipy.integrate import solve_ivp
from scipy.optimize import brentq
from resident import ResidentChemistry,reconstruction,residual

S_BOUNDS=(('0.04067409','0.04067412'),('0.07014930','0.07014933'))


def uptake(s):return s/2+s*s
def load(s):return s-10*s*s-20*s*s*s
def stationary(s):
    J=uptake(s);R=1-20*J;z=(F(1,2)+s)/R
    return (*reconstruction(z,R*s),R,s)
def stationary_residual(s):
    A,B,z,H,R,S=stationary(s)
    return residual(z,R*S)


class Community:
    def __init__(self,q=(.2,.3,.5),renewal=.05,mortality=.5):
        q=np.array(q,float)
        if q.ndim!=1 or len(q)==0 or not np.isfinite(q).all() or np.any(q<=0) or abs(sum(q)-1)>1e-12:raise ValueError('Positive probability vector required')
        if not math.isfinite(renewal) or not math.isfinite(mortality) or min(renewal,mortality)<=0:raise ValueError('Positive finite supply and mortality required')
        self.q=q;self.d=renewal;self.mu=mortality;self.resident=ResidentChemistry('0')
    def field(self,state):
        R=state[4];X=np.asarray(state[5:]);S=X.sum();z=state[2]
        return np.r_[self.resident.density(state[:4],R*S),self.d*(1-R)-R*z*S,X*(R*z-self.mu-X/self.q)]
    def observables(self,state):
        R=state[4];X=np.asarray(state[5:]);S=float(X.sum());z=state[2]
        W=float(np.sum(self.q*(X/self.q-S)**2));J=R*z*S
        p=X/S if S>0 else np.full(len(X),np.nan)
        chi=float(np.sum((p-self.q)**2/self.q)) if S>0 else None
        return dict(S=S,W=W,J=J,load=R*S,p=p,chi=chi)
    def composition(self,state):
        o=self.observables(state);p=o['p'];S=o['S']
        if S<=0:raise ValueError('Positive abundance required')
        ratio=p/self.q;kappa=float(np.sum(p*ratio));v=p*(kappa-ratio)
        pdot=S*v;chi_prime=-2*S*float(np.sum(p*(ratio-kappa)**2))
        return dict(p=p,v=v,pdot=pdot,chi_prime=chi_prime,KL_prime=-S*o['chi'],
            derivative_condition=math.inf if np.linalg.norm(v)<1e-14 else 1/np.linalg.norm(v))
    def equilibria(self):
        if self.d!=.05 or self.mu!=.5:raise ValueError('Paper root brackets require d=.05 and mu=.5')
        result=[]
        for a,b in S_BOUNDS:
            s=brentq(lambda s:float(stationary_residual(s)),float(a),float(b),xtol=1e-15)
            u=np.array(stationary(s),float);result.append(np.r_[u[:5],self.q*s])
        return result
    def jacobian(self,state):
        n=len(state)
        return np.column_stack([self.field(np.asarray(state,complex)+1e-20j*np.eye(n)[i]).imag/1e-20 for i in range(n)])
    def integrate(self,initial,duration=500.,samples=1001):
        initial=np.array(initial,float)
        if initial.shape!=(5+len(self.q),) or not np.isfinite(initial).all() or np.any(initial<=0) or duration<=0:raise ValueError('Positive finite full state and duration required')
        # Independent accumulated uptake, abundance, square abundance, dispersion, reservoir and exposure.
        n=len(initial)
        def rhs(t,u):
            state=u[:n];o=self.observables(state)
            return np.r_[self.field(state),o['J'],o['S'],o['S']**2,o['W'],state[4]]
        times=np.linspace(0,duration,samples)
        sol=solve_ivp(rhs,(0,duration),np.r_[initial,np.zeros(5)],method='Radau',rtol=2e-10,atol=2e-12,t_eval=times)
        if not sol.success:raise RuntimeError(sol.message)
        states=sol.y[:n].T;integrals=sol.y[n:].T
        if states.min()<=0:raise RuntimeError('Lost numerical positivity')
        S0=self.observables(initial)['S'];S1=self.observables(states[-1])['S'];a=integrals[-1]
        return dict(times=times,states=states,integrals=integrals,
            abundance_balance_error=float(a[0]-(S1-S0+self.mu*a[1]+a[2]+a[3])),
            reservoir_balance_error=float(a[0]-(self.d*(duration-a[4])-(states[-1,4]-initial[4]))))


@dataclass(frozen=True)
class WindowReadout:
    h:F=F(100)
    Smax:F=F(2,25)
    epsS:F=F(1,1000)
    epsW:F=F(1,10000)
    epsR:F=F(1,200)
    eQ:F=F(1,100000)
    eQR:F=F(1,100000)
    def __post_init__(self):
        for name in self.__dataclass_fields__:object.__setattr__(self,name,F(getattr(self,name)))
        if self.h<=0 or min(self.Smax,self.epsS,self.epsW,self.epsR,self.eQ,self.eQR)<0:raise ValueError('Invalid window/error ceilings')
    @property
    def errors(self):return (2*self.epsS/self.h+(F(1,2)+2*self.Smax)*self.epsS+self.epsS**2+self.epsW+self.eQ,
        self.epsR/20+2*self.epsR/self.h+self.eQR)
    def estimates(self,time,abundance,dispersion,reservoir):
        """Numerical trapezoid estimates; configured quadrature ceilings remain assumptions."""
        t,S,W,R=map(lambda x:np.asarray(x,float),(time,abundance,dispersion,reservoir))
        if t.ndim!=1 or len(t)<2 or any(v.shape!=t.shape or not np.isfinite(v).all() for v in (S,W,R)) or not np.isfinite(t).all() or np.any(np.diff(t)<=0) or abs(t[-1]-t[0]-float(self.h))>1e-9:
            raise ValueError('Matching finite traces on an increasing mesh of the declared window length required')
        h=float(self.h)
        return ((S[-1]-S[0])/h+np.trapezoid(S/2+S*S+W,t)/h,
            .05*(1-np.trapezoid(R,t)/h)-(R[-1]-R[0])/h)
    def recovery_error(self,E0,B,t0,*,admitted_sublevel=False):
        """Requires the local invariant-sublevel premise, not just proximity in a plot."""
        if not admitted_sublevel:return None
        if min(E0,B,t0)<0 or not all(math.isfinite(v) for v in (E0,B,t0)):raise ValueError('Nonnegative finite recovery inputs required')
        AJ=3*B*B*math.sqrt(2000*E0);h=float(self.h)
        return AJ*math.exp(-t0/800)*800/h*(-math.expm1(-h/800))
    def classify(self,js,jr,*,calibrated=False,two_alternatives=False,recovery_bound=None):
        if not calibrated or not two_alternatives or recovery_bound is None:return dict(outcome='unresolved',reason='Calibration, alternatives and admitted recovery bound are required')
        js,jr,dyn=map(F,(js,jr,recovery_bound))
        if dyn<0:raise ValueError('Nonnegative recovery error required')
        es,er=self.errors;bounds=[tuple(uptake(F(s)) for s in ab) for ab in S_BOUNDS];gap=bounds[1][0]-bounds[0][1]
        output=dict(E_S=es,E_R=er,recovery_error=dyn)
        if abs(js-jr)>es+er:return dict(output,outcome='unresolved',reason='Inconsistent balance measurements')
        candidates=[name for name,(lo,hi) in zip(('low','high'),bounds) if max(lo,js-es-dyn,jr-er-dyn)<=min(hi,js+es+dyn,jr+er+dyn)]
        if len(candidates)!=1 or min(es,er)+dyn>=gap/2:return dict(output,outcome='unresolved',reason='No unique alternative within the guaranteed error regime')
        return dict(output,outcome=candidates[0],reason='Conditional on supplied calibration, alternatives and recovery assumptions')