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
import json


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 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))
    print(f"scaled size {scaled_size}")

    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.2,
) -> 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
    spline_degree = 3  # Typically k=3 for cubic spline

    smoothed_x_sigma = 1
    smoothed_y_sigma = 1

    # Ensure contour has enough points
    if len(contour) < spline_degree + 1:
        raise ValueError(f"Contour must have at least {spline_degree + 1} points, but has {len(contour)} points.")

    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):
    degree = 3
    closed = True

    # Create a new DXF document with millimeters as the unit
    doc = ezdxf.new(units=ezdxf.units.MM)
    doc.units = ezdxf.units.MM  # Ensure units are millimeters
    doc.header["$INSUNITS"] = ezdxf.units.MM  # Set insertion units to millimeters

    msp = doc.modelspace()

    for contour in inflated_contours:
        try:
            resampled_contour = resample_contour(contour)
            points = [
                (x * scaling_factor, (height - y) * scaling_factor)
                for x, y in resampled_contour
            ]
            if len(points) >= 3:
                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

        except ValueError as e:
            print(f"Skipping contour: {e}")

    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
    )

    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 to_dxf(contours):
    # Create a new DXF document with millimeters as the unit
    doc = ezdxf.new(units=ezdxf.units.MM)
    doc.units = ezdxf.units.MM  # Ensure units are millimeters
    doc.header["$INSUNITS"] = ezdxf.units.MM  # Set insertion units to millimeters)
    msp = doc.modelspace()

    try:
        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
    except Exception as e:
        raise gr.Error(f"Unable to generate DXF: {e}")

    output_path = "./outputs/out.dxf"
    doc.saveas(output_path)
    return output_path

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("./best1.pt")
    res = box_detector.predict(img, conf=0.05)
    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, coin_size_mm):

    if offset < 0:
        raise gr.Error("Offset Value Can't be negative")

    try:
        reference_obj_img, scaling_box_coords = detect_reference_square(image)
    except:
        raise gr.Error("Unable to detect the COIN. Please try again with different magnification.")

    reference_obj_img = make_square(reference_obj_img)

    reference_square_mask = remove_bg(reference_obj_img)

    reference_square_mask = resize_img(reference_square_mask, (reference_obj_img.shape[1], reference_obj_img.shape[0]))

    try:
        scaling_factor= calculate_scaling_factor(
            target_image=reference_square_mask,
            reference_obj_size_mm = coin_size_mm,
            feature_detector="ORB",
        )
    except Exception as e:
        scaling_factor = None
        print(f"Error calculating scaling factor: {e}")

    # Default to a scaling factor if calculation fails
    if scaling_factor is None or scaling_factor == 0:
        scaling_factor = 0.07
        print("Using default scaling factor due to calculation error")

    orig_size = image.shape[:2]
    objects_mask = remove_bg(image)
    processed_size = objects_mask.shape[:2]

    objects_mask = exclude_scaling_box(
        objects_mask,
        scaling_box_coords,
        orig_size,
        processed_size,
        expansion_factor=1.2,
    )
    objects_mask = resize_img(objects_mask, (image.shape[1], image.shape[0]))
    
    # Ensure offset_inches is valid
    if scaling_factor != 0:
        offset_pixels = (float(offset) / float(scaling_factor)) * 2 + 1
    else:
        offset_pixels = 1  # Default value in case of invalid scaling factor

    dilated_mask = cv2.dilate(objects_mask, np.ones((int(offset_pixels), int(offset_pixels)), np.uint8))

    Image.fromarray(dilated_mask).save("./outputs/scaled_mask_new.jpg")
    outlines, contours = extract_outlines(dilated_mask)
    shrunked_img_contours = cv2.drawContours(image, contours, -1, (0, 0, 255), thickness=2)
    dxf = save_dxf_spline(contours, scaling_factor, processed_size[0])
    # dxf = to_dxf(contours)

    return (
        shrunked_img_contours,
        outlines,
        dxf,
        dilated_mask,
        scaling_factor,
    )

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

    # Language selector in UI
    ifer = gr.Interface(
        fn=predict,
        inputs=[
            gr.Image(label="Input Image", type="numpy"),
            gr.Number(label="Offset value for Mask(mm)", value=0.15),
            gr.Number(label="Diameter of reference coin(mm). Adjust according to coin.", value=20),
        
        ],
        outputs=[
            gr.Image(label="Output 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 millimeters",
            ),
        ],
        examples=[
            ["./examples/Test20.jpg", 0.15],
            ["./examples/Test21.jpg", 0.15],
            ["./examples/Test22.jpg", 0.15],
            ["./examples/Test23.jpg", 0.15],
        ],
    )
    ifer.launch(share=True)