"""Capped shifted Zipf rows, conditional sampling, and source-to-operation bounds."""
from fractions import Fraction as F
from math import prod
import numpy as np
import mpmath as mp
from polymer import Environment, FOOD, SELECTED, MARKS


def falling(n,k):return prod(range(n-k+1,n+1)) if k<=n else 0


class CappedZipf:
    def __init__(self,exponent=1.5):
        self.a=str(exponent)
        if not mp.isfinite(mp.mpf(self.a)) or mp.mpf(self.a)<=1:raise ValueError('Zipf exponent must exceed one.')

    def degree_probabilities(self,R):
        if type(R)!=int or not 2<=R<=10000:raise ValueError('Explicit row-law budget is 2..10000 channels.')
        with mp.workdps(60):
            a=mp.mpf(self.a);Z=mp.zeta(a)
            return [mp.mpf(k)**(-a)/Z for k in range(1,R)]+[mp.zeta(a,R)/Z]

    def conditional_degrees(self,R,required=0,forbidden=0):
        if any(type(v)!=int or v<0 for v in [required,forbidden]) or required+forbidden>R:raise ValueError('Disjoint feasible constraint counts required.')
        pmf=self.degree_probabilities(R)
        with mp.workdps(60):
            denominator=falling(R,required+forbidden)
            weights=[p*falling(d,required)*falling(R-d,forbidden)/denominator for d,p in enumerate(pmf)]
            evidence=sum(weights)
            if evidence<=0:raise ValueError('Conditional row event has zero probability.')
            return [p/evidence for p in weights],evidence

    def sample_row(self,R,rng,required=(),forbidden=()):
        required=tuple(required);forbidden=tuple(forbidden)
        if len(set(required))!=len(required) or len(set(forbidden))!=len(forbidden) or set(required)&set(forbidden) or any(type(r)!=int or not 0<=r<R for r in required+forbidden):
            raise ValueError('Required/forbidden coordinates must be distinct valid split indices.')
        probabilities,_=self.conditional_degrees(R,len(required),len(forbidden))
        p=np.array(list(map(float,probabilities)));p/=p.sum();d=int(rng.choice(R,p=p))
        options=np.array([r for r in range(R) if r not in set(required+forbidden)],int)
        return tuple(sorted(required+tuple(map(int,rng.choice(options,d-len(required),replace=False)))))

    def sample_environment(self,catalogue,seed=5901,food_silent_selected=False):
        rng=np.random.default_rng(seed);R=len(catalogue.splits);selected=catalogue.splits.index(SELECTED)
        # Marks remain independent of row selection, with S shared between the
        # basal coefficient and all catalytic coefficients of one split.
        S=tuple(map(float,rng.choice(MARKS,R)));B=tuple(map(float,rng.choice(MARKS,R)));incidences=[]
        for z in catalogue.words:
            if food_silent_selected and z in FOOD:continue
            row=self.sample_row(R,rng,required=(selected,) if food_silent_selected and z=='0011' else ())
            incidences.extend((z,r,float(rng.choice(MARKS))) for r in row)
        environment=Environment(S,B,tuple(incidences));environment.validate(catalogue)
        return environment

    def moments(self,n):
        if type(n)!=int or not 4<=n<=2000:raise ValueError('Analytic moment budget: n=4..2000.')
        R=(n-2)*2**(n+1)+4;X=2**(n+1)-2
        # Enough precision for subtraction of growing partial sums. These are
        # evaluations of exact source formulas, not directed probability bounds.
        with mp.workdps(70):
            a=mp.mpf(self.a);Z=mp.zeta(a)
            def hurwitz(s):
                if R<=100000:return mp.zeta(s,R)
                # Euler-Maclaurin evaluation for a huge lower endpoint. It is a
                # numerical approximation, checked against direct Hurwitz values
                # at overlapping moderate endpoints; no directed claim is made.
                N=mp.mpf(R);value=N**(1-s)/(s-1)+N**(-s)/2;rising=s
                for j in range(1,17):
                    value+=mp.bernoulli(2*j)/mp.factorial(2*j)*rising*N**(-s-2*j+1)
                    rising*=(s+2*j-1)*(s+2*j)
                return value
            def partial(s):return mp.harmonic(R-1) if s==1 else mp.zeta(s)-hurwitz(s)
            tail=hurwitz(a)
            ED=(partial(a-1)+R*tail)/Z-1
            ED2=(partial(a-2)-3*partial(a-1)+2*partial(a)+(R-1)*(R-2)*tail)/Z
            p=ED/R;q2=ED2/(R*(R-1));e=1/Z
            return dict(n=n,channels=R,molecules=X,exponent=a,mean_degree=ED,second_factorial_moment=ED2,
                        incidence=p,same_row_pair=q2,pair_ratio=q2/p,empty_row=e,source_witness=e**6*p,scaled_incidence=2**n*p,
                        cap_atom=tail/Z,conditional_cap_atom=(R-1)*tail/(Z*ED))

    def small_cap_residual(self,n):
        if type(n)!=int or not 4<=n<=100:raise ValueError('Numerical finite-converse evaluation budget: n=4..100.')
        with mp.workdps(70):
            a=mp.mpf(self.a);X=2**(n+1)-2;e=1/mp.zeta(a);p1=2**(-a)*e
            log_complement=X*mp.log(e)+mp.log1p(X*p1/e)
            return dict(probability_at_least_two=-mp.expm1(log_complement),log10_complement=log_complement/mp.log(10),
                        scope='n<=10^12: rectangle contains the entire catalogue; numerical exact-formula evaluation')


def source_operation_bounds(moments,failure_upper,residual_upper):
    """Numeric integration of supplied theorem bounds, with hypotheses external.

    A valid pointwise witness bound is required. The census coefficient is 224;
    supplied residual must cover the original converse at its scale conditions.
    """
    with mp.workdps(60):
        delta,r=mp.mpf(str(failure_upper)),mp.mpf(str(residual_upper))
        if not 0<=delta<=1 or not 0<=r<=1:raise ValueError('Probability bounds must be in [0,1].')
        mass=moments['source_witness'];lower=mass*(1-delta);upper=min(mp.mpf(1),224*moments['incidence']+r)
        return dict(success_lower=lower,success_upper=upper,
                    witness_posterior_lower=lower/upper if upper else None,
                    no_productive_posterior_upper=min(mp.mpf(1),r/lower) if lower else None,
                    evidence='Numeric source integration of supplied valid trajectory/residual bounds; not an empirical success estimate')


def screening_trials(success_lower,confidence=.95):
    with mp.workdps(60):
        s,alpha=mp.mpf(str(success_lower)),mp.mpf(str(confidence))
        if not 0<s<=1 or not 0<alpha<1:raise ValueError('Positive success lower bound and interior confidence required.')
        return 1 if s==1 else int(mp.ceil(mp.log1p(-alpha)/mp.log1p(-s)))