# Copyright (c) 2021, salesforce.com, inc.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# For full license text, see the LICENSE file in the repo root
# or https://opensource.org/licenses/BSD-3-Clause

import torch
import torch.nn.functional as F
from torch import nn


class IndependentPolicyNet(nn.Module):
    """
    Represents a policy network with separate heads for different types of actions.
    Thus, the resulting policy will take the form
    $pi(a | s) = pi_1(a_1 | s) pi_2(a_2 | s)...$
    """

    def __init__(self, state_size, action_size_list, norm_consts=None):
        super().__init__()

        self.state_size = state_size
        self.action_size_list = action_size_list
        if norm_consts is not None:
            self.norm_center, self.norm_scale = norm_consts
        else:
            self.norm_center = torch.zeros(self.state_size).cuda()
            self.norm_scale = torch.ones(self.state_size).cuda()
        self.fc1 = nn.Linear(state_size, 128)
        self.fc2 = nn.Linear(128, 128)
        # policy network head
        self.action_heads = nn.ModuleList(
            [nn.Linear(128, action_size) for action_size in action_size_list]
        )
        # value network head
        self.fc4 = nn.Linear(128, 1)

    def forward(self, x):
        assert x.shape[-1] == self.state_size  # Check if the last dimension matches

        # Normalize the model input
        new_shape = tuple(1 for _ in x.shape[:-1]) + (x.shape[-1],)
        view_center = self.norm_center.view(new_shape)
        view_scale = self.norm_scale.view(new_shape)
        x = (x - view_center) / view_scale

        # Feed forward
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        probs = [F.softmax(action_head(x), dim=-1) for action_head in self.action_heads]
        vals = self.fc4(x)
        return probs, vals


class PolicyNet(nn.Module):
    """
    The policy network class to output acton probabilities and the value function.
    """

    def __init__(self, state_size, action_size, norm_consts=None):
        super().__init__()

        self.state_size = state_size
        self.action_size = action_size
        if norm_consts is not None:
            self.norm_center, self.norm_scale = norm_consts
        else:
            self.norm_center = torch.zeros(self.state_size).cuda()
            self.norm_scale = torch.ones(self.state_size).cuda()
        self.fc1 = nn.Linear(state_size, 128)
        self.fc2 = nn.Linear(128, 128)
        # policy network head
        self.fc3 = nn.Linear(128, action_size)
        # value network head
        self.fc4 = nn.Linear(128, 1)

    def forward(self, x, actions_mask=None):
        # here, the action mask should be large negative constants for actions
        # that shouldn't be allowed.
        new_shape = tuple(1 for _ in x.shape[:-1]) + (x.shape[-1],)
        view_center = self.norm_center.view(new_shape)
        view_scale = self.norm_scale.view(new_shape)
        x = (x - view_center) / view_scale
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        if actions_mask is not None:
            probs = F.softmax(self.fc3(x) + actions_mask, dim=-1)
        else:
            probs = F.softmax(self.fc3(x), dim=-1)
        vals = self.fc4(x)
        return probs, vals


class DeterministicPolicy:
    """
    A policy class that outputs deterministic actions.
    """

    def __init__(self, state_size, action_size, action_choice):
        self.state_size = state_size
        self.action_size = action_size
        self.action_choice = action_choice
        self.actions_out = torch.zeros(action_size, device="cuda")
        self.actions_out[self.action_choice] = 1.0

    def __call__(self, x, actions_mask=None):
        return self.forward(x)

    def forward(self, x):
        # output enough copies of the delta function
        # distribution of the right size given x
        x_batch_shapes = x.shape[:-1]
        repeat_vals = x_batch_shapes + (1,)
        return self.actions_out.repeat(*repeat_vals), None