# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/01_twoD_01_cam_interfaces.ipynb.

# %% auto 0
__all__ = ['TopTracked2DRecording']

# %% ../../nbs/01_twoD_01_cam_interfaces.ipynb 3
from typing import Tuple, List, Dict, Optional, Union
from pathlib import Path
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

from ..core import TrackedRecording
from . import utils, preprocess, detect, export

# %% ../../nbs/01_twoD_01_cam_interfaces.ipynb 4
class TopTracked2DRecording(TrackedRecording):
    """
    Very customized subclass of `TrackedRecording` that was designed to run the gait analysis
    on 2D tracking data obtained from a single camera with a top-down-view on the subject.
    """

    def _load_remaining_metadata(self) -> None:
        self.fps = self._get_correct_fps()
        self.framerate = 1 / self.fps
        self.metadata = self._retrieve_metadata(filename=self.filepath.name)

    def _get_correct_fps(self) -> int:
        if self.loaded_tracking_df.shape[0] > 25_000:
            fps = 80
        else:
            fps = 30
        return fps

    def _retrieve_metadata(self, filename: str) -> Dict:
        """
        Very much dependent on the following file naming convention:
        LLL_FS-SS_YYMMDD_PPP_whatever.h5 (or .csv)
        Where:
            LLL: three digits mouse line code
            FS-SS: the subject ID (must start with capital F) with generation (e.g. F2 for second generation) and the mouse ID as two digits
            YYMMDD: the date of the recording as two digits year (YY), month (MM), and day (DD)
            PPP: the three letter string of the experimental paradigm
        For instance:
            196_F7-27_220826_OTT_whatever.h5
        """
        splits = filename.split("_")
        line_id, mouse_id, date, paradigm_id, cam_id = (
            splits[0],
            splits[1],
            splits[2],
            splits[3][0:3],
            "Top",
        )
        self._check_metadata(metadata=(line_id, mouse_id, date, paradigm_id, cam_id))
        metadata = {
            "recording_date": self.recording_date,
            "animal": f"{self.mouse_line}_{self.mouse_id}",
            "paradigm": self.paradigm,
            "cam": self.cam_id,
        }
        return metadata

    # ToDo - replace with something more generalizable that could be put to utils
    def _check_metadata(self, metadata=Tuple[str]) -> None:
        animal_line, animal_id, recording_date, paradigm, cam_id = (
            metadata[0],
            metadata[1],
            metadata[2],
            metadata[3],
            metadata[4],
        )
        self.cam_id = cam_id
        if animal_line not in self.valid_mouse_lines:
            while True:
                entered_input = input(f"Mouse line for {self.filepath}")
                if entered_input in self.valid_mouse_lines:
                    self.mouse_line = entered_input
                    break
                else:
                    print(
                        f"Entered mouse line does not match any of the defined mouse lines. \nPlease enter one of the following lines: {self.valid_mouse_lines}"
                    )
        else:
            self.mouse_line = animal_line
        if not animal_id.startswith("F"):
            while True:
                entered_input = input(f"Mouse ID for {self.filepath}")
                if entered_input.startswith("F"):
                    self.mouse_id = entered_input
                    break
                else:
                    print(f"Animal ID has to start with F. Example: F2-14")
        else:
            self.mouse_id = animal_id
        if paradigm not in self.valid_paradigms:
            while True:
                entered_input = input(f"Paradigm for {self.filepath}")
                if entered_input in self.valid_paradigms:
                    self.paradigm = entered_input
                    break
                else:
                    print(
                        f"Entered paradigm does not match any of the defined paradigms. \nPlease enter one of the following paradigms: {self.valid_paradigms}"
                    )
        else:
            self.paradigm = paradigm
        try:
            int(recording_date)
            self.recording_date = recording_date
        except:
            while True:
                entered_input = input(f"Recording date for {self.filepath}")
                try:
                    int(recording_date)
                    self.recording_date = recording_date
                    break
                except:
                    print(
                        f"Entered recording date has to be an integer in shape YYMMDD. Example: 220812"
                    )

    def preprocess_tracking(
        self,
        marker_ids_to_compute_coverage: List[str] = [
            "TailBase",
            "Snout",
        ],  # List of marker_ids on base of which the tracking coverage will be computed
        coverage_threshold: float = 0.75,  # If coverage of the above defined markers is lower than this threshold, the recording will be excluded from the analysis and therefore not further processed
        max_seconds_to_interpolate: float = 0.5,  # Maximum time interval in which consecutive nan´s will be interpolated
        likelihood_threshold: float = 0.5,  # Minimum prediction likelihood of DLC that is required to acceppt predicted marker position as valid
        marker_ids_to_compute_center_of_gravity: List[str] = [
            "TailBase",
            "Snout",
        ],  # marker ids that will be used to compute the center of gravity
        relative_maze_normalization_error_tolerance: float = 0.25,  # relative error that is tolerated when estimating the maze position for normalizing it´s position
    ) -> None:
        """
        Initiate the preprocessing of the DLC output tracking data. This includes smoothing,
        interpolation of small intervals with low likelihoods, and creation of additional markers,
        like a center of gravity. Furthermore, the code will identify the most reliable
        predicted position of the maze corners and evaluate, whether the reconstructed maze using
        these most reliable open corner positions yields a maze that is sufficiently similar to how the
        maze should look like (less than a specified relative error threshold). If it meets this
        criterion, it will use the open corner from there on to normalize the coordinate system.
        If not, it will compute everything just based on the closed corners. Either way, the
        coordinate system will be normalized to the closed left corner, applying rotation and
        shifting, and eventually convert the coordinates into cm values. As the very last step,
        self.bodyparts will be created - a dictionary that contains an `Bodypart` object for each
        marker id. Logs will save all parameters specified upon calling the function to make it
        easily traceable, which parameter settings were used to obtain a given result (see also
        the documentation of the `export_results()` method.
        """
        initial_logs_to_add = {
            "critical_markers_to_compute_coverage": marker_ids_to_compute_coverage,
            "coverage_threshold": coverage_threshold,
            "max_seconds_to_interpolate": max_seconds_to_interpolate,
            "min_likelihood_threshold": likelihood_threshold,
            "center_of_gravity_based_on": marker_ids_to_compute_center_of_gravity,
            "relative_error_tolerance_corner_detection": relative_maze_normalization_error_tolerance,
        }
        window_length = utils.get_max_odd_n_frames_for_time_interval(
            fps=self.fps, time_interval=max_seconds_to_interpolate
        )
        marker_ids_to_preprocess = preprocess.get_preprocessing_relevant_marker_ids(
            df=self.loaded_tracking_df,
            marker_ids_to_exclude=self.marker_ids_to_exclude_for_smoothing_and_interpolation,
        )
        smoothed_df = preprocess.smooth_tracked_coords_and_likelihood(
            df=self.loaded_tracking_df,
            window_length=window_length,
            marker_ids=marker_ids_to_preprocess,
            polyorder=3,
        )
        interpolated_df = preprocess.interpolate_low_likelihood_intervals(
            df=smoothed_df,
            marker_ids=marker_ids_to_preprocess,
            max_interval_length=window_length,
            framerate=self.framerate,
        )
        interpolated_df_with_cog = (
            preprocess.add_new_marker_derived_from_existing_markers(
                df=interpolated_df,
                existing_markers=marker_ids_to_compute_center_of_gravity,
                new_marker_id="CenterOfGravity",
                likelihood_threshold=likelihood_threshold,
            )
        )
        preprocessed_df = preprocess.interpolate_low_likelihood_intervals(
            df=interpolated_df_with_cog,
            marker_ids=["CenterOfGravity"],
            max_interval_length=window_length,
            framerate=self.framerate,
        )
        coverage_critical_markers = utils.compute_coverage(
            df=preprocessed_df,
            critical_marker_ids=marker_ids_to_compute_coverage,
            likelihood_threshold=likelihood_threshold,
        )
        initial_logs_to_add["coverage_critical_markers"] = coverage_critical_markers
        self._add_to_logs(logs_to_add=initial_logs_to_add)
        if coverage_critical_markers >= coverage_threshold:
            normalization_params = self._get_parameters_to_normalize_maze_coordinates(
                df=preprocessed_df,
                relative_error_tolerance=relative_maze_normalization_error_tolerance,
            )
            self.normalized_df = preprocess.normalize_df(
                df=preprocessed_df, normalization_parameters=normalization_params
            )
            self.bodyparts = preprocess.create_bodypart_objects(
                normalized_df=self.normalized_df, fps=self.fps
            )
            normalized_maze_corner_coords = self._get_normalized_maze_corners(
                normalization_parameters=normalization_params
            )
            coverage_center_of_gravity = utils.compute_coverage(
                df=self.normalized_df,
                critical_marker_ids=["CenterOfGravity"],
                likelihood_threshold=likelihood_threshold,
            )
            additional_logs_to_add = {
                "coverage_CenterOfGravity": coverage_center_of_gravity
            }
            for key, value in normalization_params.items():
                additional_logs_to_add[key] = value
            for key, value in normalized_maze_corner_coords.items():
                additional_logs_to_add[f"normalized_{key}_coords"] = value
            self._add_to_logs(logs_to_add=additional_logs_to_add)

    def _add_to_logs(
        self,
        logs_to_add: Dict,  # key-value pairs of things that shall be added to the logs of the recording
    ) -> None:
        """
        At the end of all the processing, the "logged" data will be saved in the results .xlsx file.
        It thus allows the user to keep track of which settings they used for the respective analysis.
        """
        if hasattr(self, "logs") == False:
            self.logs = {}
        for key, value in logs_to_add.items():
            assert (
                key not in self.logs.keys()
            ), f"{key} is already in self.logs.keys - adding it would result in overwriting the previous entry; please ensure unique naming"
            self.logs[key] = value

    def _get_parameters_to_normalize_maze_coordinates(
        self, df: pd.DataFrame, relative_error_tolerance: float
    ) -> Dict:
        corners = preprocess.get_corner_coords_with_likelihoods(df=df)
        translation_vector = preprocess.get_translation_vector(
            coords_to_become_origin=corners["MazeCornerClosedLeft"]["coords"]
        )
        best_result = preprocess.evaluate_maze_shape_using_open_corners(
            corners_and_likelihoods=corners, tolerance=relative_error_tolerance
        )
        if best_result["valid"]:
            side_id = best_result["side_id"]
            logs_to_add = {
                "maze_normalization_based_on": f"MazeCornerClosed{side_id}_and_MazeCornerOpen{side_id}"
            }
            conversion_factor = preprocess.get_conversion_factor_px_to_cm(
                coords_point_a=corners[f"MazeCornerClosed{side_id}"]["coords"],
                coords_point_b=corners[f"MazeCornerOpen{side_id}"]["coords"],
                distance_in_cm=50,
            )
            rotation_angle = preprocess.get_rotation_angle_with_open_corner(
                corners=corners,
                side_id=best_result["side_id"],
                translation_vector=translation_vector,
                conversion_factor=conversion_factor,
            )
        else:
            logs_to_add = {
                "maze_normalization_based_on": f"MazeCornerClosedRight_and_MazeCornerClosedLeft"
            }
            conversion_factor = preprocess.get_conversion_factor_px_to_cm(
                coords_point_a=corners["MazeCornerClosedLeft"]["coords"],
                coords_point_b=corners["MazeCornerClosedRight"]["coords"],
                distance_in_cm=4,
            )

            rotation_angle = preprocess.get_rotation_angle_with_closed_corners_only(
                corners=corners,
                translation_vector=translation_vector,
                conversion_factor=conversion_factor,
            )
        self._add_to_logs(logs_to_add=logs_to_add)
        return {
            "translation_vector": translation_vector,
            "rotation_angle": rotation_angle,
            "conversion_factor": conversion_factor,
        }

    def _get_normalized_maze_corners(self, normalization_parameters: Dict) -> Dict:
        normalized_maze_corner_coordinates = {}
        corners = preprocess.get_corner_coords_with_likelihoods(df=self.normalized_df)
        for corner_marker_id in corners.keys():
            normalized_maze_corner_coordinates[corner_marker_id] = corners[
                corner_marker_id
            ]["coords"]
        return normalized_maze_corner_coordinates

    def run_event_detection(
        self,
        bodyparts_critical_for_freezing: List[str] = [
            "Snout",
            "CenterOfGravity",
        ],  # List of marker_ids that need to be classified as "immobile" that the entire mouse is classified as immobile
        bodyparts_for_direction_front_to_back: List[str] = [
            "Snout",
            "CenterOfGravity",
        ],  # List of exactly two markers_ids that will be used to determine the orientation of the mouse; order of the markers: front, back
        immobility_max_rolling_speed: float = 2.0,  # A given bodypart will be classified as immobile in all frames, in which its rolling speed is less or equal to this value [cm per s]
        immobility_min_duration: float = 0.1,  # An immobility event will be created for each interval that exceeds this threshold [s] and in which all `bodyparts_critical_for_freezing` are classified as immobile
        freezing_min_duration: float = 0.5,  # A freezing event will be created for each interval that meets or exceeds this threshold [s] and in which all `bodyparts_critical_for_freezing` are classified as immobile
        gait_min_rolling_speed: float = 3.0,  # Frames will be classified as gait if the center of gravity meets or exceeds this rolling speed threshold [cm per s]
        gait_min_duration: float = 0.5,  # A gait event will be created for each interval that meets or exceeds this threshold [s] and in which the movement of the center of gravitiy is classified as "gait"
        gait_disruption_max_time_to_immobility: float = 0.15,  # A gait-disruption event will be created if an "immobility" event is found within this time interval [s] after a "gait" event has ended
        bodyparts_to_include_in_behavior_df=[
            "CenterOfGravity",
            "Snout",
            "TailBase",
        ],  # Bodyparts that will be included in the `behavior_df`. Should include at least all `bodyparts_critical_for_freezing`, `bodyparts_for_direction_front_to_back`, and the center of gravity
        merge_events_max_inbetween_time: float = 0.15,  # Not implemented yet;
    ) -> None:
        """
        Trigger the detection of "immobility", "freezing", "gait", and "gait disruption" events. All
        detected events will also be classified "towards open end" True / False, depending on whether
        the orientation of the mouse was towards the open end of the maze or not. For gait events,
        this classification will be based on the orientation at the very last frame, while the orientation
        at the very first frame will be used for all other events ("immobility", "freezing", and "gait
        disruption"). Since the specified parameters used for this step of the analysis are very crucial,
        they will be documented in the logs and exported once the `export_results()` method is called.
        Default parameter selection for immobility speed threshold & freezing minimum duration was based
        on: https://www.sciencedirect.com/science/article/pii/S0960982216306182
        """
        sliding_window_size = int(
            round(
                utils.get_max_odd_n_frames_for_time_interval(
                    fps=self.fps, time_interval=0.5
                )
                / 2,
                0,
            )
        )
        logs_to_add = {
            "bodyparts_checked_to_infer_immobility": bodyparts_critical_for_freezing,
            "bodypart_used_to_identify_front": bodyparts_for_direction_front_to_back[0],
            "bodypart_used_to_identify_back": bodyparts_for_direction_front_to_back[1],
            "immobility_max_rolling_speed": immobility_max_rolling_speed,
            "immobility_min_duration": immobility_min_duration,
            "freezing_min_duration": freezing_min_duration,
            "gait_min_rolling_speed": gait_min_rolling_speed,
            "gait_min_duration": gait_min_duration,
            "gait_disruption_max_time_to_immobility": gait_disruption_max_time_to_immobility,
            #'merge_events_max_inbetween_time': merge_events_max_inbetween_time,
            "sliding_window_size_to_compute_speed": sliding_window_size,
        }
        self._add_to_logs(logs_to_add=logs_to_add)
        for bodypart in self.bodyparts.values():
            bodypart.calculate_speed_and_identify_immobility(
                sliding_window_size=sliding_window_size,
                immobility_threshold=immobility_max_rolling_speed,
            )
        self.behavior_df = detect.create_behavior_df(
            normalized_df=self.normalized_df,
            bodyparts_to_include=bodyparts_to_include_in_behavior_df,
        )
        self.behavior_df = detect.add_orientation_to_behavior_df(
            behavior_df=self.behavior_df,
            all_bodyparts=self.bodyparts,
            bodyparts_for_direction_front_to_back=bodyparts_for_direction_front_to_back,
        )
        self.behavior_df = (
            detect.add_immobility_based_on_several_bodyparts_to_behavior_df(
                behavior_df=self.behavior_df,
                all_bodyparts=self.bodyparts,
                bodyparts_critical_for_freezing=bodyparts_critical_for_freezing,
            )
        )
        immobility_events = detect.get_immobility_related_events(
            behavior_df=self.behavior_df,
            fps=self.fps,
            min_interval_duration=immobility_min_duration,
            event_type="immobility_bout",
        )
        self.behavior_df = detect.add_event_bouts_to_behavior_df(
            behavior_df=self.behavior_df,
            event_type="immobility_bout",
            events=immobility_events,
        )
        freezing_events = detect.get_immobility_related_events(
            behavior_df=self.behavior_df,
            fps=self.fps,
            min_interval_duration=freezing_min_duration,
            event_type="freezing_bout",
        )
        self.behavior_df = detect.add_event_bouts_to_behavior_df(
            behavior_df=self.behavior_df,
            event_type="freezing_bout",
            events=freezing_events,
        )
        gait_events = detect.get_gait_events(
            all_bodyparts=self.bodyparts,
            fps=self.fps,
            gait_min_rolling_speed=gait_min_rolling_speed,
            gait_min_duration=gait_min_duration,
        )
        self.behavior_df = detect.add_event_bouts_to_behavior_df(
            behavior_df=self.behavior_df, event_type="gait_bout", events=gait_events
        )
        gait_disruption_events = detect.get_gait_disruption_events(
            behavior_df=self.behavior_df,
            fps=self.fps,
            gait_events=gait_events,
            gait_disruption_max_time_to_immobility=gait_disruption_max_time_to_immobility,
        )
        self.behavior_df = detect.add_event_bouts_to_behavior_df(
            behavior_df=self.behavior_df,
            event_type="gait_disruption_bout",
            events=gait_disruption_events,
        )

    def export_results(
        self,
        out_dir_path: Path,  # Filepath to the directory where the results shall be saved
    ) -> None:
        """
        Creates an .xlsx file the following tab names that contain the following information:
        - parameter_settings:
            The logged parameter settings that were passed (or used by default) during the processing.

        - immobility_bouts:
            Data that describe each individual immobility event that was detected.

        - freezing_bouts:
            Data that describe each individual freezing event that wes detected.

        - gait_bouts:
            Data that describe each individual gait event that was detected.

        - gait_discruption_bouts:
            Data that describe each individual gait disruption event that was detected.

        - session_overview:
            Data that summarizes all events of a given type (e.g. freezing) over the entire session.
            Events from a given type are also futher split according to the determined orientation and
            are thus always given as "towards_open_<event_type>", "towards_closed_<event_type>", or
            as "all_<event_type>" (e.g. "all_freezing_bouts", or "towards_open_gait_events").
        """
        dfs_to_export = {
            "immobility_bouts": export.export_immobility_related_bouts(
                df=self.behavior_df,
                event_type="immobility_bout",
                framerate=self.framerate,
            ),
            "freezing_bouts": export.export_immobility_related_bouts(
                df=self.behavior_df,
                event_type="freezing_bout",
                framerate=self.framerate,
            ),
            "gait_disruption_bouts": export.export_immobility_related_bouts(
                df=self.behavior_df,
                event_type="gait_disruption_bout",
                framerate=self.framerate,
            ),
            "gait_bouts": export.export_gait_related_bouts(
                df=self.behavior_df, event_type="gait_bout", framerate=self.framerate
            ),
        }
        dfs_to_export["session_overview"] = export.create_session_overview_df(
            dfs_to_export_with_individual_bout_dfs=dfs_to_export
        )
        dfs_to_export["parameter_settings"] = export.create_parameter_settings_df(
            logs=self.logs
        )
        self.base_output_filepath = out_dir_path.joinpath(
            f'{self.metadata["animal"]}_{self.metadata["paradigm"]}_week-{self.week_id}'
        )
        export.write_xlsx_file_to_disk(
            base_output_filepath=self.base_output_filepath, dfs_to_export=dfs_to_export
        )

    def inspect_processing(
        self,
        marker_ids_to_inspect: List[str] = [
            "CenterOfGravity"
        ],  # List of marker ids whose tracked position shall be plotted as scatter plot
        verbose: bool = False,  # Whether current progress and additional details like the coverages for each marker defined in `marker_ids_to_inspect` shall be printed
        show_plot: bool = False,  # Whether the created plot shall be displayed in the jupyter notebook
        save_plot: bool = True,  # Whether the created plot shall be saved (per default, it will be saved in the save directory that was specified when `export_results()` was called)
        show_legend: bool = False,  # Whether a legend shall be added to the plot or not
    ) -> None:
        """
        Since the normalization of the maze during the preprocessing is a very critical step, this
        functions allows the user to visually inspect the success of this step at a glanze. It will
        plot a schematic of how the maze is expected to look like in form of a black box. It will then
        also plot the "most reliable" estimate of the four maze corners after normalization and the
        valid tracking positions of all markers specified in `marker_ids_to_inspect`. To avoid over-
        loading of this plot, we recommend to pass only a single marker id, for instance the
        "CenterOfGravity". If you´d like to see some additional details, consider passing "verbose = True"
        and/or "show_legend = True".
        Note: The filepath where the plots will be saved cannot be changed & they will always be saved
        in the same directory that was specified when calling the `export_results()` method, to ensure
        that they are saved in the same folder as the .xlsx files that contain logged parameter settings.
        """
        if verbose:
            print(f"Inspection of file: {self.filepath.name}:")
            for marker_id in marker_ids_to_inspect:
                coverage = utils.compute_coverage(
                    df=self.bodyparts[marker_id].df, critical_marker_ids=[marker_id]
                )
                print(
                    f'... coverage of "{marker_id}" was at: {round(coverage*100, 2)} %'
                )
        if show_plot or save_plot:
            self._plot_selected_marker_ids_on_normalized_maze(
                marker_ids=marker_ids_to_inspect,
                show=show_plot,
                save=save_plot,
                legend=show_legend,
            )

    def _plot_selected_marker_ids_on_normalized_maze(
        self,
        marker_ids: List[str],
        show: bool = True,
        save: bool = True,
        legend: bool = False,
    ) -> None:
        fig = plt.figure(figsize=(10, 5), facecolor="white")
        ax = fig.add_subplot(111)
        for corner_marker_id in [
            "MazeCornerClosedRight",
            "MazeCornerClosedLeft",
            "MazeCornerOpenRight",
            "MazeCornerOpenLeft",
        ]:
            x, y = self.logs[f"normalized_{corner_marker_id}_coords"]
            plt.scatter(x, y, label=corner_marker_id)
        for marker_id in marker_ids:
            plt.scatter(
                self.bodyparts[marker_id].df["x"],
                self.bodyparts[marker_id].df["y"],
                alpha=0.1,
                label=marker_id,
            )
        plt.plot([0, 0, 50, 50, 0], [0, 4, 4, 0, 0], c="black")
        ax.set_aspect("equal")
        if legend:
            plt.legend()
        if save:
            assert hasattr(
                self, "base_output_filepath"
            ), 'You must run ".export_results()" first if you´d like to save the plot. Alternatively, just opt to show the plot.'
            plt.savefig(f"{self.base_output_filepath}.png", dpi=300)
        if show:
            plt.show()
        else:
            plt.close()
