import json
from pathlib import Path
import numpy as np


def free_energy(x, x_star):
    x = np.maximum(np.asarray(x, dtype=float), 1e-15)
    xs = np.asarray(x_star, dtype=float)
    v = x * np.log(x / xs) - x + xs
    return np.sum(v, axis=-1) if x.ndim else float(v)


class PositiveMassAction:
    """Positive mass-action state block with RK4 integration.

    reactants[j,i]=v_ij and products[j,i]=v'_ij. Rates are positive.
    """
    def __init__(self, reactants, products, dt=0.2, substeps=2, eps=1e-6):
        self.reactants = np.asarray(reactants, float)
        self.stoich = np.asarray(products, float) - self.reactants
        self.dt, self.substeps, self.eps = dt, substeps, eps

    def field(self, x, rates):
        x = np.maximum(np.asarray(x, float), self.eps)
        rates = np.maximum(np.asarray(rates, float), self.eps)
        monomial = np.prod(x[..., None, :] ** self.reactants, axis=-1)
        return np.sum(monomial[..., :, None] * rates[..., :, None] * self.stoich, axis=-2)

    def step(self, x, rates):
        x = np.asarray(x, float)
        h = self.dt / self.substeps
        for _ in range(self.substeps):
            k1 = self.field(x, rates)
            k2 = self.field(x + h*k1/2, rates)
            k3 = self.field(x + h*k2/2, rates)
            k4 = self.field(x + h*k3, rates)
            x = np.maximum(x + h*(k1+2*k2+2*k3+k4)/6, self.eps)
        return x

    def iss_residual(self, x, x_next, rates, rates_star, c, k):
        v = free_energy(x, x_star=1.0)
        v_next = free_energy(x_next, x_star=1.0)
        d2 = np.sum((np.asarray(rates)-np.asarray(rates_star))**2)
        return (v_next-v)/self.dt + c*v - k*d2


def euler(x, u, k, dt):
    return x + dt*(u-k*x)


def rk4_birth_death(x, u, k, dt):
    f=lambda z: u-k*z
    a=f(x); b=f(x+dt*a/2); c=f(x+dt*b/2); d=f(x+dt*c)
    return x+dt*(a+2*b+2*c+d)/6


def verify():
    k=u0=xs=1.0
    # Prediction 1: Euler multiplier is 1-k*dt, so stability boundary is dt*k=2.
    lo, hi = 0., 3.
    for _ in range(45):
        dt=(lo+hi)/2; x=1.2
        for _ in range(150): x=euler(x,u0,k,dt)
        if abs(x-xs)<1e-4: lo=dt
        else: hi=dt
    # Prediction 2: shifted equilibrium free energy is quadratic in rate amplitude.
    amps=np.array([.01,.02,.04,.08,.16,.25,.35])
    Vs=np.array([free_energy((u0+a)/k,xs) for a in amps])
    slope=np.polyfit(np.log(amps),np.log(Vs),1)[0]
    # Prediction 3: local V decays at twice the state decay rate.
    x=1.2; dt=.002; vals=[]
    for _ in range(500):
        vals.append(free_energy(x,xs)); x=rk4_birth_death(x,u0,k,dt)
    decay=np.polyfit(np.arange(100,450)*dt,np.log(np.maximum(vals[100:450],1e-30)),1)[0]
    ratios=Vs/amps**2
    # Direct mass-action formula check: X -> 2X and X -> empty gives u*x - k*x.
    block=PositiveMassAction([[0],[1]], [[1],[0]], dt=.1)
    field_error=abs(float(block.field(np.array([1.3]),np.array([u0,k]))[0])-(u0-k*1.3))
    return {'euler_boundary_observed':float(lo),'euler_boundary_predicted':2.0,
            'steady_energy_log_slope_observed':float(slope),'steady_energy_log_slope_predicted':2.0,
            'small_amplitude_V_over_a2_observed':float(np.mean(ratios[:3])),
            'small_amplitude_V_over_a2_predicted':.5,
            'nominal_log_V_rate_observed':float(decay),'nominal_log_V_rate_predicted':-2.0,
            'mass_action_field_absolute_error':float(field_error),
            'amplitudes':amps.tolist(),'steady_V':Vs.tolist(),'V_over_a2':ratios.tolist()}


def sequence_demo(seed=7):
    rng=np.random.default_rng(seed); T=50; nseq=60
    X=rng.uniform(-1,1,(nseq,T)); Y=np.zeros_like(X)
    for b in range(nseq):
        y=0.
        for t in range(T): y=.9*y+.1*X[b,t]; Y[b,t]=y
    sp=lambda z: np.logaddexp(0,z)
    def pos_loss(p):
        a,b,w,c=p; x=np.ones(nseq); se=0.
        for t in range(T):
            x += .2*(sp(a*X[:,t]+b)-x)
            se += np.sum((w*x+c-Y[:,t])**2)
        return se/(nseq*T)
    def tanh_loss(p):
        a,b,w,c=p; h=np.zeros(nseq); se=0.
        for t in range(T):
            h=np.tanh(a*X[:,t]+b*h)
            se += np.sum((w*h+c-Y[:,t])**2)
        return se/(nseq*T)
    best=(1e9,None); best2=(1e9,None)
    for _ in range(500):
        p=rng.normal(0,1,4); p[2]=rng.normal(0,2); p[3]=rng.normal(0,.3)
        z=pos_loss(p)
        if z<best[0]: best=(z,p.copy())
        p=rng.normal(0,1,4); z=tanh_loss(p)
        if z<best2[0]: best2=(z,p.copy())
    def long_eval(p,H,positive):
        a,b,w,c=p; errs=[]
        for row in X:
            state=1. if positive else 0.; yy=0.
            for t in range(H):
                q=row[t%T]+.15*rng.normal(); yy=.9*yy+.1*q
                state=(state+.2*(sp(a*q+b)-state)) if positive else np.tanh(a*q+b*state)
                errs.append((w*state+c-yy)**2)
        return float(np.mean(errs))
    return {'positive_train_mse':float(best[0]),'tanh_train_mse':float(best2[0]),
            'positive_long_noisy_mse':long_eval(best[1],400,True),
            'tanh_long_noisy_mse':long_eval(best2[1],400,False),
            'positive_params':best[1].tolist(),'tanh_params':best2[1].tolist()}


if __name__ == '__main__':
    out={'verification':verify(),'sequence_demo':sequence_demo()}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))