import torch
import numpy as np
from PIL import Image
import pymeshlab as ml
from pytorch3d.renderer import TexturesVertex
from pytorch3d.structures import Meshes
from rembg import new_session, remove
import trimesh
from typing import List, Tuple
import torch.nn.functional as F

# Constants
NEG_PROMPT = "sketch, sculpture, hand drawing, outline, single color, NSFW, lowres, bad anatomy, bad hands, text, error, missing fingers, yellow sleeves, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, (worst quality:1.4), (low quality:1.4)"

# CUDA Configuration
CUDA_PROVIDERS = [
    ('CUDAExecutionProvider', {
        'device_id': 0,
        'arena_extend_strategy': 'kSameAsRequested',
        'gpu_mem_limit': 8 * 1024 * 1024 * 1024,
        'cudnn_conv_algo_search': 'HEURISTIC',
    })
]

# Initialize rembg session
rembg_session = new_session(providers=CUDA_PROVIDERS)

# Mesh Loading and Conversion Functions
def load_mesh_with_trimesh(file_name, file_type=None):
    mesh = trimesh.load(file_name, file_type=file_type)
    if isinstance(mesh, trimesh.Scene):
        mesh = _process_trimesh_scene(mesh)
    
    vertices = torch.from_numpy(mesh.vertices).T
    faces = torch.from_numpy(mesh.faces).T
    colors = _get_mesh_colors(mesh)
    
    return vertices, faces, colors

def _process_trimesh_scene(mesh):
    from io import BytesIO
    with BytesIO() as f:
        mesh.export(f, file_type="obj")
        f.seek(0)
        mesh = trimesh.load(f, file_type="obj")
    if isinstance(mesh, trimesh.Scene):
        mesh = trimesh.util.concatenate(
            tuple(trimesh.Trimesh(vertices=g.vertices, faces=g.faces)
                for g in mesh.geometry.values()))
    return mesh

def _get_mesh_colors(mesh):
    if mesh.visual is not None and hasattr(mesh.visual, 'vertex_colors'):
        return torch.from_numpy(mesh.visual.vertex_colors)[..., :3].T / 255.
    return torch.ones_like(mesh.vertices.T) * 0.5

# Mesh Conversion Functions
def meshlab_mesh_to_py3dmesh(mesh: ml.Mesh) -> Meshes:
    verts = torch.from_numpy(mesh.vertex_matrix()).float()
    faces = torch.from_numpy(mesh.face_matrix()).long()
    colors = torch.from_numpy(mesh.vertex_color_matrix()[..., :3]).float()
    textures = TexturesVertex(verts_features=[colors])
    return Meshes(verts=[verts], faces=[faces], textures=textures)

def py3dmesh_to_meshlab_mesh(meshes: Meshes) -> ml.Mesh:
    colors_in = F.pad(meshes.textures.verts_features_packed().cpu().float(), [0,1], value=1).numpy().astype(np.float64)
    return ml.Mesh(
        vertex_matrix=meshes.verts_packed().cpu().float().numpy().astype(np.float64),
        face_matrix=meshes.faces_packed().cpu().long().numpy().astype(np.int32),
        v_normals_matrix=meshes.verts_normals_packed().cpu().float().numpy().astype(np.float64),
        v_color_matrix=colors_in)

# Normal Map Rotation Functions
def rotate_normalmap_by_angle(normal_map: np.ndarray, angle: float):
    angle_rad = np.radians(angle)
    R = np.array([
        [np.cos(angle_rad), 0, np.sin(angle_rad)],
        [0, 1, 0],
        [-np.sin(angle_rad), 0, np.cos(angle_rad)]
    ])
    return np.dot(normal_map.reshape(-1, 3), R.T).reshape(normal_map.shape)

def rotate_normals(normal_pils, return_types='np', rotate_direction=1):
    n_views = len(normal_pils)
    ret = []
    for idx, rgba_normal in enumerate(normal_pils):
        normal_np = _process_normal_map(rgba_normal, idx, n_views, rotate_direction)
        ret.append(_format_output(normal_np, return_types))
    return ret

def _process_normal_map(rgba_normal, idx, n_views, rotate_direction):
    normal_np = np.array(rgba_normal)[:, :, :3] / 255 * 2 - 1
    alpha_np = np.array(rgba_normal)[:, :, 3] / 255
    normal_np = rotate_normalmap_by_angle(normal_np, rotate_direction * idx * (360 / n_views))
    normal_np = (normal_np + 1) / 2 * alpha_np[..., None]
    return np.concatenate([normal_np * 255, alpha_np[:, :, None] * 255], axis=-1)

def _format_output(normal_np, return_types):
    if return_types == 'np':
        return normal_np
    elif return_types == 'pil':
        return Image.fromarray(normal_np.astype(np.uint8))
    else:
        raise ValueError(f"return_types should be 'np' or 'pil', but got {return_types}")

# Background Change Functions
def change_bkgd(img_pils, new_bkgd=(0., 0., 0.)):
    new_bkgd = np.array(new_bkgd).reshape(1, 1, 3)
    return [_process_image(rgba_img, new_bkgd) for rgba_img in img_pils]

def _process_image(rgba_img, new_bkgd):
    img_np = np.array(rgba_img)[:, :, :3] / 255
    alpha_np = np.array(rgba_img)[:, :, 3] / 255
    ori_bkgd = img_np[:1, :1]
    alpha_np_clamp = np.clip(alpha_np, 1e-6, 1)
    ori_img_np = (img_np - ori_bkgd * (1 - alpha_np[..., None])) / alpha_np_clamp[..., None]
    img_np = np.where(alpha_np[..., None] > 0.05, ori_img_np * alpha_np[..., None] + new_bkgd * (1 - alpha_np[..., None]), new_bkgd)
    rgba_img_np = np.concatenate([img_np * 255, alpha_np[..., None] * 255], axis=-1)
    return Image.fromarray(rgba_img_np.astype(np.uint8))

# Mesh Cleaning Function
def simple_clean_mesh(pyml_mesh: ml.Mesh, apply_smooth=True, stepsmoothnum=1, apply_sub_divide=False, sub_divide_threshold=0.25):
    ms = ml.MeshSet()
    ms.add_mesh(pyml_mesh, "cube_mesh")
    
    if apply_smooth:
        ms.apply_filter("apply_coord_laplacian_smoothing", stepsmoothnum=stepsmoothnum, cotangentweight=False)
    if apply_sub_divide:
        ms.apply_filter("meshing_repair_non_manifold_vertices")
        ms.apply_filter("meshing_repair_non_manifold_edges", method='Remove Faces')
        ms.apply_filter("meshing_surface_subdivision_loop", iterations=2, threshold=ml.PercentageValue(sub_divide_threshold))
    return meshlab_mesh_to_py3dmesh(ms.current_mesh())

# Image Processing Functions
def expand2square(pil_img, background_color):
    width, height = pil_img.size
    if width == height:
        return pil_img
    new_size = max(width, height)
    result = Image.new(pil_img.mode, (new_size, new_size), background_color)
    offset = ((new_size - width) // 2, (new_size - height) // 2)
    result.paste(pil_img, offset)
    return result

def simple_preprocess(input_image, rembg_session=rembg_session, background_color=255):
    RES = 2048
    input_image.thumbnail([RES, RES], Image.Resampling.LANCZOS)
    if input_image.mode != 'RGBA':
        image_rem = input_image.convert('RGBA')
        input_image = remove(image_rem, alpha_matting=False, session=rembg_session)

    arr = np.asarray(input_image)
    alpha = arr[:, :, -1]
    x_nonzero, y_nonzero = np.nonzero(alpha > 60)
    x_min, x_max = x_nonzero.min(), x_nonzero.max()
    y_min, y_max = y_nonzero.min(), y_nonzero.max()
    arr = arr[x_min:x_max+1, y_min:y_max+1]
    input_image = Image.fromarray(arr)
    return expand2square(input_image, (background_color, background_color, background_color, 0))

# Mesh Saving Functions
def save_py3dmesh_with_trimesh_fast(meshes: Meshes, save_glb_path, apply_sRGB_to_LinearRGB=True):
    vertices = meshes.verts_packed().cpu().float().numpy()
    triangles = meshes.faces_packed().cpu().long().numpy()
    np_color = meshes.textures.verts_features_packed().cpu().float().numpy()
    
    if save_glb_path.endswith(".glb"):
        vertices[:, [0, 2]] = -vertices[:, [0, 2]]

    if apply_sRGB_to_LinearRGB:
        np_color = srgb_to_linear(np_color)
    
    mesh = trimesh.Trimesh(vertices=vertices, faces=triangles, vertex_colors=np_color)
    mesh.remove_unreferenced_vertices()
    mesh.export(save_glb_path)
    
    if save_glb_path.endswith(".glb"):
        fix_vert_color_glb(save_glb_path)
    print(f"Saved to {save_glb_path}")

def save_glb_and_video(save_mesh_prefix: str, meshes: Meshes, with_timestamp=True, **kwargs) -> Tuple[str, str]:
    import time
    if '.' in save_mesh_prefix:
        save_mesh_prefix = ".".join(save_mesh_prefix.split('.')[:-1])
    if with_timestamp:
        save_mesh_prefix = save_mesh_prefix + f"_{int(time.time())}"
    ret_mesh = save_mesh_prefix + ".glb"
    save_py3dmesh_with_trimesh_fast(meshes, ret_mesh)
    return ret_mesh, None

# Utility Functions
def srgb_to_linear(c_srgb):
    return np.where(c_srgb <= 0.04045, c_srgb / 12.92, ((c_srgb + 0.055) / 1.055) ** 2.4).clip(0, 1.)

def fix_vert_color_glb(mesh_path):
    from pygltflib import GLTF2, Material, PbrMetallicRoughness
    obj1 = GLTF2().load(mesh_path)
    obj1.meshes[0].primitives[0].material = 0
    obj1.materials.append(Material(
        pbrMetallicRoughness = PbrMetallicRoughness(
            baseColorFactor = [1.0, 1.0, 1.0, 1.0],
            metallicFactor = 0.,
            roughnessFactor = 1.0,
        ),
        emissiveFactor = [0.0, 0.0, 0.0],
        doubleSided = True,
    ))
    obj1.save(mesh_path)

def init_target(img_pils, new_bkgd=(0., 0., 0.), device="cuda"):
    new_bkgd = torch.tensor(new_bkgd, dtype=torch.float32).view(1, 1, 3).to(device)
    imgs = torch.stack([torch.from_numpy(np.array(img, dtype=np.float32)) for img in img_pils]).to(device) / 255
    img_nps, alpha_nps = imgs[..., :3], imgs[..., 3]
    ori_bkgds = img_nps[:, :1, :1]
    
    alpha_nps_clamp = torch.clamp(alpha_nps, 1e-6, 1)
    ori_img_nps = (img_nps - ori_bkgds * (1 - alpha_nps.unsqueeze(-1))) / alpha_nps_clamp.unsqueeze(-1)
    ori_img_nps = torch.clamp(ori_img_nps, 0, 1)
    img_nps = torch.where(alpha_nps.unsqueeze(-1) > 0.05, ori_img_nps * alpha_nps.unsqueeze(-1) + new_bkgd * (1 - alpha_nps.unsqueeze(-1)), new_bkgd)

    return torch.cat([img_nps, alpha_nps.unsqueeze(-1)], dim=-1)

def save_obj_and_video(save_mesh_prefix: str, meshes: Meshes, with_timestamp=True, **kwargs) -> Tuple[str, str]:
    if '.' in save_mesh_prefix:
        save_mesh_prefix = ".".join(save_mesh_prefix.split('.')[:-1])
    if with_timestamp:
        save_mesh_prefix = save_mesh_prefix + f"_{int(time.time())}"
    ret_mesh = save_mesh_prefix + ".obj"
    
    vertices = meshes.verts_packed().cpu().float().numpy()
    triangles = meshes.faces_packed().cpu().long().numpy()
    np_color = meshes.textures.verts_features_packed().cpu().float().numpy()
    
    # Apply sRGB to LinearRGB conversion
    np_color = srgb_to_linear(np_color)
    
    mesh = trimesh.Trimesh(vertices=vertices, faces=triangles, vertex_colors=np_color)
    mesh.remove_unreferenced_vertices()
    mesh.export(ret_mesh)
    
    print(f"Saved to {ret_mesh}")
    
    return ret_mesh, None