import sys
import torch
import math
from torch import nn
if str(sys.version_info[0])+"."+str(sys.version_info[1]) >= "3.9": # Python >= 3.9
    from math import gcd
else:
    from fractions import gcd
    
import pdb

class MLP(nn.Module):
    """
    """
    def __init__(self, dim_list, activation='relu', batch_norm=True, dropout=0):
        super(MLP, self).__init__()
        self.module = self.make_mlp(dim_list, activation, batch_norm, dropout)

    def make_mlp(self, dim_list, activation='relu', batch_norm=True, dropout=0):
        layers = []
        for dim_in, dim_out in zip(dim_list[:-1], dim_list[1:]):
            layers.append(nn.Linear(dim_in, dim_out))
            if batch_norm:
                layers.append(nn.BatchNorm1d(dim_out))
            if activation == 'relu':
                layers.append(nn.ReLU())
            elif activation == 'leakyrelu':
                layers.append(nn.LeakyReLU())
            if dropout > 0:
                layers.append(nn.Dropout(p=dropout))
        return nn.Sequential(*layers)

    def forward(self, data_input):
        return self.module(data_input)

class Linear(nn.Module):
    def __init__(self, n_in, n_out, norm='BN', ng=32, act=True):
        super(Linear, self).__init__()
        assert(norm in ['GN', 'BN', 'SyncBN'])

        self.linear = nn.Linear(n_in, n_out, bias=False)

        if norm == 'GN':
            self.norm = nn.GroupNorm(gcd(ng, n_out), n_out)
        elif norm == 'BN':
            self.norm = nn.BatchNorm1d(n_out)
        else:
            exit('SyncBN has not been added!')
        
        self.fa = nn.GELU()
        self.act = act

        self.apply(self._init_weights)

    def _init_weights(self, module):
        """
        """
        classname = module.__class__.__name__
        if classname.find('Linear') != -1:
            try:
                nn.init.kaiming_normal_(module.weight)
            except:
                pdb.set_trace()

    def forward(self, x):
        out = self.linear(x)
        # out = self.norm(out)
        if self.act:
            out = self.fa(out)
        return out

class LinearRes(nn.Module):
    def __init__(self, n_in, n_out, norm='BN', ng=32):
        super(LinearRes, self).__init__()
        assert(norm in ['GN', 'BN', 'SyncBN'])

        mid_dim = math.ceil((n_in + n_out)/2)
        self.linear1 = nn.Linear(n_in, mid_dim, bias=False)
        self.linear2 = nn.Linear(mid_dim, n_out, bias=False)
        # self.relu = nn.ReLU(inplace=True)
        self.tanh = nn.Tanh()

        if norm == 'GN':
            self.norm1 = nn.GroupNorm(gcd(ng, mid_dim), mid_dim)
            self.norm2 = nn.GroupNorm(gcd(ng, n_out), n_out)
        elif norm == 'BN':
            self.norm1 = nn.BatchNorm1d(mid_dim)
            self.norm2 = nn.BatchNorm1d(n_out)
        else:   
            exit('SyncBN has not been added!')

        if n_in != n_out:
            if norm == 'GN':
                self.transform = nn.Sequential(
                    nn.Linear(n_in, n_out, bias=False),
                    nn.GroupNorm(gcd(ng, n_out), n_out))
            elif norm == 'BN':
                self.transform = nn.Sequential(
                    nn.Linear(n_in, n_out, bias=False),
                    nn.BatchNorm1d(n_out))
            else:
                exit('SyncBN has not been added!')
        else:
            self.transform = None

    def forward(self, x):
        out = self.linear1(x)
        out = self.norm1(out)
        out = self.tanh(out)
        out = self.linear2(out)
        out = self.norm2(out)

        if self.transform is not None:
            out += self.transform(x)
        else:
            out += x

        out = self.tanh(out)
        return out