#!/usr/bin/env python
import os
import itertools
import datetime
import logging
import logging.config
import argparse
import math
import statistics
from distutils.util import strtobool
from collections import defaultdict

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import ConnectionPatch
import seaborn
import joblib
import librosa
from tabulate import tabulate

import fingerprint
import ann
import util.tracks
import hough


from matplotlib.backends.backend_pgf import FigureCanvasPgf
matplotlib.backend_bases.register_backend('pdf', FigureCanvasPgf)
seaborn.set(style='ticks')
seaborn.set_context("paper")
logger = logging.getLogger(__name__)
scriptdir = os.path.dirname(os.path.realpath(__file__))
logfile = os.path.join(scriptdir, 'logging.ini')
logging.config.fileConfig(logfile, disable_existing_loggers=False)


class Match(object):
    def __init__(self, query, neighbors):
        self.query = query
        self.neighbors = neighbors


class Neighbor(object):
    def __init__(self, kp, dist):
        self.kp = kp
        self.dist = dist


class Model(object):
    def __init__(self, matcher, kps, settings, tracks=[], spectrograms=[]):
        self.matcher = matcher
        self.keypoints = kps
        self.settings = settings
        self.tracks = tracks
        self.spectrograms = spectrograms


class Result(object):
    def __init__(self, track, clusters, train, settings):
        self.track = track
        self.clusters = clusters
        self.sources = defaultdict(list)
        self.times = defaultdict(list)
        self.pitch_shift = defaultdict(list)
        self.time_stretch = defaultdict(list)
        for c in clusters:
            key = c[0].neighbors[0].kp.source
            try:
                key = key.decode()
            except AttributeError:
                key = str(key)
            self.sources[key].append(c)
            seconds = int(
                c[0].query.x * settings['hop_length'] / settings['sr']
            )
            time_q = datetime.timedelta(seconds=seconds)
            seconds = int(
                c[0].neighbors[0].kp.x*settings['hop_length']/settings['sr']
            )
            time_t = datetime.timedelta(seconds=seconds)
            self.times[key].append((time_q, time_t))

            combos = itertools.combinations(c, 2)
            stretch_factors = [
                abs(m2.query.x - m1.query.x)
                /
                abs(m2.neighbors[0].kp.x - m1.neighbors[0].kp.x)
                for m1, m2 in combos
            ]
            self.time_stretch[key].append(stretch_factors)
            pitch_factors = [
                (m.neighbors[0].kp.y - m.query.y)*12/settings['octave_bins']
                for m in c
            ]
            self.pitch_shift[key].append(pitch_factors)
        correct = {
            str(s.original): s for s in track.samples
            if str(s.original) in train
        }
        self.true_pos = [
            correct.get(key) for key in self.sources if key in correct]
        self.false_pos = [
            train.get(key) for key in self.sources if key not in correct]
        self.false_neg = [
            correct.get(key) for key in correct if key not in self.sources]


def filter_matches(matches, abs_thresh=None, ratio_thresh=None,
                   cluster_dist=20, cluster_size=1, match_orientation=True,
                   ordered=False):
    logger.info('Filtering nearest neighbors down to actual matched samples')
    if match_orientation:
        # Remove matches with differing orientations
        total = len(matches)
        for match in list(matches):
            orient = match.query.orientation
            while (match.neighbors and
                   abs(orient - match.neighbors[0].kp.orientation) > 0.2):
                match.neighbors = match.neighbors[1:]
            if len(match.neighbors) == 0:
                matches.remove(match)
            elif len(match.neighbors) < 2:
                # logger.warn('Orientation check left < 2 neighbors')
                matches.remove(match)