"""
Functions for reading/writing/manipulating klattgrid files

A klaatgrid can be used for speech synthesis/resynthesis.
For more information on the praat klattgrid:
http://www.fon.hum.uva.nl/praat/manual/KlattGrid.html

There are three kinds of data types in a klattgrid:
null tiers (contain no data points -- seem to function as headers for a
            set of regular tiers)
regular tiers
embedded tiers

In this code:
null tiers and regular tiers are both represented by KlattPointTier

embedded tiers contain tiers of tiers (3 layers)
A KlattContainerTier contains a list of KlattIntermediateTiers which
contains a list of KlattSubPointTiers.  Only the KlattSubPointTiers contain
any points

see **examples/klatt_resynthesis.py**
"""

import io
from os.path import join

from praatio.utilities import utils
from praatio import tgio


class _KlattBaseTier(object):
    def __init__(self, name):
        self.tierNameList = []  # Preserves the order of the tiers
        self.tierDict = {}
        self.name = name
        self.minTimestamp = None
        self.maxTimestamp = None

    def __eq__(self, other):
        isEqual = True
        isEqual &= self.name == other.name
        isEqual &= self.minTimestamp == other.minTimestamp
        isEqual &= self.maxTimestamp == other.maxTimestamp

        isEqual &= self.tierNameList == other.tierNameList
        if isEqual:
            for tierName in self.tierNameList:
                isEqual &= self.tierDict[tierName] == other.tierDict[tierName]

        return isEqual

    def addTier(self, tier, tierIndex=None):

        if tierIndex is None:
            self.tierNameList.append(tier.name)
        else:
            self.tierNameList.insert(tierIndex, tier.name)

        assert tier.name not in list(self.tierDict.keys())
        self.tierDict[tier.name] = tier

        minV = tier.minTimestamp
        if self.minTimestamp is None or (minV is not None and minV < self.minTimestamp):
            self.minTimestamp = minV

        maxV = tier.maxTimestamp
        if self.maxTimestamp is None or (maxV is not None and maxV > self.maxTimestamp):
            self.maxTimestamp = maxV


class KlattContainerTier(_KlattBaseTier):
    """
    Contains a set of intermediate tiers
    """

    def getAsText(self):
        outputTxt = ""
        outputTxt += "%s? <exists>\n" % self.name

        try:
            self.minTimestamp = toIntOrFloat(self.minTimestamp)
            outputTxt += "xmin = %s\nxmax = %s\n" % (
                repr(self.minTimestamp),
                repr(self.maxTimestamp),
            )
        except TypeError:
            pass

        for name in self.tierNameList:
            outputTxt += self.tierDict[name].getAsText()

        return outputTxt

    def modifySubtiers(self, tierName, modFunc):
        """
        Modify values in every tier contained in the named intermediate tier
        """
        kit = self.tierDict[tierName]
        for name in kit.tierNameList:
            subpointTier = kit.tierDict[name]
            subpointTier.modifyValues(modFunc)


class KlattIntermediateTier(_KlattBaseTier):
    """
    Has many point tiers that are semantically related (e.g. formant tiers)
    """

    def getAsText(self):
        outputTxt = ""
        headerTxt = "%s: size=%d\n" % (self.name, len(self.tierNameList))

        for name in self.tierNameList:
            outputTxt += self.tierDict[name].getAsText()

        outputTxt = headerTxt + outputTxt

        return outputTxt


class KlattPointTier(tgio.TextgridTier):
    """
    A Klatt tier not contained within another tier
    """

    def __init__(self, name, entryList, minT=None, maxT=None):

        entryList = [(float(time), label) for time, label in entryList]

        # Determine the min and max timestamps
        timeList = [time for time, label in entryList]
        if minT is not None:
            timeList.append(float(minT))
        if maxT is not None:
            timeList.append(float(maxT))

        try:
            minT = min(timeList)
            maxT = max(timeList)
        except ValueError:
            raise tgio.TimelessTextgridTierException()

        super(KlattPointTier, self).__init__(name, entryList, minT, maxT)

    def modifyValues(self, modFunc):
        newEntryList = [
            (timestamp, modFunc(float(value))) for timestamp, value in self.entryList
        ]

        self.entryList = newEntryList

    def getAsText(self):
        outputList = []
        self.minTimestamp = toIntOrFloat(self.minTimestamp)
        outputList.append("%s? <exists> " % self.name)
        outputList.append("xmin = %s" % repr(self.minTimestamp))
        outputList.append("xmax = %s" % repr(self.maxTimestamp))

        if self.name not in ["phonation", "vocalTract", "coupling", "frication"]:
            outputList.append("points: size= %d" % len(self.entryList))

        for i, entry in enumerate(self.entryList):
            outputList.append("points [%d]:" % (i + 1))
            outputList.append("    number = %s" % repr(entry[0]))
            outputList.append("    value = %s" % repr(entry[1]))

        return "\n".join(outputList) + "\n"


class KlattSubPointTier(KlattPointTier):
    """
    Tiers contained in a KlattIntermediateTier
    """

    def getAsText(self):
        outputList = []
        outputList.append("%s:" % self.name)
        self.minTimestamp = toIntOrFloat(self.minTimestamp)
        outputList.append("    xmin = %s" % repr(self.minTimestamp))
        outputList.append("    xmax = %s" % repr(self.maxTimestamp))
        outputList.append("    points: size = %d" % len(self.entryList))

        for i, entry in enumerate(self.entryList):
            outputList.append("    points [%d]:" % (i + 1))
            outputList.append("        number = %s" % repr(entry[0]))
            outputList.append("        value = %s" % repr(entry[1]))

        return "\n".join(outputList) + "\n"


class Klattgrid(tgio.Textgrid):
    def save(self, fn, minimumIntervalLength=None):
        """

        minimumIntervalLength is used for compatibility with Textgrid.save()
            but it has no impact on a Klattgrid
        """

        # Header
        outputTxt = ""
        outputTxt += 'File type = "ooTextFile"\n'
        outputTxt += 'Object class = "KlattGrid"\n\n'
        self.minTimestamp = toIntOrFloat(self.minTimestamp)
        outputTxt += "xmin = %s\nxmax = %s\n" % (
            repr(self.minTimestamp),
            repr(self.maxTimestamp),
        )

        for tierName in self.tierNameList:
            outputTxt += self.tierDict[tierName].getAsText()

        outputTxt = _cleanNumericValues(outputTxt)

        with io.open(fn, "w", encoding="utf-8") as fd:
            fd.write(outputTxt)


def openKlattgrid(fnFullPath):

    try:
        with io.open(fnFullPath, "r", encoding="utf-16") as fd:
            data = fd.read()
    except UnicodeError:
        with io.open(fnFullPath, "r", encoding="utf-8") as fd:
            data = fd.read()
    data = data.replace("\r\n", "\n")

    # Right now, can only open normal klatt grid and not short ones
    kg = _openNormalKlattgrid(data)

    return kg


def wavToKlattgrid(
    praatEXE,
    inputFullPath,
    outputFullPath,
    timeStep=0.005,
    numFormants=5,
    maxFormantFreq=5500.0,
    windowLength=0.025,
    preEmphasis=50.0,
    pitchFloor=60.0,
    pitchCeiling=600.0,
    minPitch=50.0,
    subtractMean=True,
    scriptFN=None,
):
    """
    Extracts the klattgrid from a wav file

    The default values are the same as the ones used in praat
    """

    if subtractMean is True:
        subtractMean = "yes"
    else:
        subtractMean = "no"

    if scriptFN is None:
        scriptFN = join(utils.scriptsPath, "sound_to_klattgrid.praat")

    utils.runPraatScript(
        praatEXE,
        scriptFN,
        [
            inputFullPath,
            outputFullPath,
            timeStep,
            numFormants,
            maxFormantFreq,
            windowLength,
            preEmphasis,
            pitchFloor,
            pitchCeiling,
            minPitch,
            subtractMean,
        ],
    )


def resynthesize(praatEXE, wavFN, klattFN, outputWavFN, doCascade=True, scriptFN=None):

    if doCascade:
        method = "Cascade"
    else:
        method = "Parallel"

    if scriptFN is None:
        scriptFN = join(utils.scriptsPath, "resynthesize_from_klattgrid.praat")

    #  Praat crashes on exit after resynthesis with a klaatgrid
    utils.runPraatScript(praatEXE, scriptFN, [wavFN, klattFN, outputWavFN, method])


def _openNormalKlattgrid(data):

    kg = Klattgrid()

    # Toss header
    data = data.split("\n\n", 1)[1]

    # Not sure if this is needed
    startI = data.index("points")
    startI = data.index("\n", startI)

    # Find sections
    sectionIndexList = _findIndicies(data, "<exists>")

    sectionIndexList.append(-1)

    for i in range(len(sectionIndexList) - 1):
        dataTuple = _getSectionHeader(data, sectionIndexList, i)
        name, minT, maxT, sectionData, sectionTuple = dataTuple

        # Container Tier
        if name in [
            "oral_formants",
            "nasal_formants",
            "nasal_antiformants",
            "tracheal_formants",
            "tracheal_antiformants",
            "delta_formants",
            "frication_formants",
        ]:

            kct = _proccessContainerTierInput(sectionData, name)
            kg.addTier(kct)

        else:

            # Process entries if this tier has any
            entryList = _buildEntryList(sectionTuple)
            tier = KlattPointTier(name, entryList, minT, maxT)
            kg.addTier(tier)

    return kg


def _proccessContainerTierInput(sectionData, name):
    sectionData = sectionData.split("\n", 3)[-1]

    formantIndexList = _findIndicies(sectionData, "formants")

    subFilterList = [
        "bandwidths",
        "oral_formants_amplitudes",
        "nasal_formants_amplitudes",
        "tracheal_formants_amplitudes",
        "frication_formants_amplitudes",
    ]

    # Find the index of all the different data sections
    subFilterIndexList = [
        _findIndicies(sectionData, subName) for subName in subFilterList
    ]

    # 'Formant' search query finds duplicates -- remove them
    newFormantList = []
    for value in formantIndexList:
        if all([value not in subList for subList in subFilterIndexList]):
            newFormantList.append(value)
    formantIndexList = newFormantList

    # Combine regular query with formant query
    indexListOfLists = [
        formantIndexList,
    ] + subFilterIndexList

    # Flatten index list
    masterIndexList = [value for sublist in indexListOfLists for value in sublist]
    masterIndexList.sort()

    # If an index list is last, it it needs to include '-1' to capture the
    # rest of the data
    for subList in indexListOfLists:
        try:
            val = subList[-1]
        except IndexError:
            continue
        ii = masterIndexList.index(val)  # Index of the index
        try:
            subList.append(masterIndexList[ii + 1] - 1)
        except IndexError:
            subList.append(-1)

    # Build the tier structure
    kct = KlattContainerTier(name)
    for indexList in indexListOfLists:
        if indexList == []:
            continue
        tierList = []
        for j in range(len(indexList) - 1):
            try:
                tmpTuple = _getSectionHeader(sectionData, indexList, j)
            except ValueError:
                continue
            subName, subMin, subMax, _, subTuple = tmpTuple
            subName = subName[:-1]

            entryList = _buildEntryList(subTuple)
            tier = KlattSubPointTier(subName, entryList, subMin, subMax)
            tierList.append(tier)
        kit = KlattIntermediateTier(subName.split()[0])
        for tier in tierList:
            kit.addTier(tier)
        kct.addTier(kit)

    return kct


def _findIndicies(data, keyword):
    indexList = utils.findAll(data, keyword)
    indexList = [data.rfind("\n", 0, i) for i in indexList]

    return indexList


def _getSectionHeader(data, indexList, i):
    sectionStart = indexList[i]
    sectionEnd = indexList[i + 1]
    sectionData = data[sectionStart:sectionEnd].strip()
    sectionTuple = sectionData.split("\n", 4)

    subheader, minr, maxr = sectionTuple[:3]
    name = subheader.split("?")[0].strip()
    minT = float(minr.split("=")[1].strip())
    maxT = float(maxr.split("=")[1].strip())

    tail = sectionTuple[3:]

    return name, minT, maxT, sectionData, tail


def _buildEntryList(sectionTuple):
    entryList = []
    if len(sectionTuple) > 1:  # Has points
        npoints = float(sectionTuple[0].split("=")[1].strip())
        if npoints > 0:
            entryList = _processSectionData(sectionTuple[1])

    return entryList


def _processSectionData(sectionData):

    sectionData += "\n"

    startI = 0
    tupleList = []
    while True:
        try:
            startI = sectionData.index("=", startI) + 1  # Past the equal sign
        except ValueError:  # No more data
            break

        endI = sectionData.index("\n", startI)
        time = float(sectionData[startI:endI].strip())

        startI = sectionData.index("=", endI) + 1  # Just past the '=' sign
        endI = sectionData.index("\n", startI)
        value = float(sectionData[startI:endI].strip())

        startI = endI
        tupleList.append((time, value))

    return tupleList


def _cleanNumericValues(dataStr):
    dataList = dataStr.split("\n")
    newDataList = []
    for row in dataList:
        row = row.rstrip()
        try:
            assert "min" not in row and "max" not in row

            head, tail = row.split("=")
            head = head.rstrip()
            tail = tail.strip()
            try:
                row = str(int(tail))
            except ValueError:
                tail = "%s" % tail
                if float(tail) == 0:
                    tail = "0"
            row = "%s = %s" % (head, tail)
        except (ValueError, AssertionError):
            pass
        finally:
            newDataList.append(row.rstrip())

    outputTxt = "\n".join(newDataList)

    return outputTxt


def toIntOrFloat(val):
    if float(val) - float(int(val)) == 0.0:
        val = int(val)
    else:
        val = float(val)
    return val
