import numpy as np
from torch import Tensor, FloatTensor
from kospeech.data.audio.core import load_audio
from kospeech.data.audio.augment import NoiseInjector, SpecAugment
from kospeech.data.audio.feature import MelSpectrogram, MFCC, Spectrogram, FilterBank

class AudioParser(object):
    """
    Provides inteface of audio parser.

    Note:
        Do not use this class directly, use one of the sub classes.

    Method:
        - **parse_audio()**: abstract method. you have to override this method.
        - **parse_transcript()**: abstract method. you have to override this method.
    """
    def __init__(self, dataset_path, noiseset_size, sample_rate=16000, noise_level=0.7, noise_augment=False):
        if noise_augment:
            self.noise_injector = NoiseInjector(dataset_path, noiseset_size, sample_rate, noise_level)

    def parse_audio(self, *args, **kwargs):
        raise NotImplementedError

    def parse_transcript(self, *args, **kwargs):
        raise NotImplementedError

class SpectrogramParser(AudioParser):
    """
    Parses audio file into (spectrogram / mel spectrogram / mfcc) with various options.

    Args:
        transform_method (str): which feature to use (default: mel)
        sample_rate (int): Sample rate of audio signal. (Default: 16000)
        n_mels (int):  Number of mfc coefficients to retain. (Default: 40)
        frame_length (int): frame length for spectrogram (ms) (Default : 20)
        frame_shift (int): Length of hop between STFT windows. (ms) (Default: 10)
        feature_extract_by (str): which library to use for feature extraction(default: librosa)
        del_silence (bool): flag indication whether to delete silence or not (default: True)
        input_reverse (bool): flag indication whether to reverse input or not (default: True)
        normalize (bool): flag indication whether to normalize spectrum or not (default:True)
        time_mask_para (int): Hyper Parameter for Time Masking to limit time masking length
        freq_mask_para (int): Hyper Parameter for Freq Masking to limit freq masking length
        time_mask_num (int): how many time-masked area to make
        freq_mask_num (int): how many freq-masked area to make
        sos_id (int): start of sentence token`s identification
        eos_id (int): end of sentence token`s identification
        target_dict (dict): dictionary of filename and labels
    """
    VANILLA = 0           # Not apply augmentation
    SPEC_AUGMENT = 1      # SpecAugment
    NOISE_INJECTION = 2   # Noise Injection
    HYBRID_AUGMENT = 3    # Noise Injection & SpecAugment

    def __init__(self, feature_extract_by: str = 'librosa', sample_rate: int = 16000,
                 n_mels: int = 80, frame_length: int = 20, frame_shift: int = 10,
                 del_silence: bool = False, input_reverse: bool = True,
                 normalize: bool = False,  transform_method: str = 'mel',
                 time_mask_para: int = 70, freq_mask_para: int = 12, time_mask_num: int = 2, freq_mask_num: int = 2,
                 sos_id: int = 1, eos_id: int = 2, target_dict: dict = None, noise_augment: bool = False,
                 dataset_path: str = None, noiseset_size: int = 0, noise_level: float = 0.7) -> None:
        super(SpectrogramParser, self).__init__(dataset_path, noiseset_size, sample_rate, noise_level, noise_augment)
        self.del_silence = del_silence
        self.input_reverse = input_reverse
        self.normalize = normalize
        self.sos_id = sos_id
        self.eos_id = eos_id
        self.target_dict = target_dict
        self.spec_augment = SpecAugment(time_mask_para, freq_mask_para, time_mask_num, freq_mask_num)

        if transform_method.lower() == 'mel':
            self.transforms = MelSpectrogram(sample_rate, n_mels, frame_length, frame_shift, feature_extract_by)

        elif transform_method.lower() == 'mfcc':
            self.transforms = MFCC(sample_rate, n_mels, frame_length, frame_shift, feature_extract_by)

        elif transform_method.lower() == 'spect':
            self.transforms = Spectrogram(sample_rate, frame_length, frame_shift, feature_extract_by)

        elif transform_method.lower() == 'fbank':
            self.transforms = FilterBank(sample_rate, n_mels, frame_length, frame_shift)

        else:
            raise ValueError("Unsupported feature : {0}".format(transform_method))

    def parse_audio(self, audio_path: str, augment_method: int) -> Tensor:
        """
        Parses audio.

        Args:
             audio_path (str): path of audio file
             augment_method (int): flag indication which augmentation method to use.

        Returns: feature_vector
            - **feature_vector** (torch.FloatTensor): feature from audio file.
        """
        signal = load_audio(audio_path, self.del_silence)

        if signal is None:
            return None

        if augment_method == SpectrogramParser.NOISE_INJECTION or augment_method == SpectrogramParser.HYBRID_AUGMENT:
            signal = self.noise_injector(signal)

        feature_vector = self.transforms(signal)

        if self.normalize:
            feature_vector -= feature_vector.mean()

        # Refer to "Sequence to Sequence Learning with Neural Network" paper
        if self.input_reverse:
            feature_vector = feature_vector[:, ::-1]
            feature_vector = FloatTensor(np.ascontiguousarray(np.swapaxes(feature_vector, 0, 1)))
        else:
            feature_vector = FloatTensor(feature_vector).transpose(0, 1)

        if augment_method == SpectrogramParser.SPEC_AUGMENT or augment_method == SpectrogramParser.HYBRID_AUGMENT:
            feature_vector = self.spec_augment(feature_vector)

        return feature_vector

    def parse_transcript(self, *args, **kwargs):
        raise NotImplementedError


#!/usr/bin/env python
# -*- coding: UTF-8 -*-

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow as tf

from vocab import Vocab
from lib.models.simplesrlers.base_simplesrler import BaseSimpleSrlers
from lib.linalg import linear

class SimpleSrler(BaseSimpleSrlers):

    def __call__(self, dataset, moving_params=None):
        vocabs = dataset.vocabs
        inputs = dataset.inputs
        targets = dataset.targets

        reuse = (moving_params is not None)
        self.tokens_to_keep3D = tf.expand_dims(tf.to_float(tf.greater(inputs[:, :, 0], vocabs[0].ROOT)), 2)
        self.sequence_lengths = tf.reshape(tf.reduce_sum(self.tokens_to_keep3D, [1, 2]), [-1, 1])
        self.n_tokens = tf.reduce_sum(self.sequence_lengths)
        self.moving_params = moving_params

        word_inputs = vocabs[0].embedding_lookup(inputs[:, :, 0], inputs[:, :, 1], keep_prob=self.word_keep_prob, moving_params=self.moving_params)
        pos_inputs = vocabs[1].embedding_lookup(inputs[:, :, 2], moving_params=self.moving_params)
        verb_inputs = vocabs[3].embedding_lookup(inputs[:, :, 3], moving_params=self.moving_params)
        is_verb_inputs = vocabs[4].embedding_lookup(inputs[:, :, 4], moving_params=self.moving_params)
        #self.embed_concat()
        # top_recur = tf.concat(2, [word_inputs, pos_inputs, verb_inputs, is_verb_inputs]+ context)
        ###################### Modify Here
        # top_recur = tf.concat(2, [word_inputs, pos_inputs, verb_inputs, is_verb_inputs])
        top_recur = tf.concat(2, [word_inputs, pos_inputs])
        ###################### End of Modify
        for i in xrange(self.n_recur):
          with tf.variable_scope('RNN%d' % i, reuse=reuse):
            top_recur, _ = self.RNN(top_recur)

        ##################### Modify Here
        top_recur = tf.concat(2, [top_recur, verb_inputs, is_verb_inputs])
        for i in xrange(self.n_srl_recur):
          with tf.variable_scope('SRLRNN%d' % i, reuse=reuse):
            top_recur, _ = self.RNN(top_recur)
            ##################### End of Modify

        predicate_token = tf.to_int32(tf.equal(inputs[:, :, 4], vocabs[4]['1']))

        # Mask predicate hidden representation out
        predicate_h = tf.mul(top_recur, tf.to_float(tf.expand_dims(predicate_token, 2)))
        # print(predicate_h.get_shape().as_list())
        # Reduce dimension
        predicate_h = tf.reduce_sum(predicate_h, axis=1)
        # print(predicate_h.get_shape().as_list())
        # Broadcasting
        predicate_h = tf.mul(tf.expand_dims(predicate_h, 1), self.tokens_to_keep3D)
        # print(predicate_h.get_shape().as_list())
        # Concat
        classifier_input = tf.concat(2, [predicate_h, top_recur])

        # Mask predicate embedding out
        predicate_em = tf.mul(verb_inputs, tf.to_float(tf.expand_dims(predicate_token, 2)))
        # Reduce dimension
        predicate_em = tf.reduce_sum(predicate_em, axis=1)  # b x dim
        # b x dim -> b x r x dim
        predicate_em = tf.pack([predicate_em] * len(vocabs[2]), axis=1)
        # Role Representation

        role_em_list = vocabs[2].embedding_lookup(range(len(vocabs[2])), moving_params=self.moving_params)
        batch_list = tf.to_float(tf.greater(self.sequence_lengths, -1))
        # print(batch_list.get_shape().as_list())
        # batch_list = tf.pack([batch_list]*len(vocabs[2]), axis=1)
        # print(batch_list.get_shape().as_list())
        role_em_list = tf.mul(tf.to_float(tf.expand_dims(batch_list, 2)), role_em_list)
        para_input = tf.concat(2, [predicate_em, role_em_list])

        with tf.variable_scope('SRLClassifier_para', reuse=reuse):
          para = self.MLP4SRLWeight(para_input)
          result_dist = tf.batch_matmul(classifier_input, para, adj_y=True)
          srl_output = self.output(result_dist, targets)

        '''
        with tf.variable_scope('MLP0', reuse=reuse):
          argument = self.MLP(top_mlp)
        verb_inputs = tf.expand_dims(verb_inputs, 1)
        context = []
        for i in range(self.windonw_length):
          context.append(tf.expand_dims(vocabs[0].embedding_lookup(inputs[:, -1, 5 + i * 2], inputs[:, -1, 5 + i * 2 + 1],
                                                              moving_params=self.moving_params),1))
        com = tf.concat(2, [verb_inputs]+context)

        with tf.variable_scope('MLP1', reuse=reuse):
          predicate = self.MLP(com)
        with tf.variable_scope('GraphSRL', reuse=reuse):
          srl_logits = self.bilinear_classifier_srl(argument, predicate, len(vocabs[2]), add_bias1=True)
          srl_output = self.output(srl_logits, targets)
        '''

        #decode_output = self.CRFDecode(crf_output, targets[:,:,0])

        output = {}
        output["loss"] = srl_output["loss"]
        output["n_tokens"] = self.n_tokens
        output["predictions"] = srl_output['predictions']
        output["n_correct"] = srl_output["n_correct"]
        output["sequence_lengths"] = tf.reshape(tf.to_int64(self.sequence_lengths), [-1])
        # output["prob"] = result_dist
        return output

        # decode_output = self.CRFDecode(crf_output, targets[:,:,0])
        # output = {}
        # output["loss"] = crf_output["loss"]
        # output["n_tokens"] = self.n_tokens
        # output["crf_output"] = crf_output
        # output["sequence_lengths"] = tf.reshape(tf.to_int64(self.sequence_lengths), [-1])

        # return output


# -*- coding: utf-8 -*-
"""
DJI D-Gamut Colourspace
=======================

Defines the *DJI D-Gamut* colourspace:

-   :attr:`colour.models.DJI_D_GAMUT_COLOURSPACE`.

References
----------
-   :cite:`DJI2017` : Dji. (2017). White Paper on D-Log and D-Gamut of DJI
    Cinema Color System. Retrieved from https://dl.djicdn.com/downloads/\
zenmuse+x7/20171010/D-Log_D-Gamut_Whitepaper.pdf
"""

from __future__ import division, unicode_literals

import numpy as np

from colour.colorimetry import ILLUMINANTS
from colour.models.rgb import (RGB_Colourspace, log_encoding_DJIDLog,
                               log_decoding_DJIDLog)

__author__ = 'Colour Developers'
__copyright__ = 'Copyright (C) 2013-2020 - Colour Developers'
__license__ = 'New BSD License - https://opensource.org/licenses/BSD-3-Clause'
__maintainer__ = 'Colour Developers'
__email__ = 'colour-developers@colour-science.org'
__status__ = 'Production'

__all__ = [
    'DJI_D_GAMUT_PRIMARIES', 'DJI_D_GAMUT_WHITEPOINT_NAME',
    'DJI_D_GAMUT_WHITEPOINT', 'DJI_D_GAMUT_TO_XYZ_MATRIX',
    'XYZ_TO_DJI_D_GAMUT_MATRIX', 'DJI_D_GAMUT_COLOURSPACE'
]

DJI_D_GAMUT_PRIMARIES = np.array([
    [0.71, 0.31],
    [0.21, 0.88],
    [0.09, -0.08],
])
"""
*DJI D-Gamut* colourspace primaries.

DJI_D_GAMUT_PRIMARIES : ndarray, (3, 2)
"""

DJI_D_GAMUT_WHITEPOINT_NAME = 'D65'
"""
*DJI D-Gamut* colourspace whitepoint name.

DJI_D_GAMUT_WHITEPOINT : unicode
"""

DJI_D_GAMUT_WHITEPOINT = (ILLUMINANTS['CIE 1931 2 Degree Standard Observer'][
    DJI_D_GAMUT_WHITEPOINT_NAME])
"""
*DJI D-Gamut* colourspace whitepoint.

DJI_D_GAMUT_WHITEPOINT : ndarray
"""

DJI_D_GAMUT_TO_XYZ_MATRIX = np.array([[0.6482, 0.1940,
                                       0.1082], [0.2830, 0.8132, -0.0962],
                                      [-0.0183, -0.0832, 1.1903]])
"""
*DJI D-Gamut* colourspace to *CIE XYZ* tristimulus values matrix.

DJI_D_GAMUT_TO_XYZ_MATRIX : array_like, (3, 3)
"""

XYZ_TO_DJI_D_GAMUT_MATRIX = np.array([[1.7257, -0.4314,
                                       -0.1917], [-0.6025, 1.3906, 0.1671],
                                      [-0.0156, 0.0905, 0.8489]])
"""
*CIE XYZ* tristimulus values to *DJI D-Gamut* colourspace matrix.

XYZ_TO_DJI_D_GAMUT_MATRIX : array_like, (3, 3)
"""

DJI_D_GAMUT_COLOURSPACE = RGB_Colourspace(
    'DJI D-Gamut',
    DJI_D_GAMUT_PRIMARIES,
    DJI_D_GAMUT_WHITEPOINT,
    DJI_D_GAMUT_WHITEPOINT_NAME,
    DJI_D_GAMUT_TO_XYZ_MATRIX,
    XYZ_TO_DJI_D_GAMUT_MATRIX,
    log_encoding_DJIDLog,
    log_decoding_DJIDLog,
)
DJI_D_GAMUT_COLOURSPACE.__doc__ = """
*DJI_D-Gamut* colourspace.

    References
    ----------
    :cite:`DJI2017`

DJI_D_GAMUT_COLOURSPACE : RGB_Colourspace
"""


#!/usr/bin/env python

import click
from click_option_group import optgroup
import math
import numpy as np
import pandas as pd
import os
import shutil
import torch
from torch.utils.data import DataLoader, TensorDataset

# Local imports
from architectures import DanQ, get_loss, get_metrics, get_optimizer
from selene import Trainer
from sequence import one_hot_encode_many, rc_one_hot_encoding_many
from train import _get_Xs_ys_ids, _get_data_loader

CONTEXT_SETTINGS = {
    "help_option_names": ["-h", "--help"],
}

@click.command(no_args_is_help=True, context_settings=CONTEXT_SETTINGS)
@click.argument(
    "training_file",
    type=click.Path(exists=True, resolve_path=True),
)
@click.argument(
    "validation_file",
    type=click.Path(exists=True, resolve_path=True),
)
@click.option(
    "-b", "--batch-size",
    help="Batch size.",
    type=int,
    default=2**6,
    show_default=True,
)
@click.option(
    "-d", "--debugging",
    help="Debugging mode.",
    is_flag=True,
)
@click.option(
    "-i", "--initialize",
    help="Number of models to initialize.",
    type=int,
    default=1,
    show_default=True,
)
@click.option(
    "-o", "--output-dir",
    help="Output directory.",
    type=click.Path(resolve_path=True),
    default="./",
    show_default=True,
)
@click.option(
    "-r", "--rev-complement",
    help="Reverse complement training sequences.",
    is_flag=True,
)
@optgroup.group("Optimizer")
@optgroup.option(
    "--lr",
    help="Learning rate for the `torch.optim.Adam` optimizer.",
    type=float,
    default=1e-04,
    show_default=True,
)
@optgroup.group("Selene")
@optgroup.option(
    "--checkpoint-resume",
    help="Resume training from model file.",
    type=click.Path(exists=True, resolve_path=True),
)
@optgroup.option(
    "--cpu-threads",
    help="Number of CPU threads to use.",
    type=int,
    default=1,
    show_default=True,
)
@optgroup.option(
    "--max-epochs",
    help="Number of epochs to train.",
    type=int,
    default=128,
    show_default=True,
)
@optgroup.option(
    "--patience",
    help="If no improvement, wait `n` epochs before stopping training.",
    type=int,
    default=16,
    show_default=True,
)

def main(**args):

    ##############
    # Load Data  #
    ##############

    # Get data
    Xs_train, ys_train, _ = _get_Xs_ys_ids(args["training_file"],
        args["debugging"], args["rev_complement"])
    Xs_val, ys_val, _ = _get_Xs_ys_ids(args["validation_file"],
        args["debugging"], args["rev_complement"])

    # Get DataLoaders
    dl_train = _get_data_loader(Xs_train, ys_train, args["batch_size"], True)
    dl_val = _get_data_loader(Xs_val, ys_val, args["batch_size"], True)
    data_loaders = dict({"train": dl_train, "validation": dl_val})

    ##############
    # Train      #
    ##############

    # Initialize
    sequence_length = Xs_train[0].shape[1]
    n_features = ys_train[0].shape[0]
    if np.unique(ys_train[:, 0]).size == 2:
        input_data = "binary"
    else:
        input_data = "linear"
    steps_per_epoch = math.ceil(len(dl_train.dataset) / dl_train.batch_size)

    # For each initialization...
    best_init_loss = None
    best_init_model = None
    for i in range(args["initialize"]):

        # Skip if already done
        d = os.path.join(args["output_dir"], f"init.{i}")
        if not os.path.isdir(d):

            # Create output dir
            os.makedirs(d)

            # Train for a few epoch
            _train(sequence_length, n_features, data_loaders, input_data,
                steps_per_epoch, args["lr"], 5, args["patience"],
                args["cpu_threads"], d)

        # Get best initialized
        df = pd.read_csv(os.path.join(d, "validation.txt"), sep="\t")
        init_loss = df.loss.min()
        init_model = os.path.join(d, "best_model.pth.tar")
        if best_init_model is None:
            best_init_loss = init_loss
            best_init_model = init_model
        elif init_loss < best_init_loss:
            best_init_loss = init_loss
            best_init_model = init_model

    # Train
    shutil.copy(best_init_model, args["output_dir"])
    _train(sequence_length, n_features, data_loaders, input_data,
        steps_per_epoch, args["lr"], args["max_epochs"], args["patience"],
        args["cpu_threads"], args["output_dir"], best_init_model)

def _train(sequence_length, n_features, data_loaders, input_data,
           steps_per_epoch, lr=1e-03, max_epochs=128, patience=16,
           cpu_threads=1, output_dir="./", checkpoint_resume=None):

    # Train
    model = DanQ(sequence_length, n_features)
    loss_criterion = get_loss(input_data=input_data)
    metrics = get_metrics(input_data=input_data)
    optimizer = get_optimizer(model.parameters(), lr)
    trainer = Trainer(
        model,
        data_loaders,
        loss_criterion,
        metrics,
        optimizer,
        max_steps=steps_per_epoch*max_epochs,
        patience=steps_per_epoch*patience,
        report_stats_every_n_steps=steps_per_epoch,
        output_dir=output_dir,
        cpu_n_threads=cpu_threads,
        use_cuda=torch.cuda.is_available(),
        checkpoint_resume=checkpoint_resume,
    )
    trainer.train_and_validate()

if __name__ == "__main__":
    main()


import logging
import atexit
import pyaudio
import numpy
from scipy.signal import firwin, lfilter

_audio = pyaudio.PyAudio()

def get_available_input_devices() -> dict:
    return {i: _audio.get_device_info_by_index(i)['name'] for i in range(_audio.get_device_count())}

def get_available_output_devices() -> dict:
    return {i: _audio.get_device_info_by_index(i)['name'] for i in range(_audio.get_device_count())}

class Microphone:
    def __init__(self,
                 frames_per_buffer: int,
                 frame_rate: int,
                 input_device_index: int,
                 cutoff: int = 70
                 ):
        self._stream = _audio.open(
            format=pyaudio.paInt16,
            channels=1,
            rate=frame_rate,
            input_device_index=input_device_index,
            input=True
        )
        self._frame_per_buffer = frames_per_buffer
        if cutoff == 0:
            self._read_frame_impl = self._read_frame_wo_lowpassfilter
        else:
            self._read_frame_impl = self._read_frame_with_lowpassfilter
            self._lowcutfilter = firwin(255, cutoff / (frame_rate // 2), pass_zero=False)
        atexit.register(self.close_stream)

    def read_frame(self) -> numpy.ndarray:
        '''
        Returns
        -------
        1フレーム内の音声情報
        '''
        return self._read_frame_impl()

    def _read_frame_with_lowpassfilter(self):
        wav_bytes = self._stream.read(self._frame_per_buffer)
        wav_array = numpy.frombuffer(wav_bytes, dtype=numpy.core.int16)
        wav_array = lfilter(self._lowcutfilter, 1, wav_array)
        return wav_array.astype(numpy.core.int16)

    def _read_frame_wo_lowpassfilter(self):
        wav_bytes = self._stream.read(self._frame_per_buffer)
        wav_array = numpy.frombuffer(wav_bytes, dtype=numpy.core.int16)
        return wav_array

    def close_stream(self):
        self._stream.close()

class Speaker:
    def __init__(self,
                 frames_per_buffer: int,
                 frame_rate: int,
                 output_device_index: int
                 ):
        self._stream = _audio.open(
            format=pyaudio.paInt16,
            channels=1,
            rate=frame_rate,
            frames_per_buffer=frames_per_buffer,
            output_device_index=output_device_index,
            output=True
        )
        atexit.register(self.close_stream)

    def play_frame(self, wav_data: numpy.ndarray):
        self._stream.write(wav_data.tobytes())

    def close_stream(self):
        self._stream.close()


import numpy as np
from collections import OrderedDict

import torch
import torch.optim as optim
from torch import nn
import torch.nn.functional as F
import gtimer as gt

import rlkit.torch.utils.pytorch_util as ptu
from rlkit.torch.core import np_to_pytorch_batch
from rlkit.torch.algorithms.torch_rl_algorithm import TorchRLAlgorithm
from rlkit.data_management.relabel_replay_buffer import HindsightReplayBuffer
from rlkit.data_management.path_builder import PathBuilder
from rlkit.samplers import PathSampler

class HER(TorchRLAlgorithm):
    """
    Hindsight Experience Replay. Default using TD3 for RL policy trainer.
    """

    def __init__(
        self,
        replay_buffer=None,
        her_ratio=0.8,
        relabel_type="future",
        **kwargs
    ):
        if replay_buffer is None:
            assert kwargs['max_path_length'] < kwargs['replay_buffer_size']
            replay_buffer = HindsightReplayBuffer(
                kwargs['replay_buffer_size'], kwargs['env'], random_seed=np.random.randint(10000), relabel_type=relabel_type, her_ratio=her_ratio
            )
        super().__init__(replay_buffer=replay_buffer, **kwargs)

    def get_batch(self, keys=None):
        buffer = self.replay_buffer
        batch = buffer.random_batch(self.batch_size, keys=keys)
        batch = np_to_pytorch_batch(batch)
        return batch

    def _get_action_and_info(self, observation):
        """
        Get an action to take in the environment.
        :param observation:
        :return:
        """
        self.exploration_policy.set_num_steps_total(self._n_env_steps_total)
        return self.exploration_policy.get_actions(
            observation,
        )


import os
import time
import numpy as np
import dtlpy as dl
from nsfw_detector import predict
# import keras  # ImportError: cannot import name 'get_config' from 'tensorflow.python.eager.context'
from tensorflow import keras

class ModelAdapter(dl.BaseModelAdapter):
    """
    Specific Model adapter.
    The class bind Dataloop model and snapshot entities with model code implementation
    """
    # TODO:
    #   1) docstring for your ModelAdapter
    #   2) implement the virtual methods for full adapter support
    #   3) add your _defaults

    _defaults = {
        "input_shape": 299,
        "default_weights_filename": 'nsfw.299x299.h5'  # https://s3.amazonaws.com/nsfwdetector/nsfw.299x299.h5
    }

    def __init__(self, model_entity):
        super(ModelAdapter, self).__init__(model_entity)

    # ===============================
    # NEED TO IMPLEMENT THESE METHODS
    # ===============================

    def load(self, local_path=None, **kwargs):
        """ Loads model and populates self.model with a `runnable` model
            This function is called by load_from_snapshot (download to local and then loads)
        :param local_path: `str` directory path in local FileSystem default is current dir
        """
        if local_path is None:
            local_path = '.'
        if getattr(self, 'weights_filename', False):
            fname = self.weights_filename
        else:
            fname = self.default_weights_filename

        model_path = os.path.join(local_path, fname)
        if abs(os.stat(model_path).st_size - 158652512) > 10:
            # Temp fix to mitigate the gcs background download
            print('Waiting for dowlonad..... {}'.format(os.stat(model_path).st_size))
            time.sleep(10)

        if not os.path.isfile(model_path):
            raise RuntimeError("model path {!r} does not exsist".format(model_path))

        self.model = predict.load_model(model_path)

    def save(self, local_path, **kwargs):
        """ saves configuration and weights locally

            Virtual method - need to implement

            the function is called in save_to_snapshot which first save locally and then uploads to snapshot entity

        :param local_path: `str` directory path in local FileSystem
        """
        raise NotImplementedError("Please implement 'save' method in {}".format(self.__class__.__name__))

    def train(self, data_path, dump_path, **kwargs):
        """ Train the model according to data in local_path and save the snapshot to dump_path

            Virtual method - need to implement
        :param data_path: `str` local File System path to where the data was downloaded and converted at
        :param dump_path: `str` local File System path where to dump training mid-results (checkpoints, logs...)
        """
        raise NotImplementedError("Please implement 'train' method in {}".format(self.__class__.__name__))

    def predict(self, batch, **kwargs):
        """ Model inference (predictions) on batch of images

            Virtual method - need to implement

        :param batch: `np.ndarray`
        :return: `list[dl.AnnotationCollection]` each collection is per each image / item in the batch
        """
        batch = self.prepare_batch(batch, target_size=self.input_shape)
        preds = predict.classify_nd(self.model, batch)
        self.logger.debug("Prediction: {}".format(preds))

        batch_annotations = []
        for pred in preds:
            max_label = max(pred, key=pred.get)
            prob = pred[max_label]
            batch_annotations.append(
                dl.ml.predictions_utils.add_classification( label=max_label, score=prob, adapter=self)
            )
        return batch_annotations

    def predict_video(self, item: dl.Item, with_upload=True, progress=None):
        # item.mimetype.split('/')[0] == 'video'  # image
        batch_size=4
        import cv2
        target_path = os.path.join('/tmp', item.filename[1:])
        item.download(target_path)

        cap = cv2.VideoCapture(target_path)
        ret = True
        vid_arr = []
        while(cap.isOpened() and ret):
            ret, frame = cap.read()
            if ret:
                vid_arr.append(frame)
            else:
                break

        results = []
        for i in range(0,len(vid_arr), batch_size):
            batch = self.prepare_batch(vid_arr[i:i+batch_size], target_size=self.input_shape)
            predictions = predict.classify_nd(self.model, batch)
            results += predictions

        # calculate the average of all predictions classes
        sum_results = {}
        for label in self.snapshot.label_map.keys():
            sum_results[label] = sum([res[label] for res in results])

        max_label = max(sum_results, key=sum_results.get)
        normalize_scores = {k: v/len(results) for k, v in sum_results.items()}
        score = normalize_scores[max_label]
        print("Predict: {} ({:1.2f}).  normalized scores: {}".format(max_label, score, normalize_scores))

        # classify item with the single prediction
        preds = dl.ml.predictions_utils.add_classification(label=max_label, score=score, adapter=self)
        os.remove(target_path)

        if with_upload:
            try:
                item.annotations.upload(preds)
                self._upload_predictions_collection(item, predictions=preds, cleanup=False)
            except dl.exceptions.InternalServerError as err:
                self.logger.error("Failed to upload annotations items. Error: {}".format(err))
        return preds

    @staticmethod
    def prepare_batch(batch, target_size):
        """
        resize and scals the image to the proper input of the network
        """
        from skimage.transform import resize
        loaded_images = []
        for image in batch:
            # image = keras.preprocessing.image.load_img(img_path, target_size=target_size)
            # image = keras.preprocessing.image.img_to_array(image)
            image = resize(image, (target_size, target_size))
            image /= 255
            loaded_images.append(image)
        return np.array(loaded_images)

    def convert(self, data_path, **kwargs):
        """ Convert Dataloop structure data to model structured

            Virtual method - need to implement

            e.g. take dlp dir structure and construct annotation file

        :param data_path: `str` local File System directory path where
                           we already downloaded the data from dataloop platform
        :return:
        """
        raise NotImplementedError("Please implement 'convert' method in {}".format(self.__class__.__name__))

    # =============
    # DTLPY METHODS
    # Do not change
    # =============
    # function are here to ease the traceback...
    def load_from_snapshot(self, local_path, snapshot_id, **kwargs):
        """ Loads a model from given `snapshot`.
            Reads configurations and instantiate self.snapshot
            Downloads the snapshot bucket (if available)

        :param local_path:  `str` directory path in local FileSystem to download the snapshot to
        :param snapshot_id:  `str` snapshot id
        """
        super(ModelAdapter, self).load_from_snapshot(local_path=local_path, snapshot_id=snapshot_id, **kwargs)

    def save_to_snapshot(self, local_path, snapshot_name=None, description=None, cleanup=False, **kwargs):
        """ Saves configuration and weights to new snapshot bucket
            loads only applies for remote buckets

        :param local_path: `str` directory path in local FileSystem
        :param snapshot_name: `str` name for the new snapshot
        :param description:  `str` description for the new snapshot
        :param cleanup: `bool` if True (default) remove the data from local FileSystem after upload
        :return:
        """
        super(ModelAdapter, self).save_to_snapshot(local_path=local_path, snapshot_name=snapshot_name,
                                                   description=description, cleanup=cleanup,
                                                   **kwargs)

    def prepare_trainset(self, data_path, partitions=None, filters=None, **kwargs):
        """
        Prepares train set for train.
        download the specific partition selected to data_path and preforms `self.convert` to the data_path dir

        :param data_path: `str` directory path to use as the root to the data from Dataloop platform
        :param partitions: `dl.SnapshotPartitionType` or list of partitions, defaults for all partitions
        :param filters: `dl.Filter` in order to select only part of the data
        """
        super(ModelAdapter, self).prepare_trainset(data_path=data_path, partitions=partitions,
                                                   filters=filters,
                                                   **kwargs)

    def predict_items(self, items: list, with_upload=True, cleanup=False, batch_size=16, output_shape=None, **kwargs):
        """
        Predicts all items given

        :param items: `List[dl.Item]`
        :param with_upload: `bool` uploads the predictions back to the given items
        :param cleanup: `bool` if set removes existing predictions with the same model-snapshot name (default: False)
        :param batch_size: `int` size of batch to run a single inference
        :param output_shape: `tuple` (width, height) of resize needed per image
        :return: `List[Prediction]' `Prediction is set by model.output_type
        """
        super(ModelAdapter, self).predict_items(items=items, with_upload=with_upload,
                                                cleanup=cleanup, batch_size=batch_size, output_shape=output_shape,
                                                **kwargs)


from model.texture import *
from model.obj import *
from renderer import *
import glm
from glfw.GLFW import glfwSetCursor, glfwSetCursorPos
from Display import *
from OpenGL.GL import *
from model.loader import *
from shader import *
from tool.mmath import *
from camera import *
import glfw
import numpy as np
from pynput import keyboard

WIDTH = 1366
HEIGHT = 768

display = Display()

def gridGen(size, blocksize):
    w, h = size
    bs = blocksize
    w = int(w/bs)
    h = int(h/bs)
    out = []
    vout = []
    for i in range(h+bs):
        a = []
        b = []
        for j in range(w+bs):
            a.append([i, j, 0])
            b.append(i)
            b.append(j)
            b.append(0)
        [out.append(e) for e in a]
        [vout.append(e) for e in b]
    iout = []
    for i in range(h):
        for j in range(w):
            iout.append(out.index([i, j, 0]))
            iout.append(out.index([i, j+1, 0]))
            iout.append(out.index([i+1, j, 0]))
            iout.append(out.index([i, j+1, 0]))
            iout.append(out.index([i+1, j, 0]))
            iout.append(out.index([i+1, j+1, 0]))
    return vout, iout

gridGen((2, 2), 1)

gridData, gridIndices = gridGen((2, 2), 1)

gridData = [e/4 for e in gridData]

data = [
    -0.5, -0.5, 0.0,  # 0 Bottom Left
    -0.5,  0.5, 0.0,  # 1 Top Left
    0.5,  0.5, 0.0,  # 2 Top Right
    0.5, -0.5, 0.0  # 3 Bottom Right
]

indices = [
    0, 1, 2,
    2, 0, 3
]

line_indices = [
    0, 1,
    1, 2,
    2, 3,
    3, 0
]
data = np.array(data, dtype=np.float32)
indices = np.array(indices, dtype=np.int32)
line_indices = np.array(line_indices, dtype=np.int32)
print(data)

cameraZ = 0.0

def main():
    window = display.create(WIDTH, HEIGHT, 'GL')
    shader = Shader()
    camera = Camera(window, glm.vec3(0.0, 0.0, 1.1),
                    0.0, 0.0, 0.0, 60, 0.01, -1000)
    shader.start()
    shader.putVecInUniformLocation('color', glm.vec4(0.01, 1.0, 0.01, 1.0))

    proj = camera.createProjectionMatrix()
    shader.loadProjectionMatrix(glm.value_ptr(proj))

    view = camera.createViewMatrix()
    shader.loadViewMatrix(glm.value_ptr(view))

    vb = glGenBuffers(1)
    glEnableVertexAttribArray(0)
    glBindBuffer(GL_ARRAY_BUFFER, vb)
    glBufferData(GL_ARRAY_BUFFER, data, GL_STATIC_DRAW)
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0, None)

    ivb = glGenBuffers(1)
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, ivb)
    glBufferData(GL_ELEMENT_ARRAY_BUFFER, indices, GL_STATIC_DRAW)

    glPolygonMode(GL_FRONT_AND_BACK, GL_LINE)

    while display.isNotClosed():
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        glClearColor(0.1, 0.0, 0.03, 1.0)
        shader.loadTransformationMatrix(glm.value_ptr(
            createTransformationMatrix(
                glm.vec3(0.0, 0.0, -0.1),
                glm.vec3(0.0, 0.1*glfw.get_time(), 0.0*glfw.get_time()),
                glm.vec3(1.0, 1.0, 1.0)
            )
        ))
        camera.movement()
        shader.loadViewMatrix(glm.value_ptr(camera.createViewMatrix()))
        glPointSize(3.0)
        glDrawElements(GL_TRIANGLES, len(indices)+1, GL_UNSIGNED_INT, None)
        display.update()
    shader.stop()
    glDeleteBuffers(1, [vb])
    display.close()

if __name__ == '__main__':
    main()

# --------------------------------------------------------------------------------------------------------

WIDTH = 1366
HEIGHT = 768

display = Display()

def gridGen(size, blocksize):
    w, h = size
    bs = blocksize
    w = int(w/bs)
    h = int(h/bs)
    out = []
    vout = []
    for i in range(h+bs):
        a = []
        b = []
        for j in range(w+bs):
            a.append([i, j, 0])
            b.append(i)
            b.append(j)
            b.append(0)
        [out.append(e) for e in a]
        [vout.append(e) for e in b]
    iout = []
    for i in range(h):
        for j in range(w):
            iout.append(out.index([i, j, 0]))
            iout.append(out.index([i, j+1, 0]))
            iout.append(out.index([i+1, j, 0]))
            iout.append(out.index([i, j+1, 0]))
            iout.append(out.index([i+1, j, 0]))
            iout.append(out.index([i+1, j+1, 0]))
    return vout, iout

data = [
    -0.5, -0.5, 0.0,  # 0 Bottom Left
    -0.5,  0.5, 0.0,  # 1 Top Left
    0.5,  0.5, 0.0,  # 2 Top Right
    0.5, -0.5, 0.0  # 3 Bottom Right
]

indices = [
    0, 1, 2,
    2, 0, 3
]

line_indices = [
    0, 1,
    1, 2,
    2, 3,
    3, 0
]
data = np.array(data, dtype=np.float32)
indices = np.array(indices, dtype=np.int32)
line_indices = np.array(line_indices, dtype=np.int32)

cameraZ = 0.0

def main():
    window = display.create(WIDTH, HEIGHT, 'GL')
    shader = Shader()
    glfwSetCursorPos(window, WIDTH/2, HEIGHT/2)
    camera = Camera(window, glm.vec3(0.0, 0.0, 1.1),
                    0.0, 0.0, 0.0, 60, 0.01, -1000)
    renderer = Renderer(camera, window)
    loader = Loader()
    obj = OBJ('dummy.obj')
    entity = Entity(
        #loader.loadModel(obj.loadObj(), 'tex5.jpg'),
        loader.loadRaw(data, indices),
        glm.vec3(1.0, 0.0, -1.1),
        glm.vec3(0.0, 0.0, 0.0),
        glm.vec3(1.0, 1.0, 1.0),
        mode='FILL'
    )
    entity0 = Entity(
        loader.loadModel(obj.loadObj(), 'tex3.jpg'),
        glm.vec3(0.0, 0.0, -1.0),
        glm.vec3(0.0, 0.0, 0.0),
        glm.vec3(1.0, 1.0, 1.0),
        mode='FILL',
        color=glm.vec4(1.0, 1.0, 0.0, 0.3)
    )
    renderer.process(entity, shader, hasIndex=True)
    while display.isNotClosed():
        renderer.render()
        display.update()
    display.close()

if __name__ == '__main__':
    main()


from fastai.vision.all import *
import streamlit as st
from PIL import Image
import numpy as np

import pathlib

learn = load_learner(pathlib.PosixPath('./export.pkl'))

st.write("""
         # Poisonous/Edible Mushroom Prediction
         # **DO NOT** USE IN REAL LIFE SCENARIO
         """
         )
st.write("This is an PoC on transforming models into web apps. The accuracy of the model is not the main focus,**DO NOT** use it as a real life mushroom classifier")
file = st.file_uploader("Please upload an image file", type=["jpg", "png"])

if file is None:
    st.text("Please upload an image file")
else:
    image = Image.open(file)
    image = PILImage.create(np.asarray(image))
    tfms = Normalize.from_stats(*imagenet_stats)
    image = tfms(image)
    st.image(image, use_column_width=True)
    pred_class,_,_ = learn.predict(image)

    st.write('''It is {}!
    **PREDICTION NOT ACCURATE PLEASE DO *NOT* CONSUME THE MUSHROOM**'''.format(pred_class))


# Authors: Alexandre Gramfort <alexandre.gramfort@telecom-paristech.fr>
#          Matti Hamalainen <msh@nmr.mgh.harvard.edu>
#          Martin Luessi <mluessi@nmr.mgh.harvard.edu>
#          Eric Larson <larsoner@uw.edu>
#
# License: BSD (3-clause)

import numpy as np
from copy import deepcopy

from ..surface import (fast_cross_3d, _find_nearest_tri_pt, _get_tri_supp_geom,
                       _triangle_coords)
from ..io.constants import FIFF
from ..transforms import apply_trans
from ..utils import logger
from ..parallel import parallel_func
from ..io.compensator import get_current_comp, make_compensator
from ..io.pick import pick_types

##############################################################################
# COIL SPECIFICATION

def _dup_coil_set(coils, coord_frame, t):
    """Make a duplicate"""
    if t is not None and coord_frame != t['from']:
        raise RuntimeError('transformation frame does not match the coil set')
    coils = deepcopy(coils)
    if t is not None:
        coord_frame = t['to']
        for coil in coils:
            coil['r0'] = apply_trans(t['trans'], coil['r0'])
            coil['ex'] = apply_trans(t['trans'], coil['ex'], False)
            coil['ey'] = apply_trans(t['trans'], coil['ey'], False)
            coil['ez'] = apply_trans(t['trans'], coil['ez'], False)
            coil['rmag'] = apply_trans(t['trans'], coil['rmag'])
            coil['cosmag'] = apply_trans(t['trans'], coil['cosmag'], False)
            coil['coord_frame'] = t['to']
    return coils, coord_frame

def _check_coil_frame(coils, coord_frame, bem):
    """Check to make sure the coils are in the correct coordinate frame"""
    if coord_frame != FIFF.FIFFV_COORD_MRI:
        if coord_frame == FIFF.FIFFV_COORD_HEAD:
            # Make a transformed duplicate
            coils, coord_Frame = _dup_coil_set(coils, coord_frame,
                                               bem['head_mri_t'])
        else:
            raise RuntimeError('Bad coil coordinate frame %d' % coord_frame)
    return coils, coord_frame

def _bem_lin_field_coeffs_simple(dest, normal, tri_rr, tri_nn, tri_area):
    """Simple version..."""
    out = np.zeros((3, len(dest)))
    for rr, o in zip(tri_rr, out):
        diff = dest - rr
        dl = np.sum(diff * diff, axis=1)
        x = fast_cross_3d(diff, tri_nn[np.newaxis, :])
        o[:] = tri_area * np.sum(x * normal, axis=1) / (3.0 * dl * np.sqrt(dl))
    return out

def _lin_field_coeff(s, mult, rmags, cosmags, ws, counts, func, n_jobs):
    """Use the linear field approximation to get field coefficients"""
    parallel, p_fun, _ = parallel_func(_do_lin_field_coeff, n_jobs)
    nas = np.array_split
    coeffs = parallel(p_fun(s['rr'], t, tn, ta,
                            rmags, cosmags, ws, counts, func)
                      for t, tn, ta in zip(nas(s['tris'], n_jobs),
                                           nas(s['tri_nn'], n_jobs),
                                           nas(s['tri_area'], n_jobs)))
    return mult * np.sum(coeffs, axis=0)

def _do_lin_field_coeff(rr, t, tn, ta, rmags, cosmags, ws, counts, func):
    """Actually get field coefficients (parallel-friendly)"""
    coeff = np.zeros((len(counts), len(rr)))
    bins = np.repeat(np.arange(len(counts)), counts)
    for tri, tri_nn, tri_area in zip(t, tn, ta):
        # Accumulate the coefficients for each triangle node
        # and add to the corresponding coefficient matrix
        tri_rr = rr[tri]

        # The following is equivalent to:
        #for j, coil in enumerate(coils['coils']):
        #    x = func(coil['rmag'], coil['cosmag'],
        #             tri_rr, tri_nn, tri_area)
        #    res = np.sum(coil['w'][np.newaxis, :] * x, axis=1)
        #    coeff[j][tri + off] += mult * res

        xx = func(rmags, cosmags, tri_rr, tri_nn, tri_area)
        # only loops 3x (one per direction)
        zz = np.array([np.bincount(bins, weights=x * ws,
                                   minlength=len(counts)) for x in xx])
        coeff[:, tri] += zz.T
    return coeff

def _bem_specify_coils(bem, coils, coord_frame, n_jobs):
    """Set up for computing the solution at a set of coils"""
    # Compute the weighting factors to obtain the magnetic field
    # in the linear potential approximation
    coils, coord_frame = _check_coil_frame(coils, coord_frame, bem)

    # leaving this in in case we want to easily add in the future
    #if method != 'simple':  # in ['ferguson', 'urankar']:
    #    raise NotImplementedError
    #else:
    func = _bem_lin_field_coeffs_simple

    # Process each of the surfaces
    rmags = np.concatenate([coil['rmag'] for coil in coils])
    cosmags = np.concatenate([coil['cosmag'] for coil in coils])
    counts = np.array([len(coil['rmag']) for coil in coils])
    ws = np.concatenate([coil['w'] for coil in coils])

    lens = np.cumsum(np.r_[0, [len(s['rr']) for s in bem['surfs']]])
    coeff = np.empty((len(counts), lens[-1]))
    for o1, o2, surf, mult in zip(lens[:-1], lens[1:],
                                  bem['surfs'], bem['field_mult']):
        coeff[:, o1:o2] = _lin_field_coeff(surf, mult, rmags, cosmags,
                                           ws, counts, func, n_jobs)
    # put through the bem
    sol = np.dot(coeff, bem['solution'])
    return sol

def _bem_specify_els(bem, els):
    """Set up for computing the solution at a set of electrodes"""
    sol = np.zeros((len(els), bem['solution'].shape[1]))
    # Go through all coils
    scalp = bem['surfs'][0]
    scalp['geom'] = _get_tri_supp_geom(scalp['tris'], scalp['rr'])
    inds = np.arange(len(scalp['tris']))

    # In principle this could be parallelized, but pickling overhead is huge
    # (makes it slower than non-parallel)
    for k, el in enumerate(els):
        # Go through all 'integration points'
        el_r = apply_trans(bem['head_mri_t']['trans'], el['rmag'])
        for elw, r in zip(el['w'], el_r):
            best = _find_nearest_tri_pt(inds, r, scalp['geom'], True)[2]
            # Calculate a linear interpolation between the vertex values
            tri = scalp['tris'][best]
            x, y, z = _triangle_coords(r, scalp['geom'], best)
            w = elw * np.array([(1.0 - x - y), x, y])
            amt = np.dot(w, bem['solution'][tri])
            sol[k] += amt
    return sol

#############################################################################
# FORWARD COMPUTATION

def _make_ctf_comp_coils(info, coils):
    """Get the correct compensator for CTF coils"""
    # adapted from mne_make_ctf_comp() from mne_ctf_comp.c
    logger.info('Setting up compensation data...')
    comp_num = get_current_comp(info)
    if comp_num is None or comp_num == 0:
        logger.info('    No compensation set. Nothing more to do.')
        return None

    # Need to meaningfully populate comp['set'] dict a.k.a. compset
    n_comp_ch = sum([c['kind'] == FIFF.FIFFV_MEG_CH for c in info['chs']])
    logger.info('    %d out of %d channels have the compensation set.'
                % (n_comp_ch, len(coils)))

    # Find the desired compensation data matrix
    compensator = make_compensator(info, 0, comp_num, True)
    logger.info('    Desired compensation data (%s) found.' % comp_num)
    logger.info('    All compensation channels found.')
    logger.info('    Preselector created.')
    logger.info('    Compensation data matrix created.')
    logger.info('    Postselector created.')
    return compensator

#def _bem_inf_pot(rd, Q, rp):
#    """The infinite medium potential in one direction"""
#    # NOTE: the (4.0 * np.pi) that was in the denominator has been moved!
#    diff = rp - rd
#    diff2 = np.sum(diff * diff, axis=1)
#    return np.sum(Q * diff, axis=1) / (diff2 * np.sqrt(diff2))

def _bem_inf_pots(rr, surf_rr, Q=None):
    """The infinite medium potential in all 3 directions"""
    # NOTE: the (4.0 * np.pi) that was in the denominator has been moved!
    diff = surf_rr.T[np.newaxis, :, :] - rr[:, :, np.newaxis]  # n_rr, 3, n_bem
    diff_norm = np.sum(diff * diff, axis=1)
    diff_norm *= np.sqrt(diff_norm)
    diff_norm[diff_norm == 0] = 1  # avoid nans
    if Q is None:  # save time when Q=np.eye(3) (e.g., MEG sensors)
        return diff / diff_norm[:, np.newaxis, :]
    else:  # get components in each direction (e.g., EEG sensors)
        return np.einsum('ijk,mj->imk', diff, Q) / diff_norm[:, np.newaxis, :]

# This function has been refactored to process all points simultaneously
#def _bem_inf_field(rd, Q, rp, d):
#    """Infinite-medium magnetic field"""
#    diff = rp - rd
#    diff2 = np.sum(diff * diff, axis=1)
#    x = fast_cross_3d(Q[np.newaxis, :], diff)
#    return np.sum(x * d, axis=1) / (diff2 * np.sqrt(diff2))

def _bem_inf_fields(rr, rp, c):
    """Infinite-medium magnetic field in all 3 basis directions"""
    # Knowing that we're doing all directions, the above can be refactored:
    diff = rp.T[np.newaxis, :, :] - rr[:, :, np.newaxis]
    diff_norm = np.sum(diff * diff, axis=1)
    diff_norm *= np.sqrt(diff_norm)
    diff_norm[diff_norm == 0] = 1  # avoid nans
    # This is the result of cross-prod calcs with basis vectors,
    # as if we had taken (Q=np.eye(3)), then multiplied by the cosmags (c)
    # factor, and then summed across directions
    x = np.array([diff[:, 1] * c[:, 2] - diff[:, 2] * c[:, 1],
                  diff[:, 2] * c[:, 0] - diff[:, 0] * c[:, 2],
                  diff[:, 0] * c[:, 1] - diff[:, 1] * c[:, 0]])
    return np.rollaxis(x / diff_norm, 1)

def _bem_pot_or_field(rr, mri_rr, mri_Q, mults, coils, solution, srr,
                      n_jobs, coil_type):
    """Calculate the magnetic field or electric potential

    The code is very similar between EEG and MEG potentials, so we'll
    combine them.
    """
    # multiply solution by "mults" here for simplicity
    # we can do this one in-place because it's not used elsewhere
    solution *= mults

    # Both MEG and EEG have the inifinite-medium potentials
    # This could be just vectorized, but eats too much memory, so instead we
    # reduce memory by chunking within _do_inf_pots and parallelize, too:
    parallel, p_fun, _ = parallel_func(_do_inf_pots, n_jobs)
    nas = np.array_split
    B = np.sum(parallel(p_fun(mri_rr, sr.copy(), mri_Q, sol.copy())
                        for sr, sol in zip(nas(srr, n_jobs),
                                           nas(solution.T, n_jobs))), axis=0)
    # The copy()s above should make it so the whole objects don't need to be
    # pickled...

    # Only MEG gets the primary current distribution
    if coil_type == 'meg':
        # Primary current contribution (can be calc. in coil/dipole coords)
        parallel, p_fun, _ = parallel_func(_do_prim_curr, n_jobs)
        pcc = np.concatenate(parallel(p_fun(rr, c)
                                      for c in nas(coils, n_jobs)), axis=1)
        B += pcc
        B *= 1e-7  # MAG_FACTOR from C code
    return B

def _do_prim_curr(rr, coils):
    """Calculate primary currents in a set of coils"""
    out = np.empty((len(rr) * 3, len(coils)))
    for ci, c in enumerate(coils):
        out[:, ci] = np.sum(c['w'] * _bem_inf_fields(rr, c['rmag'],
                                                     c['cosmag']), 2).ravel()
    return out

def _do_inf_pots(rr, srr, mri_Q, sol):
    """Calculate infinite potentials using chunks"""
    # The following code is equivalent to this, but saves memory
    #v0s = _bem_inf_pots(rr, srr, mri_Q)  # n_rr x 3 x n_surf_rr
    #v0s.shape = (len(rr) * 3, v0s.shape[2])
    #B = np.dot(v0s, sol)

    # We chunk the source rr's in order to save memory
    bounds = np.r_[np.arange(0, len(rr), 1000), len(rr)]
    B = np.empty((len(rr) * 3, sol.shape[1]))
    for bi in range(len(bounds) - 1):
        v0s = _bem_inf_pots(rr[bounds[bi]:bounds[bi + 1]], srr, mri_Q)
        v0s.shape = (v0s.shape[0] * 3, v0s.shape[2])
        B[3 * bounds[bi]:3 * bounds[bi + 1]] = np.dot(v0s, sol)
    return B

def _compute_forwards(src, bem, coils_list, cfs, ccoils_list, ccfs,
                      infos, coil_types, n_jobs):
    """Compute the MEG and EEG forward solutions"""
    if bem['bem_method'] != 'linear collocation':
        raise RuntimeError('only linear collocation supported')

    # Precompute some things that are used for both MEG and EEG
    rr = np.concatenate([s['rr'][s['vertno']] for s in src])
    mults = np.repeat(bem['source_mult'] / (4.0 * np.pi),
                      [len(s['rr']) for s in bem['surfs']])[np.newaxis, :]
    # The dipole location and orientation must be transformed
    mri_rr = apply_trans(bem['head_mri_t']['trans'], rr)
    mri_Q = apply_trans(bem['head_mri_t']['trans'], np.eye(3), False)
    srr = np.concatenate([s['rr'] for s in bem['surfs']])

    # Now, actually compute MEG and EEG solutions
    Bs = list()
    for coil_type, coils, cf, ccoils, ccf, info in zip(coil_types, coils_list,
                                                       cfs, ccoils_list, ccfs,
                                                       infos):
        if coils is None:  # nothing to do
            Bs.append(None)
        else:
            if coil_type == 'meg':
                # Compose a compensation data set if necessary
                compensator = _make_ctf_comp_coils(info, coils)

                # Field computation matrices...
                logger.info('')
                start = 'Composing the field computation matrix'
                logger.info(start + '...')
                solution = _bem_specify_coils(bem, coils, cf, n_jobs)
                if compensator is not None:
                    logger.info(start + ' (compensation coils)...')
                    csolution = _bem_specify_coils(bem, ccoils, ccf, n_jobs)

            elif coil_type == 'eeg':
                solution = _bem_specify_els(bem, coils)
                compensator = None

            # Do the actual calculation
            logger.info('Computing %s at %d source locations '
                        '(free orientations)...'
                        % (coil_type.upper(), len(rr)))
            # Note: this function modifies "solution" in-place
            B = _bem_pot_or_field(rr, mri_rr, mri_Q, mults, coils,
                                  solution, srr, n_jobs, coil_type)

            # Compensate if needed (only done for MEG systems w/compensation)
            if compensator is not None:
                # Compute the field in the compensation coils
                work = _bem_pot_or_field(rr, mri_rr, mri_Q, mults,
                                         ccoils, csolution, srr, n_jobs,
                                         coil_type)
                # Combine solutions so we can do the compensation
                both = np.zeros((work.shape[0], B.shape[1] + work.shape[1]))
                picks = pick_types(info, meg=True, ref_meg=False)
                both[:, picks] = B
                picks = pick_types(info, meg=False, ref_meg=True)
                both[:, picks] = work
                B = np.dot(both, compensator.T)
            Bs.append(B)

    return Bs


#!/usr/bin/env python
"""
Usage:
    nearestneighbor.py [options] (--code | --query | --both) MODEL_PATH DATA_PATH

Options:
    --azure-info=<path>        Azure authentication information file (JSON). Used to load data from Azure storage.
    --distance-metric METRIC   The distance metric to use [default: cosine]
    -h --help                  Show this screen.
    --hypers-override HYPERS   JSON dictionary overriding hyperparameter values.
    --debug                    Enable debug routines. [default: False]
    --num-nns NUM              The number of NNs to visualize [default: 2]
    --distance-threshold TH    The distance threshold above which to ignore [default: 0.2]
    --max-num-items LIMIT      The maximum number of items to use. Use zero for all. [default: 0]
"""
import json
from itertools import chain
from typing import Any, Dict, List

from docopt import docopt
from dpu_utils.utils import RichPath, run_and_debug
import numpy as np
from more_itertools import take
from scipy.spatial.distance import pdist

from pygments import highlight
from pygments.lexers import get_lexer_by_name
from pygments.formatters import TerminalFormatter

import model_restore_helper

# Condensed to square from
# http://stackoverflow.com/questions/13079563/how-does-condensed-distance-matrix-work-pdist
from utils.visutils import square_to_condensed

def to_string(code: str, language: str) -> str:
    lexer = get_lexer_by_name(language, stripall=True)
    formatter = TerminalFormatter(linenos=True)
    return highlight(code, lexer, formatter)

def run(arguments) -> None:
    azure_info_path = arguments.get('--azure-info', None)
    data_path = RichPath.create(arguments['DATA_PATH'], azure_info_path)
    assert data_path.is_dir(), "%s is not a folder" % (data_path,)

    hypers_override = arguments.get('--hypers-override')
    if hypers_override is not None:
        hypers_override = json.loads(hypers_override)
    else:
        hypers_override = {}

    model_path = RichPath.create(arguments['MODEL_PATH'], azure_info_path=azure_info_path)

    model = model_restore_helper.restore(
        path=model_path,
        is_train=False,
        hyper_overrides=hypers_override)

    num_elements_to_take = int(arguments['--max-num-items'])
    data = chain(*chain(list(f.read_by_file_suffix()) for f in data_path.iterate_filtered_files_in_dir('*.jsonl.gz')))
    if num_elements_to_take == 0:  # Take all
        data = list(data)
    else:
        assert num_elements_to_take > 0
        data = take(num_elements_to_take, data)

    num_nns = int(arguments['--num-nns'])

    if arguments['--code']:
        representations = model.get_code_representations(data)
    elif arguments['--query']:
        representations = model.get_query_representations(data)
    else:
        code_representations = model.get_code_representations(data)
        query_representations = model.get_query_representations(data)
        representations = np.concatenate([code_representations, query_representations], axis=-1)

    filtered_representations = []
    filtered_data = []  # type: List[Dict[str, Any]]
    for i, representation in enumerate(representations):
        if representation is None:
            continue
        filtered_representations.append(representation)
        filtered_data.append(data[i])

    filtered_representations = np.stack(filtered_representations, axis=0)
    flat_distances = pdist(filtered_representations, arguments['--distance-metric'])

    for i, data in enumerate(filtered_data):
        distance_from_i = np.fromiter(
            (flat_distances[square_to_condensed(i, j, len(filtered_data))] if i != j else float('inf') for j in
             range(len(filtered_data))), dtype=np.float)

        nns = [int(k) for k in np.argsort(distance_from_i)[:num_nns]]  # The first two NNs

        if distance_from_i[nns[0]] > float(arguments['--distance-threshold']):
            continue

        print('===============================================================')
        print(f"{data['repo']}:{data['path']}:{data['lineno']}")
        print(to_string(data['original_string'], language=data['language']))

        for j in range(num_nns):
            print()
            print(f'Nearest Neighbour {j+1}: {filtered_data[nns[j]]["repo"]}:{filtered_data[nns[j]]["path"]}:{filtered_data[nns[j]]["lineno"]} (distance {distance_from_i[nns[j]]})')
            print(to_string(filtered_data[nns[j]]['original_string'], language=filtered_data[nns[j]]['language']))

if __name__ == '__main__':
    args = docopt(__doc__)
    run_and_debug(lambda: run(args), args.get('--debug', False))