# 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 torch
from torch import nn


class Flatten(nn.Module):
    def forward(self, x):
        x = x.view(x.size()[0], -1)
        return x


class DQNCNN(nn.Module):
    def __init__(
        self,
        n_in_channels,
        image_dim,
        num_actions,
        hidden_dim=128,
        activation=nn.ReLU,
    ):
        super(DQNCNN, self).__init__()
        self.im_feature = nn.Sequential(
            nn.Conv2d(n_in_channels, 32, 8, 4),
            activation(),
            nn.Conv2d(32, 64, 4, 2),
            activation(),
            nn.Conv2d(64, 64, 3, 1),
            activation(),
            Flatten(),
        )

        dummy = torch.zeros((1, n_in_channels, *image_dim))
        im_feature_size = self.im_feature(dummy).data.cpu().numpy().size

        self.q_outs = nn.ModuleList()
        for num_action in num_actions:
            q_out = nn.Sequential(
                nn.Linear(im_feature_size, hidden_dim),
                nn.ReLU(),
                nn.Linear(hidden_dim, num_action),
            )
            self.q_outs.append(q_out)

        self.init()

    def init(self):
        for q_out in self.q_outs:
            # nn.init.normal_(q_out[-1].weight.data, 0.0, 1e-2)
            nn.init.constant_(q_out[-1].bias.data, 0.0)

    def forward(self, state, training=False):
        x = self.im_feature(state)
        x = [e(x) for e in self.q_outs]

        if training:
            aux_losses = {}
            return x, aux_losses
        else:
            return x


class DQNFC(nn.Module):
    def __init__(
        self,
        num_actions,
        state_size,
        hidden_dim=256,
        activation=nn.ReLU,
    ):
        super(DQNFC, self).__init__()

        self.q_outs = nn.ModuleList()
        for num_action in num_actions:
            q_out = nn.Sequential(
                nn.Linear(state_size, hidden_dim),
                activation(),
                nn.Linear(hidden_dim, hidden_dim // 2),
                activation(),
                nn.Linear(hidden_dim // 2, hidden_dim // 4),
                activation(),
                nn.Linear(hidden_dim // 4, num_action),
            )
            self.q_outs.append(q_out)

        self.init()

    def init(self):
        for q_out in self.q_outs:
            nn.init.constant_(q_out[-1].bias.data, 0.0)
            # nn.init.normal_(q_out[-1].weight.data, 0.0, 1e-3)

    def forward(self, state, training=False):
        low_dim_state = state["low_dim_states"]
        if len(low_dim_state.shape) == 1:
            low_dim_state = torch.unsqueeze(low_dim_state, 0)
            unsqueezed = True
        else:
            unsqueezed = False
        x = low_dim_state
        # x = self.feature(low_dim_state)
        x = [e(x) for e in self.q_outs]
        if unsqueezed:
            x = [torch.squeeze(e, 0) for e in x]

        if training:
            aux_losses = {}
            return x, aux_losses
        else:
            return x


class DQNWithSocialEncoder(nn.Module):
    def __init__(
        self,
        num_actions,
        state_size,
        hidden_dim=256,
        activation=nn.ReLU,
        social_feature_encoder_class=None,
        social_feature_encoder_params=None,
    ):
        super(DQNWithSocialEncoder, self).__init__()

        self.social_feature_encoder = (
            social_feature_encoder_class(**social_feature_encoder_params)
            if social_feature_encoder_class
            else None
        )

        self.state_size = state_size
        self.q_outs = nn.ModuleList()
        for num_action in num_actions:
            q_out = nn.Sequential(
                nn.Linear(state_size, hidden_dim),
                activation(),
                nn.Linear(hidden_dim, hidden_dim // 2),
                activation(),
                nn.Linear(hidden_dim // 2, hidden_dim // 4),
                activation(),
                nn.Linear(hidden_dim // 4, num_action),
            )
            self.q_outs.append(q_out)

        self.init()

    def init(self):
        for q_out in self.q_outs:
            nn.init.constant_(q_out[-1].bias.data, 0.0)

    def forward(self, state, training=False):
        low_dim_state = state["low_dim_states"]
        social_vehicles_state = state["social_vehicles"]

        aux_losses = {}
        social_feature = []
        if self.social_feature_encoder is not None:
            social_feature, social_encoder_aux_losses = self.social_feature_encoder(
                social_vehicles_state, training
            )
            aux_losses.update(social_encoder_aux_losses)
        else:
            social_feature = [e.reshape(1, -1) for e in social_vehicles_state]

        social_feature = torch.cat(social_feature, 0) if len(social_feature) > 0 else []
        x = (
            torch.cat([low_dim_state, social_feature], -1)
            if len(social_feature) > 0
            else low_dim_state
        )
        x = [e(x) for e in self.q_outs]

        if training:
            return x, aux_losses
        else:
            return x
