"""

File input/output for localization data.

"""
from __future__ import annotations
import logging
from enum import Enum

import pandas as pd

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.locan_io.locdata.rapidstorm_io import load_rapidSTORM_file, load_rapidSTORM_track_file
from locan.locan_io.locdata.thunderstorm_io import load_thunderstorm_file
from locan.locan_io.locdata.elyra_io import load_Elyra_file
from locan.locan_io.locdata.nanoimager_io import load_Nanoimager_file
from locan.locan_io.locdata.asdf_io import load_asdf_file
from locan.locan_io.locdata.decode_io import load_decode_file
from locan.locan_io.locdata.smap_io import load_SMAP_file
from locan.locan_io.locdata.smlm_io import load_SMLM_file


__all__ = ['load_txt_file', 'load_locdata']

logger = logging.getLogger(__name__)


def load_txt_file(path, sep=',', columns=None, nrows=None, property_mapping=None, convert=True, **kwargs):
    """
    Load localization data from a txt file.

    Locan column names are either supplied or read from the first line header.

    Parameters
    ----------
    path : str, bytes, os.PathLike, file-like
        File path for a localization file to load.
    sep : str
        separator between column values (Default: ',')
    columns : list of str, None
        Locan column names. If None the first line is interpreted as header (Default: None).
    nrows : int, None
        The number of localizations to load from file. None means that all available rows are loaded (Default: None).
    property_mapping : dict[str: str] or list[dict]
        Mappings between column names and locan property names
    convert : bool
        If True convert types by applying type specifications in locan.constants.PROPERTY_KEYS.
    kwargs : dict
        Other parameters passed to `pandas.read_csv()`.

    Returns
    -------
    LocData
        A new instance of LocData with all localizations.
    """
    # define columns
    if columns is None:
        dataframe = pd.read_csv(path, sep=sep, nrows=nrows,
                                **dict(dict(skiprows=0), **kwargs))

        column_keys = convert_property_names(dataframe.columns, property_mapping=property_mapping)
        dataframe.columns = column_keys
    else:
        column_keys = convert_property_names(columns, property_mapping=property_mapping)
        dataframe = pd.read_csv(path, sep=sep, nrows=nrows,
                                **dict(dict(skiprows=1, names=column_keys), **kwargs))

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

    dat = LocData.from_dataframe(dataframe=dataframe)

    dat.meta.source = metadata_pb2.EXPERIMENT
    dat.meta.state = metadata_pb2.RAW
    dat.meta.file.type = metadata_pb2.CUSTOM
    dat.meta.file.path = str(path)

    del dat.meta.history[:]
    dat.meta.history.add(name='load_txt_file',
                         parameter='path={}, sep={}, columns={}, nrows={}'.format(path, sep, columns, nrows))

    return dat


def _map_file_type_to_load_function(file_type):
    """
    Interpret user input for file_type.

    Parameters
    ----------
    file_type : int, str, locan.constants.FileType, locan.data.metadata_pb2.Metadata
        Identifier for the file type. Integer or string should be according to locan.constants.FileType.

    Returns
    -------
    callable
        Name of function for loading the localization file of `type`.
    """
    look_up_table = dict(
        load_txt_file=load_txt_file,
        load_rapidSTORM_file=load_rapidSTORM_file,
        load_Elyra_file=load_Elyra_file,
        load_thunderstorm_file=load_thunderstorm_file,
        load_asdf_file=load_asdf_file,
        load_Nanoimager_file=load_Nanoimager_file,
        load_rapidSTORM_track_file=load_rapidSTORM_track_file,
        load_SMLM_file=load_SMLM_file,
        load_decode_file=load_decode_file,
        load_SMAP_file=load_SMAP_file,
    )

    class LoadFunction(Enum):
        load_txt_file = 1
        load_rapidSTORM_file = 2
        load_Elyra_file = 3
        load_thunderstorm_file = 4
        load_asdf_file = 5
        load_Nanoimager_file = 6
        load_rapidSTORM_track_file = 7
        load_SMLM_file = 8
        load_decode_file = 9
        load_SMAP_file = 10

    try:
        if isinstance(file_type, int):
            function_name = LoadFunction(file_type).name
        elif isinstance(file_type, str):
            function_name = LoadFunction(locan.constants.FileType[file_type.upper()].value).name
        elif isinstance(file_type, locan.constants.FileType):
            function_name = LoadFunction(file_type.value).name
        elif isinstance(file_type, metadata_pb2):
            function_name = LoadFunction(file_type).name
        else:
            raise TypeError
        return look_up_table[function_name]
    except ValueError:
        raise ValueError(f'There is no load function for type {file_type}.')


def load_locdata(path, file_type=1, nrows=None, **kwargs):
    """
    Load data from localization file as specified by type.

    This function is a wrapper for read functions for the various types of SMLM data.

    Parameters
    ----------
    path : str, bytes, os.PathLike, file-like
        File path for a localization data file to load.
    file_type : int, str, locan.constants.FileType, locan.data.metadata_pb2.Metadata
        Indicator for the file type.
        Integer or string should be according to locan.constants.FileType.
    nrows : int, None
        The number of localizations to load from file. None means that all available rows are loaded.

    Returns
    -------
    LocData
        A new instance of LocData with all localizations.
    """
    return _map_file_type_to_load_function(file_type)(path=path, nrows=nrows, **kwargs)
