# Secondly bootstrapped from:
# https://github.com/cloneofsimo/lora
# whatever man I wrote it anyways

# Bootstrapped from:
# https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth.py
import itertools
import math
import os
import re
from typing import Optional, List, Literal

import torch
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.checkpoint
from diffusers import (
    AutoencoderKL,
    DDPMScheduler,
    StableDiffusionPipeline,
    UNet2DConditionModel,
)
from diffusers.optimization import get_scheduler

from PIL import Image

from tqdm.auto import tqdm
from transformers import CLIPTextModel, CLIPTokenizer
import wandb
import fire

from lora_diffusion import (
    save_all,
    prepare_clip_model_sets,
    evaluate_pipe,
)

from .dataset import PPPDataset, pp_extend
from .utils import PPPAttenProc

def get_models(
    pretrained_model_name_or_path,
    pretrained_vae_name_or_path,
    revision,
    placeholder_tokens: List[str],
    initializer_tokens: List[str],
    device="cuda:0",
):

    tokenizer = CLIPTokenizer.from_pretrained(
        pretrained_model_name_or_path,
        subfolder="tokenizer",
        revision=revision,
    )

    text_encoder = CLIPTextModel.from_pretrained(
        pretrained_model_name_or_path,
        subfolder="text_encoder",
        revision=revision,
    )

    placeholder_token_ids = []

    for token, init_tok in zip(placeholder_tokens, initializer_tokens):
        num_added_tokens = tokenizer.add_tokens(token)
        if num_added_tokens == 0:
            raise ValueError(
                f"The tokenizer already contains the token {token}. Please pass a different"
                " `placeholder_token` that is not already in the tokenizer."
            )

        placeholder_token_id = tokenizer.convert_tokens_to_ids(token)

        placeholder_token_ids.append(placeholder_token_id)

        # Load models and create wrapper for stable diffusion

        text_encoder.resize_token_embeddings(len(tokenizer))
        token_embeds = text_encoder.get_input_embeddings().weight.data
        if init_tok.startswith("<rand"):
            # <rand-"sigma">, e.g. <rand-0.5>
            sigma_val = float(re.findall(r"<rand-(.*)>", init_tok)[0])

            token_embeds[placeholder_token_id] = (
                torch.randn_like(token_embeds[0]) * sigma_val
            )
            print(
                f"Initialized {token} with random noise (sigma={sigma_val}), empirically {token_embeds[placeholder_token_id].mean().item():.3f} +- {token_embeds[placeholder_token_id].std().item():.3f}"
            )
            print(f"Norm : {token_embeds[placeholder_token_id].norm():.4f}")

        elif init_tok == "<zero>":
            token_embeds[placeholder_token_id] = torch.zeros_like(token_embeds[0])
        else:
            token_ids = tokenizer.encode(init_tok, add_special_tokens=False)
            # Check if initializer_token is a single token or a sequence of tokens
            if len(token_ids) > 1:
                raise ValueError("The initializer token must be a single token.")

            initializer_token_id = token_ids[0]
            token_embeds[placeholder_token_id] = token_embeds[initializer_token_id]

    vae = AutoencoderKL.from_pretrained(
        pretrained_vae_name_or_path or pretrained_model_name_or_path,
        subfolder=None if pretrained_vae_name_or_path else "vae",
        revision=None if pretrained_vae_name_or_path else revision,
    )
    unet = UNet2DConditionModel.from_pretrained(
        pretrained_model_name_or_path,
        subfolder="unet",
        revision=revision,
    )

    unet.set_attn_processor(PPPAttenProc())

    return (
        text_encoder.to(device),
        vae.to(device),
        unet.to(device),
        tokenizer,
        placeholder_token_ids,
    )

@torch.no_grad()
def text2img_dataloader(
    train_dataset,
    train_batch_size,
    tokenizer,
    vae,
    text_encoder,
    cached_latents: bool = False,
):

    if cached_latents:
        cached_latents_dataset = []
        for idx in tqdm(range(len(train_dataset))):
            batch = train_dataset[idx]
            # rint(batch)
            latents = vae.encode(
                batch["instance_images"].unsqueeze(0).to(dtype=vae.dtype).to(vae.device)
            ).latent_dist.sample()
            latents = latents * 0.18215
            batch["instance_images"] = latents.squeeze(0)
            cached_latents_dataset.append(batch)

    def collate_fn(examples):
        input_ids = [example["instance_prompt_ids"] for example in examples]
        pixel_values = [example["instance_images"] for example in examples]
        pixel_values = torch.stack(pixel_values)
        pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float()

        input_ids = tokenizer.pad(
            {"input_ids": input_ids[0]},
            padding="max_length",
            max_length=tokenizer.model_max_length,
            return_tensors="pt",
        ).input_ids

        batch = {
            "input_ids": input_ids,
            "pixel_values": pixel_values,
        }

        if examples[0].get("mask", None) is not None:
            batch["mask"] = torch.stack([example["mask"] for example in examples])

        return batch

    if cached_latents:

        train_dataloader = torch.utils.data.DataLoader(
            cached_latents_dataset,
            batch_size=train_batch_size,
            shuffle=True,
            collate_fn=collate_fn,
        )

        print("PTI : Using cached latent.")

    else:
        train_dataloader = torch.utils.data.DataLoader(
            train_dataset,
            batch_size=train_batch_size,
            shuffle=True,
            collate_fn=collate_fn,
        )

    return train_dataloader

def loss_step(
    batch,
    unet,
    vae,
    text_encoder,
    scheduler,
    train_inpainting=False,
    t_mutliplier=1.0,
    mixed_precision=False,
    mask_temperature=1.0,
    cached_latents: bool = False,
    is_coarse_inversion: bool = False,
):
    weight_dtype = torch.float32
    if not cached_latents:
        latents = vae.encode(
            batch["pixel_values"].to(dtype=weight_dtype).to(unet.device)
        ).latent_dist.sample()
        latents = latents * 0.18215

        if train_inpainting:
            masked_image_latents = vae.encode(
                batch["masked_image_values"].to(dtype=weight_dtype).to(unet.device)
            ).latent_dist.sample()
            masked_image_latents = masked_image_latents * 0.18215
            mask = F.interpolate(
                batch["mask_values"].to(dtype=weight_dtype).to(unet.device),
                scale_factor=1 / 8,
            )
    else:
        latents = batch["pixel_values"]

        if train_inpainting:
            masked_image_latents = batch["masked_image_latents"]
            mask = batch["mask_values"]

    noise = torch.randn_like(latents)
    bsz = latents.shape[0]

    timesteps = torch.randint(
        0,
        int(scheduler.config.num_train_timesteps * t_mutliplier),
        (bsz,),
        device=latents.device,
    )
    timesteps = timesteps.long()

    noisy_latents = scheduler.add_noise(latents, noise, timesteps)

    if train_inpainting:
        latent_model_input = torch.cat(
            [noisy_latents, mask, masked_image_latents], dim=1
        )
    else:
        latent_model_input = noisy_latents

    encoder_hidden_states = text_encoder(batch["input_ids"].to(text_encoder.device))[0]

    if not is_coarse_inversion:
        _hs = {"this_idx": 0}

        for idx in range(16):
            _hs[f"CONTEXT_TENSOR_{idx}"] = encoder_hidden_states[idx : idx + 1, :, :]
    else:
        _hs = encoder_hidden_states

    model_pred = unet(latent_model_input, timesteps, _hs).sample

    if scheduler.config.prediction_type == "epsilon":
        target = noise
    elif scheduler.config.prediction_type == "v_prediction":
        target = scheduler.get_velocity(latents, noise, timesteps)
    else:
        raise ValueError(f"Unknown prediction type {scheduler.config.prediction_type}")

    if batch.get("mask", None) is not None:

        mask = (
            batch["mask"]
            .to(model_pred.device)
            .reshape(
                model_pred.shape[0], 1, model_pred.shape[2] * 8, model_pred.shape[3] * 8
            )
        )
        # resize to match model_pred
        mask = F.interpolate(
            mask.float(),
            size=model_pred.shape[-2:],
            mode="nearest",
        )

        mask = (mask + 0.01).pow(mask_temperature)

        mask = mask / mask.max()

        model_pred = model_pred * mask

        target = target * mask

    loss = (
        F.mse_loss(model_pred.float(), target.float(), reduction="none")
        .mean([1, 2, 3])
        .mean()
    )

    return loss

def train_inversion(
    unet,
    vae,
    text_encoder,
    dataloader,
    num_steps: int,
    scheduler,
    index_no_updates,
    optimizer,
    save_steps: int,
    placeholder_token_ids,
    placeholder_tokens,
    save_path: str,
    tokenizer,
    lr_scheduler,
    test_image_path: str,
    cached_latents: bool,
    accum_iter: int = 1,
    log_wandb: bool = False,
    wandb_log_prompt_cnt: int = 10,
    class_token: str = "person",
    train_inpainting: bool = False,
    mixed_precision: bool = False,
    clip_ti_decay: bool = True,
    is_coarse_inversion: bool = False,
):

    progress_bar = tqdm(range(num_steps))
    progress_bar.set_description("Steps")
    global_step = 0

    # Original Emb for TI
    orig_embeds_params = text_encoder.get_input_embeddings().weight.data.clone()

    if log_wandb:
        preped_clip = prepare_clip_model_sets()

    index_updates = ~index_no_updates
    loss_sum = 0.0

    for epoch in range(math.ceil(num_steps / len(dataloader))):
        unet.eval()
        text_encoder.train()
        for batch in dataloader:

            lr_scheduler.step()

            with torch.set_grad_enabled(True):
                loss = (
                    loss_step(
                        batch,
                        unet,
                        vae,
                        text_encoder,
                        scheduler,
                        train_inpainting=train_inpainting,
                        mixed_precision=mixed_precision,
                        cached_latents=cached_latents,
                        is_coarse_inversion=is_coarse_inversion,
                    )
                    / accum_iter
                )

                loss.backward()
                loss_sum += loss.detach().item()

                if global_step % accum_iter == 0:
                    # print gradient of text encoder embedding
                    print(
                        text_encoder.get_input_embeddings()
                        .weight.grad[index_updates, :]
                        .norm(dim=-1)
                        .mean()
                    )
                    optimizer.step()
                    optimizer.zero_grad()

                    with torch.no_grad():

                        # normalize embeddings
                        if clip_ti_decay:
                            pre_norm = (
                                text_encoder.get_input_embeddings()
                                .weight[index_updates, :]
                                .norm(dim=-1, keepdim=True)
                            )

                            lambda_ = min(1.0, 100 * lr_scheduler.get_last_lr()[0])
                            text_encoder.get_input_embeddings().weight[
                                index_updates
                            ] = F.normalize(
                                text_encoder.get_input_embeddings().weight[
                                    index_updates, :
                                ],
                                dim=-1,
                            ) * (
                                pre_norm + lambda_ * (0.4 - pre_norm)
                            )
                            print(pre_norm)

                        current_norm = (
                            text_encoder.get_input_embeddings()
                            .weight[index_updates, :]
                            .norm(dim=-1)
                        )

                        text_encoder.get_input_embeddings().weight[
                            index_no_updates
                        ] = orig_embeds_params[index_no_updates]

                        print(f"Current Norm : {current_norm}")

                global_step += 1
                progress_bar.update(1)

                logs = {
                    "loss": loss.detach().item(),
                    "lr": lr_scheduler.get_last_lr()[0],
                }
                progress_bar.set_postfix(**logs)

            if global_step % save_steps == 0:
                save_all(
                    unet=unet,
                    text_encoder=text_encoder,
                    placeholder_token_ids=placeholder_token_ids,
                    placeholder_tokens=placeholder_tokens,
                    save_path=os.path.join(
                        save_path, f"step_inv_{global_step}.safetensors"
                    ),
                    save_lora=False,
                )
                if log_wandb:
                    with torch.no_grad():
                        pipe = StableDiffusionPipeline(
                            vae=vae,
                            text_encoder=text_encoder,
                            tokenizer=tokenizer,
                            unet=unet,
                            scheduler=scheduler,
                            safety_checker=None,
                            feature_extractor=None,
                        )

                        # open all images in test_image_path
                        images = []
                        for file in os.listdir(test_image_path):
                            if (
                                file.lower().endswith(".png")
                                or file.lower().endswith(".jpg")
                                or file.lower().endswith(".jpeg")
                            ):
                                images.append(
                                    Image.open(os.path.join(test_image_path, file))
                                )

                        wandb.log({"loss": loss_sum / save_steps})
                        loss_sum = 0.0
                        wandb.log(
                            evaluate_pipe(
                                pipe,
                                target_images=images,
                                class_token=class_token,
                                learnt_token="".join(placeholder_tokens),
                                n_test=wandb_log_prompt_cnt,
                                n_step=50,
                                clip_model_sets=preped_clip,
                            )
                        )

            if global_step >= num_steps:
                return

def train(
    instance_data_dir: str,
    pretrained_model_name_or_path: str,
    output_dir: str,
    pretrained_vae_name_or_path: str = None,
    revision: Optional[str] = None,
    use_template: Literal[None, "object", "style"] = None,
    train_inpainting: bool = False,
    placeholder_tokens: str = "",
    preserve_prefix: bool = True,
    extend_amount: int = 1,
    do_coarse_inversion: bool = True,
    do_fine_inversion: bool = True,
    initializer_tokens: Optional[str] = None,
    seed: int = 42,
    resolution: int = 512,
    color_jitter: bool = True,
    train_batch_size: int = 1,
    sample_batch_size: int = 1,
    max_train_steps_ti: int = 3000,
    save_steps: int = 100,
    gradient_accumulation_steps: int = 2,
    gradient_checkpointing: bool = False,
    clip_ti_decay: bool = True,
    learning_rate_ti: float = 1e-3,
    cached_latents: bool = False,
    lr_scheduler: str = "linear",
    lr_warmup_steps: int = 0,
    weight_decay_ti: float = 0.00,
    use_8bit_adam: bool = False,
    device="cuda:0",
    extra_args: Optional[dict] = None,
    log_wandb: bool = False,
    wandb_log_prompt_cnt: int = 10,
    wandb_project_name: str = "new_pti_project",
    wandb_entity: str = "new_pti_entity",
    proxy_token: str = "person",
    enable_xformers_memory_efficient_attention: bool = False,
):
    torch.manual_seed(seed)

    if log_wandb:
        wandb.init(
            project=wandb_project_name,
            entity=wandb_entity,
            name=f"steps_{max_train_steps_ti}_lr_{learning_rate_ti}_{instance_data_dir.split('/')[-1]}",
            reinit=True,
            config={
                **(extra_args if extra_args is not None else {}),
            },
        )

    if output_dir is not None:
        os.makedirs(output_dir, exist_ok=True)
    # print(placeholder_tokens, initializer_tokens)
    if len(placeholder_tokens) == 0:
        placeholder_tokens = []
        print("PTI : Placeholder Tokens not given, using null token")
    else:
        placeholder_tokens = placeholder_tokens.split("|")

        assert (
            sorted(placeholder_tokens) == placeholder_tokens
        ), f"Placeholder tokens should be sorted. Use something like {'|'.join(sorted(placeholder_tokens))}'"

    if extend_amount > 0:
        placeholder_tokens = placeholder_tokens[:1] + (
            pp_extend(
                placeholder_tokens[0],
                placeholder_tokens[0],
                preserve_prefix=False,
                extend_amount=extend_amount,
            )
        )

    if initializer_tokens is None:
        print("PTI : Initializer Tokens not given, doing random inits")
        initializer_tokens = ["<rand-0.017>"] * len(placeholder_tokens)
    else:
        initializer_tokens = initializer_tokens.split("|")

    assert len(initializer_tokens) == len(
        placeholder_tokens
    ), "Unequal Initializer token for Placeholder tokens."

    if proxy_token is not None:
        class_token = proxy_token
    class_token = "".join(initializer_tokens)

    print("PTI : Placeholder Tokens", placeholder_tokens)
    print("PTI : Initializer Tokens", initializer_tokens)

    # get the models
    text_encoder, vae, unet, tokenizer, placeholder_token_ids = get_models(
        pretrained_model_name_or_path,
        pretrained_vae_name_or_path,
        revision,
        placeholder_tokens,
        initializer_tokens,
        device=device,
    )

    noise_scheduler = DDPMScheduler.from_config(
        pretrained_model_name_or_path, subfolder="scheduler"
    )

    if gradient_checkpointing:
        unet.enable_gradient_checkpointing()

    if enable_xformers_memory_efficient_attention:
        from diffusers.utils.import_utils import is_xformers_available

        if is_xformers_available():
            unet.enable_xformers_memory_efficient_attention()
        else:
            raise ValueError(
                "xformers is not available. Make sure it is installed correctly"
            )

    ti_lr = learning_rate_ti

    train_dataset = PPPDataset(
        instance_data_root=instance_data_dir,
        use_template=use_template,
        is_extended=False,
        main_token=placeholder_tokens[0],
        preserve_prefix=preserve_prefix,
        extend_amount=extend_amount,
        tokenizer=tokenizer,
        size=resolution,
        color_jitter=color_jitter,
        train_inpainting=train_inpainting,
    )

    train_dataset.blur_amount = 200

    train_dataloader = text2img_dataloader(
        train_dataset,
        train_batch_size,
        tokenizer,
        vae,
        text_encoder,
        cached_latents=cached_latents,
    )

    index_no_updates = torch.arange(len(tokenizer)) != -1

    for tok_id in placeholder_token_ids:
        index_no_updates[tok_id] = False

    unet.requires_grad_(False)
    vae.requires_grad_(False)

    params_to_freeze = itertools.chain(
        text_encoder.text_model.encoder.parameters(),
        text_encoder.text_model.final_layer_norm.parameters(),
        text_encoder.text_model.embeddings.position_embedding.parameters(),
    )
    for param in params_to_freeze:
        param.requires_grad = False

    if cached_latents:
        vae = None

    ti_optimizer = optim.AdamW(
        text_encoder.get_input_embeddings().parameters(),
        lr=ti_lr,
        betas=(0.9, 0.999),
        eps=1e-08,
        weight_decay=weight_decay_ti,
    )

    lr_scheduler = get_scheduler(
        lr_scheduler,
        optimizer=ti_optimizer,
        num_warmup_steps=lr_warmup_steps,
        num_training_steps=max_train_steps_ti,
    )

    # STEP 1 : Perform Course Inversion
    if do_coarse_inversion:

        train_dataset.is_extended = False

        train_inversion(
            unet,
            vae,
            text_encoder,
            train_dataloader,
            max_train_steps_ti,
            cached_latents=cached_latents,
            accum_iter=gradient_accumulation_steps,
            scheduler=noise_scheduler,
            index_no_updates=index_no_updates,
            optimizer=ti_optimizer,
            lr_scheduler=lr_scheduler,
            save_steps=save_steps,
            placeholder_tokens=placeholder_tokens,
            placeholder_token_ids=placeholder_token_ids,
            save_path=output_dir,
            test_image_path=instance_data_dir,
            log_wandb=log_wandb,
            wandb_log_prompt_cnt=wandb_log_prompt_cnt,
            class_token=class_token,
            train_inpainting=train_inpainting,
            mixed_precision=False,
            tokenizer=tokenizer,
            clip_ti_decay=clip_ti_decay,
            is_coarse_inversion=True,
        )

    # STEP 2 : Perform Fine Inversion

    if do_fine_inversion:
        train_dataset.is_extended = True
        ti_optimizer.param_groups[0]["lr"] = ti_lr * 0.33

        train_inversion(
            unet,
            vae,
            text_encoder,
            train_dataloader,
            max_train_steps_ti,
            cached_latents=cached_latents,
            accum_iter=gradient_accumulation_steps,
            scheduler=noise_scheduler,
            index_no_updates=index_no_updates,
            optimizer=ti_optimizer,
            lr_scheduler=lr_scheduler,
            save_steps=save_steps,
            placeholder_tokens=placeholder_tokens,
            placeholder_token_ids=placeholder_token_ids,
            save_path=output_dir,
            test_image_path=instance_data_dir,
            log_wandb=log_wandb,
            wandb_log_prompt_cnt=wandb_log_prompt_cnt,
            class_token=class_token,
            train_inpainting=train_inpainting,
            mixed_precision=False,
            tokenizer=tokenizer,
            clip_ti_decay=clip_ti_decay,
            is_coarse_inversion=False,
        )

def main():
    fire.Fire(train)


#!/usr/bin/env python3
import sys
from itertools import chain
from functools import partial
from collections import namedtuple

SPACE_JOINER = ", "
LINGER_SPEC_PREFIX = "# "

IBEX_TEMPLATE = """var shuffleSequence = seq("intro", sepWith("sep", seq("practice", rshuffle({conditions}))))
var practiceItemTypes = ["practice"];

var defaults = [
    "Separator", {{
        transfer: 1000,
        normalMessage: "Please wait for the next sentence.",
        errorMessage: "Wrong. Please wait for the next sentence."
    }},
    "DashedSentence", {{
        mode: "self-paced reading"
    }},
    "AcceptabilityJudgment", {{
        as: ["1", "2", "3", "4", "5", "6", "7"],
        presentAsScale: true,
        instructions: "Use number keys or click boxes to answer.",
        leftComment: "(Bad)", rightComment: "(Good)"
    }},
    "Question", {{
        hasCorrect: true
    }},
    "Message", {{
        hideProgressBar: true
    }},
    "Form", {{
        hideProgressBar: true,
        continueOnReturn: true,
        saveReactionTime: true
    }}
];

var items = [
  ["sep", "Separator", {{ }}],
  ["setcounter", "__SetCounter__", {{ }}],
  ["intro", "Form", {{
    html: {{ include: "example_intro.html" }},
    validators: {{
      age: function (s) {{ if (s.match(/^\d+$/)) return true;
                           else return "Bad value for \u2018age\u2019"; }}
      }}
  }} ],
    ["practice", "DashedSentence",
     {{s: "This is a practice sentence to get you used to reading sentences like this."}}],
    ["practice", "DashedSentence",
     {{s: "This is another practice sentence with a practice question following it."}},
                 "Separator", {{normalMessage: "Get ready for the question..."}},
                 "Question", {{hasCorrect: false, randomOrder: false,
                              q: "How would you like to answer this question?",
                              as: ["Press 1 or click here for this answer.",
                                   "Press 2 or click here for this answer.",
                                   "Press 3 or click here for this answer."]}}],
    ["practice", "DashedSentence",
     {{s: "This is the last practice sentence before the experiment begins."}}],

    {stims}
];
"""

SPEC_TEMPLATE = """["{condition}", {item}]"""

STIM_TEMPLATE = """
[{spec}, "{stimtype}", {{s: "{sentence}"}},
"Separator", {{normalMessage: "Get ready for the question..."}}, {questions}]
"""

QUESTION_TEMPLATE = """
"Question", {{q: "{question}", as: [{answers}]}}
"""

Stim = namedtuple("Stim", ['spec', 'sentence', 'questions'])
Spec = namedtuple("Spec", ['experiment', 'condition', 'item', 'rest'])
Question = namedtuple("Question", ['question', 'answers'])

def enquote(x):
    return "\"%s\"" % x

def consume(xs):
    for _ in xs:
        pass

def isplit(xs, sep):
    xs_it = iter(xs)
    while True:
        subit = iter(partial(next, xs_it), sep)
        try:
            probe = next(subit)
        except StopIteration:
            raise StopIteration
        yield chain([probe], subit)
        consume(subit)

flat = chain.from_iterable

def write_answers(answers):
    return SPACE_JOINER.join(map(enquote, answers))

def write_question(question):
    return QUESTION_TEMPLATE.format(
        question=question.question,
        answers=write_answers(question.answers)
    )

def write_questions(questions):
    return SPACE_JOINER.join(map(write_question, questions))

def write_conditions(spec):
    return "_".join([spec.experiment, spec.condition])

def write_spec(spec):
    # special case for fillers, which have no meaningful item number
    if spec.experiment.startswith("filler"):
        return enquote(write_conditions(spec))
    else:
        return SPEC_TEMPLATE.format(
            condition=write_conditions(spec),
            item=spec.item
        )

def write_stim(stimtype, stim):
    return STIM_TEMPLATE.format(
        spec=write_spec(stim.spec),
        stimtype=stimtype,
        sentence=stim.sentence,
        questions=write_questions(stim.questions)
    )

def parse_question(line):
    assert line.startswith("? ")
    line = line.strip("? ")
    question, answer = line.rsplit(" ", 1)
    if answer.lower() == "y" or answer.lower() == "yes":
        answers = ['Yes', 'No']
    elif answer.lower() == "n" or answer.lower() == "no":
        answers = ['No', 'Yes']
    else:
        raise ValueError("Unknown answer in question: %s" % line)
    return Question(question, answers)

def parse_stim(lines):
    spec_str, sentence, *questions = lines
    spec = parse_spec(spec_str)
    return Stim(spec, sentence, map(parse_question, questions))

def split_item(lines):
    lines = iter(lines)
    first_line = next(lines)
    assert first_line.startswith(LINGER_SPEC_PREFIX)
    stim = [first_line]
    for line in lines:
        if line.startswith(LINGER_SPEC_PREFIX):
            yield stim
            stim = []
        stim.append(line)
    yield stim

def parse_spec(line):
    assert line.startswith(LINGER_SPEC_PREFIX)
    line = line.strip(LINGER_SPEC_PREFIX)
    experiment, item, conditions, *rest = line.split()
    return Spec(experiment, conditions, int(item), rest)

def parse_linger(lines):
    lines = map(str.strip, lines)
    items = isplit(lines, "")
    stims = flat(map(split_item, items))
    return map(parse_stim, stims)

def write_ibex(stimtype, stims):
    stims = list(stims)
    write = partial(write_stim, stimtype)
    stims_string = ",".join(map(write, stims))
    conditions = {write_conditions(stim.spec) for stim in stims}
    conditions_str = ", ".join(map(enquote, conditions))
    return IBEX_TEMPLATE.format(
        conditions=conditions_str,
        stims=stims_string
    )

def run(stimtype, lines):
    stims = parse_linger(lines)
    return write_ibex(stimtype, stims)

def main(filename, stimtype="DashedSentence"):
    with open(filename) as infile:
        print(run(stimtype, infile))

if __name__ == '__main__':
    main(*sys.argv[1:])


from api.module.permission.permission_repository import PermissionRepository
from api.module.permission.dto.permission_create import PermissionCreateDto
from api.utils.response_message import ResponseMessage
from api.utils.uuid import UuidUtils

class PermissionService:
    permissionRepository = PermissionRepository()

    def list(self):
        return self.permissionRepository.list()

    def create(self, item: PermissionCreateDto):
        checkName = self.permissionRepository.count_by_field("name", item.name)
        checkTag = self.permissionRepository.count_by_field("tag", item.tag)

        if(checkName >= 1):
            raise ResponseMessage.ErrorHandleResponse(400, { "name": ResponseMessage.getAlreadyCreated() })

        if(checkTag >= 1):
            raise ResponseMessage.ErrorHandleResponse(400, { "tag": ResponseMessage.getAlreadyCreated() })

        return self.permissionRepository.save(item)

    def update(self, id: str, item: PermissionCreateDto):
        self.validateId(id)
        self.validateFields(item.name, item.tag, id)

        return self.permissionRepository.update(id, item)

    def delete(self, id):
        self.validateId(id)

        return self.permissionRepository.delete_by_id(id)

    def validateId(self, id):
        validUuid = UuidUtils.validUuid(id)

        if validUuid is False:
            raise ResponseMessage.ErrorHandleResponse(400, { "id": ResponseMessage.getInvalidId() })

        checkExists = self.permissionRepository.check_already_create_by_id(id)

        if checkExists <= 0:
            raise ResponseMessage.ErrorHandleResponse(404, { "id": ResponseMessage.getNotFound() })

    def validateFields(self, name: str, tag: str, id: str):
        checkName = self.permissionRepository.count_by_field("name", name, id)
        checkTag = self.permissionRepository.count_by_field("tag", tag, id)

        if(checkName >= 1):
            raise ResponseMessage.ErrorHandleResponse(400, { "name": ResponseMessage.getAlreadyCreated() })

        if(checkTag >= 1):
            raise ResponseMessage.ErrorHandleResponse(400, { "tag": ResponseMessage.getAlreadyCreated() })


from discord.ext import commands
import discord
from random import randint
import json
from modules.libraries.roll import roll_dice

class Monsters(commands.Cog):
    def __init__(self, client):
        self.client = client

    @commands.Cog.listener()
    async def on_ready(self):
        print("loaded cog")

    @commands.command()
    async def monster(self, ctx, *, arg):
        with open('./modules/libraries/monsters.json') as f:
            monster_data = json.load(f)
            for monster in monster_data:
                if monster["title"].lower() == arg.lower():
                    embed_monster = discord.Embed(title=monster["title"])

                    for content in monster["contents"]:
                        if content != "rule":
                            if "subtitle" in content:
                                embed_monster.description = content.split('|')[1]

                            # if "armor class" in content.lower() or "hit points" in content.lower() or "speed" in content.lower():
                            #     if "text" not in content.lower():
                            #         splitContent = content.split("|")

                            #         embed_monster.add_field(name=splitContent[1], value=splitContent[2])

                    embed_monster.set_thumbnail(url=monster["background_image"])

                    await ctx.send(embed=embed_monster)

def setup(client):
    client.add_cog(Monsters(client))


"""
Pipeline to handle model fitting from data extraction to saving results
"""

import json
import os
import pickle
import shutil
import uuid
from datetime import date, datetime
from glob import glob

import numpy as np
import pandas as pd

from . import changepoint_model
from . import changepoint_preprocess
from .utils import EphysData
#import changepoint_model
#import changepoint_preprocess
#from utils import EphysData

import pymc3
import theano

MODULE_DIR = os.path.dirname(__file__)
PARAMS_PATH = os.path.join(MODULE_DIR, 'config', 'MODEL_SAVE_DIR.params')
MODEL_SAVE_DIR = open(PARAMS_PATH, 'r').read().strip()
if not os.path.exists(MODEL_SAVE_DIR):
    os.makedirs(MODEL_SAVE_DIR)
    print('Created directory: {}'.format(MODEL_SAVE_DIR))
else:
    print('Using directory: {}'.format(MODEL_SAVE_DIR))
MODEL_DATABASE_PATH = os.path.join(MODEL_SAVE_DIR, 'model_database.csv')

class FitHandler():
    """Class to handle pipeline of model fitting including:
    1) Loading data
    2) Preprocessing loaded arrays
    3) Fitting model
    4) Writing out fitted parameters to pkl file

    """

    def __init__(self,
                 data_dir,
                 taste_num,
                 region_name,
                 laser_type = None,
                 experiment_name=None,
                 model_params_path=None,
                 preprocess_params_path=None,
                 ):
        """Initialize FitHandler class

        Args:
            data_dir (str): Path to directory containing HDF5 file
            taste_num (int): Index of taste to perform fit on
                    (Corresponds to INDEX of taste in spike array, not actual dig_ins)
            region_name (str): Region on which to perform fit on
                    (must match regions in .info file)
            experiment_name (str, optional): Name given to fitted batch
                    (for metedata). Defaults to None.
            model_params_path (str, optional): Path to json file
                    containing model parameters. Defaults to None.
            preprocess_params_path (str, optional): Path to json file
                    containing preprocessing parameters. Defaults to None.

        Raises:
            Exception: If "experiment_name" is None
            Exception: If "laser_type" is not in [None, 'on', 'off']
            Exception: If "taste_num" is not integer or "all"
        """

        # =============== Check for exceptions ===============
        if experiment_name is None:
            raise Exception('Please specify an experiment name')
        if laser_type not in [None, 'on','off']:
            raise Exception('laser_type must be from [None, "on","off"]')
        if not (isinstance(taste_num,int) or taste_num == 'all'):
            raise Exception('taste_num must be an integer or "all"')

        # =============== Save relevant arguments ===============
        self.data_dir = data_dir
        self.EphysData = EphysData(self.data_dir)
        #self.data = self.EphysData.get_spikes({"bla","gc","all"})

        self.taste_num = taste_num
        self.laser_type = laser_type
        self.region_name = region_name
        self.experiment_name = experiment_name

        data_handler_init_kwargs = dict(zip(
            [   'data_dir',
                'experiment_name',
                'taste_num',
                'laser_type',
                'region_name'],
            [   data_dir,
                experiment_name,
                taste_num,
                laser_type,
                region_name]))
        self.database_handler = DatabaseHandler()
        self.database_handler.set_run_params(**data_handler_init_kwargs)

        if model_params_path is None:
            print('MODEL_PARAMS will have to be set')
        else:
            self.set_model_params(file_path=model_params_path)

        if preprocess_params_path is None:
            print('PREPROCESS_PARAMS will have to be set')
        else:
            self.set_preprocess_params(file_path=preprocess_params_path)

    ########################################
    # SET PARAMS
    ########################################

    def set_preprocess_params(self,
                              time_lims,
                              bin_width,
                              data_transform,
                              file_path=None):

        """Load given params as "preprocess_params" attribute

        Args:
            time_lims (array/tuple/list): Start and end of where to cut
                    spike train array
            bin_width (int): Bin width for binning spikes to counts
            data_transform (str): Indicator for which transformation to
                    use (refer to changepoint_preprocess)
            file_path (str, optional): Path to json file containing preprocess
                    parameters. Defaults to None.
        """

        if file_path is None:
            preprocess_params_dict = dict(zip(['time_lims', 'bin_width', 'data_transform'],
                         [time_lims, bin_width, data_transform]))
            self.preprocess_params = preprocess_params_dict
            print('Set preprocess params to: {}'.format(preprocess_params_dict))

        else:
            # Load json and save dict
            pass

    def set_model_params(self,
                         states,
                         fit,
                         samples,
                         model_kwargs = None,
                         file_path=None):

        """Load given params as "model_params" attribute

        Args:
            states (int): Number of states to use in model
            fit (int): Iterations to use for model fitting (given ADVI fit)
            samples (int): Number of samples to return from fitten model
            model_kwargs (dict) : Additional paramters for model
            file_path (str, optional): Path to json file containing
                    preprocess parameters. Defaults to None.
        """

        if file_path is None:
            model_params_dict = dict(zip(['states', 'fit', 'samples', 'model_kwargs'],
                    [states, fit, samples, model_kwargs]))
            self.model_params = model_params_dict
            print('Set model params to: {}'.format(model_params_dict))

        else:
            # Load json and save dict
            pass

    ########################################
    # SET PIPELINE FUNCS
    ########################################

    def set_preprocessor(self, preprocessing_func):
        """Manually set preprocessor for data e.g.

        FitHandler.set_preprocessor(
                    changepoint_preprocess.preprocess_single_taste)

        Args:
            preprocessing_func (func):
                    Function to preprocess data (refer to changepoint_preprocess)
        """
        self.preprocessor = preprocessing_func

    def preprocess_selector(self):
        """Function to return preprocess function based off of input flag

        Preprocessing can be set manually but it is preferred to
        go through preprocess selector

        Raises:
            Exception: If self.taste_num is neither int nor str

        """

        if isinstance(self.taste_num, int):
            self.set_preprocessor(
                changepoint_preprocess.preprocess_single_taste)
        elif self.taste_num == 'all':
            self.set_preprocessor(changepoint_preprocess.preprocess_all_taste)
        else:
            raise Exception("Something went wrong")

    def set_model_template(self, model_template):
        """Manually set model_template for data e.g.

        FitHandler.set_model(changepoint_model.single_taste_poisson)

        Args:
            model_template (func): Function to generate model template for data]
        """
        self.model_template = model_template

    def model_template_selector(self):
        """Function to set model based off of input flag

        Models can be set manually but it is preferred to go through model selector

        Raises:
            Exception: If self.taste_num is neither int nor str

        """
        if isinstance(self.taste_num, int):
            #self.set_model_template(changepoint_model.single_taste_poisson_varsig)
            self.set_model_template(
                    changepoint_model.single_taste_poisson)
        elif self.taste_num == 'all':
            self.set_model_template(changepoint_model.all_taste_poisson)
        else:
            raise Exception("Something went wrong")

    def set_inference(self, inference_func):
        """Manually set inference function for model fit e.g.

        FitHandler.set_inference(changepoint_model.advi_fit)

        Args:
            inference_func (func): Function to use for fitting model
        """
        self.inference_func = changepoint_model.advi_fit

    def inference_func_selector(self):
        """Function to return model based off of input flag

        Currently hard-coded to use "advi_fit"
        """
        self.set_inference(changepoint_model.advi_fit)

    ########################################
    # PIPELINE FUNCS
    ########################################

    def load_spike_trains(self):
        """Helper function to load spike trains from data_dir using EphysData module
        """
        full_spike_array = self.EphysData.return_region_spikes(
                region_name = self.region_name,
                laser = self.laser_type)
        if isinstance(self.taste_num, int):
            self.data = full_spike_array[self.taste_num]
        if self.taste_num == 'all':
            self.data = full_spike_array
        print(f'Loading spike trains from {self.database_handler.data_basename}, '
              f'dig_in {self.taste_num}, laser {str(self.laser_type)}')

    def preprocess_data(self):
        """Perform data preprocessing

        Will check for and complete:
        1) Raw data loaded
        2) Preprocessor selected
        """
        if 'data' not in dir(self):
            self.load_spike_trains()
        if 'preprocessor' not in dir(self):
            self.preprocess_selector()
        print('Preprocessing spike trains, '
              f'preprocessing func: <{self.preprocessor.__name__}>')
        self.preprocessed_data = \
            self.preprocessor(self.data, **self.preprocess_params)

    def create_model(self):
        """Create model and save as attribute

        Will check for and complete:
        1) Data preprocessed
        2) Model template selected
        """
        if 'preprocessed_data' not in dir(self):
            self.preprocess_data()
        if 'model_template' not in dir(self):
            self.model_template_selector()

        # In future iterations, before fitting model,
        # check that a similar entry doesn't exist

        changepoint_model.compile_wait()
        print(f'Generating Model, model func: <{self.model_template.__name__}>')
        self.model = self.model_template(self.preprocessed_data,
                                         self.model_params['states'],
                                         **self.model_params['model_kwargs'])

    def run_inference(self):
        """Perform inference on data

        Will check for and complete:
        1) Model created
        2) Inference function selected
        """
        if 'model' not in dir(self):
            self.create_model()
        if 'inference_func' not in dir(self):
            self.inference_func_selector()

        changepoint_model.compile_wait()
        print('Running inference, inference func: '
              f'<{self.inference_func.__name__}>')
        temp_outs = self.inference_func(self.model,
                                        self.model_params['fit'],
                                        self.model_params['samples'])
        varnames = ['model', 'approx', 'lambda', 'tau', 'data']
        self.inference_outs = dict(zip(varnames, temp_outs))

    def _gen_fit_metadata(self):
        """Generate metadata for fit

        Generate metadat by compiling:
        1) Preprocess parameters given as input
        2) Model parameters given as input
        3) Functions used in inference pipeline for : preprocessing,
                model generation, fitting

        Returns:
            dict: Dictionary containing compiled metadata for different
                    parts of inference pipeline
        """
        pre_params = self.preprocess_params
        model_params = self.model_params
        pre_params['preprocessor_name'] = self.preprocessor.__name__
        model_params['model_template_name'] = self.model_template.__name__
        model_params['inference_func_name'] = self.inference_func.__name__
        fin_dict = dict(zip(['preprocess', 'model'], [pre_params, model_params]))
        return fin_dict

    def _pass_metadata_to_handler(self):
        """Function to coordinate transfer of metadata to DatabaseHandler
        """
        self.database_handler.ingest_fit_data(self._gen_fit_metadata())

    def _return_fit_output(self):
        """Compile data, model, fit, and metadata to save output

        Returns:
            dict: Dictionary containing fitted model data and metadata
        """
        self._pass_metadata_to_handler()
        agg_metadata = self.database_handler.aggregate_metadata()
        return {'model_data': self.inference_outs, 'metadata': agg_metadata}

    def save_fit_output(self):
        """Save fit output (fitted data + metadata) to pkl file
        """
        if 'inference_outs' not in dir(self):
            self.run_inference()
        out_dict = self._return_fit_output()
        with open(self.database_handler.model_save_path + '.pkl', 'wb') as buff:
            pickle.dump(out_dict, buff)

        json_file_name = os.path.join(
            self.database_handler.model_save_path + '.info')
        with open(json_file_name, 'w') as file:
            json.dump(out_dict['metadata'], file, indent=4)

        self.database_handler.write_to_database()

        print('Saving inference output to : \n'
                f'{self.database_handler.model_save_dir}'
                "\n" + "================================" + '\n')

class DatabaseHandler():
    """Class to handle transactions with model database
    """

    def __init__(self):
        """Initialize DatabaseHandler class
        """
        self.unique_cols = ['exp.model_id', 'exp.save_path', 'exp.fit_date']
        self.model_database_path = MODEL_DATABASE_PATH
        self.model_save_base_dir = MODEL_SAVE_DIR

        if os.path.exists(self.model_database_path):
            self.fit_database = pd.read_csv(self.model_database_path,
                                            index_col=0)
            all_na = [all(x) for num, x in self.fit_database.isna().iterrows()]
            if all_na:
                print(f'{sum(all_na)} rows found with all NA, removing...')
                self.fit_database = self.fit_database.dropna(how='all')
        else:
            print('Fit database does not exist yet')

    def show_duplicates(self, keep='first'):
        """Find duplicates in database

        Args:
            keep (str, optional): Which duplicate to keep
                    (refer to pandas duplicated). Defaults to 'first'.

        Returns:
            pandas dataframe: Dataframe containing duplicated rows
            pandas series : Indices of duplicated rows
        """
        dup_inds = self.fit_database.drop(self.unique_cols, axis=1)\
            .duplicated(keep=keep)
        return self.fit_database.loc[dup_inds], dup_inds

    def drop_duplicates(self):
        """Remove duplicated rows from database
        """
        _, dup_inds = self.show_duplicates()
        print(f'Removing {sum(dup_inds)} duplicate rows')
        self.fit_database = self.fit_database.loc[~dup_inds]

    def check_mismatched_paths(self):
        """Check if there are any mismatched pkl files between database and directory

        Returns:
            pandas dataframe: Dataframe containing rows for which pkl file not present
            list: pkl files which cannot be matched to model in database
            list: all files in save directory
        """
        mismatch_from_database = [not os.path.exists(x + ".pkl")
                                  for x in self.fit_database['exp.save_path']]
        file_list = glob(os.path.join(self.model_save_base_dir, "*/*.pkl"))
        mismatch_from_file = [not
                              (x.split('.')[0] in list(
                                  self.fit_database['exp.save_path']))
                              for x in file_list]
        print(f"{sum(mismatch_from_database)} mismatches from database" + "\n"
              + f"{sum(mismatch_from_file)} mismatches from files")
        return mismatch_from_database, mismatch_from_file, file_list

    def clear_mismatched_paths(self):
        """Remove mismatched files and rows in database

        i.e. Remove
        1) Files for which no entry can be found in database
        2) Database entries for which no corresponding file can be found
        """
        mismatch_from_database, mismatch_from_file, file_list = \
            self.check_mismatched_paths()
        mismatch_from_file = np.array(mismatch_from_file)
        mismatch_from_database = np.array(mismatch_from_database)
        self.fit_database = self.fit_database.loc[~mismatch_from_database]
        mismatched_files = [x for x, y in zip(file_list, mismatch_from_file) if y]
        for x in mismatched_files:
            os.remove(x)
        print('==== Clearing Completed ====')

    def write_updated_database(self):
        """Can be called following clear_mismatched_entries to update current database
        """
        database_backup_dir = os.path.join(
            self.model_save_base_dir, '.database_backups')
        if not os.path.exists(database_backup_dir):
            os.makedirs(database_backup_dir)
        #current_date = date.today().strftime("%m-%d-%y")
        current_date = str(datetime.now()).replace(" ", "_")
        shutil.copy(self.model_database_path,
                    os.path.join(database_backup_dir,
                                 f"database_backup_{current_date}"))
        self.fit_database.to_csv(self.model_database_path, mode='w')

    def set_run_params(
            self,
            data_dir,
            experiment_name,
            taste_num,
            laser_type,
            region_name):
        """Store metadata related to inference run

        Args:
            data_dir (str): Path to directory containing HDF5 file
            experiment_name (str): Name given to fitted batch
                    (for metedata). Defaults to None.
            taste_num (int): Index of taste to perform fit on (Corresponds to
                    INDEX of taste in spike array, not actual dig_ins)
            laser_type (None or str): None, 'on', or 'off' (For a laser session,
                    which set of trials are wanted, None indicated return all trials)
            region_name (str): Region on which to perform fit on
                    (must match regions in .info file)
        """
        self.data_dir = data_dir
        self.data_basename = os.path.basename(self.data_dir)
        self.animal_name = self.data_basename.split("_")[0]
        self.session_date = self.data_basename.split("_")[-1]

        self.experiment_name = experiment_name
        self.model_save_dir = os.path.join(self.model_save_base_dir,
                                           experiment_name)

        if not os.path.exists(self.model_save_dir):
            os.makedirs(self.model_save_dir)

        self.model_id = str(uuid.uuid4()).split('-')[0]
        self.model_save_path = os.path.join(self.model_save_dir,
                                            self.experiment_name +
                                            "_" + self.model_id)
        self.fit_date = date.today().strftime("%m-%d-%y")

        self.taste_num = taste_num
        self.laser_type = laser_type
        self.region_name = region_name

        self.fit_exists = None

    def ingest_fit_data(self, met_dict):
        """Load external metadata

        Args:
            met_dict (dict): Dictionary of metadata from FitHandler class
        """
        self.external_metadata = met_dict

    def aggregate_metadata(self):
        """Collects information regarding data and current "experiment"

        Raises:
            Exception: If 'external_metadata' has not been ingested, that needs to be done first

        Returns:
            dict: Dictionary of metadata given to FitHandler class
        """
        if 'external_metadata' not in dir(self):
            raise Exception('Fit run metdata needs to be ingested '
                            'into data_handler first')

        data_details = dict(zip(
            ['data_dir',
             'basename',
             'animal_name',
             'session_date',
             'taste_num',
             'laser_type',
             'region_name'],
            [self.data_dir,
             self.data_basename,
             self.animal_name,
             self.session_date,
             self.taste_num,
             self.laser_type,
             self.region_name]))

        exp_details = dict(zip(
            ['exp_name',
             'model_id',
             'save_path',
             'fit_date'],
            [self.experiment_name,
             self.model_id,
             self.model_save_path,
             self.fit_date]))

        module_details = dict(zip(
            ['pymc3_version',
             'theano_version'],
            [pymc3.__version__,
                theano.__version__]
            ))

        temp_ext_met = self.external_metadata
        temp_ext_met['data'] = data_details
        temp_ext_met['exp'] = exp_details
        temp_ext_met['module'] = module_details

        return temp_ext_met

    def write_to_database(self):
        """Write out metadata to database
        """
        agg_metadata = self.aggregate_metadata()
        # Convert model_kwargs to str so that they are save appropriately
        agg_metadata['model']['model_kwargs'] = str(agg_metadata['model']['model_kwargs'])
        flat_metadata = pd.json_normalize(agg_metadata)
        if not os.path.isfile(self.model_database_path):
            flat_metadata.to_csv(self.model_database_path, mode='a')
        else:
            flat_metadata.to_csv(self.model_database_path,
                                 mode='a', header=False)

    def check_exists(self):
        """Check if the given fit already exists in database

        Returns:
            bool: Boolean for whether fit already exists or not
        """
        if self.fit_exists is not None:
            return self.fit_exists


import os
import requests
import urllib3
import click
import tqdm
import multiprocessing

# using the 1 arc-second dataset, but could use 1/3 arc-second for higher resolution
dataset = "National Elevation Dataset (NED) 1 arc-second"

# use urllib3 to download a url
def request(method, url):
    try:
        r = urllib3.request(method, url)
        return r
    except Exception as e:
        print("Failed to download", url, e)
        return None

# download a url to data/sources/
def download(url):
    filename = url.split("/")[-1]
    try:
        filepath = os.path.join("data/sources", filename)
        if not os.path.exists(filepath):
            data = request("GET", url)
            # save the data to filepath
            with open(filepath, "wb") as f:
                f.write(data.data)
    except Exception as e:
        print("Failed to save file", filename, e)

# click cli to get number of workers and bounding box
@click.command()
@click.option("--workers", default=multiprocessing.cpu_count(), help="Number of workers to use")
@click.option(
    "--bbox",
    default="-124.566244,46.864746,-116.463504,41.991794",
    help="Bounding box to download data for",
)
def cli(workers, bbox):
    print("Getting region details from the National Map API...")
    # USGS national map api endpoint to list DEMs for our bounding box
    url = (
        "https://tnmaccess.nationalmap.gov/api/v1/products?datasets={}&bbox={}".format(
            dataset, bbox
        )
    )

    response = requests.get(url)
    data = response.json()

    urls = []
    download_bytes = 0

    total_items = data["total"]
    n = len(data["items"])
    for item in data["items"]:
        urls.append(item["downloadURL"])
        download_bytes += item["sizeInBytes"]

    print("Getting list of tile URLs in bounding box region...")
    with tqdm.tqdm(total=total_items) as pbar:
        pbar.update(n)
        while n < total_items:
            url += "&offset=" + str(n)
            response = requests.get(url + "&offset={}".format(n))
            data = response.json()
            for item in data["items"]:
                urls.append(item["downloadURL"])
                download_bytes += item["sizeInBytes"]

            n += len(data["items"])
            pbar.update(len(data["items"]))

    print("Approximate total download size: {:.2f} GB".format(download_bytes / 1e9))

    # create data/sources/ dir if it doesn't already exist
    if not os.path.exists("data/sources"):
        os.makedirs("data/sources")

    # filter out any urls that already exist as saved files
    filtered_urls = [
        url
        for url in urls
        if not os.path.exists(os.path.join("data/sources", url.split("/")[-1]))
    ]

    print(
        "Skipping downloading {} files that already exist".format(
            len(urls) - len(filtered_urls)
        )
    )

    print("Downloading {} files...".format(len(filtered_urls)))

    # download the data in parallel save each file to data/sources/
    with multiprocessing.Pool(processes=workers) as pool:
        for _ in tqdm.tqdm(
            pool.imap_unordered(download, filtered_urls), total=len(filtered_urls)
        ):
            pass
        pool.close()
        pool.join()

if __name__ == "__main__":
    cli()


'''
Hacker's rank
Given a set of distinct integers, print the size of a maximal subset of  where the sum of any  numbers in  is not evenly divisible by .

Example

One of the arrays that can be created is . Another is . After testing all permutations, the maximum length solution array has  elements.

Function Description

Complete the nonDivisibleSubset function in the editor below.

nonDivisibleSubset has the following parameter(s):

int S[n]: an array of integers
int k: the divisor
Returns

int: the length of the longest subset of  meeting the criteria
Input Format

The first line contains  space-separated integers,  and , the number of values in  and the non factor.
The second line contains  space-separated integers, each an , the unique values of the set.

Constraints

All of the given numbers are distinct.
Sample Input

STDIN    Function
-----    --------
4 3      S[] size n = 4, k = 3
1 7 2 4  S = [1, 7, 2, 4]
Sample Output

3
Explanation

The sums of all permutations of two elements from  are:

1 + 7 = 8
1 + 2 = 3
1 + 4 = 5
7 + 2 = 9
7 + 4 = 11
2 + 4 = 6
Only  will not ever sum to a multiple of
'''
import math
import os
import random
import re
import sys
#
# Complete the 'nonDivisibleSubset' function below.
#
# The function is expected to return an INTEGER.
# The function accepts following parameters:
#  1. INTEGER k
#  2. INTEGER_ARRAY s
#

# timeout
def nonDivisibleSubset1(k, s):
    #sequence = []

    while True:
        count = [[i, 0] for i in range(len(s))]

        found = False
        for i in range(len(s) - 1):
            for j in range(i + 1, len(s)):
                if (s[i] + s[j]) % k == 0:
                    count[i][1] += 1
                    count[j][1] += 1
                    found = True
                    #sequence.append((i, j))

        if found:
            max_val = 0
            max_index = 0
            for i in range(len(s)):
                if count[i][1] > max_val:
                    max_index = i
                    max_val = count[i][1]
            s.pop(max_index)
        else:
            break

    return len(s)

# timeout
def nonDivisibleSubset2(k, s):
    count = [[i, []] for i in range(len(s))]

    for i in range(len(s) - 1):
        for j in range(i + 1, len(s)):
            if (s[i] + s[j]) % k == 0:
                count[i][1].append(j)
                count[j][1].append(i)

    #print(count)
    total_cnt = len(s)
    while True:
        max_value = 0
        max_index = 0

        for i in range(len(count)):
            if len(count[i][1]) > max_value:
                max_index = i
                max_value = len(count[i][1])

        #print(max_value)
        if max_value == 0:
            break

        for j in count[max_index][1]:
            count[j][1].remove(max_index)

        count[max_index][1] = []
        total_cnt -= 1

    return total_cnt

# 나머지가 1인 것은 (k-1)인 것을 더하면, 나머지가 0이됨.
# 둘 중 개수가 적은 것을 버리는 것이 유리.
def nonDivisibleSubset(k, s):
    count = [0] * k

    for a in s:
        r = a % k
        count[r] += 1

    num = len(s)

    if count[0] > 1:
        num = num - (count[0] - 1)

    if k % 2 == 0 and count[k // 2] > 0:
        num = num - (count[k // 2] - 1)
        last = k // 2 - 1
    else:
        last = k // 2

    for i in range(1, last + 1):
        if count[i] >= count[k - i]:
            num = num - (count[k - i])
        else:
            num = num - count[i]

    return num

if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')

    first_multiple_input = input().rstrip().split()

    n = int(first_multiple_input[0])

    k = int(first_multiple_input[1])

    s = list(map(int, input().rstrip().split()))

    result = nonDivisibleSubset2(k, s)

    fptr.write(str(result) + '\n')

    fptr.close()


test_cases= int(input())
count=0
sum=0
for i in range(test_cases):
    green_price,purple_price = map(int,input().split())
    Total_student= int(input())
    for _ in range(Total_student):
        question1,question2 = map(int,input().split())
        if question1 == 1:
            count+= 1
        if question2 == 1:
            sum+= 1
    if green_price>purple_price:
        if count>sum:
            print(green_price*sum + purple_price*count)
        else:
            print(green_price*count + purple_price*sum)
    else:
        if count>sum:
            print(green_price*count + purple_price*sum)
        else:
            print(green_price*sum + purple_price*count)
    count=0
    sum=0


import pandas as pd
from pathlib import Path
import os
from wordcloud import WordCloud
import matplotlib.pyplot as plt

# This script uses wordcloud
# get it via: conda install -c conda-forge wordcloud
# The general setup is that you first run this script from within RM
# Afterwards you can run this with your IDE (e.g pyCharm). This is able to pick up the dumps

#
# These folders are pointing to a tmp folder for dumps and results
#
resultfolder = os.path.join(str(Path.home()), ".RapidMiner", "python_results")
dumppath = os.path.join(resultfolder, "data.pickle")

#
# Main function called from rapidminer.
#
# data: a pandas dataframe with 3 cols: topic, word, weight
# singleOutput: If True, we generate one big output file. If true we write #topics individual files.
# calledFromRM: In RapidMiner calledFromRM is always true. Just used in the first lines.
#
# writes: Files to the resultfolder path.
# returns: data to work again on it in RM.
#
def rm_main(data, singleOutput=False, calledFromRM=True):
    name = "Topic"
    if (calledFromRM is True):
        # do rm specific things
        if os.path.exists(resultfolder) is False:
            os.makedirs(resultfolder)
        data.to_pickle(dumppath)

    # convert dataframe into something wordcloud lib can use
    d = {}
    topics = []
    for topic, a, x in data.values:
        d[topic] = {}
        if topic not in topics:
            topics.append(topic)

    for topic, a, x in data.values:
        d[topic][a] = x

    row, col = __getRowAndColSetting__(len(topics))
    plt.subplots_adjust(hspace=0.1)
    # loop over the topics and create word clouds
    # if singleOutput is True, create one big clouds picture
    # otherwise create individual charts.
    for topic in topics:
        wordcloud = WordCloud(background_color="white")
        wordcloud.generate_from_frequencies(frequencies=d[topic])
        if singleOutput is True:
            plt.subplot(row, col, topic + 1)
        plt.imshow(wordcloud, interpolation='bilinear')
        plt.title("Topic_" + str(topic))
        plt.axis("off")
        if singleOutput is False:
            plt.savefig(resultfolder + "/" + name + "_" + str(topic) + ".png")

    if singleOutput is True:
        plt.axis("off")
        plt.savefig(resultfolder + "/" + name + "_Overview_" + str(len(topics)) + ".png", quality=100)

    return data

# short helper function to get the number of rows and cols for a subplot
# we always take 4 rows and then as many as needed cols.
def __getRowAndColSetting__(nTopics):
    row = 3
    if nTopics % 3 == 0:
        col = nTopics / 3
    else:
        col = nTopics / 3 + 1
    return row, col

if __name__ == "__main__":
    print("Starting Word Cloud script in batch mode")
    print("Reading last execution data from: ", dumppath)
    data = pd.read_pickle(dumppath)
    rm_main(data, calledFromRM=False, singleOutput=True)


from django.conf.urls import url
from django.urls import path, include
from django.utils.translation import ugettext as _
from apps.dashboard.views import reorder_dashboard, dashboard_view, overview_dashboard, add_dashboard, edit_dashboard, toggle_dashboard_activation_view, delete_dashboard, delete_ajax_dashboard_modal, \
                                 overview_reversion, revert_dashboard
urlpatterns = [
    #dashboard urls
    url(r'^$', dashboard_view, name='dashboard'),
    #apps urls
    url(_('^overview$').strip(),overview_dashboard,name="overviewdashboard"),
    url(_('^add$').strip(),add_dashboard,name="adddashboard"),
    url(_('^edit/(?P<pk>[0-9a-f-]+)$').strip(),edit_dashboard,name="editdashboard"),
    url(_('^toggle-activation/(?P<pk>[0-9a-f-]+)$').strip(), toggle_dashboard_activation_view, name="activate-dashboard"),
    url(_('^delete/(?P<pk>[0-9a-f-]+)$').strip(),delete_dashboard,name="deletedashboard"),
    url(_('^delete/modal$').strip(),delete_ajax_dashboard_modal,name="deletemodaldashboard"),
    url(_('^dashboard/reorder').strip(), reorder_dashboard, name="dashboard-reorder"),

    url(_('^overview/reversion/$').strip(),overview_reversion,name="overviewreversiondashboard"),
    url(_('^revert/(?P<pk>[0-9a-f-]+)$').strip(),revert_dashboard,name="revertdashboard"),
]


from datetime import datetime

import logging
from flask import Flask
from flask import jsonify
from flask import request
from flask_cors import CORS, cross_origin

from smart_client import poaching_data, human_activity_data
from metadata import api_schema

app = Flask(__name__)
CORS(app)
logging.basicConfig()
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

@app.route('/poaching', methods=['GET'])
def poaching_endpoint(event=None, context=None):
    logger.info('Lambda function invoked /poaching')
    limit = int(request.args.get('limit', 1000))
    raw_categories = request.args.get('categories')
    start_date = _cleanup_time(request.args.get('start_date'))
    end_date = _cleanup_time(request.args.get('end_date'))

    if raw_categories:
        categories = raw_categories.split(',')
    else:
        categories = []
    return jsonify(poaching_data(max_rows=limit, categories=categories, start_date=start_date, end_date=end_date))

@app.route('/human_activity', methods=['GET'])
def human_activity_endpoint(event=None, context=None):
    logger.info('Lambda function invoked /human_activity')
    limit = int(request.args.get('limit', 1000))
    start_date = _cleanup_time(request.args.get('start_date'))
    end_date = _cleanup_time(request.args.get('end_date'))

    return jsonify(human_activity_data(max_rows=limit, start_date=start_date, end_date=end_date))

@app.route('/metadata', methods=['GET'])
def metadata(event=None, context=None):
    return jsonify(api_schema)

@app.route('/', methods=['GET'])
def index(event=None, context=None):
    logger.info('Lambda function invoked /')
    return jsonify({'ok': True})

def _cleanup_time(time_string):
    if not time_string:
        return None
    return datetime.strptime(time_string, "%Y-%m-%dT%H:%M:%S")

if __name__ == '__main__':
    app.run(debug=True)


import requests
from abstract_test import  *

class Test_Admin_Page(FsConnectionAdminWeb):
    def test_Admin_Page(self):
        self.connection = self.instance()
        self.connection.loginRequest("arao_admin_test","arao_admin")
        r = self.connection.getRequest("","user")
        assert r is not None and r.status_code == 200
        assert "Admin Dashboard" in r.text
        assert "arao_admin_test" in r.text
        assert "Total Confirmed cases" in r.text


import os
from datetime import datetime
import pandas as pd
from tm.dataset.dataset_config import SPLIT_NAME_TRAIN, SPLIT_NAME_VALIDATION
from tm.dataset.dataset_split import DS_SPLIT_METHOD_VAL_2023, DS_SPLIT_METHOD_VAL_3M9M
from tm.dataset.dataset_config import intervals_meta

def rolling_to_interval(base_dir,
                        symbol,
                        from_interval,
                        from_df,
                        rolling_to,
                        ds_split_name=None,
                        ds_split_method_name=None,
                        ):
    to_interval, times = rolling_to
    print('>', to_interval, ds_split_name)

    df2 = from_df.rolling(times, min_periods=times).agg({'time': lambda s: s.iloc[-1],
                                                         'open_time': lambda s: s.iloc[0],
                                                         'open_hour': lambda s: s.iloc[0],
                                                         'open': lambda s: s.iloc[0],
                                                         'high': 'max',
                                                         'low': 'min',
                                                         'close': lambda s: s.iloc[-1],
                                                         'volume': 'sum',
                                                         'amount': 'sum',
                                                         'buy_vol': 'sum',
                                                         'buy_amt': 'sum',
                                                         'trades': 'sum',
                                                         })
    df2 = df2.iloc[(times - 1):]
    i_span = intervals_meta[to_interval]['seconds']
    spans = df2['time'] - df2['open_time']
    span_neq = spans == i_span
    df2 = df2[span_neq]
    print('wi', (~span_neq).sum())

    def to_time_str(d):
        ts = round(d)
        dt = datetime.utcfromtimestamp(ts)
        return dt.isoformat(timespec='seconds')

    df2['time'] = df2['time'].apply(to_time_str)
    df2['open_time'] = df2['open_time'].apply(to_time_str)

    print(df2.head())

    base_filename = f'{symbol}-{to_interval}-rolling-{from_interval}'
    if ds_split_method_name is not None:
        split_target_dir = f'{base_dir}/{to_interval}-rolling-{from_interval}-{ds_split_method_name}-{ds_split_name}/{symbol}'
    else:
        split_target_dir = f'{base_dir}/{to_interval}-rolling-{from_interval}/{symbol}'

    os.makedirs(split_target_dir, exist_ok=True)

    # rows = df2.shape[0]
    for mod in range(times):
        # split_df = df2.iloc[range(mod, rows, times)]
        split_df = df2[df2['open_hour'] % times == mod]
        split_df = split_df.drop(['open_hour'], axis=1)
        split_df.to_csv(f'{split_target_dir}/{base_filename}-m{mod}.csv',
                        index=False,
                        )

    target_dir = f'{base_dir}/{to_interval}-all/{symbol}'
    if ds_split_method_name is not None:
        target_filename = f'{base_filename}-{ds_split_method_name}-{ds_split_name}.csv'
    else:
        target_filename = f'{base_filename}-all.csv'
    os.makedirs(target_dir, exist_ok=True)
    df2 = df2.drop(['open_hour'], axis=1)
    df2.to_csv(f'{target_dir}/{target_filename}',
               index=False,
               )

def build_rolling_kl(base_dir,
                     symbol,
                     from_interval,
                     to_intervals,
                     ds_split_name=None,
                     ds_split_method_name=None,
                     ):
    from_freq = from_interval.upper()

    if ds_split_name is None:
        src_filename = f'{symbol}-{from_interval}-all.csv'
    else:
        src_filename = f'{symbol}-{from_interval}-{ds_split_method_name}-{ds_split_name}.csv'
    df = pd.read_csv(f'{base_dir}/{from_interval}-all/{symbol}/{src_filename}',
                     # nrows=200,
                     index_col='time',
                     )
    print(df.head())
    print('total', df.shape[0], 'rows')
    print()

    open_time_period = pd.PeriodIndex(data=df['open_time'], freq=from_freq)
    df['open_hour'] = open_time_period.hour
    df['open_time'] = open_time_period.to_timestamp().map(pd.Timestamp.timestamp)
    df.index = pd.PeriodIndex(data=df.index, freq=from_freq)
    df = df[~df.index.duplicated()]
    df = df.ffill()
    time2 = df.index.to_timestamp().map(pd.Timestamp.timestamp)
    df.reset_index(inplace=True)
    df['time'] = time2

    for rolling_to in to_intervals:
        rolling_to_interval(base_dir,
                            symbol,
                            from_interval,
                            df,
                            rolling_to,
                            ds_split_name=ds_split_name,
                            ds_split_method_name=ds_split_method_name,
                            )
        print()

if __name__ == '__main__':
    bd = '../../data/tm/market/kline-basic'
    split_method = DS_SPLIT_METHOD_VAL_2023
    split_names = [SPLIT_NAME_TRAIN, SPLIT_NAME_VALIDATION]
    for sym in ['BTCUSDT', 'ETHUSDT']:
        for split_name in split_names:
            build_rolling_kl(base_dir=bd,
                             symbol=sym,
                             from_interval='1h',
                             to_intervals=[('4h', 4), ('1d', 24)],
                             ds_split_name=split_name,
                             ds_split_method_name=split_method,
                             )


import numpy
import theano
import theano.sandbox.rng_mrg
import theano.tensor as T

import nn
import ops
from bridge import map_key
from encoder import Encoder
from decoder import DecoderGruCond
from search import beam, select_nbest

class rnnsearch:
    def __init__(self, **option):
        # source and target embedding dim
        sedim, tedim, xposhdim, yposhdim = option["embdim"]
        # source, target and attention hidden dim
        shdim, thdim, ahdim, xposnn, yposnn, word2pos, pos2word, pos2pos = option["hidden"]
        # maxout hidden dim
        maxdim = option["maxhid"]
        # maxout part
        maxpart = option["maxpart"]
        # deepout hidden dim
        deephid = option["deephid"]
        svocab, tvocab, tagvocab = option["vocabulary"]
        sw2id, sid2w = svocab
        tw2id, tid2w = tvocab

        stag2id, ttag2id = tagvocab
        # source and target vocabulary size
        svsize, tvsize = len(sid2w), len(tid2w)
        stagsize, ttagsize = len(stag2id), len(ttag2id)

        if "scope" not in option or option["scope"] is None:
            option["scope"] = "rnnsearch"

        if "initializer" not in option:
            option["initializer"] = None

        if "regularizer" not in option:
            option["regularizer"] = None

        if "keep_prob" not in option:
            option["keep_prob"] = 1.0

        dtype = theano.config.floatX
        initializer = option["initializer"]
        regularizer = option["regularizer"]
        keep_prob = option["keep_prob"] or 1.0

        scope = option["scope"]
        decoder_scope = "decoder"

        encoder = Encoder(sedim, shdim)
        decoderType = eval("Decoder{}".format(option["decoder"]))
        decoder = decoderType(tedim, thdim, ahdim, 2 * shdim + xposhdim, dim_maxout=maxdim, max_part=maxpart, dim_readout=deephid,
                              n_y_vocab=tvsize, n_y_tagvocab=ttagsize, poshdim=yposhdim, posnndim=yposnn,
                              word2pos=word2pos, pos2word=pos2word, pos2pos=pos2pos)

        # training graph
        with ops.variable_scope(scope, initializer=initializer,
                                regularizer=regularizer, dtype=dtype):
            src_seq = T.imatrix("source_sequence")
            src_mask = T.matrix("source_sequence_mask")
            tgt_seq = T.imatrix("target_sequence")
            tgt_mask = T.matrix("target_sequence_mask")
            src_pos = T.imatrix("source_postag")
            tgt_pos = T.imatrix("target_postag")

            with ops.variable_scope("source_embedding"):
                source_embedding = ops.get_variable("embedding",
                                                    [svsize, sedim])
                source_bias = ops.get_variable("bias", [sedim])

            with ops.variable_scope("target_embedding"):
                target_embedding = ops.get_variable("embedding",
                                                    [tvsize, tedim])
                target_bias = ops.get_variable("bias", [tedim])

            with ops.variable_scope("srctag_embedding"):
                srctag_embedding = ops.get_variable("embedding", [stagsize, xposhdim])
                srctag_bias = ops.get_variable("bias", [xposhdim])

            source_inputs = nn.embedding_lookup(source_embedding, src_seq)
            target_inputs = nn.embedding_lookup(target_embedding, tgt_seq)

            source_inputs = source_inputs + source_bias
            target_inputs = target_inputs + target_bias

            states, r_states = encoder.forward(source_inputs, src_mask)
            annotation = T.concatenate([states, r_states], 2)

            with ops.variable_scope("srcpostagger"):
                tempstates = nn.feedforward(annotation, [shdim*2, xposnn], True,
                                                 scope="staggerstates", activation=T.nnet.relu)
                scores = nn.linear(tempstates, [xposnn, stagsize], True,
                                            scope="staggerscores")

                new_shape = [scores.shape[0] * scores.shape[1], -1]
                scores = scores.reshape(new_shape)
                srcposprobs = T.nnet.softmax(scores)

                srctaggerstates = T.dot(srcposprobs, srctag_embedding) + srctag_bias
                srctaggerstates = srctaggerstates.reshape([annotation.shape[0], annotation.shape[1], -1])

                idx = T.arange(src_pos.flatten().shape[0])
                ce = -T.log(srcposprobs[idx, src_pos.flatten()])
                ce = ce.reshape(src_pos.shape)
                ce = T.sum(ce * src_mask, 0)
                srcpos_cost = T.mean(ce)

            tempposkeys = T.concatenate([srctaggerstates, tempstates], -1)

            src_words_keys = map_key(annotation, 2 * shdim, ahdim, "srcwordkeys")
            src_pos_keys = map_key(tempposkeys, xposnn + xposhdim, word2pos, "srcposkeys")

            pos_words_keys = map_key(annotation, 2 * shdim, pos2word, "pos2wordkeys")
            pos_pos_keys = map_key(tempposkeys, xposnn + xposhdim, pos2pos, "pos2poskeys")

            annotation = T.concatenate([annotation, srctaggerstates], -1)

            # compute initial state for decoder
            # first state of backward encoder
            final_state = T.concatenate([r_states[0], srctaggerstates[0]], -1)
            with ops.variable_scope(decoder_scope):
                initial_state = nn.feedforward(final_state, [shdim + xposhdim, thdim],
                                               True, scope="initial",
                                               activation=T.tanh)

                _, _, transcost, _, tgtpos_cost = decoder.forward(tgt_seq, target_inputs, tgt_mask, src_words_keys, src_pos_keys,
                                                                  pos_words_keys, pos_pos_keys, src_mask,
                                                                annotation, initial_state, tgt_pos, keep_prob)

        lambx = theano.shared(numpy.asarray(option["lambda"][0], dtype), "lambdax")
        lamby = theano.shared(numpy.asarray(option["lambda"][1], dtype), "lambday")

        totalcost = transcost + lambx * srcpos_cost + lamby * tgtpos_cost
        training_inputs = [src_seq, src_mask, tgt_seq, tgt_mask, src_pos, tgt_pos]
        training_outputs = [srcpos_cost, tgtpos_cost, transcost, totalcost]

        # decoding graph
        with ops.variable_scope(scope, reuse=True):
            prev_words = T.ivector("prev_words")

            source_inputs = nn.embedding_lookup(source_embedding, src_seq)
            source_inputs = source_inputs + source_bias
            target_inputs = nn.embedding_lookup(target_embedding, tgt_seq)
            target_inputs = target_inputs + target_bias

            states, r_states = encoder.forward(source_inputs, src_mask)
            annotation = T.concatenate([states, r_states], 2)

            with ops.variable_scope("srcpostagger"):
                tempstates = nn.feedforward(annotation, [shdim*2, xposnn], True,
                                                 scope="staggerstates", activation=T.nnet.relu)
                scores = nn.linear(tempstates, [xposnn, stagsize], True,
                                            scope="staggerscores")

                new_shape = [scores.shape[0] * scores.shape[1], -1]
                scores = scores.reshape(new_shape)
                srcposprobs = T.nnet.softmax(scores)

                srctaggerstates = T.dot(srcposprobs, srctag_embedding) + srctag_bias
                srctaggerstates = srctaggerstates.reshape([annotation.shape[0], annotation.shape[1], -1])

            tempposkeys = T.concatenate([srctaggerstates, tempstates], -1)

            src_words_keys = map_key(annotation, 2 * shdim, ahdim, "srcwordkeys")
            src_pos_keys = map_key(tempposkeys, xposnn + xposhdim, word2pos, "srcposkeys")

            pos_words_keys = map_key(annotation, 2 * shdim, pos2word, "pos2wordkeys")
            pos_pos_keys = map_key(tempposkeys, xposnn + xposhdim, pos2pos, "pos2poskeys")

            annotation = T.concatenate([annotation, srctaggerstates], -1)

            # decoder
            final_state = T.concatenate([r_states[0], srctaggerstates[0]], -1)
            with ops.variable_scope(decoder_scope):
                initial_state = nn.feedforward(final_state, [shdim+xposhdim, thdim],
                                               True, scope="initial",
                                               activation=T.tanh)

            prev_inputs = nn.embedding_lookup(target_embedding, prev_words)
            prev_inputs = prev_inputs + target_bias

            cond = T.neq(prev_words, 0)
            # zeros out embedding if y is 0, which indicates <s>
            prev_inputs = prev_inputs * cond[:, None]

            with ops.variable_scope(decoder_scope):
                mask = T.ones_like(prev_words, dtype=dtype)
                next_state, context, next_pos, tgtposprob = decoder.step(prev_inputs, mask, initial_state,
                                                                         src_words_keys, src_pos_keys,
                                                                         pos_words_keys, pos_pos_keys,
                                                                        annotation, src_mask)
                if option["decoder"] == "GruSimple":
                    probs = decoder.prediction(prev_inputs, initial_state, context)
                elif option["decoder"] == "GruCond":
                    probs = decoder.prediction(prev_inputs, next_state, context, next_pos)

        # encoding
        encoding_inputs = [src_seq, src_mask]
        encoding_outputs = [annotation, initial_state, src_words_keys, src_pos_keys, pos_words_keys, pos_pos_keys, srcposprobs]
        encode = theano.function(encoding_inputs, encoding_outputs)

        if option["decoder"] == "GruSimple":
            prediction_inputs = [prev_words, initial_state, annotation,
                                 mapped_keys, src_mask]
            prediction_outputs = [probs, context]
            predict = theano.function(prediction_inputs, prediction_outputs)

            generation_inputs = [prev_words, initial_state, context]
            generation_outputs = next_state
            generate = theano.function(generation_inputs, generation_outputs)

            self.predict = predict
            self.generate = generate
        elif option["decoder"] == "GruCond":
            prediction_inputs = [prev_words, initial_state, annotation,
                                 src_words_keys, src_pos_keys, pos_words_keys, pos_pos_keys, src_mask]
            prediction_outputs = [probs, next_state, tgtposprob]
            predict = theano.function(prediction_inputs, prediction_outputs, on_unused_input='warn')
            self.predict = predict

        self.cost = totalcost
        self.inputs = training_inputs
        self.outputs = training_outputs
        self.updates = []
        self.encode = encode
        self.option = option

def beamsearch(models, seq, mask=None, beamsize=10, normalize=False,
               maxlen=None, minlen=None, arithmetic=False, dtype=None, suppress_unk=False):
    dtype = dtype or theano.config.floatX

    if not isinstance(models, (list, tuple)):
        models = [models]

    num_models = len(models)

    # get vocabulary from the first model
    option = models[0].option
    vocab = option["vocabulary"][1][1]
    eosid = option["eosid"]
    bosid = option["bosid"]
    unk_sym = models[0].option["unk"]
    unk_id = option["vocabulary"][1][0][unk_sym]

    if maxlen is None:
        maxlen = seq.shape[0] * 3

    if minlen is None:
        minlen = seq.shape[0] / 2

    # encoding source
    if mask is None:
        mask = numpy.ones(seq.shape, dtype)

    outputs = [model.encode(seq, mask) for model in models]

    annotations, states, coarse_words_keys, coarse_pos_keys, mapped_fineword_annots, mapped_finepos_annots = [], [], [], [], [], []
    for item in outputs:
        annotations.append(item[0])
        states.append(item[1])
        coarse_words_keys.append(item[2])
        coarse_pos_keys.append(item[3])
        mapped_fineword_annots.append(item[4])
        mapped_finepos_annots.append(item[5])

    initial_beam = beam(beamsize)
    size = beamsize
    # bosid must be 0
    initial_beam.candidates = [[bosid]]
    initial_beam.scores = numpy.zeros([1], dtype)

    hypo_list = []
    beam_list = [initial_beam]
    done_predicate = lambda x: x[-1] == eosid

    for k in range(maxlen):
        # get previous results
        prev_beam = beam_list[-1]
        candidates = prev_beam.candidates
        num = len(candidates)
        last_words = numpy.array(map(lambda cand: cand[-1], candidates), "int32")

        # compute context first, then compute word distribution
        batch_mask = numpy.repeat(mask, num, 1)
        batch_annots = map(numpy.repeat, annotations, [num] * num_models,
                           [1] * num_models)

        batch_mcoarseword = map(numpy.repeat, coarse_words_keys, [num] * num_models,
                            [1] * num_models)
        batch_mcoarsepos = map(numpy.repeat, coarse_pos_keys, [num] * num_models,
                            [1] * num_models)

        batch_mfineword = map(numpy.repeat, mapped_fineword_annots, [num] * num_models,
                            [1] * num_models)
        batch_mfinepos = map(numpy.repeat, mapped_finepos_annots, [num] * num_models,
                            [1] * num_models)

        # predict returns [probs, next_state, next_pos]
        outputs = [model.predict(last_words, state, annot, coarsemword, coarsempos, finemword, finempos, batch_mask)
                   for model, state, annot, coarsemword, coarsempos, finemword, finempos in
                   zip(models, states, batch_annots,
                       batch_mcoarseword, batch_mcoarsepos, batch_mfineword, batch_mfinepos)]
        prob_dists = [item[0] for item in outputs]

        # search nbest given word distribution
        if arithmetic:
            logprobs = numpy.log(sum(prob_dists) / num_models)
        else:
            # geometric mean
            logprobs = sum(numpy.log(prob_dists)) / num_models

        if suppress_unk:
            logprobs[:, unk_id] = -numpy.inf

        if k < minlen:
            logprobs[:, eosid] = -numpy.inf  # make sure eos won't be selected

        # force to add eos symbol
        if k == maxlen - 1:
            # copy
            eosprob = logprobs[:, eosid].copy()
            logprobs[:, :] = -numpy.inf
            logprobs[:, eosid] = eosprob  # make sure eos will be selected

        next_beam = beam(size)
        finished, remain_beam_indices = next_beam.prune(logprobs, done_predicate, prev_beam)

        hypo_list.extend(finished)  # completed translation
        size -= len(finished)

        if size == 0:  # reach k completed translation before maxlen
            break

        # generate next state
        candidates = next_beam.candidates
        num = len(candidates)
        last_words = numpy.array(map(lambda t: t[-1], candidates), "int32")

        if option["decoder"] == "GruSimple":
            contexts = [item[1] for item in outputs]
            states = select_nbest(states, remain_beam_indices)  # select corresponding states for each model
            contexts = select_nbest(contexts, remain_beam_indices)

            states = [model.generate(last_words, state, context)
                      for model, state, context in zip(models, states, contexts)]
        elif option["decoder"] == "GruCond":
            states = [item[1] for item in outputs]
            states = select_nbest(states, remain_beam_indices)  # select corresponding states for each model
        beam_list.append(next_beam)

    # postprocessing
    if len(hypo_list) == 0:
        score_list = [0.0]
        hypo_list = [[eosid]]
    else:
        score_list = [item[1] for item in hypo_list]
        # exclude bos symbol
        hypo_list = [item[0][1:] for item in hypo_list]

    for i, (trans, score) in enumerate(zip(hypo_list, score_list)):
        count = len(trans)
        if count > 0:
            if normalize:
                score_list[i] = score / count
            else:
                score_list[i] = score

    # sort
    hypo_list = numpy.array(hypo_list)[numpy.argsort(score_list)]
    score_list = numpy.array(sorted(score_list))

    output = []

    for trans, score in zip(hypo_list, score_list):
        trans = map(lambda x: vocab[x], trans)
        output.append((trans, score))

    return output


from sklearn.neighbors import NearestCentroid
import pandas as pd

class CentroidsService():
    def __init__(self):
        self._model_censo = None
        self._model_consumptions = None

    @property
    def model_consumptions(self):
        if self._model_consumptions is None:
            self._load_consumption_model()
        return self._model_consumptions

    @property
    def model_censo(self):
        if self._model_censo is None:
            self._load_censo_model()
        return self._model_censo

    def assign_consumptions_cluster(self, x):
        if x is None:
            return -1
        if -1 in x:
            return -2
        if all(isinstance(x, float) for x in x):
            y = [float(i) if type(i)!=tuple else float(i[0]) for i in x]
            return str(self._model_consumptions.predict([y])[0])
        return -3

    def assign_censo_cluster(self, censo):
        x = [censo.CANT_HOG, censo.EDAD_15A64, censo.HOMBRES,censo.INDIGENAS,
            censo.INMIGRANTES,censo.MATACEP,censo.v1_Casa,censo.v1_Departamento,
            censo.v3a_Albanyileria,censo.v3a_Hormigon,censo.v3b_Hormigon,censo.v3b_Metal]
        x = self._float_convertible(x)
        if x is not None:
            return str(self._model_censo.predict([x])[0])

    def _float_convertible(self, array):
        try:
            return [float(i) for i in array]
        except:
            return None

    def _load_censo_model(self):
        cen_centroids = pd.read_csv('clusteringCensoMedoids.csv', sep=';')
        cen_centroids = cen_centroids.apply(lambda x: x.str.replace(',','.'))
        cen_centroids = cen_centroids.astype(float)
        cen_centroids['cluster'] = [1, 2, 3]
        cen_centroids['cluster'] = cen_centroids['cluster'].astype('category')
        X = cen_centroids.drop(['cluster'], axis=1)
        y = cen_centroids['cluster']
        self._model_censo = NearestCentroid(metric='manhattan')
        self._model_censo.fit( X, y )

    def _load_consumption_model(self):
        cns_centroids  = pd.read_csv('clusteringConsumoMedoids.csv', sep=';')
        cns_centroids = cns_centroids.apply(lambda x: x.str.replace(',','.'))
        cns_centroids = cns_centroids.astype(float)
        cns_centroids['cluster'] = [1, 2, 3, 4]
        cns_centroids['cluster'] = cns_centroids['cluster'].astype('category')
        X = cns_centroids.drop(['cluster'], axis=1)
        y = cns_centroids['cluster']
        self._model_consumptions = NearestCentroid(metric='manhattan')
        self._model_consumptions.fit( X, y )


import os
import re

class Committer:

    def __init__(self, name: str, team_config: dict):
        self.name = name
        self.team = None
        for team, members in team_config.items():
            if name in members:
                self.team = team
                return

class ConcourseJob:

    def __init__(self, id: int, status: str):
        self.id = id
        self.status = status

    def __eq__(self, other):
        return self.id == other.id and self.status == other.status

def get_latest_job(concourse_target: str, concourse_job_name: str) -> ConcourseJob:
    latest_job = os.popen('fly -t {} builds -j {} -c 1'.format(concourse_target, concourse_job_name)).read()
    columns = re.split('\s+', latest_job)
    return ConcourseJob(columns[2], columns[3])

def get_latest_committer(bare_repo_path: str, team_config: dict) -> Committer:
    return_path = os.getcwd()
    os.chdir(bare_repo_path)
    os.system('git fetch')
    latest_committer = os.popen("git log -1 --pretty=format:'%an'").read().replace("'", "")
    os.chdir(return_path)
    return Committer(latest_committer, team_config)


from typing import List

class Solution:
    def addSpaces(self, s: str, spaces: List[int]) -> str:
        ans = s[:spaces[0]] + " "
        for i in range(1, len(spaces)):
            ans += s[spaces[i-1]:spaces[i]] + " "
        ans += s[spaces[-1]:]
        return ans

s = "LeetcodeHelpsMeLearn"
spaces = [8, 13, 15]
print(Solution().addSpaces(s, spaces))