import os
import util
import matplotlib.pyplot as plt
import seaborn as sns
import torch
import data
import numpy as np
import models
import models_continuous
import math
import random
import rendering
from pathlib import Path

# hack for https://github.com/dmlc/xgboost/issues/1715
os.environ["KMP_DUPLICATE_LIB_OK"] = "True"
device = util.get_device()

def plot_losses(
    path,
    theta_losses,
    phi_losses,
    dataset_size,
    prior_losses=None,
    accuracies=None,
    novel_proportions=None,
    new_maps=None,
    batch_split=None,
):
    util.logging.info("plot_losses")
    loss_str = ""
    num_subplots = 2 + len(
        [x for x in [accuracies, novel_proportions, new_maps] if x is not None and len(x) > 0]
    )
    fig, axs = plt.subplots(1, num_subplots, figsize=(2 * num_subplots, 3), dpi=100)
    ax = axs[0]
    ax.plot(theta_losses)
    if prior_losses is not None:
        ax.plot(np.abs(prior_losses))  # abs because I accidentally saved them negative whoops
        ax.plot(
            np.ones(len(prior_losses)) * math.log(dataset_size), color="gray", linestyle="dashed",
        )

    ax.set_xlabel("iteration (batch_split={})".format(batch_split))
    ax.set_ylabel("theta loss")
    ax.set_xticks([0, len(theta_losses) - 1])
    if len(theta_losses) > 0:
        loss_str += "theta loss: {:.2f}\n".format(np.mean(theta_losses[:-100]))

    ax = axs[1]
    ax.plot(phi_losses)
    ax.set_xlabel("iteration")
    ax.set_ylabel("phi loss")
    ax.set_xticks([0, len(phi_losses) - 1])
    if len(phi_losses) > 0:
        loss_str += "phi loss: {:.2f}\n".format(np.mean(phi_losses[:-100]))

    plot_idx = 2
    if accuracies is not None and len(accuracies) > 0:
        ax = axs[plot_idx]
        ax.plot(accuracies)
        ax.set_xlabel("iteration")
        ax.set_ylabel("accuracy")
        ax.set_xticks([0, len(accuracies) - 1])
        if len(accuracies) > 0:
            loss_str += "accuracy: {:.2f}%\n".format(np.mean(accuracies[:-100]) * 100)
        plot_idx = plot_idx + 1

    if novel_proportions is not None and len(novel_proportions) > 0:
        ax = axs[plot_idx]
        ax.plot(novel_proportions)
        ax.set_xlabel("iteration")
        ax.set_ylabel("novel_proportion")
        ax.set_xticks([0, len(novel_proportions) - 1])
        if len(novel_proportions) > 0:
            loss_str += "novel_proportion: {:.2f}%\n".format(
                np.mean(novel_proportions[:-100]) * 100
            )
        plot_idx = plot_idx + 1

    if new_maps is not None and len(new_maps) > 0:
        ax = axs[plot_idx]
        ax.plot(new_maps)
        ax.set_xlabel("iteration")
        ax.set_ylabel("new MAP?")
        ax.set_ylim(0, 0.2)
        ax.set_xticks([0, len(new_maps) - 1])
        if len(new_maps) > 0:
            loss_str += "new_map: {:.2f}%\n".format(np.mean(new_maps[:-100]) * 100)
        plot_idx = plot_idx + 1

    loss_str += "iteration: {}".format(len(theta_losses))

    with open(path + ".txt", "w") as text_file:
        text_file.write(loss_str)

    for ax in axs:
        sns.despine(ax=ax, trim=True)

    util.save_fig(fig, path)

def plot_log_ps(path, log_ps):
    util.logging.info("plot_log_ps")
    fig, ax = plt.subplots(1, 1, figsize=(6, 4), dpi=100)
    ax.plot(log_ps)
    ax.set_xlabel("iteration")
    ax.set_ylabel("log p")
    sns.despine(ax=ax, trim=True)
    util.save_fig(fig, path)

def plot_kls(path, kls):
    util.logging.info("plot_kls")
    fig, ax = plt.subplots(1, 1, figsize=(6, 4), dpi=100)
    ax.plot(kls)
    ax.set_xlabel("iteration")
    ax.set_ylabel("KL")
    sns.despine(ax=ax, trim=True)
    util.save_fig(fig, path)

def sample_from_guide(guide, obs):
    """
    Args:
        guide
        obs [batch_size, num_rows, num_cols]

    Returns:
        ids_and_on_offs [batch_size, num_arcs, 2]
    """
    if isinstance(guide, models.Guide):
        return guide.sample(obs, 1)[0]
    elif isinstance(guide, models_continuous.Guide):
        return guide.sample(obs, 1)[0][0]

def sample_from_memory(memory, generative_model, obs, obs_id):
    """
    Args:
        memory
        generative_model
        obs [batch_size, num_rows, num_cols]
        obs_id

    Returns:
        ids_and_on_offs [batch_size, num_arcs, 2]
    """
    # [batch_size, memory_size, num_arcs, 2]
    memory_latent = memory[obs_id]
    memory_latent_transposed = memory_latent.transpose(
        0, 1
    ).contiguous()  # [memory_size, batch_size, num_arcs, 2]
    if isinstance(generative_model, models.GenerativeModel):
        latent = memory_latent_transposed
    elif isinstance(generative_model, models_continuous.GenerativeModel):
        motor_noise = torch.zeros(
            [*memory_latent_transposed.shape[:-1], 3], device=memory_latent_transposed.device,
        )
        latent = memory_latent_transposed, motor_noise
    memory_log_p = generative_model.get_log_prob(latent, obs)  # [memory_size, batch_size]
    dist = torch.distributions.Categorical(
        probs=util.exponentiate_and_normalize(memory_log_p.t(), dim=1)
    )
    sampled_memory_id = dist.sample()  # [batch_size]
    return torch.gather(
        memory_latent_transposed,
        0,
        sampled_memory_id[None, :, None, None].repeat(1, 1, generative_model.num_arcs, 2),
    )[
        0
    ]  # [batch_size, num_arcs, 2]

def plot_reconstructions(
    path,
    test,
    generative_model,
    guide,
    num_reconstructions,
    dataset,
    data_location,
    memory=None,
    resolution=28,
    data_size="small",
):
    util.logging.info("plot_reconstructions")
    (
        data_train,
        data_valid,
        data_test,
        target_train,
        target_valid,
        target_test,
    ) = data.load_binarized_omniglot_with_targets(location=args.data_location, data_size=data_size)

    if test:
        data_test = torch.tensor(data_test, device=device)
        # obs_id = torch.tensor(
        #     np.random.choice(np.arange(len(data_test)), num_reconstructions), device=device
        # ).long()
        num_reconstructions = min(len(data_test), num_reconstructions)
        obs_id = torch.arange(num_reconstructions, device=device).long()
        obs = data_test[obs_id].float().view(-1, 28, 28)
        obs_id = None
    else:
        data_train = torch.tensor(data_train, device=device)
        # obs_id = torch.tensor(
        #     np.random.choice(np.arange(len(data_train)), num_reconstructions), device=device
        # ).long()
        num_reconstructions = min(len(data_train), num_reconstructions)
        obs_id = torch.arange(num_reconstructions, device=device).long()
        obs = data_train[obs_id].float().view(-1, 28, 28)

    start_point = torch.tensor([0.5, 0.5], device=device)[None].expand(num_reconstructions, -1)
    if memory is None:
        # [batch_size, num_arcs, 2]
        ids_and_on_offs = sample_from_guide(guide, obs)
    else:
        # [batch_size, num_arcs, 2]
        ids_and_on_offs = sample_from_memory(memory, generative_model, obs, obs_id)

    ids = ids_and_on_offs[..., 0]
    on_offs = ids_and_on_offs[..., 1]
    num_arcs = ids.shape[1]

    reconstructed_images = []
    for arc_id in range(num_arcs):
        reconstructed_images.append(
            models.get_image_probs(
                ids[:, : (arc_id + 1)],
                on_offs[:, : (arc_id + 1)],
                start_point,
                generative_model.get_primitives(),
                generative_model.get_rendering_params(),
                resolution,
                resolution,
            ).detach()
        )

    fig, axss = plt.subplots(
        num_reconstructions,
        num_arcs + 1,
        figsize=((num_arcs + 1) * 2, num_reconstructions * 2),
        sharex=False,
        sharey=False,
    )
    for i, axs in enumerate(axss):
        obs_ = obs[i]
        axs[0].imshow(obs_.cpu(), "Greys", vmin=0, vmax=1)
        for j, ax in enumerate(axs[1:]):
            tmp = len(axs[1:])
            char_on = on_offs[i].cpu()
            for k in range(num_arcs):
                if char_on[-k]:
                    tmp = tmp - 1
                if tmp == j:
                    ax.imshow(reconstructed_images[-k][i].cpu(), "Greys", vmin=0, vmax=1)
                    break

            if j < num_arcs:
                ax.text(
                    0.99,
                    0.99,
                    f"{ids[i, j]} {'ON' if on_offs[i, j] else 'OFF'}",
                    horizontalalignment="right",
                    verticalalignment="top",
                    transform=ax.transAxes,
                )

    for axs in axss:
        for ax in axs:
            ax.set_xticks([])
            ax.set_yticks([])
            sns.despine(ax=ax, left=True, right=True, top=True, bottom=True)

    axss[0, 0].set_title(r"Image $x$")
    axss[0, 1].set_title("Sequential\nreconstruction ->")

    util.save_fig(fig, path)

def plot_primitives(path, generative_model):
    util.logging.info("plot_primitives")
    num_primitives = generative_model.num_primitives
    ids = torch.arange(num_primitives, device=device).long().unsqueeze(1)
    on_offs = torch.ones(num_primitives, device=device).long().unsqueeze(1)
    start_point = torch.Tensor([0.5, 0.5]).unsqueeze(0).repeat(num_primitives, 1).to(device)

    primitives = generative_model.get_primitives()
    rendering_params = generative_model.get_rendering_params()
    primitives_imgs = models.get_image_probs(
        ids,
        on_offs,
        start_point,
        primitives,
        rendering_params,
        generative_model.num_rows,
        generative_model.num_cols,
    ).detach()

    num_rows = math.floor(math.sqrt(num_primitives))
    num_cols = num_rows
    fig, axss = plt.subplots(
        num_rows, num_cols, figsize=(num_cols * 2, num_rows * 2), sharex=True, sharey=True
    )

    for i, axs in enumerate(axss):
        for j, ax in enumerate(axs):
            ax.imshow(primitives_imgs[i * num_cols + j].cpu(), "Greys", vmin=0, vmax=1)
            ax.text(
                0.99,
                0.99,
                "{}".format(i * num_cols + j),
                horizontalalignment="right",
                verticalalignment="top",
                transform=ax.transAxes,
            )
            ax.set_xticks([])
            ax.set_yticks([])

    fig.tight_layout(pad=0)
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(path, bbox_inches="tight")
    util.logging.info(f"Saved to {path}")
    plt.close(fig)

def plot_prior(path, generative_model, num_samples, resolution=28):
    util.logging.info("plot_prior")
    try:
        latent, _ = generative_model.sample_latent_and_obs(num_samples=num_samples)
    except NotImplementedError:
        util.logging.info("Can't plot samples for this model")
        return
    generative_model.num_rows, generative_model.num_cols = resolution, resolution
    obs_probs = generative_model.get_obs_params(latent).sigmoid().detach()
    generative_model.num_rows, generative_model.num_cols = 28, 28

    num_rows = math.floor(math.sqrt(num_samples))
    num_cols = num_rows
    if num_rows * num_cols < num_samples:
        util.logging.info(
            "Plotting {} * {} samples instead of {} samples.".format(
                num_rows, num_cols, num_samples
            )
        )
    fig, axss = plt.subplots(num_rows, num_cols, figsize=(2 * num_cols, 2 * num_rows))
    for i, axs in enumerate(axss):
        for j, ax in enumerate(axs):
            ax.imshow(obs_probs[i * num_cols + j].cpu(), cmap="Greys", vmin=0, vmax=1)
            sns.despine(ax=ax, left=True, right=True, top=True, bottom=True)
            ax.set_xticks([])
            ax.set_yticks([])
            if isinstance(generative_model, models.GenerativeModel):
                ids_and_on_offs = latent
            elif isinstance(generative_model, models_continuous.GenerativeModel):
                ids_and_on_offs = latent[0]
            arc_ids = ids_and_on_offs[i * num_cols + j][:, 0]
            on_off_ids = ids_and_on_offs[i * num_cols + j][:, 1]

    util.save_fig(fig, path)

def plot_renderer(path):
    util.logging.info("plot_renderer")

    # Set up sweep
    num_interpolations = 20

    dx = torch.linspace(-0.4, 0.4, num_interpolations)
    dy = torch.linspace(-0.4, 0.4, num_interpolations)
    theta = torch.linspace(-0.99 * math.pi, 0.99 * math.pi, num_interpolations)
    sharpness = torch.linspace(-20, 20, num_interpolations)
    width = torch.linspace(0, 0.1, num_interpolations)
    scale = torch.linspace(0, 10, num_interpolations)
    bias = torch.linspace(-10, 10, num_interpolations)
    interpolations = [dx, dy, theta, sharpness, width]
    interpolations_str = ["dx", "dy", "theta", "sharpness", "width"]

    dx_constant = 0.2
    dy_constant = 0.2
    theta_constant = math.pi / 3
    sharpness_constant = 20.0
    width_constant = 0.01
    scale_constant = 3.0
    bias_constant = -3
    constants = [dx_constant, dy_constant, theta_constant, sharpness_constant, width_constant]

    # Plot
    num_rows, num_cols = 100, 100
    fig, axss = plt.subplots(
        7, num_interpolations, figsize=(num_interpolations * 2, 7 * 2), dpi=200
    )

    # Sweep over `dx`, `dy`, `theta`, `sharpness`, `width`
    for interpolation_id in range(5):
        # Render
        # [num_interpolations, 1, 7]
        arcs = torch.cat(
            # [x_start, y_start]
            [torch.tensor(0.5).repeat(num_interpolations, 2)]
            # keep properties before `interpolation_id` constant
            + [
                torch.tensor(constants[i]).repeat(num_interpolations, 1)
                for i in range(interpolation_id)
            ]
            # interpolate `interpolation_id`'s property
            + [interpolations[interpolation_id][:, None]]
            # keep properties after `interpolation_id` constant
            + [
                torch.tensor(constants[i]).repeat(num_interpolations, 1)
                for i in range(interpolation_id + 1, 5)
            ],
            dim=1,
        )[:, None, :]

        # [num_interpolations, 1]
        on_offs = torch.ones(num_interpolations, 1).long()

        # [2]
        rendering_params = torch.tensor([scale_constant, bias_constant])

        probs = rendering.get_probs(arcs, on_offs, rendering_params, num_rows, num_cols)

        # Plot rendered images
        for i in range(num_interpolations):
            ax = axss[interpolation_id, i]
            ax.imshow(probs[i], vmin=0, vmax=1, cmap="Greys")
            if interpolations_str[interpolation_id] == "theta":
                interpolation = f"{math.degrees(interpolations[interpolation_id][i].item()):.2f}°"
            else:
                interpolation = f"{interpolations[interpolation_id][i]:.2f}"
            ax.text(
                0.99,
                0.99,
                f"{interpolations_str[interpolation_id]} = {interpolation}",
                horizontalalignment="right",
                verticalalignment="top",
                transform=ax.transAxes,
            )

    # Sweep over `scale`, `bias`
    for i in range(num_interpolations):
        # Rendering params
        # [1, 1, 7]
        arcs = torch.cat(
            [torch.tensor(0.5).repeat(1, 2)]
            + [torch.tensor(constants[i]).repeat(1, 1) for i in range(5)],
            dim=1,
        )[:, None, :]

        # [1, 1]
        on_offs = torch.ones(1, 1).long()

        # Render `scale`
        # [2]
        rendering_params = torch.tensor([scale[i], bias_constant])
        prob = rendering.get_probs(arcs, on_offs, rendering_params, num_rows, num_cols)[0]
        ax = axss[5, i]
        ax.imshow(prob, vmin=0, vmax=1, cmap="Greys")
        ax.text(
            0.99,
            0.99,
            f"scale = {scale[i]:.2f}",
            horizontalalignment="right",
            verticalalignment="top",
            transform=ax.transAxes,
        )

        # Render `bias`
        # [2]
        rendering_params = torch.tensor([scale_constant, bias[i]])
        prob = rendering.get_probs(arcs, on_offs, rendering_params, num_rows, num_cols)[0]
        ax = axss[6, i]
        ax.imshow(prob, vmin=0, vmax=1, cmap="Greys")
        ax.text(
            0.99,
            0.99,
            f"bias = {bias[i]:.2f}",
            color="grey",
            horizontalalignment="right",
            verticalalignment="top",
            transform=ax.transAxes,
        )

    # Save fig
    for axs in axss:
        for ax in axs:
            ax.set_xticks([])
            ax.set_yticks([])

    util.save_fig(fig, path)

def plot_motor_noise(path):
    util.logging.info("plot_motor_noise")

    # Set up sweep
    num_interpolations = 5
    num_motor_noise_samples = 10

    dx = torch.linspace(-0.4, 0.4, num_interpolations)
    dy_constant = 0.2
    theta_constant = math.pi / 3
    sharpness_constant = 20.0
    width_constant = 0.01
    scale_constant = 3.0
    bias_constant = -3

    # [num_interpolations, 1 + num_motor_noise_samples, 7]
    arcs = torch.cat(
        [
            torch.ones((num_interpolations, 1 + num_motor_noise_samples, 2)) * 0.5,
            dx[:, None, None].expand(num_interpolations, 1 + num_motor_noise_samples, 1),
            torch.ones((num_interpolations, 1 + num_motor_noise_samples, 1)) * dy_constant,
            torch.ones((num_interpolations, 1 + num_motor_noise_samples, 1)) * theta_constant,
            torch.ones((num_interpolations, 1 + num_motor_noise_samples, 1)) * sharpness_constant,
            torch.ones((num_interpolations, 1 + num_motor_noise_samples, 1)) * width_constant,
        ],
        dim=-1,
    )
    # Add motor noise
    dx_motor_noise = torch.distributions.Normal(0, 0.05).sample(
        (num_interpolations, num_motor_noise_samples)
    )
    dy_motor_noise = torch.distributions.Normal(0, 0.05).sample(
        (num_interpolations, num_motor_noise_samples)
    )
    theta_motor_noise = torch.distributions.Normal(0, math.pi / 18).sample(
        (num_interpolations, num_motor_noise_samples)
    )
    arcs[:, 1:, 2:5] += torch.stack([dx_motor_noise, dy_motor_noise, theta_motor_noise], dim=-1)

    # [num_interpolations, 1 + num_motor_noise_samples, 100, 100]
    probs = rendering.get_probs(
        arcs.view(num_interpolations * (1 + num_motor_noise_samples), 1, 7),
        torch.ones(num_interpolations * (1 + num_motor_noise_samples), 1).long(),
        torch.tensor([scale_constant, bias_constant]),
        100,
        100,
    ).view(num_interpolations, 1 + num_motor_noise_samples, 100, 100)

    # Plot
    num_rows, num_cols = num_interpolations, 1 + num_motor_noise_samples
    fig, axss = plt.subplots(num_rows, num_cols, figsize=(num_cols * 2, num_rows * 2), dpi=200)

    for i in range(num_interpolations):
        for j in range(1 + num_motor_noise_samples):
            ax = axss[i, j]
            ax.imshow(probs[i, j], vmin=0, vmax=1, cmap="Greys")

    axss[0, 0].set_title("Stroke type")
    axss[0, 1].set_title(f"Stroke tokens ->")
    # Save fig
    for axs in axss:
        for ax in axs:
            ax.set_xticks([])
            ax.set_yticks([])

    util.save_fig(fig, path)

def plot_dataset(data_size, max_num_imgs=1000):
    util.logging.info("plot_dataset")
    (
        data_train,
        data_valid,
        data_test,
        target_train,
        target_valid,
        target_test,
    ) = data.load_binarized_omniglot_with_targets(location=args.data_location, data_size=data_size)

    for test, imgs in zip([True, False], [data_test[:max_num_imgs], data_train[:max_num_imgs]]):
        num_cols = 10
        num_rows = math.ceil(len(imgs) / 10)

        fig, axss = plt.subplots(num_rows, num_cols, figsize=(num_cols * 2, num_rows * 2))

        for i, (ax, img) in enumerate(zip(axss.flat, imgs)):
            ax.imshow(img, "Greys", vmin=0, vmax=1)
            ax.text(
                0.99,
                0.99,
                str(i),
                horizontalalignment="right",
                verticalalignment="top",
                transform=ax.transAxes,
            )

        for axs in axss:
            for ax in axs:
                ax.set_xticks([])
                ax.set_yticks([])
                sns.despine(ax=ax, left=True, right=True, top=True, bottom=True)

        filename = "test.pdf" if test else "train.pdf"
        path = f"save/_data/{data_size}/{filename}"
        util.save_fig(fig, path)

def main(args):
    # for data_size in ["mini", "full", "small"]:
    #     plot_dataset(data_size)
    # plot_motor_noise("save/_model/motor_noise.pdf")
    # plot_renderer("save/_model/renderer.pdf")
    dataset = "omniglot"

    if args.checkpoint_path is None:
        import sweep

        checkpoint_paths = [
            util.get_checkpoint_path(sweep_args) for sweep_args in sweep.get_sweep_argss()
        ]
        # checkpoint_paths = list(util.get_checkpoint_paths())
    else:
        checkpoint_paths = [args.checkpoint_path]

    if args.shuffle:
        random.shuffle(checkpoint_paths)

    for checkpoint_path in checkpoint_paths:
        try:
            ((generative_model, guide), optimizer, memory, stats, run_args,) = util.load_checkpoint(
                checkpoint_path, device=device
            )
        except FileNotFoundError as e:
            print(e)
            if "No such file or directory" in str(e):
                print(e)
                continue

        iteration = len(stats.theta_losses)
        (
            data_train,
            data_valid,
            data_test,
            target_train,
            target_valid,
            target_test,
        ) = data.load_binarized_omniglot_with_targets(
            location=args.data_location, data_size=run_args.data_size
        )
        dataset_size = data_train.shape[0]

        diagnostics_dir = util.get_save_dir(run_args)
        Path(diagnostics_dir).mkdir(parents=True, exist_ok=True)

        plot_losses(
            f"{diagnostics_dir}/losses.pdf",
            stats.theta_losses,
            stats.phi_losses,
            dataset_size,
            stats.prior_losses,
            stats.accuracies,
            stats.novel_proportions,
            stats.new_maps,
            1,
        )
        if args.loss_only:
            continue
        if len(stats.kls) > 0:
            plot_log_ps(f"{diagnostics_dir}/logp.pdf", stats.log_ps)
            plot_kls(f"{diagnostics_dir}/kl.pdf", stats.kls)

        for test in [False, True]:
            test_str = "test" if test else "train"

            # Guide
            plot_reconstructions(
                f"{diagnostics_dir}/reconstructions/{test_str}/guide/{iteration}.pdf",
                test,
                generative_model,
                guide,
                args.num_reconstructions,
                dataset,
                args.data_location,
                None,
                args.resolution,
                data_size=run_args.data_size,
            )

            # Memory
            if not test and memory is not None:
                plot_reconstructions(
                    f"{diagnostics_dir}/reconstructions/{test_str}/memory/{iteration}.pdf",
                    test,
                    generative_model,
                    guide,
                    args.num_reconstructions,
                    dataset,
                    args.data_location,
                    memory,
                    args.resolution,
                    data_size=run_args.data_size,
                )

        plot_primitives(f"{diagnostics_dir}/primitives/{iteration}.pdf", generative_model)
        plot_prior(
            f"{diagnostics_dir}/prior/{iteration}.pdf", generative_model, args.num_prior_samples,
        )
        print("-----------------")

def get_parser():
    import argparse

    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument("--repeat", action="store_true", help="")
    parser.add_argument("--shuffle", action="store_true", help="")

    # load checkpoint
    parser.add_argument("--checkpoint-path", type=str, default=None, help=" ")
    parser.add_argument("--data-location", default="local", help=" ")

    # plot
    parser.add_argument("--diagnostics-dir", default="diagnostics", help=" ")
    parser.add_argument("--loss-only", action="store_true")
    parser.add_argument("--resolution", type=int, default=28)
    parser.add_argument("--num-reconstructions", type=int, default=100, help=" ")
    parser.add_argument("--num-prior-samples", type=int, default=100)

    # for custom checkpoints
    parser.add_argument("--checkpoint-iteration", default=None, type=int)
    return parser

if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()

    with torch.no_grad():
        if args.repeat:
            while True:
                main(args)
        else:
            main(args)


#!/usr/bin/env python

import json
import random
import webapp2
import logging
import tdt_database
from collections import Counter
from datetime import datetime
from google.appengine.ext import db
from tdt_database import Tweet, TrackEventCounterShard, EventCauseCounterShard, CurrentEventCounter, EventDescriptionCounterShard

def get_tweets(tasks):
    start_tweets = []
    end_tweets = []
    task_list = json.loads(tasks, encoding='utf8')

    for task in task_list:
        task_data = json.loads(task)
        if task_data["tag"] == "START":
            start_tweets.append(task_data)
        else:
            end_tweets.append(task_data)

    return start_tweets, end_tweets

def update_track_event_counter(track_number, count):
    shard_number = random.randint(0, TrackEventCounterShard.SHARD_COUNT - 1)
    shard = TrackEventCounterShard.all().filter("shard_number=", shard_number).filter("track_number=", track_number).ancestor(tdt_database.counters_key()).get()

    if shard is None:
        shard = TrackEventCounterShard(
            shard_number=shard_number,
            track_number=track_number,
            parent=tdt_database.counters_key()
        )

    shard.event_count += count
    shard.put()

def update_event_cause_counter(cause, count):
    shard_number = random.randint(0, EventCauseCounterShard.SHARD_COUNT - 1)
    shard = EventCauseCounterShard.all().filter("shard_number=", shard_number).filter("cause=", cause).ancestor(tdt_database.counters_key()).get()

    if shard is None:
        shard = EventCauseCounterShard(
            shard_number=shard_number,
            cause=cause,
            parent=tdt_database.counters_key()
        )

    shard.event_count += count
    shard.put()

def update_current_event_counter(track_number, number):
    counter = CurrentEventCounter.all().filter("track_number=", track_number).ancestor(tdt_database.counters_key()).get()

    if counter is None:
        counter = CurrentEventCounter(
            parent=tdt_database.counters_key(),
            track_number=track_number
        )

    counter.balance += number
    counter.put()

def update_word_frequency_counter(word, number):
    shard_number = random.randint(0, EventDescriptionCounterShard.SHARD_COUNT - 1)
    counter = EventDescriptionCounterShard.all().filter("word=", word).ancestor(tdt_database.counters_key()).get()

    if counter is None:
        counter = EventDescriptionCounterShard(
            shard_number=shard_number,
            parent=tdt_database.counters_key(),
            word=word
        )

    counter.frequency += number
    counter.put()

def update_counters(track_event_counter, event_cause_counter, word_frequency_counter, current_event_counter):
    for track, count in track_event_counter.items():
        update_track_event_counter(track, count)

    for cause, count in event_cause_counter.items():
        update_event_cause_counter(cause, count)

    for word, count in word_frequency_counter.items():
        update_word_frequency_counter(word, count)

    for key, count in current_event_counter.items():
        update_current_event_counter(key, count)

def save_new_tweet(tweet):
    dtweet = Tweet(
        parent=tdt_database.tweets_key(),
        tweet_id=long(tweet['tweet_id']),
        publish_time=datetime.fromtimestamp(tweet['publish_time']),
        track_number=tweet['track_number'],
        tag=tweet['tag']
    )
    if 'link' in tweet:
        dtweet.link = tweet['link']
    if 'desired_end' in tweet:
        dtweet.desired_end = datetime.fromtimestamp(tweet['desired_end'])
    if 'cause' in tweet:
        dtweet.cause = tweet['cause']
    if 'description' in tweet:
        dtweet.description = tweet['description']
    dtweet.put()

def save_start_tweets(start_tweets, track_event_counter, event_cause_counter, word_frequency_counter, current_event_counter):
    for tweet in start_tweets:
        save_new_tweet(tweet)
        logging.info("Saving start tweet: " + str(tweet['tweet_id']))
        track_event_counter[tweet['track_number']] += 1
        event_cause_counter[tweet['cause']] += 1
        current_event_counter[tweet['track_number']] += 1

        for word in tweet['description'].split(" "):
            if len(word) > 3:
                word_frequency_counter[word.lower()] += 1

def save_end_tweets(end_tweets, current_event_counter):
    for tweet in end_tweets:
        save_new_tweet(tweet)
        logging.info("Saving end tweet: " + str(tweet['tweet_id']))
        current_event_counter[tweet['track_number']] -= 1

@db.transactional(xg=True)
def save_tweets(start_tweets, end_tweets):
    track_event_counter = Counter()
    event_cause_counter = Counter()
    word_frequency_counter = Counter()
    current_event_counter = Counter()

    save_start_tweets(start_tweets, track_event_counter, event_cause_counter, word_frequency_counter, current_event_counter)
    save_end_tweets(end_tweets, current_event_counter)
    update_counters(track_event_counter, event_cause_counter, word_frequency_counter, current_event_counter)

class ProcessTweets(webapp2.RequestHandler):
    def post(self):
        tasks = self.request.body
        start_tweets, end_tweets = get_tweets(tasks)

        self.response.out.write('Start tweets:
\n')
        for tweet in start_tweets:
            self.response.out.write('  %s
\n' % json.dumps(tweet, ensure_ascii=False, encoding='utf8'))
        self.response.out.write('
End tweets:
\n')
        for tweet in end_tweets:
            self.response.out.write('  %s
\n' % json.dumps(tweet, ensure_ascii=False, encoding='utf8'))
        self.response.out.write('
')

        save_tweets(start_tweets, end_tweets)
        self.response.out.write('Tweets processed successfully!
')

app = webapp2.WSGIApplication([('/process_tweets', ProcessTweets)], debug=True)


import boto3
# Create a table called Employee

dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
mytable = dynamodb.create_table(
TableName= 'Employee',KeySchema=[
    {
        'KeyType': 'HASH',
        'AttributeName': 'Id'
    },
    {
        'KeyType': 'RANGE',
        'AttributeName': 'Sal'
    }],
    AttributeDefinitions=[
    {
        'AttributeName': 'Id',
        'AttributeType': 'N'
    },
    {
        'AttributeName': 'Sal',
        'AttributeType': 'N'
    }
    ],
    ProvisionedThroughput={
        'ReadCapacityUnits': 2,
        'WriteCapacityUnits': 2
    }
)
# Wait until the table exists.
mytable.meta.client.get_waiter('table_exists').wait(TableName='Employee')
print('Table is ready, please continue to isert data.')

# Insert the data into dynamodb table
mytable.put_item(
Item={
'Id': 1,
'FirstName': 'Ramasankar',
'LastName': 'Molleti',
'Dept': 'IT',
'Sal': 5000
}
)
mytable.put_item(
Item={
'Id': 1,
'FirstName': 'Sourav',
'LastName': 'Mukherjee',
'Dept': 'IT',
'Sal': 10000
}
)
mytable.put_item(
Item={
'Id': 1,
'FirstName': 'Praveen',
'LastName': 'Kumar',
'Dept': 'Finance',
'Sal': 5000
}
)
mytable.put_item(
Item={
'Id': 1,
'FirstName': 'Suresh',
'LastName': 'Kumar',
'Dept': 'Finance',
'Sal': 12000
}
)

response = mytable.scan()

for i in response['Items']:
    print("added item:", i['Id'], ":", i['FirstName'], ":", i['LastName'], ":", i['Dept'], ":", i['Sal'])


#!/usr/bin/env python3

import argparse
import os
import sys
import time
import credb
import multiprocessing
import numpy as np
from tqdm import tqdm
from test import *

def uniform(length, ops, args):
    assert len(args) == 0
    return np.random.randint(length, size=ops)

def zipf(length, ops, args):
    assert len(args) == 1
    a = float(args[0])
    return np.random.zipf(a, size=ops)

def normal(length, ops, args):
    assert len(args) == 2
    mean = float(args[0])
    stddev = np.sqrt(float(args[0]))
    return mean + np.random.randn(ops) * stddev

def get_distribution(length, ops, args):
    args = args.split('-')
    distributions = [uniform, zipf, normal]
    d = {x.__name__.lower(): x for x in distributions}
    xs = d[args[0].lower()](length, ops, args[1:])
    offset = random.randrange(length)
    return (xs + offset) % length

def create_client(args):
    return credb.create_client('latency bench client', args.server_name, args.host, args.port)

def do_bench(args):
    c = create_client(args)
    obj = random_string(args.obj_size)
    read_distribution = iter(get_distribution(args.num_objs, args.num_ops, args.read_distribution))
    write_distribution = iter(get_distribution(args.num_objs, args.num_ops, args.write_distribution))
    res = np.empty(args.num_ops, dtype=np.float64)

    for i in tqdm(range(args.num_ops)):
        if random.random() < args.prob_write:
            key = 'latency' + str(next(write_distribution))
            st = time.perf_counter()
            c.put(key, obj)
            ed = time.perf_counter()
        else:
            key = 'latency' + str(next(read_distribution))
            st = time.perf_counter()
            c.get(key)
            ed = time.perf_counter()
        res[i] = (ed - st) * 1e6

    res.sort()
    print('mean:', np.mean(res), 'us')
    print('std :', np.std(res), 'us')
    print('50% :', res[int(args.num_ops * 0.5)], 'us')
    print('70% :', res[int(args.num_ops * 0.7)], 'us')
    print('90% :', res[int(args.num_ops * 0.9)], 'us')
    print('99% :', res[int(args.num_ops * 0.99)], 'us')
    print('99.9:', res[int(args.num_ops * 0.999)], 'us')
    print('9999:', res[int(args.num_ops * 0.9999)], 'us')
    np.save(args.output, res)
    print('written to', args.output)

def do_prepare_worker(args, start, count):
    c = create_client(args)
    obj = random_string(args.obj_size)
    for i in range(start, start + count):
        key = 'latency' + str(i)
        assert c.put(key, obj)

def do_prepare(args):
    ps = []
    num_workers = 16
    per_worker = int(np.ceil(args.num_objs / num_workers))
    for i in range(num_workers):
        p = multiprocessing.Process(target=do_prepare_worker, args=(args, per_worker*i, per_worker))
        p.start()
        ps.append(p)
    for p in ps:
        p.join()
    print('done')

if __name__ == '__main__':
    action_dict = {'bench': do_bench, 'prepare': do_prepare}
    parser = argparse.ArgumentParser()
    parser.add_argument('--server_name', type=str, default='latency-bench')
    parser.add_argument('--host', type=str, default='localhost')
    parser.add_argument('--port', type=int, default=5042)
    parser.add_argument('--num_objs', type=int, default=1000000)
    parser.add_argument('--num_ops', type=int, default=100000)
    parser.add_argument('--obj_size', type=int, default=1024)
    parser.add_argument('--prob_write', type=int, default=0)
    parser.add_argument('--read_distribution', type=str, default='zipf-1.5')
    parser.add_argument('--write_distribution', type=str, default='uniform')
    parser.add_argument('--output', type=str, default='latency.csv')
    parser.add_argument('action', choices=list(action_dict.keys()))
    args = parser.parse_args()

    action_dict[args.action](args)


import pytest

def test_remote_301_post():
    import logging

    logging.basicConfig(level="DEBUG")

    from oda_api.api import DispatcherAPI, URLRedirected
    from oda_api.token import discover_token

    api_cat_str = '{"cat_frame": "fk5", "cat_coord_units": "deg", "cat_column_list": [[17, 87], ["GRS 1915+105", "MAXI J1820+070"], [29.396455764770508, 1803.1607666015625], [288.799560546875, 275.0911865234375], [10.939922332763672, 7.185144901275635], [-32768, -32768], [2, 2], [0, 0], [0.0002800000074785203, 0.00041666667675599456]], "cat_column_names": ["meta_ID", "src_names", "significance", "ra", "dec", "NEW_SOURCE", "ISGRI_FLAG", "FLAG", "ERR_RAD"], "cat_column_descr": [["meta_ID", "<i8"], ["src_names", "<U20"], ["significance", "<f8"], ["ra", "<f8"], ["dec", "<f8"], ["NEW_SOURCE", "<i8"], ["ISGRI_FLAG", "<i8"], ["FLAG", "<i8"], ["ERR_RAD", "|O"]], "cat_lat_name": "dec", "cat_lon_name": "ra"}'

    with pytest.raises(URLRedirected) as execinfo:
        disp = DispatcherAPI(url="http://www.astro.unige.ch/mmoda/dispatch-data")
        spectrum = disp.get_product(
            instrument="isgri",
            product="isgri_spectrum",
            product_type="Real",
            osa_version="OSA11.1",
            RA="275.09142677",
            DEC="7.18535523",
            radius="8",
            T1="58193.455",
            T2="58246.892",
            T_format="mjd",
            max_pointings="50",
            token=discover_token(),
            selected_catalog=api_cat_str,
        )

    assert execinfo.value.args[0] == ('the service was moved permanently, please reinitialize DispatcherAPI '
                                      'with "https://www.astro.unige.ch/mmoda/dispatch-data/run_analysis" '
                                      '(you asked for "http://www.astro.unige.ch/mmoda/dispatch-data/run_analysis")')


from this import d
from django.shortcuts import render

from rest_framework import generics
from rest_framework.views import APIView
from rest_framework import status
from rest_framework.response import Response
from django.views.decorators.csrf import csrf_exempt
from django.http import JsonResponse
import datetime
import json

from .models import Version,Device,Data,DateSent,SensorValues
from .serializers import VersionSerializer,DeviceSerializer

# Create your views here.

def all_versions(request):
    if request.method == 'GET':
        versions = Version.objects.all()
        data = [version.vv for version in versions]
        return JsonResponse({
            'payload': data,
            'status' : 'versions extracted successfully'
        })

def all_devices(request):
    if request.method == 'GET':
        devices = Device.objects.all()
        data = [device.dd for device in devices]
        return JsonResponse({
            'payload': data,
            'status' : 'Device ids extracted successfully'
        })

def versions_date(request):
    if request.method == 'GET':
        sent_date = request.GET['date']
        sent_date = datetime.datetime.strptime(sent_date,"%d-%m-%Y")
        sent_date = sent_date.date()
        all_data = Data.objects.filter(sent_date=sent_date)
        versions = [data.version.vv for data in all_data]
        version = []
        for v in versions:
            if v not in version:
                version.append(v)

        return JsonResponse({
            'payload' : version,
            'status' : 'versions extracted successfully for required date'
        })

def highest_data(request):
    if request.method == 'GET':
        sent_date = request.GET['date']
        sent_date = datetime.datetime.strptime(sent_date,"%d-%m-%Y")
        sent_date = sent_date.date()
        all_data = Data.objects.filter(sent_date=sent_date)
        tm = 0.00
        hm = 0.00
        for data in all_data:
            if float(data.dt['tm']) > tm:
                tm = float(data.dt['tm'])
            if float(data.dt['hm']) > hm:
                hm = float(data.dt['hm'])
        output = [{'tm':tm},{'hm':hm}]
        return JsonResponse({
            'payload' : output,
            'status' : 'highest tm and hm values for given date extracted'
        })

def no_of_data(request):
    if request.method == 'GET':
        start_date = request.GET['start_date']
        start_date = datetime.datetime.strptime(start_date,"%d-%m-%Y")
        start_date = start_date.date()
        end_date = request.GET['end_date']
        end_date = datetime.datetime.strptime(end_date,"%d-%m-%Y")
        end_date = end_date.date()
        version = request.GET['vn']
        version = str(version)
        versions = Version.objects.all()
        for v in versions:
            if v.vv==version:
                version=v
                break
        all_data = Data.objects.filter(sent_date__range=[start_date, end_date],version=version)
        output = [{'datapoints':len(all_data)}]
        return JsonResponse({
            'payload' : output,
            'status' : 'number of data points for given version extracted'
        })

def no_of_days(request,n):
    if request.method == 'GET':
        n=int(n)
        data = DateSent.objects.filter(noof_versions=n)
        days = len(data)
        output = [{'days':days}]
        return JsonResponse({
            'payload' : output,
            'status' : 'required number of days extracted'
        })

def is_equal(dt,dt2):
    if str(dt.get("tm",0))==str(dt2.get("tm",0)) and str(dt.get("hm",0))==str(dt2.get("hm",0)) and str(dt.get("wd",0))==str(dt2.get("wd",0)) and str(dt.get("ws",0))==str(dt2.get("ws",0)) and str(dt.get("sm",0))==str(dt2.get("sm",0)) and str(dt.get("st",0))==str(dt2.get("st",0)) and str(dt.get("sc",0))==str(dt2.get("sc",0)) and str(dt.get("lt",0))==str(dt2.get("lt",0)) and str(dt.get("lw",0))==str(dt2.get("lw",0)) and str(dt.get("bl",0))==str(dt2.get("bl",0)) and str(dt.get("pv",0))==str(dt2.get("pv",0)) :
        return True

    return False

def update_versions():
    data = Data.objects.all()
    versions = Version.objects.all()
    data_vv = []
    version_vv = []
    for d in data:
        if d.version.vv not in data_vv:
            data_vv.append(d.version.vv)
    for v in versions:
        if v.vv not in version_vv:
            version_vv.append(v.vv)
    for v in version_vv:
        if v not in data_vv:
            version = Version.objects.get(vv=v).delete()
    return

@csrf_exempt
def convert_version(request):
    if request.method == 'POST':
        request_body_unicode = request.body.decode('utf-8')
        request_body = json.loads(request_body_unicode)
        dataepoch = request_body['dataepoch']
        ep=str(dataepoch)
        dt = request_body['dd']
        new_vn = request_body['new_vn']
        version = str(new_vn)
        versions = Version.objects.all()
        f=0
        for v in versions:
            if v.vv==version:
                version=v
                f=1
                break
        if f==0:
            return JsonResponse(
                {
                    'status' : 'Version not found.'
                },
                status=400
            )
        values = {
                "tm" : str(dt.get("tm",0)),
                "hm" : str(dt.get("hm",0)),
                "pp" : str(dt.get("pp",0)),
                "wd" : str(dt.get("wd",0)),
                "ws" : str(dt.get("ws",0)),
                "sm" : str(dt.get("sm",0)),
                "st" : str(dt.get("st",0)),
                "sc" : str(dt.get("sc",0)),
                "lt" : str(dt.get("lt",0)),
                "lw" : str(dt.get("lw",0)),
                "bl" : str(dt.get("bl",0)),
               "pv" : str(dt.get("pv",0))
        }
        datas = Data.objects.filter(ep=ep)
        u=0
        for data in datas:
            data_values=data.dt
            if(is_equal(data_values,values)):
                version1=data.version
                data.version=version
                data.save()
                sent_date = data.sent_date
                u=1
                try:
                    date = DateSent.objects.get(date=sent_date)
                    versions = date.versions
                    f=0
                    for v in  versions:
                        if v["version"] == version1.vv:
                            v["no"] = v["no"]-1
                        if v["version"] == version.vv:
                            v["no"] = v["no"]+1
                            f=1
                        if v["no"] == 0:
                            versions.remove(v)
                    if f==0:
                        versions.append({'version':version.vv,'no':1})
                    date.noof_versions = len(versions)
                    date.versions = versions
                    date.save()
                except Exception as e:
                    print(e)
                print("saved")
        if u==1:
            update_versions()

        return JsonResponse(
            {
            'status' : 'Datapoint successfully updated with new version'
            },
            status=201
        )

@csrf_exempt
def exchange_tm_hm(request):
    if request.method == 'POST':
        request_body_unicode = request.body.decode('utf-8')
        request_body = json.loads(request_body_unicode)
        dataepoch = request_body['dataepoch']
        ep=str(dataepoch)
        dt = request_body['dd']
        values = {
                "tm" : str(dt.get("tm",0)),
                "hm" : str(dt.get("hm",0)),
                "pp" : str(dt.get("pp",0)),
                "wd" : str(dt.get("wd",0)),
                "ws" : str(dt.get("ws",0)),
                "sm" : str(dt.get("sm",0)),
                "st" : str(dt.get("st",0)),
                "sc" : str(dt.get("sc",0)),
                "lt" : str(dt.get("lt",0)),
                "lw" : str(dt.get("lw",0)),
                "bl" : str(dt.get("bl",0)),
               "pv" : str(dt.get("pv",0))
        }
        datas = Data.objects.filter(ep=ep)
        for data in datas:
            data_values=data.dt
            if(is_equal(data_values,values)):
                temp = data.dt['tm']
                data.dt['tm']=data.dt['hm']
                data.dt['hm'] = str(temp)
                data.save()
                print("saved")
        return JsonResponse(
            {
            'status' : 'tm and hm in data point are successfully exchanged'
            },
            status=201
        )

@csrf_exempt
def delete_data_point(request):
    if request.method == 'DELETE':
        request_body_unicode = request.body.decode('utf-8')
        request_body = json.loads(request_body_unicode)
        dataepoch = request_body['dataepoch']
        ep=str(dataepoch)
        dt = request_body['dd']
        values = {
                "tm" : str(dt.get("tm",0)),
                "hm" : str(dt.get("hm",0)),
                "pp" : str(dt.get("pp",0)),
                "wd" : str(dt.get("wd",0)),
                "ws" : str(dt.get("ws",0)),
                "sm" : str(dt.get("sm",0)),
                "st" : str(dt.get("st",0)),
                "sc" : str(dt.get("sc",0)),
                "lt" : str(dt.get("lt",0)),
                "lw" : str(dt.get("lw",0)),
                "bl" : str(dt.get("bl",0)),
               "pv" : str(dt.get("pv",0))
        }
        datas = Data.objects.filter(ep=ep)
        u=0
        for data in datas:
            data_values=data.dt
            if(is_equal(data_values,values)):
                u=1
                sent_date = data.sent_date
                version = data.version.vv
                date=DateSent.objects.get(date=sent_date)
                versions = date.versions
                for v in  versions:
                        if v["version"] == version:
                            v["no"] = v["no"]-1

                        if v["no"] == 0:
                            versions.remove(v)

                date.noof_versions = len(versions)
                date.versions = versions
                date.save()
                data.delete()
                print("deleted.")
        if u==1:
            update_versions()
        return JsonResponse(
            {
            'status' : 'Datapoint deleted succesfully'
            },
            status=202
        )


import json
import openai
import streamlit as st

from streamlit_chat import message
from streamlit_option_menu import option_menu
from credentials import *
from cog_search import *

def update_creds():
    with open('credentials.py') as f:
        l = list(f)

    for attribute, value in creds.items():
        with open('credentials.py', 'w') as output:
            for line in l:
                    if line.startswith(attribute):
                        print('found attribute: ' + attribute + ' = "' +  value + '"\n')
                        print('about to write: ' + attribute + ' = "' +  value + '"\n')
                        output.write( attribute + ' = "' +  value + '"\n')
                    else:
                        output.write(line)
    f.close()
    output.close()

#reads from credentials.pyand puts values of creds into session_state
def set_creds():
    for attribute, value in creds.items():
        if attribute not in st.session_state:
            st.session_state[attribute] = value

#print session_state message for chat like experience
def print_messages():
        for i in range (len(st.session_state.messages) -1, -1, -1):
            msg = st.session_state.messages[i]
            if msg is not None:
                if msg["role"] == "user":
                    message(msg["content"], is_user=True, key = str(i) + "user", avatar_style = "initials", seed = "ðŸ‘¤")
                else:
                    if msg["role"] == "assistant":
                        print(msg["content"])
                        if msg["content"] == "I don't know" or msg["content"] == "I don't know." or msg['content'] == "Sorry, I don't know the answer to that question. Please try rephrasing your question.":
                            message(msg["content"], is_user=False, key = str(i) + "system", avatar_style="initials", seed = "ðŸ˜•")
                        elif msg['content'] == "How can I help you?":
                            message(msg["content"], is_user=False, key = str(i) + "system", avatar_style="initials", seed = "ðŸ™‚")
                        else:
                            message(msg["content"], is_user=False, key = str(i) + "system", avatar_style="initials", seed = "ðŸ˜‰")

#clear session_state messages - this is on the settings page
def settings_form():
    reset = st.checkbox('Reset Messages')
    if reset:
        st.write('Sure thing!')
        st.session_state.messages = [{"role":"system","content":"You are an AI assistant that helps people find information."}]
        st.session_state.messages.append({"role": "assistant", "content": "How can I help you?"})
        print("complteted reset")

#display home

#main screen
if __name__ == "__main__":

    if "messages" not in st.session_state:
        print("messages not in session state")
        st.session_state["messages"] = [{"role":"system","content":"You are an AI assistant that helps people find information."}]
        st.session_state.messages.append({"role": "assistant", "content": "How can I help you?"})

    with st.sidebar:
        set_creds()
        menu_index = ['Home', 'Settings', 'Upload file', 'Chat']
        menu_icons = ['house', 'gear', 'cloud-upload',  'chat']

        selected = option_menu("Main Menu",menu_index, icons=menu_icons, menu_icon="cast", default_index=1)

    if selected == 'Home':
        "[View the source code](https://github.com/streamlit/llm-examples/blob/main/Chatbot.py)"

        st.title("Welcome to the AI Assistant")
        with st.form("chat_input", clear_on_submit=True):
            a, b = st.columns([4, 1])
            user_input = a.text_input(
                label="Your message:",
                placeholder="What would you like to say?",
                label_visibility="collapsed",
            )
            b.form_submit_button("Send")

        openai.api_type = "azure"
        openai.api_base = creds['AZURE_OPENAI_ENDPOINT']
        openai.api_version = "2023-03-15-preview"
        openai.api_key = creds['AZURE_OPENAI_KEY']

        if user_input and AZURE_OPENAI_KEY:
            st.session_state.messages.append({"role": "user", "content": user_input})
            response = openai.ChatCompletion.create(
                engine="gpt-35-turbo",
                messages = st.session_state.messages,
                temperature=0.0,
                max_tokens=200,
                top_p=0.95,
                frequency_penalty=0,
                presence_penalty=0,
                stop=None)
            print(response)
            if response.choices[0].message.content != None:
                st.session_state.messages.append({"role": "assistant", "content": response.choices[0].message.content})

        print_messages()

    elif selected == 'Upload file':
        st.title('Upload file')
    elif selected == 'Chat':
        st.title('Chat')
        "[View the source code](https://github.com/streamlit/llm-examples/blob/main/Chatbot.py)"
        st.title("Cognitive Search & Azure OpenAI")
        with st.form("chat_input", clear_on_submit=True):
            a, b = st.columns([4, 1])
            user_input_ondata = a.text_input(
                label="Your message:",
                placeholder="What would you like to ask?",
                label_visibility="collapsed",
            )
            b.form_submit_button("Send")

        if user_input_ondata and AZURE_OPENAI_KEY:
            question = user_input_ondata
            st.session_state.messages.append({"role": "user", "content": question})
            arg = OpenAIHelper(creds['COG_SEARCH_INDEX'])
            response = arg.get_Answer(user_input_ondata)
            st.session_state.messages.append({"role": "assistant", "content": response})

        print_messages()

    elif selected == 'Settings':
        settings_form()

        with st.form("my_form"):
            st.write("Configuration Settings")

            azure_openai_endpoint = st.text_input("azure_openai_endpoint", creds["AZURE_OPENAI_ENDPOINT"])

            azure_openai_key = st.text_input("azure_openai_key", creds["AZURE_OPENAI_KEY"], type = "password")
            txt_davinci = st.text_input("txt davinici", creds["TEXT_DAVINCI"])

            cog_search_resource = st.text_input("Cog Search Resource",creds["COG_SEARCH_RESOURCE"])
            cog_search_index = st.text_input("Cog Search Index", creds["COG_SEARCH_INDEX"])
            cog_service_key   = st.text_input("Cog Search Key", creds["COG_SEARCH_KEY"], type = "password")

            storage_connection_string = st.text_input("Storage Connection String", creds["STORAGE_CONNECTION_STRING"], type="password")
            storage_account = st.text_input("Storage Account", creds["STORAGE_ACCOUNT"])
            storage_container = st.text_input("Storage Container", creds["STORAGE_CONTAINER"])
            storage_key = st.text_input("Storage Key", creds["STORAGE_KEY"], type = "password")

            submitted = st.form_submit_button("Submit")

            #don't use this to update the search index.
            if submitted:
                creds["AZURE_OPENAI_ENDPOINT"] = azure_openai_endpoint
                creds["AZURE_OPENAI_KEY"] = azure_openai_key
                creds["TEXT_DAVINCI"] = txt_davinci
                creds["COG_SEARCH_RESOURCE"] = cog_search_resource
                #creds["COG_SEARCH_INDEX"] = cog_search_index
                creds["COG_SEARCH_KEY"] = cog_service_key
                creds["STORAGE_CONNECTION_STRING"] = storage_connection_string
                creds["STORAGE_ACCOUNT"] = storage_account
                creds["STORAGE_CONTAINER"] = storage_container
                creds["STORAGE_KEY"] = storage_key
                set_creds()
                # update_creds()
                st.write("Settings updated")

        with st.form("create index"):

            st.write("Create Index")
            create_index = st.form_submit_button("SubmitCreateIndex")

        if create_index:
            cogSearch = CogSearchHelper(index = creds["COG_SEARCH_INDEX"])
            response, success = cogSearch.create_datasource()
            if success:
                st.write("Data source created")
                response, success = cogSearch.create_skillset()
                if success:
                    st.write("Skillset created")
                    response, success = cogSearch.create_index()
                    if success:
                        st.write("Index created")
                        with st.spinner(text="In progress..."):
                            response, success = cogSearch.create_indexer()
                            if success:
                                st.write("Running indexer")


import numpy as np
import matplotlib.pyplot as plt
import pyfftw
import time

from numpy import pi

# Parameters
nu = 8e-5
Lx = 2.0*pi
nx = 128
dt = 1e-2
nsteps = 5000
nthreads = 4

# Construct grid.
dx = Lx/nx
nk, nl = nx//2+1, nx
x = np.arange(0.0, Lx, dx)

k = 2.0*pi/Lx * np.arange(0.0, nk)
l = 2.0*pi/Lx * np.append(
    np.arange(0.0, nx/2.0), np.arange(-nx/2.0, 0.0))

X, Y = np.meshgrid(x, x)
K, L = np.meshgrid(k, l)
Ksq = K**2.0 + L**2.0

divsafeKsq = Ksq.copy()
divsafeKsq[0, 0] = float('Inf')
invKsq = 1.0/divsafeKsq

# Define fft
usefftw = True
if usefftw:
    effort = 'FFTW_MEASURE'
    fft2 = (lambda x: pyfftw.interfaces.numpy_fft.rfft2(x,
        threads=nthreads, planner_effort=effort))
    ifft2 = (lambda x: pyfftw.interfaces.numpy_fft.irfft2(x,
        threads=nthreads, planner_effort=effort))
else:
    fft2 = np.fft.rfft2
    ifft2 = np.fft.irfft2

# Initial condition
t = 0.0
q = np.random.rand(nx, nx)
qh = fft2(q)

# Step forward
fig = plt.figure(1)
t1 = time.time()
for step in range(nsteps):

    if step % 1000 is 0:
        plt.clf()
        plt.imshow(q)
        plt.pause(0.01)

        print("step = {:4d}, t = {:06.1f} s, wall time = {:.3f}".format(
            step, t, time.time()-t1))
        t1=time.time()

    # Calculate right hand side
    psih = -qh * invKsq

    q = ifft2(qh)
    u = ifft2(1j*L*psih)
    v = -ifft2(1j*K*psih)

    rhs = -1j*K*fft2(u*q) - 1j*L*fft2(v*q) - nu*Ksq*qh

    qh += dt*rhs
    t += dt


import pandas as pd
import numpy as np

# Для работы с матрицами
from scipy.sparse import csr_matrix

# Матричная факторизация
from implicit.als import AlternatingLeastSquares
from implicit.nearest_neighbours import ItemItemRecommender  # нужен для одного трюка
from implicit.nearest_neighbours import bm25_weight, tfidf_weight

class MainRecommender:
    """Рекоммендации, которые можно получить из ALS

    Input
    -----
    user_item_matrix: pd.DataFrame
        Матрица взаимодействий user-item
    """

    def __init__(self, data, weighting=True, show_progress=False):

        self.user_item_matrix = data  # pd.DataFrame
        self.user_item_matrix = self.user_item_matrix.astype(float)

        self.id_to_itemid, self.id_to_userid, self.itemid_to_id, self.userid_to_id = self.prepare_dicts(self.user_item_matrix)

        if weighting:
            self.user_item_matrix = bm25_weight(self.user_item_matrix.T).T

        self.model = self.fit(self.user_item_matrix, show_progress=show_progress)

        binary_user_item_matrix = data.copy()
        binary_user_item_matrix[binary_user_item_matrix > 0] = 1.
        binary_user_item_matrix = binary_user_item_matrix.astype(float)
        self.own_recommender = self.fit_own_recommender(binary_user_item_matrix, show_progress=show_progress)

    @staticmethod
    def prepare_dicts(user_item_matrix):
        """Подготавливает вспомогательные словари"""

        userids = user_item_matrix.index.values
        itemids = user_item_matrix.columns.values

        matrix_userids = np.arange(len(userids))
        matrix_itemids = np.arange(len(itemids))

        id_to_itemid = dict(zip(matrix_itemids, itemids))
        id_to_userid = dict(zip(matrix_userids, userids))

        itemid_to_id = dict(zip(itemids, matrix_itemids))
        userid_to_id = dict(zip(userids, matrix_userids))

        return id_to_itemid, id_to_userid, itemid_to_id, userid_to_id

    @staticmethod
    def fit_own_recommender(user_item_matrix, show_progress=False):
        """Обучает модель, которая рекомендует товары, среди товаров, купленных юзером"""

        own_recommender = ItemItemRecommender(K=1, num_threads=4)
        own_recommender.fit(csr_matrix(user_item_matrix).T.tocsr(), show_progress=show_progress)

        return own_recommender

    def get_own_recommends(self, user_id, n_recommends=5, filter_999999=True):
        recs = self.own_recommender.recommend(userid=self.userid_to_id[user_id],  # userid - id от 0 до N
                            user_items=csr_matrix(self.user_item_matrix).tocsr(),   # на вход user-item matrix
                            N=n_recommends, # кол-во рекомендаций
                            filter_already_liked_items=False,
                            filter_items= [self.itemid_to_id[999999]] if filter_999999 else None,
                            recalculate_user=True)

        recs = [self.id_to_itemid[rec[0]] for rec in recs]

        return recs

    @staticmethod
    def fit(user_item_matrix, n_factors=20, regularization=0.001, iterations=15, num_threads=4, show_progress=False):
        """Обучает ALS"""

        model = AlternatingLeastSquares(factors=n_factors,
                                             regularization=regularization,
                                             iterations=iterations,
                                             num_threads=num_threads)
        model.fit(csr_matrix(user_item_matrix).T.tocsr(), show_progress=show_progress)

        return model

    def get_similar_items_recommendation(self, user_id, N=5, filter_999999=True):
        """Рекомендуем товары, похожие на топ-N купленных юзером товаров"""

        own_recs = self.own_recommender.recommend(userid=self.userid_to_id[user_id],  # userid - id от 0 до N
                            user_items=csr_matrix(self.user_item_matrix).tocsr(),   # на вход user-item matrix
                            N=N, # кол-во рекомендаций
                            filter_already_liked_items=False,
                            filter_items=[self.itemid_to_id[999999]] if filter_999999 else None,
                            recalculate_user=True)

        # Для каждой найденной собственной рекомендации ищем максиально похожий товар
        # Если найденный товар уже есть в списке рекомендаций, ищем следующий
        recs = []
        for own_rec in own_recs:
            N_neigbors = 2
            success = False
            while not success:
                rec = self.model.similar_items(own_rec[0], N=N_neigbors)
                rec = self.id_to_itemid[rec[-1][0]]
                if rec not in recs:
                    recs.append(rec)
                    success = True
                else:
                    N_neigbors += 1

        if len(recs) != 5:
            print('Error: Рекоммендаций меньше, чем ожидается. user_id:', user_id)
        #assert len(recs) == N, 'Количество рекомендаций != {}'.format(N)
        return recs

    def get_similar_users_recommendation(self, user_id, N=5, filter_999999=True):
        """Рекомендуем топ-N товаров, среди купленных похожими юзерами"""

        users = self.model.similar_users(userid=self.userid_to_id[user_id], N=N)

        max_iterations = 10
        recs = []
        # Находим N похожих пользователей и берем самый популярный товар от каждого пользователя
        # Если найденный товар уже есть в списке рекомендаций, ищем следующий популярный у пользователя
        for user in users:
            N_recommends = 1
            success = False

            iterations = 0
            while not success:
                iterations += 1
                rec = self.own_recommender.recommend(userid=user[0],  # userid - id от 0 до N
                                    user_items=csr_matrix(self.user_item_matrix).tocsr(),   # на вход user-item matrix
                                    N=1, # кол-во рекомендаций
                                    filter_already_liked_items=False,
                                    filter_items=[self.itemid_to_id[999999]] if filter_999999 else None,
                                    recalculate_user=True)

                if len(rec) == 0 or iterations >= max_iterations:
                    success = True
                    break

                rec = self.id_to_itemid[rec[-1][0]]
                if rec not in recs:
                    recs.append(rec)
                    success = True
                else:
                    N_recommends += 1

        #assert len(recs) == N, 'Количество рекомендаций != {}'.format(N)
        if len(recs) != 5:
            print('Error: Рекоммендаций меньше, чем ожидается. user_id:', user_id)
        return recs


from pyramid.view import view_config
from twitter.users.utils import logged_in
from pyramid.httpexceptions import HTTPFound, HTTPNotFound
from ..models import Post

@view_config(route_name='star')
@logged_in
def star(request):
    post = Post.get_by_id(request, request.matchdict['id'])
    if post not in request.user.starred:
        request.user.starred.append(post)
    else:
        raise HTTPNotFound()
    return HTTPFound(request.route_url('board'))

@view_config(route_name='unstar')
@logged_in
def unstar(request):
    post = Post.get_by_id(request, request.matchdict['id'])
    if post in request.user.starred:
        request.user.starred.remove(post)
    else:
        raise HTTPNotFound()
    return HTTPFound(location=request.route_url('favourites'))


from flask import Blueprint
from datetime import date, datetime
import json

import config
from models import ItemSet, Category
from utils import get_warning_date, get_highlight_date
from hardware_manager import hardware_manager

cron_blueprint = Blueprint('cron_blueprint', __name__)

###############################################################################
# CRON ROUTES
###############################################################################
@cron_blueprint.route("/cron/test/", methods=['POST'])
def test_cron():
    hardware_manager.run_test()
    return "All's well!\n"

@cron_blueprint.route("/cron/reset/", methods=['GET', 'POST'])
def reset():
    hardware_manager.reset()
    return 'true'

@cron_blueprint.route("/cron/check-expiry/", methods=['POST'])
def check_expiry():
    report = build_expiry_report()
    write_expiry_report(report)
    hardware_manager.accept_report(report)

    return json.dumps(report) + "\n"

@cron_blueprint.route("/cron/check-longevity/", methods=['POST'])
def check_longevity():
    items = ItemSet.query.all()
    edible_count = 0
    drinkable_count = 0
    consumption_rate = config.FAMILY_MEMBER_COUNT * config.MEALS_PER_DAY_COUNT

    for item in items:
        if item.is_drinkable:
            drinkable_count += item.quantity
        else:
            edible_count += item.quantity

    edible_longevity = edible_count / consumption_rate
    drinkable_longevity = drinkable_count / consumption_rate

    return json.dumps({
        "drinkable_count": drinkable_count,
        "edible_count": edible_count,
        "consumption_rate": consumption_rate,
        "edible_longevity": edible_longevity,
        "drinkable_longevity": drinkable_longevity,
    }) + "\n"

def count_nonexpired(items):
    count = 0
    for item in items:
        if not item.is_ignorable:
            count += 1
    return count

def build_expiry_report():
    expired_items = ItemSet.query.filter(ItemSet.expiration < date.today()).all()
    warning_items = ItemSet.query.filter(ItemSet.expiration < get_warning_date()).filter(ItemSet.expiration >= date.today()).all()
    highlight_items = ItemSet.query.filter(ItemSet.expiration < get_highlight_date()).filter(ItemSet.expiration >= get_warning_date()).all()

    return {
        "expired": count_nonexpired(expired_items),
        "warning": count_nonexpired(warning_items),
        "highlight": count_nonexpired(highlight_items),
    }

def write_expiry_report(report):
    expired_padding = " " * (19 - len(str(report["expired"])))
    warning_padding = " " * (19 - len(str(report["warning"])))
    highlight_padding = " " * (17 - len(str(report["highlight"])))
    report_file = open("report.txt", "w")
    report_file.write("▛▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▜" + "\n")
    report_file.write("▌ " + str(datetime.now()) + "  ▐\n")
    report_file.write("▌ expired: " + str(report["expired"]) + expired_padding + "▐\n")
    report_file.write("▌ warning: " + str(report["warning"]) + warning_padding + "▐\n")
    report_file.write("▌ highlight: " + str(report["highlight"]) + highlight_padding + "▐\n")
    report_file.write("▙▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▟" + "\n")
    report_file.close()


from xivo.rest_api_helpers import APIException

class MeetingGuestSIPTemplateNotFound(APIException):
    def __init__(self, tenant_uuid):
        self.msg = (
            f'Could not find SIP template for meeting guests in tenant {tenant_uuid}'
        )
        details = {
            'tenant_uuid': tenant_uuid,
        }
        super().__init__(
            503,
            self.msg,
            'meeting-guest-sip-template-not-found',
            details=details,
            resource='meeting',
        )

    def __str__(self):
        return self.msg


for p in range(int(input())):
    (input())
    x=[int(i) for i in input().split()]
    e=0
    o=0
    for i in x:
        if i%2==0: e+=i
        else: o+=i
    if e>o:
        print("YES")
    else: print("NO")

Codeforces-Solutions/1807B - Grab the Candies.py


import requests
from datetime import datetime
from operator import methodcaller
from django.http import HttpResponse, HttpResponseRedirect, Http404
from django.contrib import messages
from django.shortcuts import render, redirect
from django.core.exceptions import ValidationError
from django.core.validators import validate_email
from rest_framework.test import APIClient
from .forms import StreamForm, LoginForm, PurgeForm
from .utils import UserDB, TweepyToSQL
from .ds_utils import Summary, embed_html
from .tasks import stream_task
from .keys import (DB, POSTGRES_USER, POSTGRES_PASS, HOST)

def home(request):
	return render(request, 'interface/body_home.html')

def demo(request):
	return render(request, 'interface/demo.html')

def streamer(request):
	if request.method == 'POST':
		form = StreamForm(request.POST)
		if form.is_valid():
			data = form
			topic = data['topic'].value()
			stream_time = int(data['stream_time'].value())
			email = data['email'].value()

			user = UserDB(DB, POSTGRES_USER, POSTGRES_PASS, HOST)
			print('Connected to Postgres: {}'.format(not bool(user.connection.closed)))

			# user.cursor.execute('DELETE FROM interface_user')
			# user.connection.commit()

			try:
				validate_email(email)
			except ValidationError as e:
				raise ValidationError('Email not valid, try again', e)

			try:
				user.to_db(email)
			except Exception as E:
				messages.error(request, 'Email {} already exists in the system, try a different email or request a new session'.format(email))
				return HttpResponseRedirect(request.path_info)

			try:
				user.get_user(email)
				user._dconnect()
				messages.success(request, 'Check your email in {} minutes to access your unique link'.format(stream_time))
				stream_task.delay(email, stream_time, topic)
				return redirect('/streamer')
			except Exception as E:
				user._dconnect()
				messages.error(request, 'Duplicate email found or something went wrong on this end with error: ' + str(E))
				return HttpResponseRedirect(request.path_info)

			# print(json.loads(requests.get('http://127.0.0.1:8000//model/?sentence='+topic).text))
	else:
		form = StreamForm()
	return render(request, 'interface/streamer.html', {'form': form})

def purge(request):
	if request.method == 'POST':
		form = PurgeForm(request.POST)
		if form.is_valid():
			data = form
			email = data['email'].value()
			user = UserDB(DB, POSTGRES_USER, POSTGRES_PASS, HOST)
			print('Connected to Postgres: {}'.format(not bool(user.connection.closed)))

			user.drop_user(email)

			to_sql = TweepyToSQL(DB, POSTGRES_USER, POSTGRES_PASS, HOST, email)
			if to_sql.connection:
				to_sql.connection.rollback()
			to_sql.purge(email)

			return redirect('/')
	else:
		form = StreamForm()
	return render(request, 'interface/purge.html', {'form': form})

def login(request):
	if request.method == 'POST':
		form = LoginForm(request.POST)
		if form.is_valid():
			data = form
			email = data['email'].value()
			uuid_in = data['uuid'].value()
			user = UserDB(DB, POSTGRES_USER, POSTGRES_PASS, HOST)
			print('Connected to Postgres: {}'.format(not bool(user.connection.closed)))

			if user.get_id(email) == []:
				messages.error(request, 'Email not found in database, redirecting to Streamer page')
				print('Authentication failure')
				return redirect('/streamer')
			elif str(user.get_id(email)[0][0]) != str(uuid_in):
				messages.error(request, 'UUID does not match, check email and try again')
				print('Authentication failure')
				return redirect('/login')
			elif str(user.get_id(email)[0][0]) == str(uuid_in):
				print('Authentication success')
				return redirect('dash/{}'.format(email.split('@')[0]))

			user._dconnect()
	else:
		form = LoginForm()
	return render(request, 'interface/login.html', {'form': form})

def dash(request):
	url = request.build_absolute_uri().split('/dash/')[1]
	filename = None

	user = UserDB(DB, POSTGRES_USER, POSTGRES_PASS, HOST)
	all_users = [x[2] for x in user.get_all()]

	spl = [x[0] for x in map(methodcaller("split", "@"), all_users)]
	idx = ''
	for e in spl:
		if e == url:
			idx = spl.index(e)
	if idx == '':
		raise Http404('User not in DB, return to the streamer page to create a new user')
	user._dconnect()

	user = all_users[idx]
	to_sql = TweepyToSQL(DB, POSTGRES_USER, POSTGRES_PASS, HOST, user)
	if to_sql.connection:
		to_sql.connection.rollback()
	filename = user.split('@')[0] + '_' + datetime.now().strftime('%H:%M:%S') + '_tweets.csv'
	to_sql.query_table(filename, user)

	su = Summary(filename=filename)

	reaction_type = ['likes', 'retweets']
	stat = ['mean', 'std', 'min', 'max']
	res = []
	for t in reaction_type:
		for s in stat:
			res.append(round(su.summary_stats(t, s), 2))

	su.wc(); wc = su.wc_dict
	top = list(zip(list(wc.keys())[:10], list(wc.values())[:10]))

	html_l = embed_html(su.get_most_likes())
	html_re = embed_html(su.get_most_re())
	lr = su.likes_retweets_plot()
	wc = su.wc_plot()

	if request.method == 'POST':
		with open(filename) as file:
			response = HttpResponse(file, content_type='text/csv')
			response['Content-Disposition'] = 'attachment; filename={}'.format(filename)
			return response

	args = {'user': url.capitalize(), 'post': filename,
		    'likes_mean': res[0], 'likes_std': res[1], 'likes_min': res[2], 'likes_max': res[3],
		    're_mean': res[4], 're_std': res[5], 're_min': res[6], 're_max': res[7],
			'top': top, 'lr': lr, 'wc': wc,
			'html_l': html_l, 'html_re': html_re}
	return render(request, 'interface/dash.html', args)

def bot(request):
	return render(request, 'interface/bot.html')

def bot_old(request):
	message = ''
	submitted = False
	url = request.build_absolute_uri().split('/old/')[1]

	if request.method == 'POST':
		client = APIClient()
		res = client.get('/model/{}/'.format(url), {'sentence': url.capitalize()}, format='json')

		message = res.content.decode("utf-8")
		submitted = True

	return render(request, 'interface/bot_old.html', {'url': url.capitalize(), 'message': message, 'submitted': submitted})

def bot_new(request):
	return render(request, 'interface/bot_new.html')


from odoo import _, api, fields, models
from datetime import datetime,date

class StockLocation(models.Model):
    _inherit = 'stock.location'

    chemical_catridge_usage = fields.Float('Chemical Catridge Usage')
    cut_off = fields.Date('Cut Off')
    equipment_ids = fields.One2many('maintenance.equipment', 'location_id', string='Equipment')

class ChemicalCatridge(models.Model):
    _name = 'chemical.catridge'
    _description = 'Chemical Catridge'
    _order = 'date desc'

    product_id = fields.Many2one('product.product', string='Product/Chemical')
    warehouse_id = fields.Many2one('stock.location', string='Location',domain=[("usage", "=", "internal")])
    date = fields.Date('Date',default=fields.Date.today)
    stock_awal = fields.Float('Beginning Stock',group_operator=False)
    penerimaan = fields.Float('Receiving',group_operator=False)
    penuangan = fields.Float(string='Pouring',group_operator=False)
    pemakaian = fields.Float(string='Usage',compute='_get_pemakaian',store=True)
    cleaning_basa = fields.Float('Alkali Cleaning Usage')
    adj_over_loss = fields.Float('Adjustment Over/(Loss)',group_operator=False)
    dosing_stroke = fields.Float('Dosing Stoke',group_operator=False)
    dosing_stroke_percent = fields.Float('Dosing Stoke (%)',group_operator=False)
    sisa_stock = fields.Float(compute='_compute_sisa_stock', string='Remaining Stock',store=True,group_operator=False)
    type = fields.Selection([('boo', 'BOO'),('oms', 'OMS')], string='type')
    sisa_tangki = fields.Integer('Remaining Tank',group_operator=False)
    sisa_tangki_kg = fields.Float(compute='_compute_sisa_tangki_kg', string='Remaining Tank (kg)',store=True,group_operator=False)

    @api.onchange('product_id','warehouse_id','date')
    def _onchange_product_id(self):
        for i in self:
            if i.product_id and i.warehouse_id and i.date:
                first_day = i.date.replace(day=1)
                # first_day = datetime.strptime(i.date, '%d %b %Y').replace(day=1)
                sisa_stock = self.env["chemical.catridge"].search([("type", "=", i.type),("product_id", "=", i.product_id.id),("warehouse_id", "=", i.warehouse_id.id),('date','<',i.date )],limit = 1, order = 'date desc')
                if sisa_stock:
                    i.stock_awal = sisa_stock.sisa_stock

    @api.depends('sisa_tangki_kg','date','warehouse_id','product_id')
    def _get_pemakaian(self):
        for i in self:
            if i.sisa_tangki_kg and i.date and i.warehouse_id and i.product_id:
                getyes = self.env["chemical.catridge"].search([("product_id", "=", i.product_id.id),("warehouse_id", "=", i.warehouse_id.id),('date','<',i.date ),('sisa_tangki_kg','>',0.0 )],limit = 1, order = 'date desc')
                if i.penuangan not in [0,False]:
                    i.pemakaian = getyes.sisa_tangki_kg
                else:
                    if getyes:
                        i.pemakaian = getyes.sisa_tangki_kg - i.sisa_tangki_kg
                    else:
                        i.pemakaian = 0
            else:
                i.pemakaian = 0

    @api.depends('sisa_tangki','product_id')
    def _compute_sisa_tangki_kg(self):
        for i in self:
            if i.sisa_tangki and i.warehouse_id:
                i.sisa_tangki_kg = i.sisa_tangki / 100 * i.warehouse_id.chemical_catridge_usage
            else:
                i.sisa_tangki_kg = 0

    @api.depends('stock_awal','penerimaan','penuangan','pemakaian','cleaning_basa','adj_over_loss','dosing_stroke')
    def _compute_sisa_stock(self):
        for i in self:
            i.sisa_stock = i.stock_awal + i.penerimaan - i.penuangan - i.cleaning_basa - i.adj_over_loss - i.dosing_stroke or 0


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Grid search algorithm that takes arrival time picks from S-files
and calculates 2D epicentres as the grid cell with lowest L2 residual.
"""

import datetime
import glob
import math
import numpy as np

def get_arrival_times(S_file):

    '''
    Parse S-file for P wave arrival times
    at recording stations.
    '''

    arrival_time_data = [[], []]
    with open(S_file) as openfile:

        r = -1
        for row in openfile:
            r += 1

            # Split row by spaces

            row_split = row[:-2].split(' ')

            # Get non-empty columns in the row

            cols = []
            for col in row_split:

                if len(col) != 0:

                    cols.append(col)

            # Collect data

            if r == 0:

                # Get date

                year = cols[0]

                if len(cols[1]) > 2:

                    month = cols[1][:-2]
                    day = cols[1][-2:]

                else:

                    month= cols[1]
                    day = cols[2]

                if len(month) < 2:
                    month = '0' + month

                if len(day) < 2:
                    day = '0' + day

                date = datetime.datetime.strptime(year + '-' + month + '-' + day,
                                                  '%Y-%m-%d')

            elif r > 3:

                if len(cols) == 0:

                    continue

                # Get arrival time data (assume no day increments)

                station = cols[0][:-2] # Ignore component details

                seconds = cols[-1]

                if len(cols[-2]) > 2:

                    minutes = cols[-2][-2:]
                    hours = cols[-2][:-2]

                else:

                    minutes = cols[-2]
                    hours = cols[-3]

                total_seconds = (int(hours) * 3600 + int(minutes) * 60 + float(seconds))

                arrival_time_data[0].append(station)
                arrival_time_data[1].append(total_seconds)

    return arrival_time_data

def relative_timing(arrival_time_data):

    '''
    Make arrival time data relative for all stations.
    '''

    reference_station = arrival_time_data[0][0]
    reference_arrival_time = arrival_time_data[1][0]

    relative_arrival_time_data = [[], []]
    for i in range(len(arrival_time_data[0])):

        relative_arrival_time_data[0].append(arrival_time_data[0][i])
        relative_arrival_time_data[1].append(arrival_time_data[1][i] - reference_arrival_time)

    return relative_arrival_time_data, reference_station

def generate_tt(xmin, xmax, ymin, ymax, xstep, ystep, velocity, station_positions):

    '''
    Generate 2D travel time grid.
    '''

    gridx, gridy = np.meshgrid(np.linspace(xmin, xmax, (xmax - xmin) / xstep + 1),
                               np.linspace(ymin, ymax, (ymax - ymin) / ystep + 1))

    # Generate travel times from each station to each grid point

    travel_times = [[[] for i in range(len(gridx))] for j in range(len(gridx[0]))]

    for i in range(len(gridx)):
        for j in range(len(gridx[i])):

            x = gridx[i][j]
            y = gridy[i][j]

            for station_position in station_positions:

                distance = calculate_distance(x, y, station_position[0], station_position[1])

                travel_times[j][i].append(distance / velocity)

    return gridx, gridy, travel_times

def calculate_distance(x1, y1, x2, y2):

    '''
    Calculate distance between two points in a two-dimensional
    cartesian coordinate system.
    '''

    distance = math.sqrt((x2 - x1)**2 + (y2 - y1)**2)

    return distance

def relative_travel_time(travel_times, stations):

    '''
    Make travel times relative to a reference station.
    '''

    relative_travel_times = [[[[] for i in range(len(gridx))] for j in range(len(gridx[0]))] for k in range(len(stations))]

    for k in range(len(stations)):
        for j in range(len(travel_times)):
            for i in range(len(travel_times[j])):
                for l in range(len(travel_times[j][i])):

                    relative_travel_times[k][j][i].append(travel_times[j][i][l] - travel_times[j][i][k])

    return relative_travel_times

def L2_residual(relative_arrival_times, relative_travel_times, stations):

    '''
    Calculate L2 residual for data and model output.
    '''

    L2_residual = 0

    for i in range(len(relative_arrival_times[0])):

        data_station = relative_arrival_times[0][i]

        for j in range(len(relative_travel_times)):

            model_station = stations[j]

            if data_station == model_station:

                # Scale by 1000 before squaring, to avoid square scaling reversals
                # with > 1, < 1 residuals.

                L2_residual += (1000000 * (relative_arrival_times[1][i] - relative_travel_times[j])**2)

    return math.sqrt(L2_residual)

# Get arrival time data

S_file_dir = '/home/sam/PAPERV2_SEISMIC_3COMP_EVENTS/'
S_files = glob.glob(S_file_dir + '*.S*')

events = []
for S_file in S_files:

    events.append(get_arrival_times(S_file))

# Define station names and positions (XYZ)

stations = ['TSNC1', 'TSNL2', 'TSNC3', 'TSNR3', 'TSNR2', 'TSNL3']

station_positions = [[0.00, 0.00],
                     [-0.53, 0.34],
                     [2.12, -0.11],
                     [1.01, -0.66],
                     [-0.77, -0.35],
                     [1.19, 0.43]]

# Set grid size and scale (km)

xmin = -4
xmax = 4
ymin = -3
ymax = 3
xstep = 0.05
ystep = 0.05

# Set grid velocity (km / s)

velocity = 1.65

# Calculate travel time grid and travel times

gridx, gridy, travel_times = generate_tt(xmin, xmax, ymin, ymax, xstep, ystep,
                                         velocity, station_positions)

# Calculate relative travel times for each station

relative_tt = relative_travel_time(travel_times, stations)

# Perform grid search

for e in range(len(events)):

    event = events[e]
    epicentre_idx = [0, 0]
    min_residual = 999999
    L2_residuals = [[[] for i in range(len(gridx))] for j in range(len(gridx[0]))]

    # Make arrival time relative to a given station

    relative_arrival_time_data, reference_station = relative_timing(event)

    # Find index of reference station in station list

    reference_station_idx = stations.index(reference_station)

    # Look through all grid cells

    for i in range(len(gridx)):

        for j in range(len(gridx[i])):

            # Make travel times relative to the reference station for the grid cell

            relative_travel_times = relative_tt[reference_station_idx][j][i]

            # Calculate the L2 residual for the grid cell

            L2_residuals[j][i] = L2_residual(relative_arrival_time_data, relative_travel_times, stations)

            if L2_residuals[j][i] < min_residual:

                min_residual = L2_residuals[j][i]
                epicentre_idx = [j, i]
                epicentre = [gridx[i][j], gridy[i][j]]

    print(S_files[e].split('/')[-1] + ',' + str(min_residual) + ',' + str(epicentre_idx) + ',' + str(epicentre))
    np.save(S_files[e].split('/')[-1] + '_residuals.npy', L2_residuals)