# VARK Simple Parking Grid Demo
# Paste this file as app.py in a Hugging Face Space (Gradio SDK)

import time
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

import gradio as gr
import matplotlib.pyplot as plt
import matplotlib.patches as patches

# -----------------------------
# Config
# -----------------------------
ROWS = 4
COLS = 5
PLATFORM_COUNT = 16

# Two external cells (22 total cells shown)
DROP_POINT = (-1, 0)              # outside the 20-cell grid, top-left
PICK_POINT = (ROWS, COLS - 1)     # outside the 20-cell grid, bottom-right

STEP_DELAY = 0.08

# -----------------------------
# Data models
# -----------------------------
@dataclass
class Platform:
    pid: int
    r: int
    c: int
    car_id: Optional[int] = None
    home: Tuple[int, int] = (0, 0)


@dataclass
class State:
    platforms: Dict[int, Platform] = field(default_factory=dict)
    occupancy: Dict[Tuple[int, int], Optional[int]] = field(default_factory=dict)
    cars: Dict[int, int] = field(default_factory=dict)  # car_id -> platform_id
    log: List[str] = field(default_factory=list)

    # animation overlays
    moving_pid: Optional[int] = None
    moving_dir: Optional[str] = None  # 'U'|'D'|'L'|'R'


# -----------------------------
# Helpers
# -----------------------------
def in_main_grid(r: int, c: int) -> bool:
    return 0 <= r < ROWS and 0 <= c < COLS


def in_extended(r: int, c: int) -> bool:
    return in_main_grid(r, c) or (r, c) == DROP_POINT or (r, c) == PICK_POINT


def manhattan(a: Tuple[int, int], b: Tuple[int, int]) -> int:
    return abs(a[0] - b[0]) + abs(a[1] - b[1])


def arrow_for_dir(d: Optional[str]) -> str:
    return {
        'U': '↑',
        'D': '↓',
        'L': '←',
        'R': '→',
        None: ''
    }.get(d, '')


# -----------------------------
# Init
# -----------------------------
def init_state() -> State:
    st = State()

    # four empty cells inside the 20-cell grid
    empty = {(0, 4), (1, 2), (3, 0), (3, 4)}

    # occupancy for all displayed cells (20 + 2)
    for r in range(-1, ROWS + 1):
        for c in range(-1, COLS + 1):
            if in_extended(r, c):
                st.occupancy[(r, c)] = None

    pid = 1
    for r in range(ROWS):
        for c in range(COLS):
            if (r, c) in empty:
                continue
            if pid > PLATFORM_COUNT:
                continue
            p = Platform(pid=pid, r=r, c=c, home=(r, c))
            st.platforms[pid] = p
            st.occupancy[(r, c)] = pid
            pid += 1

    st.log.append("System initialised: 4x5 grid (20) + 2 external cells = 22 total")
    st.log.append(f"DROP point at {DROP_POINT} (outside top-left)")
    st.log.append(f"PICK point at {PICK_POINT} (outside bottom-right)")
    return st


# -----------------------------
# Rendering (frame-by-frame)
# -----------------------------
def render(st: State):
    fig, ax = plt.subplots(figsize=(8, 6))

    def xy_for_cell(cell: Tuple[int, int]) -> Tuple[float, float]:
        if cell == DROP_POINT:
            # place above (0,0)
            return (0, ROWS)
        if cell == PICK_POINT:
            # place below bottom-right main cell
            return (COLS - 1, -1)
        r, c = cell
        return (c, ROWS - 1 - r)

    # Draw main grid
    for r in range(ROWS):
        for c in range(COLS):
            x, y = xy_for_cell((r, c))
            ax.add_patch(patches.Rectangle((x, y), 1, 1, fill=False, linewidth=1.2))

    # Draw external cells
    for cell, label in [(DROP_POINT, "DROP"), (PICK_POINT, "PICK")]:
        x, y = xy_for_cell(cell)
        ax.add_patch(patches.Rectangle((x, y), 1, 1, fill=False, linewidth=2.6))
        ax.text(x + 0.5, y + 0.86, label, ha="center", va="center", fontsize=10)

    # Blue fill when carrying a car
    for p in st.platforms.values():
        if p.car_id is None:
            continue
        x, y = xy_for_cell((p.r, p.c))
        ax.add_patch(patches.Rectangle((x, y), 1, 1, fill=True, alpha=0.25))

    # Platform labels + motion arrow (direction of next step)
    for p in st.platforms.values():
        x, y = xy_for_cell((p.r, p.c))
        moving_arrow = arrow_for_dir(st.moving_dir) if st.moving_pid == p.pid else ''

        label = f"P{p.pid}"
        if p.car_id is not None:
            label += f"\nCar {p.car_id}"
        if moving_arrow:
            label = f"{moving_arrow}\n{label}"

        ax.text(x + 0.5, y + 0.5, label, ha="center", va="center", fontsize=9)

    # bounds include external cells
    ax.set_xlim(-0.2, COLS + 0.2)
    ax.set_ylim(-1.2, ROWS + 1.2)
    ax.set_aspect("equal")
    ax.axis("off")
    return fig


# -----------------------------
# Movement logic (4-direction only)
# -----------------------------
def step_towards(src: Tuple[int, int], dst: Tuple[int, int]) -> Tuple[int, int]:
    r, c = src
    tr, tc = dst

    # deterministic: move row first then col
    if r < tr:
        return (r + 1, c)
    if r > tr:
        return (r - 1, c)
    if c < tc:
        return (r, c + 1)
    if c > tc:
        return (r, c - 1)
    return (r, c)


def direction_from_to(a: Tuple[int, int], b: Tuple[int, int]) -> Optional[str]:
    ar, ac = a
    br, bc = b
    if br == ar - 1 and bc == ac:
        return 'U'
    if br == ar + 1 and bc == ac:
        return 'D'
    if br == ar and bc == ac - 1:
        return 'L'
    if br == ar and bc == ac + 1:
        return 'R'
    return None


def move_one(st: State, pid: int, target: Tuple[int, int]):
    p = st.platforms[pid]
    cur = (p.r, p.c)

    if not in_extended(*target):
        return

    if target not in st.occupancy:
        st.occupancy[target] = None

    occ = st.occupancy[target]
    if occ is None:
        st.occupancy[cur] = None
        p.r, p.c = target
        st.occupancy[target] = pid
    else:
        # swap to avoid overlap (simple demo rule)
        other = st.platforms[occ]
        p.r, other.r = other.r, p.r
        p.c, other.c = other.c, p.c
        st.occupancy[(p.r, p.c)] = pid
        st.occupancy[(other.r, other.c)] = occ


def animate_move(st: State, pid: int, dst: Tuple[int, int]):
    p = st.platforms[pid]
    st.moving_pid = pid

    while (p.r, p.c) != dst:
        cur = (p.r, p.c)
        nxt = step_towards(cur, dst)
        st.moving_dir = direction_from_to(cur, nxt)
        move_one(st, pid, nxt)
        yield st
        time.sleep(STEP_DELAY)

    st.moving_dir = None
    st.moving_pid = None


# -----------------------------
# Business logic
# -----------------------------
def free_platform(st: State) -> Optional[int]:
    free = [p for p in st.platforms.values() if p.car_id is None]
    if not free:
        return None
    free.sort(key=lambda p: (manhattan((p.r, p.c), DROP_POINT), p.pid))
    return free[0].pid


def drop_off(st: State, car_id: int):
    if car_id in st.cars:
        st.log.append(f"Drop rejected: Car {car_id} already parked")
        yield st
        return

    if len(st.cars) >= PLATFORM_COUNT:
        st.log.append("Drop rejected: capacity full (16 cars)")
        yield st
        return

    pid = free_platform(st)
    if pid is None:
        st.log.append("Drop rejected: no empty platform available")
        yield st
        return

    st.log.append(f"Dispatch P{pid} to DROP for Car {car_id}")
    for _ in animate_move(st, pid, DROP_POINT):
        yield st

    st.platforms[pid].car_id = car_id
    st.cars[car_id] = pid
    st.log.append(f"Car {car_id} loaded onto P{pid}")

    home = st.platforms[pid].home
    st.log.append(f"Store Car {car_id}: return P{pid} to {home}")
    for _ in animate_move(st, pid, home):
        yield st

    st.log.append(f"Drop complete: Car {car_id} stored")
    yield st


def pick_up(st: State, car_id: int):
    if car_id not in st.cars:
        st.log.append(f"Pick rejected: Car {car_id} not found")
        yield st
        return

    pid = st.cars[car_id]
    st.log.append(f"Dispatch P{pid} to PICK for Car {car_id}")

    for _ in animate_move(st, pid, PICK_POINT):
        yield st

    st.platforms[pid].car_id = None
    del st.cars[car_id]
    st.log.append(f"Car {car_id} unloaded at PICK")

    home = st.platforms[pid].home
    st.log.append(f"Return P{pid} to {home}")
    for _ in animate_move(st, pid, home):
        yield st

    st.log.append("Pick complete")
    yield st


# -----------------------------
# UI handlers
# -----------------------------
def reset_ui():
    st = init_state()
    return st, render(st), "\n".join(st.log[-16:]), f"{len(st.cars)} / 16"


def ui_drop(st: State, car: str):
    try:
        cid = int(car)
        if cid <= 0:
            raise ValueError
    except Exception:
        st.log.append("Drop rejected: Car ID must be a positive integer")
        yield st, render(st), "\n".join(st.log[-16:]), f"{len(st.cars)} / 16"
        return

    for s in drop_off(st, cid):
        yield s, render(s), "\n".join(s.log[-16:]), f"{len(s.cars)} / 16"


def ui_pick(st: State, car: str):
    try:
        cid = int(car)
        if cid <= 0:
            raise ValueError
    except Exception:
        st.log.append("Pick rejected: Car ID must be a positive integer")
        yield st, render(st), "\n".join(st.log[-16:]), f"{len(st.cars)} / 16"
        return

    for s in pick_up(st, cid):
        yield s, render(s), "\n".join(s.log[-16:]), f"{len(s.cars)} / 16"


# -----------------------------
# Gradio App
# -----------------------------
with gr.Blocks(title="VARK Parking Grid Demo (Animated)") as demo:
    state = gr.State(init_state())

    gr.Markdown(
        """
VARK grid demo (animated, frame-by-frame)

- Main grid: 4 x 5 = 20 cells
- External cells: DROP (outside top-left) and PICK (outside bottom-right)
- Total shown: 22 cells
- Platforms move only up / down / left / right
- Moving platform shows an arrow for direction each frame
- Platform cell turns blue when carrying a car
        """.strip()
    )

    with gr.Row():
        car_id = gr.Textbox(label="Car ID", value="101")
        capacity = gr.Textbox(label="Capacity", value="0 / 16", interactive=False)

    with gr.Row():
        drop_btn = gr.Button("DROP OFF")
        pick_btn = gr.Button("PICK UP")
        reset_btn = gr.Button("RESET")

    plot = gr.Plot(label="Grid")
    log = gr.Textbox(label="Log", lines=16, interactive=False)

    plot.value = render(state.value)
    log.value = "\n".join(state.value.log[-16:])

    reset_btn.click(reset_ui, outputs=[state, plot, log, capacity])
    drop_btn.click(ui_drop, inputs=[state, car_id], outputs=[state, plot, log, capacity])
    pick_btn.click(ui_pick, inputs=[state, car_id], outputs=[state, plot, log, capacity])


demo.queue().launch()