# controlnet_cloth_restore_symmetry.py

import os
import cv2
import torch
import numpy as np

from PIL import Image, ImageFilter
from transformers import AutoImageProcessor, SegformerForSemanticSegmentation
from diffusers import (
    ControlNetModel,
    StableDiffusionControlNetInpaintPipeline,
    UniPCMultistepScheduler,
)


IMAGE_PATH = "person.jpg"
OUTPUT_DIR = "output_controlnet_symmetry"

SEG_MODEL = "mattmdjaga/segformer_b2_clothes"
BASE_MODEL = "runwayml/stable-diffusion-inpainting"
CONTROLNET_MODEL = "lllyasviel/control_v11p_sd15_canny"

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

os.makedirs(OUTPUT_DIR, exist_ok=True)


# =========================================================
# 유틸
# =========================================================

def save_mask(path, mask):
    Image.fromarray((mask.astype(np.uint8) * 255)).save(path)


def save_image(path, img):
    if isinstance(img, np.ndarray):
        Image.fromarray(img).save(path)
    else:
        img.save(path)


def clean_mask(mask, k=7):
    mask = mask.astype(np.uint8) * 255
    kernel = np.ones((k, k), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
    return mask > 127


def dilate(mask, k=15, iterations=1):
    kernel = np.ones((k, k), np.uint8)
    return cv2.dilate(mask.astype(np.uint8), kernel, iterations=iterations) > 0


def erode(mask, k=5, iterations=1):
    kernel = np.ones((k, k), np.uint8)
    return cv2.erode(mask.astype(np.uint8), kernel, iterations=iterations) > 0


def blur_mask(mask, radius=3):
    img = Image.fromarray((mask.astype(np.uint8) * 255))
    return img.filter(ImageFilter.GaussianBlur(radius))


def get_bbox(mask, padding=70):
    ys, xs = np.where(mask)

    if len(xs) == 0:
        raise ValueError("상의 mask가 비어 있습니다.")

    h, w = mask.shape

    x1 = max(0, xs.min() - padding)
    y1 = max(0, ys.min() - padding)
    x2 = min(w, xs.max() + padding)
    y2 = min(h, ys.max() + padding)

    return x1, y1, x2, y2


def crop_np(arr, bbox):
    x1, y1, x2, y2 = bbox
    return arr[y1:y2, x1:x2]


def resize_to_multiple_of_8(img, mask, control):
    w, h = img.size

    new_w = max(8, (w // 8) * 8)
    new_h = max(8, (h // 8) * 8)

    img = img.resize((new_w, new_h), Image.LANCZOS)
    mask = mask.resize((new_w, new_h), Image.NEAREST)
    control = control.resize((new_w, new_h), Image.LANCZOS)

    return img, mask, control


# =========================================================
# Symmetry Fitting 핵심
# =========================================================

def estimate_center_x(mask):
    """
    상의 mask의 y별 중심을 구한 뒤 median으로 중심축 추정
    """
    h, w = mask.shape
    centers = []

    for y in range(h):
        xs = np.where(mask[y])[0]
        if len(xs) > 5:
            centers.append((xs.min() + xs.max()) / 2)

    if len(centers) == 0:
        return w // 2

    return int(np.median(centers))


def mirror_index_map(width, center_x):
    """
    중심축 기준 mirror x 좌표 계산
    """
    xs = np.arange(width)
    mirror_xs = 2 * center_x - xs
    mirror_xs = np.clip(mirror_xs, 0, width - 1).astype(np.int32)
    return mirror_xs


def symmetry_fit_image_and_mask(
    image_np,
    cloth_mask,
    remove_mask,
    occlusion_mask,
):
    """
    팔/얼굴/하의로 가려진 영역을 좌우 대칭 기반으로 먼저 보정
    """

    h, w = cloth_mask.shape
    center_x = estimate_center_x(cloth_mask)

    mirror_xs = mirror_index_map(w, center_x)

    mirrored_image = image_np[:, mirror_xs, :]
    mirrored_cloth = cloth_mask[:, mirror_xs]

    # 좌우 중 어느 쪽이 더 온전한지 판단
    left_area = cloth_mask[:, :center_x].sum()
    right_area = cloth_mask[:, center_x:].sum()

    left_removed = remove_mask[:, :center_x].sum()
    right_removed = remove_mask[:, center_x:].sum()

    left_score = left_area - left_removed * 2
    right_score = right_area - right_removed * 2

    # 오른쪽이 더 온전하면 오른쪽을 왼쪽으로 복사
    # 왼쪽이 더 온전하면 왼쪽을 오른쪽으로 복사
    use_mirror = np.zeros_like(cloth_mask, dtype=bool)

    if left_score >= right_score:
        # 왼쪽 기준 → 오른쪽 복원
        use_mirror[:, center_x:] = True
    else:
        # 오른쪽 기준 → 왼쪽 복원
        use_mirror[:, :center_x] = True

    # 실제 복사할 영역:
    # 1. occlusion 영역
    # 2. 사람 제거 영역
    # 3. 대칭된 cloth mask 안쪽
    target = (occlusion_mask | remove_mask) & mirrored_cloth & use_mirror

    fitted = image_np.copy()
    fitted[target] = mirrored_image[target]

    fitted_mask = cloth_mask | (mirrored_cloth & target)

    return fitted, fitted_mask, center_x, target


def draw_center_axis(image_np, center_x):
    out = image_np.copy()
    cv2.line(out, (center_x, 0), (center_x, out.shape[0]), (255, 0, 0), 2)
    return out


# =========================================================
# ControlNet Canny
# =========================================================

def make_canny_control_image(image_np, cloth_mask, remove_mask=None):
    gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
    edges = cv2.Canny(gray, 70, 150)

    cloth_area = dilate(cloth_mask, k=20)
    edges = edges * cloth_area.astype(np.uint8)

    if remove_mask is not None:
        remove_area = dilate(remove_mask, k=30)
        edges[remove_area] = 0

    edges_3ch = np.stack([edges, edges, edges], axis=2)
    return Image.fromarray(edges_3ch)


def remove_human_parts_from_image(image_np, remove_mask):
    out = image_np.copy()
    remove_area = dilate(remove_mask, k=8)
    out[remove_area] = [255, 255, 255]
    return out


# =========================================================
# 1. 세그멘테이션
# =========================================================

print("1. segmentation")

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

H, W = image_np.shape[:2]

processor = AutoImageProcessor.from_pretrained(SEG_MODEL)
model = SegformerForSemanticSegmentation.from_pretrained(SEG_MODEL).to(DEVICE)
model.eval()

inputs = processor(images=image, return_tensors="pt").to(DEVICE)

with torch.no_grad():
    outputs = model(**inputs)

logits = torch.nn.functional.interpolate(
    outputs.logits,
    size=(H, W),
    mode="bilinear",
    align_corners=False,
)

seg = logits.argmax(dim=1)[0].cpu().numpy()


LABEL = {
    "background": 0,
    "hat": 1,
    "hair": 2,
    "sunglasses": 3,
    "upper": 4,
    "skirt": 5,
    "pants": 6,
    "dress": 7,
    "belt": 8,
    "left_shoe": 9,
    "right_shoe": 10,
    "face": 11,
    "left_leg": 12,
    "right_leg": 13,
    "left_arm": 14,
    "right_arm": 15,
    "bag": 16,
    "scarf": 17,
}

upper_mask = seg == LABEL["upper"]
face_mask = seg == LABEL["face"]
hair_mask = seg == LABEL["hair"]
arm_mask = (seg == LABEL["left_arm"]) | (seg == LABEL["right_arm"])
lower_mask = (seg == LABEL["pants"]) | (seg == LABEL["skirt"])

upper_mask = clean_mask(upper_mask, k=7)

save_mask(f"{OUTPUT_DIR}/01_upper_mask.png", upper_mask)
save_mask(f"{OUTPUT_DIR}/01_arm_mask.png", arm_mask)
save_mask(f"{OUTPUT_DIR}/01_face_mask.png", face_mask)
save_mask(f"{OUTPUT_DIR}/01_lower_mask.png", lower_mask)


# =========================================================
# 2. Crop
# =========================================================

print("2. crop")

bbox = get_bbox(upper_mask, padding=70)

crop_img = crop_np(image_np, bbox)
crop_upper = crop_np(upper_mask, bbox)
crop_arm = crop_np(arm_mask, bbox)
crop_face = crop_np(face_mask, bbox)
crop_hair = crop_np(hair_mask, bbox)
crop_lower = crop_np(lower_mask, bbox)

save_image(f"{OUTPUT_DIR}/02_crop_image.png", crop_img)
save_mask(f"{OUTPUT_DIR}/02_crop_upper.png", crop_upper)


# =========================================================
# 3. 상의 mask 보정
# =========================================================

print("3. refine upper mask")

upper_closed = clean_mask(crop_upper, k=25)
upper_limit = dilate(crop_upper, k=25)

refined_upper = upper_closed & upper_limit
refined_upper = clean_mask(refined_upper, k=9)

save_mask(f"{OUTPUT_DIR}/03_refined_upper.png", refined_upper)


# =========================================================
# 4. 사람 영역 제거 mask
# =========================================================

print("4. remove human parts")

remove_mask = crop_arm | crop_face | crop_hair | crop_lower
remove_mask = clean_mask(remove_mask, k=5)

save_mask(f"{OUTPUT_DIR}/04_remove_mask.png", remove_mask)

clean_crop_img = remove_human_parts_from_image(
    crop_img,
    remove_mask
)

save_image(f"{OUTPUT_DIR}/04_clean_crop_no_human.png", clean_crop_img)


# =========================================================
# 5. Occlusion mask
# =========================================================

print("5. occlusion mask")

upper_near = dilate(refined_upper, k=18)

arm_occ = dilate(crop_arm, k=14) & upper_near
face_occ = dilate(crop_face | crop_hair, k=14) & upper_near
lower_occ = dilate(crop_lower, k=16) & upper_near

hole_occ = refined_upper & (~crop_upper)

h, w = crop_upper.shape

bottom_band = np.zeros_like(crop_upper, dtype=bool)
bottom_band[int(h * 0.86):, :] = True
bottom_occ = bottom_band & dilate(crop_upper, k=22)

occlusion_mask = arm_occ | face_occ | lower_occ | hole_occ | bottom_occ

occlusion_mask = occlusion_mask & dilate(refined_upper, k=22)

# 상의와 겹치는 remove 영역도 복원 대상으로 포함
occlusion_mask = occlusion_mask | (remove_mask & dilate(refined_upper, k=24))

occlusion_mask = clean_mask(occlusion_mask, k=5)

save_mask(f"{OUTPUT_DIR}/05_occlusion_mask.png", occlusion_mask)


# =========================================================
# 6. Symmetry Fitting 적용
# =========================================================

print("6. symmetry fitting")

symmetry_prefill, symmetry_mask, center_x, symmetry_target = symmetry_fit_image_and_mask(
    image_np=clean_crop_img,
    cloth_mask=refined_upper,
    remove_mask=remove_mask,
    occlusion_mask=occlusion_mask,
)

save_image(f"{OUTPUT_DIR}/06_symmetry_prefill.png", symmetry_prefill)
save_mask(f"{OUTPUT_DIR}/06_symmetry_mask.png", symmetry_mask)
save_mask(f"{OUTPUT_DIR}/06_symmetry_target.png", symmetry_target)

axis_vis = draw_center_axis(crop_img, center_x)
save_image(f"{OUTPUT_DIR}/06_center_axis.png", axis_vis)


# =========================================================
# 7. ControlNet Canny 생성
# =========================================================

print("7. make control image")

control_image = make_canny_control_image(
    image_np=symmetry_prefill,
    cloth_mask=symmetry_mask,
    remove_mask=remove_mask,
)

control_image.save(f"{OUTPUT_DIR}/07_control_canny.png")


# =========================================================
# 8. ControlNet Inpainting
# =========================================================

print("8. controlnet inpainting")

inpaint_image = Image.fromarray(symmetry_prefill).convert("RGB")
inpaint_mask = blur_mask(occlusion_mask, radius=3)

inpaint_image, inpaint_mask, control_image = resize_to_multiple_of_8(
    inpaint_image,
    inpaint_mask,
    control_image,
)

inpaint_image.save(f"{OUTPUT_DIR}/08_inpaint_input.png")
inpaint_mask.save(f"{OUTPUT_DIR}/08_inpaint_mask.png")
control_image.save(f"{OUTPUT_DIR}/08_control_image.png")

controlnet = ControlNetModel.from_pretrained(
    CONTROLNET_MODEL,
    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
)

pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
    BASE_MODEL,
    controlnet=controlnet,
    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
    safety_checker=None,
)

pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to(DEVICE)

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

prompt = """
a single upper garment only,
repair only the masked missing fabric,
same garment,
same color,
same texture,
same pattern,
symmetrical clothing structure,
follow the clothing edge structure,
no human body,
no arms,
no hands,
no skin,
realistic fabric continuation
"""

negative_prompt = """
person,
human,
model,
body,
torso,
skin,
arm,
arms,
hand,
hands,
finger,
fingers,
face,
neck,
hair,
leg,
legs,
pants,
skirt,
new clothes,
different garment,
changed color,
changed pattern,
extra sleeve,
distorted sleeve,
distorted collar,
warped clothing,
asymmetrical clothing,
cartoon,
illustration,
blurry,
low quality
"""

iw, ih = inpaint_image.size

result = pipe(
    prompt=prompt,
    negative_prompt=negative_prompt,
    image=inpaint_image,
    mask_image=inpaint_mask,
    control_image=control_image,
    height=ih,
    width=iw,
    strength=0.30,
    guidance_scale=5.0,
    controlnet_conditioning_scale=0.95,
    num_inference_steps=32,
).images[0]

result.save(f"{OUTPUT_DIR}/09_controlnet_result.png")


# =========================================================
# 9. 최종 저장
# =========================================================

print("9. save product image")

result_np = np.array(result)

final_alpha = symmetry_mask | occlusion_mask
final_alpha = dilate(final_alpha, k=10)
final_alpha = clean_mask(final_alpha, k=5)

alpha_img = Image.fromarray((final_alpha.astype(np.uint8) * 255))
alpha_img = alpha_img.filter(ImageFilter.GaussianBlur(2))
alpha = np.array(alpha_img)

rh, rw = result_np.shape[:2]

if alpha.shape[:2] != (rh, rw):
    alpha = cv2.resize(alpha, (rw, rh), interpolation=cv2.INTER_LINEAR)

rgba = np.dstack([result_np, alpha])
Image.fromarray(rgba).save(f"{OUTPUT_DIR}/10_product_transparent.png")

white_bg = np.ones_like(result_np) * 255
alpha_f = alpha[..., None] / 255.0

white_result = (
    result_np * alpha_f +
    white_bg * (1 - alpha_f)
).astype(np.uint8)

Image.fromarray(white_result).save(f"{OUTPUT_DIR}/11_product_white_bg.png")

print("완료")
print(f"결과 폴더: {OUTPUT_DIR}")
print(f"추정 중심축 x = {center_x}")