import os
import sys
import cv2
import torch
import argparse
from pathlib import Path

# =========================
# 需要确保本脚本能找到 YOLOv5 源码
# 如果脚本不在 yolov5 根目录，可手动添加：
# sys.path.append("/path/to/yolov5")
# =========================

from models.common import DetectMultiBackend
from utils.general import (
    check_img_size,
    non_max_suppression,
    scale_boxes
)
from utils.torch_utils import select_device
from utils.augmentations import letterbox


def is_image_file(filename):
    image_exts = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"}
    return Path(filename).suffix.lower() in image_exts


def load_image(img_path, img_size=640, stride=32):
    """
    读取图像，并进行与YOLOv5推理一致的预处理
    返回：
        im0: 原图(BGR)
        im:  预处理后的输入张量所对应的numpy图像(CHW, RGB)
    """
    im0 = cv2.imread(str(img_path))
    if im0 is None:
        raise ValueError(f"无法读取图片: {img_path}")

    # letterbox缩放
    im = letterbox(im0, img_size, stride=stride, auto=True)[0]

    # BGR -> RGB, HWC -> CHW
    im = im[:, :, ::-1].transpose(2, 0, 1)
    im = im.copy()

    return im0, im


def save_crop(img, xyxy, save_path):
    """
    根据检测框裁剪并保存
    xyxy: [x1, y1, x2, y2]
    """
    x1, y1, x2, y2 = map(int, xyxy)

    h, w = img.shape[:2]
    x1 = max(0, min(x1, w - 1))
    x2 = max(0, min(x2, w - 1))
    y1 = max(0, min(y1, h - 1))
    y2 = max(0, min(y2, h - 1))

    if x2 <= x1 or y2 <= y1:
        return False

    crop = img[y1:y2, x1:x2]
    if crop.size == 0:
        return False

    cv2.imwrite(str(save_path), crop)
    return True


def run(
    weights,
    source,
    output_dir,
    imgsz=640,
    conf_thres=0.25,
    iou_thres=0.45,
    device='',
    max_det=1000,
    classes=None,
    class_num=5
):
    # 设备
    device = select_device(device)

    # 加载模型
    model = DetectMultiBackend(weights, device=device, dnn=False, fp16=False)
    stride, names, pt = model.stride, model.names, model.pt
    imgsz = check_img_size(imgsz, s=stride)

    # 创建输出目录和类别子目录
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    for i in range(class_num):
        (output_dir / str(i)).mkdir(parents=True, exist_ok=True)

    # 收集图片
    source = Path(source)
    if not source.exists():
        raise FileNotFoundError(f"输入文件夹不存在: {source}")

    image_files = [p for p in source.iterdir() if p.is_file() and is_image_file(p.name)]
    image_files.sort()

    if len(image_files) == 0:
        print(f"在文件夹中未找到图片: {source}")
        return

    print(f"共找到 {len(image_files)} 张图片")
    print(f"类别名称: {names}")

    total_crops = 0

    for img_idx, img_path in enumerate(image_files, 1):
        try:
            im0, im = load_image(img_path, img_size=imgsz, stride=stride)

            # 转tensor
            im = torch.from_numpy(im).to(device)
            im = im.float() / 255.0
            if im.ndimension() == 3:
                im = im.unsqueeze(0)

            # 推理
            pred = model(im)

            # NMS
            pred = non_max_suppression(
                pred,
                conf_thres,
                iou_thres,
                classes=classes,
                agnostic=False,
                max_det=max_det
            )

            det = pred[0]

            crop_count_this_image = 0

            if len(det):
                # 将框映射回原图尺寸
                det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()

                # 遍历每个检测目标
                for det_idx, (*xyxy, conf, cls) in enumerate(det):
                    cls_id = int(cls.item())

                    # 只保存 0~4 这5类
                    if cls_id < 0 or cls_id >= class_num:
                        continue

                    # 保存文件名：原图名_序号_类别_置信度.jpg
                    stem = img_path.stem
                    conf_val = float(conf.item())
                    save_name = f"{stem}_{det_idx:03d}_cls{cls_id}_{conf_val:.4f}.jpg"
                    save_path = output_dir / str(cls_id) / save_name

                    ok = save_crop(im0, xyxy, save_path)
                    if ok:
                        crop_count_this_image += 1
                        total_crops += 1

            print(f"[{img_idx}/{len(image_files)}] {img_path.name} -> 裁剪保存 {crop_count_this_image} 个目标")

        except Exception as e:
            print(f"处理失败: {img_path}，原因: {e}")

    print("\n处理完成")
    print(f"总共保存裁剪目标数: {total_crops}")
    print(f"输出目录: {output_dir.resolve()}")


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, required=True, help='训练好的YOLOv5 pt权重路径')
    parser.add_argument('--source', type=str, required=True, help='待检测图片文件夹路径')
    parser.add_argument('--output-dir', type=str, required=True, help='裁剪结果输出文件夹')
    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='推理尺寸')
    parser.add_argument('--conf-thres', type=float, default=0.25, help='置信度阈值')
    parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU阈值')
    parser.add_argument('--device', default='', help='cuda device，例如 0 或 0,1,2,3 或 cpu')
    parser.add_argument('--max-det', type=int, default=1000, help='每张图最多检测目标数')
    parser.add_argument('--class-num', type=int, default=5, help='类别总数，这里默认是5')
    parser.add_argument('--classes', nargs='+', type=int, default=None, help='只保留某些类别，例如 --classes 0 1 2')
    opt = parser.parse_args()

    # imgsz兼容处理
    if len(opt.imgsz) == 1:
        opt.imgsz = opt.imgsz[0]
    elif len(opt.imgsz) == 2:
        # yolov5这里一般接受单值，若传两个值则取最大，避免报错
        opt.imgsz = max(opt.imgsz)
    else:
        raise ValueError("imgsz 参数格式错误")

    return opt


def main(opt):
    run(
        weights=opt.weights,
        source=opt.source,
        output_dir=opt.output_dir,
        imgsz=opt.imgsz,
        conf_thres=opt.conf_thres,
        iou_thres=opt.iou_thres,
        device=opt.device,
        max_det=opt.max_det,
        classes=opt.classes,
        class_num=opt.class_num
    )


if __name__ == "__main__":
    opt = parse_opt()
    main(opt)