# Copyright 2018 Amazon.com, Inc. or its affiliates. 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.
# A copy of the License is located at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# or in the "license" file accompanying this file. This file 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.

from math import pi

import torch
from torch import nn as nn


class ImplicitQuantileModule(nn.Module):
    """
    See arXiv: 1806.06923 This module, in combination with quantile loss,
    learns how to generate the quantile of the distribution of the target.

    A quantile value, tau, is randomly generated with a Uniform([0, 1])). This
    quantile value is embedded in this module and also passed to the quantile
    loss: this should force the model to learn the appropriate quantile.
    """

    def __init__(self, in_features, output_domain_cls):
        super(ImplicitQuantileModule, self).__init__()
        self.in_features = in_features
        self.quantile_layer = QuantileLayer(in_features)
        self.output_layer = nn.Sequential(
            nn.Linear(in_features, in_features),
            nn.Softplus(),
            nn.Linear(in_features, 1),
            output_domain_cls(),
        )

    def forward(self, input_data, tau):
        embedded_tau = self.quantile_layer(tau)
        new_input_data = input_data * (
            torch.ones_like(embedded_tau) + embedded_tau
        )
        return self.output_layer(new_input_data).squeeze(-1)


class QuantileLayer(nn.Module):
    """
    Define quantile embedding layer.

    i.e. phi in the IQN paper (arXiv: 1806.06923)
    """

    def __init__(self, num_output):
        super(QuantileLayer, self).__init__()
        self.n_cos_embedding = 64
        self.num_output = num_output
        self.output_layer = nn.Sequential(
            nn.Linear(self.n_cos_embedding, self.n_cos_embedding),
            nn.PReLU(),
            nn.Linear(self.n_cos_embedding, num_output),
        )

    def forward(self, tau):
        cos_embedded_tau = self.cos_embed(tau)
        final_output = self.output_layer(cos_embedded_tau)
        return final_output

    def cos_embed(self, tau):
        integers = torch.repeat_interleave(
            torch.arange(0, self.n_cos_embedding).unsqueeze(dim=0),
            repeats=tau.shape[-1],
            dim=0,
        ).to(tau.device)
        return torch.cos(pi * tau.unsqueeze(dim=-1) * integers)
