import sys, os
import numpy as np
from pathlib import Path
from pydub import AudioSegment

def generate_audio(osufn, progress=None):
    osufn = Path(osufn)

    global_offset = 13  # milliseconds

    audiofn = "audio.mp3"

    kickfn = "kick.wav"
    hihatfn = "hihat.wav"
    clapfn = "clap.wav"

    with open(osufn) as f:
        lines = f.readlines()
        start = -1
        end = -1
        for i in range(len(lines)):
            if lines[i].startswith("AudioFilename: "):
                audiofn = osufn.parent/(' '.join(lines[i].strip().split(' ')[1:]))
            if lines[i].startswith("[TimingPoints]"):
                start = i
            elif start > 0 and lines[i].startswith("["):
                end = i
                break
        if start == -1:
            print("error: invalid .osu file")
            sys.exit()
        elif end == -1:
            end = len(lines)
        timing_points = [x.strip().split(',') for x in lines[start:end]]
        for i in range(len(timing_points)-1, -1, -1):
            if len(timing_points[i]) < 8 or not int(timing_points[i][6]):  # remove invalid lines and inherited timing points
                del timing_points[i]

    kick = AudioSegment.from_file(kickfn)
    hihat = AudioSegment.from_file(hihatfn)
    clap = AudioSegment.from_file(clapfn)

    audio = AudioSegment.silent(duration=len(AudioSegment.from_file(audiofn)))

    kicks = []
    hihats = []
    claps = []
    # offset in ms, beat length, beats per measure, sound set, idk, volume, inherited, flags
    # generate beat timestamps
    l = len(timing_points)
    if l == 1:
        # for single BPM songs you could just do
        offset = float(timing_points[0][0])
        beat_length = float(timing_points[0][1])
        beats_per_measure = int(timing_points[0][2])

        measure_length = beat_length*beats_per_measure

        while offset < 0:
            offset += measure_length

        loop = AudioSegment.silent(duration=measure_length)

        progress.emit(10)

        if beats_per_measure == 4:
            loop = loop.overlay(kick, position=0).overlay(hihat, position=beat_length).overlay(clap, position=beat_length*2).overlay(hihat, position=beat_length*3)
        elif beats_per_measure == 2:
            loop = loop.overlay(kick, position=0).overlay(hihat, position=beat_length)
        else:
            loop = loop.overlay(kick, position=0)
            for i in range(beats_per_measure - 1):
                loop = loop.overlay(hihat, position=beat_length*(i+1))

        progress.emit(40)

        t = measure_length
        for measure in range(int(len(audio) / (beat_length*beats_per_measure) + 1)):
            t += measure_length
            loop += loop

            diff_ms = len(loop) - t
            if diff_ms > 0:  # loop is longer than it should be
                loop = loop[:-diff_ms]
            elif diff_ms < 0:  # loop is shorter than it should be
                loop += AudioSegment.silent(duration=diff_ms)

            # if diff_ms:
            #     print(f"{diff_ms} detected")

        progress.emit(80)

        audio = audio.overlay(loop, position=global_offset+offset)

        progress.emit(99)
    else:
        for i in range(len(timing_points)):
            if progress is not None:
                progress.emit(int(i/l*10))
            offset = float(timing_points[i][0])
            next_offset = float(timing_points[i+1][0]) if i < len(timing_points)-1 else len(audio)
            beat_length = float(timing_points[i][1])
            beats_per_measure = int(timing_points[i][2])

            # consider speeding up by pre-rendering kick-hat-clap-hat pattern for each timing_points' beat_length
            # and then just overlaying the patterns up until next_offset and then cut leftover audio after next_offset
            # consequentially this makes BPM changes sound more abrupt but it's way faster in theory and removes the external for loops

            if beats_per_measure == 4:
                kicks += list(np.arange(offset + global_offset, next_offset + global_offset, beat_length*2))
                hihats += list(np.arange(offset + beat_length/2 + global_offset, next_offset + global_offset, beat_length))
                claps += list(np.arange(offset + beat_length + global_offset, next_offset + global_offset, beat_length*2))
            elif beats_per_measure == 2:
                kicks += list(np.arange(offset + global_offset, next_offset + global_offset, beat_length))
                hihats += list(np.arange(offset + beat_length/2 + global_offset, next_offset + global_offset, beat_length))
            else:
                kicks += list(np.arange(offset + global_offset, next_offset + global_offset, beat_length*beats_per_measure))
                hihats += list(np.arange(offset + beat_length + global_offset, next_offset + global_offset, beat_length))

        i = 0
        l = len(kicks)+len(hihats)+len(claps)
        for pos in kicks:
            audio = audio.overlay(kick, position=int(pos))
            if progress is not None:
                progress.emit(10+int(i/l*90))
            i += 1
        for pos in hihats:
            audio = audio.overlay(hihat, position=int(pos))
            if progress is not None:
                progress.emit(10+int(i/l*90))
            i += 1
        for pos in claps:
            audio = audio.overlay(clap, position=int(pos))
            if progress is not None:
                progress.emit(10+int(i/l*90))
            i += 1

    os.rename(audiofn, f"{audiofn}.bak")
    audio.export(audiofn)

    if progress is not None:
        progress.emit(100)


import copy

import cv2
from PIL import Image
from open3d.cpu.pybind.geometry import PointCloud
from open3d.cpu.pybind.utility import Vector3dVector
from open3d.cpu.pybind.visualization import draw_geometries
from torch import nn
from torchvision import models

import PIL
import torch
import torchvision.transforms as T
import numpy as np
from tqdm import tqdm

from utils.misc.input import RealSense
from utils.model.wrappers import Segmentator

if __name__ == '__main__':
    model = models.segmentation.fcn_resnet101(pretrained=False)
    model.classifier[4] = nn.Conv2d(512, 2, kernel_size=(1, 1), stride=(1, 1)) # checkpoints/seg_model_f10.4585658013820648
    model.load_state_dict(torch.load('./checkpoints/sym5/epoch23'), strict=False) # ./checkpoints/sym4/sgdCE
    model.eval()
    model.cuda()

    model = Segmentator(model, device='cuda')

    vid = cv2.VideoCapture(0)
    # camera = RealSense()

    with tqdm(total=10000) as pbar:
        i = 0
        while (True):
            ret, frame = vid.read()
            # frame, _ = camera.read()

            frame = np.array(Image.open(f'data/real/data/rgb/eCub{i}_rgb.png'))
            depth = np.array(Image.open(f'data/real/data/depth/eCub{i}_depth.png')).astype(np.uint16)
            i+=1

            # depth = np.array(Image.open('test_depth.png')).astype(np.uint16)

            # frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

            segmented, categories = model(frame)

            overlay = copy.deepcopy(frame)
            if np.any(categories == 1):
                overlay[categories == 1] = np.array([0, 0, 128])
            res = cv2.addWeighted(frame, 1, overlay, 0.5, 0)

            # cv2.imshow('frame', res)
            cv2.imshow('frame', cv2.cvtColor(res, cv2.COLOR_RGB2BGR))

            depth[categories != 1] = 0
            cv2.imshow('depth', cv2.applyColorMap(cv2.convertScaleAbs(depth, alpha=0.03), cv2.COLORMAP_JET))

            pc = PointCloud()
            pc.points = Vector3dVector(RealSense.pointcloud(depth))

            draw_geometries([pc])

            pbar.update(1)
            if cv2.waitKey(0) & 0xFF == ord('q'):
                break

    vid.release()
    cv2.destroyAllWindows()


"""This script expects something like a binary sentiment data set, such as
 that available here: `http://www.cs.cornell.edu/people/pabo/movie-review-data/`

It expects a directory structure like: `data_dir/train/{pos|neg}`
 and `data_dir/test/{pos|neg}`. Put (say) 90% of the files in the former
 and the remainder in the latter.
"""

from __future__ import unicode_literals
from __future__ import print_function
from __future__ import division

from collections import defaultdict
from pathlib import Path
import numpy
import plac

import spacy.en

def read_data(nlp, data_dir):
    for subdir, label in (('pos', 1), ('neg', 0)):
        for filename in (data_dir / subdir).iterdir():
            text = filename.open().read()
            doc = nlp(text)
            if len(doc) >= 1:
                yield doc, label

def partition(examples, split_size):
    examples = list(examples)
    numpy.random.shuffle(examples)
    n_docs = len(examples)
    split = int(n_docs * split_size)
    return examples[:split], examples[split:]

def minibatch(data, bs=24):
    for i in range(0, len(data), bs):
        yield data[i:i+bs]

class Extractor(object):
    def __init__(self, nlp, vector_length, dropout=0.3):
        self.nlp = nlp
        self.dropout = dropout
        self.vector = numpy.zeros((vector_length, ))

    def doc2bow(self, doc, dropout=None):
        if dropout is None:
            dropout = self.dropout
        bow = defaultdict(int)
        all_words = defaultdict(int)
        for word in doc:
            if numpy.random.random() >= dropout and not word.is_punct:
                bow[word.lower] += 1
            all_words[word.lower] += 1
        if sum(bow.values()) >= 1:
            return bow
        else:
            return all_words

    def bow2vec(self, bow, E):
        self.vector.fill(0)
        n = 0
        for orth_id, freq in bow.items():
            self.vector += self.nlp.vocab[self.nlp.vocab.strings[orth_id]].vector * freq
            # Apply the fine-tuning we've learned
            if orth_id < E.shape[0]:
                self.vector += E[orth_id] * freq
            n += freq
        return self.vector / n

class NeuralNetwork(object):
    def __init__(self, depth, width, n_classes, n_vocab, extracter, optimizer):
        self.depth = depth
        self.width = width
        self.n_classes = n_classes
        self.weights = Params.random(depth, width, width, n_classes, n_vocab)
        self.doc2bow = extracter.doc2bow
        self.bow2vec = extracter.bow2vec
        self.optimizer = optimizer
        self._gradient = Params.zero(depth, width, width, n_classes, n_vocab)
        self._activity = numpy.zeros((depth, width))

    def train(self, batch):
        activity = self._activity
        gradient = self._gradient
        activity.fill(0)
        gradient.data.fill(0)
        loss = 0
        word_freqs = defaultdict(int)
        for doc, label in batch:
            word_ids = self.doc2bow(doc)
            vector = self.bow2vec(word_ids, self.weights.E)
            self.forward(activity, vector)
            loss += self.backprop(vector, gradient, activity, word_ids, label)
            for w, freq in word_ids.items():
                word_freqs[w] += freq
        self.optimizer(self.weights, gradient, len(batch), word_freqs)
        return loss

    def predict(self, doc):
        actv = self._activity
        actv.fill(0)
        W = self.weights.W
        b = self.weights.b
        E = self.weights.E

        vector = self.bow2vec(self.doc2bow(doc, dropout=0.0), E)
        self.forward(actv, vector)
        return numpy.argmax(softmax(actv[-1], W[-1], b[-1]))

    def forward(self, actv, in_):
        actv.fill(0)
        W = self.weights.W; b = self.weights.b
        actv[0] = relu(in_, W[0], b[0])
        for i in range(1, self.depth):
            actv[i] = relu(actv[i-1], W[i], b[i])

    def backprop(self, input_vector, gradient, activity, ids, label):
        W = self.weights.W
        b = self.weights.b

        target = numpy.zeros(self.n_classes)
        target[label] = 1.0
        pred = softmax(activity[-1], W[-1], b[-1])
        delta = pred - target

        for i in range(self.depth, 0, -1):
            gradient.b[i] += delta
            gradient.W[i] += numpy.outer(delta, activity[i-1])
            delta = d_relu(activity[i-1]) * W[i].T.dot(delta)

        gradient.b[0] += delta
        gradient.W[0] += numpy.outer(delta, input_vector)
        tuning = W[0].T.dot(delta).reshape((self.width,)) / len(ids)
        for w, freq in ids.items():
            if w < gradient.E.shape[0]:
                gradient.E[w] += tuning * freq
        return -sum(target * numpy.log(pred))

def softmax(actvn, W, b):
    w = W.dot(actvn) + b
    ew = numpy.exp(w - max(w))
    return (ew / sum(ew)).ravel()

def relu(actvn, W, b):
    x = W.dot(actvn) + b
    return x * (x > 0)

def d_relu(x):
    return x > 0

class Adagrad(object):
    def __init__(self, lr, rho):
        self.eps = 1e-3
        # initial learning rate
        self.learning_rate = lr
        self.rho = rho
        # stores sum of squared gradients
        #self.h = numpy.zeros(self.dim)
        #self._curr_rate = numpy.zeros(self.h.shape)
        self.h = None
        self._curr_rate = None

    def __call__(self, weights, gradient, batch_size, word_freqs):
        if self.h is None:
            self.h = numpy.zeros(gradient.data.shape)
            self._curr_rate = numpy.zeros(gradient.data.shape)
        self.L2_penalty(gradient, weights, word_freqs)
        update = self.rescale(gradient.data / batch_size)
        weights.data -= update

    def rescale(self, gradient):
        if self.h is None:
            self.h = numpy.zeros(gradient.data.shape)
            self._curr_rate = numpy.zeros(gradient.data.shape)
        self._curr_rate.fill(0)
        self.h += gradient ** 2
        self._curr_rate = self.learning_rate / (numpy.sqrt(self.h) + self.eps)
        return self._curr_rate * gradient

    def L2_penalty(self, gradient, weights, word_freqs):
        # L2 Regularization
        for i in range(len(weights.W)):
            gradient.W[i] += weights.W[i] * self.rho
            gradient.b[i] += weights.b[i] * self.rho
        for w, freq in word_freqs.items():
            if w < gradient.E.shape[0]:
                gradient.E[w] += weights.E[w] * self.rho

class Params(object):
    @classmethod
    def zero(cls, depth, n_embed, n_hidden, n_labels, n_vocab):
        return cls(depth, n_embed, n_hidden, n_labels, n_vocab, lambda x: numpy.zeros((x,)))

    @classmethod
    def random(cls, depth, nE, nH, nL, nV):
        return cls(depth, nE, nH, nL, nV, lambda x: (numpy.random.rand(x) * 2 - 1) * 0.08)

    def __init__(self, depth, n_embed, n_hidden, n_labels, n_vocab, initializer):
        nE = n_embed; nH = n_hidden; nL = n_labels; nV = n_vocab
        n_weights = sum([
            (nE * nH) + nH,
            (nH * nH  + nH) * depth,
            (nH * nL) + nL,
            (nV * nE)
        ])
        self.data = initializer(n_weights)
        self.W = []
        self.b = []
        i = self._add_layer(0, nE, nH)
        for _ in range(1, depth):
            i = self._add_layer(i, nH, nH)
        i = self._add_layer(i, nL, nH)
        self.E = self.data[i : i + (nV * nE)].reshape((nV, nE))
        self.E.fill(0)

    def _add_layer(self, start, x, y):
        end = start + (x * y)
        self.W.append(self.data[start : end].reshape((x, y)))
        self.b.append(self.data[end : end + x].reshape((x, )))
        return end + x

@plac.annotations(
    data_dir=("Data directory", "positional", None, Path),
    n_iter=("Number of iterations (epochs)", "option", "i", int),
    width=("Size of hidden layers", "option", "H", int),
    depth=("Depth", "option", "d", int),
    dropout=("Drop-out rate", "option", "r", float),
    rho=("Regularization penalty", "option", "p", float),
    eta=("Learning rate", "option", "e", float),
    batch_size=("Batch size", "option", "b", int),
    vocab_size=("Number of words to fine-tune", "option", "w", int),
)
def main(data_dir, depth=3, width=300, n_iter=5, vocab_size=40000,
         batch_size=24, dropout=0.3, rho=1e-5, eta=0.005):
    n_classes = 2
    print("Loading")
    nlp = spacy.en.English(parser=False)
    train_data, dev_data = partition(read_data(nlp, data_dir / 'train'), 0.8)
    print("Begin training")
    extracter = Extractor(nlp, width, dropout=0.3)
    optimizer = Adagrad(eta, rho)
    model = NeuralNetwork(depth, width, n_classes, vocab_size, extracter, optimizer)
    prev_best = 0
    best_weights = None
    for epoch in range(n_iter):
        numpy.random.shuffle(train_data)
        train_loss = 0.0
        for batch in minibatch(train_data, bs=batch_size):
            train_loss += model.train(batch)
        n_correct = sum(model.predict(x) == y for x, y in dev_data)
        print(epoch, train_loss, n_correct / len(dev_data))
        if n_correct >= prev_best:
            best_weights = model.weights.data.copy()
            prev_best = n_correct

    model.weights.data = best_weights
    print("Evaluating")
    eval_data = list(read_data(nlp, data_dir / 'test'))
    n_correct = sum(model.predict(x) == y for x, y in eval_data)
    print(n_correct / len(eval_data))

if __name__ == '__main__':
    #import cProfile
    #import pstats
    #cProfile.runctx("main(Path('data/aclImdb'))", globals(), locals(), "Profile.prof")
    #s = pstats.Stats("Profile.prof")
    #s.strip_dirs().sort_stats("time").print_stats(100)
    plac.call(main)


import argparse
import math
import json
import numpy as np
import os
import torch

import librosa

from hparams import hparams as hp
from utils.dsp import load_wav
from utils.dsp import melspectrogram
from model_vc import Generator
from synthesis import build_model
from synthesis import wavegen

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='')
    parser.add_argument('--output', type=str, required=True, help='path to output wav(./output.wav)')
    parser.add_argument('--src-wav', type=str, required=True, help='path to src wav(./data/test/[speaker]/[filename]/0000.wav')
    parser.add_argument('--src-emb', type=str, required=True, help='path to src wav(./data/test/[speaker]/emb.npy')
    parser.add_argument('--tgt-emb', type=str, required=True, help='path to src wav(./data/test/[speaker]/emb.npy')
    parser.add_argument('--vocoder', type=str, required=True, help='path to checkpoint_step001000000_ema.pth')
    parser.add_argument('--autovc', type=str, required=True, help='checkpoints/checkpoint_step000600.pth')
    args = parser.parse_args()

    output_path = args.output
    src_wav_path = args.src_wav
    src_emb_path = args.src_emb
    tgt_emb_path = args.tgt_emb
    vocoder_checkpoint_path = args.vocoder
    autovc_checkpoint_path = args.autovc

    dim_neck = 32
    dim_emb = 256
    dim_pre = 512
    freq = 32

    device = torch.device('cpu')
    wavnet = build_model().to(device)
    checkpoint = torch.load(vocoder_checkpoint_path, map_location=device)
    wavnet.load_state_dict(checkpoint["state_dict"])

    wav = load_wav(src_wav_path)
    emb = np.load(src_emb_path)
    emb_tgt = np.load(tgt_emb_path)

    mel = melspectrogram(wav)

    pad_len = math.ceil(mel.shape[1] / 32) * 32 - mel.shape[1]
    mel = np.pad(mel, ((0,0), (0, pad_len)), mode='constant')

    mel = torch.FloatTensor(mel)
    emb = torch.FloatTensor(emb)
    emb_tgt = torch.FloatTensor(emb_tgt)

    model = Generator(dim_neck, dim_emb, dim_pre, freq)

    checkpoint = torch.load(autovc_checkpoint_path, map_location=torch.device('cpu'))
    model.load_state_dict(checkpoint['model'])
    model.eval()

    x = mel.unsqueeze(0).transpose(2,1)
    e = emb.unsqueeze(0)
    et = emb_tgt.unsqueeze(0)

    mel_outputs, mel_outputs_postnet, codes = model(x, e, et)
    mel_rec = mel_outputs_postnet.transpose(2,1).cpu().detach().numpy()[0]

    mel_rec = mel_rec[:,:-pad_len]

    c = np.transpose(mel_rec, (1, 0))
    waveform = wavegen(wavnet, device, c=c)
    librosa.output.write_wav(output_path, waveform, sr=16000)