"""
2D rendering of the pressure plate domain
"""

import math
import os
import sys
from typing import Any

import numpy as np
import six
from gymnasium import error

if "Apple" in sys.version:
    if "DYLD_FALLBACK_LIBRARY_PATH" in os.environ:
        os.environ["DYLD_FALLBACK_LIBRARY_PATH"] += ":/usr/lib"
        # (JDS 2016/04/15): avoid bug on Anaconda 2.3.0 / Yosemite


try:
    import pyglet
except ImportError as e:
    raise ImportError(
        """
    Cannot import pyglet.
    HINT: you can install pyglet directly via 'pip install pyglet'.
    But if you really just want to install all Gym dependencies and not have to think about it,
    'pip install -e .[all]' or 'pip install gym[all]' will do it.
    """
    )

try:
    from pyglet.gl import *

except ImportError as e:
    raise ImportError(
        """
    Error occured while running `from pyglet.gl import *`
    HINT: make sure you have OpenGL install. On Ubuntu, you can run 'apt-get install python-opengl'.
    If you're running on a server, you may need a virtual frame buffer; something like this should work:
    'xvfb-run -s \"-screen 0 1400x900x24\" python <your_script.py>'
    """
    )


RAD2DEG = 57.29577951308232
# # Define some colors
_BLACK = (0, 0, 0, 255)
_WHITE = (255, 255, 255, 255)
_GREEN = (0, 255, 0, 255)
_RED = (255, 0, 0, 255)


def get_display(spec: Any):
    """Convert a display specification (such as :0) into an actual Display
    object.
    Pyglet only supports multiple Displays on Linux.
    """
    if spec is None:
        return None
    elif isinstance(spec, six.string_types):
        return pyglet.canvas.Display(spec)
    else:
        raise error.Error(
            "Invalid display specification: {}. (Must be a string like :0 or None.)".format(
                spec
            )
        )


class Viewer:
    def __init__(self, world_size, reward_dim):
        display = get_display(None)
        self.frame = 0
        # Let's have the display auto-scale to a 500x500 window
        self.rows, self.cols = world_size
        self.reward_dim = reward_dim

        self.grid_size = 800 / self.rows
        self.icon_size = 20

        self.grid_width = self.cols * self.grid_size + 1
        self.right_padding = 20
        self.left_padding = 20
        self.grid_height = self.rows * self.grid_size + 1
        self.top_padding = 100
        self.bottom_padding = 100

        disp_height = self.grid_height + self.top_padding + self.bottom_padding
        disp_width = (
            self.grid_width + self.left_padding + self.right_padding
        ) * (self.cols / self.rows)
        self.window = pyglet.window.Window(
            width=int(disp_width), height=int(disp_height), display=display
        )

        self.window.on_close = self.window_closed_by_user
        self.isopen = True

        glEnable(GL_BLEND)
        glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)

        script_dir = f"{os.path.dirname(os.path.dirname(__file__))}/assets"
        pyglet.resource.path = [os.path.join(script_dir, "icons")]
        pyglet.resource.reindex()

        # self.img_apple = pyglet.resource.image("apple.png")
        self.img_agent = pyglet.resource.image("agent.png")
        self.img_wall = pyglet.resource.image("brick-wall.png")
        self.img_door = pyglet.resource.image("spiked-fence.png")
        # self.img_door_left = pyglet.resource.image('doorway_left.png')
        # self.img_door_right = pyglet.resource.image('doorway_right.png')
        self.img_plate_off = pyglet.resource.image("plate_off.png")
        self.img_plate_on = pyglet.resource.image("plate_on.png")
        self.goals = [
            pyglet.resource.image(f"rewards/{i}/chest.png")
            for i in range(self.reward_dim)
        ]
        self.goals_open = [
            pyglet.resource.image(f"rewards/{i}/open-treasure-chest.png")
            for i in range(self.reward_dim)
        ]

    def close(self):
        self.window.close()

    def window_closed_by_user(self):
        self.isopen = False
        exit()

    def set_bounds(self, left, right, bottom, top):
        assert right > left and top > bottom
        scalex = self.grid_width / (right - left)
        scaley = self.grid_height / (top - bottom)
        self.transform = Transform(
            translation=(-left * scalex, -bottom * scaley),
            scale=(scalex, scaley),
        )

    def render(self, env, return_rgb_array=False):  # type:ignore
        glClearColor(0.95, 0.95, 0.95, 1)
        self.window.clear()
        self.window.switch_to()
        self.window.dispatch_events()

        # TODO: CHANGE
        self._draw_grid()
        self._draw_goal(env)
        self._draw_players(env)
        self._draw_badges(env)
        self.draw_labels(env)
        if return_rgb_array:
            buffer = pyglet.image.get_buffer_manager().get_color_buffer()
            image_data = buffer.get_image_data()
            arr = np.frombuffer(image_data.get_data(), dtype=np.uint8)
            arr = arr.reshape(buffer.height, buffer.width, 4)
            arr = arr[::-1, :, 0:3]
        self.window.flip()
        self.frame += 1
        return arr if return_rgb_array else self.isopen

    def _draw_grid(self):
        batch = pyglet.graphics.Batch()
        _ = pyglet.shapes.Rectangle(
            self.left_padding,
            self.bottom_padding,
            self.grid_height,
            self.grid_width,
            color=(150, 150, 150, 200),
            batch=batch,
        )
        lines = []
        for r in range(self.rows + 1):
            lines.append(
                pyglet.shapes.Line(
                    x=self.left_padding,
                    y=self.bottom_padding + self.grid_size * r,
                    x2=self.left_padding + self.grid_size * self.cols,
                    y2=self.bottom_padding + self.grid_size * r,
                    color=_BLACK,
                    batch=batch,
                )
            )

        for c in range(self.cols + 1):
            lines.append(
                pyglet.shapes.Line(
                    x=self.left_padding + self.grid_size * c,
                    y=self.bottom_padding,
                    x2=self.left_padding + self.grid_size * c,
                    y2=self.bottom_padding + self.grid_size * self.rows,
                    color=_BLACK,
                    batch=batch,
                )
            )

        batch.draw()

    def _draw_players(self, env):
        players = []
        batch = pyglet.graphics.Batch()

        for player in env.agents:
            row, col = player.y, player.x
            players.append(
                pyglet.sprite.Sprite(
                    self.img_agent,
                    self.left_padding + self.grid_size * col,
                    self.bottom_padding
                    + self.grid_height
                    - self.grid_size * (row + 1),
                    batch=batch,
                )
            )
        for p in players:
            p.update(scale=self.grid_size / p.width)
        batch.draw()
        # for p in env.agents:
        #     self._draw_badge(*(p.x, p.y), p.level)

    def draw_labels(self, env):
        l = pyglet.text.Label(
            f"Total score : {list(map(int, env.collected_chests))}",
            font_name="helvetica",
            font_size=30,
            color=(15, 15, 15, 255),
            x=self.left_padding + self.grid_width // 2,
            y=self.bottom_padding + self.grid_height + 50,
            anchor_x="center",
            anchor_y="center",
        )
        l.draw()

    def _draw_goal(self, env):
        goals = []
        batch = pyglet.graphics.Batch()
        for i in range(self.reward_dim):
            for j in range(len(env.goals[i])):
                curr_goal = env.goals[i][j]
                row, col = curr_goal.y, curr_goal.x

                if curr_goal.achieved:
                    goals.append(
                        pyglet.sprite.Sprite(
                            self.goals_open[i],
                            self.left_padding + self.grid_size * col,
                            self.bottom_padding
                            + self.grid_height
                            - self.grid_size * (row + 1),
                            batch=batch,
                        )
                    )
                else:
                    goals.append(
                        pyglet.sprite.Sprite(
                            self.goals[i],
                            self.left_padding + self.grid_size * col,
                            self.bottom_padding
                            + self.grid_height
                            - self.grid_size * (row + 1),
                            batch=batch,
                        )
                    )

        for g in goals:
            g.update(scale=self.grid_size / g.width)
        batch.draw()

    def _draw_badges(self, env):
        # Agents
        for agent in env.agents:
            self._draw_badge(agent.y, agent.x, agent.max_load - sum(agent.bag))

    def _draw_badge(self, row, col, id):
        resolution = 6
        radius = self.grid_size / 5

        badge_x = col * self.grid_size + (3 / 4) * self.grid_size
        badge_y = (
            self.grid_height
            - self.grid_size * (row + 1)
            + (1 / 4) * self.grid_size
        )

        # make a circle
        verts = []
        for i in range(resolution):
            angle = 2 * math.pi * i / resolution
            x = self.left_padding + radius * math.cos(angle) + badge_x
            y = self.bottom_padding + radius * math.sin(angle) + badge_y
            verts += [(x, y)]
        lines = []
        batch = pyglet.graphics.Batch()
        for i in range(len(verts) // 2):
            lines.append(
                pyglet.shapes.Polygon(*verts, batch=batch, color=_BLACK)
            )
        _ = pyglet.text.Label(
            str(id),
            font_name="Times New Roman",
            font_size=12,
            x=self.left_padding + badge_x,
            y=self.bottom_padding + badge_y + 2,
            anchor_x="center",
            anchor_y="center",
            batch=batch,
        )
        batch.draw()
