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

import torch

from training_structures.Simple_Late_Fusion import train, test
from fusions.common_fusions import Concat
sys.path.append('/home/pliang/multibench/MultiBench/datasets/imdb')
from get_data_robust import get_dataloader, get_dataloader_robust
from robustness.all_in_one import general_train, general_test
from unimodals.common_models import MLP, VGG16, Linear, MaxOut_MLP

traindata, validdata = get_dataloader('../../../video/multimodal_imdb.hdf5', batch_size=128, vgg=True)
robustdata = get_dataloader_robust('../../../video/mmimdb', '../../../video/multimodal_imdb.hdf5',batch_size=128)

encoders=[MaxOut_MLP(512, 128, 300, linear_layer=False), MaxOut_MLP(512, 1024, 4096, 128, False)]
#encoders=[MLP(300, 512, 512), VGG16(512)]
head= Linear(256, 23).cuda()
fusion=Concat().cuda()

def trainprocess(filename):
    train(encoders,fusion,head,traindata,validdata,1000, early_stop=True,task="multilabel", regularization=True, save=filename, optimtype=torch.optim.AdamW,lr=1e-2,weight_decay=0.01, criterion=torch.nn.BCEWithLogitsLoss())
filename = general_train(trainprocess, 'mmimdb_lf')

def testprocess(model, testdata):
    return test(model,testdata,criterion=torch.nn.BCEWithLogitsLoss(),task="multilabel")
general_test(testprocess, filename, robustdata, multi_measure=True)
