"""SegEarth-OV3 segmentor — standalone, no mmseg/mmcv/mmengine dependencies."""

import os
import torch
from torch import nn
import torch.nn.functional as F
from dataclasses import dataclass, field
from typing import Any
from PIL import Image

from sam3 import build_sam3_image_model
from sam3.model.sam3_image_processor import Sam3Processor

_DEFAULT_BPE_PATH = os.path.join(
    os.path.dirname(__file__), "sam3", "assets", "bpe_simple_vocab_16e6.txt.gz",
)


# ---------------------------------------------------------------------------
# Lightweight replacements for mmengine/mmseg data structures
# ---------------------------------------------------------------------------

@dataclass
class PixelData:
    """Minimal replacement for mmengine.structures.PixelData."""
    data: torch.Tensor = None


@dataclass
class SegDataSample:
    """Minimal replacement for mmseg.structures.SegDataSample."""
    metainfo: dict = field(default_factory=dict)
    pred_sem_seg: PixelData = field(default_factory=PixelData)
    seg_logits: PixelData = field(default_factory=PixelData)

    def set_metainfo(self, meta: dict) -> None:
        self.metainfo.update(meta)

    def set_data(self, data: dict) -> None:
        for k, v in data.items():
            setattr(self, k, v)


# ---------------------------------------------------------------------------
# Segmentor
# ---------------------------------------------------------------------------

class SegEarthOV3Segmentation(nn.Module):
    def __init__(self, classname_path,
                 device=torch.device('cuda'),
                 prob_thd=0.0,
                 bg_idx=0,
                 slide_stride=0,
                 slide_crop=0,
                 confidence_threshold=0.5,
                 use_sem_seg=True,
                 use_presence_score=True,
                 use_transformer_decoder=True,
                 bpe_path=None,
                 checkpoint_path="weights/sam3/sam3.pt",
                 **kwargs):
        super().__init__()

        self.device = device
        # Initialize SAM3 model
        if bpe_path is None:
            bpe_path = _DEFAULT_BPE_PATH
        model = build_sam3_image_model(
            bpe_path=bpe_path,
            checkpoint_path=checkpoint_path,
            device=str(device),
        )
        self.processor = Sam3Processor(model, confidence_threshold=confidence_threshold, device=device)
        self.query_words, self.query_idx = get_cls_idx(classname_path)
        self.num_cls = max(self.query_idx) + 1
        self.num_queries = len(self.query_idx)
        self.query_idx = torch.Tensor(self.query_idx).to(torch.int64).to(device)

        self.prob_thd = prob_thd
        self.bg_idx = bg_idx
        self.slide_stride = slide_stride
        self.slide_crop = slide_crop
        self.confidence_threshold = confidence_threshold
        self.use_sem_seg = use_sem_seg
        self.use_presence_score = use_presence_score
        self.use_transformer_decoder = use_transformer_decoder

    def _inference_single_view(self, image):
        """Inference on a single PIL image or crop patch."""
        w, h = image.size
        seg_logits = torch.zeros((self.num_queries, h, w), device=self.device)

        with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            inference_state = self.processor.set_image(image)

            for query_idx, query_word in enumerate(self.query_words):
                self.processor.reset_all_prompts(inference_state)
                inference_state = self.processor.set_text_prompt(state=inference_state, prompt=query_word)

                if self.use_transformer_decoder:
                    if inference_state['masks_logits'].shape[0] > 0:
                        inst_len = inference_state['masks_logits'].shape[0]
                        for inst_id in range(inst_len):
                            instance_logits = inference_state['masks_logits'][inst_id].squeeze()
                            instance_score = inference_state['object_score'][inst_id]

                            if instance_logits.shape != (h, w):
                                instance_logits = F.interpolate(
                                    instance_logits.view(1, 1, *instance_logits.shape),
                                    size=(h, w),
                                    mode='bilinear',
                                    align_corners=False
                                ).squeeze()

                            seg_logits[query_idx] = torch.max(seg_logits[query_idx], instance_logits * instance_score)

                if self.use_sem_seg:
                    semantic_logits = inference_state['semantic_mask_logits']
                    if semantic_logits.shape != (h, w):
                        semantic_logits = F.interpolate(
                            semantic_logits,
                            size=(h, w),
                            mode='bilinear',
                            align_corners=False
                        ).squeeze()

                    seg_logits[query_idx] = torch.max(seg_logits[query_idx], semantic_logits)

                if self.use_presence_score:
                    seg_logits[query_idx] = seg_logits[query_idx] * inference_state["presence_score"]

        return seg_logits

    def slide_inference(self, image, stride, crop_size):
        """Inference by sliding-window with overlap using PIL cropping."""
        w_img, h_img = image.size

        if isinstance(stride, int):
            stride = (stride, stride)
        if isinstance(crop_size, int):
            crop_size = (crop_size, crop_size)

        h_stride, w_stride = stride
        h_crop, w_crop = crop_size

        preds = torch.zeros((self.num_queries, h_img, w_img), device=self.device)
        count_mat = torch.zeros((1, h_img, w_img), device=self.device)

        h_grids = max(h_img - h_crop + h_stride - 1, 0) // h_stride + 1
        w_grids = max(w_img - w_crop + w_stride - 1, 0) // w_stride + 1

        for h_idx in range(h_grids):
            for w_idx in range(w_grids):
                y1 = h_idx * h_stride
                x1 = w_idx * w_stride
                y2 = min(y1 + h_crop, h_img)
                x2 = min(x1 + w_crop, w_img)

                y1 = max(y2 - h_crop, 0)
                x1 = max(x2 - w_crop, 0)

                crop_img = image.crop((x1, y1, x2, y2))
                crop_seg_logit = self._inference_single_view(crop_img)

                preds[:, y1:y2, x1:x2] += crop_seg_logit
                count_mat[:, y1:y2, x1:x2] += 1

        assert (count_mat == 0).sum() == 0, "Error: Sparse sliding window coverage."

        preds = preds / count_mat
        return preds

    def predict(self, inputs, data_samples):
        """Run segmentation prediction.

        Args:
            inputs: [B, 3, H, W] tensor (placeholder — actual image read from img_path).
            data_samples: list[SegDataSample] with metainfo containing
                          'img_path' and 'ori_shape'.

        Returns:
            list[SegDataSample] with pred_sem_seg and seg_logits filled.
        """
        if data_samples is not None:
            batch_img_metas = [ds.metainfo for ds in data_samples]
        else:
            batch_img_metas = [
                dict(
                    ori_shape=inputs.shape[2:],
                    img_shape=inputs.shape[2:],
                    pad_shape=inputs.shape[2:],
                    padding_size=[0, 0, 0, 0])
            ] * inputs.shape[0]

        for i, meta in enumerate(batch_img_metas):
            image_path = meta.get('img_path')
            image = Image.open(image_path).convert('RGB')
            ori_shape = meta['ori_shape']

            if self.slide_crop > 0 and (self.slide_crop < image.size[0] or self.slide_crop < image.size[1]):
                seg_logits = self.slide_inference(image, self.slide_stride, self.slide_crop)
            else:
                seg_logits = self._inference_single_view(image)

            if seg_logits.shape[-2:] != ori_shape:
                seg_logits = F.interpolate(
                    seg_logits.unsqueeze(0),
                    size=ori_shape,
                    mode='bilinear',
                    align_corners=False
                ).squeeze(0)

            # Post-processing: merge synonym queries into classes.
            if self.num_cls != self.num_queries:
                seg_logits = seg_logits.unsqueeze(0)
                cls_index = nn.functional.one_hot(self.query_idx)
                cls_index = cls_index.T.view(self.num_cls, len(self.query_idx), 1, 1)
                seg_logits = (seg_logits * cls_index).max(1)[0]

            seg_pred = torch.argmax(seg_logits, dim=0)

            # Apply probability threshold.
            max_vals = seg_logits.max(0)[0]
            seg_pred[max_vals < self.prob_thd] = self.bg_idx

            data_samples[i].set_data({
                'seg_logits': PixelData(data=seg_logits),
                'pred_sem_seg': PixelData(data=seg_pred.unsqueeze(0)),
            })

        return data_samples

    @staticmethod
    def _slice_backbone_out(backbone_out, idx):
        """Slice batched backbone output to extract single-image features."""
        sliced = {}
        for key, val in backbone_out.items():
            if key == "sam2_backbone_out":
                sliced[key] = SegEarthOV3Segmentation._slice_backbone_out(val, idx) if val is not None else None
            elif isinstance(val, torch.Tensor):
                sliced[key] = val[idx:idx + 1]
            elif isinstance(val, list):
                sliced[key] = [t[idx:idx + 1] if isinstance(t, torch.Tensor) else t for t in val]
            else:
                sliced[key] = val
        return sliced

    def _cache_text_embeddings(self):
        """Pre-compute and cache text embeddings for all prompts (run once)."""
        if hasattr(self, '_text_cache'):
            return
        self._text_cache = []
        with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            for query_word in self.query_words:
                text_out = self.processor.model.backbone.forward_text(
                    [query_word], device=self.device,
                )
                # Deep-clone tensors so they survive state mutations.
                cached = {k: v.clone() if isinstance(v, torch.Tensor) else v
                          for k, v in text_out.items()}
                self._text_cache.append(cached)

    def _apply_cached_text(self, state, query_idx):
        """Inject cached text embeddings and run grounding decoder."""
        cached = self._text_cache[query_idx]
        state["backbone_out"].update(
            {k: v.clone() if isinstance(v, torch.Tensor) else v
             for k, v in cached.items()}
        )
        if "geometric_prompt" not in state:
            state["geometric_prompt"] = self.processor.model._get_dummy_prompt()
        return self.processor._forward_grounding(state)

    def predict_pil_batch(self, images: list[Image.Image],
                          ori_shapes: list[tuple[int, int]] | None = None):
        """Batch prediction: single backbone pass for all images, then per-image grounding.

        Args:
            images: list of PIL RGB images.
            ori_shapes: list of (H, W) target output shapes. Defaults to each image's size.

        Returns:
            list of (seg_pred, seg_logits) tuples.
        """
        B = len(images)
        if ori_shapes is None:
            ori_shapes = [(img.size[1], img.size[0]) for img in images]

        # Cache text embeddings once (amortized across all batches).
        self._cache_text_embeddings()

        results = []

        with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            # Single backbone pass for all images
            batch_state = self.processor.set_image_batch(images)

            for i in range(B):
                w, h = images[i].size
                seg_logits = torch.zeros((self.num_queries, h, w), device=self.device)

                # Create single-image state by slicing backbone features
                state = {
                    "original_height": batch_state["original_heights"][i],
                    "original_width": batch_state["original_widths"][i],
                    "backbone_out": self._slice_backbone_out(batch_state["backbone_out"], i),
                }

                for query_idx in range(len(self.query_words)):
                    self.processor.reset_all_prompts(state)
                    state = self._apply_cached_text(state, query_idx)

                    if self.use_transformer_decoder:
                        if state['masks_logits'].shape[0] > 0:
                            inst_len = state['masks_logits'].shape[0]
                            for inst_id in range(inst_len):
                                instance_logits = state['masks_logits'][inst_id].squeeze()
                                instance_score = state['object_score'][inst_id]

                                if instance_logits.shape != (h, w):
                                    instance_logits = F.interpolate(
                                        instance_logits.view(1, 1, *instance_logits.shape),
                                        size=(h, w),
                                        mode='bilinear',
                                        align_corners=False
                                    ).squeeze()

                                seg_logits[query_idx] = torch.max(seg_logits[query_idx], instance_logits * instance_score)

                    if self.use_sem_seg:
                        semantic_logits = state['semantic_mask_logits']
                        if semantic_logits.shape != (h, w):
                            semantic_logits = F.interpolate(
                                semantic_logits,
                                size=(h, w),
                                mode='bilinear',
                                align_corners=False
                            ).squeeze()

                        seg_logits[query_idx] = torch.max(seg_logits[query_idx], semantic_logits)

                    if self.use_presence_score:
                        seg_logits[query_idx] = seg_logits[query_idx] * state["presence_score"]

                # Post-processing
                ori_shape = ori_shapes[i]
                if seg_logits.shape[-2:] != ori_shape:
                    seg_logits = F.interpolate(
                        seg_logits.unsqueeze(0), size=ori_shape,
                        mode='bilinear', align_corners=False,
                    ).squeeze(0)

                if self.num_cls != self.num_queries:
                    seg_logits = seg_logits.unsqueeze(0)
                    cls_index = nn.functional.one_hot(self.query_idx)
                    cls_index = cls_index.T.view(self.num_cls, len(self.query_idx), 1, 1)
                    seg_logits = (seg_logits * cls_index).max(1)[0]

                seg_pred = torch.argmax(seg_logits, dim=0)
                max_vals = seg_logits.max(0)[0]
                seg_pred[max_vals < self.prob_thd] = self.bg_idx

                results.append((seg_pred.to(torch.uint8), seg_logits))

        return results

    def predict_pil(self, image: Image.Image, ori_shape: tuple[int, int] | None = None):
        """Convenience method: predict from a PIL image directly (no file path needed).

        Args:
            image: PIL RGB image.
            ori_shape: (H, W) target output shape. Defaults to image size.

        Returns:
            seg_pred: [H, W] uint8 tensor of class indices.
            seg_logits: [num_cls, H, W] float tensor.
        """
        if ori_shape is None:
            ori_shape = (image.size[1], image.size[0])  # PIL is (W, H)

        if self.slide_crop > 0 and (self.slide_crop < image.size[0] or self.slide_crop < image.size[1]):
            seg_logits = self.slide_inference(image, self.slide_stride, self.slide_crop)
        else:
            seg_logits = self._inference_single_view(image)

        if seg_logits.shape[-2:] != ori_shape:
            seg_logits = F.interpolate(
                seg_logits.unsqueeze(0), size=ori_shape,
                mode='bilinear', align_corners=False,
            ).squeeze(0)

        if self.num_cls != self.num_queries:
            seg_logits = seg_logits.unsqueeze(0)
            cls_index = nn.functional.one_hot(self.query_idx)
            cls_index = cls_index.T.view(self.num_cls, len(self.query_idx), 1, 1)
            seg_logits = (seg_logits * cls_index).max(1)[0]

        seg_pred = torch.argmax(seg_logits, dim=0)
        max_vals = seg_logits.max(0)[0]
        seg_pred[max_vals < self.prob_thd] = self.bg_idx

        return seg_pred.to(torch.uint8), seg_logits


def get_cls_idx(path):
    with open(path, 'r') as f:
        name_sets = f.readlines()
    num_cls = len(name_sets)

    class_names, class_indices = [], []
    for idx in range(num_cls):
        names_i = name_sets[idx].split(',')
        names_i = [i.strip() for i in names_i]
        class_names += names_i
        class_indices += [idx for _ in range(len(names_i))]
    class_names = [item.replace('\n', '') for item in class_names]
    return class_names, class_indices
