import os

import cv2


def mosaic_face_with_yunet(
    input_video_path: str,
    model_path: str,
    output_video_path: str,
    frames_output_dir: str,
    mosaic_scale: float = 0.1,
    score_threshold: float = 0.9,
    nms_threshold: float = 0.3,
    top_k: int = 5000,
):
    """
    YuNet (FaceDetectorYN) を使用して動画の顔のみをモザイク処理し、
    モザイク済みの動画と全フレームを出力する関数。

    Parameters
    ----------
    input_video_path : str
        入力動画ファイルのパス
    model_path : str
        YuNet の ONNX モデルファイルへのパス
    output_video_path : str
        モザイクをかけた結果の出力動画ファイルのパス
    frames_output_dir : str
        全フレーム画像の出力先ディレクトリパス
    mosaic_scale : float
        モザイクの縮小率 (例: 0.1 なら幅高さを1/10に縮小して再拡大)
    score_threshold : float
        顔検出におけるスコア閾値 (自信度がこれより低い検出結果は無視)
    nms_threshold : float
        NMS(非最大抑制)の閾値
    top_k : int
        NMS 後に保持する最大検出数
    """

    # 出力用ディレクトリ作成
    os.makedirs(frames_output_dir, exist_ok=True)

    # 動画を開く
    cap = cv2.VideoCapture(input_video_path)
    if not cap.isOpened():
        print(f"動画ファイルを開けませんでした: {input_video_path}")
        return

    # 動画情報を取得
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)

    # 動画出力の設定 (fourcc はmp4用にmp4vを使用)
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))

    # YuNet FaceDetectorYN の初期化
    # 入力サイズは推論用に一旦 (320, 320) など固定サイズを与えておき、
    # 各フレーム処理時に setInputSize で動的に変更します。
    face_detector = cv2.FaceDetectorYN.create(
        model=model_path,
        config="",
        input_size=(320, 320),  # 後ほどフレームサイズに合わせて上書き
        score_threshold=score_threshold,
        nms_threshold=nms_threshold,
        top_k=top_k,
        backend_id=cv2.dnn.DNN_BACKEND_DEFAULT,
        target_id=cv2.dnn.DNN_TARGET_CPU,
    )

    frame_count = 0

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        # フレームサイズが変化する場合は毎回 setInputSize を呼び出す
        face_detector.setInputSize((frame.shape[1], frame.shape[0]))

        # 顔検出
        # detect() は (retval, faces) のタプルを返し、faces が検出結果
        # faces の列は [x, y, w, h, confidence, x1, y1, x2, y2, ... x5, y5] (ランドマーク)
        _, faces = face_detector.detect(frame)

        if faces is not None:
            for face in faces:
                x, y, w, h = face[:4].astype(int)
                # 必要に応じて座標補正（画像範囲外をはみ出さないようにクリップ）
                x = max(0, x)
                y = max(0, y)
                w = min(w, frame.shape[1] - x)
                h = min(h, frame.shape[0] - y)

                # 顔領域を取得
                face_roi = frame[y : y + h, x : x + w]

                # 縮小
                face_roi_small = cv2.resize(
                    face_roi,
                    (max(1, int(w * mosaic_scale)), max(1, int(h * mosaic_scale))),
                    interpolation=cv2.INTER_LINEAR,
                )
                # 再拡大 (NEARESTで拡大するとドットが強めのモザイクに)
                face_roi_mosaic = cv2.resize(
                    face_roi_small, (w, h), interpolation=cv2.INTER_NEAREST
                )

                # モザイク部分をフレームに貼り付け
                frame[y : y + h, x : x + w] = face_roi_mosaic

        # フレーム画像を保存
        frame_filename = os.path.join(frames_output_dir, f"frame_{frame_count:06d}.jpg")
        cv2.imwrite(frame_filename, frame)

        # モザイク処理済みフレームを動画ファイルに書き出し
        out.write(frame)

        frame_count += 1

    cap.release()
    out.release()
    print("処理が完了しました。")


def main():
    input_video = "video/tes.mp4"
    model_path = "./face_detection_yunet_2023mar.onnx"
    output_video = "output_mosaic_yunet.mp4"
    frames_dir = "frames"

    mosaic_face_with_yunet(
        input_video_path=input_video,
        model_path=model_path,
        output_video_path=output_video,
        frames_output_dir=frames_dir,
        mosaic_scale=0.2,  # モザイクの粗さを調整
        score_threshold=0.9,  # 検出結果の信頼度が 0.9 未満のものは無視
        nms_threshold=0.3,  # NMS 閾値
        top_k=5000,  # NMS 後に保持する最大数
    )


if __name__ == "__main__":
    main()