"""
Dental Bridge 3D Animation - Patient Education Video
Renders frames using matplotlib 3D, then compiles to MP4 via ffmpeg.
"""

import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import os, subprocess

OUT_DIR = "/home/daytona/workspace/dental-bridge-video/frames"
os.makedirs(OUT_DIR, exist_ok=True)

FPS = 24
W, H = 1280, 720
DPI = 96

# ─────────────────────────────────────────────
# Geometry helpers
# ─────────────────────────────────────────────

def tooth_mesh(cx, cy, cz, w=0.7, d=0.55, h=1.1, taper=0.15, color='ivory', alpha=1.0):
    """Return a list of Poly3DCollection verts for a simple tooth block."""
    hw, hd = w/2, d/2
    # bottom face (wider)
    bx = [cx-hw, cx+hw, cx+hw, cx-hw]
    by = [cy-hd, cy-hd, cy+hd, cy+hd]
    bz = [cz]*4
    # top face (tapered)
    tx = [cx-hw+taper, cx+hw-taper, cx+hw-taper, cx-hw+taper]
    ty = [cy-hd+taper, cy-hd+taper, cy+hd-taper, cy+hd-taper]
    tz = [cz+h]*4

    verts = []
    # bottom
    verts.append(list(zip(bx, by, bz)))
    # top
    verts.append(list(zip(tx, ty, tz)))
    # 4 sides
    for i in range(4):
        ni = (i+1) % 4
        verts.append([
            (bx[i], by[i], bz[i]),
            (bx[ni], by[ni], bz[ni]),
            (tx[ni], ty[ni], tz[ni]),
            (tx[i], ty[i], tz[i]),
        ])
    return verts

def gum_pad(cx, cy, cz=0, w=0.9, d=0.7, h=0.25, color='#e8a0a0'):
    """A flat rounded gum block under a tooth."""
    hw, hd = w/2, d/2
    bx = [cx-hw, cx+hw, cx+hw, cx-hw]
    by = [cy-hd, cy-hd, cy+hd, cy+hd]
    bz = [cz]*4
    tx = [cx-hw, cx+hw, cx+hw, cx-hw]
    ty = [cy-hd, cy-hd, cy+hd, cy+hd]
    tz = [cz+h]*4
    verts = []
    verts.append(list(zip(bx, by, bz)))
    verts.append(list(zip(tx, ty, tz)))
    for i in range(4):
        ni = (i+1) % 4
        verts.append([(bx[i],by[i],bz[i]),(bx[ni],by[ni],bz[ni]),
                      (tx[ni],ty[ni],tz[ni]),(tx[i],ty[i],tz[i])])
    return verts

def draw_verts(ax, verts, facecolor, edgecolor='#888888', alpha=1.0, linewidth=0.5):
    poly = Poly3DCollection(verts, alpha=alpha, linewidth=linewidth)
    poly.set_facecolor(facecolor)
    poly.set_edgecolor(edgecolor)
    ax.add_collection3d(poly)
    return poly

def setup_ax(fig, elev=20, azim=-60, title='', subtitle=''):
    ax = fig.add_subplot(111, projection='3d')
    ax.set_facecolor('#0a1628')
    fig.patch.set_facecolor('#0a1628')
    ax.view_init(elev=elev, azim=azim)
    ax.set_xlim(-2.5, 2.5)
    ax.set_ylim(-1.5, 1.5)
    ax.set_zlim(-0.3, 2.2)
    ax.set_axis_off()
    # Title
    fig.text(0.5, 0.93, title, ha='center', va='top',
             fontsize=22, fontweight='bold', color='white',
             fontfamily='DejaVu Sans')
    if subtitle:
        fig.text(0.5, 0.87, subtitle, ha='center', va='top',
                 fontsize=14, color='#aaddff',
                 fontfamily='DejaVu Sans')
    return ax

def add_label_3d(ax, x, y, z, text, color='#ffdd88', fontsize=10):
    ax.text(x, y, z, text, color=color, fontsize=fontsize,
            fontweight='bold', ha='center', va='bottom',
            fontfamily='DejaVu Sans')

# Tooth x-positions for 5 teeth: left abutment, gap-left, GAP, gap-right, right abutment
TOOTH_XS = [-2.0, -0.9, 0.0, 0.9, 2.0]  # gap at index 2
TOOTH_COLOR = '#f5f0e8'
CROWN_COLOR = '#e8d5b0'   # slightly yellower for crowns
PONTIC_COLOR = '#d4eaf5'  # porcelain blue-white for pontic
GUM_COLOR = '#e8a0a0'
PREP_COLOR = '#c8a870'    # prepared/trimmed abutment

# ─────────────────────────────────────────────
# Scene renderers
# ─────────────────────────────────────────────

def render_scene1(frame, n_frames, azim_start=-70):
    """Scene 1: Healthy teeth, rotating intro."""
    azim = azim_start + (frame / n_frames) * 40
    fig = plt.figure(figsize=(W/DPI, H/DPI), dpi=DPI)
    ax = setup_ax(fig, elev=18, azim=azim,
                  title='Dental Bridge',
                  subtitle='A fixed solution to replace a missing tooth')
    # Draw 5 healthy teeth with gum
    for i, tx in enumerate(TOOTH_XS):
        gv = gum_pad(tx, 0, cz=0)
        draw_verts(ax, gv, GUM_COLOR, alpha=1.0)
        tv = tooth_mesh(tx, 0, cz=0.25)
        draw_verts(ax, tv, TOOTH_COLOR, alpha=1.0)
    # Ground / jawbone hint
    xs = np.linspace(-2.8, 2.8, 40)
    ys = np.linspace(-1.0, 1.0, 10)
    XX, YY = np.meshgrid(xs, ys)
    ZZ = np.zeros_like(XX) - 0.05
    ax.plot_surface(XX, YY, ZZ, color='#b07050', alpha=0.4, linewidth=0, zorder=0)
    fig.savefig(f"{OUT_DIR}/frame_{frame:05d}.png", dpi=DPI, bbox_inches='tight',
                facecolor='#0a1628')
    plt.close(fig)

def render_scene2(frame, n_frames, t):
    """Scene 2: One tooth fades/disappears — missing tooth."""
    azim = -30 + t * 10
    fig = plt.figure(figsize=(W/DPI, H/DPI), dpi=DPI)
    ax = setup_ax(fig, elev=22, azim=azim,
                  title='Step 1: Missing Tooth',
                  subtitle='A gap in your smile affects chewing, speech, and adjacent teeth')
    alpha_missing = max(0.0, 1.0 - t * 2.5)
    for i, tx in enumerate(TOOTH_XS):
        alp = alpha_missing if i == 2 else 1.0
        gv = gum_pad(tx, 0, cz=0)
        draw_verts(ax, gv, GUM_COLOR, alpha=alp)
        tv = tooth_mesh(tx, 0, cz=0.25)
        draw_verts(ax, tv, TOOTH_COLOR, alpha=alp)
    # Gap highlight after tooth gone
    if alpha_missing < 0.3:
        ax.plot([TOOTH_XS[2]-0.35, TOOTH_XS[2]+0.35],
                [0, 0], [0.25, 0.25], '--', color='#ff6655', linewidth=2, alpha=0.9)
        add_label_3d(ax, TOOTH_XS[2], 0, 1.5, 'Gap (missing tooth)', color='#ff9977', fontsize=11)
    # bone
    xs = np.linspace(-2.8, 2.8, 40)
    ys = np.linspace(-1.0, 1.0, 10)
    XX, YY = np.meshgrid(xs, ys)
    ZZ = np.zeros_like(XX) - 0.05
    ax.plot_surface(XX, YY, ZZ, color='#b07050', alpha=0.4, linewidth=0)
    fig.savefig(f"{OUT_DIR}/frame_{frame:05d}.png", dpi=DPI, bbox_inches='tight',
                facecolor='#0a1628')
    plt.close(fig)

def render_scene3(frame, n_frames, t):
    """Scene 3: Abutment teeth are prepared (trimmed)."""
    azim = -20 + t * 15
    fig = plt.figure(figsize=(W/DPI, H/DPI), dpi=DPI)
    ax = setup_ax(fig, elev=22, azim=azim,
                  title='Step 2: Preparing the Abutment Teeth',
                  subtitle='Adjacent teeth are gently shaped to support the bridge crowns')
    shrink = min(t * 1.5, 1.0)  # how much they are trimmed
    for i, tx in enumerate(TOOTH_XS):
        if i == 2:
            # gap — draw exposed gum
            gv = gum_pad(tx, 0, cz=0, color='#d08080')
            draw_verts(ax, gv, '#d08080', alpha=0.9)
            continue
        gv = gum_pad(tx, 0, cz=0)
        draw_verts(ax, gv, GUM_COLOR, alpha=1.0)
        if i in (1, 3):
            # Prepared abutments — smaller, chamfered
            taper_v = 0.05 + shrink * 0.12
            w_v = 0.7 - shrink * 0.12
            tv = tooth_mesh(tx, 0, cz=0.25, w=w_v, h=0.9, taper=taper_v)
            draw_verts(ax, tv, PREP_COLOR, alpha=1.0)
            if t > 0.5:
                add_label_3d(ax, tx, 0, 1.55, 'Abutment\n(prepared)', color='#ffdd88', fontsize=9)
        else:
            tv = tooth_mesh(tx, 0, cz=0.25)
            draw_verts(ax, tv, TOOTH_COLOR, alpha=1.0)
    xs = np.linspace(-2.8, 2.8, 40)
    ys = np.linspace(-1.0, 1.0, 10)
    XX, YY = np.meshgrid(xs, ys)
    ZZ = np.zeros_like(XX) - 0.05
    ax.plot_surface(XX, YY, ZZ, color='#b07050', alpha=0.4, linewidth=0)
    fig.savefig(f"{OUT_DIR}/frame_{frame:05d}.png", dpi=DPI, bbox_inches='tight',
                facecolor='#0a1628')
    plt.close(fig)

def render_scene4(frame, n_frames, t):
    """Scene 4: Bridge components shown floating above gap."""
    azim = -5 + t * 20
    elev = 25 + t * 5
    fig = plt.figure(figsize=(W/DPI, H/DPI), dpi=DPI)
    ax = setup_ax(fig, elev=elev, azim=azim,
                  title='Step 3: The Bridge Components',
                  subtitle='Two crowns anchor the bridge; the pontic fills the gap')
    # Teeth below (prepared abutments + gap)
    for i, tx in enumerate(TOOTH_XS):
        if i == 2:
            gv = gum_pad(tx, 0, cz=0, color='#d08080')
            draw_verts(ax, gv, '#d08080', alpha=0.9)
            continue
        gv = gum_pad(tx, 0, cz=0)
        draw_verts(ax, gv, GUM_COLOR, alpha=1.0)
        if i in (1, 3):
            tv = tooth_mesh(tx, 0, cz=0.25, w=0.58, h=0.9, taper=0.17)
            draw_verts(ax, tv, PREP_COLOR, alpha=1.0)
        else:
            tv = tooth_mesh(tx, 0, cz=0.25)
            draw_verts(ax, tv, TOOTH_COLOR, alpha=1.0)
    # Bridge floating above — lower as t progresses
    float_z = 0.6 + (1.0 - t) * 1.1
    bridge_xs = [TOOTH_XS[1], TOOTH_XS[2], TOOTH_XS[3]]
    bridge_colors = [CROWN_COLOR, PONTIC_COLOR, CROWN_COLOR]
    bridge_labels = ['Crown', 'Pontic\n(artificial tooth)', 'Crown']
    for bx, bc, bl in zip(bridge_xs, bridge_colors, bridge_labels):
        bv = tooth_mesh(bx, 0, cz=float_z, w=0.68, h=1.05, taper=0.08)
        draw_verts(ax, bv, bc, alpha=0.92)
    # Connector bars between bridge pieces
    for xi in range(2):
        x1, x2 = bridge_xs[xi], bridge_xs[xi+1]
        ax.plot([x1+0.34, x2-0.34], [0, 0],
                [float_z+0.4, float_z+0.4], '-', color='#a09070', linewidth=4, alpha=0.9)
    # Labels
    if t > 0.3:
        label_colors = ['#ffe090', '#90e0ff', '#ffe090']
        for bx, bl, lc in zip(bridge_xs, bridge_labels, label_colors):
            add_label_3d(ax, bx, 0, float_z + 1.25, bl, color=lc, fontsize=9)
    xs = np.linspace(-2.8, 2.8, 40)
    ys = np.linspace(-1.0, 1.0, 10)
    XX, YY = np.meshgrid(xs, ys)
    ZZ = np.zeros_like(XX) - 0.05
    ax.plot_surface(XX, YY, ZZ, color='#b07050', alpha=0.4, linewidth=0)
    fig.savefig(f"{OUT_DIR}/frame_{frame:05d}.png", dpi=DPI, bbox_inches='tight',
                facecolor='#0a1628')
    plt.close(fig)

def render_scene5(frame, n_frames, t):
    """Scene 5: Bridge lowering into place."""
    azim = 15 + t * 15
    fig = plt.figure(figsize=(W/DPI, H/DPI), dpi=DPI)
    ax = setup_ax(fig, elev=25, azim=azim,
                  title='Step 4: Placing the Bridge',
                  subtitle='The bridge is fitted and cemented onto the abutment teeth')
    for i, tx in enumerate(TOOTH_XS):
        if i == 2:
            gv = gum_pad(tx, 0, cz=0, color='#d08080')
            draw_verts(ax, gv, '#d08080', alpha=0.9)
            continue
        gv = gum_pad(tx, 0, cz=0)
        draw_verts(ax, gv, GUM_COLOR, alpha=1.0)
        if i in (1, 3):
            tv = tooth_mesh(tx, 0, cz=0.25, w=0.58, h=0.9, taper=0.17)
            draw_verts(ax, tv, PREP_COLOR, alpha=1.0)
        else:
            tv = tooth_mesh(tx, 0, cz=0.25)
            draw_verts(ax, tv, TOOTH_COLOR, alpha=1.0)
    # Bridge descending
    bridge_z = 0.25 + (1.0 - t) * 1.3
    bridge_xs = [TOOTH_XS[1], TOOTH_XS[2], TOOTH_XS[3]]
    bridge_colors = [CROWN_COLOR, PONTIC_COLOR, CROWN_COLOR]
    for bx, bc in zip(bridge_xs, bridge_colors):
        bv = tooth_mesh(bx, 0, cz=bridge_z, w=0.68, h=1.05, taper=0.08)
        draw_verts(ax, bv, bc, alpha=0.95)
    for xi in range(2):
        x1, x2 = bridge_xs[xi], bridge_xs[xi+1]
        ax.plot([x1+0.34, x2-0.34], [0, 0],
                [bridge_z+0.4, bridge_z+0.4], '-', color='#a09070', linewidth=4, alpha=0.9)
    # Descend arrows
    if t < 0.85:
        for bx in bridge_xs:
            ax.quiver(bx, 0, bridge_z + 1.15, 0, 0, -0.3,
                      color='#88ddff', arrow_length_ratio=0.4, linewidth=1.5, alpha=0.8)
    xs = np.linspace(-2.8, 2.8, 40)
    ys = np.linspace(-1.0, 1.0, 10)
    XX, YY = np.meshgrid(xs, ys)
    ZZ = np.zeros_like(XX) - 0.05
    ax.plot_surface(XX, YY, ZZ, color='#b07050', alpha=0.4, linewidth=0)
    fig.savefig(f"{OUT_DIR}/frame_{frame:05d}.png", dpi=DPI, bbox_inches='tight',
                facecolor='#0a1628')
    plt.close(fig)

def render_scene6(frame, n_frames, t):
    """Scene 6: Final result — bridge seated, rotating."""
    azim = 30 + t * 50
    fig = plt.figure(figsize=(W/DPI, H/DPI), dpi=DPI)
    ax = setup_ax(fig, elev=20, azim=azim,
                  title='Result: Your Dental Bridge',
                  subtitle='Restored function, appearance, and confidence!')
    for i, tx in enumerate(TOOTH_XS):
        if i == 2:
            gv = gum_pad(tx, 0, cz=0, color='#e8a0a0')
            draw_verts(ax, gv, '#e8a0a0', alpha=1.0)
            bv = tooth_mesh(tx, 0, cz=0.25, w=0.68, h=1.05, taper=0.08)
            draw_verts(ax, bv, PONTIC_COLOR, alpha=1.0)
            continue
        gv = gum_pad(tx, 0, cz=0)
        draw_verts(ax, gv, GUM_COLOR, alpha=1.0)
        if i in (1, 3):
            bv = tooth_mesh(tx, 0, cz=0.25, w=0.68, h=1.05, taper=0.08)
            draw_verts(ax, bv, CROWN_COLOR, alpha=1.0)
        else:
            tv = tooth_mesh(tx, 0, cz=0.25)
            draw_verts(ax, tv, TOOTH_COLOR, alpha=1.0)
    # Connectors
    for xi in range(2):
        x1 = TOOTH_XS[xi+1]
        x2 = TOOTH_XS[xi+2]
        ax.plot([x1+0.34, x2-0.34], [0, 0],
                [0.25+0.4, 0.25+0.4], '-', color='#a09070', linewidth=4, alpha=0.9)
    # Benefits text
    if t > 0.5:
        benefits = ['Restores chewing', 'Prevents tooth drift', 'Natural appearance']
        colors = ['#90ff90', '#90ffff', '#ffff90']
        for j, (b, c) in enumerate(zip(benefits, colors)):
            fig.text(0.72, 0.70 - j*0.08, f'✓  {b}', color=c, fontsize=12,
                     fontweight='bold', fontfamily='DejaVu Sans')
    xs = np.linspace(-2.8, 2.8, 40)
    ys = np.linspace(-1.0, 1.0, 10)
    XX, YY = np.meshgrid(xs, ys)
    ZZ = np.zeros_like(XX) - 0.05
    ax.plot_surface(XX, YY, ZZ, color='#b07050', alpha=0.4, linewidth=0)
    fig.savefig(f"{OUT_DIR}/frame_{frame:05d}.png", dpi=DPI, bbox_inches='tight',
                facecolor='#0a1628')
    plt.close(fig)

# ─────────────────────────────────────────────
# Build frame list
# ─────────────────────────────────────────────

scenes = [
    # (renderer, n_frames)
    (render_scene1, 72),   # ~3s  intro rotate
    (render_scene2, 60),   # ~2.5s missing tooth
    (render_scene3, 60),   # ~2.5s prep
    (render_scene4, 72),   # ~3s   components
    (render_scene5, 72),   # ~3s   placement
    (render_scene6, 96),   # ~4s   final
]

frame_idx = 0
for scene_fn, n in scenes:
    print(f"Rendering {scene_fn.__name__} ({n} frames)...")
    for i in range(n):
        t = i / max(n - 1, 1)
        if scene_fn == render_scene1:
            scene_fn(frame_idx, n, azim_start=-70)
        else:
            scene_fn(frame_idx, n, t)
        frame_idx += 1
    print(f"  done. total frames so far: {frame_idx}")

print(f"\nAll {frame_idx} frames rendered.")