include 'imager.g';
include 'calibrater.g';
include 'image.g';
include 'ms.g';
include 'flagger.g';
include 'sysinfo.g';

#
# 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) {
   angle:=x;
   if ( angle > 180.0 ) 
      {
      while ( angle > 180.0 ) angle -:= 360.0;	
      }
   else
      {
      if ( angle < -180.0 )
         {
         while ( angle < -180.0 ) angle +:= 360.0;
         }
      }
   return angle;
};
#
# Some setup variables
#
   pwdir:='/home/sandrock/smyers/Testing/Mar01/08Mar01/'
   fitsf:='20010204AC.UVF'
   prefix:='test.08mar01.'
   abscal:='1331+305';
   absid:=6
   angcal:='1331+305'; 
   absang:=66.00;
   polcal:=['1159+292'];
   polid:=1
   nfld:=8
   refsrc:=['0713+438', '0854+201', '0927+390', '1159+292', '1229+020', 
            '1256-057', '1310+323'];
   fieldlist:=['0713+438', '0854+201', '0927+390', '1159+292', '1229+020', 
               '1256-057', '1310+323', '1331+305'];
#
# Load the data
#
   aipsroot:=sysinfo().root();
#   fitsfilename:=spaste(aipsroot, '/data/demo/CALS00FEB.fits');
   fitsfilename:=spaste(pwdir, fitsf);
   msfilename:=spaste(prefix, 'ms');
   m:=fitstoms(fitsfile=fitsfilename, msfile=msfilename);
   m.done();
   print 'Data loaded into', msfilename;
#
# 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
#
   imagr:=imager(msfilename);
#   imagr.setjy(spwid=1);
   imagr.setjy();
#
# Solve for G Jones
#
   caltablename:=spaste(prefix, 'gcal1');
   cal:=calibrater(msfilename);
#   cal.setdata(msselect='SPECTRAL_WINDOW_ID==1');
   cal.setdata();
   cal.setapply(type='P', t=5.0);
   cal.setsolve(type='G', t=60.0, refant=8, 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';
#
# IF1
#
# Make a map of polarization calibrater
#
   calprefix:=spaste(prefix, 'abscal');
   modfilename:=spaste(calprefix, '.if1.model');
   clnfilename:=spaste(calprefix, '.if1.restored');
   resfilename:=spaste(calprefix, '.if1.residual');
   imagr.setdata(fieldid=polid, mode='channel', nchan=1, start=1, step=1, 
                  spwid=1);
   imagr.setimage(nx=256, ny=256, cellx='0.4arcsec', celly='0.4arcsec', 
                  stokes='IQUV', fieldid=polid, start=1, step=1, nchan=1,
                  mode='channel', spwid=1);
   imagr.make(modfilename);
   imagr.clean(algorithm='clark', niter=5000, 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=1, fluxdensity=[ival, qval, -uval, 0.0]);
   print 'IF1 Polarization model computed for ',polcal;
#
# Solve for D Jones from polarization calibrator
#
   poltablename:=spaste(prefix, 'dcal1');
   
   cal:= calibrater(msfilename);
# Data selection parameters
   fldid:=polid
   spid:=1
   selectr := sprintf('FIELD_ID==%d && SPECTRAL_WINDOW_ID==%d',fldid,spid)
   cal.setdata(msselect=selectr);

   cal.setapply(type='P', t=5.0);
   cal.setapply(type='G', t=0.0, table=outtablename);
   cal.setsolve(type='D', t=86400.0, preavg=600.0, table=poltablename);
   cal.solve();
#
# Correct all other sources for G Jones and D Jones
#
   cal.setdata(msselect='SPECTRAL_WINDOW_ID==1');
   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();
   print 'IF1 Polarization model applied to all sources';
#
# 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=1);
      imagr.setimage(nx=256, ny=256, cellx='0.4arcsec', celly='0.4arcsec', 
                     stokes='IQUV', fieldid=ifld, start=1, step=1, 
                     nchan=1, mode='channel', spwid=1);

      srcprefix:= spaste(prefix, fldname[ifld]);
      modelname:= spaste(srcprefix, '.if1.model');
      restoredname:= spaste(srcprefix, '.if1.restored');
      residualname:= spaste(srcprefix, '.if1.residual');
      print 'IF1 Final imaging for', fldname[ifld];

      imagr.make(modelname);
      imagr.clean(algorithm='clark', niter=5000, gain=0.1, threshold='0.0Jy',
                  model=modelname, image=restoredname, residual=residualname);
   };
#
# IF2
#
# Make a map of polarization calibrater in IF2
#
   calprefix:=spaste(prefix, 'abscal' );
   modfilename:=spaste(calprefix, '.if2.model');
   clnfilename:=spaste(calprefix, '.if2.restored');
   resfilename:=spaste(calprefix, '.if2.residual');
   imagr.setdata(fieldid=polid, mode='channel', nchan=1, start=1, step=1, 
                  spwid=2);
   imagr.setimage(nx=256, ny=256, cellx='0.4arcsec', celly='0.4arcsec', 
                  stokes='IQUV', fieldid=polid, start=1, step=1, nchan=1,
                  mode='channel', spwid=2);
   imagr.make(modfilename);
   imagr.clean(algorithm='clark', niter=5000, 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=2, fluxdensity=[ival, qval, -uval, 0.0]);
   print 'IF2 Polarization model computed for ',polcal;
#
# Solve for D Jones from polarization calibrator
#
   cal:= calibrater(msfilename);
# Data selection parameters
   fldid:=polid
   spid:=2
   selectr := sprintf('FIELD_ID==%d && SPECTRAL_WINDOW_ID==%d',fldid,spid)
   cal.setdata(msselect=selectr);

   cal.setapply(type='P', t=5.0);
   cal.setapply(type='G', t=0.0, table=outtablename);
   cal.setsolve(type='D', t=86400.0, preavg=600.0, table=poltablename);
   cal.solve();
#
# Correct all other sources for G Jones and D Jones
#
   cal.setdata(msselect='SPECTRAL_WINDOW_ID==2');
   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();
   print 'IF2 Polarization model applied to all sources';
#
# 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=2);
      imagr.setimage(nx=256, ny=256, cellx='0.4arcsec', celly='0.4arcsec', 
                     stokes='IQUV', fieldid=ifld, start=1, step=1, 
                     nchan=1, mode='channel', spwid=2);

      srcprefix:= spaste(prefix, fldname[ifld]);
      modelname:= spaste(srcprefix, '.if2.model');
      restoredname:= spaste(srcprefix, '.if2.restored');
      residualname:= spaste(srcprefix, '.if2.residual');
      print 'Final imaging for', fldname[ifld];

      imagr.make(modelname);
      imagr.clean(algorithm='clark', niter=5000, gain=0.1, threshold='0.0Jy',
                  model=modelname, image=restoredname, residual=residualname);
   };
   imagr.done();
#
# Inter-compare AIPS and AIPS++ polarization values
#
# Tabulate the values derived from AIPS/PCAL/IMSTAT (max abs).
#
   ivalc:= [=];
   qvalc:= [=];
   uvalc:= [=];
#   
# Iterate through all sources
#
   fldtablename:=spaste(msfilename, '/FIELD');
   fldtab:= table(fldtablename);
   fldname:= fldtab.getcol('NAME');

   for (field in fldname) {
#
# IF1
#
      fldprefix:=spaste(prefix, field);
      restoredname:= spaste(fldprefix, '.if1.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];
 
      ivalc['IF1'][field]:= ival;
      qvalc['IF1'][field]:= qval;
      uvalc['IF1'][field]:= uval;
#
# IF2
#
      fldprefix:=spaste(prefix, field);
      restoredname:= spaste(fldprefix, '.if2.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];
 
      ivalc['IF2'][field]:= ival;
      qvalc['IF2'][field]:= qval;
      uvalc['IF2'][field]:= uval;
   };      

#
# Print RL phase differences
#
# angle calibrator as reference
#
   qcalc1:= qvalc['IF1'][angcal];
   ucalc1:= uvalc['IF1'][angcal];
   chicalc1:= atan2 (qcalc1, ucalc1);

   qcalc2:= qvalc['IF2'][angcal];
   ucalc2:= uvalc['IF2'][angcal];
   chicalc2:= atan2 (qcalc2, ucalc2);

   for (field in fieldlist) {
      ivala1:= ivalc['IF1'][field];
      qvala1:= qvalc['IF1'][field];
      uvala1:= uvalc['IF1'][field];
      pvala1:= sqrt (qvala1^2 + uvala1^2);
      chivala1:= atan2 (qvala1, uvala1) - chicalc1 + absang;
      chivala1:=regang(chivala1);

      ivala2:= ivalc['IF2'][field];
      qvala2:= qvalc['IF2'][field];
      uvala2:= uvalc['IF2'][field];
      pvala2:= sqrt (qvala2^2 + uvala2^2);
      chivala2:= atan2 (qvala2, uvala2) - chicalc2 + absang;
      chivala2:=regang(chivala2);

      print;
      print field, "IF1 I= ", sprintf("%7.3f", ivala1), 
         " P= ", sprintf("%7.3f", pvala1), 
         " chi= ", sprintf("%7.2f", chivala1);
      print field, "IF2 I= ", sprintf("%7.3f", ivala2), 
         " P= ", sprintf("%7.3f", pvala2), 
         " chi= ", sprintf("%7.2f", chivala2);
   };
#