import torch
from torch import nn
import transformers
import tqdm
from torch.utils.data import DataLoader
import pytorch_lightning as pl
from util import *
from model_parts import *

class NeuralAlignerBase(pl.LightningModule):
    def __init__(self, **kwargs):
        super().__init__()
        self.sinkhorn_epsilon = 0.1
        self.stopThr = 1e-4
        self.numItermax = 1000
        self.ot_type = kwargs['ot_type']

        self.loss_fnc = nn.BCELoss(reduction='sum')
        self.positive_label = 1
        self.negative_label = 0

    def get_word_embeddings(self, hidden_tensors, subword_to_word_conv, s2_start_idx):
        word_num = torch.amax(subword_to_word_conv).item() + 1
        word_embeddings = torch.vstack(
            ([torch.mean(hidden_tensors[subword_to_word_conv == word_idx], dim=0) for word_idx in
              range(word_num)]))

        s1_word_embeddigs = word_embeddings[0:s2_start_idx, :]
        s2_word_embeddigs = word_embeddings[s2_start_idx:, :]

        return s1_word_embeddigs, s2_word_embeddigs

    def get_attention_per_word(self, attentions, subword_to_word_conv, s2_start_idx):
        eps = 1e-12
        word_num = torch.amax(subword_to_word_conv).item() + 1
        attentions = torch.vstack(
            ([torch.sum(attentions[subword_to_word_conv == word_idx], dim=0) for word_idx in
              range(word_num)]))
        attentions = torch.vstack(
            ([torch.sum(attentions[:, subword_to_word_conv == word_idx], dim=1) for word_idx in
              range(word_num)]))

        attentions = attentions + eps
        atten_s2t = attentions[0:s2_start_idx, s2_start_idx:word_num]
        atten_s2t = (atten_s2t.T / atten_s2t.sum(1)).T
        atten_t2s = attentions[s2_start_idx:word_num, 0:s2_start_idx]
        atten_t2s = (atten_t2s.T / atten_t2s.sum(1)).T
        atten = (atten_s2t + atten_t2s.T) / 2

        return atten

    def convert_to_word_embeddings(self, offset_mapping, token_ids, hidden_tensors, pooling):
        pooling_method = torch.mean if pooling == 'mean' else torch.amax
        word_idx = -1
        subword_to_word_conv = np.full((hidden_tensors.shape[0]), -1)
        # Bug in hugging face tokenizer? Sometimes Metaspace is inserted
        metaspace = getattr(self.tokenizer.decoder, "replacement", None)
        metaspace = self.tokenizer.decoder.prefix if metaspace is None else metaspace
        tokenizer_bug_idxes = [i for i, x in enumerate(self.tokenizer.convert_ids_to_tokens(token_ids)) if
                               x == metaspace]

        for subw_idx, offset in enumerate(offset_mapping):
            if subw_idx in tokenizer_bug_idxes:
                continue
            elif offset[0] == offset[1]:  # Special token
                continue
            elif offset[0] == 0:
                word_idx += 1
                subword_to_word_conv[subw_idx] = word_idx
            else:
                subword_to_word_conv[subw_idx] = word_idx

        word_embeddings = torch.vstack(
            ([pooling_method(hidden_tensors[subword_to_word_conv == word_idx], dim=0) for word_idx in
              range(word_idx + 1)]))

        sep_tok_indices = [i for i, x in enumerate(token_ids) if x == self.tokenizer.sep_token_id]
        s2_start_idx = subword_to_word_conv[
            sep_tok_indices[0] + np.argmax(subword_to_word_conv[sep_tok_indices[0]:] > -1)]

        s1_word_embeddigs = word_embeddings[0:s2_start_idx, :]
        s2_word_embeddigs = word_embeddings[s2_start_idx:, :]

        return s1_word_embeddigs, s2_word_embeddigs

    def convert_to_word_attentions(self, offset_mapping, token_ids, attentions, cross_attention):
        word_idx = -1
        subword_to_word_conv = np.full((attentions.shape[0]), -1)
        # Bug in hugging face tokenizer? Sometimes Metaspace is inserted
        metaspace = getattr(self.tokenizer.decoder, "replacement", None)
        metaspace = self.tokenizer.decoder.prefix if metaspace is None else metaspace
        tokenizer_bug_idxes = [i for i, x in enumerate(self.tokenizer.convert_ids_to_tokens(token_ids)) if
                               x == metaspace]

        for subw_idx, offset in enumerate(offset_mapping):
            if subw_idx in tokenizer_bug_idxes:
                continue
            elif offset[0] == offset[1]:  # Special token
                continue
            elif offset[0] == 0:
                word_idx += 1
                subword_to_word_conv[subw_idx] = word_idx
            else:
                subword_to_word_conv[subw_idx] = word_idx

        # Take mean for 'from', sum for 'to' following https://arxiv.org/pdf/1906.04341.pdf
        attentions = torch.vstack(
            ([torch.mean(attentions[subword_to_word_conv == word_idx], dim=0) for word_idx in range(word_idx + 1)]))
        attentions = torch.vstack(
            ([torch.sum(attentions[:, subword_to_word_conv == word_idx], dim=1) for word_idx in range(word_idx + 1)]))

        sep_tok_indices = [i for i, x in enumerate(token_ids) if x == self.tokenizer.sep_token_id]
        s2_start_idx = subword_to_word_conv[
            sep_tok_indices[0] + np.argmax(subword_to_word_conv[sep_tok_indices[0]:] > -1)]

        if cross_attention:
            s1_attn = attentions[0:s2_start_idx, s2_start_idx:]
            s2_attn = attentions[s2_start_idx:, 0:s2_start_idx]

        else:
            # self-attention: average 'from' and 'to' attention weights
            s1_attn = (attentions[0:s2_start_idx, 0:s2_start_idx] + attentions[0:s2_start_idx, 0:s2_start_idx].t()) / 2
            s2_attn = (attentions[s2_start_idx:, s2_start_idx:] + attentions[s2_start_idx:, s2_start_idx:].t()) / 2

        return s1_attn, s2_attn

    def batch_center_embedding(self, hidden_outputs):
        mean_vec = torch.mean(hidden_outputs.detach(), dim=(1, 0))
        hidden_outputs = hidden_outputs - mean_vec
        return hidden_outputs

    def compute_loss(self, batch, forward_outputs):
        loss = 0
        transport_matrix_list = forward_outputs['transport_matrix_list']
        for bidx in range(len(transport_matrix_list)):
            P = transport_matrix_list[bidx]
            y = torch.full_like(P, fill_value=self.negative_label, device=self.device)
            # weights = torch.ones_like(P, device=self.device)
            for pair in batch['atuples'][bidx]:
                y[pair[0], pair[1]] = self.positive_label

            loss += self.loss_fnc(P, y)

        loss = loss / len(transport_matrix_list)
        logs = {"loss": loss}
        return loss, logs

    def compute_score(self, pred, gold):
        gold = set(gold.split())
        if len(pred) > 0:
            precision = len(gold & pred) / len(pred)
        else:
            if len(gold) == 0:
                precision = 1
            else:
                precision = 0
        if len(gold) > 0:
            recall = len(gold & pred) / len(gold)
        else:
            if len(pred) == 0:
                recall = 1
            else:
                recall = 0
        if len(pred & gold) == len(gold) and len(pred & gold) == len(pred):
            acc = 1
        else:
            acc = 0

        return precision, recall, acc

    def compute_null_score(self, pred, gold, sent1, sent2):
        gold = set(gold.split())
        sent1 = sent1.split()
        sent2 = sent2.split()
        gold_s1_indices, gold_s2_indices = get_aligned_indices(gold)
        gold_null_s1_indices = set(range(len(sent1))) - gold_s1_indices
        gold_null_s2_indices = set(range(len(sent2))) - gold_s2_indices

        aligned_s1_indices, aligned_s2_indices = get_aligned_indices(pred)
        pred_null_s1_indices = set(range(len(sent1))) - aligned_s1_indices
        pred_null_s2_indices = set(range(len(sent2))) - aligned_s2_indices

        if len(pred_null_s1_indices) + len(pred_null_s2_indices) > 0:
            precision = (len(gold_null_s1_indices & pred_null_s1_indices) + len(
                gold_null_s2_indices & pred_null_s2_indices)) / (len(pred_null_s1_indices) + len(pred_null_s2_indices))
        else:
            if len(gold_null_s1_indices) + len(gold_null_s2_indices) == 0:
                precision = 1
            else:
                precision = 0
        if len(gold_null_s1_indices) + len(gold_null_s2_indices) > 0:
            recall = (len(gold_null_s1_indices & pred_null_s1_indices) + len(
                gold_null_s2_indices & pred_null_s2_indices)) / (len(gold_null_s1_indices) + len(gold_null_s2_indices))
        else:
            if len(pred_null_s1_indices) + len(pred_null_s2_indices) == 0:
                recall = 1
            else:
                recall = 0

        if len(gold_null_s1_indices & pred_null_s1_indices) == len(gold_null_s1_indices) and len(
                gold_null_s1_indices & pred_null_s1_indices) == len(pred_null_s1_indices) and len(
            gold_null_s2_indices & pred_null_s2_indices) == len(gold_null_s2_indices) and len(
            gold_null_s2_indices & pred_null_s2_indices) == len(pred_null_s2_indices):
            acc = 1.0
        else:
            acc = 0.0

        return precision, recall, acc

    def compute_total_score(self, pred, gold, sent1, sent2):
        gold = set(gold.split())
        sent1 = sent1.split()
        sent2 = sent2.split()
        gold_s1_indices, gold_s2_indices = get_aligned_indices(gold)
        gold_null_s1_indices = set(range(len(sent1))) - gold_s1_indices
        gold_null_s2_indices = set(range(len(sent2))) - gold_s2_indices

        aligned_s1_indices, aligned_s2_indices = get_aligned_indices(pred)
        pred_null_s1_indices = set(range(len(sent1))) - aligned_s1_indices
        pred_null_s2_indices = set(range(len(sent2))) - aligned_s2_indices

        precision = (len(gold & pred) + len(gold_null_s1_indices & pred_null_s1_indices) + len(
            gold_null_s2_indices & pred_null_s2_indices)) / (
                            len(pred) + len(pred_null_s1_indices) + len(pred_null_s2_indices))
        recall = (len(gold & pred) + len(gold_null_s1_indices & pred_null_s1_indices) + len(
            gold_null_s2_indices & pred_null_s2_indices)) / (
                         len(gold) + len(gold_null_s1_indices) + len(gold_null_s2_indices))

        if len(pred & gold) == len(gold) and len(pred & gold) == len(pred):
            acc = 1
        else:
            acc = 0

        return precision, recall, acc

    def evaluate_alignments(self, outputs, val_or_test, save_prediction=False):
        precision_all = []
        recall_all = []
        accuracy_all = []
        null_precision_all = []
        null_recall_all = []
        null_accuracy_all = []
        total_precision_all = []
        total_recall_all = []
        total_accuracy_all = []
        # align_thresh = self.determine_align_thresh(outputs)
        for output in outputs:
            transport_matrix_list = output['transport_matrix_list']
            golds = output['gold']
            sents1 = output['s1_sents']
            sents2 = output['s2_sents']
            for bidx, P in enumerate(transport_matrix_list):
                alignments = set()
                align_pairs = np.transpose(np.nonzero(P > self.align_hypr))
                for i_j in align_pairs:
                    alignments.add('{0}-{1}'.format(i_j[0], i_j[1]))

                p, r, acc = self.compute_score(alignments, golds[bidx])
                null_p, null_r, null_acc = self.compute_null_score(alignments, golds[bidx], sents1[bidx], sents2[bidx])
                total_p, total_r, total_acc = self.compute_total_score(alignments, golds[bidx], sents1[bidx],
                                                                       sents2[bidx])
                precision_all.append(p)
                recall_all.append(r)
                accuracy_all.append(acc)
                null_precision_all.append(null_p)
                null_recall_all.append(null_r)
                null_accuracy_all.append(null_acc)
                total_precision_all.append(total_p)
                total_recall_all.append(total_r)
                total_accuracy_all.append(total_acc)

        precision = sum(precision_all) / len(precision_all)
        recall = sum(recall_all) / len(recall_all)
        if precision + recall > 0:
            f1 = 2 * precision * recall / (precision + recall)
        else:
            f1 = 0.0
        accuracy = sum(accuracy_all) / len(accuracy_all)

        null_precision = sum(null_precision_all) / len(null_precision_all)
        null_recall = sum(null_recall_all) / len(null_recall_all)
        if null_precision + null_recall > 0:
            null_f1 = 2 * null_precision * null_recall / (null_precision + null_recall)
        else:
            null_f1 = 0.0
        null_accuracy = sum(null_accuracy_all) / len(null_accuracy_all)

        total_precision = sum(total_precision_all) / len(total_precision_all)
        total_recall = sum(total_recall_all) / len(total_recall_all)
        if total_precision + total_recall > 0:
            total_f1 = 2 * total_precision * total_recall / (total_precision + total_recall)
        else:
            total_f1 = 0.0
        total_accuracy = sum(total_accuracy_all) / len(total_accuracy_all)

        logs = {"precision": precision, "recall": recall, "f1": f1, "exact_match": accuracy,
                "null_precision": null_precision, "null_recall": null_recall, "null_f1": null_f1,
                "null_exact_match": null_accuracy, "total_precision": total_precision, "total_recall": total_recall,
                "total_f1": total_f1, "total_exact_match": total_accuracy}

        if save_prediction:
            with open(self.logger.log_dir + '/{0}_{1}_report.txt'.format(self.data, val_or_test), 'a') as fw:
                fw.write('Alignment Hyperparameter:\t{0:.4f}\n'.format(self.align_hypr))
                fw.write(
                    'Precision\tRecall\tF1\tExactMatch\t[Null]Precision\t[Null]Recall\t[Null]F1\t[Null]ExactMatch\t[Total]Precision\t[Total]Recall\t[Total]F1\t[Total]ExactMatch\n')
                fw.write(
                    '{0:.4f}\t{1:.4f}\t{2:.4f}\t{3:.4f}\t{4:.4f}\t{5:.4f}\t{6:.4f}\t{7:.4f}\t{8:.4f}\t{9:.4f}\t{10:.4f}\t{11:.4f}\n'.format(
                        precision * 100,
                        recall * 100,
                        f1 * 100,
                        accuracy * 100,
                        null_precision * 100,
                        null_recall * 100,
                        null_f1 * 100,
                        null_accuracy * 100,
                        total_precision * 100,
                        total_recall * 100,
                        total_f1 * 100,
                        total_accuracy * 100, ))

            with open(self.logger.log_dir + '/{0}_{1}_alignments_{2:.4f}.txt'.format(self.data, val_or_test,
                                                                                self.align_hypr),
                      'w') as fw:
                fw.write('Sentence_1\tSentence_2\tGold\tPrediction\n')
                for output in outputs:
                    transport_matrix_list = output['transport_matrix_list']
                    golds = output['gold']
                    s1_sents = output['s1_sents']
                    s2_sents = output['s2_sents']
                    for bidx, P in enumerate(transport_matrix_list):
                        alignments = set()
                        align_pairs = np.transpose(np.nonzero(P > self.align_hypr))
                        for i_j in align_pairs:
                            alignments.add('{0}-{1}-{2:.4f}'.format(i_j[0], i_j[1], P[i_j[0], i_j[1]]))

                        fw.write('{0}\t{1}\t{2}\t{3}\n'.format(s1_sents[bidx], s2_sents[bidx], golds[bidx],
                                                               ' '.join(alignments)))

        return logs

    def determine_sinkhorn_threshold(self, outputs, val_or_test):
        best_thresh = 0.0
        best_val_log = {"precision": 0.0, "recall": 0.0, "f1": 0.0, "exact_match": 0.0, "null_precision": 0.0,
                        "null_recall": 0.0, "null_f1": 0.0, "null_exact_match": 0.0, "total_precision": 0.0,
                        "total_recall": 0.0, "total_f1": 0.0, "total_exact_match": 0.0}

        if self.ot_type == 'uot-mm':  # No thresholding is necessary
            self.align_hypr = best_thresh = 0.0
            best_val_log = self.evaluate_alignments(outputs, val_or_test, save_prediction=False)
        else:
            # Determine threshold
            for th in np.linspace(0.0, 1.0, num=100, endpoint=True):
                self.align_hypr = th
                logs = self.evaluate_alignments(outputs, val_or_test, save_prediction=False)
                if logs['total_f1'] > best_val_log['total_f1']:
                    best_thresh = th
                    best_val_log = logs
        return best_thresh, best_val_log

    def training_step(self, batch, batch_idx):
        forward_outputs = self.forward(batch)
        loss, logs = self.compute_loss(batch, forward_outputs)
        self.log_dict({f"train_{k}": v for k, v in logs.items()})
        self.log_dict({f"batch_size": self.batch_size})
        return loss

    def bidirectional_align(self, batch):
        outputs = self.forward(batch)
        transport_matrix_list = [matrix.cpu().detach().numpy() for matrix in outputs['transport_matrix_list']]
        return transport_matrix_list

    def validation_step(self, batch, batch_idx):
        transport_matrix_list = self.bidirectional_align(batch)

        return {'transport_matrix_list': transport_matrix_list, 'gold': batch['alignments'],
                's1_sents': batch['s1_sents'], 's2_sents': batch['s2_sents']}

    def validation_epoch_end(self, outputs):
        best_thresh, _ = self.determine_sinkhorn_threshold(outputs, 'val')
        self.align_hypr = best_thresh
        logs = self.evaluate_alignments(outputs, 'val', save_prediction=False)
        self.log_dict({f"val_{k}": v for k, v in logs.items()})
        self.log_dict({f"batch_size": self.batch_size})
        return logs

    def on_test_start(self):
        # Determine best align_hypr
        # generate alignments
        eval_dict = {'transport_matrix_list': [], 'gold': [], 's1_sents': [], 's2_sents': []}
        all_transport_matrices = []
        for batch in tqdm.tqdm(self.val_dataloader(), desc='Encoding validation set'):
            batch = send_batch_to_gpu(batch, self.device)
            with torch.no_grad():
                all_transport_matrices += self.bidirectional_align(batch)

            eval_dict['gold'] += batch['alignments']
            eval_dict['s1_sents'] += batch['s1_sents']
            eval_dict['s2_sents'] += batch['s2_sents']
        eval_dict['transport_matrix_list'] = all_transport_matrices
        best_thresh, best_val_log = self.determine_sinkhorn_threshold([eval_dict], 'val')
        self.align_hypr = best_thresh
        logs = self.evaluate_alignments([eval_dict], 'val', save_prediction=True)

    def test_step(self, batch, batch_idx):
        transport_matrix_list = self.bidirectional_align(batch)

        return {'transport_matrix_list': transport_matrix_list, 'gold': batch['alignments'],
                's1_sents': batch['s1_sents'], 's2_sents': batch['s2_sents']}

    def test_epoch_end(self, outputs):
        logs = self.evaluate_alignments(outputs, 'test', save_prediction=True)
        self.log_dict({f"test_{k}": v for k, v in logs.items()})
        self.log_dict({f"batch_size": self.batch_size})

    def configure_optimizers(self):
        optimizer = transformers.Adafactor(self.parameters(), scale_parameter=False, relative_step=False,
                                           warmup_init=False, lr=self.lr)
        return optimizer

    # return the dataloader for each split
    def train_dataloader(self):
        return self.data_loader(self.data, 'train')

    def val_dataloader(self):
        return self.data_loader(self.data, 'dev')

    def test_dataloader(self):
        return self.data_loader(self.data, 'test')

    def data_loader(self, data, type):
        sents1, sents2, align_lists = load_WA_corpus(data, type, self.sure_and_possible)

        if type == 'train':
            shuffle = True
        else:
            shuffle = False

        input_s2t = self.tokenizer(text=sents1, text_pair=sents2, padding=True, truncation=True,
                                   is_split_into_words=True,
                                   return_offsets_mapping=True, return_tensors="pt")

        all_subword_to_word_map, all_s2_start_idx = get_token_word_mapping(input_s2t, self.tokenizer)

        return DataLoader(
            AlignmentDataset(sents1, sents2, input_s2t, align_lists, all_subword_to_word_map, all_s2_start_idx),
            batch_size=self.batch_size,
            shuffle=shuffle, collate_fn=alignment_collate)


import random
pwin = 0
cwin = 0
while cwin < 2 and pwin < 2 :  #人机 赢了不够2次 继续执行

    all_list = ['石头','拳头','布']
    win_list = [['石头','拳头'],['石头','布'],['拳头','布']] #前面是用户后面是电脑
    prompt = '''0 石头
    1 拳头
    2 布
    请选择'''
    computer = random.choice(all_list)
    ind = int(input(prompt))
    play = all_list[ind]
    print('play: %s,computer: %s' % (play,computer) )
    if play == computer:

                print('\033[32;1m平局\033[0m')

    elif [play,computer] in win_list:
                print('\033[31;1m赢了\033[0m')
                pwin += 1

    else:
                print('\033[32;1m失败\033[0m')

                cwin += 1

    if cwin < pwin:
        print('人赢了' )
    else:
        print('电脑赢了')


#!/usr/bin/env python
# coding=UTF-8
import cv2
import os
import numpy as np
from tqdm import tqdm
import time
import json
from collections import defaultdict
import random
import matplotlib.pyplot as plt
import sys
sys.path.append('../')
from utils import draw_rectangle

img_path = '/home/caoliwei/Dataset/visdrone/images'
annFile = '/home/caoliwei/Dataset/visdrone/test_tiny.json'

extract_num = 10

print('loading json:')
start = time.time()
data = json.load(open(annFile))
print('loading json done, time use %.3fs' % (time.time() - start ))

img_infos = data['images']
img_ids = []
imgid_imginfo_dict = {}
imgname_imgid_dict = {}
for img_info in tqdm(img_infos):
    img_id = img_info['id']
    imgid_imginfo_dict[img_id] = img_info
    file_name = img_info['file_name'].split('/')[-1]
    imgname_imgid_dict[file_name] = img_id
    img_ids.append(img_id)

annotations = data['annotations']
imgid_ann_dict = defaultdict(list)
for ann in tqdm(annotations):
    img_id = ann['image_id']
    imgid_ann_dict[img_id].append(ann)

img_names = os.listdir(img_path)
for i, img_id in enumerate(tqdm(img_ids[:extract_num])):
# for i, img_name in enumerate(tqdm(img_names[:extract_num])):
    #img_id = imgname_imgid_dict[img_name[:-12] + '.jpg'  ] # no '_resized'
    # if not '33349478' in img_name:
    #     continue
    # img_id = imgname_imgid_dict[img_name]
    img_info = imgid_imginfo_dict[img_id]
    image_name = img_info['file_name'].split('/')[-1]
    img_info_w = img_info['width']
    img_info_h = img_info['height']
    # print(image_name)
    anns = imgid_ann_dict[img_id]

    coordinates = []
    labels_text = []
    img_raw = cv2.imread(os.path.join(img_path, image_name))
    img_h, img_w = img_raw.shape[:2]

    for ann in anns:
        # 1、求坐标
        coordinate = []
        coordinate.append(ann['bbox'][0] / img_info_w * img_w )  # convert to actual img size, not img_info size
        coordinate.append(ann['bbox'][1] / img_info_h * img_h )
        coordinate.append(ann['bbox'][2] / img_info_w * img_w )
        coordinate.append(ann['bbox'][3] / img_info_h * img_h )
        coordinates.append(coordinate)

        cls_name = str(ann['category_id']  )
        labels_text.append(cls_name)

    img = draw_rectangle(coordinates, labels_text, img_raw)
    plt.figure(figsize=(5, 5), dpi=200)
    plt.imshow(img[:, :, ::-1])
    plt.show()


n, m = input().split(" ")
n = int(n)
m = int(m)
matrix = []
for i in range(n):
    matrix.append(input().split(" "))

kernel = []
for i in range(3):
    kernel.append(input().split(" "))

for i in range(0, n - 2):
    for j in range(0, m - 2):
        res = 0
        for k in range(3):
            for l in range(3):
                res += int(matrix[i + k][j + l]) * int(kernel[k][l])
        if j == m - 3:
            print(res)
        else:
            print(res, end=' ')  # 注意这里end的用法


#!/usr/bin/env python
import django
import logging
import os
import sys

from django.conf import settings
from django.test.runner import DiscoverRunner
from edc_test_utils import DefaultTestSettings
from os.path import abspath, dirname, join

app_name = 'ambition_permissions'
base_dir = dirname(abspath(__file__))

DEFAULT_SETTINGS = DefaultTestSettings(
    calling_file=__file__,
    BASE_DIR=base_dir,
    LOAD_NAVBARS=False,
    APP_NAME=app_name,
    ETC_DIR=os.path.join(base_dir, app_name, "tests", "etc"),
    SUBJECT_CONSENT_MODEL="ambition_subject.subjectconsent",
    SUBJECT_VISIT_MODEL="ambition_subject.subjectvisit",
    SUBJECT_REQUISITION_MODEL="ambition_subject.subjectrequisition",
    ADVERSE_EVENT_ADMIN_SITE="ambition_ae_admin",
    ADVERSE_EVENT_APP_LABEL="ambition_ae",
    EDC_BOOTSTRAP=3,
    EMAIL_BACKEND="django.core.mail.backends.locmem.EmailBackend",
    EMAIL_CONTACTS={
        "data_request": "someone@example.com",
        "data_manager": "someone@example.com",
        "tmg": "someone@example.com",
    },
    EMAIL_ENABLED=True,
    HOLIDAY_FILE=join(base_dir, app_name, "tests", "holidays.csv"),
    LIVE_SYSTEM=False,
    RANDOMIZATION_LIST_PATH=join(
        base_dir, app_name, "tests", "test_randomization_list.csv"),
    INSTALLED_APPS=[
        "django.contrib.admin",
        "django.contrib.auth",
        "django.contrib.contenttypes",
        "django.contrib.sessions",
        "django.contrib.messages",
        "django.contrib.staticfiles",
        "django.contrib.sites",
        "django_crypto_fields.apps.AppConfig",
        "edc_action_item.apps.AppConfig",
        "edc_adverse_event.apps.AppConfig",
        "edc_appointment.apps.AppConfig",
        "edc_auth.apps.AppConfig",
        "edc_dashboard.apps.AppConfig",
        "edc_export.apps.AppConfig",
        "edc_identifier.apps.AppConfig",
        "edc_data_manager.apps.AppConfig",
        "edc_lab.apps.AppConfig",
        "edc_lab_dashboard.apps.AppConfig",
        "edc_locator.apps.AppConfig",
        "edc_metadata.apps.AppConfig",
        "edc_metadata_rules.apps.AppConfig",
        "edc_navbar.apps.AppConfig",
        "edc_notification.apps.AppConfig",
        "edc_offstudy.apps.AppConfig",
        "edc_pharmacy.apps.AppConfig",
        "edc_registration.apps.AppConfig",
        "edc_reference.apps.AppConfig",
        "edc_permissions.apps.AppConfig",
        "edc_sites.apps.AppConfig",
        "edc_visit_schedule.apps.AppConfig",
        'ambition_ae.apps.AppConfig',
        'ambition_labs.apps.AppConfig',
        'ambition_lists.apps.AppConfig',
        'ambition_prn.apps.AppConfig',
        'ambition_rando.apps.AppConfig',
        'ambition_screening.apps.AppConfig',
        'ambition_visit_schedule.apps.AppConfig',
        'ambition_subject.apps.AppConfig',
        'ambition_permissions.apps.AppConfig',
    ],
    add_dashboard_middleware=True,
    add_lab_dashboard_middleware=True,
    use_test_urls=True,
).settings

def main():
    if not settings.configured:
        settings.configure(**DEFAULT_SETTINGS)
    django.setup()
    tags = [t.split('=')[1] for t in sys.argv if t.startswith('--tag')]
    failures = DiscoverRunner(failfast=True, tags=tags).run_tests(
        [f'{app_name}.tests'])
    sys.exit(failures)

if __name__ == "__main__":
    logging.basicConfig()
    main()


# -*- coding: utf-8 -*-
# @Author: TomLotze
# @Date:   2020-09-18 11:21
# @Last Modified by:   TomLotze
# @Last Modified time: 2020-10-15 22:36

import argparse
import numpy as np
import os
from regression import Regression
import torch
from torch.utils.data import DataLoader, random_split
import matplotlib.pyplot as plt
import pickle as pkl
from collections import Counter
from scipy.stats import ttest_rel

# Default constants
DNN_HIDDEN_UNITS_DEFAULT = '300, 32'
DROPOUT_DEFAULT = '0.0, 0.0'
LEARNING_RATE_DEFAULT = 1e-3
NR_EPOCHS_DEFAULT = 30
BATCH_SIZE_DEFAULT = 64
EVAL_FREQ_DEFAULT = 100
NEG_SLOPE_DEFAULT = 0.02
DATA_DIR_DEFAULT = "Data/"

FLAGS = None

# os.environ['KMP_DUPLICATE_LIB_OK']='True'

def train():
    """
    Performs training and evaluation of Regression model.
    """
    # Set the random seeds for reproducibility
    np.random.seed(10)
    torch.manual_seed(10)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    # Get number of units in each hidden layer
    if FLAGS.dnn_hidden_units:
        dnn_hidden_units = FLAGS.dnn_hidden_units.split(",")
        dnn_hidden_units = [int(dnn_hidden_unit_) for dnn_hidden_unit_ in dnn_hidden_units]
    else:
        dnn_hidden_units = []

    # convert dropout percentages
    dropout_probs = [float(prob) for prob in FLAGS.dropout_probs.split(',')]

    # check if length of dropout is equal to nr of hidden layers
    if len(dropout_probs) != len(dnn_hidden_units):
        dropout_len = len(dropout_probs)
        hidden_len = len(dnn_hidden_units)
        if dropout_len < hidden_len:
            for _ in range(hidden_len-dropout_len):
                dropout_probs.append(0)
        else:
            dropout_probs = dropout_probs[:hidden_len]
    # use GPU if available
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("Device :", device)

    # extract all data and divide into train, valid and split dataloaders
    dataset_filename = f"dataset_filename=MIMICS-Click.tsv_expanded=False_balance=True_impression={FLAGS.impression}_reduced_classes={FLAGS.reduced_classes}_embedder={FLAGS.embedder}.p"
    with open(os.path.join(FLAGS.data_dir, dataset_filename), "rb") as f:
        dataset = pkl.load(f)

    len_all = len(dataset)

    train_len, valid_len = int(0.7 * len_all), int(0.15 * len_all)
    test_len = len_all - train_len - valid_len
    splits = [train_len, valid_len, test_len]
    train_data, valid_data, test_data = random_split(dataset, splits)

    train_dl = DataLoader(train_data, batch_size=FLAGS.batch_size, shuffle=True, drop_last=True)
    valid_dl = DataLoader(valid_data, batch_size=FLAGS.batch_size, shuffle=True, drop_last=True)
    test_dl = DataLoader(test_data, batch_size=FLAGS.batch_size, shuffle=True, drop_last=True)

    with open(f"{FLAGS.data_dir}/test_dl.pt", "wb") as f:
        pkl.dump(test_dl, f)

     # initialize MLP and loss function
    input_size = iter(train_dl).next()[0].shape[1] # 5376 for BERT embeddings
    nn = Regression(input_size, dnn_hidden_units, dropout_probs, 1, FLAGS.neg_slope, FLAGS.batchnorm).to(device)
    loss_function = torch.nn.MSELoss()

    if FLAGS.verbose:
        print(f"neural net:\n {[param.data for param in nn.parameters()]}")

    # initialize optimizer
    if FLAGS.optimizer == "SGD":
        optimizer = torch.optim.SGD(nn.parameters(), lr=FLAGS.learning_rate,
            weight_decay=FLAGS.weightdecay,
            momentum=FLAGS.momentum)
    elif FLAGS.optimizer == "Adam":
      optimizer = torch.optim.Adam(nn.parameters(), lr=FLAGS.learning_rate,
        amsgrad=FLAGS.amsgrad, weight_decay=FLAGS.weightdecay)
    elif FLAGS.optimizer == "AdamW":
      optimizer = torch.optim.AdamW(nn.parameters(), lr=FLAGS.learning_rate,
        amsgrad=FLAGS.amsgrad, weight_decay=FLAGS.weightdecay)
    elif FLAGS.optimizer == "RMSprop":
      optimizer = torch.optim.RMSprop(nn.parameters(), lr=FLAGS.learning_rate,
        weight_decay=FLAGS.weightdecay, momentum=FLAGS.momentum)

    # initialization for plotting and metrics
    training_losses = []
    valid_losses = []

    initial_train_loss = eval_on_test(nn, loss_function, train_dl, device)
    training_losses.append(initial_train_loss)
    initial_valid_loss = eval_on_test(nn, loss_function, valid_dl, device)
    valid_losses.append(initial_valid_loss)

    # construct name for saving models and figures
    variables_string = f"regression_{FLAGS.embedder}_{FLAGS.impression}_{FLAGS.reduced_classes}_{FLAGS.optimizer}_{FLAGS.learning_rate}_{FLAGS.weightdecay}_{FLAGS.momentum}_{FLAGS.dnn_hidden_units}_{FLAGS.dropout_probs}_{FLAGS.batchnorm}_{FLAGS.nr_epochs}"

    overall_batch = 0
    min_valid_loss = 10000

    # training loop
    for epoch in range(FLAGS.nr_epochs):

        print(f"\nEpoch: {epoch}")

        for batch, (x, y) in enumerate(train_dl):
            nn.train()

            # squeeze the input, and put on device
            x = x.to(device)
            y = y.to(device)

            optimizer.zero_grad()

            # forward pass
            pred = nn(x).to(device)

            # compute loss and backpropagate
            loss = loss_function(pred, y)
            loss.backward()

            # update the weights
            optimizer.step()

            # save training loss
            training_losses.append(loss.item())

            # print(f"batch loss ({batch}): {loss.item()}")

            # get loss on validation set and evaluate
            if overall_batch % FLAGS.eval_freq == 0 and overall_batch != 0:
                valid_loss = eval_on_test(nn, loss_function, valid_dl, device)
                valid_losses.append(valid_loss)
                print(f"Training loss: {loss.item()} / Valid loss: {valid_loss}")
                if valid_loss < min_valid_loss:
                    print(f"Model is saved in epoch {epoch}, overall batch: {overall_batch}")
                    torch.save(nn.state_dict(), f"Models/Regression_{variables_string}.pt")
                    min_valid_loss = valid_loss
                    optimal_batch = overall_batch

            overall_batch += 1

    # Load the optimal model (with loweest validation loss, and evaluate on test set)
    optimal_nn = Regression(input_size, dnn_hidden_units, dropout_probs, 1, FLAGS.neg_slope, FLAGS.batchnorm).to(device)
    optimal_nn.load_state_dict(torch.load(f"Models/Regression_{variables_string}.pt"))

    test_loss, test_pred, test_true = eval_on_test(optimal_nn, loss_function, test_dl, device, verbose=FLAGS.verbose, return_preds=True)

    # save the test predictions of the regressor
    with open(f"Predictions/regression_test_preds{FLAGS.embedder}_{FLAGS.reduced_classes}_{FLAGS.impression}.pt", "wb") as f:
        pkl.dump(test_pred, f)

    print(f"Loss on test set of optimal model (batch {optimal_batch}): {test_loss}")

    significance_testing(test_pred, test_true, loss_function, FLAGS)

    if FLAGS.plotting:
        plotting(training_losses, valid_losses, test_loss, variables_string, optimal_batch, FLAGS)

def significance_testing(test_preds, test_labels, loss_fn, FLAGS):

    print("\nImpression:", FLAGS.impression)
    print("Reduced Classes:", FLAGS.reduced_classes)
    print(f"Embedder: {FLAGS.embedder}")

    print("Engagement levels:", Counter(test_labels))
    print(f"Total number of engagement levels: {len(test_labels)}\n")

    test_labels = torch.tensor(test_labels)
    test_preds = torch.tensor(test_preds)
    mean_eng = torch.mean(test_labels)
    median_eng = torch.median(test_labels)
    mode_eng = torch.mode(test_labels)[0]

    print(f"mean, median, mode: {mean_eng}, {median_eng}, {mode_eng}")

    mean = torch.full_like(test_labels, mean_eng)
    median = torch.full_like(test_labels, median_eng)
    mode = torch.full_like(test_labels, mode_eng)

    print("shapes", mean.shape, test_preds.shape, test_labels.shape)

    Model_MSE = loss_fn(test_preds.unsqueeze(1), test_labels.unsqueeze(1))
    print(f"Model loss: {Model_MSE}")
    MSE_mean = loss_fn(mean, test_labels)
    MSE_median = loss_fn(median, test_labels)
    MSE_mode = loss_fn(mode, test_labels)

    t_mean, p_mean = ttest_rel(mean, test_preds)
    t_median, p_median = ttest_rel(median, test_preds)
    t_mode, p_mode = ttest_rel(mode, test_preds)

    print(f"MSE mean Loss: {MSE_mean}, p-value: {p_mean}")
    print(f"MSE median Loss: {MSE_median}, p-value: {p_median}")
    print(f"MSE mode Loss: {MSE_mode}, p-value: {p_mode}")

def eval_on_test(nn, loss_function, dl, device, verbose=False, return_preds=False):
    """
    Find the accuracy and loss on the test set, given the current weights
    """
    all_predictions = []
    all_labels = []

    nn.eval()
    if verbose:
        print(f"neural net:\n {[param.data for param in nn.parameters()]}")

    nn.to(device)
    with torch.no_grad():
        losses = []
        for i, (x, y) in enumerate(dl):
            x = x.to(device)
            y = y.to(device)

            test_pred = nn(x).to(device)

            loss = loss_function(test_pred, y)

            losses.append(loss.item())

            if return_preds:
                all_predictions.extend(test_pred.squeeze().tolist())
                all_labels.extend(y.squeeze().tolist())

            if verbose and i == 0:
                print(test_pred)

    mean_losses = np.mean(losses)

    if not return_preds:
        return mean_losses

    return mean_losses, all_predictions, all_labels

def plotting(train_losses, valid_losses, test_loss, variables_string, optimal_batch, FLAGS):
    plt.rcParams.update({"font.size": 22})

    os.makedirs("Images", exist_ok=True)

    plt.figure(figsize=(15, 9))
    steps_all = np.arange(0, len(train_losses))
    steps_valid = np.arange(0, len(valid_losses)) * FLAGS.eval_freq

    # plot the losses
    plt.plot(steps_all, train_losses, '-', lw=2, label="Training loss")
    plt.plot(steps_valid, valid_losses, '-', lw=3, label="Validation loss")
    plt.axhline(test_loss, label="Test loss", color="red", lw=3)
    plt.axvline(optimal_batch, linestyle="dashed", color='red', label="Optimal model", lw=3)
    # plt.title('Losses over training, including final test loss using optimal model')

    plt.xlabel('Batch')
    plt.ylabel('MSE Loss')
    plt.grid(True)
    plt.legend()

    plt.tight_layout()

    fig_name = f"lossplot_{variables_string}.png"
    plt.savefig(f"Images/{fig_name}")

def print_flags():
    """
    Prints all entries in FLAGS variable.
    """
    for key, value in vars(FLAGS).items():
        print(key + ' : ' + str(value))

def main():
    """
    Main function
    """
    # Print all Flags to confirm parameter settings
    print("Training regression with following parameters:")
    print_flags()

    # Run the training operation
    train()

if __name__ == '__main__':
    # Command line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument('--dnn_hidden_units', type=str,
        default=DNN_HIDDEN_UNITS_DEFAULT,
        help='Comma separated list of number of units in each hidden layer')
    parser.add_argument('--dropout_probs', type=str, default= DROPOUT_DEFAULT,
      help='Comma separated list of dropout probabilities in each layer')
    parser.add_argument('--learning_rate', type=float,
        default=LEARNING_RATE_DEFAULT, help='Learning rate')
    parser.add_argument('--nr_epochs', type=int, default = NR_EPOCHS_DEFAULT,
      help='Number of epochs to run trainer.')
    parser.add_argument('--batch_size', type=int, default = BATCH_SIZE_DEFAULT,
      help='Batch size to run trainer.')
    parser.add_argument('--eval_freq', type=int, default=EVAL_FREQ_DEFAULT,
    help='Frequency of evaluation on the test set')
    parser.add_argument('--data_dir', type = str, default = DATA_DIR_DEFAULT,
      help='Directory for storing input data')
    parser.add_argument('--neg_slope', type=float, default=NEG_SLOPE_DEFAULT,
      help='Negative slope parameter for LeakyReLU')
    parser.add_argument('--optimizer', type=str, default="Adam",
      help='Type of optimizer')
    parser.add_argument('--amsgrad', type=int, default=1,
                        help='Boolean: Amsgrad for Adam and Adamw')
    parser.add_argument('--batchnorm', type=int, default=1,
                        help='Boolean: apply batch normalization?')
    # 0.0001 seems optimal
    parser.add_argument('--weightdecay', type=float, default=0,
      help='weight decay for optimizer')
    parser.add_argument('--momentum', type=float, default=0,
      help='momentum for optimizer')
    parser.add_argument('--embedder', type=str, default='Bert',
      help='which dataset to use: TFIDF or Bert')
    parser.add_argument('--verbose', type=int, default=0,
      help='print neural net and predictions')
    parser.add_argument('--reduced_classes', type=int, default=0,
      help='Use only 2 class dataset')
    parser.add_argument('--impression', type=int, default=1,
      help='If true, filter low impression instances out')
    parser.add_argument('--plotting', type=int, default=1,
      help='if true, plots are saved')

    FLAGS, unparsed = parser.parse_known_args()
    FLAGS.amsgrad = bool(FLAGS.amsgrad)
    FLAGS.batchnorm = bool(FLAGS.batchnorm)
    FLAGS.verbose = bool(FLAGS.verbose)
    FLAGS.reduced_classes = bool(FLAGS.reduced_classes)
    FLAGS.impression = bool(FLAGS.impression)
    FLAGS.plotting = bool(FLAGS.plotting)

    main()


from datetime import datetime
import unittest

from psycopg2.errors import UniqueViolation

from tests.implementations import create_test_database, TweetPgsqlRepository, \
    TweetStatsRepository

class TestItemToInsertStatement(unittest.TestCase):

    def test_trouble_case(self):
        repo = TweetStatsRepository()
        item = {
            'tweet_id': 'Ugx7w7FYrLQILWq-TE14AaABAg',
            'collected_at': datetime(2021, 11, 29, 3, 13, 53, 809005),
            'num_likes': 0}
        sql, _ = repo._item_to_insert_statement(item, upsert=True)
        self.assertIn('SET num_likes = ', sql)

class TestPostgresRepository(unittest.TestCase):

    def test_get_conditions_and_values(self):
        kwargs = {
            'a': None,
            'b': 1,
            'c': 'three',
        }
        conditions, values = TweetPgsqlRepository._get_conditions_and_values(
            **kwargs)
        expected_conditions = 'b = %s, c = %s'
        self.assertEqual(expected_conditions, conditions)
        expected_values = [1, 'three']
        self.assertEqual(expected_values, values)

    def test_get_selector_returns_star_if_no_projection(self):
        kwargs = {}
        selector = TweetPgsqlRepository._get_selector(**kwargs)
        expected = '*'
        self.assertEqual(expected, selector)

    def test_get_selector_returns_star_if_projection_is_none(self):
        kwargs = {'projection': None}
        selector = TweetPgsqlRepository._get_selector(**kwargs)
        expected = '*'
        self.assertEqual(expected, selector)

    def test_get_selector_returns_projection_if_specified(self):
        kwargs = {'projection': ['attr1', 'attr2']}
        selector = TweetPgsqlRepository._get_selector(**kwargs)
        expected = 'attr1,attr2'
        self.assertEqual(expected, selector)

    def test_add_inserts_item(self):
        db_name = 'test_add_inserts_item'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        tweet = {'label': None, 'tweet_id': 1, 'tweet': 'tweet1'}
        repo.add(tweet)
        item = repo.get(1)
        self.assertEqual(tweet, item)

    def test_add_ignores_duplicates_when_requested(self):
        db_name = 'test_add_ignores_duplicates_when_requested'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        tweet = {'tweet_id': 1, 'tweet': 'tweet1'}
        repo.add(tweet)
        error = False
        try:
            repo.add(tweet, ignore_duplicates=True)
        except UniqueViolation:
            error = True
        self.assertFalse(error)

    def test_add_errors_duplicates_when_requested(self):
        db_name = 'test_add_errors_duplicates_when_requested'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        tweet = {'tweet_id': 1, 'tweet': 'tweet1'}
        repo.add(tweet)
        error = False
        try:
            repo.add(tweet, ignore_duplicates=False)
        except UniqueViolation:
            error = True
        self.assertTrue(error)

    def test_add_many(self):
        db_name = 'test_add_many'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        tweet1 = {'tweet_id': 1, 'tweet': 'tweet1'}
        tweet2 = {'tweet_id': 2, 'tweet': 'tweet2'}
        repo.add_many([tweet1, tweet2])
        self.assertTrue(repo.exists(1))
        self.assertTrue(repo.exists(2))

    def test_add_many_rolls_back_on_error(self):
        db_name = 'test_add_many_rolls_back_on_error'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        tweet1 = {'tweet_id': 1, 'tweet': 'tweet1'}
        tweet2 = {'tweet_id': 1, 'tweet': 'tweet1'}
        try:
            repo.add_many([tweet1, tweet2])
        except UniqueViolation:
            pass
        self.assertFalse(repo.exists(1))

    def test_all_returns_all_items(self):
        db_name = 'test_all_returns_all_items'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        tweet1 = {'label': None, 'tweet_id': 1, 'tweet': 'tweet1'}
        tweet2 = {'label': None, 'tweet_id': 2, 'tweet': 'tweet2'}
        repo.add(tweet1)
        repo.add(tweet2)
        items = list(repo.all())
        expected = [tweet1, tweet2]
        self.assertEqual(expected, items)

    def test_count_returns_number_of_items_in_table(self):
        db_name = 'test_count_returns_number_of_items_in_table'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        tweet1 = {'tweet_id': 1, 'tweet': 'tweet1'}
        tweet2 = {'tweet_id': 2, 'tweet': 'tweet2'}
        tweet3 = {'tweet_id': 3, 'tweet': 'tweet3'}
        repo.add(tweet1)
        repo.add(tweet2)
        repo.add(tweet3)
        count = repo.count()
        self.assertEqual(3, count)

    def test_delete_removes_items(self):
        db_name = 'test_delete_removes_items'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        tweet1 = {'tweet_id': 1, 'tweet': 'tweet1'}
        repo.add(tweet1)
        self.assertTrue(repo.exists(1))
        repo.delete(dict(tweet_id=1))
        self.assertFalse(repo.exists(1))

    def test_delete_many(self):
        db_name = 'test_delete_many'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        tweet1 = {'tweet_id': 1, 'tweet': 'tweet1'}
        tweet2 = {'tweet_id': 2, 'tweet': 'tweet2'}
        repo.add_many([tweet1, tweet2])
        self.assertTrue(repo.exists(1))
        self.assertTrue(repo.exists(2))
        repo.delete_many([{'tweet_id': 1}, {'tweet_id': 2}])
        self.assertFalse(repo.exists(1))
        self.assertFalse(repo.exists(2))

    def test_exists_returns_true_when_item_exists(self):
        db_name = 'test_exists_returns_true_when_item_exists'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        tweet1 = {'tweet_id': 1, 'tweet': 'tweet1'}
        repo.add(tweet1)
        self.assertTrue(repo.exists(1))

    def test_exists_returns_false_when_item_does_not_exist(self):
        db_name = 'test_exists_returns_false_when_item_does_not_exist'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        self.assertFalse(repo.exists(1))

    def test_get_returns_none_when_item_does_not_exist(self):
        db_name = 'test_get_returns_none_when_item_does_not_exist'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        item = repo.get(1)
        self.assertIsNone(item)

    def test_get_returns_item_when_exists(self):
        db_name = 'test_get_returns_item_when_exists'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        tweet1 = {'tweet_id': 1, 'tweet': 'tweet1'}
        repo.add(tweet1)
        item = repo.get(1)
        self.assertIsInstance(item, dict)
        self.assertIsNotNone(item)

    def test_search_returns_correct_items(self):
        db_name = 'test_search_returns_correct_items'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        tweet1 = {'tweet_id': 1, 'tweet': 'tweet1'}
        tweet2 = {'tweet_id': 2, 'tweet': 'tweet1'}
        tweet3 = {'tweet_id': 3, 'tweet': 'tweet3'}
        repo.add(tweet1)
        repo.add(tweet2)
        repo.add(tweet3)
        items = list(repo.search(tweet='tweet1'))
        self.assertEqual(2, len(items))

    def test_update_single_item(self):
        db_name = 'test_update_single_item'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        tweet = {'tweet_id': 1, 'tweet': 'tweet1', 'label': 'a'}
        repo.add(tweet)
        result = repo.get(1)
        self.assertEqual('a', result['label'])
        tweet['tweet'] = 'edited tweet'
        tweet['label'] = 'b'
        repo.update(tweet, ['tweet_id'], ['tweet', 'label'])
        result = repo.get(1)
        self.assertEqual('edited tweet', result['tweet'])
        self.assertEqual('b', result['label'])

    def test_update_many_items(self):
        db_name = 'test_update_many_items'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        tweet1 = {'tweet_id': 1, 'tweet': 'tweet1', 'label': 'a'}
        tweet2 = {'tweet_id': 2, 'tweet': 'tweet2', 'label': 'a'}
        repo.add(tweet1)
        repo.add(tweet2)
        result = repo.get(1)
        self.assertEqual('a', result['label'])
        result = repo.get(2)
        self.assertEqual('a', result['label'])
        tweet1['label'] = 'b'
        tweet2['label'] = 'b'
        repo.update_many([tweet1, tweet2], ['tweet_id'], ['label'])
        result = repo.get(1)
        self.assertEqual('b', result['label'])
        result = repo.get(2)
        self.assertEqual('b', result['label'])

    def test_upsert_ok_when_not_exists(self):
        db_name = 'test_upsert_ok_when_not_exists'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        tweet = {'tweet_id': 1, 'tweet': 'tweet1', 'label': 'a'}
        repo.upsert(tweet)
        result = repo.get(1)
        self.assertEqual(tweet, result)

    def test_upsert_ok_when_exists(self):
        db_name = 'test_upsert_ok_when_exists'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        tweet = {'tweet_id': 1, 'tweet': 'tweet1', 'label': 'a'}
        repo.upsert(tweet)
        result = repo.get(1)
        self.assertEqual(tweet, result)
        error = False
        try:
            repo.upsert(tweet)
        except UniqueViolation:
            error = True
        self.assertFalse(error)
        result = repo.get(1)
        self.assertEqual(tweet, result)

    def test_upsert_many(self):
        db_name = 'test_upsert_many'
        create_test_database(db_name)
        repo = TweetPgsqlRepository(db_name=db_name)
        tweet1 = {'tweet_id': 1, 'tweet': 'tweet1', 'label': 'a'}
        repo.add(tweet1)

        tweet1['label'] = 'b'
        tweet2 = {'tweet_id': 2, 'tweet': 'tweet2', 'label': 'b'}
        repo.upsert_many([tweet1, tweet2])

        result = repo.get(1)
        self.assertEqual('b', result['label'])
        result = repo.get(2)
        self.assertEqual('b', result['label'])

    def test_upsert_with_two_primary_keys(self):
        db_name = 'test_upsert_with_two_primary_keys'
        create_test_database(db_name)
        repo = TweetStatsRepository(db_name=db_name)
        now = datetime.now()
        tweet_stats = {
            'tweet_id': 1,
            'collected_at': now,
            'num_likes': 1}
        repo.upsert(tweet_stats)
        result = repo.get(1, now)
        self.assertEqual(tweet_stats, result)


"""
@brief      test log(time=4s)
@author     Xavier Dupre
"""
import os
import unittest
import warnings
from docutils.parsers.rst import directives
from pyquickhelper.helpgen import rst2html
from pyquickhelper.pycode import ExtTestCase
from pyquickhelper.sphinxext.sphinx_runpython_extension import RunPythonDirective

class TestRunPythoncontext(ExtTestCase):

    def test_runpython_context(self):

        from docutils import nodes

        class runpythonthis_node(nodes.Structural, nodes.Element):
            pass

        class RunPythonThisDirective(RunPythonDirective):
            runpython_class = runpythonthis_node

        def visit_rp_node(self, node):
            self.add_text("
visit_rp_node
")

        def depart_rp_node(self, node):
            self.add_text("
depart_rp_node
")

        content = """
                    test a directive
                    ================

                    .. runpythonthis::
                        :showcode:
                        :store:

                        magic_store = "saved for later"

                    .. runpythonthis::
                        :restore:

                        a = 5
                        try:
                            print("restored", magic_store)
                        except Exception as e:
                            print("failed")
                            for k, v in sorted(locals().copy().items()):
                                print(k, "=", [v])
                    """.replace("                    ", "")

        tives = [("runpythonthis", RunPythonThisDirective, runpythonthis_node,
                  visit_rp_node, depart_rp_node)]

        rst = rst2html(content, writer="rst", keep_warnings=True,
                       directives=tives)
        self.assertIn("restored saved for later", rst)
        self.assertNotIn("failed", rst)

if __name__ == "__main__":
    unittest.main()


import argparse
import os
from utils import load_and_cache_examples, load_and_cache_unlabeled_examples, init_logger, load_tokenizer
from trainer import Trainer
import torch
import numpy as np
import random
import torch.nn as nn
from torch.utils.data import Subset
import json
import pickle
from transformers import AutoTokenizer
from eval import load_pred_data, save_data
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from collections import Counter

def set_seed(args):
    random.seed(args.train_seed)
    np.random.seed(args.train_seed)
    torch.manual_seed(args.train_seed)
    if args.n_gpu > 0  and torch.cuda.is_available():
        # print('yes')
        # assert 0
        torch.cuda.manual_seed_all(args.train_seed)
        torch.cuda.manual_seed(args.train_seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

def main(args):
    init_logger()
    set_seed(args)
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, use_fast=True)
    # tokenizer = load_tokenizer(args)

    train_dataset, num_labels, train_size  = load_and_cache_examples(args, tokenizer, mode="train")
    dev_dataset, num_labels,  dev_size = load_and_cache_examples(args, tokenizer, mode="dev")
    test_dataset, num_labels, test_size = load_and_cache_examples(args, tokenizer, mode="test")
    unlabeled_dataset, unlabeled_size = load_and_cache_unlabeled_examples(args, tokenizer, mode = 'unlabeled', train_size = train_size)
    contra_datasets = []

    print('number of labels:', num_labels)
    print('train_size:', train_size)
    print('dev_size:', dev_size)
    print('test_size:', test_size)
    print('unlabel_size:', unlabeled_size)

    trainer = Trainer(args, train_dataset=train_dataset, dev_dataset=dev_dataset,test_dataset=test_dataset, contra_datasets = contra_datasets, \
            unlabeled = unlabeled_dataset, num_labels = num_labels, data_size = train_size
            )
    if args.load_prev == 1:
        print(f"Load ckpt from {args.prev_ckpt}")
        trainer.load_model(path= args.prev_ckpt)

    trainer.init_model()
    if args.semi_method == 'ft':
        print("Begin Fine-Tuning!!")
        trainer.train(n_sample = len(train_dataset))
        train_pred, train_feat, train_label, unlabeled_pred, unlabeled_feat, unlabeled_label, unlabeled_pseudo = trainer.inference()
        save_data(train_pred, train_feat, train_label,  unlabeled_pred, unlabeled_feat, unlabeled_label, unlabeled_pseudo, mutual_info_bald = None, n_iter = args.round, dataset = args.task, ckpt = args.dr_model + f"_{args.dr_N}")
    else:
        print("Begin Self-Training!!")
        trainer.load_model(path= args.prev_ckpt)
        trainer.init_model()
        try:
            train_pred, train_feat, train_label, unlabeled_pred, unlabeled_feat, unlabeled_label, unlabeled_pseudo = load_pred_data(dataset = args.task, ckpt = args.dr_model, n_iter = args.round)
        except:
            train_pred, train_feat, train_label, unlabeled_pred, unlabeled_feat, unlabeled_label, unlabeled_pseudo = trainer.inference()
            save_data(train_pred, train_feat, train_label,  unlabeled_pred, unlabeled_feat, unlabeled_label, unlabeled_pseudo, mutual_info_bald = None, n_iter = args.round, dataset = args.task, ckpt = args.dr_model)
            train_pred, train_feat, train_label, unlabeled_pred, unlabeled_feat, unlabeled_label, unlabeled_pseudo = load_pred_data(dataset = args.task, ckpt = args.dr_model, n_iter = args.round)

        current_val = -np.max(unlabeled_pred, axis = -1)
        idx = np.argsort(current_val)
        cnt = Counter()
        labeled_idx = []
        for i in idx[:int(args.num_unlabeled)]:
            labeled_idx.append(i)
            cnt[np.argmax(unlabeled_pred[i])] += 1
        trainer.unlabeled = Subset(trainer.unlabeled, idx[:int(args.num_unlabeled)])
        trainer.selftrain()

    exit()

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument("--method", default='clean', type=str, help="which method to use")
    parser.add_argument("--gpu", default='0,1,2,3', type=str, help="which gpu to use")
    parser.add_argument("--n_gpu", default=1, type=int, help="which gpu to use")

    parser.add_argument("--seed", default=0, type=int, help="which seed to use")
    parser.add_argument("--train_seed", default=0, type=int, help="which seed to use")
    parser.add_argument("--task", default="agnews", type=str, help="The name of the task to train")

    parser.add_argument("--data_dir", default="../datasets", type=str,
                        help="The input data dir. Should contain the .tsv files (or other data files) for the task.")
    parser.add_argument("--model_dir", default="./model", type=str, help="Path to model")
    parser.add_argument("--eval_dir", default="./eval", type=str, help="Evaluation script, result directory")
    parser.add_argument("--train_file", default="train.tsv", type=str, help="Train file")
    parser.add_argument("--dev_file", default="dev.tsv", type=str, help="dev file")
    parser.add_argument("--test_file", default="test.tsv", type=str, help="Test file")
    parser.add_argument("--unlabel_file", default="unlabeled.tsv", type=str, help="Test file")
    parser.add_argument("--do_train", action="store_true", help="Whether to run training.")
    parser.add_argument("--do_eval", action="store_true", help="Whether to run eval on the test set.")

    parser.add_argument("--extra_dataset", default="", type=str, help="Whether to run extra eval on the test set.")

    parser.add_argument("--dev_labels", default=100, type=int, help="number of labels for dev set")
    parser.add_argument("--cache_dir", default="", type=str, help="Where do you want to store the pre-trained models downloaded from s3",)
    parser.add_argument("--output_dir", default=None, type=str, required=True, help="The output directory where the model predictions and checkpoints will be written.",)

    parser.add_argument('--logging_steps', type=int, default=20, help="Log every X updates steps.")
    parser.add_argument('--self_train_logging_steps', type=int, default=20, help="Log every X updates steps.")
    parser.add_argument('--save_steps', type=int, default=200, help="Save checkpoint every X updates steps.")
    parser.add_argument("--model_type", default="bert-base-uncased", type=str)
    parser.add_argument("--tokenizer", default="bert-base-uncased", type=str)

    parser.add_argument("--auto_load", default=1, type=int, help="Auto loading the model or not")
    parser.add_argument("--add_sep_token", action="store_true", help="Add [SEP] token at the end of the sentence")

    parser.add_argument("--num_train_epochs", default=3.0, type=float, help="Total number of training epochs to perform.")
    parser.add_argument("--max_steps", default=100, type=int, help="Training steps for initialization.")
    parser.add_argument("--weight_decay", default=1e-4, type=float, help="Weight decay if we apply some.")
    parser.add_argument("--adam_epsilon", default=1e-8, type=float, help="Epsilon for Adam optimizer.")
    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
    parser.add_argument("--warmup_steps", default=0, type=int, help="Linear warmup over warmup_steps.")
    parser.add_argument("--dropout_rate", default=0.1, type=float, help="Dropout for fully-connected layers")
    parser.add_argument("--batch_size", default=32, type=int, help="Batch size for training and evaluation.")
    parser.add_argument("--self_training_batch_size", default=32, type=int, help="Batch size for training and evaluation.")
    parser.add_argument("--eval_batch_size", default=256, type=int, help="Batch size for training and evaluation.")
    parser.add_argument("--learning_rate", default=1e-5, type=float, help="The initial learning rate for Adam.")
    parser.add_argument("--max_seq_len", default=128, type=int, help="The maximum total input sequence length after tokenization.")
    parser.add_argument("--max_seq_len_test", default=128, type=int, help="The maximum total input sequence length after tokenization.")

    parser.add_argument("--gradient_accumulation_steps", default=1, type=int, help="The maximum total input sequence length after tokenization.")
    parser.add_argument("--learning_rate_st", default=1e-5, type=float, help="The initial learning rate for self-training.")
    parser.add_argument('--self_training_max_step', type = int, default = 10000)
    parser.add_argument('--soft_label', type = int, default = 1)
    parser.add_argument('--self_training_update_period', type = int, default = 100, help = 'update period')
    parser.add_argument("--self_training_eps", default=0.6, type=float)
    parser.add_argument("--self_training_power", default=2, type=float)
    parser.add_argument("--self_training_weight", default=0.5, type=float)
    parser.add_argument("--gce_loss_q", default=0.5, type=float)

    parser.add_argument('--semi_method', type = str, default = 'ft')
    parser.add_argument('--k_cal', type = int, default = 5)

    parser.add_argument('--num_unlabeled', type = int, default = 100, help = 'number of unlabeled data.')

    parser.add_argument('--dr_model', type = str, default = 'arxiv_ckpt', help = 'the name for DR model.')
    parser.add_argument('--dr_N', type = int, default = 10, help = 'the number of retrieved documents.')

    parser.add_argument('--filter', type = int, default = 0, help = 'Whether use filter.')
    parser.add_argument('--template', type = int, default = 0, help = 'Prompt Template.')
    parser.add_argument('--round', type = int, default = 0, help = 'Prompt Template.')
    parser.add_argument('--load_prev', type = int, default = 0, help = 'whether load model from prev_round.')
    parser.add_argument('--prev_ckpt', type = str, default = '', help = 'dir of prev ckpt.')
    parser.add_argument('--train_full_labels', type = int, default = 0, help = 'whether train with full labels or not')

# semi_method
    args = parser.parse_args()
    if 'biobert' in args.model_type:
        args.model_name_or_path = 'dmis-lab/biobert-v1.1'
    elif 'scibert' in args.model_type:
        args.model_name_or_path = 'allenai/scibert_scivocab_uncased'
        args.tokenizer = 'allenai/scibert_scivocab_uncased'
    else:
        args.model_name_or_path = args.model_type
        args.tokenizer=args.model_type
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
    print(os.environ['CUDA_VISIBLE_DEVICES'])
    main(args)


from database.models import User, CarWash, Appointment, Base
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

class Database():
    def __init__(self, db_filename):
        self.engine = create_engine(db_filename)
        Base.metadata.create_all(self.engine)
        self.Session = sessionmaker(bind=self.engine)

    def add_user(self, record):
        session = self.Session()
        self.session.add(record)
        session.commit()
        session.close()

# database = Database('sqlite:///sqlite_pomoika.db')


"""

Function to estimate the required connection radii
for RGGs. This can be used to estimate the connection
radius required to obtain RGGs with a given expected degree
for RGGs with arbitary position distributions.

"""

#Import packages
import numpy as np
from scipy.spatial.distance import pdist
import periodic_distances as Cypdist
import scipy
import socket
import math
from scipy import special

def Get_Required_Kappa(Kappa_Desired , Network_Size, Position_Samples , d , Boundaries = 'S' , positions = None , pattern="uniform") :

	"""

	Function to determine the mean degree parameter required to obtain Kappa_Desired
	for an RGG with the specified conditions.

	Given a connection radius, r, the mean degree parameter k_tilde is the volume of the ball in
	d dimensions multiplied by the number of nodes. This is the expected mean degree for
	a homogenous RGG in a region with volume 1.

	This function supports generation of example point distributions for uniform in [0,1]^d with
	periodic or solid boundary conditions. For Non-uniform point distributions we must have the
	postions inputted as an argument.

	Parameters
	--------------

	Kappa_Desired : float

	Choosen mean degree

	Network_Size : int

	Number of nodes in the graph of interest.

	Position_Samples : int

	Reduced number of samples we take to estimate the distribution of distances. Given M sample positions we will obtain
	M(M-1)/2 samples from the distribution of distances.

	d : int

	dimension of the embedding space.

	boundaries : str (optional)

	choice of boundary condition. For RGGs in [0,1]^d
	we can choose between 'S' (solid) and 'P' (periodic
	or toroidal).

	positions : list

	list of of lists containing positions. Used if we
	have previously sampled the RGG positions.

	pattern : str

	Point pattern to sample node positions from:
	"uniform" and "gaussian" are possible choices.

	Returns
	----------

	d_threshold : float

	Estimated conneciton radius required
	to generate a network with the desired expected degree.

	"""

	#Make sure the required number of position samples is less than the network size:
	#An error will also occur when K >> N as the list coord will be greater than it should be.

	if Position_Samples > Network_Size :
		print("Error: position samples greater than network size")

	#If we don't get given a sample of positions then generate some:
	if positions is None and pattern == "uniform" :
		positions = np.random.uniform(0, 1.0, (Position_Samples, d))
	elif positions is None and pattern == "gaussian":
		positions = np.random.multivariate_normal(np.zeros(d), np.eye(d), size=Position_Samples)

	#Compute the array of distances (periodic distance computation sped up using Cython).
	if Boundaries == 'S' :
		Distances = pdist(positions)
	else :
		Distances = Cypdist.Periodic_Distances( positions , 1.0 )
		Distances = list(Distances)

	#Sort the distances:
	Sorted_Distances = np.sort(Distances)

	#Find the value required to get the edges above the required threshold:
	Edges_Required = (Kappa_Desired*Network_Size)/2.0

	#How far through the list to we have to be? F represents the 'fraction' through the list that we have got.
	F = (2*Edges_Required)/(Network_Size*(Network_Size-1))
	List_Coord = int(0.5*F*(Position_Samples*(Position_Samples-1)))

	#Get the 'Edge coun'th' element (but check that it isn't greater than the list length)
	if List_Coord > len(Sorted_Distances) :
		d_threshold = Sorted_Distances[-1]
	else :
		d_threshold = Sorted_Distances[List_Coord]

	return d_threshold


#!/usr/bin/env python3
"""
Import terrain height data from OS OpenData files

Data requested from OS OpenData:
OS Terrain 50 - GB | Data type: Grid | Supply format: ASCII GRID AND GML (GRID)
From: https://www.ordnancesurvey.co.uk/opendatadownload/products.html

This provides a .zip folder, which this project expects to be unzipped into:
    "./OS - terr50_gagg_gb"

Inside that is data/ and a license
Inside data/ are a lot of folders for the Ordnance Survey national grid

In each of those folders are several zip files, which contain the data for
each of the grid squares. The important one of these is the .asc file
"""

from PIL import Image
import zipfile
import os
from shutil import rmtree
from math import floor
import numpy as np

class HeightCell:
    """
    A single cell of the height map from Ordnance Survey's site
    """
    def __init__(self, dimensions, heights, exclude=[]):
        """
        Initialise this HeightCell with the properties passed to it
        dimensions is (cols, rows, xcorner, ycorner, cellsize)
        """
        cols, rows, xcorner, ycorner, cellsize = dimensions

        self.xsize = cols
        self.ysize = rows
        self.xcorner = xcorner  # Lower left corner
        self.ycorner = ycorner  # Lower left corner
        self.size = cellsize
        self.heights = heights
        self.exclude = exclude

        valid_values = list(filter(lambda x: x is not None, self.flattened))
        self.max = max(valid_values)
        self.min = min(valid_values)

    @property
    def flattened(self):
        """
        Return a flattened list of the heights in the cell
        Yes, I'm aware of the irony in the name
        """
        flat = [item for sublist in self.heights for item in sublist]
        return [h if h not in self.exclude else None for h in flat]

def importAsc(asc_file_name):
    """
    Import a .asc file from Ordnance Survey as a HeightCell
    """

    with open(asc_file_name, "r") as asc_file:
        lines = asc_file.readlines()

    # First 5 lines are dimensions
    dimensions = (int(l.split(" ")[-1]) for l in lines[:5])

    # Lines after this are the height data
    heights = [list(map(float, l.rstrip().split(" "))) for l in lines[5:]]

    # Some values may need to be excluded
    exclude = []

    height_cell = HeightCell(dimensions, heights, exclude)

    return height_cell

def saveCellAsImage(cell, image_name):
    """
    Save a HeightCell's height data to an image
    cell should be a valid HeightCell instance
    """

    def scaleBetween(x, in_min, in_max, out_min, out_max):
        return (x - in_min) * (out_max - out_min) / (in_max - in_min) + out_min

    # Scale point data so it goes 0-255
    # Excluded values are set to 0 (black)
    scaleTo255 = lambda x: scaleBetween(x, cell.min, cell.max, 0, 255) if x is not None else 0
    scaled_pixels = list(map(scaleTo255, cell.flattened))

    # Save image
    img = Image.new("L", (cell.xsize, cell.ysize))
    img.putdata(scaled_pixels)
    img.save(image_name)

def chooseSquare(base_dir):
    """
    Choose one of the 100x100km squares from the OS National Grid
    Prompts the user to enter the two-letter name of the chosen square
    """
    valid_square_names = sorted(os.listdir(base_dir))

    print("OS National Grid squares:")
    for a, b, c, d in zip(valid_square_names[::4], valid_square_names[1::4],
                          valid_square_names[2::4], valid_square_names[3::4]):
        print("{}, {}, {}, {}".format(a, b, c, d))

    square_name = input("Enter a square name, or multiple space-separated:\n")

    square_names = square_name.lower().split(" ")

    for name in square_names:
        if name not in valid_square_names:
            raise ValueError("\"{}\" not a valid square name".format(name))

    return square_names

def extractAscsFromSquare(base_dir, tmp_dir, square_name):
    """
    Extract all .asc files from the .zip files in the chosen square
    The files are extracted into tmp_dir

    Returns list of the file names
    """
    square_base_dir = os.path.join(base_dir, square_name)

    # Find out if this square is already extracted
    imported_asc_files = sorted(os.listdir(tmp_dir))
    is_selected_square = (f[:2].lower() == square_name for f in imported_asc_files)
    already_extracted = any(is_selected_square)

    # No need to extract any files - just return this square's ones
    if already_extracted:
        return filter(lambda f: f[:2].lower() == square_name,
                      imported_asc_files)

    # Otherwise get the files from the zip file
    file_names = []

    for item in os.listdir(square_base_dir):
        if not item.endswith(".zip"):
            continue    # Skip non-zip files

        item_path = os.path.join(square_base_dir, item)

        with zipfile.ZipFile(item_path, "r") as zip_ref:
            contents = zip_ref.namelist()
            asc_files = filter(lambda f: f.endswith(".asc"), contents)

            for asc_file in asc_files:
                zip_ref.extract(asc_file, tmp_dir)
                file_names.append(asc_file)

    return sorted(file_names)

def extractCellDataFromAscs(map_data_dir, asc_list):
    """
    Return HeightCell instances for all .asc files in asc_list
    The asc files should be in map_data_dir already
    """

    return [importAsc(os.path.join(map_data_dir, a)) for a in asc_list]

def getDimensions(height_cells):
    """
    Get dimensions needed for scaling
    Return dict:
        cell_size   = Cell size in metres
        cell_side   = Number of measurements (pixels) along one side of a cell
        cell_rows   = Number of rows of cells in the image
        img_width   = Width in pixels of the output image
        img_height  = Height in pixels of the output image
        min_x       = Minimum x coordinate (metres)
        min_y       = Minimum y coordinate (metres)
    """
    # Setup image grid parameters
    cell_size = 10000           # In metres
    measurement_interval = 50   # 50m per measurement
    cell_side = cell_size / measurement_interval
    cell_res = cell_side ** 2

    # Get image dimensions from real-life size
    x_corners = [c.xcorner for c in height_cells]
    y_corners = [c.ycorner for c in height_cells]
    max_x = max(x_corners) + cell_size
    max_y = max(y_corners) + cell_size
    min_x = min(x_corners)
    min_y = min(y_corners)

    ground_width = (max_x - min_x)
    ground_height = (max_y - min_y)
    cell_cols = ground_width // cell_size
    cell_rows = ground_width // cell_size
    img_width = int(ground_width * cell_side // cell_size)
    img_height = int(ground_height * cell_side // cell_size)

    return {"cell_size": cell_size,
            "cell_side": cell_side,
            "cell_rows": cell_rows,
            "img_width": img_width,
            "img_height": img_height,
            "min_x": min_x,
            "min_y": min_y,
            "max_y": max_y}

def scaleHeightData(heights):
    """
    Scale height data from its current range to 0-255 for image output
    """

    min_height = heights.min()

    # Shift any negative values up
    if min_height < 0:
        heights -= min_height

    # Scale heights to 255
    heights *= 255.0 / heights.max()
    return heights

def combineCells(height_cells):
    """
    Combine all height cells into one array and return it
    """

    dims = getDimensions(height_cells)

    # Use zero for default height (sea)
    heights_combined = np.zeros((dims["img_height"], dims["img_width"]))

    for cell in height_cells:
        # Get the start x and y indices of the cell
        cell_col = (cell.xcorner - dims["min_x"]) / dims["cell_size"]
        cell_row = (dims["max_y"] - cell.ycorner) / dims["cell_size"] - 1
        cell_start_x = int(cell_col * dims["cell_side"])
        cell_start_y = int(cell_row * dims["cell_side"])

        # Add actual height data to heights_combined
        for row, row_data in enumerate(cell.heights):
            for col, h in enumerate(row_data):
                # Get the pixel locations of this pixel
                px_col = cell_start_x + col
                px_row = cell_start_y + row

                heights_combined[px_row][px_col] = h

    return heights_combined

def makeImage(base_dir, map_data_dir, image_name, square_names, verbose=False):
    """
    Generate an image from the chosen square
    """

    if verbose:
        print("Importing map squares from .asc files")

    asc_files = []
    for square in square_names:
        # Add all .asc files from this square
        asc_files.extend(
            extractAscsFromSquare(base_dir, map_data_dir, square))

    # Import the cells as HeightCell objects
    height_cells = extractCellDataFromAscs(map_data_dir, asc_files)

    if verbose:
        print("Merging cells into one array")

    # Merge all HeightCells into one array
    heights = combineCells(height_cells)

    if verbose:
        print("Scaling to useful pixel brightness range")

    # Scale to 0-255
    heights = scaleHeightData(heights)

    if verbose:
        print("Saving image")

    # Ensure extension on image name
    if not image_name.endswith(".png"):
        image_name += ".png"

    # Save image
    img = Image.new("L", heights.shape[::-1])
    img.putdata(heights.flatten())
    img.save(image_name)

    # Remove map_data_dir so files aren't included in the next iteration
    #rmtree(map_data_dir)

    if verbose:
        print("Done!")

def interactiveMakeImage(base_dir, map_data_dir, image_name, verbose):
    """
    Interactively select grid squares and generate an image
    """

    square_names = chooseSquare(base_dir)

    image_name_new = input(
            "Image name (default=\"{}\"): ".format(image_name))

    if image_name_new != "":
        image_name = image_name_new

    makeImage(base_dir, map_data_dir, image_name, square_names, verbose)

if __name__ == "__main__":
    base_dir = os.path.join(".", "OS - terr50_gagg_gb", "data")
    map_data_dir = os.path.join(".", "map_data")
    image_name = "image"

    interactiveMakeImage(base_dir, map_data_dir, image_name, verbose=True)


# pylint: disable=R0904,C0103
""":class:`geoutils.GdeltConfig` tests.

"""
import unittest2
import tempfile
import os

import geoutils
from geosutils.files import (copy_file,
                             remove_files)
import geolib_mock

class TestGdeltDaemon(unittest2.TestCase):
    """:class:`geoutils.GdeltDaemon` test cases.
    """
    @classmethod
    def setUpClass(cls):
        conf = os.path.join('geoutils',
                            'tests',
                            'files',
                            'proxy.properties')
        cls._mock = geolib_mock.MockServer(conf)
        cls._mock.start()

        cls._conf = geoutils.GdeltConfig()
        cls._conf.threads = 1
        cls._conf.inbound_dir = tempfile.mkdtemp()

        cls._gdelt_spatial_index_table_name = 'gdelt_spatial_index'
        cls._audit_table_name = 'audit'

    def setUp(self):
        self._gdeltd = geoutils.GdeltDaemon(pidfile=None, conf=self._conf)

        self._ds = geoutils.Datastore()
        self._ds.host = 'localhost'
        self._ds.port = 42425
        self._ds.user = 'root'

    def test_init(self):
        """Initialise a GdeltDaemon object.
        """
        msg = 'Object is not a geoutils.GdeltDaemon'
        self.assertIsInstance(self._gdeltd, geoutils.GdeltDaemon, msg)

    def test_start_dry_run(self):
        """GdeltDaemon dry run.
        """
        self._ds.connect()
        self._ds.init_table(self._gdelt_spatial_index_table_name)
        self._ds.init_table(self._audit_table_name)

        gdelt_dir = os.path.join('geoutils', 'daemon', 'tests', 'files')
        gdelt_file = '20141124.export.CSV.zip'
        target_file = os.path.join(self._conf.inbound_dir, gdelt_file)
        copy_file(os.path.join(gdelt_dir, gdelt_file), target_file)

        old_dry = self._gdeltd.dry

        self._gdeltd.dry = True
        self._gdeltd._start(self._gdeltd.exit_event)

        # Clean up.
        self._gdeltd.dry = old_dry
        self._gdeltd.exit_event.clear()
        remove_files(target_file)
        self._ds.delete_table(self._gdelt_spatial_index_table_name)
        self._ds.delete_table(self._audit_table_name)

    def tearDown(self):
        del self._gdeltd
        del self._ds

    @classmethod
    def tearDownClass(cls):
        cls._mock.stop()

        del cls._gdelt_spatial_index_table_name
        del cls._audit_table_name

        os.removedirs(cls._conf.inbound_dir)
        del cls._conf


import logging

def __init_log(logfile: str):
  log_level = logging.INFO
  if logfile is None:
    log_stream = sys.stdout

  log_format = "%(asctime)s - %(levelname)-5s %(name)s %(lineno)s - %(message)s"
  log_date_format = "%Y-%m-%d %H:%M:%S"
  if logfile is None:
    logging.basicConfig(level=log_level, stream=log_stream, format=log_format, datefmt=log_date_format)
  else:
    logging.basicConfig(level=log_level, filename=logfile, format=log_format, datefmt=log_date_format)


# Dataset preprocessor to prepare the data for training
import numpy as np
import torch
import pickle
import os
import optuna
from time import time
import scipy.io as sio
import matplotlib.pyplot as plt

from NTdatasets.generic import GenericDataset
from NTdatasets.cumming.monocular import MultiDataset
import NTdatasets.conway.cloud_datasets as datasets

import ColorDataUtils.ConwayUtils as CU
import ColorDataUtils.EyeTrackingUtils as ETutils

# NDN tools
import NDNT.utils as utils # some other utilities
import NDNT.NDNT as NDN
from NDNT.modules.layers import *
from NDNT.networks import *
from time import time
from copy import deepcopy

import sys
sys.path.append('./lib')
import runner2 as r
import model as m
import experiment as exp
import plot

from NDNT.utils import imagesc   # because I'm lazy
from NDNT.utils import ss        # because I'm real lazy

class Model:
    def __init__(self, ndn_model, LLs, trial):
        self.ndn_model = ndn_model
        self.LLs = LLs
        self.trial = trial

datadir = '/home/dbutts/ColorV1/Data/'
dirname = '/home/dbutts/ColorV1/CLRworkspace/'
fn = 'Jocamo_220715_full_CC_ETCC_nofix_v08'
#fn = 'Jocamo_220727_full_CC_ETCC_nofix_v08'
num_lags=12

t0 = time()
data = datasets.ColorClouds(
    datadir=datadir, filenames=[fn], eye_config=3, drift_interval=16,
    luminance_only=True, binocular=False, include_MUs=True, num_lags=num_lags,
    trial_sample=True)
t1 = time()
print(t1-t0, 'sec elapsed')

NT = data.robs.shape[0]
NA = data.Xdrift.shape[1]
print("%d (%d valid) time points"%(NT, len(data)))
#data.valid_inds = np.arange(NT, dtype=np.int64)

lam_units = np.where(data.channel_map < 32)[0]
ETunits = np.where(data.channel_map >= 32)[0]
UTunits = np.where(data.channel_map >= 32+127)[0]

print( "%d laminar units, %d ET units"%(len(lam_units), len(ETunits)))

# Replace DFs
matdat = sio.loadmat(datadir+'Jocamo_220715_full_CC_ETCC_nofix_v08_DFextra.mat')
data.dfs = torch.tensor( matdat['XDF'][:NT, :], dtype=torch.float32 )

# Pull correct saccades
matdat = sio.loadmat(datadir+'Jocamo_220715_full_CC_ETCC_v08_ETupdate.mat')
sac_ts_all = matdat['ALLsac_bins'][0, :]

data.process_fixations( sac_ts_all )
sac_tsB = matdat['sac_binsB'][0, :]
sac_tsL = matdat['sac_binsL'][0, :]
sac_tsR = matdat['sac_binsR'][0, :]

NFIX = torch.max(data.fix_n).detach().numpy()
print(NFIX, 'fixations')
et1kHzB = matdat['et1kHzB']
et60B = matdat['et60HzB']
et60all = matdat['et60Hz_all']

print(lam_units, 'laminar units', data.robs.shape, 'robs', data.dfs.shape, 'dfs')

# Set cells-to-analyze and pull best model configuration and mus
Reff = torch.mul(data.robs[:, lam_units], data.dfs[:, lam_units]).numpy()
nspks = np.sum(Reff, axis=0)
a = np.where(nspks > 10)[0]
vallam = lam_units[a]
NCv = len(vallam)
print("%d out of %d units used"%(len(vallam), len(lam_units)))

## CONVERT LLsNULL, which is based on

# Read in previous data
dirname2 = dirname+'0715/et/'
matdat = sio.loadmat(dirname2+'LLsGLM.mat')
Dreg = matdat['Dreg']
top_cornerUT = matdat['top_corner'][:, 0]

data.set_cells(vallam)

# Load shifts and previous models
dirname2 = dirname+'0715/et/'
SHfile = sio.loadmat( dirname2 + 'BDshifts1.mat' )
fix_n = SHfile['fix_n']
shifts = SHfile['shifts']
metricsLL = SHfile['metricsLL']
metricsTH = SHfile['metricsTH']
ETshifts = SHfile['ETshifts']
ETmetrics = SHfile['ETmetrics']
Ukeeps = SHfile['Ctrain']
XVkeeps = SHfile['Cval']

top_corner_lam = [938, 515]

# Make 60x60 STAs (and GLMs)
Xshift = 0 #8+4
Yshift = 0 #-10+4
NX = 60

new_tc = np.array([top_corner_lam[0]-Xshift, top_corner_lam[1]-Yshift], dtype=np.int64)
data.draw_stim_locations(top_corner = new_tc, L=NX)

data.assemble_stimulus(top_corner=[new_tc[0], new_tc[1]], L=NX, fixdot=0, shifts=-shifts)

goodfix = np.where(ETmetrics[:,1] < 0.80)[0]
valfix = torch.zeros([ETmetrics.shape[0], 1], dtype=torch.float32)
valfix[goodfix] = 1.0
# Test base-level performance (full DFs and then modify DFs)
#DFsave = deepcopy(data2.dfs)  # this is also in data.dfs
data.dfs_out *= valfix

device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
device0 = torch.device("cpu")
dtype = torch.float32

goodfix = np.where(ETmetrics[:,1] < 0.80)[0]
valfix = torch.zeros([ETmetrics.shape[0], 1], dtype=torch.float32)
valfix[goodfix] = 1.0
# Test base-level performance (full DFs and then modify DFs)
#DFsave = deepcopy(data2.dfs)  # this is also in data.dfs
data.dfs_out *= valfix
print("%0.1f%% fixations remaining"%(100*len(goodfix)/ETmetrics.shape[0]))

lbfgs_pars = utils.create_optimizer_params(
    optimizer_type='lbfgs',
    tolerance_change=1e-8,
    tolerance_grad=1e-8,
    history_size=100,
    batch_size=20,
    max_epochs=3,
    max_iter = 500,
    device = device)

adam_parsT = utils.create_optimizer_params(
    optimizer_type='AdamW', batch_size=5, num_workers=0,
    learning_rate=0.01, early_stopping_patience=10,  # changed from 4
    optimize_graph=False, weight_decay = 0.2)
adam_parsT['device'] = device
adam_parsT['accumulated_grad_batches']=6

#set up fits
Treg = 1
Xreg = 20 # [20]
L1reg = 0.1 # [0.5]
GLreg = 10.0 # [4.0]

# drift network
drift_pars1 = NDNLayer.layer_dict(
    input_dims=[1,1,1,NA], num_filters=1, bias=False, norm_type=0, NLtype='lin')
drift_pars1['reg_vals'] = {'d2t': Dreg, 'bcs':{'d2t':0} }
# for stand-alone drift model
drift_pars1N = deepcopy(drift_pars1)
drift_pars1N['NLtype'] = 'softplus'
drift_net =  FFnetwork.ffnet_dict( xstim_n = 'Xdrift', layer_list = [drift_pars1] )

# glm net
glm_layer = Tlayer.layer_dict(
    input_dims=data.stim_dims, num_filters=1, bias=False, num_lags=num_lags,
    NLtype='lin', initialize_center = True)
glm_layer['reg_vals'] = {'d2x': Xreg, 'd2t': Treg, 'l1': L1reg, 'glocalx': GLreg,'edge_t':10}
stim_net =  FFnetwork.ffnet_dict( xstim_n = 'stim', layer_list = [glm_layer] )

# gqm net
num_subs = 2
gqm_layer = Tlayer.layer_dict(
    input_dims=data.stim_dims, num_filters=num_subs, num_inh=0, bias=False, num_lags=num_lags,
    NLtype='square', initialize_center = True)
gqm_layer['reg_vals'] = {'d2x': Xreg, 'd2t': Treg, 'l1': L1reg, 'glocalx': GLreg,'edge_t':10}
stim_qnet =  FFnetwork.ffnet_dict( xstim_n = 'stim', layer_list = [gqm_layer] )

#combine glm
comb_layer = NDNLayer.layer_dict(
    num_filters = 1, NLtype='softplus', bias=False)
comb_layer['weights_initializer'] = 'ones'

net_comb = FFnetwork.ffnet_dict(
    xstim_n = None, ffnet_n=[0,1],
    layer_list = [comb_layer], ffnet_type='add')

#combine gqm
comb2_layer = ChannelLayer.layer_dict(
    num_filters = 1, NLtype='softplus', bias=False)
comb2_layer['weights_initializer'] = 'ones'

net2_comb = FFnetwork.ffnet_dict(
    xstim_n = None, ffnet_n=[0,1,2],
    layer_list = [comb2_layer], ffnet_type='normal')
net2_comb['layer_list'][0]['bias'] = True

X2opt = np.zeros([NCv,2])
T2opt = np.zeros([NCv,2])
GLopt = np.zeros([NCv,2])
L1opt = np.zeros([NCv,2])
LLsNULL = np.zeros(NCv)
LLsLR = np.zeros([NCv,4])
LLsQR = np.zeros([NCv,4])

rvals = [0.001, 0.01, 0.1, 1, 10, 100, 1000, 10000]
rvalsG = [0.001, 0.1, 1, 10, 100, 1000] # glocal
rvalsL = [0.001, 0.01, 0.1, 1] # L1

matt_drifts = []
dan_drifts = []
matt_glms = []
dan_glms = []
for cc in range(NCv):
    matt_drift_filename = 'models/optim/matt_drift_cc'+str(cc)+'.pkl'

    # continue if the file already exists
    if os.path.isfile(matt_drift_filename):
        # load the model and continue
        print('loading model', cc)
        with open(matt_drift_filename, 'rb') as f:
            matt_drifts.append(pickle.load(f))

    matt_glm_filename = 'models/optim/matt_glm_cc'+str(cc)+'.pkl'

    # # continue if the file already exists
    if os.path.isfile(matt_glm_filename):
        # load the model and continue
        print('loading model', cc)
        with open(matt_glm_filename, 'rb') as f:
            matt_glms.append(pickle.load(f))

matt_gqms = []
for cc in range(NCv):
    gqm_filename = 'models/optim/matt_gqm_cc'+str(cc)+'.pkl'

    drift_weights = matt_drifts[cc].ndn_model.networks[0].layers[0].weight.data[:,0]

    # # continue if the file already exists
    if os.path.isfile(gqm_filename):
        # load the model and continue
        print('loading model', cc)
        with open(gqm_filename, 'rb') as f:
            matt_gqms.append(pickle.load(f))
        continue

    data.set_cells([vallam[cc]])

    LLsNULL_cc = matt_drifts[cc].LLs

    # get the best reg_vals for the GLM
    #best_reg_vals = matt_glms[cc].ndn_model.networks[0].layers[0].reg.vals
    #stim_net['layer_list'][0]['reg_vals'] = deepcopy(best_reg_vals)
    matt_gqm = NDN.NDN(ffnet_list = [stim_net, drift_net, stim_qnet, net2_comb], loss_type='poisson')
    matt_gqm.networks[0].layers[0] = deepcopy(matt_glms[cc].ndn_model.networks[0].layers[0])
    matt_gqm.block_sample=True
    matt_gqm.networks[3].layers[0].set_parameters(val=False,name='weight')
    matt_gqm.networks[1].layers[0].weight.data[:,0] = deepcopy(
        matt_drifts[cc].ndn_model.networks[0].layers[0].weight.data[:,0])
    matt_gqm.networks[1].layers[0].set_parameters(val=False)

    matt_gqms_temp = []

    def objective(trial):
        lbfgs_pars = utils.create_optimizer_params(
            optimizer_type='lbfgs',
            tolerance_change=1e-10*trial.suggest_int('tolerance_change', 10, 1e6, log=True),
            tolerance_grad=1e-10*trial.suggest_int('tolerance_grad', 10, 1e6, log=True),
            history_size=100,
            batch_size=20,
            max_epochs=3,
            max_iter = 500,
            device = device)

        # adjust the stim_net regularization
        matt_gqm.networks[0].layers[0].reg.vals['d2x'] = 1e-5*trial.suggest_int('d2x', 10, 1e10, log=True)
        matt_gqm.networks[0].layers[0].reg.vals['d2t'] = 1e-5*trial.suggest_int('d2t', 10, 1e10, log=True)
        matt_gqm.networks[0].layers[0].reg.vals['d2xt'] = 1e-5*trial.suggest_int('d2xt', 10, 1e10, log=True)
        matt_gqm.networks[0].layers[0].reg.vals['l1'] = 1e-5*trial.suggest_int('l1', 10, 1e10, log=True)
        matt_gqm.networks[0].layers[0].reg.vals['glocalx'] = 1e-5*trial.suggest_int('glocalx', 10, 1e10, log=True)
        matt_gqm.networks[0].layers[0].reg.vals['edge_t'] = 1e-5*trial.suggest_int('edge_t', 10, 1e10, log=True)

        # adjust the stim_qnet regularization
        matt_gqm.networks[2].layers[0].reg.vals['d2x'] = 1e-5*trial.suggest_int('d2x', 10, 1e10, log=True)
        matt_gqm.networks[2].layers[0].reg.vals['d2t'] = 1e-5*trial.suggest_int('d2t', 10, 1e10, log=True)
        matt_gqm.networks[2].layers[0].reg.vals['d2xt'] = 1e-5*trial.suggest_int('d2xt', 10, 1e10, log=True)
        matt_gqm.networks[2].layers[0].reg.vals['l1'] = 1e-5*trial.suggest_int('l1', 10, 1e10, log=True)
        matt_gqm.networks[2].layers[0].reg.vals['glocalx'] = 1e-5*trial.suggest_int('glocalx', 10, 1e10, log=True)
        matt_gqm.networks[2].layers[0].reg.vals['edge_t'] = 1e-5*trial.suggest_int('edge_t', 10, 1e10, log=True)

        matt_gqm.fit( data, force_dict_training=True, **lbfgs_pars)
        LL = matt_gqm.eval_models(data[data.val_blks], null_adjusted=False)[0]

        gqm_model = Model(matt_gqm, LLsNULL_cc-LL, trial)
        matt_gqms_temp.append(gqm_model)

        return LLsNULL_cc-LL

    study = optuna.create_study(direction='maximize', pruner=optuna.pruners.HyperbandPruner())

    study.optimize(objective, n_trials=40)

    matt_gqms.append(matt_gqms_temp[study.best_trial.number])

    with open(gqm_filename, 'wb') as f:
        matt_gqms[cc].trial = study
        pickle.dump(matt_gqms[cc], f)

    print(study.best_trial.number, study.best_params)


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

# Define your item pipelines here
#
# Don't forget to add your pipeline to the ITEM_PIPELINES setting
# See: https://doc.scrapy.org/en/latest/topics/item-pipeline.html
import pymongo
from scrapy.conf import settings
import redis

class ZhihucrawlPipeline(object):
    def __init__(self):
        host = settings["MONGODB_HOST"]
        port = settings["MONGODB_PORT"]
        dbname = settings["MONGODB_DBNAME"]
        sheetname = settings["MONGODB_SHEETNAME"]

        # 创建MONGODB数据库链接
        client = pymongo.MongoClient(host=host, port=port)
        # 指定数据库
        db = client[dbname]
        # 存放数据的数据库表名
        self.post = db[sheetname]

        # host = settings["REDIS_HOST"]
        # port = settings["REDIS_PORT"]

        # 创建redis数据库链接
        # client = redis.Redis(host=host, port=port)

        # 连接池管理redis连接
        # pool=redis.ConnectionPool(host=host,port=port)
        # r=redis.Redis(connection_pool=pool)

    def process_item(self, item, spider):
        data=dict(item)
        self.post.insert(data)
        return item

# from pymongo import MongoClient
#
# client=MongoClient("mongodb://127.0.0.1:27017,127.0.0.1:27021,127.0.0.1:27022",replicaset='repset')
# print( client.python.find_one())


from GoogleNews import GoogleNews
googlenews = GoogleNews()

googlenews = GoogleNews(lang='en')
googlenews = GoogleNews(period='7d')
googlenews = GoogleNews(encode='utf-8')
googlenews = GoogleNews(start='02/01/2019', end='02/28/2020')

# googlenews.set_lang('en')
# googlenews.set_period('7d')
# googlenews.set_encode('utf-8')

#googlenews.set_time_range('02/01/2020', '02/28/2020')

'''
googlenews.search('emeritus senior living')
# googlenews.get_news('APPLE')

# print(type(googlenews.get_texts()))

title_list = googlenews.get_texts()

diction = googlenews.results()
for item in diction:
    print(item)
    print("\n")
    print(item.keys())
    print("\n")
    print(item['desc'])
    print("\n\n")

# for item in title_list:
#    print(item + "\n")
'''

def get_title_text(search_string):
    googlenews.search(search_string)
    title_list = googlenews.get_texts()
    return title_list

#gets json of results
def get_results(search_string):
    googlenews.search(search_string)
    diction = googlenews.result()
    return diction