import json, logging, os, torch

import random

from config import get_configuration

import numpy as np

from classifiers.classifier_factory import ClassifierFactory

from datareader.data_reader_factory import DataReaderFactory

from cliparser import setup_parser

from evaluator.evaluator_factory import EvaluatorFactory

from evaluator.standard_evaluator import StandardEvaluator

from tokenizers.tokenizer_factory import TokenizerFactory

from evaluator.validation_metric_factory import ValidationMetricFactory

def set_random_seeds(seed):

    random.seed(seed)

    np.random.seed(seed)

    torch.manual_seed(seed)

def save_config(defconfig):

    config_path = os.path.join(defconfig['output_dir'], defconfig['experiment_name'], 'config.json')

    if not os.path.exists(

            os.path.join(defconfig['output_dir'], defconfig['experiment_name'])):

        os.makedirs(os.path.join(defconfig['output_dir'], defconfig['experiment_name']))

    with open(config_path, 'w') as f:

        json.dump(defconfig, f)

if __name__ == "__main__":

    # Load config

    args = setup_parser()

    defconfig = get_configuration(args)

    print(defconfig)

    save_config(defconfig)

    logging.info("Number of GPUS: {}".format(torch.cuda.device_count()))

    if 'seed' in defconfig:

        set_random_seeds(defconfig['seed'])

    # Load the tokenizer to use

    tokenizerFactory = TokenizerFactory()

    tokenizer = tokenizerFactory.make_tokenizer(defconfig['tokenizer'])

    # Build a classifier object to use

    classifierFactory = ClassifierFactory(defconfig)

    classifier = classifierFactory.make_classifier(defconfig['classifier'])

    # load the data

    dataReaderFactory = DataReaderFactory(defconfig)

    datareader = dataReaderFactory.make_datareader(defconfig['datareader'], tokenizer)

    if args.train:

        # Load from checkpoint if we're using one (won't do anything if were not)

        if args.train_from_checkpoint:

            classifier.load_from_checkpoint(args.train_from_checkpoint)

        # Pass the classifier to the trainer

        classifier.train(datareader)

    if args.eval:

        # setup the correct validator

        results_path = os.path.join(defconfig['output_dir'], defconfig['experiment_name'], "results")

        validator = ValidationMetricFactory().make_validator(defconfig['validation_metric'])

        evaluator = EvaluatorFactory().make_evaluator(defconfig['evaluator'], results_path, defconfig, datareader, validator)

        checkpoints_path = os.path.join(defconfig['output_dir'], defconfig['experiment_name'], "checkpoints")

        # Loop over all the checkpoints, running evaluations on all them.

        for checkpoint in os.listdir(checkpoints_path):

            # Load the checkpoint model

            classifier.load_from_checkpoint(os.path.join(checkpoints_path, checkpoint))

            evaluator.go(classifier, checkpoint)

        evaluator.test()