# Bojan Nikolic <bn204@mrao.cam.ac.uk>, <bojan@bnikolic.co.uk>
#
# Initial version January 2008
#
# Calibration of UV-FITS data

import pyfits
import numpy, numarray

import uvfits
import iofits4



def AntennaPhaseSolutions(fnamein,
                          calsource):

    "Model antenna-based phase drifts"

    """
    Returns a model for each antenna phase for each integration time
    in the supplied UV file.

    calsource is the source to use as the calibrator
    
    """

    din=uvfits.CalibrationData(fnamein)

    sc=numpy.array(uvfits.SourcesColumn(din.fin))
    ts=numpy.array(uvfits.IntegrationTimeSeries(din.fin))

    # Integrations during which we were observing the calibrator
    calints= numpy.unique( ts[sc==calsource])



    res=[]
    for ci in calints:
        print ci
        
        #phase derived form this integration only
        xmask = (ts==ci)
        iphase=uvfits.SolveForAntennaPhases( din,
                                             mask=xmask)
        
        res.append( (ci, iphase))
    return res

def RecordAntennaPhaseSolutions(fnamein,
                                fnameout,
                                calsource):

    "Write antenna phase solutions to a FITS file "

    """
    Each integration solved for separatelly
    """

    fin=pyfits.open(fnamein)
    
    r=AntennaPhaseSolutions(fnamein,
                            calsource)

    ptimes = numpy.array( [x[0] for x in r])
    phases = numarray.array( [x[1] for x in r])
    
    coldefs = [
        pyfits.Column( "time", "E", "s",
                       array=numarray.array(ptimes)),
        pyfits.Column( "APhase", "%iE" %phases.shape[1] , "rad",
                       array=phases)
        ]    

    tabout= pyfits.new_table( coldefs )
    # output section
    phdu=pyfits.PrimaryHDU( )
    iofits4.AddBzrVersionInfo(phdu)

    iofits4.Write( [phdu,tabout] ,
                   fnameout,
                   overwrite=1)    


def CorrectPhaseSimple(fnamein,
                       acalname,
                       fnameout,
                       scale=1.0):

    "Simple linear correction of phases from a calibration file"

    """
    
    scale: rescale the phase correction. Use when transfering lower
    frequency solutions .
    
    """

    fin = pyfits.open(fnamein)

    cal = pyfits.open(acalname)[1].data
    cal_times=cal.field("time")
    acals = cal.field("APhase")
    # Current lower calibration observation. Interpolate between cal_i
    # and cal_i+1
    cal_i = 0

    ts=numpy.array(uvfits.IntegrationTimeSeries(fin))
    itimes = numpy.unique(ts)

    phases=numpy.zeros( len(ts))
    for i,itime in enumerate(itimes):

        while cal_i + 1 < len(cal_times) and  cal_times[cal_i+1] < itime:
            cal_i += 1

        if cal_i +1 == len(cal_times):
            amodel= acals[cal_i]
        else:
            tdelta= cal_times[cal_i+1] - cal_times[cal_i]
            tfwd  = cal_times[cal_i+1] - itime
            tback = itime - cal_times[cal_i]

            amodel=   acals[cal_i] * ( tfwd/tdelta) + acals[cal_i+1] * ( tback/tdelta) 
            
        mask= ( ts == itime)
        A=uvfits.PhaseMatrix(fin[0].data,mask)

        if scale is 1.0 :
            phases [ mask] = -1* numarray.dot(A, numarray.array(amodel) ) 
        else:
            phases [ mask] = -1* numarray.dot(A, numarray.array(amodel) ) * scale
    
    uvfits.DePhaseDataAcross(fnamein,
                             fnameout,
                             phases)