"""World-anchored checkerboard with a camera that translates with the root."""
from pathlib import Path
import shutil

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.animation import FFMpegWriter
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import numpy as np

CHAINS = [
    [0, 2, 5, 8, 11], [0, 1, 4, 7, 10],
    [0, 3, 6, 9, 12, 15], [9, 14, 17, 19, 21], [9, 13, 16, 18, 20],
]
COLORS = ["#43d5ca", "#3899ed", "#f3c46c", "#e586ad", "#b7a0f5"]
FLOOR_COLORS = ["#e4d8bd", "#627988"]


def world_to_camera(points, root):
    """Apply the identical horizontal translation to body, floor and landmarks."""
    result = np.array(points, dtype=np.float64, copy=True)
    result[..., 0] -= root[0]
    result[..., 2] -= root[2]
    return result


def checkerboard(root, radius, tile_size=0.5):
    """Tiles retain world coordinates and parity when the camera crosses a cell."""
    x0, z0 = np.floor(np.asarray(root)[[0, 2]] / tile_size).astype(int)
    extent = int(np.ceil(3 * radius / tile_size)) + 2
    tiles, colors = [], []
    for ix in range(x0 - extent, x0 + extent + 1):
        for iz in range(z0 - extent, z0 + extent + 1):
            x, z = ix * tile_size, iz * tile_size
            world = [[x, 0, z], [x + tile_size, 0, z],
                     [x + tile_size, 0, z + tile_size], [x, 0, z + tile_size]]
            camera = world_to_camera(world, root)
            tiles.append(camera[:, [0, 2, 1]])
            colors.append(FLOOR_COLORS[(ix + iz) % 2])
    return tiles, colors


def render_mp4(joints, destination, fps=20):
    joints = np.asarray(joints, dtype=np.float64)
    if joints.ndim != 3 or joints.shape[1:] != (22, 3) or len(joints) < 2:
        raise ValueError("Rendering expects positions with shape (T,22,3).")
    if not np.isfinite(joints).all():
        raise ValueError("Invalid joint positions for rendering.")
    if not np.isfinite(fps) or not 1 <= float(fps) <= 120:
        raise ValueError("Rendering requires a frame rate between 1 and 120 fps.")
    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        try:
            import imageio_ffmpeg
            ffmpeg = imageio_ffmpeg.get_ffmpeg_exe()
        except (ImportError, RuntimeError) as error:
            raise RuntimeError("FFmpeg was not found. Install it before rendering MP4 files.") from error
    matplotlib.rcParams["animation.ffmpeg_path"] = ffmpeg
    destination = Path(destination)
    destination.parent.mkdir(parents=True, exist_ok=True)
    positions = joints.copy()
    positions[:, :, 1] -= positions[:, :, 1].min()
    local = positions.copy()
    local[:, :, 0] -= positions[:, :1, 0]
    local[:, :, 2] -= positions[:, :1, 2]
    height = max(float(positions[:, :, 1].max()), 1.0)
    radius = max(float(np.abs(local[:, :, [0, 2]]).max()) * 1.3, height * 0.85)
    fig = plt.figure(figsize=(6, 6), dpi=96, facecolor="#111c2d")
    ax = fig.add_subplot(111, projection="3d", facecolor="#111c2d", computed_zorder=False)
    ax.set_xlim(-radius, radius)
    ax.set_ylim(-radius, radius)
    ax.set_zlim(-0.02, height * 1.15)
    ax.set_box_aspect((2 * radius, 2 * radius, height * 1.15))
    ax.view_init(elev=20, azim=-65)
    ax.set_axis_off()
    fig.subplots_adjust(0, 0.05, 1, 1)
    floor = Poly3DCollection([], edgecolors="#344353", linewidths=0.3, zorder=0)
    ax.add_collection3d(floor)
    trail, = ax.plot([], [], [], color="#eeae56", lw=1.5, alpha=0.8, zorder=2)
    start, = ax.plot([], [], [], marker="o", markersize=6, color="#f07d55", zorder=2)
    lines = [ax.plot([], [], [], color=color, lw=3.2, solid_capstyle="round", zorder=3)[0]
             for color in COLORS]
    fig.text(0.5, 0.035, "World-fixed floor · 0.5 m tiles · camera follows the person",
             ha="center", color="#bdc9d8", fontsize=8,
             bbox={"facecolor": "#111c2d", "edgecolor": "none", "pad": 4})
    distance = fig.text(0.5, 0.085, "", ha="center", color="#e4d8bd", fontsize=9,
                        bbox={"facecolor": "#111c2d", "edgecolor": "none", "pad": 4})
    roots = positions[:, 0].copy()
    roots[:, 1] = 0.006
    writer = FFMpegWriter(fps=fps, codec="libx264", bitrate=1800,
                          extra_args=["-pix_fmt", "yuv420p", "-movflags", "+faststart", "-threads", "2"])
    try:
        with writer.saving(fig, str(destination), dpi=96):
            for index, frame in enumerate(positions):
                root = frame[0]
                tiles, colors = checkerboard(root, radius)
                floor.set_verts(tiles)
                floor.set_facecolor(colors)
                for line, chain in zip(lines, CHAINS):
                    points = local[index, chain]
                    line.set_data(points[:, 0], points[:, 2])
                    line.set_3d_properties(points[:, 1])
                path = world_to_camera(roots[:index + 1], root)
                outside = (np.abs(path[:, 0]) > radius) | (np.abs(path[:, 2]) > radius)
                path[outside] = np.nan
                trail.set_data(path[:, 0], path[:, 2])
                trail.set_3d_properties(path[:, 1])
                origin = world_to_camera(roots[0], root)
                visible_origin = abs(origin[0]) <= radius and abs(origin[2]) <= radius
                start.set_data([origin[0]] if visible_origin else [], [origin[2]] if visible_origin else [])
                start.set_3d_properties([origin[1]] if visible_origin else [])
                traveled = np.linalg.norm(frame[0, [0, 2]] - positions[0, 0, [0, 2]])
                distance.set_text(f"{index / fps:.2f} s   ·   displacement from start: {traveled:.2f} m")
                writer.grab_frame(facecolor=fig.get_facecolor())
    finally:
        plt.close(fig)
    if not destination.is_file() or destination.stat().st_size == 0:
        raise RuntimeError("FFmpeg did not produce the video.")
    return destination