import torch 
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import time
from math import sqrt

# Add TT Linear layer implementation
class TTLinear(nn.Module):
    def __init__(self, in_features, out_features, ranks, activation=None, bias=True):
        super(TTLinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.ranks = ranks
        self.activation = activation
        
        # Factorize input and output dimensions
        self.in_shape, self.out_shape = self.factorize_dims(in_features, out_features)
        
        # Create TT cores
        self.cores = nn.ParameterList()
        for i in range(len(self.in_shape)):
            if i == 0:
                input_rank = 1
            else:
                input_rank = ranks[i-1]
            
            if i == len(self.in_shape)-1:
                output_rank = 1
            else:
                output_rank = ranks[i]
            
            core = nn.Parameter(torch.randn(
                input_rank, self.in_shape[i], self.out_shape[i], output_rank
            ) * sqrt(2.0 / (input_rank * self.in_shape[i])))
            self.cores.append(core)
        
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))
        else:
            self.register_parameter('bias', None)
    
    def factorize_dims(self, in_dim, out_dim):
        # Simple factorization - can be customized
        factors = []
        d = in_dim
        while d > 1:
            for i in range(int(sqrt(d)), 0, -1):
                if d % i == 0:
                    factors.append(i)
                    factors.append(d // i)
                    break
            d = factors.pop()
        in_shape = factors
        
        factors = []
        d = out_dim
        while d > 1:
            for i in range(int(sqrt(d)), 0, -1):
                if d % i == 0:
                    factors.append(i)
                    factors.append(d // i)
                    break
            d = factors.pop()
        out_shape = factors
        
        return in_shape, out_shape
    
    def forward(self, x):
        batch_size = x.size(0)
        
        # Reshape input to tensor form
        x = x.view(batch_size, *self.in_shape)
        
        # Contract input with first core
        res = torch.einsum('bi,rivo->bvo', x, self.cores[0])
        
        # Contract with remaining cores
        for i in range(1, len(self.cores)):
            res = torch.einsum('bvi,rivo->bvo', res, self.cores[i])
        
        # Reshape to output dimension
        res = res.contiguous().view(batch_size, self.out_features)
        
        if self.bias is not None:
            res = res + self.bias
        
        if self.activation is not None:
            res = self.activation(res)
        
        return res


# Modify get_activation_fn to include TTLinear option
def get_activation_fn(activation):
    if activation == 'tanh':
        return F.tanh
    elif activation == 'elu':
        return F.elu
    elif activation == 'relu':
        return F.relu
    elif activation == 'selu':
        return F.selu
    elif activation == 'sigmoid':
        return F.sigmoid
    elif activation == 'gelu':
        return F.gelu
    elif activation == 'silu':
        return F.silu
    elif activation == 'mish':
        return F.mish
    elif activation == 'linear':
        return nn.Identity()
    else:
        raise NotImplementedError(f'Activation function {activation} not implemented')

# Modify g_Full to use TTLinear
class g_Full(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, depth=3, skip_depth=1, 
                 skip_layer=1, ell=2, activation='selu', use_skip=False, augment=False,
                 use_tt=False, tt_ranks=None):
        super(g_Full, self).__init__()
        
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.depth = depth
        self.ell = ell
        self.ell_input_size = input_size//self.ell
        self.augment = augment
        self.activation_fn = get_activation_fn(activation)
        self.skip_depth = skip_depth
        self.skip_layer = skip_layer
        self.use_skip = use_skip
        self.use_tt = use_tt
        self.tt_ranks = tt_ranks if tt_ranks is not None else [8, 8]  # Default ranks
        
        if self.use_skip:
            if self.use_tt:
                self.skip = nn.ModuleList([TTLinear(self.input_size + self.output_size, 
                                                  self.hidden_size, 
                                                  ranks=self.tt_ranks,
                                                  activation=self.activation_fn)])
                self.skip.extend([TTLinear(self.hidden_size, self.hidden_size, 
                                         ranks=self.tt_ranks,
                                         activation=self.activation_fn) 
                               for ii in range(1, self.skip_depth)])
            else:
                self.skip = nn.ModuleList([nn.Linear(self.input_size + self.output_size, 
                                                   self.hidden_size, bias=True)])
                self.skip.extend([nn.Linear(self.hidden_size, self.hidden_size, bias=True) 
                                for ii in range(1, self.skip_depth)])
        
        if self.use_tt:
            self.linears = nn.ModuleList([TTLinear(self.input_size, self.hidden_size, 
                                                  ranks=self.tt_ranks,
                                                  activation=self.activation_fn)])
            self.linears.extend([TTLinear(self.hidden_size, self.hidden_size, 
                                        ranks=self.tt_ranks,
                                        activation=self.activation_fn) 
                              for ii in range(1, self.depth)])
            self.linears.append(TTLinear(self.hidden_size, self.output_size, 
                                       ranks=self.tt_ranks))
        else:
            self.linears = nn.ModuleList([nn.Linear(self.input_size, self.hidden_size, bias=True)])
            self.linears.extend([nn.Linear(self.hidden_size, self.hidden_size, bias=True) 
                               for ii in range(1, self.depth)])
            self.linears.append(nn.Linear(self.hidden_size, self.output_size, bias=True))

    def __init__(self, input_size, hidden_size, output_size, depth=3, skip_depth = 1, skip_layer = 1, ell = 2, activation = 'selu', use_skip = False, augment = False):
        super(g_Full, self).__init__()

        self.input_size  = input_size


        self.hidden_size = hidden_size
        self.output_size = output_size

        self.depth = depth
        self.ell = ell
        self.ell_input_size = input_size//self.ell

        self.augment = augment
        self.activation_fn = get_activation_fn(activation)
        self.skip_depth = skip_depth
        self.skip_layer = skip_layer

        self.use_skip = use_skip 
        if self.use_skip:
            self.skip = nn.ModuleList([nn.Linear(self.input_size + self.output_size, self.hidden_size, bias=True)])
            self.skip.extend([nn.Linear(self.hidden_size, self.hidden_size, bias=True) for ii in range(1, self.skip_depth)])

        self.linears = nn.ModuleList([nn.Linear(self.input_size, self.hidden_size, bias=True)])
        self.linears.extend([nn.Linear(self.hidden_size, self.hidden_size, bias=True) for ii in range(1, self.depth)])
        self.linears.append(nn.Linear(self.hidden_size, self.output_size, bias=True))


    @staticmethod
    def get_augment(msg, ell):
        u = msg.clone()
        n = int(np.log2(ell))
        for d in range(0, n):
            num_bits = 2**d
            for i in np.arange(0, ell, 2*num_bits):
                # [u v] encoded to [u xor(u,v)]
                if len(u.shape) == 2:
                    u = torch.cat((u[:, :i], u[:, i:i+num_bits].clone() * u[:, i+num_bits: i+2*num_bits], u[:, i+num_bits:]), dim=1)
                elif len(u.shape) == 3:
                    u = torch.cat((u[:, :, :i], u[:, :, i:i+num_bits].clone() * u[:, :, i+num_bits: i+2*num_bits], u[:, :, i+num_bits:]), dim=2)

                # u[:, i:i+num_bits] = u[:, i:i+num_bits].clone() * u[:, i+num_bits: i+2*num_bits].clone
        if len(u.shape) == 3:
            return u[:, :, :-1]
        elif len(u.shape) == 2:
            return u[:, :-1]

    def forward(self, y):
                
        x = y.clone()
        for ii, layer in enumerate(self.linears):
            if ii != self.depth:
                x = self.activation_fn(layer(x))
                if self.use_skip and  ii == self.skip_layer:
                    if len(x.shape) == 3:
                        skip_input = torch.cat([y, g_Full.get_augment(y, self.ell)], dim = 2)
                    elif len(x.shape) == 2:
                        skip_input = torch.cat([y, g_Full.get_augment(y, self.ell)], dim = 1)
                    for jj, skip_layer in enumerate(self.skip):
                        skip_input = self.activation_fn(skip_layer(skip_input))
                    x = x + skip_input
            else:
                x = layer(x)
                if self.augment:
                    x = x + g_Full.get_augment(y, self.ell)
        return x

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.01)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(0.0, 0.01)
        m.bias.data.fill_(0)
    elif classname.find('Linear') != -1:
        m.weight.data.normal_(0.0, 0.01)
        try:
            m.bias.data.fill_(0.)
        except:
            pass

# Modify f_Full to use TTLinear
class f_Full(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, dropout_p=0., 
                 activation='selu', depth=3, use_norm=False, use_tt=False, tt_ranks=None):
        super(f_Full, self).__init__()
        
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.depth = depth
        self.use_norm = use_norm
        self.activation_fn = get_activation_fn(activation)
        self.use_tt = use_tt
        self.tt_ranks = tt_ranks if tt_ranks is not None else [8, 8]  # Default ranks
        
        if self.use_tt:
            self.linears = nn.ModuleList([TTLinear(self.input_size, self.hidden_size, 
                                                 ranks=self.tt_ranks,
                                                 activation=self.activation_fn)])
            for ii in range(1, self.depth):
                self.linears.append(TTLinear(self.hidden_size, self.hidden_size, 
                                           ranks=self.tt_ranks,
                                           activation=self.activation_fn))
            self.linears.append(TTLinear(self.hidden_size, self.output_size, 
                                       ranks=self.tt_ranks))
        else:
            self.linears = nn.ModuleList([nn.Linear(self.input_size, self.hidden_size, bias=True)])
            for ii in range(1, self.depth):
                self.linears.append(nn.Linear(self.hidden_size, self.hidden_size, bias=True))
            self.linears.append(nn.Linear(self.hidden_size, self.output_size, bias=True))
        
        if self.use_norm:
            self.norms = nn.ModuleList([nn.LayerNorm(self.hidden_size) 
                                      for _ in range(self.depth)])
    def __init__(self, input_size, hidden_size, output_size, dropout_p = 0., activation = 'selu', depth=3, use_norm = False):
        super(f_Full, self).__init__()
        self.input_size  = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.depth = depth
        self.use_norm = use_norm

        self.activation_fn = get_activation_fn(activation)

        self.linears = nn.ModuleList([nn.Linear(self.input_size, self.hidden_size, bias=True)])
        if self.use_norm:
            self.norms = nn.ModuleList([nn.LayerNorm(self.hidden_size)])
        for ii in range(1, self.depth):
            self.linears.append(nn.Linear(self.hidden_size, self.hidden_size, bias=True))
            if self.use_norm:
                self.norms.append(nn.LayerNorm(self.hidden_size))
        self.linears.append(nn.Linear(self.hidden_size, self.output_size, bias=True))

    def forward(self, y, aug = None):

        x = y.clone()
        for ii, layer in enumerate(self.linears):
            if ii != self.depth:
                x = layer(x)
                if not hasattr(self, 'use_norm') or not self.use_norm:
                    pass
                else:
                    x = self.norms[ii](x)
                x = self.activation_fn(x)
            else:
                x = layer(x)
        return x

def get_onehot(actions):
    inds = (0.5 + 0.5*actions).long()
    return torch.eye(2, device = inds.device)[inds].reshape(actions.shape[0], -1)
