import os
import sys
sys.path.append(os.path.abspath(os.path.dirname(os.getcwd())))

from AFC.dataloader_afc import IceLoader
from torch.utils.data import Dataset, DataLoader
import argparse
from tqdm import tqdm
from tqdm.contrib import tzip
from model.audio_visual_model_incremental import IncreAudioVisualNet
import torch
import torch.nn as nn
from torch.nn import functional as F
import matplotlib.pyplot as plt
from torch.optim.lr_scheduler import ReduceLROnPlateau, MultiStepLR
import numpy as np
from datetime import datetime
import random
from model.layers import NormedLinear, SplitNormedLinear, LSCLinear, SplitLSCLinear

from sklearn.cluster import KMeans

import copy

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

def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    torch.backends.cudnn.deterministic = True

def boolean_string(s):
    if s not in {'False', 'True'}:
        raise ValueError('Not a valid boolean string')
    return s == 'True'

def nca(similarities, targets, scale=1, margin=0.6):
    margins = torch.zeros_like(similarities)
    margins[torch.arange(margins.shape[0]), targets] = margin

    similarities = scale * (similarities - margin)

    similarities = similarities - similarities.max(1)[0].view(-1, 1)  # Stability

    disable_pos = torch.zeros_like(similarities)
    disable_pos[torch.arange(len(similarities)),
                targets] = similarities[torch.arange(len(similarities)), targets]

    numerator = similarities[torch.arange(similarities.shape[0]), targets]
    denominator = similarities - disable_pos

    losses = numerator - torch.log(torch.exp(denominator).sum(-1))

    losses = -losses
    loss = torch.mean(losses)

    return loss

def imprint_weights(args, model, step, train_loader):
    model.eval()

    class_per_step = args.class_num_per_step
    start = step * class_per_step

    features_arr = []
    targets_arr = []

    with torch.no_grad():
        print('===========================')
        print('Imprint weights')
        print('===========================')
        model = model.to(device)

        for samples in tqdm(train_loader):
            data, labels = samples
            # labels = labels.to(device)
            if args.modality == 'visual':
                visual = data
                visual = visual.to(device)
                features = model(visual=visual, out_logits=False, out_features=True, out_features_norm=True)
            elif args.modality == 'audio':
                audio = data
                audio = audio.to(device)
                features = model(audio=audio, out_logits=False, out_features=True, out_features_norm=True)
            else:
                visual = data[0]
                audio = data[1]
                visual = visual.to(device)
                audio = audio.to(device)
                features = model(visual=visual, audio=audio, out_logits=False, out_features=True, out_features_norm=True)

            features = features.detach().cpu()

            new_idx = labels >= start

            new_features = features[new_idx]
            new_target_idx = labels[new_idx] - start

            features_arr.append(new_features)
            targets_arr.append(new_target_idx)

        model = model.to('cpu')
        weights_norm = model.classifier.fc1.weight.data.norm(dim=1, keepdim=True)
        avg_weights_norm = torch.mean(weights_norm, dim=0)

        features = torch.cat(features_arr, dim=0)
        targets = torch.cat(targets_arr, dim=0)

        new_weights = []
        for c in range(class_per_step):
            class_features = features[targets == c]
            clusterizer = KMeans(n_clusters=10)
            clusterizer.fit(class_features.numpy())

            for center in clusterizer.cluster_centers_:
                new_weights.append(torch.tensor(center) * avg_weights_norm)

        new_weights = torch.stack(new_weights)
        model.classifier.fc2.weight.data = new_weights

    return model

def top_1_acc(logits, target):
    top1_res = logits.argmax(dim=1)
    top1_acc = torch.eq(target, top1_res).sum().float() / len(target)
    return top1_acc.item()

def adjust_learning_rate(args, optimizer, epoch):
    miles_list = np.array(args.milestones) - 1
    if epoch in miles_list:
        current_lr = optimizer.param_groups[0]['lr']
        new_lr = current_lr * 0.1
        print('Reduce lr from {} to {}'.format(current_lr, new_lr))
        for param_group in optimizer.param_groups:
            param_group['lr'] = new_lr

def gen_step_exemplar_set(args, model, step, train_data_set, m_per_class, old_exemplar_class_vids):
    model.eval()

    train_data_set._switch_gen_exemplar(old_exemplar_class_vids)
    gen_exemplar_loader = DataLoader(train_data_set, batch_size=args.gen_exem_batch_size, num_workers=args.num_workers,
                                     pin_memory=True, drop_last=False, shuffle=False)

    all_features = torch.Tensor([])
    # all_vids = torch.Tensor([])
    all_vids = []
    all_category_id = torch.Tensor([])
    with torch.no_grad():
        for data, vids, category_id in tqdm(gen_exemplar_loader):
            if args.modality == 'visual':
                visual = data
                visual = visual.to(device)
                features = model(visual=visual, out_logits=False, out_features=True, out_features_norm=True).detach().cpu()
            elif args.modality == 'audio':
                audio = data
                audio = audio.to(device)
                features = model(audio=audio, out_logits=False, out_features=True, out_features_norm=True).detach().cpu()
            else:
                visual = data[0]
                audio = data[1]
                visual = visual.to(device)
                audio = audio.to(device)
                features = model(visual=visual, audio=audio, out_logits=False, out_features=True, out_features_norm=True).detach().cpu()
            all_features = torch.cat((all_features, features), dim=0)
            # print(vids)
            all_vids += vids
            # all_vids = torch.cat((all_vids, vids), dim=0)
            all_category_id = torch.cat((all_category_id, category_id), dim=0)

    all_features = all_features.numpy()
    all_vids = np.array(all_vids)
    all_category_id = all_category_id.numpy()

    current_step_class_idxs = np.array(range(args.class_num_per_step * step, args.class_num_per_step * (step + 1)))

    step_all_class_exemplar_features = []
    step_all_class_exemplar_vids = []

    for class_id in current_step_class_idxs:
        class_features = all_features[all_category_id == class_id]
        class_vids = all_vids[all_category_id == class_id]
        class_mean_feature = np.mean(class_features, axis=0)

        class_exemplar_features = []
        class_exemplar_vids = []
        now_class_mean = np.zeros((1, class_features.shape[-1]))
        for i in range(m_per_class):
            if i >= len(class_features):
                class_exemplar_vids.append(None)
                class_exemplar_features.append(np.zeros((class_features.shape[-1])))
                continue
            x = class_mean_feature - (now_class_mean + class_features) / (i + 1)
            x = np.linalg.norm(x, axis=1)
            index = np.argmin(x)
            now_class_mean += class_features[index]
            class_exemplar_features.append(class_features[index])
            class_exemplar_vids.append(class_vids[index])

        step_all_class_exemplar_features.append(class_exemplar_features)
        step_all_class_exemplar_vids.append(class_exemplar_vids)

    step_all_class_exemplar_features = np.array(step_all_class_exemplar_features)
    step_all_class_exemplar_vids = np.array(step_all_class_exemplar_vids)

    if step == 0:
        exemplar_class_features = step_all_class_exemplar_features
        new_exemplar_class_vids = step_all_class_exemplar_vids
    else:
        old_exemplar_class_vids_ = old_exemplar_class_vids[..., :m_per_class]

        exemplar_class_features = []
        for old_class_id in range(len(old_exemplar_class_vids_)):
            old_class_vids = old_exemplar_class_vids_[old_class_id]
            old_class_features = np.array([all_features[all_vids==vid][0] for vid in old_class_vids])
            exemplar_class_features.append(old_class_features)
        exemplar_class_features = np.array(exemplar_class_features)

        exemplar_class_features = np.concatenate((exemplar_class_features, step_all_class_exemplar_features), axis=0)
        new_exemplar_class_vids = np.concatenate((old_exemplar_class_vids_, step_all_class_exemplar_vids), axis=0)

    train_data_set._switch_train()

    if step != 0:
        train_data_set._conbine_exemplar()

    return exemplar_class_features, new_exemplar_class_vids

def train(args, step, train_data_set, val_data_set, exemplar_class_vids):
    train_loader = DataLoader(train_data_set, batch_size=args.train_batch_size, num_workers=args.num_workers,
                              pin_memory=True, drop_last=True, shuffle=True)
    val_loader = DataLoader(val_data_set, batch_size=args.infer_batch_size, num_workers=args.num_workers,
                            pin_memory=True, drop_last=False, shuffle=False)

    step_out_class_num = (step + 1) * args.class_num_per_step
    if step == 0:
        model = IncreAudioVisualNet(args, step_out_class_num, LSC=True)
    elif step == 1:
        model = torch.load('./save/{}/{}/{}/step_{}_best_{}_model.pkl'.format(args.dataset, args.modality, args.classify, step-1, args.modality))
        in_features = model.classifier.in_features
        out_features = model.classifier.out_features

        new_classifier = SplitLSCLinear(in_features, out_features, args.class_num_per_step)
        new_classifier.fc1.weight.data = model.classifier.weight.data
        model.classifier = new_classifier

        old_model = copy.deepcopy(model)
    else:
        model = torch.load('./save/{}/{}/{}/step_{}_best_{}_model.pkl'.format(args.dataset, args.modality, args.classify, step-1, args.modality))
        in_features = model.classifier.in_features

        out_features1 = model.classifier.fc1.out_features
        out_features2 = model.classifier.fc2.out_features
        K = model.classifier.fc2.K

        new_fc = SplitLSCLinear(in_features, out_features1 + out_features2, args.class_num_per_step)

        new_fc.fc1.weight.data[:out_features1*K] = model.classifier.fc1.weight.data
        new_fc.fc1.weight.data[out_features1*K:] = model.classifier.fc2.weight.data
        model.classifier = new_fc

        old_model = copy.deepcopy(model)

    if step > 0:
        model = imprint_weights(args, model, step, train_loader)

    if torch.cuda.device_count() > 1:
        model = nn.DataParallel(model)
        if step != 0:
            old_model = nn.DataParallel(old_model)

    model = model.to(device)
    if step != 0:
        old_model = old_model.to(device)
        old_model.eval()

    if step > 0:
        ignored_params = list(map(id, model.classifier.fc1.parameters()))
        base_params = filter(lambda p: id(p) not in ignored_params, model.parameters())
        params =[{'params': base_params, 'lr': args.lr, 'weight_decay': args.weight_decay}, \
                 {'params': model.classifier.fc1.parameters(), 'lr': 0, 'weight_decay': 0}]
    else:
        params = model.parameters()

    opt = torch.optim.Adam(params, lr=args.lr, weight_decay=args.weight_decay)

    train_loss_list = []
    val_acc_list = []
    best_val_res = 0.0
    best_exemplar_class_vids = None
    best_exemplar_class_mean = None

    for epoch in range(args.max_epoches):
        train_loss = 0.0
        num_steps = 0
        model.train()
        iterator = tqdm(train_loader)

        for samples in iterator:
            data, labels = samples
            labels = labels.to(device)
            if args.modality == 'visual':
                visual = data
                visual = visual.to(device)
                out, feature = model(visual=visual, AFC_train_out=True)
                if step != 0:
                    old_out, old_feature = old_model(visual=visual, AFC_train_out=True)
            elif args.modality == 'audio':
                audio = data
                audio = audio.to(device)
                out, feature = model(audio=audio, AFC_train_out=True)
                if step != 0:
                    old_out, old_feature = old_model(audio=audio, AFC_train_out=True)
            else:
                visual = data[0]
                audio = data[1]
                visual = visual.to(device)
                audio = audio.to(device)
                out, visual_pooled_feature, audio_feature, visual_feature = model(visual=visual, audio=audio, AFC_train_out=True)

                if step != 0:
                    old_out, old_visual_pooled_feature, old_audio_feature, old_visual_feature = old_model(visual=visual, audio=audio, AFC_train_out=True)

            if step != 0:
                prev_loss_cls = nca(old_out, labels, scale=old_model.classifier.factor)
                prev_loss_cls.backward(retain_graph=True)
                if args.modality == 'audio-visual':
                    old_audio_feature_grad = old_audio_feature.grad
                    old_visual_feature_grad = old_visual_feature.grad
                    I_audio = torch.norm(old_audio_feature_grad, p='fro', dim=-1)
                    I_visual = torch.norm(old_visual_feature_grad, p='fro', dim=-1)

                    old_visual_pooled_feature_grad = old_visual_pooled_feature.grad
                    I_visual_pooled = torch.norm(old_visual_pooled_feature_grad, p='fro', dim=-1)
                    all_I = I_audio + I_visual + I_visual_pooled
                    norm_I_visual_pooled = I_visual_pooled / all_I
                    loss_I_visual_pooled = torch.mean(norm_I_visual_pooled * torch.norm(F.normalize(visual_pooled_feature) - F.normalize(old_visual_pooled_feature.detach()), p='fro', dim=-1))

                    norm_I_audio = I_audio / all_I
                    norm_I_visual = I_visual / all_I
                    loss_I_audio = torch.mean(norm_I_audio * torch.norm(F.normalize(audio_feature) - F.normalize(old_audio_feature.detach()), p='fro', dim=-1))
                    loss_I_visual = torch.mean(norm_I_visual * torch.norm(F.normalize(visual_feature) - F.normalize(old_visual_feature.detach()), p='fro', dim=-1))

                    loss_disc = loss_I_visual_pooled + loss_I_audio + loss_I_visual
                else:
                    loss_disc = torch.mean(torch.norm(F.normalize(feature) - F.normalize(old_feature), p='fro', dim=-1))
                loss_cls = nca(out, labels, scale=model.classifier.factor)
                lam_t = (step + 1) ** 0.5
                lam_disc = args.lam_disc
                loss = loss_cls + lam_disc * lam_t * loss_disc
            else:
                loss = nca(out, labels, scale=model.classifier.factor)
            model.zero_grad()
            loss.backward()
            opt.step()
            train_loss += loss.item()
            num_steps += 1
        train_loss /= num_steps
        train_loss_list.append(train_loss)
        print('Epoch:{} train_loss:{:.5f}'.format(epoch, train_loss), flush=True)

        if args.classify == 'NME':
            m = args.memory_size // (args.class_num_per_step * (step + 1))
            exemplar_class_features, new_exemplar_class_vids = gen_step_exemplar_set(args, model, step, train_data_set, m, exemplar_class_vids)
            exemplar_class_mean = np.mean(exemplar_class_features, axis=1)

        all_val_out = torch.Tensor([])
        all_val_labels = torch.Tensor([])
        model.eval()
        with torch.no_grad():
            for val_data, val_labels in tqdm(val_loader):
                if args.modality == 'visual':
                    val_visual = val_data
                    val_visual = val_visual.to(device)
                    if args.classify == 'NME':
                        val_out = model(visual=val_visual, out_logits=False, out_features=True, out_features_norm=True).detach().cpu()
                    else:
                        val_out = model(visual=val_visual).detach().cpu()
                elif args.modality == 'audio':
                    val_audio = val_data
                    val_audio = val_audio.to(device)
                    if args.classify == 'NME':
                        val_out = model(audio=val_audio, out_logits=False, out_features=True, out_features_norm=True).detach().cpu()
                    else:
                        val_out = model(audio=val_audio).detach().cpu()
                else:
                    val_visual = val_data[0]
                    val_audio = val_data[1]
                    val_visual = val_visual.to(device)
                    val_audio = val_audio.to(device)
                    if args.classify == 'NME':
                        val_out = model(visual=val_visual, audio=val_audio, out_logits=False, out_features=True, out_features_norm=True).detach().cpu()
                    else:
                        val_out = model(visual=val_visual, audio=val_audio).detach().cpu()
                        val_out = F.softmax(val_out, dim=-1).detach().cpu()
                all_val_out = torch.cat((all_val_out, val_out), dim=0)
                all_val_labels = torch.cat((all_val_labels, val_labels), dim=0)
        if args.classify == 'NME':
            all_val_pre = NME_classify(all_val_out, exemplar_class_mean)
            all_val_labels = all_val_labels.numpy()
            val_res = (all_val_pre == all_val_labels).sum() / len(all_val_labels)
        else:
            val_res = top_1_acc(all_val_out, all_val_labels)
        val_acc_list.append(val_res)
        print('Epoch:{} val_res:{:.6f} '.format(epoch, val_res), flush=True)

        if val_res > best_val_res:
            best_val_res = val_res

            if args.classify == 'NME':
                best_exemplar_class_vids = new_exemplar_class_vids
                best_exemplar_class_mean = exemplar_class_mean

            print('Saving best model at Epoch {}'.format(epoch), flush=True)
            if torch.cuda.device_count() > 1:
                torch.save(model.module, './save/{}/{}/{}/step_{}_best_{}_model.pkl'.format(args.dataset, args.modality, args.classify, step, args.modality))
            else:
                torch.save(model, './save/{}/{}/{}/step_{}_best_{}_model.pkl'.format(args.dataset, args.modality, args.classify, step, args.modality))

        plt.figure()
        plt.plot(range(len(train_loss_list)), train_loss_list, label='train_loss')
        plt.legend()
        plt.savefig('./save/fig/{}/{}/{}/train_loss_step_{}.png'.format(args.dataset, args.modality, args.classify, step))
        plt.close()

        plt.figure()
        plt.plot(range(len(val_acc_list)), val_acc_list, label='val_acc')
        plt.legend()
        plt.savefig('./save/fig/{}/{}/{}/val_acc_step_{}.png'.format(args.dataset, args.modality, args.classify, step))
        plt.close()

        if args.lr_decay:
            adjust_learning_rate(args, opt, epoch)

    if args.classify == 'LSC':
        step_best_model = torch.load('./save/{}/{}/{}/step_{}_best_{}_model.pkl'.format(args.dataset, args.modality, args.classify, step, args.modality))
        m = args.memory_size // (args.class_num_per_step * (step + 1))
        best_exemplar_class_features, best_exemplar_class_vids = gen_step_exemplar_set(args, step_best_model, step, train_data_set, m, exemplar_class_vids)
        best_exemplar_class_mean = np.mean(best_exemplar_class_features, axis=1)
    np.save('./save/{}/{}/{}/step_{}_best_{}_exemplar_class_mean.npy'.format(args.dataset, args.modality, args.classify, step, args.modality), best_exemplar_class_mean)

    return best_exemplar_class_vids, best_exemplar_class_mean

def NME_classify(features, exemplar_class_mean):
    norm_features = F.normalize(features).numpy()
    all_res = []
    for f in norm_features:
        dis = f - exemplar_class_mean
        dis = np.linalg.norm(dis, ord=2, axis=1)
        pre = np.argmin(dis)
        all_res.append(pre)
    return np.array(all_res)

def detailed_test(args, step, test_data_set, task_best_acc_list):
    print("=====================================")
    print("Start testing...")
    print("=====================================")
    # model = IncreAudioVisualNet(args=args)

    model = torch.load('./save/{}/{}/{}/step_{}_best_{}_model.pkl'.format(args.dataset, args.modality, args.classify, step, args.modality))

    exemplar_class_mean = np.load('./save/{}/{}/{}/step_{}_best_{}_exemplar_class_mean.npy'.format(args.dataset, args.modality, args.classify, step, args.modality))

    model.to(device)

    test_loader = DataLoader(test_data_set, batch_size=args.infer_batch_size, num_workers=args.num_workers,
                             pin_memory=True, drop_last=False, shuffle=False)

    all_test_out = torch.Tensor([])
    all_test_labels = torch.Tensor([])
    model.eval()
    with torch.no_grad():
        for test_data, test_labels in tqdm(test_loader):
            # test_labels = test_labels.to(device)
            if args.modality == 'visual':
                test_visual = test_data
                test_visual = test_visual.to(device)
                if args.classify == 'NME':
                    test_out = model(visual=test_visual, out_logits=False, out_features=True, out_features_norm=True).detach().cpu()
                else:
                    test_out = model(visual=test_visual).detach().cpu()
            elif args.modality == 'audio':
                test_audio = test_data
                test_audio = test_audio.to(device)
                if args.classify == 'NME':
                    test_out = model(audio=test_audio, out_logits=False, out_features=True, out_features_norm=True).detach().cpu()
                else:
                    test_out = model(audio=test_audio).detach().cpu()
            else:
                test_visual = test_data[0]
                test_audio = test_data[1]
                test_visual = test_visual.to(device)
                test_audio = test_audio.to(device)
                if args.classify == 'NME':
                    test_out = model(visual=test_visual, audio=test_audio, out_logits=False, out_features=True, out_features_norm=True).detach().cpu()
                else:
                    test_out = model(visual=test_visual, audio=test_audio).detach().cpu()
                    test_out = F.softmax(test_out, dim=-1).detach().cpu()
            # test_out = F.softmax(test_out, dim=-1).detach().cpu()
            all_test_out = torch.cat((all_test_out, test_out), dim=0)
            all_test_labels = torch.cat((all_test_labels, test_labels), dim=0)

    if args.classify == 'NME':
        all_test_pre = NME_classify(all_test_out, exemplar_class_mean)
        # all_test_labels = all_test_labels.numpy()
        test_res = (all_test_pre == all_test_labels.numpy()).sum() / len(all_test_labels)

    else:
        test_res = top_1_acc(all_test_out, all_test_labels)

    print("Incremental step {} Testing res: {:.6f}".format(step, test_res))

    old_task_acc_list = []
    for i in range(step+1):
        step_class_list = range(i*args.class_num_per_step, (i+1)*args.class_num_per_step)
        step_class_idxs = []
        for c in step_class_list:
            idxs = np.where(all_test_labels.numpy() == c)[0].tolist()
            step_class_idxs += idxs
        step_class_idxs = np.array(step_class_idxs)
        i_labels = torch.Tensor(all_test_labels.numpy()[step_class_idxs])
        if args.classify == 'NME':
            i_pred = all_test_pre[step_class_idxs]
            i_acc = (i_pred == i_labels.numpy()).sum() / len(i_labels)
        else:
            i_outs = torch.Tensor(all_test_out.numpy()[step_class_idxs])
            i_acc = top_1_acc(i_outs, i_labels)

        if i == step:
            curren_step_acc = i_acc
        else:
            old_task_acc_list.append(i_acc)
    if step > 0:
        forgetting = np.mean(np.array(task_best_acc_list) - np.array(old_task_acc_list))
        print('forgetting: {:.6f}'.format(forgetting))
        for i in range(len(task_best_acc_list)):
            task_best_acc_list[i] = max(task_best_acc_list[i], old_task_acc_list[i])
    else:
        forgetting = None
    task_best_acc_list.append(curren_step_acc)

    return forgetting

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset', type=str, default='AVE', choices=['AVE', 'ksounds', 'VGGSound_100'])
    parser.add_argument('--modality', type=str, default='visual', choices=['visual', 'audio', 'audio-visual'])
    parser.add_argument('--train_batch_size', type=int, default=128)
    parser.add_argument('--infer_batch_size', type=int, default=32)

    parser.add_argument('--gen_exem_batch_size', type=int, default=64)

    parser.add_argument('--num_workers', type=int, default=0)
    parser.add_argument('--max_epoches', type=int, default=500)
    parser.add_argument('--num_classes', type=int, default=28)
    parser.add_argument('--lr', type=float, default=1e-3)
    parser.add_argument('--weight_decay', type=float, default=1e-4)
    parser.add_argument('--lr_decay', type=boolean_string, default=False)
    parser.add_argument("--milestones", type=int, default=[500], nargs='+', help="")
    parser.add_argument('--seed', type=int, default=42)

    parser.add_argument('--class_num_per_step', type=int, default=7)

    parser.add_argument('--memory_size', type=int, default=340)

    parser.add_argument('--classify', type=str, default='NME', choices=['LSC', 'NME'])
    parser.add_argument('--lam_disc', type=float, default=4.0)

    args = parser.parse_args()
    print(args)

    total_incremental_steps = args.num_classes // args.class_num_per_step

    setup_seed(args.seed)

    print('Training start time: {}'.format(datetime.now()))

    train_set = IceLoader(args=args, mode='train', modality=args.modality)
    val_set = IceLoader(args=args, mode='val', modality=args.modality)
    test_set = IceLoader(args=args, mode='test', modality=args.modality)

    ckpts_root = './save/{}/{}/{}/'.format(args.dataset, args.modality, args.classify)
    figs_root = './save/fig/{}/{}/{}/'.format(args.dataset, args.modality, args.classify)

    if not os.path.exists(ckpts_root):
        os.makedirs(ckpts_root)
    if not os.path.exists(figs_root):
        os.makedirs(figs_root)

    task_best_acc_list = []
    step_forgetting_list = []

    exemplar_class_vids = None
    for step in range(total_incremental_steps):
        train_set.set_incremental_step(step)
        val_set.set_incremental_step(step)
        test_set.set_incremental_step(step)

        train_set._update_exemplars(exemplar_class_vids)

        print('Incremental step: {}'.format(step))

        exemplar_class_vids, exemplar_class_mean = train(args, step, train_set, val_set, exemplar_class_vids)

        step_forgetting = detailed_test(args, step, test_set, task_best_acc_list)
        if step_forgetting is not None:
            step_forgetting_list.append(step_forgetting)
    Mean_forgetting = np.mean(step_forgetting_list)
    print('Average Forgetting: {:.6f}'.format(Mean_forgetting))

    if args.dataset != 'AVE' and args.modality != 'audio':
        train_set.close_visual_features_h5()
        val_set.close_visual_features_h5()
        test_set.close_visual_features_h5()


from contextlib import contextmanager
from pathlib import Path
from unittest.mock import patch
import tempfile

import pytest

from cohortextractor.cohortextractor import list_study_definitions

@contextmanager
def dummy_repo_folder(with_definition=True):
    with tempfile.TemporaryDirectory() as tmpdir:
        folder = Path(tmpdir)
        (folder / "analysis").mkdir(exist_ok=True)
        if with_definition:
            (folder / "analysis" / "study_definition_test.py").touch()
        yield tmpdir

@patch("cohortextractor.cohortextractor.relative_dir")
def test_no_study_definition_raises(relative_dir):
    with dummy_repo_folder(with_definition=False) as dummy_repo:
        relative_dir.return_value = dummy_repo
        with pytest.raises(RuntimeError):
            list_study_definitions()

@patch("cohortextractor.cohortextractor.relative_dir")
def test_list_study_definition(relative_dir):
    with dummy_repo_folder() as dummy_repo:
        relative_dir.return_value = dummy_repo
        definitions = list_study_definitions()
        assert definitions == [("study_definition_test", "_test")]


from torchvision.datasets.utils import list_dir
from torchvision.datasets.folder import make_dataset
from torchvision.datasets.video_utils import VideoClips
from torchvision.datasets import VisionDataset
import os
import torch

class Kinetics400(VisionDataset):
    """
    `Kinetics-400 <https://deepmind.com/research/open-source/open-source-datasets/kinetics/>`_
    dataset.

    Kinetics-400 is an action recognition video dataset.
    This dataset consider every video as a collection of video clips of fixed size, specified
    by ``frames_per_clip``, where the step in frames between each clip is given by
    ``step_between_clips``.

    To give an example, for 2 videos with 10 and 15 frames respectively, if ``frames_per_clip=5``
    and ``step_between_clips=5``, the dataset size will be (2 + 3) = 5, where the first two
    elements will come from video 1, and the next three elements from video 2.
    Note that we drop clips which do not have exactly ``frames_per_clip`` elements, so not all
    frames in a video might be present.

    Internally, it uses a VideoClips object to handle clip creation.

    Args:
        root (string): Root directory of the Kinetics-400 Dataset.
        frames_per_clip (int): number of frames in a clip
        step_between_clips (int): number of frames between each clip
        transform (callable, optional): A function/transform that  takes in a TxHxWxC video
            and returns a transformed version.

    Returns:
        video (Tensor[T, H, W, C]): the `T` video frames
        audio(Tensor[K, L]): the audio frames, where `K` is the number of channels
            and `L` is the number of points
        label (int): class of the video clip
    """

    def __init__(self, root, frames_per_clip, step_between_clips=1, frame_rate=None,
                 extensions=('avi',), transform=None, num_workers=1, _video_width=0, _video_height=0,
                 _video_min_dimension=0, _audio_samples=0):
        super(Kinetics400, self).__init__(root)

        classes = list(sorted(list_dir(root)))
        class_to_idx = {classes[i]: i for i in range(len(classes))}
        self.samples = make_dataset(self.root, class_to_idx, extensions, is_valid_file=None)
        self.classes = classes
        video_list = [x[0] for x in self.samples]
        split = root.split('/')[-1].strip('/')
        metadata_filepath = os.path.join(root, 'kinetics_metadata_{}.pt'.format(split))

        if os.path.exists(metadata_filepath):
            metadata = torch.load(metadata_filepath)

        else:
            metadata = None

        self.video_clips = VideoClips(
            video_list,
            frames_per_clip,
            step_between_clips,
            frame_rate,
            metadata,
            num_workers=num_workers,
            _video_width=_video_width,
            _video_height=_video_height,
            _video_min_dimension=_video_min_dimension,
            _audio_samples=_audio_samples,
        )
        self.transform = transform
        if not os.path.exists(metadata_filepath):
            torch.save(self.video_clips.metadata, metadata_filepath)
    @property
    def metadata(self):
        return self.video_clips.metadata

    def __len__(self):
        return self.video_clips.num_clips()

    def __getitem__(self, idx):
        # video_q, audio_q, info_q, video_idx_q = self.video_clips.get_clip(idx[0])
        # video_k, audio_k, info_k, video_idx_k = self.video_clips.get_clip(idx[1])
        video, audio, info, video_idx = self.video_clips.get_clip(idx)

        video_q = self.transform['video'](video)
        video_k = self.transform['video'](video)
        audio_q = self.transform['audio'](audio)
        audio_k = self.transform['audio'](audio)

        return (video_q, video_k), (audio_q, audio_k)


import os
import sys
from dataclasses import dataclass
from itertools import product

import kernel_explorer as ke
import numpy as np
import pytest
from utils import get_gemm_basic_sizes, get_gemm_bert_sizes, get_gemm_bound, matmul, transab_to_suffix

max_batch_size = int(os.environ.get("KERNEL_EXPLORER_BATCHED_GEMM_MAX_BATCH_SIZE", 64))

def dtype_to_suffix(dtype):
    return {
        "float32": "float",
        "float16": "half",
    }[dtype]

def _test_strided_batched_gemm(
    func, dtype: str, transa: bool, transb: bool, m: int, n: int, k: int, batch: int, alpha=1.0, beta=0.0
):
    assert dtype in ["float32", "float16"]

    a_shape = (k, m) if transa else (m, k)
    b_shape = (n, k) if transb else (k, n)

    np.random.seed(0)
    a = (np.random.rand(batch, *a_shape) + 0.5).astype(dtype).astype("float64")
    b = (np.random.rand(batch, *b_shape) + 0.5).astype(dtype).astype("float64")
    tmp_a = a.swapaxes(1, 2) if transa else a
    tmp_b = b.swapaxes(1, 2) if transb else b
    intermediate_c = matmul(tmp_a, tmp_b)
    if alpha == 1.0 and beta == 0.0:  # fast path
        ref_c = intermediate_c
    else:
        ref_c = alpha * intermediate_c + beta * np.ones_like(intermediate_c)

    bounds = [get_gemm_bound(dtype, a[i], b[i], ref_c[i], transa, transb, a_b_positive=True) for i in range(batch)]

    a = a.astype(dtype)
    b = b.astype(dtype)

    my_c = np.ones((batch, m, n), dtype=dtype)
    dev_a = ke.DeviceArray(a)
    dev_b = ke.DeviceArray(b)
    dev_c = ke.DeviceArray(my_c)

    opa = ke.blas_op.T if transa else ke.blas_op.N
    opb = ke.blas_op.T if transb else ke.blas_op.N
    lda = a_shape[1]
    ldb = b_shape[1]
    ldc = n
    stride_a = m * k
    stride_b = k * n
    stride_c = m * n
    my_gemm = func(
        opa, opb, m, n, k, alpha, dev_a, lda, stride_a, dev_b, ldb, stride_b, beta, dev_c, ldc, stride_c, batch
    )

    failures = {}
    print(
        f"dtype={dtype} {transab_to_suffix((transa, transb))} m={m:<5} n={n:<5} k={k:<5} batch={batch:<3} max bound: {max(bounds)}"
    )

    for impl in my_gemm.ListOps():
        if not my_gemm.SelectOp(impl):
            continue

        # Restore C Array
        my_c.fill(1.0)
        dev_c.UpdateDeviceArray()
        my_gemm.Run()
        dev_c.UpdateHostNumpyArray()

        for i in range(batch):
            try:
                np.testing.assert_allclose(my_c[i], ref_c[i], rtol=bounds[i])
            except Exception as err:
                header = "*" * 30 + impl + "*" * 30
                print(header, bounds[i])
                print(err)
                print("*" * len(header))
                failures[impl] = str(err)

    if failures:
        raise Exception(failures)

dtypes = ["float32", "float16"]
all_transabs = list(product([True, False], repeat=2))

@pytest.mark.parametrize("batch", [1, max_batch_size])
@pytest.mark.parametrize("m, n, k", get_gemm_basic_sizes(full=False) + get_gemm_bert_sizes(full=False))
@pytest.mark.parametrize("transa, transb", all_transabs)
@pytest.mark.parametrize("dtype", dtypes)
def test_rocblas_gemm_all_cases(dtype, transa, transb, m, n, k, batch):
    wrapper_name = "RocBlasStridedBatchedGemm_" + dtype_to_suffix(dtype)
    _test_strided_batched_gemm(getattr(ke, wrapper_name), dtype, transa, transb, m, n, k, batch)

@pytest.mark.skipif(not ke.is_composable_kernel_available(), reason="ck is not enabled")
@pytest.mark.parametrize("batch", [1, max_batch_size])
@pytest.mark.parametrize("m, n, k", get_gemm_basic_sizes(full=False) + get_gemm_bert_sizes(full=False))
@pytest.mark.parametrize("transa, transb", all_transabs)
@pytest.mark.parametrize("dtype", dtypes)
def test_ck_gemm_all_cases(dtype, transa, transb, m, n, k, batch):
    wrapper_name = f"CKStridedBatchedGemm_{dtype_to_suffix(dtype)}_{transab_to_suffix((transa, transb))}"
    _test_strided_batched_gemm(getattr(ke, wrapper_name), dtype, transa, transb, m, n, k, batch)

# Tunable is basically wrapped around of rocblas and ck gemm, so no need for full tests
@pytest.mark.parametrize("batch", [1, max_batch_size])
@pytest.mark.parametrize("m, n, k", get_gemm_bert_sizes(full=False))
@pytest.mark.parametrize("transa, transb", all_transabs)
@pytest.mark.parametrize("dtype", dtypes)
def test_gemm_tunable_bert_cases(dtype, transa, transb, m, n, k, batch):
    wrapper_name = f"StridedBatchedGemmTunable_{dtype_to_suffix(dtype)}_{transab_to_suffix((transa, transb))}"
    _test_strided_batched_gemm(getattr(ke, wrapper_name), dtype, transa, transb, m, n, k, batch)

@pytest.mark.parametrize("alpha, beta", [(0.5, 0.5)])
@pytest.mark.parametrize("transa, transb", all_transabs)
@pytest.mark.parametrize("dtype", dtypes)
def test_rocblas_gemm_alpha_beta(dtype, transa, transb, alpha, beta):
    wrapper_name = "RocBlasStridedBatchedGemm_" + dtype_to_suffix(dtype)
    _test_strided_batched_gemm(
        getattr(ke, wrapper_name), dtype, transa, transb, 128, 256, 768, 16, alpha=alpha, beta=beta
    )

@pytest.mark.skipif(not ke.is_composable_kernel_available(), reason="ck is not enabled")
@pytest.mark.parametrize("alpha, beta", [(0.5, 0.5)])
@pytest.mark.parametrize("transa, transb", all_transabs)
@pytest.mark.parametrize("dtype", dtypes)
def test_ck_gemm_alpha_beta(dtype, transa, transb, alpha, beta):
    wrapper_name = f"CKStridedBatchedGemm_{dtype_to_suffix(dtype)}_{transab_to_suffix((transa, transb))}"
    _test_strided_batched_gemm(
        getattr(ke, wrapper_name), dtype, transa, transb, 256, 128, 384, 8, alpha=alpha, beta=beta
    )

@pytest.mark.parametrize("alpha, beta", [(0.5, 0.5)])
@pytest.mark.parametrize("transa, transb", all_transabs)
@pytest.mark.parametrize("dtype", dtypes)
def test_gemm_tunable_alpha_beta(dtype, transa, transb, alpha, beta):
    wrapper_name = f"StridedBatchedGemmTunable_{dtype_to_suffix(dtype)}_{transab_to_suffix((transa, transb))}"
    _test_strided_batched_gemm(
        getattr(ke, wrapper_name), dtype, transa, transb, 128, 512, 384, 4, alpha=alpha, beta=beta
    )

@dataclass
class StridedBatchedGemmMetric(ke.ComputeMetric):
    transa: bool
    transb: bool
    m: int
    n: int
    k: int
    batch: int

    def report(self):
        common = (
            f"{self.dtype} {transab_to_suffix((self.transa, self.transb))} "
            f"m={self.m:<4} n={self.n:<4} k={self.k:<4} batch={self.batch:<3} {self.name}"
        )
        if self.duration <= 0:
            return "not supported          " + common

        return f"{self.duration:>6.2f} us {self.tflops:>5.2f} tflops " + common

def profile_gemm_func(f, dtype: str, transa: bool, transb: bool, m: int, n: int, k: int, batch: int):
    a_shape = (k, m) if transa else (m, k)
    b_shape = (n, k) if transb else (k, n)

    np.random.seed(0)
    a = (np.random.rand(batch, *a_shape) * 2 - 1).astype(dtype)
    b = (np.random.rand(batch, *b_shape) * 2 - 1).astype(dtype)
    my_c = np.zeros((batch, m, n), dtype=dtype)

    dev_a = ke.DeviceArray(a)
    dev_b = ke.DeviceArray(b)
    dev_c = ke.DeviceArray(my_c)

    opa = ke.blas_op.T if transa else ke.blas_op.N
    opb = ke.blas_op.T if transb else ke.blas_op.N
    lda = a_shape[1]
    ldb = b_shape[1]
    ldc = n
    stride_a = m * k
    stride_b = k * n
    stride_c = m * n
    alpha = 1.0
    beta = 0.0
    my_gemm = f(opa, opb, m, n, k, alpha, dev_a, lda, stride_a, dev_b, ldb, stride_b, beta, dev_c, ldc, stride_c, batch)
    for impl in my_gemm.ListOps():
        duration_ms = -1
        if my_gemm.SelectOp(impl):
            duration_ms = my_gemm.Profile()
        FLOPs = batch * m * k * n * 2  # noqa: N806
        ke.report(StridedBatchedGemmMetric(impl, dtype, duration_ms, FLOPs, transa, transb, m, n, k, batch))

def profile_with_args(dtype, transa, transb, m, n, k, batch, sort):
    dtype_suffix = "_" + dtype_to_suffix(dtype)
    transab_suffix = "_" + transab_to_suffix((transa, transb))
    fn_rocblas = getattr(ke, "RocBlasStridedBatchedGemm" + dtype_suffix)
    fn_ck = getattr(ke, "CKStridedBatchedGemm" + dtype_suffix + transab_suffix)
    fn_tunable = getattr(ke, "StridedBatchedGemmTunable" + dtype_suffix + transab_suffix)
    if ke.is_hipblaslt_available():
        fn_hipblaslt = getattr(ke, "StridedBatchedGemmHipBlasLt" + dtype_suffix + transab_suffix)
    with ke.benchmark(sort):
        profile_gemm_func(fn_rocblas, dtype, transa, transb, m, n, k, batch)
        profile_gemm_func(fn_ck, dtype, transa, transb, m, n, k, batch)
        profile_gemm_func(fn_tunable, dtype, transa, transb, m, n, k, batch)
        if ke.is_hipblaslt_available():
            profile_gemm_func(fn_hipblaslt, dtype, transa, transb, m, n, k, batch)
    print()

def profile():
    for dtype in dtypes:
        for m, n, k in get_gemm_bert_sizes(full=False):
            for batch in [1, 32, 64]:
                profile_with_args(dtype, False, False, m, n, k, batch, True)

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    group = parser.add_argument_group("profile with args")
    group.add_argument("dtype", choices=dtypes)
    group.add_argument("transa", choices="NT")
    group.add_argument("transb", choices="NT")
    group.add_argument("m", type=int)
    group.add_argument("n", type=int)
    group.add_argument("k", type=int)
    group.add_argument("batch", type=int)
    group.add_argument("--sort", action="store_true")

    if len(sys.argv) == 1:
        profile()
    else:
        args = parser.parse_args()
        profile_with_args(
            args.dtype, args.transa == "T", args.transb == "T", args.m, args.n, args.k, args.batch, args.sort
        )


#!/usr/bin/python3
from subprocess import PIPE, CalledProcessError, Popen, TimeoutExpired
from subprocess import run as sub_run
import logging
import os
import shutil
import shlex
import signal
from concurrent.futures import ThreadPoolExecutor, wait

try:
    import click
except ImportError as e:
    if 'No module named \'click\'' in e.msg:
        sub_run('sudo pip3 install click'.split(), check=True)
        import click
    else:
        raise ImportError('Something other than a missing module is wrong: {}'.format(e.msg))

logging.basicConfig(format='[%(asctime)s]\t%(levelname)s: %(message)s')

def install_gazebo(log):
    '''Installs the latest version of Gazebo. Note: Only Ubuntu is officially supported.'''
    log.info('Installing Gazebo')
    # These commands are taken from the official Ubuntu install sequence
    # at http://gazebosim.org/tutorials?tut=install_ubuntu&cat=install
    install_commands = [
            'sudo sh -c \'echo "deb http://packages.osrfoundation.org/gazebo/ubuntu-stable\
 `lsb_release -cs` main" > /etc/apt/sources.list.d/gazebo-stable.list\'',
            'wget http://packages.osrfoundation.org/gazebo.key -O - | sudo apt-key add -',
            'sudo apt-get update',
            'sudo apt-get install gazebo7 libgazebo7-dev'
            ]
    for command in install_commands:
        # Yes, running in a shell has security concerns. It should be OK for this instance,
        # though - no non-hardcoded commands running, and it's easier than setting up the
        # necessary pipes and such
        sub_run(command, shell=True, check=True)
    log.info('Gazebo has been successfully installed!')

def download_assignment(log, assignment_name):
    '''Downloads and extracts the provided files for the given assignment'''
    url = 'https://rpal.cs.cornell.edu/teaching/foundations/assignments/' + assignment_name
    base_cmd = 'wget ' + url
    log.info('Downloading and extracting archive for {}'.format(assignment_name))
    def unzip(f):
        if not f.exception():
            sub_run('tar xzvf {}.tar.gz -C {}'.format(assignment_name, assignment_name).split())
        else:
            raise f.exception()

    with ThreadPoolExecutor(max_workers=1) as e:
        download_future = e.submit(lambda: sub_run(('wget ' + url + '.tar.gz').split(), check=True))
        download_future.add_done_callback(unzip)
    log.info('Downloading version signature for {}'.format(assignment_name))
    sub_run(('wget -O .sigs/{} ' + url + '.sig').format(assignment_name).split(), check=True)
    wait([download_future])
    log.info('Done downloading and extracting materials for {}'.format(assignment_name))

def remove_assignment(log, assignment_name):
    log.info('Deleting old assignment archive and directory')
    shutil.rmtree(assignment_name)
    os.remove('{}.tar.gz'.format(assignment_name))

def fetch_assignment(log, assignment_name):
    '''Fetch the most up-to-date version of the assignment specified'''
    url_format = 'https://rpal.cs.cornell.edu/teaching/foundations/assignments/{}.sig'
    # Check if assignment has already been fetched and is up to date
    if os.path.isdir(assignment_name):
        log.info('Found directory for {}'.format(assignment_name))
        log.info('Checking version signature for {}'.format(assignment_name))
        current_sig = sub_run(
                'wget -qO- {}'.format(url_format.format(assignment_name)).split(),
                stdout=PIPE,
                universal_newlines=True,
                check=True
                ).stdout
        with open('.sigs/{}'.format(assignment_name, assignment_name), 'r') as sigfile:
            old_sig = sigfile.readline()
        if old_sig != current_sig:
            log.info('Version of {} appears to be outdated. Downloading current version.'
                     .format(assignment_name))
            remove_assignment(log, assignment_name)
            os.mkdir(assignment_name)
            download_assignment(log, assignment_name)
        else:
            log.info('Version of {} is up to date; proceeding'.format(assignment_name))
    else:
        log.info('Found no existing files for {}. Downloading current version.'
                 .format(assignment_name))
        os.mkdir(assignment_name)
        if not os.path.isdir('.sigs'):
            os.mkdir('.sigs')
        download_assignment(log, assignment_name)

def build_assignment(log, assignment_name):
    # TODO: This seems unnecessary
    def ensure_dir(dirname):
        if not os.path.isdir(dirname):
            os.mkdir(dirname)
    curr_dir = os.getcwd()
    log.info('Moving into assignment directory')
    os.chdir(assignment_name)
    log.info('Ensuring existence of build directory')
    ensure_dir('build')
    os.chdir('build')
    if not os.path.isfile('Makefile'):
        log.info('No Makefile; running CMake')
        sub_run('cmake ..'.split(), check=True)
    log.info('Running make')
    sub_run('make', check=True)
    os.chdir(curr_dir)
    log.info('Building done. {} is ready to run'.format(assignment_name))

def run_gazebo(log, assignment_name, problem):
    log.info('Running test driver for {}'.format(assignment_name))
    driver_proc = Popen(
            '{}/build/{}'.format(assignment_name, problem),
            stdin=None,
            stdout=None,
            stderr=None
            )
    env  = os.environ.copy()
    env['GAZEBO_PLUGIN_PATH'] = '{}/{}/build/'.format(os.getcwd(), assignment_name)
    env['GAZEBO_MODEL_PATH'] = '{}/{}/models/'.format(os.getcwd(), assignment_name)
    log.info('Running gazebo with the world for {}'.format(assignment_name))
    gazebo_proc = Popen(
            shlex.split('gazebo --verbose {}/{}.world'.format(assignment_name, assignment_name)),
            stdin=None,
            stdout=None,
            stderr=None,
            env=env
        )
    # Run for 60 seconds max
    try:
        driver_proc.wait(timeout=60)
    except TimeoutExpired:
        os.killpg(os.getpgid(driver_proc.pid), signal.SIGTERM)
    log.info('Wrapping up execution...')
    os.killpg(os.getpgid(gazebo_proc.pid), signal.SIGTERM)
    log.info('Done with execution')

@click.group()
def cli():
    pass

# TODO: Add no-check option
@cli.command()
@click.argument('assignment')
@click.argument('problem')
def run(assignment, problem):
    '''Build (if necessary) and run the given assignment problem.'''
    log = logging.getLogger('simulator-run')
    log.setLevel(logging.INFO)
    fetch_assignment(log, assignment)
    build_assignment(log, assignment)
    run_gazebo(log, assignment, problem)

@cli.command()
def install():
    '''Install Gazebo'''
    log = logging.getLogger('simulator-install')
    log.setLevel(logging.INFO)
    log.info('Checking to see if we already have gazebo')
    if shutil.which('gazebo'):
        log.info('Gazebo is already installed! Exiting.')
        return
    else:
        log.info("Gazebo doesn't seem to be installed.")
        install_gazebo(log)

# TODO: Add no-check option
@cli.command()
@click.argument('assignment')
def build(assignment):
    '''Build the given assignment'''
    log = logging.getLogger('simulator-build')
    log.setLevel(logging.INFO)
    fetch_assignment(log, assignment)
    build_assignment(log, assignment)

@cli.command()
@click.argument('assignment')
def get(assignment):
    '''Download the given assignment'''
    log = logging.getLogger('simulator-get')
    log.setLevel(logging.INFO)
    fetch_assignment(log, assignment)

if __name__ == '__main__':
    cli()


import multiprocessing
import pandas as pd
import os
import time
from selenium import webdriver
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC

def get_details(driver, selector_value, timeout=3, locator_type=By.XPATH):
    try:
        return WebDriverWait(driver=driver, timeout=timeout).until(
            EC.presence_of_element_located((locator_type, selector_value))
        ).text
    except:
        return None

def get_job_details(url, location="Texas", keyword="Data Scientist"):
    df = pd.DataFrame(columns=[ "Company name", "Rating", "Location", "Job title",
                                "Estimated salary", "Job description",
                                "Number of employees", "Founded", "Sector",
                                "Industry", "Revenue", "Keyword"])

    selector_company_detail = '//*[@id="JDCol"]/div/article/div/div[1]/div/div/div[1]/div[3]/div[1]/div[1]/div'
    selector_rating = '//*[@id="JDCol"]/div/article/div/div[1]/div/div/div[1]/div[3]/div[1]/div[1]/div/span'
    selector_job_title = '//*[@id="JDCol"]/div/article/div/div[1]/div/div/div[1]/div[3]/div[1]/div[2]'
    selector_job_location = '//*[@id="JDCol"]/div/article/div/div[1]/div/div/div[1]/div[3]/div[1]/div[3]'
    selector_estimated_salary = '//*[@id="JDCol"]/div/article/div/div[1]/div/div/div[1]/div[3]/div[1]/div[4]/span'
    selector_job_description = 'JobDescriptionContainer'
    selector_company_size = '//*[@id="EmpBasicInfo"]/div[1]/div/div[1]/span[2]'
    selector_founded = '//*[@id="EmpBasicInfo"]/div[1]/div/div[2]/span[2]'
    selector_company_type = '//*[@id="EmpBasicInfo"]/div[1]/div/div[3]/span[2]'
    selector_industry = '//*[@id="EmpBasicInfo"]/div[1]/div/div[4]/span[2]'
    selector_sector = '//*[@id="EmpBasicInfo"]/div[1]/div/div[5]/span[2]'
    selector_revenue = '//*[@id="EmpBasicInfo"]/div[1]/div/div[6]/span[2]'
    no_found_text = "Your search for data scientist in United States does not match any open jobs. Don't worry, we can still help. Below, please find related information to help you with your job search.".lower()

    driver = webdriver.Chrome()
    driver.maximize_window()
    driver.get(url)

    timeout = 2
    time_out_2 = 15

    for i in range(2):
        try:
            WebDriverWait(driver=driver, timeout=time_out_2).until(
                EC.presence_of_element_located((By.CLASS_NAME, 'selected'))
            ).click()
            time.sleep(5)
        except:
            pass

        try:
            WebDriverWait(driver=driver, timeout=time_out_2).until(
                EC.presence_of_element_located((By.XPATH, '//*[@id="LoginModal"]/div/div/div/div[2]/button'))
            ).click()
            time.sleep(5)
        except:
            pass

    # fetching job list to loop in
    job_list = WebDriverWait(driver, time_out_2).until(
        EC.presence_of_all_elements_located((By.XPATH, '//li[@class="react-job-listing css-108gl9c eigr9kq3"]/div/div/a/div/div[2]'))
    )
    count = 0

    total_no_of_job_found = WebDriverWait(driver, time_out_2).until(
        EC.presence_of_element_located((By.XPATH, '//h1[@data-test="jobCount-H1title"]'))
    ).text.split(" ")[0]

    for job in job_list:
        company_detail = None
        company_name = None
        rating = None
        job_title = None
        job_location = None
        estimated_salary = None
        job_description = None
        company_size = None
        founded = None
        company_type = None
        industry = None
        sector = None
        revenue = None

        try:
            job.click()
            time.sleep(.5)

            company_name = get_details(driver, selector_company_detail, timeout).split("\n")[0]
            rating = get_details(driver, selector_rating, timeout)
            job_title = get_details(driver, selector_job_title, timeout)
            job_location = get_details(driver, selector_job_location, timeout)
            estimated_salary = get_details(driver, selector_estimated_salary, timeout)
            count += 1
            try:
                WebDriverWait(driver=driver, timeout=timeout).until(
                    EC.presence_of_element_located((By.XPATH, '//*[@id="JobDescriptionContainer"]/div[2]'))
                ).click()
            except:
                pass
            job_description = get_details(driver, selector_job_description, timeout, By.ID)

            # checking if CompanyContainer exist, if it doesn't, then no need to fetch data
            company_container = get_details(driver, "CompanyContainer", timeout, By.ID)
            if company_container == None:
                print(f"{count}/{len(job_list)}\tState: {location}", end="\t")
                print(company_name, rating, job_title)
                # print(company_name, rating, job_title, job_location, estimated_salary)
                new_row = {     "Company name": company_name, "Rating": rating, "Location": job_location, "Job title": job_title,
                            "Estimated salary": estimated_salary, "Job description": job_description,
                            "Number of employees": company_size, "Founded": founded, "Sector": sector,
                            "Industry": industry, "Revenue": revenue }
                df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
                continue

            company_size = get_details(driver, selector_company_size, timeout)
            founded = get_details(driver, selector_founded, timeout)
            company_type = get_details(driver, selector_company_type, timeout)
            industry = get_details(driver, selector_industry, timeout)
            sector = get_details(driver, selector_sector, timeout)
            revenue = get_details(driver, selector_revenue, timeout)
            print(f"{count}/{len(job_list)}\tState: {location}", end="\t")
            print(company_name, rating, job_title)
            # print(company_name, rating, job_title, job_location, estimated_salary, company_size, founded, company_type, industry, sector, revenue)

            new_row = {     "Company name": company_name, "Rating": rating, "Location": job_location, "Job title": job_title,
                            "Estimated salary": estimated_salary, "Job description": job_description,
                            "Number of employees": company_size, "Founded": founded, "Sector": sector,
                            "Industry": industry, "Revenue": revenue, "Keyword": keyword}

            df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)

        except Exception as e:
            print("An error occured:", str(e))

    return df

file_lock = multiprocessing.Lock()

def process_state(url, state, keyword, output_path, completed_files, total_files, total_progress):
    link_df = get_job_details(url=url, location=state, keyword=keyword)

    with file_lock:
        link_df.to_csv(output_path, mode="a", index=False, header=not os.path.exists(output_path))

    print(f"Processed jobs in {state}")
    completed_files.value += 1
    progress = completed_files.value / total_files * 100
    total_progress.value = progress
    print(f"Progress: {progress:.2f}%")

def main():
    output_path = "output.csv"
    state_url_keyword = pd.read_csv("links_merged.csv")

    cpu_count = multiprocessing.cpu_count()
    manager = multiprocessing.Manager()
    total_progress = manager.Value('f', 0.0)

    # Veri kümesini işlem sayısı kadar böl
    partitions = [state_url_keyword[i:i + cpu_count] for i in range(0, len(state_url_keyword), cpu_count)]

    pool = multiprocessing.Pool()

    for i, partition in enumerate(partitions):
        completed_files = manager.Value('i', 0)
        arguments = [(row["Url"], row["State"], row["Keyword"], output_path, completed_files, len(partition), total_progress) for index, row in partition.iterrows()]
        pool.starmap(process_state, arguments)

        completed_percentage = (i + 1) / len(partitions) * 100
        total_progress.value = completed_percentage
        print(f"Total Progress: {completed_percentage:.2f}%")
        time.sleep(3)

    pool.close()
    pool.join()

    print("All jobs processed successfully.")

if __name__ == "__main__":
    main()


from sqlalchemy import (
    Table,
    MetaData,
    Column,
    Integer,
    String,
    DateTime,
    Boolean,
)
from sqlalchemy.engine import create_engine
from sqlalchemy.orm import mapper, relationship

from domain.model import Birthday, GroupConfig, GroupHour, Staff
import config

metadata = MetaData()

birthday = Table(
    "birthday",
    metadata,
    Column("user_id", Integer, primary_key=True, autoincrement=False),
    Column("day", Integer),
    Column("month", Integer),
)

group_config = Table(
    "group_config",
    metadata,
    Column("chat_id", Integer, primary_key=True, autoincrement=False),
    Column("offset", Integer),
    Column("congrats_message", String),
)

group_hour = Table(
    "group_hour",
    metadata,
    Column("chat_id", Integer, primary_key=True, autoincrement=False),
    Column("hour", Integer, primary_key=True, autoincrement=False),
)

staff = Table(
    "staff",
    metadata,
    Column("user_id", Integer, primary_key=True, autoincrement=False),
)

def start_mappers():
    mapper(Birthday, birthday)
    mapper(GroupConfig, group_config)
    mapper(GroupHour, group_hour)
    mapper(Staff, staff)

def create_tables():
    engine = create_engine(config.get_sqlite_uri())
    metadata.create_all(engine)


import aiohttp

from userge import Config, Message, logging, pool, userge

NEKOBIN_URL = "https://nekobin.com/"

_LEVELS = {
    "debug": logging.DEBUG,
    "info": logging.INFO,
    "warning": logging.WARNING,
    "error": logging.ERROR,
    "critical": logging.CRITICAL,
}

@userge.on_cmd(
    "logs",
    about={
        "header": "check USERGE-X logs",
        "flags": {
            "-d": "get logs in document",
            "-h": "get heroku logs",
            "-l": "heroku logs lines limit : default 100",
        },
    },
    allow_channels=False,
)
async def check_logs(message: Message):
    """check logs"""
    await message.edit("`checking logs ...`")
    if "-h" in message.flags and Config.HEROKU_APP:
        limit = int(message.flags.get("-l", 100))
        logs = await pool.run_in_thread(Config.HEROKU_APP.get_log)(lines=limit)
        await message.client.send_as_file(
            chat_id=message.chat.id,
            text=logs,
            filename="userge-heroku.log",
            caption=f"userge-heroku.log [ {limit} lines ]",
        )
    elif "-d" not in message.flags:
        with open("logs/userge.log", "r") as d_f:
            text = d_f.read()
        async with aiohttp.ClientSession() as ses:
            async with ses.post(
                NEKOBIN_URL + "api/documents", json={"content": text}
            ) as resp:
                if resp.status == 201:
                    response = await resp.json()
                    key = response["result"]["key"]
                    file_ext = ".txt"
                    final_url = NEKOBIN_URL + key + file_ext
                    final_url_raw = f"{NEKOBIN_URL}raw/{key}{file_ext}"
                    reply_text = "**Here Are Your Logs** :\n"
                    reply_text += (
                        f"• [NEKO]({final_url})            • [RAW]({final_url_raw})"
                    )
                    await message.edit(reply_text, disable_web_page_preview=True)
                else:
                    await message.edit("Failed to reach Nekobin !")
                    await message.client.send_document(
                        chat_id=message.chat.id,
                        document="logs/userge.log",
                        caption="**USERGE-X Logs**",
                    )
    else:
        await message.delete()
        await message.client.send_document(
            chat_id=message.chat.id,
            document="logs/userge.log",
            caption="**USERGE-X Logs**",
        )

@userge.on_cmd(
    "setlvl",
    about={
        "header": "set logger level, default to info",
        "types": ["debug", "info", "warning", "error", "critical"],
        "usage": "{tr}setlvl [level]",
        "examples": ["{tr}setlvl info", "{tr}setlvl debug"],
    },
)
async def set_level(message: Message):
    """set logger level"""
    await message.edit("`setting logger level ...`")
    level = message.input_str.lower()
    if level not in _LEVELS:
        await message.err("unknown level !")
        return
    for logger in (logging.getLogger(name) for name in logging.root.manager.loggerDict):
        logger.setLevel(_LEVELS[level])
    await message.edit(
        f"`successfully set logger level as` : **{level.upper()}**", del_in=3
    )


from flask import request
from flask_socketio import emit
from app import db
from app.models.bro import Bro
from json import loads
from sqlalchemy import func
import random

from app.models.bro_bros import BroBros
from app.models.broup import Broup
from app.models.broup_message import BroupMessage
from datetime import datetime

from app.util.util import remove_last_occur, update_broups

def add_bro(data):
    token = data["token"]
    bros_bro_id = data["bros_bro_id"]
    logged_in_bro = Bro.verify_auth_token(token)
    if logged_in_bro:
        bro_to_be_added = Bro.query.filter_by(id=bros_bro_id).first()

        logged_in_bro_room = "room_%s" % logged_in_bro.id
        bro_added_exists = BroBros.query.filter_by(bro_id=logged_in_bro.id, bros_bro_id=bros_bro_id).first()
        if bro_added_exists:
            # If the bro object already exists, he is probably removed. We will undo the removing.
            bro_added_exists.re_join()
            db.session.add(bro_added_exists)
            db.session.commit()
            emit("message_event_add_bro_success", bro_added_exists.serialize, room=logged_in_bro_room)
        elif bro_to_be_added:
            if bro_to_be_added.id != logged_in_bro.id:
                chat_colour = '%02X%02X%02X' % (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
                bros_bro = logged_in_bro.add_bro(bro_to_be_added, chat_colour)
                bros_bro_added_bro = bro_to_be_added.add_bro(logged_in_bro, chat_colour)
                db.session.commit()
                if bros_bro:
                    emit("message_event_add_bro_success", bros_bro.serialize, room=logged_in_bro_room)
                else:
                    emit("message_event_add_bro_failed", "Bro was already added", room=logged_in_bro_room)

                if bros_bro_added_bro:
                    bro_room = "room_%s" % bro_to_be_added.id
                    emit("message_event_bro_added_you", bros_bro_added_bro.serialize, room=bro_room)
            else:
                emit("message_event_add_bro_failed", "You tried to add yourself", room=logged_in_bro_room)
        else:
            emit("message_event_add_bro_failed", "failed to add bro", room=logged_in_bro_room)
    else:
        emit("message_event_add_bro_failed", "failed to add bro", room=request.sid)

def add_broup(data):
    token = data["token"]
    logged_in_bro = Bro.verify_auth_token(token)
    if not logged_in_bro:
        emit("message_event_add_broup_failed", "failed to add broup", room=request.sid)
    else:
        participants = loads(data["participants"])

        broup_name = data["broup_name"]
        broup_name += " " + logged_in_bro.bromotion

        max_broup_id = db.session.query(func.max(Broup.broup_id)).scalar()
        if max_broup_id is None:
            max_broup_id = 0
        broup_id = max_broup_id + 1

        bro_ids = [logged_in_bro.id]
        admins = [logged_in_bro.id]
        broup = [logged_in_bro]
        for part in participants:
            bro_for_broup = Bro.query.filter_by(id=part).first()
            if bro_for_broup is None:
                emit("message_event_add_broup_failed", "failed to add broup", room=request.sid)
                return
            else:
                broup.append(bro_for_broup)
                broup_name += "" + bro_for_broup.bromotion
                bro_ids.append(bro_for_broup.id)

        broup_colour = '%02X%02X%02X' % (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
        for bro in broup:
            b = bro.add_broup(broup_name, broup_id, bro_ids, broup_colour, admins)
            room = "room_%s" % bro.id
            if b is not None:
                emit("message_event_added_to_broup", b.serialize, room=room)

        db.session.commit()

        emit("message_event_add_broup_success", "broup was added!", room=request.sid)
        return

def add_bro_to_broup(data):
    token = data["token"]
    logged_in_bro = Bro.verify_auth_token(token)
    if not logged_in_bro:
        emit("message_event_add_bro_to_broup_failed", "adding bro to broup failed", room=request.sid)
    else:
        broup_id = data["broup_id"]
        bro_id = data["bro_id"]

        broup_objects = Broup.query.filter_by(broup_id=broup_id)
        if broup_objects is None:
            emit("message_event_add_bro_to_broup_failed", "adding bro to broup failed", room=request.sid)
            return
        else:
            new_bro_for_broup = Bro.query.filter_by(id=bro_id).first()
            if not new_bro_for_broup:
                emit("message_event_add_bro_to_broup_failed", "adding bro to broup failed", room=request.sid)
                return

            bro_broup = Broup.query.filter_by(broup_id=broup_id, bro_id=bro_id).first()
            if bro_broup is not None and not bro_broup.has_left() and not bro_broup.is_removed():
                emit("message_event_add_bro_to_broup_failed", "adding bro to broup failed", room=request.sid)
            else:
                broup_name = ""
                broup_colour = ""
                broup_description = ""
                bro_ids = [bro_id]
                admins = []
                for broup in broup_objects:
                    if not broup.has_left() and not broup.is_removed():
                        broup_name = broup.get_broup_name() + "" + new_bro_for_broup.get_bromotion()
                        broup.set_broup_name(broup_name)
                        broup.add_participant(bro_id)
                        db.session.add(broup)
                        # Appending the bro ids from each of the broups
                        bro_ids.append(broup.get_bro_id())
                        # These should be the same in all broup objects
                        admins = broup.get_admins()
                        broup_colour = broup.get_broup_colour()
                        broup_description = broup.get_broup_description()

                new_bro_room = "room_%s" % new_bro_for_broup.id
                if bro_broup:
                    # If a broup object already exists than it has to be because the bro left or removed this broup.
                    # We add the bro again to the broup
                    bro_broup.set_participants(bro_ids)
                    bro_broup.set_broup_name(broup_name)
                    bro_broup.set_admins(admins)
                    bro_broup.update_description(broup_description)
                    bro_broup.update_colour(broup_colour)
                    bro_broup.rejoin()
                    db.session.add(bro_broup)
                    emit("message_event_added_to_broup", bro_broup.serialize, room=new_bro_room)
                else:
                    b = new_bro_for_broup.add_broup(broup_name, broup_id, bro_ids, broup_colour, admins, broup_description)
                    if b is not None:
                        emit("message_event_added_to_broup", b.serialize, room=new_bro_room)

                broup_message = BroupMessage(
                    sender_id=bro_id,
                    broup_id=broup_id,
                    body="%s has been added to the broup!" % new_bro_for_broup.get_full_name(),
                    text_message="",
                    timestamp=datetime.utcnow(),
                    info=True
                )
                db.session.add(broup_message)
                db.session.commit()

                broup_room = "broup_%s" % broup_id
                emit("message_event_send", broup_message.serialize, room=broup_room)

                update_broups(broup_objects)

def change_broup_remove_bro(data):
    token = data["token"]
    broup_id = data["broup_id"]
    bro_id = data["bro_id"]
    logged_in_bro = Bro.verify_auth_token(token)

    if logged_in_bro is None:
        emit("message_event_change_broup_remove_bro_failed", "token authentication failed", room=request.sid)
    else:
        broup_objects = Broup.query.filter_by(broup_id=broup_id)
        remove_broup = Broup.query.filter_by(broup_id=broup_id, bro_id=bro_id).first()
        remove_bro = Bro.query.filter_by(id=bro_id).first()
        if broup_objects is None or remove_broup is None or remove_bro is None:
            emit("message_event_change_broup_remove_bro_failed", "broup finding failed", room=request.sid)
        else:
            remove_bromotion = remove_bro.get_bromotion()
            for broup in broup_objects:
                broup_name = remove_last_occur(broup.get_broup_name(), remove_bromotion)
                broup.set_broup_name(broup_name)
                broup.remove_bro(bro_id)
                db.session.add(broup)

            # It's possible that the user who left was the only admin in the broup at the moment.
            # When there are no admins left we randomly assign a user to be admin.
            if len(remove_broup.get_admins()) == 0:
                # In this special event we will make everyone remaining an admin.
                for broup in broup_objects:
                    if not broup.has_left() and not broup.is_removed():
                        new_admin_id = broup.bro_id
                        for broup2 in broup_objects:
                            broup2.add_admin(new_admin_id)
                            db.session.add(broup2)

            # It's possible, that the user left the broup and that there are no more bros left.
            # In this case we will remove all the messages
            if len(remove_broup.get_participants()) == 0:
                messages = BroupMessage.query.filter_by(broup_id=broup_id)
                for message in messages:
                    db.session.delete(message)

            remove_broup.leave_broup()
            db.session.add(remove_broup)

            leave_message = "%s has been removed" % remove_bro.get_full_name()
            if bro_id == logged_in_bro.id:
                leave_message = "%s has left" % remove_bro.get_full_name()

            broup_message = BroupMessage(
                sender_id=bro_id,
                broup_id=broup_id,
                body=leave_message,
                text_message="",
                timestamp=datetime.utcnow(),
                info=True
            )
            db.session.add(broup_message)

            db.session.commit()

            update_broups(broup_objects)

            broup_room = "broup_%s" % broup_id
            emit("message_event_send", broup_message.serialize, room=broup_room)

            bro_room = "room_%s" % remove_broup.bro_id
            emit("message_event_chat_changed", remove_broup.serialize, room=bro_room)


from cx_Freeze import setup, Executable

# Dependencies are automatically detected, but it might need fine-tuning.
build_exe_options = {
    "packages": ["os", "streamlit", "pyarrow"],  # add other packages if needed
    "excludes": []  # add packages you want to exclude if needed
}

base = None

setup(
    name="VoucherVisionEditor",
    version="0.1",
    description="Voucher Vision Editor App",
    options={"build_exe": build_exe_options},
    executables=[Executable("run.py", base=base)]
)


import xarray
import xarray.plot
import numpy as np

def compute_rpe(initial_state_file_name='initial_state.nc',
                output_file_prefix='output_', num_files=5):
    """
    Computes the reference (resting) potential energy for the whole domain
    Parameters
    ----------
    initial_state_file_name : str
        Name of the netCDF file containing the initial state

    output_file_prefix : str
        Prefix for the netCDF file containing output of forward step

    num_files : int
        Number of files on which to compute rpe

    Returns
    -------
    rpe : numpy.ndarray
          the reference potential energy of size (num_timesteps) x (num_files)
    """
    gravity = 9.80616

    # --- Open and read vars from netCDF file
    dsInit = xarray.open_dataset(initial_state_file_name)
    nCells = dsInit.sizes['nCells']
    nEdges = dsInit.sizes['nEdges']
    nVertLevels = dsInit.sizes['nVertLevels']

    xCell = dsInit.xCell
    yCell = dsInit.yCell
    xEdge = dsInit.xEdge
    yEdge = dsInit.yEdge
    areaCell = dsInit.areaCell
    minLevelCell = dsInit.minLevelCell - 1
    maxLevelCell = dsInit.maxLevelCell - 1
    bottomDepth = dsInit.bottomDepth

    # --- Compute a few quantities relevant to domain geometry
    areaCellMatrix = np.tile(areaCell, (nVertLevels, 1)).transpose()
    bottomMax = np.max(bottomDepth.values)
    yMin = np.min(yEdge.values)
    yMax = np.max(yEdge.values)
    xMin = np.min(xEdge.values)
    xMax = np.max(xEdge.values)
    areaDomain = (yMax - yMin) * (xMax - xMin)

    vert_index = \
        xarray.DataArray.from_dict({'dims': ('nVertLevels',),
                                    'data': np.arange(nVertLevels)})

    cell_mask = np.logical_and(vert_index >= minLevelCell,
                               vert_index <= maxLevelCell)
    cell_mask = np.swapaxes(cell_mask, 0, 1)
    nCells_1D = len(areaCellMatrix[cell_mask])

    # --- Get number of timesteps, assuming all output files have the same
    # --- number of timesteps
    ds = xarray.open_dataset('{}{}.nc'.format(
                             output_file_prefix, 1))
    nt = ds.sizes['Time']
    ds.close()

    # --- Allocations
    rpe1 = np.zeros(nCells_1D)
    zMid = np.zeros(nCells_1D)
    thickness = np.zeros(nCells_1D)
    rpe = np.ones((num_files, nt))

    for n in range(num_files):

        ds = xarray.open_dataset('{}{}.nc'.format(
                                 output_file_prefix, n+1))

        xtime = ds.xtime.values
        hFull = ds.layerThickness
        densityFull = ds.density

        for tidx, t in enumerate(xtime):

            h = hFull[tidx, :, :].values
            vol = np.multiply(h, areaCellMatrix)
            density = densityFull[tidx, :, :].values
            density_1D = density[cell_mask]
            vol_1D = vol[cell_mask]

            # --- Density sorting in ascending order
            sorted_ind = np.argsort(density_1D)
            density_sorted = density_1D[sorted_ind]
            vol_sorted = vol_1D[sorted_ind]

            thickness = np.divide(vol_sorted.tolist(), areaDomain)

            # --- RPE computation
            z = np.append([0.], -np.cumsum(thickness))
            zMid = z[0:-1] - thickness/2. + bottomMax
            rpe1 = gravity * np.multiply(
                       np.multiply(density_sorted, zMid),
                       vol_sorted)

            rpe[n, tidx] = np.sum(rpe1)/np.sum(areaCell)

        ds.close()

    # --- Write rpe to text file
    rows = ['{}'.format(t.astype(str)) for t in xtime]
    with open('rpe.txt', 'w+') as csvfile:
        col_headings = 'time'
        for n in range(num_files):
            col_headings += ',output_' + str(n+1)
        col_headings += '\n'
        csvfile.writelines(col_headings)
        for tidx, t in enumerate(xtime):
            for n in range(num_files):
                rows[tidx] += ',%f' % rpe[n, tidx]
            rows[tidx] += '\n'
        csvfile.writelines(rows)
    csvfile.close()

    return rpe