from .. import ALIGNMENT_TRANSFORMER_CHECKPOINT
import numpy as np
from collections import defaultdict
from .pretrained_models import AlignmentTransformer
import torch
from .matchers import na_within
from scipy.interpolate import interp1d

################################### TEMPO MODELS ###################################


class TempoModel(object):
    """
    Base class for synchronization models

    Attributes
    ----------
    """
    def __init__(
        self,
        init_beat_period = 0.5,
        init_score_onset = 0,
        init_perf_onset = 0,
        lookback = 1    
        ):
        self.lookback = lookback
        self.beat_period = init_beat_period
        self.prev_score_onsets = [init_score_onset - 2 * lookback]
        self.prev_perf_onsets = [init_perf_onset - 2 * lookback * self.beat_period]
        self.prev_perf_onsets_at_score_onsets = defaultdict(list)
        self.prev_perf_onsets_at_score_onsets[self.prev_score_onsets[-1]].append(self.prev_perf_onsets[-1])
        self.est_onset = None
        self.score_perf_map = None
        # Count how many times has the tempo model been called
        self.counter = 0
        self.update(init_perf_onset - lookback * self.beat_period, init_score_onset - lookback)

    def predict(
        self,
        score_onset
        ):
        self.est_onset = self.score_perf_map(score_onset - (self.lookback+1)) + \
            (self.lookback+1) * self.beat_period 
        return self.est_onset

    
    def predict_ratio(
        self,
        score_onset,
        perf_onset
        ):
        self.est_onset = self.score_perf_map(score_onset - (self.lookback+1)) + \
            (self.lookback+1) * self.beat_period 
        error = perf_onset - self.est_onset
        offset_score =  score_onset  - self.prev_score_onsets[-1] 
        if offset_score > 0.0:
            return error/(offset_score * self.beat_period)
        else:
            return error

    def update(
        self,
        performed_onset,
        score_onset
        ):

        self.prev_perf_onsets_at_score_onsets[score_onset].append(performed_onset)
        if score_onset == self.prev_score_onsets[-1]:
            #     self.prev_perf_onsets[-1] = 4/5 * self.prev_perf_onsets[-1] + 1/5* performed_onset
            self.prev_perf_onsets[-1] = np.median(self.prev_perf_onsets_at_score_onsets[score_onset])
        else:
            self.prev_score_onsets.append(score_onset)
            self.prev_perf_onsets.append(performed_onset)
            
        self.score_perf_map = interp1d(self.prev_score_onsets[-100:], 
                                       self.prev_perf_onsets[-100:], 
                                       fill_value="extrapolate")
        self.beat_period = np.clip((self.score_perf_map(score_onset) - \
            self.score_perf_map(score_onset - self.lookback))/self.lookback, 0.1, 10.0)
        self.counter += 1


class DummyTempoModel(object):


    """
    Base class for synchronization models

    Attributes
    ----------
    """
    def __init__(
        self,
        init_beat_period = 0.5,
        init_score_onset = 0,
        init_perf_onset = 0,
        lookback = 1,
        func = None
    ):
        
        self.lookback = lookback
        self.beat_period = init_beat_period
        self.score_perf_map = func
        # Count how many times has the tempo model been called
        self.counter = 0

    def predict(
        self,
        score_onset
        ):
        self.est_onset = self.score_perf_map(score_onset)
        return self.est_onset
    
    # def predict_ratio(
    #     self,
    #     score_onset,
    #     perf_onset
    #     ):
    #     self.est_onset = self.score_perf_map(score_onset - self.lookback) + \
    #         self.lookback * self.beat_period 
    #     error = perf_onset - self.est_onset
    #     offset_score =  score_onset  - self.prev_score_onsets[-1] 
    #     if offset_score > 0.0:
    #         return error/(offset_score * self.beat_period)
    #     else:
    #         return error
    def update(
        self,
        performed_onset,
        score_onset
        ):
        self.counter += 1


################################### ONLINE MATCHERS ###################################


class OnlineTransformerMatcher(object):
    def __init__(self,
                 score_note_array_full
                 ):
        self.score_note_array_full = np.sort(score_note_array_full, order="onset_beat")
        self.first_p_onset = None
        self.tempo_model = None
        
        self._prev_performance_notes = list()
        self._prev_score_onset = None
        self._snote_aligned = set()
        self._pnote_aligned = set()
        self._pnote_aligned_pitch = list()
        self.alignment = []
        self.note_alignments = []
        self.time_since_nn_update = 0
        self.prepare_score()
        self.prepare_model()

    def prepare_score(self):

        self.score_note_array_no_grace = self.score_note_array_full[self.score_note_array_full["is_grace"] == False]
        self.score_by_pitch = defaultdict(list)
        unique_pitches = np.unique(self.score_note_array_full["pitch"])
        for pitch in unique_pitches:
            self.score_by_pitch[pitch] = self.score_note_array_full[self.score_note_array_full["pitch"] == pitch]

        self._prev_score_onset = self.score_note_array_full["onset_beat"][0]
        self._unique_score_onsets = np.unique(self.score_note_array_full["onset_beat"])

        # onset range for forward backward view
        self.onset_range_at_onset = dict()
        for s_id, s_onset in enumerate(self._unique_score_onsets[1:-1]):
            self.onset_range_at_onset[s_onset] = [self._unique_score_onsets[s_id], self._unique_score_onsets[s_id+2]]
        self.onset_range_at_onset[self._unique_score_onsets[0]] = [self._unique_score_onsets[0], self._unique_score_onsets[1]]
        self.onset_range_at_onset[self._unique_score_onsets[-1]] = [self._unique_score_onsets[-2], self._unique_score_onsets[-1]]

        # set of pitches at onset / map from onset to idx in unique onsets
        self.pitches_at_onset_by_id = list()
        self.id_by_onset = dict()

        for i, onset in enumerate(self._unique_score_onsets):
            self.pitches_at_onset_by_id.append(
                set(self.score_note_array_no_grace[
                    self.score_note_array_no_grace["onset_beat"] == onset
                    ]["pitch"])
                )
            self.id_by_onset[onset] = i

        # aligned notes at each onset
        self.aligned_notes_at_onset = defaultdict(list)

    def prepare_performance(self, first_onset):
        self.tempo_model = TempoModel(init_beat_period = 0.5,
                                    init_score_onset = self.score_note_array_full["onset_beat"][0],
                                    init_perf_onset = first_onset,
                                    lookback = 1)
        
    def prepare_model(self):
        self.model = AlignmentTransformer(
            token_number = 91,
            dim_model = 64,
            dim_class = 2,
            num_heads = 8,
            num_decoder_layers = 6,
            dropout_p = 0.1
            )
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        checkpoint = torch.load(ALIGNMENT_TRANSFORMER_CHECKPOINT, 
                                map_location=torch.device(self.device))
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.model.to(self.device)
        self.model.eval()

    def offline(self, performance_note_array):
        self.prepare_performance(performance_note_array[0]["onset_sec"])

        for p_note in performance_note_array[:]:
            self.online(p_note)

        for s_ID, p_ID in self.alignment:
                self.note_alignments.append({'label': 'match', 
                                        "score_id": s_ID, 
                                        "performance_id": p_ID})
        # add unmatched notes
        for score_note in self.score_note_array_full:
            if score_note["id"] not in self._snote_aligned:
                self.note_alignments.append({'label': 'deletion', 'score_id': score_note["id"]})
        
        for performance_note in performance_note_array:
            if performance_note["id"] not in self._pnote_aligned:
                self.note_alignments.append({'label': 'insertion', 'performance_id': performance_note["id"]})

        return self.note_alignments

    def online(self, performance_note, debug=False):
        self.time_since_nn_update += 1
        p_id = performance_note["id"]
        p_onset = performance_note["onset_sec"]
        p_pitch = performance_note["pitch"]
        self._prev_performance_notes.append(p_pitch)

        possible_score_notes = self.score_by_pitch[p_pitch]

        # align greedily if open note at current onset
        if p_pitch in self.pitches_at_onset_by_id[self.id_by_onset[self._prev_score_onset]]:
            best_notes = na_within(possible_score_notes, "onset_beat", 
                                    self._prev_score_onset, self._prev_score_onset,
                                    exclusion_ids=self._snote_aligned)
            if len(best_notes) > 0:
                best_note = best_notes[0]
                self.add_note_alignment(p_id, best_note["id"], p_onset, best_note["onset_beat"])
                return
        
        current_id = self.id_by_onset[self._prev_score_onset]
        s_slice = slice(np.max((current_id-7, 0)), current_id+9 )
        p_slice = slice(-8, None )
        score_seq = self.pitches_at_onset_by_id[s_slice]
        perf_seq = self._prev_performance_notes[p_slice]

        tokenized_score_seq =  tokenize(score_seq, perf_seq, dims = 7)
        out = self.model(torch.from_numpy(tokenized_score_seq).unsqueeze(0).to(self.device))
        pred_ids = torch.argsort(torch.softmax(out.squeeze(1),dim=0)[:,1], descending=True).cpu().numpy()

        top_three_notes = dict()
        for pred_id in pred_ids[:3]:
            new_pred_id = pred_id - len(perf_seq) - 1 - (current_id - np.max((current_id-7, 0)))

            pred_score_onset = self._unique_score_onsets[current_id + new_pred_id]
            possible_score_notes = self.score_by_pitch[p_pitch]
            possible_score_notes =  na_within(possible_score_notes, "onset_beat", 
                                          pred_score_onset, pred_score_onset,
                                          exclusion_ids=self._snote_aligned)

            if len(possible_score_notes) > 0:
                dist = np.abs(self.tempo_model.predict(possible_score_notes[0]["onset_beat"]) - p_onset)
                top_three_notes[dist] = possible_score_notes[0]
                
        dists = list(top_three_notes.keys())
        if len(dists) >= 1:

            best_note = top_three_notes[np.min(dists)]
            if best_note["is_grace"]:
                self.add_note_alignment(p_id, best_note["id"])
            else:
                self.add_note_alignment(p_id, best_note["id"], p_onset, best_note["onset_beat"])                        

    def add_note_alignment(self,
                           perf_id, score_id, 
                           perf_onset = None, score_onset = None
                           ):
        self.alignment.append((score_id, perf_id))
        self._snote_aligned.add(score_id)
        self._pnote_aligned.add(perf_id)
        if perf_onset is not None and score_onset is not None:
            self.aligned_notes_at_onset[score_onset].append(perf_onset)
            if score_onset >= self._prev_score_onset:
                self.tempo_model.update(perf_onset, score_onset)
                self._prev_score_onset = score_onset

    def __call__(self):
        return None


def perf_tokenizer(pitch, dims = 7):
    return np.ones((1,dims), dtype = int) * (pitch - 20)


def score_tokenizer(pitch_set, dims = 7):
    token = np.zeros((1,dims), dtype = int)
    for no, pitch in enumerate(list(pitch_set)):
        if pitch >= 21 and pitch <= 108 and no < dims:
            token[0,no] = pitch - 20
    return token


def perf_to_score_tokenizer(dims = 7):
    return np.ones((1,dims), dtype = int) *89


def end_tokenizer(dims = 7, end_dims=1):
    return np.ones((end_dims,dims), dtype = int) *90


def tokenize(score_segment, perf_segment, dims = 7):
    tokens = list()
    for perf_note in perf_segment:
        perf_token = perf_tokenizer(perf_note, dims)
        tokens.append(perf_token)
    tokens.append(perf_to_score_tokenizer(dims))
    for score_set in score_segment:
        score_token = score_tokenizer(score_set, dims)
        tokens.append(score_token)
    
    end_token = end_tokenizer(dims, 26 - len(tokens))
    tokens.append(end_token)

    return np.row_stack(tokens)


class OnlinePureTransformerMatcher(object):
    def __init__(self,
                 score_note_array_full
                 ):
        self.score_note_array_full = np.sort(score_note_array_full, order="onset_beat")
        self.first_p_onset = None
        self.tempo_model = None
        
        self._prev_performance_notes = list()
        self._prev_score_onset = None
        self._snote_aligned = set()
        self._pnote_aligned = set()
        self._pnote_aligned_pitch = list()
        self.alignment = []
        self.note_alignments = []
        self.time_since_nn_update = 0
        self.prepare_score()
        self.prepare_model()

    def prepare_score(self):

        self.score_note_array_no_grace = self.score_note_array_full[self.score_note_array_full["is_grace"] == False]
        self.score_by_pitch = defaultdict(list)
        unique_pitches = np.unique(self.score_note_array_full["pitch"])
        for pitch in unique_pitches:
            self.score_by_pitch[pitch] = self.score_note_array_full[self.score_note_array_full["pitch"] == pitch]

        self._prev_score_onset = self.score_note_array_full["onset_beat"][0]
        self._unique_score_onsets = np.unique(self.score_note_array_full["onset_beat"])

        # onset range for forward backward view
        self.onset_range_at_onset = dict()
        for s_id, s_onset in enumerate(self._unique_score_onsets[1:-1]):
            self.onset_range_at_onset[s_onset] = [self._unique_score_onsets[s_id], self._unique_score_onsets[s_id+2]]
        self.onset_range_at_onset[self._unique_score_onsets[0]] = [self._unique_score_onsets[0], self._unique_score_onsets[1]]
        self.onset_range_at_onset[self._unique_score_onsets[-1]] = [self._unique_score_onsets[-2], self._unique_score_onsets[-1]]

        # set of pitches at onset / map from onset to idx in unique onsets
        self.pitches_at_onset_by_id = list()
        self.id_by_onset = dict()

        for i, onset in enumerate(self._unique_score_onsets):
            self.pitches_at_onset_by_id.append(
                set(self.score_note_array_no_grace[
                    self.score_note_array_no_grace["onset_beat"] == onset
                    ]["pitch"])
                )
            self.id_by_onset[onset] = i

        # aligned notes at each onset
        self.aligned_notes_at_onset = defaultdict(list)

    def prepare_performance(self, first_onset, func = None):
        if func is None:
            self.tempo_model = TempoModel(init_beat_period = 0.5,
                                    init_score_onset = self.score_note_array_full["onset_beat"][0],
                                    init_perf_onset = first_onset,
                                    lookback = 3)
        else:
            self.tempo_model = DummyTempoModel(init_beat_period = 0.5,
                                    init_score_onset = self.score_note_array_full["onset_beat"][0],
                                    init_perf_onset = first_onset,
                                    lookback = 3,
                                    func = func)

    def prepare_model(self):
        self.model = AlignmentTransformer(
            token_number = 91,# 21 - 108 + 2 for padding (start_score, end) + 1 for non_pitch
            dim_model = 64,
            dim_class = 2,
            num_heads = 8,
            num_decoder_layers = 6,
            dropout_p = 0.1
            )
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        checkpoint = torch.load(ALIGNMENT_TRANSFORMER_CHECKPOINT, 
                                map_location=torch.device(self.device))
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.model.to(self.device)
        self.model.eval()

    def offline(self, performance_note_array, func = None):
        self.prepare_performance(performance_note_array[0]["onset_sec"], func)

        for p_note in performance_note_array[:]:
            self.online(p_note)

        for s_ID, p_ID in self.alignment:
                self.note_alignments.append({'label': 'match', 
                                        "score_id": s_ID, 
                                        "performance_id": p_ID})
        # add unmatched notes
        for score_note in self.score_note_array_full:
            if score_note["id"] not in self._snote_aligned:
                self.note_alignments.append({'label': 'deletion', 'score_id': score_note["id"]})
        
        for performance_note in performance_note_array:
            if performance_note["id"] not in self._pnote_aligned:
                self.note_alignments.append({'label': 'insertion', 'performance_id': performance_note["id"]})

        return self.note_alignments

    def online(self, performance_note, debug=False):
        # directly align with NN without any cautionary measures
        p_id = performance_note["id"]
        p_onset = performance_note["onset_sec"]
        p_pitch = performance_note["pitch"]
        self._prev_performance_notes.append(p_pitch)

        current_id = self.id_by_onset[self._prev_score_onset]
        s_slice = slice(np.max((current_id-7, 0)), current_id+9 )
        p_slice = slice(-8, None )
        score_seq = self.pitches_at_onset_by_id[s_slice]
        perf_seq = self._prev_performance_notes[p_slice]

        tokenized_score_seq =  tokenize(score_seq, perf_seq, dims = 7)
        out = self.model(torch.from_numpy(tokenized_score_seq).unsqueeze(0).to(self.device))
        pred_id = torch.argmax(
            torch.softmax(out.squeeze(1),dim=1)[:,1]
            ).cpu().numpy()
        new_pred_id = pred_id - len(perf_seq) - 1 - (current_id - np.max((current_id-7, 0)))

        
        pred_score_onset = self._unique_score_onsets[current_id + new_pred_id]
        possible_score_notes = self.score_by_pitch[p_pitch]
        possible_score_notes =  na_within(possible_score_notes, "onset_beat", 
                                        pred_score_onset, pred_score_onset,
                                        exclusion_ids=self._snote_aligned)

        if len(possible_score_notes) > 0:
            best_note = possible_score_notes[0]
            if best_note["is_grace"]:
                self.add_note_alignment(p_id, best_note["id"])
            else:
                self.add_note_alignment(p_id, best_note["id"], p_onset, best_note["onset_beat"])    

    def add_note_alignment(self,
                           perf_id, score_id, 
                           perf_onset = None, 
                           score_onset = None
                           ):
        self.alignment.append((score_id, perf_id))
        self._snote_aligned.add(score_id)
        self._pnote_aligned.add(perf_id)
        if perf_onset is not None and score_onset is not None:
            self.aligned_notes_at_onset[score_onset].append(perf_onset)
            if score_onset >= self._prev_score_onset:
                self.tempo_model.update(perf_onset, score_onset)
                self._prev_score_onset = score_onset

    def __call__(self):

        return None
