from __future__ import annotations

import argparse
import csv
import sys
from dataclasses import dataclass
from pathlib import Path

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


REQUIRED_COLUMNS = {"time_us", "rpm_raw", "rpm_lpf"}
ANGLE_COLUMNS = {"revolution_index", "crank_angle_deg"}


@dataclass(frozen=True)
class GpEvent:
    time_s: float
    revolution_index: int
    crank_angle_deg: float


@dataclass(frozen=True)
class RpmData:
    time_s: list[float]
    revolution_index: list[int]
    crank_angle_deg: list[float]
    rpm_raw: list[float]
    rpm_lpf: list[float]
    turn_time_s: list[float]
    gp0_falling: list[GpEvent]
    gp1_falling: list[GpEvent]


def read_rpm(path: Path) -> RpmData:
    time_s: list[float] = []
    revolution_index: list[int] = []
    crank_angle_deg: list[float] = []
    rpm_raw: list[float] = []
    rpm_lpf: list[float] = []
    turn_time_s: list[float] = []
    gp0_falling: list[GpEvent] = []
    gp1_falling: list[GpEvent] = []

    with path.open(newline="") as file:
        reader = csv.DictReader(file)
        fieldnames = set(reader.fieldnames or ())
        missing = REQUIRED_COLUMNS - fieldnames
        if missing:
            raise ValueError(f"missing columns: {', '.join(sorted(missing))}")

        has_turn = "turn" in fieldnames
        has_angle = ANGLE_COLUMNS <= fieldnames
        has_gp0 = "gp0_falling" in fieldnames
        has_gp1 = "gp1_falling" in fieldnames
        for row in reader:
            sample_time_s = int(row["time_us"]) / 1_000_000
            row_revolution_index = int(row["revolution_index"]) if has_angle and row["revolution_index"] else 0
            row_crank_angle_deg = float(row["crank_angle_deg"]) if has_angle and row["crank_angle_deg"] else 0.0

            if row["rpm_raw"] and row["rpm_lpf"]:
                time_s.append(sample_time_s)
                revolution_index.append(row_revolution_index)
                crank_angle_deg.append(row_crank_angle_deg)
                rpm_raw.append(float(row["rpm_raw"]))
                rpm_lpf.append(float(row["rpm_lpf"]))
            if has_turn and int(row["turn"]):
                turn_time_s.append(sample_time_s)
            if has_gp0 and int(row["gp0_falling"]):
                gp0_falling.append(
                    GpEvent(
                        time_s=sample_time_s,
                        revolution_index=row_revolution_index,
                        crank_angle_deg=row_crank_angle_deg,
                    )
                )
            if has_gp1 and int(row["gp1_falling"]):
                gp1_falling.append(
                    GpEvent(
                        time_s=sample_time_s,
                        revolution_index=row_revolution_index,
                        crank_angle_deg=row_crank_angle_deg,
                    )
                )

    return RpmData(
        time_s=time_s,
        revolution_index=revolution_index,
        crank_angle_deg=crank_angle_deg,
        rpm_raw=rpm_raw,
        rpm_lpf=rpm_lpf,
        turn_time_s=turn_time_s,
        gp0_falling=gp0_falling,
        gp1_falling=gp1_falling,
    )


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


def plot_time(data: RpmData, ax: plt.Axes) -> None:
    ax.plot(data.time_s, data.rpm_raw, label="raw", linewidth=0.75, alpha=0.35)
    ax.plot(data.time_s, data.rpm_lpf, label="lpf", linewidth=1.4)

    for turn_time in data.turn_time_s:
        ax.axvline(turn_time, color="0.7", linewidth=0.5, alpha=0.35)

    ax.set_xlabel("time (s)")


def selected_gp_events(data: RpmData, gp_mode: str) -> list[tuple[str, list[GpEvent], str]]:
    selected: list[tuple[str, list[GpEvent], str]] = []
    if gp_mode in {"gp0", "both"}:
        selected.append(("gp0 falling", data.gp0_falling, "tab:green"))
    if gp_mode in {"gp1", "both"}:
        selected.append(("gp1 falling", data.gp1_falling, "tab:red"))
    return selected


def plot_gp_bars(data: RpmData, ax: plt.Axes, *, x_axis: str, gp_mode: str) -> None:
    for label, events, color in selected_gp_events(data, gp_mode):
        plotted_label = False
        for event in events:
            if x_axis == "time":
                x_value = event.time_s
            elif x_axis in {"angle", "angle-overlay"}:
                if event.revolution_index < 0:
                    continue
                x_value = event.crank_angle_deg
            else:
                if event.revolution_index < 0:
                    continue
                x_value = event.revolution_index * 360.0 + event.crank_angle_deg

            ax.axvline(
                x_value,
                color=color,
                alpha=0.18,
                linewidth=5.0,
                label=label if not plotted_label else None,
            )
            plotted_label = True


def has_angle_data(data: RpmData) -> bool:
    return any(data.revolution_index) or any(data.crank_angle_deg)


def plot_angle_overlay(data: RpmData, ax: plt.Axes) -> None:
    revolutions = sorted(set(data.revolution_index))
    raw_label_used = False
    lpf_label_used = False

    for revolution in revolutions:
        indices = [
            index
            for index, value in enumerate(data.revolution_index)
            if value == revolution and value >= 0
        ]
        if len(indices) < 2:
            continue

        angles = [data.crank_angle_deg[index] for index in indices]
        raw = [data.rpm_raw[index] for index in indices]
        lpf = [data.rpm_lpf[index] for index in indices]

        ax.plot(
            angles,
            raw,
            color="0.55",
            linewidth=0.5,
            alpha=0.12,
            label="raw" if not raw_label_used else None,
        )
        ax.plot(
            angles,
            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(data: RpmData, ax: plt.Axes) -> None:
    indices = [
        index
        for index, revolution in enumerate(data.revolution_index)
        if revolution >= 0
    ]
    if not indices:
        return

    continuous_angle = [
        data.revolution_index[index] * 360.0 + data.crank_angle_deg[index]
        for index in indices
    ]
    rpm_raw = [data.rpm_raw[index] for index in indices]
    rpm_lpf = [data.rpm_lpf[index] for index in indices]

    ax.plot(continuous_angle, rpm_raw, label="raw", linewidth=0.75, alpha=0.35)
    ax.plot(continuous_angle, rpm_lpf, label="lpf", linewidth=1.4)

    max_angle = max(continuous_angle)
    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 plot_rpm(csv_path: Path, output: Path | None, *, x_axis: str, gp_mode: str) -> bool:
    try:
        data = read_rpm(csv_path)
    except (OSError, ValueError) as exc:
        print(f"Could not read {csv_path}: {exc}", file=sys.stderr)
        return False

    if not data.time_s:
        print(f"Skipping empty RPM file: {csv_path}", file=sys.stderr)
        return True
    if x_axis.startswith("angle") and not has_angle_data(data):
        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(data, ax)
    elif x_axis == "angle-continuous":
        plot_angle_continuous(data, ax)
    else:
        plot_time(data, ax)
    plot_gp_bars(data, ax, x_axis=x_axis, gp_mode=gp_mode)

    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_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Plot raw and filtered RPM CSV files.")
    parser.add_argument("csv_path", type=Path, help="RPM CSV file or directory of RPM 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="none",
        help="Overlay GP falling events as translucent vertical bars.",
    )
    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

    failures = 0
    for csv_path in csv_paths:
        output = output_path_for(csv_path, args.output, multiple)
        if multiple and output is None:
            print(f"Plotting {csv_path}")
        if not plot_rpm(csv_path, output, x_axis=args.x, gp_mode=args.gp):
            failures += 1

    return 1 if failures else 0


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