import os
import cv2
import torch
import numpy as np
from PIL import Image, ImageFilter
from diffusers import StableDiffusionInpaintPipeline


IMAGE_PATH = "shirt.png"
OUTPUT_DIR = "output_refined_cloth"
MODEL_ID = "runwayml/stable-diffusion-inpainting"

os.makedirs(OUTPUT_DIR, exist_ok=True)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


def save_img(arr, path):
    Image.fromarray(arr.astype(np.uint8)).save(path)


def largest_component(mask):
    n, labels, stats, _ = cv2.connectedComponentsWithStats(mask, 8)
    if n <= 1:
        return mask

    idx = 1 + np.argmax(stats[1:, cv2.CC_STAT_AREA])
    out = np.zeros_like(mask)
    out[labels == idx] = 255
    return out


def segment_cloth(image_np):
    gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)

    # 흰 배경 기준
    mask = (gray < 245).astype(np.uint8) * 255

    kernel = np.ones((5, 5), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1)

    return largest_component(mask)


def estimate_center_axis(mask):
    ys, xs = np.where(mask > 0)
    if len(xs) == 0:
        raise RuntimeError("옷 마스크가 비어 있습니다.")

    # 단순 평균보다 median이 안정적
    center_x = int(np.median(xs))

    return center_x


def symmetry_full_mask(mask):
    h, w = mask.shape
    center_x = estimate_center_axis(mask)

    full = mask.copy()

    for y in range(h):
        row = mask[y]
        xs = np.where(row > 0)[0]

        if len(xs) < 2:
            continue

        left = xs.min()
        right = xs.max()

        dist_left = center_x - left
        dist_right = right - center_x

        target_half = max(dist_left, dist_right)

        new_left = max(0, center_x - target_half)
        new_right = min(w - 1, center_x + target_half)

        full[y, new_left:new_right + 1] = 255

    kernel = np.ones((9, 9), np.uint8)
    full = cv2.morphologyEx(full, cv2.MORPH_CLOSE, kernel, iterations=2)

    full = largest_component(full)

    return full


def make_missing_mask(visible_mask, full_mask):
    missing = cv2.subtract(full_mask, visible_mask)

    # 너무 작은 잡영역 제거
    n, labels, stats, _ = cv2.connectedComponentsWithStats(missing, 8)

    clean = np.zeros_like(missing)

    for i in range(1, n):
        area = stats[i, cv2.CC_STAT_AREA]
        if area > 500:
            clean[labels == i] = 255

    # 너무 과도한 확장 방지
    kernel = np.ones((13, 13), np.uint8)
    clean = cv2.dilate(clean, kernel, iterations=1)

    # full mask 밖으로 나가지 않게 제한
    clean = cv2.bitwise_and(clean, full_mask)

    return clean


def mirror_fill(image_np, visible_mask, missing_mask):
    h, w = visible_mask.shape
    center_x = estimate_center_axis(visible_mask)

    filled = image_np.copy()

    missing_ys, missing_xs = np.where(missing_mask > 0)

    for y, x in zip(missing_ys, missing_xs):
        mirror_x = 2 * center_x - x

        if 0 <= mirror_x < w and visible_mask[y, mirror_x] > 0:
            filled[y, x] = image_np[y, mirror_x]

    return filled


def feather_blend(original, filled, missing_mask):
    mask = missing_mask.astype(np.float32) / 255.0

    # 경계 부드럽게
    mask = cv2.GaussianBlur(mask, (31, 31), 0)
    mask = np.clip(mask, 0, 1)

    mask_3 = np.dstack([mask, mask, mask])

    blended = original.astype(np.float32) * (1 - mask_3) + filled.astype(np.float32) * mask_3

    return np.clip(blended, 0, 255).astype(np.uint8)


def refine_with_diffusion(image_pil, mask_pil):
    dtype = torch.float16 if DEVICE == "cuda" else torch.float32

    pipe = StableDiffusionInpaintPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=dtype
    ).to(DEVICE)

    if DEVICE == "cuda":
        pipe.enable_attention_slicing()
        pipe.enable_vae_slicing()

    image_in = image_pil.resize((512, 512))
    mask_in = mask_pil.resize((512, 512))

    prompt = """
    black and white horizontal striped short sleeve t-shirt,
    same stripe pattern, same fabric texture,
    realistic product photo, white background,
    clean garment edge
    """

    negative_prompt = """
    person, body, arm, hand, face, neck, mannequin,
    logo, text, watermark,
    distorted, warped, blurry, deformed,
    extra sleeve, bad clothing shape
    """

    result = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        image=image_in,
        mask_image=mask_in,
        strength=0.35,
        guidance_scale=6.5,
        num_inference_steps=30,
        generator=torch.Generator(device=DEVICE).manual_seed(1234),
    ).images[0]

    return result.resize(image_pil.size)


# =========================
# main
# =========================

image = Image.open(IMAGE_PATH).convert("RGB")
image_np = np.array(image)

visible_mask = segment_cloth(image_np)
save_img(visible_mask, f"{OUTPUT_DIR}/01_visible_mask.png")

full_mask = symmetry_full_mask(visible_mask)
save_img(full_mask, f"{OUTPUT_DIR}/02_symmetry_full_mask.png")

missing_mask = make_missing_mask(visible_mask, full_mask)
save_img(missing_mask, f"{OUTPUT_DIR}/03_missing_mask.png")

mirror_np = mirror_fill(image_np, visible_mask, missing_mask)
save_img(mirror_np, f"{OUTPUT_DIR}/04_mirror_fill.png")

blended_np = feather_blend(image_np, mirror_np, missing_mask)
save_img(blended_np, f"{OUTPUT_DIR}/05_feather_blend.png")

diffusion_mask = Image.fromarray(missing_mask).convert("L")
diffusion_mask = diffusion_mask.filter(ImageFilter.GaussianBlur(radius=3))
diffusion_mask.save(f"{OUTPUT_DIR}/06_diffusion_mask.png")

result = refine_with_diffusion(
    Image.fromarray(blended_np),
    diffusion_mask
)

result.save(f"{OUTPUT_DIR}/07_final_result.png")

print("완료")
print(f"결과 폴더: {OUTPUT_DIR}")