"""Main script for the warp model."""
from shutil import rmtree, move
from os import listdir, path, mkdir
import cv2
import numpy as np
import tensorflow as tf
from tqdm import tqdm

from model import MyModel
from utils import dense_image_warp, combine_videos

TRAIN_EPOCHS = 100

IMAGE_SIZE = 1024
MORPH_DIRECTORY = "morph/morph.mp4"

tf.keras.mixed_precision.set_global_policy("mixed_float16")

@tf.function
def warp(origins, targets, preds_org, preds_trg):
    """
    The warp function takes in the original images, and the predicted offsets for each image.
    It then applies these offsets to create a warped version of each image.

    :param origins: Pass the original image to the warp function
    :param targets: Define the target image
    :param preds_org: Warp the origin image
    :param preds_trg: Warp the target image
    :return: The warped images
    """
    scale_org = tf.maximum(0.1, 1.0 + preds_org[..., :3] * MULT_SCALE)
    scale_trg = tf.maximum(0.1, 1.0 + preds_trg[..., :3] * MULT_SCALE)

    offset_org = preds_org[..., 3:6] * 2.0 * ADD_SCALE
    offset_trg = preds_trg[..., 3:6] * 2.0 * ADD_SCALE

    warp_org = preds_org[..., 6:8] * IMAGE_SIZE * WARP_SCALE
    warp_trg = preds_trg[..., 6:8] * IMAGE_SIZE * WARP_SCALE

    if ADD_FIRST:
        res_targets = dense_image_warp((origins + offset_org) * scale_org, warp_org)
        res_origins = dense_image_warp((targets + offset_trg) * scale_trg, warp_trg)
    else:
        res_targets = dense_image_warp(origins * scale_org + offset_org, warp_org)
        res_origins = dense_image_warp(targets * scale_trg + offset_trg, warp_trg)

    return res_targets, res_origins

def create_grid(scale):
    """
    The create_grid function creates a grid of coordinates that can be used to
    sample the output of a convolutional layer.

    :param scale: Determine the size of the grid
    :return: A grid of size (scale, scale) with values between - 1 and 1
    """
    grid = np.mgrid[0:scale, 0:scale] / (scale - 1) * 2 - 1
    grid = np.swapaxes(grid, 0, 2)
    grid = np.expand_dims(grid, axis=0)
    return grid

def produce_warp_maps(origins, targets, original_width, original_height):
    """
    The produce_warp_maps function takes two images, origins and targets, as input
    and produces a set of warp maps that can be used to transform the origins image to the target.

    :param original_width: Original width of the image
    :param original_height: Original height of the image
    :param origins: Store the original images
    :param targets: Warp the original image to the target image
    :return: The predicted maps
    """
    model = MyModel()

    loss_object = tf.keras.losses.MeanSquaredError()
    optimizer = tf.keras.optimizers.Adam(learning_rate=0.0002)

    train_loss = tf.keras.metrics.Mean(name="train_loss")
    flow_scale = IMAGE_SIZE * WARP_SCALE

    @tf.function
    def train_step(training_maps, training_origins, training_targets):
        """
        The train_step function takes in training data, and performs one step of model training.
        The loss function is also calculated and stored in a history object for later visualization.

        :param training_maps: Train the model
        :param training_origins: Calculate the loss of the origin images
        :param training_targets: Calculate the loss of the target images
        :return: The loss, so we can plot it
        """
        with tf.GradientTape() as tape:
            map_pred = model(training_maps)
            map_pred = tf.image.resize(map_pred, [IMAGE_SIZE, IMAGE_SIZE])
            res_targets_, res_origins_ = warp(
                training_origins, training_targets, map_pred[..., :8], map_pred[..., 8:]
            )
            res_map = dense_image_warp(training_maps, map_pred[:, :, :, 6:8] * flow_scale)
            res_map = dense_image_warp(res_map, map_pred[:, :, :, 14:16] * flow_scale)

            loss = (
                    loss_object(training_maps, res_map) * 1
                    + loss_object(res_targets_, training_targets) * 0.3
                    + loss_object(res_origins_, training_origins) * 0.3
            )

        gradients = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

        train_loss(loss)

    maps = create_grid(IMAGE_SIZE)
    maps = np.concatenate((maps, origins * 0.1, targets * 0.1), axis=-1).astype(np.float32)

    epoch = 0
    template = "Epoch {}, Loss: {}"

    training = tqdm(range(TRAIN_EPOCHS), desc=template.format(epoch, train_loss.result()))

    for iteration in training:
        epoch = iteration + 1

        training.set_description(template.format(epoch, train_loss.result()))
        training.refresh()

        train_step(maps, origins, targets)

        if (epoch < 100 and epoch % 10 == 0) or (epoch < 1000 and epoch % 100 == 0) or \
                (epoch % 1000 == 0):
            preds = model(maps, training=False)[:1]
            preds = tf.image.resize(preds, [IMAGE_SIZE, IMAGE_SIZE])

            res_targets, res_origins = warp(origins, targets, preds[..., :8], preds[..., 8:])
            np.save("preds.npy", preds.numpy())

            res_targets = tf.clip_by_value(res_targets, -1, 1)[0]
            res_origins = tf.clip_by_value(res_origins, -1, 1)[0]

            for res, prefix in zip([res_targets, res_origins], ["a_to_b", "b_to_a"]):
                res_img = ((res.numpy() + 1) * 127.5).astype(np.uint8)
                res_img = cv2.resize(res_img, (original_width, original_height),
                                     interpolation=cv2.INTER_AREA)
                cv2.imwrite(f"train/{prefix}_{epoch}.jpg", cv2.cvtColor(res_img, cv2.COLOR_RGB2BGR))

def use_warp_maps(origins, targets, original_width, original_height):
    """
    The use_warp_maps function takes in the original and target images,
    loads in the predicted warp maps, then outputs a video of this morphing process.

    :param original_width: Original width of the image
    :param original_height: Original height of the image
    :param origins: Get the original image
    :param targets: Store the target image
    :return: A video of the morph between two images
    """
    preds = np.load("preds.npy")

    res_img = np.zeros((IMAGE_SIZE * 2, IMAGE_SIZE * 3, 3), dtype=np.uint8)

    res_img[IMAGE_SIZE * 0: IMAGE_SIZE * 1, IMAGE_SIZE * 0: IMAGE_SIZE * 1] = \
        preds[0, :, :, 0:3]
    res_img[IMAGE_SIZE * 0: IMAGE_SIZE * 1, IMAGE_SIZE * 1: IMAGE_SIZE * 2] = \
        preds[0, :, :, 3:6]
    res_img[IMAGE_SIZE * 0: IMAGE_SIZE * 1, IMAGE_SIZE * 2: IMAGE_SIZE * 3, :2] = \
        preds[0, :, :, 6:8]
    res_img[IMAGE_SIZE * 1: IMAGE_SIZE * 2, IMAGE_SIZE * 0: IMAGE_SIZE * 1] = \
        preds[0, :, :, 8:11]
    res_img[IMAGE_SIZE * 1: IMAGE_SIZE * 2, IMAGE_SIZE * 1: IMAGE_SIZE * 2] = \
        preds[0, :, :, 11:14]
    res_img[IMAGE_SIZE * 1: IMAGE_SIZE * 2, IMAGE_SIZE * 2: IMAGE_SIZE * 3, :2] = \
        preds[0, :, :, 14:16]

    res_img = np.clip(res_img, -1, 1)
    res_img = ((res_img + 1) * 127.5).astype(np.uint8)
    cv2.imwrite("morph/maps.jpg", cv2.cvtColor(res_img, cv2.COLOR_RGB2BGR))

    org_strength = tf.reshape(tf.range(STEPS, dtype=tf.float32), [STEPS, 1, 1, 1]) / (STEPS - 1)
    trg_strength = tf.reverse(org_strength, axis=[0])

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    video = cv2.VideoWriter(MORPH_DIRECTORY, fourcc, FPS, (original_width, original_height))

    for iterations in tqdm(range(STEPS)):
        preds_org = preds * org_strength[iterations]
        preds_trg = preds * trg_strength[iterations]

        res_targets, res_origins = warp(origins, targets, preds_org[..., :8], preds_trg[..., 8:])
        res_targets = tf.clip_by_value(res_targets, -1, 1)
        res_origins = tf.clip_by_value(res_origins, -1, 1)

        results = (res_targets * trg_strength[iterations]) + \
                  (res_origins * org_strength[iterations])
        res_numpy = results.numpy()

        output_img = ((res_numpy[0] + 1) * 127.5).astype(np.uint8)
        output_img = cv2.cvtColor(output_img, cv2.COLOR_RGB2BGR)
        output_img = cv2.resize(output_img, (original_width, original_height),
                                interpolation=cv2.INTER_AREA)

        video.write(output_img)

    cv2.destroyAllWindows()
    video.release()

def match_size(image_one, image_two):
    """
    The matchSize function takes two images as input and returns the same two images,
    but with one of them resized to match the size of the other.

    :param image_one: Pass in the first image
    :param image_two: Resize the image_one parameter to match its size
    :return: The two images that have been resized to the same size
    """
    if image_one.shape[1] > image_two.shape[1]:
        image_one = cv2.resize(image_one, (image_two.shape[1], image_two.shape[0]),
                               interpolation=cv2.INTER_AREA)
    else:
        image_two = cv2.resize(image_two, (image_one.shape[1], image_one.shape[0]),
                               interpolation=cv2.INTER_AREA)
    return image_one, image_two

def driver(source, target):
    """
    The driver function takes in the source and target images
    It then resizes both images to a square of size im_sz x im_sz, converts them to RGB from BGR,
    It then reshapes the image arrays into 4D tensors, produce warp maps on these two tensors,
    which produces a set of warp maps for each image pair

    :param source: Specify the source image
    :param target: Specify the target image
    :return: The final morphed image
    """
    dom_a = cv2.imread(source, cv2.IMREAD_COLOR)
    dom_b = cv2.imread(target, cv2.IMREAD_COLOR)
    dom_a, dom_b = match_size(dom_a, dom_b)
    original_width = dom_a.shape[1]
    original_height = dom_a.shape[0]
    dom_a = cv2.cvtColor(dom_a, cv2.COLOR_BGR2RGB)
    dom_a = cv2.resize(dom_a, (IMAGE_SIZE, IMAGE_SIZE), interpolation=cv2.INTER_AREA)
    dom_a = dom_a / 127.5 - 1
    dom_b = cv2.cvtColor(dom_b, cv2.COLOR_BGR2RGB)
    dom_b = cv2.resize(dom_b, (IMAGE_SIZE, IMAGE_SIZE), interpolation=cv2.INTER_AREA)
    dom_b = dom_b / 127.5 - 1
    origins = dom_a.reshape(1, IMAGE_SIZE, IMAGE_SIZE, 3).astype(np.float32)
    targets = dom_b.reshape(1, IMAGE_SIZE, IMAGE_SIZE, 3).astype(np.float32)
    produce_warp_maps(origins, targets, original_width, original_height)
    use_warp_maps(origins, targets, original_width, original_height)

def main():
    """
    The driver function for this program.

    :return: The output file name
    """
    if path.exists("output"):
        rmtree("output")
    mkdir("output")
    if path.exists("morph"):
        rmtree("morph")
    mkdir("morph")
    image_list = [file for file in listdir("input") if not file.startswith(".")]
    image_list.sort()
    filenames = []
    for i in range(len(image_list) - 1):
        start = f"input/{image_list[i]}"
        end = f"input/{image_list[i + 1]}"
        driver(start, end)
        filename = f"output/morph{i:03d}.mp4"
        filenames.append(filename)
        move(MORPH_DIRECTORY, filename)
    if LOOP:
        driver(f"input/{image_list[-1]}", f"input/{image_list[0]}")
        filename = f"output/morph{len(image_list):03d}.mp4"
        filenames.append(filename)
        move(MORPH_DIRECTORY, filename)
    combine_videos(filenames, "output/Final", FPS, LOOP)

if __name__ == "__main__":
    WARP_SCALE = 0.075
    MULT_SCALE = 0.4
    ADD_SCALE = 0.4
    ADD_FIRST = False
    LOOP = True
    FPS = 144
    STEPS = 72
    main()


# Object Detection

# Importing the libraries
import torch
from torch.autograd import Variable
import cv2
from data import BaseTransform, VOC_CLASSES as labelmap
from ssd import build_ssd
import imageio

# Defining a function that will do the detections
def detect(frame, net, transform):
    height, width = frame.shape[:2]
    frame_t = transform(frame)[0]
    x = torch.from_numpy(frame_t).permute(2, 0, 1)
    x = Variable(x.unsqueeze(0))
    y = net(x)
    detections = y.data
    scale = torch.Tensor([width, height, width, height])
    # detections = [batch, number of classes, number of occurence, (score, x0, Y0, x1, y1)]
    for i in range(detections.size(1)):
        j = 0
        while detections[0, i, j, 0] >= 0.6:
            pt = (detections[0, i, j, 1:] * scale).numpy()
            cv2.rectangle(frame, (int(pt[0]), int(pt[1])), (int(pt[2]), int(pt[3])), (255, 0, 0), 2)
            cv2.putText(frame, labelmap[i - 1], (int(pt[0]), int(pt[1])), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), 2, cv2.LINE_AA)
            j += 1
    return frame

# Creating the SSD neural network
net = build_ssd('test')
net.load_state_dict(torch.load('ssd300_mAP_77.43_v2.pth', map_location = lambda storage, loc: storage))

# Creating the transformation
transform = BaseTransform(net.size, (104/256.0, 117/256.0, 123/256.0))

# Doing some Object Detection on a video
reader = imageio.get_reader('funny_dog.mp4')
fps = reader.get_meta_data()['fps']
writer = imageio.get_writer('output.mp4', fps = fps)
for i, frame in enumerate(reader):
    frame = detect(frame, net.eval(), transform)
    writer.append_data(frame)
    print(i)
writer.close()


# -*- coding: utf-8 -*-
"""
Created on Tue Apr 23 16:17:44 2020

@author: pedro.marlo
"""

import cdflib
import numpy as np
import logging
import copy
import os

class CdfFileData():
    def __init__(self,filePath, sensorKeyList, boolReadUnits = False):

        logger = logging.getLogger('__main__')

        #### Path of the file.
        self.filePath = filePath
        index = self.filePath.rfind(os.path.sep)
        self.fileName  = self.filePath[index+1:]

        #### List with all the sensors of the CDF. If sensorKeyList == 'all', all sensors will be saved, else: specific sensors
        if sensorKeyList == 'all':
            self.sensorNamesList = cdflib.CDF(self.filePath).cdf_info().get('zVariables')
        else:
            self.sensorNamesList = []
            for sensorName in sensorKeyList:
                if sensorName in cdflib.CDF(self.filePath).cdf_info().get('zVariables'):
                    self.sensorNamesList.append(sensorName)
                else:
                    logger.info('Sensor {} not found in the CDF {}'.format(sensorName, self.fileName))

        if boolReadUnits:
            #### Saving units in a dictionary
            self.units = {}
            self.fullName = {}
            logger.info('Reading {} sensors of {}'.format(len(self.sensorNamesList), self.fileName))
            for sensorName in self.sensorNamesList:
                if sensorName != 'GMT':
                    try:
                        self.units[sensorName] = cdflib.CDF(self.filePath).varattsget(sensorName)['UNIT']
                    except:
                        self.units[sensorName] = ''
                    try:
                        self.fullName[sensorName] = cdflib.CDF(self.filePath).varattsget(sensorName)['EG_NAME']
                    except:
                        self.fullName[sensorName] = ''
                else:
                    self.units['GMT'] = 's'
                    self.fullName['GMT'] = 'Time'
        else:
            self.units = {}
            self.fullName = {}
            for sensorName in self.sensorNamesList:
                self.fullName[sensorName] = ''
                self.units[sensorName] = ''

        #### List with all the times.
        self.time = []
        self.time_year = []

        ### Reading time value
        GMT = cdflib.CDF(self.filePath).varget('GMT')
        ### Changing numpy array to 1D array
        GMT = GMT.ravel()
        ### Last wrong values are -1, deleting them
        indexToDelete = np.where(GMT.astype(int) == -1)
        GMT = np.delete(GMT, indexToDelete)
        ### Creating attributes of time
        self.time_year = GMT
        self.time = GMT-self.time_year[0]
        ### self.timeInitial is an attribute in which range of initial time is copied. It will be always the same
        self.timeInitial = copy.deepcopy(self.time)

        #### Value of the shample rate
        time = self.time[~np.isnan(self.time)]
        self.sampleRate = int(1/(time[1]-time[0]))

        #### List with all the data of every sensor. Organized by columns.
        self.data = []
        self.dataInitial = []
        n = 0
        # All sensors are going to be read
        totalSensorCount = len(self.sensorNamesList)
        # Var to print some percentages, not every step.
        percentageToPrint = 0

        for sensorName in self.sensorNamesList:
            n += 1
            percentage = n/totalSensorCount*100
            if (percentage >= percentageToPrint) or (percentage>= 99):
                logger.info('{:.2f} % completed.'.format(percentage))
                percentageToPrint += 3
            else:
                pass
            # Returns raw data of the variable
            raw_data = cdflib.CDF(self.filePath).varget(sensorName)
            # Deleting elements that were not correct in GMT var
            raw_data = raw_data.ravel()
            raw_data = np.delete(raw_data, indexToDelete)
            # Appending every array in attribute data
            self.data.append(raw_data)
            self.dataInitial.append(raw_data)

    def CDFpos(self,sensor_name):
        '''
        Due to the fact that self.data is a list of lists, it does not have any key to call the sensors, so this method
        returns the position of an specific value.
        '''
        self.cdfpos = self.sensorNamesList.index(sensor_name)
        return self.cdfpos


import textwrap as _textwrap
import networkx as _nx
from pydna._pretty import pretty_str as _pretty_str
from pydna.dseqrecord import Dseqrecord as _Dseqrecord
from pydna.utils import rc as _rc

class Contig(_Dseqrecord):
    """This class holds information about a DNA assembly. This class is instantiated by
    the :class:`Assembly` class and is not meant to be used directly.

    """

    def __init__(self, record, *args, graph=None, nodemap=None, **kwargs):
        super().__init__(record, *args, **kwargs)
        self.graph = graph
        self.nodemap = nodemap

    @classmethod
    def from_string(cls, record: str = "", *args, graph=None, nodemap=None, **kwargs):
        obj = super().from_string(record, *args, **kwargs)
        obj.graph = graph
        obj.nodemap = nodemap
        return obj

    @classmethod
    def from_SeqRecord(cls, record, *args, graph=None, nodemap=None, **kwargs):
        obj = super().from_SeqRecord(record, *args, **kwargs)
        obj.graph = graph
        obj.nodemap = nodemap
        return obj

    def __repr__(self):
        return "Contig({}{})".format({True: "-", False: "o"}[not self.circular], len(self))

    def _repr_pretty_(self, p, cycle):
        """returns a short string representation of the object"""
        p.text("Contig({}{})".format({True: "-", False: "o"}[not self.circular], len(self)))

    def _repr_html_(self):
        return "<pre>" + self.figure() + "</pre>"

    def reverse_complement(self):
        answer = type(self)(super().reverse_complement())
        g = _nx.DiGraph()
        nm = self.nodemap
        g.add_edges_from([(nm[v], nm[u], d) for u, v, d in list(self.graph.edges(data=True))[::-1]])
        g.add_nodes_from((nm[n], d) for n, d in list(self.graph.nodes(data=True))[::-1])
        for u, v, ed in g.edges(data=True):
            ed["name"] = ed["name"][:-3] if ed["name"].endswith("_rc") else "{}_rc".format(ed["name"])[:13]
            ed["seq"] = _rc(ed["seq"])
            ln = len(ed["seq"])
            start, stop = ed["piece"].start, ed["piece"].stop
            ed["piece"] = slice(ln - stop - g.nodes[u]["length"], ln - start - g.nodes[v]["length"])
            ed["features"] = [f._flip(ln) for f in ed["features"]]
        answer.graph = g
        answer.nodemap = {v: k for k, v in self.nodemap.items()}
        return answer

    rc = reverse_complement

    def detailed_figure(self):
        """Returns a text representation of the assembled fragments.

        Linear:

        ::

            acgatgctatactgCCCCCtgtgctgtgctcta
                               TGTGCTGTGCTCTA
                               tgtgctgtgctctaTTTTTtattctggctgtatc

        Circular:

        ::

            ||||||||||||||
            acgatgctatactgCCCCCtgtgctgtgctcta
                               TGTGCTGTGCTCTA
                               tgtgctgtgctctaTTTTTtattctggctgtatc
                                                  TATTCTGGCTGTATC
                                                  tattctggctgtatcGGGGGtacgatgctatactg
                                                                       ACGATGCTATACTG

        """

        fig = ""
        fragmentposition = 0
        nodeposition = 0
        mylist = []
        for u, v, e in self.graph.edges(data=True):
            nodeposition += e["piece"].stop - e["piece"].start
            fragmentposition -= e["piece"].start
            mylist.append([fragmentposition, e["seq"]])
            mylist.append([nodeposition, v.upper()])
            fragmentposition += e["piece"].stop

        if self.circular:
            edges = list(self.graph.edges(data=True))
            nodeposition = edges[0][2]["piece"].start
            nodelength = len(v)
            mylist = [[nodeposition, "|" * nodelength]] + mylist
        else:
            mylist = mylist[:-1]

        firstpos = -1 * min(0, min(mylist)[0])

        for p, s in mylist:
            fig += "{}{}\n".format(" " * (p + firstpos), s)

        return _pretty_str(fig)

    def figure(self):
        r"""Compact ascii representation of the assembled fragments.

        Each fragment is represented by:

        ::

         Size of common 5' substring|Name and size of DNA fragment|
         Size of common 5' substring

        Linear:

        ::

          frag20| 6
                 \\/
                 /\\
                  6|frag23| 6
                           \\/
                           /\\
                            6|frag14

        Circular:

        ::

          -|2577|61
         |       \\/
         |       /\\
         |       61|5681|98
         |               \\/
         |               /\\
         |               98|2389|557
         |                       \\/
         |                       /\\
         |                       557-
         |                          |
          --------------------------

        """
        nodes = list(self.graph.nodes(data=True))
        edges = list(self.graph.edges(data=True))

        if not self.circular:
            r"""
            frag20| 6
                   \/
                   /\
                    6|frag23| 6
                             \/
                             /\
                              6|frag14
            """

            f = edges[0]

            space2 = len(f[2]["name"])

            fig = ("{name}|{o2:>2}\n" "{space2} \\/\n" "{space2} /\\\n").format(
                name=f[2]["name"], o2=nodes[1][1]["length"], space2=" " * space2
            )
            space = space2  # len(f.name)

            for i, f in enumerate(edges[1:-1]):
                name = "{o1:>2}|{name}|".format(o1=nodes[i + 1][1]["length"], name=f[2]["name"])
                space2 = len(name)

                fig += ("{space} {name}{o2:>2}\n" "{space} {space2}\\/\n" "{space} {space2}/\\\n").format(
                    name=name,
                    o2=nodes[i + 2][1]["length"],
                    space=" " * space,
                    space2=" " * space2,
                )

                space += space2

            f = edges[-1]
            fig += ("{space} {o1:>2}|{name}").format(name=f[2]["name"], o1=nodes[-2][1]["length"], space=" " * (space))

        else:  # circular
            r"""
             -|2577|61
            |       \/
            |       /\
            |       61|5681|98
            |               \/
            |               /\
            |               98|2389|557
            |                       \/
            |                       /\
            |                       557-
            |                          |
             --------------------------
            """

            nodes.append(nodes[0])
            f = edges[0]

            space = len(f[2]["name"]) + 3

            fig = (" -|{name}|{o2:>2}\n" "|{space}\\/\n" "|{space}/\\\n").format(
                name=f[2]["name"], o2=nodes[1][1]["length"], space=" " * space
            )

            for i, f in enumerate(edges[1:]):
                name = "{o1:>2}|{name}|".format(o1=nodes[i + 1][1]["length"], name=f[2]["name"])
                space2 = len(name)
                fig += ("|{space}{name}{o2:>2}\n" "|{space}{space2}\\/\n" "|{space}{space2}/\\\n").format(
                    o2=nodes[i + 2][1]["length"],
                    name=name,
                    space=" " * space,
                    space2=" " * space2,
                )
                space += space2

            fig += "|{space}{o1:>2}-\n".format(space=" " * (space), o1=nodes[0][1]["length"])
            fig += "|{space}   |\n".format(space=" " * (space))
            fig += " {space}".format(space="-" * (space + 3))
        return _pretty_str(_textwrap.dedent(fig))

if __name__ == "__main__":
    import os as _os

    cached = _os.getenv("pydna_cached_funcs", "")
    _os.environ["pydna_cached_funcs"] = ""
    import doctest

    doctest.testmod(verbose=True, optionflags=doctest.ELLIPSIS)
    _os.environ["pydna_cached_funcs"] = cached


import curses
import start, nv0, fonts

screen = curses.initscr()

curses.noecho()
curses.cbreak()
curses.start_color()
curses.curs_set(0)
screen.keypad(1)

#------------------------------------Base de Datos------------------------------------
#-----------------Configuración-----------------
animation = True #. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Animación activada SI/NO
flicker = True #. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Parpadeo activado SI/NO

color_bold = ('white') #. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Color de fondo.
color_lyrics = ('black') #. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Color de letra.

cursor = 1
#-----------------------------------------------<<<

##-------------------Opcional-------------------

CURSOR = { #. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Tipo decursores.
	1: '|> ',
	2: ' > ',
	3: '-> ',
	4: '- ',
	5: ''
}

COLOR_LYRICS = { #. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Paleta para Letra.
	'white': curses.COLOR_WHITE,
	'blue': curses.COLOR_BLUE,
	'red': curses.COLOR_RED,
	'black': curses.COLOR_BLACK,
	'yellow': curses.COLOR_YELLOW,
	'cyan': curses.COLOR_CYAN,
	'green': curses.COLOR_GREEN,
}

COLOR_BOLD = { #. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Paleta para fondo.
	'white': curses.COLOR_WHITE,
	'blue': curses.COLOR_BLUE,
	'red': curses.COLOR_RED,
	'black': curses.COLOR_BLACK,
	'yellow': curses.COLOR_YELLOW,
	'cyan': curses.COLOR_CYAN,
	'green': curses.COLOR_GREEN,
}

#-----------------------------------------------<<<

#----------------Animacion SI/NO----------------
if animation == True: #. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Animación por defecto 12ms
	animation = 12

elif animation == False:
	animation = 0
#-----------------------------------------------<<<

#----------------Parpadeo SI/NO-----------------
if flicker == True:
	flicker = curses.A_BLINK

elif flicker == False:
	flicker = curses.A_STANDOUT
#----------------------------------------------<<<

#-------------------Colores--------------------
color_lyrics = COLOR_LYRICS[color_lyrics] #. . . . . . . . . . . . . . . . . . . . . . . . . . . . . Almacena los colores.
color_bold = COLOR_BOLD[color_bold]
cursor = CURSOR[cursor]

curses.init_pair(1, color_bold, color_lyrics)
curses.init_pair(2, curses.COLOR_BLACK, curses.COLOR_RED)
#----------------------------------------------<<<
#-------------------------------------------------------------------------------------<<<

TEXT = {
	1: str('- Presione ESPACIO -'),
	2: str('~ MODO ~'),
	3: str('Eligue tu modo')
}

OPTION = {
	5: 'normal',
	4: 'Hacking',
	3: 'web',
	2: 'game',
	1: 'exit'
}

x = curses.LINES // 2
y = curses.COLS // 2

def process(num):
	screen.border()
	screen.refresh()

	rows = 5
	cols = 36

	y = (curses.LINES - rows) // 2
	x = (curses.COLS - cols) // 2

	win = curses.newwin(rows, cols, y, x)

	win.box()

	win.move(2, 11)
	win.addstr('Proximamente...', curses.A_BOLD)
	escape = False
	option = OPTION[num]

	while escape == False:
		curses.echo()
		key = win.getch(2, 2)
		win.move(2, 2)
		win.addstr('Proximamente...', curses.A_BOLD)

		if key == 10:
			opt(num, option)
			escape = True

		elif key == 27:
			opt(num, option)
			escape = True

def opt(num, opt):
	screen.border()
	escape = False
	fonts.tittle(screen, TEXT[2])

	up_down(num)

	option = str(opt)

	while escape == False:
		key = screen.getch(2, 1)

		if key == 258:
			num -= 1

			if num < 1: num = 1

			option = OPTION[num]

			screen.border()
			up_down(num)
			escape = False

		elif key == 259:
			num += 1

			if num > 5: num = 5

			option = OPTION[num]

			screen.border()
			up_down(num)
			escape = False

		elif key in [curses.KEY_ENTER, ord('\n'), 10]:
			escape = True
			curses.endwin()

			if option == ('normal'):
				nv0.main()

			elif option == ('Hacking'):
				process(4)

			elif option == ('web'):
				process(3)

			elif opt == ('game'):
				main_exit(num)

			elif option == ('exit'):
				main_exit(num)

			else:
				process(2)

		elif key == curses.KEY_RESIZE:
			screen.erase()
			screen.refresh()
			escape = True
			curses.endwin()
			fonts.error()

def group():
	screen.addstr(x - 5, 3, 'Normal       '),
	screen.addstr(x - 3, 3, 'Hacking         '),
	screen.addstr(x - 1, 3, 'Web       '),
	screen.addstr(x + 1, 3, 'Videojuegos       '),
	screen.addstr(x + 3, 3, 'Atras      ')

def up_down(num):
	fonts.tittle(screen, TEXT[2])

	if num == 5:
		screen.border()
		line = x - 5
		group()

		screen.addstr(line, 3,  '               ')
		fonts.animation(screen, '|> Normal ', line, animation, flicker)

	elif num == 4:
		screen.border()
		line = x - 3
		group()

		screen.addstr(line, 3,  '               ')
		fonts.animation_off(screen, '|> Hacking ', line, animation, flicker | curses.color_pair(2))

	elif num == 3:
		screen.border()
		line = x - 1
		group()

		screen.addstr(line, 3,  '             ')
		fonts.animation_off(screen, '|> Web ', line, animation, flicker | curses.color_pair(2))

	elif num == 2:
		screen.border()
		line = x + 1
		group()

		screen.addstr(line, 3,  '              ')
		fonts.animation_off(screen, '|> Videojuegos ', line, animation, flicker | curses.color_pair(2))

	elif num == 1:
		screen.border()
		line = x + 3
		group()

		screen.addstr(line, 3,  '      ')
		fonts.animation(screen, '|> Atras ', line, animation, flicker)

def main_space():
	z = curses.LINES // 2
	screen.addstr(z + 1, 3, '|> Seguir Demo ', curses.color_pair(1) | curses.A_STANDOUT)

	screen.border()
	screen.refresh()

	rows = 5
	cols = 36

	y = (curses.LINES - rows) // 2
	x = (curses.COLS - cols) // 2

	win = curses.newwin(rows, cols, y, x)

	win.box()

	win.move(2, 11)
	win.addstr(TEXT[3], curses.A_BOLD | curses.A_BLINK)

	escape = False

	while escape == False:
		curses.echo()
		key = win.getch(2, 2)

		if key == 10: escape = True

		elif key == 113:
			opt(5, 'normal')
			escape = True

def main():
	main_space()
	fonts.tittle(screen, TEXT[2])
	screen.erase()
	screen.refresh()
	escape = True
	opt(5, 'normal')

def main_exit(num):
	screen.erase()
	screen.refresh()
	curses.endwin()
	start.main(2)


import idd_parser
import pytest

# If this timeouts while you make modifications, you can do `pip install pytest-timeout` and use this mark here:
# @pytest.mark.timeout(2)
def test_ok_idd():
    data = idd_parser.Data()
    data.file = r"""
! Comment

\group test

OkObject,
  A1; \field Name
      \minimum 1

NotBrokenObject,
  A1, \field Name
  N2; \field Last Field
      \maximum 10

Obj2,
  A1; \field Name
"""
    idd_parser.parse_idd(data)
    assert data.schema["properties"].keys() == {"OkObject", "NotBrokenObject", "Obj2"}
    assert data.schema["properties"]["NotBrokenObject"] == {
        "patternProperties": {
            ".*": {"type": "object", "properties": {"last_field": {"type": "number", "maximum": 10.0}}}
        },
        "group": "test",
        "name": {"type": "string"},
        "legacy_idd": {
            "field_info": {
                "name": {"field_name": "Name", "field_type": "a"},
                "last_field": {"field_name": "Last Field", "field_type": "n"},
            },
            "fields": ["name", "last_field"],
            "alphas": {"fields": ["name"]},
            "numerics": {"fields": ["last_field"]},
        },
        "type": "object",
    }

def test_broken_idd():
    data = idd_parser.Data()
    data.file = r"""! Comment

\group test

OkObject,
  A1; \field Name
      \minimum 1

BrokenObject,
  A1, \field Name
  A2, \field Last Field
      \maximum 10

Obj2,
  A1; \field Name
"""
    with pytest.raises(idd_parser.MissingSemiColonException) as e:
        idd_parser.parse_idd(data)

    assert (
        str(e.value)
        == r"""In object 'BrokenObject', Missing semi-colon in field A2
Context:
  A1; \field Name
      \minimum 1

BrokenObject,
  A1, \field Name
  A2, \field Last Field
    ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
      \maximum 10

Obj2,
  A1; \field Name"""
    )

def test_broken_idd_2():
    data = idd_parser.Data()
    data.file = r"""BrokenObject,
  A1 \field Name
  A2, \field Other
"""
    with pytest.raises(idd_parser.IddParsingError) as e:
        idd_parser.parse_idd(data)
    assert (
        str(e.value)
        == r"""In object 'BrokenObject', No comma or semicolon after field A1
Context:
BrokenObject,
  A1 \field Name
    ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  A2, \field Other"""
    )

def test_broken_idd_3():
    data = idd_parser.Data()
    data.file = r"""BrokenObject,
      \obsolete
  A1; \field Name

"""
    with pytest.raises(idd_parser.IddParsingError) as e:
        idd_parser.parse_idd(data)
    assert (
        str(e.value)
        == r"""In object 'BrokenObject', expected string after /obsolete
Context:
BrokenObject,
      \obsolete
                ^~~~~~~~~~~~~~~~~~~~~~~~~
  A1; \field Name
"""
    )

def test_duplicate_a_field_number():
    data = idd_parser.Data()
    data.file = r"""BrokenObject,
  A1, \field Name
  N1, \field Number 1
  A1, \field String Field
  N2; \field Number 1
"""
    with pytest.raises(idd_parser.IddParsingError) as e:
        idd_parser.parse_idd(data)
    assert (
        str(e.value)
        == r"""In object 'BrokenObject', duplicate field number A1
Context:
BrokenObject,
  A1, \field Name
  N1, \field Number 1
  A1, \field String Field
     ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  N2; \field Number 1"""
    )

def test_duplicate_n_field_number():
    data = idd_parser.Data()
    data.file = r"""BrokenObject,
  A1, \field Name
  N1, \field Number 1
  A2, \field String Field
  N1; \field Number 1
"""
    with pytest.raises(idd_parser.IddParsingError) as e:
        idd_parser.parse_idd(data)
    assert (
        str(e.value)
        == r"""In object 'BrokenObject', duplicate field number N1
Context:
BrokenObject,
  A1, \field Name
  N1, \field Number 1
  A2, \field String Field
  N1; \field Number 1
     ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"""
    )


#Author: Nils Becker

import numpy as np
import cv2
import mediapipe as mp
import face_recognition
import base64

cap = cv2.VideoCapture("w3r1.mp4")
#out = cv2.VideoWriter('w3r1_cv.avi',cv2.VideoWriter_fourcc(*'DIVX') , 25, (1280,720))     Creates new Video
mpFaceDetection = mp.solutions.face_detection
mpDraw = mp.solutions.drawing_utils
FaceDetection = mpFaceDetection.FaceDetection(min_detection_confidence = 0.7)
frameCounter = 0
personCounter = 0
faces = {}
table = [] #Hold frames and samples for all ids, form: [[personId,[frames],[faceSamples]]]

def getFaces():
    '''
    This method decodes all saved faces
    '''
    #From FaceRecognition Lab xD
    face_bytes = dict((k, base64.b64decode(v)) for k, v in faces.items())
    face_embeddings = dict((k, np.frombuffer(v, dtype=np.float) if len(v) > 0 else None) for k, v in face_bytes.items())
    return face_embeddings

while True:
    success, frame = cap.read()
    if frame is None: #last frame of video reached
        break

    cv2.putText(frame, f'Total persons in video: {personCounter}',(10,50), cv2.FONT_HERSHEY_PLAIN, 1.5, (0,255,0), 2)
    frameCounter += 1
    newFace = True
    #Defines detection tolerance for part where participants wear masks
    if frameCounter > 1100:
        detectionCutOff = 0.2
    else:
        detectionCutOff = 0.35

    results = FaceDetection.process(frame)
    if results.detections: #faces are present
        for id, detection in enumerate(results.detections):
            bbox1 = detection.location_data.relative_bounding_box
            ih, iw, _ = frame.shape
            bbox = int(bbox1.xmin * iw), int(bbox1.ymin * ih), int(bbox1.width * iw), int(bbox1.height * ih)

            embedding = face_recognition.face_encodings(frame, known_face_locations=[bbox])
            face = base64.b64encode(embedding[0].tobytes())
            faces[personCounter] = face
            faceEmbeddings = getFaces()

            cv2.rectangle(frame, bbox, (0,255,0), 2)

            if len(table) == 0: #gets called when first face is spotted
                personCounter += 1
                table.append([personCounter,[frameCounter],[face]])
                cv2.putText(frame, f'Person {personCounter}, Score{int(detection.score[0]*100)}%', (bbox[0],bbox[1]-10), cv2.FONT_HERSHEY_PLAIN, 1.2, (0,255,0), 1)

            else:
                embeddingsOnly = [x for x in faceEmbeddings.values() if x is not None]
                #strickt tolerance, leads to more different persons detected but differnciates better between faces
                comparisons = face_recognition.compare_faces(embeddingsOnly, faceEmbeddings[personCounter], tolerance = detectionCutOff)
                #Person is already spotted ->
                #comparisons hold at least one "true" for the new face
                for i in range (len(comparisons[:-1])):
                    if comparisons[i]:
                        #Get the face closest to found face
                        distance = face_recognition.face_distance(embeddingsOnly, faceEmbeddings[personCounter])
                        nameIdx = np.argmin(distance[:-1]) + 1
                        for entry in table:
                            if entry[0] == nameIdx:
                                entry[1].append(frameCounter)
                                entry[2].append(face)
                        newFace = False
                        cv2.putText(frame, f'Person{nameIdx}, Score{int(detection.score[0]*100)}%', (bbox[0],bbox[1]-10), cv2.FONT_HERSHEY_PLAIN, 1, (0,255,0), 1)
                        break
                if (newFace):  #New Person found
                    personCounter += 1
                    cv2.putText(frame, f'Person{personCounter}, Score{int(detection.score[0]*100)}%', (bbox[0],bbox[1]-10), cv2.FONT_HERSHEY_PLAIN, 1, (0,255,0), 1)
                    table.append([personCounter,[frameCounter],[face]])

    #out.write(frame)   Write frame with boundling box etc. to new Video
    cv2.imshow("press 'q' to quit",frame)
    if cv2.waitKey(1)==ord('q'):
        break

print(f'Total persons in video:{personCounter}')

#out.release()
cap.release()
cv2.destroyAllWindows()


import os
import shutil
import numpy as np
import pickle
if __name__ == '__main__':
    import sys
    sys.path.append(os.path.join(os.path.dirname(__file__), '..'))

from tqdm import tqdm
from data import VIDEO_DIR, TRAIN_SPLIT_FILE, TEST_SPLIT_FILE, VIDEO_EXT, LABEL_EXT
from data import TRAIN_VID_DIR, TEST_VID_DIR, TRAIN_LABEL_DIR, TEST_LABEL_DIR
from data import TRAIN_SEGMENT_DIR, TRAIN_SEGMENT_DICT, TEST_SEGMENT_DIR, TEST_SEGMENT_DICT
from data import read_label_file, read_mapping_file, read_segment_from_video
from data import OLD_MAPPING_FILE, MAPPING_FILE
from data import BAD_TRAIN_SEGMENTS_FILE, BAD_TEST_SEGMENTS_FILE

def _read_split_file(split_file):
    with open(split_file) as f:
        videos = f.readlines()

    videos = videos[1:]
    videos = [video.strip() for video in videos]
    videos = [(video.split('/'))[-1] for video in videos]
    videos = [video[:-4] for video in videos]
    videos = [video.split('_') for video in videos]
    return videos

def _get_split_dict(split_file):
    split_vids = _read_split_file(split_file)
    split_dict = dict()

    for vid in split_vids:
        if vid[0] not in split_dict.keys():
            split_dict[vid[0]] = dict()

        if vid[1] == 'stereo01':
            vid[1] = 'stereo'

        part_dict = split_dict[vid[0]]
        if vid[1] not in part_dict:
            part_dict[vid[1]] = []

        camera_dict = part_dict[vid[1]]
        camera_dict.append('_'.join(vid[2:]))
    return split_dict

def _get_video_list():
    print('Retrieving list of videos...')
    videos = []

    parts = os.listdir(VIDEO_DIR)
    for part in parts:
        part_dir = os.path.join(VIDEO_DIR, part)
        cameras = os.listdir(part_dir)

        for camera in cameras:
            camera_dir = os.path.join(part_dir, camera)
            files = os.listdir(camera_dir)
            files = [file for file in files if file[-7:] != '.labels']
            videos += files
    return videos

def _print_video_exts():
    """ prints the list of video annotations, video extension is .avi """
    videos = _get_video_list()
    exts = [vid.split('.')[-1] for vid in videos]
    exts = np.unique(exts)
    print(exts)

def _split_train_test_videos():
    def split_videos(video_dir, label_dir, split_dict):
        if not os.path.exists(video_dir):
            os.makedirs(video_dir)

        if not os.path.exists(label_dir):
            os.makedirs(label_dir)

        pbar = tqdm(split_dict.items())
        for part, part_dict in pbar:
            pbar.set_postfix({'part': part})
            for camera, files in part_dict.items():
                for file in files:
                    if camera == 'stereo':
                        original_vid = os.path.join(VIDEO_DIR, part, camera, file + '_ch1' + VIDEO_EXT)
                        original_label = os.path.join(VIDEO_DIR, part, camera, file + '_ch1' + LABEL_EXT)

                        if not (os.path.exists(original_label) and os.path.exists(original_vid)):
                            original_vid = os.path.join(VIDEO_DIR, part, camera, file + '_ch0' + VIDEO_EXT)
                            original_label = os.path.join(VIDEO_DIR, part, camera, file + '_ch0' + LABEL_EXT)
                    else:
                        original_vid = os.path.join(VIDEO_DIR, part, camera, file + VIDEO_EXT)
                        original_label = os.path.join(VIDEO_DIR, part, camera, file + LABEL_EXT)

                    assert os.path.exists(original_vid), print(original_vid)
                    assert os.path.exists(original_label), print(original_label)

                    if camera == 'stereo':
                        copy_vid = os.path.join(video_dir, '.'.join([part, camera + '01', file + VIDEO_EXT]))
                        copy_label = os.path.join(label_dir, '.'.join([part, camera + '01', file + LABEL_EXT]))
                    else:
                        copy_vid = os.path.join(video_dir, '.'.join([part, camera, file + VIDEO_EXT]))
                        copy_label = os.path.join(label_dir, '.'.join([part, camera, file + LABEL_EXT]))

                    if not os.path.exists(copy_vid):
                        shutil.copy2(original_vid, copy_vid)
                    if not os.path.exists(copy_label):
                        shutil.copy2(original_label, copy_label)
    print('generating train videos and labels')
    train_split = _get_split_dict(TRAIN_SPLIT_FILE)
    split_videos(TRAIN_VID_DIR, TRAIN_LABEL_DIR, train_split)

    print('generating test videos and labels')
    test_split = _get_split_dict(TEST_SPLIT_FILE)
    split_videos(TEST_VID_DIR, TEST_LABEL_DIR, test_split)

def _extract_train_test_segments():
    def generate_segments(video_dir, label_dir, segment_dir):
        # the frames are in one-based indexing.
        if not os.path.exists(segment_dir):
            os.makedirs(segment_dir)

        files = sorted(os.listdir(video_dir))
        action_to_logit_dict, _ = read_mapping_file(MAPPING_FILE)

        all_dict = dict()
        pbar = tqdm(files)
        for file in pbar:
            pbar.set_postfix({'video': file})

            vid_file = os.path.join(video_dir, file)
            label_file = str(file)[:-len(VIDEO_EXT)] + LABEL_EXT
            label_file = os.path.join(label_dir, label_file)

            segment_windows, segment_actions = read_label_file(label_file)
            for i, segment_action in enumerate(segment_actions):
                segment_window = segment_windows[i]
                segment_logit = action_to_logit_dict[segment_action]
                segment_name = '.'.join([file[:-len(VIDEO_EXT)], str(i)])
                segment_file = os.path.join(segment_dir, segment_name + '.npy')
                segment_frames = read_segment_from_video(vid_file, segment_window)
                np.save(segment_file, segment_frames)

                segment_dict = {
                    'action': segment_action,
                    'label': segment_logit,
                    'window': segment_window,
                    'vid-file': vid_file,
                    'label-file': label_file,
                    'segment-file': segment_file,
                    'n-frames': len(segment_frames)
                }
                all_dict[segment_name] = segment_dict
        return all_dict
    train_segments_dict = generate_segments(TRAIN_VID_DIR, TRAIN_LABEL_DIR, TRAIN_SEGMENT_DIR)
    with open(TRAIN_SEGMENT_DICT, 'wb') as f:
        pickle.dump(train_segments_dict, f)
    test_segments_dict = generate_segments(TEST_VID_DIR, TEST_LABEL_DIR, TEST_SEGMENT_DIR)
    with open(TEST_SEGMENT_DICT, 'wb') as f:
        pickle.dump(test_segments_dict, f)

def _generate_mapping_file():
    def update_mapping(mapping, label_dir):
        label_files = os.listdir(label_dir)
        label_files = [os.path.join(label_dir, file) for file in label_files]
        for file in label_files:
            _, segment_labels = read_label_file(file)
            for label in segment_labels:
                if label not in mapping.keys():
                    mapping[label] = len(mapping.keys())
        return mapping
    updated_mapping, _ = read_mapping_file(OLD_MAPPING_FILE)
    updated_mapping = update_mapping(updated_mapping, label_dir=TRAIN_LABEL_DIR)
    updated_mapping = update_mapping(updated_mapping, label_dir=TEST_LABEL_DIR)
    with open(MAPPING_FILE, 'w') as f:
        for key, value in updated_mapping.items():
            f.write(' '.join([str(value), str(key)]) + '\n')

def _check_train_test_segments():
    def check_segments(all_dict_file):
        with open(all_dict_file, 'rb') as f:
            all_dict = pickle.load(f)
        bad_segments = []
        pbar = tqdm(all_dict.items())
        for segment, segment_dict in pbar:
            segment_file = segment_dict['segment-file']
            segment_frames = np.load(segment_file)
            segment_window = segment_dict['window']

            if len(segment_frames) == 0 and (segment_window[1] - segment_window[0] != 0):
                bad_segments.append(segment)
            pbar.set_postfix({'# bad segments': len(bad_segments)})
        return bad_segments
    # checking train and test segments
    bad_train_segments = check_segments(TRAIN_SEGMENT_DICT)
    with open(BAD_TRAIN_SEGMENTS_FILE, 'w') as f:
        for seg in bad_train_segments:
            f.write(seg + '\n')

    bad_test_segments = check_segments(TEST_SEGMENT_DICT)
    with open(BAD_TEST_SEGMENTS_FILE, 'w') as f:
        for seg in bad_test_segments:
            f.write(seg + '\n')

def main():
    # _split_train_test_videos()
    # _generate_mapping_file()
    # _extract_train_test_segments()
    # _check_train_test_segments()
    pass

if __name__ == '__main__':
    main()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Mar  6 18:41:45 2017

@author: EmileMathieu
"""
import numpy as np

from hmc import HMC
from metropolis import AMH
from rmhmc import RMHMC
from gibbs_sampler import auxiliary_gibbs
from tools import nextpow2, ac, CalculateESS
from iwls import iwls

if __name__ == '__main__':
    #%% Load and preprocess data
    dataset_name = 'australian'

    if dataset_name == 'heart':
        X = np.loadtxt(open("data/heart.csv", "rb"), delimiter=",")
        t = np.reshape(X[:, -1], (-1, 1))
        # Replacing classes 1 (resp. 2) by label 0 (resp. 1)
        t[t == 1] = 0
        t[t == 2] = 1
    elif dataset_name == 'australian':
        X = np.loadtxt(open("data/australian.csv", "rb"), delimiter=",")
        t = np.reshape(X[:, -1], (-1, 1))
    else:
        print("Error! Dataset must be 'australian' or 'heart'.")

    X = np.delete(X, -1, 1)
    N = X.shape[0]

    X = (X - np.tile(np.mean(X, axis=0), (N, 1))) / np.tile(np.std(X, axis=0), (N, 1))

    # Create Polynomial Basis
    XX = np.ones((N, 1))
    XX = np.hstack((XX, X))

    n_experiments = 10
    D = XX.shape[1]

    results_beta = np.zeros((n_experiments, 5000, D))
    results_time = np.zeros(n_experiments)
    for i in range(n_experiments):
        #results_beta[i], results_time[i] = iwls(XX, t)
        #results_beta[i], results_time[i] = auxiliary_gibbs(XX, t)
        #results_beta[i], results_time[i] = AMH(XX, t)
        #results_beta[i], results_time[i] = RMHMC(XX, t)
        results_beta[i], results_time[i] = HMC(XX, t)
    avg_beta_posterior = np.mean(results_beta, axis=0)
    avg_time_taken = np.mean(results_time)

    import pdb; pdb.set_trace()

    #%%
    #%matplotlib inline

    import matplotlib.pyplot as plt
    plt.plot(range(avg_beta_posterior.shape[0]), avg_beta_posterior, linewidth=0.2)

    #%%
    ACF = ac(avg_beta_posterior[:, 0], len(avg_beta_posterior[:, 0])-1)
    plt.plot(range(len(avg_beta_posterior[:, 0])-1), ACF)

    #%%
    print('ESS')
    ESS = CalculateESS(avg_beta_posterior, avg_beta_posterior.shape[0]-1)

    print('Min', np.min(ESS))
    print('Median', np.median(ESS))
    print('Mean', np.mean(ESS))
    print('Max', np.max(ESS))
    print('Time', avg_time_taken)

    print('Time per Min ESS:', round(avg_time_taken / np.min(ESS), 6))


"""

26 May 2017: DO NOT USE!
There are major problems with the gridding/interpolation scheme in here
Use project_racmo.R instead!

Project RACMO 2.3 data onto a regular grid using interpolation.
Outputs an equivalent NetCDF of the time series.

Remember to set environment variable PROCESS_DIR, with trailing slash

example for polar stereographic full time series:
reproject_racmo.py /scratch/L0data/RACMO/RACMO2.3_GRN11_runoff_monthly_1958-2015.nc /scratch/L0data/RACMO/RACMO2.3_GRN11_masks.nc 1958-01-01 2015-12-31

@author Andrew Tedstone a.j.tedstone@bristol.ac.uk
@created 2016-December-05

"""

import xarray as xr
import pyproj
from scipy import interpolate
import pandas as pd
import numpy as np
from osgeo import osr, gdal
from scipy import ndimage
import os
import argparse
import datetime as dt
import math
import georaster
import matplotlib.pyplot as plt

parser = argparse.ArgumentParser(description='Project geographic RACMO data for Arctic domain')

# Positional arguments
parser.add_argument('fn_RACMO', type=str, help='str, absolute path to RACMO netCDF file')
parser.add_argument('fn_RACMO_masks', type=str, help='str, absolute path to RACMO masks netCDF file')
parser.add_argument('date_start', type=str, help='str, start date of RACMO data in format yyyy-mm-dd')
parser.add_argument('date_end', type=str, help='str, end date of RACMO data in format yyyy-mm-dd')

parser.add_argument('-grid', type=str, dest='grid', default='pstere', help='str, grid to use, pstere or bamber')
parser.add_argument('-test', dest='test', action='store_true', help='Test script (run only for first year of data)')

args = parser.parse_args()

##############################################################################

# Set grid to 'pstere' or 'bamber'
# 'bamber' grid is GrIS only, 'pstere' will result in grid across full RACMO extent
# Use 'bamber' grid only for testing to ensure that interpolation result matches 2012 FWF paper results
# Characteristics of each grid are hard-coded into logic below
grid = args.grid

PROCESS_DIR = os.environ['PROCESS_DIR']

# ============================================================================

def pstere_lat2k0(lat):
	""" Given lat in degrees, return k0

	Implementation of eq. 1 in:
	Rollins, 2011. Computation of scale-factor and standard parallel for the
	polar stereographic projection.
	URL: http://earth-info.nga.mil/GandG/coordsys/polar_stereographic/Polar_Stereo_phi1_from_k0_memo.pdf
	Last retrieved 2012-Dec-07.

	"""

	ecc = 0.0818191908426215
	lat = np.deg2rad(70)
	k90 = np.sqrt(np.power((1. + ecc), (1. + ecc)) * np.power((1. - ecc), (1. - ecc)))
	k0_1 = ((1. + np.sin(lat)) / 2.)
	k0_2 = (k90 / np.sqrt(np.power(1. + (ecc * np.sin(lat)), (1.+ecc)) * np.power(1. - (ecc * np.sin(lat)), (1.-ecc))))
	k0 = k0_1 * k0_2

	return k0

# ============================================================================

## Determine grid projection
# 5 km grid is hard-coded
if args.grid == 'pstere':
	grid_proj = pyproj.Proj('+init=EPSG:3413')
	fn_mask_ice = PROCESS_DIR + 'mask_ice_racmo_EPSG3413_5km.tif'
	fn_mask_land = PROCESS_DIR + 'mask_land_racmo_EPSG3413_5km.tif'
	fn_mask_greenland = PROCESS_DIR + 'mask_greenland_racmo_EPSG3413_5km.tif'
	fn_dem_ice = PROCESS_DIR + 'dem_ice_racmo_EPSG3413_5km.tif'
	fn_mask_landandice = PROCESS_DIR + 'mask_landandice_racmo_EPSG3413_5km.tif'
elif args.grid == 'bamber':
	grid_proj = pyproj.Proj('+proj=sterea +lat_0=90 +lat_ts=71 +lon_0=-39 +k=1 +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs')

# Create SRS representation
srs = osr.SpatialReference()
srs.ImportFromProj4(grid_proj.srs)

## Open RACMO runoff dataset
# Times aren't encoded in format understandable by xarray
runoff = xr.open_dataset(args.fn_RACMO,
		decode_times=False, chunks={'time':15})
# Apply time dimension in format understandable by xarray
times = pd.date_range(args.date_start, args.date_end, freq='1MS')
runoff['time'] = times

## Extract and then project lat/lon coordinates to grid projection
lon = runoff.LON.values.flatten()
lat = runoff.LAT.values.flatten()
latlon = np.vstack((lat, lon)).T
(x, y) = grid_proj(lon, lat)
xy = np.vstack((x, y)).T

def grid_ceil(n):
	return math.ceil(n / 10000) * 10000

def grid_floor(n):
	return math.floor(n / 10000) * 10000

## Create grid at 5 km resolution
if grid == 'pstere':
	yi, xi = np.mgrid[grid_floor(np.min(y)):grid_ceil(np.max(y)):5000, grid_floor(np.min(x)):grid_ceil(np.max(x)):5000]
elif grid == 'bamber':
	yi, xi = np.mgrid[-3400000:-600000:5000, -800000:700000:5000]

## Calculate scaling factors for grid (after Bamber Interpolate_racmo.BAK)
# Calculate lat/lon coordinates for each point of projected grid
(grid_lon, grid_lat) = grid_proj(xi.flatten(), yi.flatten(), inverse=True)
# Convert to paired columns
grid_latlon = np.vstack((grid_lat, grid_lon)).T
# Convert to radians
grid_latlon_radians = grid_latlon / 57.29578
# Calculate latitudinal scaling
m = 2.0 * np.tan(45.0 / 57.29578 - (grid_latlon_radians[:, 0] / 2)) / np.cos(grid_latlon_radians[:, 0])
# Compute scale factor for each grid cell
k0 = pstere_lat2k0(srs.GetProjParm('latitude_of_origin'))
print(k0)
scale_factors = np.power(m * k0, 2) #np.sqrt((m * k0))
# reshape scale_factors to the same dimensions as the grid
scale_factors = scale_factors.reshape(xi.shape)

# Create lat/lon variables of projected grid
grid_lon_2d = grid_lon.reshape(xi.shape)
grid_lat_2d = grid_lat.reshape(yi.shape)

## Begin the interpolation/projection

# Are we just testing the first year or not?
if not args.test:
	process_times = times
else:
	date_end_here = (dt.datetime.strptime(args.date_start, '%Y-%m-%d') + dt.timedelta(days=364)).strftime('%Y-%m-%d')
	process_times = pd.date_range(args.date_start, date_end_here, freq='1MS')

# Set up store in which to hold projected (gridded) data
store = np.zeros((len(process_times), yi.shape[0], yi.shape[1]))
# Integer time counter for indexing to np store array
nt = 0

# Do the projection to grid
for t in process_times:

	print(t.strftime('%Y-%m'))

	# Extract runoff mmWE value at each coordinate
	r = runoff.runoff.sel(time=t).values
	# And flatten to a vector which matches the xy (grid coords) array order
	r = r.flatten()

	# Interpolate point data onto grid
	zi = interpolate.griddata(xy, r, (xi, yi), method='linear')
	# plt.figure()
	# plt.imshow(zi, origin='lower')
	# plt.colorbar()
	# plt.show()
	# CHECK:  lookup the values of geo coords in the new grid!!

	print(np.mean(zi))

	# Apply scale_factor
	zi_sf = zi / scale_factors

	# Save to store array
	store[nt, :, :] = zi_sf

	# Increment year counter
	nt += 1

## Create xarray data array
# Coordinates of the data, for xarray representation
coords = {'TIME': process_times, 'Y': yi[:,0], 'X': xi[0,:]}
# Projected runoff data
da = xr.DataArray(store, coords=coords,
	dims=['TIME', 'Y', 'X'], encoding={'dtype':np.dtype('Float32')})

# Set attributes from RACMO netCDF
da.name = runoff.runoff.long_name
da.attrs['long_name'] = runoff.runoff.long_name
da.attrs['units'] = runoff.runoff.units
da.attrs['standard_name'] = runoff.runoff.standard_name
da.attrs['grid_mapping'] = 'polar_stereographic'

## Create associated lat/lon coordinates DataArrays
coords_geo = {'Y': yi[:,0], 'X': xi[0,:]}

lon_da = xr.DataArray(grid_lon_2d, coords=coords_geo, dims=['Y', 'X'],
	encoding={'_FillValue': -9999.})
lon_da.attrs['grid_mapping'] = 'polar_stereographic'
lon_da.attrs['units'] = 'degrees'
lon_da.attrs['standard_name'] = 'longitude'

lat_da = xr.DataArray(grid_lat_2d, coords=coords_geo, dims=['Y', 'X'],
	encoding={'_FillValue': -9999.})
lat_da.attrs['grid_mapping'] = 'polar_stereographic'
lat_da.attrs['units'] = 'degrees'
lat_da.attrs['standard_name'] = 'latitude'

## Define projection
crs = xr.DataArray(0, encoding={'dtype':np.dtype('int8')})
crs.attrs['grid_mapping_name'] = 'polar_stereographic'
crs.attrs['scale_factor_at_central_origin'] = srs.GetProjParm('scale_factor')
crs.attrs['standard_parallel'] = srs.GetProjParm('latitude_of_origin')
crs.attrs['straight_vertical_longitude_from_pole'] = srs.GetProjParm('central_meridian')
crs.attrs['false_easting'] = srs.GetProjParm('false_easting')
crs.attrs['false_northing'] = srs.GetProjParm('false_northing')
# lat_0 does not have a WKT representation for polar stereographic
# http://cfconventions.org/wkt-proj-4.html
# look it up from Proj4 string directly instead
p4 = srs.ExportToProj4()
lat_0_pos = p4.find('+lat_0=')
lat_0 = float(p4[lat_0_pos+7:lat_0_pos+7+3].strip())
crs.attrs['latitude_of_projection_origin'] = lat_0

## Create the Dataset
ds = xr.Dataset({'runoff':da, 'lon': lon_da, 'lat': lat_da, 'polar_stereographic':crs})

# Main metadata
ds.attrs['Conventions'] = 'CF-1.4'
ds.attrs['history'] = 'This NetCDF generated using bitbucket atedstone/fwf/reproject_racmo.py on %s provided by JLB/BN/MvdB\n %s' %(args.fn_RACMO, runoff.history)
ds.attrs['institution'] = 'University of Bristol (Andrew Tedstone), IMAU (Brice Noel)'
ds.attrs['title'] = 'Monthly runoff in the RACMO 2.3 domain on a projected grid'
ds.attrs['source'] = 'RACMO 2.3'
ds.attrs['proj4'] = grid_proj.srs

# Additional geo-referencing
ds.attrs['nx'] = float(xi.shape[1])
ds.attrs['ny'] = float(yi.shape[0])
ds.attrs['xmin'] = float(np.round(np.min(xi), 0))
ds.attrs['ymax'] = float(np.round(np.max(yi), 0))
ds.attrs['spacing'] = 5000.

# NC conventions metadata for dimensions variables
ds.X.attrs['units'] = 'meters'
ds.X.attrs['standard_name'] = 'projection_x_coordinate'
ds.X.attrs['point_spacing'] = 'even'
ds.X.attrs['axis'] = 'X'

ds.Y.attrs['units'] = 'meters'
ds.Y.attrs['standard_name'] = 'projection_y_coordinate'
ds.Y.attrs['point_spacing'] = 'even'
ds.Y.attrs['axis'] = 'Y'

ds.TIME.attrs['standard_name'] = 'time'
ds.TIME.attrs['axis'] = 'TIME'

## Save the Dataset as a NetCDF file
if args.test:
	fn_save = args.fn_RACMO.split('/')[-1][:-3] + '_%s_TEST.nc' % args.grid
else:
	fn_save = args.fn_RACMO.split('/')[-1][:-3] + '_%s.nc' % args.grid
print(PROCESS_DIR + fn_save)
ds.to_netcdf(PROCESS_DIR + fn_save,	format='NETCDF4')

##############################################################################
## Do some simple checks if projecting the bamber grid (as per used in 2012 paper)

if args.grid == 'bamber':

	# To create this file, use Williams mc_land_mask___bamber_proj.tif:
	# gdalwarp -tr 5000 5000 -te -800000 -3400000 700000 -600000 mc_land_mask___bamber_proj.tif mc_land_mask___bamber_proj_5km.tif
	mask = georaster.SingleBandRaster('/scratch/bedmachine/mc_land_mask___bamber_proj_5km.tif')
	ice_area = np.where(mask.r == 2, 1, 0)

	#print((np.sum(store[0:12, :, :] * np.flipud(ice_area)) * (5*5)/1.0e6))

	# Bamber's value for gridded product stated as 242.888 km^3 in Interpolate_racmo.BAK
	print('1958 runoff flux:')
	r1958 = ((ds.runoff.sel(TIME=slice('1958-01-01','1958-12-01')) * np.flipud(ice_area)) * (5*5) / 1.0e6)
	print(r1958.sum())

	plt.figure()
	r1958.sum(dim='TIME').plot()

	plt.figure()
	(ds.runoff.sel(TIME=slice('1958-01-01','1958-12-01')) * (5*5) / 1.0e6).sum(dim='TIME').plot()
	plt.show()

##############################################################################
## Project masks(s) etc to geotiff
# Use all the same grid logic as defined above
print('Projecting mask(s)...')

# Set geoTransform for writing to geotiff
trans = (xi[0,0], (xi[0,1]-xi[0,0]), 0, yi[-1,0], 0, (yi[-2,0]-yi[-1,0]))

# Set up coords for saving to netcdf
masks_coords = {'Y': yi[:,0], 'X': xi[0,:]}

# Open RACMO masks (they are on same grid)
masks_ds = xr.open_dataset(args.fn_RACMO_masks)

## Mask of all ice areas; we'll need to split into numerical basins later
mask_ice_pts = masks_ds.icecon.values
# And flatten to a vector which matches the xy (grid coords) array order
mask_ice_pts = mask_ice_pts.flatten()

# Interpolate point data onto grid
mask_ice_gridded = interpolate.griddata(xy, mask_ice_pts, (xi, yi), method='nearest')

# Write to geotiff
georaster.simple_write_geotiff(
	fn_mask_ice,
	np.flipud(mask_ice_gridded),
	trans,
	proj4=grid_proj.srs,
	dtype=gdal.GDT_Byte
	)
xr_mask_ice = xr.DataArray(mask_ice_gridded, dims=('Y', 'X'), coords=masks_coords)

## Ice-only Topography
topo_ice_pts = masks_ds.topography.where(masks_ds.icecon == 1).values
# And flatten to a vector which matches the xy (grid coords) array order
topo_ice_pts = topo_ice_pts.flatten()

# Interpolate point data onto grid
topo_ice_gridded = interpolate.griddata(xy, topo_ice_pts, (xi, yi), method='linear')

# Write to geotiff
georaster.simple_write_geotiff(
	fn_dem_ice,
	np.flipud(topo_ice_gridded),
	trans,
	proj4=grid_proj.srs,
	dtype=gdal.GDT_Float64
	)
xr_topo_ice = xr.DataArray(topo_ice_gridded, dims=['Y','X'], coords=masks_coords, encoding={'dtype':np.dtype('Float32')})

## Land-and-ice mask
mask_landandice_pts = (masks_ds.LSM_noGrIS + masks_ds.Gr_land).values
mask_landandice_pts = mask_landandice_pts.flatten()
mask_landandice_gridded = interpolate.griddata(xy, mask_landandice_pts, (xi, yi), method='nearest')
mask_landandice_gridded = ndimage.binary_fill_holes(mask_landandice_gridded)

# Write to geotiff
georaster.simple_write_geotiff(
	fn_mask_landandice,
	np.flipud(mask_landandice_gridded),
	trans,
	proj4=grid_proj.srs,
	dtype=gdal.GDT_UInt16
	)
xr_mask_landandice = xr.DataArray(mask_landandice_gridded, dims=['Y','X'], coords=masks_coords)

## Land-only mask
mask_land_pts = ((masks_ds.LSM_noGrIS + masks_ds.Gr_land) - (masks_ds.icecon == 1)).values
mask_land_pts = mask_land_pts.flatten()
mask_land_gridded = interpolate.griddata(xy, mask_land_pts, (xi, yi), method='nearest')

# Write to geotiff
georaster.simple_write_geotiff(
	fn_mask_land,
	np.flipud(mask_land_gridded),
	trans,
	proj4=grid_proj.srs,
	dtype=gdal.GDT_UInt16
	)
xr_mask_land = xr.DataArray(mask_land_gridded, dims=['Y','X'], coords=masks_coords)

## Greenland land+ice
greenland_pts = masks_ds.Gr_land.values
greenland_pts = greenland_pts.flatten()
greenland_gridded = interpolate.griddata(xy, greenland_pts, (xi, yi), method='nearest')
greenland_gridded = np.where(greenland_gridded > 0, 1, 0)
# Write to geotiff
georaster.simple_write_geotiff(
	fn_mask_greenland,
	np.flipud(greenland_gridded),
	trans,
	proj4=grid_proj.srs,
	dtype=gdal.GDT_UInt16
	)
xr_mask_greenland = xr.DataArray(greenland_gridded, dims=['Y','X'], coords=masks_coords)

## The others for netcdf only
icemask_noGrIS = masks_ds.icemask_no_GrIS.values
icemask_noGrIS_gridded = interpolate.griddata(xy, icemask_noGrIS.flatten(), (xi, yi), method='nearest')
xr_mask_icemask_noGrIS = xr.DataArray(icemask_noGrIS_gridded, dims=['Y','X'], coords=masks_coords)

LSM_noGrIS = masks_ds.LSM_noGrIS.values
LSM_noGrIS_gridded = interpolate.griddata(xy, LSM_noGrIS.flatten(), (xi, yi), method='nearest')
xr_mask_LSM_noGrIS = xr.DataArray(LSM_noGrIS_gridded, dims=['Y','X'], coords=masks_coords)

Gr_land = masks_ds.Gr_land.values
Gr_land_gridded = interpolate.griddata(xy, Gr_land.flatten(), (xi, yi), method='nearest')
xr_mask_Gr_land = xr.DataArray(Gr_land_gridded, dims=['Y','X'], coords=masks_coords)

GrIS_mask = masks_ds.GrIS_mask.values
GrIS_mask_gridded = interpolate.griddata(xy, GrIS_mask.flatten(), (xi, yi), method='nearest')
xr_mask_GrIS_mask = xr.DataArray(GrIS_mask_gridded, dims=['Y','X'], coords=masks_coords)

GrIS_caps_mask = masks_ds.GrIS_caps_mask.values
GrIS_caps_mask_gridded = interpolate.griddata(xy, GrIS_caps_mask.flatten(), (xi, yi), method='nearest')
xr_mask_GrIS_caps_mask = xr.DataArray(GrIS_caps_mask_gridded, dims=['Y','X'], coords=masks_coords)

## Also put them all into a netCDF
new_ds_masks = xr.Dataset({
	'LSM_noGrIS':xr_mask_LSM_noGrIS,
	'icemask_no_GrIS':xr_mask_icemask_noGrIS,
	'Gr_land':xr_mask_Gr_land,
	'GrIS_mask':xr_mask_GrIS_mask,
	'GrIS_caps_mask':xr_mask_GrIS_caps_mask,
	'Greenland_all':xr_mask_greenland,
	'all_land_no_ice':xr_mask_land,
	'landandice':xr_mask_landandice,
	'all_ice':xr_mask_ice,
	'ice_topo':xr_topo_ice
	})

fn_save = args.fn_RACMO_masks.split('/')[-1][:-3] + '_%s.nc' % args.grid
new_ds_masks.to_netcdf(PROCESS_DIR + fn_save)

print('Finished.')

"""
Code to plot all masks:

for i in masks.variables:
	print(i)
	if i in ('X','Y'): continue
	figure(),masks[i].plot(),title(i)

"""

#############################################################################

## An example runoff file for testing runoff routing
# runoff_195807 = ds.runoff.sel(TIME='1958-07-01').where(mask_gridded == 1)
# runoff_195807 = runoff_195807 * 5*5 / 1.0e6
# georaster.simple_write_geotiff(
# 	'/home/at15963/Dropbox/work/papers/bamber_fwf/RACMO/racmo_pstere_5km_runoff_1958-07.tif',
# 	np.flipud(runoff_195807),
# 	trans,
# 	proj4=grid_proj.srs,
# 	dtype=gdal.GDT_Float64
# 	)


#!/usr/bin/env python

import argparse
import sys
import json
from decimal import Decimal
import ijson
import requests

def parse_args():
    parser = argparse.ArgumentParser(description='Load JSON documents into a CouchDB database')
    parser.add_argument('-H', '--host', dest='host', default='localhost')
    parser.add_argument('-P', dest='port', type=int, default=5984)
    parser.add_argument('-u', dest='username')
    parser.add_argument('-p', dest='password')
    parser.add_argument('-b', dest='bulk_size', type=int, default=1000)
    parser.add_argument('target_database', help='database to import data to')
    return parser.parse_args()

def default(obj):
    if isinstance(obj, Decimal):
        return float(obj)
    raise TypeError("Object of type '%s' is not JSON serializable" % type(obj).__name__)

def load(args, input_file, bulk_size=1000):
    session = requests.session()
    bulk = []
    for doc in ijson.items(input_file, 'docs.item'):
        bulk.append(doc)
        if len(bulk) == bulk_size:
            bulk_insert(args, session, bulk)
            bulk = []
    if bulk:
        bulk_insert(args, session, bulk)

def bulk_insert(args, session, docs):
    url = 'http://{}:{:d}/{}/_bulk_docs'.format(args.host, args.port, args.target_database)
    auth = (args.username, args.password or '') if args.username else None
    payload = json.dumps({'docs': docs}, separators=(',', ':'), default=default)
    response = session.post(url, data=payload, auth=auth, headers={
        'Accept': 'application/json',
        'Content-Type': 'application/json',
        'X-Couch-Full-Commit': 'true'
    })
    if response.status_code == 201:
        results = response.json()
        for result in results:
            if result['error']:
                message = 'Failed to import document {}: {}. Reason: {}\n'.format(
                    result['id'], result['error'], result['reason'])
                sys.stderr.write(message)
                sys.exit(1)
    else:
        sys.stderr.write('Failed to import documents into {}\n'.format(args.target_database))
        sys.exit(1)

def main():
    args = parse_args()
    load(args, sys.stdin, bulk_size=args.bulk_size)

main()


#This is the actual version
#Fully fuctionnal but response is wrong

import socket, sys,time, base64, zlib
m = 0
#Informations du serveur
serveur = "challenge01.root-me.org"
port = 52022

s = socket.socket(socket.AF_INET,socket.SOCK_STREAM)

#connexion
try:
    s.connect((serveur,port))
except socket.error:
    print("Connexion échoué")
    sys.exit()
print("Connexion réussi")

#Début challenge

msg_serv1 = s.recv(2048).decode()
msg_parts1 = msg_serv1.split()
msg_parts1 = msg_parts1[16]
c= base64.b64decode(msg_parts1)
a= (zlib.decompress(c)).decode()
msg2 = f"{a}\n"
while m < 5:
    s.send(msg2.encode())
    m+=1
    msg_serv2 = s.recv(2048).decode()
    print(msg_serv2)


import requests
import re

response = requests.get('http://textfiles.com/etext/FICTION/2city11.txt')
text = (response.text).lower()
wordlist = re.findall(r'\w+\'*\w*', text)
word_dict= {}
for word in wordlist:
    if word not in word_dict:
        word_dict[word] = 1
    else:
        word_dict[word] += 1
word_order = {}
for i in range(len(wordlist)-1):
    if wordlist[i] + ' ' + wordlist[i+1] not in word_order:
        word_order[wordlist[i] + ' ' + wordlist[i+1]] = 1
    else:
        word_order[wordlist[i] + ' ' + wordlist[i+1]] += 1

def print_word_count(word_dict):
    # word_dict is a dictionary where the key is the word and the value is the count
    words = list(word_dict.items()) # .items() returns a list of tuples
    words.sort(key=lambda tup: tup[1], reverse=True)  # sort largest to smallest, based on count
    for i in range(min(10, len(words))):  # print the top 10 words, or all of them, whichever is smaller
        print(words[i])
print_word_count(word_dict)
print("list of common pairs of words")
print_word_count(word_order)
while True:
    userword=input('input a word you want to find the most common following words for: ')
    if userword in wordlist:
        break
    print('That word (or nonword input) is not in the text.')
word_follow={}
for i in range(len(wordlist)-1):
    if wordlist[i] == userword and wordlist[i+1] not in word_follow:
        word_follow[wordlist[i+1]] = 1
    elif wordlist[i] == userword:
        word_follow[wordlist[i+1]] += 1
print_word_count(word_follow)


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

"""
    capture_frame.py

    Usage: python capture_frame.py {video file} {output dir}

    動画ファイルからいくつかのフレームをキャプチャして指定のディレクトリに画像ファイルを保存する。

"""

import configparser
from logging import getLogger, DEBUG
import sys

import cv2

# --------------------------------------------------------------------------
# Preprocessing
# --------------------------------------------------------------------------
logger = getLogger(__name__)
logger.setLevel(DEBUG)

# --------------------------------------------------------------------------
# constants
# --------------------------------------------------------------------------
# OpenCV
CV_CAP_PROP_POS_FRAMES = 1
CV_CAP_PROP_FRAME_WIDTH = 3
CV_CAP_PROP_FRAME_HEIGHT = 4
CV_CAP_PROP_FPS = 5
CV_CAP_PROP_FRAME_COUNT = 7

# 動画からキャプチャするフレーム数
CAPUTURE_COUNT = 9

# 保存する画像のwidth
CAPTURE_WIDTH = 200

# メタ情報の保存ファイル
METAFILE = 'meta.dat'

# --------------------------------------------------------------------------
# functions
# --------------------------------------------------------------------------
def capture_frame(video_file, output_dir):
    """動画ファイルから画像をキャプチャする

    :param video_file 動画ファイル
    :param output_dir キャプチャしたフレームを保存するディレクトリ

    :return 結果コード
        -1 error
    """
    logger.debug("start capture_frame ... ")

    capture = cv2.VideoCapture(video_file)
    if not capture.isOpened():
        logger.error("ファイルが開けませんでした")
        return -1

    # 動画の情報を取得する
    fps = capture.get(CV_CAP_PROP_FPS)
    resize_height = int(CAPTURE_WIDTH / capture.get(CV_CAP_PROP_FRAME_WIDTH) * capture.get(CV_CAP_PROP_FRAME_HEIGHT))
    all_frame_cnt = capture.get(CV_CAP_PROP_FRAME_COUNT)
    interval = all_frame_cnt / (CAPUTURE_COUNT + 1)
    all_time = all_frame_cnt / fps

    # キャプチャ
    capture_frame_list = list(map(lambda pos: interval * pos, range(1, CAPUTURE_COUNT + 1)))
    for i, capture_frame in enumerate(capture_frame_list):
        capture.set(CV_CAP_PROP_POS_FRAMES, capture_frame - 1)
        ret, frame = capture.read()
        if not ret:
            continue

        resized_frame = cv2.resize(frame, (CAPTURE_WIDTH, resize_height))
        cv2.imwrite(output_dir + "/" + str(i) + ".jpg", resized_frame)

    # メタ情報を保存する
    config = configparser.ConfigParser()
    config['metadata'] = {'filename=':video_file, 'time':str(all_time)}
    with open(output_dir + '/' + METAFILE, 'w') as configfile:
        config.write(configfile)

    logger.debug("finish capture_frame ... ")
    return 0

# --------------------------------------------------------------------------
# main
# --------------------------------------------------------------------------
if __name__ == "__main__":
    capture_frame(sys.argv[1], sys.argv[2])


import sys
sys.path.insert(0,'../tt100k/fpgan/code/caffe/python')
import caffe
import os
import flask
from flask import request
from flask import Response
import time
import logging
import json
import datetime
import werkzeug
import optparse
import tornado.wsgi
import tornado.httpserver
from PIL import Image
from io import BytesIO
import exifutil
import base64
import random
import cv2
import skimage.io as io
import numpy as np
import anno_func
import pylab as pl

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

UPLOAD_FOLDER = "/tmp/tt100k_detection_uploads"
DETECTED_IMAGE_OUTPUT_WIDTH = 512
TT100K_IMAGE_SIZE = 2048

app = flask.Flask(__name__)
def resize_image(image, width=None, height=None, inter=cv2.INTER_AREA):
    dim = None
    (h, w) = image.shape[:2]

    if width is None and height is None:
        return image

    if width is None:
        r = height / float(h)
        dim = (int(w * r), height)
    else:
        r = width / float(w)
        dim = (width, int(h * r))

    resized = cv2.resize(image, dim, interpolation = inter)

    return resized

def embed_image_html(image):
    """Creates an image embedded in HTML base64 format."""
    image_output_width = DETECTED_IMAGE_OUTPUT_WIDTH

    if image.shape[1] < image_output_width:
       image_output_width = image.shape[1]

    image_pil = image
    #image_pil = resize_image(image, width=image_output_width)
    #image_pil = Image.fromarray((255 * image).astype('uint8'))
    image_pil = Image.fromarray((image_pil).astype('uint8'))
    #image_pil = image_pil.resize((256, 256))
    stream = BytesIO()
    image_pil.save(stream, format='png')
    #data = string_buf.getvalue().encode('base64').replace('\n', '')
    data = base64.b64encode(stream.getvalue()).decode('utf-8').replace('\n', '')
    return 'data:image/png;base64,' + data

def http_error_response(error_msg, status_code):
    data = {
        'status_code': status_code,
        'msg': error_msg
    }

    js = json.dumps(data)

    res = Response(js, status=status_code, mimetype='application/json')
    logger.error(error_msg)
    return res

def http_success_response(message, result):
    data = {
        'status_code': 200,
        'msg' : message,
        'result' : result
    }

    js = json.dumps(data)

    res = Response(js, status=200, mimetype='application/json')
    return res

class ObjectDetector(object):
    def __init__(self, model_path, weights_path, annotation_path, binaryproto_path, gpu_mode=True, gpu_id=0):
        logging.info('Loading Caffe and associated files.')

        self.model_path       = model_path
        self.weights_path     = weights_path
        self.annotation_path  = annotation_path
        self.binaryproto_path = binaryproto_path
        self.gpu_mode         = gpu_mode

        self.annos = json.loads(open(self.annotation_path).read())

        if self.gpu_mode:
           caffe.set_mode_gpu()
        else:
           caffe.set_mode_cpu()

        caffe.set_device(gpu_id)
        logging.info('Using GPU ID: ', gpu_id)
        self.net = caffe.Net(model_path, weights_path, caffe.TEST)

        logging.info('Loading test_mean.binaryproto file.')
        self.mean = caffe.proto.caffe_pb2.BlobProto.FromString(open(self.binaryproto_path).read())
        self.mn = np.array(self.mean.data)
        self.mn = self.mn.reshape(self.mean.channels, self.mean.height, self.mean.width)
        self.mn = self.mn.transpose((1,2,0))

    def detect(self, image_file_path):
        logging.info('Running detect.')

        imgdata = pl.imread(image_file_path)

        if imgdata.shape[1] != TT100K_IMAGE_SIZE or imgdata.shape[0] != TT100K_IMAGE_SIZE:
           imgdata = cv2.resize(imgdata, (TT100K_IMAGE_SIZE, TT100K_IMAGE_SIZE))

        if imgdata.max() > 2:
           imgdata = imgdata/255.

        imgdata = imgdata[:,:,[2,1,0]]*255. - self.mn

        rectmap = {}
        all_rect = []

        logging.info('Running work.')
        anno_func.work(self.net, imgdata, all_rect)
        rectmap[0] = all_rect

        logging.info('Running get_refine_rects.')
        result_annos = anno_func.get_refine_rects(self.annos, rectmap)
        #detected_result = json.dumps(result_annos)
        detected_result = result_annos
        #print(detected_result)
        logging.info(detected_result)

        return detected_result

    '''
        darknet.set_gpu(0)
        self.cfg_path = cfg_path
        self.weights_path = weights_path
        self.metadata_path = metadata_path
        self.net = darknet.load_net(cfg_path, weights_path, 0)
        self.meta = darknet.load_meta(metadata_path)

    def detect(self, image_file_path):
        detected_result = darknet.detect_objects(self.net, self.meta, self.cfg_path, self.weights_path, self.metadata_path, image_file_path)
        return detected_result
    '''
@app.route('/detect_object', methods=['POST'])
def detect_object():
    current_time = time.time()

    json_result = {
        "data": []
    }

    if flask.request.files['submitImageFile'].filename == '':
        return http_error_response("There is no image file.", 412)

    try:
        imagefile = flask.request.files['submitImageFile']

        filename_ = str(datetime.datetime.now()).replace(' ', '_') + \
                    werkzeug.secure_filename(imagefile.filename)
        filename = os.path.join(UPLOAD_FOLDER, filename_)

        imagefile.save(filename)
        logging.info('Saving to %s.', filename)

        # Convert wrong orientation of uploaded image
        im_arr = exifutil.open_oriented_im(filename)

        image_file_path =  filename.encode('utf-8')
        logging.info('image_file_path is %s.', image_file_path)

        # Detect the objects on uploaded image
        detected_result = app.detector.detect(image_file_path)

        if request.args.get('embed_image') == "true":
           detected_image = showAnnsBBox(im_arr, detected_result)
           json_result["data"].append(dict(detected_embed_image=embed_image_html(detected_image)))

        json_result["data"].append(dict(detected_results=detected_result))

        if os.path.exists(filename):
            os.remove(filename)
            logging.info('Uploaded image file is removed: %s', filename)

    except Exception as err:
        logging.info('Upload the image error: %s', err)

        if os.path.exists(filename):
            os.remove(filename)
            logging.info('Uploaded image file is removed: %s', filename)

        return http_error_response(err, 400)

    elapsed_time = time.time() - current_time
    json_result["data"].append(dict(elapsed_time=elapsed_time))

    return http_success_response("success", json_result)

@app.route('/', methods=['GET'])
def index():
    json_result = {
        "deep_learning_framework":"Caffe, http://caffe.berkeleyvision.org",
        "model":"FPN-CNN model which is trained by tsinghua-tencent 100K dataset"
    }
    return http_success_response("success", json_result)

def start_tornado(app, port=5000):
    http_server = tornado.httpserver.HTTPServer(
        tornado.wsgi.WSGIContainer(app))
    http_server.listen(port)
    logging.info("Tornado server starting on port {}".format(port))
    tornado.ioloop.IOLoop.instance().start()

def start_from_terminal(app):
    """
        Parse command line options and start the server.
        """
    parser = optparse.OptionParser()
    parser.add_option(
        '-d', '--debug',
        help="enable debug mode",
        action="store_true", default=False)
    parser.add_option(
        '-p', '--port',
        help="which port to serve content on",
        type='int', default=5003)
    parser.add_option(
        '-g', '--gpu',
        help="use gpu mode",
        action='store_true', default=True)

    opts, args = parser.parse_args()

    model_path = os.path.join('../tt100k/weights', 'model.prototxt').encode('utf-8')
    weights_path = os.path.join('../tt100k/weights', 'model.caffemodel').encode('utf-8')
    annotation_path = os.path.join('../tt100k/data', 'annotations.json').encode('utf-8')
    binaryproto_path = os.path.join('../tt100k/data/lmdb', 'test_mean.binaryproto').encode('utf-8')

    logging.info('model_path is %s.', model_path)
    logging.info('weights_path is %s.', weights_path)
    logging.info('annotation_path is %s.', annotation_path)
    logging.info('binaryproto_path is %s.', binaryproto_path)

     # Initialize tt100k object detector for warm start
    app.detector = ObjectDetector(model_path, weights_path, annotation_path, binaryproto_path, opts.gpu)

    '''
    dirname = os.path.dirname(os.path.abspath(__file__))
    cfg_path = os.path.join(dirname, 'yolov3-cg.cfg').encode('utf-8')
    weights_path = os.path.join(dirname, 'yolov3-cg_final.weights').encode('utf-8')
    metadata_path = os.path.join(dirname, 'cg.data').encode('utf-8')

    logging.info('cfg_path is %s.', cfg_path)
    logging.info('weights_path is %s.', weights_path)
    logging.info('metadata_path is %s.', metadata_path)

    # Initialize tt100k object detector for warm start
    app.detector = ObjectDetector(cfg_path, weights_path, metadata_path)
    '''
    if opts.debug:
        app.run(debug=True, host='0.0.0.0', port=opts.port)
    else:
        start_tornado(app, opts.port)

if __name__ == '__main__':
    if not os.path.exists(UPLOAD_FOLDER):
        os.makedirs(UPLOAD_FOLDER)

    start_from_terminal(app)


# coding=UTF-8
#!/usr/bin/python

import time
from BB8 import bb8client
from Car import carclient
import json
import traceback

from SimpleWebSocketServer import SimpleWebSocketServer, WebSocket

class SimpleEcho(WebSocket):

    def handleMessage(self):
        # echo message back to client
        print(self.data)
        data = json.loads(self.data)
        if data["name"] == 'bb8':
            if data["type"] == 'connect':
                if bb8client.get():
                    self.sendMessage('BB8准备就绪')
                else:
                    if bb8client.reconnect():
                            self.sendMessage('BB8准备就绪')
                    else:
                            self.sendMessage('BB8准备失败，请重新启动')
            if data["type"] == 'close':
                bb8client.colse()
            if data["type"] == 'control':
                if(data["angle"] != -1):
                    bb8client.roll(100,data["angle"])
        if data["name"] == 'car':
            if data["type"] == 'connect':
                self.sendMessage('正在打开服务端并等待小车连接')
                carclient.get()
                self.sendMessage('小车准备就绪')
            if data["type"] == 'control':
                if data["angle"] == 0:
                    carclient.run('c')
                if data["angle"] == 90:
                    carclient.run('a')
                if data["angle"] == 180:
                    carclient.run('d')
                if data["angle"] == 270:
                    carclient.run('b')
                if data["angle"] == -1:
                    carclient.run('e')
                print('未识别的指令',data["angle"])

            if data["type"] == 'close':
                carclient.close()
            self.sendMessage('未识别的小车指令')
        self.sendMessage('未识别的指令')

    def handleConnected(self):
        print(self.address, 'connected')

    def handleClose(self):
        print(self.address, 'closed')

server = SimpleWebSocketServer('0.0.0.0', 888, SimpleEcho)
server.serveforever()


import time
import numpy as np
import coremltools as ct
from coremltools.models.neural_network import quantization_utils
import tensorflow as tf

MN = 3
CK = 256
HW = 512
HW_og = 512
num_conv = 50
batch_size = 1

model = tf.keras.Sequential(
    layers=[
        tf.keras.Input(shape=(HW, HW, 3,)),
    ],
    name="model",
)

model_flops = 0

model.add(tf.keras.layers.Conv2D(filters=CK,
                               kernel_size=(MN,MN),
                               activation=None,
                               use_bias=False))
conv_flops = MN * MN * CK * 3 * HW * HW

HW = int((HW - MN + 0) + 1)

for i in range(num_conv-1):
    model.add(tf.keras.layers.Conv2D(filters=CK,
                               kernel_size=(MN,MN),
                               activation=None,
                               use_bias=False))
    HW = int((HW - MN + 0) + 1)
    conv_flops = MN * MN * CK * CK * HW * HW
    model_flops += conv_flops

print("")
print("FLOPs:", model_flops)
print("")

tf_out = model.predict(np.random.rand(batch_size,HW_og,HW_og,3))
print("tf_out", tf_out.shape)

model.summary()

print("Check HW:", HW)

print(model.inputs)

print("Converting model")

mlmodel = ct.convert(model,
                     convert_to="mlprogram",
                     inputs=[ct.TensorType(shape=(batch_size,HW_og,HW_og,3))],
                     compute_precision=ct.precision.FLOAT16,
                     compute_units=ct.ComputeUnit.ALL)

#mlmodel = quantization_utils.quantize_weights(mlmodel, 16)

print("Testing model")

random_input = {"input_1": np.random.rand(batch_size,HW_og,HW_og,3)}

ane_out = mlmodel.predict(random_input)

print("ane_out", ane_out["Identity"].shape)

print("Benchmarking model")

# warmup

ane_out = mlmodel.predict(random_input)

iterations = 30

st = time.time()
for i in range(iterations):
    ane_out = mlmodel.predict(random_input)
et = time.time()

print("ane_out", ane_out["Identity"].shape)

duration = et-st
fps = batch_size*iterations/duration
tfops = fps * model_flops / 1e12

print("TFLOPS", tfops)


#!/usr/bin/env python
""" publisher.py

    A small demonstration application that uses AMQP in blocking mode.

    To use, open two console windows - one for publishing and one
    for receiving.

    Reference:
      https://www.youtube.com/watch?v=BKDoxM9KOJY

    Ni Luo  April 13, 2018
"""

import pika

url = 'amqp://rzbdnysq:al19XcdnQULK3U2xqpox0LlIevEtxKup@otter.rmq.cloudamqp.com/rzbdnysq'
params = pika.URLParameters(url)
connection = pika.BlockingConnection(params)
channel = connection.channel()
channel.basic_publish(exchange = '', routing_key = 'hello', body = 'Hello viewers!')


# -*- coding: utf-8 -*-
import six
from rest_framework.exceptions import NotAcceptable, ParseError
from django.db.models import Max, Min

ORDER_TYPES = {
    'ASC': '',
    'DESC': '-'
}
STRING_TYPES = tuple(
    [six.text_type] + list(six.string_types)
)

def get_field_name(order):
    if order[0] == '-':
        return order[1:], True
    return order, False

class ProcessOrderFilter:
    error_msgs = {
        'invalid_type': "Filter '{property}' should be {expected_types}, got - {type}",
        'malformed_order': "Malformed parameter for 'order' filter. See https://loopback.io/doc/en/lb2/Order-filter.html",
        'invalid_field_name': "Field '{field_name}' for model '{model_name}' does't exists. You can't use 'order' filter"
    }

    def __init__(self, queryset, _order):
        self.order = _order
        self.model_fields = queryset.model._meta.get_fields()
        self.model_name = queryset.model.__name__
        self.queryset = queryset
        self.order_field = None

    def filter_queryset(self):
        if not self.order:
            return self.queryset
        order = self.validate_value(self.order)
        fieldname, order_type = get_field_name(order)
        if self.is_order_by_m2m():
            # Ordering by related field creates duplicates in resultant querysets
            # https://code.djangoproject.com/ticket/18165
            # WTF django??
            if order_type:
                return self.queryset.annotate(Max(fieldname)).order_by(order+'__max')
            else:
                return self.queryset.annotate(Min(fieldname)).order_by(order+'__min')
        return self.queryset.order_by(order)

    def is_order_by_m2m(self):
        if not self.order_field:
            return False
        if getattr(self.order_field, 'is_relation', False) and getattr(self.order_field, 'm2m_field_name', None):
            return True
        return False

    def validate_value(self, value):
        if not isinstance(value, STRING_TYPES):
            raise NotAcceptable(self.error_msgs['invalid_type'].format(
                property='order',
                expected_types="<type 'str'>",
                type=type(value)
            ))

        if value.strip() == '?':
            return '?'

        if len(value.split(' ')) != 2:
            raise ParseError(self.error_msgs['malformed_order'])

        order_field, order_type = value.split(' ')
        if order_type not in ORDER_TYPES.keys():
            raise ParseError(self.error_msgs['malformed_order'])

        for p in self.model_fields:
            if p.name == order_field.split('__')[0]:
                self.order_field = p
                break

        if not self.order_field:
            raise ParseError(self.error_msgs['invalid_field_name'].format(
                field_name=order_field,
                model_name=self.model_name
            ))

        return ORDER_TYPES[order_type] + order_field


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time : 2019/1/2 12:59
# @Author : LiangJiangHao
# @Software: PyCharm

import os
import requests
import json
def returnId(idStr):
    header_dict = {'deviceType': '0', 'market': 'appstore', 'appVersion': '4.7.6','udid':'E8131C7D-8AD5-45A0-8E6A-5F97A90CA5A9'}
    url = "https://apipc.app.acfun.cn/v2/videos/%s" % idStr
    # print(url)
    res = requests.get(url, headers=header_dict,verify=False).content
    result = json.loads(res)
    # print(result)
    # f=open('1.txt','wb+')
    # f.write(res)
    # f.close()
    # print(result['vdata']['videos'][0]['videoId'])
    return result['vdata']['videos'][0]['videoId']
    # return str(result['vdata']['channelId']) + '-' + str(result['vdata']['parentChannelId'])
    # return str(result['vdata']['parentChannelId']) + '-' + str(result['vdata']['channelId'])

dmFile=open('danmu.txt','w+',encoding='utf-8')
url='http://www.acfun.cn/v/ac4447301'
# url='http://www.acfun.cn/v/ac1282697_2'
url='http://www.acfun.cn/v/ac1282697'
url='http://www.acfun.cn/v/ac1758344'
acid = url.split('/ac')[1]
# print(returnId(acid))
danmuID=returnId(acid)
danmuUrl = 'http://danmu.aixifan.com/V4/%s/0/500' % danmuID
# print(danmuUrl)
response = requests.get(danmuUrl,verify=False).content
jsonData = json.loads(response)
print(jsonData)
for danmu in jsonData:
    if len(danmu) > 0:
        for dm in danmu:
            print(dm['m'])
            dmFile.write(dm['m']+'\n')
dmFile.close()


# %%
from pymocap import MocapTrajectory, MagnetometerData
from typing import List
import matplotlib.pyplot as plt
import seaborn as sns
from rosbags.highlevel import AnyReader
from pathlib import Path
sns.set_theme(style="whitegrid")
filename = "data/imu_calib.bag"
agent = "ifo002"

# Extract data
mocap = MocapTrajectory.from_bag(filename, agent)
mag = MagnetometerData.from_bag(filename, f"/{agent}/mavros/imu/mag")

D, b, m_a = mag.calibrate(mocap)
mag_calib = mag.apply_calibration(D, b)
# %% Plotting
fig, ax = plt.subplots(1, 1)
m = mag.magnetic_field
ax.scatter(m[:, 0], m[:, 1], label="x vs y")
ax.scatter(m[:, 1], m[:, 2], label="y vs z")
ax.scatter(m[:, 0], m[:, 2], label="x vs z")
ax.legend()
ax.set_title("Raw data")
ax.axis("equal")

fig, ax = mag_calib.plot(mocap, mag_vector=m_a)
fig, ax = plt.subplots(1, 1)
m = mag_calib.magnetic_field
ax.scatter(m[:, 0], m[:, 1], label="x vs y")
ax.scatter(m[:, 1], m[:, 2], label="y vs z")
ax.scatter(m[:, 0], m[:, 2], label="x vs z")
ax.axis("equal")
ax.legend()
ax.set_title("Calibrated")
plt.show()


#!/usr/bin/env python3
from jinja2 import Environment, FileSystemLoader
from math import sin, cos, pi

"""
Compile templates into URDF robot description
"""

templates = ["hexapod.urdf", "one_leg.urdf"]

loader = FileSystemLoader(searchpath="./src/")
env = Environment(loader=loader, autoescape=True)

if __name__ == "__main__":
    # Load and compile Jinja2 templates when executed
    for t in templates:
        template = env.get_template(t + ".j2")
        with open(t, "w") as f:
            f.write(template.render(sin=sin, cos=cos, pi=pi))