"""
Distributed Learning using Pytorch's torch.distributed.launcher and
torch.nn.parallel.distributed_c10d on FfDL.
"""

import time
import argparse
import sys
import os
import threading
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data
import torch.distributed.c10d

from math import ceil
from random import Random
from torch.multiprocessing import Process
from torch.autograd import Variable
from torchvision import datasets, transforms

class Net(nn.Module):
    """ Network architecture. """

    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

def partition_dataset(batch_size, world_size):
    """ Partitioning MNIST """
    vision_data = os.environ.get("DATA_DIR") + "/data"
    dataset = datasets.MNIST(
        vision_data,
        train=True,
        download=True,
        transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307, ), (0.3081, ))
        ]))

    bsz = int(batch_size / float(world_size))
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=bsz, shuffle=True)

    return dataloader, bsz

def average_gradients(model, world_size, pg):
    """ Gradient averaging. """
    for param in model.parameters():
        torch.distributed.c10d.all_reduce(param.grad.data, pg)
        param.grad.data /= world_size

def run(rank, world_rank, world_size, group, batch_size, is_gpu):
    """ Distributed Synchronous SGD Example """
    torch.manual_seed(1234)
    size = os.environ.get("WORLD_SIZE")
    result_dir = os.environ.get("RESULT_DIR") + "/saved_model"
    train_set, bsz = partition_dataset(batch_size, world_size)
    # For GPU use
    if is_gpu:
        # device = torch.device("cuda:{}".format(0))
        # model = Net().to(device)
        model = Net().cuda()
    else:
        model = Net()
        model = model
    model = torch.nn.parallel._DistributedDataParallelC10d(model, group)
#    model = model.cuda(rank)
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)

    num_batches = ceil(len(train_set.dataset) / float(bsz))
    for epoch in range(10):
        epoch_loss = 0.0
        for data, target in train_set:
            # For GPU use
            if is_gpu:
                data, target = data.cuda(), target.cuda()
            else:
                data, target = Variable(data), Variable(target)
#            data, target = Variable(data.cuda(rank)), Variable(target.cuda(rank))
            optimizer.zero_grad()
            output = model(data)
            loss = F.nll_loss(output, target)
            epoch_loss += loss.item()
            loss.backward()
            if not (size == 1):
                average_gradients(model, world_size, group)
            optimizer.step()
        print('Process ', world_rank,
              ', epoch ', epoch, ': ',
              epoch_loss / num_batches)
    torch.save(model.state_dict(), result_dir)

# Change 'backend' to appropriate backend identifier
def init_processes(local_rank, world_rank, world_size, fn, batch_size, shared_file, is_gpu, backend):
    """ Initialize the distributed environment. """
    print("World Rank: " + str(world_rank) + "  Local Rank: " + str(local_rank)  + " connected")
    pg = torch.distributed.c10d.ProcessGroupGloo(shared_file, world_rank, world_size)
    pg.Options.timeout = 300.0 * world_size
    print("GROUP CREATED")
    fn(local_rank, world_rank, world_size, pg, batch_size, is_gpu)

def local_process(target, args):
    return Process(target=target, args=args)

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument('--batch_size', help='Specify the batch size to be used in training')
    args = parser.parse_args()

    batch_size = args.batch_size
    # Default batch size is set to 1024. When using a large numbers of learners, a larger batch
    # size is sometimes necessary to see speed improvements.
    if batch_size is None:
        batch_size = 1024
    else:
        batch_size = int(batch_size)

    start_time = time.time()
    num_gpus = int(float(os.environ.get("GPU_COUNT")))
    if num_gpus == 0:
        world_size = int(os.environ.get("NUM_LEARNERS"))
    else:
        world_size = num_gpus * int(os.environ.get("NUM_LEARNERS"))
    data_dir = "/job/" + os.environ.get("TRAINING_ID")
    processes = []

    start_time = time.time()
    shared_file = torch.distributed.c10d.FileStore(data_dir)
    processes = []

    print("SHARED FILE PATH: " + data_dir, " WORLD_SIZE: " + str(world_size))
    world_rank = int(os.environ.get("LEARNER_ID")) - 1

    if num_gpus == 0:
        args = (0, world_rank, world_size, run, batch_size, shared_file, True, 'gloo')
        p = local_process(init_processes, args)
        p.start()
        processes.append(p)
    else:
        print("Opening processes")
        for local_rank in range(0, num_gpus):
            args = (local_rank, world_rank, world_size, run, batch_size, shared_file, True, 'gloo')
            p = local_process(init_processes, args)
            print("Process Created")
            p.start()
            processes.append(p)
            print("Process Added")

    for p in processes:
        print("Waiting on Process")
        p.join()

    print("COMPLETION TIME: " + str(time.time() - start_time))

    if int(os.environ.get("LEARNER_ID")) != 1:
        while True:
            time.sleep(1000000)


# coding:utf-8
__author__ = 'Albert'

from django.conf.urls import patterns, url

from materials import views

materials_views = patterns(
    '',
    url(r'^mis/update/student/materials$',  views.update_student_material),
    url(r'^mis/material/packages$', views.list_material_packages),
    url(r'^mis/set/student/material$', views.set_student_material),
    url(r'^mis/material/upload$', views.material_upload),
    url(r'^mis/material/preview$', views.material_preview),
    url(r'^mis/unfinished/classes$', views.unfinished_classes),
)


# This code is modified from https://github.com/jakesnell/prototypical-networks

import backbone
import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np
import torch.nn.functional as F
from methods.meta_template import MetaTemplate

import utils

####
import torch
import torchvision

from torch import optim, nn
from torch.nn import *
from torchvision import transforms
import torch.utils.data as data
import torch.nn.functional as F
from torch.autograd import Variable

# from PIL import Image
# from pathlib import Path

# for ROC curve
# from sklearn import metrics
# from scipy import interp

import pandas as pd
import numpy as np

# import matplotlib.pyplot as plt #, mpld3
# import matplotlib

import math
from typing import *
import time
import datetime

# from IPython.display import display, clear_output
# from IPython.display import HTML

###
import warnings

# from torch.nn.module import Module
from torch.nn import functional as F
from torch.nn import _reduction as _Reduction

from torch import Tensor
from typing import Optional

class _Loss(nn.Module):
    reduction: str

    def __init__(self, size_average=None, reduce=None, reduction: str = 'mean') -> None:
        super(_Loss, self).__init__()
        if size_average is not None or reduce is not None:
            self.reduction = _Reduction.legacy_get_string(size_average, reduce)
        else:
            self.reduction = reduction

class _WeightedLoss(_Loss):
    def __init__(self, weight: Optional[Tensor] = None, size_average=None, reduce=None, reduction: str = 'mean') -> None:
        super(_WeightedLoss, self).__init__(size_average, reduce, reduction)
        self.register_buffer('weight', weight)

class CrossEntropyLoss(_WeightedLoss):
    __constants__ = ['ignore_index', 'reduction']
    ignore_index: int

    def __init__(self, weight: Optional[Tensor] = None, size_average=None, ignore_index: int = -100,
                 reduce=None, reduction: str = 'mean') -> None:
        super(CrossEntropyLoss, self).__init__(weight, size_average, reduce, reduction)
        self.ignore_index = ignore_index

    def forward(self, input: Tensor, target: Tensor) -> Tensor:
#         print('input',input)
#         print('target',target)
        return F.cross_entropy(input, target, weight=self.weight,
                               ignore_index=self.ignore_index, reduction=self.reduction)

class LargeMarginCosineLoss(nn.Module):
    """
    Reference:
    H. Wang et al. CosFace: Large Margin Cosine Loss for Deep Face Recognition
    https://arxiv.org/pdf/1801.09414.pdf

    Also referenced cvqluu's implementation of Angular Penalty Loss:
    https://paperswithcode.com/paper/cosface-large-margin-cosine-loss-for-deep
    """

    def __init__(self, in_features=5, out_features=2, s=64.0, m=0.35):
        super(LargeMarginCosineLoss, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.s = s
        self.m = m

        # cos(θ_j,i) = W_j^T * x_i
        self.linear = Linear(in_features, out_features, bias=False)

    def forward(self, x, targets):

        # normalize
        x = F.normalize(x, p=2, dim=1)
        for W in self.linear.parameters():
            W = F.normalize(W, p=2, dim=1)

        cos_θ = self.linear(x)
        s_cos_θ = self.s*torch.diagonal(cos_θ.transpose(0,1)[targets]-self.m)
#         print(s_cos_θ)
        try:
            cos_θj = [torch.cat((cos_θ[j,:y], cos_θ[j,(y+1):])).unsqueeze(0) for j, y in zip(len(targets), targets)] # <<<-- issue
            sum_j = torch.sum(torch.exp(self.s*torch.cat(cos_θj, dim=0)), dim=1)
        except:
            raise ValueError(cos_θ)

        result = torch.mean(torch.log(torch.exp(s_cos_θ) + sum_j) - torch.log(torch.exp(s_cos_θ)))

        return result

class ProtoNet(MetaTemplate):
    def __init__(self, model_func,  n_way, n_support):
        super(ProtoNet, self).__init__( model_func,  n_way, n_support)
#         self.loss_fn  = nn.CrossEntropyLoss()
#         self.loss_fn = LargeMarginCosineLoss()
        self.loss_fn = CrossEntropyLoss()

    def set_forward(self,x,is_feature = False):
        z_support, z_query  = self.parse_feature(x,is_feature)

        z_support   = z_support.contiguous()
        z_proto     = z_support.view(self.n_way, self.n_support, -1 ).mean(1) #the shape of z is [n_data, n_dim]
        z_query     = z_query.contiguous().view(self.n_way* self.n_query, -1 )

        dists = euclidean_dist(z_query, z_proto)
        scores = -dists

        return scores

    def set_forward_loss(self, x):
        y_query = torch.from_numpy(np.repeat(range( self.n_way ), self.n_query ))
        y_query = Variable(y_query.cuda())

        scores = self.set_forward(x)
#         print('scores', scores.size())
        loss = self.loss_fn(scores, y_query)
        if isinstance(loss, float):
            pass
#             print('loss',loss)
        else:
            pass
#             print('> loss', loss)
        return loss

def euclidean_dist( x, y):
    # x: N x D
    # y: M x D
    n = x.size(0)
    m = y.size(0)
    d = x.size(1)
    assert d == y.size(1)

    x = x.unsqueeze(1).expand(n, m, d)
    y = y.unsqueeze(0).expand(n, m, d)

    return torch.pow(x - y, 2).sum(2)


#!/usr/bin/env python
import bz2
from contextlib import contextmanager
import datetime
from deepdiff import DeepDiff
from distutils.version import LooseVersion as Version
import fire
import pathlib
from pprint import pprint
import questionary
import re
import secrets
import string
import subprocess
import sys
import time
from typing import Optional, Sequence
import yaml

dependencies = {
    "python": (
        # Command to get version
        ["python", "--version"],
        # Extract *only* the version number
        lambda v: v.split()[1],
        # It must be >= 3.7
        "3.7",
    ),
    "docker": (
        # Command to get version
        ["docker", "--version"],
        # Extract *only* the version number
        lambda v: v.split()[2][:-1],
        # It must be >= 18.06
        "18.06",
    ),
    "docker-compose": (
        # Command to get version
        ["docker-compose", "--version"],
        # Extract *only* the version number
        lambda v: re.search(r"\s*([\d.]+)", v).group(0).strip(),
        # It must be >= 1.22.0
        "1.22.0",
    ),
}

@contextmanager
def status(message):
    """
    Borrowed from https://github.com/cesium-ml/baselayer/

    :param message: message to print
    :return:
    """
    print(f"[·] {message}", end="")
    sys.stdout.flush()
    try:
        yield
    except Exception:
        print(f"\r[✗] {message}")
        raise
    else:
        print(f"\r[✓] {message}")

def deps_ok() -> bool:
    """
    Check system dependencies

    Borrowed from https://github.com/cesium-ml/baselayer/
    :return:
    """
    print("Checking system dependencies:")

    fail = []

    for dep, (cmd, get_version, min_version) in dependencies.items():
        try:
            query = f"{dep} >= {min_version}"
            with status(query):
                p = subprocess.Popen(
                    cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
                )
                out, err = p.communicate()
                try:
                    version = get_version(out.decode("utf-8").strip())
                    print(f"[{version.rjust(8)}]".rjust(40 - len(query)), end="")
                except Exception:
                    raise ValueError("Could not parse version")

                if not (Version(version) >= Version(min_version)):
                    raise RuntimeError(f"Required {min_version}, found {version}")
        except Exception as e:
            fail.append((dep, e))

    if fail:
        print()
        print("[!] Some system dependencies seem to be unsatisfied")
        print()
        print("    The failed checks were:")
        print()
        for (pkg, exc) in fail:
            cmd, get_version, min_version = dependencies[pkg]
            print(f'    - {pkg}: `{" ".join(cmd)}`')
            print("     ", exc)
        print()
        print(
            "    Please refer to https://github.com/dmitryduev/tails "
            "for installation instructions."
        )
        print()
        return False

    print("-" * 20)
    return True

def check_configs(
    config_wildcards: Sequence = ("config.*yaml", "docker-compose.*yaml")
):
    """
    - Check if config files exist
    - Offer to use the config files that match the wildcards
    - For config.yaml, check its contents against the defaults to make sure nothing is missing/wrong

    :param config_wildcards:
    :return:
    """
    path = pathlib.Path(__file__).parent.absolute()

    for config_wildcard in config_wildcards:
        config = config_wildcard.replace("*", "")
        # use config defaults if configs do not exist?
        if not (path / config).exists():
            answer = questionary.select(
                f"{config} does not exist, do you want to use one of the following"
                " (not recommended without inspection)?",
                choices=[p.name for p in path.glob(config_wildcard)],
            ).ask()
            subprocess.run(["cp", f"{path / answer}", f"{path / config}"])

        # check contents of config.yaml WRT config.defaults.yaml
        if config == "config.yaml":
            with open(path / config.replace(".yaml", ".defaults.yaml")) as config_yaml:
                config_defaults = yaml.load(config_yaml, Loader=yaml.FullLoader)
            with open(path / config) as config_yaml:
                config_wildcard = yaml.load(config_yaml, Loader=yaml.FullLoader)
            deep_diff = DeepDiff(config_wildcard, config_defaults, ignore_order=True)
            difference = {
                k: v
                for k, v in deep_diff.items()
                if k in ("dictionary_item_added", "dictionary_item_removed")
            }
            if len(difference) > 0:
                print("config.yaml structure differs from config.defaults.yaml")
                pprint(difference)
                raise KeyError("Fix config.yaml before proceeding")

def get_git_hash_date():
    """Get git date and hash

    Borrowed from SkyPortal https://skyportal.io

    :return:
    """
    hash_date = dict()
    try:
        p = subprocess.Popen(
            ["git", "log", "-1", '--format="%h %aI"'],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            cwd=pathlib.Path(__file__).parent.absolute(),
        )
    except FileNotFoundError:
        pass
    else:
        out, err = p.communicate()
        if p.returncode == 0:
            git_hash, git_date = (
                out.decode("utf-8")
                .strip()
                .replace('"', "")
                .split("T")[0]
                .replace("-", "")
                .split()
            )
            hash_date["hash"] = git_hash
            hash_date["date"] = git_date

    return hash_date

class Kowalski:
    def __init__(self, yes=False):
        """

        :param yes: answer yes to all possible requests?
        """
        self.yes = yes

    @staticmethod
    def check_containers_up(
        containers: Sequence,
        num_retries: int = 10,
        sleep_for_seconds: int = 10,
    ):
        """Check if containers in question are up and running

        :param containers: container name sequence, e.g. ("kowalski_api_1", "kowalski_mongo_1")
        :param num_retries:
        :param sleep_for_seconds: number of seconds to sleep for before retrying
        :return:
        """
        for i in range(num_retries):
            if i == num_retries - 1:
                raise RuntimeError(f"{containers} containers failed to spin up")

            command = ["docker", "ps", "-a"]
            container_list = (
                subprocess.check_output(command, universal_newlines=True)
                .strip()
                .split("\n")
            )
            print(container_list)
            if len(container_list) == 1:
                print("No containers are running, waiting...")
                time.sleep(sleep_for_seconds)
                continue

            containers_up = (
                len(
                    [
                        container
                        for container in container_list
                        if (
                            (container_name in container)
                            and (" Up " in container)
                            and ("unhealthy" not in container)
                            and ("health: starting" not in container)
                        )
                    ]
                )
                > 0
                for container_name in containers
            )

            if not all(containers_up):
                print(f"{containers} containers are not up, waiting...")
                time.sleep(sleep_for_seconds)
                continue

            break

    @staticmethod
    def check_keyfile():
        """Check if MongoDB keyfile for replica set authorization exists; generate one if not"""
        mongodb_keyfile = pathlib.Path(__file__).parent.absolute() / "mongo_key.yaml"
        if not mongodb_keyfile.exists():
            print("Generating MongoDB keyfile")
            # generate a random key that is required to be able to use authorization with replica set
            key = "".join(
                secrets.choice(string.ascii_lowercase + string.digits)
                for _ in range(32)
            )
            with open(mongodb_keyfile, "w") as f:
                f.write(key)
            command = ["chmod", "400", "mongo_key.yaml"]
            subprocess.run(command)

    @classmethod
    def up(cls, build: bool = False):
        """
        🐧🚀 Launch Kowalski

        :param build: build the containers first?
        :return:
        """
        print("Spinning up Kowalski 🐧🚀")

        config_wildcards = ["config.*yaml", "docker-compose.*yaml"]

        # check configuration
        with status("Checking configuration"):
            check_configs(config_wildcards=config_wildcards)

        cls.check_keyfile()

        if build:
            cls.build()

        command = ["docker-compose", "-f", "docker-compose.yaml", "up", "-d"]

        # start up Kowalski
        print("Starting up")
        subprocess.run(command)

    @staticmethod
    def down():
        """
        ✋ Shut down Kowalski

        :return:
        """
        print("Shutting down Kowalski")
        command = ["docker-compose", "-f", "docker-compose.yaml", "down"]

        subprocess.run(command)

    @classmethod
    def build(cls):
        """
        Build Kowalski's containers

        :return:
        """
        print("Building Kowalski")

        config_wildcards = ["config.*yaml", "docker-compose.*yaml"]

        # always use docker-compose.yaml
        command = ["docker-compose", "-f", "docker-compose.yaml", "build"]

        # check configuration
        with status("Checking configuration"):
            check_configs(config_wildcards=config_wildcards)

        # load config
        with open(
            pathlib.Path(__file__).parent.absolute() / "config.yaml"
        ) as config_yaml:
            config = yaml.load(config_yaml, Loader=yaml.FullLoader)["kowalski"]

        # get git version:
        git_hash_date = get_git_hash_date()
        version = (
            f"v{config['server']['version']}"
            f"+git{git_hash_date.get('date', datetime.datetime.utcnow().strftime('%Y%m%d'))}"
            f".{git_hash_date.get('hash', 'unknown')}"
        )
        with open(
            pathlib.Path(__file__).parent.absolute() / "version.txt", "w"
        ) as version_file:
            version_file.write(f"{version}\n")

        # check MongoDB keyfile
        cls.check_keyfile()

        subprocess.run(command)

    @staticmethod
    def seed(source: str = "./", drop: Optional[bool] = False):
        """
        Ingest catalog dumps into Kowalski

        :param source: where to look for the dumps;
                       can be a local path or a Google Cloud Storage bucket address, e.g. gs://kowalski-catalogs
        :param drop: drop existing collections with same names before ingesting?
        :return:
        """
        print("Ingesting catalog dumps into a running Kowalski instance")

        # check configuration
        with status("Checking configuration"):
            check_configs(config_wildcards=["config.*yaml"])

        with open(
            pathlib.Path(__file__).parent.absolute() / "config.yaml"
        ) as config_yaml:
            config = yaml.load(config_yaml, Loader=yaml.FullLoader)["kowalski"]

        command = [
            "docker",
            "exec",
            "-i",
            "kowalski_mongo_1",
            "mongorestore",
            f"-u={config['database']['admin_username']}",
            f"-p={config['database']['admin_password']}",
            "--authenticationDatabase=admin",
            "--archive",
        ]

        if drop:
            command.append("--drop")

        if "gs://" not in source:
            # ingesting from a local path
            path = pathlib.Path(source).absolute()

            dumps = [p.name for p in path.glob("*.dump")]

            if len(dumps) == 0:
                print(f"No catalog dumps found under {path}")
                return False

            answer = questionary.checkbox(
                "Found the following collection dumps. Which ones would you like to ingest?",
                choices=dumps,
            ).ask()

            for dump in answer:
                with open(f"{path / dump}") as f:
                    subprocess.call(command, stdin=f)

        else:
            # ingesting from Google Cloud
            path_tmp = pathlib.Path(__file__).parent / ".catalog_dumps"
            if not path_tmp.exists():
                path_tmp.mkdir(parents=True, exist_ok=True)

            ls_command = ["gsutil", "ls", source]
            catalog_list = (
                subprocess.check_output(ls_command, universal_newlines=True)
                .strip()
                .split("\n")
            )
            dumps = [dump for dump in catalog_list if "dump" in dump]

            answer = questionary.checkbox(
                "Found the following collection dumps. Which ones would you like to ingest?",
                choices=dumps,
            ).ask()

            for dump in answer:
                cp_command = [
                    "gsutil",
                    "-m",
                    "cp",
                    "-n",
                    dump,
                    str(path_tmp),
                ]
                p = subprocess.run(cp_command, check=True)
                if p.returncode != 0:
                    raise RuntimeError(f"Failed to fetch {dump}")

                path_dump = f"{path_tmp / pathlib.Path(dump).name}"
                if dump.endswith(".bz2"):
                    with bz2.BZ2File(path_dump) as f:
                        subprocess.call(command, stdin=f)
                elif dump.endswith(".gz"):
                    with open(path_dump) as f:
                        subprocess.call(command + ["--gzip"], stdin=f)
                else:
                    with open(path_dump) as f:
                        subprocess.call(command, stdin=f)

                rm_fetched = questionary.confirm(f"Remove {path_dump}?").ask()
                if rm_fetched:
                    pathlib.Path(path_dump).unlink()

    @classmethod
    def test(cls):
        """
        Run the test suite

        :return:
        """
        print("Running the test suite")

        # make sure the containers are up and running
        cls.check_containers_up(
            containers=("kowalski_ingester_1", "kowalski_api_1", "kowalski_mongo_1"),
            sleep_for_seconds=10,
        )

        test_setups = [
            {
                "part": "PGIR alert broker components",
                "container": "kowalski_ingester_1",
                "test_script": "test_alert_broker_pgir.py",
                "flaky": False,
            },
            {
                "part": "ZTF alert broker components",
                "container": "kowalski_ingester_1",
                "test_script": "test_alert_broker_ztf.py",
                "flaky": False,
            },
            {
                "part": "PGIR alert ingestion",
                "container": "kowalski_ingester_1",
                "test_script": "test_ingester_pgir.py",
                "flaky": False,
            },
            {
                "part": "ZTF alert ingestion",
                "container": "kowalski_ingester_1",
                "test_script": "test_ingester.py",
                "flaky": False,
            },
            {
                "part": "API",
                "container": "kowalski_api_1",
                "test_script": "test_api.py",
                "flaky": False,
            },
            {
                "part": "TNS monitoring",
                "container": "kowalski_ingester_1",
                "test_script": "test_tns_watcher.py",
                "flaky": True,
            },
            {
                "part": "Tools",
                "container": "kowalski_ingester_1",
                "test_script": "test_tools.py",
                "flaky": False,
            },
        ]

        failed_tests = []

        for setup in test_setups:
            print(f"Testing {setup['part']}")
            command = [
                "docker",
                "exec",
                "-i",
                setup["container"],
                "python",
                "-m",
                "pytest",
                "-s",
                setup["test_script"],
            ]
            try:
                subprocess.run(command, check=True)
            except subprocess.CalledProcessError:
                if not setup.get("flaky", False):
                    failed_tests.append(setup["part"])
                else:
                    print(f"{setup['part']} test, marked as flaky, failed.")
                continue

        if failed_tests:
            print(f"Failed tests: {failed_tests}")
            sys.exit(1)

    @staticmethod
    def develop():
        """
        Install developer tools
        """
        subprocess.run(["pip", "install", "-U", "pre-commit"], check=True)
        subprocess.run(["pre-commit", "install"], check=True)

    @classmethod
    def lint(cls):
        """
        Lint the full code base

        :return:
        """
        try:
            import pre_commit  # noqa: F401
        except ImportError:
            cls.develop()

        try:
            subprocess.run(["pre-commit", "run", "--all-files"], check=True)
        except subprocess.CalledProcessError:
            sys.exit(1)

if __name__ == "__main__":
    # check environment
    env_ok = deps_ok()
    if not env_ok:
        raise RuntimeError("Halting because of unsatisfied system dependencies")

    fire.Fire(Kowalski)


import os
import sys
import shutil
import subprocess
from flask import current_app as app
from threading import Thread
from . import socketio
from src.exceptions import InvalidSongSegmentException, ChordParseException

'''
从完整的歌曲对象分离出歌曲存储目录。

参数:
    d_song: 数据库中的歌曲所映射的字典。
返回:
    该歌曲的存储目录。
'''
def generate_song_directory(d_song):
    song_name = d_song['song_name']
    singer = d_song['singer']
    folder_name = song_name + '-' + singer
    return os.path.join(app.config['FILE_UPLOAD_DIR'], folder_name)

'''
由歌曲原唱生成歌曲打分文件。

参数:
    real_app: Flask的app对象，需要手动在调用函数中传入。
    original_file_path: 歌曲原唱文件路径。
    song_info: 格式为“<歌曲名称>-<歌手>”的歌曲信息，用于向客户端发送socket。
'''
def rate_by_original(real_app, original_file_path, song_info):

    with real_app.app_context():

        vocal_file_path = generate_vocal_file(original_file_path)
        single_track_file_path = generate_single_track_file(vocal_file_path)
        generate_rate_file(single_track_file_path)

        if song_info:
            socketio.emit('rate', song_info, namespace='/karaoke')

'''
解释和弦文件。

参数:
    real_app: Flask的app对象，需要手动在调用函数中传入。
    org_chord_path: 原始和弦文件路径。
    song_info: 格式为“<歌曲名称>-<歌手>”的歌曲信息，用于向客户端发送socket。
'''
def trans_chord(real_app, org_chord_path, song_info):

    with real_app.app_context():

        directory, _ = os.path.split(org_chord_path)
        chord_trans_path = os.path.join(directory, app.config['CHORD_TRANS_FILENAME'])

        shell_args = [app.config['CHORD_TRANSLATOR_PATH'], org_chord_path, chord_trans_path]
        process = subprocess.Popen(shell_args, cwd=app.config['CHORD_TRANS_WORKING_DIR'],
                                   shell=app.config['REQUIRE_SHELL'])
        process.wait()

        # 删除原始和弦文件
        try:
            os.remove(org_chord_path)
        except OSError:
            pass

        if song_info:
            socketio.emit('chord', song_info, namespace='/karaoke')

'''
生成片段歌词文件、音轨文件等自弹自唱相关文件。

参数:
    real_app: Flask的app对象，需要手动在调用函数中传入。
    chord_path: 解释完毕后的和弦文件路径。
    lyric_path: 完整歌词文件路径。
    original_path: 歌曲原唱文件路径。
    song_info: 格式为“<歌曲名称>-<歌手>”的歌曲信息，用于向客户端发送socket。
'''
def generate_instrument_sing_files(real_app, chord_path, lyric_path, original_path, song_info):

    start_time, end_time = read_chord(chord_path)

    lrc_thread = Thread(target=trim_lrc, args=(real_app, lyric_path, start_time, end_time,))
    track_thread = Thread(target=separate_audio_track, args=(real_app, original_path, start_time, end_time,))

    lrc_thread.start()
    track_thread.start()

    lrc_thread.join()
    track_thread.join()

    if song_info:
        socketio.emit('instrument', song_info, namespace='/karaoke')

'''
对歌曲特定片段进行音轨分离。

参数:
    real_app: Flask的app对象，需要手动在调用函数中传入。
    original_path: 歌曲原唱文件路径。
    start_time: 歌曲特定片段的开始时间。
    end_time: 歌曲特定片段的结束时间。
'''
def separate_audio_track(real_app, original_path, start_time, end_time):

    with real_app.app_context():

        start_time = start_time - app.config['BUTTON_ANI_SEC']
        if start_time  < 0:
            raise InvalidSongSegmentException

        trimmed_wav_path = trim_wav(original_path, start_time, end_time)
        generate_inst_wav(trimmed_wav_path)

        # 删除切分出的歌曲特定片段文件
        try:
            os.remove(trimmed_wav_path)
        except OSError:
            pass

'''
从歌曲原唱中分离出人声。

参数:
    original_file_path: 歌曲原唱文件路径。

返回:
    分离出的无伴奏人声文件路径。
'''
def generate_vocal_file(original_file_path):

    directory, original_filename = os.path.split(original_file_path)
    filename_without_suffix = original_filename.rsplit('.', 1)[0]

    shell_args = ['spleeter', 'separate', '-p', 'spleeter:2stems', '-o',
                  directory, original_file_path]
    if sys.platform == 'win32':
        shell_args.insert(0, 'python')
        shell_args.insert(1, '-m')

    process = subprocess.Popen(shell_args, cwd=app.config['WORKING_DIR'],
                               shell=app.config['REQUIRE_SHELL'])
    process.wait()

    return os.path.join(directory, filename_without_suffix, 'vocals.wav')

'''
将无伴奏人声文件转换为单音轨。

参数:
    vocal_file_path: 人声文件路径。

返回:
    单音轨无伴奏人声文件路径。
'''
def generate_single_track_file(vocal_file_path):

    directory, _ = os.path.split(vocal_file_path)
    single_track_file_path = os.path.join(directory, 'vocal_single.wav')

    shell_args = ['ffmpeg', '-y', '-i', vocal_file_path,
                    '-ar', '44100', '-ac', '1', single_track_file_path]
    process = subprocess.Popen(shell_args, cwd=app.config['WORKING_DIR'],
                               shell=app.config['REQUIRE_SHELL'])
    process.wait()

    return single_track_file_path

'''
由单音轨无伴奏人声文件生成歌曲打分文件。

参数:
    single_track_file_path: 单音轨无伴奏人声文件路径。
'''
def generate_rate_file(single_track_file_path):

    rm_dir, _ = os.path.split(single_track_file_path)

    if sys.platform == 'win32':     # Windows
        directory = rm_dir.rsplit('\\', 1)[0]
    else:                           # Linux 或 MacOS
        directory = rm_dir.rsplit('/', 1)[0]

    rate_file_path = os.path.join(directory, app.config['RATE_FILENAME'])
    shell_args = [app.config['RATING_PATH'], single_track_file_path,
                  '-t', '-s', '50', '-o', rate_file_path]
    process = subprocess.Popen(shell_args, cwd=app.config['WORKING_DIR'],
                               shell=app.config['REQUIRE_SHELL'])
    process.wait()

    try:
        shutil.rmtree(rm_dir)
    except OSError:
        pass

'''
从完整的歌词文件中提取出歌曲特定片段的歌词。

参数:
    real_app: Flask的app对象，需要手动在调用函数中传入。
    org_lrc_path: 完整歌词文件路径。
    start_time: 歌曲特定片段的开始时间。
    end_time: 歌曲特定片段的结束时间。
'''
def trim_lrc(real_app, org_lrc_path, start_time, end_time):

    with real_app.app_context():

        directory, _ = os.path.split(org_lrc_path)
        new_lrc_path = os.path.join(directory, app.config['LYRIC_INSTRUMENT_FILENAME'])

        shell_args = [app.config['TRIMMER_PATH'], str(start_time), str(end_time), org_lrc_path, new_lrc_path]
        process = subprocess.Popen(shell_args, cwd=app.config['WORKING_DIR'],
                                   shell=app.config['REQUIRE_SHELL'])
        process.wait()

'''
从解释完毕的和弦文件中获取歌曲片段的开始时间与结束时间。

参数:
    chord_trans_path: 解释完毕的和弦文件路径。

返回:
    歌曲特定片段的开始时间与结束时间。
'''
def read_chord(chord_trans_path):

    with open(chord_trans_path, 'r') as chord:
        line = chord.readline()
        args = line.split()

        if len(args) != 5:
            raise ChordParseException

        start_time = float(args[3]) / 1000
        end_time = float(args[4]) / 1000
        return start_time, end_time

'''
切分特定片段的歌曲原唱。

参数:
    org_wav_path: 歌曲原唱文件路径。
    start_time: 歌曲特定片段的开始时间。
    end_time: 歌曲特定片段的结束时间。

返回:
    切分出的歌曲原唱片段文件路径。
'''
def trim_wav(org_wav_path, start_time, end_time):

    directory, _ = os.path.split(org_wav_path)
    trimmed_wav_path = os.path.join(directory, app.config['TRIMMED_WAV_FILENAME'])

    duration = end_time - start_time
    shell_args = ['ffmpeg', '-ss', str(start_time), '-t', str(duration),
                  '-i', org_wav_path, trimmed_wav_path]

    process = subprocess.Popen(shell_args, shell=app.config['REQUIRE_SHELL'])
    process.wait()

    return trimmed_wav_path

'''
对歌曲原唱片段文件进行音轨分离，产生自弹自唱所需文件。

参数:
    trimmed_wav_path: 歌曲原唱片段文件路径。
'''
def generate_inst_wav(trimmed_wav_path):

    directory, _ = os.path.split(trimmed_wav_path)

    shell_args = ['spleeter', 'separate', '-p', 'spleeter:5stems',
                  '-o', directory, trimmed_wav_path]
    if sys.platform == 'win32':
        shell_args.insert(0, 'python')
        shell_args.insert(1, '-m')

    process = subprocess.Popen(shell_args, cwd=app.config['WORKING_DIR'],
                               shell=app.config['REQUIRE_SHELL'])
    process.wait()


import json
from django.db import transaction
from django.db.models import F
from django.db.models.signals import pre_save, post_save, post_delete
from django.dispatch import receiver
from django.utils import timezone
from django_celery_beat.models import PeriodicTask, IntervalSchedule

from .models import Book, BookChapter
from .tasks import book_scraper_info, book_scraper_chaps

@receiver(post_save, sender=Book)
def book_scraper_initial_signal(sender, instance, created=False, **kwargs):
    if not instance.visited and instance.visit_id:
        schedule, created = IntervalSchedule.objects.get_or_create(
            every=15,
            period=IntervalSchedule.SECONDS,
        )
        PeriodicTask.objects.create(
            one_off=True,
            interval=schedule,
            name=f'Update book: {instance.title}',
            task='novel2read.apps.books.tasks.book_scraper_info',
            args=json.dumps([instance.pk]),
        )

        if not instance.chapters_count:
            schedule, created = IntervalSchedule.objects.get_or_create(
                every=50,
                period=IntervalSchedule.SECONDS,
            )
            PeriodicTask.objects.create(
                one_off=True,
                interval=schedule,
                name=f'Update book chapters init: {instance.title}',
                task='novel2read.apps.books.tasks.book_scraper_chaps',
                args=json.dumps([instance.pk]),
            )

@receiver(post_save, sender=BookChapter)
def create_update_chapter_cid(sender, instance, created=False, **kwargs):
    if created:
        instance.book.update_chapters_count()
        instance.c_id = instance.book.chapters_count
        instance.save(update_fields=['c_id'])

@receiver(post_delete, sender=BookChapter)
def delete_update_chapter_cid(sender, instance, **kwargs):
    instance.book.update_chapters_count()
    c_id_del = instance.c_id
    book_chaps = BookChapter.objects.filter(book__slug=instance.book.slug).filter(c_id__gt=c_id_del)
    book_chaps.update(c_id=F('c_id') - 1)


from django.contrib import admin
from django.urls import path, include
from pastes import views as paste_views

urlpatterns = [
    path('', include('home.urls', 'home')),
    path('admin/', admin.site.urls),
    path('pastes/', paste_views.ShowQueryResults.as_view(), name="query_results"),
    path('pastes/<str:char_id>', paste_views.ShowPaste.as_view(), name="show_paste"),
    path('pastes/<str:char_id>/confirm_delete', paste_views.ConfirmDelete.as_view(), name="confirm_delete"),
    path('pastes/<str:char_id>/delete', paste_views.DeletePaste.as_view(), name='delete_paste'),
]


# -*- coding: utf-8 -*-
from zabbix_api import ZabbixAPI
from datetime import datetime
import time

zapi = ZabbixAPI("https://endereco.com.br/zabbix")
# Disable SSL certificate verification
zapi.session.verify = False
# Specify a timeout (in seconds)
zapi.timeout = 10.1
# informacoes de acesso
zapi.login("usuario.zabbix", "senhaUsuario")
# mostra versao do zabbix
print("Connected to Zabbix API Version %s" % zapi.api_version())

for hosts in zapi.host.get({'output': ['host','interface']}):
    print(hosts)

#item_id = 1879051

# Create a time range
#time_till = time.mktime(datetime.now().timetuple())
#time_from = time_till - 60 * 60 * 4  # 4 hours

#time_from = time.mktime(datetime.now().timetuple()) - 60 * 5  # 5 min

#historico = zapi.history.get({ 'itemids': [ item_id ],
#'history': 0,
#'output': 'extend',
#'time_from': time_from,
#'time_till': “1439250959”
#    })

# Print out each datapoint
#for point in historico:
#    print("{0}: {1}".format(datetime.fromtimestamp(int(point['clock']))
#                            .strftime("%x %X"), point['value']))

# 0 pertence a data
# 1 pertence a valor


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

__author__ = 'ipetrash'

import time
from PyQt5.QtWidgets import QApplication, QMessageBox

from print__hprof_or_big_size_file import find_files_by_dirs, DIRS

if __name__ == '__main__':
    app = QApplication([])

    while True:
        result = find_files_by_dirs(DIRS)
        if result:
            QMessageBox.warning(None, 'Warn', '\n'.join(result))

        time.sleep(5 * 60 * 60)


import argparse
import json
import os
import time

import requests

import dataclasses
import logging
from typing import *

L = logging.getLogger(__name__)

class RequestFailed(RuntimeError):
    def __init__(self, reason: str):
        super("request failed: " + reason)

class Client:
    def __init__(self, token):
        self.baseurl = "https://slack.com/api"
        self.default_header = {"content-type": "application/x-www-form-urlencoded"}
        self.token = token

    def _get(self, url, params) -> requests.Response:
        headers = self.default_header
        params["token"] = self.token
        res = requests.get(url, headers=headers, params=params)
        return self._decode_response(res)

    def _post(self, url, data) -> Any:
        headers = self.default_header
        data["token"] = self.token
        res = requests.post(url, headers=headers, data=data)
        return self._decode_response(res)

    def _decode_response(self, res: requests.Response) -> Any:
        if res.status_code != 200:
            raise RequestFailed(f"status_code isn't 200 ({res.status_code})")
        return res.json()

    def auth_test(self):
        return self._get(self.baseurl + "/auth.test", {})

    def conversations_list(self, cursor: str = None):
        params = {"types": "public_channel,private_channel,mpim"}
        if cursor is not None:
            params["cursor"] = cursor
        return self._get(self.baseurl + "/conversations.list", params)

    def conversations_members(self, channel: str, cursor: str = None):
        params = {"channel": channel}
        if cursor is not None:
            params["cursor"] = cursor
        return self._get(self.baseurl + "/conversations.members", params)

    def conversations_history(self, channel: str, cursor: str = None):
        params = {"channel": channel}
        if cursor is not None:
            params["cursor"] = cursor
        return self._get(self.baseurl + "/conversations.history", params)

    def conversations_replies(self, channel: str, ts: str, cursor: str = None):
        params = {"channel": channel, "ts": ts}
        if cursor is not None:
            params["cursor"] = cursor
        return self._get(self.baseurl + "/conversations.replies", params)

    def conversations_join(self, channel: str):
        params = {"channel": channel}
        return self._post(self.baseurl + "/conversations.join", params)

    def users_list(self, cursor: str = None):
        params = {}
        if cursor is not None:
            params["cursor"] = cursor
        return self._get(self.baseurl + "/users.list", params)

    def users_profile_set(self, user: str, key: str, value: str):
        params = {"user": user, "name": key, "value": value}
        return self._post(self.baseurl + "/users.profile.set", params)

def is_thread_parent(msg: Any) -> bool:
    return "thread_ts" in msg and msg["ts"] == msg["thread_ts"]

def get_channels(cli: Client) -> List[Any]:
    L.info("fetching channel metadata...")
    channels: List[Any] = []
    next_cursor = None
    while next_cursor != "":
        data = cli.conversations_list(next_cursor)
        if not data["ok"]:
            raise RuntimeError(f"request failed: (data={data})")
        channels += data["channels"]
        next_cursor = data["response_metadata"]["next_cursor"]

    L.info("fetching channel members...")
    for c in channels:
        L.info(f"fetching channel members for channel {c['name']}...")

        c["members"] = []
        if c["is_archived"]:
            L.info(f"channel {c['name']} is archived, skipped")
            continue

        next_cursor = None
        try:
            while next_cursor != "":
                data = cli.conversations_members(c["id"])
                if not data["ok"]:
                    raise RuntimeError(f"request failed: (channel={c}, data={data})")
                c["members"] += data["members"]
                next_cursor = data["response_metadata"]["next_cursor"]
        except Exception as e:
            pass

    return channels

def get_users(cli: Client) -> List[Any]:
    L.info("fetching user metadata...")
    users: List[Any] = []
    next_cursor = None
    while next_cursor != "":
        data = cli.users_list(next_cursor)
        if not data["ok"]:
            raise RuntimeError(f"request failed: (data={data})")
        users += data["members"]
        next_cursor = data["response_metadata"]["next_cursor"]

    return users

def get_replies(cli: Client, channel: Any, ts: str) -> List[Any]:
    messages: List[Any] = []
    next_cursor = None
    while next_cursor != "":
        data = cli.conversations_replies(channel["id"], ts, next_cursor)
        if not data["ok"]:
            raise RuntimeError(f"request failed: (data={data})")
        messages += data["messages"]
        next_cursor = data["response_metadata"]["next_cursor"] if data["has_more"] else ""

    return messages

def get_messages(cli: Client, channel: Any) -> List[Any]:
    L.info(f"fetching messages for channel {channel['name']}...")
    messages: List[Any] = []
    next_cursor = None
    while next_cursor != "":
        data = cli.conversations_history(channel["id"], next_cursor)
        if not data["ok"]:
            raise RuntimeError(f"request failed: (data={data})")
        messages += data["messages"]
        next_cursor = data["response_metadata"]["next_cursor"] if data["has_more"] else ""

    thread_broadcast_set = set()
    for msg in messages:
        if "subtype" in msg and msg["subtype"] == "thread_broadcast":
            thread_broadcast_set.add(msg["ts"])

    for msg in messages:
        if is_thread_parent(msg):
            replies = get_replies(cli, channel, msg["thread_ts"])
            msg["replies"] = []
            for reply in replies:
                if msg["ts"] == reply["ts"]:
                    continue
                msg["replies"].append({"user": reply["user"], "ts": reply["ts"]})
                if not reply["ts"] in thread_broadcast_set:
                    messages.append(reply)

    return messages

def append_download_token(msg: Any, download_token: str):
    if not "files" in msg:
        return

    for f in msg["files"]:
        if f["mimetype"].startswith("image"):
            for s in [64, 80, 360, 480, 160, 720, 800, 960, 1024]:
                try:
                    f[f"thumb_{s}"] += f"?t={download_token}"
                except Exception as e:
                    L.debug("exception occured in append_download_token, ignored...")

def output(dest: str, channels: List[Any], users: List[Any], messages: Dict[str, List[Any]], download_token: Optional[str] = None):
    os.makedirs(dest, exist_ok=True)

    with open(f"{dest}/channels.json", "w") as f:
        f.write(json.dumps(channels))

    with open(f"{dest}/users.json", "w") as f:
        f.write(json.dumps(users))

    for channel in channels:
        channel_dir = f"{dest}/{channel['name']}"
        os.makedirs(channel_dir, exist_ok=True)

        if not channel["name"] in messages:
            continue

        msgs = {}
        for msg in messages[channel["name"]]:
            if download_token is not None:
                append_download_token(msg, download_token)

            t = time.gmtime(float(msg["ts"]))
            key = f"{t.tm_year:04}-{t.tm_mon:02}-{t.tm_mday:02}"
            if not key in msgs:
                msgs[key] = []
            msgs[key].append(msg)

        for key in msgs.keys():
            msgs[key] = sorted(msgs[key], key=lambda m: float(m["ts"]))
            with open(f"{channel_dir}/{key}.json", "w") as f:
                f.write(json.dumps(msgs[key]))

def main(args: argparse.Namespace):
    logging.basicConfig(level=logging.INFO)

    cli = Client(args.bot_token)

    L.info("checking validity of token...")
    user = cli.auth_test()
    if not user["ok"]:
        raise RuntimeError("token isn't valid")

    L.info("fetching channels...")
    channels = get_channels(cli)

    L.info("fetching users...")
    users = get_users(cli)

    L.info("fetching messages...")
    messages: Dict[str, List[Any]] = {}
    for channel in channels:
        if user["user_id"] in channel["members"]:
            messages[channel["name"]] = get_messages(cli, channel)

    output(args.destination, channels, users, messages, args.download_token)

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

    parser.add_argument("--bot-token", help="token for accesssing Slack")
    parser.add_argument("--download-token", help="token for fetching assets from Slack")
    parser.add_argument("--destination", help="the output directory")

    args: argparse.Namespace = parser.parse_args()

    main(args)


import json
import logging
import os

import requests
from django.core.management.base import BaseCommand
from django.utils.http import urlencode

from public_data_collector.models import Area, Sigungu, SmallArea

class Collector:
    def __init__(self):
        super(Collector, self).__init__()
        self.base_url = 'http://api.visitkorea.or.kr/openapi/service/rest/KorService'
        self.endpoint = self.base_url
        with open('/etc/secrets/culturedata_proto/service_key.txt') as f:
            self.service_key = f.read().strip()
        self.base_query_params = {
            'pageNo': 1,
            'MobileOS': 'ETC',
            'MobileApp': 'culterdata_proto',
        }

    def send_request(self):
        return requests.get(self.endpoint)

class Command(Collector, BaseCommand):
    help = 'Collect public data'

    def __init__(self):
        super(Command, self).__init__()

    def update_setting(self, query_params):
        self.query_params = self.base_query_params.copy()
        self.query_params.update(query_params)
        self.endpoint = self.base_url \
                        + '/areaCode?ServiceKey={}&{}'.format(self.service_key, urlencode(self.query_params))

    def handle(self, *args, **options):
        self.update_setting({
            'numOfRows': 1000,
            '_type': 'json',
        })
        response = self.send_request()
        areas = json.loads(response.text)['response']['body']['items']['item']
        for area in areas:
            area_instance = Area.objects.get_or_create(
                code=int(area['code']),
                name=area['name']
            )[0]
            self.update_setting({
                'numOfRows': 1000,
                '_type': 'json',
                'areaCode': area['code'],
            })
            response = self.send_request()
            sigungus = json.loads(response.text)['response']['body']['items']['item']
            if type(sigungus) is not list:
                sigungus = [sigungus]
            for sigungu in sigungus:
                sigungu_instance = Sigungu.objects.get_or_create(
                    area=area_instance,
                    code=int(sigungu['code']),
                    name=sigungu['name']
                )[0]
                self.update_setting({
                    'numOfRows': 1000,
                    '_type': 'json',
                    'areaCode': area['code'],
                    'sigunguCode': sigungu['code']
                })
                response = self.send_request()
                if json.loads(response.text)['response']['body']['totalCount'] == 0:
                    continue
                small_areas = json.loads(response.text)['response']['body']['items']['item']
                for small_area in small_areas:
                    SmallArea.objects.get_or_create(
                        sigungu=sigungu_instance,
                        code=int(small_area['code']),
                        name=small_area['name']
                    )

        return 'collect data complete.'


from django.shortcuts import (
    render, redirect, reverse, HttpResponse, get_object_or_404
)
from django.contrib.auth.decorators import login_required
from django.contrib import messages
from django.http import JsonResponse
from django.template.loader import render_to_string
from .forms import WishlistForm
from products.models import Product
from profiles.models import UserProfile
from .models import Wishlist

@login_required
def view_wishlist(request):
    """ A view that renders the wishlist contents page """

    # basic view for displaying User wishlist page

    user = UserProfile.objects.get(user=request.user)

    # wishlist = get_object_or_404(Wishlist, user=user)

    try:
        wishlist = Wishlist.objects.get(user=user)
    except Wishlist.DoesNotExist:
        wishlist = None
        messages.warning(request, "No items in the Wishlist yet!")
    if wishlist:
    # user exist
        pass
    else:
    # user does not exist
        pass

    context={
        'wishlist': wishlist,
    }

    return render(request, 'wishlist/wishlist.html', context)

@login_required
def add_to_wishlist(request, product_id):

    product_wish = get_object_or_404(Product, pk=product_id) # get product

    wishlist, created = Wishlist.objects.get_or_create(
        user=request.user.userprofile,
        name='rick'
    )

    if wishlist.products.filter(name=product_wish).exists():
        messages.warning(request, "Item already added to wishlist")
        return redirect(reverse('view_wishlist'))
    else:
        wishlist.products.add(product_wish)
        messages.success(request, "Item added to Wishlist!")
        return redirect(reverse('view_wishlist'))

@login_required
def remove_from_wishlist(request, product_id):

    product_wish = get_object_or_404(Product, pk=product_id) # get product
    user = UserProfile.objects.get(user=request.user) # get user

    wishlist = get_object_or_404(Wishlist, user=user) # filter wishlist item from with user

    wishlist.products.remove(product_wish) # remove item
    messages.success(request, 'Product Removed From Wishlist')

    return redirect(reverse('view_wishlist'))


import os

from torch.utils import data
from torchvision import transforms
from PIL import Image

def get_transform():
    """Module for image pre-processing definition.

    You can customize this module.
    """
    normalize = transforms.Normalize(
        mean=[x / 255.0 for x in [125.3, 123.0, 113.9]],
        std=[x / 255.0 for x in [63.0, 62.1, 66.7]])
    transform = [
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        normalize
    ]
    return transforms.Compose(transform)

class CustomDataset(data.Dataset):
    """Dataset class.

    This class is used for internal NSML inference system.
    You can modify this module for improving your data loading efficiency.
    """
    def __init__(self, root, transform):
        self.data_root = os.path.join(root, 'test_data')
        self.transform = transform
        self.image_ids = [img for img in os.listdir(self.data_root)]

    def __getitem__(self, idx):
        image_id = self.image_ids[idx]
        image = Image.open(os.path.join(self.data_root, image_id))
        image = image.convert('RGB')
        image = self.transform(image)
        return image, image_id

    def __len__(self):
        return len(self.image_ids)

def data_loader(root, batch_size=64):
    """Test data loading module.

    Args:
        root: string. dataset path.
        batch_size: int.

    Returns:
        DataLoader instance
    """
    input_transform = get_transform()
    dataset = CustomDataset(root, input_transform)
    return data.DataLoader(dataset=dataset,
                           batch_size=batch_size,
                           shuffle=False)


# -*- coding: utf-8 -*-
# pylint: disable=W0621
import pytest
from unittestzero import Assert
from fixtures.server_roles import default_roles, server_roles

@pytest.mark.nondestructive
@pytest.mark.fixtureconf(server_roles=default_roles+('automate',))
@pytest.mark.usefixtures(
        "maximized",
        "setup_infrastructure_providers",
        "setup_pxe_provision",
        "mgmt_sys_api_clients")
class TestTemplateProvisioning:
    def test_linux_template_cancel(
            self,
            provisioning_start_page,
            provisioning_data_basic_only):
        '''Test Cancel button'''
        provisioning_start_page.click_on_template_item(
                provisioning_data_basic_only["template"])
        provision_pg = provisioning_start_page.click_on_continue()
        vm_pg = provision_pg.click_on_cancel()
        Assert.true(vm_pg.is_the_current_page,
                "not returned to the correct page")

    def test_linux_template_workflow(
            self,
            server_roles,
            provisioning_start_page,
            provisioning_data,
            mgmt_sys_api_clients,
            random_name):
        '''Test Basic Provisioning Workflow'''
        assert len(server_roles) == len(default_roles) + 1
        provisioning_start_page.click_on_template_item(
            provisioning_data["template"])
        provision_pg = provisioning_start_page.click_on_continue()
        self.complete_provision_pages_info(provisioning_data, provision_pg, \
            random_name)
        vm_pg = assert_vm_state(provisioning_data, provision_pg, "on", \
            random_name)
        remove_vm(provisioning_data, vm_pg, mgmt_sys_api_clients, \
            random_name)

    def complete_provision_pages_info(self,
        provisioning_data, provision_pg, random_name):
        ''' Fills in data for Provisioning tabs'''
        tab_buttons = provision_pg.tabbutton_region
        request_pg = tab_buttons.tabbutton_by_name("Request").click()
        request_pg = request_pg.fill_fields(
                "admin@example.com",
                "admin",
                "admin",
                "Adding a test note",
                "Manager Name")
        purpose_pg = tab_buttons.tabbutton_by_name("Purpose").click()
        # tree = purpose_pg.click_on_nodes(provisioning_data["node"],
        #        provisioning_data["child_node")
        catalog_pg = tab_buttons.tabbutton_by_name("Catalog").click()
        catalog_pg.fill_fields(
                provisioning_data["provision_type"],
                provisioning_data["pxe_server"],
                provisioning_data["server_image"],
                str(provisioning_data["count"]),
                '%s%s' % (provisioning_data["vm_name"], random_name),
                provisioning_data["vm_description"])
        environment_pg = tab_buttons.tabbutton_by_name("Environment").click()
        environment_pg.fill_fields(
                unicode(provisioning_data["host"]),
                unicode(provisioning_data["datastore"]))
        hardware_pg = tab_buttons.tabbutton_by_name("Hardware").click()
        network_pg = tab_buttons.tabbutton_by_name("Network").click()
        if ("PXE" in provisioning_data["provision_type"]) or \
           ("ISO" in provisioning_data["provision_type"]):
            customize_pg = tab_buttons.tabbutton_by_name("Customize").click()
            customize_pg.fill_fields(
                provisioning_data["root_password"],
                provisioning_data["address_node_value"],
                provisioning_data["customization_template"])
        schedule_pg = tab_buttons.tabbutton_by_name("Schedule").click()
        schedule_pg.fill_fields(
                provisioning_data["when_to_provision"],
                provisioning_data["power_on"],
                str(provisioning_data["time_until_retirement"]))

        services_requests_pg = schedule_pg.click_on_submit()
        Assert.true(services_requests_pg.is_the_current_page,
                "not returned to the correct page")
        Assert.equal(services_requests_pg.flash_message,
                "VM Provision Request was Submitted, "\
                "you will be notified when your VMs are ready")
        services_requests_pg.approve_request(1)
        services_requests_pg.wait_for_request_status("Last 24 Hours", \
                "Finished", 12)

def assert_vm_state(provisioning_data, current_page, \
    current_state, random_name):
    ''' Asserts that the VM is created in the expected state '''
    vm_pg = current_page.header.site_navigation_menu(
            'Infrastructure').sub_navigation_menu('Virtual Machines').click()
    vm_pg.refresh()
    vm_pg.wait_for_vm_state_change( '%s%s' % (provisioning_data["vm_name"],
            random_name), 'on', 12)
    Assert.equal(vm_pg.quadicon_region.get_quadicon_by_title(
        '%s%s' % (provisioning_data["vm_name"], random_name))\
        .current_state, current_state,
        "vm not in correct state: " + current_state)
    return vm_pg

def remove_vm(provisioning_data,
    current_page, provider_api_clients, random_name):
    '''Powers off and removes the VM'''
    vm_pg = current_page.header.site_navigation_menu(
            'Infrastructure').sub_navigation_menu('Virtual Machines').click()
    vm_pg.power_off(['%s%s' % (provisioning_data["vm_name"], random_name)])
    Assert.true(vm_pg.flash.message.startswith("Stop initiated"))
    vm_pg.wait_for_vm_state_change(
            '%s%s' % (provisioning_data["vm_name"], random_name), 'off', 12)
    Assert.equal(vm_pg.quadicon_region.get_quadicon_by_title(
            '%s%s' % (provisioning_data["vm_name"], random_name))\
            .current_state, 'off', "vm running")
    for provider in provider_api_clients.values():
        if ('%s%s' % (provisioning_data["vm_name"], random_name) + "/" +
                '%s%s' % (provisioning_data["vm_name"], random_name) + ".vmx"
                ) in provider.list_vm() or \
                '%s%s' % (provisioning_data["vm_name"], random_name) \
                in provider.list_vm():
            provider.delete_vm('%s%s' % (provisioning_data["vm_name"], \
                random_name))


import torch.nn as nn
import torch
import torch.nn.functional as F
from torch.nn.utils import spectral_norm
import mit_resnet
from torchvision import models
from non_local_embedded_gaussian import NONLocalBlock2D

MAX_LENGTH = 40

PAD_token = 0
SOS_token = 1
EOS_token = 2

class ResidualBlock(nn.Module):
    def __init__(self, input_channels, output_channels, stride=1):
        super(ResidualBlock, self).__init__()
        self.input_channels = input_channels
        self.output_channels = output_channels
        self.stride = stride
        self.bn1 = nn.BatchNorm2d(input_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv2d(input_channels, int(output_channels / 4), 1, 1, bias=False)
        self.bn2 = nn.BatchNorm2d(int(output_channels / 4))
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(int(output_channels / 4), int(output_channels / 4), 3, stride, padding=1, bias=False)
        self.bn3 = nn.BatchNorm2d(int(output_channels / 4))
        self.relu = nn.ReLU(inplace=True)
        self.conv3 = nn.Conv2d(int(output_channels / 4), output_channels, 1, 1, bias=False)
        self.conv4 = nn.Conv2d(input_channels, output_channels, 1, stride, bias=False)

    def forward(self, x):
        residual = x
        out = self.bn1(x)
        out1 = self.relu(out)
        out = self.conv1(out1)
        out = self.bn2(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn3(out)
        out = self.relu(out)
        out = self.conv3(out)
        if (self.input_channels != self.output_channels) or (self.stride != 1):
            residual = self.conv4(out1)
        out += residual
        return out

class cnn_att(nn.Module):
    # input size is 8*8
    def __init__(self, in_channels, out_channels):
        super(cnn_att, self).__init__()
        self.first_residual_blocks = ResidualBlock(in_channels, out_channels)

        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.bn = self.bn4 = nn.BatchNorm2d(in_channels, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)

        self.trunk_branches = nn.Sequential(
            ResidualBlock(in_channels, out_channels),
            ResidualBlock(in_channels, out_channels)
         )

        self.mpool1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)  # 4*4

        self.middle_2r_blocks = nn.Sequential(
            ResidualBlock(in_channels, out_channels),
            ResidualBlock(in_channels, out_channels)
        )

        # self.interpolation1 = nn.UpsamplingBilinear2d(size=size)  # 8*8

        self.conv1_1_blocks = nn.Sequential(
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=1, stride=1, bias=False),
            # nn.BatchNorm2d(out_channels),
            # nn.ReLU(inplace=True),
            # nn.Conv2d(out_channels, out_channels, kernel_size=1, stride=1, bias = False),
            nn.Sigmoid()
        )

        self.last_blocks = ResidualBlock(in_channels, out_channels)

    def forward(self, x, size):
        x = self.bn(self.conv(x))
        out_trunk = self.bn(self.conv(x))
        out_mpool1 = self.mpool1(x)
        out_middle_2r_blocks = self.bn(self.conv(out_mpool1))
        #
        # out_interp = self.interpolation1(out_middle_2r_blocks) + out_trunk
        out_interp = F.upsample(out_middle_2r_blocks, size=size, mode='bilinear', align_corners=True) + out_trunk
        # print(out_skip2_connection.data)
        # print(out_interp3.data)
        out_conv1_1_blocks = self.conv1_1_blocks(out_interp)
        out = (1 + out_conv1_1_blocks) * out_trunk
        out_last = self.bn(self.conv(out))

        return out_last

class ResnetDilated(nn.Module):
    def __init__(self, orig_resnet, dilate_scale=8):
        super(ResnetDilated, self).__init__()
        from functools import partial

        if dilate_scale == 8:
            orig_resnet.layer3.apply(
                partial(self._nostride_dilate, dilate=2))
            orig_resnet.layer4.apply(
                partial(self._nostride_dilate, dilate=4))
        elif dilate_scale == 4:
            orig_resnet.layer2.apply(
                partial(self._nostride_dilate, dilate=2))
            orig_resnet.layer3.apply(
                partial(self._nostride_dilate, dilate=2))
            orig_resnet.layer4.apply(
                partial(self._nostride_dilate, dilate=4))

        # take pretrained resnet, except AvgPool and FC
        self.conv1 = orig_resnet.conv1
        self.bn1 = orig_resnet.bn1
        self.relu1 = orig_resnet.relu1
        self.conv2 = orig_resnet.conv2
        self.bn2 = orig_resnet.bn2
        self.relu2 = orig_resnet.relu2
        self.conv3 = orig_resnet.conv3
        self.bn3 = orig_resnet.bn3
        self.relu3 = orig_resnet.relu3
        self.maxpool = orig_resnet.maxpool
        self.layer1 = orig_resnet.layer1
        self.layer2 = orig_resnet.layer2
        self.layer3 = orig_resnet.layer3
        self.layer4 = orig_resnet.layer4
        self.non_local1 = NONLocalBlock2D(64, sub_sample=True, bn_layer=True)
        self.non_local2 = NONLocalBlock2D(128, sub_sample=True, bn_layer=True)

    def _nostride_dilate(self, m, dilate):
        classname = m.__class__.__name__
        if classname.find('Conv') != -1:
            # the convolution with stride
            if m.stride == (2, 2):
                m.stride = (1, 1)
                if m.kernel_size == (3, 3):
                    m.dilation = (dilate//2, dilate//2)
                    m.padding = (dilate//2, dilate//2)
            # other convoluions
            else:
                if m.kernel_size == (3, 3):
                    m.dilation = (dilate, dilate)
                    m.padding = (dilate, dilate)

    def forward(self, x, return_feature_maps=True):
        conv_out = []

        x = self.relu1(self.bn1(self.conv1(x)))
        x = self.relu2(self.bn2(self.conv2(x)))
        x = self.relu3(self.bn3(self.conv3(x)))
        x = self.maxpool(x)

        x = self.non_local1(self.layer1(x))
        conv_out.append(x)
        x = self.non_local2(self.layer2(x))
        conv_out.append(x)
        x = self.layer3(x)
        conv_out.append(x)
        x = self.layer4(x)
        conv_out.append(x)

        if return_feature_maps:
            return conv_out
        return [x]

class CNN(nn.Module):
    def __init__(self, imgH, nc, leakyRelu=False):
        super(CNN, self).__init__()
        assert imgH % 16 == 0, 'imgH has to be a multiple of 16'

        ks = [3, 3, 3, 3, 3, 3, 2]
        ps = [1, 1, 1, 1, 1, 1, 0]
        ss = [1, 1, 1, 1, 1, 1, 1]
        nm = [64, 128, 256, 256, 512, 512, 512]

        cnn = nn.Sequential()
        self.orig_resnet = mit_resnet.__dict__['resnet18'](pretrained=False)
        self.net_encoder = ResnetDilated(self.orig_resnet,
                                         dilate_scale=8)

        self.att = cnn_att(512, 512)
        self.conv4 = nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.bn4 = nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        self.relu4 = nn.ReLU(512)
        self.conv5 = nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.relu5 = nn.ReLU(512)
        self.bn5 = nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        self.pool3 = nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 1), padding=(0, 1), dilation=1, ceil_mode=False)

        self.conv6 = nn.Conv2d(512, 1024, kernel_size=(2, 2), stride=(2, 1), padding=(0, 0))
        self.pool4 = nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 1), padding=(0, 1), dilation=1, ceil_mode=False)
        self.bn6 = nn.BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        self.relu6 = nn.ReLU(1024)

        def convRelu(i, batchNormalization=False):
            nIn = nc if i == 0 else nm[i - 1]
            nOut = nm[i]
            cnn.add_module('conv{0}'.format(i),
                           nn.Conv2d(nIn, nOut, ks[i], ss[i], ps[i]))
            if batchNormalization:
                cnn.add_module('batchnorm{0}'.format(i), nn.BatchNorm2d(nOut))
            if leakyRelu:
                cnn.add_module('relu{0}'.format(i),
                               nn.LeakyReLU(0.2, inplace=True))
            else:
                cnn.add_module('relu{0}'.format(i), nn.ReLU(True))

        convRelu(0)
        cnn.add_module('pooling{0}'.format(0), nn.MaxPool2d(2, 2))  # 64x16x64
        convRelu(1)
        cnn.add_module('pooling{0}'.format(1), nn.MaxPool2d(2, 2))  # 128x8x32
        convRelu(2, True)
        convRelu(3)
        cnn.add_module('pooling{0}'.format(2),
                       nn.MaxPool2d((2, 2), (2, 1), (0, 1)))  # 256x4x16
        convRelu(4, True)
        convRelu(5)
        cnn.add_module('pooling{0}'.format(3),
                       nn.MaxPool2d((2, 2), (2, 1), (0, 1)))  # 512x2x16
        convRelu(6, True)  # 512x1x16

        self.cnn = cnn
        self.attention = SelfAttention(1024)

    def forward(self, input):
        conv_all = self.net_encoder(input)

        conv = self.conv4(conv_all[3])
        conv = self.bn4(conv)
        conv = self.relu4(conv)
        conv = self.att(conv, conv.size()[2:])
        conv = self.conv5(conv)
        conv = self.bn5(conv)
        conv = self.relu5(conv)
        conv = self.pool3(conv)
        conv = self.conv6(conv)
        conv = self.pool4(conv)
        conv = self.bn6(conv)
        conv = self.relu6(conv)
        conv = self.attention(conv)

        b, c, h, w = conv.size()
        assert h == 1, "the height of conv must be 1"
        conv = conv.squeeze(2) # b *512 * width
        conv = conv.permute(0, 2, 1)  # [b, w, c]
        output = conv
        return output

class SelfAttention(nn.Module):

    def __init__(self, d):
        super(SelfAttention, self).__init__()

        assert d % 8 == 0
        self.projections = nn.ModuleList([
            spectral_norm(nn.Conv2d(d, d // 8, 1)),
            spectral_norm(nn.Conv2d(d, d // 8, 1)),
            spectral_norm(nn.Conv2d(d, d, 1))
        ])
        self.gamma = nn.Parameter(torch.zeros(1))  # shape [1]

    def forward(self, x):
        """
        Arguments:
            x: a float tensor with shape [b, d, h, w].
        Returns:
            a float tensor with shape [b, d, h, w].
        """
        b, d, h, w = x.size()

        q = self.projections[0](x)
        k = self.projections[1](x)
        v = self.projections[2](x)

        q = q.view(b, d // 8, h * w).permute(0, 2, 1)
        k = k.view(b, d // 8, h * w)
        v = v.view(b, d, h * w).permute(0, 2, 1)

        attention = torch.bmm(q, k)  # shape [b, h * w, h * w]
        attention = F.softmax(attention, dim=2)

        out = torch.bmm(attention, v)  # shape [b, h * w, d]
        out = out.permute(0, 2, 1).view(b, d, h, w)
        return x + self.gamma * out

class EncoderRNN(nn.Module):
    def __init__(self, input_size, hidden_size, n_layers=1, dropout=0.1):
        super(EncoderRNN, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.n_layers = n_layers
        self.gru = nn.GRU(self.input_size, self.hidden_size, num_layers=self.n_layers, bidirectional=True, dropout=(0 if n_layers == 1 else dropout))

    def forward(self, input, hidden=None):
        self.gru.flatten_parameters()
        outputs, hidden = self.gru(input, hidden)

        outputs = outputs[:,:,:self.hidden_size] + outputs[:,:,self.hidden_size:]
        return outputs, hidden

# Luong attention layer
class Attn(torch.nn.Module):
    def __init__(self, method, hidden_size):
        super(Attn, self).__init__()
        self.method = method
        if self.method not in ['dot', 'general', 'concat']:
            raise ValueError(self.method, "is not an appropriate attention method.")
        self.hidden_size = hidden_size
        if self.method == 'general':
            self.attn = torch.nn.Linear(self.hidden_size, hidden_size)
        elif self.method == 'concat':
            self.attn = torch.nn.Linear(self.hidden_size * 2, hidden_size)
            self.v = torch.nn.Parameter(torch.FloatTensor(hidden_size))

    def dot_score(self, hidden, encoder_output):
        return torch.sum(hidden * encoder_output, dim=2)

    def general_score(self, hidden, encoder_output):
        energy = self.attn(encoder_output)
        return torch.sum(hidden * energy, dim=2)

    def concat_score(self, hidden, encoder_output):
        energy = self.attn(torch.cat((hidden.expand(encoder_output.size(0), -1, -1), encoder_output), 2)).tanh()
        return torch.sum(self.v * energy, dim=2)

    def forward(self, hidden, encoder_outputs):
        # Calculate the attention weights (energies) based on the given method
        if self.method == 'general':
            attn_energies = self.general_score(hidden, encoder_outputs)
        elif self.method == 'concat':
            attn_energies = self.concat_score(hidden, encoder_outputs)
        elif self.method == 'dot':
            attn_energies = self.dot_score(hidden, encoder_outputs)

        # Transpose max_length and batch_size dimensions
        attn_energies = attn_energies.t()

        # Return the softmax normalized probability scores (with added dimension)
        return F.softmax(attn_energies, dim=1).unsqueeze(1)

class LuongAttnDecoderRNN(nn.Module):
    def __init__(self, attn_model, embedding, hidden_size, output_size, n_layers=1, dropout=0.1):
        super(LuongAttnDecoderRNN, self).__init__()

        # Keep for reference
        self.attn_model = attn_model
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.n_layers = n_layers
        self.dropout = dropout

        # Define layers
        self.embedding = embedding
        self.embedding_dropout = nn.Dropout(dropout)
        self.gru = nn.GRU(hidden_size, hidden_size, n_layers, dropout=(0 if n_layers == 1 else dropout))
        self.concat = nn.Linear(hidden_size * 2, hidden_size)
        self.out = nn.Linear(hidden_size, output_size)

        self.attn = Attn(attn_model, hidden_size)

    def forward(self, input_step, last_hidden, encoder_outputs):
        # Note: we run this one step (word) at a time
        # Get embedding of current input word
        embedded = self.embedding(input_step).view(1,-1,self.hidden_size)
        embedded = self.embedding_dropout(embedded)
        # Forward through unidirectional GRU
        self.gru.flatten_parameters()
        rnn_output, hidden = self.gru(embedded, last_hidden)
        # Calculate attention weights from the current GRU output
        attn_weights = self.attn(rnn_output, encoder_outputs)
        # Multiply attention weights to encoder outputs to get new "weighted sum" context vector
        context = attn_weights.bmm(encoder_outputs.transpose(0, 1))
        # Concatenate weighted context vector and GRU output using Luong eq. 5
        rnn_output = rnn_output.squeeze(0)
        context = context.squeeze(1)
        concat_input = torch.cat((rnn_output, context), 1)
        concat_output = torch.tanh(self.concat(concat_input))
        # Predict next word using Luong eq. 6
        output = self.out(concat_output)
        output = F.log_softmax(output, dim=1)
        # Return output and final hidden state
        return output, hidden

class model(nn.Module):
    def __init__(self, cnn, encoder, decoder):
        super(model, self).__init__()

        self.cnn = cnn
        self.encoder = encoder
        self.decoder = decoder

    def forward(self, image, max_length):

        batch_size = image.size()[0]

        input_tensor = self.cnn(image)
        input_tensor = input_tensor.permute(1, 0, 2)

        encoder_outputs, encoder_hidden = self.encoder(
            input_tensor)

        decoder_input = torch.tensor([[SOS_token] * batch_size]).cuda()
        decoder_hidden = encoder_hidden[:self.decoder.n_layers]

        decoder_outputs = []
        for di in range(max_length):
            decoder_output, decoder_hidden = self.decoder(decoder_input, decoder_hidden, encoder_outputs)
            decoder_scores, decoder_input = torch.max(decoder_output, dim=1)
            decoder_outputs.append(decoder_output)
            # loss += self.criterion(decoder_output, target_tensor[di].squeeze(1))
        decoder_outputs = torch.stack(decoder_outputs, 0)
        return decoder_outputs.permute(1, 0, 2)


#from __future__ import print_function
import cv2
import argparse
import os
import random
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torchnet.meter as meter
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.autograd import Variable
from dataset import multiPIE
from siamese_model_2nd import Siamese
from contrastive import ContrastiveLoss
import numpy as np
# import cv2
#from pycrayon import CrayonClient

#for plotting loss
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import time,math
from logger import Logger
# from models_Parsing import ParseNet
saveFile = open('/home/shumao/wyw_files/siamese_output_M_3/record.txt', 'w')
saveFile.write("niter:" + str(50000) + "\n")
saveFile.write("---lr:" + str(0.0001) + "\n")
saveFile.write("beta1:" + str(0.7) + "\n")
saveFile.write("W:-1-x-x-x-x-x-" + "\n")
saveFile.write("L1 Loss" + "\n")
saveFile.write("after load model from: train-3-28000pth")
logger = Logger('./log_1');

parser = argparse.ArgumentParser()
parser.add_argument('--batchSize', type=int, default=64, help='input batch size')
parser.add_argument('--loadSize', type=int, default=100, help='the height / width of the input image to network')
parser.add_argument('--fineSize', type=int, default=96, help='the height / width of the input image to network')
parser.add_argument('--id_num', type=int, default=200, help='Total training identity.')
parser.add_argument('--pose_num', type=int, default=9, help='Total training pose.')
parser.add_argument('--light_num', type=int, default=20, help='Total training light.')
parser.add_argument('--niter', type=int, default=50000, help='number of iterations to train for')
parser.add_argument('--lr', type=float, default=0.0001, help='learning rate, default=0.0002')
parser.add_argument('--beta1', type=float, default=0.7, help='beta1 for adam. default=0.7')
parser.add_argument('--cuda', action='store_true', help='enables cuda')
parser.add_argument('--outf', default='/home/shumao/wyw_files/siamese_output_M_3', help='folder to output images and model checkpoints')
parser.add_argument('--manualSeed', type=int, help='manual seed')
parser.add_argument('--dataPath', default='/home/shumao/dr-gan/Data_new_realigned2/setting2/train/', help='which dataset to train on')
parser.add_argument('--modelPath', default='/home/shumao/wyw_files/siamese_output_3/netS_28000.pth', help='which dataset to train on')
parser.add_argument('--save_step', type=int, default=400, help='save weights every 400 iterations ')
parser.add_argument('--labelPath', default='/home/shumao/dr-gan/Data_new_realigned2/setting2/Facedata/', help='which dataset to train on')

opt = parser.parse_args()
print(opt) # print every parser arguments
# print(opt.niter)

try:
    os.makedirs(opt.outf)
except OSError:
    pass

w_r = 1
# w_cL = 0.02
# w_cP = 0.02
# w_cI = 0.02
# w_P = 0.02
# w_L = 0.02

if opt.manualSeed is None:
    opt.manualSeed = random.randint(1, 10000)
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if opt.cuda:
    torch.cuda.manual_seed_all(opt.manualSeed)

cudnn.benchmark = True
#---------------------Load Mask-------------------
mask = np.load('mask_20.npy')
mask = mask.astype(np.float32)
M = torch.from_numpy(mask.transpose((2, 0, 1)))
FinalMask = M.expand(opt.batchSize,3,96,96)
# print m.size()
# 3x96x96

#---------------------Load DATA-------------------------
dataset_1 = multiPIE(opt.dataPath,opt.loadSize,opt.fineSize,labelPath = opt.labelPath)
# dataset_2 = multiPIE(opt.dataPath,opt.loadSize,opt.fineSize,opt.labelPath)
dataset_test = multiPIE('/home/shumao/dr-gan/comparison/',opt.loadSize,opt.fineSize,labelPath = opt.labelPath)
loader_train_1 = torch.utils.data.DataLoader(dataset=dataset_1,
                                      batch_size = opt.batchSize,
                                      shuffle=True,
                                      num_workers=4,
                                      drop_last = True)
# loader_train_2 = torch.utils.data.Dataloader(dataset=dataset_1,
#                                       batch_size = opt.batchSize,
#                                       shuffle=True,
#                                       num_workers=4)

loader_test = torch.utils.data.DataLoader(dataset=dataset_test,
                                          batch_size = 9,
                                          shuffle=False,
                                          num_workers=4)
data_train_1 = iter(loader_train_1)
# data_train_2 = iter(loader_train_2)
data_test = iter(loader_test)

#----------------------Parameters-----------------------
num_pose = opt.pose_num
num_light = opt.light_num
num_iden = opt.id_num

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') !=-1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') !=-1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

netS = Siamese()
netS.load_state_dict(torch.load(opt.modelPath))

#-----------------params freeze-----------------
for param in netS.conv11.parameters():
    param.requires_grad = False
for param in netS.conv1r.parameters():
    param.requires_grad = False
for param in netS.conv12.parameters():
    param.requires_grad = False
for param in netS.conv21.parameters():
    param.requires_grad = False
for param in netS.conv22.parameters():
    param.requires_grad = False
for param in netS.conv23.parameters():
    param.requires_grad = False
for param in netS.conv31.parameters():
    param.requires_grad = False
for param in netS.conv32.parameters():
    param.requires_grad = False
for param in netS.conv33.parameters():
    param.requires_grad = False
for param in netS.conv41.parameters():
    param.requires_grad = False
for param in netS.conv42.parameters():
    param.requires_grad = False
for param in netS.conv43.parameters():
    param.requires_grad = False
for param in netS.conv51.parameters():
    param.requires_grad = False
for param in netS.conv52.parameters():
    param.requires_grad = False
for param in netS.conv53.parameters():
    param.requires_grad = False
for param in netS.convfc.parameters():
    param.requires_grad = False

#-----------------params freeze-----------------
if(opt.cuda):
    netS.cuda()

#-------------------Loss & Optimization

optimizerS = torch.optim.Adam(filter(lambda p: p.requires_grad, netS.parameters()),lr=opt.lr, betas=(opt.beta1, 0.999))

poss_contrastive_loss = ContrastiveLoss() # load from the begining
light_contrastive_loss = ContrastiveLoss()
identity_contrastive_loss = ContrastiveLoss()
reconstructe_loss = nn.L1Loss()
pose_class_loss = nn.CrossEntropyLoss()
light_class_loss = nn.CrossEntropyLoss()

#------------------ Global Variables------------------
input_pose_1 = torch.LongTensor(opt.batchSize)
input_light_1 = torch.LongTensor(opt.batchSize)
# input_pose_2 = torch.LongTensor(opt.batchSize)
# input_light_2 = torch.LongTensor(opt.batchSize)

inputImg_1 = torch.FloatTensor(opt.batchSize, 3, opt.fineSize, opt.fineSize)
inputImg_2 = torch.FloatTensor(opt.batchSize, 3, opt.fineSize, opt.fineSize)
GT = torch.FloatTensor(opt.batchSize, 3,opt.fineSize, opt.fineSize)
same_pose = torch.FloatTensor(opt.batchSize)
same_iden = torch.FloatTensor(opt.batchSize)
same_light = torch.FloatTensor(opt.batchSize)

# w_1 = torch.FloatTensor(1)
# w_2 = torch.FloatTensor(20)
# w_3 = torch.FloatTensor(10)
# w_4 = torch.FloatTensor(10)
# w_5 = torch.FloatTensor(10)
# w_6 = torch.FloatTensor(20)
# output_pose_1_label = torch.LongTensor(opt.batchSize)
# output_pose_2_label = torch.LongTensor(opt.batchSize)
# output_light_1_label = torch.LongTensor(opt.batchSize)
# output_light_2_label = torch.LongTensor(opt.batchSize)

input_pose_1 = Variable(input_pose_1)
# input_pose_2 = Variable(input_pose_2)
input_light_1 = Variable(input_light_1)
# input_light_2 = Variable(input_light_2)

inputImg_1 = Variable(inputImg_1)
inputImg_2 = Variable(inputImg_2)
GT = Variable(GT)
same_pose = Variable(same_pose)
same_iden = Variable(same_iden)
same_light = Variable(same_light)

FinalMask = Variable(FinalMask)

# w_1 = Variable(w_1, requires_grad = False)
# w_2 = Variable(w_2, requires_grad = False)
# w_3 = Variable(w_3, requires_grad = False)
# w_4 = Variable(w_4, requires_grad = False)
# w_5 = Variable(w_5, requires_grad = False)
# w_6 = Variable(w_6, requires_grad = False)

pose_mtr = meter.ConfusionMeter(k=opt.pose_num)
light_mtr = meter.ConfusionMeter(k=opt.light_num)

if(opt.cuda):

    input_pose_1 = input_pose_1.cuda()
    # input_pose_2 = input_pose_2.cuda()
    input_light_1 = input_light_1.cuda()
    # input_light_2 = input_light_2.cuda()
    inputImg_1 = inputImg_1.cuda()
    inputImg_2 = inputImg_2.cuda()
    GT = GT.cuda()
    same_pose = same_pose.cuda()
    same_light = same_light.cuda()
    same_iden = same_iden.cuda()

    FinalMask = FinalMask.cuda()

    # w_1 = w_1.cuda()
    # w_2 = w_1.cuda()
    # w_3 = w_1.cuda()
    # w_4 = w_1.cuda()
    # w_5 = w_1.cuda()
    # w_6 = w_1.cuda()
    # poss_contrastive_loss.cuda()
    # light_contrastive_loss.cuda()
    # identity_contrastive_loss.cuda()
    pose_class_loss.cuda()
    light_class_loss.cuda()
    reconstructe_loss.cuda()

#------------------test---------

# k = 0 # for meter

err_total = 0
err_recon = 0
err_contraL = 0
err_contraP = 0
err_contraI = 0
err_classP = 0
err_classL = 0

def test(iteration, data_test, loader_test):
    try:
        images_1,po_1,li_1,GT_1,by_image,same_po,same_li,same_id = data_test.next()
    except StopIteration:
        data_test = iter(loader_test)
        images_1,po_1,li_1,GT_1,by_image,same_po,same_li,same_id = data_test.next()

    GT.data.resize_(GT_1.size()).copy_(GT_1)
    inputImg_1.data.resize_(images_1.size()).copy_(images_1)
    inputImg_2.data.resize_(by_image.size()).copy_(by_image)
    input_pose_1.data.resize_(po_1.size()).copy_(po_1)
    input_light_1.data.resize_(li_1.size()).copy_(li_1)

    output_pose_1, output_pose_2, output_light_1, output_light_2, out_f_1, out_f_2, out = netS(inputImg_1, inputImg_2)
    vutils.save_image(out.data,
        '%s/fake_samples_iteration_%03d.png' % (opt.outf, iteration), normalize=True)
    vutils.save_image(inputImg_1.data,
            '%s/input_samples_iteration_%03d.png' % (opt.outf, iteration), normalize=True)

#-------------------train----------------------
for iteration in range(1,opt.niter+1):
    running_corrects = 0
    running_corrects_light = 0
    try:
        images_1,po_1,li_1,GT_1,by_image,same_po,same_li,same_id= data_train_1.next()
    except StopIteration:
        data_train_1 = iter(loader_train_1)
        images_1,po_1,li_1,GT_1,by_image,same_po,same_li,same_id = data_train_1.next()

    GT.data.resize_(GT_1.size()).copy_(GT_1)

    inputImg_1.data.resize_(images_1.size()).copy_(images_1)
    inputImg_2.data.resize_(by_image.size()).copy_(by_image)

    input_pose_1.data.resize_(po_1.size()).copy_(po_1)
    input_light_1.data.resize_(li_1.size()).copy_(li_1)

    same_pose.data.resize_(same_po.size()).copy_(same_po)
    same_light.data.resize_(same_li.size()).copy_(same_li)
    same_iden.data.resize_(same_id.size()).copy_(same_id)
    netS.zero_grad()

    output_pose_1, output_pose_2, output_light_1, output_light_2, out_f_1, out_f_2, out = netS(inputImg_1, inputImg_2)
    #-----------------mask test area-----------------------------
    # print out.data.type()
    # print GT.data.type()
    # print FinalMask.data.type() same
    # print FinalMask.data.size() 64x3x96x96
    # print FinalMask.data.size()
    # print out.data.size()

    Final_out = FinalMask * out

    Final_GT = FinalMask * GT

    #-----------------mask test area-----------------------------
    # f_1 & f_2 variable
    # same_iden variable
    err_recon = reconstructe_loss(Final_out, Final_GT)
    err_contraI = identity_contrastive_loss(out_f_1, out_f_2, same_iden)
    err_contraP = poss_contrastive_loss(output_pose_1, output_pose_2, same_pose)
    err_contraL = light_contrastive_loss(output_light_1,output_light_2, same_light)
    err_classL = light_class_loss(output_light_1, input_light_1)
    err_classP = pose_class_loss(output_pose_1, input_pose_1)
    # print(err_recon.data.size())
    # print(err_contraL.data.size())
    # print(err_classP.data.size())
    # modify the contrastive loss function to make contrastive loss be 1Lx1L
    # contrastive loss and Softmax and Loss1 are all requires_grad
    # err_total = 1 * err_recon + 10 * err_contraP + 10 * err_contraI + 10 * err_classP + 20 * err_classL
    # err_total = err_recon + err_contraI + err_contraP + err_contraL + err_classL + err_classP
    err_total = w_r * err_recon
    err_total.backward()
    optimizerS.step()

    #----------------------Visualize-----------
    if(iteration % 200 == 0):

        pose_mtr.add(output_pose_1.data, input_pose_1.data)
        pose_trainacc = pose_mtr.value().diagonal().sum()*1.0/opt.batchSize
        pose_mtr.reset()

        light_mtr.add(output_light_1.data, input_light_1.data)
        light_trainacc = light_mtr.value().diagonal().sum()*1.0/opt.batchSize
        light_mtr.reset()
        #-----------------------------------------
        test(iteration, data_test, loader_test)
    # #pose prediction

    # preds_pose = torch.max(output_pose_1.data, 1)
    # running_corrects += torch.sum(preds == input_pose_1)
    # print('pose_accuracy: %.2f'
    #         % (running_corrects * 1.0/images.size(0)))

    # #light prediction
    # preds_light = torch.max(output_light_1.data, 1)
    # running_corrects_light += torch.sum(preds_light == input_light_1)
    # print('light_accuracy: %.2f'
    #         % (running_corrects_light * 1.0/images.size(0)))

    print('----------------------------------------')
    print('[%d/%d] Loss_S: %.4f ' %(iteration, opt.niter, err_total.data[0]))
    print('        Reco_S: %.4f ' %(err_recon.data[0]))
    print('        conL_S: %.4f ' %(err_contraL.data[0]))
    print('        conP_S: %.4f ' %(err_contraP.data[0]))
    print('        conI_S: %.4f ' %(err_contraI.data[0]))
    print('        Clas_P: %.4f ' %(err_classP.data[0]))
    print('        Clas_L: %.4f ' %(err_classL.data[0]))

    if(iteration % opt.save_step == 0):
        torch.save(netS.state_dict(), '%s/netS_%d.pth' % (opt.outf,iteration))