from __future__ import annotations

import functools
import importlib
import sys
import tempfile
import types
from pathlib import Path
from typing import Any


def _torchvision_nms():
    from torchvision.ops import nms

    return nms


def _torchvision_roi_align():
    from torchvision.ops import roi_align

    return roi_align


def _torchvision_roi_align_module():
    from torchvision.ops import RoIAlign

    return RoIAlign


def _torchvision_roi_pool_module():
    from torchvision.ops import RoIPool

    return RoIPool


def _multiscale_deformable_attention_class():
    import torch.nn as nn

    class MultiScaleDeformableAttention(nn.Module):
        """Import-only fallback for mmdet registries when running with mmcv-lite."""

        def __init__(self, *args: Any, **kwargs: Any) -> None:
            super().__init__()

        def init_weights(self) -> None:
            return None

        def forward(self, *args: Any, **kwargs: Any) -> Any:
            raise RuntimeError(
                "MultiScaleDeformableAttention requires full mmcv with compiled ops. "
                "The EgoForce demo uses RTMDet and should not execute this layer."
            )

    return MultiScaleDeformableAttention


def _unsupported_module_class(name: str):
    import torch.nn as nn

    class UnsupportedMMCVOp(nn.Module):
        def __init__(self, *args: Any, **kwargs: Any) -> None:
            super().__init__()

        def init_weights(self) -> None:
            return None

        def forward(self, *args: Any, **kwargs: Any) -> Any:
            raise RuntimeError(f"{name} requires full mmcv with compiled ops and is not used by the EgoForce RTMDet demo.")

    UnsupportedMMCVOp.__name__ = name
    return UnsupportedMMCVOp


def _unsupported_function(name: str):
    def unsupported(*args: Any, **kwargs: Any) -> Any:
        raise RuntimeError(f"{name} requires full mmcv with compiled ops and is not used by the EgoForce RTMDet demo.")

    unsupported.__name__ = name
    return unsupported


def _bbox_overlaps(
    bboxes1: Any,
    bboxes2: Any,
    mode: str = "iou",
    aligned: bool = False,
    offset: int = 0,
    eps: float = 1e-6,
) -> Any:
    import torch

    if bboxes1.numel() == 0 or bboxes2.numel() == 0:
        if aligned:
            return bboxes1.new_zeros((bboxes1.shape[0],))
        return bboxes1.new_zeros((bboxes1.shape[0], bboxes2.shape[0]))

    if aligned:
        lt = torch.maximum(bboxes1[:, :2], bboxes2[:, :2])
        rb = torch.minimum(bboxes1[:, 2:], bboxes2[:, 2:])
        wh = (rb - lt + offset).clamp(min=0)
        overlap = wh[:, 0] * wh[:, 1]
        area1 = (bboxes1[:, 2] - bboxes1[:, 0] + offset) * (bboxes1[:, 3] - bboxes1[:, 1] + offset)
        area2 = (bboxes2[:, 2] - bboxes2[:, 0] + offset) * (bboxes2[:, 3] - bboxes2[:, 1] + offset)
    else:
        lt = torch.maximum(bboxes1[:, None, :2], bboxes2[None, :, :2])
        rb = torch.minimum(bboxes1[:, None, 2:], bboxes2[None, :, 2:])
        wh = (rb - lt + offset).clamp(min=0)
        overlap = wh[..., 0] * wh[..., 1]
        area1 = ((bboxes1[:, 2] - bboxes1[:, 0] + offset) * (bboxes1[:, 3] - bboxes1[:, 1] + offset))[:, None]
        area2 = ((bboxes2[:, 2] - bboxes2[:, 0] + offset) * (bboxes2[:, 3] - bboxes2[:, 1] + offset))[None, :]

    if mode == "iof":
        union = area1
    elif mode == "giou":
        union = area1 + area2 - overlap
        if aligned:
            enclosed_lt = torch.minimum(bboxes1[:, :2], bboxes2[:, :2])
            enclosed_rb = torch.maximum(bboxes1[:, 2:], bboxes2[:, 2:])
        else:
            enclosed_lt = torch.minimum(bboxes1[:, None, :2], bboxes2[None, :, :2])
            enclosed_rb = torch.maximum(bboxes1[:, None, 2:], bboxes2[None, :, 2:])
        enclosed_wh = (enclosed_rb - enclosed_lt + offset).clamp(min=0)
        enclosed_area = enclosed_wh[..., 0] * enclosed_wh[..., 1]
        iou = overlap / union.clamp(min=eps)
        return iou - (enclosed_area - union) / enclosed_area.clamp(min=eps)
    else:
        union = area1 + area2 - overlap

    return overlap / union.clamp(min=eps)


def _nms(
    boxes: Any,
    scores: Any,
    iou_threshold: float,
    offset: int = 0,
    score_threshold: float = 0,
    max_num: int = -1,
) -> tuple[Any, Any]:
    import torch

    if boxes.numel() == 0 or scores.numel() == 0:
        keep = torch.empty((0,), dtype=torch.long, device=scores.device)
        dets = torch.cat((boxes.reshape(0, boxes.shape[-1]), scores.reshape(0, 1)), dim=1)
        return dets, keep

    if score_threshold > 0:
        valid = scores > score_threshold
        original_indices = torch.nonzero(valid, as_tuple=False).squeeze(1)
        filtered_boxes = boxes[valid]
        filtered_scores = scores[valid]
    else:
        original_indices = torch.arange(scores.numel(), device=scores.device)
        filtered_boxes = boxes
        filtered_scores = scores

    keep_local = _torchvision_nms()(filtered_boxes, filtered_scores, float(iou_threshold))
    if max_num > 0:
        keep_local = keep_local[:max_num]

    keep = original_indices[keep_local]
    dets = torch.cat((filtered_boxes[keep_local], filtered_scores[keep_local, None]), dim=1)
    return dets, keep


def _batched_nms(
    boxes: Any,
    scores: Any,
    idxs: Any,
    nms_cfg: dict[str, Any] | None,
    class_agnostic: bool = False,
) -> tuple[Any, Any]:
    import torch

    if boxes.numel() == 0 or scores.numel() == 0:
        keep = torch.empty((0,), dtype=torch.long, device=scores.device)
        dets = torch.cat((boxes.reshape(0, boxes.shape[-1]), scores.reshape(0, 1)), dim=1)
        return dets, keep

    if nms_cfg is None:
        order = scores.argsort(descending=True)
        return torch.cat((boxes[order], scores[order, None]), dim=1), order

    nms_cfg = dict(nms_cfg)
    iou_threshold = nms_cfg.pop("iou_threshold", nms_cfg.pop("iou_thr", 0.5))
    score_threshold = nms_cfg.pop("score_threshold", 0)
    max_num = nms_cfg.pop("max_num", -1)

    if class_agnostic:
        boxes_for_nms = boxes
    else:
        max_coordinate = boxes.max()
        offsets = idxs.to(boxes) * (max_coordinate + boxes.new_tensor(1))
        boxes_for_nms = boxes + offsets[:, None]

    if score_threshold > 0:
        valid = scores > score_threshold
        original_indices = torch.nonzero(valid, as_tuple=False).squeeze(1)
        boxes_for_nms = boxes_for_nms[valid]
        scores_for_nms = scores[valid]
    else:
        original_indices = torch.arange(scores.numel(), device=scores.device)
        scores_for_nms = scores

    keep_local = _torchvision_nms()(boxes_for_nms, scores_for_nms, float(iou_threshold))
    if max_num > 0:
        keep_local = keep_local[:max_num]

    keep = original_indices[keep_local]
    dets = torch.cat((boxes[keep], scores[keep, None]), dim=1)
    return dets, keep


def _nms_match(dets: Any, iou_threshold: float) -> list[Any]:
    """Pure PyTorch fallback for mmcv.ops.nms_match import paths."""
    import torch

    if dets.numel() == 0:
        return []

    boxes = dets[:, :4]
    scores = dets[:, 4]
    order = scores.argsort(descending=True)
    groups = []

    while order.numel() > 0:
        current = order[0]
        if order.numel() == 1:
            groups.append(current.reshape(1))
            break

        rest = order[1:]
        lt = torch.maximum(boxes[current, :2], boxes[rest, :2])
        rb = torch.minimum(boxes[current, 2:], boxes[rest, 2:])
        wh = (rb - lt).clamp(min=0)
        inter = wh[:, 0] * wh[:, 1]
        current_area = (boxes[current, 2] - boxes[current, 0]).clamp(min=0) * (
            boxes[current, 3] - boxes[current, 1]
        ).clamp(min=0)
        rest_area = (boxes[rest, 2] - boxes[rest, 0]).clamp(min=0) * (
            boxes[rest, 3] - boxes[rest, 1]
        ).clamp(min=0)
        iou = inter / (current_area + rest_area - inter).clamp(min=1e-6)
        matched = rest[iou > float(iou_threshold)]
        groups.append(torch.cat((current.reshape(1), matched)))
        order = rest[iou <= float(iou_threshold)]

    return groups


def _point_sample(input: Any, points: Any, align_corners: bool = False, **kwargs: Any) -> Any:
    import torch.nn.functional as F

    add_dim = False
    if points.dim() == 3:
        add_dim = True
        points = points.unsqueeze(2)

    output = F.grid_sample(input, points.mul(2).sub(1), align_corners=align_corners, **kwargs)
    if add_dim:
        output = output.squeeze(3)
    return output


def _rel_roi_point_to_rel_img_point(rois: Any, rel_roi_points: Any, img_shape: Any) -> Any:
    x1, y1, x2, y2 = rois[:, 1], rois[:, 2], rois[:, 3], rois[:, 4]
    roi_w = (x2 - x1).clamp(min=1)
    roi_h = (y2 - y1).clamp(min=1)
    img_h, img_w = img_shape[:2]
    rel_img_points = rel_roi_points.clone()
    rel_img_points[..., 0] = (x1[:, None] + rel_roi_points[..., 0] * roi_w[:, None]) / float(img_w)
    rel_img_points[..., 1] = (y1[:, None] + rel_roi_points[..., 1] * roi_h[:, None]) / float(img_h)
    return rel_img_points


def _sigmoid_focal_loss(
    pred: Any,
    target: Any,
    gamma: float = 2.0,
    alpha: float = 0.25,
    weight: Any = None,
    reduction: str = "mean",
) -> Any:
    import torch.nn.functional as F

    pred_sigmoid = pred.sigmoid()
    target = target.type_as(pred)
    pt = (1 - pred_sigmoid) * target + pred_sigmoid * (1 - target)
    focal_weight = (alpha * target + (1 - alpha) * (1 - target)) * pt.pow(gamma)
    loss = F.binary_cross_entropy_with_logits(pred, target, reduction="none") * focal_weight
    if weight is not None:
        loss = loss * weight
    if reduction == "sum":
        return loss.sum()
    if reduction == "mean":
        return loss.mean()
    return loss


def _ensure_module(module_name: str) -> types.ModuleType:
    module = sys.modules.get(module_name)
    if module is None:
        module = types.ModuleType(module_name)
        sys.modules[module_name] = module
    return module


_GRADIO_CSS_PATCH_PATH: Path | None = None


def _normalize_gradio_css_paths(css_paths: Any) -> list[str]:
    if css_paths is None:
        return []
    if isinstance(css_paths, (str, Path)):
        return [str(css_paths)]
    return [str(path) for path in css_paths]


def _persist_egoforce_gradio_css(css: str) -> str:
    global _GRADIO_CSS_PATCH_PATH

    if _GRADIO_CSS_PATCH_PATH is None:
        _GRADIO_CSS_PATCH_PATH = Path(tempfile.gettempdir()) / "egoforce-gradio-launch.css"
    _GRADIO_CSS_PATCH_PATH.write_text(css, encoding="utf-8")
    return str(_GRADIO_CSS_PATCH_PATH)


def _patch_gradio_launch() -> None:
    try:
        import gradio as gr
    except ImportError:
        return

    launch_method = getattr(gr.Blocks, "launch", None)
    if launch_method is None or getattr(launch_method, "__egoforce_runtime_patch__", False):
        return

    @functools.wraps(launch_method)
    def patched_launch(self: Any, *args: Any, **kwargs: Any) -> Any:
        css = kwargs.get("css")
        if isinstance(css, str) and css.strip() and (".egoforce-hero" in css or "#sample-video-carousel" in css):
            css_path_entries = _normalize_gradio_css_paths(kwargs.get("css_paths"))
            patched_css_path = _persist_egoforce_gradio_css(css)
            if patched_css_path not in css_path_entries:
                css_path_entries.append(patched_css_path)
            kwargs["css_paths"] = css_path_entries
            kwargs["css"] = None
        return launch_method(self, *args, **kwargs)

    setattr(patched_launch, "__egoforce_runtime_patch__", True)
    gr.Blocks.launch = patched_launch


def apply_runtime_patches() -> None:
    _patch_gradio_launch()

    try:
        mmcv = importlib.import_module("mmcv")
    except ImportError:
        mmcv = None

    ops_module = sys.modules.get("mmcv.ops")
    if ops_module is None:
        ops_module = types.ModuleType("mmcv.ops")
        sys.modules["mmcv.ops"] = ops_module
    ops_module.__path__ = []

    nms_module = _ensure_module("mmcv.ops.nms")

    roi_align_module = _ensure_module("mmcv.ops.roi_align")
    deform_conv_module = _ensure_module("mmcv.ops.deform_conv")
    modulated_deform_conv_module = _ensure_module("mmcv.ops.modulated_deform_conv")
    carafe_module = _ensure_module("mmcv.ops.carafe")
    merge_cells_module = _ensure_module("mmcv.ops.merge_cells")
    multi_scale_deform_attn_module = _ensure_module("mmcv.ops.multi_scale_deform_attn")

    deform_conv2d = _unsupported_function("deform_conv2d")
    DeformConv2d = _unsupported_module_class("DeformConv2d")
    ModulatedDeformConv2d = _unsupported_module_class("ModulatedDeformConv2d")
    MaskedConv2d = _unsupported_module_class("MaskedConv2d")
    CornerPool = _unsupported_module_class("CornerPool")
    CARAFEPack = _unsupported_module_class("CARAFEPack")
    GlobalPoolingCell = _unsupported_module_class("GlobalPoolingCell")
    SumCell = _unsupported_module_class("SumCell")
    ConcatCell = _unsupported_module_class("ConcatCell")
    MultiScaleDeformableAttention = _multiscale_deformable_attention_class()

    ops_module.nms = _nms
    ops_module.batched_nms = _batched_nms
    ops_module.nms_match = _nms_match
    ops_module.point_sample = _point_sample
    ops_module.rel_roi_point_to_rel_img_point = _rel_roi_point_to_rel_img_point
    ops_module.sigmoid_focal_loss = _sigmoid_focal_loss
    ops_module.bbox_overlaps = _bbox_overlaps
    ops_module.roi_align = _torchvision_roi_align()
    ops_module.RoIAlign = _torchvision_roi_align_module()
    ops_module.RoIPool = _torchvision_roi_pool_module()
    ops_module.deform_conv2d = deform_conv2d
    ops_module.DeformConv2d = DeformConv2d
    ops_module.ModulatedDeformConv2d = ModulatedDeformConv2d
    ops_module.MaskedConv2d = MaskedConv2d
    ops_module.CornerPool = CornerPool
    ops_module.CARAFEPack = CARAFEPack
    ops_module.GlobalPoolingCell = GlobalPoolingCell
    ops_module.SumCell = SumCell
    ops_module.ConcatCell = ConcatCell
    ops_module.MultiScaleDeformableAttention = MultiScaleDeformableAttention
    nms_module.nms = _nms
    nms_module.batched_nms = _batched_nms
    roi_align_module.roi_align = ops_module.roi_align
    roi_align_module.RoIAlign = ops_module.RoIAlign
    deform_conv_module.deform_conv2d = deform_conv2d
    deform_conv_module.DeformConv2d = DeformConv2d
    modulated_deform_conv_module.ModulatedDeformConv2d = ModulatedDeformConv2d
    carafe_module.CARAFEPack = CARAFEPack
    merge_cells_module.GlobalPoolingCell = GlobalPoolingCell
    merge_cells_module.SumCell = SumCell
    merge_cells_module.ConcatCell = ConcatCell
    multi_scale_deform_attn_module.MultiScaleDeformableAttention = MultiScaleDeformableAttention

    try:
        transformer_module = importlib.import_module("mmcv.cnn.bricks.transformer")
        if not hasattr(transformer_module, "MultiScaleDeformableAttention"):
            transformer_module.MultiScaleDeformableAttention = MultiScaleDeformableAttention
    except ImportError:
        pass

    if mmcv is not None:
        mmcv.ops = ops_module