# Bojan Nikolic <bn204@mrao.cam.ac.uk>, <bojan@bnikolic.co.uk>
#
# Initial version January 2008
#
# Do imaging simulations, using obit

import os, shutil
import re
import pickle

import pyfits
import numpy

import sys
sys.path.append("../obitsim")

import obittools

from imagingsim_utils import UVNameToDirtyBeamName

def DirUVSets(dirin):

    "Return the list of uvdata sets to be reduced in directory"

    msre=re.compile(r".*\.ms\.fits")
    res = []
    for fname in os.listdir(dirin):
        if msre.match(fname):
            res.append( os.path.join(dirin, fname) )

    return res


def DirDirtyMaps(dirin,
                 **kwargs):

    "Reduce whole observations to dirty maps"

    for fnamein in DirUVSets(dirin):
        obittools.DirtyImage( fnamein,
                              UVNameToDirtyBeamName(fnamein),
                              **kwargs)

def SwitchedDirDirtyMaps(dirin,
                         **kwargs):

    "Reduce directory of switched observations"

    #  Need to make it reduce the science target source only
    DirDirtyMaps(dirin,
                 sources=["F1"])    
                         

def DirBeamParams(dirin):

    "Find bestfitting gaussians to each full set in directory"

    res = {}
    for fnamein in DirUVSets(dirin):
        fbase=os.path.basename(fnamein)
        dname=UVNameToDirtyBeamName(fnamein)
        res[fbase] = obittools.FitGaussian( os.path.join(dirin, dname))

    ofile= os.path.join( os.path.dirname(fnamein) ,
                         "DirBeamParams.pickle")
    pickle.dump( res, open( ofile, "w"))

    
def IntegrationRanges(fnamein):
    "Return time ranges suitable for selecting individual integrations"

    il=numpy.unique(pyfits.open(fnamein)[1].data.field("date"))
    eps=1.0/24/3600/2

    res =[ (i-eps, i+eps) for i in il]
    return res

def SnapshotRangeSW(fnamein):

    "Return time ranges suitable for analysing fast-switched observations"

    """
    Just look for gaps caused by the calibrator obs
    """

    srccol=pyfits.open(fnamein)[1].data.field("source")
    date=pyfits.open(fnamein)[1].data.field("date")

    date_target= numpy.unique(date[srccol==2.0])

    ddiff  = date_target[1:] - date_target[:-1]

    # look for gaps longer than 1.5 secs
    maxdiff = 1.5 / 24/ 3600

    # Don't use last one
    gapi = numpy.nonzero((ddiff > maxdiff))[0]

    res=[]
    eps=1.0/24/3600/2.0
    for i in range(len( gapi)-1):
        tb= date_target[ gapi[i] +1 ]
        te= date_target[ gapi[i+1] ]

        res.append( (tb-eps , te+eps))

    return res


def IntBeams(fnamein,
             keepbeams=False,
             fsw=False):

    "List of dirty map beam parameters for each integration"

    """
    fsw : this is a fast switched observation, reduce accordingly
    """

    res=[]
    beampars=[]

    # Thic keeps hold of the obit input data structures so we are not
    # re-opening them all the time
    od =  obittools.ObitInputData(fnamein,
                                  nowarn=True)

    if fsw :
        tr = SnapshotRangeSW(od.fnamein)
    else:
        tr = IntegrationRanges(od.fnamein)
        
    for i, (tb, te) in enumerate(tr):
        print tb, te

        if fsw:
            sources=["F1"]
        else:
            sources=["SimField"]
            
        obittools.DirtyImage( od,
                              "t_obit.fits",
                              trange= [float(tb), float(te)],
                              sources=sources)
        
        tempbeamname = os.path.join(os.path.dirname(fnamein), "t_obit.fits")

        if keepbeams:
            shutil.copyfile( tempbeamname,
                             os.path.join(os.path.dirname(fnamein), "beam-%04i.fits" % i))
        
        fin= pyfits.open(tempbeamname)
        res.append(fin[0].data.max())
        
        beampars.append(obittools.FitGaussian( tempbeamname))
        
    return res , beampars


def DirIntBeamsCalc(dirin):
    "Find peak flux in each itegration of each file in the supplied directory"

    res={}
    for fnamein in DirUVSets(dirin):
        print fnamein
        pf = IntBeams( fnamein)
        res[os.path.basename(fnamein)] = pf
    return res

    ofile= os.path.join( os.path.dirname(fnamein) , "DirIntBeams.pickle")
    pickle.dump( res, open( ofile, "w"))        

def DirIntBeams(dirin):
    """Find peak fluxes and save to pickle"""

    res=DirIntBeamsCalc(dirin)
    ofile= os.path.join( os.path.dirname(fnamein) , "DirIntBeams.pickle")
    pickle.dump( res, open( ofile, "w"))
    
def ImagingMemoSnapshots():

    """Dirty maps for illustration of the imaging memo"""

    fin="../casasim/midsw/midswitchv3-9-50-12w.ms.fits"

    od =  obittools.ObitInputData(fin,
                                  nowarn=True)    

    sources=["F1"]

    il= IntegrationRanges(od.fnamein)

    il = il[3::100]

    for i, (tb, te) in enumerate(il):

        obittools.DirtyImage( od,
                              "beam-midsw-9-50-12w-snap%03i.fits" % (i),
                              trange= [float(tb), float(te)],
                              sources=sources)