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

# %% auto 0
__all__ = ['process_all_dlc_tracking_h5s_with_default_settings', 'get_max_odd_n_frames_for_time_interval', 'get_column_names',
           'get_interval_border_idxs', 'compute_coverage', 'get_idxs_where_all_markers_exceed_likelihood']

# %% ../../nbs/01_twoD_00_utils.ipynb 3
from typing import List, Tuple, Dict, Optional
from pathlib import Path
import pandas as pd
import numpy as np
from tqdm import tqdm
from scipy.signal import savgol_filter

from ..core import EventBout
from . import cam_interfaces

# %% ../../nbs/01_twoD_00_utils.ipynb 4
# ToDo: If additional camera interfaces will be implemented, they should be specifieable here


def process_all_dlc_tracking_h5s_with_default_settings(
    in_dir_path: Path,  # path to the input directory which contains all DLC tracking data results
    out_dir_path: Path,  # path to the output directory where all processed results will be saved
    week_id: int,  # number of weeks post injection
) -> None:
    filepaths_dlc_trackings = []
    for filepath in in_dir_path.iterdir():
        if filepath.name.endswith(".h5"):
            if "filtered" not in filepath.name:
                filepaths_dlc_trackings.append(filepath)
    for filepath in tqdm(filepaths_dlc_trackings):
        recording = cam_interfaces.TopTracked2DRecording(
            filepath=filepath, week_id=week_id
        )
        if recording.df_successfully_loaded:
            recording.preprocess_tracking()
            if (
                recording.logs["coverage_critical_markers"]
                >= recording.logs["coverage_threshold"]
            ):
                recording.run_event_detection()
                recording.export_results(out_dir_path=out_dir_path)
                recording.inspect_processing()

# %% ../../nbs/01_twoD_00_utils.ipynb 5
def get_max_odd_n_frames_for_time_interval(
    fps: int,  # frames per second of the recording
    time_interval: 0.5,  # desired maximal time interval in seconds; default = 0.5 s
) -> int:
    """
    For the savgol_filter function of scipy - which will be used during preprocessing to smooth the data -
    you need an odd integer as the window_length parameter. This function helps to find the maximum odd number
    of frames that still fit within a specified time interval at a given fps.
    """
    assert type(fps) == int, '"fps" has to be an integer!'
    frames_per_time_interval = fps * time_interval
    if frames_per_time_interval % 2 == 0:
        max_odd_frame_count = frames_per_time_interval - 1
    elif frames_per_time_interval == int(frames_per_time_interval):
        max_odd_frame_count = frames_per_time_interval
    else:
        frames_per_time_interval = int(frames_per_time_interval)
        if frames_per_time_interval % 2 == 0:
            max_odd_frame_count = frames_per_time_interval - 1
        else:
            max_odd_frame_count = frames_per_time_interval
    assert (
        max_odd_frame_count > 0
    ), f"The specified time interval is too short to fit an odd number of frames"
    return int(max_odd_frame_count)

# %% ../../nbs/01_twoD_00_utils.ipynb 6
def get_column_names(
    df: pd.DataFrame,
    column_identifiers: List[str],
    marker_ids: List[str] = ["all"],
) -> List[str]:
    matching_column_names = []
    for column_name in df.columns:
        marker_id, column_identifier = column_name.split("_")
        if marker_ids == ["all"]:
            if column_identifier in column_identifiers:
                matching_column_names.append(column_name)
        else:
            if (marker_id in marker_ids) and (column_identifier in column_identifiers):
                matching_column_names.append(column_name)
    return matching_column_names

# %% ../../nbs/01_twoD_00_utils.ipynb 7
def get_interval_border_idxs(
    all_matching_idxs: np.ndarray,
    framerate: float,
    min_interval_duration: Optional[float] = None,
    max_interval_duration: Optional[float] = None,
) -> List[Tuple[int, int]]:
    interval_border_idxs = []
    if all_matching_idxs.shape[0] >= 1:
        step_idxs = np.where(np.diff(all_matching_idxs) > 1)[0]
        step_end_idxs = np.concatenate(
            [step_idxs, np.array([all_matching_idxs.shape[0] - 1])]
        )
        step_start_idxs = np.concatenate([np.array([0]), step_idxs + 1])
        interval_start_idxs = all_matching_idxs[step_start_idxs]
        interval_end_idxs = all_matching_idxs[step_end_idxs]
        for start_idx, end_idx in zip(interval_start_idxs, interval_end_idxs):
            interval_frame_count = (end_idx + 1) - start_idx
            interval_duration = interval_frame_count * framerate
            if (min_interval_duration != None) and (max_interval_duration != None):
                append_interval = (
                    min_interval_duration <= interval_duration <= max_interval_duration
                )
            elif min_interval_duration != None:
                append_interval = min_interval_duration <= interval_duration
            elif max_interval_duration != None:
                append_interval = interval_duration <= max_interval_duration
            else:
                append_interval = True
            if append_interval:
                interval_border_idxs.append((start_idx, end_idx))
    return interval_border_idxs

# %% ../../nbs/01_twoD_00_utils.ipynb 8
def compute_coverage(
    df: pd.DataFrame, critical_marker_ids: List[str], likelihood_threshold: float = 0.5
) -> float:
    idxs_where_all_markers_exceed_likelihood_threshold = (
        get_idxs_where_all_markers_exceed_likelihood(
            df=df,
            marker_ids=critical_marker_ids,
            likelihood_threshold=likelihood_threshold,
        )
    )
    return idxs_where_all_markers_exceed_likelihood_threshold.shape[0] / df.shape[0]

# %% ../../nbs/01_twoD_00_utils.ipynb 9
def get_idxs_where_all_markers_exceed_likelihood(
    df: pd.DataFrame, marker_ids: List[str], likelihood_threshold: float = 0.5
) -> np.ndarray:
    valid_idxs_per_marker_id = []
    for marker_id in marker_ids:
        valid_idxs_per_marker_id.append(
            df.loc[df[f"{marker_id}_likelihood"] >= likelihood_threshold].index.values
        )
    shared_valid_idxs_for_all_markers = valid_idxs_per_marker_id[0]
    if len(valid_idxs_per_marker_id) > 1:
        for i in range(1, len(valid_idxs_per_marker_id)):
            shared_valid_idxs_for_all_markers = np.intersect1d(
                shared_valid_idxs_for_all_markers, valid_idxs_per_marker_id[i]
            )
    return shared_valid_idxs_for_all_markers
