# Copyright 2020 The OATomobile Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Implements the deep imitative model-based agent."""

from typing import Mapping

import numpy as np
import scipy.interpolate
import torch
import carla

import oatomobile
from oatomobile.baselines.base import SetPointAgent
from oatomobile.baselines.torch.dim.model import ImitativeModel


class DIMAgent(SetPointAgent):
    """The deep imitative model agent."""

    def __init__(self, environment: oatomobile.envs.CARLAEnv, *,
                 model: ImitativeModel, **kwargs) -> None:
        """Constructs a deep imitation model agent.

        Args:
          environment: The navigation environment to spawn the agent.
          model: The deep imitative model.
        """
        super(DIMAgent, self).__init__(environment=environment, **kwargs)

        # Determines device, accelerator.
        self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # pylint: disable=no-member
        self._model = model.to(self._device)

    def __call__(self, observation: Mapping[str, np.ndarray],
                 **kwargs) -> np.ndarray:
        """Returns the imitative prior."""

        # Prepares observation for the neural-network.
        observation["overhead_features"] = observation[
            "bird_view_camera_cityscapes"]
        for attr in observation:
            if not isinstance(observation[attr], np.ndarray):
                observation[attr] = np.atleast_1d(observation[attr])
            observation[attr] = observation[attr][None, ...].astype(np.float32)
        # Makes `goal` 2D.
        observation["goal"] = observation["goal"][..., :2]
        # Convert image to CHW.
        observation["lidar"] = np.transpose(observation["lidar"], (0, 3, 1, 2))
        # Processes observations for the `ImitativeModel`.
        observation = {
            key: torch.from_numpy(tensor).to(self._device)  # pylint: disable=no-member
            for (key, tensor) in observation.items()
        }

        observation = self._model.transform(observation)

        # Queries model.
        plan = self._model(num_steps=kwargs.get("num_steps", 20),
                           epsilon=kwargs.get("epsilon", 1.0),
                           lr=kwargs.get("lr", 5e-2),
                           **observation).detach().cpu().numpy()[0]  # [T, 2]

        # TODO(filangel): clean API.
        # Interpolates plan.
        player_future_length = 40
        increments = player_future_length // plan.shape[0]
        time_index = list(range(0, player_future_length, increments))  # [T]
        plan_interp = scipy.interpolate.interp1d(x=time_index, y=plan, axis=0)
        xy = plan_interp(np.arange(0, time_index[-1]))

        # Appends z dimension.
        z = np.zeros(shape=(xy.shape[0], 1))

        # if observation['is_at_traffic_light'] and observation['traffic_light_state'] == carla.TrafficLightState.Red:
        #    xy = np.zeros_like(xy)

        return np.c_[xy, z], -1
