#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
换衣 + 换背景 功能模块 v2
使用 MediaPipe Tasks API（新版 mediapipe 0.10.x）
"""

import os
import numpy as np
import cv2
import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision

APP_DIR = os.path.dirname(os.path.abspath(__file__))

# ─── 模型路径 ───
_MODEL_SELFIE = os.path.join(APP_DIR, "selfie_segmenter.tflite")
_MODEL_POSE = os.path.join(APP_DIR, "pose_landmarker_lite.task")

# ═══════════════════════════════════════════════
# 背景替换 - BackgroundChanger
# ═══════════════════════════════════════════════

class BackgroundChanger:
    """
    人像分割 + 背景替换
    使用 MediaPipe ImageSegmenter（DeepLabV3 / SelfieSegmenter）
    CPU 实时运行，模型仅 244KB
    """
    _shared_segmenter = None

    def __init__(self, bg_image=None):
        self.bg_image = bg_image
        self._blur_ksize = 15
        self._threshold = 0.5
        self._init_model()

    @classmethod
    def _init_model(cls):
        if cls._shared_segmenter is not None:
            return
        if not os.path.exists(_MODEL_SELFIE):
            print(f"⚠️ 自拍分割模型未找到: {_MODEL_SELFIE}")
            cls._shared_segmenter = False  # 标记为不可用
            return
        try:
            options = vision.ImageSegmenterOptions(
                base_options=python.BaseOptions(model_asset_path=_MODEL_SELFIE),
                running_mode=vision.RunningMode.IMAGE,
                output_confidence_masks=True,
            )
            cls._shared_segmenter = vision.ImageSegmenter.create_from_options(options)
        except Exception as e:
            print(f"⚠️ 背景分割模型加载失败: {e}")
            cls._shared_segmenter = False

    @property
    def ready(self):
        return (self._shared_segmenter is not None and 
                self._shared_segmenter is not False and 
                self.bg_image is not None)

    def set_background(self, bg_image):
        if bg_image is not None:
            self.bg_image = bg_image.copy()

    def get_mask(self, image):
        """获取人物分割 mask（单通道 float32, 0~1）"""
        if self._shared_segmenter is None or self._shared_segmenter is False:
            return None
        try:
            rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            mp_img = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb)
            result = self._shared_segmenter.segment(mp_img)
            if result.confidence_masks and len(result.confidence_masks) > 0:
                # 第0个通道通常是人的置信度（背景可能是第1个）
                mask = result.confidence_masks[0].numpy_view().astype(np.float32)
                return mask
            return None
        except Exception as e:
            print(f"分割出错: {e}")
            return None

    def change(self, image):
        """替换背景"""
        if self.bg_image is None:
            return image

        h, w = image.shape[:2]

        # 如果图像太小，放大到 256+ 再做分割，缩回原尺寸
        need_upscale = w < 200 or h < 200
        if need_upscale:
            sw = max(320, w * 2)
            sh = max(240, h * 2)
            proc_img = cv2.resize(image, (sw, sh))
        else:
            proc_img = image

        mask = self.get_mask(proc_img)
        if mask is None:
            return image

        # 缩回原尺寸
        if need_upscale:
            mask = cv2.resize(mask, (w, h))

        bg = cv2.resize(self.bg_image, (w, h))

        # 平滑 mask 边缘
        mask = cv2.GaussianBlur(mask, (self._blur_ksize, self._blur_ksize), 0)
        mask_3c = np.stack([mask, mask, mask], axis=2)

        result = (image.astype(np.float32) * mask_3c +
                  bg.astype(np.float32) * (1.0 - mask_3c))
        return result.clip(0, 255).astype(np.uint8)

    def change_with_bg(self, image, bg_image):
        if bg_image is None:
            return image
        self.set_background(bg_image)
        return self.change(image)


# ═══════════════════════════════════════════════
# 换衣 - ClothesChanger
# ═══════════════════════════════════════════════

class ClothesChanger:
    """
    换衣：源衣服 ➔ 姿势匹配 ➔ 透视变形 ➔ 颜色融合

    使用 MediaPipe PoseLandmarker 检测身体关键点，
    将源图中的衣服区域贴合到目标人物身上
    """
    _shared_pose = None

    # 上半身关键点索引
    LANDMARKS = {
        'left_shoulder': 11,
        'right_shoulder': 12,
        'left_hip': 23,
        'right_hip': 24,
        'left_ear': 7,
        'right_ear': 8,
        'nose': 0,
        'left_elbow': 13,
        'right_elbow': 14,
    }

    def __init__(self, source_clothes_img=None):
        self.source_img = source_clothes_img
        self._extend_ratio = 1.2
        self._feather_ksize = 15
        self._hist_match_strength = 0.6
        self._init_model()

    @classmethod
    def _init_model(cls):
        if cls._shared_pose is not None:
            return
        if not os.path.exists(_MODEL_POSE):
            print(f"⚠️ 姿态模型未找到: {_MODEL_POSE}")
            cls._shared_pose = False
            return
        try:
            options = vision.PoseLandmarkerOptions(
                base_options=python.BaseOptions(model_asset_path=_MODEL_POSE),
                running_mode=vision.RunningMode.IMAGE,
                min_pose_detection_confidence=0.3,
                min_pose_presence_confidence=0.3,
                min_tracking_confidence=0.3,
            )
            cls._shared_pose = vision.PoseLandmarker.create_from_options(options)
        except Exception as e:
            print(f"⚠️ 姿态模型加载失败: {e}")
            cls._shared_pose = False

    @property
    def ready(self):
        return (self._shared_pose is not None and 
                self._shared_pose is not False)

    def set_source_clothes(self, img):
        if img is not None:
            self.source_img = img.copy()

    def _get_landmarks(self, image):
        """获取姿态关键点数组"""
        if self._shared_pose is None or self._shared_pose is False:
            return None
        try:
            rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            mp_img = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb)
            result = self._shared_pose.detect(mp_img)
            if result.pose_landmarks and len(result.pose_landmarks) > 0:
                return result.pose_landmarks[0]
            return None
        except:
            return None

    def _get_clothes_quad(self, landmarks, img_shape):
        """
        从关键点计算衣服四边形
        返回: [(x1,y1), (x2,y2), (x3,y3), (x4,y4)]
        顺序: 左上(左肩上), 右上(右肩上), 右下(右臀), 左下(左臀)
        """
        h, w = img_shape[:2]

        def pt(idx):
            lm = landmarks[idx]
            return np.array([lm.x * w, lm.y * h], dtype=np.float32)

        ls = pt(self.LANDMARKS['left_shoulder'])
        rs = pt(self.LANDMARKS['right_shoulder'])
        lh = pt(self.LANDMARKS['left_hip'])
        rh = pt(self.LANDMARKS['right_hip'])

        # 肩膀宽度
        shoulder_width = np.linalg.norm(rs - ls)
        extend = shoulder_width * (self._extend_ratio - 1.0) / 2.0

        # 垂直方向（肩膀连线的垂直方向）
        shoulder_vec = rs - ls
        perp = np.array([-shoulder_vec[1], shoulder_vec[0]], dtype=np.float32)
        perp_len = np.linalg.norm(perp)
        if perp_len < 1:
            perp = np.array([0.0, -1.0], dtype=np.float32)
        else:
            perp /= perp_len

        # 向上延伸（覆盖脖子区域）
        ls_up = ls - perp * extend * 0.3
        rs_up = rs - perp * extend * 0.3

        # 向外延伸（覆盖肩膀外侧）
        ls_ext = ls_up + perp * extend
        rs_ext = rs_up - perp * extend

        # 臀部位置（也可向下延伸）
        lh_ext = lh
        rh_ext = rh

        quad = np.array([ls_ext, rs_ext, rh_ext, lh_ext], dtype=np.float32)
        quad[:, 0] = np.clip(quad[:, 0], 0, w - 1)
        quad[:, 1] = np.clip(quad[:, 1], 0, h - 1)
        return quad

    def _color_hist_match(self, warped, tgt_roi, mask):
        """颜色直方图匹配"""
        h, w = warped.shape[:2]
        matched = warped.copy().astype(np.float32)
        tgt_float = tgt_roi.astype(np.float32)

        for c in range(3):
            src_ch = warped[:, :, c].astype(np.float32)
            tgt_ch = tgt_roi[:, :, c].astype(np.float32)

            if mask is not None:
                pixels_src = src_ch[mask > 0.1]
                pixels_tgt = tgt_ch[mask > 0.1]
            else:
                pixels_src = src_ch.flatten()
                pixels_tgt = tgt_ch.flatten()

            if len(pixels_src) < 20 or len(pixels_tgt) < 20:
                continue

            # 直方图匹配：源 -> 目标颜色分布
            src_sorted = np.sort(pixels_src)
            tgt_sorted = np.sort(pixels_tgt)

            indices = np.searchsorted(src_sorted, src_ch)
            indices = np.clip(indices, 0, len(tgt_sorted) - 1)
            matched_ch = tgt_sorted[indices]

            # 混合原始与匹配结果
            blended = (matched_ch * self._hist_match_strength + 
                      src_ch * (1.0 - self._hist_match_strength))
            matched[:, :, c] = blended

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

    def change(self, target_img, face_bbox=None):
        """
        将源衣服换到目标人物身上
        target_img: 目标帧 (BGR)
        face_bbox: (x1,y1,x2,y2) 人脸区域（保留接口兼容，实际不需要）
        returns: 换衣后图像
        """
        if self.source_img is None or not self.ready:
            return target_img

        # 如果图像太小，临时放大再做姿态检测（pose 模型最小需要约200px宽）
        h, w = target_img.shape[:2]
        scale = 1.0
        if w < 200 or h < 200:
            scale = max(400 / w, 300 / h)
            scaled_target = cv2.resize(target_img, None, fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR)
        else:
            scaled_target = target_img

        # 检测目标姿势（在放大的图上）
        tgt_lm = self._get_landmarks(scaled_target)
        if tgt_lm is None:
            return target_img

        # 将关键点坐标缩放回原图尺寸
        if scale > 1.0:
            for lm in tgt_lm:
                lm.x /= scale
                lm.y /= scale

        # 检测源姿势（源图如果太小也放大）
        src_h, src_w = self.source_img.shape[:2]
        src_scale = 1.0
        if src_w < 200 or src_h < 200:
            src_scale = max(400 / src_w, 300 / src_h)
            scaled_src = cv2.resize(self.source_img, None, fx=src_scale, fy=src_scale, interpolation=cv2.INTER_LINEAR)
        else:
            scaled_src = self.source_img

        src_lm = self._get_landmarks(scaled_src)
        if src_lm is None:
            return target_img

        if src_scale > 1.0:
            for lm in src_lm:
                lm.x /= src_scale
                lm.y /= src_scale

        try:
            src_quad = self._get_clothes_quad(src_lm, self.source_img.shape)
            tgt_quad = self._get_clothes_quad(tgt_lm, target_img.shape)

            tgt_h, tgt_w = target_img.shape[:2]
            src_h, src_w = self.source_img.shape[:2]

            # 计算目标区域边界
            x_min = max(0, int(np.min(tgt_quad[:, 0])) - 10)
            y_min = max(0, int(np.min(tgt_quad[:, 1])) - 10)
            x_max = min(tgt_w, int(np.max(tgt_quad[:, 0])) + 10)
            y_max = min(tgt_h, int(np.max(tgt_quad[:, 1])) + 10)

            if x_max - x_min < 10 or y_max - y_min < 10:
                return target_img

            # 透视变换
            tgt_local = tgt_quad.copy()
            tgt_local[:, 0] -= x_min
            tgt_local[:, 1] -= y_min

            matrix = cv2.getPerspectiveTransform(src_quad, tgt_local)
            local_w = x_max - x_min
            local_h = y_max - y_min

            warped = cv2.warpPerspective(self.source_img, matrix, (local_w, local_h))

            # 创建羽化遮罩
            mask = np.zeros((local_h, local_w), dtype=np.uint8)
            cv2.fillConvexPoly(mask, tgt_local.astype(np.int32), 255)
            mask = cv2.GaussianBlur(mask, (self._feather_ksize, self._feather_ksize), 0)
            mask_f = mask.astype(np.float32) / 255.0

            # 颜色匹配
            tgt_roi = target_img[y_min:y_max, x_min:x_max]
            if tgt_roi.shape[:2] == warped.shape[:2]:
                warped = self._color_hist_match(warped, tgt_roi, mask_f)

            # 合成
            mask_3c = np.stack([mask_f, mask_f, mask_f], axis=2)
            roi_result = (warped.astype(np.float32) * mask_3c +
                         tgt_roi.astype(np.float32) * (1.0 - mask_3c))
            roi_result = roi_result.clip(0, 255).astype(np.uint8)

            result = target_img.copy()
            result[y_min:y_max, x_min:x_max] = roi_result
            return result

        except Exception as e:
            print(f"换衣出错: {e}")
            return target_img


# ═══════════════════════════════════════════════
# 颜色迁移版换衣（不依赖姿态检测，更稳但只改颜色）
# ═══════════════════════════════════════════════

class SimpleClothesChanger:
    """
    简易换衣：用比例估算衣服区域，做 Lab 颜色迁移
    不依赖姿态检测，适合 debug/低配
    效果：改变衣服色调接近源图，不做图案替换
    """
    def __init__(self, source_clothes_img=None):
        self.source_img = source_clothes_img
        self._strength = 0.7

    def set_source_clothes(self, img):
        if img is not None:
            self.source_img = img.copy()

    @property
    def ready(self):
        return self.source_img is not None

    def _get_dominant_colors(self, img):
        """K-Means 提取主色调"""
        samples = img.reshape(-1, 3).astype(np.float32)
        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
        _, labels, centers = cv2.kmeans(samples, 3, None, criteria, 5, cv2.KMEANS_RANDOM_CENTERS)
        counts = np.bincount(labels.flatten())
        return centers[np.argmax(counts)]

    def change(self, target_img, face_bbox=None):
        """
        颜色迁移换衣
        face_bbox: (x1,y1,x2,y2) 用于估计衣服位置
        """
        if self.source_img is None:
            return target_img

        h, w = target_img.shape[:2]

        # 估算衣服区域
        if face_bbox is not None and len(face_bbox) == 4:
            x1, y1, x2, y2 = face_bbox
            face_bottom = y2
            fw = x2 - x1
            cloth_y1 = min(h - 1, face_bottom + int(fw * 0.2))
            cloth_y2 = h
            cloth_x1 = max(0, x1 - int(fw * 0.8))
            cloth_x2 = min(w, x2 + int(fw * 0.8))
        else:
            cloth_y1 = h // 3
            cloth_y2 = h
            cloth_x1 = 0
            cloth_x2 = w

        if cloth_y1 >= cloth_y2 or cloth_x1 >= cloth_x2:
            return target_img

        # 获取源主色
        src_color = self._get_dominant_colors(self.source_img)
        if src_color is None:
            return target_img

        # Lab 颜色迁移
        result = target_img.copy()
        roi = result[cloth_y1:cloth_y2, cloth_x1:cloth_x2]
        if roi.size == 0:
            return target_img

        lab = cv2.cvtColor(roi, cv2.COLOR_BGR2LAB).astype(np.float32)

        l_mean = np.mean(lab[:, :, 0])
        a_mean = np.mean(lab[:, :, 1])
        b_mean = np.mean(lab[:, :, 2])

        src_lab = cv2.cvtColor(np.uint8([[src_color]]), cv2.COLOR_BGR2LAB)[0, 0].astype(np.float32)

        lab[:, :, 1] += (src_lab[1] - a_mean) * self._strength
        lab[:, :, 2] += (src_lab[2] - b_mean) * self._strength

        lab[:, :, 0] = np.clip(lab[:, :, 0], 0, 255)
        lab[:, :, 1] = np.clip(lab[:, :, 1], 0, 255)
        lab[:, :, 2] = np.clip(lab[:, :, 2], 0, 255)

        result[cloth_y1:cloth_y2, cloth_x1:cloth_x2] = cv2.cvtColor(lab.astype(np.uint8), cv2.COLOR_LAB2BGR)
        return result
