import re
import os
import numpy as np
import gradio as gr
import matplotlib.pyplot as plt

from io import BytesIO
from PIL import Image
from canvas import Idefics2Pipeline


def run_canvas(front_view, map_view, prompt):
    pipeline = Idefics2Pipeline.from_pretrained(
        "maum-ai/CANVAS-S"
    )
    messages = [
        {"role": "system", "content": [{"type": "text", "text": prompt}]},
        {
            "role": "user",
            "content": [{"type": "image"}, {"type": "image"}],
        },
    ]

    images = [front_view, map_view]
    pred = pipeline([messages], [images], return_traj=False)
    pred_action = re.findall(r"<ACTION_(\d+)>", pred[0])
    pred_action = np.array(pred_action, dtype=np.int64)
    pred_action_odom = pipeline.action_tokenizer.detokenize(pred_action).tolist()

    # Create a figure and axes
    fig, axes = plt.subplots(1, 1, figsize=(8, 8))

    # Scale factor for the arrow
    scale_factor = 0.2

    axes.plot(0, 0, marker="o", color="black", markersize=10)
    axes.invert_xaxis()

    for i, center in zip(pred_action, pred_action_odom):
        x, y, yaw = center
        axes.plot(y, x, marker="^", color="blue")
        axes.arrow(
            y,
            x,
            np.sin(yaw) * scale_factor,
            np.cos(yaw) * scale_factor,
            head_width=scale_factor * 0.3,
            head_length=scale_factor * 0.3,
            fc="k",
            ec="k",
        )
        axes.text(y, x, f"{i}", fontsize=20)
    axes.axis("equal")
    axes.grid(True)

    buf = BytesIO()
    fig.savefig(buf, format="png")
    buf.seek(0)  # Rewind the buffer to the beginning
    pil_img = Image.open(buf)

    return pil_img

examples_path = os.path.dirname(__file__)

examples = [
    [f"{examples_path}/src/office/0a1f277a93fed629365ac5863c20c64e_frontview_6.0.png", f"{examples_path}/src/office/0a1f277a93fed629365ac5863c20c64e_map_6.0.png", """You are an indoor food-serving robot.

You must follow these driving instructions:
1. You must avoid collisions.
2. You should prioritize reaching the final destination.
3. You should follow the Trajectory Instruction.
    a. If the Trajectory Instruction cannot be followed due to any obstacles, you should deviate to bypass the obstacle.
    b. You should try to evade any identifiable obstacles.
4. You should maintain a constant driving speed.
    a. Indoors, you should drive at a speed of 3-4km/h.
5. You must slow down(2km/h or lower) if a human or obstacle comes within 1.5m radius.
    a. You must slow down(2km/h or lower) in areas where a human could suddenly appear from a blind spot."""],
    [f"{examples_path}/src/orchard/d578264e1e51cc5b8f0e496ab381cee4_frontview_79.0.png", f"{examples_path}/src/orchard/d578264e1e51cc5b8f0e496ab381cee4_map_79.0.png", """You are an outdoor speed-sprayer robot.

You must follow these driving instructions:
1. You must avoid collisions.
2. You should prioritize reaching the final destination.
3. You should follow the Trajectory Instruction.
    a. If the Trajectory Instruction cannot be followed due to any obstacles, you should deviate to bypass the obstacle.
    b. You should try to evade any identifiable obstacles.
4. You should maintain a constant driving speed."""],
    [f"{examples_path}/src/sidewalk/2d0dde2a98083b7d60b24651d37532dc_frontview_4.0.png", f"{examples_path}/src/sidewalk/2d0dde2a98083b7d60b24651d37532dc_map_4.0.png", """You are an outdoor last mile delivery robot.

You must follow these driving instructions:
1. You must avoid collisions.
2. You should prioritize reaching the final destination.
3. You should follow the Trajectory Instruction.
    a. If the Trajectory Instruction cannot be followed due to any obstacles, you should deviate to bypass the obstacle.
    b. You should try to evade any identifiable obstacles.
4. You should maintain a constant driving speed.
5. You must drive on the sidewalk.
    a. If you need to cross the road, you must use the crosswalk."""],
    [f"{examples_path}/src/road/3cfce98ab33a3dc8d43584d5a7039cf5_frontview_8.75.png", f"{examples_path}/src/road/3cfce98ab33a3dc8d43584d5a7039cf5_map_8.75.png", """You are an outdoor self-driving robot taxi.

You must follow these driving instructions:
1. You must avoid collisions.
2. You should prioritize reaching the final destination.
3. You should follow the Trajectory Instruction.
    a. If the Trajectory Instruction cannot be followed due to any obstacles, you should deviate to bypass the obstacle.
    b. You should try to evade any identifiable obstacles.
4. You should maintain a constant driving speed.
5. You must drive on the road.
    a. You should drive according to the left-hand-traffic law.
6. You should slow down before entering intersections, speed bumps, and crosswalks."""],
]

demo = gr.Interface(
    fn = run_canvas,
    inputs = [
        gr.Image(label="front_view", type="pil"),
        gr.Image(label="map_view", type="pil"),
        gr.Textbox(label="prompt")
    ],
    outputs = gr.Image(label="generated waypoint"),
    title="CANVAS Demo",
    description="This is the demo of the [CANVAS-S](https://huggingface.co/maum-ai/CANVAS-S) model from our paper, [CANVAS: Commonsense-Aware Navigation System for Intuitive Human-Robot Interaction](https://worv-ai.github.io/canvas). In CPU, it takes a minute to predict robot's next action token.",
    examples=examples
)

demo.launch()