include 'imager.g';
include 'calibrater.g';
include 'image.g';
include 'ms.g';
include 'misc.g';
include 'flagger.g';
include 'sysinfo.g';
#
# Some setup variables - put into record par
#
par.pwdir := '/home/sandrock/smyers/Testing/2001-12-06/';
par.prefix := '20010204.xband.';          # Prefix for files generated
par.fitsf := '20010204AX.UVF';            # Name of UVFITS file
par.msf := '20010204AX.ms';               # Name of ms 
par.abscal := '1331+305';                 # Name of primary flux cal
par.absid := 6;                           # FIELD_ID of primary flux cal
par.angcal := '1331+305';                 # Polarization angle calibrator
par.absang := 66.00;                      # R-L Phase Difference of PCal
par.polcal := ['1159+292'];               # Name of d-term calibrator
par.polid :=1;                            # FIELD_ID of d-term calibrator
par.refsrc := ['0713+438', '0854+201', '0927+390', '1159+292', '1229+020', 
            '1256-057', '1310+323'];      # Names of secondary sources
par.nfld := 8;                            # Number of targets
par.fieldlist := ['0713+438', '0854+201', '0927+390', 
               '1159+292', '1229+020', '1256-057', 
               '1310+323', '1331+305'];   # Names of targets
par.gsolint := 60.0;                      # solution interval (secs) for gain
par.grefant := 8;                         # reference antenna for gain sols
par.clniter := 200;                       # Number of clean iterations desired
par.clcell := '0.4arcsec';               # image cell size (Xband)
#par.clcell := '0.02arcsec';               # image cell size (Kband)
#
# Now define functions that use the parameters
#
# Make the ms from fits
#
makemsf := function(par)
{
   m:=fitstoms(fitsfile=par.fitsf, msfile=par.msf);
   m.summary();
   m.done();
   print 'Data loaded into', par.msf;
}
#
# The automapping function itself, use par record defined above
#    
automap := function(par)
{
   msf := par.msf;
#
# Assume ms has been created - check with default os tool
#
   if( dos.fileexists(file=msf) ) {
      print 'Found MS file ',msf;
   } else {
      fail "MS file does not exist";
   }
   msfilename := msf;
#
# Copy setup variables from par into internal variables
#
   prefix := par.prefix;                 # File prefix
   pwdir := par.pwdir;			 # Working directory
   abscal := par.abscal;                 # Name of primary flux cal
   absid := par.absid;                   # FIELD_ID of primary flux cal
   angcal := par.angcal;                 # Polarization angle calibrator
   absang := par.absang;                 # R-L Phase Difference of PCal
   polcal := par.polcal;                 # Name of d-term calibrator
   polid := par.polid;                   # FIELD_ID of d-term calibrator
   refsrc := par.refsrc;                 # Names of secondary sources
   nfld := par.nfld;                     # Number of targets
   fieldlist := par.fieldlist;           # Names of targets
   gsolint := par.gsolint;               # solution interval (secs) for gain
   grefant := par.grefant;               # reference antenna for gain sols
   clniter := par.clniter;               # Number of clean iterations desired
   clcell := par.clcell;                 # Cell size
#
# Basic editing
#
#   fg:=flagger(msfilename);
#   fg.setflagmode('unflag');
#   fg.query(query='FIELD_ID >= 0');
#   fg.setflagmode('flag');
#   fg.filter(column='DATA', operation='range', comparison='Amplitude',
#             range=['0.00001Jy','1.34Jy']);
#   fg.done();
#   print 'Data flagged';
#
# Set default source model for 3C286
# Note that the MODEL_DATA columns now default to 1.0 when imager starts
#
   imagr:=imager(msfilename);
   imagr.setjy(fieldid=absid);
#
# Solve for G Jones (amp & phase selfcal)
#
   caltablename:=spaste(prefix, 'gcal1');
   cal:=calibrater(msfilename);
   cal.setdata();
# First apply parallactic angle correction (necessary)
   cal.setapply(type='P', t=5.0);
# Solve for gain, solution interval = gsolint
   cal.setsolve(type='G', t=gsolint, refant=grefant, table=caltablename);
   cal.solve();
   print 'Gain calibration complete';
#
# Establish the flux density scale
#
   outtablename:=spaste(prefix, 'gcalout1');
   cal.fluxscale(tablein=caltablename, tableout=outtablename,
                 reference=abscal, transfer=refsrc);
   print 'Flux density calibration determined from ',abscal;
#
# Correct the data
   cal.setapply(type='P', t=5.0); 
   cal.setapply(type='G', t=0.0, table=outtablename);
   cal.correct();
   cal.done();
   print 'Gain and Flux density calibration applied';
#
# Loop over IFs
#
   for ( spid in 1:2 ) {
#
# Make a map of primary calibrater, natural weighting
#
      ifstr := sprintf('.if%d',spid);
      calprefix:=spaste(prefix, 'abscal', ifstr);
      modfilename:=spaste(calprefix, '.model');
      clnfilename:=spaste(calprefix, '.restored');
      resfilename:=spaste(calprefix, '.residual');
      psffilename:=spaste(calprefix, '.psf');
      imagr.setdata(fieldid=polid, mode='channel', nchan=1, start=1, step=1, 
                  spwid=spid);
      imagr.setimage(nx=256, ny=256, cellx=clcell, celly=clcell, 
                  stokes='IQUV', fieldid=polid, start=1, step=1, nchan=1,
                  mode='channel', spwid=spid);
      imagr.make(modfilename);
#      imagr.weight('natural');
#      imagr.makeimage(type='psf',image=psffilename)
#      bmaj:=0.; bmin:=0.; bpa:=0.; 
#      imagr.fitpsf(psf=psffilename, bmaj=bmaj, bmin=bmin, bpa=bpa)
      imagr.clean(algorithm='clark', niter=clniter, gain=0.1, 
                  threshold='0.0Jy',
                  model=modfilename, image=clnfilename,
                  residual=resfilename);
#
# Extract peak (I,Q,U,V)
#
      img1:= image(modfilename);
      pmax:= array(0.0, 4);
      for (plane in 1:4) {
         pregion:= drm.box([1,1,plane], [256,256,plane]);
         img1.statistics(statsout=pstats, region=pregion);
         if (abs(pstats.min) > abs(pstats.max)) {
            pmax[plane]:= pstats.min;
         } else {
            pmax[plane]:= pstats.max;
         };
      };
      img1.done();
#
      ival:= pmax[1];
      qval:= pmax[2];
      uval:= pmax[3];
#
# Set the polarization model, applying the AIPS PCAL constraint
#
      imagr.setjy(fieldid=polid, spwid=spid, 
                  fluxdensity=[ival, qval, uval, 0.0]);
      printf('IF%d Polarization model computed for %s\n',spid,polcal);
#
# Solve for D Jones from polarization calibrator
#
      poltablename:=spaste(prefix, 'dcal1');
   
      cal:= calibrater(msfilename);
# Data selection parameters
      fldid:=polid
      selectr := sprintf('FIELD_ID==%d && SPECTRAL_WINDOW_ID==%d',fldid,spid)
      cal.setdata(msselect=selectr);
# Apply parallactic angle and previous gain corrections
      cal.setapply(type='P', t=5.0);
      cal.setapply(type='G', t=0.0, table=outtablename);
# Solve for single day-long D-terms, with 10min pre-avg time
      cal.setsolve(type='D', t=86400.0, preavg=600.0, table=poltablename);
      cal.solve();
#
# Correct all other sources for G Jones and D Jones
#
      selectr := sprintf('SPECTRAL_WINDOW_ID==%d',spid)
      cal.setdata(msselect=selectr);
      cal.setapply(type='P', t=5.0);
      cal.setapply(type='G', t=0.0, table=outtablename);
      cal.setapply(type='D', t=0.0, table=poltablename);
      cal.correct();
      cal.done();
      printf('IF%d Polarization model applied to all sources\n',spid);
#
# Make a polarization-calibrated (I,Q,U,V) image for all sources
#
      fldtablename:=spaste(msfilename, '/FIELD');
      fldtab:= table(fldtablename);
      fldname:= fldtab.getcol('NAME');
      fldtab.close();

      for (ifld in 1:nfld) {
         imagr.setdata(fieldid=ifld, mode='channel', nchan=1, 
                       start=1, step=1, spwid=spid);
         imagr.setimage(nx=256, ny=256, 
                     cellx=clcell, celly=clcell, 
                     stokes='IQUV', fieldid=ifld, start=1, step=1, 
                     nchan=1, mode='channel', spwid=spid);

         srcprefix:= spaste(prefix, fldname[ifld], ifstr);
         modelname:= spaste(srcprefix, '.model');
         restoredname:= spaste(srcprefix, '.restored');
         residualname:= spaste(srcprefix, '.residual');
         psffilename:=spaste(srcprefix, '.psf');
         printf('IF%d Final imaging for %s\n',spid,fldname[ifld]);

         imagr.make(modelname);
#         imagr.weight('natural');
#         imagr.makeimage(type='psf',image=psffilename)
#         bmaj:=0.; bmin:=0.; bpa:=0.; 
#         imagr.fitpsf(psf=psffilename, bmaj=bmaj, bmin=bmin, bpa=bpa)
         imagr.clean(algorithm='clark', niter=clniter, gain=0.1, 
                  threshold='0.0Jy', model=modelname, 
                  image=restoredname, residual=residualname);
      };
   }
   imagr.done();
#
# Initialize records to store Stokes info
#
   ivalc:= [=];
   qvalc:= [=];
   uvalc:= [=];
   chicalc:= [=];
#
# Iterate through all sources
#
   fldtablename:=spaste(msfilename, '/FIELD');
   fldtab:= table(fldtablename);
   fldname:= fldtab.getcol('NAME');

   for (field in fldname) {
      for (spid in 1:2) {
         ifstr := sprintf('.if%d',spid);
         fldprefix:=spaste(prefix, field, ifstr);
         restoredname:= spaste(fldprefix, '.restored');
#
# Measure max abs in (I,Q,U,V)
#
         img1:= image(restoredname);
         pmax:= array(0.0, 4);
         for (plane in 1:4) { 
            pregion:= drm.box([1,1,plane], [256,256,plane]);
            img1.statistics(statsout=pstats, region=pregion);
            if (abs(pstats.min) > abs(pstats.max)) {
               pmax[plane]:= pstats.min;
            } else {
               pmax[plane]:= pstats.max;
            };
         };
         img1.done();
#
         ival:= pmax[1];
         qval:= pmax[2];
         uval:= pmax[3];

         ifstr := sprintf('IF%d',spid);
         ivalc[ifstr][field]:= ival;
         qvalc[ifstr][field]:= qval;
         uvalc[ifstr][field]:= uval;
      }
   };      
#
# Print RL phase differences
#
# angle calibrator as reference
#
   for (spid in 1:2) {
      ifstr := sprintf('IF%d',spid);
      qcalc:= qvalc[ifstr][angcal];
      ucalc:= uvalc[ifstr][angcal];
      chicalc[ifstr]:= atan2 (qcalc, ucalc);
   }

   for (field in fieldlist) {
      print;
      for (spid in 1:2) {
         ifstr := sprintf('IF%d',spid);
         ivala:= ivalc[ifstr][field];
         qvala:= qvalc[ifstr][field];
         uvala:= uvalc[ifstr][field];
         pvala:= sqrt (qvala^2 + uvala^2);
         chivala:= atan2 (qvala, uvala) - chicalc[ifstr] + absang;
         chivala:=regang(chivala);

         print field, " ", ifstr, 
            " I= ", sprintf("%7.3f", ivala), 
            " P= ", sprintf("%7.3f", pvala), 
            " chi= ", sprintf("%7.2f", chivala);
      }
   };
#
# Successful return value should be T
#
   return T;
}

killant := function(msf,ants=[]) {
#
# Use flagger to flag a set of antennas
#
   if( ants == [] ) {
      print 'Blank antenna list - no flagging'
      return
   }
# Assume ms has been created - check with default os tool
   if( dos.fileexists(file=msf) ) {
      print 'Found MS file ',msf;
   } else {
      fail "MS file does not exist";
   }
   msfilename := msf;
   kfg := flagger(msfile=msf);
   kfg.setflagmode('flag');
   kfg.setantennas(ants);
   xf := kfg.query(query='FIELD_ID >= 0');
   if ( is_fail(xf) ) {
      print 'Failure on flagging';
      fail;
   }
   kfg.done();
   print 'Flagged MS file ',msf;
}

quackall := function(msf) {
#
# Quack first 10s from each scan
#
# Assume ms has been created - check with default os tool
   if( dos.fileexists(file=msf) ) {
      print 'Found MS file ',msf;
   } else {
      fail "MS file does not exist";
   }
   msfilename := msf;
   kfg := flagger(msfile=msf);
   kfg.setflagmode('flag');
   kfg.quack();
   kfg.done();
   print 'Quacked MS file ',msf;
}
#
# Utility function: atan2(x,y)
# Input cos(phi),sin(phi) Output atan in range -180 to +180
#
atan2 := function(x, y) {
   fact:= 180.0 / pi;
   if ((x > 0) & (y > 0)) angle:= atan(y/x) * fact;
   if ((x > 0) & (y < 0)) angle:= atan(y/x) * fact + 360;
   if ((x < 0) & (y > 0)) angle:= 180 - atan(y/abs(x)) * fact;
   if ((x < 0) & (y < 0)) angle:= 180 + atan(abs(y)/abs(x)) * fact;
   return angle;
};
#
# Utility function: regang(x)
# regularize angle in degrees to lie from -180 to +180
#
regang := function(x) {
   if ( x > 180.0 ) 
      {
      while ( x > 180.0 ) x -:= 360.0;	
      }
   else
      {
      if ( x < -180.0 )
         {
         while ( x < -180.0 ) x +:= 360.0;
         }
      }
   return x;
};