"""Mass-balanced equilibrium for a peptide connecting two protein partners.

All concentrations and dissociation constants use nM. Alpha is dimensionless.
The six species are A, B, L, AL, BL, and ABL. This model assumes negligible
unliganded A:B association, monovalent binding, and solution equilibrium.
"""
from __future__ import annotations
import numpy as np
from scipy.optimize import least_squares, minimize_scalar


def equilibrium(a_total: float, b_total: float, l_total: float,
                kd_a: float, kd_b: float, alpha: float) -> dict[str, float]:
    values = np.array([a_total,b_total,l_total,kd_a,kd_b,alpha],float)
    if not np.all(np.isfinite(values)) or np.any(values[:3] < 0) or np.any(values[3:] <= 0):
        raise ValueError("invalid concentrations or equilibrium constants")
    totals = values[:3]
    if np.any(totals == 0):
        # Solve the remaining binary complex analytically without subtraction loss.
        a,b,l = totals
        al=bl=0.
        if a>0 and l>0:
            s=a+l+kd_a;al=2*a*l/(s+np.sqrt(s*s-4*a*l));a-=al;l-=al
        elif b>0 and l>0:
            s=b+l+kd_b;bl=2*b*l/(s+np.sqrt(s*s-4*b*l));b-=bl;l-=bl
        return dict(A=float(a),B=float(b),L=float(l),AL=float(al),BL=float(bl),ABL=0.)
    def species(logs):
        a,b,l=np.exp(logs)
        return np.array([a,b,l,a*l/kd_a,b*l/kd_b,alpha*a*b*l/(kd_a*kd_b)])
    def residual(logs):
        a,b,l,al,bl,abl=species(logs)
        return (np.array([a+al+abl,b+bl+abl,l+al+bl+abl])-totals)/totals
    fit=least_squares(residual,np.log(totals)-1.,bounds=(np.log(totals)-80,np.log(totals)),
                      xtol=1e-12,gtol=1e-12,ftol=1e-12,max_nfev=1000)
    if np.max(np.abs(residual(fit.x)))>1e-7:
        raise RuntimeError("equilibrium solver failed mass conservation")
    return dict(zip(["A","B","L","AL","BL","ABL"],map(float,species(fit.x))))


def fit_cooperativity(totals, abl_nm, kd_a, kd_b, sd_nm=None):
    """Fit alpha with binary Kd values fixed from separate binding measurements.

    totals has shape (N,3). abl_nm and sd_nm have shape (N,). This fitter expects
    calibrated ternary concentrations. Raw BRET needs its own observation model.
    """
    t=np.asarray(totals,float);y=np.asarray(abl_nm,float)
    s=np.ones_like(y) if sd_nm is None else np.asarray(sd_nm,float)
    if t.ndim!=2 or t.shape[1]!=3 or y.shape!=(len(t),) or s.shape!=y.shape or np.any(s<=0):
        raise ValueError("invalid concentration-grid dimensions")
    if len(y)==0 or not np.all(np.isfinite(t)) or not np.all(np.isfinite(y)) or not np.all(np.isfinite(s)) or np.any(y<0):
        raise ValueError("invalid ternary data")
    def objective(log_alpha):
        pred=np.array([equilibrium(*row,kd_a,kd_b,np.exp(log_alpha))["ABL"] for row in t])
        return float(np.sum(((pred-y)/s)**2))
    result=minimize_scalar(objective,bounds=(-12.,12.),method="bounded",options={"xatol":1e-8})
    return {"alpha":float(np.exp(result.x)),"weighted_sse":float(result.fun),
            "boundary":bool(abs(result.x)>11.9),"n":len(y)}