import argparse
import dataio
import torch
import random
import textwrap

def predict(model_type, model, batch, max_steps, set_rs=None):
    if model_type == "seq2seq_rnn":
        return model.greedy_predict(
            batch.encoder_inputs, batch.encoder_input_length,
            max_steps=max_steps)
    elif model_type == "seq2seq_rnn_rs":
        rs = batch.decoder_inputs[1].data[:,0]
        if set_rs is not None:
            rs.fill_(set_rs)

        return model.greedy_predict(
            batch.encoder_inputs, batch.encoder_input_length,
            max_steps=max_steps,
            extra_inputs=rs)
    elif model_type == "seq2seq_rnn_ts":
        ts = batch.decoder_inputs[1].data[:,0]
        if set_rs is not None:
            ts.fill_(set_rs)

        return model.greedy_predict(
            batch.encoder_inputs, batch.encoder_input_length,
            max_steps=max_steps,
            extra_inputs=ts)

    else:
        raise Exception("Bad model type: {}".format(model_type))

if __name__ == "__main__":
    parser = argparse.ArgumentParser("Eval seq2seq.")

    parser.add_argument("--data", required=True)

    parser.add_argument("--model", required=True)
    parser.add_argument("--max-steps", required=False, default=100, type=int)

    parser.add_argument("--gpu", default=-1, type=int, required=False)
    parser.add_argument("--batch-size", required=False, type=int, default=4)
    parser.add_argument("--set-rs", default=None, required=False, type=int)

    args = parser.parse_args()

#    random.seed(args.seed)
#    torch.manual_seed(args.seed)
    model = torch.load(args.model, map_location=lambda storage, loc: storage)

    vocab_src = model.get_meta("encoder_vocab")
    vocab_tgt = model.get_meta("decoder_vocab")

    model_type = model.get_meta("model_type")

    if model_type == "seq2seq_rnn":
        data = dataio.read_seq2seq_dataset(
            args.data, 0, vocab_src, 1, vocab_tgt,
            skip_header=True, batch_size=args.batch_size, gpu=args.gpu)
    elif model_type == "seq2seq_rnn_rs":
        data = dataio.read_seq2seq_rs_dataset(
            args.data, 0, vocab_src, 1, vocab_tgt,
            max_steps=model.get_meta("max_steps"), skip_header=True,
            batch_size=args.batch_size, gpu=args.gpu)
    elif model_type == "seq2seq_rnn_ts":
        data = dataio.read_seq2seq_ts_dataset(
            args.data, 0, vocab_src, 1, vocab_tgt,
            max_steps=model.get_meta("max_steps"), skip_header=True,
            batch_size=args.batch_size, gpu=args.gpu)

    if args.gpu > -1:
        #with torch.cuda.device(args.gpu):
            #torch.cuda.manual_seed(args.seed)
        model.cuda(args.gpu)

    model.eval()

    import random
    random.seed(34535)

    for batch in data.iter_batch():
        inputs = vocab_src.inv_lookup_matrix(batch.encoder_inputs[0].data)
        targets = vocab_tgt.inv_lookup_matrix(batch.target.data)
#        predictions = predict(
#            model_type, model, batch, max_steps=args.max_steps,
#            set_rs=args.set_rs)
#        predicted_targets = vocab_tgt.inv_lookup_matrix(predictions)

        beam_size = 64
        if model.get_meta("model_type") == "seq2seq_rnn":
            batch_beam_predictions = model.beam_search(
                batch.encoder_inputs, batch.encoder_input_length, beam_size=beam_size)
        else:
            batch_size = batch.target.size(0)
            rs = batch.decoder_inputs[1].data.new(batch_size * beam_size)

            for b in range(batch_size):
                #rs[b * beam_size: (b+1) * beam_size].fill_(batch.decoder_inputs[1].data[b,0])
                rs[b * beam_size: (b+1) * beam_size].fill_(args.set_rs)

            batch_beam_predictions = model.beam_search(
                batch.encoder_inputs, batch.encoder_input_length, extra_inputs=rs, beam_size=beam_size)

        print("")
        for i, input_tokens in enumerate(inputs):
            input_string = "           input: {}".format(
                " ".join(input_tokens))

            print(textwrap.fill(input_string, subsequent_indent=" " * 18))
            expected_string = " expected output: {}".format(
                " ".join(targets[i]))
            print(textwrap.fill(expected_string, subsequent_indent=" " * 18))
            #predicted_string = "predicted output: {}".format(
            #    " ".join(predicted_targets[i]))
            #print(textwrap.fill(predicted_string, subsequent_indent=" " * 18))
            #print("")
            beam_predictions = batch_beam_predictions[i]

            for score, beam_prediction in beam_predictions[:10]:
                if score == float("-inf"): continue
                print(score)
                print(beam_prediction)

            #attention = predictions.attention[i]
            #attention_vals, attention_idxs = torch.sort(attention, 1, True)
            x = input()

#            if x == "p":
#                for j, token in enumerate(predicted_targets[i]):
#                    print(predicted_targets[i][:j+1])
#                    for k in range(3):
#                        if attention_vals[j,k] == 0:
#                            break
#                        weight = attention_vals[j,k]
#                        index = attention_idxs[j,k]
#
#                        print(weight, index, input_tokens[index])
#                    print("")
#
#                input()


from torch import nn, optim

from autoencoder import VAELatent
from coordconv import *
from utilities import flatten, unflatten, nan_canary

import torch.nn.functional as F

class SkippableLayerNorm(nn.Module):
    def __init__(self, norm=True):
        super().__init__()
        self.norm = norm

    def forward(self, x):
        if self.norm:
            x = F.layer_norm(x, x.size()[1:])
        return x

class SimpleLayerNorm(nn.Module):
    def __init__(self, norm=True):
        super().__init__()

    def forward(self, x):
        x = F.layer_norm(x, x.size()[1:])
        return x

class WMEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = CoordConv2d(3, 32, 4, stride=2)
        self.conv2 = CoordConv2d(32, 64, 4, stride=2)
        self.conv3 = CoordConv2d(64, 128, 4, stride=2)
        self.conv4 = CoordConv2d(128, 256, 4, stride=2)  # 2x2x256
        self.activation = F.silu

    def forward(self, x):
        x = self.conv1(x)
        x = self.activation(x)
        x = F.layer_norm(x, x.size()[1:])
        x = self.activation(self.conv2(x))
        x = F.layer_norm(x, x.size()[1:])
        x = self.activation(self.conv3(x))
        x = F.layer_norm(x, x.size()[1:])
        x = self.activation(self.conv4(x))
        x = F.layer_norm(x, x.size()[1:])
        return x

class WMDecoder(nn.Module):
    def __init__(self):
        super().__init__()
        # self.deconv1 = CoordConvTranspose2d(1024, 128, 5, stride=2)
        self.deconv1 = nn.ConvTranspose2d(1024, 128, 5, stride=2)
        self.deconv2 = CoordConvTranspose2d(128, 64, 5, stride=2)
        self.deconv3 = CoordConvTranspose2d(64, 32, 6, stride=2)
        self.deconv4 = CoordConvTranspose2d(32, 16, 6, stride=2)
        self.to_rgb = CoordConv2d(16, 3, 1, stride=1)
        self.activation = F.silu

    def forward(self, input):
        x = input
        x = F.layer_norm(x, x.size()[1:])
        x = self.deconv1(x)
        x = self.activation(x)
        x = self.activation(self.deconv2(F.layer_norm(x, x.size()[1:])))
        x = self.activation(self.deconv3(F.layer_norm(x, x.size()[1:])))
        x = F.layer_norm(x, x.size()[1:])
        x = F.sigmoid(self.deconv4(x))  # activation and norm removed so it's a fully linear layer
        x = self.to_rgb(x)
        return x

class WMLatent(nn.Module):
    def __init__(self):
        super().__init__()
        self.compress = nn.Linear(1024, 256)
        self.decompress = nn.Linear(256, 1024)

    def forward(self, x):
        x = flatten(x)
        x = self.compress(x)
        latent = x
        x = self.decompress(x)
        x = unflatten(x)
        return x, latent

class WMAutoencoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = WMEncoder()
        self.latent = WMLatent()
        self.decoder = WMDecoder()

    def forward(self, x):
        x = self.encoder(x)
        #         print(x.shape)
        x, latent = self.latent(x)
        x = self.decoder(x)
        return x, latent

    def loss(self, inputs, outputs, latents):
        return F.mse_loss(inputs, outputs)

class WM_VAE(nn.Module):
    def __init__(self, latent_size=128):
        super().__init__()
        activation = nn.ReLU
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 32, 4, stride=2),
            activation(),
            nn.Conv2d(32, 64, 4, stride=2),
            activation(),
            nn.Conv2d(64, 128, 4, stride=2),
            activation(),
            nn.Conv2d(128, 256, 4, stride=2),
            activation(),
        )

        self.latent = VAELatent(latent_size, 1024)

        self.unlatent = nn.Linear(latent_size, 1024)

        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(1024, 128, 5, stride=2),
            activation(),
            nn.ConvTranspose2d(128, 64, 5, stride=2),
            activation(),
            nn.ConvTranspose2d(64, 32, 6, stride=2),
            activation(),
            nn.ConvTranspose2d(32, 3, 6, stride=2),
            nn.Sigmoid(),
        )

    def forward(self, input):
        encoded = self.encoder(input)
        latent = self.latent(encoded)
        x = self.unlatent(latent[0])
        x = unflatten(x)
        decoded = self.decoder(x)
        return decoded, latent

"""
Inputs:
    xp: observation (pov)
    xv: extra latent (vector)
    a: action

Outputs:
    rh: reward
    xph: reconstructed pov
    xvh: reconstructed vector
    h: hidden state
    z: state representation
    zh: stochastic state representation
    vh: critic
    ah: actor

Networks:
    encoder: xp -> xp_latent
    decoder: (h, z) -> xph, xvh
    recurrent: (h, z, a) -> h
    representation: (h, xp, xv) -> z
    rep_distribution h -> zh
    reward: (h, z) -> rh
    actor: z -> ah
    critic: z -> v

"""

class VisionEncoder(nn.Module):
    def __init__(self, latent_size=128, activation=nn.ELU, embed=True, norm=False, coord=False):
        super().__init__()
        conv = nn.Conv2d
        if coord:
            conv = CoordConv2d

        layers = [
            conv(3, 32, 4, stride=2),
            activation(),
            SkippableLayerNorm(norm),
            conv(32, 64, 4, stride=2),
            activation(),
            SkippableLayerNorm(norm),
            conv(64, 128, 4, stride=2),
            activation(),
            SkippableLayerNorm(norm),
            conv(128, 256, 4, stride=2),
            activation(),
            SkippableLayerNorm(norm),
            nn.Flatten(),
        ]
        if embed:
            layers.append(nn.Linear(1024, latent_size))
#             layers.append(SkippableLayerNorm(norm))
            layers.append(nn.Tanh())

        self.encoder = nn.Sequential(*layers)

    def forward(self, input):
        x = self.encoder(input)
        return x

class VisionDecoder(nn.Module):
    def __init__(self, latent_size=128, activation=nn.ELU, norm=False, coord=False):
        super().__init__()
        conv = nn.ConvTranspose2d
        if coord:
            conv = CoordConvTranspose2d

        self.decoder = nn.Sequential(
            nn.Linear(latent_size, 1024),
            nn.Unflatten(-1, (1024, 1, 1)),
            nn.ConvTranspose2d(1024, 128, 5, stride=2),
            activation(),
            SkippableLayerNorm(norm),
            conv(128, 64, 5, stride=2),
            activation(),
            SkippableLayerNorm(norm),
            conv(64, 32, 6, stride=2),
            activation(),
            SkippableLayerNorm(norm),
            conv(32, 3, 6, stride=2),
            nn.Sigmoid(),
        )

    def forward(self, input):
        x = self.decoder(input)
        return x

class CCDecoder(nn.Module):
    def __init__(self, latent_size=128, activation=nn.SiLU, norm=True):
        super().__init__()

        self.latent_size = latent_size

        self.decoder = nn.Sequential(
            CoordConv2d(latent_size, 128, 1),
            activation(),
            SkippableLayerNorm(norm),
            CoordConv2d(128, 64, 1),
            activation(),
            SkippableLayerNorm(norm),
            CoordConv2d(64, 32, 1),
            activation(),
            SkippableLayerNorm(norm),
            CoordConv2d(32, 3, 1),
            nn.Sigmoid(),
        )

    def forward(self, input):
        x = torch.tile(input.view(-1, self.latent_size, 1, 1), (64, 64))
        x = self.decoder(x)
        return x

class VisionModel(nn.Module):
    def __init__(self, latent_size=128, activation=nn.SiLU, norm=True, coord=True):
        super().__init__()
        self.norm = norm
        self.encoder = VisionEncoder(latent_size, activation, coord=coord, norm=norm)
        self.decoder = VisionDecoder(latent_size, activation, coord=coord, norm=norm)
#         self.decoder = CCDecoder(latent_size, activation, norm=norm)

    def forward(self, input):
        encoded = self.encoder(input)
        decoded = self.decoder(encoded)
        return decoded, encoded

class VisionVAE(nn.Module):
    def __init__(self, latent_size=128, activation=nn.SiLU, norm=True, coord=True):
        super().__init__()
        self.norm = norm
        self.encoder = VisionEncoder(latent_size, activation, embed=False, coord=coord, norm=norm)
        self.latent = VAELatent(latent_size, 1024)
        self.decoder = VisionDecoder(latent_size, activation, coord=coord, norm=norm)

    #         self.decoder = CCDecoder(latent_size, activation, norm=norm)

    def forward(self, input):
        encoded = self.encoder(input)
        vae_out = self.latent(encoded)
        decoded = self.decoder(vae_out[0])
        return decoded, vae_out

class ConvNeXtBlock(nn.Module):
    def __init__(self, channels, kernel_size=7, padding=3, stride=1, activation=nn.GELU, norm=SimpleLayerNorm):
        super().__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(channels, channels, kernel_size, groups=channels, padding=padding, stride=stride),
            norm(),
            nn.Conv2d(channels, channels * 4, 1),
            activation(),
            nn.Conv2d(channels * 4, channels, 1),
        )

    def forward(self, x):
        x = x + self.layer(x)
        return x


import sys
from copy import deepcopy
from typing import Optional, Tuple

import numpy as np
import torch as to
import torch.nn as nn
from tqdm import tqdm

import pyrado
from pyrado.algorithms.base import Algorithm
from pyrado.algorithms.step_based.value_based import ValueBased
from pyrado.environments.base import Env
from pyrado.exploration.stochastic_action import EpsGreedyExplStrat
from pyrado.logger.step import StepLogger
from pyrado.policies.base import Policy
from pyrado.policies.feed_back.fnn import DiscreteActQValPolicy
from pyrado.sampling.cvar_sampler import CVaRSampler
from pyrado.sampling.parallel_rollout_sampler import ParallelRolloutSampler

class DQL(ValueBased):
    """
    Deep Q-Learning (without bells and whistles)

    .. seealso::
        [1] V. Mnih et.al., "Human-level control through deep reinforcement learning", Nature, 2015
    """

    name: str = "dql"

    def __init__(
        self,
        save_dir: pyrado.PathLike,
        env: Env,
        policy: DiscreteActQValPolicy,
        memory_size: int,
        eps_init: float,
        eps_schedule_gamma: float,
        gamma: float,
        max_iter: int,
        num_updates_per_step: int,
        target_update_intvl: Optional[int] = 5,
        num_init_memory_steps: Optional[int] = None,
        min_rollouts: Optional[int] = None,
        min_steps: Optional[int] = None,
        batch_size: int = 256,
        eval_intvl: int = 100,
        max_grad_norm: float = 0.5,
        lr: float = 5e-4,
        lr_scheduler=None,
        lr_scheduler_hparam: Optional[dict] = None,
        num_workers: int = 4,
        logger: Optional[StepLogger] = None,
        use_trained_policy_for_refill: bool = False,
    ):
        r"""
        Constructor

        :param save_dir: directory to save the snapshots i.e. the results in
        :param env: the environment which the policy operates
        :param policy: (current) Q-network updated by this algorithm
        :param memory_size: number of transitions in the replay memory buffer
        :param eps_init: initial value for the probability of taking a random action, constant if `eps_schedule_gamma=1`
        :param eps_schedule_gamma: temporal discount factor for the exponential decay of epsilon
        :param gamma: temporal discount factor for the state values
        :param max_iter: maximum number of iterations (i.e. policy updates) that this algorithm runs
        :param num_updates_per_step: number of (batched) updates per algorithm steps
        :param target_update_intvl: number of iterations that pass before updating the `qfcn_targ` network
        :param num_init_memory_steps: number of samples used to initially fill the replay buffer with, pass `None` to
                                      fill the buffer completely
        :param min_rollouts: minimum number of rollouts sampled per policy update batch
        :param min_steps: minimum number of state transitions sampled per policy update batch
        :param batch_size: number of samples per policy update batch
        :param eval_intvl: interval in which the evaluation rollouts are collected, also the interval in which the
                           logger prints the summary statistics
        :param max_grad_norm: maximum L2 norm of the gradients for clipping, set to `None` to disable gradient clipping
        :param lr: (initial) learning rate for the optimizer which can be by modified by the scheduler.
                   By default, the learning rate is constant.
        :param lr_scheduler: learning rate scheduler that does one step per epoch (pass through the whole data set)
        :param lr_scheduler_hparam: hyper-parameters for the learning rate scheduler
        :param num_workers: number of environments for parallel sampling
        :param logger: logger for every step of the algorithm, if `None` the default logger will be created
        :param use_trained_policy_for_refill: whether to use the trained policy instead of a dummy policy to refill the
                                              replay buffer after resets
        """
        if not isinstance(policy, DiscreteActQValPolicy):
            raise pyrado.TypeErr(given=policy, expected_type=DiscreteActQValPolicy)

        # Call ValueBased's constructor
        super().__init__(
            save_dir=save_dir,
            env=env,
            policy=policy,
            memory_size=memory_size,
            gamma=gamma,
            max_iter=max_iter,
            num_updates_per_step=num_updates_per_step,
            target_update_intvl=target_update_intvl,
            num_init_memory_steps=num_init_memory_steps,
            min_rollouts=min_rollouts,
            min_steps=min_steps,
            batch_size=batch_size,
            eval_intvl=eval_intvl,
            max_grad_norm=max_grad_norm,
            num_workers=num_workers,
            logger=logger,
            use_trained_policy_for_refill=use_trained_policy_for_refill,
        )

        self.qfcn_targ = deepcopy(self._policy).eval()  # will not be trained using the optimizer
        self.eps = eps_init

        # Create sampler for exploration during training
        self._expl_strat = EpsGreedyExplStrat(self._policy, eps_init, eps_schedule_gamma)
        self._sampler = ParallelRolloutSampler(
            self._env,
            self._expl_strat,
            num_workers=num_workers if min_steps != 1 else 1,
            min_steps=min_steps,
            min_rollouts=min_rollouts,
        )

        # Q-function optimizer
        self.optim = to.optim.RMSprop([{"params": self._policy.parameters()}], lr=lr)

        # Learning rate scheduler
        self._lr_scheduler = lr_scheduler
        self._lr_scheduler_hparam = lr_scheduler_hparam
        if lr_scheduler is not None:
            self._lr_scheduler = lr_scheduler(self.optim, **lr_scheduler_hparam)

    @property
    def sampler(self) -> ParallelRolloutSampler:
        return self._sampler

    @sampler.setter
    def sampler(self, sampler: ParallelRolloutSampler):
        if not isinstance(sampler, (ParallelRolloutSampler, CVaRSampler)):
            raise pyrado.TypeErr(given=sampler, expected_type=(ParallelRolloutSampler, CVaRSampler))
        self._sampler = sampler

    @staticmethod
    def loss_fcn(q_vals: to.Tensor, expected_q_vals: to.Tensor) -> to.Tensor:
        r"""
        The Huber loss function on the one-step TD error $\delta = Q(s,a) - (r + \gamma \max_a Q(s^\prime, a))$.

        :param q_vals: state-action values $Q(s,a)$, from policy network
        :param expected_q_vals: expected state-action values $r + \gamma \max_a Q(s^\prime, a)$, from target network
        :return: loss value
        """
        return nn.functional.smooth_l1_loss(q_vals, expected_q_vals)

    def update(self):
        """Update the policy's and qfcn_targ Q-function's parameters on transitions sampled from the replay memory."""
        losses = to.zeros(self.num_batch_updates)
        policy_grad_norm = to.zeros(self.num_batch_updates)

        for b in tqdm(
            range(self.num_batch_updates),
            total=self.num_batch_updates,
            desc=f"Updating",
            unit="batches",
            file=sys.stdout,
            leave=False,
        ):

            # Sample steps and the associated next step from the replay memory
            steps, next_steps = self._memory.sample(self.batch_size)
            steps.torch(data_type=to.get_default_dtype())
            next_steps.torch(data_type=to.get_default_dtype())

            # Create masks for the non-final observations
            not_done = to.from_numpy(1.0 - steps.done).to(device=self.policy.device, dtype=to.get_default_dtype())

            # Compute the state-action values Q(s,a) using the current DQN policy
            q_vals = self.expl_strat.policy.q_values_argmax(steps.observations)

            # Compute the second term of TD-error
            with to.no_grad():
                next_v_vals = self.qfcn_targ.q_values_argmax(next_steps.observations)
                expected_q_val = steps.rewards.to(self.policy.device) + not_done * self.gamma * next_v_vals

            # Compute the loss, clip the gradients if desired, and do one optimization step
            loss = DQL.loss_fcn(q_vals, expected_q_val)
            losses[b] = loss.data
            self.optim.zero_grad()
            loss.backward()
            policy_grad_norm[b] = Algorithm.clip_grad(self.expl_strat.policy, self.max_grad_norm)
            self.optim.step()

            # Update the qfcn_targ network by copying all weights and biases from the DQN policy
            if (self._curr_iter * self.num_batch_updates + b) % self.target_update_intvl == 0:
                self.qfcn_targ.load_state_dict(self.expl_strat.policy.state_dict())

        # Schedule the exploration parameter epsilon
        self.expl_strat.schedule_eps(self._curr_iter)

        # Update the learning rate if a scheduler has been specified
        if self._lr_scheduler is not None:
            self._lr_scheduler.step()

        # Logging
        with to.no_grad():
            self.logger.add_value("loss after", to.mean(losses), 4)
        self.logger.add_value("expl strat eps", self.expl_strat.eps, 4)
        self.logger.add_value("avg grad norm policy", to.mean(policy_grad_norm), 4)
        if self._lr_scheduler is not None:
            self.logger.add_value("avg lr", np.mean(self._lr_scheduler.get_last_lr()), 6)

    def reset(self, seed: Optional[int] = None):
        # Reset samplers, replay memory, exploration strategy, internal variables and the random seeds
        super().reset(seed)

        # Reset the learning rate scheduler
        if self._lr_scheduler is not None:
            self._lr_scheduler.last_epoch = -1

    def init_modules(self, warmstart: bool, suffix: str = "", prefix: str = "", **kwargs):
        # Initialize the policy
        super().init_modules(warmstart, suffix, prefix, **kwargs)

        if prefix == "":
            prefix = f"iter_{self._curr_iter - 1}"

        tpi = kwargs.get("target_param_init", None)

        if warmstart and tpi is not None:
            self.qfcn_targ.init_param(tpi)
        elif warmstart and tpi is None and self._curr_iter > 0:
            self.qfcn_targ = pyrado.load(
                "qfcn_target.pt", self.save_dir, prefix=prefix, suffix=suffix, obj=self.qfcn_targ
            )
        else:
            # Reset the target Q-function
            self.qfcn_targ.init_param()

    def save_snapshot(self, meta_info: dict = None):
        super().save_snapshot(meta_info)

        if meta_info is None:
            # This algorithm instance is not a subroutine of another algorithm
            pyrado.save(self.qfcn_targ, "qfcn_target.pt", self.save_dir, use_state_dict=True)

        else:
            # This algorithm instance is a subroutine of another algorithm
            pyrado.save(
                self.qfcn_targ,
                "qfcn_target.pt",
                self.save_dir,
                prefix=meta_info.get("prefix", ""),
                suffix=meta_info.get("suffix", ""),
                use_state_dict=True,
            )

    def load_snapshot(self, parsed_args) -> Tuple[Env, Policy, dict]:
        env, policy, extra = super().load_snapshot(parsed_args)

        # Algorithm specific
        ex_dir = self._save_dir or getattr(parsed_args, "dir", None)
        if self.name == "dql":
            extra["qfcn_target"] = pyrado.load("qfcn_target.pt", ex_dir, obj=self.qfcn_targ, verbose=True)

        return env, policy, extra


import pytest
import trio

from libp2p.tools.factories import HostFactory
from libp2p.tools.interop.utils import connect

@pytest.mark.trio
async def test_connect(security_protocol, p2pds):
    async with HostFactory.create_batch_and_listen(
        1, security_protocol=security_protocol
    ) as hosts:
        p2pd = p2pds[0]
        host = hosts[0]
        assert len(await p2pd.control.list_peers()) == 0
        # Test: connect from Py
        await connect(host, p2pd)
        assert len(await p2pd.control.list_peers()) == 1
        # Test: `disconnect` from Py
        await host.disconnect(p2pd.peer_id)
        assert len(await p2pd.control.list_peers()) == 0
        # Test: connect from Go
        await connect(p2pd, host)
        assert len(host.get_network().connections) == 1
        # Test: `disconnect` from Go
        await p2pd.control.disconnect(host.get_id())
        await trio.sleep(0.01)
        assert len(host.get_network().connections) == 0


#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
:py:mod:`rect.py` Rectify the pressure grid
-------------------------------------------

Converts a NetCDF file from sigma coordinates to pressure coordinates,
then interpolates all arrays to a rectified pressure grid. The user
can choose between linear interpolation (slow) and cubic interpolation
(really, really slow).

'''

from __future__ import division, print_function, absolute_import, unicode_literals
from .utils import array
from .constants import PREF
import numpy as np
from scipy.interpolate import interp1d
from scipy.io import netcdf
import sys

__all__ = ['Rectify']

def Rectify(file, interpolation = 'linear', rect_file = None):
  '''

  '''

  # Open the file
  f = netcdf.netcdf_file(file)

  # Get the variables we need to modify
  pfvars = []
  phvars = []
  newvars = {}
  for name in list(f.variables.keys()):
    if f.variables[name].dimensions == ('time', 'pfull', 'lat', 'lon'):
      pfvars.append(name)
    elif f.variables[name].dimensions == ('time', 'phalf', 'lat', 'lon'):
      phvars.append(name)

  # Create the true irregular pressure grid from the sigma coordinate
  # This is a 4D array (time, pressure, lat, lon). Units are mb.
  print("Computing new pressure grid...")
  pfull_ = np.stack([f.variables['ps'][:] * p / PREF for p in f.variables['pfull'][:]], axis = 1)
  phalf_ = np.stack([f.variables['ps'][:] * p / PREF for p in f.variables['phalf'][:]], axis = 1)

  # Rectify pfull
  pfull = np.nanmedian(pfull_, axis = (0,2,3))
  phalf = np.nanmedian(phalf_, axis = (0,2,3))

  # Get the dimensions
  ntime = f.variables['time'].shape[0]
  nlat = f.variables['lat'].shape[0]
  nlon = f.variables['lon'].shape[0]

  # Rectify each of the 4D phalf variables and interpolate to the pfull grid
  newvars.update({'phalf': phalf, 'phalf_unit': np.array('mb'), 'phalf_dims': np.array(['phalf']), 'phalf_desc': np.array('True half-pressure grid')})
  for n, name in enumerate(phvars):
    var_old = np.array(f.variables[name][:])
    var_new = np.zeros((ntime, len(pfull), nlat, nlon))
    for i in range(ntime):
      sys.stdout.write('\rVariable %d/%d: Processing time %d/%d...' % (n + 1, len(pfvars) + len(phvars), i + 1, ntime))
      sys.stdout.flush()
      for k in range(nlat):
        for l in range(nlon):
          var_new[i,:,k,l] = interp1d(phalf_[i,:,k,l], var_old[i,:,k,l], kind=interpolation,
                                      bounds_error=False, fill_value="extrapolate")(pfull)
    newvars.update({name: var_new})
    newvars.update({name + '_unit': np.array(f.variables[name].units.decode('utf-8'))})
    newvars.update({name + '_dims': np.array(['time', 'pfull', 'lat', 'lon'])})
    newvars.update({name + '_desc': np.array(f.variables[name].long_name.decode('utf-8'))})

  # Rectify each of the 4D pfull variables
  newvars.update({'pfull': pfull, 'pfull_unit': np.array('mb'), 'pfull_dims': np.array(['pfull']), 'pfull_desc': np.array('True pressure grid')})
  for n, name in enumerate(pfvars):
    var = np.array(f.variables[name][:])
    for i in range(ntime):
      sys.stdout.write('\rVariable %d/%d: Processing time %d/%d...' % (len(phvars) + n + 1, len(pfvars) + len(phvars), i + 1, ntime))
      sys.stdout.flush()
      for k in range(nlat):
        for l in range(nlon):
          var[i,:,k,l] = interp1d(pfull_[i,:,k,l], var[i,:,k,l], kind=interpolation,
                                  bounds_error=False, fill_value="extrapolate")(pfull)
    newvars.update({name: var})
    newvars.update({name + '_unit': np.array(f.variables[name].units.decode('utf-8'))})
    newvars.update({name + '_dims': np.array(['time', 'pfull', 'lat', 'lon'])})
    newvars.update({name + '_desc': np.array(f.variables[name].long_name.decode('utf-8'))})

  # Close original file
  f.close()

  # Create a new NetCDF file with just the rectified grids
  if rect_file is None:
    rect_file = '%s.rect.npz' % file[:-3]
  print("\nSaving to `%s`..." % rect_file)
  np.savez(rect_file, **newvars)

  f.close()
  print("Done!")


import sys

import folium
import joblib
import numpy as np
import tensorflow as tf
from matplotlib import pyplot as plt
from sklearn.metrics.pairwise import polynomial_kernel
from tqdm import tqdm

from lstmwgan import WGANGP
from preprocessing.utils import decodeTrajectories

latent_dim = 100
n_epochs = 30
n_samples = 100

if __name__ == '__main__':
    real_trajectories = np.load('data/preprocessed/train_gl.npy', allow_pickle=True)
    real_trajectories = real_trajectories[:n_samples]
    gen_inputs = []
    random_latent_vectors = noise = tf.random.normal((tf.shape(real_trajectories)[0], latent_dim), 0, 1)
    gen_inputs.append(real_trajectories)
    gen_inputs.append(random_latent_vectors)
    gan = WGANGP()
    gan.generator.load_weights('parameters/G_model_gl_' + str(n_epochs) + '.h5')  # params/G_model_100.h5
    generated_trajectories = gan.generator(gen_inputs).numpy()
    scaler = joblib.load("preprocessing/scaler.pkl")
    gen_lat, gen_lng = decodeTrajectories(generated_trajectories, scaler)
    points = list(zip(gen_lat, gen_lng))
    # my_map = folium.Map(location=[46.3615142, 6.399388], zoom_start=10)
    my_map = folium.Map(location=[39.9075, 116.39723], zoom_start=14)
    folium.PolyLine(points, color="red", weight=2.5, opacity=1).add_to(my_map)
    my_map.save(f"MDC_gen.html")
    print("Plot created")

    # -----------------HEATMAP-------------------
    xLabel = "Longitude"
    yLabel = "Latitude"
    real_lat, real_lng = decodeTrajectories(real_trajectories, scaler)
    heatmap, xedges, yedges = np.histogram2d(gen_lng, gen_lat, bins=20)
    heatmap2, xedges2, yedges2 = np.histogram2d(real_lng, real_lat, bins=20)

    extent = [0, 50, 0, 50]
    fig, (ax1, ax2) = plt.subplots(ncols=2)
    im1 = ax1.imshow(heatmap.T, extent=extent, origin='lower')
    im2 = ax2.imshow(heatmap2.T, extent=extent, origin='lower')
    plt.show()

    # -----------------MMD-------------------
    """
    def compute_kernel(x, y):
        x_size = tf.shape(x)[0]
        y_size = tf.shape(y)[0]
        dim = tf.shape(x)[1]
        tiled_x = tf.tile(tf.reshape(x, tf.stack([x_size, 1, dim])), tf.stack([1, y_size, 1]))
        tiled_y = tf.tile(tf.reshape(y, tf.stack([1, y_size, dim])), tf.stack([x_size, 1, 1]))
        return tf.exp(-tf.reduce_mean(tf.square(tiled_x - tiled_y), axis=2) / tf.cast(dim, tf.float32))

    def compute_mmd(x, y):  # [batch_size, z_dim] [batch_size, z_dim]
        x_kernel = compute_kernel(x, x)
        y_kernel = compute_kernel(y, y)
        xy_kernel = compute_kernel(x, y)
        return tf.reduce_mean(x_kernel) + tf.reduce_mean(y_kernel) - 2 * tf.reduce_mean(xy_kernel)

    real_trajectories = np.float32(real_trajectories)
    mmds = []
    for index in range(0, n_samples):
        mmd = compute_mmd(real_trajectories[index], generated_trajectories[index])
        mmds.append(mmd.numpy())
    mean_mmd = np.mean(mmds)
    print(mean_mmd)
"""

    def polynomial_mmd_averages(codes_g, codes_r, n_subsets=50, subset_size=1000,
                                ret_var=True, output=sys.stdout, **kernel_args):
        m = min(codes_g.shape[0], codes_r.shape[0])
        mmds = np.zeros(n_subsets)
        if ret_var:
            vars = np.zeros(n_subsets)
        choice = np.random.choice

        with tqdm(range(n_subsets), desc='MMD', file=output) as bar:
            for i in bar:
                g = codes_g[choice(len(codes_g), subset_size, replace=False)]
                r = codes_r[choice(len(codes_r), subset_size, replace=False)]
                o = polynomial_mmd(g, r, **kernel_args, var_at_m=m, ret_var=ret_var)
                if ret_var:
                    mmds[i], vars[i] = o
                else:
                    mmds[i] = o
                bar.set_postfix({'mean': mmds[:i + 1].mean()})
        return (mmds, vars) if ret_var else mmds

    def polynomial_mmd(codes_g, codes_r, degree=3, gamma=None, coef0=1,
                       var_at_m=None, ret_var=True):
        # use  k(x, y) = (gamma <x, y> + coef0)^degree
        # default gamma is 1 / dim
        X = codes_g
        Y = codes_r

        K_XX = polynomial_kernel(X, degree=degree, gamma=gamma, coef0=coef0)
        K_YY = polynomial_kernel(Y, degree=degree, gamma=gamma, coef0=coef0)
        K_XY = polynomial_kernel(X, Y, degree=degree, gamma=gamma, coef0=coef0)

        return _mmd2_and_variance(K_XX, K_XY, K_YY,
                                  var_at_m=var_at_m, ret_var=ret_var)

    def _sqn(arr):
        flat = np.ravel(arr)
        return flat.dot(flat)

    def _mmd2_and_variance(K_XX, K_XY, K_YY, unit_diagonal=False,
                           mmd_est='unbiased', block_size=1024,
                           var_at_m=None, ret_var=True):
        # based on
        # https://github.com/dougalsutherland/opt-mmd/blob/master/two_sample/mmd.py
        # but changed to not compute the full kernel matrix at once
        m = K_XX.shape[0]
        assert K_XX.shape == (m, m)
        assert K_XY.shape == (m, m)
        assert K_YY.shape == (m, m)
        if var_at_m is None:
            var_at_m = m

        # Get the various sums of kernels that we'll use
        # Kts drop the diagonal, but we don't need to compute them explicitly
        if unit_diagonal:
            diag_X = diag_Y = 1
            sum_diag_X = sum_diag_Y = m
            sum_diag2_X = sum_diag2_Y = m
        else:
            diag_X = np.diagonal(K_XX)
            diag_Y = np.diagonal(K_YY)

            sum_diag_X = diag_X.sum()
            sum_diag_Y = diag_Y.sum()

            sum_diag2_X = _sqn(diag_X)
            sum_diag2_Y = _sqn(diag_Y)

        Kt_XX_sums = K_XX.sum(axis=1) - diag_X
        Kt_YY_sums = K_YY.sum(axis=1) - diag_Y
        K_XY_sums_0 = K_XY.sum(axis=0)
        K_XY_sums_1 = K_XY.sum(axis=1)

        Kt_XX_sum = Kt_XX_sums.sum()
        Kt_YY_sum = Kt_YY_sums.sum()
        K_XY_sum = K_XY_sums_0.sum()

        if mmd_est == 'biased':
            mmd2 = ((Kt_XX_sum + sum_diag_X) / (m * m)
                    + (Kt_YY_sum + sum_diag_Y) / (m * m)
                    - 2 * K_XY_sum / (m * m))
        else:
            assert mmd_est in {'unbiased', 'u-statistic'}
            mmd2 = (Kt_XX_sum + Kt_YY_sum) / (m * (m - 1))
            if mmd_est == 'unbiased':
                mmd2 -= 2 * K_XY_sum / (m * m)
            else:
                mmd2 -= 2 * (K_XY_sum - np.trace(K_XY)) / (m * (m - 1))

        if not ret_var:
            return mmd2

        Kt_XX_2_sum = _sqn(K_XX) - sum_diag2_X
        Kt_YY_2_sum = _sqn(K_YY) - sum_diag2_Y
        K_XY_2_sum = _sqn(K_XY)

        dot_XX_XY = Kt_XX_sums.dot(K_XY_sums_1)
        dot_YY_YX = Kt_YY_sums.dot(K_XY_sums_0)

        m1 = m - 1
        m2 = m - 2
        zeta1_est = (
                1 / (m * m1 * m2) * (
                _sqn(Kt_XX_sums) - Kt_XX_2_sum + _sqn(Kt_YY_sums) - Kt_YY_2_sum)
                - 1 / (m * m1) ** 2 * (Kt_XX_sum ** 2 + Kt_YY_sum ** 2)
                + 1 / (m * m * m1) * (
                        _sqn(K_XY_sums_1) + _sqn(K_XY_sums_0) - 2 * K_XY_2_sum)
                - 2 / m ** 4 * K_XY_sum ** 2
                - 2 / (m * m * m1) * (dot_XX_XY + dot_YY_YX)
                + 2 / (m ** 3 * m1) * (Kt_XX_sum + Kt_YY_sum) * K_XY_sum
        )
        zeta2_est = (
                1 / (m * m1) * (Kt_XX_2_sum + Kt_YY_2_sum)
                - 1 / (m * m1) ** 2 * (Kt_XX_sum ** 2 + Kt_YY_sum ** 2)
                + 2 / (m * m) * K_XY_2_sum
                - 2 / m ** 4 * K_XY_sum ** 2
                - 4 / (m * m * m1) * (dot_XX_XY + dot_YY_YX)
                + 4 / (m ** 3 * m1) * (Kt_XX_sum + Kt_YY_sum) * K_XY_sum
        )
        var_est = (4 * (var_at_m - 2) / (var_at_m * (var_at_m - 1)) * zeta1_est
                   + 2 / (var_at_m * (var_at_m - 1)) * zeta2_est)

        return mmd2, var_est

    output = {}
    generated_trajectories = np.reshape(generated_trajectories,
                                        [generated_trajectories.shape[0], generated_trajectories.shape[1]])
    real_trajectories = np.reshape(real_trajectories,
                                   [real_trajectories.shape[0], real_trajectories.shape[1]])
    ret = polynomial_mmd_averages(
        generated_trajectories, real_trajectories, degree=3, gamma=None,
        coef0=1, ret_var=False,
        n_subsets=100, subset_size=100)
    if False:
        output['mmd2'], output['mmd2_var'] = mmd2s, vars = ret
    else:
        output['mmd2'] = mmd2s = ret
    print("mean MMD^2 estimate:", mmd2s.mean())
    # print("std MMD^2 estimate:", mmd2s.std())
    # print("MMD^2 estimates:", mmd2s, sep='\n')


# 提出する時にはこの部分は含めない
import io,sys
with open("input.txt") as txt_opend:
    TXT_INPUT = txt_opend.read()
sys.stdin=io.StringIO(TXT_INPUT)
# --------------------------------------------------------
N,M = map(int, input().split())

for k in range(10000000):
    if -(-k//N) >= M:
        print(k)
        break


from django.contrib.auth import get_user_model
from django.contrib.auth.models import Permission
from django.test import TestCase

from lava_scheduler_app.models import (
    GroupDevicePermission,
    GroupDeviceTypePermission,
    TestJob,
)
from lava_server.backends import GroupPermissionBackend, is_object_supported
from tests.lava_scheduler_app.test_submission import TestCaseWithFactory

User = get_user_model()

class BackendAuthTest(TestCaseWithFactory):
    def setUp(self):
        super().setUp()
        self.group = self.factory.make_group()
        self.user = self.factory.make_user()
        self.user.groups.add(self.group)

        self.definition = self.factory.make_job_data_from_file(
            "qemu-pipeline-first-job.yaml"
        )
        self.device_type = self.factory.make_device_type(name="qemu")
        self.device = self.factory.make_device(
            device_type=self.device_type, hostname="qemu-1"
        )
        self.job = TestJob.from_yaml_and_user(self.definition, self.user)
        self.backend = GroupPermissionBackend()

    def tearDown(self):
        super().tearDown()
        GroupDeviceTypePermission.objects.all().delete()
        GroupDevicePermission.objects.all().delete()

    def test_is_object_supported(self):
        self.assertTrue(is_object_supported(self.device))
        self.assertTrue(is_object_supported(self.device_type))
        self.assertFalse(is_object_supported(self.user))
        self.assertFalse(is_object_supported(self.group))
        self.assertFalse(is_object_supported(self.job))

    def test_is_object_supported_none(self):
        self.assertFalse(is_object_supported(None))

    def test_get_all_permissions(self):
        GroupDevicePermission.objects.assign_perm(
            "view_device", self.group, self.device
        )
        self.assertEqual(
            {"view_device"},
            set(self.backend.get_all_permissions(self.user, self.device)),
        )
        GroupDevicePermission.objects.assign_perm(
            "change_device", self.group, self.device
        )
        self.assertEqual(
            {"change_device", "submit_to_device", "view_device"},
            set(self.backend.get_all_permissions(self.user, self.device)),
        )

        GroupDeviceTypePermission.objects.assign_perm(
            "view_devicetype", self.group, self.device_type
        )
        self.assertEqual(
            {"view_devicetype"},
            set(self.backend.get_all_permissions(self.user, self.device_type)),
        )
        GroupDeviceTypePermission.objects.assign_perm(
            "change_devicetype", self.group, self.device_type
        )
        self.assertEqual(
            {"change_devicetype", "view_devicetype", "submit_to_devicetype"},
            self.backend.get_all_permissions(self.user, self.device_type),
        )

    def test_has_perm(self):
        GroupDevicePermission.objects.assign_perm(
            "change_device", self.group, self.device
        )
        self.assertTrue(
            self.backend.has_perm(
                self.user, "lava_scheduler_app.change_device", self.device
            )
        )
        self.assertTrue(
            self.backend.has_perm(
                self.user, "lava_scheduler_app.submit_to_device", self.device
            )
        )
        self.assertFalse(
            self.backend.has_perm(
                self.user, "lava_scheduler_app.view_devicetype", self.device_type
            )
        )

    def test_has_global_perm(self):
        user = self.factory.make_user()
        user.user_permissions.add(Permission.objects.get(codename="change_device"))
        self.assertTrue(
            self.backend.has_perm(user, "lava_scheduler_app.change_device", self.device)
        )

    def test_has_perm_wrong_app_label(self):
        with TestCase.assertRaises(self, ValueError):
            self.backend.has_perm(
                self.user, "lava_results_app.change_device", self.device
            )


# -*- coding: utf-8 -*-
# @Author: durban.zhang
# @Date:   2019-11-11 11:25:10
# @Last Modified by:   durban.zhang
# @Last Modified time: 2019-12-06 09:15:41

import os
from flask import (
    Blueprint, flash, g, redirect,
    request, render_template,
    url_for,
    current_app,
    send_from_directory
)

from werkzeug import secure_filename
from werkzeug.exceptions import abort
from baby.views.auth import login_required
from baby.db import get_db
from baby.exception import InvalidUsage

bp = Blueprint('blog', __name__, url_prefix='/post')

def get_posts(page, pagesize):
    db = get_db()
    posts = db.execute(
        'SELECT p.id, pc.category_id, c.name AS category_name,'
        ' p.title, p.body, p.created, p.author_id, u.username'
        ' FROM post p'
        ' JOIN user u ON p.author_id = u.id'
        ' JOIN post_category pc ON pc.post_id = p.id'
        ' JOIN category c ON c.id = pc.category_id'
        ' ORDER BY p.created DESC'
        ' LIMIT ?, ?', ((page - 1) * pagesize, pagesize)
    ).fetchall()

    for n, p in enumerate(posts):
        o = dict(p)

        o['tag'] = db.execute(
            'SELECT name FROM tag WHERE post_id = ?', (o['id'],)).fetchall()

        current_app.logger.info(o)
        posts[n] = o

    prev_page = (page - 1) if (page - 1) else ''
    next_page = (page + 1) if len(posts) >= pagesize else ''

    return posts, page, prev_page, next_page, pagesize

@bp.route('/item')
def index_item():
    page = request.args.get('page', 1, type=int)
    page = int(page)
    pagesize = 5

    posts, page, prev_page, next_page, pagesize = get_posts(page, pagesize)

    return render_template(
        'blog/index_item.j2',
        posts=posts,
        page=page,
        prev_page=prev_page,
        next_page=next_page,
        pagesize=pagesize
    )

@bp.route('/')
def index():
    page = request.args.get('page', 1, type=int)
    page = int(page)
    pagesize = 5

    posts, page, prev_page, next_page, pagesize = get_posts(page, pagesize)

    return render_template(
        'blog/index.j2',
        posts=posts,
        page=page,
        prev_page=prev_page,
        next_page=next_page,
        pagesize=pagesize
    )

@bp.route('/create', methods=['GET', 'POST'])
@login_required
def create():
    if request.method == 'POST':
        title = request.form['title'] if 'title' in request.form else ''
        body = request.form['body'] if 'body' in request.form else ''
        date = request.form['date'] if 'date' in request.form else ''
        tag = request.form['tag'] if 'tag' in request.form else ''
        category_id = request.form[
            'category_id'] if 'category_id' in request.form else ''

        error = None

        if not title:
            error = 'Title is required'
        elif not category_id:
            error = 'Category is required'
        elif not tag:
            error = 'Tag is required'

        if error is not None:
            flash(error)
        else:
            db = get_db()
            cur = db.cursor()
            cur.execute(
                'INSERT INTO post (title, body, author_id, the_date)'
                ' VALUES (?,?,?,?)',
                (title, body, g.user['id'], date)
            )

            db.commit()

            # tag insert
            insert_id = cur.lastrowid
            tags = tag.split(',')

            cur.executemany(
                'INSERT INTO tag (post_id, name)'
                ' VALUES (?,?)', tag_generator(insert_id, tags)
            )

            db.commit()

            # category insert
            cur.execute(
                'INSERT INTO post_category (post_id, category_id)'
                ' VALUES (?,?)',
                (insert_id, category_id)
            )

            db.commit()

            return redirect(url_for('blog.index'))

    category = get_category()

    return render_template(
        'blog/create.j2',
        category=category
    )

def tag_generator(insert_id, tags):
    for c in tags:
        yield (insert_id, c)

def get_category():
    db = get_db()
    category = db.execute('SELECT * FROM category ORDER BY id DESC').fetchall()

    return category

def get_post(id, check_author=True):
    db = get_db()

    post = db.execute(
        'SELECT p.id, pc.category_id,'
        ' the_date, title, body, p.created, author_id, username'
        ' FROM post p'
        ' JOIN user u ON p.author_id = u.id'
        ' JOIN post_category pc ON pc.post_id = p.id'
        ' WHERE p.id = ?',
        (id,)
    ).fetchone()

    if post is None:
        abort(404, "Post id {0} doesn't exist".format(id))

    if check_author and post['author_id'] != g.user['id']:
        abort(403)

    post = dict(post)

    tags = db.execute(
        'SELECT name FROM tag WHERE post_id = ?', (post['id'],)).fetchall()

    post['tag'] = ','.join((x['name'] for x in tags))

    current_app.logger.info(post)
    return post

@bp.route('/<int:id>/update', methods=['GET', 'POST'])
@login_required
def update(id):
    post = get_post(id)

    if request.method == 'POST':
        title = request.form['title'] if 'title' in request.form else ''
        body = request.form['body'] if 'body' in request.form else ''
        date = request.form['date'] if 'date' in request.form else ''
        tag = request.form['tag'] if 'tag' in request.form else ''
        category_id = request.form[
            'category_id'] if 'category_id' in request.form else ''
        error = None

        if not title:
            error = 'Title is required'
        elif not category_id:
            error = 'Category is request'
        elif not tag:
            error = 'Tag is required'

        if error is not None:
            flash(error)
        else:
            db = get_db()
            cur = db.cursor()
            cur.execute(
                'UPDATE post SET title = ?, body = ?, the_date = ?'
                ' WHERE id = ?',
                (title, body, date, id)
            )

            db.commit()

            # 清空原来的关联
            # 建立新的关联
            #

            cur.execute(
                'DELETE FROM tag WHERE post_id = ?', (id,)
            )

            tags = tag.split(',')
            cur.executemany(
                'INSERT INTO tag (post_id, name)'
                ' VALUES (?,?)', tag_generator(id, tags)
            )

            cur.execute(
                'UPDATE post_category SET category_id = ?'
                ' WHERE post_id = ?',
                (category_id, id,))

            db.commit()

            return redirect(url_for('blog.index'))

    category = get_category()

    return render_template(
        'blog/update.j2',
        post=post,
        category=category
    )

@bp.route('/<int:id>/delete', methods=['GET', 'POST'])
@login_required
def delete(id):
    get_post(id)
    db = get_db()
    db.execute(
        'DELETE FROM post where id = ?', (id,)
    )
    db.commit()
    return redirect(url_for('blog.index'))

@bp.route('/<int:id>/show', methods=['GET'])
def show(id):
    post = get_post(id, False)

    return render_template('blog/show.j2', post=post)

@bp.route('/tag/<string:name>', methods=['GET'])
def tag(name):
    db = get_db()
    posts = db.execute(
        'SELECT p.id, title, body, p.created, author_id, username'
        ' FROM tag AS t'
        ' JOIN post AS p ON p.id = t.post_id'
        ' JOIN user u ON p.author_id = u.id'
        ' WHERE t.name = ?'
        ' ORDER BY p.created DESC', (name,)
    ).fetchall()

    for n, p in enumerate(posts):
        o = dict(p)

        o['tag'] = db.execute(
            'SELECT name FROM tag WHERE post_id = ?', (o['id'],)).fetchall()

        current_app.logger.info(o)
        posts[n] = o

        return render_template('blog/tag.j2', posts=posts)

@bp.route('/read', methods=['GET'])
def read_sum():
    raise InvalidUsage('read sum function developing...', status_code=404)

def allowed_file(filename):
    return '.' in filename and \
        filename.rsplit('.', 1)[1].lower() in current_app.config[
            'UPLOAD_ALLOWED_EXTENSIONS']

def allowed_length(length):
    return length <= current_app.config[
        'UPLOAD_MAX_LENGTH']

@bp.route('/icon', methods=['GET', 'POST'])
@login_required
def icon():
    if request.method == 'POST':
        print(request.files)
        if 'icon' not in request.files:
            flash('No file part')
            return redirect(request.url)

        file = request.files['icon']
        if file.filename == '':
            flash('No selected file')
            return redirect(request.url)

        if file and \
                allowed_file(file.filename) and \
                allowed_length(file.content_length):
            filename = secure_filename(file.filename)
            file.save(os.path.join(current_app.config['UPLOAD_DIR'], filename))
            return redirect(url_for('blog.uploaded_icon', filename=filename))
    else:
        return render_template('blog/icon.j2')

@bp.route('/uploaded/icon/<string:filename>', methods=['GET'])
def uploaded_icon(filename):
    return send_from_directory(current_app.config['UPLOAD_DIR'], filename)


# Origin: RhinoSecurityLabs, Pacu <https://github.com/RhinoSecurityLabs/pacu>
# Purpose: Utilities and additional session contextualized config settings
#

import contextlib
import os

from typing import Optional, cast, Callable, Generator, IO, Any
from pathlib import Path
from stratustryke import settings

active_session: Optional[Callable] = None

class StratustrykeException(Exception):
    pass

class FrameworkError(Exception):
	def __init__(self, msg):
		self.msg = msg

	def __str__(self):
		return repr(self.msg)

class FrameworkConfigurationError(FrameworkError):
	pass

class FrameworkRuntimeError(FrameworkError):
	pass

def strip_lines(text: str) -> str:
    '''
    Remove newline characters and tabs from supplied string.
    :param: text <str> string to remove newlines and tabs from.
    :rtype: <str> input string without newlines or tabs
    '''
    out = []
    for line in text.splitlines():
        out.append(line.split('\t '))
    return ' '.join(out)

def home_dir() -> Path:
    '''
    Get stratustryke app home directory
    :rtype: <Path> home directory of stratustryke app
    '''
    path = Path('~/.local/share/stratustryke').expanduser().absolute()
    os.makedirs(path, exist_ok=True, mode=0o700)
    return path

def stratustryke_dir() -> Path:
    '''
    Get directory stratustryke resides in
    :rtype: <Path> Path to directory stratustryke is in
    '''
    # e.g., /home/user/Tools/Git/Strautstryke/strautstryke
    return Path(__file__).parents[1]

def sqlite_filepath() -> Path:
    path = home_dir()/'stratustryke.sqlite'
    return path

def session_dir() -> Path:
    '''
    Returns directory for the current stratustryke session
    :rtype: <Path> path to the current stratustryke session's directory
    '''
    if not active_session:
        raise UserWarning('No session_name designated.')

    # e.g., /home/user/.local/share/stratustryke/mysession
    p = (home_dir()/cast(Callable, active_session)().name).absolute()
    # Create session dir if it doesn't exist
    os.makedirs(p, exist_ok=True)
    return p

def downloads_dir() -> Path:
    '''
    Returns directory downloads should be sent to for the current session
    :rtype: <Path> Path to current session's download directory
    '''
    # e.g., /home/user/.local/share/stratustryke/mysession/downloads
    p = (session_dir()/'downloads').absolute()
    # Create download dir if it doesn't exist
    os.makedirs(p, exist_ok=True)
    return p

def custom_modules_dir() -> Path:
    '''
    Returns path to directory where users can place custom modules
    :rtype: <Path> Path to custom user modules directory
    '''
    p = (stratustryke_dir()/'user_modules').absolute()
    os.makedirs(p, exist_ok=True)
    return p

def module_data_dir(mod: str) -> Path:
    '''
    Returns directory for specific modules data
    :param: mod: <str> module name
    :rtype: <Path> Path to module-specific directory for a session
    '''
    # e.g., /home/user/.stratustryke/mysession/modules/aws_whoami
    p = (home_dir()/'modules'/mod).absolute()
    os.makedirs(p, exist_ok=True)
    return p

@contextlib.contextmanager
def save(file_name: str, mode: str = 'w', header: Optional[str] = None, **kwargs) -> Generator[IO[Any], None, None]:
    '''
    Save text contents to a session's download directory
    :param: file_name: <str> name of file to write
    :param: mode: <str> file I/O mode [default: \'w\']
    :param: header: <str> string to write before text content (useful for CSVs)
    :rtype: <Generator> File I/O generator
    '''
    p = Path(downloads_dir()) / file_name
    p.parent.mkdir(parents=True, exist_ok=True, mode=0o700)

    with open(str(p), mode, **kwargs) as f:
        if header and not p.exists():
            f.write(header + '\n')
        try:
            yield f
        finally:
            f.close()


from setuptools import setup, find_packages
import codecs
import os

here = os.path.abspath(os.path.dirname(__file__))

with codecs.open(os.path.join(here, "SIMPLEREADME.md"), encoding="utf-8") as fh:
    long_description = "\n" + fh.read()

VERSION = '0.0.8'
DESCRIPTION = 'A library that helps you build dlt pipelines by side stepping the dlt library and making your code interactive'

# Setting up
setup(
    name="dlt_sidestep",
    version=VERSION,
    author="Michael Gardner",
    author_email="gardnmi@gmail.com",
    description=DESCRIPTION,
    long_description_content_type="text/markdown",
    long_description=long_description,
    packages=find_packages(),
    install_requires=[''],
    keywords=['python', 'databricks', 'dlt', 'delta live tables'],
    classifiers=[
        "Development Status :: 5 - Production/Stable",
        "Intended Audience :: Developers",
        "Programming Language :: Python :: 3",
        "Operating System :: Unix",
        "Operating System :: MacOS :: MacOS X",
        "Operating System :: Microsoft :: Windows",
    ]
)


from __future__ import print_function, unicode_literals
# Author & maintainer: BurritoBazooka <burritobazooka@gmail.com>
# License:
######################################
# The MIT License (MIT)

# Copyright (c) 2016 BurritoBazooka

# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject
# to the following conditions:

# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE
# FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
######################################
import hexchat
import threading
import time
import itertools
import re
import argparse
import traceback
from collections import defaultdict
from functools import wraps

import sys, os
sys.path.append(os.path.join(os.path.dirname(__file__), "submodules"))

# helper modules that should've come with this one. Place them in "submodules" directory, in the same directory as this script.
import inputbox as ibx
import pastebins

# dependency: xerox, for handling clipboards.
try:
    import xerox
    HAVE_XEROX = True
except ImportError:
    print("\002[Floodcontrol]\017\t", "xerox not found, we won't be able to use clipboards. https://pypi.python.org/pypi/xerox https://github.com/kennethreitz/xerox")
    HAVE_XEROX = False

__module_name__ = str("Floodcontrol")
__module_version__ = str("0.1")
__module_description__ = str("Sends your floody messages to a pastebin service and gives you the URL.")

"""HexChat plugin for intercepting floody messages from the user, and redirecting them to a pastebin service."""

def get_max_lines():
    # TODO: Should be more for PMs, and less for large channels.
    # TODO: These defaults could be overridden by users, per channel/window.
    return 2

KEYS = {"enter":        ['65293', '0', '\r', '1'],
        "ctrl+enter":   ['65293', '4', '\r', '1'],
        "altgr+enter":  ['65293', '128', '\r', '1'],
        "alt+enter":    ['65293', '8', '\r', '1'],
        "space":        ['32', '0', ' ', '1'],
        "tab":          ['65289', '0', '', '0']
        }

# TODO
# Possible additional options & examples:
# pastebin: pastebin.com, dpaste.de, codepad, random
# expiry: 1 day, 1 week, shortest, longest, infinite (we'd pick the closest possible)
# name
#

DEBUG = False

def print_fc(*args, **kwargs):
    return print("\002[Floodcontrol]\017\t", *args, **kwargs)

def print_debug(*args, **kwargs):
    if DEBUG:
        return print("\002[FC_debug]\017\t", *args, **kwargs)

def is_mainthread():
    try:
        return threading.main_thread() is threading.current_thread()
    except AttributeError: # < py3.4
        return isinstance(threading.current_thread(), threading._MainThread)

def default_returnvalue(returnvalue, error_returnvalue):
    # Decorator for always returning a default value if None is returned by the
    # decorated function, or if the decorated function raises an error.
    #
    # This is to prevent HexChat behaviour where if a command raises an exception,
    # the event isn't eaten. We use this to EAT_HEXCHAT by default.
    def dec(f):
        @wraps(f)
        def wrapped(*args, **kwargs):
            try:
                r = f(*args, **kwargs)
            except Exception as e:
                for line in traceback.format_exc().splitlines():
                    print_fc(line)
                return error_returnvalue

            if r is None:
                return returnvalue
            else:
                return r
        return wrapped
    return dec

class FloodcontrolError(BaseException):
    pass

class PluginConfigError(FloodcontrolError):
    pass

class ArgparseException(FloodcontrolError):
    pass

class ArgparseError(ArgparseException):
    pass

class ArgparseHelpMessage(ArgparseException):
    pass

class ArgparseExit(ArgparseError):
    pass

class NoExitParser(argparse.ArgumentParser):
    # This is so that argparse doesn't make HexChat exit during help and error messages.
    raise_on_next_fail = False # TODO: override __init__ and put this in there
    def exit(self, *args, **kwargs):
        if self.raise_on_next_fail:
            self.raise_on_next_fail = False
            raise ArgparseExit(args, kwargs)
        else:
            print_debug("(ArgparseExit)", args, kwargs)
    def error(self, *args, **kwargs):
        if self.raise_on_next_fail:
            self.raise_on_next_fail = False
            raise ArgparseError(args, kwargs)
        else:
            print_fc("(ArgparseError)", args, kwargs)
    def print_help(self, *args, **kwargs):
        if self.raise_on_next_fail:
            self.raise_on_next_fail = False
            raise ArgparseHelpMessage
        else:
            self.print_help_for_real(*args, **kwargs)
            print_debug("(ArgparseHelpMessage)")
    def print_help_for_real(self, *args, **kwargs):
        super(NoExitParser, self).print_help(*args, **kwargs)

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

######## Configuration

default_config = {
    "service": "pb",
    "expiry": "1 week",
    "exposure": "unlisted",
    "syntax": "text",
    "name": "Pasted by Hexchat Floodcontrol",
    "autopaste": "on"
}

def make_argparser_and_args_after_config():
    global argparser, ARG_GROUPS, ALL_ARGS
    argparser, ARG_GROUPS, ALL_ARGS = make_argparser_and_args()

def set_option(key, value):
    hkey = "{}_{}".format(__module_name__, key)

    value = str(value)

    success = hexchat.set_pluginpref(hkey, value)
    if not success:
        raise PluginConfigError((key, value))
    make_argparser_and_args_after_config()

def get_option(key):
    hkey = "{}_{}".format(__module_name__, key)
    configsetting = hexchat.get_pluginpref(hkey)
    if configsetting is not None:
        return configsetting
    else:
        return default_config.get(key)

def del_option(key):
    hkey = "{}_{}".format(__module_name__, key)

    success = hexchat.del_pluginpref(hkey)
    if not success:
        raise PluginConfigError(key)
    make_argparser_and_args_after_config()

# def set_max_lines_for_channel(value=None, channel=None, network=None):
#     # TODO: This and the next function are supposed to set the preferred maximum
#     # lines before we should ask the user whether to pastebin.
#     if channel is None:
#         channel = hexchat.get_info("channel")
#     if network is None:
#         network = hexchat.get_info("network")
#     key = "maxlines_{}_{}".format(channel, network)
#
#     if value is not None:
#         set_option(key, value)
#     else:
#         del_option(key)
#
# def get_max_lines_for_channel(channel=None, network=None):
#     key = "maxlines_{}_{}".format(channel, network)
#     # TODO

###

# TODO: Have a single "set_option_cmd" function, with validator functions for each valid option.

@default_returnvalue(hexchat.EAT_HEXCHAT, hexchat.EAT_ALL)
def set_service_cmd(words, words_eol, *args):
    apis = pastebins.get_api_names()
    default_api = default_config.get('service')
    available_str = "Available pastebins (or 'default' for {}): {}".format(default_api, apis)
    if len(words) == 1:
        print_fc(available_str)
        print_fc("Current pastebin: {}".format(get_option("service")))
    else:
        if words[1] in apis:
            set_option("service", words[1])
        elif words[1] == "default":
            del_option("service")
        else:
            print_fc("We don't have that pastebin. {}".format(available_str))
            return
        print_fc("Pastebin is now: {}".format(get_option("service")))

# TODO: test this.
@default_returnvalue(hexchat.EAT_HEXCHAT, hexchat.EAT_ALL)
def set_option_cmd(words, words_eol, userdata, *args):
    option_name = userdata
    try:
        option_value = words[1]
    except IndexError:
        option_value = None

    option_setters = {"service": set_service_cmd}
    # We don't validate the following because different APIs might handle them [in]correctly.
    # TODO: add validators to the PastebinAPI objects?
    option_rest = {"expiry", "exposure", "syntax", "name"}
    if option_name in option_rest.union(option_funcs):
        if callable(option_funcs.get(option_name)):
            return option_funcs[option_name](words, words_eol)
        else:
            if option_value is None:
                print_fc("Current {} setting: {}".format(option_name, get_option(option_name)))
            else:
                if option_value == "default":
                    del_option(option_name)
                else:
                    set_option(option_name, option_value)
                print_fc("{} setting is now: {}".format(option_name, get_option(option_name)))

######## Pastebin service functions

# Passed to the `pastebins` submodule for use in HTTP requests:
USER_AGENT = "hexchat-{}/{}".format(__module_name__, __module_version__)

def _pastebin(contents, **kwargs):
    pastebin = pastebins.get_api_by_name(kwargs['service'])
    print_fc("Sending paste to {}...".format(kwargs['service']))
    start_time = time.time()
    result = pastebin.write(contents, **kwargs)
    if result[1] is not None:
        print_fc("Response from {}:".format(kwargs['service']))
        print_fc(result[1])
        print_fc("Time taken: {}".format(time.time() - start_time))
    return result[0]

def pastebin(callback, contents, **kwargs):
    for argname in ARG_GROUPS['api']:
        if argname not in kwargs:
            kwargs[argname] = get_option(argname)

    callback(_pastebin(contents, **kwargs))

def pastebin_thread(callback, *args, **kwargs):
    print_debug("pastebin_thread", callback, args, kwargs)
    thread = threading.Thread(target=pastebin, args=(callback,)+args, kwargs=kwargs)
    thread.start()
    return thread

@default_returnvalue(hexchat.EAT_HEXCHAT, hexchat.EAT_ALL)
def add_shellcommand_pastebin_cmd(words, words_eol, *args):
    # TODO: Write to config in order to make the setting permanent.
    words = words + ([None] * (4 - len(words)))
    apiname, shellcommand, shellcommand_read = words[1:4]

    pastebins.add_shellcommand_pastebin(apiname, shellcommand, shellcommand_read)

    print_fc("Shell command Pastebin API added. Name: {}, Command to write: {}, Command to read: {}".format(apiname, shellcommand, shellcommand_read))

@default_returnvalue(hexchat.EAT_HEXCHAT, hexchat.EAT_ALL)
def remove_shellcommand_pastebin_cmd(words, words_eol, *args):
    apiname = words[1]
    api = pastebins.get_api_by_name(apiname)
    if isinstance(api, pastebins.ShellCommandPastebin):
        api.remove_api()
    else:
        raise FloodControlError("{} is not a ShellCommandPastebin, cannot remove.".format(apiname))
    print_fc("Shell command Pastebin API added. Name: {}".format(apiname))

######## Line counting and floodiness estimation functions

# The following two functions are based on split_up_text in hexchat/src/common/outbound.c.
def split_up_text(text, maxlen, cmd_length):
    if len(text) > maxlen:
        # Try splitting at last space:
        spaceindex = text.rfind(" ", 0, maxlen)
        if spaceindex >= 0:
            # Only split if last word is of sane length:
            if maxlen != spaceindex and maxlen - spaceindex < 20:
                maxlen = spaceindex + 1

        split_text = text[:maxlen]

        return split_text

# This and the previous function is based on split_up_text in hexchat/src/common/outbound.c.
def splits_up_text(text, cmd_length):
    maxlen = 512 # rfc 2812
    maxlen -= 3 # :, !, @
    maxlen -= cmd_length
    maxlen -= len(hexchat.get_info("nick"))
    maxlen -= len(hexchat.get_info("channel"))

    # Hexchat's Python interface doesn't provide for a nice way to get our own
    # username and hostname. Writing a hack to do it isn't worth it for this case.
    maxlen -= 9 + 65 # max length of username@hostname

    l = 0
    splits = []
    while True:
        split_text = split_up_text(text[l:], maxlen, cmd_length)
        if split_text is None:
            splits.append(text[l:])
            break
        l += len(split_text)
        splits.append(split_text)
    return splits

def linecount(message=None, cmd_length=len(" PRIVMSG : \r\n")):
    if message is None:
        message = hexchat.get_info("inputbox")

    split_msg = message.splitlines()

    approximate = False
    linecount = 0
    for line in split_msg:
        lines_in_line = len(splits_up_text(line, cmd_length))
        if lines_in_line > 1:
            approximate = True
        linecount += lines_in_line

    return linecount, approximate

commands_options = { # Values may be dicts or callables which return dicts when called with inputbox text.
    "msg": {
        "cmdlength": len(" PRIVMSG : \r\n"),
        "paramcount": 1
        },
    "notice": {
        "cmdlength": len(" NOTICE : \r\n"),
        "paramcount": 1
        },
    "me": {
        "cmdlength": len(" PRIVMSG : \001ACTION \001\r\n"),
        "paramcount": 0
        }
}

commands_options['say'] = {"cmdlength": commands_options['msg']['cmdlength'], 'paramcount': 0}
commands_options['m'] = commands_options['msg']

def add_special_preprocessors():
    # For when a command's paramcount depends on what the command's parameters
    # are, or something dynamic like that.
    # Special preprocessors should never return another callable. Instead, they
    # could use get_opts_for_cmd which would do the recursion.
    def get_opts_for_wrappercmd(inputbox):
        # for wrapper commands like /allchan, /allchanl, and /allserv
        msg_split = inputbox.split(" ")
        command = msg_split[1]
        options = get_opts_for_cmd(command, inputbox).copy()
        if 'paramcount' in options:
            options['paramcount'] += 1
        return options
    for command in ['allchan', 'allchanl', 'allserv', 'doat']:
        commands_options[command] = get_opts_for_wrappercmd

add_special_preprocessors()

def get_opts_for_cmd(command, inputbox):
    print_debug("get_opts_for_cmd", command, inputbox, sep="|")
    options = commands_options.get(command.lower())
    if callable(options):
        options = options(inputbox)
    return options

######## Argument parsing and input/output for /fc_paste command and do_paste function

def find_content_in_args(words, words_eol, parser, key_to_watch="content"):
    # Works around argparse's magic to find the raw unparsed content to pastebin.
    # The alternative was to create an ArgumentParser sublcass and then override
    # or wrap a _semiprivate method. This would have been bad because the
    # _semiprivately named methods are unspecified and might change, per convention.

    # The problem this solves can be demonstrated in the example:
    # /argparse --validflag -v2 content to paste goes -- here
    # Would normally output: "content to paste goes here", where we would want
    # "content to paste goes -- here", also parsing out any quoted statements
    # which we needed to have raw.

    # What we actually do is parse an ever-increasing part of the args until we
    # find an arg we don't recognise. At this point, it is assumed the rest
    # of the arguments are content to paste (called 'extra' in this function).

    command_params = words[1:]
    parsed_before = []
    helpmessage_position = None
    for i in range(len(command_params)):
        to_parse = command_params[:i+1]
        print_debug("command_params to_parse", to_parse)
        try:
            parser.raise_on_next_fail = True
            parsed, extra = parser.parse_known_args(to_parse)
            parser.raise_on_next_fail = False
        except ArgparseHelpMessage as e:
            helpmessage_position = i
            continue
        except ArgparseException as e:
            print_debug("ArgparseException:", repr(e))
            continue
        parsed = vars(parsed)
        print_debug("e_parsed", parsed)
        print_debug("extra", extra)

        if len(parsed[key_to_watch]) > 0 or len(extra) > 0:
            if len(parsed_before) == 0:
                if i > 0:
                    return None, words_eol[i+1]
                else:
                    return parsed, words_eol[i+1]
            else:
                return parsed_before[-1], words_eol[i+1]

        parsed_before.append(parsed)

    print_debug("parsed_before", parsed_before)

    # No extras found at this point.

    if helpmessage_position is not None and len(parsed_before) == 0:
        parser.print_help_for_real()
        return None, None

    return parsed_before[-1], None

def output_from_argparse(output, parsed_args): # TODO: argparse functions need better name.
    parsed_args = parsed_args.copy()

    commands = parsed_args.get('to_command')
    if commands is None:
        commands = []

    if parsed_args.get('say'):
        commands.append("say")

    for c in commands:
        hexchat.command(" ".join((c, output)))
    if parsed_args.get('guard_inputbox_cmd') and (parsed_args.get('to_inputbox') or parsed_args.get('to_inputbox_replace')):
        ibx.set(" ".join((parsed_args['guard_inputbox_cmd'][0], output)))
    elif parsed_args.get("to_inputbox"):
        ibx.add_at_cursor(output)
    elif parsed_args.get("to_inputbox_replace"):
        ibx.set(output)
    if parsed_args.get("to_clipboard") and HAVE_XEROX:
        xerox.copy(output)

def get_input_from_argparse(callback, parsed_args):
    print_debug(parsed_args)
    parsed_args = parsed_args.copy()
    source = parsed_args.get('source')
    print_debug("get_input args", parsed_args)

    if parsed_args.get("confirm"):
        del parsed_args['confirm']
        def confirmed_cb(b):
            if b:
                get_input_from_argparse(callback, parsed_args)
        if is_mainthread():
            send_getbool_to_callback(confirmed_cb)
        else:
            hexchat.hook_timer(20, send_getbool_to_callback, confirmed_cb)
        return

    if source == "inputbox":
        if not parsed_args['guard_inputbox_cmd']:
            callback(ibx.get())
        else:
            callback(parsed_args['guard_inputbox_cmd'][1])

    elif source == "clipboard" and HAVE_XEROX:
        callback(xerox.paste())

    elif source == "window":
        # Ask main thread to getstr and then give it to our callback.
        if is_mainthread():
            send_getstr_to_callback(callback)
        else:
            hexchat.hook_timer(20, send_getstr_to_callback, callback)

    elif parsed_args['content']:
        callback(parsed_args['content'])

    else:
        raise FloodcontrolError("Could not get input. Requested source: {}".format(source))

hexchat_callback_handlers = defaultdict(list)

def send_getstr_to_callback(callback):
    # TODO: More generalised callback-command thing?
    cmdname = ".floodcontrol_cb_getstr"
    @default_returnvalue(hexchat.EAT_HEXCHAT, hexchat.EAT_ALL)
    def callback_and_unhook(words, words_eol, *args):
        o = words_eol[1]
        if o != "(paste here)":
            threading.Thread(target=callback, args=(o,)).start()
        {hexchat.unhook(handler) for handler in hexchat_callback_handlers[cmdname]}
    {hexchat.unhook(handler) for handler in hexchat_callback_handlers[cmdname]}
    handler = hexchat.hook_command(cmdname, callback_and_unhook, userdata=time.time())
    hexchat_callback_handlers[cmdname].append(handler)
    hexchat.command(" ".join(("GETSTR", '"(paste here)"', cmdname, '"Floodcontrol pastebin input:"')))

def send_getbool_to_callback(callback):
    # TODO: More generalised callback-command thing?
    cmdname = ".floodcontrol_cb_getbool"
    @default_returnvalue(hexchat.EAT_HEXCHAT, hexchat.EAT_ALL)
    def callback_and_unhook(words, words_eol, *args):
        o = bool(int(words_eol[1]))
        threading.Thread(target=callback, args=(o,)).start()
        {hexchat.unhook(handler) for handler in hexchat_callback_handlers[cmdname]}
    {hexchat.unhook(handler) for handler in hexchat_callback_handlers[cmdname]}
    handler = hexchat.hook_command(cmdname, callback_and_unhook, userdata=time.time())
    hexchat_callback_handlers[cmdname].append(handler)
    hexchat.command(" ".join(("GETBOOL", cmdname, '"Floodcontrol"', '"Confirm sending to pastebin?"')))

argparser = ARG_GROUPS = ALL_ARGS = None
def make_argparser_and_args():
    argparser = NoExitParser("fc_paste") # TODO: find way for this name to be set dynamically by the command running it.

    help_xerox_missing = "Does nothing apart from telling the user about the missing 'xerox' dependency. Would do something if we had the 'xerox' dependency."

    inputgroup = argparser.add_argument_group(title="Input sources")
    ia = []
    ia.append(inputgroup.add_argument("content", nargs="*", help="Ignored if a --from-* argument is specified. The content we should send to the pastebin. Not reliable when running the command from HexChat's inputbox because of how line breaks behave. Use a --from-* argument instead."))
    ia.append(inputgroup.add_argument("-fi", "--from-inputbox", action="store_const", help="If specified, will retrieve HexChat inputbox contents and send it to the pastebin.", dest="source", const="inputbox"))
    if HAVE_XEROX:
        ia.append(inputgroup.add_argument("-fb", "--from-clipboard", action="store_const", help="If specified, will retrieve clipboard and send it to the pastebin.", dest="source", const="clipboard"))
    else:
        ia.append(inputgroup.add_argument("-fb", "--from-clipboard", action="store_true", help=help_xerox_missing, dest="tried_clipboard"))
    ia.append(inputgroup.add_argument("-fw", "--from-window", action="store_const", help="If specified, will create a popup window using /GETSTR and send the response to the pastebin.", dest="source", const="window"))
    ia.append(inputgroup.add_argument("-C", "--confirm", action="store_true", help="If specified, we will first use /GETBOOL to get user confirmation before sending data to the patebin service."))

    # IDEA: Paste from file path: -ff /home/user/blah.txt

    # Pastebin API options. The interpretation and functionality of these depend completely on the chosen API.
    # IDEA: Custom HelpFormatter for bold and italics in IRC.

    bin_api_group = argparser.add_argument_group(title="Pastebin API arguments")
    ba = []
    ba.append(bin_api_group.add_argument("-p", "--service", default=get_option("service"), help="The name of the pastebin API desired. (default: %(default)s)"))
    ba.append(bin_api_group.add_argument("-n", "--name", default=get_option("name"), help="The name of your paste. (default: %(default)s)"))
    ba.append(bin_api_group.add_argument("-e", "--expiry", default=get_option("expiry"), help="Length of time we should tell the pastebin API to keep the paste. (default: %(default)s)"))
    ba.append(bin_api_group.add_argument("-s", "--syntax", default=get_option("syntax"), help="Try to get the pastebin to use syntax highlighting, e.g. 'python' or 'html'. (default: %(default)s)"))
    ba.append(bin_api_group.add_argument("-x", "--exposure", default=get_option("exposure"), help="Privacy setting for this paste. (default: %(default)s)"))

    # TODO: Make this work:
    #ba.append(bin_api_group.add_argument("-c", "--custom-opt", action="append", help="Custom options to give to the Pastebin API (as strings). For instance, some APIs might be coded to support a maximum read count. Use the syntax \"key=value\" including quotes."))

    outputgroup = argparser.add_argument_group(title="Output destinations")
    oa = []
    oa.append(outputgroup.add_argument("-tc", "--to-command", action="append", metavar="COMMAND", help="The command that should be ran with the result, e.g. '--to-command \"/msg OtherPerson\" content goes here' will send the Pastebin URL as '/msg OtherPerson http://example.com/abcd'."))
    oa.append(outputgroup.add_argument("-ti", "--to-inputbox", action="store_true", help="If specified, will add the pastebin URL to HexChat's inputbox at the current cursor position."))
    oa.append(outputgroup.add_argument("-tir", "--to-inputbox-replace", action="store_true", help="If specified, will set HexChat's inputbox to the pastebin URL."))
    if HAVE_XEROX:
        oa.append(outputgroup.add_argument("-tb", "--to-clipboard", action="store_true", help="If specified, will add the pastebin URL to the clipboard."))
    else:
        oa.append(outputgroup.add_argument("-tb", "--to-clipboard", action="store_true", help=help_xerox_missing, dest="tried_clipboard"))
    oa.append(outputgroup.add_argument("-S", "--say", action="store_true", help="If specified, will output pastebin URL to current window. Same as '--to-command say'"))

    ma = []
    # The following arg gets mutated into a tuple, when used through do_paste_cmd: (command, rest_of_message)
    ma.append(argparser.add_argument("--guard-inputbox-cmd", action="store_true", help="If specified, will take care not to send known commands to the pastebin, and preserve it when we get the URL back. For instance, \"/msg Burrito testing\" would only send \"testing\" and preserve \"/msg Burrito\" in the inputbox."))

    ARG_GROUPS = {}
    ARG_GROUPS['input'] = {arg.dest for arg in ia}
    ARG_GROUPS['api'] = {arg.dest for arg in ba}
    ARG_GROUPS['output'] = {arg.dest for arg in oa}
    ARG_GROUPS['misc'] = {arg.dest for arg in ma}
    ALL_ARGS = set(itertools.chain(*ARG_GROUPS.values()))

    return argparser, ARG_GROUPS, ALL_ARGS
argparser, ARG_GROUPS, ALL_ARGS = make_argparser_and_args()

@default_returnvalue(hexchat.EAT_HEXCHAT, hexchat.EAT_ALL)
def do_paste_cmd(words, words_eol, userdata, *args, **kwargs):
    # IDEA/TODO: Paster object to store options and argparser?
    # Example:
    # /fc_paste -fb -n "Paste name" -p pb -e "1 week" -s text -x unlisted -S
    # /fc_paste -tc "command" Paste content goes here
    # /fc_paste -fb -ti -C
    # /fc_paste -h

    parsed_args = kwargs.get("parsed_args")

    print_debug("words", words)
    for w in words_eol:
        print_debug("words_eol:", w)

    if parsed_args is None:
        parsed_args, content = find_content_in_args(words, words_eol, argparser)
    if parsed_args is None:
        print_debug("parsed_args is None even after find_content_in_args")
        return
    parsed_args['content'] = content

    if parsed_args['guard_inputbox_cmd'] and parsed_args['source'] == "inputbox":
        # TODO: could split into separate function and also use in do_paste
        doprocess, message, fullcommand, _ = preprocess_inputbox(ibx.get())
        if doprocess:
            parsed_args['guard_inputbox_cmd'] = fullcommand, message
        else:
            parsed_args['guard_inputbox_cmd'] = False
    print_debug("parsed", parsed_args)

    do_paste(**parsed_args)

def add_config_to_options(options, keys=ALL_ARGS, filter_keys=None):
    # TODO: take a look at dict.get and dict.setdefault, we might be able to use it here
    if filter_keys is not None:
        options = {k: v for k, v in options.items() if k in filter_keys}
    else:
        options = options.copy()

    for arg in keys:
        if arg not in options:
            options[arg] = get_option(arg)
    return options

def do_paste(**options): # unknown options are discarded, see PASTECMD_ARGS
    print_debug("ALL_ARGS", ALL_ARGS)
    options = {k: v for k, v in options.items() if k in ALL_ARGS}
    api_options = add_config_to_options(options, ARG_GROUPS['api'], ARG_GROUPS['api'])
    if options.get('tried_clipboard'):
        print_fc("We cannot use the clipboard, because we do not have the 'xerox' dependency. https://pypi.python.org/pypi/xerox https://github.com/kennethreitz/xerox")
    def paste_content_cb(content):
        # Receives paste content, sends it to the pastebin. Pastebin API will reply through url_cb.
        pastebin_thread(url_cb, content, **api_options)
    def url_cb(url):
        # Takes in pastebin URL, outputs URL to necessary places.
        output_from_argparse(url, options)
    thread = threading.Thread(target=get_input_from_argparse, args=(paste_content_cb, options))
    thread.start()
    return thread

######## Inputbox interaction and autopaste functions

@default_returnvalue(hexchat.EAT_HEXCHAT, hexchat.EAT_ALL)
def inputbox_autopaste_on_limit(*args):
    if len(args) < 4:
        doprocess, message, fullcommand, cmd_length = preprocess_inputbox(ibx.get())
    else:
        doprocess, message, fullcommand, cmd_length = args
    options = {"guard_inputbox_cmd": (fullcommand, message),
               "source": "inputbox",
               "to_inputbox": True}
    do_paste(**options)

re_words = re.compile(r"(\S*\s)")

def preprocess_inputbox(inputbox):
    # Return values:
    # doprocess, message, fullcommand, cmd_length
    # TODO: only return at end.
    # "doprocess" means we can understand what needs to be parsed in the inputbox so that we can have content to pastebin.
    # "message" is the content to pastebin, without the command.
    # "fullcommand" is the command we have parsed from the inputbox.
    #   e.g. "/msg OtherUser"
    # "cmd_legth" is how much space this command would take up in the raw message to the server.
    cmdprefix = hexchat.get_prefs("input_command_char") # default is "/" as in "/msg"
    if inputbox.startswith(cmdprefix + cmdprefix) or not inputbox.startswith(cmdprefix):
        # ('//msg blah' is the same as '/say /msg blah')
        cmd_length = commands_options['say']['cmdlength']
        return True, inputbox, "", cmd_length
    else:
        boxsplit = [word for word in re.split(re_words, inputbox) if word]
        print_debug("boxsplit", boxsplit)
        command = boxsplit[0]
        command = command.lstrip("/").rstrip()
        options = get_opts_for_cmd(command, inputbox)
        if options is not None:
            paramcount = options['paramcount']
            print_debug("paramcount", paramcount)
            re_cmd_and_msg = r'^((?:[^\s]*\s*){%d}[^\s]*)\s(.*)' % paramcount # by perreal https://stackoverflow.com/a/17060122/3143160
            # fullcommand here would be "/msg OtherUser" in the example of "/msg OtherUser message"
            split_cmd_and_msg = [word for word in re.split(re_cmd_and_msg, inputbox) if word]
            fullcommand = split_cmd_and_msg[0]
            message = inputbox[len(fullcommand)+1:]
            fullcommand = fullcommand.strip()

            cmd_length = options['cmdlength']
            return True, message, fullcommand, cmd_length
    return False, None, None, None

######## Callbacks and hooks

# IDEA: User types: "blah blah blah %Po%longtextgoeshere%Pc% blah blah"
#  And it gets turned into: "blah blah blah http://e.g.bin/jdkfai blah blah"
# Alternatively, user could enter a keyboard shortcut and we insert a placeholder for them.
# Like a Unicode non-character character.

_mode = 0
def mode(*args):
    global _mode
    if len(args) > 0:
        print_debug("mode", _mode, "to", args[0])
        _mode = args[0]
    return _mode

def keypress_cb(key, *args):
    # TODO: Split into functions. State information should be held as object attributes?

    inputbox = hexchat.get_info("inputbox")

    if key == KEYS["enter"] and len(inputbox) > 0:
        doprocess, message, fullcommand, cmd_length = preprocess_inputbox(inputbox)
        if message is not None:
            print_debug(doprocess, len(message), message[:50], fullcommand, sep="|")
        else:
            print_debug(doprocess, message, fullcommand, sep="|")
        if doprocess:
            # boolean "doprocess" here means that we can understand what we need to
            # extract from the inputbox, and how to send the replacement.
            nlines, approx = linecount(message, cmd_length)
            nprefix = "~" * approx

            if mode() == 0 and nlines > get_max_lines():
                print_fc("That message looks like a flood ({}{} lines). To really send it, "
                "press \002Alt+Enter\017. To send it through a pastebin, "
                "press \002Enter\017 again.".format(nprefix, nlines))
                mode(1)
                return hexchat.EAT_ALL
            elif mode() == 1 and nlines > get_max_lines():
                mode(0)
                return inputbox_autopaste_on_limit(doprocess, message, fullcommand, cmd_length)
            elif not nlines > get_max_lines():
                mode(0)
                return hexchat.EAT_NONE

    elif key in (KEYS["alt+enter"], KEYS["altgr+enter"]) or key[-1] != '0': # key[-1] is the key's character length
        mode(0)

def debug_keypress_cb(key, *args):
    print_debug(key)

debug_handler = None
@default_returnvalue(hexchat.EAT_HEXCHAT, hexchat.EAT_ALL)
def toggle_debug(*args):
    global debug_handler
    global DEBUG

    if not debug_handler:
        print_fc("enabling debug")
        DEBUG = True
        debug_handler = hexchat.hook_print("Key Press", debug_keypress_cb)
    else:
        print_fc("disabling debug")
        DEBUG = False
        hexchat.unhook(debug_handler)
        debug_handler = None

autopaste_handler = None
@default_returnvalue(hexchat.EAT_HEXCHAT, hexchat.EAT_ALL)
def toggle_autopaste(*args):
    global autopaste_handler
    words = args[0]
    bools = {True: ("1", "on", "true", "t"), False: ("0", "off", "false", "f")}
    try:
        requested = words[1].lower()
    except IndexError:
        requested = ""

    print_debug("toggle_autopaste requested", requested)
    if not autopaste_handler and requested not in bools[False]:
        print_fc("enabling autopaste")
        set_option("autopaste", "on")
        autopaste_handler = hexchat.hook_print("Key Press", keypress_cb)
    elif autopaste_handler and requested not in bools[True]:
        print_fc("disabling autopaste")
        set_option("autopaste", "off")
        hexchat.unhook(autopaste_handler)
        autopaste_handler = None

if __name__ == "__main__":

    hexchat.hook_command("fc_debug", toggle_debug, help="Toggle debug messages on/off for Floodcontrol.")

    hexchat.hook_command("fc_autopaste", toggle_autopaste, help="Usage: fc_autopaste [on|off]\nToggle autopaste on/off. Autopaste will notify you to confirm sending your floody message to the pastebin service automatically when you try to send it.")
    toggle_autopaste(["", get_option("autopaste")])
    hexchat.hook_command("fc_setpastebin", set_service_cmd, help="Set which pastebin to use for Floodcontrol. Provide no parameters to see a list of available pastebins. Provide 'default' to set back to default.")

    hexchat.hook_command("fc_paste", do_paste_cmd, help="Try /fc_paste -h")

    hexchat.hook_command("fc_rm_shellpastebin", remove_shellcommand_pastebin_cmd)
    hexchat.hook_command("fc_add_shellpastebin", add_shellcommand_pastebin_cmd) # TODO: Help message for this and previous one.


#!/usr/bin/python
from datetime import datetime
import time
import argparse
import logging
import sys
from os import getenv
import requests
from petadmin import Environment

log = logging.getLogger(__name__)
env = Environment('check_status')

env.configure_logger(log)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('target', action='store', help='Target of check: local or remote')
    parser.add_argument('-fail', help='Simulate a fail', action='store_true')

    args = parser.parse_args()

    log.info('Running %s', ' '.join(sys.argv))

    target = args.target

    if target == 'local':
        url = 'http://192.168.0.200/wordpress/wp-content/plugins/crowbank/status.php'
    else:
        url = 'http://dev.crowbankkennels.co.uk/wp-content/plugins/crowbank/status.php'

    sql = "select ws_state, ws_lasttransfer from vwwebsite_state where ws_iscurrent = 1"
    cursor = env.get_cursor()

    cursor.execute(sql)

    last_fail = 0

    for row in cursor:
        previous_status = row[0]
        previous_lasttransfer = row[1]

    now = datetime.now()
    r = requests.get(url)
    status = r.json()

    age = status['age']
    lasttransfer_text = status['lasttransfer']['date']
    system_status = status['status']

    in_date_format = '%Y-%m-%d %H:%M:%S.%f'
    date_format = '%Y-%m-%d %H:%M:%S'
    lasttransfer = datetime.strptime(lasttransfer_text, in_date_format)
    lasttransfer_text = lasttransfer.strftime(date_format)

    new_status = 'Loaded'

    if age > 1200:
        new_status = 'Obsolete'
        log.error('Data age on %s system is %d seconds, last transfer on %s' % (target, age, lasttransfer))

    if system_status != 'Loaded':
        new_status = system_status
        log.error('Status on %s system is %s' % (target, system_status))

    if new_status == 'Loaded':
        if args.fail:
            new_status = 'Fail'
            log.error('Simulated fail on %s system' % target)
        else:
            log.info('Status on %s system is good. Last transfer, %d seconds ago, at %s' % (target, age, lasttransfer))

    if target == 'remote':
        send_email = (new_status != previous_status) or (lasttransfer != previous_lasttransfer)

        sql = "execute pset_website_state '%s', '%s', %d, '%s'" % (new_status, r, send_email, lasttransfer_text)

        env.execute(sql)

        if send_email:
            body = 'Remote system status changed from %s to %s. Last successful transfer was %s, %d seconds ago' % (previous_status, new_status, lasttransfer_text, age)
            subject = 'Remote system updated to %s' % new_status

            send_to = env.email_logs
            env.send_email(send_to, body, subject, body)

    env.close()

if __name__ == '__main__':
    main()


import pytest
import requests

@pytest.mark.duckduckgo
@pytest.mark.restapi
def test_duckduckgo():

    #Arrange
   # url = 'https://api.duckduckgo.com/?q=python+programming&format=json'
    url = 'https://pocketfactory-qa.ssidecisions.com/'

    #Act
    response = requests.get(url)
    #body = response.json()

    #Assert
    assert response.status_code == 200
    #assert 'Python' in body['AbstractText']