###############################################################################
#
#  Project:  Embedded Learning Library (ELL)
#  File:     imageConverter.py
#  Authors:  Chris Lovett
#
###############################################################################

import sys
import numpy as np
import cv2
import argparse

def resize_image(image, newSize):
    if (image.shape[0] > image.shape[1]):  # Tall(more rows than cols)
        rowStart = int((image.shape[0] - image.shape[1]) / 2)
        rowEnd = rowStart + image.shape[1]
        colStart = 0
        colEnd = image.shape[1]
    else:  # Wide(more cols than rows)
        rowStart = 0
        rowEnd = image.shape[0]
        colStart = int((image.shape[1] - image.shape[0]) / 2)
        colEnd = colStart + image.shape[0]

        cropped = image[rowStart:rowEnd, colStart:colEnd]
        resized = cv2.resize(cropped, newSize)
        return resized

def prepare_image_for_predictor(image, newSize, scale, bgr2rgb):
    if newSize[0] and newSize[1]:
        resized = resize_image(image, newSize)
    if bgr2rgb != 0:
        resized = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
    resized = resized.astype(np.float).ravel()
    if scale != 0:
        resized *= scale
    return resized

def save_raw(name, data):
    with open(name, 'wb') as f:
        f.write(bytearray(data))
        f.close()

def strtobool(arg):
    if arg.lower() in ['t', 'true', 'yes', '1']:
        return True
    return False

def main(argv):
    """
    This is a plugin for the debugCompiler tool.  It expects a list of string arguments as input and it returns the
    image preprocessed and prepared for input to an ELL test model.  The result has to be a list of floats.
    """
    arg_parser = argparse.ArgumentParser("imageConverter takes an image as input and converts it to an array of \
floating point numbers")
    arg_parser.add_argument("--bgr", default="False",
                            help="specify True if input data should be in BGR format (default False)")
    arg_parser.add_argument("--width", type=int, default=0, help="resize the image to this new width (default None)")
    arg_parser.add_argument("--height", type=int, default=0, help="resize the image to this new height (default None)")
    arg_parser.add_argument("--scale", type=float, default=0,
                            help="scale the floating point numbers by this amount (default None)")
    arg_parser.add_argument("--image", default=None, required=True,
                            help="name of file containing image as png, jpg, etc (default None)")
    arg_parser.add_argument("--outputData", "-od", default=None,
                            help="name of file to write the resulting raw floating point data (default None)")

    args = arg_parser.parse_args(argv)

    image = cv2.imread(args.image)
    if image is None:
        print("Error reading image {}".format(args.image))

    resized = prepare_image_for_predictor(image, (args.height, args.width), args.scale, args.bgr)

    if args.outputData:
        save_raw(args.outputData, resized)

    return list(resized)

if __name__ == "__main__":
    args = sys.argv
    main(args[1:])


"""
Adapted from
    - https://github.com/PyTorchLightning/lightning-bolts/blob/master/pl_bolts/models/self_supervised/swav/swav_module.py
Bugs to fix:
- STL10 augmentation
- weight_init, bn3 and so on.
"""

import os
import math

import torch
from torch import nn
from torch import distributed as dist
import pytorch_lightning as pl
from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint

from backbone import resnet18, resnet50, SwAV
from optimizer import LARS
from lr_scheduler import linear_warmup_decay
from dataloader import get_swav_transform, SwAVFinetuneTransform, STL10DataModule, CIFAR10DataModule
from callback import SSLOnlineEvaluator

def exclude_from_wt_decay(named_params, weight_decay):
    skip_list = ["bias", "bn"]
    params = []
    excluded_params = []
    for name, param in named_params:
        if not param.requires_grad:
            continue
        elif any(layer_name in name for layer_name in skip_list):
            excluded_params.append(param)
        else:
            params.append(param)
    return [
        {"params": params, "weight_decay": weight_decay},
        {"params": excluded_params, "weight_decay": 0.0},
    ]

class Net(pl.LightningModule):
    def __init__(
        self,
        # ---- Model -----
        backbone_arch,
        first_conv,
        maxpool1,
        hidden_mlp,
        feature_dim,
        num_prototypes,
        # ----- Optimizer/schedular -----
        optimizer,
        lr,
        exclude_bn_bias,
        weight_decay,
        warmup_epochs,
        max_epochs,
        # start_lr,
        # final_lr,
        # ----- Queue -----
        num_samples,
        queue_length,
        queue_path,
        queue_start_epoch,
        # ----- Miscellaneous -----
        dataset,
        total_gpus,
        batch_size,
        crops_for_assign,
        num_crops,
        sinkhorn_iterations,
        epsilon,
        temperature,
        freeze_prototypes_epochs,
    ):
        """
        Args:
            crops_for_assign: list of crop ids for computing assignment
            num_crops: number of global and local crops, ex: [2, 6]

            queue_length: set queue when batch size is small,
                must be divisible by total batch-size (i.e. total_gpus * batch_size),
                set to 0 to remove the queue
            queue_path: folder within the logs directory
            queue_start_epoch: start uing the queue after this epoch
        """

        super().__init__()
        self.total_batch_size = batch_size * total_gpus
        # self.queue_length = queue_length
        # self.queue_path = queue_path
        # self.queue_start_epoch = queue_start_epoch

        self.save_hyperparameters()

        if backbone_arch == "resnet18":
            backbone = resnet18
        elif backbone_arch == "resnet50":
            backbone = resnet50
        else:
            raise ValueError("Backbone unknown.")
        self.model = SwAV(
            backbone=backbone(first_conv=first_conv, maxpool1=maxpool1),
            hid_dim=hidden_mlp,
            out_dim=feature_dim,
            normalize=True,
            num_prototypes=num_prototypes,
        )

        self.get_assignments = (
            self.distributed_sinkhorn if total_gpus > 1 else self.sinkhorn
        )

    def setup(self, stage):
        assert self.hparams.queue_length % self.total_batch_size == 0
        self.queue = None
        if self.hparams.queue_length > 0:
            # if not os.path.exists(queue_folder):
            #     os.makedirs(queue_folder)
            self.hparams.queue_path = os.path.join(
                self.logger.log_dir,
                self.hparams.queue_path,
                "queue" + str(self.trainer.global_rank) + ".pth",
            )
            if os.path.isfile(self.hparams.queue_path):
                self.queue = torch.load(self.hparams.queue_path)["queue"]
            self.queue_use = False

    def forward(self, x):
        return self.model.backbone(x)

    def training_step(self, batch, batch_idx):
        loss = self.shared_step(batch)
        self.log("train_loss", loss)
        return loss

    def validation_step(self, batch, batch_idx):
        loss = self.shared_step(batch)
        self.log("val_loss", loss)

    def configure_optimizers(self):
        steps_per_epoch = math.ceil(
            self.hparams.num_samples // self.total_batch_size
        )
        if self.hparams.exclude_bn_bias:
            params = exclude_from_wt_decay(
                self.named_parameters(), weight_decay=self.hparams.weight_decay
            )
        else:
            params = self.parameters()

        if self.hparams.optimizer == "lars":
            optimizer = LARS(
                params,
                lr=self.hparams.lr,
                momentum=0.9,
                weight_decay=self.hparams.weight_decay,
                trust_coefficient=0.001,
            )
        elif self.hparams.optimizer == "adam":
            optimizer = torch.optim.Adam(
                params, lr=self.hparams.lr, weight_decay=self.hparams.weight_decay
            )

        warmup_steps = steps_per_epoch * self.hparams.warmup_epochs
        total_steps = steps_per_epoch * self.hparams.max_epochs

        return dict(
            optimizer=optimizer,
            lr_scheduler=torch.optim.lr_scheduler.LambdaLR(
                optimizer,
                linear_warmup_decay(warmup_steps, total_steps, cosine=True),
            ),
            interval="step",
            frequency=1,
        )

    # def optimizer_step(
    #     self,
    #     epoch,
    #     batch_idx,
    #     optimizer,
    #     optimizer_idx,
    #     optimizer_closure,
    #     on_tpu,
    #     using_native_amp,
    #     using_lbfgs,
    # ):
    #     for param_group in optimizer.param_groups:
    #         param_group["lr"] = self.lr_schedule[self.trainer.global_step]
    #     self.log(
    #         "lr",
    #         self.lr_schedule[self.trainer.global_step],
    #         on_step=True,
    #         on_epoch=False,
    #     )

    #     super().optimizer_step(
    #         epoch,
    #         batch_idx,
    #         optimizer,
    #         optimizer_idx,
    #         optimizer_closure,
    #         on_tpu,
    #         using_native_amp,
    #         using_lbfgs,
    #     )

    def on_after_backward(self) -> None:
        if self.current_epoch < self.hparams.freeze_prototypes_epochs:
            for name, p in self.model.named_parameters():
                if "prototypes" in name:
                    p.grad = None
                    break
            else:
                raise RuntimeError("No prototype tensor found in your model.")

    def on_train_epoch_start(self):
        if self.hparams.queue_length > 0:
            if self.current_epoch == self.hparams.queue_start_epoch and self.queue is None:
                # Init the queue
                self.queue = torch.zeros(
                    len(self.crops_for_assign),
                    self.hparams.queue_length // (self.hparams.total_gpus),
                    self.hparams.feature_dim,
                )
                if self.hparams.total_gpus > 0:
                    self.queue = self.queue.cuda()

    def on_train_epoch_end(self, outputs) -> None:
        if self.queue is not None:
            torch.save({"queue": self.queue}, self.queue_path)

    def shared_step(self, batch):
        if self.hparams.dataset == "stl10":
            batch = batch[0]  # unlabeled data
        images, _ = batch
        images = images[:-1]

        # Step 1: normalize the prototypes
        with torch.no_grad():
            w = self.model.head.prototype.weight.data.clone()
            w = nn.functional.normalize(w, dim=1, p=2)
            self.model.head.prototype.weight.copy_(w)

        # Step 2: multi forward pass
        embedding, outputs = self.model(images)
        embedding = embedding.detach()

        # Step 3: SwAV loss
        batch_size = self.hparams.batch_size
        loss = 0
        for i, crop_id in enumerate(self.hparams.crops_for_assign):
            with torch.no_grad():
                out = outputs[batch_size * crop_id : batch_size * (crop_id + 1)]

                # Step 4: use the queue
                if self.queue is not None:
                    if not self.queue_use:
                        if not torch.all(self.queue[i, -1, :] == 0):
                            # Only concatenate the queue when it is filled.
                            self.queue_use = True
                    if self.queue_use:
                        out = torch.cat(
                            (self.queue[i] @ self.model.head.prototype.weight.t(), out)
                        )
                    # fill the queue
                    self.queue[i, batch_size:] = self.queue[i, :-batch_size].clone()
                    self.queue[i, :batch_size] = embedding[
                        crop_id * batch_size : (crop_id + 1) * batch_size
                    ]

                # Step 5: get assignments
                q = torch.exp(out / self.hparams.epsilon).t()
                q = self.get_assignments(q, self.hparams.sinkhorn_iterations)[
                    -batch_size:
                ]
            sub_loss = 0
            for v in torch.arange(sum(self.hparams.num_crops)):
                if v == crop_id:
                    continue
                p = (
                    outputs[batch_size * v: batch_size * (v + 1)]
                    / self.hparams.temperature
                ).softmax(dim=1)
                sub_loss -= (q * torch.log(p)).sum(dim=1).mean()
            loss += sub_loss / (sum(self.hparams.num_crops) - 1)
        return loss / len(self.hparams.crops_for_assign)

    # def get_linear_warmup_cosine_annealing_scheduler(self):
    #     global_batch_size = self.hparams.total_gpus * self.hparams.batch_size
    #     step_per_epoch = torch.ceil(self.hparams.num_samples // global_batch_size)

    #     warmup_lrs = torch.linspace(
    #         self.hparams.start_lr, self.hparams.lr, self.warmup_epochs * step_per_epoch
    #     )

    #     cosine_steps = step_per_epoch * (
    #         self.hparams.max_epochs - self.hparams.warmup_epochs
    #     )
    #     # Angular frequency of cosine function. Thus the period equals 2 * cosine_steps.
    #     w = math.pi / cosine_steps
    #     # Amplitude of cosine function
    #     a = 0.5 * (self.hparams.lr - self.hparams.final_lr)
    #     cosine_lrs = [
    #         0.5 * (self.hparams.lr + self.hparams.final_lr) + a * torch.cos(w * step)
    #         for step in range(cosine_steps)
    #     ]
    #     return torch.cat([warmup_lrs, cosine_lrs])

    def sinkhorn(self, Q, num_iters):
        with torch.no_grad():
            Q /= Q.sum()
            K, B = Q.shape
            u = torch.zeros(K, device=self.device)
            r = torch.ones(K, device=self.device) / K
            c = torch.ones(B, device=self.device) / B
            for _ in range(num_iters):
                u = Q.sum(dim=1)
                Q *= (r / u).unsqueeze(1)
                Q *= (c / Q.sum(dim=0)).unsqueeze(0)
            return (Q / Q.sum(dim=0, keepdim=True)).t().float()

    def distributed_sinkhorn(self, Q, nmb_iters):
        with torch.no_grad():
            sum_ = Q.sum()
            dist.all_reduce(sum_)

            K, B = Q.shape
            r = torch.ones(K).cuda(non_blocking=True) / K
            c = torch.ones(B).cuda(non_blocking=True) / (self.hparams.total_gpus * B)

            for _ in range(nmb_iters):
                sum_1 = Q.sum(dim=1)
                dist.all_reduce(sum_1)
                Q *= (r / sum_1).unsqueeze(1)
                Q *= (c / Q.sum(dim=0)).unsqueeze(0)
            return (Q / Q.sum(dim=0, keepdim=True)).t().float()

if __name__ == "__main__":
    pl.seed_everything(42)
    # [32, 16] for CIFAR10, [96, 36] for STL10 [224, 96] for ImageNet
    SIZE_CROPS = [32, 16]
    CROPS_FOR_ASSIGN = [0, 1]
    NUM_CROPS = [2, 1]  # [2, 1] for CIFAR10, [2, 4] for STL10, [2, 6] for ImageNet.
    # [0.33, 0.10] for CIFAR10 & STL10, [0.14, 0.05] for ImageNet
    MIN_SCALE_CROPS = [0.33, 0.10]
    MAX_SCALE_CROPS = [1, 0.33]
    USE_GAUSSIAN_BLUR = True  # False for CIFAR10, True for ImageNet & STL10
    JITTER_STRENGTH = 1.0
    DATA_DIR = "/home/tiankang/wusuowei/data/"
    BATCH_SIZE = 128
    DATAMODULE = CIFAR10DataModule
    transforms = get_swav_transform(
        means=DATAMODULE.means,
        stds=DATAMODULE.stds,
        size_crops=SIZE_CROPS,
        num_crops=NUM_CROPS,
        min_scale_crops=MIN_SCALE_CROPS,
        max_scale_crops=MAX_SCALE_CROPS,
        use_gaussian_blur=USE_GAUSSIAN_BLUR,
        jitter_strength=JITTER_STRENGTH,
    )
    datamodule = DATAMODULE(DATA_DIR, BATCH_SIZE, *transforms)

    BACKBONE_ARCH = "resnet50"
    FIRST_CONV = False  # False for CIFAR10 & STL10, True for ImageNet
    MAXPOOL1 = False  # False for CIFAR10 & STL10, True for ImageNet
    HIDDEN_MLP = 2048
    FEATURE_DIM = 128
    NUM_PROTOTYPES = [512]

    OPTIMIZER = "adam"  # 'adam' for CIFAR10 & STL10, 'lars' for ImageNet
    LR = 1e-3
    WEIGHT_DECAY = 1e-6
    EXCLUDE_BN_BIAS = True
    WARMUP_EPOCHS = 10
    MAX_EPOCHS = 100
    MAX_STEPS = -1

    NUM_SAMPLES = 50000  # 95_000 for STL10 mixed dataloader
    QUEUE_PATH = "queue"
    QUEUE_LENGTH = 0
    QUEUE_START_EPOCH = 15

    GPUS = [9]
    TOTAL_GPUS = len(GPUS) * 1  # I only have one node, sadly
    SINKHORN_ITERATIONS = 3
    EPSILON = 0.05
    TEMPERTURE = 0.1
    FREEZE_PROTOTYPES_EPOCHS = 1

    model = Net(
        backbone_arch=BACKBONE_ARCH,
        first_conv=FIRST_CONV,
        maxpool1=MAXPOOL1,
        hidden_mlp=HIDDEN_MLP,
        feature_dim=FEATURE_DIM,
        num_prototypes=NUM_PROTOTYPES,
        optimizer=OPTIMIZER,
        lr=LR,
        weight_decay=WEIGHT_DECAY,
        exclude_bn_bias=EXCLUDE_BN_BIAS,
        warmup_epochs=WARMUP_EPOCHS,
        max_epochs=MAX_EPOCHS,
        num_samples=NUM_SAMPLES,
        queue_length=QUEUE_LENGTH,
        queue_path=QUEUE_PATH,
        queue_start_epoch=QUEUE_START_EPOCH,
        dataset=DATAMODULE.name,
        total_gpus=TOTAL_GPUS,
        batch_size=BATCH_SIZE,
        crops_for_assign=CROPS_FOR_ASSIGN,
        num_crops=NUM_CROPS,
        sinkhorn_iterations=SINKHORN_ITERATIONS,
        epsilon=EPSILON,
        temperature=TEMPERTURE,
        freeze_prototypes_epochs=FREEZE_PROTOTYPES_EPOCHS,
    )

    callbacks = [
        LearningRateMonitor(logging_interval="step"),
        ModelCheckpoint(save_last=True, save_top_k=1, monitor="val_loss"),
        SSLOnlineEvaluator(
            dataset=DATAMODULE.name,
            z_dim=model.model.backbone.out_dim,
            hid_dim=None,
            num_classes=DATAMODULE.num_classes,
            drop_p=0.0,
        ),
    ]
    CHECKPOINT = ""
    FINETUNE = False
    if not CHECKPOINT:
        trainer = pl.Trainer(
            deterministic=True,
            max_epochs=MAX_EPOCHS,
            max_steps=None if MAX_STEPS == -1 else MAX_STEPS,
            gpus=GPUS,
            num_nodes=1,
            distributed_backend="ddp" if TOTAL_GPUS > 1 else None,
            sync_batchnorm=True if TOTAL_GPUS > 1 else False,
            precision=16,
            callbacks=callbacks,
            # fast_dev_run=True,
        )

        trainer.fit(model, datamodule=datamodule)
    else:
        if not FINETUNE:
            trainer = pl.Trainer(
                deterministic=True,
                resume_from_checkpoint=CHECKPOINT,
                max_epochs=MAX_EPOCHS,
                max_steps=None if MAX_STEPS == -1 else MAX_STEPS,
                gpus=GPUS,
                num_nodes=1,
                distributed_backend="ddp" if TOTAL_GPUS > 1 else None,
                sync_batchnorm=True if TOTAL_GPUS > 1 else False,
                precision=32,
                callbacks=callbacks,
                # fast_dev_run=True,
            )
            trainer.fit(model, datamodule=datamodule)
        else:
            raise NotImplementedError


import sys
import pygame
from pygame import *

import Bird
from Bird import *

# Create a const variables for game
WIDTH = 1000
HEIGHT = 500
FPS = 30
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (100, 200, 255)

# Function which init game and screen for game
def run():
    pygame.init()
    pygame.mixer.init()
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    pygame.display.set_caption("Pixel Birds")
    clock = pygame.time.Clock()
    screen.fill(BLUE)
    pygame.display.flip()

    bird = Bird(500, 200)
    left = right = False
    up = False

    def render_screen():
        pygame.display.update()
        screen.fill(BLUE)

    def render_bird():
        bird.update(left, right, up)
        bird.draw(screen)

    # Cycle for game event
    game_running = True
    while game_running:
        clock.tick(FPS)
        render_bird()
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                game_running = False
            elif event.type == KEYDOWN and event.key == K_LEFT:
                left = True
            elif event.type == KEYDOWN and event.key == K_RIGHT:
                right = True
            elif event.type == KEYDOWN and event.key == K_UP:
                up = True
            elif event.type == KEYUP and event.key == K_UP:
                up = False
            elif event.type == KEYUP and event.key == K_LEFT:
                left = False
            elif event.type == KEYUP and event.key == K_RIGHT:
                right = False
            elif event.type == QUIT:
                return sys.exit()
        render_screen()

run()


from src.engine.blocks import Block, Convex, Flat
import pytest
import mock

def test_block_class_object_cant_be_created():
    with pytest.raises(TypeError):
        Block()

def test_convex_class_object_cant_be_created():
    with pytest.raises(TypeError):
        Convex()

def test_flat_class_object_cant_be_created():
    with pytest.raises(TypeError):
        Flat()

def test_block_class_has_static_empty_method_texture():
    Block.__abstractmethods__ = set()
    assert Block.texture() is None

def test_block_class_has_empty_method_interact_with_snake():
    Block.__abstractmethods__ = set()
    snake = mock.Mock()
    assert Block().interact_with_snake(snake) is None

def test_block_class_set_field():
    field = mock.Mock()
    Block.__abstractmethods__ = set()
    block = Block()
    block.set_field(field)
    assert block.field == field

def test_flat_class_has_empty_method_interact_with_convex():
    Flat.__abstractmethods__ = set()
    convex = mock.Mock()
    assert Flat().interact_with_convex(convex) is None

def test_convex_class_has_empty_method_move():
    Convex.__abstractmethods__ = set()
    direction = mock.Mock()
    assert Convex().move(direction) is None

def test_convex_class_has_empty_method_check_snake_move():
    Convex.__abstractmethods__ = set()
    assert Convex().check_snake_move() is None

def test_convex_class_has_empty_method_check_move():
    Convex.__abstractmethods__ = set()
    direction = mock.Mock()
    assert Convex().check_move(direction) is None

def test_block_self_draw_when_drawing_same_size_as_remembered():
    Block().__abstractmethods__ = set()
    frame = mock.Mock()
    position = mock.Mock()
    displayed_texture = mock.Mock()
    block = Block()
    block.displayed_side_length = 10
    block.displayed_texture = displayed_texture
    block.self_draw(frame, position, 10)
    assert block.displayed_texture == displayed_texture
    frame.blit.assert_called_once_with(displayed_texture, position)

def test_convex_that_is_not_alive_self_draw_wont_draw_in_frame_and_will_be_removed_from_its_field():
    Convex().__abstractmethods__ = set()
    convex = Convex()
    convex.is_alive = False
    frame = mock.Mock()
    position = mock.Mock()
    field = mock.Mock()
    convex.field = field
    convex.self_draw(frame, position, 77)
    field.remove_convex.assert_called_once()
    assert not frame.blit.called

def test_convex_that_is_alive_self_draw_will_draw_in_frame_when_remembered_size_equals_current_one():
    Convex().__abstractmethods__ = set()
    convex = Convex()
    convex.is_alive = True
    frame = mock.Mock()
    position = mock.Mock()
    field = mock.Mock()
    convex.field = field
    convex.displayed_side_length = 1234
    convex.self_draw(frame, position, 1234)
    frame.blit.assert_called_once()


import logging
import warnings
from copy import deepcopy
from datetime import datetime, timedelta
from datetime import time
from pathlib import Path
from tkinter.filedialog import askopenfilenames, askopenfilename
from tkinter.messagebox import showerror

import numpy as np
import pandas as pd
import utm
from tqdm import tqdm

from src.magrnd.ground_one.data_processing.consts import GZ_DIR, EXPECTED_TIMESTAMP_LEN, IA_SAMPLE_RATE_HZ, WINDOW_TITLE, \
    VALIDATED_STR
from src.magrnd.ground_one.loaders.file_identifier import identify_scan
from src.magrnd.ground_one.loaders.ia_file_encoding_fixer import fix_encoding
from src.mag_utils.mag_utils.functional import validator
from src.mag_utils.mag_utils.functional.mag_eili_converter import magmap_file_reformatter
from src.mag_utils.mag_utils.loader import blackwidow, base
from src.mag_utils.mag_utils.loader import main_load
from src.mag_utils.mag_utils.scans import AluminumManScan
from src.mag_utils.mag_utils.scans import magScan, HorizontalScan

warnings.filterwarnings("ignore")

def mag_utils_wrapper_func(path_to_scan_file: Path, scan_type):
    SCAN_TYPE_MAPPING = {"widow": blackwidow,
                         "base": base}

    # create gz dir if does not exist
    if not GZ_DIR.exists():
        GZ_DIR.mkdir()

    # hide logging
    logger = logging.getLogger("mag_utils.functional.logger")
    logger.setLevel(logging.INFO)

    # perform validation of data
    validated_file_path = GZ_DIR / "validated.txt"

    # scan_type_object = base

    validator.validate(file_to_validate_path=str(path_to_scan_file),
                       validated_file_path=str(validated_file_path),
                       log_file_path=str(GZ_DIR / "faulty_lines_log.txt"),
                       file_type=scan_type)

    scan = SCAN_TYPE_MAPPING[scan_type].load(validated_file_path.__str__())

    # delete validated file
    validated_file_path.unlink()
    scan.file_name = str(path_to_scan_file)
    return scan

def auto_format_time(time_str: str):
    if time_str is None:
        return
    # remove 1900-01-01 from old GZ IA files
    if '1900-' in time_str:
        time_str = time_str[11:]
    if ":" not in time_str:
        # determine if time is valid, and if not make it valid
        zeros_string = '0' * (6 - len(time_str[:time_str.rindex('.')]))
        time_str = (zeros_string + time_str)
        ms_s = "" if "." not in time_str else time_str[time_str.index(".") + 1:]

        # extract time data from string
        microsecond = int(
            float(ms_s) * 10 ** (EXPECTED_TIMESTAMP_LEN - len(
                ms_s))) if ms_s else 0  # for taking only first 6 digits after the decimal
        hour, minute, second = time_str[:2], time_str[2:4], time_str[4:time_str.index(".")]
    else:
        hour, minute, second = tuple(time_str.split(":"))
        second, microsecond = tuple(second.split(".")) if "." in second else (second, 0)

    return time(hour=int(hour),
                minute=int(minute),
                second=int(second),
                microsecond=int(microsecond))

def load_gz(path_to_scan_file: Path):
    # save original format time function to undo change later
    original_format_time_function = blackwidow.format_time

    # set to use custom function
    blackwidow.format_time = auto_format_time

    # load scan
    scan = blackwidow.lab_txt_format_loader(str(path_to_scan_file))

    # return format time to original state
    blackwidow.format_time = original_format_time_function

    return scan

def load_widow(path_to_scan_file: Path):
    return mag_utils_wrapper_func(path_to_scan_file, "widow")

def load_base(path_to_scan_file: Path):
    return mag_utils_wrapper_func(path_to_scan_file, "base")

def load_ia_raw(path_to_scan_file: Path, path_to_calib_file=None):  # todo: work with mike

    """
    :param path: str scan file path
    :return: ia matrix of bx,by,height,time
    """

    ## LOAD IA DATA ##

    # read all gps data
    gps_df = pd.read_csv(path_to_scan_file, delimiter=',', skiprows=3, names=list(range(15)), skipfooter=2,
                         error_bad_lines=True)
    gps_df = gps_df.iloc[gps_df[1].dropna().index, :]

    # read all sensor data
    b_df = pd.read_csv(path_to_scan_file, delimiter=' ', skiprows=3, names=list(range(8)), skipfooter=2,
                       error_bad_lines=True)
    b_df = b_df.iloc[b_df[7].dropna().index, :]

    gps_df.iloc[:, 2] = gps_df.iloc[:, 2].interpolate()
    gps_df.iloc[:, 4] = gps_df.iloc[:, 4].interpolate()

    height = gps_df.iloc[:, 9].values

    # handle time format
    time = gps_df.iloc[:, 1].values
    time_s = [t.zfill(8) for t in time.astype(str)]
    t_seconds = np.array([int(t[:2]) * 3600 + int(t[2:4]) * 60 + float(t[4:]) for t in
                          time_s])  # time in seconds and not H:M:S

    # find if there is a day transfer within scan
    diff_t = np.diff(t_seconds)
    if any(diff_t < 0):
        next_day_index, = np.where(diff_t < 0)[0]
        t_seconds[next_day_index + 1:] += 3600 * 24

    gps_index = b_df.iloc[:, 2].values

    vectorized_translate_gps = np.vectorize(blackwidow.translate_gps)
    x_utm, y_utm = vectorized_translate_gps(gps_df.iloc[:, 4].values, gps_df.iloc[:, 2].values)

    sensor_df_new = pd.DataFrame({'bx': b_df.iloc[:, 5],
                                  'by': b_df.iloc[:, 6],
                                  'bz': b_df.iloc[:, 7],
                                  'gps_index': gps_index})
    gps_df_new = pd.DataFrame(
        {'x': x_utm,
         'y': y_utm,
         'height': height,
         'time': t_seconds})

    sensor_df_new = sensor_df_new.append({'x': None, 'y': None, 'height': None, 'time': None},
                                         ignore_index=True)

    # merging gps and b ia data and interpolating data
    groups = sensor_df_new.groupby('gps_index')
    ia_data_matrix = pd.DataFrame()
    index_for_adding_gps = 0

    for group in tqdm(groups):
        if index_for_adding_gps >= len(gps_df_new['x']):
            pass
        else:
            index_for_group = group[1].index[0]
            group[1]['x'][index_for_group] = gps_df_new['x'][index_for_adding_gps]
            group[1]['y'][index_for_group] = gps_df_new['y'][index_for_adding_gps]
            group[1]['height'][index_for_group] = gps_df_new['height'][index_for_adding_gps]
            group[1]['time'][index_for_group] = gps_df_new['time'][index_for_adding_gps]

        ia_data_matrix = pd.concat([ia_data_matrix, group[1]])
        index_for_adding_gps += 1
    ia_data_matrix.drop(columns='gps_index', inplace=True)
    ia_data_matrix = ia_data_matrix.astype('float64').interpolate()

    # translate numbers back to time
    ia_data_matrix["time"] = [
        (datetime(year=1970, month=1, day=1, hour=0, minute=0, second=0) + timedelta(seconds=t)).time() for t in
        ia_data_matrix["time"]]

    ## CALIBRATE
    from src.magrnd.ground_one.data_processing.ia_calibration import apply_ia_calib
    calib_xyz = apply_ia_calib(ia_data_matrix, path_to_calib_file)
    ia_data_matrix["B"] = np.linalg.norm(calib_xyz, axis=1)

    if path_to_calib_file is None:
        return ia_data_matrix

    return magScan(file_name=str(path_to_scan_file),
                   x=ia_data_matrix['x'],
                   y=ia_data_matrix['y'],
                   a=ia_data_matrix['height'],
                   b=ia_data_matrix['B'],
                   time=ia_data_matrix['time'],
                   is_base_removed=True)

def load_ia_calibrated(path_to_scan_file: Path):
    df = pd.read_csv(path_to_scan_file, delimiter=r"\s+")
    vectorized_from_latlon = np.vectorize(utm.from_latlon)
    gps = vectorized_from_latlon(df["Latitude"], df["Longitude"], 36)

    return magScan(file_name=str(path_to_scan_file),
                   x=gps[0],
                   y=gps[1],
                   a=df['Altitude'],
                   b=df['MagneticField'],
                   time=df['Time'],
                   is_base_removed=False)

def universal_loader(path_to_scan_file: Path, path_to_calib_file: Path = None):
    # fix file encoding if needed
    fix_encoding(path_to_scan_file, path_to_calib_file)

    scan_type = identify_scan(path_to_scan_file)
    if scan_type == "widow":
        return load_widow(path_to_scan_file)
    elif scan_type == "base":
        return load_base(path_to_scan_file)
    elif scan_type == "ia_raw":
        # if path to calib file was not passed, ask user to provide it
        if path_to_calib_file is None:
            path_to_calib_file = Path(askopenfilename(title="Select calibration file"))
        scan = load_ia_raw(path_to_scan_file, path_to_calib_file)
        scan.sampling_rate = IA_SAMPLE_RATE_HZ
        return scan
    elif scan_type == "ia_calibrated":
        return load_ia_calibrated(path_to_scan_file)
    elif scan_type == "gz":
        return load_gz(path_to_scan_file)
    else:
        showerror(WINDOW_TITLE, "File type not recognized. Please try again")
        raise NotImplementedError

def load(scan_paths: list = None):
    if scan_paths is None:
        scan_paths = [Path(file_path) for file_path in
                      askopenfilenames(title="Select file",
                                       filetypes=[("All Files", "*.*")])]

    # if one file is selected, make it the scan path, else opt for vertical scan loading
    if isinstance(scan_paths, Path):
        scan_paths = [scan_paths]

    if not len(scan_paths):
        raise ValueError("Invalid number of files selected.")

    loaded_scans = []
    for scan_path in scan_paths:
        # handle mag eili files
        if scan_path.suffix == ".dat":
            scan_path = magmap_file_reformatter(scan_path)

        # try loading thru main_load, if unsuccessful, use universal loader
        try:
            scan = main_load.load(str(scan_path), validate=True)

            if isinstance(scan, AluminumManScan):
                raise Exception("Aluminum man is not supported in mag_utils yet!")

            # removing "validated" from the file name
            if VALIDATED_STR in scan.file_name:
                validated_index = scan.file_name.rindex(VALIDATED_STR)
                scan.file_name = scan.file_name[:validated_index] + scan.file_name[
                                                                    validated_index + len(VALIDATED_STR):]
        except Exception as e:
            print(f"Encountered error: {e}, falling back to universal loader..")

        scan = universal_loader(scan_path)

        loaded_scans.append(scan)

    # check if given more than one HorizontaScan if so, join all scans together. if given single scan return only it.
    if isinstance(loaded_scans[0], HorizontalScan):
        if len(loaded_scans) > 1:
            appended_scans = deepcopy(loaded_scans[0])
            for scan in loaded_scans[1:]:
                appended_scans = appended_scans.append(scan)
            return appended_scans

    if len(loaded_scans) == 1:
        return loaded_scans[0]

    return loaded_scans


from flask import Flask, request
app = Flask("ccserver")
datadict = dict()

def log(text:str, filename="log.txt"):
    with open(filename, "a") as f:
        f.write("\n" + text)

@app.route("/up", methods=["GET"])
def up():
    source = request.args.get("source")
    target = request.args.get("target")
    data = request.args.get("data")
    log(data, "up.txt")
    if target in list(datadict.keys()):
        datadict[target].append({source:data})
    else:
        datadict[target] = [{source:data}]
    log(str(datadict), "log.txt")
    return ""

@app.route("/down", methods=["GET"])
def down():
    source = request.args.get("source")
    target = request.args.get("target")
    try:
        for i, source_data in enumerate(datadict[target]):
            if list(source_data.keys())[0] == source:
                requested_data = datadict[target].pop(i)
                return source_data[source]
    except Exception as e:
        log(str(e), "down.txt")
        return ""

app.run("0.0.0.0", port=1234, debug=True)


import cv2

import serial
import time

def write_read(message):
    if arduino_output:
        arduino.write(bytes(str(message), 'utf-8'))
        time.sleep(0.05)
        data = arduino.readline()
        return data
    return ""

def set_pos(x, y):
    out = str(x)+str(y)+"\n"
    value = write_read(out)
    print(f"write :{out}")
    print(f"read :{value}")

# ^ pip install opencv-contrib-python

# Enable camera

arduino_output = False  # constant for arduino communication, false to test with no arduino
debug = True  # print vars? display image?

if arduino_output:
    arduino = serial.Serial(port='COM20', baudrate=115200, timeout=.1)  # change com port depending on the com port

resX = 1000  # constants for camera res (does not need to match real res)
resY = 500

cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
cap.set(3, resX)
cap.set(4, resY)

# import cascade file for facial recognition
faceCascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml")

# if you want to detect any object for example eyes, use one more layer of classifier as below:
eyeCascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_eye_tree_eyeglasses.xml")

dxd = 0
dyd = 0

eye_conf_threshold = 3  # const
face_conf_threshold = 3  # const

eye_coords = []
targetX, targetY = 0, 0
dx, dy = 0, 0
loop_counter = 0

while True:
    success, img = cap.read()
    try:
        imgGray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    except:  # Fix pep thing Vojta
        raise Exception("Camera not found!")

    # Getting corners around the face

    # detecting eyes
    eyes = eyeCascade.detectMultiScale3(
        imgGray,
        outputRejectLevels=True)
    faces = faceCascade.detectMultiScale3(
        imgGray,
        scaleFactor=1.1,
        minNeighbors=5,
        minSize=(30, 30),
        outputRejectLevels=True)

    eye_coords = []
    face_coords = []
    # print(faces)
    # print(f"confidence = {eyes[2][0][0]}")
    #
    # for (x, y, w, h) in faces[0]:
    #     face_coords.append((x + (w/2), y + (h/2)))
    #     cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)

    for counter in range(len(eyes[2])):
        if eyes[2][counter][0] > eye_conf_threshold:
            ex, ey, ew, eh = eyes[0][counter]
            eye_coords.append((ex + (ew / 2), ey + (eh / 2)))
            img = cv2.rectangle(img, (ex, ey), (ex + ew, ey + eh), (255, 0, 0), 3)

    eye_in_face_coords = (-1, -1)

    for counter in range(len(faces[2])):
        if faces[2][counter][0] > face_conf_threshold:
            fx, fy, fw, fh = faces[0][counter]
            for eye_coord in eye_coords:
                currentX = eye_coord[0]
                currentY = eye_coord[1]
                if fx < currentX < fx+fw and fy < currentY < fy+fh:
                    eye_in_face_coords = (currentX, currentY)
            eye_coords.append((fx + (fw / 2), fy + (fh / 2)))
            img = cv2.rectangle(img, (fx, fy), (fx + fw, fy + fh), (255, 0, 0), 3)

    eyesX = []
    eyesY = []
    if eye_in_face_coords != (-1, -1):
        targetX = eye_in_face_coords[0]
        targetY = eye_in_face_coords[1]
    elif len(eye_coords) > 0:
        for eye in eye_coords:
            eyesX.append(eye[0])
            eyesY.append(eye[1])
        targetX = int(sum(eyesX) / len(eyesX))
        targetY = int(sum(eyesY) / len(eyesY))
    else:  # if no target currently on screen, do not move camera
        targetX = (resX/2)
        targetY = (resY/2)

    weightedX = int(float((resX/2) - targetX)/(resX/100))  # expected range: [-100,100]
    weightedY = int(float((resY/2) - targetY)/(resY/100))  # expected range: [-100,100]

    if loop_counter < 50:
        loop_counter += 1
    else:
        loop_counter = 0

    if -10 > weightedX:
        dxd = 2  # 2 means -1, but using only one char
    elif weightedX > 10:
        dxd = 1
    else:
        dxd = 0

    if -10 > weightedY:
        dyd = 1
    elif weightedY > 10:
        dyd = 2  # 2 means -1, but using only one char
    else:
        dyd = 0

    if abs(weightedX) > loop_counter:
        dx = dxd
    else:
        dx = 0
    if abs(weightedY) > loop_counter:
        dy = dyd
    else:
        dy = 0
    if debug:
        print(f"weightedX: {weightedX}\nweightedY: {weightedY}")
        print(f"dxd: {dxd}\ndyd: {dyd}")
        print(f"dx: {dx}\ndy: {dy}")
        print(f"targetX: {targetX}\ntargetY: {targetY}")
        set_pos(dx, dy)
        img = cv2.rectangle(img, (int(targetX-5), int(targetY-5)), (int(targetX+5), int(targetY+5)), (255, 255, 0), 3)
        if eye_in_face_coords != (-1, -1):
            eye_in_face_x = eye_in_face_coords[0]
            eye_in_face_y = eye_in_face_coords[1]
            img = cv2.rectangle(img, (int(eye_in_face_x - 5), int(eye_in_face_y - 5)),
                                (int(eye_in_face_x + 5), int(eye_in_face_y + 5)), (255, 255, 255), 3)
        cv2.imshow('face_detect', img)

    if cv2.waitKey(10) & 0xFF == ord('q'):
        break
cap.release()
cv2.destroyWindow('face_detect')


from __future__ import print_function

import os
import re
import sys
import math
import glob
import time
import pandas
import random
import logging
import weakref
import schedule
import argparse
import datetime
import threading
import collections
import cv2
import csv
from PIL import Image
from csv import writer
import tensorflow as tf
from collections import deque
from agents.rl_agent import *
from agents.vision_agent import *
from agents.followline_agent import *
from carla_birdeye_view import BirdViewProducer, BirdViewCropType, PixelDimensions

try:
    import pygame
    from pygame.locals import KMOD_CTRL
    from pygame.locals import KMOD_SHIFT
    from pygame.locals import K_0
    from pygame.locals import K_9
    from pygame.locals import K_BACKQUOTE
    from pygame.locals import K_BACKSPACE
    from pygame.locals import K_COMMA
    from pygame.locals import K_DOWN
    from pygame.locals import K_ESCAPE
    from pygame.locals import K_F1
    from pygame.locals import K_LEFT
    from pygame.locals import K_PERIOD
    from pygame.locals import K_RIGHT
    from pygame.locals import K_SLASH
    from pygame.locals import K_SPACE
    from pygame.locals import K_TAB
    from pygame.locals import K_UP
    from pygame.locals import K_a
    from pygame.locals import K_c
    from pygame.locals import K_d
    from pygame.locals import K_h
    from pygame.locals import K_m
    from pygame.locals import K_p
    from pygame.locals import K_q
    from pygame.locals import K_r
    from pygame.locals import K_s
    from pygame.locals import K_w
except ImportError:
    raise RuntimeError('cannot import pygame, make sure pygame package is installed')

try:
    import numpy as np
except ImportError:
    raise RuntimeError(
        'cannot import numpy, make sure numpy package is installed')

gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

# TOWN 02:
# 12 - left turn at the gas station
# 20 - right turn at the gas station

# 17 - left turn at the gas station
# 22 - right turn at the gas station
global SPAWNPOINT
SPAWNPOINT = 192

# ==============================================================================
# -- find carla module ---------------------------------------------------------
# ==============================================================================
try:
    sys.path.append(glob.glob('../carla/dist/*%d.%d-%s.egg' % (
        sys.version_info.major,
        7,
        'win-amd64' if os.name == 'nt' else 'linux-x86_64'))[0])
except IndexError:
    pass

import carla
from carla import ColorConverter as cc

# ==============================================================================
# -- Global functions ----------------------------------------------------------
# ==============================================================================

def find_weather_presets():
    rgx = re.compile('.+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)')
    name = lambda x: ' '.join(m.group(0) for m in rgx.finditer(x))
    presets = [x for x in dir(carla.WeatherParameters) if re.match('[A-Z].+', x)]
    return [(getattr(carla.WeatherParameters, x), name(x)) for x in presets]

def get_actor_display_name(actor, truncate=250):
    name = ' '.join(actor.type_id.replace('_', '.').title().split('.')[1:])
    return (name[:truncate-1] + u'\u2026') if len(name) > truncate else name

# ==============================================================================
# -- World ---------------------------------------------------------------
# ==============================================================================

class World(object):
    def __init__(self, carla_world, hud, args):
        self.world = carla_world
        self.hud = hud
        self.vehicle = None
        self.collision_sensor = None
        self.camera_manager = None
        self._weather_presets = find_weather_presets()
        self._weather_index = 0
        self.restart()
        self.world.on_tick(hud.on_world_tick)
        self.args = args

    def restart(self):
        # Keep same camera config if the camera manager exists.
        cam_index = self.camera_manager._index if self.camera_manager is not None else 0
        cam_pos_index = self.camera_manager._transform_index if self.camera_manager is not None else 0

        blueprint = self.world.get_blueprint_library().find('vehicle.tesla.model3')
        blueprint.set_attribute('role_name', 'hero')
        if blueprint.has_attribute('color'):
            color = random.choice(blueprint.get_attribute('color').recommended_values)
            blueprint.set_attribute('color', color)

        # Spawn the vehicle.
        if self.vehicle is not None:
            spawn_point = self.vehicle.get_transform()
            spawn_point.location.z += 2.0
            spawn_point.rotation.roll = 0.0
            spawn_point.rotation.pitch = 0.0
            self.destroy()

            spawn_points = self.world.get_map().get_spawn_points()
            spawn_point = spawn_points[SPAWNPOINT]
            self.vehicle = self.world.spawn_actor(blueprint, spawn_point)

        while self.vehicle is None:
            spawn_points = self.world.get_map().get_spawn_points()
            spawn_point = spawn_points[SPAWNPOINT]
            random_starting_point = bool(random.getrandbits(1))
            """if random_starting_point:
                random_value = random.randint(20, 50)
                spawn_point.rotation.yaw += random_value
                print(f"RIGHT: {random_value}")
            else:
                random_value = random.randint(20, 50)
                spawn_point.rotation.yaw -= random_value
                print(f"LEFT: {random_value}")"""

            self.vehicle = self.world.spawn_actor(blueprint, spawn_point)

        # Set up the sensors.
        self.collision_sensor = CollisionSensor(self.vehicle, self.hud)
        self.camera_manager = CameraManager(self.vehicle, self.hud)
        self.camera_manager._transform_index = cam_pos_index
        self.camera_manager.set_sensor(cam_index, notify=False)
        self.camera_manager.set_sensor_2(500, notify=False)
        actor_type = get_actor_display_name(self.vehicle)
        self.hud.notification(actor_type)

    def next_weather(self, reverse=False):
        self._weather_index += -1 if reverse else 1
        self._weather_index %= len(self._weather_presets)
        preset = self._weather_presets[self._weather_index]
        self.hud.notification('Weather: %s' % preset[1])
        self.vehicle.get_world().set_weather(preset[0])

    def tick(self, clock):
        self.hud.tick(self, clock)

    def render(self, display):
        self.camera_manager.render(display)
        self.hud.render(display)

    def destroy(self):
        actors = [
            self.camera_manager.sensor,
            self.collision_sensor.sensor,
            self.vehicle]
        for actor in actors:
            if actor is not None:
                actor.destroy()

# ==============================================================================
# -- HUD -----------------------------------------------------------------
# ==============================================================================

class HUD(object):
    def __init__(self, width, height):
        self.dim = (width, height)
        font = pygame.font.Font(pygame.font.get_default_font(), 20)
        fonts = [x for x in pygame.font.get_fonts() if 'mono' in x]
        default_font = 'ubuntumono'
        mono = default_font if default_font in fonts else fonts[0]
        mono = pygame.font.match_font(mono)
        self._font_mono = pygame.font.Font(mono, 14)
        self._notifications = FadingText(font, (width, 40), (0, height - 40))
        self.help = HelpText(pygame.font.Font(mono, 24), width, height)
        self.server_fps = 0
        self.frame_number = 0
        self.simulation_time = 0
        self._show_info = True
        self._info_text = []
        self._server_clock = pygame.time.Clock()

    def on_world_tick(self, timestamp):
        self._server_clock.tick()
        self.server_fps = self._server_clock.get_fps()
        self.frame_number = timestamp.frame_count
        self.simulation_time = timestamp.elapsed_seconds

    def tick(self, world, clock):
        t = world.vehicle.get_transform()
        v = world.vehicle.get_velocity()
        c = world.vehicle.get_control()
        heading = 'N' if abs(t.rotation.yaw) < 89.5 else ''
        heading += 'S' if abs(t.rotation.yaw) > 90.5 else ''
        heading += 'E' if 179.5 > t.rotation.yaw > 0.5 else ''
        heading += 'W' if -0.5 > t.rotation.yaw > -179.5 else ''
        colhist = world.collision_sensor.get_collision_history()
        collision = [colhist[x + self.frame_number - 200] for x in range(0, 200)]
        max_col = max(1.0, max(collision))
        collision = [x / max_col for x in collision]
        vehicles = world.world.get_actors().filter('vehicle.*')
        self._info_text = [
            'Vehicle: % 20s' % get_actor_display_name(world.vehicle, truncate=20),
            'Simulation time: % 12s' % datetime.timedelta(seconds=int(self.simulation_time)),
            '',
            'Speed:   % 15.0f km/h' % (3.6 * math.sqrt(v.x**2 + v.y**2 + v.z**2)),
            '',
            'Throttle: %16.2f' % c.throttle,
            ('Throttle:', c.throttle, 0.0, 1.0),
            'Steer: %19.2f' % c.steer,
            ('Steer:', c.steer, -1.0, 1.0),
            '',
            'Collision:',
            collision,
            '',
            'Number of vehicles: % 8d' % len(vehicles)
        ]
        if len(vehicles) > 1:
            self._info_text += ['Nearby vehicles:']
            distance = lambda l: math.sqrt((l.x - t.location.x)**2 + (l.y - t.location.y)**2 + (l.z - t.location.z)**2)
            vehicles = [(distance(x.get_location()), x) for x in vehicles if x.id != world.vehicle.id]
            for d, vehicle in sorted(vehicles):
                if d > 200.0:
                    break
                vehicle_type = get_actor_display_name(vehicle, truncate=22)
                self._info_text.append('% 4dm %s' % (d, vehicle_type))

        self._notifications.tick(world, clock)

    def toggle_info(self):
        self._show_info = not self._show_info

    def notification(self, text, seconds=2.0):
        self._notifications.set_text(text, seconds=seconds)

    def error(self, text):
        self._notifications.set_text('Error: %s' % text, (255, 0, 0))

    def render(self, display):
        if self._show_info:
            info_surface = pygame.Surface((220, self.dim[1]))
            info_surface.set_alpha(100)
            display.blit(info_surface, (0, 0))
            v_offset = 4
            bar_h_offset = 100
            bar_width = 106
            for item in self._info_text:
                if v_offset + 18 > self.dim[1]:
                    break
                if isinstance(item, list):
                    if len(item) > 1:
                        points = [(x + 8, v_offset + 8 + (1.0 - y) * 30) for x, y in enumerate(item)]
                        pygame.draw.lines(display, (255, 136, 0), False, points, 2)
                    item = None
                    v_offset += 18
                elif isinstance(item, tuple):
                    if isinstance(item[1], bool):
                        rect = pygame.Rect((bar_h_offset, v_offset + 8), (6, 6))
                        pygame.draw.rect(display, (255, 255, 255), rect, 0 if item[1] else 1)
                    else:
                        rect_border = pygame.Rect((bar_h_offset, v_offset + 8), (bar_width, 6))
                        pygame.draw.rect(display, (255, 255, 255), rect_border, 1)
                        f = (item[1] - item[2]) / (item[3] - item[2])
                        if item[2] < 0.0:
                            rect = pygame.Rect((bar_h_offset + f * (bar_width - 6), v_offset + 8), (6, 6))
                        else:
                            rect = pygame.Rect((bar_h_offset, v_offset + 8), (f * bar_width, 6))
                        pygame.draw.rect(display, (255, 255, 255), rect)
                    item = item[0]
                if item:  # At this point has to be a str.
                    surface = self._font_mono.render(item, True, (255, 255, 255))
                    display.blit(surface, (8, v_offset))
                v_offset += 18
        self._notifications.render(display)
        self.help.render(display)

# ==============================================================================
# -- FadingText ----------------------------------------------------------------
# ==============================================================================

class FadingText(object):
    def __init__(self, font, dim, pos):
        self.font = font
        self.dim = dim
        self.pos = pos
        self.seconds_left = 0
        self.surface = pygame.Surface(self.dim)

    def set_text(self, text, color=(255, 255, 255), seconds=2.0):
        text_texture = self.font.render(text, True, color)
        self.surface = pygame.Surface(self.dim)
        self.seconds_left = seconds
        self.surface.fill((0, 0, 0, 0))
        self.surface.blit(text_texture, (10, 11))

    def tick(self, _, clock):
        delta_seconds = 1e-3 * clock.get_time()
        self.seconds_left = max(0.0, self.seconds_left - delta_seconds)
        self.surface.set_alpha(500.0 * self.seconds_left)

    def render(self, display):
        display.blit(self.surface, self.pos)

# ==============================================================================
# -- HelpText ------------------------------------------------------------------
# ==============================================================================

class HelpText(object):
    def __init__(self, font, width, height):
        lines = __doc__.split('\n')
        self.font = font
        self.dim = (680, len(lines) * 22 + 12)
        self.pos = (0.5 * width - 0.5 * self.dim[0], 0.5 * height - 0.5 * self.dim[1])
        self.seconds_left = 0
        self.surface = pygame.Surface(self.dim)
        self.surface.fill((0, 0, 0, 0))
        for n, line in enumerate(lines):
            text_texture = self.font.render(line, True, (255, 255, 255))
            self.surface.blit(text_texture, (22, n * 22))
            self._render = False
        self.surface.set_alpha(220)

    def toggle(self):
        self._render = not self._render

    def render(self, display):
        if self._render:
            display.blit(self.surface, self.pos)

# ==============================================================================
# -- CollisionSensor -----------------------------------------------------------
# ==============================================================================

class CollisionSensor(object):
    def __init__(self, parent_actor, hud):
        self.sensor = None
        self._history = []
        self._parent = parent_actor
        self._hud = hud
        world = self._parent.get_world()
        bp = world.get_blueprint_library().find('sensor.other.collision')
        self.sensor = world.spawn_actor(bp, carla.Transform(), attach_to=self._parent)
        # We need to pass the lambda a weak reference to self to avoid circular
        # reference.
        weak_self = weakref.ref(self)
        self.sensor.listen(lambda event: CollisionSensor._on_collision(weak_self, event))

    def get_collision_history(self):
        history = collections.defaultdict(int)
        for frame, intensity in self._history:
            history[frame] += intensity
        return history

    @staticmethod
    def _on_collision(weak_self, event):
        self = weak_self()
        if not self:
            return
        actor_type = get_actor_display_name(event.other_actor)
        self._hud.notification('Collision with %r, id = %d' % (actor_type, event.other_actor.id))
        impulse = event.normal_impulse
        intensity = math.sqrt(impulse.x ** 2 + impulse.y ** 2 + impulse.z ** 2)
        self._history.append((event.frame_number, intensity))
        if len(self._history) > 4000:
            self._history.pop(0)

# ==============================================================================
# -- CameraManager -------------------------------------------------------------
# ==============================================================================

class CameraManager(object):
    def __init__(self, parent_actor, hud):
        self.sensor = None
        self._surface = None
        self._second_surface = None
        self._parent = parent_actor
        self._hud = hud
        self._recording = False
        self._camera_transforms = [
            carla.Transform(carla.Location(x=2.5, z=0.8)),
            carla.Transform(carla.Location(x=-10, z=4))]
        self._transform_index = 1
        self._sensors = [
            ['sensor.camera.rgb', cc.Raw, 'Camera RGB'],
            ['sensor.camera.depth', cc.Raw, 'Camera Depth (Raw)'],
            ['sensor.camera.depth', cc.Depth, 'Camera Depth (Gray Scale)'],
            ['sensor.camera.depth', cc.LogarithmicDepth, 'Camera Depth (Logarithmic Gray Scale)'],
            ['sensor.camera.semantic_segmentation', cc.Raw, 'Camera Semantic Segmentation (Raw)'],
            ['sensor.camera.semantic_segmentation', cc.CityScapesPalette,
             'Camera Semantic Segmentation (CityScapes Palette)'],
            ['sensor.lidar.ray_cast', None, 'Lidar (Ray-Cast)']]
        world = self._parent.get_world()
        bp_library = world.get_blueprint_library()
        for item in self._sensors:
            bp = bp_library.find(item[0])
            if item[0].startswith('sensor.camera'):
                bp.set_attribute('image_size_x', str(hud.dim[0]))
                bp.set_attribute('image_size_y', str(hud.dim[1]))
            item.append(bp)
        self._index = None
        self.image = None
        self.image2 = None

    def toggle_camera(self):
        self._transform_index = (self._transform_index + 1) % len(self._camera_transforms)
        self.sensor.set_transform(self._camera_transforms[self._transform_index])

    def set_sensor(self, index, notify=True):
        index = index % len(self._sensors)
        needs_respawn = True if self._index is None \
            else self._sensors[index][0] != self._sensors[self._index][0]
        if needs_respawn:
            if self.sensor is not None:
                self.sensor.destroy()
                self._surface = None
            self.sensor = self._parent.get_world().spawn_actor(
                self._sensors[index][-1],
                self._camera_transforms[self._transform_index],
                attach_to=self._parent)
            # We need to pass the lambda a weak reference to self to avoid
            # circular reference.
            weak_self = weakref.ref(self)
            self.sensor.listen(lambda image: CameraManager._parse_image(weak_self, image))
        if notify:
            self._hud.notification(self._sensors[index][2])
        self._index = 2

    def set_sensor_2(self, index, notify=True):
        index = index % len(self._sensors)
        needs_respawn = True if self._index is None \
            else self._sensors[index][0] != self._sensors[self._index][0]
        self.sensor = self._parent.get_world().spawn_actor(
            self._sensors[index][-1],
            self._camera_transforms[1],
            attach_to=self._parent)
            # We need to pass the lambda a weak reference to self to avoid
            # circular reference.
        weak_self = weakref.ref(self)
        self.sensor.listen(lambda image: CameraManager._parse_image_2(weak_self, image))
        self._index = 3

    def next_sensor(self):
        self.set_sensor(self._index + 1)

    def toggle_recording(self):
        self._recording = not self._recording
        self._hud.notification('Recording %s' % ('On' if self._recording else 'Off'))

    def render(self, display):
        if self._surface is not None:
            display.blit(self._surface, (0, 0))
            display.blit(self._second_surface, (0, self._hud.dim[1]))

    @staticmethod
    def _parse_image(weak_self, data):
        self = weak_self()
        if not self:
            return

        img = np.reshape(np.copy(data.raw_data), (data.height, data.width, 4))
        img = img[:,:,:3]
        img = img[:, :, ::-1]

        self.image = img

    @staticmethod
    def _parse_image_2(weak_self, data):
        self = weak_self()
        if not self:
            return

        img = np.reshape(np.copy(data.raw_data), (data.height, data.width, 4))
        img = img[:,:,:3]
        img = img[:, :, ::-1]

        self.image2 = img
        j = Image.fromarray(img)
        j.save("output2.png")

def atoi(text):
    return int(text) if text.isdigit() else text

def natural_keys(text):
    '''
    alist.sort(key=natural_keys) sorts in human order
    http://nedbatchelder.com/blog/200712/human_sorting.html
    (See Toothy's implementation in the comments)
    '''
    return [ atoi(c) for c in re.split(r'(\d+)', text) ]

def append_list_as_row(file_name, list_of_elem):
    # Open file in append mode
    with open(file_name, 'a+', newline='') as write_obj:
        # Create a writer object from csv module
        csv_writer = writer(write_obj)
        # Add contents of list as last row in the csv file
        csv_writer.writerow(list_of_elem)

def line_prepender(filename, line):
    """with open(filename, "r") as infile:
        reader = list(csv.reader(infile))
        reader.insert(0, line)

    with open(filename, "w") as outfile:
        writer = csv.writer(outfile)
        for line in reader:
            writer.writerow(line)"""
    with open(filename, 'a+', newline='') as write_obj:
        # Create a writer object from csv module
        csv_writer = writer(write_obj)
        # Add contents of list as last row in the csv file
        csv_writer.writerow(line)

def save_instance_to_dataset(world, image, image_counter, first_image_taken, prev_image_num):
    file_path = "_%d/" % (SPAWNPOINT)
    if (first_image_taken == 0) and (not os.path.exists('./_%d/data.csv' % SPAWNPOINT)):
        line_prepender(file_path + 'data.csv', ['image', 'throttle', 'steer'])
        first_image_taken = 1
    if (first_image_taken == 0) and (os.path.exists('./_%d/data.csv' % SPAWNPOINT)):
        files_list = glob.glob(file_path + '*.png')
        files_list.sort(key=natural_keys)
        first_image_taken = 1
        prev_image_num = int(files_list[-1].split('/')[1].split('.')[0])

    file_name = "%d.png" % (image_counter + prev_image_num)
    j = Image.fromarray(image)
    j.save(file_path + file_name, compress_level=1)

    row_contents = [file_name, world.vehicle.get_control().throttle, world.vehicle.get_control().steer]
    # Append a list as new line to an old csv file
    append_list_as_row(file_path + 'data.csv', row_contents)

    return first_image_taken, prev_image_num

# ==============================================================================
# -- game_loop() ---------------------------------------------------------
# ==============================================================================

def game_loop(args):
    pygame.init()
    pygame.font.init()
    world = None

    print(SPAWNPOINT)
    try:
        client = carla.Client(args.host, args.port)
        client.set_timeout(4.0)

        display = pygame.display.set_mode(
            (args.width, args.height + args.height),
            pygame.HWSURFACE | pygame.DOUBLEBUF)

        hud = HUD(args.width, args.height)
        world = World(client.load_world('Town01_Opt'), hud, args)

        # weather = world.world.get_weather()
        # weather.sun_altitude_angle = -30
        # weather.fog_density = 65
        # weather.fog_distance = 10
        # weather.wetness = 100
        # world.world.set_weather(weather)
        # vehicle_light_state = carla.VehicleLightState(carla.VehicleLightState.Position | carla.VehicleLightState.LowBeam)
        # world.vehicle.set_light_state(vehicle_light_state)

        if args.agent == "rl":
            agent = RLAgent(world.vehicle)
            agent.model.summary()
            agent.model.predict(np.ones((1, 66, 200, 3)))
        elif args.agent == "fl":
            agent = FLAgent(world.vehicle)
        elif args.agent == "fr":
            agent = Agent(world.vehicle)
        else:
            agent = None
            world.vehicle.set_autopilot(True)

            traffic_manager = client.get_trafficmanager()
            route = ["Straight"]*1000
            traffic_manager.set_route(world.vehicle, route)
            traffic_manager.ignore_lights_percentage(world.vehicle, 100)
            traffic_manager.ignore_signs_percentage(world.vehicle, 100)
            traffic_manager.keep_right_rule_percentage(world.vehicle, 0)

        if args.rec != "":
            if not os.path.exists('./_%d/' % SPAWNPOINT):
                os.makedirs('./_%d/' % SPAWNPOINT)

        birdview_producer = BirdViewProducer(
            client,
            target_size=PixelDimensions(width=150, height=336),
            pixels_per_meter=4,
            crop_type=BirdViewCropType.FRONT_AREA_ONLY
        )

        clock = pygame.time.Clock()
        fps = deque(maxlen=60)
        # media = deque(maxlen=1000)
        image_counter = 0
        first_image_taken = 0
        recording = False
        noise = False
        noise_value = 0
        prev_image_number = 0
        while True:
            # as soon as the server is ready continue!
            if not world.world.wait_for_tick(10.0):
                continue
            step_start = time.time()
            if (type(world.camera_manager.image).__module__ == np.__name__):
                # If we use autopilot of carla and no agent to control the car
                if agent == None:
                    world.camera_manager._surface = pygame.surfarray.make_surface(world.camera_manager.image.swapaxes(0, 1))

                    """#Get the traffic light affecting a vehicle
                    if world.vehicle.is_at_traffic_light():
                        traffic_light = world.vehicle.get_traffic_light()
                        if traffic_light.get_state() == carla.TrafficLightState.Red or traffic_light.get_state() == carla.TrafficLightState.Yellow:
                            traffic_light.set_state(carla.TrafficLightState.Green)
                            # traffic_light.set_green_time(4.0)"""

                    image_counter = image_counter + 1
                # The agent selected will be used to control the car
                else:
                    steer, throttle, image = agent.run_step(world.camera_manager.image)
                    world.camera_manager._surface = pygame.surfarray.make_surface(image.swapaxes(0, 1))
                    world.camera_manager._second_surface = pygame.surfarray.make_surface(world.camera_manager.image2.swapaxes(0, 1))
                    world.vehicle.apply_control(carla.VehicleControl(throttle=float(throttle), steer=float(steer)))

                    image_counter = image_counter + 1

                """if (noise == False) and (image_counter % 300 == 0):
                    noise = True
                elif (noise == True) and (image_counter % 17 == 0):
                    noise = False
                    noise_value = 0

                # Apply noise if activated
                control = world.vehicle.get_control()
                location = world.vehicle.get_location()
                if noise:
                    noise_value = random.uniform(-0.05, 0.05)
                    control.steer += noise_value
                    world.vehicle.apply_control(control)"""

                for event in pygame.event.get():
                    if event.type == pygame.QUIT:
                        quit()

                    """elif event.type == pygame.KEYDOWN:
                        if event.key == K_r and recording == False:
                            recording = True
                        elif event.key == K_r and recording == True:
                            recording = False"""

                """if (139 < location.x < 163) and (-194 < location.y < 193.5): # Map03 right turn
                    continue"""
                """if (134 < location.x < 170) and (-207.89 < location.y < -204): # Map03 left turn
                    continue
                else:"""
                """if abs(control.steer) >= 0.3: # 0.09
                    print(f"GIROOOOOOO - {image_counter}")
                    recording = True
                else:
                    recording = False"""

                # if (args.rec == True) and (image_counter % 10 == 0):
                if recording and args.rec == "rgb":
                    first_image_taken, prev_image_number = save_instance_to_dataset(world, world.camera_manager.image,
                    image_counter, first_image_taken, prev_image_number)

                    if len(glob.glob("_%d/" % (SPAWNPOINT) + '*.png')) > 5000:
                        break

                elif recording and args.rec == "bird":
                    birdview = birdview_producer.produce(agent_vehicle=world.vehicle)
                    rgb = BirdViewProducer.as_rgb(birdview)
                    first_image_taken = save_instance_to_dataset(world, rgb, image_counter, first_image_taken)

            frame_time = time.time() - step_start
            fps.append(frame_time)

            #media.append(hud.server_fps)
            # if len(media) == 1000:
            #     print(f"MEDIA DE FPS DE SERVIDOR -----------------------------------------> {sum(media) / len(media)}")

            world.tick(clock)
            if args.debug == True:
                print(f"Server: {hud.server_fps}")
                print(f"Client: {len(fps)/sum(fps):>4.1f} FPS | {frame_time*1000} ms")

            world.render(display)
            pygame.display.flip()

    finally:
        if world is not None:
            world.destroy()

        pygame.quit()

# ==============================================================================
# -- main() --------------------------------------------------------------
# ==============================================================================

def main():
    argparser = argparse.ArgumentParser(
        description='CARLA Manual Control Client')
    argparser.add_argument(
        '-v', '--verbose',
        action='store_true',
        dest='debug',
        help='print debug information')
    argparser.add_argument(
        '--host',
        metavar='H',
        default='127.0.0.1',
        help='IP of the host server (default: 127.0.0.1)')
    argparser.add_argument(
        '-p', '--port',
        metavar='P',
        default=2000,
        type=int,
        help='TCP port to listen to (default: 2000)')
    argparser.add_argument(
        '--res',
        metavar='WIDTHxHEIGHT',
        default='640x480',
        help='window resolution (default: 640x480)')

    argparser.add_argument("-r", "--rec",
                           choices=["rgb", "bird"],
                           help="select which view to record",
                           default="rgb")

    argparser.add_argument("-a", "--agent", type=str,
                           choices=["rl", "basic", "fl", "fr"],
                           help="select which agent to run",
                           default="basic")
    args = argparser.parse_args()

    args.width, args.height = [int(x) for x in args.res.split('x')]

    log_level = logging.DEBUG if args.debug else logging.INFO
    logging.basicConfig(format='%(levelname)s: %(message)s', level=log_level)

    logging.info('listening to server %s:%s', args.host, args.port)

    print(__doc__)

    try:

        game_loop(args)

    except KeyboardInterrupt:
        print('\nCancelled by user. Bye!')
    except Exception as error:
        logging.exception(error)

if __name__ == '__main__':
    """start = time.time()
    spawnpoints = [8, 9]
    for spawn in spawnpoints:
        SPAWNPOINT = spawn
        main()
    end = time.time() - start
    print(f"THE END: time -> {end}")"""

    main()


"""Add host_switched field to game

Revision ID: 394a2bf76d9a
Revises: 29d69456e7ef
Create Date: 2021-01-09 12:11:21.868366

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = '394a2bf76d9a'
down_revision = '29d69456e7ef'
branch_labels = None
depends_on = None

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('game', sa.Column('host_switched', sa.Boolean(), default=False, nullable=False))
    op.create_unique_constraint(None, 'gamelog', ['id'])
    # ### end Alembic commands ###

def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint(None, 'gamelog', type_='unique')
    op.drop_column('game', 'host_switched')
    # ### end Alembic commands ###


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

import os

from tensorboard.backend.event_processing import event_accumulator
import pandas as pd
import argparse

def _to_csv(export_path, save_data, columns, index=False):
    df = pd.DataFrame(data=save_data, columns=columns)
    df.to_csv(export_path, index=index)

def single_event_scalars_export_tool(in_path, ex_path):
    event_data = event_accumulator.EventAccumulator(in_path)
    event_data.Reload()

    keys = event_data.scalars.Keys()
    print("keys list: {}".format(keys))

    export_items = ["step", "wall_time", "value"]
    for key in keys:
        print("process key: {}".format(key))
        export_names = ["{}_{}".format(key, item) for item in export_items]
        save_data = list()
        for e in event_data.Scalars(key):
            temp_data = [e.step, e.wall_time, e.value]
            save_data.append(temp_data)

        save_name = "{}.csv".format(key.replace("/", "_"))
        save_path = os.path.join(ex_path, save_name)
        _to_csv(save_path, save_data, export_names)
        print("export data of {} done...".format(key))

def multi_event_scalars_export_tool(in_path, ex_path):
    root, dirs, _ = next(os.walk(in_path))
    for _dir in dirs:
        print("process summary: {}".format(_dir))
        event_in_path = os.path.join(root, _dir, "tensorboard_log")
        event_ex_path = os.path.join(ex_path, _dir)
        if not os.path.exists(event_ex_path):
            os.makedirs(event_ex_path)
        single_event_scalars_export_tool(event_in_path, event_ex_path)

def single_event_images_export_tool():
    pass

def multi_event_images_export_tool():
    pass

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Export tensorboard data")
    parser.add_argument("--in_path", type=str, help="tensorboard file location")
    parser.add_argument("--ex_path", type=str, default="./", help="export path")
    parser.add_argument("--state", type=str, default="single", help="single summary or multi summaries")
    parser.add_argument("--ex_type", type=str, default="scales", help="export data type")
    args = parser.parse_args()

    if args.state == "single":
        if args.ex_type == "scales":
            single_event_scalars_export_tool(args.in_path, args.ex_path)
        else:
            raise KeyError
    elif args.state == "multi":
        if args.ex_type == "scales":
            multi_event_scalars_export_tool(args.in_path, args.ex_path)
        else:
            raise KeyError


import io
import subprocess
import nltk
import sys
from finnpos.lib.bin import finnpos_ratna_feats as ratna

FINNPOS_FOLDER          = './finnpos/lib'
BIN_FOLDER              = '{0}/bin'.format(FINNPOS_FOLDER)
MODEL_FOLDER            = '{0}/share/finnpos/ftb_omorfi_model'.format(FINNPOS_FOLDER)
FREQ_WORDS              = '{0}/freq_words'.format(MODEL_FOLDER)
OMORFI_MODEL            = '{0}/ftb.omorfi.model'.format(MODEL_FOLDER)

FINNPOS_RATNA_FEATS     = '{0}/finnpos-ratna-feats.py'.format(BIN_FOLDER)
FINNPOS_LABEL           = '{0}/finnpos-label'.format(BIN_FOLDER)

POSTAG                  = '{0}/postag.py'.format(BIN_FOLDER)

def create_command():
    return 'echo "Samoin" | bash -c "python3 {ratna} {words}"'.format(# | "{label} {model}"'.format(#) | bash -c "{label} {model}'.format(
        ratna=FINNPOS_RATNA_FEATS,
        words=FREQ_WORDS
        #label=FINNPOS_LABEL,
        #model=OMORFI_MODEL
    )
    '''return 'bash -c "py {0}/test.py"'.format(
        FINNPOS_FOLDER
    )
    return 'echo Sanoin | py {0} {1} | {2} {3} | py {4}"'.format(
        FINNPOS_RATNA_FEATS,
        FREQ_WORDS,
        FINNPOS_LABEL,
        OMORFI_MODEL,
        POSTAG
    )'''

def bash_command(cmd, stdin=None, stdout=None):
    '''if stdin is not None and stdout is not None:
        return subprocess.Popen(['bash', '-c', cmd], stdin=stdin, stdout=stdout)
    elif stdin is not None:
        return subprocess.Popen(['bash', '-c', cmd], stdin=stdin)
    elif stdout is not None:
        return subprocess.Popen(['bash', '-c', cmd], stdout=stdout)
    else:'''
    return subprocess.Popen(['bash', '-c', 'bash -c "{cmd}"'.format(cmd=cmd)])#"python3 {0} {1}"'.format(FINNPOS_RATNA_FEATS, FREQ_WORDS)], shell=False)

def postag(token):

    print('Create ratnaProc')
    ratnaProc       = subprocess.Popen([
            'bash',
            '-c',
            'py {0} {1}'.format( FINNPOS_RATNA_FEATS, FREQ_WORDS ),
        ],
        stdin=subprocess.PIPE,
        stderr=sys.stderr,
        stdout=subprocess.PIPE
    )

    #ratnaOut        = ratnaProc.communicate()[0].decode('latin1')
    #print('ratnaOut', ratnaOut)

    print('Create labelProc')
    labelProc       = subprocess.check_output([
            'bash',
            '-c',
            'bash -c "{0} {1}"'.format( FINNPOS_LABEL, OMORFI_MODEL ),
        ],
        stdin=ratnaProc.stdout,
        stderr=sys.stderr
    )
    print('Write to the ratna')
    ratnaProc.stdin.write( bytes(token, 'latin1') )
    print('Wait for the ratna')
    ratnaProc.wait()
    print('Wait for the ratna done')
    #labelProc.stdin.write( bytes(ratnaOut, 'latin1'))
    #out = labelProc.communicate()

    print('Closing processes')
    #labelProc.stdin.close()

    #print('out', out)

    #print('ratnaComm', ratnaProc.communicate()[0].decode('latin1'))

    return

def test(tokens):

    with open(FREQ_WORDS, encoding='utf-8') as file:
        freq_words      = file.read().split('\n')

        stdin       = io.StringIO('\n'.join(tokens))

        print('Loading labelProc')
        labelProc   = subprocess.check_output([
            'bash',
            '-c',
            'bash -c "{0} {1}"'.format( FINNPOS_LABEL, OMORFI_MODEL )
        ], shell=True, stderr=sys.stderr)
        print('labelProc', labelProc)

        print('Pass tokens to the ratna')
        ratna.main('finnpos', 'STDIN', stdin, 'STDOUT', labelProc, sys.stderr, freq_words)

        print('Wait for labelProc')
        #labelProc.wait()
        #finnpos.finnpos_ratna_feats.main('finnpos', 'STDIN', sys.stdin, 'STDOUT', sys.stdout, sys.stderr, freq_words)
        #finnpos.finnpos_ratna_feats.main('finnpos-ratna-feats', "STDIN", stdin, "STDOUT", stdout, stderr, freq_words)

    '''

    if len(argv) != 2:
        stderr.write("cat data | %s freq_word_list\n" % argv[0])
        exit(1)

    freq_words = set(open(argv[1]).read().split('\n'))

    exit(main(argv[0], "STDIN", stdin, "STDOUT", stdout, stderr, freq_words))
    def main(pname, iname, ifile, oname, ofile, olog, freq_words):
    '''

def lemmatize(val):

    print('Lemmatizing data')

    #test(val)

    tokens              = nltk.word_tokenize(val)

    test(tokens)

    #print(postag(tokens[0]))
    #cmd                 = create_command()
    #print('Run', cmd)
    #bash_command(cmd)

    #print(postag(tokens[0]))
    #tags                = [postag(token) for token in tokens]

    #print(tags)

    #tokenized_input     = '\n'.join(tokens)

    #cmd                 = create_command('Sanoin')

    #print(cmd)

    #ps1                 = bash_command(cmd, stdin=subprocess.PIPE)
    #ps1.stdin.buffer.write('testaus')
    #print(ps1.communicate()[0])

    #ps1.stdin.write("tokens")

    #[ps1.stdin.writelines(token) for token in tokens]
    #ps1.stdin.write()

    #bash_command('a="Apples and oranges" && echo "${a/oranges/grapes}"')
    #subprocess.Popen('bash -c "ls"')

    #cmd         = 'bash -c "echo Sanoin | python3 ./lib/bin/finn-pos-ratna-feats.py"'
    #cmd          = 'ls ./finnpos/lib/bin/'
    #cmd             = 'bash -c "echo Sanoin | python3 ./lib/bin/finn-pos-ratna-feats.py ./lib/share/finnpos/ftb_omorfi_model/freq_words"'

    #ps = subprocess.Popen(cmd, shell=True)#, '"python3 ./finnpos-ratna-feats.py"'), stdout=subprocess.PIPE)
    #output = subprocess.check_output(('grep', 'process_name'), stdin=ps.stdout)
    #ps.wait()

    #subprocess.Popen(['bash', '-c', 'python3', 'finnpos-ratna-feats.py'], stdin=tokenized_input, stdout=subprocess.PIPE)

    #bash -c "cat tokenized_input | python3 finnpos-ratna-feats.py ../share/finnpos/ftb_omorfi_model/freq_words | ./finnpos-label ../share/finnpos/ftb_omorfi_model/ftb.omorfi.model | python3 postag.py > tagged_input"

    #print(tokenized_input)


# Run from prefect directory (after terraform vars gen) using
# prefect run --name sim-prep --param name=florance --param year=2018

# For logging, use `logger = prefect.context.get("logger")` within tasks
import argparse
import warnings
import pathlib

from prefect import case
from prefect.utilities import graphql
from prefect.client import Client
from prefect.tasks.control_flow import merge

from conf import PREFECT_PROJECT_NAME, INIT_FINI_LOCK
from tasks.params import (
        param_storm_name, param_storm_year,
        param_use_rdhpcs, param_use_rdhpcs_post, param_run_id,
        param_use_parametric_wind, param_subset_mesh, param_ensemble,
        param_mesh_hmax,
        param_mesh_hmin_low, param_mesh_rate_low,
        param_mesh_trans_elev,
        param_mesh_hmin_high, param_mesh_rate_high,
        param_ensemble_n_perturb, param_hr_prelandfall,
        param_ensemble_sample_rule,
        param_past_forecast,
        param_wind_coupling,
        )
from tasks.data import (
        task_copy_s3_data,
        task_init_run,
        task_final_results_to_s3,
        task_cleanup_run,
        task_cache_to_s3,
        task_cleanup_efs)
from tasks.utils import (
        task_check_param_true,
        task_bundle_params,
        task_get_flow_run_id,
        task_get_run_tag,
        FLock)
from flows.jobs.ecs import (
        make_flow_generic_ecs_task,
        make_flow_solve_ecs_task
        )
from flows.jobs.pw import(
        make_flow_mesh_rdhpcs_pw_task,
        make_flow_mesh_rdhpcs,
        make_flow_solve_rdhpcs_pw_task,
        make_flow_solve_rdhpcs)
from flows.utils import LocalAWSFlow, flow_dependency

# TODO: Later add build image and push to ECS logic into Prefect workflow

# TODO: Use subprocess.run to switch backend here
# TODO: Create user config file to be session based? https://docs-v1.prefect.io/core/concepts/configuration.html#environment-variables

def _check_project():
    client = Client()
    print(f"Connecting to {client.api_server}...")

    qry = graphql.parse_graphql({'query': {'project': ['name']}})
    rsp = client.graphql(qry)

    prj_names = [i['name'] for i in rsp['data']['project']]
    if PREFECT_PROJECT_NAME in prj_names:
        print(f"Project {PREFECT_PROJECT_NAME} found on {client.api_server}!")
        return

    print(f"Creating project {PREFECT_PROJECT_NAME} on {client.api_server}...")
    client.create_project(project_name=PREFECT_PROJECT_NAME)
    print("Done!")

def _make_workflow():
    # Create flow objects
    flow_sim_prep_info_aws = make_flow_generic_ecs_task("sim-prep-info-aws")
    flow_sim_prep_mesh_aws = make_flow_generic_ecs_task("sim-prep-mesh-aws")
    flow_sim_prep_setup_aws = make_flow_generic_ecs_task("sim-prep-setup-aws")
    flow_mesh_rdhpcs_pw_task = make_flow_mesh_rdhpcs_pw_task()
    flow_mesh_rdhpcs = make_flow_mesh_rdhpcs(flow_mesh_rdhpcs_pw_task)
    flow_schism_single_run_aws = make_flow_generic_ecs_task("schism-run-aws-single")
    flow_schism_ensemble_run_aws = make_flow_solve_ecs_task(flow_schism_single_run_aws)
    flow_solve_rdhpcs_pw_task = make_flow_solve_rdhpcs_pw_task()
    flow_solve_rdhpcs = make_flow_solve_rdhpcs(flow_solve_rdhpcs_pw_task)
    flow_sta_html_aws = make_flow_generic_ecs_task("viz-sta-html-aws")
    flow_cmb_ensemble_aws = make_flow_generic_ecs_task(
        "viz-cmb-ensemble-aws"
    )
    flow_ana_ensemble_aws = make_flow_generic_ecs_task(
        "viz-ana-ensemble-aws"
    )

    with LocalAWSFlow("end-to-end") as flow_main:

        result_flow_run_id = task_get_flow_run_id()

        result_run_tag = task_get_run_tag(
            param_storm_name, param_storm_year, result_flow_run_id)

        result_is_rdhpcs_on = task_check_param_true(param_use_rdhpcs)
        result_is_ensemble_on = task_check_param_true(param_ensemble)
        result_is_rdhpcspost_on = task_check_param_true(param_use_rdhpcs_post)

        with FLock(INIT_FINI_LOCK, task_args={'name': 'Sync init'}):
            result_copy_task = task_copy_s3_data()
            result_init_run = task_init_run(
                    result_run_tag, upstream_tasks=[result_copy_task])

        result_bundle_params_1 = task_bundle_params(
                name=param_storm_name,
                year=param_storm_year,
                rdhpcs=param_use_rdhpcs,
                rdhpcs_post=param_use_rdhpcs_post,
                run_id=result_flow_run_id,
                parametric_wind=param_use_parametric_wind,
                ensemble=param_ensemble,
                hr_before_landfall=param_hr_prelandfall,
                past_forecast=param_past_forecast,
                couple_wind=param_wind_coupling,
                )

        result_bundle_params_2 = task_bundle_params(
                name=param_storm_name,
                year=param_storm_year,
                rdhpcs=param_use_rdhpcs,
                run_id=result_flow_run_id,
                subset_mesh=param_subset_mesh,
                mesh_hmax=param_mesh_hmax,
                mesh_hmin_low=param_mesh_hmin_low,
                mesh_rate_low=param_mesh_rate_low,
                mesh_cutoff=param_mesh_trans_elev,
                mesh_hmin_high=param_mesh_hmin_high,
                mesh_rate_high=param_mesh_rate_high
                )

        result_bundle_params_3 = task_bundle_params(
                name=param_storm_name,
                year=param_storm_year,
                run_id=result_flow_run_id,
                parametric_wind=param_use_parametric_wind,
                ensemble=param_ensemble,
                ensemble_num_perturbations=param_ensemble_n_perturb,
                hr_before_landfall=param_hr_prelandfall,
                couple_wind=param_wind_coupling,
                ensemble_sample_rule=param_ensemble_sample_rule,
        )

        after_sim_prep_info = flow_dependency(
                flow_name=flow_sim_prep_info_aws.name,
                upstream_tasks=[result_init_run],
                parameters=result_bundle_params_1)

        # TODO: Meshing based-on original track for now
        # TODO: If mesh each track: diff mesh

        with case(result_is_rdhpcs_on, True):
            after_sim_prep_mesh_b1 = flow_dependency(
                    flow_name=flow_mesh_rdhpcs.name,
                    upstream_tasks=[after_sim_prep_info],
                    parameters=result_bundle_params_2)
        with case(result_is_rdhpcs_on, False):
            after_sim_prep_mesh_b2 = flow_dependency(
                    flow_name=flow_sim_prep_mesh_aws.name,
                    upstream_tasks=[after_sim_prep_info],
                    parameters=result_bundle_params_2)
        after_sim_prep_mesh = merge(after_sim_prep_mesh_b1, after_sim_prep_mesh_b2)

        after_sim_prep_setup = flow_dependency(
                flow_name=flow_sim_prep_setup_aws.name,
                upstream_tasks=[after_sim_prep_mesh],
                parameters=result_bundle_params_3)

        with case(result_is_rdhpcs_on, True):
            after_run_schism_b1 = flow_dependency(
                    flow_name=flow_solve_rdhpcs.name,
                    upstream_tasks=[after_sim_prep_setup],
                    parameters=result_bundle_params_1)
        with case(result_is_rdhpcs_on, False):
            after_run_schism_b2 = flow_dependency(
                flow_name=flow_schism_ensemble_run_aws.name,
                upstream_tasks=[after_sim_prep_setup],
                parameters=result_bundle_params_1)
        after_run_schism = merge(after_run_schism_b1, after_run_schism_b2)

        with case(result_is_ensemble_on, True):
            with case(result_is_rdhpcspost_on, False):
                after_cmb_ensemble = flow_dependency(
                    flow_name=flow_cmb_ensemble_aws.name,
                    upstream_tasks=[after_run_schism],
                    parameters=result_bundle_params_1)
                after_ana_ensemble = flow_dependency(
                    flow_name=flow_ana_ensemble_aws.name,
                    upstream_tasks=[after_cmb_ensemble],
                    parameters=result_bundle_params_1)

            with case(result_is_rdhpcspost_on, True):
                # TODO:
                pass

        with case(result_is_ensemble_on, False):
            after_sta_html = flow_dependency(
                flow_name=flow_sta_html_aws.name,
                upstream_tasks=[after_run_schism],
                parameters=result_bundle_params_1)
        after_gen_viz = merge(after_ana_ensemble, after_sta_html)

        # TODO: Make this a separate flow?
        after_results_to_s3 = task_final_results_to_s3(
                param_storm_name, param_storm_year, result_run_tag,
                upstream_tasks=[after_gen_viz])

        after_cleanup_run = task_cleanup_run(
                result_run_tag, upstream_tasks=[after_results_to_s3])

        with FLock(INIT_FINI_LOCK, upstream_tasks=[after_cleanup_run], task_args={'name': 'Sync cleanup'}):
            after_cache_storage = task_cache_to_s3(
                    upstream_tasks=[after_cleanup_run])
            task_cleanup_efs(
                    result_run_tag,
                    upstream_tasks=[after_cache_storage])

    flow_main.set_reference_tasks([after_cleanup_run])

    all_flows = [
        flow_sim_prep_info_aws,
        flow_sim_prep_mesh_aws,
        flow_sim_prep_setup_aws,
        flow_mesh_rdhpcs_pw_task,
        flow_mesh_rdhpcs,
        flow_schism_single_run_aws,
        flow_schism_ensemble_run_aws,
        flow_solve_rdhpcs_pw_task,
        flow_solve_rdhpcs,
        flow_sta_html_aws,
        flow_cmb_ensemble_aws,
        flow_ana_ensemble_aws,
        flow_main
    ]

    return all_flows

def _regiser(flows):
    # Register unregistered flows
    for flow in flows:
        flow.register(project_name=PREFECT_PROJECT_NAME)

def _viz(flows, out_dir, flow_names):
    flow_dict = {f.name: f for f in flows}
    for flow_nm in flow_names:
        flow = flow_dict.get(flow_nm)
        if flow is None:
            warnings.warn(f'Flow with the name {flow_nm} NOT found!')
        flow.visualize(filename=out_dir/flow.name, format='dot')

def _list(flows):
    flow_names = [f.name for f in flows]
    print("\n".join(flow_names))

def _main(args):

    _check_project()
    all_flows = _make_workflow()
    if args.command in ["register", None]:
        _regiser(all_flows)

    elif args.command == "visualize":
        _viz(all_flows, args.out_dir, args.flowname)

    elif args.command == "list":
        _list(all_flows)

    else:
        raise ValueError("Invalid command!")

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command")

    reg_parser = subparsers.add_parser('register')
    viz_parser = subparsers.add_parser('visualize')
    list_parser = subparsers.add_parser('list')

    viz_parser.add_argument('flowname', nargs='+')
    viz_parser.add_argument(
        '--out-dir', '-d', type=pathlib.Path, default='.')

    _main(parser.parse_args())