#-------------------------------------#
#               ReAmper               #
#   Created by Daniel Lazorcak 2023   #
#=====================================#


from PyQt5.QtCore import QThread
import numpy as np
import cv2 as cv
import my_object_tracker
import sys
from sound_proces import *
from PyQt5.QtCore import QUrl
from PyQt5.QtWidgets import QApplication, QWidget, QPushButton, QVBoxLayout, QFileDialog, QCheckBox
from PyQt5.QtMultimedia import QMediaPlayer, QMediaContent
from PyQt5.QtMultimediaWidgets import QVideoWidget
import traceback
from PyQt5.QtCore import Qt
import time

class TrackedObject:
    def __init__(self) -> None:
        self.id = 0
        self.name = ""
        self.path = []
        self.size = []
    def write_file(self):
        f = open(f"trackerLog_{self.id}.txt","w")
        f.write(f"ID: {self.id}\nName: {self.name}\nPath: \n{self.path}\nSize: \n{self.size}  ")
    
    @staticmethod
    def CalculatePath(object_id: int, detected_objects):
        path = []
        for i in range(len(detected_objects)):
            if int(detected_objects[i][1]) == object_id:
                box = detected_objects[i][3]
                x = box[0]+(box[2]-box[0])/2
                y = box[1]+(box[3]-box[1])/2
                path.append([int(detected_objects[i][0]), (x,y)])
            else:
                path.append([int(detected_objects[i][0]), (0,0)])
        return path

    @staticmethod
    def CalculateSize(object_id: int, detected_objects):
        size = []
        for i in range(len(detected_objects)):
            if int(detected_objects[i][1]) == object_id:
                box = detected_objects[i][3]
                size.append([int(detected_objects[i][0]), (box[3]-box[1])*(box[2]-box[0])])
            else:
                size.append([int(detected_objects[i][0]), 0])
        return size

    @staticmethod
    def CreateObjectMap(detected_objects):
        ids = []
        objects = []
        for i in range(len(detected_objects)):
            if int(detected_objects[i][1]) in ids:
                pass
            else:
                ids.append(detected_objects[i][1])
                temp_object = TrackedObject()
                temp_object.id = int(detected_objects[i][1])
                temp_object.name = str(detected_objects[i][2])
                temp_object.path = TrackedObject.CalculatePath(temp_object.id, detected_objects)
                temp_object.size = TrackedObject.CalculateSize(temp_object.id, detected_objects)
                objects.append(temp_object)

        return objects

def Scene(objects):
    scene_type = []
    persons = 0
    cars = 0
    variance = 0
    nature = 0

    for i in range(len(objects)):
        if objects[i].name == "person": persons += 1
        if objects[i].name == "car": cars += 1
        if objects[i].name == "tree": nature += 1
        if objects[i].name == "plant": nature += 1
        if objects[i].name == "potted plant": nature += 1

    Pvar = persons / (persons + cars + nature)
    Cvar = cars / (persons + cars + nature)
    Nvar = nature / (persons + cars + nature)
    scene_type = [Pvar, Cvar, Nvar]
    return scene_type

class MethodThread(QThread):
    def __init__(self, file_path, library, video_player, stype):
        super().__init__()
        t1 = time.time()
        detected_objects = my_object_tracker.YOLO_DeepSort_Tracking(file_path, "./Output/output.mp4") 
        t2 = time.time()
        d = t2-t1
        print(f"object tracker: {d}")
        
        t1 = time.time()
        objects = TrackedObject.CreateObjectMap(detected_objects)
        t2 = time.time()
        d = t2-t1
        print(f"Create object map: {d}")

        scene_type = Scene(objects)
        output_data= []
        for i in range(len(objects)):
            try:
                t1 = time.time()
                output_data.append(place_audio(objects[i], create_empty_audio_samples(output_info), library, output_info, stype))
                t2 = time.time()
                d = t2-t1
                print(f"audio_placing: {d}")

            except Exception as e:
                traceback.print_exc()
                print("Problem w placing object ID: "+str(objects[i].id))
        sum_audio(output_data, file_path, scene_type, library)
        video_player.set_media(f"{os.getcwd()}/synthesized.mp4")

class VideoPlayer(QWidget):
    def __init__(self):
        super().__init__()

        self.media_player = QMediaPlayer(self)
        self.video_widget = QVideoWidget(self)
        self.media_player.setVideoOutput(self.video_widget)
        self.layout = QVBoxLayout(self)
        self.layout.addWidget(self.video_widget)
        self.play_button = QPushButton('Play', self)
        self.play_button.clicked.connect(self.play)
        self.layout.addWidget(self.play_button)
        self.pause_button = QPushButton('Pause', self)
        self.pause_button.clicked.connect(self.pause)
        self.layout.addWidget(self.pause_button)
        self.stop_button = QPushButton('Stop', self)
        self.stop_button.clicked.connect(self.stop)
        self.layout.addWidget(self.stop_button)

    def play(self):
        self.media_player.play()

    def pause(self):
        self.media_player.pause()

    def stop(self):
        self.media_player.stop()

    def set_media(self, file_path):
        media_content = QMediaContent(QUrl.fromLocalFile(file_path))
        self.media_player.setMedia(media_content)


class MainWindow(QWidget):
    def __init__(self):
        super().__init__()

        self.setFixedSize(1000,700)
        self.setWindowTitle("ReAmper")

        self.library = SoundLibrary("sounds")
        self.stype = "soundlibrary/sounds/"

        self.video_player = VideoPlayer()
        self.layout = QVBoxLayout(self)
        self.layout.addWidget(self.video_player)
        self.open_button = QPushButton('Open Video', self)
        self.open_button.clicked.connect(self.open_file_dialog)
        self.layout.addWidget(self.open_button)

        self.checkbox = QCheckBox()
        self.checkbox.setChecked(False)


        self.layout.addWidget(self.checkbox)
        self.checkbox.setText("Melodic")
        self.checkbox.stateChanged.connect(self.on_checkbox_state_changed)

        self.predict_button = QPushButton("Perform Synthesis", self)
        self.predict_button.clicked.connect(self.perform_prediction)
        self.layout.addWidget(self.predict_button)      

        self.stype = "soundlibrary/sounds"

    def on_checkbox_state_changed(self, state):
            if state == Qt.Checked:
                self.stype = "soundlibrary/melodic/"
                self.library = SoundLibrary("melodic")
            else:
                self.stype = "soundlibrary/sounds/"
                self.library = SoundLibrary("sounds")

    def open_file_dialog(self):
        file_dialog = QFileDialog(self)
        file_dialog.setNameFilter('Video files (*.mp4 *.avi *.mkv *.mov)')
        file_dialog.setViewMode(QFileDialog.Detail)
        if file_dialog.exec_() == QFileDialog.Accepted:
            self.file_path = file_dialog.selectedFiles()[0]
            self.video_player.set_media(self.file_path)
            self.video_player.play()
            global output_info
            vid = cv.VideoCapture(self.file_path)
            fs = 44100
            output_info = {
                "total_frames" : int(vid.get(cv.CAP_PROP_FRAME_COUNT)),
                "framerate" : int(vid.get(cv.CAP_PROP_FPS)),
                "samplerate" : fs,
                "dimensions" : (int(vid.get(cv.CAP_PROP_FRAME_WIDTH)), int(vid.get(cv.CAP_PROP_FRAME_HEIGHT)))
            }
    
    def perform_prediction(self):
            self.thread = MethodThread(self.file_path, self.library, self.video_player, self.stype)
            self.thread.start()
            self.thread.wait(5000)
            self.thread.quit()

class SoundLibrary:
    def __init__(self, type):
        f = open("data/classes/coco.names", "r")
        names = f.read()
        names = names.split("\n")
        if(type =="melodic"):
            self.path = "soundlibrary/melodic"
        if(type =="sounds"):
            self.path = "soundlibrary/sounds"

        masterpath = os.path.dirname(__file__)
        dirpath = os.path.join(masterpath, self.path)

        self.samples = {}
        for i in names:
             self.samples[i] = []
        try:
            os.path.isdir(dirpath)
        except FileNotFoundError:
            print("File does not exist")
        else:
            dirlist = os.listdir(dirpath)
            for i in range(len(dirlist)):
                for j in range(len(names)):
                    try: 
                        dirlist[i].index(names[j])
                        self.samples[names[j]].append(dirlist[i])
                    except:
                        pass  
        self.samples["ambience"] = (["p.wav", "c.wav", "n.wav"])
        print(f"Sound Library loaded: {self.samples}")

if __name__ == "__main__":
    app = QApplication(sys.argv)
    app.setStyleSheet(" QPushButton { background-color: rgba(0, 128, 0, 0.5); color: white; font-weight: bold; border: 1px solid green; padding: 10px; }")
    window = MainWindow()
    window.show()
    app.exec_()
