import warnings

import numpy as np
import os
import shutil
import sys

import hcai_datasets.hcai_nova_dynamic.utils.nova_types as nt
import hcai_datasets.hcai_nova_dynamic.utils.nova_data_utils as ndu
import hcai_datasets.hcai_nova_dynamic.utils.nova_anno_utils as nau
from hcai_dataset_utils.dataset_iterable import DatasetIterable

from hcai_datasets.hcai_nova_dynamic.utils.nova_data_utils import (
    AudioData,
    VideoData,
    StreamData,
)
from hcai_datasets.hcai_nova_dynamic.utils.nova_anno_utils import (
    DiscreteAnnotation,
    ContinuousAnnotation,
    FreeAnnotation,
)

from hcai_datasets.hcai_nova_dynamic.nova_db_handler import NovaDBHandler
from hcai_datasets.hcai_nova_dynamic.utils.nova_utils import *


class HcaiNovaDynamicIterable(DatasetIterable):
    def __init__(
            self,
            *,
            db_config_path=None,
            db_config_dict=None,
            dataset=None,
            nova_data_dir=None,
            sessions=None,
            annotator=None,
            schemes=None,
            roles=None,
            data_streams=None,
            start=None,
            end=None,
            left_context="0s",
            right_context="0s",
            frame_size=None,
            stride=None,
            flatten_samples=False,
            supervised_keys=None,
            add_rest_class=True,
            lazy_loading=False,
            **kwargs
    ):
        """
        Initialize the HcaiNovaDynamic dataset builder
        Args:
          nova_data_dir: the directory to look for data. same as the directory specified in the nova gui.
          frame_size: the framesize to look at. the matching annotation will be calculated as majority vote from all annotations that are overlapping with the timeframe.
          left_context: additional data to pass to the classifier on the left side of the frame.
          right_context: additional data to pass to the classifier on the left side of the frame.
          stride: how much a frame is moved to calculate the next sample. equals framesize by default.
          flatten_samples: if set to True samples with the same annotation scheme but from different roles will be treated as separate samples. only <scheme> is used for the keys.
          supervised_keys: if specified the dataset can be used with "as_supervised" set to True. Should be in the format <role>.<scheme>. if flatten_samples is true <role> will be ignored.
          add_rest_class: when set to True an additional restclass will be added to the end the label list
          db_config_path: path to a configfile whith the nova database config.
          db_config_dict: dictionary with the nova database config. can be used instead of db_config_path. if both are specified db_config_dict is used.
          dataset: the name of the dataset. must match the dataset name in the nova database.
          sessions: list of sessions that should be loaded. must match the session names in nova.
          annotator: the name of the annotator that labeld the session. must match annotator names in nova.
          schemes: list of the annotation schemes to fetch
          roles: list of roles for which the annotation should be loaded.
          data_streams: list datastreams for which the annotation should be loaded. must match stream names in nova.
          start: optional start time_ms. use if only a specific chunk of a session should be retreived.
          end: optional end time_ms. use if only a specifc chunk of a session should be retreived.
          **kwargs: arguments that will be passed through to the dataset builder
        """
        self.dataset = dataset
        self.nova_data_dir = nova_data_dir
        self.sessions = sessions
        self.roles = roles
        self.schemes = schemes
        self.data_streams = data_streams
        self.annotator = annotator
        self.left_context_ms = ndu.parse_time_string_to_ms(left_context)
        self.right_context_ms = ndu.parse_time_string_to_ms(right_context)
        self.frame_size_ms = ndu.parse_time_string_to_ms(frame_size) if frame_size else None
        self.stride_ms = (
            ndu.parse_time_string_to_ms(stride) if stride else self.frame_size_ms
        )
        self.start_ms = ndu.parse_time_string_to_ms(start)
        if not self.start_ms:
            self.start_ms = 0

        self.end_ms = ndu.parse_time_string_to_ms(end)
        if not self.end_ms:
            self.end_ms = float("inf")

        self.flatten_samples = flatten_samples
        self.add_rest_class = add_rest_class
        self.lazy_loading = lazy_loading
        self.nova_db_handler = NovaDBHandler(db_config_path, db_config_dict)

        # Retrieving meta information from the database
        mongo_schemes = self.nova_db_handler.get_schemes(
            dataset=dataset, schemes=schemes
        )
        mongo_data = self.nova_db_handler.get_data_streams(
            dataset=dataset, data_streams=data_streams
        )
        self.annos, self.anno_schemes = self._populate_label_info_from_mongo_doc(mongo_schemes)
        self.data_info, self.data_schemes = self._populate_data_info_from_mongo_doc(mongo_data)

        # setting supervised keys
        if supervised_keys and self.flatten_samples:
            if supervised_keys[0] not in self.data_streams:
                # remove <role> of supervised keys
                _, data_stream = split_role_key(supervised_keys[0])
                if not data_stream in self.data_streams:
                    print(
                        "Warning: Cannot find supervised key '{}' in datastreams".format(
                            data_stream
                        )
                    )
                    raise Warning("Unknown data_stream")
                else:
                    supervised_keys[0] = data_stream
            if supervised_keys[1] not in self.schemes:
                # remove <role> of supervised keys
                _, scheme = split_role_key(supervised_keys[1])
                if not scheme in schemes:
                    print(
                        "Warning: Cannot find supervised key '{}' in schemes".format(
                            scheme
                        )
                    )
                    raise Warning("Unknown scheme")
                else:
                    supervised_keys[1] = scheme

        self.supervised_keys = tuple(supervised_keys) if supervised_keys else None

        self._iterable = self._yield_sample()

    def _populate_label_info_from_mongo_doc(self, mongo_schemes):
        """
        Setting self.annos
        Args:
          mongo_schemes:

        Returns:

        """
        annos = {}
        anno_schemes = {}

        # List of all combinations from roles and schemes that occur in the retrieved data.
        for scheme in mongo_schemes:
            for role in self.roles:
                label_id = merge_role_key(role=role, key=scheme["name"])
                scheme_type = nt.string_to_enum(nt.AnnoTypes, scheme["type"])
                scheme_name = scheme["name"]
                scheme_valid = scheme["isValid"]
                anno_schemes[label_id] = scheme_type

                if scheme_type == nt.AnnoTypes.DISCRETE:
                    labels = scheme["labels"]
                    annos[label_id] = nau.DiscreteAnnotation(
                        role=role,
                        add_rest_class=True,
                        scheme=scheme_name,
                        is_valid=scheme_valid,
                        labels=labels,
                    )

                elif scheme_type == nt.AnnoTypes.CONTINUOUS:
                    min_val = scheme["min"]
                    max_val = scheme["max"]
                    sr = scheme["sr"]
                    annos[label_id] = nau.ContinuousAnnotation(
                        role=role,
                        scheme=scheme_name,
                        is_valid=scheme_valid,
                        min_val= min_val,
                        max_val= max_val,
                        sr=sr,
                    )

                elif scheme_type == nt.AnnoTypes.DISCRETE_POLYGON:
                    labels = scheme["labels"]
                    sr = scheme["sr"]
                    annos[label_id] = nau.DiscretePolygonAnnotation(
                        role=role,
                        scheme=scheme_name,
                        is_valid=scheme_valid,
                        labels=labels,
                        sr=sr,
                    )

                elif scheme_type == nt.AnnoTypes.FREE:
                    annos[label_id] = nau.FreeAnnotation(
                        role=role, scheme=scheme_name, is_valid=scheme_valid
                    )

                else:
                    raise ValueError("Invalid label type {}".format(scheme["type"]))

        return annos, anno_schemes

    def _populate_data_info_from_mongo_doc(self, mongo_data_streams):
        """
        Setting self.data
        Args:
          mongo_schemes:

        Returns:

        """
        data_info = {}
        data_schemes = {}

        for data_stream in mongo_data_streams:
            for role in self.roles:
                sample_stream_name = (
                        role + "." + data_stream["name"] + "." + data_stream["fileExt"]
                )
                sample_data_path = os.path.join(
                    self.nova_data_dir,
                    self.dataset,
                    self.sessions[0],
                    sample_stream_name,
                )
                dtype = nt.string_to_enum(nt.DataTypes, data_stream["type"])


                if dtype == nt.DataTypes.VIDEO:
                    data = VideoData(
                        role=role,
                        name=data_stream["name"],
                        file_ext=data_stream["fileExt"],
                        sr=data_stream["sr"],
                        is_valid=data_stream["isValid"],
                        sample_data_path=sample_data_path,
                        lazy_loading=self.lazy_loading,
                    )
                elif dtype == nt.DataTypes.AUDIO:
                    data = AudioData(
                        role=role,
                        name=data_stream["name"],
                        file_ext=data_stream["fileExt"],
                        sr=data_stream["sr"],
                        is_valid=data_stream["isValid"],
                        sample_data_path=sample_data_path,
                        lazy_loading=self.lazy_loading,
                    )
                elif dtype == nt.DataTypes.FEATURE:
                    data = StreamData(
                        role=role,
                        name=data_stream["name"],
                        sr=data_stream["sr"],
                        data_type=dtype,
                        is_valid=data_stream["isValid"],
                        sample_data_path=sample_data_path,
                        lazy_loading=self.lazy_loading,
                    )
                else:
                    raise ValueError("Invalid data type {}".format(data_stream["type"]))

                data_id = merge_role_key(role=role, key=data_stream["name"])
                data_info[data_id] = data
                data_schemes[data_id] = dtype

        return data_info, data_schemes

    def _load_annotation_for_session(self, session, time_to_ms=False):
        for label_id, anno in self.annos.items():
            mongo_anno = self.nova_db_handler.get_annos(
                self.dataset, anno.scheme, session, self.annotator, anno.role
            )
            anno.set_annotation_from_mongo_doc(mongo_anno, time_to_ms=time_to_ms)

    def _open_data_reader_for_session(self, session):
        for data_id, data in self.data_info.items():
            data_path = os.path.join(
                self.nova_data_dir, self.dataset, session, data_id + "." + data.file_ext
            )
            data.open_file_reader(data_path)

    def _yield_sample(self):
        """Yields examples."""

        # Needed to sort the samples later and assure that the order is the same as in nova. Necessary for CML till the tfds can be returned in order.
        sample_counter = 1

        for session in self.sessions:

            # Gather all data we need for this session
            self._load_annotation_for_session(session, time_to_ms=True)
            self._open_data_reader_for_session(session)

            session_info = self.nova_db_handler.get_session_info(self.dataset, session)[
                0
            ]
            dur = session_info["duration"]
            # If are loading any datastreams we check if any datastream is shorter than the duration stored in the database suggests
            if self.data_info:
                dur = min(*[v.dur for k, v in self.data_info.items()], dur)

            if not dur:
                raise ValueError("Session {} has no duration.".format(session))

            dur_ms = int(dur * 1000)

            # If framesize is not specified we return the whole session as one junk
            if self.frame_size_ms is None:
                self.frame_size_ms = min(dur_ms, self.end_ms - self.start_ms)
                self.stride_ms = self.frame_size_ms

            # Starting position of the first frame in seconds
            #c_pos_ms = self.left_context_ms + self.start_ms
            c_pos_ms = max(self.left_context_ms, self.start_ms)

            # Generate samples for this session
            while c_pos_ms + self.stride_ms + self.right_context_ms <= min(
                    self.end_ms, dur_ms
            ):
                frame_start_ms = c_pos_ms - self.left_context_ms
                frame_end_ms = c_pos_ms + self.frame_size_ms + self.right_context_ms


                key = (
                        session
                        + "_"
                        + str(frame_start_ms / 1000)
                        + "_"
                        + str(frame_end_ms / 1000)
                )

                labels_for_frame = [
                    {k: v.get_label_for_frame(frame_start_ms, frame_end_ms)}
                    for k, v in self.annos.items()
                ]
                data_for_frame = []

                for k, v in self.data_info.items():
                    sample = v.get_sample(frame_start_ms, frame_end_ms)
                    if sample.shape[0] == 0:
                        print(f'Sample{frame_start_ms}-{frame_end_ms} is empty')
                        c_pos_ms += self.stride_ms
                        continue
                    data_for_frame.append({k: sample})

                sample_dict = {}

                for l in labels_for_frame:
                    sample_dict.update(l)

                for d in data_for_frame:
                    sample_dict.update(d)

                if self.flatten_samples:

                    # grouping labels and data according to roles
                    for role in self.roles:
                        # filter dictionary to contain values for role
                        sample_dict_for_role = dict(
                            filter(lambda elem: role in elem[0], sample_dict.items())
                        )

                        # remove role from dictionary keys
                        sample_dict_for_role = dict(
                            map(
                                lambda elem: (elem[0].replace(role + ".", ""), elem[1]),
                                sample_dict_for_role.items(),
                            )
                        )

                        sample_dict_for_role["frame"] = (
                                str(sample_counter) + "_" + key + "_" + role
                        )
                        # yield key + '_' + role, sample_dict_for_role
                        yield sample_dict_for_role
                        sample_counter += 1
                    c_pos_ms += self.stride_ms

                else:
                    sample_dict["frame"] = str(sample_counter) + "_" + key
                    yield sample_dict
                    c_pos_ms += self.stride_ms
                    sample_counter += 1

            # closing file readers for this session
            for k, v in self.data_info.items():
                v.close_file_reader()

    def __iter__(self):
        return self._iterable

    def __next__(self):
        return self._iterable.__next__()

    def get_output_info(self):

        def map_label_id(lid):
            if self.flatten_samples and not lid == 'frame':
                return split_role_key(lid)[-1]
            return lid

        return {
                    # Adding fake framenumber label for sorting
                    'frame': {"dtype":np.str, "shape":(1,)},
                    **{map_label_id(k): v.get_info()[1] for k, v in self.annos.items()},
                    **{map_label_id(k): v.get_info()[1] for k, v in self.data_info.items()}
                }

