# 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

import pandas as pd
import matplotlib.pyplot as plt
from pathlib import Path
import argparse


parser = argparse.ArgumentParser()
parser.add_argument("directory", type=Path, help="Source data directory")

args = parser.parse_args()

INPUT_CSV: Path = args.directory


# Width of the artificial pulse, in microseconds.
# Increase this if the pulses are still hard to see.
PULSE_WIDTH_US = 500  # 5 ms

# Optional: plot only a smaller time window.
# Use None to plot the whole recording.
START_US = None
END_US = None


def build_pulse_trace(edge_times_us, pulse_width_us):
    """
    Convert falling-edge timestamps into a drawable square pulse trace.

    Each edge becomes:
        low before edge
        high from edge to edge + pulse_width_us
        low after that
    """
    x = []
    y = []

    for t in edge_times_us:
        x.extend([t, t, t + pulse_width_us, t + pulse_width_us])
        y.extend([0, 1, 1, 0])

    return x, y


def main():
    df = pd.read_csv(INPUT_CSV)

    # If the CSV has an index/time column from df_final.to_csv(),
    # use the first column as time in microseconds.
    time_col = df.columns[0]

    # If your CSV columns are: index, crank, cam
    time_us = df[time_col]

    crank_edges = time_us[df["crank"] == 1].to_numpy()
    cam_edges = time_us[df["cam"] == 1].to_numpy()

    if START_US is not None:
        crank_edges = crank_edges[crank_edges >= START_US]
        cam_edges = cam_edges[cam_edges >= START_US]

    if END_US is not None:
        crank_edges = crank_edges[crank_edges <= END_US]
        cam_edges = cam_edges[cam_edges <= END_US]

    crank_x, crank_y = build_pulse_trace(crank_edges, PULSE_WIDTH_US)
    cam_x, cam_y = build_pulse_trace(cam_edges, PULSE_WIDTH_US)

    # Convert microseconds to seconds for a more readable x-axis
    crank_x = [x / 1_000_000 for x in crank_x]
    cam_x = [x / 1_000_000 for x in cam_x]

    # Offset cam vertically so both traces are readable
    crank_y = [y * 0.8 + 0 for y in crank_y]
    cam_y = [y * 0.8 + 1.2 for y in cam_y]

    plt.figure(figsize=(14, 5))

    plt.plot(crank_x, crank_y, label="Crank")
    plt.plot(cam_x, cam_y, label="Cam")

    plt.yticks([0.4, 1.6], ["Crank", "Cam"])
    plt.xlabel("Time [s]")
    plt.title(f"Crank and Cam Falling Edges, pulse width = {PULSE_WIDTH_US} µs")
    plt.grid(True)
    plt.legend()

    if START_US is not None or END_US is not None:
        start_s = START_US / 1_000_000 if START_US is not None else None
        end_s = END_US / 1_000_000 if END_US is not None else None
        plt.xlim(start_s, end_s)

    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()