from cmu_graphics import *
from pydub import AudioSegment
from pydub import effects

### THIS IS UGLY CODE, DON'T LOOK TOO CLOSELY AT IT
# There are little to no safeguards for code, so this code will more often
# than not crash if something is inputted slightly wrong

# To use, press buttons for the function you would like to use

# AudioSegment objects are stored in the top left
# To select different audio segment objects, use the right and left arrow keys

# Parameters for functions are displayed in the top right
# To type to different parameters, use the up and down arrow keys
# To execute a change, press the green button in the bottom right

# Executing a change creates a new AudioSegment object, usually with a long and
# ugly name based on the function that created it

def screen_onAppStart(app):
    app.audioList = []
    app.audioToName = dict()
    app.currentAudio = 0
    app.selectedEffect = None
    app.selectedInfo = 0
    app.userInput = ['', '', '', '']
    app.errorMessage = None
    pass

def screen_onKeyPress(app, key):
    twoArg = [1, 8]
    threeArg = [2, 4]
    if (key == 'up'): 
        if (app.selectedEffect in twoArg):
            if app.selectedInfo == 0: app.selectedInfo =1
            else:
                app.selectedInfo -= 1
        elif (app.selectedEffect in threeArg):
            if (app.selectedInfo == 0): app.selectedInfo = 2
            else:
                app.selectedInfo -= 1
    elif (key == 'down'): 
        if (app.selectedEffect in twoArg):
            if app.selectedInfo == 1: app.selectedInfo =0
            else:
                app.selectedInfo += 1
        elif (app.selectedEffect in threeArg):
            if (app.selectedInfo == 2): app.selectedInfo = 0
            else:
                app.selectedInfo += 1
    elif (key == 'left'): 
        if app.currentAudio > 1: app.currentAudio-= 1
        else:
            app.currentAudio = len(app.audioList)
    elif (key == 'right'): 
        if app.currentAudio == 0: app.currentAudio = 0
        elif (app.currentAudio == len(app.audioList)): app.currentAudio = 1
        else:
            app.currentAudio += 1
    elif key == 'backspace':
        if len(app.userInput) > 0:
            app.userInput[app.selectedInfo] = app.userInput[app.selectedInfo][:-1]
    elif key.isdigit() or key.isalpha() or (key in "+./-_"):
        app.userInput[app.selectedInfo] += key
    pass

def screen_onMousePress(app, mouseX, mouseY):
    if 450<=mouseX<=550 and 390<=mouseY<=410:
        execute(app)
    elif mouseY > 490 and mouseY < 510:
        if mouseX > 50 and mouseX < 550:
            button = (mouseX-50) // 100
            if (button != app.selectedEffect):
                app.selectedEffect = button
                app.userInput = ['', '', '', '']
                app.selectedInfo = 0
    elif mouseY > 540 and mouseY < 560:
        if mouseX > 50 and mouseX < 550:
            button = 5 + (mouseX - 50) // 100
            if (button != app.selectedEffect):
                app.selectedEffect = button
                app.userInput = ['', '', '', '']
                app.selectedInfo = 0
    else:
        app.selectedEffect = None
        app.userInput = ['', '', '', '']
        app.selectedInfo = 0
    pass

def screen_redrawAll(app):
    drawLabel(f'{app.currentAudio} of {len(app.audioList)}', 50, 20)
    drawLabel("Audio Files", 50, 75)
    if (app.currentAudio != 0):
        drawLabel(f'Audio duration: {app.audioList[app.currentAudio-1].duration_seconds} seconds',
                  20, 35, align = 'left')
    for i in range(len(app.audioList)):
        color = 'black'
        if i == app.currentAudio - 1:
            color = 'blue'
        drawLabel(app.audioToName[app.audioList[i]], 50, 100 + 25 * i, fill = color)
    if app.errorMessage != None:
        drawLabel(app.errorMessage, 300, 300, fill = 'red', size = 25)
    labels = ["Import Audio", "Add Audio", "Layer Audio", "Change Volume",
              "Fade Audio", "Reverse Audio", "Change Speed",
              "Adjust Audio Pitch", "Clip Audio", "Export Audio"]
    
    for i in range(len(labels)):
        if i == app.selectedEffect:
            color = 'lightBlue'
        else:
            color = "white"
        if (i < len(labels)/2):
            y = 500
        else:
            y = 550
        x = i % (len(labels)/2)
        drawRect(100 + 100* x, y, 100, 20, fill = color, border = 'black',
                 align = 'center')
        drawLabel(labels[i], 100 + 100* x, y)
    color1 = 'black'
    color2 = 'black'
    color3 = 'black'
    if app.selectedInfo == 0:
        color1 = 'blue'
    elif app.selectedInfo == 1:
        color2 = 'blue'
    elif app.selectedInfo == 2:
        color3 = 'blue'
    if app.selectedEffect != None: 
        drawLabel("Parameters", 500, 50)
        drawRect(500, 400, 100, 20, align = "center", fill = "lightGreen")
        drawLabel("Execute Change", 500, 400)
    if app.selectedEffect == 0:
        drawLabel("File Name: ", 500, 75, fill = color1)
        drawLabel(f'{app.userInput[0]}', 500, 90, fill = color1)
    elif app.selectedEffect == 1:
        drawLabel("Audio 1: ", 500, 75, fill =color1)
        drawLabel(f'{app.userInput[0]}', 500, 90, fill = color1)
        drawLabel("Audio 2: ", 500, 110, fill = color2)
        drawLabel(f'{app.userInput[1]}', 500, 125, fill = color2)
    elif app.selectedEffect == 2:
        drawLabel("Audio 1: ", 500, 75, fill = color1)
        drawLabel(f'{app.userInput[0]}', 500, 90, fill = color1)
        drawLabel("Audio 2: ", 500, 110, fill = color2)
        drawLabel(f'{app.userInput[1]}', 500, 125, fill = color2)
        drawLabel("Offset: ", 500, 145, fill = color3)
        drawLabel(f'{app.userInput[2]}', 500, 160, fill = color3)
    elif app.selectedEffect == 3:
        drawLabel("Volume Change: ", 500, 75, fill = color1)
        drawLabel(f'{app.userInput[0]}', 500, 90, fill = color1)
    elif app.selectedEffect == 4:
        drawLabel("Fade Change: ", 500, 75, fill = color1)
        drawLabel(f'{app.userInput[0]}', 500, 90, fill = color1)
        drawLabel("Start Time: ", 500, 110, fill = color2)
        drawLabel(f'{app.userInput[1]}', 500, 125, fill = color2)
        drawLabel("End Time: ", 500, 145, fill = color3)
        drawLabel(f'{app.userInput[2]}', 500, 160, fill = color3)
    elif app.selectedEffect == 5:
        drawLabel("Audio: ", 500, 75, fill = color1)
        drawLabel(f'{app.userInput[0]}', 500, 90, fill = color1)
    elif app.selectedEffect == 6:
        drawLabel("Change Speed By: ", 500, 75, fill = color1)
        drawLabel(f'{app.userInput[0]}', 500, 90, fill = color1)
    elif app.selectedEffect == 7:
        drawLabel("Half Step Difference: ", 500, 75, fill = color1)
        drawLabel(f'{app.userInput[0]}', 500, 90, fill = color1)
    elif app.selectedEffect == 8:
        drawLabel("Clip Start: ", 500, 75, fill = color1)
        drawLabel(f'{app.userInput[0]}', 500, 90, fill = color1)
        drawLabel("Clip End: ", 500, 110, fill = color2)
        drawLabel(f'{app.userInput[1]}', 500, 125, fill = color2)
    elif app.selectedEffect == 9:
        drawLabel("File Name: ", 500, 75, fill = color1)
        drawLabel(f'{app.userInput[0]}', 500, 90, fill = color1)
    pass

def execute(app):
    if app.selectedEffect == None: pass
    elif app.selectedEffect == 0:
        importFileExecute(app)
    elif app.selectedEffect == 1:
        addAudioExecute(app)
    elif app.selectedEffect == 2:
        layerAudioExecute(app)
    elif app.selectedEffect == 3:
        changeVolExecute(app)
    elif app.selectedEffect == 4:
        fadeExecute(app)
    elif app.selectedEffect == 5:
        reverseExecute(app)
    elif app.selectedEffect == 6:
        changeSpeedExecute(app)
    elif app.selectedEffect == 7:
        changePitchExecute(app)
    elif app.selectedEffect == 8:
        clipExecute(app)
    elif app.selectedEffect == 9:
        exportExecute(app)
    app.selectedEffect = None
    app.selectedInfo = 0
    pass

def searchAudio(app, soundName):
    for audio in app.audioToName:
        if app.audioToName[audio] == soundName:
            return audio
    return AudioSegment.empty()

def importFileExecute(app):
    fileName = app.userInput[0]
    audio = fileToAudioSeg(fileName)
    app.audioList.append(audio)
    app.currentAudio = len(app.audioList)
    app.audioToName[audio] = fileName.split(".")[0]

def addAudioExecute(app):
    audioName1 = app.userInput[0]
    audioName2 = app.userInput[1]
    audio1 = searchAudio(app, audioName1)
    audio2 = searchAudio(app, audioName2)
    newAudio = addAudio(audio1, audio2)
    app.audioList.append(newAudio)
    app.audioToName[newAudio] = f'{audioName1}+{audioName2}'
    pass

def layerAudioExecute(app):
    audioName1 = app.userInput[0]
    audioName2 = app.userInput[1]
    offset = int(app.userInput[2])
    audio1 = searchAudio(app, audioName1)
    audio2 = searchAudio(app, audioName2)
    newAudio = layerAudio(audio1, audio2, offset)
    app.audioList.append(newAudio)
    app.audioToName[newAudio] = f'{audioName2}on{audioName1}at{offset}'
    pass

def changeVolExecute(app):
    dbChange = app.userInput[0]
    if dbChange[0] == "-":
        dbChangeInt = -1
    elif dbChange[0] == "+":
        dbChangeInt = 1
    try:
        dbChangeInt *= int(dbChange[1:])
        sound = app.audioList[app.currentAudio-1]
        newAudio = changeVolume(sound, dbChangeInt)
        app.audioList.append(newAudio)
        prevAudioName = app.audioToName[app.audioList[app.currentAudio-1]]
        app.audioToName[newAudio] = f'{prevAudioName}{dbChange}dB'
    except:
        app.errorMessage = "Input the sign of decibel gain followed by the magnitude!"
    pass

def fadeExecute(app):
    dVolStr = app.userInput[0]
    if dVolStr[0] == '+':
        dVol = 1
    elif dVolStr[0] == '-':
        dVol = -1
    try:
        dVol *= int(dVolStr[1:])
        startTime = 1000 * int(app.userInput[1])
        endTime = 1000 * int(app.userInput[2])
        oldAudio = app.audioList[app.currentAudio-1]
        newAudio = oldAudio.fade(to_gain=dVol, start=startTime, duration=endTime-startTime)
        app.audioList.append(newAudio)
        prevAudioName = app.audioToName[app.audioList[app.currentAudio-1]]
        app.audioToName[newAudio] = f'{prevAudioName}-fade{dVol}from{startTime/1000}to{endTime/1000}'
    except:
        app.errorMessage = '''Remember to type a sign for the dB gain and only input integers for \nthe start and end times!'''
    pass

def reverseExecute(app):
    oldAudio = app.audioList[app.currentAudio-1]
    newAudio = reverse(oldAudio)
    app.audioList.append(newAudio)
    prevAudioName = app.audioToName[app.audioList[app.currentAudio-1]]
    app.audioToName[newAudio] = f'{prevAudioName}-reverse'
    pass

def changeSpeedExecute(app):
    oldAudio = app.audioList[app.currentAudio-1]
    speed = float(app.userInput[0])
    newAudio = speedup(oldAudio, speed)
    app.audioList.append(newAudio)
    prevAudioName = app.audioToName[app.audioList[app.currentAudio-1]]
    app.audioToName[newAudio] = f'{prevAudioName}x{speed}'
    pass

def changePitchExecute(app):
    keyStr = app.userInput[0]
    if keyStr[0] == '+':
        keyChange = 1
    elif keyStr[0] == '-':
        keyChange = -1
    try:
        keyChange *= int(keyStr[1:])
        oldAudio = app.audioList[app.currentAudio-1]
        newAudio = changePitch(oldAudio, keyChange)
        app.audioList.append(newAudio)
        prevAudioName = app.audioToName[app.audioList[app.currentAudio-1]]
        app.audioToName[newAudio] = f'{prevAudioName}{keyStr}halfsteps'
    except:
        app.errorMessage = "input the sign in front of the half step change!"
    pass

def clipExecute(app):
    startStr = app.userInput[0]
    endStr = app.userInput[1]
    try:
        startTime = 1000 * float(startStr)
        endTime = 1000 * float(endStr)
        oldAudio = app.audioList[app.currentAudio-1]
        newAudio = clip(oldAudio, startTime, endTime)
        app.audioList.append(newAudio)
        prevAudioName = app.audioToName[oldAudio]
        app.audioToName[newAudio] = f'{prevAudioName}-from{startTime/1000}to{endTime/1000}'
    except:
        app.errorMessage = '''Remember to only input integers for the start and end times!'''
    pass

def exportExecute(app):
    sound = app.audioList[app.currentAudio-1]
    fileName = app.userInput[0]
    if fileName == '': 
        sound.export(f'{app.audioToName[sound]}.mp3', format = 'mp3')
    elif '.mp3' in fileName:
        sound.export(f'{fileName}', format = 'mp3')
    elif '.wav' in fileName:
        sound.export(f'{fileName}', format = 'wav')
    else:
        app.errorMessage = "Assign this file a type (mp3 or wav)!"
    pass

def fileToAudioSeg(fileName):
    # turns an mp3 or wav file into an AudioSegment object
    if ".mp3" in fileName:
        return AudioSegment.from_mp3(fileName)
    else:
        return AudioSegment.from_wav(fileName)

def addAudio(sound1, sound2):
    # adds sound2 to the end of sound1
    return sound1 + sound2

def layerAudio(sound1, sound2, offset):
    # layers sound1 onto sound2 <offset> milliseconds after sound1 begins
    return sound1.overlay(sound2, position = offset)

def changeVolume(sound, gain):
    # increases the volume of sound by <gain> dB
    # positive means louder, negative means quieter
    return sound + gain

def duration(sound):
    # returns length of the audio file in seconds
    return sound.duration_seconds

def fade(sound, gain, start, end):
    # fades the audio from original volume (0 dB) to the to_gain dB 
    # from the <start> in milliseconds continuing for <duration> milliseconds
    return sound.fade(to_gain = gain, start = start, duration = end - start)

def reverse(sound):
    # reverses audio
    return sound.reverse()

def speedup(sound, speed):
    # speeds up audio. DOES NOT SLOW DOWN AUDIO
    return sound.speedup(playback_speed = speed)
    
def export(sound, fileName, mp3 = True):
    # exports audio as the fileName given
    return sound.export(fileName, format = 'mp3')

def changePitch(sound, halfSteps):
    # there are 12 half steps in an octave
    octaves = halfSteps/12
    # creates a new sample rate at the new frequency based on how fast
    # the new clip is supposed to be
    newRate = int(sound.frame_rate * (2**octaves))
    # uses the raw data of the sound and crunches/stretches it 
    # to change frequency in order to change the pitch
    newPitch = sound._spawn(sound.raw_data, overrides = {"frame_rate": newRate})
    # sets a new frame rate
    newPitch.set_frame_rate(44100)

    # new sound is sped up/slowed down depending on whether the pitch was 
    # shifted up or down
    return newPitch

def clip(sound, start, end):
    # returns a clip of the audio from <start> milliseconds to <end> milliseconds
    return sound[start:end]

def main():
    runAppWithScreens(initialScreen = 'screen', width=600, height = 600)

main()