#include <emscripten.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <math.h>

// Define facial landmark structure
typedef struct {
  float x;
  float y;
  float z;
  int type;  // Type of landmark (e.g., eye, nose, mouth, etc.)
} FacialLandmark;

// Define vertex structure
typedef struct {
  float x;
  float y;
  float z;
  float nx;  // Normal x component
  float ny;  // Normal y component 
  float nz;  // Normal z component
  float importance;  // Importance value for mesh simplification
} Vertex;

// Define triangle structure
typedef struct {
  int v1;
  int v2;
  int v3;
} Triangle;

// Define mesh structure
typedef struct {
  Vertex* vertices;
  int vertexCount;
  Triangle* triangles;
  int triangleCount;
  FacialLandmark* landmarks;
  int landmarkCount;
} Mesh;

// Global mesh state for current processing
Mesh* currentMesh = NULL;

/**
 * Initialize a new mesh for processing
 */
EMSCRIPTEN_KEEPALIVE
int initMesh(float* vertexData, int vertexCount, int* triangleData, int triangleCount) {
  // Free any existing mesh
  if (currentMesh != NULL) {
    free(currentMesh->vertices);
    free(currentMesh->triangles);
    if (currentMesh->landmarks != NULL) {
      free(currentMesh->landmarks);
    }
    free(currentMesh);
  }
  
  // Allocate new mesh
  currentMesh = (Mesh*)malloc(sizeof(Mesh));
  if (currentMesh == NULL) {
    return 0;  // Allocation failed
  }
  
  // Allocate and copy vertices
  currentMesh->vertices = (Vertex*)malloc(vertexCount * sizeof(Vertex));
  if (currentMesh->vertices == NULL) {
    free(currentMesh);
    currentMesh = NULL;
    return 0;  // Allocation failed
  }
  
  currentMesh->vertexCount = vertexCount;
  
  // Copy vertex data (assuming packed format: x,y,z,nx,ny,nz)
  for (int i = 0; i < vertexCount; i++) {
    Vertex* v = &currentMesh->vertices[i];
    v->x = vertexData[i * 6];
    v->y = vertexData[i * 6 + 1];
    v->z = vertexData[i * 6 + 2];
    v->nx = vertexData[i * 6 + 3];
    v->ny = vertexData[i * 6 + 4];
    v->nz = vertexData[i * 6 + 5];
    v->importance = 0.0f;  // Initialize importance to 0
  }
  
  // Allocate and copy triangles
  currentMesh->triangles = (Triangle*)malloc(triangleCount * sizeof(Triangle));
  if (currentMesh->triangles == NULL) {
    free(currentMesh->vertices);
    free(currentMesh);
    currentMesh = NULL;
    return 0;  // Allocation failed
  }
  
  currentMesh->triangleCount = triangleCount;
  
  // Copy triangle data
  for (int i = 0; i < triangleCount; i++) {
    Triangle* t = &currentMesh->triangles[i];
    t->v1 = triangleData[i * 3];
    t->v2 = triangleData[i * 3 + 1];
    t->v3 = triangleData[i * 3 + 2];
  }
  
  // Initialize landmarks to NULL
  currentMesh->landmarks = NULL;
  currentMesh->landmarkCount = 0;
  
  return 1;  // Success
}

/**
 * Set facial landmarks for the current mesh
 */
EMSCRIPTEN_KEEPALIVE
int setLandmarks(float* landmarkData, int landmarkCount) {
  if (currentMesh == NULL) {
    return 0;  // No mesh initialized
  }
  
  // Free existing landmarks if any
  if (currentMesh->landmarks != NULL) {
    free(currentMesh->landmarks);
  }
  
  // Allocate landmarks
  currentMesh->landmarks = (FacialLandmark*)malloc(landmarkCount * sizeof(FacialLandmark));
  if (currentMesh->landmarks == NULL) {
    return 0;  // Allocation failed
  }
  
  currentMesh->landmarkCount = landmarkCount;
  
  // Copy landmark data (assuming packed format: x,y,z,type)
  for (int i = 0; i < landmarkCount; i++) {
    FacialLandmark* l = &currentMesh->landmarks[i];
    l->x = landmarkData[i * 4];
    l->y = landmarkData[i * 4 + 1];
    l->z = landmarkData[i * 4 + 2];
    l->type = (int)landmarkData[i * 4 + 3];
  }
  
  return 1;  // Success
}

/**
 * Calculate distance between two points in 3D space
 */
float distance(float x1, float y1, float z1, float x2, float y2, float z2) {
  float dx = x2 - x1;
  float dy = y2 - y1;
  float dz = z2 - z1;
  return sqrtf(dx*dx + dy*dy + dz*dz);
}

/**
 * Calculate importance values for vertices based on landmarks
 */
EMSCRIPTEN_KEEPALIVE
int calculateVertexImportance() {
  if (currentMesh == NULL) {
    return 0;  // No mesh initialized
  }
  
  // Initialize all vertices with base importance
  for (int i = 0; i < currentMesh->vertexCount; i++) {
    currentMesh->vertices[i].importance = 0.1f;
  }
  
  // Increase importance for vertices near landmarks
  for (int i = 0; i < currentMesh->landmarkCount; i++) {
    FacialLandmark* landmark = &currentMesh->landmarks[i];
    
    // Importance factor based on landmark type
    float importanceFactor = 1.0f;
    switch (landmark->type) {
      case 0:  // Eyes landmarks
      case 1:  // Eye corners
        importanceFactor = 2.0f;
        break;
      case 2:  // Nose
        importanceFactor = 1.5f;
        break;
      case 3:  // Mouth
      case 4:  // Lip corners
        importanceFactor = 2.0f;
        break;
      case 5:  // Jawline
        importanceFactor = 1.2f;
        break;
      default:
        importanceFactor = 1.0f;
    }
    
    // Assign importance to vertices based on distance to landmark
    for (int j = 0; j < currentMesh->vertexCount; j++) {
      Vertex* vertex = &currentMesh->vertices[j];
      
      float dist = distance(vertex->x, vertex->y, vertex->z, 
                           landmark->x, landmark->y, landmark->z);
      
      // Exponential falloff of importance based on distance
      // Closer vertices get higher importance
      float importance = importanceFactor * expf(-dist * 5.0f);
      
      // Keep maximum importance value if affected by multiple landmarks
      if (importance > vertex->importance) {
        vertex->importance = importance;
      }
    }
  }
  
  // Adjust importance based on curvature and edges
  calculateCurvatureImportance();
  
  return 1;  // Success
}

/**
 * Calculate curvature and adjust vertex importance accordingly
 */
void calculateCurvatureImportance() {
  if (currentMesh == NULL || currentMesh->vertexCount == 0) {
    return;
  }
  
  // For each vertex, approximate curvature by comparing with neighbors
  for (int i = 0; i < currentMesh->triangleCount; i++) {
    Triangle* tri = &currentMesh->triangles[i];
    Vertex* v1 = &currentMesh->vertices[tri->v1];
    Vertex* v2 = &currentMesh->vertices[tri->v2];
    Vertex* v3 = &currentMesh->vertices[tri->v3];
    
    // Calculate edge lengths
    float edge1 = distance(v1->x, v1->y, v1->z, v2->x, v2->y, v2->z);
    float edge2 = distance(v2->x, v2->y, v2->z, v3->x, v3->y, v3->z);
    float edge3 = distance(v3->x, v3->y, v3->z, v1->x, v1->y, v1->z);
    
    // Calculate average edge length for this triangle
    float avgEdge = (edge1 + edge2 + edge3) / 3.0f;
    
    // Calculate normal difference as a measure of curvature
    float normalDiff1 = fabsf(v1->nx * v2->nx + v1->ny * v2->ny + v1->nz * v2->nz - 1.0f);
    float normalDiff2 = fabsf(v2->nx * v3->nx + v2->ny * v3->ny + v2->nz * v3->nz - 1.0f);
    float normalDiff3 = fabsf(v3->nx * v1->nx + v3->ny * v1->ny + v3->nz * v1->nz - 1.0f);
    
    // Adjust importance based on curvature (normal difference)
    float curvatureImportance1 = normalDiff1 * 2.0f;
    float curvatureImportance2 = normalDiff2 * 2.0f;
    float curvatureImportance3 = normalDiff3 * 2.0f;
    
    // Apply curvature importance
    if (curvatureImportance1 > v1->importance) v1->importance = curvatureImportance1;
    if (curvatureImportance2 > v2->importance) v2->importance = curvatureImportance2;
    if (curvatureImportance3 > v3->importance) v3->importance = curvatureImportance3;
  }
}

/**
 * Simplify the mesh to the target vertex count while preserving facial features
 */
EMSCRIPTEN_KEEPALIVE
float* simplifyMesh(int targetVertexCount, int* resultVertexCount, int* resultTriangleCount) {
  if (currentMesh == NULL || targetVertexCount >= currentMesh->vertexCount) {
    // Nothing to simplify
    *resultVertexCount = currentMesh->vertexCount;
    *resultTriangleCount = currentMesh->triangleCount;
    return NULL;
  }
  
  // Ensure we have importance values calculated
  calculateVertexImportance();
  
  // Create array of vertex indices
  int* vertexIndices = (int*)malloc(currentMesh->vertexCount * sizeof(int));
  for (int i = 0; i < currentMesh->vertexCount; i++) {
    vertexIndices[i] = i;
  }
  
  // Sort vertices by importance (using bubble sort for simplicity in this example)
  // In a real implementation, use a more efficient sorting algorithm
  for (int i = 0; i < currentMesh->vertexCount - 1; i++) {
    for (int j = 0; j < currentMesh->vertexCount - i - 1; j++) {
      if (currentMesh->vertices[vertexIndices[j]].importance < 
          currentMesh->vertices[vertexIndices[j + 1]].importance) {
        // Swap
        int temp = vertexIndices[j];
        vertexIndices[j] = vertexIndices[j + 1];
        vertexIndices[j + 1] = temp;
      }
    }
  }
  
  // Select vertices to keep
  int* vertexMap = (int*)malloc(currentMesh->vertexCount * sizeof(int));
  for (int i = 0; i < currentMesh->vertexCount; i++) {
    vertexMap[i] = -1;  // -1 means vertex is removed
  }
  
  // Keep the most important vertices
  int keptVertexCount = 0;
  for (int i = 0; i < targetVertexCount && i < currentMesh->vertexCount; i++) {
    int originalIndex = vertexIndices[i];
    vertexMap[originalIndex] = keptVertexCount++;
  }
  
  // Count triangles that will be kept (those with all vertices preserved)
  int keptTriangleCount = 0;
  for (int i = 0; i < currentMesh->triangleCount; i++) {
    Triangle* tri = &currentMesh->triangles[i];
    if (vertexMap[tri->v1] != -1 && vertexMap[tri->v2] != -1 && vertexMap[tri->v3] != -1) {
      keptTriangleCount++;
    }
  }
  
  // Allocate result array for simplified mesh data
  // Format: [vertexCount, triangleCount, v1x, v1y, v1z, v1nx, v1ny, v1nz, ..., t1v1, t1v2, t1v3, ...]
  float* result = (float*)malloc((2 + keptVertexCount * 6 + keptTriangleCount * 3) * sizeof(float));
  
  // Write vertex and triangle counts
  result[0] = (float)keptVertexCount;
  result[1] = (float)keptTriangleCount;
  
  // Copy kept vertices to result array
  int resultOffset = 2;
  for (int i = 0; i < currentMesh->vertexCount; i++) {
    if (vertexMap[i] != -1) {
      Vertex* v = &currentMesh->vertices[i];
      result[resultOffset++] = v->x;
      result[resultOffset++] = v->y;
      result[resultOffset++] = v->z;
      result[resultOffset++] = v->nx;
      result[resultOffset++] = v->ny;
      result[resultOffset++] = v->nz;
    }
  }
  
  // Copy kept triangles to result array, remapping vertex indices
  for (int i = 0; i < currentMesh->triangleCount; i++) {
    Triangle* tri = &currentMesh->triangles[i];
    if (vertexMap[tri->v1] != -1 && vertexMap[tri->v2] != -1 && vertexMap[tri->v3] != -1) {
      result[resultOffset++] = (float)vertexMap[tri->v1];
      result[resultOffset++] = (float)vertexMap[tri->v2];
      result[resultOffset++] = (float)vertexMap[tri->v3];
    }
  }
  
  // Clean up
  free(vertexIndices);
  free(vertexMap);
  
  // Return counts through pointer parameters
  *resultVertexCount = keptVertexCount;
  *resultTriangleCount = keptTriangleCount;
  
  return result;
}

/**
 * Free any allocated mesh data
 */
EMSCRIPTEN_KEEPALIVE
void freeMesh() {
  if (currentMesh != NULL) {
    if (currentMesh->vertices != NULL) {
      free(currentMesh->vertices);
    }
    if (currentMesh->triangles != NULL) {
      free(currentMesh->triangles);
    }
    if (currentMesh->landmarks != NULL) {
      free(currentMesh->landmarks);
    }
    free(currentMesh);
    currentMesh = NULL;
  }
}

/**
 * Clean up when module is unloaded
 */
EMSCRIPTEN_KEEPALIVE
void cleanUp() {
  freeMesh();
}