# import torch
# import numpy as np
# from PIL import Image
# from model import UNetInpaint
# import io

# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# model = UNetInpaint(input_channels=4, output_channels=3)
# model.load_state_dict(torch.load("inpainting_model_best.pth", map_location=device))
# model.eval().to(device)

# def preprocess(image: Image.Image, mask: Image.Image):
#     image = image.convert("RGB").resize((256, 256))  # Resize if needed
#     mask = mask.convert("L").resize((256, 256))

#     image = np.array(image).astype(np.float32) / 255.0
#     mask = np.array(mask).astype(np.float32) / 255.0
#     mask = (mask > 0.5).astype(np.float32)
#     mask = np.expand_dims(mask, axis=-1)

#     image = np.transpose(image, (2, 0, 1))
#     mask = np.transpose(mask, (2, 0, 1))

#     image = torch.tensor(image, dtype=torch.float32)
#     mask = torch.tensor(mask, dtype=torch.float32)

#     image = image * (1.0 - mask)
#     input_tensor = torch.cat([image, mask], dim=0).unsqueeze(0).to(device)

#     return input_tensor

# def predict(image: Image.Image, mask: Image.Image) -> Image.Image:
#     input_tensor = preprocess(image, mask)

#     with torch.no_grad():
#         output = model(input_tensor).squeeze(0).cpu().numpy().transpose(1, 2, 0)
#         output = np.clip(output, 0, 1)
#         out_img = Image.fromarray((output * 255).astype(np.uint8), mode="RGB")
#         return out_img


# from model import UNetInpaint
# from PIL import Image
# import torch
# import numpy as np
# import io

# model = UNetInpaint()
# model.load_state_dict(torch.load("model.pth", map_location="cpu"))
# model.eval()

# def predict(image: bytes, mask: bytes) -> bytes:
#     image = Image.open(io.BytesIO(image)).convert("RGB")
#     mask = Image.open(io.BytesIO(mask)).convert("L")

#     # preprocess
#     image_np = np.array(image).astype(np.float32) / 255.0
#     mask_np = np.array(mask).astype(np.float32) / 255.0
#     mask_np = (mask_np > 0.5).astype(np.float32)

#     mask_np = np.expand_dims(mask_np, axis=-1)
#     image_np = np.transpose(image_np, (2, 0, 1))
#     mask_np = np.transpose(mask_np, (2, 0, 1))

#     image_tensor = torch.tensor(image_np) * (1 - torch.tensor(mask_np))
#     input_tensor = torch.cat([image_tensor, torch.tensor(mask_np)], dim=0).unsqueeze(0)

#     with torch.no_grad():
#         output = model(input_tensor).squeeze(0).numpy().transpose(1, 2, 0)
#         output = (np.clip(output, 0, 1) * 255).astype(np.uint8)
#         result = Image.fromarray(output)

#     buffer = io.BytesIO()
#     result.save(buffer, format="PNG")
#     return buffer.getvalue()

# from model import UNetInpaint
# import torch
# import numpy as np
# from PIL import Image

# model = UNetInpaint(input_channels=4, output_channels=3)
# model.load_state_dict(torch.load("inpainting_model_best_wgt.pth", map_location="cpu"))

# model.eval()

# def get_output(image_pil, mask_pil):

#     if isinstance(mask_pil, np.ndarray):
#         mask_pil = Image.fromarray(mask_pil.astype(np.uint8))

#     # Now safely convert to grayscale
#     mask = np.array(mask_pil.convert("L")).astype(np.float32) / 255.0

#     # Same for the image if needed
#     if isinstance(image_pil, np.ndarray):
#         image_pil = Image.fromarray(image_pil.astype(np.uint8))

#     image = np.array(image_pil).astype(np.float32) / 255.0

#     mask = (mask > 0.5).astype(np.float32)

#     mask = np.expand_dims(mask, axis=-1)
#     image = np.transpose(image, (2, 0, 1))
#     mask = np.transpose(mask, (2, 0, 1))

#     image = torch.tensor(image) * (1 - torch.tensor(mask))
#     input_tensor = torch.cat([image, torch.tensor(mask)], dim=0).unsqueeze(0)

#     with torch.no_grad():
#         output = model(input_tensor).squeeze(0).numpy().transpose(1, 2, 0)
#         output = (np.clip(output, 0, 1) * 255).astype(np.uint8)

#     return Image.fromarray(output)


import torch
import numpy as np
import cv2
from PIL import Image
from model import UNetInpaint  # assuming your model class is here

# Load model (adjust paths and class accordingly)
device = torch.device('cpu')
model = UNetInpaint(input_channels=4, output_channels=3)
model.load_state_dict(torch.load("inpainting_model_best_wgt.pth", map_location="cpu"))
model.eval()

def preprocess(image, mask):
    # Resize or process if needed
    # image = cv2.resize(image, (256, 256)) 
    # mask = cv2.resize(mask, (256, 256))
    # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    if mask.shape[2]:
        mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)

    

    image = image.astype(np.float32) / 255.0
    mask = mask.astype(np.float32) / 255.0
    mask = (mask > 0.5).astype(np.float32)
    # mask = cv2.resize(mask, (image.shape[0], image.shape[1]))

    image = cv2.resize(image, (512, 512))
    mask = cv2.resize(mask, (512, 512))


    mask = np.expand_dims(mask, axis=-1)

    image = np.transpose(image, (2, 0, 1))
    mask = np.transpose(mask, (2, 0, 1))


    image = torch.tensor(image, dtype=torch.float32)
    mask = torch.tensor(mask, dtype=torch.float32)


    image = image * (1.0 - mask)
    input = torch.cat([image, mask], dim=0)
    input = input.unsqueeze(0).to(device)

    with torch.no_grad():
        output = model(input)

    output = output.squeeze(0).cpu().numpy().transpose(1, 2, 0)
    output = np.clip(output, 0, 1)
    out_img = Image.fromarray((output * 255).astype(np.uint8), mode='RGB')
    
    return out_img


def get_output(image_np, mask_np):
    input_tensor = preprocess(image_np, mask_np)
    # with torch.no_grad():
    #     output = model(input_tensor)
    # result = postprocess(output)
    return input_tensor