import numpy as np
from numpy.random import normal
import matplotlib.pyplot as plt

# ---------------------------
# Physical constants
# ---------------------------
g = 9.81
rho = 1.2
Cd = 0.47
r = 0.02
A = np.pi * r**2
m = 0.0027

k_drag = 0.5 * rho * Cd * A / m
k_mag = 4e-4    # Magnus coefficient (tuned)

# Table + net
table_length = 2.74
x_net = table_length / 2
h_net = 0.1525
y_thresh = h_net + r   # clearance threshold


# ---------------------------
# 2D flight ODE
# ---------------------------
def f(state, omega):
    x, y, vx, vy = state
    v = np.sqrt(vx**2 + vy**2)

    ax_drag = -k_drag * v * vx
    ay_drag = -k_drag * v * vy

    ax_mag =   k_mag * omega * vy
    ay_mag =  -k_mag * omega * vx

    return np.array([
        vx,
        vy,
        ax_drag + ax_mag,
        -g + ay_drag + ay_mag
    ])


# ---------------------------
# Single trajectory (NO BOUNCE)
# ---------------------------
def simulate_trajectory(S, theta_deg, omega, dt=0.001, t_max=2.0):
    theta = np.deg2rad(theta_deg)

    # initial conditions
    vx = S * np.cos(theta)
    vy = S * np.sin(theta)
    x, y = 0.0, 0.3     # contact height ~30 cm

    prev_x, prev_y = x, y
    cleared_net = False
    hit_net = False

    for _ in np.arange(0, t_max, dt):

        # ------------- NET CHECK -------------
        if (prev_x - x_net) * (x - x_net) <= 0 and (x != prev_x):
            tau = (x_net - prev_x) / (x - prev_x)
            y_cross = prev_y + tau * (y - prev_y)

            if y_cross <= y_thresh:
                return x_net, "net"      # hit net
            else:
                cleared_net = True

        # ------------- GROUND (TABLE) IMPACT -------------
        if y <= 0:
            # Ball lands
            if x < x_net:
                return x, "undershoot"      # lands on own side
            elif x <= table_length:
                return x, "valid"           # lands on opponent side
            else:
                return x, "overshoot"       # lands beyond table

        # ------------- INTEGRATE (RK4) -------------
        state = np.array([x, y, vx, vy])
        k1 = f(state, omega)
        k2 = f(state + 0.5*dt*k1, omega)
        k3 = f(state + 0.5*dt*k2, omega)
        k4 = f(state + dt*k3, omega)
        state += (dt/6)*(k1 + 2*k2 + 2*k3 + k4)

        prev_x, prev_y = x, y
        x, y, vx, vy = state

        # If ball goes far past table without landing → overshoot
        if x > table_length + 0.5:
            return x, "overshoot"

    return x, "unknown"


# ---------------------------
# Monte Carlo
# ---------------------------
def monte_carlo(n=2000):
    landings = []
    outcomes = {"valid": 0, "net": 0, "undershoot": 0, "overshoot": 0}

    for _ in range(n):
        # Random shot parameters
        S = normal(8.0, 0.4)
        ang = normal(12.0, 2.0)
        spin = normal(150, 20)

        x_land, outcome = simulate_trajectory(S, ang, spin)

        landings.append(x_land)
        outcomes[outcome] += 1

    return np.array(landings), outcomes


# ---------------------------
# Run + Plot
# ---------------------------
landings, outcomes = monte_carlo(2000)

print(outcomes)
print("Valid shot probability:", outcomes["valid"] / 2000)

plt.hist(landings, bins=80, density=True)
plt.axvline(x_net, color='r', linestyle='--', label='Net')
plt.axvline(table_length, color='k', linestyle='--', label="End of table")
plt.title("Landing distribution (PDF approx)")
plt.xlabel("Landing x-position (m)")
plt.ylabel("Probability density")
plt.legend()
plt.show()
import gradio as gr

def run_sim():
    landings, outcomes = monte_carlo(2000)

    result_str = f"""
Valid: {outcomes['valid']}
Net: {outcomes['net']}
Undershoot: {outcomes['undershoot']}
Overshoot: {outcomes['overshoot']}
"""

    return result_str

demo = gr.Interface(
    fn=run_sim,
    inputs=[],
    outputs="text",
    title="Table Tennis Monte Carlo Simulator"
)

demo.launch()