from __future__ import annotations

import argparse
import os
from pathlib import Path
from typing import Dict, Tuple
import sys

import torch
from tqdm import tqdm


THIS_DIR = Path(__file__).resolve().parent
REPO_ROOT = THIS_DIR.parents[2]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from quantization.escnet.data.calib import build_calib_loader
from quantization.escnet.methods.registry import METHODS
from quantization.escnet.quant.config import QuantConfig
from quantization.escnet.utils.io import load_quant_checkpoint
from quantization.escnet.utils.seed import set_seed


def _import_load_config():
    from model_sources.escnet.ESCNet_model.config import load_config
    return load_config


def _import_escnet():
    try:
        from model_sources.escnet.ESCNet_model.models.ESCNet import ESCNet  # type: ignore
        return ESCNet
    except Exception as e:
        raise RuntimeError(
            "Failed to import ESCNet. Run this from repo root and ensure ./model_sources/escnet/ESCNet_model exists."
        ) from e


def _try_import_dataset():
    from model_sources.escnet.ESCNet_model.dataset import MyData  # type: ignore
    return MyData


def _try_import_save_tensor_img():
    try:
        from model_sources.escnet.ESCNet_model.utils import save_tensor_img
        return save_tensor_img
    except Exception:
        # fallback
        from PIL import Image
        import numpy as np

        def save_tensor_img(t: torch.Tensor, out_path: str) -> None:
            arr = (t.detach().float().cpu().numpy().squeeze() * 255.0).clip(0, 255).astype("uint8")
            Image.fromarray(arr).save(out_path)

        return save_tensor_img


@torch.no_grad()
def run_infer(model: torch.nn.Module, loader, pred_root: str, device: torch.device) -> None:
    import cv2

    save_tensor_img = _try_import_save_tensor_img()
    model.eval()
    Path(pred_root).mkdir(parents=True, exist_ok=True)

    for batch in tqdm(loader, total=len(loader)):
        inputs = batch[0].to(device, non_blocking=True)
        label_paths = batch[-1]

        out_edge, scaled_preds = model(inputs)
        pred_lvl = (scaled_preds[-1].sigmoid() >= 0.5).float()

        for idx_sample in range(pred_lvl.shape[0]):
            res = torch.nn.functional.interpolate(
                pred_lvl[idx_sample].unsqueeze(0),
                size=cv2.imread(label_paths[idx_sample], cv2.IMREAD_GRAYSCALE).shape[:2],
                mode="bilinear",
                align_corners=True,
            )
            out_path = os.path.join(pred_root, os.path.basename(label_paths[idx_sample]))
            save_tensor_img(res, out_path)


def main() -> None:
    parser = argparse.ArgumentParser("ESCNet quant inference")
    parser.add_argument("--config", type=str, default="configs/escnet/main/config.yaml")
    parser.add_argument("--quant_ckpt", type=str, required=True, help="Quant checkpoint exported by quantize_escnet.py")
    parser.add_argument("--pred_root", type=str, default="preds_quant")
    parser.add_argument("--test_dir", type=str, default="", help="Override dataset dir (must contain Imgs/GT)")
    parser.add_argument("--device", type=str, default="cuda:0")
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--batch_size", type=int, default=-1, help="Override config.batch_size_valid")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--verbose", action="store_true")
    args = parser.parse_args()

    set_seed(args.seed)

    load_config = _import_load_config()
    ESCNet = _import_escnet()
    MyData = _try_import_dataset()

    cfg = load_config(args.config)

    state, qcfg, meta = load_quant_checkpoint(args.quant_ckpt)

    device = torch.device(args.device if torch.cuda.is_available() else "cpu")

    model = ESCNet(cfg, pretrained=False)
    model.eval()

    # wrap with same qcfg then load quant state
    METHODS["ptq_qdq"](model, qcfg)
    missing, unexpected = model.load_state_dict(state, strict=False)
    if args.verbose:
        print(f"[load] missing={len(missing)}, unexpected={len(unexpected)}")

    model.to(device)

    test_dir = args.test_dir or getattr(cfg, "test_dir")
    bs = args.batch_size if args.batch_size > 0 else int(getattr(cfg, "batch_size_valid", 1))

    loader = torch.utils.data.DataLoader(
        dataset=MyData(cfg, test_dir, image_size=int(getattr(cfg, "img_size", 384)), is_train=False),
        batch_size=bs,
        shuffle=False,
        num_workers=int(args.num_workers),
        pin_memory=True,
    )

    run_infer(model, loader, pred_root=args.pred_root, device=device)
    print(f"[done] preds -> {args.pred_root}")


if __name__ == "__main__":
    main()