import layers
import keras
import soundfile as sf
import librosa
import resampy
import os
import numpy as np
import sys
from typing import NamedTuple

class MelGenerator(keras.utils.Sequence):
    def __init__(self, n_ids, batch_size, DATA_ROOT):
        self.n_ids = n_ids
        self.batch_size = batch_size
        self.mel_basedir = os.path.join(DATA_ROOT, 'melspec/')
        with np.load(os.path.join(DATA_ROOT, 'openmic-2018.npz')) as OPENMIC:
            self.Y_true, self.Y_mask, self.sample_key = OPENMIC['Y_true'], OPENMIC['Y_mask'], OPENMIC['sample_key']

    def __len__(self):
        return int(np.ceil(len(self.n_ids) / float(self.batch_size)))

    def __getitem__(self, idx):
        batch_n_ids = self.n_ids[idx * self.batch_size:(idx + 1) * self.batch_size]
        return self._data_from_ids(batch_n_ids)

    def _data_from_ids(self, n_ids):
        x_ids = self.sample_key[n_ids]
        x_batch = []
        y_batch = np.multiply(self.Y_true[n_ids], self.Y_mask[n_ids]) > 0.5
        for x_id in x_ids:
            mel_subdir = os.path.join(self.mel_basedir, x_id[:3])
            melpath = os.path.join(mel_subdir, x_id+'.npz')
            with np.load(melpath) as f:
                x_batch.append(f['melspec'].T)
        x_batch = np.array(x_batch)
        x_batch = np.expand_dims(x_batch, axis=-1)

        return x_batch.astype('float32'), y_batch

    def on_epoch_end(self):
        np.random.shuffle(self.n_ids)

class MelParam(NamedTuple):
    sr: int
    n_fft: int
    hop_length: int

def save_melspec(x_id,
                 audio_basedir='/beegfs/qx244/ds/openmic-2018/audio/',
                 mel_basedir='/beegfs/qx244/ds/openmic-2018/melspec/',
                 recompute_if_exist=False,
                 param=MelParam(sr=44100,n_fft=2048,hop_length=512)):
    # create corresponding .npz files that stores the mel_spectrogram

    # check if feature dir exist
    audio_subdir = os.path.join(audio_basedir, x_id[:3])
    mel_subdir = os.path.join(mel_basedir, x_id[:3])
    if not os.path.isdir(mel_subdir):
        os.mkdir(mel_subdir)
    excerpt_path = os.path.join(audio_subdir, x_id+'.ogg')
    mel_path = os.path.join(mel_subdir, x_id+'.npz')

    if (not recompute_if_exist) and os.path.isfile(mel_path):
        print('{} already exist.'.format(mel_path))

    else:
        y, sr = sf.read(excerpt_path)
        # Make Mono
        if y.ndim == 2:
            y = y.mean(axis=-1)
        if sr != param.sr:
#             print('resampling excerpt {} from original sr {}'.format(x_id, sr))
            y = resampy.resample(y, sr, param.sr)
        melspec = librosa.feature.melspectrogram(y=y, sr=param.sr,
                                                 n_fft=param.n_fft, hop_length=param.hop_length)
        if melspec.shape != (128,862):
            print("bad shape! ID:{}, shape: {}".format(x_id, melspec.shape))
        np.savez(mel_path, x_id=x_id, melspec=melspec)

def construct_crnnL3_smp_tom(input_shape=(None, 128, 1), n_classes=20):
    '''
    Modified from Openmic's original models.py file
    CRNN with L3-style conv encoder
    Parameters
    ----------
    input_shape (defaults to (862, 128, 1).... but is this even what I think it is?)
    alpha (for softmax... large => max pooling)

    Returns
    -------
    Keras model
    '''

    input_feature = keras.Input(shape=input_shape)

    # Apply batch normalization
    x_bn = keras.layers.BatchNormalization()(input_feature)

    # BLOCK 1
    conv1 = keras.layers.Convolution2D(64, (3, 3),
                                   padding='same',
                                   activation='relu',
                                   kernel_initializer='he_normal')(x_bn)
    bn1 = keras.layers.BatchNormalization()(conv1)
    conv2 = keras.layers.Convolution2D(64, (3, 3),
                                   padding='same',
                                   activation='relu',
                                   kernel_initializer='he_normal')(bn1)
    bn2 = keras.layers.BatchNormalization()(conv2)
    pool2 = keras.layers.MaxPooling2D((2,2), padding='valid')(bn2)

    # BLOCK 2
    conv3 = keras.layers.Convolution2D(128, (3, 3),
                                   padding='same',
                                   activation='relu',
                                   kernel_initializer='he_normal')(pool2)
    bn3 = keras.layers.BatchNormalization()(conv3)
    conv4 = keras.layers.Convolution2D(128, (3, 3),
                                   padding='same',
                                   activation='relu',
                                   kernel_initializer='he_normal')(bn3)
    bn4 = keras.layers.BatchNormalization()(conv4)
    pool4 = keras.layers.MaxPooling2D((2, 2), padding='valid')(bn4)

    # BLOCK 3
    conv5 = keras.layers.Convolution2D(256, (3, 3),
                                   padding='same',
                                   activation='relu',
                                   kernel_initializer='he_normal')(pool4)
    bn5 = keras.layers.BatchNormalization()(conv5)
    conv6 = keras.layers.Convolution2D(256, (3, 3),
                                   padding='same',
                                   activation='relu',
                                   kernel_initializer='he_normal')(bn5)
    bn6 = keras.layers.BatchNormalization()(conv6)
    pool6 = keras.layers.MaxPooling2D((2, 2), padding='valid')(bn6)

    # BLOCK 4
    conv7 = keras.layers.Convolution2D(512, (3, 3),
                                       padding='same',
                                       activation='relu',
                                       kernel_initializer='he_normal')(pool6)
    bn7 = keras.layers.BatchNormalization()(conv7)
    conv8 = keras.layers.Convolution2D(512, (3, 3),
                                       padding='same',
                                       activation='relu',
                                       kernel_initializer='he_normal')(bn7)
    bn8 = keras.layers.BatchNormalization()(conv8)
    pool8 = keras.layers.MaxPooling2D((2, 2), padding='valid')(bn8)

    # CONV SQUEEZE
    conv_sq = keras.layers.Convolution2D(1024, (1, 8),
                                         padding='valid',
                                         activation='relu',
                                         kernel_initializer='he_normal')(pool8)
    bn8 = keras.layers.BatchNormalization()(conv_sq)
    sq2 = layers.SqueezeLayer(axis=-2)(bn8)

    # RNN
    # First recurrent layer: a 128-dim bidirectional gru
#     rnn1 = keras.layers.Bidirectional(keras.layers.GRU(128,
#                                                        return_sequences=True))(sq2)

    p0 = keras.layers.Dense(n_classes, activation='sigmoid')

    p_dynamic = keras.layers.TimeDistributed(p0, name='dynamic/tags')(sq2) #rnn1

    p_static = layers.AutoPool(axis=1,
                               name='static/tags')(p_dynamic)

    model = keras.models.Model(input_feature,
                               p_static)

    return model


import torch.nn as nn
from torchsummary import summary

class ResNet_FeatureExtractor(nn.Module):

    def __init__(self, input_channel, output_channel=512):
        super(ResNet_FeatureExtractor, self).__init__()
        self.ConvNet = ResNet(input_channel, output_channel, BasicBlock, [1, 2, 5, 3])

    def forward(self, input):
        return self.ConvNet(input)

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = self._conv3x3(inplanes, planes)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = self._conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def _conv3x3(self, in_planes, out_planes, stride=1):
        "3x3 convolution with padding"
        return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                         padding=1, bias=False)

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            residual = self.downsample(x)
        out += residual
        out = self.relu(out)

        return out

class ResNet(nn.Module):

    def __init__(self, input_channel, output_channel, block, layers):
        super(ResNet, self).__init__()

        self.output_channel_block = [int(output_channel / 4), int(output_channel / 2), output_channel, output_channel]

        self.inplanes = int(output_channel / 8)
        self.conv0_1 = nn.Conv2d(input_channel, int(output_channel / 16),
                                 kernel_size=3, stride=1, padding=1, bias=False)
        self.bn0_1 = nn.BatchNorm2d(int(output_channel / 16))
        self.conv0_2 = nn.Conv2d(int(output_channel / 16), self.inplanes,
                                 kernel_size=3, stride=1, padding=1, bias=False)
        self.bn0_2 = nn.BatchNorm2d(self.inplanes)
        self.relu = nn.ReLU(inplace=True)

        self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.layer1 = self._make_layer(block, self.output_channel_block[0], layers[0])
        self.conv1 = nn.Conv2d(self.output_channel_block[0], self.output_channel_block[
                               0], kernel_size=3, stride=1, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(self.output_channel_block[0])

        self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.layer2 = self._make_layer(block, self.output_channel_block[1], layers[1], stride=1)
        self.conv2 = nn.Conv2d(self.output_channel_block[1], self.output_channel_block[
                               1], kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(self.output_channel_block[1])

        self.maxpool3 = nn.MaxPool2d(kernel_size=2, stride=(2, 1), padding=(0, 1))
        self.layer3 = self._make_layer(block, self.output_channel_block[2], layers[2], stride=1)
        self.conv3 = nn.Conv2d(self.output_channel_block[2], self.output_channel_block[
                               2], kernel_size=3, stride=1, padding=1, bias=False)
        self.bn3 = nn.BatchNorm2d(self.output_channel_block[2])

        self.layer4 = self._make_layer(block, self.output_channel_block[3], layers[3], stride=1)
        self.conv4_1 = nn.Conv2d(self.output_channel_block[3], self.output_channel_block[
                                 3], kernel_size=2, stride=(2, 1), padding=(0, 1), bias=False)
        self.bn4_1 = nn.BatchNorm2d(self.output_channel_block[3])
        self.conv4_2 = nn.Conv2d(self.output_channel_block[3], self.output_channel_block[
                                 3], kernel_size=2, stride=1, padding=0, bias=False)
        self.bn4_2 = nn.BatchNorm2d(self.output_channel_block[3])

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv0_1(x)
        x = self.bn0_1(x)
        x = self.relu(x)
        x = self.conv0_2(x)
        x = self.bn0_2(x)
        x = self.relu(x)

        x = self.maxpool1(x)
        x = self.layer1(x)
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)

        x = self.maxpool2(x)
        x = self.layer2(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu(x)

        x = self.maxpool3(x)
        x = self.layer3(x)
        x = self.conv3(x)
        x = self.bn3(x)
        x = self.relu(x)

        x = self.layer4(x)
        x = self.conv4_1(x)
        x = self.bn4_1(x)
        x = self.relu(x)
        x = self.conv4_2(x)
        x = self.bn4_2(x)
        x = self.relu(x)

        return x

class BidirectionalLSTM(nn.Module):

    def __init__(self, nIn, nHidden, nOut):
        super(BidirectionalLSTM, self).__init__()

        self.rnn = nn.LSTM(nIn, nHidden, bidirectional=True)
        self.embedding = nn.Linear(nHidden * 2, nOut)

    def forward(self, input):
        recurrent, _ = self.rnn(input)
        T, b, h = recurrent.size()
        t_rec = recurrent.view(T * b, h)

        output = self.embedding(t_rec)  # [T * b, nOut]
        output = output.view(T, b, -1)

        return output

class CRNN(nn.Module):

    def __init__(self, imgH, nc, nclass, nh, n_rnn=2, leakyRelu=False):
        super(CRNN, self).__init__()
        assert imgH % 16 == 0, 'imgH has to be a multiple of 16'

        cnn = ResNet_FeatureExtractor(1, 512)
        self.cnn = cnn
        self.rnn = nn.Sequential(
            BidirectionalLSTM(512, nh, nh),
            BidirectionalLSTM(nh, nh, nclass))

    def forward(self, input):
        # feature extraction
        conv = self.cnn(input)
        b, c, h, w = conv.size()
        assert h == 1, "the height of conv must be 1"
        conv = conv.squeeze(2)
        conv = conv.permute(2, 0, 1)  # [w, b, c]

        # sequence modelings
        output = self.rnn(conv)

        return output

if __name__ == "__main__":
    model = CRNN(32, 1, 92, 256)
    summary(model, (1, 32, 100))


import sys

input = sys.stdin.readline
n, m = map(int, input().rstrip().split())

graph = [list(map(str, input())) for _ in range(n)]

from collections import deque

direction = [[1, 0], [1, -1], [1, 1], [0, 1], [0, -1], [-1, -1], [-1, 0], [-1, 1]]
queue = deque()
visit = [[False] * m for _ in range(n)]

def check(i, j):
    cnt = 0
    for dy, dx in direction:
        ny, nx = i + dy, j + dx
        if 0 <= ny < n and 0 <= nx < m:
            if graph[ny][nx] == '.': cnt += 1
    if graph[i][j] <= cnt:
        return True
    else:
        return False

for i in range(n):
    for j in range(m):

        if ord('1') <= ord(graph[i][j]) <= ord('9'):
            graph[i][j] = int(graph[i][j])
            if check(i, j):
                queue.append([i, j])
                visit[i][j] = True
days = 0
next_queue = deque()

while True:
    if not queue: break
    days += 1
    while queue:

        y, x = queue.popleft()
        graph[y][x] = '.'

        for dy, dx in direction:
            ny, nx = y + dy, x + dx

            if 0 <= ny < n and 0 <= nx < m and ord('1') <= ord(str(graph[ny][nx])) <= ord('8') and not visit[ny][nx]:
                if check(ny, nx):
                    next_queue.append([ny, nx])
                    visit[ny][nx] = True

    queue = next_queue
    next_queue = deque()

print(days)


import pandas as pd
from IPython.display import HTML

def side_by_side(df1, df2, name1='', name2=''):
    if isinstance(df1, pd.Series):
        df1 = df1.to_frame(name=df1.name)
    if isinstance(df2, pd.Series):
        df2 = df2.to_frame(name=df2.name)
    inline = 'style="display: float; max-width:50%" class="table"'
    q = '''

{}

{}

    '''.format(df1.style.set_table_attributes(inline)
                  .set_caption(name1).render(),
               df2.style.set_table_attributes(inline)
                  .set_caption(name2).render())
    return HTML(q)


"""@package NMF-visualization

Compare different spectrogram normalization methods for NMF.
"""

import numpy as np

from tabulate import tabulate

from lib.NMF import betaDivergence
from lib.NMF import frobenius
from lib.NMF import KLD
from lib.NMF import NMF
from lib.normalize import maxNormalize
from lib.normalize import RMSnormalize
from lib.normalize import sumNormalize
from lib.spectrogram import magnitudeSpectrogram
from lib.utils import createDir

AUDIO_PATH = "data/audio/"
SPECTROGRAM_PATH = "data/spectrograms/"

if (__name__ == "__main__"):

    excerpt = "bassoon-solo"
    hopLen = 10   # ms
    Fs = 44100

    try:
        print ("Loading spectrogram.")
        S = np.load(SPECTROGRAM_PATH + excerpt + ".npy")
    except:
        print ("Spectrogram not found. Calcualting from audio.")
        path = AUDIO_PATH + excerpt + ".wav"
        x, Fs = librosa.load(path, sr = Fs, mono = True)
        S = magnitudeSpectrogram(x, Fs)
        path = SPECTROGRAM_PATH + excerpt + ".npy"
        createDir(path)
        np.save(path, S)

    table = []
    headers = ["Normalization", "Frobenius norm", "Beta divergence", "KLD"]
    fmt = ".3g"

    # No normalization.
    print ("Calculating NMF with no normalization.")
    V = S
    H, W = NMF(V, verbose = True)
    Vestimate = W*H
    frobeniusError = frobenius(V, Vestimate)
    betaError = betaDivergence(V, Vestimate)
    kldError = KLD(V, Vestimate)
    table += [["None", frobeniusError, betaError, kldError]]

    # Sum normalization.
    print ("Calculating NMF with sum normalization.")
    V = sumNormalize(S, axis = 1)
    H, W = NMF(V, verbose = True)
    Vestimate = W*H
    frobeniusError = frobenius(V, Vestimate)
    betaError = betaDivergence(V, Vestimate)
    kldError = KLD(V, Vestimate)
    table += [["Sum", frobeniusError, betaError, kldError]]

    # Max normalization.
    print ("Calculating NMF with max normalization.")
    V = maxNormalize(S, axis = 1)
    H, W = NMF(V, verbose = True)
    Vestimate = W*H
    frobeniusError = frobenius(V, Vestimate)
    betaError = betaDivergence(V, Vestimate)
    kldError = KLD(V, Vestimate)
    table += [["Max", frobeniusError, betaError, kldError]]

    # RMS normalization.
    print ("Calculating NMF with RMS normalization.")
    V = RMSnormalize(S, axis = 1)
    H, W = NMF(V, verbose = True)
    Vestimate = W*H
    frobeniusError = frobenius(V, Vestimate)
    betaError = betaDivergence(V, Vestimate)
    kldError = KLD(V, Vestimate)
    table += [["RMS", frobeniusError, betaError, kldError]]

    print(tabulate(table, headers, tablefmt = "github", floatfmt = fmt))


# https://adventofcode.com/2021/day/1

from more_itertools import quantify

from adventofcode.utils import load_list

measurements = load_list(parser=int)

def part_1() -> int:
    # Take the measurements and create pairs one index apart, then map them to a bool
    # which is True if b is greater than a. `quantify` counts the True values
    return quantify(b > a for a, b in zip(measurements, measurements[1:]))

print(part_1())

def part_2(window_size: int) -> int:
    count = 0
    # Get the sum of the first `window_size` numbers
    cur_sum = sum(measurements[:window_size])

    # Start iterating at index `window_size`
    for index, measurement in enumerate(measurements[window_size:]):
        # Add the current measurement and subtract the one `window_size` ago.
        # For example, with
        #   [..., 100, 101, 102, 103, 104, ...]
        #              -------------
        # The current sum is 306. We subtract 100 and add 104 to get 310
        next_sum = cur_sum + measurement - measurements[index]

        if next_sum > cur_sum:
            count += 1

        cur_sum = next_sum

    return count

print(part_2(window_size=3))


# https://www.acmicpc.net/problem/1260
from collections import deque, defaultdict

def dfs(x):
    print(x, end = ' ')
    for i in tree[x]:
        if not visit_s[i]:
            visit_s[i] = True
            dfs(i)
    return

N, M, V = map(int, input().split(" "))
tree = defaultdict(list)

for _ in range(M):
    start, end = map(int, input().split(" "))
    tree[start].append(end)
    tree[end].append(start)

for t in tree:
    tree[t].sort()
# print(tree)

q = deque([V])
s = deque([V])
visit_q = [False]*(N+1)
visit_s = [False]*(N+1)
visit_q[V], visit_s[V] = True, True

dfs(V)
print()

while q:
    now = q.popleft()
    print(now, end=' ')
    for idx in tree[now]:
        if not visit_q[idx]:
            visit_q[idx] = True
            q.append(idx)


import datetime
import os
import numpy as np
import pandas as pd
import pytest
import shutil
import tempfile

dask = pytest.importorskip('dask')

from cyavro import dask_reader, write_avro_file_from_dataframe

@pytest.yield_fixture
def tempdir():
    d = tempfile.mkdtemp()
    try:
        yield d
    finally:
        shutil.rmtree(d)

N = 1000000
d = np.random.randint(0, 65000, size=N).astype('timedelta64')
t = d + np.array([datetime.datetime.now()]).astype('datetime64[ms]')
data = pd.DataFrame({'a': np.random.randint(0, 65000, size=N),
                     'b': np.random.random(size=N),
                     #'c': d,
                     #'d': t,
                     'e': ['a'] * N})

@pytest.mark.parametrize('block', ([1024, 1024**2, 1024**3]))
@pytest.mark.parametrize('mode', (['scan', 'seek', 'none']))
@pytest.mark.parametrize('bs', ([1024**2, 10*1024**3]))
def test_roundtrip(tempdir, block, mode, bs):
    fn = os.path.join(tempdir, 'test.avro')
    write_avro_file_from_dataframe(data, fn, block_size=block)
    out = dask_reader.read_avro(fn, block_finder=mode).compute()
    pd.util.testing.assert_frame_equal(data, out)

@pytest.mark.parametrize('block', ([1024, 1024**2, 1024**3]))
@pytest.mark.parametrize('mode', (['scan', 'seek', 'none']))
@pytest.mark.parametrize('bs', ([1024**2, 10*1024**3]))
def test_roundtrip_multifile(tempdir, block, mode, bs):
    fn1 = os.path.join(tempdir, 'test1.avro')
    write_avro_file_from_dataframe(data, fn1, block_size=block)
    fn2 = os.path.join(tempdir, 'test2.avro')
    write_avro_file_from_dataframe(data, fn2, block_size=block)
    expected = pd.concat([data, data]).reset_index(drop=True)
    fn = os.path.join(tempdir, 'test*.avro')
    out = dask_reader.read_avro(fn, block_finder=mode).compute()
    out.reset_index(drop=True, inplace=True)  # why does cyavro compute integers?
    pd.util.testing.assert_frame_equal(expected, out)


import os
import sys
import csv
import time
import torch
import logging
import dataclasses

import numpy      as np

from typing                                import Dict, List, Optional
from filelock                              import FileLock
from dataclasses                           import dataclass, field
from sklearn.metrics                       import f1_score, recall_score, precision_score
from sklearn.metrics                       import matthews_corrcoef
from torch.utils.data.dataset              import Dataset

from transformers                          import AutoTokenizer, EvalPrediction
from transformers                          import AutoConfig, AutoModelForSequenceClassification
from transformers.tokenization_utils       import PreTrainedTokenizer
from transformers.tokenization_roberta     import RobertaTokenizer, RobertaTokenizerFast
from transformers.data.processors.glue     import glue_convert_examples_to_features
from transformers.data.processors.utils    import InputFeatures
from transformers.data.processors.utils    import DataProcessor, InputExample, InputFeatures
from transformers.tokenization_xlm_roberta import XLMRobertaTokenizer
import six
from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    set_seed,
)

logger = logging.getLogger(__name__)

## From /src/transformers/data/metrics/__init__.py
def simple_accuracy(preds, labels):
    return (preds == labels).mean()
def acc_and_f1(preds, labels):
    acc = simple_accuracy(preds, labels)
    mcc = matthews_corrcoef(labels, preds) # labels and preds instead of other way!
    return_dict = {
        "acc" : acc,
        "mcc" : mcc,
    }
    for average in [ "binary", "micro", "macro", "weighted"] :
        f1        = f1_score       (y_true=labels, y_pred=preds, average=average)
        precision = precision_score(y_true=labels, y_pred=preds, average=average)
        recall    = recall_score   (y_true=labels, y_pred=preds, average=average)
        return_dict[ "f1_"        + average ] = f1
        return_dict[ "precision_" + average ] = precision
        return_dict[ "recall_"    + average ] = recall
    return return_dict

# ColaProcessor from src/transformers/data/processors/glue.py
class Processor(DataProcessor):
    """Processor for the CoLA data set (GLUE version)."""

    def get_example_from_tensor_dict(self, tensor_dict):
        """See base class."""
        return InputExample(
            tensor_dict["idx"].numpy(),
            tensor_dict["data"].numpy().decode("utf-8"),
            None,
            str(tensor_dict["label"].numpy()),
        )

    def get_train_examples(self, data_dir):
        """See base class."""
        return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")

    def get_dev_examples(self, data_dir):
        """See base class."""
        return self._create_examples(self._read_tsv(os.path.join(data_dir, "val.tsv")), "dev")

    def get_labels(self):
        """See base class."""

        return [str(x) for x in range(2)]

    def _create_examples(self, lines, set_type):
        """Creates examples for the training and dev sets."""
        examples = []
        for (i, line) in enumerate(lines):
            guid = "%s-%s" % (set_type, i)

            text_a = line[3]
            # label = line[5]
            if 'Security' not in line[5]:
                label = str(0)
            else:
                label = str(1)
            examples.append(InputExample(guid=guid, text_a=text_a, text_b=None, label=label))
        print(self.get_labels())
        return examples

# GlueDataTrainingArguments from: src/transformers/data/datasets/glue.py
@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.

    Using `HfArgumentParser` we can turn this class
    into argparse arguments to be able to specify them on
    the command line.
    """

    name: str = field(metadata={"help": "A name for this task (Cache folders will be based on this"})
    data_dir: str = field(
        metadata={"help": "The input data dir. Should contain the .tsv files (or other data files) for the task."}
    )
    max_seq_length: int = field(
        default=128,
        metadata={
            "help": "The maximum total input sequence length after tokenization. Sequences longer "
            "than this will be truncated, sequences shorter will be padded."
        },
    )
    overwrite_cache: bool = field(
        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
    )

# From: src/transformers/data/datasets/glue.py
class GlueDataset(Dataset):
    """
    This will be superseded by a framework-agnostic approach
    soon.
    """

    args: DataTrainingArguments
    output_mode: str
    features: List[InputFeatures]

    def __init__(
        self,
        args: DataTrainingArguments,
        tokenizer: PreTrainedTokenizer,
        limit_length: Optional[int] = None,
        evaluate=False,
    ):
        self.args = args
        processor = Processor()
        self.output_mode = 'classification'
        # Load data features from cache or dataset file
        cached_features_file = os.path.join(
            args.data_dir,
            "cached_{}_{}_{}_{}".format(
                "dev" if evaluate else "train", tokenizer.__class__.__name__, str(args.max_seq_length), args.name,
            ),
        )
        print(cached_features_file)
        # Make sure only the first process in distributed training processes the dataset,
        # and the others will use the cache.
        lock_path = cached_features_file + ".lock"
        with FileLock(lock_path):

            if os.path.exists(cached_features_file) and not args.overwrite_cache:
                start = time.time()
                self.features = torch.load(cached_features_file)
                logger.info(
                    f"Loading features from cached file {cached_features_file} [took %.3f s]", time.time() - start
                )
            else:
                logger.info(f"Creating features from dataset file at {args.data_dir}")
                label_list = processor.get_labels()

                # if args.task_name in ["mnli", "mnli-mm"] and tokenizer.__class__ in (
                #     RobertaTokenizer,
                #     RobertaTokenizerFast,
                #     XLMRobertaTokenizer,
                # ):
                #     # HACK(label indices are swapped in RoBERTa pretrained model)
                #     label_list[1], label_list[2] = label_list[2], label_list[1]

                examples = (
                    processor.get_dev_examples(args.data_dir)
                    if evaluate
                    else processor.get_train_examples(args.data_dir)
                )
                if limit_length is not None:
                    examples = examples[:limit_length]
                self.features = glue_convert_examples_to_features(
                    examples,
                    tokenizer,
                    max_length=args.max_seq_length,
                    label_list=label_list,
                    output_mode='classification',
                )
                start = time.time()
                torch.save(self.features, cached_features_file)
                # ^ This seems to take a lot of time so I want to investigate why and how we can improve.
                logger.info(
                    "Saving features into cached file %s [took %.3f s]", cached_features_file, time.time() - start
                )

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

    def __getitem__(self, i) -> InputFeatures:
        return self.features[i]

@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """

    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )
    cache_dir: Optional[str] = field(
        default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
    )
    class_weights: Optional[str] = field(
        default=None, metadata={"help": "Comma seperated list of class weights. Weight for Class 0, Weight for Class 1"}
    )

class Modifiedargments:

    def __init__(self, datadir):
        self.name = 'CoLA'
        self.data_dir =  datadir
        self.max_seq_length = 128
        self.overwrite_cache = False

def main():
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))

    print(os.path.abspath(sys.argv[1]))

    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    if (
        os.path.exists(training_args.output_dir)
        and os.listdir(training_args.output_dir)
        and training_args.do_train
        and not training_args.overwrite_output_dir
    ):
        raise ValueError(
            f"Output directory ({training_args.output_dir}) already exists and is not empty. Use --overwrite_output_dir to overcome."
        )

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,
    )
    logger.warning(
        "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
        training_args.local_rank,
        training_args.device,
        training_args.n_gpu,
        bool(training_args.local_rank != -1),
        training_args.fp16,
    )
    logger.info("Training/evaluation parameters %s", training_args)

    # Set seed
    set_seed(training_args.seed)

    num_labels  = 2 ## Defined in processor. Change labels before changing this (also hardcoded elsewhere) .
    output_mode = 'classification'

    # Load pretrained model and tokenizer
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.

    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        num_labels=2, ## Defined in processor. Change labels before changing this (also hardcoded elsewhere) .
        finetuning_task='cola', ## Not sure why we need this!
        cache_dir=model_args.cache_dir,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
    )
    model = AutoModelForSequenceClassification.from_pretrained(
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
        config=config,
        cache_dir=model_args.cache_dir,
    )

    #############################################################################
    #                            Set Class Weights                             ##
    #############################################################################
    class_weights = None
    if not model_args.class_weights is None :
        class_weights = model_args.class_weights.lstrip().rstrip().split( ',' )
        class_weights = [ float(i) for i in class_weights ]
        logger.info(
            "Using Class Weights {}".format( ','.join( [ str(i) for i in class_weights ] ) )
        )
        class_weights = torch.tensor( class_weights, dtype=torch.float, device=training_args.device )
    model.set_class_weights( class_weights )
    #############################################################################

    # Get datasets
    aspect = 'Security'
    print(data_args)
    train_path = '/workspace/RoBERTa_textclassification/aspect_classifier/costsensitivebert/Cost-Sensitive_Bert_and_Transformers/examples/glue_data/CoLA/'+aspect+'/0/'
    eval_path = '/workspace/RoBERTa_textclassification/aspect_classifier/costsensitivebert/Cost-Sensitive_Bert_and_Transformers/examples/glue_data/CoLA/'+aspect+'/0/'
    train_dataset = GlueDataset(Modifiedargments(datadir = train_path), tokenizer=tokenizer) if training_args.do_train else None
    eval_dataset = GlueDataset(Modifiedargments(datadir = eval_path), tokenizer=tokenizer, evaluate=True) if training_args.do_eval else None

    def compute_metrics(p: EvalPrediction) -> Dict:
        if output_mode == "classification":
            preds = np.argmax(p.predictions, axis=1)
        elif output_mode == "regression":
            preds = np.squeeze(p.predictions)
        # return glue_compute_metrics(data_args.task_name, preds, p.label_ids)
        with open( training_args.output_dir + 'predictions.csv', 'w' ) as fh:
            writer = csv.writer( fh )
            preds  = [ [i] for i in preds ]
            writer.writerows( preds )
        return acc_and_f1(preds, p.label_ids)

    # Initialize our Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        compute_metrics=compute_metrics,
    )

    # Training
    if training_args.do_train:
        trainer.train(
            model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
        )
        trainer.save_model()
        # For convenience, we also re-save the tokenizer to the same directory,
        # so that you can share your model easily on huggingface.co/models =)
        if trainer.is_world_master():
            tokenizer.save_pretrained(training_args.output_dir)

    # Evaluation
    results = {}
    if training_args.do_eval and training_args.local_rank in [-1, 0]:
        logger.info("*** Evaluate ***")

        # Loop to handle MNLI double evaluation (matched, mis-matched)
        eval_datasets = [eval_dataset]
        # if data_args.task_name == "mnli":
        #     mnli_mm_data_args = dataclasses.replace(data_args, task_name="mnli-mm")
        #     eval_datasets.append(GlueDataset(mnli_mm_data_args, tokenizer=tokenizer, evaluate=True))

        for eval_dataset in eval_datasets:
            result = trainer.evaluate(eval_dataset=eval_dataset)

            output_eval_file = os.path.join(
                training_args.output_dir, f"eval_results_classification.txt"
            )
            with open(output_eval_file, "w") as writer:
                logger.info("***** Eval results classification *****")
                for key, value in result.items():
                    logger.info("  %s = %s", key, value)
                    writer.write("%s = %s\n" % (key, value))

            results.update(result)

    return results

def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()

if __name__ == "__main__":
    main()


from flask import  render_template,session
from flask import Flask,redirect,url_for
from forms import Loginform
from pessoalform import pessoalform
from parentescoform import parentescoform
from trabalhoform import trabalhoform
import os.path

import models
import contracheques

app = Flask(__name__)
app.config.from_object(__name__)

# Set the secret key to some random bytes. Keep this really secret!
app.secret_key = b'_5#y2L"F4Q8z\n\xec]/'

@app.route("/login/", methods=["GET","POST"])
def login():
    global form
    form = Loginform()
    nome = str(form.inputEmail.data).upper()
    result = models.conn.execute(models.s, nome=form.inputEmail.data).fetchone()
    if form.validate_on_submit() and nome == result[2] and form.inputPassword.data.strip()== result[3]:
        global session
        session["username"] = result[0]
        print(session["username"])
        fotos()
        print(result[1])
        return redirect(url_for('index'))

    else:
        print(form.errors)
        return render_template('login.html',
                           form=form)

@app.route("/index/")
def index():
    return render_template('SIDENAVIGATION.html')

@app.route("/construcao/")
def construcao():
    return render_template('construcao.html')

@app.route("/contracheque/")
def contracheque():
    if session["username"]:
        comp = "052018"
        file = comp+"/"+session["username"]+".pdf"
        return redirect(url_for('static', filename=file))
    else:
        print('Matricula diferente!')
    return render_template('construcao.html')

@app.route("/fotos/")
def fotos():
    photos = "fotos/"+session["username"]+".jpg"
    try:
        with open('static/'+photos, 'r') as f:
            return redirect(url_for('static', filename=photos))
    except IOError:
        return redirect(url_for('static', filename="fotos/loginvazio.jpg"))

@app.route("/cadastro/")
def cadastro():
    #return render_template('Side_Cadastro.html')
    return render_template('construcao.html')

@app.route("/pessoal/", methods=['GET', 'POST'])
def pessoal():
    form = pessoalform()
    return render_template('pessoal.html',form=form)

@app.route("/parentesco/", methods=['GET', 'POST'])
def parentesco():
    form = parentescoform()
    return render_template('parentesco.html',form=form)

@app.route("/trabalho/", methods=['GET', 'POST'])
def trabalho():
    form = trabalhoform()
    return render_template('trabalho.html',form=form)

app.run(threaded=True,debug=True, use_reloader=True,host='0.0.0.0')


#!/usr/bin/env python3
# -*- coding: UTF8 -*-

# Author: Guillaume Bouvier -- guillaume.bouvier@pasteur.fr
# https://research.pasteur.fr/en/member/guillaume-bouvier/
# 2021-04-30 11:09:18 (UTC+0200)

import os
import glob
import numpy as np
import scipy.spatial.distance as distance
from pymol import cmd
scriptdir = os.path.dirname(os.path.realpath(__file__))
cmd.set('fetch_path', cmd.exp_path('~/pdb'), quiet=1)

def get_chain(mobile, target):
    coords = cmd.get_coords(f'{mobile} and polymer.protein')
    coords_ref = cmd.get_coords(f'{target} and polymer.protein')
    cdistmat = distance.cdist(coords, coords_ref)
    myspace = {'chains': []}
    # cmd.iterate(f'{mobile} and polymer.protein', 'chains.append(chain)', space=myspace)
    cmd.iterate_state(1, f'{mobile} and polymer.protein', 'chains.append(chain)', space=myspace)
    chains_out = np.asarray(myspace['chains'])
    sel = cdistmat.min(axis=1) <= 2.
    # print(sel.shape, chains_out.shape, coords.shape)
    chains_unique, count = np.unique(chains_out[sel], return_counts=True)
    chain_aligned = chains_unique[count.argmax()]
    return chain_aligned

if __name__ == '__main__':
    import argparse
    # argparse.ArgumentParser(prog=None, usage=None, description=None, epilog=None, parents=[], formatter_class=argparse.HelpFormatter, prefix_chars='-', fromfile_prefix_chars=None, argument_default=None, conflict_handler='error', add_help=True, allow_abbrev=True, exit_on_error=True)
    parser = argparse.ArgumentParser(description='')
    # parser.add_argument(name or flags...[, action][, nargs][, const][, default][, type][, choices][, required][, help][, metavar][, dest])
    parser.add_argument('-l', '--list', type=str,
                        default=f'{scriptdir}/data/missing_pdbs_testset.txt')
    args = parser.parse_args()

    with open(args.list, 'r') as apofile:
        lines = apofile.readlines()
        n = len(lines)
        visited = set()
        for i, line in enumerate(lines):
            line = line.strip()
            line = os.path.splitext(line)[0]
            uniprot, apo, holo = line.split('_')
            if holo in visited:
                continue
            if apo == holo:
                continue
            visited.add(holo)
            print(f'id: {i+1}/{n}')
            print(f'uniprot: {uniprot}')
            print(f'holo: {holo}')
            print(f'apo: {apo}')
            holopdb = f'PL-database/{holo[0]}/{holo[1]}/{holo[2]}/{holo[3]}/{holo}-?-{uniprot}.pdb'
            holopdb = glob.glob(holopdb)
            assert len(holopdb) == 1
            holopdb = holopdb[0]
            cmd.load(holopdb, 'holo')
            cmd.fetch(apo, name='apo', type='pdb')
            # cmd.remove("not alt ''+A")
            # cmd.alter("all", "alt=''")
            rmsd = cmd.align(mobile='apo and polymer.protein',
                             target='holo and polymer.protein')[0]
            print(f'rmsd: {rmsd:.3f}')
            chain_aligned = get_chain('apo', 'holo')
            print(f'chain_aligned: {chain_aligned}')
            outfilename = f'{uniprot}_{apo}_{holo}.pdb'
            cmd.save(f'apo/{outfilename}', selection=f'apo and chain {chain_aligned}')
            print(f'apofile: {outfilename}')
            print()
            cmd.reinitialize()
            cmd.set('fetch_path', cmd.exp_path('~/pdb'), quiet=1)


from PIL import Image, ImageDraw
import numpy as np

input_image = Image.open("edge-detection2.jpg")
input_pixels = input_image.load()

vert = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]])
hor = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
sobel = np.array([[1, 2, 2], [1, 1, 0], [2, 4, 1]])

kernel = vert

offset = len(kernel) // 2
output_image = Image.new("RGB", input_image.size)
draw = ImageDraw.Draw(output_image)
for x in range(offset, input_image.width - offset):
    for y in range(offset, input_image.height - offset):
        acc = [0, 0, 0]
        for a in range(len(kernel)):
            for b in range(len(kernel)):
                xn = x + a - offset
                yn = y + b - offset
                pixel = input_pixels[xn, yn]
                acc[0] += pixel[0] * kernel[a][b]
                acc[1] += pixel[1] * kernel[a][b]
                acc[2] += pixel[2] * kernel[a][b]

        draw.point((x, y), (int(acc[0]), int(acc[1]), int(acc[2])))
output_image.show()
output_image.save('Vertical Edge Detection-Cube.jpg')


#import stuff
from numpy import *
from random import normalvariate, random, uniform
from scipy.optimize import fmin

def get_likelihood( params, subj_trials, res_index, rew_index, n_options):

    alpha = params[0]
    B_value = params[1]
    stay_prob = params[2]
    mixture = params[3]
    neg_log_likelihood = 0.

    if (alpha < 0.0) or (alpha > 1.0): return inf
    if B_value < 0.0: return inf
    if (stay_prob < 0.0) or (stay_prob > 0.5): return inf
    if (mixture < 0.0) or (mixture > 1.0): return inf

    Q = array([0.0]*n_options)

    last_choice = -1

    for trial in subj_trials:

        probs = [0.25, 0.25, 0.25, 0.25]
        sw_probs = [(1-stay_prob)/(n_options-1), (1-stay_prob)/(n_options-1), (1-stay_prob)/(n_options-1), (1-stay_prob)/(n_options-1)]

        numerators = [ exp(i*B_value) for i in Q]
        denominator = sum(numerators)

        RL_probs = [ i/denominator for i in numerators]

        choice = int(trial[res_index])
        payoff = float(trial[rew_index])

        sw_probs[choice] = stay_prob

        for i in range(len(probs)):
            probs[i] = (RL_probs[i]*mixture) + (sw_probs[i]*(1-mixture))

        if trial[2] == '1':
            prob = 0.25
        else:
        	prob = probs[choice]

        neg_log_likelihood += -log( prob )

        Q[choice] += alpha * ( payoff - Q[choice] )
        last_choice = choice

#         print RL_probs
#         print sw_probs
#         print mixture
#         print probs
#         print '*************'

    return neg_log_likelihood

def do_fit(subj_data, res_index, rew_index, n_options):

    #print subj_data

    subj = subj_data[0][0]

    tries = 15

    outputs = []

    for i in range(tries):

        # sets the starting parameters by picking them randomly
        # params: alpha,  temp
        iparams = [ uniform(0,1),  uniform(0, 3), uniform(0,1), uniform(0,1)]
        #print iparams

        res = fmin(get_likelihood, iparams, args=(subj_data, res_index, rew_index, n_options,), full_output=True) #(function to input, parameters to start with, other inputs to the function, output shit)

        outputs.append([res[1], res[0]])

    minLLH = 1000

    for i in range(len(outputs)):
            if outputs[i][0] < minLLH:
                minLLH = outputs[i][0]
                params = outputs[i][1]

    return [minLLH, params]