# TODO Workaround to cyclic import
# HELP NEEDED HERE
# This file is normally inside the nn/ directory. In the precision.py file
# when we import nn, the nn/__init__.py gets called. Currently all other classes
# (ie Conv2d, AvgPool2d, etc) behave well when imported in the nn/__init__.py
# But due to the RNN Classes requiring the import of FixedPrecisionTensor (line 14)
# and precision.py requiring the import of nn module, there is a problem.
import numpy as np
import torch
from torch import nn
from torch.nn import init

from syft.frameworks.torch.tensors.interpreters.additive_shared import AdditiveSharingTensor
from syft.frameworks.torch.tensors.interpreters.precision import FixedPrecisionTensor
from syft.generic.pointers.pointer_tensor import PointerTensor


class RNNCellBase(nn.Module):
    """
    Cell to be used as base for all RNN cells, including GRU and LSTM
    This class overrides the torch.nn.RNNCellBase
    Only Linear and Dropout layers are used to be able to use MPC
    """

    def __init__(self, input_size, hidden_size, bias, num_chunks, nonlinearity=None):
        super(RNNCellBase, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.bias = bias
        self.num_chunks = num_chunks
        self.nonlinearity = nonlinearity
        self.fc_xh = nn.Linear(input_size, self.num_chunks * hidden_size, bias=bias)
        self.fc_hh = nn.Linear(hidden_size, self.num_chunks * hidden_size, bias=bias)
        self.reset_parameters()

    def reset_parameters(self):
        """
        This method initializes or reset all the parameters of the cell.
        The paramaters are initiated following a uniform distribution.
        """
        std = 1.0 / np.sqrt(self.hidden_size)
        for w in self.parameters():
            init.uniform_(w, -std, std)

    def init_hidden(self, input):
        """
        This method initializes a hidden state when no hidden state is provided
        in the forward method. It creates a hidden state with zero values.
        """
        h = torch.zeros(input.shape[0], self.hidden_size, dtype=input.dtype, device=input.device)
        if input.has_child() and isinstance(input.child, PointerTensor):
            h = h.send(input.child.location)
        if input.has_child() and isinstance(input.child, FixedPrecisionTensor):
            h = h.fix_precision()
            child = input.child
            if isinstance(child.child, AdditiveSharingTensor):
                crypto_provider = child.child.crypto_provider
                owners = child.child.locations
                h = h.share(*owners, crypto_provider=crypto_provider)
        return h


class RNNCell(RNNCellBase):
    """
    Python implementation of RNNCell with tanh or relu non-linearity for MPC
    This class overrides the torch.nn.RNNCell
    """

    def __init__(self, input_size, hidden_size, bias=True, nonlinearity="tanh"):
        super(RNNCell, self).__init__(input_size, hidden_size, bias, num_chunks=1)

        if nonlinearity == "tanh":
            self.nonlinearity = torch.tanh
        elif nonlinearity == "relu":
            self.nonlinearity = torch.relu
        else:
            raise ValueError(f"Unknown nonlinearity: {nonlinearity}")

    def forward(self, x, h=None):

        if h is None:
            h = self.init_hidden(x)
        h_ = self.nonlinearity(self.fc_xh(x) + self.fc_hh(h))

        return h_


class GRUCell(RNNCellBase):
    """
    Python implementation of GRUCell for MPC
    This class overrides the torch.nn.GRUCell
    """

    def __init__(self, input_size, hidden_size, bias=True, nonlinearity=None):
        super(GRUCell, self).__init__(input_size, hidden_size, bias, num_chunks=3)

    def forward(self, x, h=None):

        if h is None:
            h = self.init_hidden(x)

        gate_x = self.fc_xh(x)
        gate_h = self.fc_hh(h)
        x_r, x_z, x_n = gate_x.chunk(self.num_chunks, 1)
        h_r, h_z, h_n = gate_h.chunk(self.num_chunks, 1)

        resetgate = torch.sigmoid(x_r + h_r)
        updategate = torch.sigmoid(x_z + h_z)
        newgate = torch.tanh(x_n + (resetgate * h_n))

        h_ = newgate + updategate * (h - newgate)

        return h_


class LSTMCell(RNNCellBase):
    """
    Python implementation of LSTMCell for MPC
    This class overrides the torch.nn.LSTMCell
    """

    def __init__(self, input_size, hidden_size, bias=True, nonlinearity=None):
        super(LSTMCell, self).__init__(input_size, hidden_size, bias, num_chunks=4)

    def reset_parameters(self):
        super(LSTMCell, self).reset_parameters()

        # Bias of forget gate should be initialize with 1 or 2
        # Ref: http://proceedings.mlr.press/v37/jozefowicz15.pdf
        incr_bias = 1.0 / self.hidden_size
        init.constant_(self.fc_xh.bias[self.hidden_size : 2 * self.hidden_size], incr_bias)
        init.constant_(self.fc_hh.bias[self.hidden_size : 2 * self.hidden_size], incr_bias)

    def forward(self, x, hc=None):

        if hc is None:
            hc = (self.init_hidden(x), self.init_hidden(x))
        h, c = hc

        gate_x = self.fc_xh(x)
        gate_h = self.fc_hh(h)

        x_i, x_f, x_c, x_o = gate_x.chunk(self.num_chunks, 1)
        h_i, h_f, h_c, h_o = gate_h.chunk(self.num_chunks, 1)

        inputgate = torch.sigmoid(x_i + h_i)
        forgetgate = torch.sigmoid(x_f + h_f)
        cellgate = torch.tanh(x_c + h_c)
        outputgate = torch.sigmoid(x_o + h_o)

        c_ = torch.mul(forgetgate, c) + torch.mul(inputgate, cellgate)

        h_ = torch.mul(outputgate, torch.tanh(c_))

        return h_, c_


class RNNBase(nn.Module):
    """
    Module to be used as base for all RNN modules, including GRU and LSTM
    This class overrides the torch.nn.RNNBase
    Only Linear and Dropout layers are used to be able to use MPC
    """

    def __init__(
        self,
        input_size,
        hidden_size,
        num_layers,
        bias,
        batch_first,
        dropout,
        bidirectional,
        base_cell,
        nonlinearity=None,
    ):
        super(RNNBase, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.bias = bias
        self.batch_first = batch_first
        self.dropout = float(dropout)
        self.bidirectional = bidirectional
        self.num_directions = 2 if bidirectional else 1
        self.is_lstm = base_cell is LSTMCell
        self.nonlinearity = nonlinearity

        # Dropout layers
        # TODO: implement a nn.Dropout class for PySyft
        # Link to issue: https://github.com/OpenMined/PySyft/issues/2500

        # Build RNN layers
        self.rnn_forward = nn.ModuleList()
        for layer in range(self.num_layers):
            if layer == 0:
                self.rnn_forward.append(base_cell(input_size, hidden_size, bias, nonlinearity))
            else:
                self.rnn_forward.append(base_cell(hidden_size, hidden_size, bias, nonlinearity))

        if self.bidirectional:
            self.rnn_backward = nn.ModuleList()
            for layer in range(self.num_layers):
                if layer == 0:
                    self.rnn_backward.append(base_cell(input_size, hidden_size, bias, nonlinearity))
                else:
                    self.rnn_backward.append(
                        base_cell(hidden_size, hidden_size, bias, nonlinearity)
                    )

    def forward(self, x, h=None):
        # If batch_first == True, swap axis with seq_len
        # At the end of the process we swap it back to the original structure
        if self.batch_first:
            x, h = self._swap_axis(x, h)

        # If it is LSTM, get hidden and cell states
        if h is not None and self.is_lstm:
            h, c = h

        batch_size = x.shape[1]
        seq_len = x.shape[0]

        # Initiate states if needed
        if h is None:
            h = self._init_hidden(x)
            c = self._init_hidden(x) if self.is_lstm else None

        # If bidirectional==True, split states in two, each one for each direction
        if self.bidirectional:
            h = h.contiguous().view(self.num_layers, 2, batch_size, self.hidden_size)
            h_for = h[:, 0, :, :]
            h_back = h[:, 1, :, :]
            if self.is_lstm:
                c = c.contiguous().view(self.num_layers, 2, batch_size, self.hidden_size)
                c_for = c[:, 0, :, :]
                c_back = c[:, 1, :, :]
            else:
                c_for = c_back = None

        else:
            h_for = h
            c_for = c if self.is_lstm else None

        # Run through rnn in the forward direction
        output = x.new(seq_len, batch_size, self.hidden_size).zero_()
        for t in range(seq_len):
            h_for, c_for = self._apply_time_step(x, h_for, c_for, t)
            output[t, :, :] = h_for[-1, :, :]

        hidden = h_for
        cell = c_for  # None if it is not an LSTM

        # Run through rnn in the backward direction if bidirectional==True
        if self.bidirectional:
            output_back = x.new(seq_len, batch_size, self.hidden_size).zero_()
            for t in range(seq_len - 1, -1, -1):
                h_back, c_back = self._apply_time_step(x, h_back, c_back, t, reverse_direction=True)
                output_back[t, :, :] = h_back[-1, :, :]

            # Concatenate both directions
            output = torch.cat((output, output_back), dim=-1)
            hidden = torch.cat((hidden, h_back), dim=0)
            if self.is_lstm:
                cell = torch.cat((cell, c_back), dim=0)

        # If batch_first == True, swap axis back to get original structure
        if self.batch_first:
            output = torch.transpose(output, 0, 1)
            hidden = torch.transpose(hidden, 0, 1)
            if self.is_lstm:
                cell = torch.transpose(cell, 0, 1)

        hidden = (hidden, cell) if self.is_lstm else hidden

        return output, hidden

    def _swap_axis(self, x, h):
        """
        This method swap the axes for batch_size and seq_len. It is used when
        batch_first==True.
        """
        x = torch.transpose(x, 0, 1)
        if h is not None:
            if self.is_lstm:
                h, c = h
                h = torch.transpose(h, 0, 1)
                c = torch.transpose(c, 0, 1)
                h = (h, c)
            else:
                h = torch.transpose(h, 0, 1)
        return x, h

    def _init_hidden(self, input):
        """
        This method initializes a hidden state when no hidden state is provided
        in the forward method. It creates a hidden state with zero values for each
        layer of the network.
        """
        h = torch.zeros(
            self.num_layers * self.num_directions,
            input.shape[1],
            self.hidden_size,
            dtype=input.dtype,
            device=input.device,
        )
        if input.has_child() and isinstance(input.child, PointerTensor):
            h = h.send(input.child.location)
        if input.has_child() and isinstance(input.child, FixedPrecisionTensor):
            h = h.fix_precision()
            child = input.child
            if isinstance(child.child, AdditiveSharingTensor):
                crypto_provider = child.child.crypto_provider
                owners = child.child.locations
                h = h.share(*owners, crypto_provider=crypto_provider)
        return h

    def _apply_time_step(self, x, h, c, t, reverse_direction=False):
        """
        Apply RNN layers at time t, given input and previous hidden states
        """
        rnn_layers = self.rnn_backward if reverse_direction else self.rnn_forward

        h_next = h.new(h.shape).zero_()
        c_next = c.new(c.shape).zero_() if self.is_lstm else None

        for layer in range(self.num_layers):
            if layer == 0:
                if self.is_lstm:
                    h_next[layer, :, :], c_next[layer, :, :] = rnn_layers[layer](
                        x[t, :, :], (h[layer, :, :], c[layer, :, :])
                    )
                else:
                    h_next[layer, :, :] = rnn_layers[layer](x[t, :, :], h[layer, :, :])
            else:
                if self.is_lstm:
                    h_next[layer, :, :], c_next[layer, :, :] = rnn_layers[layer](
                        h_next[layer - 1, :, :].clone(), (h[layer, :, :], c[layer, :, :])
                    )
                else:
                    h_next[layer, :, :] = rnn_layers[layer](
                        h_next[layer - 1, :, :].clone(), h[layer, :, :]
                    )

        return h_next, c_next


class RNN(RNNBase):
    """
    Python implementation of RNN for MPC
    This class overrides the torch.nn.RNN
    """

    def __init__(
        self,
        input_size,
        hidden_size,
        num_layers=1,
        nonlinearity="tanh",
        bias=True,
        batch_first=False,
        dropout=0,
        bidirectional=False,
    ):

        super(RNN, self).__init__(
            input_size,
            hidden_size,
            num_layers,
            bias,
            batch_first,
            dropout,
            bidirectional,
            RNNCell,
            nonlinearity,
        )


class GRU(RNNBase):
    """
    Python implementation of GRU for MPC
    This class overrides the torch.nn.GRU
    """

    def __init__(
        self,
        input_size,
        hidden_size,
        num_layers=1,
        bias=True,
        batch_first=False,
        dropout=0,
        bidirectional=False,
    ):

        super(GRU, self).__init__(
            input_size, hidden_size, num_layers, bias, batch_first, dropout, bidirectional, GRUCell
        )


class LSTM(RNNBase):
    """
    Python implementation of LSTM for MPC
    This class overrides the torch.nn.LSTM
    """

    def __init__(
        self,
        input_size,
        hidden_size,
        num_layers=1,
        bias=True,
        batch_first=False,
        dropout=0,
        bidirectional=False,
    ):

        super(LSTM, self).__init__(
            input_size, hidden_size, num_layers, bias, batch_first, dropout, bidirectional, LSTMCell
        )
