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


import cv2 as cv
import numpy as np
from hmmlearn import hmm
from midiutil import MIDIFile
import matplotlib.pyplot as plt
import mido
import sys
from PyQt5.QtWidgets import QApplication, QWidget, QPushButton, QLineEdit, QGridLayout, QFileDialog, QMainWindow, QVBoxLayout, QHBoxLayout, QLabel
from PyQt5.QtCore import Qt
import time


def display_images(images):
    """
    Funkcia na zobrazenie obrázkov pomocou PyPlot
    """

    fig, axes = plt.subplots(1, len(images), figsize=(20, 4))
    for i in range(len(images)):
        axes[i].imshow(images[i])
    plt.show()

def set_min_distance(arr, mindist):
    """
    Funkcia pre prepočet signálového poľa podľa nastavenej minimálnej vzdialenosti medzi 1
    """

    prev_idx = -mindist - 1
    for i in range(len(arr)):
        if arr[i] == 1:
            if i - prev_idx < mindist:
                arr[i] = 0
            else:
                prev_idx = i
    return arr

def sinc_interpolation(signal, new_length, treshold):
    """
    Interpolácia signálu pomocou interpolácie sinc na zadanú novú dĺžku.
    """
    old_length = len(signal)
    t_old = np.linspace(0, 1, old_length)
    t_new = np.linspace(0, 1, new_length)
    kernel = np.sinc((t_new[:, None] - t_old[None, :]) * old_length)
    interpolated = np.dot(kernel, signal)

    interpolated = (interpolated + 1) / 2

    return (interpolated >= treshold).astype(int)


class VideoAnalyzer():
    """
    Trieda VideoAnalyzer
    Obsahuje funckie pre načítanie videa, analýzu obrazu: tresholding, region grow, watershed segmentation a edge based segmentation.
    Volanie týchto funkcií: load_video(), tresholding(), region_grow(), watershed_segmentation(), edge_based_segmentation()
    """
    def __init__(self, textfield):
        self.frames = []
        self.video_data = []
        self.audio = []
        self.textfield = textfield

    def load_video(self, path: str):
        print("Loading video...")
        cap = cv.VideoCapture(path)
        self.video_data = {
            "total_frames": int(cap.get(cv.CAP_PROP_FRAME_COUNT)), 
            "fps" : int(cap.get(cv.CAP_PROP_FPS)),
            "width" : int(cap.get(3)), 
            "height" : int(cap.get(4))}
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            self.frames.append(cv.resize(frame,(256,256)))

    def dct(self):
        self.textfield.setText("Performing DCT...")
        dct_array = []
        for i in self.frames:
            i = cv.cvtColor(i, cv.COLOR_RGB2GRAY)
            dct = cv.dct(np.float32(i))
            dct_array.append(dct)

        averaged = np.sum(dct_array, axis=0) / len(dct_array)
 
        upsum = np.sum(averaged[0:100], axis=0)
        midsum = np.sum(averaged[100:200], axis=0)
        lowsum = np.sum(averaged[200:255], axis=0)

        self.top_control = upsum / max([max(upsum), abs(min(upsum))])
        self.mid_control = midsum / max([max(midsum), abs(min(midsum))])
        self.low_control = lowsum / max([max(lowsum), abs(min(lowsum))])

    def thresholding(self, limits):
        self.textfield.setText("Thresholding...")
        out = []
        for img in self.frames:
            gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
            ret, thresh = cv.threshold(gray, limits[0], limits[1], cv.THRESH_BINARY)
            thresh = cv.applyColorMap(thresh, cv.COLORMAP_JET)
            out.append(cv.resize(thresh, (1, 8)))
        
        return out
    

    def watershed_segmentation(self):
        self.textfield.setText("Watersheding...")
        out = []
        for img in self.frames:
            gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
            ret, thresh = cv.threshold(gray, 0, 255, cv.THRESH_BINARY_INV+cv.THRESH_OTSU)
            kernel = np.ones((3,3),np.uint8)
            open = cv.morphologyEx(thresh,cv.MORPH_OPEN,kernel, iterations = 2)
            dtransform = cv.distanceTransform(open,cv.DIST_L2, 5)
            ret, sfg = cv.threshold(dtransform,0.1*dtransform.max(),255,0)
            sfg = np.uint8(sfg)
            unknown = cv.subtract(open,sfg)
            ret, markers = cv.connectedComponents(sfg)
            markers = markers+1
            markers[unknown==255] = 0
            markers = cv.watershed(img,markers)
            markers = cv.normalize(markers,markers, 0, 255, cv.NORM_MINMAX, cv.CV_8UC1)
            markers = cv.applyColorMap(markers, cv.COLORMAP_JET)
            out.append(cv.resize(markers, (1, 7)))
            # display_images([markers, cv.resize(markers, (1, 7))])
            # cv.waitKey(0)

        return out

    def edge_based_segmentation(self):
        self.textfield.setText("Edge...")
        out = []
        for img in self.frames:
            gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
            edges = cv.Canny(gray, 100, 200)
            edges = cv.applyColorMap(edges, cv.COLORMAP_JET)
            out.append(cv.resize(edges, (1, 7)))
        return out
    
    def HMM(self):
        self.textfield.setText("HMM...")
        out = []
        colors = np.array([[255, 0, 0], [0, 255, 0], [0, 0, 255], [255, 255, 0], [255, 0, 255]])
        for frame in self.frames:
            frame = cv.resize(frame, (24, 24))
            gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
            gray = gray / 255.0
            features = gray.reshape(-1, 1)
            model = hmm.GaussianHMM(n_components=7, covariance_type='diag')
            model.fit(features)
            state_sequence = model.predict(features)
            label_image = state_sequence.reshape(gray.shape)
            colored_image = np.zeros((label_image.shape[0], label_image.shape[1], 3), dtype=np.uint8)
            for i in range(colors.shape[0]):
                colored_image[label_image == i, :] = colors[i, :]
            out.append(cv.resize(colored_image, (1, 7)))

        return out

    def kmeans(self):
        self.textfield.setText("KMeans...")
        out = []
        for img in self.frames:
            w, h, d = img.shape
            image_array = img.reshape((w * h, d)).astype(np.float32)
            K = 7
            criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 10, 1.0)
            flags = cv.KMEANS_RANDOM_CENTERS
            _, labels, _ = cv.kmeans(image_array, K, None, criteria, attempts=10, flags=flags)
            label_image = np.reshape(labels, (w, h))
            o = (label_image * 50).astype(np.uint8)
            o = cv.applyColorMap(o, cv.COLORMAP_JET)
            out.append(cv.resize(o, (1, 7)))
            # display_images([o, cv.resize(o, (1, 7))])
            # cv.waitKey(0)

        return out


class Synthesis():
    """
    Trieda syntetizátora, vypočítava hudobné hodnoty zo segmentovaného tenzora a generuje MIDI súbor
    """

    def __init__(self, input_data, textfield, inlist):

        self.inlist = inlist
        self.textfield = textfield
        self.t = input_data[0]
        self.w = input_data[1]
        self.e = input_data[2]
        self.k = input_data[3]
        self.h = input_data[4]

        self.melody = []
        self.harmony = [] 
        self.rythm = []
        self.octaves = [] 

        self.scale = [36, 38, 40, 41, 43, 45, 47]
        self.intervals = [-3, -2, -1, 1, 2, 3, 4]
        self.drums = [36, 38, 42, 43, 45, 48, 51, 57]

    def create_melody_map(self):
        map = []
        for i in range(len(self.k)):
            self.k[i] = cv.cvtColor(self.k[i], cv.COLOR_RGB2GRAY)
            self.k[i] = cv.normalize(self.k[i], None, alpha=0, beta=7, norm_type=cv.NORM_MINMAX)
            self.k[i] = np.uint8(self.k[i])
            self.k[i] = self.k[i] * 127 / 7
            m = max(self.k[i])
            for j in range(len(self.k[i])):
                if self.k[i][j] == m:
                    map.append([self.scale[j], int(self.k[i][j]), i])
                    break

        self.mmap = map
        return map
        
    def create_harmony_map(self):
        map = []
        for i in range(len(self.h)):
            self.h[i] = cv.cvtColor(self.h[i], cv.COLOR_RGB2GRAY)
            self.h[i] = cv.normalize(self.h[i], None, alpha=0, beta=7, norm_type=cv.NORM_MINMAX)
            self.h[i] = np.uint8(self.h[i])
            self.h[i] = self.h[i] * 127 / 7
            for j in range(len(self.h[i])):
                map.append([self.scale[j], int(self.h[i][j]), i])
        self.hmap = map
        return map

    def create_rythm_map(self):
        map = []
        for i in range(len(self.t)):
            self.t[i] = cv.cvtColor(self.t[i], cv.COLOR_RGB2GRAY)
            self.t[i] = cv.normalize(self.t[i], None, alpha=0, beta=8, norm_type=cv.NORM_MINMAX)
            self.t[i] = np.uint8(self.t[i])
            self.t[i] = self.t[i] * 127 / 8
            for j in range(len(self.t[i])):
                map.append([self.drums[j], int(self.t[i][j]), i])
        self.rmap = map
        return map

    def create_interval_map(self):
        map = []
        for i in range(len(self.w)):
            self.w[i] = cv.cvtColor(self.w[i], cv.COLOR_RGB2GRAY)
            self.w[i] = cv.normalize(self.w[i], None, alpha=0, beta=7, norm_type=cv.NORM_MINMAX)
            self.w[i] = np.uint8(self.w[i])
            self.w[i] = self.w[i] * 127 / 7
            for j in range(len(self.w[i])):
                map.append([self.intervals[j], int(self.w[i][j]), i])
        self.imap = map
        return map
        
    def create_octave_map(self):
        map = []
        for i in range(len(self.e)):
            self.e[i] = cv.cvtColor(self.e[i], cv.COLOR_RGB2GRAY)
            self.e[i] = cv.normalize(self.e[i], None, alpha=0, beta=4, norm_type=cv.NORM_MINMAX)
            self.e[i] = np.uint8(self.e[i])
            self.e[i] = self.e[i] * 12
            for j in range(len(self.e[i])):
                map.append([int(self.e[i][j]), i])
        self.omap = map
        return map

    def create_midi_file(self, analyzer):
        self.textfield.setText("Creating MIDI File...")
        top_control = sinc_interpolation(analyzer.top_control, len(self.rmap), 0.5)
        mid_control = sinc_interpolation(analyzer.mid_control, len(self.rmap), 0.5)
        low_control = sinc_interpolation(analyzer.low_control, len(self.rmap), 0.5)
        top_control = set_min_distance(top_control, self.inlist[1])
        mid_control = set_min_distance(mid_control, self.inlist[0])
        low_control = set_min_distance(low_control, self.inlist[2])

        midi_file = mido.MidiFile(type=1)
        drum_track = mido.MidiTrack()
        piano_track = mido.MidiTrack()
        clarinet_track = mido.MidiTrack()
        midi_file.tracks.extend([drum_track, piano_track, clarinet_track])

        drum_track.append(mido.Message('program_change', program=0, time=0))
        piano_track.append(mido.Message('program_change', program=0, time=0))
        clarinet_track.append(mido.Message('program_change', program=71, time=0))

        self.midi = MIDIFile(numTracks = 3, deinterleave=False)
        tempo = mido.bpm2tempo(160)
        drum_track.append(mido.MetaMessage('set_tempo', tempo=tempo, time=0))

        for i in range(len(self.rmap)):
            drum_track.append(mido.Message('note_on', channel=9, note=int(self.rmap[i][0]), velocity=int(self.rmap[i][1]*mid_control[i]), time=self.rmap[i][2]))
        for i in range(len(self.hmap)):
            piano_track.append(mido.Message('note_on', channel=0, note=int(self.hmap[i][0]+int(self.omap[i][0])), velocity=int(self.hmap[i][1]*top_control[i]), time=self.hmap[i][2]))
        for i in range(len(self.mmap)):
            clarinet_track.append(mido.Message('note_on', channel=1, note=int(self.mmap[i][0] + self.imap[i][0]) + 24, velocity=int(self.mmap[i][1]*low_control[i]), time=self.mmap[i][2]))
        
        midi_file.save('output.mid')
        
        self.textfield.setText("Finished...")
        return midi_file


def the_sum_of_it_all(arr, interval):
    """
    Funkcia, ktorá vytvára z celkového počtu obrázkov počet menší, podľa vstupného intervalu
    """

    summed = [[], [], [], [], []]
    a = 0
    
    for i in range(len(arr)):
        j=0
        while j < len(arr[i]):
            for x in range(interval+1):
                try:
                    a += arr[i][j+x]
                except:
                    pass
            summed[i].append(a)
            a = 0
            j += interval

    return summed

def connect_images(arr):
    """
    Funkcia pre spojenie obrázkov vo vstupnej premennej
    """

    img = np.zeros((arr[0].shape[0], arr[0].shape[1]*len(arr),3))
    for i in range(len(arr)):
        x1 = i 
        x2 = i+1
        img[:, x1:x2] = arr[i]
    return img



class MyWindow(QMainWindow):
    """
    Trieda okna GUI
    """

    def __init__(self):
        super().__init__()

        self.setFixedSize(700, 400)
        self.setWindowTitle('SegMentor')

        self.var1 = 5
        self.var2 = 5
        self.var3 = 5
        self.var4 = 1

        self.num1_label = QLabel("Drum Intensity:", self)
        self.num1 = QLineEdit(str(self.var1))
        self.num1.setFixedWidth(50)
        self.num1.setMaxLength(3)

        self.num2_label = QLabel("Piano Intensity:", self)
        self.num2 = QLineEdit(str(self.var2))
        self.num2.setFixedWidth(50)
        self.num2.setMaxLength(3)

        self.num3_label = QLabel("Clarinet Intensity:", self)
        self.num3 = QLineEdit(str(self.var3))
        self.num3.setFixedWidth(50)
        self.num3.setMaxLength(3)

        self.num4_label = QLabel("Averaging:", self)
        self.num4 = QLineEdit(str(self.var4))
        self.num4.setFixedWidth(50)
        self.num4.setMaxLength(3)

        self.nbox = QWidget()
        self.inbox = QHBoxLayout(self.nbox)
        self.inbox.addWidget(self.num1_label)
        self.inbox.addWidget(self.num1)
        self.inbox.addWidget(self.num2_label)
        self.inbox.addWidget(self.num2)
        self.inbox.addWidget(self.num3_label)
        self.inbox.addWidget(self.num3)
        self.inbox.addWidget(self.num4_label)
        self.inbox.addWidget(self.num4)

        self.load_button = QPushButton('Load Video', self)
        self.load_button.setGeometry(10, 10, 100, 30)
        self.load_button.clicked.connect(self.load_video)

        self.synth_button = QPushButton('Start Synthesis', self)
        self.synth_button.setGeometry(120, 10, 100, 30)
        self.synth_button.clicked.connect(self.start_synthesis)

        self.textfield = QLineEdit(self)
        self.textfield.setGeometry(10, 50, 480, 30)
        self.textfield.setReadOnly(True)

        self.save_button = QPushButton('Save File', self)
        self.save_button.setGeometry(10, 90, 100, 30)
        self.save_button.clicked.connect(self.save_video)

        self.exit_button = QPushButton('Exit', self)
        self.exit_button.setGeometry(120, 90, 100, 30)
        self.exit_button.clicked.connect(self.exit_app)
        
        central_widget = QWidget(self)
        self.setCentralWidget(central_widget)
        self.layout = QVBoxLayout(central_widget)

        self.layout.addWidget(self.nbox)

        self.layout.addWidget(self.load_button)
        self.layout.addWidget(self.synth_button)
        self.layout.addWidget(self.textfield)
        self.layout.addWidget(self.save_button)
        self.layout.addWidget(self.exit_button)

        self.setLayout(self.layout)

    def load_video(self):
        self.filename, _ = QFileDialog.getOpenFileName(self, 'Load Video', '', 'Video Files (*.mp4 *.avi *.mov)')
        self.textfield.setText(f'Loaded video file: {self.filename}')

    def start_synthesis(self):
        DRUMS_MINDIST = int(self.num1.text())
        HARM_MINDIST = int(self.num2.text())
        MELOD_MINDIST = int(self.num3.text())
        AVERAGING = int(self.num4.text())
        self.textfield.setText(f'Starting synthesis with numbers: {self.num1.text()}, {self.num2.text()}, {self.num3.text()}')

        self.midi_file = Segmentor(self.filename, self.textfield, [DRUMS_MINDIST, HARM_MINDIST, MELOD_MINDIST, AVERAGING])

    def save_video(self):
        try:
            filename, _ = QFileDialog.getSaveFileName(self, 'Save MIDI file', '', '(*.mid)')
            self.textfield.setText(f'Saved synthesized video to file: {filename}')
            self.midi_file.save(f'{filename}')
        except:
            pass


    def exit_app(self):
        sys.exit()

def Segmentor(path, g, inlist):
    """
    Funkcia riadiaca celý proces syntézy
    """
    print("Hello, this is SegMenter, a video to audio synthesizer powered by OpenCV")
    analyzer = VideoAnalyzer(g)
    analyzer.load_video(path)
    
    st = time.time()
    analyzer.dct()
    et = time.time()
    time_dct = st - et

    st = time.time()
    t = analyzer.thresholding([200,250])
    et = time.time()
    time_T = st - et

    st = time.time()
    w = analyzer.watershed_segmentation()
    et = time.time()
    time_W = st - et

    st = time.time()
    e = analyzer.edge_based_segmentation()
    et = time.time()
    time_E = st - et

    st = time.time()
    k = analyzer.kmeans()
    et = time.time()
    time_k = st - et

    st = time.time()
    h = analyzer.HMM()
    et = time.time()
    time_h = st - et

    sum = the_sum_of_it_all([t,w,e,k,h], inlist[3])
    s = Synthesis([sum[0],sum[1],sum[2],sum[3],sum[4]], g, inlist)

    st = time.time()
    h = s.create_harmony_map()
    et = time.time()
    time_hh = st - et

    st = time.time()
    i = s.create_interval_map()
    et = time.time()
    time_ii = st - et

    st = time.time()
    m = s.create_melody_map()
    et = time.time()
    time_mm = st - et

    st = time.time()
    o = s.create_octave_map()
    et = time.time()
    time_oo = st - et

    st = time.time()
    r = s.create_rythm_map() 
    et = time.time()
    time_rr = st - et

    print(f"DCT: {time_dct}   T: {time_T} W: {time_W} E: {time_E} K: {time_k} H: {time_h} HARMONY: {time_hh} INTERVAL: {time_ii} MELODY: {time_mm} OCTAVE: {time_oo} RYTHM: {time_rr}")
    print(f"TOTAL: {time_dct+time_T+time_W+time_E+time_k+time_h+time_hh+time_ii+time_mm+time_oo+time_rr}")
    midi_file = s.create_midi_file(analyzer)
    print("Segmentation Successfull")
    return midi_file



if __name__ == '__main__':
    app = QApplication(sys.argv)
    app.setStyleSheet("QMainWindow {background-image: url('bg.png'); background-repeat: no-repeat; background-position: center;} QPushButton { background-color: rgba(50, 50, 128, 0.5); color: white; font-weight: bold; border: 2px solid black; padding: 10px;} QLabel {font-weight: bold; color:white;}}")

    win = MyWindow()
    win.show()
    sys.exit(app.exec_())