# Author: bbrighttaer
# Project: jova
# Date: 7/22/19
# Time: 10:47 AM
# File: train_joint.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import argparse
import copy
import json
import logging
import random
import time
from datetime import datetime as dt

import numpy as np
import torch
import torch.nn as nn
import torch.optim.lr_scheduler as sch
from soek import *
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

import jova.metrics as mt
import jova.utils.io
from jova import cuda
from jova.data import batch_collator, get_data, load_proteins, DtiDataset
from jova.data.data import featurize_datasets
from jova.metrics import compute_model_performance
from jova.nn.layers import GraphConvLayer, GraphPool, GraphGather2D, Unsqueeze, ElementwiseBatchNorm
from jova.nn.models import GraphConvSequential, create_fcn_layers, NwayForward, JointAttention, \
    WeaveModel, Prot2Vec, ProteinRNN, ProteinCNN, ProteinCNN2D, GraphNeuralNet2D
from jova.trans import undo_transforms
from jova.utils import Trainer
from jova.utils.args import FcnArgs, WeaveGatherArgs, WeaveLayerArgs
from jova.utils.io import load_pickle
from jova.utils.math import ExpAverage, Count
from jova.utils.tb import TBMeanTracker
from jova.utils.train_helpers import count_parameters, GradStats, FrozenModels, ViewsReg

logging.basicConfig(format='%(asctime)s %(message)s', level=logging.ERROR, filename='error.log')

currentDT = dt.now()
date_label = currentDT.strftime("%Y_%m_%d__%H_%M_%S")

seeds = [1, 8, 64]

check_data = False

torch.cuda.set_device(2)

views_reg = ViewsReg()

def create_ecfp_net(hparams):
    model = nn.Sequential(Unsqueeze(dim=0))
    return model

def create_prot_net(hparams, model_type, protein_profile, protein_embeddings, frozen_models_hook=None):
    # assert protein_profile is not None, "Protein profile has to be supplied"
    # assert protein_embeddings is not None, "Pre-trained protein embeddings are required"

    # model_type = hparams["prot"]["model_type"].lower()
    assert (model_type in ViewsReg.all_prot_views), "Valid protein types: {}".format(
        str(ViewsReg.all_prot_views))
    if model_type == "rnn":
        # pt_embeddings = create_torch_embeddings(frozen_models_hook, protein_embeddings)
        model = nn.Sequential(Prot2Vec(protein_profile=protein_profile,
                                       vocab_size=hparams["prot"]["vocab_size"],
                                       embedding_dim=hparams["prot"]["embedding_dim"],
                                       batch_first=False),
                              ProteinRNN(in_dim=hparams["prot"]["embedding_dim"] * hparams["prot"]["window"],
                                         hidden_dim=hparams["prot"]["rnn_hidden_state_dim"],
                                         dropout=hparams["dprob"],
                                         batch_first=False))
    elif model_type == "p2v":
        # pt_embeddings = create_torch_embeddings(frozen_models_hook, protein_embeddings)
        model = Prot2Vec(protein_profile=protein_profile,
                         vocab_size=hparams["prot"]["vocab_size"],
                         embedding_dim=hparams["prot"]["embedding_dim"],
                         batch_first=False)
    elif model_type == "pcnn":
        # pt_embeddings = create_torch_embeddings(frozen_models_hook, protein_embeddings)
        model = nn.Sequential(Prot2Vec(protein_profile=protein_profile,
                                       vocab_size=hparams["prot"]["vocab_size"],
                                       embedding_dim=hparams["prot"]["embedding_dim"],
                                       batch_first=False),
                              ProteinCNN(dim=hparams["prot"]["embedding_dim"],
                                         window=hparams["prot"]["window"],
                                         activation="relu",
                                         pooling_dim=0,
                                         num_layers=hparams["prot"]["pcnn_num_layers"]),
                              Unsqueeze(dim=0))
    elif model_type == "pcnn2d":
        # pt_embeddings = create_torch_embeddings(frozen_models_hook, protein_embeddings)
        model = nn.Sequential(Prot2Vec(protein_profile=protein_profile,
                                       vocab_size=hparams["prot"]["vocab_size"],
                                       embedding_dim=hparams["prot"]["embedding_dim"],
                                       batch_first=False),
                              ProteinCNN2D(dim=hparams["prot"]["embedding_dim"],
                                           window=hparams["prot"]["window"],
                                           num_layers=hparams["prot"]["pcnn_num_layers"]))
    elif model_type == "psc":
        model = nn.Sequential(Unsqueeze(dim=0))
    return model

def create_weave_net(hparams):
    weave_args = (
        WeaveLayerArgs(n_atom_input_feat=75,
                       n_pair_input_feat=14,
                       n_atom_output_feat=50,
                       # n_atom_output_feat=hparams["weave"]["dim"],
                       n_pair_output_feat=50,
                       n_hidden_AA=50,
                       n_hidden_PA=50,
                       n_hidden_AP=50,
                       n_hidden_PP=50,
                       update_pair=hparams["weave"]["update_pairs"],
                       activation='relu',
                       batch_norm=True,
                       dropout=hparams["dprob"]
                       ),
        WeaveLayerArgs(n_atom_input_feat=50,
                       n_pair_input_feat=14,
                       n_atom_output_feat=hparams["weave"]["dim"],
                       n_pair_output_feat=50,
                       n_hidden_AA=50,
                       n_hidden_PA=50,
                       n_hidden_AP=50,
                       n_hidden_PP=50,
                       update_pair=hparams["weave"]["update_pairs"],
                       batch_norm=True,
                       dropout=hparams["dprob"],
                       activation='relu'),
    )
    wg_args = WeaveGatherArgs(conv_out_depth=hparams["weave"]["dim"], gaussian_expand=True,
                              n_depth=hparams["weave"]["dim"])
    weave_model = WeaveModel(weave_args, weave_gath_arg=wg_args, weave_type='2D')
    model = nn.Sequential(weave_model)
    return model

def create_gconv_net(hparams):
    dim = hparams["gconv"]["dim"]
    gconv_model = GraphConvSequential(GraphConvLayer(in_dim=75, out_dim=64),
                                      nn.BatchNorm1d(64),
                                      nn.ReLU(),
                                      GraphPool(),

                                      GraphConvLayer(in_dim=64, out_dim=64),
                                      nn.BatchNorm1d(64),
                                      nn.ReLU(),
                                      GraphPool(),

                                      nn.Linear(in_features=64, out_features=dim),
                                      nn.BatchNorm1d(dim),
                                      nn.ReLU(),
                                      nn.Dropout(hparams["dprob"]),
                                      GraphGather2D(activation='nonsat'))

    model = nn.Sequential(gconv_model)
    return model

def create_gnn_net(hparams):
    dim = hparams["gnn"]["dim"]
    gnn_model = GraphNeuralNet2D(num_fingerprints=hparams["gnn"]["fingerprint_size"], embedding_dim=dim,
                                 num_layers=hparams["gnn"]["num_layers"])
    return nn.Sequential(gnn_model,
                         nn.Linear(dim, dim),
                         ElementwiseBatchNorm(dim),
                         nn.ReLU(),
                         nn.Dropout(hparams["dprob"]))

def create_integrated_net(hparams, protein_profile, protein_embeddings):
    # Convenient way of keeping track of models to be frozen during (or at the initial stages) training.
    frozen_models = FrozenModels()

    # N-way forward propagation
    views = {}
    seg_dims = []

    net_creators = {'ecfp8': create_ecfp_net,
                    'weave': create_weave_net,
                    'gconv': create_gconv_net,
                    'gnn': create_gnn_net}

    # compound models
    for cview in views_reg.c_views:
        views[cview] = net_creators[cview](hparams)
        seg_dims.append(hparams[cview]["dim"])

    # protein models
    for m_type in hparams["prot"]["model_types"]:
        views[m_type] = create_prot_net(hparams, m_type, protein_profile, protein_embeddings, frozen_models)
        if m_type == "psc":
            seg_dims.append(hparams["prot"]["psc_dim"])
        elif m_type == "rnn":
            seg_dims.append(hparams["prot"]["rnn_hidden_state_dim"])
        elif m_type == "p2v":
            seg_dims.append(hparams["prot"]["embedding_dim"] * hparams["prot"]["window"])
        elif m_type in ["pcnn", "pcnn2d"]:
            seg_dims.append(hparams["prot"]["embedding_dim"])

    layers = [NwayForward(models=views.values())]

    layers.append(JointAttention(d_dims=seg_dims, latent_dim=hparams["latent_dim"], num_heads=hparams["attn_heads"],
                                 num_layers=hparams["attn_layers"], dprob=hparams["dprob"]))

    p = len(views) * hparams["latent_dim"]
    for dim in hparams["lin_dims"]:
        layers.append(nn.Linear(p, dim))
        layers.append(nn.BatchNorm1d(dim))
        layers.append(nn.ReLU())
        layers.append(nn.Dropout(hparams["dprob"]))
        p = dim

    # Output layer
    layers.append(nn.Linear(in_features=p, out_features=hparams["output_dim"]))

    # Build model
    model = nn.Sequential(*layers)

    return model, frozen_models

def create_discriminator_net(hparams):
    fcn_args = []
    p = hparams["neigh_dist"]
    layers = hparams["disc_hdims"]
    if not isinstance(layers, list):
        layers = [layers]
    for dim in layers:
        conf = FcnArgs(in_features=p,
                       out_features=dim,
                       activation='relu',
                       batch_norm=True,
                       dropout=hparams["dprob"])
        fcn_args.append(conf)
        p = dim
    fcn_args.append(FcnArgs(in_features=p, out_features=1, activation="sigmoid"))
    layers = create_fcn_layers(fcn_args)
    model = nn.Sequential(*layers)
    return model

class JovaGAN(Trainer):

    @staticmethod
    def initialize(hparams, train_dataset, val_dataset, test_dataset, protein_profile, protein_embeddings,
                   cuda_devices=None, mode="regression"):

        # create networks
        generator, frozen_models = create_integrated_net(hparams, protein_profile, protein_embeddings)
        discriminator = create_discriminator_net(hparams)
        print("Number of trainable parameters: generator={}, discriminator={}".format(count_parameters(generator),
                                                                                count_parameters(discriminator)))

        try:
            if cuda:
                generator = generator.cuda()
                discriminator = discriminator.cuda()
        except:  # todo: Remove this try-catch after hyperparameter search
            generator = nn.Sequential(nn.Linear(10, 10))
            discriminator = nn.Sequential(nn.Linear(10, 10))

        # data loaders
        train_data_loader = DataLoader(dataset=train_dataset,
                                       batch_size=hparams["tr_batch_size"],
                                       shuffle=True,
                                       collate_fn=lambda x: x)
        val_data_loader = DataLoader(dataset=val_dataset,
                                     batch_size=hparams["val_batch_size"],
                                     shuffle=False,
                                     collate_fn=lambda x: x)
        test_data_loader = None
        if test_dataset is not None:
            test_data_loader = DataLoader(dataset=test_dataset,
                                          batch_size=hparams["test_batch_size"],
                                          shuffle=False,
                                          collate_fn=lambda x: x)

        # filter optimizer arguments
        optimizer_disc = optimizer_gen = None
        for suffix in ["_gen", "_disc"]:
            key = "optimizer{}".format(suffix)

            # optimizer configuration
            optimizer = {
                "adadelta": torch.optim.Adadelta,
                "adagrad": torch.optim.Adagrad,
                "adam": torch.optim.Adam,
                "adamax": torch.optim.Adamax,
                "asgd": torch.optim.ASGD,
                "rmsprop": torch.optim.RMSprop,
                "Rprop": torch.optim.Rprop,
                "sgd": torch.optim.SGD,
            }.get(hparams[key].lower(), None)
            assert optimizer is not None, "{} optimizer could not be found"

            optim_kwargs = dict()
            optim_key = hparams[key]
            for k, v in hparams.items():
                if "optimizer{}__".format(suffix) in k:
                    attribute_tup = k.split("__")
                    if optim_key == attribute_tup[1] or attribute_tup[1] == "global":
                        optim_kwargs[attribute_tup[2]] = v
            if suffix == "_gen":
                optimizer_gen = optimizer(generator.parameters(), **optim_kwargs)
            else:
                optimizer_disc = optimizer(discriminator.parameters(), **optim_kwargs)

        # metrics
        metrics = [mt.Metric(mt.rms_score, np.nanmean),
                   mt.Metric(mt.concordance_index, np.nanmean),
                   mt.Metric(mt.pearson_r2_score, np.nanmean)]
        return (generator, discriminator), (optimizer_gen, optimizer_disc), \
               {"train": train_data_loader,
                "val": val_data_loader,
                "test": test_data_loader}, metrics, hparams["prot"]["model_types"], \
               hparams["weighted_loss"], hparams["neigh_dist"], frozen_models

    @staticmethod
    def data_provider(fold, flags, data_dict):
        if not flags['cv']:
            print("Training scheme: train, validation" + (", test split" if flags['test'] else " split"))
            train_dataset = DtiDataset(x_s=[data[1][0].X for data in data_dict.values()],
                                       y_s=[data[1][0].y for data in data_dict.values()],
                                       w_s=[data[1][0].w for data in data_dict.values()])
            valid_dataset = DtiDataset(x_s=[data[1][1].X for data in data_dict.values()],
                                       y_s=[data[1][1].y for data in data_dict.values()],
                                       w_s=[data[1][1].w for data in data_dict.values()])
            test_dataset = None
            if flags['test']:
                test_dataset = DtiDataset(x_s=[data[1][2].X for data in data_dict.values()],
                                          y_s=[data[1][2].y for data in data_dict.values()],
                                          w_s=[data[1][2].w for data in data_dict.values()])
            data = {"train": train_dataset, "val": valid_dataset, "test": test_dataset}
        else:
            train_dataset = DtiDataset(x_s=[data[1][fold][0].X for data in data_dict.values()],
                                       y_s=[data[1][fold][0].y for data in data_dict.values()],
                                       w_s=[data[1][fold][0].w for data in data_dict.values()])
            valid_dataset = DtiDataset(x_s=[data[1][fold][1].X for data in data_dict.values()],
                                       y_s=[data[1][fold][1].y for data in data_dict.values()],
                                       w_s=[data[1][fold][1].w for data in data_dict.values()])
            test_dataset = DtiDataset(x_s=[data[1][fold][2].X for data in data_dict.values()],
                                      y_s=[data[1][fold][2].y for data in data_dict.values()],
                                      w_s=[data[1][fold][2].w for data in data_dict.values()])
            data = {"train": train_dataset, "val": valid_dataset, "test": test_dataset}
        return data

    @staticmethod
    def evaluate(eval_dict, y, y_pred, w, metrics, tasks, transformers):
        eval_dict.update(compute_model_performance(metrics, y_pred.cpu().detach().numpy(), y, w, transformers,
                                                   tasks=tasks))
        # scoring
        rms = np.nanmean(eval_dict["nanmean-rms_score"])
        ci = np.nanmean(eval_dict["nanmean-concordance_index"])
        r2 = np.nanmean(eval_dict["nanmean-pearson_r2_score"])
        score = np.nanmean([ci, r2]) - rms
        return score

    @staticmethod
    def train(models, optimizers, data_loaders, metrics, prot_model_types, weighted_loss, neigh_dist,
              frozen_models, transformers_dict, prot_desc_dict, tasks, n_iters=5000, sim_data_node=None,
              epoch_ckpt=(2, 1.0), tb_writer=None, is_hsearch=False):
        generator, discriminator = models
        optimizer_gen, optimizer_disc = optimizers

        tb_writer = tb_writer()
        start = time.time()
        best_model_wts = generator.state_dict()
        best_score = -10000
        best_epoch = -1
        terminate_training = False
        e_avg = ExpAverage(.01)
        n_epochs = n_iters // len(data_loaders["train"])
        grad_stats = GradStats(nn.ModuleList(models), beta=0.)

        # learning rate decay schedulers
        scheduler_gen = sch.StepLR(optimizer_gen, step_size=500, gamma=0.01)
        scheduler_disc = sch.StepLR(optimizer_disc, step_size=500, gamma=0.01)

        # pred_loss functions
        prediction_criterion = nn.MSELoss()
        adversarial_loss = nn.BCELoss()

        # sub-nodes of sim data resource
        loss_lst = []
        train_loss_node = DataNode(label="training_loss", data=loss_lst)
        metrics_dict = {}
        metrics_node = DataNode(label="validation_metrics", data=metrics_dict)
        val_scores_lst = []
        val_scores_node = DataNode(label="validation_score", data=val_scores_lst)
        train_scores_lst = []
        train_scores_node = DataNode(label="train_score", data=train_scores_lst)
        gen_loss_lst = []
        gen_loss_node = DataNode(label="generator_loss", data=gen_loss_lst)
        dis_loss_lst = []
        dis_loss_node = DataNode(label="discriminator_loss", data=dis_loss_lst)

        # add sim data nodes to parent node
        if sim_data_node:
            sim_data_node.data = [train_loss_node, train_scores_node, metrics_node, val_scores_node, gen_loss_node,
                                  dis_loss_node]

        try:
            # Main training loop
            tb_idx = {'train': Count(), 'val': Count(), 'test': Count()}
            for epoch in range(n_epochs):
                if terminate_training:
                    print("Terminating training...")
                    break
                for phase in ["train", "val" if is_hsearch else "test"]:
                    # ensure these models are frozen at all times
                    for m in frozen_models:
                        m.eval()

                    if phase == "train":
                        print("Training....")
                        # Training mode
                        generator.train()
                        discriminator.train()
                    else:
                        print("Validation...")
                        # Evaluation mode
                        generator.eval()

                    data_size = 0.
                    epoch_losses = []
                    epoch_scores = []

                    # Iterate through mini-batches
                    i = 0
                    with TBMeanTracker(tb_writer, 10) as tracker:
                        with grad_stats:
                            for batch in tqdm(data_loaders[phase]):
                                spec = {v: True for v in views_reg.c_views}
                                batch_size, data = batch_collator(batch, prot_desc_dict, spec, cuda_prot=True)

                                # organize the data for each view.
                                Xs = {}
                                Ys = {}
                                Ws = {}
                                for view_name in data:
                                    view_data = data[view_name]
                                    if view_name == "gconv":
                                        x = ((view_data[0][0], batch_size), view_data[0][1], view_data[0][2])
                                        Xs["gconv"] = x
                                    else:
                                        Xs[view_name] = view_data[0]
                                    Ys[view_name] = np.array([k for k in view_data[1]], dtype=np.float)
                                    Ws[view_name] = np.array([k for k in view_data[2]], dtype=np.float)

                                optimizer_gen.zero_grad()
                                optimizer_disc.zero_grad()

                                # forward propagation
                                # track history if only in train
                                with torch.set_grad_enabled(phase == "train"):
                                    Ys = {k: Ys[k] for k in Ys}
                                    # Ensure corresponding pairs
                                    for j in range(1, len(Ys.values())):
                                        assert (list(Ys.values())[j - 1] == list(Ys.values())[j]).all()

                                    y = Ys[list(Xs.keys())[0]]
                                    w = Ws[list(Xs.keys())[0]]

                                    # protein data in batch
                                    protein_xs = []
                                    for m_type in prot_model_types:
                                        if m_type in views_reg.embedding_based_views:
                                            protein_xs.append(Xs[list(Xs.keys())[0]][2])
                                        else:  # m_type == "psc":
                                            protein_xs.append(Xs[list(Xs.keys())[0]][1])

                                    # compound data in batch
                                    X = [Xs[v][0] for v in views_reg.c_views]

                                    # merge compound and protein list
                                    X = X + protein_xs

                                    outputs = generator(X)
                                    target = torch.from_numpy(y).float()
                                    weights = torch.from_numpy(w).float()
                                    valid = torch.ones_like(target).float()
                                    fake = torch.zeros_like(target).float()
                                    if cuda:
                                        target = target.cuda()
                                        weights = weights.cuda()
                                        valid = valid.cuda()
                                        fake = fake.cuda()
                                    outputs = outputs * weights
                                    pred_loss = prediction_criterion(outputs, target)

                                # fail fast
                                if str(pred_loss.item()) == "nan":
                                    terminate_training = True
                                    break

                                eval_dict = {}
                                score = JovaGAN.evaluate(eval_dict, y, outputs, w, metrics, tasks,
                                                         transformers_dict[list(Xs.keys())[0]])
                                # TBoard info
                                tracker.track("%s/loss" % phase, pred_loss.item(), tb_idx[phase].IncAndGet())
                                tracker.track("%s/score" % phase, score, tb_idx[phase].i)
                                for k in eval_dict:
                                    tracker.track('{}/{}'.format(phase, k), eval_dict[k], tb_idx[phase].i)

                                if phase == "train":
                                    # GAN stuff
                                    f_xx, f_yy = torch.meshgrid(outputs.squeeze(), outputs.squeeze())
                                    predicted_diffs = torch.abs(f_xx - f_yy).sort(dim=1)[0][:, : neigh_dist]
                                    r_xx, r_yy = torch.meshgrid(target.squeeze(), target.squeeze())
                                    real_diffs = torch.abs(r_xx - r_yy).sort(dim=1)[0][:, :neigh_dist]

                                    # generator
                                    gen_loss = adversarial_loss(discriminator(predicted_diffs), valid)
                                    gen_loss_lst.append(gen_loss.item())
                                    loss = pred_loss + weighted_loss * gen_loss
                                    tracker.track("gan/gen_loss", gen_loss.item(), tb_idx[phase].i)
                                    tracker.track("train/comp_loss", loss.item(), tb_idx[phase].i)
                                    loss.backward()
                                    optimizer_gen.step()

                                    # discriminator
                                    true_loss = adversarial_loss(discriminator(real_diffs), valid)
                                    fake_loss = adversarial_loss(discriminator(predicted_diffs.detach()), fake)
                                    discriminator_loss = (true_loss + fake_loss) / 2.
                                    tracker.track("gan/disc_loss", discriminator_loss.item(), tb_idx[phase].i)
                                    dis_loss_lst.append(discriminator_loss.item())
                                    discriminator_loss.backward()
                                    optimizer_disc.step()

                                    # for epoch stats
                                    epoch_losses.append(pred_loss.item())

                                    # for sim data resource
                                    loss_lst.append(pred_loss.item())

                                    print("\tEpoch={}/{}, batch={}/{}, pred_loss={:.4f}, D loss={:.4f}, "
                                          "G loss={:.4f}".format(epoch + 1, n_epochs, i + 1, len(data_loaders[phase]),
                                                                 pred_loss.item(), discriminator_loss, gen_loss))

                                    # monitor train data score
                                    train_scores_lst.append(score)
                                else:
                                    # for epoch stats
                                    epoch_scores.append(score)

                                    # for sim data resource
                                    val_scores_lst.append(score)
                                    for m in eval_dict:
                                        if m in metrics_dict:
                                            metrics_dict[m].append(eval_dict[m])
                                        else:
                                            metrics_dict[m] = [eval_dict[m]]

                                        print("\nEpoch={}/{}, batch={}/{}, "
                                              "evaluation results= {}, score={}".format(epoch + 1, n_epochs, i + 1,
                                                                                len(data_loaders[phase]),
                                                                                eval_dict, score))

                                i += 1
                                data_size += batch_size
                    # End of mini=batch iterations.
                    if phase == "train":
                        ep_loss = np.nanmean(epoch_losses)
                        e_avg.update(ep_loss)
                        if epoch % (epoch_ckpt[0] - 1) == 0 and epoch > 0:
                            if e_avg.value > epoch_ckpt[1]:
                                terminate_training = True
                        print("\nPhase: {}, avg task pred_loss={:.4f}, ".format(phase, np.nanmean(epoch_losses)))

                        # Adjust the learning rate.
                        scheduler_gen.step()
                        scheduler_disc.step()
                    else:
                        mean_score = np.mean(epoch_scores)
                        if best_score < mean_score:
                            best_score = mean_score
                            best_model_wts = copy.deepcopy(generator.state_dict())
                            best_epoch = epoch
        except RuntimeError as e:
            print(str(e))

        duration = time.time() - start
        print('\nModel training duration: {:.0f}m {:.0f}s'.format(duration // 60, duration % 60))
        try:
            generator.load_state_dict(best_model_wts)
        except RuntimeError as e:
            print(str(e))
        return {'model': generator, 'score': best_score, 'epoch': best_epoch}

    @staticmethod
    def evaluate_model(model, model_dir, model_name, data_loaders, metrics, prot_model_types,
                       transformers_dict, prot_desc_dict, tasks, sim_data_node=None):
        # load saved model and put in evaluation mode
        model.load_state_dict(jova.utils.io.load_model(model_dir, model_name))
        model.eval()

        print("Model evaluation...")
        start = time.time()
        n_epochs = 1

        # sub-nodes of sim data resource
        # loss_lst = []
        # train_loss_node = DataNode(label="training_loss", data=loss_lst)
        metrics_dict = {}
        metrics_node = DataNode(label="validation_metrics", data=metrics_dict)
        scores_lst = []
        scores_node = DataNode(label="validation_score", data=scores_lst)
        predicted_vals = []
        true_vals = []
        model_preds_node = DataNode(label="model_predictions", data={"y": true_vals,
                                                                     "y_pred": predicted_vals})

        # add sim data nodes to parent node
        if sim_data_node:
            sim_data_node.data = [metrics_node, scores_node, model_preds_node]

        # Main evaluation loop
        for epoch in range(n_epochs):

            for phase in ["test"]:
                # Iterate through mini-batches
                i = 0
                for batch in tqdm(data_loaders[phase]):
                    spec = {v: True for v in views_reg.c_views}
                    batch_size, data = batch_collator(batch, prot_desc_dict, spec, cuda_prot=True)

                    # organize the data for each view.
                    Xs = {}
                    Ys = {}
                    Ws = {}
                    for view_name in data:
                        view_data = data[view_name]
                        if view_name == "gconv":
                            x = ((view_data[0][0], batch_size), view_data[0][1], view_data[0][2])
                            Xs["gconv"] = x
                        else:
                            Xs[view_name] = view_data[0]
                        Ys[view_name] = np.array([k for k in view_data[1]], dtype=np.float)
                        Ws[view_name] = np.array([k for k in view_data[2]], dtype=np.float)

                    # forward propagation
                    with torch.set_grad_enabled(False):
                        Ys = {k: Ys[k].astype(np.float) for k in Ys}
                        # Ensure corresponding pairs
                        for j in range(1, len(Ys.values())):
                            assert (list(Ys.values())[j - 1] == list(Ys.values())[j]).all()

                        y_true = Ys[list(Xs.keys())[0]]
                        w = Ws[list(Xs.keys())[0]]

                        # protein data in batch
                        protein_xs = []
                        for m_type in prot_model_types:
                            if m_type in views_reg.embedding_based_views:
                                protein_xs.append(Xs[list(Xs.keys())[0]][2])
                            else:  # m_type == "psc":
                                protein_xs.append(Xs[list(Xs.keys())[0]][1])

                        # compound data in batch
                        X = [Xs[v][0] for v in views_reg.c_views]

                        # merge compound and protein list
                        X = X + protein_xs
                        y_predicted = model(X)

                        # apply transformers
                        predicted_vals.extend(undo_transforms(y_predicted.cpu().detach().numpy(),
                                                              transformers_dict["gconv"]).squeeze().tolist())
                        true_vals.extend(
                            undo_transforms(y_true, transformers_dict["gconv"]).astype(np.float).squeeze().tolist())

                    eval_dict = {}
                    score = JovaGAN.evaluate(eval_dict, y_true, y_predicted, w, metrics, tasks,
                                             transformers_dict["gconv"])

                    # for sim data resource
                    scores_lst.append(score)
                    for m in eval_dict:
                        if m in metrics_dict:
                            metrics_dict[m].append(eval_dict[m])
                        else:
                            metrics_dict[m] = [eval_dict[m]]

                    print("\nEpoch={}/{}, batch={}/{}, "
                          "evaluation results= {}, score={}".format(epoch + 1, n_epochs, i + 1,
                                                                    len(data_loaders[phase]),
                                                                    eval_dict, score))

                    i += 1
                # End of mini=batch iterations.

        duration = time.time() - start
        print('\nModel evaluation duration: {:.0f}m {:.0f}s'.format(duration // 60, duration % 60))

def main(flags):
    print("JOVA sims:", flags.jova)

    # Load protein data
    prot_desc_dict, prot_seq_dict = load_proteins(flags['prot_desc_path'])
    prot_profile = load_pickle(file_name=flags.prot_profile)
    prot_vocab = load_pickle(file_name=flags.prot_vocab)
    pretrained_embeddings = None  # load_numpy_array(flags.protein_embeddings)
    flags["prot_vocab_size"] = len(prot_vocab)

    # Ensures all possible compound data in each seed is present. Helps with maintaining the random number generator
    # state during splitting to avoid sample mismatch across views.
    featurize_datasets(flags.jova, views_reg.feat_dict, flags, prot_seq_dict, seeds)
    for v_arg in flags.jova:
        views_reg.parse_views(v_arg)
        comp_lbl = views_reg.c_views
        comp_lbl = '_'.join(comp_lbl)
        flags["prot_model_types"] = views_reg.p_views
        sim_label = "integrated_view_attn_gan_" + ('_'.join(flags["prot_model_types"])) + '_' + comp_lbl
        print("CUDA={}, view={}".format(cuda, sim_label))

        # Simulation data resource tree
        split_label = "warm" if flags["split_warm"] else "cold_target" if flags["cold_target"] else "cold_drug" if \
            flags["cold_drug"] else "None"
        dataset_lbl = flags["dataset_name"]
        # node_label = "{}_{}_{}_{}_{}".format(dataset_lbl, sim_label, split_label, "eval" if flags["eval"] else "train",
        #                                      date_label)
        node_label = json.dumps({'model_family': 'jova-gan',
                                 'dataset': dataset_lbl,
                                 'cviews': '-'.join(views_reg.c_views),
                                 'pviews': '-'.join(views_reg.p_views),
                                 'split': split_label,
                                 'mode': "eval" if flags["eval"] else "train",
                                 'seeds': '-'.join([str(s) for s in seeds]),
                                 'date': date_label})
        sim_data = DataNode(label=node_label)
        nodes_list = []
        sim_data.data = nodes_list

        num_cuda_dvcs = torch.cuda.device_count()
        cuda_devices = None if num_cuda_dvcs == 1 else [i for i in range(1, num_cuda_dvcs)]

        # For searching over multiple seeds
        hparam_search = None

        for seed in seeds:
            summary_writer_creator = lambda: SummaryWriter(log_dir="tb_jovagan"
                                                                   "/{}_{}_{}/".format(sim_label, seed,
                                                                                dt.now().strftime(
                                                                                "%Y_%m_%d__%H_%M_%S")))

            # for data collection of this round of simulation.
            data_node = DataNode(label="seed_%d" % seed)
            nodes_list.append(data_node)

            random.seed(seed)
            np.random.seed(seed)
            torch.manual_seed(seed)
            torch.cuda.manual_seed_all(seed)

            # load data
            print('-------------------------------------')
            print('Running on dataset: %s' % dataset_lbl)
            print('-------------------------------------')

            data_dict = dict()
            transformers_dict = dict()

            # Data
            for cview in views_reg.c_views:
                data_dict[cview] = get_data(views_reg.feat_dict[cview], flags, prot_sequences=prot_seq_dict, seed=seed)
                transformers_dict[cview] = data_dict[cview][2]
                flags["gnn_fingerprint"] = data_dict[cview][3]

            tasks = data_dict[list(data_dict.keys())[0]][0]
            # multi-task or single task is determined by the number of tasks w.r.t. the dataset loaded
            flags["tasks"] = tasks

            trainer = JovaGAN()

            if flags["cv"]:
                k = flags["fold_num"]
                print("{}, {}-Prot: Training scheme: {}-fold cross-validation".format(tasks, sim_label, k))
            else:
                k = 1
                print("{}, {}-Prot: Training scheme: train, validation".format(tasks, sim_label)
                      + (", test split" if flags['test'] else " split"))

            if check_data:
                verify_multiview_data(data_dict)
            else:
                if flags["hparam_search"]:
                    print("Hyperparameter search enabled: {}".format(flags["hparam_search_alg"]))

                    # arguments to callables
                    extra_init_args = {"mode": "regression",
                                       "cuda_devices": cuda_devices,
                                       "protein_profile": prot_profile,
                                       "protein_embeddings": pretrained_embeddings}
                    extra_data_args = {"flags": flags,
                                       "data_dict": data_dict}
                    total_iterations = 3000
                    extra_train_args = {"transformers_dict": transformers_dict,
                                        "prot_desc_dict": prot_desc_dict,
                                        "tasks": tasks,
                                        "is_hsearch": True,
                                        "n_iters": total_iterations,
                                        "tb_writer": summary_writer_creator}

                    hparams_conf = get_hparam_config(flags)
                    if hparam_search is None:
                        search_alg = {"random_search": RandomSearchCV,
                                      "bayopt_search": BayesianOptSearchCV}.get(flags["hparam_search_alg"],
                                                                                BayesianOptSearchCV)
                        min_opt = "gp"
                        hparam_search = search_alg(hparam_config=hparams_conf,
                                                   num_folds=k,
                                                   initializer=trainer.initialize,
                                                   data_provider=trainer.data_provider,
                                                   train_fn=trainer.train,
                                                   save_model_fn=jova.utils.io.save_model,
                                                   init_args=extra_init_args,
                                                   data_args=extra_data_args,
                                                   train_args=extra_train_args,
                                                   data_node=data_node,
                                                   split_label=split_label,
                                                   sim_label=sim_label,
                                                   minimizer=min_opt,
                                                   dataset_label=dataset_lbl,
                                                   results_file="{}_{}_dti_{}_{}-{}.csv".format(
                                                       flags["hparam_search_alg"], sim_label, date_label,
                                                       total_iterations, min_opt))

                    stats = hparam_search.fit(model_dir="models", model_name="".join(tasks), max_iter=20, seed=seed)
                    print(stats)
                    print("Best params = {}".format(stats.best()))
                else:
                    invoke_train(trainer, tasks, data_dict, transformers_dict, flags, prot_desc_dict, data_node,
                                 sim_label,
                                 prot_profile, pretrained_embeddings, summary_writer_creator)

        # save simulation data resource tree to file.
        sim_data.to_json(path="./analysis/")

def invoke_train(trainer, tasks, data_dict, transformers_dict, flags, prot_desc_dict, data_node, view, protein_profile,
                 protein_embeddings, tb_writer):
    hyper_params = default_hparams_bopt(flags)
    # Initialize the model and other related entities for training.
    if flags["cv"]:
        folds_data = []
        data_node.data = folds_data
        data_node.label = data_node.label + "cv"
        for k in range(flags["fold_num"]):
            k_node = DataNode(label="fold-%d" % k)
            folds_data.append(k_node)
            start_fold(k_node, data_dict, flags, hyper_params, prot_desc_dict, tasks, trainer,
                       transformers_dict, view, protein_profile, protein_embeddings, tb_writer, k)
    else:
        start_fold(data_node, data_dict, flags, hyper_params, prot_desc_dict, tasks, trainer,
                   transformers_dict, view, protein_profile, protein_embeddings, tb_writer)

def start_fold(sim_data_node, data_dict, flags, hyper_params, prot_desc_dict, tasks, trainer,
               transformers_dict, view, protein_profile, protein_embeddings, tb_writer=None, k=None):
    data = trainer.data_provider(k, flags, data_dict)
    model, optimizer, data_loaders, metrics, prot_model_types, weighted_loss, \
    neigh_dist, frozen_models = trainer.initialize(hparams=hyper_params,
                                                   train_dataset=data["train"],
                                                   val_dataset=data["val"],
                                                   test_dataset=data["test"],
                                                   protein_profile=protein_profile,
                                                   protein_embeddings=protein_embeddings)
    if flags["eval"]:
        trainer.evaluate_model(model[0], flags["model_dir"], flags["eval_model_name"],
                               data_loaders, metrics, transformers_dict,
                               prot_desc_dict, tasks, sim_data_node=sim_data_node)
    else:
        # Train the model
        results = trainer.train(model, optimizer, data_loaders, metrics, prot_model_types,
                                weighted_loss, neigh_dist, frozen_models, transformers_dict, prot_desc_dict,
                                tasks, max_iter=10000, sim_data_node=sim_data_node, tb_writer=tb_writer)
        model, score, epoch = results['model'], results['score'], results['epoch']
        # Save the model.
        split_label = "warm" if flags["split_warm"] else "cold_target" if flags["cold_target"] else "cold_drug" if \
            flags["cold_drug"] else "None"
        jova.utils.io.save_model(model, flags["model_dir"],
                                 "{}_{}_{}_{}_{}_{:.4f}".format(flags["dataset"], view, flags["model_name"],
                                                                split_label, epoch,
                                                                score))

def default_hparams_rand(flags):
    return {
        "prot_dim": 8421,
        "fp_dim": 1024,
        "gconv_dim": 128,
        "hdims": [3795, 2248, 2769, 2117],

        # weight initialization
        "kaiming_constant": 5,

        # dropout regs
        "dprob": 0.0739227,

        "tr_batch_size": 256,
        "val_batch_size": 512,
        "test_batch_size": 512,

        # optimizer params
        "optimizer": "rmsprop",
        "optimizer__sgd__weight_decay": 1e-4,
        "optimizer__sgd__nesterov": True,
        "optimizer__sgd__momentum": 0.9,
        "optimizer__sgd__lr": 1e-3,

        "optimizer__adam__weight_decay": 1e-4,
        "optimizer__adam__lr": 1e-3,

        "optimizer__rmsprop__lr": 0.000235395,
        "optimizer__rmsprop__weight_decay": 0.000146688,
        "optimizer__rmsprop__momentum": 0.00622082,
        "optimizer__rmsprop__centered": False
    }

def default_hparams_bopt(flags):
    """
    protein model types:
    --------------------------------------------------------------------------------------------
    short name  | full name
    ------------|-------------------------------------------------------------------------------
    psc         | Protein Sequence Composition
    ------------|-------------------------------------------------------------------------------
    p2v         | Protein to Vector / Embeddings using n-gram amino acid 'words'.
    ------------|-------------------------------------------------------------------------------
    rnn         | Uses embeddings and an RNN variant (e.g. LSTM) to learn protein features.
    ------------|-------------------------------------------------------------------------------
    pcnn        | Protein CNN: https://academic.oup.com/bioinformatics/article/35/2/309/5050020
                | The final output is a 1D vector for each protein in a batch.
    ------------|-------------------------------------------------------------------------------
    pcnn2d      | A variant of PCNN that returns a 2D tensor for each protein in a batch.
    ------------|-------------------------------------------------------------------------------
    NOTE: All protein models, except 'psc' and 'p2v', use embeddings from the :class:Prot2Vec module.
    """
    return {
        "prot_vocab_size": flags["prot_vocab_size"],
        "attn_heads": 4,
        "attn_layers": 1,
        "lin_dims": [1424, 64],
        "output_dim": len(flags.tasks),
        "latent_dim": 256,
        "disc_hdims": [365, ],

        # weight initialization
        "kaiming_constant": 5,

        "weighted_loss": 0.1,

        # dropout
        "dprob": 0.1,

        "neigh_dist": 56,

        "tr_batch_size": 256,
        "val_batch_size": 128,
        "test_batch_size": 128,

        # optimizer params
        "optimizer_gen": "adamax",
        "optimizer_gen__global__weight_decay": 0.0009,
        "optimizer_gen__global__lr": 0.0006,
        "optimizer_disc": "sgd",
        "optimizer_disc__global__weight_decay": 1.0,
        "optimizer_disc__global__lr": 0.0001,
        "prot": {
            "model_types": flags["prot_model_types"],
            "vocab_size": flags["prot_vocab_size"],
            "window": 11,
            "pcnn_num_layers": 2,
            "embedding_dim": 10,
            "psc_dim": 8421,
            "rnn_hidden_state_dim": 10
        },
        "weave": {
            "dim": 50,
            "update_pairs": False,
        },
        "gconv": {
            "dim": 512,
        },
        "ecfp8": {
            "dim": 1024,
        },
        "gnn": {
            "fingerprint_size": len(flags["gnn_fingerprint"]) if flags["gnn_fingerprint"] else 0,
            "num_layers": 3,
            "dim": 100,
        }
    }

def get_hparam_config(flags):
    config = {
        "attn_heads": CategoricalParam([1, 2, 4, 8, 16]),
        "attn_layers": DiscreteParam(min=1, max=2),
        "lin_dims": DiscreteParam(min=64, max=1048, size=DiscreteParam(min=1, max=2)),
        "latent_dim": ConstantParam(256),
        "output_dim": ConstantParam(len(flags.tasks)),
        "disc_hdims": DiscreteParam(min=100, max=1048, size=2),

        # weight initialization
        "kaiming_constant": ConstantParam(5),

        "weighted_loss": RealParam(0.1, max=0.5),

        # dropout
        "dprob": RealParam(0.1, max=0.5),
        "neigh_dist": DiscreteParam(min=4, max=64),

        "tr_batch_size": ConstantParam(256),
        "val_batch_size": ConstantParam(128),
        "test_batch_size": ConstantParam(128),

        # optimizer params
        "optimizer_gen": CategoricalParam(choices=["sgd", "adam", "adadelta", "adagrad", "adamax", "rmsprop"]),
        "optimizer_gen__global__weight_decay": LogRealParam(),
        "optimizer_gen__global__lr": LogRealParam(),
        "optimizer_disc": CategoricalParam(choices=["sgd", "adam", "adadelta", "adagrad", "adamax", "rmsprop"]),
        "optimizer_disc__global__weight_decay": LogRealParam(),
        "optimizer_disc__global__lr": LogRealParam(),

        "prot": DictParam({
            "model_types": ConstantParam(flags["prot_model_types"]),
            "vocab_size": ConstantParam(flags["prot_vocab_size"]),
            "window": ConstantParam(11),
            "psc_dim": ConstantParam(8421)
        }),
        "ecfp8": DictParam({"dim": ConstantParam(1024)}),
        "weave": DictParam({}),
        "gconv": DictParam({}),
        "gnn": DictParam({})
    }

    # prot additions
    for pv in views_reg.p_views:
        # common props
        if pv in views_reg.embedding_based_views:
            config["prot"]["embedding_dim"] = DiscreteParam(min=5, max=50)

        # view-specific
        if pv in views_reg.pcnn_views:
            config["prot"]["pcnn_num_layers"] = DiscreteParam(min=1, max=4)
        elif pv == 'rnn':
            config["prot"]["rnn_hidden_state_dim"] = DiscreteParam(min=5, max=50)

    # comp additions
    for cv in views_reg.c_views:
        # common prop
        if cv in views_reg.graph_based_views:
            config[cv]["dim"] = DiscreteParam(min=64, max=512)

        # view-specific
        if cv == "gnn":
            config[cv]["fingerprint_size"] = ConstantParam(len(flags["gnn_fingerprint"]))
            config[cv]["num_layers"] = DiscreteParam(1, 4)
        elif cv == "weave":
            config[cv]["update_pairs"] = ConstantParam(False)
    return config

def verify_multiview_data(data_dict):
    ecfp8_data = data_dict["ecfp8"][1][0]
    gconv_data = data_dict["gconv"][1][0]
    corr = []
    for i in range(100):
        print("-" * 100)
        ecfp8 = "mol={}, prot={}, y={}".format(ecfp8_data.X[i][0].smiles, ecfp8_data.X[i][1].get_name(),
                                               ecfp8_data.y[i])
        print("ecfp8:", ecfp8)
        gconv = "mol={}, prot={}, y={}".format(gconv_data.X[i][0].smiles, gconv_data.X[i][1].get_name(),
                                               gconv_data.y[i])
        print("gconv:", gconv)
        print('#' * 100)
        corr.append(ecfp8 == gconv)
    print(corr)

class Flags(object):
    # enables using either object referencing or dict indexing to retrieve user passed arguments of flag objects.
    def __getitem__(self, item):
        return self.__dict__[item]

    def __setitem__(self, key, value):
        setattr(self, key, value)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="DTI with jova model training.")

    parser.add_argument("--dataset_name",
                        type=str,
                        default="davis",
                        help="Dataset name.")
    parser.add_argument("--dataset_file",
                        type=str,
                        help="Dataset file.")

    # Either CV or standard train-val(-test) split.
    scheme = parser.add_mutually_exclusive_group()
    scheme.add_argument("--fold_num",
                        default=-1,
                        type=int,
                        choices=range(3, 11),
                        help="Number of folds for cross-validation")
    scheme.add_argument("--test",
                        action="store_true",
                        help="Whether a test set should be included in the data split")

    parser.add_argument("--splitting_alg",
                        choices=["random", "scaffold", "butina", "index", "task"],
                        default="random",
                        type=str,
                        help="Data splitting algorithm to use.")
    parser.add_argument('--filter_threshold',
                        type=int,
                        default=6,
                        help='Threshold such that entities with observations no more than it would be filtered out.'
                        )
    parser.add_argument('--cold_drug',
                        default=False,
                        help='Flag of whether the split will leave "cold" drugs in the test data.',
                        action='store_true'
                        )
    parser.add_argument('--cold_target',
                        default=False,
                        help='Flag of whether the split will leave "cold" targets in the test data.',
                        action='store_true'
                        )
    parser.add_argument('--cold_drug_cluster',
                        default=False,
                        help='Flag of whether the split will leave "cold cluster" drugs in the test data.',
                        action='store_true'
                        )
    parser.add_argument('--predict_cold',
                        default=False,
                        help='Flag of whether the split will leave "cold" entities in the test data.',
                        action='store_true')
    parser.add_argument('--split_warm',
                        default=False,
                        help='Flag of whether the split will not leave "cold" entities in the test data.',
                        action='store_true'
                        )
    parser.add_argument('--model_dir',
                        type=str,
                        default='./model_dir',
                        help='Directory to store the log files in the training process.'
                        )
    parser.add_argument('--model_name',
                        type=str,
                        default='model-{}'.format(date_label),
                        help='Directory to store the log files in the training process.'
                        )
    parser.add_argument('--prot_desc_path',
                        action='append',
                        help='A list containing paths to protein descriptors.'
                        )
    parser.add_argument('--prot_profile',
                        type=str,
                        help='A resource for retrieving embedding indexing profile of proteins.'
                        )
    parser.add_argument('--prot_vocab',
                        type=str,
                        help='A resource containing all N-gram segments/words constructed from the protein sequences.'
                        )
    parser.add_argument('--prot_embeddings',
                        type=str,
                        dest="protein_embeddings",
                        help='Numpy array file containing the pretrained protein "words" embeddings'
                        )
    parser.add_argument('--no_reload',
                        action="store_false",
                        dest='reload',
                        help='Whether datasets will be reloaded from existing ones or newly constructed.'
                        )
    # parser.add_argument('--data_dir',
    #                     type=str,
    #                     default='../../data/',
    #                     help='Root folder of data (Davis, KIBA, Metz) folders.')
    parser.add_argument("--hparam_search",
                        action="store_true",
                        help="If true, hyperparameter searching would be performed.")
    parser.add_argument("--hparam_search_alg",
                        type=str,
                        default="bayopt_search",
                        help="Hyperparameter search algorithm to use. One of [bayopt_search, random_search]")
    parser.add_argument("--eval",
                        action="store_true",
                        help="If true, a saved model is loaded and evaluated using CV")
    parser.add_argument("--eval_model_name",
                        default=None,
                        type=str,
                        help="The filename of the model to be loaded from the directory specified in --model_dir")
    parser.add_argument("--jova",
                        action="append",
                        type=str,
                        help="A combination of compound and protein views for simulation. "
                             "The format is: comp1-compN;prot1-protN\nor instance, for a combination "
                             "of the PSC and RNN protein views with ECFP8 and GraphConv views of a compound, "
                             "the argument would be:\tecfp8-gconv;psc-rnn\n"
                             "Available compound views:[ecfp8,weave,gconv,gnn]\n Available protien views:"
                             "[psc,rnn,p2v,pcnn, pcnn2d]")

    args = parser.parse_args()
    flags = Flags()
    args_dict = args.__dict__
    for arg in args_dict:
        setattr(flags, arg, args_dict[arg])
    setattr(flags, "cv", True if flags.fold_num > 2 else False)
    main(flags)


from django.shortcuts import render, redirect
from django.core.urlresolvers import reverse
from django.utils.translation import activate
from product.models import Category
from news.models import News
from django.http import HttpResponse
from general.models import ImageSlide
import re

def index(request):
    slides = ImageSlide.objects.filter(active=True)
    return render(request, 'index.html', {'slides': slides})

def home(request):
    news = News.objects.order_by('-created_at')[:4]
    categories = Category.objects.all()
    slides = ImageSlide.objects.filter(active=True)
    return render(request, 'home.html', {'news': news, 'categories': categories, 'slides': slides})

def locale(request):
    path = request.GET.get('path', '')
    locale = request.GET.get('locale', 'zh-tw')
    if path:
        if '/zh-cn' in path:
            match = re.match('/zh-cn(.+)', path)
            path = match.groups()[0]
        if '/zh-tw' in path:
            match = re.match('/zh-tw(.+)', path)
            path = match.groups()[0]
        if '/en' in path:
            match = re.match('/en(.+)', path)
            path = match.groups()[0]
        return redirect('/%s%s' % (locale, path))
    else:
        activate(locale)
        return redirect(reverse('home'))


"""
Вам дан массив в описанном формате, без пробелов и переводов строки. Ваша задача —отформатировать его следующим образом:

Каждая строка и каждая фигурная скобка должны находиться в отдельной строке.
Запятые находятся на той же строке, что и элемент (строка или фигурная скобка) перед ней.
После открывающейся фигурной скобки отступ увеличивается на 2 пробела.
Перед закрывающейся фигурной скобки отступ уменьшается на 2 пробела.
Изучите примеры, чтобы лучше понять, как должен выглядеть отформатированный массив.

Входные данные
Первая строка ввода содержит строку s, представляющую собой описание массива. Длина строки не превышает 1000 символов.

Выходные данные
Выведите отформатированный массив.

Примеры
входные данныеСкопировать
{odin,dva,tri}
выходные данныеСкопировать
{
  odin,
  dva,
  tri
}
входные данныеСкопировать
{}
выходные данныеСкопировать
{
}
входные данныеСкопировать
{stroka}
выходные данныеСкопировать
{
  stroka
}
входные данныеСкопировать
{a,b,{c,d},e,{}}
выходные данныеСкопировать
{
  a,
  b,
  {
    c,
    d
  },
  e,
  {
  }
}
"""
s = '{a,b,{{{e,k}},{c,d,e},{}}}'
spaces = 0
r = ''
eol = False
for c in s:
    if c == '{':
        r += '{\n'
        eol = True
        spaces += 2
        r += spaces * ' '
    elif c == '}':
        r += '\n'
        spaces -= 2
        r += spaces * ' '
        r += '}'
    elif c == ',':
        r += ',\n'
        r += spaces * ' '
    else:
        r += c
pr = r.split('\n')
for i in pr:
    if i.lstrip() == '':
        pr.remove(i)
print('\n'.join(pr))


import PySimpleGUI as sg
from PySimpleGUI.PySimpleGUI import DEFAULT_BUTTON_COLOR
from src.ventanas import menu_inicio as inicio
from src.auxiliares import clases as cla
from src.auxiliares import funciones_generales as funcs
from src.auxiliares import registro_eventos as reg

def evento_partida(tiempo, filas, columnas, coincidencias, nick,
                   nivel, dificultad, nombre_evento, estado=None):
    """ Se reciben los parametros relevantes para guardar
    un registro con informacion de un evento
    al comenzar una nueva partida """

    cant_palabras = (filas * columnas) / coincidencias
    edad = funcs.get_info_usuario_actual()['Edad']
    genero = funcs.get_info_usuario_actual()['Género']

    evento_inicio = reg.Evento()
    evento_inicio.tiempo = tiempo
    evento_inicio.cant_palabras = cant_palabras
    evento_inicio.nick = nick
    evento_inicio.genero = genero
    evento_inicio.edad = edad
    evento_inicio.nivel = nivel
    evento_inicio.dificultad = dificultad
    if nombre_evento == 'fin_partida':
        evento_inicio.nombre_evento = 'fin_partida'
        evento_inicio.estado = estado
        evento_inicio.partida = reg.numero_partida()
    else:
        evento_inicio.nombre_evento = 'inicio_partida'
        evento_inicio.partida = reg.numero_partida() + 1

    evento_inicio.cargar_evento()

def main(siguiente=False, otros_puntos=0):
    """ Inicializa la ventana de juego. Se llaman a diversos métodos
    para lectura de json de usuario actual, fecha y hora, entre otros,
    y así se recopilan los datos necesarios para instanciar una partida,
    con sus fichas, y cargarlas en una ventana implementada en este módulo,
    donde se desarrolla la jugada en sí """

    (dificultad, nivel, nivel_inicial, nivel_estatico, coincidencias,
     filas, columnas, tipo_ficha, tiempo, nick) = cla.Partida.mise_en_place()

    partida = cla.Partida(dificultad, nivel, nivel_inicial, nivel_estatico, coincidencias,
                          filas, columnas, tipo_ficha, tiempo, nick, siguiente)

    verde = '#10ae23'
    verde_claro = '#11da29'
    amarillo = '#f1e910'
    gris = '#343434'
    rojo = '#d91644'
    blanco = '#FFFFFF'
    defecto = DEFAULT_BUTTON_COLOR
    pausa_momentanea = False

    if otros_puntos > 0 and not nivel_estatico:
        partida.inc_puntos(otros_puntos)
    color_dific = [rojo if dificultad == 'Difícil'
                   else amarillo if dificultad == 'Normal'
                   else verde_claro]

    top_windows = [sg.Column([
                    [sg.Text(nick, text_color='#4185e1', background_color='#173155',
                             font=('Arial', 14), size=(12,1), justification='c')],
                    [sg.Text(dificultad + '  -  ' + str(nivel), text_color=color_dific,
                             background_color='#173155',
                             font=('Arial', 14), size=(12,1), justification='c')]]),
                    sg.Column([[sg.Text('')]]),
                    sg.Column([
                    [sg.Text(('Tiempo: ' + partida.get_tiempo_string()),
                            justification='left', size=(13,1),
                            key='reloj', font=('Verdana', 12)),
                    sg.Text('Puntos: ' + str(partida.get_puntos()),
                            key='puntos', size=(13,1),
                            justification='left', font=('Verdana', 12))],
                    [sg.Text('Encontrada: ', key='encontrada', size=(26,1),
                            justification='left', font=('Verdana', 12))]])
                  ]
    bottom_window = [
                    [sg.Button('Comenzar', key='comenzar', size=(14,1),
                                font=('Times New Roman', 20), mouseover_colors=(blanco, verde)),
                    sg.Button('Pausar', key='pausar', size=(14,1), mouseover_colors=(blanco, amarillo),
                              font=('Times New Roman', 20), visible=False)],
                    [sg.Button('Reiniciar', key = 'reiniciar', mouseover_colors=(blanco, amarillo),
                               size= (11,1), font=('', 12)),
                    sg.Button('Salir', key='salir', mouseover_colors=(blanco, rojo),
                              size= (11,1), font=('', 12))]
                    ]
    layout = [top_windows, partida.get_fichas(), bottom_window]

    window = sg.Window('MemPyG28',
                       layout, element_justification='c', margins=(30,5))
    partida.window = window
    mensajes = funcs.get_info_usuario_actual()['Configuración']['Mensajes']
    comenzado = False

    while True:
        if siguiente and not partida._nivel_estatico:
            sg.popup('Siguiente nivel', auto_close=True, auto_close_duration=2,
                     line_width=50, no_titlebar=True, background_color=gris,
                     text_color=amarillo, font=('Verdana', 14), button_type=5,
                     button_color=gris, custom_text='Vamos')
            siguiente = False

        event, values = window.read(timeout=10)

        if event == sg.WINDOW_CLOSED:
            break
        if event == 'salir':
            if comenzado and not partida._paused:
                funcs.pausar_juego(partida)
                pausa_momentanea = True
            salir = sg.popup('Salir al menu inicio?',
                             custom_text=('Salir', 'Cancelar'), no_titlebar=True,
                             background_color='Black', text_color='White')
            if salir == 'Salir':
                if comenzado:
                    evento_partida(partida.get_tiempo_string(), filas, columnas, coincidencias,
                                   nick, nivel, dificultad, 'fin_partida', 'abandonada')
                window.close()
                inicio.main()
            elif comenzado and pausa_momentanea:
                funcs.despausar_juego(partida)
                pausa_momentanea = False

        elif event == 'reiniciar':
            if comenzado and not partida._paused:
                funcs.pausar_juego(partida)
                pausa_momentanea = True
            reiniciar = sg.Popup('Si reiniciás la partida\n'
                        + 'vas a perder los puntos\n' + '¿Está seguro?',
                        custom_text=('Reiniciar', 'Cancelar'), no_titlebar=True,
                        background_color='Black', text_color='White')
            if reiniciar == 'Reiniciar':
                partida.reiniciar_juego()
            elif comenzado and pausa_momentanea:
                funcs.despausar_juego(partida)
                pausa_momentanea = False

        if not comenzado:
            if event == 'comenzar':
                evento_partida(partida.get_tiempo_string(), filas, columnas, coincidencias,
                               nick, nivel, dificultad, 'inicio_partida')
                comenzado = True
                window['comenzar'].Update(visible = False)
                window['pausar'].Update(visible = True)
                partida.resetear_tiempo()
        else:
            if (not partida.paused()):
                if event == 'pausar':
                    funcs.pausar_juego(partida)
                    defecto = sg.theme_button_color()
                    window['pausar'].Update(text='Despausar', button_color=rojo)

                event = event.replace('__TIMEOUT__', '') if type(event) == str else event

                partida.check_tiempo_turno()
                partida.actualizar_tiempo()
                if partida.tiempo_restante() < 0:
                    mensajes = funcs.get_info_usuario_actual()['Configuración']['Mensajes']
                    partida.terminar_juego(mensajes['Derrota'], False)
                    evento_partida(partida.get_tiempo_string(), filas, columnas, coincidencias,
                                      nick, nivel, dificultad, 'fin_partida', 'timeout')
                window.FindElement('reloj').Update('Tiempo: ' + partida.get_tiempo_string())

                if partida._tiempo_final and partida.tiempo_restante()<=10:
                    partida._tiempo_final=False
                    window.FindElement('reloj').Update(font=('Verdana', 12),
                                                       text_color= rojo)

                if type(event) == tuple:
                    y, x = event
                    partida.progreso_partida(y, x)
                    partida.actualizar_puntos()
                if len(partida.get_encontradas()) > 0:
                    ultima = partida.get_encontradas()[-1]
                    if (window['encontrada'].DisplayText != ('Encontrada: ' + ultima)):
                        window.FindElement('encontrada').Update('Encontrada: ' + ultima)
            else:
                if event == 'pausar':
                    funcs.despausar_juego(partida)
                    window['pausar'].Update(text='Pausar', button_color=defecto)

    window.close()

if __name__ == '__main__':

    main()


from database.Database import Database
from instagram_util.InstaUtil import InstaUtil

PATH_TEMPLATE = "{}{}/"

dbAPI = Database()

for project in dbAPI.get_projects():

    post = dbAPI.get_next_insta_post(project.project_id)
    insta_util = InstaUtil(project.insta_config)
    path = PATH_TEMPLATE.format(project.reddit_config["path"], post.post_id )
    try:
        insta_util.publish_post(path)
        dbAPI.record_upload_insta(post.post_id, project.project_id)
    except Exception as e:
        print(e)


from __future__ import absolute_import

import six

from datetime import datetime

from sentry.models import EventAttachment, File
from sentry.testutils import APITestCase

class EventAttachmentsTest(APITestCase):
    def test_simple(self):
        self.login_as(user=self.user)

        project = self.create_project()

        release = self.create_release(project, self.user)

        group = self.create_group(project=project, first_release=release)

        event1 = self.create_event(
            event_id='a',
            group=group,
            datetime=datetime(2016, 8, 13, 3, 8, 25),
            tags={'sentry:release': release.version}
        )
        event2 = self.create_event(
            event_id='b',
            group=group,
            datetime=datetime(2016, 8, 13, 3, 8, 25),
            tags={'sentry:release': release.version}
        )

        attachment1 = EventAttachment.objects.create(
            event_id=event1.event_id,
            group_id=event1.group_id,
            project_id=event1.project_id,
            file=File.objects.create(
                name='hello.png',
                type='image/png',
            ),
            name='hello.png',
        )

        EventAttachment.objects.create(
            event_id=event2.event_id,
            group_id=event2.group_id,
            project_id=event2.project_id,
            file=File.objects.create(
                name='hello.png',
                type='image/png',
            ),
            name='hello.png',
        )

        path = '/api/0/projects/{}/{}/events/{}/attachments/'.format(
            event1.project.organization.slug,
            event1.project.slug,
            event1.id,
        )

        with self.feature('organizations:event-attachments'):
            response = self.client.get(path)

        assert response.status_code == 200, response.content
        assert len(response.data) == 1
        assert response.data[0]['id'] == six.text_type(attachment1.id)


from demos.eaglesong.common import *
from demos.eaglesong.example_14_task_1 import StatusMenu
from kaia.infra.tasks import TaskProcessor, SqlSubprocTaskProcessor, SubprocessConfig, TaskResult
from kaia.eaglesong.amenities.dispatcher import Dispatcher
from kaia.eaglesong.amenities.task_filter import CreateTaskAndWait, TaskFilter

class TaskBot(Dispatcher):
    def __init__(self, processor: TaskProcessor):
        super(TaskBot, self).__init__()
        self.add_skill('talk', PushdownFilter(Routine.ensure(self.talk)))
        self.add_skill('status', PushdownFilter(Routine.ensure(self.status)))
        self.add_skill('fallback', PushdownFilter(Routine.ensure(self.fallback)))
        self.processor = processor

    def talk(self, context: BotContext):
        yield 'Enter number N to start a task that takes N seconds. At any time, use /status to control the task.'
        while True:
            yield Listen.for_types(str, int)
            ticks = context.input
            yield CreateTaskAndWait(dict(ticks=ticks))
            response = context.input #type: TaskResult
            if response.success:
                yield response.result
            else:
                yield str(response.error)

    def fallback(self, context: BotContext):
        yield 'Oops'

    def status(self, context):
        tasks = self.processor.get_active_tasks()
        if len(tasks)==0:
            yield 'No active tasks'
        else:
            yield StatusMenu(self.processor, tasks[0])

    def dispatch(self, context: BotContext):
        for name, skill in self.active_skills.items():
            if skill.expect(context.input):
                return name
        if IBotInput.is_content_input(context.input):
            if context.input == '/status':
                if 'status' in self.active_skills:
                    del self.active_skills['status']
                return 'status'
            if 'talk' not in self.active_skills:
                return 'talk'
            return 'fallback'

bot = Bot(
    "task2",
    proc_factory=lambda proc: Routine.interpretable(
        TaskBot(proc),
        PushdownFilter,
        TaskFilter.with_processor(proc)
    ),
    timer=True
)

if __name__ == '__main__':
    bot.processor = SqlSubprocTaskProcessor(SubprocessConfig(
        'kaia.infra.tasks:Waiting',
        [1]
    ))
    bot.processor.activate()
    run(bot)


import sys

opt = open(sys.argv[1],"r",encoding="latin1")

lines = opt.read().splitlines()

total_count = len(lines)
tp = [0,0]
classified_count = [0,0]
actual_count = [0,0]

for line in lines:
    x = line.split("/")
    label = x[0].strip().lower()
    actual = x[-2].strip().lower()

    if label == "ham":
        classified_count[0] += 1
    else:
        classified_count[1] += 1

    if actual == "ham":
        actual_count[0] += 1
    elif actual == "spam":
        actual_count[1] += 1

    if label == actual:
        if label =="ham":
            tp[0] += 1
        else:
            tp[1] += 1

precision = []
if classified_count[0] !=0 :
    precision.append(tp[0]/classified_count[0])
else:
    precision.append(1)

if classified_count[1] !=0 :
    precision.append(tp[1]/classified_count[1])
else:
    precision.append(1)

recall = []
if actual_count[0] != 0:
    recall.append(tp[0] / actual_count[0])
else:
    recall.append(1)

if actual_count[1] != 0:
    recall.append(tp[1] / actual_count[1])
else:
    recall.append(1)

f_score = []
f_score.append((2*precision[0]*recall[0])/(precision[0]+recall[0]))
f_score.append((2*precision[1]*recall[1])/(precision[1]+recall[1]))

print(precision)
print(recall)
print(f_score)


# This Python file uses the following encoding: utf-8
from PySide6 import QtWidgets, QtCharts, QtGui, QtCore

class GraphsWidget(QtWidgets.QWidget):
    def __init__(self, rows, columns):
        super(GraphsWidget, self).__init__()

        self.graphGrid = QtWidgets.QGridLayout()
        self.charts = {}
        self.rows = rows
        self.columns = columns
        self.setLayout(self.graphGrid)

    def init_chart(self, chart_title):
        chart_view = QtCharts.QChartView()
        chart_view.setRenderHint(QtGui.QPainter.Antialiasing)
        chart_view.chart().legend().setVisible(False)
        chart_view.chart().setTitle(chart_title)
        chart_view.chart().setAnimationOptions(QtCharts.QChart.AllAnimations)
        return chart_view

    def add_graph(self, graph_name, X_list, Y_list):
        self.charts[graph_name] = self.init_chart(graph_name)
        current_index = list(self.charts.keys()).index(graph_name)
        current_row = current_index // self.rows
        current_column = current_index % self.columns
        self.graphGrid.addWidget(self.charts[graph_name], current_row, current_column)

        series = QtCharts.QLineSeries()
        for x, y in zip(X_list, Y_list):
            series.append(x, y)
        self.charts[graph_name].chart().addSeries(series)

        x_axis, y_axis = QtCharts.QValueAxis(), QtCharts.QValueAxis()
        x_axis.setTitleText("Time, sec")
        y_axis.setTitleText(graph_name)
        self.charts[graph_name].chart().createDefaultAxes()
        self.charts[graph_name].chart().setAxisX(x_axis, series)
        self.charts[graph_name].chart().setAxisY(y_axis, series)


# import re
# def count_words(text, words):
#     count_result = 0
#     for word in words:
#         if len(re.findall(word,text.lower())) != 0:
#             count_result += 1
#     return count_result

def count_words(text, words):
    for w in words:
        print(w)
    return sum(w in text.lower() for w in words)

if __name__ == '__main__':
    #These "asserts" using only for self-checking and not necessary for auto-testing
    assert count_words("How aresjfhdskfhskd you?", {"how", "are", "you", "hello"}) == 3, "Example"
    assert count_words("Bananas, give me bananas!!!", {"banana", "bananas"}) == 2, "BANANAS!"
    assert count_words("Lorem ipsum dolor sit amet, consectetuer adipiscing elit.",
                       {"sum", "hamlet", "infinity", "anything"}) == 1, "Weird text"
    print("Coding complete? Click 'Check' to review your tests and earn cool rewards!")


import argparse
import logging
import os

import numpy as np
import math, copy, time
import torch
from torch import nn
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data.sampler import RandomSampler
from tqdm import tqdm
from IPython import embed
from itertools import chain
import torch.nn.functional as functional
import utils
from utils import EarlyStopping
import PDTransformer as transformer
from evaluate import *
from dataloader import *

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from torch.cuda.amp import autocast , GradScaler
import random
from opt import OpenAIAdam

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
logger = logging.getLogger('PDTrans.Train')

parser = argparse.ArgumentParser()
parser.add_argument('--dataset', default='elect', help='Name of the dataset')
parser.add_argument('--data-folder', default='PDTrans/data', help='Parent dir of the dataset')
parser.add_argument('--model-name', default='PDTrans/output_elect', help='Directory containing params.json')
parser.add_argument('--relative-metrics', action='store_true', help='Whether to normalize the metrics by label scales')
parser.add_argument('--save-best', action='store_true', help='Whether to save best ND to param_search.txt')
parser.add_argument('--restore-file', default=None, help='Optional, name of the file in --model_dir containing weights to reload before training')

parser.add_argument('--beta', type=float, default=10, help='hyperparameter of loss function')
parser.add_argument('--gamma', type=float, default=10, help='hyperparameter of loss function')
parser.add_argument('--lr', type=float, default=0.0002, help='learning rate')

def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True

def train(model,autoencoder,
        train_loader,
        optimizer,
        scaler,
        params: utils.Params,
        epoch) -> None:
    '''Train the model on one epoch by batches.
    Args:
        model: (torch.nn.Module) the neural network
        optimizer: (torch.optim) optimizer for parameters of model
        train_loader: load train data and labels
        test_loader: load test data and labels
        params: (Params) hyperparameters
        epoch: (int) the current training epoch
    '''
    model.train()
    autoencoder.train()
    loss_epoch = np.zeros(len(train_loader))
    train_loss = []

    for i, (train_batch, idx, labels_batch) in enumerate(tqdm(train_loader)):
        batch_size = train_batch.shape[0]
        train_batch = train_batch.to(torch.float32).to(params.device)
        labels_batch = labels_batch.to(torch.float32).to(params.device)
        idx = idx.unsqueeze(-1).to(params.device)

        mu=None; logvar=None
        distribution_mu, distribution_sigma = model.forward(train_batch, idx, mu, logvar)

        Recon_input = train_batch[:,1:params.predict_start+1,0]
        recon_x, recon_y, recon_z, mu, logvar = autoencoder(Recon_input,distribution_mu)
        loss_re = args.beta*transformer.kl_loss( mu, logvar) + args.gamma*transformer.loss_fn(recon_x, distribution_sigma, labels_batch, params.predict_start)
        # loss_r.append(loss_re.item())

        optimizer.zero_grad()
        loss =  loss_re + transformer.loss_fn(distribution_mu, distribution_sigma, labels_batch, params.predict_start)

        cuda_exist = torch.cuda.is_available()
        if cuda_exist:
            with autocast():
                scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        else:
            loss.backward()
            optimizer.step()

        # g_loss = loss.item() / params.predict_steps
        g_loss = loss.item() / params.train_window
        loss_epoch[i] = g_loss
        train_loss.append(loss.item())

        if i % 500 == 0:
            logger.info("R_loss: {} ; loss: {}".format(loss_re, loss))

    return np.average(train_loss),loss_epoch

def train_and_evaluate(model,autoencoder,
                       train_loader,valid_loader,test_loader: DataLoader,
                       optimizer,scheduler,
                       scaler,
                       params: utils.Params,
                       restore_file: str = None) -> None:

    '''Train the model and evaluate every epoch.
    Args:
        model: (torch.nn.Module) the Deep AR model
        train_loader: load train data and labels
        optimizer: (torch.optim) optimizer for parameters of model
        params: (Params) hyperparameters
        restore_file: (string) optional- name of file to restore from (without its extension .pth.tar)
    '''
    early_stopping = EarlyStopping(patience=100, verbose=True)

    # reload weights from restore_file if specified
    if restore_file is not None:
        restore_path = os.path.join(params.model_dir, restore_file + '.pth.tar')
        logger.info('Restoring parameters from {}'.format(restore_path))
        utils.load_checkpoint(restore_path, model, optimizer)

    logger.info('begin training and evaluation')
    best_valid_q50 = float('inf')
    best_valid_q90 = float('inf')
    best_MAPE = float('inf')
    train_len = len(train_loader)
    valid_len = len(valid_loader)

    q50_valid = np.zeros(params.num_epochs)
    q90_valid = np.zeros(params.num_epochs)
    MAPE_valid = np.zeros(params.num_epochs)

    loss_train_epoch = np.zeros((train_len * params.num_epochs))
    loss_valid = np.zeros((valid_len * params.num_epochs))

    valid_loss = []
    logger.info("PDTransformer have {} paramerters in total".format(sum(x.numel() for x in model.parameters()) + sum(x.numel() for x in autoencoder.parameters())))

    Avg_loss_train = np.zeros(params.num_epochs)
    Avg_loss_valid = np.zeros(params.num_epochs)

    for epoch in range(params.num_epochs):
        logger.info('Epoch {}/{}'.format(epoch + 1, params.num_epochs))
        # valid_metrics = evaluate(model, valid_loader, params, epoch)

        train_loss,loss_train_epoch[epoch * train_len:(epoch + 1) * train_len] = \
            train(model,autoencoder,train_loader,optimizer, \
                scaler,params,epoch)
        logger.info('NLL-loss:{}'.format(train_loss))
        if epoch<20:
            scheduler.step()
        logger.info('Learning rate:{}'.format(optimizer.state_dict()['param_groups'][0]['lr']))

        valid_metrics = evaluate(model,autoencoder, valid_loader, params, epoch)
        test_metrics = evaluate(model,autoencoder, test_loader, params, epoch)
        loss_valid[epoch * valid_len:(epoch + 1) * valid_len] = valid_metrics['test_loss']

        Avg_loss_train[epoch] = train_loss
        Avg_loss_valid[epoch] = valid_metrics['test_loss']

        q50_valid[epoch] = valid_metrics['q50']
        q90_valid[epoch] = valid_metrics['q90']
        # MAPE_valid[epoch] = valid_metrics['MAPE']

        valid_loss.append(valid_metrics['q50'])

        is_best = q50_valid[epoch] <= best_valid_q50

        # Save weights
        utils.save_checkpoint({'epoch': epoch + 1,
                               'state_dict': model.state_dict(),
                               'auto_dict': autoencoder.state_dict(),
                               'optim_dict': optimizer.state_dict()},
                                epoch=epoch,
                                is_best=is_best,
                                checkpoint=params.model_dir)

        if is_best:
            logger.info('- Found new best Q90/Q50')
            best_valid_q50 = q50_valid[epoch]
            best_json_path = os.path.join(params.model_dir, 'metrics_valid_best_weights.json')
            utils.save_dict_to_json(valid_metrics, best_json_path)
            best_valid_ND = epoch

        utils.save_loss(loss_train_epoch[epoch * train_len:(epoch + 1) * train_len], args.dataset + '_' + str(epoch) +'-th_epoch_loss', params.plot_dir)
        utils.save_loss(loss_valid[epoch * valid_len:(epoch + 1) * valid_len], args.dataset + '_' + str(epoch) +'-th_epoch_valid_loss', params.plot_dir)

        last_json_path = os.path.join(params.model_dir, 'metrics_valid_last_weights.json')
        utils.save_dict_to_json(valid_metrics, last_json_path)
        early_stopping(valid_loss[-1], model)
        if early_stopping.early_stop:
            print("Early stopping")
            # save weights
            utils.save_checkpoint({'epoch': epoch + 1,
                                'state_dict': model.state_dict(),
                                'auto_dict': autoencoder.state_dict(),
                                'optim_dict': optimizer.state_dict()},
                                filepath=params.model_dir)
            break

    loss_train_epoch = np.array(loss_train_epoch)
    plt.figure()
    plt.figure(figsize=(20, 10))
    plt.plot(Avg_loss_train[1:], color= 'red',label='loss')
    # plt.plot(Avg_loss_valid[1:], color= 'blue',label='loss')
    plt.xlabel('Epoch')
    plt.savefig(params.model_dir+'/DeFigs/Loss.png')
    plt.close()

    if args.save_best:
        f = open('./param_search.txt', 'w')
        f.write('-----------\n')
        list_of_params = list(params.__dict__.keys())
        print_params = ''
        for param in list_of_params:
            param_value = getattr(params, param)
            print_params += f'{param}: {param_value:.2f}'
        print_params = print_params[:-1]
        f.write(print_params + '\n')
        f.write('Best ND: ' + str(best_valid_ND) + '\n')
        logger.info(print_params)
        logger.info(f'Best ND: {best_valid_ND}')
        f.close()

if __name__ == '__main__':
    # setup_seed(12)
    # Load the parameters from json file
    args = parser.parse_args()

    model_dir = args.model_name
    json_path = os.path.join(model_dir, 'params.json')
    data_dir = os.path.join(args.data_folder, args.dataset)
    assert os.path.isfile(json_path), f'No json configuration file found at {json_path}'
    params = utils.Params(json_path)
    params.relative_metrics = args.relative_metrics
    params.model_dir = model_dir
    params.plot_dir = os.path.join(model_dir, 'figures')
    params.dataset = args.dataset

    # create missing directories
    try:
        os.mkdir(params.plot_dir)
    except FileExistsError:
        pass

    # use GPU if available
    params.ngpu = torch.cuda.device_count()

    params.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    logger.info('Using Cuda...')
    c = copy.deepcopy

    attn = transformer.MultiHeadedAttention(params)
    ff = transformer.PositionwiseFeedForward(params.d_model, d_ff=params.d_ff, dropout=params.dropout)
    position = transformer.PositionalEncoding(params.d_model, dropout=params.dropout)
    emb = transformer.Embedding(params, position)
    ############################################################
    encoder  = transformer.Encoder(params, transformer.EncoderLayer(params, c(attn), c(ff), dropout=params.dropout))
    decoder = transformer.Decoder(params, transformer.DecoderLayer(params, c(attn), c(attn), c(ff), dropout=params.dropout))
    generator = transformer.Generator(params)
    model = transformer.EncoderDecoder(params= params, emb = emb, encoder = encoder, decoder = decoder, generator = generator)
    autoencoder = transformer.AutoEncoder(params)
    autoencoder.to(params.device)

    # if torch.cuda.device_count() > 1:
    #     model = nn.DataParallel(model)
    #     discriminator = nn.DataParallel(discriminator)

    model.to(params.device)

    utils.set_logger(os.path.join(model_dir, 'train.log'))

    logger.info('Loading the datasets...')

    train_set = TrainDataset(data_dir, args.dataset, params.num_class)
    valid_set = ValidDataset(data_dir, args.dataset, params.num_class)
    test_set = TestDataset(data_dir, args.dataset, params.num_class)
    train_loader = DataLoader(train_set, batch_size=params.batch_size, sampler=RandomSampler(train_set), num_workers=4)
    valid_loader = DataLoader(valid_set, batch_size=params.predict_batch, sampler=RandomSampler(valid_set), num_workers=4)
    test_loader = DataLoader(test_set, batch_size=params.predict_batch, sampler=RandomSampler(test_set), num_workers=4)
    logger.info('Loading complete.')

    n_updates_total = (train_set.__len__() // params.batch_size) * params.num_epochs

    optimizer = optim.Adam(chain(model.parameters(),autoencoder.parameters()), lr = args.lr)

    scheduler =  torch.optim.lr_scheduler.StepLR(optimizer,step_size=2,gamma = 0.8)

    scaler = GradScaler()
    # Train the model
    logger.info('Starting training for {} epoch(s)'.format(params.num_epochs))
    train_and_evaluate(model,autoencoder,
                       train_loader,valid_loader,test_loader,
                       optimizer,scheduler,
                       scaler,
                       params,
                       args.restore_file)


from cli_maker.cli_function import CliFunction
from cli_maker.line_parser import parse_line
from textwrap import dedent
import sys
import inspect

class Cli(object):
    def __init__(self) -> None:
        super().__init__()
        self._functions = {}

    def cli_function(self, names=None):
        def decorator(func):
            sig = inspect.signature(func)
            arguments = {}
            flags = set()
            for key, param in sig.parameters.items():
                if param.default != inspect.Parameter.empty:
                    if param.annotation == bool:
                        flags.add(key)
                    else:
                        if param.annotation == inspect.Parameter.empty or param.annotation == any:
                            arguments[key] = str
                        else:
                            arguments[key] = param.annotation

            f = CliFunction(func, arguments, flags)
            self._functions[f.name] = f
            if names != None:
                for n in names:
                    self._functions[n] = f

            def wrapper(*args, **kwargs):
                return func(*args, **kwargs)

            return wrapper
        return decorator

    def _show_all_functions(self):
        all_names = sorted(list(set([f.name for f in self._functions.values()])))
        length = len(max(all_names, key=lambda x: len(x)))
        print("Commands:")
        for name in all_names:
            if self._functions[name].func.__doc__:
                splt = dedent(self._functions[name].func.__doc__).split('\n')
                i = 0
                desc = splt[i]
                while not desc:
                    i += 1
                    if i == len(splt):
                        desc = "\n".join(splt)
                        break
                    desc = splt[i]
                print(f"  {name.ljust(length, ' ')}  {desc}")
            else:
                print(f"  {name.ljust(length, ' ')}  ")

    def run(self):
        @self.cli_function(names=["h"])
        def help(fname):
            """Shows help about a specific function."""
            if fname not in self._functions:
                print(f"Unkown function: {fname}")
                return
            elif fname == "help" or fname == "h":
                self._show_all_functions()
                return

            names = [k for k in self._functions.keys() if self._functions[k].name == fname]
            original_name = self._functions[names[0]].name
            names.remove(original_name)
            if len(names) == 0:
                print(original_name + "\n" + "-"*len(original_name), end="")
            else:
                print(f"{original_name} (also {', '.join(names)})\n" + "-"*len(f"{original_name} (also {', '.join(names)})"), end="")
            print(dedent(self._functions[original_name].func.__doc__) if self._functions[original_name].func.__doc__ else "\nNo documentation found.")

        try:
            fname, regular, arguments, flags = parse_line(sys.argv[1:])
        except TypeError:
            self._show_all_functions()
            return
        if fname not in self._functions:
            print(f"Unknown function: {fname}")
            return

        try:
            self._functions[fname](regular, arguments, flags)
        except Exception as e:
            print(e)


import rrdtool
import os
import time

def grafica(filename, time_start, time_end, title, var, db, name):
    if not os.path.exists(os.path.join(os.getcwd(), "graficas")):
        os.mkdir(os.path.join(os.getcwd(), "graficas"))

    try:
        rrdtool.graph(filename, "--start", str(time_start), "--end", str(time_end),
                      "--vertical-label=Bytes/s", "--title=" + title,
                      "DEF:variable=" + db + ":" + var + ":AVERAGE",
                      "CDEF:escala=variable,8,*",
                      "LINE3:escala#0000FF:" + name)
        return "OK"
    except Exception as e:
        print(e)
        return "Error al generar la grafica :" + e.args[0]

datos = [
    ["ifInMulticastPkts", "ifIn.png", "Paquetes multicast de entrada"],
    ["ipInReceives", "ipIn.png", "Tráfico de entrada"],
    ["icmpOutEchoReps", "icmpOut.png", "Respuestas ICMP de salida"],
    ["tcpOutSegs", "tcpOut.png", "Paquetes TCP de salida"],
    ["udpInOverflows", "udpIn.png", "Paquetes UDP de entrada"]
]

#for dato in datos:
#    # Obtener el timestamp actual
#    now = int(time.time())
#    # Restar 10 minutos (600 segundos)
#    start = now - 600#

#    grafica(
#        os.path.join(os.getcwd(), "graficas", dato[1]),
#        start, now, datos,
#        dato[0],
#        os.path.join(os.getcwd(), "RRD", "1.rrd"),
#        dato[2])