#!/usr/bin/env python3
# -*- coding: UTF-8 -*-

import os
import numpy as np
import torch
from PIL import Image, ImageEnhance
import cv2
from argparse import ArgumentParser
from tqdm import tqdm

def read_image_safe(path):
    try:
        img_array = np.fromfile(path, dtype=np.uint8)
        img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
        return img
    except Exception as e:
        print(f"Error reading image {path}: {e}")
        return None

# 参数设置
parser = ArgumentParser()
parser.add_argument("--img_path_ori", type=str, required=True, help="输入原始图像的目录路径")
parser.add_argument("--img_path_depth", type=str, required=True, help="输入3d灰度图像的目录路径")
parser.add_argument("--outdir", type=str, required=True, help="输出处理后的 PLY 文件的目录")
parser.add_argument("--appendix", type=str, default='_depth', help="深度图后缀")
parser.add_argument("--suffix", type=str, default='', help="深度图前缀")
parser.add_argument("--label-type", type=str, default='Disparity', help="标注图格式：Depth 或 Disparity")
parser.add_argument("--focal-length", type=float, default=-1, help="相机焦距")
parser.add_argument("--FoV", type=float, default=60, help="相机视场角")

# === 新增调整参数 ===
parser.add_argument("--z-scale", type=float, default=1.0, help="[新增] Z轴拉伸倍数，数值越大，模型越立体/修长")
parser.add_argument("--brightness", type=float, default=1.0, help="[新增] 亮度调整倍数 (例如 1.2 为变亮20%)")
parser.add_argument("--saturation", type=float, default=1.0, help="[新增] 饱和度调整倍数 (例如 1.2 为颜色更鲜艳)")
parser.add_argument("--gamma", type=float, default=1.0, help="[新增] Gamma校正 (数值<1.0 提亮暗部细节)")
# ===================

args = parser.parse_args()

if not os.path.exists(args.outdir):
    os.makedirs(args.outdir)

if os.path.isdir(args.img_path_ori):
    file_paths_ori = [os.path.join(args.img_path_ori, f) for f in os.listdir(args.img_path_ori) if f.endswith('.png') or f.endswith('.jpg')]
    file_paths_ori.sort()
else:
    file_paths_ori = [args.img_path_ori]

if os.path.isdir(args.img_path_depth):
    depth_file_path = args.img_path_depth
else:
    depth_file_path = os.path.dirname(args.img_path_depth)

def init_image_coor(height, width):
    x_row = np.arange(0, width)
    y_col = np.arange(0, height)
    x = np.tile(x_row, (height, 1))
    y = np.tile(y_col, (width, 1)).T
    x = torch.from_numpy(x.astype(np.float32)).unsqueeze(0) - width / 2.0
    y = torch.from_numpy(y.astype(np.float32)).unsqueeze(0) - height / 2.0
    return x, y

def adjust_image_color(pil_image, brightness=1.0, saturation=1.0, gamma=1.0):
    """处理图像颜色：亮度 -> 饱和度 -> Gamma"""
    # 1. 调整亮度
    if brightness != 1.0:
        enhancer = ImageEnhance.Brightness(pil_image)
        pil_image = enhancer.enhance(brightness)
    
    # 2. 调整饱和度
    if saturation != 1.0:
        enhancer = ImageEnhance.Color(pil_image)
        pil_image = enhancer.enhance(saturation)
    
    # 转为numpy处理Gamma
    img_array = np.array(pil_image).astype(np.float32)
    
    # 3. Gamma 校正 (I_new = I_old ^ gamma)
    # 通常 gamma < 1.0 会让暗部变亮
    if gamma != 1.0:
        img_array = img_array / 255.0
        img_array = np.power(img_array, gamma)
        img_array = img_array * 255.0
    
    return np.clip(img_array, 0, 255).astype(np.uint8)

for file_path_ori in tqdm(file_paths_ori):
    # 读取原始RGB图像 (PIL)
    rgb_pil = Image.open(file_path_ori).convert('RGB')
    
    # === [新增] 在生成点云前优化颜色 ===
    # 应用亮度、饱和度和Gamma调整
    rgb_array = adjust_image_color(rgb_pil, args.brightness, args.saturation, args.gamma)
    # =================================

    # 用于边缘检测的灰度图处理 (保持原有逻辑不变)
    rgb_pil_ = read_image_safe(file_path_ori)
    if rgb_pil_ is None: continue
    imgray = cv2.cvtColor(rgb_pil_, cv2.COLOR_BGR2GRAY)
    ret, thresh = cv2.threshold(imgray, 5, 255, 0)
    contours, hierarchy = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    black_edge = np.ones_like(imgray, dtype=np.uint8)
    cv2.drawContours(black_edge, contours, -1, color=0, thickness=cv2.FILLED)
    black_edge = np.array(black_edge) 

    # 读取深度图
    base_name = os.path.splitext(os.path.basename(file_path_ori))[0]
    # 尝试匹配 png 或 jpg 后缀
    possible_exts = ['.png', '.jpg', '.jpeg']
    depth_file = None
    for ext in possible_exts:
        temp_path = os.path.join(depth_file_path, args.suffix + base_name + args.appendix + ext)
        if os.path.exists(temp_path):
            depth_file = temp_path
            break
            
    if depth_file and os.path.exists(depth_file):
        depth_pil = Image.open(depth_file).convert('L')
        depth_array = np.array(depth_pil)
        
        # 处理黑边
        depth_array[black_edge == 1] = 255 
        if args.label_type == "Disparity":
            depth_array[black_edge == 1] = 0

        w, h = depth_pil.size
        u_u0, v_v0 = init_image_coor(h, w)
        depth_tensor = torch.tensor(depth_array / 255.0).unsqueeze(0).unsqueeze(0)
        
        if args.label_type == "Disparity":
            depth_tensor = 1 - depth_tensor 

        # 计算焦距
        if args.focal_length != -1:
            focal_length = args.focal_length
        else:
            focal_length = (depth_array.shape[0] // 2 / np.tan((args.FoV / 2.0) * np.pi / 180))
        
        # 计算坐标
        x = u_u0 * depth_tensor / focal_length
        y = v_v0 * depth_tensor / focal_length
        
        # === [新增] 应用 Z轴缩放 ===
        # 乘以 args.z_scale 来拉伸深度
        z = depth_tensor * args.z_scale
        # ==========================

        data_ply = np.stack([
            x.reshape(-1), 
            y.reshape(-1), 
            z.reshape(-1), 
            rgb_array[:, :, 0].reshape(-1), 
            rgb_array[:, :, 1].reshape(-1), 
            rgb_array[:, :, 2].reshape(-1)
        ]).reshape(6, -1).T

        # 生成PLY内容
        points = [f"{p[0]:.4f} {p[1]:.4f} {p[2]:.4f} {int(p[3])} {int(p[4])} {int(p[5])} 0\n" for p in data_ply]
        ply_name = base_name + ".ply"
        ply_path = os.path.join(args.outdir, ply_name)
        
        with open(ply_path, "w") as file:
            file.write(f'''ply
format ascii 1.0
element vertex {len(points)}
property float x
property float y
property float z
property uchar red
property uchar green
property uchar blue
property uchar alpha
end_header
{''.join(points)}''')
    else:
        print(f'未找到对应的深度图: {base_name}')