from os import listdir
from os.path import isfile, isdir, join, isabs
import json
import pandas as pd
from numpy import ndarray
import numpy as np
from ..selections import store_selection
from ..pandas import subset_by_DataFrame
import warnings

class Dataset():
    """Dataset object allow to facilitate manipulation of data generated by bacmman.

    Parameters
    ----------
    path : str
        path to the folder containing the configuration file
    data_path : str
        path to the folder containing the data files (optional, if None -> path is used)
    filter : str or callable
        filter on the name of the dataset (optional)

    Attributes
    ----------
    config_name : str
        dataset name
    object_class_names : list of str
        names of the object classes of the dataset. By default names are found in the configuration file, but can be renamed using the set_object_class_name method
    data : dict
        datatable for each object class
    path : str
        path to the folder containing the configuration file
    data_path : str
        (optional) path to the folder containing the data files. if None, path is used. If data_path is not absolute it is considered as relative to path
    filter: callable or str
        filter on dataset name. if str: test if filter is contained in dataset name

    """
    def __init__(self, path:str, data_path:str = None, filter=None, name:str=None, raise_error:bool=True):
        self.path = path
        if data_path is not None:
            self.data_path = data_path if isabs(data_path) else join(path, data_path)
        else:
            self.data_path = path
        self.config_name = get_dataset_name(path, filter)
        if self.config_name is not None: # inspect config file
            cf = join(path, self.config_name+"_config.json")
            with open(cf, errors='ignore') as f:
                try:
                    conf = json.load(f)
                    oc = conf["structures"]
                    if isinstance(oc, list):
                        self.object_class_names = [c["name"] for c in oc]  # new format
                    elif "list" in oc:
                        self.object_class_names = [c["name"] for c in oc["list"]]
                    else:
                        raise IOError(f"Invalid configuration format: could not find object class names in {oc}")
                    # ensure unicity of names
                    seen = set()
                    self.object_class_names = [_increment_name(x, seen) for x in self.object_class_names]
                except Exception as e:
                    print(f"Error trying to load configuration file: {f}")
                    raise e
        else:
            self.object_class_names = None
            if raise_error:
                raise IOError(f"Invalid dataset directory : {path}")
        self.name = self.config_name if name is None else name
        self.data = {}
        self.selections = None

    def set_object_class_name(self, old_name, new_name:str):
        """Modifies the name of an object class.

        Parameters
        ----------
        old_name : str or int
            either name or index of the old object class to modify
        new_name : str
            new name of the object class

        """
        old_name = self._get_object_class_index(old_name)
        assert old_name < len(self.object_class_names), f"invalid object class idx, should be <{len(self.object_class_names)}"
        if self.object_class_names[old_name]!=new_name:
            assert new_name not in self.object_class_names, "invalid new object_class name: already exists"
            self.object_class_names[old_name] = new_name

    def _get_object_class_name(self, object_class):
        if not isinstance(object_class, str):
            assert object_class<len(self.object_class_names), f"invalid object_class index, should be < {len(self.object_class_names)}"
            return self.object_class_names[object_class]
        else:
            return object_class

    def _get_object_class_index(self, object_class):
        if isinstance(object_class, str):
            assert object_class in self.object_class_names, f"invalid object_class name, should be in {self.object_class_names}"
            return self.object_class_names.index(object_class)
        else:
            return object_class

    def _get_data_file_path(self, object_class):
        return join(self.data_path, f"{self.config_name}_{self._get_object_class_index(object_class)}.csv")

    def _open_data(self, object_class, add_dataset_name_column=False, **kwargs):
        default_dtype = {'Position': str, 'PositionIdx':np.int16, 'Frame':np.int16, "Idx":np.int16,
                         'Indices':str, 'TrackHeadIndices':str, 'ParentTrackHeadIndices':str, 'Prev':str,
                         'Next':str, 'TrackErrorNext':bool, 'TrackErrorPrev':bool}
        if 'dtype' in kwargs and kwargs['dtype'] is not None:
            kwargs['dtype'].update(default_dtype)
        else:
            kwargs['dtype'] = default_dtype
        data = pd.read_csv(self._get_data_file_path(object_class), sep=';', **kwargs) #
        if add_dataset_name_column:
            data["DatasetName"] = self.name
        return data

    def get_data(self, object_class, selection=None, copy:bool=True, cache:bool=True, **kwargs):
        """Open measurement table of an object class as pandas dataframe.
           the file is supposed to be located in the dataset path.
           Note that when calling this method, the dataframe is stored in the dataset object. call _open_data to open the dataframe without storing it.

        Parameters
        ----------
        object_class : str or int
            either name or index of the object class
        selection: str or dataframe
            str : name of a selection contained in this dataset
            dataframe : dataframe that contains columns
        copy: bool
            only used if selection is None. returned a copy of the dataframe to avoid modifying the cached one
        cache: bool
            cache the dataframe
        Returns
        -------
        pandas dataframe
            Description of returned object.

        """
        object_class_name = self._get_object_class_name(object_class)
        object_class_idx = self._get_object_class_index(object_class)
        if object_class_name not in self.data:
            data = self._open_data(object_class, **kwargs)
            if cache:
                self.data[object_class_name] = data
        else:
            data = self.data[object_class_name]
        if selection is not None:
            if isinstance(selection, str):
                selections = self.get_selections()
                if selections is None:
                    raise ValueError(f"No selections found in dataset {self}")
                selection = selections[(selections.ObjectClassIdx == object_class_idx) & (selections.SelectionName == selection)]
                if selection.shape[0]==0:
                    warnings.warn(f"Selection {selection} not found for object class: {object_class}")
            elif isinstance(selection, (list, tuple)):
                selections = self.get_selections()
                selections = selections[selections.ObjectClassIdx == object_class_idx]
                all_selections = [selections[selections.SelectionName == s] if isinstance(s, str) else s for s in selection]
                for s, sname in zip(all_selections, selection):
                    if isinstance(sname, str) and s.shape[0]==0:
                        warnings.warn(f"Selection: {sname} not found for object class: {object_class}")
                selection = pd.concat(all_selections)
                selection = selection.drop_duplicates(["Position", "Indices"])
            return subset_by_DataFrame(data, selection, on=["Position", "Indices"])
        else:
            return data.copy() if copy else data

    def _get_selections_file_path(self):
        return join(self.data_path, f"{self.config_name}_Selections.csv")

    def _open_selections(self, add_dataset_name_column=False, **kwargs):
        default_dtype = {'Position': 'str', 'PositionIdx':np.int16, 'ObjectClassIdx':np.int16, 'Indices':'str', 'Frame':np.int16, 'SelectionName':'str'}
        if 'dtype' in kwargs and kwargs['dtype'] is not None:
            kwargs['dtype'].update(default_dtype)
        else:
            kwargs['dtype'] = default_dtype
        fp = self._get_selections_file_path()
        if not isfile(fp):
            columns = ['Position', 'PositionIdx', 'ObjectClassIdx', 'Indices', 'Frame', 'SelectionName']
            if add_dataset_name_column:
                columns.append('DatasetName')
            return pd.DataFrame(columns=columns)
        else :
            data = pd.read_csv(fp, sep=';', **kwargs) #
        if add_dataset_name_column:
            data["DatasetName"] = self.name
        return data

    def append_selections_to_data(self, selections=None):
        """For each selection, adds a boolean column to the dataFrame with the name of the selection, with True where the object defined by (Position, Indices) is included in the selection. Object class is infered from the selection name.

        Parameters
        ----------
        selections : list / tuple / ndarray of str, or str
            Selection names to add. If None, all selections will be added

        Returns None

        """
        all_selections = self.get_selections()
        if all_selections is None:
            warnings.warn(f"No selections found in dataset: {self.name}")
            return
        if selections is None:
            selections = all_selections["SelectionName"].unique()
            print(type(selections))
        if isinstance(selections, str):
            selections = [selections]
        if not isinstance(selections, (tuple, list, ndarray)):
            raise ValueError("selections should be a str, tuple or list of str, or ndarray")
        for sel_name in selections:
            sel = all_selections[all_selections["SelectionName"] == sel_name]
            if sel.shape[0] == 0:
                warnings.warn(f"Selection: {sel_name} not found in exported selections")
                continue
            oc_idx = sel["ObjectClassIdx"].iloc[0]
            data = self.get_data(oc_idx)
            if data is None:
                warnings.warn(f"No data found for object class: {oc_idx}, skipping selection: {sel_name}")
            else:
                merge = data.filter(["Position", "Indices"], axis=1)
                merge['copy_index'] = data.index
                merge = merge.merge(sel[["Position", "Indices"]], how="inner", copy=False)["copy_index"]
                data[sel_name] = False
                data.loc[merge, sel_name] = True

    def get_selections(self, **kwargs):
        if self.selections is None:
            self.selections = self._open_selections(**kwargs)
        return self.selections

    def store_selection(self, selection, object_class, name:str, **kwargs):
        """Save a selection (subset of objects) to BACMMAN software.

        Parameters
        ----------
        object_class : str or int
            either name or index of the object class
        selection : pandas dataframe
            dataframe containing indices and positions of objects to include in the selection
        name : name of the selection
        **kwargs : dict
            extra arguments passed to save_selection method, such as port, python_proxy_port, address

        """
        object_class_idx = self._get_object_class_index(object_class)
        store_selection(selection, self.config_name, objectClassIdx=object_class_idx, selectionName=name, dsPath=self.path, **kwargs)

    def __getitem__(self, item):
        return self.get_data(item)

    def __str__(self):
        return f"{self.name} oc={self.object_class_names} path={self.path}"+(f"data path={self.data_path}" if self.data_path!=self.path else "")

class DatasetList(Dataset):
    def __init__(self, dataset_list:list = None, path:str = None, data_path:str = None, config_path:str=None, filter = None, object_class_name_mapping:dict = None):
        if dataset_list is not None:
            self.datasets = {d.name:d for d in dataset_list}
        elif path is not None:
            if config_path is not None:
                assert data_path is None, "data_path is not taken into account when config_path is provided"
                dataset_list = [Dataset(path=config_path, data_path=join(path, f), filter=filter, name = f, raise_error=False) for f
                                in listdir(path) if isdir(join(path, f))]
                self.datasets = {d.name: d for d in dataset_list if d.name is not None}
            else:
                if data_path is not None:
                    assert not isabs(data_path), "data_path must be relative"
                dataset_list = [Dataset(path=join(path, f), data_path=data_path, filter=filter, raise_error=False) for f in listdir(path) if isdir(join(path, f))]
                self.datasets = {d.name:d for d in dataset_list if d.name is not None}
            if object_class_name_mapping is not None:
                for d in self.datasets.values():
                    for i, n in enumerate(d.object_class_names):
                        if n in object_class_name_mapping and object_class_name_mapping[n] not in d.object_class_names:
                            d.object_class_names[i] = object_class_name_mapping[n]

        assert len(self.datasets)>1, "no datasets where found"
        # object_class names are the common object_class names between all object_class names
        self.object_class_names = None
        for d in self.datasets.values():
            if self.object_class_names is None:
                self.object_class_names = d.object_class_names
            else:
                inter = set(self.object_class_names).intersection(d.object_class_names)
                self.object_class_names = [n for n in self.object_class_names if n in inter]
        assert self.object_class_names is not None, "no object classes found"
        assert len(self.object_class_names)>0, f"no object_class names in common between datasets : {[str(d) for d in self.datasets.values()]}"
        self.data = {}

    def __getitem__(self, item):
        return self.datasets[item]

    def set_object_class_name(self, old_object_class, new_object_class_name):
        old_object_class_idx = self._get_object_class_index(old_object_class)
        old_object_class_name = self._get_object_class_name(old_object_class)
        assert old_object_class_idx < len(self.object_class_names), f"invalid object_class idx, should be <{len(self.object_class_names)}"
        self.object_class_names[old_object_class_idx] = new_object_class_name
        for d in self.datasets.values():
            d.set_object_class_name(old_object_class_name, new_object_class_name)

    def _open_data(self, object_class, **kwargs):
        object_class = self._get_object_class_name(object_class)
        return pd.concat([d.get_data(object_class, add_dataset_name_column=True, copy=False, cache=False, **kwargs) for d in self.datasets.values()]) # do not cache both in each dataset and datasetList

    def store_selection(self, selection, object_class, name:str, dataset_column="DatasetName", **kwargs):
        object_class_name = self._get_object_class_name(object_class)
        for ds_name, sel in selection.groupby(dataset_column):
            if ds_name not in self.datasets:
                print(f"Error: dataset {ds_name} not in datasets: {str(self)}")
            else:
                self.datasets[ds_name].store_selection(sel, object_class_name, name, **kwargs)

    def __str__(self):
        return f"{[n for n in self.datasets.keys()]} oc={self.object_class_names}"

# util functions
def get_dataset_name(path, filter):
    # must contain a file ending by _config.json
    if isinstance(filter, str):
        _filter = lambda file_name:filter in file_name
    else:
        _filter = filter
    if isdir(path):
        for f in listdir(path):
            if f.endswith("_config.json") and (filter is None or _filter(f[:-12])) and isfile(join(path, f)):
                return f[:-12]
    return None

def _increment_name(name, seen):
    if name in seen:
        new_name = name+"_1"
        inc = 2
        while new_name in seen:
            new_name = f"{name}_{inc}"
            inc += 1
        name = new_name
    seen.add(name)
    return name

