#!/usr/bin/env python3
#
# frep-C.py
#    functional representation to C++ multi-threaded solver
#
# usage:
#    pcb.py | frep-C.py [dpi [filename]]
#
# Neil Gershenfeld 11/26/22
#
# This work may be reproduced, modified, distributed,
# performed, and displayed for any purpose, but must
# acknowledge this project. Copyright is retained and
# must be preserved. The work is provided as is; no
# warranty is provided, and users accept all liability.
#

#
# import
#
import json,sys,os
from numpy import *
from PIL import Image
#
# read input
#
frep = json.load(sys.stdin)
#
# check arguments
#
if (frep['type'] != 'RGB'):
   print('types other than RGB not (yet) supported')
   sys.exit()
if (len(sys.argv) == 1):
   print('output to out.png at 300 DPI')
   filename = 'out.png'
   dpi = 300
elif (len(sys.argv) == 2):
   dpi = sys.argv[1]
   filename = 'out.png'
   print('output to out.png at '+dpi+'DPI')
   dpi = int(dpi)
elif (len(sys.argv) == 3):
   dpi = sys.argv[1]
   filename = sys.argv[2]
   print('output to '+filename+' at '+dpi+' DPI')
   dpi = int(dpi)
#
# define variables
#
xmin = frep['xmin']
xmax = frep['xmax']
ymin = frep['ymin']
ymax = frep['ymax']
zmin = min(frep['layers'])
zmax = max(frep['layers'])
units = float(frep['mm_per_unit'])
delta = (25.4/dpi)/units
fn = frep['function']
fn = fn.replace('math.','')
layers = str(frep['layers'][0])
nlayers = len(frep['layers'])
for i in range(1,nlayers):
   layers += ','+str(frep['layers'][i])
#
# write program
#
file = open("frep-C.cpp",'w')
file.write(
f"""
#include <iostream>
#include <cmath>
#include <thread>
#include <png.h>
//
using namespace std;
//
float xmin = {xmin};
float xmax = {xmax};
float ymin = {ymin};
float ymax = {ymax};
float zmin = {zmin};
float zmax = {zmax};
float units = {units};
int dpi = {dpi};
float delta = (25.4/dpi)/units;
int nx = (xmax-xmin)/delta;
int ny = (ymax-ymin)/delta;
int *m = (int*) calloc(nx*ny,sizeof(int));
float layers[] = {{{layers}}};
int nlayers = {nlayers};
int nthreads = std::thread::hardware_concurrency();
//
int fn(float X,float Y,float Z) {{
   return ({fn});
   }}
//
void calc(int nx,int ny,int nthreads,int thread) {{
   int intensity;
   for (int layer = 0; layer < nlayers; ++layer) {{
      float z = layers[layer];
      if (thread == 0)
         cout << "   z = " << z << endl;
      if (zmin == zmax)
         intensity = 0xffffff;
      else
         intensity = ((int) (255*(z-zmin)/(zmax-zmin))) | (255 << 8) | (255 << 16);
      int iystart = thread*ny/nthreads;
      int iyend = (thread+1)*ny/nthreads;
      for (int iy = iystart; iy < iyend; ++iy) {{
         float y = ymin+iy*delta;
         for (int ix = 0; ix < nx; ++ix) {{
            float x = xmin+ix*delta;
            m[iy*nx+ix] += intensity & fn(x,y,z);
            }}
         }}
      }}
   }}
//
int main(int argc, char** argv) {{
   cout << "   calculate " << nx << "x" << ny << " with " << nthreads << " threads" << endl;
   std::thread threads[nthreads];
   for (int i = 0; i < nthreads; ++i)
      threads[i] = std::thread(calc,nx,ny,nthreads,i);
   for (int i = 0; i < nthreads; ++i)
      threads[i].join();
   //
   FILE *file;
   file = fopen("{filename}","wb");
	png_structp pngfile;
	png_infop pnginfo;
	png_bytep pngrow;
	pngfile = png_create_write_struct(PNG_LIBPNG_VER_STRING,NULL,NULL,NULL);
	pnginfo = png_create_info_struct(pngfile);
	png_init_io(pngfile,file);
	png_set_IHDR(pngfile,pnginfo,nx,ny,
	   8,PNG_COLOR_TYPE_RGBA,PNG_INTERLACE_NONE,
	   PNG_COMPRESSION_TYPE_BASE,PNG_FILTER_TYPE_BASE);
   png_set_pHYs(pngfile,pnginfo,(png_uint_32) {1000*dpi/25.4},(png_uint_32) {1000*dpi/25.4},PNG_RESOLUTION_METER);
	png_write_info(pngfile,pnginfo);
	pngrow = (png_bytep) malloc(4*nx*sizeof(png_byte));
	for (int iy = (ny-1); iy >= 0; --iy) {{
		for (int ix = 0; ix < nx; ++ix) {{
         pngrow[4*ix] = m[iy*nx+ix] & 255;
         pngrow[4*ix+1] = (m[iy*nx+ix] >> 8) & 255; 
         pngrow[4*ix+2] = (m[iy*nx+ix] >> 16) & 255;
         pngrow[4*ix+3] = 255;
         }}
		png_write_row(pngfile,pngrow);
	   }}
	png_write_end(pngfile,NULL);
	fclose(file);
   }}
""")
file.close()
#
# compile
#
print("compile ...")
#os.system("time g++ frep-C.cpp -o frep-C -lm -lpng -O -ffast-math -pthread")
os.system("time clang++ frep-C.cpp -o frep-C -lm -lpng -O -ffast-math -pthread -Wno-bool-operation")
#
# execute
#
print("execute ...")
os.system("time ./frep-C")
#
# clean up
#
os.system("rm -f frep-C.cpp frep-C")