import sys
import os
sys.path.append(os.getcwd())
from training_structures.MVAE_mixed import train_MVAE,test_MVAE
from fusions.MVAE import ProductOfExperts
from unimodals.common_models import MLP
from unimodals.MVAE import LeNetEncoder,DeLeNet
from torch import nn
import torch
from objective_functions.recon import elbo_loss,sigmloss1dcentercrop
from datasets.avmnist.get_data import get_dataloader

from private_test_scripts.all_in_one import all_in_one_train,all_in_one_test

traindata, validdata, testdata = get_dataloader('/data/yiwei/avmnist/_MFAS/avmnist')

classes=10
n_latent=200
fuse=ProductOfExperts((1,40,n_latent))


channels=6
encoders=[LeNetEncoder(1,channels,3,n_latent).cuda(),LeNetEncoder(1,channels,5,n_latent).cuda()]
decoders=[DeLeNet(1,channels,3,n_latent).cuda(),DeLeNet(1,channels,5,n_latent).cuda()]
head=MLP(n_latent,40,classes).cuda()
elbo=elbo_loss([sigmloss1dcentercrop(28,34),sigmloss1dcentercrop(112,130)],[1.0,1.0],0.0)
def trpr():
    train_MVAE(encoders,decoders,head,fuse,traindata,validdata,elbo,20)
#all_in_one_train(trpr,[encoders[0],encoders[1],decoders[0],decoders[1],head])
mvae=torch.load('best1.pt')
head=torch.load('best2.pt')
def tepr():
    test_MVAE(mvae,head,testdata)
all_in_one_test(tepr,[encoders[0],encoders[1],head])
