import argparse
import cv2
from ultralytics import YOLO
from FastSAM.tools import *
from groundingdino.util.inference import load_model, load_image, predict, annotate, Model
from torchvision.ops import box_convert
import ast

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_path", type=str, default="./FastSAM/FastSAM-x.pt", help="model"
    )
    parser.add_argument(
        "--img_path", type=str, default="./images/dogs.jpg", help="path to image file"
    )
    parser.add_argument(
        "--text", type=str, default="the black dog.", help="text prompt for GroundingDINO"
    )
    parser.add_argument("--imgsz", type=int, default=1024, help="image size")
    parser.add_argument(
        "--iou",
        type=float,
        default=0.9,
        help="iou threshold for filtering the annotations",
    )
    parser.add_argument(
        "--conf", type=float, default=0.4, help="object confidence threshold"
    )
    parser.add_argument(
        "--output", type=str, default="./output/", help="image save path"
    )
    parser.add_argument(
        "--randomcolor", type=bool, default=True, help="mask random color"
    )
    parser.add_argument(
        "--point_prompt", type=str, default="[[0,0]]", help="[[x1,y1],[x2,y2]]"
    )
    parser.add_argument(
        "--point_label",
        type=str,
        default="[0]",
        help="[1,0] 0:background, 1:foreground",
    )
    parser.add_argument("--box_prompt", type=str, default="[0,0,0,0]", help="[x,y,w,h]")
    parser.add_argument(
        "--better_quality",
        type=str,
        default=False,
        help="better quality using morphologyEx",
    )
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument(
        "--device", type=str, default=device, help="cuda:[0,1,2,3,4] or cpu"
    )
    parser.add_argument(
        "--retina",
        type=bool,
        default=True,
        help="draw high-resolution segmentation masks",
    )
    parser.add_argument(
        "--withContours", type=bool, default=False, help="draw the edges of the masks"
    )
    return parser.parse_args()


def main(args):

    # Image Path
    img_path = args.img_path
    text = args.text

    # path to save img
    save_path = args.output
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    basename = os.path.basename(args.img_path).split(".")[0]

    # Build Fast-SAM Model
    # ckpt_path = "/comp_robot/rentianhe/code/Grounded-Segment-Anything/FastSAM/FastSAM-x.pt"
    model = YOLO(args.model_path)

    results = model(
        args.img_path,
        imgsz=args.imgsz,
        device=args.device,
        retina_masks=args.retina,
        iou=args.iou,
        conf=args.conf,
        max_det=100,
    )


    # Build GroundingDINO Model
    groundingdino_config = "GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py"
    groundingdino_ckpt_path = "./groundingdino_swint_ogc.pth"

    image_source, image = load_image(img_path)
    model = load_model(groundingdino_config, groundingdino_ckpt_path)

    boxes, logits, phrases = predict(
        model=model,
        image=image,
        caption=text,
        box_threshold=0.3,
        text_threshold=0.25,
        device=args.device,
    )


    # Grounded-Fast-SAM

    ori_img = cv2.imread(img_path)
    ori_h = ori_img.shape[0]
    ori_w = ori_img.shape[1]

    # Save each frame due to the post process from FastSAM
    boxes = boxes * torch.Tensor([ori_w, ori_h, ori_w, ori_h])
    print(f"Detected Boxes: {len(boxes)}")
    boxes = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").cpu().numpy().tolist()
    for box_idx in range(len(boxes)):
        mask, _ = box_prompt(
            results[0].masks.data,
            boxes[box_idx],
            ori_h,
            ori_w,
        )
        annotations = np.array([mask])
        img_array = fast_process(
            annotations=annotations,
            args=args,
            mask_random_color=True,
            bbox=boxes[box_idx],
        )
        cv2.imwrite(os.path.join(save_path, basename + f"_{str(box_idx)}_caption_{phrases[box_idx]}.jpg"), cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR))


if __name__ == "__main__":
    args = parse_args()
    main(args)