import cv2
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
import trimesh
from pythreejs import *
from IPython.display import display
import collada

class Topografia:
    def __init__(self) -> None:
        pass

    def contraste(self, img: np.ndarray) -> np.ndarray:
        return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    def porcentage(self, img: np.ndarray) -> float:
        gray = self.contraste(img)
        dim = np.shape(gray)
        cells = dim[0]*dim[1]
        return np.sum(gray) / (cells * 255)
    
    def model3D(self, grades: np.array) -> np.ndarray:
        n = grades.shape[0]
        x, y = np.meshgrid(np.arange(n), np.arange(n))
        z = grades
        print("A")

        vertices = np.stack((x.ravel(), y.ravel(), z.ravel()), axis=-1)
        print("B")

        i, j = np.meshgrid(np.arange(n - 1), np.arange(n - 1))
        i, j = i.ravel(), j.ravel()
        print("C")

        idx1 = i * n + j
        idx2 = idx1 + 1
        idx3 = idx1 + n
        idx4 = idx3 + 1
        print("D")

        faces = np.stack((idx1, idx2, idx3, idx2, idx4, idx3), axis=-1).reshape(-1, 3)
        print("E")

        mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
        # Convertir los vértices y caras a los formatos necesarios para pycollada
        collada_faces = faces.flatten().astype(int)
        collada_vertices = vertices.flatten().astype(float)

        # Crear la escena COLLADA
        mesh_data = collada.Collada()
        
        # Definir el formato de los vértices
        vert_src = collada.source.FloatSource("verts-array", collada_vertices, ('X', 'Y', 'Z'))
        geom = collada.geometry.Geometry(mesh_data, "geometry0", "mygeom", [vert_src])
        
        # Definir las caras
        input_list = collada.source.InputList()
        input_list.addInput(0, 'VERTEX', "#verts-array")
        triset = geom.createTriangleSet(collada_faces, input_list)
        geom.primitives.append(triset)
        
        mesh_data.geometries.append(geom)

        # Crear un efecto simple
        effect = collada.material.Effect("effect0", [], "phong", diffuse=(0.8, 0.8, 0.8), specular=(0.5, 0.5, 0.5))
        mesh_data.effects.append(effect)
        
        # Crear el material usando el efecto
        mat = collada.material.Material("material0", "mymaterial", effect=effect)
        mesh_data.materials.append(mat)
        geomnode = collada.scene.GeometryNode(geom, [])
        node = collada.scene.Node("node0", children=[geomnode])
        myscene = collada.scene.Scene("myscene", [node])
        mesh_data.scenes.append(myscene)
        mesh_data.scene = myscene

        # Exportar a .dae
        mesh_data.write('surface.dae')
        print("Exportación a surface.dae completada.")
        return


        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')

        # Graficar la superficie
        surf = ax.plot_surface(x, y, grades, cmap=cm.coolwarm, linewidth=0, antialiased=False)
        fig.colorbar(surf, ax=ax, shrink=0.6, aspect=5)

        # Añadir etiquetas a los ejes
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')

        plt.savefig("app/log/top.png")


if __name__ == '__main__':
    # import sys
    # import os
    # sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
    # from app.lib.colorimetria import Colorimetria

    from colorimetria import Colorimetria

    c = Colorimetria()
    path = "base/img.png"
    img = cv2.imread(path)
    matrix = c.get_grade_image(img)+1

    t = Topografia()
    t.model3D(matrix)
    exit()



    cv2.imwrite("log/topografia.png", t.contraste())
    print(t.porcentage())