# Bojan Nikolic  <bn204@mrao.cam.ac.uk>, <bojan@bnikolic.co.uk>
# October 2007
#
# Try to simplify some of the alma simulator code

import numpy
from itertools import izip

from almaconfparse import *

def ExportUVFits(fnamein,
                 fnameout):

    "Wrap the casa command of the same name"

    if os.access(fnameout, os.F_OK ) :
        os.remove(fnameout)

    default(exportuvfits)
    exportuvfits(fnamein, fnameout)



def SetupSM(smi,
            antennalist   ,
            startfreq,
            chanwidth,
            nchan,
            **kwargs):
    
    "Setup an instance of the casa simulator tool to simple configuration"

    """
    antennalist: filename that contains the antenna poistions
    startfreq  : start frequency of the only spectral window
    """
    
    x, y, z, d, names  = AntennaFileToLists( antennalist)    
    
    posalma=me.observatory('ALMA')
    smi.setconfig(telescopename='ALMA',
                  x=x.tolist(), 
                  y=y.tolist(),
                  z=z.tolist(), 
                  dishdiameter=d.tolist(), 
                  mount=['alt-az'],
                  antname=names,
                  coordsystem='local',
                  referencelocation=posalma);

    smi.setspwindow(spwname="simband",
                    freq= "%fGHz" % startfreq,
                    deltafreq="%fGHz" % chanwidth, 
                    freqresolution="%fGHz" % chanwidth,
                    nchannels=nchan,
                    stokes='XX YY');

    smi.setfeed(mode='perfect X Y',
                pol=['']);

    smi.setlimits(shadowlimit=0.01,
                  elevationlimit='10deg');

    # NO autocorrelations
    smi.setauto(0.0);                

def Simulate(project,
             complist,
             antennalist,
             startfreq,
             chanwidth,
             nchan,
             direction,
             integrationtime,
             observationtime,
             refdate="2008/06/21/03:25:00",
             export=True,
             **kwargs):
    """
    All frequencies are in gigaherz
    direction is a string
    """

    msfile=project+'.ms'

    sm.open(msfile)

    SetupSM( sm,
             antennalist,
             startfreq,
             chanwidth,
             nchan)

    sm.setfield(sourcename="SimField",
                sourcedirection=direction)

    reftime = me.epoch('TAI', refdate);
    sm.settimes(integrationtime="%fs" % integrationtime,
                usehourangle=True, 
                referencetime=reftime)

    # Actually getting to observation
    sm.observe(sourcename="SimField",
               spwname="simband",
               starttime="0s",
               stoptime="%fs"%observationtime)

    sm.setdata(fieldid=[0])
    sm.setvp()
    sm.predict(imagename=[""], complist=complist)
    sm.close()
    sm.done()

    if export:
        default(exportuvfits)
        exportuvfits(msfile, msfile+".fits")
    

    

    
def SimulateSwitched(project,
                     complist_fname,
                     directionl,
                     obstimel  ,
                     integrationtime,
                     obsstart,
                     conf = None):

    "Simulate a switched set of observations"

    """
    obsstart is time of the start of observation

    conf is the alma configuration
    """

    msfile=project+'.ms'
    sm.open(msfile)
    
    # To begin with set up with defaults
    SetupSM( sm,
             antennalist=Conf(conf),
             startfreq=300,
             chanwidth=8,
             nchan=1)

    reftime = me.epoch('TAI', obsstart);
    sm.settimes(integrationtime="%fs" % integrationtime,
                usehourangle=False, 
                referencetime=reftime)

    runtime = 0
    for fi , (direction, obstime) in enumerate(izip( directionl, obstimel )):

        print  fi , (direction, obstime)

        cfname = ( "F%i" % fi)

        if  ( (fi % 2) == 0):
            sm.setfield(sourcename=cfname,
                        sourcedirection=direction,
                        calcode="C")            
        else:
            sm.setfield(sourcename=cfname,
                        sourcedirection=direction,)

        
        sm.observe(sourcename=cfname,
                   spwname="simband",
                   starttime="%fs" % runtime,
                   stoptime="%fs"%  (runtime+obstime) )
        runtime += obstime

        sm.setdata(fieldid=[fi])
        sm.setvp()
        sm.predict(imagename=[""], complist=complist_fname)

    sm.close()
    sm.done()

    ExportUVFits(msfile, msfile+".fits")


def SimulateSwitchedV2(project,
                       complist_fname,
                       directionl,
                       obstimel  ,
                       integrationtime,
                       obsstart,
                       cycles=1,
                       conf = None):

    "Simulate a switched set of observations"

    """
    In this version only two sources are possible, calibrator and
    target.

    so directionl and obstimel must be lists with two elements

    obsstart is time of the start of observation

    conf is the alma configuration
    """

    msfile=project+'.ms'
    sm.open(msfile)
    
    # To begin with set up with defaults
    SetupSM( sm,
             antennalist=Conf(conf),
             startfreq=300,
             chanwidth=8,
             nchan=1)

    reftime = me.epoch('TAI', obsstart);
    sm.settimes(integrationtime="%fs" % integrationtime,
                usehourangle=False, 
                referencetime=reftime)

    runtime = 0
    for cno in range(cycles):

        fi = cno % 2
        direction= directionl[fi]
        obstime  = obstimel[fi]

        cfname = ( "F%i" % fi)

        if  (fi == 0):
            sm.setfield(sourcename=cfname,
                        sourcedirection=direction,
                        calcode="C")            
        else:
            sm.setfield(sourcename=cfname,
                        sourcedirection=direction,)

        
        sm.observe(sourcename=cfname,
                   spwname="simband",
                   starttime="%fs" % runtime,
                   stoptime="%fs"%  (runtime+obstime) )
        runtime += obstime

        sm.setdata(fieldid=[fi])
        sm.setvp()
        sm.predict(imagename=[""], complist=complist_fname)

    sm.close()
    sm.done()

    ExportUVFits(msfile, msfile+".fits")