from __future__ import annotations

import importlib
import os
import pkgutil

import onnx
import onnxsim
import torch
from rfdetr.deploy._onnx.optimizer import OnnxOptimizer
from rfdetr.models.transformer import gen_encoder_output_proposals
from rfdetr.util import box_ops
from torch import nn
from transformers.utils import torch_int


def import_all_submodules(package_name):
    """Recursively import all submodules of a package"""
    package = importlib.import_module(package_name)

    for loader, name, is_pkg in pkgutil.walk_packages(
        package.__path__, package.__name__ + "."
    ):
        try:
            importlib.import_module(name)
        except Exception as e:
            print(f"Failed to import {name}: {e}")


class RFDETRPatcher:
    def __init__(self, module):
        self.module = module
        import_all_submodules("rfdetr")

    def __call__(self):
        # Patch dinov2 backbone export function to exclude position embedding from the export
        self.module.models.backbone.dinov2.DinoV2.export = export_backbone

        # patch main export function to enable dynamic batches and change output names
        self.module.main.export = export_main

        # patch model output to enable yolo like output
        self.module.models.lwdetr.LWDETR.forward_export = forward_export

        # patch window attention antialias
        self.module.models.backbone.dinov2_with_windowed_attn.WindowedDinov2WithRegistersEmbeddings.interpolate_pos_encoding = (
            interpolate_pos_encoding
        )

        # patch static shapes in transformer encoder
        self.module.models.transformer.Transformer.forward = forward

        # patch model export name
        self.module.deploy.export.export_onnx = export_onnx_main

        # patch onnx simplify name
        self.module.deploy.export.onnx_simplify = onnx_simplify_main


# RFDETR functions


def export_backbone(self):
    self._export = True


def export_main(
    self,
    output_dir="output",
    infer_dir=None,
    simplify=False,
    backbone_only=False,
    opset_version=17,
    verbose=True,
    force=False,
    shape=None,
    batch_size=1,
    **kwargs,
):
    """Export the trained model to ONNX format"""
    print(f"Exporting model to ONNX format")
    try:
        from rfdetr.deploy.export import make_infer_image
    except ImportError:
        print(
            "It seems some dependencies for ONNX export are missing. Please run `pip install rfdetr[onnxexport]` and try again."
        )
        raise

    device = self.device
    model = deepcopy(self.model.to("cpu"))
    model.to(device)

    os.makedirs(output_dir, exist_ok=True)
    output_dir = Path(output_dir)
    if shape is None:
        shape = (self.resolution, self.resolution)
    else:
        if shape[0] % 14 != 0 or shape[1] % 14 != 0:
            raise ValueError("Shape must be divisible by 14")

    input_tensors = make_infer_image(infer_dir, shape, batch_size, device).to(
        device
    )
    input_names = ["images"]
    output_names = (
        ["features"]
        if backbone_only
        else ["num_dets", "det_boxes", "det_scores", "det_classes"]
    )
    dynamic_axes = {
        "images": {0: "batch"},
    }
    self.model.eval()
    with torch.no_grad():
        if backbone_only:
            features = model(input_tensors)
            dynamic_axes.update(
                {
                    "features": {0: "batch"},
                }
            )
            print(f"PyTorch inference output shape: {features.shape}")
        else:
            outputs = model(input_tensors)
            dynamic_axes.update(
                {
                    "num_dets": {0: "batch"},
                    "det_boxes": {0: "batch"},
                    "det_scores": {0: "batch"},
                    "det_classes": {0: "batch"},
                }
            )
            dets = outputs["pred_boxes"]
            labels = outputs["pred_logits"]
            print(
                f"PyTorch inference output shapes - Boxes: {dets.shape}, Labels: {labels.shape}"
            )
    model.cpu()
    input_tensors = input_tensors.cpu()

    # Export to ONNX
    output_file = export_onnx(
        output_dir=output_dir,
        model=model,
        input_names=input_names,
        input_tensors=input_tensors,
        output_names=output_names,
        dynamic_axes=dynamic_axes,
        backbone_only=backbone_only,
        verbose=verbose,
        opset_version=opset_version,
    )

    print(f"Successfully exported ONNX model to: {output_file}")

    if simplify:
        sim_output_file = onnx_simplify(
            onnx_dir=output_file,
            input_names=input_names,
            input_tensors=input_tensors,
            force=force,
        )
        print(f"Successfully simplified ONNX model to: {sim_output_file}")

    print("ONNX export completed successfully")
    self.model = self.model.to(device)


def forward_export(self, tensors):
    srcs, _, poss = self.backbone(tensors)
    # only use one group in inference
    refpoint_embed_weight = self.refpoint_embed.weight[: self.num_queries]
    query_feat_weight = self.query_feat.weight[: self.num_queries]

    hs, ref_unsigmoid, hs_enc, ref_enc = self.transformer(
        srcs, None, poss, refpoint_embed_weight, query_feat_weight
    )

    if self.bbox_reparam:
        outputs_coord_delta = self.bbox_embed(hs)
        outputs_coord_cxcy = (
            outputs_coord_delta[..., :2] * ref_unsigmoid[..., 2:]
            + ref_unsigmoid[..., :2]
        )
        outputs_coord_wh = (
            outputs_coord_delta[..., 2:].exp() * ref_unsigmoid[..., 2:]
        )
        outputs_coord = torch.concat(
            [outputs_coord_cxcy, outputs_coord_wh], dim=-1
        )
    else:
        outputs_coord = (self.bbox_embed(hs) + ref_unsigmoid).sigmoid()
    outputs_class = self.class_embed(hs)
    prob = outputs_class.sigmoid()
    # Static topk with explicit self.num_queries
    topk_values, topk_indexes = torch.topk(
        prob.view(outputs_class.shape[0], -1), self.num_queries, dim=1
    )
    scores = topk_values
    b_size = prob.shape[0]
    # Fix type conversion - apply to the result, not the divisor
    topk_boxes = topk_indexes // outputs_class.shape[2]
    labels = (
        (topk_indexes % outputs_class.shape[2])
        .view(b_size, self.num_queries)
        .to(torch.int32)
    )
    boxes = box_ops.box_cxcywh_to_xyxy(outputs_coord)
    boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4))

    # Create scale factor with explicit reshape to ensure static shape
    img_h, img_w = tensors.shape[2].unsqueeze(dim=0).expand(
        1, self.num_queries
    ), tensors.shape[3].unsqueeze(dim=0).expand(1, self.num_queries)
    # Ensure static shapes with reshape operations
    scale_fct = torch.cat([img_w, img_h, img_w, img_h], dim=1).view(
        b_size, self.num_queries, 4
    )
    # Ensure boxes has static shape by forcing dimension
    boxes = boxes * scale_fct
    boxes = boxes.reshape(b_size, self.num_queries, 4)
    scores = scores.reshape(b_size, self.num_queries)

    # Create num_dets with static shape
    num_dets = torch.full(
        (b_size, 1), self.num_queries, dtype=torch.int32, device=boxes.device
    ).view((b_size, 1))

    # Return with explicit static shapes
    return num_dets, boxes, scores, labels


def interpolate_pos_encoding(
    self, embeddings: torch.Tensor, height: int, width: int
) -> torch.Tensor:
    """
    This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
    resolution images. This implementation supports torch.jit tracing while maintaining backwards compatibility
    with the original implementation.

    Adapted from:
    - https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
    - https://github.com/facebookresearch/dinov2/blob/main/dinov2/models/vision_transformer.py
    """
    num_patches = embeddings.shape[1] - 1
    num_positions = self.position_embeddings.shape[1] - 1

    # Skip interpolation for matching dimensions (unless tracing)
    if (
        not torch.jit.is_tracing()
        and num_patches == num_positions
        and height == width
    ):
        return self.position_embeddings

    # Handle class token and patch embeddings separately
    class_pos_embed = self.position_embeddings[:, 0]
    patch_pos_embed = self.position_embeddings[:, 1:]
    dim = embeddings.shape[-1]

    # Calculate new dimensions
    height = height // self.config.patch_size
    width = width // self.config.patch_size

    # Reshape for interpolation
    sqrt_num_positions = torch_int(num_positions**0.5)
    patch_pos_embed = patch_pos_embed.reshape(
        1, sqrt_num_positions, sqrt_num_positions, dim
    )
    patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)

    # Store original dtype for restoration after interpolation
    target_dtype = patch_pos_embed.dtype

    # Interpolate at float32 precision
    patch_pos_embed = nn.functional.interpolate(
        patch_pos_embed.to(dtype=torch.float32),
        size=(
            torch_int(height),
            torch_int(width),
        ),  # Explicit size instead of scale_factor
        mode="bicubic",
        align_corners=False,
        antialias=False,
    ).to(dtype=target_dtype)

    # Validate output dimensions if not tracing
    if not torch.jit.is_tracing():
        if (
            int(height) != patch_pos_embed.shape[-2]
            or int(width) != patch_pos_embed.shape[-1]
        ):
            raise ValueError(
                "Width or height does not match with the interpolated position embeddings"
            )

    # Reshape back to original format
    patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)

    # Combine class and patch embeddings
    return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)


def forward(self, srcs, masks, pos_embeds, refpoint_embed, query_feat):
    src_flatten = []
    mask_flatten = [] if masks is not None else None
    lvl_pos_embed_flatten = []
    spatial_shapes = torch.empty(
        (len(srcs), 2), device=srcs[0].device, dtype=torch.long
    )
    valid_ratios = [] if masks is not None else None
    for lvl, (src, pos_embed) in enumerate(zip(srcs, pos_embeds)):
        bs, c, h, w = src.shape
        spatial_shapes[lvl, 0] = h
        spatial_shapes[lvl, 1] = w

        src = src.flatten(2).transpose(1, 2)  # bs, hw, c
        pos_embed = pos_embed.flatten(2).transpose(1, 2)  # bs, hw, c
        lvl_pos_embed_flatten.append(pos_embed)
        src_flatten.append(src)
        if masks is not None:
            mask = masks[lvl].flatten(1)  # bs, hw
            mask_flatten.append(mask)
    memory = torch.cat(src_flatten, 1)  # bs, \sum{hxw}, c
    if masks is not None:
        mask_flatten = torch.cat(mask_flatten, 1)  # bs, \sum{hxw}
        valid_ratios = torch.stack([self.get_valid_ratio(m) for m in masks], 1)
    lvl_pos_embed_flatten = torch.cat(
        lvl_pos_embed_flatten, 1
    )  # bs, \sum{hxw}, c
    level_start_index = torch.cat(
        (spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1])
    )

    if self.two_stage:
        output_memory, output_proposals = gen_encoder_output_proposals(
            memory,
            mask_flatten,
            spatial_shapes,
            unsigmoid=not self.bbox_reparam,
        )
        # group detr for first stage
        refpoint_embed_ts, memory_ts, boxes_ts = [], [], []
        group_detr = self.group_detr if self.training else 1
        for g_idx in range(group_detr):
            output_memory_gidx = self.enc_output_norm[g_idx](
                self.enc_output[g_idx](output_memory)
            )

            enc_outputs_class_unselected_gidx = self.enc_out_class_embed[
                g_idx
            ](output_memory_gidx)
            if self.bbox_reparam:
                enc_outputs_coord_delta_gidx = self.enc_out_bbox_embed[g_idx](
                    output_memory_gidx
                )
                enc_outputs_coord_cxcy_gidx = (
                    enc_outputs_coord_delta_gidx[..., :2]
                    * output_proposals[..., 2:]
                    + output_proposals[..., :2]
                )
                enc_outputs_coord_wh_gidx = (
                    enc_outputs_coord_delta_gidx[..., 2:].exp()
                    * output_proposals[..., 2:]
                )
                enc_outputs_coord_unselected_gidx = torch.concat(
                    [enc_outputs_coord_cxcy_gidx, enc_outputs_coord_wh_gidx],
                    dim=-1,
                )
            else:
                enc_outputs_coord_unselected_gidx = (
                    self.enc_out_bbox_embed[g_idx](output_memory_gidx)
                    + output_proposals
                )  # (bs, \sum{hw}, 4) unsigmoid

            topk = min(
                self.num_queries, enc_outputs_class_unselected_gidx.shape[-2]
            )
            topk_proposals_gidx = torch.topk(
                enc_outputs_class_unselected_gidx.max(-1)[0], topk, dim=1
            )[
                1
            ]  # bs, nq

            refpoint_embed_gidx_undetach = torch.gather(
                enc_outputs_coord_unselected_gidx,
                1,
                topk_proposals_gidx.unsqueeze(-1).repeat(1, 1, 4),
            )  # unsigmoid
            # for decoder layer, detached as initial ones, (bs, nq, 4)
            refpoint_embed_gidx = refpoint_embed_gidx_undetach.detach()

            # get memory tgt
            tgt_undetach_gidx = torch.gather(
                output_memory_gidx,
                1,
                topk_proposals_gidx.unsqueeze(-1).repeat(1, 1, self.d_model),
            )

            refpoint_embed_ts.append(refpoint_embed_gidx)
            memory_ts.append(tgt_undetach_gidx)
            boxes_ts.append(refpoint_embed_gidx_undetach)
        # concat on dim=1, the nq dimension, (bs, nq, d) --> (bs, nq, d)
        refpoint_embed_ts = torch.cat(refpoint_embed_ts, dim=1)
        # (bs, nq, d)
        memory_ts = torch.cat(memory_ts, dim=1)  # .transpose(0, 1)
        boxes_ts = torch.cat(boxes_ts, dim=1)  # .transpose(0, 1)

    tgt = query_feat.unsqueeze(0).repeat(bs, 1, 1)
    refpoint_embed = refpoint_embed.unsqueeze(0).repeat(bs, 1, 1)
    if self.two_stage:
        ts_len = refpoint_embed_ts.shape[-2]
        refpoint_embed_ts_subset = refpoint_embed[..., :ts_len, :]
        refpoint_embed_subset = refpoint_embed[..., ts_len:, :]

        if self.bbox_reparam:
            refpoint_embed_cxcy = (
                refpoint_embed_ts_subset[..., :2] * refpoint_embed_ts[..., 2:]
            )
            refpoint_embed_cxcy = (
                refpoint_embed_cxcy + refpoint_embed_ts[..., :2]
            )
            refpoint_embed_wh = (
                refpoint_embed_ts_subset[..., 2:].exp()
                * refpoint_embed_ts[..., 2:]
            )
            refpoint_embed_ts_subset = torch.concat(
                [refpoint_embed_cxcy, refpoint_embed_wh], dim=-1
            )
        else:
            refpoint_embed_ts_subset = (
                refpoint_embed_ts_subset + refpoint_embed_ts
            )

        refpoint_embed = torch.concat(
            [refpoint_embed_ts_subset, refpoint_embed_subset], dim=-2
        )

    hs, references = self.decoder(
        tgt,
        memory,
        memory_key_padding_mask=mask_flatten,
        pos=lvl_pos_embed_flatten,
        refpoints_unsigmoid=refpoint_embed,
        level_start_index=level_start_index,
        spatial_shapes=spatial_shapes,
        valid_ratios=(
            valid_ratios.to(memory.dtype)
            if valid_ratios is not None
            else valid_ratios
        ),
    )
    if self.two_stage:
        if self.bbox_reparam:
            return hs, references, memory_ts, boxes_ts
        else:
            return hs, references, memory_ts, boxes_ts.sigmoid()
    return hs, references, None, None


def export_onnx_main(
    output_dir,
    model,
    input_names,
    input_tensors,
    output_names,
    dynamic_axes,
    backbone_only=False,
    verbose=True,
    opset_version=17,
):
    export_name = "model"
    output_file = os.path.join(output_dir, f"{export_name}.onnx")

    # Prepare model for export
    if hasattr(model, "export"):
        model.export()

    torch.onnx.export(
        model,
        input_tensors,
        output_file,
        input_names=input_names,
        output_names=output_names,
        export_params=True,
        keep_initializers_as_inputs=False,
        do_constant_folding=True,
        verbose=verbose,
        opset_version=opset_version,
        dynamic_axes=dynamic_axes,
    )

    print(f"\nSuccessfully exported ONNX model: {output_file}")
    return output_file


def onnx_simplify_main(onnx_dir: str, input_names, input_tensors, force=True):
    sim_onnx_dir = onnx_dir

    if isinstance(input_tensors, torch.Tensor):
        input_tensors = [input_tensors]

    print(f"start simplify ONNX model: {onnx_dir}")
    opt = OnnxOptimizer(onnx_dir)
    opt.info("Model: original")
    opt.common_opt()
    opt.info("Model: optimized")
    opt.save_onnx(sim_onnx_dir)
    input_dict = {
        name: tensor.detach().cpu().numpy()
        for name, tensor in zip(input_names, input_tensors)
    }
    model_opt, check_ok = onnxsim.simplify(
        onnx_dir, check_n=3, input_data=input_dict, dynamic_input_shape=False
    )
    if check_ok:
        onnx.save(model_opt, sim_onnx_dir)
    else:
        raise RuntimeError("Failed to simplify ONNX model.")
    print(f"Successfully simplified ONNX model: {sim_onnx_dir}")
    return sim_onnx_dir
