import logging
import logging
import multiprocessing
import os
import pathlib
import pickle
import time
import zipfile
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
import yaml
from sklearn.preprocessing import OneHotEncoder
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm

from dataloader import load_graph_adj_mtx, load_graph_node_features
from model import GCN, NodeAttnMap, UserEmbeddings, Time2Vec, CategoryEmbeddings, FuseEmbeddings, TransformerModel, \
    GraphSage,GRUModel
from param_parser import parameter_parser
from utils import increment_path, calculate_laplacian_matrix, zipdir, top_k_acc_last_timestep, \
    mAP_metric_last_timestep, MRR_metric_last_timestep, maksed_mse_loss, adj_list, split_list, random_walk_with_restart,\
    split_list_by_ratio

def train(args):
    args.save_dir = increment_path(Path(args.project) / args.name, exist_ok=args.exist_ok, sep='-')
    if not os.path.exists(args.save_dir): os.makedirs(args.save_dir)

    # Setup logger
    for handler in logging.root.handlers[:]:
        logging.root.removeHandler(handler)
    logging.basicConfig(level=logging.DEBUG,
                        format='%(asctime)s %(message)s',
                        datefmt='%Y-%m-%d %H:%M:%S',
                        filename=os.path.join(args.save_dir, f"log_training.txt"),
                        filemode='w')
    console = logging.StreamHandler()
    console.setLevel(logging.INFO)
    formatter = logging.Formatter('%(asctime)s %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
    console.setFormatter(formatter)
    logging.getLogger('').addHandler(console)
    logging.getLogger('matplotlib.font_manager').disabled = True

    # Save run settings
    logging.info(args)
    with open(os.path.join(args.save_dir, 'args.yaml'), 'w') as f:
        yaml.dump(vars(args), f, sort_keys=False)

    # Save python code
    zipf = zipfile.ZipFile(os.path.join(args.save_dir, 'code.zip'), 'w', zipfile.ZIP_DEFLATED)
    zipdir(pathlib.Path().absolute(), zipf, include_format=['.py'])
    zipf.close()

    # %% ====================== Load data ======================
    # Read check-in train data
    train_df = pd.read_csv(args.data_train)
    val_df = pd.read_csv(args.data_val)
    if args.data_train!='dataset/dataset_tsmc2014/NYC_train.csv':
        train_df['timestamp'] = pd.to_datetime(train_df['local_time']).astype('int64') // 10 ** 9
        val_df['timestamp'] = pd.to_datetime(val_df['local_time']).astype('int64') // 10 ** 9

    # Build POI graph (built from train_df)
    print('Loading POI graph...')
    raw_A = load_graph_adj_mtx(args.data_adj_mtx)
    raw_X = load_graph_node_features(args.data_node_feats,
                                     args.feature1,
                                     args.feature2,
                                     args.feature3,
                                     args.feature4)
    logging.info(
        f"raw_X.shape: {raw_X.shape}; "
        f"Four features: {args.feature1}, {args.feature2}, {args.feature3}, {args.feature4}.")
    logging.info(f"raw_A.shape: {raw_A.shape}; Edge from row_index to col_index with weight (frequency).")
    num_pois = raw_X.shape[0]

    # One-hot encoding poi categories
    logging.info('One-hot encoding poi categories id')
    one_hot_encoder = OneHotEncoder()
    cat_list = list(raw_X[:, 1])
    one_hot_encoder.fit(list(map(lambda x: [x], cat_list)))
    one_hot_rlt = one_hot_encoder.transform(list(map(lambda x: [x], cat_list))).toarray()
    num_cats = one_hot_rlt.shape[-1]
    X = np.zeros((num_pois, raw_X.shape[-1] - 1 + num_cats), dtype=np.float32)
    X[:, 0] = raw_X[:, 0]
    X[:, 1:num_cats + 1] = one_hot_rlt
    X[:, num_cats + 1:] = raw_X[:, 2:]
    logging.info(f"After one hot encoding poi cat, X.shape: {X.shape}")
    logging.info(f'POI categories: {list(one_hot_encoder.categories_[0])}')
    # Save ont-hot encoder
    with open(os.path.join(args.save_dir, 'one-hot-encoder.pkl'), "wb") as f:
        pickle.dump(one_hot_encoder, f)

    # Normalization
    print('Laplician matrix...')
    A = calculate_laplacian_matrix(raw_A, mat_type='hat_rw_normd_lap_mat')

    # POI id to index
    nodes_df = pd.read_csv(args.data_node_feats)
    poi_ids = nodes_df['node_name/poi_id'].tolist()
    poi_id2idx_dict = dict(zip(poi_ids, range(len(poi_ids))))

    # Cat id to index
    cat_ids = list(set(nodes_df[args.feature2].tolist()))
    cat_id2idx_dict = dict(zip(cat_ids, range(len(cat_ids))))

    # Poi idx to cat idx
    poi_idx2cat_idx_dict = {}
    for i, row in nodes_df.iterrows():
        poi_idx2cat_idx_dict[poi_id2idx_dict[row['node_name/poi_id']]] = \
            cat_id2idx_dict[row[args.feature2]]

    # User id to index
    user_ids = [str(each) for each in list(set(train_df['user_id'].to_list()))]
    user_id2idx_dict = dict(zip(user_ids, range(len(user_ids))))

    # Print user-trajectories count
    traj_list = list(set(train_df['trajectory_id'].tolist()))
    print('load adj_list')
    if os.path.exists(os.path.join(args.adj_path, 'adj.pkl')):
        with open(os.path.join(args.adj_path, 'adj.pkl'), 'rb') as f:  # 打开pickle文件
            adj = pickle.load(f)  # 读取字典
        with open(os.path.join(args.adj_path, 'dis.pkl'), 'rb') as f:  # 打开pickle文件
            dis = pickle.load(f)  # 读取字典
    else:
        adj, dis = adj_list(raw_A, raw_X, args.geo_dis)
        with open(os.path.join(args.adj_path, 'adj.pkl'), 'wb') as f:
            pickle.dump(adj, f)  # 把字典写入pickle文件
        with open(os.path.join(args.adj_path, 'dis.pkl'), 'wb') as f:
            pickle.dump(dis, f)  # 把字典写入pickle文件

    # %% ====================== Define Dataset ======================
    class TrajectoryDatasetTrain(Dataset):
        def __init__(self, train_df):
            self.df = train_df
            self.traj_seqs = []  # traj id: user id + traj no.
            self.input_seqs = []
            self.label_seqs = []

            for traj_id in tqdm(set(train_df['trajectory_id'].tolist())):
                traj_df = train_df[train_df['trajectory_id'] == traj_id]
                poi_ids = traj_df['POI_id'].to_list()
                poi_idxs = [poi_id2idx_dict[each] for each in poi_ids]
                time_feature = traj_df[args.time_feature].to_list()
                ts = traj_df['timestamp'].to_list()

                input_seq = []
                label_seq = []
                for i in range(len(poi_idxs) - 1):
                    input_seq.append((poi_idxs[i], time_feature[i],ts[i]))
                    label_seq.append((poi_idxs[i + 1], time_feature[i + 1],ts[i+1]))

                if len(input_seq) < args.short_traj_thres:
                    continue

                self.traj_seqs.append(traj_id)
                self.input_seqs.append(input_seq)
                self.label_seqs.append(label_seq)

        def __len__(self):
            assert len(self.input_seqs) == len(self.label_seqs) == len(self.traj_seqs)
            return len(self.traj_seqs)

        def __getitem__(self, index):
            return (self.traj_seqs[index], self.input_seqs[index], self.label_seqs[index])

    class TrajectoryDatasetVal(Dataset):
        def __init__(self, df):
            self.df = df
            self.traj_seqs = []
            self.input_seqs = []
            self.label_seqs = []

            for traj_id in tqdm(set(df['trajectory_id'].tolist())):
                user_id = traj_id.split('_')[0]

                # Ignore user if not in training set
                if user_id not in user_id2idx_dict.keys():
                    continue

                # Ger POIs idx in this trajectory
                traj_df = df[df['trajectory_id'] == traj_id]
                poi_ids = traj_df['POI_id'].to_list()
                poi_idxs = []
                time_feature = traj_df[args.time_feature].to_list()
                ts = traj_df['timestamp'].to_list()

                for each in poi_ids:
                    if each in poi_id2idx_dict.keys():
                        poi_idxs.append(poi_id2idx_dict[each])
                    else:
                        # Ignore poi if not in training set
                        continue

                # Construct input seq and label seq
                input_seq = []
                label_seq = []
                for i in range(len(poi_idxs) - 1):
                    input_seq.append((poi_idxs[i], time_feature[i],ts[i]))
                    label_seq.append((poi_idxs[i + 1], time_feature[i + 1],ts[i+1]))

                # Ignore seq if too short
                if len(input_seq) < args.short_traj_thres:
                    continue

                self.input_seqs.append(input_seq)
                self.label_seqs.append(label_seq)
                self.traj_seqs.append(traj_id)

        def __len__(self):
            assert len(self.input_seqs) == len(self.label_seqs) == len(self.traj_seqs)
            return len(self.traj_seqs)

        def __getitem__(self, index):
            return (self.traj_seqs[index], self.input_seqs[index], self.label_seqs[index])

    # %% ====================== Define dataloader ======================
    print('Prepare dataloader...')
    train_dataset = TrajectoryDatasetTrain(train_df)
    val_dataset = TrajectoryDatasetVal(val_df)

    train_loader = DataLoader(train_dataset,
                              batch_size=args.batch,
                              shuffle=True, drop_last=False,
                              pin_memory=True, num_workers=args.workers,
                              collate_fn=lambda x: x)
    val_loader = DataLoader(val_dataset,
                            batch_size=args.batch,
                            shuffle=False, drop_last=False,
                            pin_memory=True, num_workers=args.workers,
                            collate_fn=lambda x: x)

    class produceSampleProcess(multiprocessing.Process):
        def __init__(self, tasks, queues, adj_list, restart_prob, num_walks, threshold, adjOrdis, stop_event, id):
            super().__init__()
            self.tasks = tasks
            self.queues = queues
            self.threshold = threshold
            self.adjOrdis = adjOrdis
            self.stop_event = stop_event
            self.id = id
            self.adj_list = adj_list
            self.restart_prob = restart_prob
            self.num_walks = num_walks
            self.count_dict={key:threshold for key in tasks}
            self.missing_dict={key:0 for key in tasks}

        def run(self):
            while True:
                for node in self.tasks:
                    q = self.queues[node]
                    if q.qsize() < self.threshold/2:
                        for _ in range(self.count_dict[node]-q.qsize()):
                            random_walk = random_walk_with_restart(self.adj_list, node, self.restart_prob, self.num_walks,
                                                               self.adjOrdis)
                            q.put(random_walk)
                        self.missing_dict[node]+=1
                        if self.missing_dict[node]>2:
                            self.missing_dict[node]=0
                            self.count_dict[node]+=self.threshold
                if self.stop_event.is_set():
                    break
            print(self.adjOrdis, self.id, 'quit')

    threshold = 10  # 队列大小阈值
    adj_queues = {node: multiprocessing.Queue() for node in range(num_pois)}  # 创建多个队列
    dis_queues = {node: multiprocessing.Queue() for node in range(num_pois)}  # 创建多个队列
    tasks = split_list([i for i in range(num_pois)],int(args.cpus/2))
    stop_event = multiprocessing.Event()

    for idx, task in enumerate(tasks):
        ap = produceSampleProcess(tasks=task, queues=adj_queues, adj_list=adj, restart_prob=args.restart_prob,
                                  num_walks=args.num_walks,
                                  threshold=threshold, adjOrdis='adj', stop_event=stop_event, id=idx)
        ap.start()
        dp = produceSampleProcess(tasks=task, queues=dis_queues, adj_list=dis, restart_prob=args.restart_prob,
                                  num_walks=args.num_walks,
                                  threshold=threshold, adjOrdis='dis', stop_event=stop_event, id=idx)
        dp.start()

    # %% ====================== Build Models ======================
    # Model1: POI embedding model
    if isinstance(X, np.ndarray):
        X = torch.from_numpy(X)
        A = torch.from_numpy(A)
    X = X.to(device=args.device, dtype=torch.float)
    A = A.to(device=args.device, dtype=torch.float)

    poi_embed_model = GraphSage(X=X,  embed_dim=args.poi_embed_dim, adj=adj, dis=dis,
                                device=args.device, restart_prob=args.restart_prob, num_walks=args.num_walks,
                                dropout=args.dropout, adj_queues=adj_queues, dis_queues=dis_queues)

    # %% Model2: User embedding model, nn.embedding
    num_users = len(user_id2idx_dict)
    user_embed_model = UserEmbeddings(num_users, args.user_embed_dim)

    # %% Model3: Time Model
    time_embed_model = Time2Vec('sin', out_dim=args.time_embed_dim)

    # %% Model4: Category embedding model
    cat_embed_model = CategoryEmbeddings(num_cats, args.cat_embed_dim)

    # %% Model5: Embedding fusion models
    embed_fuse_model1 = FuseEmbeddings(args.user_embed_dim, args.poi_embed_dim)
    embed_fuse_model2 = FuseEmbeddings(args.time_embed_dim, args.cat_embed_dim)

    # %% Model6: Sequence model
    args.seq_input_embed = args.poi_embed_dim + args.user_embed_dim + args.time_embed_dim + args.cat_embed_dim
    seq_model = TransformerModel(num_pois,
                                 num_cats,
                                 args.seq_input_embed,
                                 args.transformer_nhead,
                                 args.transformer_nhid,
                                 args.transformer_nlayers,
                                 dropout=args.transformer_dropout)

    # Define overall loss and optimizer
    optimizer = optim.Adam(params=list(poi_embed_model.parameters()) +
                                  list(user_embed_model.parameters()) +
                                  list(time_embed_model.parameters()) +
                                  list(cat_embed_model.parameters()) +
                                  list(embed_fuse_model1.parameters()) +
                                  list(embed_fuse_model2.parameters()) +
                                  list(seq_model.parameters()),
                           lr=args.lr,
                           weight_decay=args.weight_decay)

    criterion_poi = nn.CrossEntropyLoss(ignore_index=-1)  # -1 is padding
    criterion_cat = nn.CrossEntropyLoss(ignore_index=-1)  # -1 is padding
    criterion_time = maksed_mse_loss

    lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, 'min', verbose=True, factor=args.lr_scheduler_factor)

    # %% Tool functions for training
    def input_traj_to_embeddings(sample, poi_embeddings, embedding_index):
        # Parse sample
        traj_id = sample[0]
        input_seq = [each[0] for each in sample[1]]
        input_seq_time = [each[1] for each in sample[1]]
        input_seq_cat = [poi_idx2cat_idx_dict[each] for each in input_seq]

        # User to embedding
        user_id = traj_id.split('_')[0]
        user_idx = user_id2idx_dict[user_id]
        input = torch.LongTensor([user_idx]).to(device=args.device)
        user_embedding = user_embed_model(input)
        user_embedding = torch.squeeze(user_embedding)

        poi_idxs=[embedding_index+idx for idx in range(len(input_seq))]
        sample_poi_embedding=poi_embeddings[poi_idxs].squeeze()
        sample_time_embedding=time_embed_model(torch.tensor(input_seq_time,dtype=torch.float).to(device=args.device)).squeeze()
        sample_cat_embedding=cat_embed_model(torch.tensor(input_seq_cat,dtype=torch.long).to(device=args.device)).squeeze()
        sample_fused_embedding1=embed_fuse_model1(user_embedding.expand(len(input_seq),-1),sample_poi_embedding)
        sample_fused_embedding2=embed_fuse_model2(sample_time_embedding,sample_cat_embedding)
        sample_concat_embedding=torch.cat((sample_fused_embedding1,sample_fused_embedding2),dim=-1)

        # POI to embedding and fuse embeddings
        '''
        input_seq_embed = []
        for idx in range(len(input_seq)):
            poi_embedding = poi_embeddings[embedding_index+idx]
            poi_embedding = torch.squeeze(poi_embedding).to(device=args.device)

            # Time to vector
            time_embedding = time_embed_model(
                torch.tensor([input_seq_time[idx]], dtype=torch.float).to(device=args.device))
            time_embedding = torch.squeeze(time_embedding).to(device=args.device)

            # Categroy to embedding
            cat_idx = torch.LongTensor([input_seq_cat[idx]]).to(device=args.device)
            cat_embedding = cat_embed_model(cat_idx)
            cat_embedding = torch.squeeze(cat_embedding)

            # Fuse user+poi embeds
            fused_embedding1 = embed_fuse_model1(user_embedding, poi_embedding)
            fused_embedding2 = embed_fuse_model2(time_embedding, cat_embedding)

            # Concat time, cat after user+poi
            concat_embedding = torch.cat((fused_embedding1, fused_embedding2), dim=-1)

            # Save final embed
            input_seq_embed.append(concat_embedding)
        '''
        return sample_concat_embedding

    # %% ====================== Train ======================
    poi_embed_model = poi_embed_model.to(device=args.device)
    user_embed_model = user_embed_model.to(device=args.device)
    time_embed_model = time_embed_model.to(device=args.device)
    cat_embed_model = cat_embed_model.to(device=args.device)
    embed_fuse_model1 = embed_fuse_model1.to(device=args.device)
    embed_fuse_model2 = embed_fuse_model2.to(device=args.device)
    seq_model = seq_model.to(device=args.device)

    # %% Loop epoch
    # For plotting
    train_epochs_top1_acc_list = []
    train_epochs_top5_acc_list = []
    train_epochs_top10_acc_list = []
    train_epochs_top20_acc_list = []
    train_epochs_mAP20_list = []
    train_epochs_mrr_list = []
    train_epochs_loss_list = []
    train_epochs_poi_loss_list = []
    train_epochs_time_loss_list = []
    train_epochs_cat_loss_list = []
    val_epochs_top1_acc_list = []
    val_epochs_top5_acc_list = []
    val_epochs_top10_acc_list = []
    val_epochs_top20_acc_list = []
    val_epochs_mAP20_list = []
    val_epochs_mrr_list = []
    val_epochs_loss_list = []
    val_epochs_poi_loss_list = []
    val_epochs_time_loss_list = []
    val_epochs_cat_loss_list = []
    # For saving ckpt
    max_val_score = -np.inf

    for epoch in range(args.epochs):
        logging.info(f"{'*' * 50}Epoch:{epoch:03d}{'*' * 50}\n")
        poi_embed_model.train()
        user_embed_model.train()
        time_embed_model.train()
        cat_embed_model.train()
        embed_fuse_model1.train()
        embed_fuse_model2.train()
        seq_model.train()

        train_batches_top1_acc_list = []
        train_batches_top5_acc_list = []
        train_batches_top10_acc_list = []
        train_batches_top20_acc_list = []
        train_batches_mAP20_list = []
        train_batches_mrr_list = []
        train_batches_loss_list = []
        train_batches_poi_loss_list = []
        train_batches_time_loss_list = []
        train_batches_cat_loss_list = []
        #src_mask = seq_model.generate_square_subsequent_mask(args.batch).to(args.device)
        # Loop batch
        for b_idx, batch in enumerate(train_loader):
            #if len(batch) != args.batch:
                #src_mask = seq_model.generate_square_subsequent_mask(len(batch)).to(args.device)

            # For padding
            batch_input_seqs = []
            batch_seq_lens = []
            batch_seq_embeds = []
            batch_seq_labels_poi = []
            batch_seq_labels_time = []
            batch_seq_labels_cat = []

            pois = [each[0] for sample in batch for each in sample[1]]
            poi_embeddings = poi_embed_model(torch.tensor(pois).to(args.device))
            # Convert input seq to embeddings

            embedding_index=0
            for sample in batch:
                # sample[0]: traj_id, sample[1]: input_seq, sample[2]: label_seq
                traj_id = sample[0]
                input_seq = [each[0] for each in sample[1]]
                label_seq = [each[0] for each in sample[2]]
                input_seq_time = [each[1] for each in sample[1]]
                label_seq_time = [each[1] for each in sample[2]]
                label_seq_cats = [poi_idx2cat_idx_dict[each] for each in label_seq]
                input_seq_embed = input_traj_to_embeddings(sample, poi_embeddings,embedding_index)
                batch_seq_embeds.append(input_seq_embed)
                batch_seq_lens.append(len(input_seq))
                batch_input_seqs.append(input_seq)
                batch_seq_labels_poi.append(torch.LongTensor(label_seq))
                batch_seq_labels_time.append(torch.FloatTensor(label_seq_time))
                batch_seq_labels_cat.append(torch.LongTensor(label_seq_cats))
                embedding_index+=len(input_seq)

            # Pad seqs for batch training
            batch_padded = pad_sequence(batch_seq_embeds, batch_first=True, padding_value=-1)
            label_padded_poi = pad_sequence(batch_seq_labels_poi, batch_first=True, padding_value=-1)
            label_padded_time = pad_sequence(batch_seq_labels_time, batch_first=True, padding_value=-1)
            label_padded_cat = pad_sequence(batch_seq_labels_cat, batch_first=True, padding_value=-1)

            # Feedforward
            x = batch_padded.to(device=args.device, dtype=torch.float)
            y_poi = label_padded_poi.to(device=args.device, dtype=torch.long)
            y_time = label_padded_time.to(device=args.device, dtype=torch.float)
            y_cat = label_padded_cat.to(device=args.device, dtype=torch.long)
            y_pred_poi, y_pred_time, y_pred_cat = seq_model(x)

            loss_poi = criterion_poi(y_pred_poi.transpose(1, 2), y_poi)
            loss_time = criterion_time(torch.squeeze(y_pred_time), y_time)
            loss_cat = criterion_cat(y_pred_cat.transpose(1, 2), y_cat)

            # Final loss
            loss = loss_poi + loss_time * args.time_loss_weight + loss_cat
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # Performance measurement
            top1_acc = 0
            top5_acc = 0
            top10_acc = 0
            top20_acc = 0
            mAP20 = 0
            mrr = 0
            batch_label_pois = y_poi.detach().cpu().numpy()
            batch_pred_pois = y_pred_poi.detach().cpu().numpy()
            batch_pred_times = y_pred_time.detach().cpu().numpy()
            batch_pred_cats = y_pred_cat.detach().cpu().numpy()
            for label_pois, pred_pois, seq_len in zip(batch_label_pois, batch_pred_pois, batch_seq_lens):
                label_pois = label_pois[:seq_len]  # shape: (seq_len, )
                pred_pois = pred_pois[:seq_len, :]  # shape: (seq_len, num_poi)
                top1_acc += top_k_acc_last_timestep(label_pois, pred_pois, k=1)
                top5_acc += top_k_acc_last_timestep(label_pois, pred_pois, k=5)
                top10_acc += top_k_acc_last_timestep(label_pois, pred_pois, k=10)
                top20_acc += top_k_acc_last_timestep(label_pois, pred_pois, k=20)
                mAP20 += mAP_metric_last_timestep(label_pois, pred_pois, k=20)
                mrr += MRR_metric_last_timestep(label_pois, pred_pois)
            train_batches_top1_acc_list.append(top1_acc / len(batch_label_pois))
            train_batches_top5_acc_list.append(top5_acc / len(batch_label_pois))
            train_batches_top10_acc_list.append(top10_acc / len(batch_label_pois))
            train_batches_top20_acc_list.append(top20_acc / len(batch_label_pois))
            train_batches_mAP20_list.append(mAP20 / len(batch_label_pois))
            train_batches_mrr_list.append(mrr / len(batch_label_pois))
            train_batches_loss_list.append(loss.detach().cpu().numpy())
            train_batches_poi_loss_list.append(loss_poi.detach().cpu().numpy())
            train_batches_time_loss_list.append(loss_time.detach().cpu().numpy())
            train_batches_cat_loss_list.append(loss_cat.detach().cpu().numpy())

            # Report training progress
            if (b_idx % (40)) == 0:
                sample_idx = 0
                logging.info(f'Epoch:{epoch}, batch:{b_idx}, '
                             f'train_batch_loss:{loss.item():.2f}, '
                             f'train_batch_top1_acc:{top1_acc / len(batch_label_pois):.2f}, '
                             f'train_move_loss:{np.mean(train_batches_loss_list):.2f}\n'
                             f'train_move_poi_loss:{np.mean(train_batches_poi_loss_list):.2f}\n'
                             f'train_move_time_loss:{np.mean(train_batches_time_loss_list):.2f}\n'
                             f'train_move_top1_acc:{np.mean(train_batches_top1_acc_list):.4f}\n'
                             f'train_move_top5_acc:{np.mean(train_batches_top5_acc_list):.4f}\n'
                             f'train_move_top10_acc:{np.mean(train_batches_top10_acc_list):.4f}\n'
                             f'train_move_top20_acc:{np.mean(train_batches_top20_acc_list):.4f}\n'
                             f'train_move_mAP20:{np.mean(train_batches_mAP20_list):.4f}\n'
                             f'train_move_MRR:{np.mean(train_batches_mrr_list):.4f}\n'
                             f'traj_id:{batch[sample_idx][0]}\n'
                             f'input_seq: {batch[sample_idx][1]}\n'
                             f'label_seq:{batch[sample_idx][2]}\n'
                             f'pred_seq_poi:{list(np.argmax(batch_pred_pois, axis=2)[sample_idx][:batch_seq_lens[sample_idx]])} \n'
                             f'label_seq_cat:{[poi_idx2cat_idx_dict[each[0]] for each in batch[sample_idx][2]]}\n'
                             f'pred_seq_cat:{list(np.argmax(batch_pred_cats, axis=2)[sample_idx][:batch_seq_lens[sample_idx]])} \n'
                             f'label_seq_time:{list(batch_seq_labels_time[sample_idx].numpy()[:batch_seq_lens[sample_idx]])}\n'
                             f'pred_seq_time:{list(np.squeeze(batch_pred_times)[sample_idx][:batch_seq_lens[sample_idx]])} \n' +
                             '=' * 100)

        # train end --------------------------------------------------------------------------------------------------------
        poi_embed_model.eval()
        user_embed_model.eval()
        time_embed_model.eval()
        cat_embed_model.eval()
        embed_fuse_model1.eval()
        embed_fuse_model2.eval()
        seq_model.eval()
        val_batches_top1_acc_list = []
        val_batches_top5_acc_list = []
        val_batches_top10_acc_list = []
        val_batches_top20_acc_list = []
        val_batches_mAP20_list = []
        val_batches_mrr_list = []
        val_batches_loss_list = []
        val_batches_poi_loss_list = []
        val_batches_time_loss_list = []
        val_batches_cat_loss_list = []
        #src_mask = seq_model.generate_square_subsequent_mask(args.batch).to(args.device)
        for vb_idx, batch in enumerate(val_loader):
            #if len(batch) != args.batch:
                #src_mask = seq_model.generate_square_subsequent_mask(len(batch)).to(args.device)

            # For padding
            batch_input_seqs = []
            batch_seq_lens = []
            batch_seq_embeds = []
            batch_seq_labels_poi = []
            batch_seq_labels_time = []
            batch_seq_labels_cat = []

            pois = [each[0] for sample in batch for each in sample[1]]
            poi_embeddings = poi_embed_model(torch.tensor(pois).to(args.device))
            embedding_index=0
            # Convert input seq to embeddings
            for sample in batch:
                traj_id = sample[0]
                input_seq = [each[0] for each in sample[1]]
                label_seq = [each[0] for each in sample[2]]
                input_seq_time = [each[1] for each in sample[1]]
                label_seq_time = [each[1] for each in sample[2]]
                label_seq_cats = [poi_idx2cat_idx_dict[each] for each in label_seq]
                input_seq_embed = input_traj_to_embeddings(sample, poi_embeddings,embedding_index)
                batch_seq_embeds.append(input_seq_embed)
                batch_seq_lens.append(len(input_seq))
                batch_input_seqs.append(input_seq)
                batch_seq_labels_poi.append(torch.LongTensor(label_seq))
                batch_seq_labels_time.append(torch.FloatTensor(label_seq_time))
                batch_seq_labels_cat.append(torch.LongTensor(label_seq_cats))
                embedding_index+=len(input_seq)

            # Pad seqs for batch training
            batch_padded = pad_sequence(batch_seq_embeds, batch_first=True, padding_value=-1)
            label_padded_poi = pad_sequence(batch_seq_labels_poi, batch_first=True, padding_value=-1)
            label_padded_time = pad_sequence(batch_seq_labels_time, batch_first=True, padding_value=-1)
            label_padded_cat = pad_sequence(batch_seq_labels_cat, batch_first=True, padding_value=-1)

            # Feedforward
            x = batch_padded.to(device=args.device, dtype=torch.float)
            y_poi = label_padded_poi.to(device=args.device, dtype=torch.long)
            y_time = label_padded_time.to(device=args.device, dtype=torch.float)
            y_cat = label_padded_cat.to(device=args.device, dtype=torch.long)
            y_pred_poi, y_pred_time, y_pred_cat = seq_model(x)

            # Calculate loss
            loss_poi = criterion_poi(y_pred_poi.transpose(1, 2), y_poi)
            loss_time = criterion_time(torch.squeeze(y_pred_time), y_time)
            loss_cat = criterion_cat(y_pred_cat.transpose(1, 2), y_cat)
            loss = loss_poi + loss_time * args.time_loss_weight + loss_cat

            # Performance measurement
            top1_acc = 0
            top5_acc = 0
            top10_acc = 0
            top20_acc = 0
            mAP20 = 0
            mrr = 0
            batch_label_pois = y_poi.detach().cpu().numpy()
            batch_pred_pois = y_pred_poi.detach().cpu().numpy()
            batch_pred_times = y_pred_time.detach().cpu().numpy()
            batch_pred_cats = y_pred_cat.detach().cpu().numpy()
            for label_pois, pred_pois, seq_len in zip(batch_label_pois, batch_pred_pois, batch_seq_lens):
                label_pois = label_pois[:seq_len]  # shape: (seq_len, )
                pred_pois = pred_pois[:seq_len, :]  # shape: (seq_len, num_poi)
                top1_acc += top_k_acc_last_timestep(label_pois, pred_pois, k=1)
                top5_acc += top_k_acc_last_timestep(label_pois, pred_pois, k=5)
                top10_acc += top_k_acc_last_timestep(label_pois, pred_pois, k=10)
                top20_acc += top_k_acc_last_timestep(label_pois, pred_pois, k=20)
                mAP20 += mAP_metric_last_timestep(label_pois, pred_pois, k=20)
                mrr += MRR_metric_last_timestep(label_pois, pred_pois)
            val_batches_top1_acc_list.append(top1_acc / len(batch_label_pois))
            val_batches_top5_acc_list.append(top5_acc / len(batch_label_pois))
            val_batches_top10_acc_list.append(top10_acc / len(batch_label_pois))
            val_batches_top20_acc_list.append(top20_acc / len(batch_label_pois))
            val_batches_mAP20_list.append(mAP20 / len(batch_label_pois))
            val_batches_mrr_list.append(mrr / len(batch_label_pois))
            val_batches_loss_list.append(loss.detach().cpu().numpy())
            val_batches_poi_loss_list.append(loss_poi.detach().cpu().numpy())
            val_batches_time_loss_list.append(loss_time.detach().cpu().numpy())
            val_batches_cat_loss_list.append(loss_cat.detach().cpu().numpy())

            # Report validation progress
            if (vb_idx % (20)) == 0:
                sample_idx = 0
                logging.info(f'Epoch:{epoch}, batch:{vb_idx}, '
                             f'val_batch_loss:{loss.item():.2f}, '
                             f'val_batch_top1_acc:{top1_acc / len(batch_label_pois):.2f}, '
                             f'val_move_loss:{np.mean(val_batches_loss_list):.2f} \n'
                             f'val_move_poi_loss:{np.mean(val_batches_poi_loss_list):.2f} \n'
                             f'val_move_time_loss:{np.mean(val_batches_time_loss_list):.2f} \n'
                             f'val_move_top1_acc:{np.mean(val_batches_top1_acc_list):.4f} \n'
                             f'val_move_top5_acc:{np.mean(val_batches_top5_acc_list):.4f} \n'
                             f'val_move_top10_acc:{np.mean(val_batches_top10_acc_list):.4f} \n'
                             f'val_move_top20_acc:{np.mean(val_batches_top20_acc_list):.4f} \n'
                             f'val_move_mAP20:{np.mean(val_batches_mAP20_list):.4f} \n'
                             f'val_move_MRR:{np.mean(val_batches_mrr_list):.4f} \n'
                             f'traj_id:{batch[sample_idx][0]}\n'
                             f'input_seq:{batch[sample_idx][1]}\n'
                             f'label_seq:{batch[sample_idx][2]}\n'
                             f'pred_seq_poi:{list(np.argmax(batch_pred_pois, axis=2)[sample_idx][:batch_seq_lens[sample_idx]])} \n'
                             f'label_seq_cat:{[poi_idx2cat_idx_dict[each[0]] for each in batch[sample_idx][2]]}\n'
                             f'pred_seq_cat:{list(np.argmax(batch_pred_cats, axis=2)[sample_idx][:batch_seq_lens[sample_idx]])} \n'
                             f'label_seq_time:{list(batch_seq_labels_time[sample_idx].numpy()[:batch_seq_lens[sample_idx]])}\n'
                             f'pred_seq_time:{list(np.squeeze(batch_pred_times)[sample_idx][:batch_seq_lens[sample_idx]])} \n' +
                             '=' * 100)
        # valid end --------------------------------------------------------------------------------------------------------

        # Calculate epoch metrics
        epoch_train_top1_acc = np.mean(train_batches_top1_acc_list)
        epoch_train_top5_acc = np.mean(train_batches_top5_acc_list)
        epoch_train_top10_acc = np.mean(train_batches_top10_acc_list)
        epoch_train_top20_acc = np.mean(train_batches_top20_acc_list)
        epoch_train_mAP20 = np.mean(train_batches_mAP20_list)
        epoch_train_mrr = np.mean(train_batches_mrr_list)
        epoch_train_loss = np.mean(train_batches_loss_list)
        epoch_train_poi_loss = np.mean(train_batches_poi_loss_list)
        epoch_train_time_loss = np.mean(train_batches_time_loss_list)
        epoch_train_cat_loss = np.mean(train_batches_cat_loss_list)
        epoch_val_top1_acc = np.mean(val_batches_top1_acc_list)
        epoch_val_top5_acc = np.mean(val_batches_top5_acc_list)
        epoch_val_top10_acc = np.mean(val_batches_top10_acc_list)
        epoch_val_top20_acc = np.mean(val_batches_top20_acc_list)
        epoch_val_mAP20 = np.mean(val_batches_mAP20_list)
        epoch_val_mrr = np.mean(val_batches_mrr_list)
        epoch_val_loss = np.mean(val_batches_loss_list)
        epoch_val_poi_loss = np.mean(val_batches_poi_loss_list)
        epoch_val_time_loss = np.mean(val_batches_time_loss_list)
        epoch_val_cat_loss = np.mean(val_batches_cat_loss_list)

        # Save metrics to list
        train_epochs_loss_list.append(epoch_train_loss)
        train_epochs_poi_loss_list.append(epoch_train_poi_loss)
        train_epochs_time_loss_list.append(epoch_train_time_loss)
        train_epochs_cat_loss_list.append(epoch_train_cat_loss)
        train_epochs_top1_acc_list.append(epoch_train_top1_acc)
        train_epochs_top5_acc_list.append(epoch_train_top5_acc)
        train_epochs_top10_acc_list.append(epoch_train_top10_acc)
        train_epochs_top20_acc_list.append(epoch_train_top20_acc)
        train_epochs_mAP20_list.append(epoch_train_mAP20)
        train_epochs_mrr_list.append(epoch_train_mrr)
        val_epochs_loss_list.append(epoch_val_loss)
        val_epochs_poi_loss_list.append(epoch_val_poi_loss)
        val_epochs_time_loss_list.append(epoch_val_time_loss)
        val_epochs_cat_loss_list.append(epoch_val_cat_loss)
        val_epochs_top1_acc_list.append(epoch_val_top1_acc)
        val_epochs_top5_acc_list.append(epoch_val_top5_acc)
        val_epochs_top10_acc_list.append(epoch_val_top10_acc)
        val_epochs_top20_acc_list.append(epoch_val_top20_acc)
        val_epochs_mAP20_list.append(epoch_val_mAP20)
        val_epochs_mrr_list.append(epoch_val_mrr)

        # Monitor loss and score
        monitor_loss = epoch_val_loss
        monitor_score = np.mean(epoch_val_top1_acc * 4 + epoch_val_top20_acc)

        # Learning rate schuduler
        lr_scheduler.step(monitor_loss)

        # Print epoch results
        logging.info(f"Epoch {epoch}/{args.epochs}\n"
                     f"train_loss:{epoch_train_loss:.4f}, "
                     f"train_poi_loss:{epoch_train_poi_loss:.4f}, "
                     f"train_time_loss:{epoch_train_time_loss:.4f}, "
                     f"train_cat_loss:{epoch_train_cat_loss:.4f}, "
                     f"train_top1_acc:{epoch_train_top1_acc:.4f}, "
                     f"train_top5_acc:{epoch_train_top5_acc:.4f}, "
                     f"train_top10_acc:{epoch_train_top10_acc:.4f}, "
                     f"train_top20_acc:{epoch_train_top20_acc:.4f}, "
                     f"train_mAP20:{epoch_train_mAP20:.4f}, "
                     f"train_mrr:{epoch_train_mrr:.4f}\n"
                     f"val_loss: {epoch_val_loss:.4f}, "
                     f"val_poi_loss: {epoch_val_poi_loss:.4f}, "
                     f"val_time_loss: {epoch_val_time_loss:.4f}, "
                     f"val_cat_loss: {epoch_val_cat_loss:.4f}, "
                     f"val_top1_acc:{epoch_val_top1_acc:.4f}, "
                     f"val_top5_acc:{epoch_val_top5_acc:.4f}, "
                     f"val_top10_acc:{epoch_val_top10_acc:.4f}, "
                     f"val_top20_acc:{epoch_val_top20_acc:.4f}, "
                     f"val_mAP20:{epoch_val_mAP20:.4f}, "
                     f"val_mrr:{epoch_val_mrr:.4f}")

        # Save train/val metrics for plotting purpose
        with open(os.path.join(args.save_dir, 'metrics-train.txt'), "w") as f:
            print(f'train_epochs_loss_list={[float(f"{each:.4f}") for each in train_epochs_loss_list]}', file=f)
            print(f'train_epochs_poi_loss_list={[float(f"{each:.4f}") for each in train_epochs_poi_loss_list]}', file=f)
            print(f'train_epochs_time_loss_list={[float(f"{each:.4f}") for each in train_epochs_time_loss_list]}',
                  file=f)
            print(f'train_epochs_cat_loss_list={[float(f"{each:.4f}") for each in train_epochs_cat_loss_list]}', file=f)
            print(f'train_epochs_top1_acc_list={[float(f"{each:.4f}") for each in train_epochs_top1_acc_list]}', file=f)
            print(f'train_epochs_top5_acc_list={[float(f"{each:.4f}") for each in train_epochs_top5_acc_list]}', file=f)
            print(f'train_epochs_top10_acc_list={[float(f"{each:.4f}") for each in train_epochs_top10_acc_list]}',
                  file=f)
            print(f'train_epochs_top20_acc_list={[float(f"{each:.4f}") for each in train_epochs_top20_acc_list]}',
                  file=f)
            print(f'train_epochs_mAP20_list={[float(f"{each:.4f}") for each in train_epochs_mAP20_list]}', file=f)
            print(f'train_epochs_mrr_list={[float(f"{each:.4f}") for each in train_epochs_mrr_list]}', file=f)
        with open(os.path.join(args.save_dir, 'metrics-val.txt'), "w") as f:
            print(f'val_epochs_loss_list={[float(f"{each:.4f}") for each in val_epochs_loss_list]}', file=f)
            print(f'val_epochs_poi_loss_list={[float(f"{each:.4f}") for each in val_epochs_poi_loss_list]}', file=f)
            print(f'val_epochs_time_loss_list={[float(f"{each:.4f}") for each in val_epochs_time_loss_list]}', file=f)
            print(f'val_epochs_cat_loss_list={[float(f"{each:.4f}") for each in val_epochs_cat_loss_list]}', file=f)
            print(f'val_epochs_top1_acc_list={[float(f"{each:.4f}") for each in val_epochs_top1_acc_list]}', file=f)
            print(f'val_epochs_top5_acc_list={[float(f"{each:.4f}") for each in val_epochs_top5_acc_list]}', file=f)
            print(f'val_epochs_top10_acc_list={[float(f"{each:.4f}") for each in val_epochs_top10_acc_list]}', file=f)
            print(f'val_epochs_top20_acc_list={[float(f"{each:.4f}") for each in val_epochs_top20_acc_list]}', file=f)
            print(f'val_epochs_mAP20_list={[float(f"{each:.4f}") for each in val_epochs_mAP20_list]}', file=f)
            print(f'val_epochs_mrr_list={[float(f"{each:.4f}") for each in val_epochs_mrr_list]}', file=f)

if __name__ == '__main__':
    args = parameter_parser()
    # The name of node features in NYC/graph_X.csv
    args.feature1 = 'checkin_cnt'
    args.feature2 = 'poi_catid'
    args.feature3 = 'latitude'
    args.feature4 = 'longitude'
    train(args)


#!/usr/bin/env python
import unittest
import lastwatch

class FilenameTest(unittest.TestCase):
    def parse_and_check(self, path, expected_result):
        parser = lastwatch.FilenameParser(path)
        self.assertEqual(parser.parse(), expected_result)

    def test_filename1(self):
        path = '/music/Moonspell/The Antidote/01. The Antidote.ogg'

        result = {
            'album': 'The Antidote',
            'title': 'The Antidote',
            'number': '01',
            'artist': 'Moonspell',
        }

        self.parse_and_check(path, result)

    def test_filename2(self):
        path = '/music/Linkin Park - A Thousand Suns (2010)/'
        path += '08. Linkin Park - Waiting for the End.ogg'

        result = {
            'album': 'A Thousand Suns',
            'title': 'Waiting for the End',
            'number': '08',
            'artist': 'Linkin Park',
        }

        self.parse_and_check(path, result)

if __name__ == '__main__':
    unittest.main()


from functools import partial

from src.models.scatter import FastScatterW1, Scatter, Scatter_Diffuse_Second, ScatterW1

class FastScatterTransform:
    def __init__(self, device="cpu", agg="moment", alpha=0, power=2, cheb_order=None):
        self.device = device
        self.agg = agg
        self.alpha = alpha
        self.power = power
        self.cheb_order = cheb_order
        assert power == 1 or power == 2
        if power == 2:
            self.scatter_fn = Scatter
        elif cheb_order is not None and cheb_order != "None":
            print("Fast scatterw1")
            self.scatter_fn = partial(FastScatterW1, cheb_order=cheb_order)
        else:
            self.scatter_fn = ScatterW1

    def __call__(self, data):
        data.x = self.scatter_fn(in_channels=data.x.shape[1], agg=self.agg, alpha=self.alpha)(data)
        return data

class FastScatterTransformSort:
    def __init__(self, device="cpu"):
        self.device = device

    def __call__(self, data):
        data.x = Scatter_Diffuse_Second(data.x.shape[1])(data)
        return data


from django.shortcuts import render, redirect
from .models import Image

def index(request):
	images = Image.objects.all()
	context = { 'images' : images}
	return render(request, 'index.html', context)

def uploader(request):
	newImage = Image()
	newImage.title = request.POST['title']
	newImage.about = request.POST['about']
	newImage.img = request.FILES['img']
	newImage.save()
	return redirect('home')


'''
    Given a string s, return the longest palindromic substring in s. Print multiple longest palindrome if any

Example 1:
Input: s = "babad"
Output: "bab"
Note: "aba" is also a valid answer.

Example 2:
Input: s = "cbbd"
Output: "bb"

Example 3:
Input: s = "geekforgeegeegeekeeg"
Output: "eegeegee"

Constraints:
    1 <= s.length <= 1000
    s consist of only digits and English letters.
'''

import sys
def solution(s):

    if len(s) == 1:
        return s

    def is_palindrome(ss):
        i = 0
        j = len(ss)-1
        while i < j:
            if ss[i] == ss[j]:
                i+=1
                j-=1
                continue
            else:
                return False
        return True

    _max_palindrome = ''
    all_results = list()
    map = dict()
    for i in range(len(s)):
        for j in range(i,len(s)):
            if is_palindrome(s[i:j+1]):
                if len(s[i:j+1]) >= len(_max_palindrome):
                    _max_palindrome = s[i:j+1]
                    all_results.append(_max_palindrome)

    map = { s: len(s) for s in all_results}
    _max_length = max(map.values())
    for key , value in map.items():
        if value == _max_length:
            print(key)

solution('abcaba')


import nltk
from sklearn.preprocessing import MultiLabelBinarizer
from nltk.tokenize import word_tokenize

# Более реалистичной является ситуация, когда есть данные, где каждое наблюдение
# содержит твит, и мы хотим преобразовать эти предложения в признаки отдельных
# частей речи (например, признак с 1, если присутствует собственное существительное,
# и 0 в противном случае):

if __name__ == '__main__':
    tweets = ["I am eating a burrito for breakfast",
              "Political science is an amazing field",
              "San Francisco is an awesome city"]

    tagged_tweets = []

    # Пометить каждое слово и каждый твит
    for tweet in tweets:
        tweet_tag = nltk.pos_tag(word_tokenize(tweet))
        tagged_tweets.append([tag for word, tag in tweet_tag])

    # Применить кодирование с одним активным состоянием, чтобы
    # конвертировать метки в признаки
    one_hot_multi = MultiLabelBinarizer()
    X = one_hot_multi.fit_transform(tagged_tweets)

    print(X, "\n")

    # Показать имена признаков
    print(one_hot_multi.classes_)


# from config_node import *
import paho.mqtt.client as mqtt
import config as cfg

from siisthing import SIISThing

from hardware.servo import ServoMotor

class SIISLock(SIISThing):
    def __init__(self, name: str = "lock_1"):
        SIISThing.__init__(self, name)
        self.last_state: str = ""

        self.device: ServoMotor = ServoMotor(cfg.pin)
        self.connect()

    def on_connect(self, client: mqtt.Client, userdata, flags, rc):
        # Extend the on_connect method of SIISThing
        SIISThing.on_connect(self, client, userdata, flags, rc)
        self.client.publish(self.state_topic, payload=self.last_state, qos=1, retain=True)

    def set_state(self, state: str) -> None:
        if state == "ON":
            self.device.on()
        elif state == "OFF":
            self.device.off()
        else:
            print("Invalid state")
        self.last_state = state

    def on_message(self, client: mqtt.Client, userdata, message: mqtt.MQTTMessage):
        # Check if this is a message that sets the lock
        if message.topic == self.set_topic or message.topic == self.scheduler_topic:
            # Read the target state
            lock: str = message.payload.decode("utf-8")
            print("Setting the lock to %s" % lock)
            self.last_state = lock
            # Echo it back
            response = lock
            client.publish(self.state_topic, response, qos=1, retain=True)
        else:
            # Pass the message down
            SIISThing.on_message(self, client, userdata, message)

    def start(self):
        self.connect()
        self.client.loop_forever()

# Start polling the files
if __name__ == "__main__":
    lock = SIISLock()
    lock.start()


import logging
import numpy as np
import glob
import os
import h5py
import warnings

from intan2kwik.core.h5 import tables
from intan2kwik.core import reading
#from intan2kwik.core.reading import rh_search_string

logger = logging.getLogger('intan2kwik.kwd')

def h5_unicode_hack(x):
    if isinstance(x, str):
        x = x.encode('utf8')
    return x

def list_chan_names(header, include_channels):
    ch_groups = [c + '_channels' for c in include_channels]

    c_names = []
    for chgrp in ch_groups:
        ch_list = header[chgrp]
        c_names += [ch['custom_channel_name'] for ch in ch_list]
    return c_names

def list_chan_bit_volts(header, include_channels):
    v_multipliers = [0.195, 50.354e-6]
    if header['eval_board_mode'] == 1:
        v_multipliers[1] = 152.59e-6

    ch_groups = [c + '_channels' for c in include_channels]
    b_volts = []

    for chgrp in ch_groups:
        ch_list = header[chgrp]
        if 'adc' in chgrp:
            b_value = v_multipliers[1]
        elif 'amplifier' in chgrp:
            b_value = v_multipliers[0]
        b_volts += [b_value for ch in ch_list]

    return np.array(b_volts, dtype=np.float)

def rhd_data_block(rhd_file, include_chans, times_too=False):
    read_block = reading.read_data(rhd_file)
    # identify channels, make numy array with all the block and a list of metadata for each
    block_data = np.vstack([read_block['{}_data'.format(ch_grp)] for ch_grp in include_chans])
    block_t = np.vstack([read_block['t_{}'.format(ch_grp)] for ch_grp in include_chans])

    if times_too:
        ret_value = (block_data, block_t)
    else:
        ret_value = block_data

    return ret_value

def rhd_rec_to_table(rhd_list, parent_group, chan_groups_wishlist):
    '''
    :param rhd_list: (list of strings) paths of files to include in this rec
    :param parent_group: (h5 object group) parent group for this rec
    :param chan_groups_wishlist: (flat ndarray/list) channel groups to get from the rhd files
    :return:
    '''
    # make the table
    # read the blocks and append them to the table
    # set the attributes of the table
    logger.info('Appending {} files to data table in {}'.format(len(rhd_list),
                                                                parent_group))

    total_samples_in = 0
    total_samples_in_dig =0

    last_t = 0
    s_f = parent_group.attrs.get('sample_rate')

    for i, rhd_file in enumerate(rhd_list):
        logger.info('Reading file {0}/{1}'.format(i+1, len(rhd_list)))
        read_block = reading.read_data(rhd_file)
        if i == 0:
            # filter include groups, warn if a group wasn't in the data and remove it from the list
            include_chan_groups = []
            for g in chan_groups_wishlist:
                if '{}_data'.format(g) in read_block.keys():
                    include_chan_groups.append(g)
                else:
                    logger.warn('Channel data group {} is not in the recordings'.format(g))
        # The main data (neural chans and dac chans; from include_chans)
        block_data = np.vstack([read_block['{}_data'.format(ch_grp)] for ch_grp in include_chan_groups])
        block_t = np.vstack([read_block['t_{}'.format(ch_grp)] for ch_grp in include_chan_groups])
        save_block = block_data.T.astype(np.int32) - 32768
        # assuming the first block in include_chans is amplifier
        save_t = block_t[0].reshape([-1, 1])
        name_t = 't_{}'.format(include_chan_groups[0])

        # The digital channels

        try:
            dig_in_data = read_block['board_dig_in_data'].T.astype(np.short)
            dig_in_t = read_block['t_dig'].reshape([-1, 1])
            has_digital_in = True
        except KeyError as k:
            logger.warn('No digital channels')
            has_digital_in = False

        # data comes in as uint16
        if i == 0:
            logger.info('Creating tables of neural/adc data')
            dset = tables.unlimited_rows_data(parent_group, 'data',
                                              save_block.astype(np.int16))
            tset = tables.unlimited_rows_data(parent_group, name_t, save_t)

            if has_digital_in:
                logger.info('Creating tables of digital data')
                dset_dig = tables.unlimited_rows_data(parent_group, 'dig_in',
                                                      dig_in_data.astype(np.short))
                tset_dig = tables.unlimited_rows_data(parent_group, 't_dig', dig_in_t)

        else:
            tables.append_rows(dset, save_block.astype(np.int16))
            tables.append_rows(tset, save_t)

            if has_digital_in:
                tables.append_rows(dset_dig, dig_in_data.astype(np.short))
                tables.append_rows(tset_dig, dig_in_t)

            # assert time continuity
            more_control_d_samples = (save_t[0] * s_f - total_samples_in)
            logger.info('Delta cum_samples/cum_t is {}'.format(more_control_d_samples))

            control_dt = save_t[0] - last_t
            control_delta_samples = int(np.round(control_dt * s_f))
            # logger.info('Delta samples between rhd files is {}'.format(control_delta_samples))
            if not control_delta_samples == 1:
                warn_msg = 'Skipped a beat i rhd files diff is {}s'.format(control_dt)
                warnings.warn(warn_msg, RuntimeWarning)

        last_t = save_t[-1]
        total_samples_in += save_block.shape[0]

        if has_digital_in:
            last_t_dig = dig_in_t[-1]
            total_samples_in_dig += last_t_dig.shape[0]

    # only atrribute for table is valid_samples
    dset.attrs.create('valid_samples', np.ones(save_block.shape[1]) * total_samples_in)

    if has_digital_in:
        dset_dig.attrs.create('valid_samples', np.ones(dig_in_data.shape[1]) * total_samples_in_dig)

def create_data_grp(rec_grp, intan_hdr, include_channels, rec):
    logger.debug('Creating data group for this rec {}'.format(rec))

    data_grp = rec_grp.create_group('{}'.format(rec))
    # append the metadata to this data group
    data_grp.attrs.create('bit_depth', 16)
    data_grp.attrs.create('sample_rate', intan_hdr['sample_rate'])
    data_grp.attrs.create('name', rec)
    data_grp.attrs.create('start_sample', 0)
    data_grp.attrs.create('start_time', 1)

    all_chan_names = list_chan_names(intan_hdr, include_channels)
    all_chan_names_uni = [h5_unicode_hack(x) for x in all_chan_names]
    all_bit_volts = list_chan_bit_volts(intan_hdr, include_channels)
    n_chan = len(all_chan_names)
    all_rates = np.ones(n_chan) * intan_hdr['sample_rate']

    # create application data
    app_data_grp = data_grp.create_group('application_data')
    app_data_grp.attrs.create('is_multiSampleRate_data', 0)
    app_data_grp.attrs.create('channels_sample_rate', all_rates)
    app_data_grp.attrs.create('channel_bit_volts', all_bit_volts)
    app_data_grp.attrs.create('channel_names', all_chan_names_uni)

    return data_grp

def which_board(folder):
    rhx_files = glob.glob(os.path.join(folder, '*.rh[d,s]'))
    if len(rhx_files) < 1:
        raise RuntimeError('No .rh[d, x] found in folder {}'.format(folder))
    return os.path.split(rhx_files[0])[-1].split('.')[-1]

def neural_ports_meta(folder):
    """
    :param folder: (string) folder with .rh? files
    :return: (tuple) of dictionaries:
        chans: {port: list of channels}
        count: {port: channel count
    """
    board = which_board(folder)
    all_rhx_files = glob.glob(os.path.join(folder, '*.{}'.format(board)))
    all_rhx_files.sort()
    intan_read = reading.read_data(all_rhx_files[0])
    chans, count = used_neural_channels(intan_read)
    return chans, count

def intan_to_kwd(folder, dest_file_path, rec=0, include_channels=None, board='auto'):
    """
    :param folder: (string) folder where the .rh? files are
    :param dest_file_path: (string) dest of the kw? files
    :param rec: (int)
    :param include_channels: (flat ndarray/list)
    :param board: (str) 'rhd' or 'rhs' for rhd2000 or rhs2000, 'auto' for detecting from extension
    :return:
    """
    # make the .kwd file
    # make the /recording/0 group
    # dump header to application data
    # run rhd_rec_to table to create the table in the group
    if include_channels is None:
        include_channels = ['amplifier', 'board_adc']
    logger.info('reading intan chans data across all of rec {0}'.format(folder))

    if board=='auto':
        board = which_board(folder)
    all_rhx_files = glob.glob(os.path.join(folder, '*.{}'.format(board)))
    all_rhx_files.sort()
    logger.info('Found {} .{} files to process'.format(len(all_rhx_files), board))

    # attributes from the header
    first_header = reading.read_intan_header(all_rhx_files[0])
    # logger.info('First header {}'.format(first_header))
    v_multipliers = [0.195, 50.354e-6]
    if first_header['eval_board_mode'] == 1:
        v_multipliers[1] = 152.59e-6

    # chan_names =

    # list all the chunk files:
    logger.debug('dest file: {}'.format(dest_file_path))
    with h5py.File(dest_file_path, 'a') as kwd_file:
        logger.debug('Creating the /recordings/{}'.format(rec))
        rec_grp = kwd_file.require_group('recordings')
        data_grp = create_data_grp(rec_grp, first_header, include_channels, rec)
        # create application data with all metadata
        logger.debug('Creating data table and going throug the recs')
        rhd_rec_to_table(all_rhx_files, data_grp, include_channels)

    return first_header


#baseball
import random

comnum = random.sample(range(1,10),3)
print(comnum)
strike = 0
count = 0

while strike != 3:
    strike = 0
    ball = 0
    guessnum = list(input("please choose three numbers: "))
    print(guessnum)
    for i in guessnum:
        for j in comnum:
            if int(i) == j:
                if guessnum.index(i) == comnum.index(j):
                    strike += 1
                else:
                    ball+= 1
    count+= 1
    print("you have {} strike(s), {} ball(s), and tried {} times,".format(strike, ball, count))


import re
from GameObjects import Location, Collectable, Pocket
from GameEngine import screen_clear, timed_print, close_game, opensite, im_show, copyFile
from random import choice
import networkx as nx
from matplotlib.pyplot import show

"""Build all the locations"""

Main_Gate = Location(
    "Main Gate",
    "A Nice Big Arch!\n"
    "Students are coming back from their Friday night party at 12am...")
East_Gate = Location(
    "East Gate",
    "A train passes by...")
T_Point = Location(
    "T point",
    "Hunh, I can see something weird in this mirror, Oh that's just me...")
Admin = Location(
    "Admin",
    "You enter the lovely automatic sliding doors, a cool blast of AC...\n"
    "You see the IISER-M model and remember that 4 more hostels need to be built.\n"
    "But there are no funds...")
CAF = Location(
    "CAF",
    "Very sensitive machines which may be disturbed by students "
    "eating in the cafeteria\n"
    "You hear someone shouting at someone else in the distance, "
    "but do not record it, for it would be a violation of privacy")
Library = Location("Library", "So quiet...")
Computer_Center = Location("Computer Center",
                           "A few computers have not been turned off")
Health_Center = Location("Health Center", "A very dark and scary place...")

Outside_LHC = Location("Outside LHC", "The LHC may be locked at this time...")
LHC = Location("LHC", "Ah the good old LHC...")
LH1 = Location('LH1', "There's a BIG flowchart on the board!")
LH2 = Location('LH2', "There are so many computers.")
LH3 = Location(
    'LH3',
    "You can still hear the debates that happened here a few hours ago.")
LH4 = Location('LH4', "Nothing interesting")
LH5 = Location('LH5', "Someone has left their bottle outside.")
LH6 = Location(
    'LH6',
    'Someone forgot their umbrella. Some differential equations are written on the board')
LH7 = Location(
    'LH7',
    "The watchman shoos you away. The stage is too precious to even repair...")

Rotunda = Location(
    'Rotunda',
    'You see a gang of MS19 making a lot of noise...')
H5 = Location(
    'H5',
    'Caretaker\'s Office is locked. Why is it always locked? Hmm...\n'
    "It's dark outside, you may need a torch")
H5_SR = Location('H5 SR', "Back to the start...")
H6 = Location('H6', 'Hos... The rest of the banner is torn...\n'
              "It's dark outside, you may need a torch")
H6_SR = Location(
    'H6 SR',
    'You see a couple snogging. They see you. It\'s so awkward')
H7 = Location('H7', 'Someone is taking their clothes to wash...\n'
              "It's dark outside, you may need a torch")
H7_SR = Location(
    'H7 SR',
    'Those Dance Club people are still dancing in the Dance Room')
H8 = Location(
    'H8',
    'Some MS18 kids are loitering around...\n'
    "It's dark outside, you may need a torch")
H8_SR = Location(
    'H8 SR',
    'The ACs are off, it\'s so hot...\nOh wait, they are on... '
    'Why is it so hot then? Oh the ACs aren\'t working')

Stadium = Location('Stadium', 'A lapwing shouts in the distance...')
BB_Court = Location('BB Court', 'Why the fuck do lapwings make so much noise?')
VH = Location(
    'VH',
    'Dollar building. The terrace may be locked')
VH_Terrace = Location(
    'VH Terrace',
    'Haaah! Nice and windy.\n'
    'You see a couple of friends talking about the struggles of life...')
Shopping_Complex = Location(
    'Shopping Complex',
    'It\'s so empty and quiet. Aunty must have gone home...')

AB1 = Location(
    'AB1',
    'You buy a nice coffee for 10rs, coz neither of you have change')
AB1_1F = Location('AB1 1F', 'People playing Badminton...')
AB1_2F = Location('AB1 2F', 'Eerily dark and quiet')
AB1_3F = Location(
    'AB1 3F',
    'One lone MS13 is studying there. Why haven\'t they graduated yet?')
AB1_4F = Location('AB1 4F', 'Eerily dark and quiet')
EBL_Lab = Location(
    'EBL',
    'So many dead flies in flasks.\nIt\'s like a horror movie')
AB1_5F = Location(
    'AB1 5F',
    'Eerily dark and quiet, but there\'s some news article about '
    'Plastic Firecrackers...')
AB2 = Location(
    'AB2',
    'ERROR! The app devs didn\'t care enough to make 2 AB\'s.\n'
    'There\'s nothing here. Go away, shoo.')
Gazebo = Location(
    'Gazebo',
    'There\'s some MS17 party going on there. Oh god so much cake wasted.\n'
    'Talking about wasting stuff, they\'re fixing the covers...')
Behind_AB = Location(
    'Behind AB',
    'Aww! 3 cute puppies run up to you.\n'
    r"""
           A_         __         __
      (___/ "`;  [___()"`;  \___()"`;
      /,    /`   /,    /`   /,    /`
      \\"--\\   //"--//     \\"--\\

    """)
Animal_Center = Location(
    'Animal Center',
    'The everlasting humm of the machines...\n' +
    'What are these machines? What happens here? SO MANY QUESTIONS!')
Room69 = Location(
    'Room No 69',
    'Why is it numbered 69 if it is the only room here?\n' +
    'I have found this place unlocked for the first time. Should explore it a bit.')

"""Connect the Locations to make a traversable map"""
Main_Gate.append_locations([Behind_AB, Admin, CAF])
East_Gate.append_locations([Animal_Center, Behind_AB])
T_Point.append_locations([CAF, Outside_LHC, Admin, Health_Center])
Health_Center.append_locations([T_Point, BB_Court])
Admin.append_locations([Main_Gate, CAF, T_Point])
CAF.append_locations([Admin, Main_Gate, T_Point])
Library.append_locations([Outside_LHC, Gazebo, Computer_Center])
Computer_Center.append_locations([Library])

Outside_LHC.append_locations([T_Point, Library, Rotunda, Animal_Center])
LHC.append_locations([LH1, LH2, LH3, LH4, LH5, LH6, LH7, Outside_LHC])
LH1.append_locations([LHC])
LH2.append_locations([LHC])
LH3.append_locations([LHC])
LH4.append_locations([LHC])
LH5.append_locations([LHC])
LH6.append_locations([LHC])
LH7.append_locations([LHC])

Rotunda.append_locations([Outside_LHC, H5, H6, H7, H8, Stadium])
H5.append_locations([H5_SR])
H6.append_locations([Rotunda, H6_SR])
H7.append_locations([Rotunda, H7_SR])
H8.append_locations([Rotunda, H8_SR])
H5_SR.append_locations([H5])
H6_SR.append_locations([H6])
H7_SR.append_locations([H7])
H8_SR.append_locations([H8])

Stadium.append_locations([Rotunda, BB_Court, Shopping_Complex])
BB_Court.append_locations([Health_Center, Stadium])
Shopping_Complex.append_locations([Stadium, VH])
VH.append_locations([Shopping_Complex])
VH_Terrace.append_locations([VH])

AB1.append_locations([Gazebo, AB2, Behind_AB, AB1_1F])
AB1_1F.append_locations([AB1, AB1_2F])
AB1_2F.append_locations([AB1_1F, AB1_3F])
AB1_3F.append_locations([AB1_2F, AB1_4F])
AB1_4F.append_locations([AB1_3F, AB1_5F, EBL_Lab])
EBL_Lab.append_locations([AB1_4F])
AB1_5F.append_locations([AB1_4F])
AB2.append_locations([Gazebo, Behind_AB, AB1])
Gazebo.append_locations([Library, AB1, AB2, Behind_AB])
Behind_AB.append_locations([Main_Gate, East_Gate, AB1, AB2, Gazebo])
Animal_Center.append_locations([Outside_LHC, East_Gate, ])
Room69.append_locations([Animal_Center])

"""Build Actual clues"""

def clue8Action():
    print("Is half the year over already?!")
    im_show("calender.jpg")
    ans = input("Enter the correct year: ")
    if ans == "2036":
        return True
    return False

clue8 = Collectable(
    "unusually clean sheet of paper",
    "taped to the pillar, above the wash basin",
    "Seems kinda out of place between all these Melody chocolate posters from 2007",
    hidden=True,
    action=clue8Action,
    nextnotes=None,
    onComplete="Congratulations! You have successfully reached the end of this treasure hunt\n"
               "You can go back to dying of boredom now. Bye.",  # something else
    onFail="Tsk.Tsk. So close yet so far. Try again.")

def clue7Action():
    print("There is a file. Save it in your PC.")
    a = input("Would you like to proceed?[Y/N]")
    if(a == "Y"):
        _ = copyFile("image.jpg")
    b = input("Enter the number if you have found it: ")
    if(b == "690"):
        Pocket.append("Note - Truly, 69 is an important number")
        Animal_Center.append_locations([Room69])
        return True
    return False

clue7 = Collectable(
    "Found a USB drive",
    "Hanging from the ceiling by means of a thread!\n"
    "Either the developer was too sleepy to think of anything better "
    "or this is a weird place.",
    "Should plug it in my pc to check what it is",
    hidden=True,
    action=clue7Action,
    nextnotes=[clue8],
    onComplete="That was some smart work! Truly 69 is an important number!",
    onFail="What's the first thing you do when you don't know what a file is?")

def clue6Action():
    print("""
            The quest is to find the direction of your life...\n
            Two roads... first towards the accelerator and second towards jordon\n
            And let your octaves fill the keanu reeves movie\n
            Apply some contrast of free will and maybe lose the negativity ...\n
            You should be good to go now\n
            """)

    ans = input("Enter your password: ")
    if ans == "22":
        Pocket.append(
            "Clue - H-ate me or love me, you need to Study with me")
        Pocket.append("Key to VH Terrace")
        VH.append_locations([VH_Terrace])
        return True
    return False

clue6 = Collectable(
    "wooden box",
    "between H6 & H8",
    "It has a paper stuck on it...\n"
    "Man these treasure hunt people need to stop dipping paper in coffee to make it seem old..\n"
    "There's also a two digit number lock!\n",
    hidden=True,
    action=clue6Action,
    nextnotes=[clue7],
    onComplete="The box clicked open! Let's see what's inside."
    "A key to VH Terrace and a Note-\n\n"
    """
    +=================================+
    |                                 |
    |    H-ate me or love me,         |
    |    you need to Study with me    |
    |                                 |
    +=================================+

    """
    "They seriously need to drop the coffee trick",
    onFail="The box didn't open.\n"
    "Hehe seems like you're directionally challenged :)\n"
    "Look at it from a different perspective maybe")

def clue5Action():
    i = 0
    question_prompt = ["Cell's leader says 'Carry on!' suspiciously...",
                       "Is the Spanish Crown alive?", "Hiding in a plaza having a handicap",
                       "Chris, who left the hotel, heads the Public Relations Forum at Caribou Biosciences"]
    answers = ["KARYON", "CORONAVIRUS", "ZAHAVI", "CRISPR"]
    while i < len(question_prompt):
        print(question_prompt[i])
        ans = input("Enter your answer: ")
        if ans.upper() == answers[i]:
            i += 1
        else:
            return False
    Pocket.append("Bio Pop quiz answer sheet - Should show friends in Rotunda")
    return True

clue5 = Collectable(
    "Bio Pop Quiz",
    "tucked inside the bottle...?",
    "Why would someone do that?\n"
    "A small sheet with a few questions (that will definitely be counted for your consolidated grades :-))",
    True,
    clue5Action,
    nextnotes=[clue6],
    onComplete="Yay!! You answered everything correctly!\n"
    "*happy biologist noises*\n"
    "Maybe you should go tell your friends in Rotunda",
    onFail="OOPS. Looks like you'll need to start all over again")

def clue4Action():
    dst = copyFile("mutations.pdf")
    if dst:
        Pocket.append("mutations.pdf in " + dst)
        return True
    return False

clue4 = Collectable(
    "BIO211 experiment draft",
    "in the dustbin",
    "Maybe this holds some mysteries to the next location?",
    hidden=True,
    action=clue4Action,
    nextnotes=[clue5],
    onComplete="Read the copied file, the secret lies within",
    onFail="File copy error, retry")

def clue3Action():
    print("PENALTY at 45 degrees")
    im_show("cube.jpg")
    print("Message: C.Amjg md eclcrgaq")  # change this please
    a = input("What does that mean? ")
    if a == "E.Coli of genetics":  # change this please
        Pocket.append("Clue - E.Coli of genetics")
        return True
    return False

clue3 = Collectable(  # put the clue in lhc rn, change it to anywhere
    "Poster",
    "on the softboard",
    "It's the poster of the Cubing competition held in college some time back",
    hidden=True,
    action=clue3Action,
    nextnotes=[clue4],
    onComplete="Hmm? Is that important?",
    onFail="Just search the net man")

def clue2Action():
    opensite("clue2.html")
    a = input("What's the password? ")
    if a == "acceptme":
        Pocket.append("Clue - Large Hadron Collider")
        return True
    return False

clue2 = Collectable(
    "Quantum Computer",
    "in the BIG BIG hall",
    "Another webpage on the way, on a quantum computer!\n"
    "They used a quantum computer for opening a webpage?\n"
    "Perfect utilization of a non-existent resource",
    hidden=True,
    action=clue2Action,
    nextnotes=[clue3],  # made clue 3 next
    onComplete="Yeah, this was a joke. Now head to the Large Hadron Collider..?",
    onFail="Noooo! This one is dumb too.")

def clue1Action():
    opensite('clue1.html')
    a = input("What's the key? ")
    if a == "TUr!ng":
        Pocket.append("Clue - Costly Machines, easy clues...")
        return True
    return False

clue1 = Collectable(
    "website",
    "running on one of the computers",
    "This website seems like a clue...",
    hidden=False,
    action=clue1Action,
    nextnotes=[clue2],
    onComplete="Costly Machines, easy clues...",
    onFail="No, it's the first question. It's dumber than you think!")

def lhckey_action():
    if input("Pick it?[Y N] ") == "Y":
        Outside_LHC.append_locations([LHC])
        Pocket.append("LHC Key")
        return True
    return False

lhckey = Collectable(
    "LHC Key",
    "hanging by the security",
    "It looks like a normal key. Maybe you can finally see what's inside LHC?",
    hidden=False,
    action=lhckey_action,
    onComplete="Now get out of here before the security guard comes back",
    onFail="There was a nice art exhibition in LHC I heard?")

""" Build FLUFF """
hostel_cctv = Collectable(
    "Notice",
    "on the Soft Board",
    """
    ========================================
    |                                      |
    |   YOU ARE UNDER CCTV SURVIELLENCE!   |
    |                                      |
    ========================================
    """)

def fake_clue_action():
    print("The password is 5")  # some message
    return input("What's the password? ") == "5"  # some prompt

fake_clue = Collectable(
    "Fake Clue",
    "under some table",  # some specific place
    "This is a test clue",  # message of a clue
    False,  # Not hidden
    fake_clue_action,  # action object
    [hostel_cctv],  # next note
    "Yay! Look for a treat in the Hostels",
    "Noo, you dumb?")

def biswas_car_action():
    global here
    next_loc = choice(
        [Behind_AB, East_Gate, Animal_Center, LHC, BB_Court, Stadium, Shopping_Complex])
    print("You followed the car and ended up in " + next_loc.name + "!")
    here = next_loc
    return True

biswas_car = Collectable(
    "red sedan",
    "coming from Animal Facility road",
    "It's approaching almost quasistatically...\n"
    "Why would someone be driving around at midnight?! This is pretty shady."
    "Probably not important. It's not like they're going to run over anybody.\nRight?\n\nRIGHT?",
    hidden=False,
    action=biswas_car_action,
    onComplete="This had better be worth-- \nWAIT! THE DRIVER IS ASLEEP?! *facepalm*",
    onFail="Unreachable. Report to devs")
biswas_car.nextnotes = [biswas_car]  # Work around to skip the try again

stadium_test = Collectable(
    name="Rocket Pad",
    spot="in the stadium",
    desc="Lets fly the rocket.\n"
    "It went so high that you can't even see it!!!\n\n\n"
    "Oh no! It's falling back! RUNNN!",
    hidden=True)

def printer_3d_action():
    print("Printer shows a message 'Please input the blueprint'")
    print("You should wait for the MS19 kids to finish...\n\nOr you could just print the blueprints now...")
    if input("Do you want to print the blueprint? [Y N] ") == "Y":
        Pocket.append("3D Printed Rocket")
        return True
    return False

printer_3d = Collectable(
    name="3D printer",
    spot="outside physics lab",
    desc="You go near the 3D printer and see a bunch of MS19 students sleeping around it.\n\n"
    "They are probably tired from waiting their turn at the printer.",
    hidden=True,
    action=printer_3d_action,
    nextnotes=[stadium_test],
    onComplete="The Printer started making weird noises.\n\n\n"
    "You got a 3d printed rocket!\n"
    "You should try flying it in some open place with a lot of grass.\n"
    "Better get out of here before the MS19 kids wake up.",
    onFail="Lets get outta here")

def samrat_inno_action():
    im_show("rocket.jpg")
    print("Maybe I can 3D print this and see if it works")
    if input("Do you want to steal - no, borrow - this? [Y N] ") == "Y":
        Pocket.append("Rocket.jpg")
        return True
    return False

samrat_inno = Collectable(
    name="a Blueprint",
    spot="on the table",
    desc="Innovation #42054....\n\n",
    hidden=False,
    action=samrat_inno_action,
    nextnotes=[printer_3d],
    onComplete="We will have to find a 3D printer...",
    onFail="Nope, it's someone else's property")

def art_competition_action():
    print("- Shivali")
    im_show("art1.jpg")
    print("- Saloni Sahu")
    im_show("art2.jpg")
    print("- Shivali")
    im_show("art3.jpg")
    return False

art_competition = Collectable(
    name="multiple art pieces",
    spot="on soft boards",
    desc="There was an art competition held a few days ago.\n"
    "These are some of the beautiful paintings some people have made!",
    hidden=False,
    action=art_competition_action,
    onComplete="Unreachable. Report to devs",
    onFail="Damn! Those were so pretty! I wish I could make things like that")

def room69_action():
    im_show('69.jpg')
    Pocket.append("Clue - Go to Shopping Complex")
    print("Let's take a picture of this.")
    a = input("Do you want to continue[Y/N]")
    if(a == "Y"):
        _ = copyFile("69.jpg")
    return False

room69poster = Collectable(
    name='lot of cult articles',
    spot='everywhere, even on the ceiling',
    desc='There is a poster on the wall. Should read it.',
    hidden=False,
    action=room69_action,
    onFail="That is some weird thing...\n"
    "More than what meets the eye...\n"
    "Could it be a clue?\n\n"
    "Eh screw it... I'm hungry now... I should check if Shopping Complex is open",
    onComplete="Unreachable. Report to devs.")

def torch_action():
    print("You turn it on and it's very bright!\n")
    if input("Do you want to take it [Y N]? ") == "Y":
        H5.append_locations([Rotunda])
        Pocket.append("A red Torch")
        return True
    return False

torch = Collectable(
    "Torch",
    "on the table",
    "It's a nice red color.",
    action=torch_action,
    hidden=False,
    onComplete="You can now see in the dark and outside!",
    onFail="You should really take it...")

def mmystery_board_action():
    im_show("board.jpg")
    return False

mmystery_board = Collectable(
    name="flowchart",
    spot="on the board",
    desc="These Murder Mystery people are out of their mind.",
    hidden=False,
    action=mmystery_board_action,
    onComplete="Unreachable. Report to devs",
    onFail=""
)

def mmystery_poster_action():
    im_show("mmystery.jpg")
    return False

mmystery_poster = Collectable(
    name="poster for murder mystery",
    spot="on the table",
    desc="There's a sign under the poster. '- Bhavneet Singh'",
    hidden=False,
    action=mmystery_poster_action,
    onComplete="Unreachable. Report to devs",
    onFail="This guy can make nice posters!"
)

def vht_action():
    im_show("trails.jpg")
    im_show("vht.jpg")

vht = Collectable(
    "photo album of the night sky",
    "on the ground",
    "Such lovely images!\n"
    "These were taken by Kaustubh Verma and Amit Yogi!",
    hidden=False,
    action=vht_action,
    onComplete="Unreachable. Report to devs",
    onFail="These guys can take nice pictures!"
)

"""Add Notes to Locations"""

Computer_Center.append_collectable([clue1])
CAF.append_collectable([clue2])
EBL_Lab.append_collectable([clue4])
LHC.append_collectable([art_competition, clue3])
LH5.append_collectable([clue5])
Rotunda.append_collectable([clue6])
H8_SR.append_collectable([clue7])
Shopping_Complex.append_collectable([clue8])

VH_Terrace.append_collectable([vht])

H5.append_collectable([hostel_cctv])
H6.append_collectable([hostel_cctv])
H7.append_collectable([hostel_cctv])
H8.append_collectable([hostel_cctv])

AB1_3F.append_collectable([printer_3d])
AB1_5F.append_collectable([samrat_inno])
Stadium.append_collectable([stadium_test])

Admin.append_collectable([lhckey])
H5_SR.append_collectable([torch])
T_Point.append_collectable([biswas_car])
Room69.append_collectable([room69poster])
LH1.append_collectable([mmystery_board, mmystery_poster])

def makemap(root):
    g: nx.DiGraph = nx.DiGraph()

    def makenodes(root, l):
        if root not in l:
            l.append(root)
            for loc2 in root.locations:
                g.add_edge(root.name, loc2.name)
                makenodes(loc2, l)
    makenodes(root, [])
    nx.draw_kamada_kawai(
        g,
        with_labels=True,
        node_size=750,
        font_size=7,
        width=2,
        edge_color='#6B6B6B')
    show()

hello_banner = '''Welcome to the Virtual Treasure Hunt!
The Treasure Hunt will go on until you finish it.
Though this is an individual game, it is best served with a side of video call with your friends.
But remember that they can hide things from you and win themselves. :P

You can use the internet, but it is not necessary to have a continuous connectivity.
Please remember that if you exit from this app (by pressing the X or by
shutting down your computer), you will HAVE to restart.
You can minimize this safely, however.
We suggest that you also keep a notepad window open to note down information.

Some clues may seem too technical, but they're generally straight forward.
Just search the internet if you get stuck!

Go forth and do your best!'''

here = H5_SR

def systest():
    screen_clear()
    print("This is a system test. Report to devs if any errors are shown.")
    if input("Enter to Start systest") == "skip":
        return
    print()
    print("You will be shown a map. This test will be a success if you can see the map. Close the map to continue...")
    input("Enter to continue")
    makemap(Rotunda)
    print("makemap success")
    print()

    print("You will be shown an image. This test will be a success if you can see the image. Close the image to continue...")
    input("Enter to continue")
    im_show("owl.jpg")
    print("im_show success")
    print()

    print("You will be asked to save a file. Success if file is saved correctly")
    input("Enter to continue")
    if copyFile("test.txt"):
        print("copyFile success")
    else:
        print("copyFile fail")
    print()

    print("You will be redirected to a test site")
    input("Enter to continue")
    opensite("test.html")
    print("If website shown correctly, opensite success")
    print()

    input("This concludes systest. Enter to continue")

def mainloop():
    global here
    while True:
        screen_clear()
        msg = ""
        print(here)
        print("\n\n")
        print("If you want to search around, type search or ls")
        print("If you want to view map, type map")
        print("If you want to view pocket, type pocket")
        print("If you want to go somewhere else, "
              "type goto _location_ or cd _location_")
        print()

        inp = input("Your choice: ")
        match_obj = re.search(r"^(goto|cd) .*", inp)
        if re.search(r"^(search|ls)$", inp):
            msg = here.search()
        elif match_obj:
            new_loc = inp[len(match_obj.group(1)) + 1:]
            msg, here = here.goto(new_loc)
        elif re.search(r"^map$", inp):
            print("Close map to continue")
            makemap(here)
        elif re.search(r"^pocket$", inp):
            Pocket.show()
        elif re.search(r"^exit$", inp):
            if input("Type exit again to exit, anything else to go back: ") == "exit":
                close_game()
        else:
            msg = "Invalid Input"
        if (msg):
            timed_print(msg, 3)

screen_clear()
print(hello_banner)
input("Press Enter to start systest")
systest()
screen_clear()
print("""    How to play -
        Use `ls` or `search` to look around for clues and other helpful items
        Use `goto location name` or `cd location name` to move to connecting locations
        Use `map` to open up the map. Note that as you collect more items, new locations open up!
        Use `pocket` to view a list of items collected. This is very helpful if you get stuck.

    Go ahead and enjoy! Good Luck!
""")
input("Press Enter to start the Game!")
mainloop()


import doctest
import os
import sys
import unittest

class TestLoader(unittest.TestLoader):

    # unittest.main does not allow multiple "--start-directory"
    # options, but we can make it scan multiple separate directories
    # by overriding discover().  This allows us to have a "tests"
    # directory that's separate from "switch_model".
    #
    # We don't want to scan for *.py files in the parent directory in
    # case any of those are throwaway scripts that have unexpected
    # effects when imported.
    def discover(self, start_dir, pattern, top_level_dir):
        test_suite = unittest.TestSuite()
        for subdir in ("switch_model", "tests"):
            test_suite.addTests(
                super(TestLoader, self).discover(
                    os.path.join(top_level_dir, subdir), pattern, top_level_dir
                )
            )
        return test_suite

    # The unittest module does not have built-in support for finding
    # doctests.  In order to run the doctests, we need a custom
    # TestLoader that overrides loadTestsFromModule().
    def loadTestsFromModule(self, module, **kwargs):
        test_suite = super(TestLoader, self).loadTestsFromModule(module, **kwargs)

        docstring = module.__doc__
        if not docstring:
            # Work around a misfeature whereby doctest complains if a
            # module contains no docstrings.
            module.__doc__ = "Placeholder docstring"
        test_suite.addTests(doctest.DocTestSuite(module))
        if not docstring:
            # Restore the original, in case this matters.
            module.__doc__ = docstring
        return test_suite

def main():
    script_dir = os.path.join(os.getcwd(), os.path.dirname(__file__))
    # print('old argv: {}'.format(sys.argv))
    argv = [
        sys.argv[0],
        "discover",
        "--top-level-directory",
        script_dir,
        "--pattern",
        "*.py",
    ] + sys.argv[1:]
    unittest.TestProgram(testLoader=TestLoader(), argv=argv, module=None)

if __name__ == "__main__":
    main()


from oslo.config import cfg

import nova.openstack.common.rpc.proxy

rpcapi_opts = [
    cfg.StrOpt('console_topic',
               default='console',
               help='the topic console proxy nodes listen on'),
]

CONF = cfg.CONF
CONF.register_opts(rpcapi_opts)

class ConsoleAPI(nova.openstack.common.rpc.proxy.RpcProxy):
    '''Client side of the console rpc API.

    API version history:

        1.0 - Initial version.
        1.1 - Added get_backdoor_port()
    '''

    #
    # NOTE(russellb): This is the default minimum version that the server
    # (manager) side must implement unless otherwise specified using a version
    # argument to self.call()/cast()/etc. here.  It should be left as X.0 where
    # X is the current major API version (1.0, 2.0, ...).  For more information
    # about rpc API versioning, see the docs in
    # openstack/common/rpc/dispatcher.py.
    #
    BASE_RPC_API_VERSION = '1.0'

    def __init__(self, topic=None):
        topic = topic if topic else CONF.console_topic
        super(ConsoleAPI, self).__init__(
                topic=topic,
                default_version=self.BASE_RPC_API_VERSION)

    def add_console(self, ctxt, instance_id):
        self.cast(ctxt, self.make_msg('add_console', instance_id=instance_id))

    def remove_console(self, ctxt, console_id):
        self.cast(ctxt, self.make_msg('remove_console', console_id=console_id))

    def get_backdoor_port(self, ctxt, host):
        return self.call(ctxt, self.make_msg('get_backdoor_port'),
                         version='1.1')


# Mostly Stolen From CS230.
# `https://github.com/cs230-stanford/cs230-code-examples/blob/master/pytorch/vision/utils.py`
# Utility Functions for logging and Saving model.
import os
import torch
import numpy as np
'''
 ##################
 # Example Usage  #
 ##################
state = {
    'epoch': epoch + 1,
    'state_dict': model.state_dict(),
    'optim_dict' : optimizer.state_dict()
}
utils.save_checkpoint(
    state,
    checkpoint=model_dir
)
'''

def save_checkpoint(state, checkpoint='../Data/FullData', name='best_model.tar'):
    """
    Saves the best model and training parameters at
    `checkpoint + 'best_model.tar'``.
    Args:
        state:
            (dict) contains model's state_dict,
            may contain other keys such as epoch, optimizer state_dict
        checkpoint:
            (string) folder where parameters are to be saved
    """
    filepath = os.path.join(checkpoint, name)
    if not os.path.exists(checkpoint):
        print("Checkpoint Directory does not exist! Making directory {}".format(checkpoint))
        os.mkdir(checkpoint)
    torch.save(state, filepath)

def load_checkpoint(checkpoint, model, optimizer=None):
    """
    Loads model parameters (state_dict) from file_path.
    If optimizer is provided, loads state_dict of
    optimizer assuming it is present in checkpoint.
    Args:
        checkpoint:
            (string) filename which needs to be loaded
        model:
            (torch.nn.Module) model for which the parameters are loaded
        optimizer:
            (torch.optim) optional: resume optimizer from checkpoint
    """
    if not os.path.exists(checkpoint):
        raise("File doesn't exist {}".format(checkpoint))
    checkpoint = torch.load(
        checkpoint,
        map_location=lambda storage, loc: storage
    )
    model.load_state_dict(checkpoint['state_dict'])

    if optimizer:
        optimizer.load_state_dict(checkpoint['optim_dict'])

class SubtructMeanImage(object):
    '''
        subtract the mean image from the image
    '''
    def __init__(self, path_to_mean_image):
        self.mean_image = torch.from_numpy(
            np.load(path_to_mean_image)
        ).permute(2, 0, 1).float()
    def __call__(self, image_tensor):
        image =  image_tensor - self.mean_image
        return image

def show_image(pytorch_tensor):
    import torchvision.transforms.functional as F
    a = F.to_pil_image(pytorch_tensor)
    return a


from brownie import (
    accounts,
    exceptions,
    UnstoppableLender,
    DamnValuableToken,
    ReceiverUnstoppable,
)
from web3 import Web3
import pytest

# Pool has 1M * 10**18 tokens
TOKENS_IN_POOL = Web3.toWei("1000000", "ether")
INITIAL_ATTACKER_BALANCE = Web3.toWei("100", "ether")

def before():
    # Setup scenario
    deployer = accounts[0]
    attacker = accounts[1]
    randomUser = accounts[2]
    dvt_token = DamnValuableToken.deploy({"from": deployer})
    lender = UnstoppableLender.deploy(dvt_token.address, {"from": deployer})

    # approve lender for deployer
    dvt_token.approve(lender.address, TOKENS_IN_POOL, {"from": deployer})
    # setup lender with 1,000,000 DVT
    lender.depositTokens(TOKENS_IN_POOL, {"from": deployer})
    # setup attacker with 100 DVT
    dvt_token.transfer(attacker.address, INITIAL_ATTACKER_BALANCE, {"from": deployer})

    # assert balances and that users can execute flash loans.
    assert dvt_token.balanceOf(lender.address) == TOKENS_IN_POOL
    assert dvt_token.balanceOf(attacker.address) == INITIAL_ATTACKER_BALANCE
    receiver = ReceiverUnstoppable.deploy(lender.address, {"from": randomUser})
    receiver.executeFlashLoan(10, {"from": randomUser})

def run_exploit():
    # remove pass and add exploit code here
    # attacker = accounts[1] - account to be used for exploit
    pass

def after():
    randomUser = accounts[2]
    # Confirm other users now cannot execute flash loans
    with pytest.raises(exceptions.VirtualMachineError):
        ReceiverUnstoppable[-1].executeFlashLoan(10, {"from": randomUser})

def test_unstoppable_challenge():
    before()
    run_exploit()
    after()


import pygame

class Output:
    def __init__(self):
        self.scorenum = 0
        self.score = 0
        self.count = 3
        self.img = pygame.image.load('images/live.png')
        self.img = pygame.transform.scale(self.img, (70, 70))
        self.dead = 0 # 1 / 0

    def show(self, screen):
        if self.count > 0:
            font = pygame.font.SysFont('Helvetica', 30, True)
            s = self.scorenum
            ScoreRender = font.render('SCORE:{}'.format(s), False, (255, 255, 255))
            for i in range(self.count):
                screen.blit(self.img, (3 + i * 32, 50))
            screen.blit(ScoreRender, (13, 13))
        else:
            font = pygame.font.SysFont('Helvetica', 20, True)
            screen.blit(font.render('GAME    OVER!', False, (255, 0, 0)), (330, 353))
            pygame.display.flip()
            over = False
            cnt = 0
            while not over:
                for event in pygame.event.get():  # Получение всех событий
                    if event.type == pygame.QUIT:  # Событие выхода
                        over = True
                if cnt == 4000:
                    over = True
                pygame.time.wait(1)
                cnt += 1
            #time.sleep(5)
            self.dead = 1

    def seed_score(self):
        self.scorenum += 10
        self.score += 10

    def bigseed_score(self):
        self.scorenum += 50
        self.score += 50

    def berry_score(self):
        self.scorenum += 100

    def ghost(self, n):
        self.scorenum += 100*2**n
        self.score += 100*2**n

    def death(self):
        self.count = self.count - 1


from afbench import Benchmark, BenchmarkInfo

import numpy as np
import matplotlib.pyplot as plt

fft_info = BenchmarkInfo('benchmarks/fft.json')
print('Available benchmarks: {}'.format(fft_info.benchmark_names))
for bench in fft_info.benchmark_names:
    print('{} parameters:'.format(bench))
    for param in fft_info.params(bench):
        print('{}: {}'.format(param, fft_info.paramvals(bench, 'f32', param)))

fft1 = Benchmark(
    filepath='benchmarks/fft.json',
    name='fft1',
    filters={
        'dim1': 1,
        'dim2': 1,
        'fft_dim': 1
    }
)
fft1_x_vals = fft1.collect_param_vals('dim0', 'f32')
fft1_y_vals = fft1.collect_real_times('f32')

fft2 = Benchmark(
    filepath='benchmarks/fft.json',
    name='fft2',
    filters={
        'dim1': 16,
        'dim2': 1,
        'fft_dim': 2
    }
)
fft2_x_vals = fft2.collect_param_vals('dim0', 'f32')
for i in range(len(fft2_x_vals)):
    fft2_x_vals[i] *= 16
fft2_y_vals = fft2.collect_real_times('f32')

np_fft1_x = np.asarray(fft1_x_vals)
np_fft1_y = np.asarray(fft1_y_vals)
np_fft2_x = np.asarray(fft2_x_vals)
np_fft2_y = np.asarray(fft2_y_vals)

plt.plot(np_fft1_x, np_fft1_y, 'bo-',
         np_fft2_x, np_fft2_y, 'go-'
)
plt.xscale('log')
plt.yscale('log')
plt.show()


from random import randint
name = ""
for i in range(100):
	name = name + str(randint(0, 9))
name = name + ".txt"
try:
	f = open(name, "r")
	data = f.read()
	f.close()
except FileNotFoundError:
	print("This file no have")


from fastapi.testclient import TestClient
from sqlalchemy.orm import Session
from sqlalchemy import and_
from app import models

def test_desire_create(client: TestClient, db: Session, seller_bearer_token: str):
    cfa_image_id = client.post("/cfa-image/create", json={
        "count": 100,
        "description": "test cfa-image",
        "title": "test cfa-image"
    }, headers={"Authorization": f"Bearer {seller_bearer_token}"}).json()["id"]

    response_desire_create = client.post("/desire/create", json={
        "cfa_image_id": cfa_image_id,
        "count": 100,
        "price": 100,
    })

    assert response_desire_create.is_success

def test_desire_sell(client: TestClient, seller_bearer_token: str):
    cfa_image_id = client.post("/cfa-image/create", json={
        "count": 100,
        "description": "test cfa-image",
        "title": "test cfa-image"
    }, headers={"Authorization": f"Bearer {seller_bearer_token}"}).json()["id"]

    response_desire_create = client.post("/desire/create", json={
        "cfa_image_id": cfa_image_id,
        "count": 100,
        "price": 100,
    }).json()

    client.post("/user/deposit", json={"value": 1000000})
    pre_user = client.get("user/profile").json()

    response = client.post(f"/desire/sell/{response_desire_create['id']}",
                           headers={"Authorization": f"Bearer {seller_bearer_token}"},
                           json={"count": 100},
                           )
    assert response.is_success
    after_user = client.get("user/profile").json()
    assert pre_user["balance"] - 100 * 100 == after_user["balance"]

def test_desire_cancel(client: TestClient, seller_bearer_token):
    cfa_image_id = client.post("/cfa-image/create", json={
        "count": 100,
        "description": "test cfa-image",
        "title": "test cfa-image"
    }).json()["id"]

    response_desire_create = client.post("/desire/create", json={
        "cfa_image_id": cfa_image_id,
        "count": 100,
        "price": 100,
    }).json()
    user_desires = client.get(f"/desire/list/{cfa_image_id}").json()
    desire_id = user_desires[0]["id"]

    # cancel by another person
    response = client.post(f"/desire/cancel/{desire_id}", headers={"Authorization": f"Bearer {seller_bearer_token}"})
    assert not response.is_success

    # cancel by another person
    response = client.post(f"/desire/cancel/{desire_id}")
    assert response.is_success

    assert len(user_desires) - 1 == len(client.get(f"/desire/list/{cfa_image_id}").json())

def test_desire_list(client: TestClient, user_id: int):
    cfa_image_id = client.post("/cfa-image/create", json={
        "count": 111,
        "description": "test cfa-image",
        "title": "test cfa-image"
    }).json()["id"]

    response_desire_create = client.post("/desire/create", json={
        "cfa_image_id": cfa_image_id,
        "count": 100,
        "price": 100,
    }).json()

    response_desire_create = client.post("/desire/create", json={
        "cfa_image_id": cfa_image_id,
        "count": 10,
        "price": 10,
    }).json()

    response_desire_create = client.post("/desire/create", json={
        "cfa_image_id": cfa_image_id,
        "count": 1,
        "price": 1,
    }).json()

    desire_list_response = client.get(f"desire/list/{cfa_image_id}")

    assert desire_list_response.is_success
    assert len(desire_list_response.json()) == 3


from aoc_get import get_input

inp = get_input()
nums = [int(s) for s in inp.split(',')]
nums.sort()

# p1
med = nums[len(nums) // 2]
p1 = sum(abs(num - med) for num in nums)
print(p1)

# p2
def cost(x1, x2):
    diff = abs(x1 - x2)
    return diff * (diff + 1) // 2

avg = sum(nums) // len(nums)
p2 = sum(cost(num, avg) for num in nums)
print(p2)