#  Copyright (c) 2022.
#  ProrokLab (https://www.proroklab.org/)
#  All rights reserved.
import typing
from typing import Dict, Callable, List

import torch
from torch import Tensor

from vmas import render_interactively
from vmas.simulator.core import Agent, Landmark, World, Sphere, Entity
from vmas.simulator.scenario import BaseScenario
from vmas.simulator.utils import Color, ScenarioUtils

if typing.TYPE_CHECKING:
    from vmas.simulator.rendering import Geom


class Scenario(BaseScenario):
    def make_world(self, batch_dim: int, device: torch.device, **kwargs):
        self.n_agents = kwargs.get("n_agents", 4)
        self.lidar_range = kwargs.get("lidar_range", 0.35)
        self.agent_radius = kwargs.get("agent_radius", 0.1)
        self.comms_range = kwargs.get("comms_range", 0.55)

        self.shared_rew = kwargs.get("shared_rew", True)
        self.pos_shaping_factor = kwargs.get("pos_shaping_factor", 1)
        self.final_reward = kwargs.get("final_reward", 0.01)

        self.agent_collision_penalty = kwargs.get("agent_collision_penalty", -1)

        known_colors = [
            Color.GREEN,
            Color.RED,
            Color.GRAY,
            Color.WHITE,
            Color.BLUE,
        ]
        colors = torch.randn(
            (max(self.n_agents - len(known_colors), 0), 3), device=device
        )

        self.min_istance_between_entities = self.agent_radius * 2 + 0.05
        self.world_semidim = 1
        self.min_collision_distance = 0.005

        entity_filter_agents: Callable[[Entity], bool] = lambda e: isinstance(e, Agent)

        # Make world
        world = World(batch_dim, device, substeps=2)

        # Add agents
        for i in range(self.n_agents):
            color = (
                known_colors[i]
                if i < len(known_colors)
                else colors[i - len(known_colors)]
            )

            # Constraint: all agents have same action range and multiplier
            agent = Agent(
                name=f"agent {i}",
                collide=True,
                color=color,
                alpha=1,
                shape=Sphere(radius=self.agent_radius),
                # sensors=[
                #     Lidar(
                #         world,
                #         angle_start=0.05,
                #         angle_end=2 * torch.pi + 0.05,
                #         n_rays=12,
                #         max_range=self.lidar_range,
                #         entity_filter=entity_filter_agents,
                #     ),
                # ],
            )
            world.add_agent(agent)

            # Add goals
            goal = Landmark(
                name=f"goal {i}",
                collide=False,
                color=color,
                shape=Sphere(radius=0.04),
            )
            world.add_landmark(goal)
            agent.goal = goal

        return world

    def reset_world_at(self, env_index: int = None):
        ScenarioUtils.spawn_entities_randomly(
            self.world.agents + self.world.landmarks,
            self.world,
            env_index,
            self.min_istance_between_entities,
            (-self.world_semidim, self.world_semidim),
            (-self.world_semidim, self.world_semidim),
        )
        for i, landmark in enumerate(self.world.landmarks):
            # landmark.set_pos(
            #     torch.zeros(
            #         (1, self.world.dim_p)
            #         if env_index is not None
            #         else (self.world.batch_dim, self.world.dim_p),
            #         device=self.world.device,
            #         dtype=torch.float32,
            #     ).uniform_(
            #         -self.world_semidim,
            #         self.world_semidim,
            #     ),
            #     batch_index=env_index,
            # )

            if env_index is None:
                self.world.agents[i].pos_shaping = (
                    torch.linalg.vector_norm(
                        self.world.agents[i].state.pos
                        - self.world.agents[i].goal.state.pos,
                        dim=1,
                    )
                    * self.pos_shaping_factor
                )
            else:
                self.world.agents[i].pos_shaping[env_index] = (
                    torch.linalg.vector_norm(
                        self.world.agents[i].state.pos[env_index]
                        - self.world.agents[i].goal.state.pos[env_index]
                    )
                    * self.pos_shaping_factor
                )

    def reward(self, agent: Agent):
        is_first = agent == self.world.agents[0]

        if is_first:
            self.pos_rew = torch.zeros(
                self.world.batch_dim, device=self.world.device, dtype=torch.float32
            )
            self.final_rew = torch.zeros(self.world.batch_dim, device=self.world.device)

            for a in self.world.agents:
                self.pos_rew += self.agent_reward(a)
                a.agent_collision_rew = torch.zeros(
                    (self.world.batch_dim,), device=self.world.device
                )

            self.all_goal_reached = torch.all(
                torch.stack([a.on_goal for a in self.world.agents], dim=-1), dim=-1
            )

            self.final_rew[self.all_goal_reached] = self.final_reward

            for i, a in enumerate(self.world.agents):
                for j, b in enumerate(self.world.agents):
                    if i <= j:
                        continue
                    if self.world.collides(a, b):
                        distance = self.world.get_distance(a, b)
                        a.agent_collision_rew[
                            distance <= self.min_collision_distance
                        ] += self.agent_collision_penalty
                        b.agent_collision_rew[
                            distance <= self.min_collision_distance
                        ] += self.agent_collision_penalty

        pos_reward = self.pos_rew if self.shared_rew else agent.pos_rew
        return pos_reward + self.final_rew + agent.agent_collision_rew

    def agent_reward(self, agent: Agent):
        agent.distance_to_goal = torch.linalg.vector_norm(
            agent.state.pos - agent.goal.state.pos,
            dim=-1,
        )
        agent.on_goal = agent.distance_to_goal < agent.goal.shape.radius

        pos_shaping = agent.distance_to_goal * self.pos_shaping_factor
        agent.pos_rew = agent.pos_shaping - pos_shaping
        agent.pos_shaping = pos_shaping
        return agent.pos_rew

    def observation(self, agent: Agent):
        return torch.cat(
            [
                agent.state.pos,
                agent.state.vel,
                agent.state.pos - agent.goal.state.pos,
                # agent.sensors[0].measure(),
            ],
            dim=-1,
        )

    def info(self, agent: Agent) -> Dict[str, Tensor]:
        return {
            "pos_rew": self.pos_rew if self.shared_rew else agent.pos_rew,
            "final_rew": self.final_rew,
            "agent_collisions": agent.agent_collision_rew,
        }

    def extra_render(self, env_index: int = 0) -> "List[Geom]":
        from vmas.simulator import rendering

        geoms: List[Geom] = []

        # Communication lines
        for i, agent1 in enumerate(self.world.agents):
            for j, agent2 in enumerate(self.world.agents):
                if j <= i:
                    continue
                agent_dist = torch.linalg.vector_norm(
                    agent1.state.pos - agent2.state.pos, dim=-1
                )
                if agent_dist[env_index] <= self.comms_range:
                    color = Color.BLACK.value
                    line = rendering.Line(
                        (agent1.state.pos[env_index]),
                        (agent2.state.pos[env_index]),
                        width=1,
                    )
                    xform = rendering.Transform()
                    line.add_attr(xform)
                    line.set_color(*color)
                    geoms.append(line)

        return geoms


if __name__ == "__main__":
    render_interactively(
        __file__,
        control_two_agents=True,
    )
