import re
from copy import deepcopy
from datetime import datetime
from pathlib import Path
from typing import Optional, Tuple

import numpy as np
from hdmf.backends.hdf5 import H5DataIO
from pynwb import NWBFile

from ....basetemporalalignmentinterface import BaseTemporalAlignmentInterface
from ....tools import get_module
from ....utils import (
    DeepDict,
    FilePathType,
    calculate_regular_series_rate,
    get_base_schema,
)


class LightningPoseDataInterface(BaseTemporalAlignmentInterface):
    """Data interface for Lightning Pose datasets."""

    display_name = "Lightning Pose"
    keywords = ("pose estimation", "video")
    associated_suffixes = (".csv", ".mp4")
    info = "Interface for handling a single stream of lightning pose data."

    def get_metadata_schema(self) -> dict:
        metadata_schema = super().get_metadata_schema()
        metadata_schema["properties"]["Behavior"] = get_base_schema(tag="Behavior")

        metadata_schema["properties"]["Behavior"].update(
            required=["PoseEstimation"],
            properties=dict(
                PoseEstimation=dict(
                    type="object",
                    required=["name"],
                    properties=dict(
                        name=dict(type="string", default="PoseEstimation"),
                        description=dict(type="string"),
                        scorer=dict(type="string"),
                        source_software=dict(type="string", default="LightningPose"),
                    ),
                    patternProperties={
                        "^(?!(name|description|scorer|source_software)$)[a-zA-Z0-9_]+$": dict(
                            title="PoseEstimationSeries",
                            type="object",
                            properties=dict(name=dict(type="string"), description=dict(type="string")),
                            minProperties=1,
                            additionalProperties=False,
                        )
                    },
                    minProperties=2,
                    additionalProperties=False,
                )
            ),
        )

        return metadata_schema

    def __init__(
        self,
        file_path: FilePathType,
        original_video_file_path: FilePathType,
        labeled_video_file_path: Optional[FilePathType] = None,
        verbose: bool = True,
    ):
        """
        Interface for writing pose estimation data from the Lightning Pose algorithm.

        Parameters
        ----------
        file_path : a string or a path
            Path to the .csv file that contains the predictions from Lightning Pose.
        original_video_file_path : a string or a path
            Path to the original video file (.mp4).
        labeled_video_file_path : a string or a path, optional
            Path to the labeled video file (.mp4).
        verbose : bool, default: True
            controls verbosity. ``True`` by default.
        """
        from neuroconv.datainterfaces.behavior.video.video_utils import (
            VideoCaptureContext,
        )

        self._vc = VideoCaptureContext

        self.file_path = Path(file_path)
        assert self.file_path.exists(), f"The file '{self.file_path}' does not exist."
        self.original_video_file_path = Path(original_video_file_path)
        assert (
            self.original_video_file_path.exists()
        ), f"The original video file '{self.original_video_file_path}' does not exist."

        super().__init__(
            verbose,
            file_path=file_path,
            original_video_file_path=original_video_file_path,
            labeled_video_file_path=labeled_video_file_path,
        )

        # dimension is width by height
        self.dimension = self._get_original_video_shape()

        pose_estimation_data = self._load_source_data()
        _, self.scorer_name = pose_estimation_data.columns.get_level_values(0).drop_duplicates()
        self.pose_estimation_data = pose_estimation_data[self.scorer_name]
        self.keypoint_names = self.pose_estimation_data.columns.get_level_values(0).drop_duplicates().tolist()

        self._times = None

    def _load_source_data(self):
        import pandas as pd

        # The order of the header is "scorer", "bodyparts", "coords"
        pose_estimation_data = pd.read_csv(self.file_path, header=[0, 1, 2])
        return pose_estimation_data

    def _get_original_video_shape(self) -> Tuple[int, int]:
        with self._vc(file_path=str(self.original_video_file_path)) as video:
            video_shape = video.get_frame_shape()
        # image size of the original video is in height x width
        return video_shape[0], video_shape[1]

    def get_original_timestamps(self, stub_test: bool = False) -> np.ndarray:
        max_frames = 10 if stub_test else None
        with self._vc(file_path=str(self.original_video_file_path)) as video:
            timestamps = video.get_video_timestamps(max_frames=max_frames)
        return timestamps

    def get_timestamps(self, stub_test: bool = False) -> np.ndarray:
        max_frames = 10 if stub_test else None
        if self._times is None:
            return self.get_original_timestamps(stub_test=stub_test)

        timestamps = self._times if not stub_test else self._times[:max_frames]
        return timestamps

    def set_aligned_timestamps(self, aligned_timestamps: np.ndarray):
        self._times = aligned_timestamps

    def get_metadata(self) -> DeepDict:
        metadata = super().get_metadata()

        # Update the session start time if folder structure is saved in the format: YYYY-MM-DD/HH-MM-SS
        pattern = r"(?P<date_time>\d{4}-\d{2}-\d{2}/\d{2}-\d{2}-\d{2})"
        # Convert the file path parts to a string with forward slashes
        file_path = "/".join(self.file_path.parts)
        match = re.search(pattern, file_path)
        if match and "session_start_time" not in metadata["NWBFile"]:
            datetime_str = match.group("date_time")
            session_start_time = datetime.strptime(datetime_str, "%Y-%m-%d/%H-%M-%S")
            metadata["NWBFile"].update(session_start_time=session_start_time)

        metadata["Behavior"]["PoseEstimation"].update(
            name="PoseEstimation",
            description="Contains the pose estimation series for each keypoint.",
            scorer=self.scorer_name,
            source_software="LightningPose",
        )
        for keypoint_name in self.keypoint_names:
            keypoint_name_without_spaces = keypoint_name.replace(" ", "")
            pose_estimation_series_metadata = {
                keypoint_name: dict(
                    name=f"PoseEstimationSeries{keypoint_name_without_spaces}",
                    description=f"The estimated position (x, y) of {keypoint_name} over time.",
                )
            }
            metadata["Behavior"]["PoseEstimation"].update(pose_estimation_series_metadata)

        return metadata

    def add_to_nwbfile(
        self,
        nwbfile: NWBFile,
        metadata: Optional[dict] = None,
        reference_frame: Optional[str] = None,
        confidence_definition: Optional[str] = None,
        stub_test: Optional[bool] = False,
    ) -> None:
        """
        Add the pose estimation data to the nwbfile.

        Parameters
        ----------
        nwbfile : NWBFile
            The nwbfile to which the pose estimation data is added.
        metadata : dict, optional
            The metadata for the pose estimation data.
        reference_frame : str, optional
            The description defining what the (0, 0) coordinate corresponds to.
        confidence_definition : str, optional
            The description of how the confidence was computed, e.g., 'Softmax output of the deep neural network'.
        stub_test : bool, default: False
        """
        from ndx_pose import PoseEstimation, PoseEstimationSeries

        metadata_copy = deepcopy(metadata)

        # The parameters for the pose estimation container
        pose_estimation_metadata = metadata_copy["Behavior"]["PoseEstimation"]

        behavior = get_module(nwbfile, "behavior")

        pose_estimation_name = pose_estimation_metadata["name"]
        if pose_estimation_name in behavior.data_interfaces:
            raise ValueError(f"The nwbfile already contains a data interface with the name '{pose_estimation_name}'.")

        if "Videos" not in metadata_copy["Behavior"]:
            original_video_name = str(self.original_video_file_path)
        else:
            original_video_name = metadata_copy["Behavior"]["Videos"][0]["name"]

        pose_estimation_kwargs = dict(
            name=pose_estimation_metadata["name"],
            description=pose_estimation_metadata["description"],
            source_software=pose_estimation_metadata["source_software"],
            scorer=pose_estimation_metadata["scorer"],
            original_videos=[original_video_name],
            dimensions=[self.dimension],
        )

        pose_estimation_data = self.pose_estimation_data if not stub_test else self.pose_estimation_data.head(n=10)
        timestamps = self.get_timestamps(stub_test=stub_test)
        rate = calculate_regular_series_rate(series=timestamps)
        if rate:
            pose_estimation_series_kwargs = dict(rate=rate, starting_time=timestamps[0])
        else:
            assert len(timestamps) == len(
                pose_estimation_data
            ), f"The length of timestamps ({len(timestamps)}) and pose estimation data ({len(pose_estimation_data)}) must be equal."
            pose_estimation_series_kwargs = dict(timestamps=H5DataIO(data=timestamps, compression="gzip"))

        pose_estimation_series = []
        for keypoint_name in self.keypoint_names:
            pose_estimation_series_data = pose_estimation_data[keypoint_name]
            keypoint_name_without_spaces = keypoint_name.replace(" ", "")

            pose_estimation_series_kwargs.update(
                name=pose_estimation_metadata[keypoint_name_without_spaces]["name"],
                description=pose_estimation_metadata[keypoint_name_without_spaces]["description"],
                data=pose_estimation_series_data[["x", "y"]].values,
                confidence=pose_estimation_series_data["likelihood"].values,
                reference_frame=reference_frame or "(0,0) is unknown.",
                unit="px",
            )

            if confidence_definition:
                pose_estimation_series_kwargs.update(confidence_definition=confidence_definition)

            pose_estimation_series.append(PoseEstimationSeries(**pose_estimation_series_kwargs))

        pose_estimation_kwargs.update(
            pose_estimation_series=pose_estimation_series,
        )

        if self.source_data["labeled_video_file_path"]:
            if "Videos" not in metadata_copy["Behavior"]:
                labeled_video_name = str(self.source_data["labeled_video_file_path"])
            else:
                labeled_video_name = metadata_copy["Behavior"]["Videos"][1]["name"]

            pose_estimation_kwargs.update(labeled_videos=[labeled_video_name])

        # Create the container for pose estimation
        pose_estimation = PoseEstimation(**pose_estimation_kwargs)

        behavior.add(pose_estimation)
