import os, shutil
import numpy
import numpy.fft as FFT
import math

from asap.scantable import scantable
from asap.parameters import rcParams
from asap.logging import asaplog, asaplog_post_dec
from asap.selector import selector
from asap.asapgrid import asapgrid2
from asap._asap import SBSeparator

class sbseparator:
    """
    The sbseparator class is defined to separate SIGNAL and IMAGE
    sideband spectra observed by frequency-switching technique.
    It also helps supressing emmission of IMAGE sideband.
    *** WARNING *** THIS MODULE IS EXPERIMENTAL
    Known issues:
    - Frequencies of IMAGE sideband cannot be reconstructed from
      information in scantable in sideband sparation. Frequency
      setting of SIGNAL sideband is stored in output table for now.
    - Flag information (stored in FLAGTRA) is ignored.

    Example:
        # Create sideband separator instance whith 3 input data
        sbsep = sbseparator(['test1.asap', 'test2.asap', 'test3.asap'])
        # Set reference IFNO and tolerance to select data
        sbsep.set_frequency(5, 30, frame='TOPO')
        # Set direction tolerance to select data in unit of radian
        sbsep.set_dirtol(1.e-5)
        # Set rejection limit of solution
        sbsep.set_limit(0.2)
        # Solve image sideband as well
        sbsep.set_both(True)
        # Invoke sideband separation
        sbsep.separate('testout.asap', overwrite = True)
    """
    def __init__(self, infiles):
        self.intables = None
        self.signalShift = []
        self.imageShift = []
        self.dsbmode = True
        self.getboth = False
        self.rejlimit = 0.2
        self.baseif = -1
        self.freqtol = 10.
        self.freqframe = ""
        self.solveother = False
        self.dirtol = [1.e-5, 1.e-5] # direction tolerance in rad (2 arcsec)
        #self.lo1 = 0.

        self.tables = []
        self.nshift = -1
        self.nchan = -1

        self.set_data(infiles)
        
        self._separator = SBSeparator(infiles)

    @asaplog_post_dec
    def set_data(self, infiles):
        """
        Set data to be processed.

        infiles  : a list of filenames or scantables
        """
        if not (type(infiles) in (list, tuple, numpy.ndarray)):
            infiles = [infiles]
        if isinstance(infiles[0], scantable):
            # a list of scantable
            for stab in infiles:
                if not isinstance(stab, scantable):
                    asaplog.post()
                    raise TypeError, "Input data is not a list of scantables."

            #self._separator._setdata(infiles)
            self._reset_data()
            self.intables = infiles
        else:
            # a list of filenames
            for name in infiles:
                if not os.path.exists(name):
                    asaplog.post()
                    raise ValueError, "Could not find input file '%s'" % name
            
            #self._separator._setdataname(infiles)
            self._reset_data()
            self.intables = infiles

        asaplog.push("%d files are set to process" % len(self.intables))


    def _reset_data(self):
        del self.intables
        self.intables = None
        self.signalShift = []
        #self.imageShift = []
        self.tables = []
        self.nshift = -1
        self.nchan = -1

    @asaplog_post_dec
    def set_frequency(self, baseif, freqtol, frame=""):
        """
        Set IFNO and frequency tolerance to select data to process.

        Parameters:
          - reference IFNO to process in the first table in the list
          - frequency tolerance from reference IF to select data
          frame  : frequency frame to select IF
        """
        self._reset_if()
        self.baseif = baseif
        if isinstance(freqtol,dict) and freqtol["unit"] == "Hz":
            if freqtol['value'] > 0.:
                self.freqtol = freqtol
            else:
                asaplog.post()
                asaplog.push("Frequency tolerance should be positive value.")
                asaplog.post("ERROR")
                return
        else:
            # torelance in channel unit
            if freqtol > 0:
                self.freqtol = float(freqtol)
            else:
                asaplog.post()
                asaplog.push("Frequency tolerance should be positive value.")
                asaplog.post("ERROR")
                return
        self.freqframe = frame

    def _reset_if(self):
        self.baseif = -1
        self.freqtol = 0
        self.freqframe = ""
        self.signalShift = []
        #self.imageShift = []
        self.tables = []
        self.nshift = 0
        self.nchan = -1

    @asaplog_post_dec
    def set_dirtol(self, dirtol=[1.e-5,1.e-5]):
        """
        Set tolerance of direction to select data
        """
        # direction tolerance in rad
        if not (type(dirtol) in [list, tuple, numpy.ndarray]):
            dirtol = [dirtol, dirtol]
        if len(dirtol) == 1:
            dirtol = [dirtol[0], dirtol[0]]
        if len(dirtol) > 1:
            self.dirtol = dirtol[0:2]
        else:
            asaplog.post()
            asaplog.push("Invalid direction tolerance. Should be a list of float in unit radian")
            asaplog.post("ERROR")
            return
        asaplog.post("Set direction tolerance [%f, %f] (rad)" % \
                     (self.dirtol[0], self.dirtol[1]))

    @asaplog_post_dec
    def set_shift(self, mode="DSB", imageshift=None):
        """
        Set shift mode and channel shift of image band.

        mode       : shift mode ['DSB'|'SSB'(='2SB')]
                     When mode='DSB', imageshift is assumed to be equal
                     to the shift of signal sideband but in opposite direction.
        imageshift : a list of number of channel shift in image band of
                     each scantable. valid only mode='SSB'
        """
        if mode.upper().startswith("D"):
            # DSB mode
            self.dsbmode = True
            self.imageShift = []
        else:
            if not imageshift:
                raise ValueError, "Need to set shift value of image sideband"
            self.dsbmode = False
            self.imageShift = imageshift
            asaplog.push("Image sideband shift is set manually: %s" % str(self.imageShift))

    @asaplog_post_dec
    def set_both(self, flag=False):
        """
        Resolve both image and signal sideband when True.
        """
        self.getboth = flag
        if self.getboth:
            asaplog.push("Both signal and image sidebands are solved and output as separate tables.")
        else:
            asaplog.push("Only signal sideband is solved and output as an table.")

    @asaplog_post_dec
    def set_limit(self, threshold=0.2):
        """
        Set rejection limit of solution.
        """
        #self._separator._setlimit(abs(threshold))
        self.rejlimit = threshold
        asaplog.push("The threshold of rejection is set to %f" % self.rejlimit)


    @asaplog_post_dec
    def set_solve_other(self, flag=False):
        """
        Calculate spectra by subtracting the solution of the other sideband
        when True.
        """
        self.solveother = flag
        if flag:
            asaplog.push("Expert mode: solution are obtained by subtraction of the other sideband.")

    def set_lo1(self,lo1):
        """
        Set LO1 frequency to calculate frequency of image sideband.

        lo1 : LO1 frequency in float
        """
        lo1val = -1.
        if isinstance(lo1, dict) and lo1["unit"] == "Hz":
            lo1val = lo1["value"]
        else:
            lo1val = float(lo1)
        if lo1val <= 0.:
            asaplog.push("Got negative LO1 frequency. It will be ignored.")
            asaplog.post("WARN")
        else:
            self._separator.set_lo1(lo1val)


    def set_lo1root(self, name):
        """
        Set MS name which stores LO1 frequency of signal side band.
        It is used to calculate frequency of image sideband.

        name : MS name which contains 'ASDM_SPECTRALWINDOW' and
               'ASDM_RECEIVER' tables.
        """
        self._separator.set_lo1root(name)


    @asaplog_post_dec
    def separate(self, outname="", overwrite=False):
        """
        Invoke sideband separation.

        outname   : a name of output scantable
        overwrite : overwrite existing table
        """
        # List up valid scantables and IFNOs to convolve.
        #self._separator._separate()
        self._setup_shift()
        #self._preprocess_tables()

        ### TEMPORAL ###
        self._separator._get_asistb_from_scantb(self.tables[0])
        ################

        nshift = len(self.tables)
        signaltab = self._grid_outtable(self.tables[0])
        if self.getboth:
            imagetab = signaltab.copy()

        rejrow = []
        for irow in xrange(signaltab.nrow()):
            currpol = signaltab.getpol(irow)
            currbeam = signaltab.getbeam(irow)
            currdir = signaltab.get_directionval(irow)
            spec_array, tabidx = self._get_specarray(polid=currpol,\
                                                     beamid=currbeam,\
                                                     dir=currdir)
            #if not spec_array:
            if len(tabidx) == 0:
                asaplog.post()
                asaplog.push("skipping row %d" % irow)
                rejrow.append(irow)
                continue
            signal = self._solve_signal(spec_array, tabidx)
            signaltab.set_spectrum(signal, irow)

            # Solve image side side band
            if self.getboth:
                image = self._solve_image(spec_array, tabidx)
                imagetab.set_spectrum(image, irow)

        # TODO: Need to remove rejrow form scantables here
        signaltab.flag_row(rejrow)
        if self.getboth:
            imagetab.flag_row(rejrow)
        
        if outname == "":
            outname = "sbsepareted.asap"
        signalname = outname + ".signalband"
        if os.path.exists(signalname):
            if not overwrite:
                raise Exception, "Output file '%s' exists." % signalname
            else:
                shutil.rmtree(signalname)
        signaltab.save(signalname)

        if self.getboth:
            # Warnings
            asaplog.post()
            asaplog.push("Saving IMAGE sideband.")
            #asaplog.push("Note, frequency information of IMAGE sideband cannot be properly filled so far. (future development)")
            #asaplog.push("Storing frequency setting of SIGNAL sideband in FREQUENCIES table for now.")
            #asaplog.post("WARN")

            imagename = outname + ".imageband"
            if os.path.exists(imagename):
                if not overwrite:
                    raise Exception, "Output file '%s' exists." % imagename
                else:
                    shutil.rmtree(imagename)
            # Update frequency information
            self._separator.set_imgtable(imagetab)
            self._separator.solve_imgfreq()
            imagetab.save(imagename)


    def _solve_signal(self, data, tabidx=None):
        if not tabidx:
            tabidx = range(len(data))

        tempshift = []
        dshift = []
        if self.solveother:
            for idx in tabidx:
                tempshift.append(-self.imageShift[idx])
                dshift.append(self.signalShift[idx] - self.imageShift[idx])
        else:
            for idx in tabidx:
                tempshift.append(-self.signalShift[idx])
                dshift.append(self.imageShift[idx] - self.signalShift[idx])

        shiftdata = numpy.zeros(data.shape, numpy.float)
        for i in range(len(data)):
            shiftdata[i] = self._shiftSpectrum(data[i], tempshift[i])
        ifftdata = self._Deconvolution(shiftdata, dshift, self.rejlimit)
        result_image = self._combineResult(ifftdata)
        if not self.solveother:
            return result_image
        result_signal = self._subtractOtherSide(shiftdata, dshift, result_image)
        return result_signal


    def _solve_image(self, data, tabidx=None):
        if not tabidx:
            tabidx = range(len(data))

        tempshift = []
        dshift = []
        if self.solveother:
            for idx in tabidx:
                tempshift.append(-self.signalShift[idx])
                dshift.append(self.imageShift[idx] - self.signalShift[idx])
        else:
            for idx in tabidx:
                tempshift.append(-self.imageShift[idx])
                dshift.append(self.signalShift[idx] - self.imageShift[idx])

        shiftdata = numpy.zeros(data.shape, numpy.float)
        for i in range(len(data)):
            shiftdata[i] = self._shiftSpectrum(data[i], tempshift[i])
        ifftdata = self._Deconvolution(shiftdata, dshift, self.rejlimit)
        result_image = self._combineResult(ifftdata)
        if not self.solveother:
            return result_image
        result_signal = self._subtractOtherSide(shiftdata, dshift, result_image)
        return result_signal

    @asaplog_post_dec
    def _grid_outtable(self, table):
        # Generate gridded table for output (Just to get rows)
        gridder = asapgrid2(table)
        gridder.setIF(self.baseif)

        cellx = str(self.dirtol[0])+"rad"
        celly = str(self.dirtol[1])+"rad"
        dirarr = numpy.array(table.get_directionval()).transpose()
        mapx = dirarr[0].max() - dirarr[0].min()
        mapy = dirarr[1].max() - dirarr[1].min()
        centy = 0.5 * (dirarr[1].max() + dirarr[1].min())
        nx = max(1, numpy.ceil(mapx*numpy.cos(centy)/self.dirtol[0]))
        ny = max(1, numpy.ceil(mapy/self.dirtol[0]))

        asaplog.push("Regrid output scantable with cell = [%s, %s]" % \
                     (cellx, celly))
        gridder.defineImage(nx=nx, ny=ny, cellx=cellx, celly=celly)
        gridder.setFunc(func='box', convsupport=1)
        gridder.setWeight(weightType='uniform')
        gridder.grid()
        return gridder.getResult()

    @asaplog_post_dec
    def _get_specarray(self, polid=None, beamid=None, dir=None):
        ntable = len(self.tables)
        spec_array = numpy.zeros((ntable, self.nchan), numpy.float)
        nspec = 0
        asaplog.push("Start data selection by POL=%d, BEAM=%d, direction=[%f, %f]" % (polid, beamid, dir[0], dir[1]))
        tabidx = []
        for itab in range(ntable):
            tab = self.tables[itab]
            # Select rows by POLNO and BEAMNO
            try:
                tab.set_selection(pols=[polid], beams=[beamid])
                if tab.nrow() > 0: tabidx.append(itab)
            except: # no selection
                asaplog.post()
                asaplog.push("table %d - No spectrum ....skipping the table" % (itab))
                asaplog.post("WARN")
                continue

            # Select rows by direction
            spec = numpy.zeros(self.nchan, numpy.float)
            selrow = []
            for irow in range(tab.nrow()):
                currdir = tab.get_directionval(irow)
                if (abs(currdir[0] - dir[0]) > self.dirtol[0]) or \
                   (abs(currdir[1] - dir[1]) > self.dirtol[1]):
                    continue
                selrow.append(irow)
            if len(selrow) == 0:
                asaplog.post()
                asaplog.push("table %d - No spectrum ....skipping the table" % (itab))
                asaplog.post("WARN")
                continue
            else:
                seltab = tab.copy()
                seltab.set_selection(selector(rows=selrow))

            if tab.nrow() > 1:
                asaplog.push("table %d - More than a spectrum selected. averaging rows..." % (itab))
                tab = seltab.average_time(scanav=False, weight="tintsys")
            else:
                tab = seltab

            spec_array[nspec] = tab._getspectrum()
            nspec += 1

        if nspec != ntable:
            asaplog.post()
            #asaplog.push("Some tables has no spectrum with POL=%d BEAM=%d. averaging rows..." % (polid, beamid))
            asaplog.push("Could not find corresponding rows in some tables.")
            asaplog.push("Number of spectra selected = %d (table: %d)" % (nspec, ntable))
            if nspec < 2:
                asaplog.push("At least 2 spectra are necessary for convolution")
                asaplog.post("ERROR")
                return False, tabidx

        return spec_array[0:nspec], tabidx


    @asaplog_post_dec
    def _setup_shift(self):
        ### define self.tables, self.signalShift, and self.imageShift
        if not self.intables:
            asaplog.post()
            raise RuntimeError, "Input data is not defined."
        #if self.baseif < 0:
        #    asaplog.post()
        #    raise RuntimeError, "Reference IFNO is not defined."
        
        byname = False
        #if not self.intables:
        if isinstance(self.intables[0], str):
            # A list of file name is given
            if not os.path.exists(self.intables[0]):
                asaplog.post()
                raise RuntimeError, "Could not find '%s'" % self.intables[0]
            
            stab = scantable(self.intables[0],average=False)
            ntab = len(self.intables)
            byname = True
        else:
            stab = self.intables[0]
            ntab = len(self.intables)

        if len(stab.getbeamnos()) > 1:
            asaplog.post()
            asaplog.push("Mult-beam data is not supported by this module.")
            asaplog.post("ERROR")
            return

        valid_ifs = stab.getifnos()
        if self.baseif < 0:
            self.baseif = valid_ifs[0]
            asaplog.post()
            asaplog.push("IFNO is not selected. Using the first IF in the first scantable. Reference IFNO = %d" % (self.baseif))

        if not (self.baseif in valid_ifs):
            asaplog.post()
            errmsg = "IF%d does not exist in the first scantable" %  \
                     self.baseif
            raise RuntimeError, errmsg

        asaplog.push("Start selecting tables and IFNOs to solve.")
        asaplog.push("Checking frequency of the reference IF")
        unit_org = stab.get_unit()
        coord = stab._getcoordinfo()
        frame_org = coord[1]
        stab.set_unit("Hz")
        if len(self.freqframe) > 0:
            stab.set_freqframe(self.freqframe)
        stab.set_selection(ifs=[self.baseif])
        spx = stab._getabcissa()
        stab.set_selection()
        basech0 = spx[0]
        baseinc = spx[1]-spx[0]
        self.nchan = len(spx)
        # frequency tolerance
        if isinstance(self.freqtol, dict) and self.freqtol['unit'] == "Hz":
            vftol = abs(self.freqtol['value'])
        else:
            vftol = abs(baseinc * float(self.freqtol))
            self.freqtol = dict(value=vftol, unit="Hz")
        # tolerance of frequency increment
        inctol = abs(baseinc/float(self.nchan))
        asaplog.push("Reference frequency setup (Table = 0, IFNO = %d):  nchan = %d, chan0 = %f Hz, incr = %f Hz" % (self.baseif, self.nchan, basech0, baseinc))
        asaplog.push("Allowed frequency tolerance = %f Hz ( %f channels)" % (vftol, vftol/baseinc))
        poltype0 = stab.poltype()

        self.tables = []
        self.signalShift = []
        if self.dsbmode:
            self.imageShift = []

        for itab in range(ntab):
            asaplog.push("Table %d:" % itab)
            tab_selected = False
            if itab > 0:
                if byname:
                    stab = scantable(self.intables[itab],average=False)
                else:
                    stab = self.intables[itab]
                unit_org = stab.get_unit()
                coord = stab._getcoordinfo()
                frame_org = coord[1]
                stab.set_unit("Hz")
                if len(self.freqframe) > 0:
                    stab.set_freqframe(self.freqframe)

            # Check POLTYPE should be equal to the first table.
            if stab.poltype() != poltype0:
                asaplog.post()
                raise Exception, "POLTYPE should be equal to the first table."
            # Multiple beam data may not handled properly
            if len(stab.getbeamnos()) > 1:
                asaplog.post()
                asaplog.push("table contains multiple beams. It may not be handled properly.")
                asaplog.push("WARN")

            for ifno in stab.getifnos():
                stab.set_selection(ifs=[ifno])
                spx = stab._getabcissa()
                if (len(spx) != self.nchan) or \
                   (abs(spx[0]-basech0) > vftol) or \
                   (abs(spx[1]-spx[0]-baseinc) > inctol):
                    continue
                tab_selected = True
                seltab = stab.copy()
                seltab.set_unit(unit_org)
                seltab.set_freqframe(frame_org)
                self.tables.append(seltab)
                self.signalShift.append((spx[0]-basech0)/baseinc)
                if self.dsbmode:
                    self.imageShift.append(-self.signalShift[-1])
                asaplog.push("- IF%d selected: sideband shift = %16.12e channels" % (ifno, self.signalShift[-1]))
            stab.set_selection()
            stab.set_unit(unit_org)
            stab.set_freqframe(frame_org)
            if not tab_selected:
                asaplog.post()
                asaplog.push("No data selected in Table %d" % itab)
                asaplog.post("WARN")

        asaplog.push("Total number of IFs selected = %d" % len(self.tables))
        if len(self.tables) < 2:
            asaplog.post()
            raise RuntimeError, "At least 2 IFs are necessary for convolution!"

        if not self.dsbmode and len(self.imageShift) != len(self.signalShift):
            asaplog.post()
            errmsg = "User defined channel shift of image sideband has %d elements, while selected IFNOs are %d" % (len(self.imageShift), len(self.signalShift))
            errmsg += "\nThe frequency tolerance (freqtol) you set may be too small."
            raise RuntimeError, errmsg

        self.signalShift = numpy.array(self.signalShift)
        self.imageShift = numpy.array(self.imageShift)
        self.nshift = len(self.tables)

    @asaplog_post_dec
    def _preprocess_tables(self):
        ### temporary method to preprocess data
        ### Do time averaging for now.
        for itab in range(len(self.tables)):
            self.tables[itab] = self.tables[itab].average_time(scanav=False, weight="tintsys")

#     @asaplog_post_dec
#     def _setup_image_freq(self, table):
#         # User defined coordinate
#         # Get from associated MS
#         # Get from ASDM
#         lo1 = -1.
#         if self.lo1 > 0.:
#             asaplog.push("Using user defined LO1 frequency %e16.12 [Hz]" % self.lo1)
#             lo1 = self.lo1
#         else:
#             print "NOT IMPLEMENTED YET!!!"

#     def save(self, outfile, outform="ASAP", overwrite=False):
#         if not overwrite and os.path.exists(outfile):
#             raise RuntimeError, "Output file '%s' already exists" % outfile
# 
#         #self._separator._save(outfile, outform)

#     def done(self):
#         self.close()

#     def close(self):
#         pass
#         #del self._separator
    


########################################################################
    def _Deconvolution(self, data_array, shift, threshold=0.00000001):
        FObs = []
        Reject = 0
        nshift, nchan = data_array.shape
        nspec = nshift*(nshift-1)/2
        ifftObs  = numpy.zeros((nspec, nchan), numpy.float)
        for i in range(nshift):
           F = FFT.fft(data_array[i])
           FObs.append(F)
        z = 0
        for i in range(nshift):
            for j in range(i+1, nshift):
                Fobs = (FObs[i]+FObs[j])/2.0
                dX = (shift[j]-shift[i])*2.0*math.pi/float(self.nchan)
                #print 'dX,i,j=',dX,i,j
                for k in range(1,self.nchan):
                    if math.fabs(math.sin(dX*k)) > threshold:
                        Fobs[k] += ((FObs[i][k]-FObs[j][k])/2.0/(1.0-math.cos(dX*k))*math.sin(dX*k))*1.0j
                    else: Reject += 1
                ifftObs[z] = FFT.ifft(Fobs)
                z += 1
        print 'Threshold=%s Reject=%d' % (threshold, Reject)
        return ifftObs

    def _combineResult(self, ifftObs):
        nspec = len(ifftObs)
        sum = ifftObs[0]
        for i in range(1,nspec):
            sum += ifftObs[i]
        return(sum/float(nspec))

    def _subtractOtherSide(self, data_array, shift, Data):
        sum = numpy.zeros(len(Data), numpy.float)
        numSP = len(data_array)
        for i in range(numSP):
            SPsub = data_array[i] - Data
            sum += self._shiftSpectrum(SPsub, -shift[i])
        return(sum/float(numSP))

    def _shiftSpectrum(self, data, Shift):
        Out = numpy.zeros(self.nchan, numpy.float)
        w2 = Shift % 1
        w1 = 1.0 - w2
        for i in range(self.nchan):
            c1 = int((Shift + i) % self.nchan)
            c2 = (c1 + 1) % self.nchan
            Out[c1] += data[i] * w1
            Out[c2] += data[i] * w2
        return Out.copy()