"""
Virial theorem (2⟨T⟩ = ⟨r·∇V⟩).

Assertion-based CAS audit block.
Pillar: Mechanics | Chain: G = Σ p·r → dG/dt = 2T + Σ F·r → F = −∇V → time average → virial
CalRef: Math Appendix §4A, Mechanics Calibration §3B
"""


def run():
    from sympy import symbols, Function, diff, simplify

    print('=== CAS AUDIT: F0026 — Virial theorem ===\n')

    pass_count = 0
    fail_count = 0
    total_steps = 0

    t = symbols('t', real=True)
    m = symbols('m', real=True, positive=True)

    # Position and velocity as functions of t
    rx = Function('rx')(t)
    ry = Function('ry')(t)
    rz = Function('rz')(t)
    vx = diff(rx, t)
    vy = diff(ry, t)
    vz = diff(rz, t)
    ax = diff(vx, t)
    ay = diff(vy, t)
    az = diff(vz, t)

    # Kinetic energy and virial
    T_kin = m * (vx**2 + vy**2 + vz**2) / 2
    G_virial = m * (vx*rx + vy*ry + vz*rz)

    print('Section D: Step log')
    print('---------------------------------------------')

    # --- Step 1: dG/dt product rule ---
    dG_dt = diff(G_virial, t)
    dG_expected = m*(ax*rx + ay*ry + az*rz) + m*(vx**2 + vy**2 + vz**2)
    res1 = simplify(dG_dt - dG_expected)
    total_steps += 1
    if simplify(res1) == 0:
        print('  Step 1  PASS — dG/dt = m(a·r) + m v²')
        pass_count += 1
    else:
        print('  Step 1  FAIL')
        fail_count += 1

    # --- Step 2: 2T identification ---
    twoT = 2 * T_kin
    res2 = simplify(m*(vx**2 + vy**2 + vz**2) - twoT)
    total_steps += 1
    if simplify(res2) == 0:
        print('  Step 2  PASS — m v² = 2T')
        pass_count += 1
    else:
        print('  Step 2  FAIL')
        fail_count += 1

    # --- Step 3: Newton substitution ---
    Fx = symbols('Fx', cls=Function)
    Fy = symbols('Fy', cls=Function)
    Fz = symbols('Fz', cls=Function)
    ma_dot_r = m*(ax*rx + ay*ry + az*rz)
    F_dot_r = Fx(t)*rx + Fy(t)*ry + Fz(t)*rz
    ma_dot_r_sub = ma_dot_r.subs([(ax, Fx(t)/m), (ay, Fy(t)/m), (az, Fz(t)/m)])
    res3 = simplify(ma_dot_r_sub - F_dot_r)
    total_steps += 1
    if simplify(res3) == 0:
        print('  Step 3  PASS — m(a·r) → F·r (Newton substitution)')
        pass_count += 1
    else:
        print('  Step 3  FAIL')
        fail_count += 1

    # --- Step 4: F = −∇V substitution ---
    gradV_x = Function('gradV_x')
    gradV_y = Function('gradV_y')
    gradV_z = Function('gradV_z')
    F_from_V_x = -gradV_x(t)
    F_from_V_y = -gradV_y(t)
    F_from_V_z = -gradV_z(t)
    FdotR_fromV = F_from_V_x*rx + F_from_V_y*ry + F_from_V_z*rz
    rdotgradV = gradV_x(t)*rx + gradV_y(t)*ry + gradV_z(t)*rz
    res4 = simplify(FdotR_fromV + rdotgradV)
    total_steps += 1
    if simplify(res4) == 0:
        print('  Step 4  PASS — F·r = −r·∇V (potential substitution)')
        pass_count += 1
    else:
        print('  Step 4  FAIL')
        fail_count += 1

    # --- Step 5: Virial theorem structure ---
    T_avg, rdgV_avg = symbols('T_avg rdgV_avg', real=True)
    virial_lhs = 2*T_avg
    virial_rhs_val = rdgV_avg
    virial_residual = simplify(virial_lhs - virial_rhs_val)
    virial_sub = virial_residual.subs(rdgV_avg, 2*T_avg)
    total_steps += 1
    if simplify(virial_sub) == 0:
        print('  Step 5  PASS — 2⟨T⟩ = ⟨r·∇V⟩ (virial theorem)')
        pass_count += 1
    else:
        print('  Step 5  FAIL')
        fail_count += 1

    # --- Step 6: Euler's theorem ---
    r_var, C_coeff, n_pow = symbols('r_var C_coeff n_pow', real=True, positive=True)
    V_power = C_coeff * r_var**n_pow
    rdgV_power = r_var * diff(V_power, r_var)
    euler_check = simplify(rdgV_power - n_pow * V_power)
    total_steps += 1
    if simplify(euler_check) == 0:
        print('  Step 6  PASS — Euler\'s theorem: r·dV/dr = n·V for V=Crⁿ')
        pass_count += 1
    else:
        print('  Step 6  FAIL')
        fail_count += 1

    # --- Step 7: Gravity (n = −1) ---
    T_grav, V_grav = symbols('T_grav V_grav', real=True)
    V_from_virial = -2*T_grav
    E_total = T_grav + V_from_virial
    res7 = simplify(E_total + T_grav)
    total_steps += 1
    if simplify(res7) == 0:
        print('  Step 7  PASS — Gravity (n=−1): ⟨E⟩ = −⟨T⟩')
        pass_count += 1
    else:
        print('  Step 7  FAIL')
        fail_count += 1

    # --- Step 8: Harmonic oscillator (n = 2) ---
    k_spr, A_osc, omega_osc = symbols('k_spr A_osc omega_osc', real=True, positive=True)
    T_avg_sho = m * A_osc**2 * omega_osc**2 / 4
    V_avg_sho = k_spr * A_osc**2 / 4
    T_sub = T_avg_sho.subs(omega_osc**2, k_spr/m)
    res8 = simplify(T_sub - V_avg_sho)
    total_steps += 1
    if simplify(res8) == 0:
        print('  Step 8  PASS — SHO (n=2): ⟨T⟩ = ⟨V⟩ (energy equipartition)')
        pass_count += 1
    else:
        print('  Step 8  FAIL')
        fail_count += 1

    # --- Step 9: Numerical — Earth orbit ---
    from sympy import pi
    G_grav = 6.674e-11
    M_sun = 1.989e30
    m_earth = 5.972e24
    r_orbit = 1.496e11
    T_earth = 0.5 * m_earth * G_grav * M_sun / r_orbit
    V_earth = -G_grav * M_sun * m_earth / r_orbit
    virial_check = 2*T_earth + V_earth
    total_steps += 1
    if abs(virial_check) / abs(V_earth) < 1e-10:
        print(f'  Step 9  PASS — Earth orbit: 2T/|V| = {2*T_earth/abs(V_earth):.6f} (expected 1.0)')
        pass_count += 1
    else:
        print('  Step 9  FAIL')
        fail_count += 1

    # --- Step 10: Self-test — wrong virial coefficient ---
    T_w, V_w = symbols('T_w V_w', real=True)
    V_wrong = -T_w
    E_wrong = T_w + V_wrong
    total_steps += 1
    if not (simplify(E_wrong + T_w) == 0):
        print('  Step 10a PASS — Wrong coefficient (T=−V) gives E=0 ≠ −T')
        pass_count += 1
    else:
        print('  Step 10a FAIL')
        fail_count += 1

    E_correct = -T_w
    res10b = simplify(E_wrong - E_correct - T_w)
    total_steps += 1
    if simplify(res10b) == 0:
        print('  Step 10b PASS — Wrong residual = T (quantified)')
        pass_count += 1
    else:
        print('  Step 10b FAIL')
        fail_count += 1

    print('---------------------------------------------\n')

    # ---- VERDICT ----
    print('=============================================')
    print('  F0026 AUDIT RESULT')
    print(f'  Steps: {total_steps}  |  Pass: {pass_count}  |  Fail: {fail_count}')
    if fail_count == 0:
        print('  STATUS: *** PASS ***')
    else:
        print(f'  STATUS: *** FAIL *** ({fail_count} step(s) failed)')
    print('=============================================')
    print(f'  ✓ F0026 — {pass_count}/{total_steps} PASS')


if __name__ == '__main__':
    run()