"""
GAI-TP-002 Rev B -- reduced physical model for EAF shell cooling.

Computes, for every cooling circuit on a developed shell map:
  incident irradiation G       (arc columns + melt surface, foam-attenuated, cavity-augmented)
  net absorbed flux q''        (radiative balance at the fire-side surface)
  frozen-slag skull thickness  (self-consistent: outer face sits at slag solidus)
  panel hot-face temperature   (series resistance chain to coolant bulk)
  coolant-side state           (Dittus-Boelter, per-circuit bulk rise, boiling margin)
  energy closure               (integral of reconstructed field vs summed circuit extraction)

All parameters are declared in PARAMS and printed with the results.
"""
import numpy as np
from scipy.optimize import brentq
import json

SIG = 5.670374419e-8

P = dict(
    # ---- furnace geometry -------------------------------------------------
    R_shell        = 3.20,   # m, shell inner radius
    R_melt         = 3.00,   # m, melt/slag surface radius
    r_electrode    = 1.10,   # m, electrode pitch-circle radius
    theta_elec     = [90., 210., 330.],   # deg
    z_rows         = {'A': 1.20, 'B': 2.40, 'C': 3.60},   # m above sill
    h_row          = 1.20,   # m, panel height
    z_roof         = 4.90,   # m
    r_roof_in      = 2.40,
    r_roof_out     = 3.20,
    n_sectors      = 12,

    # ---- thermal sources --------------------------------------------------
    P_furnace      = 98.0e6, # W, electrical input
    f_arc_rad      = 0.28,   # fraction of input radiated from the arc columns
    L_arc          = 0.45,   # m, arc length
    T_foam_top     = 1748.,  # K, radiating temperature of a covered (foamed) surface
    T_melt_bare    = 2123.,  # K, radiating temperature of exposed slag/steel
    h_ref_foam     = 0.15,   # m, foam depth over which coverage becomes effective
    eps_melt       = 0.90,
    f_cavity       = 1.18,   # single-bounce cavity augmentation on incident irradiation

    # ---- foamy slag -------------------------------------------------------
    h_foam_base    = 0.35,   # m
    theta_door     = 180.,   # deg
    dh_door        = 0.32,   # m depression at the door
    sigma_door     = 38.,    # deg
    kappa_foam     = 4.0,    # 1/m extinction through foam
    h_foam_min     = 0.20,   # m, coverage below which skull delivery fails
    T_freeboard    = 1873.,  # K, freeboard gas temperature
    eps_freeboard  = 0.15,   # dusty CO/CO2 emissivity

    # ---- surfaces ---------------------------------------------------------
    eps_skull      = 0.85,
    eps_panel      = 0.80,
    T_solidus      = 1523.,  # K, frozen-slag solidus (1250 C)
    k_skull        = 1.50,   # W/m-K
    t_skull_max    = 0.050,  # m, mechanical sloughing limit
    z_splash       = 3.00,   # m, upper limit of molten-slag supply
    t_dust         = 10.0e-3,# m, dust/oxide crust outside the splash zone
    k_dust         = 0.90,   # W/m-K

    # ---- panel construction ----------------------------------------------
    k_wall         = 45.0,   # W/m-K, steel
    t_wall         = 0.025,  # m, effective conduction path
    k_deposit      = 1.00,   # W/m-K, water-side scale
    psi_area       = 1.20,   # wetted/face area ratio for the tube panel

    # ---- coolant ----------------------------------------------------------
    Q_total        = 1080.,  # m3/h, total circulation
    T_inlet        = 35.0,   # C
    p_circuit      = 6.0,    # bar(a)  -> T_sat approx 158.8 C
    T_sat          = 158.8,  # C
    D_tube         = 0.060,  # m
    rho_w          = 990., cp_w = 4180., mu_w = 5.5e-4, k_w = 0.640, Pr_w = 3.55,

    # ---- off-gas duct -----------------------------------------------------
    T_offgas       = 1523.,  # K
    h_offgas       = 45.,    # W/m2-K
    eps_offgas     = 0.25,

    # ---- limits -----------------------------------------------------------
    T_hotface_limit = 400.,  # C
    t_deposit_limit = 0.85e-3,  # m
    flow_bias_sd    = 0.012,    # 1s relative flowmeter bias
    rtd_noise_sd    = 0.12,     # K, 1s RTD noise per sensor
)

SECT = np.arange(0, 360, 30).astype(float)
ROWS = ['A', 'B', 'C']

# ---- prescribed degradation (the quantity the estimator infers) -----------
DEPOSIT = {k: 0.0 for r in ROWS for k in [f'{r}-{int(s):03d}' for s in SECT]}
DEPOSIT.update({'A-300': 0.62e-3, 'B-330': 0.25e-3, 'A-330': 0.10e-3})
for s in SECT:
    DEPOSIT[f'D-{int(s):03d}'] = 0.0
ABSTAINED = ['D-240', 'D-270']


def h_foam(theta_deg):
    d = np.abs(((theta_deg - P['theta_door'] + 180.) % 360.) - 180.)
    return P['h_foam_base'] - P['dh_door'] * np.exp(-(d / P['sigma_door']) ** 2)


def T_surface_melt(theta_deg):
    """Radiating temperature of the melt surface: foam top where covered,
    exposed slag/steel where coverage has collapsed."""
    hf = h_foam(theta_deg)
    f = np.exp(-hf / P['h_ref_foam'])
    return P['T_foam_top'] + (P['T_melt_bare'] - P['T_foam_top']) * f


def build_elements():
    """Sub-discretised receiving elements. Returns dict of arrays."""
    els = {'pos': [], 'nrm': [], 'area': [], 'circuit': []}
    n_th, n_z = 6, 5                       # sub-elements per shell panel
    for row in ROWS:
        zc = P['z_rows'][row]
        for s in SECT:
            cid = f'{row}-{int(s):03d}'
            th_edges = np.linspace(s - 15., s + 15., n_th + 1)
            z_edges = np.linspace(zc - P['h_row'] / 2, zc + P['h_row'] / 2, n_z + 1)
            for i in range(n_th):
                th = np.deg2rad(0.5 * (th_edges[i] + th_edges[i + 1]))
                dth = np.deg2rad(th_edges[i + 1] - th_edges[i])
                for j in range(n_z):
                    z = 0.5 * (z_edges[j] + z_edges[j + 1])
                    dz = z_edges[j + 1] - z_edges[j]
                    els['pos'].append([P['R_shell'] * np.cos(th), P['R_shell'] * np.sin(th), z])
                    els['nrm'].append([-np.cos(th), -np.sin(th), 0.])
                    els['area'].append(P['R_shell'] * dth * dz)
                    els['circuit'].append(cid)
    # roof ring, normal pointing down into the furnace
    n_r = 3
    r_edges = np.linspace(P['r_roof_in'], P['r_roof_out'], n_r + 1)
    for s in SECT:
        cid = f'D-{int(s):03d}'
        th_edges = np.linspace(s - 15., s + 15., n_th + 1)
        for i in range(n_th):
            th = np.deg2rad(0.5 * (th_edges[i] + th_edges[i + 1]))
            dth = np.deg2rad(th_edges[i + 1] - th_edges[i])
            for j in range(n_r):
                r = 0.5 * (r_edges[j] + r_edges[j + 1])
                dr = r_edges[j + 1] - r_edges[j]
                els['pos'].append([r * np.cos(th), r * np.sin(th), P['z_roof']])
                els['nrm'].append([0., 0., -1.])
                els['area'].append(r * dth * dr)
                els['circuit'].append(cid)
    for k in ('pos', 'nrm', 'area'):
        els[k] = np.array(els[k], float)
    els['circuit'] = np.array(els['circuit'])
    return els


def irradiation(els):
    """Incident irradiation G [W/m2] from arc columns and the melt surface."""
    pos, nrm = els['pos'], els['nrm']
    G = np.zeros(len(pos))

    # ---- arc columns: isotropic line sources, foam-attenuated -------------
    n_arc = 15
    P_per_pt = P['f_arc_rad'] * P['P_furnace'] / (3 * n_arc)
    z_src = np.linspace(0.02, P['L_arc'], n_arc)
    th_rx = np.degrees(np.arctan2(pos[:, 1], pos[:, 0])) % 360.
    hf_rx = h_foam(th_rx)
    for th_e in P['theta_elec']:
        xe = P['r_electrode'] * np.cos(np.deg2rad(th_e))
        ye = P['r_electrode'] * np.sin(np.deg2rad(th_e))
        for zs in z_src:
            d_vec = np.array([xe, ye, zs]) - pos
            d2 = np.sum(d_vec ** 2, axis=1)
            d = np.sqrt(d2)
            cos_rx = np.sum(d_vec * nrm, axis=1) / d
            cos_rx = np.clip(cos_rx, 0., None)
            tau = np.exp(-P['kappa_foam'] * np.clip(hf_rx - zs, 0., None))
            G += P_per_pt * cos_rx / (4 * np.pi * d2) * tau

    # ---- melt surface: diffuse patches at the local foam height ----------
    n_r, n_t = 8, 36
    r_ed = np.linspace(0.0, P['R_melt'], n_r + 1)
    t_ed = np.linspace(0., 360., n_t + 1)
    for i in range(n_r):
        r = 0.5 * (r_ed[i] + r_ed[i + 1]); dr = r_ed[i + 1] - r_ed[i]
        for j in range(n_t):
            tdeg = 0.5 * (t_ed[j] + t_ed[j + 1]); dt = np.deg2rad(t_ed[j + 1] - t_ed[j])
            dA = r * dr * dt
            E_melt = P['eps_melt'] * SIG * T_surface_melt(tdeg) ** 4
            src = np.array([r * np.cos(np.deg2rad(tdeg)), r * np.sin(np.deg2rad(tdeg)),
                            h_foam(tdeg)])
            d_vec = src - pos
            d2 = np.sum(d_vec ** 2, axis=1); d = np.sqrt(d2)
            cos_rx = np.clip(np.sum(d_vec * nrm, axis=1) / d, 0., None)
            cos_tx = np.clip(-d_vec[:, 2] / d, 0., None)       # melt normal is +z
            G += E_melt * cos_tx * cos_rx / (np.pi * d2) * dA

    G += P['eps_freeboard'] * SIG * P['T_freeboard'] ** 4
    return G * P['f_cavity']


def coolant_side():
    """Face-referred coolant heat transfer coefficient and per-circuit mass flow."""
    n_circ = 3 * P['n_sectors'] + P['n_sectors']          # shell + roof
    m_dot = P['Q_total'] / 3600. * P['rho_w'] / (n_circ + 4)   # + duct circuits
    A_tube = np.pi * P['D_tube'] ** 2 / 4
    v = m_dot / (P['rho_w'] * A_tube)
    Re = P['rho_w'] * v * P['D_tube'] / P['mu_w']
    Nu = 0.023 * Re ** 0.8 * P['Pr_w'] ** 0.4
    h_w = Nu * P['k_w'] / P['D_tube']
    return m_dot, v, Re, h_w, h_w * P['psi_area']


def _bare(G, R_extra, R_chain, T_bulk_K, eps):
    """Radiative balance on an exposed surface backed by R_extra + R_chain."""
    R = R_extra + R_chain
    f = lambda Ts: eps * G - eps * SIG * Ts ** 4 - (Ts - T_bulk_K) / R
    Ts = brentq(f, T_bulk_K + 1e-6, 3000.)
    return (Ts - T_bulk_K) / R, Ts


def solve_element(G, R_chain, T_bulk_K, slag_supply=True, R_out=0.):
    """Conjugate solve. Returns (q_net, T_surface_K, t_skull)."""
    a_s, a_p = P['eps_skull'], P['eps_panel']
    if slag_supply:
        # hypothesis 1: frozen skull with its outer face at the slag solidus
        q_s = a_s * G - a_s * SIG * P['T_solidus'] ** 4
        if q_s > 0:
            t_sk = P['k_skull'] * ((P['T_solidus'] - T_bulk_K) / q_s - R_chain)
            if 0 < t_sk <= P['t_skull_max']:
                return q_s, P['T_solidus'], t_sk
            if t_sk > P['t_skull_max']:
                # skull at the sloughing limit; surface sits below the solidus
                q, Ts = _bare(G, P['t_skull_max'] / P['k_skull'], R_chain, T_bulk_K, a_s)
                return q, Ts, P['t_skull_max']
        else:
            q, Ts = _bare(G, P['t_skull_max'] / P['k_skull'], R_chain, T_bulk_K, a_s)
            return q, Ts, P['t_skull_max']
    # hypothesis 2: bare panel behind any dust/oxide layer
    q, Ts = _bare(G, R_out, R_chain, T_bulk_K, a_p)
    return q, Ts, 0.0


def run(scenario='current'):
    global DEPOSIT
    keep_d, keep_dep, keep_h = P['dh_door'], dict(DEPOSIT), P['h_foam_base']
    if scenario == 'baseline':
        P['dh_door'] = 0.0
        DEPOSIT = {k: 0.0 for k in DEPOSIT}
    elif scenario == 'foaming_collapse':
        P['dh_door'] = 0.0
        P['h_foam_base'] = 0.10
        DEPOSIT = {k: 0.0 for k in DEPOSIT}
    els = build_elements()
    G = irradiation(els)
    m_dot, v, Re, h_w, h_eff = coolant_side()

    circuits = sorted(set(els['circuit']))
    res = {}
    for cid in circuits:
        m = els['circuit'] == cid
        A = els['area'][m].sum()
        G_bar = np.average(G[m], weights=els['area'][m])
        t_dep = DEPOSIT.get(cid, 0.0)
        z_el = np.average(els['pos'][m, 2], weights=els['area'][m])
        th_el = float(cid.split('-')[1])
        slag = (z_el <= P['z_splash']) and (h_foam(th_el) >= P['h_foam_min'])
        R_out = 0. if z_el <= P['z_splash'] else P['t_dust'] / P['k_dust']
        R_chain = 1. / h_eff + t_dep / P['k_deposit'] + P['t_wall'] / P['k_wall']

        # two passes: bulk temperature depends on absorbed heat
        T_bulk = P['T_inlet'] + 5.
        for _ in range(6):
            q, Ts, t_sk = solve_element(G_bar, R_chain, T_bulk + 273.15, slag, R_out)
            dT_w = q * A / (m_dot * P['cp_w'])
            T_bulk = P['T_inlet'] + dT_w / 2.

        dT_film = q / h_eff
        dT_dep = q * t_dep / P['k_deposit']
        dT_wall = q * P['t_wall'] / P['k_wall']
        T_hot = T_bulk + dT_film + dT_dep + dT_wall
        T_inner_wall = T_bulk + dT_film
        res[cid] = dict(area=A, G=G_bar, q=q, T_surface=Ts - 273.15, t_skull=t_sk,
                        t_deposit=t_dep, T_bulk=T_bulk, dT_water=dT_w,
                        dT_film=dT_film, dT_dep=dT_dep, dT_wall=dT_wall,
                        T_hotface=T_hot,
                        boiling_margin=P['T_sat'] - T_inner_wall,
                        Q=q * A,
                        abstained=cid in ABSTAINED)

    # ---- off-gas duct circuits -------------------------------------------
    A_duct = 3.4
    for i, nm in enumerate(['E-000', 'E-001', 'E-002', 'E-003']):
        Tg = P['T_offgas'] - i * 70.
        R_chain = 1. / h_eff + P['t_wall'] / P['k_wall'] + (0.05e-3 * i) / P['k_deposit']
        T_bulk = P['T_inlet'] + 4.
        for _ in range(6):
            f = lambda Ts: (P['h_offgas'] * (Tg - Ts)
                            + P['eps_offgas'] * SIG * (Tg ** 4 - Ts ** 4)
                            - (Ts - (T_bulk + 273.15)) / R_chain)
            Ts = brentq(f, T_bulk + 273.2, Tg - 1.)
            q = (Ts - (T_bulk + 273.15)) / R_chain
            T_bulk = P['T_inlet'] + (q * A_duct / (m_dot * P['cp_w'])) / 2.
        res[nm] = dict(area=A_duct, G=np.nan, q=q, T_surface=Ts - 273.15, t_skull=0.,
                       t_deposit=0.05e-3 * i, T_bulk=T_bulk,
                       dT_water=q * A_duct / (m_dot * P['cp_w']),
                       dT_film=q / h_eff, dT_dep=q * (0.05e-3 * i) / P['k_deposit'],
                       dT_wall=q * P['t_wall'] / P['k_wall'],
                       T_hotface=T_bulk + q * R_chain,
                       boiling_margin=P['T_sat'] - (T_bulk + q / h_eff),
                       Q=q * A_duct, abstained=False)

    # ---- closure ----------------------------------------------------------
    # The reconstructed field and the measured circuit extraction are independent
    # quantities in service. The measured side is simulated here by applying
    # flowmeter bias and RTD noise of the magnitudes declared in PARAMS.
    rng = np.random.default_rng(84732)
    Q_field = sum(r['q'] * r['area'] for r in res.values())
    Q_meas = 0.
    for cid, r in res.items():
        bias = rng.normal(0., P['flow_bias_sd'])
        dT_obs = r['dT_water'] + rng.normal(0., P['rtd_noise_sd'] * np.sqrt(2))
        r['Q_measured'] = m_dot * (1 + bias) * P['cp_w'] * dT_obs
        Q_meas += r['Q_measured']
    closure = 100. * (Q_meas - Q_field) / Q_field

    meta = dict(m_dot=m_dot, v=v, Re=Re, h_w=h_w, h_eff=h_eff,
                Q_field=Q_field, Q_circuits=Q_meas, closure_pct=closure,
                cooling_fraction=100. * Q_field / P['P_furnace'])
    P['dh_door'] = keep_d
    DEPOSIT = keep_dep
    P['h_foam_base'] = keep_h
    return res, meta


if __name__ == '__main__':
    base, mbase = run('baseline')
    res, meta = run('current')
    print(f"baseline (full coverage, no deposit): {mbase['Q_field']/1e6:.2f} MW "
          f"= {mbase['cooling_fraction']:.1f} % of input")
    print(f"current  (door band + A-300 deposit): {meta['Q_field']/1e6:.2f} MW "
          f"= {meta['cooling_fraction']:.1f} % of input")
    print(f"attributable excess loss: {(meta['Q_field']-mbase['Q_field'])/1e6:.2f} MW\n")
    print(f"coolant: mdot {meta['m_dot']:.2f} kg/s  v {meta['v']:.2f} m/s  "
          f"Re {meta['Re']:.3g}  h_w {meta['h_w']:.0f}  h_eff {meta['h_eff']:.0f} W/m2K")
    print(f"field integral {meta['Q_field']/1e6:.2f} MW  "
          f"circuit sum {meta['Q_circuits']/1e6:.2f} MW  "
          f"closure {meta['closure_pct']:+.2f} %  "
          f"cooling {meta['cooling_fraction']:.1f} % of input\n")
    hdr = f"{'circuit':9s}{'G kW/m2':>9s}{'q kW/m2':>9s}{'skull mm':>10s}{'Thot C':>9s}{'boil K':>8s}"
    print(hdr); print('-' * len(hdr))
    for cid in sorted(res):
        r = res[cid]
        print(f"{cid:9s}{r['G']/1e3:9.0f}{r['q']/1e3:9.0f}{r['t_skull']*1e3:10.1f}"
              f"{r['T_hotface']:9.0f}{r['boiling_margin']:8.0f}")
    json.dump({'baseline': {k: {kk: (None if isinstance(vv, float) and np.isnan(vv) else vv)
                                for kk, vv in v.items()} for k, v in base.items()},
               'meta_baseline': mbase,
               'res': {k: {kk: (None if isinstance(vv, float) and np.isnan(vv) else vv)
                           for kk, vv in v.items()} for k, v in res.items()},
               'meta': meta, 'params': {k: v for k, v in P.items()
                                        if not isinstance(v, (dict, list))}},
              open('/home/claude/eaf_results.json', 'w'), indent=1)
