from moETM.train import Trainer_moETM_for_cross_prediction, Train_moETM_for_cross_prediction
from dataloader import load_nips_dataset_rna_protein_dataset, prepare_nips_dataset, data_process_moETM_cross_prediction
from moETM.build_model import build_moETM
import pandas as pd
import gc
import os
import numpy as np
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

import warnings
warnings.filterwarnings('ignore')

# Load dataset
mod_file_path = "./data/GSE194122_openproblems_neurips2021_cite_BMMC_processed.h5ad"
gene_encoding = pd.read_csv('./useful_file/gene_coding_nips_rna_protein.csv')
protein_encoding = pd.read_csv('./useful_file/protein_coding_nips_rna_protein.csv')

adata_mod1, adata_mod2 = load_nips_dataset_rna_protein_dataset(mod_file_path, gene_encoding, protein_encoding)
gc.collect()

# Prepare dataset
adata_mod1, adata_mod2 = prepare_nips_dataset(adata_mod1, adata_mod2)

n_total_sample = adata_mod1.shape[0]

X_mod1_train_T, X_mod2_train_T, batch_index_train_T, X_mod1_test_T, X_mod2_test_T, batch_index_test_T, test_adata_mod1, train_adata_mod1, test_mod1_sum, test_mod2_sum= data_process_moETM_cross_prediction(adata_mod1, adata_mod2, n_sample=np.int(np.floor(n_total_sample*0.8)))

num_batch = len(batch_index_train_T.unique())
input_dim_mod1 = X_mod1_train_T.shape[1]
input_dim_mod2 = X_mod2_train_T.shape[1]
train_num = X_mod1_train_T.shape[0]

num_topic = 200
emd_dim = 400
encoder_mod1, encoder_mod2, decoder, optimizer = build_moETM(input_dim_mod1, input_dim_mod2, num_batch, num_topic=num_topic, emd_dim=emd_dim)

direction = 'another_to_rna'   # Or another_to_rna
trainer = Trainer_moETM_for_cross_prediction(encoder_mod1, encoder_mod2, decoder, optimizer, direction)

Total_epoch = 500
batch_size = 2000
Train_set = [X_mod1_train_T, X_mod2_train_T, batch_index_train_T]
Test_set = [X_mod1_test_T, X_mod2_test_T, batch_index_test_T, test_adata_mod1, test_mod1_sum, test_mod2_sum]
Train_moETM_for_cross_prediction(trainer, Total_epoch, train_num, batch_size, Train_set, Test_set)