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

# %% auto 0
__all__ = ['export_immobility_related_bouts', 'get_all_bout_ids', 'get_bout_duration_per_bout_id',
           'get_column_values_at_event_borders', 'get_distance_covered_per_event',
           'get_interval_start_and_end_time_per_event', 'export_gait_related_bouts', 'create_session_overview_df',
           'get_bout_id_splits_depending_on_direction', 'get_column_name_from_substring',
           'add_results_to_session_overview', 'create_parameter_settings_df', 'write_xlsx_file_to_disk']

# %% ../../nbs/01_twoD_04_export.ipynb 3
from typing import List, Tuple, Dict, Optional
from pathlib import Path
import pandas as pd
import numpy as np

# %% ../../nbs/01_twoD_04_export.ipynb 4
def export_immobility_related_bouts(
    df: pd.DataFrame, event_type: str, framerate: float
) -> pd.DataFrame:
    results_per_event = {
        "bout_id": [],
        "duration": [],
        "CenterOfGravity_x_at_bout_start": [],
        "towards_open_at_bout_start": [],
        "distance_covered_cm": [],
        "start_time": [],
        "end_time": [],
    }
    results_per_event["bout_id"] = get_all_bout_ids(df=df, event_type=event_type)
    if len(results_per_event["bout_id"]) >= 1:
        results_per_event["duration"] = get_bout_duration_per_bout_id(
            df=df, event_type=event_type, event_ids=results_per_event["bout_id"]
        )
        x_positions_center_of_gravity_at_interval_borders = (
            get_column_values_at_event_borders(
                df=df,
                event_type=event_type,
                event_ids=results_per_event["bout_id"],
                column_name="CenterOfGravity_x",
            )
        )
        results_per_event[
            "CenterOfGravity_x_at_bout_start"
        ] = x_positions_center_of_gravity_at_interval_borders[:, 0]
        direction_towards_open_at_interval_borders = get_column_values_at_event_borders(
            df=df,
            event_type=event_type,
            event_ids=results_per_event["bout_id"],
            column_name="facing_towards_open_end",
        )
        results_per_event[
            "towards_open_at_bout_start"
        ] = direction_towards_open_at_interval_borders[:, 0]
        results_per_event["distance_covered_cm"] = get_distance_covered_per_event(
            df=df,
            event_type=event_type,
            event_ids=results_per_event["bout_id"],
            marker_id="CenterOfGravity",
        )
        bout_start_and_end_time = get_interval_start_and_end_time_per_event(
            df=df,
            event_type=event_type,
            event_ids=results_per_event["bout_id"],
            framerate=framerate,
        )
        results_per_event["start_time"] = bout_start_and_end_time[:, 0]
        results_per_event["end_time"] = bout_start_and_end_time[:, 1]
    return pd.DataFrame(data=results_per_event)

# %% ../../nbs/01_twoD_04_export.ipynb 5
def get_all_bout_ids(df: pd.DataFrame, event_type: str) -> np.ndarray:
    return df[f"{event_type}_id"].dropna().unique()

# %% ../../nbs/01_twoD_04_export.ipynb 6
def get_bout_duration_per_bout_id(
    df: pd.DataFrame, event_type: str, event_ids: List[float]
) -> List[float]:
    durations = []
    for event_id in event_ids:
        durations.append(
            df.loc[df[f"{event_type}_id"] == event_id, f"{event_type}_duration"].iloc[0]
        )
    return durations

# %% ../../nbs/01_twoD_04_export.ipynb 7
def get_column_values_at_event_borders(
    df: pd.DataFrame, event_type: str, event_ids: List[float], column_name: str
) -> np.ndarray:
    values_at_interval_borders = []
    for event_id in event_ids:
        start_value = df.loc[df[f"{event_type}_id"] == event_id, column_name].iloc[0]
        end_value = df.loc[df[f"{event_type}_id"] == event_id, column_name].iloc[-1]
        values_at_interval_borders.append((start_value, end_value))
    return np.asarray(values_at_interval_borders)

# %% ../../nbs/01_twoD_04_export.ipynb 8
def get_distance_covered_per_event(
    df: pd.DataFrame, event_type: str, event_ids: List[float], marker_id: str
) -> List[float]:
    distances_per_event = []
    for event_id in event_ids:
        df_tmp = df.loc[df[f"{event_type}_id"] == event_id].copy()
        distances_per_event.append(
            (
                (
                    df_tmp[f"{marker_id}_x"].diff() ** 2
                    + df_tmp[f"{marker_id}_y"].diff() ** 2
                )
                ** 0.5
            )
            .cumsum()
            .iloc[-1]
        )
    return distances_per_event

# %% ../../nbs/01_twoD_04_export.ipynb 9
def get_interval_start_and_end_time_per_event(
    df: pd.DataFrame, event_type: str, event_ids: List[float], framerate: float
) -> np.ndarray:
    interval_border_idxs = []
    for event_id in event_ids:
        start_time, end_time = (
            df.loc[df[f"{event_type}_id"] == event_id].index.values[[0, -1]] * framerate
        )
        interval_border_idxs.append((start_time, end_time))
    return np.asarray(interval_border_idxs)

# %% ../../nbs/01_twoD_04_export.ipynb 10
def export_gait_related_bouts(
    df: pd.DataFrame, event_type: str, framerate: float
) -> pd.DataFrame:
    # ToDo: very similar to "export_immobility_related_bouts" - can they be combined to one generalized version?
    results_per_event = {
        "bout_id": [],
        "duration": [],
        "CenterOfGravity_x_at_bout_end": [],
        "towards_open_at_bout_end": [],
        "distance_covered_cm": [],
        "start_time": [],
        "end_time": [],
    }
    results_per_event["bout_id"] = get_all_bout_ids(df=df, event_type=event_type)
    if len(results_per_event["bout_id"]) >= 1:
        results_per_event["duration"] = get_bout_duration_per_bout_id(
            df=df, event_type=event_type, event_ids=results_per_event["bout_id"]
        )
        x_positions_center_of_gravity_at_interval_borders = (
            get_column_values_at_event_borders(
                df=df,
                event_type=event_type,
                event_ids=results_per_event["bout_id"],
                column_name="CenterOfGravity_x",
            )
        )
        results_per_event[
            "CenterOfGravity_x_at_bout_end"
        ] = x_positions_center_of_gravity_at_interval_borders[:, 1]
        direction_towards_open_at_interval_borders = get_column_values_at_event_borders(
            df=df,
            event_type=event_type,
            event_ids=results_per_event["bout_id"],
            column_name="facing_towards_open_end",
        )
        results_per_event[
            "towards_open_at_bout_end"
        ] = direction_towards_open_at_interval_borders[:, 1]
        results_per_event["distance_covered_cm"] = get_distance_covered_per_event(
            df=df,
            event_type=event_type,
            event_ids=results_per_event["bout_id"],
            marker_id="CenterOfGravity",
        )
        bout_start_and_end_time = get_interval_start_and_end_time_per_event(
            df=df,
            event_type=event_type,
            event_ids=results_per_event["bout_id"],
            framerate=framerate,
        )
        results_per_event["start_time"] = bout_start_and_end_time[:, 0]
        results_per_event["end_time"] = bout_start_and_end_time[:, 1]
    return pd.DataFrame(data=results_per_event)

# %% ../../nbs/01_twoD_04_export.ipynb 11
def create_session_overview_df(
    dfs_to_export_with_individual_bout_dfs: Dict[str, pd.DataFrame]
) -> pd.DataFrame:
    session_overview = {
        "bout_type": [],
        "total_bouts_count": [],
        "total_duration": [],
        "total_distance_covered": [],
        "mean_duration": [],
        "mean_distance_covered": [],
        "mean_CenterOfGravity_x": [],
    }
    for tab_name, df in dfs_to_export_with_individual_bout_dfs.items():
        bout_ids_split_depending_on_direction = (
            get_bout_id_splits_depending_on_direction(df=df)
        )
        for (
            split_id,
            relevant_bout_ids,
        ) in bout_ids_split_depending_on_direction.items():
            session_overview = add_results_to_session_overview(
                session_overview=session_overview,
                df=df,
                event_type=tab_name,
                event_prefix=split_id,
                bout_ids=relevant_bout_ids,
            )
    return pd.DataFrame(data=session_overview)

# %% ../../nbs/01_twoD_04_export.ipynb 12
def get_bout_id_splits_depending_on_direction(
    df: pd.DataFrame,
) -> Dict[str, List[float]]:
    towards_open_column_name = get_column_name_from_substring(
        all_columns=list(df.columns), substring="towards_open"
    )
    bout_ids_split_by_direction = {
        "all": list(df["bout_id"].unique()),
        "towards_open": list(
            df.loc[df[towards_open_column_name] == True, "bout_id"].unique()
        ),
        "towards_closed": list(
            df.loc[df[towards_open_column_name] != True, "bout_id"].unique()
        ),
    }
    return bout_ids_split_by_direction

# %% ../../nbs/01_twoD_04_export.ipynb 13
def get_column_name_from_substring(all_columns: List[str], substring: str) -> str:
    matching_column_names = [
        column_name for column_name in all_columns if substring in column_name
    ]
    assert (
        len(matching_column_names) == 1
    ), f"There should be exactly one match for {substring} - however, {len(matching_column_names)} were found: [{matching_column_names}]."
    return matching_column_names[0]

# %% ../../nbs/01_twoD_04_export.ipynb 14
def add_results_to_session_overview(
    session_overview: Dict,
    df: pd.DataFrame,
    event_type: str,
    event_prefix: str,
    bout_ids: List[float],
) -> Dict:
    session_overview["bout_type"].append(f"{event_prefix}_{event_type}")
    if len(bout_ids) > 0:
        session_overview["total_bouts_count"].append(len(bout_ids))
        session_overview["total_duration"].append(
            df.loc[df["bout_id"].isin(bout_ids), "duration"].cumsum().iloc[-1]
        )
        session_overview["total_distance_covered"].append(
            df.loc[df["bout_id"].isin(bout_ids), "distance_covered_cm"]
            .cumsum()
            .iloc[-1]
        )
        session_overview["mean_duration"].append(
            df.loc[df["bout_id"].isin(bout_ids), "duration"].mean()
        )
        session_overview["mean_distance_covered"].append(
            df.loc[df["bout_id"].isin(bout_ids), "distance_covered_cm"].mean()
        )
        center_of_gravity_x_column_name = get_column_name_from_substring(
            all_columns=list(df.columns), substring="CenterOfGravity_x"
        )
        session_overview["mean_CenterOfGravity_x"].append(
            df.loc[df["bout_id"].isin(bout_ids), center_of_gravity_x_column_name].mean()
        )
    else:
        session_overview["total_bouts_count"].append(0)
        session_overview["total_duration"].append(0)
        session_overview["total_distance_covered"].append(0)
        session_overview["mean_duration"].append(np.nan)
        session_overview["mean_distance_covered"].append(np.nan)
        session_overview["mean_CenterOfGravity_x"].append(np.nan)
    return session_overview

# %% ../../nbs/01_twoD_04_export.ipynb 15
def create_parameter_settings_df(logs: Dict) -> pd.DataFrame:
    logged_settings = {"parameter": [], "specified_value": []}
    for parameter, value in logs.items():
        logged_settings["parameter"].append(parameter)
        logged_settings["specified_value"].append(value)
    return pd.DataFrame(data=logged_settings)

# %% ../../nbs/01_twoD_04_export.ipynb 16
def write_xlsx_file_to_disk(
    base_output_filepath: Path, dfs_to_export: Dict[str, pd.DataFrame]
) -> None:
    writer = pd.ExcelWriter(f"{base_output_filepath}.xlsx", engine="xlsxwriter")
    for tab_name, df in dfs_to_export.items():
        df.to_excel(writer, sheet_name=tab_name)
    writer.save()
