"""
This file contains a data-generator that works on the data as it is provided by the ETL.
"""

from __future__ import absolute_import
import os
import logging
import numpy as np
import glob2 as glob
import random
import uuid
import progressbar
import multiprocessing as mp
from pyntcloud import PyntCloud
import pickle
from . import utils

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s - %(pathname)s: line %(lineno)d'))
logger.addHandler(handler)


class ETLDataGenerator(object):
    """
    This class generates data for training.
    """

    def __init__(
        self,
        dataset_path,
        input_type,
        #output_targets,
        sequence_length=0,
        image_target_shape=(160, 90),
        voxelgrid_target_shape=(32, 32, 32),
        voxel_size_meters=0.01,
        voxelgrid_random_rotation=False,
        pointcloud_target_size=32000,
        pointcloud_random_rotation=False
    ):
        """
        Initializes a DataGenerator.

        Args:
            dataset_path (string): Where the raw data is.
            input_type (string): Specifies how the input-data for the Neural Network looks like.
                                 Either 'image', 'pointcloud', 'voxgrid'.
            output_targets (list of strings): A list of targets for the Neural Network.
                                 For example *['height', 'weight']*.
            sequence_length (int): Specifies the lenght of the sequences. 0 would yield no sequence at all.
            image_target_shape (2D tuple of ints): Target shape of the images.
            voxelgrid_target_shape (3D tuple of ints): Target shape of the voxelgrids.
            voxel_size_meters (float): Size of the voxels. That is, edge length.
            voxelgrid_random_rotation (bool): If True voxelgrids will be rotated randomly.
            pointcloud_target_size (int): Target size of the pointclouds.
            pointcloud_random_rotation (bool): If True pointclouds will be rotated randomly.

        """

        # Preconditions.
        assert os.path.exists(dataset_path), "dataset_path must exist: " + str(dataset_path)
        assert isinstance(input_type, str), "input_type must be string: " + str(input_type)
        #assert isinstance(output_targets, list), "output_targets must be list: " + str(output_targets)
        if input_type == "image":
            assert len(image_target_shape) == 2, "image_target_shape must be 2-dimensional: " + \
                str(image_target_shape)
        if input_type == "voxelgrid":
            assert len(voxelgrid_target_shape) == 3, "voxelgrid_target_shape must be 3-dimensional: " + \
                str(voxelgrid_target_shape)

        # Assign the instance-variables.
        self.dataset_path = dataset_path
        self.input_type = input_type
        #self.output_targets = output_targets
        self.sequence_length = sequence_length
        self.image_target_shape = image_target_shape
        self.voxelgrid_target_shape = voxelgrid_target_shape
        self.voxel_size_meters = voxel_size_meters
        self.voxelgrid_random_rotation = voxelgrid_random_rotation
        self.pointcloud_target_size = pointcloud_target_size
        self.pointcloud_random_rotation = pointcloud_random_rotation

        # Find all QR-codes.
        self._find_qrcodes()
        assert self.qrcodes is not None, "No QR-codes found!"

        # Prepare the data.
        self._prepare_qrcodes_dictionary()

        # Create some caches.
        self.image_cache = {}
        self.voxelgrid_cache = {}
        self.pointcloud_cache = {}

        # Check if paths are fine.
        if self.input_type == "image":
            assert self.all_jpg_paths != []
        elif self.input_type == "voxelgrid" or self.input_type == "pointcloud":
            assert self.all_pcd_paths != []
        else:
            raise Exception("Unexpected: " + self.input_type)

    def _find_qrcodes(self):
        """
        Finds all QR-codes.

        Each individual is represented via a unique QR-code. This method extracts the set of QR-codes.
        """

        # Retrieve the QR-codes from the folders.
        paths = glob.glob(os.path.join(self.dataset_path, "*"))
        self.qrcodes = sorted([path.split("/")[-1] for path in paths])

    def _prepare_qrcodes_dictionary(self):

        self.all_pcd_paths = []
        self.all_jpg_paths = []

        self.qrcodes_dictionary = {}
        for qrcode in self.qrcodes:
            self.qrcodes_dictionary[qrcode] = []
            #logger.info("QR-code:", qrcode)
            measurement_paths = glob.glob(os.path.join(self.dataset_path, qrcode, "*"))
            for measurement_path in measurement_paths:
                # Getting PCDs.
                pcd_paths = glob.glob(os.path.join(measurement_path, "pcd", "*.pcd"))

                # Getting JPGs.
                jpg_paths = glob.glob(os.path.join(measurement_path, "jpg", "*.jpg"))

                # Loading the targets.
                target_path = os.path.join(measurement_path, "target.txt")
                target_file = open(target_path, "r")
                targets = np.array([float(value) for value in target_file.read().split(",")])
                target_file.close()

                # Done.
                self.qrcodes_dictionary[qrcode].append((pcd_paths, jpg_paths, targets))
                self.all_pcd_paths.extend(pcd_paths)
                self.all_jpg_paths.extend(jpg_paths)

    def analyze_files(self):

        for qrcode in self.qrcodes:
            logger.info("QR-code: %s", qrcode)
            for pcd_paths, jpg_paths, targets in self.qrcodes_dictionary[qrcode]:
                logger.info(" %d PCD-files for targets %s", len(pcd_paths), targets)
                logger.info(" %d JPG-files for targets %s", len(jpg_paths), targets)

        logger.info("Total PCD-files: %d", len(self.all_pcd_paths))
        logger.info("Total JPG-files: %d", len(self.all_jpg_paths))

        #measurement_paths = glob.glob(os.path.join(self.dataset_path, qrcode, "*"))
        #for measurement_path in measurement_paths:
        #    logger.info("  %s", measurement_path.split("/")[-1])
        #    pcd_paths = glob.glob(os.path.join(measurement_path, "*.pcd"))
        #    logger.info("    %d PCDs", len(pcd_paths))
        #logger.info("Number of JPGs: %d", len(self.jpg_paths))
        #logger.info("Number of PCDs: %d", len(self.pcd_paths))
        #logger.info("Number of JSONs (personal): %d", len(self.json_paths_personal))
        #logger.info("Number of JSONs (measures): %d", len(self.json_paths_measures))

    def generate(self, size, qrcodes_to_use=None, verbose=False, yield_file_paths=False, multiprocessing_jobs=1):

        if qrcodes_to_use is None:
            qrcodes_to_use = self.qrcodes

        # Main loop.
        while True:

            # Use only a single process.
            if multiprocessing_jobs == 1:
                yield generate_data(self, size, qrcodes_to_use, verbose, yield_file_paths, None)

            # Use multiple processes.
            elif multiprocessing_jobs > 1:

                # Create chunks of almost equal size.
                subset_sizes = [0] * multiprocessing_jobs
                subset_sizes[0:multiprocessing_jobs - 1] = [size
                                                            // multiprocessing_jobs] * (multiprocessing_jobs - 1)
                subset_sizes[multiprocessing_jobs - 1] = size - \
                    sum(subset_sizes[0:multiprocessing_jobs - 1])
                subset_sizes = [s for s in subset_sizes if s > 0]
                assert sum(subset_sizes) == size

                # Create an output_queue.
                output_queue = mp.Queue()

                # Create the processes.
                processes = []
                for subset_size in subset_sizes:
                    process_target = generate_data
                    process_args = (self, subset_size, qrcodes_to_use,
                                    verbose, yield_file_paths, output_queue)
                    process = mp.Process(target=process_target, args=process_args)
                    processes.append(process)

                # Start the processes.
                for p in processes:
                    p.start()

                # Exit the completed processes.
                for p in processes:
                    p.join()

                # Get process results from the output queue
                output_paths = [output_queue.get() for p in processes]

                # Merge data.
                x_inputs_arrays = []
                y_outputs_arrays = []
                file_paths_arrays = []
                for output_path in output_paths:
                    # Read data from file and delete it.
                    result_values = pickle.load(open(output_path, "rb"))
                    assert len(result_values[0]) > 0
                    assert len(result_values[1]) > 0
                    os.remove(output_path)

                    # Gather the data into arrays.
                    if x_inputs_arrays != []:
                        assert result_values[0].shape[1:] == x_inputs_arrays[-1].shape[1:], str(
                            result_values[0].shape) + " vs " + str(x_inputs_arrays[-1].shape)
                    if y_outputs_arrays != []:
                        assert result_values[1].shape[1:] == y_outputs_arrays[-1].shape[1:], str(
                            result_values[1].shape) + " vs " + str(y_outputs_arrays[-1].shape)
                    x_inputs_arrays.append(result_values[0])
                    y_outputs_arrays.append(result_values[1])
                    if yield_file_paths is True:
                        file_paths_arrays.append(result_values[2])
                    else:
                        file_paths_arrays.append([])

                x_inputs = np.concatenate(x_inputs_arrays)
                y_outputs = np.concatenate(y_outputs_arrays)
                file_paths = np.concatenate(file_paths_arrays)
                assert len(x_inputs) == size
                assert len(y_outputs) == size

                # Done.
                if yield_file_paths is False:
                    yield x_inputs, y_outputs
                else:
                    yield x_inputs, y_outputs, file_paths

            else:
                raise Exception("Unexpected value for 'multiprocessing_jobs' "
                                + str(multiprocessing_jobs))

    def _load_pointcloud(self, pcd_path, preprocess=True, augmentation=True):

        pointcloud = self.pointcloud_cache.get(pcd_path, [])
        if pointcloud is None:
            pointcloud = PyntCloud.from_file(pcd_path).points.values

            if self.pointcloud_target_size is not None and preprocess is True:
                pointcloud = np.array(pointcloud)[:, 0:3]  # Drop confidence.
                pointcloud = pointcloud[:self.pointcloud_target_size]
                if len(pointcloud) < self.pointcloud_target_size:
                    zeros = np.zeros((self.pointcloud_target_size - len(pointcloud), 4))
                    pointcloud = np.concatenate([pointcloud, zeros])

            if self.pointcloud_random_rotation is True and augmentation is True:
                numpy_points = pointcloud[:, 0:3]
                numpy_points = self._rotate_point_cloud(numpy_points)
                pointcloud[:, 0:3] = numpy_points

            #self.pointcloud_cache[pcd_path] = pointcloud
        return pointcloud

    def _load_voxelgrid(self, pcd_path, preprocess=True, augmentation=True):
        voxelgrid = self.voxelgrid_cache.get(pcd_path, [])
        if voxelgrid == []:

            # Load the pointcloud.
            point_cloud = PyntCloud.from_file(pcd_path)
            if self.voxelgrid_random_rotation is True and augmentation is True:
                points = point_cloud.points
                numpy_points = points.values[:, 0:3]
                numpy_points = self._rotate_point_cloud(numpy_points)
                points.iloc[:, 0:3] = numpy_points
                point_cloud.points = points

            # Create voxelgrid from pointcloud.
            voxelgrid_id = point_cloud.add_structure(
                "voxelgrid",
                size_x=self.voxel_size_meters,
                size_y=self.voxel_size_meters,
                size_z=self.voxel_size_meters)
            voxelgrid = point_cloud.structures[voxelgrid_id].get_feature_vector(mode="density")

            # Do the preprocessing.
            if preprocess is True:
                voxelgrid = utils.ensure_voxelgrid_shape(voxelgrid, self.voxelgrid_target_shape)
                assert voxelgrid.shape == self.voxelgrid_target_shape

            #self.voxelgrid_cache[pcd_path] = voxelgrid # TODO cache is turned off because of you know why...

        return voxelgrid

    def _rotate_point_cloud(self, point_cloud):

        rotation_angle = np.random.uniform() * 2 * np.pi
        cosval = np.cos(rotation_angle)
        sinval = np.sin(rotation_angle)
        rotation_matrix = np.array([[cosval, sinval, 0],
                                    [-sinval, cosval, 0],
                                    [0, 0, 1]])

        rotated_data = np.zeros(point_cloud.shape, dtype=np.float32)
        for k in range(point_cloud.shape[0]):

            shape_pc = point_cloud[k, ...]
            rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix)

        return rotated_data


def create_datagenerator_from_parameters(dataset_path, dataset_parameters):
    logger.info("Creating data-generator...")
    datagenerator = ETLDataGenerator(
        dataset_path=dataset_path,
        input_type=dataset_parameters["input_type"],
        #output_targets=dataset_parameters["output_targets"],
        sequence_length=dataset_parameters.get("sequence_length", 0),
        voxelgrid_target_shape=dataset_parameters.get("voxelgrid_target_shape", None),
        voxel_size_meters=dataset_parameters.get("voxel_size_meters", None),
        voxelgrid_random_rotation=dataset_parameters.get("voxelgrid_random_rotation", None),
        pointcloud_target_size=dataset_parameters.get("pointcloud_target_size", None),
        pointcloud_random_rotation=dataset_parameters.get("pointcloud_random_rotation", None)
    )
    #datagenerator.print_statistics()
    return datagenerator


def get_dataset_path(root_path="../data/etl"):
    if os.path.exists("etldatasetpath.txt"):
        with open("etldatasetpath.txt", "r") as file:
            dataset_path = file.read().replace("\n", "")
    else:
        # Finding the latest.
        dataset_paths = glob.glob(os.path.join(root_path, "*"))
        dataset_paths = [
            dataset_path for dataset_path in dataset_paths if os.path.isdir(dataset_path)]
        dataset_path = list(reversed(sorted(dataset_paths)))[0]

    return dataset_path


def generate_data(class_self, size, qrcodes_to_use, verbose, yield_file_paths, output_queue):
    if verbose is True:
        logger.info("Generating QR-codes to be used: %s", qrcodes_to_use)

    assert size != 0

    x_inputs = []
    y_outputs = []
    file_paths = []

    if verbose is True:
        bar = progressbar.ProgressBar(max_value=size)
    while len(x_inputs) < size:

        # Get a random QR-code.
        qrcode = random.choice(qrcodes_to_use)

        # Get targets and paths randomly.
        if qrcode not in class_self.qrcodes_dictionary.keys():
            continue
        pcd_paths, jpg_paths, targets = random.choice(class_self.qrcodes_dictionary[qrcode])

        # No pointclouds.
        if len(pcd_paths) == 0:
            continue

        # Get a sample.
        x_input = None
        y_output = None
        file_path = None

        # Get the input. Not dealing with sequences.
        if class_self.sequence_length == 0:
            x_input, file_path = get_input(class_self, jpg_paths, pcd_paths)

        # Get the input. Dealing with sequences here.
        else:
            count = 0
            x_input, file_path = [], []
            while count != class_self.sequence_length:
                x, f = get_input(class_self, jpg_paths, pcd_paths)
                if x is not None and f is not None:
                    x_input.append(x)
                    file_path.append(f)
                    count += 1
            x_input = np.array(x_input)
            file_path = np.array(file_path)

        # Set the output.
        y_output = targets

        # Got a proper sample.
        if x_input is not None and y_output is not None and file_path is not None:
            x_inputs.append(x_input)
            y_outputs.append(y_output)
            file_paths.append(file_path)

        assert len(x_inputs) == len(y_outputs)
        assert len(y_outputs) == len(file_paths)

        if verbose is True:
            bar.update(len(x_inputs))

    if verbose is True:
        bar.finish()

    assert len(x_inputs) == size
    assert len(y_outputs) == size

    # Turn everything into ndarrays.
    x_inputs = np.array(x_inputs)
    y_outputs = np.array(y_outputs)

    # Prepare result values.
    assert len(x_inputs) == size
    assert len(y_outputs) == size
    if yield_file_paths is False:
        return_values = (x_inputs, y_outputs)
    else:
        return_values = (x_inputs, y_outputs, file_paths)

    # This is used in multiprocessing. Creates a pickle file and puts the data there.
    if output_queue is not None:
        output_path = uuid.uuid4().hex + ".p"
        pickle.dump(return_values, open(output_path, "wb"))
        output_queue.put(output_path)
    else:
        return return_values


def get_input(class_self, jpg_paths, pcd_paths):
    # Get a random image.
    if class_self.input_type == "image":
        if len(jpg_paths) == 0:
            logger.info("777")
            return None, None
        jpg_path = random.choice(jpg_paths)
        image = class_self._load_image(jpg_path)
        file_path = jpg_path
        x_input = image

    # Get a random voxelgrid.
    elif class_self.input_type == "voxelgrid":
        if len(pcd_paths) == 0:
            return None, None
        pcd_path = random.choice(pcd_paths)
        try:
            voxelgrid = class_self._load_voxelgrid(pcd_path)
            file_path = pcd_path
            x_input = voxelgrid
        except Exception as e:
            print(e)
            return None, None

    # Get a random pointcloud.
    elif class_self.input_type == "pointcloud":
        if len(pcd_paths) == 0:
            return None, None
        pcd_path = random.choice(pcd_paths)
        try:
            pointcloud = class_self._load_pointcloud(pcd_path)
            file_path = pcd_path
            x_input = pointcloud
        except Exception as e:
            print(e)
            return None, None

    # Should not happen.
    else:
        raise Exception("Unknown input_type: " + class_self.input_type)

    return x_input, file_path
