# Copyright (C) 2026 Hector van der Aa <hector@h3cx.dev>
# Copyright (C) 2026 Association Exergie <association.exergie@gmail.com>
# SPDX-License-Identifier: GPL-3.0-or-later

from __future__ import annotations

import argparse
import sys
from bisect import bisect_left
from dataclasses import dataclass
from pathlib import Path

import matplotlib.pyplot as plt
import pandas as pd
from matplotlib.ticker import MultipleLocator


GP_COLUMNS = ("gp0_falling", "gp1_falling")
SPARK_COLUMNS = ("spark_1", "spark_2", "spark_3")
PULSE_STYLES = {
    "gp0_falling": ("gp0 falling", "tab:green", 5.0, 0.18),
    "gp1_falling": ("gp1 falling", "tab:red", 5.0, 0.18),
    "spark_1": ("spark 1", "tab:blue", 2.4, 0.65),
    "spark_2": ("spark 2", "tab:orange", 2.4, 0.65),
    "spark_3": ("spark 3", "tab:purple", 2.4, 0.65),
}


@dataclass(frozen=True)
class EncoderPoint:
    time_us: int
    continuous_angle_deg: float


def read_csv(path: Path) -> pd.DataFrame:
    df = pd.read_csv(path)
    if "time_us" not in df.columns:
        raise ValueError("missing column: time_us")
    df["time_us"] = df["time_us"].astype("int64")
    return df


def encoder_points(df: pd.DataFrame) -> list[EncoderPoint]:
    required = {"time_us", "revolution_index", "crank_angle_deg"}
    if not required <= set(df.columns):
        return []

    points: list[EncoderPoint] = []
    angle_df = df.dropna(subset=["revolution_index", "crank_angle_deg"])
    for row in angle_df.itertuples():
        revolution_index = int(row.revolution_index)
        if revolution_index < 0:
            continue
        continuous_angle = revolution_index * 360.0 + float(row.crank_angle_deg)
        points.append(EncoderPoint(int(row.time_us), continuous_angle))

    deduped: dict[int, EncoderPoint] = {}
    for point in sorted(points, key=lambda item: item.time_us):
        deduped[point.time_us] = point
    return list(deduped.values())


def interpolate_angle(points: list[EncoderPoint], time_us: int) -> float | None:
    if not points:
        return None

    times = [point.time_us for point in points]
    index = bisect_left(times, time_us)
    if index < len(points) and points[index].time_us == time_us:
        return points[index].continuous_angle_deg
    if index == 0 or index == len(points):
        return None

    before = points[index - 1]
    after = points[index]
    span = after.time_us - before.time_us
    if span <= 0:
        return before.continuous_angle_deg

    fraction = (time_us - before.time_us) / span
    return before.continuous_angle_deg + fraction * (
        after.continuous_angle_deg - before.continuous_angle_deg
    )


def iter_csv_paths(path: Path) -> list[Path]:
    if path.is_dir():
        return sorted(path.rglob("*.csv"))
    return [path]


def has_angle_data(points: list[EncoderPoint]) -> bool:
    return len(points) >= 2


def rpm_rows(df: pd.DataFrame) -> pd.DataFrame:
    if "rpm_raw" not in df.columns or "rpm_lpf" not in df.columns:
        return pd.DataFrame()
    return df.dropna(subset=["rpm_raw", "rpm_lpf"])


def plot_time(df: pd.DataFrame, ax: plt.Axes) -> None:
    rpm_df = rpm_rows(df)
    if not rpm_df.empty:
        time_s = rpm_df["time_us"] / 1_000_000.0
        ax.plot(time_s, rpm_df["rpm_raw"], label="raw", linewidth=0.75, alpha=0.35)
        ax.plot(time_s, rpm_df["rpm_lpf"], label="lpf", linewidth=1.4)
    ax.set_xlabel("time (s)")


def plot_angle_overlay(df: pd.DataFrame, ax: plt.Axes) -> None:
    rpm_df = rpm_rows(df)
    if rpm_df.empty:
        return

    raw_label_used = False
    lpf_label_used = False
    for revolution in sorted(rpm_df["revolution_index"].dropna().astype(int).unique()):
        if revolution < 0:
            continue
        rev_df = rpm_df[rpm_df["revolution_index"].astype(int) == revolution]
        ax.plot(
            rev_df["crank_angle_deg"],
            rev_df["rpm_raw"],
            color="0.55",
            linewidth=0.5,
            alpha=0.12,
            label="raw" if not raw_label_used else None,
        )
        ax.plot(
            rev_df["crank_angle_deg"],
            rev_df["rpm_lpf"],
            linewidth=0.8,
            alpha=0.28,
            label="lpf per rev" if not lpf_label_used else None,
        )
        raw_label_used = True
        lpf_label_used = True

    ax.set_xlim(0, 360)
    ax.set_xlabel("crank angle (deg)")


def plot_angle_continuous(df: pd.DataFrame, ax: plt.Axes) -> None:
    rpm_df = rpm_rows(df)
    if rpm_df.empty:
        return

    angle_df = rpm_df.dropna(subset=["revolution_index", "crank_angle_deg"])
    angle_df = angle_df[angle_df["revolution_index"].astype(int) >= 0]
    if angle_df.empty:
        return

    continuous_angle = (
        angle_df["revolution_index"].astype(float) * 360.0
        + angle_df["crank_angle_deg"].astype(float)
    )
    ax.plot(continuous_angle, angle_df["rpm_raw"], label="raw", linewidth=0.75, alpha=0.35)
    ax.plot(continuous_angle, angle_df["rpm_lpf"], label="lpf", linewidth=1.4)

    max_angle = float(continuous_angle.max())
    turn_angle = 0.0
    while turn_angle <= max_angle:
        ax.axvline(turn_angle, color="0.7", linewidth=0.5, alpha=0.35)
        turn_angle += 360.0

    ax.set_xlim(0, max_angle)
    ax.xaxis.set_major_locator(MultipleLocator(360.0))
    ax.xaxis.set_minor_locator(MultipleLocator(180.0))
    ax.set_xlabel("continuous crank angle (deg)")


def selected_pulse_columns(
    df: pd.DataFrame,
    gp_mode: str,
    spark_columns: list[str],
) -> list[str]:
    columns: list[str] = []
    if gp_mode in {"gp0", "both"} and "gp0_falling" in df.columns:
        columns.append("gp0_falling")
    if gp_mode in {"gp1", "both"} and "gp1_falling" in df.columns:
        columns.append("gp1_falling")
    columns.extend(column for column in spark_columns if column in df.columns)
    return columns


def pulse_x_value(
    row: pd.Series,
    *,
    x_axis: str,
    points: list[EncoderPoint],
) -> float | None:
    time_us = int(row["time_us"])
    if x_axis == "time":
        return time_us / 1_000_000.0

    continuous_angle = interpolate_angle(points, time_us)
    if continuous_angle is None:
        return None
    if x_axis in {"angle", "angle-overlay"}:
        return continuous_angle % 360.0
    return continuous_angle


def plot_pulse_bars(
    df: pd.DataFrame,
    ax: plt.Axes,
    *,
    x_axis: str,
    gp_mode: str,
    spark_columns: list[str],
    points: list[EncoderPoint],
) -> None:
    for column in selected_pulse_columns(df, gp_mode, spark_columns):
        label, color, linewidth, alpha = PULSE_STYLES[column]
        pulse_df = df.loc[df[column].fillna(0).astype(int) == 1]
        plotted_label = False
        for _, row in pulse_df.iterrows():
            x_value = pulse_x_value(row, x_axis=x_axis, points=points)
            if x_value is None:
                continue
            ax.axvline(
                x_value,
                color=color,
                alpha=alpha,
                linewidth=linewidth,
                label=label if not plotted_label else None,
            )
            plotted_label = True


def plot_file(
    csv_path: Path,
    output: Path | None,
    *,
    x_axis: str,
    gp_mode: str,
    spark_columns: list[str],
) -> bool:
    try:
        df = read_csv(csv_path)
    except (OSError, ValueError) as exc:
        print(f"Could not read {csv_path}: {exc}", file=sys.stderr)
        return False

    points = encoder_points(df)
    if x_axis.startswith("angle") and not has_angle_data(points):
        print(f"Could not plot crank angle for {csv_path}: missing angle columns", file=sys.stderr)
        return False

    fig, ax = plt.subplots(figsize=(12, 5))
    if x_axis in {"angle", "angle-overlay"}:
        plot_angle_overlay(df, ax)
    elif x_axis == "angle-continuous":
        plot_angle_continuous(df, ax)
    else:
        plot_time(df, ax)

    plot_pulse_bars(
        df,
        ax,
        x_axis=x_axis,
        gp_mode=gp_mode,
        spark_columns=spark_columns,
        points=points,
    )

    ax.set_ylabel("RPM")
    ax.set_title(str(csv_path))
    ax.grid(alpha=0.25)
    if x_axis == "angle-continuous":
        ax.grid(which="minor", axis="x", alpha=0.15, linewidth=0.6)
    ax.legend()
    fig.tight_layout()

    if output is not None:
        output.parent.mkdir(parents=True, exist_ok=True)
        fig.savefig(output, dpi=160)
        print(f"Wrote {output}")
        plt.close(fig)
    else:
        plt.show()
        plt.close(fig)

    return True


def output_path_for(csv_path: Path, output: Path | None, multiple: bool) -> Path | None:
    if output is None:
        return None
    if not multiple:
        return output
    return output / f"{csv_path.stem}.png"


def parse_spark_columns(value: str) -> list[str]:
    if value == "none":
        return []
    if value == "all":
        return list(SPARK_COLUMNS)

    columns = [column.strip() for column in value.split(",") if column.strip()]
    if not columns:
        raise argparse.ArgumentTypeError("must specify at least one spark column")

    invalid = sorted(set(columns) - set(SPARK_COLUMNS))
    if invalid:
        valid = ", ".join(("none", "all", *SPARK_COLUMNS))
        raise argparse.ArgumentTypeError(
            f"invalid spark column(s): {', '.join(invalid)}; choose from {valid}"
        )

    deduped: list[str] = []
    for column in columns:
        if column not in deduped:
            deduped.append(column)
    return deduped


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Plot algo1 RPM output CSV files.")
    parser.add_argument("csv_path", type=Path, help="Algo1 CSV file or directory of CSV files.")
    parser.add_argument("-o", "--output", type=Path, help="Optional image output path or directory.")
    parser.add_argument(
        "--x",
        choices=("time", "angle", "angle-overlay", "angle-continuous"),
        default="time",
        help="Plot RPM against time, overlaid 0-360 crank angle, or continuous crank angle.",
    )
    parser.add_argument(
        "--gp",
        choices=("none", "gp0", "gp1", "both"),
        default="both",
        help="Overlay GP falling events as vertical bars.",
    )
    parser.add_argument(
        "--sparks",
        type=parse_spark_columns,
        default="all",
        help="Overlay spark events: none, all, one spark, or a comma-separated subset.",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()

    if not args.csv_path.exists():
        print(f"Path does not exist: {args.csv_path}", file=sys.stderr)
        return 1

    csv_paths = iter_csv_paths(args.csv_path)
    if not csv_paths:
        print(f"No CSV files found under {args.csv_path}", file=sys.stderr)
        return 1

    multiple = len(csv_paths) > 1
    if multiple and args.output is not None and args.output.suffix:
        print("When plotting a directory, --output must be a directory.", file=sys.stderr)
        return 1

    spark_columns = (
        args.sparks if isinstance(args.sparks, list) else parse_spark_columns(args.sparks)
    )
    failures = 0
    for csv_path in csv_paths:
        output = output_path_for(csv_path, args.output, multiple)
        if not plot_file(
            csv_path,
            output,
            x_axis=args.x,
            gp_mode=args.gp,
            spark_columns=spark_columns,
        ):
            failures += 1

    return 1 if failures else 0


if __name__ == "__main__":
    sys.exit(main())