# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['EventBout', 'TrackedRecording', 'Bodypart']

# %% ../nbs/00_core.ipynb 3
from abc import ABC, abstractmethod
from pathlib import Path, PosixPath
from typing import List, Tuple, Dict, Optional

import pandas as pd
import numpy as np

# %% ../nbs/00_core.ipynb 4
class EventBout:
    """
    Analysis of the `gait_analysis` package is mainly composed of identifying different
    types of events and their combined analysis thereafter. This class represents the
    core interface for these analyses, as it bundles all relevant information of such
    events.

    Attributes:
        self.event_type(str): type of event (e.g. "freezing" or "gait")
        self.event_id(int): index of the occurance of this event type in a recording in chronological order
        self.start_idx(int): frame index in which this event bout starts
        self.end_idx(int): frame index at which this event bout ends
        self.fps(int): frames-per-second of the recording
        self.duration(float): duration of the event bout in s
    """

    def __init__(
        self,
        event_type: str,  # type of event (e.g. "immobility" or "gait")
        event_id: int,  # evnt index (e.g. 0 for the first occurance of this event type in a recording)
        start_idx: int,  # index of the frame at which this event bout starts
        end_idx: int,  # index of the frame at which this event bout ends
        fps: int,  # frames-per-second value of the corresponding video recording
    ) -> None:
        self.event_type = event_type
        self.id = event_id
        self.start_idx = start_idx
        self.end_idx = end_idx
        self.fps = fps
        self.duration = ((end_idx + 1) - start_idx) / fps

# %% ../nbs/00_core.ipynb 5
"""
ToDo:
- reading of metadata and checking whether paradigm_id, week_id, .. are valid needs to be
  integrated with the Configs approach, that was already designed for the 3D pipeline
  
- the abstract methods like "preprocess" or "run_behavrioral_analyses" should probably also take
  some "configs" as input to make it consistent for all potential different subclasses? 
  Yet, at least they are never returning anything, so this would not change.
"""


class TrackedRecording(ABC):
    """
    While the analysis depends mainly on the identification of `EventBout`s and their integrated analyses,
    `TrackedRecording`s are the data objects on which these events are identified. They are implemented here
    as an abstract base class, as they will require slightly different functionalities depending on whether
    a 2D or 3D tracking needs to be handled.
    """

    @property
    def valid_paradigms(self) -> List[str]:
        return ["OTR", "OTT", "OTE"]

    @property
    def valid_mouse_lines(self) -> List[str]:
        return ["194", "195", "196", "206", "209"]

    @property
    def valid_recording_weeks(self) -> List[int]:
        return [1, 4, 8, 12, 14]

    @property
    def marker_ids_to_exclude_for_smoothing_and_interpolation(self) -> List[str]:
        return [
            "LED5",
            "MazeCornerClosedRight",
            "MazeCornerClosedLeft",
            "MazeCornerOpenRight",
            "MazeCornerOpenLeft",
        ]

    def __init__(
        self,
        filepath: Path,  # the filepath to the output of DLC (.h5 or .csv file)
        week_id: int,  # the experimental week post injection in which the recording was performed. Has to be an element of the self.valid_week_ids list
    ) -> None:

        assert type(filepath) == PosixPath, '"filepath" has to be a pathlib.Path object'
        assert (
            week_id in self.valid_recording_weeks
        ), f'"week_id" = {week_id} is not listed in "valid_recording_weeks": {self.valid_recording_weeks}'
        (
            self.df_successfully_loaded,
            self.loaded_tracking_df,
        ) = self._attempt_loading_df_from_dlc_output(filepath=filepath)
        if self.df_successfully_loaded:
            self.filepath = filepath
            self.week_id = week_id
            self._load_remaining_metadata()

    @abstractmethod
    def _load_remaining_metadata(self) -> None:
        pass

    @abstractmethod
    def preprocess_tracking(self) -> None:
        pass

    @abstractmethod
    def run_event_detection(self) -> None:
        pass

    @abstractmethod
    def export_results(self) -> None:
        pass

    @abstractmethod
    def inspect_processing(self) -> None:
        pass

    def _attempt_loading_df_from_dlc_output(
        self, filepath: Path  # filepath to the dlc tracking result (.h5 or .csv file)
    ) -> Tuple[bool, Optional[pd.DataFrame]]:
        assert filepath.name.endswith(".csv") or filepath.name.endswith(
            ".h5"
        ), "The filepath you specified is not referring to a .csv or a .h5 file!"
        try:
            if filepath.name.endswith(".csv"):
                df = pd.read_csv(filepath, low_memory=False)
                df = df.drop("scorer", axis=1)
                df.columns = df.iloc[0, :] + "_" + df.iloc[1, :]
                df = df.drop([0, 1], axis=0)
                df = df.reset_index()
                df = df.drop("index", axis=1)
                df = df.astype(float)
            else:
                df = pd.read_hdf(filepath)
                target_column_names = []
                for marker_id, data_id in zip(
                    df.columns.get_level_values(1), df.columns.get_level_values(2)
                ):
                    target_column_names.append(f"{marker_id}_{data_id}")
                df.columns = target_column_names
                df = df.astype(float)
            successfully_loaded = True
        except:
            successfully_loaded = False
            df = None
        return successfully_loaded, df

# %% ../nbs/00_core.ipynb 6
class Bodypart:
    """
    Bundles all information and functions needed to process the data of a single bodypart (i.e. marker id).
    Requires that the coordinates have already been normalized (translated & rotated) and converted to cm.
    """

    @property
    def exclusion_criteria(self) -> Dict:
        return {
            "likelihood_threshold": 0.5,
            "min_x": -5,
            "max_x": 55,
            "min_y": -3,
            "max_y": 7,
        }

    def __init__(
        self,
        bodypart_id: str,  # The ID of the bodypart, i.e. the marker_id
        df: pd.DataFrame,  # The normalized (rotated & translated) & converted (in cm) DataFrame that holds the DLC tracking information about this marker
        fps: int,  # frames-per-second at which this recording was acquired
    ) -> None:
        self.id = bodypart_id
        sliced_df = self._slice_df(df=df)
        self.df = self._apply_exclusion_criteria(
            df=sliced_df, exclusion_criteria=self.exclusion_criteria
        )
        self.fps = fps
        self.framerate = 1 / fps

    def _slice_df(self, df: pd.DataFrame) -> pd.DataFrame:
        df_input = df.copy()
        data = {
            "x": df_input.loc[:, f"{self.id}_x"],
            "y": df_input.loc[:, f"{self.id}_y"],
            "likelihood": df_input.loc[:, f"{self.id}_likelihood"],
        }
        return pd.DataFrame(data=data)

    def _apply_exclusion_criteria(
        self, df: pd.DataFrame, exclusion_criteria: Dict
    ) -> None:
        df.loc[
            df["likelihood"] < exclusion_criteria["likelihood_threshold"], :
        ] = np.nan
        for coord in ["x", "y"]:
            df.loc[df[coord] < exclusion_criteria[f"min_{coord}"], :] = np.nan
            df.loc[df[coord] > exclusion_criteria[f"max_{coord}"], :] = np.nan
        return df

    def calculate_speed_and_identify_immobility(
        self, sliding_window_size: int, immobility_threshold: float
    ) -> None:
        self._add_speed_to_df()
        self._add_rolling_speed_to_df(sliding_window_size=sliding_window_size)
        self._add_immobility_to_df(immobility_threshold=immobility_threshold)

    def _add_speed_to_df(self) -> None:
        self.df.loc[:, "speed_cm_per_s"] = (
            self.df.loc[:, "x"].diff() ** 2 + self.df.loc[:, "y"].diff() ** 2
        ) ** 0.5 / self.framerate

    def _add_rolling_speed_to_df(self, sliding_window_size: int) -> None:
        min_periods = int(sliding_window_size * 0.66)
        self.df.loc[:, "rolling_speed_cm_per_s"] = (
            self.df.loc[:, "speed_cm_per_s"]
            .rolling(sliding_window_size, min_periods=min_periods, center=True)
            .mean()
        )

    def _add_immobility_to_df(self, immobility_threshold: float) -> None:
        self.df.loc[:, "immobility"] = False
        self.df.loc[
            self.df["rolling_speed_cm_per_s"] < immobility_threshold, "immobility"
        ] = True
