# original code: https://github.com/dyhan0920/PyramidNet-PyTorch/blob/master/train.py
import argparse
import os
import shutil
import time

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models
from torch.autograd import Variable
import models.resnet as RN
import models.pyramidnet as PYRM
import utils
import numpy as np
from mixup import mixup_graph

from multiprocessing import Pool
import warnings

warnings.filterwarnings("ignore")

model_names = sorted(
    name for name in models.__dict__
    if name.islower() and not name.startswith("__") and callable(models.__dict__[name]))

def str2bool(v):
    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.')

parser = argparse.ArgumentParser(
    description='Cutmix PyTorch CIFAR-10, CIFAR-100 and ImageNet-1k Training')
parser.add_argument('--net_type',
                    default='resnet',
                    type=str,
                    help='networktype: resnet, and pyamidnet')
parser.add_argument('-j',
                    '--workers',
                    default=4,
                    type=int,
                    metavar='N',
                    help='number of data loading workers (default: 4)')
parser.add_argument('--epochs',
                    default=300,
                    type=int,
                    metavar='N',
                    help='number of total epochs to run')
parser.add_argument('-b',
                    '--batch_size',
                    default=64,
                    type=int,
                    metavar='N',
                    help='mini-batch size (default: 256)')
parser.add_argument('--lr',
                    '--learning-rate',
                    default=0.25,
                    type=float,
                    metavar='LR',
                    help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M', help='momentum')
parser.add_argument('--weight-decay',
                    '--wd',
                    default=1e-4,
                    type=float,
                    metavar='W',
                    help='weight decay (default: 1e-4)')
parser.add_argument('--print-freq',
                    '-p',
                    default=1,
                    type=int,
                    metavar='N',
                    help='print frequency (default: 10)')
parser.add_argument('--depth', default=50, type=int, help='depth of the network (default: 32)')
parser.add_argument('--no-bottleneck',
                    dest='bottleneck',
                    action='store_false',
                    help='to use basicblock for CIFAR datasets (default: bottleneck)')
parser.add_argument('--dataset',
                    dest='dataset',
                    default='imagenet',
                    type=str,
                    help='dataset (options: cifar10, cifar100, and imagenet)')
parser.add_argument('--no-verbose',
                    dest='verbose',
                    action='store_false',
                    help='to print the status at every iteration')
parser.add_argument('--alpha',
                    default=240,
                    type=float,
                    help='number of new channel increases per depth (default: 300)')
parser.add_argument('--expname', default='', type=str, help='name of experiment')
parser.add_argument('--beta', default=1.0, type=float, help='hyperparameter beta')
parser.add_argument('--mixup_prob', default=1.0, type=float, help='mixup probability')
parser.add_argument('--padding', default=4, type=int, help='padding for transform_train')

parser.add_argument('--method',
                    default='puzzle',
                    type=str,
                    help='mixup type',
                    choices=('vanilla', 'cut', 'puzzle'))
# For Puzzle Mix
parser.add_argument('--mp', type=int, default=8, help='multi-process for graphcut (CPU)')
parser.add_argument('--n_labels', type=int, default=3, help='label space size')
parser.add_argument('--neigh_size',
                    type=int,
                    default=4,
                    help='neighbor size for computing distance beteeen image regions')

parser.add_argument('--p_beta', type=float, default=1.2, help='label smoothness')
parser.add_argument('--p_gamma', type=float, default=0.5, help='data local smoothness')
parser.add_argument('--p_eta', type=float, default=0.2, help='prior term')
parser.add_argument('--transport', type=str2bool, default=True, help='whether to use transport')
parser.add_argument('--t_eps', type=float, default=0.8, help='transport cost coefficient')

parser.add_argument('--clean_lam', type=float, default=0.0, help='clean regularization')
parser.add_argument('--seed', type=int, default=-1, help='random seed')

parser.set_defaults(bottleneck=True)
parser.set_defaults(verbose=True)

best_err1 = 100
best_err5 = 100

def main():
    global args, best_err1, best_err5
    args = parser.parse_args()

    if args.seed >= 0:
        np.random.seed(args.seed)
        torch.manual_seed(args.seed)
        torch.cuda.manual_seed(args.seed)
    cudnn.benchmark = True

    # Save path
    args.expname += args.method
    if args.transport:
        args.expname += '_tp'
    args.expname += '_prob_' + str(args.mixup_prob)
    if args.clean_lam > 0:
        args.expname += '_clean_' + str(args.clean_lam)
    if args.seed >= 0:
        args.expname += '_seed' + str(args.seed)
    print("Model is saved at {}".format(args.expname))

    # Dataset and loader
    if args.dataset.startswith('cifar'):
        mean = [x / 255.0 for x in [125.3, 123.0, 113.9]]
        std = [x / 255.0 for x in [63.0, 62.1, 66.7]]
        normalize = transforms.Normalize(mean=mean, std=std)

        transform_train = transforms.Compose([
            transforms.RandomCrop(32, padding=args.padding),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ])
        transform_test = transforms.Compose([transforms.ToTensor(), normalize])

        if args.dataset == 'cifar100':
            train_loader = torch.utils.data.DataLoader(datasets.CIFAR100('~/Datasets/cifar100/',
                                                                         train=True,
                                                                         download=True,
                                                                         transform=transform_train),
                                                       batch_size=args.batch_size,
                                                       shuffle=True,
                                                       num_workers=args.workers,
                                                       pin_memory=True)
            val_loader = torch.utils.data.DataLoader(datasets.CIFAR100('~/Datasets/cifar100/',
                                                                       train=False,
                                                                       transform=transform_test),
                                                     batch_size=args.batch_size // 4,
                                                     shuffle=True,
                                                     num_workers=args.workers,
                                                     pin_memory=True)
            numberofclass = 100
        elif args.dataset == 'cifar10':
            train_loader = torch.utils.data.DataLoader(datasets.CIFAR10('../data',
                                                                        train=True,
                                                                        download=True,
                                                                        transform=transform_train),
                                                       batch_size=args.batch_size,
                                                       shuffle=True,
                                                       num_workers=args.workers,
                                                       pin_memory=True)
            val_loader = torch.utils.data.DataLoader(datasets.CIFAR10('../data',
                                                                      train=False,
                                                                      transform=transform_test),
                                                     batch_size=args.batch_size,
                                                     shuffle=True,
                                                     num_workers=args.workers,
                                                     pin_memory=True)
            numberofclass = 10
        else:
            raise Exception('unknown dataset: {}'.format(args.dataset))

    elif args.dataset == 'imagenet':
        traindir = os.path.join('/data/readonly/ImageNet-Fast/imagenet/train')
        valdir = os.path.join('/data/readonly/ImageNet-Fast/imagenet/val')
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]
        normalize = transforms.Normalize(mean=mean, std=std)
        jittering = utils.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4)
        lighting = utils.Lighting(alphastd=0.1,
                                  eigval=[0.2175, 0.0188, 0.0045],
                                  eigvec=[[-0.5675, 0.7192, 0.4009], [-0.5808, -0.0045, -0.8140],
                                          [-0.5836, -0.6948, 0.4203]])

        train_dataset = datasets.ImageFolder(
            traindir,
            transforms.Compose([
                transforms.RandomResizedCrop(224),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                jittering,
                lighting,
                normalize,
            ]))
        train_sampler = None

        train_loader = torch.utils.data.DataLoader(train_dataset,
                                                   batch_size=args.batch_size,
                                                   shuffle=(train_sampler is None),
                                                   num_workers=args.workers,
                                                   pin_memory=True,
                                                   sampler=train_sampler)
        val_transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            normalize,
        ])
        val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(valdir, val_transform),
                                                 batch_size=args.batch_size // 4,
                                                 shuffle=False,
                                                 num_workers=args.workers,
                                                 pin_memory=True)
        numberofclass = 1000
        args.neigh_size = min(args.neigh_size, 2)

    else:
        raise Exception('unknown dataset: {}'.format(args.dataset))

    # Model
    print("=> creating model '{}'".format(args.net_type))
    if args.net_type == 'resnet':
        model = RN.ResNet(args.dataset, args.depth, numberofclass, args.bottleneck)  # for ResNet
    elif args.net_type == 'pyramidnet':
        model = PYRM.PyramidNet(args.dataset, args.depth, args.alpha, numberofclass,
                                args.bottleneck)
    else:
        raise Exception('unknown network architecture: {}'.format(args.net_type))

    pretrained = "runs/{}/{}".format(args.expname, 'checkpoint.pth.tar')
    if os.path.isfile(pretrained):
        print("=> loading checkpoint '{}'".format(pretrained))
        checkpoint = torch.load(pretrained)
        checkpoint['state_dict'] = dict(
            (key[7:], value) for (key, value) in checkpoint['state_dict'].items())
        model.load_state_dict(checkpoint['state_dict'])
        cur_epoch = checkpoint['epoch'] + 1
        best_err1 = checkpoint['best_err1']
        print("=> loaded checkpoint '{}'(epoch: {}, best err1: {}%)".format(
            pretrained, cur_epoch, checkpoint['best_err1']))
    else:
        cur_epoch = 0
        print("=> no checkpoint found at '{}'".format(pretrained))

    model = torch.nn.DataParallel(model).cuda()
    print('the number of model parameters: {}'.format(
        sum([p.data.nelement() for p in model.parameters()])))

    # define loss function (criterion) and optimizer
    criterion = nn.CrossEntropyLoss().cuda()
    criterion_batch = nn.CrossEntropyLoss(reduction='none').cuda()

    optimizer = torch.optim.SGD(model.parameters(),
                                args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay,
                                nesterov=True)
    if os.path.isfile(pretrained):
        optimizer.load_state_dict(checkpoint['optimizer'])
        print("optimizer is loaded!")

    mean_torch = torch.tensor(mean, dtype=torch.float32).reshape(1, 3, 1, 1).cuda()
    std_torch = torch.tensor(std, dtype=torch.float32).reshape(1, 3, 1, 1).cuda()
    if args.mp > 0:
        mp = Pool(args.mp)
    else:
        mp = None

    # Start training and validation
    for epoch in range(cur_epoch, args.epochs):
        adjust_learning_rate(optimizer, epoch)

        # train for one epoch
        train_loss = train(train_loader, model, criterion, criterion_batch, optimizer, epoch,
                           mean_torch, std_torch, mp)
        # evaluate on validation set
        err1, err5, val_loss = validate(val_loader, model, criterion, epoch)

        # remember best prec@1 and save checkpoint
        is_best = err1 <= best_err1
        best_err1 = min(err1, best_err1)
        if is_best:
            best_err5 = err5

        print('Current best accuracy (top-1 and 5 error):', best_err1, best_err5)
        save_checkpoint(
            {
                'epoch': epoch,
                'arch': args.net_type,
                'state_dict': model.state_dict(),
                'best_err1': best_err1,
                'best_err5': best_err5,
                'optimizer': optimizer.state_dict(),
            }, is_best)

    print('Best accuracy (top-1 and 5 error):', best_err1, best_err5)

def train(train_loader, model, criterion, criterion_batch, optimizer, epoch, mean, std, mp=None):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()

    # switch to train mode
    model.train()

    end = time.time()
    current_LR = get_learning_rate(optimizer)[0]
    for i, (input, target) in enumerate(train_loader):
        optimizer.zero_grad()
        # measure data loading time
        data_time.update(time.time() - end)

        input = input.cuda()
        target = target.cuda()

        r = np.random.rand(1)
        if args.beta > 0 and r < args.mixup_prob and args.method == 'cut':
            # generate mixed sample
            lam = np.random.beta(args.beta, args.beta)
            rand_index = torch.randperm(input.size()[0]).cuda()
            target_a = target
            target_b = target[rand_index]
            bbx1, bby1, bbx2, bby2 = rand_bbox(input.size(), lam)
            input[:, :, bbx1:bbx2, bby1:bby2] = input[rand_index, :, bbx1:bbx2, bby1:bby2]
            # adjust lambda to exactly match pixel ratio
            lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (input.size()[-1] * input.size()[-2]))
            # compute output
            output = model(input)
            loss = criterion(output, target_a) * lam + criterion(output, target_b) * (1. - lam)

        elif args.beta > 0 and r < args.mixup_prob and args.method == 'puzzle':
            # calculate saliency map
            input_var = Variable(input, requires_grad=True)
            if args.clean_lam == 0:
                model.eval()
                output = model(input_var)
                loss_clean = criterion(output, target)
                loss_clean.backward(retain_graph=False)
                optimizer.zero_grad()
                model.train()
            else:
                # gradient regularization
                output = model(input_var)
                loss_clean = args.clean_lam * criterion(output, target)
                loss_clean.backward(retain_graph=True)

            unary = torch.sqrt(torch.mean(input_var.grad**2, dim=1))

            # perform mixup
            alpha = np.random.beta(args.beta, args.beta)
            rand_index = torch.randperm(input.size()[0]).cuda()
            block_num = 2**np.random.randint(1, 5)
            with torch.no_grad():
                input, lam = mixup_graph(input,
                                         unary,
                                         rand_index,
                                         block_num=block_num,
                                         alpha=alpha,
                                         beta=args.p_beta,
                                         gamma=args.p_gamma,
                                         eta=args.p_eta,
                                         neigh_size=args.neigh_size,
                                         n_labels=args.n_labels,
                                         mean=mean,
                                         std=std,
                                         transport=args.transport,
                                         t_eps=args.t_eps,
                                         dataset=args.dataset,
                                         mp=mp)
            # calculate loss
            output = model(input)
            loss = lam * criterion_batch(output, target) + (1 - lam) * criterion_batch(
                output, target[rand_index])
            loss = torch.mean(loss)

        else:
            # compute output
            output = model(input)
            loss = criterion(output, target)

        # measure accuracy and record loss
        err1, err5 = accuracy(output.data, target, topk=(1, 5))

        losses.update(loss.item(), input.size(0))
        top1.update(err1.item(), input.size(0))
        top5.update(err5.item(), input.size(0))

        # compute gradient and do SGD step
        loss.backward()
        optimizer.step()

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % args.print_freq == 0 and args.verbose == True:
            print('Epoch: [{0}/{1}][{2}/{3}]\t'
                  'LR: {LR:.6f}\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  'Top 1-err {top1.val:.4f} ({top1.avg:.4f})\t'
                  'Top 5-err {top5.val:.4f} ({top5.avg:.4f})'.format(epoch,
                                                                     args.epochs,
                                                                     i,
                                                                     len(train_loader),
                                                                     LR=current_LR,
                                                                     batch_time=batch_time,
                                                                     data_time=data_time,
                                                                     loss=losses,
                                                                     top1=top1,
                                                                     top5=top5))

    print(
        '* Epoch: [{0}/{1}]\t Top 1-err {top1.avg:.3f}  Top 5-err {top5.avg:.3f}\t Train Loss {loss.avg:.3f}'
        .format(epoch, args.epochs, top1=top1, top5=top5, loss=losses))

    return losses.avg

def rand_bbox(size, lam):
    W = size[2]
    H = size[3]
    cut_rat = np.sqrt(1. - lam)
    cut_w = np.int(W * cut_rat)
    cut_h = np.int(H * cut_rat)

    # uniform
    cx = np.random.randint(W)
    cy = np.random.randint(H)

    bbx1 = np.clip(cx - cut_w // 2, 0, W)
    bby1 = np.clip(cy - cut_h // 2, 0, H)
    bbx2 = np.clip(cx + cut_w // 2, 0, W)
    bby2 = np.clip(cy + cut_h // 2, 0, H)

    return bbx1, bby1, bbx2, bby2

def validate(val_loader, model, criterion, epoch):
    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()

    # switch to evaluate mode
    model.eval()

    end = time.time()
    for i, (input, target) in enumerate(val_loader):
        target = target.cuda()

        output = model(input)
        loss = criterion(output, target)

        # measure accuracy and record loss
        err1, err5 = accuracy(output.data, target, topk=(1, 5))

        losses.update(loss.item(), input.size(0))

        top1.update(err1.item(), input.size(0))
        top5.update(err5.item(), input.size(0))

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % args.print_freq == 0 and args.verbose == True:
            print('Test (on val set): [{0}/{1}][{2}/{3}]\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  'Top 1-err {top1.val:.4f} ({top1.avg:.4f})\t'
                  'Top 5-err {top5.val:.4f} ({top5.avg:.4f})'.format(epoch,
                                                                     args.epochs,
                                                                     i,
                                                                     len(val_loader),
                                                                     batch_time=batch_time,
                                                                     loss=losses,
                                                                     top1=top1,
                                                                     top5=top5))

    print(
        '* Epoch: [{0}/{1}]\t Top 1-err {top1.avg:.3f}  Top 5-err {top5.avg:.3f}\t Test Loss {loss.avg:.3f}'
        .format(epoch, args.epochs, top1=top1, top5=top5, loss=losses))
    return top1.avg, top5.avg, losses.avg

def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
    directory = "runs/%s/" % (args.expname)
    if not os.path.exists(directory):
        os.makedirs(directory)
    filename = directory + filename
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, 'runs/%s/' % (args.expname) + 'model_best.pth.tar')

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def adjust_learning_rate(optimizer, epoch):
    """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
    if args.dataset.startswith('cifar'):
        lr = args.lr * (0.1**(epoch // (args.epochs * 0.5))) * (0.1**(epoch //
                                                                      (args.epochs * 0.75)))
    elif args.dataset == ('imagenet'):
        if args.epochs == 300:
            lr = args.lr * (0.1**(epoch // 75))
        else:
            lr = args.lr * (0.1**(epoch // 30))

    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

def get_learning_rate(optimizer):
    lr = []
    for param_group in optimizer.param_groups:
        lr += [param_group['lr']]
    return lr

def accuracy(output, target, topk=(1, )):
    """Computes the precision@k for the specified values of k"""
    maxk = max(topk)
    batch_size = target.size(0)

    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.reshape(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
        wrong_k = batch_size - correct_k
        res.append(wrong_k.mul_(100.0 / batch_size))

    return res

if __name__ == '__main__':
    main()


import numpy as np
import torch
import requests
import matplotlib.pyplot as plt
import seaborn as sns # for data visualization
# plt.rcParams['font.family'] = 'Noto Sans CJK JP'
plt.rcParams.update({
  # 'font.size': 5,
  'font.family': 'Noto Sans CJK JP'
})
t = torch
GPU_OK = t.cuda.is_available()

def position_encoding_ddd(t, i, d):
  k = int(i/2)
  omiga = 1 / np.power(10000, 2 * k / d)
  even = (i / 2).is_integer()
  return np.sin(omiga * t) if even else np.cos(omiga * t)

# seq: (seq_len, feature)
# return: (seq_len, feature)
def position_encoding(seq):
  embs = []
  for t, data in enumerate(seq):
    d = data.shape[0]
    pos_emb = [position_encoding_ddd(t, i, d) for i in range(0, d)]
    pos_emb = torch.tensor(pos_emb)
    embs.append(pos_emb)
  embs = torch.stack(embs)
  return embs.cuda() if GPU_OK else embs

def position_matrix(seq_len, feature):
  d = feature
  embs = []
  for j in range(seq_len):
    pos_emb = [position_encoding_ddd(j, i, d) for i in range(d)]
    pos_emb = torch.tensor(pos_emb)
    embs.append(pos_emb)
  embs = torch.stack(embs)
  return embs.cuda() if GPU_OK else embs

# def request_my_logger(dic, desc = 'No describe'):
#   try:
#     url = "https://hookb.in/b9xlr2GnnjC3DDogQ0jY"
#     dic['desc'] = desc
#     requests.post(url, json=dic)
#   except:
#     print('Something went wrong in request_my_logger()')

# head_scores: (batch, head, seq_len, seq_len)
# info: (seq_len)
def draw_head_attention(head_scores, info, cls_pos = 0, path='dd.png', desc = ''):
  mat = []
  head_scores = head_scores[0]
  head, seq_len, _ = head_scores.shape
  avg_scores = head_scores.sum(dim=0) / head # (seq_len, seq_len)
  avg_scores = avg_scores[cls_pos] # (seq_len)
  xs = info
  ys = []
  for i in range(head):
    ys.append(f'head_{i}')
    score = head_scores[i] # (seq_len, seq_len)
    score = score[cls_pos] # (seq_len)
    mat.append(score)
  mat.append(avg_scores)
  ys += ['avg']
  mat = [row.tolist() for row in mat]
  for row in mat:
    row.pop(cls_pos)
  output_heatmap(np.transpose(mat), ys, xs, path, desc)

def output_heatmap(mat, xs, ys, path = 'dd.png', desc = ''):
  if len(ys) > 16:
    print(f'Warning: too long sequence: {len(ys)}')
    # sns.set(font_scale=0.5)
  else:
    # 5sns.set(font_scale=1.0)
    pass
  plt.clf()
  sns.heatmap(mat, xticklabels=xs, yticklabels=ys)
  plt.text(0 , -0.5, desc)
  plt.savefig(path)

def fit_sigmoided_to_label(out):
  assert len(out.shape) == 2
  results = []
  for item in out:
    assert item >= 0 and item <= 1
    if item < 0.5:
      results.append(0)
    else:
      results.append(1)
  return t.LongTensor(results)

def cal_prec_rec_f1_v2(results, targets):
  TP = 0
  FP = 0
  FN = 0
  TN = 0
  for guess, target in zip(results, targets):
    if guess == 1:
      if target == 1:
        TP += 1
      elif target == 0:
        FP += 1
    elif guess == 0:
      if target == 1:
        FN += 1
      elif target == 0:
        TN += 1
  prec = TP / (TP + FP) if (TP + FP) > 0 else 0
  rec = TP / (TP + FN) if (TP + FN) > 0 else 0
  f1 = (2 * prec * rec) / (prec + rec) if (prec + rec) != 0 else 0
  balanced_acc_factor1 = TP / (TP + FN) if (TP + FN) > 0 else 0
  balanced_acc_factor2 = TN / (FP + TN) if (FP + TN) > 0 else 0
  balanced_acc = (balanced_acc_factor1 + balanced_acc_factor2) / 2
  return prec, rec, f1, balanced_acc


import os
import rasterio
import numpy as np

from raster_utils.spectral_indices import ndvi, msi, ndmi

main_path = r"G:\Mi unidad\Tesis_5"
dest_path = r"C:\Users\DELL\PycharmProjects\Tesis\spectral_indices_images"

parcel_ids = [i.split('_')[1] for i in os.listdir(main_path)]

for root, dirs, files in os.walk(main_path):
    for dir_ in dirs:
        dir_path = os.path.join(root, dir_)
        for file in os.listdir(dir_path):
            file_path = os.path.join(dir_path, file)
            file_name = f"{file.split('.')[0]}_indices"
            with rasterio.open(file_path) as src:
                profile = src.profile
                profile.update(dtype=rasterio.float32, count=3)
                ndvi_band = ndvi(src)
                msi_band = msi(src)
                ndmi_band = ndmi(src)

                stacked = np.stack((ndvi_band, msi_band, ndmi_band))

                with rasterio.open(os.path.join(dest_path, f"{file_name}.tif"), 'w', **profile) as dst:
                    dst.descriptions = (['ndvi', 'msi', 'ndmi'])
                    dst.write(stacked.astype(rasterio.float32))


from collections import defaultdict

dic = defaultdict(int)
total_cnt = 0

while True:
    words = input()
    if not words:
        break

    dic[words] += 1
    total_cnt += 1

tree_name = sorted(list(dic.keys()))
print(tree_name)
for name in tree_name:
    print("{} {:.4f}".format(name, (dic[name]/total_cnt) * 100))


from conv2d import Conv2d
from linear import Linear
from torch import nn

def construct_features(cfg,
                       padding,
                       module_names,
                       batch_norm=False,
                       slot_machine=False,
                       k=8,
                       greedy_selection=True
                       ):
    """
    Constructs convolutional layers for a neural network
    :param cfg: The configuration of the model. A list of the layers
    :param padding: The padding to apply to layers
    :param module_names: The module names
    :param batch_norm: (bool), Whether to apply batch normalization or not
    :param slot_machine: (bool), constructs a module for weight updates or slot_machines
    :param k: (int), The number of options per weight
    :param greedy_selection:(bool) the selection method for if model is slot machine
    :return: model: a sequential module of convolutional layers
    """
    model = nn.Sequential()
    in_channels = 3
    for i, v in enumerate(cfg):
        if v == 'M':
            model.add_module(module_names[i], nn.MaxPool2d(kernel_size=2, stride=2))
        else:
            if slot_machine:
                model.add_module(module_names[i],
                                 Conv2d(in_channels, v, k,
                                        kernel_size=3,
                                        padding=padding[i],
                                        greedy_selection=greedy_selection
                                        )
                                 )
            else:
                model.add_module(module_names[i],
                                 nn.Conv2d(in_channels, v,
                                           kernel_size=3,
                                           padding=padding[i],
                                           bias=False)
                                 )

            if batch_norm:
                model.add_module(f"bn_{i}", nn.BatchNorm2d(v, affine=False))

            model.add_module(f"relu_{i}", nn.ReLU(inplace=True))
            in_channels = v

    return model

def construct_classifier(cfg,
                         module_names,
                         in_features,
                         slot_machine=False,
                         k=8,
                         greedy_selection=True
                         ):
    """
    Constructs a sequential model of  fully-connected layers
    :param cfg:(List) The configuration of the model
    :param module_names: (List) The names of the layers
    :param in_features: (int) The number of input features to first fully-connected layer
    :param slot_machine: (bool)  constructs a module for weight updates or slot_machines
    :param k:(int), the number of options per weight if model is a slot machine
    :param greedy_selection: (bool), use greedy selection if model is slot machine
    :return: model: a sequential module of fully-connected layers
    """
    model = nn.Sequential()
    for i, v in enumerate(cfg):
        if v == 'D':
            model.add_module(module_names[i], nn.Dropout(p=0.5))
        elif v == "relu":
            model.add_module(module_names[i], nn.ReLU(inplace=True))
        else:
            if slot_machine:
                model.add_module(module_names[i],Linear(in_features, v, k, greedy_selection))
            else:
                model.add_module(module_names[i], nn.Linear(in_features, v, bias=False))
            in_features = v

    return model


from twisted.internet.defer import succeed

from txzookeeper import ZookeeperClient
from juju.state.initialize import StateHierarchy

from juju.control import admin
from .common import ControlToolTest

class AdminInitializeTest(ControlToolTest):

    def test_initialize(self):
        """The admin cli dispatches the initialize method with arguments."""

        client = self.mocker.patch(ZookeeperClient)
        hierarchy = self.mocker.patch(StateHierarchy)

        self.setup_cli_reactor()
        client.connect()
        self.mocker.result(succeed(client))

        hierarchy.initialize()
        self.mocker.result(succeed(True))
        client.close()
        self.capture_stream('stderr')
        self.setup_exit(0)
        self.mocker.replay()

        admin(["initialize",
               "--instance-id", "foobar",
               "--admin-identity", "admin:genie",
               "--provider-type", "dummy"])


# Databricks notebook source
# MAGIC %md
# MAGIC ### Pyspark Data types and its length
# MAGIC * All datatypes we can import from  __`from pyspark.sql.types import *`__

# COMMAND ----------

# MAGIC %md
# MAGIC

# MAGIC

# MAGIC <colgroup>
# MAGIC <col width="3%">
# MAGIC <col width="60%">
# MAGIC <col width="37%">
# MAGIC </colgroup>
# MAGIC

# MAGIC
Data type
# MAGIC Value type
# MAGIC API to access or create data type
# MAGIC

# MAGIC

# MAGIC

# MAGIC
ByteType
# MAGIC integer  Note: Numbers are converted to 1-byte signed =
# MAGIC integer numbers at runtime.  Make sure sure that numbers are within the ran=
# MAGIC ge of -128 to 127.
# MAGIC ByteType()
# MAGIC

# MAGIC
ShortType
# MAGIC integer  Note: Numbers are converted to 2-byte signed =
# MAGIC integer numbers at runtime.  Make sure sure that numbers are within the ran=
# MAGIC ge of -32768 to 32767.
# MAGIC ShortType()
# MAGIC

# MAGIC
IntegerType
# MAGIC integer
# MAGIC IntegerType()
# MAGIC

# MAGIC
LongType
# MAGIC integer  Note: Numbers are converted to 8-byte signed =
# MAGIC integer numbers at runtime.  Make sure sure that numbers are within the ran=
# MAGIC ge of -9223372036854775808 to 9223372036854775807.  Otherwise, convert data=
# MAGIC  to decimal.Decimal and use DecimalType.
# MAGIC LongType()
# MAGIC

# MAGIC
FloatType
# MAGIC numeric  Note: Numbers are converted to 4-byte single-=
# MAGIC precision floating point numbers at runtime.
# MAGIC FloatType()
# MAGIC

# MAGIC
DoubleType
# MAGIC float
# MAGIC DobleType()
# MAGIC

# MAGIC
DecimalType
# MAGIC decimal.Decimal
# MAGIC DecimalType()
# MAGIC

# MAGIC
StringType
# MAGIC character
# MAGIC StringType()
# MAGIC

# MAGIC
BinaryType
# MAGIC raw
# MAGIC BinaryType()
# MAGIC

# MAGIC
BooleanType
# MAGIC logical
# MAGIC BooleanType()
# MAGIC

# MAGIC
TimestampType
# MAGIC datetime.datetime
# MAGIC TimestampType()
# MAGIC

# MAGIC
DateType
# MAGIC datetime.date
# MAGIC DateType()
# MAGIC

# MAGIC
ArrayType
# MAGIC vector or list
# MAGIC list(type=array, elementType=elementType, contain=
# MAGIC sNull=[containsNull]) Note: The default value of contain=
# MAGIC sNull is TRUE.
# MAGIC

# MAGIC
MapType
# MAGIC environment
# MAGIC list(type=map, keyType=keyType, valueType=value=
# MAGIC Type, valueContainsNull=[valueContainsNull]) Note: The d=
# MAGIC efault value of valueContainsNull is TRUE.
# MAGIC

# MAGIC
StructType
# MAGIC named list
# MAGIC list(type struct, fields=fields) Note:</s=trong> fields is a Seq of StructFields. Also, two fields with the same name=
# MAGIC  are not allowed.
# MAGIC

# MAGIC
StructField
# MAGIC The value type of the data type of this field (For example, integer for=
# MAGIC  a StructField with the data type IntegerType)
# MAGIC list(name=name, type=dataType, nullable=[nullable]) Note:=
# MAGIC  The default value of nullable is TRUE.
# MAGIC

# MAGIC

# MAGIC

# MAGIC

# COMMAND ----------

#**************************************************
#* Importing pyspark required Data Types
#***************************************************/
from pyspark.sql.types import StructField, StructType, StringType,LongType,IntegerType,DoubleType

# COMMAND ----------

#**************************************************
#* Creating Pyspark Schema For Channel Table
#***************************************************/
def channels_schema():
  channels_schema = StructType([StructField('CHANNEL_ID', IntegerType(), False),
                     StructField('CHANNEL_DESC', StringType(), False),
                     StructField('CHANNEL_CLASS', StringType(), False),
                     StructField('CHANNEL_CLASS_ID', IntegerType(), False),
                     StructField('CHANNEL_TOTAL', StringType(), False),
                     StructField('CHANNEL_TOTAL_ID', IntegerType(), False),
                     StructField('_corrupt_record', StringType(), True)])
  return channels_schema

# COMMAND ----------

#**************************************************
#* Creating Pyspark Schema For Countries Table
#***************************************************/
def countries_schema():
  countries_schema = StructType([StructField('COUNTRY_ID', IntegerType(), False),
                     StructField('COUNTRY_ISO_CODE', StringType(), False),
                     StructField('COUNTRY_NAME', StringType(), False),
                     StructField('COUNTRY_SUBREGION', StringType(), False),
                     StructField('COUNTRY_SUBREGION_ID', IntegerType(), False),
                     StructField('COUNTRY_REGION', StringType(), False),
                     StructField('COUNTRY_REGION_ID', IntegerType(), False),
					 StructField('COUNTRY_TOTAL', StringType(), False),
					 StructField('COUNTRY_TOTAL_ID', IntegerType(), False),
					 StructField('COUNTRY_NAME_HIST', StringType(), False),
                     StructField('_corrupt_record', StringType(), False)])
  return countries_schema

# COMMAND ----------

#**************************************************
#* Creating Pyspark Schema For Customers Table
#***************************************************
def customers_schema():
  customers_schema = StructType([StructField('CUST_ID', IntegerType(), False),
                     StructField('CUST_FIRST_NAME', StringType(), False),
                     StructField('CUST_LAST_NAME', StringType(), False),
                     StructField('CUST_GENDER', StringType(), False),
                     StructField('CUST_YEAR_OF_BIRTH', IntegerType(), False),
                     StructField('CUST_MARITAL_STATUS', StringType(), False),
                     StructField('CUST_STREET_ADDRESS', StringType(), False),
					 StructField('CUST_POSTAL_CODE', IntegerType(), False),
					 StructField('CUST_CITY', StringType(), False),
					 StructField('CUST_CITY_ID', StringType(), False),
					 StructField('CUST_STATE_PROVINCE', StringType(), False),
					 StructField('CUST_STATE_PROVINCE_ID', StringType(), False),
					 StructField('COUNTRY_ID', IntegerType(), False),
					 StructField('CUST_MAIN_PHONE_NUMBER', StringType(), False),
					 StructField('CUST_INCOME_LEVEL', StringType(), False),
					 StructField('CUST_CREDIT_LIMIT', StringType(), False),
					 StructField('CUST_EMAIL', StringType(), False),
					 StructField('CUST_TOTAL', StringType(), False),
					 StructField('CUST_TOTAL_ID', StringType(), False),
					 StructField('CUST_SRC_ID', StringType(), False),
					 StructField('CUST_EFF_FROM', StringType(), False),
					 StructField('CUST_EFF_TO', StringType(), False),
					 StructField('CUST_VALID', StringType(), False),
                     StructField('_corrupt_record', StringType(), False)])
  return customers_schema

# COMMAND ----------

#**************************************************
#* Creating Pyspark Schema For product Table
#***************************************************
def product_schema():
  product_schema = StructType([StructField('PROD_ID', IntegerType(), False),
                     StructField('PROD_NAME', StringType(), False),
                     StructField('PROD_DESC', StringType(), False),
                     StructField('PROD_SUBCATEGORY', StringType(), False),
                     StructField('PROD_SUBCATEGORY_ID', StringType(), False),
                     StructField('PROD_SUBCATEGORY_DESC', StringType(), False),
                     StructField('PROD_CATEGORY', StringType(), False),
					 StructField('PROD_CATEGORY_ID', StringType(), False),
					 StructField('PROD_CATEGORY_DESC', StringType(), False),
					 StructField('PROD_WEIGHT_CLASS', StringType(), False),
					 StructField('PROD_UNIT_OF_MEASURE', StringType(), False),
					 StructField('PROD_PACK_SIZE', StringType(), False),
					 StructField('SUPPLIER_ID', StringType(), False),
					 StructField('PROD_STATUS', StringType(), False),
					 StructField('PROD_LIST_PRICE', StringType(), False),
					 StructField('PROD_MIN_PRICE', StringType(), False),
					 StructField('PROD_TOTAL', StringType(), False),
					 StructField('PROD_TOTAL_ID', StringType(), False),
					 StructField('PROD_SRC_ID', StringType(), False),
					 StructField('PROD_EFF_FROM', StringType(), False),
					 StructField('PROD_EFF_TO', StringType(), False),
					 StructField('PROD_VALID', StringType(), False),
                     StructField('_corrupt_record', StringType(), False)])
  return product_schema

# COMMAND ----------

#**************************************************
#* Creating Pyspark Schema For promotions Table
#***************************************************
def promotions_schema():
  promotions_schema = StructType([StructField('PROMO_ID', IntegerType(), False),
                     StructField('PROMO_NAME', StringType(), False),
                     StructField('PROMO_SUBCATEGORY', StringType(), False),
                     StructField('PROMO_SUBCATEGORY_ID', IntegerType(), False),
                     StructField('PROMO_CATEGORY', StringType(), False),
                     StructField('PROMO_CATEGORY_ID', IntegerType(), False),
                     StructField('PROMO_COST', IntegerType(), False),
                     StructField('PROMO_BEGIN_DATE', StringType(), False),
                     StructField('PROMO_END_DATE', StringType(), False),
                     StructField('PROMO_TOTAL', StringType(), False),
                     StructField('PROMO_TOTAL_ID', IntegerType(), False),
                     StructField('_corrupt_record', StringType(), False)])
  return promotions_schema

# COMMAND ----------

#**************************************************
#* Creating Pyspark Schema For times Table
#***************************************************
def times_schema():
  times_schema = StructType([StructField('TIME_ID',StringType(),True),
                           StructField('DAY_NAME',StringType(),True),
                           StructField('DAY_NUMBER_IN_WEEK',IntegerType(),True),
                           StructField('DAY_NUMBER_IN_MONTH',IntegerType(),True),
                           StructField('CALENDAR_WEEK_NUMBER',IntegerType(),True),
                           StructField('FISCAL_WEEK_NUMBER',IntegerType(),True),
                           StructField('WEEK_ENDING_DAY',StringType(),True),
                           StructField('WEEK_ENDING_DAY_ID',IntegerType(),True),
                           StructField('CALENDAR_MONTH_NUMBER',IntegerType(),True),
                           StructField('FISCAL_MONTH_NUMBER',IntegerType(),True),
                           StructField('CALENDAR_MONTH_DESC',StringType(),True),
                           StructField('CALENDAR_MONTH_ID',IntegerType(),True),
                           StructField('FISCAL_MONTH_DESC',StringType(),True),
                           StructField('FISCAL_MONTH_ID',IntegerType(),True),
                           StructField('DAYS_IN_CAL_MONTH',IntegerType(),True),
                           StructField('DAYS_IN_FIS_MONTH',IntegerType(),True),
                           StructField('END_OF_CAL_MONTH',StringType(),True),
                           StructField('END_OF_FIS_MONTH',StringType(),True),
                           StructField('CALENDAR_MONTH_NAME',StringType(),True),
                           StructField('FISCAL_MONTH_NAME',StringType(),True),
                           StructField('CALENDAR_QUARTER_DESC',StringType(),True),
                           StructField('CALENDAR_QUARTER_ID',IntegerType(),True),
                           StructField('FISCAL_QUARTER_DESC',StringType(),True),
                           StructField('FISCAL_QUARTER_ID',IntegerType(),True),
                           StructField('DAYS_IN_CAL_QUARTER',IntegerType(),True),
                           StructField('DAYS_IN_FIS_QUARTER',IntegerType(),True),
                           StructField('END_OF_CAL_QUARTER',StringType(),True),
                           StructField('END_OF_FIS_QUARTER',StringType(),True),
                           StructField('CALENDAR_QUARTER_NUMBER',IntegerType(),True),
                           StructField('FISCAL_QUARTER_NUMBER',IntegerType(),True),
                           StructField('CALENDAR_YEAR',IntegerType(),True),
                           StructField('CALENDAR_YEAR_ID',IntegerType(),True),
                           StructField('FISCAL_YEAR',IntegerType(),True),
                           StructField('FISCAL_YEAR_ID',IntegerType(),True),
                           StructField('DAYS_IN_CAL_YEAR',IntegerType(),True),
                           StructField('DAYS_IN_FIS_YEAR',IntegerType(),True),
                           StructField('END_OF_CAL_YEAR',StringType(),True),
                           StructField('END_OF_FIS_YEAR',StringType(),True),
                           StructField('_corrupt_record', StringType(), False)])
  return times_schema

# COMMAND ----------

#**************************************************
#* Creating Pyspark Schema For sales transaction Table
#***************************************************
def sales_schema():
  sales_schema = StructType([StructField('PROD_ID', IntegerType(), True),
                     StructField('CUST_ID', IntegerType(), True),
                     StructField('TIME_ID', StringType(), True),
                     StructField('CHANNEL_ID', IntegerType(), True),
                     StructField('PROMO_ID', IntegerType(), True),
                     StructField('QUANTITY_SOLD', DoubleType(), True),
                     StructField('AMOUNT_SOLD', DoubleType(), True),
                     StructField('_corrupt_record', StringType(), True)])
  return sales_schema

# COMMAND ----------

#**************************************************
#* Creating Pyspark Schema For costs transaction Table
#***************************************************
def costs_schema():
  costs_schema = StructType([StructField('PROD_ID',IntegerType(),True),
                               StructField('TIME_ID',StringType(),True),
                               StructField('PROMO_ID',IntegerType(),True),
                               StructField('CHANNEL_ID',IntegerType(),True),
                               StructField('UNIT_COST',DoubleType(),True),
                               StructField('UNIT_PRICE',DoubleType(),True),
                               StructField('_corrupt_record', StringType(), True)])
  return costs_schema


"""
File containing code that related to n-level extraction of address
"""

from typing import List, Tuple, Dict
from copy import deepcopy

import pandas as pd
from flashtext import KeywordProcessor
from halo import Halo

from preprocessing_pgp.address.utils import (
    flatten_list,
    remove_substr,
    create_dependent_query
)
from preprocessing_pgp.address.const import (
    METHOD_REFER_DICT,
    LOCATION_ENRICH_DICT
)

class LevelExtractor:
    """
    Class contains information of extract level and functions to process on each level

    * Level 1: City, Countryside
    * Level 2: District
    * Level 3: Ward
    """

    def __init__(self) -> None:
        self.avail_levels = METHOD_REFER_DICT.keys()
        self.avail_methods = flatten_list(METHOD_REFER_DICT.values())
        self.keyword_refer_dict =\
            dict(zip(
                self.avail_methods,
                [self._generate_keyword_processor(method)
                 for method in self.avail_methods]
            ))

    def _get_level_methods(self, level) -> List:
        """
        _summary_

        Parameters
        ----------
        level : _type_
            _description_

        Returns
        -------
        List
            _description_
        """
        return METHOD_REFER_DICT[level]

    def _get_level_keywords(self, level: int = 1) -> List:
        """
        Function to get all keywords of each specific level

        Parameters
        ----------
        level : int, optional
            The number of level (1-3) to generate method keywords, by default 1

        Returns
        -------
            The output list of keywords for each of the level's methods
        """
        level_kws = []
        for method in self._get_level_methods(level):
            method_kw =\
                self._generate_keyword_processor(method)
            level_kws.append(method_kw)

        return level_kws

    def _generate_keyword_processor(self, method: str) -> KeywordProcessor:
        """
        Function to generate `KeywordProcessor` object
        to the specific `method`
        using the generated dictionary of keywords

        Parameters
        ----------
        method : str
            The method string refer to the dictionary column name

        Returns
        -------
        KeywordProcessor
            `KeywordProcessor` object that contains all the keywords of the method
        """
        unique_method_kws =\
            LOCATION_ENRICH_DICT[method].unique().tolist()

        keyword_processor = KeywordProcessor(case_sensitive=True)
        keyword_processor.add_keywords_from_list(unique_method_kws)

        return keyword_processor

    def _get_match_keyword(
        self,
        query: str,
        kw_processor: KeywordProcessor
    ) -> str:
        """
        Function to get the last keyword found in `query` string

        Parameters
        ----------
        query : str
            The input query string containing the address
        kw_processor : KeywordProcessor
            The keyword processor for specific method in dictionary (norm, abbrev, ...)

        Returns
        -------
        str
            Last match keyword found in query, if not found return None
        """
        found_keywords = kw_processor.extract_keywords(query)

        if len(found_keywords) > 0:
            return found_keywords[-1]

        return None

    def extract_all_levels(self, address: str) -> Tuple[Dict, str]:
        """
        Traverse through all possible level from lowest to highest
        to check and find `best pattern` in each level
        and `return` as a `dictionary of pattern` found at each level with `remaining address`

        Parameters
        ----------
        address : str
            The unified address that has been `lowered` and `unidecode`

        Returns
        -------
        Tuple[Dict, str]
            The output contains:
            * a dictionary of `pattern` found at each level
            * `remained address`
            * a dictionary of `best pattern` found at each level
        """

        remained_address = deepcopy(address)
        found_patterns =\
            dict(zip(
                self.avail_levels,
                [None]*len(self.avail_levels)
            ))
        best_patterns = deepcopy(found_patterns)

        dependents = []
        for level in self.avail_levels:
            (level_pattern,
             remained_address,
             level_method,
             level_best_pattern) =\
                self._extract_by_level(remained_address, level, *dependents)

            found_patterns[level] = level_pattern
            best_patterns[level] = level_best_pattern
            dependents.append((level_pattern, level_method))

        return found_patterns, remained_address, best_patterns

    def __is_query_exist(
        self,
        query: str
    ) -> bool:
        """
        Helper function to check whether the query is possibly found in location data
        """
        n_result = LOCATION_ENRICH_DICT.query(query).shape[0]

        return n_result > 0

    def __make_dependent_query(
        self,
        *dependents
    ) -> str:
        """
        Helper to make search query from dependents
        """
        query = ''
        # Making dependent queries
        if len(dependents) > 0:
            dependent_queries = []
            for d_term, d_method in dependents:
                if d_term is not None:
                    term_query = f'{d_method} == "{d_term}"'
                    dependent_queries.append(term_query)

            query = create_dependent_query(*dependent_queries)

        return query

    def __is_correct_dependent(
        self,
        *dependents
    ) -> bool:
        """
        Helper function to check whether the dependencies are all correct
        """
        query = self.__make_dependent_query(*dependents)

        if query != '':
            return self.__is_query_exist(query)

        return False

    def _extract_by_level(
        self,
        address: str,
        level: int,
        *dependents
    ) -> Tuple[str, str, str, str]:
        """
        Extract address with list of `method`
        at the specific `level`
        to find the best matched pattern

        Parameters
        ----------
        address : str
            The unified address that has been `lowered` and `unidecode`
        level : int
            The level which is currently traversed from

        Returns
        -------
        Tuple[str, str, str, str]
            The output contains:
            * `pattern` found within specific level
            * `remained address`
            * `method` which retrieved the pattern
            * `best pattern` unified by `pattern`
        """
        level_methods = self._get_level_methods(level)

        for method in level_methods:
            pattern_found, remained_address =\
                self._extract_by_method(address, method)

            # Found something then return
            if (pattern_found is not None)\
                    and (len(pattern_found) > 0):
                if self.__is_correct_dependent(*dependents, (pattern_found, method)):
                    best_pattern = self.__trace_match_pattern(
                        pattern_found,
                        method,
                        *dependents
                    )
                    return pattern_found, remained_address, method, best_pattern

        return None, address, None, None

    def _extract_by_method(
        self,
        address: str,
        method: str
    ) -> Tuple[str, str]:
        """
        Extract the address with `keywords` of the specific method

        Parameters
        ----------
        address : str
            The unified address that has been `lowered` and `unidecode`
        method : str
            The method of level which is currently being explored

        Returns
        -------
        Tuple[str, str]
            The output contains:
            * `pattern` found within specific method
            * `remained address`
        """
        method_kw = self.keyword_refer_dict[method]
        match_pattern = self._get_match_keyword(address, method_kw)
        remained_address = remove_substr(address, match_pattern)

        return match_pattern, remained_address

    def __trace_match_pattern(
        self,
        pattern: str,
        method: str,
        *dependents
    ) -> str:
        """
        Helper function to trace back best pattern found
        """
        query = self.__make_dependent_query(*dependents, (pattern, method))

        if not self.__is_query_exist(query):
            return None

        level_col = method[:3]

        best_match = self.__trace_best_match(query, level_col)

        return best_match

    def __trace_best_match(
        self,
        query: str,
        level_col: str
    ) -> str:
        """
        Helper to get the best match consist of:

        1. Exist
        2. Unique found
        """
        found_terms = LOCATION_ENRICH_DICT.query(query)[level_col].unique()

        n_found = found_terms.shape[0]
        if n_found == 0 or n_found > 1:
            return None

        return found_terms[0]

@Halo(
    text='Extracting address',
    color='cyan',
    spinner='dots7',
    text_color='magenta'
)
def extract_vi_address_by_level(
    data:  pd.DataFrame,
    address_col: str
) -> pd.DataFrame:
    """
    Function to extract vietnamese address into each specific level found by pattern

    Parameters
    ----------
    data : pd.DataFrame
        Raw data containing the address
    address_col : str
        The raw address column that need cleansing and unifying

    Returns
    -------
    pd.DataFrame
        Final data with new columns:

        * `level 1`: city, countryside found
        * `level 2`: district found
        * `level 3`: ward found
        * `remained address`: the remaining in the address
    """

    extractor = LevelExtractor()

    extracted_data = data.copy()

    extracted_results =\
        extracted_data[address_col].apply(
            extractor.extract_all_levels
        )

    for level in extractor.avail_levels:
        extracted_data[f'level {level}'] =\
            [result[0][level]
             for result in extracted_results]

        extracted_data[f'best level {level}'] =\
            [result[-1][level]
             for result in extracted_results]

    extracted_data['remained address'] =\
        [result[1]
         for result in extracted_results]

    return extracted_data


import math
import torch
from torch import nn
from timm.models.layers import DropPath, trunc_normal_
from einops import rearrange

from .build import MODEL_REGISTRY

class LearnedPosMapT(nn.Module):
    def __init__(self, win_size,gamma=1):
        super().__init__()
        self.gamma = gamma
        self.win_size =win_size
        self.token_proj_n_bias = nn.Parameter(torch.zeros(1 ,self.win_size,1))
        self.rel_locl_init(self.win_size,register_name='window')
        self.init_bias_table()

    def rel_locl_init(self,win_size, register_name='window'):

        h= win_size
        w= win_size
        self.register_parameter(f'{register_name}_relative_position_bias_table' ,nn.Parameter(
            torch.zeros(2*h-1, self.gamma)))

        coords_h = torch.arange(h)
        coords_w = torch.arange(w)
        relative_coords = coords_h.unsqueeze(1) - coords_w.unsqueeze(0)
        relative_coords = relative_coords.unsqueeze(0).permute(1, 2, 0).contiguous()
        relative_coords[:, :, 0] += h - 1
        relative_position_index = relative_coords
        self.register_buffer(f"{register_name}_relative_position_index", relative_position_index)

    def init_bias_table(self):
        for k,v in self.named_modules():
            if 'relative_position_bias_table' in k:
                trunc_normal_(v.weight, std=.02)

    def forward_pos(self):
        posmap = self.window_relative_position_bias_table[self.window_relative_position_index.view(-1)].view(
            self.win_size, self.win_size, -1)
        posmap =  posmap.permute(2, 0, 1).contiguous()  #weizhi seq seq
        return posmap

    def forward(self,x,weight=None,mask=None):
        posmap = self.forward_pos()
        x = rearrange(x,'b w n (v s) -> b w n v s', s = self.gamma)
        win_weight = rearrange(posmap,'(s b) m n  -> b m n s', s = self.gamma)  #m,b都是seq_len = win_size*wen_size

        x = torch.einsum('wmns,bwnvs->bwmvs',win_weight,x) + self.token_proj_n_bias.unsqueeze(0).unsqueeze(-1)
        x = rearrange(x,'b w n v s -> b w n (v s)')

        return x

class PoSGUT(nn.Module):

    def __init__(self, dim, win_size,chunks=2, norm_layer=nn.LayerNorm,quadratic=True, gamma=16,pos_only=True):
        super().__init__()
        self.chunks = chunks
        self.gate_dim = dim // chunks
        self.pos=LearnedPosMapT(win_size,gamma=gamma)

    def forward(self, x,mask = None):
           # B W N C
        if self.chunks==1:
            u = x
            v = x
        else:
            x_chunks = x.chunk(2, dim=-1)
            u = x_chunks[0]
            v = x_chunks[1]
        u = self.pos(u)
        u = u * v
        return u

class PosMLPLayerT(nn.Module):
    def __init__(self, dim, win_size,shift = False, gate_unit=PoSGUT, #num_blocks=1,
    chunks = 2,  drop=0., act_layer=nn.GELU,norm_layer=nn.LayerNorm,gamma=8):
        super().__init__()
        chunks = chunks
        self.dim =dim
        self.norm = norm_layer(dim)
        self.gate_unit = gate_unit(dim, win_size = win_size,chunks=chunks, norm_layer=norm_layer,gamma=gamma)
        self.act = act_layer()
        self.fc2 = nn.Linear(dim,dim)
        self.fc1 = nn.Linear(dim , dim*2)
        self.drop = nn.Dropout(drop)

    def forward_gate(self,x,mask= None):
        x = self.act(self.fc1(self.norm(x)))
        # 1 c b n
        x = self.gate_unit(x,mask=mask)
        x = self.drop(self.fc2(x))
        return x

    def forward_noshift(self,x):
                # Input : x (1,b,n,c)
        return self.forward_gate(x)

    def forward(self,x):
        x = self.forward_noshift(x)
        return x

class LearnedPosMap(nn.Module):
    def __init__(self, win_size,gamma=1):
        super().__init__()
        self.gamma = gamma
        self.win_size =win_size
        self.seq_len = win_size*win_size
        self.token_proj_n_bias = nn.Parameter(torch.zeros(1 ,self.seq_len,1))
        self.rel_locl_init(self.win_size,register_name='window')
        self.init_bias_table()

    def rel_locl_init(self,win_size, register_name='window'):

        h= win_size
        w= win_size
        self.register_parameter(f'{register_name}_relative_position_bias_table' ,nn.Parameter(
            torch.zeros((2 * h - 1) * (2 * w - 1), self.gamma)))

        coords_h = torch.arange(h)
        coords_w = torch.arange(w)
        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))
        coords_flatten = torch.flatten(coords, 1)
        relative_coords = coords_flatten.unsqueeze(2) - coords_flatten.  unsqueeze(1)
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()
        relative_coords[:, :, 0] += h - 1
        relative_coords[:, :, 1] += w - 1
        relative_coords[:, :, 0] *= 2 * w - 1

        relative_position_index = relative_coords.sum(-1)
        self.register_buffer(f"{register_name}_relative_position_index", relative_position_index)

    def init_bias_table(self):
        for k,v in self.named_modules():
            if 'relative_position_bias_table' in k:
                trunc_normal_(v.weight, std=.02)

    def forward_pos(self):
        posmap = self.window_relative_position_bias_table[self.window_relative_position_index.view(-1)].view(
            self.seq_len, self.seq_len, -1)
        posmap =  posmap.permute(2, 0, 1).contiguous()
        return posmap

    def forward(self,x,weight=None,mask=None):
        posmap = self.forward_pos()
        x = rearrange(x,'b w n (v s) -> b w n v s', s = self.gamma)             #x是b,win_size,win_size,v,s
        win_weight = rearrange(posmap,'(s b) m n  -> b m n s', s = self.gamma)  #m,b都是seq_len = win_size*wen_size
        if weight is not None:
            win_weight = win_weight + weight.unsqueeze(-1)
        else:
            win_weight = win_weight
        x = torch.einsum('wmns,bwnvs->bwmvs',win_weight,x) + self.token_proj_n_bias.unsqueeze(0).unsqueeze(-1)
        x = rearrange(x,'b w n v s -> b w n (v s)')

        return x

class PoSGU(nn.Module):

    def __init__(self, dim, win_size,chunks=2, norm_layer=nn.LayerNorm,quadratic=True, gamma=16,pos_only=True):
        super().__init__()
        self.chunks = chunks
        self.gate_dim = dim // chunks
        self.seq_len=win_size*win_size
        self.quadratic = quadratic
        self.pos_only=pos_only
        self.pos=LearnedPosMap(win_size,gamma=gamma)

        if not self.pos_only:
            self.token_proj_n_weight = nn.Parameter(torch.zeros(1, self.seq_len, self.seq_len))
            trunc_normal_(self.token_proj_n_weight,std=1e-6)

    def forward(self, x,mask = None):
        if self.chunks==1:
            u = x
            v = x
        else:
            x_chunks = x.chunk(2, dim=-1)
            u = x_chunks[0]
            v = x_chunks[1]
        if not self.pos_only and not self.quadratic:
            u = self.pos(u,self.token_proj_n_weight)
        else:
            u =self.pos(u,mask=mask)

        u = u * v
        return u

class PosMLPLayer(nn.Module):
    def __init__(self, dim, win_size,shift = False, gate_unit=PoSGU, num_blocks=1,
    chunks = 2, drop=0., act_layer=nn.GELU,norm_layer=nn.LayerNorm,gamma=8):
        super().__init__()
        chunks = chunks
        shift_size = 0
        self.shift = shift
        # if shift and num_blocks != 1 and layer_idx %2 ==1:
        #     shift_size = win_size//2
        #     _logger.info(f"shifted layer {layer_idx}")
        self.shift_size = shift_size
        self.dim =dim
        self.norm = norm_layer(dim)
        self.window_size = win_size
        self.gate_unit = gate_unit(dim, win_size = win_size,chunks=chunks, norm_layer=norm_layer,gamma=gamma)
        self.act = act_layer()
        self.fc2 = nn.Linear(dim , dim)
        self.fc1 = nn.Linear(dim , dim*2)
        self.drop = nn.Dropout(drop)

    def forward_gate(self,x,mask= None):
        x = self.act(self.fc1(self.norm(x)))
        # 1 c b n
        x = self.gate_unit(x,mask=mask)
        x = self.drop(self.fc2(x))

        return x

    def forward_noshift(self,x):
                # Input : x (1,b,n,c)
        return self.forward_gate(x)

    def forward_shift(self,x):
        # @swin transformer
        # @http://arxiv.org/abs/2103.14030

        _, H, W, C = x.size()

        if self.shift_size > 0:
            # print("shift,here")
            shifted_x= torch.roll(
                x,
                shifts=(-self.shift_size, -self.shift_size),
                dims=(1, 2))
            img_mask = torch.zeros((1, H, W, 1), device=x.device)
            h_slices = (slice(0, -self.window_size),
                        slice(-self.window_size,
                              -self.shift_size), slice(-self.shift_size, None))
            w_slices = (slice(0, -self.window_size),
                        slice(-self.window_size,
                              -self.shift_size), slice(-self.shift_size, None))
            cnt = 0
            for h in h_slices:
                for w in w_slices:
                    img_mask[:, h, w, :] = cnt
                    cnt += 1

            # 1,nW, window_size*window_size, 1
            mask_windows = self.window_partition(img_mask)
            attn_mask = mask_windows.unsqueeze(2) - mask_windows.unsqueeze(3)
            attn_mask = attn_mask.masked_fill(attn_mask != 0,
                                              float(-200.0)).masked_fill(
                                                  attn_mask == 0, float(0.0)).squeeze()
        else:
            shifted_x = x
            attn_mask =None
        # _,NM,N,C
        shifted_x = window_partition(shifted_x,self.window_size,self.window_size)
        out = self.forward_gate(shifted_x,mask = attn_mask)
        shifted_x = window_reverse(out, H, W)
        # reverse cyclic shift
        if self.shift_size > 0:
            x = torch.roll(
                shifted_x,
                shifts=(self.shift_size, self.shift_size),
                dims=(1, 2))
        else:
            x = shifted_x
        return x

    def forward(self,x):
        if self.shift:
            x = self.forward_shift(x)
        else:
            x = self.forward_noshift(x)
        return x

def window_reverse(x, H, W):

    B,_,N,C=x.size()
    Wh= int(math.sqrt(N))
    Ww = Wh
    x= x.view(B, -1, Wh,  Ww, C)
    x = x.view(B, H // Wh, W // Ww, Wh, Ww, -1)
    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
    return x

def window_partition( x,Wh,Ww):

    B, H, W, C = x.shape
    x = x.view(B, H //Wh, Wh, W // Ww,Ww, C)
    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous()
    windows = windows.view(B,-1, Wh*Ww, C)
    return windows

class PosMLPLevelS(nn.Module):
    """ Single hierarchical level of a Nested Transformer
    """
    def __init__(
            self,gate_unit, win_size,  embed_dim, gate_layer=PosMLPLayer,
            drop_rate=0., gamma = 8,
            norm_layer=None, act_layer=None,  shift = False):

        super().__init__()
        self.win_size=win_size
        self.shift = shift
        self.encoder = gate_layer(dim= embed_dim,win_size=win_size,gate_unit=gate_unit,
                                drop = drop_rate,shift =shift ,gamma = gamma,
                                norm_layer=norm_layer,act_layer=act_layer)
    def forward(self, x):
        """
        expects x as (1, C, H, W)
        """
        B,T,H,W,C = x.shape
        bt = B*T
        x = x.reshape(bt,H,W,C) #BT,H,W,Cs
        if self.shift :
            x = self.encoder(x)  # (1, H, W, C)
        else:
            # M,H,W,C = x.shape
            x = window_partition(x,self.win_size,self.win_size)  # (1, T, N, C)
            x = self.encoder(x)  # (1, T, N, C)
            x = window_reverse(x, H, W)  # (1, H, W, C)
        x = x.reshape(B,T,H,W,C)
        return x

class PosMLPLevelT(nn.Module):
    """ Single hierarchical level of a Nested Transformer
    """
    def __init__(
            self,gate_unit, win_size_t, embed_dim, gate_layer=PosMLPLayerT,
             drop_rate=0., gamma = 8,
            norm_layer=None, act_layer=None,  shift = False):

        super().__init__()
        self.win_size_t=win_size_t
        self.shift = shift
        self.encoder = gate_layer(dim= embed_dim,win_size=win_size_t,gate_unit=gate_unit,
                                 drop = drop_rate,shift = shift ,gamma = gamma,
                                 norm_layer=norm_layer,act_layer=act_layer)
    def forward(self, x):
        """
        expects x as (B,T,H,W,C)
        """
        B, T, H, W, C = x.shape
        hw = H*W
        x = x.reshape(B,T,hw,C).permute(0,2,1,3)
        x = self.encoder(x)
        x = x.permute(0,2,1,3).reshape(B,T,H,W,C)
        return x

def conv_3xnxn(inp, oup, kernel_size=3, stride=3):
    return nn.Conv3d(inp, oup, (1, kernel_size, kernel_size), (1, stride, stride), (0, 1, 1))

def conv_1xnxn(inp, oup, kernel_size=3, stride=3):
    return nn.Conv3d(inp, oup, (1, kernel_size, kernel_size), (1, stride, stride), (0, 1, 1))

class PermutatorBlock(nn.Module):
    def __init__(self, dim,
                 drop_path=0.,act_layer=None, norm_layer=None, gate_unit_t = PoSGUT,
                 gate_unit=PoSGU, win_size=None,win_size_t = None, gate_layer=PosMLPLayer,gate_layer_t = PosMLPLayerT,
                 drop_rate=0., gamma = 1,
                 shift = False):
        super().__init__()

        self.win_size=win_size
        self.shift = shift
        self.t_fc = PosMLPLevelT(gate_unit_t,win_size_t,embed_dim=dim, gate_layer=gate_layer_t,gamma=gamma,
                                drop_rate=drop_rate, norm_layer=norm_layer, act_layer=act_layer)

        self.fc = PosMLPLevelS(gate_unit,win_size, embed_dim=dim,gate_layer=gate_layer,gamma=gamma,shift = self.shift,
                         drop_rate=drop_rate, norm_layer=norm_layer, act_layer=act_layer)
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

    def forward(self, x):
        x = x + self.drop_path(self.fc(x) + self.t_fc(x))

        return x

class PatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """

    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        self.proj1 = conv_3xnxn(in_chans, embed_dim//2, kernel_size=3, stride=2)
        self.norm1= nn.BatchNorm3d(embed_dim//2)
        self.act=nn.GELU()
        self.proj2 = conv_1xnxn(embed_dim//2, embed_dim, kernel_size=3, stride=2)
        self.norm2 = nn.BatchNorm3d(embed_dim)

    def forward(self, x):
        x = self.proj1(x)
        x = self.norm1(x)
        x = self.act(x)
        x = self.proj2(x)
        x = self.norm2(x)
        return x

class Downsample(nn.Module):
    """ Image to Patch Embedding
    """

    def __init__(self, in_embed_dim, out_embed_dim, patch_size):
        super().__init__()
        self.proj = conv_1xnxn(in_embed_dim, out_embed_dim, kernel_size=3, stride=2)
        self.norm=nn.LayerNorm(out_embed_dim)

    def forward(self, x):
        x = x.permute(0, 4, 1, 2, 3)
        x = self.proj(x)  # B, C, T, H, W
        x = x.permute(0, 2, 3, 4, 1)
        x=self.norm(x)
        return x

@MODEL_REGISTRY.register()
class PosMLP_Video(nn.Module):
    """ PosMLP_Video
    """

    def __init__(self, cfg):
        super().__init__()

        win_size = cfg.WIN_SIZE
        win_size_t = cfg.WIN_SIZE_T
        gamma = cfg.POSGAMMA
        shift = cfg.POSMLP_VIDEO.SHIFT
        num_classes = cfg.MODEL.NUM_CLASSES
        img_size = cfg.DATA.TRAIN_CROP_SIZE
        in_chans = cfg.DATA.INPUT_CHANNEL_NUM[0]
        layers = cfg.POSMLP_VIDEO.LAYERS
        transitions = cfg.POSMLP_VIDEO.TRANSITIONS
        embed_dims = cfg.POSMLP_VIDEO.EMBED_DIMS
        patch_size = cfg.POSMLP_VIDEO.PATCH_SIZE
        drop_path_rate = cfg.POSMLP_VIDEO.DROP_DEPTH_RATE
        norm_layer = nn.LayerNorm

        self.num_classes = num_classes

        self.patch_embed1 = PatchEmbed(img_size=img_size, patch_size=patch_size, in_chans=in_chans,
                                       embed_dim=embed_dims[0])
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(layers))]

        # stage1
        self.blocks1 = nn.ModuleList([])
        for i in range(layers[0]):
            self.blocks1.append(
                PermutatorBlock(embed_dims[0], drop_path=dpr[i],shift = shift,gamma = gamma[0],
                                act_layer=nn.GELU,norm_layer=nn.LayerNorm, win_size=win_size[0], win_size_t=win_size_t[0],
                                )
            )
        if transitions[0] or embed_dims[0] != embed_dims[1]:
            patch_size = 2 if transitions[0] else 1
            self.patch_embed2 = Downsample(embed_dims[0], embed_dims[1], patch_size)
        else:
            self.patch_embed2 = nn.Identity()
        # stage2
        self.blocks2 = nn.ModuleList([])
        for i in range(layers[1]):
            self.blocks2.append(
                PermutatorBlock(embed_dims[1], drop_path=dpr[i + layers[0]], shift = shift,gamma = gamma[1],
                                 act_layer=nn.GELU,norm_layer=nn.LayerNorm, win_size=win_size[1], win_size_t=win_size_t[1],
                                 )
            )
        if transitions[1] or embed_dims[1] != embed_dims[2]:
            patch_size = 2 if transitions[1] else 1
            self.patch_embed3 = Downsample(embed_dims[1], embed_dims[2], patch_size)
        else:
            self.patch_embed3 = nn.Identity()
        # stage3
        self.blocks3 = nn.ModuleList([])
        for i in range(layers[2]):
            self.blocks3.append(
               PermutatorBlock(embed_dims[2],
                                drop_path=dpr[i + layers[0] + layers[1]],
                                shift = shift,
                                act_layer=nn.GELU,
                                norm_layer=nn.LayerNorm,
                                win_size=win_size[2],
                                win_size_t=win_size_t[2],
                                gamma = gamma[2]
                                )
            )
        if transitions[2] or embed_dims[2] != embed_dims[3]:
            patch_size = 2 if transitions[2] else 1
            self.patch_embed4 = Downsample(embed_dims[2], embed_dims[3], patch_size)
        else:
            self.patch_embed4 = nn.Identity()

        # stage4
        self.blocks4 = nn.ModuleList([])
        for i in range(layers[3]):
            self.blocks4.append(
                PermutatorBlock(embed_dims[3], gamma = gamma[3],drop_path=dpr[i + layers[0]+layers[1]+layers[2]], shift = shift,
                                 act_layer=nn.GELU,norm_layer=nn.LayerNorm, win_size=win_size[3], win_size_t=win_size_t[3],
                                 )
            )

        self.norm = norm_layer(embed_dims[-1])

        # Classifier head
        self.head = nn.Linear(embed_dims[-1], num_classes) if num_classes > 0 else nn.Identity()
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def get_pretrained_model(self, cfg):
        if cfg.POSMLP_VIDEO.PRETRAIN_PATH:
            checkpoint = torch.load(cfg.POSMLP_VIDEO.PRETRAIN_PATH, map_location='cpu')
            if self.num_classes != 1000:
                del checkpoint['head.weight']
                del checkpoint['head.bias']
            return checkpoint
        else:
            return None

    def forward_features(self, x):
        x = self.patch_embed1(x[0])
        # B,C,T,H,W -> B,T,H,W,C
        x = x.permute(0, 2, 3,4, 1)
        for blk in self.blocks1:
            x = blk(x)
        x = self.patch_embed2(x)
        for blk in self.blocks2:
            x = blk(x)
        x = self.patch_embed3(x)
        for blk in self.blocks3:
            x = blk(x)
        x = self.patch_embed4(x)
        for blk in self.blocks4:
            x = blk(x)
        B, T,H, W, C = x.shape
        x = x.reshape(B, -1, C)
        return x

    def forward(self, x):
        x = self.forward_features(x)
        x = self.norm(x)
        return self.head(x.mean(1))


from googlesearch import search
import requests
from bs4 import BeautifulSoup
import csv

def join_results(hits_list):
    hits = ''
    for x in hits_list:
        hits += str(x)
    return hits

def grab_links(site, word):
    links_list = []
    for l in search(site+word, tld='com', lang='en', num=10, start=0, stop=10, pause=2.0):
        links_list.append(l)
    return links_list

def grab_results(site, word):
    with requests.Session() as s:
        url = f"http://www.google.com/search?q={site}+{word}"
        headers = {
            "referer":"referer: https://www.google.com/",
            "user-agent":"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
                         " (KHTML, like Gecko) Chrome/89.0.4389.114 Safari/537.36"
            }
        s.post(url, headers=headers)
        response = s.get(url, headers=headers)
        soup = BeautifulSoup(response.content, 'html.parser')
        result_list = soup.find_all(id='result-stats')[0].text
        results_string = join_results([int(s) for s in result_list.split() if s.isdigit()])
    return f"{word}: {results_string}"

keywords = open('keywords.txt', 'r').read().split(" ")
query = 'site:https://www.searchenginejournal.com/'

for w in keywords:
    links = grab_links(query, w)
    with open("links.csv", "a", encoding='utf-8') as w_file:
        file_writer = csv.writer(w_file, delimiter=",", lineterminator="\n")
        for link in links:
            file_writer.writerow(link)
    results = grab_results(query, w)
    with open("results.csv", "a", encoding='utf-8') as w_file:
        file_writer = csv.writer(w_file, delimiter=",", lineterminator="\n")
        for result in results:
            file_writer.writerow(result)


import glob
import os.path
from os.path import basename
from typing import Dict, Optional

from PIL import Image, ImageFile

import tifffile

import torch
from torch import Tensor
from torch.utils.data import Dataset

ImageFile.LOAD_TRUNCATED_IMAGES = True

class DaVinciEnsembleDataset(Dataset):
    """Dataset class for loading the Hamlyn da Vinci images.

    Given the root of the dataset path, this class will find all left and
    right `.png` images and collect each pair as a dictionary of Tensors.

    If there are any missing image IDs from either left or right folders,
    the pair is ignored.

    Note:
        Transforms must be able to handle dictionaries containing left and
        right views as separate entries.

    Args:
        root (str): Path to the root of the dataset directory.
        split (str): The folder in the dataset to use. Must be "train" or
            "test".
        transform (Optional[object], optional): The transforms to apply to
            each image pair while loading. Defaults to None.
        limit (Optional[int], optional): The maximum number of images to load.
            Loads all images if None. Defaults to None.
    """
    LEFT_PATH = 'image_0'
    RIGHT_PATH = 'image_1'
    IMAGE_EXT = 'png'
    ENSEMBLE_EXT = 'tiff'

    def __init__(self, davinci_path: str, ensemble_path: str, split: str,
                 transform: Optional[object] = None,
                 limit: Optional[int] = None) -> None:

        if split not in ('train', 'test'):
            raise ValueError('Split must be either "train" or "test".')

        left_glob = os.path.join(davinci_path, split, self.LEFT_PATH,
                                 f'*.{self.IMAGE_EXT}')

        right_glob = os.path.join(davinci_path, split, self.RIGHT_PATH,
                                  f'*.{self.IMAGE_EXT}')

        ensemble_glob = os.path.join(ensemble_path, split,
                                     f'*.{self.ENSEMBLE_EXT}')

        left = glob.glob(left_glob)
        right = glob.glob(right_glob)

        ensemble = glob.glob(ensemble_glob)

        left_names = set(map(self.name, left))
        right_names = set(map(self.name, right))
        ensemble_names = set(map(self.name, ensemble))

        stereo_missing = left_names.symmetric_difference(right_names)
        ensemble_missing = left_names.symmetric_difference(ensemble_names)

        missing = set.union(stereo_missing, ensemble_missing)

        if len(missing) > 0:
            print(f'Missing {len(missing):,} images from the dataset.')
            left = [i for i in left if self.name(i) not in missing]
            right = [i for i in right if self.name(i) not in missing]
            ensemble = [i for i in ensemble if self.name(i) not in missing]
            print(f'Dataset reduced to {len(left):,} images.')

        self.lefts = sorted(left[:limit])
        self.rights = sorted(right[:limit])

        self.ensembles = sorted(ensemble[:limit])

        self.transform = transform

    @staticmethod
    def name(path: str) -> str:
        return os.path.splitext(basename(path))[0]

    def __getitem__(self, idx: int) -> Dict[str, Tensor]:
        """Retrieve a single sample from the dataset.

        Args:
            idx (int): The index of the sample in the dataset.

        Returns:
            Dict[str, Tensor]: The left and right images packaged as a
                dictionary containing `left` and `right` keys.
        """
        left_path = self.lefts[idx]
        right_path = self.rights[idx]

        ensemble_path = self.ensembles[idx]

        left = Image.open(left_path).convert('RGB')
        right = Image.open(right_path).convert('RGB')

        ensemble_prediction = tifffile.imread(ensemble_path)

        image_pair = {
            'left': left, 'right': right,
            'ensemble': torch.from_numpy(ensemble_prediction)
        }

        if self.transform is not None:
            image_pair = self.transform(image_pair)

        return image_pair

    def __len__(self) -> int:
        return len(self.lefts)


from django.shortcuts import render, redirect
from django.contrib.auth import login,logout,authenticate
from django.contrib.auth.forms import UserChangeForm
from .forms import RegisterForm, ProfileForm

# Create your views here.
def register(request):
    if request.method == "POST":
        form = RegisterForm(request.POST)
        if form.is_valid():
            user = form.save()
            login(request,user)
            return redirect("/home")
    else:
        form = RegisterForm()
    return render(request,'register.html',{"form":form})

# def login(request):
def usereditview(request):
    if request.method == "POST":
        form2 = ProfileForm(request.POST, instance=request.user)
        if form2.is_valid():
            form2.save()
            return redirect("/home")
    else:
        form2 = ProfileForm()
    return render(request,'edit_profile.html',{"form2":form2})


import os

import click

from openpype.modules import (
    OpenPypeModule,
    ITrayModule,
    IPluginPaths,
)
from openpype.lib import Logger

FTRACK_MODULE_DIR = os.path.dirname(os.path.abspath(__file__))
_URL_NOT_SET = object()

class FtrackModule(
    OpenPypeModule,
    ITrayModule,
    IPluginPaths,
):
    name = "ftrack"
    enabled = True

    def initialize(self, settings):
        ftrack_settings = settings[self.name]

        self._settings_ftrack_url = ftrack_settings["ftrack_server"]
        self._ftrack_url = _URL_NOT_SET

        current_dir = os.path.dirname(os.path.abspath(__file__))

        # User event handler paths
        user_event_handlers_paths = [
            os.path.join(current_dir, "event_handlers_user")
        ]

        # Prepare attribute
        self.user_event_handlers_paths = user_event_handlers_paths
        self.tray_module = None

        # TimersManager connection
        self.timers_manager_connector = None
        self._timers_manager_module = None

    def get_ftrack_url(self):
        """Resolved ftrack url.

        Resolving is trying to fill missing information in url and tried to
        connect to the server.

        Returns:
            Union[str, None]: Final variant of url or None if url could not be
                reached.
        """

        if self._ftrack_url is _URL_NOT_SET:
            self._ftrack_url = resolve_ftrack_url(
                self._settings_ftrack_url,
                logger=self.log
            )
        return self._ftrack_url

    ftrack_url = property(get_ftrack_url)

    @property
    def settings_ftrack_url(self):
        """Ftrack url from settings in a format as it is.

        Returns:
            str: Ftrack url from settings.
        """

        return self._settings_ftrack_url

    def get_global_environments(self):
        """Ftrack's global environments."""

        return {
            "FTRACK_SERVER": self.ftrack_url
        }

    def get_plugin_paths(self):
        """Ftrack plugin paths."""
        return {
            "publish": [os.path.join(FTRACK_MODULE_DIR, "plugins", "publish")]
        }

    def get_launch_hook_paths(self):
        """Implementation for applications launch hooks."""

        return os.path.join(FTRACK_MODULE_DIR, "launch_hooks")

    def modify_application_launch_arguments(self, application, env):
        if not application.use_python_2:
            return

        self.log.info("Adding Ftrack Python 2 packages to PYTHONPATH.")

        # Prepare vendor dir path
        python_2_vendor = os.path.join(FTRACK_MODULE_DIR, "python2_vendor")

        # Add Python 2 modules
        python_paths = [
            # `python-ftrack-api`
            os.path.join(python_2_vendor, "ftrack-python-api", "source"),
        ]

        # Load PYTHONPATH from current launch context
        python_path = env.get("PYTHONPATH")
        if python_path:
            python_paths.append(python_path)

        # Set new PYTHONPATH to launch context environments
        env["PYTHONPATH"] = os.pathsep.join(python_paths)

    def connect_with_modules(self, enabled_modules):
        for module in enabled_modules:
            if not hasattr(module, "get_ftrack_event_handler_paths"):
                continue

            try:
                paths_by_type = module.get_ftrack_event_handler_paths()
            except Exception:
                continue

            if not isinstance(paths_by_type, dict):
                continue

            for key, value in paths_by_type.items():
                if not value:
                    continue

                if key not in ("server", "user"):
                    self.log.warning(
                        "Unknown event handlers key \"{}\" skipping.".format(
                            key
                        )
                    )
                    continue

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

                if key == "user":
                    self.user_event_handlers_paths.extend(value)

    def create_ftrack_session(self, **session_kwargs):
        import ftrack_api

        if "server_url" not in session_kwargs:
            session_kwargs["server_url"] = self.ftrack_url

        api_key = session_kwargs.get("api_key")
        api_user = session_kwargs.get("api_user")
        # First look into environments
        # - both AYON tray and ftrack event server should have set them
        # - ftrack event server may crash when credentials are tried to load
        #   from keyring
        if not api_key or not api_user:
            api_key = os.environ.get("FTRACK_API_KEY")
            api_user = os.environ.get("FTRACK_API_USER")

        if not api_key or not api_user:
            from .lib import credentials
            cred = credentials.get_credentials()
            api_user = cred.get("username")
            api_key = cred.get("api_key")

        session_kwargs["api_user"] = api_user
        session_kwargs["api_key"] = api_key
        return ftrack_api.Session(**session_kwargs)

    def tray_init(self):
        from .tray import FtrackTrayWrapper

        self.tray_module = FtrackTrayWrapper(self)
        # Module is it's own connector to TimersManager
        self.timers_manager_connector = self

    def tray_menu(self, parent_menu):
        return self.tray_module.tray_menu(parent_menu)

    def tray_start(self):
        return self.tray_module.validate()

    def tray_exit(self):
        self.tray_module.tray_exit()

    def set_credentials_to_env(self, username, api_key):
        os.environ["FTRACK_API_USER"] = username or ""
        os.environ["FTRACK_API_KEY"] = api_key or ""

    # --- TimersManager connection methods ---
    def start_timer(self, data):
        if self.tray_module:
            self.tray_module.start_timer_manager(data)

    def stop_timer(self):
        if self.tray_module:
            self.tray_module.stop_timer_manager()

    def register_timers_manager(self, timer_manager_module):
        self._timers_manager_module = timer_manager_module

    def timer_started(self, data):
        if self._timers_manager_module is not None:
            self._timers_manager_module.timer_started(self.id, data)

    def timer_stopped(self):
        if self._timers_manager_module is not None:
            self._timers_manager_module.timer_stopped(self.id)

    def get_task_time(self, project_name, asset_name, task_name):
        session = self.create_ftrack_session()
        query = (
            'Task where name is "{}"'
            ' and parent.name is "{}"'
            ' and project.full_name is "{}"'
        ).format(task_name, asset_name, project_name)
        task_entity = session.query(query).first()
        if not task_entity:
            return 0
        hours_logged = (task_entity["time_logged"] / 60) / 60
        return hours_logged

    def get_credentials(self):
        # type: () -> tuple
        """Get local Ftrack credentials."""
        from .lib import credentials

        cred = credentials.get_credentials(self.ftrack_url)
        return cred.get("username"), cred.get("api_key")

    def cli(self, click_group):
        click_group.add_command(cli_main)

def _check_ftrack_url(url):
    import requests

    try:
        result = requests.get(url, allow_redirects=False)
    except requests.exceptions.RequestException:
        return False

    if (result.status_code != 200 or "FTRACK_VERSION" not in result.headers):
        return False
    return True

def resolve_ftrack_url(url, logger=None):
    """Checks if Ftrack server is responding."""

    if logger is None:
        logger = Logger.get_logger(__name__)

    url = url.strip("/ ")
    if not url:
        logger.error("Ftrack URL is not set!")
        return None

    if not url.startswith("http"):
        url = "https://" + url

    ftrack_url = None
    if url and _check_ftrack_url(url):
        ftrack_url = url

    if not ftrack_url and not url.endswith("ftrackapp.com"):
        ftrackapp_url = url + ".ftrackapp.com"
        if _check_ftrack_url(ftrackapp_url):
            ftrack_url = ftrackapp_url

    if not ftrack_url and _check_ftrack_url(url):
        ftrack_url = url

    if ftrack_url:
        logger.debug("Ftrack server \"{}\" is accessible.".format(ftrack_url))

    else:
        logger.error("Ftrack server \"{}\" is not accessible!".format(url))

    return ftrack_url

@click.group(FtrackModule.name, help="Ftrack module related commands.")
def cli_main():
    pass

{'ftrack_api': 'ftrack_api', 'credentials': 'lib.credentials', 'FtrackTrayWrapper': 'tray.FtrackTrayWrapper'}.name


import logging
import yaml
import os

class IndentDumper(yaml.Dumper):
    def increase_indent(self, flow=False, indentless=False):
        return super(IndentDumper, self).increase_indent(flow, False)

# 获取当前文件的绝对路径，然后构建config.yml文件的路径
current_dir = os.path.dirname(os.path.abspath(__file__))
config_file_path = os.path.join(current_dir, '../../conf/webhooks.yml')

def getconfig(filename=config_file_path):
    if os.path.exists(filename):
        with open(filename, 'r') as file:
            return yaml.safe_load(file)
    else:
        default_config = []
        saveconfig(default_config)
        return default_config

def saveconfig(config_data, filename=config_file_path):
    with open(filename, 'w') as file:
        yaml.dump(config_data, file, default_flow_style=False,
                  allow_unicode=True, sort_keys=False, Dumper=IndentDumper)

def getWebhooksByEvent(event_name):
    '''根据事件类型获取webhooks'''
    allWebhooks = getconfig()
    webhooks = []
    for webhook in allWebhooks:
        if event_name in webhook.get('bindEvents', []):
            webhooks.append(webhook)
    return webhooks

def getAllWebhookEvents():
    '''获取所有webhook事件'''
    allWebhooks = getconfig()
    events = []
    for webhook in allWebhooks:
        events.extend(webhook.get('bindEvents', []))
    return list(set(events))


"""
this code is adapted from https://github.com/black7375/korean_ocr_using_pororo

Apache License 2.0 @yunwoong7
Apache License 2.0 @black7375
"""

"""
Pororo task-specific factory class

    isort:skip_file

"""

import logging
from typing import Optional
from .tasks.utils.base import PororoTaskBase

import torch

from .tasks import (
    PororoOcrFactory,
)

SUPPORTED_TASKS = {
    "ocr": PororoOcrFactory,
}

LANG_ALIASES = {
    "english": "en",
    "eng": "en",
    "korean": "ko",
    "kor": "ko",
    "kr": "ko",
    "chinese": "zh",
    "chn": "zh",
    "cn": "zh",
    "japanese": "ja",
    "jap": "ja",
    "jp": "ja",
    "jejueo": "je",
    "jje": "je",
}

logging.getLogger("transformers").setLevel(logging.WARN)
logging.getLogger("fairseq").setLevel(logging.WARN)
logging.getLogger("sentence_transformers").setLevel(logging.WARN)
logging.getLogger("youtube_dl").setLevel(logging.WARN)
logging.getLogger("pydub").setLevel(logging.WARN)
logging.getLogger("librosa").setLevel(logging.WARN)

class Pororo:
    r"""
    This is a generic class that will return one of the task-specific model classes of the library
    when created with the `__new__()` method

    """

    def __new__(
        cls,
        task: str,
        lang: str = "en",
        model: Optional[str] = None,
        **kwargs,
    ) -> PororoTaskBase:
        if task not in SUPPORTED_TASKS:
            raise KeyError(
                "Unknown task {}, available tasks are {}".format(
                    task,
                    list(SUPPORTED_TASKS.keys()),
                )
            )

        lang = lang.lower()
        lang = LANG_ALIASES[lang] if lang in LANG_ALIASES else lang

        # Get device information from torch API
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Instantiate task-specific pipeline module, if possible
        task_module = SUPPORTED_TASKS[task](
            task,
            lang,
            model,
            **kwargs,
        ).load(device)

        return task_module

    @staticmethod
    def available_tasks() -> str:
        """
        Returns available tasks in Pororo project

        Returns:
            str: Supported task names

        """
        return "Available tasks are {}".format(list(SUPPORTED_TASKS.keys()))

    @staticmethod
    def available_models(task: str) -> str:
        """
        Returns available model names correponding to the user-input task

        Args:
            task (str): user-input task name

        Returns:
            str: Supported model names corresponding to the user-input task

        Raises:
            KeyError: When user-input task is not supported

        """
        if task not in SUPPORTED_TASKS:
            raise KeyError(
                "Unknown task {} ! Please check available models via `available_tasks()`".format(
                    task
                )
            )

        langs = SUPPORTED_TASKS[task].get_available_models()
        output = f"Available models for {task} are "
        for lang in langs:
            output += f"([lang]: {lang}, [model]: {', '.join(langs[lang])}), "
        return output[:-2]


import base64

from odoo import http
from odoo.http import content_disposition, request

class BankRecoReport(http.Controller):
    @http.route(
        "/web/binary/jic_attendance_report",
        type="http",
        auth="user",
        sitemap=False,
    )
    def download_jic_attendance_report(self, wizard_id, **kw):
        model = request.env["hr.attendance.report.wizard"]
        case = model.browse(int(wizard_id))
        filename = "jic_report_attendance.xlsx"
        if case and case.report_file:
            file_data = base64.b64decode(case.report_file)
            return request.make_response(
                file_data,
                headers=[
                    ("Content-Disposition", content_disposition(filename)),
                    ("Content-Type", "application/vnd.ms-excel"),
                    ("Content-Length", len(filename)),
                ],
                cookies={"fileToken": ""},
            )