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")

[((515, 611), 'datasets.imdb.get_data.get_dataloader', 'get_dataloader', (['"""../video/multimodal_imdb.hdf5"""', '"""../video/mmimdb"""'], {'vgg': '(True)', 'batch_size': '(128)'}), "('../video/multimodal_imdb.hdf5', '../video/mmimdb', vgg=True,\n    batch_size=128)\n", (529, 611), False, 'from datasets.imdb.get_data import get_dataloader\n'), ((1227, 1533), 'training_structures.Supervised_Learning.train', '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': '(0.005)', 'weight_decay': '(0.01)', 'objective': 'recon_loss'}), "(encoders, fuse, head, traindata, validdata, 1000, decoders +\n    intermediates, early_stop=True, task='multilabel', objective_args_dict=\n    {'decoders': decoders, 'intermediates': intermediates}, save=filename,\n    optimtype=torch.optim.AdamW, lr=0.005, weight_decay=0.01, objective=\n    recon_loss)\n", (1232, 1533), False, 'from training_structures.Supervised_Learning import train, test\n'), ((51, 62), 'os.getcwd', 'os.getcwd', ([], {}), '()\n', (60, 62), False, 'import os\n'), ((842, 872), 'unimodals.common_models.Linear', 'Linear', (['(n_latent // 2)', 'classes'], {}), '(n_latent // 2, classes)\n', (848, 872), False, 'from unimodals.common_models import Linear, MLP, MaxOut_MLP\n'), ((1196, 1224), 'torch.nn.BCEWithLogitsLoss', 'torch.nn.BCEWithLogitsLoss', ([], {}), '()\n', (1222, 1224), False, 'import torch\n'), ((1546, 1566), 'torch.load', 'torch.load', (['filename'], {}), '(filename)\n', (1556, 1566), False, 'import torch\n'), ((1646, 1674), 'torch.nn.BCEWithLogitsLoss', 'torch.nn.BCEWithLogitsLoss', ([], {}), '()\n', (1672, 1674), False, 'import torch\n'), ((661, 669), 'fusions.common_fusions.Concat', 'Concat', ([], {}), '()\n', (667, 669), False, 'from fusions.common_fusions import Concat\n'), ((671, 713), 'unimodals.common_models.MLP', 'MLP', (['(2 * n_latent)', 'n_latent', '(n_latent // 2)'], {}), '(2 * n_latent, n_latent, n_latent // 2)\n', (674, 713), False, 'from unimodals.common_models import Linear, MLP, MaxOut_MLP\n'), ((730, 772), 'unimodals.common_models.MaxOut_MLP', 'MaxOut_MLP', (['(512)', '(512)', '(300)', 'n_latent', '(False)'], {}), '(512, 512, 300, n_latent, False)\n', (740, 772), False, 'from unimodals.common_models import Linear, MLP, MaxOut_MLP\n'), ((782, 826), 'unimodals.common_models.MaxOut_MLP', 'MaxOut_MLP', (['(512)', '(1024)', '(4096)', 'n_latent', '(False)'], {}), '(512, 1024, 4096, n_latent, False)\n', (792, 826), False, 'from unimodals.common_models import Linear, MLP, MaxOut_MLP\n'), ((891, 914), 'unimodals.common_models.MLP', 'MLP', (['n_latent', '(600)', '(300)'], {}), '(n_latent, 600, 300)\n', (894, 914), False, 'from unimodals.common_models import Linear, MLP, MaxOut_MLP\n'), ((923, 948), 'unimodals.common_models.MLP', 'MLP', (['n_latent', '(2048)', '(4096)'], {}), '(n_latent, 2048, 4096)\n', (926, 948), False, 'from unimodals.common_models import Linear, MLP, MaxOut_MLP\n'), ((974, 1017), 'unimodals.common_models.MLP', 'MLP', (['n_latent', '(n_latent // 2)', '(n_latent // 2)'], {}), '(n_latent, n_latent // 2, n_latent // 2)\n', (977, 1017), False, 'from unimodals.common_models import Linear, MLP, MaxOut_MLP\n'), ((1039, 1082), 'unimodals.common_models.MLP', 'MLP', (['n_latent', '(n_latent // 2)', '(n_latent // 2)'], {}), '(n_latent, n_latent // 2, n_latent // 2)\n', (1042, 1082), False, 'from unimodals.common_models import Linear, MLP, MaxOut_MLP\n')]