"""Unified 'master' CV simulator parameterized by theta_super.

One Backward-Euler reaction-diffusion solver whose limits reproduce the
diffusional-family corner mechanisms (Nernst, BV, MHC, EC, EC', CE, EE, ECE).
Reuses the existing grid / diffusion / flux helpers from generate_dataset_diffec
so the numerical scheme matches the established datasets exactly.

Stage M1a: ET1 only (Nernst / BV / MHC). Chemical steps (M1b) and a real
reversible 2nd electron transfer with tracked species C (M1c) are added next.

ET1 rate law (theta_super slots log10_K0_1, log10_reorg_e_1, asym_1):
  - reversible (K0 large)        -> Nernst boundary condition
  - log10_reorg_e_1 >= BV cutoff -> Butler-Volmer, alpha = asym_1
  - else                         -> symmetric Marcus-Hush (existing MHC integral)
The smooth asymmetric-Marcus unification of BV<->MHC is a focused follow-up;
this interface (et1_rates) is where it will plug in.
"""
from __future__ import annotations

import numpy as np
import scipy.linalg
from scipy.sparse import csr_matrix
from scipy.sparse.linalg import spsolve

from generate_dataset_diffec import (
    gen_grid, ini_conc, ini_coeff, calc_abc_linear, calc_flux, calc_mhc_rates,
    DELTA_X, DELTA_THETA, EXPANDING_GRID_FACTOR, SIMULATION_SPACE_MULTIPLE,
)
from generate_extended_mechanisms import (
    _make_potential_waveform, _clamp_bv_rates,
)
from manifold.theta_super import (
    SLOT_IDX, REVERSIBLE_K0, BV_REORG_SENTINEL, FLOOR_LOGRATE,
)

_ON = FLOOR_LOGRATE + 0.5  # a log-rate above this is an "active" step

# Unified continuous ET law: smooth interpolation in log10(reorg) between the
# symmetric Marcus-Hush-Chidsey integral (finite lambda) and the Butler-Volmer
# law (lambda -> large). This removes the old hard reorg-sentinel branch, so the
# Nernst <-> BV <-> MHC electron-transfer sub-type becomes ONE continuous axis
# (the manifold ET coordinate). The crossover sits at high lambda (~10^2.3), where
# symmetric MHC already approaches BV(alpha=0.5), so the two physical limits are
# reproduced and the transition between them is gentle. BV carries the transfer
# coefficient (asymmetry) alpha; symmetric MHC corresponds to alpha=0.5.
LAMBDA_MID = 2.0      # log10(reorg) midpoint of the BV<->MHC crossover
LAMBDA_WIDTH = 0.12   # log10(reorg) crossover width
_W_EPS = 5e-3         # fast-path cutoffs (skip the unused limb; keeps the BV
                      # sentinel and the MHC sampling range on their exact limbs,
                      # with the blend confined to the high-lambda BV~MHC region)


def _bv_weight(log10_reorg):
    """Smooth BV weight: ->1 at large reorg (BV), ->0 at small reorg (MHC)."""
    return 1.0 / (1.0 + np.exp(-(log10_reorg - LAMBDA_MID) / LAMBDA_WIDTH))


def _unpack(theta_super):
    v = np.asarray(theta_super, dtype=np.float64).reshape(-1)
    return {nm: v[i] for nm, i in SLOT_IDX.items()}


# The legacy calc_mhc_rates (degree-50 Gauss-Hermite) is non-smooth in reorg
# (~9% step noise) and inflates/overflows at large reorg, which showed up as a
# pseudo-seam in the blended ET axis. We use the SAME MHC integrand at a higher,
# exp-clipped quadrature degree (smooth + stable) and cap the reorg fed to it to
# its reliable range (beyond ~this, symmetric MHC ~ BV(alpha=0.5) and the BV limb
# dominates the blend, so the cap is physically invisible).
_REORG_MHC_MAX = 50.0
_MHC_DEG = 128
_MHC_PTS, _MHC_WTS = np.polynomial.hermite.hermgauss(_MHC_DEG)


def _mhc_rates_smooth(Theta, K0, reorg):
    """Symmetric MHC (K_red, K_ox) via the legacy integrand at high, clipped
    Gauss-Hermite degree. Matches calc_mhc_rates but smooth and overflow-safe."""
    s = np.sqrt(reorg)
    base = reorg * (_MHC_PTS * 2.0 / s - 1.0)         # = 2*sqrt(reorg)*x - reorg

    def I_red(theta):
        arg = np.clip(-(base - theta), -500.0, 500.0)
        return np.sum(_MHC_WTS * (2.0 * s / (1.0 + np.exp(arg))))

    def I_ox(theta):
        arg = np.clip(-base - theta, -500.0, 500.0)
        return np.sum(_MHC_WTS * (-2.0 * s / (1.0 + np.exp(arg))))

    return K0 * I_red(Theta) / I_red(0.0), K0 * I_ox(Theta) / I_ox(0.0)


def et_rates_unified(Theta, K0, log10_reorg, asym):
    """Unified (K_red, K_ox) electron-transfer rates: smooth blend of Butler-
    Volmer (transfer coefficient = asym) and symmetric Marcus-Hush-Chidsey,
    weighted by reorg so BV is the large-lambda limit and MHC the finite-lambda
    one. Used for both ET1 and ET2 so the ET sub-type axis is continuous."""
    w = _bv_weight(log10_reorg)
    if w > 1.0 - _W_EPS:                       # pure BV limit (fast path)
        alpha = float(asym); beta = 1.0 - alpha
        K_red = K0 * np.exp(np.clip(-alpha * Theta, -500, 500))
        K_ox = K0 * np.exp(np.clip(beta * Theta, -500, 500))
        return _clamp_bv_rates(K_red, K_ox)
    reorg = min(10.0 ** log10_reorg, _REORG_MHC_MAX)
    if w < _W_EPS:                              # pure symmetric MHC (fast path)
        return _mhc_rates_smooth(Theta, K0, reorg)
    # blended regime
    alpha = float(asym); beta = 1.0 - alpha
    K_red_bv = K0 * np.exp(np.clip(-alpha * Theta, -500, 500))
    K_ox_bv = K0 * np.exp(np.clip(beta * Theta, -500, 500))
    K_red_bv, K_ox_bv = _clamp_bv_rates(K_red_bv, K_ox_bv)
    K_red_mhc, K_ox_mhc = _mhc_rates_smooth(Theta, K0, reorg)
    return (w * K_red_bv + (1.0 - w) * K_red_mhc,
            w * K_ox_bv + (1.0 - w) * K_ox_mhc)


def et1_rates(Theta, log10_K0, log10_reorg, asym):
    """ET1 rate (K_red, K_ox) and an 'is_reversible' flag via the unified law.

    Returns (K_red, K_ox, reversible). When reversible (K0 large), the caller
    uses the Nernst boundary condition instead of a Robin (kinetic) surface row.
    """
    if log10_K0 >= REVERSIBLE_K0 - 0.5:
        return None, None, True
    K_red, K_ox = et_rates_unified(Theta, 10.0 ** log10_K0, log10_reorg, asym)
    return K_red, K_ox, False


def _run_ee_3species(theta_super, sigma=1.0, C_A_bulk=1.0, C_B_bulk=0.0,
                     C_C_bulk=0.0, theta_i=20.0, theta_v=-20.0, cycles=1,
                     irreversible_et2=False):
    """EE (two sequential electron transfers) with a tracked 3rd species C and
    a real, reversible 2nd electron transfer B + e- <-> C.

    State vector length 3n: A in [0:n], B in [n:2n], C in [2n:3n]. Surface at
    A[n-1], B[n], C[2n]; far bulk at A[0], B[2n-1], C[3n-1].

    irreversible_et2=True zeros K_ox_2 (no reverse 2nd ET); in that limit C is a
    passive bystander and A/B/flux match the existing run_ee_simulation exactly.
    """
    p = _unpack(theta_super)
    K0_1 = 10.0 ** p["log10_K0_1"]; alpha_1 = float(p["asym_1"]); lam1 = float(p["log10_reorg_e_1"])
    K0_2 = 10.0 ** p["log10_K0_2"]; alpha_2 = float(p["asym_2"]); lam2 = float(p["log10_reorg_e_2"])
    E0_2 = float(p["E0_2_offset"])
    dA = 10.0 ** p["log10_dA"]; dB = 10.0 ** p["log10_dB"]; dC = 10.0 ** p["log10_dC"]

    deltaT = DELTA_THETA / sigma
    maxT = cycles * 2.0 * abs(theta_v - theta_i) / sigma
    E = _make_potential_waveform(theta_i, theta_v, cycles)
    total_steps = len(E)

    X_grid, n = gen_grid(0.0, DELTA_X, SIMULATION_SPACE_MULTIPLE * np.sqrt(maxT),
                         EXPANDING_GRID_FACTOR)
    conc = np.zeros(3 * n)
    conc[0:n] = C_A_bulk; conc[n:2 * n] = C_B_bulk; conc[2 * n:3 * n] = C_C_bulk
    conc_d = conc.copy()

    z = np.zeros(n)
    aA, bA, cA = calc_abc_linear(n, X_grid, deltaT, z.copy(), z.copy(), z.copy(), dA)
    aB, bB, cB = calc_abc_linear(n, X_grid, deltaT, z.copy(), z.copy(), z.copy(), dB)
    aC, bC, cC = calc_abc_linear(n, X_grid, deltaT, z.copy(), z.copy(), z.copy(), dC)

    fluxes = np.zeros(total_steps)
    X0 = X_grid[1] - X_grid[0]
    # Constant (rate-independent) entries assembled ONCE: interior diffusion bands,
    # far-field Dirichlet, and the -1 surface-neighbour links. Only the 7 surface
    # rate entries change per step (built sparse -> O(nnz) per step, not O(n^3)).
    rA = np.arange(n - 2, 0, -1); rB = np.arange(n + 1, 2 * n - 1)
    rC = np.arange(2 * n + 1, 3 * n - 1)
    rows_c = np.concatenate([rA, rA, rA, rB, rB, rB, rC, rC, rC,
                             [0, 2 * n - 1, 3 * n - 1, n - 1, n, 2 * n]])
    cols_c = np.concatenate([rA - 1, rA, rA + 1, rB - 1, rB, rB + 1, rC - 1, rC, rC + 1,
                             [0, 2 * n - 1, 3 * n - 1, n - 2, n + 1, 2 * n + 1]])
    vals_c = np.concatenate([cA[1:n - 1], bA[1:n - 1], aA[1:n - 1],
                             aB[1:n - 1], bB[1:n - 1], cB[1:n - 1],
                             aC[1:n - 1], bC[1:n - 1], cC[1:n - 1],
                             [1.0, 1.0, 1.0, -1.0, -1.0, -1.0]]).astype(np.float64)
    s_rows = np.array([n - 1, n - 1, n, n, n, 2 * n, 2 * n])
    s_cols = np.array([n - 1, n, n - 1, n, 2 * n, n, 2 * n])
    all_rows = np.concatenate([rows_c, s_rows]); all_cols = np.concatenate([cols_c, s_cols])

    for idx in range(total_steps):
        Theta = E[idx]; Theta2 = Theta - E0_2
        K_red_1, K_ox_1 = et_rates_unified(Theta, K0_1, lam1, alpha_1)
        K_red_2, K_ox_2 = et_rates_unified(Theta2, K0_2, lam2, alpha_2)
        if irreversible_et2:
            K_ox_2 = 0.0
        s_vals = np.array([1.0 + X0 / dA * K_red_1, -X0 / dA * K_ox_1,
                           -X0 / dB * K_red_1, 1.0 + X0 / dB * (K_ox_1 + K_red_2),
                           -X0 / dB * K_ox_2,
                           -X0 / dC * K_red_2, 1.0 + X0 / dC * K_ox_2])
        A_matrix = csr_matrix((np.concatenate([vals_c, s_vals]), (all_rows, all_cols)),
                              shape=(3 * n, 3 * n))

        conc_d[:] = conc[:]
        conc_d[n - 1] = 0.0; conc_d[n] = 0.0; conc_d[2 * n] = 0.0
        conc_d[0] = C_A_bulk; conc_d[2 * n - 1] = C_B_bulk; conc_d[3 * n - 1] = C_C_bulk

        conc = spsolve(A_matrix, conc_d)
        flux_1 = calc_flux(conc, n, dA, X_grid)
        flux_2 = K_red_2 * conc[n] - K_ox_2 * conc[2 * n]
        fluxes[idx] = flux_1 + flux_2

    return {"potential": E, "flux": fluxes, "n": n}


def _run_ece_3species(theta_super, sigma=1.0, C_A_bulk=1.0, C_B_bulk=0.0,
                      C_C_bulk=0.0, theta_i=20.0, theta_v=-20.0, cycles=1):
    """True ECE: ET1 (A + e- <-> B), chemical B -> C (rate kc), ET2 (C + e- <-> D).

    3-species state (A, B, C); D untracked (irreversible 2nd ET product, D=0).
    ET1 acts on A/B at the surface; the chemical step B->C is operator-split in
    the bulk (same physical depth, B and C share orientation); ET2 consumes C at
    the surface. flux = flux_1(ET1) + flux_2(ET2 on C).

    (The existing run_ece_simulation is EC-with-decay - its 2nd ET is dead code -
    so this is a corrected mechanism, validated by physics not by matching it.)
    """
    p = _unpack(theta_super)
    K0_1 = 10.0 ** p["log10_K0_1"]; alpha_1 = float(p["asym_1"]); lam1 = float(p["log10_reorg_e_1"])
    K0_2 = 10.0 ** p["log10_K0_2"]; alpha_2 = float(p["asym_2"]); lam2 = float(p["log10_reorg_e_2"])
    E0_2 = float(p["E0_2_offset"])
    kc = 10.0 ** p["log10_kc"]
    dA = 10.0 ** p["log10_dA"]; dB = 10.0 ** p["log10_dB"]; dC = 10.0 ** p["log10_dC"]

    deltaT = DELTA_THETA / sigma
    decay_kc = np.exp(-kc * deltaT)
    maxT = cycles * 2.0 * abs(theta_v - theta_i) / sigma
    E = _make_potential_waveform(theta_i, theta_v, cycles)
    total_steps = len(E)

    X_grid, n = gen_grid(0.0, DELTA_X, SIMULATION_SPACE_MULTIPLE * np.sqrt(maxT),
                         EXPANDING_GRID_FACTOR)
    conc = np.zeros(3 * n)
    conc[0:n] = C_A_bulk; conc[n:2 * n] = C_B_bulk; conc[2 * n:3 * n] = C_C_bulk
    conc_d = conc.copy()

    z = np.zeros(n)
    aA, bA, cA = calc_abc_linear(n, X_grid, deltaT, z.copy(), z.copy(), z.copy(), dA)
    aB, bB, cB = calc_abc_linear(n, X_grid, deltaT, z.copy(), z.copy(), z.copy(), dB)
    aC, bC, cC = calc_abc_linear(n, X_grid, deltaT, z.copy(), z.copy(), z.copy(), dC)

    fluxes = np.zeros(total_steps)
    X0 = X_grid[1] - X_grid[0]
    # Constant entries assembled once (see _run_ee_3species). ECE surface couples
    # ET1 on A/B and ET2 consuming C; 5 rate entries change per step.
    rA = np.arange(n - 2, 0, -1); rB = np.arange(n + 1, 2 * n - 1)
    rC = np.arange(2 * n + 1, 3 * n - 1)
    rows_c = np.concatenate([rA, rA, rA, rB, rB, rB, rC, rC, rC,
                             [0, 2 * n - 1, 3 * n - 1, n - 1, n, 2 * n]])
    cols_c = np.concatenate([rA - 1, rA, rA + 1, rB - 1, rB, rB + 1, rC - 1, rC, rC + 1,
                             [0, 2 * n - 1, 3 * n - 1, n - 2, n + 1, 2 * n + 1]])
    vals_c = np.concatenate([cA[1:n - 1], bA[1:n - 1], aA[1:n - 1],
                             aB[1:n - 1], bB[1:n - 1], cB[1:n - 1],
                             aC[1:n - 1], bC[1:n - 1], cC[1:n - 1],
                             [1.0, 1.0, 1.0, -1.0, -1.0, -1.0]]).astype(np.float64)
    s_rows = np.array([n - 1, n - 1, n, n, 2 * n])
    s_cols = np.array([n - 1, n, n - 1, n, 2 * n])
    all_rows = np.concatenate([rows_c, s_rows]); all_cols = np.concatenate([cols_c, s_cols])

    for idx in range(total_steps):
        Theta = E[idx]; Theta2 = Theta - E0_2
        K_red_1, K_ox_1 = et_rates_unified(Theta, K0_1, lam1, alpha_1)
        K_red_2, K_ox_2 = et_rates_unified(Theta2, K0_2, lam2, alpha_2)
        s_vals = np.array([1.0 + X0 / dA * K_red_1, -X0 / dA * K_ox_1,
                           -X0 / dB * K_red_1, 1.0 + X0 / dB * K_ox_1,
                           1.0 + X0 / dC * K_red_2])
        A_matrix = csr_matrix((np.concatenate([vals_c, s_vals]), (all_rows, all_cols)),
                              shape=(3 * n, 3 * n))

        conc_d[:] = conc[:]
        conc_d[n - 1] = 0.0; conc_d[n] = 0.0; conc_d[2 * n] = 0.0
        conc_d[0] = C_A_bulk; conc_d[2 * n - 1] = C_B_bulk; conc_d[3 * n - 1] = C_C_bulk

        conc = spsolve(A_matrix, conc_d)
        flux_1 = calc_flux(conc, n, dA, X_grid)
        flux_2 = K_red_2 * conc[2 * n]
        fluxes[idx] = flux_1 + flux_2

        # operator-split chemical B -> C (same physical depth; shared orientation)
        removed = conc[n:2 * n - 1] * (1.0 - decay_kc)
        conc[n:2 * n - 1] -= removed
        conc[2 * n:3 * n - 1] += removed
        conc[2 * n - 1] = C_B_bulk

    return {"potential": E, "flux": fluxes, "n": n}


def run_cv_master(theta_super, sigma=1.0, C_A_bulk=None, C_B_bulk=0.0,
                  C_Y_bulk=1.0, theta_i=20.0, theta_v=-20.0, cycles=1):
    """Master CV simulation (M1a ET1 + M1b chemical steps).

    Active chemical steps (from theta_super log-rates above the off-floor):
      - following chemical B->Y   (log10_kc)   : conc[n:2n-1] *= exp(-kc*dt)
      - catalytic regeneration B->A (log10_kcat): operator-split B->A transfer
      - preceding equilibrium Y<->A (log10_kf, log10_Keq): source toward C_A_eq,
        with bulk/initial C_A = Keq*C_Y/(1+Keq).
    """
    p = _unpack(theta_super)

    # ET2 active -> 3-species path. With a chemical step it is ECE; without, EE.
    if p["log10_K0_2"] > _ON:
        ca = 1.0 if C_A_bulk is None else C_A_bulk
        if p["log10_kc"] > _ON:
            return _run_ece_3species(
                theta_super, sigma=sigma, C_A_bulk=ca, C_B_bulk=C_B_bulk,
                theta_i=theta_i, theta_v=theta_v, cycles=cycles)
        return _run_ee_3species(
            theta_super, sigma=sigma, C_A_bulk=ca, C_B_bulk=C_B_bulk,
            theta_i=theta_i, theta_v=theta_v, cycles=cycles)

    dA = 10.0 ** p["log10_dA"]
    dB = 10.0 ** p["log10_dB"]

    # chemical-step activity + rates
    kc_on = p["log10_kc"] > _ON
    kcat_on = p["log10_kcat"] > _ON
    pre_on = p["log10_kf"] > _ON
    kc = 10.0 ** p["log10_kc"] if kc_on else 0.0
    kcat = 10.0 ** p["log10_kcat"] if kcat_on else 0.0
    kf = 10.0 ** p["log10_kf"] if pre_on else 0.0
    Keq = 10.0 ** p["log10_Keq"]
    C_A_eq = Keq * C_Y_bulk / (1.0 + Keq) if pre_on else None
    if C_A_bulk is None:
        C_A_bulk = C_A_eq if pre_on else 1.0

    deltaT = DELTA_THETA / sigma
    decay_kc = np.exp(-kc * deltaT)
    decay_kcat = np.exp(-kcat * deltaT)
    maxT = cycles * 2.0 * abs(theta_v - theta_i) / sigma
    nTimeSteps = int(2 * abs(theta_v - theta_i) / DELTA_THETA) + 1
    Esteps = np.arange(nTimeSteps)
    E = np.where(Esteps < nTimeSteps / 2.0,
                 theta_i - DELTA_THETA * Esteps,
                 theta_v + DELTA_THETA * (Esteps - nTimeSteps / 2.0))
    E = np.tile(E, cycles)
    total_steps = len(E)

    maxX = SIMULATION_SPACE_MULTIPLE * np.sqrt(maxT)
    X_grid, n = gen_grid(0.0, DELTA_X, maxX, EXPANDING_GRID_FACTOR)
    conc, conc_d = ini_conc(n, C_A_bulk, C_B_bulk)

    A_matrix, aA, bA, cA, aB, bB, cB = ini_coeff(n)
    aA, bA, cA = calc_abc_linear(n, X_grid, deltaT, aA, bA, cA, dA)
    aB, bB, cB = calc_abc_linear(n, X_grid, deltaT, aB, bB, cB, dB)
    X0 = X_grid[1] - X_grid[0]

    fluxes = np.zeros(total_steps)
    for idx in range(total_steps):
        Theta = E[idx]
        K_red, K_ox, reversible = et1_rates(
            Theta, p["log10_K0_1"], p["log10_reorg_e_1"], p["asym_1"])

        # --- assemble matrix (interior diffusion bands, reused layout) ---
        A_matrix[:] = 0.0
        rows_A = np.arange(n - 2, 0, -1)
        A_matrix[rows_A, rows_A - 1] = cA[1:n - 1]
        A_matrix[rows_A, rows_A] = bA[1:n - 1]
        A_matrix[rows_A, rows_A + 1] = aA[1:n - 1]
        rows_B = np.arange(n + 1, 2 * n - 1)
        A_matrix[rows_B, rows_B - 1] = aB[1:n - 1]
        A_matrix[rows_B, rows_B] = bB[1:n - 1]
        A_matrix[rows_B, rows_B + 1] = cB[1:n - 1]

        # --- bulk far-field Dirichlet rows ---
        A_matrix[0, 0] = 1.0
        A_matrix[0, 1] = 0.0
        A_matrix[2 * n - 1, 2 * n - 1] = 1.0
        A_matrix[2 * n - 1, 2 * n - 2] = 0.0

        # --- surface boundary rows ---
        if reversible:  # Nernst
            A_matrix[n - 1, n - 1] = 1.0
            A_matrix[n, n - 2] = -dA
            A_matrix[n, n - 1] = dA
            A_matrix[n, n] = dB
            A_matrix[n, n + 1] = -dB
        else:           # Robin (BV / MHC)
            A_matrix[n - 1, n - 2] = -1.0
            A_matrix[n - 1, n - 1] = 1.0 + X0 / dA * K_red
            A_matrix[n - 1, n] = -X0 / dA * K_ox
            A_matrix[n, n - 1] = -X0 / dB * K_red
            A_matrix[n, n] = 1.0 + X0 / dB * K_ox
            A_matrix[n, n + 1] = -1.0

        # --- RHS ---
        conc_d[:] = conc[:]
        if reversible:
            conc_d[n - 1] = 1.0 / (1.0 + np.exp(-Theta))
            conc_d[n] = 0.0
        else:
            conc_d[n - 1] = 0.0
            conc_d[n] = 0.0
        conc_d[0] = C_A_bulk
        conc_d[2 * n - 1] = C_B_bulk

        conc = scipy.linalg.solve(A_matrix, conc_d)
        fluxes[idx] = calc_flux(conc, n, dA, X_grid)

        # --- operator-split chemical steps (after flux, before next step) ---
        if kc_on:  # following chemical B -> Y
            conc[n:2 * n - 1] *= decay_kc
            conc[2 * n - 1] = C_B_bulk
        if kcat_on:  # catalytic regeneration B -> A
            b_interior = conc[n + 2:2 * n - 1].copy()
            amount = b_interior * (1.0 - decay_kcat)
            conc[n + 2:2 * n - 1] -= amount
            conc[1:n - 2] += amount[::-1]
        if pre_on:  # preceding equilibrium Y <-> A
            for i in range(2, n - 1):
                c_a = conc[n - 1 - i]
                conc[n - 1 - i] = max(c_a + kf * (C_A_eq - c_a) * deltaT, 0.0)

    return {"potential": E, "flux": fluxes, "n": n}