"""Guided screen-based gaze exercises for Pupil Capture."""

from __future__ import annotations

import argparse
import csv
from datetime import datetime
import json
import math
from pathlib import Path
import queue
import random
import statistics
import threading
import time
import tkinter as tk
from collections import deque

from monitor_utils import format_monitor_list
from show_apriltag_surface import SurfaceTagWindow, enable_windows_dpi_awareness, validate_args
from visualize_surface_gaze import GazeSample, SurfaceReceiver


DEFAULT_VIEWING_DISTANCE_CM = 60.0
TARGET_RADIUS = 80
HIT_RADIUS = 140
SAFE_MARGIN_X = 260
SAFE_MARGIN_Y = 180
DISPERSION_WINDOW = 5
MAX_CONTINUOUS_GAP_S = 0.5


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run four guided gaze exercises.")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=50020)
    parser.add_argument("--surface", default="screen")
    parser.add_argument("--confidence", type=float, default=0.9)
    parser.add_argument("--screen-diagonal", type=float, default=None,
                        help="physical screen diagonal in inches "
                             "(default: auto-detected from the selected monitor)")
    parser.add_argument("--viewing-distance", type=float, default=DEFAULT_VIEWING_DISTANCE_CM,
                        help="eye-to-screen distance in cm (default: 60)")
    parser.add_argument("--subject-id", default="anon",
                        help="identifier for the person being recorded (default: anon)")
    parser.add_argument("--session-tag", default="",
                        help="free-text label for this setup, e.g. laptop_50cm (default: none)")
    parser.add_argument("--rounds", type=int, default=3,
                        help="number of complete 1-2-3-4 exercise rounds (default: 3)")
    parser.add_argument("--animation-fps", type=int, default=120,
                        help="target animation refresh rate (default: 120)")
    parser.add_argument("--cursor", action="store_true",
                        help="show the live gaze cursor (hidden by default)")
    parser.add_argument("--tag-ids", nargs=4, type=int, default=(0, 1, 2, 3))
    parser.add_argument("--tag-size", type=int, default=200)
    parser.add_argument("--margin", type=int, default=0)
    parser.add_argument("--background", default="#10141c")
    parser.add_argument("--windowed", action="store_true")
    parser.add_argument("--monitor", type=int, default=None,
                        help="Windows display number to use, e.g. 2 for \\\\.\\DISPLAY2 "
                             "(default: primary)")
    parser.add_argument("--list-monitors", action="store_true",
                        help="print detected monitors (id, resolution, position, diagonal) and exit")
    return parser.parse_args()


class ExerciseWindow(SurfaceTagWindow):
    BLOCKS = ("fixation", "saccade", "colors", "pursuit")

    def __init__(self, root: tk.Tk, args: argparse.Namespace) -> None:
        super().__init__(root, args)
        self.args = args
        self.width = self.monitor.width
        self.height = self.monitor.height
        self.safe_left, self.safe_right = SAFE_MARGIN_X, self.width - SAFE_MARGIN_X
        self.safe_top, self.safe_bottom = SAFE_MARGIN_Y, self.height - SAFE_MARGIN_Y
        self.pursuit_amp_x = min(760, self.width / 2 - SAFE_MARGIN_X)
        self.pursuit_amp_y = min(390, self.height / 2 - SAFE_MARGIN_Y)
        if args.screen_diagonal is not None:
            self.screen_diagonal = args.screen_diagonal
        elif self.monitor.diagonal_in is not None:
            self.screen_diagonal = self.monitor.diagonal_in
        else:
            raise SystemExit(
                "error: could not auto-detect the diagonal for this monitor; "
                "pass --screen-diagonal explicitly"
            )
        self.frame_interval_ms = max(1, round(1000 / args.animation_fps))
        self.messages: queue.Queue = queue.Queue(maxsize=256)
        self.stop_event = threading.Event()
        self.receiver = SurfaceReceiver(
            self.messages, self.stop_event, args.host, args.port,
            args.surface, args.confidence
        )
        self.receiver.start()
        self.gaze: tuple[float, float] | None = None
        self.rng = random.Random()

        self.block_index = -1
        self.round = 1
        self.block = "welcome"
        self.trial = 0
        self.trial_started = 0.0
        self.block_started = 0.0
        self.target = (self.width / 2, self.height / 2)
        self.distractor: tuple[float, float] | None = None
        self.correct_color = "#39e681"
        self.hit_started: float | None = None
        self.acquired = False
        self.errors: list[float] = []
        self.hits = 0
        self.samples = 0
        self.summary_lines: list[str] = []
        diagonal_px = math.hypot(self.width, self.height)
        self.cm_per_px = self.screen_diagonal * 2.54 / diagonal_px

        self.history: deque[tuple[float, float, float]] = deque(maxlen=DISPERSION_WINDOW)
        self.prev_sample: tuple[float, float, float] | None = None
        self.prev_velocity_px_s: float | None = None

        results_dir = Path(__file__).resolve().parent / "results"
        results_dir.mkdir(exist_ok=True)
        stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.session_id = f"{args.subject_id}_{stamp}"
        self.result_path = results_dir / f"gaze_exercises_{stamp}.csv"
        self.result_file = self.result_path.open("w", newline="", encoding="utf-8")
        self.writer = csv.writer(self.result_file)
        self.writer.writerow([
            "subject_id", "session_id", "pupil_timestamp", "round", "block", "trial",
            "trial_elapsed_s", "target_x", "target_y", "gaze_x", "gaze_y", "confidence",
            "error_px", "error_cm", "error_deg", "hit", "velocity_px_s", "velocity_deg_s",
            "accel_px_s2", "dispersion_px", "event_label"
        ])

        meta_path = self.result_path.with_name(self.result_path.stem + "_meta.json")
        meta_path.write_text(json.dumps({
            "subject_id": args.subject_id,
            "session_id": self.session_id,
            "session_tag": args.session_tag,
            "recorded_at": datetime.now().isoformat(timespec="seconds"),
            "monitor_id": self.monitor.display_id,
            "screen_diagonal_in": self.screen_diagonal,
            "viewing_distance_cm": args.viewing_distance,
            "confidence_threshold": args.confidence,
            "rounds": args.rounds,
            "animation_fps": args.animation_fps,
            "screen_width_px": self.width,
            "screen_height_px": self.height,
            "result_csv": self.result_path.name,
        }, indent=2), encoding="utf-8")

        self.target_item = self.canvas.create_oval(0, 0, 0, 0, state="hidden", tags="exercise")
        self.target_center = self.canvas.create_oval(0, 0, 0, 0, state="hidden", tags="exercise")
        self.distractor_item = self.canvas.create_oval(0, 0, 0, 0, state="hidden", tags="exercise")
        self.gaze_item = self.canvas.create_oval(
            0, 0, 0, 0, fill="#ffd43b", outline="white", width=2,
            state="hidden", tags="exercise"
        )
        self.heading = self.canvas.create_text(
            0, 70, fill="white", font=("Segoe UI", 25, "bold"), tags="exercise"
        )
        self.instructions = self.canvas.create_text(
            0, 125, fill="#d7deea", font=("Segoe UI", 16),
            justify="center", tags="exercise"
        )
        self.network_status = self.canvas.create_text(
            0, 20, fill="#ffca58", font=("Segoe UI", 11), tags="exercise"
        )
        root.bind("<space>", self.on_space)
        root.protocol("WM_DELETE_WINDOW", self.close)
        self.show_welcome()
        root.after(self.frame_interval_ms, self.tick)

    def sx(self, x: float) -> float:
        return x / self.width * self.canvas.winfo_width()

    def sy(self, y: float) -> float:
        return y / self.height * self.canvas.winfo_height()

    def draw(self, event: tk.Event) -> None:
        super().draw(event)
        self.canvas.coords(self.heading, event.width / 2, 65)
        self.canvas.coords(self.instructions, event.width / 2, 115)
        self.canvas.coords(self.network_status, event.width / 2, 18)
        self.position_shapes()
        self.canvas.tag_raise("exercise")

    def close(self, _event: tk.Event | None = None) -> None:
        self.stop_event.set()
        if not self.result_file.closed:
            self.result_file.close()
        super().close(_event)

    def show_welcome(self) -> None:
        self.block = "welcome"
        self.hide_targets()
        self.canvas.itemconfigure(self.heading, text="Gaze exercise sequence")
        self.canvas.itemconfigure(
            self.instructions,
            text=(f"The following sequence repeats {self.args.rounds} times in 1-2-3-4 order:\n\n"
                  "1. Hold gaze on 10 targets for 3 seconds each\n"
                  "2. Move quickly to 10 appearing targets\n"
                  "3. Choose the instructed color in 20 trials\n"
                  "4. Follow a moving target for 10 seconds × 3\n\n"
                  "Sit comfortably, keep your head natural, then press SPACE."),
        )

    def on_space(self, _event: tk.Event | None = None) -> None:
        if self.block in {"welcome", "between"}:
            self.start_next_block()

    def start_next_block(self) -> None:
        self.block_index += 1
        if self.block_index >= len(self.BLOCKS):
            if self.round >= self.args.rounds:
                self.finish()
                return
            self.round += 1
            self.block_index = 0
        self.block = self.BLOCKS[self.block_index]
        self.trial = 0
        self.errors.clear()
        self.hits = self.samples = 0
        self.block_started = self.trial_started = time.monotonic()
        self.history.clear()
        self.prev_sample = None
        self.prev_velocity_px_s = None
        self.new_trial()

    def random_target(self) -> tuple[float, float]:
        return self.rng.uniform(self.safe_left, self.safe_right), self.rng.uniform(self.safe_top, self.safe_bottom)

    def new_trial(self) -> None:
        self.trial_started = time.monotonic()
        self.hit_started = None
        self.acquired = False
        self.distractor = None
        if self.block == "fixation":
            self.target = self.random_target()
            self.set_text(f"Round {self.round}/{self.args.rounds} — Steady fixation — {self.trial + 1}/10", "Keep looking at the center dot")
        elif self.block == "saccade":
            self.target = self.random_target()
            self.set_text(f"Round {self.round}/{self.args.rounds} — Quick movement — {self.trial + 1}/10", "Look at the new target as quickly as possible")
        elif self.block == "colors":
            self.target = self.random_target()
            while True:
                self.distractor = self.random_target()
                if math.dist(self.target, self.distractor) > 500:
                    break
            self.correct_color = "#39e681"
            self.set_text(
                f"Round {self.round}/{self.args.rounds} — Always look at GREEN",
                f"Ignore red  •  Trial {self.trial + 1}/20",
            )
        elif self.block == "pursuit":
            self.set_text(f"Round {self.round}/{self.args.rounds} — Smooth pursuit — {self.trial + 1}/3", "Follow the moving target with your eyes")
        self.position_shapes()

    def set_text(self, heading: str, instructions: str) -> None:
        self.canvas.itemconfigure(self.heading, text=heading)
        self.canvas.itemconfigure(self.instructions, text=instructions)

    def hide_targets(self) -> None:
        for item in (self.target_item, self.target_center, self.distractor_item):
            self.canvas.itemconfigure(item, state="hidden")

    def position_circle(self, item: int, point: tuple[float, float], radius: int) -> None:
        x, y = self.sx(point[0]), self.sy(point[1])
        rx = radius / self.width * self.canvas.winfo_width()
        ry = radius / self.height * self.canvas.winfo_height()
        self.canvas.coords(item, x - rx, y - ry, x + rx, y + ry)

    def position_shapes(self) -> None:
        if self.block not in self.BLOCKS:
            return
        color = self.correct_color if self.block == "colors" else "#55aaff"
        self.position_circle(self.target_item, self.target, TARGET_RADIUS)
        self.position_circle(self.target_center, self.target, 7)
        self.canvas.itemconfigure(self.target_item, fill=color, outline="white", width=3, state="normal")
        self.canvas.itemconfigure(self.target_center, fill="white", outline="", state="normal")
        if self.block == "colors" and self.distractor is not None:
            other = "#ff5263" if self.correct_color == "#39e681" else "#39e681"
            self.position_circle(self.distractor_item, self.distractor, TARGET_RADIUS)
            self.canvas.itemconfigure(self.distractor_item, fill=other, outline="white", width=3, state="normal")
        else:
            self.canvas.itemconfigure(self.distractor_item, state="hidden")

    def process_network(self) -> list[tuple[GazeSample, tuple[float, float]]]:
        samples: list[tuple[GazeSample, tuple[float, float]]] = []
        try:
            while True:
                kind, value = self.messages.get_nowait()
                if kind == "gaze":
                    sample: GazeSample = value  # type: ignore[assignment]
                    gaze = (
                        sample.x_norm * self.width,
                        (1.0 - sample.y_norm) * self.height,
                    )
                    samples.append((sample, gaze))
                else:
                    self.canvas.itemconfigure(self.network_status, text=str(value))
        except queue.Empty:
            pass
        if samples:
            self.gaze = samples[-1][1]
            self.position_circle(self.gaze_item, self.gaze, 14)
            if self.args.cursor:
                self.canvas.itemconfigure(self.gaze_item, state="normal")
        return samples

    def record_sample(
        self,
        sample: GazeSample,
        gaze: tuple[float, float],
        now: float,
    ) -> None:
        error = math.dist(gaze, self.target)
        error_cm = error * self.cm_per_px
        error_deg = math.degrees(2.0 * math.atan2(error_cm / 2.0, self.args.viewing_distance))
        hit = error <= HIT_RADIUS
        self.errors.append(error)
        self.samples += 1
        self.hits += int(hit)

        if hit and self.block != "pursuit":
            self.acquired = True

        velocity_px_s = accel_px_s2 = None
        if self.prev_sample is not None:
            prev_t, prev_x, prev_y = self.prev_sample
            dt = sample.timestamp - prev_t
            if 0 < dt <= MAX_CONTINUOUS_GAP_S:
                dist_px = math.dist((prev_x, prev_y), gaze)
                velocity_px_s = dist_px / dt
                if self.prev_velocity_px_s is not None:
                    accel_px_s2 = (velocity_px_s - self.prev_velocity_px_s) / dt
        self.prev_sample = (sample.timestamp, gaze[0], gaze[1])
        self.prev_velocity_px_s = velocity_px_s

        velocity_deg_s = None
        if velocity_px_s is not None:
            velocity_cm_s = velocity_px_s * self.cm_per_px
            velocity_deg_s = math.degrees(
                2.0 * math.atan2(velocity_cm_s / 2.0, self.args.viewing_distance)
            )

        self.history.append(gaze + (sample.timestamp,))
        dispersion_px = None
        if len(self.history) >= 2:
            xs = [p[0] for p in self.history]
            ys = [p[1] for p in self.history]
            dispersion_px = (max(xs) - min(xs)) + (max(ys) - min(ys))

        if self.block == "pursuit":
            event_label = "pursuit"
        else:
            event_label = "fixation" if self.acquired else "saccade"

        self.writer.writerow([
            self.args.subject_id, self.session_id, f"{sample.timestamp:.6f}", self.round,
            self.block, self.trial + 1, f"{now - self.trial_started:.3f}",
            f"{self.target[0]:.1f}", f"{self.target[1]:.1f}", f"{gaze[0]:.1f}", f"{gaze[1]:.1f}",
            f"{sample.confidence:.3f}", f"{error:.1f}", f"{error_cm:.3f}", f"{error_deg:.3f}",
            int(hit),
            "" if velocity_px_s is None else f"{velocity_px_s:.1f}",
            "" if velocity_deg_s is None else f"{velocity_deg_s:.3f}",
            "" if accel_px_s2 is None else f"{accel_px_s2:.1f}",
            "" if dispersion_px is None else f"{dispersion_px:.1f}",
            event_label,
        ])

    def tick(self) -> None:
        gaze_samples = self.process_network()
        now = time.monotonic()
        if self.block in self.BLOCKS:
            elapsed = now - self.trial_started
            if self.block == "pursuit":
                phase = elapsed / 10.0 * math.tau
                self.target = (
                    self.width / 2 + self.pursuit_amp_x * math.sin(phase),
                    self.height / 2 + self.pursuit_amp_y * math.sin(phase * 2),
                )
                self.position_shapes()
            for sample, gaze in gaze_samples:
                self.record_sample(sample, gaze, now)

            if self.block == "fixation" and elapsed >= 3.0:
                self.advance_trial(10)
            elif self.block == "saccade":
                on_target = self.gaze is not None and math.dist(self.gaze, self.target) <= HIT_RADIUS
                if on_target:
                    self.hit_started = self.hit_started or now
                else:
                    self.hit_started = None
                if (self.hit_started and now - self.hit_started >= 0.15) or elapsed >= 3.0:
                    self.advance_trial(10)
            elif self.block == "colors" and elapsed >= 1.75:
                self.advance_trial(20)
            elif self.block == "pursuit" and elapsed >= 10.0:
                self.advance_trial(3)

        if self.root.winfo_exists():
            self.root.after(self.frame_interval_ms, self.tick)

    def advance_trial(self, total: int) -> None:
        self.trial += 1
        if self.trial >= total:
            self.end_block()
        else:
            self.new_trial()

    def end_block(self) -> None:
        mean_error = statistics.fmean(self.errors) if self.errors else float("nan")
        hit_rate = self.hits / self.samples * 100 if self.samples else 0.0
        mean_cm = mean_error * self.cm_per_px
        mean_deg = math.degrees(2.0 * math.atan2(mean_cm / 2.0, self.args.viewing_distance))
        summary = (f"round {self.round} {self.block}: mean error {mean_error:.0f} px / {mean_cm:.2f} cm / "
                   f"{mean_deg:.2f} deg, on-target {hit_rate:.0f}%")
        self.summary_lines.append(summary)
        self.block = "between"
        self.hide_targets()
        self.result_file.flush()
        self.set_text("Block complete", f"{summary}\n\nPress SPACE for the next exercise")

    def finish(self) -> None:
        self.block = "finished"
        self.hide_targets()
        self.result_file.flush()
        self.set_text(
            "All exercises complete",
            "\n".join(self.summary_lines) + f"\n\nResults saved to:\n{self.result_path}\n\nPress Esc to close",
        )


def main() -> None:
    args = parse_args()
    if args.list_monitors:
        print(format_monitor_list())
        return
    validate_args(args)
    if not 0 <= args.confidence <= 1:
        raise SystemExit("error: --confidence must be between 0 and 1")
    if args.screen_diagonal is not None and args.screen_diagonal <= 0:
        raise SystemExit("error: --screen-diagonal must be greater than 0")
    if args.viewing_distance <= 0:
        raise SystemExit("error: --viewing-distance must be greater than 0")
    if args.rounds < 1:
        raise SystemExit("error: --rounds must be at least 1")
    if args.animation_fps < 1:
        raise SystemExit("error: --animation-fps must be at least 1")
    enable_windows_dpi_awareness()
    root = tk.Tk()
    ExerciseWindow(root, args)
    root.mainloop()


if __name__ == "__main__":
    main()
