import sys
import getopt

from datetime import datetime
from typing import Optional

import wandb

import datasets
import torch
from torch.utils.data import DataLoader
from transformers import (
    AutoConfig,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    get_linear_schedule_with_warmup,
)
from pytorch_lightning import LightningDataModule, LightningModule, Trainer, seed_everything
from pytorch_lightning.loggers import WandbLogger

class GLUEDataModule(LightningDataModule):
    task_text_field_map = {
        "cola": ["sentence"],
        "sst2": ["sentence"],
        "mrpc": ["sentence1", "sentence2"],
        "qqp": ["question1", "question2"],
        "stsb": ["sentence1", "sentence2"],
        "mnli": ["premise", "hypothesis"],
        "qnli": ["question", "sentence"],
        "rte": ["sentence1", "sentence2"],
        "wnli": ["sentence1", "sentence2"],
        "ax": ["premise", "hypothesis"],
    }

    glue_task_num_labels = {
        "cola": 2,
        "sst2": 2,
        "mrpc": 2,
        "qqp": 2,
        "stsb": 1,
        "mnli": 3,
        "qnli": 2,
        "rte": 2,
        "wnli": 2,
        "ax": 3,
    }

    loader_columns = [
        "datasets_idx",
        "input_ids",
        "token_type_ids",
        "attention_mask",
        "start_positions",
        "end_positions",
        "labels",
    ]

    def __init__(
        self,
        model_name_or_path: str,
        task_name: str = "mrpc",
        max_seq_length: int = 128,
        train_batch_size: int = 32,
        eval_batch_size: int = 32,
        **kwargs,
    ):
        super().__init__()
        self.model_name_or_path = model_name_or_path
        self.task_name = task_name
        self.max_seq_length = max_seq_length
        self.train_batch_size = train_batch_size
        self.eval_batch_size = eval_batch_size

        self.text_fields = self.task_text_field_map[task_name]
        self.num_labels = self.glue_task_num_labels[task_name]
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path, use_fast=True)

    def setup(self, stage: str):
        self.dataset = datasets.load_dataset("glue", self.task_name)

        for split in self.dataset.keys():
            self.dataset[split] = self.dataset[split].map(
                self.convert_to_features,
                batched=True,
                remove_columns=["label"],
            )
            self.columns = [c for c in self.dataset[split].column_names if c in self.loader_columns]
            self.dataset[split].set_format(type="torch", columns=self.columns)

        self.eval_splits = [x for x in self.dataset.keys() if "validation" in x]

    def prepare_data(self):
        datasets.load_dataset("glue", self.task_name)
        AutoTokenizer.from_pretrained(self.model_name_or_path, use_fast=True)

    def train_dataloader(self):
        return DataLoader(self.dataset["train"], batch_size=self.train_batch_size, shuffle=True)

    def val_dataloader(self):
        if len(self.eval_splits) == 1:
            return DataLoader(self.dataset["validation"], batch_size=self.eval_batch_size)
        elif len(self.eval_splits) > 1:
            return [DataLoader(self.dataset[x], batch_size=self.eval_batch_size) for x in self.eval_splits]

    def test_dataloader(self):
        if len(self.eval_splits) == 1:
            return DataLoader(self.dataset["test"], batch_size=self.eval_batch_size)
        elif len(self.eval_splits) > 1:
            return [DataLoader(self.dataset[x], batch_size=self.eval_batch_size) for x in self.eval_splits]

    def convert_to_features(self, example_batch, indices=None):
        # Either encode single sentence or sentence pairs
        if len(self.text_fields) > 1:
            texts_or_text_pairs = list(zip(example_batch[self.text_fields[0]], example_batch[self.text_fields[1]]))
        else:
            texts_or_text_pairs = example_batch[self.text_fields[0]]

        # Tokenize the text/text pairs
        features = self.tokenizer.batch_encode_plus(
            texts_or_text_pairs, max_length=self.max_seq_length, pad_to_max_length=True, truncation=True
        )

        # Rename label to labels to make it easier to pass to model forward
        features["labels"] = example_batch["label"]

        return features

class GLUETransformer(LightningModule):
    def __init__(
        self,
        model_name_or_path: str,
        num_labels: int,
        task_name: str,
        learning_rate: float = 2e-5,
        adam_epsilon: float = 1e-8,
        warmup_steps: int = 0,
        weight_decay: float = 0.0,
        train_batch_size: int = 32,
        eval_batch_size: int = 32,
        eval_splits: Optional[list] = None,
        **kwargs,
    ):
        super().__init__()

        self.save_hyperparameters()

        self.config = AutoConfig.from_pretrained(model_name_or_path, num_labels=num_labels)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name_or_path, config=self.config)
        self.metric = datasets.load_metric(
            "glue", self.hparams.task_name, experiment_id=datetime.now().strftime("%d-%m-%Y_%H-%M-%S")
        )
        self.save_hyperparameters()

    def forward(self, **inputs):
        return self.model(**inputs)

    def training_step(self, batch, batch_idx):
        outputs = self(**batch)
        loss = outputs[0]
        self.log("train_loss", loss)
        return loss

    def validation_step(self, batch, batch_idx, dataloader_idx=0):
        outputs = self(**batch)
        val_loss, logits = outputs[:2]

        if self.hparams.num_labels > 1:
            preds = torch.argmax(logits, axis=1)
        elif self.hparams.num_labels == 1:
            preds = logits.squeeze()

        labels = batch["labels"]

        return {"loss": val_loss, "preds": preds, "labels": labels}

    def validation_epoch_end(self, outputs):
        if self.hparams.task_name == "mnli":
            for i, output in enumerate(outputs):
                # matched or mismatched
                split = self.hparams.eval_splits[i].split("_")[-1]
                preds = torch.cat([x["preds"] for x in output]).detach().cpu().numpy()
                labels = torch.cat([x["labels"] for x in output]).detach().cpu().numpy()
                loss = torch.stack([x["loss"] for x in output]).mean()
                self.log(f"val_loss_{split}", loss, prog_bar=True)
                split_metrics = {
                    f"{k}_{split}": v for k, v in self.metric.compute(predictions=preds, references=labels).items()
                }
                self.log_dict(split_metrics, prog_bar=True)
            return loss

        preds = torch.cat([x["preds"] for x in outputs]).detach().cpu().numpy()
        labels = torch.cat([x["labels"] for x in outputs]).detach().cpu().numpy()
        loss = torch.stack([x["loss"] for x in outputs]).mean()
        self.log("val_loss", loss, prog_bar=True)
        self.log_dict(self.metric.compute(predictions=preds, references=labels), prog_bar=True)

    def configure_optimizers(self):
        """Prepare optimizer and schedule (linear warmup and decay)"""
        model = self.model
        no_decay = ["bias", "LayerNorm.weight"]
        optimizer_grouped_parameters = [
            {
                "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
                "weight_decay": self.hparams.weight_decay,
            },
            {
                "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
                "weight_decay": 0.0,
            },
        ]
        optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=self.hparams.learning_rate, eps=self.hparams.adam_epsilon)

        scheduler = get_linear_schedule_with_warmup(
            optimizer,
            num_warmup_steps=self.hparams.warmup_steps,
            num_training_steps=self.trainer.estimated_stepping_batches,
        )
        scheduler = {"scheduler": scheduler, "interval": "step", "frequency": 1}
        return [optimizer], [scheduler]

def load_arguments(argv, file_path, learning_rate, epsilon, warmup_steps, entity, project_name):
  arg_help = "{0} -f <file_path> -l <learning_rate> -e <epsilon> -w <warmup_steps> -y <entity> -p <project_name>".format(argv[0])

  try:
    opts, args = getopt.getopt(argv[1:], "hf:l:e:w:y:p:", ["help", "file_path=",
      "learning_rate=", "epsilon=", "warmup_steps=", "entity=", "project_name="])
  except:
      print(arg_help)
      sys.exit(2)

  for opt, arg in opts:
    if opt in ("-h", "--help"):
      print(arg_help)  # print the help message
      sys.exit(2)
    elif opt in ("-f", "--file_path"):
      file_path = arg
    elif opt in ("-l", "--learning_rate"):
      print(f'l: {arg}')
      learning_rate = float(arg)
    elif opt in ("-e", "--epsilon"):
      print(f'e: {arg}')
      epsilon = float(arg)
    elif opt in ("-w", "--warmup_steps"):
      print(f'w: {arg}')
      warmup_steps = int(arg)
    elif opt in ("-y", "--entity"):
        entity = arg
    elif opt in ("-p", "--project_name"):
        project_name = arg

  return file_path, learning_rate, epsilon, warmup_steps, entity, project_name

if __name__ == "__main__":
  seed_everything(42)

  learning_rate: float = 2e-5
  epsilon: float = 1e-8
  warmup_steps: int = 0
  file_path = "model.pth"
  entity = "mlops"
  project_name = "project-2"

  file_path, learning_rate, epsilon, warmup_steps, entity, project_name = load_arguments(sys.argv, file_path, learning_rate, epsilon, warmup_steps, entity, project_name)

  wandb.init(entity=entity, project=project_name)

  wandb_logger = WandbLogger(project=project_name, log_model="all")

  dm = GLUEDataModule(
      model_name_or_path="distilbert-base-uncased",
      task_name="mrpc",
  )
  dm.setup("fit")
  model = GLUETransformer(
      model_name_or_path="distilbert-base-uncased",
      num_labels=dm.num_labels,
      eval_splits=dm.eval_splits,
      task_name=dm.task_name,
      learning_rate=learning_rate,
      adam_epsilon=epsilon,
      warmup_steps=warmup_steps
  )

  trainer = Trainer(
      max_epochs=3,
      accelerator="auto",
      devices=1 if torch.cuda.is_available() else None,
      logger=wandb_logger
  )

  trainer.fit(model, datamodule=dm)

  torch.save(model.state_dict(), file_path)
  wandb.finish()


import os
from tempfile import NamedTemporaryFile

import pytest
import toml
import yaml

from script_args_parser import ArgumentsParser

@pytest.fixture
def toml_file_path():
    mappings = {
        'first_arg': {
            'type': 'str',
            'description': 'Some fancy description',
            'cli_arg': '--cli-option-name',
            'env_var': 'ENV_VAR_NAME',
            'default_value': 'first_default'
        }
    }
    toml_file = NamedTemporaryFile(mode='w', delete=False)
    try:
        toml.dump(mappings, toml_file)
        toml_file.close()
        yield toml_file.name
    finally:
        os.unlink(toml_file.name)

def test_yaml(toml_file_path):
    values = {
        'string_arg': 'some (not so) random string',
        'int_arg': 123,
        'int_arg_as_string': '123',
        'path_arg': '.',
        'list_arg': [123, 156, 1956],
        'list_of_tuples_arg': [[123, 156], [12, 14], [145, 140]],
    }
    yaml_file = NamedTemporaryFile(mode='w', delete=False)
    try:
        yaml.dump(values, yaml_file)
        yaml_file.close()
        parser = ArgumentsParser.from_files(toml_file_path, [], yaml_config=yaml_file.name)
        assert parser.user_values == values
    finally:
        os.unlink(yaml_file.name)


import warnings
warnings.filterwarnings("ignore", category = UserWarning, module = "matplotlib")

import matplotlib.pyplot as pl
import numpy as np
import sklearn.learning_curve as curves
from sklearn.tree import DecisionTreeRegressor
from sklearn.cross_validation import ShuffleSplit, train_test_split

def ModelLearning(X, y):
    cv = ShuffleSplit(X.shape[0], n_iter = 10, test_size = 0.2, random_state = 0)
    train_sizes = np.rint(np.linspace(1, X.shape[0]*0.8 - 1, 9)).astype(int)
    fig = pl.figure(figsize=(10,7))
    for k, depth in enumerate([1,3,6,10]):
        regressor = DecisionTreeRegressor(max_depth = depth)
        sizes, train_scores, test_scores = curves.learning_curve(regressor, X, y, \
            cv = cv, train_sizes = train_sizes, scoring = 'r2')
        train_std = np.std(train_scores, axis = 1)
        train_mean = np.mean(train_scores, axis = 1)
        test_std = np.std(test_scores, axis = 1)
        test_mean = np.mean(test_scores, axis = 1)
        ax = fig.add_subplot(2, 2, k+1)
        ax.plot(sizes, train_mean, 'o-', color = 'r', label = 'Training Score')
        ax.plot(sizes, test_mean, 'o-', color = 'g', label = 'Testing Score')
        ax.fill_between(sizes, train_mean - train_std, \
            train_mean + train_std, alpha = 0.15, color = 'r')
        ax.fill_between(sizes, test_mean - test_std, \
            test_mean + test_std, alpha = 0.15, color = 'g')

        ax.set_title('max_depth = %s'%(depth))
        ax.set_xlabel('Number of Training Points')
        ax.set_ylabel('Score')
        ax.set_xlim([0, X.shape[0]*0.8])
        ax.set_ylim([-0.05, 1.05])

    ax.legend(bbox_to_anchor=(1.05, 2.05), loc='lower left', borderaxespad = 0.)
    fig.suptitle('Learning Performances', fontsize = 16, y = 1.03)
    fig.tight_layout()
    fig.show()

def ModelComplexity(X, y):
    cv = ShuffleSplit(X.shape[0], n_iter = 10, test_size = 0.2, random_state = 0)
    max_depth = np.arange(1,11)
    train_scores, test_scores = curves.validation_curve(DecisionTreeRegressor(), X, y, \
        param_name = "max_depth", param_range = max_depth, cv = cv, scoring = 'r2')
    train_mean = np.mean(train_scores, axis=1)
    train_std = np.std(train_scores, axis=1)
    test_mean = np.mean(test_scores, axis=1)
    test_std = np.std(test_scores, axis=1)
    pl.figure(figsize=(7, 5))
    pl.title('Complexity Performance')
    pl.plot(max_depth, train_mean, 'o-', color = 'r', label = 'Training Score')
    pl.plot(max_depth, test_mean, 'o-', color = 'g', label = 'Validation Score')
    pl.fill_between(max_depth, train_mean - train_std, \
        train_mean + train_std, alpha = 0.15, color = 'r')
    pl.fill_between(max_depth, test_mean - test_std, \
        test_mean + test_std, alpha = 0.15, color = 'g')

    pl.legend(loc = 'lower right')
    pl.xlabel('Maximum Depth')
    pl.ylabel('Score')
    pl.ylim([-0.05,1.05])
    pl.show()


from typing import List

class Solution:
    def findMaxConsecutiveOnes(self, nums: List[int]) -> int:

        if(len(nums) == 0):
            return 0

        if(1 not in nums):
            return 0

        if(0 not in nums):
            return len(nums)

        maxCnt = 0
        tempMaxCnt = 0
        len_ = len(nums)
        for i in range(len_):
            if(nums[i] == 0):
                maxCnt = max(tempMaxCnt, maxCnt)
                tempMaxCnt = 0
            else:
                if (nums[i] == 1):
                    tempMaxCnt += 1

        maxCnt = max(tempMaxCnt, maxCnt)
        return maxCnt

    # sliding window
    def findMaxConsecutiveOnes(self, nums: List[int]) -> int:

        maxCnt = 0
        window_start = 0
        for window_end in range(len(nums)):
            if(nums[window_end] == 1):
                maxCnt = max(maxCnt, window_end - window_start + 1)
            else:
                window_start = window_end + 1

        return maxCnt

s = Solution()
l = [1,1,0,1,1,1]
maxCnt = s.findMaxConsecutiveOnes(l)
print(maxCnt)


from urllib import response
import requests, base64, subprocess, argparse, datetime

def login():
    url = 'http://127.0.0.1:5000/soap/server_sa'
    header = {'SOAPAction': 'urn:NETGEAR-ROUTER:service:DeviceConfig:1#SOAPLogin'}
    data = '<Username>admin</Username><Password>password123</Password>'

    response = requests.post(url, headers=header, data=data)
    cookie = response.cookies.get_dict()["sess_id"]
    #print(cookie)
    return cookie

def sendUpdate(loginCookie, xml):
    timeout = 10
    url = 'http://127.0.0.1:5000/soap/server_sa'
    header = {'SOAPAction': 'urn:NETGEAR-ROUTER:service:DeviceConfig:1#UpdateNewFirmware'}
    cookie = {'sess_id': loginCookie}
    data = xml

    try:
        response = requests.post(url, data=data, headers=header, cookies=cookie, timeout=timeout)
        print(response.text)
    except requests.exceptions.ConnectionError:
        print("Connection aborted :(\n")
        if (crashTest()):
            dumpCrash(xml)

    except requests.exceptions.ReadTimeout:
        dumpCrash(xml)

def craftXML(payload):
    xml = '<?xml version="1.0"?>\r\n'
    xml += '<SOAP-ENV:Envelope xmlns:SOAP-ENV="http://schemas.xmlsoap.org/soap/envelope/" SOAP-ENV:encodingStyle="http://schemas.xmlsoap.$\r\n'
    xml += '<SOAP-ENV:Body>\r\n'
    xml += 'UpdateNewFirmware\r\n'
    xml += '<NewFirmware>\r\n'
    xml += base64.b64encode(payload).decode("utf-8")
    xml += '\r\n</NewFirmware>\r\n'
    xml += '</SOAP-ENV:Body>\r\n'
    xml += '</SOAP-ENV:Envelope>\r\n'

    return xml

def crashTest():
    url = 'http://127.0.0.1:5000/soap/server_sa'
    try:
        requests.get(url)
        return False
    except:
        return True

def dumpCrash(xml):
    print("\nFOUND POTENTIAL CRASH\n")
    ts = datetime.datetime.now().strftime("%m-%d-%Y_%H:%M:%S") + ".dmp"
    print("Saving dump at ./"+ts)
    f = open(ts, "w+")
    f.write(xml)
    f.close()
    exit(1)

def readPayload(path):
    with open(path, 'rb') as f:
        return f.read()

def main(firmware_path, nofuzz, seed):
    cookie = login()

    if nofuzz:
        payload = readPayload(firmware_path)
        xml = craftXML(payload)
        sendUpdate(cookie, xml)

    else:
        while(True):
            command = 'radamsa '

            if seed:
                command += '--seed ' + "% s" % seed + ' '
                seed += 1

            payload = subprocess.check_output(command + firmware_path, shell=True)
            xml = craftXML(payload)
            sendUpdate(cookie, xml)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Fuzz the upnpd daemon using radamsa by making http requests")
    parser.add_argument("-f", "--firmware", help="Firmware in .chk format", nargs=1, required=True)
    parser.add_argument("-s", "--seed", help="Initial seed for Radamsa", nargs=1, required=False, type=int)
    parser.add_argument("--nofuzz", help="Send the given firmware without fuzzing", action="store_true")
    args = parser.parse_args()

    seed = ''
    if args.seed:
        seed = args.seed[0]

    main(args.firmware[0], args.nofuzz, seed)


import random

import pytest
from brownie.test import given, strategy

@pytest.mark.math
class TestBitmap:
    @pytest.fixture(scope="module", autouse=True)
    def mockBitmap(self, MockBitmap, accounts):
        return accounts[0].deploy(MockBitmap)

    @given(bitmap=strategy("bytes32"))
    def test_is_bit_set(self, mockBitmap, bitmap):
        index = random.randint(1, 256)
        bitmask = list("".zfill(256))
        bitmask[index - 1] = "1"
        bitmask = "".join(bitmask)

        result = mockBitmap.isBitSet(bitmap, index)
        computedResult = (int(bitmask, 2) & int(bitmap.hex(), 16)) != 0
        assert result == computedResult

    @given(bitmap=strategy("bytes32"))
    def test_set_bit_on(self, mockBitmap, bitmap):
        index = random.randint(1, 256)
        newBitmap = mockBitmap.setBit(bitmap, index, True)
        assert mockBitmap.isBitSet(newBitmap, index)

    @given(bitmap=strategy("bytes32"))
    def test_set_bit_off(self, mockBitmap, bitmap):
        index = random.randint(1, 256)
        newBitmap = mockBitmap.setBit(bitmap, index, False)
        assert not mockBitmap.isBitSet(newBitmap, index)

    @given(bitmap=strategy("bytes32"))
    def test_total_bits_set(self, mockBitmap, bitmap):
        total = mockBitmap.totalBitsSet(bitmap)
        bitstring = "{:08b}".format(int(bitmap.hex(), 16))
        computedTotal = len([x for x in filter(lambda x: x == "1", list(bitstring))])

        assert total == computedTotal

    @given(bitmap=strategy("bytes32"))
    def test_msb_and_bit_num(self, mockBitmap, bitmap):
        msb = mockBitmap.getMSB(bitmap)
        bitNum = mockBitmap.getNextBitNum(bitmap)

        bitstring = "{:0256b}".format(int(bitmap.hex(), 16))
        indexes = [i for (i, b) in enumerate(list(bitstring)) if b == "1"]
        if len(indexes) == 0:
            assert msb == 0
            assert bitNum == 0
        else:
            assert msb == (255 - min(indexes))
            assert bitNum == (min(indexes) + 1)


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

"""
@date: 2020/9/23 下午2:35
@file: norm_helper.py
@author: zj
@description:
"""

import numpy as np
import torch
import torch.nn as nn
import torch.distributed as dist

from tsn.util import logging

def simple_group_split(world_size, rank, num_groups):
    # world_size: number of all processes
    # rank: current process ID
    # num_groups: number of groups in total, e.g. world_size=8 and you want to use 4 GPUs in a syncBN group, so num_groups=2
    groups = []
    rank_list = np.split(np.arange(world_size), num_groups)
    rank_list = [list(map(int, x)) for x in rank_list]
    for i in range(num_groups):
        groups.append(dist.new_group(rank_list[i]))
    group_size = world_size // num_groups

    logger = logging.setup_logging(__name__)
    logger.info(
        "Rank no.{} start sync BN on the process group of {}".format(rank, rank_list[rank // group_size]))
    return groups[rank // group_size]

def convert_sync_bn(model, process_group, device):
    # convert all BN layers in the model to syncBN
    for _, (child_name, child) in enumerate(model.named_children()):
        if isinstance(child, nn.modules.batchnorm._BatchNorm):
            m = nn.SyncBatchNorm.convert_sync_batchnorm(child, process_group)
            m = m.to(device=device)
            setattr(model, child_name, m)
        else:
            convert_sync_bn(child, process_group, device)

def get_norm(type):
    if type == 'BatchNorm2d':
        return nn.BatchNorm2d
    elif type == 'BatchNorm3d':
        return nn.BatchNorm3d
    elif type == 'SyncBatchNorm':
        return nn.SyncBatchNorm
    else:
        raise ValueError(f'{type} does not exists')


from django.utils.translation import gettext_lazy
from corehq.apps.reports.filters.select import MonthFilter
from corehq.apps.userreports.util import get_table_name
from memoized import memoized
from sqlagg.columns import SimpleColumn
from corehq.apps.reports.filters.base import BaseDrilldownOptionFilter, BaseSingleOptionFilter
from corehq.apps.reports.sqlreport import SqlData, DatabaseColumn
from corehq.apps.reports.filters.dates import DatespanFilter
from corehq.apps.reports.standard import DatespanMixin

class HierarchySqlData(SqlData):

    @property
    def table_name(self):
        return get_table_name(self.config['domain'], 'location_hierarchy')

    @property
    def filters(self):
        return []

    @property
    def group_by(self):
        return ['doc_id']

    @property
    def columns(self):
        return [
            DatabaseColumn('block', SimpleColumn('block')),
            DatabaseColumn('district', SimpleColumn('district')),
            DatabaseColumn('first_name', SimpleColumn('first_name')),
            DatabaseColumn('last_name', SimpleColumn('last_name')),
            DatabaseColumn('user_id', SimpleColumn('doc_id'))
        ]

class DrillDownOptionFilter(BaseDrilldownOptionFilter):
    label = gettext_lazy("Hierarchy")
    slug = "hierarchy"

    @property
    def filter_context(self):
        controls = []
        for level, label in enumerate(self.rendered_labels):
            controls.append({
                'label': label[0],
                'slug': label[1],
                'level': level,
            })

        return {
            'option_map': self.drilldown_map,
            'controls': controls,
            'selected': self.selected,
            'use_last': self.use_only_last,
            'notifications': self.final_notifications,
            'empty_text': self.drilldown_empty_text,
            'is_empty': not self.drilldown_map
        }

    @classmethod
    def get_labels(cls):
        return [('District', 'district'), ('Block', 'block'), ('AF', 'af')]

    @classmethod
    def _get_label_value(cls, request, label):
        slug = str(label[1])
        val = request.GET.getlist('%s_%s' % (cls.slug, str(label[1])))
        return {
            'slug': slug,
            'value': val,
        }

    @property
    @memoized
    def drilldown_map(self):
        def make_drilldown(hierarchy):
            hierarchy = [{
                "val": current[0] if isinstance(current, tuple) else current,
                "text": current[1] if isinstance(current, tuple) else current,
                "next": make_drilldown(next_level) if next_level else []
            } for current, next_level in hierarchy.items()]

            return sorted(hierarchy, key=lambda r: r['text'])

        return make_drilldown(self.get_hierarchy())

    @property
    def data_source(self):
        return HierarchySqlData

    def get_hierarchy(self):
        hierarchy = {}
        for location in self.data_source(config={'domain': self.domain}).get_data():
            district = location['district']
            block = location['block']
            user = ("%s %s" % (location['first_name'] or '', location['last_name'] or '')).strip()
            user_id = location['doc_id']
            if not (district and block and user):
                continue
            hierarchy[district] = hierarchy.get(district, {})
            hierarchy[district][block] = hierarchy[district].get(block, {})
            hierarchy[district][block][(user_id, user)] = None
        return hierarchy

class SampleFormatFilter(BaseSingleOptionFilter):
    slug = 'sf'
    label = gettext_lazy('Report type')
    default_text = "Format-1 for ASHA Sanginis"

    @property
    def options(self):
        return [
            ('sf2', 'Format-2 Consolidation of the Functionality numbers'),
            ('sf3', 'Format-3 Block Consolidation of the functionality status'),
            ('sf4', 'Format-4 Block Consolidation of the functionality status'),
            ('sf5', 'Format-5 Functionality of ASHAs in blocks')
        ]

class ASHAMonthFilter(MonthFilter):
    label = gettext_lazy("Last Reporting month of the quarter")

class LanguageFilter(BaseSingleOptionFilter):
    slug = 'lang'
    label = gettext_lazy('Language')
    default_text = "English"

    @property
    def options(self):
        return [('hin', gettext_lazy('Hindi'))]


#!/usr/bin/env python
import argparse
import atexit
import logging
import os
import shutil
import sys
from pathlib import Path

from checkov.arm.runner import Runner as arm_runner
from checkov.cloudformation.runner import Runner as cfn_runner
from checkov.common.bridgecrew.platform_integration import bc_integration
from checkov.common.bridgecrew.image_scanning.image_scanner import image_scanner
from checkov.common.goget.github.get_git import GitGetter
from checkov.common.runners.runner_registry import RunnerRegistry, OUTPUT_CHOICES
from checkov.common.util.banner import banner as checkov_banner
from checkov.common.util.consts import DEFAULT_EXTERNAL_MODULES_DIR
from checkov.common.util.docs_generator import print_checks
from checkov.common.util.runner_dependency_handler import RunnerDependencyHandler
from checkov.common.util.type_forcers import convert_str_to_bool
from checkov.terraform.runner import Runner as tf_graph_runner
from checkov.helm.runner import Runner as helm_runner
from checkov.kubernetes.runner import Runner as k8_runner
from checkov.logging_init import init as logging_init
from checkov.runner_filter import RunnerFilter
from checkov.serverless.runner import Runner as sls_runner
from checkov.terraform.plan_runner import Runner as tf_plan_runner
from checkov.dockerfile.runner import Runner as dockerfile_runner
from checkov.version import version

outer_registry = None

logging_init()
logger = logging.getLogger(__name__)
checkov_runner_module_names = ['cfn', 'tf', 'k8', 'sls', 'arm', 'tf_plan', 'helm']
checkov_runners = ['cloudformation', 'terraform', 'kubernetes', 'serverless', 'arm', 'terraform_plan', 'helm', 'dockerfile']

# Check runners for necessary system dependencies.
runnerDependencyHandler = RunnerDependencyHandler(checkov_runner_module_names, globals())
runnerDependencyHandler.validate_runner_deps()

def run(banner=checkov_banner, argv=sys.argv[1:]):
    parser = argparse.ArgumentParser(description='Infrastructure as code static analysis')
    add_parser_args(parser)
    args = parser.parse_args(argv)

    # bridgecrew uses both the urllib3 and requests libraries, while checkov uses the requests library.
    # Allow the user to specify a CA bundle to be used by both libraries.
    bc_integration.setup_http_manager(args.ca_certificate)

    # Disable runners with missing system dependencies
    args.skip_framework = runnerDependencyHandler.disable_incompatible_runners(args.skip_framework)

    runner_filter = RunnerFilter(framework=args.framework, skip_framework=args.skip_framework, checks=args.check, skip_checks=args.skip_check,
                                 download_external_modules=convert_str_to_bool(args.download_external_modules),
                                 external_modules_download_path=args.external_modules_download_path,
                                 evaluate_variables=convert_str_to_bool(args.evaluate_variables), runners=checkov_runners)
    if outer_registry:
        runner_registry = outer_registry
        runner_registry.runner_filter = runner_filter
    else:
        runner_registry = RunnerRegistry(banner, runner_filter, tf_graph_runner(), cfn_runner(), k8_runner(), sls_runner(),
                                         arm_runner(), tf_plan_runner(), helm_runner(),dockerfile_runner())
    if args.version:
        print(version)
        return

    if args.bc_api_key == '':
        parser.error('The --bc-api-key flag was specified but the value was blank. If this value was passed as a secret, you may need to double check the mapping.')
    elif args.bc_api_key:
        logger.debug(f'Using API key ending with {args.bc_api_key[-8:]}')

        if args.repo_id is None:
            parser.error("--repo-id argument is required when using --bc-api-key")
        if len(args.repo_id.split('/')) != 2:
            parser.error("--repo-id argument format should be 'organization/repository_name' E.g "
                         "bridgecrewio/checkov")

        source = os.getenv('BC_SOURCE', 'cli')
        source_version = os.getenv('BC_SOURCE_VERSION', version)
        logger.debug(f'BC_SOURCE = {source}, version = {source_version}')
        try:
            bc_integration.setup_bridgecrew_credentials(bc_api_key=args.bc_api_key, repo_id=args.repo_id,
                                                        skip_fixes=args.skip_fixes,
                                                        skip_suppressions=args.skip_suppressions,
                                                        source=source, source_version=source_version, repo_branch=args.branch)
            excluded_paths = bc_integration.get_excluded_paths()
            runner_filter.excluded_paths = excluded_paths
        except Exception as e:
            logger.error('An error occurred setting up the Bridgecrew platform integration. Please check your API token and try again.', exc_info=True)
            return
    else:
        logger.debug('No API key found. Scanning locally only.')

    guidelines = {}
    if not args.no_guide:
        guidelines = bc_integration.get_guidelines()
    if args.check and args.skip_check:
        parser.error("--check and --skip-check can not be applied together. please use only one of them")
        return
    if args.list:
        print_checks(framework=args.framework)
        return
    external_checks_dir = get_external_checks_dir(args)
    url = None

    if args.directory:
        exit_codes = []
        for root_folder in args.directory:
            file = args.file
            scan_reports = runner_registry.run(root_folder=root_folder, external_checks_dir=external_checks_dir,
                                               files=file, guidelines=guidelines, bc_integration=bc_integration)
            if bc_integration.is_integration_configured():
                bc_integration.persist_repository(root_folder)
                bc_integration.persist_scan_results(scan_reports)
                url = bc_integration.commit_repository(args.branch)

            exit_codes.append(runner_registry.print_reports(scan_reports, args, url))

        exit_code = 1 if 1 in exit_codes else 0
        return exit_code
    elif args.file:
        scan_reports = runner_registry.run(external_checks_dir=external_checks_dir, files=args.file,
                                           guidelines=guidelines, bc_integration=bc_integration)
        if bc_integration.is_integration_configured():
            files = [os.path.abspath(file) for file in args.file]
            root_folder = os.path.split(os.path.commonprefix(files))[0]
            bc_integration.persist_repository(root_folder)
            bc_integration.persist_scan_results(scan_reports)
            url = bc_integration.commit_repository(args.branch)
        return runner_registry.print_reports(scan_reports, args, url)
    elif args.docker_image:
        if args.bc_api_key is None:
            parser.error("--bc-api-key argument is required when using --docker-image")
            return
        if args.dockerfile_path is None:
            parser.error("--dockerfile-path argument is required when using --docker-image")
            return
        if args.branch is None:
            parser.error("--branch argument is required when using --docker-image")
            return
        image_scanner.scan(args.docker_image, args.dockerfile_path)
    else:
        print(f"{banner}")

        bc_integration.onboarding()

def add_parser_args(parser):
    parser.add_argument('-v', '--version',
                        help='version', action='store_true')
    parser.add_argument('-d', '--directory', action='append',
                        help='IaC root directory (can not be used together with --file).')
    parser.add_argument('-f', '--file', action='append',
                        help='IaC file(can not be used together with --directory)')
    parser.add_argument('--external-checks-dir', action='append',
                        help='Directory for custom checks to be loaded. Can be repeated')
    parser.add_argument('--external-checks-git', action='append',
                        help='Github url of external checks to be added. \n you can specify a subdirectory after a '
                             'double-slash //. \n cannot be used together with --external-checks-dir')
    parser.add_argument('-l', '--list', help='List checks', action='store_true')
    parser.add_argument('-o', '--output', nargs='?', choices=OUTPUT_CHOICES,
                        default='cli',
                        help='Report output format')
    parser.add_argument('--no-guide', action='store_true',
                        default=False,
                        help='do not fetch bridgecrew guide in checkov output report')
    parser.add_argument('--quiet', action='store_true',
                        default=False,
                        help='in case of CLI output, display only failed checks')
    parser.add_argument('--compact', action='store_true',
                        default=False,
                        help='in case of CLI output, do not display code blocks')
    parser.add_argument('--framework', help='filter scan to run only on a specific infrastructure code frameworks',
                        choices=checkov_runners + ["all"],
                        default='all')
    parser.add_argument('--skip-framework', help='filter scan to skip specific infrastructure code frameworks. \n'
                                                 'will be included automatically for some frameworks if system dependencies are missing.',
                        choices=checkov_runners,
                        default=None)
    parser.add_argument('-c', '--check',
                        help='filter scan to run only on a specific check identifier(allowlist), You can '
                             'specify multiple checks separated by comma delimiter', default=None)
    parser.add_argument('--skip-check',
                        help='filter scan to run on all check but a specific check identifier(denylist), You can '
                             'specify multiple checks separated by comma delimiter', default=None)
    parser.add_argument('-s', '--soft-fail',
                        help='Runs checks but suppresses error code', action='store_true')
    parser.add_argument('--bc-api-key', help='Bridgecrew API key')
    parser.add_argument('--docker-image', help='Scan docker images by name or ID. Only works with --bc-api-key flag')
    parser.add_argument('--dockerfile-path', help='Path to the Dockerfile of the scanned docker image')
    parser.add_argument('--repo-id',
                        help='Identity string of the repository, with form <repo_owner>/<repo_name>')
    parser.add_argument('-b', '--branch',
                        help="Selected branch of the persisted repository. Only has effect when using the --bc-api-key flag",
                        default='master')
    parser.add_argument('--skip-fixes',
                        help='Do not download fixed resource templates from Bridgecrew. Only has effect when using the --bc-api-key flag',
                        action='store_true')
    parser.add_argument('--skip-suppressions',
                        help='Do not download preconfigured suppressions from the Bridgecrew platform. Code comment suppressions will still be honored. '
                             'Only has effect when using the --bc-api-key flag',
                        action='store_true')
    parser.add_argument('--download-external-modules',
                        help="download external terraform modules from public git repositories and terraform registry",
                        default=os.environ.get('DOWNLOAD_EXTERNAL_MODULES', False))
    parser.add_argument('--external-modules-download-path',
                        help="set the path for the download external terraform modules",
                        default=DEFAULT_EXTERNAL_MODULES_DIR)
    parser.add_argument('--evaluate-variables',
                        help="evaluate the values of variables and locals",
                        default=True)
    parser.add_argument('-ca', '--ca-certificate',
                        help='custom CA (bundle) file', default=None)

def get_external_checks_dir(args):
    external_checks_dir = args.external_checks_dir
    if args.external_checks_git:
        git_getter = GitGetter(args.external_checks_git[0])
        external_checks_dir = [git_getter.get()]
        atexit.register(shutil.rmtree, str(Path(external_checks_dir[0]).parent))
    return external_checks_dir

if __name__ == '__main__':
    exit(run())


import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import seaborn as sns
sns.set()

def graficar(dominio, imagen, titulo, x, y):
    plt.title(titulo)
    plt.xlabel(x)
    plt.ylabel(y)
    plt.grid(True)
    plt.plot(dominio, imagen, '-g')

def ajuste_lineal(x, A, B):
    return A * x + B

def ajuste_exp(x, A, B):
    return A * np.exp(x * B)

dias, casos_totales = np.loadtxt('covid.csv', delimiter=';', skiprows=1, unpack=True)

# np.set_printoptions(suppress=True)
# print(dias)
# print(casos_totales)

dia_cuarentena = 18

dias_antes_cuarentena = dias[:dia_cuarentena]
casos_antes_cuarentena = casos_totales[:dia_cuarentena]

dominio = dias_antes_cuarentena
imagen = casos_antes_cuarentena

popt, pcov = curve_fit(ajuste_exp, dominio, imagen)
A, C = popt
new_var_x = np.linspace(min(dias), max(dias), 1000)
new_var_y = ajuste_exp(new_var_x, A, C)

graficar(dias, casos_totales,
         "Gráfico de los casos totales en Argentina",
         "Días",
         "Casos totales")
plt.axvline(x=18, ls='--')
plt.text(10, 8000, 'Comienzo de la cuarentena')

plt.plot(new_var_x, new_var_y, ':r', label='Ajuste exp')
plt.legend()
plt.show()


import datetime
import os
import numpy as np
import torch
from torch import nn
import matplotlib
from tqdm import tqdm
matplotlib.use('Agg')
import scipy.signal
from matplotlib import pyplot as plt
from torch.utils.tensorboard import SummaryWriter
from utils.metrics import itm_eval

class LossHistory():
    def __init__(self, log_dir, model, input_shape):
        self.log_dir    = log_dir
        self.losses     = []
        self.val_loss   = []

        os.makedirs(self.log_dir)
        try:
            self.writer     = SummaryWriter(self.log_dir)
            # dummy_input     = torch.randn(2, 3, input_shape[0], input_shape[1])
            # text_input      = ["OK", "OK"]
            # self.writer.add_graph(model, [dummy_input, text_input])
        except:
            pass

    def append_loss(self, epoch, loss, val_loss):
        if not os.path.exists(self.log_dir):
            os.makedirs(self.log_dir)

        self.losses.append(loss)
        self.val_loss.append(val_loss)

        with open(os.path.join(self.log_dir, "epoch_loss.txt"), 'a') as f:
            f.write(str(loss))
            f.write("\n")
        with open(os.path.join(self.log_dir, "epoch_val_loss.txt"), 'a') as f:
            f.write(str(val_loss))
            f.write("\n")

        self.writer.add_scalar('loss', loss, epoch)
        self.writer.add_scalar('val_loss', val_loss, epoch)
        self.loss_plot()

    def loss_plot(self):
        iters = range(len(self.losses))

        plt.figure()
        plt.plot(iters, self.losses, 'red', linewidth = 2, label='train loss')
        plt.plot(iters, self.val_loss, 'coral', linewidth = 2, label='val loss')
        try:
            if len(self.losses) < 25:
                num = 5
            else:
                num = 15
            plt.plot(iters, scipy.signal.savgol_filter(self.losses, num, 3), 'green', linestyle = '--', linewidth = 2, label='smooth train loss')
            plt.plot(iters, scipy.signal.savgol_filter(self.val_loss, num, 3), '#8B4513', linestyle = '--', linewidth = 2, label='smooth val loss')
        except:
            pass
        plt.grid(True)
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.legend(loc="upper right")
        plt.savefig(os.path.join(self.log_dir, "epoch_loss.png"))
        plt.cla()
        plt.close("all")

def is_parallel(model):
    # Returns True if model is of type DP or DDP
    return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)

def de_parallel(model):
    # De-parallelize a model: returns single-GPU model if model is of type DP or DDP
    return model.module if is_parallel(model) else model

class EvalCallback():
    def __init__(self, net, gen_val, log_dir, cuda, batch_size=32, eval_flag=True, period=1):
        super(EvalCallback, self).__init__()

        self.net                = net
        self.gen_val            = gen_val
        self.log_dir            = log_dir
        self.cuda               = cuda
        self.batch_size         = batch_size
        self.eval_flag          = eval_flag
        self.period             = period

        self.txt_r1 = []
        self.txt_r5 = []
        self.img_r1 = []
        self.img_r5 = []
        self.epoches = []

    def on_epoch_end(self, epoch, model_eval):
        if epoch % self.period == 0 and self.eval_flag:
            self.net = de_parallel(model_eval)

            i_features = []
            t_features = []
            for iteration, batch in tqdm(enumerate(self.gen_val)):
                images, texts = batch
                with torch.no_grad():
                    if self.cuda:
                        images  = images.cuda()

                    images_feature = self.net.encode_image(images)
                    i_features.append(images_feature)

            texts       = self.gen_val.dataset.text
            num_text    = len(texts)
            for i in tqdm(range(0, num_text, self.batch_size)):
                text = texts[i: min(num_text, i + self.batch_size)]
                with torch.no_grad():
                    texts_feature = self.net.encode_text(text)
                    t_features.append(texts_feature)

            i_features = torch.cat(i_features, 0)
            t_features = torch.cat(t_features, 0)

            i_features  = i_features / i_features.norm(dim=-1, keepdim=True)
            t_features  = t_features / t_features.norm(dim=-1, keepdim=True)

            logits_per_image    = i_features @ t_features.t()
            logits_per_text     = logits_per_image.t()

            logits_per_image    = logits_per_image.cpu().numpy()
            logits_per_text     = logits_per_text.cpu().numpy()

            itm_results = itm_eval(logits_per_image, logits_per_text, self.gen_val.dataset.txt2img, self.gen_val.dataset.img2txt)

            self.txt_r1.append(itm_results['txt_r1'])
            self.txt_r5.append(itm_results['txt_r5'])
            self.img_r1.append(itm_results['img_r1'])
            self.img_r5.append(itm_results['img_r5'])
            self.epoches.append(epoch)

            with open(os.path.join(self.log_dir, "epoch_R@1_R@5_R@10.txt"), 'a') as f:
                f.write(str(itm_results))
                f.write("\n")

            plt.figure()
            plt.plot(self.epoches, self.txt_r1, 'red', linewidth = 2, label='txt_r1')
            plt.plot(self.epoches, self.txt_r5, 'green', linewidth = 2, label='txt_r5')
            plt.plot(self.epoches, self.img_r1, 'blue', linewidth = 2, label='img_r1')
            plt.plot(self.epoches, self.img_r5, 'pink', linewidth = 2, label='img_r5')

            plt.grid(True)
            plt.xlabel('Epoch')
            plt.ylabel('Recall')
            plt.title('A Recall Curve')
            plt.legend(loc="upper right")

            plt.savefig(os.path.join(self.log_dir, "epoch_recall.png"))
            plt.cla()
            plt.close("all")
            print(itm_results)
            print("Get recall done.")


# Licensed under a 3-clause BSD style license - see LICENSE.rst
from __future__ import absolute_import, division, print_function, unicode_literals
import numpy as np
from astropy.utils import lazyproperty
import astropy.units as u
from ..utils.fitting import Fit
from ..stats import cash
from ..maps import Map, MapAxis

__all__ = ["MapFit", "MapEvaluator"]

class MapFit(Fit):
    """Perform sky model likelihood fit on maps.

    This is the first go at such a class. It's geared to the
    `~gammapy.spectrum.SpectrumFit` class which does the 1D spectrum fit.

    Parameters
    ----------
    model : `~gammapy.cube.models.SkyModel`
        Fit model
    counts : `~gammapy.maps.WcsNDMap`
        Counts cube
    exposure : `~gammapy.maps.WcsNDMap`
        Exposure cube
    background : `~gammapy.maps.WcsNDMap`
        Background Cube
    mask : `~gammapy.maps.WcsNDMap`
        Mask to apply for the fit. All the pixels that contain 1 or True are included
        in the fit, all others are ignored.
    psf : `~gammapy.cube.PSFKernel`
        PSF kernel
    edisp : `~gammapy.irf.EnergyDispersion`
        Energy dispersion
    """

    def __init__(
        self, model, counts, exposure, background=None, mask=None, psf=None, edisp=None
    ):
        if mask is not None and mask.data.dtype != np.dtype("bool"):
            raise ValueError("mask data must have dtype bool")

        self._model = model
        self.counts = counts
        self.exposure = exposure
        self.background = background
        self.mask = mask
        self.psf = psf
        self.edisp = edisp

        self.evaluator = MapEvaluator(
            model=self._model,
            exposure=exposure,
            background=self.background,
            psf=self.psf,
            edisp=self.edisp,
        )

    @property
    def stat(self):
        """Likelihood per bin given the current model parameters"""
        npred = self.evaluator.compute_npred()
        return cash(n_on=self.counts.data, mu_on=npred)

    def total_stat(self, parameters):
        """Total likelihood given the current model parameters"""
        self._model.parameters = parameters
        if self.mask:
            stat = self.stat[self.mask.data]
        else:
            stat = self.stat
        return np.sum(stat, dtype=np.float64)

class MapEvaluator(object):
    """Sky model evaluation on maps.

    This is a first attempt to compute flux as well as predicted counts maps.

    The basic idea is that this evaluator is created once at the start
    of the analysis, and pre-computes some things.
    It it then evaluated many times during likelihood fit when model parameters
    change, re-using pre-computed quantities each time.
    At the moment it does some things, e.g. cache and re-use energy and coordinate grids,
    but overall it is not an efficient implementation yet.

    For now, we only make it work for 3D WCS maps with an energy axis.
    No HPX, no other axes, those can be added later here or via new
    separate model evaluator classes.

    We should discuss how to organise the model and IRF evaluation code,
    and things like integrations and convolutions in a good way.

    Parameters
    ----------
    model : `~gammapy.cube.models.SkyModel`
        Sky model
    exposure : `~gammapy.maps.Map`
        Exposure map
    background : `~gammapy.maps.Map`
        background map
    psf : `~gammapy.cube.PSFKernel`
        PSF kernel
    edisp : `~gammapy.irf.EnergyDispersion`
        Energy dispersion
    """

    def __init__(
        self, model=None, exposure=None, background=None, psf=None, edisp=None
    ):
        self.model = model
        self.exposure = exposure
        self.background = background
        self.psf = psf
        self.edisp = edisp

    @lazyproperty
    def geom(self):
        """This will give the energy axes in e_true"""
        return self.exposure.geom

    @lazyproperty
    def geom_image(self):
        return self.geom.to_image()

    @lazyproperty
    def energy_center(self):
        """True energy axis bin centers (`~astropy.units.Quantity`)"""
        energy_axis = self.geom.get_axis_by_name("energy")
        energy = energy_axis.center * energy_axis.unit
        return energy[:, np.newaxis, np.newaxis]

    @lazyproperty
    def energy_edges(self):
        """Energy axis bin edges (`~astropy.units.Quantity`)"""
        energy_axis = self.geom.get_axis_by_name("energy")
        energy = energy_axis.edges * energy_axis.unit
        return energy[:, np.newaxis, np.newaxis]

    @lazyproperty
    def energy_bin_width(self):
        """Energy axis bin widths (`astropy.units.Quantity`)"""
        return np.diff(self.energy_edges, axis=0)

    @lazyproperty
    def lon_lat(self):
        """Spatial coordinate pixel centers.

        Returns ``lon, lat`` tuple of `~astropy.units.Quantity`.
        """
        lon, lat = self.geom_image.get_coord()
        return (u.Quantity(lon, "deg", copy=False), u.Quantity(lat, "deg", copy=False))

    @lazyproperty
    def lon(self):
        return self.lon_lat[0]

    @lazyproperty
    def lat(self):
        return self.lon_lat[1]

    @lazyproperty
    def solid_angle(self):
        """Solid angle per pixel"""
        return self.geom.solid_angle()

    @lazyproperty
    def bin_volume(self):
        """Map pixel bin volume (solid angle times energy bin width)."""
        omega = self.solid_angle
        de = self.energy_bin_width
        return omega * de

    def compute_dnde(self):
        """Compute model differential flux at map pixel centers.

        Returns
        -------
        model_map : `~gammapy.maps.Map`
            Sky cube with data filled with evaluated model values.
            Units: ``cm-2 s-1 TeV-1 deg-2``
        """
        coord = (self.lon, self.lat, self.energy_center)
        dnde = self.model.evaluate(*coord)
        return dnde

    def compute_flux(self):
        """Compute model integral flux over map pixel volumes.

        For now, we simply multiply dnde with bin volume.
        """
        dnde = self.compute_dnde()
        volume = self.bin_volume
        flux = dnde * volume
        return flux

    def apply_exposure(self, flux):
        """Compute npred cube

        For now just divide flux cube by exposure
        """
        npred = (flux * self.exposure.quantity).to_value("")
        return self.exposure.copy(data=npred)

    def apply_psf(self, npred):
        """Convolve npred cube with PSF"""
        return npred.convolve(self.psf)

    def apply_edisp(self, npred):
        """Convolve map data with energy dispersion.

        Parameters
        ----------
        npred : `~gammapy.maps.Map`
            Predicted counts in true energy bins

        Returns
        ---------
        npred_reco : `~gammapy.maps.Map`
            Predicted counts in reco energy bins
        """
        loc = npred.geom.get_axis_index_by_name("energy")
        data = np.rollaxis(npred.data, loc, len(npred.data.shape))
        data = np.dot(data, self.edisp.pdf_matrix)
        data = np.rollaxis(data, -1, loc)
        e_reco_axis = MapAxis.from_edges(
            self.edisp.e_reco.bins, unit=self.edisp.e_reco.unit
        )
        geom_ereco = self.exposure.geom.to_image().to_cube(axes=[e_reco_axis])
        npred = Map.from_geom(geom_ereco, unit="")
        npred.data = data
        return npred

    def compute_npred(self):
        """
        Evaluate model predicted counts.

        Returns
        -------
        npred.data : ~numpy.ndarray
            array of the predicted counts in each bin (in reco energy)
        """
        flux = self.compute_flux()
        npred = self.apply_exposure(flux)
        if self.psf is not None:
            npred = self.apply_psf(npred)
        if self.edisp is not None:
            npred = self.apply_edisp(npred)
        if self.background:
            npred.data += self.background.data
        return npred.data


# -*- coding: UTF-8 -*-
import logging
import zlib
import pickle

import re
from copy import copy

from adventure.game import Game
from adventure import load_advent_dat
from moneypenny.handler import HandlerType
from moneypenny.conf import settings

log = logging.getLogger(__name__)

def _make_suspend(game, path):
    _path = copy(path)
    self = game  # ignore propper binding

    def suspend(verb, obj=None):
        savefile = open(_path, 'wb')
        r = self.random_generator  # must replace live object with static state
        self.random_state = r.getstate()
        f = self.i_suspend
        try:
            del self.random_generator
            del self.i_suspend
            del self.t_suspend
            savefile.write(zlib.compress(pickle.dumps(self), 9))
        finally:
            self.random_generator = r
            self.i_suspend = f
            self.t_suspend = f
            savefile.close()
        self.write('Game saved')

    return suspend

class Adventure(metaclass=HandlerType):
    regex = re.compile('adventure (?P<cmd>start|continue)', re.I)

    def __init__(self, event, conversation):
        match = self.regex.match(event['body'])
        if not match:
            if Adventure in conversation.data:
                game = conversation.data[Adventure]
                words = re.findall('\w+', event['body'])
                event.reply(game.do_command(words).strip()).send()
                if not game.is_finished:
                    conversation.handler = Adventure
            else:
                log.warn('An adventure session was lost.')
                event.reply('You have to start an adventure first.').send()
        elif match.group('cmd') == 'start':
            game = Game()
            load_advent_dat(game)
            savefile = settings['textadventure']['file'].format(conversation.jid)
            game.i_suspend = _make_suspend(game, savefile)
            game.t_suspend = game.i_suspend
            game.start()
            conversation.handler = Adventure
            event.reply(game.output.strip()).send()
            conversation.data[Adventure] = game
        else:
            savefile = settings['textadventure']['file'].format(conversation.jid)
            game = Game.resume(savefile)
            game.i_suspend = _make_suspend(game, savefile)
            game.t_suspend = game.i_suspend
            event.reply('GAME RESTORED').send()
            conversation.handler = Adventure
            conversation.data[Adventure] = game


from aiodocker.util import task
from aiodocker.exceptions import NotFound, ConflictError, ValidationError
from aiodocker.exceptions import ServerError, UnexpectedError
from aiodocker.formatters import from_search_results
from aiodocker.helpers import stream_json
from aiodocker.util import parse_name
import logging

log = logging.getLogger(__name__)

class RegistryEndpoint:

    def __init__(self, api, *, registry):
        self.api = api
        self.registry = registry

    @task
    def pull(self, image, *, tag=None):
        """
        Pull image from registry.

        Parameters:
            image (str): in the form of {repo}:{tag}
            tag (str): optionnal. if not set,
                       it will use tag in image parameters
        """
        if tag and ':' in image:
            raise ValidationError('tag cannot be declared '
                                  'into image and tag parameters')
        elif tag is None:
            image, tag = parse_name(image)

        path = '/images/create'
        params = {
            'fromImage': image,
            'registry': self.registry,
        }
        if tag:
            params['tag'] = tag

        response = yield from self.api.post(path, params=params)
        if response.status == 200:
            for data in (yield from stream_json(response)):
                log.info(data)
            return True

        data = yield from response.text()
        if response.status == 500:
            raise ServerError(data)
        raise UnexpectedError(response.status, data)

    @task
    def push(self, image, *, tag=None):
        """
        Push image to registry.

        Parameters:
            image (str): in the form of {repo}:{tag}
            tag (str): optionnal. if not set,
                       it will use tag in image parameters
        """
        if tag and ':' in image:
            raise ValidationError('tag cannot be declared '
                                  'into image and tag parameters')
        elif tag is None:
            image, tag = parse_name(image)

        params = {}
        if tag:
            params['tag'] = tag

        if self.registry:
            path = '/images/%s/%s/push' % (self.registry, image)
        else:
            path = '/images/%s/push' % image

        response = yield from self.api.post(path, params=params)
        if response.status == 200:
            data = yield from response.json()
            return data

        data = yield from response.text()
        if response.status == 404:
            raise NotFound(data)
        elif response.status == 500:
            raise ServerError(data)
        raise UnexpectedError(response.status, data)

    @task
    def tag(self, image, *, tag=None, force=False):
        """
        Tag image into registry.

        Parameters:
            image (str): in the form of {repo}:{tag}
            tag (str): optionnal. if not set,
                       it will use tag in image parameters
            force (bool): force tag even if conflicts?
        """
        if tag and ':' in image:
            raise ValidationError('tag cannot be declared '
                                  'into image and tag parameters')
        elif tag is None:
            image, tag = parse_name(image)

        path = '/images/%s/tag' % image
        params = {}
        if tag:
            params['tag'] = tag

        response = yield from self.api.post(path, params=params)
        if response.status == 201:
            return True

        data = yield from response.text()
        if response.status == 400:
            raise ValidationError(data)
        elif response.status == 404:
            raise NotFound(data)
        elif response.status == 409:
            raise ConflictError(data)
        elif response.status == 500:
            raise ServerError(data)
        raise UnexpectedError(response.status, data)

class DockerHubEndpoint(RegistryEndpoint):
    """
    https://hub.docker.com endpoint
    """

    def __init__(self, api):
        self.api = api
        self.registry = None

    @task
    def search(self, term):
        """Search for images"""
        path = '/images/search'
        params = {
            'term': term
        }

        response = yield from self.api.get(path, params=params)
        if response.status == 200:
            data = yield from response.json()
            return from_search_results(data)

        data = yield from response.text()
        if response.status == 500:
            raise ServerError(data)
        raise UnexpectedError(response.status, data)


#!/usr/bin/python3
"""ecursive function that queries the Reddit API"""

import requests

def recurse(subreddit, hot_list=[], count=0, next_page=None):
    """return list containing titles of all hot articles"""
    headers = {"User-Agent": "0x16. API_advanced-m00n-die"}

    url = "https://www.reddit.com/r/{}/hot/.json".format(subreddit)

    params = {"limit": 50, "next_page": next_page, "count": count}

    res = requests.get(url, headers=headers,
                       params=params, allow_redirects=False)

    if res.status_code != 200:
        return None
    json_res = res.json().get("data")
    next_page = json_res.get("next_page")
    count += json_res.get("dist")
    children = json_res.get("children")

    for child in children:
        title = child.get("data").get("title")
        hot_list.append(title)

    if next_page is not None:
        return recurse(subreddit, hot_list, count, next_page)
        return hot_list


# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
The identities module is part of the nmeta suite

It provides an abstraction for participants (identities), using
a MongoDB database for storage and data retention maintenance.

Identities are identified via  TBD....

There are methods (see class docstring) that provide harvesting
of identity metadata and various retrieval searches
"""
import sys
import struct

#*** For packet methods:
import socket

#*** Import dpkt for packet parsing:
import dpkt

#*** mongodb Database Import:
import pymongo
from pymongo import MongoClient

#*** For timestamps:
import datetime

#*** For logging configuration:
from baseclass import BaseClass

#*** For Regular Expression searches:
import re

#*** For hashing of identities:
import hashlib

#*** How long in seconds to cache ARP responses for (in seconds):
ARP_CACHE_TIME = 14400
#*** DHCP lease time to use if none present (in seconds):
DHCP_DEFAULT_LEASE_TIME = 3600

class Identities(BaseClass):
    """
    An object that represents identity metadata

    Main function used to harvest identity metadata:
    (assumes class instantiated as an object called 'ident')

        ident.harvest(pkt, flow.packet)
            Passed a raw packet and packet metadata from flow object.
            Check a packet_in event and harvest any relevant identity
            indicators to metadata

    Functions available for Classifiers:
    (assumes class instantiated as an object called 'ident')

        ident.findbymac(mac_address)
            Look up identity object for a MAC address

        ident.findbynode(host_name)
            Look up identity object by host name (aka node)
            Additionally, can set:
                regex=True       Treat service_name as a regular expression
                harvest_type=    Specify what type of harvest (i.e. DHCP)

        ident.findbyservice(service_name)
            Look up identity object by service name
            Additionally, can set:
                regex=True        Treat service_name as a regular expression
                harvest_type=     Specify what type of harvest (i.e. DNS_A)
                ip_address=       Look for specific IP address

    See function docstrings for more information
    """

    def __init__(self, config, policy):
        """
        Initialise an instance of the Identities class
        """
        self.policy = policy
        #*** Required for BaseClass:
        self.config = config
        #*** Set up Logging with inherited base class method:
        self.configure_logging(__name__, "identities_logging_level_s",
                                       "identities_logging_level_c")
        #*** Get parameters from config:
        mongo_addr = config.get_value("mongo_addr")
        mongo_port = config.get_value("mongo_port")
        mongo_dbname = self.config.get_value("mongo_dbname")
        #*** Max bytes of the identities capped collection:
        identities_max_bytes = config.get_value("identities_max_bytes")
        #*** How far back in time to go back looking for an identity:
        self.identity_time_limit = datetime.timedelta \
                              (seconds=config.get_value("identity_time_limit"))
        #*** Max bytes of the dhcp_messages capped collection:
        dhcp_messages_max_bytes = config.get_value("dhcp_messages_max_bytes")
        #*** How far back in time to go back looking for an dhcp message:
        self.dhcp_messages_time_limit = datetime.timedelta \
                         (seconds=config.get_value("dhcp_messages_time_limit"))

        #*** Start mongodb:
        self.logger.info("Connecting to MongoDB database...")
        mongo_client = MongoClient(mongo_addr, mongo_port)

        #*** Connect to MongoDB nmeta database:
        db_nmeta = mongo_client[mongo_dbname]

        #*** Delete (drop) previous identities collection if it exists:
        self.logger.debug("Deleting previous identities MongoDB collection...")
        db_nmeta.identities.drop()

        #*** Create the identities collection, specifying capped option
        #*** with max size in bytes, so MongoDB handles data retention:
        self.identities = db_nmeta.create_collection('identities', capped=True,
                                            size=identities_max_bytes)

        #*** Index to improve look-up performance:
        self.identities.create_index([('valid_from', pymongo.DESCENDING),
                                        ('valid_to', pymongo.DESCENDING),
                                        ('ip_address', pymongo.ASCENDING),
                                        ('harvest_type', pymongo.ASCENDING)
                                        ],
                                        unique=False)

        #*** Index to improve MAC address look-up performance:
        self.identities.create_index([('mac_address', pymongo.ASCENDING),
                             ('valid_from', pymongo.DESCENDING)], unique=False)

        #*** Index to improve Node (host_name) look-up performance:
        self.identities.create_index([('host_name', pymongo.ASCENDING),
                             ('valid_from', pymongo.DESCENDING)], unique=False)

        #*** Index to improve Service look-up performance:
        self.identities.create_index([('service_name', pymongo.ASCENDING),
                             ('valid_from', pymongo.DESCENDING)], unique=False)

        #*** Delete (drop) previous dhcp_messages collection if it exists:
        self.logger.debug("Deleting previous dhcp_messages MongoDB "
                                                               "collection...")
        db_nmeta.dhcp_messages.drop()

        #*** Create the dhcp_messages collection, specifying capped option
        #*** with max size in bytes, so MongoDB handles data retention:
        self.dhcp_messages = db_nmeta.create_collection('dhcp_messages',
                                     capped=True, size=dhcp_messages_max_bytes)

        #*** Index dhcp_messages to improve look-up performance:
        self.dhcp_messages.create_index([('ingest_time', pymongo.DESCENDING),
                                        ('transaction_id', pymongo.ASCENDING),
                                        ('message_type', pymongo.ASCENDING)
                                        ],
                                        unique=False)

    class Identity(object):
        """
        An object that represents an individual Identity Indicator
        """
        def __init__(self):
            #*** Initialise identity variables:
            self.dpid = 0
            self.in_port = 0
            self.mac_address = ""
            self.ip_address = ""
            self.harvest_type = 0
            self.harvest_time = 0
            self.host_name = ""
            self.host_type = ""
            self.host_os = ""
            self.host_desc = ""
            self.service_name = ""
            self.service_alias = ""
            self.user_id = ""
            self.valid_from = ""
            self.valid_to = ""
            self.id_hash = ""
            self.location_logical = ""
            self.location_physical = ""

        def dbdict(self):
            """
            Return a dictionary object of identity metadata
            parameters for storing in the database
            """
            return self.__dict__

    class DHCPMessage(object):
        """
        An object that represents an individual DHCP message.
        Used for storing DHCP state by recording DHCP events
        """
        def __init__(self):
            #*** Initialise identity variables:
            self.dpid = 0
            self.in_port = 0
            self.ingest_time = 0
            self.eth_src = 0
            self.eth_dst = 0
            self.ip_src = 0
            self.ip_dst = 0
            self.tp_src = 0
            self.tp_dst = 0
            self.transaction_id = 0
            self.message_type = 0
            self.host_name = ""
            self.ip_assigned = 0
            self.ip_dhcp_server = 0
            self.lease_time = 0

        def dbdict(self):
            """
            Return a dictionary object of dhcp message
            parameters for storing in the database
            """
            return self.__dict__

    def harvest(self, pkt, flow_pkt):
        """
        Passed a raw packet and packet metadata from flow object.
        Check a packet_in event and harvest any relevant identity
        indicators to metadata
        """
        #*** ARP:
        if flow_pkt.eth_type == 2054:
            self.harvest_arp(pkt, flow_pkt)

        #*** DHCP:
        elif flow_pkt.eth_type == 2048 and flow_pkt.proto == 17 and \
                              (flow_pkt.tp_dst == 67 or flow_pkt.tp_dst == 68):
            self.harvest_dhcp(flow_pkt)

        #*** LLDP:
        elif flow_pkt.eth_type == 35020:
            self.harvest_lldp(flow_pkt)

        #*** DNS:
        elif (flow_pkt.proto == 6 or flow_pkt.proto == 17) and \
                        (flow_pkt.tp_src == 53 or flow_pkt.tp_src == 53):
            self.harvest_dns(flow_pkt)

        else:
            #*** Not an identity indicator
            return 0

    def harvest_arp(self, pkt, flow_pkt):
        """
        Harvest ARP identity metadata into database.
        Passed packet-in metadata from flow object.
        Check ARP reply and harvest identity
        indicators to metadata
        """
        self.logger.debug("Harvesting metadata from ARP request")

        eth = dpkt.ethernet.Ethernet(pkt)
        pkt_arp = eth.arp
        if pkt_arp:
            #*** It's an ARP, but is it a reply (opcode 2) for IPv4?:
            if pkt_arp.op == 2 and pkt_arp.pro == 2048:
                #*** Instantiate an instance of Indentity class:
                ident = self.Identity()
                ident.dpid = flow_pkt.dpid
                ident.in_port = flow_pkt.in_port
                ident.mac_address = mac_addr(pkt_arp.sha)
                ident.ip_address = socket.inet_ntoa(pkt_arp.spa)
                ident.harvest_type = 'ARP'
                ident.harvest_time = flow_pkt.timestamp
                ident.valid_from = flow_pkt.timestamp
                ident.valid_to = flow_pkt.timestamp + \
                                    datetime.timedelta(0, ARP_CACHE_TIME)
                ident.id_hash = self._hash_identity(ident)
                ident.location_logical = self.policy.locations.get_location \
                                                    (ident.dpid, ident.in_port)
                db_dict = ident.dbdict()
                #*** Write ARP identity metadata to database collection:
                self.logger.debug("writing db_dict=%s", db_dict)
                self.identities.insert_one(db_dict)
        return 1

    def harvest_dhcp(self, flow_pkt):
        """
        Harvest DHCP identity metadata into database.
        Passed packet-in metadata from flow object.
        Check LLDP TLV fields and harvest any relevant identity
        indicators to metadata
        """
        self.logger.debug("Harvesting metadata from DHCP request")
        dhcp_hostname = ""
        dhcp_leasetime = 0
        #*** Use dpkt to parse UDP DHCP data:
        try:
            pkt_dhcp = dpkt.dhcp.DHCP(flow_pkt.payload)
        except:
            exc_type, exc_value, exc_traceback = sys.exc_info()
            self.logger.error("DHCP extraction failed "
                        "Exception %s, %s, %s",
                         exc_type, exc_value, exc_traceback)
            return 0
        #*** Turn DHCP options list of tuples into a dictionary:
        dhcp_opts = dict(pkt_dhcp.opts)
        self.logger.debug("dhcp_opts=%s", dhcp_opts)
        #*** Get the type of the DHCP message:
        try:
            dhcp_type = ord(dhcp_opts[53])
        except:
            exc_type, exc_value, exc_traceback = sys.exc_info()
            self.logger.error("DHCP type extraction failed "
                        "Exception %s, %s, %s",
                         exc_type, exc_value, exc_traceback)
            return 0
        #*** Do stuff based on the DHCP message type:
        if dhcp_type == dpkt.dhcp.DHCPDISCOVER:
            self.logger.debug("Matched DHCPDISCOVER, TBD - not handled")
        elif dhcp_type == dpkt.dhcp.DHCPOFFER:
            self.logger.debug("Matched DHCPOFFER, TBD - not handled")
        elif dhcp_type == dpkt.dhcp.DHCPREQUEST:
            self.logger.debug("Matched DHCPREQUEST")
            if dpkt.dhcp.DHCP_OPT_HOSTNAME in dhcp_opts:
                #*** Instantiate an instance of DHCP class:
                self.dhcp_msg = self.DHCPMessage()
                self.dhcp_msg.dpid = flow_pkt.dpid
                self.dhcp_msg.in_port = flow_pkt.in_port
                self.dhcp_msg.ingest_time = flow_pkt.timestamp
                self.dhcp_msg.eth_src = flow_pkt.eth_src
                self.dhcp_msg.eth_dst = flow_pkt.eth_dst
                self.dhcp_msg.ip_src = flow_pkt.ip_src
                self.dhcp_msg.ip_dst = flow_pkt.ip_dst
                self.dhcp_msg.tp_src = flow_pkt.tp_src
                self.dhcp_msg.tp_dst = flow_pkt.tp_dst
                self.dhcp_msg.transaction_id = hex(pkt_dhcp.xid)
                self.dhcp_msg.host_name = str(dhcp_opts
                                                 [dpkt.dhcp.DHCP_OPT_HOSTNAME])
                self.dhcp_msg.message_type = 'DHCPREQUEST'
                #*** Record DHCP event to db collection:
                db_dict = self.dhcp_msg.dbdict()
                #*** Write DHCP message to db collection:
                self.logger.debug("writing dhcp_messages db_dict=%s", db_dict)
                self.dhcp_messages.insert_one(db_dict)
                return 1
        elif dhcp_type == dpkt.dhcp.DHCPDECLINE:
            self.logger.debug("Matched DHCPDECLINE, TBD - not handled")
        elif dhcp_type == dpkt.dhcp.DHCPACK:
            self.logger.debug("Matched DHCPACK")
            xid = hex(pkt_dhcp.xid)
            #*** Look up dhcp db collection for DHCPREQUEST:
            db_data = {'transaction_id': xid,
                        'message_type': 'DHCPREQUEST'}
            #*** Filter by documents that are still within 'best before' time:
            db_data['ingest_time'] = {'$gte': datetime.datetime.now() -
                                                 self.dhcp_messages_time_limit}
            #*** Run db search:
            result = self.dhcp_messages.find(db_data).sort('ingest_time', -1) \
                                                                      .limit(1)
            if result.count():
                result0 = list(result)[0]
                self.logger.debug("Found DHCPREQUEST for DHCPACK")
                #*** Found a DHCP Request for the ACK, record results:
                #*** Instantiate an instance of DHCP class:
                self.dhcp_msg = self.DHCPMessage()
                self.dhcp_msg.dpid = flow_pkt.dpid
                self.dhcp_msg.in_port = flow_pkt.in_port
                self.dhcp_msg.ingest_time = flow_pkt.timestamp
                self.dhcp_msg.eth_src = flow_pkt.eth_src
                self.dhcp_msg.eth_dst = flow_pkt.eth_dst
                self.dhcp_msg.ip_src = flow_pkt.ip_src
                self.dhcp_msg.ip_dst = flow_pkt.ip_dst
                self.dhcp_msg.tp_src = flow_pkt.tp_src
                self.dhcp_msg.tp_dst = flow_pkt.tp_dst
                self.dhcp_msg.transaction_id = hex(pkt_dhcp.xid)
                self.dhcp_msg.ip_assigned = \
                           socket.inet_ntoa(struct.pack(">L", pkt_dhcp.yiaddr))
                if dpkt.dhcp.DHCP_OPT_LEASE_SEC in dhcp_opts:
                    self.dhcp_msg.lease_time = struct.unpack('>L', dhcp_opts
                                            [dpkt.dhcp.DHCP_OPT_LEASE_SEC])[0]
                    self.logger.debug("Found dhcp_leasetime=%s",
                                                      self.dhcp_msg.lease_time)
                else:
                    self.dhcp_msg.lease_time = DHCP_DEFAULT_LEASE_TIME
                    self.logger.debug("Using default dhcp_leasetime=%s",
                                                      self.dhcp_msg.lease_time)
                self.dhcp_msg.message_type = 'DHCPACK'
                #*** Record DHCP event to db collection:
                db_dict = self.dhcp_msg.dbdict()
                #*** Write DHCP message to db collection:
                self.logger.debug("writing dhcp_messages db_dict=%s", db_dict)
                self.dhcp_messages.insert_one(db_dict)
                #*** Instantiate an instance of Identity class:
                ident = self.Identity()
                ident.dpid = flow_pkt.dpid
                ident.in_port = flow_pkt.in_port
                ident.mac_address = flow_pkt.eth_dst
                ident.ip_address = self.dhcp_msg.ip_assigned
                ident.harvest_type = 'DHCP'
                ident.host_name = result0['host_name']
                ident.harvest_time = flow_pkt.timestamp
                ident.valid_from = flow_pkt.timestamp
                #*** Calculate validity:
                ident.valid_to = flow_pkt.timestamp + \
                                datetime.timedelta(0, self.dhcp_msg.lease_time)
                ident.id_hash = self._hash_identity(ident)
                ident.location_logical = self.policy.locations.get_location \
                                                    (ident.dpid, ident.in_port)
                db_dict = ident.dbdict()
                #*** Write DHCP identity metadata to db collection:
                self.logger.debug("writing db_dict=%s", db_dict)
                self.identities.insert_one(db_dict)
                return 1
            else:
                self.logger.debug("Prev DHCP host_name not found")
                return 0
        elif dhcp_type == dpkt.dhcp.DHCPNAK:
            self.logger.debug("Matched DHCPNAK, TBD - not handled")
        elif dhcp_type == dpkt.dhcp.DHCPRELEASE:
            self.logger.debug("Matched DHCPRELEASE, TBD - not handled")
        elif dhcp_type == dpkt.dhcp.DHCPINFORM:
            self.logger.debug("Matched DHCPINFORM, TBD - not handled")
        else:
            self.logger.debug("Unknown DHCP option 53 value: %s", dhcp_type)
            return 0

    def harvest_lldp(self, flow_pkt):
        """
        Harvest LLDP identity metadata into database.
        Passed packet-in metadata from flow object.
        Check LLDP TLV fields and harvest any relevant identity
        indicators to metadata
        """
        self.logger.debug("Checking LLDP for metadata")
        payload = flow_pkt.payload
        lldp_dict = self._parse_lldp_detail(payload)
        if not len(lldp_dict):
            self.logger.warning("Failed to parse LLDP")
            return 0
        self.logger.debug("LLDP parsed %s", lldp_dict)
        #*** Instantiate an instance of Indentity class:
        ident = self.Identity()
        if 'system_name' in lldp_dict:
            ident.host_name = lldp_dict['system_name']
        if 'system_desc' in lldp_dict:
            ident.host_desc = lldp_dict['system_desc']
        if 'TTL' in lldp_dict:
            ttl = lldp_dict['TTL']
        else:
            #*** TBD, handle this better:
            ttl = 60
        ident.dpid = flow_pkt.dpid
        ident.in_port = flow_pkt.in_port
        ident.mac_address = flow_pkt.eth_src
        ident.harvest_type = 'LLDP'
        ident.harvest_time = flow_pkt.timestamp
        ident.valid_from = flow_pkt.timestamp
        #*** valid to based on LLDP TTL:
        ident.valid_to = flow_pkt.timestamp + \
                                    datetime.timedelta(0, ttl)
        #*** Try looking up an IP for the LLDP source MAC:
        ident2 = self.findbymac(ident.mac_address)
        if 'ip_address' in ident2:
            ident.ip_address = ident2['ip_address']
            self.logger.debug("Found ip=%s for LLDP flow_hash=%s",
                                    ident.ip_address, flow_pkt.flow_hash)
        else:
            self.logger.debug("Could not find IP for LLDP flow_hash=%s",
                                    flow_pkt.flow_hash)
        ident.id_hash = self._hash_identity(ident)
        ident.location_logical = self.policy.locations.get_location \
                                                    (ident.dpid, ident.in_port)
        #*** Write LLDP identity metadata to db collection:
        db_dict = ident.dbdict()
        self.logger.debug("writing db_dict=%s", db_dict)
        self.identities.insert_one(db_dict)
        return 1

    def harvest_dns(self, flow_pkt):
        """
        Harvest DNS identity metadata into database.
        Passed packet-in metadata from flow object.
        Check DNS answer(s) and harvest any relevant identity
        indicators to metadata
        """
        self.logger.debug("Checking DNS for metadata")
        #*** Use dpkt to parse DNS:
        try:
            pkt_dns = dpkt.dns.DNS(flow_pkt.payload)
        except:
            exc_type, exc_value, exc_traceback = sys.exc_info()
            self.logger.error("DNS extraction failed "
                        "Exception %s, %s, %s",
                         exc_type, exc_value, exc_traceback)
            return 0
        answers = pkt_dns.an
        for answer in answers:
            if answer.type == 1:
                #*** DNS A Record:
                ident = self.Identity()
                ident.dpid = flow_pkt.dpid
                ident.in_port = flow_pkt.in_port
                ident.harvest_type = 'DNS_A'
                ident.ip_address = socket.inet_ntoa(answer.rdata)
                ident.service_name = answer.name
                ident.harvest_time = flow_pkt.timestamp
                ident.valid_from = flow_pkt.timestamp
                ident.valid_to = flow_pkt.timestamp + \
                                    datetime.timedelta(0, answer.ttl)
                ident.id_hash = self._hash_identity(ident)
                ident.location_logical = self.policy.locations.get_location \
                                                    (ident.dpid, ident.in_port)
                db_dict = ident.dbdict()
                #*** Write DNS identity metadata to database collection:
                self.logger.debug("writing db_dict=%s", db_dict)
                self.identities.insert_one(db_dict)
            elif answer.type == 5:
                #*** DNS CNAME Record:
                ident = self.Identity()
                ident.dpid = flow_pkt.dpid
                ident.in_port = flow_pkt.in_port
                ident.harvest_type = 'DNS_CNAME'
                ident.service_name = answer.name
                ident.service_alias = answer.cname
                ident.harvest_time = flow_pkt.timestamp
                ident.valid_from = flow_pkt.timestamp
                ident.valid_to = flow_pkt.timestamp + \
                                    datetime.timedelta(0, answer.ttl)
                ident.id_hash = self._hash_identity(ident)
                ident.location_logical = self.policy.locations.get_location \
                                                    (ident.dpid, ident.in_port)
                db_dict = ident.dbdict()
                #*** Write DNS identity metadata to database collection:
                self.logger.debug("writing db_dict=%s", db_dict)
                self.identities.insert_one(db_dict)
            else:
                #*** Not a type that we handle yet
                self.logger.debug("Unhandled DNS answer type=%s", answer.type)

    def findbymac(self, mac_addr, test=0):
        """
        Passed a MAC address and reverse search identities collection
        returning first match as a dictionary version of
        an Identity class, or empty dictionary if not found

        Setting test=1 returns database query execution statistics
        """
        db_data = {'mac_address': mac_addr}
        if not test:
            result = self.identities.find(db_data).sort('valid_from', -1).limit(1)
        else:
            return self.identities.find(db_data).sort('valid_from', -1).limit(1).explain()
        if result.count():
            result0 = list(result)[0]
            self.logger.debug("found result=%s len=%s", result0, len(result0))
            return result0
        else:
            self.logger.debug("mac_addr=%s not found", mac_addr)
            return {}

    def findbynode(self, host_name, harvest_type='any', regex=False, test=0):
        """
        Find by node name
        Pass it the name of the node to search for. Additionally,
        can set:
          regex=True       Treat service_name as a regular expression
          harvest_type=    Specify what type of harvest (i.e. DHCP)
        Returns a dictionary version of an Identity class, or 0 if not found

        Setting test=1 returns database query execution statistics
        """
        db_data = {'host_name': host_name}
        if harvest_type != 'any':
            #*** Filter by harvest type:
            db_data['harvest_type'] = harvest_type
        if regex:
            #*** Regular expression search on service name:
            regx = re.compile(host_name)
            db_data['host_name'] = regx
        #*** Filter by documents that are still within 'best before' time:
        db_data['valid_to'] = {'$gte': datetime.datetime.now()}
        #*** Run db search:
        if not test:
            result = self.identities.find(db_data).sort('valid_from', -1).limit(1)
        else:
            return self.identities.find(db_data).sort('valid_from', -1).limit(1).explain()
        if result.count():
            result0 = list(result)[0]
            self.logger.debug("found result=%s len=%s", result0, len(result0))
            return result0
        else:
            self.logger.debug("host_name=%s not found", host_name)
            return 0

    def findbyservice(self, service_name, harvest_type='any', regex=False,
                        ip_address='any', test=0):
        """
        Find by service name
        Pass it the name of the service to search for. Additionally,
        can set:
          regex=True        Treat service_name as a regular expression
          harvest_type=     Specify what type of harvest (i.e. DNS_A)
          ip_address=       Look for specific IP address
        Returns boolean

        Setting test=1 returns database query execution statistics
        """
        db_data = {'service_name': service_name}
        if harvest_type != 'any':
            #*** Filter by harvest type:
            db_data['harvest_type'] = harvest_type
        if ip_address != 'any':
            #*** Filter by IP address:
            db_data['ip_address'] = ip_address
        if regex:
            #*** Regular expression search on service name:
            regx = re.compile(service_name)
            db_data['service_name'] = regx
        #*** Filter by documents that are still within 'best before' time:
        db_data['valid_to'] = {'$gte': datetime.datetime.now()}
        #*** Run db search:
        if not test:
            result = self.identities.find(db_data).sort('valid_from', -1).limit(1)
        else:
            return self.identities.find(db_data).sort('valid_from', -1).limit(1).explain()
        if result.count():
            result0 = list(result)[0]
            self.logger.debug("found result=%s len=%s", result0, len(result0))
            return result0
        else:
            self.logger.debug("service_name=%s not found", service_name)
            return 0

    def _hash_identity(self, ident):
        """
        Generate a hash of the current identity used for deduplication
        where the same identity is received periodically, or from multiple
        sources.
        """
        hash_result = hashlib.md5()
        id_tuple = (ident.harvest_type,
                    ident.host_name,
                    ident.service_name,
                    ident.user_id)
        id_tuple_as_string = str(id_tuple)
        hash_result.update(id_tuple_as_string)
        return hash_result.hexdigest()

    #=================== PRIVATE ==============================================
    def _parse_lldp_detail(self, lldpPayload):
        """
        Parse basic LLDP parameters from an LLDP packet payload
        """
        result = {}

        while lldpPayload:
            tlv_header = struct.unpack("!H", lldpPayload[:2])[0]
            tlv_type = tlv_header >> 9
            tlv_len = (tlv_header & 0x01ff)
            lldpDU = lldpPayload[2:tlv_len + 2]
            if tlv_type == 0:
                #*** TLV type 0 is end of TLVs so break the while loop:
                break
            else:
                tlv_subtype = struct.unpack("!B", lldpDU[0:1]) \
                                                    if tlv_type is 2 else ""
                startbyte = 1 if tlv_type is 2 else 0
                tlv_datafield = lldpDU[startbyte:tlv_len]
            #*** Pull out values from specific TLVs:
            if tlv_type == 3:
                result['TTL'] = struct.unpack("!h", tlv_datafield)[0]
            elif tlv_type == 4:
                result['port_desc'] = tlv_datafield
            elif tlv_type == 5:
                result['system_name'] = tlv_datafield
            elif tlv_type == 6:
                result['system_desc'] = tlv_datafield
            else:
                pass
            lldpPayload = lldpPayload[2 + tlv_len:]
        return result

def mac_addr(address):
    """
    Convert a MAC address to a readable/printable string
    """
    return ':'.join('%02x' % ord(b) for b in address)


import numpy as np
from collections import defaultdict
from enum import Enum

class AgentType(Enum):
    SARSA=0
    SARSAMAX=1
    EXPECTED_SARSA=2

class Agent:

    def __init__(self, nA=6):
        """ Initialize agent.

        Params
            - nA: number of actions available to the agent
        """
        self.nA = nA
        self.Q = defaultdict(lambda: np.zeros(self.nA))
        self.epochs = 0 #number of iterations
        self.epsilon = 1.0
        self.alpha = 0.01
        self.gamma = 1.0
        # the following type could be "SARSA", "SARSAmax", or "Expected SARSA"
        self.TYPE = AgentType.EXPECTED_SARSA

    def __epsilon_greedy(self, state):
        """
        Gets the policy using epsilon-greedy algorithm
        Args:
            - q_state: the state of the agent
        Returns:
            - pi: the policy of the agent in that state using epsilon-greedy algorithm
        """
        pi = np.full((self.nA), 1.0*self.epsilon/(self.nA)) #equiprobable-random values
        # get maximum action given state
        a_max = np.argmax(self.Q[state])
        # update the value of that maximum action
        pi[a_max] += 1.0 - self.epsilon
        return pi

    def select_action(self, state):
        """
        Given the state, select an action using epsilon-greedy policy
        Params
            - state: the current state of the environment
        Returns
            - action: an integer, compatible with the task's action space
        """
        policy = self.__epsilon_greedy(state)
        action = np.random.choice(np.arange(self.nA), p=policy)
        return action

    def step(self, state, action, reward, next_state, done):
        """ Update the agent's knowledge, using the most recently sampled tuple.
        Params
            - state: the previous state of the environment
            - action: the agent's previous choice of action
            - reward: last reward received
            - next_state: the current state of the environment
            - done: whether the episode is complete (True or False)
        """
        #update epsilon
        self.epochs += 1
        self.epsilon = 1/self.epochs
        self.epsilon = max(self.epsilon, 0.005) # constant after 200 epochs
        # update state-action function
        if self.TYPE == AgentType.SARSA:
            next_action = self.select_action(state)
            self.Q[state][action] = ( (1-self.alpha)*self.Q[state][action] )\
                                   + ( self.alpha*(reward + self.gamma*self.Q[next_state][next_action]) )
        elif self.TYPE == AgentType.SARSAMAX:
            self.Q[state][action] = ( (1-self.alpha)*self.Q[state][action] )\
                                   + ( self.alpha*(reward + self.gamma*np.max(self.Q[next_state])) )
        elif self.TYPE == AgentType.EXPECTED_SARSA:
            policy = self.__epsilon_greedy(next_state)
            self.Q[state][action] = ( (1-self.alpha)*self.Q[state][action] )\
                                   + ( self.alpha*(reward + self.gamma*sum(policy * self.Q[next_state])) )

"""
SARSA:
Episode 20000/20000 || Best average reward -5.684
SARSAMAX:
Episode 20000/20000 || Best average reward 9.3398
EXPECTED_SARSA
Episode 20000/20000 || Best average reward 9.3153
"""

06 - Reinforcement Learning/Lesson 07 - Taxi Miniproject/agent_SOLVED.py


'''
classifier for specific tasks
'''

from types import new_class
import torch
import torch.nn as nn

class MTNet(nn.Module):
    def __init__(self, params=dict()):
        super(MTNet, self).__init__()
        '''
        input: h_mix (input_embed_dim), one-dimensional embedding
        '''
        self.input_embed_dim = int(params['embedding_final'])
        self.n_class = int(params['sample_class_num'])
        self.classifier = nn.Linear(self.input_embed_dim, self.n_class)

    def forward(self, x):
        logits = self.classifier(x).unsqueeze(0)
        Y_hat = torch.topk(logits, 1, dim=-1)[1] # shape: (1,1)
        hazards = torch.sigmoid(logits) # shape: (1, n_class)
        S = torch.cumprod(1 - hazards, dim=-1) # shape: (1, n_class)
        return (hazards, S, Y_hat, logits)

# params = {
#     'embedding_coattn_dim':256,
#     'sample_class_num':4
# }
# model = MTNet(params = params)
# input = torch.rand(size=(256,))
# hazards, S, Y_hat = model.forward(input)
# print('ok')
# print('end')