backend/tracker.py

"""
tracker.py — Core detection + tracking engine
Counting : FIRST-SEEN (chaque track_id compté une seule fois dès sa 1re apparition)
CSV log  : suit exactement le schéma SCHEMA_EXAMPLE.csv
           frame, timestamp_sec, scene_name, group_id, video_name, track_id,
           class_name, confidence, bbox_x1, bbox_y1, bbox_x2, bbox_y2,
           cx, cy, frame_width, frame_height, crossed_line, direction, speed_px_s
"""

import cv2
import csv
import json
import math
import re
import uuid
import numpy as np
from pathlib import Path
from datetime import datetime
from collections import defaultdict
from ultralytics import YOLO

# ─── Classes COCO traffic ─────────────────────────────────────────────────────
TRAFFIC_CLASSES = {
    0: "person",
    1: "bicycle",
    2: "car",
    3: "motorbike",
    5: "bus",
    7: "truck",
}

CLASS_COLORS = {
    "person":    (0,   200, 255),
    "bicycle":   (50,  255, 50),
    "car":       (255, 165, 0),
    "motorbike": (255, 50,  200),
    "bus":       (0,   100, 255),
    "truck":     (180, 0,   255),
}

DEFAULT_CLASSES = list(TRAFFIC_CLASSES.values())

# Colonnes CSV — ordre exact du schéma
CSV_COLUMNS = [
    "frame", "timestamp_sec", "scene_name", "group_id", "video_name",
    "track_id", "class_name", "confidence",
    "bbox_x1", "bbox_y1", "bbox_x2", "bbox_y2",
    "cx", "cy", "frame_width", "frame_height",
    "crossed_line", "direction", "speed_px_s",
]


class TrafficTracker:
    def __init__(
        self,
        model_path: str = "best.pt",
        selected_classes: list = None,
        conf_threshold: float = 0.5,
        scene_name: str = "scene_01",
        group_id: str = "group_05",
        video_name: str = "video.mp4",
        output_dir: str = "logs",
    ):
        # ── Modèle ────────────────────────────────────────────────────────────
        # Pour utiliser un modèle fine-tuné :
        #   model_path = "runs/detect/traffic_finetune/weights/best.pt"
        self.model            = YOLO(model_path)

        self.selected_classes = selected_classes or DEFAULT_CLASSES
        self.conf             = conf_threshold
        self.scene_name       = scene_name
        self.group_id         = group_id          # <— nouveau champ schéma
        self.video_name       = video_name        # <— nouveau champ schéma
        self.output_dir       = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        # State
        self.session_id   = str(uuid.uuid4())[:8]
        self.frame_index  = 0
        self.fps          = 30.0
        self.frame_width  = 0
        self.frame_height = 0

        # Tracking
        self.track_history:   dict = defaultdict(list)   # id -> [(cx,cy), ...]
        self.counted_ids:     set  = set()
        self.count_per_class: dict = defaultdict(int)

        # Logs
        self.detection_log: list = []   # chaque ligne = une détection
        self.frame_stats:   list = []   # résumé par frame

        # Video writer
        self.video_writer      = None
        self.output_video_path = None

        # Filtre classes COCO
        self._class_ids = [
            cid for cid, name in TRAFFIC_CLASSES.items()
            if name in self.selected_classes
        ]

    # ── Get next order number ───────────────────────────────────────────────────
    def _get_next_order_number(self) -> int:
        """
        Retourne le prochain numéro d'ordre pour les fichiers de logs.
        Cherche les fichiers existants avec le pattern Group_X_Y_NNN_*.
        """
        try:
            prefix = f"{self.group_id}_{self.scene_name}"
            pattern = f"{prefix}_*"
            order_pattern = re.compile(rf"^{re.escape(prefix)}_(\d{{3}})(?:_|\.|$)")
            existing = list(self.output_dir.glob(pattern))
            if not existing:
                return 1
            numbers = []
            for f in existing:
                match = order_pattern.match(f.name)
                if match:
                    numbers.append(int(match.group(1)))
            return max(numbers, default=0) + 1
        except Exception:
            return 1

    def setup_frame_source(self, width: int, height: int, fps: float = 5.0):
        self.fps = fps or 5.0
        self.frame_width = width
        self.frame_height = height

    # ── Setup ──────────────────────────────────────────────────────────────────
    def setup_video(self, cap: cv2.VideoCapture, save_output: bool = True):
        self.fps          = cap.get(cv2.CAP_PROP_FPS) or 30.0
        self.frame_width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        if save_output:
            # Générer le numéro d'ordre (incrément des sessions)
            video_order = self._get_next_order_number()
            out_name = f"{self.group_id}_{self.scene_name}_{video_order:03d}_annotated.mp4"
            self.output_video_path = str(self.output_dir / out_name)
            fourcc = cv2.VideoWriter_fourcc(*"mp4v")
            self.video_writer = cv2.VideoWriter(
                self.output_video_path, fourcc,
                self.fps, (self.frame_width, self.frame_height)
            )

    # ── Process one frame ──────────────────────────────────────────────────────
    def process_frame(self, frame: np.ndarray) -> tuple:
        self.frame_index += 1
        timestamp_sec = round(self.frame_index / self.fps, 3)

        results = self.model.track(
            frame,
            persist=True,
            conf=self.conf,
            classes=self._class_ids,
            tracker="bytetrack.yaml",
            verbose=False,
        )

        annotated          = frame.copy()
        frame_detections   = []
        per_class_in_frame = defaultdict(int)
        any_object_visible = False

        result = results[0]

        if result.boxes is not None and len(result.boxes) > 0:
            for box in result.boxes:
                cls_id   = int(box.cls[0])
                cls_name = TRAFFIC_CLASSES.get(cls_id, "unknown")
                if cls_name not in self.selected_classes:
                    continue

                conf_val        = float(box.conf[0])
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                track_id        = int(box.id[0]) if box.id is not None else -1
                cx, cy          = (x1 + x2) // 2, (y1 + y2) // 2

                any_object_visible = True
                per_class_in_frame[cls_name] += 1

                # ── Vitesse (px/s) ────────────────────────────────────────────
                # Distance euclidienne entre position courante et précédente,
                # multipliée par fps pour obtenir des px/s.
                speed_px_s = 0.0
                if track_id != -1 and self.track_history[track_id]:
                    prev_cx, prev_cy = self.track_history[track_id][-1]
                    dist = math.hypot(cx - prev_cx, cy - prev_cy)
                    speed_px_s = round(dist * self.fps, 1)

                # ── Comptage first-seen ───────────────────────────────────────
                is_new = False
                if track_id != -1 and track_id not in self.counted_ids:
                    self.counted_ids.add(track_id)
                    self.count_per_class[cls_name] += 1
                    is_new = True

                # Mise à jour historique
                if track_id != -1:
                    self.track_history[track_id].append((cx, cy))

                # ── Colonnes schéma : crossed_line & direction ─────────────────
                # On ne dessine plus de ligne physique sur la vidéo, mais on
                # conserve les colonnes dans le CSV (vide / false par défaut).
                # Si tu veux réactiver le croisement de ligne, tu peux définir
                # une ligne ici et compléter la logique.
                crossed_line = False
                direction    = ""

                # ── Ligne CSV (schéma exact) ───────────────────────────────────
                row = {
                    "frame":        self.frame_index,
                    "timestamp_sec": timestamp_sec,
                    "scene_name":   self.scene_name,
                    "group_id":     self.group_id,
                    "video_name":   self.video_name,
                    "track_id":     track_id,
                    "class_name":   cls_name,
                    "confidence":   round(conf_val, 3),
                    "bbox_x1":      x1,
                    "bbox_y1":      y1,
                    "bbox_x2":      x2,
                    "bbox_y2":      y2,
                    "cx":           cx,
                    "cy":           cy,
                    "frame_width":  self.frame_width,
                    "frame_height": self.frame_height,
                    "crossed_line": str(crossed_line).lower(),   # "true"/"false"
                    "direction":    direction,
                    "speed_px_s":   speed_px_s,
                }
                frame_detections.append(row)
                self.detection_log.append(row)

                # ── Bounding box ───────────────────────────────────────────────
                color     = CLASS_COLORS.get(cls_name, (255, 255, 255))
                thickness = 3 if is_new else 2
                cv2.rectangle(annotated, (x1, y1), (x2, y2), color, thickness)

                id_str = f"#{track_id}" if track_id != -1 else "#?"
                label  = f"{cls_name} {id_str} {conf_val:.2f}"
                (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 1)
                cv2.rectangle(annotated, (x1, y1 - th - 8), (x1 + tw + 4, y1), color, -1)
                cv2.putText(annotated, label, (x1 + 2, y1 - 4),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 0, 0), 1)

                # Trail
                if track_id != -1:
                    trail = self.track_history[track_id][-20:]
                    for i in range(1, len(trail)):
                        alpha = i / len(trail)
                        tc    = tuple(int(c * alpha) for c in color)
                        cv2.line(annotated, trail[i - 1], trail[i], tc, 2)

        self._draw_counters(annotated, any_object_visible)

        if self.video_writer is not None:
            self.video_writer.write(annotated)

        frame_stat = {
            "frame":       self.frame_index,
            "timestamp":   timestamp_sec,
            "scene":       self.scene_name,
            "detections":  len(frame_detections),
            "per_class":   dict(per_class_in_frame),
            "any_visible": any_object_visible,
            "cumulative":  dict(self.count_per_class),
        }
        self.frame_stats.append(frame_stat)

        return annotated, frame_stat

    # ── Overlay compteurs ──────────────────────────────────────────────────────
    def _draw_counters(self, frame: np.ndarray, any_visible: bool):
        overlay = frame.copy()
        cv2.rectangle(overlay, (0, 0),
                      (260, 30 + 22 * len(self.selected_classes) + 30),
                      (20, 20, 20), -1)
        cv2.addWeighted(overlay, 0.6, frame, 0.4, 0, frame)

        cv2.putText(frame, "TRAFFIC COUNTER", (10, 20),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)
        y = 42
        for cls_name in self.selected_classes:
            count = self.count_per_class.get(cls_name, 0)
            color = CLASS_COLORS.get(cls_name, (255, 255, 255))
            cv2.putText(frame, f" {cls_name:<12} {count:>4}", (8, y),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
            y += 22

        if not any_visible:
            cv2.putText(frame, "NO OBJECTS DETECTED", (10, y + 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (80, 80, 80), 1)

    # ── Sauvegarde des logs ────────────────────────────────────────────────────
    def save_logs(self) -> dict:
        if self.video_writer is not None:
            self.video_writer.release()
            self.video_writer = None

        # Générer le numéro d'ordre pour les logs
        log_order = self._get_next_order_number()
        prefix = f"{self.group_id}_{self.scene_name}_{log_order:03d}"

        # ── CSV principal — schéma exact ──────────────────────────────────────
        csv_path = self.output_dir / f"{prefix}_detections.csv"
        with open(csv_path, "w", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=CSV_COLUMNS)
            writer.writeheader()
            for row in self.detection_log:
                # S'assurer que toutes les colonnes sont présentes
                writer.writerow({col: row.get(col, "") for col in CSV_COLUMNS})

        # ── JSONL brut (optionnel, pour debug) ────────────────────────────────
        jsonl_path = self.output_dir / f"{prefix}_detections.jsonl"
        with open(jsonl_path, "w", encoding="utf-8") as f:
            for row in self.detection_log:
                f.write(json.dumps(row) + "\n")

        # ── Summary JSON ──────────────────────────────────────────────────────
        summary  = self.get_summary()
        sum_path = self.output_dir / f"{prefix}_summary.json"
        with open(sum_path, "w", encoding="utf-8") as f:
            json.dump(summary, f, indent=2)

        # ── Frame stats CSV (résumé par frame) ────────────────────────────────
        fstats_path = self.output_dir / f"{prefix}_frame_stats.csv"
        if self.frame_stats:
            with open(fstats_path, "w", newline="", encoding="utf-8") as f:
                writer = csv.DictWriter(
                    f, fieldnames=self.frame_stats[0].keys()
                )
                writer.writeheader()
                writer.writerows(self.frame_stats)

        return {
            "detections_csv":   str(csv_path),
            "detections_jsonl": str(jsonl_path),
            "summary":          str(sum_path),
            "frame_stats":      str(fstats_path),
            "annotated_video":  self.output_video_path or "",
        }

    # ── Summary ─────────────────────────────────────────────────────────────────
    def get_summary(self) -> dict:
        total_duration = self.frame_index / self.fps
        total_objects  = sum(self.count_per_class.values())

        buckets: dict = defaultdict(int)
        for stat in self.frame_stats:
            bucket = int(stat["timestamp"] // 10)
            buckets[bucket] += stat["detections"]
        temporal = [{"bucket_10s": k, "detections": v} for k, v in sorted(buckets.items())]

        return {
            "scene":                 self.scene_name,
            "group_id":              self.group_id,
            "video_name":            self.video_name,
            "session_id":            self.session_id,
            "processed_at":          datetime.now().isoformat(),
            "total_frames":          self.frame_index,
            "duration_sec":          round(total_duration, 2),
            "fps":                   round(self.fps, 2),
            "resolution":            [self.frame_width, self.frame_height],
            "selected_classes":      self.selected_classes,
            "total_unique_objects":  total_objects,
            "count_per_class":       dict(self.count_per_class),
            "annotated_video":       self.output_video_path or "",
            "temporal_distribution": temporal,
        }