# Copyright (C) 2026 Hector van der Aa <hector@h3cx.dev>
# Copyright (C) 2026 Pierre Barbier <pierrebarbier741@gmail.com>
# 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 pandas as pd


DEFAULT_DATA_ROOTS = (
    Path("~/projects/Exergie/Engine-Recorder/recordings-rpm").expanduser(),
    Path("~/projects/Exergie/Engine-Recorder/Python/recordings_rpm").expanduser(),
)
GP_COLUMNS = ("gp0_falling", "gp1_falling")
SPARK_COLUMNS = ("spark_1", "spark_2", "spark_3")
INTERPOLATED_COLUMNS = (
    "algo1_continuous_angle_deg",
    "algo1_revolution_index",
    "algo1_crank_angle_deg",
)


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


@dataclass(frozen=True)
class SparkPulse:
    time_us: int
    column: str


def resolve_input_path(path: Path) -> Path:
    expanded = path.expanduser()
    if expanded.exists():
        return expanded

    for root in DEFAULT_DATA_ROOTS:
        candidate = root / path
        if candidate.exists():
            return candidate

    searched = ", ".join(str(root) for root in DEFAULT_DATA_ROOTS)
    raise FileNotFoundError(f"{path} does not exist; also searched {searched}")


def read_rpm(path: Path) -> pd.DataFrame:
    df = pd.read_csv(path)
    missing = {"time_us", *GP_COLUMNS} - set(df.columns)
    if missing:
        raise ValueError(f"missing columns: {', '.join(sorted(missing))}")

    df["time_us"] = df["time_us"].astype("int64")
    return df


def gp_times(df: pd.DataFrame, column: str) -> list[int]:
    pulse_rows = df.loc[df[column].fillna(0).astype(int) == 1]
    return sorted(int(row.time_us) for row in pulse_rows.itertuples())


def generate_sparks(
    crank_times: list[int],
    cam_times: list[int],
    *,
    spark_advance_deg: float,
) -> list[SparkPulse]:
    events: list[tuple[int, int, str]] = []
    events.extend((time_us, 0, "cam") for time_us in cam_times)
    events.extend((time_us, 1, "crank") for time_us in crank_times)
    events.sort()

    sync_ok = False
    cycle_ctr = 0
    cam_ctr = 0
    exhaust_timestamp = 0
    intake_timestamp = 0
    sparks: list[SparkPulse] = []

    for time_us, _, event_type in events:
        if event_type == "cam":
            if cycle_ctr < 4:
                sync_ok = False
            if not sync_ok and cycle_ctr == 4:
                cam_ctr += 1
            if cam_ctr >= 2:
                sync_ok = True
            cycle_ctr = 0
            continue

        cycle_ctr += 1
        if cycle_ctr >= 5:
            sync_ok = False

        if cycle_ctr == 2:
            exhaust_timestamp = time_us
        elif cycle_ctr == 3:
            intake_timestamp = time_us
        elif cycle_ctr == 4 and sync_ok:
            d1a = time_us - intake_timestamp
            d1b = intake_timestamp - exhaust_timestamp
            predicted_interval = 2 * d1a - d1b
            if d1a <= 0 or predicted_interval <= 0:
                continue

            spark_1_time = int(time_us + d1a * (45.0 - spark_advance_deg) / 180.0)
            spark_2_time = int(time_us + d1b * (45.0 - spark_advance_deg) / 180.0)

            t_s = (
                d1a * (135.0 + spark_advance_deg) / 180.0
                + predicted_interval * (45.0 - spark_advance_deg) / 180.0
            )
            t_a = d1a / 2.0 + t_s / 2.0
            spark_3_time = int(t_a * (45.0 - spark_advance_deg) / 180.0 + time_us)

            sparks.extend(
                (
                    SparkPulse(spark_1_time, "spark_1"),
                    SparkPulse(spark_2_time, "spark_2"),
                    SparkPulse(spark_3_time, "spark_3"),
                )
            )

    return sorted(set(sparks), key=lambda pulse: (pulse.time_us, pulse.column))


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
        angle = revolution_index * 360.0 + float(row.crank_angle_deg)
        points.append(
            EncoderPoint(time_us=int(row.time_us), continuous_angle_deg=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 add_spark_rows(df: pd.DataFrame, sparks: list[SparkPulse]) -> pd.DataFrame:
    for column in SPARK_COLUMNS:
        df[column] = 0
    for column in INTERPOLATED_COLUMNS:
        if column not in df.columns:
            df[column] = pd.NA

    rows_by_time = {
        int(row.time_us): index for index, row in enumerate(df.itertuples())
    }
    new_rows_by_time: dict[int, dict[str, object]] = {}

    for pulse in sparks:
        if pulse.time_us in rows_by_time:
            df.at[rows_by_time[pulse.time_us], pulse.column] = 1
            continue

        if pulse.time_us not in new_rows_by_time:
            row = {column: pd.NA for column in df.columns}
            row["time_us"] = pulse.time_us
            for column in SPARK_COLUMNS:
                row[column] = 0
            new_rows_by_time[pulse.time_us] = row
        new_rows_by_time[pulse.time_us][pulse.column] = 1

    if new_rows_by_time:
        new_rows = list(new_rows_by_time.values())
        df = pd.concat(
            [df, pd.DataFrame(new_rows, columns=df.columns)], ignore_index=True
        )

    return df.sort_values("time_us", kind="stable").reset_index(drop=True)


def add_interpolated_angles(
    df: pd.DataFrame, points: list[EncoderPoint]
) -> pd.DataFrame:
    for index, row in df.iterrows():
        needs_angle = bool(sum(int(row[column] or 0) for column in SPARK_COLUMNS))
        if not needs_angle:
            continue

        continuous_angle = interpolate_angle(points, int(row["time_us"]))
        if continuous_angle is None:
            continue

        revolution_index = int(continuous_angle // 360.0)
        crank_angle = continuous_angle - revolution_index * 360.0
        df.at[index, "algo1_continuous_angle_deg"] = round(continuous_angle, 6)
        df.at[index, "algo1_revolution_index"] = revolution_index
        df.at[index, "algo1_crank_angle_deg"] = round(crank_angle, 6)

    return df


def output_path_for(input_path: Path, output_path: Path | None) -> Path:
    if output_path is not None:
        return output_path.expanduser()

    stem = input_path.stem
    if stem.endswith("_rpm"):
        stem = stem[: -len("_rpm")]
    return input_path.with_name(f"{stem}_algo1.csv")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run algo1 crank/cam angle prediction against one RPM dataset."
    )
    parser.add_argument(
        "file",
        type=Path,
        help="RPM CSV file. Bare names are searched in recorder RPM folders.",
    )
    parser.add_argument("-o", "--output", type=Path, help="Output CSV path.")
    parser.add_argument(
        "--crank-gp",
        choices=GP_COLUMNS,
        default="gp0_falling",
        help="GP falling column to treat as crank pulses.",
    )
    parser.add_argument(
        "--cam-gp",
        choices=GP_COLUMNS,
        default="gp1_falling",
        help="GP falling column to treat as cam pulses.",
    )
    parser.add_argument(
        "--spark-advance",
        type=float,
        default=25.0,
        help="Spark advance in crank degrees, matching the original algo1 constant.",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if args.crank_gp == args.cam_gp:
        print(
            "--crank-gp and --cam-gp must select different GP columns.", file=sys.stderr
        )
        return 1

    try:
        input_path = resolve_input_path(args.file)
        df = read_rpm(input_path)
    except (OSError, ValueError) as exc:
        print(f"Could not read input: {exc}", file=sys.stderr)
        return 1

    crank_times = gp_times(df, args.crank_gp)
    cam_times = gp_times(df, args.cam_gp)
    sparks = generate_sparks(
        crank_times,
        cam_times,
        spark_advance_deg=args.spark_advance,
    )

    output_df = add_spark_rows(df.copy(), sparks)
    output_df = add_interpolated_angles(output_df, encoder_points(df))

    output_path = output_path_for(input_path, args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_df.to_csv(output_path, index=False)

    print(
        f"Wrote {output_path} "
        f"crank={args.crank_gp}({len(crank_times)}) "
        f"cam={args.cam_gp}({len(cam_times)}) "
        f"sparks={len(sparks)}"
    )
    return 0


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