"""

File input/output for localization data.

"""
from __future__ import annotations

import logging
from enum import Enum

import pandas as pd

import locan.constants
from locan.data import metadata_pb2
from locan.data.locdata import LocData
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.elyra_io import load_Elyra_file
from locan.locan_io.locdata.nanoimager_io import load_Nanoimager_file
from locan.locan_io.locdata.rapidstorm_io import (
    load_rapidSTORM_file,
    load_rapidSTORM_track_file,
)
from locan.locan_io.locdata.smap_io import load_SMAP_file
from locan.locan_io.locdata.smlm_io import load_SMLM_file
from locan.locan_io.locdata.thunderstorm_io import load_thunderstorm_file
from locan.locan_io.locdata.utilities import (
    convert_property_names,
    convert_property_types,
)

__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)
