from collections import OrderedDict

import networkx as nx

import dwave_networkx as dnx

from neal import SimulatedAnnealingSampler
from dwave.system.composites import EmbeddingComposite
from dwave.system.samplers import DWaveSampler
import dwave.cloud.exceptions

from dwave_structural_imbalance_demo.mmp_network import global_signed_social_network

# compatibility for python 2/3
if dnx._PY2:
    def iteritems(d): return d.iteritems()
else:
    def iteritems(d): return d.items()

class GlobalSignedSocialNetwork(object):
    """A class encapsulating access to graphs from the Stanford Militants Mapping Project.

    Args:
        qpu (bool, optional):
            Specifies whether structural imblance problems will be solved on the QPU or CPU. Defaults to True if
            dwave-system is installed, False otherwise.

    Examples:
        >>> import dwave_structural_imbalance_demo as sbdemo
        >>> gssn = sbdemo.GlobalSignedSocialNetwork()
        >>> nld_before = gssn.get_node_link_data('Syria', 2013)
        >>> nld_before['nodes'][0]
        {'id': 1, 'map': 'Aleppo'}
        >>> nld_before['links'][0]
        {'event_description': 'Ahrar al-Sham and the Islamic State coordinated an attack on Alawite villages in the Latakia governorate that killed 190 civilians.',
         'event_id': '1821',
         'event_type': 'all',
         'event_year': 2013,
         'sign': 1,
         'source': 1,
         'target': 523}
        >>> nld_after = gssn.solve_structural_imbalance('Syria', 2013)
        >>> nld_after['nodes'][0]
        {'color': 0, 'id': 1, 'map': 'Aleppo'}
        >>> nld_after['links'][0]
        {'event_description': 'Ahrar al-Sham and the Islamic State coordinated an attack on Alawite villages in the Latakia governorate that killed 190 civilians.',
         'event_id': '1821',
         'event_type': 'all',
         'event_year': 2013,
         'frustrated': False,
         'sign': 1,
         'source': 1,
         'target': 523}

    """

    def __init__(self, qpu):
        maps = dict()
        maps['Global'] = global_signed_social_network()

        # The Syria subregion
        syria_groups = set()
        for v, data in maps['Global'].nodes(data=True):
            if 'map' not in data:
                continue
            if data['map'] in {'Syria', 'Aleppo'}:
                syria_groups.add(v)
        maps['Syria'] = maps['Global'].subgraph(syria_groups)

        # The Iraq subregion
        iraq_groups = set()
        for v, data in maps['Global'].nodes(data=True):
            if 'map' not in data:
                continue
            if data['map'] == 'Iraq':
                iraq_groups.add(v)
        maps['Iraq'] = maps['Global'].subgraph(iraq_groups)

        self._maps = maps
        self._qpu = qpu
        self._init_sampler()

    def _init_sampler(self):
        """Allows for re-init in case a solver goes offline."""

        if self._qpu:
            # select the first available sampler in the `DW_2000Q` class
            self._sampler = EmbeddingComposite(DWaveSampler(solver=dict(qpu=True)))
        else:
            self._sampler = SimulatedAnnealingSampler()

        self._sampler_args = {}
        if 'num_reads' in self._sampler.parameters:
            self._sampler_args['num_reads'] = 50
        if 'answer_mode' in self._sampler.parameters:
            self._sampler_args['answer_mode'] = 'histogram'
        if 'chain_strength' in self._sampler.parameters:
            self._sampler_args['chain_strength'] = 2.0

    def _get_graph(self, subregion='Global', year=None):
        G = self._maps[subregion]
        if year is not None:
            if not isinstance(year, int):
                raise ValueError("year must be int")
            filtered_edges = ((u, v) for u, v, a in G.edges(data=True) if a['event_year'] <= year)
            G = G.edge_subgraph(filtered_edges)
        return G

    def get_node_link_data(self, subregion='Global', year=None):
        """Accessor for Stanford Militants Mapping Project node link data.

        Args:
            subregion (str, optional):
                Filter graph by subregion. One of ['Global', 'Syria', 'Iraq']. Defaults to 'Global' (entire network).
            year (int, optional):
                Filter graph by year. Returns only events in or before year. Defaults to None (no filter applied).

        Returns:
            A dictionary with node-link formatted data. Conforms to dwave_structural_imbalance_demo.json_schema.

        """

        G = self._get_graph(subregion, year)
        return {"results": [nx.node_link_data(G)]}

    def solve_structural_imbalance(self, subregion='Global', year=None):
        """Solves specified Stanford Militants Mapping Project structural imbalance problem and returns annotated graph.

        If self._qpu is True (set during object initialization), this function will first attempt to embed the entire
        problem on the hardware graph using EmbeddingComposite. Failing this, it will fallback on QBSolv to decompose
        the problem. If self._qpu is False, this function will use ExactSolver for problems with less than 20 nodes.
        For problems with 20 more more nodes, it will use QBSolv to solve the problem classically.

        Args:
            subregion (str, optional):
                Filter graph by subregion. One of ['Global', 'Syria', 'Iraq']. Defaults to 'Global' (entire network).
            year (int, optional):
                Filter graph by year. Returns only events in or before year. Defaults to None (no filter applied).

        Returns:
            A dictionary with node-link formatted data. Conforms to dwave_structural_imbalance_demo.json_schema.
            Optional property 'color' is set for each item in 'nodes'. Optional property 'frustrated' is set for each
            item in 'links'.

        """

        G_in = self._get_graph(subregion, year)
        if len(G_in) == 0:
            raise ValueError("Filtered network has no nodes to solve problem on")

        h, J = dnx.social.structural_imbalance_ising(G_in)

        # <10% of the time it will fail to find an embedding, so keep trying
        while True:
            try:
                # use the sampler to find low energy states
                response = self._sampler.sample_ising(h, J, **self._sampler_args)
                break
            except ValueError:
                pass
            except dwave.cloud.exceptions.SolverOfflineError:
                # if solver goes offline while sampling (or while in queue),
                # retry with another (online) solver
                self._init_sampler()

        # histogram answer_mode should return counts for unique solutions
        if 'num_occurrences' not in response.data_vectors:
            response.data_vectors['num_occurrences'] = [1] * len(response)

        # should equal num_reads
        total = sum(response.data_vectors['num_occurrences'])

        results_dict = OrderedDict()

        for sample, num_occurrences in response.data(['sample', 'num_occurrences']):
            # spins determine the color
            colors = {v: (spin + 1) // 2 for v, spin in iteritems(sample)}

            key = tuple(colors.values())
            if key in results_dict:
                results_dict[key].graph["numOfOccurrences"] += num_occurrences
                results_dict[key].graph["percentageOfOccurrences"] = 100 * \
                    results_dict[key].graph["numOfOccurrences"] / total
            else:
                G = G_in.copy()
                # frustrated edges are the ones that are violated
                frustrated_edges = {}
                for u, v, data in G.edges(data=True):
                    sign = data['sign']
                    if sign > 0 and colors[u] != colors[v]:
                        frustrated_edges[(u, v)] = data
                    elif sign < 0 and colors[u] == colors[v]:
                        frustrated_edges[(u, v)] = data
                    # else: not frustrated or sign == 0, no relation to violate
                for edge in G.edges:
                    G.edges[edge]['frustrated'] = edge in frustrated_edges
                for node in G.nodes:
                    G.nodes[node]['color'] = colors[node]
                G.graph['numOfOccurrences'] = num_occurrences
                G.graph['percentageOfOccurrences'] = 100 * num_occurrences / total
                results_dict[key] = G

        output = {'results': [nx.node_link_data(result) for result in results_dict.values()], 'numberOfReads': total}
        if 'timing' in response.info:
            output['timing'] = {"actual": {"qpuProcessTime": response.info['timing']['qpu_access_time']}}
        return output


import json

from django.http.response import HttpResponse
from django.shortcuts import render

from CRM.Core.BaseCrmManager import RequestProcessException
from CRM.Core.SystemGroupManagement import VIPGroupManagement
from CRM.Decorators.Permission import multi_check
from CRM.context_processors.Utils import check_ajax
from CRM.models import ServiceGroups

__author__ = 'FAM10'

@multi_check(need_staff=True, perm='CRM.view_vip_group', methods=('GET',))
def view_vip_groups(request):
    if not check_ajax(request):
        groups = ServiceGroups.objects.filter(is_deleted=False)
        return render(request, 'vip/VIPGroupManagement.html', {'groups': groups})
    gm = VIPGroupManagement(request)
    return HttpResponse(gm.get_all())

@multi_check(need_staff=True, perm='CRM.add_vipgroups', methods=('POST',), disable_csrf=True)
def add_new_vip_group(request):
    vm = VIPGroupManagement(request)
    vm.set_post()
    try:
        vm.update()
        return HttpResponse('200')
    except RequestProcessException as e:
        return e.get_response()

@multi_check(True, True, True, False, perm='CRM.delete_vipgroups', methods=('GET',))
def delete_vip_group(request):
    try:
        vm = VIPGroupManagement(request)
        vm.delete()
        return HttpResponse('200')
    except RequestProcessException as e:
        return e.get_response()

@multi_check(need_staff=True, perm='CRM.view_vip_group', methods=('GET',))
def get_vip_group_detail(request):
    try:
        vm = VIPGroupManagement(request)
        x = vm.get_single_pk(True)
        return HttpResponse(json.dumps({'name': x.name, 'group_id': x.group_id, 'pk': x.pk}))
    except RequestProcessException as e:
        return e.get_response()


"""
    brief:
        - post-processes and plots the results of the training using PPO, either for MF-DRL or MB-DRL
        - plots the results of the controlled case using the best policy in comparison to the uncontrolled case

    dependencies:
        - 'ppo_data_loader.py' for handling the loading, sorting and merging of all training data
        - 'analyze_frequency_spectrum.py' for plotting the frequency spectrum of cl- and cd of PPO-training and the
           final results

    prerequisites:
        - execution of the 'run_training.py' function in the 'test_training' directory in order to conduct a training
          and generate trajectories within the CFD environment (https://github.com/OFDataCommittee/drlfoam)
        - this training can either be model-free or model-based

    optional:
        - execution of simulation for the best policy from training, also results of a simulation without control
        - in this case, the results of the training are not averaged over multiple seed values, since the final policy
          corresponds to a specific training of one seed value
"""
import re
from os.path import join

import pandas as pd
import matplotlib.pyplot as plt

from typing import Union
from os import mkdir, path
from matplotlib.patches import Circle, Rectangle

from ppo_data_loader import *
from analyze_frequency_spectrum import analyze_frequencies_final_result, analyze_frequencies_ppo_training,\
    analyze_frequencies_probes_final_result

def plot_coefficients_vs_episode(settings: dict, cd_mean: Union[list, pt.Tensor], cd_std: Union[list, pt.Tensor],
                                 cl_mean: Union[list, pt.Tensor], cl_std: Union[list, pt.Tensor],
                                 actions_mean: Union[list, pt.Tensor] = None,
                                 actions_std: Union[list, pt.Tensor] = None,
                                 n_cases: int = 1, plot_action: bool = False,
                                 ylabel: list = ["$\\bar{c}_L$", "$\\bar{c}_D$", "$\\bar{\omega}$"]) -> None:
    """
    plot cl, cd and actions (if specified) depending on the episode (training)

    :param settings: dict containing all the paths etc.
    :param cd_mean: mean cd received over the training periode
    :param cd_std: corresponding standard deviation of cd throughout the training periode
    :param cl_mean: mean cl received over the training periode
    :param cl_std: corresponding standard deviation of cl throughout the training periode
    :param actions_mean: mean actions (omega) done over the training periode
    :param actions_std: corresponding standard deviation of the actions done over the training periode
    :param n_cases: number of cases to compare (= number of imported data)
    :param plot_action: if 'True' cl, cd and actions will be plotted, otherwise only cl and cd will be plotted
    :param ylabel: ylabels for plots [cl, cd, action]
    :return: None
    """
    if plot_action:
        fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(6, 3))
        n_subfig = 3
    else:
        fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(6, 3))
        n_subfig = 2

    for c in range(n_cases):
        for i in range(n_subfig):
            if i == 0:
                ax[i].plot(range(len(cl_mean[c])), cl_mean[c], color=settings["color"][c], label=settings["legend"][c])
                ax[i].fill_between(range(len(cl_mean[c])), cl_mean[c] - cl_std[c], cl_mean[c] + cl_std[c],
                                   color=settings["color"][c], alpha=0.3)
                ax[i].set_ylabel(ylabel[0])

            elif i == 1:
                ax[i].plot(range(len(cd_mean[c])), cd_mean[c], color=settings["color"][c])
                ax[i].fill_between(range(len(cd_mean[c])), cd_mean[c] - cd_std[c], cd_mean[c] + cd_std[c],
                                   color=settings["color"][c], alpha=0.3)
                ax[i].set_ylabel(ylabel[1])

            elif plot_action:
                ax[i].plot(range(len(actions_mean[c])), actions_mean[c], color=settings["color"][c])
                ax[i].fill_between(range(len(actions_mean[c])), actions_mean[c] - actions_std[c],
                                   actions_mean[c] + actions_std[c], color=settings["color"][c], alpha=0.3)
                ax[i].set_ylabel(ylabel[2])

            ax[i].set_xlabel("$e$")

    fig.tight_layout()
    fig.legend(loc="upper center", framealpha=1.0, ncol=3)
    fig.subplots_adjust(wspace=0.35, top=0.88)
    plt.savefig(join(settings["main_load_path"], settings["path_controlled"], "plots", "coefficients_vs_episode.png"),
                dpi=340)
    plt.show(block=False)
    plt.pause(2)
    plt.close("all")

def plot_rewards_vs_episode(settings: dict, reward_mean: Union[list, pt.Tensor], reward_std: Union[list, pt.Tensor],
                            n_cases: int = 0) -> None:
    """
    plots the mean rewards received throughout the training periode and the corresponding standard deviation

    :param settings: dict containing all the paths etc.
    :param reward_mean: mean rewards received over the training periode
    :param reward_std: corresponding standard deviation of the rewards received over the training periode
    :param n_cases: number of cases to compare (= number of imported data)
    :return: None
    """
    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(12, 6))
    for c in range(n_cases):
        ax.plot(range(len(reward_mean[c])), reward_mean[c], color=settings["color"][c], label=settings["legend"][c])
        ax.fill_between(range(len(reward_mean[c])), reward_mean[c] - reward_std[c], reward_mean[c] + reward_std[c],
                        color=settings["color"][c], alpha=0.3)

    ax.set_ylabel("$\\bar{r}$")
    ax.set_xlabel("$e$")
    ax.set_xlim(0, max([len(i) for i in reward_mean]))
    fig.tight_layout()
    ax.legend(loc="lower right", framealpha=1.0, ncol=2)
    fig.subplots_adjust(wspace=0.2)
    plt.savefig(join(settings["main_load_path"], settings["path_controlled"], "plots", "rewards_vs_episode.png"),
                dpi=340)
    plt.show(block=False)
    plt.pause(2)
    plt.close("all")

def plot_cl_cd_alpha_beta(settings: dict, controlled_cases: Union[list, pt.Tensor],
                          uncontrolled_case: Union[list, pt.Tensor] = None, plot_coeffs=True, factor: int = 10) -> None:
    """
    plot either cl and cd vs. time or alpha and beta vs. time

    :param settings: dict containing all the paths etc.
    :param controlled_cases: results from the loaded cases with active flow control
    :param uncontrolled_case: reference case containing results from uncontrolled flow past cylinder
    :param plot_coeffs: 'True' means cl and cd will be plotted, otherwise alpha and beta will be plotted wrt to time
    :param factor: factor for converting the physical time into non-dimensional time, t^* = u * t / d
    :return: None
    """
    fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(14, 6))

    if plot_coeffs:
        keys = ["t", "cl", "cd"]
        save_name = "comparison_cl_cd"
        ax[1].set_ylim(2.95, 3.25)  # Re = 100
        # ax[1].set_ylim(2, 3.75)          # Re = 500
        n_cases = range(len(settings["case_name"]) + 1)
        ylabels = ["$c_L$", "$c_D$"]
        x_min = 0
    else:
        keys = ["t", "alpha", "beta"]
        save_name = "comparison_alpha_beta"
        ylabels = ["$\\alpha$", "$\\beta$"]
        n_cases = range(1, len(settings["case_name"]) + 1)
        x_min = 4 * factor      # control starts at t = 4s, so there are no alpha & beta available for t < 4s

    for c in n_cases:
        for i in range(2):
            if i == 0:
                if c == 0:
                    ax[i].plot(uncontrolled_case[keys[0]] * factor, uncontrolled_case[keys[1]], color="black",
                               label="uncontrolled")
                else:
                    ax[i].plot(controlled_cases[c - 1][keys[0]] * factor, controlled_cases[c - 1][keys[1]],
                               color=settings["color"][c - 1], label=settings["legend"][c - 1])
                ax[i].set_ylabel(ylabels[0], fontsize=13)
            else:
                if c == 0:
                    ax[i].plot(uncontrolled_case[keys[0]] * factor, uncontrolled_case[keys[2]], color="black")
                else:
                    ax[i].plot(controlled_cases[c - 1][keys[0]] * factor, controlled_cases[c - 1][keys[2]],
                               color=settings["color"][c - 1])
                ax[i].set_ylabel(ylabels[1], fontsize=13)

            ax[1].set_xlabel("$t^*$", fontsize=14)
            ax[i].set_xlim(x_min, controlled_cases[0]["t"].iloc[-1] * factor)
    fig.tight_layout()
    fig.legend(loc="upper center", framealpha=1.0, fontsize=10, ncol=2)
    fig.subplots_adjust(wspace=0.2, top=0.84)
    plt.savefig(join(settings["main_load_path"], settings["path_controlled"], "plots", f"{save_name}.png"), dpi=340)
    plt.show(block=False)
    plt.pause(2)
    plt.close("all")

def plot_omega(settings: dict, controlled_cases: Union[list, pt.Tensor], factor: int = 10) -> None:
    """
    plot omega (actions) vs. time

    :param settings: dict containing all the paths etc.
    :param controlled_cases: results from the loaded cases with active flow control
    :param factor: factor for converting the physical time into non-dimensional time, t^* = u * t / d
    :return: None
    """
    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(8, 5))
    for c in range(len(settings["case_name"])):
        ax.plot(controlled_cases[c]["t"] * factor, controlled_cases[c]["omega"], color=settings["color"][c],
                label=settings["legend"][c])

    ax.set_ylabel("$\omega$", fontsize=13)
    ax.set_xlabel("$t^*$", fontsize=13)
    fig.tight_layout()
    fig.subplots_adjust(top=0.91)
    plt.legend(loc="upper right", framealpha=1.0, ncol=1)
    plt.savefig(join(settings["main_load_path"], settings["path_controlled"], "plots", "omega_controlled_case.png"),
                dpi=340)
    plt.show(block=False)
    plt.pause(2)
    plt.close("all")

def plot_variance_of_beta_dist(settings: dict, var_beta_dist: Union[list, pt.Tensor], n_cases: int = 0) -> None:
    """
    plots the mean rewards received throughout the training periode and the corresponding standard deviation

    :param settings: dict containing all the paths etc.
    :param var_beta_dist: computed variance of the beta-function wrt episode
    :param n_cases: number of cases to compare (= number of imported data)
    :return: None
    """
    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(12, 6))
    for c in range(n_cases):
        ax.plot(range(len(var_beta_dist[c])), var_beta_dist[c], color=settings["color"][c], label=settings["legend"][c])

    ax.set_ylabel("$mean$ $variance$ $of$ $beta-distribution$", fontsize=13)
    ax.set_xlabel("$e$", fontsize=13)
    ax.legend(loc="upper right", framealpha=1.0, fontsize=10, ncol=2)
    plt.savefig(join(settings["main_load_path"], settings["path_controlled"], "plots", "var_beta_distribution.png"),
                dpi=340)
    plt.show(block=False)
    plt.pause(2)
    plt.close("all")

def plot_cl_cd_trajectories(settings: dict, data: list, number: int, e: int = 1, factor: int = 10) -> None:
    """
    plots the trajectory of cl and cd for different episodes of the training, meant to use for either comparing MF-
    trajectories to trajectories generated by the environment models or comparing trajectories from environment models
    run with different settings to each other

    :param settings: setup containing all the paths etc.
    :param data: trajectory data to plot
    :param number: number of the trajectory within the data set (either within the episode or in total)
    :param e: episode number
    :param factor: factor for converting the physical time into non-dimensional time, t^* = u * t / d
    :return: None
    """
    fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(15, 5))
    epochs = pt.tensor(list(range(len(data[0]["cl"][1, :, number])))) / factor
    for n in range(len(data)):
        for i in range(3):
            try:
                if i == 0:
                    # 2nd episode is always MB (if MB-DRL was used)
                    ax[i].plot(epochs, data[n]["cl"][e, :, number], color=settings["color"][n],
                               label=f"{settings['legend'][n]}, episode {e + 1}")
                    ax[i].set_ylabel("$c_L$", fontsize=13)
                elif i == 1:
                    ax[i].plot(epochs, data[n]["cd"][e, :, number], color=settings["color"][n])
                    ax[i].set_ylabel("$c_D$", fontsize=13)
                else:
                    ax[i].plot(epochs, data[n]["actions"][e, :, number], color=settings["color"][n])
                    ax[i].set_ylabel("$omega$", fontsize=13)
                ax[i].set_xlabel("$t^*$", fontsize=13)
            except IndexError:
                print("omit plotting trajectories of failed cases")
    fig.tight_layout()
    fig.legend(loc="upper right", framealpha=1.0, fontsize=10, ncol=2)
    fig.subplots_adjust(wspace=0.25, top=0.90)
    plt.savefig(join(settings["main_load_path"], settings["path_controlled"], "plots",
                     f"comparison_traj_cl_cd_{e}.png"), dpi=340)
    plt.show(block=False)
    plt.pause(2)
    plt.close("all")

def plot_total_reward(settings: dict, reward_mean: list, reward_std: list, n_cases: int) -> None:
    """
    plot the total rewards of the complete training for each case

    :param settings: dict containing all the paths etc.
    :param reward_mean: mean total rewards received in the training
    :param reward_std: corresponding standard deviation of the total rewards received in the training
    :param n_cases: number of cases to compare (= number of imported data)
    :return: None
    """
    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(8, 5))
    for c in range(n_cases):
        ax.errorbar(c + 1, reward_mean[c], yerr=reward_std[c], color=settings["color"][c], fmt="o", capsize=5,
                    label=settings["legend"][c])

    ax.set_ylabel("$total$ $reward$", usetex=True, fontsize=13)
    ax.set_xlabel("$case$ $number$", usetex=True, fontsize=13)
    ax.set_xticks(range(1, n_cases + 1, 1))
    ax.legend(loc="upper right", framealpha=1.0, fontsize=10, ncol=1)
    plt.grid(which="major", axis="y", linestyle="--", alpha=0.85, color="black", lw=1)
    plt.savefig(join(settings["main_load_path"], settings["path_controlled"], "plots", "total_rewards.png"), dpi=340)
    plt.show(block=False)
    plt.pause(2)
    plt.close("all")

def plot_numerical_setup(settings: dict) -> None:
    """
    plot the domain of the cylinder case

    :param settings: setup containing the import and save path
    :return: None
    """
    pattern = r"\d.\d+ \d.\d+ \d.\d+"
    path = "".join([settings["main_load_path"], settings["path_controlled"], settings["case_name"][0],
                    settings["path_final_results"]])
    with open("".join([path, settings["path_to_probes"]]), "r") as f:
        loc = f.readlines()

    # get coordinates of probes, omit appending empty lists and map strings to floats
    coord = [re.findall(pattern, line) for line in loc if re.findall(pattern, line)]
    pos_probes = pt.tensor([list(map(float, i[0].split())) for i in coord])

    # get coordinates of domain and cylinder
    with open("".join([path, "system/blockMeshDict"]), "r") as f:
        loc = f.readlines()

    # structure in blockMeshDict always the same: [lengthX 2.2, lengthY 0.41, cylinderX 0.2, cylinderY 0.2, radius 0.05]
    l, h, pos_x, pos_y, r = [float(loc[i].strip(";\n").split()[1]) for i in range(16, 21)]

    # plot cylinder, probe locations and annotate the domain
    fig, ax = plt.subplots(figsize=(10, 3))
    ax.plot(pos_probes[:, 0], pos_probes[:, 1], linestyle="None", marker="o", color="red", label="probes")
    # dummy point for legend
    ax.scatter(-10, -10, marker="o", color="black", alpha=0.4, label="cylinder")
    circle = Circle((pos_x, pos_y), radius=r, color="black", alpha=0.4)
    rectangle = Rectangle((0, 0), width=l, height=h, edgecolor="black", linewidth=2, facecolor="none")
    ax.add_patch(circle)
    ax.add_patch(rectangle)
    fig.legend(loc="lower right", framealpha=1.0, fontsize=10, ncol=2)
    ax.set_xlim(0, l)
    ax.set_ylim(0, h)
    ax.set_xticks([])
    ax.set_yticks([])

    # annotate inlet & outlet
    plt.arrow(-0.05, -0.05, 0.1, 0.0, color="black", head_width=0.02, clip_on=False)
    plt.arrow(-0.05, -0.05, 0.0, 0.1, color="black", head_width=0.02, clip_on=False)
    plt.arrow(-0.1, h * 2 / 3 + 0.025, 0.075, -0.05, color="black", head_width=0.015, clip_on=False)
    plt.arrow(-0.1 + l, h * 2 / 3, 0.075, -0.05, color="black", head_width=0.015, clip_on=False)

    plt.annotate("$inlet$", (-0.17, h * 2 / 3 + 0.05), annotation_clip=False, usetex=True, fontsize=13)
    plt.annotate("$\\frac{x}{d}$", (0.1, -0.065), annotation_clip=False, usetex=True, fontsize=16)
    plt.annotate("$\\frac{y}{d}$", (-0.1, 0.065), annotation_clip=False, usetex=True, fontsize=16)
    plt.annotate("$outlet$", (-0.2 + l, h * 2 / 3 + 0.01), annotation_clip=False, usetex=True, fontsize=13)

    # annotate the dimensions & position of the domain
    pos = {"xy": [(0, h + 0.04), (0, h), (l, h), (pos_x - r - 0.01, pos_y - 0.1), (l, h), (l, 0), (l + 0.04, h),
                  (pos_x, pos_y + 0.9 * r), (0, pos_y)],
           "xytxt": [(l, h + 0.04), (0, h + 0.075), (l, h + 0.075), (pos_x + r + 0.01, pos_y - 0.1), (l + 0.075, h),
                     (l + 0.075, 0),
                     (l + 0.04, 0), (pos_x, h), (pos_x - 0.9 * r, pos_y)],
           "style": [("<->", "-"), ("-", "--"), ("-", "--"), ("<->", "-"), ("-", "--"), ("-", "--"), ("<->", "-"),
                     ("<->", "-"), ("<->", "-")]
           }
    for i in range(len(pos["style"])):
        plt.annotate("", xy=pos["xy"][i], xytext=pos["xytxt"][i],
                     arrowprops=dict(arrowstyle=pos["style"][i][0], color="black", linestyle=pos["style"][i][1]),
                     annotation_clip=False)

    plt.annotate(f"${l / (2 * r)}$", (l / 2, h + 0.07), annotation_clip=False, usetex=True, fontsize=12)
    plt.annotate(f"${h / (2 * r)}$", (l + 0.07, h / 2), annotation_clip=False, usetex=True, fontsize=12)
    plt.annotate("$d$", (pos_x - r / 4, pos_y - 3 * r), usetex=True, fontsize=12)
    plt.annotate("${:.2f}$".format((h - (pos_y + r)) / (2 * r)), (pos_x + 0.025, pos_y + 2.25 * r), usetex=True,
                 fontsize=12)
    plt.annotate("${:.2f}$".format((pos_x - r) / (2 * r)), (pos_x - 3.25 * r, pos_y + 0.5 * r), usetex=True,
                 fontsize=12)

    ax.plot((pos_x - r, pos_x - r), (pos_y, pos_y - 0.15), color="black", linestyle="--", lw=1)
    ax.plot((pos_x + r, pos_x + r), (pos_y, pos_y - 0.15), color="black", linestyle="--", lw=1)
    ax.plot(pos_x, pos_y, marker="+", color="black")

    ax.set_aspect("equal")
    fig.tight_layout()
    fig.subplots_adjust(left=0.1, right=0.92)
    plt.savefig(join(settings["main_load_path"], settings["path_controlled"], "plots", "domain_setup.png"), dpi=340)
    plt.show(block=False)
    plt.pause(2)
    plt.close("all")

def plot_train_validation_loss(settings: dict, mse_train: Union[list, pt.Tensor], mse_val: Union[list, pt.Tensor],
                               mse_train_cd: Union[list, pt.Tensor], mse_val_cd: Union[list, pt.Tensor],
                               std_dev_train: Union[list, pt.Tensor], std_dev_val: Union[list, pt.Tensor],
                               std_dev_train_cd: Union[list, pt.Tensor], std_dev_val_cd: Union[list, pt.Tensor],
                               case: int = 1) -> None:
    """
    plots the avg. train- and validation loss and the corresponding std. deviation of the environment models wrt to
    epochs

    :param settings: path where the plot should be saved
    :param mse_train: tensor containing the (mean) training loss of the cl-p env. model
    :param mse_val: tensor containing the (mean) validation loss of the cl-p env. model
    :param mse_train_cd: tensor containing the (mean) training loss of the cd env. model
    :param mse_val_cd: tensor containing the (mean) validation loss of the cd env. model
    :param std_dev_train: tensor containing the (std. deviation) training loss of the cl-p env. model
    :param std_dev_val: tensor containing the (std. deviation) validation loss of the cl-p env. model
    :param std_dev_train_cd: tensor containing the (std. deviation) training loss of the cd env. model
    :param std_dev_val_cd: tensor containing the (std. deviation) validation loss of the cd env. model
    :param case: name to append for savin img
    :return: None
    """
    fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(14, 6))
    for i in range(2):
        if i == 0:
            x = range(len(mse_train))
            ax[i].plot(x, mse_train, color="blue")
            ax[i].plot(x, mse_val, color="red")
            ax[i].fill_between(x, mse_val - std_dev_val, mse_val + std_dev_val, color="red", alpha=0.3)
            ax[i].fill_between(x, mse_train - std_dev_train, mse_train + std_dev_train, color="blue", alpha=0.3)
            ax[i].set_ylabel("$MSE$ $loss$", usetex=True, fontsize=13)
            ax[i].set_xlabel("$epoch$ $number$", usetex=True, fontsize=13)
            ax[i].set_title("$environment$ $model$ $for$ $c_L$ $\&$ $p_i$", usetex=True, fontsize=14)
            ax[i].set_yscale("log")

        else:
            x = range(len(mse_train_cd))
            ax[i].plot(x, mse_train_cd, color="blue", label="training loss")
            ax[i].plot(x, mse_val_cd, color="red", label="validation loss")
            ax[i].fill_between(x, mse_train_cd - std_dev_train_cd, mse_train_cd + std_dev_train_cd, color="blue",
                               alpha=0.3)
            ax[i].set_xlabel("$epoch$ $number$", usetex=True, fontsize=13)
            ax[i].set_title("$environment$ $model$ $for$ $c_D$", usetex=True, fontsize=14)
            ax[i].set_yscale("log")
            ax[i].set_ylabel("$MSE$ $loss$", usetex=True, fontsize=13)

    ax[1].legend(loc="upper right", framealpha=1.0, fontsize=10, ncol=2)
    fig.tight_layout()
    fig.subplots_adjust(wspace=0.2)
    plt.savefig(join(settings["main_load_path"], settings["path_controlled"], "plots",
                     "train_val_losses_case{case}.png"),
                dpi=340)
    plt.show(block=False)
    plt.pause(2)
    plt.close("all")

def plot_mean_std_trajectories(settings: dict, data: list, factor: int = 10) -> None:
    """
    plots the trajectory of cl and cd for different episodes of the training, meant to use for either comparing MF-
    trajectories to trajectories generated by the environment models or comparing trajectories from environment models
    run with different settings to each other

    :param settings: setup containing all the paths etc.
    :param data: trajectory data to plot
    :param factor: factor for converting the physical time into non-dimensional time, t^* = u * t / d
    :return: None
    """
    fig, ax = plt.subplots(nrows=4, ncols=2, figsize=(6, 8), sharey="col", sharex="all")
    epochs = pt.tensor(list(range(len(data[0]["cl"][1, :, 0])))) / factor
    e = [24, 74, 124, 199]
    for n in range(len(data)):
        for k in range(4):          # k = rows
            for i in range(2):      # i = cols
                if i == 0:
                    mean_tmp = pt.mean(data[n]["cd"][e[k], :, :], dim=1)
                    std_tmp = pt.std(data[n]["cd"][e[k], :, :], dim=1)
                    if k == 0:
                        # 2nd episode is always MB (if MB-DRL was used)
                        ax[k][i].plot(epochs, mean_tmp, color=settings["color"][n], label=settings['legend'][n])
                        ax[k][i].fill_between(epochs, mean_tmp - std_tmp, mean_tmp + std_tmp, color=settings["color"][n],
                                              alpha=0.3)
                    else:
                        ax[k][i].plot(epochs, mean_tmp, color=settings["color"][n])
                        ax[k][i].fill_between(epochs, mean_tmp - std_tmp, mean_tmp + std_tmp, color=settings["color"][n],
                                              alpha=0.3)

                    ax[k][i].set_ylabel("$\\bar{c}_D$")
                else:
                    mean_tmp = pt.mean(data[n]["cl"][e[k], :, :], dim=1)
                    std_tmp = pt.std(data[n]["cl"][e[k], :, :], dim=1)
                    ax[k][i].plot(epochs, mean_tmp, color=settings["color"][n])
                    ax[k][i].fill_between(epochs, mean_tmp - std_tmp, mean_tmp + std_tmp, color=settings["color"][n],
                                          alpha=0.3)
                    if i == 1:
                        ax[k][i].set_ylabel("$\\bar{c}_L$")

                ax[k][i].set_xlim(0, data[0]["cl"].size()[1] / factor)
    ax[-1][0].set_xlabel("$t^*$")
    ax[-1][1].set_xlabel("$t^*$")
    fig.tight_layout()
    fig.legend(loc="upper center", framealpha=1.0, ncol=2)
    fig.subplots_adjust(wspace=0.3, top=0.9)
    plt.savefig(join(settings["main_load_path"], settings["path_controlled"], "plots",
                     "comparison_traj_cd_mean_std.png"), dpi=340)
    plt.show(block=False)
    plt.pause(2)
    plt.close("all")

if __name__ == "__main__":
    # Setup
    setup = {
        "main_load_path": r"/home/janis/Hiwi_ISM/results_drlfoam_MB/",
        "path_to_probes": r"postProcessing/probes/0/p",  # path to the file containing trajectories of probes
        "path_uncontrolled": r"run/uncontrolled_re100/",  # path to uncontrolled reference case
        "path_controlled": r"run/final_routine_AWS/",
        "path_final_results": r"results_best_policy/",  # path to the results using the best policy
        "case_name": ["e200_r10_b10_f8_MF/", "e200_r10_b10_f8_MB_1model/",
                      "e200_r10_b10_f8_MB_5models/", "e200_r10_b10_f8_MB_5models_threshold40/",
                      "e200_r10_b10_f8_MB_10models_threshold50/", "e200_r10_b10_f8_MB_10models_threshold30/"],
        # "case_name": ["e200_r10_b10_f8_MF/seed4/", "e200_r10_b10_f8_MB_1model/seed2/",
        #               "e200_r10_b10_f8_MB_5models/seed1/", "e200_r10_b10_f8_MB_5models_threshold40/seed1/",
        #               "e200_r10_b10_f8_MB_10models_threshold50/seed0/",
        #               "e200_r10_b10_f8_MB_10models_threshold30/seed4/"],
        "e_trajectory": [4, 9, 24, 49, 74, 99, 124, 149, 174, 199],   # episodes trajectories (cl & cd, not avg.)
        "n_probes": 12,  # number of probes placed in flow field
        "avg_over_cases": True,  # if cases should be averaged over, e.g. different seeds
        "plot_final_res": False,  # if the final policy already ran in openfoam, plot the results
        "param_study": False,  # flag if parameter study, only used for generating legend entries automatically
        "mark_e_cfd": False,  # flag if CFD episodes should be marked (in case of avg., 1st seed is taken)
        "color": ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2', '#7f7f7f',
                  '#bcbd22', '#17becf'],  # default color cycle
        "legend": ["MF", "MB, $N_{m} = 1$", "MB, $N_{m} = 5, N_{thr} = 3$", "MB, $N_{m} = 5, N_{thr} = 2$",
                   "MB, $N_{m} = 10, N_{thr} = 5$", "MB, $N_{m} = 10, N_{thr} = 3$"]
    }

    # create directory for plots
    if not path.exists(join(setup["main_load_path"], setup["path_controlled"], "plots")):
        mkdir(join(setup["main_load_path"], setup["path_controlled"], "plots"))

    # use latex fonts
    plt.rcParams.update({"text.usetex": True})

    # load all the data
    all_data = load_all_data(setup)

    # average the trajectories episode-wise
    averaged_data = average_results_for_each_case(all_data)

    # for parameter study: generate legend entries automatically
    if setup["param_study"]:
        if setup["avg_over_cases"]:
            setup["legend"] = [f"b = {averaged_data['buffer_size'][c]}, l = {averaged_data['len_traj'][c]} s" for c in
                               range(len(averaged_data["len_traj"]))]
        else:
            setup["legend"] = [f"seed = {c}" for c in range(len(averaged_data["len_traj"]))]

    # print info amount CFD episodes, assuming 1st case is MF
    for i in range(len(averaged_data["MF_episodes"])):
        print(f"{averaged_data['MF_episodes'][i]} CFD episodes for case {i}")

    # plots for avg. the trajectories for cl & cd and plot then vs. t
    plot_mean_std_trajectories(setup, all_data)

    # plot variance of the beta-distribution wrt episodes
    plot_variance_of_beta_dist(setup, averaged_data["var_beta_fct"], n_cases=len(setup["case_name"]))

    # plot mean rewards wrt to episode
    plot_rewards_vs_episode(setup, reward_mean=averaged_data["mean_rewards"], reward_std=averaged_data["std_rewards"],
                            n_cases=len(setup["case_name"]))

    # plot mean cl and cd wrt to episode
    plot_coefficients_vs_episode(setup, cd_mean=averaged_data["mean_cd"], cd_std=averaged_data["std_cd"],
                                 cl_mean=averaged_data["mean_cl"], cl_std=averaged_data["std_cl"],
                                 n_cases=len(setup["case_name"]), plot_action=False)

    # plot total rewards received in the training
    plot_total_reward(setup, averaged_data["tot_mean_rewards"], averaged_data["tot_std_rewards"],
                      n_cases=len(setup["case_name"]))

    # compare trajectories over the course of the training
    for e in setup["e_trajectory"]:
        plot_cl_cd_trajectories(setup, all_data, number=1, e=e)

    # do frequency analysis of the cd- and cl-trajectories wrt episode number for each case
    for case in range(len(setup["case_name"])):
        analyze_frequencies_ppo_training(setup, all_data[case], case=case + 1)

        # also plot training- and validation losses of the environment models, if MB-DRL was used
        if "train_loss_cd" in all_data[case]:
            keys = 2 * ["train_loss_cl_p", "val_loss_cl_p", "train_loss_cd", "val_loss_cd"]
            key = [f"mean_" + k if idx < 4 else f"std_" + k for idx, k in enumerate(keys)]
            plot_train_validation_loss(setup, averaged_data["losses"][case][key[0]],
                                       averaged_data["losses"][case][key[1]], averaged_data["losses"][case][key[2]],
                                       averaged_data["losses"][case][key[3]], averaged_data["losses"][case][key[4]],
                                       averaged_data["losses"][case][key[5]], averaged_data["losses"][case][key[6]],
                                       averaged_data["losses"][case][key[7]], case + 1)

    # if the cases are run in openfoam using the trained network (using the best policy), plot the results
    if setup["plot_final_res"]:
        # plot the numerical setup for one case, assuming it's the same for all cases
        plot_numerical_setup(setup)

        # import the trajectory of the uncontrolled case
        uncontrolled = pd.read_csv(join(setup["main_load_path"], setup["path_uncontrolled"], "postProcessing", "forces",
                                        "0", "coefficient.dat"), skiprows=13, header=0,
                                   sep=r"\s+", usecols=[0, 1, 2], names=["t", "cd", "cl"])
        p_uncontrolled = pd.read_csv("".join([setup["main_load_path"], setup["path_uncontrolled"],
                                              setup["path_to_probes"]]), skiprows=setup["n_probes"] + 1, header=0,
                                     names=["t"] + [f"probe_{i}" for i in range(setup["n_probes"])], sep=r"\s+")

        controlled, p_controlled, traj = [], [], []
        for case in range(len(setup["case_name"])):
            # import the trajectories of the controlled cases
            controlled.append(pd.read_csv(join(setup["main_load_path"], setup["path_controlled"],
                                               setup["case_name"][case], setup["path_final_results"], "postProcessing",
                                               "forces", "0", "coefficient.dat"), skiprows=13, header=0,
                                          sep=r"\s+", usecols=[0, 1, 2], names=["t", "cd", "cl"]))

            traj.append(pd.read_csv(join(setup["main_load_path"], setup["path_controlled"], setup["case_name"][case],
                                         setup["path_final_results"], "trajectory.csv"), header=0, sep=r",",
                                    usecols=[0, 1, 2, 3], names=["t", "omega", "alpha", "beta"]))

            p_controlled.append(pd.read_csv(join(setup["main_load_path"], setup["path_controlled"],
                                                 setup["case_name"][case], setup["path_final_results"],
                                                 setup["path_to_probes"]), skiprows=setup["n_probes"] + 1,
                                            header=0, names=["t"] + [f"probe_{i}" for i in range(setup["n_probes"])],
                                            sep=r"\s+"))

        # plot cl and cd of the controlled cases vs. the uncontrolled cylinder flow
        plot_cl_cd_alpha_beta(setup, controlled, uncontrolled, plot_coeffs=True)

        # plot omega of the controlled cases
        plot_omega(setup, traj)

        # plot alpha and beta of the controlled cases
        plot_cl_cd_alpha_beta(setup, traj, plot_coeffs=False)

        # analyze frequency spectrum of cl- and cd-trajectories, therefore insert empty list int traj. data, so the idx
        # matches with the other data (since uncontrolled case hass no alpha, beta, omega)
        traj.insert(0, [])
        analyze_frequencies_final_result(setup, uncontrolled, controlled, traj)

        # analyze frequency spectrum of probes
        analyze_frequencies_probes_final_result(setup, p_uncontrolled, p_controlled, n_probes=setup["n_probes"])


# Licensed under a 3-clause BSD style license - see LICENSE.rst

from astropy import units
from astropy.io import fits
import numpy as np
import pytest

from sofia_redux.scan.configuration.configuration import Configuration
from sofia_redux.scan.flags.mounts import Mount
from sofia_redux.scan.info.instrument import InstrumentInfo
from sofia_redux.scan.source_models.beams.instant_focus import InstantFocus

class TestInstrumentInfo(object):
    def test_init(self):
        info = InstrumentInfo()
        assert info.name is None
        assert info.gain == 1.0
        assert info.mount == Mount.UNKNOWN
        assert info.log_id == 'inst'
        assert np.isnan(info.resolution)
        assert info.get_size_unit() == units.arcsec

    def test_set_configuration(self):
        info = InstrumentInfo()
        assert info.configuration is None
        config = Configuration()
        info.set_configuration(config)
        assert info.configuration is config

    def test_set_mount(self):
        info = InstrumentInfo()

        info.set_mount(2)
        assert info.mount == Mount.CASSEGRAIN

        info.set_mount('LEFT_NASMYTH')
        assert info.mount == Mount.LEFT_NASMYTH

        info.set_mount(Mount.RIGHT_NASMYTH)
        assert info.mount == Mount.RIGHT_NASMYTH

        with pytest.raises(ValueError) as err:
            info.set_mount('bad')
        assert 'bad is not a valid Mount' in str(err)

        with pytest.raises(ValueError) as err:
            info.set_mount([1, 2, 3])
        assert 'is not a valid Mount' in str(err)

    def test_get_source_size(self):
        info = InstrumentInfo()
        assert np.isnan(info.get_source_size())

        info.configuration = Configuration()
        assert np.isnan(info.get_source_size())

        # configure source size
        info.configuration.set_option('sourcesize', 10)
        assert np.isnan(info.get_source_size())

        # configure resolution
        info.resolution = 2 * units.arcsec
        assert info.get_source_size() == np.hypot(10, 2) * units.arcsec

    def test_get_stability(self):
        info = InstrumentInfo()
        assert info.get_stability() == 10 * units.s

        info.configuration = Configuration()
        assert info.get_stability() == 10 * units.s

        info.configuration.set_option('stability', 20)
        assert info.get_stability() == 20 * units.s

    def test_get_point_size(self):
        info = InstrumentInfo()
        assert np.isnan(info.get_point_size())
        info.resolution = 10 * units.arcsec
        assert info.get_point_size() == 10 * units.arcsec

    def test_get_spectral_size(self):
        info = InstrumentInfo()
        assert np.isnan(info.get_spectral_size())
        assert info.get_spectral_size().unit == 'um'

    def test_get_spectral_unit(self):
        assert InstrumentInfo.get_spectral_unit() == 'um'

    def test_get_data_unit(self):
        info = InstrumentInfo()
        assert info.get_data_unit() == units.count

        info.configuration = Configuration()
        assert info.get_data_unit() == units.count

        info.configuration.set_option('dataunit', 'Jy')
        assert info.get_data_unit() == units.Jy

    def test_jansky_per_beam(self):
        info = InstrumentInfo()
        assert info.jansky_per_beam() == 1 * units.Jy / units.beam

        info.configuration = Configuration()
        assert info.jansky_per_beam() == 1 * units.Jy / units.beam

        info.configuration.set_option('jansky', 100)
        assert info.jansky_per_beam() == 100 * units.Jy / units.beam

        info.configuration.set_option('jansky.inverse', True)
        assert info.jansky_per_beam() == .01 * units.Jy / units.beam

        info.configuration.set_option('dataunit', 'mJy')
        assert info.jansky_per_beam() == 1e-5 * units.Jy / units.beam

    def test_kelvin(self):
        info = InstrumentInfo()
        assert np.isnan(info.kelvin())

        info.configuration = Configuration()
        assert np.isnan(info.kelvin())

        info.configuration.set_option('k2jy', 10)
        assert info.kelvin() == 10 * units.Kelvin

        info.configuration.set_option('kelvin', 100)
        assert info.kelvin() == 100 * units.Kelvin

    def test_edit_image_header(self):
        info = InstrumentInfo()
        info.name = 'test'

        hdr = fits.Header()
        info.edit_image_header(hdr)
        assert hdr['INSTRUME'] == 'test'
        assert hdr['V2JY'] == 1
        assert len(hdr) == 2

    def test_validate_scan(self, populated_scan):
        info = InstrumentInfo()
        info.configuration = Configuration()
        inst = populated_scan.info.instrument

        # nothing configured: no op
        info.validate_scan(populated_scan)
        assert np.isnan(inst.frequency)
        assert inst.resolution == 10 * units.arcsec
        assert inst.gain == 1.0

        # configure wavelength, resolution, gain
        info.configuration.set_option('wavelength', 5)
        info.configuration.set_option('resolution', 2)
        info.configuration.set_option('gain', 0.5)
        info.validate_scan(populated_scan)
        assert np.isclose(inst.frequency, 5.9958e13 * units.Hz)
        assert inst.resolution == 2 * units.arcsec
        assert inst.gain == 0.5

        # configure frequency directly instead
        info.configuration.set_option('frequency', 5)
        info.validate_scan(populated_scan)
        assert np.isclose(inst.frequency, 5 * units.Hz)

    def test_get_focus_string(self):
        msg = InstrumentInfo.get_focus_string(None)
        assert msg == 'No instant focus'

        focus = InstantFocus()
        focus.x = 1 * units.cm
        focus.x_weight = 1 * units.cm ** -2
        msg = InstrumentInfo.get_focus_string(focus)
        assert msg == '\n  Focus.dX --> 1.0 cm +- 1.0 cm'

        focus.x = 1
        focus.x_weight = 1
        msg = InstrumentInfo.get_focus_string(focus)
        assert msg == '\n  Focus.dX --> 1.0 mm +- 1.0 mm'


import os
import yaml
from typing import Any

from prefect.tasks.shell import ShellTask
from prefect.utilities.tasks import defaults_from_attrs

class DbtShellTask(ShellTask):
    """
    Task for running dbt commands. It will create a profiles.yml file prior to running dbt commands.

    This task inherits all configuration options from the
    [ShellTask](https://docs.prefect.io/api/latest/tasks/shell.html#shelltask).

    Args:
        - command (string, optional): dbt command to be executed; can also be
            provided post-initialization by calling this task instance
        - dbt_kwargs (dict, optional): keyword arguments used to populate the profiles.yml file
            (e.g.  `{'type': 'snowflake', 'threads': 4, 'account': '...'}`); can also be
            provided at runtime
        - env (dict, optional): dictionary of environment variables to use for
            the subprocess; can also be provided at runtime
        - environment (string, optional): The default target your dbt project will use
        - overwrite_profiles (boolean, optional): flag to indicate whether existing
            profiles.yml file should be overwritten; defaults to `False`
        - profile_name (string, optional): Profile name used for populating the profile name of
            profiles.yml
        - profiles_dir (string, optional): path to directory where the profile.yml file will be
            contained
        - set_profiles_envar (boolean, optional): flag to indicate whether DBT_PROFILES_DIR
            should be set to the provided profiles_dir; defaults to `True`
        - helper_script (str, optional): a string representing a shell script, which
            will be executed prior to the `command` in the same process. Can be used to
            change directories, define helper functions, etc. when re-using this Task
            for different commands in a Flow; can also be provided at runtime
        - shell (string, optional): shell to run the command with; defaults to "bash"
        - return_all (bool, optional): boolean specifying whether this task should return all
            lines of stdout as a list, or just the last line as a string; defaults to `False`
        - log_stderr (bool, optional): boolean specifying whether this task
            should log the output from stderr in the case of a non-zero exit code;
            defaults to `False`
        - **kwargs: additional keyword arguments to pass to the Task constructor

    Example:
        ```python
        from prefect import Flow
        from prefect.tasks.dbt import DbtShellTask

        with Flow(name="dbt_flow") as f:
            task = DbtShellTask(
                profile_name='default',
                environment='test',
                dbt_kwargs={
                    'type': 'snowflake',
                    'threads': 1,
                    'account': 'account.us-east-1'
                },
                overwrite_profiles=True,
                profiles_dir=test_path
            )(command='dbt run')

        out = f.run()
        ```
    """

    def __init__(
        self,
        command: str = None,
        profile_name: str = None,
        env: dict = None,
        environment: str = None,
        overwrite_profiles: bool = False,
        profiles_dir: str = None,
        set_profiles_envar: bool = True,
        dbt_kwargs: dict = None,
        helper_script: str = None,
        shell: str = "bash",
        return_all: bool = False,
        log_stderr: bool = False,
        **kwargs: Any
    ):
        self.command = command
        self.profile_name = profile_name
        self.environment = environment
        self.overwrite_profiles = overwrite_profiles
        self.profiles_dir = profiles_dir
        self.set_profiles_envar = set_profiles_envar
        self.dbt_kwargs = dbt_kwargs or {}
        super().__init__(
            **kwargs,
            command=command,
            env=env,
            helper_script=helper_script,
            shell=shell,
            return_all=return_all,
            log_stderr=log_stderr
        )

    @defaults_from_attrs("command", "env", "helper_script", "dbt_kwargs")
    def run(
        self,
        command: str = None,
        env: dict = None,
        helper_script: str = None,
        dbt_kwargs: dict = None,
    ) -> str:
        """
        If no profiles.yml file is found or if overwrite_profiles flag is set to True, this
        will first generate a profiles.yml file in the profiles_dir directory. Then run the dbt
        cli shell command.

        Args:
            - command (string): shell command to be executed; can also be
                provided at task initialization. Any variables / functions defined in
                `self.helper_script` will be available in the same process this command
                runs in
            - env (dict, optional): dictionary of environment variables to use for
                the subprocess
            - helper_script (str, optional): a string representing a shell script, which
                will be executed prior to the `command` in the same process. Can be used to
                change directories, define helper functions, etc. when re-using this Task
                for different commands in a Flow
             - dbt_kwargs(dict, optional): keyword arguments used to populate the profiles.yml file

        Returns:
            - stdout (string): if `return_all` is `False` (the default), only the last line of
                stdout is returned, otherwise all lines are returned, which is useful for
                passing result of shell command to other downstream tasks. If there is no
                output, `None` is returned.

        Raises:
            - prefect.engine.signals.FAIL: if command has an exit code other
                than 0
        """
        DEFAULT_PROFILES_DIR = os.path.join(os.path.expanduser("~"), ".dbt")
        profiles_exists = False
        if os.getenv("DBT_PROFILES_DIR"):
            dbt_profiles_dir = os.path.expanduser(
                os.getenv("DBT_PROFILES_DIR", DEFAULT_PROFILES_DIR)
            )
            profiles_exists = os.path.exists(
                os.path.join(dbt_profiles_dir, "profiles.yml")
            )
        elif self.profiles_dir:
            profiles_exists = os.path.exists(
                os.path.join(self.profiles_dir, "profiles.yml")
            )
        else:
            profiles_exists = os.path.exists(
                os.path.join(DEFAULT_PROFILES_DIR, "profiles.yml")
            )

        dbt_kwargs = {**self.dbt_kwargs, **(dbt_kwargs or {})}

        if self.overwrite_profiles or not profiles_exists:
            profile = {
                self.profile_name: {
                    "outputs": {self.environment: dbt_kwargs},
                    "target": self.environment,
                }
            }

            if not self.profiles_dir:
                try:
                    os.mkdir(DEFAULT_PROFILES_DIR)
                except OSError:
                    self.logger.warning(
                        "Creation of directory %s has failed" % DEFAULT_PROFILES_DIR
                    )
                profile_path = os.path.join(DEFAULT_PROFILES_DIR, "profiles.yml")
                self.profiles_dir = DEFAULT_PROFILES_DIR
            else:
                profile_path = os.path.join(self.profiles_dir, "profiles.yml")

            with open(profile_path, "w+") as yaml_file:
                yaml.dump(profile, yaml_file, default_flow_style=False)

        if self.set_profiles_envar:
            os.environ["DBT_PROFILES_DIR"] = self.profiles_dir

        return super(DbtShellTask, self).run(
            command=command, env=env, helper_script=helper_script
        )


import asyncio
import hci_packets as hci
import link_layer_packets as ll
import math
import random
import unittest
from dataclasses import dataclass
from hci_packets import ErrorCode, FragmentPreference
from py.bluetooth import Address
from py.controller import ControllerTest
from typing import List

def make_advertising_event_properties(properties: int) -> hci.AdvertisingEventProperties:
    return hci.AdvertisingEventProperties(connectable=(properties & 0x1) != 0,
                                          scannable=(properties & 0x2) != 0,
                                          directed=(properties & 0x4) != 0,
                                          high_duty_cycle=(properties & 0x8) != 0,
                                          legacy=(properties & 0x10) != 0,
                                          anonymous=(properties & 0x20) != 0,
                                          tx_power=(properties & 0x40) != 0)

@dataclass
class TestRound:
    advertising_event_properties: int
    data_length: int
    fragment_preference: FragmentPreference
    duration: int
    max_extended_advertising_events: int

class Test(ControllerTest):
    # Test parameters.
    LL_advertiser_advInterval_MIN = 0x200
    LL_advertiser_advInterval_MAX = 0x200
    LL_advertiser_Adv_Channel_Map = 0x7
    LL_initiator_connInterval = 0x200
    LL_initiator_connPeripheralLatency = 0x200
    LL_initiator_connSupervisionTimeout = 0x200

    # LL/DDI/ADV/BV-47-C [Extended Advertising, Non-Connectable – LE 1M PHY]
    async def test(self):
        controller = self.controller

        # 1. The Upper Tester sends an HCI_LE_Read_Maximum_Advertising_Data_Length command to the
        # IUT and expects the IUT to return a Maximum_Advertising_Data_Length between 0x001F and
        # 0x0672. The Upper Tester stores the Maximum_Advertising_Data_Length for future use.
        controller.send_cmd(hci.LeReadMaximumAdvertisingDataLength())

        event = await self.expect_cmd_complete(hci.LeReadMaximumAdvertisingDataLengthComplete)
        maximum_advertising_data_length = event.maximum_advertising_data_length

        # Test rounds.
        test_rounds = [
            TestRound(0x0, 0, FragmentPreference.CONTROLLER_MAY_FRAGMENT, 0x0, 0x0),
            TestRound(0x0, 31, FragmentPreference.CONTROLLER_MAY_FRAGMENT, 0x0, 0x0),
            TestRound(0x0, 474, FragmentPreference.CONTROLLER_MAY_FRAGMENT, 0x0, 0x0),
            TestRound(0x0, 711, FragmentPreference.CONTROLLER_MAY_FRAGMENT, 0x0, 0x0),
            TestRound(0x0, 948, FragmentPreference.CONTROLLER_MAY_FRAGMENT, 0x0, 0x0),
            TestRound(0x0, maximum_advertising_data_length, FragmentPreference.CONTROLLER_MAY_FRAGMENT, 0x0, 0x0),
            TestRound(0x0, maximum_advertising_data_length, FragmentPreference.CONTROLLER_SHOULD_NOT, 0x0, 0x0),
            TestRound(0x4, 0, FragmentPreference.CONTROLLER_MAY_FRAGMENT, 0x0, 0x0),
            TestRound(0x4, 251, FragmentPreference.CONTROLLER_MAY_FRAGMENT, 0x0, 0x0),
            TestRound(0x4, maximum_advertising_data_length, FragmentPreference.CONTROLLER_MAY_FRAGMENT, 0x0, 0x0),
            TestRound(0x0, 0, FragmentPreference.CONTROLLER_MAY_FRAGMENT, 0x1f4, 0x0),
            TestRound(0x4, 0, FragmentPreference.CONTROLLER_MAY_FRAGMENT, 0x1f4, 0x0),
            TestRound(0x0, 0, FragmentPreference.CONTROLLER_MAY_FRAGMENT, 0x0, 0x32),
            TestRound(0x4, 0, FragmentPreference.CONTROLLER_MAY_FRAGMENT, 0x0, 0x32),
        ]

        # 14. Repeat steps 2–13 for each Round shown in Table 4.6
        for test_round in test_rounds:
            await self.steps_2_13(maximum_advertising_data_length, **vars(test_round))

    async def steps_2_13(self, maximum_advertising_data_length: int, advertising_event_properties: int,
                         data_length: int, fragment_preference: FragmentPreference, duration: int,
                         max_extended_advertising_events: int):
        controller = self.controller
        advertising_event_properties = make_advertising_event_properties(advertising_event_properties)

        # 2. If the Data Length listed in Table 4.6 for the current Round is less than or equal to the
        # Maximum_Advertising_Data_Length proceed to step 3, otherwise skip to step 14.
        if data_length > maximum_advertising_data_length:
            return

        # 3. The Upper Tester sends an HCI_LE_Set_Extended_Advertising_Parameters command to the
        # IUT using all supported advertising channels and a selected advertising interval between the
        # minimum and maximum advertising intervals supported. Advertising_Event_Properties parameter
        # shall be set to the value specified in Table 4.6 for this round. The Primary_Advertising_PHY and
        # Secondary_Advertising_PHY shall be set to the values specified in Table 4.5. If the
        # Advertising_Event_Properties value for this Round specifies directed advertising, the
        # Peer_Address_Type shall be set to 0x00 (Public Device Address), and the Peer_Address shall be
        # set to the Lower Tester’s address.
        controller.send_cmd(
            hci.LeSetExtendedAdvertisingParameters(advertising_handle=0,
                                                   advertising_event_properties=advertising_event_properties,
                                                   primary_advertising_interval_min=self.LL_advertiser_advInterval_MIN,
                                                   primary_advertising_interval_max=self.LL_advertiser_advInterval_MAX,
                                                   primary_advertising_channel_map=self.LL_advertiser_Adv_Channel_Map,
                                                   own_address_type=hci.OwnAddressType.PUBLIC_DEVICE_ADDRESS,
                                                   advertising_filter_policy=hci.AdvertisingFilterPolicy.ALL_DEVICES,
                                                   primary_advertising_phy=hci.PrimaryPhyType.LE_1M))

        await self.expect_evt(
            hci.LeSetExtendedAdvertisingParametersComplete(status=ErrorCode.SUCCESS, num_hci_command_packets=1))

        # 4. The Upper Tester sends one or more HCI_LE_Set_Extended_Advertising_Data commands to the
        # IUT with values according to Table 4.6 and using random octets from 1 to 254 as the payload. If
        # the Data Length is greater than 251 the Upper Tester shall send multiple commands using one
        # Operation 0x01 (First fragment) command, followed by zero or more Operation 0x00
        # (Intermediate Fragment) commands, and a final Operation 0x02 (Last fragment) command.
        # Otherwise the Upper Tester shall send a single command using Operation 0x03 (Complete Data).
        advertising_data = [random.randint(1, 254) for n in range(data_length)]
        num_fragments = math.ceil(data_length / 251) or 1  # Make sure to set the advertising data if it is empty.
        for n in range(num_fragments):
            fragment_offset = 251 * n
            fragment_length = min(251, data_length - fragment_offset)
            if num_fragments == 1:
                operation = hci.Operation.COMPLETE_ADVERTISEMENT
            elif n == 0:
                operation = hci.Operation.FIRST_FRAGMENT
            elif n == num_fragments - 1:
                operation = hci.Operation.LAST_FRAGMENT
            else:
                operation = hci.Operation.INTERMEDIATE_FRAGMENT

            controller.send_cmd(
                hci.LeSetExtendedAdvertisingData(advertising_handle=0,
                                                 operation=operation,
                                                 advertising_data=advertising_data[fragment_offset:fragment_offset +
                                                                                fragment_length]))

            await self.expect_evt(
                hci.LeSetExtendedAdvertisingDataComplete(status=ErrorCode.SUCCESS, num_hci_command_packets=1))

        # 5. The Upper Tester enables advertising using the HCI_LE_Set_Extended_Advertising_Enable
        # command. The Duration[0] parameter shall be set to the value specified in Table 4.6 for this
        # round. The Max_Extended_Advertising_Events[0] parameter shall be set to the value specified in
        # Table 4.6 for this round.
        controller.send_cmd(
            hci.LeSetExtendedAdvertisingEnable(enable=hci.Enable.ENABLED,
                                               enabled_sets=[
                                                   hci.EnabledSet(
                                                       advertising_handle=0,
                                                       duration=duration,
                                                       max_extended_advertising_events=max_extended_advertising_events)
                                               ]))

        await self.expect_evt(
            hci.LeSetExtendedAdvertisingEnableComplete(status=ErrorCode.SUCCESS, num_hci_command_packets=1))

        # 6. The Lower Tester receives an ADV_EXT_IND packet from the IUT with AdvMode set to 00b. The
        # ADV_EXT_IND PDU shall not include the SuppInfo, SyncInfo, TxPower, ACAD, or AdvData
        # fields. If advertising data was set in step 4, the ADV_EXT_IND PDU shall include the AuxPtr field;
        # otherwise, the ADV_EXT_IND PDU may include the AuxPtr field. If the AuxPtr field is included,
        # the ADV_EXT_IND PDU shall also include the ADI field with the SID set to the value used in step
        # 3; otherwise that field shall not be included.

        # 7. If the AuxPtr is absent, skip to step 10.

        # 8. The Lower Tester utilizes the AuxPtr field to listen for an AUX_ADV_IND PDU on the secondary
        # advertising channel with the AdvMode field set to 00b. The AUX_ADV_IND PDU shall not include
        # the SuppInfo, SyncInfo, or TxPower fields. The AUX_ADV_IND PDU shall include the ADI field
        # matching the ADI field from step 6. If the AUX_ADV_IND PDU does not contain all the data
        # submitted in step 4 (if any), it shall include an AuxPtr field.

        # 9. If the AUX_ADV_IND PDU contains an AuxPtr field, the Lower Tester utilizes it to listen for an
        # AUX_CHAIN_IND PDU with the AdvMode field set to 00b. The AUX_CHAIN_IND PDU shall
        # include the ADI field matching the ADI field from step 6 and the AdvData field containing
        # additional data submitted in step 4. The AUX_CHAIN_IND PDU shall not include the AdvA,
        # TargetA, SuppInfo, TxPower, or SyncInfo fields. If the AUX_CHAIN_IND PDU contains an AuxPtr
        # field this step is repeated until an AUX_CHAIN_IND PDU is received with no AuxPtr field and all
        # data has been received.
        repeat = max_extended_advertising_events or 3
        for n in range(repeat):
            await self.expect_ll(
                ll.LeExtendedAdvertisingPdu(source_address=controller.address,
                                            advertising_address_type=ll.AddressType.PUBLIC,
                                            target_address_type=ll.AddressType.PUBLIC,
                                            connectable=int(advertising_event_properties.connectable),
                                            scannable=int(advertising_event_properties.scannable),
                                            directed=int(advertising_event_properties.directed),
                                            sid=0,
                                            tx_power=0,
                                            primary_phy=ll.PrimaryPhyType.LE_1M,
                                            secondary_phy=ll.SecondaryPhyType.NO_PACKETS,
                                            advertising_data=advertising_data))

        # 10. If the Max_Extended_Advertising_Events was set to a value different than 0, repeat steps 6–9
        # until the IUT stops advertising. Afterwards, the Lower Tester confirms that the IUT did not send
        # more than Max_Extended_Advertising_Events advertising events. Upper Tester shall receive LE
        # Advertising Set Terminated event with ErrorCode 0x43. Skip to step 13.
        if max_extended_advertising_events > 0:
            try:
                # Note: The test should timeout waiting for an advertising event
                # past Max Extended Advertising Events count.
                await asyncio.wait_for(self.controller.receive_ll(), timeout=1)
                self.assertTrue(False)
            except asyncio.exceptions.TimeoutError:
                pass

            await self.expect_evt(
                hci.LeAdvertisingSetTerminated(
                    status=ErrorCode.ADVERTISING_TIMEOUT,
                    advertising_handle=0,
                    connection_handle=0,
                    num_completed_extended_advertising_events=max_extended_advertising_events))

        # 11. Otherwise if Duration was set to a value different than 0, repeat steps 6–9 until the amount of
        # time specified for Duration has elapsed. Afterwards, the Lower Tester confirms that the IUT does
        # not start any additional advertising events. Upper Tester shall receive LE Advertising Set
        # Terminated event with ErrorCode 0x3C. Skip to step 13.
        elif duration > 0:
            try:
                # Note: The test should timeout waiting for a directed advertising event
                # past the direct advertising timeout.
                end_time = asyncio.get_running_loop().time() + duration / 100
                while asyncio.get_running_loop().time() < end_time:
                    await asyncio.wait_for(self.controller.receive_ll(), timeout=1)
                self.assertTrue(False)
            except asyncio.exceptions.TimeoutError:
                pass

            await self.expect_evt(
                hci.LeAdvertisingSetTerminated(status=ErrorCode.ADVERTISING_TIMEOUT,
                                               advertising_handle=0,
                                               connection_handle=0,
                                               num_completed_extended_advertising_events=0))

        # 12. Otherwise, repeat steps 6–9 until a number of advertising intervals (10) have been detected.

        # 13. The Upper Tester disables advertising using the HCI_LE_Set_Extended_Advertising_Enable
        # command.
        controller.send_cmd(hci.LeSetExtendedAdvertisingEnable(enable=hci.Enable.DISABLED, enabled_sets=[]))

        await self.expect_evt(
            hci.LeSetExtendedAdvertisingEnableComplete(status=ErrorCode.SUCCESS, num_hci_command_packets=1))


#! /usr/bin/python3
import iota.harness.api as api

def __get_flap_cfg_cmd(intf, oper):
    flap_base_cmd = "ifconfig %s %s"
    flap_cmd = flap_base_cmd % (intf, oper)
    return flap_cmd

def __verify_response(resp, oper):
    if resp is None:
        api.Logger.critical("Failed to set state on host interface - No response")
        return False
    for cmd in resp.commands:
        if cmd.exit_code != 0:
            api.Logger.critical("Failed to set state on interface")
            api.PrintCommandResults(cmd)
            #return False
    return True

def __set_state_all_host_intfs(oper):
    req = api.Trigger_CreateExecuteCommandsRequest()
    for wl in api.GetWorkloads():
        flap_cmd = __get_flap_cfg_cmd(wl.parent_interface, oper)
        api.Trigger_AddCommand(req, wl.node_name, wl.workload_name, flap_cmd)

    resp = api.Trigger(req)
    return __verify_response(resp, oper)

def __flap_host_intf():
    if __set_state_all_host_intfs("down") == False:
        return api.types.status.FAILURE

    if __set_state_all_host_intfs("up") == False:
        return api.types.status.FAILURE

    return api.types.status.SUCCESS

def flapHostIntf(tc): #use this api if invoking as a background task
    api.Logger.info("FLAPPING Host Intfs")
    return __flap_host_intf()

def Main(step): #directly use utils.hostflap as step in iota suite
    api.Logger.info("FLAPPING Host Intfs")
    return __flap_host_intf()

if __name__ == '__main__':
    Main(None)


import os, tqdm, random, pickle

import torch
import torchvision

from torch.autograd import Variable
from torchvision.transforms import CenterCrop, ToTensor, Compose, Lambda, Resize
from torchvision.datasets import coco
from torch.nn.functional import binary_cross_entropy, relu, nll_loss
from torch.nn import Embedding
from torch.optim import Adam

import nltk

from argparse import ArgumentParser

from collections import defaultdict, Counter, OrderedDict

import util, models

from util import PAD, SOS, EOS, UNK, EXTRA_SYMBOLS
from enum import Enum

from tensorboardX import SummaryWriter

REP = 3
TEMPS = [0.0, 0.1, 1.0]

class Mode(Enum):
    independent = 'independent'
    coupled = 'coupled'
    style = 'style'

    def __str__(self):
        return self.value

def go(arg):

    tbw = SummaryWriter(log_dir=arg.tb_dir)

    transform = Compose([
        Lambda(lambda x: CenterCrop(min(x.size))(x)),
        Resize(size=(arg.img_size, arg.img_size)),
        ToTensor()])

    imdir = arg.data_dir + os.sep + 'val2017'
    anfile = arg.data_dir + os.sep + 'annotations' + os.sep + 'captions_val2017.json'

    coco_data = coco.CocoCaptions(root=imdir, annFile=anfile, transform=transform)

    ## Make a dictionary

    util.ensure(arg.cache_dir)
    if os.path.isfile(arg.cache_dir + os.sep + 'i2w.pkl'):
        with open(arg.cache_dir + os.sep + 'i2w.pkl', 'rb') as file:
            i2w = pickle.load(file)
        with open(arg.cache_dir + os.sep + 'w2i.pkl', 'rb') as file:
            w2i = pickle.load(file)
        print('Word indices loaded.')
    else:
        print('Creating word indices') # Why is this so slow?

        dist = Counter()
        for i in tqdm.trange(len(coco_data)):
            for caption in coco_data[i][1]:
                dist.update(util.tokenize(caption))

        vocab = dist.most_common(arg.max_vocab - len(EXTRA_SYMBOLS))

        i2w = EXTRA_SYMBOLS + [w[0] for w in vocab]
        w2i = {word:ix for ix, word in enumerate(i2w)}

        with open(arg.cache_dir + os.sep + 'i2w.pkl', 'wb') as file:
            pickle.dump(i2w, file)
        with open(arg.cache_dir + os.sep + 'w2i.pkl', 'wb') as file:
            pickle.dump(w2i, file)

    vocab_size = len(i2w)
    print('vocabulary size', vocab_size)
    print('top 100 words:', i2w[:100])

    def decode(indices):

        sentence = ''
        for id in indices:
                # if id == PAD:
                #     break
                sentence += i2w[id] + ' '

        return sentence

    ## Set up the models
    embedding = torch.nn.Embedding(num_embeddings=vocab_size, embedding_dim=arg.embedding_size)

    if arg.mode != Mode.style:
        img_enc = models.ImEncoder(in_size=(arg.img_size, arg.img_size), zsize=arg.latent_size)
        img_dec = models.ImDecoder(in_size=(arg.img_size, arg.img_size), zsize=arg.latent_size)

        seq_enc = models.SeqEncoder(vocab_size=vocab_size, embedding=embedding, zsize=arg.latent_size)
        seq_dec = models.SeqDecoder(vocab_size=vocab_size, embedding=embedding, zsize=arg.latent_size)

        mods = [img_enc, img_dec, seq_enc, seq_dec]
    else:
        img_enc = models.ImEncoder(in_size=(arg.img_size, arg.img_size), zsize=arg.latent_size)
        img_sty = models.ImEncoder(in_size=(arg.img_size, arg.img_size), zsize=arg.latent_size)
        img_dec = models.ImDecoder(in_size=(arg.img_size, arg.img_size), zsize=arg.latent_size * 2)

        seq_enc = models.SeqEncoder(vocab_size=vocab_size, embedding=embedding, zsize=arg.latent_size)
        seq_sty = models.SeqEncoder(vocab_size=vocab_size, embedding=embedding, zsize=arg.latent_size)
        seq_dec = models.SeqDecoder(vocab_size=vocab_size, embedding=embedding, zsize=arg.latent_size * 2)

        mods = [img_enc, img_dec, img_sty, seq_enc, seq_dec, seq_sty]

    if torch.cuda.is_available():
        for model in mods:
            model.cuda()

    #- The standard dataloader approach doesn't seem to work with the captions, so we'll do our own batching.
    #  It's a little slower, probably, but it won't be the bottleneck
    params = []
    for model in mods:
        params.extend(model.parameters())
    optimizer = Adam(params, lr=arg.lr)

    instances_seen = 0

    for e in range(arg.epochs):
        print('epoch', e)
        for fr in tqdm.trange(0, len(coco_data), arg.batch_size):
            if arg.instance_limit is not None and fr > arg.instance_limit:
                break

            to = min(len(coco_data), fr + arg.batch_size)

            images = []
            captions = []

            for i in range(fr, to):
                images.append(coco_data[i][0].unsqueeze(0))
                captions.append(random.choice(coco_data[i][1])) # we choose one of the available captions at random

            imbatch = torch.cat(images, dim = 0)
            b, c, w, h = imbatch.size()

            capbatch = [] # to integer sequence
            for caption in captions:
                capbatch.append(util.intseq(util.tokenize(caption), w2i))

            capbatch, lengths = util.pad(capbatch)

            # Created shifted versions
            b, s = capbatch.size()

            # Input for the decoder
            cap_teacher = torch.cat([torch.ones(b, 1, dtype=torch.long), capbatch], dim=1)
            cap_out     = torch.cat([capbatch, torch.zeros(b, 1, dtype=torch.long)], dim=1)

            lengths = torch.LongTensor(lengths)

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

                capbatch = capbatch.cuda()
                cap_teacher = cap_teacher.cuda()
                cap_out = cap_out.cuda()

                lengths = lengths.cuda()

            imbatch = Variable(imbatch)
            capbatch = Variable(capbatch)
            cap_teacher = Variable(cap_teacher)
            cap_out = Variable(cap_out)
            lengths = Variable(lengths)

            zimg = img_enc(imbatch)
            zcap = seq_enc(capbatch, lengths)

            kl_img = util.kl_loss(*zimg)
            kl_cap = util.kl_loss(*zcap)

            zimg_sample = util.sample(*zimg)
            zcap_sample = util.sample(*zcap)

            if arg.mode == Mode.style:
                zimg_sty = img_sty(imbatch)
                zcap_sty = seq_sty(capbatch, lengths)

                kl_img_sty = util.kl_loss(*zimg_sty)
                kl_cap_sty = util.kl_loss(*zcap_sty)

                zimg_sample_sty = util.sample(*zimg_sty)
                zcap_sample_sty = util.sample(*zcap_sty)

                zimg_sample = torch.cat([zimg_sample, zimg_sample_sty], dim=1)
                zcap_sample = torch.cat([zcap_sample, zcap_sample_sty], dim=1)

            rec_imgimg = img_dec(zimg_sample)
            rl_imgimg = binary_cross_entropy(rec_imgimg, imbatch, reduce=False).view(b, -1).sum(dim=1)

            rec_capcap = seq_dec(zcap_sample, cap_teacher, lengths + 1).transpose(1, 2)
            rl_capcap = nll_loss(rec_capcap, cap_out, reduce=False).view(b, -1).sum(dim=1)

            if arg.mode != Mode.independent:
                rec_capimg = img_dec(zcap_sample)
                rl_capimg = binary_cross_entropy(rec_capimg, imbatch, reduce=False).view(b, -1).sum(dim=1)

                rec_imgcap = seq_dec(zimg_sample, cap_teacher, lengths + 1).transpose(1, 2)
                rl_imgcap = nll_loss(rec_imgcap, cap_out, reduce=False).view(b, -1).sum(dim=1)

            loss_img = rl_imgimg + kl_img
            loss_cap = rl_capcap + kl_cap

            if arg.mode == Mode.coupled:
                loss_img = loss_img + rl_capimg + kl_img
                loss_cap = loss_cap + rl_imgcap + kl_cap

            if arg.mode == Mode.style:
                loss_img = loss_img + kl_img_sty
                loss_cap = loss_cap + kl_cap_sty

            loss = loss_img.mean() + loss_cap.mean()

            #- backward pass
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            instances_seen += b

            tbw.add_scalar('score/img/kl', float(kl_img.mean()), instances_seen)
            tbw.add_scalar('score/imgimg/rec', float(rl_imgimg.mean()), instances_seen)
            tbw.add_scalar('score/cap/kl', float(kl_cap.mean()), instances_seen)
            tbw.add_scalar('score/capcap/rec', float(rl_capcap.mean()), instances_seen)
            tbw.add_scalar('score/loss', float(loss), instances_seen)

            if arg.mode != Mode.independent:
                tbw.add_scalar('score/capimg/rec', float(rl_capimg.mean()), instances_seen)
                tbw.add_scalar('score/imgcap/rec', float(rl_imgcap.mean()), instances_seen)

        # Interpolate
        zpairs = []
        for r in range(REP):

            print('Interpolation, repeat', r)

            l = arg.latent_size if arg.mode != Mode.style else arg.latent_size * 2
            z1, z2 = torch.randn(2, l)
            if torch.cuda.is_available():
                z1, z2 = z1.cuda(), z2.cuda()

            zpairs.append((z1, z2))

            zs = util.slerp(z1, z2, 10)

            print('== sentences (temp={}) =='.format(TEMPS[r]))
            sentences = seq_dec.sample(z=zs, temperature=TEMPS[r])

            for s in sentences:
                print('   ', decode(s))

        print('== images ==')

        util.interpolate(zpairs, img_dec, name='interpolate.{}'.format(e))

if __name__ == "__main__":

    ## Parse the command line options
    parser = ArgumentParser()

    parser.add_argument("-m", "--mode",
                        dest="mode",
                        help="Mode. independent: trains fully separate autoencoders for images and language. shared: couples the latent space of the two autoencoders. style: uses separate encoders to capture style.",
                        default=Mode.independent, type=Mode, choices=list(Mode))

    parser.add_argument("-e", "--epochs",
                        dest="epochs",
                        help="Number of epochs.",
                        default=150, type=int)

    parser.add_argument("-b", "--batch-size",
                        dest="batch_size",
                        help="Size of the batches.",
                        default=32, type=int)

    parser.add_argument("-L", "--latent-size",
                        dest="latent_size",
                        help="Size of the latent representations.",
                        default=32, type=int)

    parser.add_argument("-E", "--embedding-size",
                        dest="embedding_size",
                        help="Size of the embeddings.",
                        default=300, type=int)

    parser.add_argument("--limit",
                        dest="instance_limit",
                        help="Limit on the number of instances seen per batch (for debugging).",
                        default=None, type=int)

    parser.add_argument("-l", "--learn-rate",
                        dest="lr",
                        help="Learning rate.",
                        default=0.001, type=float)

    parser.add_argument("-I", "--image-size",
                        dest="img_size",
                        help="Size in pixels of on eof the sides of the (square) images.",
                        default=128, type=int)

    parser.add_argument("-w", "--max-vocab",
                        dest="max_vocab",
                        help="Maximum vocabulary.",
                        default=25000, type=int)

    parser.add_argument("-D", "--data-directory",
                        dest="data_dir",
                        help="Data directory",
                        default='./data', type=str)

    parser.add_argument("-T", "--tb-directory",
                        dest="tb_dir",
                        help="Tensorboard directory",
                        default='./runs/score', type=str)

    parser.add_argument("-C", "--cache-directory",
                        dest="cache_dir",
                        help="Dir for cache files (delete the dir to reconstruct)",
                        default='./cache', type=str)

    options = parser.parse_args()

    print('OPTIONS', options)

    go(options)


import numpy as np
import cv2
import math

# First we create a color space containing hist_bin^3 bins. And we also need to
# adjust position inputs to make sure they lie in the frame. For every pixel in
# this small region, we judge which bin the RGB pixel should be in and sum them
# up for the final histogram and then normalize it.
def color_histogram(xmin, ymin, xmax, ymax, frame, hist_bin) -> np.array:
    hist = np.zeros(hist_bin ** 3)
    divide = np.linspace(0, 255, hist_bin + 1)
    xmin = np.maximum(xmin, 0)
    ymin = np.maximum(ymin, 0)
    xmax = np.minimum(xmax, frame.shape[1])
    ymax = np.minimum(ymax, frame.shape[0])

    # hist[i][j][k] = np.sum(frame[:,:] == [i,j,k])
    for x in range(math.ceil(xmin), math.floor(xmax)):
        for y in range(math.ceil(ymin), math.floor(ymax)):
            R = frame[y, x, 0]
            G = frame[y, x, 1]
            B = frame[y, x, 2]
            for i in range(hist_bin):
                if divide[i] <= R < divide[i + 1]:
                    R_index = i
                if divide[i] <= G < divide[i + 1]:
                    G_index = i
                if divide[i] <= B < divide[i + 1]:
                    B_index = i
            RGB_index = (hist_bin ** 2) * R_index + hist_bin * G_index + B_index
            hist[RGB_index] = hist[RGB_index] + 1

    return hist/sum(hist)


# CORRECT SOLUTION, START FROM THE GATES:
def fill_nearest_gate_distance(nlines, ncolumns, grid):
	visited = [[False for _ in range(ncolumns)] for _ in range(nlines)]
	directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]

	def dfs(i, j, steps):
		if (
			i < 0 or i >= nlines
			or j < 0 or j >= ncolumns
			or visited[i][j]
			or grid[i][j] == -1
			or grid[i][j] == 0
			or grid[i][j] <= steps
		): return
		visited[i][j] = True
		grid[i][j] = steps
		for direction in directions:
			dfs(i + direction[0], j + direction[1], steps + 1)
		visited[i][j] = False

	for i in range(nlines):
		for j in range(ncolumns):
			if grid[i][j] == 0:
				for direction in directions:
					dfs(i + direction[0], j + direction[1], 1)
	print(grid)

grid = [
	[float('inf'),  -1,             0,              float('inf')],
	[float('inf'),  float('inf'),   float('inf'),   -1],
	[float('inf'),  -1,             float('inf'),   -1],
	[0,             -1,             float('inf'),   float('inf')],
]

result = [
	[3, -1, 0, 1],
	[2, 2, 1, -1],
	[1, -1, 2, -1],
	[0, -1, 3, 4],
]

fill_nearest_gate_distance(4, 4, grid)
assert str(grid) == str(result)


import torch
import numpy as np
import multiprocessing

from cosypose.lib3d.transform_ops import invert_T
from .bullet_scene_renderer import BulletSceneRenderer

def init_renderer(urdf_ds, preload=True):
    renderer = BulletSceneRenderer(urdf_ds=urdf_ds,
                                   preload_cache=preload,
                                   background_color=(0, 0, 0))
    return renderer

def worker_loop(worker_id, in_queue, out_queue, object_set, preload=True):
    renderer = init_renderer(object_set, preload=preload)
    while True:
        kwargs = in_queue.get()
        if kwargs is None:
            return
        obj_infos = kwargs['obj_infos']
        cam_infos = kwargs['cam_infos']
        render_depth = kwargs['render_depth']
        is_valid = np.isfinite(obj_infos[0]['TWO']).all() \
            and np.isfinite(cam_infos[0]['TWC']).all() \
            and np.isfinite(cam_infos[0]['K']).all()
        if is_valid:
            cam_obs = renderer.render_scene(cam_infos=cam_infos, obj_infos=obj_infos,
                                            render_depth=render_depth)
            images = np.stack([d['rgb'] for d in cam_obs])
            depth = np.stack([d['depth'] for d in cam_obs]) if render_depth else None
        else:
            res = cam_infos[0]['resolution']
            images = np.zeros((1, min(res), max(res), 3), dtype=np.uint8)
            depth = np.zeros((1, min(res), max(res)), dtype=np.float32)
        out_queue.put((kwargs['data_id'], images, depth))

class BulletBatchRenderer:
    def __init__(self, object_set, n_workers=8, preload_cache=True):
        self.object_set = object_set
        self.n_workers = n_workers
        self.init_plotters(preload_cache)

    def render(self, obj_infos, TCO, K, resolution=(240, 320), render_depth=False):
        TCO = torch.as_tensor(TCO).detach()
        TOC = invert_T(TCO).cpu().numpy()
        K = torch.as_tensor(K).cpu().numpy()
        bsz = len(TCO)
        assert TCO.shape == (bsz, 4, 4)
        assert K.shape == (bsz, 3, 3)

        # NOTE: Could be faster with pytorch 3.8's sharedmemory
        for n in np.arange(bsz):
            obj_info = dict(
                name=obj_infos[n]['name'],
                TWO=np.eye(4)
            )
            cam_info = dict(
                resolution=resolution,
                K=K[n],
                TWC=TOC[n],
            )
            kwargs = dict(cam_infos=[cam_info], obj_infos=[obj_info], render_depth=render_depth)
            if self.n_workers > 0:
                kwargs['data_id'] = n
                self.in_queue.put(kwargs)
            else:
                cam_obs = self.plotters[0].render_scene(**kwargs)
                images = np.stack([d['rgb'] for d in cam_obs])
                depth = np.stack([d['depth'] for d in cam_obs]) if render_depth else None
                self.out_queue.put((n, images, depth))

        images = [None for _ in np.arange(bsz)]
        depths = [None for _ in np.arange(bsz)]
        for n in np.arange(bsz):
            data_id, im, depth = self.out_queue.get()
            images[data_id] = im[0]
            if render_depth:
                depths[data_id] = depth[0]
        images = torch.as_tensor(np.stack(images, axis=0)).pin_memory().cuda(non_blocking=True)
        images = images.float().permute(0, 3, 1, 2) / 255

        if render_depth:
            depths = torch.as_tensor(np.stack(depths, axis=0)).pin_memory().cuda(non_blocking=True)
            depths = depths.float()
            return images, depths
        else:
            return images

    def init_plotters(self, preload_cache):
        self.plotters = []
        self.in_queue = multiprocessing.Queue()
        self.out_queue = multiprocessing.Queue()

        if self.n_workers > 0:
            for n in range(self.n_workers):
                plotter = multiprocessing.Process(target=worker_loop,
                                                  kwargs=dict(worker_id=n,
                                                              in_queue=self.in_queue,
                                                              out_queue=self.out_queue,
                                                              preload=preload_cache,
                                                              object_set=self.object_set))
                plotter.start()
                self.plotters.append(plotter)
        else:
            self.plotters = [init_renderer(self.object_set, preload_cache)]

    def stop(self):
        if self.n_workers > 0:
            for p in self.plotters:
                self.in_queue.put(None)
            for p in self.plotters:
                p.join()
                p.terminate()
        self.in_queue.close()
        self.out_queue.close()

    def __del__(self):
        self.stop()


import requests
from urllib.request import Request, urlopen
from bs4 import BeautifulSoup
from csv import writer

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

req = Request('https://academic.oup.com/journals/search-results?page=1&q=data+mining&SearchSourceType=1&allJournals=1', headers={'User-Agent': 'Mozilla/5.0'})
page = urlopen(req).read()
soup = BeautifulSoup(page, 'html.parser')

#print(page.status_code)
title=[] #List to store title of the article
author_complete=[] #List to store all the author names
year=[] #List to store published year
link = [] # List to store link of the title
result = soup.find('div',{'id' : 'main'})
if result is not None:
    jobs = result.find_all('div',class_='sr-list al-article-box al-normal clearfix')
    for job in jobs:
        t = job.find('h4',class_='sri-title customLink al-title')
        li = job.find('a',{'class':'article-link'})
        if None in (t,link):
            continue
        title.append(t.text.strip())
        link.append(li.get('href'))
        writter = job.find('div',class_='sri-authors al-authors-list')
        author_complete.append(writter.text.strip())
        date = job.find('div',class_='sri-date al-pub-date')
        year.append(date.text.strip())

for i in range(len(title)):
    yr = year[i].split(':')
    lst = [title[i],link[i],yr[1],author_complete[i]]
    append_list_as_row('knimbus_oxford.csv',lst)


# check items with same check id
def filter_same_checkid(ver_list, am_err_list):
    am_err_list.sort(key=lambda x: x.check_id, reverse=False)
    i = 0
    total = len(am_err_list)
    while i < total:
        cur_cid = am_err_list[i].check_id
        cur_am = 0
        cur_start = i
        # collect amount of same cid
        for j in range(cur_start, total):
            if cur_cid != am_err_list[j].check_id:
                break
            cur_am += am_err_list[j].amount
            i += 1
        # compare amount
        if cur_cid not in ver_map:
            continue
        ver_pos = ver_map[cur_cid]
        v_am = ver_list[ver_pos].amount
        if not math.isclose(v_am, cur_am, abs_tol=1e-5):
            continue
        # mark verified
        ver_list[ver_pos].map_id = 0
        ver_rid = ver_list[ver_pos].rid
        for j in range(cur_start, i):
            am_err_list[j].map_id = ver_rid


import os
from concurrent.futures import ThreadPoolExecutor
from functools import partial
from time import time

from downloader import setup_download_dir, get_links, download_link, get_logger

logger = get_logger(__name__)

def main():
    ts = time()
    client_id = os.getenv('IMGUR_CLIENT_ID')
    if not client_id:
        raise Exception("Couldn't find IMGUR_CLIENT_ID environment variable!")
    download_dir = setup_download_dir()
    links = get_links(client_id)
    dir_links = [(download_dir, link) for link in links]

    # By placing the executor inside a with block, the executors shutdown method
    # will be called cleaning up threads.
    #
    # By default, the executor sets number of workers to 5 times the number of
    # CPUs.
    with ThreadPoolExecutor() as executor:

        # Executes fn concurrently using threads on the links iterable. The
        # timeout is for the entire process, not a single call, so downloading
        # all images must complete within 30 seconds.
        executor.map(download_link, dir_links, timeout=30)

    logger.info(f'Downloading {len(links)} images took {time() - ts} seconds')

if __name__ == '__main__':
    main()


from adafruit_circuitplayground.express import circuit
import board
import neopixel
import time

pixels = neopixel.NeoPixel(board.NEOPIXEL, 10)

while True:
	temp = circuit.temperature
	if temp >= 26.00:
		pixels.fill((255,0,0))
	elif temp < 26.00:
		pixels.fill((0,0,255))
	pixels.write()
	time.sleep(0.5)
	pixels.fill((0,0,0))
	print("Temperatura: %.2f °C" %temp)
	time.sleep(1)


# -*- coding: utf-8 -*-
"""
Created on Tue Nov 11 17:04:48 2014

@author: tong
"""

import numpy as np

def sample(n):
    #%% compute uniform samples on sphere
#    The drawback is that there is always some patterns in the histogram
    #http://blog.marmakoide.org/?p=1

    golden_angle = np.pi * (3 - np.sqrt(5))
    theta = golden_angle * np.arange(n)
    z = np.linspace(1 - 1.0 / n, 1.0 / n - 1, n)
    radius = np.sqrt(1 - z * z)

    d = np.zeros((n, 3))
    d[:,0] = radius * np.cos(theta)
    d[:,1] = radius * np.sin(theta)
    d[:,2] = z

    return d


import copy
import numpy as np

def load_channels(base_path):
    # User-defined channels to use
    with open(base_path+'Experiment_Information/Channels.txt','r') as f:
        Channels = f.read().split('\n')

    Channels = [m for n,m in enumerate(Channels) if m!='']

    Marker_Names = copy.deepcopy(Channels)

    # If the user defines one of the values to None, it is not included.
    Channels = [n for n,m in enumerate(Channels) if m!='None']
    return Channels, Marker_Names