# vmspy のインストール
# - 使用している Python のバージョンに対応する vmspy パッケージをダウンロードします
!unzip vmspy.zip
!./vmspy/install-vmspy-collab.sh

# YOLO のインストール（Ultralytics）
%pip install ultralytics
import ultralytics
from ultralytics import YOLO

# YOLO モデルをロード（軽量版）
model = YOLO('yolo11n.pt')

# 必要なモジュールをインポート
from tarfile import NUL  # null の戻り値として使用（通常は None の使用を推奨）
import cv2
import time
from shapely.geometry import Point
from shapely.geometry.polygon import Polygon

# 入力フレームの解像度（正規化に使用）
input_width = 640
input_height = 640

# live_video オブジェクトの設定（VMS ストリーミングインターフェース）
live_video = vmspy.live_video()
live_video.set_keyframe_only(True)                 # パフォーマンス最適化のためキーフレームのみ使用
live_video.set_image_size(input_width, input_height)
live_video.set_pixel_format("BGR")                 # OpenCV は BGR フォーマットを使用
live_video.init("url.to.vmsserver", 3300, "admin", "admin")

# 監視エリアの定義（正規化座標: 0.0 ～ 1.0）
# このエリアは VMS サーバー上にも可視化されます
monitoring_zone_points = [
    (0.3380152329749104, 0.3314180107526881),
    (0.42672491039426524, 0.3404905913978495),
    (0.41261200716845886, 0.9241599462365592),
    (0.1323700716845878, 0.9060147849462368)
]

monitoring_zone = Polygon(monitoring_zone_points)

# VMS 画面上に監視エリアを描画
live_video.draw_polygon(monitoring_zone_points, border_color="purple", border_thickness=5)

# VMS 上の検知表示設定
live_video.set_detection_class(model.names)
live_video.set_detection_box(border_color="blue", border_thickness=5)
live_video.set_detection_caption(
    show_object_id=True,
    show_class_name=True,
    show_confidence=False
)

# 設定パラメータ（秒単位）
config = {
    "parking_violation_sec": 10,   # 駐車違反と判定するまでの時間
    "periodic_check_sec": 10,      # 再通知の間隔
    "remove_criteria_sec": 10      # 非アクティブな追跡オブジェクトを削除するまでの時間
}

# オブジェクトの保持とイベント発生を管理する簡易トラッキングロジック
class TrackingHistory():
    def __init__(self):
        self.history = {}

    def update(self, key, frame_info):
        """
        指定されたオブジェクト（track_id）の追跡情報を更新します。
        通知条件を満たした場合、最初の frame_info を返します。
        """
        timestamp = frame_info["timestamp"]

        if key in self.history:
            track_item = self.history[key]
            track_item["timestamp_last"] = timestamp

            # 再通知のタイミングかどうかを確認
            if timestamp > track_item["timestamp_report"]:
                track_item["timestamp_report"] = timestamp + config["periodic_check_sec"]
                print(f"{key}: REPORTED / n={len(self.history)}")
                return track_item["frameinfo0"]
        else:
            # 初めて検出されたオブジェクト
            self.history[key] = {
                "frameinfo0": frame_info.copy(),  # 初回フレーム情報を保存
                "timestamp_first": timestamp,
                "timestamp_last": timestamp,
                "timestamp_report": timestamp + config["parking_violation_sec"]
            }
            print(f"{key}: CREATED / n={len(self.history)}")

        return NUL  # 通知対象なし

    def remove_old_items(self, frame_info):
        """
        最近更新されていないオブジェクトを削除します。
        """
        current_timestamp = frame_info["timestamp"]
        min_sec = config["remove_criteria_sec"]

        keys_to_remove = [
            key for key, value in self.history.items()
            if current_timestamp - value["timestamp_last"] > min_sec
        ]

        for key in keys_to_remove:
            sec = current_timestamp - self.history[key]["timestamp_first"]
            del self.history[key]
            print(f"{key}: REMOVED ({sec}sec) / n={len(self.history)}")

        return len(keys_to_remove)

tracking = TrackingHistory()

# チャンネルからライブ映像を開始
ch_no = 1
live_video.start(ch_no)

count = 0
while count < 10000:
    count += 1

    # フレームとメタデータを取得
    (frame_image, frame_info) = live_video.get_frame()

    # 終了条件：フレームが取得できない場合
    if frame_image.size == 0 and count > 2:
        print("\nEnd of stream")
        break

    # オブジェクト検知 / トラッキングの実行
    timestamp_detect_start = time.time_ns()
    results = model.track(frame_image, persist=True, verbose=False)
    # 代替: results = model.predict(frame_image, verbose=False)
    time_detect_elapsed = (time.time_ns() - timestamp_detect_start) / 1_000_000_000

    # 検知結果の処理
    for box in results[0].boxes:

        # トラッキング ID を取得
        if box.id is not None:
            track_id = int(box.id[0].item())
        else:
            continue  # トラッキング ID がない場合はスキップ

        # バウンディングボックスを正規化座標に変換
        xyxy = box.xyxy[0].tolist()
        x = xyxy[0] / input_width
        y = xyxy[1] / input_height
        w = (xyxy[2] - xyxy[0]) / input_width
        h = (xyxy[3] - xyxy[1]) / input_height

        class_id = int(box.cls[0].item())
        conf = box.conf[0].item()

        # オブジェクト中心が監視エリア内にあるか確認
        if monitoring_zone.contains(Point(x + w / 2, y + h / 2)):
            name = model.names[class_id]

            # VMS に検知情報を送信
            live_video.input_detection(
                x, y, w, h,
                class_id=class_id,
                object_id=track_id,
                confidence=conf
            )

            # トラッキング状態を更新
            frame_info0 = tracking.update(track_id, frame_info)

            # 違反条件を満たした場合、イベントを通知
            if frame_info0 != NUL:
                duration = float(frame_info["timestamp"] - frame_info0["timestamp"])
                print(f"duration: {duration}")

                live_video.report_event(
                    frame_info0,
                    x, y, w, h,
                    duration=duration,
                    event_type_id=0,
                    class_id=class_id,
                    object_id=track_id
                )

    # 古いトラッキングデータの削除
    tracking.remove_old_items(frame_info)

    # すべての検知結果を VMS に送信
    live_video.send_detections(frame_info)