"""Trial-session runner for the head-yaw accuracy experiment.

Run it in a second terminal while ``python main.py`` records ``*_face_frames.csv``.
The runner only writes *wall-clock markers* (``data/analysis/accuracy/trials_*.csv``);
``src.accuracy_report`` later joins those markers with the face log by timestamp, so
the live pipeline is not touched.

1. Generate a randomized plan (blocked by distance, shuffled inside each block)::

    python -m src.accuracy_session plan --distances 1,2,3 --repeats 3 --seed 1 --subject S01

2. Start the system (voice prompt off), then run the session::

    python main.py --disable-voice --camera-position-note accuracy
    python -m src.accuracy_session run --plan data/analysis/accuracy/plan.csv

   Per trial: baseline (look straight ahead) -> cue -> turn the head to the target and
   hold -> cue -> return. Prompt keys: Enter=start, r=redo previous, x=discard previous,
   s=skip this condition, q=quit. ``--auto-interval 6`` removes the Enter key between
   trials (single-person operation); a distance change waits for Enter, then optionally
   uses ``--block-interval`` to give the subject time to move before the next trial.
   Re-run with ``--resume --trials <same file>`` to continue an interrupted session.

The sign of ``true_yaw_deg`` is whatever you define on the floor sheet; the report
checks the regression slope and tells you when ``--flip-sign`` is needed.
"""

from __future__ import annotations

import argparse
import csv
import importlib
import logging
import random
import sys
import time
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Callable, Sequence

DEFAULT_DIRECTORY = Path("data/analysis/accuracy")
INTERVAL_AUDIO = {
    "block": Path("assets/audio/block.wav"),
    "auto": Path("assets/audio/auto_interval.wav"),
}
TRIAL_AUDIO = {
    "trial_start": Path("assets/audio/1.wav"),
    "turn": Path("assets/audio/2.wav"),
    "trial_end": Path("assets/audio/3.wav"),
}
# Angle announcement is composed from separate sign and number WAV files.
# Prepare these files under assets/audio/ using exactly these names.
ANGLE_SIGN_AUDIO = {
    "plus": Path("assets/audio/plus.wav"),
    "minus": Path("assets/audio/minus.wav"),
}
ANGLE_NUMBER_AUDIO = {
    0: Path("assets/audio/0.wav"),
    5: Path("assets/audio/5.wav"),
    10: Path("assets/audio/10.wav"),
    15: Path("assets/audio/15.wav"),
    20: Path("assets/audio/20.wav"),
    30: Path("assets/audio/30.wav"),
}
LOGGER = logging.getLogger(__name__)
DEFAULT_ANGLES = (0.0, 5.0, -5.0, 10.0, -10.0, 15.0, -15.0, 20.0, -20.0, 30.0, -30.0)
PLAN_FIELDS = ["seq", "block", "condition_id", "distance_m", "true_yaw_deg", "pitch_deg", "repeat", "subject"]
TRIAL_FIELDS = ["session_id", "seq", "block", "condition_id", "distance_m", "true_yaw_deg", "pitch_deg", "repeat",
                "subject", "status", "t_baseline_start", "t_turn_cue", "t_hold_end", "baseline_sec", "hold_sec"]
FINAL_STATUSES = {"ok", "skipped"}


@dataclass(frozen=True)
class PlanRow:
    seq: int
    block: int
    condition_id: str
    distance_m: float
    true_yaw_deg: float
    pitch_deg: float
    repeat: int
    subject: str = ""


def condition_id(distance: float, yaw: float, pitch: float) -> str:
    return f"d{distance:g}_y{yaw + 0.0:+g}_p{pitch + 0.0:+g}"


def parse_floats(text: str) -> list[float]:
    return [float(part) for part in text.replace(" ", "").split(",") if part]


def build_plan(
    distances: Sequence[float], angles: Sequence[float], repeats: int, *, seed: int | None = None,
    subject: str = "", pitch_levels: Sequence[float] = (), pitch_distances: Sequence[float] = (),
    pitch_angles: Sequence[float] = (0.0, 10.0, -10.0),
) -> list[PlanRow]:
    """One block per distance; trials are shuffled inside a block (moving the subject is slow)."""
    rng = random.Random(seed)
    rows: list[PlanRow] = []
    seq = 0
    for block, distance in enumerate(distances, start=1):
        combos = [(a, 0.0) for a in angles]
        if distance in pitch_distances:
            combos += [(a, p) for p in pitch_levels if p != 0 for a in pitch_angles]
        items = [(a, p, r) for a, p in combos for r in range(1, repeats + 1)]
        rng.shuffle(items)
        for a, p, r in items:
            seq += 1
            rows.append(PlanRow(seq, block, condition_id(distance, a, p), distance, a + 0.0, p + 0.0, r, subject))
    return rows


def write_plan(rows: Sequence[PlanRow], path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8-sig", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=PLAN_FIELDS)
        writer.writeheader()
        for row in rows:
            writer.writerow({field: getattr(row, field) for field in PLAN_FIELDS})


def read_plan(path: Path) -> list[PlanRow]:
    with path.open(encoding="utf-8-sig", newline="") as stream:
        return [
            PlanRow(int(r["seq"]), int(r["block"]), r["condition_id"], float(r["distance_m"]),
                    float(r["true_yaw_deg"]), float(r["pitch_deg"]), int(r["repeat"]), r.get("subject", ""))
            for r in csv.DictReader(stream)
        ]


def load_done(path: Path) -> set[int]:
    """Sequence numbers whose *last* log row is final (ok / skipped)."""
    last: dict[int, str] = {}
    with path.open(encoding="utf-8-sig", newline="") as stream:
        for row in csv.DictReader(stream):
            last[int(row["seq"])] = row["status"]
    return {seq for seq, status in last.items() if status in FINAL_STATUSES}


def _now_iso() -> str:
    return datetime.now().astimezone().isoformat(timespec="milliseconds")


def _beep(kind: str) -> None:
    try:
        import winsound
        winsound.Beep({"start": 880, "turn": 1320, "end": 440}.get(kind, 660), 200)
    except (ImportError, RuntimeError):
        print("\a", end="", flush=True)


def _sign(value: float) -> str:
    return f"{value + 0.0:+g}"


class IntervalAudioPlayer:
    """Minimal pygame-backed audio player for accuracy-session interval cues."""

    def __init__(self, backend: object | None = None) -> None:
        self._backend = backend
        self._sounds: dict[str, object] = {}
        self._initialized = False
        self._closed = False

        if self._backend is None:
            try:
                self._backend = importlib.import_module("pygame")
            except ImportError:
                LOGGER.warning("pygame is not available; interval audio cues are disabled")
                return

        self._load_sounds()

    def _load_sounds(self) -> None:
        if self._backend is None or self._closed:
            return
        try:
            mixer = self._backend.mixer  # type: ignore[attr-defined]
            if not mixer.get_init():
                mixer.init()
            self._initialized = True
            for kind, path in INTERVAL_AUDIO.items():
                if not path.is_file():
                    LOGGER.warning("Interval audio file does not exist: %s", path)
                    continue
                try:
                    self._sounds[kind] = mixer.Sound(str(path))
                except Exception as exc:
                    LOGGER.error("Failed to load interval audio %s: %s", path, exc)
            for kind, path in TRIAL_AUDIO.items():
                if not path.is_file():
                    LOGGER.warning("Trial audio file does not exist: %s", path)
                    continue
                try:
                    self._sounds[kind] = mixer.Sound(str(path))
                except Exception as exc:
                    LOGGER.error("Failed to load trial audio %s: %s", path, exc)
            for sign, path in ANGLE_SIGN_AUDIO.items():
                if not path.is_file():
                    LOGGER.warning("Angle sign audio file does not exist (%s): %s", sign, path)
                    continue
                try:
                    self._sounds[f"angle_sign:{sign}"] = mixer.Sound(str(path))
                except Exception as exc:
                    LOGGER.error("Failed to load angle sign audio (%s) %s: %s", sign, path, exc)
            for number, path in ANGLE_NUMBER_AUDIO.items():
                if not path.is_file():
                    LOGGER.warning("Angle number audio file does not exist (%d): %s", number, path)
                    continue
                try:
                    self._sounds[f"angle_number:{number}"] = mixer.Sound(str(path))
                except Exception as exc:
                    LOGGER.error("Failed to load angle number audio (%d) %s: %s", number, path, exc)
        except Exception as exc:
            LOGGER.error("Failed to initialize pygame mixer for interval audio: %s", exc)
            self._initialized = False
            self._sounds.clear()

    def play(self, kind: str) -> bool:
        if self._closed or kind not in self._sounds:
            return False
        try:
            sound = self._sounds[kind]
            sound.play()  # type: ignore[attr-defined]
            return True
        except Exception as exc:
            LOGGER.error("Failed to play interval audio %s: %s", kind, exc)
            return False

    def play_angle(self, angle: float) -> bool:
        """Announce the target yaw as separate sign + number WAV files."""
        value = float(angle)
        if value == 0:
            return self.play("angle_number:0")

        number = int(abs(value)) if abs(value).is_integer() else None
        if number is None:
            LOGGER.warning("No integer angle audio is defined for %+g deg", value)
            return False

        sign = "plus" if value > 0 else "minus"
        sign_played = self.play(f"angle_sign:{sign}")
        # Wait until the sign audio has finished before starting the number.
        # pygame Sound.play() is asynchronous, so use the sound length when available.
        sign_sound = self._sounds.get(f"angle_sign:{sign}")
        if sign_played and sign_sound is not None:
            try:
                duration_ms = max(0, int(float(sign_sound.get_length()) * 1000))  # type: ignore[attr-defined]
                if duration_ms:
                    time.sleep(duration_ms / 1000.0)
            except Exception:
                pass
        number_played = self.play(f"angle_number:{number}")
        return sign_played and number_played

    def close(self) -> None:
        if self._closed:
            return
        if self._backend is not None and self._initialized:
            try:
                mixer = self._backend.mixer  # type: ignore[attr-defined]
                if mixer.get_init():
                    mixer.quit()
            except Exception as exc:
                LOGGER.warning("Failed to close pygame mixer for interval audio: %s", exc)
        self._closed = True
        self._sounds.clear()


class SessionRunner:
    """Interactive loop; all I/O and time sources are injectable for tests."""

    def __init__(
        self, plan: Sequence[PlanRow], trials_path: Path, session_id: str, *, baseline_sec: float = 3.0,
        hold_sec: float = 5.0, turn_grace_sec: float = 3.0, auto_interval: float = 0.0, block_interval: float = 0.0, input_fn: Callable[[str], str] = input,
        out: Callable[[str], object] = print, sleep: Callable[[float], object] = time.sleep,
        clock: Callable[[], str] = _now_iso, beep: Callable[[str], object] = _beep,
        interval_audio: IntervalAudioPlayer | None = None,
    ) -> None:
        self.plan, self.trials_path, self.session_id = list(plan), trials_path, session_id
        self.baseline_sec, self.hold_sec, self.turn_grace_sec = baseline_sec, hold_sec, turn_grace_sec
        self.auto_interval, self.block_interval = auto_interval, block_interval
        self.input_fn, self.out, self.sleep, self.clock, self.beep = input_fn, out, sleep, clock, beep
        self.interval_audio = interval_audio or IntervalAudioPlayer()

    def run(self, resume: bool = False) -> int:
        path = self.trials_path
        exists = path.is_file() and path.stat().st_size > 0
        if exists and not resume:
            self.out(f"trials file already exists: {path} (use --resume or another --trials)")
            self.interval_audio.close()
            return 2
        pending = [row for row in self.plan if row.seq not in (load_done(path) if exists else set())]
        if not pending:
            self.out("nothing to do: every planned condition is already finished")
            self.interval_audio.close()
            return 0
        path.parent.mkdir(parents=True, exist_ok=True)
        estimate = len(pending) * (self.baseline_sec + self.hold_sec + 4.0) / 60.0
        self.out(f"{len(pending)} trials pending (rough duration {estimate:.0f} min). session={self.session_id}")
        try:
            with path.open("a", encoding="utf-8-sig", newline="") as stream:
                writer = csv.DictWriter(stream, fieldnames=TRIAL_FIELDS)
                if not exists:
                    writer.writeheader()
                return self._loop(pending, writer, stream)
        finally:
            self.interval_audio.close()

    def _record(self, row: PlanRow, status: str, t0: str = "", t1: str = "", t2: str = "") -> dict[str, object]:
        return {"session_id": self.session_id, "seq": row.seq, "block": row.block, "condition_id": row.condition_id,
                "distance_m": row.distance_m, "true_yaw_deg": row.true_yaw_deg, "pitch_deg": row.pitch_deg,
                "repeat": row.repeat, "subject": row.subject, "status": status, "t_baseline_start": t0,
                "t_turn_cue": t1, "t_hold_end": t2, "baseline_sec": self.baseline_sec, "hold_sec": self.hold_sec}

    def _loop(self, pending: list[PlanRow], writer: csv.DictWriter, stream) -> int:  # type: ignore[type-arg]
        def log(row: PlanRow, status: str, *stamps: str) -> None:
            writer.writerow(self._record(row, status, *stamps))
            stream.flush()

        index, last_index, current_block, finished = 0, None, None, 0
        first_trial_move_pending = True
        while index < len(pending):
            row = pending[index]
            new_block = row.block != current_block
            block_change = current_block is not None and new_block
            if new_block:
                current_block = row.block
                self.out(f"\n=== block {row.block}: set the subject at {row.distance_m:g} m, then continue ===")
            self.out(f"[{index + 1}/{len(pending)}] seq {row.seq}  distance {row.distance_m:g} m  "
                     f"yaw {_sign(row.true_yaw_deg)} deg  pitch {_sign(row.pitch_deg)} deg  (repeat {row.repeat})")
            try:
                if self.auto_interval > 0 and not new_block:
                    self.interval_audio.play("auto")
                    self.interval_audio.play_angle(row.true_yaw_deg)
                    self._countdown(self.auto_interval)
                    command = ""
                else:
                    command = self.input_fn("Enter=start r=redo-prev x=discard-prev s=skip q=quit > ").strip().lower()
            except (KeyboardInterrupt, EOFError):
                command = "q"
            if command == "q":
                break
            if command in {"r", "x"}:
                if last_index is None:
                    self.out("no previous trial in this run")
                    continue
                log(pending[last_index], "redo" if command == "r" else "discarded")
                finished -= 1
                if command == "r":
                    index, last_index, current_block = last_index, None, None
                else:
                    last_index = None
                continue
            if command == "s":
                log(row, "skipped")
                index, last_index = index + 1, None
                continue
            if first_trial_move_pending:
                self.interval_audio.play("block")
                self._countdown(10.0, label="initial move")
                first_trial_move_pending = False
            elif block_change and self.block_interval > 0:
                self.interval_audio.play("block")
                self.interval_audio.play_angle(row.true_yaw_deg)
                self._countdown(self.block_interval, label="moving to next block")
            try:
                stamps = self._execute(row)
            except KeyboardInterrupt:
                log(row, "aborted")
                self.out("aborted; the interrupted trial is marked and can be re-run with --resume")
                return 130
            log(row, "ok", *stamps)
            finished += 1
            last_index, index = index, index + 1
        self.out(f"session ended: {finished} trials recorded in this run -> {self.trials_path}")
        return 0

    def _countdown(self, seconds: float, *, label: str = "starting") -> None:
        remaining = seconds
        while remaining > 0:
            self.out(f"  {label}: {remaining:.0f} s ...")
            self.sleep(min(1.0, remaining))
            remaining -= 1.0

    def _execute(self, row: PlanRow) -> tuple[str, str, str]:
        t_base = self.clock()
        self.interval_audio.play("trial_start")
        self.out(f"  baseline {self.baseline_sec:g} s: face the camera and stay still")
        self.sleep(self.baseline_sec)
        t_cue = self.clock()
        self.interval_audio.play("turn")
        self.out(f"  TURN to {_sign(row.true_yaw_deg)} deg (head only) — you have {self.turn_grace_sec:g} s to turn")
        self.sleep(self.turn_grace_sec)
        self.out(f"  HOLD {self.hold_sec:g} s")
        self.sleep(self.hold_sec)
        t_end = self.clock()
        #self.interval_audio.play("trial_end")
        self.out("  return to the front")
        return t_base, t_cue, t_end


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Accuracy experiment: plan generation and trial session")
    sub = parser.add_subparsers(dest="command", required=True)
    plan = sub.add_parser("plan", help="write a randomized condition plan CSV")
    plan.add_argument("--distances", default="1,2,3", help="metres, comma separated (one block each)")
    plan.add_argument("--angles", default=",".join(f"{a:g}" for a in DEFAULT_ANGLES), help="true yaw angles in degrees")
    plan.add_argument("--repeats", type=int, default=3)
    plan.add_argument("--seed", type=int, default=None)
    plan.add_argument("--subject", default="")
    plan.add_argument("--pitch-levels", default="", help="extra pitch conditions, e.g. -10,10 (optional)")
    plan.add_argument("--pitch-distances", default="2", help="distances that also get the pitch conditions")
    plan.add_argument("--pitch-angles", default="0,10,-10", help="yaw angles used with the extra pitch levels")
    plan.add_argument("--out", type=Path, default=DEFAULT_DIRECTORY / "plan.csv")
    run = sub.add_parser("run", help="run the interactive trial session")
    run.add_argument("--plan", type=Path, default=DEFAULT_DIRECTORY / "plan.csv")
    run.add_argument("--trials", type=Path, default=None, help="output markers CSV (default: trials_<session>.csv)")
    run.add_argument("--session-id", default=None)
    run.add_argument("--baseline-sec", type=float, default=3.0)
    run.add_argument("--hold-sec", type=float, default=5.0)
    run.add_argument("--turn-grace-sec", type=float, default=3.0, help="seconds allowed to complete the head turn after the turn cue")
    run.add_argument("--auto-interval", type=float, default=0.0, help="seconds between trials within a block; 0 = wait for Enter")
    run.add_argument("--block-interval", type=float, default=0.0, help="seconds to allow moving after Enter at a block change; 0 = no extra wait")
    run.add_argument("--resume", action="store_true")
    return parser


def main(argv: Sequence[str] | None = None, out: Callable[[str], object] = print) -> int:
    args = _build_parser().parse_args(argv)
    if args.command == "plan":
        rows = build_plan(parse_floats(args.distances), parse_floats(args.angles), args.repeats, seed=args.seed,
                          subject=args.subject, pitch_levels=parse_floats(args.pitch_levels),
                          pitch_distances=parse_floats(args.pitch_distances),
                          pitch_angles=parse_floats(args.pitch_angles))
        write_plan(rows, args.out)
        out(f"{len(rows)} trials, {len({r.block for r in rows})} blocks -> {args.out}")
        return 0
    if not args.plan.is_file():
        print(f"plan not found: {args.plan} (run the 'plan' command first)", file=sys.stderr)
        return 2
    session_id = args.session_id or datetime.now().strftime("acc_%Y%m%d_%H%M%S")
    trials = args.trials or DEFAULT_DIRECTORY / f"trials_{session_id}.csv"
    if args.resume and args.trials is None:
        print("--resume needs --trials <existing file>", file=sys.stderr)
        return 2
    runner = SessionRunner(read_plan(args.plan), trials, session_id, baseline_sec=args.baseline_sec,
                           hold_sec=args.hold_sec, turn_grace_sec=args.turn_grace_sec, auto_interval=args.auto_interval,
                           block_interval=args.block_interval)
    return runner.run(resume=args.resume)


if __name__ == "__main__":
    raise SystemExit(main())