import cv2, os, json
import numpy as np
from math import atan2, cos, sin, sqrt, pi
from sklearn.cluster import DBSCAN
from sklearn.preprocessing import StandardScaler
from sklearn.cluster import OPTICS, cluster_optics_dbscan
from tqdm import tqdm
from glob import glob

def nms(boxes, thresh):
    if len(boxes) == 0:
        return []

    pick = []
    x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
    area = (x2 - x1 + 1) * (y2 - y1 + 1)
    idxs = np.argsort(y2)

    while len(idxs) > 0:
        last = len(idxs) - 1
        i = idxs[last]
        pick.append(i)

        xx1 = np.maximum(x1[i], x1[idxs[:last]])
        yy1 = np.maximum(y1[i], y1[idxs[:last]])
        xx2 = np.minimum(x2[i], x2[idxs[:last]])
        yy2 = np.minimum(y2[i], y2[idxs[:last]])

        w = np.maximum(0, xx2 - xx1 + 1)
        h = np.maximum(0, yy2 - yy1 + 1)

        overlap = (w * h) / area[idxs[:last]]
        idxs = np.delete(idxs, np.concatenate(([last], np.where(overlap > thresh)[0])))

    return boxes[pick]

def find_pts_on_line(og, slope, d):
    cx, cy = og
    x1 = float(cx - d / ((1 + slope**2)**0.5))
    y1 = float(cy - slope * cx + x1 * slope)
    return (x1, y1)

def cluster(img_path, im, save_dir):
    img = cv2.imread(img_path)
    imgray = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)

    ret, binary_map = cv2.threshold(imgray, 127, 255, 0)

    nlabels, labels, stats, centroids = cv2.connectedComponentsWithStats(binary_map, None, None, None, 8, cv2.CV_32S)
    areas = stats[1:, cv2.CC_STAT_AREA]
    result = np.zeros((labels.shape), np.uint8)
    for i in range(0, nlabels - 1):
        if areas[i] >= 250:
            result[labels == i + 1] = 255
    re_copy = result.copy()

    edgeimg = cv2.Canny(result, 10, 150)

    size = np.size(result)
    skel = np.zeros(result.shape, np.uint8)
    element = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))

    while True:
        open_ = cv2.morphologyEx(result, cv2.MORPH_OPEN, element)
        temp = cv2.subtract(result, open_)
        eroded = cv2.erode(result, element)
        skel = cv2.bitwise_or(skel, temp)
        result = eroded.copy()
        if cv2.countNonZero(result) == 0:
            break

    nlabels, labels, stats, centroids = cv2.connectedComponentsWithStats(skel, None, None, None, 8, cv2.CV_32S)
    areas = stats[1:, cv2.CC_STAT_AREA]
    skel = np.zeros((labels.shape), np.uint8)
    for i in range(0, nlabels - 1):
        if areas[i] >= 2:
            skel[labels == i + 1] = 255

    # Lưu riêng cho mỗi ảnh
    base_name = os.path.splitext(im)[0]
    os.makedirs(save_dir, exist_ok=True)
    cv2.imwrite(os.path.join(save_dir, f"Skeleton_{base_name}.png"), skel)

    combined_img = cv2.addWeighted(skel, 1, edgeimg, 1, 0)
    cv2.imwrite(os.path.join(save_dir, f"SkeletonContour_{base_name}.png"), combined_img)

    filter_size = (10, 10)
    white_pixels = np.where(skel == 255)
    x_coords, y_coords = white_pixels[1], white_pixels[0]

    x1 = x_coords - filter_size[0] // 2
    y1 = y_coords - filter_size[1] // 2
    x2 = x_coords + filter_size[0] // 2
    y2 = y_coords + filter_size[1] // 2

    white_regions = np.column_stack((x1, y1, x2, y2))
    white_regions = nms(white_regions, thresh=0.1)

    skeleton_image = cv2.cvtColor(skel.copy(), cv2.COLOR_GRAY2BGR)
    filtered_image = cv2.cvtColor(re_copy.copy(), cv2.COLOR_GRAY2BGR)
    center_points, directions = [], []

    def get_direction2(bbox_pixels):
        nonzero_indices = np.column_stack(np.nonzero(bbox_pixels))
        nonzero_indices = np.float32(nonzero_indices)
        if len(nonzero_indices) >= 2:
            mean, eigenvectors = cv2.PCACompute(nonzero_indices, mean=None)
            cntr = ((mean[0, 1]), (mean[0, 0]))
            return eigenvectors[0], cntr
        else:
            return (0, 0), (0, 0)

    for coor in white_regions:
        x1, y1, x2, y2 = coor
        bbox_pixels = skel[int(y1):int(y2), int(x1):int(x2)]
        direction, mean = get_direction2(bbox_pixels)
        directions.append(direction)
        center_points.append((mean[0] + x1, mean[1] + y1))
        cv2.rectangle(skeleton_image, (int(x1), int(y1)), (int(x2), int(y2)), (255, 255, 255), 1)
        cv2.circle(skeleton_image, (int(mean[0] + x1), int(mean[1] + y1)), 1, (255, 0, 0), -1)

    cv2.imwrite(os.path.join(save_dir, f"Result_{base_name}.png"), skeleton_image)

    pts_group, bbox_group = [], []
    for idx, pts in enumerate(center_points):
        if 640 > pts[0] > 0 and 480 > pts[1] > 0:
            pts_group.append([int(pts[0]), int(pts[1])])
            x1, y1, x2, y2 = white_regions[idx]
            bbox_group.append([int(x1), int(y1), int(x2), int(y2)])

    return pts_group, bbox_group


mask_dir = "datasets/seg_train"
save_json_dir = "datasets"
save_img_dir = os.path.join(save_json_dir, "output")

patterns = ['*.png', '*.jpg', '*.jpeg', '*.PNG', '*.JPG', '*.JPEG']
files = []
for p in patterns:
    files.extend(glob(os.path.join(mask_dir, p)))
files = sorted(set(files))
print(f"Found {len(files)} files in {mask_dir}")

file_dict = {}
bbox_dict = {}

for filepath in tqdm(files):
    filename = os.path.basename(filepath)
    pts, bbox = cluster(filepath, filename, save_img_dir)
    if len(pts) != 0:
        file_dict[filename] = pts
        bbox_dict[filename] = bbox

with open(os.path.join(save_json_dir, 'train_seg_points.json'), 'w') as json_file:
    json.dump(file_dict, json_file)

with open(os.path.join(save_json_dir, 'train_bbox_points.json'), 'w') as json_file:
    json.dump(bbox_dict, json_file)