import torch
import torchvision
import torchvision.transforms as transforms
import os
import torch.nn.functional as F
from torch.autograd import Variable
import torch.nn as nn
import torch.optim as optim
from alex_stn import STNAlexNet
from lenet5 import LeNet5
import argparse
from lenet5_stn import LeNet5STN
from alexnet import AlexNet
from hdf5dataset import Hdf5Dataset, Hdf5DatasetMPI
from capsnet import CapsuleNet, CapsuleLoss, augmentation
from torch.optim.lr_scheduler import ReduceLROnPlateau
from convnets import ConvNetA, ConvNetB, ConvNetC, ConvNetD, ConvNetE, ConvNetF, ConvNetG, ConvNetH, ConvNetI, ConvNetJ, ConvNetK, ConvNetL, ConvNetM, ConvNetN
from datetime import datetime
import torchvision.models as models
from rq import get_current_job
import numpy as np
from PIL import ImageFile

ImageFile.LOAD_TRUNCATED_IMAGES = True

def train_generic_model(model_name="alexnet",
                        dataset="custom",
                        num_classes=-1,
                        batch_size=8,
                        is_transform=1,
                        num_workers=2,
                        lr_decay=1,
                        l2_reg=0,
                        hdf5_path="dataset-224x224.hdf5",
                        trainset_dir="./TRAIN_data_224",
                        testset_dir="./TEST_data_224",
                        convert_grey=False,
                        learning_rate=0.001,
                        num_epochs=4000):
    CHKPT_PATH = "./checkpoint_{}.PTH".format(model_name)
    print("CUDA:")
    print(torch.cuda.is_available())
    if is_transform:

        trans_ls = []
        if convert_grey:
            trans_ls.append(transforms.Grayscale(num_output_channels=1))
        trans_ls.extend(
            [
                transforms.Resize((512, 512)),
                # transforms.RandomCrop((224, 224)),
                # transforms.Grayscale(num_output_channels=1),
                transforms.ToTensor(),
                # transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
            ]
        )
        transform = transforms.Compose(trans_ls)
    else:
        transform = None

    print("DATASET FORMAT: {}".format(dataset))
    print("TRAINSET PATH: {}".format(trainset_dir))
    print("TESTSET PATH: {}".format(testset_dir))
    print("HDF5 PATH: {}".format(hdf5_path))
    if dataset == "custom":
        trainset = torchvision.datasets.ImageFolder(root=trainset_dir, transform=transform)
        train_size = len(trainset)
        testset = torchvision.datasets.ImageFolder(root=testset_dir, transform=transform)
        test_size = len(testset)
    elif dataset == "cifar":
        trainset = torchvision.datasets.CIFAR10(root="CIFAR_TRAIN_data", train=True, download=True, transform=transform)
        train_size = len(trainset)
        testset = torchvision.datasets.CIFAR10(root="CIFAR_TEST_data", train=False, download=True, transform=transform)
        test_size = len(testset)
    elif dataset == "hdf5":
        if num_workers == 1:
            trainset = Hdf5Dataset(hdf5_path, transform=transform, is_test=False)
        else:
            trainset = Hdf5DatasetMPI(hdf5_path, transform=transform, is_test=False)
        train_size = len(trainset)
        if num_workers == 1:
            testset = Hdf5Dataset(hdf5_path, transform=transform, is_test=True)
        else:
            testset = Hdf5DatasetMPI(hdf5_path, transform=transform, is_test=True)
        test_size = len(testset)

    train_loader = torch.utils.data.DataLoader(trainset,
                                               batch_size=batch_size,
                                               shuffle=True,
                                               num_workers=num_workers
                                               )

    test_loader = torch.utils.data.DataLoader(testset,
                                              batch_size=batch_size,
                                              shuffle=True,
                                              num_workers=num_workers
                                              )
    if model_name == "alexnet":
        net = AlexNet(num_classes=num_classes)
    elif model_name == "lenet5":
        net = LeNet5(num_classes=num_classes)
    elif model_name == "stn-alexnet":
        net = STNAlexNet(num_classes=num_classes)
    elif model_name == "stn-lenet5":
        net = LeNet5STN(num_classes=num_classes)
    elif model_name == "capsnet":
        net = CapsuleNet(num_classes=num_classes)
    elif model_name == "convneta":
        net = ConvNetA(num_classes=num_classes)
    elif model_name == "convnetb":
        net = ConvNetB(num_classes=num_classes)
    elif model_name == "convnetc":
        net = ConvNetC(num_classes=num_classes)
    elif model_name == "convnetd":
        net = ConvNetD(num_classes=num_classes)
    elif model_name == "convnete":
        net = ConvNetE(num_classes=num_classes)
    elif model_name == "convnetf":
        net = ConvNetF(num_classes=num_classes)
    elif model_name == "convnetg":
        net = ConvNetG(num_classes=num_classes)
    elif model_name == "convneth":
        net = ConvNetH(num_classes=num_classes)
    elif model_name == "convneti":
        net = ConvNetI(num_classes=num_classes)
    elif model_name == "convnetj":
        net = ConvNetJ(num_classes=num_classes)
    elif model_name == "convnetk":
        net = ConvNetK(num_classes=num_classes)
    elif model_name == "convnetl":
        net = ConvNetL(num_classes=num_classes)
    elif model_name == "convnetm":
        net = ConvNetM(num_classes=num_classes)
    elif model_name == "convnetn":
        net = ConvNetN(num_classes=num_classes)
    elif model_name == "resnet18":
        net = models.resnet18(pretrained=False, num_classes=num_classes)

    print(net)

    if torch.cuda.is_available():
        net = net.cuda()

    if model_name == "capsnet":
        criterion = CapsuleLoss()
    else:
        criterion = nn.CrossEntropyLoss()

    optimizer = optim.SGD(net.parameters(),
                          lr=learning_rate,
                          momentum=0.9,
                          weight_decay=l2_reg
                          )

    if lr_decay:
        scheduler = ReduceLROnPlateau(optimizer, 'min')

    best_acc = 0
    from_epoch = 0

    if os.path.exists(CHKPT_PATH):
        print("Checkpoint Found: {}".format(CHKPT_PATH))
        state = torch.load(CHKPT_PATH)
        net.load_state_dict(state['state_dict'])
        optimizer.load_state_dict(state['optimizer'])
        best_acc = state['best_accuracy']
        from_epoch = state['epoch']

    for epoch in range(from_epoch, num_epochs):
        #print("Epoch: {}/{}".format(epoch + 1, NUM_EPOCHS))
        epoch_loss = 0
        correct = 0
        for i, data in enumerate(train_loader, 0):
            #print("Train \t Epoch: {}/{} \t Batch: {}/{}".format(epoch + 1,
            #                                            NUM_EPOCHS,
            #                                            i + 1,
            #                                            ceil(train_size / BATCH_SIZE)))
            inputs, labels = data
            inputs, labels = Variable(inputs).type(torch.FloatTensor),\
                             Variable(labels).type(torch.LongTensor)

            if model_name == "capsnet":
                inputs = augmentation(inputs)
                ground_truth =  torch.eye(num_classes).index_select(dim=0, index=labels)

            if torch.cuda.is_available():
                inputs = inputs.cuda()
                labels = labels.cuda()

            optimizer.zero_grad()

            if model_name == "capsnet":
                classes, reconstructions = net(inputs, ground_truth)
                loss = criterion(inputs, ground_truth, classes, reconstructions)
            else:
                outputs = net(inputs)
                loss = criterion(outputs, labels)

            loss.backward()
            optimizer.step()
            epoch_loss += loss.data[0]
            if model_name != "capsnet":
                log_outputs = F.softmax(outputs, dim=1)
            else:
                log_outputs = classes
            pred = log_outputs.data.max(1, keepdim=True)[1]
            correct += pred.eq(labels.data.view_as(pred)).sum()

        print("Epoch: {} \t Training Loss: {:.4f} \t Training Accuracy: {:.2f} \t {}/{}".format(
            epoch+1,
            epoch_loss/train_size,
            100 * correct/train_size,
            correct,
            train_size
        )
        )

        loss_train = epoch_loss
        correct_train = correct

        correct = 0
        test_loss = 0
        for i, data in enumerate(test_loader, 0):
            # print("Test \t Epoch: {}/{} \t Batch: {}/{}".format(epoch + 1,
            #                                             NUM_EPOCHS,
            #                                             i + 1,
            #                                             ceil(test_size / BATCH_SIZE)))
            inputs, labels = data
            inputs, labels = Variable(inputs).type(torch.FloatTensor), Variable(labels).type(torch.LongTensor)

            if model_name == "capsnet":
                inputs = augmentation(inputs)
                ground_truth =  torch.eye(num_classes).index_select(dim=0, index=labels)

            if torch.cuda.is_available():
                inputs = inputs.cuda()
                labels = labels.cuda()

            if model_name == "capsnet":
                classes, reconstructions = net(inputs)
                loss = criterion(inputs, ground_truth, classes, reconstructions)
            else:
                outputs = net(inputs)
                loss = criterion(outputs, labels)

            test_loss += loss.data[0]

            if model_name != "capsnet":
                log_outputs = F.softmax(outputs, dim=1)
            else:
                log_outputs = classes

            pred = log_outputs.data.max(1, keepdim=True)[1]
            correct += pred.eq(labels.data.view_as(pred)).sum()
        print("Epoch: {} \t Testing Loss: {:.4f} \t Testing Accuracy: {:.2f} \t {}/{}".format(
            epoch+1,
            test_loss/test_size,
            100 * correct/test_size,
            correct,
            test_size
        )
        )
        if correct >= best_acc:
            if not os.path.exists("./models"):
                os.mkdir("./models")
            torch.save(net.state_dict(), "./models/model-{}-{}-{}-{}-val-acc-{:.2f}-train-{}-test-{}-epoch-{}.pb".format(
                model_name,
                dataset,
                hdf5_path,
                str(datetime.now()),
                100 * correct / test_size,
                trainset_dir.replace(" ", "_").replace("/", "_"),
                testset_dir.replace(" ", "_").replace("/", "_"),
                epoch+1
            )
                       )
        best_acc = max(best_acc, correct)

        # save checkpoint path
        state = {
            'epoch': epoch,
            'state_dict': net.state_dict(),
            'optimizer': optimizer.state_dict(),
            'best_accuracy': best_acc
        }
        torch.save(state, CHKPT_PATH)

        if lr_decay:
            # Note that step should be called after validate()
            scheduler.step(test_loss)

        job = get_current_job()
        job.meta['progress'] = {
            "epoch": int(epoch+1),
            "best_accuracy": float(best_acc/test_size),
            "testing": {
                "loss": float(test_loss / test_size),
                "accuracy": float(100 * correct / test_size),
                "correct": int(correct),
                "total": int(test_size)
            },
            "training": {
                "loss": float(loss_train / train_size),
                "accuracy": float(100 * correct_train / train_size),
                "correct": int(correct_train),
                "total": int(train_size)
            },
        }

        job.save_meta()
        with open("{}.npy".format(job.id), 'a') as f_handle:
            f_handle.write("{},{},{},{}\n".format(
                float(loss_train / train_size),
                float(100 * correct_train / train_size),
                float(test_loss / test_size),
                float(100 * correct / test_size)
            ))
        f_handle.close()

    print('Finished Training')

    print("")
    print("")

def wrap_train(
        q,
        model_name,
        dataset,
        num_classes,
        batch_size,
        is_transform,
        num_workers,
        lr_decay,
        l2_reg,
        hdf5_path,
        trainset_dir,
        testset_dir,
        convert_grey,
        learning_rate,
        num_epochs
):
    qj = q.enqueue(
        train_generic_model,
        args=(
            model_name,
            dataset,
            num_classes,
            batch_size,
            is_transform,
            num_workers,
            lr_decay,
            l2_reg,
            hdf5_path,
            trainset_dir,
            testset_dir,
            convert_grey,
            learning_rate,
            num_epochs,
        ),
        timeout="3000h"
        )
    return qj.id


"""
AI driven entity.
"""
from enum import Enum

import pygame

from map import Map
from config import CONFIG
from consts import Screens, IMAGE_FOLDER, BASE_SPEED, GRAVITY_CONST
from animation import Animation, BotAnimationType

class BotMoveType(Enum):
    UP = 0
    RIGHT = 1
    LEFT= 2
    STAND = 3

class Bot:

    def __init__(self, x_pos=0, y_pos=0, x_speed=0, y_speed=0):
        self.bot_move = BotMoveType.RIGHT
        self.x_pos = x_pos
        self.y_pos = y_pos
        self.x_speed = x_speed
        self.y_speed = y_speed
        self.max_speed = 2
        self.jump_count = 0
        self.jump_limit = 1
        self.jump_power = 7
        self.camera_x = 0
        self.camera_y = 0
        self.width = 64
        self.height = 64
        self.animation = Animation(IMAGE_FOLDER + 'player.png', self.width, self.height)

    def collision_detection(self, map):

        # Move on keypress
        if self.bot_move == BotMoveType.RIGHT:
            self.x_speed += BASE_SPEED
        elif self.bot_move == BotMoveType.LEFT:
            self.x_speed -= BASE_SPEED
        else:
            self.x_speed = 0

        # Check if the bottom of the player is masked
        on_ground = Map.is_masked_h_line(map, 4, self.x_pos, self.y_pos + self.height + self.y_speed, self.width / 2)
        if on_ground:
            self.y_speed = 0
            self.jump_count = 0
        elif self.y_speed < 25:  # Gravity
            self.y_speed += GRAVITY_CONST

        # Jumping
        if on_ground and self.jump_count < self.jump_limit and self.bot_move == BotMoveType.UP:
            self.y_speed -= self.jump_power
            self.jump_count += 1
            self.animation.animation_type = BotAnimationType.FRONT

        # Check the sides of the bot
        if Map.masked_top_v_line(map, 4, self.x_pos - self.width, self.y_pos, self.height) == 0:
            self.y_speed = 0
            if self.x_speed < 0:
                self.x_speed = 0
        elif Map.masked_top_v_line(map, 4, self.x_pos + self.width, self.y_pos, self.height) == 0:
            self.y_speed = 0
            if self.x_speed > 0:
                self.x_speed = 0

        # Check max speed
        if self.x_speed > self.max_speed:
            self.x_speed = self.max_speed
        elif self.x_speed < -self.max_speed:
            self.x_speed = -self.max_speed

        # Apply changes
        self.x_pos += self.x_speed
        self.y_pos += self.y_speed

        # Change animation on direction change
        if self.x_speed < 0:
            self.animation.animation_type = BotAnimationType.LEFT
        elif self.x_speed > 0:
            self.animation.animation_type = BotAnimationType.RIGHT
        elif self.x_speed == 0 and self.y_speed == 0:
            self.animation.animation_type = BotAnimationType.STAND

    def update_bot_state(self):
        self.bot_move = BotMoveType.UP  # Just jump for now

    def display(self, screen, map, player):
        self.camera_x = player.camera_x
        self.camera_y = player.camera_y

        self.collision_detection(map)
        self.update_bot_state()
        self.animation.play(screen, self)


import sys

def board_max_value(board, N):
    row_count, col_count = 1, 1
    # 행 개수 확인
    for i in range(N):
        temp = 1
        for j in range(1, N):
            if board[i][j-1] == board[i][j]:
                temp += 1
            else:
                row_count = max(row_count, temp)
                temp = 1
        row_count = max(row_count, temp)

    # 열 개수 확인
    for i in range(N):
        temp = 1
        for j in range(1, N):
            if board[j-1][i] == board[j][i]:
                temp += 1
            else:
                col_count = max(col_count, temp)
                temp = 1
        col_count = max(col_count, temp)

    return max(row_count, col_count)

# 0. N x N 행렬에 사탕을 채워넣는다.
N = int(sys.stdin.readline())
board = list([""] * N for _ in range(N))
d = [(1, 0), (-1, 0), (0, -1), (0, 1)]

for i in range(N):
    line = sys.stdin.readline().strip()
    for j in range(N):
        board[i][j] = line[j]

max_count = board_max_value(board, N)

# 1. 사탕의 색이 다른 인접한 두칸을 고른다.
for i in range(N):
    for j in range(N):
        for k in range(len(d)):
            dx = i + d[k][0]
            dy = j + d[k][1]
            # 2. 인접한 두칸을 교환하고 개수확인 후 다시 교환
            if 0 <= dx < N and 0 <= dy < N:
                if board[i][j] != board[dx][dy]:
                    board[i][j], board[dx][dy] = board[dx][dy], board[i][j]
                    max_count = max(max_count, board_max_value(board, N))
                    board[i][j], board[dx][dy] = board[dx][dy], board[i][j]

print(max_count)


# 2306. Naming a Company
# https://leetcode.com/problems/naming-a-company/

from collections import defaultdict

class Solution:
    def distinctNames(self, ideas):
        l=[set() for _ in range(26)]
        for i in ideas:
            l[ord(i[0])-97].add(i[1:])
        ans=0
        for i in range(25):
            for j in range(i+1,26):
                common=len(l[i]&l[j])
                ans+=(2*(len(l[i])-common)*(len(l[j])-common))
        return ans

        '''
        l=defaultdict(set)
        for i in ideas:
            l[i[0]].add(i[1:])
        ans=0
        for i,x in l.items():
            for j,y in l.items():
                if i>=j:
                    continue
                common=len(x&y)
                ans+=((len(x)-common)*(len(y)-common))
        return 2*ans
        '''

Weekly Contest/Weekly Contest 297/4_Naming a Company.py


'''Chico tem 1,50 m e cresce 2 centímetros por ano, enquanto Juca tem 1,10m e
cresce 3 centímetros por ano. Construir um algoritmo que calcule e imprima
quantos anos serão necessários para que Juca seja maior que Chico.'''

C = 1.50
J = 1.10
ano = 0

while(J < C):
    ano += 1
    J += 0.03
    C += 0.02

print("Anos:", ano)


from ctypes import (
    c_void_p,
    c_int,
    c_uint32,
    c_uint64,
    CFUNCTYPE,
    Structure,
    POINTER,
    byref,
    cast,
)
from enum import IntEnum

from ..ocf import OcfLib
from .data import Data

class IoDir(IntEnum):
    READ = 0
    WRITE = 1

class IoOps(Structure):
    pass

class Io(Structure):
    START = CFUNCTYPE(None, c_void_p)
    HANDLE = CFUNCTYPE(None, c_void_p, c_void_p)
    END = CFUNCTYPE(None, c_void_p, c_int)

    _instances_ = {}
    _fields_ = [
        ("_addr", c_uint64),
        ("_flags", c_uint64),
        ("_bytes", c_uint32),
        ("_class", c_uint32),
        ("_dir", c_uint32),
        ("_io_queue", c_void_p),
        ("_start", START),
        ("_priv1", c_void_p),
        ("_priv2", c_void_p),
        ("_handle", HANDLE),
        ("_end", END),
    ]

    @classmethod
    def from_pointer(cls, ref):
        c = cls.from_address(ref)
        cls._instances_[ref] = c
        OcfLib.getInstance().ocf_io_set_cmpl_wrapper(byref(c), None, None, c.c_end)
        return c

    @classmethod
    def get_instance(cls, ref):
        return cls._instances_[cast(ref, c_void_p).value]

    def del_object(self):
        del type(self)._instances_[cast(byref(self), c_void_p).value]

    def put(self):
        OcfLib.getInstance().ocf_io_put(byref(self))

    def get(self):
        OcfLib.getInstance().ocf_io_get(byref(self))

    @staticmethod
    @END
    def c_end(io, err):
        Io.get_instance(io).end(err)

    @staticmethod
    @START
    def c_start(io):
        Io.get_instance(io).start()

    @staticmethod
    @HANDLE
    def c_handle(io, opaque):
        Io.get_instance(io).handle(opaque)

    def end(self, err):
        try:
            self.callback(err)
        except:  # noqa E722
            pass

        self.del_object()
        self.put()

    def submit(self):
        return OcfLib.getInstance().ocf_volume_submit_io(byref(self))

    def submit_flush(self):
        return OcfLib.getInstance().ocf_volume_submit_flush(byref(self))

    def submit_discard(self):
        return OcfLib.getInstance().ocf_volume_submit_discard(byref(self))

    def submit_flush(self):
        return OcfLib.getInstance().ocf_volume_submit_flush(byref(self))

    def submit_discard(self):
        return OcfLib.getInstance().ocf_volume_submit_discard(byref(self))

    def set_data(self, data: Data, offset: int = 0):
        self.data = data
        OcfLib.getInstance().ocf_io_set_data(byref(self), data, offset)

IoOps.SET_DATA = CFUNCTYPE(c_int, POINTER(Io), c_void_p, c_uint32)
IoOps.GET_DATA = CFUNCTYPE(c_void_p, POINTER(Io))

IoOps._fields_ = [("_set_data", IoOps.SET_DATA), ("_get_data", IoOps.GET_DATA)]

lib = OcfLib.getInstance()
lib.ocf_io_set_cmpl_wrapper.argtypes = [POINTER(Io), c_void_p, c_void_p, Io.END]

lib.ocf_io_set_data.argtypes = [POINTER(Io), c_void_p, c_uint32]
lib.ocf_io_set_data.restype = c_int

lib.ocf_volume_submit_io.argtypes = [POINTER(Io)]
lib.ocf_volume_submit_io.restype = None

lib.ocf_volume_submit_flush.argtypes = [POINTER(Io)]
lib.ocf_volume_submit_flush.restype = None

lib.ocf_volume_submit_discard.argtypes = [POINTER(Io)]
lib.ocf_volume_submit_discard.restype = None


from collections import defaultdict

num_samples = len(snakemake.input)-1
dist_threshold = 10000

GeneID_to_location = {}
with open(snakemake.input[-1], "r") as input_coords_file:
	for line in input_coords_file:
		GeneID, transcript, chrom, start, stop = line.strip().split()
		#for flybase_dyak
		if "-" in GeneID:
			GeneID = GeneID.split("-")[0]
		start, stop = int(start), int(stop)
		if GeneID in GeneID_to_location:
			previous_chrom, previous_start, previous_stop = GeneID_to_location[GeneID]
			GeneID_to_location[GeneID] = [chrom, min(previous_start, start), max(previous_stop, stop)]
		else:
			GeneID_to_location[GeneID] = [chrom, start, stop]

geneid_to_diffs = defaultdict(list)
geneid_to_sames = defaultdict(list)
for i in range(num_samples):
	geneid_to_cluster_found = {}
	geneid_to_anchor_range = {}
	geneid_to_best_AS = {}
	with open(snakemake.input[i], "r") as input_AS_file:
		for line in input_AS_file:
			geneid, transcript_range, anchor_range, anchor_AS, strand = line.strip().split()
			if geneid not in geneid_to_cluster_found:
				geneid_to_cluster_found[geneid] = False
			if geneid not in geneid_to_best_AS:
				geneid_to_best_AS[geneid] = 0
			anchor_AS = int(anchor_AS)
			best_AS = geneid_to_best_AS[geneid]
			if anchor_AS > best_AS:
				geneid_to_cluster_found[geneid] = True
				geneid_to_anchor_range[geneid] = anchor_range
				geneid_to_best_AS[geneid] = anchor_AS
	for geneid in geneid_to_best_AS:
		cluster_found = geneid_to_cluster_found[geneid]
		if cluster_found:
			gene_chrom, gene_start, gene_stop = GeneID_to_location[geneid]
			anchor_range = geneid_to_anchor_range[geneid]
			anchor_chrom, anchor_startstop = anchor_range.split(":")
			anchor_start, anchor_stop = anchor_startstop.split("-")
			anchor_start, anchor_stop = int(anchor_start), int(anchor_stop)
			if gene_chrom != anchor_chrom or anchor_start > gene_stop + dist_threshold or gene_start > anchor_stop + dist_threshold:
				diff = True
				same = False
			else:
				same = True
				diff = False
			geneid_to_diffs[geneid].append(diff)
			geneid_to_sames[geneid].append(same)

with open(snakemake.output[0], "w") as output_alldiff_file, open(snakemake.output[1], "w") as output_allsame_file, open(snakemake.output[2], "w") as output_somesome_file:
	for geneid in geneid_to_diffs:
		sames = geneid_to_sames[geneid]
		diffs = geneid_to_diffs[geneid]
		if all(diffs):
			output_alldiff_file.write(f"{geneid}\n")
		elif all(sames):
			output_allsame_file.write(f"{geneid}\n")
		else:
			output_somesome_file.write(f"{geneid}\n")


from urllib.request import urlopen

from scrapy_news.spiders.newssite_util import nytimes_url, nytimes_page_info
from scrapy_news.spiders.util import get_page_addr, get_snapshot_number

def test_nytimes_url():
  true_tests = [
      'https://nytimes.com/pages/a/b/c.html',
      'https://www.nytimes.com/pages/a/b/c.html',
      'https://www.nytimes.com/2010/01/01/a.html',
      'https://www.nytimes.com/aponline/2010/01/01/a.html',
      'https://www.nytimes.com/aponline/2010/1/01/a.html',
      'https://www.nytimes.com/aponline/2010/1/1/a.html',
  ]

  false_tests = [
      'https://www.nytimes.com/pages', 'https://www.nytimes.com/201/1/1/a.html',
      'https://www.nytimes.com/interactive/a.html'
  ]

  print('testing test_nytimes_url')
  for test in true_tests:
    print(nytimes_url(test) == True)

  for test in false_tests:
    print(nytimes_url(test) == False)

def test_nytimes_page_info():

  print('testing nytimes_page_info')

  url = 'http://web.archive.org/web/20100101160117/http://nytimes.com/'
  response = urlopen(url)
  print(
      nytimes_page_info(response.read().decode('utf-8', 'ignore'), url)[1:] == (
          True, 20100101, False, None))

  url = 'http://web.archive.org/web/20091231025110/http://www.nytimes.com/pages/health/index.html'
  response = urlopen(url)
  print(
      nytimes_page_info(response.read().decode('utf-8', 'ignore'), url)[1:] == (
          True, 20091230, False, None))

  url = 'http://web.archive.org/web/20100105042000/http://www.nytimes.com:80/2010/01/01/world/asia/01khost.html?hp'
  response = urlopen(url)
  print(
      nytimes_page_info(response.read().decode('utf-8', 'ignore'), url)[1:] == (
          True, None, True, 20091231))

  url = 'http://web.archive.org/web/20091231034340/http://www.nytimes.com/glogin?URI=http://www.nytimes.com/2009/12/30/health/nutrition/30recipehealth.html&OQ=_rQ3D1Q26refQ3Dhealth&OP=67dc4632Q2FQ20pQ5EEQ20!)HQ7De))B5Q205ggPQ20-5Q20Q2FgQ20_Q5E@yB_Q20FnBeuBu)FQ20Q2FgeQ5EHumQ5E_Q5E@yB_X_BDy'
  response = urlopen(url)
  print(
      nytimes_page_info(response.read().decode('utf-8', 'ignore'), url)[1:] == (
          True, None, False, None))

def test_get_page_addr():
  domain = "nytimes.com"

  print('testing get_page_addr')

  print(
      get_page_addr("https://www.nytimes.com/pages/1.html", domain) ==
      'pages_1.html')
  print(
      get_page_addr("https://www.nytimes.com/2010/01/01/1.html", domain) ==
      '2010_01_01_1.html')
  print(
      get_page_addr("https://nytimes.com/2010/01/01/1.html", domain) ==
      '2010_01_01_1.html')

  print(get_page_addr("https://www.nytimes.com/", domain) == '')
  print(get_page_addr("https://www.nytimes.com", domain) == '')

  print(get_page_addr("https://www.nytimes.com:80/", domain) == '')
  print(get_page_addr("https://www.nytimes.com:80", domain) == '')

  print(get_page_addr("https://www.", domain) is None)

def test_get_snapshot_number():

  print('testing get_snapshot_number')

  print(
      get_snapshot_number(
          'https://web.archive.org/web/20100101000000/https://nytimes.com') ==
      '20100101000000')

  print(
      get_snapshot_number(
          'https://web.archive/web/20100101000000/https://nytimes.com') is None)

  print(
      get_snapshot_number(
          'https://web.archive.org/web/20100101000000id_/https://nytimes.com')
      is None)

  print(
      get_snapshot_number(
          'https://web.archive.org/web/20100101000000*/https://nytimes.com') is
      None)

if __name__ == '__main__':
  test_nytimes_url()
  test_get_page_addr()
  test_get_snapshot_number()
  test_nytimes_page_info()


import uuid

import test_v3

class CatalogTestCase(test_v3.RestfulTestCase):
    """Test service & endpoint CRUD"""

    def setUp(self):
        super(CatalogTestCase, self).setUp()

        self.service_id = uuid.uuid4().hex
        self.service = self.new_service_ref()
        self.service['id'] = self.service_id
        self.catalog_api.create_service(
            self.service_id,
            self.service.copy())

        self.endpoint_id = uuid.uuid4().hex
        self.endpoint = self.new_endpoint_ref(service_id=self.service_id)
        self.endpoint['id'] = self.endpoint_id
        self.catalog_api.create_endpoint(
            self.endpoint_id,
            self.endpoint.copy())

    # service crud tests

    def test_create_service(self):
        """POST /services"""
        ref = self.new_service_ref()
        r = self.post(
            '/services',
            body={'service': ref})
        return self.assertValidServiceResponse(r, ref)

    def test_list_services(self):
        """GET /services"""
        r = self.get('/services')
        self.assertValidServiceListResponse(r, ref=self.service)

    def test_list_services_xml(self):
        """GET /services (xml data)"""
        r = self.get('/services', content_type='xml')
        self.assertValidServiceListResponse(r, ref=self.service)

    def test_get_service(self):
        """GET /services/{service_id}"""
        r = self.get('/services/%(service_id)s' % {
            'service_id': self.service_id})
        self.assertValidServiceResponse(r, self.service)

    def test_update_service(self):
        """PATCH /services/{service_id}"""
        service = self.new_service_ref()
        del service['id']
        r = self.patch('/services/%(service_id)s' % {
            'service_id': self.service_id},
            body={'service': service})
        self.assertValidServiceResponse(r, service)

    def test_delete_service(self):
        """DELETE /services/{service_id}"""
        self.delete('/services/%(service_id)s' % {
            'service_id': self.service_id})

    # endpoint crud tests

    def test_list_endpoints(self):
        """GET /endpoints"""
        r = self.get('/endpoints')
        self.assertValidEndpointListResponse(r, ref=self.endpoint)

    def test_list_endpoints_xml(self):
        """GET /endpoints (xml data)"""
        r = self.get('/endpoints', content_type='xml')
        self.assertValidEndpointListResponse(r, ref=self.endpoint)

    def test_create_endpoint(self):
        """POST /endpoints"""
        ref = self.new_endpoint_ref(service_id=self.service_id)
        r = self.post(
            '/endpoints',
            body={'endpoint': ref})
        self.assertValidEndpointResponse(r, ref)

    def assertValidErrorResponse(self, response):
        self.assertTrue(response.status in [400])

    def test_create_endpoint_400(self):
        """POST /endpoints"""
        ref = self.new_endpoint_ref(service_id=self.service_id)
        ref["region"] = "0" * 256
        self.post('/endpoints', body={'endpoint': ref}, expected_status=400)

    def test_get_endpoint(self):
        """GET /endpoints/{endpoint_id}"""
        r = self.get(
            '/endpoints/%(endpoint_id)s' % {
                'endpoint_id': self.endpoint_id})
        self.assertValidEndpointResponse(r, self.endpoint)

    def test_update_endpoint(self):
        """PATCH /endpoints/{endpoint_id}"""
        ref = self.new_endpoint_ref(service_id=self.service_id)
        del ref['id']
        r = self.patch(
            '/endpoints/%(endpoint_id)s' % {
                'endpoint_id': self.endpoint_id},
            body={'endpoint': ref})
        self.assertValidEndpointResponse(r, ref)

    def test_delete_endpoint(self):
        """DELETE /endpoints/{endpoint_id}"""
        self.delete(
            '/endpoints/%(endpoint_id)s' % {
                'endpoint_id': self.endpoint_id})

    def test_create_endpoint_on_v2(self):
        # clear the v3 endpoint so we only have endpoints created on v2
        self.delete(
            '/endpoints/%(endpoint_id)s' % {
                'endpoint_id': self.endpoint_id})

        # create a v3 endpoint ref, and then tweak it back to a v2-style ref
        ref = self.new_endpoint_ref(service_id=self.service['id'])
        del ref['id']
        del ref['interface']
        ref['publicurl'] = ref.pop('url')
        ref['internalurl'] = None
        # don't set adminurl to ensure it's absence is handled like internalurl

        # create the endpoint on v2 (using a v3 token)
        r = self.admin_request(
            method='POST',
            path='/v2.0/endpoints',
            token=self.get_scoped_token(),
            body={'endpoint': ref})
        endpoint_v2 = r.body['endpoint']

        # test the endpoint on v3
        r = self.get('/endpoints')
        endpoints = self.assertValidEndpointListResponse(r)
        self.assertEqual(len(endpoints), 1)
        endpoint_v3 = endpoints.pop()

        # these attributes are identical between both API's
        self.assertEqual(endpoint_v3['region'], ref['region'])
        self.assertEqual(endpoint_v3['service_id'], ref['service_id'])
        self.assertEqual(endpoint_v3['description'], ref['description'])

        # a v2 endpoint is not quite the same concept as a v3 endpoint, so they
        # receive different identifiers
        self.assertNotEqual(endpoint_v2['id'], endpoint_v3['id'])

        # v2 has a publicurl; v3 has a url + interface type
        self.assertEqual(endpoint_v3['url'], ref['publicurl'])
        self.assertEqual(endpoint_v3['interface'], 'public')

        # tests for bug 1152632 -- these attributes were being returned by v3
        self.assertNotIn('publicurl', endpoint_v3)
        self.assertNotIn('adminurl', endpoint_v3)
        self.assertNotIn('internalurl', endpoint_v3)

        # test for bug 1152635 -- this attribute was being returned by v3
        self.assertNotIn('legacy_endpoint_id', endpoint_v3)


import json
import string

#create a list of tv shows or movies and add them to a json file
# read json file in if possible to make sure there are no duplicates in show name

class tv_movie:
    def __init__(self, name, isMovie, season):
        self.name = name
        self.isMovie = isMovie
        if(isMovie): #movies do not have seasons
            self.season = 0
        else:
            self.season = season

    def convert(self): #convert the class into a dictionary
        dictionary = {
            "Name" : self.name,
            "Movie" : str(self.isMovie),
            "Seasons Watched" : self.season
        }
        return dictionary

def validate_season():
    while True:
        num = input("How many seasons have you seen: ")
        try:
            #convert length to int
            value = int(num)
            if value > 0:
                break
            #non integer answers
        except ValueError:
            print("You did not enter a number. Please enter a number \n")
    return value

exists = input("Do you already have a file containing the movies and tv shows you have seen?\n")
while exists.lower() != "yes" and exists.lower() != "no":
    exists = input("Do you already have a file containing the movies and tv shows you have seen?\n")

#list of tv shows and movies watched
tvm_list = []

file_path = "tv_movie_list.json"
# if the file exists it should be saved under the name tv_movie_list.json
if exists.lower() == "yes":
    with open(file_path, "r", encoding = 'utf-8') as f:
        tvm_list = json.load(f)

more = "yes"

while more.lower() != "no":
    # start building the list of movies and tv shows
    tvm_name = input("\nEnter the name of a movie or tv show you have seen: ").capitalize()
    tvm_movie = input("Is it a movie?  ")
    while tvm_movie.lower() != "yes" and tvm_movie.lower() != "no":
        tvm_movie = input("Is it a movie?  ")

    if tvm_movie == "no":
        tvm_season = validate_season()
        tvm_bool = False
    else:
        tvm_season = 0
        tvm_bool = True
    #create the new entry
    new_tvm = tv_movie(tvm_name, tvm_bool, tvm_season)
    #add new entry to list of movies
    tvm_list.append(new_tvm.convert())
    more = input("Are there any more TV shows or Movies you would like to add? ")

with open(file_path, "w") as outfile:
    json.dump(tvm_list, outfile)

print("done")


""" Tests for add_flyer_views of the ecommerce app. """
import datetime

from django.core.urlresolvers import reverse
from django.template.defaultfilters import date as date_filter

from coupon.models import Slot, FlyerPlacement
from coupon.service.flyer_service import next_flyer_date
from ecommerce.service.calculate_current_price import calculate_current_price
from ecommerce.tests.ecommerce_test_case import EcommerceTestCase
from market.models import Site

class TestAddFlyerDates(EcommerceTestCase):
    """ Test class for add-flyer-dates view. """

    fixtures = ['test_advertiser', 'test_coupon', 'test_consumer',
        'test_geolocation']
    urls = 'urls_local.urls_2'

    def test_get_add_flyer_dates(self):
        """ Assert a logged in advertiser can get to the purchase flyer-by-date
        page with the default checked and previous days in the month not
        checked. """
        self.prep_for_flyer_purchase()
        response = self.client.get(reverse('add-flyer-dates',
            args=[self.slot.id]))
        self.assertContains(response,
            self.advertiser.site.get_or_set_consumer_count())
        self.assertContains(response,
            'checked="checked" name="%s" value="%s"' %
                (str(date_filter(next_flyer_date(), "F j")),
                str(next_flyer_date())))
        if int(date_filter(next_flyer_date(), "j")) > 7:
            previous_date = next_flyer_date() + datetime.timedelta(days=-7)
            self.assertContains(response,
                'disabled="disabled" type="checkbox" ' +
                'name="%s" value="%s" id="id_%s" />' % (
                (str(date_filter(previous_date, "F j")),
                 str(previous_date),
                 str(date_filter(previous_date, "F j")))))

    def test_post_add_flyer_dates(self):
        """ Assert when a logged in advertiser selects dates from the
        flyer-by-date page, they get pushed to the ecommerce page with the
        appropriate order items in the order summary display."""
        self.prep_for_flyer_purchase()
        first_purchase_date = next_flyer_date() + datetime.timedelta(days=14)
        second_purchase_date = next_flyer_date() + datetime.timedelta(days=28)
        first_purchase_date_key = date_filter(first_purchase_date, "F j")
        second_purchase_date_key = date_filter(second_purchase_date, "F j")
        kwargs = {
            str(first_purchase_date_key): str(first_purchase_date),
            str(second_purchase_date_key): str(second_purchase_date),
            'subdivision_consumer_count': 100}
        response = self.client.post(reverse('add-flyer-dates',
            args=[self.slot.id]), kwargs, follow=True)
        self.assertTemplateUsed(response,
            'ecommerce/display_checkout_coupon_purchase.html')
        self.common_flyer_purchase_asserts(response, first_purchase_date,
            second_purchase_date)

    def test_sold_market_flyer_date(self):
        """ Assert when the next_flyer_date() is sold out, the user can not
        select the date to purchase flyers on the sold out date.
        """
        count = 1
        while count <= 10:
            slot = Slot(site_id=2,
                business_id=count,
                renewal_rate=99,
                is_autorenew=True,
                end_date=datetime.date.today() + datetime.timedelta(
                    days=count))
            slot.save()
            flyer_placement = FlyerPlacement(
                site_id=2,
                slot=slot,
                send_date=next_flyer_date())
            flyer_placement.save()
            count += 1
        self.prep_for_flyer_purchase()
        response = self.client.get(reverse('add-flyer-dates',
            args=[self.slot.id]))
        self.assertContains(response, 'disabled="disabled" type="checkbox" ' +
            'name="%s" value="%s" id="id_%s" />' % (
            (str(date_filter(next_flyer_date(), "F j")),
            str(next_flyer_date()),
            str(date_filter(next_flyer_date(), "F j")))))

    def test_with_no_dates_selected(self):
        """ Test Add flyer purchase form submit when no dates selected. """
        kwargs = {'subdivision_consumer_count': 220}
        self.prep_for_flyer_purchase()
        response = self.client.post(reverse('add-flyer-dates',
            args=[self.slot.id]), kwargs)
        self.assertEqual(response.status_code, 200)
        self.assertTemplateUsed(response,
            'include/frm/frm_add_flyer_dates.html')
        self.assertContains(response,
            'You must choose at least one date to send out a flyer')

    def test_post_no_consumer_count(self):
        """ Test Add flyer purchase form submit when subdivision_consumer_count
        hidden field is missing.
        """
        kwargs = {}
        self.prep_for_flyer_purchase()
        response = self.client.post(reverse('add-flyer-dates',
            args=[self.slot.id]), kwargs, follow=True)
        self.assertEqual(response.redirect_chain[0][1], 302)
        self.assertEqual(response.request['PATH_INFO'],
            '/hudson-valley/contact-us/')

    def test_post_bad_consumer_count(self):
        """ Test Add flyer purchase form submit when subdivision_consumer_count
        hidden field has alphanumeric value.
        """
        kwargs = {'subdivision_consumer_count': 'a120'}
        self.prep_for_flyer_purchase()
        response = self.client.post(reverse('add-flyer-dates',
            args=[self.slot.id]), kwargs, follow=True)
        self.assertEqual(response.redirect_chain[0][1], 302)
        self.assertEqual(response.request['PATH_INFO'],
            '/hudson-valley/contact-us/')

class TestAddFlyerByMap(EcommerceTestCase):
    """ Test case for add_flyer_views views. """

    fixtures = ['test_advertiser', 'test_coupon', 'test_consumer']
    urls = 'urls_local.urls_2'

    def test_show_add_flyers(self):
        """ Assert advertiser with valid slot gets good add flyers form. """
        self.prep_for_flyer_purchase()
        response = self.client.get(reverse('add-flyer-by-map',
            args=[self.slot.id]), follow=True)
        self.assertEqual(response.status_code, 200)
        self.assertContains(response, 'form id="frm_add_flyer_by_map"')
        self.assertContains(response,
            'Consumers have requested Hudson Valley coupons by email')
        self.assertContains(response, 'Include zip codes within')

    def test_process_add_flyers(self):
        """
        Assert POST with a valid slot displays the flyers purchase form.
        """
        kwargs = {'zip_array':'1234,12345',
                  'county_array':'',
                  'subdivision_consumer_count':'123'}
        self.prep_for_flyer_purchase()
        response = self.client.post(reverse('add-flyer-by-map',
            args=[self.slot.id]), kwargs, follow=True)
        self.assertContains(response,
            'frm_add_flyer_dates')

class TestAddFlyerEntireMarket(EcommerceTestCase):
    """ Test case for add_flyer_views views. """

    fixtures = ['test_advertiser', 'test_coupon', 'test_consumer']
    urls = 'urls_local.urls_2'

    def test_show_add_flyers(self):
        """ Assert advertiser with valid slot gets good add flyers form. """
        self.prep_for_flyer_purchase()
        response = self.client.get(reverse('buy-market-flyer',
            args=[self.slot.id]), follow=True)
        self.assertEqual(response.status_code, 200)
        self.assertContains(response, 'frm_add_flyer_dates"')
        key_list = ['add_slot_choice', 'add_flyer_choice', 'flyer_dates_list']
        for key in key_list:
            this_key = self.client.session.get(key, None)
            self.assertEqual(this_key, None)
        subdivision_dict = self.client.session.get('subdivision_dict', None)
        self.assertEqual(self.client.session['locked_flyer_price'],
            calculate_current_price(product_id=1,
                consumer_count=Site.objects.get(
                    id=2).get_or_set_consumer_count()))
        self.assertEqual(self.client.session['locked_consumer_count'],
            Site.objects.get(id=2).get_or_set_consumer_count())
        self.assertEqual(subdivision_dict['county_array'], ())
        self.assertEqual(subdivision_dict['city_array'], ())
        self.assertEqual(subdivision_dict['zip_array'], ())
        self.assertEqual(subdivision_dict['subdivision_consumer_count'],
            Site.objects.get(id=2).get_or_set_consumer_count())


import mojeklasy

def main():
    point2Da = mojeklasy.Punkt2D(1, 5);
    point2Db = mojeklasy.Punkt2D(11, 5);
    point3D = mojeklasy.Punkt3D(1, 2, 3);

    point2Da.drukuj();

    print()
    point2Db.drukuj();

    print()
    point3D.drukuj();

    print()
    odcinek = mojeklasy.Odcinek(point2Da, point2Db);
    odcinek.licz()

if __name__ == '__main__':
    main();


import asyncio

global_data = {}

def process_data(data):
    command, args = data.split(' ', 1)
    if command == 'put':
        name, value, time = args.split(' ', 2)
        try:
            value = float(value)
            time = int(time)
        except ValueError:
            resp = 'error\nwrong command\n\n'
            return resp
        if name not in global_data:
            global_data[name] = {}
        global_data[name][time] = value
        resp = 'ok\n\n'
        return resp
    if command == 'get':
        if len(args.split(' ')) != 1:
            resp = 'error\nwrong command\n\n'
            return resp
        if args == '*\n':
            resp = 'ok'
            for name in global_data:
                for key in global_data[name]:
                    resp += '\n' + name + ' ' + \
                        str(global_data[name][key]) + ' ' + str(key)
            resp += '\n\n'

        else:
            args = args.replace('\n', '')
            if args not in global_data:
                resp = 'ok\n\n'
            else:
                resp = 'ok'
                for key in global_data[args]:
                    resp += '\n' + args + ' ' + \
                        str(global_data[args][key]) + ' ' + str(key)
                resp += '\n\n'
        print("get", resp)
        return resp

    resp = 'error\nwrong command\n\n'
    return resp

class ClientServerProtocol(asyncio.Protocol):
    def connection_made(self, transport):
        # print("Connection made ", transport) # это неважно
        self.transport = transport

    def data_received(self, data):
        print("data received ", data)  # данные пришли
        resp = process_data(data.decode())  # обработали сформировари ответ
        self.transport.write(resp.encode())  # отправили

def run_server(host, port):
    loop = asyncio.get_event_loop()
    coro = loop.create_server(
        ClientServerProtocol,
        host, port
    )

    server = loop.run_until_complete(coro)

    try:
        loop.run_forever()
    except KeyboardInterrupt:
        #print('shut this beatiful server up')
        pass

    server.close()
    loop.run_until_complete(server.wait_closed())
    loop.close()

#run_server('127.0.0.1', 8888)


from typing import Optional, Sequence, Union

import numpy as np
import pandas as pd
import torch as to
from matplotlib import pyplot as plt

import pyrado
from pyrado.sampling.bootstrapping import bootstrap_ci
from pyrado.utils.data_types import merge_dicts

def draw_dts(
    dts_policy: np.ndarray, dts_step: np.ndarray, dts_remainder: np.ndarray, y_top_lim: Optional[float] = None
):
    r"""
    Create a figure and draw the time intervals $\Delta_t$ of various parts of one time step.

    :param dts_policy: time it took to compute the policy's action
    :param dts_step: time it took to perform the
    :param dts_remainder: time it took to execute all remaining commands (e.g. soring the data)
    :param y_top_lim: upper bound for the y-axis in ms, no limit by default
    """
    fig, axs = plt.subplots(2, 1, figsize=(6, 8), tight_layout=True)

    x = np.arange(0, len(dts_policy))
    y1 = 1000 * dts_policy
    y2 = 1000 * dts_step
    y3 = 1000 * dts_remainder
    labels = [r"$\Delta$t policy [ms]", r"$\Delta$t step [ms]", r"$\Delta$t remainder [ms]"]

    axs[0].plot(x, y1, label=labels[0])
    axs[0].plot(x, y2, label=labels[1])
    axs[0].plot(x, y3, label=labels[2])
    axs[0].axhline(y=2, color="k")  # at 2 ms we exceed the sampling frequency of 500 Hz
    axs[0].legend(loc="upper right")
    axs[0].set_title("individual plots")

    axs[1].stackplot(x, y1, y2, y3, labels=labels)
    axs[1].axhline(y=2, color="k")  # at 2 ms we exceed the sampling frequency of 500 Hz
    axs[1].legend(loc="upper right")
    axs[1].set_title("stack plot")

    if y_top_lim is not None:
        for ax in axs:
            ax.set_ylim(top=y_top_lim)

def draw_curve_from_data(
    plot_type: str,
    ax: plt.Axes,
    data: Union[list, np.ndarray, to.Tensor, pd.DataFrame],
    x_grid: Union[list, np.ndarray, to.Tensor],
    ax_calc: int,
    x_label: Optional[Union[str, Sequence[str]]] = None,
    y_label: Optional[str] = None,
    curve_label: Optional[str] = None,
    area_label: Optional[str] = "",
    vline_level: Optional[float] = None,
    vline_label: str = "approx. solved",
    title: Optional[str] = None,
    show_legend: bool = True,
    cmp_kwargs: dict = None,
    plot_kwargs: dict = None,
    legend_kwargs: dict = None,
) -> plt.Figure:
    """
    Create a box or violin plot for a list of data arrays or a pandas DataFrame.
    The plot is neither shown nor saved.

    .. note::
        If you want to have a tight layout, it is best to pass axes of a figure with `tight_layout=True` or
        `constrained_layout=True`.

        If you want to order the 4th element to the 2nd position in terms of colors use
        .. code-block:: python

            palette.insert(1, palette.pop(3))

    :param plot_type: tye of 1-dim plot: `mean_std`, `min_mean_max`, or `ci_on_mean`
    :param ax: axis of the figure to plot on
    :param data: data to plot,me.g. a time series
    :param x_grid: values to plot the data over, e.g. time
    :param ax_calc: axis of the data array to calculate the mean, min and max, or std over
    :param x_label: labels for the categories on the x-axis, if `data` is not given as a `DataFrame`
    :param y_label: label for the y-axis, pass `None` to set no label
    :param curve_label: label of the (1-dim) curve, pass `None` for no label
    :param area_label: label of the (transparent) area, pass `None` for no label and "" for the default label
    :param vline_level: if not `None` (default) add a vertical line at the given level
    :param vline_label: label for the vertical line
    :param show_legend: if `True` the legend is shown, useful when handling multiple subplots
    :param title: title displayed above the figure, set to None to suppress the title
    :param cmp_kwargs: keyword arguments forwarded to functions computing the statistics of interest
    :param plot_kwargs: keyword arguments forwarded to the plotting` functions
    :param legend_kwargs: keyword arguments forwarded to pyplot's `legend()` function, e.g. `loc='best'`
    :return: handle to the resulting figure
    """
    plot_type = plot_type.lower()
    if plot_type not in ["mean_std", "min_mean_max", "ci_on_mean"]:
        raise pyrado.ValueErr(given=plot_type, eq_constraint="mean_std, min_mean_max, ci_on_mean")
    if not isinstance(data, (list, to.Tensor, np.ndarray, pd.DataFrame)):
        raise pyrado.TypeErr(given=data, expected_type=[list, to.Tensor, np.ndarray, pd.DataFrame])

    # Set defaults which can be overwritten by passing plot_kwargs
    cmp_kwargs = merge_dicts(
        [dict(num_reps=1000, confidence_level=0.9, bias_correction=False, studentized=False), cmp_kwargs]
    )

    if isinstance(data, pd.DataFrame):
        data = data.to_numpy()
    elif isinstance(data, list):
        data = np.array(data)
    elif isinstance(data, to.Tensor):
        data = data.detach().cpu().numpy()

    # Extract features from data
    data_mean = np.mean(data, axis=ax_calc)
    df = pd.DataFrame()
    df = df.assign(mean=data_mean)
    if plot_type == "mean_std":
        data_std = np.std(data, axis=ax_calc)
        df = df.assign(std=data_std)

    elif plot_type == "min_mean_max":
        data_min = np.min(data, axis=ax_calc)
        data_max = np.max(data, axis=ax_calc)
        df = df.assign(min=data_min)
        df = df.assign(max=data_max)

    elif plot_type == "ci_on_mean":
        _, data_lo, data_up = bootstrap_ci(
            data.T if ax_calc == 1 else data,
            stat_fcn=np.mean,
            num_reps=cmp_kwargs["num_reps"],
            alpha=cmp_kwargs["confidence_level"],
            ci_sides=2,
            bias_correction=cmp_kwargs["bias_correction"],
            studentized=cmp_kwargs["studentized"],
            seed=0,
        )
        df = df.assign(ci_lo=data_lo)
        df = df.assign(ci_up=data_up)

    # Forward the actual plotting
    return draw_curve(
        plot_type,
        ax,
        df,
        x_grid,
        x_label,
        y_label,
        curve_label,
        area_label,
        vline_level,
        vline_label,
        title,
        show_legend,
        plot_kwargs,
        legend_kwargs,
    )

def draw_curve(
    plot_type: str,
    ax: plt.Axes,
    data: pd.DataFrame,
    x_grid: Union[list, np.ndarray, to.Tensor],
    x_label: Optional[Union[str, Sequence[str]]] = None,
    y_label: Optional[str] = None,
    curve_label: Optional[str] = None,
    area_label: Optional[str] = "",
    vline_level: Optional[float] = None,
    vline_label: str = "approx. solved",
    title: Optional[str] = None,
    show_legend: bool = True,
    plot_kwargs: dict = None,
    legend_kwargs: dict = None,
) -> plt.Figure:
    """
    Create a box or violin plot for a list of data arrays or a pandas DataFrame.
    The plot is neither shown nor saved.

    .. note::
        If you want to have a tight layout, it is best to pass axes of a figure with `tight_layout=True` or
        `constrained_layout=True`.

        If you want to order the 4th element to the 2nd position in terms of colors use
        .. code-block:: python

            palette.insert(1, palette.pop(3))

    :param plot_type: tye of 1-dim plot: `mean_std`, `min_mean_max`, or `ci_on_mean`
    :param ax: axis of the figure to plot on
    :param data: pandas DataFrame containing the columns `mean`, `std`, `min`, and `max` depending on the `plot_type`
    :param x_grid: values to plot the data over, e.g. time
    :param x_label: labels for the categories on the x-axis, if `data` is not given as a `DataFrame`
    :param y_label: label for the y-axis, pass `None` to set no label
    :param curve_label: label of the (1-dim) curve, pass `None` for no label
    :param area_label: label of the (transparent) area, pass `None` for no label and "" for the default label
    :param vline_level: if not `None` (default) add a vertical line at the given level
    :param vline_label: label for the vertical line
    :param show_legend: if `True` the legend is shown, useful when handling multiple subplots
    :param title: title displayed above the figure, set to None to suppress the title
    :param plot_kwargs: keyword arguments forwarded to the plotting` functions
    :param legend_kwargs: keyword arguments forwarded to pyplot's `legend()` function, e.g. `loc='best'`
    :return: handle to the resulting figure
    """
    plot_type = plot_type.lower()
    if plot_type not in ["mean_std", "std", "min_mean_max", "ci_on_mean"]:
        raise pyrado.ValueErr(given=plot_type, eq_constraint="mean_std, min_mean_max, or ci_on_mean")
    if not isinstance(data, pd.DataFrame):
        raise pyrado.TypeErr(given=data, expected_type=pd.DataFrame)
    if x_label is not None and not isinstance(x_label, str):
        raise pyrado.TypeErr(given=x_label, expected_type=str)
    if y_label is not None and not isinstance(y_label, str):
        raise pyrado.TypeErr(given=y_label, expected_type=str)

    # Set defaults which can be overwritten by passing plot_kwargs
    plot_kwargs = merge_dicts([dict(alpha=0.3, num_std=2), plot_kwargs])
    num_std = plot_kwargs.pop("num_std")  # pop here since some plotting functions do not expect num_std
    legend_kwargs = dict() if legend_kwargs is None else legend_kwargs
    # palette = sns.color_palette() if palette is None else palette

    # Preprocess
    if isinstance(x_grid, list):
        x_grid = np.array(x_grid)
    elif isinstance(x_grid, to.Tensor):
        x_grid = x_grid.detach().cpu().numpy()

    # Plot
    if plot_type in ["mean_std", "std"]:
        if not ("mean" in data.columns and "std" in data.columns):
            raise pyrado.KeyErr(keys="'mean' and 'std'", container=data)
        if area_label == "":
            area_label = rf"$\pm {num_std}$ std"
        ax.fill_between(
            x_grid,
            data["mean"] - num_std * data["std"],
            data["mean"] + num_std * data["std"],
            label=area_label,
            **plot_kwargs,
        )

    elif plot_type == "min_mean_max":
        if not ("mean" in data.columns and "min" in data.columns and "max" in data.columns):
            raise pyrado.KeyErr(keys="'mean' and 'min' and 'max'", container=data)
        if area_label == "":
            area_label = r"min \& max"
        ax.fill_between(x_grid, data["min"], data["max"], label=area_label, **plot_kwargs)

    elif plot_type == "ci_on_mean":
        if not ("mean" in data.columns and "ci_lo" in data.columns and "ci_up" in data.columns):
            raise pyrado.KeyErr(keys="'mean' and 'ci_lo' and 'ci_up'", container=data)
        if area_label == "":
            area_label = r"conf. intvl."
        ax.fill_between(x_grid, data["ci_lo"], data["ci_up"], label=area_label, **plot_kwargs)

    # Plot the mean last for proper z-ordering
    if plot_type != "std":
        plot_kwargs["alpha"] = 1
        ax.plot(x_grid, data["mean"], label=curve_label, **plot_kwargs)

    # Postprocess
    if vline_level is not None:
        # Add dashed line to mark a threshold
        ax.axhline(vline_level, c="k", ls="--", lw=1.0, label=vline_label)

    if x_label is None:
        ax.get_xaxis().set_ticks([])

    if y_label is not None:
        ax.set_ylabel(y_label)

    if show_legend:
        ax.legend(**legend_kwargs)

    if title is not None:
        ax.set_title(title)

    return plt.gcf()


from random import choice
import numpy as np
from numpy import load
from numpy import expand_dims
from sklearn.preprocessing import LabelEncoder
from sklearn.preprocessing import Normalizer
from sklearn.svm import SVC
from matplotlib import pyplot
import pickle

# load faces
data = np.array(pickle.load(open('data/datasets/Xtest.pickle', 'rb')))
testX_faces = data
# load face embeddings
data = load('embeddedFaceData.npz')
trainX, trainy, testX, testy = data['arr_0'], data['arr_1'], data['arr_2'], data['arr_3']
# normalize input vectors
in_encoder = Normalizer(norm='l2')
trainX = in_encoder.transform(trainX)
testX = in_encoder.transform(testX)
# label encode targets
out_encoder = LabelEncoder()
out_encoder.fit(trainy)
trainy = out_encoder.transform(trainy)
testy = out_encoder.transform(testy)
model = SVC(kernel='linear', probability=True)
model.fit(trainX, trainy)

# test model on 5 random examples from the test dataset
for i in range(5):
    selection = choice([j for j in range(testX.shape[0])])
    randomImage = testX_faces[selection]
    face_embed = testX[selection]
    face_class = testy[selection]
    face_name = out_encoder.inverse_transform([face_class])
    # prediction for the face
    samples = expand_dims(face_embed, axis=0)
    yhat_class = model.predict(samples)
    yhat_prob = model.predict_proba(samples)
    # get name
    class_index = yhat_class[0]
    class_probability = yhat_prob[0,class_index] * 100
    predict_names = out_encoder.inverse_transform(yhat_class)
    print(f'Predicted: {predict_names[0]} ({class_probability:.3f})')
    print(f'Expected: {face_name[0]}')
    pyplot.imshow(randomImage)
    title = '%s (%.3f)' % (predict_names[0], class_probability)
    pyplot.title(title)
    pyplot.show()


import os

# os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'

import argparse
import cv2
from os.path import join as pjoin
import json
import math
import numpy as np
from torch.utils.data import Dataset, DataLoader
import torch
from PIL import Image
import matplotlib.pyplot as plt
import random
import os
from torchvision import transforms, utils
import scipy.io as io
import resnext101_wsl
from DepthNet4 import *
from ConvTrans import *
from DepthNet4 import Decoder
from DepthNet4 import Decoder_noAO
import matplotlib.pyplot as plt
from torch.nn import DataParallel
import h5py
import glob
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True

def img_loader(path):

    image = cv2.imread(path)
    if image.ndim == 2:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) / 255.0

    return image

mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]

seq_len= 12
small_seq_len = 2

devices = [f'cuda:{c}' for c in range(torch.cuda.device_count())]

device_cnn_encoder = torch.device(devices[0])
device_fmnet_encoder = torch.device(devices[0])
device_fmnet_decoder = torch.device(devices[-1])
device_cnn_decoder = torch.device(devices[-1])

cnn_encoder = resnext101_wsl.resnext101_wsl()  #resnet.resnet50() #resnext101_wsl.resnext101_wsl()
fmnet_Encoder = fmnet_encoder(num_hidden = [2048],batch_size = 1 ,img_width = 20,img_height = 15 ,input_length = seq_len ,small_length = small_seq_len ,encoder_depth = 6, device = device_fmnet_encoder)
fmnet_Decoder = fmnet_decoder(num_hidden = [2048],batch_size = 1 ,img_width = 20,img_height = 15 ,input_length = seq_len ,small_length = small_seq_len ,encoder_depth = 1, device = device_fmnet_decoder)
cnn_decoder = Decoder(inchannels=[256, 512, 1024, 2048], midchannels=[256, 256, 256, 512], upfactors=[2, 2, 2, 2], outchannels=1)

checkpoint = torch.load('./checkpoint/nyu_epoch_20.pth',map_location = 'cpu')    #18
cnn_encoder.load_state_dict(checkpoint['cnn_encoder'])
cnn_decoder.load_state_dict(checkpoint['cnn_decoder'])
fmnet_Encoder.load_state_dict(checkpoint['Mae_Encoder'])
fmnet_Decoder.load_state_dict(checkpoint['Mae_Decoder'])

cnn_encoder.to(device_cnn_encoder)
fmnet_Encoder.to(device_fmnet_encoder)
fmnet_Decoder.to(device_fmnet_decoder)
cnn_decoder.to(device_cnn_decoder)

base_dir = './demo/'
image_dir = base_dir + '/rgb/'
image_seq = glob.glob(image_dir+'*.png')
image_seq.sort()

for j in range(seq_len):
    img = img_loader(image_seq[j])
    img = (img - mean) / std
    img = np.transpose(img, (2, 0, 1))
    img = torch.Tensor(np.ascontiguousarray(img).astype(np.float32))
    img = img.unsqueeze(0)
    if j==0:
        img_seq = img
    else:
        img_seq=torch.cat([img_seq,img],dim=0)

with torch.no_grad():
    img_seq = img_seq.to(device_cnn_encoder)
    img_seq = img_seq.squeeze(0)
    spatial_feat0, spatial_feat1, spatial_feat2,outputs = cnn_encoder(img_seq)
    outputs = outputs.to(device_fmnet_encoder)
    outputs, choose_frames = fmnet_Encoder(outputs)
    outputs = outputs.to(device_fmnet_decoder)
    outputsall = fmnet_Decoder(outputs, choose_frames)
    spatial_feat0 = spatial_feat0.to(device_cnn_decoder)
    spatial_feat1 = spatial_feat1.to(device_cnn_decoder)
    spatial_feat2 = spatial_feat2.to(device_cnn_decoder)
    outputsall = outputsall.to(device_cnn_decoder)
    outputsall = cnn_decoder([spatial_feat0, spatial_feat1, spatial_feat2, outputsall])[0]
    outputsall = F.relu(outputsall).squeeze()

for k in range(outputsall.shape[0]):
    vis = outputsall[k].detach().cpu().numpy().squeeze()
    vis[vis!=0] = 1.0 / vis[vis!=0]
    plt.imsave('./demo/results/'+str(k+1)+'.png',vis, cmap='inferno',vmin =np.min(vis) , vmax = np.max(vis))


import streams
from constants import CKeys, CColors
import kalman
import drawing
import aruco
import depthai as dai
import cv2
import numpy as np
import sys
import time
from pdb import set_trace

# Function to calculate the camera intrinsic matrix for a cropped field of view
def getCroppedCameraIntrinsics(calibData, cameraId, topLeft = dai.Point2f()):

    # Image height and width
    img_width, img_height = cameraId.getResolutionSize()

    # Get the default intrinsic matrix (uncropped)
    camera_matrix = calibData.getCameraIntrinsics(cameraId.getBoardSocket(), resizeWidth=img_width, resizeHeight=img_height)

    # Update the translation components of the matrix according to the crop location
    camera_matrix[0][2] -= topLeft.x * img_width
    camera_matrix[1][2] -= topLeft.y * img_height

    return camera_matrix

# Printing options
np.set_printoptions(precision=2)
np.set_printoptions(floatmode = "fixed")

# Limit value of num to be within range (v0, v1)
def clamp(num, v0, v1):
    return max(v0, min(num, v1))

# Exposure controls
expTime = 2000
expMin = 1
expMax = 33000
expStep = 500  # us

# Sensitivity controls
sensIso = 1600
sensMin = 100
sensMax = 1600
isoStep = 50

# Step size for crop location ('W','A','S','D' controls)
cropLocationStepSize = 0.02

# Output video path
out_video_path = None

# Draw tracking?
draw_tracking_status = False

# Draw corner numbers?
draw_corner_nums = False

# Draw edges of each tag?
draw_aruco_edges = False

# Draw the ID of each tag?
draw_aruco_ids = False

# Draw the cube pose?
draw_aruco_axes = True

# Draw xyz coordinates?
draw_xyz = True

# Which camera's coordinate system to report state in
master_camera = "mono_right"

# Which camera to use ("mono_left," "mono_right," or "color")
camera_type = "mono_right"

 # Target position in the image (homogeneous coordinates, e.g., [x,y,z,1])
xyz_target = np.array([0,0,0.5,1], dtype=np.float32)

# For plotting
tracker_color = (0,255,255) # Line / text color
tracker_thickness = 2 # Line thickness
aruco_edge_color = (0,0,255) # Line / text color
aruco_edge_thickness = 2 # Line thickness
frames_to_plot = [0,1]

# For plotting coordinates
xyz_xo = 0
xyz_dy = 35
xyz_yo = int(-3.5 * xyz_dy)
xyz_color = (0,255,255)
xyz_size = 0.8

# Aruco stuff
arucoDict = cv2.aruco.Dictionary_get(cv2.aruco.DICT_4X4_50)
arucoParams = cv2.aruco.DetectorParameters_create()
arucoParams.cornerRefinementMethod=cv2.aruco.CORNER_REFINE_CONTOUR

# # Aruco tag size in meters
# # Small cube uses 40 mm tags
# # Big cube uses 80 mm tags
aruco_tag_size_meters = 0.040 # 2 inch cube
# aruco_tag_size_meters = 0.080 # 4 inch cube
# aruco_tag_size_meters = 0.187325

# # Length of each side of the aruco cube in meters.
# # 2 inch = 0.0508 m
# # 4 inch = 0.1016 m
aruco_cube_size_meters = 0.0508 # 2 inch cube
# aruco_cube_size_meters = 0.1016 # 4 inch cube
# aruco_cube_size_meters = 0

# Create aruco board object (for the 4" cube)
# Order of tags in board_id should be: [front, right, back, left, top, bottom]
board_ids = np.array([0, 1, 2, 3, 5, 4]) # 2 inch cube. There is actually no tag 5; need it to specify cube
# board_ids = np.array([5, 6, 7, 8, 9, 10]) # 4 inch cube

# Board rotation
board_rvec = np.array([0,0,np.pi])
# board_rvec = np.array([0,0,0])

# Make the aruco cube
board = aruco.create_aruco_cube(board_ids = board_ids,
    aruco_dict = arucoDict,
    cube_width_m = aruco_cube_size_meters,
    tag_width_m = aruco_tag_size_meters,
    rvec = board_rvec)

# DepthAI Pipeline
pipeline = dai.Pipeline()

# Number of cameras
cam_fps = 40
fpsStepSize = 5
video_fps = 40

# The reference coordinate system is cameraBoardSockets[0]
cameraBoardSockets = [dai.CameraBoardSocket.LEFT, dai.CameraBoardSocket.RIGHT]

# Initialize lists of cameras and output links
cameras = []
xOutImages = []
imageManips = []

# Camera control node
controlIn = pipeline.create(dai.node.XLinkIn)
controlIn.setStreamName('control')

# Camera config node
configIn = pipeline.create(dai.node.XLinkIn)
configIn.setStreamName('config')
sendCamConfig = False

# Initial crop range
topLeft = dai.Point2f(0.2, 0.2)
bottomRight = dai.Point2f(0.8, 0.8)

# Set up pipeline
for i in range(len(cameraBoardSockets)):

    # Mono camera nodes
    cam = pipeline.create(dai.node.MonoCamera)
    cam.setResolution(dai.MonoCameraProperties.SensorResolution.THE_800_P)
    cam.setFps(cam_fps)
    cam.setBoardSocket(cameraBoardSockets[i])

    # Image manipulation nodes
    imageManip = pipeline.create(dai.node.ImageManip)

    # Output nodes
    xOutImage = pipeline.create(dai.node.XLinkOut)
    xOutImage.input.setBlocking(False)
    xOutImage.input.setQueueSize(1)
    xOutImage.setStreamName("imageOut" + str(i))

    # Cropping etc
    imageManip.initialConfig.setCropRect(topLeft.x, topLeft.y, bottomRight.x, bottomRight.y)
    imageManip.setMaxOutputFrameSize(cam.getResolutionHeight()*cam.getResolutionWidth()*3)

    # Link camera control / config
    controlIn.out.link(cam.inputControl)
    configIn.out.link(imageManip.inputConfig)

    # Link stuff
    cam.out.link(imageManip.inputImage)
    imageManip.out.link(xOutImage.input)

    # Append camera to list of cams
    cameras.append(cam)
    xOutImages.append(xOutImage)

# Set up the connection for streaming out data
socket = 5555
pub = streams.Publisher(socket=socket)
pub.connect()

# Connect and start pipeline
with dai.Device(pipeline) as device:

    # Control and config queues
    controlQueue = device.getInputQueue(controlIn.getStreamName())
    configQueue = device.getInputQueue(configIn.getStreamName())

    # Get camera calibration info
    calibData = device.readCalibration()

    # Get camera extrinsics relative to first camera in list
    camera_extrinsics = []
    for i, cam in enumerate(cameras):
        if i > 0:
            extrinsicMatrix = np.array(calibData.getCameraExtrinsics(cameras[0].getBoardSocket(), cam.getBoardSocket()))

            # calibData outputs the translation part of the extrinsic matrix in centimeters.
            # In this house we work in meters!
            extrinsicMatrix[0:3, 3] /= 100
        else:
            extrinsicMatrix = np.eye(4)
        camera_extrinsics.append(extrinsicMatrix)

    # 3x3 matrix of [[fx, 0, cx], [0, fy, cy], [0, 0, 1]]
    # Get camera intrinsic matrices and distortion parameters
    camera_matrices = []
    camera_distortions = []
    for cam in cameras:
        img_width, img_height = cam.getResolutionSize()

        # Intrinsic matrix
        camera_matrix = np.array(getCroppedCameraIntrinsics(calibData, cam, topLeft = topLeft))

        # Distortion coeffs
        camera_distortion = np.array(calibData.getDistortionCoefficients(cam.getBoardSocket()))[:-2]

        # Append matrices to lists
        camera_matrices.append(camera_matrix)
        camera_distortions.append(camera_distortion)

    # Open a video writer object
    if out_video_path is not None:

        # Figure out the video resolution
        videoWidthPixels  = int((bottomRight.x - topLeft.x) * cameras[0].getResolutionWidth())
        videoHeightPixels = int((bottomRight.y - topLeft.y) * cameras[0].getResolutionHeight())
        video_resolution = (videoWidthPixels, videoHeightPixels)

        # Open the videowriter
        outVideo = cv2.VideoWriter(out_video_path, cv2.VideoWriter_fourcc('M','J','P','G'), video_fps, video_resolution)
    else:
        outVideo = None
        video_resolution = (0,0)

    # Listen for data on the device xLinkOut queue
    image_out_queues = []
    for xOutImage in xOutImages:
        qImageOut = device.getOutputQueue(name = xOutImage.getStreamName(), maxSize=1, blocking=False)
        image_out_queues.append(qImageOut)

    prev_frame_time = 0
    new_frame_time = 0
    rmat_kalman_prev = np.eye(3)
    rvec_kalman = np.zeros(3)
    debug = False
    plotFiltered = True

    # Set up the kalman filter
    dt = 1 / cam.getFps() # Time step

    # Measurement noise
    kf, dt_idx = kalman.kalmanFilter6DOFConstantVelocityMultiCam(dt = dt)
    H_ref = kf.measurementMatrix
    R_ref_diag = np.diag(kf.measurementNoiseCov)

    # Initialize camera controls
    ctrl = dai.CameraControl()
    ctrl.setManualExposure(expTime, sensIso)
    controlQueue.send(ctrl)

    # Initialize counter
    count = 0

    # Set up plotting windows
    for i, boardSock in enumerate(cameraBoardSockets):
        if i in frames_to_plot:
            # How far to shift image windows
            dx_image_windows = int((bottomRight.x - topLeft.x) * cameras[i].getResolutionWidth())
            fTitle = str(boardSock)
            cv2.namedWindow(fTitle)
            cv2.moveWindow(fTitle, i * dx_image_windows, 0)

    # Main processing loop
    while True:

        # Update time step
        new_frame_time = time.time()
        if prev_frame_time > 0:
            dt = new_frame_time - prev_frame_time
        prev_frame_time = new_frame_time

        # Update kalman state predictions (calculate xk_km1)
        kf.transitionMatrix[dt_idx] = dt
        kf.predict()

        # Initialize vectors for frames (images) and pose measurements
        frames = []
        measurements = []

        # Get all the frames before we do any processing so that they're closely spaced in time
        for i, queue in enumerate(image_out_queues):

            # Get the frame from the queue
            frame = queue.get().getCvFrame()
            if len(frame.shape) == 2:
                frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
            frames.append(frame)

        # Process each frame
        for i, frame in enumerate(frames):
            cameraMatrix = camera_matrices[i]
            distCoeffs    = camera_distortions[i]

            # Measure the position and pose of the aruco board
            valid, measurement, detectedCorners, detectedIds = aruco.getPoseMeasurement(frame, board, cameraMatrix=cameraMatrix, distCoeffs = distCoeffs, arucoDict = arucoDict, parameters=arucoParams)

            # Estimate the board pose if any tags were detected
            if len(detectedCorners) > 0 and i in frames_to_plot:

                # Draw IDs?
                if draw_aruco_ids is True:
                    frame = drawing.draw_ids(frame, detectedCorners, detectedIds)

                if draw_aruco_edges is True:
                    frame = drawing.draw_edges(frame, detectedCorners)

                if draw_corner_nums is True:
                    frame = drawing.draw_corners(frame, detectedCorners)

                if draw_tracking_status is True:
                    frame = drawing.draw_tracking(frame, detectedCorners)

            if plotFiltered is False and i in frames_to_plot:
                cv2.putText(frame, "Kalman filter OFF",
                    (2, 25), cv2.FONT_HERSHEY_TRIPLEX, 1, (255,255,255))

            # Transform measurement to reference camera's coordinate system
            if valid is True:

                # Transform the measurement into the reference coordinate system (cameras[0])
                measurement = aruco.transformMeasurement(measurement, np.linalg.inv(camera_extrinsics[i]))

                # If the two rvecs differ by >= pi then one of them has flipped.
                # The actual angle between the vectors is smaller than the norm of their differences.
                # Flip the new measurement to point in the same direction as the reference measurement
                if len(measurements) > 0:
                    measurement = aruco.alignPose(measurement, measurements[0])

                # Make the measurements array
                measurements.append(measurement)

            # Draw FPS on frame
            if i in frames_to_plot:
                frame = drawing.draw_fps(frame, 1/dt)

            # Upate frames vector
            frames[i] = frame

        # Update the kalman filter's measurement matrix according to
        # the number of valid measurements that were recorded.
        if len(measurements) > 0:

            # Set the measurement matrix and meas. noise matrix for the number of valid measurements
            H = np.concatenate([H_ref for x in measurements], axis=0)
            R = np.diag(np.concatenate([R_ref_diag for x in measurements], axis=0))
            measurement = np.concatenate(measurements)

            # Set kalman filter parameters for this iteration
            kf.measurementMatrix = H
            kf.measurementNoiseCov = R

            # Correct state prediction (calculate xk_k)
            kf.correct(measurement)

        # Read the pose from the kf state vector
        tvec = kf.statePost[0:3]
        rvec = kf.statePost[3:6]

        # Error covariance trace
        errCovTrace = np.trace(kf.errorCovPost)

        # Create the data object
        dataframe = streams.DataFrame(
            tvec = tvec, rvec = rvec,
            kfErrorCov = errCovTrace,
            timestamp = new_frame_time,
            framenumber = count)
        pub.send_frame(data=dataframe)

        # Draw the axes on the rames
        for i, frame in enumerate(frames):

            if i in frames_to_plot:

                # Transform cube coordinates into current camera's frame of ref
                pose_mat = np.eye(4)
                pose_mat[0:3, 0:3], _ = cv2.Rodrigues(rvec)
                pose_mat[0:3, 3] = np.transpose(tvec)

                # Target point in the board coordinates
                xyz_target_board_frame_homogeneous = np.linalg.inv(pose_mat) @ xyz_target

                # XYZ of the target in drone-centered coordinate system, in mm
                xyz_t_d = 1000 * xyz_target_board_frame_homogeneous / xyz_target_board_frame_homogeneous[-1]

                # Coordinates to plot
                xyz = np.squeeze(tvec)

                # Draw xyz
                if draw_xyz:
                    imagePoints, _ = cv2.projectPoints(objectPoints = tvec.astype(np.float32), rvec = np.array([0,0,0], dtype=np.float32), tvec = np.array([0,0,0], dtype=np.float32), cameraMatrix = camera_matrix, distCoeffs = camera_distortion)
                    xy = np.squeeze(imagePoints)
                    yc = int(xy[1])
                    xc = int(xy[0])

                    try:
                        cv2.putText(frame, "X: {:.0f} mm".format(xyz[0]*1000), (int(xc + xyz_xo), int(yc) + xyz_yo ), cv2.FONT_HERSHEY_TRIPLEX,         xyz_size, xyz_color)
                        cv2.putText(frame, "Y: {:.0f} mm".format(xyz[1]*1000), (int(xc + xyz_xo), int(yc) + xyz_yo + xyz_dy), cv2.FONT_HERSHEY_TRIPLEX,     xyz_size, xyz_color)
                        cv2.putText(frame, "Z: {:.0f} mm".format(xyz[2]*1000), (int(xc + xyz_xo), int(yc) + xyz_yo + 2 * xyz_dy), cv2.FONT_HERSHEY_TRIPLEX, xyz_size, xyz_color)
                    except Exception as err:
                        print(str(err))

                # Draw text saying that we're plotting filtered data
                if plotFiltered is True:
                    cv2.putText(frame, "Kalman filter ON", (2, 25), cv2.FONT_HERSHEY_TRIPLEX, 1, (255,255,255))

                if draw_aruco_axes is True and i == 0:
                    if plotFiltered is True:
                        frame = drawing.draw_pose(frame, camera_matrices[i], camera_distortions[i], rvec, tvec, 2*aruco_tag_size_meters)
                    else:
                        if len(measurements) > 0:
                            tvec = np.array(measurements[0][:3])
                            rvec = np.array(measurements[0][3:])
                            frame = drawing.draw_pose(frame, camera_matrices[i], camera_distortions[i], rvec, tvec, 2*aruco_tag_size_meters)

                # Display the image
                frame_title = str(cameraBoardSockets[i])
                cv2.imshow(frame_title, frame)
            else:
                frame_title = str(cameraBoardSockets[i])
                cv2.imshow(frame_title, frame)

        # Write the video frame
        if outVideo is not None:

            # If the size of the video and the size of the frame aren't equal, close the videoWriter and make a new one of the correct size. Any previously-written frames are lost!
            if frames[0].shape[:2] != video_resolution[::-1]:
                outVideo.release()
                videoWidthPixels = frames[0].shape[1]
                videoHeightPixels = frames[0].shape[0]
                video_resolution = (videoWidthPixels, videoHeightPixels)
                outVideo = cv2.VideoWriter(out_video_path, cv2.VideoWriter_fourcc('M','J','P','G'), video_fps, video_resolution)
            outVideo.write(frames[0]) # Write the frame to the output video

        key=cv2.waitKey(1)
        if key == ord('q'):
            break
        elif key == ord(' '):
            debug = not debug
            if debug is True:
                print("Debugging ON")
            else:
                print("Debugging OFF")
        elif key == ord('f'):
            plotFiltered = not plotFiltered
            if plotFiltered:
                print("Kalman filter ON")
            else:
                print("Kalman filter OFF")
        elif key == ord('e'):
            print("Autoexposure enable")
            ctrl = dai.CameraControl()
            ctrl.setAutoExposureEnable()
            controlQueue.send(ctrl)

        # Set exposure and ISO
        elif key in [ord('i'), ord('o'), ord('k'), ord('l')]:
            if key == ord('i'): expTime -= expStep
            if key == ord('o'): expTime += expStep
            if key == ord('k'): sensIso -= isoStep
            if key == ord('l'): sensIso += isoStep
            expTime = clamp(expTime, expMin, expMax)
            sensIso = clamp(sensIso, sensMin, sensMax)
            print("Setting manual exposure, time:", expTime, "iso:", sensIso)
            ctrl = dai.CameraControl()
            ctrl.setManualExposure(expTime, sensIso)
            controlQueue.send(ctrl)

        # Move crop up
        elif key == ord('w'):
            if topLeft.y - cropLocationStepSize >= 0:
                topLeft.y -= cropLocationStepSize
                bottomRight.y -= cropLocationStepSize
                sendCamConfig = True

        # Move crop left
        elif key == ord('a'):
            if topLeft.x - cropLocationStepSize >= 0:
                topLeft.x -= cropLocationStepSize
                bottomRight.x -= cropLocationStepSize
                sendCamConfig = True

        # Move crop down
        elif key == ord('s'):
            if bottomRight.y + cropLocationStepSize <= 1:
                topLeft.y += cropLocationStepSize
                bottomRight.y += cropLocationStepSize
                sendCamConfig = True

        # Move crop right
        elif key == ord('d'):
            if bottomRight.x + cropLocationStepSize <= 1:
                topLeft.x += cropLocationStepSize
                bottomRight.x += cropLocationStepSize
                sendCamConfig = True

        # Make image bigger
        elif key == ord('='):
            topLeft.x -= cropLocationStepSize
            topLeft.y -= cropLocationStepSize
            bottomRight.x += cropLocationStepSize
            bottomRight.y += cropLocationStepSize
            sendCamConfig = True

        # Make image smaller
        elif key == ord('-'):
            topLeft.x += cropLocationStepSize
            topLeft.y += cropLocationStepSize
            bottomRight.x -= cropLocationStepSize
            bottomRight.y -= cropLocationStepSize
            sendCamConfig = True

        # Update camera config (cropping)
        if sendCamConfig is True:

            # Clamp the image corner coordinates so the size can never go to 0
            topLeft.x = clamp(topLeft.x, 0.0, 0.9)
            topLeft.y = clamp(topLeft.y, 0.0, 0.9)
            bottomRight.x = clamp(bottomRight.x, topLeft.x + 0.1, 1.0)
            bottomRight.y = clamp(bottomRight.y, topLeft.y + 0.1, 1.0)

            # Configure the image crop
            cfg = dai.ImageManipConfig()
            cfg.setCropRect(topLeft.x, topLeft.y, bottomRight.x, bottomRight.y)

            # Send the image crop command to the camera
            configQueue.send(cfg)

            # Update the camera intrinsic matrices according to the new crop
            camera_matrices = []
            for cam in cameras:
                topLeftPixel = dai.Point2f(topLeft.x * img_width, topLeft.y * img_height)
                camera_matrix = np.array(getCroppedCameraIntrinsics(calibData, cam, topLeft = topLeft))
                camera_matrices.append(camera_matrix)

            # No more cam config sending
            sendCamConfig = False

        # Update loop counter
        count += 1

# Close the video
if out_video_path is not None:
    outVideo.release()

# GTFO
print("Exiting")