#!/usr/bin/env python3
"""BEF-2026-R: revised design (corrections that maximise the near-light-speed shell term and the positron yield).
Computes v1 (as stated) and v2 (revised) side by side."""
import math, json
pi = math.pi
c = 2.99792458e8; G = 6.67430e-11; hbar = 1.054571817e-34; h = 2*pi*hbar
mu0 = 4e-7*pi; eps0 = 8.8541878128e-12; me = 9.1093837015e-31; e = 1.602176634e-19
sigmaT = 6.6524587e-29; day = 86400.0
def gamma_of(V): return 1 + V*e/(me*c**2)
def vol_ell(a, b): return 4/3*pi*a*b*b
def ell_area(a, b, p=1.6075):
    return 4*pi*((2*(a*b)**p + b**(2*p))/3)**(1/p)
def half_perimeter(a, b, n=20000):
    s = 0.
    for i in range(n):
        t0, t1 = pi*i/n, pi*(i+1)/n
        s += math.hypot(b*math.sin(t1)-b*math.sin(t0), a*math.cos(t1)-a*math.cos(t0))
    return s
def sphere_layers(r_in, layers):
    r = r_in; out = []; m = 0.
    for name, t, rho in layers:
        v = 4/3*pi*((r+t)**3 - r**3); out.append((name, t, rho, v, v*rho/1000.)); m += v*rho; r += t
    return out, m, r
def phi_mass(M, r): return -G*M/(c**2*r)
def phi_energy(E, r, fac=1.0): return -G*E*fac/(c**4*r)

def model(v):
    R = {}
    rev = (v == "v2")
    # ---------- geometry ----------
    a_out, b_out = 6.0, 5.0
    layers_in = [("SiC-SiC", 0.90, 2700.), ("W-steel", 0.90, 12000.), ("Pb", 1.00, 11340.), ("water", 0.30, 1000.)]
    a, b, m_in = a_out, b_out, 0.
    for name, t, rho in layers_in:
        v1 = vol_ell(a, b); a -= t; b -= t; m_in += (v1 - vol_ell(a, b))*rho
    V_cav = vol_ell(a, b)
    coil_layer_t = 0.5 if rev else 0.0           # 5 T REBCO guide/storage coil layer on the inner capsule surface
    A_surf = ell_area(a_out, b_out)
    m_coil = A_surf*coil_layer_t*6000. if rev else 0.
    r_sheath = a_out + coil_layer_t + 0.5          # sheath ~0.5 m outside the (coil) surface, polar radius
    layers_sph = [("W-steel", 0.90, 12000.), ("water", 0.50, 1000.), ("SiC-SiC", 1.50, 2700.)]
    lo, m_out, r_out_ext = sphere_layers(8.0, layers_sph)
    r_extra_in = 14.0 if rev else 11.0
    le, m_ext, r_ext_ext = sphere_layers(r_extra_in, layers_sph)
    gap_extra = r_extra_in - r_out_ext
    R.update(m_in_t=(m_in+m_coil)/1e3, m_coil_t=m_coil/1e3, m_out_t=m_out/1e3, m_ext_t=m_ext/1e3, r_out_ext=r_out_ext,
             r_ext_ext=r_ext_ext, r_extra_in=r_extra_in, gap_extra=gap_extra, A_surf=A_surf, V_cav=V_cav,
             V_gap1=4/3*pi*8**3 - vol_ell(a_out, b_out), V_gap2=4/3*pi*r_extra_in**3 - 4/3*pi*r_out_ext**3)
    # ---------- beams ----------
    V1 = 30e6 if rev else 25e6; V2 = 30e6
    I1 = 20e-3 if rev else 2e-3; I2 = 20e-3 if rev else 2e-3
    g1, g2 = gamma_of(V1), gamma_of(V2)
    P1, P2, P_L = V1*I1, V2*I2, 20e3
    lam = 1064e-9; omega0 = 2*pi*c/lam; Eph = hbar*omega0
    L1 = half_perimeter(a_out+coil_layer_t+0.5, b_out+coil_layer_t+0.5)
    L2 = pi*(r_out_ext + (1.5 if rev else 0.05))
    t1, t2 = L1/c, L2/c
    E1, E2, EL = P1*t1, P2*t2, P_L*t1
    p1 = math.sqrt(g1**2-1)*me*c
    R.update(V1_MV=V1/1e6, V2_MV=V2/1e6, I_mA=I1*1e3, gamma1=g1, gamma2=g2, beta1=math.sqrt(1-1/g1**2),
             P1_kW=P1/1e3, P2_kW=P2/1e3, L1=L1, L2=L2, t1_ns=t1*1e9, t2_ns=t2*1e9,
             E1_mJ=E1*1e3, E2_mJ=E2*1e3, EL_mJ=EL*1e3, Ne_s=I1/e, E_field_gap2_MVm=V2/gap_extra/1e6)
    B_pole = 10.0 if rev else 2.0; B_guide = 5.0 if rev else 0.0
    R.update(B_pole=B_pole, B_guide=B_guide, rL_pole_cm=p1/(e*B_pole)*100, rL_guide_cm=(p1/(e*B_guide)*100 if B_guide else None))
    # ---------- electron storage layer (v2) ----------
    if rev:
        theta = math.radians(10.)                 # pitch angle in the guide field
        U_B = B_guide**2/(2*mu0)
        P_sync = 2*sigmaT*c*U_B*(g1**2)*math.sin(theta)**2
        E_e = g1*me*c**2
        tau_sync = E_e/P_sync
        tau_conf = tau_sync                       # confinement limited by synchrotron cooling (pitch 10 deg)
        E_store = (P1+P2)*tau_conf                # steady stored energy = beam power x confinement time
        n_B = eps0*B_guide**2/(2*g1*me)           # Brillouin density
        N_store = E_store/E_e
        layer_t = 0.04; V_layer = A_surf*layer_t
        n_store = N_store/V_layer
        R.update(P_sync_W=P_sync, tau_conf_s=tau_conf, E_store_kJ=E_store/1e3, N_store=N_store, n_store=n_store,
                 n_Brillouin=n_B, brillouin_fraction=n_store/n_B, fill_time_s=N_store/((I1+I2)/e))
        # SRF photon shell
        f_rf = 1.3e9; Q = 1e10; U_srf = 20e3
        P_2K = 2*pi*f_rf*U_srf/Q
        R.update(U_srf_kJ=U_srf/1e3, P_2K_kW=P_2K/1e3, P_cryo_srf_MW=P_2K*700/1e6, tau_srf_s=Q/(2*pi*f_rf))
    else:
        E_store = 0.; U_srf = 0.; tau_conf = 0.
    # ---------- magnets ----------
    B_ring = 20.0 if rev else 10.0; R_ring = 13.0; a_coil = 0.5
    I_ring = B_ring*2*pi*a_coil/mu0
    E_ring = B_ring**2/(2*mu0)*pi*a_coil**2*2*pi*R_ring
    z_c = r_ext_ext + 1.0
    B_center = 2*mu0*I_ring*R_ring**2/(2*(R_ring**2+z_c**2)**1.5)
    E_pole = 2*(B_pole**2/(2*mu0))*1.0
    E_guide = (B_guide**2/(2*mu0))*A_surf*coil_layer_t if rev else 0.
    E_es = 0.5*4*pi*eps0*r_out_ext*r_extra_in/gap_extra*V2**2
    R.update(B_ring=B_ring, I_ring_MA=I_ring/1e6, E_ring_GJ=E_ring/1e9, B_center=B_center, E_pole_MJ=E_pole/1e6,
             E_guide_GJ=E_guide/1e9, E_es_MJ=E_es/1e6, z_c=z_c)
    # ---------- rotation ----------
    rpm = 120. if rev else 700.; om = rpm*2*pi/60
    m_ring = 2000e3 if rev else 1000e3
    I_plat = 2/3*m_ext*r_ext_ext**2 + m_ring*R_ring**2
    KE = 0.5*I_plat*om**2; v_rim = om*r_ext_ext
    Omega_LT = 4/3*G*(m_ext+m_ring)*om/(c**2*r_ext_ext)
    R.update(rpm=rpm, v_rim=v_rim, a_cf_g=om**2*r_ext_ext/9.81, KE_TJ=KE/1e12, KE_MWh=KE/3.6e9,
             hoop_SiC_MPa=2700*v_rim**2/1e6, hoop_W_MPa=12000*v_rim**2/1e6, Omega_LT=Omega_LT)
    # ---------- potential budget ----------
    Phi = {}
    Phi["wall_inner"] = phi_mass(m_in+m_coil, 4.0)
    Phi["wall_outer"] = phi_mass(m_out, 9.45)
    Phi["wall_extra"] = phi_mass(m_ext, r_extra_in+1.45)
    Phi["rotation"] = phi_energy(KE, r_ext_ext)
    Phi["rings"] = phi_energy(2*E_ring, math.hypot(R_ring, z_c))
    Phi["guide_coil"] = phi_energy(E_guide, a_out+0.25) if rev else 0.
    Phi["pole_magnets"] = phi_energy(E_pole, 3.0)
    Phi["electrostatic"] = phi_energy(E_es, 8.0)
    Phi["transit_arc1"] = phi_energy(E1, r_sheath, 2.0)
    Phi["transit_laser"] = phi_energy(EL, r_sheath, 2.0)
    Phi["transit_arc2"] = phi_energy(E2, r_out_ext+1.5 if rev else r_out_ext+0.05, 2.0)
    Phi["stored_electrons"] = phi_energy(E_store, r_sheath, 2.0) if rev else 0.
    Phi["srf_photon_shell"] = phi_energy(U_srf, a_out+coil_layer_t+1.0, 2.0) if rev else 0.
    Phi_walls = Phi["wall_inner"]+Phi["wall_outer"]+Phi["wall_extra"]
    Phi_fields = Phi["rings"]+Phi["guide_coil"]+Phi["pole_magnets"]+Phi["electrostatic"]
    Phi_s = Phi["transit_arc1"]+Phi["transit_laser"]+Phi["transit_arc2"]+Phi["stored_electrons"]+Phi["srf_photon_shell"]
    R["Phi"] = Phi; R["Phi_walls"] = Phi_walls; R["Phi_fields"] = Phi_fields; R["Phi_s"] = Phi_s
    R["dt_walls_as_day"] = -Phi_walls*day*1e18; R["dt_s_day"] = -Phi_s*day
    R["E_sheath_total_J"] = E1+E2+EL+E_store+U_srf
    P_dep = P1+P2+P_L
    R["Phi_acc"] = phi_energy(P_dep*day, 9.0)
    # curvature
    u_sheath = (E1+EL+E_store+U_srf)/R["V_gap1"]
    R["K_sheath"] = 8*pi*G*u_sheath/c**4
    # ---------- Borekci metric ----------
    R["xi_L1"] = (1-1/g1**2)/(2*(-Phi_s))
    R["xi_bounds"] = {k: s/(-Phi_s) for k, s in (("CSAC", 1e-12), ("Hmaser", 1e-15), ("optical", 1e-18), ("lockin", 5e-19))}
    R["Gamma_in_L2_minus1"] = -Phi_s
    # ---------- chopper ----------
    T_cycle = 300.; tau_cut = 1.0 if rev else 1e-3
    kappa0 = 200/omega0; T_sw = 1e-9
    n1D = kappa0/(4*T_sw) + kappa0**2/(16*T_sw**2); N_modes = pi*0.05**2/lam**2
    N_DCE = n1D*N_modes
    R.update(T_cycle=T_cycle, tau_cut=tau_cut, duty_loss=tau_cut/T_cycle, N_DCE=N_DCE, E_DCE_J=N_DCE*Eph,
             mod_depth_store=(1-math.exp(-tau_cut/tau_conf)) if rev else 1.0)
    Q_sh = R["E_sheath_total_J"]*r_sheath**2/c**2
    R["h_chop"] = 2*G*(Q_sh/(0.1 if rev else 1e-6)**2)/(c**4*r_sheath)   # v2: shell decays in ~0.1 s
    # ---------- positrons ----------
    Ne_day = ((I1+I2)/e)*day
    Y = {"conservative": 1.5e-2 if rev else 1.0e-2, "optimistic": 3.0e-2 if rev else 2.3e-2}
    eps_mod = 1e-2 if rev else 1e-3; eps_trap = 0.2 if rev else 0.1
    R["Ne_day"] = Ne_day
    R["pos_fast"] = {k: Ne_day*y for k, y in Y.items()}
    R["pos_fast_ng"] = {k: Ne_day*y*me*1e12 for k, y in Y.items()}
    R["pos_slow_trapped"] = {k: Ne_day*y*eps_mod*eps_trap for k, y in Y.items()}
    R["pos_slow_ng"] = {k: Ne_day*y*eps_mod*eps_trap*me*1e12 for k, y in Y.items()}
    R["value_fast_MUSD"] = {k: Ne_day*y*me*1e3*62.5e12/1e6 for k, y in Y.items()}
    R["value_trapped_MUSD"] = {k: Ne_day*y*eps_mod*eps_trap*me*1e3*62.5e12/1e6 for k, y in Y.items()}
    R["pos_rate_fast_per_s"] = {k: (I1+I2)/e*y for k, y in Y.items()}
    R["pos_rate_slow_per_s"] = {k: (I1+I2)/e*y*eps_mod for k, y in Y.items()}
    R["anode_power_kW"] = (P1+P2)/1e3
    # ---------- power & cost ----------
    if rev:
        P_wall = {"arcs": P1+P2, "laser": P_L/0.3, "cryo_rings_20T": 4e6, "cryo_guide_coil": 1e6, "cryo_srf": R["P_cryo_srf_MW"]*1e6,
                  "cryo_pole": 0.4e6, "rotation_drive_120rpm": 0.3e6, "anode_cooling_vacuum_control": 1.5e6, "trap_line_moderator": 0.3e6}
        capex = (3.0e9, 15e9)
    else:
        P_wall = {"arcs": P1+P2, "laser": P_L/0.3, "cryo_rings": 2e6, "cryo_inner": 0.2e6, "rotation_drive": 3e6, "vacuum_cooling_control": 1e6}
        capex = (2.4e9, 12e9)
    P_tot = sum(P_wall.values()); E_day_MWh = P_tot*day/3.6e9; E_spin_MWh = KE/3.6e9
    cost_energy = (E_day_MWh+E_spin_MWh)*1e3*0.10; opex = 500e3 if rev else 400e3
    R.update(P_wall_MW={k: v/1e6 for k, v in P_wall.items()}, P_tot_MW=P_tot/1e6, E_day_MWh=E_day_MWh, E_spin_MWh=E_spin_MWh,
             cost_energy_MUSD=cost_energy/1e6, opex_MUSD=opex/1e6, amort_MUSD=(capex[0]/2000/1e6, capex[1]/2000/1e6), capex_B=(capex[0]/1e9, capex[1]/1e9))
    fee = 5e6
    R["result_fee_only_MUSD"] = (fee - cost_energy - opex)/1e6
    R["result_fee_fastvalue_MUSD"] = {k: (fee + R["value_fast_MUSD"][k]*1e6 - cost_energy - opex)/1e6 for k in Y}
    R["result_fee_trapped_MUSD"] = {k: (fee + R["value_trapped_MUSD"][k]*1e6 - cost_energy - opex)/1e6 for k in Y}
    return R

out = {"v1": model("v1"), "v2": model("v2")}
json.dump(out, open("/home/claude/bef/results_v2.json", "w"), indent=1, default=str)
for v in ("v1", "v2"):
    print("=" * 30, v)
    for k, val in out[v].items():
        print(f"{k}: {val}")
