#!/usr/bin/env slsh
%
% Copyright (C) 2007 John E. Davis
%
% fztopng is free software; you can redistribute it and/or modify it
% under the terms of the GNU General Public License as published by
% the Free Software Foundation; either version 2 of the License, or
% (at your option) any later version.
%
% fztopng is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
% General Public License for more details.
%
require ("png");
require ("cmdopt");

public variable RGB_METHOD_HSV = 1;
public variable RGB_METHOD_GRADIENT = 2;
public variable RGB_METHOD_XOR = 3;

public variable RGB_Method = RGB_METHOD_HSV;
public variable Min_Saturation = 0.3;
public variable Max_Saturation = 1.0;
public variable Min_Value = 0.6;
public variable Max_Value = 1.0;
public variable Min_Hue = 0.0;
public variable Max_Hue = 360.0;

private variable Version = "0.1.0";

define make_complex_image (zreal, zimag)
{
   variable n = length (zreal);
   variable m = length (zimag);
   variable z = Complex_Type[m,n];
   
   zimag *= 1i;
   variable i;
   _for i (0, n-1, 1)
     z[*,i] = zimag;
   
   _for i (0, m-1, 1)
     z[i,*] += zreal;

   return z;
}

private define arg(z)
{
   return atan2 (Imag(z), Real(__tmp(z)));    %  runs from -PI to PI
}

% returns a value 0->360
private define hue (z)
{
   return ((arg(__tmp(z)) * (180.0/PI)) + 360.0) mod 360.0;
}

% returns a value 0->1
private define saturation (z)
{
   return (2/PI)*atan(abs(2.0*z));
}

private define value (z)
{
   return (1.0 - (2/PI)*atan(abs(z/8.0)));
}

private define z_to_hsv (z)
{
   %return hue(z), 

   variable h, s, v;

   h = hue (z);
   % model 1
   if (1)
     {
	variable r = abs(__tmp(z));
	r = log(1.0+__tmp(r));
	v = ((1.0+sin((2*PI)*r))/2.0);
	s = ((1.0+cos((2*PI)*r))/2.0);
     }
   else
     {
	s = saturation(z);
	v = value(z);
     }

   variable sat_min = Min_Saturation, sat_max = Max_Saturation;
   variable val_min = Min_Value, val_max = Max_Value;

   v = val_min + (val_max - val_min) * __tmp(v);
   s = sat_min + (sat_max - sat_min) * __tmp(s);

   % Min_Hue could be negative.  So do it this way
   h = ((Min_Hue + 360.0) + (Max_Hue-Min_Hue) * __tmp(h)/360.0) mod 360.0;
   return h, s, v;
}
   
private define rgb_to_rgb (r, g, b)
{
   return ((r & 0xFF) shl 16) + ((g & 0xFF) shl 8) + (b & 0xFF);
}

private define uchar (x)
{
   return typecast (x, UChar_Type);
}

private define z_to_rgb_via_hsv (z)
{
   variable h, s, v;
   (h, s, v) = z_to_hsv (__tmp(z));

   h = __tmp(h)/60.0;
   v = 255.99999 * __tmp(v);

   variable i = int(floor(h)) mod 6;
   variable f = h - i;
   variable p = uchar (v * (1.0 - s));
   variable q = uchar (v * (1.0 - (s * f)));
   variable t = uchar (v * (1.0 - (s * (1.0 - f))));
   v = uchar (v);

   variable r = @p;
   variable g = @p;
   variable b = @p;

   variable j = where (i == 0);
   r[j] = v[j]; g[j] = t[j]; b[j] = p[j];
   
   j = where (i == 1);
   r[j] = q[j]; g[j] = v[j]; b[j] = p[j];

   j = where (i == 2);
   r[j] = p[j]; g[j] = v[j]; b[j] = t[j];
   
   j = where (i == 3);
   r[j] = p[j]; g[j] = q[j]; b[j] = v[j];
   
   j = where (i == 4);
   r[j] = t[j]; g[j] = p[j]; b[j] = v[j];

   j = where (i == 5);
   r[j] = v[j]; g[j] = p[j]; b[j] = q[j];

   return ((r & 0xFF) shl 16) + ((g & 0xFF) shl 8) + (b & 0xFF);
}

private define z_to_rgb_via_grad (z, a0, a1)
{
   variable c0, c1;
   variable t = hue(__tmp(z))/360.0;
   variable mask = 0xFF0000, n = 16;
   variable c = 0;
   loop (3)
     {
	c0 = ((a0 & mask) shr n);
	c1 = ((a1 & mask) shr n);
	c |= (uchar ((c1-c0)*t + c0) shl n);
	n -= 8;
	mask = mask shr 8;
     }
   return c;
}

private variable f_of_z_func;
private define iterate (z, n)
{
   loop (n)
     z = (@f_of_z_func) (__tmp(z));
   return z;
}

private define exit_usage ()
{
   () = fprintf (stderr, "Version %S Usage: %s [options] -f func | func.sl\n", Version, __argv[0]);
   foreach 
     (["Options:\n",
       " -x|--real xlo:xhi:#N        Grid for x=Real(z) (default: -2:2:#512)\n",
       " -y|--imag ylo:yhi:#N        Grid for y=Imag(z) (default: -2:2:#512)\n",
       " --iter N                    Iterate N type (default=1)\n",
       " --hsv                       Use HSV method (default)\n",
       " --hue=float,float           Min/Max values of Hue ([0->360])\n",
       " --sat=float,float           Min/Max values of Saturation ([0->1])\n",
       " --val=float,float           Min/Max values of Value ([0->1])\n",
       " --grad                      Use gradient method\n",
       " --rgb=int,int               Min/Max RGB values for gradient method\n",
       " -o file.png                 Output file (default: func.png)\n",
       " -h|--help                   This message\n",
       "\n",
       "Note: The file func.sl must contain a function called f_of_z.\n",
       "If the -f expr form is used, it must involve z, e.g., -f sin(z)\n"
      ])
     {
	variable line = ();
	() =fputs (line, stderr);
     }
   exit (1);
}

define slsh_main ()
{
   variable prefs = struct
     {
	xgrid = [-1.5:1.5:#512],
	ygrid = [-1.5:1.5:#512], 
	method = RGB_METHOD_HSV,
	iter = 0,
	hue = [0.0, 360.0],
	saturation = [0.3, 1.0],
	value = [0.6, 1.0],
	rgb = [0x4C4CFF, 0xFFFF00],
     };

   variable xgrid = NULL, ygrid = NULL, method = NULL, iter = NULL,
     rgb = NULL, sat = NULL, val = NULL, hue = NULL;

   variable outfile = NULL;
   variable func = NULL, func_file = NULL;
   variable c = cmdopt_new ();
   c.add ("x|real", &xgrid; type="str");
   c.add ("y|imag", &ygrid; type="str");
   c.add ("iter", &iter; type="int");
   c.add ("xor", &method; default=RGB_METHOD_XOR);
   c.add ("hsv", &method; default=RGB_METHOD_HSV);
   c.add ("grad", &method; default=RGB_METHOD_GRADIENT);
   c.add ("rgb", &rgb; type="str");
   c.add ("sat", &sat; type="str");
   c.add ("val", &val; type="str");
   c.add ("hue", &hue; type="str");
   c.add ("o", &outfile; type="str");
   c.add ("f", &func; type="str");
   variable i = c.process (__argv, 1);
   
   if (i == __argc)
     {
	if (func == NULL)
	  exit_usage ();
     }
   else if (i + 1 != __argc)
     exit_usage ();
   else func_file = __argv[i];

   if (outfile == NULL)
     {
	if (func_file == NULL)
	  {
	     () = fprintf (stderr, "No output png file specified\n");
	     exit (1);
	  }
	outfile = path_basename_sans_extname (func_file) + ".png";
     }
   
   if (func_file != NULL)
     {
	ifnot (path_is_absolute (func_file))
	  func_file = path_concat (getcwd (), func_file);
	if (NULL == stat_file (func_file))
	  {
	     () = fprintf (stderr, "Function file %s does not exist\n", func_file);
	     exit (1);
	  }
	try
	  {
	     () = evalfile (func_file);
	  }
	catch AnyError: 
	  {
	     () = fprintf (stderr, "An error occured evaluating %s:\n\n", func_file);
	     throw;
	  }
     }
   if (func != NULL)
     {
	try
	  {
	     eval ("define f_of_z(z) { $func; }"$);
	  }
	catch AnyError:
	  {
	     () = fprintf (stderr, "An error occured evaluating %s:\n\n", func);
	     throw;
	  }
     }

   f_of_z_func = __get_reference ("f_of_z");
   if (f_of_z_func == NULL)
     {
	() = fprintf (stderr, "%s did not define a public function called f_of_z\n", func_file);
	exit (1);
     }

   variable set_prefs = __get_reference ("set_options");
   if (set_prefs != NULL)
     (@set_prefs)(prefs);

   if (xgrid == NULL) xgrid = prefs.xgrid;
   if (ygrid == NULL) ygrid = prefs.ygrid;
   if (rgb == NULL) rgb = prefs.rgb;
   if (hue == NULL) hue = prefs.hue;
   if (sat == NULL) sat = prefs.saturation;
   if (val == NULL) val = prefs.value;
   if (method == NULL) method = prefs.method;
   if (iter == NULL) iter = prefs.iter;

   try
     {
	if (Array_Type != typeof(xgrid)) xgrid = eval ("[$xgrid]"$);
	if (Array_Type != typeof(ygrid)) ygrid = eval ("[$ygrid]"$);
	if (Array_Type != typeof(rgb)) rgb = eval ("[$rgb]"$);
	if (Array_Type != typeof(hue)) hue = eval ("[$hue]"$);
	if (Array_Type != typeof(sat)) sat = eval ("[$sat]"$);
	if (Array_Type != typeof(val)) val = eval ("[$val]"$);
     }
   catch AnyError:
     {
	() = fprintf (stderr, "Unable to expressions\n");
	exit (1);
     }

   try
     {
	Min_Saturation = sat[0];
	Max_Saturation = sat[-1];
	Min_Value = val[0];
	Max_Value = val[-1];
	Min_Hue = hue[0];
	Max_Hue = hue[-1];
	if ((_typeof (rgb) != Int_Type) || (length(rgb) != 2))
	  {
	     () = fprintf (stderr, "Expecting 2 element integer arrays for the rgb option\n");
	     exit (1);
	  }
     }
   catch AnyError:
     {
	() = fprintf (stderr, "Expecting 2 element arrays for hue, sat, val, and rgb options\n");
	exit (1);
     }

   variable z = make_complex_image (xgrid, ygrid);
   z = iterate (__tmp(z), iter+1);

   if (method == RGB_METHOD_XOR)
     {
	rgb = (z_to_rgb_via_grad (z, rgb[0], rgb[1]) xor z_to_rgb_via_hsv(__tmp(z)));
     }
   else if (method == RGB_METHOD_HSV)
     rgb = z_to_rgb_via_hsv(__tmp(z));
   else
     rgb = z_to_rgb_via_grad (z, rgb[0], rgb[1]);

   png_write_flipped (outfile, rgb);
}