"""Finite-reagent sandwich assays: source models, exact margins, outer inference.

Run: python example.py --output outputs
Synthetic paper inputs; no assay calibration or clinical interpretation is inferred.
"""
from __future__ import annotations
from dataclasses import dataclass, replace, asdict
from fractions import Fraction as F
from pathlib import Path
import argparse
import csv
import hashlib
import json
import math
import platform
import numpy as np
from scipy.integrate import solve_ivp
from scipy.optimize import brentq
from scipy.stats import norm

# EDITABLE INPUTS: final reaction concentrations in nM; raw signal in separate units.
C, D, K, J = 1., 1., 1., 1.
DILUTION = 10.
NATIVE_AVAILABILITY = 1.
ON_RATES = (0.001, 0.001)  # nM^-1 s^-1 = 10^6 M^-1 s^-1
NATIVE_CEILING = F(100)  # External operating-domain assumption, never inferred.
READING_ERROR = F(1, 10000)  # Simultaneous absolute allowance, NOT a CV.
ENCLOSURE_DEPTH = 14
CELL_BUDGET = 40000
CALIBRATION_ESTABLISHED = False
CARRY_OVER = 0.001
MANUSCRIPT_SHA256 = '9d7e87368f1db8e2548ecfc878dab2b4d56b42faef68d1734bc9262526c8116a'


@dataclass(frozen=True)
class BindingSite:
    capacity: float
    dissociation: float

    def __post_init__(self):
        if not all(math.isfinite(v) and v > 0 for v in (self.capacity, self.dissociation)):
            raise ValueError('Positive finite capacity and dissociation constant required')

    def occupancy(self, u):
        u = np.asarray(u, dtype=float)
        if np.any(~np.isfinite(u)) or np.any(u < 0):
            raise ValueError('Nonnegative finite accessible concentration required')
        c, k = self.capacity, self.dissociation
        # Positive terms avoid subtracting the two almost equal quadratic roots.
        return 2*c / (u+c+k+np.sqrt((u-c)**2+2*k*(u+c)+k*k))

    def captured(self, u):
        return np.asarray(u)*self.occupancy(u)

    def elasticity(self, u):
        p = self.occupancy(u)
        return np.asarray(u)*(1-p)**2/(np.asarray(u)*(1-p)**2+self.dissociation)


@dataclass(frozen=True)
class SandwichSource:
    capture: BindingSite
    detector: BindingSite

    def signal(self, u):
        return np.asarray(u)*self.capture.occupancy(u)*self.detector.occupancy(u)

    def species(self, u):
        p, q = self.capture.occupancy(u), self.detector.occupancy(u)
        return np.array([u*(1-p)*(1-q), u*p*(1-q), u*(1-p)*q, u*p*q])

    def log_slope(self, u):
        return 1-self.capture.elasticity(u)-self.detector.elasticity(u)

    def peak(self):
        scale = max(self.capture.capacity+self.capture.dissociation,
                    self.detector.capacity+self.detector.dissociation)
        return brentq(self.log_slope, scale*1e-10, scale*1e10)

    def paired(self, x, rho=1., dilution=10., gain=1., drift=1.):
        if not 0 < rho <= 1 or min(dilution, gain, drift) <= 0:
            raise ValueError('Invalid preparation or gains')
        u = rho*x
        return gain*self.signal(u), gain*drift*self.signal(u/dilution)


@dataclass(frozen=True)
class SequentialSource:
    source: SandwichSource
    carry_over: float = 0.

    def __post_init__(self):
        if not math.isfinite(self.carry_over) or not 0 <= self.carry_over <= 1:
            raise ValueError('Carry-over fraction must lie in [0,1]')

    def signal(self, u):
        w = self.source.capture.captured(u)
        return w*self.source.detector.occupancy(w+self.carry_over*(np.asarray(u)-w))

    @property
    def plateau(self):
        return self.source.detector.captured(self.source.capture.capacity)

    @staticmethod
    def wash_limit(detector_min: F, upper: F, retained_fraction: F):
        if detector_min <= 0 or upper <= 0 or not 0 < retained_fraction <= 1:
            raise ValueError('Invalid wash specification')
        return detector_min*(1/retained_fraction-1)/upper


@dataclass(frozen=True)
class BindingKinetics:
    source: SandwichSource
    on_capture: float
    on_detector: float

    def __post_init__(self):
        if not all(math.isfinite(v) and v > 0 for v in (self.on_capture, self.on_detector)):
            raise ValueError('Positive finite association rates required')

    def rates(self, u):
        return [on*math.sqrt((u-s.capacity)**2+2*s.dissociation*(u+s.capacity)+s.dissociation**2)
                for s, on in [(self.source.capture, self.on_capture), (self.source.detector, self.on_detector)]]

    def attenuation(self, u, t):
        a, b = self.rates(u)
        return -np.expm1(-a*np.asarray(t)) * -np.expm1(-b*np.asarray(t))

    def uniform_rate(self):
        return min(on*max(s.dissociation, 2*math.sqrt(s.capacity*s.dissociation))
                   for s, on in [(self.source.capture, self.on_capture), (self.source.detector, self.on_detector)])

    def trajectory(self, u, times, full=True):
        """Independent full four-species mass action or reduced site equations; unbound start."""
        if u < 0 or not np.all(np.diff(times) > 0) or times[0] != 0:
            raise ValueError('Nonnegative u and increasing times starting at zero required')
        c, k = self.source.capture.capacity, self.source.capture.dissociation
        d, j = self.source.detector.capacity, self.source.detector.dissociation
        kc, kd = self.on_capture, self.on_detector
        def rhs(t, y):
            if not full:
                p, q = y
                return [kc*((c-u*p)*(1-p)-k*p), kd*((d-u*q)*(1-q)-j*q)]
            a00, a10, a01, a11 = y
            cf, df = c-a10-a11, d-a01-a11
            f0, f1 = kc*(cf*a00-k*a10), kc*(cf*a01-k*a11)
            g0, g1 = kd*(df*a00-j*a01), kd*(df*a10-j*a11)
            return [-f0-g0, f0-g1, g0-f1, f1+g1]
        y0 = [u, 0, 0, 0] if full else [0, 0]
        sol = solve_ivp(rhs, (0, times[-1]), y0, t_eval=times, method='Radau', rtol=2e-10, atol=1e-13)
        if not sol.success:
            raise RuntimeError(sol.message)
        if full:
            return sol.y.T
        p, q = sol.y
        return np.array([u*(1-p)*(1-q), u*p*(1-q), u*(1-p)*q, u*p*q]).T


@dataclass(frozen=True)
class Interval:
    lo: F
    hi: F

    def __post_init__(self):
        object.__setattr__(self, 'lo', F(self.lo)); object.__setattr__(self, 'hi', F(self.hi))
        if self.lo > self.hi:
            raise ValueError('Reversed interval')

    @classmethod
    def point(cls, v):
        return cls(F(str(v)), F(str(v)))


@dataclass(frozen=True)
class CalibrationBox:
    capture: Interval = Interval(F(9,10), F(11,10))
    detector: Interval = Interval(F(9,10), F(11,10))
    K: Interval = Interval(F(9,10), F(11,10))
    J: Interval = Interval(F(9,10), F(11,10))
    dilution: Interval = Interval(F(9), F(11))
    drift: Interval = Interval(F(19,20), F(21,20))
    rho: Interval = Interval(F(4,5), F(1))
    attenuation: F = F(1)

    def __post_init__(self):
        if min(getattr(self, n).lo for n in ['capture','detector','K','J','dilution','drift','rho']) <= 0:
            raise ValueError('Positive calibration bounds required')
        if self.rho.hi > 1 or not 0 < self.attenuation <= 1:
            raise ValueError('Invalid availability or attenuation')

    @property
    def A(self): return self.capture.hi+self.K.hi
    @property
    def H(self): return self.detector.hi+self.J.hi
    @property
    def product_min(self): return self.capture.lo*self.detector.lo

    def relative_error(self, eta):
        if not 0 <= eta < 1: raise ValueError('Relative envelope must lie in [0,1)')
        return replace(self, drift=Interval(self.drift.lo*(1-eta)/(1+eta), self.drift.hi*(1+eta)/(1-eta)))

    def availability_drift(self, ratio: Interval):
        if ratio.lo <= 0: raise ValueError('Positive availability ratio required')
        return replace(self, dilution=Interval(self.dilution.lo/ratio.hi, self.dilution.hi/ratio.lo))

    def response_bounds(self, lo, hi):
        lower = self.attenuation*self.product_min*lo/((hi+self.A)*(hi+self.H))
        upper = min(hi, self.capture.hi*self.detector.hi*hi/((lo+self.K.lo)*(lo+self.J.lo)))
        return lower, upper


@dataclass(frozen=True)
class MarginCertificate:
    box: CalibrationBox
    low: F
    high: Interval
    margin: F

    def __post_init__(self):
        if self.low < 0 or self.high.lo <= 0 or self.margin < 0:
            raise ValueError('Invalid certificate range or margin')

    def polynomial(self, u, d):
        b = self.box; c = b.product_min
        return c*b.drift.lo*b.attenuation*d*u*u-(c+self.margin*u)*(u+b.A*d)*(u+b.H*d)

    def coefficients(self, d):
        b = self.box; c = b.product_min; m = self.margin
        a = [-c*b.A*b.H*d*d, -c*(b.A+b.H)*d-m*b.A*b.H*d*d,
             c*b.drift.lo*b.attenuation*d-c-m*(b.A+b.H)*d, -m]
        derivative = lambda u: a[1]+2*a[2]*u+3*a[3]*u*u
        l, h = self.high.lo, self.high.hi; width = h-l
        return [self.polynomial(l,d), self.polynomial(l,d)+width*derivative(l)/3,
                self.polynomial(h,d)-width*derivative(h)/3, self.polynomial(h,d)]

    def check(self):
        b = self.box
        coeffs = {str(d): self.coefficients(d) for d in (b.dilution.lo, b.dilution.hi)}
        low_ok = b.drift.hi*(self.low+b.A)*(self.low+b.H) <= b.attenuation*b.product_min*b.dilution.lo
        return dict(low_sign=low_ok, high_margin=all(v >= 0 for row in coeffs.values() for v in row),
                    coefficients=coeffs, meaning='Sufficient exact certificate; a failed coefficient is inconclusive')


@dataclass(frozen=True)
class ObservationBudget:
    neat_error: F = READING_ERROR
    diluted_error: F = READING_ERROR
    gain_min: F = F(1)

    def __post_init__(self):
        if min(self.neat_error, self.diluted_error) < 0 or self.gain_min <= 0:
            raise ValueError('Nonnegative errors and positive gain required')

    def bounds(self, margin):
        e = self.neat_error+self.diluted_error
        return e, self.gain_min*margin-e

    def report(self, z: F, certificate: MarginCertificate, promised=False, calibrated=False):
        check = certificate.check(); a, b = self.bounds(certificate.margin)
        if not calibrated:
            return dict(status='UNRESOLVED', reason='Calibration premises have not been established')
        if not check['low_sign'] or not check['high_margin']:
            return dict(status='UNRESOLVED', reason='Sufficient source certificate failed')
        exclusions = dict(excludes_low=z>a, excludes_high=z<b)
        # Both exclusions may hold in the guard band. A caller's binary promise then contradicts the record.
        if promised and all(exclusions.values()):
            status = 'MODEL_OR_PROMISE_INCOMPATIBLE'
        elif promised and a < b:
            status = 'HIGH_UNDER_PROMISE' if z > a else 'LOW_UNDER_PROMISE'
        else:
            status = 'ONE_SIDED_EXCLUSIONS_ONLY'
        return dict(status=status, **exclusions, contrast_bounds=(a,b))

    @staticmethod
    def gaussian_allowance(bias, sigma, alpha):
        if min(bias,sigma) < 0 or not 0 < alpha < 1: raise ValueError('Invalid precision contract')
        return bias+norm.ppf(1-alpha/4)*sigma  # Two readings, four tails, no independence needed.


@dataclass(frozen=True)
class OuterInference:
    box: CalibrationBox
    budget: ObservationBudget
    certificate: MarginCertificate | None = None
    low_label: F = F(1,10)
    high_label: F = F(25)

    def feasible(self, lo, hi, readings):
        # Independent interval envelopes relax reagent sharing. Intersect the SAME gain for both readings.
        b = self.box
        bounds = [b.response_bounds(lo,hi), b.response_bounds(lo/b.dilution.hi,hi/b.dilution.lo)]
        bounds[1] = (bounds[1][0]*b.drift.lo, bounds[1][1]*b.drift.hi)
        gl, gu = self.budget.gain_min, None
        for obs, (sl,sh), eps in zip(readings,bounds,[self.budget.neat_error,self.budget.diluted_error]):
            if obs is None: continue
            yl, yh = obs.lo-eps, obs.hi+eps
            if yh < 0 or (sh == 0 and yl > 0): return False
            if sh > 0: gl = max(gl,yl/sh)
            if sl > 0: gu = yh/sl if gu is None else min(gu,yh/sl)
            if gu is not None and gl > gu: return False
        cert = self.certificate
        if cert is not None and all(v is not None for v in readings):
            zlo, zhi = readings[1].lo-readings[0].hi, readings[1].hi-readings[0].lo
            a, c = self.budget.bounds(cert.margin)
            if hi <= cert.low and zlo > a: return False
            if lo >= cert.high.lo and hi <= cert.high.hi and zhi < c: return False
        return True

    def enclose(self, readings, ceiling=NATIVE_CEILING, depth=ENCLOSURE_DEPTH,
                calibrated=False, cell_budget=CELL_BUDGET):
        if len(readings) != 2: raise ValueError('Exactly two reading intervals, or None for missing')
        if not isinstance(depth,int) or depth < 0 or cell_budget < 1: raise ValueError('Invalid work budget')
        if ceiling is None or not calibrated:
            return dict(status='UNRESOLVED', intervals=[['0','unbounded']], reason='Finite ceiling and established calibration required')
        if ceiling < 0: raise ValueError('Nonnegative ceiling required')
        if self.certificate is not None:
            if self.certificate.box != self.box or not all(self.certificate.check()[v] for v in ['low_sign','high_margin']):
                raise ValueError('Contrast cuts require a matching verified certificate')
        cells = [(F(0), F(ceiling)*self.box.rho.hi)]; completed = 0
        # If the budget is reached, retain the complete previous covering. Never return a partial search as complete.
        for _ in range(depth):
            if 2*len(cells) > cell_budget: break
            cells = [(a,b) for lo,hi in cells for a,b in [(lo,(lo+hi)/2),((lo+hi)/2,hi)] if self.feasible(a,b,readings)]
            completed += 1
        merged = []
        for lo,hi in cells:
            a,b = lo/self.box.rho.hi, min(F(ceiling),hi/self.box.rho.lo)
            if merged and a <= merged[-1][1]: merged[-1] = (merged[-1][0],max(b,merged[-1][1]))
            else: merged.append((a,b))
        status = ('MODEL_OR_MEASUREMENT_INCOMPATIBLE' if not merged else
                  'LOW_CONDITIONAL' if merged[-1][1] <= self.low_label else
                  'HIGH_CONDITIONAL' if merged[0][0] >= self.high_label else 'UNRESOLVED')
        return dict(status=status,intervals=merged,depth=completed,requested_depth=depth,
                    ceiling=ceiling,meaning='Outer relaxation; retained states need not be feasible')


def paper_certificates():
    unit = CalibrationBox()
    cap = replace(unit,capture=Interval(9,11),detector=Interval(9,11),K=Interval(F(9,100),F(11,100)),J=Interval(F(9,100),F(11,100)))
    unequal = replace(cap,detector=Interval(18,22),J=Interval(F(18,100),F(22,100)))
    specs = [('main',unit,F(2,5),20,100,F(1,25)), ('smaller_guard',unit,F(2,5),12,100,F(1,100)),
             ('range_1000',unit,F(2,5),20,1000,F(1,200)), ('range_10000',unit,F(2,5),20,10000,F(1,2000)),
             ('availability_drift',unit.availability_drift(Interval(F(19,20),F(21,20))),F(1,10),20,100,F(3,100)),
             ('transient',replace(unit,attenuation=F(9,10)),F(1,10),20,100,F(3,100)),
             ('relative_error',unit.relative_error(F(3,100)),F(1,10),20,100,F(3,100)),
             ('high_capacity',cap,F(15),100,1000,F(2,5)), ('capacity_guard',cap,F(15),60,1000,F(3,20)),
             ('unequal_capacity',unequal,F(20),200,1000,F(4,5))]
    return {name:MarginCertificate(box,low,Interval(l,h),m) for name,box,low,l,h,m in specs}


def main():
    parser = argparse.ArgumentParser(); parser.add_argument('--output',default='outputs'); args=parser.parse_args()
    out=Path(args.output); out.mkdir(parents=True,exist_ok=True)
    source=SandwichSource(BindingSite(C,K),BindingSite(D,J)); certificates=paper_certificates()
    checks={name:cert.check() for name,cert in certificates.items()}
    assert all(v['low_sign'] and v['high_margin'] for v in checks.values())
    budget=ObservationBudget(); inference=OuterInference(certificates['main'].box,budget,certificates['main'])
    pt=Interval.point
    scenarios={'blank':[pt(0),pt(0)],'high_50':list(map(pt,source.paired(50,NATIVE_AVAILABILITY,DILUTION))),
               'guard_1':list(map(pt,source.paired(1))), 'collision_low':list(map(pt,source.paired(200/2499))),
               'collision_high':list(map(pt,source.paired(2499/50))), 'neat_only':[pt(F(49,2550)),None],
               'negative':[pt(-1),pt(0)], 'censored_dilution':[pt(0),Interval(F(-1,10000),F(1,10000))]}
    enclosures={name:inference.enclose(v,calibrated=True) for name,v in scenarios.items()}
    enclosures['blank_without_ceiling']=inference.enclose(scenarios['blank'],ceiling=None,calibrated=True)
    observed_status=inference.enclose(scenarios['high_50'],calibrated=CALIBRATION_ESTABLISHED)
    kinetics=BindingKinetics(source,*ON_RATES); times=np.linspace(0,3/kinetics.uniform_rate(),161)
    kinetic_rows=[]; maxgap=0.
    for u in [.05,2,50]:
        full=kinetics.trajectory(u,times); reduced=kinetics.trajectory(u,times,full=False)
        maxgap=max(maxgap,float(np.max(np.abs(full-reduced))))
        kinetic_rows.extend(zip([u]*len(times),times,full[:,3],kinetics.attenuation(u,times)*source.signal(u),[float(source.signal(u))]*len(times)))
    highsource=SandwichSource(BindingSite(10,.1),BindingSite(10,.1))
    sequential=SequentialSource(highsource,CARRY_OVER)
    seqrows=[[u,float(SequentialSource(highsource).signal(u)),float(sequential.signal(u))] for u in [1,10,100,1000,1e6]]
    u=np.geomspace(.01,1e6,601)
    sourcerows=list(zip(u,source.signal(u),source.signal(u/DILUTION),source.signal(u/DILUTION)/source.signal(u)))
    results=dict(inputs=dict(C=C,D=D,K=K,J=J,dilution=DILUTION,rho=NATIVE_AVAILABILITY,ceiling=NATIVE_CEILING,error=READING_ERROR),
                 certificates=checks,conditional_enclosures=enclosures,configured_record=observed_status,
                 peak=float(source.peak()),collision=dict(neat=F(49,2550),accessible_low=F(200,2499),accessible_high=F(2499,50)),
                 kinetic_max_full_reduced_gap=maxgap,near_unit_box_sufficient_seconds=3/(.001*1.8),
                 precision_allowance=budget.gaussian_allowance(.001,.002,.01),
                 wash_limit=SequentialSource.wash_limit(F(10),F(1000),F(9,10)),
                 blank_tail_from=F(121,100)*11/READING_ERROR,
                 uniform_tail_margin_constant=F(7371,8000),
                 masking_worlds=[dict(x=F(1,10),rho=F(1)),dict(x=F(100),rho=F(1,1000))],
                 scope='Fresh exact rational certificates and conservative inference; numerical source/ODE; Lean not rerun')
    def dump(name,data): (out/name).write_text(json.dumps(data,indent=2,default=lambda v:str(v) if isinstance(v,F) else float(v))+'\n',encoding='utf-8')
    def table(name,header,rows):
        with (out/name).open('w',newline='') as f:
            w=csv.writer(f);w.writerow(header);w.writerows(rows)
    dump('results.json',results)
    table('source.csv',['accessible','neat','diluted','ratio'],sourcerows)
    table('kinetics.csv',['accessible','seconds','full_A11','attenuation_lower','equilibrium'],kinetic_rows)
    table('sequential.csv',['accessible_nM','ideal_wash','carry_over'],seqrows)
    table('certificates.csv',['case','dilution','b0','b1','b2','b3'],[[name,d,*v] for name,check in checks.items() for d,v in check['coefficients'].items()])
    lines=[f'{len(checks)} exact continuum certificates passed (80 Bernstein coefficients).',
           'Blank with external ceiling 100: '+str(enclosures['blank']['intervals'])+'; without ceiling: unresolved.',
           f'Full/reduced kinetic discrepancy: {maxgap:.3g}; box-wide sufficient incubation: {results["near_unit_box_sufficient_seconds"] / 60:.3f} minutes.',
           f'Configured record: {observed_status["status"]}. All numerical inputs are synthetic.']
    (out/'console.txt').write_text('\n'.join(lines)+'\n');print('\n'.join(lines))
    plot(out,source,certificates,kinetic_rows,highsource)
    digest=lambda p:hashlib.sha256(p.read_bytes()).hexdigest()
    dump('run_metadata.json',dict(manuscript_sha256=MANUSCRIPT_SHA256,python=platform.python_version(),source_sha256=digest(Path(__file__)),
                               output_sha256={p.name:digest(p) for p in sorted(out.iterdir()) if p.is_file() and p.name!='run_metadata.json'}))


def plot(out,source,certs,kinetic_rows,highsource):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained')
    u=np.geomspace(.01,1000,501)
    axs[0].semilogx(u,source.signal(u),label='Neat');axs[0].semilogx(u,source.signal(u/10),label='Tenfold dilution')
    for x in [200/2499,2499/50]:axs[0].scatter(x,source.signal(x),color='k',s=20,zorder=4)
    axs[0].set(title='Undiluted and diluted assay response curves',xlabel='Accessible concentration (nM)',ylabel='Doubly occupied analyte (nM)');axs[0].legend(fontsize=8)
    for name in ['main','range_1000','range_10000']:
        c=certs[name];axs[1].loglog([25,float(c.high.hi)],[float(c.margin/4)]*2,lw=3,label=f'U={c.high.hi}')
    ceilings=np.geomspace(25,10000,200);axs[1].loglog(ceilings,7371/(32000*ceilings),'k--',label='Uniform sufficient bound')
    axs[1].set(title='Allowable measurement error versus\nconcentration range',xlabel='Declared native ceiling (nM)',ylabel='Strict per-reading error ceiling');axs[1].legend(fontsize=8)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'dilution.png',dpi=180);fig.savefig(out/'dilution.svg');plt.close(fig)
    fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained')
    for u in [.05,2,50]:
        rows=np.array([r for r in kinetic_rows if r[0]==u]);axs[0].plot(rows[:,1]/60,rows[:,2]/rows[:,4],label=f'u={u:g}')
    t=np.linspace(0,max(r[1] for r in kinetic_rows),200)
    rate=BindingKinetics(source,*ON_RATES).uniform_rate();axs[0].plot(t/60,(-np.expm1(-rate*t))**2,'k--',label='Uniform source bound')
    axs[0].axhline(.9,color='gray',ls=':');axs[0].set(title='Approach to equilibrium signal during\nincubation',xlabel='Time (minutes)',ylabel='Signal / equilibrium signal');axs[0].legend(fontsize=8)
    u=np.geomspace(.1,1e6,501)
    for lam in [0,.001,.01]:axs[1].loglog(u,SequentialSource(highsource,lam).signal(u),label=f'Carry-over {lam:g}')
    axs[1].loglog(u,highsource.signal(u),'k--',label='Simultaneous')
    axs[1].set(title='Sequential assay response with reagent\ncarry-over',xlabel='Accessible concentration (nM)',ylabel='Signal source (nM)');axs[1].legend(fontsize=8)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'protocol.png',dpi=180);fig.savefig(out/'protocol.svg');plt.close(fig)


if __name__ == '__main__': main()