"""
Real-world 3D container geometries for sphere packing.

Each venue is modeled with its actual cross-section shape:
  - Box: elevator, porta-potty, coffin, hot tub (square Jacuzzi)
  - Circular fuselage: Boeing 747 (cylinder with floor cut)
  - Rounded-roof box: MINI Cooper, Tokyo subway
  - Domed box: K6 phone booth

Dimensions sourced from manufacturer specs and engineering references.
"""

import numpy as np
import pyvista as pv


# ============================================================
# Container definitions with REAL dimensions and shapes
# ============================================================

VENUES = {
    "Elevator (Otis Gen2)": {
        "type": "box",
        "dims": {
            # Otis Gen2 passenger elevator, 3500 lb (1588 kg) model
            # Interior: 6'8-5/16" x 5'5-5/8" x 8'0" (2.04m x 1.67m x 2.44m)
            # EN81-20: 3.41 m^2 floor area -> ~1600kg / 21 persons
            # Source: Otis Elevator Planning and Selection Guide (US market)
            "Lx": 2.04,  # 6'8-5/16" width
            "Ly": 1.67,  # 5'5-5/8" depth
            "Lz": 2.44,  # 8'0" height
        },
        "description": "Otis Gen2 passenger elevator, 3500 lb (1588 kg) capacity",
        "color": "#AAAAAA",
        "rated_capacity": 21,
        "source": "Otis Elevator Planning and Selection Guide",
    },

    "Boeing 747-400": {
        "type": "fuselage",
        "dims": {
            # Boeing 747-400 main deck
            # Fuselage interior diameter: 6.13m (20'1", single-circle approx)
            # Main deck cabin length: 57.0m (bulkhead to bulkhead)
            # Floor height: cargo hold ~1.12m clear + 0.15m floor beams
            #   + lower fuselage structure ~1.43m = 2.70m above circle bottom
            # Floor sits 0.365m below fuselage centerline
            # Note: 747 constant-section is approximately circular (not double-bubble)
            # Source: Boeing 747-400 Airport Planning Document, AeroCorner, b737.org.uk
            "length": 57.0,
            "fuselage_radius": 3.065,  # 6.13m diameter / 2
            "floor_height": 2.70,     # main deck floor above fuselage bottom
        },
        "description": "Boeing 747-400 main deck (circular fuselage, single-class config)",
        "color": "#DDDDDD",
        "rated_capacity": 660,
        "source": "Boeing 747-400 Airport Planning Document",
    },

    "K6 Phone Booth": {
        "type": "domed_box",
        "dims": {
            # British K6 red telephone kiosk (Giles Gilbert Scott, 1935)
            # Exterior base: 3'x3' (0.91m); walls ~5-6cm thick
            # Interior: ~0.80m x 0.80m x 2.0m + domed roof ~0.30m
            # Source: Design Museum, GPO specification, the-telephone-box.co.uk
            "Lx": 0.80,
            "Ly": 0.80,
            "Lz": 2.0,
            "dome_height": 0.30,
        },
        "description": "K6 red telephone kiosk with domed roof",
        "color": "#CC2222",
        "rated_capacity": 1,
        "world_record": 25,   # South Africa, 1959 (phone booth stuffing craze)
        "source": "Design Museum / GPO Specification",
    },

    "Hot Tub (Jacuzzi J-345)": {
        "type": "box",
        "dims": {
            # Jacuzzi J-345 6-person hot tub (square)
            # Exterior: 84" x 84" x 36" (2.13m x 2.13m x 0.91m)
            # Interior opening: ~60" x 60" (1.52m x 1.52m) per spec sheet drawing
            # Max depth (seat H): 33.5" (0.85m)
            # Source: Jacuzzi J-345 spec sheet (Aqua Quip PDF)
            "Lx": 1.52,   # 60" interior opening
            "Ly": 1.52,   # 60" interior opening
            "Lz": 0.85,   # 33.5" max depth
        },
        "description": "Jacuzzi J-345 6-person hot tub, interior shell opening",
        "color": "#4488BB",
        "rated_capacity": 6,
        "source": "Jacuzzi J-345 Spec Sheet (aquaquip.com)",
    },

    "Porta-Potty": {
        "type": "box",
        "dims": {
            # Standard portable restroom (PolyJohn PJN3)
            # Exterior: 43.5"W x 47"D x 91"H
            # Interior: 41"W x 41"D x 82"H (square interior)
            # Source: PolyJohn PJN3 product specifications
            "Lx": 1.04,   # 41"
            "Ly": 1.04,   # 41"
            "Lz": 2.08,   # 82"
        },
        "description": "PolyJohn PJN3 standard portable restroom interior",
        "color": "#2E8B57",
        "rated_capacity": 1,
        "source": "PolyJohn PJN3 Product Specifications",
    },

    "Standard Coffin": {
        "type": "box",
        "dims": {
            # Standard American metal casket interior dimensions
            # Interior: 78" x 23" x 15" (1.98m x 0.58m x 0.38m)
            # Exterior: 84" x 28" x 23"
            # Source: Overnight Caskets, Discount Caskets dimension guides
            "Lx": 1.98,   # 78" length
            "Ly": 0.58,   # 23" width
            "Lz": 0.38,   # 15" interior depth
        },
        "description": "Standard American metal casket, interior dimensions",
        "color": "#5C4033",
        "rated_capacity": 1,
        "source": "Overnight Caskets / Discount Caskets dimension guides",
    },

    "King's Chamber (Great Pyramid)": {
        "type": "box",
        "dims": {
            # King's Chamber, Great Pyramid of Giza
            # Interior: 10.47m x 5.23m x 5.82m (20 x 10 x 11.18 royal cubits)
            # Granite-lined rectangular chamber
            # Source: Petrie (1883), "The Pyramids and Temples of Gizeh"
            "Lx": 10.47,
            "Ly": 5.23,
            "Lz": 5.82,
        },
        "description": "King's Chamber, Great Pyramid of Giza (granite-lined)",
        "color": "#C4A35A",
        "rated_capacity": 1,  # 1 deceased pharaoh
        "source": "Petrie (1883), The Pyramids and Temples of Gizeh",
    },

    "Volkswagen Beetle (Classic)": {
        "type": "rounded_box",
        "dims": {
            # Classic VW Beetle (Type 1, 1960s) passenger compartment, seats removed
            # Interior: ~1.70m long x 1.26m wide x 1.10m tall
            # Exterior shoulder room: 49.6" (1.26m)
            # Curved roof, ~10cm sag at edges
            # Source: VW Beetle owner's manual, classicandsportscar.com
            "Lx": 1.70,
            "Ly": 1.26,
            "Lz": 1.10,
            "roof_sag": 0.10,
        },
        "description": "VW Beetle (Type 1) passenger cabin, seats removed",
        "color": "#87CEEB",
        "rated_capacity": 4,
        "world_record": 20,   # Guinness World Record (New Beetle, 2010)
        "source": "VW Beetle Specifications / Guinness World Records",
    },

    "ISS Destiny Lab": {
        "type": "cylinder",
        "dims": {
            # ISS Destiny Laboratory Module (US Lab)
            # Interior: 8.53m (28 ft) long, 4.27m (14 ft) diameter
            # No floor — microgravity, full cylinder usable
            # Source: NASA ISS Reference Guide
            "radius": 2.135,   # 4.27m / 2
            "height": 8.53,
        },
        "description": "ISS Destiny Laboratory Module (microgravity, no floor)",
        "color": "#B0C4DE",
        "rated_capacity": 3,   # typical crew working in the module
        "source": "NASA ISS Reference Guide",
    },
}


# ============================================================
# Containment checks: is a sphere fully inside the container?
# ============================================================

def sphere_in_box(cx, cy, cz, r, Lx, Ly, Lz, **kw):
    """Check if sphere at (cx,cy,cz) with radius r fits in box [0,Lx]x[0,Ly]x[0,Lz]."""
    return (cx - r >= 0 and cx + r <= Lx and
            cy - r >= 0 and cy + r <= Ly and
            cz - r >= 0 and cz + r <= Lz)


def sphere_in_domed_box(cx, cy, cz, r, Lx, Ly, Lz, dome_height=0.0, **kw):
    """Box with an ellipsoidal dome cap on top."""
    if not (cx - r >= 0 and cx + r <= Lx and
            cy - r >= 0 and cy + r <= Ly and
            cz - r >= 0):
        return False
    if cz + r <= Lz:
        return True
    # Dome region: ellipsoidal cap centered at (Lx/2, Ly/2, Lz)
    # Check that the TOP of the sphere is below the dome surface at (cx, cy)
    a = min(Lx, Ly) / 2  # dome base radius
    h = dome_height
    dx_norm = (cx - Lx / 2) / a if a > 0 else 0
    dy_norm = (cy - Ly / 2) / a if a > 0 else 0
    r_sq = dx_norm**2 + dy_norm**2
    if r_sq >= 1:
        return False  # outside dome base circle
    dome_z = Lz + h * np.sqrt(1 - r_sq)
    return cz + r <= dome_z


def sphere_in_cylinder(cx, cy, cz, r, radius, height, **kw):
    """Cylinder centered at (radius, radius) on XY plane, height along Z."""
    dx = cx - radius
    dy = cy - radius
    dist_from_axis = np.sqrt(dx**2 + dy**2)
    return (dist_from_axis + r <= radius and
            cz - r >= 0 and cz + r <= height)


def sphere_in_fuselage(cx, cy, cz, r, length, fuselage_radius, floor_height, **kw):
    """Cylindrical fuselage along X-axis with floor cut.
    Circle center at Y=fuselage_radius, Z=fuselage_radius.
    Floor at Z = floor_height.
    """
    if cx - r < 0 or cx + r > length:
        return False
    if cz - r < floor_height:
        return False
    dy = cy - fuselage_radius
    dz = cz - fuselage_radius
    dist = np.sqrt(dy**2 + dz**2)
    return dist + r <= fuselage_radius


def sphere_in_rounded_box(cx, cy, cz, r, Lx, Ly, Lz, roof_sag=0.0, **kw):
    """Box with arched roof (parabolic sag from center to edges along Y)."""
    if not (cx - r >= 0 and cx + r <= Lx and
            cy - r >= 0 and cy + r <= Ly and
            cz - r >= 0):
        return False
    y_norm = 2 * (cy - Ly / 2) / Ly  # -1 to 1
    local_ceiling = Lz - roof_sag * y_norm**2
    return cz + r <= local_ceiling


CONTAINMENT_CHECKS = {
    "box": sphere_in_box,
    "domed_box": sphere_in_domed_box,
    "cylinder": sphere_in_cylinder,
    "fuselage": sphere_in_fuselage,
    "rounded_box": sphere_in_rounded_box,
}


# ============================================================
# Sphere packing into arbitrary containers
# ============================================================

def get_bounding_box(venue):
    """Get the axis-aligned bounding box for a venue."""
    dims = venue["dims"]
    vtype = venue["type"]
    if vtype == "box":
        return dims["Lx"], dims["Ly"], dims["Lz"]
    elif vtype == "domed_box":
        return dims["Lx"], dims["Ly"], dims["Lz"] + dims.get("dome_height", 0)
    elif vtype == "cylinder":
        d = 2 * dims["radius"]
        return d, d, dims["height"]
    elif vtype == "fuselage":
        d = 2 * dims["fuselage_radius"]
        return dims["length"], d, d
    elif vtype == "rounded_box":
        return dims["Lx"], dims["Ly"], dims["Lz"]
    return 1, 1, 1


def get_volume(venue):
    """Compute the actual interior volume of a venue."""
    dims = venue["dims"]
    vtype = venue["type"]
    if vtype == "box":
        return dims["Lx"] * dims["Ly"] * dims["Lz"]
    elif vtype == "domed_box":
        box_vol = dims["Lx"] * dims["Ly"] * dims["Lz"]
        dome_r = min(dims["Lx"], dims["Ly"]) / 2
        dome_vol = (2/3) * np.pi * dome_r**2 * dims.get("dome_height", 0)
        return box_vol + dome_vol
    elif vtype == "cylinder":
        return np.pi * dims["radius"]**2 * dims["height"]
    elif vtype == "fuselage":
        R = dims["fuselage_radius"]
        fh = dims["floor_height"]
        theta = 2 * np.arccos((fh - R) / R)
        seg_area = R**2 / 2 * (theta - np.sin(theta))
        return seg_area * dims["length"]
    elif vtype == "rounded_box":
        sag = dims.get("roof_sag", 0)
        box_vol = dims["Lx"] * dims["Ly"] * dims["Lz"]
        lost = dims["Lx"] * dims["Ly"] * sag / 3
        return box_vol - lost
    return 0


def _fcc_with_offset(radius, check_fn, dims, Lx, Ly, Lz, ox, oy, oz):
    """Generate FCC lattice with given origin offset, return centers list."""
    d = 2 * radius
    layer_height = d * np.sqrt(2/3)
    row_offset = d * np.sqrt(3) / 2

    centers = []
    iz = 0
    z = oz
    while z < Lz + radius:
        if z > -radius:
            iy = 0
            y_base = oy + (iz % 2) * (row_offset / 3)
            y = y_base
            while y < Ly + radius:
                if y > -radius:
                    x_base = ox + ((iy + iz) % 2) * radius
                    x = x_base
                    while x < Lx + radius:
                        if x > -radius and check_fn(x, y, z, radius, **dims):
                            centers.append([x, y, z])
                        x += d
                iy += 1
                y += row_offset
        iz += 1
        z += layer_height

    return centers


def _cubic_with_offset(radius, check_fn, dims, Lx, Ly, Lz, ox, oy, oz):
    """Generate cubic lattice with given origin offset, return centers list."""
    d = 2 * radius
    centers = []
    z = oz
    while z < Lz + radius:
        if z > -radius:
            y = oy
            while y < Ly + radius:
                if y > -radius:
                    x = ox
                    while x < Lx + radius:
                        if x > -radius and check_fn(x, y, z, radius, **dims):
                            centers.append([x, y, z])
                        x += d
                y += d
        z += d
    return centers


def pack_spheres_fcc_in_venue(radius, venue, n_shifts=5):
    """Pack spheres in FCC lattice, searching over grid offsets for best fit."""
    check_fn = CONTAINMENT_CHECKS[venue["type"]]
    dims = venue["dims"]
    bb = get_bounding_box(venue)
    Lx, Ly, Lz = bb

    d = 2 * radius
    layer_height = d * np.sqrt(2/3)
    row_offset = d * np.sqrt(3) / 2

    best_centers = []
    shifts = np.linspace(0, 1, n_shifts, endpoint=False)
    for sx in shifts:
        for sy in shifts:
            for sz in shifts:
                ox = radius + sx * d
                oy = radius + sy * row_offset
                oz = radius + sz * layer_height
                c = _fcc_with_offset(radius, check_fn, dims, Lx, Ly, Lz,
                                     ox, oy, oz)
                if len(c) > len(best_centers):
                    best_centers = c

    return len(best_centers), np.array(best_centers) if best_centers else np.empty((0, 3))


def pack_spheres_cubic_in_venue(radius, venue, n_shifts=5):
    """Pack spheres in simple cubic lattice, searching over grid offsets."""
    check_fn = CONTAINMENT_CHECKS[venue["type"]]
    dims = venue["dims"]
    bb = get_bounding_box(venue)
    Lx, Ly, Lz = bb

    d = 2 * radius

    best_centers = []
    shifts = np.linspace(0, 1, n_shifts, endpoint=False)
    for sx in shifts:
        for sy in shifts:
            for sz in shifts:
                ox = radius + sx * d
                oy = radius + sy * d
                oz = radius + sz * d
                c = _cubic_with_offset(radius, check_fn, dims, Lx, Ly, Lz,
                                       ox, oy, oz)
                if len(c) > len(best_centers):
                    best_centers = c

    return len(best_centers), np.array(best_centers) if best_centers else np.empty((0, 3))


# ============================================================
# PyVista mesh generation for container visualization
# ============================================================

def _fuselage_cross_section(R, floor_height, n_arc=48):
    """Generate 2D cross-section for circular fuselage with floor.
    Returns Nx2 array of (y, z) points forming a closed polygon.
    """
    cos_val = np.clip((floor_height - R) / R, -1, 1)
    theta_floor = np.arccos(cos_val)

    angles = np.linspace(np.pi + theta_floor, -(np.pi + theta_floor) + 2 * np.pi,
                         n_arc)
    arc_y = R + R * np.cos(angles)
    arc_z = R + R * np.sin(angles)

    pts = np.column_stack([arc_y, arc_z])
    floor_left_y = R + R * np.cos(np.pi + theta_floor)
    floor_right_y = R + R * np.cos(-theta_floor)
    pts = np.vstack([pts, [floor_right_y, floor_height], [floor_left_y, floor_height]])

    return pts


def make_container_mesh(venue):
    """Create a pyvista mesh for the container shape."""
    dims = venue["dims"]
    vtype = venue["type"]

    if vtype == "box":
        return pv.Box(bounds=[0, dims["Lx"], 0, dims["Ly"], 0, dims["Lz"]])

    elif vtype == "domed_box":
        Lx, Ly, Lz = dims["Lx"], dims["Ly"], dims["Lz"]
        dome_h = dims.get("dome_height", 0)
        dome_r = min(Lx, Ly) / 2

        box = pv.Box(bounds=[0, Lx, 0, Ly, 0, Lz])

        n_rings = 10
        n_circ = 24
        pts = []
        for i in range(n_rings + 1):
            t = i / n_rings
            scale = np.sqrt(1 - t**2) if t < 1 else 0
            z = Lz + dome_h * t
            for j in range(n_circ):
                angle = 2 * np.pi * j / n_circ
                y = Ly / 2 + dome_r * scale * np.cos(angle)
                x = Lx / 2 + dome_r * scale * np.sin(angle)
                pts.append([x, y, z])
        pts = np.array(pts)

        faces = []
        for i in range(n_rings):
            for j in range(n_circ):
                j_next = (j + 1) % n_circ
                a = i * n_circ + j
                b = i * n_circ + j_next
                c = (i + 1) * n_circ + j_next
                d = (i + 1) * n_circ + j
                faces.append([4, a, b, c, d])

        face_arr = np.hstack(faces)
        dome_mesh = pv.PolyData(pts, faces=face_arr)
        return box.merge(dome_mesh)

    elif vtype == "cylinder":
        r = dims["radius"]
        h = dims["height"]
        cyl = pv.Cylinder(center=(r, r, h/2), direction=(0, 0, 1),
                           radius=r, height=h, resolution=48)
        return cyl

    elif vtype == "fuselage":
        R = dims["fuselage_radius"]
        L = dims["length"]
        fh = dims["floor_height"]

        cs = _fuselage_cross_section(R, fh, n_arc=32)
        n_cs = len(cs)

        n_rings = max(8, int(L / 2))
        x_vals = np.linspace(0, L, n_rings)

        all_pts = []
        for x in x_vals:
            ring = np.column_stack([np.full(n_cs, x), cs[:, 0], cs[:, 1]])
            all_pts.append(ring)
        all_pts = np.vstack(all_pts)

        faces = []
        for r_idx in range(n_rings - 1):
            off0 = r_idx * n_cs
            off1 = (r_idx + 1) * n_cs
            for i in range(n_cs):
                j = (i + 1) % n_cs
                faces.append([4, off0 + i, off0 + j, off1 + j, off1 + i])

        front_face = [n_cs] + list(range(n_cs))
        last_off = (n_rings - 1) * n_cs
        back_face = [n_cs] + list(range(last_off, last_off + n_cs))
        faces.append(front_face)
        faces.append(back_face)

        face_arr = np.hstack(faces)
        mesh = pv.PolyData(all_pts, faces=face_arr)
        return mesh

    elif vtype == "rounded_box":
        Lx, Ly, Lz = dims["Lx"], dims["Ly"], dims["Lz"]
        sag = dims.get("roof_sag", 0)

        ny = 8
        y_vals = np.linspace(0, Ly, ny)
        y_norm = 2 * (y_vals - Ly / 2) / Ly
        roof_z = Lz - sag * y_norm**2

        pts = []
        for i, y in enumerate(y_vals):
            rz = roof_z[i]
            pts.extend([[0, y, 0], [Lx, y, 0], [Lx, y, rz], [0, y, rz]])
        pts = np.array(pts)

        faces = []
        for i in range(ny - 1):
            base = i * 4
            nxt = (i + 1) * 4
            faces.append([4, base, base+1, nxt+1, nxt])
            faces.append([4, base+1, base+2, nxt+2, nxt+1])
            faces.append([4, base+2, base+3, nxt+3, nxt+2])
            faces.append([4, base+3, base, nxt, nxt+3])

        faces.append([4, 0, 1, 2, 3])
        last = (ny - 1) * 4
        faces.append([4, last, last+1, last+2, last+3])

        face_arr = np.hstack(faces)
        mesh = pv.PolyData(pts, faces=face_arr)
        return mesh

    return pv.Box(bounds=[0, 1, 0, 1, 0, 1])