# Imports from PCygni
from Utils import Utils
from LineSpectrumData import LineSpectrumData

# Imports from delivery/script
import ScriptEnvironment
import ScriptLineAnalysis
import ScriptUtilities
from FileReader import FileReader
from cassisStats import printStats

# Imports from cassis
from eu.omp.irap.cassis.gui             import PanelView
from eu.omp.irap.cassis.file            import FileReaderCassis
from eu.omp.irap.cassis.common.axes     import XAxisVelocity
from eu.omp.irap.cassis.database.access import AccessDataBase
from eu.omp.irap.cassis.common          import UtilArrayList
from eu.omp.irap.cassis.gui.plot.tools  import ToolsUtil

# Imports from Java
from java.io import File
from javax.swing import JOptionPane

# Imports from Python
import math
import os
import sys


# FUNCTIONS {{{1
def warning(msg):
	print("WARNING: %s" % msg)


def askConfirmation(msg, title="WARNING"):
	answer = JOptionPane.showConfirmDialog(None,
										msg + "\nDo you want to continue?",
										title,
										JOptionPane.YES_NO_OPTION)

	if (answer != JOptionPane.YES_OPTION):
		sys.exit("Exited before the computation.")


# COMPUTATIONS {{{2

def computeChi2Reduced(chi2s, ns, dof, nInd):
    """Return the reduced chi2
    Usage:
        chi2Reduced = computeChi2Reduced(chi2s, ns, dof, nInd)

    Inputs
    ------
    chi2s : list
        the list of chi2
    ns : list
        the list of number of points taken into account in the computation of chi2
        (one for each spectrum)
    dof : int
        the degree of freedom
    nInd : int
        the number of independent points

    Returns
    -------
    chi2Reduced : float
        the reduced chi2

    """
    nbOfSpectra = len(chi2s)
    indexes = range(nbOfSpectra)
    chi2Reduced = sum(chi2s[i] / (ns[i] / nInd - dof) for i in indexes) / nbOfSpectra
    return chi2Reduced


def computeDegreeOfFreedom(component):
    """Return the degree of freedom of the given component"""
    minima = [getattr(component, name).min for name in component.names]
    maxima = [getattr(component, name).max for name in component.names]
    extendedRanges = filter(lambda i: minima[i] != maxima[i], range(len(minima)))
    degreeOfFreedom = len(extendedRanges)
    return degreeOfFreedom


# MISC {{{2

def selectLinesWithEnoughPoints(nbOfPointsInLines, independenceFactor, nbOfIndependentPointsMin):
    """Return the lines with enough points for the computation of the reduced chi2Min

    Indeed:
        chi2MinReduced = sum(chi2s*weight/((nbOfPoints/nbOfIndepPoints)-dof)/nbOfLines)

    which means that if nbOfPoints/nbOfIndepPoints = dof, the chi2MinReduced can't be computed
    because of a division by 0 (because there are not enough points).
    This function ensures that the number of points is sufficient for
    a relevant (and a possible) computation

    Usage
    -----
    linesWithEnoughPoints = selectLinesWithEnoughPoints(nbOfPointsInLines,
                                                        independenceFactor,
                                                        nbOfIndependentPointsMin)

    Parameters
    ----------
    nbOfPointsInLines : dict
        a dictionary whose keys are the indexes of lines and the values
        the corresponding number of points
    independenceFactor : int
        the factor of independence selected by the user that
        allows to determine the number of independent points
        (nbOfIndependentPoints = nbOfPointsInLines / independenceFactor)
    nbOfIndependentPointsMin : int
        the miminum number of independent points per line
        (this number must be greater than the dof)

    Returns
    -------
    linesWithEnoughPoints : list
        the list of indexes of lines whose number of points satifisfies the criterion
        nbOfIndependentPoints > nbOfIndependentPointsMin

    """
    linesWithEnoughPoints = []
    for indexOfLine, nbOfPoints in nbOfPointsInLines.items():
        # If the criterion is satisfied and the line is not already in the
        # list, append the index of this line
        if math.floor(nbOfPoints / independenceFactor) > nbOfIndependentPointsMin and\
            indexOfLine not in linesWithEnoughPoints:
            linesWithEnoughPoints.append(indexOfLine)
    return sorted(linesWithEnoughPoints)


def combinations(iterable, r):
    # combinations('ABCD', 2) --> AB AC AD BC BD CD
    # combinations(range(4), 3) --> 012 013 023 123
    pool = tuple(iterable)
    n = len(pool)
    if r > n:
        return
    indices = range(r)
    yield tuple(pool[i] for i in indices)
    while True:
        for i in reversed(range(r)):
            if indices[i] != i + n - r:
                break
        else:
            return
        indices[i] += 1
        for j in range(i+1, r):
            indices[j] = indices[j-1] + 1
        yield tuple(pool[i] for i in indices)


def getModelIndexes(nbOfModels):
    """
    Return the values of the model every modulo N where N < nbOfModels < 10*N

    """
    values = []
    #
    for i in range(0, 10):
        min = pow(10, i)
        max = pow(10, i + 1)
        if min == 10:
            mod = 100
        else:
            mod = min / 10
        if min < nbOfModels+1 <= max:
            values = filter(lambda x: x % min == 0 or (x - min) % mod == 0, range(nbOfModels))
    return values


# CHECKS {{{2

class LinesValidityError(Exception):
	"""Raised when all indexes of lines are not valid."""
	pass

def checkAllIndexesOfLinesAreValid(inputData, linesAndRanges, name):
    invalidIndexesOfLines = filter(lambda x: x not in range(1, len(inputData)+1), linesAndRanges.keys())
    if len(invalidIndexesOfLines) > 0:
        msg = "All the lines in %s must be valid lines. Lines: %r are not correct." % (name, invalidIndexesOfLines)
        msg += "\nPlease choose your lines between %d and %d" % (1, len(inputData)+1)
        raise LinesValidityError(msg)


def checkAllIndexesOfLinesAreEqual(allLinesAndRanges):
    # Make combinations of indexes of lines grouping them 2 by 2 to compare them
    allIndexesOfLines = [linesAndRanges.keys() for linesAndRanges in allLinesAndRanges]
    combins = combinations(allIndexesOfLines, 2)
    allIndexesOfLinesAreEqual = all(sorted(c[0]) == sorted(c[1]) for c in combins)

    if not allIndexesOfLinesAreEqual:
        msg = "The lines of selectedLines, rmsLines, calLines must all be the same.\n"
        msg += "Please check your inputs. "
        raise LinesEqualityError(msg)


def checkFileExists(filename):
    if not os.path.exists(filename):
        msg = "The input filename %s does not exist." % filename
        raise ValueError(msg)
    else:
        return filename


# ANALYSIS {{{2

def analyzeSelectedLines(selectedLines, spectra):
    if len(selectedLines.keys()) == 1 and "*-*" in selectedLines:
        indexes = range(1, len(spectra) + 1)
        selectedRange = selectedLines["*-*"]
        selectedLines = dict((str(x), selectedRange) for x in indexes)
    return selectedLines


def analyzeValuesSelected(values, rangeSelected):
    """Return the values selected in the given range

    Usage:
        ------
        analyzeValuesSelected(values, rangeSelected)

    Parameters
    ----------
    values : an iterable
    all the values
    rangeSelected : iterable
    a value range [minValue, maxValue]

    Returns
    -------
    valuesSelected

    """
    if len(rangeSelected) == 0:
        return []

    valuesSelected = filter(lambda f: min(rangeSelected) <= f <= max(rangeSelected), values)
    return valuesSelected


def analyzeLines(linesAndRanges):
    newLinesAndRanges = {}
    invalidFormatMsg = "%r is not a valid input. Please use the correct format:"
    invalidFormatMsg += """
    - {'m'  : freq1}
    - {'m-n'  : freq1}
    - {'*-*'  : freq1}
    - {'m-n,p': freq1}
    - {'m,n-p': freq1}
    - {'m,p'  : freq1}"""

    # Loop through the selected lines
    for selectedLines, freqs in linesAndRanges.items():

        # Invalid format
        try:
            pair_of_lines = selectedLines.split(",")
        except AttributeError:
            raise InvalidFormatError(invalidFormatMsg % selectedLines)
        #
        # Loop through the selected ranges of lines
        for selection in selectedLines.split(","):

            # Invalid format
            try:
                pair_of_lines = selection.split("-")
            except AttributeError:
                raise InvalidFormatError(invalidFormatMsg % selectedLines)
            #
            try:
                if len(pair_of_lines) == 2:
                    range_of_lines = range(int(pair_of_lines[0]), int(pair_of_lines[1])+1)
                elif len(pair_of_lines) == 1:
                    range_of_lines = [int(pair_of_lines[0])]
                else:
                    msg = "Not a valid input."
                    raise InvalidFormatError(msg)
            except AttributeError:
                raise AttributeError(invalidFormatMsg % selectedLines)

            # Loop through the range of lines
            for line in range_of_lines:
                newLinesAndRanges[line] = freqs

    return newLinesAndRanges


def analyzeLinesAndRanges(inputData, linesAndRanges, name):
    newLinesAndRanges = analyzeLines(linesAndRanges)
    checkAllIndexesOfLinesAreValid(inputData, newLinesAndRanges, name)
    return newLinesAndRanges


# JPARAMETERS {{{1

class JParameters(object):
    """Parameters Jb, Jc, Jf, Jr, Jcr"""
    H  = 6.62606957e-34
    K  = 1.3806488e-23
    Tb = 2.7

    def __init__(self, physicalModel, nu):
        """Initialize JParameters
        Inputs 
        ------
        physicalModel : PCygniPhysicalModel
            A physical model
        nu : float
            the frequency in Hertz
        """
        self.physicalModel = physicalModel
        self.nu = nu
        self.calculateJParameters()

    def calculateJParameters(self):
        T0 = self.H * self.nu / self.K
        self.Jb  = T0 / (math.exp(T0 / self.Tb) - 1.0)
        self.Jf  = T0 / (math.exp(T0 / self.physicalModel.Tf) - 1.0)
        self.Jc  = T0 / (math.exp(T0 / self.physicalModel.Tc) - 1.0)
        self.Jr  = T0 / (math.exp(T0 / self.physicalModel.Tr) - 1.0)
        self.Jcr = self.physicalModel.Fi * self.Jc + (1.0 - self.physicalModel.Fi) * self.Jr


# PCYGNIMODEL {{{1

class PCygniModel(object):
    H  = 6.62606957e-34
    K  = 1.3806488e-23
    TB = 2.7

    def __init__(self, tuningRange, tuningBand, moltag, template,
                 eup, telescope, tmb2ta, inputFile, outputFile,
                 aijMin=0, plotTitle="", 
                 selectedLines={}, rmsLines={}, calLines={},
                 warning=True):
        self.tuningRange  = tuningRange
        self.tuningBand   = tuningBand
        self.moltag       = moltag
        self.template     = template
        self.eup          = eup
        self.telescope    = telescope
        self.tmb2ta       = tmb2ta
        self.inputFile    = checkFileExists(inputFile)
        self.outputFile   = outputFile
        self.aijMin       = aijMin
        self.lineAnalysis = ScriptLineAnalysis.LineAnalysis()
        self.spectra      = self.getSpectra()
        self.chi2Min      = 1E9
        self.chi2MinReduced    = 1E9
        self.bestPhysicalModel = None
        self.component = None
        self.plotTitle = os.path.basename(self.outputFile)
        self.chi2Mins  = []
        self.frequenciesSelected = []
        self.overSampling = 1
        self.warning = warning

        # Analyze the selectedLines
        selectedLines = analyzeSelectedLines(selectedLines, self.spectra)

        # ANALYZE AND CHECK THE LINES INDEXES AND RANGES PROVIDED
        # =======================================================
        self.selectedLines = analyzeLinesAndRanges(self.spectra, selectedLines, "selectedLines")
        self.rmsLines      = analyzeLinesAndRanges(self.spectra, rmsLines, "rmsLines")
        self.calLines      = analyzeLinesAndRanges(self.spectra, calLines, "calLines")
        checkAllIndexesOfLinesAreEqual([self.selectedLines, self.rmsLines, self.calLines])

        self.selectedNumLines = sorted(self.selectedLines.keys())
        print("Number of lines selected: %d" % len(self.selectedLines))

    def createSpecies(self):
        species = self.lineAnalysis.species(template = self.template, tag = self.moltag)
        return species

    def runLineAnalysis(self):
        """Return the data and the LineAnalysis model"""
        species = self.createSpecies()
        cassisSpectrum = FileReaderCassis.createCassisSpectrumFromFile(File(self.inputFile))
        res, model = self.lineAnalysis.run(cassisSpectrum = cassisSpectrum,
                                      tuningRange=self.tuningRange,
                                      tuningBand=self.tuningBand,
                                      plottingBand="Signal",
                                      plottingSort="Frequency",
                                      species=species, 
                                      eup=self.eup,
                                      aijMin=self.aijMin,
                                      telescope=self.telescope,
                                      tmb2ta=self.tmb2ta,
                                      overSampling=1, 
                                      rms=0.0,
                                      components = [], 
                                      warning=False)
        return res, model

    def getSpectra(self):
        """Return the spectra contained in the data"""
        res, _ = self.runLineAnalysis()
        return res.getFileSpectrumList()

    def calculateDeltaTb(self, velocity, pm, jParams):
        """Calculate the deltaTb for the given physical model"""

        tauf = pm.tau * math.exp( -0.5 * (velocity - (pm.vlsr + pm.Vin))**2 / pm.fwhm**2 )
        taur = pm.tau * math.exp( -0.5 * (velocity - (pm.vlsr - pm.Vin))**2 / pm.fwhm**2 )
        deltaTb = (jParams.Jf - jParams.Jcr) * (1.0 - math.exp(-tauf)) +\
                  (1.0 - pm.Fi) * (jParams.Jr - jParams.Jb) * (1.0 - math.exp(-taur - tauf))
        return deltaTb

    def calculateDeltaTbForSpectrum(self, numLineUser, pmIndex, overSampling=1):
        spectrum = self.spectra[numLineUser-1]
        pm = self.physicalModels[pmIndex]
        #
        lineRestFrequency = spectrum.getListOfLines().get(0).getObsFrequency()
        jParams = JParameters(pm, lineRestFrequency * 1E6)

        velocity = self.getVelocityForSpectrum(numLineUser, overSampling)

        return [self.calculateDeltaTb(v, pm, jParams) for v in velocity]

    def getVelocityForSpectrum(self, numLineUser, overSampling=1):
        spectrum = self.spectra[numLineUser-1]
        
        # Convert frequency to velocity
        lineRestFrequency = spectrum.getListOfLines().get(0).getObsFrequency()
        xaxis = XAxisVelocity.getXAxisVelocity()
        xaxis.setCoeff(spectrum.getVlsr(), lineRestFrequency)

        velocity = map(xaxis.convertFromMhzFreq, spectrum.getFrequenciesSignal())
        if overSampling != 1:
            velocity = UtilArrayList.oversample(velocity, overSampling)

        return velocity

    def getFrequencyForSpectrum(self, numLineUser, overSampling=1):
        """docstring for getFrequencyForSpectrum"""
        frequency = self.spectra[numLineUser-1].getFrequenciesSignal()
        if overSampling != 1:
            frequency = UtilArrayList.oversample(frequency, overSampling)
        return frequency
        

    def calculateChi2(self, dataIntensity, modelIntensity, rms, cal, continuum=0):
        """Return the chi2 value for a channel"""
        return (dataIntensity - modelIntensity)**2 / (rms**2 + cal**2 * (dataIntensity - continuum)**2)

    def calculateChi2ForSpectrum(self, numLineUser, physicalModelIndex):
        """chi2ForSpecrtum = calculateChi2ForSpectrum(x, yData, yModel, rms, cals, bounds)
        
        Parameters
        ----------
        x : list
            the velocity or frequency scale
        yData : list
            the intensity of the data spectrum
        yModel : list
            the intensity of the model spectrum
        rms : list
            the rms values for different regions of the spectrum
        cals : list
            the calibration values for different regions of the spectrum
        bounds : list
            the list of bounds on which the computation must be performed 
            (it determines the regions such as a region are the points within the
            min and max values of this bound)
            Example:
            bounds = [[b1, b2], [b3, b4]]
            rms    = [r1, r2]
            cals   = [c1, c2]

        """
        deltaTb       = self.calculateDeltaTbForSpectrum(numLineUser, physicalModelIndex)
        velocity      = self.getVelocityForSpectrum(numLineUser)
        spectrum      = self.spectra[numLineUser-1]
        intensity     = spectrum.getIntensities()
        rms           = self.rmsLines[numLineUser]
        cals          = self.calLines[numLineUser]
        velocityRange = self.selectedLines[numLineUser]
        #
        n    = 0
        chi2 = 0
        #
        for i in xrange(len(velocity)):
            if velocityRange[0] <= velocity[i] <= velocityRange[1]:
                chi2 += self.calculateChi2(intensity[i], deltaTb[i], rms, cals)
                n += 1

        return (chi2, n)

    def warnNoValueSelected(self, i):
        msg = "Line %d - No range selected\n" % i
        if self.warning:
            askConfirmation(msg)
        else:
            warning(msg)

    def computeBestPhysicalModel(self, component):
        self.component = component

        print("Computation of the physical models...")
        self.physicalModels = component.createPhysicalModel()
        nbOfPhysicalModels = len(self.physicalModels)
        print("Number of models: %d" % nbOfPhysicalModels)
        print("\nComputing the best physical model:")
        self.degreeOfFreedom = computeDegreeOfFreedom(component)

        # Loop through the models
        modelIndexes = getModelIndexes(nbOfPhysicalModels)
        nbsOfChannels = [s.getSize() for s in self.spectra]
        self.nbOfChannels = sum(nbsOfChannels)

        self.nbOfPointsInLines = {}
        for numLineUser in self.selectedNumLines:
            velocity = self.getVelocityForSpectrum(numLineUser)

            # TODO Add computation of the continuum

            # Determination of the lines to be used
            self.velocitiesSelected = analyzeValuesSelected(velocity, self.selectedLines[numLineUser])

            if len(self.velocitiesSelected) == 0:
                self.warnNoValueSelected(numLineUser)

            self.nbOfPointsInLines[numLineUser] = len(self.velocitiesSelected)

        # TODO Remove hardcoded number of independent points
        self.nbIndPoints = 1
        self.selectedLinesWithEnoughPoints = selectLinesWithEnoughPoints( self.nbOfPointsInLines,
                                      self.nbIndPoints,
                                      self.degreeOfFreedom + 2)

        # Each element of chi2Totals is the sum of the chi2s for all the
        # spectra for each model
        self.chi2Totals = []
        #
        for m in xrange(nbOfPhysicalModels):

            if m in modelIndexes:
                print("* model %d/%d" % (m, nbOfPhysicalModels)),

            s         = 0
            chi2Total = 0
            chi2s     = []
            if m == 0: nbOfChannelsForChi2s = []
            #
            for numLineUser in self.selectedLinesWithEnoughPoints:
                chi2, nbOfChannelsForChi2 = self.calculateChi2ForSpectrum(numLineUser, m)
                chi2Total += chi2
                chi2s.append(chi2)
                if m == 0: nbOfChannelsForChi2s.append(nbOfChannelsForChi2)
                s += 1

            self.chi2Totals.append(chi2Total)

            if m in modelIndexes:
                print("\tchi2Min = %f" % self.chi2Min)

            # Determination of the chi2Min the best model index
            if chi2Total < self.chi2Min:
                self.chi2Min   = chi2Total
                self.chi2Mins  = chi2s
                self.bestPhysicalModelIndex = m

        self.bestPhysicalModel = self.physicalModels[self.bestPhysicalModelIndex]

        # Computation of the reduced minimal chi2
        self.nbOfChannelsForChi2s = nbOfChannelsForChi2s
        self.chi2MinReduced = computeChi2Reduced(self.chi2Mins, self.nbOfChannelsForChi2s, self.degreeOfFreedom, 1)

        print()
        print("-" * 80)
        print("Initial number of models: %d" % nbOfPhysicalModels)
        print("Number of lines used: %d" % len(self.chi2Mins))
        print("Minimal chi2: %.2f" % self.chi2Min)
        print("Reduced minimal chi2: %.2f" % self.chi2MinReduced)
        print("Degree of freedom: %d" % self.degreeOfFreedom)
        print("Best parameters:", self.bestPhysicalModel)

        print()
        print("Extra information:")
        print("* Lines used: %r" % self.selectedLinesWithEnoughPoints)
        print("* Minimal chi2s per line: %r" % [float("%.2f" % chi2Min) for chi2Min in self.chi2Mins])
        print("* List of number of points for chi2s: %r" % self.nbOfChannelsForChi2s)
        print("-" * 80)

    def createHeader(self):
        header =  "chi2Min=%f\t" % self.chi2Min
        header += "Chi2MinReduced=%f\t" % self.chi2MinReduced
        header += "Dof=%d\t" % self.degreeOfFreedom
        header += "NbPoints=%d\t" % sum(self.nbOfChannelsForChi2s)
        header += "PlotTitle=%s\n" % self.plotTitle
        return header

    def getParameterNames(self):
        return [x for x in Utils.ORDERED_PARAMETER_NAMES if x in self.component.names]
        
    def getColumnNames(self):
        return "#\t%s\tchi2\n" % "\t".join(self.getParameterNames())

    def overSampleModel(self, overSampling):
        resampledFrequencies = []
        resampledVelocities  = []
        resampledDeltaTbs    = []
        #
        for i in range(len(self.spectra)):
            velocity           = self.getVelocityForSpectrum(i+1)
            resampledVelocity  = self.getVelocityForSpectrum(i+1, overSampling)
            resampledFrequency = self.getFrequencyForSpectrum(i+1, overSampling)
            resampledDeltaTb   = self.calculateDeltaTbForSpectrum(i+1, self.bestPhysicalModelIndex, overSampling)
            resampledFrequencies.append(resampledFrequency)
            resampledVelocities.append(resampledVelocity)
            resampledDeltaTbs.append(resampledDeltaTb)
        return (resampledFrequencies, resampledVelocities, resampledDeltaTbs)
        
    def writeLineSpectrumFile(self, filename, frequency, velocity, deltaTb):
        # Assumption: All the vlsrs of the spectra are identical
        vlsr = self.spectra[0].getVlsr()
        lineSpectrumData = LineSpectrumData(frequency, velocity, deltaTb, vlsr)
        lineSpectrumData.saveToFile(filename)

    # FILE WRITING {{{2

    def saveLineSpectrumFile(self, filename, overSampling=1):

        basename, extension = os.path.splitext(filename)

        if extension not in (".lis", ".fus"):
            filename += ".lis"

        print("Oversampling = %d" % overSampling)
        
        if overSampling != 1:
            frequency, velocity, deltaTb = self.overSampleModel(overSampling)
        else:
            indexesOfSpectra = range(len(self.spectra))
            frequency = [self.spectra[i].getFrequenciesSignal() for i in indexesOfSpectra]
            velocity  = [self.getVelocityForSpectrum(i+1)     for i in indexesOfSpectra]
            deltaTb   = [self.calculateDeltaTbForSpectrum(i+1, self.bestPhysicalModelIndex) for i in indexesOfSpectra]

        self.writeLineSpectrumFile(filename, frequency, velocity, deltaTb)
        print("* Saved the Line Spectrum File in %s." % filename)



    # PLOTS {{{2

    def plotBestPhysicalModel(self, overSampling=1):
        """Plot the current spectra in LineAnalysis"""
        lineSpectrumControl = self.lineAnalysis.plot() 
        lineSpectrumFilename = self.outputFile + ".lis"
        self.overSampling = overSampling

        print("Plotting the best physical model")
        self.saveLineSpectrumFile(lineSpectrumFilename, overSampling)
        self.overlay(lineSpectrumControl, lineSpectrumFilename)
        print("OK\n")

    def overlay(self, lineSpectrumControl, datafile):
        """Overlay the datafile in the current Line Spectrum"""
        if os.path.exists(datafile):
            f = File(datafile)
        else:
            msg = "%s doesn't exist" % datafile
            raise ValueError(msg)
        lineSpectrumControl.openOverlayDataFile(f)

    def plotChi2Contour(self):
        print("The Chi2Contour functionality was removed.")
