import os
from pathlib import Path
from typing import List, Union
from PIL import Image
import numpy as np
import torch
from torchvision import transforms
from ultralytics import YOLOWorld, YOLO
from ultralytics.engine.results import Results
from ultralytics.utils.plotting import save_one_box
from transformers import AutoModelForImageSegmentation
import cv2
import ezdxf
import gradio as gr
import gc
from scalingtestupdated import calculate_scaling_factor



def yolo_detect(
    image: Union[str, Path, int, Image.Image, list, tuple, np.ndarray, torch.Tensor],
    classes: List[str],
) -> np.ndarray:
    drawer_detector = YOLOWorld("yolov8x-worldv2.pt")
    drawer_detector.set_classes(classes)
    results: List[Results] = drawer_detector.predict(image)
    boxes = []
    for result in results:
        boxes.append(
            save_one_box(result.cpu().boxes.xyxy, im=result.orig_img, save=False)
        )

    del drawer_detector

    return boxes[0]


def remove_bg(image: np.ndarray) -> np.ndarray:
    birefnet = AutoModelForImageSegmentation.from_pretrained(
    "zhengpeng7/BiRefNet", trust_remote_code=True
    )

    device = "cpu"
    torch.set_float32_matmul_precision(["high", "highest"][0])

    birefnet.to(device)
    birefnet.eval()
    transform_image = transforms.Compose(
        [
            transforms.Resize((1024, 1024)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    )

    image = Image.fromarray(image)
    input_images = transform_image(image).unsqueeze(0).to("cpu")

    # Prediction
    with torch.no_grad():
        preds = birefnet(input_images)[-1].sigmoid().cpu()
    pred = preds[0].squeeze()

    # Show Results
    pred_pil = transforms.ToPILImage()(pred)
    # Scale proportionally with max length to 1024 for faster showing
    scale_ratio = 1024 / max(image.size)
    scaled_size = (int(image.size[0] * scale_ratio), int(image.size[1] * scale_ratio))
    
    del birefnet

    return np.array(pred_pil.resize(scaled_size))

def exclude_scaling_box(image: np.ndarray, bbox: np.ndarray, orig_size: tuple, processed_size: tuple, expansion_factor: float = 1.5) -> np.ndarray:
    # Unpack the bounding box
    x_min, y_min, x_max, y_max = map(int, bbox)

    # Calculate scaling factors
    scale_x = processed_size[1] / orig_size[1]  # Width scale
    scale_y = processed_size[0] / orig_size[0]  # Height scale

    # Adjust bounding box coordinates
    x_min = int(x_min * scale_x)
    x_max = int(x_max * scale_x)
    y_min = int(y_min * scale_y)
    y_max = int(y_max * scale_y)

    # Calculate expanded box coordinates
    box_width = x_max - x_min
    box_height = y_max - y_min
    expanded_x_min = max(0, int(x_min - (expansion_factor - 1) * box_width / 2))
    expanded_x_max = min(image.shape[1], int(x_max + (expansion_factor - 1) * box_width / 2))
    expanded_y_min = max(0, int(y_min - (expansion_factor - 1) * box_height / 2))
    expanded_y_max = min(image.shape[0], int(y_max + (expansion_factor - 1) * box_height / 2))

    # Black out the expanded region
    image[expanded_y_min:expanded_y_max, expanded_x_min:expanded_x_max] = 0

    return image


def extract_outlines(binary_image: np.ndarray) -> np.ndarray:
    """
    Extracts and draws the outlines of masks from a binary image.
    Args:
        binary_image: Grayscale binary image where white represents masks and black is the background.
    Returns:
        Image with outlines drawn.
    """
    # Detect contours from the binary image
    contours, _ = cv2.findContours(
        binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )

    # Create a blank image to draw contours
    outline_image = np.zeros_like(binary_image)
    # Smooth the contours
    smoothed_contours = []
    for contour in contours:
        # Calculate epsilon for approxPolyDP
        epsilon = 0.002 * cv2.arcLength(contour, True)
        # Approximate the contour with fewer points
        smoothed_contour = cv2.approxPolyDP(contour, epsilon, True)
        smoothed_contours.append(smoothed_contour)

    # Draw the contours on the blank image
    cv2.drawContours(
        outline_image, smoothed_contours, -1, (255), thickness=1
    )  # White color for outlines

    return cv2.bitwise_not(outline_image), smoothed_contours


def shrink_bbox(image: np.ndarray, shrink_factor: float):
    """
    Crops the central 80% of the image, maintaining proportions for non-square images.
    Args:
        image: Input image as a NumPy array.
    Returns:
        Cropped image as a NumPy array.
    """
    height, width = image.shape[:2]
    center_x, center_y = width // 2, height // 2

    # Calculate 80% dimensions
    new_width = int(width * shrink_factor)
    new_height = int(height * shrink_factor)

    # Determine the top-left and bottom-right points for cropping
    x1 = max(center_x - new_width // 2, 0)
    y1 = max(center_y - new_height // 2, 0)
    x2 = min(center_x + new_width // 2, width)
    y2 = min(center_y + new_height // 2, height)

    # Crop the image
    cropped_image = image[y1:y2, x1:x2]
    return cropped_image


# def to_dxf(outlines):
#     upper_range_tuple = (200)
#     lower_range_tuple = (0)

#     doc = ezdxf.new('R2010')
#     msp = doc.modelspace() 
#     masked_jpg = cv2.inRange(outlines,lower_range_tuple, upper_range_tuple)

#     for i in range(0,masked_jpg.shape[0]):
#         for j in range(0,masked_jpg.shape[1]):
#             if masked_jpg[i][j] == 255:
#                 msp.add_line((j,masked_jpg.shape[0] - i), (j,masked_jpg.shape[0] - i))

#     doc.saveas("./outputs/out.dxf")
#     return "./outputs/out.dxf"

def to_dxf(contours):
    doc = ezdxf.new()
    msp = doc.modelspace()

    for contour in contours:
        points = [(point[0][0], point[0][1]) for point in contour]
        msp.add_lwpolyline(points, close=True)  # Add a polyline for each contour
    
    doc.saveas("./outputs/out.dxf")
    return "./outputs/out.dxf"

def smooth_contours(contour):
    epsilon = 0.01 * cv2.arcLength(contour, True)  # Adjust factor (e.g., 0.01)
    return cv2.approxPolyDP(contour, epsilon, True)


def scale_image(image: np.ndarray, scale_factor: float) -> np.ndarray:
    """
    Resize image by scaling both width and height by the same factor.

    Args:
        image: Input numpy image
        scale_factor: Factor to scale the image (e.g., 0.5 for half size, 2 for double size)

    Returns:
        np.ndarray: Resized image
    """
    if scale_factor <= 0:
        raise ValueError("Scale factor must be positive")

    current_height, current_width = image.shape[:2]

    # Calculate new dimensions
    new_width = int(current_width * scale_factor)
    new_height = int(current_height * scale_factor)

    # Choose interpolation method based on whether we're scaling up or down
    interpolation = cv2.INTER_AREA if scale_factor < 1 else cv2.INTER_CUBIC

    # Resize image
    resized_image = cv2.resize(
        image, (new_width, new_height), interpolation=interpolation
    )

    return resized_image


def detect_reference_square(img) -> np.ndarray:
    box_detector = YOLO("./last.pt")
    res = box_detector.predict(img)
    del box_detector
    return save_one_box(res[0].cpu().boxes.xyxy, res[0].orig_img, save=False), res[0].cpu().boxes.xyxy[0
                                                                                                       ]

def predict(image):
    drawer_img = yolo_detect(image, ["box"])
    shrunked_img = shrink_bbox(drawer_img, 0.8)
    # Detect the scaling reference square
    reference_obj_img, scaling_box_coords = detect_reference_square(shrunked_img)
    reference_obj_img_scaled = shrink_bbox(reference_obj_img, 1.2)
    try:
        scaling_factor = calculate_scaling_factor(
            reference_image_path="./Reference_ScalingBox.jpg",
            target_image=reference_obj_img_scaled,
            feature_detector="SIFT",
        )
    except:
        scaling_factor = 1.0
    # Save original size before `remove_bg` processing
    orig_size = shrunked_img.shape[:2]
    # Generate foreground mask and save its size
    objects_mask = remove_bg(shrunked_img)
    processed_size = objects_mask.shape[:2]
    # Exclude scaling box region from objects mask
    objects_mask = exclude_scaling_box(
        objects_mask, scaling_box_coords, orig_size, processed_size, expansion_factor=3.0
    )
    # Scale the object mask according to scaling factor
    # objects_mask_scaled = scale_image(objects_mask, scaling_factor)
    Image.fromarray(objects_mask).save("./outputs/scaled_mask_new.jpg")
    outlines, contours = extract_outlines(objects_mask)
    dxf = to_dxf(contours)

    return outlines, dxf, objects_mask, scaling_factor, reference_obj_img_scaled



if __name__ == "__main__":
    os.makedirs("./outputs", exist_ok=True)

    ifer = gr.Interface(
        fn=predict,
        inputs=[gr.Image(label="Input Image")],
        outputs=[
            gr.Image(label="Ouput Image"),
            gr.File(label="DXF file"),
            gr.Image(label="Mask"),
            gr.Textbox(label="Scaling Factor(mm)", placeholder="Every pixel is equal to mentioned number in mm(milimeter)"),
            gr.Image(label="Image used for calculating scaling factor")
        ],
        examples=["./examples/Test20.jpg", "./examples/Test21.jpg", "./examples/Test22.jpg", "./examples/Test23.jpg"]
    )
    ifer.launch(share=True)