#!/usr/bin/env python
# -*- coding: utf-8 -*-
# File              : layers.py
# Author            : Sun Fu <cstsunfu@gmail.com>
# Date              : 21.03.2018
# Last Modified Date: 11.05.2018
# Last Modified By  : Sun Fu <cstsunfu@gmail.com>
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from torch.nn import Conv1d, Linear
from typing import Callable

# ------------------------------------------------------------------------------
# Modules
# ------------------------------------------------------------------------------

class Answer(nn.Module):
    """A bilinear attention layer over a sequence X w.r.t y:
    * o_i = softmax(x_i'Wy) for x_i in X.

    Optionally don't normalize output weights.
    """
    RNN_TYPES = {'lstm': nn.LSTM, 'gru': nn.GRU, 'rnn': nn.RNN}

    def __init__(self, x_size, y_size, opt, first=False, identity=True):
        super(Answer, self).__init__()
        self.gru = torch.nn.GRU(x_size, y_size, batch_first=True)
        self.linear = nn.Linear(y_size, x_size)
        # self.linearx = nn.Linear(x_size, x_size)
        self.first = first
        if self.first is True:
            self.doc_rnn = StackedBRNN(
                input_size=x_size * 2,
                hidden_size=opt['hidden_size'],
                num_layers=opt['doc_layers'],
                dropout_rate=opt['dropout_rnn'],
                dropout_output=opt['dropout_rnn_output'],
                concat_layers=opt['concat_rnn_layers'],
                rnn_type=self.RNN_TYPES[opt['rnn_type']],
                padding=opt['rnn_padding'],
            )

    def forward(self, x, y, x_mask):
        """
        x = batch * len * h1
        y = batch * h2
        x_mask = batch * len
        """
        Wy = self.linear(y)
        xWy = x.bmm(Wy.unsqueeze(2)).squeeze(2)
        xWy.data.masked_fill_(x_mask.data, -float('inf'))
        x_new = x

        if self.first or self.training is not True:
            alpha = F.softmax(xWy, dim=1)
            if self.first is True:
                x_new = alpha.unsqueeze(2).expand(x.size(0), x.size(1), x.size(2)) * x
                x_new = torch.cat((x_new, x), 2)
                x_new = self.doc_rnn(x_new, x_mask)
        if self.training:
            alpha = F.log_softmax(xWy, dim=1)

        return alpha, x_new

class CharEmbeddingLayer(nn.Module):
    def __init__(self, char_single_embedding_dim, char_embedding_dim, char_embedding):
        super(CharEmbeddingLayer, self).__init__()
        self.char_single_embedding_dim = char_single_embedding_dim
        self.char_embedding_dim = char_embedding_dim
        # self.char_vocab_size = char_vocab_size

        self.embedding_lookup = nn.Embedding(char_embedding.size(0), char_embedding.size(1), padding_idx=0)
        self.embedding_lookup.weight.data[2:, :] = char_embedding[2:, :]
    # def __init__(self, embedding_dim, num_filters=100, ngram_filter_sizes=(2, 3, 4, 5), conv_layer_activation='relu', output_dim=50):
        self.cnn = CNN(embedding_dim=char_single_embedding_dim, output_dim=char_embedding_dim)

    def forward(self, text_char, text_char_mask):
        batch_size, max_length, max_word_length = text_char.size()
        # embedding look up
        text_char = text_char.contiguous().view(batch_size * max_length, max_word_length)
        text_char_mask = text_char_mask.view(batch_size * max_length, max_word_length)

        text_char = self.embedding_lookup(text_char)
        # dropout + cnn
        # if self.dropout < 1.0:
        #     assert is_train is not None
        #     if is_train:
        #         text_char = nn.Dropout(p=self.dropout)(text_char)
        text_char = self.cnn(text_char, text_char_mask)
        text_char = text_char.contiguous().view(batch_size * max_length * self.char_embedding_dim, -1)
        text_char = nn.functional.relu(text_char)

        text_char = torch.max(text_char, 1)[0]
        text_char = text_char.contiguous().view(batch_size, max_length, self.char_embedding_dim)
        return text_char

class CNN(nn.Module):

    """Docstring for CNN"""

    def __init__(self, embedding_dim, num_filters=100, ngram_filter_sizes=(2, 3, 4, 5), conv_layer_activation='relu', output_dim=50):
        """TODO: to be defined1. """
        nn.Module.__init__(self)

        self.embedding_dim = embedding_dim
        self.num_filters = num_filters
        self.ngram_filter_sizes = ngram_filter_sizes
        self.conv_layer_activation = torch.nn.functional.relu
        self.output_dim = output_dim

        self.convolution_layers = [Conv1d(in_channels=self.embedding_dim,
                                          out_channels=self.num_filters,
                                          kernel_size=ngram_size).cuda()
                                   for ngram_size in self.ngram_filter_sizes]

        maxpool_output_dim = self.num_filters * len(self.ngram_filter_sizes)
        if self.output_dim:
            self.projection_layer = Linear(maxpool_output_dim, self.output_dim)
        else:
            self.projection_layer = None
            self.output_dim = maxpool_output_dim

    def forward(self, tokens, mask):
        """TODO: Docstring for forward.

        :tokens: TODO
        :mask: TODO
        :returns: TODO

        """
        if mask is not None:
            tokens = tokens * mask.unsqueeze(-1).float()

        tokens = torch.transpose(tokens, 1, 2)

        filter_outputs = []
        for convolution_layer in self.convolution_layers:
            filter_outputs.append(
                self.conv_layer_activation(convolution_layer(tokens)).max(dim=2)[0]
            )

        maxpool_output = torch.cat(filter_outputs, dim=1) if len(filter_outputs) > 1 else filter_outputs[0]

        if self.projection_layer:
            result = self.projection_layer(maxpool_output)
        else:
            result = maxpool_output
        return result

class Highway(torch.nn.Module):
    """
    A `Highway layer <https://arxiv.org/abs/1505.00387>`_ does a gated combination of a linear
    transformation and a non-linear transformation of its input.  :math:`y = g * x + (1 - g) *
    f(A(x))`, where :math:`A` is a linear transformation, :math:`f` is an element-wise
    non-linearity, and :math:`g` is an element-wise gate, computed as :math:`sigmoid(B(x))`.

    This module will apply a fixed number of highway layers to its input, returning the final
    result.

    Parameters
    ----------
    input_dim : ``int``
        The dimensionality of :math:`x`.  We assume the input has shape ``(batch_size,
        input_dim)``.
    num_layers : ``int``, optional (default=``1``)
        The number of highway layers to apply to the input.
    activation : ``Callable[[torch.Tensor], torch.Tensor]``, optional (default=``torch.nn.functional.relu``)
        The non-linearity to use in the highway layers.
    """
    def __init__(self,
                 input_dim: int,
                 num_layers: int = 2,
                 activation: Callable[[torch.Tensor], torch.Tensor] = torch.nn.functional.relu) -> None:
        super(Highway, self).__init__()
        self._input_dim = input_dim
        self._layers = torch.nn.ModuleList([torch.nn.Linear(input_dim, input_dim * 2)
                                            for _ in range(num_layers)])
        self._activation = activation
        for layer in self._layers:
            # We should bias the highway layer to just carry its input forward.  We do that by
            # setting the bias on `B(x)` to be positive, because that means `g` will be biased to
            # be high, to we will carry the input forward.  The bias on `B(x)` is the second half
            # of the bias vector in each Linear layer.
            layer.bias[input_dim:].data.fill_(1)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:  # pylint: disable=arguments-differ
        current_input = inputs
        for layer in self._layers:
            projected_input = layer(current_input)
            linear_part = current_input
            # NOTE: if you modify this, think about whether you should modify the initialization
            # above, too.
            nonlinear_part = projected_input[:, (0 * self._input_dim):(1 * self._input_dim)]
            gate = projected_input[:, (1 * self._input_dim):(2 * self._input_dim)]
            nonlinear_part = self._activation(nonlinear_part)
            gate = torch.nn.functional.sigmoid(gate)
            current_input = gate * linear_part + (1 - gate) * nonlinear_part
        return current_input

class TimeDistributed(torch.nn.Module):
    """
    Given an input shaped like ``(batch_size, time_steps, [rest])`` and a ``Module`` that takes
    inputs like ``(batch_size, [rest])``, ``TimeDistributed`` reshapes the input to be
    ``(batch_size * time_steps, [rest])``, applies the contained ``Module``, then reshapes it back.

    Note that while the above gives shapes with ``batch_size`` first, this ``Module`` also works if
    ``batch_size`` is second - we always just combine the first two dimensions, then split them.
    """
    def __init__(self, module):
        super(TimeDistributed, self).__init__()
        self._module = module

    def forward(self, *inputs):  # pylint: disable=arguments-differ
        reshaped_inputs = []
        for input_tensor in inputs:
            input_size = input_tensor.size()
            if len(input_size) <= 2:
                raise RuntimeError("No dimension to distribute: " + str(input_size))

            # Squash batch_size and time_steps into a single axis; result has shape
            # (batch_size * time_steps, input_size).
            squashed_shape = [-1] + [x for x in input_size[2:]]
            reshaped_inputs.append(input_tensor.contiguous().view(*squashed_shape))

        reshaped_outputs = self._module(*reshaped_inputs)

        # Now get the output back into the right shape.
        # (batch_size, time_steps, [hidden_size])
        new_shape = [input_size[0], input_size[1]] + [x for x in reshaped_outputs.size()[1:]]
        outputs = reshaped_outputs.contiguous().view(*new_shape)

        return outputs

class StackedBRNN(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers,
                 dropout_rate=0, dropout_output=False, rnn_type=nn.LSTM,
                 concat_layers=False, padding=False):
        super(StackedBRNN, self).__init__()
        self.padding = padding
        self.dropout_output = dropout_output
        self.dropout_rate = dropout_rate
        self.num_layers = num_layers
        self.concat_layers = concat_layers
        self.rnns = nn.ModuleList()
        for i in range(num_layers):
            input_size = input_size if i == 0 else 2 * hidden_size
            self.rnns.append(rnn_type(input_size, hidden_size,
                                      num_layers=1,
                                      bidirectional=True))

    def forward(self, x, x_mask):
        """Can choose to either handle or ignore variable length sequences.
        Always handle padding in eval.
        """
        # No padding necessary.
        if x_mask.data.sum() == 0:
            return self._forward_unpadded(x, x_mask)
        # Pad if we care or if its during eval.
        if self.padding or not self.training:
            return self._forward_padded(x, x_mask)
        # We don't care.
        return self._forward_unpadded(x, x_mask)

    def _forward_unpadded(self, x, x_mask):
        """Faster encoding that ignores any padding."""
        # Transpose batch and sequence dims
        x = x.transpose(0, 1)

        # Encode all layers
        outputs = [x]
        for i in range(self.num_layers):
            rnn_input = outputs[-1]

            # Apply dropout to hidden input
            if self.dropout_rate > 0:
                rnn_input = F.dropout(rnn_input,
                                      p=self.dropout_rate,
                                      training=self.training)
            # Forward
            rnn_output = self.rnns[i](rnn_input)[0]
            outputs.append(rnn_output)

        # Concat hidden layers
        if self.concat_layers:
            output = torch.cat(outputs[1:], 2)
        else:
            output = outputs[-1]

        # Transpose back
        output = output.transpose(0, 1)

        # Dropout on output layer
        if self.dropout_output and self.dropout_rate > 0:
            output = F.dropout(output,
                               p=self.dropout_rate,
                               training=self.training)
        return output

    def _forward_padded(self, x, x_mask):
        """Slower (significantly), but more precise,
        encoding that handles padding."""
        # Compute sorted sequence lengths
        lengths = x_mask.data.eq(0).long().sum(1).squeeze()
        _, idx_sort = torch.sort(lengths, dim=0, descending=True)
        _, idx_unsort = torch.sort(idx_sort, dim=0)

        lengths = list(lengths[idx_sort])
        idx_sort = Variable(idx_sort)
        idx_unsort = Variable(idx_unsort)

        # Sort x
        x = x.index_select(0, idx_sort)

        # Transpose batch and sequence dims
        x = x.transpose(0, 1)

        # Pack it up
        rnn_input = nn.utils.rnn.pack_padded_sequence(x, lengths)

        # Encode all layers
        outputs = [rnn_input]
        for i in range(self.num_layers):
            rnn_input = outputs[-1]

            # Apply dropout to input
            if self.dropout_rate > 0:
                dropout_input = F.dropout(rnn_input.data,
                                          p=self.dropout_rate,
                                          training=self.training)
                rnn_input = nn.utils.rnn.PackedSequence(dropout_input,
                                                        rnn_input.batch_sizes)
            outputs.append(self.rnns[i](rnn_input)[0])

        # Unpack everything
        for i, o in enumerate(outputs[1:], 1):
            outputs[i] = nn.utils.rnn.pad_packed_sequence(o)[0]

        # Concat hidden layers or take final
        if self.concat_layers:
            output = torch.cat(outputs[1:], 2)
        else:
            output = outputs[-1]

        # Transpose and unsort
        output = output.transpose(0, 1)
        output = output.index_select(0, idx_unsort)

        # Pad up to original batch sequence length
        if output.size(1) != x_mask.size(1):
            padding = torch.zeros(output.size(0),
                                  x_mask.size(1) - output.size(1),
                                  output.size(2)).type(output.data.type())
            output = torch.cat([output, Variable(padding)], 1)

        # Dropout on output layer
        if self.dropout_output and self.dropout_rate > 0:
            output = F.dropout(output,
                               p=self.dropout_rate,
                               training=self.training)
        return output

class SeqAttnMatch(nn.Module):
    """Given sequences X and Y, match sequence Y to each element in X.
    * o_i = sum(alpha_j * y_j) for i in X
    * alpha_j = softmax(y_j * x_i)
    """
    def __init__(self, input_size, identity=False):
        super(SeqAttnMatch, self).__init__()
        # self.dropout_rate = dropout_rate
        if not identity:
            self.linear = nn.Linear(input_size, input_size)
        else:
            self.linear = None
        self.bilstm1 = nn.LSTM(input_size, input_size//2, num_layers=1, batch_first=True, bidirectional=True)
        self.bilstm2 = nn.LSTM(input_size, input_size//2, num_layers=1, batch_first=True, bidirectional=True)

    def forward(self, x, y, x_mask, y_mask, dropout_rate=0.0):
        """Input shapes:
            x = batch * len1 * h
            y = batch * len2 * h
            x_mask = batch * len1
            y_mask = batch * len2
        Output shapes:
            matched_seq = batch * len1 * h
        """
        # Project vectors

        E1_D = x
        E1_Q = y

        # Compute scores
        # size batch * len1 * len2
        A = E1_D.bmm(E1_Q.transpose(2, 1))
        # size batch * len2 * len1
        A_t = E1_Q.bmm(E1_D.transpose(2, 1))
        # A_t = A.transpose(2, 1)

        # Mask padding
        y_mask = y_mask.unsqueeze(1).expand(A.size())
        A.data.masked_fill_(y_mask.data, -float('inf'))

        x_mask = x_mask.unsqueeze(1).expand(A_t.size())
        A_t.data.masked_fill_(y_mask.data, -float('inf'))

        # Normalize with softmax
        S1_D_alpha_flat = F.softmax(A.view(-1, E1_Q.size(1)), dim=1)
        S1_D_alpha = S1_D_alpha_flat.view(-1, E1_D.size(1), E1_Q.size(1))

        # Take weighted average
        S1_D = S1_D_alpha.bmm(E1_Q)

        S1_Q_alpha_flat = F.softmax(A_t.view(-1, E1_D.size(1)), dim=1)
        S1_Q_alpha = S1_Q_alpha_flat.view(-1, E1_Q.size(1), E1_D.size(1))
        S1_Q = S1_Q_alpha.bmm(E1_D)

        C1_D = S1_D_alpha.bmm(S1_Q)

        E2_D = self.bilstm2(S1_D)[0]
        E2_Q = self.bilstm2(S1_Q)[0]
        E2_D = nn.functional.dropout(E2_D, p=dropout_rate, training=self.training)
        E2_Q = nn.functional.dropout(E2_Q, p=dropout_rate, training=self.training)

        # size batch * len1 * len2
        A2 = E2_D.bmm(E2_Q.transpose(2, 1))
        # size batch * len2 * len1
        A_t2 = E2_Q.bmm(E2_D.transpose(2, 1))
        A2.data.masked_fill_(y_mask.data, -float('inf'))
        A_t2.data.masked_fill_(y_mask.data, -float('inf'))
        # Normalize with softmax
        S2_D_alpha_flat = F.softmax(A2.view(-1, E1_Q.size(1)), dim=1)
        S2_D_alpha = S2_D_alpha_flat.view(-1, E1_D.size(1), E1_Q.size(1))
        S2_D = S2_D_alpha.bmm(E2_Q)
        S2_Q_alpha_flat = F.softmax(A_t.view(-1, E2_D.size(1)), dim=1)
        S2_Q_alpha = S2_Q_alpha_flat.view(-1, E2_Q.size(1), E2_D.size(1))
        S2_Q = S2_Q_alpha.bmm(E2_D)
        C2_D = S2_D_alpha.bmm(S2_Q)

        return E2_D, S1_D, S2_D, C1_D, C2_D

class BilinearSeqAttn(nn.Module):
    """A bilinear attention layer over a sequence X w.r.t y:
    * o_i = softmax(x_i'Wy) for x_i in X.

    Optionally don't normalize output weights.
    """
    def __init__(self, x_size, y_size, identity=False):
        super(BilinearSeqAttn, self).__init__()
        if not identity:
            self.linear = nn.Linear(y_size, x_size)
        else:
            self.linear = None

    def forward(self, x, y, x_mask):
        """
        x = batch * len * h1
        y = batch * h2
        x_mask = batch * len
        """
        Wy = self.linear(y) if self.linear is not None else y
        xWy = x.bmm(Wy.unsqueeze(2)).squeeze(2)
        xWy.data.masked_fill_(x_mask.data, -float('inf'))
        if self.training:
            # In training we output log-softmax for NLL
            alpha = F.log_softmax(xWy, dim=1)
        else:
            # ...Otherwise 0-1 probabilities
            alpha = F.softmax(xWy, dim=1)
        return alpha

class LinearSeqAttn(nn.Module):
    """Self attention over a sequence:
    * o_i = softmax(Wx_i) for x_i in X.
    """
    def __init__(self, input_size):
        super(LinearSeqAttn, self).__init__()
        self.linear = nn.Linear(input_size, 1)

    def forward(self, x, x_mask):
        """
        x = batch * len * hdim
        x_mask = batch * len
        """
        x_flat = x.view(-1, x.size(-1))
        scores = self.linear(x_flat).view(x.size(0), x.size(1))
        scores.data.masked_fill_(x_mask.data, -float('inf'))
        alpha = F.softmax(scores, dim=1)
        return alpha

# ------------------------------------------------------------------------------
# Functional
# ------------------------------------------------------------------------------

def uniform_weights(x, x_mask):
    """Return uniform weights over non-masked input."""
    alpha = Variable(torch.ones(x.size(0), x.size(1)))
    if x.data.is_cuda:
        alpha = alpha.cuda()
    alpha = alpha * x_mask.eq(0).float()
    alpha = alpha / alpha.sum(1).expand(alpha.size())
    return alpha

def weighted_avg(x, weights):
    """x = batch * len * d
    weights = batch * len
    """
    return weights.unsqueeze(1).bmm(x).squeeze(1)


#!/usr/bin/env python4
# -*- coding: utf-8 -*-
#  python/test_mouse.py Author "Nathan Wycoff <nathanbrwycoff@gmail.com>" Date 02.14.2020

import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
exec(open("python/lib.py").read())

np.random.seed(123)

# For the  mouse problem, groups of neurons fire to indicate the task parameters. We specify their properties here.
signal_duration = 0.05
break_duration = 0.05
cue_duration = 0.15
spikes_per_signal = 1
neur_per_group = 10
n_signals = 1

R = 2
H = 40
#P = 100
P = 2
Q = 1
N = 100 # Keep even
NN = 100

feed_align = False
#feed_align = True
t_steps = 50
burn_in = t_steps//2

# Variational Inference Params
prior_var = 1e0 * 1/H
#prior_var = 0
data_var = 1e-2#Should be 1 when doing classification.
#data_var = 1e-0#Should be 1 when doing classification.
#V = 10
V = 2
V_test = 10

# Optimizer
#learn_rate = 5e-4
#learn_rate = 0#5e-5
learn_rate = 1e-3
#mb_size = 50
mb_size = 50
#epochs = 500 * mb_size
epochs = 100 * mb_size
# Keep multiple of mb_size
#epochs = 100 * mb_size

#epochs = 5000 * mb_size

## Make problem
#X_vec = np.concatenate([np.random.uniform(low = -1, high = -0.75, size = N//2),  np.random.uniform(low = 0.75, high = 1, size = N//2)])
#x_val = np.random.uniform(low = -1, high = 1, size = N)

x_val = np.random.uniform(low = -1, high = 1, size = N)
X_vec = np.copy(x_val)
X_vec = np.stack([X_vec, np.ones(N)], axis = 1)

#x_test = np.random.uniform(low = -10, high = 10, size = NN)
x_test = np.random.uniform(low = -3, high = 3, size = N)
X_test = np.copy(x_test)
X_test = np.stack([X_test, np.ones(NN)], axis = 1)

#x_val  = np.random.uniform(low = 0, high = 1, size = N)
#x_val = np.random.uniform(low = 0.1, high = 10, size = N)
#X_vec = np.stack([np.sin(np.ceil((p+0.1)/2)*2*np.pi*x_val) if p % 2 == 0 else  np.cos(np.ceil((p+0.1)/2)*2*np.pi*x_val) for p in range(P-1)], axis = 1)
#X_vec = np.append(X_vec, np.ones([N,1]), axis = 1)

#y_vec = np.repeat(0, N)
#y_vec = (x_val > 0)
a = 10
y_vec = np.random.uniform(size=N) < (np.exp(a*x_val) / (np.exp(a*x_val)+1))
#y_vec = np.sqrt(np.cos(2*x_val + np.pi) + 1.2)
#y_vec -= np.mean(y_vec)

t_eps = 0.05
betas = np.array([2 for i in range(H)])

#TODO: No Xavier? Also, make sure it aligns for FA.
sf = 1
LAMBDA_in = np.zeros([H,P]) + np.abs(np.random.normal(size=[H,P], scale = sf*np.sqrt(2/P)))
LAMBDA_rec = np.zeros([H,H]) + np.random.normal(size=[H,H], scale = sf*np.sqrt(2 / H))
LAMBDA_out = np.zeros([Q,H]) + np.random.normal(size=[Q,H], scale = sf*np.sqrt(2/H))

#isig = [1.0,0.001,1.0]
#isig = [1.0,0.1,1.0]
#LAMBDA_in = np.zeros([H,P]) + np.random.normal(size=[H,P], scale = isig[0])
#LAMBDA_rec = np.zeros([H,H]) + np.random.normal(size=[H,H], scale = isig[1])
#LAMBDA_out = np.zeros([Q,H]) + np.random.normal(size=[Q,H], scale = isig[2])

#phi_std = 1e-8
phi_std = prior_var
PHI_in = np.sqrt(prior_var) * np.ones([H,P]) + np.random.normal(size=[H,P], scale = phi_std)
PHI_rec = np.sqrt(prior_var) * np.ones([H,H]) + np.random.normal(size=[H,H], scale = phi_std)
PHI_out = np.sqrt(prior_var) * np.ones([Q,H]) + np.random.normal(size=[Q,H], scale = phi_std)

#RAND_MAT = np.zeros([Q,H]) + np.random.normal(size=[Q,H], scale = isig)
#RAND_MAT = np.zeros([Q,H,V]) + np.random.normal(size=[Q,H,V], scale = stdi)

isig = 1.
RAND_MAT = np.zeros([Q,H,V]) + np.random.normal(size=[Q,H,V], scale = isig)
#isig = 1.
#RAND_MAT = np.zeros([Q,H]) + np.random.normal(size=[Q,H], scale = isig)

#opt = Adam(learn_rate = learn_rate)
opt = GD(learn_rate = learn_rate)
# NOTE: mb_size is set to 1 on purpose; otherwise last_gradients won't be updated every iter.
snn = ALifNeuralNetwork(R, H, P, Q, t_eps = t_eps, mb_size = 1, costfunc = crossentropy, optimizer = opt, thresh = 1)

# Rewrite the L function so that it avoids weight transport using FA.
def fa_L(nn, yt, target_t):
    d = nn.eprop_funcs['dEdy'](nn, yt, target_t)
    if feed_align:
        return (d * RAND_MAT[:,:,nn.RAND_ID]).flatten()
        #return (d * RAND_MAT).flatten()
    else:
        return (d * nn.trainable_params['out']).flatten()

# Overwrite existing L function
snn.eprop_funcs['L'] = fa_L

# Want our own optimizer that is not interferred with anywhere else.
#vb_opt = Adam(learn_rate = learn_rate)
vb_opt = GD(learn_rate = learn_rate)

guess = np.empty(epochs, dtype = float)
truths = np.empty(epochs, dtype = float)
mean_costs = np.zeros(epochs)
mb_costs = np.zeros(epochs//mb_size)
est_probs = np.zeros(epochs) - 1.
LAMI_grad = np.zeros_like(LAMBDA_in)
LAMR_grad = np.zeros_like(LAMBDA_rec)
LAMO_grad = np.zeros_like(LAMBDA_out)

PHII_grad = np.zeros_like(LAMBDA_in)
PHIR_grad = np.zeros_like(LAMBDA_rec)
PHIO_grad = np.zeros_like(LAMBDA_out)
for epoch in tqdm(range(epochs)):
    n_of_epoch = np.random.choice(N)

    y_out = np.empty([V,t_steps])
    costs = np.empty(V)

    # Define inputs for this datum
    def get_xt(self, ip, ti):
        xt = X_vec[n_of_epoch].reshape([P])
        return xt
    target = np.repeat(y_vec[n_of_epoch], t_steps)

    in_params = {'get_xt' : get_xt}

    for v in range(V):
        # Sample our weights from the variational distribution.
        THETA_in = np.random.normal(loc = LAMBDA_in, scale = np.abs(PHI_in))
        THETA_rec = np.random.normal(loc = LAMBDA_rec, scale = np.abs(PHI_rec))
        THETA_out = np.random.normal(loc = LAMBDA_out, scale = np.abs(PHI_out))
        snn.trainable_params = {'in' : THETA_in, 'rec' : THETA_rec, 'out' : THETA_out}

        #TODO: Improve the way states and spikes are reset, as well as RAND_ID
        snn.net_params['spikes'] = []
        snn.RAND_ID = v

        ret = snn.run(t_steps = t_steps, in_params = in_params, in_spikes = None, target = target, train = True, save_states = True, save_traces = True)

        #ret = snn.run(t_steps = t_steps, in_spikes = in_spikes, target = target, train = True, save_states = True, save_traces = True)
        snn.reset_states()

        costs[v] = ret['cost']
        y_out[v,:] = ret['y']

        # Propogate and accumulate gradients to variational parameters.
        # Just variational mean for now.
        LAMI_grad += snn.last_grads[0] / V
        LAMR_grad += snn.last_grads[1] / V
        LAMO_grad += snn.last_grads[2] / V

        PHII_grad += (snn.last_grads[0] * (THETA_in - LAMBDA_in) / PHI_in) / V
        PHIR_grad += (snn.last_grads[1] * (THETA_rec - LAMBDA_rec) / PHI_rec) / V
        PHIO_grad += (snn.last_grads[2] * (THETA_out - LAMBDA_out) / PHI_out) / V

    #NOTE: What if we averaged after the zero comparison? May be more robust.
    #ind_decs = np.mean(y_out[:,int(cue_time / t_eps):], axis = 1) > 0
    ev_y_out = np.mean(y_out, axis = 0)
    e_y_out = np.mean(ev_y_out[burn_in:])

    guess[epoch] = e_y_out
    truths[epoch] = target[-1]

    ## Fold in data variance.
    #TODO: If we re-use this block, we may need to account for mb_size.
    #LAMI_grad *= epochs / data_var
    #LAMR_grad *= epochs / data_var
    #LAMO_grad *= epochs / data_var
    #PHII_grad *= epochs / data_var
    #PHIR_grad *= epochs / data_var
    #PHIO_grad *= epochs / data_var

    # Add in prior grads
    if (epoch+1) % mb_size == 0:
        # Because this code is being ran epoch / mb_size times, we need to account for that
        const = (N / mb_size) / data_var
        constinv = 1/const
        LAMI_grad += constinv * (LAMBDA_in / np.square(PHI_in))
        LAMR_grad += constinv * (LAMBDA_rec / np.square(PHI_rec))
        LAMO_grad += constinv * (LAMBDA_out / np.square(PHI_out))
        PHII_grad += constinv * (PHI_in * (1/prior_var - 1/np.square(PHI_in)))
        PHIR_grad += constinv * (PHI_rec * (1/prior_var - 1/np.square(PHI_rec)))
        PHIO_grad += constinv * (PHI_out * (1/prior_var - 1/np.square(PHI_out)))

        # Feed the accumulated gradients to our optimizer.
        cur_LAMBDAs = [LAMBDA_in, LAMBDA_rec, LAMBDA_out]
        cur_PHIs = [PHI_in, PHI_rec, PHI_out]
        cur_params = cur_LAMBDAs + cur_PHIs
        grads_LAMBDAs = [LAMI_grad, LAMR_grad, LAMO_grad]
        grads_PHIs = [LAMI_grad, LAMR_grad, LAMO_grad]
        grads_all = grads_LAMBDAs + grads_PHIs
        #TODO: The use of vb_opt here is due to the fact that I am updating the weights in the snn as well as out here. A short terms solution would be to have a flag that allows us to compute gradients and a separate flag that allows us to do the updating in the snn.run call; right now "train" does both or neither of these.
        LAMBDA_in, LAMBDA_rec, LAMBDA_out, PHI_in, PHI_rec, PHI_out = vb_opt.apply_gradients(cur_params, grads_all)
        #LAMBDA_in, LAMBDA_rec, LAMBDA_out = vb_opt.apply_gradients(cur_LAMBDAs, grads_LAMBDAs)
        LAMI_grad = np.zeros_like(LAMBDA_in)
        LAMR_grad = np.zeros_like(LAMBDA_rec)
        LAMO_grad = np.zeros_like(LAMBDA_out)

        PHII_grad = np.zeros_like(LAMBDA_in)
        PHIR_grad = np.zeros_like(LAMBDA_rec)
        PHIO_grad = np.zeros_like(LAMBDA_out)

        mb_costs[epoch // mb_size] = np.mean(mean_costs[(epoch-mb_size+1):(epoch+1)])

    mean_costs[epoch] = np.mean(costs)

model_preds = np.empty(N)
n_of_epoch = np.random.choice(N)

y_out = np.empty([V,t_steps])
costs = np.empty(V)

THETA_in_draws = [np.random.normal(loc = LAMBDA_in, scale = np.abs(PHI_in)) for _ in range(V_test)]
THETA_rec_draws  = [np.random.normal(loc = LAMBDA_rec, scale = np.abs(PHI_rec)) for _ in range(V_test)]
THETA_out_draws  = [np.random.normal(loc = LAMBDA_out, scale = np.abs(PHI_out)) for _ in range(V_test)]

yv = np.empty([V_test,NN])
for n in tqdm(range(NN)):
    # Define inputs for this datum
    def get_xt(self, ip, ti):
        xt = X_test[n].reshape([P])
        return xt
    target = np.repeat(y_vec[n_of_epoch], t_steps)

    in_params = {'get_xt' : get_xt}

    for v in range(V_test):
        # Sample our weights from the variational distribution.
        THETA_in = THETA_in_draws[v]
        THETA_rec = THETA_rec_draws[v]
        THETA_out = THETA_out_draws[v]
        snn.trainable_params = {'in' : THETA_in, 'rec' : THETA_rec, 'out' : THETA_out}

        #TODO: Improve the way states and spikes are reset, as well as RAND_ID
        snn.net_params['spikes'] = []
        snn.RAND_ID = v

        ret = snn.run(t_steps = t_steps, in_params = in_params, in_spikes = None, target = target, train = True, save_states = True, save_traces = True)

        yv[v,n] = np.mean(ret['y'][:,burn_in:])

fig = plt.figure()
plt.subplot(2,1,1)
plt.scatter(x_val, y_vec)
for v in range(V_test):
    plt.scatter(x_test, yv[v,:], color = 'orange', alpha = 0.5)
plt.title("Data")
plt.subplot(2,1,2)
plt.title("Data")
plt.plot(mb_costs)
plt.savefig("temp.pdf")
plt.close()


#!/usr/bin/python3
import matplotlib.pyplot as plt
import numpy as np
import sys

if (len(sys.argv) != 2):
    sys.exit("Expect 1 argument: Usage udc_process.py <log file>")

x, y1, y2, y3, y4, y5 = np.genfromtxt(sys.argv[1], delimiter=',', unpack=True)
#plt.plot(x,y1, label='Recall 1 in 10')

#approach1 - just plot one at a time
approach1 = False

if approach1:
    plt.xlabel('Number of Epochs')
    plt.ylabel('Recall@1')
    plt.title('Recall@1')
    plt.grid()
    plt.legend()
    plt.show()
else:
    fig, ax = plt.subplots()
    plt.grid()
    plt.xlabel('Number of Epochs')
    plt.ylabel('Recall Value')

    #Limit Range of x axis values
    ax.set_xlim(right=100000)

    #Plot the recall values for @1, @2, @5 and loss with epochs.
    ax.plot(x, y1, 'k--', color='green', label='Recall@1')
    ax.plot(x, y2, 'k:', color='blue', label='Recall@2')
    ax.plot(x, y3, 'k', color='orange', label='Recall@5')
    ax.plot(x, y5*0.25, 'k', color='black' , label='Loss')

    legend = ax.legend(loc='lower right', shadow=True, fontsize='large')

    # Put a nicer background color on the legend.
    #legend.get_frame().set_facecolor('C0')
    plt.show()


from django.urls import path
from .views import (employee,
                     department,
                     company,
                     show,
                     edit,
                     updateEmployee,
                     destroyEmployee,
                     destroyDepartment,
                     destroyCompany,
                     updateDepartment,
                     updateCompany,
                     home)

app_name = 'core'

urlpatterns = [
   path('',home),
   path('company', company),
   path('company/employee', employee),
   path('company/department', department),
   path('show', show),
   path('edit/<id>/', edit),
   path('updateEmployee/<id>', updateEmployee),
   path('updateDepart/<id>', updateDepartment),
   path('updateCompany/<id>', updateCompany),
   path('deleteemp/<id>', destroyEmployee),
   path('deletedep/<id>',destroyDepartment),
   path('deletecomp/<id>',destroyCompany)

   # path('process-list-full', ProcessList.as_view(), name='process-list-full'),
   # path('process-list-partner', ProcessListPartner.as_view(), name='process-list-partner'),
   # path('process-list-owner', ProcessListOwner.as_view(), name='process-list-owner'),
   # path('process-create', ProcessCreate.as_view(), name='process-create'),
   # path('process-detail/<uuid:pk>', ProcessDetail.as_view(), name='process-detail'),
   # path('process-update/<uuid:pk>', ProcessUpdate.as_view(), name='process-update'),
   # path('process-delete/<uuid:pk>', ProcessDelete.as_view(), name='process-delete'),

]


import re

import pytest
from django.urls import include, path

from drf_spectacular.generators import SchemaGenerator
from tests import assert_schema

transforms = [
    # User model first_name differences
    lambda x: re.sub(r'(first_name:\n *type: string\n *maxLength:) 150', r'\g<1> 30', x, re.M),
    # User model username validator tail
    lambda x: x.replace(r'+$', r'+\Z'),
]

@pytest.mark.contrib('dj_rest_auth', 'allauth')
def test_rest_auth(no_warnings):
    urlpatterns = [
        path('rest-auth/', include('dj_rest_auth.urls')),
        path('rest-auth/registration/', include('dj_rest_auth.registration.urls')),
    ]

    generator = SchemaGenerator(patterns=urlpatterns)
    schema = generator.get_schema(request=None, public=True)

    assert_schema(
        schema, 'tests/contrib/test_rest_auth.yml', transforms=transforms
    )


import numpy as np
import cv2
import mahotas

class Preprocessor:
    def __init__(self, image_path):
        # 读取图像并截取
        self.image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), -1)

    def find_roi(self):
        # 转灰度图
        if len(self.image.shape)>2:
            gray = cv2.cvtColor(self.image, cv2.COLOR_BGR2GRAY)
        else:
            gray = self.image

        # 去噪
        blurred = cv2.GaussianBlur(gray, (7, 7), 0)
        blurred = cv2.bilateralFilter(blurred, 17, 16, 32)
        # cv2.imshow('Blurred', blurred)

        # 二值化
        T = mahotas.thresholding.otsu(blurred)
        # print("Otsu’s threshold: {}".format(T))
        thresh = blurred.copy()
        thresh[thresh > T] = 255
        thresh[thresh < 255] = 0
        cv2.imshow('Thresh', thresh)

        # 边缘检测
        edged = cv2.Canny(thresh, 30, 150)
        (_, cnts, _) = cv2.findContours(edged.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)

        # 轮廓去噪，去除噪声轮廓
        '''有bug，如果寻找所有轮廓，则contour不连续，外接框无法框定最大范围
        如果只保留最大轮廓，contour又不完整，外接框只框定部分范围'''
        contour = []
        for each in cnts:
            # 去除长度小于10的轮廓，看做噪声
            if each.shape[0] > 10:
                contour.append(each)

        # 显示轮廓
        chara = self.image.copy()
        cv2.drawContours(chara, contour, -1, (0, 255, 0), 2)
        cv2.imshow("Contour", chara)
        # cv2.waitKey(0)

        # 轮廓的外接矩形框
        (x, y, w, h) = cv2.boundingRect(np.array(contour))

        # 显示原图中的字符区域
        ROI = thresh[y:y + h, x:x + w]
        cv2.imshow("ROI", ROI)
        cv2.waitKey(0)

        return ROI


import datetime
#特定の日付からオリンピックまでの日数を調べる
def calc_days(y,m,d):
    olympic=datetime.datetime(2021,7,23).timestamp()
    target=datetime.datetime(y,m,d).timestamp()
    #1日あたりの秒数
    perday=24*60*60
    days=int(olympic-target)//perday
    s=f'{y}/{m}/{d}から{days}日後'
    print(s)
calc_days(2018,12,1)
t=datetime.date.today()
calc_days(t.year,t.month,t.day)


import db_connection
import configparser

import build_query
from setup import super_logger

config = configparser.ConfigParser()
config.read('config.ini')

credentials = '/Volumes/Data/sling-stream-recover-eb60c169ef8a.json'
projectid = 'sling-stream-recover'

reportingdate = config.get('reporting_date', 'reportingdate')

# Schema
bq_project = config.get('bigquery', 'project_id')

bqschema = config['schema']['bqschema']
redshiftbq = config['schema']['redshiftschema']

# Dany delete
credentials = '/Volumes/Data/sling-stream-recover-eb60c169ef8a.json'
projectid = 'sling-stream-recover'

bqconn = db_connection.big_query(credentials, projectid)

def delete_data(uuid):
    lines = []
    queries = []
    with open("bqtables.txt") as file_in:

        for line in file_in:
            lines.append(line.strip())

    with open('lisquery.txt') as file:
        for quer in file:
            queries.append(quer.strip())

        for i in lines:
            super_logger.info(f'*****************{i}**********************')
            print('*****************{i}**********************')

            for j in queries:
                selectstatement = ' '.join(j.split()[1:2])
                query = str(j).format(bqschema.format(str(i)),
                                      "'" + uuid + "'") + ';'

                super_logger.info(f'BigQuery:{query}:')
                print('BigQuery:{query}:')

                # Get the result from Bq
                BQ_query = build_query.selectquery_bq(query, bqconn)
                result_BQ = BQ_query
                super_logger.info(f'BigQuery:{selectstatement}:{result_BQ}')
                print('BigQuery:{selectstatement}:{result_BQ}')

if __name__ == '__main__':
    delete_data("uuid")


"""
Implement some utils used to convert image and it's corresponding label into tfrecords
"""
import numpy as np
import tensorflow as tf
import os
import re
from data_prepare import char_dict
import config

def sparse_tensor_to_str(spares_tensor: tf.SparseTensor):
    """
    :param spares_tensor:
    :return: a str
    """
    indices = spares_tensor.indices
    values = spares_tensor.values
    # values = np.array([self.__ord_2_index_map[str(tmp)] for tmp in values])
    dense_shape = spares_tensor.dense_shape

    number_lists = np.ones(dense_shape, dtype=values.dtype) * (-1)
    char_lists = []
    res = []
    for i, index in enumerate(indices):
        number_lists[index[0], index[1]] = values[i]
    print(number_lists)
    for number_list in number_lists:
        char_lists.append([char_dict.int_to_char[val] for val in number_list if val >= 0])
    for char_list in char_lists:
        res.append(''.join(char_list))
    return res

def read_features(tfrecords_dir, num_epochs, flag):
    """
    :param tfrecords_dir:
    :param num_epochs:
    :param flag: 'Train', 'Test', 'Validation'
    :return:
    """

    assert os.path.exists(tfrecords_dir)

    if not isinstance(flag, str):
        raise ValueError('flag should be a str in [\'Train\', \'Test\', \'Val\']')
    if flag.lower() not in ['train', 'test', 'val']:
        raise ValueError('flag should be a str in [\'Train\', \'Test\', \'Val\']')

    if flag.lower() == 'train':
        re_patten = r'^train_feature_\d{0,15}_\d{0,15}\.tfrecords\Z'
    elif flag.lower() == 'test':
        re_patten = r'^test_feature_\d{0,15}_\d{0,15}\.tfrecords\Z'
    else:
        re_patten = r'^val_feature_\d{0,15}_\d{0,15}\.tfrecords\Z'

    tfrecords_list = [os.path.join(tfrecords_dir, tmp) for tmp in os.listdir(tfrecords_dir) if re.match(re_patten, tmp)]

    print('tfrecords_list: ', tfrecords_list)

    filename_queue = tf.train.string_input_producer(tfrecords_list)
    reader = tf.TFRecordReader()
    _, serialized_example = reader.read(filename_queue)

    features = tf.parse_single_example(serialized_example,
                                       features={
                                           'images': tf.FixedLenFeature((), tf.string),
                                           'imagenames': tf.FixedLenFeature([1], tf.string),
                                           'labels': tf.VarLenFeature(tf.int64),
                                           # 'labels': tf.FixedLenFeature([], tf.int64),
                                       })
    image = tf.decode_raw(features['images'], tf.uint8)
    images = tf.reshape(image, [config.img_h, config.img_w, config.img_ch])
    labels = features['labels']
    # labels = tf.one_hot(indices=labels, depth=10) #config.cfg.TRAIN.CLASSES_NUMS)
    labels = tf.cast(labels, tf.int32)
    # labels = tf.keras.utils.to_categorical(labels, num_classes=10)
    imagenames = features['imagenames']
    return images, labels, imagenames

def cnn_net(inputs):
    # cnn -----------------------------------------
    conv1 = tf.layers.conv2d(inputs=inputs,
                             filters=128,
                             kernel_size=(3, 3),
                             padding='same',
                             # kernel_initializer='he_normal',
                             name='conv1')
    bn1 = tf.layers.batch_normalization(inputs=conv1,
                                        name='bn1')
    pool1 = tf.layers.max_pooling2d(inputs=bn1,
                                    pool_size=(2, 2),
                                    strides=(2, 2),
                                    name='pool1')

    pool2 = tf.layers.max_pooling2d(inputs=pool1,
                                    pool_size=(1, 2),
                                    strides=(1, 2),
                                    name='pool1')  # (bs, 16, 32, 128)
    # map to sequence -----------------------------------
    act1 = tf.nn.relu(features=pool2,
                      name='act1')

    tsp1 = tf.transpose(a=act1, perm=(0, 2, 1, 3))  # (bs, 32, 16, 128)

    rs1 = tf.reshape(tensor=tsp1, shape=(-1, 32, 2048), name='reshape1')  # (bs, 25, 2048)

    # rnn -----------------------------------------------
    # fw_cell_list = [tf.contrib.rnn.BasicLSTMCell(num_units=n_hidden, forget_bias=1.0) for n_hidden in list_n_hidden]
    # bw_cell_list = [tf.contrib.rnn.BasicLSTMCell(num_units=n_hidden, forget_bias=1.0) for n_hidden in list_n_hidden]
    # rnn_out, _, _ = tf.contrib.rnn.static_bidirectional_rnn(cell_fw=fw_cell_list,
    #                                           cell_bw=bw_cell_list,
    #                                           inputs=rs1,
    #                                           dtype=tf.float32)
    n_hidden = 256
    # first bi_rnn layer
    fw_cell1 = tf.contrib.rnn.BasicLSTMCell(num_units=n_hidden, forget_bias=1.0)
    bw_cell1 = tf.contrib.rnn.BasicLSTMCell(num_units=n_hidden, forget_bias=1.0)
    rnn_out1, state1 = tf.nn.bidirectional_dynamic_rnn(cell_fw=fw_cell1,
                                                       cell_bw=bw_cell1,
                                                       inputs=rs1,
                                                       dtype=tf.float32,
                                                       scope='bi_rnn1')
    # fw_out, bw_out = rnn_out
    concat1 = tf.concat(values=rnn_out1, axis=2, name='concat1')  # (bs, time_steps, 255+255)

    # second bi_rnn layer
    fw_cell2 = tf.contrib.rnn.BasicLSTMCell(num_units=n_hidden, forget_bias=1.0)
    bw_cell2 = tf.contrib.rnn.BasicLSTMCell(num_units=n_hidden, forget_bias=1.0)
    rnn_out2, state2 = tf.nn.bidirectional_dynamic_rnn(cell_fw=fw_cell2,
                                                       cell_bw=bw_cell2,
                                                       inputs=concat1,
                                                       dtype=tf.float32,
                                                       scope='bi_rnn2')
    # fw_out, bw_out = rnn_out
    concat2 = tf.concat(values=rnn_out2, axis=2, name='concat2')  # (bs, time_steps, 255+255)

    # TODO: add dropout layer

    # fc -------------------------------
    shape = tf.shape(concat2)
    fc_input = tf.reshape(tensor=concat2, shape=(shape[0] * shape[1], shape[2]))  # (bs*time_steps, 255+255)

    num_classes = 6764
    fc_output = tf.layers.dense(inputs=fc_input,
                                units=num_classes,
                                name='fc')
    rs_fc_output = tf.reshape(tensor=fc_output,
                              shape=(shape[0], shape[1], -1))  # (bs, time_steps, num_classes)

    act_output = tf.nn.softmax(logits=rs_fc_output, axis=2, name='activate')

    return act_output

if __name__ == '__main__':
    imgs, labels, img_names = read_features('/data/data/crnn_tfrecords',
                                            num_epochs=None,
                                            flag='Train')

    inputdata, input_labels, input_imagenames = tf.train.shuffle_batch(tensors=[imgs, labels, img_names],
                                                                       batch_size=4,
                                                                       capacity=32 + 4 * 4,
                                                                       min_after_dequeue=32,
                                                                       num_threads=1)
    input_data = tf.cast(inputdata, dtype=tf.float32)
    net_output = cnn_net(input_data)

    sess = tf.Session()

    init = tf.global_variables_initializer()
    sess.run(init)

    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)

    with sess.as_default():
        _net_output = sess.run([net_output])
        print(_net_output[0].shape)

        # imgs_val, labels_val, img_names_val = sess.run([inputdata, input_labels, input_imagenames])
        # print(type(labels_val))
        # print(labels_val)
        # print(img_names_val)
        # scene_labels = sparse_tensor_to_str(labels_val)
        # print('scene_labels_len: ', len(scene_labels))
        # for index, scene_label in enumerate(scene_labels):
        #     img_name = img_names_val[index][0].decode('utf-8')
        #     print('{:s} --- {:s}'.format(img_name, scene_label))

        coord.request_stop()
        coord.join(threads=threads)


# -*- coding: utf-8 -*-
# Copyright 2019 Huairuo.ai.
# Author:
#
# cython: language_level=3
#

import json
import re
import os
import numpy as np
import torch
from collections import Counter
import lib.config as config
from lib.CMRC_QADataSet import CMRC_QADataSet
from lib.QADataSet import QADataSet
from lib.config import logger
import importlib
from my_py_toolkit.file.file_toolkit import *
from lib.utils import *
from my_py_toolkit.decorator.decorator import fn_timer
from pytorch_transformers import BertModel
from tqdm import tqdm
import collections
from lib.tokenizations import official_tokenization as tokenization

SPIECE_UNDERLINE = '▁'

def evaluate_valid_result(valid_result, exact_match_total=0, f1_total=0, total=0):
  # f1 = exact_match = total = 0
  for item in valid_result:
    total += 1
    ground_truths = item.get("label_answer")
    prediction = item.get("predict_answer")
    exact_match_total += metric_max_over_ground_truths(exact_match_score, prediction,
                                                       ground_truths)
    f1_total += metric_max_over_ground_truths(f1_score, prediction, ground_truths)
  exact_match = 100.0 * exact_match_total / total
  f1 = 100.0 * f1_total / total
  return exact_match_total, f1_total, exact_match, f1

def convert_tokens(eval_file, qa_id, pp1, pp2):
  answer_dict = {}
  remapped_dict = {}
  for qid, p1, p2 in zip(qa_id, pp1, pp2):
    context = eval_file[str(qid)]["context"]
    spans = eval_file[str(qid)]["spans"]
    uuid = eval_file[str(qid)]["uuid"]
    l = len(spans)
    if p1 >= l or p2 >= l:
      ans = ""
    else:
      start_idx = spans[p1][0]
      end_idx = spans[p2][1]
      ans = context[start_idx: end_idx]
    answer_dict[str(qid)] = ans
    remapped_dict[uuid] = ans
  return answer_dict, remapped_dict

def evaluate(eval_file, answer_dict):
  f1 = exact_match = total = 0
  for key, value in answer_dict.items():
    total += 1
    ground_truths = eval_file[key]["answers"]
    prediction = value
    exact_match += metric_max_over_ground_truths(exact_match_score, prediction,
                                                 ground_truths)
    f1 += metric_max_over_ground_truths(f1_score, prediction, ground_truths)
  exact_match = 100.0 * exact_match / total
  f1 = 100.0 * f1 / total
  return {'exact_match': exact_match, 'f1': f1}

def normalize_answer(s):
  def remove_articles(text):
    return re.sub(r'\b(a|an|the)\b', ' ', text)

  def white_space_fix(text):
    return ' '.join(text.split())

  def remove_punc(text):
    exclude = set(string.punctuation)
    return ''.join(ch for ch in text if ch not in exclude)

  def lower(text):
    return text.lower()

  return white_space_fix(remove_articles(remove_punc(lower(s))))

def f1_score(prediction, ground_truth):
  """

  Args:
    prediction(str):
    ground_truth(str):

  Returns:

  """
  # prediction_tokens = normalize_answer(prediction).split()
  # ground_truth_tokens = normalize_answer(ground_truth).split()
  prediction_tokens = list(prediction)
  ground_truth_tokens = list(ground_truth)
  common = Counter(prediction_tokens) & Counter(ground_truth_tokens)
  num_same = sum(common.values())
  if num_same == 0:
    return 0
  precision = 1.0 * num_same / len(prediction_tokens)
  recall = 1.0 * num_same / len(ground_truth_tokens)
  f1 = (2 * precision * recall) / (precision + recall)
  return f1

def exact_match_score(prediction, ground_truth):
  # return (normalize_answer(prediction) == normalize_answer(ground_truth))
  return prediction == ground_truth

def metric_max_over_ground_truths(metric_fn, prediction, ground_truths):
  """
  计算准确率.
  Args:
    metric_fn:
    prediction(str):
    ground_truths(str):

  Returns:

  """
  # todo: 代码有问题需要修改
  # scores_for_ground_truths = []
  # for ground_truth in ground_truths:
  #   score = metric_fn(prediction, ground_truth)
  #   scores_for_ground_truths.append(score)
  # return max(scores_for_ground_truths)
  return metric_fn(prediction, ground_truths)

def convert_valid_result(Cwids, Qwids, y1s, y2s, p1s, p2s, dataset, ids):
  """"""
  result = []
  for Cwid, Qwid, y1, y2, p1, p2, id in zip(Cwids, Qwids, y1s, y2s, p1s, p2s, ids):
    Cw_text = dataset.get_origin_data(id, "context")
    Qw_text = dataset.get_origin_data(id, "question")
    Aw_text = dataset.get_origin_data(id, "answer_text")
    # y1, y2 = int(y1), int(y2)
    p1, p2 = int(p1), int(p2)
    result.append({
      "context": Cw_text,
      "question": Qw_text,
      "labelled_answer": Aw_text,
      "predict_answer": Cw_text[p1: p2]
    })
  return result

def convert_valid_result_baseline(y1s, y2s, p1s, p2s, dataset, ids):
  """"""
  result = []
  for y1, y2, p1, p2, id in zip(y1s, y2s, p1s, p2s, ids):
    input_feature = dataset.input_features(id)
    # Cw_text = dataset.get_origin_data(id, "context")
    # Qw_text = dataset.get_origin_data(id, "question")
    # Aw_text = dataset.get_origin_data(id, "answer_text")
    # y1, y2 = int(y1), int(y2)
    p1, p2 = int(p1), int(p2)
    result.append({
      "context": Cw_text,
      "question": Qw_text,
      "labelled_answer": Aw_text,
      "predict_answer": Cw_text[p1: p2]
    })
  return result

def load_model(model_dir, check_point):
  """

  Returns:

  """
  model_path = os.path.join(model_dir, f"model_{check_point}.pt")
  return torch.load(model_path)

def load_bert(model_path, device, use_pretrained_bert=True, bert_conf=None,
              use_segments_embeddings=True, local_bert_class=None):
  """
  加载与训练的 bert.
  Args:
    model_path:
    device:

  Returns:

  """
  if use_pretrained_bert:
    return BertModel.from_pretrained(model_path).to(device)
  else:
    return local_bert_class(bert_conf, use_segments_embeddings).to(device)

@fn_timer(logger)
def get_model(package, name, class_name):
  """
  Gets models.
  Returns:

  """
  model_class = load_class(package, name, class_name)
  if not config.is_continue:
    return model_class(config.bert_path, config.device, config.dropout)
  else:
    logger.info(f"Continue train, continue_checkpoint: {config.continue_checkpoint}")
    model_path = os.path.join(config.model_dir,
                              f"model_{config.continue_checkpoint}.pkl")
    if not config.is_only_save_params:
      model = torch.load(model_path, map_location=config.device)
    else:
      model = model_class(config.bert_path, config.device, config.dropout)
      model.load_state_dict(torch.load(model_path))
    return model

@fn_timer(logger)
def get_optimizer(base_lr, params):
  """
  Gets models.
  Returns:

  """
  if not config.is_continue:
    return torch.optim.Adam(lr=base_lr, betas=(config.beta1, config.beta2),
                         eps=1e-8, weight_decay=3e-7, params=params)
  else:
    logger.info(f"Continue train, continue_checkpoint: {config.continue_checkpoint}")
    optimizer_path = os.path.join(config.model_dir,
                              f"optimizer_{config.continue_checkpoint}.pkl")
    if not config.is_only_save_params:
      optimizer = torch.load(optimizer_path, map_location=config.device)
    else:
      optimizer = torch.optim.Adam(lr=base_lr, betas=(config.beta1, config.beta2),
                         eps=1e-8, weight_decay=3e-7, params=params)
      optimizer.load_state_dict(torch.load(optimizer_path))
    return optimizer

def save_model(model, optmizer, steps=0):
  model_path = os.path.join(config.model_dir, f"model_{str(steps)}.pkl")
  optimizer_path = os.path.join(config.model_dir, f"optimizer_{str(steps)}.pkl")
  make_path_legal(model_path)
  make_path_legal(optimizer_path)
  if not config.is_only_save_params:
    torch.save(model, model_path)
    torch.save(optmizer, optimizer_path)
    # torch.save(model.embedding.trainable_embedding,
    #            config.embedding_trainable_model)
    # torch.save(data_set.trainable_embedding, config.embedding_trainable_model)
  else:
    torch.save(model.state_dict(), model_path)
    torch.save(optmizer.state_dict(), optimizer_path)
    # torch.save(model.embedding.trainable_embedding,
    #            config.embedding_trainable_model)
    # torch.save(data_set.trainable_embedding, config.embedding_trainable_model)

@fn_timer(logger)
def get_dataset(data_type="train", mode="train"):
  """

  Args:
    data_type(str): train, dev, trial

  Returns:

  """
  data_dir = config.dataset_dir.get(config.dataset_name)
  data_path = [path for path in get_file_paths(data_dir) if data_type in get_file_name(path)][0]
  if "cmcr" in config.dataset_name:
    dataset = CMRC_QADataSet(data_path,
                             config.is_train, mode)
    return dataset
  else:
    dataset = QADataSet(data_path,
                        config.batch_size, mode)
    return dataset

def record_info(losses=[], f1=[], em=[], valid_result={}, epoch=0,
                r_type="train", is_continue=True):
  """
  记录训练中的 loss, f1, em 值.
  Args:
    losses:
    f1:
    em:
    r_type:

  Returns:

  """
  dir_name = f"./log/{r_type}/"
  losses = [str(v) for v in losses]
  f1 = [str(v) for v in f1]
  em = [str(v) for v in em]
  if losses:
    write2file(",".join(losses), f"{dir_name}losses.txt", is_continue=is_continue)
  if f1:
    write2file(",".join(f1), f"{dir_name}f1.txt", is_continue=is_continue)
  if em:
    write2file(",".join(em), f"{dir_name}em.txt", is_continue=is_continue)

  if valid_result:
    write2file(json.dumps(valid_result, ensure_ascii=False, indent=2),
               f"{dir_name}valid_result_{epoch}.json", is_continue=is_continue)

def corresponds_index(origin, token):
  """
  将 BertTokenizer.tokenize 出的 token 与原始字符串对应回去，
  主要处理情况：对英文切割不好(eg: luminoso 被拆分为：'lu', '##min', '##os', '##o',)，
  在使用了原始句子下标(index)的情况下，
  BertTokenizer.tokenize 后结果 index 与原句对应不上， index 需要纠正。
  Args:
    origin(list):
    token(list):

  Returns:

  """
  origin_index_in_token = []
  index_in_token = 0
  token_len = len(token)
  for index, char in enumerate(origin):
    if index_in_token > token_len - 1 or re.search("[a-zA-Z\d]", char):
      origin_index_in_token.append(-1)
      continue

    for j in range(index_in_token, token_len):
      if char == token[j]:
        origin_index_in_token.append(j)
        index_in_token = j
        break

    while len(origin_index_in_token) < index + 1:
      origin_index_in_token.append(-1)

  return origin_index_in_token

def transfer_index(origin_content, tokenize_content, *indexes):
  """
  将 origin content 中的 index 映射到 tokenize_content 中的 Index.
  TODO: 目前英文单词后面紧结数字，返回 index 不对, 待修复。
  Args:
    origin_content:
    tokenize_content:
    *indexes:

  Returns:

  """
  text_index_in_token = corresponds_index(origin_content, tokenize_content)
  result = []
  for i in indexes:
    if text_index_in_token[i] > -1:
      result.append(text_index_in_token[i])
    else:
      first_non_negative_index = get_first_non_negative_index(text_index_in_token[:i], reverse=True)
      if first_non_negative_index > -1:
        result.append(text_index_in_token[first_non_negative_index] + 1)
      else:
        result.append(-1)
  return result

def get_trainable_embedding(vocab_size):
  """
  获取可训练 embedding.
  Returns:

  """
  embedding = torch.Tensor(vocab_size, config.embedding_trainable_dim)
  embedding = torch.nn.Parameter(torch.nn.init.normal_(embedding).to(config.device))
  return embedding
  # if not config.is_continue or not os.path.exists(config.embedding_trainable_model):
  #   # torch.nn.init.normal_(torch.Tensor(3,3)).requires_grad_(True)
  #   embedding = torch.Tensor(vocab_size, config.embedding_trainable_dim)
  #   embedding = torch.nn.Parameter(torch.nn.init.normal_(embedding)).to(config.device).requires_grad_(True)
  #   return embedding
  # else:
  #   return torch.load(config.embedding_trainable_model).to(config.device).requires_grad_(True)

def adapt_conv2D(input_tensor, origin_shape):
  """
  将输入转换为适合二维 conv （卷积）的维度.
  词向量输入： batch_size, length, word_dim, 二维卷积需要的维度: batch_size, length, word_dim, channel,
  将输入增加一维。
  Args:
    input_tensor:

  Returns:

  """
  origin_shape = list(origin_shape)
  if len(origin_shape) == 3:
    origin_shape.append(1)
    input_tensor = input_tensor.contiguous().view(*origin_shape)

  return input_tensor

def adapt_normal(input_tensor, origin_shape):
  """
  调整 normal 操作所需的 input tensor shape.
  Args:
    input_tensor:

  Returns:

  """
  origin_shape = list(origin_shape)
  input_tensor = input_tensor.contiguous().view(origin_shape[0], -1)

  return input_tensor

def adapt_line(input_tensor, origin_shape):
  """
  调整 normal 操作所需的 input tensor shape.
  Args:
    input_tensor:

  Returns:

  """
  origin_shape = list(origin_shape)
  new_shape = origin_shape
  if new_shape[-1] != 1:
    new_shape.append(1)
  input_tensor = input_tensor.contiguous().view(*new_shape)

  return input_tensor

def get_steps(data_set_type, mode="train"):
  """

  Args:
    data_set_type:
    mode:train, debug

  Returns:

  """
  if mode == "debug":
    return 1

  if data_set_type == "train":
    return config.num_steps

  if data_set_type == "valid":
    return config.val_num_steps

  if data_set_type == "test":
    return config.test_num_steps

def get_vocab_size(vocab_file):
  with open(vocab_file, "r", encoding="utf-8") as f:
    context = f.read()
    return len(context.split("\n"))

def is_leaf(value):
  keys = "".join([k for k,_ in value.items()])
  if "." in keys:
    return False
  else:
    return True

def process_leaf(value):
  value = [v for k,v in value.items()]
  value = [reshape_tensor(v,[-1]) for v in value if v is not None]
  if value:
    value = torch.cat(value, dim=0)
  return value

def handler_gradient(gradient):
  result ={}
  leaf = True
  for k,v in gradient.items():
    if not isinstance(v, dict):
      result[k] = v
      continue

    if isinstance(list(v.values())[0], dict):
      v = handler_gradient(v)
      leaf = False

    if leaf and is_leaf(v):
      v = process_leaf(v)
      result[k] = v
    else:
      for sub_key, sub_v in v.items():
        result[f"{k}.{sub_key}"] = sub_v

  return result

def is_last_layer(value):
  if not isinstance(value, dict):
    return True
  for k,v in value.items():
    if isinstance(v, dict):
      return False
  return True

def transfer_multi_layer_dict(dict_value):
  """
  将多层 dict 转换为 1 层 dict.
  Args:
    dict_value:

  Returns:

  """
  result = {}
  for key, value in dict_value.items():
    if not isinstance(value, dict):
      result[key] = value
      continue

    if not is_last_layer(value):
      value = transfer_multi_layer_dict(value)

    for sub_key, sub_value in value.items():
      result[f"{key}.{sub_key}"] = sub_value

  return result

def covert_position_embeddings(position_embeddings, predict, log_softmax):
  for index, v in enumerate(predict):
    position_embeddings[index][v] += 1
  position_embeddings = log_softmax(log_softmax)
  return position_embeddings

def convert_one_hot(input_array, class_num):
  result = torch.zeros((len(input_array), class_num))
  for index, v in enumerate(input_array):
    result[index][v] = 1
  return result

def truncated_normal_(tensor, mean=0, std=0.09):
  with torch.no_grad():
    size = tensor.shape
    tmp = tensor.new_empty(size + (4,)).normal_()
    valid = (tmp < 2) & (tmp > -2)
    ind = valid.max(-1, keepdim=True)[1]
    tensor.data.copy_(tmp.gather(-1, ind).squeeze(-1))
    tensor.data.mul_(std).add_(mean)
    return tensor

def _improve_answer_span(doc_tokens, input_start, input_end, tokenizer,
                         orig_answer_text):
  """Returns tokenized answer spans that better match the annotated answer."""

  # The SQuAD annotations are character based. We first project them to
  # whitespace-tokenized words. But then after WordPiece tokenization, we can
  # often find a "better match". For example:
  #
  #   Question: What year was John Smith born?
  #   Context: The leader was John Smith (1895-1943).
  #   Answer: 1895
  #
  # The original whitespace-tokenized answer will be "(1895-1943).". However
  # after tokenization, our tokens will be "( 1895 - 1943 ) .". So we can match
  # the exact answer, 1895.
  #
  # However, this is not always possible. Consider the following:
  #
  #   Question: What country is the top exporter of electornics?
  #   Context: The Japanese electronics industry is the lagest in the world.
  #   Answer: Japan
  #
  # In this case, the annotator chose "Japan" as a character sub-span of
  # the word "Japanese". Since our WordPiece tokenizer does not split
  # "Japanese", we just use "Japanese" as the annotation. This is fairly rare
  # in SQuAD, but does happen.
  tok_answer_text = " ".join(tokenizer.tokenize(orig_answer_text))

  for new_start in range(input_start, input_end + 1):
    for new_end in range(input_end, new_start - 1, -1):
      text_span = " ".join(doc_tokens[new_start:(new_end + 1)])
      if text_span == tok_answer_text:
        return (new_start, new_end)

  return (input_start, input_end)

def _check_is_max_context(doc_spans, cur_span_index, position):
  """Check if this is the 'max context' doc span for the token."""

  # Because of the sliding window approach taken to scoring documents, a single
  # token can appear in multiple documents. E.g.
  #  Doc: the man went to the store and bought a gallon of milk
  #  Span A: the man went to the
  #  Span B: to the store and bought
  #  Span C: and bought a gallon of
  #  ...
  #
  # Now the word 'bought' will have two scores from spans B and C. We only
  # want to consider the score with "maximum context", which we define as
  # the *minimum* of its left and right context (the *sum* of left and
  # right context will always be the same, of course).
  #
  # In the example the maximum context for 'bought' would be span C since
  # it has 1 left context and 3 right context, while span B has 4 left context
  # and 0 right context.
  best_score = None
  best_span_index = None
  for (span_index, doc_span) in enumerate(doc_spans):
    end = doc_span.start + doc_span.length - 1
    if position < doc_span.start:
      continue
    if position > end:
      continue
    num_left_context = position - doc_span.start
    num_right_context = end - position
    score = min(num_left_context, num_right_context) + 0.01 * doc_span.length
    if best_score is None or score > best_score:
      best_score = score
      best_span_index = span_index

  return cur_span_index == best_span_index

def convert2example(input_file, is_training=False,
                  repeat_limit=3):
  with open(input_file, 'r') as f:
    train_data = json.load(f)
    train_data = train_data['data']

  def _is_chinese_char(cp):
    if ((cp >= 0x4E00 and cp <= 0x9FFF) or  #
      (cp >= 0x3400 and cp <= 0x4DBF) or  #
      (cp >= 0x20000 and cp <= 0x2A6DF) or  #
      (cp >= 0x2A700 and cp <= 0x2B73F) or  #
      (cp >= 0x2B740 and cp <= 0x2B81F) or  #
      (cp >= 0x2B820 and cp <= 0x2CEAF) or
      (cp >= 0xF900 and cp <= 0xFAFF) or  #
      (cp >= 0x2F800 and cp <= 0x2FA1F)):  #
      return True

    return False

  def is_fuhao(c):
    if c == '。' or c == '，' or c == '！' or c == '？' or c == '；' or c == '、' or c == '：' or c == '（' or c == '）' \
      or c == '－' or c == '~' or c == '「' or c == '《' or c == '》' or c == ',' or c == '」' or c == '"' or c == '“' or c == '”' \
      or c == '$' or c == '『' or c == '』' or c == '—' or c == ';' or c == '。' or c == '(' or c == ')' or c == '-' or c == '～' or c == '。' \
      or c == '‘' or c == '’':
      return True
    return False

  def _tokenize_chinese_chars(text):
    """Adds whitespace around any CJK character."""
    output = []
    for char in text:
      cp = ord(char)
      if _is_chinese_char(cp) or is_fuhao(char):
        if len(output) > 0 and output[-1] != SPIECE_UNDERLINE:
          output.append(SPIECE_UNDERLINE)
        output.append(char)
        output.append(SPIECE_UNDERLINE)
      else:
        output.append(char)
    return "".join(output)

  def is_whitespace(c):
    if c == " " or c == "\t" or c == "\r" or c == "\n" or ord(
      c) == 0x202F or c == SPIECE_UNDERLINE:
      return True
    return False

  # to examples
  examples = []
  mis_match = 0
  for article in tqdm(train_data):
    for para in article['paragraphs']:
      context = para['context']
      context_chs = _tokenize_chinese_chars(context)
      doc_tokens = []
      char_to_word_offset = []
      prev_is_whitespace = True
      for c in context_chs:
        if is_whitespace(c):
          prev_is_whitespace = True
        else:
          if prev_is_whitespace:
            doc_tokens.append(c)
          else:
            doc_tokens[-1] += c
          prev_is_whitespace = False
        if c != SPIECE_UNDERLINE:
          char_to_word_offset.append(len(doc_tokens) - 1)

      for qas in para['qas']:
        qid = qas['id']
        ques_text = qas['question']
        ans_text = qas['answers'][0]['text']

        start_position_final = None
        end_position_final = None
        if is_training:
          count_i = 0
          start_position = qas['answers'][0]['answer_start']

          end_position = start_position + len(ans_text) - 1
          while context[
                start_position:end_position + 1] != ans_text and count_i < repeat_limit:
            start_position -= 1
            end_position -= 1
            count_i += 1

          while context[start_position] == " " or context[
            start_position] == "\t" or \
            context[start_position] == "\r" or context[
            start_position] == "\n":
            start_position += 1

          start_position_final = char_to_word_offset[start_position]
          end_position_final = char_to_word_offset[end_position]

          if doc_tokens[start_position_final] in {"。", "，", "：", ":", ".",
                                                  ","}:
            start_position_final += 1

          actual_text = "".join(
            doc_tokens[start_position_final:(end_position_final + 1)])
          cleaned_answer_text = "".join(
            tokenization.whitespace_tokenize(ans_text))

          if actual_text != cleaned_answer_text:
            print(actual_text, 'V.S', cleaned_answer_text)
            mis_match += 1
            # ipdb.set_trace()

        examples.append({'doc_tokens': doc_tokens,
                         'orig_answer_text': ans_text,
                         'qid': qid,
                         'question': ques_text,
                         'answer': ans_text,
                         'start_position': start_position_final,
                         'end_position': end_position_final})
  return examples

def convert2feature(examples, tokenizer, output_file=None,
                    is_training=False,  max_query_length=64,
                  max_seq_length=512, doc_stride=128):
  # to features
  features = []
  unique_id = 1000000000
  for (example_index, example) in enumerate(tqdm(examples)):
    query_tokens = tokenizer.tokenize(example['question'])
    if len(query_tokens) > max_query_length:
      query_tokens = query_tokens[0:max_query_length]

    tok_to_orig_index = []
    orig_to_tok_index = []
    all_doc_tokens = []
    for (i, token) in enumerate(example['doc_tokens']):
      orig_to_tok_index.append(len(all_doc_tokens))
      sub_tokens = tokenizer.tokenize(token)
      for sub_token in sub_tokens:
        tok_to_orig_index.append(i)
        all_doc_tokens.append(sub_token)

    tok_start_position = None
    tok_end_position = None
    if is_training:
      tok_start_position = orig_to_tok_index[
        example['start_position']]  # 原来token到新token的映射，这是新token的起点
      if example['end_position'] < len(example['doc_tokens']) - 1:
        tok_end_position = orig_to_tok_index[example['end_position'] + 1] - 1
      else:
        tok_end_position = len(all_doc_tokens) - 1
      (tok_start_position, tok_end_position) = _improve_answer_span(
        all_doc_tokens, tok_start_position, tok_end_position, tokenizer,
        example['orig_answer_text'])

    # The -3 accounts for [CLS], [SEP] and [SEP]
    max_tokens_for_doc = max_seq_length - len(query_tokens) - 3

    doc_spans = []
    _DocSpan = collections.namedtuple("DocSpan", ["start", "length"])
    start_offset = 0
    while start_offset < len(all_doc_tokens):
      length = len(all_doc_tokens) - start_offset
      if length > max_tokens_for_doc:
        length = max_tokens_for_doc
      doc_spans.append(_DocSpan(start=start_offset, length=length))
      if start_offset + length == len(all_doc_tokens):
        break
      start_offset += min(length, doc_stride)

    for (doc_span_index, doc_span) in enumerate(doc_spans):
      tokens = []
      token_to_orig_map = {}
      token_is_max_context = {}
      segment_ids = []
      tokens.append("[CLS]")
      segment_ids.append(0)
      for token in query_tokens:
        tokens.append(token)
        segment_ids.append(0)
      tokens.append("[SEP]")
      segment_ids.append(0)

      for i in range(doc_span.length):
        split_token_index = doc_span.start + i
        token_to_orig_map[len(tokens)] = tok_to_orig_index[split_token_index]
        is_max_context = _check_is_max_context(doc_spans, doc_span_index,
                                               split_token_index)
        token_is_max_context[len(tokens)] = is_max_context
        tokens.append(all_doc_tokens[split_token_index])
        segment_ids.append(1)
      tokens.append("[SEP]")
      segment_ids.append(1)

      input_ids = tokenizer.convert_tokens_to_ids(tokens)

      # The mask has 1 for real tokens and 0 for padding tokens. Only real
      # tokens are attended to.
      input_mask = [1] * len(input_ids)

      # Zero-pad up to the sequence length.
      while len(input_ids) < max_seq_length:
        input_ids.append(0)
        input_mask.append(0)
        segment_ids.append(0)

      assert len(input_ids) == max_seq_length
      assert len(input_mask) == max_seq_length
      assert len(segment_ids) == max_seq_length

      start_position = None
      end_position = None
      if is_training:
        # For training, if our document chunk does not contain an annotation
        # we throw it out, since there is nothing to predict.
        if tok_start_position == -1 and tok_end_position == -1:
          start_position = 0  # 问题本来没答案，0是[CLS]的位子
          end_position = 0
        else:  # 如果原本是有答案的，那么去除没有答案的feature
          out_of_span = False
          doc_start = doc_span.start  # 映射回原文的起点和终点
          doc_end = doc_span.start + doc_span.length - 1

          if not (
            tok_start_position >= doc_start and tok_end_position <= doc_end):  # 该划窗没答案作为无答案增强
            out_of_span = True
          if out_of_span:
            start_position = 0
            end_position = 0
          else:
            doc_offset = len(query_tokens) + 2
            start_position = tok_start_position - doc_start + doc_offset
            end_position = tok_end_position - doc_start + doc_offset

      features.append({'unique_id': unique_id,
                       'example_index': example_index,
                       'doc_span_index': doc_span_index,
                       'tokens': tokens,
                       'token_to_orig_map': token_to_orig_map,
                       'token_is_max_context': token_is_max_context,
                       'input_ids': input_ids,
                       'input_mask': input_mask,
                       'segment_ids': segment_ids,
                       'start_position': start_position,
                       'end_position': end_position})
      unique_id += 1

  print('features num:', len(features))
  json.dump(features, open(output_file, 'w'))
  return features

def json2features(input_file, output_files, tokenizer, is_training=False,
                  repeat_limit=3, max_query_length=64,
                  max_seq_length=512, doc_stride=128):
  with open(input_file, 'r') as f:
    train_data = json.load(f)
    train_data = train_data['data']

  def _is_chinese_char(cp):
    if ((cp >= 0x4E00 and cp <= 0x9FFF) or  #
      (cp >= 0x3400 and cp <= 0x4DBF) or  #
      (cp >= 0x20000 and cp <= 0x2A6DF) or  #
      (cp >= 0x2A700 and cp <= 0x2B73F) or  #
      (cp >= 0x2B740 and cp <= 0x2B81F) or  #
      (cp >= 0x2B820 and cp <= 0x2CEAF) or
      (cp >= 0xF900 and cp <= 0xFAFF) or  #
      (cp >= 0x2F800 and cp <= 0x2FA1F)):  #
      return True

    return False

  def is_fuhao(c):
    if c == '。' or c == '，' or c == '！' or c == '？' or c == '；' or c == '、' or c == '：' or c == '（' or c == '）' \
      or c == '－' or c == '~' or c == '「' or c == '《' or c == '》' or c == ',' or c == '」' or c == '"' or c == '“' or c == '”' \
      or c == '$' or c == '『' or c == '』' or c == '—' or c == ';' or c == '。' or c == '(' or c == ')' or c == '-' or c == '～' or c == '。' \
      or c == '‘' or c == '’':
      return True
    return False

  def _tokenize_chinese_chars(text):
    """Adds whitespace around any CJK character."""
    output = []
    for char in text:
      cp = ord(char)
      if _is_chinese_char(cp) or is_fuhao(char):
        if len(output) > 0 and output[-1] != SPIECE_UNDERLINE:
          output.append(SPIECE_UNDERLINE)
        output.append(char)
        output.append(SPIECE_UNDERLINE)
      else:
        output.append(char)
    return "".join(output)

  def is_whitespace(c):
    if c == " " or c == "\t" or c == "\r" or c == "\n" or ord(
      c) == 0x202F or c == SPIECE_UNDERLINE:
      return True
    return False

  # to examples
  examples = []
  mis_match = 0
  for article in tqdm(train_data):
    for para in article['paragraphs']:
      context = para['context']
      context_chs = _tokenize_chinese_chars(context)
      doc_tokens = []
      char_to_word_offset = []
      prev_is_whitespace = True
      for c in context_chs:
        if is_whitespace(c):
          prev_is_whitespace = True
        else:
          if prev_is_whitespace:
            doc_tokens.append(c)
          else:
            doc_tokens[-1] += c
          prev_is_whitespace = False
        if c != SPIECE_UNDERLINE:
          char_to_word_offset.append(len(doc_tokens) - 1)

      for qas in para['qas']:
        qid = qas['id']
        ques_text = qas['question']
        ans_text = qas['answers'][0]['text']

        start_position_final = None
        end_position_final = None
        if is_training:
          count_i = 0
          start_position = qas['answers'][0]['answer_start']

          end_position = start_position + len(ans_text) - 1
          while context[
                start_position:end_position + 1] != ans_text and count_i < repeat_limit:
            start_position -= 1
            end_position -= 1
            count_i += 1

          while context[start_position] == " " or context[
            start_position] == "\t" or \
            context[start_position] == "\r" or context[
            start_position] == "\n":
            start_position += 1

          start_position_final = char_to_word_offset[start_position]
          end_position_final = char_to_word_offset[end_position]

          if doc_tokens[start_position_final] in {"。", "，", "：", ":", ".",
                                                  ","}:
            start_position_final += 1

          actual_text = "".join(
            doc_tokens[start_position_final:(end_position_final + 1)])
          cleaned_answer_text = "".join(
            tokenization.whitespace_tokenize(ans_text))

          if actual_text != cleaned_answer_text:
            print(actual_text, 'V.S', cleaned_answer_text)
            mis_match += 1
            # ipdb.set_trace()

        examples.append({'doc_tokens': doc_tokens,
                         'orig_answer_text': ans_text,
                         'qid': qid,
                         'question': ques_text,
                         'answer': ans_text,
                         'start_position': start_position_final,
                         'end_position': end_position_final})

  print('examples num:', len(examples))
  print('mis_match:', mis_match)
  os.makedirs('/'.join(output_files[0].split('/')[0:-1]), exist_ok=True)
  json.dump(examples, open(output_files[0], 'w'))

  # to features
  features = []
  unique_id = 1000000000
  for (example_index, example) in enumerate(tqdm(examples)):
    query_tokens = tokenizer.tokenize(example['question'])
    if len(query_tokens) > max_query_length:
      query_tokens = query_tokens[0:max_query_length]

    tok_to_orig_index = []
    orig_to_tok_index = []
    all_doc_tokens = []
    for (i, token) in enumerate(example['doc_tokens']):
      orig_to_tok_index.append(len(all_doc_tokens))
      sub_tokens = tokenizer.tokenize(token)
      for sub_token in sub_tokens:
        tok_to_orig_index.append(i)
        all_doc_tokens.append(sub_token)

    tok_start_position = None
    tok_end_position = None
    if is_training:
      tok_start_position = orig_to_tok_index[
        example['start_position']]  # 原来token到新token的映射，这是新token的起点
      if example['end_position'] < len(example['doc_tokens']) - 1:
        tok_end_position = orig_to_tok_index[example['end_position'] + 1] - 1
      else:
        tok_end_position = len(all_doc_tokens) - 1
      (tok_start_position, tok_end_position) = _improve_answer_span(
        all_doc_tokens, tok_start_position, tok_end_position, tokenizer,
        example['orig_answer_text'])

    # The -3 accounts for [CLS], [SEP] and [SEP]
    max_tokens_for_doc = max_seq_length - len(query_tokens) - 3

    doc_spans = []
    _DocSpan = collections.namedtuple("DocSpan", ["start", "length"])
    start_offset = 0
    while start_offset < len(all_doc_tokens):
      length = len(all_doc_tokens) - start_offset
      if length > max_tokens_for_doc:
        length = max_tokens_for_doc
      doc_spans.append(_DocSpan(start=start_offset, length=length))
      if start_offset + length == len(all_doc_tokens):
        break
      start_offset += min(length, doc_stride)

    for (doc_span_index, doc_span) in enumerate(doc_spans):
      tokens = []
      token_to_orig_map = {}
      token_is_max_context = {}
      segment_ids = []
      tokens.append("[CLS]")
      segment_ids.append(0)
      for token in query_tokens:
        tokens.append(token)
        segment_ids.append(0)
      tokens.append("[SEP]")
      segment_ids.append(0)

      for i in range(doc_span.length):
        split_token_index = doc_span.start + i
        token_to_orig_map[len(tokens)] = tok_to_orig_index[split_token_index]
        is_max_context = _check_is_max_context(doc_spans, doc_span_index,
                                               split_token_index)
        token_is_max_context[len(tokens)] = is_max_context
        tokens.append(all_doc_tokens[split_token_index])
        segment_ids.append(1)
      tokens.append("[SEP]")
      segment_ids.append(1)

      input_ids = tokenizer.convert_tokens_to_ids(tokens)

      # The mask has 1 for real tokens and 0 for padding tokens. Only real
      # tokens are attended to.
      input_mask = [1] * len(input_ids)

      # Zero-pad up to the sequence length.
      while len(input_ids) < max_seq_length:
        input_ids.append(0)
        input_mask.append(0)
        segment_ids.append(0)

      assert len(input_ids) == max_seq_length
      assert len(input_mask) == max_seq_length
      assert len(segment_ids) == max_seq_length

      start_position = None
      end_position = None
      if is_training:
        # For training, if our document chunk does not contain an annotation
        # we throw it out, since there is nothing to predict.
        if tok_start_position == -1 and tok_end_position == -1:
          start_position = 0  # 问题本来没答案，0是[CLS]的位子
          end_position = 0
        else:  # 如果原本是有答案的，那么去除没有答案的feature
          out_of_span = False
          doc_start = doc_span.start  # 映射回原文的起点和终点
          doc_end = doc_span.start + doc_span.length - 1

          if not (
            tok_start_position >= doc_start and tok_end_position <= doc_end):  # 该划窗没答案作为无答案增强
            out_of_span = True
          if out_of_span:
            start_position = 0
            end_position = 0
          else:
            doc_offset = len(query_tokens) + 2
            start_position = tok_start_position - doc_start + doc_offset
            end_position = tok_end_position - doc_start + doc_offset

      features.append({'unique_id': unique_id,
                       'example_index': example_index,
                       'doc_span_index': doc_span_index,
                       'tokens': tokens,
                       'token_to_orig_map': token_to_orig_map,
                       'token_is_max_context': token_is_max_context,
                       'input_ids': input_ids,
                       'input_mask': input_mask,
                       'segment_ids': segment_ids,
                       'start_position': start_position,
                       'end_position': end_position})
      unique_id += 1

  print('features num:', len(features))
  json.dump(features, open(output_files[1], 'w'))

def convert_pre_res(input_ids, pre_start, pre_end, ori_start, ori_end, probabilities, tokenizer):
  """"""
  result = []
  for input, p_start, p_end, o_start, o_end, probability in zip(input_ids, pre_start, pre_end, ori_start, ori_end, probabilities):
    o_start, o_end, probability = o_start.tolist(), o_end.tolist(), probability.tolist()
    tokens = tokenizer.convert_ids_to_tokens(input.tolist())
    question = "".join(tokens[:tokens.index("[SEP]")])
    context = "".join(tokens[tokens.index("[SEP]"):])
    label_answer = "".join(
      tokens[o_start:o_end + 1])
    predict_answer = "".join(tokens[p_start:p_end + 1])
    cur_res = {
      "context": context,
      "question": question,
      "label_answer": label_answer,
      "predict_answer": predict_answer,
      "is_correct": label_answer == predict_answer,
      "probability": float(probability)
    }
    result.append(cur_res)

  return result


import time
from keras.callbacks import Callback
import numpy as np
import pandas as pd
import os

def progress_bar(current_value, max_value):
    progress = ((current_value+1)/max_value)*100
    if progress>98: progress=100
    print('\r[{0}{1}] {2:.1f}%'.format('#'*int(progress/2), ' '*(50-int(progress/2)), progress), end='')

def play_bell():
    import winsound
    duration = 200  # millisecond
    freq = 440  # Hz
    for i in range(5):
        winsound.Beep(int(freq*(i/2+1)), duration)

class LossHistory(Callback):
    def __init__(self, number_of_epochs):
        self.number_of_epochs = number_of_epochs
        self.current_epoch = 0
    def on_train_begin(self, logs={}):
        self.initial_time = time.time()
    def on_batch_end(self, batch, logs={}):
        if logs['batch']==0:
            NUMBER_OF_DIESIS = 20
            self.current_epoch += 1
            progress = self.current_epoch/self.number_of_epochs
            diesis = np.round(progress*NUMBER_OF_DIESIS).astype('int')
            eta = (time.time()-self.initial_time) * (self.number_of_epochs/self.current_epoch -1)
            remaining_time = time.strftime("%H hours, %M min, %S sec", time.gmtime(eta))
            print('\r[{}{}] acc: {:.3f} eta: {}'.format('#'*diesis, '-'*(NUMBER_OF_DIESIS-diesis), logs['acc'], remaining_time), end='')

def add_grid_and_save(grid):

    # save data in temporary dataframe if cross validation is not over
    if len(grid['fit_outs'][1]) < grid['skf_n_splits'][1]:
        temp_grid_file_path = '{}/model_data/grid_tmp.pkl'.format(grid['root_path'][1])
        grid_df = pd.DataFrame()
        grid_df = grid_df.append({key:grid[key][1] for key in grid.keys()}, ignore_index=True)
        grid_df.to_pickle(temp_grid_file_path)

    else:
        grid_file_path = '{}/model_data/grid_{}.pkl'.format(grid['root_path'][1], grid['version'][1])
        if os.path.isfile(grid_file_path):
            grid_df = pd.read_pickle(grid_file_path)
        else:
            grid_df = pd.DataFrame()

        grid_df = grid_df.append({key:grid[key][1] for key in grid.keys()}, ignore_index=True)
        grid_df.to_pickle(grid_file_path)


import time
from appium import webdriver

class DouYinApp():
    def __init__(self):
        self.params = {}
        self.params['deviceName'] = '127.0.0.1:62001 device'
        self.params['platformVersion'] = '5.1.1'
        self.params['automationName'] = 'Appium'
        self.params['platformName'] = 'Android'
        self.params['autoAcceptAlerts'] = 'true'
        self.params['noReset'] = 'true'
        self.params['appPackage'] = 'com.ss.android.ugc.aweme'
        self.params['appActivity'] = 'com.ss.android.ugc.aweme.live.LivePlayActivity'
        self.driver = webdriver.Remote('http://localhost:4723/wd/hub', desired_capabilities=self.params)

    def swipDown(self, n):
        L = self.driver.get_window_size()
        x = L['width'] * 0.5
        y1 = L['height'] * 0.8
        y2 = L['height'] * 0.3
        num = 0
        # while num < n:
        while True:
            self.driver.swipe(x, y1, x, y2, 6000)
            # num += 1

    def swipLeft(self):
        L = self.driver.get_window_size()
        x1 = L['width'] * 0.8
        x2 = L['width'] * 0.3
        y1 = L['height'] * 0.5

        self.driver.swipe(x1, y1, x2, y1, 500)

    def login(self):
        self.driver.find_element_by_xpath('//android.widget.TextView[@text="我"]').click()
        time.sleep(0.5)
        self.driver.find_element_by_xpath('//android.widget.TextView[@text="密码登录"]').click()
        time.sleep(0.5)
        self.driver.find_element_by_xpath('//android.widget.EditText[@text="请输入手机号"]').send_keys('18595794414')
        time.sleep(0.5)
        self.driver.find_element_by_id('com.ss.android.ugc.aweme:id/dfk').send_keys('525748037')
        time.sleep(0.5)
        self.driver.find_element_by_xpath('//android.widget.CheckBox').click()
        time.sleep(0.5)
        self.driver.find_element_by_xpath('//android.widget.TextView[@text="登录"]').click()
        time.sleep(1)
        self.driver.find_element_by_id('com.ss.android.ugc.aweme:id/ep8').click()
        print('登录成功')

    def guanggao(self):
        try:
            guanggao = self.driver.find_element_by_id('com.ss.android.ugc.aweme:id/eb')
        except:
            guanggao = ''
        return guanggao

    def stream(self):
        self.driver.find_element_by_id('com.ss.android.ugc.aweme:id/cnt').click()

    def more(self):
        self.driver.find_element_by_xpath('//android.widget.TextView[@text="更多直播"]').click()

    def run(self, n):
        try:
            guanggao = self.guanggao()
            if guanggao != '':
                guanggao.click()
            print('--------跳过广告--------')
            time.sleep(1)

            # print('开始登录')
            # self.login()
            # time.sleep(5)

            self.stream()
            print('----------直播---------')
            time.sleep(3)

            self.swipLeft()
            print('----------更多---------')
            time.sleep(2)

            print('--------开始爬取-----------')
            self.swipDown(n)

        except BaseException as e:
            print(e)

if __name__ == '__main__':
    douyin = DouYinApp()
    douyin.run(100)