"""
modules/planarian_tracker.py

Détection et suivi multi-individus de planaires dans un tube.
Supporte de 1 à MAX_PLANARIANS planaires par tube.

Stratégie :
    - Soustraction de fond MOG2 (léger sur Raspberry Pi 4)
    - Détection de tous les contours valides (surface >= min_area_px)
    - Association frame-à-frame par distance euclidienne minimale
      via algorithme hongrois (scipy.optimize.linear_sum_assignment)
    - Un état inter-frame indépendant par individu (PlanarianState)
    - Retourne une liste de résultats, un par individu suivi

Created on 25 avr. 2026
@author: denis
"""

import cv2
import numpy as np

from scipy.optimize import linear_sum_assignment
import logging


logger = logging.getLogger(__name__)

# Nombre maximum de planaires suivis simultanément par tube
MAX_PLANARIANS = 10

# Distance maximale en pixels entre deux positions consécutives
# pour qu'une association soit acceptée (évite les sauts aberrants)
MAX_ASSOC_DIST_PX = 80

# Couleurs d'annotation par individu (BGR)
# Cycle automatique si plus de planaires que de couleurs
INDIVIDUAL_COLORS = [
    (255, 255,   0),  # cyan
    (  0, 165, 255),  # orange
    (255,   0, 255),  # magenta
    (  0, 255, 255),  # jaune
    (128,   0, 255),  # violet
    (  0, 255, 128),  # vert clair
    (255, 128,   0),  # bleu clair
    (  0, 128, 255),  # orange foncé
    (128, 255,   0),  # vert-jaune
    (255,   0, 128),  # rose
]

# Couleur du contour principal (individu le plus grand)
COLOR_LARGEST  = (255, 255,   0)   # cyan
COLOR_OTHER    = (  0, 255,   0)   # vert
COLOR_CENTER   = (  0,   0, 255)   # rouge


# ---------------------------------------------------------------------------
# État inter-frame d'un individu
# ---------------------------------------------------------------------------

class PlanarianState:
    """
    Mémorise la position et le timestamp de la dernière détection
    pour un planaire individuel.

    Un PlanarianState par slot (index 0 à max_planarians-1).
    Quand un slot n'est pas associé à un contour sur plusieurs frames
    consécutives, il est marqué comme perdu (lost).
    """

    # Nombre de frames sans détection avant de considérer l'individu perdu
    MAX_LOST_FRAMES = 5

    def __init__(self, idx: int):
        """
        Args:
            idx : index de l'individu (0-based)
        """
        self.idx        = idx
        self.cx         = None
        self.cy         = None
        self.ts         = None
        self.lost       = 0     # compteur de frames sans détection
        self.active     = False # vrai si l'individu a été détecté au moins une fois

    def update(self, cx: int, cy: int, ts: float):
        """
        Met à jour la position suite à une association réussie.

        Args:
            cx, cy : position du centre de masse en pixels
            ts     : timestamp de la frame
        """
        self.cx     = cx
        self.cy     = cy
        self.ts     = ts
        self.lost   = 0
        self.active = True

    def mark_lost(self):
        """Incrémente le compteur de perte — appelé quand aucun contour n'est associé."""
        self.lost += 1

    @property
    def is_lost(self) -> bool:
        """Retourne True si l'individu est considéré perdu (trop de frames sans détection)."""
        return self.lost >= self.MAX_LOST_FRAMES

    def compute_speed(self, cx: int, cy: int, ts: float, tube_axis: str) -> tuple:
        """
        Calcule la vitesse instantanée depuis la position précédente.

        Args:
            cx, cy    : position courante en pixels
            ts        : timestamp courant
            tube_axis : "vertical" ou "horizontal"

        Returns:
            tuple (speed_px_s, axial_speed) ou (0.0, 0.0) si état vide
        """
        if self.cx is None or self.ts is None:
            return 0.0, 0.0

        dt = ts - self.ts
        if dt <= 0:
            return 0.0, 0.0

        dx          = cx - self.cx
        dy          = cy - self.cy
        speed_px_s  = float(np.sqrt(dx**2 + dy**2) / dt)
        axial_speed = float((dy / dt) if tube_axis == "vertical" else (dx / dt))

        return speed_px_s, axial_speed

    def reset(self):
        """Réinitialise l'état de cet individu."""
        self.cx     = None
        self.cy     = None
        self.ts     = None
        self.lost   = 0
        self.active = False


# ---------------------------------------------------------------------------
# Tracker multi-individus
# ---------------------------------------------------------------------------

class PlanarianTracker:
    """
    Détection et suivi multi-individus de planaires dans un tube.

    Instancié une fois par caméra active, réutilisé frame à frame.
    Utilise la soustraction de fond MOG2 — léger sur Raspberry Pi 4.
    Association frame-à-frame par algorithme hongrois (distance euclidienne).

    Usage :
        tracker = PlanarianTracker(tube_axis="vertical", max_planarians=3)
        while capturing:
            frame_out, results = tracker.process(frame, ts)
            # results : liste de dicts, un par individu détecté
            for r in results:
                metrics.update(r, planarian_id=r["planarian_id"])
    """

    # Nombre de frames d'initialisation MOG2 ignorées (fond non appris)
    WARMUP_FRAMES = 10

    def __init__(
        self,
        tube_axis:          str   = "vertical",
        min_area_px:        int   = 20,
        max_area_ratio:     float = 0.10,
        max_planarians:     int   = 1,
        merge_kernel_size:  int   = 15,
        min_contour_dist_px:int   = 40,
        draw_contours:       bool = True,
    ):
        """
        Args:
            tube_axis      : axe principal du tube — "vertical" (cy) ou "horizontal" (cx)
            min_area_px    : surface minimale d'un contour pour être considéré valide (px²)
            max_area_ratio : surface maximale d'un contour en fraction de la frame (défaut 10%)
                             filtre les faux positifs du fond non encore appris par MOG2
            max_planarians      : nombre maximum de planaires à suivre simultanément (1-10)
            merge_kernel_size   : taille du kernel elliptique de fusion des fragments (px).
                                  Régler ≈ largeur du planaire en pixels. Défaut : 15.
            min_contour_dist_px : distance min entre deux contours pour les considérer
                                  comme individus distincts. Défaut : 40px.
        """
        self.tube_axis      = tube_axis
        self.min_area_px    = min_area_px
        self.max_area_ratio = max_area_ratio
        self.max_planarians = max(1, min(max_planarians, MAX_PLANARIANS))
        self.draw_contours  = draw_contours

        # Un état inter-frame par slot individu
        self._states = [PlanarianState(i) for i in range(self.max_planarians)]

        # Taille du kernel de fusion morphologique (pixels) —
        # doit être proche de la largeur du planaire en pixels.
        # Trop petit : fragments non fusionnés → IDs multiples.
        # Trop grand : deux planaires proches fusionnés en un seul.
        self.merge_kernel_size    = merge_kernel_size

        # Distance minimale en pixels entre deux contours distincts.
        # En-dessous : le plus petit est considéré comme fragment du plus grand.
        self.min_contour_dist_px  = min_contour_dist_px

        # Soustracteur de fond adaptatif MOG2
        self._bg_sub = self._make_bg_sub()

        # Compteur de frames d'initialisation — MOG2 retourne du bruit
        # pendant les premières WARMUP_FRAMES frames
        self._warmup_count = 0

    @staticmethod
    def _make_bg_sub():
        """Crée et retourne un soustracteur de fond MOG2."""
        return cv2.createBackgroundSubtractorMOG2(
            history      = 50,
            varThreshold = 25,
            detectShadows= False,
        )

    def reset(self):
        """
        Réinitialise l'état inter-frame complet.
        À appeler lors du changement de puits.
        """
        for s in self._states:
            s.reset()
        self._bg_sub       = self._make_bg_sub()
        self._warmup_count = 0
        # Les paramètres morphologiques (merge_kernel_size, min_contour_dist_px)
        # sont conservés — ils ne dépendent pas du puits

    # ------------------------------------------------------------------ #
    # Interface principale
    # ------------------------------------------------------------------ #

    def process(self, frame: np.ndarray, ts: float) -> tuple:
        """
        Analyse une frame, associe les contours aux individus connus,
        dessine les annotations et retourne les métriques.

        Args:
            frame : image BGR (numpy array)
            ts    : timestamp de la frame (float, secondes epoch)

        Returns:
            tuple (frame_annotée, results)

            frame_annotée : copie BGR avec contours, croix et textes
            results       : liste de dicts — un dict par planaire actif détecté.
                            Chaque dict contient :
                                planarian_id  int     index de l'individu (0-based)
                                detected      bool    True si détecté cette frame
                                cx, cy        int     centre de masse en pixels
                                area_px       int     surface du contour (px²)
                                speed_px_s    float   vitesse totale (px/s)
                                axial_speed   float   vitesse axiale (px/s)
                                axial_pos     float   position axiale normalisée (0-1)
                                timestamp     float   ts de la frame
        """
        frame_out = frame.copy()
        h, w      = frame.shape[:2]

        # --- Extraction du premier plan ---
        gray    = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        fg_mask = self._bg_sub.apply(gray)

        # --- Morphologie : fusion des fragments du corps ---
        # Un planaire ondulant est souvent segmenté en plusieurs contours
        # (tête, milieu, queue). La dilatation fusionne les fragments proches
        # avant la détection des contours.
        # Kernel 3×3 : supprime le bruit fin (OPEN)
        # Kernel merge_kernel : fusionne les fragments du corps (CLOSE + DILATE)
        noise_kernel = np.ones((3, 3), np.uint8)
        merge_kernel  = cv2.getStructuringElement(
            cv2.MORPH_ELLIPSE, (self.merge_kernel_size, self.merge_kernel_size)
        )
        fg_mask = cv2.morphologyEx(fg_mask, cv2.MORPH_OPEN,   noise_kernel)
        fg_mask = cv2.morphologyEx(fg_mask, cv2.MORPH_CLOSE,  merge_kernel)
        fg_mask = cv2.dilate(fg_mask, merge_kernel, iterations=1)

        # Warmup MOG2 : les premières WARMUP_FRAMES frames retournent du bruit
        # (fond non encore appris) — on les alimente mais on ne détecte rien
        self._warmup_count += 1
        if self._warmup_count <= self.WARMUP_FRAMES:
            return frame_out, []

        contours, _ = cv2.findContours(
            fg_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )

        # Surface maximale admissible : fraction de la frame
        # Filtre les faux positifs du fond (contours couvrant toute l'image)
        max_area_px = h * w * self.max_area_ratio

        # Filtrage : surface min ET surface max, triés par surface décroissante
        valid = sorted(
            [c for c in contours
             if self.min_area_px <= cv2.contourArea(c) <= max_area_px],
            key=cv2.contourArea,
            reverse=True,
        )

        # --- Suppression des fragments résiduels trop proches ---
        # Si plusieurs contours valides ont leur centre à moins de
        # min_contour_dist_px les uns des autres, seul le plus grand est conservé.
        # Évite les cas où la fusion morphologique est incomplète.
        filtered = []
        for c in valid:
            M = cv2.moments(c)
            if M["m00"] == 0:
                continue
            cx_c = int(M["m10"] / M["m00"])
            cy_c = int(M["m01"] / M["m00"])
            too_close = False
            for kept in filtered:
                Mk = cv2.moments(kept)
                if Mk["m00"] == 0:
                    continue
                cx_k = int(Mk["m10"] / Mk["m00"])
                cy_k = int(Mk["m01"] / Mk["m00"])
                dist = np.sqrt((cx_c - cx_k)**2 + (cy_c - cy_k)**2)
                if dist < self.min_contour_dist_px:
                    too_close = True
                    break
            if not too_close:
                filtered.append(c)

        # Limiter au nombre maximum de planaires attendus
        valid = filtered[:self.max_planarians]

        # --- Calcul des centres de masse des contours détectés ---
        detections = []   # liste de (cx, cy, area, contour)
        for c in valid:
            M = cv2.moments(c)
            if M["m00"] == 0:
                continue
            cx   = int(M["m10"] / M["m00"])
            cy   = int(M["m01"] / M["m00"])
            area = cv2.contourArea(c)
            detections.append((cx, cy, int(area), c))

        # --- Association hongroise détections → slots individus ---
        assignments = self._hungarian_assign(detections)

        # --- Mise à jour des états et construction des résultats ---
        results = []

        for slot_idx, det_idx in assignments.items():
            state = self._states[slot_idx]

            if det_idx is None:
                # Aucune détection associée à ce slot
                state.mark_lost()
                if state.active and not state.is_lost:
                    # L'individu était suivi : on retourne un résultat "perdu"
                    results.append(self._lost_result(slot_idx, ts))
                continue

            cx, cy, area, contour = detections[det_idx]

            # Calcul de la vitesse depuis la position précédente
            speed_px_s, axial_speed = state.compute_speed(cx, cy, ts, self.tube_axis)

            axial_pos = (cy / h) if self.tube_axis == "vertical" else (cx / w)

            # Mise à jour de l'état
            state.update(cx, cy, ts)

            
            if self.draw_contours:
                
                # Annotation visuelle
                color = INDIVIDUAL_COLORS[slot_idx % len(INDIVIDUAL_COLORS)]
                cv2.drawContours(frame_out, [contour], -1, color, 2)        
                self._draw_center(frame_out, cx, cy, slot_idx, speed_px_s, axial_pos, color)

            results.append({
                "planarian_id": slot_idx,
                "detected":     True,
                "cx":           cx,
                "cy":           cy,
                "area_px":      area,
                "speed_px_s":   round(speed_px_s,  3),
                "axial_speed":  round(axial_speed,  3),
                "axial_pos":    round(axial_pos,    4),
                "timestamp":    ts,
            })

        # Marquer les slots non présents dans les assignments comme perdus
        assigned_slots = set(assignments.keys())
        for state in self._states:
            if state.idx not in assigned_slots:
                state.mark_lost()
        return frame_out, results

    # ------------------------------------------------------------------ #
    # Association hongroise
    # ------------------------------------------------------------------ #

    def _hungarian_assign(self, detections: list) -> dict:
        """
        Associe les détections courantes aux slots individus connus
        via l'algorithme hongrois (coût = distance euclidienne).

        Contrainte : une association n'est acceptée que si la distance
        est inférieure à MAX_ASSOC_DIST_PX (évite les sauts aberrants).

        Args:
            detections : liste de (cx, cy, area, contour)

        Returns:
            dict {slot_idx: det_idx | None}
            det_idx = None si aucune détection assignée à ce slot
        """
        n_slots = self.max_planarians
        n_dets  = len(detections)

        if n_dets == 0:
            # Aucune détection : tous les slots sont "perdus"
            return {i: None for i in range(n_slots)}

        # Slots actifs (déjà vus au moins une fois et non perdus)
        active_slots = [s for s in self._states if s.active and not s.is_lost]

        if not active_slots:
            # Première frame ou tous perdus : attribution séquentielle simple
            assignment = {}
            for i in range(n_slots):
                assignment[i] = i if i < n_dets else None
            return assignment

        # --- Construction de la matrice de coût (distance euclidienne) ---
        cost = np.full((len(active_slots), n_dets), fill_value=1e6)

        for si, state in enumerate(active_slots):
            for di, (cx, cy, _, _) in enumerate(detections):
                dist = np.sqrt((cx - state.cx)**2 + (cy - state.cy)**2)
                cost[si, di] = dist

        # --- Algorithme hongrois ---
        row_ind, col_ind = linear_sum_assignment(cost)

        # Construire le dict d'association
        assignment = {i: None for i in range(n_slots)}

        assigned_dets = set()
        for ri, ci in zip(row_ind, col_ind):
            if cost[ri, ci] <= MAX_ASSOC_DIST_PX:
                slot_idx = active_slots[ri].idx
                assignment[slot_idx] = ci
                assigned_dets.add(ci)

        # --- Nouvelles détections non assignées → slots inactifs libres ---
        free_slots = [s for s in self._states if not s.active or s.is_lost]
        new_dets   = [di for di in range(n_dets) if di not in assigned_dets]

        for state, det_idx in zip(free_slots, new_dets):
            assignment[state.idx] = det_idx

        return assignment

    # ------------------------------------------------------------------ #
    # Dessin des annotations
    # ------------------------------------------------------------------ #

    def _draw_center(
        self,
        frame:      np.ndarray,
        cx:         int,
        cy:         int,
        idx:        int,
        speed_px_s: float,
        axial_pos:  float,
        color:      tuple,
    ):
        """
        Dessine la croix, le cercle et le label de vitesse/position
        pour un individu.

        Args:
            frame      : image à annoter (en place)
            cx, cy     : centre de masse en pixels
            idx        : index de l'individu
            speed_px_s : vitesse en px/s
            axial_pos  : position axiale normalisée
            color      : couleur BGR de l'individu
        """
        cross = 8
        cv2.line(frame, (cx - cross, cy), (cx + cross, cy), COLOR_CENTER, 1)
        cv2.line(frame, (cx, cy - cross), (cx, cy + cross), COLOR_CENTER, 1)
        cv2.circle(frame, (cx, cy), 12, color, 1)

        # Badge numéro individu
        cv2.circle(frame, (cx + 14, cy - 14), 8, color, -1)
        cv2.putText(
            frame, str(idx),
            (cx + 10, cy - 10),
            cv2.FONT_HERSHEY_SIMPLEX, 0.35, (0, 0, 0), 1, cv2.LINE_AA,
        )

        # Texte vitesse + position axiale
        label = (
            f"#{idx} v={speed_px_s:.1f}px/s ax={axial_pos:.2f}"
            if speed_px_s > 0
            else f"#{idx} ax={axial_pos:.2f}"
        )
        cv2.putText(
            frame, label,
            (max(cx - 60, 0), max(cy - 22, 12)),
            cv2.FONT_HERSHEY_SIMPLEX, 0.35, (255, 255, 255), 1, cv2.LINE_AA,
        )

    # ------------------------------------------------------------------ #
    # Résultats vides / perdus
    # ------------------------------------------------------------------ #

    def _lost_result(self, planarian_id: int, ts: float) -> dict:
        """
        Retourne un résultat pour un individu temporairement non détecté.

        Args:
            planarian_id : index de l'individu
            ts           : timestamp de la frame courante

        Returns:
            dict avec detected=False et les dernières coordonnées connues
        """
        state = self._states[planarian_id]
        return {
            "planarian_id": planarian_id,
            "detected":     False,
            "cx":           state.cx or 0,
            "cy":           state.cy or 0,
            "area_px":      0,
            "speed_px_s":   0.0,
            "axial_speed":  0.0,
            "axial_pos":    0.0,
            "timestamp":    ts,
        }
