import numpy as np
import cv2
from PIL import Image
from typing import Union
from mappings import (
    DENSE_INDEX_MAP,
    ATR_MAPPING,
    LIP_MAPPING,
    MASK_CLOTH_PARTS,
    MASK_DENSE_PARTS,
    PROTECT_BODY_PARTS,
    PROTECT_CLOTH_PARTS
)

def get_person_height(mask):
    mask = (mask > 127).astype(np.uint8)
    nonzero_pixels = np.any(mask, axis=1)
    if not np.any(nonzero_pixels):
        return 50
    ymin, ymax = np.where(nonzero_pixels)[0][[0, -1]]
    person_height = max(50, ymax - ymin)
    return person_height

def calculate_kernels(person_height):
    dilate_size = max(1, person_height // 30)
    dilate_size = dilate_size if dilate_size % 2 == 1 else dilate_size + 1
    dilate_kernel = np.ones((dilate_size, dilate_size), np.uint8)
    kernel_size = max(5, person_height // 15)
    kernel_size = kernel_size if kernel_size % 2 == 1 else kernel_size + 1
    return dilate_kernel, kernel_size

def part_mask_of(part: Union[str, list], parse: np.ndarray, mapping: dict):
    if isinstance(part, str):
        part = [part]
    mask = np.zeros_like(parse, dtype=np.uint8)
    for p in part:
        if p not in mapping:
            continue
        val = mapping[p]
        if isinstance(val, list):
            for i in val:
                mask[parse == i] = 1
        else:
            mask[parse == val] = 1
    return mask

def hull_mask(mask_area: np.ndarray):
    mask_area = (mask_area * 255).astype(np.uint8)
    _, binary = cv2.threshold(mask_area, 127, 255, cv2.THRESH_BINARY)
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    hull_mask_result = np.zeros_like(binary)
    for cnt in contours:
        hull = cv2.convexHull(cnt)
        cv2.fillPoly(hull_mask_result, [hull], 255)
    return (hull_mask_result > 127).astype(np.uint8)

def compute_strong_protect_area(densepose_mask, schp_lip_mask, schp_atr_mask, dilate_kernel):
    hands_protect_area = part_mask_of(['hands', 'feet', 'face'], densepose_mask, DENSE_INDEX_MAP)
    hands_protect_area = cv2.dilate(hands_protect_area, dilate_kernel, iterations=1)
    arms_legs_schp = (
        part_mask_of(['Left-arm', 'Right-arm', 'Left-leg', 'Right-leg'], schp_atr_mask, ATR_MAPPING) |
        part_mask_of(['Left-arm', 'Right-arm', 'Left-leg', 'Right-leg'], schp_lip_mask, LIP_MAPPING)
    )
    hands_protect_area = hands_protect_area & arms_legs_schp
    face_protect_area = part_mask_of('Face', schp_lip_mask, LIP_MAPPING)
    return (hands_protect_area | face_protect_area).astype(np.uint8)

def compute_weak_protect_area(schp_lip_mask, schp_atr_mask, strong_protect_area, part):
    body_protect_area = (
        part_mask_of(PROTECT_BODY_PARTS[part], schp_lip_mask, LIP_MAPPING) |
        part_mask_of(PROTECT_BODY_PARTS[part], schp_atr_mask, ATR_MAPPING)
    )
    hair_protect_area = (
        part_mask_of(['Hair'], schp_lip_mask, LIP_MAPPING) |
        part_mask_of(['Hair'], schp_atr_mask, ATR_MAPPING)
    )
    cloth_protect_area = part_mask_of(PROTECT_CLOTH_PARTS[part]['ATR'], schp_atr_mask, ATR_MAPPING)
    accessory_parts = ['Hat', 'Glove', 'Sunglasses', 'Bag', 'Left-shoe', 'Right-shoe', 'Scarf', 'Socks']
    accessory_protect_area = (
        part_mask_of(accessory_parts, schp_lip_mask, LIP_MAPPING) |
        part_mask_of(accessory_parts, schp_atr_mask, ATR_MAPPING)
    )
    weak_area = (body_protect_area | cloth_protect_area | hair_protect_area |
                 strong_protect_area | accessory_protect_area)
    return weak_area.astype(np.uint8)

def compute_mask_area(densepose_mask, schp_lip_mask, schp_atr_mask,
                      weak_protect_area, strong_protect_area,
                      dilate_kernel, part):
    strong_mask_area = (
        part_mask_of(MASK_CLOTH_PARTS[part], schp_lip_mask, LIP_MAPPING) |
        part_mask_of(MASK_CLOTH_PARTS[part], schp_atr_mask, ATR_MAPPING)
    )
    background_area = (
        part_mask_of(['Background'], schp_lip_mask, LIP_MAPPING) &
        part_mask_of(['Background'], schp_atr_mask, ATR_MAPPING)
    )
    mask_dense_area = part_mask_of(MASK_DENSE_PARTS[part], densepose_mask, DENSE_INDEX_MAP)
    mask_dense_area = cv2.resize(mask_dense_area.astype(np.uint8), None, fx=0.25, fy=0.25, interpolation=cv2.INTER_NEAREST)
    mask_dense_area = cv2.dilate(mask_dense_area, dilate_kernel, iterations=2)
    mask_dense_area = cv2.resize(mask_dense_area.astype(np.uint8), None, fx=4, fy=4, interpolation=cv2.INTER_NEAREST)
    target_shape = densepose_mask.shape
    def resize_mask(mask, target_shape):
        return cv2.resize(mask.astype(np.uint8), (target_shape[1], target_shape[0]), interpolation=cv2.INTER_NEAREST)
    weak_protect_area = resize_mask(weak_protect_area, target_shape)
    background_area = resize_mask(background_area, target_shape)
    mask_dense_area = resize_mask(mask_dense_area, target_shape)
    base_area = np.ones_like(densepose_mask, dtype=np.uint8)
    mask_area = (base_area & (~weak_protect_area) & (~background_area)) | mask_dense_area
    mask_area = hull_mask(mask_area)
    mask_area = (mask_area & (~weak_protect_area)).astype(np.uint8)
    return mask_area, background_area

def safe_dilate(mask, kernel, iterations, protect_mask):
    dilated_mask = mask.copy()
    for _ in range(iterations):
        temp_dilated = cv2.dilate(dilated_mask, kernel, iterations=1)
        temp_dilated[protect_mask > 0] = 0
        dilated_mask = temp_dilated
    return dilated_mask

def finalize_mask(mask_area, kernal_size, strong_mask_area, strong_protect_area, dilate_kernel):
    mask_area = (mask_area * 255).astype(np.uint8)
    mask_area = cv2.GaussianBlur(mask_area, (kernal_size, kernal_size), 0)
    mask_area = np.where(mask_area < 25, 0, 1).astype(np.uint8)
    mask_area = (mask_area | strong_mask_area) & (~strong_protect_area)
    mask_area = safe_dilate(mask_area, dilate_kernel, iterations=1, protect_mask=strong_protect_area)
    return Image.fromarray((mask_area * 255).astype(np.uint8))

def keep_largest_connected_component(mask_np):
    mask_uint8 = (mask_np > 128).astype(np.uint8)
    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask_uint8, connectivity=8)
    if num_labels <= 1:
        return mask_uint8 * 255
    largest_label = 1 + np.argmax(stats[1:, cv2.CC_STAT_AREA])
    cleaned_mask = np.zeros_like(mask_uint8)
    cleaned_mask[labels == largest_label] = 1
    return cleaned_mask * 255