# Bojan Nikolic <bn204@mrao.cam.ac.uk>, <bojan@bnikolic.co.uk>
#
# Initial version January 2008
#
# Tools for using Obit

import os

import pyfits

import obitsetup
import uvrngtobin

import Obit, UV, UVImager, Image, ImageMosaic, OSystem, History, OErr

# Fitting
import ImageFit , FitModel, FitRegion


def ObitSys(datapath,
            nowarn=False):

    "Set up an instance of Obit"

    """
    As things stand can only really access one directory 
    """

    # Check the data path exists, create if not
    if not os.path.exists(datapath):
        os.makedirs(datapath)

    err=OErr.OErr()

    if nowarn:
        Obit.SetLogLevel(err.me,3)
    
    ObitSys=OSystem.OSystem ("Imager", 1, 100,
                             1, [datapath], # AIPS dirs
                             1, [datapath], # FITS directories
                             1, 0, err)

    OErr.printErrMsg(err, "Error with Obit startup")    

    return err, ObitSys

def TestSameDir(f1,
                f2):

    return ( os.path.dirname(f1) ==
             os.path.dirname(f2) )

def CvtMaybe(fnamein):

    "If random group fits file convert to bintable"

    h=pyfits.open(fnamein)[0].header

    if h.has_key("NAXIS1") and h["NAXIS1"]==0:
        fnameout= os.path.splitext(fnamein)[0]+".bin.fits"
        uvrngtobin.Cvt( fnamein, fnameout,
                        overwrite=1,
                        setdate=True,
                        corrupt_pol=False)
        return fnameout
    else:
        return fnamein
    
class ObitInputData:

    "A simple class to persist Obit structures related to input data"

    def __init__(self, fnamein,
                 nowarn=False):
        self.fnamein=CvtMaybe(fnamein)

        indir=os.path.dirname(self.fnamein)
        self.err, self.OS = ObitSys( indir,
                                     nowarn=nowarn)

        self.inData  = UV.newPFUV("Input data",
                                  os.path.basename(self.fnamein),
                                  1,
                                  True,
                                  self.err)        
    
def DirtyImage(fnamein,
               fnameout,
               pixsize=None,
               trange =None,
               sources=["SimField"]):

    "Make a dirty image"

    """

    pixsize : pixel size in arcseconds (same for both axes)

    trange : time range of data to image (type list, unit days, if
             both zero that means no selection )
    
    """

    if type(fnamein)==str:
        ObInD = ObitInputData(fnamein)
    else:
        ObInD = fnamein

    if not TestSameDir(ObInD.fnamein,
                       fnameout):
        print "Input/output directories should be the same"




    Input = UVImager.UVCreateImagerInput()
    Input['InData']      = ObInD.inData
    Input['Name']   = 'Image'  # Have to pretend it's AIPS
    Input['Class']  = 'Class'
    Input['Seq']    = 1
    Input['Disk']   = 1
    Input['doBeam'] = True    
    Input['Type']   = 0     # FITS
    Input['doFull']    = False
    Input['Sources']   = sources
    Input['NField'] = 1

    Input['doCalSelect'] = True
    Input['Stokes']      = "RR"
    
    Input['doCalib']     = -1

    if not (pixsize is None) :
        Input['xCells']= Input['yCells']= pixsize

    if not (trange is None ) :
        Input['timeRange']   = trange

    myImager = UVImager.PUVCreateImager(ObInD.err,
                                        "myMosaic",
                                        Input)


    UVImager.PUVImage(myImager,
                      ObInD.err,
                      doWeight=False,
                      doBeam=True)    


    outMosaic = UVImager.PGetMosaic(myImager, ObInD.err)

    tmpImage  = ImageMosaic.PGetImage (outMosaic, 0, ObInD.err)

    outImage  = Image.newPFImage( "Output image",
                                  os.path.basename(fnameout),
                                  1,
                                  False,
                                  ObInD.err)

    Image.PClone( tmpImage,
                  outImage,
                  ObInD.err)   # Same structure etc.

    Image.PCopy (tmpImage,
                 outImage,
                 ObInD.err)
    
    
def FitGaussian(fnamein):

    "Fit a gaussian to supplied FITS image"

    indir = os.path.dirname(fnamein)
    
    err, OS = ObitSys( indir,
                       nowarn=False)

    im=Image.newPFImage("TEST",
                        os.path.basename(fnamein),
                        1,
                        True,
                        err,
                        True)

    Image.POpen( im, 1 , err)

    # Should be ok to fit the whole image
    fin = pyfits.open(fnamein)
    if len(fin[0].data.shape) == 4:
        # stokes, etc
        s1, s2 , nx,ny= fin[0].data.shape
        cd1, cd2      = [ fin[0].header["CDELT%i" %i ] for i in [1,2] ]
        rp1, rp2      = [ fin[0].header["CRPIX%i" %i ] for i in [1,2] ]
    else:
        raise "don't know how to process this image"

    # Fit models
    # Parms to gausmodel: major, minor, pos angle  (all in degrees)
    # DeltaX, DeltaY in pixels
    fm = [FitModel.FitModel(type=FitModel.GaussMod,
                            Peak=0,
                            DeltaX=nx/2+1,
                            DeltaY=ny/2+1,
                            parms=[ 5.0, 5.0, 0.0] ) ]
    # Fit regions
    fr= FitRegion.FitRegion(corner=[0,0], dim=[nx,ny],
                            RMSResid=0.00,
                            models=fm)


    # Fit inputs 
    fitp=ImageFit.FitInput()

    fitp["fitImage"]= im
    fitp["fitRegion"]= fr
    fitp["MaxIter"]=100
    
    ifit=ImageFit.ImageFit("TEST")
    
    ImageFit.PFit( ifit, err, fitp)

    print fm[0].DeltaX, fm[0].DeltaY
    
    res = ( fm[0].Peak ,
            (rp1-fm[0].DeltaX)* cd1,
            (rp2-fm[0].DeltaY)* cd2 ,
            fm[0].parms[0]*cd1 ,
            fm[0].parms[1]*cd2, fm[0].parms[2],
            fr.RMSResid,
            fr.peakResid)
    return res