"""Literal finite-bath count law, pulse histories, metered feeds and density diagnostics."""
from dataclasses import dataclass,asdict
from fractions import Fraction as Q
import copy
import numpy as np
from scipy.integrate import solve_ivp
from chemistry import Reactor,Intervention,MolecularPulse,counts,restart,dot,A,B,I,Y,LEDGER,ceildiv


@dataclass(frozen=True)
class Bath:
    capacity: int
    fuel: int
    waste: int

    def __post_init__(self):
        if any(type(x) is not int for x in (self.capacity,self.fuel,self.waste)) or self.capacity<=0 or min(self.fuel,self.waste)<0:raise ValueError('Positive integer capacity and nonnegative integer bath counts required.')

    @property
    def inventory(self):return self.fuel+self.waste

    def step(self,label):
        j=int(label==18)-int(label==19)
        return Bath(self.capacity,self.fuel-j,self.waste+j)


class FiniteBathReactor:
    def __init__(self,release='20',drive='1/50'):
        self.internal=Reactor(release,drive);self.r=self.internal.r;self.d=self.internal.d

    def admitted(self,bath):return 19<=self.r<=21 and self.d>=0 and self.d*bath.inventory<=Q(bath.capacity,25)

    def propensities(self,N,V,bath):
        N=counts(N)
        if type(V) is not int or V<=0:raise ValueError('Positive integer copy scale required.')
        a=[c.propensity(N,V) for c in self.internal.channels]
        a[18]*=Q(bath.fuel,bath.capacity);a[19]*=Q(bath.waste,bath.capacity)
        return a

    def rates(self,N,V,bath):
        a=self.internal.rates(N,V);a[18]*=bath.fuel/bath.capacity;a[19]*=bath.waste/bath.capacity
        return a

    def exact_drift(self,N,V,bath,weights):
        return sum(a*dot(weights,c.jump) for a,c in zip(self.propensities(N,V,bath),self.internal.channels))


@dataclass
class FoodMeter:
    u: int
    w: int

    def __post_init__(self):
        if any(type(x) is not int or x<0 for x in (self.u,self.w)):raise ValueError('Nonnegative integer stocks required.')

    def pay(self,u,w):
        if u>self.u or w>self.w:return False
        self.u-=u;self.w-=w;return True


class CountMission:
    def __init__(self,reactor):self.reactor=reactor

    def run(self,initial,V,bath,cycles,controller,seed,event_budget=200000,meter=None):
        if type(V) is not int or V<=0 or type(cycles) is not int or cycles<0 or type(event_budget) is not int or event_budget<0:raise ValueError('Invalid count mission.')
        N=counts(initial);rng=np.random.default_rng(seed);history=[];pulse=MolecularPulse();original_bath=bath;used=0;forward=reverse=0
        min_f=max_f=bath.fuel
        for cycle in range(cycles):
            intervention=controller(copy.deepcopy(history));before=N;incoming=bath;draw=pulse.sample(N,V,intervention,rng)
            if meter and not meter.pay(*draw.doses):return dict(status='food_shutdown_at_pulse',history=history,failed_cycle=cycle+1,total_events=used)
            N=draw.start;t=0.;v=[0,0,*draw.doses,0,0,0,0,0];start_used=used
            while t<4:
                rates=self.reactor.rates(N,V,bath);total=float(rates.sum());next_t=t+rng.exponential(1/total)
                if next_t>4:break
                if used>=event_budget:return dict(status='event_budget_exhausted',history=history,partial_cycle=cycle+1,total_events=used,success=None)
                if next_t<=t:raise ArithmeticError('Time precision exhausted; lower the sampled copy scale.')
                label=min(int(np.searchsorted(np.cumsum(rates),rng.random()*total,side='right')),19);t=next_t
                if meter and label in (10,11) and not meter.pay(int(label==10),int(label==11)):return dict(status='food_shutdown_at_feed',history=history,failed_cycle=cycle+1,total_events=used)
                c=self.reactor.internal.channels[label]
                if self.reactor.propensities(N,V,bath)[label]<=0:raise ArithmeticError('Unavailable event sampled.')
                N=tuple(n+x for n,x in zip(N,c.jump));bath=bath.step(label)
                v=[a+b for a,b in zip(v,(*c.marks(t>=3),int(label==18),int(label==19)))];used+=1
                forward+=int(label==18);reverse+=int(label==19)
                assert bath.fuel==original_bath.fuel-forward+reverse and bath.waste==original_bath.waste+forward-reverse
                assert abs(bath.fuel-original_bath.fuel)<=forward+reverse
                min_f=min(min_f,bath.fuel);max_f=max(max_f,bath.fuel)
            net=dot(I,N)-dot(I,before)+v[5]+dot(I,draw.withdrawn)+dot(I,draw.lost)
            assert net==v[6] and bath.inventory==original_bath.inventory
            success=restart(N,V) and v[0]>=ceildiv(V,56) and v[1]>=ceildiv(V,1080) and v[2]<=5*V and v[3]<=5*V and v[4]<=V//5
            history.append(dict(cycle=cycle+1,start=list(before),incoming_bath=asdict(incoming),pulse=asdict(draw),endpoint=list(N),bath=asdict(bath),events=used-start_used,counters=dict(zip((*LEDGER,'forward','reverse'),v)),success=bool(success),sharp_success=bool(success and v[4]<=V//10),inventory_residual=v[6]-net))
        return dict(status='complete',history=history,total_events=used,all_success=all(r['success'] for r in history),all_sharp_success=all(r['sharp_success'] for r in history),bath_prefix_fuel_min=min_f,bath_prefix_fuel_max=max_f,total_forward=forward,total_reverse=reverse,net_service=forward-reverse,remaining_food=asdict(meter) if meter else None)


class HistoryController:
    """An example history policy: retain more after a failed returned cycle."""
    def __call__(self,history):return Intervention(q='3/4' if history and not history[-1]['success'] else '1/4')


class DensityMission:
    """Eight changing concentrations plus nine counters; the bath is never reset.

    Bath entries are f/R and p/R, internal entries N/V. Thus the bath derivatives
    contain V/R. Mean pulses and continuous doses are density illustrations.
    """
    def __init__(self,reactor,capacity_per_V,total_activity):
        if capacity_per_V<=0 or total_activity<0:raise ValueError('Invalid reservoir density.')
        self.reactor=reactor;self.sigma=float(capacity_per_V);self.total=float(total_activity)

    def run(self,initial,cycles,intervention=Intervention(),method='Radau'):
        y=np.r_[np.array(initial,float),np.zeros(9)];history=[];trajectory=[]
        if len(y)!=17 or y[:8].min()<0 or abs(y[6]+y[7]-self.total)>1e-12:raise ValueError('Initial internal six plus two bath activities required.')
        for k in range(cycles):
            before=y[:6].copy();incoming=y[6:8].copy();withdrawn=(1-float(intervention.q))*before;lost=float(intervention.q)*(1-np.array(intervention.survival,float))*before
            y[:6]*=float(intervention.q)*np.array(intervention.survival,float);doses=1-float(intervention.q)+np.array(intervention.refill,float);y[:2]+=doses;y[8:]=0;y[10:12]=doses
            for lo,hi in [(0,3),(3,4)]:
                marks=np.array([(*c.marks(lo==3),int(c.label==18),int(c.label==19)) for c in self.reactor.internal.channels],float)
                def rhs(_,state):
                    rates=self.reactor.internal.rates(state[:6]);rates[18]*=state[6];rates[19]*=state[7];net=rates[18]-rates[19]
                    return np.r_[rates@self.reactor.internal.jumps,-net/self.sigma,net/self.sigma,rates@marks]
                sol=solve_ivp(rhs,(lo,hi),y,method=method,rtol=2e-10,atol=1e-13,t_eval=np.linspace(lo,hi,61))
                if not sol.success:raise RuntimeError(sol.message)
                trajectory.extend(dict(cycle=k+1,time=4*k+float(t),fuel_activity=float(z[6]),waste_activity=float(z[7]),free_template=float(z[2]),template_inventory=float(dot(I,z[:6]))) for t,z in zip(sol.t,sol.y.T));y=sol.y[:,-1]
            v=y[8:];res=v[6]-(dot(I,y[:6])-dot(I,before)+v[5]+dot(I,withdrawn)+dot(I,lost))
            bath_res=self.sigma*(incoming[0]-y[6])-(v[7]-v[8])
            history.append(dict(cycle=k+1,endpoint=y[:8].tolist(),counters=dict(zip((*LEDGER,'forward','reverse'),v.tolist())),inventory_residual=float(res),bath_residual=float(bath_res)))
        return dict(history=history,trajectory=trajectory)