import datetime
import pickle
from geneformer import Classifier
import os 
current_date = datetime.datetime.now()
datestamp = f"{str(current_date.year)[-2:]}{current_date.month:02d}{current_date.day:02d}{current_date.hour:02d}{current_date.minute:02d}{current_date.second:02d}"
datestamp_min = f"{str(current_date.year)[-2:]}{current_date.month:02d}{current_date.day:02d}"

output_prefix = "tf_dosage_sens_test"
output_dir = f"/path/to/output_dir/{datestamp}"
os.makedirs(output_dir)

with open("/path/to/Genecorpus-30M/dosage_sensitivity_TFs.pickle", "rb") as fp:
    gene_class_dict = pickle.load(fp)

cc = Classifier(classifier="gene",
                gene_class_dict = gene_class_dict,
                max_ncells = 10_000,
                freeze_layers = 4,
                num_crossval_splits = 5,
                forward_batch_size=200,
                nproc=16)


cc.prepare_data(input_data_file="/path/to/Genecorpus-30M/dosage_sensitive_tfs",
                output_directory=output_dir,
                output_prefix=output_prefix)

all_metrics = cc.validate(model_directory="/home/Geneformer",
                          prepared_input_data_file=f"{output_dir}/{output_prefix}_labeled.dataset",
                          id_class_dict_file=f"{output_dir}/{output_prefix}_id_class_dict.pkl",
                          output_directory=output_dir,
                          output_prefix=output_prefix)