import textwrap
from pathlib import Path

from pyfomod import fomod, parser

PACKAGE_PATH = Path(__file__).parent / "package_test"
INFO_PATH = Path(PACKAGE_PATH, "fomod", "info.xml")
CONF_PATH = Path(PACKAGE_PATH, "fomod", "moduleconfig.xml")

def test_preserve_data(tmp_path):
    with INFO_PATH.open() as info_file:
        orig_info = info_file.read()
    with CONF_PATH.open() as conf_file:
        orig_conf = conf_file.read()
    root = parser.parse(str(PACKAGE_PATH))
    parser.write(root, str(tmp_path))
    info_path = tmp_path / "fomod" / "info.xml"
    with info_path.open() as info_file:
        new_info = info_file.read()
    conf_path = tmp_path / "fomod" / "moduleconfig.xml"
    with conf_path.open() as conf_file:
        new_conf = conf_file.read()
    assert orig_info == new_info
    assert orig_conf == new_conf

def test_parse(tmp_path):
    root = parser.parse(str(PACKAGE_PATH))
    tuple_root = parser.parse((str(INFO_PATH), str(CONF_PATH)))
    assert root.to_string() == tuple_root.to_string()
    assert root._info.to_string() == tuple_root._info.to_string()
    content = textwrap.dedent(
        """\
            <config>
            <!-- comment -->
            </config>"""
    )
    (tmp_path / "fomod").mkdir()
    conf_path = tmp_path / "fomod" / "moduleconfig.xml"
    with conf_path.open("w") as conf_file:
        conf_file.write(content)
        conf_file.write("\n")
    expected = [
        fomod.ValidationWarning(
            "Syntax Error",
            "Element 'config': Missing child element(s). "
            "Expected is ( moduleName ).",
            None,
            critical=True,
        ),
        fomod.ValidationWarning(
            "Comment Detected",
            "There are comments in this fomod, they will be ignored.",
            None,
            critical=True,
        ),
        fomod.ValidationWarning(
            "Missing Info",
            "Info.xml is missing from the fomod subfolder.",
            None,
            critical=True,
        ),
    ]
    warnings = []
    parser.parse((None, str(conf_path)), warnings=warnings)
    assert warnings == expected


#!/usr/bin/python3

"""
Augur.py

APP: Inquisition
DESC: A library used to fetch and receive OSINT data used for further learning
CREATION_DATE: 2017-05-14

"""

# MODULES
# | Native

# | Third-Party

# | Custom
from lib.inquisit.Inquisit import Inquisit

# METADATA
__author__ = 'Joshua Carlson-Purcell'
__copyright__ = 'Copyright 2017, CarlsoNet'
__license__ = 'MIT'
__version__ = '1.0.0-alpha'
__maintainer__ = 'Joshua Carlson-Purcell'
__email__ = 'jcarlson@carlso.net'
__status__ = 'Development'

class Augur(Inquisit):
    """
    OSINT aggreegation utility
    """

    def __init__(self, cfg):
        Inquisit.__init__(self, cfg, lgrName=__name__)


import time
import os
import numpy as np
import torch
import random
from random import choice
import torch.nn as nn
import torch.utils.data as data
from torchvision import transforms
from torch.autograd import Variable

import itertools
from tqdm import tqdm as tqdm
import datetime
from tensorboardX import SummaryWriter

from utils import *
import config_decomposer as config
from DataLoader_Star import Dataset

from model.model_decomposer import Discriminator_CAM as Discriminator
from model.model_decomposer import Discriminator as Discriminator_Domain
from model.model_decomposer import Integrator, Intrinsic_Encoder, Style_Encoder
from loss.loss import MMD

def weights_init(m):
    for each in m.modules():
        if isinstance(each, nn.Conv2d):
            nn.init.xavier_uniform_(each.weight.data)
            if each.bias is not None:
                each.bias.data.zero_()
        elif isinstance(each, nn.BatchNorm2d):
            each.weight.data.fill_(1)
        #     each.bias.data.zero_()
        # elif isinstance(each, nn.Linear):
        #     nn.init.xavier_uniform_(each.weight.data)
        #     each.bias.data.zero_()

def label2onehot(labels, c):
    batch_size = labels.size(0)
    out = torch.zeros(batch_size)
    out[np.arange(batch_size)] = c
    return out.long().cuda()

if __name__ == '__main__':

    if config.train['random_seed'] is None:
        config.train['random_seed'] = random.randint(1, 10000)
    print('Random seed: ', config.train['random_seed'])
    random.seed(config.train['random_seed'])
    torch.manual_seed(config.train['random_seed'])

    if not os.path.exists(config.train['logs']):
        os.mkdir(config.train['logs'])
    if not os.path.exists(config.train['checkpoint']):
        os.mkdir(config.train['checkpoint'])
    if not os.path.exists(config.train['out']):
        os.mkdir(config.train['out'])

    if config.train['cuda']:
        torch.cuda.manual_seed_all(config.train['random_seed'])

    # Dataloader
    trainset = Dataset()
    trainLoader = data.DataLoader(trainset, batch_size=config.train['batch_size'], shuffle=True,
                                  num_workers=config.train['num_workers'])

    # Model
    D_DOMAIN = Discriminator_Domain()
    D_VIS = Discriminator()
    D_NIR = Discriminator()
    D_SKETCH = Discriminator()
    Integrator = Integrator()
    Style_Encoder = Style_Encoder()
    Intrinsic_Encoder = Intrinsic_Encoder()

    # Losses
    l1_loss = torch.nn.L1Loss()
    mse_loss = torch.nn.MSELoss()
    cross_entropy = torch.nn.CrossEntropyLoss()
    mmd = MMD()

    # Optimizer
    optimizer_D_DOMAIN = torch.optim.Adam(D_DOMAIN.parameters(), lr=config.train['lr_D'],
                                          betas=(config.train['beta1_D'], config.train['beta2_D']),
                                          weight_decay=config.train['weight_decay_D'])
    optimizer_D_VIS = torch.optim.Adam(D_VIS.parameters(), lr=config.train['lr_D'],
                                       betas=(config.train['beta1_D'], config.train['beta2_D']),
                                       weight_decay=config.train['weight_decay_D'])
    optimizer_D_NIR = torch.optim.Adam(D_NIR.parameters(), lr=config.train['lr_D'],
                                       betas=(config.train['beta1_D'], config.train['beta2_D']),
                                       weight_decay=config.train['weight_decay_D'])
    optimizer_D_SKETCH = torch.optim.Adam(D_SKETCH.parameters(), lr=config.train['lr_D'],
                                          betas=(config.train['beta1_D'], config.train['beta2_D']),
                                          weight_decay=config.train['weight_decay_D'])
    # optimizer_G = torch.optim.Adam(itertools.chain(Style_Encoder.parameters(), Intrinsic_Encoder.parameters(), Integrator.parameters()),
    #                                lr=config.train['lr_G'],
    #                                betas=(config.train['beta1_G'], config.train['beta2_G']),
    #                                weight_decay=config.train['weight_decay_G'])
    optimizer_Style = torch.optim.Adam(Style_Encoder.parameters(),
                                       lr = config.train['lr_G'],
                                       betas=(config.train['beta1_G'],config.train['beta2_G']),
                                       weight_decay=config.train['weight_decay_G'])
    optimizer_Intrinsic = torch.optim.Adam(Intrinsic_Encoder.parameters(),
                                       lr=config.train['lr_G'],
                                       betas=(config.train['beta1_G'], config.train['beta2_G']),
                                       weight_decay=config.train['weight_decay_G'])
    optimizer_Integrator = torch.optim.Adam(Integrator.parameters(),
                                       lr=config.train['lr_G'],
                                       betas=(config.train['beta1_G'], config.train['beta2_G']),
                                       weight_decay=config.train['weight_decay_G'])

    if config.train['resume_Style_Encoder']:
        checkpoint = torch.load(config.train['resume_Style_Encoder'])
        Style_Encoder.load_state_dict(checkpoint['state_dict_Style_Encoder'])
        start_epoch = checkpoint['epoch_Style_Encoder'] + 1
        optim_checkpoint = torch.load(config.train['resume_optim_Style'])
        optimizer_Style.load_state_dict(optim_checkpoint)
        for state in optimizer_Style.state.values():
            for k, v in state.items():
                if isinstance(v, torch.Tensor):
                    state[k] = v.cuda()

        checkpoint = torch.load(config.train['resume_Intrinsic_Encoder'])
        Intrinsic_Encoder.load_state_dict(checkpoint['state_dict_Intrinsic_Encoder'])
        start_epoch = checkpoint['epoch_Intrinsic_Encoder'] + 1
        optim_checkpoint = torch.load(config.train['resume_optim_Intrinsic'])
        optimizer_Intrinsic.load_state_dict(optim_checkpoint)
        for state in optimizer_Intrinsic.state.values():
            for k, v in state.items():
                if isinstance(v, torch.Tensor):
                    state[k] = v.cuda()

        checkpoint = torch.load(config.train['resume_Integrator'])
        Integrator.load_state_dict(checkpoint['state_dict_Integrator'])
        start_epoch = checkpoint['epoch_Integrator'] + 1
        optim_checkpoint = torch.load(config.train['resume_optim_Integrator'])
        optimizer_Integrator.load_state_dict(optim_checkpoint)
        for state in optimizer_Integrator.state.values():
            for k, v in state.items():
                if isinstance(v, torch.Tensor):
                    state[k] = v.cuda()

    else:
        Integrator.apply(weights_init)
        Intrinsic_Encoder.apply(weights_init)
        Style_Encoder.apply(weights_init)
        start_epoch = 0

    if config.train['resume_D_VIS']:
        checkpoint = torch.load(config.train['resume_D_VIS'])
        D_VIS.load_state_dict(checkpoint['state_dict_D_VIS'])

        checkpoint = torch.load(config.train['resume_D_NIR'])
        D_NIR.load_state_dict(checkpoint['state_dict_D_NIR'])

        checkpoint = torch.load(config.train['resume_D_SKETCH'])
        D_SKETCH.load_state_dict(checkpoint['state_dict_D_SKETCH'])

        checkpoint = torch.load(config.train['resume_D_DOMAIN'])
        D_DOMAIN.load_state_dict(checkpoint['state_dict_D_DOMAIN'])

        # start_epoch = checkpoint['epoch_D_DOMAIN'] + 1

        optim_checkpoint_VIS = torch.load(config.train['resume_optim_D_VIS'])
        optimizer_D_VIS.load_state_dict(optim_checkpoint_VIS)

        for state in optimizer_D_VIS.state.values():
            for k, v in state.items():
                if isinstance(v, torch.Tensor):
                    state[k] = v.cuda()

        optim_checkpoint_NIR = torch.load(config.train['resume_optim_D_NIR'])
        optimizer_D_NIR.load_state_dict(optim_checkpoint_NIR)

        for state in optimizer_D_NIR.state.values():
            for k, v in state.items():
                if isinstance(v, torch.Tensor):
                    state[k] = v.cuda()

        optim_checkpoint_SKETCH = torch.load(config.train['resume_optim_D_SKETCH'])
        optimizer_D_SKETCH.load_state_dict(optim_checkpoint_SKETCH)

        for state in optimizer_D_SKETCH.state.values():
            for k, v in state.items():
                if isinstance(v, torch.Tensor):
                    state[k] = v.cuda()

        optim_checkpoint_DOMAIN = torch.load(config.train['resume_optim_D_DOMAIN'])
        optimizer_D_DOMAIN.load_state_dict(optim_checkpoint_DOMAIN)

        for state in optimizer_D_DOMAIN.state.values():
            for k, v in state.items():
                if isinstance(v, torch.Tensor):
                    state[k] = v.cuda()

    else:
        D_VIS.apply(weights_init)
        D_NIR.apply(weights_init)
        D_SKETCH.apply(weights_init)
        start_epoch = 0

    if config.train['cuda']:
        D_VIS.cuda()
        D_NIR.cuda()
        D_SKETCH.cuda()
        D_DOMAIN.cuda()
        Intrinsic_Encoder.cuda()
        Style_Encoder.cuda()
        Integrator.cuda()
        l1_loss.cuda()
        mse_loss.cuda()
        cross_entropy.cuda()

    # ################################       LR Schedulers           ################################
    # lr_schedule_G = torch.optim.lr_scheduler.LambdaLR(optimizer_G, lr_lambda=LambdaLR(config.train['epochs'],
    #                                                                                config.train['lr_G_decay_rate'],
    #                                                                                config.train['lr_G_decay_epoch'],
    #                                                                                config.train['lr_G']).step)

    lr_schedule_Style = torch.optim.lr_scheduler.LambdaLR(optimizer_Style, lr_lambda=LambdaLR(config.train['epochs'],
                                                                                config.train['lr_G_decay_rate'],
                                                                                config.train['lr_G_decay_epoch'],
                                                                                config.train['lr_G']).step)

    lr_schedule_Intrinsic = torch.optim.lr_scheduler.LambdaLR(optimizer_Intrinsic, lr_lambda=LambdaLR(config.train['epochs'],
                                                                                config.train['lr_G_decay_rate'],
                                                                                config.train['lr_G_decay_epoch'],
                                                                                config.train['lr_G']).step)

    lr_schedule_Integrator = torch.optim.lr_scheduler.LambdaLR(optimizer_Integrator, lr_lambda=LambdaLR(config.train['epochs'],
                                                                                config.train['lr_G_decay_rate'],
                                                                                config.train['lr_G_decay_epoch'],
                                                                                config.train['lr_G']).step)

    lr_schedule_D_VIS = torch.optim.lr_scheduler.LambdaLR(optimizer_D_VIS, lr_lambda=LambdaLR(config.train['epochs'],
                                                                                config.train['lr_D_decay_rate'],
                                                                                config.train['lr_D_decay_epoch'],
                                                                                config.train['lr_D']).step)

    lr_schedule_D_NIR = torch.optim.lr_scheduler.LambdaLR(optimizer_D_NIR, lr_lambda=LambdaLR(config.train['epochs'],
                                                                                config.train['lr_D_decay_rate'],
                                                                                config.train['lr_D_decay_epoch'],
                                                                                config.train['lr_D']).step)

    lr_schedule_D_SKETCH = torch.optim.lr_scheduler.LambdaLR(optimizer_D_SKETCH, lr_lambda=LambdaLR(config.train['epochs'],
                                                                                config.train['lr_D_decay_rate'],
                                                                                config.train['lr_D_decay_epoch'],
                                                                                config.train['lr_D']).step)

    lr_schedule_D_DOMAIN = torch.optim.lr_scheduler.LambdaLR(optimizer_D_DOMAIN, lr_lambda=LambdaLR(config.train['epochs'],
                                                                                config.train['lr_D_decay_rate'],
                                                                                config.train['lr_D_decay_epoch'],
                                                                                config.train['lr_D']).step)

    prev_time = time.time()

    count = int(len(trainLoader) // config.train['save_img_num'])

    vis_buffer = ReplayBuffer()
    nir_buffer = ReplayBuffer()
    sketch_buffer = ReplayBuffer()

    print('***************** Save Image Iteration: {} ********************'.format(count))
    # writer_loss_G_epochs = SummaryWriter(config.train['logs'] + 'epochs/loss_G')
    writer_loss_Intrinsic_epochs = SummaryWriter(config.train['logs'] + 'epochs/loss_Intrinsic')
    writer_loss_Style_epochs = SummaryWriter(config.train['logs'] + 'epochs/loss_Style')
    writer_loss_Integrator_epochs = SummaryWriter(config.train['logs'] + 'epochs/loss_Integrator')
    writer_loss_D_VIS_epochs = SummaryWriter(config.train['logs'] + 'epochs/loss_D_VIS')
    writer_loss_D_NIR_epochs = SummaryWriter(config.train['logs'] + 'epochs/loss_D_NIR')
    writer_loss_D_SKETCH_epochs = SummaryWriter(config.train['logs'] + 'epochs/loss_D_SKETCH')
    writer_loss_D_DOMAIN_epochs = SummaryWriter(config.train['logs'] + 'epochs/loss_D_DOMAIN')

    # Training
    if config.train['if_train']:
        for epoch in range(start_epoch, config.train['epochs']):
            # writer_loss_G_steps = SummaryWriter(config.train['logs'] + 'steps/loss_G')
            writer_loss_Intrinsic_steps = SummaryWriter(config.train['logs'] + 'steps/loss_Intrinsic')
            writer_loss_Style_steps = SummaryWriter(config.train['logs'] + 'steps/loss_Style')
            writer_loss_Integrator_steps = SummaryWriter(config.train['logs'] + 'steps/loss_Integrator')
            writer_loss_D_VIS_steps = SummaryWriter(config.train['logs'] + 'steps/loss_D_VIS')
            writer_loss_D_NIR_steps = SummaryWriter(config.train['logs'] + 'steps/loss_D_NIR')
            writer_loss_D_SKETCH_steps = SummaryWriter(config.train['logs'] + 'steps/loss_D_SKETCH')
            writer_loss_D_DOMAIN_steps = SummaryWriter(config.train['logs'] + 'steps/loss_D_DOMAIN')

            D_VIS.train()
            D_NIR.train()
            D_SKETCH.train()
            D_DOMAIN.train()
            Integrator.train()
            Style_Encoder.train()
            Intrinsic_Encoder.train()

            batch_time = AverageMeter()
            data_time = AverageMeter()
            # losses_G = AverageMeter()
            losses_Intrinsic = AverageMeter()
            losses_Style = AverageMeter()
            losses_Integrator = AverageMeter()
            losses_D_VIS = AverageMeter()
            losses_D_NIR = AverageMeter()
            losses_D_SKETCH = AverageMeter()
            losses_D_DOMAIN = AverageMeter()

            lr_G = LambdaLR(config.train['epochs'], config.train['lr_G_decay_rate'],
                            config.train['lr_G_decay_epoch'], config.train['lr_G']).step(epoch)
            lr_D = LambdaLR(config.train['epochs'], config.train['lr_D_decay_rate'],
                            config.train['lr_D_decay_epoch'], config.train['lr_D']).step(epoch)

            end_time = time.time()

            bar = Bar('Processing: ', max=len(trainLoader))

            for i, batch in enumerate(trainLoader):
                for k in batch:
                    batch[k] = batch[k].cuda()

                step = 0
                src_list = ['sketch', 'nir', 'vis']
                # for src in ['sketch', 'nir', 'vis']:
                if True:
                    src = choice(src_list)
                    if src == 'vis':
                        c_src = 0
                        D_X = D_VIS
                        optimizer_D_X = optimizer_D_VIS
                        src_buffer = vis_buffer
                        extra = 1.0
                        losses_D_X = losses_D_VIS
                        writer_loss_D_X_steps = writer_loss_D_VIS_steps

                    elif src == 'nir':
                        c_src = 1
                        D_X = D_NIR
                        optimizer_D_X = optimizer_D_NIR
                        src_buffer = nir_buffer
                        extra = 1.0
                        losses_D_X = losses_D_NIR
                        writer_loss_D_X_steps = writer_loss_D_NIR_steps

                    elif src == 'sketch':
                        c_src = 2
                        D_X = D_SKETCH
                        optimizer_D_X = optimizer_D_SKETCH
                        src_buffer = sketch_buffer
                        extra = 3.0
                        losses_D_X = losses_D_SKETCH
                        writer_loss_D_X_steps = writer_loss_D_SKETCH_steps

                    x_real = batch[src]
                    target_list = ['vis', 'sketch', 'nir']
                    while src in target_list:
                        target_list.remove(src)

                    # for tgt in target_list:
                    if True:
                        tgt = choice(target_list)
                        if tgt == 'vis':
                            c_tgt = 0
                            D_Y = D_VIS
                            optimizer_D_Y = optimizer_D_VIS
                            tgt_buffer = vis_buffer
                            losses_D_Y = losses_D_VIS
                            writer_loss_D_Y_steps = writer_loss_D_VIS_steps

                        elif tgt == 'nir':
                            c_tgt = 1
                            D_Y = D_NIR
                            optimizer_D_Y = optimizer_D_NIR
                            tgt_buffer = nir_buffer
                            losses_D_Y = losses_D_NIR
                            writer_loss_D_Y_steps = writer_loss_D_NIR_steps

                        elif tgt == 'sketch':
                            c_tgt = 2
                            D_Y = D_SKETCH
                            optimizer_D_Y = optimizer_D_SKETCH
                            tft_buffer = sketch_buffer
                            extra = 1.0
                            losses_D_Y = losses_D_SKETCH
                            writer_loss_D_Y_steps = writer_loss_D_SKETCH_steps

                        y_real = batch[tgt]

                        data_time.update(time.time() - end_time)

                        # =================================================================================== #
                        #                               3. Train the generator                                #
                        # =================================================================================== #
                        if (i + 0) % config.train['generator_steps'] == 0:
                            x_intrinsic = Intrinsic_Encoder(x_real)
                            y_intrinsic = Intrinsic_Encoder(y_real)

                            x_style, x_style_logit = Style_Encoder(x_real)
                            y_style, y_style_logit = Style_Encoder(y_real)

                            x_rec, x_rec_heatmap, x_cam_logit = Integrator(x_intrinsic, x_style)
                            y_rec, y_rec_heatmap, y_cam_logit = Integrator(y_intrinsic, y_style)

                            xy, xy_heatmap, xy_cam_logit = Integrator(x_intrinsic, y_style)
                            yx, yx_heatmap, yx_cam_logit = Integrator(y_intrinsic, x_style)

                            xy_intrinsic = Intrinsic_Encoder(xy)
                            yx_intrinsic = Intrinsic_Encoder(yx)
                            # yx_intrinsic = Intrinsic_Encoder(x_real)

                            xy_style, xy_style_logit = Style_Encoder(xy)
                            yx_style, yx_style_logit = Style_Encoder(yx)
                            # yx_style, yx_style_logit = Style_Encoder(x_real)

                            x_recon, x_recon_heatmap, x_recon_cam_logit = Integrator(xy_intrinsic, yx_style)
                            y_recon, y_recon_heatmap, y_recon_cam_logit = Integrator(yx_intrinsic, xy_style)

                            # x_recon, x_recon_heatmap, x_recon_cam_logit = Integrator(xy_intrinsic, x_style)
                            # y_recon, y_recon_heatmap, y_recon_cam_logit = Integrator(yx_intrinsic, y_style)

                            # intrinsic loss
                            g_intrinsic = mse_loss(x_intrinsic, xy_intrinsic) + mse_loss(y_intrinsic, yx_intrinsic)

                            # reconstruction loss
                            g_recon =l1_loss(x_real, x_rec) +l1_loss(y_real, y_rec)

                            # cycle consistency loss
                            g_cycle =l1_loss(x_real, x_recon) +l1_loss(y_real, y_recon)

                            # mmd loss
                            g_mmd = mmd(x_intrinsic, y_intrinsic)
                                    # + mmd(xy_intrinsic, yx_intrinsic) + mmd(x_intrinsic, yx_intrinsic) + mmd(y_intrinsic, xy_intrinsic)

                            # domain_cls1 = D_DOMAIN(x_intrinsic, y_intrinsic)
                            # domain_cls2 = D_DOMAIN(xy_intrinsic, yx_intrinsic)
                            # domain_cls3 = D_DOMAIN(x_intrinsic, yx_intrinsic)
                            # domain_cls4 = D_DOMAIN(y_intrinsic, xy_intrinsic)
                            # g_domain = mse_loss(domain_cls1, torch.ones_like(domain_cls1).cuda()) \
                            #            + mse_loss(domain_cls2, torch.ones_like(domain_cls2).cuda()) \
                            #            + mse_loss(domain_cls2, torch.ones_like(domain_cls2).cuda()) \
                            #            + mse_loss(domain_cls2, torch.ones_like(domain_cls2).cuda())

                            # adv loss & cam loss
                            # x_rec_adv, x_rec_adv_cam_logit, _ = D_X(x_rec)
                            # y_rec_adv, y_rec_adv_cam_logit, _ = D_Y(y_rec)

                            xy_adv, xy_adv_cam_logit, _ = D_Y(xy)
                            yx_adv, yx_adv_cam_logit, _ = D_X(yx)

                            g_adv = mse_loss(xy_adv, torch.ones_like(xy_adv).cuda()) + mse_loss(yx_adv, torch.ones_like(yx_adv).cuda())
                            g_cam = mse_loss(xy_adv_cam_logit, torch.ones_like(xy_adv_cam_logit).cuda()) \
                                    + mse_loss(yx_adv_cam_logit, torch.ones_like(yx_adv_cam_logit).cuda())

                            # style cls loss
                            x_cls_label = label2onehot(x_style_logit, c_src)
                            y_cls_label = label2onehot(y_style_logit, c_tgt)

                            g_cls = cross_entropy(x_style_logit, x_cls_label) + cross_entropy(y_style_logit, y_cls_label) \
                                    + cross_entropy(xy_style_logit, y_cls_label) * 0.0 + cross_entropy(yx_style_logit, x_cls_label) * 0.0

                            # total loss
                            # g_loss = g_recon * 10.0 * 255.0 + g_cycle * 10.0 * 255.0 + g_adv * 1.0 * 0.0 + g_cam * 10.0 * 0.0 + g_cls * 1.0 + g_domain * 0.0 + g_intrinsic * 1.0 * 0.0

                            # optimizer_G.zero_grad()
                            # g_loss.backward()
                            # optimizer_G.step()

                            intrinsic_loss = g_recon * 10.0 + g_cycle * 10.0 + g_adv * 0.1 + g_cam * 1.0 + g_mmd * 0.0 + g_intrinsic * 0.1
                            style_loss = g_cls
                            integrator_loss = g_recon * 10.0 + g_cycle * 10.0 + g_adv * 0.1 + g_cam * 1.0

                            optimizer_Intrinsic.zero_grad()
                            intrinsic_loss.backward(retain_graph=True)
                            optimizer_Intrinsic.step()

                            optimizer_Style.zero_grad()
                            style_loss.backward(retain_graph=True)
                            optimizer_Style.step()

                            optimizer_Integrator.zero_grad()
                            integrator_loss.backward()
                            optimizer_Integrator.step()

                        # =================================================================================== #
                        #                             2. Train the discriminator                              #
                        # =================================================================================== #
                        # x_rec_adv, x_rec_adv_cam_logit, _ = D_X(x_rec.detach())
                        # y_rec_adv, y_rec_adv_cam_logit, _ = D_Y(y_rec.detach())

                        xy_adv, xy_adv_cam_logit, _ = D_Y(xy.detach())
                        yx_adv, yx_adv_cam_logit, _ = D_X(yx.detach())

                        d_X_adv_fake = mse_loss(yx_adv, torch.zeros_like(yx_adv).cuda())

                        d_X_cam_fake = mse_loss(yx_adv_cam_logit, torch.zeros_like(yx_adv_cam_logit))

                        d_Y_adv_fake = mse_loss(xy_adv, torch.zeros_like(xy_adv).cuda())

                        d_Y_cam_fake = mse_loss(xy_adv_cam_logit, torch.zeros_like(xy_adv_cam_logit).cuda())

                        d_loss_x_fake = d_X_adv_fake * 1.0 + d_X_cam_fake * 10.0
                        d_loss_y_fake = d_Y_adv_fake * 1.0 + d_Y_cam_fake * 10.0

                        # Backward and optimize.
                        optimizer_D_X.zero_grad()
                        d_loss_x_fake.backward()
                        optimizer_D_X.step()

                        optimizer_D_Y.zero_grad()
                        d_loss_y_fake.backward()
                        optimizer_D_Y.step()

                        x_adv, x_adv_cam_logit, _ = D_X(x_real)
                        y_adv, y_adv_cam_logit, _ = D_Y(y_real)

                        d_X_adv_real = mse_loss(x_adv, torch.ones_like(x_adv).cuda())

                        d_X_cam_real = mse_loss(x_adv_cam_logit, torch.ones_like(x_adv_cam_logit))

                        d_Y_adv_real = mse_loss(y_adv, torch.ones_like(y_adv).cuda())

                        d_Y_cam_real = mse_loss(y_adv_cam_logit, torch.ones_like(y_adv_cam_logit))

                        d_loss_x_real = d_X_adv_real * 1.0 + d_X_cam_real * 10.0
                        d_loss_y_real = d_Y_adv_real * 1.0 + d_Y_cam_real * 10.0

                        # Backward and optimize.
                        optimizer_D_X.zero_grad()
                        d_loss_x_real.backward()
                        optimizer_D_X.step()

                        optimizer_D_Y.zero_grad()
                        d_loss_y_real.backward()
                        optimizer_D_Y.step()

                        d_loss_x = d_X_adv_real + d_X_adv_fake
                        d_loss_y = d_Y_adv_real + d_Y_adv_fake

                        '''
                        # =================================================================================== #
                        #                          3. Train the domain discriminator                          #
                        # =================================================================================== #
                        x_intrinsic = Intrinsic_Encoder(x_real)
                        y_intrinsic = Intrinsic_Encoder(y_real)

                        x_style, x_style_logit = Style_Encoder(x_real)
                        y_style, y_style_logit = Style_Encoder(y_real)

                        # x_rec, x_rec_heatmap, x_cam_logit = Integrator(x_intrinsic, x_style)
                        # y_rec, y_rec_heatmap, y_cam_logit = Integrator(y_intrinsic, y_style)

                        xy, xy_heatmap, xy_cam_logit = Integrator(x_intrinsic, y_style)
                        yx, yx_heatmap, yx_cam_logit = Integrator(y_intrinsic, x_style)

                        xy_intrinsic = Intrinsic_Encoder(xy)
                        yx_intrinsic = Intrinsic_Encoder(yx)

                        if i % 10 == 0:
                            domain_cls1 = D_DOMAIN(x_intrinsic, y_intrinsic)
                            domain_cls2 = D_DOMAIN(xy_intrinsic, yx_intrinsic)
                            domain_cls3 = D_DOMAIN(x_intrinsic, yx_intrinsic)
                            domain_cls4 = D_DOMAIN(y_intrinsic, xy_intrinsic)
                            d_domain = mse_loss(domain_cls1, torch.zeros_like(domain_cls1).cuda()) \
                                       + mse_loss(domain_cls2, torch.zeros_like(domain_cls2).cuda()) \
                                       + mse_loss(domain_cls2, torch.zeros_like(domain_cls2).cuda()) \
                                       + mse_loss(domain_cls2, torch.zeros_like(domain_cls2).cuda())
                            d_domain1 = d_domain * 1.0

                            optimizer_D_DOMAIN.zero_grad()
                            d_domain1.backward()
                            optimizer_D_DOMAIN.step()
                        '''

                        # =================================================================================== #
                        #                                 4. Miscellaneous                                    #
                        # =================================================================================== #
                        # losses_G.update(g_loss.data.cpu().numpy(), config.train['batch_size'])
                        losses_Intrinsic.update(intrinsic_loss.data.cpu().numpy(), config.train['batch_size'])
                        losses_Style.update(style_loss.data.cpu().numpy(), config.train['batch_size'])
                        losses_Integrator.update(integrator_loss.data.cpu().numpy(), config.train['batch_size'])
                        losses_D_X.update(d_loss_x.data.cpu().numpy(), config.train['batch_size'])
                        losses_D_Y.update(d_loss_y.data.cpu().numpy(), config.train['batch_size'])
                        # losses_D_DOMAIN.update(d_domain.data.cpu().numpy(), config.train['batch_size'])

                        # SummaryWriter
                        # writer_loss_G_steps.add_scalar('steps/loss_G', losses_G.avg, i)
                        writer_loss_Intrinsic_steps.add_scalar('steps/loss_Intrinsic', losses_Intrinsic.avg, i)
                        writer_loss_Style_steps.add_scalar('steps/loss_Style', losses_Style.avg, i)
                        writer_loss_Integrator_steps.add_scalar('steps/loss_G', losses_Integrator.avg, i)
                        writer_loss_D_X_steps.add_scalar('steps/loss_D_' + src, losses_D_X.avg, i)
                        writer_loss_D_Y_steps.add_scalar('steps/loss_D_' + tgt, losses_D_Y.avg, i)
                        writer_loss_D_DOMAIN_steps.add_scalar('steps/loss_D_DOMAIN', losses_D_DOMAIN.avg, i)

                        # Determine approximate time left
                        batch_time.update(time.time() - end_time)
                        end_time = time.time()
                        batches_done = epoch * len(trainLoader) + i
                        batches_left = config.train['epochs'] * len(trainLoader) - batches_done
                        time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))
                        prev_time = time.time()

                        bar.suffix = 'Epoch/Step: {epoch}/{step} | LR_G: {lr_G:.8f} | LR_D: {lr_D:.8f}' \
                                     '\n' \
                                     'Loss_intrinsic: {loss_intrinsic:.6f} | Loss_style: {loss_style:.6f} | Loss_integrator: {loss_integrator:.6f} | G_Recon: {g_recon:.6f} | G_Cyc: {g_cycle:.6f} | G_Adv: {g_adv:.6f} | G_Cam: {g_cam:.6f} ' \
                                     '| G_Cls: {g_cls:.6f} | G_Intrinsic: {g_intrinsic:.6f}' \
                                     '\n' \
                                     'Loss_D_VIS: {loss_D_VIS:.6f} | Loss_D_NIR: {loss_D_NIR:.6f} | Loss_D_SKETCH: {loss_D_SKETCH:.6f} | Loss_D_DOMAIN: {loss_D_DOMAIN:.6f} | ETA: {time_left}' \
                                     '\n' \
                                     '***************************************************************************'.format(
                            step=i,
                            epoch=epoch,
                            lr_G=lr_G,
                            lr_D=lr_D,
                            # loss_G=losses_G.avg,
                            loss_intrinsic=losses_Intrinsic.avg,
                            loss_style=losses_Style.avg,
                            loss_integrator=losses_Integrator.avg,
                            g_recon=g_recon.data.cpu().numpy(),
                            g_cycle=g_cycle.data.cpu().numpy(),
                            g_adv=g_adv.data.cpu().numpy(),
                            g_cam=g_cam.data.cpu().numpy(),
                            g_cls=g_cls.data.cpu().numpy(),
                            # g_domain=g_domain.data.cpu().numpy(),
                            g_intrinsic=g_intrinsic.data.cpu().numpy(),
                            loss_D_VIS=losses_D_VIS.avg,
                            loss_D_NIR=losses_D_NIR.avg,
                            loss_D_SKETCH=losses_D_SKETCH.avg,
                            loss_D_DOMAIN=g_mmd.data.cpu().numpy(),
                            time_left=time_left,
                        )
                        print(bar.suffix)
                        # step += 1

                        # Save Image
                        if i % count == 0:
                            fake_y_single = x_real.detach().cpu().numpy()[0]
                            fake_y_single = tensor2numpy(fake_y_single)
                            fake_y_single_name = 'real_{}_{}_{}.png'.format(src, epoch, i // count)
                            save_image_single(fake_y_single, config.train['out'] + fake_y_single_name)

                            fake_y_single = y_real.detach().cpu().numpy()[0]
                            fake_y_single = tensor2numpy(fake_y_single)
                            fake_y_single_name = 'real_{}_{}_{}.png'.format(tgt, epoch, i // count)
                            save_image_single(fake_y_single, config.train['out'] + fake_y_single_name)

                            fake_y_single = x_rec.detach().cpu().numpy()[0]
                            fake_y_single = tensor2numpy(fake_y_single)
                            fake_y_single_name = 'recon_{}_{}_{}.png'.format(src, epoch, i // count)
                            save_image_single(fake_y_single, config.train['out'] + fake_y_single_name)

                            fake_y_single = x_recon.detach().cpu().numpy()[0]
                            fake_y_single = tensor2numpy(fake_y_single)
                            fake_y_single_name = 'cycle_{}_{}_{}.png'.format(src, epoch, i // count)
                            save_image_single(fake_y_single, config.train['out'] + fake_y_single_name)

                            fake_y_single = xy.detach().cpu().numpy()[0]
                            fake_y_single = tensor2numpy(fake_y_single)
                            fake_y_single_name = '{}2{}_{}_{}.png'.format(src, tgt, epoch, i // count)
                            save_image_single(fake_y_single, config.train['out'] + fake_y_single_name)

                            # heatmap_single = xy_heatmap.detach().cpu().numpy()[0]
                            # heatmap_single = cam(tensor2numpy(heatmap_single), 112)
                            # heatmap_single = heatmap_single * 0.5 + fake_y_single * 0.5
                            # heatmap_single_name = '{}2{}_heatmap_{}_{}.png'.format(src, tgt, epoch, i // count)
                            # save_image_single(heatmap_single, config.train['out'] + heatmap_single_name)

                            fake_y_single = y_rec.detach().cpu().numpy()[0]
                            fake_y_single = tensor2numpy(fake_y_single)
                            fake_y_single_name = 'recon_{}_{}_{}.png'.format(tgt, epoch, i // count)
                            save_image_single(fake_y_single, config.train['out'] + fake_y_single_name)

                            fake_y_single = y_recon.detach().cpu().numpy()[0]
                            fake_y_single = tensor2numpy(fake_y_single)
                            fake_y_single_name = 'cycle_{}_{}_{}.png'.format(tgt, epoch, i // count)
                            save_image_single(fake_y_single, config.train['out'] + fake_y_single_name)

                            fake_y_single = yx.detach().cpu().numpy()[0]
                            fake_y_single = tensor2numpy(fake_y_single)
                            fake_y_single_name = '{}2{}_{}_{}.png'.format(tgt, src, epoch, i // count)
                            save_image_single(fake_y_single, config.train['out'] + fake_y_single_name)

                            # heatmap_single = yx_heatmap.detach().cpu().numpy()[0]
                            # heatmap_single = cam(tensor2numpy(heatmap_single), 112)
                            # heatmap_single = heatmap_single * 0.5 + fake_y_single * 0.5
                            # heatmap_single_name = '{}2{}_heatmap_{}_{}.png'.format(tgt, src, epoch, i // count)
                            # save_image_single(heatmap_single, config.train['out'] + heatmap_single_name)

            # SummaryWriter
            # writer_loss_G_epochs.add_scalar('epochs/loss_G', losses_G.avg, epoch)
            writer_loss_Integrator_epochs.add_scalar('epochs/loss_Integrator', losses_Integrator.avg, epoch)
            writer_loss_Intrinsic_epochs.add_scalar('epochs/loss_Intrinsic', losses_Intrinsic.avg, epoch)
            writer_loss_Style_epochs.add_scalar('epochs/loss_Style', losses_Style.avg, epoch)
            writer_loss_D_SKETCH_epochs.add_scalar('epochs/loss_D_SKETCH', losses_D_SKETCH.avg, epoch)
            writer_loss_D_VIS_epochs.add_scalar('epochs/loss_D_VIS', losses_D_VIS.avg, epoch)
            writer_loss_D_NIR_epochs.add_scalar('epochs/loss_D_NIR', losses_D_NIR.avg, epoch)
            writer_loss_D_DOMAIN_epochs.add_scalar('epochs/loss_D_DOMAIN', losses_D_DOMAIN.avg, epoch)

            # lr_schedule_G.step()
            lr_schedule_Intrinsic.step()
            lr_schedule_Style.step()
            lr_schedule_Integrator.step()
            lr_schedule_D_VIS.step()
            lr_schedule_D_NIR.step()
            lr_schedule_D_SKETCH.step()
            lr_schedule_D_DOMAIN.step()
            if epoch % 10 == 0:
                date = '20190924'

                torch.save({
                    'state_dict_Intrinsic_Encoder': Intrinsic_Encoder.state_dict(),
                    'epoch_Intrinsic_Encoder': epoch,
                }, os.path.join(config.train['checkpoint'], 'Intrinsic_Encoder_' + date + '.pth'))

                torch.save({
                    'state_dict_Style_Encoder': Style_Encoder.state_dict(),
                    'epoch_Style_Encoder': epoch,
                }, os.path.join(config.train['checkpoint'], 'Style_Encoder_' + date + '.pth'))
                torch.save({
                    'state_dict_Integrator': Integrator.state_dict(),
                    'epoch_Integrator': epoch,
                }, os.path.join(config.train['checkpoint'], 'Integrator_' + date + '.pth'))
                torch.save({
                    'state_dict_D_VIS': D_VIS.state_dict(),
                    'epoch_D_VIS': epoch
                }, os.path.join(config.train['checkpoint'], 'D_VIS_' + date + '.pth'))
                torch.save({
                    'state_dict_D_NIR': D_NIR.state_dict(),
                    'epoch_D_NIR': epoch
                }, os.path.join(config.train['checkpoint'], 'D_NIR_' + date + '.pth'))
                torch.save({
                    'state_dict_D_SKETCH': D_SKETCH.state_dict(),
                    'epoch_D_SKETCH': epoch
                }, os.path.join(config.train['checkpoint'], 'D_SKETCH_' + date + '.pth'))
                # torch.save({
                #     'state_dict_D_DOMAIN': D_DOMAIN.state_dict(),
                #     'epoch_D_DOMAIN': epoch
                # }, os.path.join(config.train['checkpoint'], 'D_DOMAIN_' + date + '.pth'))
                # torch.save(optimizer_G.state_dict(),
                #            os.path.join(config.train['checkpoint'], 'optimizer_G_' + date + '.pth'))
                torch.save(optimizer_Intrinsic.state_dict(),
                           os.path.join(config.train['checkpoint'], 'optimizer_Intrinsic_' + date + '.pth'))
                torch.save(optimizer_Style.state_dict(),
                           os.path.join(config.train['checkpoint'], 'optimizer_Style_' + date + '.pth'))
                torch.save(optimizer_Integrator.state_dict(),
                           os.path.join(config.train['checkpoint'], 'optimizer_Integrator_' + date + '.pth'))
                torch.save(optimizer_D_VIS.state_dict(),
                           os.path.join(config.train['checkpoint'], 'optimizer_D_VIS_' + date + '.pth'))
                torch.save(optimizer_D_NIR.state_dict(),
                           os.path.join(config.train['checkpoint'], 'optimizer_D_NIR_' + date + '.pth'))
                torch.save(optimizer_D_SKETCH.state_dict(),
                           os.path.join(config.train['checkpoint'], 'optimizer_D_SKETCH_' + date + '.pth'))
                # torch.save(optimizer_D_DOMAIN.state_dict(),
                #            os.path.join(config.train['checkpoint'], 'optimizer_D_DOMAIN_' + date + '.pth'))


from . import api as api_bp
from flask import g, request, jsonify, current_app
import time
import requests
import base64
import random

from servelite.exceptions import std_res
from servelite.auth import verify, signature
from servelite.Reger.utils import JWT_required
from servelite.Reger.models import Servers, session_factory, API
from servelite.utils import *

@api_bp.route('/heartbeat')
@JWT_required
def heartbeat():
    """
    接受心跳包的接口
    :return:
    """
    session = session_factory()
    server = current_app.reger.get_server(g.payload['hash'])
    if server is None:
        return std_res('400', 'server_test hash is not registered')
    server.check_time = time.time()
    session.commit()
    return std_res(200, 'success')

@api_bp.route('/register-service', methods=['POST'])
def register_service():
    """
    服务注册接口
    该接口接受前端的服务注册请求，向服务发送确认请求。服务返回公钥的签名字符串以确认身份。
    得到确认后，该服务进入初始化队列，由服务管理进程完成对服务各项数据的初始化。

    note: 先实现同步版本
    :return:
    """
    name = request.json.get('name') or request.values.get('name')
    auth = request.json.get('auth') or request.values.get('auth')
    description = request.json.get('description') or request.values.get('description')
    host = request.json.get('host') or request.values.get('host')
    port = request.json.get('port') or request.values.get('port')
    # 值校验在前端完成
    url = 'http://' + host + ':' + str(port)

    # 服务确认
    try:
        # todo: 使用统一配置文件代替hard code
        data = {
            'pubkey': current_app.reger.config['pubkey'],
            'port': current_app.reger.config['port']
        }
        r = requests.post(url + '/communicate/confirm', params=data)
    except requests.exceptions.Timeout:
        return std_res(400, 'confirm time out')
    except requests.exceptions.ConnectionError:
        return std_res(400, 'can not connect to server')

    if r.status_code != 200:
        return std_res(400, 'confirm request respond with code {}'.format(r.status_code))

    r = r.json()
    sign = r.get('signature')
    pubkey = r.get('pubkey')
    if sign is None:
        return std_res(400, 'confirm request should respond with signature but nothing now')
    if not verify(pubkey, sign, pubkey, decode_base64=True):
        return std_res(403, 'signature should be decoded to pubkey. Check if your pubkey is right!')
    serve_hash = md5(pubkey)
    # 存在性
    server = g.session.query(Servers).filter_by(hash=serve_hash).first()
    # 重名检查
    while server is None \
            and g.session.query(Servers).filter_by(name=name).first():
        name = name + '_' + str(random.randint(1000, 9999))
    if not server:
        server = Servers(
            name=name,
            hash=serve_hash,
            server_key=md5(pubkey[:4] + str(time.time())),
            status=1,
        )
        server.update_auth_by_name(auth)
        g.session.add(server)
    server.description = description
    server.host = host
    server.port = port
    server.pubkey = pubkey
    try:
        g.session.commit()
    except Exception as e:
        return std_res(400, 'commit failed', error=str(e))

    # 服务初始化
    if server.status == 2:
        return std_res(400, 'already complete')
    try:
        with open(current_app.reger.config['PRIVKEY_FILE'], 'r') as f:
            privkey = f.read()
    except FileNotFoundError:
        return std_res(400, 'rsa file is not found')
    try:
        init_data = [
            {'pos': 'reg_api.' + api_name, 'data': rule} for api_name, rule in current_app.reger.config['API'].items()
        ]
        init_data.append({'pos': 'hash', 'data':server.hash})
        data = {'signature': decode_bytes(base64.b64encode(signature(current_app.reger.config['pubkey'], privkey))),
                'data': init_data}
        r = requests.post(url + '/communicate/update', json=data)
    except requests.exceptions.Timeout:
        return std_res(400, 'service info update time out')
    if r.json()['code'] == 200:
        server.status = 2
        try:
            g.session.commit()
        except Exception as e:
            return std_res(400, 'commit failed', error=str(e))
        return std_res(200, 'register complete')
    else:
        return std_res(400, 'info update failed with msg: {}'.format(r.json()['msg']), error=r.json().get('error'))

@api_bp.route('/register-api', methods=['POST', 'GET'])
@JWT_required
def register_api():
    """
    api注册接口
    该接口接受API注册请求。
    :return:
    """
    session = session_factory()

    serve_hash = g.payload['hash']
    hash = request.values.get('hash', md5(serve_hash + time.time()))
    rule = request.values.get('rule', '/')

    # 存在性
    api = session.query(API).filter_by(hash=hash).first()
    if api is None:
        api = API(
            hash=hash,
            server_hash=serve_hash,
            rule=rule,
        )
        session.add(api)
    api.name = request.values.get('name', hash)
    api.description = request.values.get('description', '')
    api.params = request.values.get('params', '')
    api.res_example = request.values.get('res_example', '')
    api.sensitivity = int(request.values.get('sensitivity', 5))
    api.note = request.values.get('note', '')
    api.status = int(request.values.get('status', 1))
    try:
        session.commit()
    except Exception as e:
        return std_res(400, 'commit failed', error=str(e))
    return std_res(200, 'success', api_id=api.id)

@api_bp.route('/info/<key>')
def info(key):
    """
    信息查询接口
    :param key:
    :return:
    """
    value = current_app.reger.info.find(key)
    if value is None:
        return std_res(400, '{} is not found'.format(key))
    return jsonify(value)

@api_bp.route('/servers')
def servers():
    result = [row2dict(s, ['server_key', 'auth'],
                       author=s.author.user_name) for s in g.session.query(Servers)]
    return std_res(200, 'success', result=result)

@api_bp.route('/apis')
def get_apis():
    session = session_factory()

    result = []
    for server in session.query(Servers):  # type: Servers
        apis_of_server = {
            'server_name': server.name,
            'apis': []
        }
        for api in server.apis:
            apis_of_server['apis'].append(row2dict(api))
        result.append(apis_of_server)
    return std_res(200, 'success', result=result)


import argparse
import logging

import numpy as np
import torch
from tensorboardX import SummaryWriter
from torch.nn import functional as F

import datasets.datasetfactory as df
import datasets.task_sampler as ts
import model.modelfactory as mf
from experiment.experiment import experiment
from model.meta_learner import MetaLearingClassification

logger = logging.getLogger('experiment')

def main(args):
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    np.random.seed(args.seed)
    my_experiment = experiment(args.name, args, "../results/", commit_changes=args.commit)
    writer = SummaryWriter(my_experiment.path + "tensorboard")

    logger = logging.getLogger('experiment')

    args.classes = list(range(963))

    dataset = df.DatasetFactory.get_dataset(args.dataset, background=True, train=True)
    dataset_test = df.DatasetFactory.get_dataset(args.dataset, background=True, train=False)

    iterator_test = torch.utils.data.DataLoader(dataset_test, batch_size=5,
                                                shuffle=True, num_workers=1)

    iterator_train = torch.utils.data.DataLoader(dataset, batch_size=5,
                                                 shuffle=True, num_workers=1)

    logger.info("Train set length = %d", len(iterator_train) * 5)
    logger.info("Test set length = %d", len(iterator_test) * 5)
    sampler = ts.SamplerFactory.get_sampler(args.dataset, args.classes, dataset, dataset_test)

    config = mf.ModelFactory.get_model("na", args.dataset)

    if torch.cuda.is_available():
        device = torch.device('cuda')
    else:
        device = torch.device('cpu')

    maml = MetaLearingClassification(args, config).to(device)

    for name, param in maml.named_parameters():
        param.learn = True
    for name, param in maml.net.named_parameters():
        param.learn = True

    frozen_layers = []
    for temp in range(args.rln * 2):
        frozen_layers.append("net.vars." + str(temp))

    for name, param in maml.named_parameters():
        logger.info(name)
        if name in frozen_layers:
            logger.info("Freeezing name %s", str(name))
            param.learn = False

    # Update the classifier
    list_of_params = list(filter(lambda x: x.learn, maml.parameters()))
    list_of_names = list(filter(lambda x: x[1].learn, maml.named_parameters()))

    for a in list_of_names:
        logger.info("Unfrozen layers for rep learning = %s", a[0])

    for step in range(args.epoch):

        t1 = np.random.choice(args.classes, np.random.randint(1, args.tasks + 1), replace=False)

        d_traj_iterators = []
        for t in t1:
            d_traj_iterators.append(sampler.sample_task([t]))

        d_rand_iterator = sampler.get_complete_iterator()

        x_spt, y_spt, x_qry, y_qry = maml.sample_training_data(d_traj_iterators, d_rand_iterator,
                                                               steps=args.update_step)
        if torch.cuda.is_available():
            x_spt, y_spt, x_qry, y_qry = x_spt.cuda(), y_spt.cuda(), x_qry.cuda(), y_qry.cuda()

        accs, loss = maml(x_spt, y_spt, x_qry, y_qry)

        if step % 40 == 0:
            writer.add_scalar('/metatrain/train/accuracy', accs[-1], step)
            logger.info('step: %d \t training acc %s', step, str(accs))
        if step % 300 == 0:
            correct = 0
            torch.save(maml.net, my_experiment.path + "learner.model")
            for img, target in iterator_test:
                with torch.no_grad():
                    img = img.to(device)
                    target = target.to(device)
                    logits_q = maml.net(img, vars=None, bn_training=False, feature=False)
                    pred_q = F.softmax(logits_q, dim=1).argmax(dim=1)
                    correct += torch.eq(pred_q, target).sum().item() / len(img)
            writer.add_scalar('/metatrain/test/classifier/accuracy', correct / len(iterator_test), step)
            logger.info("Test Accuracy = %s", str(correct / len(iterator_test)))
            correct = 0
            for img, target in iterator_train:
                with torch.no_grad():
                    img = img.to(device)
                    target = target.to(device)
                    logits_q = maml.net(img, vars=None, bn_training=False, feature=False)
                    pred_q = (logits_q).argmax(dim=1)
                    correct += torch.eq(pred_q, target).sum().item() / len(img)

            logger.info("Train Accuracy = %s", str(correct / len(iterator_train)))
            writer.add_scalar('/metatrain/train/classifier/accuracy', correct / len(iterator_train), step)

if __name__ == '__main__':
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--epoch', type=int, help='epoch number', default=40000)
    argparser.add_argument('--seed', type=int, help='Seed for random', default=10000)
    argparser.add_argument('--seeds', type=int, nargs='+', help='n way', default=[10])
    argparser.add_argument('--tasks', type=int, help='meta batch size, namely task num', default=1)
    argparser.add_argument('--meta_lr', type=float, help='meta-level outer learning rate', default=1e-4)
    argparser.add_argument('--update_lr', type=float, help='task-level inner update learning rate', default=0.01)
    argparser.add_argument('--update_step', type=int, help='task-level inner update steps', default=10)
    argparser.add_argument('--name', help='Name of experiment', default="mrcl_classification")
    argparser.add_argument('--dataset', help='Name of experiment', default="omniglot")
    argparser.add_argument("--commit", action="store_true")
    argparser.add_argument("--rln", type=int, default=6)
    args = argparser.parse_args()

    args.name = "/".join([args.dataset, str(args.meta_lr).replace(".", "_"), args.name])
    print(args)
    main(args)


import argparse
import os
import time
from torch.backends import cudnn

from data import *
from trainer import *
import _pickle as pickle

def main(args):
	cuda.set_device(int(args.gpu_num))
	cudnn.benchmark = True

	start_time = time.time()

	# if os.path.isfile(args.data_path + 'vocab.pb'):
	# 	vocab = pickle.load(open(args.data_path, 'rb'))
	# else:
	# vocab = Vocab()
	# vocab.create_vocab_for_translation(args.train_path + 'train_all')
	# pickle.dump(vocab, open(args.data_path + 'vocab.pb', 'wb'))

	# Load dataset
	if os.path.isfile(args.data_path + 'prepro.pb'):
		dp = pickle.load(open(args.data_path + 'prepro.pb', 'rb'))
	else:
		dp = Preprocess(args.train_path + 'train.src', args.train_path + 'train.tgt', args.val_path + 'val.src', args.val_path + 'val.tgt', args.train_path + 'train_all')
		dp.preprocess_train()
		dp.preprocess_val()
		pickle.dump(dp, open(args.data_path + 'prepro.pb', 'wb'))

	print ("Elapsed Time: %1.3f \n"  %(time.time() - start_time))

	# print(dp.vocab.itos)
	# print(dp.train_src[0])
	# print(dp.train_tgt[0])
	# training_dataset = list(zip(dp.train_src, dp.train_tgt, dp.train_srclng, dp.train_tgtlng))
	# import random; random.seed(1234);
	# random.shuffle(training_dataset)
	MAX_SAMPLES = 40040
	# dp.train_src[:], dp.train_tgt[:], dp.train_srclng[:], dp.train_tgtlng[:] = zip(*training_dataset)
	dp.train_src, dp.train_tgt, dp.train_srclng, dp.train_tgtlng = dp.train_src[:MAX_SAMPLES], dp.train_tgt[:MAX_SAMPLES], dp.train_srclng[:MAX_SAMPLES], dp.train_tgtlng[:MAX_SAMPLES]

	print(dp.vocab.itos)
	print(dp.train_src[0])
	print(dp.train_tgt[0])
	print ("=========== Data Stat ===========")
	print ("Train: ", len(dp.train_src))
	print ("val: ", len(dp.val_tgt))
	print ("val: ", len(set(dp.train_src_sentlen)))
	print ("=================================")

	trainer = Trainer(dp, args)
	trainer.train()

if __name__ == '__main__':

	parser = argparse.ArgumentParser()

	# Language setting
	parser.add_argument('--dataset', type=str, default='europarl')
	parser.add_argument('--src-lang', action="append", type=str, default=[], dest='src_lang',
	                    help='list of type of lexical relations to capture. Options | syn | hyp | mer')
	parser.add_argument('--trg_lang', type=str, default='en')
	parser.add_argument('--mode', type=str, default='probe', help='probe|train')
	parser.add_argument('--max_len', type=int, default=20)
	parser.add_argument('--cuda', action='store_true',
	                    help='use CUDA')

	parser.add_argument('--gpu', type=int, default=0,
	                    help='use gpu x')
	parser.add_argument('--log-interval', type=int, default=100)
	# Model hyper-parameters
	parser.add_argument('--lr', type=float, default=0.001)
	parser.add_argument('--grad_clip', type=float, default=1)
	parser.add_argument('--num_layer', type=int, default=2)
	parser.add_argument('--embed_dim',  type=int, default=256)
	parser.add_argument('--hidden_dim', type=int, default=512)

	# Training setting
	parser.add_argument('--batch_size', type=int, default=32)
	parser.add_argument('--num_epoch', type=int, default=100)

	# Path
	parser.add_argument('--data_path', type=str, default='./data/')
	parser.add_argument('--train_path', type=str, default='./data/training/')
	parser.add_argument('--val_path', type=str, default='./data/dev/')

	# Dir.
	parser.add_argument('--log', type=str, default='log')
	parser.add_argument('--sample', type=str, default='sample')

	# Misc.
	parser.add_argument('--gpu_num', type=int, default=0)

	args = parser.parse_args()
	print (args)
	main(args)


from datetime import datetime, timedelta

def to_datetime(now, string):
    if len(string) != 6:
        return None

    # if any field is '*' then make it equal to current time
    if string[0] == '*':
        secs = now.second
    else:
        secs = string[0]
    if string[1] == '*':
        mins = now.minute
    else:
        mins = string[1]
    if string[2] == '*':
        hours = now.hour
    else:
        hours = string[2]
    if string[3] == '*':
        dom = now.day
    else:
        dom = string[3]
    if string[4] == '*':
        mon = now.month
    else:
        mon = string[4]
    if string[5] == '*':
        dow = now.weekday()
    else:
        dow = string[5]

    return (datetime(now.year,mon,dom,hours,mins,secs))

def crontime_in_range(now, start, end):
    dt_start = to_datetime(now, start)
    dt_end = to_datetime(now, end,)
    if dt_start == None or dt_end == None:
        return False

    # take care of weekday
    dw_start = start[5]
    dw_end = end[5]
    if dw_start != '*' and dw_end != '*':
        if dw_start == dw_end != now.weekday():
            return False
        elif dw_start > dw_end:
            dt_end = dt_end + timedelta(weeks=1)
            if now.weekday() < dw_start:
                now = now + timedelta(weeks=1)
        else:
            if dw_start <= now.weekday() <= dw_end:
                dt_end = dt_end + timedelta(days=dw_end - dw_start)

    if dt_start > dt_end:
        dt_end = dt_end + timedelta(days=365)
        if now < dt_start:
            now = now + timedelta(days=365)

    if dt_start <= now <= dt_end:
        return True

    return False


from pyface.qt import QtGui, QtCore
#%matplotlib qt
import mne
import glob
import sys
import numpy as np
import os.path as op
import pandas as pd
# import matplotlib
# matplotlib.use('Qt4Agg')
import matplotlib.pyplot as plt
from mayavi import mlab
from scipy.spatial.distance import euclidean
from intra_extra_fx import find_stim_coords, get_stim_params, plot_source_space, save_results
from intra_extra_info import study_path, subjects_dir
import re
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.patches as mpatches
import nibabel as nib
pd.set_option('display.expand_frame_repr', False)

def make_fwd_solution(cond_fname, fs_subj, study_path):
    eeg_epo = mne.read_epochs(cond_fname)
    img_type = 'anony' if fs_subj.find('_an') > 0 else 'orig'
    subj = fs_subj.strip('_an')

    trans_fname = op.join(study_path, 'source_stim', subj, 'source_files', img_type,'%s_fid-trans.fif' % fs_subj)
    src_fname = op.join(study_path, 'source_stim', subj, 'source_files', img_type, '%s-oct5-src.fif' % fs_subj)
    bem_fname = op.join(subjects_dir, fs_subj, 'bem', '%s-bem-sol.fif' % fs_subj)

    fwd = mne.make_forward_solution(eeg_epo.info, trans_fname, src_fname, bem_fname)

    trans = mne.read_trans(trans_fname)
    mne.viz.plot_alignment(eeg_epo.info, trans, subject=subj, surfaces=['head', 'brain'], subjects_dir=subjects_dir)
    # mlab.show()

    mne.write_forward_solution(op.join(study_path, 'source_stim', subj, 'source_files', img_type,
                                       '%s-fwd.fif' % fs_subj), fwd, overwrite=True)
    return fwd

def source_loc(evoked, cov, fwd, subj, img_type, study_path, plot=False):
    cond = evoked.info['description']
    stim_coords = find_stim_coords(cond, subj, study_path)
    stim_params = get_stim_params(cond)

    inv = mne.minimum_norm.make_inverse_operator(evoked.info, fwd, cov, loose=0.2, depth=0.8)
    method = "dSPM"
    snr = 300.
    lambda2 = 1. / snr ** 2
    stc = mne.minimum_norm.apply_inverse(evoked, inv, lambda2, method=method, pick_ori=None)
    stc.save(op.join(study_path, 'source_stim', subj, 'source_files', img_type,'stc', '%s_%s' %(subj, cond)))

    # plot_source_space(subj, study_path, subjects_dir)

    hemi = 'lh' if stim_params['is_left'] else 'rh'
    contact_nr = int(re.findall('\d+', stim_params['ch'])[0])
    view = 'medial' if contact_nr < 5 else 'lateral'  # improve view selection

    vertno_max, time_max = stc.get_peak(hemi=hemi)
    stc_max = np.max(stc.data)

    if plot:
        brain_inf = stc.plot(surface='inflated', hemi=hemi, subjects_dir=subjects_dir, subject=fs_subj,
                             clim=dict(kind='value', lims=[0, stc_max*0.75, stc_max]),
                             initial_time=time_max, time_unit='s', alpha=1, background='w', foreground='k')

        brain_inf.add_foci(vertno_max, coords_as_verts=True, hemi=hemi, color='blue',
                           scale_factor=0.8)

        brain_inf.add_foci(stim_coords['surf_ori'], map_surface='pial', hemi=hemi, color='green', scale_factor=0.8)

        brain_inf.show_view(view)  # mlab.view(130, 90)
        fig_fname = op.join(study_path, 'source_stim', subj, 'figures', 'distributed', '%s_inf_foci.eps' % cond)
        #brain_inf.save_image(fig_fname)
        mlab.show()
        #mlab.clf()

    max_mni = mne.vertex_to_mni(vertno_max, hemis=0, subject=subj, subjects_dir=subjects_dir)  # mni coords of vertex (ojo vertex resolution ico-5)
    stim_mni = stim_coords['mni']  # mni coords of stimulation
    dist_mni = euclidean(max_mni, stim_mni)

    print(dist_mni)
    return dist_mni

    # brain_pial = stc.plot(surface='pial', hemi=hemi, subjects_dir=subjects_dir,
    #                      clim=dict(kind='value', lims=[0, stc_max*0.75, stc_max]),
    #                      initial_time=time_max, time_unit='s', alpha=1, background='w', foreground='k')
    #
    # brain_pial.add_foci(vertno_max, coords_as_verts=True, hemi=hemi, color='blue',
    #                    scale_factor=1)
    #
    # brain_pial.add_foci(stim_coords['surf'], hemi='lh', color='green', scale_factor=1)
    #
    # brain_pial.show_view('medial')  # mlab.view(130, 90)
    # fig_fname = op.join(study_path, 'source_stim', subj, 'figures', 'distributed', '%s_pial_foci.eps' % cond)
    # brain_pial.save_image(fig_fname)
    # mlab.show()

def vol_source_loc(evoked, cov, fs_subj, study_path, plot=True):
    cond = evoked.info['description']

    bem_fname = op.join(subjects_dir, fs_subj, 'bem', '%s-bem-sol.fif' % fs_subj)
    trans_fname = op.join(study_path, 'source_stim', subj, 'source_files', img_type, '%s_static-trans.fif' % fs_subj)

    mri_file = op.join(subjects_dir, subj, 'mri', 'T1.mgz')
    v_source = mne.setup_volume_source_space(subject=fs_subj, mri=mri_file, bem=bem_fname, subjects_dir=subjects_dir,
                                             pos=2.)

    fwd = mne.make_forward_solution(evoked.info, trans_fname, v_source, bem_fname,
                                mindist=5.0,  # ignore sources<=5mm from innerskull
                                meg=False, eeg=True,
                                n_jobs=2)

    mne.write_forward_solution(op.join(study_path, 'source_stim', subj, 'source_files', '%s_vol_source_space_2mm'),
                               fwd)

    inv = mne.minimum_norm.make_inverse_operator(evoked.info, fwd, cov, loose=1, depth=1)
    method = "sLORETA"
    snr = 30.
    lambda2 = 1. / snr ** 2
    stc = mne.minimum_norm.apply_inverse(evoked, inv, lambda2, method=method, pick_ori=None)

    stc.crop(-0.005, 0.005)

    from nilearn.plotting import plot_stat_map
    from nilearn.image import index_img
    import nilearn as ni

    nii_fname = op.join(study_path, 'source_stim', subj, 'source_files', img_type, 'stc', '%s_%s_stc.nii' % (subj, cond))
    img = mne.save_stc_as_volume(nii_fname, stc.copy().crop(-0.003, 0.003), fwd['src'], dest='mri')

    t0_img = index_img(img, 3)
    act_coord = ni.plotting.find_xyz_cut_coords(t0_img)
    #act_coord = [57, 122, 87]
    thr = np.percentile(t0_img.get_data(), 99.9)
    st_map = plot_stat_map(t0_img, mri_file, threshold=thr, display_mode='ortho', cmap=plt.cm.plasma)
    st_map.add_markers(np.reshape(act_coord, (-1, 3)), marker_color='g', marker_size=50)
    fname_fig = op.join(study_path, 'source_stim', subj, 'figures', 'volume', '%s_%s_st_map.pdf' % (subj, cond))
    st_map.savefig(fname_fig)

    dist_vol = euclidean()

def dipole_source_loc(evoked, cov,  fs_subj, study_path, plot=False):
    evo_crop = evoked.copy().crop(-0.005, 0.005)
    subj = fs_subj.strip('_an') if fs_subj.find('_an') > 0 else fs_subj
    img_type = 'anony' if fs_subj.find('_an') > 0 else 'orig'

    seeg_ch_info = pd.read_csv(op.join(study_path, 'physio_data', subj, 'chan_info', '%s_seeg_ch_info.csv' % subj))

    # ori_to_an_fname = op.join(study_path, 'freesurfer_subjects', subj, 'mri', 'ori_to_an_trans.lta')  # mri to mri
    # ori_to_an_trans = np.genfromtxt(ori_to_an_fname, skip_header=8, skip_footer=18)

    trans_fname = op.join(study_path, 'source_stim', subj, 'source_files', img_type, '%s_static-trans.fif' % fs_subj)
    bem_fname = op.join(subjects_dir, fs_subj, 'bem', '%s-bem-sol.fif' % fs_subj)

    cond = eeg_epo.info['description']
    stim_coords = find_stim_coords(cond, subj, study_path)

    # Fit a dipole
    dip = mne.fit_dipole(evo_crop, cov, bem_fname, trans_fname, min_dist=5, n_jobs=2)[0]

    from mne.beamformer import rap_music
    dip = rap_music(evo_crop, fwd, cov, n_dipoles=1, return_residual=True, verbose=True)[0][0]

    import mne.transforms as tra
    from scipy import linalg
    trans = mne.read_trans(trans_fname)

    surf_to_head = linalg.inv(trans['trans'])

    # stim_point = tra.apply_trans(surf_to_head, stim_coords['surf']/1e3)
    stim_point = stim_coords['surf_ori']  # get point for plot in mm
    #dip.pos[np.argmax(dip.gof)] = tra.apply_trans(surf_to_head, stim_coords['surf_ori']/1e3)  # check stim loc (apply affine in m)

    #stim_point = tra.apply_trans(ori_to_an_trans, stim_point)

    dip_surf = tra.apply_trans(trans['trans'], dip.pos[np.argmax(dip.gof)]) * 1e3  # transform from head to surface
    dist_surf = euclidean(dip_surf, stim_point)  # compute distance in surface space
    print(dist_surf)

    if plot:
        # Plot the result in 3D brain with the MRI image.
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        dip.plot_locations(trans_fname, fs_subj, subjects_dir, mode='orthoview', coord_frame='mri', ax=ax, show_all=True,
                           idx='gof')
        ax.scatter(stim_point[0], stim_point[1], stim_point[2])
        ax.plot([stim_point[0], -128], [stim_point[1], stim_point[1]], [stim_point[2], stim_point[2]], color='g')
        ax.plot([stim_point[0], stim_point[0]], [stim_point[1], -128], [stim_point[2], stim_point[2]], color='g')
        ax.plot([stim_point[0], stim_point[0]], [stim_point[1], stim_point[1]], [stim_point[2], -128], color='g')
        ax.text2D(0.05, 0.90, 'distance: %i mm \nstim coords = %0.1f %0.1f %0.1f' % (dist_surf, stim_point[0], stim_point[1], stim_point[2]),
                  transform=ax.transAxes)
        red_patch = mpatches.Patch(color='red')
        green_patch = mpatches.Patch(color='green')
        fig.legend(handles=[red_patch, green_patch], labels=['dipole', 'electrode'])

        fig.savefig(op.join(study_path, 'source_stim', subj, 'figures', 'dipole', '%s_dip_15mm.png' % cond))
        plt.close()

        from mne.viz import plot_dipole_locations, plot_dipole_amplitudes
        plot_dipole_amplitudes(dip)
        plot_dipole_locations(dip, trans, subj, subjects_dir=subjects_dir)

    return dist_surf

if __name__ == '__main__':
    subj = 'S5'
    img_type = 'orig'
    fs_subj = subj + '_an' if img_type == 'anony' else subj

    epo_path = op.join(study_path, 'source_stim', subj, 'epochs', 'fif')
    conds = glob.glob(epo_path + '/*-epo.fif')

    fwd_fname = op.join(study_path, 'source_stim', subj, 'source_files', img_type, '%s-fwd.fif' % fs_subj)

    if not op.isfile(fwd_fname):
        fwd = make_fwd_solution(conds[0], fs_subj, study_path)
    else:
        fwd = mne.read_forward_solution(fwd_fname)

    for cond_fname in conds:
        eeg_epo = mne.read_epochs(cond_fname)
        eeg_epo.interpolate_bads(reset_bads=True)
        # eeg_epo.filter(None, 40)
        cov = mne.compute_covariance(eeg_epo, method='shrunk', tmin=-0.5, tmax=-0.3)  # use method='auto' for final computation
        evoked = eeg_epo.average()

        evoked.plot(spatial_colors=True)
        mne.viz.plot_evoked_topo(evoked)
        evoked.plot_topomap(np.linspace(-0.005, 0.005, 11))

        dist_dis = source_loc(evoked, cov, fwd, subj, img_type, study_path, plot=False)
        dist_dip = dipole_source_loc(evoked, cov,  fs_subj, study_path, plot=True)

        cond = eeg_epo.info['description']
        stim_params = get_stim_params(cond)

        save_results(subj, study_path, stim_params, evoked, dist_dis, dist_dip, img_type)


# coding=utf-8
import pandas as pd
import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np

data = pd.read_csv('/home/drproduck/Downloads/ibm.csv')
for c in data.columns:
    s = {x for x in data[c]}
    print(c, len(s))
data.Attrition = [1 if x == 'Yes' else 0 for x in data.Attrition]
print(data.columns)
data.drop(['Over18', 'EmployeeCount', 'StandardHours'], axis=1, inplace=True)
for c in data.columns:
    if all(isinstance(n, str) for n in data[c]):
        s = {x for x in data[c]}
        t = list(s)
        dic = {}
        for z in range(len(s)):
            dic[t[z]] = z
        print(dic)
        data[c] = [dic[x] for x in data[c]]

plt.figure(1)
plt.hist([data.BusinessTravel[data.Attrition==1],data.BusinessTravel[data.Attrition==0]], color=['green','red'], label=['quit', 'stay'])

matrix = data.as_matrix()

attrition = matrix[:,1:2]
matrix = np.append(matrix, attrition, axis=1)
matrix = np.delete(matrix, 1, axis=1)
print(np.shape(matrix))
import collections
c = collections.Counter()
for x in matrix[:,31]:
    c[x] += 1
print(c)

np.save('matrix_weighted', matrix)


# !usr/bin/env python
# -*- coding: utf-8 -*-
#
# Licensed under a 3-clause BSD license.
#
# @Author: Brian Cherinka
# @Date:   2017-08-06 22:12:44
# @Last modified by:   Brian Cherinka
# @Last Modified time: 2017-08-30 15:10:16

from __future__ import print_function, division, absolute_import
from functools import wraps
from sciserver.exceptions import SciServerError, SciServerAPIError
from sciserver import config
import requests

def checkAuth(func):
    ''' Decorator that checks if a token has been generated

    Function Decorator to check if a token has already been generated.
    If not it raises an error and tells you to log in.  Otherwise it
    returns the function and proceeds as normal.

    Returns:
        The decorated function

    Raises:
        SciServerError: User token is not defined. First log into SciServer.

    Example:
        >>>
        >>> @checkauth
        >>> def my_function():
        >>>     return 'I am working function'
        >>>

    '''

    @wraps(func)
    def wrapper(*args, **kwargs):
        token = config.get_token()
        if not token:
            raise SciServerError('User token is not defined. First log into SciServer.')
        else:
            return func(*args, **kwargs)
    return wrapper

def check_response(response, errmsg='Error'):
    ''' Checks the response

    Checks the response from a given request.  Raises an error
    for any response with a non 200 status_code .

    Parameters:
        response:
            the Python requests response object
        errmsg (str):
            A custom error message for the type of request being sent

    Returns:
        the HTTP response

    Raises:
        SciServerApiError: upon any HTTPError caused by a bad status code

    '''

    try:
        isbad = response.raise_for_status()
    except requests.HTTPError as http:
        err = response.content.decode()
        raise SciServerAPIError('{0}\n {1}: {2}'.format(http, errmsg, err))
    else:
        assert isbad is None, 'Http status code should not be bad'
        assert response.ok is True, 'Ok status should be true'
        return response

def make_header(content_type='application/json', accept_header='text/plain'):
    ''' Makes a request header

    Makes a request header to be passed along.  Includes the Content-Type,
    Accept arguments.  Also includes an X-Auth-Token is a token is present.

    Parameters:
        content_type (str):
            The request header Content-Type (default: application/json)
        accept_header (str):
            The request header Accept (default: application/json)

    Returns:
        headers (dict):
            The dictionary to be used as a request header

    '''

    headers = {'Content-Type': content_type, 'Accept': accept_header}

    # check for auth token
    token = config.get_token()
    if token is not None and token != "":
        headers['X-Auth-Token'] = token

    return headers

def send_request(url, reqtype='get', data=None, content_type='application/json',
                 acceptHeader='text/plain', errmsg='Error', stream=None):
    ''' Sends a request to the server

    Parameters:
        url (str):
            The url path for the request
        reqtype (str):
            The type of request to perform.  Default is get.  Choices are get, post, put, delete
        data ({str|dict}):
            Optional data to send in the request
        content_type (str):
            the header Content-Type argument.  Default is application/json
        acceptHeader (str):
            the header Accept argument. Default is text/plain
        errmsg (str):
            custom error message in case of faults
        stream (bool):
            optional.  if False, the response content will be immediately downloaded

    Returns:
        response:
            The HTTP response object

    Raises:
        SciServerError: upon any error to occur when sending the requests

    Example:
        >>> # to send a simple get request
        >>> response = send_request(url, errmsg='Error when executing a sql query.')
        >>>
        >>> # to send a post request
        >>> response = send_request(jobsURL, reqtype='post', data=data, content_type='application/json',
        >>>        acceptHeader='application/json', errmsg='Error when submitting job on queue')

    '''

    headers = make_header(content_type=content_type, accept_header=acceptHeader)

    # send the request
    try:
        if reqtype == 'get':
            response = requests.get(url, headers=headers, stream=stream)
        elif reqtype == 'post':
            response = requests.post(url, data=data, headers=headers, stream=stream)
        elif reqtype == 'put':
            response = requests.put(url, data=data, headers=headers, stream=stream)
        elif reqtype == 'delete':
            response = requests.delete(url, headers=headers, stream=stream)
    except Exception as e:
        raise SciServerError("A requests error occurred attempting to send: {0}".format(e))
    else:
        resp = check_response(response, errmsg=errmsg)
        return resp

class Task(object):
    ''' This class describes a SciServer Task '''

    def __init__(self, name=None, use_base=None, component='SciServer'):
        self.use_base = use_base
        self.component = component
        self.set_name(name)

    @property
    def base_name(self):
        if config.isSciServerComputeEnvironment():
            base = "Compute.SciScript-Python.{0}".format(self.component)
        else:
            base = "SciScript-Python.{0}".format(self.component)
        return base

    def set_name(self, name):
        if name:
            if self.use_base:
                self.name = '{0}.{1}'.format(self.base_name, name)
            else:
                self.name = name
        else:
            self.name = self.base_name


import pygame

def load_image(name):
    try:
        image = pygame.image.load(name)
    except (pygame.error, message):
        print("Cannot load image: " + name)
        raise (SystemExit, message)
    image = image.convert()
    return image, image.get_rect()

class Earth(pygame.sprite.Sprite):

    def __init__(self, gs):
        pygame.sprite.Sprite.__init__(self)
        self.gs = gs
        self.img_path = "./globe.png"
        self.red_globe = "./globe_red100.png"
        self.image, self.rect = load_image(self.img_path)
        self.x = self.gs.width-(self.rect.width/4)
        self.y = self.gs.height-(self.rect.height/4)
        self.rect.centerx = self.x
        self.rect.centery = self.y
        self.radius = self.rect.width/2
        self.max_hp = 500
        self.hp = 500
        self.exploding = 0
        self.exp_frame = 0

    def tick(self):
        self.explode()

    def hit(self):
        self.hp -= 100
        if self.hp <= 0:
            self.exploding = 1
        elif self.hp < .41*self.max_hp:
            self.image, self.rect = load_image(self.red_globe)
            self.rect.centerx = self.x
            self.rect.centery = self.y

    def explode(self):
        if self.exploding and self.exp_frame <= 16:
            if self.exp_frame < 10:
                self.image, self.rect = load_image("./explosion/frames00" + str(self.exp_frame) + "a.png")
            else:
                self.image, self.rect = load_image("./explosion/frames0" + str(self.exp_frame) + "a.png")

            self.rect.centerx = self.x
            self.rect.centery = self.y
            self.exp_frame += 1
            if self.exp_frame == 16:
                self.gs.end = 1