"""SAM3 mask tracking for the demo, through `transformers`.

Points: clicks on the first frame -> Sam3TrackerVideoModel propagates one object.
Text:   a short noun phrase -> Sam3VideoModel finds every matching instance in
        every frame; their union is the mask.
"""

import numpy as np
import torch
from PIL import Image
from transformers import (
    Sam3TrackerVideoModel,
    Sam3TrackerVideoProcessor,
    Sam3VideoModel,
    Sam3VideoProcessor,
)

SAM3_ID = "facebook/sam3"
DTYPE = torch.bfloat16


class Sam3Tracker:
    def __init__(self, device="cuda", token=None):
        self.device = device
        self.tracker = Sam3TrackerVideoModel.from_pretrained(SAM3_ID, token=token).to(device, dtype=DTYPE).eval()
        self.tracker_proc = Sam3TrackerVideoProcessor.from_pretrained(SAM3_ID, token=token)
        self.video_model = Sam3VideoModel.from_pretrained(SAM3_ID, token=token).to(device, dtype=DTYPE).eval()
        self.video_proc = Sam3VideoProcessor.from_pretrained(SAM3_ID, token=token)

    # ---------------------------------------------------------------- points
    def _point_session(self, frames, points, labels):
        sess = self.tracker_proc.init_video_session(
            video=[Image.fromarray(f) for f in frames], inference_device=self.device, dtype=DTYPE)
        self.tracker_proc.add_inputs_to_inference_session(
            inference_session=sess, frame_idx=0, obj_ids=1,
            input_points=[[[[float(x), float(y)] for x, y in points]]],
            input_labels=[[[int(l) for l in labels]]],
        )
        return sess

    def _to_mask(self, pred_masks, sess):
        m = self.tracker_proc.post_process_masks(
            [pred_masks], original_sizes=[[sess.video_height, sess.video_width]], binarize=True)[0]
        return (m[0, 0].cpu().numpy() > 0).astype(np.uint8) * 255

    @torch.inference_mode()
    def first_frame_mask(self, frame, points, labels):
        sess = self._point_session([frame], points, labels)
        out = self.tracker(inference_session=sess, frame_idx=0)
        return self._to_mask(out.pred_masks, sess)

    @torch.inference_mode()
    def track_points(self, frames, points, labels):
        sess = self._point_session(frames, points, labels)
        self.tracker(inference_session=sess, frame_idx=0)
        masks = [None] * len(frames)
        for out in self.tracker.propagate_in_video_iterator(sess):
            masks[out.frame_idx] = self._to_mask(out.pred_masks, sess)
        empty = np.zeros(frames[0].shape[:2], np.uint8)
        return [m if m is not None else empty for m in masks]

    # ------------------------------------------------------------------ text
    @torch.inference_mode()
    def track_text(self, frames, text):
        sess = self.video_proc.init_video_session(
            video=[Image.fromarray(f) for f in frames], inference_device=self.device, dtype=DTYPE)
        sess = self.video_proc.add_text_prompt(inference_session=sess, text=text)
        h, w = frames[0].shape[:2]
        masks = [np.zeros((h, w), np.uint8) for _ in frames]
        for out in self.video_model.propagate_in_video_iterator(inference_session=sess, max_frame_num_to_track=len(frames)):
            res = self.video_proc.postprocess_outputs(sess, out)
            m = res["masks"]
            if m is not None and len(m) > 0:
                masks[out.frame_idx] = (m.any(0).cpu().numpy() > 0).astype(np.uint8) * 255
        return masks