import torch
import torch.nn as nn
import pickle
from torch.autograd import Variable
import numpy as np
import pypcd


class Encoder_End2End_Annotated(torch.jit.ScriptModule):
    __constants__ = ['encoder']

    def __init__(self):
        super(Encoder_End2End_Annotated, self).__init__()
        self.encoder = nn.Sequential(nn.Linear(16053, 256), nn.PReLU(), #adds dropouts are not expected
                                    nn.Linear(256, 256), nn.PReLU(),
                                 	nn.Linear(256, 60))

    @torch.jit.script_method
    def forward(self, x):
        x = self.encoder(x)
        return x


class MLP_NN_Annotated(torch.jit.ScriptModule):
    __constants__ = ['fc']
    def __init__(self):
        super(MLP_NN_Annotated, self).__init__()
        self.fc = nn.Sequential(
                	nn.Linear(74, 1280), nn.PReLU(), nn.Dropout(),
                	nn.Linear(1280, 896), nn.PReLU(), nn.Dropout(),
                	nn.Linear(896, 512), nn.PReLU(), nn.Dropout(),
                	nn.Linear(512, 384), nn.PReLU(), nn.Dropout(),
                	nn.Linear(384, 256), nn.PReLU(), nn.Dropout(),
                	nn.Linear(256, 128), nn.PReLU(), nn.Dropout(),
                	nn.Linear(128, 64), nn.PReLU(), nn.Dropout(),
                	nn.Linear(64, 32), nn.PReLU(),
                	nn.Linear(32, 7))

    @torch.jit.script_method
    def forward(self, x):
        out = self.fc(x)
        return out




# Creates the script
encoder = Encoder_End2End_Annotated()
MLP = MLP_NN_Annotated()

#modified to load weights
device = torch.device('cpu')
cae_filename = 'cae_encoder_140.pkl'
mlp_filename = 'mlp_PReLU_ae_dd140.pkl'
encoder.load_state_dict(torch.load(cae_filename, map_location=device))

# Print model's state_dict
print("Model's state_dict:")
for param_tensor in MLP.state_dict():
    print(param_tensor, "\t", MLP.state_dict()[param_tensor].size())


#print(MLP.state_dict())


MLP.load_state_dict(torch.load(mlp_filename, map_location=device))


encoder.save("encoder_annotated.pt")
encoder.save("mlp_annotated.pt")