# MIT License
#
# Copyright (C) 2021. Huawei Technologies Co., Ltd. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
import argparse
import copy
import glob
import json
import os
import random
import re
import timeit
from ast import literal_eval
from collections import defaultdict

import dill
import gym
import numpy as np
import ray
import torch
from matplotlib import pyplot as plt

from smarts.core.agent import Agent, AgentSpec
from smarts.core.agent_interface import AgentInterface, AgentType, NeighborhoodVehicles
from smarts.core.controllers import ActionSpaceType
from ultra.scenarios.common.social_vehicle_definitions import get_social_vehicle_color
from ultra.scenarios.common.visualization import (
    convert_to_gif,
    draw_intersection,
    profile_vehicles,
)


class DefaultPolicy(Agent):
    def act(self, obs):
        lane_index = 0
        num_trajectory_points = min([10, len(obs.waypoint_paths[lane_index])])
        desired_speed = 20
        trajectory = [
            [
                obs.waypoint_paths[lane_index][i].pos[0]
                for i in range(num_trajectory_points)
            ],
            [
                obs.waypoint_paths[lane_index][i].pos[1]
                for i in range(num_trajectory_points)
            ],
            [
                obs.waypoint_paths[lane_index][i].heading
                for i in range(num_trajectory_points)
            ],
            [desired_speed for i in range(num_trajectory_points)],
        ]
        return trajectory


class BaseAnalysis:
    def __init__(self):
        self.social_vehicles_states = {}
        self.social_vehicles_ids = set()
        self.finished_social_vehicles_ids = set()
        self.analysis = {}

    def reset_scenario_cache(self):
        self.social_vehicles_ids = {}
        self.finished_social_vehicles_ids = set()
        self.social_vehicles_states = defaultdict(
            lambda: {
                "position": [0, 0],
                "route": (),
                "in_junction": [],
                "behavior": None,
                "steps": 0,
                "stop_step": 0,
                "start_step": None,
                "end_step": None,
                "edge": None,
                "speeds": [0],
                "accels": [0],
            }
        )

    def save_histogram(self, data, figure_name, title, x_range=None):
        plt.figure()

        n, bins, patches = plt.hist(
            x=data, bins=600, color="#0504aa", alpha=0.7, rwidth=0.85
        )
        plt.grid(axis="y", alpha=0.75)
        plt.xlabel(title)
        plt.ylabel("Frequency")
        maxfreq = n.max()
        # Set a clean upper y-axis limit.
        plt.ylim(ymax=np.ceil(maxfreq / 10) * 10 if maxfreq % 10 else maxfreq + 10)
        if x_range is not None:
            plt.xlim(x_range)

        plt.savefig(figure_name)
        plt.close()

    def get_agent(self, ego, policy, max_episode_steps):
        observation_adapter = None
        if ego:
            config = get_agent_config_by_type(policy)
            agent_spec = AgentSpec(
                interface=config["interface"],
                policy_params=dict(config["policy"], checkpoint_dir=ego_model),
                policy_builder=config["policy_class"],
            )
            agent_spec.interface.max_episode_steps = max_episode_steps
            observation_adapter = IntersectionAdapter(
                agent_id="AGENT_007",
                social_vehicle_config=config["social_vehicle_config"],
                timestep_sec=config["env"]["timestep_sec"],
                **config["other"],
            )
        else:
            # Lane Following agent
            agent_spec = AgentSpec(
                interface=AgentInterface(
                    max_episode_steps=max_episode_steps,  # 10 mins
                    waypoints=True,
                    action=ActionSpaceType.Lane,
                    debug=False,
                    neighborhood_vehicles=NeighborhoodVehicles(radius=2000),
                ),
                agent_builder=DefaultPolicy,
            )
        return agent_spec, observation_adapter

    def process_social_vehicles(self, vehicles, dt, step):
        stopwatcher_state = None
        stopwatcher_exit = False
        current_vehicles = set()
        for v in vehicles:
            v_id = v.id
            is_new = False
            behavior_key, _ = get_social_vehicle_color(v.id)

            if v_id not in self.social_vehicles_ids:
                self.social_vehicles_ids[v_id] = len(self.social_vehicles_ids)
                is_new = True
            current_vehicles.add(self.social_vehicles_ids[v_id])
            # -----------------------------------
            # 1- Find route information
            # -----------------------------------
            _route = tuple(re.findall("edge-((....|.....)-..)", v_id))
            route = (_route[0][0], _route[1][0])

            previous_edge = self.social_vehicles_states[self.social_vehicles_ids[v_id]][
                "edge"
            ]
            current_edge = re.findall("edge-(....|.....|junction*)-", v.edge_id)
            current_edge = current_edge[0] if len(current_edge) == 1 else None

            # -----------------------------------
            # 2- start caching
            # -----------------------------------
            if is_new:
                self.social_vehicles_states[self.social_vehicles_ids[v_id]][
                    "start_step"
                ] = step
                self.social_vehicles_states[self.social_vehicles_ids[v_id]][
                    "behavior"
                ] = behavior_key
                self.social_vehicles_states[self.social_vehicles_ids[v_id]][
                    "route"
                ] = route

            # -----------------------------------
            # 3- apply updates
            # -----------------------------------
            # update position
            self.social_vehicles_states[self.social_vehicles_ids[v_id]]["position"] = [
                v.position[0],
                v.position[1],
            ]

            # stopped time
            if v.speed < 0.01:
                self.social_vehicles_states[self.social_vehicles_ids[v_id]][
                    "stop_step"
                ] += 1

            # collect speed and accel:
            accel = (
                v.speed
                - self.social_vehicles_states[self.social_vehicles_ids[v_id]]["speeds"][
                    -1
                ]
            ) / dt
            self.social_vehicles_states[self.social_vehicles_ids[v_id]][
                "speeds"
            ].append(v.speed)
            self.social_vehicles_states[self.social_vehicles_ids[v_id]][
                "accels"
            ].append(accel)

            # in junction time range
            if "junction" in v.edge_id:
                n = len(
                    self.social_vehicles_states[self.social_vehicles_ids[v_id]][
                        "speeds"
                    ]
                )
                if not self.social_vehicles_states[self.social_vehicles_ids[v_id]][
                    "in_junction"
                ]:
                    self.social_vehicles_states[self.social_vehicles_ids[v_id]][
                        "in_junction"
                    ] = [n, n]
                else:
                    self.social_vehicles_states[self.social_vehicles_ids[v_id]][
                        "in_junction"
                    ][0] = min(
                        self.social_vehicles_states[self.social_vehicles_ids[v_id]][
                            "in_junction"
                        ][0],
                        n,
                    )
                    self.social_vehicles_states[self.social_vehicles_ids[v_id]][
                        "in_junction"
                    ][1] = max(
                        self.social_vehicles_states[self.social_vehicles_ids[v_id]][
                            "in_junction"
                        ][1],
                        n,
                    )

            self.social_vehicles_states[self.social_vehicles_ids[v_id]]["steps"] += 1
            if current_edge:
                self.social_vehicles_states[self.social_vehicles_ids[v_id]][
                    "edge"
                ] = current_edge

            # -------------------------------------
            # 3- end_time
            # -----------------------------------
            self.social_vehicles_states[self.social_vehicles_ids[v_id]][
                "end_step"
            ] = step

            # check for stopwatcher vehicle
            if (
                "stopwatcher"
                in self.social_vehicles_states[self.social_vehicles_ids[v_id]][
                    "behavior"
                ]
            ):
                stopwatcher_state = self.social_vehicles_states[
                    self.social_vehicles_ids[v_id]
                ]

        # cleanup expired vehicles
        for v in list(self.social_vehicles_states.keys()):
            if v not in current_vehicles:
                self.finished_social_vehicles_ids.add(v)
                if "stopwatcher" in self.social_vehicles_states[v]["behavior"]:
                    stopwatcher_exit = True

        return stopwatcher_state, stopwatcher_exit

    # Number of GPUs should be splited between remote functions.
    def run(
        self,
        scenarios,
        timestep_sec,
        ego,
        policy,
        max_episode_steps,
        video_rate,
        custom_traci_functions,
        end_by_stopwatcher,
        analyze_func,
        init_time_skip=20,
    ):
        agent_spec, observation_adapter = self.get_agent(ego, policy, max_episode_steps)
        agent_id = "AGENT_007"
        ego_less = False if ego else True
        env = gym.make(
            "smarts.env:hiway-v0",
            scenarios=scenarios,
            agent_specs={agent_id: agent_spec},
            headless=True,
            visdom=False,
            timestep_sec=timestep_sec,
            sumo_headless=True,
            seed=0,
            endless_traffic=False,
        )
        visited_scenario = set()
        while len(visited_scenario) < len(scenarios):
            agent = agent_spec.build_agent()
            observations = env.reset()
            scenario_path = env._smarts.scenario.root_filepath

            # only run scenario once
            if scenario_path not in visited_scenario:
                with open(f"{scenario_path}/metadata.json", "r") as metadata_rd:
                    metadata = json.load(metadata_rd)
                visited_scenario.add(scenario_path)
                all_waypoints = [
                    [linked_wp.wp.pos[0], linked_wp.wp.pos[1]]
                    for linked_wp in env._smarts.waypoints._linked_waypoints
                ]
                dones = {"__all__": False}
                intersection_name, intersection_tag = (
                    metadata["intersection"]["name"],
                    metadata["intersection"]["type"],
                )
                total_vehicles = int(metadata["total_vehicles"])
                print(
                    f"Processing Scenario {len(visited_scenario)}/{len(scenarios)}: {scenario_path} "
                )
                print(f"Total number of vehicles: {total_vehicles}")
                print(f"Running egoless: {ego_less}")

                self.reset_scenario_cache()
                step, episode_time = 0, 0.0
                simulation_start = timeit.default_timer()
                stopwatcher_exit = False
                stopwatcher_logged = False
                stopwatcher_max_steps = 0
                images, simulation_data = [], {}
                while not dones["__all__"]:
                    if step % 200 == 0:
                        print("step", step)
                    # todo ------ add ego agent
                    agent_obs = observations[agent_id]
                    start = agent_obs.ego_vehicle_state.mission.start
                    goal = agent_obs.ego_vehicle_state.mission.goal
                    path = agent_obs.waypoint_paths[0]
                    agent_action = "slow_down"
                    observations, rewards, dones, infos = env.step(
                        {agent_id: agent_action}
                    )
                    # ------------
                    (
                        stopwatcher_state,
                        stopwatcher_exit,
                    ) = self.process_social_vehicles(
                        agent_obs.neighborhood_vehicle_states,
                        timestep_sec,
                        step,
                    )
                    if stopwatcher_state:
                        has_stopwatcher = True
                        stopwatcher_max_steps = max(
                            stopwatcher_max_steps, stopwatcher_state["steps"]
                        )
                        if not stopwatcher_logged:
                            print("stopwatcher detected!")
                            stopwatcher_logged = True
                    if (
                        step > (init_time_skip / timestep_sec)
                        and step % int(1 / timestep_sec) == 0
                    ):
                        for func in custom_traci_functions:
                            func(
                                env._smarts.traffic_sim._traci_conn,
                                simulation_data,
                                last_step=step,
                            )

                    step += 1
                    episode_time += timestep_sec

                    if len(visited_scenario) % video_rate == 0:
                        images.append(
                            draw_intersection(
                                ego_position=agent_obs.ego_vehicle_state.position,
                                social_vehicle_states=agent_obs.neighborhood_vehicle_states,
                                goal_path=path,
                                all_waypoints=all_waypoints,
                                step=step,
                                lookaheads_positions=[
                                    waypoint.pos for waypoint in path
                                ],
                                goal=goal.position[0:2],
                                start=start.position[0:2],
                                intersection_tag=intersection_tag,
                                finished_vehicles=self.finished_social_vehicles_ids,
                            )
                        )
                    if ego_less and stopwatcher_exit and end_by_stopwatcher:
                        break
                    if len(self.finished_social_vehicles_ids) == total_vehicles:
                        break
                simulation_end = timeit.default_timer()
                if ego_less and stopwatcher_max_steps > 0:
                    print("Saved stopwatcher steps to", scenario_path)
                    with open(f"{scenario_path}/max_steps.txt", "w") as file_wr:
                        file_wr.write(str(stopwatcher_max_steps))
                simulation_time = simulation_end - simulation_start
                print(analyze_func)
                analyze_func(episode_time, stopwatcher_max_steps, simulation_data)
                print(
                    f"Episode finished in {simulation_time}s, {episode_time}s simulation"
                )

                if len(visited_scenario) % video_rate == 0:
                    print("Converting images to video...")
                    convert_to_gif(
                        images=images,
                        save_dir=scenario_path,
                        name=f"{intersection_name}-{metadata['flow_id']}",
                    )

                print("--------------------------------")

        env.close()
