#!/usr/bin/env python3
"""
Tatiana's embroidered patch composite — v2.
Places FF logo below the flag patch on the astronaut's lit upper left arm.
Focus: making it look physically real, not digital.
"""

from PIL import Image, ImageDraw, ImageFilter, ImageEnhance
import numpy as np
import random

random.seed(42)
np.random.seed(42)

# Load images
base = Image.open("v1-03-mid-full-body.png").convert("RGBA")
logo_raw = Image.open("ff-logo-mark.png").convert("RGBA")

print(f"Base: {base.size}")
print(f"Logo raw: {logo_raw.size}")

base_arr = np.array(base)

# ============================================================
# 1. DETERMINE PATCH PLACEMENT
#    Below the flag patch on astronaut's left upper arm.
#    Flag center ~(675, 645) in full image.
#    We'll place patch below: around x=645, y=680
# ============================================================
PASTE_X = 642
PASTE_Y = 682
PATCH_W = 72
PATCH_H = 56

# ============================================================
# 2. SAMPLE LOCAL ARM COLOR for realistic tinting
# ============================================================
arm_region = base_arr[PASTE_Y:PASTE_Y+PATCH_H, PASTE_X:PASTE_X+PATCH_W, :3]
avg_arm = arm_region.mean(axis=(0,1))
print(f"Arm avg color: {avg_arm.round(1)}")

# Suit fabric is mostly white/off-white, slightly warm
# Patch should be off-white matching suit tone
patch_r = min(255, int(avg_arm[0] * 0.97 + 6))
patch_g = min(255, int(avg_arm[1] * 0.97 + 4))
patch_b = min(255, int(avg_arm[2] * 0.97 + 2))
patch_base_color = (patch_r, patch_g, patch_b, 255)
print(f"Patch base color: {patch_base_color}")

# ============================================================
# 3. BUILD PATCH WITH FABRIC TEXTURE
# ============================================================
patch = Image.new("RGBA", (PATCH_W, PATCH_H), patch_base_color)
patch_arr = np.array(patch, dtype=np.float32)

# Add subtle fabric weave texture (fine noise)
noise = np.random.normal(0, 4, (PATCH_H, PATCH_W, 3)).astype(np.float32)
# Also add a coarser texture grain
coarse_noise = np.random.normal(0, 2, (PATCH_H // 4, PATCH_W // 4, 3))
coarse_noise = np.kron(coarse_noise, np.ones((4, 4, 1)))[:PATCH_H, :PATCH_W, :]

patch_arr[:,:,:3] = np.clip(patch_arr[:,:,:3] + noise + coarse_noise, 0, 255)
patch_arr[:,:,3] = 255

# Slight darkening at edges (patch has natural shadow from stitching thickness)
for edge_d in range(4):
    alpha_factor = 0.96 - edge_d * 0.015
    patch_arr[edge_d, :, :3] *= alpha_factor
    patch_arr[PATCH_H-1-edge_d, :, :3] *= alpha_factor
    patch_arr[:, edge_d, :3] *= alpha_factor
    patch_arr[:, PATCH_W-1-edge_d, :3] *= alpha_factor

patch = Image.fromarray(np.clip(patch_arr, 0, 255).astype(np.uint8), "RGBA")

# ============================================================
# 4. PROCESS LOGO — convert to embroidery look
# ============================================================
logo_area_w = PATCH_W - 16
logo_area_h = PATCH_H - 14

# Maintain aspect ratio
logo_copy = logo_raw.copy()
logo_copy.thumbnail((logo_area_w, logo_area_h), Image.LANCZOS)
lw, lh = logo_copy.size
print(f"Logo sized: {lw}x{lh}")

logo_data = np.array(logo_copy).astype(np.float32)
r, g, b, a = logo_data[:,:,0], logo_data[:,:,1], logo_data[:,:,2], logo_data[:,:,3]

# Luminance of each pixel
lum = (r * 0.299 + g * 0.587 + b * 0.114)

# Dark pixels = logo symbol
logo_mask = (lum < 140) & (a > 30)

# Embroidery thread color — very dark, with slight warm undertone
# Threads have subtle sheen variation
thread_base = np.array([18, 15, 10], dtype=np.float32)

# Build embroidery layer
emb_arr = np.zeros((lh, lw, 4), dtype=np.float32)
emb_arr[logo_mask, :3] = thread_base
# Add thread sheen — slight horizontal variation (simulates thread direction)
sheen = np.sin(np.arange(lh)[:, None] * 1.8) * 6
emb_arr[:, :, 0] += logo_mask * sheen
emb_arr[:, :, 1] += logo_mask * sheen * 0.8
emb_arr[:, :, 2] += logo_mask * sheen * 0.6
emb_arr[logo_mask, 3] = 225  # slightly translucent for thread look
emb_arr[~logo_mask, 3] = 0

emb_img = Image.fromarray(np.clip(emb_arr, 0, 255).astype(np.uint8), "RGBA")

# Slight blur — thread edges aren't sharp like print
emb_img = emb_img.filter(ImageFilter.GaussianBlur(0.4))

# Center on patch
lx = (PATCH_W - lw) // 2
ly = (PATCH_H - lh) // 2
patch.paste(emb_img, (lx, ly), emb_img)

# ============================================================
# 5. STITCH BORDER — realistic looking
# ============================================================
patch_arr2 = np.array(patch, dtype=np.float32)
draw = ImageDraw.Draw(patch)

# Thread color for border (matches embroidery color, slightly lighter)
stitch_color = (25, 22, 16, 200)
stitch_inner = (40, 36, 28, 160)

# Outer edge: a row of stitches (dashes at 45-degree for border feel)
stitch_len = 3
stitch_gap = 2
border_margin = 3

# Top & bottom
for x in range(border_margin, PATCH_W - border_margin, stitch_len + stitch_gap):
    xe = min(x + stitch_len, PATCH_W - border_margin)
    draw.line([(x, border_margin), (xe, border_margin)], fill=stitch_color, width=1)
    draw.line([(x, PATCH_H - border_margin - 1), (xe, PATCH_H - border_margin - 1)], fill=stitch_color, width=1)

# Left & right  
for y in range(border_margin, PATCH_H - border_margin, stitch_len + stitch_gap):
    ye = min(y + stitch_len, PATCH_H - border_margin)
    draw.line([(border_margin, y), (border_margin, ye)], fill=stitch_color, width=1)
    draw.line([(PATCH_W - border_margin - 1, y), (PATCH_W - border_margin - 1, ye)], fill=stitch_color, width=1)

# Inner border line (chain stitch inside edge)
for x in range(border_margin + 3, PATCH_W - border_margin - 3, stitch_len + stitch_gap + 1):
    xe = min(x + stitch_len, PATCH_W - border_margin - 3)
    draw.line([(x, border_margin + 3), (xe, border_margin + 3)], fill=stitch_inner, width=1)
    draw.line([(x, PATCH_H - border_margin - 4), (xe, PATCH_H - border_margin - 4)], fill=stitch_inner, width=1)
for y in range(border_margin + 3, PATCH_H - border_margin - 3, stitch_len + stitch_gap + 1):
    ye = min(y + stitch_len, PATCH_H - border_margin - 3)
    draw.line([(border_margin + 3, y), (border_margin + 3, ye)], fill=stitch_inner, width=1)
    draw.line([(PATCH_W - border_margin - 4, y), (PATCH_W - border_margin - 4, ye)], fill=stitch_inner, width=1)

# ============================================================
# 6. APPLY LIGHTING from suit surface
# ============================================================
# Sample local lighting gradient from base image
base_rgb_arr = np.array(base.convert("RGB"), dtype=np.float32)
arm_patch = base_rgb_arr[PASTE_Y:PASTE_Y+PATCH_H, PASTE_X:PASTE_X+PATCH_W]

# Create luminance map of the suit surface to replicate its lighting
arm_lum = (arm_patch[:,:,0] * 0.299 + arm_patch[:,:,1] * 0.587 + arm_patch[:,:,2] * 0.114)
arm_lum_norm = arm_lum / 255.0  # 0.0 to 1.0

# Light adjustment: brighten where suit is bright, darken where suit is dark
patch_final_arr = np.array(patch, dtype=np.float32)

# Mix: 60% flat patch + 40% lit-adjusted
light_factor = arm_lum_norm[:, :, np.newaxis]
light_adjust = 0.7 + 0.5 * light_factor  # range ~0.7–1.2
patch_final_arr[:,:,:3] = np.clip(patch_final_arr[:,:,:3] * light_adjust, 0, 255)
patch_final = Image.fromarray(patch_final_arr.astype(np.uint8), "RGBA")

# Slight overall blur to integrate (removes "sharp cutout" look)
patch_final = patch_final.filter(ImageFilter.GaussianBlur(0.3))

# ============================================================
# 7. COMPOSITE: SHADOW + PATCH onto BASE
# ============================================================
result = base.copy()

# Shadow layer (subtle, directional - light comes from upper right in this image)
shadow_layer = Image.new("RGBA", base.size, (0, 0, 0, 0))
shadow_patch = Image.new("RGBA", (PATCH_W + 6, PATCH_H + 6), (0, 0, 0, 0))
# Build gradient shadow
for i in range(6):
    alpha = int(35 * (1 - i / 6.0))
    shadow_sub = Image.new("RGBA", (PATCH_W + 6 - i*2, PATCH_H + 6 - i*2), (5, 5, 5, alpha))
    shadow_patch.paste(shadow_sub, (i, i), shadow_sub)
shadow_patch = shadow_patch.filter(ImageFilter.GaussianBlur(1.5))
shadow_layer.paste(shadow_patch, (PASTE_X - 1, PASTE_Y + 1), shadow_patch)
result = Image.alpha_composite(result, shadow_layer)

# Patch layer
patch_layer = Image.new("RGBA", base.size, (0, 0, 0, 0))
patch_layer.paste(patch_final, (PASTE_X, PASTE_Y), patch_final)
result = Image.alpha_composite(result, patch_layer)

# ============================================================
# 8. SAVE
# ============================================================
result_rgb = result.convert("RGB")
result_rgb.save("tatiana-patch-v2.jpg", quality=96)
print(f"Saved: tatiana-patch-v2.jpg")

# Debug crop — wider view to show context
debug = result_rgb.crop((550, 580, 800, 800))
debug_big = debug.resize((debug.width * 2, debug.height * 2), Image.NEAREST)
debug_big.save("tatiana-patch-v2-debug.jpg", quality=95)
print(f"Saved debug: tatiana-patch-v2-debug.jpg ({debug.size} → 2x)")