import torch
import torch.nn as nn

from expRNN.exp_numpy import expm, expm_frechet
from expRNN.initialization import henaff_init
import numpy as np

import snorm
class modrelu(nn.Module):
    def __init__(self, features):
        # For now we just support square layers
        super(modrelu, self).__init__()
        self.features = features
        self.b = nn.Parameter(torch.Tensor(self.features))
        self.reset_parameters()

    def reset_parameters(self):
        self.b.data.uniform_(-0.01, 0.01)

    def forward(self, inputs):
        norm = torch.abs(inputs)
        biased_norm = norm + self.b
        magnitude = nn.functional.relu(biased_norm)
        phase = torch.sign(inputs)

        return phase * magnitude


class ExpRNN(nn.Module):
    def __init__(self, input_size, hidden_size, nonlinearity=modrelu, exponential="exact", skew_initializer=henaff_init,input_initializer=nn.init.xavier_normal_):
        super(ExpRNN, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        # Can be optimised for size but it's not really worth it in most applications
        self.log_recurrent_kernel = nn.Parameter(torch.Tensor(self.hidden_size, self.hidden_size))
        self.recurrent_kernel = nn.Parameter(torch.Tensor(self.hidden_size, self.hidden_size))
        self.input_kernel = nn.Linear(self.input_size, self.hidden_size, bias=False)
        self.gamma = nn.Parameter(torch.ones(1))
        
        self.input_initializer = input_initializer
        self.skew_initializer = skew_initializer
        self.params = [self.input_kernel.weight]
        self.orthogonal_params = [self.log_recurrent_kernel]
        self.ortho = False
        self.exact = exponential == "exact"
        if self.exact:
            self.exp = expm
        else:
            self.exp = exponential

        if nonlinearity == 'relu':
            self.nonlinearity = nn.ReLU()
        elif nonlinearity == 'modrelu':
            self.nonlinearity = modrelu(hidden_size)
            self.params.append(self.nonlinearity.b)
        elif nonlinearity == 'tanh':
            self.nonlinearity = nn.Tanh()
        elif nonlinearity == 'sigmoid':
            self.nonlinearity = nn.Sigmoid()
        else:
            self.nonlinearity = None

        self.reset_parameters()

    def _A(self, gradients):
        A = self.log_recurrent_kernel
        if not gradients:
            A = A.data
        A = A.triu(diagonal=1)
        return A - A.t()

    def _B(self, gradients):
        return self.exp(self._A(gradients))

    def _norm(self):
        norm = snorm.spectral_norm(self.recurrent_kernel.data)
        return norm

        
    def _clip2(self):
        norm = snorm.spectral_norm(self.recurrent_kernel.data)
        #print(f'spectral norm{norm}')
        eps = 1e-4
        inv_norm = 1 / norm
        min_value = 1 / (norm + eps)
        max_value = 1 / (norm - eps)
        self.gamma.data = torch.clamp(self.gamma, min=min_value, max=inv_norm)
        
    def reset_parameters(self):
        print(self.skew_initializer)
        self.log_recurrent_kernel.data = \
                torch.as_tensor(self.skew_initializer(self.hidden_size),
                        dtype=self.log_recurrent_kernel.dtype,
                        device=self.log_recurrent_kernel.device)
        self.recurrent_kernel.data = self._B(gradients=False)
        self.input_initializer(self.input_kernel.weight.data)

    def orthogonal_step(self, optim):
        if self.exact:
            A = self._A(gradients=False)
            B = self.recurrent_kernel.data
            G = self.recurrent_kernel.grad.data
            BtG = B.t().mm(G)
            grad = 0.5*(BtG - BtG.t())
            frechet_deriv = B.mm(expm_frechet(-A, grad))

            # Account for the triangular representation of the skew-symmetric matrix
            # The gradient with respect to the parameter x_i,j is the upper part of the triangular matrix
            # minus the lower part of the triangular matrix.
            # The gradient needn't be triangular, but we make it triangular for consistency
            self.log_recurrent_kernel.grad = (frechet_deriv - frechet_deriv.t()).triu(diagonal=1)
        else:
            orth_param = self._B(gradients=True)
            self.log_recurrent_kernel.grad = \
                torch.autograd.grad([orth_param], self.log_recurrent_kernel,
                                    grad_outputs=(self.recurrent_kernel.grad,))[0]
        optim.step()
        self.recurrent_kernel.data = self._B(gradients=False)
        self.recurrent_kernel.grad.data.zero_()

    def forward(self, input, hidden=None): 
        if hidden is None:
            hidden = input.new_zeros(input.size(0), self.hidden_size, requires_grad=True)
            #init = torch.mm(self.input_kernel(input), torch.inverse(self.recurrent_kernel.data - torch.eye(self.recurrent_kernel.data.size(0) , device= input.device)))
            #hidden = init.detach().clone() 
            #hidden.requires_grad=True
            self.first_hidden = hidden

        input = self.input_kernel(input)

        hidden = hidden.matmul(self.recurrent_kernel)

        out = input + hidden

        if self.nonlinearity:
            return self.nonlinearity(out)
        else:
            return out

    def extra_repr(self):
        return 'input_size={}, hidden_size={}'.format(self.input_size, self.hidden_size)
class ExpRNN2(nn.Module):
    def __init__(self, input_size, hidden_size, slen, nonlinearity=modrelu, exponential="exact", skew_initializer=henaff_init,input_initializer=nn.init.xavier_normal_):
        super(ExpRNN2, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        # Can be optimised for size but it's not really worth it in most applications
        self.log_recurrent_kernel = nn.Parameter(torch.Tensor(self.hidden_size, self.hidden_size))
        self.recurrent_kernel = nn.Parameter(torch.Tensor(self.hidden_size, self.hidden_size))
        self.input_kernel = nn.Linear(self.input_size, self.hidden_size, bias=False)
        self.gamma = nn.Parameter(torch.rand(slen))
        
        self.input_initializer = input_initializer
        self.skew_initializer = skew_initializer
        self.params = [self.input_kernel.weight]
        self.orthogonal_params = [self.log_recurrent_kernel]
        self.ortho = False
        self.exact = exponential == "exact"
        if self.exact:
            self.exp = expm
        else:
            self.exp = exponential

        if nonlinearity == 'relu':
            self.nonlinearity = nn.ReLU()
        elif nonlinearity == 'modrelu':
            self.nonlinearity = modrelu(hidden_size)
            self.params.append(self.nonlinearity.b)
        elif nonlinearity == 'tanh':
            self.nonlinearity = nn.Tanh()
        elif nonlinearity == 'sigmoid':
            self.nonlinearity = nn.Sigmoid()
        else:
            self.nonlinearity = None

        self.reset_parameters()

    def _A(self, gradients):
        A = self.log_recurrent_kernel
        if not gradients:
            A = A.data
        A = A.triu(diagonal=1)
        return A - A.t()

    def _B(self, gradients):
        return self.exp(self._A(gradients))

    def _norm(self):
        norm = snorm.spectral_norm(self.recurrent_kernel.data)
        return norm

    def get_alpha(self):
        return self.gamma.clone().detach().cpu().numpy()
        
        
    def reset_parameters(self):
        print(self.skew_initializer)
        self.log_recurrent_kernel.data = \
                torch.as_tensor(self.skew_initializer(self.hidden_size),
                        dtype=self.log_recurrent_kernel.dtype,
                        device=self.log_recurrent_kernel.device)
        self.recurrent_kernel.data = self._B(gradients=False)
        self.input_initializer(self.input_kernel.weight.data)

    def orthogonal_step(self, optim):
        if self.exact:
            A = self._A(gradients=False)
            B = self.recurrent_kernel.data
            G = self.recurrent_kernel.grad.data
            BtG = B.t().mm(G)
            grad = 0.5*(BtG - BtG.t())
            frechet_deriv = B.mm(expm_frechet(-A, grad))

            # Account for the triangular representation of the skew-symmetric matrix
            # The gradient with respect to the parameter x_i,j is the upper part of the triangular matrix
            # minus the lower part of the triangular matrix.
            # The gradient needn't be triangular, but we make it triangular for consistency
            self.log_recurrent_kernel.grad = (frechet_deriv - frechet_deriv.t()).triu(diagonal=1)
        else:
            orth_param = self._B(gradients=True)
            self.log_recurrent_kernel.grad = \
                torch.autograd.grad([orth_param], self.log_recurrent_kernel,
                                    grad_outputs=(self.recurrent_kernel.grad,))[0]
        optim.step()
        self.recurrent_kernel.data = self._B(gradients=False)
        self.recurrent_kernel.grad.data.zero_()

    def forward(self, input, hidden=None, i = 1): 
        if hidden is None:
            hidden = input.new_zeros(input.size(0), self.hidden_size, requires_grad=True)
            #init = torch.mm(self.input_kernel(input), torch.inverse(self.recurrent_kernel.data - torch.eye(self.recurrent_kernel.data.size(0) , device= input.device)))
            #hidden = init.detach().clone() 
            #hidden.requires_grad=True
            self.first_hidden = hidden

        input = self.input_kernel(input)

        hidden = hidden.matmul(self.recurrent_kernel)

        h = input + hidden

        if self.nonlinearity:
            h = (1- torch.exp(-self.gamma[i] * h)) * self.nonlinearity(h)
        else:
            h = (1- torch.exp(-self.gamma[i] * h)) * h
        return h

    def extra_repr(self):
        return 'input_size={}, hidden_size={}'.format(self.input_size, self.hidden_size)
