import os
from pathlib import Path
from typing import List, Union
from PIL import Image
import ezdxf.units
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
from scipy.interpolate import splprep, splev
from scipy.ndimage import gaussian_filter1d

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]),
    ]
)


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:
    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: Image = transforms.ToPILImage()(pred)
    print(pred_pil)
    # 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))

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


def make_square(img: np.ndarray):
    # Get dimensions
    height, width = img.shape[:2]

    # Find the larger dimension
    max_dim = max(height, width)

    # Calculate padding
    pad_height = (max_dim - height) // 2
    pad_width = (max_dim - width) // 2

    # Handle odd dimensions
    pad_height_extra = max_dim - height - 2 * pad_height
    pad_width_extra = max_dim - width - 2 * pad_width

    # Create padding with edge colors
    if len(img.shape) == 3:  # Color image
        # Pad the image
        padded = np.pad(
            img,
            (
                (pad_height, pad_height + pad_height_extra),
                (pad_width, pad_width + pad_width_extra),
                (0, 0),
            ),
            mode="edge",
        )
    else:  # Grayscale image
        padded = np.pad(
            img,
            (
                (pad_height, pad_height + pad_height_extra),
                (pad_width, pad_width + pad_width_extra),
            ),
            mode="edge",
        )

    return padded


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 resample_contour(contour):
    # ---------------------------------------------------------------------------------------- #
    # Get all the parameters at the start:
    num_points = 1000
    smoothing_factor = 5

    smoothed_x_sigma = 1
    smoothed_y_sigma = 1
    # ---------------------------------------------------------------------------------------- #
    contour = contour[:, 0, :]

    tck, _ = splprep([contour[:, 0], contour[:, 1]], s=smoothing_factor)

    u = np.linspace(0, 1, num_points)
    resampled_points = splev(u, tck)

    smoothed_x = gaussian_filter1d(resampled_points[0], sigma=smoothed_x_sigma)
    smoothed_y = gaussian_filter1d(resampled_points[1], sigma=smoothed_y_sigma)

    return np.array([smoothed_x, smoothed_y]).T


def save_dxf_spline(inflated_contours, scaling_factor, height):
    # ---------------------------------------------------------------------------------------- #
    # Get all the parameters at the start:
    degree = 3
    closed = True
    # ---------------------------------------------------------------------------------------- #

    doc = ezdxf.new(units=0)
    doc.units = ezdxf.units.IN
    doc.header["$INSUNITS"] = ezdxf.units.IN

    msp = doc.modelspace()

    for contour in inflated_contours:
        resampled_contour = resample_contour(contour)
        points = [
            (x * scaling_factor, (height - y) * scaling_factor)
            for x, y in resampled_contour
        ]
        if len(points) >= 3:
            # Manually Closing the Contour in case it hasn't been closed by the contours before.
            if np.linalg.norm(np.array(points[0]) - np.array(points[-1])) > 1e-2:
                points.append(points[0])

            spline = msp.add_spline(points, degree=degree)
            spline.closed = closed

    # Step 14: Save the DXF file
    dxf_filepath = os.path.join("./outputs", "out.dxf")
    doc.saveas(dxf_filepath)
    return dxf_filepath


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_NONE
    )

    # smooth_contours_list = []
    # for contour in contours:
    #     smooth_contours_list.append(smooth_contours(contour))
    # Create a blank image to draw contours
    outline_image = np.zeros_like(binary_image)

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

    return cv2.bitwise_not(outline_image), 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 resize_img(img: np.ndarray, resize_dim):
    return np.array(Image.fromarray(img).resize(resize_dim))


def predict(image, offset_inches):
    try:
        drawer_img = yolo_detect(image, ["box"])
        shrunked_img = make_square(shrink_bbox(drawer_img, 0.8))
    except:
        raise gr.Error("Unable to DETECT DRAWER, please take another picture with different magnification level!")

    # Detect the scaling reference square
    try:
        reference_obj_img, scaling_box_coords = detect_reference_square(shrunked_img)
    except:
        raise gr.Error("Unable to DETECT REFERENCE BOX, please take another picture with different magnification level!")

    # reference_obj_img_scaled = shrink_bbox(reference_obj_img, 1.2)
    # make the image sqaure so it does not effect the size of objects
    reference_obj_img = make_square(reference_obj_img)
    reference_square_mask = remove_bg(reference_obj_img)

    # make the mask same size as org image
    reference_square_mask = resize_img(
        reference_square_mask, (reference_obj_img.shape[1], reference_obj_img.shape[0])
    )

    try:
        scaling_factor = calculate_scaling_factor(
            reference_image_path="./Reference_ScalingBox.jpg",
            target_image=reference_square_mask,
            feature_detector="ORB",
        )
    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,
    )
    objects_mask = resize_img(
        objects_mask, (shrunked_img.shape[1], shrunked_img.shape[0])
    )
    offset_pixels = (offset_inches / scaling_factor) * 2 + 1
    dilated_mask = cv2.dilate(
        objects_mask, np.ones((int(offset_pixels), int(offset_pixels)), np.uint8)
    )

    # Scale the object mask according to scaling factor
    # objects_mask_scaled = scale_image(objects_mask, scaling_factor)
    Image.fromarray(dilated_mask).save("./outputs/scaled_mask_new.jpg")
    outlines, contours = extract_outlines(dilated_mask)
    shrunked_img_contours = cv2.drawContours(
        shrunked_img, contours, -1, (0, 0, 255), thickness=2
    )
    dxf = save_dxf_spline(contours, scaling_factor, processed_size[0])

    return (
        cv2.cvtColor(shrunked_img_contours, cv2.COLOR_BGR2RGB),
        outlines,
        dxf,
        dilated_mask,
        scaling_factor,
    )


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

    ifer = gr.Interface(
        fn=predict,
        inputs=[
            gr.Image(label="Input Image"),
            gr.Number(label="Offset value for Mask(inches)", value=0.075),
        ],
        outputs=[
            gr.Image(label="Ouput Image"),
            gr.Image(label="Outlines of Objects"),
            gr.File(label="DXF file"),
            gr.Image(label="Mask"),
            gr.Textbox(
                label="Scaling Factor(mm)",
                placeholder="Every pixel is equal to mentioned number in inches",
            ),
        ],
        examples=[
            ["./examples/Test20.jpg", 0.075],
            ["./examples/Test21.jpg", 0.075],
            ["./examples/Test22.jpg", 0.075],
            ["./examples/Test23.jpg", 0.075],
        ],
    )
    ifer.launch(share=True)