import os, sys, argparse, time, math
import torch
import torch.nn as nn

from dataset.cifar import CIFAR10Instance, CIFAR100Instance

from helper.pretrain import init
from util import TwoCropTransform, TwoCropTransformWithBase
from util import save_model, AverageMeter, BYOL

from models import model_dict
from optimizer import LR_Scheduler

from torchvision import transforms, datasets
from tensorboardX import SummaryWriter

def parse_option():
    parser = argparse.ArgumentParser('argument for training')

    parser.add_argument('--print_freq', type=int, default=10,
                        help='print frequency')
    parser.add_argument('--save_freq', type=int, default=100,
                        help='save frequency')
    parser.add_argument('--batch_size', type=int, default=256,
                        help='batch size')
    parser.add_argument('--num_workers', type=int, default=16,
                        help='num of workers to use')
    parser.add_argument('--epochs', type=int, default=300,
                        help='number of training epochs')
    parser.add_argument('--init_epochs', type=int, default=30,
                        help='init training for two-stage methods')

    # optimizer
    parser.add_argument('--lr', type=float, default=0.03,
                        help='learning rate')
    parser.add_argument('--wd', type=float, default=5e-4,
                        help='weight decay')
    parser.add_argument('--momentum', type=float, default=0.9,
                        help='momentum for SGD')

    # model / dataset
    parser.add_argument('--model', type=str, default='ResNet50')
    parser.add_argument('--dataset', type=str, default='cifar100',
                        choices=['cifar10', 'cifar100', 'path'])

    # custom dataset
    parser.add_argument('--mean', type=str, help='mean of dataset in path in form of str tuple')
    parser.add_argument('--std', type=str, help='std of dataset in path in form of str tuple')
    parser.add_argument('--data_folder', type=str, default=None,
                        help='path to custom dataset')
    parser.add_argument('--img_size', type=int, default=32,
                        help='parameter for RandomResizedCrop')

    # distillation
    parser.add_argument('--distill', type=str, default='OURS',
                        help='distillation method')
    parser.add_argument('--teacher', type=str, default='teacher.pth',
                        help='path to teacher')
    parser.add_argument('--beta', type=float, default=1.,
                        help='weight for distillation loss term')

    # NCE distillation
    parser.add_argument('--feat_dim', default=128, type=int,
                        help='feature dimension')
    parser.add_argument('--mode', default='exact', type=str,
                        choices=['exact', 'relax'], help='mode for CRD')
    parser.add_argument('--nce_k', default=16384, type=int,
                        help='number of negative samples for NCE')
    parser.add_argument('--nce_t', default=0.07, type=float,
                        help='temperature parameter for softmax')
    parser.add_argument('--nce_m', default=0.5, type=float,
                        help='momentum for non-parametric updates')

    # hint layer
    parser.add_argument('--hint_layer', default=2, type=int,
                        choices=[0, 1, 2, 3, 4], help='hint layer for FitNet')

    # OURS
    parser.add_argument('--weighted', action='store_true',
                        help='weighted loss')
    parser.add_argument('--weighting', type=str, default='naive',
                        choices=['naive', 'l1_normalized', 'l2_normalized', 'softmax'], help='weighting method')

    # SimCLR Loss
    parser.add_argument('--temp', type=float, default=0.5,
                        help='temperature for SimCLR loss function')

    # others
    parser.add_argument('--memo', type=str, default='',
                        help='additional memo for ckpt path')

    opt = parser.parse_args()

    # set different learning rate from these 4 models
    if opt.model in ['MobileNetV2', 'ShuffleV1', 'ShuffleV2']:
        opt.lr = 0.01

    # check if dataset is path that passed required arguments
    if opt.dataset == 'path':
        assert opt.data_folder is not None \
            and opt.mean is not None \
            and opt.std is not None

    # set the path according to the environment
    opt.model_path = './save/BYOL/{}_models'.format(opt.dataset)
    opt.tb_path = './save/BYOL/{}_tensorboards'.format(opt.dataset)

    opt.model_name = '{}_{}_lr_{}_decay_{}_bsz_{}_beta_{}{}'.\
        format(opt.dataset, opt.model,
            opt.lr, opt.wd, opt.batch_size, opt.beta, opt.memo)

    opt.tb_folder = os.path.join(opt.tb_path, opt.model_name)
    if not os.path.isdir(opt.tb_folder):
        os.makedirs(opt.tb_folder)

    opt.save_folder = os.path.join(opt.model_path, opt.model_name)
    if not os.path.isdir(opt.save_folder):
        os.makedirs(opt.save_folder)

    if opt.data_folder is None:
        opt.data_folder = './data'

    # device setting
    opt.device = torch.device('cpu')
    if torch.cuda.is_available():
        opt.device = torch.device('cuda')

    return opt

def set_loader(opt):

    # construct data loader
    if opt.dataset == 'cifar10':
        mean = (0.4914, 0.4822, 0.4465)
        std = (0.2023, 0.1994, 0.2010)
    elif opt.dataset == 'cifar100':
        mean = (0.5071, 0.4867, 0.4408)
        std = (0.2675, 0.2565, 0.2761)
    elif opt.dataset == 'path':
        mean = eval(opt.mean)
        std = eval(opt.std)
    else:
        raise ValueError('dataset not supported: {}'.format(opt.dataset))

    normalize = transforms.Normalize(mean=mean, std=std)

    train_transform = transforms.Compose([
        transforms.RandomResizedCrop(size=32, scale=(0.2, 1.)),
        transforms.RandomHorizontalFlip(),
        transforms.RandomApply([
            transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)
        ], p=0.8),
        transforms.RandomGrayscale(p=0.2),
        transforms.ToTensor(),
        normalize])

    if opt.dataset == 'cifar10':
        train_dataset = CIFAR10Instance(root=opt.data_folder,
                        transform=TwoCropTransform(train_transform),
                        download=True)
    elif opt.dataset == 'cifar100':
        train_dataset = CIFAR100Instance(root=opt.data_folder,
                        transform=TwoCropTransform(train_transform),
                        download=True)
    elif opt.dataset == 'path':
        train_dataset = datasets.ImageFolder(root=opt.data_folder,
                        transform=TwoCropTransform(train_transform))
    else:
        raise ValueError(opt.dataset)

    n_data = len(train_dataset)

    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=opt.batch_size,
        num_workers=opt.num_workers, pin_memory=True)

    return train_loader, n_data

def set_model(opt, train_loader, n_data):
    if opt.model == 'ResNet50':
        opt.s_dim = 2048
    elif opt.model == 'ResNet18':
        opt.s_dim = 512

    model = BYOL(opt, model_dict[opt.model])
    model = model.to(opt.device)
    return model

def set_optimizer(opt, model, num_iter):
    optimizer = torch.optim.SGD(model.parameters(),
                 lr=opt.lr, momentum=opt.momentum, weight_decay=opt.wd)

    # cosine decay scheduler + warmup
    lr_scheduler = LR_Scheduler(
        optimizer,
        warmup_epochs=10, warmup_lr=0,
        num_epochs=800, base_lr=opt.lr,
        final_lr=0,
        iter_per_epoch=num_iter, # len(train_loader)
    )

    return optimizer, lr_scheduler

def train(train_loader, model, optimizer, lr_scheduler, epoch, opt):
    """ one epoch training """

    model.train()

    losses_byol = AverageMeter()

    end = time.time()
    for idx, ((x1, x2), _, index) in enumerate(train_loader):
        x1, x2 = x1.to(opt.device), x2.to(opt.device)

        # ===================forward=====================
        loss_byol = model((x1, x2))

        loss = loss_byol
        losses_byol.update(loss_byol.item(), len(index))

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        model.step(epoch / opt.epochs)
        lr_scheduler.step()

        if (idx + 1) % opt.print_freq == 0:
            print('Train: [{0}/{1}][{2}/{3}]'
                  ' BYOL loss {loss_byol.val:.6f} ({loss_byol.avg:.6f})'
                  ' | lr {lr:.6f} | time :{time:.2f}s'.format(
                   epoch, opt.epochs, idx + 1, len(train_loader),
                   time=time.time() - end, loss_byol=losses_byol,
                   lr=lr_scheduler.get_lr()))

            end = time.time()
            sys.stdout.flush()

        #torch.cuda.empty_cache()

    return losses_byol.avg

def main():
    opt = parse_option()
    train_loader, n_data = set_loader(opt)
    model = set_model(opt, train_loader, n_data)
    optimizer, lr_scheduler = set_optimizer(opt, model, len(train_loader))

    writer = SummaryWriter(logdir=opt.tb_folder)

    for epoch in range(1, opt.epochs+1):
        end = time.time()
        loss_byol = train(train_loader, model, optimizer, lr_scheduler, epoch, opt)
        print('epoch {}, total time {:.2f}s'.format(epoch, time.time()-end))

        writer.add_scalar('train loss', loss_byol, epoch)
        writer.add_scalar('learning_rate', lr_scheduler.get_lr(), epoch)

        if epoch % opt.save_freq == 0:
            save_file = os.path.join(
                opt.save_folder, 'ckpt_epoch_{epoch}.pth'.format(epoch=epoch))
            save_model(model.model, optimizer, opt, epoch, save_file)

    # save the last model
    save_file = os.path.join(
        opt.save_folder, 'last.pth')
    save_model(model.model, optimizer, opt, opt.epochs, save_file)

if __name__ == '__main__':
    main()


import graphlib as gl
import os
from concurrent.futures import ThreadPoolExecutor
from concurrent.futures import TimeoutError as ConcurrentTimeoutError
from typing import Dict, List, Optional, Set, Union

from cosmpy.aerial.client import LedgerClient
from cosmpy.aerial.contract import LedgerContract
from cosmpy.aerial.wallet import LocalWallet, Wallet
from cosmpy.crypto.address import Address
from cosmpy.crypto.keypairs import PrivateKey
from jenesis.config import Config, Deployment, Profile
from jenesis.contracts import Contract
from jenesis.contracts.detect import detect_contracts
from jenesis.contracts.monkey import MonkeyContract
from jenesis.keyring import (LocalInfo, query_keychain_item,
                             query_keychain_items)
from jenesis.tasks import Task, TaskStatus
from jenesis.tasks.monitor import run_tasks

# Recursive function to insert deployed contract address into instantiation msg
def insert(data: Union[Dict, List], contract_name: str, address: str):
    for key, value in data.items() if isinstance(data, dict) else enumerate(data):
        if value == ("$" + contract_name):
            data[key] = address
        elif isinstance(value, (dict, list)):
            insert(value, contract_name, address)

def insert_address(contract_address_names: List[str], deployment: Deployment, profile: Profile) -> dict:

    deployment_names = list(profile.deployments.keys())
    init_data = deployment.init

    # iterate over the addresses to insert
    for name in contract_address_names:
        if name in deployment_names:
            deployed_contract_address = profile.deployments[name].address

            assert deployed_contract_address is not None, f"Contract {name} address not found"

            # insert deployed contract address in instantiation msg
            insert(init_data, name, str(deployed_contract_address))

    return init_data

def retreive_init_addresses(data: Union[Dict, List], contract_names: List[str], init_addresses: List[str]):
    for _, value in data.items() if isinstance(data, dict) else enumerate(data):
        if isinstance(value, str) and len(value) > 0:
            if value[0] == "$" and value[1:] in contract_names:
                init_addresses.append(value[1:])
                contract_names.remove(value[1:])
        elif isinstance(value, (dict, list)):
            retreive_init_addresses(value, contract_names,init_addresses)
    return init_addresses

def load_keys(key_names: Set[str], cfg: Config) -> Dict[str, PrivateKey]:
    keys = {}
    available_key_names = set(query_keychain_items(cfg.keyring_backend))
    for key_name in key_names:

        if key_name not in available_key_names:
            print(f"Key not found: {key_name}")
            continue

        info = query_keychain_item(key_name, cfg.keyring_backend)
        if not isinstance(info, LocalInfo):
            print(f"Failed to retrieve local key: {key_name}")
            continue

        keys[key_name] = PrivateKey(info.private_key)
    return keys

class DeployContractTask(Task):
    def __init__(self, project_path: str, cfg: Config, profile: Profile, contract: Contract,
                 deployment: Deployment, client: LedgerClient, wallet: Wallet):
        self._project_path = project_path
        self._cfg = cfg
        self._profile = profile
        self._contract = contract
        self._deployment = deployment
        self._client = client
        self._wallet = wallet

        # task status
        self._status = TaskStatus.IDLE
        self._status_text = ''
        self._logs = ''

        # background worker
        self._executor = ThreadPoolExecutor(max_workers=1)

        # state machine state
        self._state = 'idle'
        self.ledger_contract = None  # type: Optional[LedgerContract]
        self.contract_address = None  # type: Optional[Address]
        self._future = None

    @property
    def name(self) -> str:
        return self._deployment.name

    @property
    def status(self) -> TaskStatus:
        return self._status

    @property
    def status_text(self) -> str:
        return self._status_text

    @property
    def logs_text(self) -> str:
        return self._logs

    def poll(self):
        if self._state == 'idle':
            self._schedule_build_contract()
        elif self._state == 'wait-for-ledger-contract':
            self._wait_for_ledger_contract()
        elif self._state == 'schedule-deployment':
            self._schedule_deploy_contract()
        elif self._state == 'wait-for-deployment':
            self._wait_for_contract_deployment()
        elif self._state == 'failed':
            self._failed()
        elif self._state == 'complete':
            self._complete()
        else:
            assert False, "bad state"

    def _schedule_build_contract(self):
        assert self._future is None

        def action():
            return MonkeyContract(
                self._contract,
                self._client,
                code_id=self._deployment.code_id,
                address=self._deployment.address
            )

        self._future = self._executor.submit(action)
        self._state = 'wait-for-ledger-contract'
        self._status = TaskStatus.IN_PROGRESS
        self._status_text = '(1/2) Determining contract parameters...'

    def _wait_for_ledger_contract(self):
        try:
            self.ledger_contract = self._future.result(timeout=0.05)
            self._state = 'schedule-deployment'
            self._future = None
        except ConcurrentTimeoutError:
            pass

    def _schedule_deploy_contract(self):
        assert self._future is None
        assert self.ledger_contract is not None

        def action():
            return self.ledger_contract._deploy( #  pylint: disable=W0212
                args=self._deployment.init,
                sender=self._wallet,
                admin_address=self._wallet.address(),
                funds=self._deployment.init_funds,
            )

        self._future = self._executor.submit(action)
        self._state = 'wait-for-deployment'
        self._status = TaskStatus.IN_PROGRESS
        self._status_text = '(2/2) Deploying contract...'

    def _wait_for_contract_deployment(self):
        try:
            self.contract_address = self._future.result(timeout=0.05)
            self._state = 'complete'
            self._future = None
        except ConcurrentTimeoutError:
            pass

    def _complete(self):
        # update the configuration and save it to disk
        self._cfg.update_deployment(self._profile.name, self._deployment.name,
                                    self.ledger_contract.digest.hex(),
                                    self.ledger_contract.code_id, self.contract_address)
        self._cfg.save(self._project_path)

        self._finished(True)

    def _failed(self):
        self._finished(False)

    def _finished(self, success: bool):
        self._executor.shutdown(wait=True)
        if success:
            self._status = TaskStatus.COMPLETE
        else:
            self._status = TaskStatus.FAILED
        self._status_text = ''

    def teardown(self):
        pass

def deploy_contracts(cfg: Config, project_path: str, deployer_key: Optional[str], profile_name: Optional[str] = None):
    if profile_name is None:
        profile_name = cfg.get_default_profile()

    profile = cfg.profiles[profile_name]

    project_contracts = {contract.name: contract for contract in detect_contracts(project_path)}
    deployments = profile.deployments

    init_addresses = {}

    # iterate over the addresses to insert
    for (name, deployment) in deployments.items():
        init_data = deployment.init
        deployment_names = list(deployments.keys())

        # retreive init addresses for each contract
        contract_init_addresses = retreive_init_addresses(init_data, deployment_names, [])

        # add all contract init addresses in one dictionary
        init_addresses[name] = set(contract_init_addresses)

    # sort the contracts to be deployed in the right order
    sorter = gl.TopologicalSorter(init_addresses)
    deployment_order = list(sorter.static_order())

    # load all the keys required for this operation
    key_names = {deployment.deployer_key for deployment in deployments.values()} | {deployer_key}
    keys = load_keys(key_names, cfg)

    for deployment_name in deployment_order:
        deployment = deployments[deployment_name]

        # ensure specified contract is in project
        if deployment.contract not in project_contracts:
            print(f"Contract {deployment_name} not found in project")
            continue
        contract = project_contracts[deployment.contract]

        # ensure that contract has been compiled first
        if not os.path.isfile(contract.binary_path):
            print(f"No contract binary found for {contract.name}. Please run 'jenesis compile' first.")
            continue

        if deployer_key is not None:

            if deployer_key not in keys:
                print(f"Skipping {deployment_name}: deployer key {deployer_key} not available")
                continue

            deployment.deployer_key = deployer_key
            Config.update_key(os.getcwd(), profile_name, deployment_name, deployer_key)

        if deployment.address is not None:
            if not deployment.is_configuration_out_of_date():
                print(f"Skipping {deployment_name}: configuration is up to date")
                continue

            if contract.digest() == deployment.digest:
                print(f"Skipping {deployment_name}: digest has not changed")
                continue

        client = LedgerClient(profile.network)

        deployment.address = None  # clear the old address

        contract_address_names = list(init_addresses[deployment_name])

        deployment.init = insert_address(contract_address_names, deployment, profile)

        # lookup the wallet key
        wallet = LocalWallet(keys[deployment.deployer_key])

        # create the deployment task
        task = DeployContractTask(
            project_path,
            cfg,
            profile,
            contract,
            deployment,
            client,
            wallet,
        )

        # run the deployment task
        run_tasks([task])


#!/usr/bin/env python

import sys
import json
import subprocess as sp

exclude = []
wifi_icons = ['\uf6aa', '\uf6ab', '\uf1eb']

def read_fwf(text_table):
    """
    only works when the first row is columns with a single word
    """
    cols_row = text_table.splitlines()[0]
    columns = cols_row.split()
    cols_idxs = []
    for i, (col, colnxt) in enumerate(zip(columns, columns[1:])):
        cols_idxs.append((0 if i == 0 else cols_row.index(f' {col}')+1,
                     cols_row.index(f' {colnxt}')+1))
    # cols_widths = [(cols_row.index(c), cols_row.index(cnxt)) for c, cnxt in zip(columns, columns[1:])]
    cols_idxs.append((cols_row.index(f' {columns[-1]}')+1, -1))
    final_dict = {
        col: [l[col_idx[0] : col_idx[1]].strip() for l in
              text_table.splitlines()[1:]]
        for col, col_idx in zip(columns, cols_idxs)
    }
    return final_dict

def get_devices_status():
    devices_table = sp.getoutput('nmcli device')
    return read_fwf(devices_table)

def get_wifi_networks(ifname):
    cmd = f'nmcli dev wifi list --rescan no ifname {ifname}'
    wifi_aps = sp.getoutput(cmd)
    return read_fwf(wifi_aps)

def output_wrapper():
    try:
        devs = get_devices_status()
    except IndexError:
        print('(label :class "dimmed" :text "loading...")')
        return
    if list(devs.keys())[0] == 'Error:':
        print('(label :class "dimmed" :text "loading...")')
        return
    output = '(eventbox :onclick "rofi -show nmmenu -modi nmmenu:~/Scripts/RofiMenus/nmcli-menu.py &" (box :space-evenly false :spacing 5 {}))'
    common_device = '(label :tooltip {dev_info} :class "{color}-icon" :text "{icon}")'
    eww_devs = []
    for dev_name, dev_type, dev_state, dev_con in zip(
            devs['DEVICE'], devs['TYPE'], devs['STATE'], devs['CONNECTION']):
        if dev_name == 'lo' or dev_name in exclude:
            continue
        if dev_type == 'wifi' and dev_state == 'connected':
            wifi_nets = get_wifi_networks(dev_name)
            wifi_signal = int(wifi_nets['SIGNAL'][wifi_nets['SSID'].index(dev_con)])

            # wifi_icon = r'󰤟' if wifi_signal < 25 else r'󰤢' if wifi_signal < 50 else '󰤥' if wifi_signal < 75 else '󰤨'
            eww_devs.append(common_device.format(int_idx=devs['DEVICE'].index(dev_name), color='green', icon='\uf1eb', dev_info=f'"{dev_name}: {wifi_signal}% {dev_con}"'))
            # output += bars
        elif dev_type == 'wifi' and dev_state == 'disconnected':
            eww_devs.append(common_device.format(int_idx=devs['DEVICE'].index(dev_name), color='orange', icon='\uf1eb', dev_info=f'"{dev_name}: {dev_con}"'))
        elif dev_type == 'wifi' and dev_state in ['disabled', 'unavailable']:
            eww_devs.append(common_device.format(int_idx=devs['DEVICE'].index(dev_name), color='red', icon='\uf1eb', dev_info=f'"{dev_name}: {dev_con}"'))
        elif dev_type == 'wifi' and dev_state == 'unmanaged':
            eww_devs.append(common_device.format(int_idx=devs['DEVICE'].index(dev_name), color='grey', icon='\uf1eb', dev_info=f'"{dev_name}: {dev_con}"'))
        elif dev_type == 'ethernet' and dev_state == 'connected':
            eww_devs.append(common_device.format(int_idx=devs['DEVICE'].index(dev_name), color='green', icon='\uf796', dev_info=f'"{dev_name}: {dev_con}"'))
        elif dev_type == 'ethernet' and dev_state == 'connecting':
            eww_devs.append(common_device.format(int_idx=devs['DEVICE'].index(dev_name), color='yellow', icon='\uf796', dev_info=f'"{dev_name}: {dev_con}"'))
        elif dev_type == 'ethernet' and dev_state in ['unavailable', 'disconnected']:
            eww_devs.append(common_device.format(int_idx=devs['DEVICE'].index(dev_name), color='orange', icon='\uf796', dev_info=f'"{dev_name}: {dev_con}"'))
        elif dev_type == 'ethernet' and dev_state == 'unmanaged':
            eww_devs.append(common_device.format(int_idx=devs['DEVICE'].index(dev_name), color='red', icon='\uf796', dev_info=f'"{dev_name}: {dev_con}"'))
        elif dev_type == 'tun' and dev_state in ['connected (externally)', 'connected']:
            eww_devs.append(common_device.format(int_idx=devs['DEVICE'].index(dev_name), color='green', icon='\uf6ff', dev_info=f'"{dev_name}: {dev_con}"'))
        elif dev_type == 'tun' and dev_state == 'disconnected':
            eww_devs.append(common_device.format(int_idx=devs['DEVICE'].index(dev_name), color='orange', icon='\uf6ff', dev_info=f'"{dev_name}: {dev_con}"'))

    print(output.format(' '.join(eww_devs)))
    sys.stdout.flush()

nm_monitor = sp.Popen('nmcli monitor', stdout=sp.PIPE, text=True, shell=True)

while True:
    nm_monitor.stdout.readline()
    output_wrapper()


import cv2
import pickle #нужен для того, чтобы сохранить позиции парковочных мест

try: # если файл такой есть, то забираем его данные в posList, если же такого файла нет, то создаем новый posList.
    # это нужно для того, чтобы при каждом новом запуске сразу же загружались прошлая версия выделений.
    with open('CarParkPos', 'rb') as f:
        posList = pickle.load(f)
except:
    posList = []
width, height = 107, 48

def mouseClick(events, x, y, flags, params):
    if events == cv2.EVENT_LBUTTONDOWN:
        posList.append((x, y))
    if events == cv2.EVENT_RBUTTONDOWN:
        for i, pos in enumerate(posList):
            x1, y1 = pos
            if x1<x<x1+width and y1<y<y1+height:
                posList.pop(i)

    # чтобы сохранять то, что нарисовали wb - read and write    dump - сваливать
    with open('CarParkPos', 'wb') as f:
        pickle.dump(posList, f) # сваливаем posList в файл f

while True:
    img = cv2.imread('carParkImg.png')  # считываем изображени
    for pos in posList:
        cv2.rectangle(img, pos, (pos[0]+width, pos[1]+height), (255, 0, 255), 2)
    cv2.imshow("Image", img)
    cv2.setMouseCallback("Image", mouseClick) # первым аргументом передаем название изображения, по которому клацаем,
    # а вторым параметром передаем функцию, которая будет вызываться, когда мы кликаем по чему-либо
    cv2.waitKey(1)


import logging
from pathlib import Path
from time import sleep
from typing import List, Optional

from tomtom_api import config, log
from tomtom_api.client import TomtomClient
from tomtom_api.priority_queue.models.database import PriorityQueueDB
from tomtom_api.priority_queue.models.queue import QueueItem
from tomtom_api.priority_queue.models.status import QueueItemStatus
from tomtom_api.traffic_stats import N_CONCURRENT_JOB_IN_PROGRESS
from tomtom_api.traffic_stats.models.status import (TomtomJobState,
                                                    TomtomReportType)
from tomtom_api.utils.daemon import Daemon

DAEMON_LOG_FILE = config.path.home / 'daemon.log'

class PriorityQueueDaemon(Daemon):
    def __init__(self,
                 pid_file: Optional[Path] = None,
                 stdin: Path = Path('/dev/null'),
                 stdout: Path = Path('/dev/null'),
                 stderr: Path = Path('/dev/null')
                 ):
        pid_file = config.path.home / 'daemon-tomtom-api.pid' if pid_file is None else Path(pid_file)
        pid_file.parent.mkdir(parents=True, exist_ok=True)
        super().__init__(pid_file, stdin, stdout, stderr)

    def run(self):
        # add a file log handler since we won't see anything in STDOUT
        # because all standard input are redirected to /dev/null
        file_handler = logging.FileHandler(DAEMON_LOG_FILE, encoding='utf-8')
        file_handler.setFormatter(log.handlers[0].formatter)
        file_handler.setLevel(config.log.level.upper())
        log.addHandler(file_handler)

        # init client
        client = TomtomClient()

        # init db
        db = PriorityQueueDB()

        # states for "In progress"
        in_progress_states = [TomtomJobState.NEW,
                              TomtomJobState.MAPMATCHED,
                              TomtomJobState.MAPMATCHING,
                              TomtomJobState.READING_GEOBASE,
                              TomtomJobState.CALCULATIONS,
                              TomtomJobState.SCHEDULED]

        while True:
            log.debug('Daemonic loop starting.')
            sleep(config.queue.loop_duration)
            db.read()
            try:
                log.debug('Checking the number of tomtom jobs that are IN PROGRESS.')
                in_progress_jobs = client.search_jobs(state=in_progress_states)
                n = in_progress_jobs.total_elements

                if n < N_CONCURRENT_JOB_IN_PROGRESS:
                    log.info(f'There are {N_CONCURRENT_JOB_IN_PROGRESS - n} available spots on tomtom side.')

                    # 1: find old submitted job(s) and mark them as completed
                    tomtom_job_in_progress_ids = [c.job_id for c in in_progress_jobs.content]
                    submitted_items: List[QueueItem] = db.get_filtered_items(status=[QueueItemStatus.SUBMITTED])

                    # the items that should be marked as complete are the one that are tagged by ourselves as
                    # "submitted" but that are not (anymore) in the tomtom "in progress" jobs.
                    should_be_complete_items: List[QueueItem] = [
                        j for j in submitted_items
                        if j.tomtom_job_id not in tomtom_job_in_progress_ids
                    ]

                    # then mark all the jobs that should be complete as "complete".
                    for item in should_be_complete_items:
                        item.complete(client)
                        log.info(f'Completed item {item.uid} (tomtom job {item.tomtom_job_id})')
                    db.update(should_be_complete_items, True)

                    # 2: find next job(s) and submit them
                    for item in db.get_next(N_CONCURRENT_JOB_IN_PROGRESS - n):
                        item.submit(client)
                        log.info(f'Submitted item {item.uid} as tomtom job {item.tomtom_job_id}.')
                        db.update([item], True)

            # We don't want the daemon to exit because of some unmanaged exception
            except Exception as e:
                log.critical(str(e))


from reportlab.lib.pagesizes import letter,QRCode
from reportlab.lib import colors
from reportlab.platypus import Frame, PageTemplate,Image
from reportlab.lib.units import cm,mm
from reportlab.platypus import (Table, TableStyle, BaseDocTemplate)
from reportlab.platypus.flowables import TopPadder

def create_pdf():

    I = Image('qr1.png')
    I.drawHeight = 22*mm*I.drawHeight / I.drawWidth
    I.drawWidth = 22*mm

    # Create a frame
    CatBox_frame = Frame(
        x1=0 * mm,  # From left
        y1=0 * mm,  # From bottom
        height=75 * mm,
        width=110 * mm,
        leftPadding=0 * mm,
        bottomPadding=0 * mm,
        rightPadding=0 * mm,
        topPadding=0 * mm,
        showBoundary=1,
        id='CatBox_frame')

    table_data=[[[I], [I], [I]],
                [[I], [I], [I]],
                [[I], [I], [I]]]
    # Create a table
    CatBox = Table(table_data, 36.6* mm, 25 * mm, vAlign='BOTTOM')

    # Style the table
    CatBox.setStyle(TableStyle(
            [('ALIGN',(0,0),(-1,-1),'CENTER'),
             ('VALIGN',(0,0),(-1,-1),'MIDDLE')
            ]))

    # Trying to tell the table to be a bottom align object (when later put in frame)
    CatBox.Align = 'BOTTOM'
    CatBox.vAlign = 'BOTTOM'

    # Building the story
    #story = [CatBox] # adding CatBox table (alternative, story.add(CatBox))
    story = [TopPadder(CatBox)]

    # Establish a document
    doc = BaseDocTemplate("BottomAlignTable.pdf", pagesize=QRCode)

    # Creating a page template
    frontpage = PageTemplate(id='FrontPage',
                             frames=[CatBox_frame]
                             )
    # Adding the story to the template and template to the document
    doc.addPageTemplates(frontpage)

    # Building doc
    doc.build(story)

# ----------------------------------------------------------------------
if __name__ == "__main__":
    create_pdf() # Printing the pdf


import re
import nltk
from nltk.classify import *

class ClassifierHelper:
    #start __init__
    def __init__(self, featureListFile):
        self.wordFeatures = []
        # Read feature list
        inpfile = open(featureListFile, 'r')
        line = inpfile.readline()
        while line:
            self.wordFeatures.append(line.strip())
            line = inpfile.readline()
    #end

    #start extract_features
    def extract_features(self, document):
        document_words = set(document)
        features = {}
        for word in self.wordFeatures:
            word = self.replaceTwoOrMore(word)
            word = word.strip('\'"?,.')
            features['contains(%s)' % word] = (word in document_words)
        return features
    #end

    #start replaceTwoOrMore
    def replaceTwoOrMore(self, s):
        # pattern to look for three or more repetitions of any character, including
        # newlines.
        pattern = re.compile(r"(.)\1{1,}", re.DOTALL)
        return pattern.sub(r"\1\1", s)
    #end

    #start process_tweet
    def process_tweet(self, tweetu):
        #Conver to lower case
        tweet=str(tweetu)
        tweet = tweet.lower()
        #Convert https?://* to URL
        tweet = re.sub('((www\.[^\s]+)|(https?://[^\s]+))','URL',tweet)
        #Convert @username to AT_USER
        tweet = re.sub('@[^\s]+','AT_USER',tweet)
        #Remove additional white spaces
        tweet = re.sub('[\s]+', ' ', tweet)
        #Replace #word with word
        tweet = re.sub(r'#([^\s]+)', r'\1', tweet)
        #trim
        tweet = tweet.strip()
        #remove first/last " or 'at string end
        tweet = tweet.rstrip('\'"')
        tweet = tweet.lstrip('\'"')
        return tweet
    #end

    #start is_ascii
    def is_ascii(self, word):
        return all(ord(c) < 128 for c in word)
    #end
#end class


from fastapi import APIRouter
from config.db import engine
from models.tables import usuarios
from schemas.Users import UserRegister
from config.cryptography import f

newUser = APIRouter()

@newUser.post('/register')
def registrarUsuario(user:UserRegister):
    conexion = engine.connect()
    nuevoUsuario = {"nombre": user.nombre, "correo_electronico": user.email.lower()}
    nuevoUsuario["contrasena"] = f.encrypt(user.contrasena.encode("utf-8"))

    user = conexion.execute(usuarios.select().where(usuarios.c.correo_electronico == nuevoUsuario["correo_electronico"])).first()

    if user:
        message = "El email ingresado ya existe en el sistema."
    else:
        resultado = conexion.execute(usuarios.insert().values(nuevoUsuario))
        if resultado:
            message = 204
    conexion.close()
    engine.dispose()
    return message


import glob
import gzip
import json
import uuid
import argparse

def parseOptions():
    parser = argparse.ArgumentParser(
        description="parse the data dictionaries of unharmonized dbGaP studies"
    )
    parser.add_argument("-d", help="directory of dbGaP study", required=True)
    parser.add_argument("-p", help="phs number of dbGaP study", required=True)
    parser.add_argument("-c", help="consent group number", required=True)
    parser.add_argument("-o", help="output directory for data files", required=True)

    args = parser.parse_args()

    return args

def writeProgramNode(output):
    # writing node for program with name of dbGaP
    programData = []
    program = {}
    program["type"] = "program"
    program["submitter_id"] = "dbGaP"
    programData.append(program)
    fileName = output + "/program.json"
    with open(fileName, "w+") as o:
        o.write(json.dumps(programData))

def writeProjectNode(phs, output):
    # writing node for study with one phs number
    projectData = []
    project = {}
    project["type"] = "project"
    project["submitter_id"] = phs.split(".")[0]
    project["programs"] = {}
    project["programs"]["submitter_id"] = "dbGaP"
    projectData.append(project)
    fileName = output + "/project.json"
    with open(fileName, "w+") as o:
        o.write(json.dumps(projectData))

def writeStudyNode(phs, output):
    # writing node for study with one phs number
    studyData = []
    study = {}
    study["type"] = "Study"
    study["submitter_id"] = phs
    study["projects"] = {}
    study["projects"]["submitter_id"] = phs.split(".")[0]
    studyData.append(study)
    fileName = output + "/Study.json"
    with open(fileName, "w+") as o:
        o.write(json.dumps(studyData))

def parseSubjects(fileName, nodeName, consent, output):
    with gzip.open(fileName, "r") as data:
        header = True
        nodeData = []
        consent_var = ""
        for line in data:
            try:
                line = line.decode("utf-8")
            except Exception as e:
                print("hey utf-8 didn't work so we are going to try iso")
                line = line.decode("iso-8859-1")
            if line[0] == "#" or len(line) < 2:
                continue
            elif header:
                heads = line.replace("\n", "").split("\t")
                for h in heads:
                    if "consent" in h.lower() or h.lower() == "gencons":
                        consent_var = h
                        break
                heads.append("submitter_id")
                header = False
            else:
                entity = {}
                entity["type"] = "subject"
                data = line.replace("\n", "").split("\t")
                size = len(data)
                for i in range(size):
                    entity[heads[i]] = data[i]
                entity["submitter_id"] = str(uuid.uuid4())
                entity["studies"] = {}
                entity["studies"]["submitter_id"] = phs
                if entity[consent_var] == str(consent):
                    nodeData.append(entity)

        # write the file for subject
        writeFile = output + "/" + nodeName + ".json"
        with open(writeFile, "w+") as o:
            o.write(json.dumps(nodeData))

        return nodeData

def parseSample(fileName, nodeName, subjects, output):
    with gzip.open(fileName, "r") as data:
        header = True
        nodeData = []
        subject_var = ""
        for line in data:
            try:
                line = line.decode("utf-8")
            except Exception as e:
                print("hey utf-8 didn't work so we are going to try iso")
                line = line.decode("iso-8859-1")
            if line[0] == "#" or len(line) < 2:
                continue
            elif header:
                heads = line.replace("\n", "").split("\t")
                heads.append("submitter_id")
                header = False
            else:
                entity = {}
                entity["type"] = "sample"
                lineData = line.replace("\n", "").split("\t")
                size = len(lineData)
                for i in range(size):
                    entity[heads[i]] = lineData[i]
                entity["submitter_id"] = str(uuid.uuid4())
                entity["studies"] = {}
                entity["studies"]["submitter_id"] = phs
                if entity["dbGaP_Subject_ID"] in subjects:
                    nodeData.append(entity)

        print("Number of sample records:", len(nodeData))

        # write the file for sample
        writeFile = output + "/" + nodeName + ".json"
        with open(writeFile, "w+") as o:
            o.write(json.dumps(nodeData))

        return

def parsePedigree(fileName, nodeName, subjects, output):
    with gzip.open(fileName, "r") as data:
        header = True
        nodeData = []
        for line in data:
            try:
                line = line.decode("utf-8")
            except Exception as e:
                print("hey utf-8 didn't work so we are going to try iso")
                line = line.decode("iso-8859-1")
            if line[0] == "#" or len(line) < 2:
                continue
            elif header:
                heads = line.replace("\n", "").split("\t")
                heads.append("submitter_id")
                header = False
            else:
                entity = {}
                entity["type"] = "pedigree"
                lineData = line.replace("\n", "").split("\t")
                size = len(lineData)
                for i in range(size):
                    entity[heads[i]] = lineData[i]
                entity["submitter_id"] = str(uuid.uuid4())
                entity["studies"] = {}
                entity["studies"]["submitter_id"] = phs
                if entity["dbGaP_Subject_ID"] in subjects:
                    nodeData.append(entity)

        print("Number of pedigree records:", len(nodeData))

        # write the file for pedigree
        writeFile = output + "/" + nodeName + ".json"
        with open(writeFile, "w+") as o:
            o.write(json.dumps(nodeData))

        return

def parseDataFile(fileName, nodeName, output):
    with gzip.open(fileName, "r") as data:
        header = True
        nodeData = []
        subject_var = ""
        for line in data:
            try:
                line = line.decode("utf-8")
            except Exception as e:
                print("hey utf-8 didn't work so we are going to try iso")
                line = line.decode("iso-8859-1")
            if line[0] == "#" or len(line) < 2:
                continue
            elif header:
                heads = line.replace("\n", "").split("\t")
                heads.append("submitter_id")
                header = False
            else:
                entity = {}
                entity["type"] = nodeName
                lineData = line.replace("\n", "").split("\t")
                size = len(lineData)
                for i in range(size):
                    entity[heads[i]] = lineData[i]
                entity["submitter_id"] = str(uuid.uuid4())
                entity["studies"] = {}
                entity["studies"]["submitter_id"] = phs

                nodeData.append(entity)

        # write the file for all the nodes
        writeFile = output + "/" + nodeName + ".json"
        with open(writeFile, "w+") as o:
            for chunk in json.JSONEncoder(indent=2).iterencode(nodeData):
                o.write(chunk)
            # o.write(json.dumps(nodeData))
        return

# parse argumens and return if required ones are not given
args = parseOptions()
location = args.d
phs = args.p
consent = args.c
output = args.o

# use glob to find all files from location that are .txt.gz files
if location.endswith("/"):
    fileFind = location + "*.txt.gz"
else:
    fileFind = location + "/*.txt.gz"
files = glob.glob(fileFind)

print("Outputting to", output, "directory")

# Write initial study node. Only contains one field of submitter id that is the phs of the project
print("Writing data file for: Program")
writeProgramNode(output)
print("Writing data file for: Project")
writeProjectNode(phs, output)
print("Writing data file for: Study")
writeStudyNode(phs, output)

nodes = {}
dataFiles = []

#
for f in files:
    dataFile = f.replace(location, "")
    thing = dataFile.split(".")
    if "c" + str(consent) in thing[5]:
        file = {}
        file["fileName"] = f
        file["nodeName"] = thing[6]
        dataFiles.append(file)
    elif "Subject" in thing[5] and thing[6] == "MULTI":
        subject = {}
        subject["file"] = f
        subject["name"] = thing[5]
        nodes["subject"] = subject
    elif "Sample" in thing[5] and thing[6] == "MULTI":
        sample = {}
        sample["file"] = f
        sample["name"] = thing[5]
        nodes["sample"] = sample
    elif "Pedigree" in thing[5] and thing[6] == "MULTI":
        pedigree = {}
        pedigree["file"] = f
        pedigree["name"] = thing[5]
        nodes["pedigree"] = pedigree

# parse and write subject file
print("Writing data file for: Subject")
subjects = parseSubjects(
    nodes["subject"]["file"], nodes["subject"]["name"], consent, output
)

# grab the subject ids that are in the consent that we are looking for
subject_ids = []
for s in subjects:
    subject_ids.append(s["dbGaP_Subject_ID"])

# print(subject_ids)

print("Number of subjects with matching consent: ", str(len(subject_ids)))

# parse and write sample file
parseSample(nodes["sample"]["file"], nodes["sample"]["name"], subject_ids, output)

# parse and write pedigree file if it exists
if "pedigree" in nodes.keys():
    parsePedigree(
        nodes["pedigree"]["file"], nodes["pedigree"]["name"], subject_ids, output
    )

print("Writing the rest of the data files")
for d in dataFiles:
    parseDataFile(d["fileName"], d["nodeName"], output)


import numpy as np

from mindquantum.core.circuit import Circuit, SwapParts
from mindquantum.core.gates import H, PhaseShift
from mindquantum.utils.type_value_check import _check_input_type

def _rn(k):
    return PhaseShift(2 * np.pi / 2**k)

def _qft_unit(qubits):
    circ = Circuit(H.on(qubits[0]))
    for index, ctrl_qubit in enumerate(qubits[1:]):
        circ += _rn(index + 2).on(qubits[0], ctrl_qubit)
    return circ

def qft(qubits):
    """
    Quantum fourier transform (QFT).

    The function of the quantum Fourier transform is similar to that of the classical Fourier transform.

    Note:
        Please refer to Nielsen, M., & Chuang, I. (2010) for more information.

    Args:
        qubits (list[int]): Qubits you want to apply quantum fourier transform.

    Examples:
        >>> from mindquantum.algorithm.library import qft
        >>> print(qft([0, 1]).get_qs(ket=True))
        1/2¦00⟩
        1/2¦01⟩
        1/2¦10⟩
        1/2¦11⟩

    Returns:
        Circuit, circuit that can do fourier transform.
    """
    _check_input_type('qubits', (list, range), qubits)
    circuit = Circuit()
    n_qubits = len(qubits)
    for i in range(n_qubits):
        circuit += _qft_unit(qubits[i:])
    if n_qubits > 1:
        part1 = []
        part2 = []
        for j in range(n_qubits // 2):
            part1.append(qubits[j])
            part2.append(qubits[n_qubits - 1 - j])
        circuit += SwapParts(part1, part2)
    return circuit


#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""Create the data for the LSTM.
"""

import os
import sys
import argparse
import numpy as np
import h5py
import itertools
from collections import defaultdict
import pickle
import sys
import re
import codecs
from itertools import izip

class Indexer:
    def __init__(self, symbols = ["*blank*","<unk>","<s>","</s>"]):
        # This needs to be changed so that it loads the proper mappings from word to characters
        self.vocab = defaultdict(int)
        self.PAD = symbols[0]
        self.UNK = symbols[1]
        self.BOS = symbols[2]
        self.EOS = symbols[3]
        self.d = {self.PAD: 1, self.UNK: 2, self.BOS: 3, self.EOS: 4}

    def add_w(self, ws):
        for w in ws:
            if w not in self.d:
                self.d[w] = len(self.d) + 1

    def convert(self, w):
        return self.d[w] if w in self.d else self.d[self.UNK]

    def convert_sequence(self, ls):
        return [self.convert(l) for l in ls]

    def clean(self, s):
        s = s.replace(self.PAD, "")
        # s = s.replace(self.UNK, "")
        s = s.replace(self.BOS, "")
        s = s.replace(self.EOS, "")
        return s

    def write(self, outfile):
        out = open(outfile, "w")
        items = [(v, k) for k, v in self.d.iteritems()]
        items.sort()
        for v, k in items:
            print >>out, k.encode('utf-8'), v
        out.close()

    # This sorts the vocab according to frequency count and reduces the
    # vocab to a certain specified amount
    def prune_vocab(self, k):
        vocab_list = [(word, count) for word, count in self.vocab.iteritems()]
        vocab_list.sort(key = lambda x: x[1], reverse=True)
        k = min(k, len(vocab_list))
        self.pruned_vocab = {pair[0]:pair[1] for pair in vocab_list[:k]}
        for word in self.pruned_vocab:
            if word not in self.d:
                self.d[word] = len(self.d) + 1

    def load_vocab(self, vocab_file):
        self.d = {}
        for line in open(vocab_file, 'r'):
            v, k = line.decode("utf-8").strip().split()
            self.d[v] = int(k)

def pad(ls, length, symbol):
    if len(ls) >= length:
        return ls[:length]
    return ls + [symbol] * (length -len(ls))

def get_data(args):
    src_indexer = Indexer(["<blank>","<unk>","<s>","</s>"])
    target_indexer = Indexer(["<blank>","<unk>","<s>","</s>"])

    def make_vocab(srcfile, targetfile, seqlength, max_word_l=0, chars=0):
        num_sents = 0
        sent_thrown_out = 0
        for _, (src_orig, targ_orig) in \
                enumerate(itertools.izip(open(srcfile,'r'), open(targetfile,'r'))):
            src_orig = src_indexer.clean(src_orig.decode("utf-8").strip())
            targ_orig = target_indexer.clean(targ_orig.decode("utf-8").strip())
            targ = targ_orig.strip().split()
            src = src_orig.strip().split()
            if len(targ) > seqlength or len(src) > seqlength or len(targ) < 1 or len(src) < 1:
                sent_thrown_out = sent_thrown_out + 1
                continue
            num_sents += 1
            for word in targ:
                target_indexer.vocab[word] += 1

            for word in src:
                src_indexer.vocab[word] += 1
        print('Number of sentences thrown out', sent_thrown_out)
        return max_word_l, num_sents

    def convert(srcfile, targetfile, batchsize, seqlength, outfile, num_sents,
                max_word_l, max_sent_l=0,chars=0, unkfilter=0):

        newseqlength = seqlength + 2 #add 2 for EOS and BOS
        targets = np.zeros((num_sents, newseqlength), dtype=int)
        target_output = np.zeros((num_sents, newseqlength), dtype=int)
        sources = np.zeros((num_sents, newseqlength), dtype=int)
        source_lengths = np.zeros((num_sents,), dtype=int)
        target_lengths = np.zeros((num_sents,), dtype=int)
        dropped = 0
        sent_id = 0
        for _, (src_orig, targ_orig) in \
                enumerate(itertools.izip(open(srcfile,'r'), open(targetfile,'r'))):
            # Loading the sentences
            src_orig = src_indexer.clean(src_orig.decode("utf-8").strip())
            targ_orig = target_indexer.clean(targ_orig.decode("utf-8").strip())
            targ = [target_indexer.BOS] + targ_orig.strip().split() + [target_indexer.EOS]
            src =  [src_indexer.BOS] + src_orig.strip().split() + [src_indexer.EOS]
            max_sent_l = max(len(targ), len(src), max_sent_l)

            # We're bounding the length of a sequence for a file
            if len(targ) > newseqlength or len(src) > newseqlength or len(targ) < 3 or len(src) < 3:
                dropped += 1
                continue

            targ = pad(targ, newseqlength+1, target_indexer.PAD)
            for word in targ:
                word = word if word in target_indexer.d else target_indexer.UNK
            targ = target_indexer.convert_sequence(targ)
            targ = np.array(targ, dtype=int)

            src = pad(src, newseqlength, src_indexer.PAD)
            src = src_indexer.convert_sequence(src)
            src = np.array(src, dtype=int)

            # Drops all unknown characters
            if unkfilter > 0:
                targ_unks = float((targ[:-1] == 2).sum())
                src_unks = float((src == 2).sum())
                if unkfilter < 1: #unkfilter is a percentage if < 1
                    targ_unks = targ_unks/(len(targ[:-1])-2)
                    src_unks = src_unks/(len(src)-2)
                if targ_unks > unkfilter or src_unks > unkfilter:
                    dropped += 1
                    continue

            targets[sent_id] = np.array(targ[:-1],dtype=int)
            target_lengths[sent_id] = (targets[sent_id] != 1).sum()
            target_output[sent_id] = np.array(targ[1:],dtype=int)
            sources[sent_id] = np.array(src, dtype=int)
            source_lengths[sent_id] = (sources[sent_id] != 1).sum()

            sent_id += 1
            if sent_id % 100000 == 0:
                print("{}/{} sentences processed".format(sent_id, num_sents))
                if sent_id > 3000000:
                    break

        #break up batches based on source lengths
        source_lengths = source_lengths[:sent_id]
        source_sort = np.argsort(source_lengths)

        sources = sources[source_sort]
        targets = targets[source_sort]
        target_output = target_output[source_sort]
        target_l = target_lengths[source_sort]
        source_l = source_lengths[source_sort]

        curr_l = 0
        l_location = [] #idx where sent length changes

        for j,i in enumerate(source_sort):
            if source_lengths[i] > curr_l:
                curr_l = source_lengths[i]
                l_location.append(j+1)
        l_location.append(len(sources))

        #get batch sizes
        curr_idx = 1
        batch_idx = [1]
        nonzeros = []
        batch_l = []
        batch_w = []
        target_l_max = []
        for i in range(len(l_location)-1):
            while curr_idx < l_location[i+1]:
                curr_idx = min(curr_idx + batchsize, l_location[i+1])
                batch_idx.append(curr_idx)
        for i in range(len(batch_idx)-1):
            batch_l.append(batch_idx[i+1] - batch_idx[i])
            batch_w.append(source_l[batch_idx[i]-1])
            nonzeros.append((target_output[batch_idx[i]-1:batch_idx[i+1]-1] != 1).sum().sum())
            target_l_max.append(max(target_l[batch_idx[i]-1:batch_idx[i+1]-1]))

        # Write output
        f = h5py.File(outfile, "w")

        f["source"] = sources
        f["target"] = targets
        f["target_output"] = target_output
        f["target_l"] = np.array(target_l_max, dtype=int)
        f["target_l_all"] = target_l
        f["batch_l"] = np.array(batch_l, dtype=int)
        f["batch_w"] = np.array(batch_w, dtype=int)
        f["batch_idx"] = np.array(batch_idx[:-1], dtype=int)
        f["target_nonzeros"] = np.array(nonzeros, dtype=int)
        f["source_size"] = np.array([len(src_indexer.d)])
        f["target_size"] = np.array([len(target_indexer.d)])

        print("Saved {} sentences (dropped {} due to length/unk filter)".format(
            len(f["source"]), dropped))
        f.close()
        return max_sent_l

    print("First pass through data to get vocab...")
    max_word_l, num_sents_train = make_vocab(args.srcfile, args.targetfile,
                                             args.seqlength, 0, 0)
    print("Number of sentences in training: {}".format(num_sents_train))
    max_word_l, num_sents_valid = make_vocab(args.srcvalfile, args.targetvalfile,
                                             args.seqlength, max_word_l, 0)
    print("Number of sentences in valid: {}".format(num_sents_valid))

    if args.subtle:
        max_word_l, num_sents_subtle = make_vocab(args.srcsubtlefile, args.targetsubtlefile,
                                                 args.seqlength, max_word_l, 0)
        print("Number of sentences in subtle: {}".format(num_sents_subtle))

    #prune and write vocab
    src_indexer.prune_vocab(args.srcvocabsize)
    target_indexer.prune_vocab(args.targetvocabsize)
    if args.srcvocabfile != '':
        # You can try and load vocabulary here for both the source and target files
        print('Loading pre-specified source vocab from ' + args.srcvocabfile)
        src_indexer.load_vocab(args.srcvocabfile)
    if args.targetvocabfile != '':
        print('Loading pre-specified target vocab from ' + args.targetvocabfile)
        target_indexer.load_vocab(args.targetvocabfile)

    src_indexer.write(args.outputfile + ".src.dict")
    target_indexer.write(args.outputfile + ".targ.dict")

    print("Source vocab size: Original = {}, Pruned = {}".format(len(src_indexer.vocab),
                                                          len(src_indexer.d)))
    print("Target vocab size: Original = {}, Pruned = {}".format(len(target_indexer.vocab),
                                                          len(target_indexer.d)))

    max_sent_l = 0
    max_sent_l = convert(args.srcvalfile, args.targetvalfile, args.batchsize, args.seqlength,
                         args.outputfile + "-val.hdf5", num_sents_valid,
                         max_word_l, max_sent_l, 0, args.unkfilter)
    max_sent_l = convert(args.srcfile, args.targetfile, args.batchsize, args.seqlength,
                         args.outputfile + "-train.hdf5", num_sents_train, max_word_l,
                         max_sent_l, 0, args.unkfilter)
    print("Max sent length: {}".format(max_sent_l))

    if args.subtle:
        max_sent_l = convert(args.srcsubtlefile, args.targetsubtlefile, args.batchsize, args.subtle_seqlength,
                             args.outputfile + "-subtle.hdf5", num_sents_subtle, max_word_l,
                             max_sent_l, 0, args.unkfilter)

    print("Max sent length for subtle: {}".format(max_sent_l))

def format_data(args):
    '''
        Formats the data so it can be passed into get_data function
    '''

    def clean_word_mappings(data_dict):
        '''
            Function makes necessary changes to the word_mappings
            so that they match up with the MovieTriples design.
            The changes to the mappings are then applied to the
            training validation and test set of MovieTriples.

            Function could be changed to remove some of the hard-coding
        '''

        # Words that need to be altered
        bad_words = ['<unk>', '<s>', '</s>', '.', "'"]

        #the new mappings that need to be added
        new_mappings = [('<blank>', 1, 0, 0), ('<unk>', 2, 190588, 89059), ('<s>', 3, 588827, 785135), ('</s>', 4, 588827, 785135), ('.', 10003, 855616, 192250), ("'", 10004, 457542, 160249), ('<t>', 10005, 0, 0)]

        # Bad incides
        del_ind = []
        for i in range(len(data_dict['MovieTriples']['word_mappings'])):
            word_mapping = data_dict['MovieTriples']['word_mappings'][i]
            if word_mapping[0] in bad_words:
                    del_ind.append(i)

        # Delete the bad indicies in reverse order
        del_ind.sort(reverse=True)
        for ind in del_ind:
            del data_dict['MovieTriples']['word_mappings'][ind]

        # Add the new mappings
        for n_m in new_mappings:
            data_dict['MovieTriples']['word_mappings'].append(n_m)

        # Sanity check
        check_mappings = range(1, len(data_dict['MovieTriples']['word_mappings']) + 1)
        for word_mapping in data_dict['MovieTriples']['word_mappings']:
            check_mappings.remove(word_mapping[1])
        assert check_mappings == []

        # Make changes to the dataset
        data_sets = [data_dict['MovieTriples']['train_set'], data_dict['MovieTriples']['valid_set'],
                     data_dict['MovieTriples']['test_set']]
        if args.subtle:
            data_sets.append(data_dict['MovieTriples']['subtle_set'])

        for i in range(len(data_sets)):
            for j in range(len(data_sets[i])):
                line = data_sets[i][j]
                for k in range(len(line)):
                    ind = line[k]
                    if ind == 3:
                        line[k] = 10003
                    elif ind == 4:
                        line[k] = 10004
                    elif ind == 0:
                        line[k] = 2
                    elif ind == 1:
                        line[k] = 3
                    elif ind == 2:
                        line[k] = 4
                data_sets[i][j] = line

        return data_dict

    def load_data():
        '''
            Assuming all inputs files are .pkl files

            Loads all data files into a data_dict thats first layer
            is the name of the data folder that is processed, and the
            second layer is the variable name of the loaded file.

            i.e.

            data_dict['MovieTriples']['train_set'] = pickle.load(...)
        '''
        data_dict = {}
        for data_set in args.input_files:
            data_dict[data_set] = {}
            for var_name, data_file in args.input_files[data_set].iteritems():
                with open('%s%s/%s' % (args.input_directory, data_set, data_file)) as f:
                    data_dict[data_set][var_name] = pickle.load(f)

        return data_dict

    def write_indicies_to_file(filename, y):
        '''
            Write the contents of y into filename
        '''
        f =  open(filename, 'w')
        for context in y:
            for ind in context:
                f.write(str(ind) + ' ')
            f.write('\n')
        f.close()

    def write_vocab_to_file(filename):
        '''
            Write the vocabulary to filename
        '''
        with open(filename, 'w') as f:
            for i in range(1, len(indices_to_word)+1):
                f.write(indices_to_word[i] + ' ' + str(i) + '\n')

    def write_words_to_file(filename, indices_dict):
        '''
            Write the examples to files as words, removing special
            indices
        '''
        lst = []
        for context in indices_dict:
            context_words = []
            for ind in context:
                if ind not in special_indices:
                    context_words.append(indices_to_word[ind])
            lst.append(' '.join(context_words))

        f =  open(filename, 'w')
        for context in lst:
            f.write(str(context) + ' \n')
        f.close()

    # Load in datafiles
    data_dict = load_data()

    # Fix the wordmappings for Movie Triples
    data_dict = clean_word_mappings(data_dict)

    # indices -> word
    indices_to_word = {}
    for word_ex in data_dict['MovieTriples']['word_mappings']:
        indices_to_word[word_ex[1]] = word_ex[0]

    # word -> indices
    word_to_indices = {}
    for word_ex in data_dict['MovieTriples']['word_mappings']:
        word_to_indices[word_ex[0]] = word_ex[1]

    # Apply above basic parsing to all contexts and outputs

    PADDING = word_to_indices['<blank>']
    END_OF_CONV = word_to_indices['<t>']

    full_context = []
    full_output = []
    max_len_context = 0
    max_len_output = 0

    pattern = [word_to_indices['</s>'], word_to_indices['<s>']]

    data_set_contexts = []
    data_set_outputs = []
    data_sets = [data_dict['MovieTriples']['train_set'], data_dict['MovieTriples']['valid_set'],
                    data_dict['MovieTriples']['test_set']]

    if args.subtle:
        data_sets.append(data_dict['MovieTriples']['subtle_set'])

    for j in range(len(data_sets)):
        data_set = data_sets[j]
        PADDING = word_to_indices['<blank>']
        END_OF_CONV = word_to_indices['<t>']

        full_context = []
        full_output = []

        # Make the subtle dataset a bit shorter
        if j != 3:
            max_len_context = 52
            max_len_output = 52
        else:
            # import bpdb; bpdb.set_trace()
            max_len_context = 22
            max_len_output = 22

        for i in range(len(data_set)):
            break_pt = []
            for ind in range(len(data_set[i]))[::-1]:
                if pattern == data_set[i][ind:ind+2]:
                    break_pt.append(ind)

            if break_pt == []:
                continue

            context = data_set[i][:break_pt[0]]
            output = data_set[i][break_pt[0]+2:]

            context = context + [word_to_indices['</s>']]
            output = [word_to_indices['<s>']] + output

            # There is only one utterance in the subtle set
            if j != 3:
                # Start of sentence and end of sentence is ONLY used at the end
                # We create a new character that represents the start and end of a conversation
                context = context[:break_pt[1]] + [END_OF_CONV] + context[break_pt[1]+2:]

            # Cap the target and src length at 302 words to make computation simpler, goes up to ~1500
            if len(context) > max_len_context:
                continue
            if len(output) > max_len_output:
                continue

            full_context.append(context)
            full_output.append(output)

        # Add padding to all contexts and outputs
        for i in range(len(full_context)):
            full_context[i] = full_context[i] + [PADDING] * (max_len_context - len(full_context[i]))
            full_output[i] = full_output[i] + [PADDING] * (max_len_output - len(full_output[i]))

        data_set_contexts.append(full_context)
        data_set_outputs.append(full_output)

    train_full_context = data_set_contexts[0]
    train_full_output = data_set_outputs[0]
    valid_full_context = data_set_contexts[1]
    valid_full_output = data_set_outputs[1]

    if args.subtle:
        subtle_full_context = data_set_contexts[3]
        subtle_full_output = data_set_outputs[3]

    if not os.path.exists(args.output_directory):
        os.makedirs(args.output_directory)

    write_indicies_to_file(args.srcfile_ind, train_full_context)
    write_indicies_to_file(args.targetfile_ind, train_full_output)
    write_indicies_to_file(args.srcvalfile_ind, valid_full_context)
    write_indicies_to_file(args.targetvalfile_ind, valid_full_output)

    if args.subtle:
        write_indicies_to_file(args.srcsubtlefile_ind, subtle_full_context)
        write_indicies_to_file(args.targetsubtlefile_ind, subtle_full_output)

    write_vocab_to_file(args.targetvocabfile)
    write_vocab_to_file(args.srcvocabfile)

    special_indices = [1, 2, 3, 4,]

    write_words_to_file(args.srcfile, train_full_context)
    write_words_to_file(args.srcvalfile, valid_full_context)
    write_words_to_file(args.targetfile, train_full_output)
    write_words_to_file(args.targetvalfile, valid_full_output)

    if args.subtle:
        write_words_to_file(args.srcsubtlefile, subtle_full_context)
        write_words_to_file(args.targetsubtlefile, subtle_full_output)

    # Additional embeddings
    np.random.seed(9844)
    additional_vectors = np.random.uniform(-0.1, 0.1, (2, 300))
    additional_vectors = additional_vectors / np.linalg.norm(additional_vectors)

    emb_wordvec_upd = np.vstack((data_dict['MovieTriples']['emb_wordvec'][0], additional_vectors))
    emb_wordvec_upd[10001] = data_dict['MovieTriples']['emb_wordvec'][0][3][:]
    emb_wordvec_upd[10002] = data_dict['MovieTriples']['emb_wordvec'][0][4][:]
    emb_wordvec_upd[3] = data_dict['MovieTriples']['emb_wordvec'][0][10001][:]
    emb_wordvec_upd[4] = data_dict['MovieTriples']['emb_wordvec'][0][10002][:]
    emb_wordvec_upd = np.roll(emb_wordvec_upd, 1, axis=0)
    f = h5py.File(args.output_directory + 'word_vecs.hdf5', 'w')
    f['word_vecs'] = emb_wordvec_upd
    f.close()

    print('Done formatting the data from Movie Triples dataset')

def main(arguments):
    '''
        Main function parses the command arguments and sets up the
        paths to each of the input/output files. Output files can be set
        via command line while input files should be hard coded into the
        preprocessing file
    '''
    parser = argparse.ArgumentParser(
                description=__doc__,
                formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    # Directories of input raw files and output processed files
    parser.add_argument('--input_directory', help="Path to the folder that contains all of "
                                                   "the raw data files to be used for preprocessing",
                                                   default='../data/')
    parser.add_argument('--output_directory', help="Path to the folder that will hold all of "
                                                    "the datafiles generated by this program",
                                                default='data/')
    parser.add_argument('--outputfile', help="Prefix of the output file names. ", type=str, default='conv')

    # Vocabularys
    parser.add_argument('--srcvocabsize', help="Size of source vocabulary, constructed "
                                                "by taking the top X most frequent words. "
                                                " Rest are replaced with special UNK tokens.",
                                                type=int, default=10000)
    parser.add_argument('--targetvocabsize', help="Size of target vocabulary, constructed "
                                                "by taking the top X most frequent words. "
                                                "Rest are replaced with special UNK tokens.",
                                                type=int, default=10000)
    parser.add_argument('--srcvocabfile', help="If working with a preset vocab, "
                                          "then including this will ignore srcvocabsize and use the"
                                          "vocab provided here.",
                                          type = str, default='src.dict')
    parser.add_argument('--targetvocabfile', help="If working with a preset vocab, "
                                         "then including this will ignore targetvocabsize and "
                                         "use the vocab provided here.",
                                          type = str, default='targ.dict')

    # Filenames of processed files filled with words
    parser.add_argument('--srcfile', help="Filename of source training data, "
                                           "where each line represents a single "
                                           "source/target sequence. Will be placed in "
                                           "the args.output_directory folder",
                                           default='train_src_words.txt')
    parser.add_argument('--targetfile', help="Path to target training data, "
                                           "where each line represents a single "
                                           "source/target sequence. Will be placed in "
                                           "the args.output_directory folder",
                                           default='train_targ_words.txt')
    parser.add_argument('--srcvalfile', help="Filename of source validation data.", default='dev_src_words.txt')
    parser.add_argument('--targetvalfile', help="Filename of target validation data.", default='dev_targ_words.txt')
    parser.add_argument('--srcsubtlefile', help="Filename of source subtle data.", default='subtle_src_words.txt')
    parser.add_argument('--targetsubtlefile', help="Filename of target subtle data.", default='subtle_targ_words.txt')

    # Filenames of processed files filled with indices
    parser.add_argument('--srcfile_ind', help="Filename of source training data, "
                                           "where each line represents a single "
                                           "source/target sequence. Will be placed in "
                                           "the args.output_directory folder",
                                           default='train_src_indices.txt')
    parser.add_argument('--targetfile_ind', help="Path to target training data, "
                                           "where each line represents a single "
                                           "source/target sequence. Will be placed in "
                                           "the args.output_directory folder",
                                           default='train_targ_indices.txt')
    parser.add_argument('--srcvalfile_ind', help="Filename of source validation data.", default='dev_src_indices.txt')
    parser.add_argument('--targetvalfile_ind', help="Filename of target validation data.", default='dev_targ_indices.txt')
    parser.add_argument('--srcsubtlefile_ind', help="Filename of source subtle data.", default='subtle_src_indices.txt')
    parser.add_argument('--targetsubtlefile_ind', help="Filename of target subtle data.", default='subtle_targ_indices.txt')

    # Preprocess Modifications
    parser.add_argument('--batchsize', help="Size of each minibatch.", type=int, default=64)
    parser.add_argument('--seqlength', help="Maximum sequence length. Sequences longer "
                                               "than this are dropped.", type=int, default=50)
    parser.add_argument('--subtle_seqlength', help="Maximum sequence length for subtle data. Sequences longer "
                                               "than this are dropped.", type=int, default=15)
    parser.add_argument('--maxwordlength', help="For the character models, words are "
                                           "(if longer than maxwordlength) or zero-padded "
                                            "(if shorter) to maxwordlength", type=int, default=35)
    parser.add_argument('--unkfilter', help="Ignore sentences with too many UNK tokens. "
                                       "Can be an absolute count limit (if > 1) "
                                       "or a proportional limit (0 < unkfilter < 1).",
                                          type = float, default = 0)
    parser.add_argument('--subtle', help="Preprocess the subtle dataset, takes substantially longer"
                                         "because subtle dataset is ~5,000,000 lines", default=False)

    args = parser.parse_args(arguments)

    # Dictionary holding the locations of the input files
    args.input_files = { 'MovieTriples' : { 'train_set' : 'Training.triples.pkl',
                                           'valid_set' : 'Validation.triples.pkl',
                                           'test_set' : 'Test.triples.pkl',
                                           'emb_wordvec' : 'Word2Vec_WordEmb.pkl',
                                           'emb_mt' : 'MT_WordEmb.pkl',
                                           'word_mappings' : 'Training.dict.pkl'
                                          }
                        }

    # Append on output directory to the output files
    output_files = ['srcfile', 'targetfile', 'srcvalfile', 'targetvalfile', 'srcfile_ind', 'targetfile_ind',
                    'srcvalfile_ind', 'targetvalfile_ind', 'outputfile', 'srcvocabfile', 'targetvocabfile']

    if args.subtle:
        args.input_files['MovieTriples']['subtle_set'] = 'Subtle_Dataset.triples.pkl'
        output_files = output_files + ['srcsubtlefile_ind', 'targetsubtlefile_ind', 'srcsubtlefile', 'targetsubtlefile']

    for o_f in output_files:
        setattr(args, o_f, args.output_directory + getattr(args, o_f))

    # print args.output_directory
    # Call preprocessing code
    format_data(args)
    get_data(args)

if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))


# from https://github.com/mcordts/cityscapesScripts/blob/master/cityscapesscripts/helpers/labels.py

from __future__ import print_function, absolute_import, division
from collections import namedtuple
import numpy as np
import tensorflow as tf

#--------------------------------------------------------------------------------
# Definitions
#--------------------------------------------------------------------------------

# a label and all meta information
Label = namedtuple( 'Label' , [

    'name'        , # The identifier of this label, e.g. 'car', 'person', ... .
                    # We use them to uniquely name a class

    'id'          , # An integer ID that is associated with this label.
                    # The IDs are used to represent the label in ground truth images
                    # An ID of -1 means that this label does not have an ID and thus
                    # is ignored when creating ground truth images (e.g. license plate).
                    # Do not modify these IDs, since exactly these IDs are expected by the
                    # evaluation server.

    'trainId'     , # Feel free to modify these IDs as suitable for your method. Then create
                    # ground truth images with train IDs, using the tools provided in the
                    # 'preparation' folder. However, make sure to validate or submit results
                    # to our evaluation server using the regular IDs above!
                    # For trainIds, multiple labels might have the same ID. Then, these labels
                    # are mapped to the same class in the ground truth images. For the inverse
                    # mapping, we use the label that is defined first in the list below.
                    # For example, mapping all void-type classes to the same ID in training,
                    # might make sense for some approaches.
                    # Max value is 255!

    'category'    , # The name of the category that this label belongs to

    'categoryId'  , # The ID of this category. Used to create ground truth images
                    # on category level.

    'hasInstances', # Whether this label distinguishes between single instances or not

    'ignoreInEval', # Whether pixels having this class as ground truth label are ignored
                    # during evaluations or not

    'color'       , # The color of this label
    ] )

#--------------------------------------------------------------------------------
# A list of all labels
#--------------------------------------------------------------------------------

# Please adapt the train IDs as appropriate for your approach.
# Note that you might want to ignore labels with ID 255 during training.
# Further note that the current train IDs are only a suggestion. You can use whatever you like.
# Make sure to provide your results using the original IDs and not the training IDs.
# Note that many IDs are ignored in evaluation and thus you never need to predict these!

labels = [
    #       name                     id    trainId   category            catId     hasInstances   ignoreInEval   color
    Label(  'unlabeled'            ,  0 ,      255 , 'void'            , 0       , False        , True         , (  0,  0,  0) ),
    Label(  'ego vehicle'          ,  1 ,      255 , 'void'            , 0       , False        , True         , (  0,  0,  0) ),
    Label(  'rectification border' ,  2 ,      255 , 'void'            , 0       , False        , True         , (  0,  0,  0) ),
    Label(  'out of roi'           ,  3 ,      255 , 'void'            , 0       , False        , True         , (  0,  0,  0) ),
    Label(  'static'               ,  4 ,      255 , 'void'            , 0       , False        , True         , (  0,  0,  0) ),
    Label(  'dynamic'              ,  5 ,      255 , 'void'            , 0       , False        , True         , (111, 74,  0) ),
    Label(  'ground'               ,  6 ,      255 , 'void'            , 0       , False        , True         , ( 81,  0, 81) ),
    Label(  'road'                 ,  7 ,        0 , 'flat'            , 1       , False        , False        , (160,160,160) ), # 128, 64,128
    Label(  'sidewalk'             ,  8 ,        1 , 'flat'            , 1       , False        , False        , (100,100,100) ), # 244, 35,232
    Label(  'parking'              ,  9 ,      255 , 'flat'            , 1       , False        , True         , (250,170,160) ),
    Label(  'rail track'           , 10 ,      255 , 'flat'            , 1       , False        , True         , (230,150,140) ),
    Label(  'building'             , 11 ,        2 , 'construction'    , 2       , False        , False        , ( 70, 70, 70) ),
    Label(  'wall'                 , 12 ,        3 , 'construction'    , 2       , False        , False        , (102,102,156) ),
    Label(  'fence'                , 13 ,        4 , 'construction'    , 2       , False        , False        , (190,153,153) ),
    Label(  'guard rail'           , 14 ,      255 , 'construction'    , 2       , False        , True         , (180,165,180) ),
    Label(  'bridge'               , 15 ,      255 , 'construction'    , 2       , False        , True         , (150,100,100) ),
    Label(  'tunnel'               , 16 ,      255 , 'construction'    , 2       , False        , True         , (150,120, 90) ),
    Label(  'pole'                 , 17 ,        5 , 'object'          , 3       , False        , False        , (153,153,153) ),
    Label(  'polegroup'            , 18 ,      255 , 'object'          , 3       , False        , True         , (153,153,153) ),
    Label(  'traffic light'        , 19 ,        6 , 'object'          , 3       , False        , False        , (250,170, 30) ),
    Label(  'traffic sign'         , 20 ,        7 , 'object'          , 3       , False        , False        , (220,220,  0) ),
    Label(  'vegetation'           , 21 ,        8 , 'nature'          , 4       , False        , False        , (107,142, 35) ),
    Label(  'terrain'              , 22 ,        9 , 'nature'          , 4       , False        , False        , (152,251,152) ),
    Label(  'sky'                  , 23 ,       10 , 'sky'             , 5       , False        , False        , ( 70,130,180) ),
    Label(  'person'               , 24 ,       11 , 'human'           , 6       , True         , False        , (255,255,  0) ), # 220, 20, 60
    Label(  'rider'                , 25 ,       12 , 'human'           , 6       , True         , False        , (255,  0,  0) ),
    Label(  'car'                  , 26 ,       13 , 'vehicle'         , 7       , True         , False        , (  0,  0,142) ),
    Label(  'truck'                , 27 ,       14 , 'vehicle'         , 7       , True         , False        , (  0,  0, 70) ),
    Label(  'bus'                  , 28 ,       15 , 'vehicle'         , 7       , True         , False        , (  0, 60,100) ),
    Label(  'caravan'              , 29 ,      255 , 'vehicle'         , 7       , True         , True         , (  0,  0, 90) ),
    Label(  'trailer'              , 30 ,      255 , 'vehicle'         , 7       , True         , True         , (  0,  0,110) ),
    Label(  'train'                , 31 ,       16 , 'vehicle'         , 7       , True         , False        , (  0, 80,100) ),
    Label(  'motorcycle'           , 32 ,       17 , 'vehicle'         , 7       , True         , False        , (  0,  0,230) ),
    Label(  'bicycle'              , 33 ,       18 , 'vehicle'         , 7       , True         , False        , (119, 11, 32) ),
    Label(  'license plate'        , -1 ,       -1 , 'vehicle'         , 7       , False        , True         , (  0,  0,142) ),
]

eval_labels = [label for label in labels if not label.ignoreInEval]

keep_labels = [not label.ignoreInEval for label in labels]

label_colors = {label.id: label.color for label in labels}

model2eval_id = {}
for i, label in enumerate(eval_labels):
    model2eval_id[i] = label.id

def remap_mask_eval_id(mask):
    eval_mask = np.copy(mask)
    ids = np.unique(eval_mask)
    for mask_id in ids:
        eval_mask[mask == mask_id] = model2eval_id[mask_id]

    return eval_mask

def id_to_rgb(seg_id, id2rgb_dict):
    # to be used with np.apply_along_axis with the segmentation ID array
    seg_id = seg_id if isinstance(seg_id, int) else seg_id[0]
    rgb = id2rgb_dict[seg_id]
    return np.array([rgb[0], rgb[1], rgb[2]])

def id2rgb_seg_img(id_seg_img, id2rgb_dict):
    return np.array(np.apply_along_axis(id_to_rgb, 2, id_seg_img[:, :, np.newaxis], id2rgb_dict)).astype(np.uint8)

def out_array_to_eval(out_array):
    mask = np.argmax(out_array, axis=-1)
    mask_remap = remap_mask_eval_id(mask)
    return mask_remap

def out_array_to_label_img(out_array):
    mask_remap = out_array_to_eval(out_array)
    mask_img = id2rgb_seg_img(mask_remap, label_colors)

    return mask_img.astype(int)


import pandas as pd
import os
import matplotlib.pyplot as plt
plt.style.use('seaborn-deep')
plt.rcParams['font.size'] = 30
plt.rcParams['axes.labelsize'] = 30
plt.rcParams['legend.fontsize'] = 30
# plt.rcParams['legend.labelspacing'] = 0.02
# plt.rcParams['figure.figsize'] = (11, 9)

RESULT_ROOT = "../../results"

EXP_NAME = "sim_eval_reproduce"
# EXP_NAME1 = "sim_eval_bbr_reproduce"

fig, axes = plt.subplots(2, 2, figsize=(16, 8.5))
for dim, xlabel, ax in zip(['bandwidth', 'T_s', 'queue', 'loss'],
                           ['Bandwidth (Mbps)', 'Bandwidth change interval (s)',
                            'Queue (Packets)', 'Random packet loss rate (%)'], axes.flatten()):
    df = pd.read_csv(os.path.join(RESULT_ROOT, EXP_NAME, "sim_eval_vary_{}_bbr_with_cubic.csv".format(dim)))
    assert isinstance(df, pd.DataFrame)
    if dim == 'delay':
        vals_to_test = df[dim] * 2
    elif dim == 'loss':
        vals_to_test = df[dim] * 100
    else:
        vals_to_test = df[dim]
    genet_avg_rewards = df['genet_avg_rewards']
    genet_low_bnd = df['genet_low_bnd']
    genet_up_bnd = df['genet_up_bnd']
    bbr_avg_rewards = df['bbr_avg_rewards']
    bbr_low_bnd = df['bbr_low_bnd']
    bbr_up_bnd = df['bbr_up_bnd']
    cubic_avg_rewards = df['cubic_avg_rewards']
    cubic_low_bnd = df['cubic_low_bnd']
    cubic_up_bnd = df['cubic_up_bnd']
    udr_small_avg_rewards = df['udr_small_avg_rewards']
    udr_small_low_bnd = df['udr_small_low_bnd']
    udr_small_up_bnd = df['udr_small_up_bnd']
    udr_mid_avg_rewards = df['udr_mid_avg_rewards']
    udr_mid_low_bnd = df['udr_mid_low_bnd']
    udr_mid_up_bnd = df['udr_mid_up_bnd']
    udr_large_avg_rewards = df['udr_large_avg_rewards']
    udr_large_low_bnd = df['udr_large_low_bnd']
    udr_large_up_bnd = df['udr_large_up_bnd']

    ax.plot(vals_to_test, genet_avg_rewards, color='C2',
            linewidth=4, alpha=1, linestyle='-', label="GENET")
    # genet_low_bnd = np.array(genet_avg_rewards) - \
    #     np.array(genet_reward_errs)
    # genet_up_bnd = np.array(genet_avg_rewards) + \
    #     np.array(genet_reward_errs)
    ax.fill_between(vals_to_test, genet_low_bnd,
                    genet_up_bnd, color='C2', alpha=0.1)

    ax.plot(vals_to_test, bbr_avg_rewards, color='C0',
            linestyle='-', linewidth=4, alpha=1, label="BBR")
    # bbr_low_bnd = np.array(bbr_avg_rewards) - \
    #     np.array(bbr_reward_errs)
    # bbr_up_bnd = np.array(bbr_avg_rewards) + \
    #     np.array(bbr_reward_errs)
    ax.fill_between(vals_to_test, bbr_low_bnd, bbr_up_bnd, color='C0', alpha=0.1)

    ax.plot(vals_to_test, cubic_avg_rewards, color='C0',
            linestyle=':', linewidth=4, alpha=1, label="TCP Cubic")
    # cubic_low_bnd = np.array(cubic_avg_rewards) - \
    #     np.array(cubic_reward_errs)
    # cubic_up_bnd = np.array(cubic_avg_rewards) + \
    #     np.array(cubic_reward_errs)
    ax.fill_between(vals_to_test, cubic_low_bnd,
                    cubic_up_bnd, color='C0', alpha=0.1)

    ax.plot(vals_to_test, udr_small_avg_rewards, color='grey',
            linewidth=4, linestyle='-', label="UDR-1")
    # udr_small_low_bnd = np.array(
    #     udr_small_avg_rewards) - np.array(udr_small_reward_errs)
    # udr_small_up_bnd = np.array(
    #     udr_small_avg_rewards) + np.array(udr_small_reward_errs)
    ax.fill_between(vals_to_test, udr_small_low_bnd,
                    udr_small_up_bnd, color='grey', alpha=0.1)

    ax.plot(vals_to_test, udr_mid_avg_rewards, color='grey',
            linewidth=4, linestyle='--', label="UDR-2")
    # udr_mid_low_bnd = np.array(
    #     udr_mid_avg_rewards) - np.array(udr_mid_reward_errs)
    # udr_mid_up_bnd = np.array(udr_mid_avg_rewards) + \
    #     np.array(udr_mid_reward_errs)
    ax.fill_between(vals_to_test, udr_mid_low_bnd,
                    udr_mid_up_bnd, color='grey', alpha=0.1)

    ax.plot(vals_to_test, udr_large_avg_rewards, color='grey',
            linewidth=4, linestyle=':', label="UDR-3")
    # udr_large_low_bnd = np.array(
    #     udr_large_avg_rewards) - np.array(udr_large_reward_errs)
    # udr_large_up_bnd = np.array(
    #     udr_large_avg_rewards) + np.array(udr_large_reward_errs)
    ax.fill_between(vals_to_test, udr_large_low_bnd,
                    udr_large_up_bnd, color='grey', alpha=0.1)
    ax.set_xlim(min(vals_to_test), max(vals_to_test))
    ax.set_xlabel(xlabel)
    if dim == 'bandwidth' or dim == 'queue':
        ax.set_ylabel("Test Reward")
# plt.legend(bbox_to_anchor=(0,1.02,1,0.2), loc="lower left",
#                 mode="expand", borderaxespad=0, ncol=3)

line_labels = ["GENET", "BBR", "TCP Cubic", "UDR-1", "UDR-2", "UDR-3"]

lines_labels = [ax.get_legend_handles_labels() for ax in fig.axes]
lines, labels = [sum(lol, []) for lol in zip(*lines_labels)]

fig.subplots_adjust(hspace=0.5)
fig.legend(handles=lines,     # The line objects
           labels=line_labels,   # The labels for each line
           loc="upper center",   # Position of legend
           borderaxespad=0.1,    # Small spacing around legend box
            ncol=6, columnspacing=0.2, handletextpad=0.3,
            handlelength=1.2)
# plt.tight_layout()
    #     writer.writerow([dim, 'genet_avg_rewards', 'genet_low_bnd', 'genet_up_bnd',
    #         'bbr_avg_rewards', 'bbr_low_bnd', 'bbr_up_bnd',
    #         'cubic_avg_rewards', 'cubic_low_bnd', 'cubic_up_bnd',
    #         'udr_small_avg_rewards', 'udr_small_low_bnd', 'udr_small_up_bnd'
    #         'udr_mid_avg_rewards', 'udr_mid_low_bnd', 'udr_mid_up_bnd'
    #         'udr_large_avg_rewards', 'udr_large_low_bnd', 'udr_large_up_bnd'])
    #     writer.writerows(zip(vals_to_test,
    #             genet_avg_rewards, genet_low_bnd, genet_up_bnd,
    #             bbr_avg_rewards, bbr_low_bnd, bbr_up_bnd,
    #             cubic_avg_rewards, cubic_low_bnd, cubic_up_bnd,
    #             udr_small_avg_rewards, udr_small_low_bnd, udr_small_up_bnd,
    #             udr_mid_avg_rewards, udr_mid_low_bnd, udr_mid_up_bnd,
    #             udr_large_avg_rewards, udr_large_low_bnd, udr_large_up_bnd))
save_dir = os.path.join(RESULT_ROOT, EXP_NAME, "sim_eval_vary_summary.png")
save_dir = os.path.join(RESULT_ROOT, EXP_NAME, "sim_eval_vary_summary.pdf")
save_dir = os.path.join("../../figs/sim_eval_vary_summary.pdf")
plt.savefig(save_dir)