import numpy as np
import torch
import random
from meta_train import mmdPreModel
from collections import namedtuple
import joblib
from transformers import RobertaTokenizer, RobertaModel


def api_init():

	random.seed(0)
	np.random.seed(0)
	torch.manual_seed(0)
	torch.cuda.manual_seed(0)
	torch.cuda.manual_seed_all(0)
	torch.backends.cudnn.benchmark = False
	torch.backends.cudnn.deterministic = True

	model_name = 'roberta-base-openai-detector'
	model_path_api = f'.'
	token_num, hidden_size = 100, 768

	Config = namedtuple('Config', ['in_dim', 'hid_dim', 'dropout', 'out_dim', 'token_num'])
	config = Config(
			in_dim=hidden_size,
			token_num=token_num,
			hid_dim=512,
			dropout=0.2,
			out_dim=300,)
	
	net = mmdPreModel(config=config, num_mlp=0, transformer_flag=True, num_hidden_layers=1)

	# load the features and models
	feature_ref_for_test_filename = f'{model_path_api}/feature_ref_for_test.pt'
	model_filename = f'{model_path_api}/logistic_regression_model.pkl'
	net_filename = f'{model_path_api}/net.pt'
	
	load_ref_data =  torch.load(feature_ref_for_test_filename,map_location=torch.device('cpu')) # cpu
	loaded_model = joblib.load(model_filename) # cpu
	checkpoint = torch.load(net_filename,map_location=torch.device('cpu'))
	net.load_state_dict(checkpoint['net'])
	sigma, sigma0_u, ep  = checkpoint['sigma'], checkpoint['sigma0_u'], checkpoint['ep']

	# generic generative model
	cache_dir = ".cache"
	base_tokenizer = RobertaTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
	base_model = RobertaModel.from_pretrained(model_name, output_hidden_states=True, cache_dir=cache_dir)

	# whether load the model to gpu
	gpu_using = False

	DEVICE = torch.device("cpu")
	if gpu_using:
		DEVICE = torch.device("cuda:0")
	net = net.to(DEVICE)
	sigma, sigma0_u, ep  = sigma.to(DEVICE), sigma0_u.to(DEVICE), ep.to(DEVICE)
	load_ref_data = load_ref_data.to(DEVICE)
	base_model = base_model.to(DEVICE)	
	num_ref = 5000
	feature_ref = load_ref_data[np.random.permutation(load_ref_data.shape[0])][:num_ref].to(DEVICE)

	return base_model, base_tokenizer, net, feature_ref, sigma, sigma0_u, ep, loaded_model, DEVICE