"""Finite-resource amplification: exact small witnesses and numerical design tools.

Run python example.py --output outputs. Inputs are illustrative paper values,
not fitted clinical parameters. Probabilities are per reaction.
"""
from __future__ import annotations
import argparse
import csv
from dataclasses import dataclass, asdict, replace
from fractions import Fraction as F
import hashlib
import json
import math
from pathlib import Path
import platform
import time

import numpy as np
import scipy
from scipy.linalg import expm
from scipy.optimize import brentq
from scipy.integrate import quad
from scipy.special import gammainc, gammaincc, gammaln, digamma
from scipy.stats import poisson

# USER INPUTS: paper preset --------------------------------------------------
THRESHOLD = 5                  # positive integer active units at detection
CAPACITIES = (5, 6, 7, 8, 9, 10)  # integer resource inventory >= threshold
BACKGROUND = '0.01'            # units/minute, strictly positive; exact decimal
GROWTH = '1'                   # inverse minutes, strictly positive
LOADING_MEAN = '4'             # Poisson mean activated units, nonnegative
BLANK_LIMIT = '0.01'           # allowed blank-positive probability, in (0,1)
MISS_LIMIT = '0.05'            # allowed loaded-negative probability, in (0,1)
GRID_STEP = 0.1                # minutes, positive; observation grid anchored at 0
CLOCK_MULTIPLIER = 2.0         # positive; scales every rate equally
HEADROOMS = (0, 1, 2, 3, 5, 10, 20)  # nonnegative units remaining at detection
LARGE_THRESHOLD = 100_000_000  # theorem evaluator requires h >= 10**8
# Numerical controls, separate from physical inputs.
ROOT_TOLERANCE = 1e-11
QUADRATURE_TOLERANCE = 2e-11
POISSON_TAIL_TOLERANCE = 1e-13
MAX_DENSE_THRESHOLD = 200      # avoid accidental huge matrix allocation
PAPER_SHA256 = 'f27612271a1fbd5bf5c8c9381d22b0974f07ece63d0f0fea8dcd520a7f17ee4c'
# --------------------------------------------------------------------------


@dataclass(frozen=True)
class BirthSource:
    threshold: int = THRESHOLD
    capacity: int = 10
    background: F = F(BACKGROUND)
    growth: F = F(GROWTH)

    def __post_init__(self):
        for name in ('threshold','capacity'):
            value=getattr(self,name)
            if not isinstance(value,int) or isinstance(value,bool) or value < 1:
                raise ValueError(f'{name} must be a positive integer.')
        if self.capacity < self.threshold:
            raise ValueError('Capacity must be at least the detection threshold.')
        for name in ('background','growth'):
            value=F(str(getattr(self,name)))
            if value <= 0: raise ValueError(f'{name} must be positive.')
            object.__setattr__(self,name,value)

    def rates(self):
        if self.threshold > MAX_DENSE_THRESHOLD:
            raise ValueError('Finite-state budget exceeded; use LimitLaw or the large-threshold theorem evaluator.')
        return tuple((self.background+self.growth*z)*(1-F(z,self.capacity)) for z in range(self.threshold))


@dataclass(frozen=True)
class PoissonLoading:
    mean: F = F(LOADING_MEAN)

    def __post_init__(self):
        value=F(str(self.mean))
        if value < 0: raise ValueError('Loading mean must be nonnegative.')
        object.__setattr__(self,'mean',value)

    def transient_weights(self, threshold):
        # Do not renormalize: loads at or above threshold are positive at t=0.
        return poisson.pmf(np.arange(threshold),float(self.mean))


class FiniteChain:
    """Row generator Q; exp(t Q) maps a terminal indicator backward to each start."""
    def __init__(self, source: BirthSource, loading=PoissonLoading(), rate_factors=None):
        self.source,self.loading=source,loading
        rates=np.array(list(map(float,source.rates())))
        if rate_factors is not None:
            factors=np.asarray(rate_factors,float)
            if factors.shape!=rates.shape or not np.all(np.isfinite(factors)) or np.any(factors<=0):
                raise ValueError('One finite positive factor per transient-state rate is required.')
            rates *= factors
        h=source.threshold
        self.generator=np.zeros((h+1,h+1))
        self.generator[np.arange(h),np.arange(h)]=-rates
        self.generator[np.arange(h),np.arange(1,h+1)]=rates
        self.uncalled=np.r_[np.ones(h),0.]
        self.weights=loading.transient_weights(h)

    def survival(self,t):
        if not np.isfinite(t) or t < 0: raise ValueError('Deadline must be finite and nonnegative.')
        return expm(t*self.generator)@self.uncalled

    def errors(self,t):
        s=self.survival(t)
        return float(1-s[0]),float(self.weights@s[:-1])


@dataclass(frozen=True)
class ErrorLimits:
    blank: float = float(BLANK_LIMIT)
    miss: float = float(MISS_LIMIT)

    def __post_init__(self):
        if not 0 < self.blank < 1 or not 0 < self.miss < 1:
            raise ValueError('Error limits must lie strictly between zero and one.')

    def window(self, chain: FiniteChain, grid_step=GRID_STEP):
        if not np.isfinite(grid_step) or grid_step<=0: raise ValueError('Grid step must be positive.')
        end=1.
        while chain.errors(end)[0] < self.blank:
            end*=2
            if end>1e12: raise ArithmeticError('Cannot bracket the blank boundary within time budget.')
        latest=brentq(lambda t:chain.errors(t)[0]-self.blank,0,end,xtol=ROOT_TOLERANCE)
        if chain.errors(0)[1] <= self.miss:
            earliest=0.
        else:
            end=max(end,1.)
            while chain.errors(end)[1] > self.miss:
                end*=2
                if end>1e12: raise ArithmeticError('Cannot bracket the sensitivity boundary within time budget.')
            earliest=brentq(lambda t:chain.errors(t)[1]-self.miss,0,end,xtol=ROOT_TOLERANCE)
        # Root locations are numerical. A near-tie is explicitly unresolved.
        margin=latest-earliest
        status='unresolved' if abs(margin)<10*ROOT_TOLERANCE else ('feasible' if margin>0 else 'empty')
        frames=[]
        if status=='feasible':
            for k in range(math.ceil(earliest/grid_step),math.floor(latest/grid_step)+1):
                t=k*grid_step
                blank,miss=chain.errors(t)
                if blank<=self.blank and miss<=self.miss: frames.append(t)
        return dict(capacity=chain.source.capacity,earliest=earliest,latest=latest,
                    width=max(0.,margin),status=status,grid_frames=frames,
                    blank_boundary_miss=chain.errors(latest)[1])


@dataclass(frozen=True)
class Interval:
    """Exact rational endpoints, with sign-aware arithmetic."""
    lo: F
    hi: F

    def __post_init__(self):
        if self.lo>self.hi: raise ValueError('Reversed interval.')

    def scale(self,c):
        return Interval(min(c*self.lo,c*self.hi),max(c*self.lo,c*self.hi))

    def __add__(self,other): return Interval(self.lo+other.lo,self.hi+other.hi)

    def __mul__(self,other):
        products=[a*b for a in (self.lo,self.hi) for b in (other.lo,other.hi)]
        return Interval(min(products),max(products))

    def record(self):
        return dict(lower_fraction=str(self.lo),upper_fraction=str(self.hi),
                    approximate_lower=float(self.lo),approximate_upper=float(self.hi),
                    width=float(self.hi-self.lo))


def exp_enclosure(x: F, degree=24):
    """Taylor enclosure with exact rational remainder, scaling and reciprocation.

    For 0<=y<=1, all positive omitted terms have ratio <= y/(degree+2).
    Sum the resulting geometric majorant; no floating-point exponent is used.
    """
    x=F(x); y=abs(x); squares=0
    while y>1: y/=2; squares+=1
    term=F(1); total=term
    for k in range(1,degree+1): term*=y/k; total+=term
    tail=(term*y/(degree+1))/(1-y/(degree+2))
    lo,hi=total,total+tail
    for _ in range(squares): lo*=lo; hi*=hi
    if x<0: lo,hi=1/hi,1/lo
    # Exact outward rounding keeps certificates small and inspectable.
    scale=10**40
    return Interval(F((lo.numerator*scale)//lo.denominator,scale),
                    F(-((-hi.numerator*scale)//hi.denominator),scale))


class ExactWitness:
    """Rational spectral decomposition; intended for small distinct-rate sources."""
    def __init__(self,source: BirthSource,loading=PoissonLoading()):
        self.source,self.loading=source,loading
        self.q=source.rates(); h=source.threshold
        if h>20: raise ValueError('Exact spectral certificate budget is threshold <=20.')
        if len(set(self.q))!=h:
            raise ValueError('Spectral witness requires distinct rates; FiniteChain also handles repeated rates.')
        self.v=[[F(0) for _ in range(h)] for _ in range(h)]
        for i in reversed(range(h)):
            self.v[i][i]=1-sum(self.v[j][i] for j in range(i+1,h))
            for z in reversed(range(i)):
                self.v[i][z]=self.q[z]/(self.q[z]-self.q[i])*self.v[i][z+1]
        # Independent algebraic verification of decomposition and Q v = -q_i v.
        for z in range(h):
            if sum(v[z] for v in self.v)!=1: raise ArithmeticError('Decomposition failed.')
        for i,v in enumerate(self.v):
            for z in range(h):
                if self.q[z]*((v[z+1] if z+1<h else 0)-v[z]) != -self.q[i]*v[z]:
                    raise ArithmeticError('Eigenvector identity failed.')

    def errors(self,t):
        t=F(str(t))
        if t<0: raise ValueError('Deadline must be nonnegative.')
        h=self.source.threshold
        exps=[exp_enclosure(-q*t) for q in self.q]
        s=[sum_intervals([e.scale(self.v[i][z]) for i,e in enumerate(exps)]) for z in range(h)]
        blank=Interval(1-s[0].hi,1-s[0].lo)
        polynomial=sum_intervals([s[n].scale(self.loading.mean**n/math.factorial(n)) for n in range(h)])
        return blank,exp_enclosure(-self.loading.mean)*polynomial


def sum_intervals(values):
    out=Interval(F(0),F(0))
    for value in values: out=out+value
    return out


class LimitLaw:
    """Large-h law, separate from finite-h probabilities; a = b/g."""
    def __init__(self,a=float(F(BACKGROUND)/F(GROWTH)),loading=float(LOADING_MEAN),headroom=None):
        if not np.isfinite(a) or a<=0 or not np.isfinite(loading) or loading<0:
            raise ValueError('Positive a and nonnegative finite loading required.')
        if headroom is not None and (not isinstance(headroom,int) or headroom<0):
            raise ValueError('Headroom must be a nonnegative integer, or None for growing reserve.')
        self.a,self.loading,self.headroom=a,loading,headroom
        self.last=int(poisson.ppf(1-POISSON_TAIL_TOLERANCE,loading))
        self.weights=poisson.pmf(np.arange(self.last+1),loading)
        self.tail=float(poisson.sf(self.last,loading))
        self.max_quad_error=0.

    def errors(self,y):
        if not np.isfinite(y) or y<=0: raise ValueError('Positive finite gamma-product cutoff required.')
        shapes=self.a+np.arange(self.last+1)
        if self.headroom is None:
            return float(gammaincc(self.a,y)),float(self.weights@gammainc(shapes,y))
        shape=self.headroom+1
        def density(v):
            return math.exp((shape-1)*math.log(v)-v-gammaln(shape)) if v>0 else 0.
        def blank_integrand(v): return density(v)*gammaincc(self.a,y/v) if v>0 else 0.
        def miss_integrand(v): return density(v)*float(self.weights@gammainc(shapes,y/v)) if v>0 else 0.
        blank,e1=quad(blank_integrand,0,np.inf,epsabs=QUADRATURE_TOLERANCE,epsrel=QUADRATURE_TOLERANCE,limit=200)
        miss,e2=quad(miss_integrand,0,np.inf,epsabs=QUADRATURE_TOLERANCE,epsrel=QUADRATURE_TOLERANCE,limit=200)
        self.max_quad_error=max(self.max_quad_error,e1,e2)
        return blank,miss

    def boundary(self,limits=ErrorLimits()):
        lo,hi=1.,1.
        while self.errors(lo)[0]<limits.blank: lo/=2
        while self.errors(hi)[0]>limits.blank: hi*=2
        y=brentq(lambda x:self.errors(x)[0]-limits.blank,lo,hi,xtol=ROOT_TOLERANCE)
        blank,miss=self.errors(y)
        margin=limits.miss-miss
        return dict(headroom='growing reserve' if self.headroom is None else str(self.headroom),
                    cutoff_y=y,blank=blank,miss=miss,poisson_tail_bound=self.tail,
                    quadrature_error_estimate=self.max_quad_error,
                    status='unresolved' if abs(margin)<10*(self.tail+self.max_quad_error+ROOT_TOLERANCE)
                    else ('feasible limit' if margin>0 else 'empty limit'))


def large_threshold_bound(h=LARGE_THRESHOLD):
    """Evaluate the manuscript's proven allowances; do not re-certify its scalar bounds.

    Four scalar constants below are imported from the paper's Appendix C.
    Rational bookkeeping is exact; deadlines via digamma are floating point.
    """
    if not isinstance(h,int) or h<10**8: raise ValueError('This theorem requires integer h >= 10**8.')
    a=.01; k=h//2; ell=h-k
    # Harmonic/digamma identities evaluate the paper's sums without h states.
    harmonic=digamma(2*h+1)-digamma(h+1)
    mu_plus=(2*h*harmonic-a*(digamma(h+a)-digamma(a)))/(2*h+a)
    mu_minus=digamma(h+1)-digamma(h-k+1)+digamma(h+a)-digamma(k+a)
    c=F(10000)+F(20000,199)**2
    dp=c/h; dm=F(401,100)*c/h
    bounds=dict(doubled_blank_upper=F('0.009372')+dp,
                doubled_miss_upper=F('0.048731')+dp+F(3,10**13),
                saturated_blank_lower=F('0.01087')-dm,
                saturated_miss_lower=F('0.056383')-dm)
    return dict(threshold=h,assumptions='b/g=.01, loading=4, blank limit=.01, miss limit=.05; paper scalar bounds imported',
                doubled_deadline=(math.log(h+a)+mu_plus-math.log(.3))/float(GROWTH),
                saturated_separating_deadline=h/(h+a)*(math.log(k+a)+math.log(ell+1)+mu_minus-math.log(.15))/float(GROWTH),
                bounds={key:str(value) for key,value in bounds.items()},
                approximate_bounds={key:float(value) for key,value in bounds.items()})


def witnesses():
    rows=[]
    for capacity,t in [(5,'5'),(6,'4'),(7,'3.67'),(8,'3.45'),(8,'3.4'),(8,'3.5'),(9,'3.3'),(10,'3.15'),(10,'3.3')]:
        source=BirthSource(5,capacity,F('0.01'),F(1))
        blank,miss=ExactWitness(source,PoissonLoading(F(4))).errors(t)
        fb,fm=FiniteChain(source,PoissonLoading(F(4))).errors(float(t))
        if max(abs(fb-float(blank.lo)),abs(fm-float(miss.lo)))>1e-11:
            raise ArithmeticError('Matrix exponential disagrees with rational enclosure.')
        rows.append(dict(capacity=capacity,time=t,blank=blank.record(),miss=miss.record(),
                         feasible=blank.hi<=F(1,100) and miss.hi<=F(1,20),
                         separating=blank.lo>F(1,100) and miss.lo>F(1,20)))
    source=BirthSource(5,10,F('0.01'),F(1))
    robust_blank=ExactWitness(source).errors('3.2742')[0]
    robust_miss=ExactWitness(source,PoissonLoading(F('3.99'))).errors('3.1262')[1]
    if not robust_blank.hi<F(1,100) or not robust_miss.hi<F(1,20):
        raise ArithmeticError('Robust endpoint inequalities failed.')
    return dict(preset='Fixed paper witnesses, independent of exploration inputs',rows=rows,
                robust_blank=robust_blank.record(),robust_miss=robust_miss.record())


def write_csv(path,rows):
    with path.open('w',newline='',encoding='utf-8') as stream:
        writer=csv.DictWriter(stream,fieldnames=list(rows[0])); writer.writeheader(); writer.writerows(rows)


def plot(curves,windows,limits,output):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    colors=['#0072B2','#D55E00','#007F5F','#7B3294','#3B3B3B','#A66B00']
    with plt.rc_context({'font.size':11,'axes.spines.top':False,'axes.spines.right':False,
                         'figure.facecolor':'white','savefig.facecolor':'white','svg.fonttype':'none'}):
        fig,ax=plt.subplots(figsize=(7.2,4.8),layout='constrained')
        ax.fill_between([0,100*float(BLANK_LIMIT)],0,100*float(MISS_LIMIT),facecolor='#eeeeee')
        for index,capacity in enumerate(CAPACITIES):
            rows=[r for r in curves if r['capacity']==capacity]
            ax.plot([100*r['blank'] for r in rows],[100*r['miss'] for r in rows],color=colors[index%len(colors)],linestyle=['-','--','-.'][index%3],label=f'Capacity {capacity}')
        ax.axvline(100*float(BLANK_LIMIT),color='#444444',linestyle=':')
        ax.axhline(100*float(MISS_LIMIT),color='#444444',linestyle=':')
        ax.set(xlabel='Blank-positive probability (%)',ylabel='Loaded-negative probability (%)',
               xlim=(0,2),ylim=(0,15))
        ax.legend(fontsize=9,ncol=2); ax.grid(color='#e3e6e8',linewidth=.5)
        fig.savefig(output/'tradeoff.png',dpi=220);fig.savefig(output/'tradeoff.svg');plt.close(fig)
        fig,ax=plt.subplots(figsize=(7.2,4.5),layout='constrained')
        feasible=[r for r in windows if r['status']=='feasible']
        left=min(r['earliest'] for r in (feasible or windows))-.1
        right=max(r['latest'] for r in (feasible or windows))+.1
        ax.set_xlim(left,right)
        for index,row in enumerate(windows):
            if row['status']=='feasible':
                ax.plot([row['earliest'],row['latest']],[row['capacity']]*2,linewidth=7,color=colors[index%len(colors)])
                ax.plot(row['grid_frames'],[row['capacity']]*len(row['grid_frames']),'o',color='#202124',markersize=6)
            else:
                ax.text(left+.02,row['capacity'],'No continuous window' if row['status']=='empty' else 'Unresolved',fontsize=9,va='center')
        ax.set(xlabel='Deadline (minutes)',ylabel='Resource capacity',yticks=list(CAPACITIES))
        ax.grid(axis='x',color='#e3e6e8');ax.set_ylim(min(CAPACITIES)-.6,max(CAPACITIES)+.6)
        fig.savefig(output/'windows.png',dpi=220);fig.savefig(output/'windows.svg');plt.close(fig)
        fig,ax=plt.subplots(figsize=(7.2,4.5),layout='constrained')
        finite=[r for r in limits if r['headroom']!='growing reserve']
        ax.plot([int(r['headroom']) for r in finite],[100*r['miss'] for r in finite],'o-',color='#0072B2',label='Fixed headroom limit')
        ax.axhline(100*limits[-1]['miss'],color='#D55E00',linestyle='--',label='Growing reserve limit')
        ax.axhline(100*float(MISS_LIMIT),color='#444444',linestyle=':',label='Miss limit')
        ax.set(xlabel='Unused units at threshold (fixed headroom m)',ylabel='Limiting miss at blank boundary (%)',ylim=(0,None))
        ax.legend(fontsize=9);ax.grid(axis='y',color='#e3e6e8')
        fig.savefig(output/'headroom.png',dpi=220);fig.savefig(output/'headroom.svg');plt.close(fig)


def main():
    parser=argparse.ArgumentParser(description=__doc__);parser.add_argument('--output',type=Path,default=Path('outputs'))
    args=parser.parse_args();args.output.mkdir(parents=True,exist_ok=True);start=time.perf_counter()
    inputs=dict(threshold=THRESHOLD,capacities=CAPACITIES,background=BACKGROUND,growth=GROWTH,
                loading=LOADING_MEAN,blank_limit=BLANK_LIMIT,miss_limit=MISS_LIMIT,grid_step=GRID_STEP,
                clock_multiplier=CLOCK_MULTIPLIER,headrooms=HEADROOMS,quadrature_tolerance=QUADRATURE_TOLERANCE,
                root_tolerance=ROOT_TOLERANCE,poisson_tail_tolerance=POISSON_TAIL_TOLERANCE)
    intro='Resolved inputs: '+json.dumps(inputs);print(intro,flush=True)
    windows=[];curves=[];clock_error=0.
    for capacity in CAPACITIES:
        source=BirthSource(capacity=capacity);chain=FiniteChain(source)
        window=ErrorLimits().window(chain);windows.append(window)
        scaled=FiniteChain(source,rate_factors=[CLOCK_MULTIPLIER]*THRESHOLD)
        for t in np.linspace(0,max(window['earliest'],window['latest'])*1.2,160):
            blank,miss=chain.errors(float(t));curves.append(dict(capacity=capacity,time=float(t),blank=blank,miss=miss))
        for t in (0.,1.,3.):
            clock_error=max(clock_error,max(abs(a-b) for a,b in zip(scaled.errors(t),chain.errors(CLOCK_MULTIPLIER*t))))
    limits=[LimitLaw(headroom=m).boundary() for m in (*HEADROOMS,None)]
    certs=witnesses()
    paper_parameters=(F(BACKGROUND)/F(GROWTH)==F(1,100) and F(LOADING_MEAN)==4
                      and F(BLANK_LIMIT)==F(1,100) and F(MISS_LIMIT)==F(1,20))
    summary=dict(inputs=inputs,windows=windows,limit_laws=limits,clock_identity_error=clock_error,
                 initially_positive=float(poisson.sf(THRESHOLD-1,float(LOADING_MEAN))),
                 large_threshold=large_threshold_bound() if paper_parameters else 'Paper theorem parameters not satisfied.',
                 evidence='Rational small-threshold enclosures; numerical finite windows and limit quadratures; imported large-threshold scalar theorem constants. No Lean compilation.')
    write_csv(args.output/'tradeoff.csv',curves);write_csv(args.output/'windows.csv',windows)
    write_csv(args.output/'limits.csv',limits)
    (args.output/'exact_witnesses.json').write_text(json.dumps(certs,indent=2)+'\n',encoding='utf-8')
    plot(curves,windows,limits,args.output)
    transcript=json.dumps(summary,indent=2,allow_nan=False)
    (args.output/'summary.json').write_text(transcript+'\n',encoding='utf-8')
    (args.output/'console.txt').write_text(intro+'\n'+transcript+'\n',encoding='utf-8')
    metadata=dict(paper_sha256=PAPER_SHA256,source_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
                  python=platform.python_version(),numpy=np.__version__,scipy=scipy.__version__,platform=platform.platform(),
                  processor=platform.processor(),elapsed_seconds=time.perf_counter()-start,
                  command='python example.py --output outputs',seed_policy='No random sampling.',
                  output_sha256={p.name:hashlib.sha256(p.read_bytes()).hexdigest() for p in sorted(args.output.iterdir())
                                 if p.is_file() and p.name!='run_metadata.json'})
    (args.output/'run_metadata.json').write_text(json.dumps(metadata,indent=2)+'\n',encoding='utf-8')
    print(transcript)


if __name__=='__main__': main()