from __future__ import annotations

import argparse
import csv
import statistics
import sys
from bisect import bisect_right
from dataclasses import dataclass
from pathlib import Path


INPUT_COLUMNS = {"time_us", "pulse"}
GP_COLUMNS = ("gp0_falling", "gp1_falling")
OUTPUT_COLUMNS = (
    "time_us",
    "pulse_index",
    "revolution_index",
    "crank_angle_deg",
    "turn",
    *GP_COLUMNS,
    "rpm_raw",
    "rpm_lpf",
)
DEFAULT_PPR = 256


@dataclass(frozen=True)
class Events:
    pulse_times: list[int]
    turn_times: list[int]
    gp0_falling: list[int]
    gp1_falling: list[int]


@dataclass(frozen=True)
class RpmSample:
    time_us: int
    pulse_index: int
    revolution_index: int
    crank_angle_deg: float
    turn: int
    rpm_raw: float
    rpm_lpf: float


@dataclass(frozen=True)
class OutputRow:
    time_us: int
    pulse_index: int | None
    revolution_index: int | None
    crank_angle_deg: float | None
    turn: int
    gp0_falling: int
    gp1_falling: int
    rpm_raw: float | None
    rpm_lpf: float | None


def read_events(path: Path) -> Events:
    pulse_times: list[int] = []
    turn_times: list[int] = []
    gp0_falling: list[int] = []
    gp1_falling: list[int] = []

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

        has_turn = "turn" in fieldnames
        for row in reader:
            time_us = int(row["time_us"])
            if int(row["pulse"]):
                pulse_times.append(time_us)
            if has_turn and int(row["turn"]):
                turn_times.append(time_us)
            if "gp0_falling" in fieldnames and int(row["gp0_falling"]):
                gp0_falling.append(time_us)
            if "gp1_falling" in fieldnames and int(row["gp1_falling"]):
                gp1_falling.append(time_us)

    return Events(
        pulse_times=sorted(pulse_times),
        turn_times=sorted(turn_times),
        gp0_falling=sorted(gp0_falling),
        gp1_falling=sorted(gp1_falling),
    )


def median_or_none(values: list[float]) -> float | None:
    if not values:
        return None
    return float(statistics.median(values))


def despike(values: list[float], *, window: int, sigma: float) -> tuple[list[float], int]:
    if len(values) < 3:
        return values[:], 0

    radius = max(1, window // 2)
    cleaned = values[:]
    replaced = 0

    for index, value in enumerate(values):
        start = max(0, index - radius)
        end = min(len(values), index + radius + 1)
        local = values[start:end]
        median = median_or_none(local)
        if median is None:
            continue

        deviations = [abs(sample - median) for sample in local]
        mad = median_or_none(deviations) or 0.0
        threshold = sigma * 1.4826 * mad

        if threshold <= 0.0:
            threshold = max(1.0, abs(median) * 0.25)

        if abs(value - median) > threshold:
            cleaned[index] = median
            replaced += 1

    return cleaned, replaced


def ema(values: list[float], alpha: float) -> list[float]:
    if not values:
        return []

    filtered = [values[0]]
    previous = values[0]
    for value in values[1:]:
        previous = alpha * value + (1.0 - alpha) * previous
        filtered.append(previous)

    return filtered


def nearest_index(times: list[int], target: int) -> int | None:
    if not times:
        return None

    lo = 0
    hi = len(times)
    while lo < hi:
        mid = (lo + hi) // 2
        if times[mid] < target:
            lo = mid + 1
        else:
            hi = mid

    candidates = []
    if lo < len(times):
        candidates.append(lo)
    if lo > 0:
        candidates.append(lo - 1)

    return min(candidates, key=lambda index: abs(times[index] - target))


def robust_median_interval(times: list[int]) -> float | None:
    intervals = [b - a for a, b in zip(times, times[1:]) if b > a]
    if not intervals:
        return None

    intervals = sorted(intervals)
    if len(intervals) >= 20:
        trim = len(intervals) // 20
        intervals = intervals[trim : len(intervals) - trim]

    return float(statistics.median(intervals))


def choose_turn_phase(pulse_times: list[int], turn_times: list[int], ppr: int) -> int:
    if not pulse_times or not turn_times:
        return 0

    pulse_interval = robust_median_interval(pulse_times)
    if pulse_interval is None:
        return 0

    max_distance_us = max(1_000, int(pulse_interval * 10))
    scores = [0.0] * ppr

    for turn_time in turn_times:
        pulse_index = nearest_index(pulse_times, turn_time)
        if pulse_index is None:
            continue
        distance = abs(pulse_times[pulse_index] - turn_time)
        if distance <= max_distance_us:
            scores[pulse_index % ppr] += 1.0 - distance / max_distance_us

    best_score = max(scores)
    if best_score <= 0.0:
        return 0

    return scores.index(best_score)


def revolution_index_for(pulse_index: int, phase: int, ppr: int) -> int:
    return (pulse_index - phase) // ppr


def calculate_rpm(
    events: Events,
    *,
    ppr: int,
    filter_window: int,
    filter_alpha: float,
    hampel_sigma: float,
) -> tuple[list[RpmSample], int]:
    raw_rows: list[tuple[int, int, int, float, int, float]] = []
    pulse_times = events.pulse_times
    phase = choose_turn_phase(pulse_times, events.turn_times, ppr)
    degrees_per_pulse = 360.0 / ppr

    for pulse_index, (previous_time, current_time) in enumerate(zip(pulse_times, pulse_times[1:]), start=1):
        delta_us = current_time - previous_time
        if delta_us <= 0:
            continue

        rpm = 60_000_000.0 / (delta_us * ppr)
        angle_pulse = (pulse_index - phase) % ppr
        angle_deg = angle_pulse * degrees_per_pulse
        revolution_index = revolution_index_for(pulse_index, phase, ppr)
        turn = 1 if angle_pulse == 0 else 0
        raw_rows.append((current_time, pulse_index, revolution_index, angle_deg, turn, rpm))

    raw_rpm = [row[5] for row in raw_rows]
    despiked_rpm, replaced = despike(raw_rpm, window=filter_window, sigma=hampel_sigma)
    filtered_rpm = ema(despiked_rpm, filter_alpha)

    samples = [
        RpmSample(
            time_us=time_us,
            pulse_index=pulse_index,
            revolution_index=revolution_index,
            crank_angle_deg=angle_deg,
            turn=turn,
            rpm_raw=rpm_raw,
            rpm_lpf=rpm_lpf,
        )
        for (time_us, pulse_index, revolution_index, angle_deg, turn, rpm_raw), rpm_lpf in zip(
            raw_rows,
            filtered_rpm,
        )
    ]

    return samples, replaced


def pulse_position_for_time(
    pulse_times: list[int],
    time_us: int,
    *,
    phase: int,
    ppr: int,
) -> tuple[int, int, float]:
    pulse_index = bisect_right(pulse_times, time_us) - 1
    if pulse_index < 0:
        pulse_index = 0

    angle_pulse = (pulse_index - phase) % ppr
    revolution_index = revolution_index_for(pulse_index, phase, ppr)
    angle_deg = angle_pulse * (360.0 / ppr)

    return pulse_index, revolution_index, angle_deg


def build_output_rows(
    events: Events,
    samples: list[RpmSample],
    *,
    ppr: int,
) -> list[OutputRow]:
    phase = choose_turn_phase(events.pulse_times, events.turn_times, ppr)
    rows: dict[tuple[int, int], OutputRow] = {}

    for sample in samples:
        rows[(sample.time_us, 0)] = OutputRow(
            time_us=sample.time_us,
            pulse_index=sample.pulse_index,
            revolution_index=sample.revolution_index,
            crank_angle_deg=sample.crank_angle_deg,
            turn=sample.turn,
            gp0_falling=0,
            gp1_falling=0,
            rpm_raw=sample.rpm_raw,
            rpm_lpf=sample.rpm_lpf,
        )

    for channel_index, gp_times in enumerate((events.gp0_falling, events.gp1_falling), start=1):
        for gp_time in gp_times:
            pulse_index, revolution_index, angle_deg = pulse_position_for_time(
                events.pulse_times,
                gp_time,
                phase=phase,
                ppr=ppr,
            )
            existing_key = (gp_time, 0)
            if existing_key in rows:
                existing = rows[existing_key]
                rows[existing_key] = OutputRow(
                    time_us=existing.time_us,
                    pulse_index=existing.pulse_index,
                    revolution_index=existing.revolution_index,
                    crank_angle_deg=existing.crank_angle_deg,
                    turn=existing.turn,
                    gp0_falling=1 if channel_index == 1 else existing.gp0_falling,
                    gp1_falling=1 if channel_index == 2 else existing.gp1_falling,
                    rpm_raw=existing.rpm_raw,
                    rpm_lpf=existing.rpm_lpf,
                )
                continue

            rows[(gp_time, channel_index)] = OutputRow(
                time_us=gp_time,
                pulse_index=pulse_index,
                revolution_index=revolution_index,
                crank_angle_deg=angle_deg,
                turn=0,
                gp0_falling=1 if channel_index == 1 else 0,
                gp1_falling=1 if channel_index == 2 else 0,
                rpm_raw=None,
                rpm_lpf=None,
            )

    return [rows[key] for key in sorted(rows)]


def write_rpm(path: Path, rows: list[OutputRow]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="") as file:
        writer = csv.writer(file)
        writer.writerow(OUTPUT_COLUMNS)
        for row in rows:
            writer.writerow(
                (
                    row.time_us,
                    "" if row.pulse_index is None else row.pulse_index,
                    "" if row.revolution_index is None else row.revolution_index,
                    "" if row.crank_angle_deg is None else f"{row.crank_angle_deg:.6f}",
                    row.turn,
                    row.gp0_falling,
                    row.gp1_falling,
                    "" if row.rpm_raw is None else f"{row.rpm_raw:.3f}",
                    "" if row.rpm_lpf is None else f"{row.rpm_lpf:.3f}",
                )
            )


def output_path_for(input_path: Path, input_root: Path, output_root: Path) -> Path:
    relative = input_path.relative_to(input_root)
    stem = relative.stem
    if stem.endswith("_cleaned"):
        stem = stem[: -len("_cleaned")]
    return output_root / relative.with_name(f"{stem}_rpm.csv")


def process_file(
    input_path: Path,
    output_path: Path,
    *,
    ppr: int,
    filter_window: int,
    filter_alpha: float,
    hampel_sigma: float,
) -> tuple[int, int]:
    events = read_events(input_path)
    samples, replaced = calculate_rpm(
        events,
        ppr=ppr,
        filter_window=filter_window,
        filter_alpha=filter_alpha,
        hampel_sigma=hampel_sigma,
    )
    rows = build_output_rows(events, samples, ppr=ppr)
    write_rpm(output_path, rows)
    return len(samples), replaced


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Convert cleaned pulse event CSVs into raw and filtered RPM CSVs."
    )
    parser.add_argument(
        "--input-root",
        type=Path,
        default=Path("recordings_cleaned"),
        help="Folder to search recursively for cleaned CSV files.",
    )
    parser.add_argument(
        "--output-root",
        type=Path,
        default=Path("recordings_rpm"),
        help="Folder where RPM CSV files are written.",
    )
    parser.add_argument("--ppr", type=int, default=DEFAULT_PPR, help="Pulse encoder pulses per revolution.")
    parser.add_argument(
        "--filter-window",
        type=int,
        default=5,
        help="Odd sample count for local despiking before low-pass filtering.",
    )
    parser.add_argument(
        "--filter-alpha",
        type=float,
        default=0.45,
        help="EMA alpha for the low-pass RPM column. Higher follows transients faster.",
    )
    parser.add_argument(
        "--hampel-sigma",
        type=float,
        default=4.0,
        help="Local median absolute deviation threshold for replacing single-sample outliers.",
    )
    return parser.parse_args()


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

    if not 0.0 < args.filter_alpha <= 1.0:
        print("--filter-alpha must be in the range (0, 1].", file=sys.stderr)
        return 1
    if args.filter_window < 3:
        print("--filter-window must be at least 3.", file=sys.stderr)
        return 1

    input_root = args.input_root
    output_root = args.output_root
    if not input_root.exists():
        print(f"Input folder does not exist: {input_root}", file=sys.stderr)
        return 1

    csv_paths = sorted(path for path in input_root.rglob("*.csv") if not path.name.endswith("_rpm.csv"))
    if not csv_paths:
        print(f"No CSV files found under {input_root}", file=sys.stderr)
        return 1

    failures = 0
    for input_path in csv_paths:
        output_path = output_path_for(input_path, input_root, output_root)
        try:
            samples, replaced = process_file(
                input_path,
                output_path,
                ppr=args.ppr,
                filter_window=args.filter_window,
                filter_alpha=args.filter_alpha,
                hampel_sigma=args.hampel_sigma,
            )
        except (OSError, ValueError) as exc:
            failures += 1
            print(f"Skipping {input_path}: {exc}", file=sys.stderr)
            continue

        print(
            f"{input_path} -> {output_path} "
            f"samples={samples} despiked={replaced}"
        )

    return 1 if failures else 0


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