#!/usr/bin/env python3
"""Audit EasyCom EOT and wearer interruption/backchannel candidates.

This is a read-only, TurnBench-inspired candidate audit.  EasyCom does not
ship human-consensus end-of-turn (EOT), interruption, or backchannel labels,
so every automatically derived category is explicitly named a proxy.  The
script reuses the participant-resolved transcript timeline and the conservative
interaction rules in :mod:`derive_easycom_interaction_events`.

The primary output is a count table.  With ``--output-dir`` the script also
writes candidate JSONL/CSV files suitable for manual adjudication.  It never
modifies the source dataset.
"""

from __future__ import annotations

import argparse
import csv
import json
from collections import Counter
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Sequence

from derive_easycom_interaction_events import (
    Event,
    Rules,
    Utterance,
    build_clips,
    classify_utterance,
    load_utterances,
)
from easycom_manifest import ALL_SESSIONS, TEST_SESSIONS


FLOOR_CLAIM_EVENT_TYPES = {
    "ordinary_switch",
    "overlap_takeover",
    "interruption_proxy_high",
}
NON_FLOOR_EVENT_TYPES = {
    "backchannel_proxy_high",
    "backchannel_candidate",
    "overlap_hold",
    "nonverbal_feedback",
    "nonverbal_candidate",
}
AMBIGUOUS_EVENT_TYPES = {"interruption_candidate"}

INT_PROXY_GROUPS = {
    "takeover_proxy": {"overlap_takeover", "interruption_proxy_high"},
    "backchannel_or_hold_proxy": {
        "backchannel_proxy_high",
        "overlap_hold",
        "nonverbal_feedback",
    },
    "ambiguous_overlap_or_feedback": {
        "interruption_candidate",
        "backchannel_candidate",
        "nonverbal_candidate",
    },
}


@dataclass(frozen=True)
class EOTCandidate:
    session: str
    split: str
    category: str
    confidence: str
    holder_uid: str
    holder_participant_id: int
    holder_role: str
    boundary_seconds: float
    pause_seconds: float | None
    incoming_uid: str | None
    incoming_participant_id: int | None
    incoming_role: str | None
    incoming_start_seconds: float | None
    incoming_event_type: str | None
    reason: str


def _event_map(
    utterances: Sequence[Utterance],
    *,
    camera_wearer_id: int,
    rules: Rules,
) -> dict[str, Event]:
    return {
        row.uid: classify_utterance(
            row,
            utterances,
            camera_wearer_id=camera_wearer_id,
            rules=rules,
        )
        for row in utterances
    }


def _next_same_speaker(
    holder: Utterance,
    utterances: Sequence[Utterance],
    *,
    epsilon: float,
) -> Utterance | None:
    candidates = [
        row
        for row in utterances
        if row.uid != holder.uid
        and row.segment_id == holder.segment_id
        and row.participant_id == holder.participant_id
        and row.dense_start_seconds >= holder.dense_end_seconds - epsilon
    ]
    return min(candidates, key=lambda row: (row.dense_start_seconds, row.uid), default=None)


def _same_speaker_continues_across_boundary(
    holder: Utterance,
    utterances: Sequence[Utterance],
    *,
    epsilon: float,
) -> bool:
    boundary = holder.dense_end_seconds
    return any(
        row.uid != holder.uid
        and row.segment_id == holder.segment_id
        and row.participant_id == holder.participant_id
        and row.dense_start_seconds < boundary - epsilon
        and row.dense_end_seconds > boundary + epsilon
        for row in utterances
    )


def _incoming_before_resume(
    holder: Utterance,
    utterances: Sequence[Utterance],
    *,
    resume: Utterance | None,
    epsilon: float,
) -> list[Utterance]:
    boundary = holder.dense_end_seconds
    resume_start = float("inf") if resume is None else resume.dense_start_seconds
    candidates = []
    for row in utterances:
        if row.segment_id != holder.segment_id or row.participant_id == holder.participant_id:
            continue
        active_across_boundary = (
            row.dense_start_seconds < boundary - epsilon
            and row.dense_end_seconds > boundary + epsilon
        )
        starts_after_boundary = row.dense_start_seconds >= boundary - epsilon
        if not (active_across_boundary or starts_after_boundary):
            continue
        if row.dense_start_seconds > resume_start + epsilon:
            continue
        candidates.append(row)
    return sorted(candidates, key=lambda row: (max(boundary, row.dense_start_seconds), row.uid))


def derive_eot_candidates(
    utterances: Sequence[Utterance],
    *,
    camera_wearer_id: int,
    rules: Rules,
    min_pause_seconds: float = 0.1,
    max_pause_seconds: float = 2.0,
    epsilon: float = 1e-6,
) -> list[EOTCandidate]:
    """Derive conservative EOT positive/negative candidates.

    A positive proxy requires another participant to claim the floor no later
    than ``rules.max_switch_gap_seconds`` after the holder stops.  A hard
    negative is a pause of ``[min_pause_seconds, max_pause_seconds]`` after
    which the same holder resumes with no intervening floor claim.  Long gaps,
    terminal boundaries, annotation continuations, and disputed overlap are
    kept as ambiguous/excluded rather than silently becoming labels.
    """

    if min_pause_seconds < 0 or max_pause_seconds < min_pause_seconds:
        raise ValueError("Require 0 <= min_pause_seconds <= max_pause_seconds")
    event_by_uid = _event_map(
        utterances,
        camera_wearer_id=camera_wearer_id,
        rules=rules,
    )
    results: list[EOTCandidate] = []

    for holder in utterances:
        split = "test" if holder.session in TEST_SESSIONS else "train"
        holder_role = "wearer" if holder.participant_id == camera_wearer_id else "other"
        boundary = holder.dense_end_seconds

        def make(
            category: str,
            confidence: str,
            reason: str,
            *,
            pause: float | None = None,
            incoming: Utterance | None = None,
        ) -> EOTCandidate:
            incoming_event = None if incoming is None else event_by_uid[incoming.uid]
            return EOTCandidate(
                session=holder.session,
                split=split,
                category=category,
                confidence=confidence,
                holder_uid=holder.uid,
                holder_participant_id=holder.participant_id,
                holder_role=holder_role,
                boundary_seconds=boundary,
                pause_seconds=pause,
                incoming_uid=None if incoming is None else incoming.uid,
                incoming_participant_id=(
                    None if incoming is None else incoming.participant_id
                ),
                incoming_role=(
                    None
                    if incoming is None
                    else ("wearer" if incoming.participant_id == camera_wearer_id else "other")
                ),
                incoming_start_seconds=(
                    None if incoming is None else incoming.dense_start_seconds
                ),
                incoming_event_type=(
                    None if incoming_event is None else incoming_event.event_type
                ),
                reason=reason,
            )

        if _same_speaker_continues_across_boundary(
            holder,
            utterances,
            epsilon=epsilon,
        ):
            results.append(make(
                "eot_excluded_internal_boundary",
                "excluded",
                "another annotation for the same speaker continues across this endpoint",
            ))
            continue

        resume = _next_same_speaker(holder, utterances, epsilon=epsilon)
        pause = None if resume is None else max(0.0, resume.dense_start_seconds - boundary)
        if (
            resume is not None
            and pause <= min_pause_seconds + epsilon
            and (resume.is_cross_clip_continuation or holder.ends_at_clip_boundary)
        ):
            results.append(make(
                "eot_excluded_file_continuation",
                "excluded",
                "same-speaker activity continues across an adjacent file boundary",
                pause=pause,
                incoming=resume,
            ))
            continue

        incoming_rows = _incoming_before_resume(
            holder,
            utterances,
            resume=resume,
            epsilon=epsilon,
        )
        floor_claims = [
            row
            for row in incoming_rows
            if event_by_uid[row.uid].event_type in FLOOR_CLAIM_EVENT_TYPES
        ]
        ambiguous = [
            row
            for row in incoming_rows
            if event_by_uid[row.uid].event_type in AMBIGUOUS_EVENT_TYPES
        ]
        unresolved = [
            row
            for row in incoming_rows
            if event_by_uid[row.uid].event_type
            not in FLOOR_CLAIM_EVENT_TYPES
            | NON_FLOOR_EVENT_TYPES
            | AMBIGUOUS_EVENT_TYPES
            | {"cross_clip_continuation"}
        ]

        if floor_claims:
            incoming = floor_claims[0]
            gap = incoming.dense_start_seconds - boundary
            if gap <= rules.max_switch_gap_seconds + epsilon:
                incoming_type = event_by_uid[incoming.uid].event_type
                category = (
                    "eot_positive_overlap_takeover_proxy"
                    if incoming.dense_start_seconds < boundary - epsilon
                    else "eot_positive_handover_proxy"
                )
                results.append(make(
                    category,
                    "timing_proxy",
                    f"incoming {incoming_type} claims the floor within the allowed gap",
                    pause=max(0.0, gap),
                    incoming=incoming,
                ))
            else:
                results.append(make(
                    "eot_ambiguous_long_gap_before_handover",
                    "ambiguous",
                    "another speaker eventually claims the floor, but only after a long gap",
                    pause=gap,
                    incoming=incoming,
                ))
            continue

        if ambiguous or unresolved:
            incoming = sorted(
                [*ambiguous, *unresolved],
                key=lambda row: (max(boundary, row.dense_start_seconds), row.uid),
            )[0]
            incoming_type = event_by_uid[incoming.uid].event_type
            gap = incoming.dense_start_seconds - boundary
            if incoming.dense_start_seconds < boundary - epsilon:
                category = "eot_ambiguous_overlap"
            elif gap > rules.max_switch_gap_seconds + epsilon:
                category = "eot_ambiguous_long_gap_before_handover"
            else:
                category = "eot_ambiguous_short_or_unresolved_incoming"
            results.append(make(
                category,
                "ambiguous",
                f"an unresolved incoming event ({incoming_type}) occurs before the holder resumes",
                pause=max(0.0, gap),
                incoming=incoming,
            ))
            continue

        if resume is not None:
            if pause < min_pause_seconds - epsilon:
                category = "eot_excluded_annotation_seam"
                confidence = "excluded"
                reason = "same speaker resumes too quickly to form a meaningful pause"
            elif pause <= max_pause_seconds + epsilon:
                category = "eot_negative_mid_turn_pause_proxy"
                confidence = "timing_proxy"
                reason = "same speaker resumes within the pause limit with no floor claim"
            else:
                category = "eot_ambiguous_long_same_speaker_gap"
                confidence = "ambiguous"
                reason = "same speaker resumes after a long gap; turn continuity is unknown"
            results.append(make(
                category,
                confidence,
                reason,
                pause=pause,
                incoming=resume,
            ))
            continue

        # TurnBench counts a speaker's final turn as EOT.  We deliberately do
        # not: EasyCom clips and redacted source gaps make the endpoint
        # semantically unknown without manual review.
        results.append(make(
            "eot_excluded_terminal_or_gap",
            "excluded",
            "no later same-speaker or floor-claim evidence in this contiguous segment",
        ))

    return results


def _serialize_event(event: Event) -> dict:
    payload = asdict(event)
    payload["target_of_speech"] = list(event.target_of_speech)
    return payload


def _write_rows(path: Path, rows: list[dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    if path.suffix == ".jsonl":
        with path.open("w", encoding="utf-8") as handle:
            for row in rows:
                handle.write(json.dumps(row, ensure_ascii=False) + "\n")
        return
    fieldnames = list(rows[0]) if rows else []
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        if fieldnames:
            writer.writeheader()
            for row in rows:
                encoded = {
                    key: json.dumps(value, ensure_ascii=False)
                    if isinstance(value, (list, tuple, dict))
                    else value
                    for key, value in row.items()
                }
                writer.writerow(encoded)


def _print_counts(title: str, counts: Counter) -> None:
    print(title)
    for name, count in sorted(counts.items()):
        print(f"  {name:<44} {count:>6}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dataset-root", type=Path, required=True)
    parser.add_argument("--sessions", nargs="+", default=list(ALL_SESSIONS))
    parser.add_argument("--camera-wearer-id", type=int, default=2)
    parser.add_argument("--output-dir", type=Path)
    parser.add_argument("--min-pause-seconds", type=float, default=0.1)
    parser.add_argument("--max-pause-seconds", type=float, default=2.0)
    parser.add_argument("--context-seconds", type=float, default=1.5)
    parser.add_argument("--backchannel-max-seconds", type=float, default=1.0)
    parser.add_argument("--max-switch-gap-seconds", type=float, default=0.6)
    parser.add_argument("--min-overlap-seconds", type=float, default=0.1)
    parser.add_argument("--takeover-candidate-seconds", type=float, default=0.4)
    parser.add_argument("--takeover-high-seconds", type=float, default=0.8)
    parser.add_argument("--max-contiguous-source-gap-seconds", type=float, default=1.0)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    unknown = sorted(set(args.sessions) - set(ALL_SESSIONS))
    if unknown:
        raise ValueError(f"Unknown EasyCom sessions: {unknown}")
    rules = Rules(
        context_seconds=args.context_seconds,
        backchannel_max_seconds=args.backchannel_max_seconds,
        max_switch_gap_seconds=args.max_switch_gap_seconds,
        min_overlap_seconds=args.min_overlap_seconds,
        takeover_candidate_seconds=args.takeover_candidate_seconds,
        takeover_high_seconds=args.takeover_high_seconds,
        max_contiguous_source_gap_seconds=args.max_contiguous_source_gap_seconds,
    )

    all_eot: list[EOTCandidate] = []
    all_events: list[Event] = []
    for session in args.sessions:
        clips = build_clips(args.dataset_root.resolve(), session, rules)
        utterances = load_utterances(session, clips)
        events = list(_event_map(
            utterances,
            camera_wearer_id=args.camera_wearer_id,
            rules=rules,
        ).values())
        eot = derive_eot_candidates(
            utterances,
            camera_wearer_id=args.camera_wearer_id,
            rules=rules,
            min_pause_seconds=args.min_pause_seconds,
            max_pause_seconds=args.max_pause_seconds,
        )
        all_events.extend(events)
        all_eot.extend(eot)
        print(f"{session}: utterances={len(utterances)} eot_boundaries={len(eot)}")

    for split in ("train", "test", "all"):
        selected_eot = [row for row in all_eot if split == "all" or row.split == split]
        selected_events = [row for row in all_events if split == "all" or row.split == split]
        print(f"\n[{split}]")
        _print_counts(
            "Other-floor EOT candidates",
            Counter(row.category for row in selected_eot if row.holder_role == "other"),
        )
        _print_counts(
            "All-holder EOT candidates",
            Counter(row.category for row in selected_eot),
        )
        int_counts = Counter()
        for event in selected_events:
            if event.participant_role != "wearer":
                continue
            for group, names in INT_PROXY_GROUPS.items():
                if event.event_type in names:
                    int_counts[group] += 1
                    int_counts[f"{group}/{event.event_type}"] += 1
                    break
        _print_counts("Wearer takeover/backchannel candidates", int_counts)

    if args.output_dir is not None:
        output = args.output_dir.resolve()
        eot_rows = [
            asdict(row)
            for row in all_eot
            if row.holder_role == "other"
            and not row.category.startswith("eot_excluded_")
        ]
        int_rows = []
        candidate_names = set().union(*INT_PROXY_GROUPS.values())
        for event in all_events:
            if event.participant_role == "wearer" and event.event_type in candidate_names:
                row = _serialize_event(event)
                row["proxy_group"] = next(
                    group for group, names in INT_PROXY_GROUPS.items()
                    if event.event_type in names
                )
                int_rows.append(row)
        _write_rows(output / "easycom_eot_candidates.jsonl", eot_rows)
        _write_rows(output / "easycom_eot_candidates.csv", eot_rows)
        _write_rows(output / "easycom_wearer_int_candidates.jsonl", int_rows)
        _write_rows(output / "easycom_wearer_int_candidates.csv", int_rows)
        split_summaries = {}
        candidate_names = set().union(*INT_PROXY_GROUPS.values())
        for split in ("train", "test", "all"):
            selected_eot = [
                row for row in all_eot if split == "all" or row.split == split
            ]
            selected_events = [
                event for event in all_events if split == "all" or event.split == split
            ]
            proxy_groups = Counter()
            for event in selected_events:
                if event.participant_role != "wearer":
                    continue
                for group, names in INT_PROXY_GROUPS.items():
                    if event.event_type in names:
                        proxy_groups[group] += 1
                        break
            split_summaries[split] = {
                "other_floor_eot_counts": dict(sorted(Counter(
                    row.category for row in selected_eot if row.holder_role == "other"
                ).items())),
                "wearer_proxy_group_counts": dict(sorted(proxy_groups.items())),
                "wearer_event_type_counts": dict(sorted(Counter(
                    event.event_type
                    for event in selected_events
                    if event.participant_role == "wearer"
                    and event.event_type in candidate_names
                ).items())),
            }
        summary = {
            "schema_version": 1,
            "dataset": "EasyCom",
            "semantic_warning": (
                "All EOT, takeover, interruption, and backchannel categories are "
                "timing/transcript-derived proxies, not human-consensus labels."
            ),
            "sessions": args.sessions,
            "rules": asdict(rules),
            "min_pause_seconds": args.min_pause_seconds,
            "max_pause_seconds": args.max_pause_seconds,
            "splits": split_summaries,
        }
        output.mkdir(parents=True, exist_ok=True)
        (output / "summary.json").write_text(
            json.dumps(summary, indent=2, ensure_ascii=False) + "\n",
            encoding="utf-8",
        )
        print(f"\nSaved proxy audit to {output}")


if __name__ == "__main__":
    main()