#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (C) 2018 Brno University of Technology FIT
# Author: Jan Profant <jan.profant@phonexia.com>
# All Rights Reserved

import os
import logging
import pickle
import multiprocessing

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

from vbdiar.features.segments import get_frames_from_time
from vbdiar.embeddings.embedding import extract_embeddings
from vbdiar.utils import mkdir_p
from vbdiar.utils.utils import Utils

logger = logging.getLogger(__name__)


def process_files(fns, speakers_dict, features_extractor, embedding_extractor,
                  audio_dir, wav_suffix, in_rttm_dir, rttm_suffix, min_length, n_jobs=1):
    """

    Args:
        fns:
        speakers_dict:
        features_extractor:
        embedding_extractor:
        audio_dir:
        wav_suffix:
        in_rttm_dir:
        rttm_suffix:
        min_length:
        n_jobs:

    Returns:

    """
    kwargs = dict(speakers_dict=speakers_dict, features_extractor=features_extractor,
                  embedding_extractor=embedding_extractor, audio_dir=audio_dir, wav_suffix=wav_suffix,
                  in_rttm_dir=in_rttm_dir, rttm_suffix=rttm_suffix, min_length=min_length)
    if n_jobs == 1:
        ret = _process_files((fns, kwargs))
    else:
        pool = multiprocessing.Pool(n_jobs)
        ret = pool.map(_process_files, ((part, kwargs) for part in Utils.partition(fns, n_jobs)))
    return ret


def _process_files(dargs):
    """

    Args:
        dargs:

    Returns:

    """
    fns, kwargs = dargs
    ret = []
    for fn in fns:
        ret.append(process_file(file_name=fn, **kwargs))
    return ret


def process_file(file_name, speakers_dict, features_extractor, embedding_extractor,
                 audio_dir, wav_suffix, in_rttm_dir, rttm_suffix, min_length):
    """ Extract embeddings for all defined speakers.

    Args:
        file_name (string_types): path to input audio file
        speakers_dict (dict): dictionary containing all embedding across speakers
        features_extractor (Any):
        embedding_extractor (Any):
        audio_dir (string_types):
        wav_suffix (string_types):
        in_rttm_dir (string_types):
        rttm_suffix (string_types):
        min_length (float):

    Returns:
        dict: updated dictionary with speakers
    """
    logger.info('Processing file `{}`.'.format(file_name.split()[0]))
    # extract features from whole audio
    features = features_extractor.audio2features(os.path.join(audio_dir, '{}{}'.format(file_name, wav_suffix)))

    # process utterances of the speakers
    features_dict = {}
    with open(f'{os.path.join(in_rttm_dir, file_name)}{rttm_suffix}') as f:
        for line in f:
            start_time, dur = int(float(line.split()[3]) * 1000), int(float(line.split()[4]) * 1000)
            speaker = line.split()[7]
            if dur > min_length:
                end_time = start_time + dur
                start, end = get_frames_from_time(int(start_time)), get_frames_from_time(int(end_time))
                if speaker not in features_dict:
                    features_dict[speaker] = {}
                
                assert 0 <= start < end, \
                    f'Incorrect timing for extracting features, start: {start}, size: {features.shape[0]}, end: {end}.'
                if end >= features.shape[0]:
                    end = features.shape[0] - 1
                features_dict[speaker][(start_time, end_time)] = features[start:end]
    for speaker in features_dict:
        embedding_set = extract_embeddings(features_dict[speaker], embedding_extractor)
        embeddings_long = embedding_set.get_all_embeddings()
        if speaker not in speakers_dict.keys():
            speakers_dict[speaker] = embeddings_long
        else:
            speakers_dict[speaker] = np.concatenate((speakers_dict[speaker], embeddings_long), axis=0)
    return speakers_dict


class Normalization(object):
    """ Speaker normalization S-Norm. """
    embeddings = None
    in_emb_dir = None

    def __init__(self, norm_list, audio_dir=None, in_rttm_dir=None, in_emb_dir=None,
                 out_emb_dir=None, min_length=None, features_extractor=None, embedding_extractor=None,
                 plda=None, wav_suffix='.wav', rttm_suffix='.rttm', n_jobs=1):
        """ Initialize normalization object.

        Args:
            norm_list (string_types): path to normalization list
            audio_dir (string_types|None): path to audio directory
            in_rttm_dir (string_types|None): path to directory with rttm files
            in_emb_dir (str|None): path to directory with i-vectors
            out_emb_dir (str|None): path to directory for storing embeddings
            min_length (int): minimal length for extracting embeddings
            features_extractor (Any): object for feature extraction
            embedding_extractor (Any): object for extracting embedding
            plda (PLDA|None): plda model object
            wav_suffix (string_types): suffix of wav files
            rttm_suffix (string_types): suffix of rttm files
        """
        if audio_dir:
            self.audio_dir = os.path.abspath(audio_dir)
        self.norm_list = norm_list
        if in_rttm_dir:
            self.in_rttm_dir = os.path.abspath(in_rttm_dir)
        else:
            raise ValueError('It is required to have input rttm files for normalization.')
        self.features_extractor = features_extractor
        self.embedding_extractor = embedding_extractor
        self.plda = plda
        self.wav_suffix = wav_suffix
        self.rttm_suffix = rttm_suffix
        if in_emb_dir:
            self.in_emb_dir = os.path.abspath(in_emb_dir)
        if out_emb_dir:
            self.out_emb_dir = os.path.abspath(out_emb_dir)
        self.min_length = min_length
        self.n_jobs = n_jobs
        if self.in_emb_dir is None:
            self.embeddings = self.extract_embeddings()
        else:
            self.embeddings = self.load_embeddings()
        self.mean = np.mean(self.embeddings, axis=0)

    def __iter__(self):
        current = 0
        while current < len(self.embeddings):
            yield self.embeddings[current]
            current += 1

    def __getitem__(self, key):
        return self.embeddings[key]

    def __setitem__(self, key, value):
        self.embeddings[key] = value

    def __len__(self):
        return len(self.embeddings)

    def extract_embeddings(self):
        """ Extract normalization embeddings using averaging.

        Returns:
            Tuple[np.array, np.array]: vectors for individual speakers, global mean over all speakers
        """
        speakers_dict, fns = {}, []
        with open(self.norm_list) as f:
            for line in f:
                if len(line.split()) > 1:  # number of speakers is defined
                    line = line.split()[0]
                else:
                    line = line.replace(os.linesep, '')
                fns.append(line)

        speakers_dict = process_files(fns, speakers_dict=speakers_dict, features_extractor=self.features_extractor,
                                      embedding_extractor=self.embedding_extractor, audio_dir=self.audio_dir,
                                      wav_suffix=self.wav_suffix, in_rttm_dir=self.in_rttm_dir,
                                      rttm_suffix=self.rttm_suffix, min_length=self.min_length, n_jobs=self.n_jobs)
        assert len(speakers_dict) == len(fns)
        # all are the same
        merged_speakers_dict = speakers_dict[0]

        if self.out_emb_dir:
            for speaker in merged_speakers_dict:
                out_path = os.path.join(self.out_emb_dir, f'{speaker}.pkl')
                mkdir_p(os.path.dirname(out_path))
                with open(out_path, 'wb') as f:
                    pickle.dump(merged_speakers_dict[speaker], f, pickle.HIGHEST_PROTOCOL)

        for speaker in merged_speakers_dict:
            merged_speakers_dict[speaker] = np.mean(merged_speakers_dict[speaker], axis=0)

        return np.array(list(merged_speakers_dict.values()))

    def load_embeddings(self):
        """ Load normalization embeddings from pickle files.

        Returns:
            np.array: embeddings per speaker
        """
        embeddings, speakers = [], set()
        with open(self.norm_list) as f:
            for file_name in f:
                if len(file_name.split()) > 1:  # number of speakers is defined
                    file_name = file_name.split()[0]
                else:
                    file_name = file_name.replace(os.linesep, '')
                with open('{}{}'.format(os.path.join(self.in_rttm_dir, file_name), self.rttm_suffix)) as fp:
                    for line in fp:
                        speakers.add(line.split()[7])

        logger.info('Loading pickled normalization embeddings from `{}`.'.format(self.in_emb_dir))
        for speaker in speakers:
            embedding_path = os.path.join(self.in_emb_dir, '{}.pkl'.format(speaker))
            if os.path.isfile(embedding_path):
                logger.info('Loading normalization pickle file `{}`.'.format(speaker))
                with open(embedding_path, 'rb') as f:
                    # append mean from speaker's embeddings
                    speaker_embeddings = pickle.load(f)
                    embeddings.append(np.mean(speaker_embeddings, axis=0))
            else:
                logger.warning('No pickle file found for `{}` in `{}`.'.format(speaker, self.in_emb_dir))
        return np.array(embeddings)

    def s_norm(self, test, enroll):
        """ Run speaker normalization (S-Norm) on cached embeddings.

        Args:
            test (np.array): test embedding
            enroll (np.array): enroll embedding

        Returns:
            float: hypothesis
        """
        if self.plda:
            a = self.plda.score(test, self.embeddings).T
            b = self.plda.score(enroll, self.embeddings).T
            c = self.plda.score(enroll, test).T
        else:
            a = cosine_similarity(test, self.embeddings).T
            b = cosine_similarity(enroll, self.embeddings).T
            c = cosine_similarity(enroll, test).T
        scores = []
        for ii in range(test.shape[0]):
            test_scores = []
            for jj in range(enroll.shape[0]):
                test_mean, test_std = np.mean(a.T[ii]), np.std(a.T[ii])
                enroll_mean, enroll_std = np.mean(b.T[jj]), np.std(b.T[jj])
                s = c[ii][jj]
                test_scores.append((((s - test_mean) / test_std + (s - enroll_mean) / enroll_std) / 2))
            scores.append(test_scores)
        return np.array(scores)

[((21, 9, 21, 36), 'logging.getLogger', 'logging.getLogger', ({(21, 27, 21, 35): '__name__'}, {}), '(__name__)', False, 'import logging\n'), ((49, 15, 49, 43), 'multiprocessing.Pool', 'multiprocessing.Pool', ({(49, 36, 49, 42): 'n_jobs'}, {}), '(n_jobs)', False, 'import multiprocessing\n'), ((110, 24, 110, 87), 'vbdiar.embeddings.embedding.extract_embeddings', 'extract_embeddings', ({(110, 43, 110, 65): 'features_dict[speaker]', (110, 67, 110, 86): 'embedding_extractor'}, {}), '(features_dict[speaker], embedding_extractor)', False, 'from vbdiar.embeddings.embedding import extract_embeddings\n'), ((164, 20, 164, 52), 'numpy.mean', 'np.mean', (), '', True, 'import numpy as np\n'), ((244, 15, 244, 35), 'numpy.array', 'np.array', ({(244, 24, 244, 34): 'embeddings'}, {}), '(embeddings)', True, 'import numpy as np\n'), ((273, 15, 273, 31), 'numpy.array', 'np.array', ({(273, 24, 273, 30): 'scores'}, {}), '(scores)', True, 'import numpy as np\n'), ((115, 37, 115, 102), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((143, 29, 143, 55), 'os.path.abspath', 'os.path.abspath', ({(143, 45, 143, 54): 'audio_dir'}, {}), '(audio_dir)', False, 'import os\n'), ((146, 31, 146, 59), 'os.path.abspath', 'os.path.abspath', ({(146, 47, 146, 58): 'in_rttm_dir'}, {}), '(in_rttm_dir)', False, 'import os\n'), ((155, 30, 155, 57), 'os.path.abspath', 'os.path.abspath', ({(155, 46, 155, 56): 'in_emb_dir'}, {}), '(in_emb_dir)', False, 'import os\n'), ((157, 31, 157, 59), 'os.path.abspath', 'os.path.abspath', ({(157, 47, 157, 58): 'out_emb_dir'}, {}), '(out_emb_dir)', False, 'import os\n'), ((212, 44, 212, 90), 'numpy.mean', 'np.mean', (), '', True, 'import numpy as np\n'), ((236, 15, 236, 45), 'os.path.isfile', 'os.path.isfile', ({(236, 30, 236, 44): 'embedding_path'}, {}), '(embedding_path)', False, 'import os\n'), ((206, 27, 206, 75), 'os.path.join', 'os.path.join', ({(206, 40, 206, 56): 'self.out_emb_dir', (206, 58, 206, 74): 'f"""{speaker}.pkl"""'}, {}), "(self.out_emb_dir, f'{speaker}.pkl')", False, 'import os\n'), ((261, 16, 261, 56), 'sklearn.metrics.pairwise.cosine_similarity', 'cosine_similarity', ({(261, 34, 261, 38): 'test', (261, 40, 261, 55): 'self.embeddings'}, {}), '(test, self.embeddings)', False, 'from sklearn.metrics.pairwise import cosine_similarity\n'), ((262, 16, 262, 58), 'sklearn.metrics.pairwise.cosine_similarity', 'cosine_similarity', ({(262, 34, 262, 40): 'enroll', (262, 42, 262, 57): 'self.embeddings'}, {}), '(enroll, self.embeddings)', False, 'from sklearn.metrics.pairwise import cosine_similarity\n'), ((263, 16, 263, 47), 'sklearn.metrics.pairwise.cosine_similarity', 'cosine_similarity', ({(263, 34, 263, 40): 'enroll', (263, 42, 263, 46): 'test'}, {}), '(enroll, test)', False, 'from sklearn.metrics.pairwise import cosine_similarity\n'), ((50, 67, 50, 95), 'vbdiar.utils.utils.Utils.partition', 'Utils.partition', ({(50, 83, 50, 86): 'fns', (50, 88, 50, 94): 'n_jobs'}, {}), '(fns, n_jobs)', False, 'from vbdiar.utils.utils import Utils\n'), ((94, 17, 94, 53), 'os.path.join', 'os.path.join', ({(94, 30, 94, 41): 'in_rttm_dir', (94, 43, 94, 52): 'file_name'}, {}), '(in_rttm_dir, file_name)', False, 'import os\n'), ((207, 24, 207, 49), 'os.path.dirname', 'os.path.dirname', ({(207, 40, 207, 48): 'out_path'}, {}), '(out_path)', False, 'import os\n'), ((209, 20, 209, 90), 'pickle.dump', 'pickle.dump', ({(209, 32, 209, 61): 'merged_speakers_dict[speaker]', (209, 63, 209, 64): 'f', (209, 66, 209, 89): 'pickle.HIGHEST_PROTOCOL'}, {}), '(merged_speakers_dict[speaker], f, pickle.HIGHEST_PROTOCOL)', False, 'import pickle\n'), ((240, 41, 240, 55), 'pickle.load', 'pickle.load', ({(240, 53, 240, 54): 'f'}, {}), '(f)', False, 'import pickle\n'), ((268, 38, 268, 54), 'numpy.mean', 'np.mean', ({(268, 46, 268, 53): 'a.T[ii]'}, {}), '(a.T[ii])', True, 'import numpy as np\n'), ((268, 56, 268, 71), 'numpy.std', 'np.std', ({(268, 63, 268, 70): 'a.T[ii]'}, {}), '(a.T[ii])', True, 'import numpy as np\n'), ((269, 42, 269, 58), 'numpy.mean', 'np.mean', ({(269, 50, 269, 57): 'b.T[jj]'}, {}), '(b.T[jj])', True, 'import numpy as np\n'), ((269, 60, 269, 75), 'numpy.std', 'np.std', ({(269, 67, 269, 74): 'b.T[jj]'}, {}), '(b.T[jj])', True, 'import numpy as np\n'), ((241, 38, 241, 73), 'numpy.mean', 'np.mean', (), '', True, 'import numpy as np\n'), ((229, 40, 229, 81), 'os.path.join', 'os.path.join', ({(229, 53, 229, 69): 'self.in_rttm_dir', (229, 71, 229, 80): 'file_name'}, {}), '(self.in_rttm_dir, file_name)', False, 'import os\n')]