#!/usr/bin/env python
#
# math_stl
#    compile and evaluate .math to STL
#
# todo
#    run-length code
#    prune triangulation
#
# Neil Gershenfeld
# CBA MIT 10/29/11
#
# (c) Massachusetts Institute of Technology 2011
# Permission granted for experimental and personal use;
# license for commercial sale available from MIT.
#
#

import sys
import string
import Image
import os
from numpy import *

#
# command line args
#

if (not((len(sys.argv) == 3) | (len(sys.argv) == 4))):
   print "command line: math_stl in.math out.stl [resolution]"
   print "   in.math = input math string file"
   print "   out.stl = output STL image"
   print "   resolution = pixels per mm (optional, default 1)"
   sys.exit()

resolution = "1"
if (len(sys.argv) == 4):
   resolution = sys.argv[3]

#
# read .math file
#

input_file_name = sys.argv[1]
output_file_name = sys.argv[2]
input_file = open(input_file_name,'r')
(file_type,) = string.split(input_file.readline())
(units,) = string.split(input_file.readline())
units = float(units)
(dx,dy,dz) = string.split(input_file.readline())
(xmin,ymin,zmin) = string.split(input_file.readline())
math_string = input_file.readline()
input_file.close()
print "read "+input_file_name
print "   math string type: "+file_type
print "  ",units,"mm per unit"
print "   dx: "+dx+", dy: "+dy+", dz: "+dz
print "   xmin: "+xmin+", ymin: "+ymin+", zmin: "+zmin

#
# write evaluation program
#

program = """//
// math_stl_eval.c
//    math string to STL evaluation generated by math_stl
//
// Neil Gershenfeld
// CBA MIT 10/29/11
//
// (c) Massachusetts Institute of Technology 2011
// Permission granted for experimental and personal use;
// license for commercial sale available from MIT.
//

#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <stdint.h>
#include <png.h>

//
// data structures
//

struct fab_vars {
   unsigned char empty;
   unsigned char interior;
   unsigned char edge;
   unsigned char north;
   unsigned char south;
   unsigned char east;
   unsigned char west;
   unsigned char stop;
   unsigned char corner;
   unsigned char corner2;
   unsigned char direction;
   unsigned int nx,ny,nz;
   unsigned int bit_depth;
   unsigned int word_size;
   double dx,dy,dz;
   double xmin,ymin,zmin;
   unsigned char **array;
   uint32_t **rgb0,**rgb1;
   short int **xptr, **yptr;
   struct fab_path_type *path;
   struct fab_mesh_type *mesh;
   png_bytep *row_pointers;
   png_structp png_ptr;
   png_infop info_ptr;
   };

struct fab_vars init_vars() {
   //
   // fab_vars constructor
   //
   struct fab_vars v;
   //
   v.empty = 0;
   v.interior = 1;
   v.edge = (1 << 1);
   v.north = (1 << 2);
   v.west = (2 << 2);
   v.east = (3 << 2);
   v.south = (4 << 2);
   v.stop = (5 << 2);
   v.corner = (6 << 2);
   v.corner2 = (7 << 2);
   v.direction = (7 << 2);
   //
   return v;
   }

main(int argc, char **argv) {
   //
   // local vars
   //
   struct fab_vars v = init_vars(&v);
   double units = 25.4;
   double X,Y,Z;
   double dx,dy,dz;
   v.dx = units*"""+dx+""";
   v.dy = units*"""+dy+""";
   v.dz = units*"""+dz+""";
   v.xmin = units*"""+xmin+""";
   v.ymin = units*"""+ymin+""";
   v.zmin = units*"""+zmin+""";
   double resolution="""+resolution+""";
   v.nx = resolution*v.dx;
   v.ny = resolution*v.dy;
   v.nz = resolution*v.dz;
   int x,y,z;
   uint32_t val;
   //
   // create RGB arrays
   //
   v.rgb0 = malloc(v.ny*sizeof(uint32_t *));
   for (y = 0; y < v.ny; ++y)
      v.rgb0[y] = (uint32_t*) malloc(v.nx*sizeof(uint32_t));
   v.rgb1 = malloc(v.ny*sizeof(uint32_t *));
   for (y = 0; y < v.ny; ++y)
      v.rgb1[y] = (uint32_t*) malloc(v.nx*sizeof(uint32_t));
   //
   // create triangle array
   //
   uint32_t triangle_count=0;
   struct triangle_type {
      int v0x,v0y,v0z;
      int v1x,v1y,v1z;
      int v2x,v2y,v2z;
      struct triangle_type *next;
      };
   struct triangle_type *t,*tstart,*tnew;
   t = malloc(sizeof(struct triangle_type));
   tstart = t;
   //
   // loop over layers
   //
   for (z = 0; z < v.nz; ++z) {
      Z = (v.zmin + v.dz*z/(v.nz-1.0))/units;
      printf("   Z = %f\\n",Z);
      //
      // evaluate new layer
      //
      for (y = 0; y < v.ny; ++y) {
         Y = (v.ymin + v.dy*(v.ny-y-1)/(v.ny-1.0))/units;
         for (x = 0; x < v.nx; ++x) {
            X = (v.xmin + v.dx*x/(v.nx-1.0))/units;
            v.rgb0[y][x] = v.rgb1[y][x];
            v.rgb1[y][x] = """+math_string+""";
            }
         }
      //
      // check voxel faces
      //
      for (y = 0; y < (v.ny-1); ++y) {
         for (x = 0; x < (v.nx-1); ++x) {
            //
            // left
            //
            if ((v.rgb0[y][x] != 0) && (v.rgb0[y+1][x] != 0) && (v.rgb1[y][x] != 0) && (v.rgb1[y+1][x] != 0)) {
               if ((v.rgb0[y][x+1] == 0) || (v.rgb0[y+1][x+1] == 0) || (v.rgb1[y][x+1] == 0) || (v.rgb1[y+1][x+1] == 0)) {
                  triangle_count += 1;
                  (*t).v0x = x;
                  (*t).v0y = y;
                  (*t).v0z = z;
                  (*t).v1x = x;
                  (*t).v1y = y;
                  (*t).v1z = z+1;
                  (*t).v2x = x;
                  (*t).v2y = y+1;
                  (*t).v2z = z+1;
                  tnew = malloc(sizeof(struct triangle_type));
                  (*t).next = tnew;
                  t = tnew;
                  triangle_count += 1;
                  (*t).v0x = x;
                  (*t).v0y = y;
                  (*t).v0z = z;
                  (*t).v1x = x;
                  (*t).v1y = y+1;
                  (*t).v1z = z+1;
                  (*t).v2x = x;
                  (*t).v2y = y+1;
                  (*t).v2z = z;
                  tnew = malloc(sizeof(struct triangle_type));
                  (*t).next = tnew;
                  t = tnew;
                  }
               }
            //
            // right 
            //
            if ((v.rgb0[y][x+1] != 0) && (v.rgb0[y+1][x+1] != 0) && (v.rgb1[y][x+1] != 0) && (v.rgb1[y+1][x+1] != 0)) {
               if ((v.rgb0[y][x] == 0) || (v.rgb0[y+1][x] == 0) || (v.rgb1[y][x] == 0) || (v.rgb1[y+1][x] == 0)) {
                  triangle_count += 1;
                  (*t).v0x = x+1;
                  (*t).v0y = y;
                  (*t).v0z = z;
                  (*t).v1x = x+1;
                  (*t).v1y = y+1;
                  (*t).v1z = z+1;
                  (*t).v2x = x+1;
                  (*t).v2y = y;
                  (*t).v2z = z+1;
                  tnew = malloc(sizeof(struct triangle_type));
                  (*t).next = tnew;
                  t = tnew;
                  triangle_count += 1;
                  (*t).v0x = x+1;
                  (*t).v0y = y;
                  (*t).v0z = z;
                  (*t).v1x = x+1;
                  (*t).v1y = y+1;
                  (*t).v1z = z;
                  (*t).v2x = x+1;
                  (*t).v2y = y+1;
                  (*t).v2z = z+1;
                  tnew = malloc(sizeof(struct triangle_type));
                  (*t).next = tnew;
                  t = tnew;
                  }
               }
            //
            // front
            //
            if ((v.rgb0[y][x] != 0) && (v.rgb0[y][x+1] != 0) && (v.rgb1[y][x] != 0) && (v.rgb1[y][x+1] != 0)) {
               if ((v.rgb0[y+1][x] == 0) || (v.rgb0[y+1][x+1] == 0) || (v.rgb1[y+1][x] == 0) || (v.rgb1[y+1][x+1] == 0)) {
                  triangle_count += 1;
                  (*t).v0x = x;
                  (*t).v0y = y;
                  (*t).v0z = z;
                  (*t).v1x = x+1;
                  (*t).v1y = y;
                  (*t).v1z = z;
                  (*t).v2x = x+1;
                  (*t).v2y = y;
                  (*t).v2z = z+1;
                  tnew = malloc(sizeof(struct triangle_type));
                  (*t).next = tnew;
                  t = tnew;
                  triangle_count += 1;
                  (*t).v0x = x;
                  (*t).v0y = y;
                  (*t).v0z = z;
                  (*t).v1x = x+1;
                  (*t).v1y = y;
                  (*t).v1z = z+1;
                  (*t).v2x = x;
                  (*t).v2y = y;
                  (*t).v2z = z+1;
                  tnew = malloc(sizeof(struct triangle_type));
                  (*t).next = tnew;
                  t = tnew;
                  }
               }
            //
            // back
            //
            if ((v.rgb0[y+1][x] != 0) && (v.rgb0[y+1][x+1] != 0) && (v.rgb1[y+1][x] != 0) && (v.rgb1[y+1][x+1] != 0)) {
               if ((v.rgb0[y][x] == 0) || (v.rgb0[y][x+1] == 0) || (v.rgb1[y][x] == 0) || (v.rgb1[y][x+1] == 0)) {
                  triangle_count += 1;
                  (*t).v0x = x;
                  (*t).v0y = y+1;
                  (*t).v0z = z;
                  (*t).v1x = x+1;
                  (*t).v1y = y+1;
                  (*t).v1z = z;
                  (*t).v2x = x+1;
                  (*t).v2y = y+1;
                  (*t).v2z = z+1;
                  tnew = malloc(sizeof(struct triangle_type));
                  (*t).next = tnew;
                  t = tnew;
                  triangle_count += 1;
                  (*t).v0x = x;
                  (*t).v0y = y+1;
                  (*t).v0z = z;
                  (*t).v1x = x+1;
                  (*t).v1y = y+1;
                  (*t).v1z = z+1;
                  (*t).v2x = x;
                  (*t).v2y = y+1;
                  (*t).v2z = z+1;
                  tnew = malloc(sizeof(struct triangle_type));
                  (*t).next = tnew;
                  t = tnew;
                  }
               }
            //
            // top 
            //
            if ((v.rgb0[y][x] != 0) && (v.rgb0[y][x+1] != 0) && (v.rgb0[y+1][x] != 0) && (v.rgb0[y+1][x+1] != 0)) {
               if ((v.rgb1[y][x] == 0) || (v.rgb1[y][x+1] == 0) || (v.rgb1[y+1][x] == 0) || (v.rgb1[y+1][x+1] == 0)) {
                  triangle_count += 1;
                  (*t).v0x = x;
                  (*t).v0y = y;
                  (*t).v0z = z;
                  (*t).v1x = x+1;
                  (*t).v1y = y+1;
                  (*t).v1z = z;
                  (*t).v2x = x;
                  (*t).v2y = y+1;
                  (*t).v2z = z;
                  tnew = malloc(sizeof(struct triangle_type));
                  (*t).next = tnew;
                  t = tnew;
                  triangle_count += 1;
                  (*t).v0x = x;
                  (*t).v0y = y;
                  (*t).v0z = z;
                  (*t).v1x = x+1;
                  (*t).v1y = y;
                  (*t).v1z = z;
                  (*t).v2x = x+1;
                  (*t).v2y = y+1;
                  (*t).v2z = z;
                  tnew = malloc(sizeof(struct triangle_type));
                  (*t).next = tnew;
                  t = tnew;
                  }
               }
            //
            // bottom 
            //
            if ((v.rgb1[y][x] != 0) && (v.rgb1[y][x+1] != 0) && (v.rgb1[y+1][x] != 0) && (v.rgb1[y+1][x+1] != 0)) {
               if ((v.rgb0[y][x] == 0) || (v.rgb0[y][x+1] == 0) || (v.rgb0[y+1][x] == 0) || (v.rgb0[y+1][x+1] == 0)) {
                  triangle_count += 1;
                  (*t).v0x = x;
                  (*t).v0y = y;
                  (*t).v0z = z+1;
                  (*t).v1x = x+1;
                  (*t).v1y = y+1;
                  (*t).v1z = z+1;
                  (*t).v2x = x+1;
                  (*t).v2y = y;
                  (*t).v2z = z+1;
                  tnew = malloc(sizeof(struct triangle_type));
                  (*t).next = tnew;
                  t = tnew;
                  triangle_count += 1;
                  (*t).v0x = x;
                  (*t).v0y = y;
                  (*t).v0z = z+1;
                  (*t).v1x = x;
                  (*t).v1y = y+1;
                  (*t).v1z = z+1;
                  (*t).v2x = x+1;
                  (*t).v2y = y+1;
                  (*t).v2z = z+1;
                  tnew = malloc(sizeof(struct triangle_type));
                  (*t).next = tnew;
                  t = tnew;
                  }
               }
            }
         }
      }
   //
   // write STL
   //
   char header[80];
   uint16_t attribute;
   int i;
   t = tstart;
   float v0x,v0y,v0z;
   float v1x,v1y,v1z;
   float v2x,v2y,v2z;
   FILE *output_file;
   output_file = fopen(\""""+output_file_name+"""\","wb");
   fwrite(header,1,80,output_file);
   fwrite(&triangle_count,4,1,output_file);
   float zero = 0;
   for (i = 0; i < triangle_count; ++i) {
      fwrite(&zero,4,1,output_file);
      fwrite(&zero,4,1,output_file);
      fwrite(&zero,4,1,output_file);
      v0x = (v.xmin + v.dx*((*t).v0x)/(v.nx-1.0))/units;
      fwrite(&v0x,4,1,output_file);
      v0y = (v.ymin + v.dy*((*t).v0y)/(v.ny-1.0))/units;
      fwrite(&v0y,4,1,output_file);
      v0z = (v.zmin + v.dz*((*t).v0z)/(v.nz-1.0))/units;
      fwrite(&v0z,4,1,output_file);
      v1x = (v.xmin + v.dx*((*t).v1x)/(v.nx-1.0))/units;
      fwrite(&v1x,4,1,output_file);
      v1y = (v.ymin + v.dy*((*t).v1y)/(v.ny-1.0))/units;
      fwrite(&v1y,4,1,output_file);
      v1z = (v.zmin + v.dz*((*t).v1z)/(v.nz-1.0))/units;
      fwrite(&v1z,4,1,output_file);
      v2x = (v.xmin + v.dx*((*t).v2x)/(v.nx-1.0))/units;
      fwrite(&v2x,4,1,output_file);
      v2y = (v.ymin + v.dy*((*t).v2y)/(v.ny-1.0))/units;
      fwrite(&v2y,4,1,output_file);
      v2z = (v.zmin + v.dz*((*t).v2z)/(v.nz-1.0))/units;
      fwrite(&v2z,4,1,output_file);
      fwrite(&attribute,2,1,output_file);
      t = (*t).next;
      }
   fclose(output_file);
   }
"""

eval_file = open("math_stl_eval.c",'w')
eval_file.write(program)
eval_file.close()

#
# compile
#

print "compile"
if os.uname()[0] == 'Darwin':
	os.system("gcc -O math_stl_eval.c -o math_stl_eval -lm -I/usr/X11/include -L/usr/X11/lib -lpng")
else:
	os.system("gcc -O math_stl_eval.c -o math_stl_eval -lm -lpng")


#
# execute
#

print "execute"
os.system("./math_stl_eval")
os.system("rm ./math_stl_eval.c")
os.system("rm ./math_stl_eval")