import argparse
import itertools
import json
import multiprocessing as mp
import os
import pickle
import random
import re
import string
import sys
import time
import json
import math
import copy
from collections import Counter, OrderedDict

import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader, RandomSampler, TensorDataset
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm
from transformers import AdamW, BertModel, BertTokenizer, get_linear_schedule_with_warmup

from util import load_pickle, save_pickle, count_parameters, compute_metrics, compute_metrics_from_logits
import logging

logging.basicConfig(level = logging.INFO, \
                    format = '%(asctime)s  %(levelname)-5s %(message)s', \
                    datefmt =  "%Y-%m-%d-%H-%M-%S")

def cprint(*args):
    text = ""
    for arg in args:
        text += "{0} ".format(arg)
    logging.info(text)

def tokenize_conversations(data, tokenizer, max_sent_len):
    new_data = []
    for conv in tqdm(data):
        new_conv = []
        for i, (speaker, sents) in enumerate(conv):
            # each utterance has been segmented into multiple sentences
            if i==0:
                word_limit = 90
            else:
                word_limit = max_sent_len

            tokenized_sent = []
            for sent in sents:
                tokenized = tokenizer.tokenize(sent)
                if len(tokenized_sent) + len(tokenized) <= word_limit:
                    tokenized_sent.extend(tokenized)
                else:
                    break
            if len(tokenized_sent) == 0:
                tokenized_sent = tokenized[:word_limit]
            new_conv.append((speaker, tokenized_sent))
        new_data.append(new_conv)
    return new_data

def tokenize_personas(data, tokenizer, all_speakers, num_personas):
    # average: each speaker corresponds to a list of tokens, separated by [SEP] between sents
    # memnet: each speaker corresponds to a 2D list of tokens
    new_data = {}
    for k, sents in tqdm(data):
        if k in all_speakers:
            tokenized_words = []
            for sent in sents[:num_personas]:
                tokenized_words.extend(tokenizer.tokenize(sent) + ["[SEP]"])
            if len(tokenized_words) > 1:
                tokenized_words.pop() # remove the last [SEP]
                new_data[k] = tokenized_words
            else:
                new_data[k] = ["."]
    return new_data

def create_context_and_response(data):
    new_data = []
    for conv in tqdm(data):
        context = []
        for s, ts in conv[:-1]:
            context.extend(ts + ["[SEP]"])
        context.pop() # pop the last [SEP]
        response = conv[-1][1]
        # if len(context) > 0 and len(response) > 0:
        #     new_data.append((context, response, conv[-1][0]))
        if len(response) == 0:
            response.extend(" ")
        new_data.append((context, response, conv[-1][0]))
    return new_data

def convert_conversations_to_ids(data, persona, tokenizer, max_seq_len, max_sent_len, num_personas):
    def pad_tokens(tokens, max_len, sentence_type, num_personas=0, response_ids=None):
        # note token_type_ids to differentiate context utterances
        # speaker A has 0, speaker B has 1, response is speaker B and has 1, persona has 1
        # persona does not have positional embedding
        if sentence_type == "persona" and num_personas > 0:
            # filter persona sentences that appeared in response_ids
            if response_ids is not None:
                response_sent = " ".join(tokenizer.convert_ids_to_tokens(response_ids, skip_special_tokens=True))

                all_persona_sent_ids = []
                for t_id in tokens:
                    if t_id in [101]:
                        sent_ids = []
                    if t_id in [102]:
                        all_persona_sent_ids.append(sent_ids)
                        sent_ids = []
                    if t_id not in tokenizer.all_special_ids:
                        sent_ids.append(t_id)

                # convert ids to tokens
                filtered_tokens = []
                for sent_ids in all_persona_sent_ids:
                    sent = " ".join(tokenizer.convert_ids_to_tokens(sent_ids))
                    if sent not in response_sent:
                        filtered_tokens.extend(sent_ids + [tokenizer.convert_tokens_to_ids("[SEP]")])
                filtered_tokens.insert(0, tokenizer.convert_tokens_to_ids("[CLS]"))

                tokens = filtered_tokens

            # remove additional persona sentences
            persona_sent_count = 0
            truncated_tokens = []
            for token_id in tokens:
                if token_id == tokenizer.convert_tokens_to_ids("[SEP]"):
                    persona_sent_count += 1
                    if persona_sent_count == num_personas:
                        break
                truncated_tokens.append(token_id)
            tokens = truncated_tokens

        # clip too much tokens
        while len(tokens) >= max_len and sentence_type == "context":
            pos = tokens.index(tokenizer.convert_tokens_to_ids('[SEP]')) + 1
            tokens = tokens[pos:]
            pos = tokens.index(tokenizer.convert_tokens_to_ids('[SEP]')) + 1
            tokens = tokens[pos:]
            tokens.insert(0, tokenizer.convert_tokens_to_ids("[CLS]"))
        assert max_len >= len(tokens)
        attention_mask = [1]*len(tokens)
        padding_length = max_len - len(tokens)
        attention_mask = attention_mask + ([0] * padding_length)

        if sentence_type == "context":
            token_type_ids = []
            token_type = 0
            for token_id in tokens:
                token_type_ids.append(token_type)
                if token_id == tokenizer.convert_tokens_to_ids("[SEP]"):
                    token_type = int(1-token_type)
            token_type_ids = token_type_ids + [0] * padding_length
        else:
            token_type_ids = [0] * max_len

        tokens = tokens + [0] * padding_length
        return tokens, attention_mask, token_type_ids

    all_context_ids = []
    all_context_attention_mask = []
    all_context_token_type_ids = []
    all_response_ids = []
    all_response_attention_mask = []
    all_response_token_type_ids = []
    all_persona_ids = []
    all_persona_attention_mask = []
    all_persona_token_type_ids = []
    max_persona_len = 23*num_personas+1
    context_lens = []
    for context, response, speaker in tqdm(data):
        context_ids = tokenizer.encode(context, add_special_tokens=True) # convert to token ids, add [cls] and [sep] at beginning and end
        response_ids = tokenizer.encode(response, add_special_tokens=True)
        context_lens.append(len(context_ids))

        context_ids, context_attention_mask, context_token_type_ids = pad_tokens(context_ids, max_seq_len, "context")
        response_ids, response_attention_mask, response_token_type_ids = pad_tokens(response_ids, max_sent_len+2, "response")

        all_context_ids.append(context_ids)
        all_context_attention_mask.append(context_attention_mask)
        all_context_token_type_ids.append(context_token_type_ids)
        all_response_ids.append(response_ids)
        all_response_attention_mask.append(response_attention_mask)
        all_response_token_type_ids.append(response_token_type_ids)

        if persona is not None:
            persona_ids = tokenizer.encode(persona[speaker], add_special_tokens=True)
            persona_ids, persona_attention_mask, persona_token_type_ids = pad_tokens(persona_ids, max_persona_len, "persona", num_personas, response_ids)
            # persona_ids, persona_attention_mask, persona_token_type_ids = pad_tokens(persona_ids, max_persona_len, "persona", num_personas)
            all_persona_ids.append(persona_ids)
            all_persona_attention_mask.append(persona_attention_mask)
            all_persona_token_type_ids.append(persona_token_type_ids)

    # (num_examples, max_seq_len)
    all_context_ids = torch.tensor(all_context_ids, dtype=torch.long)
    all_context_attention_mask = torch.tensor(all_context_attention_mask, dtype=torch.long)
    all_context_token_type_ids = torch.tensor(all_context_token_type_ids, dtype=torch.long)

    # (num_examples, max_sent_len)
    all_response_ids = torch.tensor(all_response_ids, dtype=torch.long)
    all_response_attention_mask = torch.tensor(all_response_attention_mask, dtype=torch.long)
    all_response_token_type_ids = torch.tensor(all_response_token_type_ids, dtype=torch.long)

    if persona is not None:
        # (num_examples, max_persona_len)
        all_persona_ids = torch.tensor(all_persona_ids, dtype=torch.long)
        all_persona_attention_mask = torch.tensor(all_persona_attention_mask, dtype=torch.long)
        all_persona_token_type_ids = torch.tensor(all_persona_token_type_ids, dtype=torch.long)

    cprint(all_context_ids.shape, all_context_attention_mask.shape, all_context_token_type_ids.shape)
    cprint(all_response_ids.shape, all_response_attention_mask.shape, all_response_token_type_ids.shape)

    if persona is not None:
        cprint(all_persona_ids.shape, all_persona_attention_mask.shape, all_persona_token_type_ids.shape)
        dataset = TensorDataset(all_context_ids, all_context_attention_mask, all_context_token_type_ids, \
            all_response_ids, all_response_attention_mask, all_response_token_type_ids, \
                all_persona_ids, all_persona_attention_mask, all_persona_token_type_ids)
    else:
        dataset = TensorDataset(all_context_ids, all_context_attention_mask, all_context_token_type_ids, \
            all_response_ids, all_response_attention_mask, all_response_token_type_ids)

    cprint("context lens stats: ", min(context_lens), max(context_lens), \
        np.mean(context_lens), np.std(context_lens))
    return dataset

def match(model, matching_method, x, y, x_mask, y_mask):
    # Multi-hop Co-Attention
    # x: (batch_size, m, hidden_size)
    # y: (batch_size, n, hidden_size)
    # x_mask: (batch_size, m)
    # y_mask: (batch_size, n)
    assert x.dim() == 3 and y.dim() == 3
    assert x_mask.dim() == 2 and y_mask.dim() == 2
    assert x_mask.shape == x.shape[:2] and y_mask.shape == y.shape[:2]
    m = x.shape[1]
    n = y.shape[1]

    attn_mask = torch.bmm(x_mask.unsqueeze(-1), y_mask.unsqueeze(1)) # (batch_size, m, n)
    attn = torch.bmm(x, y.transpose(1,2)) # (batch_size, m, n)
    model.attn = attn
    model.attn_mask = attn_mask

    x_to_y = torch.softmax(attn * attn_mask + (-5e4) * (1-attn_mask), dim=2) # (batch_size, m, n)
    y_to_x = torch.softmax(attn * attn_mask + (-5e4) * (1-attn_mask), dim=1).transpose(1,2) # # (batch_size, n, m)

    # x_attended, y_attended = None, None # no hop-1
    x_attended = torch.bmm(x_to_y, y) # (batch_size, m, hidden_size)
    y_attended = torch.bmm(y_to_x, x) # (batch_size, n, hidden_size)

    # x_attended_2hop, y_attended_2hop = None, None # no hop-2
    y_attn = torch.bmm(y_to_x.mean(dim=1, keepdim=True), x_to_y) # (batch_size, 1, n) # true important attention over y
    x_attn = torch.bmm(x_to_y.mean(dim=1, keepdim=True), y_to_x) # (batch_size, 1, m) # true important attention over x

    # truly attended representation
    x_attended_2hop = torch.bmm(x_attn, x).squeeze(1) # (batch_size, hidden_size)
    y_attended_2hop = torch.bmm(y_attn, y).squeeze(1) # (batch_size, hidden_size)

    # # hop-3
    # y_attn, x_attn = torch.bmm(x_attn, x_to_y), torch.bmm(y_attn, y_to_x) # (batch_size, 1, n) # true important attention over y
    # x_attended_3hop = torch.bmm(x_attn, x).squeeze(1) # (batch_size, hidden_size)
    # y_attended_3hop = torch.bmm(y_attn, y).squeeze(1) # (batch_size, hidden_size)
    # x_attended_2hop = torch.cat([x_attended_2hop, x_attended_3hop], dim=-1)
    # y_attended_2hop = torch.cat([y_attended_2hop, y_attended_3hop], dim=-1)

    x_attended = x_attended, x_attended_2hop
    y_attended = y_attended, y_attended_2hop

    return x_attended, y_attended

def aggregate(model, aggregation_method, x, x_mask):
    # x: (batch_size, seq_len, emb_size)
    # x_mask: (batch_size, seq_len)
    assert x.dim() == 3 and x_mask.dim() == 2
    assert x.shape[:2] == x_mask.shape
    # batch_size, seq_len, emb_size = x.shape

    if aggregation_method == "mean":
        return (x * x_mask.unsqueeze(-1)).sum(dim=1)/x_mask.sum(dim=-1, keepdim=True).clamp(min=1) # (batch_size, emb_size)

    if aggregation_method == "max":
        return x.masked_fill(x_mask.unsqueeze(-1)==0, -5e4).max(dim=1)[0] # (batch_size, emb_size)

    if aggregation_method == "mean_max":
        return torch.cat([(x * x_mask.unsqueeze(-1)).sum(dim=1)/x_mask.sum(dim=-1, keepdim=True).clamp(min=1), \
            x.masked_fill(x_mask.unsqueeze(-1)==0, -5e4).max(dim=1)[0]], dim=-1) # (batch_size, 2*emb_size)

def fuse(model, matching_method, aggregation_method, batch_x_emb, batch_y_emb, batch_persona_emb, \
    batch_x_mask, batch_y_mask, batch_persona_mask, batch_size, num_candidates):

    batch_x_emb, batch_y_emb_context = match(model, matching_method, batch_x_emb, batch_y_emb, batch_x_mask, batch_y_mask)
    # batch_x_emb: ((batch_size*num_candidates, m, emb_size), (batch_size*num_candidates, emb_size))
    # batch_y_emb_context: (batch_size*num_candidates, n, emb_size), (batch_size*num_candidates, emb_size)

    # hop 2 results
    batch_x_emb_2hop = batch_x_emb[1]
    batch_y_emb_context_2hop = batch_y_emb_context[1]

    # mean_max aggregation for the 1st hop result
    batch_x_emb = aggregate(model, aggregation_method, batch_x_emb[0], batch_x_mask) # batch_x_emb: (batch_size*num_candidates, 2*emb_size)
    batch_y_emb_context = aggregate(model, aggregation_method, batch_y_emb_context[0], batch_y_mask) # batch_y_emb_context: (batch_size*num_candidates, 2*emb_size)

    if batch_persona_emb is not None:
        batch_persona_emb, batch_y_emb_persona = match(model, matching_method, batch_persona_emb, batch_y_emb, batch_persona_mask, batch_y_mask)
        # batch_persona_emb: (batch_size*num_candidates, m, emb_size), (batch_size*num_candidates, emb_size)
        # batch_y_emb_persona: (batch_size*num_candidates, n, emb_size), (batch_size*num_candidates, emb_size)

        batch_persona_emb_2hop = batch_persona_emb[1]
        batch_y_emb_persona_2hop = batch_y_emb_persona[1]

        # # no hop-1
        # return torch.bmm(torch.cat([batch_x_emb_2hop, batch_persona_emb_2hop], dim=-1).unsqueeze(1), \
        #             torch.cat([batch_y_emb_context_2hop, batch_y_emb_persona_2hop], dim=-1)\
        #                 .unsqueeze(-1)).reshape(batch_size, num_candidates)

        batch_persona_emb = aggregate(model, aggregation_method, batch_persona_emb[0], batch_persona_mask) # batch_persona_emb: (batch_size*num_candidates, 2*emb_size)
        batch_y_emb_persona = aggregate(model, aggregation_method, batch_y_emb_persona[0], batch_y_mask) # batch_y_emb_persona: (batch_size*num_candidates, 2*emb_size)

        # # no hop-2
        # return torch.bmm(torch.cat([batch_x_emb, batch_persona_emb], dim=-1).unsqueeze(1), \
        #             torch.cat([batch_y_emb_context, batch_y_emb_persona], dim=-1)\
        #                 .unsqueeze(-1)).reshape(batch_size, num_candidates)
        return torch.bmm(torch.cat([batch_x_emb, batch_x_emb_2hop, batch_persona_emb, batch_persona_emb_2hop], dim=-1).unsqueeze(1), \
                    torch.cat([batch_y_emb_context, batch_y_emb_context_2hop, batch_y_emb_persona, batch_y_emb_persona_2hop], dim=-1)\
                        .unsqueeze(-1)).reshape(batch_size, num_candidates)
    else:
        return torch.bmm(torch.cat([batch_x_emb, batch_x_emb_2hop], dim=-1).unsqueeze(1), \
                    torch.cat([batch_y_emb_context, batch_y_emb_context_2hop], dim=-1)\
                        .unsqueeze(-1)).reshape(batch_size, num_candidates)

def dot_product_loss(batch_x_emb, batch_y_emb):
    """
        if batch_x_emb.dim() == 2:
            # batch_x_emb: (batch_size, emb_size)
            # batch_y_emb: (batch_size, emb_size)

        if batch_x_emb.dim() == 3:
            # batch_x_emb: (batch_size, batch_size, emb_size), the 1st dim is along examples and the 2nd dim is along candidates
            # batch_y_emb: (batch_size, emb_size)
    """
    batch_size = batch_x_emb.size(0)
    targets = torch.arange(batch_size, device=batch_x_emb.device)

    if batch_x_emb.dim() == 2:
        dot_products = batch_x_emb.mm(batch_y_emb.t())
    elif batch_x_emb.dim() == 3:
        dot_products = torch.bmm(batch_x_emb, batch_y_emb.unsqueeze(0).repeat(batch_size, 1, 1).transpose(1,2))[:, targets, targets] # (batch_size, batch_size)

    # dot_products: [batch, batch]
    log_prob = F.log_softmax(dot_products, dim=1)
    loss = F.nll_loss(log_prob, targets)
    nb_ok = (log_prob.max(dim=1)[1] == targets).float().sum()
    return loss, nb_ok

def evaluate_epoch(data_iter, models, num_personas, gradient_accumulation_steps, device, dataset, epoch, \
    apply_interaction, matching_method, aggregation_method, tokenizer, info_, num_cand, rank1_k, rank2_k):
    epoch_loss = []
    ok = 0
    total = 0
    recall = []
    MRR = []

    rank_KoK_context = []
    rank_KoK_response = []
    rank_KoK_logits = []
    rank_KoK_persona = []
    rank_KoK_info = []

    if len(models) == 1:
        if num_personas == 0:
            context_model, response_model = models[0], models[0]
        else:
            context_model, response_model, persona_model = models[0], models[0], models[0]
    if len(models) == 2:
        context_model, response_model = models
    if len(models) == 3:
        context_model, response_model, persona_model = models
    ##################################################################
    print("rank 1 processing...")
    count = 0
    ##################################################################
    for batch_idx, batch, in enumerate(data_iter):
        batch = tuple(t.to(device) for t in batch)
        batch_y = {"input_ids": batch[3], "attention_mask": batch[4], "token_type_ids": batch[5]}
        ##################################################################
        count += 1
        input_batch_x = []
        input_batch_y = []
        input_batch_persona = []
        for k in batch[0]:
            tokens = tokenizer.convert_ids_to_tokens(k.cpu())
            tokens.remove('[CLS]')
            while '[PAD]' in tokens: tokens.remove('[PAD]')
            tokens = [item.replace("[SEP]", "    ") for item in tokens]
            strings = tokenizer.convert_tokens_to_string(tokens)
            # input_batch_x.append(' '.join(strings))
            input_batch_x.append(strings)
        for k in batch[3]:
            # tokens = tokenizer.convert_ids_to_tokens(k.cpu())
            # strings = tokenizer.convert_tokens_to_string(tokens)
            # t = ' '.join(tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(k.cpu())))
            tokens = tokenizer.convert_ids_to_tokens(k.cpu())
            tokens.remove('[CLS]')
            try:
                while '[PAD]' in tokens: tokens.remove('[PAD]')
            except ValueError:
                print(tokens)
            try:
                while '[SEP]' in tokens: tokens.remove('[SEP]')
            except ValueError:
                print(tokens)
            strings = tokenizer.convert_tokens_to_string(tokens)
            # input_batch_y.append(t)
            input_batch_y.append(strings)
        for k in batch[6]:
            # input_batch_persona.append(' '.join(tokenizer.convert_ids_to_tokens(k.cpu())))
            tokens = tokenizer.convert_ids_to_tokens(k.cpu())
            tokens.remove('[CLS]')
            # tokens.remove('[PAD]')
            # tokens.remove('[SEP]')
            while '[PAD]' in tokens: tokens.remove('[PAD]')
            while '[SEP]' in tokens: tokens.remove('[SEP]')
            strings = tokenizer.convert_tokens_to_string(tokens)
            input_batch_persona.append(strings)
        ##################################################################
        has_persona = len(batch) > 6

        # get context embeddings in chunks due to memory constraint
        batch_size = batch[0].shape[0]
        chunk_size = 20
        num_chunks = math.ceil(batch_size/chunk_size)

        if apply_interaction:
            # batch_x_mask = batch[0].ne(0).float()
            # batch_y_mask = batch[3].ne(0).float()
            batch_x_mask = batch[1].float()
            batch_y_mask = batch[4].float()

            batch_x_emb = []
            batch_x_pooled_emb = []
            with torch.no_grad():
                for i in range(num_chunks):
                    mini_batch_x = {
                        "input_ids": batch[0][i*chunk_size: (i+1)*chunk_size],
                        "attention_mask": batch[1][i*chunk_size: (i+1)*chunk_size],
                        "token_type_ids": batch[2][i*chunk_size: (i+1)*chunk_size]
                        }
                    mini_output_x = context_model(**mini_batch_x)
                    batch_x_emb.append(mini_output_x[0]) # [(chunk_size, seq_len, emb_size), ...]
                    batch_x_pooled_emb.append(mini_output_x[1])
                batch_x_emb = torch.cat(batch_x_emb, dim=0) # (batch_size, seq_len, emb_size)
                batch_x_pooled_emb = torch.cat(batch_x_pooled_emb, dim=0)
                emb_size = batch_x_emb.shape[-1]
            if has_persona:
                # batch_persona_mask = batch[6].ne(0).float()
                batch_persona_mask = batch[7].float()
                batch_persona_emb = []
                batch_persona_pooled_emb = []
                with torch.no_grad():
                    for i in range(num_chunks):
                        mini_batch_persona = {
                            "input_ids": batch[6][i*chunk_size: (i+1)*chunk_size],
                            "attention_mask": batch[7][i*chunk_size: (i+1)*chunk_size],
                            "token_type_ids": batch[8][i*chunk_size: (i+1)*chunk_size]
                            }
                        mini_output_persona = persona_model(**mini_batch_persona)
                        # [(chunk_size, emb_size), ...]
                        batch_persona_emb.append(mini_output_persona[0])
                        batch_persona_pooled_emb.append(mini_output_persona[1])
                    batch_persona_emb = torch.cat(batch_persona_emb, dim=0)
                    batch_persona_pooled_emb = torch.cat(batch_persona_pooled_emb, dim=0)

            with torch.no_grad():
                output_y = response_model(**batch_y)
                batch_y_emb = output_y[0]

            batch_size, sent_len, emb_size = batch_y_emb.shape

            # interaction
            # context-response attention
            num_candidates = batch_size

            with torch.no_grad():
                # evaluate per example
                logits = []
                for i in range(batch_size):
                    x_emb = batch_x_emb[i:i+1].repeat_interleave(num_candidates, dim=0) # (num_candidates, context_len, emb_size)
                    x_mask = batch_x_mask[i:i+1].repeat_interleave(num_candidates, dim=0) # (batch_size*num_candidates, context_len)
                    persona_emb, persona_mask = None, None
                    if has_persona:
                        persona_emb = batch_persona_emb[i:i+1].repeat_interleave(num_candidates, dim=0)
                        persona_mask = batch_persona_mask[i:i+1].repeat_interleave(num_candidates, dim=0)
                    logits_single = fuse(context_model, matching_method, aggregation_method, \
                        x_emb, batch_y_emb, persona_emb, x_mask, batch_y_mask, persona_mask, 1, num_candidates).reshape(-1)

                    logits.append(logits_single)
                logits = torch.stack(logits, dim=0)

                # compute loss
                targets = torch.arange(batch_size, dtype=torch.long, device=batch[0].device)
                loss = F.cross_entropy(logits, targets)
            ##################################################################
            input_logits = []
            for l in logits.cpu():
                row = []
                for r in l:
                    row.append(r.item())
                input_logits.append(row)
            rank_arr = input_logits[0]

            dict = {"context": input_batch_x, "response": input_batch_y, "logits": rank_arr, "persona": input_batch_persona}
            df = pd.DataFrame(dict)
            df = df.sort_values(by=['logits'], axis=0, ascending=False)

            # rank_KoK.append(df[0:rank1_k])
            for i in range(rank1_k):
                rank_KoK_context.append(df.iloc[i, 0])
                rank_KoK_response.append(df.iloc[i, 1])
                rank_KoK_logits.append(df.iloc[i, 2])
                rank_KoK_persona.append(df.iloc[i, 3])
            input_logits_argmax = []
            for l in logits:
                input_logits_argmax.append(l.cpu().float().argmax().item())
            ####################################################################
            num_ok = (targets.long() == logits.float().argmax(dim=1)).sum()
            valid_recall, valid_MRR = compute_metrics_from_logits(logits, targets)
        else:
            batch_x_emb = []
            with torch.no_grad():
                for i in range(num_chunks):
                    mini_batch_x = {
                        "input_ids": batch[0][i*chunk_size: (i+1)*chunk_size],
                        "attention_mask": batch[1][i*chunk_size: (i+1)*chunk_size],
                        "token_type_ids": batch[2][i*chunk_size: (i+1)*chunk_size]
                        }
                    mini_output_x = context_model(**mini_batch_x)
                    batch_x_emb.append(mini_output_x[0].mean(dim=1)) # [(chunk_size, emb_size), ...]
                batch_x_emb = torch.cat(batch_x_emb, dim=0) # (batch_size, emb_size)
                emb_size = batch_x_emb.shape[-1]

            if has_persona:
                batch_persona_emb = []
                with torch.no_grad():
                    for i in range(num_chunks):
                        mini_batch_persona = {
                            "input_ids": batch[6][i*chunk_size: (i+1)*chunk_size],
                            "attention_mask": batch[7][i*chunk_size: (i+1)*chunk_size],
                            "token_type_ids": batch[8][i*chunk_size: (i+1)*chunk_size]
                            }
                        mini_output_persona = persona_model(**mini_batch_persona)

                        # [(chunk_size, emb_size), ...]
                        batch_persona_emb.append(mini_output_persona[0].mean(dim=1))

            with torch.no_grad():
                batch_persona_emb = torch.cat(batch_persona_emb, dim=0)
                batch_x_emb = (batch_x_emb + batch_persona_emb)/2
                output_y = response_model(**batch_y)
                batch_y_emb = output_y[0].mean(dim=1)

            # compute loss
            loss, num_ok = dot_product_loss(batch_x_emb, batch_y_emb)
            valid_recall, valid_MRR = compute_metrics(batch_x_emb, batch_y_emb)

        ok += num_ok.item()
        total += batch[0].shape[0]

        # compute valid recall
        recall.append(valid_recall)
        MRR.append(valid_MRR)

        if gradient_accumulation_steps > 1:
            loss = loss / gradient_accumulation_steps
        epoch_loss.append(loss.item())

        # if batch_idx%print_every == 0:
        #     cprint("loss: ", np.mean(epoch_loss[-print_every:]))
        #     cprint("valid recall: ", np.mean(recall[-print_every:], axis=0))
        #     cprint("valid MRR: ", np.mean(MRR[-print_every:], axis=0))
    print("info_ len: ", len(info_))
    # info_idx = 0
    # for i in range(0, len(info_)):
    #     if info_idx >= len(info_):
    #         break
    #     num_exp_persona = info_[info_idx][2]
    #     for j in range(num_exp_persona):
    #         if info_idx + j >= len(info_):
    #             break
    #         rank_KoK_info.append(info_[info_idx])
    #     info_idx += num_exp_persona
    for i in range(0, len(info_), num_cand):
        rank_KoK_info.append(info_[i])
    acc = ok/total
    # compute recall for validation dataset
    recall = np.mean(recall, axis=0)
    MRR = np.mean(MRR)
    rank1_dict = {'context': rank_KoK_context, 'response': rank_KoK_response,
            'logits': rank_KoK_logits, 'persona': rank_KoK_persona, 'info': rank_KoK_info}

    print("rank 2 processing...")
    c = 0
    rank2_context = []
    rank2_response = []
    rank2_logits = []
    rank2_info = []

    print(len(rank1_dict['context']))
    print(len(rank1_dict['response']))
    print(len(rank1_dict['logits']))
    print(len(rank1_dict['persona']))
    print(len(rank1_dict['info']))

    while c < len(rank1_dict['response']):
        num_exp_persona = rank1_dict['info'][c][2]

        context = []
        response = []
        logits = []
        info = []
        for d in range(num_exp_persona):
            if c + d >= len(rank1_dict['response']):
                break
            context.append(rank1_dict['context'][c + d])
            response.append(rank1_dict['response'][c + d])
            logits.append(rank1_dict['logits'][c + d])
            info.append(rank1_dict['info'][c + d])
        c += num_exp_persona
        rank2_dict = {'context': context, 'response': response, 'logits': logits, 'info': info}
        df = pd.DataFrame(rank2_dict)
        df = df.sort_values(by=['logits'], axis=0, ascending=False)
        for e in range(rank2_k):
            if e >= len(response):
                break
            rank2_context.append(df.iloc[e, 0])
            rank2_response.append(df.iloc[e, 1])
            rank2_logits.append(df.iloc[e, 2])
            rank2_info.append(df.iloc[e, 3])
    # while c < len(rank1_dict['response']):
    #     num_exp_persona = rank1_dict['info'][c][2]
    #     response = []
    #     logits = []
    #     info = []
    #     rank2_dict = {'response': response[c:(c + num_exp_persona)], 'logits': logits[c:(c + num_exp_persona)],
    #                   'info': info[c:(c + num_exp_persona)]}
    #     c += num_exp_persona
    #     df = pd.DataFrame(rank2_dict)
    #     df = df.sort_values(by=['logits'], axis=0, ascending=False)
    #     end = min(num_exp_persona, rank2_k)
    #     for e in range(end):
    #         if e >= len(response):
    #             break
    #         rank2_response.append(df.iloc[e, 0])
    #         rank2_logits.append(df.iloc[e, 1])
    #         rank2_info.append(df.iloc[e, 2])
    rank2_dict = {'context': rank2_context, 'response': rank2_response, 'logits': rank2_logits, 'info': rank2_info}

    return np.mean(epoch_loss), (acc, recall, MRR), rank1_dict, rank2_dict

def main(config, progress):
    # save config
    with open("./log/configs.json", "a") as t:
        json.dump(config, t)
        t.write("\n")
    cprint("*"*80)
    cprint("Experiment progress: {0:.2f}%".format(progress*100))
    cprint("*"*80)
    metrics = {}

    # data hyper-params
    train_path = config["train_path"]
    valid_path = config["valid_path"]
    test_path = config["test_path"]
    dataset = train_path.split("/")[1]
    test_mode = bool(config["test_mode"])
    load_model_path = config["load_model_path"]
    save_model_path = config["save_model_path"]
    personaChat_file_path = config["personaChat_file_path"]
    num_candidates = config["num_candidates"]
    num_personas = config["num_personas"]
    persona_path = config["persona_path"]
    max_sent_len = config["max_sent_len"]
    max_seq_len = config["max_seq_len"]
    PEC_ratio = config["PEC_ratio"]
    train_ratio = config["train_ratio"]
    rank1_k = config["rank1_k"]
    rank2_k = config["rank2_k"]

    if PEC_ratio != 0 and train_ratio != 1:
        raise ValueError("PEC_ratio or train_ratio not qualified!")

    # model hyper-params
    config_id = config["config_id"]
    model = config["model"]
    shared = bool(config["shared"])
    apply_interaction = bool(config["apply_interaction"])
    matching_method = config["matching_method"]
    aggregation_method = config["aggregation_method"]
    output_hidden_states = False

    # training hyper-params
    batch_size = config["batch_size"]
    epochs = config["epochs"]
    warmup_steps = config["warmup_steps"]
    gradient_accumulation_steps = config["gradient_accumulation_steps"]
    lr = config["lr"]
    weight_decay = 0
    seed = config["seed"]
    device = torch.device(config["device"])
    fp16 = bool(config["fp16"])
    fp16_opt_level = config["fp16_opt_level"]

    # set seed
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)

    if test_mode and load_model_path == "":
        raise ValueError("Must specify test model path when in test mode!")

    # load data
    cprint("Loading conversation data...")
    train = load_pickle(train_path)
    train = train[0]
    valid = load_pickle(valid_path)
    test = load_pickle(test_path)
    valid_path = test_path
    valid = test[0]
    valid_info = test[1]
    print("valid len: ", len(valid))
    cprint("sample train data: ", train[0])
    cprint("sample valid data: ", valid[0])

    # tokenization
    cprint("Tokenizing ...")
    tokenizer = BertTokenizer.from_pretrained(model)
    cached_tokenized_train_path = train_path.replace(".pkl", "_tokenized.pkl")
    cached_tokenized_valid_path = valid_path.replace(".pkl", "_tokenized.pkl")
    if os.path.exists(cached_tokenized_train_path):
        cprint("Loading tokenized dataset from ", cached_tokenized_train_path)
        train = load_pickle(cached_tokenized_train_path)
    else:
        train = tokenize_conversations(train, tokenizer, max_sent_len)
        cprint("Saving tokenized dataset to ", cached_tokenized_train_path)
        save_pickle(train, cached_tokenized_train_path)

    if os.path.exists(cached_tokenized_valid_path):
        cprint("Loading tokenized dataset from ", cached_tokenized_valid_path)
        valid = load_pickle(cached_tokenized_valid_path)

    else:
        valid = tokenize_conversations(valid, tokenizer, max_sent_len)
        cprint("Saving tokenized dataset to ", cached_tokenized_valid_path)
        save_pickle(valid, cached_tokenized_valid_path)

    persona = None
    if num_personas > 0:
        cprint("Tokenizing persona sentences...")
        cached_tokenized_persona_path = persona_path.replace(".pkl", "_tokenized.pkl")
        if os.path.exists(cached_tokenized_persona_path):
            cprint("Loading tokenized persona from file...")
            persona = load_pickle(cached_tokenized_persona_path)
        else:
            cprint("Loading persona data...")
            persona = load_pickle(persona_path)
            all_speakers = set([s for conv in load_pickle(config["train_path"])[0] + \
                load_pickle(config["valid_path"])[0] + load_pickle(config["test_path"])[0] for s, sent in conv])
            cprint("Tokenizing persona data...")
            persona = tokenize_personas(persona, tokenizer, all_speakers, num_personas)
            cprint("Saving tokenized persona to file...")
            save_pickle(persona, cached_tokenized_persona_path)
        cprint("Persona dataset statistics (after tokenization):", len(persona))
        cprint("Sample tokenized persona:", list(persona.values())[0])

    cprint("Sample tokenized data: ")
    cprint(train[0])
    cprint(valid[0])

    # select subsets of training and validation data for casualconversation
    cprint(dataset)

    # create context and response
    train = create_context_and_response(train)
    valid = create_context_and_response(valid)
    cprint("Sample context and response: ")
    cprint(train[0])
    cprint(valid[0])

    # convert to token ids
    cprint("Converting conversations to ids: ")
    valid_dataset = convert_conversations_to_ids(valid, persona, tokenizer, \
        max_seq_len, max_sent_len, num_personas)
    ####################################################
    valid_dataloader = DataLoader(valid_dataset, batch_size=num_candidates)

    # create model
    cprint("Building model...")
    model = BertModel.from_pretrained(model, output_hidden_states=output_hidden_states)
    cprint(model)
    cprint("number of parameters: ", count_parameters(model))

    if shared:
        cprint("number of encoders: 1")
        models = [model]
    else:
        if num_personas == 0:
            cprint("number of encoders: 2")
            # models = [model, copy.deepcopy(model)]
            models = [model, pickle.loads(pickle.dumps(model))]
        else:
            cprint("number of encoders: 3")
            # models = [model, copy.deepcopy(model), copy.deepcopy(model)]
            models = [model, pickle.loads(pickle.dumps(model)), pickle.loads(pickle.dumps(model))]

    cprint("Loading weights from ", load_model_path)
    model.load_state_dict(torch.load(load_model_path))
    models = [model]

    for i, model in enumerate(models):
        cprint("model {0} number of parameters: ".format(i), count_parameters(model))
        model.to(device)

    # optimization
    amp = None
    if fp16:
        from apex import amp

    no_decay = ["bias", "LayerNorm.weight"]
    optimizers = []
    schedulers = []
    for i, model in enumerate(models):
        optimizer_grouped_parameters = [
            {
                "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
                "weight_decay": weight_decay,
            },
            {"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], "weight_decay": 0.0},
        ]
        optimizer = AdamW(optimizer_grouped_parameters, lr=lr, eps=1e-8)

        if fp16:
            model, optimizer = amp.initialize(model, optimizer, opt_level=fp16_opt_level)
            models[i] = model
        optimizers.append(optimizer)

    # evaluation
    for model in models:
        model.eval()
    valid_iterator = tqdm(valid_dataloader, desc="Iteration")
    valid_loss, (valid_acc, valid_recall, valid_MRR), rank1_dict, rank2_dict = \
        evaluate_epoch(valid_iterator, models, num_personas, gradient_accumulation_steps, device, dataset, 0,
                       apply_interaction, matching_method, aggregation_method, tokenizer, valid_info, num_candidates, rank1_k, rank2_k)
    cprint("test loss: {0:.4f}, test acc: {1:.4f}, test recall: {2}, test MRR: {3:.4f}"
        .format(valid_loss, valid_acc, valid_recall, valid_MRR))

    directory = "./ranking_dataset"
    try:
        if not os.path.exists(directory):
            os.makedirs(directory)
    except OSError:
        print('Error: Creating directory. ' + directory)

    print(len(rank1_dict['context']))
    print(len(rank1_dict['response']))
    print(len(rank1_dict['logits']))
    print(len(rank1_dict['persona']))
    print(len(rank1_dict['info']))

    df = pd.DataFrame(rank1_dict)
    df.to_csv("./ranking_dataset/rank_1.tsv", sep='\t', na_rep='NaN')
    with open("./ranking_dataset/rank_1.json", "w") as p:
        json.dump(rank1_dict, p)

    df = pd.DataFrame(rank2_dict)
    df.to_csv("./ranking_dataset/rank_2.tsv", sep='\t', na_rep='NaN')
    with open("./ranking_dataset/rank_2.json", "w") as p:
        json.dump(rank2_dict, p)

    cached_personaChat_file_path = personaChat_file_path.replace(".json", "_ranked_" + str(rank2_k) + ".json")
    with open(cached_personaChat_file_path, 'w') as p:
        data = afterProcess(personaChat_file_path, rank2_dict, rank2_k)
        json.dump(data, p)
    print("ranking finished")
    sys.exit()

def afterProcess(read_path, dict, rank_k):
    with open(read_path, 'r') as p:
        data = json.load(p)
        idx = 0
        info = dict['info']
        print(len(info))
        for k in range(len(info)):
            i, j, num_exp_persona, pair, tv = info[idx][0], info[idx][1], info[idx][2], info[idx][3], info[idx][4]
            end = min(rank_k, num_exp_persona)
            if idx + end >= len(info):
                break
            if tv == 't':
                aug = []
                for l in range(end):
                    aug.append({'exp_persona': data['train'][i]['utterances'][j]['exp_persona'][info[idx + l][3]],
                                'new_responses': dict['response'][idx + l]})
                del data['train'][i]['utterances'][j]['new_responses']
                del data['train'][i]['utterances'][j]['exp_persona']
                data['train'][i]['utterances'][j]['aug'] = aug
                idx += end
            elif tv == 'v':
                aug = []
                for l in range(end):
                    aug.append({'exp_persona': data['valid'][i]['utterances'][j]['exp_persona'][info[idx + l][3]],
                                'new_responses': dict['response'][idx + l]})
                del data['valid'][i]['utterances'][j]['new_responses']
                del data['valid'][i]['utterances'][j]['exp_persona']
                data['valid'][i]['utterances'][j]['aug'] = aug
                idx += end
        for dialogue in data['train']:
            for context in dialogue['utterances']:
                if "exp_persona" in context or "new_responses" in context and "aug" not in context:
                    try:
                        del context["exp_persona"]
                    except:
                        pass
                    try:
                        del context["new_responses"]
                    except:
                        pass
                    context['aug'] = [{"exp_persona": "", "new_responses": ""}]
        for dialogue in data['valid']:
            for context in dialogue['utterances']:
                if "exp_persona" in context or "new_responses" in context and "aug" not in context:
                    try:
                        del context["exp_persona"]
                    except:
                        pass
                    try:
                        del context["new_responses"]
                    except:
                        pass
                    context['aug'] = [{"exp_persona": "", "new_responses": ""}]
    return data

def clean_config(configs):
    cleaned_configs = []
    for config in configs:
        if config not in cleaned_configs:
            cleaned_configs.append(config)
    return cleaned_configs

def merge_metrics(metrics):
    avg_metrics = {"score" : 0}
    num_metrics = len(metrics)
    for metric in metrics:
        for k in metric:
            if k != "config":
                avg_metrics[k] += np.array(metric[k])

    for k, v in avg_metrics.items():
        avg_metrics[k] = (v/num_metrics).tolist()

    return avg_metrics

if __name__ == "__main__":
    mp.set_start_method('spawn', force=True)
    parser = argparse.ArgumentParser(description="Model for Transformer-based Dialogue Generation with Controlled Emotion")
    parser.add_argument('--config', help='Config to read details', required=True)
    parser.add_argument('--note', help='Experiment note', default="")
    args = parser.parse_args()
    cprint("Experiment note: ", args.note)
    with open(args.config) as configfile:
        config = json.load(configfile) # config is now a python dict

    # pass experiment config to main
    parameters_to_search = OrderedDict() # keep keys in order
    other_parameters = {}
    keys_to_omit = ["kernel_sizes"] # keys that allow a list of values
    for k, v in config.items():
        # if value is a list provided that key is not device, or kernel_sizes is a nested list
        if isinstance(v, list) and k not in keys_to_omit:
            parameters_to_search[k] = v
        elif k in keys_to_omit and isinstance(config[k], list) and isinstance(config[k][0], list):
            parameters_to_search[k] = v
        else:
            other_parameters[k] = v

    if len(parameters_to_search) == 0:
        config_id = time.perf_counter()
        config["config_id"] = config_id
        cprint(config)
        output = main(config, progress=1)
        cprint("-"*80)
        cprint(output["config"])
        cprint(output["epoch"])
        cprint(output["score"])
        cprint(output["recall"])
        cprint(output["MRR"])
    else:
        all_configs = []
        for i, r in enumerate(itertools.product(*parameters_to_search.values())):
            specific_config = {}
            for idx, k in enumerate(parameters_to_search.keys()):
                specific_config[k] = r[idx]

            # merge with other parameters
            merged_config = {**other_parameters, **specific_config}
            all_configs.append(merged_config)

        #   cprint all configs
        for config in all_configs:
            config_id = time.perf_counter()
            config["config_id"] = config_id
            logging.critical("config id: {0}".format(config_id))
            cprint(config)
            cprint("\n")

        # multiprocessing
        num_configs = len(all_configs)
        # mp.set_start_method('spawn')
        pool = mp.Pool(processes=config["processes"])
        results = [pool.apply_async(main, args=(x,i/num_configs)) for i,x in enumerate(all_configs)]
        outputs = [p.get() for p in results]

        # if run multiple models using different seed and get the averaged result
        if "seed" in parameters_to_search:
            all_metrics = []
            all_cleaned_configs = clean_config([output["config"] for output in outputs])
            for config in all_cleaned_configs:
                metrics_per_config = []
                for output in outputs:
                    if output["config"] == config:
                        metrics_per_config.append(output)
                avg_metrics = merge_metrics(metrics_per_config)
                all_metrics.append((config, avg_metrics))
            # log metrics
            cprint("Average evaluation result across different seeds: ")
            for config, metric in all_metrics:
                cprint("-"*80)
                cprint(config)
                cprint(metric)

            # save to log
            with open("./log/{0}.txt".format(time.perf_counter()), "a+") as f:
                for config, metric in all_metrics:
                    f.write(json.dumps("-"*80) + "\n")
                    f.write(json.dumps(config) + "\n")
                    f.write(json.dumps(metric) + "\n")

        else:
            for output in outputs:
                cprint("-"*80)
                cprint(output["config"])
                cprint(output["score"])
                cprint(output["recall"])
                cprint(output["MRR"])
                cprint("Best result at epoch {0}: ".format(output["epoch"]))
                cprint(output["recall"][output["epoch"]], output["MRR"][output["epoch"]])

            # save to log
            with open("./log/{0}.txt".format(time.perf_counter()), "a+") as f:
                for output in outputs:
                    f.write(json.dumps("-"*80) + "\n")
                    f.write(json.dumps(output) + "\n")


import torch.nn.functional as F
from optimize import Optimizer

def train(args, model, device, train_loader, epoch):
    model.train()

    optimizer = Optimizer(model)

    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)

        output, out_fc0 = model(data)
        outputs = [out_fc0, output]
        loss = F.nll_loss(output, target)

        optimizer.step(loss, outputs)

        if batch_idx % args.log_interval == 0:
            print(
                f"Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({format(100.0 * batch_idx / len(train_loader), '.1f')}%)]\tLoss: {format(loss.item(), '.6f')}"
            )
            if args.dry_run:
                break


"""
This module provides services which are used
to send requests directly to the database
"""

from fastapi.encoders import jsonable_encoder
from datetime import datetime
from pymongo import MongoClient, TEXT
from pydantic import EmailStr

from services.common import update_dict
from schemas.mlflow_server import (
    MlflowServer,
    MlflowServerPut,
    MlflowServerPatch,
    CommonMeta,
)

class MlflowServerService:
    COLLECTION_NAME = "mlflow_servers"

    @classmethod
    async def _create_indexes(cls, db: MongoClient) -> None:
        await db[MlflowServerService.COLLECTION_NAME].create_index(
            [("name", TEXT)],
            name="unique_name",
            default_language="english",
            unique=True,
        )

    @classmethod
    async def get_mlflow_servers(
        cls, db: MongoClient, skip: int = 0, limit: int = 100
    ) -> list[MlflowServer]:
        """
        Returns a list of mlflow server data, with optional pagination

        :param db: Database session
        :param skip: (optional) the number of records to skip (default: 0)
        :param limit: (optional) the maximum number of records to retrieve
        (default: 100)

        :return: a list of mlflow server data, where skip < mlflow_server_id < limit
        """
        mlflow_servers = (
            await db[MlflowServerService.COLLECTION_NAME]
            .find()
            .skip(skip)
            .limit(limit)
            .to_list(limit)
        )
        return [MlflowServer(**mlflow_server) for mlflow_server in mlflow_servers]

    @classmethod
    async def get_mlflow_server_by_name(
        cls, db: MongoClient, name: str
    ) -> MlflowServer | None:
        """
        Returns the mlflow server data found by gate name

        :param db: Database session
        :param name: the mlflow server name to retrieve

        :return: the mlflow server data corresponding to the given name
        or None if not found
        """
        mlflow_server = await db[MlflowServerService.COLLECTION_NAME].find_one(
            {"name": name}
        )
        return MlflowServer(**mlflow_server) if mlflow_server else None

    @classmethod
    async def put_mlflow_server(
        cls,
        db: MongoClient,
        data: MlflowServerPut,
        user: EmailStr,
    ) -> MlflowServer:
        """
        Inserts a new mlflow server record into the database

        :param db: Database session
        :param data: the mlflow server data to insert

        :return: the newly-inserted mlflow server record
        """
        creation_time = datetime.utcnow()
        db_mlflow_server = MlflowServer(
            **data.dict(exclude=set(["meta"])),
            meta=CommonMeta(
                created_by=user,
                created_at=creation_time,
                updated_by=user,
                updated_at=creation_time,
                tags=data.meta.tags,
            )
        )
        result = await db[MlflowServerService.COLLECTION_NAME].insert_one(
            jsonable_encoder(db_mlflow_server)
        )
        return MlflowServer(
            **(
                await db[MlflowServerService.COLLECTION_NAME].find_one(
                    {"_id": result.inserted_id}
                )
            )
        )

    @classmethod
    async def patch_mlflow_server(
        cls,
        db: MongoClient,
        name: str,
        data: MlflowServerPatch,
        user: EmailStr,
    ) -> MlflowServer:
        """
        Updates an existing mlflow server record in the database

        :param db: Database session
        :param name: the mlflow server name to patch
        :param data: the mlflow server data to update

        :return: the updated mlflow server record
        """
        db_mlflow_server = await db[MlflowServerService.COLLECTION_NAME].find_one(
            {"name": name}
        )
        db_mlflow_server.pop("_id")

        update_dict(db_mlflow_server, data.dict(exclude_unset=True))
        db_mlflow_server["meta"]["updated_by"] = user
        db_mlflow_server["meta"]["updated_at"] = datetime.utcnow()
        db_mlflow_server = await db[MlflowServerService.COLLECTION_NAME].update_one(
            {"name": name}, {"$set": jsonable_encoder(db_mlflow_server)}
        )
        return MlflowServer(
            **(await db[MlflowServerService.COLLECTION_NAME].find_one({"name": name}))
        )

    @classmethod
    async def delete_mlflow_server(cls, db: MongoClient, name: str) -> None:
        """
        Deletes a mlflow server record from the database

        :param db: Database sessioner
        :param name: the mlflow server name to delete

        :return: a json with a "detail" key indicating success
        """
        await db[MlflowServerService.COLLECTION_NAME].delete_one({"name": name})


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
""" Build, test and install configuration """

import os
import sys
from glob import glob
from setuptools import setup, find_packages, Command

class CheckPylint(Command):
    """ setup.py command checking code-correctness with pylint and pep8 """
    user_options = []
    description = "Check code using pylint and pep8"

    def initialize_options(self):
        pass

    def finalize_options(self):
        pass

    def run(self):
        files = glob('*.py') + glob('*/*.py')

        output_format = sys.stdout.isatty() and "colorized" or "text"

        cmd = "/usr/bin/pylint "
        cmd += "--output-format=%s " % output_format
        cmd += " ".join(files)
        os.system(cmd + " --rcfile tests/pylint.cfg")

        print("running pep8")
        cmd = "/usr/bin/pep8 "
        cmd += " ".join(files)
        os.system(cmd + " --config tests/pep8.cfg")

setup(
    name='washred',
    version='0.0.2',
    description='Washington Redskins',
    long_description=''.join(open('README.rst').readlines()),
    keywords='git, github, gerrit, gerrithub, tryout',
    author='Martin Kletzander',
    author_email='nert.pinx+github@gmail.com',
    license='GPLv3',
    packages=find_packages(),
    classifiers=[
        'Development Status :: 3 - Alpha',
        'Intended Audience :: Developers',
        'License :: OSI Approved :: GNU General Public License v3 (GPLv3)',
        'Operating System :: POSIX :: Linux',
        'Programming Language :: Python',
        ],
    scripts=['washred.py'],
    cmdclass={'pylint': CheckPylint},
)


# -*- coding: utf-8  -*-

from __future__ import absolute_import
from __future__ import unicode_literals
from __future__ import print_function

from pywiktionary.wiktionary import Wiktionary
import json
import argparse

def json_option(parser):
	parser.add_argument(
		'--json', action='store_true',
		help='Output in machine readable json. This may contain additional information.')

def cli():
	# Make "wikitionary | cat" work with unicode output
	# http://stackoverflow.com/questions/2276200/changing-default-encoding-of-python#17628350
	# An alternative approach is to use PYTHONIOENCODING
	import sys
	import io
	sys.stdout = io.open(sys.stdout.fileno(), mode='w', encoding='utf8', buffering=1)

	wiki = Wiktionary()
	PARSER = argparse.ArgumentParser(description='Fetch information from wikitionary')
	PARSER.add_argument('word', type=str)
	PARSER.add_argument(
		'--language',
		type=str.title,
		help='Display entry for this language (the default is English)')

	json_option(PARSER)
	args = PARSER.parse_args()

	args.language = "English" if args.language is None else args.language
	result = wiki.lookup(args.word, lang=args.language)

#	if args.language == 'All':
#		language_entries = result.items()
#	elif args.language:
#		language_entries = [(args.language, result[args.language])]
#	elif 'English' in result:
#		language_entries = [('English', result['English'])]
#	else:
#		language_entries = result.items()
	language_entries = [(args.language, result[args.language])]

	if args.json:
		print(json.dumps(result, indent=4))
	else:
		for language, language_entry in language_entries:
			pronunciation = format_pronunciation(language_entry)
			if pronunciation:
				print(language)
				print(indent(pronunciation))

def indent(s):
	return '\n'.join(['    ' + l for l in  s.split('\n')])

def format_pronunciation(entry):
	if 'Pronunciation' not in entry:
		return None

	result = []

	for pronunciation in entry['Pronunciation']:
		if not 'IPA' in pronunciation:
			continue
		else:
			accent = pronunciation.get('Accent', 'Standard')
			if isinstance(accent, list):
				accent = ', '.join(accent)
			result.append(accent)
			variants, _ = pronunciation['IPA']
			for variant in variants:
				result.append('    ' + variant)

	if not result:
		return None

	return '\n'.join(result)


import argparse
import os
import tqdm
import numpy as np
from kitti_transformations import read_kitti_calibration_file, project_velo_to_image

MASRKCNN_TO_KITTI_CLASSES = {
    'person': 'Pedestrian',
    'car': 'Car'
}
USED_CLASSES = set(MASRKCNN_TO_KITTI_CLASSES.keys())
VELO_X_UNIT_VECTOR = np.array([1.0, 0.0])

def read_pointcloud(pointcloud_file_path: str) -> np.ndarray:
    points = np.fromfile(pointcloud_file_path, dtype=np.float32)
    return points.reshape(-1, 4)

def filter_forward_points(points: np.ndarray) -> np.ndarray:
    return points[points[:, 0] >= 0.0]

def save_frustum_points(labelled_points: np.ndarray, class_name: str, output_path: str):
    np.savez(output_path, points=labelled_points, class_name=class_name)

def get_arguments() -> dict:
    parser = argparse.ArgumentParser()

    parser.add_argument(
        'detections', type=str,
        help='Path to directory containing all applicable detections'
    )

    parser.add_argument(
        'kitti', type=str,
        help='Path to base KITTI data directory that contains folders for images, pointclouds, etc.'
    )

    parser.add_argument(
        'output', type=str,
        help='Path where to save KITTI format labels from input detections'
    )

    return parser.parse_args()

if __name__ == '__main__':
    args = get_arguments()

    detections_dir = args.detections
    kitti_dir = args.kitti
    output_dir = args.output
    os.makedirs(output_dir, exist_ok=True)

    pointcloud_dir = os.path.join(kitti_dir, 'velodyne')
    calibration_dir = os.path.join(kitti_dir, 'calib')

    detection_files = os.listdir(detections_dir)

    running_detection_ind = 0

    for scene in tqdm.tqdm(detection_files):

        detections_file_path = os.path.join(detections_dir, scene)
        detections = np.load(detections_file_path)

        masks = detections['masks']
        classes = detections['classes']
        bboxes = detections['bboxes']
        scores = detections['scores']
        if len(classes) == 0:
            continue

        scene_id = scene.split('.')[0]

        calibration_file_path = os.path.join(calibration_dir, scene_id + '.txt')
        calibration_data = read_kitti_calibration_file(calibration_file_path)

        pointcloud_file_path = os.path.join(pointcloud_dir, scene_id + '.bin')
        points = read_pointcloud(pointcloud_file_path)
        points = filter_forward_points(points)
        points = points[:, :3]

        original_point_indices = np.arange(0, len(points))

        projected_points = project_velo_to_image(points, calibration_data)  # Points are x, y
        projected_points = np.rint(projected_points).astype(np.int32)

        points = np.c_[points, np.zeros(len(points))]

        image_x_len = masks.shape[1]
        image_y_len = masks.shape[0]

        points_on_image = (projected_points[:, 0] >= 0) & (projected_points[:, 0] < image_x_len) \
            & (projected_points[:, 1] >= 0) & (projected_points[:, 1] < image_y_len)
        projected_points = projected_points[points_on_image]
        original_point_indices = original_point_indices[points_on_image]

        for ind in range(masks.shape[2]):
            class_name = classes[ind]
            if class_name not in USED_CLASSES:
                continue

            current_points = points.copy()
            mask = masks[:, :, ind]
            bbox_2d = bboxes[ind]

            # Label points for which the projection is on the mask
            points_matching_mask = mask[projected_points[:, 1], projected_points[:, 0]]
            mask_points_original_ind = original_point_indices[points_matching_mask]
            current_points[mask_points_original_ind, 3] = 1.0

            # Keep only the points for which the projection falls into the bounding box
            y_min, x_min, y_max, x_max = bbox_2d
            projected_points_in_bbox = (projected_points[:, 0] >= x_min) & (projected_points[:, 0] <= x_max) & \
                (projected_points[:, 1] >= y_min) & (projected_points[:, 1] <= y_max)

            bbox_points_original_ind = original_point_indices[projected_points_in_bbox]
            frustum_points = current_points[bbox_points_original_ind]

            if y_max - y_min < 25 or len(frustum_points) < 100:
                continue

            # Sanity check
            assert len(projected_points_in_bbox) == len(points_matching_mask), 'Different amount of mask points in projected points data'
            assert projected_points_in_bbox.sum() >= points_matching_mask.sum(), 'Bounding box has fewer points than mask'
            assert (points_matching_mask & projected_points_in_bbox).sum() == points_matching_mask.sum(), 'Mask is not subset of 2D bounding box'

            output_filename = f'{scene_id}_{ind}'
            output_filepath = os.path.join(output_dir, output_filename)
            save_frustum_points(frustum_points, class_name, output_filepath)

    print('Done!')


from pathlib import Path

from aiogram import Bot, Dispatcher
from aiogram.types import BotCommand, BotCommandScope, BotCommandScopeAllPrivateChats
from aiogram.client.session.aiohttp import AiohttpSession
from ujson import loads, dumps
from celery import Celery

from .types import Settings

BASE_DIR = Path(__file__).resolve().parent.parent
SETTINGS = Settings()

bot = Bot(
    token=SETTINGS.BOT_TOKEN.get_secret_value(),
    session=AiohttpSession(json_loads=loads, json_dumps=dumps),
    parse_mode='HTML'
)

dp = Dispatcher(disable_fsm=True)

URLS = (
    ('/filter?sort=4', 'https://cars.av.by'),
    ('/filter?page=2&sort=4', 'https://cars.av.by'),
    ('/filter?sort=4', 'https://truck.av.by'),
    ('/filter?sort=4', 'https://moto.av.by'),
    ('/filter?page=2&sort=4', 'https://moto.av.by'),
    ('/filter?sort=4', 'https://bus.av.by'),
)

celery = Celery()
celery.config_from_object(SETTINGS, namespace='CELERY')  # Ищет необходимые ему аттрибуты в объхекте settings которые начинаются с CELERY
celery.autodiscover_tasks(packages=['src'])

celery.conf.beat_schedule = {
    'parser': {
        'task': 'src.tasks.task_parser',
        'schedule': 1800
    }
}


import torch
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
from classifier.simpleClassifier import simpleClassifier
from classifier.vggClassifier import vggClassifier
import os
from collections import OrderedDict
import torch
from torch.autograd import Variable
import numpy as np

def load_model(model, path, name, mode = "parallel"):
	state_dict = torch.load(os.path.join(path, name))["model"]
	new_state_dict = OrderedDict()

	for k, v in state_dict.items():
		name = ""
		if mode == "single" and k.startswith("module."):
			name = k[7:]
		elif mode == "parallel" and not k.startswith("module."):
			name = "module."+k
		else:
			name = k
		new_state_dict[name] = v

	model.load_state_dict(new_state_dict)

def feat_ext(ext_type, cuda = True):
	temp_model = None
	model = None
	num_layers = 0

	if ext_type == "simpleClassifier42":
		temp_model = simpleClassifier()
		load_model(temp_model, "/data/classifier/trainedModels/", "default_modelsimpleClassifier42.pth", mode="single")
		num_layers = 3

	elif ext_type == "simpleClassifier58":
		temp_model = simpleClassifier()
		load_model(temp_model, "/data/classifier/trainedModels/", "default_modelsimpleClassifier58.pth", mode="single")
		num_layers = 3

	elif ext_type == "vggClassifier67":
		temp_model = vggClassifier()
		load_model(temp_model, "/data/classifier/trainedModels/", "default_modelvggClassifier67.pth", mode="single")
		num_layers = 3

	for param in temp_model.parameters():
		param.requires_grad = False

	model = torch.nn.Sequential(*(temp_model.features[i] for i in range(num_layers)))

	if cuda:
		model.cuda()

	return model

def applyFeatureTransformation(img, ext_type):
	feat_extractor = feat_ext(ext_type, cuda = True)
	img = img.astype(np.float32)
	img = np.expand_dims(img, axis=2)
	img = np.transpose(img, [2,0,1])
	img = np.expand_dims(img, axis=0)
	img = Variable(torch.from_numpy(img))
	img = img.cuda()
	feat_extractor = feat_ext(ext_type, cuda=True)
	feature_space_img = feat_extractor(img)
	feature_space_img = (feature_space_img.data).cpu().numpy()
	feature_space_img = feature_space_img[0,0,:,:]
	return feature_space_img


""" VCL Coreset Module"""
import numpy as np
import torch
import torch.utils.data as data
import torch.optim as optim
import torchvision.transforms as transforms

import trainer
import mnist

def fit(args, model, device, optimizer, loss_fn,  coresets, dataset, labels_list, task_id):
    # Dataloader
    train_loader = trainer.get_loader(mnist.getTrain(dataset), args, device, 'train')
    val_loader = trainer.get_loader(mnist.getVal(dataset), args, device, 'val')
    # Log Best Accuracy
    best_val_acc = 0
    all_test_accs = []
    # Early Stopping
    early_stop = 0

    # Training loop
    for epoch in range(1, args.epochs + 1):
        # Prepare model for current task
        model.set_range(labels_list[task_id])
        trainer.train(args, model, device, train_loader, optimizer, epoch, loss_fn)
        _, val_acc = trainer.test(args, model, device, val_loader, loss_fn, val=True)
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            best_state = model.state_dict()
            early_stop = 0
        else:
            early_stop += 1
            if early_stop >= args.early_stop_after:
                break

        # Evaluate all tasks on test set
        test_accs = test_all_tasks(coresets, args, model, device,
            loss_fn, labels_list, 'test')
        all_test_accs.append(test_accs)

    return best_state, all_test_accs

def get_dataset(args, task_id, split):
    trans = transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                       ])
    if args.dataset=='splitMNIST':
        Dataset = mnist.SplitMNIST
    elif args.dataset=='fashionMNIST':
        Dataset = mnist.FashionMNIST
    elif args.dataset=='permutedMNIST':
        Dataset = mnist.PermutedMNIST

    if split == 'train':
        trainval = Dataset('data', task_id, train=True, download=True, transform=trans)
        dataset = mnist.getTrain(trainval)
    elif split == 'val':
        trainval = Dataset('data', task_id, train=True, download=True, transform=trans)
        dataset = mnist.getVal(trainval)
    elif split == 'trainval':
        dataset = Dataset('data', task_id, train=True, download=True, transform=trans)
    else:
        dataset = Dataset('data', task_id, train=False, download=True, transform=trans)
    return dataset

def test_all_tasks(coresets, args, model, device, loss_fn, labels_list, split):
    def fit_coreset(coreset, labels, args, model, device):
        # Get Inference Model
        final_model, _ = trainer.get_model(args, device)
        final_model.load_state_dict(model.state_dict())
        final_model.set_range(labels)

        # Get Loss Function and Optimizer and Dataloader
        loss_fn = trainer.get_loss_fn(args, device, final_model, model)
        optimizer = optim.Adam(final_model.parameters(), lr=args.lr)
        coreset_loader = trainer.get_loader(coreset, args, device, 'coreset')

        for epoch in range(args.coreset_epochs):
            trainer.train(args, final_model, device, coreset_loader, optimizer, epoch, loss_fn, verbose=False)

        return final_model

    # Coreset Evaluation Loop
    test_accs = np.zeros(len(labels_list))
    for j, labels in enumerate(labels_list):
        if args.coreset_size > 0 and len(coresets) > j:
            # Get Inference Model by Tuning on Coreset
            final_model = fit_coreset(coresets[j], labels, args, model, device)
        else:
            # No Coreset
            final_model = model

        # Evaluate on the test set
        testset = get_dataset(args, j, split)
        test_loader = trainer.get_loader(testset, args, device, split)
        final_model.set_range(labels)
        _, test_accs[j] = trainer.test(args, final_model, device, test_loader, loss_fn)

    return test_accs

def run_vcl(args, device, labels_list):
    task_final_accs = np.zeros((5, 5)) # Test accuracy after each task ends
    all_accs = [] # Test accuracy after every epoch

    coresets = [] # Coresets
    coreset_method = Coreset.rand_from_batch

    # Pretraining
    print ("=============Pretraining ==================")
    model, name = trainer.get_model(args, device, mle=True)
    loss_fn = trainer.get_loss_fn(args, device, model)
    optimizer = optim.Adam(model.parameters(), lr=args.lr)
    dataset = get_dataset(args, 0, 'trainval')
    fit(args, model, device, optimizer, loss_fn, coresets, dataset, labels_list, 0)

    # train all tasks incrementally
    for task_id, labels in enumerate(labels_list):
        print ("===========TASK {}: {}=============".format(task_id + 1, labels))
        # Model
        old_model = model
        model, model_name = trainer.get_model(args, device)
        model.load_state_dict(old_model.state_dict())
        # Loss Function and Optimizer
        loss_fn = trainer.get_loss_fn(args, device, model, old_model)
        optimizer = optim.Adam(model.parameters(), lr=args.lr)

        # Dataset and Coreset
        dataset = get_dataset(args, task_id, 'trainval')
        if coreset_method is not None:
            coresets, dataset = coreset_method(coresets, dataset, args.coreset_size)

        # Fit
        best_state, test_accs = fit(args, model, device, optimizer, loss_fn, coresets, dataset, labels_list, task_id)
        task_final_accs[task_id] = test_accs[-1]
        all_accs.extend(test_accs)

        # save model
        path = trainer.get_model_path(labels, model_name)
        torch.save(best_state, path)

    return task_final_accs, all_accs

class Coreset():
    """ Random coreset selection """
    @staticmethod
    def rand_from_batch(coresets, dataset, coreset_size):
        sizes = [coreset_size, len(dataset) - coreset_size]
        new_coreset, dataset = data.random_split(dataset, sizes)
        coresets.append(new_coreset)

        return coresets, dataset

    """ K-center coreset selection """
    @staticmethod
    def k_center(coresets, dataset, coreset_size):
        # Select K centers from dataset and add to current coreset
        dists = np.full(len(dataset), np.inf)
        current_id = 0
        dists = Coreset.update_distance(dists, dataset, current_id)
        idx = [ current_id ]

        for i in range(1, coreset_size):
            current_id = np.argmax(dists)
            dists = Coreset.update_distance(dists, dataset, current_id)
            idx.append(current_id)

        keep_idx = [i for i in range(len(dataset)) if i not in idx]
        new_coreset = data.Subset(dataset, idx)
        coresets.append(new_coreset)
        dataset = data.Subset(dataset, keep_idx)

        return coresets, dataset

    @staticmethod
    def update_distance(dists, dataset, current_id):
        current_datapoint, _ = dataset[current_id]
        for i in range(len(dataset)):
            data, _ = dataset[i]
            current_dist = np.linalg.norm(data - current_datapoint)
            dists[i] = np.minimum(current_dist, dists[i])
        return dists


#!/usr/bin/python3

# This is the top-level BREW V1 CPU

import sys
from pathlib import Path
import itertools

sys.path.append(str(Path(__file__).parent / ".." / ".." / ".." / "silicon"))
sys.path.append(str(Path(__file__).parent / ".." / ".." / ".." / "silicon" / "unit_tests"))

try:
    from .brew_types import *
    from .brew_utils import *

    from .pipeline import Pipeline
    from .bus_if import BusIf
    from .cpu_dma import CpuDma
    from .synth import *
    from .assembler import *
    from .apb_timer import ApbSimpleTimer

except ImportError:
    from brew_types import *
    from brew_utils import *

    from pipeline import Pipeline
    from bus_if import BusIf
    from cpu_dma import CpuDma
    from synth import *
    from assembler import *
    from apb_timer import ApbSimpleTimer

class BrewV1Top(GenericModule):
    clk               = ClkPort()
    rst               = RstPort()

    # DRAM interface
    dram              = Output(ExternalBusIf)

    # External dma-request
    drq               = Input(Unsigned(4))

    n_int             = Input(logic)

    def construct(self, nram_base: int = 0x0, has_multiply: bool = True, has_shift: bool = True, page_bits: int = 7):
        self.nram_base = nram_base
        self.has_multiply = has_multiply
        self.has_shift = has_shift
        self.page_bits = page_bits

        self.csr_cpu_task_mode_page      = 0x8000
        self.csr_cpu_scheduler_mode_page = 0x0000

        self.csr_mach_arch      = 0x00
        self.csr_capability     = 0x01
        self.csr_pmem_base_ofs  = 0x80
        self.csr_pmem_limit_ofs = 0x81
        self.csr_dmem_base_ofs  = 0x82
        self.csr_dmem_limit_ofs = 0x83
        self.csr_ecause_ofs     = 0x00
        self.csr_eaddr_ofs      = 0x01

        self.csr_mach_arch_reg  = self.csr_cpu_task_mode_page + self.csr_mach_arch
        self.csr_capability_reg = self.csr_cpu_task_mode_page + self.csr_capability
        self.csr_pmem_base_reg  = self.csr_cpu_scheduler_mode_page + self.csr_pmem_base_ofs
        self.csr_pmem_limit_reg = self.csr_cpu_scheduler_mode_page + self.csr_pmem_limit_ofs
        self.csr_dmem_base_reg  = self.csr_cpu_scheduler_mode_page + self.csr_dmem_base_ofs
        self.csr_dmem_limit_reg = self.csr_cpu_scheduler_mode_page + self.csr_dmem_limit_ofs
        self.csr_ecause_reg     = self.csr_cpu_scheduler_mode_page + self.csr_ecause_ofs
        self.csr_eaddr_reg      = self.csr_cpu_scheduler_mode_page + self.csr_eaddr_ofs

    def body(self):
        bus_if = BusIf(nram_base=self.nram_base)
        pipeline = Pipeline(has_multiply=self.has_multiply, has_shift=self.has_shift, page_bits=self.page_bits)
        dma = CpuDma()
        timer = ApbSimpleTimer()

        # Things that need CSR access
        ecause     = Wire(EnumNet(exceptions))
        ecause_clear_pulse = Wire()
        eaddr      = Wire(BrewAddr)
        pmem_base  = Wire(BrewMemBase)
        pmem_limit = Wire(BrewMemBase)
        dmem_base  = Wire(BrewMemBase)
        dmem_limit = Wire(BrewMemBase)

        self.csrs_cpu_task_mode = {
            self.csr_mach_arch:  RegMapEntry("csr_mach_arch",  (RegField("32'b0",access="R"),), "Machine architecture register"),
            self.csr_capability: RegMapEntry("csr_capability", (RegField("32'b0",access="R"),), "Machine capability register")
        }

        self.csrs_cpu_scheduler_mode = {
            self.csr_ecause_ofs:     RegMapEntry("csr_ecause",     (RegField(ecause,access="R"),), "Exception cause register", read_pulse=ecause_clear_pulse),
            self.csr_eaddr_ofs:      RegMapEntry("csr_eaddr",      (RegField(eaddr,access="R"),), "Exception address register"),

            self.csr_pmem_base_ofs:  RegMapEntry("csr_pmem_base",  (RegField(pmem_base,  start_bit=10),), "Program memory base register"),
            self.csr_pmem_limit_ofs: RegMapEntry("csr_pmem_limit", (RegField(pmem_limit, start_bit=10),), "Program memory limit register"),
            self.csr_dmem_base_ofs:  RegMapEntry("csr_dmem_base",  (RegField(dmem_base,  start_bit=10),), "Data memory base register"),
            self.csr_dmem_limit_ofs: RegMapEntry("csr_dmem_limit", (RegField(dmem_limit, start_bit=10),), "Data memory limit register"),
        }

        fetch_to_bus = Wire(BusIfRequestIf)
        bus_to_fetch = Wire(BusIfResponseIf)
        mem_to_bus = Wire(BusIfRequestIf)
        bus_to_mem = Wire(BusIfResponseIf)
        dma_to_bus = Wire(BusIfDmaRequestIf)
        bus_to_dma = Wire(BusIfDmaResponseIf)
        csr_if = Wire(CsrIf)
        self.cpu_task_mode_csr_if = Wire(CsrIf)
        self.cpu_scheduler_mode_csr_if = Wire(CsrIf)
        bus_if_reg_if = Wire(CsrIf)
        dma_reg_if = Wire(CsrIf)
        timer_reg_if = Wire(CsrIf)

        # BUS INTERFACE
        ###########################
        bus_if.fetch_request <<= fetch_to_bus
        bus_to_fetch <<= bus_if.fetch_response
        bus_if.mem_request <<= mem_to_bus
        bus_to_mem <<= bus_if.mem_response
        bus_if.dma_request <<= dma_to_bus
        bus_to_dma <<= bus_if.dma_response

        self.dram <<= bus_if.dram

        bus_if.reg_if <<= bus_if_reg_if

        # DMA
        ###########################
        dma_to_bus <<= dma.bus_req_if
        dma.bus_rsp_if <<= bus_to_dma

        dma.drq <<= self.drq

        dma.reg_if <<= dma_reg_if

        # Timer
        ############################
        timer.bus_if <<= timer_reg_if

        # PIPELINE
        ############################
        fetch_to_bus <<= pipeline.fetch_to_bus
        pipeline.bus_to_fetch <<= bus_to_fetch
        mem_to_bus <<= pipeline.mem_to_bus
        pipeline.bus_to_mem <<= bus_to_mem
        csr_if <<= pipeline.csr_if

        pipeline.ecause_clear_pulse <<= ecause_clear_pulse
        ecause <<= pipeline.ecause
        eaddr  <<= pipeline.eaddr
        pipeline.pmem_base  <<= pmem_base
        pipeline.pmem_limit <<= pmem_limit
        pipeline.dmem_base  <<= dmem_base
        pipeline.dmem_limit <<= dmem_limit

        pipeline.interrupt <<= ~self.n_int | ~timer.n_int

        event_fetch_wait_on_bus = pipeline.event_fetch_wait_on_bus
        event_decode_wait_on_rf = pipeline.event_decode_wait_on_rf
        event_mem_wait_on_bus   = pipeline.event_mem_wait_on_bus
        event_branch_taken      = pipeline.event_branch_taken
        event_branch            = pipeline.event_branch
        event_load              = pipeline.event_load
        event_store             = pipeline.event_store
        event_execute           = pipeline.event_execute
        event_bus_idle          = bus_if.event_bus_idle
        event_fetch             = pipeline.event_fetch
        event_fetch_drop        = pipeline.event_fetch_drop
        event_inst_word         = pipeline.event_inst_word

        # CSR address decode
        #############################
        csr_cpu_task_mode_psel      = csr_if.psel & (csr_if.paddr[15:8] == 0x80)
        csr_cpu_scheduler_mode_psel = csr_if.psel & (csr_if.paddr[15:8] == 0x00)
        csr_event_psel              = csr_if.psel & (csr_if.paddr[15:8] == 0x81)
        csr_bus_if_psel             = csr_if.psel & (csr_if.paddr[15:8] == 0x02)
        csr_dma_psel                = csr_if.psel & (csr_if.paddr[15:8] == 0x03)
        csr_timer_psel              = csr_if.psel & (csr_if.paddr[15:8] == 0x04)

        top_level_prdata = Wire(Unsigned(32))
        top_level_pready = Wire(logic)

        # CSR bus routing
        #############################
        dma_reg_if.pwrite  <<= csr_if.pwrite
        dma_reg_if.psel    <<= csr_dma_psel
        dma_reg_if.penable <<= csr_if.penable
        dma_reg_if.paddr   <<= csr_if.paddr[4:0] # FIXME: This used to be 3:0 which generated a silent error: There are actually 17 decoded registers, yet we happily generated the decoder logic with a 4-bit address
        dma_reg_if.pwdata  <<= csr_if.pwdata

        bus_if_reg_if.pwrite  <<= csr_if.pwrite
        bus_if_reg_if.psel    <<= csr_bus_if_psel
        bus_if_reg_if.penable <<= csr_if.penable
        bus_if_reg_if.paddr   <<= csr_if.paddr[3:0]
        bus_if_reg_if.pwdata  <<= csr_if.pwdata

        timer_reg_if.pwrite  <<= csr_if.pwrite
        timer_reg_if.psel    <<= csr_timer_psel
        timer_reg_if.penable <<= csr_if.penable
        timer_reg_if.paddr   <<= csr_if.paddr[1:0]
        timer_reg_if.pwdata  <<= csr_if.pwdata

        self.cpu_task_mode_csr_if.pwrite  <<= csr_if.pwrite
        self.cpu_task_mode_csr_if.psel    <<= csr_cpu_task_mode_psel
        self.cpu_task_mode_csr_if.penable <<= csr_if.penable
        self.cpu_task_mode_csr_if.paddr   <<= csr_if.paddr[3:0]
        self.cpu_task_mode_csr_if.pwdata  <<= csr_if.pwdata

        self.cpu_scheduler_mode_csr_if.pwrite  <<= csr_if.pwrite
        self.cpu_scheduler_mode_csr_if.psel    <<= csr_cpu_scheduler_mode_psel
        self.cpu_scheduler_mode_csr_if.penable <<= csr_if.penable
        self.cpu_scheduler_mode_csr_if.paddr   <<= csr_if.paddr[7:0]
        self.cpu_scheduler_mode_csr_if.pwdata  <<= csr_if.pwdata

        event_prdata = Wire(BrewData)

        csr_if.prdata <<= SelectOne(
            csr_dma_psel,                dma_reg_if.prdata,
            csr_timer_psel,              timer_reg_if.prdata,
            csr_bus_if_psel,             bus_if_reg_if.prdata,
            csr_event_psel,              event_prdata,
            csr_cpu_task_mode_psel,      self.cpu_task_mode_csr_if.prdata,
            csr_cpu_scheduler_mode_psel, self.cpu_scheduler_mode_csr_if.prdata
        )
        csr_if.pready <<= SelectOne(
            csr_dma_psel,                dma_reg_if.pready,
            csr_timer_psel,              timer_reg_if.pready,
            csr_bus_if_psel,             bus_if_reg_if.pready,
            csr_event_psel,              1,
            csr_cpu_task_mode_psel,      self.cpu_task_mode_csr_if.pready,
            csr_cpu_scheduler_mode_psel, self.cpu_scheduler_mode_csr_if.pready
        )

        # EVENT COUNTERS
        #############################

        event_counter_size = 32
        event_counter_cnt = 8

        event_cnts = []
        event_selects = []
        event_regs = []
        event_enabled = Wire(logic)
        event_write_strobe = csr_event_psel &  csr_if.pwrite & csr_if.penable

        for i in range(event_counter_cnt):
            event_cnt = Wire(Unsigned(event_counter_size))
            event_select = Wire(Unsigned(5))
            event = Select(event_select,
                1,
                event_fetch_wait_on_bus,
                event_decode_wait_on_rf,
                event_mem_wait_on_bus,
                event_branch_taken,
                event_branch,
                event_load,
                event_store,
                event_load | event_store,
                event_execute,
                event_bus_idle,
                event_fetch,
                event_fetch_drop,
                event_inst_word
            )
            event_cnt <<= Reg((event_cnt + event)[event_cnt.get_num_bits()-1:0], clock_en=event_enabled)
            setattr(self, f"event_cnt_{i}", event_cnt)
            setattr(self, f"event_select_{i}", event_select)
            event_cnts.append(event_cnt)
            event_selects.append(event_select)
            event_regs.append(event_select)
            event_regs.append(event_cnt)
            del event
            del event_cnt
            del event_select

        event_addr = csr_if.paddr[4:0]
        event_prdata <<= Reg(Select(
            event_addr,
            event_enabled, # Global enable register
            0,             # Reserved
            *event_regs,   # Pairs of selector/counter registers
        ))
        event_enabled <<= Reg(csr_if.pwdata[0], clock_en=(event_addr == 0) & event_write_strobe) # Global enable register
        for i, (event_cnt, event_select) in enumerate(zip(event_cnts, event_selects)):
            event_select <<= Reg(csr_if.pwdata[event_select.get_num_bits()-1:0], clock_en=(event_addr == i*2+2) & event_write_strobe) # Selector registers
            # Counters are read-only.
        del event_cnt
        del event_select

        # CSRs
        #####################

        '''
        APB signalling

                        <-- read -->      <-- write ->
            CLK     \__/^^\__/^^\__/^^\__/^^\__/^^\__/^^\__/
            psel    ___/^^^^^^^^^^^\_____/^^^^^^^^^^^\______
            penable _________/^^^^^\___________/^^^^^\______
            pready  ---------/^^^^^\-----------/^^^^^\------
            pwrite  ---/^^^^^^^^^^^\-----\___________/------
            paddr   ---<===========>-----<===========>------
            prdata  ---------<=====>------------------------
            pwdata  ---------------------<===========>------
            csr_rs  ___/^^^^^^^^^^^\________________________
            csr_wr  ___________________________/^^^^^\______
        '''

        create_apb_reg_map(self.csrs_cpu_task_mode,      self.cpu_task_mode_csr_if)
        create_apb_reg_map(self.csrs_cpu_scheduler_mode, self.cpu_scheduler_mode_csr_if)

def gen():
    def top():
        return BrewV1Top(nram_base=0x0, has_multiply=True, has_shift=True, page_bits=7)

    back_end = SystemVerilog()
    back_end.yosys_fix = True
    back_end.support_cast = False
    netlist = Build.generate_rtl(top, "brew_v1_top.sv", back_end=back_end)
    top_level_name = netlist.get_module_class_name(netlist.top_level)
    cyclone_v_device = "5CEBA4U15C7"
    max_10_device = "10M04SAE144C7G"
    flow = QuartusFlow(
        target_dir="q_brew_v1_top",
        top_level="brew_v1",
        source_files=("brew_v1.sv", "brew_v1_top.sv"),
        constraint_files=("brew_v1.sdc",),
        clocks=(("clk", 100),),
        device=max_10_device,
        project_name="BREW_V1"
    )
    flow.generate()
    flow.run()

if __name__ == "__main__":
    gen()


from flask import redirect, render_template, request, url_for
from application import app, db

from application.controllers.models import Controller
from application.cables.models import Cable

@app.route("/")
def index():
    cs = Controller.findAll()
    cables = []
    for c in Cable.query.all():
        con_a = Controller.query.get(c.controller_a_id)
        con_b = Controller.query.get(c.controller_b_id)
        cables.append({
            "id":c.id,
            "name":c.name,
            "x1":con_a.x,
            "y1":con_a.y,
            "x2":con_b.x,
            "y2":con_b.y
        })
    return render_template("index.html", controllers = Controller.findAll(), cables = cables)


import tensorflow as tf
from keras.callbacks import CSVLogger, ModelCheckpoint
import keras
import os
from scipy.io import wavfile
import numpy as np
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
from models import  LSTM
from glob import glob
import warnings

class DataGenerator(tf.keras.utils.Sequence):
    def __init__(self, wav_paths, labels, sr, dt, n_classes,
                 batch_size=32, shuffle=True):
        self.wav_paths = wav_paths
        self.labels = labels
        self.sr = sr
        self.dt = dt
        self.n_classes = n_classes
        self.batch_size = batch_size
        self.shuffle = True
        self.on_epoch_end()

    def __len__(self):
        return int(np.floor(len(self.wav_paths) / self.batch_size))

    def __getitem__(self, index):
        indexes = self.indexes[index*self.batch_size:(index+1)*self.batch_size]

        wav_paths = [self.wav_paths[k] for k in indexes]
        labels = [self.labels[k] for k in indexes]

        X = np.empty((self.batch_size, int(self.sr*self.dt), 1), dtype=np.float32)
        Y = np.empty((self.batch_size, self.n_classes), dtype=np.float32)

        for i, (path, label) in enumerate(zip(wav_paths, labels)):
            rate, wav = wavfile.read(path)
            X[i,] = wav.reshape(-1, 1)
            Y[i,] = keras.utils.np_utils.to_categorical(label, num_classes=self.n_classes)

        return X, Y

    def on_epoch_end(self):
        self.indexes = np.arange(len(self.wav_paths))
        if self.shuffle:
            np.random.shuffle(self.indexes)

def train(direct,timeIn,sampleRate,model,batchSize):
    src_root = direct
    sr = sampleRate
    dt = timeIn
    batch_size = batchSize
    model_type = model
    params = {'N_CLASSES':len(os.listdir(direct)),
              'SR':sr,
              'DT':dt}
    models = {'lstm':  LSTM(**params)}

    assert model_type in models.keys(), '{} not an available model'.format(model_type)
    csv_path = os.path.join('logs', '{}_history.csv'.format(model_type))

    wav_paths = glob('{}/**'.format(src_root), recursive=True)
    wav_paths = [x.replace(os.sep, '/') for x in wav_paths if '.wav' in x]

    classes = sorted(os.listdir(direct))
    le = LabelEncoder()
    le.fit(classes)
    labels = [os.path.split(x)[0].split('/')[-1] for x in wav_paths]
    labels = le.transform(labels)
    wav_train, wav_val, label_train, label_val = train_test_split(wav_paths,labels,test_size=0.1,random_state=0)

    assert len(label_train) >= batchSize, 'Number of train samples must be >= batch_size'
    if len(set(label_train)) != params['N_CLASSES']:
        warnings.warn('Found {}/{} classes in training data. Increase data size or change random_state.'.format(len(set(label_train)), params['N_CLASSES']))
    if len(set(label_val)) != params['N_CLASSES']:
        warnings.warn('Found {}/{} classes in validation data. Increase data size or change random_state.'.format(len(set(label_val)), params['N_CLASSES']))

    tg = DataGenerator(wav_train, label_train, sr, dt,
                       params['N_CLASSES'], batch_size=batch_size)
    vg = DataGenerator(wav_val, label_val, sr, dt,
                       params['N_CLASSES'], batch_size=batch_size)
    model = models[model_type]
    cp = ModelCheckpoint('models/{}.h5'.format(model_type), monitor='val_loss',
                         save_best_only=True, save_weights_only=False,
                         mode='auto', save_freq='epoch', verbose=1)
    csv_logger = CSVLogger(csv_path, append=False)
    model.fit(tg, validation_data=vg,
              epochs=30, verbose=1,
              callbacks=[csv_logger, cp])

if __name__ == '__main__':
    model='lstm'
    direct='clean'
    batchSize=16
    timeIn= 1.0
    sampleRate=16000
    train(direct, timeIn, sampleRate,model,batchSize)


from typing import Tuple, Dict, List, Iterable

from abianalysis import Hierarchy
from pylineage.node import TreeNode

Color = Tuple[float, float, float]
Position = Tuple[int, int]
Stack = Dict[Position, TreeNode]
Side = List[int]
FRONTAL = [1, 0]
SAGITTAL = [1, 2]
HORIZONTAL = [0, 2]

def trim_sides(pos_dict):
    """Returns copy dictionary where all positions are shifted so that
    the smallest x and y value are now zero.

    :param pos_dict: Dictionary where all keys are 2D positions.

    """
    minx = min(x for x, _ in pos_dict)
    miny = min(y for _, y in pos_dict)
    return {(x - minx, y - miny): node for (x, y), node in pos_dict.items()}

def view_stack_fn(depth: int, side: Side, trim: bool = True):
    def _get_ancestor(node: TreeNode):
        try:
            return node.ancestor_at_depth(depth)
        except ValueError:
            return node

    def _get_pos(node: Hierarchy):
        return tuple(map(int, node.position / node.volume.voxel_size))

    def _project_pos(pos):
        return tuple(pos[s] for s in side)

    def _get_placed_ancestors(nodes: Iterable[Hierarchy]):
        return {_get_pos(node): _get_ancestor(node) for node in nodes}

    def _count_stack(placed_nodes):
        counts_per_pos = {}
        for pos, node in placed_nodes.items():
            proj_pos = _project_pos(pos)
            counts = counts_per_pos.setdefault(proj_pos, {node: 0})
            counts[node] = counts.get(node, 0) + 1
        return {p: max(counts, key=counts.__getitem__)
                for p, counts in counts_per_pos.items()}

    def _get_stack(nodes):
        placed_ancestors = _get_placed_ancestors(nodes)
        stack = _count_stack(placed_ancestors)

        if trim:
            stack = trim_sides(stack)

        return stack

    return _get_stack

def view_stack(nodes, depth: int, side: Side, trim=True) -> Stack:
    """Create a stacked view of a hierarchy at a specific depth, from
    one side.

    :param nodes:  A map from positions (in 3d) to leaf hierarchy nodes.
        Positions without a node are ignored.
    :param depth:  The depth of the hierarchy to assign to the stack
    :param side:  list of indices deciding the view. For example,
        [1, 0] is a front view, and [1,2] is a side view.

    :return: A dictionary from projected positions to the most common
        Hierarchy node in the collapsed direction
    """
    return view_stack_fn(depth, side, trim)(nodes)

def stack_to_image(stack: Stack,
                   color_dict: Dict[TreeNode, Color]):
    """

    :param Stack stack:
    :param Dict[Hierarchy, Color] color_dict:
    :return:
    """
    x, y = np.array(list(stack), dtype=int).T
    mix, miy = x.min(), y.min()
    arr = np.zeros((x.max() - mix + 1, y.max() - miy + 1, 4))
    for x, y in zip(x, y):
        arr[x - mix, y - miy] = [*color_dict[stack[x, y]], 1]
    return arr


#!/usr/bin/env python3

import sys
import glob

sys.path.insert(0, 'pycoin')
from pycoin.block import Block
from pycoin.tx.script.tools import get_opcode, opcode_list
from pycoin.tx.script.check_signature import parse_signature_blob, delete_subscript, bin_script
from pycoin.tx.script.der import UnexpectedDER
from pycoin.serialize import b2h_rev, h2b, h2b_rev, stream_to_bytes
from pycoin.ecdsa.secp256k1 import secp256k1_generator
from pycoin.encoding import sec_to_public_pair, EncodingError, secret_exponent_to_wif

tx_db = dict()
for filename in glob.glob("*.bin"):
    block = Block.parse(open(filename, "rb"))
    for tx in block.txs:
        tx_db[tx.hash()] = tx

def opcode_data(script):
    opcodes = []
    pc = 0
    while pc < len(script):
        try: opcode, data, pc = get_opcode(script, pc)
        except: return None
        opcodes.append(data)
    return opcodes

for filename in glob.glob("*.bin"):
    block = Block.parse(open(filename, "rb"))
    for tx in block.txs:
        for tx_in_idx, tx_in in enumerate(tx.txs_in):
            if tx.is_coinbase(): continue
            tx.unspents_from_db(tx_db, ignore_missing=True)
            if not tx.is_signature_ok(tx_in_idx): continue
            tx_out_script = tx.unspents[tx_in_idx].script
            to_sign = []
            def signature_for_hash_type_f(hash_type, script):
                h = tx.signature_hash(script, tx_in_idx, hash_type)
                to_sign.append(h)
                return h
            if not tx_in.verify(tx_out_script, signature_for_hash_type_f,
                lock_time=tx.lock_time, tx_version=tx.version): continue
            if len(to_sign) < 1: continue
            signature_hash = to_sign[0]

            stack = opcode_data(tx_in.script)
            if stack is None or len(stack) != 2: continue
            pair_blob = stack.pop()
            sig_blob = stack.pop()
            try: sig_pair, signature_type = parse_signature_blob(sig_blob)
            except: continue
            public_pair = sec_to_public_pair(pair_blob)
            r, s = sig_pair
            x, y = public_pair

            print("Address:", tx_in.bitcoin_address())
            print("r:", r)
            print("s:", s)
            print("x:", x)
            print("y:", y)
            print("h:", signature_hash)
            # print("Verifies:", secp256k1_generator.verify(public_pair, signature_hash, sig_pair))


import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt

img = cv.imread('/Users/randheerkumar/Downloads/1.jpg')

"""plt.imshow(img)
plt.show()"""

"""bgr to hsv
hsv=cv.cvtColor(img, cv.COLOR_BGR2HSV)
bgr to lab
lab=cv.cvtColor(img, cv.COLOR_BGR2LAB)"""

height=400
width=400
dim=(width,height)

resized=cv.resize(img, dim)

"""b,g,r=cv.split(resized)
cv.imshow('blue', b)
cv.imshow('green', g)
cv.imshow('red', r)"""

blur3=cv.blur(resized, (3,3))
blur5=cv.blur(resized, (3,3))
blur9=cv.blur(resized, (3,3))
cv.imshow('blur3', blur3)
cv.imshow('blur5', blur5)
cv.imshow('blur9', blur9)

cv.imshow('original', resized)
#cv.imwrite('/Users/randheerkumar/Downloads/lab.jpg', lab)
cv.waitKey(20000)
cv.destroyAllWindows()