""" Code to align science frame files

Context : SRP
Module  : SRPImageMapping.py
Version : 1.7.0
Author  : Stefano Covino
Date    : 29/10/2014
E-mail  : stefano.covino@brera.inaf.it
URL:    : http://www.merate.mi.astro.it/utenti/covino
Purpose : Derive rototraslastion parameters for FITS images.

Usage   : SRPImageMapping [-v] [-h] [-f] -i arg1 [-l arg2] [-m arg3] [-n arg4] [-o] [-p] [-t]
            -f Filter reference object by means of their FWHM
            -i Input FITS file list
            -m Minimum number of stars in common area for matching (default 5)
            -n Number of objects for matching for matching search (default 10)
            -o Save files with object positions
            -t Force pure translation
            -p Integer pixel shift for pure translation
            -r Maximum tolerance (pixel, default 0.0)
            -l Search deepness (defaault 2, same paramter as for ESO-eclipse peak)
            
            The script identifies common objects by means of a triangle match.


History : (28/10/2008) First version.
        : (13/11/2008) Quicker execution.
        : (17/11/2008) Better FWHM filtering.
        : (27/11/2008) Logical improvement.
        : (14/08/2009) Possibility to choose translation only.
        : (08/09/2009) Minor correction.
        : (11/09/2009) pyfits transition.
        : (28/05/2010) Possibility to force integer shifts in pure translation.
        : (04/06/2010) Possibility to give maximum tolerance for association.
        : (22/06/2010) Minor correction.
        : (27/08/2010) Deepness level selectable.
        : (27/09/2010) Minor correction.
        : (14/10/2010) Better import style.
        : (06/08/2011) Better cosmetics.
        : (05/03/2012) Import path updated.
        : (21/05/2012) Better imoprt style.
        : (21/02/2014) Only optparse.
        : (27/10/2014) sextractor source extraction added.
"""



import os, os.path, string, sys, math
from optparse import OptionParser
import SRP.SRPConstants as SRPConstants
import SRP.SRPFiles as SRPFiles
import SRP.SRPUtil as SRPUtil
from SRP.SRPSystem.Pipe import Pipe
from SRP.SRPSystem.Which import Which
from SRP.SRPFrames.SexObjectClass import SexObjects
from SRP.SRPStatistics.AverIterSigmaClipp import AverIterSigmaClipp
import pyfits
#import scipy.optimize


NMINMAG = 20



def totMatch (pars):
    x0 = pars[0]
    y0 = pars[1]
    alpha = pars[2]
    nmatch = 0
    mtchlst = []
    for i in listmag1:
        for l in listmag2:
            l.NX,l.NY = SRPUtil.rotoTrasla((l.X,l.Y),x0,y0,alpha,hX,hY)
            if (gxmin <= l.NX <= gxmax and gymin <= l.NY <= gymax) and not (l in mtchlst):
                td = math.sqrt((l.NX-i.X)**2+(l.NY-i.Y)**2)
                if td < fwhm:
                    nmatch = nmatch + 1
                    mtchlst.append(l)
                    break
    return nmatch



def computeAngles(l1,l2,ris,MINANG=5.0):
    x1r = l1[ris[0]].X
    y1r = l1[ris[0]].Y
    x2r = l1[ris[1]].X
    y2r = l1[ris[1]].Y
    x3r = l1[ris[2]].X
    y3r = l1[ris[2]].Y
    x1t = l2[ris[3]].X
    y1t = l2[ris[3]].Y
    x2t = l2[ris[4]].X
    y2t = l2[ris[4]].Y
    x3t = l2[ris[5]].X
    y3t = l2[ris[5]].Y
    m1r = (y2r-y1r)/(x2r-x1r)
    m2r = (y3r-y1r)/(x3r-x1r)
    tangr = math.fabs((m1r-m2r)/(1+m1r*m2r))
    m1t = (y2t-y1t)/(x2t-x1t)
    m2t = (y3t-y1t)/(x3t-x1t)
    tangt = math.fabs((m1t-m2t)/(1+m1t*m2t))
    angr = SRPUtil.rad2deg(math.atan(tangr))
    angt = SRPUtil.rad2deg(math.atan(tangt))
#        print ris, angr, angt, tangr, tangt
    if math.fabs(angr-angt) < MINANG:
        return True
    else:
        return False


def commonPerc (sx,sy,ang,hx,hy,rx,ry):
    npix = 0
    for i in range(2*hx):
        for l in range(2*hy):
            nx,ny = SRPUtil.rotoTrasla((i+1,l+1),sx,sx,ang,hx,hy)
            if 1 <= nx <= 2*rx and 1 <= ny <= 2*ry:
                npix = npix + 1
    return npix/float((2*rx*2*ry))




def lookforpars (list1,list2,fwhm,reflistX,reflistY,matchp,hx,hy,minang=1.0,PureShift=False):
    area = 1.0
    nm = 0
    mpars = 0,0,0
    for r in range(len(list1)):
        for rr in range(len(list1)):
            for rrr in range(len(list1)):
                for o in range(len(list2)):
                    for oo in range(len(list2)):
                        for ooo in range(len(list2)):
                            if r != rr and r != rrr and rr != rrr and o != oo and o != ooo and oo != ooo:
                                dist1 = math.sqrt((list1[r].X-list1[rr].X)**2+(list1[r].Y-list1[rr].Y)**2)
                                dist2 = math.sqrt((list2[o].X-list2[oo].X)**2+(list2[o].Y-list2[oo].Y)**2)
                                dist3 = math.sqrt((list1[r].X-list1[rrr].X)**2+(list1[r].Y-list1[rrr].Y)**2)
                                dist4 = math.sqrt((list2[o].X-list2[ooo].X)**2+(list2[o].Y-list2[ooo].Y)**2)
                                if math.fabs(dist1-dist2) < fwhm/2.35 and math.fabs(dist3-dist4) < fwhm/2.35:
                                    DX = list2[o].X-list2[oo].X
                                    DY = list2[o].Y-list2[oo].Y
                                    DNX = list1[r].X-list1[rr].X
                                    DNY = list1[r].Y-list1[rr].Y
                                    try:
                                        sina = (DNY-DY*DNX/DX)/(DX+DY**2/DX)
                                        cosa = (DNY+DX*DNX/DY)/(DY+DX**2/DY)
                                    except ZeroDivisionError:
                                        break
                                    sangles1 = SRPUtil.rad2deg(math.atan2(sina,cosa))
                                    DX = list2[o].X-list2[ooo].X
                                    DY = list2[o].Y-list2[ooo].Y
                                    DNX = list1[r].X-list1[rrr].X
                                    DNY = list1[r].Y-list1[rrr].Y
                                    try:
                                        sina = (DNY-DY*DNX/DX)/(DX+DY**2/DX)
                                        cosa = (DNY+DX*DNX/DY)/(DY+DX**2/DY)
                                    except ZeroDivisionError:
                                        break
                                    sangles2 = SRPUtil.rad2deg(math.atan2(sina,cosa))
                                    if PureShift:
                                        condition = ((math.fabs(sangles1) < minang) and (math.fabs(sangles2) < minang))
                                    else:
                                        condition = (math.fabs(sangles1-sangles2) < minang)
                                    if condition:
                                        if PureShift:
                                            sangles = 0.0
                                        else:
                                            sangles = (sangles1 + sangles2)/2.0
                                        sx = list1[rr].X - hx - (list2[oo].X-hx)*math.cos(SRPUtil.deg2rad(sangles)) + (list2[oo].Y-hy)*math.sin(SRPUtil.deg2rad(sangles))
                                        sy = list1[rr].Y - hy - (list2[oo].X-hx)*math.sin(SRPUtil.deg2rad(sangles)) - (list2[oo].Y-hy)*math.cos(SRPUtil.deg2rad(sangles))
#                                                                                print r,rr,rrr,o,oo,ooo,sx,sy,sangles,sangles1,sangles2
                                        inizio = [sx,sy,sangles]
                                        mpars = inizio
                                        nm = totMatch(inizio)
                                        area = commonPerc (sx,sy,sangles,hx,hy,reflistX,reflistY)
#                                                                                print nm
                                        if nm >= matchp:
                                            return area,nm,mpars[0],mpars[1],mpars[2],True
    return area,nm,mpars[0],mpars[1],mpars[2],False




parser = OptionParser(usage="usage: %prog [-v] [-h] [-e/-s] [-f] -i arg1 [-l arg2] [-m arg3] [-n arg4] [-o] [-p] [-t]", version="%prog 1.6.6")
parser.add_option("-e", "--eclipse", action="store_true", dest="eclipse", help="Source extraxcted by eclipse")
parser.add_option("-f", "--fwhmfilter", action="store_true", dest="fwhmf", help="Filter for FWHM value")
parser.add_option("-i", "--inputlist", action="store", nargs=1, type="string", dest="fitsfilelist", help="Input FITS file list")
parser.add_option("-l", "--level", action="store", nargs=1, type="float", dest="level", default=2.0, help="Search deepness.")
parser.add_option("-m", "--matchstars", action="store", nargs=1, type="int", dest="matchp", default=5, help="Minimum number of stars in common area for matching")
parser.add_option("-n", "--nobj", action="store", type="int", dest="nobj", default=10, help="Number of objects for matching search")
parser.add_option("-o", "--outfiles", action="store_true", dest="outfiles", help="Save files with object positions")
parser.add_option("-p", "--pureint", action="store_true", dest="pure", default=False, help="Integer pixel shift for pure translation")
parser.add_option("-r", "--radius", action="store", nargs=1, type="float", dest="radius", default=0.0, help="Max tolerance (pixel)")
parser.add_option("-s", "--sex", action="store_true", dest="sex", help="Source extraxcted by sextractor")
parser.add_option("-t", "--translation", action="store_true", dest="transla", default=False, help="Force pure translation")
parser.add_option("-v", "--verbose", action="store_true", dest="verbose", help="Fully describe operations")
(options, args) = parser.parse_args()


if options.radius < 0.0:
    parser.error("Tolerance must be positive.")
    sys.exit(SRPConstants.SRPExitFailure)


if options.fitsfilelist:
    if os.path.isfile(options.fitsfilelist):
        f = SRPFiles.SRPFile(SRPConstants.SRPLocalDir,options.fitsfilelist,SRPFiles.ReadMode)
        f.SRPOpenFile()
        if options.verbose:
            print "Input FITS file list is: %s." % options.fitsfilelist
        flist = []
        rlist = []
        nentr = 0
        while True:
            dt = f.SRPReadFile()
            if dt != '' and len(dt) > 2:
                flist.append(string.split(string.strip(dt))[0])
                rlist.append(string.join(string.split(string.strip(dt))[1:]))
                nentr = nentr + 1
                if not os.path.isfile(flist[nentr-1]):
                    parser.error("Input FITS file %s not found" % flist[nentr-1])
                if options.verbose:
                    print "FITS file selected: %s" % string.split(string.strip(dt))[0]
            else:
                break
        f.SRPCloseFile()
        if not 2 <= options.matchp <= 19:
            parser.error("Matching number of stars must be in range [2..19].")
        if not options.nobj >= 3:
            parser.error("Number of matching objects too low.")
        if options.level < 0:
            parser.error("Search deepness must be positive.")
        # Reference frame
        reffile = flist[0]
        if options.verbose:
            print "FITS reference frame: %s" % reffile
        if Which(SRPConstants.SRPpeak) == None:
            parser.error("%s package not found." % SRPConstants.SRPpeak)
        starlist = []
        #
        eclipse = False
        sex = False
        if options.eclipse and options.sex:
            parser.error("Only one extraction method can be chosen.")
        elif options.eclipse and not options.sex:
            eclipse = True
        elif not options.eclipse and options.sex:
            sex = True
        else:
            sex = True
        #
        if options.verbose:
            if eclipse:
                print "Data extraction with eclipse..."
            elif sex:
                print "Data extraction with sextractor..."
        #
        for i in range(len(flist)):
            if eclipse:
                stlist = []
                argl = " -f '5 10 15' -F -P '5 10 15' -m clip -k %s " % options.level
                cmd = SRPConstants.SRPpeak+argl+flist[i]
                res = Pipe(cmd)
                if res == None:
                    print "FITS file %s can not be processed." % flist[i]
                    sys.exit(SRPConstants.SRPExitFailure)
                for l in res.split(os.linesep):
                    if len(l) > 2:
                        try:
                            stlist.append(SRPUtil.PeakData(l.split()))
                        except ValueError:
                            pass
            elif sex:
                stlist = []
                d = SexObjects(flist[i],options.level)
                d.FindSexObjects()
                if len(d.ListEntries) == 0:
                    print "FITS file %s can not be processed." % flist[i]
                    sys.exit(SRPConstants.SRPExitFailure)
                d.ListEntries.sort()
                d.ListEntries.reverse()
                for l in d.ListEntries:
                    stlist.append(SRPUtil.PeakData((l.Id,l.X,l.Y,l.npix,1.,1.,1.,1.,l.peak,1.,1,l.ellip,l.flux)))
            #
            hhfits = pyfits.open(flist[i])
            hh = hhfits[0].header
            irange = SRPUtil.getRange(hh)
            grange = SRPUtil.getGoodRange(irange,1.0)
            if i == 0:      # ref frame
                gxmin = grange[0]
                gxmax = grange[1]
                gymin = grange[2]
                gymax = grange[3]
            halfsizeX, halfsizeY = irange[1]/2, irange[3]/2
            del hh
            hhfits.close()
            stlistgood = []
            # remove objects close to the edges
            for l in stlist:
                if grange[0] < l.X < grange[1] and grange[2] < l.Y < grange[3]:
                    stlistgood.append(l)
            # compute FWHM
            sgclist = []
            for l in stlistgood:
                if l.FWHM > 0.0:
                    sgclist.append((l.FWHM,1.0))
#                        fwhm = SRPUtil.averageSigmaClippingIter(sgclist)[0]
            fwhm = AverIterSigmaClipp(sgclist)[0]
            if eclipse:
                if fwhm == None:
                    fwhm = 5.0
                if options.radius > 0.0:
                    fwhm = options.radius
                cmd = SRPConstants.SRPpeak+" -f '%.1f %.1f %.1f' -F -P '%.1f %.1f %.1f' -m clip -k %s  %s" % (fwhm, 2*fwhm, 3*fwhm, fwhm, 2*fwhm, 3*fwhm, options.level, flist[i])
                res = Pipe(cmd)
                if res == None:
                    print "FITS file %s can not be processed." % flist[i]
                    sys.exit(SRPConstants.SRPExitFailure)
                stlist = []
                for l in res.split(os.linesep):
                    if len(l) > 2:
                        try:
                            stlist.append(SRPUtil.PeakData(l.split()))
                        except ValueError:
                            pass
                stlistgood = []
                # remove objects close to the edges
                for l in stlist:
                    if options.fwhmf:
                        if grange[0] < l.X < grange[1] and grange[2] < l.Y < grange[3] and fwhm*0.75 < l.FWHM < fwhm*1.25:
                            stlistgood.append(l)
                    else:
                        if grange[0] < l.X < grange[1] and grange[2] < l.Y < grange[3]:
                            stlistgood.append(l)
            #
            starlist.append((stlistgood,fwhm,flist[i],(halfsizeX,halfsizeY)))
        # Begin fit
        if options.verbose:
            print "Roto-traslation parameters search:"
        root,ext = os.path.splitext(options.fitsfilelist)
        g = SRPFiles.SRPFile(SRPConstants.SRPLocalDir, root+SRPConstants.SRPMapFile+ext, SRPFiles.WriteMode)
        g.SRPOpenFile()
#                msg = "# Filename\tX shift\tY Shift\tRotation Angle\tX rotation center\tY rotation center\tFWHM\tCommon area\tMatched stars\tComment"
#                if options.verbose:
#                        print msg
#                g.SRPWriteFile(msg+os.linesep)
        reflist = starlist[0]
        reffwhm = reflist[1]
        reflist[0].sort()
        reflist[0].reverse()
        for i in starlist:
            if options.verbose:
                print "Processing file: %s" % i[2]
            i[0].sort()
            i[0].reverse()
            nstars = options.nobj
            nmagstars = NMINMAG
            sx = 10.0
            sy = -10.0
            sangles = 0.0
            if len(reflist[0]) < nstars or len(i[0]) < nstars:
                if len(reflist[0]) <= len(i[0]):
                    nstars = len(reflist[0])
                else:
                    nstars = len(i[0])
            if len(reflist[0]) < nmagstars or len(i[0]) < nmagstars:
                if len(reflist[0]) <= len(i[0]):
                    nmagstars = len(reflist[0])
                else:
                    nmagstars = len(i[0])
            list1 = []
            list2 = []
            listmag1 = []
            listmag2 = []
            hX = i[3][0]
            hY = i[3][1]
            fwhm = math.sqrt(reffwhm**2+i[1]**2)
            for l in range(nstars):
                list1.append(reflist[0][l])
                list2.append(i[0][l])
            for l in range(nmagstars):
                listmag1.append(reflist[0][l])
                listmag2.append(i[0][l])
            # look for lines
            ris = lookforpars(list1,list2,fwhm,reflist[3][0],reflist[3][1],options.matchp,hX,hY,1.0,options.transla)
            sX,sY,sANG = ris[2],ris[3],ris[4]
            if options.transla and options.pure:
                sX,sY = round(sX),round(sY)
            area,nm = ris[0],ris[1]
            if ris[5] == True:
                matchflag = "Matched"
            else:
                matchflag = "No match"
            msg = "%s\t%.2f\t%.2f\t%.5f\t%.1f\t%.1f\t%.2f\t%.1f\t%d\t%s" % (i[2], sX, sY, SRPUtil.angleRange(sANG), hX, hY, i[1], 100*area, nm, matchflag)
            if options.outfiles:
                root,ext = os.path.splitext(reflist[2])
                f = file(root+'.pos','w')
                f.write("serv_type: catalog"+os.linesep)
                f.write("long_name: SRP catalog for file %s" % root+'.pos'+os.linesep)
                f.write("short_name: %s" % root+'.pos'+os.linesep)
                f.write("url: ./%s" % root+'.pos'+os.linesep)
                f.write("id_col: 0"+os.linesep)
                f.write("x_col: 13"+os.linesep)
                f.write("y_col: 14"+os.linesep)
                f.write("symbol: {} circle 4"+os.linesep)
                f.write("Id\tX\tY\tNpix\tMean\tDev\tMed\tMin\tMax\tFWHMX\tFWHMY\tFWHM\tFlux\tNX\tNY"+os.linesep)
                f.write("---------"+os.linesep)
                for ii in listmag1:
                    f.write(str(ii))
                f.write("EOD"+os.linesep)
                f.close()
                root,ext = os.path.splitext(i[2])
                f = file(root+'.pos','w')
                f.write("serv_type: catalog"+os.linesep)
                f.write("long_name: SRP catalog for file %s" % root+'.pos'+os.linesep)
                f.write("short_name: %s" % root+'.pos'+os.linesep)
                f.write("url: ./%s" % root+'.pos'+os.linesep)
                f.write("id_col: 0"+os.linesep)
                f.write("x_col: 13"+os.linesep)
                f.write("y_col: 14"+os.linesep)
                f.write("symbol: {} circle 4"+os.linesep)
                f.write("Id\tX\tY\tNpix\tMean\tDev\tMed\tMin\tMax\tFWHMX\tFWHMY\tFWHM\tFlux\tNX\tNY"+os.linesep)
                f.write("---------"+os.linesep)
                for ll in listmag2:
                    f.write(str(ll))
                f.write("EOD"+os.linesep)
                f.close()
            if options.verbose:
                print msg
            g.SRPWriteFile(str(msg)+os.linesep)
        g.SRPCloseFile()
    else:
        parser.error("Input FITS file list %s not found" % options.fitsfilelist)
else:
    parser.print_help()
