import sys
import os
import torch
from modules.python.TextColor import TextColor
from modules.python.models.predict_cpu import predict_cpu
from modules.python.models.predict_gpu import predict_gpu
from modules.python.FileManager import FileManager
from os.path import isfile, join
from os import listdir
"""
The Call Consensus method generates base predictions for images generated through MarginPolish. This script reads
hdf5 files generated by MarginPolish and produces another Hdf5 file that holds all predictions. The generated hdf5 file
is given to stitch.py which then stitches the segments using an alignment which gives us a polished sequence.

The algorithm is described here:

  1) INPUTS:
    - directory path to the image files generated by MarginPolish
    - model path directing to a trained model
    - batch size for mini-batch prediction
    - num workers for mini-batch processing threads
    - output directory path to where the output hdf5 will be saved
    - gpu mode indicating if GPU will be used
  2) METHOD:
    - Call predict function that loads the neural network and generates base predictions and saves it into a hdf5 file
        - Loads the model
        - Iterates over the input images in minibatch
        - For each image uses a sliding window method to slide of the image sequence
        - Aggregate the predictions to get sequence prediction for the entire image sequence
        - Save all the predictions to a file
  3) OUTPUT:
    - A hdf5 file containing all the base predictions   
"""


def get_file_paths_from_directory(directory_path):
    """
    Returns all paths of files given a directory path
    :param directory_path: Path to the directory
    :return: A list of paths of files
    """
    file_paths = [os.path.abspath(join(directory_path, file)) for file in listdir(directory_path)
                  if isfile(join(directory_path, file)) and file[-2:] == 'h5']
    return file_paths


def call_consensus(image_dir, model_path, batch_size, num_workers, threads, output_dir, output_prefix, gpu_mode,
                   device_ids, callers):
    """
    This method provides an interface too call the predict method that generates the prediction hdf5 file
    :param image_dir: Path to directory where all MarginPolish images are saved
    :param model_path: Path to a trained model
    :param batch_size: Batch size for minibatch processing
    :param num_workers: Number of workers for minibatch processing
    :param threads: Number of threads for pytorch
    :param output_dir: Path to the output directory
    :param output_prefix: Prefix of the output HDF5 file
    :param gpu_mode: If true, predict method will use GPU.
    :param device_ids: List of CUDA devices to use.
    :param callers: Total number of callers.
    :return:
    """
    # check the model file
    if not os.path.isfile(model_path):
        sys.stderr.write(TextColor.RED + "ERROR: CAN NOT LOCATE MODEL FILE.\n" + TextColor.END)
        exit(1)

    # check the input directory
    if not os.path.isdir(image_dir):
        sys.stderr.write(TextColor.RED + "ERROR: CAN NOT LOCATE IMAGE DIRECTORY.\n" + TextColor.END)
        exit(1)

    # check batch_size
    if batch_size <= 0:
        sys.stderr.write(TextColor.RED + "ERROR: batch_size NEEDS TO BE >0.\n" + TextColor.END)
        exit(1)

    # check num_workers
    if num_workers < 0:
        sys.stderr.write(TextColor.RED + "ERROR: num_workers NEEDS TO BE >=0.\n" + TextColor.END)
        exit(1)

    # check number of threads
    if threads <= 0:
        sys.stderr.write(TextColor.RED + "ERROR: THREAD NEEDS TO BE >=0.\n" + TextColor.END)
        exit(1)

    output_dir = FileManager.handle_output_directory(output_dir)

    # create a filename for the output file
    output_filename = os.path.join(output_dir, output_prefix)

    # inform the output directory
    sys.stderr.write(TextColor.GREEN + "INFO: " + TextColor.END + "OUTPUT FILE: " + output_filename + "\n")

    if gpu_mode:
        # Make sure that GPU is
        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 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 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)

        sys.stderr.write(TextColor.GREEN + "INFO: AVAILABLE GPU DEVICES: " + str(device_ids) + "\n" + TextColor.END)
        threads_per_caller = 0
    else:
        # calculate how many threads each caller can use
        threads_per_caller = int(threads / callers)
        device_ids = []

    # chunk the inputs
    input_files = get_file_paths_from_directory(image_dir)

    # generate file chunks to process in parallel
    file_chunks = [[] for i in range(callers)]
    for i in range(0, len(input_files)):
        file_chunks[i % callers].append(input_files[i])

    # get the file chunks
    file_chunks = [file_chunks[i] for i in range(len(file_chunks)) if len(file_chunks[i]) > 0]

    callers = len(file_chunks)

    if gpu_mode:
        # Distributed GPU setup
        predict_gpu(file_chunks, output_filename, model_path, batch_size, callers, device_ids, num_workers)
    else:
        # distributed CPU setup, call the prediction function
        predict_cpu(file_chunks, output_filename, model_path, batch_size,
                    callers, threads_per_caller, num_workers)

    # notify the user that process has completed successfully
    sys.stderr.write(TextColor.GREEN + "INFO: " + TextColor.END + "PREDICTION GENERATED SUCCESSFULLY.\n")
