import os
import zipfile
import json
import numpy as np
import cv2
from pathlib import Path

out_dir = Path("sample_cv_test_datasets")
out_dir.mkdir(exist_ok=True)

def create_sample_img(filename, color=(0, 200, 100)):
    img = np.zeros((300, 300, 3), dtype=np.uint8)
    cv2.rectangle(img, (50, 50), (250, 250), color, -1)
    cv2.putText(img, "TEST DATASET", (60, 160), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
    cv2.imwrite(filename, img)

# 1. Instance Segmentation ZIP Dataset
seg_dir = out_dir / "sample_instance_segmentation"
seg_dir.mkdir(exist_ok=True)
create_sample_img(str(seg_dir / "car_sample.jpg"), (220, 50, 50))
coco_seg = {
    "images": [{"id": 1, "file_name": "car_sample.jpg", "height": 300, "width": 300}],
    "categories": [{"id": 1, "name": "vehicle_segmentation"}],
    "annotations": [{
        "id": 1, "image_id": 1, "category_id": 1,
        "segmentation": [[50, 50, 250, 50, 250, 250, 50, 250]],
        "bbox": [50, 50, 200, 200], "area": 40000, "iscrowd": 0
    }]
}
with open(seg_dir / "_annotations.coco.json", "w") as f:
    json.dump(coco_seg, f, indent=2)

with zipfile.ZipFile(out_dir / "Instance_Segmentation_Test_Dataset.zip", "w") as z:
    z.write(seg_dir / "car_sample.jpg", "car_sample.jpg")
    z.write(seg_dir / "_annotations.coco.json", "_annotations.coco.json")

# 2. Keypoint Detection ZIP Dataset
pose_dir = out_dir / "sample_keypoint_pose"
pose_dir.mkdir(exist_ok=True)
create_sample_img(str(pose_dir / "person_pose.jpg"), (50, 100, 250))
coco_pose = {
    "images": [{"id": 1, "file_name": "person_pose.jpg", "height": 300, "width": 300}],
    "categories": [{"id": 1, "name": "human_pose", "keypoints": ["nose", "left_eye", "right_eye", "left_shoulder", "right_shoulder"]}],
    "annotations": [{
        "id": 1, "image_id": 1, "category_id": 1,
        "keypoints": [150, 100, 2, 140, 90, 2, 160, 90, 2, 100, 180, 2, 200, 180, 2],
        "bbox": [50, 50, 200, 200], "num_keypoints": 5, "area": 40000, "iscrowd": 0
    }]
}
with open(pose_dir / "_annotations.coco.json", "w") as f:
    json.dump(coco_pose, f, indent=2)

with zipfile.ZipFile(out_dir / "Keypoint_Detection_Test_Dataset.zip", "w") as z:
    z.write(pose_dir / "person_pose.jpg", "person_pose.jpg")
    z.write(pose_dir / "_annotations.coco.json", "_annotations.coco.json")

# 3. OCR Text Extraction ZIP Dataset
ocr_dir = out_dir / "sample_ocr_dataset"
ocr_dir.mkdir(exist_ok=True)
create_sample_img(str(ocr_dir / "ocr_text.jpg"), (10, 180, 220))
coco_ocr = {
    "images": [{"id": 1, "file_name": "ocr_text.jpg", "height": 300, "width": 300}],
    "categories": [{"id": 1, "name": "text_block"}],
    "annotations": [{
        "id": 1, "image_id": 1, "category_id": 1,
        "text": "TEST DATASET",
        "bbox": [60, 140, 180, 40], "area": 7200, "iscrowd": 0
    }]
}
with open(ocr_dir / "_annotations.coco.json", "w") as f:
    json.dump(coco_ocr, f, indent=2)

with zipfile.ZipFile(out_dir / "OCR_Text_Test_Dataset.zip", "w") as z:
    z.write(ocr_dir / "ocr_text.jpg", "ocr_text.jpg")
    z.write(ocr_dir / "_annotations.coco.json", "_annotations.coco.json")

print("Generated sample benchmark test datasets cleanly.")