import argparse
import json
import logging
import math
import os
import random
import time
import warnings
warnings.filterwarnings('ignore')

os.environ["HOME"] = "/home/nsml"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import numpy as np
from sklearn.model_selection import KFold

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
from torch.optim.lr_scheduler import LambdaLR
import transformers
from transformers import BertTokenizer, get_linear_schedule_with_warmup

from model import TransformerModel, EncoderDecoder, BertGpt2, TransformerS2Model
from tokenizer import CharTokenizer, CustomTokenizer
from dataset import TextDataset, collate_fn, collate_fn_pretrain, collate_fn_transformer, collate_fn_pretrain_transformer, collate_fn_mlm_pretrain_transformer
from data_loader import read_strings
from meter import Meter
from evaluation import em, gleu

import larva
from larva import LarvaTokenizer, LarvaModel, list_larva_models

import nsml
from nsml import DATASET_PATH

logging.basicConfig(format='%(asctime)s -  %(message)s', datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO)
logger = logging.getLogger(__name__)

def get_args():
    parser = argparse.ArgumentParser()
    # data
    parser.add_argument("--data_dir", type=str, default=os.path.join(DATASET_PATH, 'train'))
    parser.add_argument("--val_ratio", type=float, default=0.2)
    # parser.add_argument("--num_val_data", type=int, default=1000)

    # tokenizer
    parser.add_argument("--limit_alphabet", type=int, default=6000)
    parser.add_argument("--tokenizer_vocab_size", type=int, default=6000)

    # model
    parser.add_argument("--model", type=str, default='transformer')
    parser.add_argument("--vocab_size", type=int, default=1600)
    parser.add_argument("--hidden_size", type=int, default=512)
    parser.add_argument("--num_attention_heads", type=int, default=8)
    parser.add_argument("--num_encoder_layers", type=int, default=6)
    parser.add_argument("--num_decoder_layers", type=int, default=6)
    parser.add_argument("--intermediate_size", type=int, default=1024)
    parser.add_argument("--ffn_size", type=int, default=2048)
    parser.add_argument("--share_decoder_input_output_embed", type=str2bool, default=False)

    # bert
    parser.add_argument("--bert", type=str, default='original')
    parser.add_argument("--bert_hidden_size", type=int, default=768)
    parser.add_argument("--bert_num_attention_heads", type=int, default=12)
    parser.add_argument("--num_hidden_layers", type=int, default=12)
    parser.add_argument("--bert_intermediate_size", type=int, default=3072)

    # training
    parser.add_argument("--kfold", type=str2bool, default=False)
    parser.add_argument("--n_splits", type=int, default=5)

    parser.add_argument("--mlm_pre", type=str2bool, default=False)
    parser.add_argument("--num_pretrain_epochs", type=int, default=5)

    parser.add_argument("--train_all", type=str2bool, default=False)

    parser.add_argument("--max_seq_length", type=int, default=128)
    parser.add_argument("--train_batch_size", type=int, default=128)
    parser.add_argument("--eval_batch_size", type=int, default=64)

    parser.add_argument("--lr", type=float, default=5e-4)
    parser.add_argument("--tokenizer", type=str, default='char')
    parser.add_argument("--optimizer", type=str, default='Adam')
    parser.add_argument("--scheduler", type=str, default='lambda')
    parser.add_argument("--adam_betas", type=str, default="(0.9, 0.98)")
    parser.add_argument("--eps", type=float, default=1e-9)
    parser.add_argument("--weight_decay", type=float, default=0.0001)
    parser.add_argument("--max_grad_norm", type=float, default=1.0)
    parser.add_argument("--dropout", type=float, default=0.3)

    parser.add_argument("--max_steps", type=int, default=10000)
    parser.add_argument("--num_warmup_steps", type=int, default=4000)
    parser.add_argument("--warmup_ratio", type=float, default=0.1)
    parser.add_argument("--log_interval", type=int, default=100)
    parser.add_argument("--eval_interval", type=int, default=1000)

    parser.add_argument('--seed', type=int, default=2021)
    parser.add_argument('--num_workers', type=int, default=20)

    # nsml
    parser.add_argument('--mode', type=str, default="train")
    parser.add_argument('--pause', type=int, default=0)
    parser.add_argument('--iteration', type=str, default="0")

    args = parser.parse_args()
    return args

def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)

def str2bool(v : str) -> bool:
    """convert string argument to boolean
    Args:
        v (str)
    Raises:
        argparse.ArgumentTypeError: [Boolean value expected]
    Returns:
        bool
    """
    if isinstance(v, bool):
        return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

def postprocess_state(sentence : str) -> str:
    """TRADE, SUMBT postprocessing
    Args:
        state (List[str]): state prediction
    Returns:
        List[str]: postprocessing state
    """
    sentence = sentence.replace(" : ", ":").replace(" , ", ", ").replace('( ', '(').replace(' )', ')').replace(' & ', '&').replace(' = ', '=')
    sentence = sentence.replace(" % ", "%").replace(' ~ ', '~').replace(' ^ ', '^')
    if sentence.endswith(' ~'):
        sentence = sentence.replace(' ~', '~')
    if sentence.endswith(' ^^'):
        sentence = sentence.replace(' ^^', '^^')
    if sentence.endswith(' ^'):
        sentence = sentence.replace(' ^', '^')
    if sentence.endswith('......'):
        sentence = sentence.replace('......', ' ......')
    sentence = sentence.replace(') 에', ')에').replace('곳 (', '곳(').replace('부터~트', '부터~ 트').replace('# 정왕동', '#정왕동')
    sentence = sentence.replace('쨘 -', '쨘-').replace('해드리겠습니다!', '해드리겠습니다 !').replace('6 / 6', '6/6').replace('6 / 4', '6/4')
    sentence = sentence.replace('> ㅋ', '>ㅋ').replace('이상~헤', '이상~ 헤').replace('6 / 6', '6/6').replace('6 / 4', '6/4')

    return sentence

def get_lr(optimizer : transformers.optimization) -> float:
    """Get learning_rate.

    Args:
        optimizer (transformers.optimization).

    Return:
        Learning_Rate (float).
    """
    for param_group in optimizer.param_groups:
        return param_group['lr']

def calc_loss(model, batch, args):
    src, tgt, src_mask, tgt_mask, tgt_label = batch
    if args.model == 'transformer':
        output = model(src=src, tgt=tgt, src_key_padding_mask=~src_mask, tgt_key_padding_mask=~tgt_mask)
    elif args.model in ['encoderdecoder', 'BertGpt2']:
        output = model(src=src, tgt=tgt, src_key_padding_mask=src_mask, tgt_key_padding_mask=tgt_mask)
    elif args.model == 'bertfuse':
        output = model(src_tokens=src, prev_output_tokens=tgt, bert_input=src)

    if args.tokenizer in ['char', 'pre_char', 'kfold_char']:
        mask_tgt_label = tgt_label.clone()
        mask_tgt_label[tgt_label.eq(1)] = -100
    else:
        mask_tgt_label = tgt_label.clone()
        mask_tgt_label[tgt_label.eq(0)] = -100

    raw_loss = F.cross_entropy(output.view(-1, output.size(-1)), mask_tgt_label.view(-1), reduction='none')
    # raw_loss = F.cross_entropy(output.view(-1, output.size(-1)), tgt_label.view(-1), reduction='none')

    if args.model == 'transformer':
        bsz = tgt.size(1)
        raw_loss = raw_loss.view(-1, bsz)
    else:
        bsz = tgt.size(0)
        raw_loss = raw_loss.view(bsz, -1)

    if args.model == 'transformer':
        loss = (raw_loss * tgt_mask.float()).sum(0).mean()
    else:
        loss = (raw_loss * tgt_mask.float()).sum(-1).mean()

    items = [loss.data.item(), bsz, tgt_mask.sum().item()]
    return loss, items

def calc_loss1(model, batch, args):
    src, tgt, src_mask, tgt_mask, tgt_label = batch

    print(f"source padded is {src}, source padded's shape is {src.shape}")
    print(f"target padded before -1 is {tgt}, target padded before -1 's shape is {tgt.shape}")
    print(f"source padding mask is {src_mask}, source padding mask's shape is {src_mask.shape}")
    print(f"target padding mask is {tgt_mask}, target padding mask's shape is {tgt_mask.shape}")
    print(f"target padded after 1 is {tgt_label}, target padded after 1 's shape is {tgt_label.shape}")

    if args.model == 'transformer':
        output = model.forward1(src=src, tgt=tgt, src_key_padding_mask=~src_mask, tgt_key_padding_mask=~tgt_mask)
    elif args.model in ['encoderdecoder', 'BertGpt2']:
        output = model.forward1(src=src, tgt=tgt, src_key_padding_mask=src_mask, tgt_key_padding_mask=tgt_mask)
    elif args.model == 'bertfuse':
        output = model(src_tokens=src, prev_output_tokens=tgt, bert_input=src)

    print(f"output shape is {output.shape}")

    if args.tokenizer in ['char', 'pre_char', 'kfold_char']:
        mask_tgt_label = tgt_label.clone()
        mask_tgt_label[tgt_label.eq(1)] = -100
    else:
        mask_tgt_label = tgt_label.clone()
        mask_tgt_label[tgt_label.eq(0)] = -100

    print(f"cross entropy first input is {output.view(-1, output.size(-1)).shape}, second input is {tgt_label.view(-1).shape}")

    raw_loss = F.cross_entropy(output.view(-1, output.size(-1)), mask_tgt_label.view(-1), reduction='none')
    # raw_loss = F.cross_entropy(output.view(-1, output.size(-1)), tgt_label.view(-1), reduction='none')

    if args.model == 'transformer':
        bsz = tgt.size(1)
        raw_loss = raw_loss.view(-1, bsz)
    else:
        bsz = tgt.size(0)
        raw_loss = raw_loss.view(bsz, -1)

    print(f"raw_loss is {raw_loss}, raw_loss 's shape is {raw_loss.shape}")
    if args.model == 'transformer':
        loss = (raw_loss * tgt_mask.float()).sum(0).mean()
    else:
        print(f"sum of raw_loss is {(raw_loss * tgt_mask.float()).sum(-1)}, sum of raw_loss's shape is {(raw_loss * tgt_mask.float()).sum(-1).shape}")
        loss = (raw_loss * tgt_mask.float()).sum(-1).mean()

    print(f"loss is {loss}, loss 's shape is {loss.shape}")

    items = [loss.data.item(), bsz, tgt_mask.sum().item()]
    return loss, items

def evaluate(model, data_loader, args):
    model.eval()
    meter = Meter()
    with torch.no_grad():
        for batch in data_loader:
            batch = tuple(t.to(args.device) for t in batch)
            _, items = calc_loss(model, batch, args)
            meter.add(*items)
    return meter.average(), meter.print_str(False)

def correct(model, tokenizer, test_data, args):
    model.eval()
    prediction = []
    for i in range(0, len(test_data), args.eval_batch_size):
        batch = test_data[i:i + args.eval_batch_size]

        if args.tokenizer in ['char', 'pre_char', 'kfold_char']:
            # src_token_ids = [tokenizer(x) for x in batch]
            src_token_ids = [[2] + tokenizer(x) + [3] for x in batch]
        elif args.tokenizer in ['dsksd', 'plus', 'wordpiece']:
            src_token_ids = tokenizer(batch).input_ids

        if args.model == 'transformer':
            src_seq_length = [len(x) for x in src_token_ids]
            src_max_seq_length = max(src_seq_length)
            src_padded = []
            src_padding_mask = []
            for x in src_token_ids:
                x = x[:src_max_seq_length]
                src_pad_length = src_max_seq_length - len(x)
                src_padded.append(x + [tokenizer.pad_token_id] * src_pad_length)
                src_padding_mask.append([1] * len(x) + [0] * src_pad_length)

            src_padded = torch.tensor(src_padded).t().contiguous().to(args.device)
            src_padding_mask = torch.tensor(src_padding_mask).bool().t().to(args.device)

            memory = model(src=src_padded, src_key_padding_mask=~src_padding_mask)

            tgt_token_ids = [[2] for _ in batch]
            end = [False for _ in batch]
            for l in range(src_max_seq_length + 20):
                tgt = torch.tensor(tgt_token_ids).t().contiguous().to(args.device)
                output = model(tgt=tgt, memory=memory, memory_key_padding_mask=~src_padding_mask)
                top1 = output[-1].argmax(-1).tolist()
                for i, tok in enumerate(top1):
                    if tok == 3 or l >= src_seq_length[i] + 20:
                        end[i] = True
                    tgt_token_ids[i].append(tok if not end[i] else 3)
                if all(end):
                    break
        elif args.model in ['encoderdecoder', 'BertGpt2']:
            src_seq_length = [len(x) for x in src_token_ids]
            src_max_seq_length = max(src_seq_length)

            src_padded = []
            for x in src_token_ids:
                x = x[:src_max_seq_length]
                src_pad_length = src_max_seq_length - len(x)
                src_padded.append(x + [tokenizer.pad_token_id] * src_pad_length)

            src_padded = torch.LongTensor(src_padded).contiguous().to(args.device)
            src_padding_mask = src_padded.ne(tokenizer.pad_token_id).to(args.device)

            memory = model.encoder(
                input_ids=src_padded,
                attention_mask=src_padding_mask,
            )

            tgt_token_ids = [[2] for _ in batch]
            end = [False for _ in batch]
            for l in range(src_max_seq_length + 20):
                tgt_padded = torch.LongTensor(tgt_token_ids).contiguous().to(args.device)
                tgt_padding_mask = tgt_padded.ne(tokenizer.pad_token_id).to(args.device)

                # output = model(src=src_padded, tgt=tgt_padded, src_key_padding_mask=src_padding_mask, tgt_key_padding_mask=tgt_padding_mask)
                output = model(tgt=tgt_padded, memory=memory, src_key_padding_mask=src_padding_mask, tgt_key_padding_mask=tgt_padding_mask)
                top1 = output[:,-1,:].argmax(-1).tolist()

                for i, tok in enumerate(top1):
                    if tok == 3 or l >= src_seq_length[i] + 20:
                        end[i] = True
                    tgt_token_ids[i].append(tok if not end[i] else 3)
                if all(end):
                    break
        elif args.model == 'bertfuse':
            src_seq_length = [len(x) for x in src_token_ids]
            src_max_seq_length = max(src_seq_length)

            src_padded = []
            for x in src_token_ids:
                x = x[:src_max_seq_length]
                src_pad_length = src_max_seq_length - len(x)
                src_padded.append(x + [tokenizer.pad_token_id] * src_pad_length)

            src_padded = torch.LongTensor(src_padded).contiguous().to(args.device)
            bert_encoder_padding_mask = src_padded.eq(tokenizer.pad_token_id)

            bert_encoder_out =  model.bert_encoder(src_padded, attention_mask=~bert_encoder_padding_mask)
            bert_encoder_out = bert_encoder_out.last_hidden_state

            bert_encoder_out = bert_encoder_out.permute(1,0,2).contiguous()
            bert_encoder_out = {
                'bert_encoder_out': bert_encoder_out,
                'bert_encoder_padding_mask': bert_encoder_padding_mask,
            }

            encoder_out = model.encoder(src_padded, bert_encoder_out=bert_encoder_out)

            tgt_token_ids = [[2] for _ in batch]
            end = [False for _ in batch]
            for l in range(src_max_seq_length + 20):
                tgt_padded = torch.LongTensor(tgt_token_ids).contiguous().to(args.device)

                output = model.decoder(tgt_padded, encoder_out=encoder_out, bert_encoder_out=bert_encoder_out)[0]

                top1 = output[:,-1,:].argmax(-1).tolist()

                for i, tok in enumerate(top1):
                    if tok == 3 or l >= src_seq_length[i] + 20:
                        end[i] = True
                    tgt_token_ids[i].append(tok if not end[i] else 3)
                if all(end):
                    break

        if args.tokenizer in ['char', 'pre_char', 'kfold_char']:
            # prediction.extend([''.join([tokenizer.i2c[tok] for tok in x if tok >= 5]) for x in tgt_token_ids])
            prediction.extend([postprocess_state(''.join([tokenizer.i2c[tok] for tok in x if tok >= 5])) for x in tgt_token_ids])
        elif args.tokenizer in ['dsksd', 'plus', 'wordpiece']:
            prediction.extend([postprocess_state(tokenizer.decode([tok for tok in x if tok >= 5])) for x in tgt_token_ids])
    return prediction

def mlm_pretrain(args : argparse.Namespace, model : nn.Module, loader : DataLoader, optimizer : transformers.optimization, scheduler,
                tokenizer : BertTokenizer, loss_fnc_pretrain : nn.CrossEntropyLoss, epoch : int) -> None:
    """Masked Language Model Pretrain
    Args:
        args (argparse.Namespace)
        model (nn.Module)
        loader (DataLoader)
        optimizer (transformers.optimization)
        tokenizer (BertTokenizer)
        loss_fnc_pretrain (nn.CrossEntropyLoss): mlm pretrain loss
        epoch (int)
    """
    model.train()
    for step, batch in enumerate(loader):
        src, tgt, src_mask, tgt_mask, tgt_label = tuple(t.to(args.device) for t in batch)

        logits, labels = model.forward_pretrain(src, src_mask, tokenizer)
        loss = loss_fnc_pretrain(logits.view(-1, args.vocab_size), labels.view(-1))

        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
        optimizer.step()
        optimizer.zero_grad()
        scheduler.step()

        if step % 100 == 0:
            current_lr = get_lr(optimizer)
            epochs = args.num_pretrain_epochs
            print(
                f"[{epoch + 1}/{epochs}] [{step}/{len(loader)}] loss : {loss.item()}  lr : {current_lr}"
            )

def train(model, tokenizer, train_data, valid_data, corpuses, args, index = None):
    global best_val_gleu

    train_dataset = TextDataset(train_data)

    print(f"\nfirst train_dataset is {train_dataset[0]}")

    if args.model == 'transformer':
        if args.tokenizer in ['char', 'pre_char', 'kfold_char']:
            collate_func = collate_fn
        else:
            collate_func = collate_fn_pretrain
    elif args.model in ['encoderdecoder', 'BertGpt2', 'bertfuse']:
        if args.tokenizer in ['char', 'pre_char', 'kfold_char']:
            collate_func = collate_fn_transformer
        elif args.tokenizer in ['plus', 'wordpiece']:
            collate_func = collate_fn_pretrain_transformer

    train_dataloader = DataLoader(
        train_dataset,
        sampler=RandomSampler(train_dataset),
        batch_size=args.train_batch_size,
        num_workers=args.num_workers,
        collate_fn=lambda x: collate_func(x, tokenizer, args.max_seq_length),
    )

    if valid_data is not None:
        valid_dataset = TextDataset(valid_data)
        valid_dataloader = DataLoader(
            valid_dataset,
            sampler=SequentialSampler(valid_dataset),
            batch_size=args.eval_batch_size,
            num_workers=args.num_workers,
            collate_fn=lambda x: collate_func(x, tokenizer, args.max_seq_length),
        )

        valid_noisy = [x['noisy'] for x in valid_data]
        valid_clean = [x['clean'] for x in valid_data]

    epochs = (args.max_steps - 1) // len(train_dataloader) + 1
    optimizer = getattr(torch.optim, args.optimizer)(
        model.parameters(),
        lr=args.lr,
        betas=eval(args.adam_betas),
        eps=args.eps,
        weight_decay=args.weight_decay
    )

    if args.scheduler == 'lambda':
        lr_lambda = lambda x: x / args.num_warmup_steps if x <= args.num_warmup_steps else (x / args.num_warmup_steps) ** -0.5
        scheduler = LambdaLR(optimizer, lr_lambda)
    elif args.scheduler == 'linear':
        warmup_steps = args.max_steps * args.warmup_ratio
        scheduler = get_linear_schedule_with_warmup(
            optimizer, num_warmup_steps=warmup_steps, num_training_steps=args.max_steps
        )

    if args.mlm_pre and (index is None or index == 0):
        print('mlm_pretrain!!!')
        pre_optimizer = getattr(torch.optim, args.optimizer)(
            model.parameters(),
            lr=args.lr,
            betas=eval(args.adam_betas),
            eps=args.eps,
            weight_decay=args.weight_decay
        )
        loss_fnc_pretrain = nn.CrossEntropyLoss()

        if args.tokenizer == 'char':
            t_total = len(train_dataloader) * args.num_pretrain_epochs
            warmup_steps = t_total * args.warmup_ratio
            pre_scheduler = get_linear_schedule_with_warmup(
                pre_optimizer, num_warmup_steps=warmup_steps, num_training_steps=t_total
            )
            for epoch in range(args.num_pretrain_epochs):
                mlm_pretrain(args, model, train_dataloader, pre_optimizer, pre_scheduler, tokenizer, loss_fnc_pretrain, epoch)
        elif args.tokenizer == 'pre_char':
            corpus_dataset = TextDataset(corpuses)
            corpus_dataloader = DataLoader(
                corpus_dataset,
                sampler=RandomSampler(corpus_dataset),
                batch_size=args.train_batch_size,
                num_workers=args.num_workers,
                collate_fn=lambda x: collate_fn_pretrain(x, tokenizer, args.max_seq_length),
            )
            t_total = len(corpus_dataloader) * args.num_pretrain_epochs
            warmup_steps = t_total * args.warmup_ratio
            pre_scheduler = get_linear_schedule_with_warmup(
                pre_optimizer, num_warmup_steps=warmup_steps, num_training_steps=t_total
            )
            for epoch in range(args.num_pretrain_epochs):
                mlm_pretrain(args, model, corpus_dataloader, pre_optimizer, pre_scheduler, tokenizer, loss_fnc_pretrain, epoch)
        elif args.tokenizer == 'wordpiece':
            corpus_dataset = TextDataset(corpuses)
            corpus_dataloader = DataLoader(
                corpus_dataset,
                sampler=RandomSampler(corpus_dataset),
                batch_size=args.train_batch_size,
                num_workers=args.num_workers,
                collate_fn=lambda x: collate_fn_mlm_pretrain_transformer(x, tokenizer, args.max_seq_length),
            )
            t_total = len(corpus_dataloader) * args.num_pretrain_epochs
            warmup_steps = t_total * args.warmup_ratio
            pre_scheduler = get_linear_schedule_with_warmup(
                pre_optimizer, num_warmup_steps=warmup_steps, num_training_steps=t_total
            )
            for epoch in range(args.num_pretrain_epochs):
                mlm_pretrain(args, model, corpus_dataloader, pre_optimizer, pre_scheduler, tokenizer, loss_fnc_pretrain, epoch)

    step = 0
    if index is None or index == 0:
        best_val_gleu = -float("inf")
    meter = Meter()

    for epoch in range(1, epochs + 1):
        for batch in train_dataloader:
            model.train()
            step += 1
            batch = tuple(t.to(args.device) for t in batch)

            if step < 4:
                loss, items = calc_loss1(model, batch, args)
            else:
                loss, items = calc_loss(model, batch, args)

            if step < 4:
                print(f"first item loss is {items[0]}")
                print(f"second item bsz is {items[1]}")
                print(f"third item target mask sum is {items[2]}")
                print()

            meter.add(*items)

            loss.backward()
            if args.max_grad_norm > 0:
                nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
            optimizer.step()
            model.zero_grad()
            scheduler.step()

            if step % args.log_interval == 0:
                lr = scheduler.get_lr()[0]
                loss_sent, loss_token = meter.average()

                logger.info(f' [{step:5d}] lr {lr:.6f} | {meter.print_str(True)}')
                nsml.report(
                    step=step,
                    scope=locals(),
                    summary=True,
                    train__lr=lr,
                    train__loss_sent=loss_sent,
                    train__token_ppl=math.exp(loss_token)
                )
                meter.init()

            if step % args.eval_interval == 0 and valid_data is not None:
                start_eval = time.time()
                (val_loss, val_loss_token), valid_str = evaluate(model, valid_dataloader, args)
                prediction = correct(model, tokenizer, valid_noisy, args)

                for (i, (pred, clean)) in enumerate(zip(prediction, valid_clean)):
                    if i == 5:
                        break
                    print(f"{i + 1}'s pred sentence is {pred}, clean sentence is {clean}")

                val_em = em(prediction, valid_clean)
                val_gleu = gleu(prediction, valid_clean)

                logger.info('-' * 89)
                logger.info(f' [{step:6d}] valid | {valid_str} | em {val_em:5.2f} | gleu {val_gleu:5.2f}')
                logger.info('-' * 89)
                nsml.report(
                    step=step,
                    scope=locals(),
                    summary=True,
                    valid__loss_sent=val_loss,
                    valid__token_ppl=math.exp(val_loss_token),
                    valid__em=val_em,
                    valid__gleu=val_gleu
                )

                if val_gleu > best_val_gleu:
                    if args.kfold:
                        if index == args.n_splits - 1:
                            print('------------save model------------')
                            nsml.save(f"model_{epoch}")
                    else:
                        print('------------save model------------')
                        best_val_gleu = val_gleu
                        nsml.save("best")
                meter.start += time.time() - start_eval

            elif step % args.eval_interval == 0 and valid_data is None:
                if epoch >= 5:
                    print('------------save model------------')
                    nsml.save(f"model_{epoch}")

            if step >= args.max_steps:
                break
        if step >= args.max_steps:
            break

def bind_nsml(model, tokenizer=None, args=None):
    def save(path, **kwargs):
        torch.save(model.state_dict(), open(os.path.join(path, 'model.pt'), 'wb'))
        if tokenizer is not None and args.tokenizer in ['char', 'pre_char', 'kfold_char']:
            tokenizer.save(os.path.join(path, 'vocab.txt'))
        elif args.tokenizer == 'wordpiece':
            tokenizer.save(os.path.join(path, 'wordpiece_tokenizer'))

    def load(path, **kwargs):
        model.load_state_dict(torch.load(open(os.path.join(path, 'model.pt'), 'rb'),
                                         map_location=lambda storage, loc: storage))
        if tokenizer is not None and args.tokenizer in ['char', 'pre_char', 'kfold_char']:
            tokenizer.load(os.path.join(path, 'vocab.txt'))
        elif args.tokenizer == 'wordpiece':
            tokenizer.load(os.path.join(path, 'wordpiece_tokenizer'))

    def infer(test_data, **kwargs):
        '''
        :param test_data: list of noisy sentences
        :return: list of corrected sentences
        '''
        return correct(model, tokenizer, test_data, args)

    import nsml
    nsml.bind(save, load, infer)

def main():
    args = get_args()
    logger.info(f"args: {json.dumps(args.__dict__, indent=2, sort_keys=True)}")

    args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    args.n_gpu = torch.cuda.device_count()
    set_seed(args)

    if args.tokenizer in ['char', 'pre_char', 'kfold_char']:
        tokenizer = CharTokenizer([])
    elif args.tokenizer == 'dsksd':
        tokenizer = BertTokenizer.from_pretrained('dsksd/bert-ko-small-minimal')
        print(f"vocab size is {len(tokenizer)}")
        args.vocab_size = len(tokenizer)
    elif args.tokenizer == 'plus':
        tokenizer = LarvaTokenizer.from_pretrained("larva-kor-plus-base-cased")
        print(f"vocab size is {len(tokenizer)}")
        args.vocab_size = len(tokenizer)
    elif args.tokenizer == 'wordpiece':
        tokenizer = CustomTokenizer(args)

    if args.model == 'transformer':
        model = TransformerModel(
            vocab_size=args.vocab_size,
            hidden_size=args.hidden_size,
            num_attention_heads=args.num_attention_heads,
            num_encoder_layers=args.num_encoder_layers,
            num_decoder_layers=args.num_decoder_layers,
            intermediate_size=args.intermediate_size,
            dropout=args.dropout,
        ).to(args.device)
    elif args.model == 'encoderdecoder':
        model = EncoderDecoder(
            vocab_size=args.vocab_size,
            hidden_size=args.hidden_size,
            num_attention_heads=args.num_attention_heads,
            num_encoder_layers=args.num_encoder_layers,
            num_decoder_layers=args.num_decoder_layers,
            intermediate_size=args.intermediate_size,
            dropout=args.dropout,
        ).to(args.device)
    elif args.model == 'BertGpt2':
        model = BertGpt2(
            vocab_size=args.vocab_size,
            hidden_size=args.hidden_size,
            num_attention_heads=args.num_attention_heads,
            num_encoder_layers=args.num_encoder_layers,
            num_decoder_layers=args.num_decoder_layers,
            intermediate_size=args.intermediate_size,
            dropout=args.dropout,
        ).to(args.device)
    elif args.model == 'bertfuse':
        model = TransformerS2Model(tokenizer, args=args).to(args.device)

    logger.info(f"# of model parameters: {sum(p.numel() for p in model.parameters()) * 1e-6:.2f}M")

    bind_nsml(model, tokenizer, args)
    if args.pause:
        nsml.paused(scope=locals())

    if args.mode == "train":
        noisy_sents = read_strings(os.path.join(args.data_dir, "train_data", "train_data"))
        annotations = read_strings(os.path.join(args.data_dir, "train_data", "train_annotation"))
        corpuses = read_strings(os.path.join(args.data_dir, "train_data", "train_corpus"))
        clean_sents = read_strings(os.path.join(args.data_dir, "train_label"))

        print(f"noisy_sents length is {len(noisy_sents)}")
        print(f"annotations length is {len(annotations)}")
        print(f"clean_sents length is {len(clean_sents)}")
        print(f"corpuses length is {len(corpuses)}")

        for (i, noisy_sent) in enumerate(noisy_sents[:5]):
            print(f"{i + 1}'s noisy sentence is {noisy_sent}")

        for (i, clean_sent) in enumerate(clean_sents[:5]):
            print(f"{i + 1}'s clean sentence is {clean_sent}")

        for (i, annotation) in enumerate(annotations[:5]):
            print(f"{i + 1}'s annotation is {annotation}")

        pairs = [{"noisy": noisy, "clean": clean} for noisy, clean in zip(noisy_sents, clean_sents)]

        if args.train_all:
            train_data = pairs
            valid_data = None
            logger.info(f"# of train data: {len(train_data)}")
        else:
            train_data, valid_data = pairs[:-int(len(pairs) * args.val_ratio)], pairs[-int(len(pairs) * args.val_ratio):]
            logger.info(f"# of train data: {len(train_data)}")
            logger.info(f"# of valid data: {len(valid_data)}")

        print(f"train_data is {train_data[0]}")

        train_sents = [x['noisy'] for x in train_data] + [x['clean'] for x in train_data]

        print(f"first train_sentences is {train_sents[0]}")
        print(f"last train_sentences is {train_sents[-1]}")

        if args.tokenizer == 'char':
            tokenizer = CharTokenizer.from_strings(train_sents, args.vocab_size)
        elif args.tokenizer == 'pre_char':
            tokenizer = CharTokenizer.from_strings(corpuses + clean_sents + noisy_sents, args.vocab_size)
        elif args.tokenizer == 'kfold_char':
            tokenizer = CharTokenizer.from_strings(clean_sents + noisy_sents, args.vocab_size)
        elif args.tokenizer == 'wordpiece':
            tokenizer.get_tokenizer()
            print(f"vocab size is {len(tokenizer)}")
            args.vocab_size = len(tokenizer)

        bind_nsml(model, tokenizer, args)

    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model, dim=1)

    if args.mode == "train":
        if not args.kfold:
            train(model, tokenizer, train_data, valid_data, corpuses, args)
        else:
            kfold = KFold(n_splits=args.n_splits)
            for index, (train_idx, valid_idx) in enumerate(kfold.split(range(len(pairs)))):
                train_data, valid_data = [pairs[idx] for idx in train_idx], [pairs[idx] for idx in valid_idx]
                train(model, tokenizer, train_data, valid_data, corpuses, args, index)

if __name__ == "__main__":
    main()


"""
/*
 * @Author: Yang Zhong
 * @Date: 2021-10-09 13:46:53
 * @Last Modified by: Yang Zhong
 * @Last Modified time: 2021-10-29 21:09:02
 */
"""
import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.optim as opt
from typing import List, Dict, Union
from torch.nn import functional as F
from .utils import scatter_plot
import numpy as np
import os
import pandas as pd

class Model(pl.LightningModule):
    def __init__(
            self,
            representation: nn.Module,
            output: nn.Module,
            losses: List[Dict],
            validation_metrics: List[Dict],
            lr: float = 1e-3,
            lr_decay: float = 0.1,
            lr_patience: int = 100,
            lr_monitor="training/total_loss",
            epsilon: float = 1e-8,
            beta1: float = 0.99,
            beta2: float = 0.999,
            amsgrad: bool = True,
            max_points_to_scatter: int = 100000,
            post_processing: callable = None
            ):
        super().__init__()

        self.representation = representation
        self.output_module = output

        self.losses = losses
        self.metrics = validation_metrics

        self.lr = lr
        self.lr_decay = lr_decay
        self.lr_patience = lr_patience
        self.lr_monitor = lr_monitor

        self.epsilon = epsilon
        self.beta1 = beta1
        self.beta2 = beta2
        self.amsgrad = amsgrad

        self.max_points_to_scatter = max_points_to_scatter
        # post_processing is used to calculate some physical quantities that rely on gradient backpropagation
        self.post_processing = post_processing

        #self.save_hyperparameters()

        # For gradients
        self.requires_dr = self.output_module.derivative

    def calculate_loss(self, batch, result, mode):
        loss = torch.tensor(0.0, device=self.device)
        for loss_dict in self.losses:
            loss_fn = loss_dict["metric"]

            if "target" in loss_dict.keys():
                pred = result[loss_dict["prediction"]]
                target = batch[loss_dict["target"]]

                loss_i = loss_fn(pred, target)
            else:
                loss_i = loss_fn(result[loss_dict["prediction"]])
            loss += loss_dict["loss_weight"] * loss_i

            if hasattr(loss_fn, "name"):
                lossname = loss_fn.name
            else:
                lossname = type(loss_fn).__name__.split(".")[-1]

            self.log(
                mode
                + "/"
                + lossname
                + "_"
                + loss_dict["prediction"],
                loss_i,
                on_step=False,
                on_epoch=True,
            )

        return loss

    def training_step(self, data, batch_idx):
        self._enable_grads(data)
        pred = self(data)
        loss = self.calculate_loss(data, pred, 'training')
        self.log("training/total_loss", loss, on_step=False, on_epoch=True)
        # self.check_param()
        return loss

    def validation_step(self, data, batch_idx):
        if self.requires_dr:
            torch.set_grad_enabled(True)
        else:
            torch.set_grad_enabled(False)
        self._enable_grads(data)
        pred = self(data)
        val_loss = self.calculate_loss(
            data, pred, 'validation').detach().item()
        self.log("validation/total_loss", val_loss,
                 on_step=False, on_epoch=True)
        self.log_metrics(data, pred, 'validation')
        outputs_pred, outputs_target = {}, {}
        for loss_dict in self.losses:
            outputs_pred[loss_dict["prediction"]] = pred[loss_dict["prediction"]].detach().cpu().numpy()
            outputs_target[loss_dict["target"]] = data[loss_dict["target"]].detach().cpu().numpy()
        return {'pred': outputs_pred, 'target': outputs_target}

    def validation_epoch_end(self, validation_step_outputs):
        for loss_dict in self.losses:
            if "target" in loss_dict.keys():
                pred = np.concatenate([out['pred'][loss_dict["prediction"]]
                                 for out in validation_step_outputs])
                target = np.concatenate([out['target'][loss_dict["target"]]
                                    for out in validation_step_outputs])
                if (pred.dtype == np.complex64) and (target.dtype == np.complex64):
                    lossname = type(loss_dict['metric']).__name__.split(".")[-1]
                    if lossname.lower() == 'abs_mae':
                        pred = np.absolute(pred)
                        target = np.absolute(target)
                    else:
                        pred = np.concatenate([pred.real, pred.imag], axis=-1)
                        target = np.concatenate([target.real, target.imag], axis=-1)
                # Control the number of scatter points to plot
                if pred.size > self.max_points_to_scatter:
                    random_state = np.random.RandomState(seed=42)
                    perm = list(random_state.permutation(np.arange(pred.size)))
                    pred = pred.reshape(-1)[perm[:self.max_points_to_scatter]]
                    target = target.reshape(-1)[perm[:self.max_points_to_scatter]]
                figure = scatter_plot(pred.reshape(-1), target.reshape(-1))
                figname = 'PredVSTarget_' + loss_dict['prediction']
                self.logger.experiment.add_figure(
                    'validation/'+figname, figure, global_step=self.global_step)
            else:
                pass

    def test_step(self, data, batch_idx):
        if self.requires_dr:
            torch.set_grad_enabled(True)
        else:
            torch.set_grad_enabled(False)
        self._enable_grads(data)

        if self.post_processing is not None:
            pred = self.post_processing(data)
            if type(self.post_processing).__name__.split(".")[-1].lower() == 'epc_output':
                proessed_values = {'epc_mat': pred['epc_mat'].detach().cpu().numpy()}
            else:
                raise NotImplementedError
        else:
            pred = self(data)
            proessed_values = None

        loss = self.calculate_loss(data, pred, 'test').detach().item()
        self.log("test/total_loss", loss, on_step=False, on_epoch=True)
        self.log_metrics(data, pred, "test")
        outputs_pred, outputs_target = {}, {}
        for loss_dict in self.losses:
            outputs_pred[loss_dict["prediction"]] = pred[loss_dict["prediction"]].detach().cpu().numpy()
            outputs_target[loss_dict["target"]] = data[loss_dict["target"]].detach().cpu().numpy()
        return {'pred': outputs_pred, 'target': outputs_target, 'processed_values': proessed_values}

    def test_epoch_end(self, test_step_outputs):
        for loss_dict in self.losses:
            if "target" in loss_dict.keys():
                pred = np.concatenate([out['pred'][loss_dict["prediction"]]
                                 for out in test_step_outputs])
                target = np.concatenate([out['target'][loss_dict["target"]]
                                    for out in test_step_outputs])

                if not os.path.exists(self.trainer.logger.log_dir):
                    os.makedirs(self.trainer.logger.log_dir)

                np.save(os.path.join(
                    self.trainer.logger.log_dir, 'prediction_'+loss_dict["prediction"]+'.npy'), pred)
                np.save(os.path.join(self.trainer.logger.log_dir,
                        'target_'+loss_dict["target"]+'.npy'), target)

                # plot
                if (pred.dtype == np.complex64) and (target.dtype == np.complex64):
                    lossname = type(loss_dict['metric']).__name__.split(".")[-1]
                    if lossname.lower() == 'abs_mae':
                        pred = np.absolute(pred)
                        target = np.absolute(target)
                    else:
                        pred = np.concatenate([pred.real, pred.imag], axis=-1)
                        target = np.concatenate([target.real, target.imag], axis=-1)

                # Control the number of scatter points to plot
                if pred.size > self.max_points_to_scatter:
                    random_state = np.random.RandomState(seed=42)
                    perm = list(random_state.permutation(np.arange(pred.size)))
                    pred = pred.reshape(-1)[perm[:self.max_points_to_scatter]]
                    target = target.reshape(-1)[perm[:self.max_points_to_scatter]]

                figure = scatter_plot(pred.reshape(-1), target.reshape(-1))
                figname = 'PredVSTarget_' + loss_dict['prediction']
                self.logger.experiment.add_figure(
                    'test/'+figname, figure, global_step=self.global_step)
            else:
                pass

        if self.post_processing is not None:
            if type(self.post_processing).__name__.split(".")[-1].lower() == 'epc_output':
                processed_values = np.concatenate([out['processed_values']["epc_mat"]
                                        for out in test_step_outputs])
                np.save(os.path.join(
                    self.trainer.logger.log_dir, 'processed_values_'+'epc_mat'+'.npy'), processed_values)

    def forward(self, data):
        # torch.set_grad_enabled(True)
        self._enable_grads(data)
        representation = self.representation(data)
        pred = self.output_module(data, representation)
        return pred

    def log_metrics(self, batch, result, mode):
        for metric_dict in self.metrics:
            loss_fn = metric_dict["metric"]

            if "target" in metric_dict.keys():
                pred = result[metric_dict["prediction"]]

                target = batch[metric_dict["target"]]

                loss_i = loss_fn(
                    pred, target
                ).detach().item()
            else:
                loss_i = loss_fn(
                    result[metric_dict["prediction"]]).detach().item()

            if hasattr(loss_fn, "name"):
                lossname = loss_fn.name
            else:
                lossname = type(loss_fn).__name__.split(".")[-1]

            self.log(
                mode
                + "/"
                + lossname
                + "_"
                + metric_dict["prediction"],
                loss_i,
                on_step=False,
                on_epoch=True,
            )

    def configure_optimizers(
            self,
    ):
        optimizer = opt.AdamW(self.parameters(), lr=self.lr, eps=self.epsilon, betas=(self.beta1, self.beta2), weight_decay=0.0, amsgrad=True)
        scheduler = {
            "scheduler": opt.lr_scheduler.ReduceLROnPlateau(
                optimizer,
                factor=self.lr_decay,
                patience=self.lr_patience,
                threshold=1e-6,
                cooldown=self.lr_patience // 2,
                min_lr=1e-6,
            ),
            "monitor": self.lr_monitor,
            "interval": "epoch",
            "frequency": 1,
            "strict": True,
        }
        return [optimizer], [scheduler]

    def _enable_grads(self, data):
        if self.requires_dr:
            data.pos.requires_grad_()

    def check_param(self):
        for name, parms in self.named_parameters():
            print('-->name:', name, '-->grad_requirs:', parms.requires_grad,
                  '-->grad_value:', parms.grad)


import os
import sys
import argparse
import numpy as np
from datetime import datetime

import wandb
import torch
from torch import nn
import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger

user = os.getlogin()
sys.path.append(f"/home/{user}/workspace")
from utils import yaml_read  # noqa: E402
from brats18_project.data_utils import Brats18DataModule  # noqa: E402
from brats18_project.models import SegModel, ClassifierModel  # noqa: E402

sys.path.append(f"/home/doronser/workspace/MedicalZooPytorch/")
from lib.losses3D import DiceLoss  # noqa: E402
from torch.nn import CrossEntropyLoss  # noqa: E402
# import lib.medzoo as md  # noqa: E402
import brats18_project.models as model_zoo  # noqa: E402

# for reproducibility
torch.manual_seed(42)
np.random.seed(42)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--cfg", type=str, help="path to YAML config file")
    parser.add_argument("-d", "--debug", action="store_true", default=False)
    args = parser.parse_args()

    assert os.path.exists(args.cfg), f"cfg file {args.cfg} does not exist!"
    cfg = yaml_read(args.cfg, easy_dict=True)

    data_module = Brats18DataModule(cfg.data)

    model_cls = getattr(model_zoo, cfg.model.name)
    net = model_cls(**cfg.model.kwargs)

    if cfg.loss.name == 'DiceLoss':
        model = SegModel(net=net, criterion=DiceLoss(classes=4),
                         optimizer_params=cfg.optimizer, scheduler_params=cfg.scheduler)
    elif cfg.loss.name == 'CrossEntropyLoss':
        model = ClassifierModel(net=net, criterion=nn.CrossEntropyLoss(),
                                optimizer_params=cfg.optimizer, scheduler_params=cfg.scheduler)
    elif cfg.loss.name == 'BarlowTwinsLoss':
        model = model_zoo.BarlowTwins(encoder=net, batch_size=cfg.data.batch_size,
                                      learning_rate=cfg.optimizer.lr, num_training_samples=257,
                                      max_epochs=cfg.trainer.epochs, lambda_coeff=cfg.loss.lambda_coeff)
    else:
        raise ValueError("Unknown loss")

    if args.debug:
        trainer = pl.Trainer(fast_dev_run=True)
    else:
        # define pl loggers and callbacks
        time_suffix = datetime.now().strftime(r"%y%m%d_%H%M%S")
        wandb_logger = WandbLogger(project="intro2dl_brats2018", name=f"{cfg.name}_{time_suffix}")
        wandb.config.update(cfg)
        ckpt_clbk = pl.callbacks.ModelCheckpoint(dirpath=os.path.join(cfg.ckpt_path, f"{cfg.name}_{time_suffix}"),
                                                 filename=cfg.name + '_epoch{epoch:02d}', auto_insert_metric_name=False,
                                                 save_top_k=-1, every_n_epochs=10)
        lr_clbk = pl.callbacks.LearningRateMonitor(logging_interval='step')
        trainer = pl.Trainer(
            accelerator='gpu',
            devices=[cfg.trainer.gpu],
            max_epochs=cfg.trainer.epochs,
            callbacks=[ckpt_clbk, lr_clbk],
            logger=wandb_logger
        )
    trainer.fit(model=model, datamodule=data_module)


# **************************************************************************
# *
# * Authors:     Yaiza Rancel (cyrancel@cnb.csic.es) [1]
# *              Pablo Conesa (pconse@cnb.csic.es) [1]
# *              J.M. De la Rosa Trevin (delarosatrevin@scilifelab.se) [2]
# *
# * [1] Unidad de  Bioinformatica of Centro Nacional de Biotecnologia , CSIC
# * [2] SciLifeLab, Stockholm University
# *
# * This program is free software; you can redistribute it and/or modify
# * it under the terms of the GNU General Public License as published by
# * the Free Software Foundation; either version 2 of the License, or
# * (at your option) any later version.
# *
# * This program is distributed in the hope that it will be useful,
# * but WITHOUT ANY WARRANTY; without even the implied warranty of
# * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# * GNU General Public License for more details.
# *
# * You should have received a copy of the GNU General Public License
# * along with this program; if not, write to the Free Software
# * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
# * 02111-1307  USA
# *
# *  All comments concerning this program package may be sent to the
# *  e-mail address 'scipion@cnb.csic.es'
# *
# **************************************************************************
import os
import importlib
import pkgutil
import inspect
import pkg_resources
from email import message_from_string
from collections import OrderedDict
from abc import ABCMeta, abstractmethod

import pyworkflow.utils as pwutils
from pyworkflow.install import Environment

class Domain:
    """
    Class to represent the application domain.
    It will allow to specify new objects, protocols, viewers and wizards
    through the registration of new plugins.
    """

    # The following classes should be defined in subclasses of Domain
    _name = None
    _protocolClass = None
    _objectClass = None
    _viewerClass = None
    _wizardClass = None
    _baseClasses = {}  # Update this with the base classes of the Domain

    # Dictionaries to store different type of objects
    _plugins = {}
    _protocols = {}
    _objects = {}
    _viewers = {}
    _wizards = {}

    @classmethod
    def __getSubmodule(cls, name, subname):
        try:
            return importlib.import_module('%s.%s' % (name, subname))
        except Exception as e:
            msg = str(e)
            # FIXME: The following is a quick and dirty way to filter
            # when the submodule is not present
            if msg != 'No module named %s' % subname:
                pwutils.pluginNotFound("%s.%s" % (name, subname), msg)
            return None

    @classmethod
    def registerPlugin(cls, name):
        """ Register a new plugin. This function should only be called
        when creating a class with __metaclass__=PluginMeta that will
        trigger this.
        """
        m = importlib.import_module(name)
        cls._plugins[name] = m  # Register the name to as a plugin
        # TODO: Load subclasses (protocols, viewers, wizards)

        # Define variables
        m.Plugin._defineVariables()

        # Load bibtex
        m._bibtex = {}
        bib = cls.__getSubmodule(name, 'bibtex')
        if bib is not None:
            try:
                m._bibtex = pwutils.parseBibTex(bib.__doc__)
            except Exception:
                pass

        return m

    @classmethod
    def __isPlugin(cls, m):
        """ Return True if the module is a Scipion plugin. """
        return m.__name__ in cls._plugins

    @classmethod
    def getPlugins(cls):
        """ Return existing plugins for this Domain. """
        loaded = getattr(cls, '_pluginsLoaded', False)
        if not loaded:  # Load plugin only once
            for p, name, isModule in pkgutil.iter_modules():
                if isModule:
                    try:
                        importlib.import_module(name)
                        # NOTE: After importing the modules they will
                        # automatically being register if they has
                        # called Domain.registerPlugin(__name__)
                    except Exception:
                        pass
            cls._pluginsLoaded = True
        return dict(cls._plugins)

    @classmethod
    def getPlugin(cls, name):
        """ Load a given plugin name. """
        m = importlib.import_module(name)

        if not cls.__isPlugin(m):
            raise Exception("Invalid plugin '%s'. "
                            "Class Plugin with __metaclass__=PluginMeta "
                            "not found" % name)
        return m

    @classmethod
    def refreshPlugin(cls, name):
        """ Refresh a given plugin name. """
        import types

        plugin = cls.getPlugin(name)
        if plugin is not None:
            fn = plugin.__file__
            fn_dir = os.path.dirname(fn) + os.sep
            module_visit = {plugin}
            module_visit_path = {fn}
            del fn

            def get_recursive_molules(module):
                """
                Get all plugin modules recursively
                """
                for module_child in vars(module).values():
                    if isinstance(module_child, types.ModuleType):
                        fn_child = getattr(module_child, "__file__", None)
                        if (fn_child is not None) and fn_child.startswith(
                                fn_dir):
                            if fn_child not in module_visit_path:
                                module_visit.add(module_child)
                                module_visit_path.add(fn_child)
                                get_recursive_molules(module_child)

            get_recursive_molules(plugin)
            # reload all plugin modules

            while module_visit:
                for module in module_visit:
                    try:
                        reload(module)
                        module_visit.remove(module)
                        break
                    except Exception as ex:
                        pass

    @classmethod
    def __getSubclasses(cls, submoduleName, BaseClass,
                        updateBaseClasses=False):
        """ Load all detected subclasses of a given BaseClass.
        Params:
            updateBaseClasses: if True, it will try to load classes from the
                Domain submodule that was not imported as globals()
        """
        subclasses = getattr(cls, '_%s' % submoduleName)

        if not subclasses:  # Only discover subclasses once
            if updateBaseClasses:
                sub = cls.__getSubmodule(cls.getName(), submoduleName)
                if sub is not None:
                    for name in cls.getModuleClasses(sub):
                        attr = getattr(sub, name)
                        if inspect.isclass(attr) and issubclass(attr, BaseClass):
                            cls._baseClasses[name] = attr

            for pluginName, plugin in cls.getPlugins().iteritems():
                sub = cls.__getSubmodule(pluginName, submoduleName)
                if sub is not None:
                    for name in cls.getModuleClasses(sub):
                        attr = getattr(sub, name)
                        if inspect.isclass(attr) and issubclass(attr, BaseClass):

                            # Check if the class already exists (to prevent
                            # naming collisions)
                            if name in subclasses:
                                # Get already added class plugin
                                pluginCollision = subclasses[name]._package.__name__
                                print("ERROR: Name collision (%s) detected "
                                      "while discovering %s.%s.\n"
                                      " It conflicts with %s" %
                                      (name, pluginName, submoduleName,
                                       pluginCollision))
                            else:
                                # Set this special property used by Scipion
                                attr._package = plugin
                                subclasses[name] = attr
            subclasses.update(
                pwutils.getSubclasses(BaseClass, cls._baseClasses))

        return subclasses

    @classmethod
    def getModuleClasses(cls, module):

        # Dir was used before but dir returns all imported elements
        # included those imported to be BaseClasses.
        # return dir(module)

        # Get the module name
        moduleName = module.__name__

        # Get any module class
        for name, declaredClass in inspect.getmembers(module, inspect.isclass):
            if moduleName in declaredClass.__module__:
                yield name

    @classmethod
    def getProtocols(cls):
        """ Return all Protocol subclasses from all plugins for this domain.
        """
        return cls.__getSubclasses('protocols', cls._protocolClass)

    @classmethod
    def getObjects(cls):
        """ Return all EMObject subclasses from all plugins for this domain.
        """
        return cls.__getSubclasses('objects', cls._objectClass)

    @classmethod
    def getViewers(cls):
        """ Return all Viewer subclasses from all plugins for this domain.
        """
        return cls.__getSubclasses('viewers', cls._viewerClass,
                                   updateBaseClasses=True)

    @classmethod
    def getWizards(cls):
        """ Return all Wizard subclasses from all plugins for this domain.
        """
        return cls.__getSubclasses('wizards', cls._wizardClass)

    @classmethod
    def getName(cls):
        return cls._name

class Plugin:
    __metaclass__ = ABCMeta

    _vars = {}
    _homeVar = ''  # Change in subclasses to define the "home" variable
    _pathVars = []
    _supportedVersions = []

    @classmethod
    def _defineVar(cls, varName, defaultValue):
        """ Internal method to define variables from the environment. """
        cls._vars[varName] = os.environ.get(varName, defaultValue)

    @classmethod
    def _defineEmVar(cls, varName, defaultValue):
        """ Shortcut method to define variables by prepending EM_ROOT
        to the default value.
        """
        cls._defineVar(varName,
                       os.path.join(Environment.getEm(defaultValue)))

    @classmethod
    @abstractmethod
    def getEnviron(cls):
        """ Setup the environment variables needed to launch programs. """
        pass

    @classmethod
    @abstractmethod
    def _defineVariables(cls):
        """ Method to define variables and their default values.
        It will use the method _defineVar that will take a variable value
        from the environment or from an optional default value.

        This method is not supposed to be called from client code,
        except from the Domain class when registering a Plugin.
        """
        pass

    @classmethod
    @abstractmethod
    def defineBinaries(cls, env):
        """ Define required binaries in the given Environment. """
        pass

    @classmethod
    def getVar(cls, varName, defaultValue=None):
        """ Return the value of a given variable. """
        return cls._vars.get(varName, defaultValue)

    @classmethod
    def getVars(cls):
        """ Return the value of a given variable. """
        return cls._vars

    @classmethod
    def getHome(cls, *paths):
        """ Return a path from the "home" of the package
         if the _homeVar is defined in the plugin. """
        home = cls.getVar(cls._homeVar)
        return os.path.join(home, *paths) if home else ''

    @classmethod
    def getSupportedVersions(cls):
        """ Return the list of supported binary versions. """
        return cls._supportedVersions

    @classmethod
    def getActiveVersion(cls, home=None, versions=None):
        """ Return the version of the binaries that are currently active.
        In the current implementation it will be inferred from the *_HOME
        variable, so it should contain the version number in it. """
        # FIXME: (JMRT) Use the basename might aleviate the issue with matching
        # the binaries version, but we might consider to find a better solution
        home = os.path.basename(home or cls.getHome())
        versions = versions or cls.getSupportedVersions()

        for v in versions:
            if v in home:
                return v

        return ''

    @classmethod
    def validateInstallation(cls):
        """
        Check if the binaries are properly installed and if not, return
        a list with the error messages.

        The default implementation will check if the _pathVars exists.
        """
        missing = ["%s: %s" % (var, cls.getVar(var))
                   for var in cls._pathVars if not os.path.exists(cls.getVar(var))]

        return (["Missing variables:"] + missing) if missing else []

class PluginInfo:
    """
    Information related to a given plugin when it is installed via PIP
    """
    def __init__(self, name):
        try:
            dist = pkg_resources.get_distribution(name)
            lines = [l for l in dist._get_metadata(dist.PKG_INFO)]
            tuples = message_from_string('\n'.join(lines))

        except Exception:
            print("Plugin %s seems is not a pip module yet. "
                  "No metadata found" % name)
            tuples = message_from_string('Author: plugin in development mode?')

        self._name = name
        self._metadata = OrderedDict()

        for v in tuples.items():
            if v[0] == 'Keywords':
                break
            self._metadata[v[0]] = v[1]

    def getAuthor(self):
        return self._metadata.get('Author', "")

    def getAuthorEmail(self):
        return self._metadata.get('Author-email', '')

    def getHomePage(self):
        return self._metadata.get('Home-page', '')

    def getKeywords(self):
        return self._metadata.get('Keywords', '')


import pytest
from pygwidget.mixins import TrigonometryMixin

class TestTrigonometryMixin(object):
    mixin = TrigonometryMixin()

    @pytest.mark.parametrize('start_degree, stop_degree, expected', [
        (0, 90, 90),
        (0, 180, 180),
        (150, 180, 30),
        (330, 90, 120),
        (330, 150, 180),
        (300, 150, 210),
        (315, 225, 270),
        (330, 30, 60),
        (330, 330, 0),
        (330, 315, 345),
        (330, 360, 30),
        (330, 0, 30),
    ])
    def test_degree_distance(self, start_degree, stop_degree, expected):
        result = self.mixin.anti_clockwise_degree_distance(
            start_degree, stop_degree)
        assert expected == result

    @pytest.mark.parametrize(
        'start_degree, stop_degree, scales, scale, expected', [
            (0,  180, [0, 1, 2], 1, 90),
            (0, 180, range(0, 100, 10), 45, 90),
            (330, 180, range(0, 100, 10), 90, 330),
            (45, 135, [0, 1, 2], 1, 90),
            (330, 210, range(0, 9), 1, 180),
            (325, 215, range(0, 210, 20), 100, 90)
        ])
    def test_find_scale_degree(self, start_degree, stop_degree, scales, scale,
                               expected):
        result = self.mixin.find_scale_degree(
            start_degree, stop_degree, scales, scale)
        assert expected == result

    @pytest.mark.parametrize('degree, expected', [
        (90, 0),
        (0, -90),
        (180, 90),
        (45, -45),
        (135, 45),
        (225, 135),
        (270, -180),
        (315, -135),
        (-45, -135),
        (330, -120),
        (-30, -120)
    ])
    def test_find_rotation_angle(self, degree, expected):
        assert self.mixin.find_rotation_angle(degree) == expected


from django import forms
from . import models

class Rsvp_Form(forms.ModelForm):

    class Meta:
        model = models.rsvp
        fields = [
            'full_name',
            'email',
            'please_check_if_you_are_able_to_attend',
            'number_attending',
            'message_for_the_happy_couple'
        ]


"""
CPU implementation of backend methods for water condensation/evaporation
"""
import math
from functools import lru_cache

import numba
import numpy as np

from PySDM.backends.impl_common.backend_methods import BackendMethods
from PySDM.backends.impl_numba import conf
from PySDM.backends.impl_numba.toms748 import toms748_solve
from PySDM.backends.impl_numba.warnings import warn

class CondensationMethods(BackendMethods):
    # pylint: disable=unused-argument,too-many-locals
    @staticmethod
    def condensation(
        *,
        solver,
        n_cell,
        cell_start_arg,
        water_mass,
        v_cr,
        multiplicity,
        vdry,
        idx,
        rhod,
        thd,
        water_vapour_mixing_ratio,
        dv,
        prhod,
        pthd,
        predicted_water_vapour_mixing_ratio,
        kappa,
        f_org,
        rtol_x,
        rtol_thd,
        timestep,
        counters,
        cell_order,
        RH_max,
        success,
        cell_id,
    ):
        n_threads = min(numba.get_num_threads(), n_cell)
        CondensationMethods._condensation(
            solver=solver,
            n_threads=n_threads,
            n_cell=n_cell,
            cell_start_arg=cell_start_arg.data,
            water_mass=water_mass.data,
            v_cr=v_cr.data,
            multiplicity=multiplicity.data,
            vdry=vdry.data,
            idx=idx.data,
            rhod=rhod.data,
            thd=thd.data,
            water_vapour_mixing_ratio=water_vapour_mixing_ratio.data,
            dv_mean=dv,
            prhod=prhod.data,
            pthd=pthd.data,
            predicted_water_vapour_mixing_ratio=predicted_water_vapour_mixing_ratio.data,
            kappa=kappa.data,
            f_org=f_org.data,
            rtol_x=rtol_x,
            rtol_thd=rtol_thd,
            timestep=timestep,
            counter_n_substeps=counters["n_substeps"].data,
            counter_n_activating=counters["n_activating"].data,
            counter_n_deactivating=counters["n_deactivating"].data,
            counter_n_ripening=counters["n_ripening"].data,
            cell_order=cell_order,
            RH_max=RH_max.data,
            success=success.data,
        )

    @staticmethod
    @numba.njit(**{**conf.JIT_FLAGS, **{"cache": False}})
    def _condensation(
        *,
        solver,
        n_threads,
        n_cell,
        cell_start_arg,
        water_mass,
        v_cr,
        multiplicity,
        vdry,
        idx,
        rhod,
        thd,
        water_vapour_mixing_ratio,
        dv_mean,
        prhod,
        pthd,
        predicted_water_vapour_mixing_ratio,
        kappa,
        f_org,
        rtol_x,
        rtol_thd,
        timestep,
        counter_n_substeps,
        counter_n_activating,
        counter_n_deactivating,
        counter_n_ripening,
        cell_order,
        RH_max,
        success,
    ):
        # pylint: disable=too-many-locals
        for thread_id in numba.prange(n_threads):  # pylint: disable=not-an-iterable
            for i in range(thread_id, n_cell, n_threads):
                cell_id = cell_order[i]

                cell_start = cell_start_arg[cell_id]
                cell_end = cell_start_arg[cell_id + 1]
                n_sd_in_cell = cell_end - cell_start
                if n_sd_in_cell == 0:
                    continue

                dthd_dt = (pthd[cell_id] - thd[cell_id]) / timestep
                d_water_vapour_mixing_ratio__dt = (
                    predicted_water_vapour_mixing_ratio[cell_id]
                    - water_vapour_mixing_ratio[cell_id]
                ) / timestep
                drhod_dt = (prhod[cell_id] - rhod[cell_id]) / timestep
                md = (prhod[cell_id] + rhod[cell_id]) / 2 * dv_mean

                (
                    success_in_cell,
                    water_vapour_mixing_ratio_new,
                    thd_new,
                    substeps_hint,
                    n_activating,
                    n_deactivating,
                    n_ripening,
                    RH_max_in_cell,
                ) = solver(
                    water_mass,
                    v_cr,
                    multiplicity,
                    vdry,
                    idx[cell_start:cell_end],
                    kappa,
                    f_org,
                    thd[cell_id],
                    water_vapour_mixing_ratio[cell_id],
                    rhod[cell_id],
                    dthd_dt,
                    d_water_vapour_mixing_ratio__dt,
                    drhod_dt,
                    md,
                    rtol_x,
                    rtol_thd,
                    timestep,
                    counter_n_substeps[cell_id],
                )
                counter_n_substeps[cell_id] = substeps_hint
                counter_n_activating[cell_id] = n_activating
                counter_n_deactivating[cell_id] = n_deactivating
                counter_n_ripening[cell_id] = n_ripening
                RH_max[cell_id] = RH_max_in_cell
                success[cell_id] = success_in_cell
                predicted_water_vapour_mixing_ratio[
                    cell_id
                ] = water_vapour_mixing_ratio_new
                pthd[cell_id] = thd_new

    @staticmethod
    def make_adapt_substeps(
        *, jit_flags, timestep, step_fake, dt_range, fuse, multiplier, within_tolerance
    ):
        if not isinstance(multiplier, int):
            raise ValueError()
        if dt_range[1] > timestep:
            dt_range = (dt_range[0], timestep)
        if dt_range[0] == 0:
            raise NotImplementedError()
        n_substeps_max = math.floor(timestep / dt_range[0])
        n_substeps_min = math.ceil(timestep / dt_range[1])

        @numba.njit(**jit_flags)
        def adapt_substeps(args, n_substeps, thd, rtol_thd):
            n_substeps = np.maximum(n_substeps_min, n_substeps // multiplier)
            success = False
            for burnout in range(fuse + 1):
                if burnout == fuse:
                    return warn(
                        "burnout (long)",
                        __file__,
                        context=(
                            "thd",
                            thd,
                        ),
                        return_value=(0, False),
                    )
                thd_new_long, success = step_fake(args, timestep, n_substeps)
                if success:
                    break
                n_substeps *= multiplier
            for burnout in range(fuse + 1):
                if burnout == fuse:
                    return warn("burnout (short)", __file__, return_value=(0, False))
                thd_new_short, success = step_fake(
                    args, timestep, n_substeps * multiplier
                )
                if not success:
                    return warn("short failed", __file__, return_value=(0, False))
                dthd_long = thd_new_long - thd
                dthd_short = thd_new_short - thd
                error_estimate = np.abs(dthd_long - multiplier * dthd_short)
                thd_new_long = thd_new_short
                if within_tolerance(error_estimate, thd, rtol_thd):
                    break
                n_substeps *= multiplier
                if n_substeps > n_substeps_max:
                    break
            return np.minimum(n_substeps_max, n_substeps), success

        return adapt_substeps

    @staticmethod
    def make_step_fake(jit_flags, step_impl):
        @numba.njit(**jit_flags)
        def step_fake(args, dt, n_substeps):
            dt /= n_substeps
            _, thd_new, _, _, _, _, success = step_impl(*args, dt, 1, True)
            return thd_new, success

        return step_fake

    @staticmethod
    def make_step(jit_flags, step_impl):
        @numba.njit(**jit_flags)
        def step(args, dt, n_substeps):
            return step_impl(*args, dt, n_substeps, False)

        return step

    @staticmethod
    def make_step_impl(
        *,
        jit_flags,
        phys_pvs_C,
        phys_lv,
        calculate_ml_old,
        calculate_ml_new,
        phys_T,
        phys_p,
        phys_pv,
        phys_dthd_dt,
        phys_D,
        phys_K,
        const,
    ):
        @numba.njit(**jit_flags)
        def step_impl(  # pylint: disable=too-many-arguments,too-many-locals
            water_mass,
            v_cr,
            multiplicity,
            vdry,
            cell_idx,
            kappa,
            f_org,
            thd,
            water_vapour_mixing_ratio,
            rhod,
            dthd_dt_pred,
            d_water_vapour_mixing_ratio__dt_predicted,
            drhod_dt,
            m_d,
            rtol_x,
            timestep,
            n_substeps,
            fake,
        ):
            timestep /= n_substeps
            ml_old = calculate_ml_old(water_mass, multiplicity, cell_idx)
            count_activating, count_deactivating, count_ripening = 0, 0, 0
            RH_max = 0
            success = True
            for _ in range(n_substeps):
                # note: no example yet showing that the trapezoidal scheme brings any improvement
                thd += timestep * dthd_dt_pred / 2
                water_vapour_mixing_ratio += (
                    timestep * d_water_vapour_mixing_ratio__dt_predicted / 2
                )
                rhod += timestep * drhod_dt / 2

                T = phys_T(rhod, thd)
                p = phys_p(rhod, T, water_vapour_mixing_ratio)
                pv = phys_pv(p, water_vapour_mixing_ratio)
                lv = phys_lv(T)
                pvs = phys_pvs_C(T - const.T0)
                RH = pv / pvs
                DTp = phys_D(T, p)
                KTp = phys_K(T, p)
                (
                    ml_new,
                    success_within_substep,
                    n_activating,
                    n_deactivating,
                    n_ripening,
                ) = calculate_ml_new(
                    timestep,
                    fake,
                    T,
                    p,
                    RH,
                    water_mass,
                    v_cr,
                    multiplicity,
                    vdry,
                    cell_idx,
                    kappa,
                    f_org,
                    lv,
                    pvs,
                    DTp,
                    KTp,
                    rtol_x,
                )
                dml_dt = (ml_new - ml_old) / timestep
                d_water_vapour_mixing_ratio__dt_corrected = -dml_dt / m_d
                dthd_dt_corr = phys_dthd_dt(
                    rhod=rhod,
                    thd=thd,
                    T=T,
                    d_water_vapour_mixing_ratio__dt=d_water_vapour_mixing_ratio__dt_corrected,
                    lv=lv,
                )

                thd += timestep * (dthd_dt_pred / 2 + dthd_dt_corr)
                water_vapour_mixing_ratio += timestep * (
                    d_water_vapour_mixing_ratio__dt_predicted / 2
                    + d_water_vapour_mixing_ratio__dt_corrected
                )
                rhod += timestep * drhod_dt / 2
                ml_old = ml_new
                count_activating += n_activating
                count_deactivating += n_deactivating
                count_ripening += n_ripening
                RH_max = max(RH_max, RH)
                success = success and success_within_substep
            return (
                water_vapour_mixing_ratio,
                thd,
                count_activating,
                count_deactivating,
                count_ripening,
                RH_max,
                success,
            )

        return step_impl

    @staticmethod
    def make_calculate_ml_old(jit_flags):
        @numba.njit(**jit_flags)
        def calculate_ml_old(water_mass, multiplicity, cell_idx):
            result = 0
            for drop in cell_idx:
                if water_mass[drop] > 0:
                    result += multiplicity[drop] * water_mass[drop]
            return result

        return calculate_ml_old

    @staticmethod
    def make_calculate_ml_new(  # pylint: disable=too-many-statements,too-many-locals
        *,
        jit_flags,
        dx_dt,
        volume_of_x,
        x,
        phys_r_dr_dt,
        phys_RH_eq,
        phys_sigma,
        radius,
        phys_lambdaK,
        phys_lambdaD,
        phys_dk_D,
        phys_dk_K,
        volume_to_mass,
        mass_to_volume,
        within_tolerance,
        max_iters,
        RH_rtol,
        const,
    ):
        @numba.njit(**jit_flags)
        def minfun(  # pylint: disable=too-many-arguments
            x_new, x_old, timestep, kappa, f_org, rd3, temperature, RH, lv, pvs, D, K
        ):
            volume = volume_of_x(x_new)
            RH_eq = phys_RH_eq(
                radius(volume),
                temperature,
                kappa,
                rd3,
                phys_sigma(temperature, volume, const.PI_4_3 * rd3, f_org),
            )
            r_dr_dt = phys_r_dr_dt(RH_eq, temperature, RH, lv, pvs, D, K)
            return x_old - x_new + timestep * dx_dt(x_new, r_dr_dt)

        @numba.njit(**jit_flags)
        def calculate_ml_new(  # pylint: disable=too-many-arguments,too-many-statements,too-many-locals
            timestep,
            fake,
            T,
            p,
            RH,
            water_mass,
            v_cr,
            multiplicity,
            vdry,
            cell_idx,
            kappa,
            f_org,
            lv,
            pvs,
            DTp,
            KTp,
            rtol_x,
        ):  # pylint: disable=too-many-branches
            result = 0
            n_activating = 0
            n_deactivating = 0
            n_activated_and_growing = 0
            success = True
            lambdaK = phys_lambdaK(T, p)
            lambdaD = phys_lambdaD(DTp, T)
            for drop in cell_idx:
                v_drop = mass_to_volume(water_mass[drop])
                if v_drop < 0:
                    continue
                x_old = x(v_drop)
                r_old = radius(v_drop)
                x_insane = x(vdry[drop] / 100)
                rd3 = vdry[drop] / const.PI_4_3
                sgm = phys_sigma(T, v_drop, vdry[drop], f_org[drop])
                RH_eq = phys_RH_eq(r_old, T, kappa[drop], rd3, sgm)
                if not within_tolerance(np.abs(RH - RH_eq), RH, RH_rtol):
                    Dr = phys_dk_D(DTp, r_old, lambdaD)
                    Kr = phys_dk_K(KTp, r_old, lambdaK)
                    args = (
                        x_old,
                        timestep,
                        kappa[drop],
                        f_org[drop],
                        rd3,
                        T,
                        RH,
                        lv,
                        pvs,
                        Dr,
                        Kr,
                    )
                    r_dr_dt_old = phys_r_dr_dt(RH_eq, T, RH, lv, pvs, Dr, Kr)
                    dx_old = timestep * dx_dt(x_old, r_dr_dt_old)
                else:
                    dx_old = 0.0
                if dx_old == 0:
                    x_new = x_old
                else:
                    a = x_old
                    b = max(x_insane, a + dx_old)
                    fa = minfun(a, *args)
                    fb = minfun(b, *args)

                    counter = 0
                    while not fa * fb < 0:
                        counter += 1
                        if counter > max_iters:
                            if not fake:
                                warn(
                                    "failed to find interval",
                                    __file__,
                                    context=(
                                        "T",
                                        T,
                                        "p",
                                        p,
                                        "RH",
                                        RH,
                                        "a",
                                        a,
                                        "b",
                                        b,
                                        "fa",
                                        fa,
                                        "fb",
                                        fb,
                                    ),
                                )
                            success = False
                            break
                        b = max(x_insane, a + math.ldexp(dx_old, counter))
                        fb = minfun(b, *args)

                    if not success:
                        break
                    if a != b:
                        if a > b:
                            a, b = b, a
                            fa, fb = fb, fa

                        x_new, iters_taken = toms748_solve(
                            minfun,
                            args,
                            a,
                            b,
                            fa,
                            fb,
                            rtol_x,
                            max_iters,
                            within_tolerance,
                        )
                        if iters_taken in (-1, max_iters):
                            if not fake:
                                warn("TOMS failed", __file__)
                            success = False
                            break
                    else:
                        x_new = x_old

                mass_new = volume_to_mass(volume_of_x(x_new))
                mass_cr = volume_to_mass(v_cr[drop])
                result += multiplicity[drop] * mass_new
                if not fake:
                    if mass_new > mass_cr and mass_new > water_mass[drop]:
                        n_activated_and_growing += multiplicity[drop]
                    if mass_new > mass_cr > water_mass[drop]:
                        n_activating += multiplicity[drop]
                    if mass_new < mass_cr < water_mass[drop]:
                        n_deactivating += multiplicity[drop]
                    water_mass[drop] = mass_new
            n_ripening = n_activated_and_growing if n_deactivating > 0 else 0
            return result, success, n_activating, n_deactivating, n_ripening

        return calculate_ml_new

    # pylint disable=unused-argument
    def make_condensation_solver(
        self,
        timestep,
        n_cell,
        *,
        dt_range,
        adaptive,
        fuse,
        multiplier,
        RH_rtol,
        max_iters,
    ):
        return CondensationMethods.make_condensation_solver_impl(
            fastmath=self.formulae.fastmath,
            phys_pvs_C=self.formulae.saturation_vapour_pressure.pvs_Celsius,
            phys_lv=self.formulae.latent_heat.lv,
            phys_r_dr_dt=self.formulae.drop_growth.r_dr_dt,
            phys_RH_eq=self.formulae.hygroscopicity.RH_eq,
            phys_sigma=self.formulae.surface_tension.sigma,
            radius=self.formulae.trivia.radius,
            phys_T=self.formulae.state_variable_triplet.T,
            phys_p=self.formulae.state_variable_triplet.p,
            phys_pv=self.formulae.state_variable_triplet.pv,
            phys_dthd_dt=self.formulae.state_variable_triplet.dthd_dt,
            phys_lambdaK=self.formulae.diffusion_kinetics.lambdaK,
            phys_lambdaD=self.formulae.diffusion_kinetics.lambdaD,
            phys_dk_D=self.formulae.diffusion_kinetics.D,
            phys_dk_K=self.formulae.diffusion_kinetics.K,
            phys_diff_D=self.formulae.diffusion_thermics.D,
            phys_diff_K=self.formulae.diffusion_thermics.K,
            within_tolerance=self.formulae.trivia.within_tolerance,
            dx_dt=self.formulae.condensation_coordinate.dx_dt,
            volume=self.formulae.condensation_coordinate.volume,
            x=self.formulae.condensation_coordinate.x,
            mass_to_volume=self.formulae.particle_shape_and_density.mass_to_volume,
            volume_to_mass=self.formulae.particle_shape_and_density.volume_to_mass,
            timestep=timestep,
            dt_range=dt_range,
            adaptive=adaptive,
            fuse=fuse,
            multiplier=multiplier,
            RH_rtol=RH_rtol,
            max_iters=max_iters,
            const=self.formulae.constants,
        )

    @staticmethod
    @lru_cache()
    def make_condensation_solver_impl(
        *,
        fastmath,
        phys_pvs_C,
        phys_lv,
        phys_r_dr_dt,
        phys_RH_eq,
        phys_sigma,
        radius,
        phys_T,
        phys_p,
        phys_pv,
        phys_dthd_dt,
        phys_lambdaK,
        phys_lambdaD,
        phys_dk_D,
        phys_dk_K,
        phys_diff_D,
        phys_diff_K,
        within_tolerance,
        dx_dt,
        volume,
        x,
        volume_to_mass,
        mass_to_volume,
        timestep,
        dt_range,
        adaptive,
        fuse,
        multiplier,
        RH_rtol,
        max_iters,
        const,
    ):
        # pylint: disable=too-many-locals
        jit_flags = {
            **conf.JIT_FLAGS,
            **{"parallel": False, "cache": False, "fastmath": fastmath},
        }

        calculate_ml_old = CondensationMethods.make_calculate_ml_old(jit_flags)
        calculate_ml_new = CondensationMethods.make_calculate_ml_new(
            jit_flags=jit_flags,
            dx_dt=dx_dt,
            volume_of_x=volume,
            x=x,
            phys_r_dr_dt=phys_r_dr_dt,
            phys_RH_eq=phys_RH_eq,
            phys_sigma=phys_sigma,
            radius=radius,
            phys_lambdaK=phys_lambdaK,
            phys_lambdaD=phys_lambdaD,
            phys_dk_D=phys_dk_D,
            phys_dk_K=phys_dk_K,
            volume_to_mass=volume_to_mass,
            mass_to_volume=mass_to_volume,
            within_tolerance=within_tolerance,
            max_iters=max_iters,
            RH_rtol=RH_rtol,
            const=const,
        )
        step_impl = CondensationMethods.make_step_impl(
            jit_flags=jit_flags,
            phys_pvs_C=phys_pvs_C,
            phys_lv=phys_lv,
            calculate_ml_old=calculate_ml_old,
            calculate_ml_new=calculate_ml_new,
            phys_T=phys_T,
            phys_p=phys_p,
            phys_pv=phys_pv,
            phys_dthd_dt=phys_dthd_dt,
            phys_D=phys_diff_D,
            phys_K=phys_diff_K,
            const=const,
        )
        step_fake = CondensationMethods.make_step_fake(jit_flags, step_impl)
        adapt_substeps = CondensationMethods.make_adapt_substeps(
            jit_flags=jit_flags,
            timestep=timestep,
            step_fake=step_fake,
            dt_range=dt_range,
            fuse=fuse,
            multiplier=multiplier,
            within_tolerance=within_tolerance,
        )
        step = CondensationMethods.make_step(jit_flags, step_impl)

        @numba.njit(**jit_flags)
        def solve(  # pylint: disable=too-many-arguments
            water_mass,
            v_cr,
            multiplicity,
            vdry,
            cell_idx,
            kappa,
            f_org,
            thd,
            water_vapour_mixing_ratio,
            rhod,
            dthd_dt,
            d_water_vapour_mixing_ratio__dt,
            drhod_dt,
            m_d,
            rtol_x,
            rtol_thd,
            timestep,
            n_substeps,
        ):
            args = (
                water_mass,
                v_cr,
                multiplicity,
                vdry,
                cell_idx,
                kappa,
                f_org,
                thd,
                water_vapour_mixing_ratio,
                rhod,
                dthd_dt,
                d_water_vapour_mixing_ratio__dt,
                drhod_dt,
                m_d,
                rtol_x,
            )
            success = True
            if adaptive:
                n_substeps, success = adapt_substeps(args, n_substeps, thd, rtol_thd)
            if success:
                (
                    water_vapour_mixing_ratio,
                    thd,
                    n_activating,
                    n_deactivating,
                    n_ripening,
                    RH_max,
                    success,
                ) = step(args, timestep, n_substeps)
            else:
                n_activating, n_deactivating, n_ripening, RH_max = -1, -1, -1, -1
            return (
                success,
                water_vapour_mixing_ratio,
                thd,
                n_substeps,
                n_activating,
                n_deactivating,
                n_ripening,
                RH_max,
            )

        return solve


import os
import webbrowser
from dash import Dash, html, Input, Output, callback_context

app = Dash(__name__)

app.layout = html.Div([
    html.Button('Button 1', id='btn-nclicks-1', n_clicks=0),
    html.Button('Button 2', id='btn-nclicks-2', n_clicks=0),
    html.Button('Button 3', id='btn-nclicks-3', n_clicks=0),
    html.Div(id='container-button-timestamp')
])

@app.callback(
    Output('container-button-timestamp', 'children'),
    Input('btn-nclicks-1', 'n_clicks'),
    Input('btn-nclicks-2', 'n_clicks'),
    Input('btn-nclicks-3', 'n_clicks')
)
def displayClick(btn1, btn2, btn3):
    changed_id = [p['prop_id'] for p in callback_context.triggered][0]
    if 'btn-nclicks-1' in changed_id:
        msg = 'Button 1 was most recently clicked'
    elif 'btn-nclicks-2' in changed_id:
        msg = 'Button 2 was most recently clicked'
    elif 'btn-nclicks-3' in changed_id:
        msg = 'Button 3 was most recently clicked'
    else:
        msg = 'None of the buttons have been clicked yet'
    return html.Div(msg)

# Run app
PORT = 1100
def main():

    # The reloader has not yet run - open the browser
    if not os.environ.get("WERKZEUG_RUN_MAIN"):
        webbrowser.open_new("http://localhost:{}".format(PORT))

    # Otherwise, continue as normal
    app.run(host="127.0.0.1", port=PORT)

if __name__=='__main__':
    main()

Example Implementations/PyArduino V2/More Examples/SimpleButton.py


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# (c) oVoIndia | oVo-HxBots | RachelSaini

import asyncio, os, requests, time
from requests import post
from multiupload import anjana
from telethon.sync import events, Button
from multiupload.fsub import *
from multiupload.utils import downloader, humanbytes
from config import Config

@anjana.on(events.NewMessage(pattern='^/fileio'))
async def fileio(event):
	user_id = event.sender_id
	if event.is_private and not await check_participant(user_id, f'@{Config.CHNAME}', event):
		return
	if event.reply_to_msg_id:
		pass
	else:
		return await event.edit("Please Reply to File")

	async with anjana.action(event.chat_id, 'typing'):
		await asyncio.sleep(2)
	msg = await event.reply("**Processing...**")
	amjana = await event.get_reply_message()

	## FILEIO LIMITATION CHECKING
	if amjana.file.size > 1048576:
		return await event.edit("Oof.. File size too Large. FileIO Limitation is 100MB")

	## LOGGING TO A CHANNEL
	xx = await event.get_chat()
	reqmsg = f'''Req User: [{xx.first_name}](tg://user?id={xx.id})
FileName: {amjana.file.name}
FileSize: {humanbytes(amjana.file.size)}
#FILEIO'''
	await anjana.send_message(Config.LOG_CHANNEL, reqmsg)

	result = await downloader(
		f"downloads/{amjana.file.name}",
		amjana.media.document,
		msg,
		time.time(),
		f"**🏷 Downloading...**\n➲ **File Name:** {amjana.file.name}",
	)

	async with anjana.action(event.chat_id, 'document'):
		await msg.edit("Now Uploading to FileIO")
		url = f"https://file.io/"
		r = post(url, files={'file': open(f'{result.name}','rb')})
	await anjana.action(event.chat_id, 'cancel')

	hmm = f'''File Uploaded successfully !!
Server: FileIO

**~ File name:** __{amjana.file.name}__
**~ File size:** __{humanbytes(amjana.file.size)}__
NOTE: Once the download is complete, The file will be deleted from our servers.'''
	await msg.edit(hmm, buttons=(
		[Button.url('📦 Download', r.json()['link'])],
		[Button.url('Support Chat 💭', 't.me/hxsupport')]
		))

	os.remove(result.name)


import pygame

"""
The dragonfly sprite
Attributes:
    Attributes:
    window_size: the window size as a tuple
    size: size of the image
    x,y: the position of the dragonfly
    vx: the horizontal speed of the dragonfly
    vy: the vertical speed of the dragonfly
"""

# Derive your class from the Sprite super class

class Dragonfly(pygame.sprite.Sprite):
    def __init__(self, x, y, vx, vy):
        """
        :param x,y: the position of the dragonfly
        :param vx: the speed of the dragonfly
        :return: void
        """
        super().__init__();
        self.image = pygame.image.load("dragonfly.png").convert_alpha()
        self.size = self.image.get_size()

        # Required for collision detection
        # HINT: You will need this for the lab
        self.rect = self.image.get_rect()

        self.rect.x = x
        self.rect.y = y
        if x == 700:
            self.vx = -vx
            self.image = pygame.transform.rotate(self.image, 90)
        elif x == 0:
            self.vx = vx
            self.image = pygame.transform.rotate(self.image, -90)
        self.vy = vy

    def draw(self, screen):
        screen.blit(self.image, (self.rect.x, self.rect.y))

    def update(self):
        t_collide = self.rect.y + self.vy < 175
        b_collide = self.rect.y + self.image.get_height() + self.vy > 525

        # Check collision on top and bottom sides of screen
        if t_collide or b_collide:
            self.vy *= -1

        self.rect.x += self.vx
        self.rect.y += self.vy

        if self.rect.x > 750 or self.rect.x < -50:
            self.kill()

    def check_collide(self, group):
        collided = pygame.sprite.spritecollide(self, group, False)
        if len(collided) > 1:
            self.vx *= -1
            self.vy *= -1