import argparse
import math
import random
import os
import torch

from lpips_pytorch import LPIPS, lpips

from torch import nn, autograd, optim
from torch.nn import functional as F
from torchvision import utils, models
from tqdm import tqdm
from PIL import Image, ImageDraw, ImageColor

from dataset import get_dataloader
from distributed import (
    get_rank,
    synchronize,
    reduce_loss_dict,
    reduce_sum,
    get_world_size,
)
from op import conv2d_gradfix
from non_leaking import augment, AdaptiveAugment
try:
    import wandb

except ImportError:
    wandb = None

torch.cuda.empty_cache()

@torch.no_grad()
def get_image(tensor, **kwargs):
    grid = utils.make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
    im = Image.fromarray(ndarr)
    return im

@torch.no_grad()
def concat_image_by_height(im1, im2):
    dst = Image.new('RGB', (im1.width + im2.width, im1.height))
    dst.paste(im1, (0, 0))
    dst.paste(im2, (im1.width, 0))
    return dst

def requires_grad(model, flag=True):
    for p in model.parameters():
        p.requires_grad = flag

def accumulate(model1, model2, decay=0.999):
    par1 = dict(model1.named_parameters())
    par2 = dict(model2.named_parameters())

    for k in par1.keys():
        par1[k].data.mul_(decay).add_(par2[k].data, alpha=1 - decay)

def sample_data(loader):
    while True:
        for batch in loader:
            yield batch

def d_logistic_loss(real_pred, fake_pred):
    real_loss = F.softplus(-real_pred)
    fake_loss = F.softplus(fake_pred)

    return real_loss.mean() + fake_loss.mean()

def d_r1_loss(real_pred, real_img):
    with conv2d_gradfix.no_weight_gradients():
        grad_real, = autograd.grad(
            outputs=real_pred.sum(), inputs=real_img, create_graph=True
        )
    grad_penalty = grad_real.pow(2).reshape(grad_real.shape[0], -1).sum(1).mean()

    return grad_penalty

def g_nonsaturating_loss(fake_pred):
    loss = F.softplus(-fake_pred).mean()

    return loss

def g_path_regularize(fake_img, latents, mean_path_length, decay=0.01):
    noise = torch.randn_like(fake_img) / math.sqrt(
        fake_img.shape[2] * fake_img.shape[3]
    )
    grad, = autograd.grad(
        outputs=(fake_img * noise).sum(), inputs=latents, create_graph=True
    )
    path_lengths = torch.sqrt(grad.pow(2).sum(2).mean(1))

    path_mean = mean_path_length + decay * (path_lengths.mean() - mean_path_length)

    path_penalty = (path_lengths - path_mean).pow(2).mean()

    return path_penalty, path_mean.detach(), path_lengths

def make_noise(batch, latent_dim, n_noise, device):
    if n_noise == 1:
        return torch.randn(batch, latent_dim, device=device)

    noises = torch.randn(n_noise, batch, latent_dim, device=device).unbind(0)

    return noises

def mixing_noise(batch, latent_dim, prob, device):
    if prob > 0 and random.random() < prob:
        return make_noise(batch, latent_dim, 2, device)

    else:
        return [make_noise(batch, latent_dim, 1, device)]

def set_grad_none(model, targets):
    for n, p in model.named_parameters():
        if n in targets:
            p.grad = None

def save_sample_images(g_ema, args, itr_idx, sample_z, sample_z_label=None):
    with torch.no_grad():
        filename = os.path.join(args.output_path, f"sample/{str(itr_idx).zfill(6)}.png")

        g_ema.eval()

        nof_images = sample_z.shape[0]
        minv = -1
        maxv = 1
        if sample_z_label is None or sample_z_label.shape[2] != 2:
            sample, _ = g_ema([sample_z])
            minv = torch.min(sample)
            maxv = torch.max(sample)

            utils.save_image(
                sample,
                filename,
                nrow=int(nof_images ** 0.5),
                normalize=True,
                value_range=(minv, maxv))
            return
        else:
            # For Conditional GAN, draw borders around each image according to its label
            sample, _ = g_ema([sample_z], labels=sample_z_label)
            minv = torch.min(sample)
            maxv = torch.max(sample)
            sample_z_label = sample_z_label.cpu().squeeze()
            im = get_image(
                sample,
                nrow=int(nof_images ** 0.5),
                normalize=True,
                value_range=(minv, maxv),
                padding=0)
            border_size = 4
            draw = ImageDraw.Draw(im)
            # we have to iterate, each image get a padding fill value of its own
            for i in range(nof_images):

                label = sample_z_label[i].numpy()[0].item()

                if label == 0:
                    row = i // 4
                    column = i % 4
                    # img = sample[i, :]
                    draw.rectangle(
                        [column*256,row*256,
                        (column+1)*256,(row+1)*256],
                        outline=ImageColor.getrgb("darkred"),
                        width=border_size)

            im.save(filename)

def save_real_vs_encoded(generator, args, i, sample_images, sample_labels):
    with torch.no_grad():
        latents = encoder(sample_images)
        encoded_imgs, _ = generator([latents], labels=sample_labels, input_is_latent=True)

        # save real vs. encoded images
        minv = torch.min(encoded_imgs)
        maxv = torch.max(encoded_imgs)
        encoded_image = get_image(
            encoded_imgs,
            nrow=int(args.batch ** 0.5),
            normalize=True,
            value_range=(minv, maxv)
        )

        # save real vs. encoded images
        minv = torch.min(sample_images)
        maxv = torch.max(sample_images)
        real_image = get_image(
            sample_images,
            nrow=int(args.batch ** 0.5),
            normalize=True,
            value_range=(minv, maxv)
        )
        filename_img = os.path.join(args.output_path, f"sample/real_vs_encoded_{str(i).zfill(6)}.png")
        combined = concat_image_by_height(encoded_image, real_image)
        combined.save(filename_img)

def train_encoder(args, loader, generator, g_optim, g_ema, device, classifier, encoder, e_optim, ckpt):

    loader = sample_data(loader)

    pbar = range(args.iter)

    if get_rank() == 0:
        pbar = tqdm(pbar, initial=args.start_iter, dynamic_ncols=True, smoothing=0.01)

    class_loss = torch.tensor(0.0, device=device)
    reconstruct_loss = torch.tensor(0.0, device=device)
    loss_dict = {}

    g_module = generator
    e_module = encoder

    accum = 0.5 ** (32 / (10 * 1000))

    sample_z = torch.randn(args.n_sample, args.latent, device=device)

    for idx in pbar:
        i = idx + args.start_iter

        if i > args.iter:
            print("Done!")

            break

        real_img, real_labels = next(loader)
        real_img, real_labels = real_img.to(device), real_labels.to(device)

        real_encoded = encoder(real_img)
        real_logits = classifier(real_img)

        _, encoded_labels = torch.max(real_logits, 1)
        encoded_labels = F.one_hot(encoded_labels, num_classes=args.classifier_nof_classes)
        fake_img, _ = generator([real_encoded], labels=encoded_labels, input_is_latent=True)
        fake_logits = classifier(fake_img)

        logsoft = torch.nn.LogSoftmax(dim = 1)
        class_loss = torch.nn.KLDivLoss(reduction='batchmean', log_target=True)(logsoft(fake_logits), logsoft(real_logits))

        loss_dict["class_loss"] = class_loss

        fake_encoded = encoder(fake_img)
        reconstruct_loss_x = F.l1_loss(fake_img, real_img)
        reconstruct_loss_w = F.l1_loss(fake_encoded, real_encoded)
        reconstruct_loss_lpips = lpips(fake_img/fake_img.max(), real_img/real_img.max(), net_type='alex', version='0.1').flatten()

        reconstruct_loss = reconstruct_loss_x + reconstruct_loss_w + reconstruct_loss_lpips
        loss_dict["reconstruct_loss"] = reconstruct_loss
        loss_dict["reconstruct_loss_x"] = reconstruct_loss_x
        loss_dict["reconstruct_loss_w"] = reconstruct_loss_w
        loss_dict["reconstruct_loss_lpips"] = reconstruct_loss_lpips

        generator.zero_grad()
        encoder.zero_grad()
        total_loss = class_loss + reconstruct_loss
        loss_dict["total_loss"] = total_loss
        total_loss.backward()
        e_optim.step()
        g_optim.step()
        accumulate(g_ema, g_module, accum)

        class_loss_val = loss_dict["class_loss"].mean().item()
        reconstruct_loss_val = loss_dict["reconstruct_loss"].mean().item()
        reconstruct_loss_x_val = loss_dict["reconstruct_loss_x"].mean().item()
        reconstruct_loss_w_val = loss_dict["reconstruct_loss_w"].mean().item()
        reconstruct_loss_lpips_val = loss_dict["reconstruct_loss_lpips"].mean().item()
        total_loss_val = loss_dict["total_loss"].mean().item()

        if get_rank() == 0:
            print(  f"class: {class_loss_val:.6f}; "
                    f"reconstruct: {reconstruct_loss_val:.6f} "
                    f"reconstruct_x: {reconstruct_loss_x_val:.6f} "
                    f"reconstruct_w: {reconstruct_loss_w_val:.6f} "
                    f"reconstruct_lpips: {reconstruct_loss_lpips_val:.6f} "
                    f"total_loss: {total_loss_val:.6f} "
                    )

            if wandb and args.wandb:
                wandb.log(
                    {
                        "Class Loss": class_loss_val,
                        "Reconstruct Loss": reconstruct_loss_val,
                        "Reconstruct x": reconstruct_loss_x_val,
                        "Reconstruct w": reconstruct_loss_w_val,
                        "Reconstruct LPIPS": reconstruct_loss_lpips_val,
                        "Total Loss": total_loss_val
                    }
                )
                wandb.watch(generator)
                wandb.watch(encoder)

            if i == 0:
                continue

            if i % args.save_samples_every == 0:
                with torch.no_grad():
                    encoded_image = get_image(
                        fake_img,
                        nrow=int(args.batch ** 0.5),
                        normalize=True,
                        scale_each=True
                    )
                    real_image = get_image(
                        real_img,
                        nrow=int(args.batch ** 0.5),
                        normalize=True,
                        scale_each=True
                    )
                    filename_img = os.path.join(args.output_path, f"sample/img_{str(i).zfill(6)}.png")
                    combined = concat_image_by_height(encoded_image, real_image)
                    combined.save(filename_img)

            if i % args.save_checkpoint_every == 0:
                filename = os.path.join(args.output_path, f"checkpoint/{str(i).zfill(6)}.pt")
                torch.save(
                    {
                        "g": g_module.state_dict(),
                        "e": e_module.state_dict(),
                        "g_ema": g_ema.state_dict(),
                        "g_optim": g_optim.state_dict(),
                        "e_optim": e_optim.state_dict(),
                        "args": args,
                    },
                    filename,
                )

def train(args, loader, generator, discriminator, g_optim, d_optim, g_ema, device, ckpt):
    loader = sample_data(loader)

    pbar = range(args.iter)

    cgan = args.cgan

    if get_rank() == 0:
        pbar = tqdm(pbar, initial=args.start_iter, dynamic_ncols=True, smoothing=0.01)

    mean_path_length = 0

    d_loss_val = 0
    r1_loss = torch.tensor(0.0, device=device)
    g_loss_val = 0
    path_loss = torch.tensor(0.0, device=device)
    weighted_path_loss = torch.tensor(0.0, device=device)
    path_lengths = torch.tensor(0.0, device=device)
    mean_path_length_avg = 0
    loss_dict = {}

    if args.distributed:
        g_module = generator.module
        d_module = discriminator.module

    else:
        g_module = generator
        d_module = discriminator

    accum = 0.5 ** (32 / (10 * 1000))
    ada_aug_p = args.augment_p if args.augment_p > 0 else 0.0
    r_t_stat = 0

    if args.augment and args.augment_p == 0:
        ada_augment = AdaptiveAugment(args.ada_target, args.ada_length, 8, device)

    sample_z = torch.randn(args.n_sample, args.latent, device=device)
    if cgan:
        sample_z_label = torch.randint(0, args.classifier_nof_classes, (args.n_sample, 1), device=device)
        sample_z_label = F.one_hot(sample_z_label, num_classes=args.classifier_nof_classes)

    for idx in pbar:
        i = idx + args.start_iter

        if i > args.iter:
            print("Done!")

            break

        if cgan:
            real_img, real_label = next(loader)
            real_img, real_label = real_img.to(device), real_label.to(device)
            # real_label = F.one_hot(real_label, num_classes=args.classifier_nof_classes)
        else:
            real_img = next(loader)
            real_img = real_img.to(device)

        requires_grad(generator, False)
        requires_grad(discriminator, True)

        # Generate image from noise
        noise = mixing_noise(args.batch, args.latent, args.mixing, device)
        if cgan:
            # Get one-hot labels from noise
            random_label = torch.randint(0, args.classifier_nof_classes, (args.batch, 1), device=device)
            random_label = F.one_hot(random_label, num_classes=args.classifier_nof_classes)
            fake_img, _ = generator(noise, labels=random_label, input_is_latent=False)
        else:
            fake_img, _ = generator(noise)

        if args.augment:
            real_img_aug, _ = augment(real_img, ada_aug_p)
            fake_img, _ = augment(fake_img, ada_aug_p)
        else:
            real_img_aug = real_img

        if cgan:
            fake_pred = discriminator(fake_img, random_label)
            real_pred = discriminator(real_img_aug, real_label)
        else:
            fake_pred = discriminator(fake_img)
            real_pred = discriminator(real_img_aug)

        d_loss = d_logistic_loss(real_pred, fake_pred)

        loss_dict["d"] = d_loss
        loss_dict["real_score"] = real_pred.mean()
        loss_dict["fake_score"] = fake_pred.mean()

        discriminator.zero_grad()
        d_loss.backward()
        d_optim.step()

        if args.augment and args.augment_p == 0:
            ada_aug_p = ada_augment.tune(real_pred)
            r_t_stat = ada_augment.r_t_stat

        d_regularize = i % args.d_reg_every == 0

        if d_regularize:
            real_img.requires_grad = True

            if args.augment:
                real_img_aug, _ = augment(real_img, ada_aug_p)

            else:
                real_img_aug = real_img

            if cgan:
                real_pred = discriminator(real_img_aug, real_label)
            else:
                real_pred = discriminator(real_img_aug)

            r1_loss = d_r1_loss(real_pred, real_img)

            discriminator.zero_grad()
            (args.r1 / 2 * r1_loss * args.d_reg_every + 0 * real_pred[0]).backward()

            d_optim.step()

            # we're using this image again
            real_img.requires_grad = False

        loss_dict["r1"] = r1_loss

        requires_grad(generator, True)
        requires_grad(discriminator, False)

        # Generate fake image from random noise and random label
        noise = mixing_noise(args.batch, args.latent, args.mixing, device)
        if cgan:
            # Get one-hot labels from noise
            random_label = torch.randint(0, args.classifier_nof_classes, (args.batch, 1), device=device)
            random_label = F.one_hot(random_label, num_classes=args.classifier_nof_classes)
            fake_img, _ = generator(noise, labels=random_label, input_is_latent=False)
        else:
            fake_img, _ = generator(noise)

        if args.augment:
            fake_img, _ = augment(fake_img, ada_aug_p)

        if cgan:
            fake_pred = discriminator(fake_img, random_label)
        else:
            fake_pred = discriminator(fake_img)

        g_loss = g_nonsaturating_loss(fake_pred)

        loss_dict["g"] = g_loss

        generator.zero_grad()
        g_loss.backward()
        g_optim.step()

        g_regularize = i % args.g_reg_every == 0

        if g_regularize:
            path_batch_size = max(1, args.batch // args.path_batch_shrink)
            noise = mixing_noise(path_batch_size, args.latent, args.mixing, device)

            if cgan:
                random_label = torch.randint(0, args.classifier_nof_classes, (path_batch_size, 1), device=device)
                random_label = F.one_hot(random_label, num_classes=args.classifier_nof_classes)
                fake_img, latents = generator(noise, labels=random_label, input_is_latent=False, return_latents=True)
            else:
                fake_img, latents = generator(noise, return_latents=True)

            path_loss, mean_path_length, path_lengths = g_path_regularize(
                fake_img, latents, mean_path_length
            )

            generator.zero_grad()
            weighted_path_loss = args.path_regularize * args.g_reg_every * path_loss

            if args.path_batch_shrink:
                weighted_path_loss += 0 * fake_img[0, 0, 0, 0]

            weighted_path_loss.backward()

            g_optim.step()

            mean_path_length_avg = (
                reduce_sum(mean_path_length).item() / get_world_size()
            )

        loss_dict["path"] = path_loss
        loss_dict["path_length"] = path_lengths.mean()

        accumulate(g_ema, g_module, accum)

        loss_reduced = reduce_loss_dict(loss_dict)

        d_loss_val = loss_reduced["d"].mean().item()
        g_loss_val = loss_reduced["g"].mean().item()
        r1_val = loss_reduced["r1"].mean().item()
        path_loss_val = loss_reduced["path"].mean().item()
        real_score_val = loss_reduced["real_score"].mean().item()
        fake_score_val = loss_reduced["fake_score"].mean().item()
        path_length_val = loss_reduced["path_length"].mean().item()

        if get_rank() == 0:
            pbar.set_description(
                (
                    f"d: {d_loss_val:.4f}; g: {g_loss_val:.4f}; r1: {r1_val:.4f}; "
                    f"path: {path_loss_val:.4f}; mean path: {mean_path_length_avg:.4f}; "
                    f"augment: {ada_aug_p:.4f}"
                )
            )

            if wandb and args.wandb:
                wandb.log(
                    {
                        "Generator": g_loss_val,
                        "Discriminator": d_loss_val,
                        "Augment": ada_aug_p,
                        "Rt": r_t_stat,
                        "R1": r1_val,
                        "Path Length Regularization": path_loss_val,
                        "Mean Path Length": mean_path_length,
                        "Real Score": real_score_val,
                        "Fake Score": fake_score_val,
                        "Path Length": path_length_val,
                    }
                )

            if i == 0:
                continue

            if i % args.save_samples_every == 0:
                save_sample_images(g_ema, args, i, sample_z, sample_z_label if cgan else None)

            if i % args.save_checkpoint_every == 0:
                filename = os.path.join(args.output_path, f"checkpoint/{str(i).zfill(6)}.pt")
                torch.save(
                     {
                         "g": g_module.state_dict(),
                         "d": d_module.state_dict(),
                         "g_ema": g_ema.state_dict(),
                         "g_optim": g_optim.state_dict(),
                         "d_optim": d_optim.state_dict(),
                         "args": args,
                         "ada_aug_p": ada_aug_p
                     },
                     filename,
                )

def train_stylechexplain(args, loader, generator, discriminator, g_optim, d_optim, g_ema, device, classifier, encoder, e_optim, ckpt):
    loader = sample_data(loader)

    pbar = range(args.iter)

    if get_rank() == 0:
        pbar = tqdm(pbar, initial=args.start_iter, dynamic_ncols=True, smoothing=0.01)

    mean_path_length = 0

    d_loss_val = 0
    r1_loss = torch.tensor(0.0, device=device)
    g_loss_val = 0
    weighted_path_loss = torch.tensor(0.0, device=device)
    path_loss = torch.tensor(0.0, device=device)
    path_lengths = torch.tensor(0.0, device=device)
    class_loss = torch.tensor(0.0, device=device)
    reconstruct_loss = torch.tensor(0.0, device=device)
    mean_path_length_avg = 0
    loss_dict = {}

    if args.distributed:
        g_module = generator.module
        d_module = discriminator.module
        e_module = encoder.module

    else:
        g_module = generator
        d_module = discriminator
        e_module = encoder

    accum = 0.5 ** (32 / (10 * 1000))
    ada_aug_p = args.augment_p if args.augment_p > 0 else 0.0
    r_t_stat = 0

    if args.augment and args.augment_p == 0:
        ada_augment = AdaptiveAugment(args.ada_target, args.ada_length, 8, device)

    sample_z = torch.randn(args.n_sample, args.latent, device=device)
    sample_z_label = torch.randint(0, args.classifier_nof_classes, (args.n_sample, 1), device=device)
    sample_z_label = F.one_hot(sample_z_label, num_classes=args.classifier_nof_classes)

    sample_images, sample_labels = next(iter(loader))
    sample_images, sample_labels = sample_images.to(device), sample_labels.to(device)
    for idx in pbar:
        i = idx + args.start_iter

        if i > args.iter:
            print("Done!")

            break

        real_img, real_labels = next(loader)
        real_img, real_labels = real_img.to(device), real_labels.to(device)

        #real_labels = F.one_hot(real_labels, num_classes=args.classifier_nof_classes)

        # First, train discriminator
        requires_grad(generator, False)
        requires_grad(encoder, False)
        requires_grad(discriminator, True)

        '''
        if args.augment:
            real_img_aug, _ = augment(real_img, ada_aug_p)
            fake_img, _ = augment(fake_img, ada_aug_p)

        else:
            real_img_aug = real_img
        '''

        # Generate image from encoded_real
        encoded_real = encoder(real_img)
        logits = classifier(real_img)
        _, encoded_labels = torch.max(logits, 1)
        encoded_labels = F.one_hot(encoded_labels, num_classes=args.classifier_nof_classes)

        fake_img_by_encoded, _ = generator([encoded_real], labels=encoded_labels, input_is_latent=True)

        # Generate image from noise
        noise = mixing_noise(args.batch, args.latent, args.mixing, device)
        random_label = torch.randint(0, args.classifier_nof_classes, (args.batch, 1), device=device)
        random_label = F.one_hot(random_label, num_classes=args.classifier_nof_classes)

        fake_img_by_z, _ = generator(noise, labels=random_label, input_is_latent=False)

        # Discriminator calls
        fake_pred_by_encoded = discriminator(fake_img_by_encoded.detach(), encoded_labels.detach())
        fake_pred_by_z = discriminator(fake_img_by_z.detach(), random_label)
        real_pred = discriminator(real_img, real_labels)

        # Loss
        d_loss_1 = d_logistic_loss(real_pred, fake_pred_by_encoded)
        d_loss_2 = d_logistic_loss(real_pred, fake_pred_by_z)

        d_loss = d_loss_1 + d_loss_2

        loss_dict["d"] = d_loss
        loss_dict["real_score"] = real_pred.mean()
        loss_dict["fake_score"] = fake_pred_by_encoded.mean()
        loss_dict["fake_by_z_score"] = fake_pred_by_z.mean()

        discriminator.zero_grad()
        d_loss.backward()
        d_optim.step()

        if args.augment and args.augment_p == 0:
            ada_aug_p = ada_augment.tune(real_pred)
            r_t_stat = ada_augment.r_t_stat

        d_regularize = i % args.d_reg_every == 0

        if d_regularize:
            real_img.requires_grad = True

            if args.augment:
                real_img_aug, _ = augment(real_img, ada_aug_p)

            else:
                real_img_aug = real_img

            real_pred = discriminator(real_img_aug, real_labels)
            r1_loss = d_r1_loss(real_pred, real_img)

            discriminator.zero_grad()
            (args.r1 / 2 * r1_loss * args.d_reg_every + 0 * real_pred[0]).backward()

            d_optim.step()

            # we're using this image again
            real_img.requires_grad = False

        loss_dict["r1"] = r1_loss

        # Train Generator
        requires_grad(generator, True)
        requires_grad(encoder, False)
        requires_grad(discriminator, False)

        '''
        if args.augment:
            fake_img, _ = augment(fake_img, ada_aug_p)
        '''

        # Generate fake image from random noise and random label
        noise = mixing_noise(args.batch, args.latent, args.mixing, device)
        random_label = torch.randint(0, args.classifier_nof_classes, (args.batch, 1), device=device)
        random_label = F.one_hot(random_label, num_classes=args.classifier_nof_classes)
        fake_img_by_z, _ = generator(noise, labels=random_label, input_is_latent=False)
        fake_pred_by_z = discriminator(fake_img_by_z, random_label)

        # Generate image from encoded_real
        encoded_real = encoder(real_img)
        logits = classifier(real_img)
        _, encoded_labels = torch.max(logits, 1)
        encoded_labels = F.one_hot(encoded_labels, num_classes=args.classifier_nof_classes)
        fake_img_by_encoded, _ = generator([encoded_real.detach()], labels=encoded_labels.detach(), input_is_latent=True)
        fake_pred_by_encoded = discriminator(fake_img_by_encoded, encoded_labels)

        g_loss = (g_nonsaturating_loss(fake_pred_by_z) + g_nonsaturating_loss(fake_pred_by_encoded)) / 2
        loss_dict["g"] = g_loss

        generator.zero_grad()
        g_loss.backward()
        g_optim.step()

        g_regularize = i % args.g_reg_every == 0

        if g_regularize:
            path_batch_size = max(1, args.batch // args.path_batch_shrink)

            noise = mixing_noise(path_batch_size, args.latent, args.mixing, device)
            random_label = torch.randint(0, args.classifier_nof_classes, (path_batch_size, 1), device=device)
            random_label = F.one_hot(random_label, num_classes=args.classifier_nof_classes)

            fake_img_path_reg, latents = generator(noise, labels=random_label, input_is_latent=False, return_latents=True)

            path_loss, mean_path_length, path_lengths = g_path_regularize(
                fake_img_path_reg, latents, mean_path_length
            )

            weighted_path_loss = args.path_regularize * args.g_reg_every * path_loss

            if args.path_batch_shrink:
                weighted_path_loss += 0 * fake_img_path_reg[0, 0, 0, 0]

            generator.zero_grad()
            weighted_path_loss.backward()
            g_optim.step()

            mean_path_length_avg = (
                reduce_sum(mean_path_length).item() / get_world_size()
            )

        loss_dict["path"] = path_loss
        loss_dict["path_length"] = path_lengths.mean()

        # Train Encoder & Generator
        requires_grad(generator, True)
        requires_grad(encoder, True)
        requires_grad(discriminator, False)

        # Generate image from real_encoded
        real_encoded = encoder(real_img)
        real_logits = classifier(real_img)
        _, encoded_labels = torch.max(real_logits, 1)
        encoded_labels = F.one_hot(encoded_labels, num_classes=args.classifier_nof_classes)
        fake_img, _ = generator([real_encoded], labels=encoded_labels, input_is_latent=True)
        fake_logits = classifier(fake_img)

        # KLDivLoss expects input to be log(prob) and target to be prob.
        logsoft = torch.nn.LogSoftmax(dim=1)
        class_loss = torch.nn.KLDivLoss(reduction='batchmean', log_target=True)(logsoft(fake_logits),
                                                                                logsoft(real_logits))

        loss_dict["class_loss"] = class_loss

        fake_encoded = encoder(fake_img)
        reconstruct_loss_x = F.l1_loss(fake_img, real_img)
        reconstruct_loss_w = F.l1_loss(fake_encoded, real_encoded)
        reconstruct_loss_lpips = lpips(fake_img/fake_img.max(), real_img/real_img.max(), net_type='alex', version='0.1').flatten()

        reconstruct_loss = reconstruct_loss_x + reconstruct_loss_w + reconstruct_loss_lpips
        loss_dict["reconstruct_loss"] = reconstruct_loss
        loss_dict["reconstruct_loss_x"] = reconstruct_loss_x
        loss_dict["reconstruct_loss_w"] = reconstruct_loss_w
        loss_dict["reconstruct_loss_lpips"] = reconstruct_loss_lpips

        encoder.zero_grad()
        generator.zero_grad()

        total_loss = class_loss + reconstruct_loss
        total_loss.backward()

        e_optim.step()
        g_optim.step()

        # update g_ema module
        accumulate(g_ema, g_module, accum)

        loss_reduced = reduce_loss_dict(loss_dict)

        d_loss_val = loss_dict["d"].mean().item()
        g_loss_val = loss_dict["g"].mean().item()
        r1_val = loss_dict["r1"].mean().item()
        path_loss_val = loss_dict["path"].mean().item()
        real_score_val = loss_dict["real_score"].mean().item()
        fake_score_val = loss_dict["fake_score"].mean().item()
        path_length_val = loss_dict["path_length"].mean().item()
        class_loss_val = loss_dict["class_loss"].mean().item()
        reconstruct_loss_val = loss_dict["reconstruct_loss"].mean().item()
        reconstruct_loss_x_val = loss_dict["reconstruct_loss_x"].mean().item()
        reconstruct_loss_w_val = loss_dict["reconstruct_loss_w"].mean().item()
        reconstruct_loss_lpips_val = loss_dict["reconstruct_loss_lpips"].mean().item()

        if get_rank() == 0:
            pbar.set_description(
                (
                    f"d: {d_loss_val:.2f}; g: {g_loss_val:.2f}; r1: {r1_val:.2f}; "
                    f"path: {path_loss_val:.2f}; mean path: {mean_path_length_avg:.2f}; "
                    f"augment: {ada_aug_p:.1f}; class: {class_loss_val:.2f}; "
                    f"reconstruct: {reconstruct_loss_val:.1f} "
                )
            )
            '''
            print(f"d: {d_loss_val:.2f}; g: {g_loss_val:.2f}; r1: {r1_val:.2f}; "
                    f"path: {path_loss_val:.2f}; mean path: {mean_path_length_avg:.2f}; "
                    f"augment: {ada_aug_p:.1f}; class: {class_loss_val:.2f}; "
                    f"reconstruct: {reconstruct_loss_val:.1f} "
                    f"reconstruct_x: {reconstruct_loss_x_val:.1f} "
                    f"reconstruct_w: {reconstruct_loss_w_val:.1f} "
                    f"reconstruct_lpips: {reconstruct_loss_lpips_val:.1f} "
                    )
            '''

            if wandb and args.wandb:
                wandb.log(
                    {
                        "Generator": g_loss_val,
                        "Discriminator": d_loss_val,
                        "Augment": ada_aug_p,
                        "Rt": r_t_stat,
                        "R1": r1_val,
                        "Path Length Regularization": path_loss_val,
                        "Mean Path Length": mean_path_length,
                        "Real Score": real_score_val,
                        "Fake Score": fake_score_val,
                        "Path Length": path_length_val,
                        "Class Loss": class_loss_val,
                        "Reconstruct Loss": reconstruct_loss_val,
                        "Reconstruct x": reconstruct_loss_x_val,
                        "Reconstruct w": reconstruct_loss_w_val,
                        "Reconstruct LPIPS": reconstruct_loss_lpips_val,
                    }
                )

            if i == 0:
                continue

            if i % args.save_samples_every == 0:
                save_sample_images(g_ema, args, i, sample_z, sample_z_label)
                save_real_vs_encoded(generator, args, i, sample_images, sample_labels)

            if i % args.save_checkpoint_every == 0:
                filename = os.path.join(args.output_path, f"checkpoint/{str(i).zfill(6)}.pt")
                torch.save(
                    {
                        "g": g_module.state_dict(),
                        "d": d_module.state_dict(),
                        "e": e_module.state_dict(),
                        "g_ema": g_ema.state_dict(),
                        "g_optim": g_optim.state_dict(),
                        "d_optim": d_optim.state_dict(),
                        "e_optim": e_optim.state_dict(),
                        "args": args,
                        "ada_aug_p": ada_aug_p
                    },
                    filename,
                )

if __name__ == "__main__":
    device = "cuda"

    parser = argparse.ArgumentParser(description="StyleGAN2 trainer")

    parser.add_argument("dataset_path", type=str, help="path to the lmdb dataset")
    parser.add_argument('--arch', type=str, default='stylegan2', help='model architectures (stylegan2 | autoencoder | StyleEx)')
    parser.add_argument(
        "--iter", type=int, default=800000, help="total training iterations"
    )
    parser.add_argument(
        "--batch", type=int, default=16, help="batch sizes for each gpus"
    )
    parser.add_argument('--output_path',
                        type=str,
                        help='Path into which artifacts are saved (sampled images, checkpoint')

    parser.add_argument(
        "--n_sample",
        type=int,
        default=16,
        help="number of the samples generated during training",
    )
    parser.add_argument(
        "--size", type=int, default=256, help="image sizes for the model"
    )
    parser.add_argument(
        "--r1", type=float, default=10, help="weight of the r1 regularization"
    )
    parser.add_argument(
        "--path_regularize",
        type=float,
        default=2,
        help="weight of the path length regularization",
    )
    parser.add_argument(
        "--path_batch_shrink",
        type=int,
        default=2,
        help="batch size reducing factor for the path length regularization (reduce memory consumption)",
    )
    parser.add_argument(
        "--d_reg_every",
        type=int,
        default=16,
        help="interval of the applying r1 regularization",
    )
    parser.add_argument(
        "--g_reg_every",
        type=int,
        default=4,
        help="interval of the applying path length regularization",
    )
    parser.add_argument(
        "--mixing", type=float, default=0.9, help="probability of latent code mixing"
    )
    parser.add_argument(
        "--ckpt",
        type=str,
        default=None,
        help="path to the checkpoints to resume training",
    )
    parser.add_argument("--lr", type=float, default=0.002, help="learning rate")
    parser.add_argument(
        "--channel_multiplier",
        type=int,
        default=2,
        help="channel multiplier factor for the model. config-f = 2, else = 1",
    )
    parser.add_argument(
        "--wandb", action="store_true", help="use weights and biases logging"
    )
    parser.add_argument(
        "--augment", action="store_true", help="apply non leaking augmentation"
    )
    parser.add_argument(
        "--augment_p",
        type=float,
        default=0,
        help="probability of applying augmentation. 0 = use adaptive augmentation",
    )
    parser.add_argument(
        "--ada_target",
        type=float,
        default=0.6,
        help="target augmentation probability for adaptive augmentation",
    )
    parser.add_argument(
        "--ada_length",
        type=int,
        default=500 * 1000,
        help="target duraing to reach augmentation probability for adaptive augmentation",
    )
    parser.add_argument(
        "--ada_every",
        type=int,
        default=256,
        help="probability update interval of the adaptive augmentation",
    )

    ####
    # Stylechexplain parameters
    ####
    parser.add_argument(
        "--classifier_nof_classes",
        type=int,
        default=2,
        help="For Stylechexplain or Conditional StyleGAN2, number of classifier target classes",
    )
    parser.add_argument(
        "--classifier_ckpt",
        type=str,
        default=None,
        help="Path for pretrained classifier checkpoint",
    )
    parser.add_argument(
        "--cgan",
        default=False,
        action="store_true",
        help="Train Conditional GAN, auto set to True for StyleEx")
    parser.add_argument(
        "--encoder_ckpt",
        type=str,
        default=None,
        help="Path for pretrained encoder checkpoint",
    )

    parser.add_argument(
        "--compare_to_healthy", action="store_true",
        help="dataset consists of one specific class and no finding/healthy samples."
    )
    parser.add_argument('--filter_label', type=str)

    args = parser.parse_args()

    # Default params
    args.latent = 512
    args.n_mlp = 8
    args.embedding_size = 10  # embedding for labels in Generator and Discriminator

    # control how often samples and checkpoints are saved
    args.save_samples_every = 500
    args.save_checkpoint_every = 500

    n_gpu = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
    args.distributed = n_gpu > 1

    if args.distributed:
        args.local_rank = int(os.environ["LOCAL_RANK"])
        torch.cuda.set_device(args.local_rank)
        torch.distributed.init_process_group(backend="nccl", init_method="env://")
        synchronize()

    # make sure required output dirs exist
    dirs = ['checkpoint', 'sample']
    for d in dirs:
        pth = os.path.join(args.output_path, d)
        if not os.path.exists(pth):
            os.makedirs(pth, exist_ok=True)

    args.start_iter = 0

    if get_rank() == 0 and wandb is not None and args.wandb:
        wandb.init(project=args.arch, resume=True)

    loader = get_dataloader(args)

    supported_archs = ['stylegan2', 'autoencoder', 'StyleEx']
    if args.arch not in supported_archs:
        raise ValueError("Only one of %s is supported architecture" % supported_archs)

    # struct to control build of different components in architecture
    build = dict()
    if args.arch == 'stylegan2':
        from model import Generator, Discriminator
        build = {"G": True, "D": True, "E": False, "C": False}

    elif args.arch == 'autoencoder':
        from model import Generator, Encoder
        args.cgan = True
        build = {"G": True, "D": False, "E": True, "C": True}

    elif args.arch == 'StyleEx':
        from model import Generator, Discriminator, Encoder
        args.cgan = True
        build = {"G": True, "D": True, "E": True, "C": True}
    else:
        print("Architecture %s not supported, exiting.." % args.arch)
        exit()

    ckpt = None
    if args.ckpt is not None:
        print("load model:", args.ckpt)

        ckpt = torch.load(args.ckpt, map_location=lambda storage, loc: storage)

        try:
            ckpt_name = os.path.basename(args.ckpt)
            args.start_iter = int(os.path.splitext(ckpt_name)[0])

        except ValueError:
            pass

    # Build Generator
    if build["G"]:
        generator = Generator(
            args.size, args.latent, args.n_mlp, channel_multiplier=args.channel_multiplier,
            conditional_gan=args.cgan, nof_classes=args.classifier_nof_classes, embedding_size=args.embedding_size
        ).to(device)

        g_ema = Generator(
            args.size, args.latent, args.n_mlp, channel_multiplier=args.channel_multiplier,
            conditional_gan=args.cgan, nof_classes=args.classifier_nof_classes, embedding_size=args.embedding_size
        ).to(device)
        g_ema.eval()
        accumulate(g_ema, generator, 0)

        g_reg_ratio = args.g_reg_every / (args.g_reg_every + 1)
        g_optim = optim.Adam(
            generator.parameters(),
            lr=args.lr * g_reg_ratio,
            betas=(0 ** g_reg_ratio, 0.99 ** g_reg_ratio),
        )

        if ckpt:
            generator.load_state_dict(ckpt["g"])
            g_ema.load_state_dict(ckpt["g_ema"])
            g_optim.load_state_dict(ckpt["g_optim"])

        if args.distributed:
            generator = nn.parallel.DistributedDataParallel(
                generator,
                device_ids=[args.local_rank],
                output_device=args.local_rank,
                broadcast_buffers=False,
            )

    # Build Discriminator
    if build["D"]:
        discriminator = Discriminator(
            args.size, channel_multiplier=args.channel_multiplier, conditional_gan=args.cgan
        ).to(device)

        d_reg_ratio = args.d_reg_every / (args.d_reg_every + 1)
        d_optim = optim.Adam(
            discriminator.parameters(),
            lr=args.lr * d_reg_ratio,
            betas=(0 ** d_reg_ratio, 0.99 ** d_reg_ratio),
        )

        if ckpt:
            discriminator.load_state_dict(ckpt["d"])
            d_optim.load_state_dict(ckpt["d_optim"])

        if args.distributed:
            discriminator = nn.parallel.DistributedDataParallel(
                discriminator,
                device_ids=[args.local_rank],
                output_device=args.local_rank,
                broadcast_buffers=False,
            )

    # Build Classifier
    if build["C"]:
        if args.classifier_ckpt is None:
            print("Pretrianed classifier needed!")
            exit()

        # classifier = init_ffhq_classifier(device, args.classifier_ckpt)
        classifier = models.densenet121(pretrained=True)
        classifier.classifier = torch.nn.Linear(1024, args.classifier_nof_classes)
        classifier.to(device)

        print("loading classifier:", args.classifier_ckpt)
        classifier.load_state_dict(torch.load(args.classifier_ckpt))

        # Classifier shouldn't be trained
        classifier.eval()
        requires_grad(classifier, False)

    # Build Encoder
    if build["E"]:
        encoder = Encoder(
            args.size, channel_multiplier=args.channel_multiplier, output_channels=args.latent
        ).to(device)
        e_optim = optim.Adam(
            encoder.parameters(),
            lr=args.lr
        )

        if args.encoder_ckpt is not None:
            # Encoder ckpt given separately
            print("load encoder model:", args.encoder_ckpt)
            e_ckpt = torch.load(args.encoder_ckpt, map_location=lambda storage, loc: storage)
        else:
            # Encoder given as part of general ckpt
            e_ckpt = ckpt

        if e_ckpt is not None and "e" in e_ckpt:
            encoder.load_state_dict(e_ckpt["e"])
            e_optim.load_state_dict(e_ckpt["e_optim"])

        if args.distributed:
            encoder = nn.parallel.DistributedDataParallel(
                encoder,
                device_ids=[args.local_rank],
                output_device=args.local_rank,
                broadcast_buffers=False,
            )

    if args.arch == "stylegan2":
        train(args, loader, generator, discriminator, g_optim, d_optim, g_ema, device, ckpt)

    elif args.arch == "autoencoder":
        train_encoder(args, loader, generator, g_optim, g_ema, device, classifier, encoder, e_optim, ckpt)

    elif args.arch == "StyleEx":
        train_stylechexplain(args, loader, generator, discriminator, g_optim, d_optim, g_ema, device,
                             classifier, encoder, e_optim, ckpt)


"""
Test LTI web components of the application
"""
import json
import logging
import os
import time
import unittest

from mock import patch

import xsiftx.config
from xsiftx.config import get_config, get_consumer, XsiftxNoConfigException
from xsiftx.util import get_sifters
from xsiftx.lti.decorators import LTI_STAFF_ROLES
import xsiftx.web

class TestLTIWebApp(unittest.TestCase):
    """
    Test class for validating LTI interface
    to xsiftx.
    """
    # pylint: disable=e1103,r0904

    def setUp(self):
        """
        grab application test client
        """
        # pylint: disable=C0103
        self.client = xsiftx.web.app.test_client()
        self.settings = get_config()

    def _oauth_request(self, params={}, consumer_id=0):
        """
        Returns the data needed for an oauth request
        and adds those to the passed dictionary to
        keep out of the way.
        """
        # pylint: disable=w0102
        consumer = self.settings['consumers'][consumer_id]
        params.update({
            'oauth_consumer_key': consumer['key'],
            'oauth_signature_method': 'PLAINTEXT',
            'oauth_version': '1.0',
            'oauth_signature': '{0}&'.format(consumer.get('secret', '')),
            'oauth_timestamp': str(int(time.time())),
            'oauth_nonce': '123456789',
            'context_id': 'MITx/A.we/some'
        })
        return params

    def _run_sifter(self, sifter='test_sifters'):
        """
        Run a valid sifter and return json from response
        """
        test_url = 'api/v0.1/run'
        response = self.client.post(
            test_url,
            data=self._oauth_request({'sifter': sifter}, 1)
        )
        self.assertEqual(response.status_code, 200)
        return json.loads(response.data)

    def test_conf_settings(self):
        """
        Test and validate configuration and test rigging
        """
        self.assertEqual(
            self.settings['consumers'],
            [
                {'key': 'test_course1',
                 'secret': 'test_secret1',
                 'allowed_sifters': ['test_sifters']},
                {'key': 'test_course2',
                 'secret': 'test_secret2', },
                {'key': 'testo', },
            ]
        )

        # Test consumer search
        consumer = get_consumer('test_course1')
        self.assertEqual(consumer['secret'], 'test_secret1')

        # Test config exception
        conf_save = os.environ['XSIFTX_CONFIG']
        config_path_save = xsiftx.config.CONFIG_PATHS
        del os.environ['XSIFTX_CONFIG']

        xsiftx.config.CONFIG_PATHS = ['/dev/null', ]
        with self.assertRaisesRegexp(XsiftxNoConfigException,
                                     'No configuration found'):
            settings = get_config()

        # Test config by path
        xsiftx.config.CONFIG_PATHS = [conf_save, ]
        settings = get_config()
        self.assertEqual(settings['awesome_canary_key'], 'test')

        os.environ['XSIFTX_CONFIG'] = conf_save
        xsiftx.config.CONFIG_PATHS = config_path_save

    @patch('xsiftx.lti.oauthstore.log')
    def test_no_consumers(self, mock_logging):
        """
        Make sure we handle a lack of conumers altogether well
        """
        saved_consumers = xsiftx.config.settings['consumers']
        del xsiftx.config.settings['consumers']

        response = self.client.post('/', data=self._oauth_request())
        self.assertEquals(response.status_code, 401)
        self.assertTrue(mock_logging.critical.called_with(
            "No consumers defined in settings."
            "Have you created a configuration file?"
        ))

        # Restore
        xsiftx.config.settings['consumers'] = saved_consumers

    @patch('xsiftx.lti.oauthstore.log')
    def test_consumer_without_secret(self, mock_logging):
        """
        Test a consumer that is missing a secret.
        """
        keyless_consumer = 2
        response = self.client.post(
            '/',
            data=self._oauth_request(
                {'oauth_secret': 'test&'},
                keyless_consumer
            )
        )
        self.assertEquals(response.status_code, 401)
        self.assertTrue(mock_logging.critical.called_with(
            'Consumer {0}, is missing secret'
            'in settings file, and needs correction.'.format(
                self.settings['consumers'][keyless_consumer]['key']
            )
        ))

    def test_lti_authentication(self):
        """
        Test out that LTI authentication is working properly
        and that the decorator is raising the proper exceptions
        """
        # Setup empty request
        response = self.client.get('/')
        self.assertEqual(response.status_code, 401)
        self.assertTrue(
            'This page requires a valid oauth session or request' in
            response.data
        )

        # Test bad key
        params = self._oauth_request()
        params['oauth_consumer_key'] = 'silly'
        response = self.client.post('/', data=params)
        self.assertEqual(response.status_code, 401)
        self.assertTrue(
            'OAuth error: Please check your key' in
            response.data
        )

        # Test bad sig
        params = self._oauth_request()
        params['oauth_signature'] = 'silly&'
        response = self.client.post('/', data=params)
        self.assertEqual(response.status_code, 401)
        self.assertTrue(
            'OAuth error: Please check your key' in
            response.data
        )

        # Valid OAuth, but missing authorization
        response = self.client.post('/', data=self._oauth_request())
        self.assertEqual(response.status_code, 401)
        self.assertTrue(
            'User does not have a role' in
            response.data
        )

        # Test https header
        with xsiftx.web.app.test_client() as client:
            response = client.post(
                '/', data=self._oauth_request(),
                headers={'x-forwarded-proto': 'https'}
            )
        self.assertEqual(response.status_code, 401)
        self.assertTrue(
            'User does not have a role' in
            response.data
        )

    def test_lti_staff_decorator(self):
        """
        Make sure that our authorization of role is happening
        """
        params = self._oauth_request({'roles': 'Student'})
        response = self.client.post('/', data=params)
        self.assertEqual(response.status_code, 401)
        self.assertTrue(
            'You are not in a staff level role. Access is restricted ' in
            response.data
        )

    def test_reauthorization(self):
        """
        Make sure that if my role changes, my access changes as well.
        """
        params = self._oauth_request({'roles': 'Student'})
        response = self.client.post('/', data=params)
        self.assertEqual(response.status_code, 401)
        self.assertTrue(
            'You are not in a staff level role. Access is restricted ' in
            response.data
        )

        params = self._oauth_request({'roles': 'Instructor'})
        response = self.client.post('/', data=params)
        self.assertEqual(response.status_code, 200)

    def test_index(self):
        """
        Several tests of the index page
        """
        params = self._oauth_request({'roles': LTI_STAFF_ROLES[0]})
        response = self.client.post('/', data=params)
        self.assertEqual(response.status_code, 200)
        self.assertTrue(
            'Available Sifters' in
            response.data
        )

    def test_sifter_acl(self):
        """
        Test with consumer that has access to all sifters
        and ensure they all get loaded
        """
        # Test limited ACL
        with xsiftx.web.app.test_client() as client:
            consumer_id = 0
            params = self._oauth_request(
                {'roles': LTI_STAFF_ROLES[0]},
                consumer_id
            )
            response = client.post('/', data=params)
            self.assertEqual(response.status_code, 200)
            self.assertTrue(
                'Available Sifters' in
                response.data
            )
            # Make sure we don't have extra sifters
            self.assertTrue(
                all(sifter in response.data
                    for sifter in
                    self.settings['consumers'][consumer_id]['allowed_sifters'])
            )

        # Test no ACL
        with xsiftx.web.app.test_client() as client:
            consumer_id = 1
            params = self._oauth_request(
                {'roles': LTI_STAFF_ROLES[0]},
                consumer_id
            )
            response = client.post('/', data=params)
            self.assertEqual(response.status_code, 200)
            self.assertTrue(
                'Available Sifters' in
                response.data
            )
            # Make sure we don't have extra sifters
            self.assertTrue(
                all(sifter in response.data
                    for sifter in get_sifters())
            )

    def test_api_error(self):
        """
        Excercise the API error handler to make sure it returns
        json.
        """

        test_url = '/api/v0.1/run'

        # Test that LTI error on API url returns JSON instead of html
        response = self.client.post(test_url)
        self.assertEqual(response.status_code, 401)
        json.loads(response.data)

        # Test no sifter
        response = self.client.post(
            test_url,
            data=self._oauth_request()
        )
        self.assertEqual(response.status_code, 400)
        # Should raise if we don't get json
        reply_json = json.loads(response.data)
        self.assertTrue(
            u'This api call requires the parameter "sifter".' ==
            reply_json['message']
        )

        # Test invalid sifter
        response = self.client.post(
            test_url,
            data=self._oauth_request({'sifter': 'not_a_real_thing'})
        )
        self.assertEqual(response.status_code, 400)
        # Should raise if we don't get json
        reply_json = json.loads(response.data)
        self.assertTrue(
            u'You have specified an invalid sifter.' ==
            reply_json['message']
        )
        self.assertIsNotNone(reply_json.get('available_sifters', None))

    def test_run_sifter(self):
        """
        Test out the run command
        """
        reply_json = self._run_sifter()
        # Even though the job should fail, we should get
        # back the job id and such
        self.assertTrue(reply_json['tasks'][0]['status'] == 'PENDING')

    def test_task_status_and_delete(self):
        """
        Test task status API point and validate response json
        """
        update_url = '/api/v0.1/update_task_status'
        # Create several runs
        num_runs = 10
        for _ in range(num_runs):
            reply_json = self._run_sifter('xqanalyze')

        # Check we have the right number
        response = self.client.put(
            update_url,
            data=self._oauth_request()
        )
        reply_json = json.loads(response.data)
        self.assertTrue(len(reply_json['tasks']), num_runs)

        # Now delete them
        delete_url = '/api/v0.1/clear_complete_tasks'
        response = self.client.delete(
            delete_url,
            data=self._oauth_request()
        )
        reply_json = json.loads(response.data)
        self.assertTrue(len(reply_json['tasks']), 0)

        # And double check they are gone
        # Check we have the right number
        response = self.client.put(
            update_url,
            data=self._oauth_request()
        )
        reply_json = json.loads(response.data)
        self.assertTrue(len(reply_json['tasks']), 0)

    def test_logging_level(self):
        """
        Tests to make sure logging config happens and handles
        bad output.
        """
        # Check it is the right setting from test config
        self.assertEqual(self.settings['log_level'], 'debug')
        level = xsiftx.web.set_log_level()
        # Check that setting is not applied
        self.assertEqual(logging.DEBUG, level)

        # Set to invalid level and assert the exception
        xsiftx.config.settings['log_level'] = 'awesome'
        with self.assertRaisesRegexp(ValueError, 'Invalid log level.*'):
            level = xsiftx.web.set_log_level()
            self.assertIsNone(level)

        # Make sure it is ok not to have a level set
        del xsiftx.config.settings['log_level']
        level = xsiftx.web.set_log_level()
        self.assertIsNone(level)

        # Restore setting
        xsiftx.config.settings['log_level'] = 'debug'


from kafka import KafkaProducer
import json
import argparse
import os
from common import Common

common = Common("client")
common.load()

parser = argparse.ArgumentParser()
parser.add_argument("topic", help="Name for the topic to submit to.")
parser.add_argument("name", help="Name for the task to submit.")
parser.add_argument("style", help="File name of style.")
parser.add_argument("--args", default='-backend cudnn -batch_size 1 -model starling', help="Arguments to use for training.")
parser.add_argument("--iters", type=int, default=250, help="Iterations per cycle.")
parser.add_argument("--cycles", type=int, default=200, help="Number of cycles to train for.")
parser.add_argument("--learning_rate", type=float, default=4e-2, help="Initial learning rate.")
parser.add_argument("--learning_decay", type=float, default=0.8, help="Discount factor for learning rate by period.")
parser.add_argument("--learning_period", type=int, default=25, help="Number of cycles between learning rate discounts.")

args = parser.parse_args()

producer = KafkaProducer(
        value_serializer=lambda m: json.dumps(m).encode("utf-8"))

print ("Ready.")

task =  {
            "args": args.args,
            "style": args.style,
            "starting_checkpoint": "",
            "iterations_per_cycle": args.iters,
            "cycles_remaining": args.cycles,
            "cycles_completed": 0,
            "base_learning_rate": args.learning_rate,
            "learning_decay": args.learning_decay,
            "learning_period": args.learning_period
        }

common.log("Uploading %s" % args.style)

common.upload("style", args.style)
common.log("Upload complete; submitting job named %s" % args.name)

producer.send(args.topic, key=args.name.encode("utf-8"), value=task)
common.log("Job submitted")


class Solution:
    def merge(self, nums1: list[int], m: int, nums2: list[int], n: int) -> None:
        ptr1, ptr2 = 0, 0
        temp_array = []
        while ptr1 < m and ptr2 < n:
            if nums1[ptr1] <= nums2[ptr2]:
                temp_array.append(nums1[ptr1])
                ptr1 += 1
            elif nums2[ptr2] < nums1[ptr1]:
                temp_array.append(nums2[ptr2])
                ptr2 +=1

        if ptr1 >= m:
            temp_array.extend(nums2[ptr2:n])

        if ptr2 >= n:
            temp_array.extend(nums1[ptr1:m])

        print(temp_array)

        nums1[:] = temp_array[:]

        """
        Do not return anything, modify nums1 in-place instead.
        """


#!/usr/bin/python
# Reducer for the King James Bible
# wyu@ateneo.edu

import sys
import pandas as pd

# create a lookup table for the books
ref = open("biblebooks.txt").readlines()
book_nums = [line.split()[1] for line in ref]
book_names = [' '.join(line.split()[2::]) for line in ref]
book_ref = pd.DataFrame({
    'book name': book_names,
    'book number': book_nums,
    'count': 0
})

for input_line in sys.stdin:
    input_line = input_line.strip()
    key, value = input_line.split("\t", 1)
    value = int(value)

    book_ref.loc[book_ref['book number'] == key, 'count'] += value

    this_key = book_ref.loc[book_ref['book number'] == key]['book name'].to_list()[0] # gets the book name
    running_total = book_ref.loc[book_ref['book number']== key]['count'].to_list()[0] # gets the current count
    print("%s\t%d" % (this_key, running_total)) # displays the running total

top_books = book_ref.nlargest(n=5, columns=['count']).drop('book number', axis=1) # gets the top 5 books according to word count
print("\n\n\nThe Top Books of the King James Bible according to word count:")
print(top_books.to_string(index=False))

#Submitted by: Andrea Pilapil and Mari Valle


from random import randint, choice
import pygame as pg

class ShroomFilter(object):
    def __init__(self, size):
        size = pg.display.get_surface().get_size()
        self.surface = pg.Surface(size).convert_alpha()
        #self.surface.set_alpha(50)
        self.colors = [(0, 0, 127, 127), (127, 0, 0, 127), (127, 0, 127, 127),
                            (0, 127, 0, 127), (127, 127, 127, 127)]
        self.ticks = 0

    def update(self):
       self.ticks += 1

    def draw(self, surface):
        if not self.ticks % 5:
            self.surface.fill(choice(self.colors))

        surface.blit(self.surface, (0, 0))


#!/usr/bin/env python
import rospy

import math
import tf2_ros
import geometry_msgs.msg

def start_server():
    rospy.init_node('narkin_listener')
    tfBuffer = tf2_ros.Buffer()
    listener = tf2_ros.TransformListener(tfBuffer)
    rate = rospy.Rate(9.0)
    while not rospy.is_shutdown():
        try:
            frame_info = tfBuffer.lookup_transform('world_frame', 'base_footprint', rospy.Time())
            break
        except:
            print('Could not get world frame')
            rate.sleep()
            continue

    frame_names = tfBuffer.all_frames_as_string()
    print(frame_names)
    parent_frame= 'world_frame'
    child_frame = 'current_link'
    translation = [0,0,0]
    rotation    = [0,0,0,0]
    translation[0] = frame_info.transform.translation.x
    translation[1] = frame_info.transform.translation.y
    translation[2] = frame_info.transform.translation.z
    rotation[0] = frame_info.transform.rotation.w
    rotation[1] = frame_info.transform.rotation.z
    rotation[2] = frame_info.transform.rotation.y
    rotation[3] = frame_info.transform.rotation.z
    print(translation)

if __name__ == '__main__':
    start_server()


import requests
import pandas
import os
from bs4 import BeautifulSoup

cwd_original = os.getcwd()
output_files = cwd_original + '/data/2022/2022'

url_links = []
issue = 1
while issue != 13:
    url = f"https://link.springer.com/journal/256/volumes-and-issues/51-{issue}"
    issue = issue + 1
    url_links.append(url)

for link in range(len(url_links)):
    dates = ['-Jan','-Feb','-Mar','-Apr','-May','-Jun','-Jul','-Aug','-Sep','-Oct','-Nov','-Dec']
    page = requests.get(url_links[link])
    soup = BeautifulSoup(page.text, 'html.parser')
    Category = soup.find_all('li',class_='c-meta__item c-meta__item--block-sm-max')
    Categories = []
    for item in Category:
        C = item.text
        C = C.strip()
        C = C.replace('\n','')
        C = C.replace('Content type: ','')
        Categories.append(C)
    len(Categories)

    Title = soup.find_all('h3',class_='c-card__title')
    Titles = []
    for item in Title:
        T = item.text
        T = T.strip()
        T = T.replace('\n','')
        Titles.append(T)
    len(Titles)
    Author = soup.find_all('div',class_='u-mb-8')
    Authors = []
    for item in Author:
        A = item.text
        A = A.strip()
        A = A.replace('\n',' ')
        Authors.append(A)
    len(Authors)
    A = [x for x in Categories if not x.startswith('Published')]
    B = [x for x in A if not x.startswith('Pages')]
    list0 = [x for x in B if not x.startswith('Open Access')]
    list1 = Titles
    list2 = Authors
    df = pandas.DataFrame({'Article Class': pandas.Series(list0), 'Article Titles': pandas.Series(list1), 'List of Article Author Names': pandas.Series(list2)})

    df['List of Article Author Names'] = df['List of Article Author Names'].replace(r'(([a-z])([A-Z]))', r'\2,\3', regex=True)
    df['First Author'] = df['List of Article Author Names'].str.rsplit(',').str[0]
    df['Last Author'] = df['List of Article Author Names'].str.rsplit(',').str[-1]

    df.to_csv(output_files + dates[link] + '.csv', index=False)
    print('...finished with link: ' + url_links[link])


Import('BuildEnv')
import sys

def MapBuildDir(list):
    return map(lambda x: '#/' + Dir('.').path + '/../' + x, list)

env = BuildEnv.Clone()

env.Prepend(CPPPATH = env['TOP'])
env.Prepend(CPPPATH = env['TOP'] + '/dns')
env.Prepend(CPPPATH = env['TOP'] + '/dns/cmn')
env.Prepend(CPPPATH = env['TOP'] + '/ifmap')
env.Prepend(CPPPATH = env['TOP'] + '/schema')
env.Prepend(CPPPATH = env['TOP'] + '/sandesh/library')
env.Prepend(CPPPATH = env['TOP'] + '/bgp')
env.Prepend(CPPPATH = env['TOP'] + '/discovery/client')
env.Prepend(CPPPATH = env['TOP'] + '/http/client')

env.Prepend(LIBS=['ifmap_vnc', 'ifmap_server', 'ifmap_common', 'net', 'ifmapio',
                  'xmpp', 'ds', 'sandeshvns', 'peer_sandesh',
                  'httpc', 'curl', 'process_info',
                  'db', 'io', 'base', 'xml', 'pugixml', 'libxml2',
                  'boost_regex', 'boost_chrono', 'boost_program_options',
                  'crypto', 'ssl',
                  'boost_filesystem'])
env.Prepend(LIBS=['dns_cfg', 'cmn', 'mgr', 'agent_xmpp', 'bind_interface', 'dns_uve'])
env.Append(LIBS=['sandesh', 'http', 'http_parser'])

if sys.platform != 'darwin':
    env.Append(LIBS=['rt'])

libs = MapBuildDir(['base', 'db', 'ifmap', 'io', 'net', 'xml', 'schema', 'xmpp', 'bgp',
                    'dns/bind', 'dns/cfg', 'dns/cmn', 'dns/mgr', 'dns/agent', 'dns/uve',
                    'discovery/client', 'http/client'])
env.Append(LIBPATH=libs)

if sys.platform != 'darwin':
    env.Prepend(LINKFLAGS =
                     ['-Wl,--whole-archive',
                      '-lbase', '-lifmap_server', '-rdynamic',
                      '-Wl,--no-whole-archive'])

contrail_dns = env.Program(target = 'contrail-dns', source = ['main.cc'])
env.Alias('contrail-dns', contrail_dns)
env.Default(contrail_dns)

# Documentation
doc_files = []
doc_files += env['BASE_DOC_FILES']
doc_files += env['IO_DOC_FILES']
doc_files += env['DB_DOC_FILES']
doc_files += env['SANDESH_DOC_FILES']
doc_files += env['IFMAP_CLIENT_DOC_FILES']
doc_files += env['DISCOVERY_CLIENT_DOC_FILES']
doc_files += env['XMPP_DOC_FILES']
doc_files += env.SandeshGenDoc('uve/uve.sandesh')
doc_files += env.SandeshGenDoc('bind/bind.sandesh')

# Please update sandesh/common/vns.sandesh on process name change
env.Alias('install', env.Install(env['INSTALL_BIN'], contrail_dns))
env.Alias('install', env.Install(env['INSTALL_CONF'], 'contrail-dns.conf'))
env.Alias('install', env.Install(env['INSTALL_CONF'] +
    '/dns', 'applynamedconfig.py'))
env.Alias('install', env.Install(env['INSTALL_CONF'] +
    '/supervisord_control_files/', 'contrail-dns.ini'))
env.Alias('install', env.Install(env['INSTALL_CONF'] +
    '/supervisord_control_files/', 'contrail-named.ini'))
env.Alias('install', env.InstallAs(env['INSTALL_INITD'] +
    '/contrail-dns', 'contrail-dns.initd.supervisord'))
env.Alias('install', env.InstallAs(env['INSTALL_INITD'] +
    '/contrail-named', 'contrail-named.initd.supervisord'))
env.Alias('install', env.InstallAs(env['INSTALL_SYSTEMD'] +
    '/contrail-dns.service', 'systemd/contrail-dns.service'))
env.Alias('install', env.InstallAs(env['INSTALL_SYSTEMD'] +
    '/contrail-named.service', 'systemd/contrail-named.service'))
env.Alias('install', env.Install(env['INSTALL_MESSAGE_DOC'] +
    '/contrail-dns/', doc_files))

DnsEnv = env.Clone()
DnsEnv.CppDisableExceptions()
Export('DnsEnv')

subdirs = [
           'agent',
           'bind',
           'cfg',
           'cmn',
           'mgr',
           'uve',
           'test',
           'scripts',
          ]

for dir in subdirs:
    env.SConscript(dir + '/SConscript', exports='DnsEnv', duplicate=0)


# author: Fei Gao
#
# Interleaving String
#
# Given s1, s2, s3, find whether s3 is formed by the interleaving of s1 and s2.
# For example,
# Given:
# s1 = "aabcc",
# s2 = "dbbca",
# When s3 = "aadbbcbcac", return true.
# When s3 = "aadbbbaccc", return false.

import collections

class Solution:
    # @return a boolean
    def isInterleave(self, s1, s2, s3):
        if not s1 or not s2:
            return s3 == s1 + s2
        if len(s1) + len(s2) != len(s3):
            return False
        l1 = len(s1) + 1
        l2 = len(s2) + 1
        dp = [[False for _ in range(l2)] for _ in range(l1)]
        for i in range(l1):
            dp[i][0] = (s1[:i] == s3[:i])
        for j in range(l2):
            dp[0][j] = (s2[:j] == s3[:j])
        for i in range(1, l1):
            for j in range(1, l2):
                dp[i][j] = (dp[i-1][j] and s1[i-1] == s3[i+j-1]) or \
                    (dp[i][j-1] and s2[j-1] == s3[i+j-1])
        return dp[l1-1][l2-1]

def main():
    solver = Solution()
    tests = [['aabcc', 'dbbca', 'aadbbcbcac'],
             ['aabcc', 'dbbca', 'aadbbbaccc'],
             ["abbbbbbcabbacaacccababaabcccabcacbcaabbbacccaaaaaababbbacbb",
              "ccaacabbacaccacababbbbabbcacccacccccaabaababacbbacabbbbabc",
              "cacbabbacbbbabcbaacbbaccacaacaacccabababbbababcccbabcabbac" + \
              "cabcccacccaabbcbcaccccaaaaabaaaaababbbbacbbabacbbacabbbbabc"]]
    for test in tests:
        print(test)
        print(' ->')
        result = solver.isInterleave(*test)
        print(result)
        print('~'*10)

if __name__ == '__main__':
    main()
    pass


import os
import urllib.parse
import shutil
from time import gmtime, strftime
from distutils.dir_util import copy_tree

#Config

output_folder = 'output'
design_folder = 'design'
articles_folder = 'articles'
temp_folder = 'temp'
final_folder = '../docs'

#Functions

def createFolder(folder):
	try:
		# Create target Directory
		os.mkdir(folder)
		print("Directory " , folder ,  " Created ")
	except FileExistsError:
		print("Directory " , folder ,  " already exists")
	return

def articleUrl(article):
	article = article.replace('\'', "_")
	article = article.replace(' ', "_")
	article = urllib.parse.quote_plus(article[8:])
	article = article.replace('%', "")
	return article;

def articleTitle(article):
	return article[:-5][8:]

def mergePages(pages, outfile):
	with open(outfile, 'w') as outf:
		for fname in pages:
			with open(fname) as infile:
				for line in infile:
					outf.write(line)
	print(outfile)
	return

def deleteFolder(folder):
	if os.path.exists(folder):
		shutil.rmtree(folder)
	print("Directory " , folder ,  " deleted ")
	return

deleteFolder(temp_folder)
createFolder(temp_folder)
createFolder(output_folder)
output_folder = output_folder + '/' + strftime("%Y%m%d-%H%M%S", gmtime())
createFolder(output_folder)

#Making menu

articles = os.listdir("articles")

f = open(temp_folder + "/articlelist","w+")
f.write("
")
for i in articles:
	f.write("
" + articleTitle(i) + "
")
	print(articleTitle(i))
f.write("
")
f.close()

#Making index

filenames = [design_folder + '/header.html',
			design_folder + '/index1.html',
			temp_folder + '/articlelist',
			design_folder + '/index2.html',
			articles_folder + '/' + articles[-1],
			design_folder + '/footer.html']

mergePages(filenames, output_folder + '/index.html' )

#Making about

filenames = [design_folder + '/header.html',
			design_folder + '/about.html',
			design_folder + '/footer.html']

mergePages(filenames, output_folder + '/about.html' )

#Making articles

for article in articles:

	filenames = [design_folder + '/header.html',
	articles_folder + '/' + article,
	design_folder + '/footer.html']

	mergePages(filenames, output_folder + '/'+ articleUrl(article))

#Copying css and images

shutil.copyfile(design_folder + '/css.css', output_folder + '/css.css')
print('css.css')
shutil.copyfile(design_folder + '/phoque.png', output_folder + '/phoque.png')
print('phoque.png')
shutil.copyfile(design_folder + '/favicon.png', output_folder + '/favicon.png')
print('favicon.png')

#Clean up
deleteFolder(temp_folder)
deleteFolder(final_folder)
createFolder(final_folder)
copy_tree(output_folder, final_folder)


import requests
from requests_html import HTML
import concurrent.futures
import time
from decorators import my_time, timer
import itertools
import logging

step = 20
fromc = 421
toc = 720

counter = itertools.count()
t1 = time.time()

#Directory to save the downloaded files
dir = r'D:\One_Piece\\' #include an extra '\'
homepage = requests.get('https://www.mangapanda.com/one-piece')
titles = HTML(html = homepage.text)
titles = titles.find('td')
titles = titles[22:-4:2]

def getchapter(chapter):
    """To change chapter number into desired format for saving"""
    chapter = str(chapter)
    if int(chapter) < 10:
        chapter = '00' + chapter
    elif int(chapter) < 100:
        chapter = '0' + chapter
    return chapter

def getpage(page):
    """To change pages number into desired format for saving"""
    page = str(page)
    if int(page) < 10:
        page = '0' + page
    return page

def download_img(source, name):
    # print(f'Source: {source}')
    # print(f'Name: {name}')
    img = requests.get(source)

    with open(name,'wb') as file:
        file.write(img.content)

    pages = next(counter)

@my_time
def download_chapters(chapter):
    site = 'https://www.mangapanda.com'
    link = '/one-piece/'+str(chapter)
    mangalink = requests.get(site+link)
    html = HTML(html = mangalink.text)
    article = html.find('div#selectpage')
    pages = int(article[0].text.split()[-1])
    for page in range(1,pages+1):
        title = titles[chapter-1].text.split(': ')[-1]
        if title.endswith(':'):
            title = title[:-5]
        # print('Parsing Chapter '+str(chapter)+'. '+title+' Page '+str(page))
        if page != 1:
            mangalink = requests.get(site+link)
            html = HTML(html = mangalink.text)
        image = html.find('div#imgholder',first = True)
        img_src = image.find('img',first = True)
        img_src = img_src.attrs['src']
        # img = requests.get(img_src)
        img_name = dir+getchapter(chapter)+' '+title+' '+getpage(page)+'.jpg'
        # download_img(img_src, img_name)
        with concurrent.futures.ThreadPoolExecutor() as runner:
            runner.submit(download_img, img_src, img_name)

        nextpage = image.find('a', first = True)
        link = nextpage.attrs['href']
    # print(f'---------Chapter {chapter} completed downloading----------')

def download(fromc = 1, toc = 1, step = 10):
    start = end =  fromc
    if toc < fromc:
        toc = fromc
    while end <= toc:
        end += step
        if end > toc:
            end = toc
        chaps = [i for i in range(start, end + 1)]
        print(f'Downloading Chapters {start} to {end}')
        with concurrent.futures.ThreadPoolExecutor() as executor:
            executor.map(download_chapters, chaps)
            # print('Download complete')
        start += step
        if end == toc:
            break

download(fromc, toc, step = step)
print ('--------------------------------Done--------------------------------------------------------------------')

t2 = time.time()
pages = next(counter) - 1

line = f'Downloaded from Chapter {fromc} to Chapter {toc} at {step} Chapters at a time'
timer.info(line)
line = f'A total {pages-1} pages downloaded in {t2-t1} seconds'
timer.info(line)


import json
import os

def read_txt(sentences):
    sentences = sentences.split("\n")
    pairs = []
    for pair in sentences:
        if pair:
            rec = pair.split("\t")
            pairs.append([int(rec[0]), int(rec[1])])
    return pairs

def read_precision(sentences, map_dict):
    sentences = sentences.split("\n")
    pairs = []
    all_precision = []
    for pair in sentences:
        if pair:
            rec = pair.split("\t")
            if rec:
                r1 = map_dict[0][rec[0]]
                r2 = map_dict[1][rec[1]]
                item = [r1, r2]
                if item not in all_precision:
                    all_precision.append(item)
    return pairs, all_precision

def read_dir(path):
    block_file_path = os.path.join(path, "block_map.json")
    prec_file_path = os.path.join(path, "precision.txt")

    if os.path.exists(block_file_path):
        with open(block_file_path, "r+", encoding="utf-8") as f:
            map_dict = json.loads(f.read())

        if os.path.exists(prec_file_path):
            with open(prec_file_path, "r+", encoding="utf-8") as f:
                s = f.read()
                pairs, precisions = read_precision(s, map_dict)
            return precisions
        else:
            print(prec_file_path, "doesn't exist")
    else:
        print(block_file_path, "doesn't exist")
    return None

def cal_prec(dataset, blocks):
    path = "../../data_block/%s/%d" % (dataset, blocks)

    all_precision = []
    for i in range(1, blocks+1):
        block_path = os.path.join(path, str(i))
        precisions = read_dir(block_path)
        all_precision += precisions

    with open("../../data_block/%s/%d/all_precision.json" % (dataset, blocks), "w+", encoding="utf-8") as f:
        f.write(json.dumps(all_precision))


from sqlalchemy import create_engine, update, Column, Float, String, Integer, TIMESTAMP, ForeignKey, func
from sqlalchemy.orm import sessionmaker, relationship, backref
from sqlalchemy.ext.declarative import declarative_base
from datetime import datetime

# docker inspect <container_id>
engine = create_engine('mysql+mysqlconnector://root:root@172.17.0.2:3306/projeto', echo=True)

Session = sessionmaker(bind=engine)
Base = declarative_base()
Base.metadata.create_all(engine)

class Compras(Base):
    __tablename__ = 'compras'
    id = Column(Integer, primary_key=True, autoincrement=True)
    timestamp = Column(TIMESTAMP(timezone=False), nullable=False, default=datetime.now())

class Produtos(Base):
    __tablename__ = 'produtos'
    id = Column(Integer, primary_key=True, autoincrement=True)
    nome = Column(String(100))
    preco = Column(Float)

class Transacoes(Base):
    __tablename__ = 'transacoes'
    id = Column(Integer, primary_key=True, autoincrement=True)
    quantidade = Column(Integer)
    valor_total = Column(Float)
    id_compra = Column(Integer, ForeignKey("compras.id"))
    id_produto = Column(Integer, ForeignKey("produtos.id"))

def initialise() -> None:
    Base.metadata.create_all(engine)

def cadastro_compra(compra: dict) -> None:
    session = Session()
    nova_compra = Compras()
    session.add(nova_compra)
    for produto in compra["produtos"]:
        novo_produto = session.query(Produtos).filter(Produtos.id == produto["id"]).one_or_none()
        if not novo_produto:
            novo_produto = Produtos(
                id = produto["id"],
                nome = produto["nome"],
                preco =produto["preco"]
            )
            session.add(novo_produto)

        nova_transacao = Transacoes(
            quantidade = produto["quantidade"],
            valor_total = produto["quantidade"] * produto["preco"],
            id_compra = nova_compra.id,
            id_produto = novo_produto.id
        )
        session.add(nova_transacao)
    session.commit()
    session.close()
    return

def relatorio_por_produto() -> dict:
    session = Session()
    vendas = {
        "vendas_por_produto": []
    }
    for venda in session.query(Transacoes).join(Produtos).with_entities(
        Produtos.nome,
        Transacoes.id_produto,
        func.sum(Transacoes.quantidade).cast(Integer).label('quantidade_total'),
        func.sum(Transacoes.valor_total).label('valor_total')
    ).group_by(Transacoes.id_produto, Produtos.nome).all():
        vendas["vendas_por_produto"].append(dict(venda))
    return vendas

def relatorio_total_vendas() -> dict:
    session = Session()
    vendas = {
        "total_vendas": None
    }

    for venda in session.query(Transacoes).with_entities(
        func.sum(Transacoes.quantidade).cast(Integer).label('qtd_produtos_vendidos'),
        func.sum(Transacoes.valor_total).label('valor_total')
    ).all():
        vendas["total_vendas"] = dict(venda)

    vendas["total_vendas"]["qtd_vendas"] = session.query(Compras).count()

    return vendas

def relatorio_rank_vendidos() -> dict:
    session = Session()
    vendas = {
        "rank_vendidos": []
    }

    for venda in session.query(Transacoes).join(Produtos).with_entities(
        Produtos.id,
        Produtos.nome,
        func.sum(Transacoes.quantidade).cast(Integer).label('qtd_produtos_vendidos'),
        func.sum(Transacoes.valor_total).label('valor_total')
    ).group_by(Produtos.id, Produtos.nome).order_by(func.sum(Transacoes.quantidade).cast(Integer).desc()).all():
        vendas["rank_vendidos"].append(dict(venda))

    return vendas

def limpar_vendas():
    session = Session()
    session.query(Transacoes).delete()
    session.query(Compras).delete()
    session.commit()
    session.close()

def consultar_vendas():
    session = Session()
    compras = {
        "compras": []
    }
    for compra in session.query(Compras).all():
        compras["compras"].append({
            "id": compra.__dict__["id"],
            "items": [],
            "valor_total": 0
        })
    for transacao in session.query(Transacoes).all():
        for compra in compras["compras"]:
            if transacao.__dict__["id_compra"] == compra["id"]:
                compra["items"].append({
                    "id": transacao.__dict__["id_produto"],
                    "quantidade": transacao.__dict__["quantidade"]
                })
                compra["valor_total"] += transacao.__dict__["valor_total"]
    for produto in session.query(Produtos).all():
        for compra in compras["compras"]:
            for item in compra["items"]:
                if item["id"] == produto.__dict__["id"]:
                    item["nome"] = produto.__dict__["nome"]
                    item["preco"] = produto.__dict__["preco"]

    session.close()
    return compras


#!/usr/bin/env python
'''
usage: certificates [-h] [--output OUTPUT] participants template

positional arguments:
  participants          csv filaname containing participants
  template              certificate template in svg format used to build

optional arguments:
  -h, --help            show this help message and exit
  --output OUTPUT, -o OUTPUT
                        destination of the generated certificates
'''
import csv
import os
from pathlib import Path

def import_from_csv(csv_filename):
    with open(csv_filename) as file:
        table = csv.DictReader(file)
        for row in table:
            yield row

def make_certificates(participants_path, template_path, output_path):
    # svg used as template
    base_svg = Path(template_path)
    # retrieve participants list
    participants = import_from_csv(participants_path)
    # Create output directory if not exists
    output_directory = Path(output_path)
    output_directory.mkdir(exist_ok=True)

    for participant in participants:
        new_svg = base_svg.read_text(encoding='utf-8').format(**participant)
        output_filename = output_directory / f'{participant["name"]}.svg'
        output_svg = Path(output_filename)
        output_svg.write_text(new_svg)
        png_filename = str(output_filename).replace("svg", "png")
        os.system(f'inkscape -z -e "{png_filename}" "{output_filename}"')
        output_svg.unlink()


# nohup python -u activeAuthorExtraction.py > activeAuthorExtractionFoRCAS.log 2>&1 &

import jsonlines
from tqdm import tqdm
import pickle as pk
# 抽取比较活跃的作者

# 考虑作者位次的作者发文数量，必须大于1
def AuthorCount(YearPublist, authorRank):
# yearPublist  {2023:{'rank1':set(), 'rankLast':set(), 'others':set() }, ...}
    count = 0
    if authorRank != 'both':
        for year, publist in YearPublist.items():
            if len(publist[authorRank]) > 0:
                count += 1
    else:
        for year, publist in YearPublist.items():
            if len(publist['rank1']) > 0 or len(publist['rankLast']) > 0:
                count += 1
    return count
# 按时间排序
def AuthorPublistTimeRank(YearPublist):
    YearPublistRanked = {}
    for k,v in sorted(YearPublist.items()):
        # 如果一个人的发文序列中，有不存在的年份，直接给这个人drop掉（数量应该极少）
        if int(k) < 1801 or int(k) > 2021:
            return ''
        YearPublistRanked[k] = v
    return YearPublistRanked

# 开始年份
def commonstartYearpublist(publist):
    return int(list(publist.keys())[0])
# 休止年份
def commonendYearpublist(publist):
    return int(list(publist.keys())[-1])
# 年份跨度
def commonspanYearpublist(publist):
    return commonendYearpublist(publist) - commonstartYearpublist(publist)

# 这里抽出来所有曾经至少2年，以第一作者或末位作者身份发表过论文的作者，全部的发文序列。
with jsonlines.open('../../../DataCrossBoundaryPerspective_InterdisciplinaryResearch/activeAuthorSeqForCAS.jsonl', mode='w') as writer:
    with jsonlines.open('../../../DataCrossBoundaryPerspective_InterdisciplinaryResearch/CASauthorSeq.jsonl', mode='r') as reader:
        for lines in tqdm(reader):
            for authorID, YearPublist in lines.items():
                # 至少有2年以第一作者身份或末位作者身份，发论文
                if AuthorCount(YearPublist,'both') > 1:
                    #按时间顺序排序，同时丢掉离谱的时间
                    temp = AuthorPublistTimeRank(YearPublist)
                    if temp == '':
                        continue
                    # 大于80显然不是正常的学术生涯轨迹，去掉
                    if commonspanYearpublist(temp) < 81:
                        writer.write({authorID: temp})
# print(authornumPubYearsFreq)
# print(Rank1authornumPubYearsFreq)
# print(RankLasTauthornumPubYearsFreq)
# print(RankIMportantAuthornumPubYearsFreq)

# pk.dump(authornumPubYearsFreq, open('../../../DataCrossBoundaryPerspective_InterdisciplinaryResearch/authornumPubYearsFreq.pk', 'wb'))
# pk.dump(Rank1authornumPubYearsFreq, open('../../../DataCrossBoundaryPerspective_InterdisciplinaryResearch/Rank1authornumPubYearsFreq.pk', 'wb'))
# pk.dump(RankLasTauthornumPubYearsFreq, open('../../../DataCrossBoundaryPerspective_InterdisciplinaryResearch/RankLasTauthornumPubYearsFreq.pk', 'wb'))
# pk.dump(RankIMportantAuthornumPubYearsFreq, open('../../../DataCrossBoundaryPerspective_InterdisciplinaryResearch/RankIMportantAuthornumPubYearsFreq.pk', 'wb'))


import pandas as pd
import requests
import wfdb
import os
from multiprocessing import cpu_count
from multiprocessing.pool import ThreadPool
from tqdm import tqdm
import scipy

from DataHandlers.DataProcessUtilities import *

base_url = "https://physionet.org/files/icentia11k-continuous-ecg/1.0/"
dataset_path = r"C:\Users\daniel\Documents\CambridgeSoftwareProjects\ecg-signal-quality\Datasets\Icentia11k"

data_max_chunk_size = 100000

def download_file(args):
    url, path = args
    try:
        r = requests.get(url, params={"download":True})
        if r.status_code != 200:
            print(f"Error: {r.reason}")
            return
        with open(path, "wb") as f:
            f.write(r.content)
    except requests.exceptions.RequestException as e:
        print(e)

def find_files(patient_range, types=["atr", "hea", "dat"], n_per_patient = 50):
    urls = []
    filepaths = []
    record_names = []

    for i in range(patient_range[0], patient_range[1]):
        p_num = str(i).rjust(5, "0")
        path_ext = f"p{p_num[:2]}/p{p_num}/"

        # make the patient directories
        if not os.path.isdir(os.path.join(dataset_path, f"p{p_num[:2]}")):
            os.mkdir(os.path.join(dataset_path, f"p{p_num[:2]}"))
        if not os.path.isdir(os.path.join(dataset_path, f"p{p_num[:2]}", f"p{p_num}")):
            os.mkdir(os.path.join(dataset_path, f"p{p_num[:2]}", f"p{p_num}"))

        for j in range(n_per_patient):
            s_num = str(j).rjust(2, "0")
            record_names.append(os.path.join(dataset_path,
                                               f"p{p_num[:2]}",
                                               f"p{p_num}",
                                               f"p{p_num}_s{s_num}"))
            for t in types:
                section_file = f"p{p_num}_s{s_num}.{t}"
                record_url = base_url + path_ext + section_file
                record_filepath = os.path.join(dataset_path,
                                               f"p{p_num[:2]}",
                                               f"p{p_num}",
                                               f"p{p_num}_s{s_num}.{t}")

                urls.append(record_url)
                filepaths.append(record_filepath)

    return urls, filepaths, record_names

def download_parallel(urls, filepaths):
    cpus = cpu_count()
    for _ in tqdm(ThreadPool(cpus - 1).imap_unordered(download_file, zip(urls, filepaths)), total=len(urls)):
        pass

def split_recording(data, cut_len, data_len=None):
    cut_samp = cut_len * data.fs
    if data_len is None:
        data_len = data.data.shape[0]
    splits = np.arange(0, data_len, cut_samp)

    data_splits = []

    afib = False
    aflutter = False

    # Note here the split array includes 0 so we get an empty array in data_splits[0],
    # which is filtered so the rest of the diags align with the data

    for i, (start, end) in enumerate(zip(splits, splits[1:])):
        ann_notes = data.notes[np.logical_and(data.r_peaks_original >= start, data.r_peaks_original < end)]

        if "(N" in ann_notes:
            afib = False
            aflutter = False
        elif "(AFIB" in ann_notes:
            afib = True
            aflutter = False
        elif "(AFL" in ann_notes:
            afib = False
            aflutter = True

        data_split = data.copy()
        data_split.data = data_split.data[start:end]
        data_split.r_peaks_original = data_split.r_peaks_original[np.logical_and(data.r_peaks_original >= start, data.r_peaks_original < end)] - start
        data_split.r_peaks = data_split.r_peaks[np.logical_and(data.r_peaks >= start, data.r_peaks < end)] - start
        data_split.notes = ann_notes
        data_split.annotations = data.symbols[np.logical_and(data.r_peaks_original >= start, data.r_peaks_original < end)]

        if afib:
            data_split["class_index"] = 1
        elif aflutter:
            data_split["class_index"] = 2
        else:
            data_split["class_index"] = 0

        data_split["length"] = cut_samp
        data_splits.append(data_split)

    return pd.DataFrame(data_splits)

def load_dataset_scratch(record_names, split_len, load_data=True):
    long_df = []

    for n in record_names:
        try:
            ann = wfdb.rdann(n, "atr")
            if load_data:
                rec = wfdb.rdrecord(n)
            else:
                rec = wfdb.rdheader(n)
                data = np.array([])

            adc_gain = rec.adc_gain[0]
            fs = rec.fs

            ptID = int(rec.file_name[0][1:6])
            sample_id = int(rec.file_name[0][8:10])

            r_peaks = np.array(ann.sample)
            symbols = np.array(ann.symbol)
            notes = np.array(ann.aux_note)

            long_df.append(pd.Series({"data": data,
                                      "adc_gain": adc_gain,
                                      "fs": fs,
                                      "r_peaks_original": r_peaks,
                                      "symbols": symbols,
                                      "notes": notes,
                                      "ptID": ptID,
                                      "sampleID": sample_id}))
        except (FileNotFoundError, FileExistsError):
            print(f"Error reading {n}")

    print("Completed reading")

    long_df = pd.DataFrame(long_df)
    long_df = long_df.dropna()
    print(long_df.data.map(lambda x: x.shape))

    print("Processing Data")
    if load_data:
        long_df = process_data(long_df)
    print(long_df.head())

    short_dfs = []

    print("Cutting to short segments")
    for i, long_sample in long_df.iterrows():
        short_dfs.append(split_recording(long_sample, split_len))

    short_df = pd.concat(short_dfs, ignore_index=True)
    print(short_df.head())
    return short_df

def process_data(ecg_data, f_low=0.67, f_high=30, resample_rate=300):
    # perform band pass filtering, notch filtering (for power line interference)
    fs_list = pd.unique(ecg_data["fs"])
    bandpass_dict = {fs: scipy.signal.butter(3, [f_low, f_high], 'bandpass', fs=fs, output='sos') for fs in fs_list}
    notch_dict = {fs: scipy.signal.butter(3, [48, 52], 'bandstop', fs=fs, output='sos') for fs in fs_list}

    ecg_data["data"] = ecg_data.apply(lambda x: filter_and_norm(x["data"], bandpass_dict[x["fs"]]), axis=1)
    ecg_data["data"] = ecg_data.apply(lambda x: filter_and_norm(x["data"], notch_dict[x["fs"]]), axis=1)

    ecg_data["data"] = ecg_data.apply_parallel(lambda x: resample(x["data"], resample_rate, x["fs"]), axis=1)
    ecg_data["length"] = ecg_data["data"].map(lambda x: x.shape[-1])
    ecg_data["fs"] = resample_rate

    # Scale the R peak positions to the new samplerate
    ecg_data["r_peaks_original"] = ecg_data["r_peaks_original"] * resample_rate/ecg_data["fs"]

    # Get r peak positions and heartrate
    ecg_data["r_peaks"] = ecg_data.apply_parallel(get_r_peaks)
    ecg_data["heartrate"] = ecg_data.apply(lambda e: (len(e["r_peaks"]) / (e["length"] / e["fs"])) * 60,
                                                       axis=1)
    # Get the rri feature
    ecg_data["rri_feature"] = (ecg_data["r_peaks"] / resample_rate).map(lambda x: get_rri_feature(x, 60))
    fewer_5_beats = ecg_data["rri_feature"].map(lambda x: np.sum(x == 0) > 55)
    ecg_data = ecg_data[~fewer_5_beats]
    ecg_data["rri_len"] = ecg_data["rri_feature"].map(lambda x: x[x > 0].shape[-1])
    ecg_data["rri_feature"] = normalise_rri_feature(ecg_data["rri_feature"])

    return ecg_data

def load_dataset_pickle(f_name):
    try:
        end_path = f"{f_name}.pk"
        ecg_data = pd.read_pickle(os.path.join(dataset_path, end_path))
        return ecg_data
    except (OSError, FileNotFoundError) as e:
        print(e)
    return

def load_dataset(save_name="dataframe", force_reload=False, pt_range=(0, 11000), sections_per_pt=50, split_size=30, download=False, load_data=True):
    dataset = None

    if not force_reload:
        dataset = load_dataset_pickle(save_name)

    if dataset is None:
        urls, filenames, record_names = find_files(pt_range, n_per_patient=sections_per_pt)

        if download:
            download_parallel(urls, filenames)

        dataset = load_dataset_scratch(record_names, split_size, load_data)
        dataset.to_pickle(os.path.join(dataset_path, f"{save_name}.pk"))

    dataset = dataset[(dataset["ptID"] < pt_range[1]) & (dataset["ptID"] >= pt_range[0])]
    return dataset

if __name__ == "__main__":
    pass


from core.k8mars_control import *
import os
from pathlib import Path

class PathInfo:
    """
    Path of directory always ends with '/'
    """
    ### Working Directory ###
    work_dir = Path(os.getcwd())
    k8mars_dir = Path(os.path.abspath(os.path.join(os.getcwd(), '..')))
    yaml_templates_dir = work_dir / 'yaml_templates/'
    cache_folder = work_dir / 'k8marsCache/'

    ### data location ###

class Wzlk8toolkitMain:
    def __init__(self):
        self.__manager = K8marsJobManager()

    def __display_main_menu(self):
        print('1) Create a new job.')
        print('2) Existing job.')
        print('3) Exit.')

    def __select_main_menu(self):
        item = input('please select: ')
        if item == '1':
            self.__create_job()
        if item == '2':
            self.__manager.show_jobs_by_index()
            selected_job_obj = self.__manager.select_job()
            kubectl_controller_obj = KubectlController(selected_job_obj)
            kubectl_controller_obj.show_options_for_job()
            kubectl_controller_obj.select_option_for_job()
            # show list of jobs with index : done
            # menu for entry points,
            # edit detail of yaml file
            # send kubectl cmd to cluster
            # copy dataset into persistent volume
            pass
        if item == '3':
            exit()

    def main(self):
        while True:
            self.__display_main_menu()
            self.__select_main_menu()

    def __create_job(self):
        K8marsJobCreator_obj = K8marsJobCreator()
        K8marsJobCreator_obj.write_yaml_file()
        self.__manager.add_new_job(K8marsJobCreator_obj)

if __name__ == '__main__':
    main = Wzlk8toolkitMain()
    main.main()


from asyncio.windows_events import NULL
import json
import tqdm

filename = 'trs\szwb-0.txt'
json_filename = 'json\szwb-0-record-dict-utf8.json'
null_record = {
    'title': NULL,
    'author': NULL,
    'time': NULL,
    'content': NULL,
}
record_list = []

with open(filename, encoding='UTF-8') as file:

    json_fp = open(json_filename, "a", encoding='UTF-8')
    print('open with success')
    data_record = {
        'title': NULL,
        'author': NULL,
        'time': NULL,
        'content': NULL,
    }
    for line in tqdm.tqdm(file):

        if line == '<REC>\n' and data_record != null_record:
            record_list.append(data_record)

            data_record = {
                'title': NULL,
                'author': NULL,
                'time': NULL,
                'content': NULL,
            }
        if line.startswith('<title>='):
            data_record['title'] = line.removeprefix('<title>=')
        if line.startswith('<authors>='):
            data_record['author'] = line.removeprefix('<authors>=')
        if line.startswith('<pub_time>='):
            data_record['time'] = line[11:21].replace('/', '-')
        if line.startswith('<content>='):
            data_record['content'] = line.removeprefix('<content>=')
    json.dump(record_list, json_fp, indent=4, separators=(',', ': '))


import io
import sys

_INPUT = """\
6
4
1 3 2 4
10
5 8 2 2 1 6 7 2 9 10
"""

sys.stdin = io.StringIO(_INPUT)
case_no=int(input())
for __ in range(case_no):
  N=int(input())
  A=list(map(lambda x: int(x)-1,input().split()))
  n=len([1 for i in range(N) if A[i]==i])
  ans=n*(n-1)//2
  m=0
  for i in range(N):
    if A[i]!=i and A[A[i]]==i: m+=1
  print(ans+m//2)


import heapq
import sys
input = sys.stdin.readline

N = int(input().strip())

left = []  # 최대힙
right = [] # 최소힙
for _ in range(N):
    number = int(input().strip())

    heapq.heappush(left, -number) if len(left) == len(right) else heapq.heappush(right, number)
    # 힙큐에 left right 길이를 맞춰주며 넣어준다.

    if right and right[0] < -left[0]:
        # left 제일 큰거랑 right 제일 큰거랑 비교해서 left보다 right가 작다면 서로 바꿔준다.
        to_left = heapq.heappop(right)
        to_right = heapq.heappop(left)
        # 각각 최소힙, 최대힙에서 뽑아온 값이니 반대로 넣어줄 때 부호에 유의한다.
        heapq.heappush(left, -to_left)
        heapq.heappush(right, -to_right)

    print(-left[0]) # 출력은 항상 left의 최대힙값이다.

백준/Gold/1655. 가운데를 말해요/가운데를 말해요.py