# All of the omitted stimuli have a duration of 250ms as defined
# by the Visual Behavior team. For questions about duration contact that
# team.
import inspect
import logging
import os
from typing import List, Optional, Union

from allensdk.brain_observatory.behavior.behavior_session import (
    BehaviorSession,
)
from allensdk.brain_observatory.behavior.data_files import (
    BehaviorStimulusFile,
)
from allensdk.brain_observatory.behavior.data_objects.metadata.behavior_metadata.date_of_acquisition import (  # noqa: E501
    DateOfAcquisition,
)
from allensdk.brain_observatory.behavior.image_api import Image, ImageApi
from allensdk.brain_observatory.session_api_utils import sessions_are_equal
from allensdk.core import (
    DataObject,
    JsonReadableInterface,
    NwbReadableInterface,
    NwbWritableInterface,
)
from allensdk.core.auth_config import LIMS_DB_CREDENTIAL_MAP
from allensdk.internal.api import db_connection_creator
from pynwb import NWBHDF5IO, NWBFile, ProcessingModule
from pynwb.base import Images
from pynwb.image import GrayscaleImage


def get_column_name(table_cols: list, possible_names: set) -> str:
    """
    This function returns a column name, given a table with unknown
    column names and a set of possible column names which are expected.
    The table column name returned should be the only name contained in
    the "expected" possible names.
    :param table_cols: the table columns to search for the possible name within
    :param possible_names: the names that could exist within the data columns
    :return: the first entry of the intersection between the possible names
             and the names of the columns of the stimulus table
    """

    column_set = set(table_cols)
    column_names = list(column_set.intersection(possible_names))
    if not len(column_names) == 1:
        raise KeyError(
            "Table expected one name column in intersection, found:"
            f" {column_names}"
        )
    return column_names[0]


def get_image(nwbfile: NWBFile, name: str, module: str) -> Image:
    nwb_img = nwbfile.processing[module].get_data_interface("images")[name]
    data = nwb_img.data
    resolution = nwb_img.resolution  # px/cm
    spacing = [resolution * 10, resolution * 10]

    img = ImageApi.serialize(data, spacing, "mm")
    img = ImageApi.deserialize(img=img)
    return img


def add_image_to_nwb(nwbfile: NWBFile, image_data: Image, image_name: str):
    """
    Adds image given by image_data with name image_name to nwbfile

    Parameters
    ----------
    nwbfile
        nwbfile to add image to
    image_data
        The image data
    image_name
        Image name

    Returns
    -------
    None
    """
    module_name = "ophys"
    description = "{} image at pixels/cm resolution".format(image_name)

    data, spacing, unit = image_data

    assert spacing[0] == spacing[1] and len(spacing) == 2 and unit == "mm"

    if module_name not in nwbfile.processing:
        ophys_mod = ProcessingModule(module_name, "Ophys processing module")
        nwbfile.add_processing_module(ophys_mod)
    else:
        ophys_mod = nwbfile.processing[module_name]

    image = GrayscaleImage(
        image_name, data, resolution=spacing[0] / 10, description=description
    )

    if "images" not in ophys_mod.containers:
        images = Images(name="images")
        ophys_mod.add_data_interface(images)
    else:
        images = ophys_mod["images"]
    images.add_image(image)


class NWBWriter:
    """Base class for writing NWB files"""

    def __init__(
        self,
        nwb_filepath: str,
        session_data: dict,
        serializer: Union[
            JsonReadableInterface, NwbReadableInterface, NwbWritableInterface
        ],
    ):
        """

        Parameters
        ----------
        nwb_filepath: path to write nwb
        session_data: dict representation of data to instantiate `serializer`
            and write nwb
        serializer: The class to use to read `session_data` and write nwb.
            Must implement `JsonReadableInterface`, `NwbReadableInterface`,
            `NwbWritableInterface`
        """
        self._serializer = serializer
        self._session_data = session_data
        self._nwb_filepath = nwb_filepath
        self.nwb_filepath_inprogress = nwb_filepath + ".inprogress"
        self._nwb_filepath_error = nwb_filepath + ".error"

        logging.basicConfig(
            format="%(asctime)s - %(process)s - %(levelname)s - %(message)s",
            level=logging.INFO,
        )

        # Clean out files from previous runs:
        for filename in [
            self.nwb_filepath_inprogress,
            self._nwb_filepath_error,
            nwb_filepath,
        ]:
            if os.path.exists(filename):
                os.remove(filename)

    @property
    def nwb_filepath(self) -> str:
        """Path to write nwb file"""
        return self._nwb_filepath

    def write_nwb(
        self,
        id_column_name: str = "behavior_session_id",
        skip_metadata: Optional[List[str]] = None,
        skip_stim: Optional[List[str]] = None,
        **kwargs,
    ):
        """Tries to write nwb to disk. If it fails, the filepath has ".error"
        appended

        Parameters
        ----------
        id_column_name : str
            Name of the id column to pull from metadata.
        skip_metadata : list of str
            Name of key in session_data to skip when comparing to session.
        skip_stim : list of str
            Name of key in stimulus file to skip when comparing to session.
        kwargs: kwargs sent to `from_nwb`, `to_nwb`

        """
        from_lims_kwargs = {
            k: v
            for k, v in kwargs.items()
            if k in inspect.signature(self._serializer.from_lims).parameters
        }
        lims_session = self._serializer.from_lims(
            self._session_data[id_column_name], **from_lims_kwargs
        )
        lims_session = self._update_session(lims_session, **kwargs)

        try:
            nwbfile = self._write_nwb(session=lims_session, **kwargs)
            self._compare_metadata(
                input_id=self._session_data[id_column_name],
                input_session=lims_session,
                skip_metadata=skip_metadata,
            )
            self._compare_stimulus_file(
                input_id=self._session_data[id_column_name],
                input_session=lims_session,
                skip_stim=skip_stim,
            )
            self._compare_sessions(
                nwbfile=nwbfile, lims_session=lims_session, **kwargs
            )
            os.rename(self.nwb_filepath_inprogress, self._nwb_filepath)
        except Exception as e:
            if os.path.isfile(self.nwb_filepath_inprogress):
                os.rename(
                    self.nwb_filepath_inprogress, self._nwb_filepath_error
                )
            raise e

    def _update_session(
        self, lims_session: BehaviorSession, **kwargs
    ) -> BehaviorSession:
        """Call session methods to update certain values within the session.

        Should be used as part of a datarelease to resolve known data issues.
        """
        return lims_session

    def _write_nwb(self, session: BehaviorSession, **kwargs) -> NWBFile:
        """

        Parameters
        ----------
        session_data
        kwargs: kwargs to pass to `to_nwb`

        Returns
        -------

        """
        to_nwb_kwargs = {
            k: v
            for k, v in kwargs.items()
            if k in inspect.signature(self._serializer.to_nwb).parameters
        }
        nwbfile = session.to_nwb(**to_nwb_kwargs)

        with NWBHDF5IO(self.nwb_filepath_inprogress, "w") as nwb_file_writer:
            nwb_file_writer.write(nwbfile)
        return nwbfile

    def _compare_metadata(
        self,
        input_id: int,
        input_session: BehaviorSession,
        skip_metadata: List[str],
    ):
        """Compare data in the metadata table and behavior/experiment session
        to each other. Raise on conflicts.

        Parameters
        ----------
        input_id : int
            Id of the session/experiment to compare.
        input_session : BehaviorSession or BehaviorOphysExperiment
            Session object we are creating an NWB file for.
        skip_metadata : list of strings
            Names of metadata keys to skip during comparison.
        """
        if skip_metadata is None:
            skip_metadata = []
        error_message = ""
        # Test BehaviorSession object metadata against the metadata from the
        # behavior session table.
        bs_metadata = input_session.metadata
        for key, bs_val in self._session_data.items():
            if key in skip_metadata:
                logging.info(f"Skipping metadata table {key} comparison...")
                continue
            if bs_val != bs_metadata[key]:
                error_message += (
                    f"Value for {key} does not match for id={input_id} "
                    f"when comparing session object metadata and associated "
                    "metadata table.\n"
                    f"\tObject data={bs_metadata[key]};\n"
                    f"\tTable data={bs_val}.\n"
                )
        if len(error_message) > 0:
            raise ValueError(error_message)

    def _compare_stimulus_file(
        self,
        input_id: int,
        input_session: BehaviorSession,
        skip_stim: List[str],
    ):
        """Compare data in the stimulus file and loaded session/experiment to
        each other. Raise on conflicts.

        Parameters
        ----------
        input_id : int
            Id of the session/experiment to compare.
        input_session : BehaviorSession or BehaviorOphysExperiment
            Session object we are creating an NWB file for.
        skip_stim : list of strings
            Names of stimulus file keys to skip during comparison.
        """
        if skip_stim is None:
            skip_stim = []
        error_message = ""
        behavior_session_id = input_session.behavior_session_id
        db_conn = db_connection_creator(
            fallback_credentials=LIMS_DB_CREDENTIAL_MAP
        )
        stimulus_file = BehaviorStimulusFile.from_lims(
            db=db_conn, behavior_session_id=behavior_session_id
        ).validate()
        stim_file_methods = dir(stimulus_file)
        for key, bs_val in input_session.metadata.items():
            if key in skip_stim:
                logging.info(f"Skipping stimulus file {key} comparison...")
                continue
            if key in stim_file_methods:
                stim_value = getattr(stimulus_file, key)
                if key == "date_of_acquisition":
                    stim_value = DateOfAcquisition(stim_value).value
                if bs_val != stim_value:
                    error_message += (
                        f"Value for {key} does not match for id={input_id} "
                        "when comparing session object metadata and data from "
                        "the behavior stimulus pickle file.\n"
                        f"\tObject data={bs_val};\n"
                        f"\tStim data={stim_value}\n"
                    )
        if len(error_message) > 0:
            raise ValueError(error_message)

    def _compare_sessions(
        self, nwbfile: NWBFile, lims_session: DataObject, **kwargs
    ):
        kwargs = {
            k: v
            for k, v in kwargs.items()
            if k in inspect.signature(self._serializer.from_nwb).parameters
        }
        nwb_session = self._serializer.from_nwb(nwbfile, **kwargs)
        assert sessions_are_equal(lims_session, nwb_session, reraise=True)
