import os
import numpy as np
from pathlib import Path
from typing import Optional, List, Dict

import imageio
from pynwb import NWBFile
from pynwb.image import IndexSeries

from allensdk.brain_observatory.behavior.data_files import BehaviorStimulusFile
from allensdk.core import DataObject
from allensdk.core import \
    NwbReadableInterface
from allensdk.brain_observatory.behavior.data_files.stimulus_file import \
    StimulusFileReadableInterface
from allensdk.core import \
    NwbWritableInterface
from allensdk.brain_observatory.behavior.data_objects.stimuli.presentations \
    import \
    Presentations
from allensdk.brain_observatory.behavior.stimulus_processing import \
    get_stimulus_templates
from allensdk.brain_observatory.behavior.data_objects.stimuli \
    .stimulus_templates import \
    StimulusTemplate, StimulusTemplateFactory, StimulusMovieTemplateFactory
from allensdk.brain_observatory.behavior.write_nwb.extensions\
    .stimulus_template.ndx_stimulus_template import \
    StimulusTemplateExtension
from allensdk.internal.core.lims_utilities import safe_system_path


class Templates(DataObject, StimulusFileReadableInterface,
                NwbReadableInterface, NwbWritableInterface):
    def __init__(self, templates: Dict[str, StimulusTemplate]):
        super().__init__(name='stimulus_templates', value=templates)
        # Grab the keys from the input dictionary. The "images" key is assumed
        # to be the key in the dictionary that does not have "movie" in its key
        # name. For VBO and VBN releases, there should only be at most 2
        # keys in the dictionary.
        image_template_keys = [
            key for key in templates.keys()
            if 'movie' not in key.lower()]
        self._image_template_key = None

        error_message = ""
        if len(image_template_keys) == 1:
            self._image_template_key = image_template_keys[0]
        elif len(image_template_keys) > 1:
            error_message += (
                "Found multiple image StimulusTemplates "
                f"{image_template_keys}. ")

        movie_template_keys = [
            key for key in templates.keys()
            if 'movie' in key.lower()]
        self._fingerprint_movie_template_key = None
        if len(movie_template_keys) == 1:
            self._fingerprint_movie_template_key = movie_template_keys[0]
        elif len(movie_template_keys) > 1:
            error_message += (
                "Found multiple fingerprint movie StimulusTemplates "
                f"{movie_template_keys}. ")
        if len(error_message) > 0:
            error_message += (
                "This is not currently supported. "
                "Please limit input to one image template and/or one "
                "fingerprint movie template.")
            raise NotImplementedError(error_message)

    @classmethod
    def from_stimulus_file(
            cls, stimulus_file: BehaviorStimulusFile,
            limit_to_images: Optional[List] = None,
            load_stimulus_movie: bool = False) -> "Templates":
        """Get stimulus templates (movies, scenes) for behavior session."""

        # TODO: Eventually the `grating_images_dict` should be provided by the
        #       BehaviorLimsExtractor/BehaviorJsonExtractor classes.
        #       - NJM 2021/2/23

        gratings_dir = "/allen/programs/braintv/production/visualbehavior"
        gratings_dir = os.path.join(gratings_dir,
                                    "prod5/project_VisualBehavior")
        grating_images_dict = {
            "gratings_0.0": {
                "warped": np.asarray(imageio.imread(
                    safe_system_path(os.path.join(gratings_dir,
                                                  "warped_grating_0.png")))),
                "unwarped": np.asarray(imageio.imread(
                    safe_system_path(os.path.join(
                        gratings_dir, "masked_unwarped_grating_0.png"))))
            },
            "gratings_90.0": {
                "warped": np.asarray(imageio.imread(
                    safe_system_path(os.path.join(gratings_dir,
                                                  "warped_grating_90.png")))),
                "unwarped": np.asarray(imageio.imread(
                    safe_system_path(os.path.join(
                        gratings_dir, "masked_unwarped_grating_90.png"))))
            },
            "gratings_180.0": {
                "warped": np.asarray(imageio.imread(
                    safe_system_path(os.path.join(gratings_dir,
                                                  "warped_grating_180.png")))),
                "unwarped": np.asarray(imageio.imread(
                    safe_system_path(os.path.join(
                        gratings_dir, "masked_unwarped_grating_180.png"))))
            },
            "gratings_270.0": {
                "warped": np.asarray(imageio.imread(
                    safe_system_path(os.path.join(gratings_dir,
                                                  "warped_grating_270.png")))),
                "unwarped": np.asarray(imageio.imread(
                    safe_system_path(os.path.join(
                        gratings_dir, "masked_unwarped_grating_270.png"))))
            }
        }

        pkl = stimulus_file.data
        stim_template = get_stimulus_templates(
            pkl=pkl,
            grating_images_dict=grating_images_dict,
            limit_to_images=limit_to_images)
        t = {stim_template.image_set_name: stim_template}

        has_fingerprint_stimulus = (
                "fingerprint" in pkl["items"]["behavior"]["items"]
        )
        if has_fingerprint_stimulus and load_stimulus_movie:
            movie_data = np.load(
                Path(pkl['items']['behavior']['items'][
                    'fingerprint']['static_stimulus']['movie_path'])
            )
            movie_template = StimulusMovieTemplateFactory.from_unprocessed(
                movie_name="natural_movie_one",
                movie_frames=movie_data,
            )
            t[movie_template.image_set_name] = movie_template

        return Templates(templates=t)

    @classmethod
    def from_nwb(cls, nwbfile: NWBFile) -> "Templates":
        templates = {}
        for image_set_name, image_data in nwbfile.stimulus_template.items():

            image_attributes = [
                {'image_name': image_name}
                for image_name in image_data.control_description
            ]
            templates[image_set_name] = StimulusTemplateFactory.from_processed(
                image_set_name=image_set_name,
                image_attributes=image_attributes,
                warped=image_data.data[:],
                unwarped=image_data.unwarped[:]
            )
        return Templates(templates=templates)

    def to_nwb(self, nwbfile: NWBFile,
               stimulus_presentations: Presentations) -> NWBFile:
        for key, stimulus_templates in self.value.items():

            unwarped_images = []
            warped_images = []
            image_names = []
            for image_name, image_data in stimulus_templates.items():
                image_names.append(image_name)
                unwarped_images.append(image_data.unwarped)
                warped_images.append(image_data.warped)

            image_index = np.zeros(len(image_names))
            image_index[:] = np.nan

            visual_stimulus_image_series = \
                StimulusTemplateExtension(
                    name=stimulus_templates.image_set_name,
                    data=warped_images,
                    unwarped=unwarped_images,
                    control=list(range(len(image_names))),
                    control_description=image_names,
                    unit='NA',
                    format='raw',
                    timestamps=image_index)

            nwbfile.add_stimulus_template(visual_stimulus_image_series)

        if 'image_index' in stimulus_presentations.value \
                and self._image_template_key is not None:
            nwbfile = self._add_image_index_to_nwb(
                nwbfile=nwbfile, presentations=stimulus_presentations)

        return nwbfile

    def _add_image_index_to_nwb(
            self, nwbfile: NWBFile, presentations: Presentations):
        """Adds the image index and start_time for all stimulus templates
        to NWB"""
        stimulus_templates = self.value[self._image_template_key]
        presentations = presentations.value

        nwb_template = nwbfile.stimulus_template[
            stimulus_templates.image_set_name]
        stimulus_name = 'image_set' \
            if 'image_set' in presentations else 'stimulus_name'
        stimulus_index = presentations[
            presentations[stimulus_name] == nwb_template.name]

        image_index = IndexSeries(
            name=nwb_template.name,
            data=stimulus_index['image_index'].values,
            unit='None',
            indexed_timeseries=nwb_template,
            timestamps=stimulus_index['start_time'].values)
        nwbfile.add_stimulus(image_index)
        return nwbfile

    @property
    def image_template_key(self) -> str:
        """
        Name of the image template in template dictionary.
        """
        return self._image_template_key

    @property
    def fingerprint_movie_template_key(self) -> str:
        """
        Name of the fingerprint movie template in template dictionary.
        """
        return self._fingerprint_movie_template_key
