import torch.nn as nn
import torch
import torch.nn.functional as F
import numpy as np
from pathlib import Path
import cv2
import trimesh
import os

from model.archs.decoders.shape_texture_net import TetTexNet
from model.archs.unet import UNetPP
from util.renderer import Renderer
from model.archs.mlp_head import SdfMlp, RgbMlp
import xatlas
from instant_texture import Converter


class Dummy:
    pass

class CRM(nn.Module):
    def __init__(self, specs):
        super(CRM, self).__init__()

        self.specs = specs
        # configs
        input_specs = specs["Input"]
        self.input = Dummy()
        self.input.scale = input_specs['scale']
        self.input.resolution = input_specs['resolution']
        self.tet_grid_size = input_specs['tet_grid_size']
        self.camera_angle_num = input_specs['camera_angle_num']

        self.arch = Dummy()
        self.arch.fea_concat = specs["ArchSpecs"]["fea_concat"]
        self.arch.mlp_bias = specs["ArchSpecs"]["mlp_bias"]

        self.dec = Dummy()
        self.dec.c_dim = specs["DecoderSpecs"]["c_dim"]
        self.dec.plane_resolution = specs["DecoderSpecs"]["plane_resolution"]

        self.geo_type = specs["Train"].get("geo_type", "flex") # "dmtet" or "flex"

        self.unet2 = UNetPP(in_channels=self.dec.c_dim)

        mlp_chnl_s = 3 if self.arch.fea_concat else 1  # 3 for queried triplane feature concatenation
        self.decoder = TetTexNet(plane_reso=self.dec.plane_resolution, fea_concat=self.arch.fea_concat)

        if self.geo_type == "flex":
            self.weightMlp = nn.Sequential(
                            nn.Linear(mlp_chnl_s * 32 * 8, 512),
                            nn.SiLU(),
                            nn.Linear(512, 21))         
            
        self.sdfMlp = SdfMlp(mlp_chnl_s * 32, 512, bias=self.arch.mlp_bias) 
        self.rgbMlp = RgbMlp(mlp_chnl_s * 32, 512, bias=self.arch.mlp_bias)
        # self.renderer = Renderer(tet_grid_size=self.tet_grid_size, camera_angle_num=self.camera_angle_num,
        #                          scale=self.input.scale, geo_type = self.geo_type)


        self.spob = True if specs['Pretrain']['mode'] is None else False  # whether to add sphere
        self.radius = specs['Pretrain']['radius']  # used when spob

        self.denoising = True
        from diffusers import DDIMScheduler
        self.scheduler = DDIMScheduler.from_pretrained("stabilityai/stable-diffusion-2-1-base", subfolder="scheduler")

    def decode(self, data, triplane_feature2):
        if self.geo_type == "flex":
            tet_verts = self.renderer.flexicubes.verts.unsqueeze(0)
            tet_indices = self.renderer.flexicubes.indices

        dec_verts = self.decoder(triplane_feature2, tet_verts)
        out = self.sdfMlp(dec_verts)

        weight = None
        if self.geo_type == "flex":
            grid_feat = torch.index_select(input=dec_verts, index=self.renderer.flexicubes.indices.reshape(-1),dim=1)
            grid_feat = grid_feat.reshape(dec_verts.shape[0], self.renderer.flexicubes.indices.shape[0], self.renderer.flexicubes.indices.shape[1] * dec_verts.shape[-1])
            weight = self.weightMlp(grid_feat)
            weight = weight * 0.1

        pred_sdf, deformation = out[..., 0], out[..., 1:]
        if self.spob:
            pred_sdf = pred_sdf + self.radius - torch.sqrt((tet_verts**2).sum(-1))

        _, verts, faces = self.renderer(data, pred_sdf, deformation, tet_verts, tet_indices, weight= weight)
        return verts[0].unsqueeze(0), faces[0].int()

    def export_mesh(self, data, out_dir, tri_fea_2 = None):
        verts = data['verts']
        faces = data['faces']

        dec_verts = self.decoder(tri_fea_2, verts.unsqueeze(0))
        colors = self.rgbMlp(dec_verts).squeeze().detach().cpu().numpy()
        print(f"[DEBUG] Raw colors min: {np.min(colors)}, max: {np.max(colors)}")
        # Expect predicted colors value range from [-1, 1]
        colors = (colors * 0.5 + 0.5).clip(0, 1)
        print(f"[DEBUG] Normalized colors min: {np.min(colors)}, max: {np.max(colors)}")

        verts = verts[..., [0, 2, 1]]
        verts[..., 0]*= -1
        verts[..., 2]*= -1
        verts = verts.squeeze().cpu().numpy()
        faces = faces[..., [2, 1, 0]][..., [0, 2, 1]]#[..., [1, 0, 2]]
        faces = faces.squeeze().cpu().numpy()#faces[..., [2, 1, 0]].squeeze().cpu().numpy()

        # export the final mesh
        with torch.no_grad():
            mesh = trimesh.Trimesh(verts, faces, vertex_colors=colors, process=False) # important, process=True leads to seg fault...
            mesh.export(f'{out_dir}.obj')

    def export_mesh_wt_uv(self, data, out_dir, tri_fea_2=None):
        """Modified to use InstantTexture for GLB output"""
        # First export the mesh with vertex colors to a temporary OBJ file
        obj_path = f"{out_dir}.obj"
        self.export_mesh(data, obj_path[:-4], tri_fea_2) # Pass base name without .obj
        
        # Then convert to textured GLB using InstantTexture
        glb_path = f"{out_dir}.glb"
        
        converter = Converter(texture_size=512)
        converter.convert(obj_path, output_mesh_path=glb_path)
        
        # Clean up intermediate OBJ and MTL files
        for ext in [".obj", ".mtl", ".png"]:
            try:
                os.remove(f"{out_dir}{ext}")
            except OSError:
                pass
            
        return glb_path