import torch
import sys
import os

sys.path.append(os.getcwd())

from utils.helper_modules import Sequential2
from unimodals.common_models import Linear, MLP, MaxOut_MLP
from datasets.imdb.get_data import get_dataloader
from fusions.common_fusions import Concat
from objective_functions.objectives_for_supervised_learning import MFM_objective
from objective_functions.recon import sigmloss1d
from training_structures.Supervised_Learning import train, test

filename = "best_mfm.pt"
traindata, validdata, testdata = get_dataloader(
    "../video/multimodal_imdb.hdf5", "../video/mmimdb", vgg=True, batch_size=128)

classes = 23
n_latent = 512
fuse = Sequential2(Concat(), MLP(2*n_latent, n_latent, n_latent//2)).cuda()
encoders = [MaxOut_MLP(512, 512, 300, n_latent, False).cuda(
), MaxOut_MLP(512, 1024, 4096, n_latent, False).cuda()]
head = Linear(n_latent//2, classes).cuda()

decoders = [MLP(n_latent, 600, 300).cuda(), MLP(n_latent, 2048, 4096).cuda()]
intermediates = [MLP(n_latent, n_latent//2, n_latent//2).cuda(),
                 MLP(n_latent, n_latent//2, n_latent//2).cuda()]

recon_loss = MFM_objective(2.0, [sigmloss1d, sigmloss1d], [
                           1.0, 1.0], criterion=torch.nn.BCEWithLogitsLoss())

train(encoders, fuse, head, traindata, validdata, 1000, decoders+intermediates, early_stop=True, task="multilabel",
      objective_args_dict={"decoders": decoders, "intermediates": intermediates}, save=filename, optimtype=torch.optim.AdamW, lr=5e-3, weight_decay=0.01, objective=recon_loss)

print("Testing:")
model = torch.load(filename).cuda()
test(model, testdata, method_name="MFM", dataset="imdb",
     criterion=torch.nn.BCEWithLogitsLoss(), task="multilabel")

[((16, 33, 17, 81), 'datasets.imdb.get_data.get_dataloader', 'get_dataloader', (), '', False, 'from datasets.imdb.get_data import get_dataloader\n'), ((33, 0, 34, 175), 'training_structures.Supervised_Learning.train', 'train', (), '', False, 'from training_structures.Supervised_Learning import train, test\n'), ((5, 16, 5, 27), 'os.getcwd', 'os.getcwd', ({}, {}), '()', False, 'import os\n'), ((24, 7, 24, 35), 'unimodals.common_models.Linear', 'Linear', ({(24, 14, 24, 25): 'n_latent // 2', (24, 27, 24, 34): 'classes'}, {}), '(n_latent // 2, classes)', False, 'from unimodals.common_models import Linear, MLP, MaxOut_MLP\n'), ((31, 48, 31, 76), 'torch.nn.BCEWithLogitsLoss', 'torch.nn.BCEWithLogitsLoss', ({}, {}), '()', False, 'import torch\n'), ((37, 8, 37, 28), 'torch.load', 'torch.load', ({(37, 19, 37, 27): 'filename'}, {}), '(filename)', False, 'import torch\n'), ((39, 15, 39, 43), 'torch.nn.BCEWithLogitsLoss', 'torch.nn.BCEWithLogitsLoss', ({}, {}), '()', False, 'import torch\n'), ((21, 19, 21, 27), 'fusions.common_fusions.Concat', 'Concat', ({}, {}), '()', False, 'from fusions.common_fusions import Concat\n'), ((21, 29, 21, 67), 'unimodals.common_models.MLP', 'MLP', ({(21, 33, 21, 43): '2 * n_latent', (21, 45, 21, 53): 'n_latent', (21, 55, 21, 66): 'n_latent // 2'}, {}), '(2 * n_latent, n_latent, n_latent // 2)', False, 'from unimodals.common_models import Linear, MLP, MaxOut_MLP\n'), ((22, 12, 22, 54), 'unimodals.common_models.MaxOut_MLP', 'MaxOut_MLP', ({(22, 23, 22, 26): '(512)', (22, 28, 22, 31): '(512)', (22, 33, 22, 36): '(300)', (22, 38, 22, 46): 'n_latent', (22, 48, 22, 53): '(False)'}, {}), '(512, 512, 300, n_latent, False)', False, 'from unimodals.common_models import Linear, MLP, MaxOut_MLP\n'), ((23, 3, 23, 47), 'unimodals.common_models.MaxOut_MLP', 'MaxOut_MLP', ({(23, 14, 23, 17): '(512)', (23, 19, 23, 23): '(1024)', (23, 25, 23, 29): '(4096)', (23, 31, 23, 39): 'n_latent', (23, 41, 23, 46): '(False)'}, {}), '(512, 1024, 4096, n_latent, False)', False, 'from unimodals.common_models import Linear, MLP, MaxOut_MLP\n'), ((26, 12, 26, 35), 'unimodals.common_models.MLP', 'MLP', ({(26, 16, 26, 24): 'n_latent', (26, 26, 26, 29): '(600)', (26, 31, 26, 34): '(300)'}, {}), '(n_latent, 600, 300)', False, 'from unimodals.common_models import Linear, MLP, MaxOut_MLP\n'), ((26, 44, 26, 69), 'unimodals.common_models.MLP', 'MLP', ({(26, 48, 26, 56): 'n_latent', (26, 58, 26, 62): '(2048)', (26, 64, 26, 68): '(4096)'}, {}), '(n_latent, 2048, 4096)', False, 'from unimodals.common_models import Linear, MLP, MaxOut_MLP\n'), ((27, 17, 27, 56), 'unimodals.common_models.MLP', 'MLP', ({(27, 21, 27, 29): 'n_latent', (27, 31, 27, 42): '(n_latent // 2)', (27, 44, 27, 55): '(n_latent // 2)'}, {}), '(n_latent, n_latent // 2, n_latent // 2)', False, 'from unimodals.common_models import Linear, MLP, MaxOut_MLP\n'), ((28, 17, 28, 56), 'unimodals.common_models.MLP', 'MLP', ({(28, 21, 28, 29): 'n_latent', (28, 31, 28, 42): '(n_latent // 2)', (28, 44, 28, 55): '(n_latent // 2)'}, {}), '(n_latent, n_latent // 2, n_latent // 2)', False, 'from unimodals.common_models import Linear, MLP, MaxOut_MLP\n')]