import argparse
import csv
import json
import os
import re
from datetime import datetime
import numpy as np
from os import path
from scipy.spatial.distance import cdist

# Matt Lang, 2021, langm@mail.muni.cz
# This script converts the camera data generated by Antarstick
# Input: Cross reference JSON, and camera JSON and results csv per camera.
#
# Output: JSON database compatible with Antarstick inspection app.
#
# IMPORTANT: Script now works only for two camera per location setup!


class CameraLookup:
    # Connect ID of the camera to the name.
    table = []

    def __init__(self):
        pass

    def addCamera(self, ID, label):
        self.table.append((ID, label))

    def getID(self, label):
        for element in self.table:
            if element[1] == label:
                return element[0]
        return None

    def getLabel(self, ID):
        for element in self.table:
            if element[0] == ID:
                return element[1]
        return None

    def getIDs(self):
        return list(map(lambda x: x[0], self.table))

    def getLabels(self):
        return list(map(lambda x: x[1], self.table))


class StickLookup:
    table = {}
    # globalID: [(localID, label, camera)]

    def __init__(self):
        pass

    def addStick(self, globalID, localID, label, camera):
        if not globalID in self.table:
            self.table[globalID] = []
        self.table[globalID].append((localID, label, camera))

    def getGlobalID(self, label, camera):
        for stick in self.table:
            views = self.table[stick]
            for view in views:
                if view[1] == label and view[2] == camera:
                    return stick
        return None

    def getLocalID(self, globalID, camera):
        for stick in self.table[globalID]:
            if stick[2] == camera:
                return stick[0]

    def getGlobalIDs(self):
        return list(self.table)

    def getLabelsByCamera(self, camera):
        res = []
        for entry in self.table:
            for stick in self.table[entry]:
                if stick[2] == camera:
                    res.append(stick[1])
        return res

    def globalIDExists(self, globalID):
        if globalID in self.table:
            return True
        return False

    def getLabelsByGlobalID(self, globalID):
        if globalID not in self.table:
            raise KeyError(
                'Stick with globalID {globalID} not found in lookup table.')
        result = {}
        for view in self.table[globalID]:
            result[view[2]] = view[1]
        return result

    def buildUnifiedLabels(self, prefix):
        result = {}
        uid = 1
        for id, infos in self.table.items():
            labels = {}
            for _, _, cam in infos:
                labels[cam] = prefix + str(uid)
            uid += 1
            result[id] = labels
        return result


def readResultsCsv(csvFilePath: str) -> dict:
    data = {}
    with open(csvFilePath, encoding='utf-8') as csvf:
        csvReader = csv.DictReader(csvf)
        for rows in csvReader:
            key = rows['image_name']
            data[key] = rows
        return data


def readJson(jsonFilePath: str) -> dict:
    with open(jsonFilePath, 'r', encoding='utf-8') as jsonf:
        data = json.load(jsonf)
        return data


def writeJson(jsonFilePath: str, database: dict):
    with open(jsonFilePath, 'w', encoding='utf-8') as jsonf:
        json.dump(database, jsonf, indent=2)


def getClosestTimestamp(timestamps0, timestamps1, maxDelta=-1):
    # TODO: Implement max delta.

    # Retype list of strings to numbers
    list0 = list(
        map(lambda x: datetime.fromisoformat(x).timestamp(), timestamps0))
    list1 = list(
        map(lambda x: datetime.fromisoformat(x).timestamp(), timestamps1))

    array0 = np.transpose(np.array(list0, ndmin=2))
    array1 = np.transpose(np.array(list1, ndmin=2))

    # Compute distance matrix
    matrix = cdist(array0, array1)
    index0 = matrix.argmin(axis=1)
    index1 = matrix.argmin(axis=0)

    return(list(index0), list(index1))


def main(config):
    # Open and read cross reference JSON
    # Look for the first JSON file in the folder.
    # That should correspond to cross reference file.
    xRef = readJson(config.dataset_json)

    # Get lookup table of camera labels and IDs
    cameraLUT = CameraLookup()
    for folder in xRef['cameras_ids']:
        cameraLUT.addCamera(xRef['cameras_ids'][folder], path.basename(folder))

    # Open and read camera.json for each camera
    # Open and read results.csv for each camera
    cameras = {}
    results = {}
    for folder in cameraLUT.getLabels():
        cameras[folder] = readJson(config.camera_files[folder]["camera_json"])
        results[folder] = readResultsCsv(config.camera_files[folder]["results_csv"])

    # Get list of timestamps from results
    timestampsList0 = list(
        map(lambda x: results[cameraLUT.getLabels()[0]][x]['date_time'], results[cameraLUT.getLabels()[0]]))
    timestampsList1 = list(
        map(lambda x: results[cameraLUT.getLabels()[-1]][x]['date_time'], results[cameraLUT.getLabels()[-1]]))

    # Find closest timestamp of the other camera to fill reciprocal timestamps
    (indexesList0, indexesList1) = getClosestTimestamp(
        timestampsList0, timestampsList1)
    timestampLists = {cameraLUT.getLabels()[0]: timestampsList0, cameraLUT.getLabels()[
        -1]: timestampsList1}
    timestampIndexes = {cameraLUT.getLabels()[0]: indexesList0, cameraLUT.getLabels()[
        -1]: indexesList1}
    imageNameLists = {cameraLUT.getLabels(
    )[0]: list(results[cameraLUT.getLabels()[0]]), cameraLUT.getLabels(
    )[-1]: list(results[cameraLUT.getLabels()[-1]])}

    # Get lookup table of global stick IDs to their labels
    stickLUT = StickLookup()
    for cam in cameras:
        sticks = cameras[cam]['sticks']
        for stick in sticks:
            # Global stick ID should have been an unique identifier of the stick
            # across all image datasets. Global ID has to be created when parsing
            # the dataset for viewing.

            # Unique identifier of the _IMAGE_ of a stick in one dataset
            localID = stick['id']
            # Label of the _IMAGE_ of a stick in one dataset
            label = stick['label']
            globalID = xRef['stick_local_to_global_ids'][str(
                cameraLUT.getID(cam))][str(localID)]  # Unique identifier of the _IMAGE_ of a stick in all datasets

            # Find corresponding globalID of the stick in the second camera
            globalIDSecondCamera = None
            for relation in xRef['stick_views_map']:
                if relation[0] == globalID:
                    globalIDSecondCamera = relation[1]
                    break
                if relation[1] == globalID:
                    globalIDSecondCamera = relation[0]
                    break

            # Create new entry in the lookup table or find an existing one for the stick
            if stickLUT.globalIDExists(globalIDSecondCamera):
                stickLUT.addStick(globalIDSecondCamera,
                                  localID, label, cam)
            else:
                stickLUT.addStick(globalID, localID, label, cam)

    # Create new array and fill it with cameras
    outDatabase = []

    # Find max snow height for metadata and chart height.
    maxSnowHeight = 0
    # Find which view is larger to determine the primary view.
    maxHeightPx = {}

    for cam in cameraLUT.getLabels():
        outCamera = {}
        outDatabase.append(outCamera)
        outCamera['id'] = cam
        # Fill each camera with images
        for image in results[cam]:
            inImage = results[cam][image]
            outImage = outCamera[image] = {}

            outImage['processedState'] = int(inImage['state'])  # -1: skipped, 0: unprocessed, 1: processed
            outImage['dateTime'] = inImage['date_time'].replace(' ', 'T')
            outImage['imageQuality'] = float(inImage['image_quality'])
            outImage['weatherConditions'] = int(inImage['weather_conditions'])
            outImage['isDay'] = inImage['is_day'].lower() == 'true'
            outImage['reciprocalImageName'] = {}
            outImage['approvedByUser'] = False      # Whether a user approved data of sticks and snow hight in the image
            outImage['discardedByUser'] = False     # Whether a user decided to discard the image and all stick data in it
            outImage['usedForTraining'] = False     # Whether the image was sent to the training data set for the sticks and snow recognition algo

            # Find reciprocal timestamps using
            # timestamIndexes and timestampLists
            for otherCam in cameraLUT.getLabels():
                if otherCam == cam:
                    continue
                thisTimePosition = timestampLists[cam].index(
                    inImage['date_time'])
                otherTimePosition = timestampIndexes[cam][thisTimePosition]
                outImage['reciprocalImageName'][otherCam] = imageNameLists[otherCam][otherTimePosition]

            # Create sticks in each image
            sticks = outImage['sticks'] = {}
            if config.debug:
                print('Save sticks '+str(stickLUT.getLabelsByCamera(cam)) +
                      ' from camera '+str(cam)+'.')
            for stick in stickLUT.getLabelsByCamera(cam):
                stickGlobalID = stickLUT.getGlobalID(stick, cam)
                outStick = sticks[stickGlobalID] = {}

                # Get top, bottom, snow_height and visible
                # parameters for each stick
                outStick['top'] = list(
                    map(int, re.findall('\d+', inImage[stick + '_top'])))
                outStick['bottom'] = list(
                    map(int, re.findall('\d+', inImage[stick + '_bottom'])))
                outStick['snowHeight'] = int(inImage[stick + '_snow_height'])
                outStick['visible'] = 'True' in inImage[stick + '_visible']

                # The classifier tool always halves the processed images in
                # both axes; therefore all coordinates into the image must
                # be multiplied by 2 to get them for the original image size.
                outStick['top'][0] *= 2
                outStick['top'][1] *= 2
                outStick['bottom'][0] *= 2
                outStick['bottom'][1] *= 2

                outStick['correctedByUser'] = False   # Whether a user manually updated this stick location or snow height in the image
                outStick['discardedByUser'] = False     # Whether a user decided to discard the stick from the image
                outStick['createdByUser'] = False       # Whether a user manually created the stick into the image

                # Store maximum snow height, for metadata.
                if maxSnowHeight < outStick['snowHeight']:
                    maxSnowHeight = outStick['snowHeight']

                # Store maximum stick height, for metadata.
                if stickGlobalID not in maxHeightPx:
                    maxHeightPx[stickGlobalID] = (
                        int(inImage[stick + '_height_px']), cam)
                else:
                    if maxHeightPx[stickGlobalID][0] < int(inImage[stick + '_height_px']):
                        maxHeightPx[stickGlobalID] = (
                            int(inImage[stick + '_height_px']), cam)

    # Create metadata dictionary
    outMetadata = {}
    outMetadata['start'] = config.year + "-01-01T00:00:00Z"
    outMetadata['end'] = config.year + "-12-31T23:59:59Z"
    outMetadata['location'] = {
        "name": config.location,
        "latitude": config.latitude,
        "longitude": config.longitude
    }
    outMetadata['maxSnowHeight'] = maxSnowHeight
    outMetadata['cameras'] = {}
    for camID in cameraLUT.getLabels():
        outMetadata['cameras'][camID] = {'x': 0, 'y': 0}
    unifiedLabels = None if config.unify_stick_labels is None else stickLUT.buildUnifiedLabels(config.unify_stick_labels)
    outMetadata['sticks'] = {}
    for stickID in stickLUT.getGlobalIDs():
        labels = stickLUT.getLabelsByGlobalID(stickID) if unifiedLabels is None else unifiedLabels[stickID]
        primaryCameraName = maxHeightPx[stickID][1]
        stickLength = float(cameras[primaryCameraName]["sticks"][stickLUT.getLocalID(stickID, primaryCameraName)]["length_cm"])
        outMetadata['sticks'][stickID] = {
            'x': 0,
            'y': 0,
            'lengthCm': stickLength,
            'labels': labels,
            'primaryCamera': primaryCameraName
        }

    writeJson(config.classifier_results_json,
              {'metadata': outMetadata, 'data': outDatabase})


def build_script_config():
    parser = argparse.ArgumentParser(
        description="The script converts output data (results.csv and camera.json) from\n"
                    "the Antarstick stick and snow recognition tool into a single json\n"
                    "file (classifier_results.json) which is then used by the Antarstick-js\n"
                    "visualisation tool.",
        )
    parser.add_argument(
        "location", type=str, default=None,
        help="A location where the data were collected, e.g. \"Monolith Lake\", \"Keller Stream\"."
        )
    parser.add_argument(
        "latitude", type=float, default=None,
        help="A latitude where the data were collected, e.g. -63.80056941579183."
        )
    parser.add_argument(
        "longitude", type=float, default=None,
        help="A longitude where the data were collected, e.g. -57.88234470557961."
        )
    parser.add_argument(
        "year", type=str, default=None,
        help="A year when the data were collected at the given location, e.g. 2011."
        )
    parser.add_argument(
        "--root-data-dir", type=str, default=None,
        help="A root data directory. If not specified, then the script searches for the 'antarstic-data' "
             "directory along the path of the current working directory."
        )
    parser.add_argument(
        "--unify-stick-labels", type=str,
        help="When present, the stick are given new unique names starting with a given prefix (passed to this option). "
             "Most importantly, if some stick appears in multiple cameras, then this option assures that the stick "
             "will get the same unique name in all those cameras."
        )
    parser.add_argument(
        "--debug", action='store_true',
        help="Whether to print debug messages or not."
        )
    config = parser.parse_args()

    if config.root_data_dir is None:
        start_path = os.path.abspath(os.getcwd())
        config.root_data_dir = start_path
        while len(config.root_data_dir) > 0 and not os.path.isdir(os.path.join(config.root_data_dir, "antarstick-data")):
            config.root_data_dir = os.path.dirname(config.root_data_dir)
        config.root_data_dir = os.path.join(config.root_data_dir, "antarstick-data")
        if not os.path.isdir(config.root_data_dir):
            raise Exception("Cannot find the data root directory 'antarstick-data' along the directory: " + start_path)
    else:
        config.root_data_dir = os.path.abspath(config.root_data_dir)
        if not os.path.basename(config.root_data_dir).startswith("antarstick-data") or not os.path.isdir(config.root_data_dir):
            raise Exception("The passed data root directory '" + config.root_data_dir + "' is not valid data root.")

    config.data_dir = os.path.join(config.root_data_dir, config.location, config.year)
    if not os.path.isdir(config.data_dir):
        raise Exception("The data directory '" + config.data_dir + "' is not valid.")

    config.classifier_results_json = os.path.join(config.data_dir, "classifier_results.json")
    # if os.path.isfile(config.classifier_results_json):
    #     raise Exception("Nothing to do. The results file exists: " + os.path.join(config.data_dir, "classifier_results.json"))

    config.camera_images_root_dir = os.path.join(config.data_dir, "camera_images")
    if not os.path.isdir(config.camera_images_root_dir):
        raise Exception("The camera images directory '" + config.camera_images_root_dir + "' is not valid.")

    config.dataset_json = os.path.join(config.camera_images_root_dir, "dataset.json")
    if not os.path.isfile(config.dataset_json):
        raise Exception("Cannot find file '" + config.dataset_json + "'.")

    config.camera_files = {}
    with open(config.dataset_json, "r") as f:
        dataset = json.load(f)
        for p in dataset["cameras_ids"]:
            camera_name = os.path.basename(p)
            config.camera_files[camera_name] = {
                "camera_json": os.path.join(config.camera_images_root_dir, camera_name, "camera.json"),
                "results_csv": os.path.join(config.camera_images_root_dir, camera_name, "_results", "results.csv")
            }
            for _, q in config.camera_files[camera_name].items():
                if not os.path.isfile(q):
                    raise Exception("Cannot find file '" + q + "'.")

    return config


if __name__ == "__main__":
    main(build_script_config())
