import torch
import spaces
import gradio as gr
import os
import numpy as np
import trimesh
import mcubes
import imageio
from torchvision.utils import save_image
from PIL import Image
from transformers import AutoModel, AutoConfig
from rembg import remove, new_session
from functools import partial
from kiui.op import recenter
import kiui
from gradio_litmodel3d import LitModel3D
import shutil

def find_cuda():
    # 检查 CUDA_HOME 或 CUDA_PATH 环境变量是否已设置
    cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')

    if cuda_home and os.path.exists(cuda_home):
        return cuda_home

    # 在系统 PATH 中搜索 nvcc 可执行文件
    nvcc_path = shutil.which('nvcc')

    if nvcc_path:
        # 删除“bin/nvcc”部分，获取 CUDA 安装路径
        cuda_path = os.path.dirname(os.path.dirname(nvcc_path))
        return cuda_path

    return None

cuda_path = find_cuda()

if cuda_path:
    print(f"CUDA 已安装在：{cuda_path}")
else:
    print("未找到已安装的 CUDA 路径")

# 从 HF 加载预训练模型
class LRMGeneratorWrapper:
    def __init__(self):
        self.config = AutoConfig.from_pretrained("yanranxiaoxi/image-upscale", trust_remote_code=True, token=os.environ.get('MODEL_ACCESS_TOKEN'))
        self.model = AutoModel.from_pretrained("yanranxiaoxi/image-upscale", trust_remote_code=True, token=os.environ.get('MODEL_ACCESS_TOKEN'))
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model.to(self.device)
        self.model.eval()

    def forward(self, image, camera):
        return self.model(image, camera)

model_wrapper = LRMGeneratorWrapper()

# 处理输入图像
def preprocess_image(image, source_size):
    session = new_session("isnet-general-use")
    rembg_remove = partial(remove, session=session)
    image = np.array(image)
    image = rembg_remove(image)
    mask = rembg_remove(image, only_mask=True)
    image = recenter(image, mask, border_ratio=0.20)
    image = torch.tensor(image).permute(2, 0, 1).unsqueeze(0) / 255.0
    if image.shape[1] == 4:
        image = image[:, :3, ...] * image[:, 3:, ...] + (1 - image[:, 3:, ...])
    image = torch.nn.functional.interpolate(image, size=(source_size, source_size), mode='bicubic', align_corners=True)
    image = torch.clamp(image, 0, 1)
    return image

def get_normalized_camera_intrinsics(intrinsics: torch.Tensor):
    fx, fy = intrinsics[:, 0, 0], intrinsics[:, 0, 1]
    cx, cy = intrinsics[:, 1, 0], intrinsics[:, 1, 1]
    width, height = intrinsics[:, 2, 0], intrinsics[:, 2, 1]
    fx, fy = fx / width, fy / height
    cx, cy = cx / width, cy / height
    return fx, fy, cx, cy

def build_camera_principle(RT: torch.Tensor, intrinsics: torch.Tensor):
    fx, fy, cx, cy = get_normalized_camera_intrinsics(intrinsics)
    return torch.cat([
        RT.reshape(-1, 12),
        fx.unsqueeze(-1), fy.unsqueeze(-1), cx.unsqueeze(-1), cy.unsqueeze(-1),
    ], dim=-1)

def _default_intrinsics():
    fx = fy = 384
    cx = cy = 256
    w = h = 512
    intrinsics = torch.tensor([
        [fx, fy],
        [cx, cy],
        [w, h],
    ], dtype=torch.float32)
    return intrinsics

def _default_source_camera(batch_size: int = 1):
    canonical_camera_extrinsics = torch.tensor([[
        [0, 0, 1, 1],
        [1, 0, 0, 0],
        [0, 1, 0, 0],
    ]], dtype=torch.float32)
    canonical_camera_intrinsics = _default_intrinsics().unsqueeze(0)
    source_camera = build_camera_principle(canonical_camera_extrinsics, canonical_camera_intrinsics)
    return source_camera.repeat(batch_size, 1)

def _center_looking_at_camera_pose(camera_position: torch.Tensor, look_at: torch.Tensor = None, up_world: torch.Tensor = None):
    """
    camera_position: (M, 3)
    look_at: (3)
    up_world: (3)
    return: (M, 3, 4)
    """
    # 默认情况下，从原点向上为 pos-z
    if look_at is None:
        look_at = torch.tensor([0, 0, 0], dtype=torch.float32)
    if up_world is None:
        up_world = torch.tensor([0, 0, 1], dtype=torch.float32)
    look_at = look_at.unsqueeze(0).repeat(camera_position.shape[0], 1)
    up_world = up_world.unsqueeze(0).repeat(camera_position.shape[0], 1)

    z_axis = camera_position - look_at
    z_axis = z_axis / z_axis.norm(dim=-1, keepdim=True)
    x_axis = torch.cross(up_world, z_axis)
    x_axis = x_axis / x_axis.norm(dim=-1, keepdim=True)
    y_axis = torch.cross(z_axis, x_axis)
    y_axis = y_axis / y_axis.norm(dim=-1, keepdim=True)
    extrinsics = torch.stack([x_axis, y_axis, z_axis, camera_position], dim=-1)
    return extrinsics

def compose_extrinsic_RT(RT: torch.Tensor):
    """
    从 RT 生成标准形式的外差矩阵。
    分批输入/输出。
    """
    return torch.cat([
        RT,
        torch.tensor([[[0, 0, 0, 1]]], dtype=torch.float32).repeat(RT.shape[0], 1, 1).to(RT.device)
        ], dim=1)

def _build_camera_standard(RT: torch.Tensor, intrinsics: torch.Tensor):
    """
    RT: (N, 3, 4)
    intrinsics: (N, 3, 2), [[fx, fy], [cx, cy], [width, height]]
    """
    E = compose_extrinsic_RT(RT)
    fx, fy, cx, cy = get_normalized_camera_intrinsics(intrinsics)
    I = torch.stack([
        torch.stack([fx, torch.zeros_like(fx), cx], dim=-1),
        torch.stack([torch.zeros_like(fy), fy, cy], dim=-1),
        torch.tensor([[0, 0, 1]], dtype=torch.float32, device=RT.device).repeat(RT.shape[0], 1),
    ], dim=1)
    return torch.cat([
        E.reshape(-1, 16),
        I.reshape(-1, 9),
    ], dim=-1)

def _default_render_cameras(batch_size: int = 1):
    M = 80
    radius = 1.5
    elevation = 0
    camera_positions = []
    rand_theta = np.random.uniform(0, np.pi/180)
    elevation = np.radians(elevation)
    for i in range(M):
        theta = 2 * np.pi * i / M + rand_theta
        x = radius * np.cos(theta) * np.cos(elevation)
        y = radius * np.sin(theta) * np.cos(elevation)
        z = radius * np.sin(elevation)
        camera_positions.append([x, y, z])
    camera_positions = torch.tensor(camera_positions, dtype=torch.float32)
    extrinsics = _center_looking_at_camera_pose(camera_positions)

    render_camera_intrinsics = _default_intrinsics().unsqueeze(0).repeat(extrinsics.shape[0], 1, 1)
    render_cameras = _build_camera_standard(extrinsics, render_camera_intrinsics)
    return render_cameras.unsqueeze(0).repeat(batch_size, 1, 1)

@spaces.GPU
def generate_mesh(image, source_size=512, render_size=384, mesh_size=512, export_mesh=False, export_video=False, fps=30):
    image = preprocess_image(image, source_size).to(model_wrapper.device)
    source_camera = _default_source_camera(batch_size=1).to(model_wrapper.device)

    with torch.no_grad():
        planes = model_wrapper.forward(image, source_camera)

        if export_mesh:
            grid_out = model_wrapper.model.synthesizer.forward_grid(planes=planes, grid_size=mesh_size)
            vtx, faces = mcubes.marching_cubes(grid_out['sigma'].float().squeeze(0).squeeze(-1).cpu().numpy(), 1.0)
            vtx = vtx / (mesh_size - 1) * 2 - 1
            vtx_tensor = torch.tensor(vtx, dtype=torch.float32, device=model_wrapper.device).unsqueeze(0)
            vtx_colors = model_wrapper.model.synthesizer.forward_points(planes, vtx_tensor)['rgb'].float().squeeze(0).cpu().numpy()
            vtx_colors = (vtx_colors * 255).astype(np.uint8)
            mesh = trimesh.Trimesh(vertices=vtx, faces=faces, vertex_colors=vtx_colors)

            mesh_path = "xiaoxis_mesh.obj"
            mesh.export(mesh_path, 'obj')

            return None, mesh_path

        if export_video:
            render_cameras = _default_render_cameras(batch_size=1).to(model_wrapper.device)
            frames = []
            chunk_size = 1
            for i in range(0, render_cameras.shape[1], chunk_size):
                frame_chunk = model_wrapper.model.synthesizer(
                    planes,
                    render_cameras[:, i:i + chunk_size],
                    render_size,
                    render_size,
                    0,
                    0
                )
                frames.append(frame_chunk['images_rgb'])

            frames = torch.cat(frames, dim=1)
            frames = frames.squeeze(0)
            frames = (frames.permute(0, 2, 3, 1).cpu().numpy() * 255).astype(np.uint8)

            video_path = "xiaoxis_video.mp4"
            imageio.mimwrite(video_path, frames, fps=fps)

            return None, video_path

        return planes, None

    return None, None

def step_1_generate_planes(image):
    planes, _ = generate_mesh(image)
    return planes

def step_2_generate_obj(image):
    _, mesh_path = generate_mesh(image, export_mesh=True)
    return mesh_path, mesh_path

def step_3_generate_video(image):
    _, video_path = generate_mesh(image, export_video=True)
    return video_path, video_path

# 从 assets 文件夹中设置示例文件，并限制最多读取 10 个文件
example_folder = "assets"
examples = [os.path.join(example_folder, f) for f in os.listdir(example_folder) if f.endswith(('.png', '.jpg', '.jpeg', '.webp'))][:10]


with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown("""
        # 图像升维计算模型：EMU Video 的衍生尝试

        我们利用视频扩散模型作为多视图数据生成器，从而促进可扩展 3D 生成模型的学习。以下展示了视频扩散模型作为多视图数据引擎的潜力，能够生成无限规模的合成数据以支持可扩展的训练。我们提出的模型从合成数据中学习，在生成 3D 资产方面表现出卓越的性能。

        除了当前状态之外，我们的模型还具有高度可扩展性，并且可以根据合成数据和 3D 数据的数量进行扩展，为 3D 生成模型铺平了新的道路。
        """)

    with gr.Row():
        with gr.Column():
            img_input = gr.Image(type="pil", label="输入图像")
            examples_component = gr.Examples(examples=examples, inputs=img_input, outputs=None, examples_per_page=5)
            generate_mesh_button = gr.Button("生成模型")
            generate_video_button = gr.Button("生成视频")
        with gr.Column():
            model_output = LitModel3D(
                clear_color=[0, 0, 0, 0],  # 可调整背景颜色，以获得更好的对比度
                label="模型可视化",
                scale=1.0,
                tonemapping="aces",        # 可使用 aces 色调映射，使灯光更逼真
                exposure=1.1,              # 可调节曝光以控制亮度
                contrast=1.1,              # 可略微增加对比度，以获得更好的深度
                camera_position=(0, 0, 2), # 将设置初始摄像机位置，使模型居中
                zoom_speed=0.5,            # 将调整变焦速度，以便更好地控制
                pan_speed=0.5,             # 将调整摇摄速度，以便更好地控制
                interactive=False          # 这样用户就可以与模型进行交互
            )

    with gr.Row():
        with gr.Column():
            obj_file_output = gr.File(label="下载 .obj 文件")
            video_file_output = gr.File(label="下载视频")
        with gr.Column():
            video_output = gr.Video(label="360° 视频")


    # 清除输出
    def clear_model_viewer():
        """在加载新模型前重置 Gradio。"""
        update_output = gr.update(value=None)
        return update_output, update_output

    # 首先清除输出的数据
    img_input.change(clear_model_viewer, inputs=None, outputs=[model_output, video_output])

    # 然后生成模型和视频
    generate_mesh_button.click(step_2_generate_obj, inputs=img_input, outputs=[obj_file_output, model_output])
    generate_video_button.click(step_3_generate_video, inputs=img_input, outputs=[video_file_output, video_output])

demo.launch(
    auth=(os.environ.get('AUTH_USERNAME'), os.environ.get('AUTH_PASSWORD'))
)