import argparse
import os
import warnings

import cv2
import numpy as np
import torch
import torch.cuda.amp as amp
import torch.nn.functional as F
from loguru import logger

import utils.config as config
from model import build_net
from utils.grasp_eval import detect_grasps
from utils.misc import setup_logger

from infer_custom_image_nodepth import (
    default_save_path,
    inverse_affine,
    preprocess_image,
    save_prediction_overlay,
)

warnings.filterwarnings("ignore")
cv2.setNumThreads(0)


ALLOWED_MISSING_PREFIXES = (
    "llava_model.",
    "module.llava_model.",
)


def get_parser():
    parser = argparse.ArgumentParser(
        description="Run non-depth MaPLe custom inference on one RGB image and sentence"
    )
    parser.add_argument("--config", required=True, type=str, help="config file")
    parser.add_argument("--image", required=True, type=str, help="input RGB image path")
    parser.add_argument("--sentence", required=True, type=str, help="referring expression")
    parser.add_argument("--save", default="", type=str, help="output visualization path")
    parser.add_argument("--num-grasps", default=5, type=int, help="number of grasp boxes to draw")
    parser.add_argument("--mask-threshold", default=0.35, type=float, help="mask threshold")
    parser.add_argument(
        "--gate-grasps-by-mask",
        action="store_true",
        help="also gate final grasp peak detection by the predicted instance mask",
    )
    parser.add_argument(
        "--opts",
        default=None,
        nargs=argparse.REMAINDER,
        help="override settings in the config.",
    )
    cli_args = parser.parse_args()

    cfg = config.load_cfg_from_cfg_file(cli_args.config)
    if cli_args.opts is not None:
        cfg = config.merge_cfg_from_list(cfg, cli_args.opts)

    cfg.custom_image = cli_args.image
    cfg.custom_sentence = cli_args.sentence
    cfg.custom_save = cli_args.save
    cfg.custom_num_grasps = cli_args.num_grasps
    cfg.custom_mask_threshold = cli_args.mask_threshold
    cfg.custom_gate_grasps_by_mask = cli_args.gate_grasps_by_mask
    return cfg


def _normalize_state_dict_for_model(model, state_dict):
    model_keys = model.state_dict().keys()
    ckpt_keys = state_dict.keys()
    model_has_module = len(model_keys) > 0 and next(iter(model_keys)).startswith("module.")
    ckpt_has_module = len(ckpt_keys) > 0 and next(iter(ckpt_keys)).startswith("module.")

    if model_has_module and not ckpt_has_module:
        return {"module." + k: v for k, v in state_dict.items()}
    if not model_has_module and ckpt_has_module:
        return {k.replace("module.", "", 1): v for k, v in state_dict.items()}
    return state_dict


def _is_allowed_missing_key(key):
    return key.startswith(ALLOWED_MISSING_PREFIXES)


def load_checkpoint_weights_checked(model, resume_path):
    if not resume_path or not os.path.isfile(resume_path):
        raise ValueError(f"=> resume failed! no checkpoint found at '{resume_path}'")

    logger.info("=> loading checkpoint '{}'".format(resume_path))
    checkpoint = torch.load(resume_path, map_location="cpu")
    state_dict = checkpoint.get("state_dict", checkpoint)
    state_dict = _normalize_state_dict_for_model(model, state_dict)

    incompatible = model.load_state_dict(state_dict, strict=False)
    missing = list(incompatible.missing_keys)
    unexpected = list(incompatible.unexpected_keys)

    bad_missing = [k for k in missing if not _is_allowed_missing_key(k)]
    if missing:
        shown = missing[:20]
        suffix = " ..." if len(missing) > len(shown) else ""
        logger.info(f"Missing keys ({len(missing)}): {shown}{suffix}")
    if unexpected:
        shown = unexpected[:20]
        suffix = " ..." if len(unexpected) > len(shown) else ""
        logger.info(f"Unexpected keys ({len(unexpected)}): {shown}{suffix}")

    if bad_missing or unexpected:
        raise RuntimeError(
            "Checkpoint/config mismatch for MaPLe inference. "
            f"Non-LLaVA missing keys: {bad_missing[:20]} "
            f"(total {len(bad_missing)}); unexpected keys: {unexpected[:20]} "
            f"(total {len(unexpected)})."
        )

    if "epoch" in checkpoint:
        logger.info(f"=> loaded checkpoint '{resume_path}' (epoch {checkpoint['epoch']})")
    else:
        logger.info("=> loaded checkpoint '{}'".format(resume_path))
    return checkpoint


def _ensure_maple_inference_args(args):
    network_name = str(args.network_name).lower()
    if "maple" not in network_name:
        raise ValueError(
            f"This custom inference script is for MaPLe checkpoints, got network_name={args.network_name!r}."
        )
    if "3dcoord" in network_name or "depthcoord" in network_name:
        raise ValueError(
            "This script is for ordinary non-depth MaPLe checkpoints. "
            "Use a depth-aware custom inference script for maple_3dcoord checkpoints."
        )
    if getattr(args, "stage1", False) and not getattr(args, "stage2", False):
        raise ValueError(
            "The config is still in stage1 mask-only mode. Pass opts like "
            "`stage1 False stage2 True use_gt_obj_masks False` for the withft checkpoint."
        )


def _unpack_maple_prediction(pred):
    if not isinstance(pred, (tuple, list)) or len(pred) < 5:
        raise RuntimeError(
            "Expected MaPLe stage2 inference to return five maps "
            "(mask, qua, sin, cos, wid). Did you pass `stage1 False stage2 True`?"
        )
    return pred[0], pred[1], pred[2], pred[3], pred[4]


@logger.catch
def main():
    args = get_parser()
    _ensure_maple_inference_args(args)

    args.output_dir = os.path.join(args.output_folder, args.exp_name)
    setup_logger(args.output_dir, distributed_rank=0, filename="custom_infer_maple.log", mode="a")
    logger.info(args)

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required for this inference script.")

    orig_img, image, inv_mat = preprocess_image(args.custom_image, args.input_size)
    h, w = orig_img.shape[:2]

    model, _ = build_net(args)
    model = torch.nn.DataParallel(model).cuda()
    load_checkpoint_weights_checked(model, args.resume)
    model.eval()

    image = image.cuda(non_blocking=True)
    text = [args.custom_sentence]
    with torch.no_grad(), amp.autocast():
        pred, _ = model(image, text)

    (
        ins_mask_logit,
        grasp_qua_logit,
        grasp_sin_logit,
        grasp_cos_logit,
        grasp_wid_logit,
    ) = _unpack_maple_prediction(pred)

    ins_mask_pred = torch.sigmoid(ins_mask_logit)
    grasp_qua_mask_pred = torch.sigmoid(grasp_qua_logit)
    grasp_sin_mask_pred = grasp_sin_logit
    grasp_cos_mask_pred = grasp_cos_logit
    grasp_wid_mask_pred = torch.sigmoid(grasp_wid_logit)

    if ins_mask_pred.shape[-2:] != image.shape[-2:]:
        pred_size = image.shape[-2:]
        ins_mask_pred = F.interpolate(
            ins_mask_pred, size=pred_size, mode="bicubic", align_corners=True
        ).squeeze(1)
        grasp_qua_mask_pred = F.interpolate(
            grasp_qua_mask_pred, size=pred_size, mode="bicubic", align_corners=True
        ).squeeze(1)
        grasp_sin_mask_pred = F.interpolate(
            grasp_sin_mask_pred, size=pred_size, mode="bicubic", align_corners=True
        ).squeeze(1)
        grasp_cos_mask_pred = F.interpolate(
            grasp_cos_mask_pred, size=pred_size, mode="bicubic", align_corners=True
        ).squeeze(1)
        grasp_wid_mask_pred = F.interpolate(
            grasp_wid_mask_pred, size=pred_size, mode="bicubic", align_corners=True
        ).squeeze(1)
    else:
        ins_mask_pred = ins_mask_pred.squeeze(1)
        grasp_qua_mask_pred = grasp_qua_mask_pred.squeeze(1)
        grasp_sin_mask_pred = grasp_sin_mask_pred.squeeze(1)
        grasp_cos_mask_pred = grasp_cos_mask_pred.squeeze(1)
        grasp_wid_mask_pred = grasp_wid_mask_pred.squeeze(1)

    pred_mask_prob = inverse_affine(ins_mask_pred[0].cpu().numpy(), inv_mat, w, h)
    pred_mask = pred_mask_prob > args.custom_mask_threshold
    grasp_qua_np = inverse_affine(grasp_qua_mask_pred[0].cpu().numpy(), inv_mat, w, h)
    grasp_sin_np = inverse_affine(grasp_sin_mask_pred[0].cpu().numpy(), inv_mat, w, h)
    grasp_cos_np = inverse_affine(grasp_cos_mask_pred[0].cpu().numpy(), inv_mat, w, h)
    grasp_wid_np = inverse_affine(grasp_wid_mask_pred[0].cpu().numpy(), inv_mat, w, h)

    logger.info(
        "Pred stats: mask_logit min={:.4f}, max={:.4f}, mean={:.4f}; "
        "mask_prob min={:.4f}, max={:.4f}, mean={:.4f}, area@{:.2f}={:.4f}; "
        "qua min={:.4f}, max={:.4f}, mean={:.4f}; wid min={:.4f}, max={:.4f}, mean={:.4f}; "
        "sin mean={:.4f}, cos mean={:.4f}".format(
            float(ins_mask_logit.min().detach().cpu()),
            float(ins_mask_logit.max().detach().cpu()),
            float(ins_mask_logit.mean().detach().cpu()),
            float(np.nanmin(pred_mask_prob)),
            float(np.nanmax(pred_mask_prob)),
            float(np.nanmean(pred_mask_prob)),
            float(args.custom_mask_threshold),
            float(np.nanmean(pred_mask.astype(np.float32))),
            float(np.nanmin(grasp_qua_np)),
            float(np.nanmax(grasp_qua_np)),
            float(np.nanmean(grasp_qua_np)),
            float(np.nanmin(grasp_wid_np)),
            float(np.nanmax(grasp_wid_np)),
            float(np.nanmean(grasp_wid_np)),
            float(np.nanmean(grasp_sin_np)),
            float(np.nanmean(grasp_cos_np)),
        )
    )

    grasp_qua_for_detect = grasp_qua_np
    if getattr(args, "custom_gate_grasps_by_mask", False):
        if pred_mask.any():
            grasp_qua_for_detect = grasp_qua_np * pred_mask.astype(grasp_qua_np.dtype)
            logger.info("Using predicted mask to gate final grasp peak detection.")
        else:
            grasp_qua_for_detect = np.zeros_like(grasp_qua_np)
            logger.warning(
                "Predicted mask is empty at the current threshold; "
                "mask-gated grasp detection will return no peaks."
            )

    pred_grasps, _ = detect_grasps(
        grasp_qua_for_detect,
        grasp_sin_np,
        grasp_cos_np,
        grasp_wid_np,
        args.custom_num_grasps,
    )
    if pred_grasps:
        logger.info(f"Pred grasps [cx, cy, width, height, theta_deg]: {pred_grasps}")
    else:
        logger.warning("No grasp peaks found by detect_grasps.")

    save_path = args.custom_save or default_save_path(args)
    save_prediction_overlay(orig_img, pred_mask, pred_grasps, args.custom_sentence, save_path)
    logger.info(f"Saved MaPLe custom inference visualization: {save_path}")
    print(save_path)


if __name__ == "__main__":
    main()