#!/usr/bin/env isis-script
% This script wobble-corrects a suzaku attitude file.
%
%  Copyright (c) 2008 Massachusetts Institute of Technology
%
%  Author:  John E. Davis <davis@space.mit.edu>
%
%  Permission to use, copy, modify, distribute, and sell this software
%  and its documentation for any purpose is hereby granted without fee,
%  provided that the above copyright notice appear in all copies and
%  that both that copyright notice and this permission notice appear in
%  the supporting documentation, and that the name of the Massachusetts
%  Institute of Technology not be used in advertising or publicity
%  pertaining to distribution of the software without specific, written
%  prior permission.  The Massachusetts Institute of Technology makes
%  no representations about the suitability of this software for any
%  purpose.  It is provided "as is" without express or implied warranty.
%  
%  THE MASSACHUSETTS INSTITUTE OF TECHNOLOGY DISCLAIMS ALL WARRANTIES
%  WITH REGARD TO THIS SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF
%  MERCHANTABILITY AND FITNESS.  IN NO EVENT SHALL THE MASSACHUSETTS
%  INSTITUTE OF TECHNOLOGY BE LIABLE FOR ANY SPECIAL, INDIRECT OR
%  CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
%  OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
%  NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION
%  WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.  }}}

if (_slang_version < 20100)
{
   () = fputs ("This script requires at least version 2.1.0 of the S-Lang library.\n"
	       + "The latest version of the library may be found at\n"
	       + "<http://www.jedsoft.org/slang/>\n",
	       stderr);
   exit (1);
}

require ("wcsfuns");
require ("cmdopt");

private variable Script_Version_String = "0.2.0";

private variable Default_DT = 100.0;   %  secs
private variable Default_Radius = 150.0;
private variable Verbose = 1;
private variable Pixels_Per_ArcSecond = 250.0/260.0;

autoload ("ds9_new", "ds9object");

% Some GTI functions taken from my gtifuns package.

define gti_new (start, stop)
{
   variable a = struct
     {
	start, stop
     };
   a.start = start;
   a.stop = stop;
   return a;
}

define gti_filter_lc (time, rate, min_rate, max_rate)
{
   variable g = ((rate >= min_rate)
		 and (rate < max_rate));
   
   g = [g, 0];
   time = [time, time[-1] + (time[-1] - time[-2])];

   variable d = g - shift(g,-1);
   variable tstart = time[where (d == 1)];
   variable tstop = time[where (d == -1)];
   
   return tstart, tstop;
}
   
define gti_lc_compute_exposure ()
{
   if (_NARGS == 0)
     {
	usage ("dt = gti_lc_compute_exposure (gti, lc_times)");
     }
   
   variable gti, tgrid;
   (gti, tgrid) = ();
   
   variable tstart = gti.start;
   variable tstop = gti.stop;
   variable m = length (tstart);
   variable old_grid = Double_Type[2*m];
   old_grid[[0::2]] = tstart;
   old_grid[[1::2]] = tstop;
   variable vals = Double_Type[2*m];
   vals[[0::2]] = tstop-tstart;

   return rebin (tgrid, make_hi_grid(tgrid), 
                 old_grid, make_hi_grid(old_grid), __tmp(vals));
}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

private define debug ()
{
   variable args = __pop_args (_NARGS);
   variable vlevel = 1;
   if (Verbose >= vlevel)
     {
	() = fprintf (stderr, "\n");
	() = fprintf (stderr, __push_args (args));
	() = fflush (stderr);
     }
}

private define exit_version ()
{
   () = fprintf (stdout, "Version: %S\n", Script_Version_String);
   exit (0);
}

private define exit_usage ()
{
   variable fp = stderr;
   () = fprintf (fp, "Usage: %s [options] evtfile inattfile outattfile\n", __argv[0]);
   variable opts = 
     [
      "\n",
      "   evtfile: The name of the event file\n",
      " inattfile: The input attitude file.\n",
      "outattfile: The output attitude file.\n",
      "\n",
      "Options:\n",
      " --ds9              Display before/after PSF images using ds9\n",
      " --dt=val           Time bin size (secs) [$Default_DT]\n"$,
      " --cel=RA,DEC,R     Extraction region with R in arc-seconds\n",
      " --sky=X,Y,R        Extraction region in sky pixel coordinates\n",
      " --gtifile=val      File with GTIs (default is to use evtfile)\n",
      " --lcfile=FILE      Write light-curve info to this file\n",
      " --radius=val       Extraction radius (pixels) [$Default_Radius]\n"$,
      " -q, --quiet        Run quietly\n",
      " -v, --version      Print version\n",
      " -h, --help         This message\n",
      "\n",
      "Note: The extraction regions are circular and specified as a center and\n",
      "      radius.  If an extraction region (--cel or --sky) is not given, then\n",
      "      ds9 will be used to obtain one interactively.\n",
     ];
   foreach (opts)
     {
	variable opt = ();
	() = fputs (opt, fp);
     }
   exit (1);
}

private define parse_cel_position (str, is_dec)
{
   ifnot (is_substr (str, ":"))
     return atof (str);

   variable hr, min, sec;
   variable n = sscanf (str, "%g:%g:%g", &hr, &min, &sec);
   if (n != 3)
     {
	() = fprintf (stderr, "Expecting decimal or HH:MM:SS format for the celestial position\n");
	exit (1);
     }
   variable x;

   variable s = sign (hr);
   hr = abs(hr);
   if (is_dec)
     x = hr + (min + sec/60.0)/60.0;
   else
     x = (hr + (min + sec/60.0)/60.0)*360.0/24;

   return x * s;
}

private define parse_cel_extracton_region (str)
{
   str = strtrans (str, ",", " ");
   variable x, y, r;
   if (3 != sscanf (str, "%s %s %lf", &x, &y, &r))
     {
	() = fprintf (stderr, "Expecting 3 values for the extraction region parameter\n");
	exit (1);
     }
   x = parse_cel_position (x, 0);
   y = parse_cel_position (y, 1);
   r *= Pixels_Per_ArcSecond;
   
   return x, y, r;
}

private define parse_sky_extracton_region (str)
{
   str = strtrans (str, ",", " ");
   variable x, y, r;
   if (3 != sscanf (str, "%lf %lf %lf", &x, &y, &r))
     {
	() = fprintf (stderr, "Expecting 3 values for the extraction region parameter\n");
	exit (1);
     }
   return (x,y,r);
}

private define read_gti_table (file)
{
   variable fp;

   try
     {
	fp = fits_open_file (file + "[STDGTI]","r");
     }
   catch FitsError:
     {
	try
	  {
	     fp = fits_open_file (file + "[GTI]","r");
	  }
	catch FitsError:
	  {
	     () = fprintf (stderr, "Unable to open the GTI or STDGTI extension in %s\n", file);
	     exit (1);
	  }
     }

   variable tstop, tstart;
   (tstart, tstop) = fits_read_col (fp, "start", "stop");
   return gti_new (tstart, tstop);
}


private define compute_lc (evt, gtis, dt)
{
   variable tgrid = [min(evt.time):max(evt.time)+dt:dt];
   variable exposure = gti_lc_compute_exposure (gtis, tgrid);

   variable nbins = length (tgrid);
   variable revi;
   variable counts = histogram (evt.time, tgrid, make_hi_grid(tgrid), &revi);

   variable lc = struct
     {
	time_min = tgrid,
	time_max = [tgrid[[1:]], tgrid[-1]],
	exposure = exposure,
	counts = counts,
	x = Double_Type[nbins],
	y = Double_Type[nbins],
     };

   debug ("Computing the source centroid position along the lightcurve\r");

   variable i, j, x = evt.x, y = evt.y, lc_x = lc.x, lc_y = lc.y;
   _for i (0, nbins-1, 1)
     {
	j = revi[i];
	lc_x[i] = mean(x[j]);
	lc_y[i] = mean(y[j]);
     }
   i = where (exposure == 0);
   lc_x[i] = _NaN;
   lc_y[i] = _NaN;

   return lc;
}

private define check_fits_error (s)
{
   if (s == 0)
     {
	_fits_clear_errmsg ();
	return;
     }
   variable errmsg = _fits_get_errstatus (s);
   throw FitsError, strjoin ([fits_read_errmsgs (), errmsg], "\n");
}

private define write_att_file (in, out, asp, history)
{
   variable fpout = fits_open_file (out, "c");
   variable fpin = fits_open_file (in, "r");
   check_fits_error (_fits_copy_hdu (fpin, fpout, 0));
   check_fits_error (_fits_movrel_hdu (fpin, 1));
   check_fits_error (_fits_copy_hdu (fpin, fpout, 10+length(history)));
   fits_close_file (fpin);
   
   foreach (history)
     {
	variable h = ();
	check_fits_error (_fits_write_history (fpout, h));
     }

   variable col = fits_get_colnum (fpout, "EULER");
   check_fits_error (_fits_write_col (fpout, col, 1, 1, asp.euler));
   check_fits_error (_fits_write_chksum (fpout));
   fits_close_file (fpout);
}

private define open_ds9 ()
{
   variable e, ds9 = NULL; 
   try (e)
     {
	ds9 = ds9_new (sprintf ("ds9.%d-a", getpid()); verbose=5, timeout=90);
     }
   catch AnyError:
     {
	() = fprintf (stderr, "** Warning: Unable to open ds9: %S\n", e.message);
     }
   return ds9;
}

private define get_extraction_region_from_ds9 (ds9, evt, wcs)
{
   variable rl = slsh_readline_new ("");
   variable prompt = "Mark a SINGLE circular extraction region, then press RETURN";
   variable regs, i;

   forever
     {
	ds9.put_regions (NULL);
	() = slsh_readline (rl, prompt);
	regs = ds9.get_regions (;system="physical", format="ds9");
	i = where (0==array_map(Int_Type, &strncmp, regs, "circle(", 7));
	if (length (i) == 1)
	  {
	     variable x, y, r;
	     variable reg = regs[i[0]];
	     if (3 != sscanf (reg, "circle(%lf,%lf,%lf)", &x, &y, &r))
	       {
		  () = fprintf (stderr, "Failed to parse %S as a circle.\n", reg[i]);
		  print (regs);
		  exit (1);
	       }
	     return x, y, r;
	  }

	() = fprintf (stderr, "Found %d region circles -- expecting a SINGLE circle.  Try again\n",
		     length (i));
     }
}

private define display_events_in_ds9 (ds9, evt, wcs)
{
   variable xgrid = [int(floor(min(evt.x))):int(ceil(max(evt.x)))];
   variable ygrid = [int(floor(min(evt.y))):int(ceil(max(evt.y)))];
   variable img = histogram2d (evt.y, evt.x, ygrid, xgrid);
   variable img_wcs = fitswcs_bin_wcs (wcs, xgrid, ygrid);
   variable phys_wcs = fitswcs_new_img_wcs (xgrid, ygrid);
   phys_wcs.wcsname = "PHYSICAL";

   ds9.put_array (img);
   ds9.set_wcs (img_wcs);
   ds9.set_wcs (phys_wcs; alt='p', append);
}

private define correct_att_file (inattfile, outattfile, wcs, lc, lc_t, dx, dy)
{
   debug ("Reading the attitude file %s\r", inattfile);
   variable asp = fits_read_table (inattfile, "TIME", "EULER");
   variable ok_times = where (lc.time_min[0] <= asp.time <= lc.time_max[-1]);

   variable asp_time = asp.time[ok_times];
   variable asp_ra = asp.euler[ok_times, 0];
   variable asp_dec = 90.0 - asp.euler[ok_times, 1];

   debug ("Computing new Euler angles\r");
   variable x, y;
   (x, y) = wcsfuns_project (wcs, asp_ra, asp_dec);
   x = __tmp(x) + interpol (asp_time, lc_t, dx);
   y = __tmp(y) + interpol (asp_time, lc_t, dy);
   (asp_ra, asp_dec) = wcsfuns_deproject (wcs, x, y);
   
   variable i = wherenot (isnan(asp_ra) or isnan(asp_dec));
   ok_times = ok_times[i];
   asp.euler[ok_times, 0] = (asp_ra[i] + 360.0) mod 360.0;
   asp.euler[ok_times, 1] = 90.0 - asp_dec[i];

   debug ("Writing new attitude file: %s\r", outattfile);

   variable history = ["This file was created by the isis-script v$_isis_version_string: "$,
		       sprintf ("%s v%s", __argv[0], Script_Version_String),
		       sprintf ("Command line: %s", strjoin (__argv, " "))
		      ];

   write_att_file (inattfile, outattfile, asp, history);
}


define slsh_main ()
{
   variable dt = Default_DT;
   variable gtifile = NULL;
   variable lcfile = NULL;
   variable sky_opt = NULL;
   variable cel_opt = NULL;
   variable use_ds9 = 0;

   variable c = cmdopt_new ();
   c.add("ds9", &use_ds9; default=1);
   c.add("dt", &dt;type="float");
   c.add("cel",&cel_opt; type="str");
   c.add("sky", &sky_opt; type="str");
   c.add("gtifile", &gtifile; type="str");
   c.add("lcfile", &lcfile; type="str");
   c.add("h|help", &exit_usage);
   c.add("q|quiet", &Verbose; default=0);
   c.add("v|version", &exit_version);

   variable i = c.process (__argv, 1);

   if (i + 3 != __argc)
     exit_usage ();

   variable x0 = NULL, y0 = NULL, radius = NULL, is_celestial = 0;
   if (cel_opt != NULL)
     {
	(x0, y0, radius) = parse_cel_extracton_region (cel_opt);
	is_celestial = 1;
     }
   if (sky_opt != NULL)
     {
	(x0, y0, radius) = parse_sky_extracton_region (sky_opt);
	is_celestial = 0;
     }
   
   variable evtfile = __argv[i],
     inattfile = __argv[i+1],
     outattfile = __argv[i+2];

   if (stat_file (outattfile) != NULL)
     {
	() = fprintf (stderr, "%s exists -- it will not be overwritten\n", outattfile);
	exit (1);
     }

   if (gtifile == NULL)
     gtifile = evtfile;
   else if (gtifile == "NONE")
     gtifile = NULL;

   variable ds9 = NULL;

   if (use_ds9 || (radius == NULL))
     {
	debug ("Starting ds9\r");
	ds9 = open_ds9 ();
	if ((ds9 == NULL) && (radius == NULL))
	  {
	     () = fprintf (stderr, "The ds9 module is not available.  You must specify an extraction region.\n");
	     exit (1);
	  }
	ds9.set_cmap ("sls");
	ds9.set_scale ("sqrt");
     }

   debug ("Reading %s\r", evtfile);
   variable evt = fits_read_table (evtfile, "time", "x", "y");
   variable wcs = fitswcs_get_column_wcs (evtfile, ["X", "Y"]);

   struct_filter (evt, where ((1 <= evt.x <= 1536) 
			      and (1 <= evt.y <= 1536)); dim=0);

   if (ds9 != NULL)
     display_events_in_ds9 (ds9, evt, wcs);

   if (x0 == NULL)
     {
	(x0, y0, radius) = get_extraction_region_from_ds9 (ds9, evt, wcs);
     }
   else
     {
	if (is_celestial)
	  (x0, y0) = wcsfuns_project (wcs, x0, y0);
	if (ds9 != NULL)
	  {
	     ds9.put_regions (["circle($x0,$y0,$radius)"$]);
	  }
     }

   debug ("Using (%g,%g) for the source position\r", x0, y0);
   debug ("Filtering the events\r");

   ifnot (any (where (where(hypot (evt.x-x0,evt.y-y0)<radius))))
     {
	() = fprintf (stderr, "No events in the filtered event list -- check the coordinates\n");
	exit (1);
     }

   debug ("Reading the GTIs from %s\r", gtifile);
   variable gtis = read_gti_table (gtifile);

   debug ("Computing the light curve\r");
   variable lc = compute_lc (evt, gtis, dt);
   if (lcfile != NULL)
     {
	debug ("Saving the light curve to %s\r", lcfile);
	fits_write_binary_table (lcfile, "LCSTATS", lc);
     }
   % Done with the gtis -- get rid of them
   gtis = NULL;

   variable lc_t = 0.5*(lc.time_min + lc.time_max);
   i = wherenot (isnan(lc.x));
   variable dx = mean (lc.x[i]) - lc.x;
   variable dy = mean (lc.y[i]) - lc.y;

   correct_att_file (inattfile, outattfile, wcs, lc, lc_t, dx, dy);

   if (ds9 != NULL)
     {
	debug ("Computing the corrected PSF\r");
	variable x = evt.x; evt.x = NULL;
	variable y = evt.y; evt.y = NULL;
	evt.x = __tmp(x) + interpol (evt.time, lc_t, dx);
	evt.y = __tmp(y) + interpol (evt.time, lc_t, dy);
	ds9.send ("tile yes");
	ds9.clear(2);
	display_events_in_ds9 (ds9, evt, wcs);
     }

   debug ("Done\n");
}