
"""
PKTron Physics — Classical Mechanics
=====================================
Kinematics, dynamics, energy, momentum, projectile motion,
simple harmonic oscillator, orbital mechanics — all float32 safe.
"""
import numpy as np
_R = np.float32

# ── Kinematics ────────────────────────────────────────────────────────────

def displacement(v0, t, a=0.0):
    """s = v0*t + 0.5*a*t^2"""
    return _R(v0)*_R(t) + _R(0.5)*_R(a)*_R(t)**2

def final_velocity(v0, a, t):
    """v = v0 + a*t"""
    return _R(v0) + _R(a)*_R(t)

def velocity_squared(v0, a, s):
    """v^2 = v0^2 + 2*a*s"""
    return _R(v0)**2 + 2*_R(a)*_R(s)

def projectile_range(v0, angle_deg, g=9.81):
    """Horizontal range of projectile."""
    th = np.radians(_R(angle_deg))
    return _R(v0)**2 * np.sin(2*th) / _R(g)

def projectile_max_height(v0, angle_deg, g=9.81):
    """Maximum height of projectile."""
    th = np.radians(_R(angle_deg))
    return (_R(v0)*np.sin(th))**2 / (2*_R(g))

def projectile_trajectory(v0, angle_deg, g=9.81, n_points=100):
    """Return (x, y) trajectory arrays."""
    th    = np.radians(_R(angle_deg))
    T     = 2*_R(v0)*np.sin(th)/_R(g)
    t     = np.linspace(_R(0), T, n_points, dtype=_R)
    x     = _R(v0)*np.cos(th)*t
    y     = _R(v0)*np.sin(th)*t - _R(0.5)*_R(g)*t**2
    return x, np.maximum(y, _R(0))

# ── Dynamics ─────────────────────────────────────────────────────────────

def newton_second_law(mass, acceleration):
    """F = m * a"""
    return _R(mass)*np.asarray(acceleration, dtype=_R)

def gravitational_force(m1, m2, r, G=6.674e-11):
    """F = G*m1*m2/r^2"""
    return _R(G)*_R(m1)*_R(m2)/_R(r)**2

def work_done(force, displacement, angle_deg=0.0):
    """W = F * d * cos(theta)"""
    return _R(force)*_R(displacement)*np.cos(np.radians(_R(angle_deg)))

def kinetic_energy(mass, velocity):
    """KE = 0.5 * m * v^2"""
    return _R(0.5)*_R(mass)*_R(velocity)**2

def potential_energy(mass, height, g=9.81):
    """PE = m * g * h"""
    return _R(mass)*_R(g)*_R(height)

def linear_momentum(mass, velocity):
    """p = m * v (vector-safe)"""
    return _R(mass)*np.asarray(velocity, dtype=_R)

def impulse(force, dt):
    """J = F * dt"""
    return _R(force)*_R(dt)

def elastic_collision(m1, v1, m2, v2):
    """1D elastic collision — returns (v1_f, v2_f)."""
    m1, v1, m2, v2 = _R(m1),_R(v1),_R(m2),_R(v2)
    v1f = ((m1-m2)*v1 + 2*m2*v2)/(m1+m2)
    v2f = ((m2-m1)*v2 + 2*m1*v1)/(m1+m2)
    return float(v1f), float(v2f)

def inelastic_collision(m1, v1, m2, v2):
    """Perfectly inelastic collision — returns v_final."""
    return float((_R(m1)*_R(v1)+_R(m2)*_R(v2))/(_R(m1)+_R(m2)))

# ── Simple Harmonic Oscillator ────────────────────────────────────────────

def sho_position(A, omega, t, phi=0.0):
    """x(t) = A*cos(omega*t + phi)"""
    return _R(A)*np.cos(_R(omega)*np.asarray(t,dtype=_R)+_R(phi))

def sho_velocity(A, omega, t, phi=0.0):
    """v(t) = -A*omega*sin(omega*t + phi)"""
    return -_R(A)*_R(omega)*np.sin(_R(omega)*np.asarray(t,dtype=_R)+_R(phi))

def sho_period(k, mass):
    """T = 2*pi*sqrt(m/k)"""
    return _R(2)*_R(np.pi)*np.sqrt(_R(mass)/_R(k))

def pendulum_period(length, g=9.81):
    """T = 2*pi*sqrt(L/g)  (small-angle)"""
    return _R(2)*_R(np.pi)*np.sqrt(_R(length)/_R(g))

# ── Orbital Mechanics ─────────────────────────────────────────────────────

def orbital_velocity(M, r, G=6.674e-11):
    """v = sqrt(G*M/r)  (circular orbit)"""
    return float(np.sqrt(_R(G)*_R(M)/_R(r)))

def orbital_period(M, r, G=6.674e-11):
    """T = 2*pi*sqrt(r^3 / G*M)  (Kepler III)"""
    return float(2*_R(np.pi)*np.sqrt(_R(r)**3/(_R(G)*_R(M))))

def escape_velocity(M, r, G=6.674e-11):
    """v_esc = sqrt(2*G*M/r)"""
    return float(np.sqrt(2*_R(G)*_R(M)/_R(r)))

def simulate_orbit(M, r0, v0, dt=60.0, n_steps=10000, G=6.674e-11):
    """
    Simulate 2D orbit using Velocity Verlet.
    Returns (x_arr, y_arr) float32.
    """
    x, y   = _R(r0), _R(0)
    vx, vy = _R(0), _R(v0)
    xs, ys = np.zeros(n_steps, dtype=_R), np.zeros(n_steps, dtype=_R)
    G, M, dt = _R(G), _R(M), _R(dt)
    for i in range(n_steps):
        r3 = (x**2+y**2)**_R(1.5)
        ax = -G*M*x/r3; ay = -G*M*y/r3
        vx += ax*dt; vy += ay*dt
        x  += vx*dt; y  += vy*dt
        xs[i]=x; ys[i]=y
    return xs, ys
