'''
Tools for reading/writing .wav files built on top of the built-in wave
module.
'''

import wave
import audio_config
import struct
import math

# Try to import simpleaudio but don't complain if it doesn't exist
try:
    import simpleaudio
except Exception:
    simpleaudio = None


class AudioFormatError(Exception):
    '''An error generated when loading an audio file that doesn't match
    the settings in audio_config.py.'''
    pass


#----------------#
# Managing files #
#----------------#

def openWaveFile(sourcePath):
    '''Opens a .wav file and returns the sample source (result of
    `wave.open`) paired with the number of samples in the file, after
    making sure that the sample width, sample rate, and number of
    channels agree with the values in `audio_config.py`. Prints
    information on the file properties, and raises an AudioFormatError
    if those properties don't match the required values.'''
    print("Opening audio file: '{}'".format(sourcePath))
    source = wave.open(sourcePath, mode='rb')
    numChannels = source.getnchannels()
    sampleWidth = source.getsampwidth() # byte depth
    frameRate = source.getframerate()
    numFrames = source.getnframes()

    # Ensure 16-bit and 44100 sample Rate
    if sampleWidth != audio_config.SAMPLE_WIDTH:
        raise AudioFormatError("Audio files must be 16-bit files")

    if frameRate != audio_config.SAMPLE_RATE:
        raise AudioFormatError("Audio files must be 44100kHz")

    if numChannels != audio_config.NUM_CHANNELS:
        raise AudioFormatError(
            "Audio files must be in stereo. Mono and multichannel audio"
          + "not allowed"
        )

    print("Number of channels for source:", numChannels)
    print("Bit Depth:", sampleWidth * 8)
    print("Frame Rate:", frameRate)
    print("Number of Frames", numFrames)
    print()

    return source, numFrames


def openDestinationFile(destinationPath, source):
    '''Opens a .wav file for writing to and configures it based on the
    properties of the given sample source (a `wave.open` result).'''
    destination = wave.open(destinationPath, mode='wb')

    destination.setnchannels(source.getnchannels())
    destination.setsampwidth(source.getsampwidth())
    destination.setframerate(source.getframerate())

    return destination


#---------------------#
# Frames and channels #
#---------------------#

def getFrames(source, framePointer):
    '''Takes a source which is an open file descriptor to a wavefile and
    a frame position to read from in the wave file and reads
    FRAME_CHUNK_SIZE number of frames from the wave file.  The frames
    are returned as a tuple. Recall that .wav files are interleaved so
    the tuple will contain a sample from the left channel followed by
    the right channel if a stereo file.  Additionally, samples will be
    16-bit signed integers so will need to be scaled to +/- 1 before
    processing. Returns both the data read (as a tuple of integers) and
    the number of samples read, which may be fewer than FRAME_CHUNK_SIZE
    if the frame pointer is close to the end of the file.'''
    source.setpos(framePointer)

    # equivalent to number of samples * num of channels * num of bytes/sample
    frames = source.readframes(audio_config.FRAME_CHUNK_SIZE)

    # samples * num of channels
    numOfSamples = round(len(frames) / audio_config.SAMPLE_WIDTH)
    # TODO: Respect the sample width in this format string?!?
    fmt = "<" + "h" * numOfSamples
    return struct.unpack(fmt, frames), numOfSamples


def separateChannels(frames, numChannels):
    '''Given a sequence of frames as tuple of interleaved samples, this
    function creates a more usable sequences of frames by creating a
    list of channels. Therefore the frames are de-interleaved such that
    each channel is in its own list. The sequence of frames are returned
    as a list of channels.

    It also converts each sample from a 16-bit signed integer to a
    floating point number on [-1, 1].

    If for whatever reason the number of frames is not a multiple of the
    number of channels, then extra frames will be discarded so that the
    number of frames in each channel is the same.
    '''
    channels = []
    limit = numChannels * (len(frames) // numChannels)
    for channelNum in range(numChannels):
        channelData = [
            frames[i] / pow(2, 15)
            for i in range(limit)
            if i % numChannels == channelNum
        ]
        channels.append(channelData)
    return channels


def reconstructFrames(channels):
    '''Takes a list of channels constituting frames to be written to a
    new wave file and converts the data from samples of +/-1 to signed
    16-bit integers. Interleaves the samples into a single list.

    If the channel lengths aren't equal, the result will only be as long
    as the shortest channel.'''
    reconstructedFrames = []
    for sampleNum in range(min(len(ch) for ch in channels)):
        for channelNum in range(len(channels)):
            sample = channels[channelNum][sampleNum] * pow(2, 15)
            sample = round(sample)
            reconstructedFrames.append(sample)
    return reconstructedFrames


#---------------------------#
# Applying effect functions #
#---------------------------#

def processEffect(
    chunks,
    prevChunks,
    effect,
    channelMode="mono",
    checkLength=None
):
    '''
    Given a pair of chunks (lists or tuples of samples) and an effect to
    apply, applies that affect to the chunks. Also needs to know the
    pair of previous chunks for look-back effects.

    The channel mode should be "mono" (the default) or "stereo", and the
    first one or two parameters of the effect function will be given
    chunks to work with, based on this value. In "mono" mode, the effect
    will be applied twice independently to the chunks from each channel,
    in "stereo" mode it will be applied once and given the chunks from
    the left and right channels in that order. In stereo mode, if the
    result is a list or tuple of exactly two lists or tuples, they will
    be treated as the left and right channel results, but if it's a list
    or tuple of numbers, it will be copied to both left and right
    channels.

    Any subsequent parameters beyond the chunk inputs will be given
    individual samples from the previous chunks, with the earliest
    applicable sample being supplied as the first of those arguments. So
    for example, if the effect function has 5 arguments and the mode is
    mono, it would receive the last 4 samples of the previous chunk (in
    that channel). The last argument would be the final sample, the
    second-to-last argument the second-to-last sample, etc. In stereo
    mode, the same function would receive only 3 samples from the
    previous chunk, since the first two arguments supplied would be the
    left and right chunks.

    If a checkLength value is provided, a ValueError will be raised if
    the length of the result of the effect function isn't equal to the
    given value.

    Returns a tuple containing the processed chunks for both channels.
    '''
    effectArgCount = effect.__code__.co_argcount
    # Unpack left and right channels from frames
    leftChunk, rightChunk = chunks
    prevLeftChunk, prevRightChunk = prevChunks
    if channelMode == "mono":
        # Previous samples to supply
        if effectArgCount > 1:
            prevLeftSamples = prevLeftChunk[-(effectArgCount - 1):]
            prevRightSamples = prevRightChunk[-(effectArgCount - 1):]
        else:
            prevLeftSamples = []
            prevRightSamples = []

        # Apply effect to each channel
        resultLeft = effect(leftChunk, *prevLeftSamples)
        resultRight = effect(rightChunk, *prevRightSamples)

    else:
        # Number of previous samples to supply
        if effectArgCount > 2:
            prevLeftSamples = prevLeftChunk[-(effectArgCount - 2):]
            prevRightSamples = prevRightChunk[-(effectArgCount - 2):]
        else:
            prevLeftSamples = []
            prevRightSamples = []

        # Apply effect to both channels
        result = effect(
            leftChunk,
            rightChunk,
            *prevLeftSamples,
            *prevRightSamples
        )
        if (
            isinstance(result, (list, tuple))
        and len(result) == 2
        and isinstance(result[0], (list, tuple))
        ):
            resultLeft, resultRight = result
        else:
            # mono result copied to each channel
            resultLeft = result
            resultRight = result

    # Error checking
    if any(
        sample > 1.0 or sample < -1.0
        for chunk in [resultLeft, resultRight]
        for sample in chunk
    ):
        raise ValueError(
            "Your effect returned a sample value that was too"
          + "large (> 1.0) or too small (< -1.0)."
        )

    if (
        checkLength is not None
    and (
            len(resultLeft) != checkLength
         or len(resultRight) != checkLength
        )
    ):
        mismatched = (
            len(resultLeft)
            if len(resultLeft) != checkLength
            else len(resultRight)
        )
        raise ValueError(
            (
                "Your effect returned a list of {} samples, when it"
              + "should have returned {} samples."
            ).format(
                mismatched,
                checkLength
            )
        )
    return (resultLeft, resultRight)


def processFile(
    sourcePath,
    destinationPath,
    effect,
    channelMode="mono",
    lengthFunction=None,
    reverseOrder=False,
    localizer=None
):
    '''
    Loads data from the fie at the given source path, applies the given
    effect, and saves a processed file at the given destination path.
    Also, if the `simpleaudio` module is available, this will play the
    processed audio directly using the system's speakers; if not, it
    prints out a message indicating that `simpleaudio` could be
    found.

    The effect function should be suitable for handling by
    `processEffect`, and the channelMode argument will be passed through
    to that function. If a `lengthFunction` is supplied, its result when
    applied to the length of each chunk being processed will be used as
    the expected length for `processEffect`, resulting in a ValueError
    if there's a mismatch.

    If reverseOrder is set to True, frames from the source file will be
    processed in reverse order and written to the destination file in
    that opposite order. The individual samples within the frames will
    say in their source ordering however.

    If localizer is provided, it must be a function. If it's a function,
    it will be given the effect function, along with the start sample
    number, the number of samples in the chunk about to be processed,
    and the total number of samples in the file being processed as
    arguments. Its return value (per chunk) will be used in place of the
    original effect function.
    '''
    source, numFrames = openWaveFile(sourcePath)
    destination = openDestinationFile(destinationPath, source)
    wavedata = b""

    # necessary for "reverse"
    frameRange = range(0, numFrames, audio_config.FRAME_CHUNK_SIZE)
    if reverseOrder:
        frameRange = reversed(frameRange)

    # use silence as the precursor to the start (or successor to the end)
    prevChunks = [
        [0] * audio_config.FRAME_CHUNK_SIZE,
        [0] * audio_config.FRAME_CHUNK_SIZE
    ]
    # [(lc: n - 1), (rc: n - 1), (lc: n - 2, rc: n - 2)]

    # process each FRAME_CHUNK_SIZE window of data in the wave file
    for framePointer in frameRange:
        # Get audio_config.FRAME_CHUNK_SIZE frames from wave file
        # Last window will be <= FRAME_CHUNK_SIZE
        chunks, numOfSamples = getFrames(source, framePointer)
        if numOfSamples == 0 and numFrames != framePointer + numOfSamples:
            print(
                "Warning: Your .wav file appears to be corrupted! Are"
              + " you sure your download was completed without errors?"
            )
            print(
                (
                    "Your .wav file says it contains {} frames, but we"
                    " only found {}."
                ).format(numFrames, framePointer + numOfSamples)
            )
        splitChunks = separateChannels(
            chunks,
            audio_config.NUM_CHANNELS
        )

        # Compute our required result length if a lengthFunction was
        # supplied
        if lengthFunction is not None:
            requiredLength = lengthFunction(len(splitChunks[0]))

        # Localize our effect
        if localizer is not None:
            local = localizer(
                effect,
                framePointer,
                numOfSamples,
                numFrames
            )
        else:
            local = effect

        # Get process effect on window
        processedChunks = processEffect(
            splitChunks,
            prevChunks,
            local,
            channelMode,
            requiredLength
        )

        # Advance prevChunks
        prevChunks = splitChunks

        # Reconstruct window into list of interleaved samples
        outputFrames = reconstructFrames(processedChunks)

        # Write data
        fmt = "<" + "h" * len(outputFrames) # generally same as numOfSamples
        coded = struct.pack(fmt, *outputFrames) # * is the unpacking operator
        destination.writeframes(coded)
        if simpleaudio is not None:
            wavedata += coded

    destination.close()
    source.close()

    if simpleaudio is None:
        print(
            "Note: if you install the `simpleaudio` package via the"
          + " 'Manage packages' option in the 'Tools' menu, the"
          + " the processed result will be played directly in Python."
        )
    else:
        print("Playing the result directly using `simpleaudio`...")
        wo = simpleaudio.WaveObject(
            wavedata,
            audio_config.NUM_CHANNELS,
            audio_config.SAMPLE_WIDTH,
            audio_config.SAMPLE_RATE
        )
        wo.play().wait_done()


#--------------------------------#
# Length and localizer functions #
#--------------------------------#

def sameLength(x):
    """
    If used as a lengthFunction, requires that a filter effect's result
    is the same length as the input chunk.
    """
    return x


def halfLength(x):
    """
    If used as a lengthFunction, requires that a filter effect's result
    is the half the length of the input chunk (rounded up).
    """
    return math.ceil(x / 2)


def doubleLength(x):
    """
    If used as a lengthFunction, requires that a filter effect's result
    is the double the length of the input chunk.
    """
    return 2 * x


def progressFractionLocalizer(
    effect,
    framePointer,
    framesInChunk,
    totalFrames
):
    """
    If used as a localizer, curries the provided effect function by
    filling it its last two arguments with the fractional position in the
    file being processed of the start and end of the chunk being
    processed. These will be numbers between 0 and 1, starting with 0 at
    the beginning of the first chunk and ending with 1 at the end of the
    last chunk. If processing a file in reverse, the progress fraction
    will still be computed from the beginning of the file, not the end.

    Note that if there is 1 frame in the chunk, the start and end
    progress fractions will be the same, and if there are zero frames
    in the chunk, they'll both be set to the progress fraction for the
    frame the chunk was supposed to start on.
    """
    # Compute start and end progress fractions
    start = framePointer / totalFrames
    if framesInChunk == 0:
        # Needed otherwise end could be negative in this case
        end = start
    else:
        end = ((framePointer + framesInChunk) - 1) / totalFrames

    # Return a curried version of the effect
    return lambda *args: effect(*(args + (start, end)))
