import gradio as gr
import numpy as np
import plotly.graph_objects as go
import time

G = 6.674e-11
c = 3e8
M_sun = 1.989e30

def black_hole_advanced(mass):
    # internal time (auto animation feel)
    t = time.time() % 100

    M = mass * M_sun
    rs = (2 * G * M) / (c**2) / 1000  # km

    fig = go.Figure()

    # ===============================
    # EVENT HORIZON (PURE BLACK)
    # ===============================
    u = np.linspace(0, 2*np.pi, 60)
    v = np.linspace(0, np.pi, 60)

    x = rs * np.outer(np.cos(u), np.sin(v))
    y = rs * np.outer(np.sin(u), np.sin(v))
    z = rs * np.outer(np.ones_like(u), np.cos(v))

    fig.add_surface(
        x=x, y=y, z=z,
        colorscale=[[0, "black"], [1, "black"]],
        showscale=False,
        opacity=1.0
    )

    # ===============================
    # ACCRETION DISK
    # ===============================
    r_disk = np.linspace(1.7*rs, 5*rs, 80)
    theta = np.linspace(0, 2*np.pi, 160)

    R, T = np.meshgrid(r_disk, theta)

    Xd = R * np.cos(T)
    Yd = R * np.sin(T)
    Zd = 0.05 * rs * np.sin(T * 3)

    disk_intensity = np.exp(-R / (3*rs))

    fig.add_surface(
        x=Xd, y=Yd, z=Zd,
        surfacecolor=disk_intensity,
        colorscale="Inferno",
        opacity=0.9,
        showscale=False
    )

    # ===============================
    # AUTO INFALLING STREAMS (NO UI CONTROL)
    # ===============================
    n_particles = 60
    angles = np.linspace(0, 2*np.pi, n_particles)
    radii = np.linspace(2.5*rs, 6*rs, n_particles)

    radii = radii - (t * 0.03 * rs)
    radii = np.mod(radii - 1.1*rs, 5*rs) + 1.1*rs

    xp = radii * np.cos(angles + t * 0.25)
    yp = radii * np.sin(angles + t * 0.25)
    zp = np.random.normal(0, 0.035*rs, n_particles)

    fig.add_trace(go.Scatter3d(
        x=xp,
        y=yp,
        z=zp,
        mode="markers",
        marker=dict(
            size=2,
            color="orange",
            opacity=0.5
        ),
        showlegend=False
    ))

    # ===============================
    # LIGHT BENDING
    # ===============================
    # ===============================
    # CORRECTED GRAVITATIONAL LENSING
    # ===============================
    y_vals = np.linspace(-4*rs, 4*rs, 9)

    for y0 in y_vals:
        x_ray = np.linspace(-8*rs, 8*rs, 500)
        y_ray, z_ray = [], []

        for x0 in x_ray:
            r = np.sqrt(x0**2 + y0**2)

            if r <= 1.02 * rs:
                y_ray.append(np.nan)
                z_ray.append(np.nan)
                continue

            # Strong relativistic-style bending
            b = abs(y0) + 0.2 * rs              # impact parameter
            bend_strength = (rs / r) * (rs / b)

            y_deflect = bend_strength * y0 * 0.9
            z_deflect = bend_strength * rs * 0.7

            y_ray.append(y0 - y_deflect)
            z_ray.append(z_deflect)

        fig.add_trace(go.Scatter3d(
            x=x_ray,
            y=y_ray,
            z=z_ray,
            mode="lines",
            line=dict(width=2),
            opacity=0.5,
            showlegend=False
        ))

    # ===============================
    # LAYOUT
    # ===============================
    fig.update_layout(
        title="Advanced 3D Black Hole Simulation",
        scene=dict(
            xaxis=dict(visible=False),
            yaxis=dict(visible=False),
            zaxis=dict(visible=False),
            aspectmode="data",
            camera=dict(eye=dict(x=1.6, y=1.6, z=1.1))
        ),
        margin=dict(l=0, r=0, b=0, t=40),
        paper_bgcolor="black",
        plot_bgcolor="black"
    )

    return fig


# ===============================
# CLEAN INTERFACE (ONE CONTROL ONLY)
# ===============================
gr.Interface(
    fn=black_hole_advanced,
    inputs=gr.Slider(5, 50, value=15, label="Black Hole Mass (Solar Masses)"),
    outputs=gr.Plot(label="3D Black Hole"),
    title="Advanced Black Hole Simulator",
    description="Pure black event horizon with automatic infalling matter and accretion disk."
).launch()