import sys
import torch

# Custom generator for our dataset
from modules.python.models.train import train
from modules.python.models.train_distributed import train_distributed
from modules.python.Options import TrainOptions
from modules.python.FileManager import FileManager
from modules.python.TextColor import TextColor

"""
The train module of HELEN trains a deep neural network to perform a multi-task classification. It takes a set of
labeled images from MarginPolish and trains the model to predict a base and the run-length of the base using a 
gated recurrent unit (GRU) based model. This script is the interface to the training module. 
"""


class TrainModule:
    """
    Train module that provides an interface to the train method of HELEN.
    """
    def __init__(self, train_file, test_file, gpu_mode, device_ids, max_epochs, batch_size, num_workers,
                 retrain_model, retrain_model_path, model_dir, stats_dir):
        self.train_file = train_file
        self.test_file = test_file
        self.gpu_mode = gpu_mode
        self.device_ids = device_ids
        self.model_dir = model_dir
        self.epochs = max_epochs
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.retrain_model = retrain_model
        self.retrain_model_path = retrain_model_path
        self.stats_dir = stats_dir
        self.hidden_size = TrainOptions.HIDDEN_SIZE
        self.gru_layers = TrainOptions.GRU_LAYERS
        self.learning_rate = 0.0001
        self.weight_decay = 0

    def train_model(self):
        # train a model
        train(self.train_file,
              self.test_file,
              self.batch_size,
              self.epochs,
              self.gpu_mode,
              self.num_workers,
              self.retrain_model,
              self.retrain_model_path,
              self.gru_layers,
              self.hidden_size,
              self.learning_rate,
              self.weight_decay,
              self.model_dir,
              self.stats_dir,
              not_hyperband=True)

    def train_model_gpu(self):
        """
        DO DISTRIBUTED GPU INFERENCE. THIS MODE WILL ENABLE ONE MODEL PER GPU
        """
        if not torch.cuda.is_available():
            sys.stderr.write(TextColor.RED + "ERROR: TORCH IS NOT BUILT WITH CUDA.\n" + TextColor.END)
            sys.stderr.write(TextColor.RED + "SEE TORCH CAPABILITY:\n$ python3\n"
                                             ">>> import torch \n"
                                             ">>> torch.cuda.is_available()\n If true then cuda is avilable"
                             + TextColor.END)
            exit(1)

        # Now see which devices to use
        if self.device_ids is None:
            total_gpu_devices = torch.cuda.device_count()
            sys.stderr.write(TextColor.GREEN + "INFO: TOTAL GPU AVAILABLE: " + str(total_gpu_devices) + "\n" + TextColor.END)
            device_ids = [i for i in range(0, total_gpu_devices)]
            callers = total_gpu_devices
        else:
            device_ids = [int(i) for i in self.device_ids.split(',')]
            for device_id in device_ids:
                major_capable, minor_capable = torch.cuda.get_device_capability(device=device_id)
                if major_capable < 0:
                    sys.stderr.write(TextColor.RED + "ERROR: GPU DEVICE: " + str(device_id) + " IS NOT CUDA CAPABLE.\n" + TextColor.END)
                    sys.stderr.write(TextColor.GREEN + "Try running: $ python3\n"
                                                       ">>> import torch \n"
                                                       ">>> torch.cuda.get_device_capability(device="
                                     + str(device_id) + ")\n" + TextColor.END)
                else:
                    sys.stderr.write(TextColor.GREEN + "INFO: CAPABILITY OF GPU#" + str(device_id)
                                     + ":\t" + str(major_capable) + "-" + str(minor_capable) + "\n" + TextColor.END)
            callers = len(device_ids)

        if callers == 0:
            sys.stderr.write(TextColor.RED + "ERROR: NO GPU AVAILABLE BUT GPU MODE IS SET\n" + TextColor.END)
            exit()

        # train a model
        train_distributed(self.train_file,
                          self.test_file,
                          self.batch_size,
                          self.epochs,
                          self.gpu_mode,
                          self.num_workers,
                          self.retrain_model,
                          self.retrain_model_path,
                          self.gru_layers,
                          self.hidden_size,
                          self.learning_rate,
                          self.weight_decay,
                          self.model_dir,
                          self.stats_dir,
                          device_ids,
                          callers,
                          train_mode=True)


def train_interface(train_dir, test_dir, gpu_mode, device_ids, epoch_size, batch_size, num_workers, output_dir,
                    retrain_model, retrain_model_path):
    """
    Interface to perform training
    :param train_dir: Path to directory containing training images
    :param test_dir: Path to directory containing training images
    :param gpu_mode: GPU mode
    :param device_ids: Device IDs of devices to use for GPU inference
    :param epoch_size: Number of epochs to train on
    :param batch_size: Batch size
    :param num_workers: Number of workers for data loading
    :param output_dir: Path to directory to save model
    :param retrain_model: If you want to retrain an existing model
    :param retrain_model_path: Path to the model you want to retrain
    :return:
    """
    model_out_dir, stats_dir = FileManager.handle_train_output_directory(output_dir)
    tm = TrainModule(train_dir,
                     test_dir,
                     gpu_mode,
                     device_ids,
                     epoch_size,
                     batch_size,
                     num_workers,
                     retrain_model,
                     retrain_model_path,
                     model_out_dir,
                     stats_dir)

    if gpu_mode:
        tm.train_model_gpu()
    else:
        tm.train_model()
