#!/usr/bin/env python3
from __future__ import annotations

import argparse
from collections import deque
from pathlib import Path

import cv2
import numpy as np
import torch
from torch import nn

from fall_detection.features import normalize_pose_sequence
from fall_detection.models import PoseGRU, PoseTCN
from fall_detection.pose_extractor import MediaPipePoseExtractor

CONNECTIONS = [
    (11, 12), (11, 13), (13, 15), (12, 14), (14, 16),
    (11, 23), (12, 24), (23, 24), (23, 25), (25, 27),
    (24, 26), (26, 28), (27, 31), (28, 32),
]


def draw_pose(frame: np.ndarray, pose: np.ndarray | None) -> None:
    if pose is None:
        return
    height, width = frame.shape[:2]
    for first, second in CONNECTIONS:
        if min(pose[first, 3], pose[second, 3]) < 0.35:
            continue
        p1 = (int(pose[first, 0] * width), int(pose[first, 1] * height))
        p2 = (int(pose[second, 0] * width), int(pose[second, 1] * height))
        cv2.line(frame, p1, p2, (70, 230, 70), 2)
    for point in pose:
        if point[3] >= 0.35:
            cv2.circle(frame, (int(point[0] * width), int(point[1] * height)), 2, (0, 255, 255), -1)


def load_model(path: Path, device: torch.device) -> tuple[nn.Module, dict]:
    checkpoint = torch.load(path, map_location=device, weights_only=False)
    if checkpoint.get("architecture", "gru") == "tcn":
        model = PoseTCN(
            input_size=checkpoint["input_size"],
            channels=checkpoint["tcn_channels"],
            dropout=checkpoint["dropout"],
        ).to(device)
    else:
        model = PoseGRU(
            input_size=checkpoint["input_size"],
            hidden_size=checkpoint["hidden_size"],
            num_layers=checkpoint["num_layers"],
            dropout=checkpoint["dropout"],
        ).to(device)
    model.load_state_dict(checkpoint["state_dict"])
    model.eval()
    return model, checkpoint


def main() -> None:
    parser = argparse.ArgumentParser(description="Run fall detection on a video")
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--model", type=Path, default=Path("artifacts/experiments/urfd/pose_tcn/model.pt"))
    parser.add_argument("--output", type=Path)
    parser.add_argument("--target-fps", type=float, default=10.0)
    parser.add_argument("--smoothing", type=int, default=5)
    parser.add_argument("--alert-consecutive", type=int, default=3)
    args = parser.parse_args()
    output = args.output or args.input.with_name(args.input.stem + "_detected.mp4")
    output.parent.mkdir(parents=True, exist_ok=True)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model, checkpoint = load_model(args.model, device)
    sequence_length = int(checkpoint["sequence_length"])
    feature_mean = np.asarray(checkpoint["feature_mean"], dtype=np.float32)
    feature_std = np.asarray(checkpoint["feature_std"], dtype=np.float32)
    threshold = float(checkpoint["threshold"])

    capture = cv2.VideoCapture(str(args.input))
    if not capture.isOpened():
        raise SystemExit(f"Cannot open {args.input}")
    source_fps = capture.get(cv2.CAP_PROP_FPS) or 25.0
    width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
    writer = cv2.VideoWriter(
        str(output), cv2.VideoWriter_fourcc(*"mp4v"), source_fps, (width, height)
    )
    sample_every = max(1, int(round(source_fps / args.target_fps)))
    pose_window: deque[np.ndarray] = deque(maxlen=sequence_length)
    probability_window: deque[float] = deque(maxlen=args.smoothing)
    probability = 0.0
    consecutive = 0
    frame_index = 0

    with MediaPipePoseExtractor() as extractor:
        while True:
            ok, frame = capture.read()
            if not ok:
                break
            current_pose = None
            display_pose = None
            if frame_index % sample_every == 0:
                is_wide_urfd = frame.shape[1] / frame.shape[0] > 2.2
                pose_frame = frame[:, frame.shape[1] // 2 :] if is_wide_urfd else frame
                current_pose = extractor.process_rgb(cv2.cvtColor(pose_frame, cv2.COLOR_BGR2RGB))
                if current_pose is None:
                    current_pose = pose_window[-1].copy() if pose_window else np.zeros((33, 4), dtype=np.float32)
                pose_window.append(current_pose)
                display_pose = current_pose.copy()
                if is_wide_urfd:
                    display_pose[:, 0] = 0.5 + 0.5 * display_pose[:, 0]
                if len(pose_window) >= sequence_length:
                    features = normalize_pose_sequence(
                        np.stack(pose_window), checkpoint["visibility_threshold"]
                    )
                    feature_indices = np.asarray(
                        checkpoint.get("feature_indices", np.arange(features.shape[1])),
                        dtype=np.int64,
                    )
                    features = features[:, feature_indices]
                    features = (features - feature_mean) / feature_std
                    tensor = torch.from_numpy(features[None]).float().to(device)
                    with torch.no_grad():
                        probability = float(torch.sigmoid(model(tensor)).item())
                    probability_window.append(probability)
                    probability = float(np.mean(probability_window))
                    consecutive = consecutive + 1 if probability >= threshold else 0

            draw_pose(frame, display_pose)
            alert = consecutive >= args.alert_consecutive
            color = (0, 0, 255) if alert else (0, 180, 0)
            status = "CANH BAO: TE NGA" if alert else "BINH THUONG"
            cv2.rectangle(frame, (10, 10), (385, 78), (0, 0, 0), -1)
            cv2.putText(frame, status, (20, 38), cv2.FONT_HERSHEY_SIMPLEX, 0.72, color, 2)
            cv2.putText(
                frame,
                f"P(fall)={probability:.2f}  threshold={threshold:.2f}",
                (20, 66),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.50,
                (255, 255, 255),
                1,
            )
            writer.write(frame)
            frame_index += 1

    capture.release()
    writer.release()
    print(f"Saved annotated video to {output}")


if __name__ == "__main__":
    main()