import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from .arrays import to_np
from .video import save_video

from obstacles import PILLARS, NOVEL_CIRCULAR, NOVEL_PLANAR, CONSTRAINT_SETS


def save_single_trajectory_image(
    trajectory, filepath, constraint="", obstacle_margin=0.0, draw_obstacle_margin=False
):
    plotter = AvoidingTrajectoryPlotter(
        constraint=constraint,
        obstacle_margin=obstacle_margin,
        draw_obstacle_margin=draw_obstacle_margin,
    )
    fig, ax = plotter.setup_figure()
    plotter.plot_single_trajectory(trajectory, ax, style="multiple_real")
    plotter.add_environment_elements(ax)
    plotter.apply_legend(ax)
    plotter.save_figure(fig, filepath)


def save_trajectory_sequence(
    trajectories,
    base_path,
    folder_name,
    constraint="",
    obstacle_margin=0.0,
    draw_obstacle_margin=False,
):
    import os

    img_path = os.path.join(base_path, folder_name)
    os.makedirs(img_path, exist_ok=True)

    for step_i, trajectory in enumerate(trajectories):
        imgpath = os.path.join(img_path, f"{step_i}.png")
        save_single_trajectory_image(
            trajectory, imgpath, constraint, obstacle_margin, draw_obstacle_margin
        )


def save_prediction_overlay_plot(
    predicted_trajectories,
    real_trajectory,
    filepath,
    title=None,
    constraint="",
    obstacle_margin=0.0,
    draw_obstacle_margin=False,
):
    plotter = AvoidingTrajectoryPlotter(
        constraint=constraint,
        obstacle_margin=obstacle_margin,
        draw_obstacle_margin=draw_obstacle_margin,
    )
    fig, ax = plotter.plot_overlaid_predictions_and_real(
        predicted_trajectories,
        real_trajectory,
        title=title,
    )
    plotter.save_figure(fig, filepath)


def plot2img(fig, remove_margins=True):

    from matplotlib.backends.backend_agg import FigureCanvasAgg

    if remove_margins:
        fig.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=0, hspace=0)

    canvas = FigureCanvasAgg(fig)
    canvas.draw()
    img_as_string, (width, height) = canvas.print_to_buffer()
    return np.fromstring(img_as_string, dtype="uint8").reshape((height, width, 4))


class AvoidingRenderer:

    def __init__(
        self,
        env_name,
        env,
        constraint="",
        obstacle_margin=0.0,
        draw_obstacle_margin=False,
    ):
        self.env_name = env_name
        self.env = env
        self.constraint = constraint
        self.obstacle_margin = obstacle_margin
        self.draw_obstacle_margin = draw_obstacle_margin

    def renders(self, observations):
        plotter = AvoidingTrajectoryPlotter(
            constraint=self.constraint,
            obstacle_margin=self.obstacle_margin,
            draw_obstacle_margin=self.draw_obstacle_margin,
        )
        fig, ax = plotter.setup_figure()
        plotter.plot_single_trajectory(observations, ax, style="multiple_real")
        plotter.add_environment_elements(ax)
        plotter.apply_legend(ax)
        img = plot2img(fig, remove_margins=False)
        plt.close(fig)
        return img

    def render_diffusion(self, savepath, diffusion_path, **video_kwargs):
        diffusion_path = to_np(diffusion_path)
        n_diffusion_steps = diffusion_path.shape[0]
        frame = []
        for t in range(n_diffusion_steps):
            print(f"[ utils/renderer ] Diffusion: {t} / {n_diffusion_steps}")
            img = self.renders(diffusion_path[t])
            frame.append(img)
        frame = np.stack(frame, axis=0)
        save_video(savepath, frame, **video_kwargs)


class AvoidingTrajectoryPlotter:

    def __init__(self, constraint="", obstacle_margin=0.0, draw_obstacle_margin=False):
        self.constraint = constraint
        self.obstacle_margin = obstacle_margin
        self.draw_obstacle_margin = draw_obstacle_margin
        self.ax_limits = [[0.2, 0.8], [-0.3, 0.4]]
        self.figsize = (9, 9)
        self.fontsize_label = 16
        self.fontsize_ticks = 14
        self.fontsize_legend = 16
        self._legend_flags = {
            "target": False,
            "obstacle": False,
            "start": False,
            "end": False,
        }

    def setup_figure(self):
        fig, ax = plt.subplots(figsize=self.figsize)
        ax.set_xlim(self.ax_limits[0])
        ax.set_ylim(self.ax_limits[1])
        ax.set_facecolor([1, 1, 0.9])
        ax.set_aspect("equal", adjustable="box")

        fig.subplots_adjust(left=0.06, right=0.985, bottom=0.06, top=0.985)
        ax.xaxis.labelpad = 2
        ax.yaxis.labelpad = 2
        ax.tick_params(axis="both", which="major", pad=1, labelsize=self.fontsize_ticks)
        ax.set_xlabel("X Position", fontsize=self.fontsize_label)
        ax.set_ylabel("Y Position", fontsize=self.fontsize_label)
        ax.grid(True, alpha=0.3)
        return fig, ax

    def add_environment_elements(self, ax):
        target_color = [0.4, 1, 0.4]

        x_min, x_max = self.ax_limits[0]
        y_min_target = 0.35
        y_max = self.ax_limits[1][1]
        if y_max > y_min_target:
            target_patch = patches.Rectangle(
                (x_min, y_min_target),
                x_max - x_min,
                y_max - y_min_target,
                facecolor=target_color,
                alpha=0.35,
                label="Target" if not self._legend_flags["target"] else "",
                zorder=0.5,
            )
            ax.add_patch(target_patch)
            self._legend_flags["target"] = True

        polytopic_constraints = self._get_polytopic_constraints(ax)
        self._add_polytopic_constraints(ax, polytopic_constraints)
        self._add_target_circles(ax)

    def _get_polytopic_constraints(self, ax):
        if self.constraint == "":
            return []
        elif self.constraint not in CONSTRAINT_SETS:
            raise ValueError(f"Unsupported constraint: {self.constraint}")

        for obs in NOVEL_CIRCULAR:
            center = list(obs.center)
            ax.add_patch(
                patches.Circle(
                    center,
                    obs.radius,
                    color="b",
                    alpha=0.3,
                    label="Obstacle" if not self._legend_flags["obstacle"] else "",
                )
            )
            self._legend_flags["obstacle"] = True

            if self.obstacle_margin > 0 and self.draw_obstacle_margin:
                self._add_margin_ring(ax, center, obs.radius, self.obstacle_margin, "b")

        return [
            [list(obs.point_a), list(obs.point_b), "below"]
            for obs in NOVEL_PLANAR
        ]

    def _add_polytopic_constraints(self, ax, polytopic_constraints):
        if self.constraint == "":
            return

        for i, constraint in enumerate(polytopic_constraints):
            mat = np.vstack((constraint[:2], np.zeros(2)))

            slope = (constraint[1][1] - constraint[0][1]) / (
                constraint[1][0] - constraint[0][0]
            )

            if slope > 0 and constraint[2] == "below":
                mat[2] = np.array([self.ax_limits[0][0], self.ax_limits[1][1]])
            elif slope < 0 and constraint[2] == "below":
                mat[2] = np.array([self.ax_limits[0][1], self.ax_limits[1][1]])
            elif slope > 0 and constraint[2] == "above":
                mat[2] = np.array([self.ax_limits[0][1], self.ax_limits[1][0]])
            elif slope < 0 and constraint[2] == "above":
                mat[2] = np.array([self.ax_limits[0][0], self.ax_limits[1][0]])

            clipped = self._clip_polygon_below_y(mat, 0.35)
            if len(clipped) >= 3:
                ax.add_patch(
                    patches.Polygon(
                        np.array(clipped),
                        color="b",
                        alpha=0.2,
                        label="Obstacle" if not self._legend_flags["obstacle"] else "",
                    )
                )
            self._legend_flags["obstacle"] = True

            if self.obstacle_margin > 0 and self.draw_obstacle_margin:
                self._add_margin_parallelogram(ax, constraint, "b")

    def _add_target_circles(self, ax):
        for idx, pillar in enumerate(PILLARS):
            center = list(pillar.center)
            radius = pillar.radius
            ax.add_patch(
                patches.Circle(
                    center,
                    radius,
                    color="r",
                    alpha=0.7,
                    label="Obstacle" if not self._legend_flags["obstacle"] else "",
                )
            )
            self._legend_flags["obstacle"] = True
            if self.obstacle_margin > 0 and self.draw_obstacle_margin:
                if not self._is_pillar_covered_by_obstacle(center):
                    self._add_margin_ring(ax, center, radius, self.obstacle_margin, "r")

    def _is_pillar_covered_by_obstacle(self, pillar_center):
        if self.constraint not in CONSTRAINT_SETS:
            return False
        for obs in NOVEL_CIRCULAR:
            distance = np.sqrt(
                (pillar_center[0] - obs.center[0]) ** 2
                + (pillar_center[1] - obs.center[1]) ** 2
            )
            if distance < obs.radius:
                return True
        return False

    def _add_margin_ring(self, ax, center, inner_radius, margin, color):
        outer_radius = inner_radius + margin

        outer_circle = patches.Circle(
            center, outer_radius, color=color, alpha=0.15, zorder=1
        )
        ax.add_patch(outer_circle)

        border_circle = patches.Circle(
            center,
            outer_radius,
            fill=False,
            edgecolor=color,
            alpha=0.7,
            linewidth=1.5,
            linestyle="--",
            zorder=3,
        )
        ax.add_patch(border_circle)

    def _add_margin_parallelogram(self, ax, constraint, color):
        p1, p2 = constraint[:2]
        slope = (p2[1] - p1[1]) / (p2[0] - p1[0])
        intercept = p2[1] - slope * p2[0]

        normal_x = -slope / np.sqrt(1 + slope**2)
        normal_y = 1 / np.sqrt(1 + slope**2)

        if constraint[2] == "below":
            margin_offset_x = -normal_x * self.obstacle_margin
            margin_offset_y = -normal_y * self.obstacle_margin
        else:
            margin_offset_x = normal_x * self.obstacle_margin
            margin_offset_y = normal_y * self.obstacle_margin

        original_mat = np.vstack((constraint[:2], np.zeros(2)))

        if slope > 0 and constraint[2] == "below":
            original_mat[2] = np.array([self.ax_limits[0][0], self.ax_limits[1][1]])
        elif slope < 0 and constraint[2] == "below":
            original_mat[2] = np.array([self.ax_limits[0][1], self.ax_limits[1][1]])
        elif slope > 0 and constraint[2] == "above":
            original_mat[2] = np.array([self.ax_limits[0][1], self.ax_limits[1][0]])
        elif slope < 0 and constraint[2] == "above":
            original_mat[2] = np.array([self.ax_limits[0][0], self.ax_limits[1][0]])

        margin_mat = original_mat.copy()
        margin_mat[0] += [margin_offset_x, margin_offset_y]
        margin_mat[1] += [margin_offset_x, margin_offset_y]

        margin_vertices = np.vstack([original_mat[:2], margin_mat[:2][::-1]])

        ax.add_patch(
            patches.Polygon(
                margin_vertices,
                color=color,
                alpha=0.1,
            )
        )

    def _clip_polygon_below_y(self, vertices, y_max):
        verts = [np.array(v, dtype=float).tolist() for v in vertices]
        if len(verts) == 0:
            return []
        output = []
        prev = verts[-1]
        prev_inside = prev[1] <= y_max
        for curr in verts:
            curr_inside = curr[1] <= y_max
            if curr_inside:
                if not prev_inside:
                    denom = curr[1] - prev[1]
                    if denom != 0:
                        t = (y_max - prev[1]) / denom
                        x_int = prev[0] + t * (curr[0] - prev[0])
                        output.append([x_int, y_max])
                output.append(curr)
            else:
                if prev_inside:
                    denom = curr[1] - prev[1]
                    if denom != 0:
                        t = (y_max - prev[1]) / denom
                        x_int = prev[0] + t * (curr[0] - prev[0])
                        output.append([x_int, y_max])
            prev, prev_inside = curr, curr_inside
        return output

    def plot_single_trajectory(
        self, trajectory, ax, style="actual", label_prefix="", show_labels=True
    ):
        x_coords = trajectory[:, 2]
        y_coords = trajectory[:, 3]

        if style == "actual":
            ax.plot(
                x_coords,
                y_coords,
                "k-",
                linewidth=3,
                label=f"{label_prefix}Actual Trajectory" if show_labels else "",
            )
            ax.plot(
                x_coords[0],
                y_coords[0],
                "go",
                markersize=10,
                markeredgecolor="darkgreen",
                markeredgewidth=1,
                label="Start" if not self._legend_flags["start"] else "",
            )
            self._legend_flags["start"] = True
            ax.plot(
                x_coords[-1],
                y_coords[-1],
                "bo",
                markersize=10,
                markeredgecolor="navy",
                markeredgewidth=1,
                label="End" if not self._legend_flags["end"] else "",
            )
            self._legend_flags["end"] = True
            if len(x_coords) > 2:
                ax.plot(
                    x_coords[1:-1],
                    y_coords[1:-1],
                    "ko",
                    markersize=3,
                    alpha=0.7,
                )

        elif style == "predicted":
            ax.plot(
                x_coords,
                y_coords,
                "m-",
                linewidth=1.5,
                alpha=0.4,
                label=f"{label_prefix}Predicted Trajectories" if show_labels else "",
            )
            ax.plot(
                x_coords[0],
                y_coords[0],
                "ms",
                markersize=5,
                alpha=0.5,
            )
            ax.plot(
                x_coords[-1],
                y_coords[-1],
                "m^",
                markersize=5,
                alpha=0.5,
            )
            if len(x_coords) > 2:
                ax.plot(
                    x_coords[1:-1],
                    y_coords[1:-1],
                    "mx",
                    markersize=3,
                    alpha=0.4,
                )

        elif style == "multiple_real":
            ax.plot(
                x_coords,
                y_coords,
                "k-",
                linewidth=2,
                label=f"{label_prefix}Trajectory" if show_labels else "",
            )
            ax.plot(
                x_coords[0],
                y_coords[0],
                "go",
                markersize=8,
                label="Start" if not self._legend_flags["start"] else "",
            )
            self._legend_flags["start"] = True
            ax.plot(
                x_coords[-1],
                y_coords[-1],
                "bo",
                markersize=8,
                label="End" if not self._legend_flags["end"] else "",
            )
            self._legend_flags["end"] = True
            if len(x_coords) > 2:
                ax.plot(
                    x_coords[1:-1],
                    y_coords[1:-1],
                    "co",
                    markersize=4,
                    alpha=0.6,
                )

    def plot_overlaid_predictions_and_real(
        self,
        predicted_trajectories,
        real_trajectory,
        title=None,
    ):
        fig, ax = self.setup_figure()

        for i, trajectory in enumerate(predicted_trajectories):
            self.plot_single_trajectory(
                trajectory, ax, style="predicted", show_labels=(i == 0)
            )

        self.plot_single_trajectory(real_trajectory, ax, style="actual")

        if title:
            ax.set_title(title, fontsize=self.fontsize_label)

        self.add_environment_elements(ax)

        self.apply_legend(ax)

        return fig, ax

    def _trajectory_violates(self, trajectory):
        """Check if any point in the trajectory violates constraints."""
        if self.constraint == "":
            return False
        from obstacles import check_violation
        for t in range(len(trajectory)):
            if check_violation(trajectory[t], self.constraint) > 0:
                return True
        return False

    def plot_multiple_real_trajectories(self, trajectories, title=None):
        fig, ax = self.setup_figure()

        for i, trajectory in enumerate(trajectories):
            x_coords = trajectory[:, 2]
            y_coords = trajectory[:, 3]

            violated = self._trajectory_violates(trajectory)
            color = "gray" if violated else "k"

            ax.plot(
                x_coords,
                y_coords,
                color=color,
                linestyle="-",
                linewidth=2,
                alpha=0.7,
                label="Trajectory" if i == 0 else "",
            )
            ax.plot(
                x_coords[0],
                y_coords[0],
                "go",
                markersize=8,
                label="Start" if not self._legend_flags["start"] else "",
            )
            self._legend_flags["start"] = True
            ax.plot(
                x_coords[-1],
                y_coords[-1],
                "bo",
                markersize=8,
                label="End" if not self._legend_flags["end"] else "",
            )
            self._legend_flags["end"] = True
            if len(x_coords) > 2:
                ax.plot(x_coords[1:-1], y_coords[1:-1], "co", markersize=4, alpha=0.6)

        self.add_environment_elements(ax)

        self.apply_legend(ax)

        return fig, ax

    def save_figure(self, fig, filepath, dpi=150, bbox_inches=None):
        fig.savefig(filepath, dpi=dpi, bbox_inches=bbox_inches)
        plt.close(fig)
        print(f"Saved trajectory plot to: {filepath}")

    def apply_legend(self, ax):

        from matplotlib.legend_handler import HandlerTuple

        handles, labels = ax.get_legend_handles_labels()

        new_handles = []
        new_labels = []
        seen = set()
        for h, l in zip(handles, labels):
            if not l:
                continue
            if l.lower() == "obstacle":
                continue
            if l not in seen:
                seen.add(l)
                new_handles.append(h)
                new_labels.append(l)

        purple_proxy = patches.Rectangle((0, 0), 1, 1, facecolor="b", alpha=0.2)
        red_proxy = patches.Circle((0.5, 0.5), 0.5, facecolor="r", alpha=0.7)
        composite = (purple_proxy, red_proxy)

        new_handles.append(composite)
        new_labels.append("Obstacle")

        ax.legend(
            new_handles,
            new_labels,
            handler_map={tuple: HandlerTuple(ndivide=None)},
            loc="lower left",
            fontsize=self.fontsize_legend,
        )