from __future__ import annotations

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


INPUT_COLUMNS = {"time_us", "turn", "pulse"}
GP_COLUMNS = ("gp0_falling", "gp1_falling")
OUTPUT_COLUMNS = ("time_us", "turn", "pulse", *GP_COLUMNS)
DEFAULT_PPR = 256


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


@dataclass(frozen=True)
class CleanResult:
    turns: list[int]
    pulses: list[int]
    gp0_falling: list[int]
    gp1_falling: list[int]
    pulse_noise_removed: int
    pulses_inserted: int
    turn_noise_removed: int
    turn_phase: int | None


def positive_intervals(times: list[int]) -> list[int]:
    return [b - a for a, b in zip(times, times[1:]) if b > a]


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


def robust_median_interval(times: list[int]) -> float | None:
    intervals = positive_intervals(times)
    if not intervals:
        return None

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

    return float(statistics.median(intervals))


def read_events(path: Path) -> Events:
    turns: list[int] = []
    pulses: 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))}")

        for row in reader:
            time_us = int(row["time_us"])
            if int(row["turn"]):
                turns.append(time_us)
            if int(row["pulse"]):
                pulses.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(
        turns=sorted(turns),
        pulses=sorted(pulses),
        gp0_falling=sorted(gp0_falling),
        gp1_falling=sorted(gp1_falling),
    )


def remove_close_pulses(pulses: list[int]) -> tuple[list[int], int]:
    if len(pulses) < 3:
        return sorted(set(pulses)), 0

    median_interval = robust_median_interval(pulses)
    if median_interval is None:
        return sorted(set(pulses)), 0

    min_interval = max(2, int(median_interval * 0.35))
    cleaned: list[int] = []
    removed = 0

    for pulse in sorted(pulses):
        if cleaned and pulse <= cleaned[-1]:
            removed += 1
            continue
        if cleaned and pulse - cleaned[-1] < min_interval:
            removed += 1
            continue
        cleaned.append(pulse)

    return cleaned, removed


def local_interval(intervals: list[int], index: int, fallback: float) -> float:
    start = max(0, index - 16)
    end = min(len(intervals), index + 17)
    neighbors = intervals[start:index] + intervals[index + 1 : end]
    plausible = [dt for dt in neighbors if 0.25 * fallback <= dt <= 4.0 * fallback]
    return median_or_none(plausible) or fallback


def interpolate_small_gaps(
    pulses: list[int],
    *,
    max_missing_pulses: int,
    gap_tolerance: float,
) -> tuple[list[int], int]:
    if len(pulses) < 3:
        return pulses, 0

    fallback = robust_median_interval(pulses)
    if fallback is None:
        return pulses, 0

    intervals = positive_intervals(pulses)
    result = [pulses[0]]
    inserted = 0

    for index, (start, end) in enumerate(zip(pulses, pulses[1:])):
        gap = end - start
        expected = local_interval(intervals, index, fallback)
        pulse_count = round(gap / expected)

        if 2 <= pulse_count <= max_missing_pulses + 1:
            corrected_interval = gap / pulse_count
            error = abs(corrected_interval - expected) / expected
            if error <= gap_tolerance:
                for step in range(1, pulse_count):
                    result.append(round(start + corrected_interval * step))
                    inserted += 1

        result.append(end)

    return result, inserted


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

    index = bisect_left(times, target)
    candidates = []
    if index < len(times):
        candidates.append(index)
    if index > 0:
        candidates.append(index - 1)

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


def remove_close_turns(turns: list[int], min_separation_us: int) -> tuple[list[int], int]:
    cleaned: list[int] = []
    removed = 0

    for turn in sorted(turns):
        if cleaned and turn - cleaned[-1] < min_separation_us:
            removed += 1
            continue
        cleaned.append(turn)

    return cleaned, removed


def choose_turn_phase(turns: list[int], pulses: list[int], ppr: int) -> int | None:
    if not turns or len(pulses) < ppr:
        return None

    pulse_interval = robust_median_interval(pulses)
    if pulse_interval is None:
        return None

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

    for turn in turns:
        index = nearest_index(pulses, turn)
        if index is None:
            continue
        distance = abs(pulses[index] - turn)
        if distance > max_distance_us:
            continue
        scores[index % ppr] += 1.0 - distance / max_distance_us

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

    return scores.index(best_score)


def synthesize_turns(turns: list[int], pulses: list[int], ppr: int) -> tuple[list[int], int, int | None]:
    if len(pulses) < ppr:
        return turns, 0, None

    pulse_interval = robust_median_interval(pulses) or 0.0
    expected_turn_interval = max(1, int(pulse_interval * ppr))
    debounced_turns, removed = remove_close_turns(
        turns,
        max(1, int(expected_turn_interval * 0.45)),
    )
    phase = choose_turn_phase(debounced_turns, pulses, ppr)
    if phase is None:
        phase = 0

    match_window_us = max(1_000, int(expected_turn_interval * 0.15))
    synthesized: list[int] = []

    for pulse_index in range(phase, len(pulses), ppr):
        pulse_time = pulses[pulse_index]
        turn_index = nearest_index(debounced_turns, pulse_time)
        if turn_index is not None and abs(debounced_turns[turn_index] - pulse_time) <= match_window_us:
            synthesized.append(debounced_turns[turn_index])
        else:
            synthesized.append(pulse_time)

    return synthesized, removed, phase


def clean_events(
    events: Events,
    *,
    ppr: int,
    max_missing_pulses: int,
    gap_tolerance: float,
) -> CleanResult:
    pulses, pulse_noise_removed = remove_close_pulses(events.pulses)
    pulses, pulses_inserted = interpolate_small_gaps(
        pulses,
        max_missing_pulses=max_missing_pulses,
        gap_tolerance=gap_tolerance,
    )
    turns, turn_noise_removed, turn_phase = synthesize_turns(events.turns, pulses, ppr)

    return CleanResult(
        turns=turns,
        pulses=pulses,
        gp0_falling=events.gp0_falling,
        gp1_falling=events.gp1_falling,
        pulse_noise_removed=pulse_noise_removed,
        pulses_inserted=pulses_inserted,
        turn_noise_removed=turn_noise_removed,
        turn_phase=turn_phase,
    )


def write_events(path: Path, result: CleanResult) -> None:
    rows: dict[int, list[int]] = {}
    for turn in result.turns:
        rows.setdefault(turn, [0, 0, 0, 0])[0] = 1
    for pulse in result.pulses:
        rows.setdefault(pulse, [0, 0, 0, 0])[1] = 1
    for gp0 in result.gp0_falling:
        rows.setdefault(gp0, [0, 0, 0, 0])[2] = 1
    for gp1 in result.gp1_falling:
        rows.setdefault(gp1, [0, 0, 0, 0])[3] = 1

    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="") as file:
        writer = csv.writer(file)
        writer.writerow(OUTPUT_COLUMNS)
        for time_us in sorted(rows):
            writer.writerow((time_us, *rows[time_us]))


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


def process_file(
    input_path: Path,
    output_path: Path,
    *,
    ppr: int,
    max_missing_pulses: int,
    gap_tolerance: float,
) -> CleanResult:
    events = read_events(input_path)
    result = clean_events(
        events,
        ppr=ppr,
        max_missing_pulses=max_missing_pulses,
        gap_tolerance=gap_tolerance,
    )
    write_events(output_path, result)
    return result


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Clean rotary turn and pulse event CSVs from recordings/."
    )
    parser.add_argument(
        "--input-root",
        type=Path,
        default=Path("recordings"),
        help="Folder to search recursively for CSV files.",
    )
    parser.add_argument(
        "--output-root",
        type=Path,
        default=Path("recordings_cleaned"),
        help="Folder where cleaned CSV files are written.",
    )
    parser.add_argument("--ppr", type=int, default=DEFAULT_PPR, help="Pulse encoder pulses per turn.")
    parser.add_argument(
        "--max-missing-pulses",
        type=int,
        default=8,
        help="Largest pulse gap to interpolate. Larger gaps are treated as recording breaks.",
    )
    parser.add_argument(
        "--gap-tolerance",
        type=float,
        default=0.45,
        help="Allowed fractional error when deciding whether a gap is missing pulses.",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    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("_cleaned.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:
            result = process_file(
                input_path,
                output_path,
                ppr=args.ppr,
                max_missing_pulses=args.max_missing_pulses,
                gap_tolerance=args.gap_tolerance,
            )
        except (OSError, ValueError) as exc:
            failures += 1
            print(f"Skipping {input_path}: {exc}", file=sys.stderr)
            continue

        phase = "unknown" if result.turn_phase is None else str(result.turn_phase)
        print(
            f"{input_path} -> {output_path} "
            f"pulses={len(result.pulses)} "
            f"turns={len(result.turns)} "
            f"removed_pulses={result.pulse_noise_removed} "
            f"inserted_pulses={result.pulses_inserted} "
            f"removed_turns={result.turn_noise_removed} "
            f"phase={phase}"
        )

    return 1 if failures else 0


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