# Based on audio_to_midi_melodia github repo by Justin Salamon <justin.salamon@nyu.edu>
import time
import argparse
import os
import numpy as np
from scipy.signal import medfilt
#import __init__
import pyaudio
import sys
import aubio
from aubio import sink
from songmatch import *
from song import *
from operator import itemgetter
from pprint import PrettyPrinter
pp = PrettyPrinter(indent=4)
from wavplayer import WavPlayer
from scipy.io.wavfile import read
import wave

def convert_durations(ds):
    start_time = ds[0][0]
    durations = []
    for i,d in enumerate(ds):
        if i+1 < len(ds):
            durations.append(d[1]-start_time)
    return durations

# if vocal input, set wf = None
song_input = '../test_recording/mary_xiuyan_quiet.wav'
wf = None#read(song_input)[1]
wf_streamout = None#wave.open(song_input, 'rb')

# pyaudio params
buffer_size = 512
pyaudio_format = pyaudio.paFloat32
n_channels = 1
samplerate = 44100
playrate = 44100
seconds_per_sample = buffer_size / samplerate

# setup aubio
tolerance = 0.2
win_s = 4096 # fft size
hop_s = buffer_size # hop size
pitch_o = aubio.pitch("default", win_s, hop_s, samplerate)
pitch_o.set_unit("midi")
pitch_o.set_tolerance(tolerance)

# setup aubio recording (for debugging)
record_output_fn = 'test.wav'
g = sink(record_output_fn, samplerate=samplerate)
recorded_output = []

# pitch detection parameters
thresh_high = 100
thresh_low = 1
start_len = 2
pitch_diff_thresh = 1.5
note_min_len = 5
pitches = [[]]
durations = [[]] # durations = [[t_start, t_end]...]
seq = []
start = True
UDS = False
start_note = None
time_counter = 0 # time recorded so far (in seconds)
keydiff = None
temporatio = None
startpt = None
detected = False
matched_song = ""

# song database
#allSongNames = ["twinkle","london_bridge","three_blind_mice","boat","lullaby","mary_had_a_little_lamb"]
allSongNames = ["let_it_go"]#, "twinkle","london_bridge","three_blind_mice","boat","lullaby","mary_had_a_little_lamb"]
songdb = SongDatabase(allSongNames)
songdb.preprocessMelodies()
songs = songdb.getAllMelody()
start_notes = songdb.getAllFirstNode()
timestamps = songdb.getAllTimestamps()

# song matcher
song_matcher = SongsMatchNew(songs, timestamps)

# load all songs
allWavs = {}
for song in allSongNames:
    allWavs[song] = songdb.getWAV(song)
player = WavPlayer(**allWavs)

input("Press Enter to continue...")

# initialise pyaudio
p = pyaudio.PyAudio()

def append_new(l, n, d, tc, note_min_len=5):
    if len(l[-1]) <= note_min_len:
        l.pop()
        d.pop()
    else:
        l[-1] = np.mean(l[-1])
        d[-1].append(tc)
    l.append(n)
    d.append([tc])
    return l, d

def process_audio(in_data, frame_count, time_info, status):
    ''' callback function for pyaudio'''
    global start
    global pitches
    global seq
    global start_note
    global time_counter
    global durations
    global keydiff
    global temporatio
    global startpt
    global wf
    global detected
    global matched_song
    global player
    global song_matcher
    global recorded_output

    time_counter += seconds_per_sample
    if wf is not None:
        signal = np.array(wf[:frame_count], dtype=np.float32)
        if len(signal) < frame_count:
            if detected:
                player.stop()
            return (in_data, pyaudio.paComplete)
        wf = wf[frame_count:]
        in_data = wf_streamout.readframes(frame_count)
    else:
        signal = np.frombuffer(in_data, dtype=np.float32)
    pitch = pitch_o(signal)[0]
    confidence = pitch_o.get_confidence()

    # process pitch information
    if (pitch < thresh_low or pitch > thresh_high) and not start:
        start = True
        pitches, durations = append_new(pitches, [], durations, time_counter)
    if pitch > thresh_low and pitch < thresh_high:
        start = False
        if len(pitches[-1]) < start_len and abs(pitch - pitches[-1]) > pitch_diff_thresh:
                pitches[-1] = [pitch]
                durations[-1] = [time_counter] # reset note-start time
        else:
            if abs(pitch - np.mean(pitches[-1])) > pitch_diff_thresh:
                pitches, durations = append_new(pitches, [pitch], durations, time_counter)
            else:
                pitches[-1].append(pitch)
                if not durations[-1]: # if start note
                    durations[-1] = [time_counter]

    # check if user stopped
    if len(durations) > 1:
        if (time_counter - durations[-1][0]) > 3.0:
            print("######### restarting #########")
            # reset all data
            pitches = [[]]
            durations = [[]]
            seq = []
            start = True
            start_note = None
            time_counter = 0
            song_matcher = SongsMatchNew(songs, timestamps)
            if detected:
                player.stop()
                keydiff = None
                temporatio = None
                startpt = None
                detected = False
                matched_song = ""

    # if note ends
    if len(pitches) > 2 :
        if not seq: # if seq empty
            start_note = pitches[-3]
        if UDS:
            if (pitches[-2] - pitches[-3]) > 1.0:
                seq.append("U")
            elif abs(pitches[-2] - pitches[-3]) <= 1.0:
                seq.append("S")
            else:
                seq.append("D")
        else:
            seq.append(int(pitches[-2]-pitches[-3]))
        pitches.pop(0)

        # add the obtained note to song_matcher to get probability
        song_matcher.addNote([seq[-1]])
        scores = song_matcher.getProbDic()
        best_song = sorted(scores.items(), key=itemgetter(1))[-1][0]
        #pp.pprint(scores)
        if max(scores.values()) > 0.8 and best_song != "Others" and not detected: # if confident enought about song
            matched_song = best_song
            converted_durations = convert_durations(durations)
            keydiff, temporatio, startpt = song_matcher.getKeyTempo(matched_song, start_notes[matched_song], start_note, converted_durations)
            player.curr_file = matched_song
            print(round(startpt*playrate))
            player.navigate(round(startpt*playrate))
            player.time_stretch = temporatio #/ (2**(keydiff/4/12))
            #player.pitch_shift = keydiff/4
            player.play()
            print("+++++++++++++")
            print("key difference: %f" %keydiff)
            print("tempo ratio: %f" %temporatio)
            print("start point: %f" %startpt)
            print("song: %s" %matched_song)
            print("+++++++++++++")
            detected = True
        if detected:
            converted_durations = convert_durations(durations)
            keydiff, temporatio, startpt = song_matcher.getKeyTempo(matched_song, start_notes[matched_song], start_note, converted_durations)
            player.curr_file = matched_song

    g(signal, len(signal))
    return (in_data, pyaudio.paContinue)

# open stream
if wf is None:
    stream = p.open(format=pyaudio_format,
                    channels=n_channels,
                    rate=samplerate,
                    input=True,
                    frames_per_buffer=buffer_size,
                    stream_callback=process_audio)
else:
    stream = p.open(format=p.get_format_from_width(wf_streamout.getsampwidth()),
                    channels=wf_streamout.getnchannels(),
                    rate=wf_streamout.getframerate(),
                    output=True,
                    frames_per_buffer=buffer_size,
                    stream_callback=process_audio)

print("*** starting recording")
stream.start_stream()
while stream.is_active():
    try:
        time.sleep(0.1)
    except KeyboardInterrupt:
        print("*** Ctrl+C pressed, exiting")
        break

print("*** done recording")
stream.stop_stream()
stream.close()
p.terminate()

print(start_note)
print(durations)
print(seq)
print(keydiff, temporatio, startpt)

song_matcher = SongsMatchNew(songs)
song_matcher.addNotes(seq)
scores = song_matcher.getProbDic()
print("+++++++++++++++++++++++++++++++")
print(sorted(scores.items(), key=itemgetter(1))[-1][0])
print("+++++++++++++++++++++++++++++++")


import torch
import torch.nn.functional as F
from torch.autograd import Variable
from torch.nn.utils import clip_grad_norm_

import numpy as np
from utils.utils import ensure_shared_grads

from params import params

def preprocess(obs, gpu_id):
    state = np.concatenate([obs[0], obs[1], obs[2], obs[3]]) / 255.
    state = torch.from_numpy(state).unsqueeze(0).float()

    with torch.cuda.device(gpu_id):
        state = state.cuda()

    return state

def postprocess(logit):
    prob = F.softmax(logit, dim=-1)
    log_prob = F.log_softmax(logit, dim=-1)
    entropy = -(log_prob * prob).sum(1)

    action = prob.multinomial(1).data

    log_prob = log_prob.gather(1, Variable(action))
    action = action.cpu().numpy()

    return action, entropy, log_prob

class run_agent(object):
    def __init__(self, model, gpu_id, goal_storage=None):
        self.model = model
        self.gpu_id = gpu_id
        self.goal_storage = goal_storage
        self.hx, self.cx = None, None
        self.eps_len = 0
        self.values = []
        self.log_probs = []
        self.rewards = []
        self.entropies = []
        self.done = False
        self.n_update = 0
        self.instruction = None

    def action_train(self, obs, instruction):
        state = preprocess(obs, self.gpu_id)

        target_preds, value, logit, self.hx, self.cx = self.model(state, instruction, self.hx, self.cx)

        action, entropy, log_prob = postprocess(logit)

        self.eps_len += 1
        return np.squeeze(action, axis=0), entropy, value, log_prob

    def action_test(self, observation, instruction):
        with torch.cuda.device(self.gpu_id):
            with torch.no_grad():
                state = preprocess(observation, self.gpu_id)
                target_preds, value, logit, self.hx, self.cx = self.model(state, instruction, self.hx, self.cx)

                prob = F.softmax(logit, dim=-1)
                action = prob.max(1)[1].data.cpu().numpy()

        self.eps_len += 1
        return action

    def synchronize(self, shared_model, instruction):
        self.instruction = instruction

        with torch.cuda.device(self.gpu_id):
            self.model.load_state_dict(shared_model.state_dict())
            self.cx = Variable(torch.zeros(1, 256).cuda())
            self.hx = Variable(torch.zeros(1, 256).cuda())
        self.eps_len = 0

    def put_reward(self, reward, entropy, value, log_prob):
        self.rewards.append(reward)
        self.entropies.append(entropy)
        self.values.append(value)
        self.log_probs.append(log_prob)

    def clear_actions(self):
        self.values.clear()
        self.log_probs.clear()
        self.rewards.clear()
        self.entropies.clear()

    def get_reward_len(self):
        return len(self.rewards)

    def training_goal_discriminator(self):
        batch = self.goal_storage.sample(params.goal_batch_size)
        batch_state, label = np.array(batch.goal), np.array(batch.label)

        batch_state = torch.from_numpy(batch_state).float()
        label = torch.tensor(label)

        with torch.cuda.device(self.gpu_id):
            batch_state = Variable(torch.FloatTensor(batch_state)).cuda()
            label = label.cuda()

        target_preds, _, _, _, _ = self.model(batch_state, self.instruction, self.hx, self.cx, use_ActorCritic=False)

        values, indices = target_preds.max(1)
        accuracy = torch.mean((indices.squeeze() == label).float())
        crossentropy_loss = F.cross_entropy(target_preds, label.long())

        return crossentropy_loss, accuracy

    def training(self, next_obs, shared_model, shared_optimizer, rank):
        self.n_update += 1
        self.cx = Variable(self.cx.data)
        self.hx = Variable(self.hx.data)

        R = torch.zeros(1, 1)
        if not self.done:
            state = preprocess(next_obs, self.gpu_id)
            target_preds, value, logit, self.hx, self.cx = self.model(state, self.instruction, self.hx, self.cx)

            R = value.data

        with torch.cuda.device(self.gpu_id):
            R = R.cuda()
        R = Variable(R)
        self.values.append(R)

        policy_loss = 0
        value_loss = 0
        gae = torch.zeros(1, 1)
        with torch.cuda.device(self.gpu_id):
            gae = gae.cuda()

        for i in reversed(range(len(self.rewards))):
            R = params.gamma * R + self.rewards[i]
            advantage = R - self.values[i]
            value_loss = value_loss + advantage.pow(2)  # 0.5 *

            # Generalized Advantage Estimation
            delta_t = params.gamma * self.values[i + 1].data - self.values[i].data + self.rewards[i]

            gae = gae * params.gamma * params.tau + delta_t

            policy_loss = policy_loss - self.log_probs[i] * Variable(gae) - params.entropy_coef \
                          * self.entropies[i]

        crossentropy_loss, accuracy = self.training_goal_discriminator()

        shared_optimizer.zero_grad()
        loss = policy_loss + params.value_loss_coef * value_loss + params.goal_coef * crossentropy_loss
        loss.backward()
        clip_grad_norm_(self.model.parameters(), params.clip_grad_norm)
        ensure_shared_grads(self.model, shared_model, gpu=self.gpu_id >= 0)

        shared_optimizer.step()

        if self.n_update % 1000 == 0:
            print("rank: {}, tr_loss: {}, accuracy: {}".format(rank, crossentropy_loss.data, accuracy))

        with torch.cuda.device(self.gpu_id):
            self.model.load_state_dict(shared_model.state_dict())

        self.clear_actions()


import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as img
import cv2
import os
from core.utils import draw_gt_bbox

###################################################################################
## This script preorganizes the Taggings of Yolov3 baseline 2 format [f_data_in_path]
## into a desired [patch_w, patch_h] cropped Tagging file [f_data_out_path] and images [write_image_path]
## main dir of cropped outputs

###################################################################################

## Inputs:
# the Tagging to be cropped:
f_data_in_path = '/home/tamar/RecceLite_code_packages/yolo3_baseline2/data/individual_datasets_without_24/recce_data_Tagging_1_2_3_4_5.txt'

## Outputs:
write_image_path = '/home/tamar/DBs/Reccelite/CroppedDB/croppedImgs_1_2_3_4_5_Th06_reg/'
# the output of this script: the cropped Tagging
f_data_out_path = '/home/tamar/DBs/Reccelite/CroppedDB/croppedImgs_1_2_3_4_5_Th06_reg.txt'

patch_w = 1*800
patch_h = 1*640

cropTH = 0.6 # above this percentage of the bbox included in the path - clip the box. else - ignore the bbox.

###################################################################################

if not os.path.exists(write_image_path):
    os.makedirs(write_image_path)
with open(f_data_in_path, 'r') as f_in:
    txt = f_in.readlines()
    annotations = [line.strip() for line in txt if len(line.strip().split()[1:]) != 0]
new_line_flag = False
# loop over each img in original list: crop each img and add the cropped_img txt to the new datatxt
with open(f_data_out_path, 'w') as f_out:
    for img_ind in range(len(annotations)):
        annotation = annotations[img_ind]
        line = annotation.split()
        image_path = line[0]
        if not os.path.exists(image_path):
            raise KeyError("%s does not exist ... " % image_path)
        print(image_path)
        bboxes = np.array([list(map(float, box.split(','))) for box in line[1:]])
        image = np.array(cv2.imread(image_path))
        print("original image size is:", image.shape)
        image_format = image_path.split('.')[-1]
        image_name = image_path.split('.')[0].split('/')[-1]
        orig_image_h = image.shape[0]
        orig_image_w = image.shape[1]
        # the number of cropped parts:
        div_w = int(np.ceil(orig_image_w/patch_w)) # num of w sections in w
        div_h = int(np.ceil(orig_image_h/patch_h)) # num of sections in h
        # draw_gt_bbox(image, bboxes)
        for ind_w in range(div_w):
            for ind_h in range(div_h):
                new_img_flag = True
                cropped_bboxes = []
                cropped_image_name = image_name + '_' + str(ind_w) + '_' + str(ind_h) + '.' + image_format
                # patches that are not last row  or last column
                if (ind_w!=div_w-1 and ind_h!=div_h-1):
                    minW = ind_w*patch_w; minH = ind_h*patch_h;
                    maxW = (ind_w + 1)*patch_w; maxH = (ind_h + 1)*patch_h
                # patch that is last row and last column
                elif (ind_w == div_w - 1 and ind_h == div_h - 1):
                    minW = orig_image_w - patch_w; minH = orig_image_h - patch_h;
                    maxW = orig_image_w; maxH = orig_image_w - 1
                # patches that are last row
                elif ind_w != div_w - 1:
                    minW = ind_w*patch_w; minH =  orig_image_h - patch_h
                    maxW = (ind_w + 1)*patch_w; maxH = orig_image_w - 1
                # patches that are last column
                else:
                    minW =orig_image_w - patch_w; minH = ind_h*patch_h;
                    maxW = orig_image_w; maxH = (ind_h + 1)*patch_h
                cropped_image = image[minH:maxH, minW:maxW]
                cv2.imwrite(write_image_path + cropped_image_name, cropped_image)
                for box in bboxes:
                    cropped1, croppedOne = [], []                    # boxes that are completely out of this patch: ignore
                    if box[0]>=maxW or box[1]>=maxH or box[2]<=minW or (box[3]<=minH):
                        continue
                    # boxes that are entirely included in patch: keep
                    elif box[0]>=minW and box[1]>=minH and box[2]<=maxW and box[3]<=maxW:
                        # print('entire')
                        cropped1 = box
                    # boxes that are partially included: check the extent of overlap
                    else:
                        # print('overlap')
                        areaBox = (box[2] - box[0]) * (box[3] - box[1])
                        npbox = np.array(box)
                        nppatch = np.array([minW, minH, maxW-1, maxH-1])
                        left_up = np.maximum(npbox[:2], nppatch[:2])
                        right_down = np.minimum(npbox[2:4], nppatch[2:4])
                        inter_section = np.maximum(right_down - left_up, 0.0)
                        inter_area = inter_section[0] * inter_section[1]
                        if inter_area >= cropTH*areaBox:
                            cropped1 = (np.concatenate([left_up, right_down])).tolist()
                            np.array(cropped1.append(box[4]))
                            cropped1 = np.array(cropped1)
                            # print('Successoverlap')
                        else:
                            # print('Failoverlap')
                            continue
                    # adjust coordinates
                    croppedOne[:2] = cropped1[:2] - [minW, minH]
                    croppedOne[2:4] = cropped1[2:4] - [minW, minH]
                    croppedOne.append(cropped1[4])
                    # print(ind_w, ind_h)
                    # print(minW, minH, maxW, maxH)
                    # print(box)
                    # print(cropped1)
                    # print(croppedOne)
                    # write box to new dataTxt:
                    if new_img_flag:
                        if new_line_flag:
                            f_out.write('\n')
                        new_line_flag = True
                        f_out.write(write_image_path + cropped_image_name + ' ')
                    new_img_flag = False
                    f_out.write(str(croppedOne[0]) + ',' + str(croppedOne[1]) + ',' + str(croppedOne[2]) + ',' + str(croppedOne[3]) + ',' + str(croppedOne[4]) + ' ')
                    cropped_bboxes.append(croppedOne)
                # print(image_name, 'minW: ', minW, 'minH: ', minH, 'maxW: ', maxW, 'maxH: ', maxH, 'len(cropped_bboxes): ', len(cropped_bboxes))
                # draw_gt_bbox(cropped_image, cropped_bboxes)
                pass