import math
from argparse import Namespace

import torch
import torch.nn as nn
import torch.nn.functional as F
import models
from models import register
import numpy as np

class ExpansionNet(nn.Module):
    def __init__(self, args):
        super(ExpansionNet, self).__init__()
        self.args = args
        self.in_dim = args.in_dim
        self.out_dim = args.out_dim
        self.hidden_list = args.hidden_list
        layers = []
        lastv = self.in_dim
        hidden_list = self.hidden_list
        out_dim = self.out_dim
        for hidden in hidden_list:
            layers.append(nn.Linear(lastv, hidden))
            layers.append(nn.ReLU())
            lastv = hidden
        layers.append(nn.Linear(lastv, out_dim))
        self.layers = nn.Sequential(*layers)

    def forward(self, x):
        b, _, c = x.shape
        x = x.view(-1, c)
        logits = self.layers(x)
        out = nn.functional.normalize(logits, dim=1)
        return out.view(b,_,self.out_dim)


@register('ExpansionNet')
def make_ExpansionNet(in_dim=580,out_dim=10,hidden_list=None):
    args = Namespace()
    args.in_dim = in_dim
    args.out_dim = out_dim
    args.hidden_list = hidden_list
    return ExpansionNet(args)