"""
.. _pipeline_val_misorientation_results:

Full-pipeline performance under simulated mounting errors
=========================================================

This analysis compares the free-living full-pipeline validation results across
simulated lower-back sensor mounting orientations. It uses the result files
generated by :ref:`pipeline_val_misorientation_gen`.

The comparison contains the default full pipeline without per-GS reorientation
and a full-mode reorientation variant where both cohort-specific sub-pipelines
use :class:`~mobgap.gait_sequences.GsdIonescu`.

The default regular-walking GSD, :class:`~mobgap.gait_sequences.GsdIluz`, is
therefore present in the default pipeline only. It is not evaluated with
reorientation enabled because ``GsdIluz`` requires body-frame input, while
per-GS reorientation requires gait sequence detection to work on the unknown
sensor frame before correction.

The reorientation-enabled variant intentionally uses ``correction_mode="full"``
instead of the default ``trust_gravity`` mode. This analysis creates one result
for every simulated orientation class per recording, i.e. an equal-prevalence
stress test. This makes the specific front-back flip class that
``trust_gravity`` intentionally does not correct much more common than expected
in a realistic low-error setting. Under this validation setup, ``full`` is the
appropriate correction mode. For realistic expected orientation-error
prevalence, ``trust_gravity`` can outperform ``full``; see the dedicated
reorientation validation analysis for that trade-off.

The analysis remains split into regular-walking and impaired-walking cohort
groups because the official full-pipeline sub-pipelines differ downstream.

"""

# %%
# Compared Pipelines
# ------------------
from pathlib import Path

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from mobgap.data.validation_results import ValidationResultLoader
from mobgap.plotting import (
    calc_min_max_with_margin,
    make_square,
    move_legend_outside,
)
from mobgap.re_orientation.pipeline import REORIENTATION_LABELS
from mobgap.utils.misc import get_env_var

algorithms = {
    "Official_MobiliseD_Pipeline": "Default pipeline",
    "Official_MobiliseD_Pipeline__gsd_ionescu_reorientation": (
        "GsdIonescu + full reorientation"
    ),
}
algorithm_order = list(algorithms.values())
orientation_order = list(REORIENTATION_LABELS)

regular_walking_cohorts = ["HA", "COPD", "CHF"]
impaired_walking_cohorts = ["MS", "PD", "PFF"]

dmos = {
    "walking_speed_mps": ("Walking speed", "m/s"),
    "stride_length_m": ("Stride length", "m"),
    "cadence_spm": ("Cadence", "steps/min"),
}

# %%
# Why Full Reorientation Mode Here?
# ---------------------------------
# The full-pipeline misorientation validation is not meant to model realistic
# mounting-error prevalence. Instead, it deliberately creates an equal number of
# datasets for the identity orientation and for every simulated misorientation.
# With eight simulated orientation classes, the validation therefore corresponds
# to the prevalence example below.
equal_prevalence_example = pd.DataFrame(
    {
        "Prevalence in this validation": [
            1 / len(orientation_order),
            (len(orientation_order) - 1) / len(orientation_order),
            1 / len(orientation_order),
        ],
        "Interpretation": [
            "correctly oriented recordings",
            "recordings with any simulated misorientation",
            "front-back flip class ignored by trust_gravity",
        ],
    },
    index=[
        "Identity orientation",
        "Any misorientation",
        "trust_gravity-ignored PA flip",
    ],
)
equal_prevalence_example.style.format(
    {"Prevalence in this validation": "{:.1%}"}
)

# %%
# Under this equal-prevalence stress test, the PA flip class that
# ``trust_gravity`` intentionally skips is frequent enough that the ``full``
# mode is the more informative validation target. In realistic studies, the
# expected misorientation rate is usually much lower; in that regime,
# ``trust_gravity`` can outperform ``full`` because it avoids unnecessary
# data-driven front-back corrections for correctly oriented gait sequences.

# %%
# Loading The Free-Living Results
# -------------------------------
# The result files add ``orientation`` to the usual TVS index. The standard
# validation loader keeps unknown CSV columns as data columns, so we can still
# use it for local and remote loading.
local_data_path = (
    Path(get_env_var("MOBGAP_VALIDATION_DATA_PATH")) / "results"
    if int(get_env_var("MOBGAP_VALIDATION_USE_LOCAL_DATA", 0))
    else None
)
__RESULT_VERSION = "main"
loader = ValidationResultLoader(
    "full_pipeline_misorientation",
    result_path=local_data_path,
    version=__RESULT_VERSION,
)


def load_misorientation_results() -> pd.DataFrame:
    results = []
    for folder_name, algorithm_label in algorithms.items():
        data = (
            loader.load_single_results(folder_name, "free_living")
            .reset_index()
            .assign(algorithm=algorithm_label)
        )
        results.append(data)
    return pd.concat(results, ignore_index=True)


def combined_error_long(
    data: pd.DataFrame, *, value_suffix: str
) -> pd.DataFrame:
    results = []
    id_vars = [
        "algorithm",
        "orientation",
        "cohort",
        "participant_id",
        "time_measure",
        "recording",
        "recording_name",
        "recording_name_pretty",
    ]
    for dmo, (dmo_label, unit) in dmos.items():
        column = f"combined__{dmo}__{value_suffix}"
        results.append(
            data[[*id_vars, column]]
            .rename(columns={column: "error"})
            .assign(dmo=dmo_label, unit=unit)
        )
    return pd.concat(results, ignore_index=True)


free_living_results = load_misorientation_results()
combined_errors = combined_error_long(free_living_results, value_suffix="error")
combined_abs_errors = combined_error_long(
    free_living_results, value_suffix="abs_error"
)

# %%
# Plot Helpers
# ------------
sns.set_context("talk")


def plot_combined_error_boxplots(
    data: pd.DataFrame, *, title: str, ylabel_prefix: str
) -> None:
    fig, axes = plt.subplots(
        1, 3, figsize=(20, 7), sharex=True, constrained_layout=True
    )
    for ax, (dmo_label, unit) in zip(axes, dmos.values()):
        plot_data = data[data["dmo"] == dmo_label]
        sns.boxplot(
            data=plot_data,
            x="algorithm",
            y="error",
            hue="orientation",
            order=algorithm_order,
            hue_order=orientation_order,
            showmeans=True,
            ax=ax,
        )
        ax.axhline(0, color="0.4", linewidth=1, linestyle=":", zorder=-50)
        ax.set_title(dmo_label)
        ax.set_xlabel("")
        ax.set_ylabel(f"{ylabel_prefix} [{unit}]")
        ax.tick_params(axis="x", rotation=20)
        ax.grid(True, axis="y", alpha=0.3)
        ax.legend(title="Orientation")

    fig.suptitle(title)
    move_legend_outside(fig, axes[-1], ncol=4)
    plt.show()


def matched_wb_count_table(data: pd.DataFrame) -> pd.DataFrame:
    matched_wb_counts = data.assign(
        n_matched_wbs=data["matched__n_matched_wbs"].fillna(0).astype(int)
    )
    table = (
        matched_wb_counts.pivot_table(
            index="algorithm",
            columns="orientation",
            values="n_matched_wbs",
            aggfunc="sum",
            sort=False,
        )
        .reindex(index=algorithm_order, columns=orientation_order)
        .fillna(0)
        .astype(int)
    )
    return table.rename_axis(index="Algorithm", columns="Orientation")


def plot_matched_wb_counts(data: pd.DataFrame, *, title: str) -> None:
    plot_data = (
        data.assign(
            n_matched_wbs=data["matched__n_matched_wbs"].fillna(0).astype(int)
        )
        .groupby(["algorithm", "orientation"], sort=False)["n_matched_wbs"]
        .sum()
        .reset_index()
    )
    fig, ax = plt.subplots(figsize=(16, 7), constrained_layout=True)
    sns.barplot(
        data=plot_data,
        x="orientation",
        y="n_matched_wbs",
        hue="algorithm",
        order=orientation_order,
        hue_order=algorithm_order,
        ax=ax,
    )
    ax.set_title(title)
    ax.set_xlabel("Orientation")
    ax.set_ylabel("# matched WBs")
    ax.tick_params(axis="x", rotation=25)
    ax.grid(True, axis="y", alpha=0.3)
    move_legend_outside(fig, ax, ncol=2)
    plt.show()


def identity_walking_speed_error_comparison(data: pd.DataFrame) -> pd.DataFrame:
    id_vars = [
        "cohort",
        "participant_id",
        "time_measure",
        "recording",
        "recording_name",
        "recording_name_pretty",
    ]
    error_col = "combined__walking_speed_mps__error"
    identity_data = data[data["orientation"] == "identity"]
    baseline_errors = (
        identity_data[identity_data["algorithm"] == "Default pipeline"][
            [*id_vars, error_col]
        ]
        .rename(columns={error_col: "baseline_error"})
        .copy()
    )
    reorientation_errors = (
        identity_data[
            identity_data["algorithm"] == "GsdIonescu + full reorientation"
        ][[*id_vars, error_col]]
        .rename(columns={error_col: "reorientation_error"})
        .copy()
    )
    return baseline_errors.merge(reorientation_errors, on=id_vars, how="inner")


def plot_identity_walking_speed_error_correlation(data: pd.DataFrame) -> None:
    comparison = identity_walking_speed_error_comparison(data)
    cohort_groups = [
        ("Regular-walking cohorts", regular_walking_cohorts),
        ("Impaired-walking cohorts", impaired_walking_cohorts),
    ]
    fig, axes = plt.subplots(1, 2, figsize=(14, 7), constrained_layout=True)
    legend_handles = {}
    for ax, (title, cohorts) in zip(axes, cohort_groups):
        plot_data = comparison[comparison["cohort"].isin(cohorts)].dropna(
            subset=["baseline_error", "reorientation_error"]
        )
        sns.scatterplot(
            data=plot_data,
            x="baseline_error",
            y="reorientation_error",
            hue="cohort",
            ax=ax,
            s=80,
        )
        min_max = calc_min_max_with_margin(
            plot_data["baseline_error"],
            plot_data["reorientation_error"],
        )
        make_square(ax, min_max)
        ax.set_title(title)
        ax.set_xlabel("Default-pipeline identity WS error [m/s]")
        ax.set_ylabel("Full-reorientation identity WS error [m/s]")
        ax.grid(True, alpha=0.3)
        handles, labels = ax.get_legend_handles_labels()
        legend_handles.update(dict(zip(labels, handles)))
        ax.get_legend().remove()

    fig.suptitle("Identity-orientation combined walking-speed errors")
    fig.legend(
        legend_handles.values(),
        legend_handles.keys(),
        loc="outside lower center",
        ncol=3,
        frameon=False,
    )
    plt.show()


def cohort_group_data(data: pd.DataFrame, cohorts: list[str]) -> pd.DataFrame:
    return data[data["cohort"].isin(cohorts)]


# %%
# Matched Walking Bout Counts
# ---------------------------
# These tables and plots show the total number of matched walking bouts per
# pipeline variant and simulated orientation. They are split by cohort group
# because the official healthy and impaired sub-pipelines differ downstream.
matched_wb_counts_regular = matched_wb_count_table(
    cohort_group_data(free_living_results, regular_walking_cohorts)
)
matched_wb_counts_regular  # noqa: B018

# %%
matched_wb_counts_impaired = matched_wb_count_table(
    cohort_group_data(free_living_results, impaired_walking_cohorts)
)
matched_wb_counts_impaired  # noqa: B018

# %%
plot_matched_wb_counts(
    cohort_group_data(free_living_results, regular_walking_cohorts),
    title="Matched WBs in regular-walking cohorts",
)

# %%
plot_matched_wb_counts(
    cohort_group_data(free_living_results, impaired_walking_cohorts),
    title="Matched WBs in impaired-walking cohorts",
)

# %%
# Identity-Orientation Walking-Speed Error Correlation
# ----------------------------------------------------
# This plot checks whether enabling full-mode reorientation changes walking-speed
# performance when the simulated mounting is already correct.
plot_identity_walking_speed_error_correlation(free_living_results)

# %%
# Regular-Walking Cohorts
# -----------------------
# ``GsdIluz`` is the default GSD for these cohorts. It remains part of the
# default-pipeline baseline, but the full-mode reorientation variant uses
# ``GsdIonescu`` because ``GsdIluz`` cannot run before per-GS reorientation.
plot_combined_error_boxplots(
    cohort_group_data(combined_errors, regular_walking_cohorts),
    title="Combined DMO errors in regular-walking cohorts",
    ylabel_prefix="Combined error",
)

# %%
plot_combined_error_boxplots(
    cohort_group_data(combined_abs_errors, regular_walking_cohorts),
    title="Combined absolute DMO errors in regular-walking cohorts",
    ylabel_prefix="Combined absolute error",
)

# %%
# Impaired-Walking Cohorts
# ------------------------
# These cohorts already use ``GsdIonescu`` in the default pipeline, so the main
# change for the modified variant is enabling full-mode per-GS reorientation.
plot_combined_error_boxplots(
    cohort_group_data(combined_errors, impaired_walking_cohorts),
    title="Combined DMO errors in impaired-walking cohorts",
    ylabel_prefix="Combined error",
)

# %%
plot_combined_error_boxplots(
    cohort_group_data(combined_abs_errors, impaired_walking_cohorts),
    title="Combined absolute DMO errors in impaired-walking cohorts",
    ylabel_prefix="Combined absolute error",
)