#!/usr/bin/python3

from __future__ import unicode_literals, print_function, division
from io import open
import unicodedata
import string
import re
import random
import os
import torch
import torch.nn as nn
from torch.autograd import Variable
from torch import optim
import torch.nn.functional as F
import time
import math
import argparse
from settings import hparams
import tokenize_weak

'''
Some code came from:

https://github.com/A-Jacobson/minimal-nmt/blob/master/nmt_tutorial.ipynb

This code is originally written by Sean Robertson and can be found at the following site:

http://pytorch.org/tutorials/intermediate/seq2seq_translation_tutorial.html#sphx-glr-intermediate-seq2seq-translation-tutorial-py

The MIT License (MIT)

Copyright (c) 2017 Sean Robertson

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.

'''

use_cuda = torch.cuda.is_available()
SOS_token = 0
EOS_token = 1
MAX_LENGTH = hparams['tokens_per_sentence']

eng_prefixes = [
    "i am ", "i'm ",
    "he is", "he's ",
    "she is", "she's",
    "you are", "you're ",
    "we are", "we're ",
    "they are", "they're "
]
teacher_forcing_ratio = 0.5

class LuongAttention(nn.Module):
    """
    LuongAttention from Effective Approaches to Attention-based Neural Machine Translation
    https://arxiv.org/pdf/1508.04025.pdf
    """

    def __init__(self, dim):
        super(LuongAttention, self).__init__()
        self.W = nn.Linear(dim, dim, bias=False)

    def score(self, decoder_hidden, encoder_out):
        # linear transform encoder out (seq, batch, dim)
        encoder_out = self.W(encoder_out)
        # (batch, seq, dim) | (2, 15, 50)
        encoder_out = encoder_out.permute(1, 0, 2)
        # (2, 15, 50) @ (2, 50, 1)
        return encoder_out @ decoder_hidden.permute(1, 2, 0)

    def forward(self, decoder_hidden, encoder_out):
        energies = self.score(decoder_hidden, encoder_out)
        mask = F.softmax(energies, dim=1)  # batch, seq, 1
        context = encoder_out.permute(
            1, 2, 0) @ mask  # (2, 50, 15) @ (2, 15, 1)
        context = context.permute(2, 0, 1)  # (seq, batch, dim)
        mask = mask.permute(2, 0, 1)  # (seq2, batch, seq1)
        return context, mask

class EncoderBiRNN(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(EncoderBiRNN, self).__init__()
        self.hidden_size = hidden_size
        self.embedding = nn.Embedding(input_size, hidden_size)
        self.bi_gru = nn.GRU(hidden_size, hidden_size, num_layers=1, batch_first=False,bidirectional=True)

    def forward(self, input, hidden):
        embedded = self.embedding(input).view(1, 1, -1)
        output = embedded
        bi_output, bi_hidden = self.bi_gru(output,hidden)
        bi_output = (bi_output[:, :, :self.hidden_size] +
                       bi_output[:, :, self.hidden_size:])
        return bi_output, bi_hidden

    def initHidden(self):
        result = Variable(torch.zeros(2, 1, self.hidden_size))
        if use_cuda:
            return result.cuda()
        else:
            return result

class EncoderRNN(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(EncoderRNN, self).__init__()
        self.hidden_size = hidden_size

        self.embedding = nn.Embedding(input_size, hidden_size)
        self.gru = nn.GRU(hidden_size, hidden_size)

    def forward(self, input, hidden):
        embedded = self.embedding(input).view(1, 1, -1)
        output = embedded
        output, hidden = self.gru(output, hidden)
        return output, hidden

    def initHidden(self):
        result = Variable(torch.zeros(1, 1, self.hidden_size))
        if use_cuda:
            return result.cuda()
        else:
            return result

class DecoderRNN(nn.Module):
    def __init__(self, hidden_size, output_size):
        super(DecoderRNN, self).__init__()
        self.hidden_size = hidden_size

        self.embedding = nn.Embedding(output_size, hidden_size)
        self.gru = nn.GRU(hidden_size, hidden_size)
        self.out = nn.Linear(hidden_size, output_size)
        self.softmax = nn.LogSoftmax(dim=1)

    def forward(self, input, hidden):
        output = self.embedding(input).view(1, 1, -1)
        output = F.relu(output)
        output, hidden = self.gru(output, hidden)
        output = self.softmax(self.out(output[0]))
        return output, hidden

    def initHidden(self):
        result = Variable(torch.zeros(1, 1, self.hidden_size))
        if use_cuda:
            return result.cuda()
        else:
            return result

class AttnDecoderRNN(nn.Module):
    def __init__(self, hidden_size, output_size, dropout_p=0.1, max_length=MAX_LENGTH):
        super(AttnDecoderRNN, self).__init__()
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.dropout_p = dropout_p
        self.max_length = max_length

        self.embedding = nn.Embedding(self.output_size, self.hidden_size)
        self.attn = nn.Linear(self.hidden_size * 2, self.max_length)
        self.attn_combine = nn.Linear(self.hidden_size * 2, self.hidden_size)
        self.dropout = nn.Dropout(self.dropout_p)
        self.gru = nn.GRU(self.hidden_size, self.hidden_size)
        self.out = nn.Linear(self.hidden_size, self.output_size)

    def forward(self, input, hidden, encoder_outputs):
        embedded = self.embedding(input).view(1, 1, -1)
        embedded = self.dropout(embedded)

        attn_weights = F.softmax(
            self.attn(torch.cat((embedded[0], hidden[0]), 1)), dim=1)
        attn_applied = torch.bmm(attn_weights.unsqueeze(0),
                                 encoder_outputs.unsqueeze(0))

        output = torch.cat((embedded[0], attn_applied[0]), 1)
        output = self.attn_combine(output).unsqueeze(0)

        output = F.relu(output)
        output, hidden = self.gru(output, hidden)

        output = F.log_softmax(self.out(output[0]), dim=1)
        return output, hidden, attn_weights

    def initHidden(self):
        result = Variable(torch.zeros(1, 1, self.hidden_size))
        if use_cuda:
            return result.cuda()
        else:
            return result

class Decoder(nn.Module):
    def __init__(self, target_vocab_size, embed_dim, hidden_dim, n_layers, dropout):
        super(Decoder, self).__init__()
        self.n_layers = n_layers
        self.embed = nn.Embedding(target_vocab_size, embed_dim, padding_idx=1)
        self.attention = LuongAttention(hidden_dim)
        self.gru = nn.GRU(embed_dim + hidden_dim, hidden_dim, n_layers,
                          dropout=dropout)
        self.out = nn.Linear(hidden_dim * 2, target_vocab_size)

    def forward(self, output, encoder_out, decoder_hidden):
        """
        decodes one output frame
        """
        embedded = self.embed(output)  # (1, batch, embed_dim)
        context, mask = self.attention(decoder_hidden[:-1], encoder_out)  # 1, 1, 50 (seq, batch, hidden_dim)
        rnn_output, decoder_hidden = self.gru(torch.cat([embedded, context], dim=2),
                                              decoder_hidden)
        output = self.out(torch.cat([rnn_output, context], 2))
        return output, decoder_hidden, mask

class Lang:
    def __init__(self, name):
        self.name = name
        self.word2index = {}
        self.word2count = {}
        self.index2word = {0: hparams['sol'], 1: hparams['eol']}
        self.n_words = 2  # Count SOS and EOS

    def addSentence(self, sentence):
        for word in sentence.split(' '):
            self.addWord(word)

    def addWord(self, word):
        if word not in self.word2index:
            self.word2index[word] = self.n_words
            self.word2count[word] = 1
            self.index2word[self.n_words] = word
            self.n_words += 1
        else:
            self.word2count[word] += 1

################################
class NMT:
    def __init__(self):
        self.model_1 = None
        self.model_2 = None
        self.opt_1 = None
        self.opt_2 = None
        self.best_loss = None

        self.input_lang = None
        self.output_lang = None
        self.vocab_lang = None

        self.print_every = hparams['steps_to_stats']
        self.epochs = hparams['epochs']
        self.hidden_size = hparams['units']

        self.train_fr = None
        self.train_to = None

        self.do_train = False
        self.do_infer = False
        self.do_review = False
        self.do_train_long = False
        self.do_interactive = False
        self.do_convert = False

        self.printable = ''

        parser = argparse.ArgumentParser(description='Train some NMT values.')
        parser.add_argument('--mode', help='mode of operation. (train, infer, review, long, interactive)')
        parser.add_argument('--printable', help='a string to print during training for identification.')
        parser.add_argument('--basename', help='base filename to use if it is different from settings file.')
        parser.add_argument('--autoencode', help='(broken) enable auto encode from the command line.', action='store_true')
        parser.add_argument('--train-all', help='(broken) enable training of the embeddings layer from the command line',
                            action='store_true')
        parser.add_argument('--convert-weights',help='convert weights', action='store_true')

        self.args = parser.parse_args()
        self.args = vars(self.args)
        # print(self.args)

        if self.args['printable'] is not None:
            self.printable = str(self.args['printable'])
        if self.args['mode'] == 'train': self.do_train = True
        if self.args['mode'] == 'infer': self.do_infer = True
        if self.args['mode'] == 'review': self.do_review = True
        if self.args['mode'] == 'long': self.do_train_long = True
        if self.args['mode'] == 'interactive': self.do_interactive = True
        if self.args['basename'] is not None:
            hparams['base_filename'] = self.args['basename']
            print(hparams['base_filename'], 'set name')
        if self.args['autoencode'] == True: hparams['autoencode'] = True
        if self.args['train_all'] == True:
            # hparams['embed_train'] = True
            self.trainable = True
        else:
            self.trainable = False
        if self.args['convert_weights'] == True: self.do_convert = True

    def task_normal_train(self):
        self.train_fr = hparams['data_dir'] + hparams['train_name'] + '.' + hparams['src_ending']
        self.train_to = hparams['data_dir'] + hparams['train_name'] + '.' + hparams['tgt_ending']
        pass

    def task_review_weights(self, pairs, stop_at_fail=False):
        num = 0 # hparams['base_file_num']
        for i in range(100):
            local_filename = hparams['save_dir'] + hparams['base_filename'] + '.'+ str(num) + '.pth.tar'
            if os.path.isfile(local_filename):
                ''' load weights '''

                print('==============================')
                print('here:',local_filename)
                self.load_checkpoint(local_filename)
                choice = random.choice(pairs)
                print(choice[0])
                out, _ =self.evaluate(None,None,choice[0])
                print(out)
            else:
                if stop_at_fail: break
            num = 10 * self.print_every * i
        pass

    def task_train_epochs(self,num=0):
        lr = hparams['learning_rate']
        if num == 0:
            num = hparams['epochs']
        for i in range(num):
            self.printable = ' epoch #' + str(i+1)
            self.trainIters(None, None, 75000, print_every=self.print_every, learning_rate=lr)
            self.save_checkpoint()
        pass

    def task_interactive(self):

        print('-------------------')
        while True:
            line = input("> ")
            line = tokenize_weak.format(line)
            print(line)
            out , _ =self.evaluate(None, None, line)
            print(out)

    def task_convert(self):
        hparams['base_filename'] += '.small'
        self.save_checkpoint(is_best=False,converted=True)

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

    def open_sentences(self, filename):
        t_yyy = []
        with open(filename, 'r') as r:
            for xx in r:
                t_yyy.append(xx)
        return t_yyy

    def asMinutes(self,s):
        m = math.floor(s / 60)
        s -= m * 60
        return '%dm %ds' % (m, s)

    def timeSince(self,since, percent):
        now = time.time()
        s = now - since
        es = s / (percent)
        rs = es - s
        return '%s (- %s)' % (self.asMinutes(s), self.asMinutes(rs))

    def unicodeToAscii(self,s):
        return ''.join(
            c for c in unicodedata.normalize('NFD', s)
            if unicodedata.category(c) != 'Mn'
        )

    def normalizeString(self,s):
        s = self.unicodeToAscii(s.lower().strip())
        s = re.sub(r"([.!?])", r" \1", s)
        s = re.sub(r"[^a-zA-Z.!?]+", r" ", s)
        return s

    def readLangs(self,lang1, lang2, reverse=False, load_vocab_file=None):
        print("Reading lines...")
        pairs = []
        if not self.do_interactive:

            l_in = self.open_sentences(hparams['data_dir'] + lang1)
            l_out = self.open_sentences(hparams['data_dir'] + lang2)

            #pairs = []
            for i in range(len(l_in)):
                #print(i)
                if i < len(l_out):
                    line = [ l_in[i].strip('\n'), l_out[i].strip('\n') ]
                    pairs.append(line)

        # Reverse pairs, make Lang instances
        if load_vocab_file is not None:
            self.vocab_lang = Lang(load_vocab_file)
            pass

        if reverse:
            pairs = [list(reversed(p)) for p in pairs]
            self.input_lang = Lang(lang2)
            self.output_lang = Lang(lang1)
        else:
            self.input_lang = Lang(lang1)
            self.output_lang = Lang(lang2)

        if hparams['autoencode'] == True:
            pairs = [ [p[0], p[0]] for p in pairs]
            self.output_lang = self.input_lang

        return self.input_lang, self.output_lang, pairs

    def filterPair(self,p):
        ends = False
        for j in eng_prefixes:
            if p[1].startswith(j): ends = True

        return len(p[0].split(' ')) < MAX_LENGTH and \
            len(p[1].split(' ')) < MAX_LENGTH and ends or True #\
            #p[1].startswith(eng_prefixes) or True

    def filterPairs(self,pairs):
        return [pair for pair in pairs if self.filterPair(pair)]

    def prepareData(self,lang1, lang2, reverse=False, omit_unk=False):
        if hparams['vocab_name'] is not None:
            v_name = hparams['data_dir'] + hparams['vocab_name']
        else:
            v_name = None
        self.input_lang, self.output_lang, pairs = self.readLangs(lang1, lang2,
                                                                  reverse,
                                                                  load_vocab_file=v_name)
        print("Read %s sentence pairs" % len(pairs))
        pairs = self.filterPairs(pairs)
        print("Trimmed to %s sentence pairs" % len(pairs))
        print("Counting words...")
        if v_name is not None:
            v = self.open_sentences(self.vocab_lang.name)
            for word in v:
                self.vocab_lang.addSentence(word.strip())
                #print(word)
            self.input_lang = self.vocab_lang
            self.output_lang = self.vocab_lang
            new_pairs = []
            for p in range(len(pairs)):
                a = []
                b = []
                for word in pairs[p][0].split(' '):
                    if word in self.vocab_lang.word2index:
                        a.append(word)
                    elif not omit_unk:
                        a.append(hparams['unk'])
                for word in pairs[p][1].split(' '):
                    if word in self.vocab_lang.word2index:
                        b.append(word)
                    elif not omit_unk:
                        b.append(hparams['unk'])
                new_pairs.append([' '.join(a), ' '.join(b)])
            pairs = new_pairs

        else:
            for pair in pairs:
                self.input_lang.addSentence(pair[0])
                self.output_lang.addSentence(pair[1])

        print("Counted words:")
        print(self.input_lang.name, self.input_lang.n_words)
        print(self.output_lang.name, self.output_lang.n_words)
        return self.input_lang, self.output_lang, pairs

    def indexesFromSentence(self,lang, sentence):
        s = sentence.split(' ')
        sent = []
        for word in s:
            if word == hparams['eol']: word = EOS_token
            elif word == hparams['sol']: word = SOS_token
            else: word = lang.word2index[word]
            sent.append(word)
        if len(sent) >= MAX_LENGTH:
            sent = sent[:MAX_LENGTH]
            sent[-1] = EOS_token
            #print(sent,'<<<<')
        return sent

        #return [lang.word2index[word] for word in sentence.split(' ')]

    def variableFromSentence(self,lang, sentence, add_eol=False):
        indexes = self.indexesFromSentence(lang, sentence)
        if add_eol: indexes.append(EOS_token)
        result = Variable(torch.LongTensor(indexes).view(-1, 1))
        if use_cuda:
            return result.cuda()
        else:
            return result

    def variablesFromPair(self,pair):
        input_variable = self.variableFromSentence(self.input_lang, pair[0])
        target_variable = self.variableFromSentence(self.output_lang, pair[1])
        return (input_variable, target_variable)

    def make_state(self, converted=False):
        if not converted:
            z = [
                {
                    'epoch':0,
                    'arch': None,
                    'state_dict': self.model_1.state_dict(),
                    'best_prec1': None,
                    'optimizer': self.opt_1.state_dict(),
                    'best_loss': self.best_loss
                },
                {
                    'epoch':0,
                    'arch':None,
                    'state_dict':self.model_2.state_dict(),
                    'best_prec1':None,
                    'optimizer': self.opt_2.state_dict(),
                    'best_loss': self.best_loss
                }
            ]
        else:
            z = [
                {
                    'epoch': 0,
                    'arch': None,
                    'state_dict': self.model_1.state_dict(),
                    'best_prec1': None,
                    'optimizer': None , # self.opt_1.state_dict(),
                    'best_loss': self.best_loss
                },
                {
                    'epoch': 0,
                    'arch': None,
                    'state_dict': self.model_2.state_dict(),
                    'best_prec1': None,
                    'optimizer': None, # self.opt_2.state_dict(),
                    'best_loss': self.best_loss
                }
            ]
        #print(z)
        return z
        pass

    def save_checkpoint(self, state=None, is_best=True, num=0, converted=False, extra=''):
        if state is None:
            state = self.make_state(converted=converted)
            if converted: print(converted, 'is converted.')
        basename = hparams['save_dir'] + hparams['base_filename']
        torch.save(state, basename + extra + '.' + str(num)+ '.pth.tar')
        if is_best:
            os.system('cp '+ basename + '.' + str(num) + '.pth.tar' + ' '  +
                      basename + '.best.pth.tar')

    def load_checkpoint(self, filename=None):
        if True:
            basename = hparams['save_dir'] + hparams['base_filename'] + '.best.pth.tar'
            if filename is not None: basename = filename
            if os.path.isfile(basename):
                print("loading checkpoint '{}'".format(basename))
                checkpoint = torch.load(basename)
                #print(checkpoint)
                try:
                    self.best_loss = checkpoint[0]['best_loss']
                except:
                    print('no best loss saved with checkpoint')
                    pass
                self.model_1.load_state_dict(checkpoint[0]['state_dict'])
                if self.opt_1 is not None:
                    self.opt_1.load_state_dict(checkpoint[0]['optimizer'])

                self.model_2.load_state_dict(checkpoint[1]['state_dict'])
                if self.opt_2 is not None:
                    self.opt_2.load_state_dict(checkpoint[1]['optimizer'])

                print("loaded checkpoint '"+ basename + "' ")
            else:
                print("no checkpoint found at '"+ basename + "'")

    def _mod_hidden(self, encoder_hidden):
        z = torch.cat((encoder_hidden,), 2)[0].view(1, 1, self.hidden_size )
        #print(z.size())
        return z

    def _shorten(self,sentence):
        #assume input is list already
        saved = [hparams['eol']]
        punct = ['.','!','?']
        out = []
        for i in sentence:
            if i in saved: break
            if i != hparams['sol']:
                out.append(i)
            if i in punct: break
            saved.append(i)
        return ' '.join(out)

    def train(self,input_variable, target_variable, encoder, decoder, encoder_optimizer, decoder_optimizer, criterion, max_length=MAX_LENGTH):
        encoder_hidden = Variable(torch.zeros(2, 1, self.hidden_size)) #encoder.initHidden()

        encoder_optimizer.zero_grad()
        decoder_optimizer.zero_grad()

        input_length = input_variable.size()[0]
        target_length = target_variable.size()[0]

        if input_length >= hparams['tokens_per_sentence'] : input_length = hparams['tokens_per_sentence']
        if target_length >= hparams['tokens_per_sentence'] : target_length = hparams['tokens_per_sentence']

        encoder_outputs = Variable(torch.zeros(max_length, encoder.hidden_size  ))
        encoder_outputs = encoder_outputs.cuda() if use_cuda else encoder_outputs

        loss = 0

        for ei in range(input_length):
            encoder_output, encoder_hidden = encoder(
                input_variable[ei], encoder_hidden)
            encoder_outputs[ei] = encoder_output[0][0]

        decoder_input = Variable(torch.LongTensor([[SOS_token]]))
        decoder_input = decoder_input.cuda() if use_cuda else decoder_input

        decoder_hidden = encoder_hidden
        decoder_hidden =self._mod_hidden(encoder_hidden) # torch.cat((encoder_hidden, encoder_hidden),2)[0].view(1,1,512)

        use_teacher_forcing = True if random.random() < teacher_forcing_ratio else False

        if use_teacher_forcing:
            # Teacher forcing: Feed the target as the next input
            for di in range(target_length):
                #print(di,'di', decoder_hidden.size(),'<', encoder_outputs.size())
                decoder_output, decoder_hidden, decoder_attention = decoder(
                    decoder_input, decoder_hidden, encoder_outputs)
                loss += criterion(decoder_output, target_variable[di])
                decoder_input = target_variable[di]  # Teacher forcing

        else:
            # Without teacher forcing: use its own predictions as the next input
            for di in range(target_length):
                #print(di,'di', decoder_hidden.size(),'<', encoder_outputs.size() )

                decoder_output, decoder_hidden, decoder_attention = decoder( #encoder_outputs, decoder_input, decoder_hidden)
                    decoder_input, decoder_hidden, encoder_outputs)
                topv, topi = decoder_output.data.topk(1)
                ni = topi[0][0]

                decoder_input = Variable(torch.LongTensor([[ni]]))
                decoder_input = decoder_input.cuda() if use_cuda else decoder_input

                loss += criterion(decoder_output, target_variable[di])
                if ni == EOS_token:
                    break

        loss.backward()

        encoder_optimizer.step()
        decoder_optimizer.step()

        return loss.data[0] / target_length

    def trainIters(self, encoder, decoder, n_iters, print_every=1000, plot_every=100, learning_rate=0.01):
        if (encoder is not None and decoder is not None and
                self.model_1 is None and self.model_2 is None):
            self.model_1 = encoder
            self.model_2 = decoder
        else:
            encoder = self.model_1
            decoder = self.model_2

        start = time.time()
        plot_losses = []
        save_num = 0
        print_loss_total = 0  # Reset every print_every
        plot_loss_total = 0  # Reset every plot_every

        encoder_optimizer = optim.SGD(encoder.parameters(), lr=learning_rate)
        decoder_optimizer = optim.SGD(decoder.parameters(), lr=learning_rate)
        training_pairs = [self.variablesFromPair(random.choice(pairs))
                          for i in range(n_iters)]
        criterion = nn.NLLLoss()

        if self.opt_1 is None and self.opt_2 is None:
            self.opt_1 = encoder_optimizer
            self.opt_2 = decoder_optimizer

        self.load_checkpoint()

        for iter in range(1, n_iters + 1):
            training_pair = training_pairs[iter - 1]
            input_variable = training_pair[0]
            target_variable = training_pair[1]

            loss = self.train(input_variable, target_variable, encoder,
                         decoder, encoder_optimizer, decoder_optimizer, criterion)
            print_loss_total += loss
            plot_loss_total += loss

            if iter % print_every == 0:
                print_loss_avg = print_loss_total / print_every
                print_loss_total = 0
                print('iter = '+str(iter)+ ', num of iters = '+str(n_iters) + ' ' + self.printable)
                if iter % (print_every * 10) == 0:
                    save_num +=1
                    if (self.best_loss is None or print_loss_avg <= self.best_loss or save_num > 5):
                        save_num = 0
                        extra = ''
                        if hparams['autoencode'] == True: extra = '.autoencode'
                        self.save_checkpoint(num=iter,extra=extra)
                        self.best_loss = print_loss_avg
                        print('======= save file '+ extra+' ========')
                print('%s (%d %d%%) %.4f' % (self.timeSince(start, iter / n_iters),
                                             iter, iter / n_iters * 100, print_loss_avg))
                choice = random.choice(pairs)
                print('src:',choice[0])
                print('ref:',choice[1])
                words, _ = self.evaluate(None, None, choice[0])
                #print(choice)
                print('ans:',words)
                print('try:',self._shorten(words))
                print("-----")

            if iter % plot_every == 0 and False:
                plot_loss_avg = plot_loss_total / plot_every
                plot_losses.append(plot_loss_avg)
                plot_loss_total = 0

    def evaluate(self, encoder, decoder, sentence, max_length=MAX_LENGTH):
        if (encoder is not None and decoder is not None and
                self.model_1 is None and self.model_2 is None):
            self.model_1 = encoder
            self.model_2 = decoder
        else:
            encoder = self.model_1
            decoder = self.model_2

        input_variable = self.variableFromSentence(self.input_lang, sentence)
        input_length = input_variable.size()[0]
        encoder_hidden = encoder.initHidden()

        if input_length >= max_length : input_length = max_length

        encoder_outputs = Variable(torch.zeros(max_length, encoder.hidden_size  ))
        encoder_outputs = encoder_outputs.cuda() if use_cuda else encoder_outputs

        for ei in range(input_length):
            encoder_output, encoder_hidden = encoder(input_variable[ei],
                                                     encoder_hidden)

            encoder_outputs[ei] = encoder_outputs[ei] + encoder_output[0][0]

        decoder_input = Variable(torch.LongTensor([[SOS_token]]))  # SOS
        decoder_input = decoder_input.cuda() if use_cuda else decoder_input

        decoder_hidden = encoder_hidden
        decoder_hidden = self._mod_hidden(encoder_hidden) # torch.cat((encoder_hidden, encoder_hidden), 2)[0].view(1, 1, 512)

        decoded_words = []
        decoder_attentions = torch.zeros(max_length, max_length)

        for di in range(max_length):
            decoder_output, decoder_hidden, decoder_attention = decoder( #encoder_outputs, decoder_input, decoder_hidden)
                decoder_input, decoder_hidden, encoder_outputs)
            #decoder_input, decoder_hidden, encoder_outputs)

            decoder_attentions[di] = decoder_attention.data
            topv, topi = decoder_output.data.topk(1)
            ni = topi[0][0]
            if ni == EOS_token:
                xxx = hparams['eol']
                decoded_words.append(xxx)
                print('eol found.')
                break
            else:
                decoded_words.append(self.output_lang.index2word[ni])

            decoder_input = Variable(torch.LongTensor([[ni]]))
            decoder_input = decoder_input.cuda() if use_cuda else decoder_input

        return decoded_words, decoder_attentions[:di + 1]

if __name__ == '__main__':

    n = NMT()

    n.task_normal_train()

    n.input_lang, n.output_lang, pairs = n.prepareData(n.train_fr, n.train_to, reverse=False, omit_unk=True)
    #print(random.choice(pairs))

    #n.model_1 = EncoderRNN(n.input_lang.n_words, n.hidden_size)

    n.model_1 = EncoderBiRNN(n.input_lang.n_words, n.hidden_size )
    n.model_2 = AttnDecoderRNN(n.hidden_size , n.output_lang.n_words, dropout_p=0.1)
    #n.model_2 = Decoder(n.output_lang.n_words, 100 ,n.hidden_size, 1 ,dropout=0.1)

    if use_cuda:
        n.model_1 = n.model_1.cuda()
        n.model_2 = n.model_2.cuda()

    if n.do_train:
        lr = hparams['learning_rate']
        n.trainIters(None, None, 75000, print_every=n.print_every, learning_rate=lr)
    if n.do_train_long:
        n.task_train_epochs()

    if n.do_interactive:
        n.load_checkpoint()
        n.task_interactive()

    if n.do_review:
        n.task_review_weights(pairs,stop_at_fail=False)

    if n.do_convert:
        n.load_checkpoint()
        n.task_convert()

    if n.do_infer:
        n.load_checkpoint()
        choice = random.choice(pairs)[0]
        print(choice)
        words, _ = n.evaluate(None,None,choice)
        print(words)

    if False:
        print(n.model_1.state_dict())


from django.contrib import messages
from django.contrib.auth.decorators import login_required
from django.shortcuts import render, redirect, get_object_or_404
from django.utils.translation import gettext as _

from snippets.forms import SnippetForm, CommentForm
from snippets.models import Snippet, Comment

def top(request):
    snippets = Snippet.objects.all()
    context = {"snippets": snippets}
    return render(request, "snippets/top.html", context)

@login_required
def snippet_new(request):
    if request.method == 'POST':
        form = SnippetForm(request.POST)
        if form.is_valid():
            snippet = form.save(commit=False)
            snippet.created_by = request.user
            snippet.save()
            messages.add_message(request, messages.SUCCESS,
                                 _("Success to create a new snippet."))
            return redirect(snippet_detail, snippet_id=snippet.pk)
        else:
            messages.add_message(request, messages.ERROR,
                                 _("Failed to create a new snippet."))
    else:
        form = SnippetForm()
    return render(request, "snippets/snippet_new.html", {'form': form})

@login_required
def snippet_edit(request, snippet_id):
    snippet = get_object_or_404(Snippet, pk=snippet_id)
    if snippet.created_by_id != request.user.id:
        return HttpResponseForbidden(_("You cannot to edit this snippet."))

    if request.method == "POST":
        form = SnippetForm(request.POST, instance=snippet)
        if form.is_valid():
            form.save()
            messages.add_message(request, messages.SUCCESS,
                                 _("Success to update the snippet."))
            return redirect('snippet_detail', snippet_id=snippet_id)
        else:
            messages.add_message(request, messages.ERROR,
                                 _("Failed to update the snippet."))
    else:
        form = SnippetForm(instance=snippet)
    return render(request, 'snippets/snippet_edit.html', {'form': form})

@login_required
def snippet_detail(request, snippet_id):
    snippet = get_object_or_404(Snippet, pk=snippet_id)
    comments = Comment.objects.filter(commented_to=snippet_id).all()
    comment_form = CommentForm()

    return render(request, "snippets/snippet_detail.html", {
        'snippet': snippet,
        'comments': comments,
        'comment_form': comment_form,
    })

@login_required
def comment_new(request, snippet_id):
    snippet = get_object_or_404(Snippet, pk=snippet_id)

    form = CommentForm(request.POST)
    if form.is_valid():
        comment = form.save(commit=False)
        comment.commented_to = snippet
        comment.commented_by = request.user
        comment.save()
        messages.add_message(request, messages.SUCCESS,
                             _("Success to post a comment."))
    else:
        messages.add_message(request, messages.ERROR,
                             _("Failed to post a comment."))
    return redirect('snippet_detail', snippet_id=snippet_id)


import os
import torch.nn
import importlib

def search_algorithm_getter(search_algorithm_type,
                            data,
                            search_parameter,
                            gnn_parameter,
                            search_space):

    search_algorithm_class = "Search"
    search_algorithm_module = importlib.import_module("automsr.search_algorithm." +
                                                      search_algorithm_type +
                                                      "." +
                                                      "search_algorithm")

    search_algorithm_obj = getattr(search_algorithm_module,
                                   search_algorithm_class)

    search_algorithm = search_algorithm_obj(data,
                                            search_parameter,
                                            gnn_parameter,
                                            search_space)
    return search_algorithm

def optimizer_getter(opt_type,
                     gnn_model,
                     optimizer_parameter_dict):

    optimizer_class = "Optimizer"
    optimizer_module = importlib.import_module("automsr.model.optimizer_function" +
                                               "." +
                                               opt_type)
    optimizer_obj = getattr(optimizer_module, optimizer_class)
    optimizer_function = optimizer_obj(gnn_model,
                                       optimizer_parameter_dict).function()

    return optimizer_function

def loss_getter(loss_type):

    loss_class = "Loss"
    loss_module = importlib.import_module("automsr.model.loss_function" +
                                          "." +
                                          loss_type)
    loss_obj = getattr(loss_module, loss_class)
    loss_function = loss_obj()

    return loss_function

def evaluator_getter(evaluator_type):

    evaluator_class = "Evaluator"
    evaluator_module = importlib.import_module("automsr.model.evaluator_function" +
                                               "." +
                                               evaluator_type)
    evaluator_obj = getattr(evaluator_module, evaluator_class)
    evaluator_function = evaluator_obj()

    return evaluator_function

def downstream_task_model_getter(downstream_task_type,
                                 gnn_embedding_dim,
                                 graph_data):

    downstream_task_model_class = "DownstreamTask"
    downstream_task_model_module = importlib.import_module("automsr.model.downstream_task_model" +
                                                           "." +
                                                           downstream_task_type)
    downstream_task_model_obj = getattr(downstream_task_model_module, downstream_task_model_class)

    downstream_task = downstream_task_model_obj(gnn_embedding_dim, graph_data)

    return downstream_task

def attention_getter(heads,
                     output_dim):

    search_space_path = os.path.split(os.path.realpath(__file__))[0] + "/search_space/attention"
    attention_list = [attention for attention in os.listdir(search_space_path) if attention not in "__pycache__"
                      and attention not in "README.md"]
    attention_dict = torch.nn.ModuleDict()
    for attention in attention_list:
        attention_class = "Attention"
        attention_module = importlib.import_module("automsr.search_space.attention" +
                                                    "." + attention[:-3])
        attention_obj = getattr(attention_module, attention_class)
        attention_function = attention_obj(heads, output_dim)
        attention_dict[attention[:-3]] = attention_function

    return attention_dict

def aggregation_getter():

    search_space_path = os.path.split(os.path.realpath(__file__))[0]+ "/search_space/aggregation"
    aggregation_list = [aggregation for aggregation in os.listdir(search_space_path) if aggregation not in "__pycache__"
                        and aggregation not in "README.md"]
    aggregation_dict = {}
    for aggregation in aggregation_list:
        aggregation_class = "Aggregation"
        aggregation_module = importlib.import_module("automsr.search_space.aggregation" +
                                                     "." + aggregation[:-3])
        aggregation_obj = getattr(aggregation_module, aggregation_class)
        aggregation_function = aggregation_obj()
        aggregation_dict[aggregation[:-3]] = aggregation_function

    return aggregation_dict

def activation_getter():

    search_space_path = os.path.split(os.path.realpath(__file__))[0] + "/search_space/activation"
    activation_list = [activation for activation in os.listdir(search_space_path) if activation not in "__pycache__"
                       and activation not in "README.md"]
    activation_dict = {}
    for activation in activation_list:
        activation_class = "Activation"
        activation_module = importlib.import_module("automsr.search_space.activation" +
                                                    "." + activation[:-3])
        activation_obj = getattr(activation_module, activation_class)
        activation_function = activation_obj().function()
        activation_dict[activation[:-3]] = activation_function

    return activation_dict

if __name__=="__main__":
    pass


from argparse import ArgumentParser
import yaml
import os
import pandas as pd
import json
from collections import defaultdict

def arguments():
    """
    Creates and returns the arg parser for Potato on the command line
    """
    parser = ArgumentParser()
    parser.set_defaults(show_path=False, show_similarity=False)

    parser.add_argument("multitask_config_file")

    return parser.parse_args()

def main():
    args = arguments()

    # load multitask annotation config
    with open(args.multitask_config_file, "rt") as f:
        multitask_config = yaml.safe_load(f)

    # create basic folders
    for folder in ["surveyflow", "annotation_output", "data_files", "configs", "htmls"]:
        cur_path = multitask_config["base_dir"] + folder
        if not os.path.exists(cur_path):
            os.makedirs(cur_path)
            print("Created directory: %s" % (cur_path))

    # build mapping dict
    key2text = defaultdict(dict)

    if "multitask_guideline_file" in multitask_config and os.path.exists(
        multitask_config["multitask_guideline_file"]
    ):
        # load multitask annotation guideline
        multitask_guideline_df = pd.read_csv(multitask_config["multitask_guideline_file"])
        for i, row in multitask_guideline_df.iterrows():
            if type(row["key"]) != str or row["key"][0] != "[" or row["key"][-1] != "]":
                continue
            for task in multitask_config["tasks"]:
                key2text[row["key"]][task] = (
                    row[task] if type(row[task]) == str else row[multitask_config["base_task"]]
                )

    for task in multitask_config["tasks"]:
        # load basic config file
        with open(multitask_config["base_config_file"], "rt") as f:
            page = f.read()
            for key in key2text:
                page = page.replace(key, key2text[key][task])

            surveyflow_output_path = multitask_config["surveyflow_output_path"].replace(
                "[TASK]", task
            )
            page = page.replace(
                multitask_config["surveyflow_path"], surveyflow_output_path + task + "-"
            )
            page = page.replace("[TASK]", task)

            config = yaml.safe_load(page)

        config["output_annotation_dir"] = multitask_config["output_annotation_dir"].replace(
            "[TASK]", task
        )
        if not os.path.exists(config["output_annotation_dir"]):
            os.makedirs(config["output_annotation_dir"])

        """
        # setup the site_dir path for each task
        config['path_under_site_dir'] = multitask_config['path_under_site_dir'].replace("[TASK]", task)
        config["site_dir"] += config['path_under_site_dir']
        if not os.path.exists(config["site_dir"]):
            os.makedirs(config["site_dir"])
        """

        config["annotation_task_name"] = multitask_config["annotation_task_name"].replace(
            "[TASK]", task
        )
        config["data_files"] = [
            it.replace("[TASK]", task) for it in multitask_config["data_files"]
        ]
        # config["prestudy"] = [it.replace("[TASK]", task) for it in multitask_config['prestudy']]

        with open(multitask_config["base_dir"] + "configs/%s.yaml" % task, "wt") as f:
            json.dump(config, f, indent=4)

        # create the dir for surveyflow_output_path
        if not os.path.exists(surveyflow_output_path):
            os.makedirs(surveyflow_output_path)

        # setup surveyflow pages
        surveyflow_files = os.listdir(multitask_config["surveyflow_path"])
        for file in surveyflow_files:
            if os.path.isdir(multitask_config["surveyflow_path"] + file):
                continue
            with open(multitask_config["surveyflow_path"] + file, "r") as f:
                page = f.read()
            for key in key2text:
                page = page.replace(key, key2text[key][task])
            page = page.replace("[TASK]", task)
            # print(surveyflow_output_path)
            with open(surveyflow_output_path + task + "-" + file, "wt") as f:
                f.write(page)

        # for key in ["surveyflow_output_path", ]

if __name__ == "__main__":
    main()


# ClearML - Hydra Example
#
import hydra

from omegaconf import OmegaConf

from clearml import Task

@hydra.main(config_path="config_files", config_name="config", version_base=None)
def my_app(cfg):
    # type (DictConfig) -> None
    task = Task.init(project_name="examples", task_name="Hydra configuration")
    logger = task.get_logger()
    logger.report_text("You can view your full hydra configuration under Configuration tab in the UI")
    print(OmegaConf.to_yaml(cfg))

if __name__ == "__main__":
    my_app()


########################
#file: zbw_animShift.py
#Author: zeth willie
#Contact: zeth@catbuks.com, www.williework.blogspot.com
#Date Modified: 4/27/13
#To Use: type in python window  "import zbw_animShift as ash; ash.animShift()"
#Notes/Descriptions: Used to move animation curves based on changes to a base object at a certain frame
########################

import maya.cmds as cmds

class AnimShift(object):

    def __init__(self):

        self.widgets = {}
        self.frame = 1

        #add in the vars for if there's a key on base
        self.txk=False
        self.tyk=False
        self.tzk=False
        self.rxk=False
        self.ryk=False
        self.rzk=False
        self.sxk=False
        self.syk=False
        self.szk=False

        self.baseObj = ""

        self.shiftUI()

    def shiftUI(self, *args):
        """ui for the class"""

        if cmds.window("shiftWin", exists=True):
            cmds.deleteUI("shiftWin")

        self.widgets["win"] = cmds.window("shiftWin", t="zbw_animShift", w=330, h=500)
        self.widgets["mainCLO"] = cmds.columnLayout()

        self.widgets["timeRBG"] = cmds.radioButtonGrp(l="Range Options:", nrb=3, l1="All", l2="Time Slider", l3="Range", sl=1, cw=([1,100], [2,50],[3,100],[4,50]), cal=(1,"left"), cc=self.enableRange)
        self.widgets["rangeFFG"] = cmds.floatFieldGrp(l="Start/End Frames", nf=2, v1=1, v2=10, en=False, cal=((1,"left"),(2,"left")), cw=((1,100), (2,50),(3,50)))

        cmds.separator(h=5, style="single")

        self.widgets["getBaseBut"]  = cmds.button(l="1. Set Base Object and Frame To Drive Shift", w=330, h=40, bgc=(.6,.8,.6), c=self.getBase)
        self.widgets["baseTFG"] = cmds.textFieldGrp(l="Base Object", cal=((1,"left"), (2,"left")), cw=((1,70),(2,250)), ed=False)
        self.widgets["baseFrameFFG"] = cmds.floatFieldGrp(l="Base Frame", cal = ((1, "left"), (2,"left")), cw=((1,100), (2,50)), pre=2, en=False)

        cmds.separator(h=10)

        self.widgets["valueRCL"] = cmds.rowColumnLayout(nc=3, cw=([1,110],[2,110],[3,110]))
        cmds.text("Orig Values")
        cmds.text("Modified Values")
        cmds.text("Difference Value")

        cmds.separator(style="single", h=5)
        cmds.separator(style="single", h=5)
        cmds.separator(style="single", h=5)

        self.widgets["origTxFFG"] = cmds.floatFieldGrp(l="tx", pre=5, cal=([1,"left"], [2, "left"]), en=False, cw=([1, 30], [2,75]))
        self.widgets["modTxFFG"] = cmds.floatFieldGrp(l="tx", pre=5, cal=([1,"left"], [2, "left"]), en=False, cw=([1, 30], [2,75]))
        self.widgets["difTxFFG"] = cmds.floatFieldGrp(l="tx", pre=5, cal=([1,"left"], [2, "left"]), cw=([1, 30], [2,75]))

        self.widgets["origTyFFG"] = cmds.floatFieldGrp(l="ty", pre=5, cal=([1,"left"], [2, "left"]), en=False, cw=([1, 30], [2,75]))
        self.widgets["modTyFFG"] = cmds.floatFieldGrp(l="ty", pre=5, cal=([1,"left"], [2, "left"]),  en=False,cw=([1, 30], [2,75]))
        self.widgets["difTyFFG"] = cmds.floatFieldGrp(l="ty", pre=5, cal=([1,"left"], [2, "left"]), cw=([1, 30], [2,75]))

        self.widgets["origTzFFG"] = cmds.floatFieldGrp(l="tz", pre=5, cal=([1,"left"], [2, "left"]), en=False, cw=([1, 30], [2,75]))
        self.widgets["modTzFFG"] = cmds.floatFieldGrp(l="tz", pre=5, cal=([1,"left"], [2, "left"]), en=False, cw=([1, 30], [2,75]))
        self.widgets["difTzFFG"] = cmds.floatFieldGrp(l="tz", pre=5, cal=([1,"left"], [2, "left"]), cw=([1, 30], [2,75]))

        cmds.separator(style="single", h=5)
        cmds.separator(style="single", h=5)
        cmds.separator(style="single", h=5)

        self.widgets["origRxFFG"] = cmds.floatFieldGrp(l="rx", pre=5, cal=([1,"left"], [2, "left"]), en=False, cw=([1, 30], [2,75]))
        self.widgets["modRxFFG"] = cmds.floatFieldGrp(l="rx", pre=5, cal=([1,"left"], [2, "left"]), en=False, cw=([1, 30], [2,75]))
        self.widgets["difRxFFG"] = cmds.floatFieldGrp(l="rx", pre=5, cal=([1,"left"], [2, "left"]), cw=([1, 30], [2,75]))

        self.widgets["origRyFFG"] = cmds.floatFieldGrp(l="ry", pre=5, cal=([1,"left"], [2, "left"]), en=False, cw=([1, 30], [2,75]))
        self.widgets["modRyFFG"] = cmds.floatFieldGrp(l="ry", pre=5, cal=([1,"left"], [2, "left"]), en=False, cw=([1, 30], [2,75]))
        self.widgets["difRyFFG"] = cmds.floatFieldGrp(l="ry", pre=5, cal=([1,"left"], [2, "left"]), cw=([1, 30], [2,75]))

        self.widgets["origRzFFG"] = cmds.floatFieldGrp(l="rz", pre=5, cal=([1,"left"], [2, "left"]), en=False, cw=([1, 30], [2,75]))
        self.widgets["modRzFFG"] = cmds.floatFieldGrp(l="rz", pre=5, cal=([1,"left"], [2, "left"]), en=False, cw=([1, 30], [2,75]))
        self.widgets["difRzFFG"] = cmds.floatFieldGrp(l="rz", pre=5, cal=([1,"left"], [2, "left"]), cw=([1, 30], [2,75]))

        cmds.separator(style="single", h=5)
        cmds.separator(style="single", h=5)
        cmds.separator(style="single", h=5)

        self.widgets["origSxFFG"] = cmds.floatFieldGrp(l="sx", pre=5, cal=([1,"left"], [2, "left"]), en=False, cw=([1, 30], [2,75]))
        self.widgets["modSxFFG"] = cmds.floatFieldGrp(l="sx", pre=5, cal=([1,"left"], [2, "left"]), en=False, cw=([1, 30], [2,75]))
        self.widgets["difSxFFG"] = cmds.floatFieldGrp(l="sx", pre=5, cal=([1,"left"], [2, "left"]), cw=([1, 30], [2,75]))

        self.widgets["origSyFFG"] = cmds.floatFieldGrp(l="sy", pre=5, cal=([1,"left"], [2, "left"]), en=False, cw=([1, 30], [2,75]))
        self.widgets["modSyFFG"] = cmds.floatFieldGrp(l="sy", pre=5, cal=([1,"left"], [2, "left"]), en=False, cw=([1, 30], [2,75]))
        self.widgets["difSyFFG"] = cmds.floatFieldGrp(l="sy", pre=5, cal=([1,"left"], [2, "left"]), cw=([1, 30], [2,75]))

        self.widgets["origSzFFG"] = cmds.floatFieldGrp(l="sz", pre=5, cal=([1,"left"], [2, "left"]), en=False, cw=([1, 30], [2,75]))
        self.widgets["modSzFFG"] = cmds.floatFieldGrp(l="sz", pre=5, cal=([1,"left"], [2, "left"]), en=False,cw=([1, 30], [2,75]))
        self.widgets["difSzFFG"] = cmds.floatFieldGrp(l="sz", pre=5, cal=([1,"left"], [2, "left"]), cw=([1, 30], [2,75]))

        cmds.separator(style="single", h=5)
        cmds.separator(style="single", h=5)
        cmds.separator(style="single", h=5)

        cmds.setParent(self.widgets["mainCLO"])

        #button to capture and calculate changes to the anim curves
        self.widgets["captureBut"] = cmds.button(l="2. Calculate and Capture Changed Values", w=330, h=30, bgc=(.8,.8,.6), c=self.captureChanges)

        cmds.separator(style="single", h=5)

        self.widgets["moveBut"] = cmds.button(l="3. Shift The Anim Curves For Selected Objects!", w=330, h=30, bgc=(.6,.6,.8), c=self.shiftAnim)

        cmds.separator(style="single", h=5)

        self.widgets["doneText"] = cmds.text(w=330, l="No Status/All Cleared. Select a base object!", vis=True, bgc=(.5,.5,.5), al="center")
        cmds.separator(style="single", h=10)

        #2 column layout to spread these two buttons
        self.widgets["bottomButtonRCL"] = cmds.rowColumnLayout(nc=2, w=330, cw=((1,250), (2,80)))
        self.widgets["clearButton"] = cmds.button(l="Clear All Values!", w=330, h=30, bgc=(.9,.5,.5), c=self.clearAll)
        # self.widgets["restoreButton"] = cmds.button(l="Restore Base", w=80, h=30, bgc=(.8,.5,.5), c=self.restoreBase)
        cmds.showWindow(self.widgets["win"])
        cmds.window(self.widgets["win"], e=True, w=330, h=500)

    #####functions#####

    def getBase(self, *args):

        sel = cmds.ls(sl=True, type="transform", l=True)

        if len(sel)>1:
            cmds.warning("You've selected more than one object. Only one object can be the base of the animation shift")
        else:
            obj = sel[0]
            #put sel in the base obj tfg and get the frame
            cmds.textFieldGrp(self.widgets["baseTFG"], e=True, tx=obj)
            self.frame = cmds.currentTime(q=True)
            cmds.floatFieldGrp(self.widgets["baseFrameFFG"], e=True, v1=self.frame)

            btx = cmds.getAttr("%s.tx"%obj)
            bty = cmds.getAttr("%s.ty"%obj)
            btz = cmds.getAttr("%s.tz"%obj)

            brx = cmds.getAttr("%s.rx"%obj)
            bry = cmds.getAttr("%s.ry"%obj)
            brz = cmds.getAttr("%s.rz"%obj)

            bsx = cmds.getAttr("%s.sx"%obj)
            bsy = cmds.getAttr("%s.sy"%obj)
            bsz = cmds.getAttr("%s.sz"%obj)

            cmds.floatFieldGrp(self.widgets["origTxFFG"], e=True, v1=btx)
            cmds.floatFieldGrp(self.widgets["origTyFFG"], e=True, v1=bty)
            cmds.floatFieldGrp(self.widgets["origTzFFG"], e=True, v1=btz)

            cmds.floatFieldGrp(self.widgets["origRxFFG"], e=True, v1=brx)
            cmds.floatFieldGrp(self.widgets["origRyFFG"], e=True, v1=bry)
            cmds.floatFieldGrp(self.widgets["origRzFFG"], e=True, v1=brz)

            cmds.floatFieldGrp(self.widgets["origSxFFG"], e=True, v1=bsx)
            cmds.floatFieldGrp(self.widgets["origSyFFG"], e=True, v1=bsy)
            cmds.floatFieldGrp(self.widgets["origSzFFG"], e=True, v1=bsz)

            if cmds.keyframe("%s.tx"%(obj), q=True, t=(self.frame, self.frame)):
                self.txk = True
            if cmds.keyframe("%s.ty"%(obj), q=True, t=(self.frame, self.frame)):
                self.tyk = True
            if cmds.keyframe("%s.tz"%obj, q=True, t=(self.frame, self.frame)):
                self.tzk = True

            if cmds.keyframe("%s.rx"%(obj), q=True, t=(self.frame, self.frame)):
                self.rxk = True
            if cmds.keyframe("%s.ry"%(obj), q=True, t=(self.frame, self.frame)):
                self.ryk = True
            if cmds.keyframe("%s.rz"%obj, q=True, t=(self.frame, self.frame)):
                self.rzk = True

            if cmds.keyframe("%s.sx"%(obj), q=True, t=(self.frame, self.frame)):
                self.sxk = True
            if cmds.keyframe("%s.sy"%(obj), q=True, t=(self.frame, self.frame)):
                self.syk = True
            if cmds.keyframe("%s.sz"%obj, q=True, t=(self.frame, self.frame)):
                self.szk = True

            #catch the base object  and frame for later use
            self.baseObj = obj

            cmds.text(self.widgets["doneText"], e=True, l='"Base Obj Set! Base Frame is: %s"'%self.frame, vis=True, bgc=(.5,.6,.5))

    def captureChanges(self, *args):
        """method to capture value changes"""
        obj = cmds.textFieldGrp(self.widgets["baseTFG"], q=True, tx=True)
        frame = cmds.floatFieldGrp(self.widgets["baseFrameFFG"], q=True, v1=True)
        cmds.currentTime(frame, edit=True)

        if obj:
            #poplulate the second column with the new values from the base object
            mtx = cmds.getAttr("%s.tx"%obj)
            mty = cmds.getAttr("%s.ty"%obj)
            mtz = cmds.getAttr("%s.tz"%obj)

            mrx = cmds.getAttr("%s.rx"%obj)
            mry = cmds.getAttr("%s.ry"%obj)
            mrz = cmds.getAttr("%s.rz"%obj)

            msx = cmds.getAttr("%s.sx"%obj)
            msy = cmds.getAttr("%s.sy"%obj)
            msz = cmds.getAttr("%s.sz"%obj)

            cmds.floatFieldGrp(self.widgets["modTxFFG"], e=True, v1=mtx)
            cmds.floatFieldGrp(self.widgets["modTyFFG"], e=True, v1=mty)
            cmds.floatFieldGrp(self.widgets["modTzFFG"], e=True, v1=mtz)

            cmds.floatFieldGrp(self.widgets["modRxFFG"], e=True, v1=mrx)
            cmds.floatFieldGrp(self.widgets["modRyFFG"], e=True, v1=mry)
            cmds.floatFieldGrp(self.widgets["modRzFFG"], e=True, v1=mrz)

            cmds.floatFieldGrp(self.widgets["modSxFFG"], e=True, v1=msx)
            cmds.floatFieldGrp(self.widgets["modSyFFG"], e=True, v1=msy)
            cmds.floatFieldGrp(self.widgets["modSzFFG"], e=True, v1=msz)

            #get orig values
            otx = cmds.floatFieldGrp(self.widgets["origTxFFG"], q=True, v1=True)
            oty = cmds.floatFieldGrp(self.widgets["origTyFFG"], q=True, v1=True)
            otz = cmds.floatFieldGrp(self.widgets["origTzFFG"], q=True, v1=True)

            orx = cmds.floatFieldGrp(self.widgets["origRxFFG"], q=True, v1=True)
            ory = cmds.floatFieldGrp(self.widgets["origRyFFG"], q=True, v1=True)
            orz = cmds.floatFieldGrp(self.widgets["origRzFFG"], q=True, v1=True)

            osx = cmds.floatFieldGrp(self.widgets["origSxFFG"], q=True, v1=True)
            osy = cmds.floatFieldGrp(self.widgets["origSyFFG"], q=True, v1=True)
            osz = cmds.floatFieldGrp(self.widgets["origSzFFG"], q=True, v1=True)

            #now subtract them and put the value in the diff column
            dtx = mtx - otx
            dty = mty - oty
            dtz = mtz - otz

            drx = mrx - orx
            dry = mry - ory
            drz = mrz - orz

            dsx = msx - osx
            dsy = msy - osy
            dsz = msz - osz

            cmds.floatFieldGrp(self.widgets["difTxFFG"], e=True, v1=dtx)
            cmds.floatFieldGrp(self.widgets["difTyFFG"], e=True, v1=dty)
            cmds.floatFieldGrp(self.widgets["difTzFFG"], e=True, v1=dtz)

            cmds.floatFieldGrp(self.widgets["difRxFFG"], e=True, v1=drx)
            cmds.floatFieldGrp(self.widgets["difRyFFG"], e=True, v1=dry)
            cmds.floatFieldGrp(self.widgets["difRzFFG"], e=True, v1=drz)

            cmds.floatFieldGrp(self.widgets["difSxFFG"], e=True, v1=dsx)
            cmds.floatFieldGrp(self.widgets["difSyFFG"], e=True, v1=dsy)
            cmds.floatFieldGrp(self.widgets["difSzFFG"], e=True, v1=dsz)

            cmds.text(self.widgets["doneText"], e=True, l='"Shift Values Captured! Now Select Objs to Shift!"', vis=True, bgc=(.6,.6,.5))
        else:
            cmds.warning("You don't have anything in the base object field!")

    def shiftAnim(self, *args):
        """method that does the shifting"""
#TO-DO----------------possible to-do? Deal with locked or already driven channels?
#TO-DO----------------eval the curve at that point, then set the value in world space, then eval the curve again, then find the difference and move the curve that much (how will this work for rotation?)

        self.getRange()
        #print "start Frame: %s. . . endFrame: %s"%(self.startF, self.endF)

        base = cmds.textFieldGrp(self.widgets["baseTFG"], q=True, tx=True)

        #get vals of changes for each
        self.dtx = cmds.floatFieldGrp(self.widgets["difTxFFG"], q=True, v1=True)
        self.dty = cmds.floatFieldGrp(self.widgets["difTyFFG"], q=True, v1=True)
        self.dtz = cmds.floatFieldGrp(self.widgets["difTzFFG"], q=True, v1=True)

        self.drx = cmds.floatFieldGrp(self.widgets["difRxFFG"], q=True, v1=True)
        self.dry = cmds.floatFieldGrp(self.widgets["difRyFFG"], q=True, v1=True)
        self.drz = cmds.floatFieldGrp(self.widgets["difRzFFG"], q=True, v1=True)

        self.dsx = cmds.floatFieldGrp(self.widgets["difSxFFG"], q=True, v1=True)
        self.dsy = cmds.floatFieldGrp(self.widgets["difSyFFG"], q=True, v1=True)
        self.dsz = cmds.floatFieldGrp(self.widgets["difSzFFG"], q=True, v1=True)

        sel = cmds.ls(sl=True, type="transform", l=True)

        if sel:
            for obj in sel:
                if self.startF:
                    #if there is a frame range, do the stuff with time here
                    cmds.keyframe(obj, at=("tx"), r=True, vc=self.dtx, t=(self.startF, self.endF))
                    cmds.keyframe(obj, at=("ty"), r=True, vc=self.dty, t=(self.startF, self.endF))
                    cmds.keyframe(obj, at=("tz"), r=True, vc=self.dtz, t=(self.startF, self.endF))

                    cmds.keyframe(obj, at=("rx"), r=True, vc=self.drx, t=(self.startF, self.endF))
                    cmds.keyframe(obj, at=("ry"), r=True, vc=self.dry, t=(self.startF, self.endF))
                    cmds.keyframe(obj, at=("rz"), r=True, vc=self.drz, t=(self.startF, self.endF))

                    cmds.keyframe(obj, at=("sx"), r=True, vc=self.dsx, t=(self.startF, self.endF))
                    cmds.keyframe(obj, at=("sy"), r=True, vc=self.dsy, t=(self.startF, self.endF))
                    cmds.keyframe(obj, at=("sz"), r=True, vc=self.dsz, t=(self.startF, self.endF))

                else:
                    #else just do it on the curve with no time argument
                    cmds.keyframe(obj, at=("tx"), r=True, vc=self.dtx)
                    cmds.keyframe(obj, at=("ty"), r=True, vc=self.dty)
                    cmds.keyframe(obj, at=("tz"), r=True, vc=self.dtz)

                    cmds.keyframe(obj, at=("rx"), r=True, vc=self.drx)
                    cmds.keyframe(obj, at=("ry"), r=True, vc=self.dry)
                    cmds.keyframe(obj, at=("rz"), r=True, vc=self.drz)

                    cmds.keyframe(obj, at=("sx"), r=True, vc=self.dsx)
                    cmds.keyframe(obj, at=("sy"), r=True, vc=self.dsy)
                    cmds.keyframe(obj, at=("sz"), r=True, vc=self.dsz)

                if obj==base:
                    #this section is to a) clear the new key if there wasn't one before or b) offset the new key BACK to the rest of the newly moved curve (by neg delta)
                    if self.txk:
                        #offset the value BACK at self.frame
                        cmds.keyframe(obj, at="tx", r=True, vc=-self.dtx, t=(self.frame, self.frame))
                    else:
                        #delete the key at self.frame if it exists
                        cmds.cutKey(obj, at="tx", t=(self.frame, self.frame), cl=True)
                    if self.tyk:
                        cmds.keyframe(obj, at="ty", r=True, vc=-self.dty, t=(self.frame, self.frame))
                    else:
                        cmds.cutKey(obj, at="ty", t=(self.frame, self.frame), cl=True)
                    if self.tzk:
                        cmds.keyframe(obj, at="tz", r=True, vc=-self.dtz, t=(self.frame, self.frame))
                    else:
                        cmds.cutKey(obj, at="tz", t=(self.frame, self.frame), cl=True)

                    if self.rxk:
                        cmds.keyframe(obj, at="rx", r=True, vc=-self.drx, t=(self.frame, self.frame))
                    else:
                        cmds.cutKey(obj, at="rx", t=(self.frame, self.frame), cl=True)
                    if self.ryk:
                        cmds.keyframe(obj, at="ry", r=True, vc=-self.dry, t=(self.frame, self.frame))
                    else:
                        cmds.cutKey(obj, at="ry", t=(self.frame, self.frame), cl=True)
                    if self.rzk:
                        cmds.keyframe(obj, at="rz", r=True, vc=-self.drz, t=(self.frame, self.frame))
                    else:
                        cmds.cutKey(obj, at="rz", t=(self.frame, self.frame), cl=True)

                    if self.sxk:
                        cmds.keyframe(obj, at="sx", r=True, vc=-self.dsx, t=(self.frame, self.frame))
                    else:
                        cmds.cutKey(obj, at="sx", t=(self.frame, self.frame), cl=True)
                    if self.syk:
                        cmds.keyframe(obj, at="sy", r=True, vc=-self.dsy, t=(self.frame, self.frame))
                    else:
                        cmds.cutKey(obj, at="sy", t=(self.frame, self.frame), cl=True)
                    if self.szk:
                        cmds.keyframe(obj, at="sz", r=True, vc=-self.dsz, t=(self.frame, self.frame))
                    else:
                        cmds.cutKey(obj, at="sz", t=(self.frame, self.frame), cl=True)
        else:
            cmds.warning("You've unselected stuff and you need some selections to shift!")

        #clear base object (so we don't double up on keyframe corrections)
        # cmds.textFieldGrp(self.widgets["baseTFG"], e=True, tx="")
        size = len(sel)
        if size == 1:
            cmds.text(self.widgets["doneText"], e=True, l='"%s Shifted! Select Other Objs to Shift or Clear!"'%sel[0], vis=True, bgc=(.5,.5,.6))
        elif size > 1:
            cmds.text(self.widgets["doneText"], e=True, l="%s Objects Shifted! Select Others to Shift or Clear!", vis=True, bgc=(.5,.5,.6))
        else:
            cmds.text(self.widgets["doneText"], e=True, l="You Need To Select Something To Shift!", vis=True, bgc=(.5,.5,.6))

    def clearAll(self, *args):
        """clears all the fields"""
        fieldNames = ["origTxFFG", "origTyFFG", "origTzFFG", "origSxFFG", "origSyFFG", "origSzFFG", "origRxFFG", "origRyFFG", "origRzFFG", "modTxFFG", "modTyFFG", "modTzFFG", "modRxFFG", "modRyFFG", "modRzFFG", "modSxFFG", "modSyFFG", "modSzFFG", "difTxFFG", "difTyFFG", "difTzFFG", "difRxFFG", "difRyFFG", "difRzFFG", "difSxFFG", "difSyFFG", "difSzFFG"]
        cmds.textFieldGrp(self.widgets["baseTFG"], e=True, tx="")
        for field in fieldNames:
            cmds.floatFieldGrp(self.widgets[field], e=True, v1=0)
        cmds.floatFieldGrp(self.widgets["baseFrameFFG"], e=True, v1=0)
        cmds.text(self.widgets["doneText"], e=True, l="No Status/All Cleared. Select a base object!", bgc=(.5,.5,.5))

    # def restoreBase(self, *args):
    # 	"""restores the base obj to the text field"""

    # 	cmds.textFieldGrp(self.widgets["baseTFG"], e=True, tx=self.baseObj)
    # 	cmds.floatFieldGrp(self.widgets["baseFrameFFG"], e=True, v1=self.frame)

    def enableRange(self, *args):
        """toggles the frame range option on/off depending on what's selected in the range radio grp"""
        #get the value of the radio button grp
        sel = cmds.radioButtonGrp(self.widgets["timeRBG"], q=True, sl=True)

        if sel == 3:
            cmds.floatFieldGrp(self.widgets["rangeFFG"], e=True, en=True)
        else:
            cmds.floatFieldGrp(self.widgets["rangeFFG"], e=True, en=False)

    def getRange(self, *args):
        """this finds the frame range and returns the StartFrame and EndFrame"""
        #get time RBG
        sel = cmds.radioButtonGrp(self.widgets["timeRBG"], q=True, sl=True)
        #get start and end frames (or "all")
        if sel == 3:
            #get values from ffg's
            self.startF = cmds.floatFieldGrp(self.widgets["rangeFFG"], q=True, v1=True)
            self.endF = cmds.floatFieldGrp(self.widgets["rangeFFG"], q=True, v2=True)

        if sel==1:
            self.startF = []
            self.endF = []
        if sel==2:
            #capture values from time slider
            self.getSliderRange()

    def getSliderRange(self, *args):
        """gets framerange in current scene and returns start and end frames"""
        #get timeslider range start
        self.startF = cmds.playbackOptions(query=True, min=True)
        self.endF = cmds.playbackOptions(query=True, max=True)

def animShift(*args):
    thisShift = AnimShift()


import warnings

import numpy as np
from scipy.ndimage import gaussian_filter1d
from scipy.spatial.distance import cdist
from scipy.stats import chi2
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA
from sklearn.metrics import silhouette_score
from sklearn.neighbors import NearestNeighbors

def isi_violations(spike_train, min_time, max_time, isi_threshold, min_isi=0):
    """Calculate ISI violations for a spike train.
    Based on metric described in Hill et al. (2011) J Neurosci 31: 8699-8705
    modified by Dan Denman from cortex-lab/sortingQuality GitHub by Nick Steinmetz
    Inputs:
    -------
    spike_train : array of spike times
    min_time : minimum time for potential spikes
    max_time : maximum time for potential spikes
    isi_threshold : threshold for isi violation
    min_isi : threshold for duplicate spikes
    Outputs:
    --------
    fpRate : rate of contaminating spikes as a fraction of overall rate
        A perfect unit has a fpRate = 0
        A unit with some contamination has a fpRate < 0.5
        A unit with lots of contamination has a fpRate > 1.0
    num_violations : total number of violations
    """
    duplicate_spikes = np.where(np.diff(spike_train) <= min_isi)[0]

    spike_train = np.delete(spike_train, duplicate_spikes + 1)
    isis = np.diff(spike_train)

    num_spikes = len(spike_train)
    num_violations = sum(isis < isi_threshold)
    violation_time = 2 * num_spikes * (isi_threshold - min_isi)
    total_rate = firing_rate(spike_train, min_time, max_time)
    violation_rate = num_violations / violation_time
    fpRate = violation_rate / total_rate

    return fpRate, num_violations

def presence_ratio(spike_train, min_time, max_time, num_bins=100):
    """Calculate fraction of time the unit is present within an epoch.
    Inputs:
    -------
    spike_train : array of spike times
    min_time : minimum time for potential spikes
    max_time : maximum time for potential spikes
    Outputs:
    --------
    presence_ratio : fraction of time bins in which this unit is spiking
    """

    h, b = np.histogram(spike_train, np.linspace(min_time, max_time, num_bins))

    return np.sum(h > 0) / num_bins

def firing_rate(spike_train, min_time=None, max_time=None):
    """Calculate firing rate for a spike train.
    If no temporal bounds are specified, the first and last spike time are used.
    Inputs:
    -------
    spike_train : numpy.ndarray
        Array of spike times in seconds
    min_time : float
        Time of first possible spike (optional)
    max_time : float
        Time of last possible spike (optional)
    Outputs:
    --------
    fr : float
        Firing rate in Hz
    """

    if min_time is not None and max_time is not None:
        duration = max_time - min_time
    else:
        duration = np.max(spike_train) - np.min(spike_train)

    fr = spike_train.size / duration

    return fr

def amplitude_cutoff(amplitudes, num_histogram_bins=500, histogram_smoothing_value=3):
    """ Calculate approximate fraction of spikes missing from a distribution of amplitudes
    Assumes the amplitude histogram is symmetric (not valid in the presence of drift)
    Inspired by metric described in Hill et al. (2011) J Neurosci 31: 8699-8705
    Input:
    ------
    amplitudes : numpy.ndarray
        Array of amplitudes (don't need to be in physical units)
    Output:
    -------
    fraction_missing : float
        Fraction of missing spikes (0-0.5)
        If more than 50% of spikes are missing, an accurate estimate isn't possible
    """

    h, b = np.histogram(amplitudes, num_histogram_bins, density=True)

    pdf = gaussian_filter1d(h, histogram_smoothing_value)
    support = b[:-1]

    peak_index = np.argmax(pdf)
    G = np.argmin(np.abs(pdf[peak_index:] - pdf[0])) + peak_index

    bin_size = np.mean(np.diff(support))
    fraction_missing = np.sum(pdf[G:]) * bin_size

    fraction_missing = np.min([fraction_missing, 0.5])

    return fraction_missing

def mahalanobis_metrics(all_pcs, all_labels, this_unit_id):
    # def mahalanobis_metrics(pcs_for_this_unit, pcs_for_other_units):
    """ Calculates isolation distance and L-ratio (metrics computed from Mahalanobis distance)
    Based on metrics described in Schmitzer-Torbert et al. (2005) Neurosci 131: 1-11
    Inputs:
    -------
    all_pcs : numpy.ndarray (num_spikes x PCs)
        2D array of PCs for all spikes
    all_labels : numpy.ndarray (num_spikes x 0)
        1D array of cluster labels for all spikes
    this_unit_id : Int
        number corresponding to unit for which these metrics will be calculated
    Outputs:
    --------
    isolation_distance : float
        Isolation distance of this unit
    l_ratio : float
        L-ratio for this unit
    """

    pcs_for_this_unit = all_pcs[all_labels == this_unit_id, :]
    pcs_for_other_units = all_pcs[all_labels != this_unit_id, :]
    mean_value = np.expand_dims(np.mean(pcs_for_this_unit, 0), 0)

    try:
        VI = np.linalg.inv(np.cov(pcs_for_this_unit.T))
    except np.linalg.linalg.LinAlgError:  # case of singular matrix
        return np.nan, np.nan

    mahalanobis_other = np.sort(cdist(mean_value,
                                      pcs_for_other_units,
                                      'mahalanobis', VI=VI)[0])

    mahalanobis_self = np.sort(cdist(mean_value,
                                     pcs_for_this_unit,
                                     'mahalanobis', VI=VI)[0])

    n = np.min([pcs_for_this_unit.shape[0], pcs_for_other_units.shape[0]])  # number of spikes

    if np.isnan(mahalanobis_other.any()):
        warnings.warn(f'Nans in mahalanobis_other for cluster {this_unit_id}. l_ratio will be nan')
    if n >= 2:

        dof = pcs_for_this_unit.shape[1]  # number of features

        l_ratio = np.sum(1 - chi2.cdf(pow(mahalanobis_other, 2), dof)) / mahalanobis_self.shape[0]
        isolation_distance = pow(mahalanobis_other[n - 1], 2)

    else:
        l_ratio = np.nan
        isolation_distance = np.nan

    return isolation_distance, l_ratio

def lda_metrics(all_pcs, all_labels, this_unit_id):
    """ Calculates d-prime based on Linear Discriminant Analysis
    Based on metric described in Hill et al. (2011) J Neurosci 31: 8699-8705
    Inputs:
    -------
    all_pcs : numpy.ndarray (num_spikes x PCs)
        2D array of PCs for all spikes
    all_labels : numpy.ndarray (num_spikes x 0)
        1D array of cluster labels for all spikes
    this_unit_id : Int
        number corresponding to unit for which these metrics will be calculated
    Outputs:
    --------
    d_prime : float
        Isolation distance of this unit
    l_ratio : float
        L-ratio for this unit
    """

    X = all_pcs

    y = np.zeros((X.shape[0],), dtype='bool')
    y[all_labels == this_unit_id] = True

    # Quick sanity check:
    from sklearn.utils.multiclass import unique_labels
    n_components = 1
    classes_ = unique_labels(y)
    max_components = min(len(classes_) - 1, X.shape[1])
    assert n_components <= max_components
    # Continue with calculation
    lda = LDA(n_components=1)

    X_flda = lda.fit_transform(X, y)

    flda_this_cluster = X_flda[np.where(y)[0]]
    flda_other_cluster = X_flda[np.where(np.invert(y))[0]]

    d_prime = (np.mean(flda_this_cluster) - np.mean(flda_other_cluster)) / np.sqrt(
        0.5 * (np.std(flda_this_cluster) ** 2 + np.std(flda_other_cluster) ** 2))

    return d_prime

def nearest_neighbors_metrics(all_pcs, all_labels, this_unit_id, max_spikes_for_nn, n_neighbors):
    """ Calculates unit contamination based on NearestNeighbors search in PCA space
    Based on metrics described in Chung, Magland et al. (2017) Neuron 95: 1381-1394
    Inputs:
    -------
    all_pcs : numpy.ndarray (num_spikes x PCs)
        2D array of PCs for all spikes
    all_labels : numpy.ndarray (num_spikes x 0)
        1D array of cluster labels for all spikes
    this_unit_id : Int
        number corresponding to unit for which these metrics will be calculated
    max_spikes_for_nn : Int
        number of spikes to use (calculation can be very slow when this number is >20000)
    n_neighbors : Int
        number of neighbors to use
    Outputs:
    --------
    hit_rate : float
        Fraction of neighbors for target cluster that are also in target cluster
    miss_rate : float
        Fraction of neighbors outside target cluster that are in target cluster
    """

    total_spikes = all_pcs.shape[0]
    ratio = max_spikes_for_nn / total_spikes
    this_unit = all_labels == this_unit_id

    X = np.concatenate((all_pcs[this_unit, :], all_pcs[np.invert(this_unit), :]), 0)

    n = np.sum(this_unit)

    if ratio < 1:
        inds = np.arange(0, X.shape[0] - 1, 1 / ratio).astype('int')
        X = X[inds, :]
        n = int(n * ratio)

    nbrs = NearestNeighbors(n_neighbors=n_neighbors, algorithm='ball_tree').fit(X)
    distances, indices = nbrs.kneighbors(X)

    this_cluster_inds = np.arange(n)

    this_cluster_nearest = indices[:n, 1:].flatten()
    other_cluster_nearest = indices[n:, 1:].flatten()

    hit_rate = np.mean(this_cluster_nearest < n)
    miss_rate = np.mean(other_cluster_nearest < n)

    return hit_rate, miss_rate

def silhouette_score_inner_loop(i, cluster_ids, cluster_labels, all_pcs):
    """
    Helper to loop over cluster_ids in one dimension. We dont want to loop over both dimensions in parallel-
    that will create too much worker overhead
    Args:
        i: index of first dimension
        cluster_ids: iterable of cluster ids
    Returns: scores for dimension j
    """
    scores_1d = []
    assert (cluster_labels.shape[0] ==
            all_pcs.shape[0]), f'Found {cluster_labels.shape[0]} cluster_ids but {all_pcs.shape[0]} pcs!'
    for j in cluster_ids:
        if j > i:
            inds = np.in1d(cluster_labels, np.array([i, j]))
            X = all_pcs[inds, :]
            labels = cluster_labels[inds]
            # len(np.unique(labels))=1 Can happen if total_spikes is low:
            if (len(labels) > 2) and (len(np.unique(labels)) > 1):
                scores_1d.append(silhouette_score(X, labels))
            else:
                scores_1d.append(np.nan)
        else:
            scores_1d.append(np.nan)
    return scores_1d

def calculate_silhouette_score(spike_clusters,
                               spike_templates,
                               total_units,
                               pc_features,
                               pc_feature_ind,
                               total_spikes,
                               do_parallel=True):
    """
    :param spike_clusters:
    :param pc_features:
    :param pc_feature_ind:
    :param total_spikes:
    :return:
    """
    import warnings
    if total_spikes > spike_clusters.shape[0]:
        total_spikes = spike_clusters.shape[0]  # take all spikes

    cluster_ids = np.unique(spike_clusters)
    random_spike_inds = np.random.permutation(spike_clusters.size)
    random_spike_inds = random_spike_inds[:total_spikes]
    num_pc_features = pc_features.shape[1]
    num_channels = np.max(pc_feature_ind) + 1
    all_pcs = np.zeros((total_spikes, num_channels * num_pc_features))

    for idx, i in enumerate(random_spike_inds):
        unit_id = spike_templates[i]
        channels = pc_feature_ind[unit_id, :]

        for j in range(0, num_pc_features):
            all_pcs[idx, channels + num_channels * j] = pc_features[i, j, :]

    cluster_labels = spike_clusters[random_spike_inds]

    SS = np.empty((total_units, total_units))
    SS[:] = np.nan

    # Build lists
    if do_parallel:
        from joblib import Parallel, delayed
        scores = Parallel(n_jobs=-1, verbose=1)(delayed(silhouette_score_inner_loop)
                                                (i, cluster_ids, cluster_labels, all_pcs) for i in cluster_ids)
    else:
        scores = [silhouette_score_inner_loop(i, cluster_ids, cluster_labels, all_pcs) for i in cluster_ids]

    # Fill the 2d array
    for i, col_score in enumerate(scores):
        if np.isnan(col_score).all():
            warnings.warn(f'Cluster {i} has all-nan col_score. It will have nan silhouette score')
        for j, one_score in enumerate(col_score):
            SS[i, j] = one_score

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        a = np.nanmin(SS, 0)
        b = np.nanmin(SS, 1)
    return np.array([np.nanmin([a, b]) for a, b in zip(a, b)])

def calculate_drift_metrics(spike_times,
                            spike_clusters,
                            spike_templates,
                            pc_features,
                            pc_feature_ind,
                            interval_length,
                            min_spikes_per_interval,
                            do_parallel=True):
    """

    :param spike_times:
    :param spike_clusters:
    :param total_units:
    :param pc_features:
    :param pc_feature_ind:
    :param interval_length:
    :param min_spikes_per_interval:
    :return:
    """

    def get_spike_depths(spike_clusters, pc_features, pc_feature_ind):

        """
        Calculates the distance (in microns) of individual spikes from the probe tip
        This implementation is based on Matlab code from github.com/cortex-lab/spikes
        Input:
        -----
        spike_clusters : numpy.ndarray (N x 0)
            Cluster IDs for N spikes
        pc_features : numpy.ndarray (N x channels x num_PCs)
            PC features for each spike
        pc_feature_ind  : numpy.ndarray (M x channels)
            Channels used for PC calculation for each unit
        Output:
        ------
        spike_depths : numpy.ndarray (N x 0)
            Distance (in microns) from each spike waveform from the probe tip
        """

        pc_features_copy = np.copy(pc_features)
        pc_features_copy = np.squeeze(pc_features_copy[:, 0, :])
        pc_features_copy[pc_features_copy < 0] = 0
        pc_power = pow(pc_features_copy, 2)
        if np.isnan(pc_power).any():
            nan_clusters = spike_clusters[np.isnan(pc_power)]
            warnings.warn(f'pc_power for {nan_clusters} are nan')

        spike_feat_ind = pc_feature_ind[spike_clusters, :]
        spike_depths = np.sum(spike_feat_ind * pc_power, 1) / np.sum(pc_power, 1)

        return spike_depths * 10

    def calc_one_cluster(cluster_id):
        """
        Helper to calculate drift for one cluster
        Args:
            cluster_id:

        Returns:
            max_drift, cumulative_drift
        """
        in_cluster = spike_clusters == cluster_id
        times_for_cluster = spike_times[in_cluster]
        depths_for_cluster = depths[in_cluster]

        median_depths = []

        for t1, t2 in zip(interval_starts, interval_ends):

            in_range = (times_for_cluster > t1) * (times_for_cluster < t2)

            if np.sum(in_range) >= min_spikes_per_interval:
                median_depths.append(np.median(depths_for_cluster[in_range]))
            else:
                median_depths.append(np.nan)

        with warnings.catch_warnings():
            warnings.simplefilter("ignore")  # RuntimeWarning: All-NaN slice encountered
            median_depths = np.array(median_depths)
            max_drift = np.around(np.nanmax(median_depths) - np.nanmin(median_depths), 2)
            cumulative_drift = np.around(np.nansum(np.abs(np.diff(median_depths))), 2)
        return max_drift, cumulative_drift

    max_drifts = []
    cumulative_drifts = []

    depths = get_spike_depths(spike_templates, pc_features, pc_feature_ind)

    interval_starts = np.arange(np.min(spike_times), np.max(spike_times), interval_length)
    interval_ends = interval_starts + interval_length

    cluster_ids = np.unique(spike_clusters)

    if do_parallel:
        from joblib import Parallel, delayed
        meas = Parallel(n_jobs=-1, verbose=1)(delayed(calc_one_cluster)(cluster_id)
                                              for cluster_id in cluster_ids)
    else:
        meas = [calc_one_cluster(cluster_id) for cluster_id in cluster_ids]

    for max_drift, cumulative_drift in meas:
        max_drifts.append(max_drift)
        cumulative_drifts.append(max_drift)
    return np.array(max_drifts), np.array(cumulative_drifts)


import site
site.addsitedir('/home/mkotak/atomic_architects/venv/jax_env/lib/python3.10/site-packages')
import time
import haiku as hk
import jax
import jax.numpy as jnp
import jaxlib

import e3nn_jax
from e3nn_jax._src.utils.jit import jit_code

# Constants
irreps_string = "1e"
batch = 10

# Training params
n = 1000
warmup = -1

irreps_in1_jax = e3nn_jax.Irreps(irreps_string)
irreps_in2_jax = e3nn_jax.Irreps(irreps_string)
irreps_out_jax = e3nn_jax.Irreps(irreps_string)

def k():
    k.key, x = jax.random.split(k.key)
    return x

k.key = jax.random.PRNGKey(0)

@hk.without_apply_rng
@hk.transform
def tp_jax(x1, x2):
    return e3nn_jax.haiku.FullyConnectedTensorProduct(irreps_out_jax)(x1, x2)

inputs_jax = (e3nn_jax.normal(irreps_in1_jax, k(), (batch,)),
              e3nn_jax.normal(irreps_in2_jax, k(), (batch,)))

w = tp_jax.init(k(), *inputs_jax)
w, inputs_jax = jax.tree_util.tree_map(jax.device_put, (w, inputs_jax))

f  = tp_jax.apply
f2 = f
f = jax.value_and_grad(
    lambda w, x1, x2: sum(jnp.sum(jnp.tanh(out)) for out in
                          jax.tree_util.tree_leaves(f2(w, x1, x2)))
)

f = jax.jit(f)

for _ in range(20):
    z = f(w, *inputs_jax)
    jax.tree_util.tree_map(lambda x: x.block_until_ready(), z)

with jax.profiler.trace("./profile/e3nn-jax"):
    z = f(w, *inputs_jax)
    jax.tree_util.tree_map(lambda x: x.block_until_ready(), z)

# with open("xla.txt", "wt") as file:
#     file.write(jit_code(f, w, *inputs_jax))


import pytest

from .pages.basket_page import BasketPage
from .pages.product_page import ProductPage

class TestProductPage:

    @pytest.mark.parametrize('promo_offer',
                             ["offer0", "offer1", "offer2", "offer3", "offer4", "offer5", "offer6",
                              pytest.param("offer7", marks=pytest.mark.xfail),
                              "offer8", "offer9"])
    def test_guest_can_add_product_to_basket(self, browser, promo_offer):
        link = 'http://selenium1py.pythonanywhere.com/catalogue/coders-at-work_207/'
        promo_link = '?promo='

        page = ProductPage(browser, link + promo_link + promo_offer)
        page.open()
        page.add_to_basket()
        page.solve_quiz_and_get_code()
        page.check_success_alert_present()
        page.check_product_and_alert_same()
        page.check_product_and_basket_price_same()

    def test_guest_should_see_login_link_on_product_page(self, browser):
        link = "http://selenium1py.pythonanywhere.com/catalogue/the-city-and-the-stars_95/"
        page = ProductPage(browser, link)
        page.open()
        page.check_login_link()

    def test_guest_can_go_to_login_page_from_product_page(self, browser):
        link = "http://selenium1py.pythonanywhere.com/catalogue/the-city-and-the-stars_95/"
        page = ProductPage(browser, link)
        page.open()
        page.check_login_link()
        page.go_to_login_page()
        assert page.is_string_in_current_url('login'), 'Not login page'

    def test_guest_cant_see_product_in_basket_opened_from_product_page(self, browser):
        link = "http://selenium1py.pythonanywhere.com/catalogue/the-city-and-the-stars_95/"
        page = ProductPage(browser, link)
        page.open()
        page.go_to_basket()
        basket_page = BasketPage(browser, browser.current_url)
        basket_page.check_basket_empty_message()
        basket_page.check_basket_not_have_checkout()

    @pytest.mark.xfail
    def test_guest_cant_see_success_message_after_adding_product_to_basket(self, browser):
        link = "http://selenium1py.pythonanywhere.com/catalogue/coders-at-work_207/"
        page = ProductPage(browser, link)
        page.open()
        page.add_to_basket()
        page.should_not_be_success_message()

    def test_guest_cant_see_success_message(self, browser):
        link = "http://selenium1py.pythonanywhere.com/catalogue/coders-at-work_207/"
        page = ProductPage(browser, link)
        page.open()
        page.should_not_be_success_message()

    @pytest.mark.xfail
    def test_message_dissapeared_after_adding_product_to_basket(self, browser):
        link = "http://selenium1py.pythonanywhere.com/catalogue/coders-at-work_207/"
        page = ProductPage(browser, link)
        page.open()
        page.add_to_basket()
        page.should_dissapear_success_message()


from collections import Counter
def decToBinary(n):
  a=[]
  for i in range(3, -1, -1):
    k = n >> i;
    if (k & 1):
      a.append(1)
    else:
      a.append(0)
  return a
def convert(list):
  s = [str(i) for i in list]
  res = int("".join(s))

  return(res)
# Driver Code
s=int(input())
p = map(int,input().strip().split(' '))
num0=[]
num1=[]
for i in p:
  num0.append(decToBinary(i).count(0))
  num1.append(decToBinary(i).count(1))
print(num0)
print(num1)
num10 = []
for i in range(len(num0)):
  num10.append(num0[i]+num1[i])
print(num10)
c=Counter(num10)
numerr = []
if(convert(c.values())>1):
  numerr = decToBinary(convert(c.values()))
else:
  numerr=decToBinary(0)
numstr=""
for j in numerr:
  numstr += str(j)
print(numstr)


import requests
from datetime import datetime
from decouple import config

user_input = input("Tell me which exercise you did: ")
today = datetime.today()
today_str = today.strftime("%d/%m/%Y")
time = datetime.now()
time_str = time.strftime("%H:%M:%S")

# Nutritionix API
nu_api_id = config('NU_API_ID')
nu_api_key = config('NU_API_KEY')
nu_api_endpoint = "https://trackapi.nutritionix.com/v2/natural/exercise"
nu_api_header = {
    "x-app-id": nu_api_id,
    "x-app-key": nu_api_key,
    "Content-Type": "application/json"
}
nu_api_body = {
    "query": user_input
}
exercise_response = requests.post(nu_api_endpoint, headers=nu_api_header, json=nu_api_body)
exercise_data = exercise_response.json()['exercises']

# Sheety API
# Query key must be the singular of spreadsheet name
sh_api_endpoint = config('SH_API_ENDPOINT')
sh_api_header = {
    "Content-Type": "application/json"
}
for exercise in exercise_data:
    sh_api_body = {
        "workout": {
            "date": today_str,
            "time": time_str,
            "exercise": exercise['name'].title(),
            "duration": exercise['duration_min'],
            "calories": exercise['nf_calories']
        }
    }
    sheet_response = requests.post(sh_api_endpoint, headers=sh_api_header, json=sh_api_body)
    sheet_response.raise_for_status()


from django.contrib import messages
from django.contrib.auth.decorators import login_required
from django.shortcuts import get_object_or_404, render, render_to_response

from . import forms
from . import models
# Create your views here.

def availability_page(request):
    items = models.Available.objects.all()
    return render_to_response('availability.html', {'items': items})

def activity_list(request):
    activities = models.Activity.objects.all()
    return render(request, 'activities/activity_list.html', {'activities': activities})

def activity_detail(request, pk):
    activity = get_object_or_404(models.Activity, pk=pk)
    return render(request, 'activities/activity_detail.html', {'activity': activity})

def package_detail(request, activity_pk, package_pk):
    package = get_object_or_404(models.PartyPackage, activity_id=activity_pk, pk=package_pk)
    return render(request, 'activities/package_detail.html', {'package': package})

@login_required
def activity_create(request):
    form = forms.ActivityForm(request.POST or None)
    if form.is_valid():
        instance = form.save(commit=False)
        instance.save()
        messages.add_message(request, messages.SUCCESS,
                             "Activity added!")
    context = {
        'form': form
    }

    return render(request, "activities/activity_form.html", context)

@login_required
def availability_create(request):
    form = forms.AvailableForm(request.POST or None)
    if form.is_valid():
        instance = form.save(commit=False)
        instance.save()
        messages.add_message(request,  messages.SUCCESS,
                             "Person added!")

    context = {
            'form': form,
    }
    return render(request, "activities/available_form.html", context)

@login_required
def package_create(request):
    form = forms.PartyPackageForm(request.POST or None)
    if form.is_valid():
        instance = form.save(commit=False)
        instance.save()
        messages.add_message(request, messages.SUCCESS,
                             "Packageg added!")
    return render(request, "activities/package_form.html", {"form": form})