import FreeCAD
import Mesh
import numpy as np

# Cargar la malla
mesh = Mesh.Mesh("ruta/a/tu/malla.stl")

# Extraer vértices y facetas
vertices = [vertex.Point for vertex in mesh.Points]
facets = [facet.PointIndices for facet in mesh.Facets]

# Crear una representación de la malla
mesh_representation = {
    "vertices": vertices,
    "facets": facets
}


# Calcular la dirección del flujo de agua
def calculate_slope_direction(facet):
    v1 = np.array(FreeCAD.Vector(facet.Points[0]))
    v2 = np.array(FreeCAD.Vector(facet.Points[1]))
    v3 = np.array(FreeCAD.Vector(facet.Points[2]))

    edge1 = v2 - v1
    edge2 = v3 - v1

    normal = np.cross(edge1, edge2)
    normal[2] = abs(normal[2])
    normal = normal / np.linalg.norm(normal)

    slope_direction = np.array([0.0, 0.0, 0.0])
    slope_direction[:2] = np.mean([edge1[:2], edge2[:2]], axis=0)
    slope_direction[2] = -(normal[:2].dot(slope_direction[:2]) / normal[2])
    slope_direction = slope_direction / np.linalg.norm(slope_direction)

    return slope_direction


flow_directions = []
for ind, facet in enumerate(facets):
    direction = calculate_slope_direction(facet)
    flow_directions.append(direction)
    if ind == 1000:
        break

# Determinar cuencas hidrológicas
basins = [None] * len(vertices)


def find_basin(vertex_index, flow_directions, facets):
    if basins[vertex_index] is not None:
        return basins[vertex_index]

    for facet, direction in zip(facets, flow_directions):
        if vertex_index in facet:
            v_indices = [idx for idx in facet if idx != vertex_index]
            min_vertex_index = v_indices[0]
            min_height = vertices[min_vertex_index][2]

            for vi in v_indices[1:]:
                if vertices[vi][2] < min_height:
                    min_height = vertices[vi][2]
                    min_vertex_index = vi

            if vertices[vertex_index][2] > min_height:
                basins[vertex_index] = find_basin(min_vertex_index, flow_directions, facets)
                return basins[vertex_index]

    basins[vertex_index] = vertex_index
    return vertex_index


for i in range(len(vertices)):
    find_basin(i, flow_directions, facets)

# Calcular acumulación de agua
water_accumulation = [0] * len(vertices)


def accumulate_water(vertex_index, flow_directions, facets, water_accumulation):
    if water_accumulation[vertex_index] > 0:
        return water_accumulation[vertex_index]

    for facet, direction in zip(facets, flow_directions):
        if vertex_index in facet:
            destination_vertex = None
            min_height = float('inf')
            for vi in facet:
                if vertices[vi][2] < min_height:
                    min_height = vertices[vi][2]
                    destination_vertex = vi
            if destination_vertex is not None:
                water_accumulation[vertex_index] += accumulate_water(destination_vertex, flow_directions, facets,
                                                                     water_accumulation) + 1

    return water_accumulation[vertex_index]


for i in range(len(vertices)):
    accumulate_water(i, flow_directions, facets, water_accumulation)

# Identificar puntos de drenaje
drainage_points = []
threshold = 10


def identify_drainage_points(vertices, water_accumulation, threshold):
    for i, water in enumerate(water_accumulation):
        if water >= threshold:
            drainage_points.append(vertices[i])


identify_drainage_points(vertices, water_accumulation, threshold)

for point in drainage_points:
    print("Punto de drenaje en:", point)