"""

File input/output for localization data in SMLM files.

File specifications are provided at https://github.com/imodpasteur/smlm-file-format/blob/master/specification.md.

Code is adapted from https://github.com/imodpasteur/smlm-file-format/blob/master/implementations/Python/smlm_file.py.
(MIT license)
"""

import time
import logging
import zipfile
import json

import numpy as np
import pandas as pd
from google.protobuf import json_format

from locan.locan_io.locdata import manifest_pb2
from locan.data.locdata import LocData
import locan.constants
from locan.data import metadata_pb2
from locan.locan_io.locdata.utilities import convert_property_types, convert_property_names
from locan.utils.format import _time_string

__all__ = ['manifest_format_from_locdata', 'manifest_file_info_from_locdata', 'manifest_from_locdata',
           'save_SMLM', 'load_SMLM_manifest', 'load_SMLM_header', 'load_SMLM_file']

logger = logging.getLogger(__name__)


dtype2length = {'int8': 1, 'uint8': 1, 'int16': 2, 'uint16': 2,
                'int32': 4, 'uint32': 4, 'int64': 8, 'uint64': 8,
                'float32': 4, 'float64': 8}


def manifest_format_from_locdata(locdata):
    """
    Prepare a manifest["format"] protobuf from locdata.
    The manifest holds metadata for smlm files.

    Parameters
    ----------
    locdata : LocData
        The LocData object.

    Returns
    -------
    locan.locan_io.locdata.manifest_pb2.Format
        The manifest format
    """
    format = manifest_pb2.Format()

    # required by
    format.name = "smlm-table(binary)-0"
    format.type = manifest_pb2.Type.TABLE
    format.mode = manifest_pb2.Mode.BINARY
    format.columns = len(locdata.data.columns)
    format.headers.extend(locdata.data.columns)
    format.dtype.extend(
        [manifest_pb2.Dtype.Value(str(dt).upper())
         for dt in locdata.data.dtypes.values]
    )
    format.shape.extend([1] * len(locdata.data.columns))
    format.units.extend([""] * len(locdata.data.columns))  # todo: add correct units

    # optional by specifications
    format.description = "localization table generated from locan.LocData"

    return format


def manifest_file_info_from_locdata(locdata):
    """
    Prepare a manifest["file_info"] protobuf from locdata.
    The manifest holds metadata for smlm files.

    Parameters
    ----------
    locdata : LocData
        The LocData object.

    Returns
    -------
    locan.locan_io.locdata.manifest_pb2.FileInfo
        The manifest file information
    """
    file_info = manifest_pb2.FileInfo()

    # required by specifications
    file_info.name = "table-0.bin"
    file_info.type = manifest_pb2.TABLE
    file_info.format = "smlm-table(binary)-0"
    file_info.channel = "default"
    file_info.rows = len(locdata)
    for column in locdata.data.columns:
        file_info.offset[column] = 0

    # optional by specifications
    for column in locdata.data.columns:
        file_info.min[column] = locdata.data[column].min()
    for column in locdata.data.columns:
        file_info.max[column] = locdata.data[column].max()
    # file_info.exposure = 20

    return file_info


def manifest_from_locdata(locdata, return_json_string=False):
    """
    Prepare a manifest protobuf from locdata.
    The manifest holds metadata for smlm files.

    Parameters
    ----------
    locdata : LocData
        The LocData object.
    return_json_string : bool
        Flag for returning json string

    Returns
    -------
    locan.locan_io.locdata.manifest_pb2.Manifest, str
        The manifest
    """
    manifest = manifest_pb2.Manifest()
    format = manifest_format_from_locdata(locdata)
    file_info = manifest_file_info_from_locdata(locdata)

    # required by specifications
    manifest.format_version = "0.2"
    manifest.formats["smlm-table(binary)-0"].CopyFrom(format)
    manifest.files.append(file_info)

    # recommended by specifications
    manifest.name = locdata.meta.identifier
    manifest.description = locdata.meta.comment
    manifest.tags.append("SMLM")
    manifest.thumbnail = ""
    manifest.sample = ""
    manifest.labeling = ""
    manifest.date = _time_string(time.time())

    # optional by specifications
    manifest.author = ""
    manifest.citation = ""
    manifest.email = ""
    manifest.locdata_meta.CopyFrom(locdata.meta)

    if return_json_string:
        json_string = json_format.MessageToJson(manifest,
                                                preserving_proto_field_name=True,
                                                including_default_value_fields=False)
        json_string = _change_upper_to_lower_keys(json_string)
        return json_string
    else:
        return manifest


def _change_upper_to_lower_keys(json_string: str) -> str:
    """ Switch selected key words in json_string from upper to lower letters."""
    json_string = json_string.replace("BINARY", "binary")
    json_string = json_string.replace("TEXT", "text")
    json_string = json_string.replace("TABLE", "table")
    json_string = json_string.replace("IMAGE", "image")
    json_string = json_string.replace("INT", "int")
    json_string = json_string.replace("UINT", "uint")
    json_string = json_string.replace("FLOAT", "float")
    return json_string


def save_SMLM(locdata, path, manifest=None):
    """
    Save LocData attributes in a SMLM single-molecule
    localization (zip) file with manifest.json (version 0.2).

    In the smlm file format we store metadata as a human-
    readable manifest. The data is stored as byte string.

    Note
    ----
    Only selected LocData attributes are saved.
    These are: 'data', 'columns', 'meta'.

    Parameters
    ----------
    locdata : LocData
        The LocData object to be saved.
    path : str, bytes, os.PathLike, file-like
        File path including file name to save to.
    manifest : locan.locan_io.locdata.manifest_pb2.Manifest
        Protobuf with manifest to use instead of an autogenerated manifest.
    """
    if manifest is None:
        manifest_json = manifest_from_locdata(locdata, return_json_string=True)
    elif isinstance(manifest, str):
        manifest_json = manifest
    elif isinstance(manifest, manifest_pb2.Manifest):
        manifest_json = json_format.MessageToJson(manifest,
                                                  preserving_proto_field_name=True,
                                                  including_default_value_fields=False
                                                  )
        manifest_json = _change_upper_to_lower_keys(manifest_json)
    else:
        raise TypeError("Type of manifest is not correct.")

    byte_string = b"".join([locdata.data[column][i].tobytes()
                            for i in range(len(locdata))
                            for column in locdata.data.columns
                            ])

    # save zip file
    with zipfile.ZipFile(path, 'w', compression=zipfile.ZIP_DEFLATED, allowZip64=True) as zf:
        # write manifest
        zf.writestr("manifest.json", manifest_json)
        # write files
        zf.writestr("table-0.bin", byte_string)


def load_SMLM_manifest(path):
    """
    Read manifest.json (version 0.2) from a SMLM single-molecule localization (zip) file.

    Parameters
    ----------
    path : str, bytes, os.PathLike, file-like
        File path for a SMLM file to load.

    Returns
    -------
    dictionary
        manifest in json format
    """
    zf = zipfile.ZipFile(path, 'r')
    file_names = zf.namelist()
    if "manifest.json" not in file_names:
        raise Exception('invalid file: no manifest.json found in the smlm file.')
    manifest = json.loads(zf.read("manifest.json"))
    assert manifest['format_version'] == '0.2'
    return manifest


def load_SMLM_header(path):
    """
    Read header (manifest) from a SMLM single-molecule localization file and identify column names.

    Parameters
    ----------
    path : str, bytes, os.PathLike, file-like
        File path for a SMLM file to load.

    Returns
    -------
    list(str)
        A list of dataset property keys as derived from the SMLM identifiers.
    """
    zf = zipfile.ZipFile(path, 'r')
    file_names = zf.namelist()
    if "manifest.json" not in file_names:
        raise Exception('invalid file: no manifest.json found in the smlm file.')
    manifest = json.loads(zf.read("manifest.json"))
    assert manifest['format_version'] == '0.2'

    locdata_columns_list = []
    for file_info in manifest['files']:
        if file_info['type'] != "table":
            logger.info('ignore file with type: %s', file_info['type'])
        else:
            name = file_info['name']
            logger.info(f'loading table {name}....')
            format_key = file_info['format']
            file_format = manifest['formats'][format_key]
            if name not in file_names:
                logger.error('ERROR: Did not find %s in zip file', file_info['name'])
            else:
                headers = file_format['headers']

                column_keys = convert_property_names(properties=headers, property_mapping=locan.constants.SMLM_KEYS)
                locdata_columns_list.append(column_keys)

    if len(locdata_columns_list) == 1:
        return locdata_columns_list[0]
    else:
        return locdata_columns_list


def load_SMLM_file(path, nrows=None, convert=True):
    """
    Load data from a SMLM single-molecule localization file.

    Parameters
    ----------
    path : str, bytes, os.PathLike, file-like
        File path for a SMLM file to load.
    nrows : int, None
        The number of localizations to load from file. None means that all available rows are loaded.
    convert : bool
        If True convert types by applying type specifications in locan.constants.PROPERTY_KEYS.

    Returns
    -------
    LocData, list(LocData)
        A new instance of LocData with all localizations. Returns a list of LocData if multiple tables are found.
    """
    zf = zipfile.ZipFile(path, 'r')
    file_names = zf.namelist()
    if "manifest.json" not in file_names:
        raise Exception('invalid file: no manifest.json found in the smlm file.')
    manifest = json.loads(zf.read("manifest.json"))
    assert manifest['format_version'] == '0.2'

    locdatas = []
    for file_info in manifest['files']:
        if file_info['type'] != "table":
            logger.info('ignore file with type: %s', file_info['type'])
        else:
            name = file_info['name']
            logger.debug(f'start loading {name} ...')
            format_key = file_info['format']
            file_format = manifest['formats'][format_key]
            if file_format['mode'] != 'binary':
                raise Exception(f"format mode {file_format['mode']} not supported.")
            else:
                try:
                    table_file = zf.read(file_info['name'])
                except KeyError:
                    logger.error('ERROR: Did not find %s in zip file', file_info['name'])
                    continue
                else:
                    logger.debug(f'loading {len(table_file)} bytes')
                    headers = file_format['headers']
                    dtype = file_format['dtype']
                    shape = file_format['shape']
                    cols = len(headers)
                    rows = nrows if nrows is not None else int(file_info['rows'])

                    logger.debug('columns: %s', headers)
                    logger.debug('rows: %s, columns: %s', rows, cols)
                    assert len(headers) == len(dtype) == len(shape)

                    rowLen = sum(dtype2length[dtype[i]] for i, header in enumerate(headers))
                    tableDict = {}
                    byteOffset = 0
                    for i, header in enumerate(headers):
                        tableDict[header] = np.ndarray((rows,), buffer=table_file, dtype=dtype[i], offset=byteOffset,
                                                       order='C', strides=(rowLen,))
                        byteOffset += dtype2length[dtype[i]]
                    logger.debug(f'finished loading {name}')

                    dataframe = pd.DataFrame.from_dict(tableDict)

                    column_keys = {}
                    for header in headers:
                        if header in locan.constants.SMLM_KEYS:
                            column_keys[header] = locan.constants.SMLM_KEYS[header]
                        elif header in locan.constants.RAPIDSTORM_KEYS:
                            column_keys[header] = locan.constants.RAPIDSTORM_KEYS[header]
                        else:
                            logger.warning(f'Column {header} is not a Locan property standard.')
                            column_keys[header] = header
                    dataframe.rename(columns=column_keys, inplace=True)

                    if convert:
                        dataframe = convert_property_types(dataframe, types=locan.constants.PROPERTY_KEYS)

                    locdata = LocData.from_dataframe(dataframe=dataframe)

                    # todo: check on taking correct metadata
                    locdata.meta.source = metadata_pb2.EXPERIMENT
                    locdata.meta.state = metadata_pb2.RAW
                    locdata.meta.file.type = metadata_pb2.SMLM
                    locdata.meta.file.path = str(path)

                    for property_ in sorted(list(set(dataframe.columns).intersection(
                            {'position_x', 'position_y', 'position_z'}))):
                        locdata.meta.localization_properties.add(name=property_, unit="nm", type="float")

                    del locdata.meta.history[:]
                    locdata.meta.history.add(name='load_SMLM_file', parameter=f'path={path}, nrows={nrows}')

                    locdatas.append(locdata)

        if len(locdatas) == 1:
            return locdatas[0]
        else:
            return locdatas
