import os
import time, pickle
import argparse
import numpy as np
import torch
import torch.nn as nn
from torch.cuda import amp
import torch.nn.functional as F
from torch.utils.data import DataLoader
from utils.dataloader import POI_Loader, MF_Loader, RNN_Dataset, STP_Dataloader
from utils.metric import *
from model.poiformer import TADGAT
from model.baseline import *
from model.hst_lstm import HSTLSTM_CLS
from model.stgcn import STGCN_CLS
from model.stp_udgat import STP_UDGAT

class FixNoiseStrategy():

    def __init__(self, hidden_size):
        self.hidden_size = hidden_size
        mu = 0
        sd = 1/self.hidden_size
        self.h0 = torch.randn(self.hidden_size, requires_grad=False) * sd + mu

    def on_init(self, user_len, device):
        hs = []
        for i in range(user_len):
            hs.append(self.h0)
        if not st_flag:
            return torch.stack(hs, dim=0).view(1, user_len, self.hidden_size).to(device)
        else:
            return torch.stack(hs, dim=0).to(device)

    def on_reset(self):
        return self.h0

class LstmStrategy():

    def __init__(self, hidden_size):
        self.hidden_size=hidden_size
        self.h_strategy = FixNoiseStrategy(hidden_size)
        self.c_strategy = FixNoiseStrategy(hidden_size)

    def on_init(self, user_len, device):
        h = self.h_strategy.on_init(user_len, device)
        c = self.c_strategy.on_init(user_len, device)
        return (h,c)

    def on_reset(self):
        h = self.h_strategy.on_reset()
        c = self.c_strategy.on_reset()
        return (h,c)

# class CELoss(nn.Module):
#     def __init__(self, label_smooth=None):
#         super().__init__()
#         self.label_smooth = label_smooth

#     def forward(self, pred, target):
#         eps = 1e-12
#         class_num = pred.size(-1)
#         if self.label_smooth is not None:
#             logprobs = F.log_softmax(pred, dim=1)
#             target = F.one_hot(target, pred.size(-1)).to(device)
#             target = torch.clamp(target.float(), min=self.label_smooth/(class_num-1), max=1.0-self.label_smooth)
#             loss = -1*torch.sum(target*logprobs, 1)

#         else:
#             loss = -1.*pred.gather(1, target.unsqueeze(-1)) + torch.log(torch.exp(pred+eps).sum(dim=1))
#         return loss.mean()

parser = argparse.ArgumentParser()
parser.add_argument('--mode', type=str, default='train', help='running mode, choose between [train, test]')
parser.add_argument('--dataset', type=str, help='dataset chosen for training')
parser.add_argument('--model', type=str, default='poiformer', help='model under verification in this run')
parser.add_argument('--checkpoint_dir', '-cd', type=str, default='checkpoint', help='path of checkpoint file')
parser.add_argument('--gpu_id', type=int, default=-1, help='gpu id chosen to run train/test')
parser.add_argument('--layer_num', type=int, default=1, help='number of stacked model layer')
parser.add_argument('--max_len', type=int, default=50, help='maximum length of item/user sequence')
parser.add_argument('--ubias_num', type=int, default=180, help='number of user ego-nets attention bias')
parser.add_argument('--ibias_num', type=int, default=360, help='number of item ego-nets attention bias')
parser.add_argument('--input_dim', type=int, default=32, help='dimension of input feature')
parser.add_argument('--hidden_channel', type=int, default=32, help='dimension of hidden layer')
parser.add_argument('--weight_decay', type=float, default=0, help='hyperparameter of weight decay')
parser.add_argument('--lr', type=float, default=0.01, help='learning rate')
parser.add_argument('--batch_size', type=int, default=256, help='batch size')
parser.add_argument('--test_bs', type=int, default=1024, help='batch size for test and validation')
parser.add_argument('--neg_sample_num', type=int, default=19, help='number of negative samples during training/validation')
parser.add_argument('--eval_per_n', type=int, default=1000, help='evaluate per n steps')
parser.add_argument('--epoch_num', type=int, default=30, help='training epoch number')
parser.add_argument('--ssl_temp', type=float, default=0.5, help='temperature hyperparameter for MTL')
parser.add_argument('--mtl_coef_1', type=float, default=0.5, help='coefficient of auxiliary task 1')
parser.add_argument('--mtl_coef_2', type=float, default=0.5, help='coefficient of auxiliary task 2')

args, _ = parser.parse_known_args()
checkpoint_path = '%s/%s.pth' % (args.checkpoint_dir, args.model)

# use_aux_init = args.add_aux_init == 1
# use_gt = args.add_gt == 1

# Print current experimental setting
print("Run %s" % time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(time.time())))
print("Current process ID: %d" % os.getpid())
print("Current Model: %s" % args.model)
print("Current Dataset: %s" % args.dataset)
print("Running Mode: %s" % args.mode)
print("------HyperParameter Settings------")
print("Number of Stacking Layer: %d" % args.layer_num)
print("In Channel: %d" % args.input_dim)
print("Hidden Channel: %d" % args.hidden_channel)
# print("Use Graph Embedding for POI Init: %s" % str(use_aux_init))
# print('Use Graph Transformer Layer: %s' % str(use_gt))
if args.mode == 'train':
    print("Maximum Sequence Length: %d" % args.max_len)
    print("Learning Rate: %f" % args.lr)
    print("Weight Decay: %f" % args.weight_decay)
    print("Batch Size: %d" % args.batch_size)
    print("Epoch Number: %d" % args.epoch_num)
    print("Evaluate per %d step" % args.eval_per_n)
    print("Multi-task coefficient: %.2f, %.2f" % (args.mtl_coef_1, args.mtl_coef_2))
    print("\n\n")
else:
    print("Batch Size: %d" % args.batch_size)
    print("\n\n")

if args.gpu_id != -1:
    device = torch.device('cuda:%d' % args.gpu_id if torch.cuda.is_available() else 'cpu')
else:
    device = torch.device('cpu')
enable_amp = True if "cuda" in device.type else False
enable_amp = False

# TODO: ucenter need 80 day slot, icenter need 150 day slot for gowalla
if 'TADGAT' in args.model:
    if args.mode == 'train':
        train_loader = POI_Loader(args.max_len, args.layer_num, args.ubias_num, args.ibias_num, args.neg_sample_num, dataset=args.dataset, mode='train', batch_size=args.batch_size, shuffle=True)    #, num_workers=1)
        print('train data loaded')
        val_loader = POI_Loader(args.max_len, args.layer_num, args.ubias_num, args.ibias_num, args.neg_sample_num, dataset=args.dataset, mode='val', batch_size=args.batch_size, shuffle=False)   #, num_workers=1)
        print('valid data loaded')
        user_num, item_num = train_loader.user_num, train_loader.poi_num
        # test_loader = POI_Loader(args.max_len, args.layer_num, args.ubias_num, args.ibias_num, args.neg_sample_num, mode='test', batch_size=args.batch_size, shuffle=False)  #, num_workers=1)
    else:
        test_loader = POI_Loader(args.max_len, args.layer_num, args.ubias_num, args.ibias_num, args.neg_sample_num, dataset=args.dataset, mode='test', batch_size=args.batch_size, shuffle=True)  #, num_workers=1)
        user_num, item_num = test_loader.user_num, test_loader.poi_num
        print('test data loaded')
elif args.model == 'MF' or args.model == 'LightGCN':
    if args.mode == 'train':
        train_loader = MF_Loader('train', batch_size=args.batch_size, dataset=args.dataset, shuffle=True)
        val_loader = MF_Loader('val', batch_size=args.batch_size, dataset=args.dataset, shuffle=False)
        user_num, item_num = train_loader.user_num, train_loader.poi_num
        edge_index = pickle.load(open('gowalla_processed/train_lgcn_eindex.pkl', 'rb'))
        edge_index = torch.tensor(edge_index, dtype=torch.int64).to(device)
    else:
        test_loader = MF_Loader('test', batch_size=args.batch_size, dataset=args.dataset, shuffle=False)
        user_num, item_num = test_loader.user_num, test_loader.poi_num
        edge_index = pickle.load(open('gowalla_processed/val_lgcn_eindex.pkl', 'rb'))
        edge_index = torch.tensor(edge_index, dtype=torch.int64).to(device)
    print('Data Loaded')
elif 'LSTM' in args.model or args.model == 'STGCN':
    st_flag = 'ST' in args.model and args.model != 'LSTM'
    if args.mode == 'train':
        train_dataset = RNN_Dataset(15, args.batch_size, st_flag, args.dataset, 'train')
        val_dataset = RNN_Dataset(15, args.batch_size, st_flag, args.dataset, 'val')
        train_loader, val_loader = DataLoader(train_dataset, batch_size=1), DataLoader(val_dataset, batch_size=1)
        user_num, item_num = train_dataset.user_num, train_dataset.poi_num
    else:
        test_dataset = RNN_Dataset(15, args.batch_size, st_flag, args.dataset, 'test')
        test_loader = DataLoader(test_dataset, batch_size=1)
        user_num, item_num = test_dataset.user_num, test_dataset.poi_num
elif args.model == 'STP_UDGAT':
    if args.mode == 'train':
        train_loader = STP_Dataloader('train', args.dataset, batch_size=args.batch_size, shuffle=True)
        val_loader = STP_Dataloader('val', args.dataset, batch_size=args.batch_size, shuffle=False)
        user_num, item_num = train_loader.user_num, train_loader.poi_num
    else:
        test_loader = STP_Dataloader('test', args.dataset, batch_size=args.batch_size, shuffle=False)
        user_num, item_num = test_loader.user_num, test_loader.poi_num

print(user_num, item_num)
neg_sample_num = args.neg_sample_num
if 'TADGAT' in args.model:
    model = TADGAT(args.input_dim, args.hidden_channel, args.layer_num, user_num, item_num, args.ubias_num, args.ibias_num, ssl_temp=args.ssl_temp, aux_feat=None)
elif args.model == 'MF':
    model = MF(args.input_dim, user_num, item_num)
elif 'LSTM' in args.model:
    if st_flag:
        model = HSTLSTM_CLS(args.input_dim, item_num, 168, 300)
    else:
        print('LSTM')
        model = LSTM_Basic(args.input_dim, item_num)
elif args.model == 'STGCN':
    model = STGCN_CLS(args.input_dim, item_num)
elif args.model == 'STP_UDGAT':
    model = STP_UDGAT(args.input_dim, args.hidden_channel, user_num, item_num)
elif args.model == 'LightGCN':
    model = LightGCN(args.input_dim, args.layer_num, user_num, item_num)
model = model.to(device)

if args.mode == 'train':
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[10, 20, 30], gamma=0.4)
    scaler = amp.GradScaler(enabled=enable_amp)

    # criterion = CELoss(label_smooth=0.05)
    criterion = nn.CrossEntropyLoss(label_smoothing=0.05)

def train(best_ndcg, best_hr):
    model.train()
    step = 0
    train_loader.soc_edges = train_loader.soc_edges.to(device)
    for all_out in train_loader:
        optimizer.zero_grad()
        adjs, node_idx, orig_seqs, time_diff, seq_lens, center_nid, sampled_edges = all_out
        # adjs, node_idx, soc_adjs, soc_node_idx, orig_seqs, time_diff, seq_lens, center_nid, soc_cnid, sampled_edges = all_out
        adjs = adjs.to(device)
        node_idx = node_idx.to(device)
        center_nid = center_nid.to(device)

        # soc_adjs = soc_adjs.to(device)
        # soc_node_idx = soc_node_idx.to(device)

        # orig_seqs = orig_seqs.to(device)
        # time_diff = time_diff.to(device)
        # seq_lens = seq_lens.to(device)

        # sampled_idx = sampled_edges[:, 1:]
        # sampled_idx = sampled_idx.to(device)
        # label = torch.zeros(sampled_edges.size(0), dtype=torch.int64)
        label = sampled_edges[:, 1]
        # sampled_idx = torch.arange(item_num) + user_num
        # sampled_idx = sampled_idx.repeat(label.size(0)).reshape(-1, item_num)
        # sampled_idx = sampled_idx.to(device)
        label = label.to(device)

        # lat_emb_idx, lon_emb_idx = lat_emb_idx.to(device), lon_emb_idx.to(device)

        with amp.autocast(enabled=enable_amp):
            logits = model(adjs, node_idx, orig_seqs, time_diff, seq_lens, center_nid, device, train_loader.soc_edges)
            # logits = model(adjs, node_idx, orig_seqs, time_diff, seq_lens, center_nid, device, soc_adjs, soc_node_idx, soc_cnid)

            # Code for ranking-based loss
            # loss = 0.
            # for k in range(logits.size(0)):
            #     pos_logit = logits[k, label[k]]
            #     neg_logit = torch.cat([logits[k, :label[k]], logits[k, (label[k]+1):]])
            #     # pos_loss = F.logsigmoid(pos_logit - pref_bound[k])
            #     # neg_loss = F.logsigmoid(pref_bound[k] - neg_logit).mean()
            #     pos_loss = F.logsigmoid(pos_logit)
            #     neg_loss = F.logsigmoid(-neg_logit).mean()
            #     loss += - pos_loss - neg_loss
            # loss /= logits.size(0)
            loss = criterion(logits, label) # + args.mtl_coef_1 * mtl_loss_1
            # new_end = time.time()
            # print(new_end - end)
        if np.isnan(loss.detach().cpu()):
            print("NaN detected!")
            exit(1)
        label = label.cpu()
        if step % 200 == 0:
            label = label.numpy()
            logits = logits.detach()
            sorted_idx = torch.argsort(logits, dim=-1, descending=True).cpu().numpy()
            # hr_1 += hit_rate(sorted_idx, label, k=1)
            # hr_5 += hit_rate(sorted_idx, label, k=5)
            # hr_10 += hit_rate(sorted_idx, label, k=10)
            hr_1 = hit_rate(sorted_idx, label, k=1) / args.batch_size
            hr_5 = hit_rate(sorted_idx, label, k=5) / args.batch_size
            hr_10 = hit_rate(sorted_idx, label, k=10) / args.batch_size
            print("step %d : loss %.4f, HR@1 %.4f, HR@5 %.4f, HR@10 %.4f" % (step, float(loss), hr_1, hr_5, hr_10))
        # start = time.time()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        # end = time.time()
        # print(end - start)

        # if step % 20 == 0:
        #     hr_1 /= len(all_out) ; hr_5 /= len(all_out) ; hr_10 /= len(all_out)
        #     print("step %d : loss %.4f, HR@1 %.4f, HR@5 %.4f, HR@10 %.4f" % (step, float(loss), hr_1, hr_5, hr_10))

        # if step != 0 and step % args.eval_per_n == 0:
        #     start = time.time()
        #     hr_1, hr_5, hr_10, mrr, ndcg = test('val')
        #     end = time.time()
        #     print("time consumption: %.4f" % (end-start))
        #     print("Test result on validation set: HR@1 %.4f, HR@5 %.4f, HR@10 %.4f, MRR: %.4f, NDCG@10: %.4f" % (hr_1, hr_5, hr_10, mrr, ndcg))
        #     if ndcg > best_ndcg:
        #         best_ndcg = ndcg
        #         print("New best model saved")
        #         torch.save(model.state_dict(), checkpoint_path)
        step += 1
        # end = time.time()
        # print(end - all_start)
    scheduler.step()

    # if step % args.eval_per_n > 20:
    start = time.time()
    hr_1, hr_5, hr_10, mrr, ndcg = test('val')
    end = time.time()
    print("time consumption: %.4f" % (end - start))
    print("Test result on validation set: HR@1 %.4f, HR@5 %.4f, HR@10 %.4f, MRR: %.4f, NDCG: %.4f" % (hr_1, hr_5, hr_10, mrr, ndcg))
    if mrr > best_ndcg: # or hr_10 > best_hr:
        # if ndcg > best_ndcg:
        #     best_ndcg = ndcg
        # else:
        #     best_hr = hr_10
        best_ndcg = mrr
        print("New best model saved")
        torch.save(model.state_dict(), checkpoint_path)
    return best_ndcg, best_hr

def train_stp(best_ndcg, best_hr):
    model.train()
    step = 0
    train_loader.u_eindex = train_loader.u_eindex.to(device)
    for out in train_loader:
        optimizer.zero_grad()
        s_adj, t_adj, f_adj, s_rw_adj, t_rw_adj, f_rw_adj, pp_adj, anchors, label = out
        s_adj, t_adj, f_adj = s_adj.to(device), t_adj.to(device), f_adj.to(device)
        s_rw_adj, t_rw_adj, f_rw_adj = s_rw_adj.to(device), t_rw_adj.to(device), f_rw_adj.to(device)
        pp_adj = pp_adj.to(device)
        anchors = anchors.to(device)
        label = label.to(device)
        with amp.autocast(enabled=enable_amp):
            logits = model(s_adj, t_adj, f_adj, s_rw_adj, t_rw_adj, f_rw_adj, pp_adj, train_loader.u_eindex, anchors)
            loss = criterion(logits, label)
        label = label.cpu()
        if step % 200 == 0:
            label = label.numpy()
            logits = logits.detach()
            sorted_idx = torch.argsort(logits, dim=-1, descending=True).cpu().numpy()
            hr_1 = hit_rate(sorted_idx, label, k=1) / args.batch_size
            hr_5 = hit_rate(sorted_idx, label, k=5) / args.batch_size
            hr_10 = hit_rate(sorted_idx, label, k=10) / args.batch_size
            print("step %d : loss %.4f, HR@1 %.4f, HR@5 %.4f, HR@10 %.4f" % (step, float(loss), hr_1, hr_5, hr_10))
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        step += 1
    scheduler.step()
    start = time.time()
    hr_1, hr_5, hr_10, mrr, ndcg = test_stp('val')
    end = time.time()
    print("time consumption: %.4f" % (end - start))
    print("Test result on validation set: HR@1 %.4f, HR@5 %.4f, HR@10 %.4f, MRR: %.4f, NDCG: %.4f" % (hr_1, hr_5, hr_10, mrr, ndcg))
    if ndcg > best_ndcg or hr_10 > best_hr:
        if ndcg > best_ndcg:
            best_ndcg = ndcg
        else:
            best_hr = hr_10
        print("New best model saved")
        torch.save(model.state_dict(), checkpoint_path)
    return best_ndcg, best_hr

def train_mf(best_ndcg, best_hr):
    model.train()
    step = 0
    for out in train_loader:
        optimizer.zero_grad()
        center_uid, seq = out
        with amp.autocast(enabled=enable_amp):
            if args.model == 'LightGCN':
                loss, logits = model(edge_index, center_uid, seq)
            else:
                loss, logits = model(center_uid, seq)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        if step % 5 == 0:
            sorted_idx = torch.argsort(logits, dim=-1, descending=True).cpu().numpy()
            hr_1 = top_k_acc(sorted_idx, seq, k=1) / args.batch_size
            hr_5 = top_k_acc(sorted_idx, seq, k=5) / args.batch_size
            hr_10 = top_k_acc(sorted_idx, seq, k=10) / args.batch_size
            print("step %d : loss %.4f, HR@1 %.4f, HR@5 %.4f, HR@10 %.4f" % (step, float(loss), hr_1, hr_5, hr_10))
        step += 1
    scheduler.step()
    start = time.time()
    hr_1, hr_5, hr_10, mrr, ndcg = test_mf('val')
    end = time.time()
    print("time consumption: %.4f" % (end - start))
    print("Test result on validation set: HR@1 %.4f, HR@5 %.4f, HR@10 %.4f, MRR: %.4f, NDCG: %.4f" % (hr_1, hr_5, hr_10, mrr, ndcg))
    if hr_10 > best_hr:
        if ndcg > best_ndcg:
            best_ndcg = ndcg
        else:
            best_hr = hr_10
        print("New best model saved")
        torch.save(model.state_dict(), checkpoint_path)
    return best_ndcg, best_hr

def train_rnn(best_ndcg, best_hr):
    model.train()
    step = 0
    h0_strategy = LstmStrategy(args.input_dim)
    h = h0_strategy.on_init(args.batch_size, device)
    for out in train_loader:
        optimizer.zero_grad()
        with torch.autograd.set_detect_anomaly(True):
            if not st_flag:
                seqs, lbls, reset_h = out
                seqs, lbls = seqs.to(device).squeeze(), lbls.to(device).squeeze()
            else:
                seqs, lbls, t_l, t_u, t_coef, s_l, s_u, s_coef, reset_h = out
                seqs, lbls = seqs.to(device).squeeze().T, lbls.to(device).squeeze().T
                # Note: the unsqueeze is for STGCN only
                t_l, t_u, t_coef = t_l.to(device).squeeze(), t_u.to(device).squeeze(), t_coef.to(device).squeeze()
                s_l, s_u, s_coef = s_l.to(device).squeeze(), s_u.to(device).squeeze(), s_coef.to(device).squeeze()
            lbls = lbls.reshape(-1)
            for j, reset in enumerate(reset_h):
                if reset:
                    hc = h0_strategy.on_reset()
                    if st_flag:
                        h[0][j] = hc[0]
                        h[1][j] = hc[1]
                    else:
                        h[0][0, j] = hc[0]
                        h[1][0, j] = hc[1]
            with amp.autocast(enabled=enable_amp):
                if not st_flag:
                    logits, h = model(seqs, h)
                else:
                    if args.model == 'ST_LSTM':
                        logits, h = model(seqs, t_l, t_u, s_l, s_u, t_coef, s_coef, h)
                    elif args.model == 'STGCN':
                        logits, h = model(seqs, t_l, s_l, h)
                loss = criterion(logits, lbls)
            scaler.scale(loss).backward(retain_graph=True)
            scaler.step(optimizer)
            scaler.update()
        h = (h[0].detach(), h[1].detach())
        if step % 100 == 0:
            sorted_idx = torch.argsort(logits, dim=-1, descending=True).cpu().numpy()
            hr_1 = hit_rate(sorted_idx, lbls, k=1) / logits.size(0)
            hr_5 = hit_rate(sorted_idx, lbls, k=5) / logits.size(0)
            hr_10 = hit_rate(sorted_idx, lbls, k=10) / logits.size(0)
            print("step %d : loss %.4f, HR@1 %.4f, HR@5 %.4f, HR@10 %.4f" % (step, float(loss), hr_1, hr_5, hr_10))
        step += 1
    scheduler.step()
    start = time.time()
    hr_1, hr_5, hr_10, mrr, ndcg = test_rnn('val')
    end = time.time()
    print("time consumption: %.4f" % (end - start))
    print("Test result on validation set: HR@1 %.4f, HR@5 %.4f, HR@10 %.4f, MRR: %.4f, NDCG: %.4f" % (hr_1, hr_5, hr_10, mrr, ndcg))
    if ndcg > best_ndcg or hr_10 > best_hr:
        if ndcg > best_ndcg:
            best_ndcg = ndcg
        else:
            best_hr = hr_10
        print("New best model saved")
        torch.save(model.state_dict(), checkpoint_path)
    return best_ndcg, best_hr

@torch.no_grad()
def test_mf(mode):
    print('Start Test...')
    model.eval()
    hr_1_tot, hr_5_tot, hr_10_tot = [], [], []
    if mode == 'test':
        model.load_state_dict(torch.load(checkpoint_path))
    cur_loader = val_loader if mode == 'val' else test_loader
    mrr = []
    for out in cur_loader:
        center_uid, seq = out
        if args.model == 'LightGCN':
            logits = model(edge_index, center_uid, seq)
        else:
            logits = model(center_uid, seq)
        sorted_idx = torch.argsort(logits, dim=-1, descending=True).cpu().numpy()
        hr_1 = top_k_acc(sorted_idx, seq, k=1) / args.batch_size
        hr_5 = top_k_acc(sorted_idx, seq, k=5) / args.batch_size
        hr_10 = top_k_acc(sorted_idx, seq, k=10) / args.batch_size
        hr_1_tot.append(hr_1)
        hr_5_tot.append(hr_5)
        hr_10_tot.append(hr_10)

        y_true = np.zeros_like(sorted_idx)
        for i in range(sorted_idx.shape[0]):
            for j in range(len(seq[i])):
                y_true[i, seq[i][j]] = 1
        for i in range(sorted_idx.shape[0]):
            y_true[i] = y_true[i][sorted_idx[i]]

        rr_score = y_true / (np.arange(np.shape(y_true)[1]) + 1)
        mrr.append(np.sum(rr_score) / np.sum(y_true))
    return np.mean(hr_1_tot), np.mean(hr_5_tot), np.mean(hr_10_tot), np.mean(mrr), 0

@torch.no_grad()
def test_rnn(mode):
    print('Start Test...')
    hr_1_tot, hr_5_tot, hr_10_tot = [], [], []
    mrr = []
    if mode == 'test':
        model.load_state_dict(torch.load(checkpoint_path))
    cur_loader = val_loader if mode == 'val' else test_loader
    h0_strategy = LstmStrategy(args.input_dim)
    h = h0_strategy.on_init(args.batch_size, device)
    for out in cur_loader:
        if not st_flag:
            seqs, lbls, reset_h = out
            seqs = seqs.to(device)
            seqs, lbls = seqs.squeeze(), lbls.squeeze()
        else:
            seqs, lbls, t_l, t_u, t_coef, s_l, s_u, s_coef, reset_h = out
            seqs, lbls = seqs.to(device).squeeze().T, lbls.squeeze().T
            t_l, t_u, t_coef = t_l.to(device).squeeze(), t_u.to(device).squeeze(), t_coef.to(device).squeeze()
            s_l, s_u, s_coef = s_l.to(device).squeeze(), s_u.to(device).squeeze(), s_coef.to(device).squeeze()
        lbls = lbls.reshape(-1)
        for j, reset in enumerate(reset_h):
            if reset:
                hc = h0_strategy.on_reset()
                if st_flag:
                    h[0][j] = hc[0]
                    h[1][j] = hc[1]
                else:
                    h[0][0, j] = hc[0]
                    h[1][0, j] = hc[1]
        if st_flag:
            if args.model == 'STGCN':
                logits, h = model(seqs, t_l, s_l, h)
            else:
                logits, h = model(seqs, t_l, t_u, s_l, s_u, t_coef, s_coef, h)
        else:
            logits, h = model(seqs, h)
        sorted_idx = torch.argsort(logits, dim=-1, descending=True).cpu().numpy()
        label = lbls.numpy()
        hr_1 = hit_rate(sorted_idx, label, k=1) / logits.size(0)
        hr_5 = hit_rate(sorted_idx, label, k=5) / logits.size(0)
        hr_10 = hit_rate(sorted_idx, label, k=10) / logits.size(0)
        hr_1_tot.append(hr_1)
        hr_5_tot.append(hr_5)
        hr_10_tot.append(hr_10)

        y_true = np.zeros_like(sorted_idx)
        for i in range(sorted_idx.shape[0]):
            y_true[i, label[i]] = 1
        for i in range(sorted_idx.shape[0]):
            y_true[i] = y_true[i][sorted_idx[i]]

        rr_score = y_true / (np.arange(np.shape(y_true)[1]) + 1)
        mrr.append(np.sum(rr_score) / np.sum(y_true))
    return np.mean(hr_1_tot), np.mean(hr_5_tot), np.mean(hr_10_tot), np.mean(mrr), 0

@torch.no_grad()
def test_stp(mode):
    print('Start Test...')
    model.eval()
    hr_1_tot, hr_5_tot, hr_10_tot = [], [], []
    mrr = []
    if mode == 'test':
        model.load_state_dict(torch.load(checkpoint_path))
    cur_loader = val_loader if mode == 'val' else test_loader
    cur_loader.u_eindex = cur_loader.u_eindex.to(device)
    for out in cur_loader:
        s_adj, t_adj, f_adj, s_rw_adj, t_rw_adj, f_rw_adj, pp_adj, anchors, label = out
        s_adj, t_adj, f_adj = s_adj.to(device), t_adj.to(device), f_adj.to(device)
        s_rw_adj, t_rw_adj, f_rw_adj = s_rw_adj.to(device), t_rw_adj.to(device), f_rw_adj.to(device)
        pp_adj = pp_adj.to(device)
        anchors = anchors.to(device)
        logits = model(s_adj, t_adj, f_adj, s_rw_adj, t_rw_adj, f_rw_adj, pp_adj, cur_loader.u_eindex, anchors)
        sorted_idx = torch.argsort(logits, dim=-1, descending=True).cpu().numpy()
        label = label.numpy()
        hr_1 = hit_rate(sorted_idx, label, k=1) / logits.size(0)
        hr_5 = hit_rate(sorted_idx, label, k=5) / logits.size(0)
        hr_10 = hit_rate(sorted_idx, label, k=10) / logits.size(0)
        hr_1_tot.append(hr_1)
        hr_5_tot.append(hr_5)
        hr_10_tot.append(hr_10)
        y_true = np.zeros_like(sorted_idx)
        for i in range(sorted_idx.shape[0]):
            y_true[i, label[i]] = 1
        for i in range(sorted_idx.shape[0]):
            y_true[i] = y_true[i][sorted_idx[i]]

        rr_score = y_true / (np.arange(np.shape(y_true)[1]) + 1)
        mrr.append(np.sum(rr_score) / np.sum(y_true))

    return np.mean(hr_1_tot), np.mean(hr_5_tot), np.mean(hr_10_tot), np.mean(mrr), 0

@torch.no_grad()
def test(mode):
    print('Start Test...')
    model.eval()
    hr_1_tot, hr_5_tot, hr_10_tot, mrr, ndcg_tot = [], [], [], [], []

    # TODO: uncomment the checkpoint loading operation during real test
    if mode == 'test':
        model.load_state_dict(torch.load(checkpoint_path))
    cur_loader = val_loader if mode == 'val' else test_loader
    cur_loader.soc_edges = cur_loader.soc_edges.to(device)
    # step = 0
    for n, all_out in enumerate(cur_loader):
        # adjs, node_idx, center_nid, sampled_edges = all_out
        # adjs, node_idx, soc_adjs, soc_node_idx, orig_seqs, time_diff, seq_lens, center_nid, soc_cnid, sampled_edges = all_out
        adjs, node_idx, orig_seqs, time_diff, seq_lens, center_nid, sampled_edges = all_out

        # if mode == 'val':
        #     sampled_idx = sampled_edges[:, 1:]
        #     label = torch.zeros(sampled_edges.size(0), dtype=torch.int64)
        # else:
        label = sampled_edges[:, 1]
        # sampled_idx = torch.arange(item_num) + user_num
        # sampled_idx = sampled_idx.repeat(label.size(0)).reshape(-1, item_num)
        # adjs = [adj.to(device) for adj in adjs]
        adjs = adjs.to(device)
        node_idx = node_idx.to(device)
        center_nid = center_nid.to(device)

        # soc_adjs = soc_adjs.to(device)
        # soc_node_idx = soc_node_idx.to(device)

        # orig_seqs = orig_seqs.to(device)
        # time_diff = time_diff.to(device)
        # seq_lens = seq_lens.to(device)

        # logits = model(adjs, node_idx, center_nid, device)
        logits = model(adjs, node_idx, orig_seqs, time_diff, seq_lens, center_nid, device, cur_loader.soc_edges)# , soc_adjs, soc_node_idx, soc_cnid)

        # full_logits.append(logits)
        # full_labels.append(label)
        # if mode == 'test':
        #     full_logits.append(logits)

        # logits = torch.cat(full_logits, dim=0)
        # label = torch.cat(full_labels)
        # random_noise = torch.rand(logits.numpy().shape) * 1e-7
        # logits = logits + random_noise
        sorted_idx = torch.argsort(logits, dim=-1, descending=True).cpu().numpy()
        label = label.numpy()
        # if mode == 'test':
        hr_1 = hit_rate(sorted_idx, label, k=1) / logits.size(0)
        hr_5 = hit_rate(sorted_idx, label, k=5) / logits.size(0)
        hr_10 = hit_rate(sorted_idx, label, k=10) / logits.size(0)
        # else:
        #     hr_1 = hit_rate(sorted_idx, label, k=1)
        #     hr_5 = hit_rate(sorted_idx, label, k=5)
        #     hr_10 = hit_rate(sorted_idx, label, k=10)
        # hr_50 = hit_rate(sorted_idx, k=50)
        hr_1_tot.append(hr_1)
        hr_5_tot.append(hr_5)
        hr_10_tot.append(hr_10)
        y_true = np.zeros_like(sorted_idx)
        for i in range(sorted_idx.shape[0]):
            y_true[i, label[i]] = 1
        for i in range(sorted_idx.shape[0]):
            y_true[i] = y_true[i][sorted_idx[i]]

        rr_score = y_true / (np.arange(np.shape(y_true)[1]) + 1)
        mrr.append(np.sum(rr_score) / np.sum(y_true))

        # y_true = np.zeros_like(sorted_idx)
        # for i in range(len(label)):
        #     y_true[i, label[i]] = 1

        # mrr_true = np.take(y_true, sorted_idx)
        # rr_score = mrr_true / (np.arange(np.shape(mrr_true)[1]) + 1)
        # mrr += np.sum(np.sum(rr_score, axis=1) / np.sum(mrr_true, axis=1))

        # ndcg = ndcg_score(y_true, sorted_idx, k=10)
        # if np.nan not in ndcg:
        #     # ndcg_tot += np.mean(ndcg)
        #     ndcg_tot += np.sum(ndcg)
        # # del full_logits, full_labels
        # # full_logits, full_labels = [], []

        # if n % 100 == 0 and n != 0:
        #     print(n)
        # # step += 1

    return np.mean(hr_1_tot), np.mean(hr_5_tot), np.mean(hr_10_tot), np.mean(mrr), 0 # mrr / step, ndcg_tot / step

np.random.seed(1)
torch.manual_seed(1)
if args.mode == 'train':
    ndcg = 0.
    hr_10 = 0.
    # start = time.time()
    # hr_1, hr_5, hr_10, mrr, ndcg = test_stp('val')
    # end = time.time()
    # print("time consumption: %.4f" % (end-start))
    # print("Test result on val set: HR@1 %.4f, HR@5 %.4f, HR@10 %.4f, MRR: %.4f, NDCG@10: %.4f" % (hr_1, hr_5, hr_10, mrr, ndcg))
    for epoch in range(1, args.epoch_num + 1):
        print("Currently epoch %d:" % epoch)
        print("Learning rate: %f" % scheduler.get_lr()[0])
        if 'TADGAT' in args.model:
            ndcg, hr_10 = train(ndcg, hr_10)
        elif args.model == 'MF' or args.model == 'LightGCN':
            ndcg, hr_10 = train_mf(ndcg, hr_10)
        elif 'LSTM' in args.model or args.model == 'STGCN':
            train_dataset.reset()
            ndcg, hr_10 = train_rnn(ndcg, hr_10)
        elif args.model == 'STP_UDGAT':
            ndcg, hr_10 = train_stp(ndcg, hr_10)
else:
    start = time.time()
    if 'TADGAT' in args.model:
        hr_1, hr_5, hr_10, mrr, ndcg = test('test')
    elif args.model == 'MF' or args.model == 'LightGCN':
        hr_1, hr_5, hr_10, mrr, ndcg = test_mf('test')
    elif 'LSTM' in args.model or args.model == 'STGCN':
        hr_1, hr_5, hr_10, mrr, ndcg = test_rnn('test')
    elif args.model == 'STP_UDGAT':
        hr_1, hr_5, hr_10, mrr, ndcg = test_stp('test')
    end = time.time()
    print("time consumption: %.4f" % (end-start))
    print("Test result on test set: HR@1 %.4f, HR@5 %.4f, HR@10 %.4f, MRR: %.4f, NDCG@10: %.4f" % (hr_1, hr_5, hr_10, mrr, ndcg))


# -*- coding: utf-8 -*-
"""
Created on Mon May 26 17:56 2023

@author: kenna
"""

# System modules
from HdfsAudioStore.workers import FileWorker
from HdfsAudioStore.audioHandler import AudioHandler
from HdfsAudioStore.hbaseAudio import HBaseAudioStore
from HdfsAudioStore.model import AudioModel
from HdfsAudioStore.hbaseAudio import Columns

class DatabaseWorker:
    """
    Class to support HBaseAudioStore with storing audio files from s3 in it
    """

    def __init__(self, table_name: str, host: str, port: int, createTable = False):

        # Compose supporting modules
        self.fileWorker = FileWorker.FileWorker()
        self.audioHandler = AudioHandler.AudioHandler()
        self.hbaseColumns = Columns.ColumnFamilyEnum

        # Define HBase config
        self.table_name = table_name
        self.host = host
        self.port = port
        self.hbaseAudioStore = HBaseAudioStore.HBaseAudioDataStore(self.table_name, self.host, self.port)
        self.createTable = createTable
        if createTable is True:
            self.hbaseAudioStore.create_table()

    def importTrack(self, trackPath: str, owner: str, audio_id: str = None):
        """
        Import track from s3 into HBase table
        """

        # Download track
        audioModel = self.fileWorker.fetchAudioFile(trackPath, owner)
        # print("Audio Model Type:\t" + type(audioModel))

        # Handle row key
        dict = audioModel.getTrackMetaData().toDict()["TrackMetaData"]
        if audio_id is None:
            audio_id = str(dict["Owner"].replace(" ", "") + "-" + dict["TrackName"])
            # print("RowKey:\t" + rowKey)
        print(f'DatabaseWorker: Using rowKey = {audio_id}')

        # Post audio signal
        self.hbaseAudioStore.put_audio_data(audio_id, audioModel)

        # Post audio metadata
        self.hbaseAudioStore.put_audio_metadata(audio_id, audioModel)

        # Post track metadata
        self.hbaseAudioStore.put_track_metadata(audio_id, audioModel)

    def handleColumnFamily(self, column: str):
        """
        Check if column choice is valid
        """
        return self.hbaseColumns.queryEnum(column)

    # What columns? Supported by enum => AudioSignal, AudioMeta, TrackMeta, ALL
    def getTrack(self, audioId: str, column: str):
        """
        Fetch audio model from DB
        """

        # Explode if invalid
        isValid = self.handleColumnFamily(column)
        if isValid is None:
            return None

        # Fetch required data
        if self.hbaseColumns.isAudio(column):
            audioSignal = self.hbaseAudioStore.get_audio_data(audioId)
            audioSignal = self.audioHandler.rebuildAudio(audioSignal)
            return audioSignal

        # Audio meta
        elif self.hbaseColumns.isAudioMeta(column):
            audioMeta = self.hbaseAudioStore.get_audio_metadata(audioId)
            print(f'Audio MetaData Results:\n {audioMeta}')
            audioMeta = self.audioHandler.makeAudioMeta(
                audioMeta["SamplingRate"],
                audioMeta["FrameCount"],
                audioMeta["Duration"]
            )
            return audioMeta

        # Track meta
        elif self.hbaseColumns.isTrackMeta(column):
            trackMeta = self.hbaseAudioStore.get_track_metadata(audioId)
            print(f'Track Meta Data Results:\n {trackMeta}')
            trackMeta = self.audioHandler.makeTrackMeta(
                trackMeta["TrackName"],
                trackMeta["Owner"],
                trackMeta["Timestamp"]
            )
            return trackMeta

        # Fetch all
        else:

            # Construct components
            audioSignal = self.hbaseAudioStore.get_audio_data(audioId)
            audioSignal = self.audioHandler.rebuildAudio(audioSignal)
            audioMeta = self.hbaseAudioStore.get_audio_metadata(audioId)
            audioMeta = self.audioHandler.makeAudioMeta(
                audioMeta["SamplingRate"],
                audioMeta["FrameCount"],
                audioMeta["Duration"]
            )
            trackMeta = self.hbaseAudioStore.get_track_metadata(audioId)
            trackMeta = self.audioHandler.makeTrackMeta(
                trackMeta["TrackName"],
                trackMeta["Owner"],
                trackMeta["Timestamp"]
            )

            # Construct & return audio model
            audioModel = AudioModel.AudioModel(audioSignal, audioMeta, trackMeta)
            return audioModel

    def fetchTrack(self, trackName: str, owner: str, outPath: str):
        """
        Fetch wav from DB
        """

        # Get track
        audioId = str(owner.replace(" ", "") + "-" + trackName)
        audioModel = self.getTrack(audioId, "All")

        # Write track
        #  only if output
        self.audioHandler.writeAudio(audioModel, outPath)


from PyQt5.QtWidgets import QWidget
from PyQt5 import uic
from PyQt5.Qt import QSize
from PyQt5.Qt import QByteArray
from PyQt5.Qt import QIcon
from PyQt5.QtWidgets import QMenu
from mc.app.Settings import Settings
from mc.common.globalvars import gVar
from PyQt5.Qt import Qt
from mc.bookmarks.bookmarksimport.BookmarksImportDialog import BookmarksImportDialog
from mc.bookmarks.bookmarksexport.BookmarksExportDialog import BookmarksExportDialog
from mc.history.HistoryManager import HistoryManager
from mc.bookmarks.BookmarksManager import BookmarksManager
from mc.lib3rd.FancyTabWidget import FancyTabWidget

class BrowsingLibrary(QWidget):
    def __init__(self, window, parent=None):
        '''
        @param: window BrowserWindow
        @param: parent QWidget
        '''
        super().__init__(parent)
        self._ui = uic.loadUi('mc/other/BrowsingLibrary.ui', self)
        self._historyManager = HistoryManager(window)  # HistoryManager
        self._bookmarksManager = BookmarksManager(window)  # BookmarksManager

        settings = Settings()
        settings.beginGroup('BrowsingLibrary')
        self.resize(settings.value('size', QSize(760, 470)))
        self._historyManager.restoreState(settings.value('historyState', QByteArray()))
        settings.endGroup()

        gVar.appTools.centerWidgetOnScreen(self)

        historyIcon = QIcon()
        historyIcon.addFile(':/icons/other/bighistory.svg', QSize(), QIcon.Normal)
        historyIcon.addFile(':/icons/other/bighistory-selected.svg', QSize(), QIcon.Selected)

        bookmarksIcon = QIcon()
        bookmarksIcon.addFile(':/icons/other/bigstar.svg', QSize(), QIcon.Normal)
        bookmarksIcon.addFile(':/icons/other/bigstar-selected.svg', QSize(), QIcon.Selected)

        self._ui.tabs.AddTab(self._historyManager, historyIcon, _('History'))
        self._ui.tabs.AddTab(self._bookmarksManager, bookmarksIcon, _('Bookmarks'))
        self._ui.tabs.SetMode(FancyTabWidget.Mode_LargeSidebar)
        self._ui.tabs.setFocus()

        m = QMenu()
        m.addAction(_('Import Bookmarks...'), self._importBookmarks)
        m.addAction(_('Export Bookmarks...'), self._exportBookmarks)
        self._ui.importExport.setMenu(m)

        self._ui.tabs.CurrentChanged.connect(self._ui.searchLine.clear)
        self._ui.searchLine.textChanged.connect(self._search)

        gVar.appTools.setWmClass('Browsing Library', self)

    def showHistory(self, window):
        '''
        @param: window BrowserWindow
        '''
        self._ui.tabs.SetCurrentIndex(0)
        self.show()
        self._historyManager.setMainWindow(window)

        self.raise_()
        self.activateWindow()

    def showBookmarks(self, window):
        '''
        @param: window BrowserWindow
        '''
        self._ui.tabs.SetCurrentIndex(1)
        self.show()
        self._bookmarksManager.setMainWindow(window)

        self.raise_()
        self.activateWindow()

    def historyManager(self):
        '''
        @return: HistoryManager
        '''
        return self._historyManager

    def bookmarksManager(self):
        '''
        @return: BookmarksManager
        '''
        return self._bookmarksManager

    # private Q_SLOTS:
    def _search(self):
        if self._ui.tabs.current_index() == 0:
            self._historyManager.search(self._ui.searchLine.text())
        else:
            self._bookmarksManager.search(self._ui.searchLine.text())

    def _importBookmarks(self):
        dialog = BookmarksImportDialog(self)
        dialog.open()

    def _exportBookmarks(self):
        dialog = BookmarksExportDialog(self)
        dialog.open()

    # private:
    def closeEvent(self, event):
        '''
        @param: event QCloseEvent
        '''
        settings = Settings()
        settings.beginGroup('BrowsingLibrary')
        settings.setValue('size', self.size())
        settings.setValue('historySize', self._historyManager.saveState())
        settings.endGroup()
        event.accept()

    def keyPressEvent(self, event):
        '''
        @param: event QKeyEvent
        '''
        evtKey = event.key()
        evtMod = event.modifiers()
        if evtKey == Qt.Key_Escape or \
                (evtKey == Qt.Key_W and evtMod == Qt.ControlModifier):
            self.close()

        super().keyPressEvent(event)


# A program that prompts the user to enter an integer and
# determines whether the number is even or odd,
# taking into account the case of entering incorrect data.

# try:
#     user_number = input("Enter number: ")
#     # user_number = int(input("Enter number: "))
#     if user_number % 2 == 0:
#         print("Your number is even")
#     else:
#         print("Your number is odd")
# except ValueError:
#     print('You did not enter an integer')
# except TypeError:
#     print('You didn`t enter an integer')

#############################

# def check_even(number):
#     try:
#         return 'odd mumber' if int(number)%2 else "even number"
#     except ValueError as v:
#         print('Error', v)
#     except TypeError as t:
#         print('Error', t)

# print(check_even(62))

###########################

# A program to calculate the divide of two numbers entered
# by the user sequentially through a comma, to predict the case of division by zero,
# cases of syntactic errors and cases of other exceptional situations.
# Use the else and finally blocks.

# def divide_two_numbers(a,b):
#     try:
#         # Якщо input, то потрібно зі str переобразити в int
#         # a = int(a)
#         # b = int(b)
#         return a/b
#     except ZeroDivisionError as e:
#         return 'Error', e
#     except SyntaxError as e:
#         return 'Error', e
#     except:
#         return 'There are errors in the code'
#     else:
#         return 'There are no errors in the code'
#     finally:
#         return 'The end'

# print(divide_two_numbers(8,4))

#######################

# try:
#     # eval перетворює зі стрінга в інтеджер та розділяє по комі на окремі зміні
#     num1, num2 = eval(input('Enter two number, separated by a comma: '))
#     result = num1 / num2
#     print('Result is', result)

#######################

# class CustomError(Exception):
#     def __init__(self, data):
#         self.data = data

#     def __str__(self):
#         return repr(self.data)

# total_score = int(input("Enter expert score: "))
# try:
#     num_of_group = int(input("Enter number of your group: "))
#     if num_of_group < 1:
#         raise CustomError("Number of your group can't be less than 1")
# except CustomError as e:
#     print("We obtain error:", e.data)

#######################

# HOMEWORK
# A program that prompts the user to enter their age,
# and then displays a message stating whether the age is even or odd.
# The program must provide the ability to enter a negative number, and in this case generate an exception.
# The master code should call a function that processes the information entered.

# class CustomError(Exception):
#     def __init__(self, data):
#         self.data = data

#     def __str__(self):
#         return repr(self.data)

# try:
#     user_age = int(input("Enter your age: "))
#     if user_age < 0:
#         raise CustomError("You enter no correct number")
#     elif user_age % 2 == 0:
#         print('Your age is even')
#     else:
#         print('Your age is odd')

# except CustomError as e:
#     print("We obtain error:", e.data)

#######################

# HOMEWORK
# Write a program that analyzes the entered number and, depending on the number,
# gives the day of the week that corresponds to this number (1 is Monday, 2 is Tuesday, etc.).
# Take into account cases of entering numbers from 8 and more, as well as cases of entering non-numerical data.

def days_of_week(user_number):
    try:
        days_week = {1:"понеділок", 2:"вівторок", 3:"середа", 4:"четвер", 5:"п'ятниця", 6:"субота", 7:"неділя"}
        if user_number <= 0 or user_number > 7:
            raise IndexError("You enter no correct number")
        print(days_week[user_number])
    except ValueError:
        print("Wrong Input")
    except TypeError as e:
            print("We obtain error:", e)
    except IndexError as e:
        print("We obtain error:", e)

days_of_week(2)

###########################################
# def day_of_the_week (day_number):
#     try:
#         week_day_name = \
#             {1: 'Monday',
#              2: 'Tuesday',
#              3: 'Wednesday',
#              4: 'Thursday',
#              5: 'Friday',
#              6: 'Saturday',
#              7: 'Sunday'}

#         return week_day_name[int(day_number)]
#     except KeyError as e:
#         print('This key as not in list ', e)
#     except ValueError as e:
#         print('invalid literal', e)

# print(day_of_the_week(input('Enter day number: ')))


def swap(x,y,arr):
    arr[x] = arr[x]^arr[y]
    arr[y] = arr[x]^arr[y]
    arr[x] = arr[x]^arr[y]
def over(l,r,arr):
    i = 0
    while l+i < r-i:
        swap(l+i,r-i,arr)
        i += 1

def solve(n,arr):
    result = []
    i = 0
    while i < len(arr):
        min = i
        j = i
        while j < len(arr):
            if arr[j] < arr[min]:
                min = j
            j += 1
        result.append(min)
        over(i,min,arr)
        i += 1
    for x in range(n):
        print(result[x]+1,end=" ")
    #print()

#main-----
n = int(input())
arr = [int(x) for x in input().split(" ")]
solve(n,arr)


"""在线策略系统"""
from   const import Valid_Col
import pandas as pd
from   policy.base_policy import BasePolicy
from   callback import error_callback
import itertools
import numpy as np
from data.hive_sql import *
from solvers.base_solver import *

class MPSFilterPolicy(BasePolicy):
    """mps解空间过滤策略"""

    def run(self):
        pred_candidate = self.df
        ecrdf = self.dao.get_wyc_ecr()
        crdf = self.dao.get_pred_cr()
        '''读取城市分框信息'''
        filter_pred_candidate_list = []
        print("城市分框数量：", len(self.dao.city_group_list))
        if len(self.dao.city_group_list) > 0:
            for group_info in self.dao.city_group_list:
                sdate, edate = group_info['start_date'], group_info['end_date']
                if type(group_info['city_list']) == str:
                    cis = list(map(int, group_info['city_list'].split(',')))
                elif type(group_info['city_list']) == list:
                    cis = group_info['city_list']
                city_bucket_budget = group_info['budget_upperbound']
                print("城市分框：", "开始日期：", sdate, "结束日期：", edate, "预算：", city_bucket_budget)
                print("城市分框 cis:", cis)

                '''生成每个ecr分框的预算'''
                group_ecrdf = ecrdf.query("city_id in @cis and stat_date >= '%s' and stat_date <= '%s'" % (sdate, edate))
                bin_costdf = self.gen_ecr_bucket_budget(group_ecrdf, sdate, edate, cis, city_bucket_budget)
                if len(bin_costdf) == 0:
                    continue
                else:
                    group_crdf = crdf.query("city_id in @cis and stat_date >= '%s' and stat_date <= '%s'" % (sdate, edate))
                    '''生成ecr-cr分框的锚点补贴率'''
                    cr_ratedf = self.gen_cr_bucket_rate(group_crdf, group_ecrdf, bin_costdf)
                    '''过滤候选集'''
                    cr_ratedf['min_combo_rate'] = cr_ratedf['avg_rate'].apply(lambda x: x * 0.8).astype(float)
                    cr_ratedf['max_combo_rate'] = cr_ratedf['avg_rate'].apply(lambda x: x * 1.2).astype(float)
                    # 临时兜底约束
                    cr_ratedf['max_combo_rate'] = cr_ratedf.apply(lambda x: 0.01 if x.city_id in (8, 25) else x.max_combo_rate, axis=1)
                    cr_ratedf['min_combo_rate'] = cr_ratedf.apply(lambda x: 0.0 if x.city_id in (8, 25) else x.min_combo_rate, axis=1)
                    cr_ratedf['max_combo_rate'] = cr_ratedf.apply(lambda x: 0.02 if x.city_id in (14, 16, 22) else x.max_combo_rate, axis=1)
                    cr_ratedf['min_combo_rate'] = cr_ratedf.apply(lambda x: 0.0 if x.city_id in (14, 16, 22) else x.min_combo_rate, axis=1)

                    group_candidate = cr_ratedf.merge(pred_candidate, on=['city_id', 'stat_date'], how='left')
                    print("城市分框，候选集数量：", len(group_candidate))
                    filter_group_candidate = group_candidate.query("is_candidate==1 and crate>=min_combo_rate and crate<=max_combo_rate")
                    print("城市分框，过滤后的候选集数量：", len(filter_group_candidate))
                    filter_pred_candidate_list.append(filter_group_candidate)
        filter_pred_candidate = pd.concat(filter_pred_candidate_list) # 过滤后的候选集
        print('过滤后的候选集数量:', len(filter_pred_candidate))
        print(filter_pred_candidate.head(5))
        return filter_pred_candidate

    def gen_ecr_bucket_budget(self, ecrdf, start_date, end_date, city_list, city_bucket_budget):
        '''生成ecr分框预算'''
        ecrdf['ecr_bin'] = ecrdf['wyc_ecr'].apply(self.gen_ecr_bucket)
        ecrdf = ecrdf[~ecrdf['ecr_bin'].isna()]

        dates = list(pd.date_range(start_date, end_date).astype(str))
        fkhf_list = list(itertools.product(city_list, dates, np.arange(0, 0.11, 0.0005)))
        ratedf = pd.DataFrame(fkhf_list, columns=['city_id', 'stat_date', 'fkhf'])
        print("ratedf:", ratedf.head(5))
        print("citys:", city_list)

        df = pd.merge(ecrdf, ratedf, on=['city_id', 'stat_date'])
        df['ecr_dgmv_ratio'] = df.apply(self.gen_ecr_dgmv_ratio, axis=1)
        gmvdf = self.df[['city_id', 'stat_date', 'total_gmv']].drop_duplicates() # candidate表
        opti_data = pd.merge(df, gmvdf, on=['city_id', 'stat_date'], how='left')

        opti_data['delta_gmv'] = opti_data['ecr_dgmv_ratio'] * opti_data['total_gmv']
        opti_data['delta_cost'] = opti_data['fkhf'] * opti_data['total_gmv']
        opti_data[Valid_Col] = 1

        print("ecr bucket opti_data number:", len(opti_data))
        if len(opti_data) == 0:
            return pd.DataFrame([], columns=['ecr_bin', 'budget'])
        else:
            # 整数规划求解
            solver = BaseSolver(self.CONF, opti_data, is_bp=False, city_must_cost=True, mode='city_bucket', total_budget=city_bucket_budget)
            resdf = solver.solve()
            bin_costdf = resdf.groupby(['ecr_bin'])[['delta_cost']].agg(sum).reset_index().rename(columns={'delta_cost': 'budget'})
            print("ecr分框的预算:")
            print(bin_costdf)
            return bin_costdf

    def gen_ecr_bucket(self, x):
        if x <= 0.59:  # x > 0.15
            return 0
        elif x > 0.59 and x <= 0.635:
            return 1
        elif x > 0.635 and x <= 0.66:
            return 2
        elif x > 0.66 and x <= 0.685:
            return 3
        elif x > 0.685: # and x <= 0.78
            return 4

    def gen_ecr_dgmv_ratio(self,x):
        w = 0.65
        if x.ecr_bin == 0:
            return 0.744 * (x.fkhf ** (w))
        elif x.ecr_bin == 1:
            return 0.69 * (x.fkhf ** (w))
        elif x.ecr_bin == 2:
            return 0.64 * (x.fkhf ** (w))
        elif x.ecr_bin == 3:
            return 0.597 * (x.fkhf ** (w))
        elif x.ecr_bin == 4:
            return 0.542 * (x.fkhf ** (w))

    def gen_cr_bucket(self,x):
        if x > 0 and x <= 0.5:
            return 0
        elif x > 0.5 and x <= 0.6:
            return 1
        elif x > 0.6 and x <= 0.68:
            return 2
        elif x > 0.68 and x <= 0.72:
            return 3
        elif x > 0.72 and x <= 0.76:
            return 4
        elif x > 0.76 and x <= 0.8:
            return 5
        elif x > 0.8 and x <= 0.85:
            return 6
        elif x > 0.85:
            return 7

    def gen_cr_bucket_rate(self, score, ecrdf, bin_costdf):
        cr_col = 'score'
        score = score[score[cr_col] > 0]
        score['bin'] = score[cr_col].apply(self.gen_cr_bucket)
        score['bin_val'] = score['bin'].apply(lambda x: (x + 1) / 8 if x <= 3 else (0.2 * (x + 1) if x <= 5 else 0.25 * x))

        gmvdf = self.df[['city_id', 'stat_date', 'total_gmv']].drop_duplicates()  # candidate表
        weightdf = pd.merge(score, gmvdf, on=['city_id', 'stat_date'])
        weightdf['weight'] = weightdf.apply(lambda x: x['total_gmv'] * x['bin_val'], axis=1)

        df_li = []
        ecr_bins = ecrdf.ecr_bin.unique()
        print("ecr bins:", ecr_bins)
        for ecr_bin in ecr_bins:
            tmpdf = ecrdf.query("ecr_bin==%d" % ecr_bin)[['city_id', 'stat_date']]
            tmpdf2 = pd.merge(weightdf, tmpdf)
            if len(tmpdf2) > 0:
                bin_budget = bin_costdf.query("ecr_bin==%d" % ecr_bin)['budget'].iloc[0] # ecr分框预算
                avg_rate = (bin_budget) / (tmpdf2['weight'].sum())  # 泛快呼返
                print("avg_rate:", avg_rate)
                tmpdf2['avg_rate'] = tmpdf2['bin_val'].apply(lambda x: avg_rate * x)
                tmpdf2['ecr_bin'] = ecr_bin
                df_li.append(tmpdf2)
            else:
                continue
        cr_ratedf = pd.concat(df_li, ignore_index=True)
        return cr_ratedf


from django.shortcuts import render
from . import models
from packages.models import Package,PackageStudent
from django.urls import reverse
from django.http import HttpResponseRedirect
from django.http import HttpResponseRedirect,Http404
from django.shortcuts import get_object_or_404

# Create your views here.

def TestView(request,package_id):
    if request.method == 'POST':
        marks = 0
        test = models.Test.objects.filter(package_id=package_id)
        for i in range(0,len(test)):
            test_id = test[i].id - 1
            if request.POST.get(str(test_id+1)) == test[i].correct:
                marks = marks + test[i].marks
            else:
                continue

        packagestudent = get_object_or_404(PackageStudent,package=package_id,student=request.user)

        packagestudent.marks = marks
        packagestudent.submission = True
        packagestudent.save()
        return HttpResponseRedirect(reverse('tests:test_successful'))
    else:
        try:
            ps = PackageStudent.objects.get(package_id=package_id,student_id=request.user.id,package__q_visibility=True,submission=False)
            data = ps.package.package_test.all()
        except:
            return HttpResponseRedirect(reverse('tests:test_denied'))

        package_end_date = get_object_or_404(Package,id=package_id).end_time
        # package_start_date = get_object_or_404(Package,id=package_id).start_time
        return render(request,'tests/test.html',{'form':data,'end_date':str(package_end_date)})


import numpy as np
from matplotlib import pyplot as plt

# Load the data
filename = 'randomData.txt'
coord = np.dtype(dict(names=('x', 'y', 'z'), \
                      formats=('float', 'float', 'float') ) )
data = np.loadtxt(filename, dtype=coord, delimiter=',')

# create the subplots
f, axs = plt.subplots(2,2)
axs[0][0].plot(data['x'], data['y'], 'bo')
axs[0][1].plot(data['x'], data['z'], 'ro')
axs[1][0].plot(data['y'], data['z'], 'go')
axs[1][1].plot(data['x'], data['y'], 'bo', data['y'], data['z'], 'ro', \
               data['x'], data['z'], 'go' )


from twisted.trial import unittest
from twisted.internet.address import IPv4Address

from txyam.factory import ConnectingMemCacheProtocol, MemCacheClientFactory

class FactoryTest(unittest.TestCase):

    def test_buildProtocol(self):
        f = MemCacheClientFactory()
        addy = IPv4Address('TCP', 'ahost', 1234)
        p = f.buildProtocol(addy, timeOut=123)
        self.assertIsInstance(p, ConnectingMemCacheProtocol)
        self.assertEqual(p.persistentTimeOut, 123)
        self.assertEqual(str(p), "memcache[%s]" % addy)


##
# File:    ContentDefinition.py
# Author:  J. Westbrook
# Date:    8-May-2018
# Version: 0.001 Initial version
#
# Updates:
#  23-May-2018 jdw add method to return vanilla data Type dictionary/revise api -
#   5-Jun-2018 jdw replace IoAdapter() with MarshalUtil() -
#  10-Jun-2018 jdw move the base table definitions in terms of content classes
#  16-Jun-2018 jdw standardize method name for getAttributeDataTypeD()
#  26-Nov-2018 jdw turn down boundary value level to standard ddl values
#   3-Feb-2019 jdw get all child items using self.__dApi.getFullDescendentList()
#  31-Mar-2019 jdw include parent details in __getAttributeFeatures()
#   6-Jun-2019 jdw take dictionary API as an argument.
#  22-Aug-2019 jdw unify naming conventions dictSubset->databaseName and *ForSubset() -> *ForDatabase()
#   5-Sep-2019 jdw add extended DDL metadata.
#  24-Jan-2022 dwp Exclude all categories beginning with "ma_" from being mandatory
#                  (temporarily hardcoded here until new configuration file section added to achieve same effect)
#
#
##
"""
Assemble configuration and dictionary metadata required to build/load database schema definitions ...

"""
__docformat__ = "restructuredtext en"
__author__ = "John Westbrook"
__email__ = "jwest@rcsb.rutgers.edu"
__license__ = "Apache 2.0"

import logging
import textwrap
from collections import OrderedDict

from mmcif.api.PdbxContainers import CifName

logger = logging.getLogger(__name__)

class ContentDefinition(object):
    """Assemble configuration and dictionary metadata required to build/load database schema definitions ..."""

    def __init__(self, dictApi, contentDefHelper=None, databaseName=None, **kwargs):
        """
        Args:
            dictApi (object): instance of DictionaryApi() class
            contentHelper (object, optional): an instance of a contentHelper().
            databaseName (string, optional): name of a content database (e.g. 'chem_comp', 'bird', 'bird_family', 'pdbx')

        """
        self.__dApi = dictApi
        self.__databaseName = databaseName
        self.__setup(contentDefHelper, databaseName)
        _ = kwargs
        #

    def __setup(self, contentDefHelper, databaseName):
        #
        iTypeCodes = []
        iQueryStrings = []
        unitCardinalityList = []
        dataSelectFilterD = OrderedDict()

        itemTransformD = OrderedDict()
        self.__categoryContentClasses = OrderedDict()
        self.__attributeContentClasses = OrderedDict()
        self.__intEnumD = OrderedDict()
        iterableD = OrderedDict()
        #
        self.__categoryList = sorted(self.__dApi.getCategoryList())
        self.__categorySchema = {catName: sorted(self.__dApi.getAttributeNameList(catName)) for catName in self.__categoryList}
        ##
        _, attributeContextIndex = self.__indexContexts(self.__categorySchema)
        self.__keyReplaceItems = attributeContextIndex["RCSB_CATEGORY_KEY_SUBSTITION"] if "RCSB_CATEGORY_KEY_SUBSTITION" in attributeContextIndex else []
        self.__keyReplaceCategoryD = OrderedDict()
        for catName, atName in self.__keyReplaceItems:
            self.__keyReplaceCategoryD.setdefault(catName, []).append(atName)
        #
        logger.debug("Primary key replacements: %r", self.__keyReplaceItems)
        logger.debug("Primary key replacement category index: %r", self.__keyReplaceCategoryD)
        #
        if contentDefHelper and databaseName:
            self.__intEnumD = {(tD["CATEGORY_NAME"], tD["ATTRIBUTE_NAME"]): True for tD in contentDefHelper.getInternalEnumItems(databaseName)}
            logger.debug("Internal enum items %r", self.__intEnumD)
            #
            cardD = contentDefHelper.getCardinalityKeyItem(databaseName)
            logger.debug("Cardinality attribute %r", cardD.items())
            #
            unitCardinalityList = self.__getUnitCardinalityCategories([cardD])
            unitCardinalityList.extend(contentDefHelper.getCardinalityCategoryExtras())
            logger.debug("Cardinality categories %r", unitCardinalityList)
            #
            self.__categoryContentClasses = contentDefHelper.getCategoryContentClasses(databaseName)
            logger.debug("categoryContentClasses %r", self.__categoryContentClasses)
            #
            self.__attributeContentClasses = contentDefHelper.getAttributeContentClasses(databaseName)
            logger.debug("attributeContentClasses %r", self.__attributeContentClasses)
            #
            self.__sliceParentItemsD = contentDefHelper.getDatabaseSliceParents(databaseName)
            self.__sliceUnitCardinalityD = OrderedDict()
            self.__sliceCategoryExtrasD = OrderedDict()
            for sliceName, pDL in self.__sliceParentItemsD.items():
                logger.debug("Slicename %s parents %r", sliceName, pDL)
                #
                # Some categories are included in a slice even if they are unconnected to the slice parent.
                self.__sliceCategoryExtrasD[sliceName] = contentDefHelper.getSliceCategoryExtras(databaseName, sliceName)
                logger.debug("Slice %s extra categories %r", sliceName, self.__sliceCategoryExtrasD[sliceName])
                #
                self.__sliceUnitCardinalityD[sliceName] = self.__getUnitCardinalityCategories(pDL)
                logger.debug("Slicename %s unit cardinality categories %r", sliceName, self.__sliceUnitCardinalityD[sliceName])
                #
                self.__sliceUnitCardinalityD[sliceName].extend(contentDefHelper.getSliceCardinalityCategoryExtras(databaseName, sliceName))
                logger.debug("Slicename %s unit cardinality categories %r", sliceName, self.__sliceUnitCardinalityD[sliceName])
            #
            #
            subCategoryD = {}
            for catName, atNameList in self.__categorySchema.items():
                scL = []
                for atName in atNameList:
                    scL.extend(self.__dApi.getItemSubCategoryIdList(catName, atName))
                scL = list(OrderedDict.fromkeys(scL)) if scL else []
                subCategoryD[catName] = scL
            logger.debug("Subcategory category dictionary %r", subCategoryD)
            #
            self.__categoryFeatures = {catName: self.__getCategoryFeatures(catName, unitCardinalityList, subCategoryD) for catName in self.__categoryList}
            iTypeCodes = contentDefHelper.getTypeCodes("iterable")
            emiTypeCodes = contentDefHelper.getTypeCodes("embedded_iterable")
            logger.debug("emiTypeCodes %r", emiTypeCodes)
            iQueryStrings = contentDefHelper.getQueryStrings("iterable")
            logger.debug("iterable types %r iterable query %r", iTypeCodes, iQueryStrings)
            iterableD = self.__getIterables(iTypeCodes, iQueryStrings, contentDefHelper)
            embeddedIterableD = self.__getEmbeddedIterables(emiTypeCodes)
            logger.debug("iterableD %r", iterableD.items())
            logger.debug("embeddedIterableD %r", embeddedIterableD.items())
            dataSelectFilterD = contentDefHelper.getDatabaseSelectionFilters(databaseName)
            itemTransformD = contentDefHelper.getItemTransformD()
            logger.debug("itemTransformD %r", itemTransformD.items())
            #
            self.__selectionFiltersD = contentDefHelper.getDatabaseSelectionFilters(databaseName)
            self.__sliceParentFiltersD = contentDefHelper.getDatabaseSliceParentFilters(databaseName)
        else:
            logger.debug("Dictionary helper not loaded for schema %r", databaseName, stack_info=True)
            self.__selectionFiltersD = {}
            self.__sliceParentItemsD = {}
            self.__sliceParentFiltersD = {}
            self.__sliceUnitCardinalityD = {}
            self.__sliceCategoryExtrasD = {}
            embeddedIterableD = {}
            logger.warning("Missing dictionary helper method or schema %r", databaseName)

        #
        self.__methodD = self.__getMethodInfo()
        self.__attributeFeatures = OrderedDict(
            {catName: self.__getAttributeFeatures(catName, iterableD, embeddedIterableD, itemTransformD, self.__methodD) for catName in self.__categoryList}
        )
        self.__attributeDataTypeD = OrderedDict({catName: self.__getAttributeTypeD(catName) for catName in self.__categoryList})
        self.__dataSelectFilterD = dataSelectFilterD
        self.__sliceD = self.__getSliceChildren(self.__sliceParentItemsD)

    def getSelectionFiltersForDatabase(self):
        try:
            return self.__selectionFiltersD
        except Exception:
            pass
        return {}

    def getSliceParentFiltersForDatabase(self):
        try:
            return self.__sliceParentFiltersD
        except Exception:
            pass
        return {}

    def getSliceParentItemsForDatabase(self):
        try:
            return self.__sliceParentItemsD
        except Exception:
            pass
        return {}

    def getSliceCategoryExtrasForDatabase(self):
        try:
            return self.__sliceCategoryExtrasD
        except Exception:
            pass
        return {}

    def getSliceUnitCardinalityForDatabase(self):
        try:
            return self.__sliceUnitCardinalityD
        except Exception:
            pass
        return {}

    def __getContentClasses(self, catName, atName=None, wildCardAtName="__all__"):
        """Return a list of contexts for input category and optional attribute.

        Handle the special case of unspecified attributes interpreted as wildcard.

        Return:
          contextList (list): list of context names
        """
        cL = []
        try:
            if atName is None:
                cL = self.__categoryContentClasses[catName] if catName in self.__categoryContentClasses else []
            elif (catName, wildCardAtName) in self.__attributeContentClasses:
                cL = self.__attributeContentClasses[(catName, wildCardAtName)]
            else:
                cL = self.__attributeContentClasses[(catName, atName)] if (catName, atName) in self.__attributeContentClasses else []
        except Exception as e:
            logger.exception("Failing catName %s atName %s with %s", catName, atName, str(e))
        #
        return cL

    def __getSliceChildren(self, sliceParentD):
        """Internal method to build data structure containing the parent-child relationships for the
        input slice parent construction.

        """
        retD = OrderedDict()
        for sliceName, sliceParents in sliceParentD.items():
            sD = OrderedDict()
            for pD in sliceParents:
                parentCategoryName = pD["CATEGORY_NAME"]
                parentAttributeName = pD["ATTRIBUTE_NAME"]
                #
                sD[parentCategoryName] = [{"PARENT_CATEGORY_NAME": parentCategoryName, "PARENT_ATTRIBUTE_NAME": parentAttributeName, "CHILD_ATTRIBUTE_NAME": parentAttributeName}]
                #
                # childItems = self.__dApi.getFullChildList(parentCategoryName, parentAttributeName)
                childItems = self.__dApi.getFullDescendentList(parentCategoryName, parentAttributeName)
                # logger.info("Slice parent %s %s  %r" % (parentCategoryName, parentAttributeName, childItems))
                for childItem in childItems:
                    atName = CifName.attributePart(childItem)
                    catName = CifName.categoryPart(childItem)
                    # Ignore children in the parent category
                    if catName == parentCategoryName:
                        continue
                    if catName not in sD:
                        sD[catName] = []
                    sD[catName].append({"PARENT_CATEGORY_NAME": parentCategoryName, "PARENT_ATTRIBUTE_NAME": parentAttributeName, "CHILD_ATTRIBUTE_NAME": atName})
                # Sort the list of dictionaries for each category
                for catName in sD:
                    sD[catName] = sorted(sD[catName], key=lambda k: (k["PARENT_CATEGORY_NAME"], k["PARENT_ATTRIBUTE_NAME"], k["CHILD_ATTRIBUTE_NAME"]))

            retD[sliceName] = sD
        return retD

    def __getIterables(self, iTypeCodes, iQueryStrings, dictHelper):
        itD = OrderedDict()
        #
        typD = {d["TYPE_CODE"]: d["DELIMITER"] for d in iTypeCodes}

        for catName in self.__categoryList:
            for atName in self.__categorySchema[catName]:
                typeCode = self.__dApi.getTypeCode(catName, atName)
                typeCodeAlt = self.__dApi.getTypeCodeAlt(catName, atName)
                if typeCode in typD or typeCodeAlt in typD:
                    itD[(catName, atName)] = typD[typeCode]
                    continue
                description = self.__dApi.getDescription(catName, atName)
                for qs in iQueryStrings:
                    if qs in description:
                        itD[(catName, atName)] = dictHelper.getDelimiter(catName, atName)
        logger.debug("iterableD: %r", list(itD.items()))
        return itD

    def __getEmbeddedIterables(self, iTypeCodes):
        itD = OrderedDict()
        #
        typD = {d["TYPE_CODE"]: d["DELIMITER"] for d in iTypeCodes}

        for catName in self.__categoryList:
            for atName in self.__categorySchema[catName]:
                typeCode = self.__dApi.getTypeCode(catName, atName)
                typeCodeAlt = self.__dApi.getTypeCodeAlt(catName, atName)
                if typeCode in typD or typeCodeAlt in typD:
                    itD[(catName, atName)] = typD[typeCode]
                    continue
        logger.debug("embedded iterableD: %r", list(itD.items()))
        return itD

    def __getMethodInfo(self):
        methodD = OrderedDict()
        methodIndex = self.__dApi.getMethodIndex()
        for _, mrL in methodIndex.items():
            for mr in mrL:
                mId = mr.getId()
                catName = mr.getCategoryName()
                atName = mr.getAttributeName()
                mType = mr.getType()
                logger.debug("mId %r catName %r atName %r mType %r", mId, catName, atName, mType)
                if (catName, atName) not in methodD:
                    methodD[(catName, atName)] = []
                methDef = self.__dApi.getMethod(mId)
                if methDef:
                    mLang = methDef.getLanguage()
                    mCode = methDef.getCode()
                    mImplement = methDef.getImplementation()
                    dD = {"METHOD_LANGUAGE": mLang, "METHOD_IMPLEMENT": mImplement, "METHOD_TYPE": mType, "METHOD_CODE": mCode}
                    methodD[(catName, atName)].append(dD)
                else:
                    logger.error("Missing method definition for %s", mId)
        ##
        logger.debug("Method dictionary %r", methodD)
        return methodD

    def getMethodImplementation(self, catName, atName, methodCodes=None):
        try:
            mC = methodCodes if methodCodes else ["calculation"]
            if (catName, atName) in self.__methodD:
                for mD in self.__methodD[(catName, atName)]:
                    if mD["METHOD_CODE"].lower() in mC:
                        return mD["METHOD_IMPLEMENT"]
            return None

        except Exception as e:
            logger.exception("Failing catName %s atName %s with %s", catName, atName, str(e))
        return None

    def getSelectionFilters(self):
        """
        Returns:  relevant data selection filters for in a dictionary by filter name.
        """
        return self.__dataSelectFilterD

    def getCategories(self):
        """
        Returns:  list of dictionary categories
        """
        #
        return self.__categoryList

    def getSchemaNames(self):
        return self.__categorySchema

    def getCategoryAttributes(self, catName):
        """

        Returns: list of attributes in the input category
        """
        return self.__categorySchema[catName]

    def getCategoryFeatures(self, catName):
        try:
            return self.__categoryFeatures[catName]
        except Exception as e:
            logger.error("Missing category %s %s", catName, str(e))
        return {}

    def getAttributeFeatures(self, catName):
        try:
            return self.__attributeFeatures[catName]
        except Exception as e:
            logger.error("Missing category %s %s", catName, str(e))
        return {}

    def getAttributeDataTypeD(self):
        return self.__attributeDataTypeD

    def getSliceNames(self):
        try:
            return list(self.__sliceD.keys())
        except Exception as e:
            logger.error("Failing with %s", str(e))
        return []

    def getSliceAttributes(self, sliceName, catName):
        try:
            return self.__sliceD[sliceName][catName]
        except Exception as e:
            logger.debug("Failing to access sliceName %s catName %s with %s", sliceName, catName, str(e))
        return []

    def getSliceCategories(self, sliceName):
        try:
            return list(self.__sliceD[sliceName].keys())
        except Exception as e:
            logger.debug("Failing to access sliceName %s with %s", sliceName, str(e))
        return []

    def __getAttributeTypeD(self, catName):
        aD = {}
        for atName in self.__categorySchema[catName]:
            aD[atName] = self.__dApi.getTypeCode(catName, atName)
        return aD

    def __assignEnumTypes(self, enumList, pT):
        rL = []
        if enumList:
            if pT != "numb":
                return enumList
            else:
                isFloat = False
                for enum in enumList:
                    if "." in enum:
                        isFloat = True
                        break
                if isFloat:
                    # a rare case
                    rL = [float(enum) for enum in enumList]
                else:
                    rL = [int(enum) for enum in enumList]
                return rL
        else:
            return enumList

    def __assignExampleTypes(self, catName, atName, examList, pT):

        try:
            rL = []
            if examList:
                if pT != "numb":
                    return examList
                else:
                    isFloat = False
                    for enum in examList:
                        if "." in enum:
                            isFloat = True
                            break
                    if isFloat:
                        # a rare case
                        rL = [float(enum) for enum in examList]
                    else:
                        rL = [int(enum) for enum in examList]
                    return rL

            else:
                return examList
        except Exception:
            logger.exception("catName %r atName %r primitive %r failing for examList %r", catName, atName, pT, examList)
        return examList

    def __assignExampleTupTypes(self, catName, atName, examTupList, pT):
        try:
            rL = []
            if examTupList:
                if pT != "numb":
                    for exam, detail in examTupList:
                        exam = str(exam).strip() if exam else exam
                        detail = str(detail).strip() if detail else detail
                        rL.append((exam, detail))
                    return rL
                else:
                    isFloat = False
                    for exam, _ in examTupList:
                        if "." in exam:
                            isFloat = True
                            break
                    if isFloat:
                        for exam, detail in examTupList:
                            try:
                                if isFloat:
                                    rL.append((float(exam), detail))
                                else:
                                    rL.append((int(exam), detail))
                            except Exception:
                                pass
                    return rL
            else:
                return examTupList
        except Exception as e:
            logger.exception("catName %r atName %r (%r) %r %s", catName, atName, pT, examTupList, str(e))
        #
        return examTupList

    def __assignEnumTupTypes(self, enumTupList, pT):
        rL = []
        if enumTupList:
            if pT != "numb":
                return enumTupList
            else:
                isFloat = False
                for enum, _, _, _ in enumTupList:
                    if "." in enum:
                        isFloat = True
                        break
                if isFloat:
                    for enum, detail, brief, units in enumTupList:
                        try:
                            if isFloat:
                                rL.append((float(enum), detail, brief, units))
                            else:
                                rL.append((int(enum), detail, brief, units))
                        except Exception:
                            pass
                return rL
        else:
            return enumTupList

    def __hasEnumDetails(self, enumTupList):
        dCount = 0
        for _, detail, brief, units in enumTupList:
            if detail or brief or units:
                dCount += 1
        return dCount > 0

    def __itemNameToDictList(self, itemNameList):
        rL = []
        for itemName in list(OrderedDict.fromkeys(itemNameList)):
            atName = CifName.attributePart(itemName)
            catName = CifName.categoryPart(itemName)
            rL.append({"CATEGORY": catName, "ATTRIBUTE": atName})
        return rL

    def __getAttributeFeatures(self, catName, iterableD, embeddedIterableD, itemTransformD, methodD):
        """
        Args:
            catName (string): Category name
            iterableD (tuple, optional): iterable dictionary type codes
            iQueryStrings (list, optional): search strings applied to item descriptions to identify iterable candidates
            itemTransformD (dict): dictionary of data transform filters  itd[(catName,atName)] = [f1,f2,...]

        Returns:
            dict: attribute features

             cL = self.getCategoryContextList(catName)
        """
        aD = {}

        #
        # keyAtNames = [CifName.attributePart(kyItem) for kyItem in self.__dApi.getCategoryKeyList(catName)]
        keyAtNames = [CifName.attributePart(kyItem) for kyItem in self.__getCategoryKeysWithReplacement(catName)]
        for atName in self.__categorySchema[catName]:
            itemName = CifName.itemName(catName, atName)
            fD = {
                "CATEGORY_NAME": catName,
                "ATTRIBUTE_NAME": atName,
                "TYPE_CODE": None,
                "TYPE_CODE_ALT": None,
                "IS_MANDATORY": False,
                "CHILD_ITEMS": [],
                "CHILDREN": [],
                "ROOT_PARENT_ITEM": None,
                "ROOT_PARENT": None,
                "PARENT": None,
                "DESCRIPTION": None,
                "DESCRIPTION_ANNOTATED": [],
                "IS_KEY": False,
                "ITERABLE_DELIMITER": None,
                "EMBEDDED_ITERABLE_DELIMITER": None,
                "FILTER_TYPES": [],
                "IS_CHAR_TYPE": False,
                "METHODS": [],
                "CONTENT_CLASSES": [],
                "UNITS": None,
                "ENUMS": None,
                "ENUMS_ANNOTATED": None,
                "SEARCH_CONTEXTS": None,
            }
            fD["TYPE_CODE"] = self.__dApi.getTypeCode(catName, atName)
            fD["TYPE_CODE_ALT"] = self.__dApi.getTypeCodeAlt(catName, atName)
            fD["IS_MANDATORY"] = True if str(self.__dApi.getMandatoryCode(catName, atName)).lower() in ["y", "yes"] else False
            fD["DESCRIPTION"] = textwrap.dedent(self.__dApi.getDescription(catName, atName)).lstrip().rstrip()
            #
            fD["DESCRIPTION_ANNOTATED"] = [{"text": fD["DESCRIPTION"], "context": "dictionary"}]
            tS = self.__dApi.getDescriptionPdbx(catName, atName)
            if tS:
                fD["DESCRIPTION_ANNOTATED"].append({"text": textwrap.dedent(tS).lstrip().rstrip(), "context": "deposition"})
            #
            fD["UNITS"] = self.__dApi.getUnits(catName, atName)

            #
            fD["CHILD_ITEMS"] = self.__dApi.getFullChildList(catName, atName)
            fD["CHILDREN"] = self.__itemNameToDictList(self.__dApi.getFullChildList(catName, atName))
            #
            pItemName = self.__dApi.getUltimateParent(catName, atName)
            pName = pItemName if pItemName != itemName else None
            fD["ROOT_PARENT_ITEM"] = pName

            fD["ROOT_PARENT"] = self.__itemNameToDictList([pName])[0] if pName else None
            #
            pL = self.__dApi.getFullParentList(catName, atName, stripSelfParent=True)
            if pL:
                rL = self.__itemNameToDictList(pL)
                fD["PARENT"] = rL[0] if rL else None
                if len(rL) > 1:
                    logger.warning("Unexpected multiple parent definition for %s %s : %r", catName, atName, rL)
            #
            # logger.debug("catName %s atName %s : parent %r root_parent %r", catName, atName, fD['PARENT'], fD['ROOT_PARENT'])
            #
            fD["IS_KEY"] = atName in keyAtNames
            pType = self.__dApi.getTypePrimitive(catName, atName)
            fD["IS_CHAR_TYPE"] = str(pType).lower() in ["char", "uchar"]
            #
            fD["ITERABLE_DELIMITER"] = iterableD[(catName, atName)] if (catName, atName) in iterableD else None
            fD["EMBEDDED_ITERABLE_DELIMITER"] = embeddedIterableD[(catName, atName)] if (catName, atName) in embeddedIterableD else None
            #
            fD["FILTER_TYPES"] = itemTransformD[(catName, "__all__")] if (catName, "__all__") in itemTransformD else []
            fD["FILTER_TYPES"] = itemTransformD[(catName, atName)] if (catName, atName) in itemTransformD else fD["FILTER_TYPES"]
            #
            fD["METHODS"] = methodD[(catName, atName)] if (catName, atName) in methodD else []
            fD["CONTENT_CLASSES"] = self.__getContentClasses(catName, atName)
            if (catName, atName) in self.__intEnumD:
                fD["ENUMS"] = sorted(self.__assignEnumTypes(self.__dApi.getEnumListPdbx(catName, atName), pType))
                logger.debug("Using internal enums for %s %s %d", catName, atName, len(fD["ENUMS"]))
                enumTupList = self.__dApi.getEnumListAltWithFullDetails(catName, atName)
            else:
                fD["ENUMS"] = sorted(self.__assignEnumTypes(self.__dApi.getEnumList(catName, atName), pType))
                enumTupList = self.__dApi.getEnumListWithFullDetails(catName, atName)
            #
            if self.__hasEnumDetails(enumTupList):
                #
                fD["ENUMS_ANNOTATED"] = []
                for eTup in self.__assignEnumTupTypes(enumTupList, pType):
                    teD = {"value": eTup[0]}
                    if eTup[1]:
                        teD["detail"] = eTup[1]
                    if eTup[2]:
                        teD["name"] = eTup[2]
                    if eTup[3]:
                        teD["units"] = eTup[3]
                    fD["ENUMS_ANNOTATED"].append(teD)
            # -----
            fD["EXAMPLES"] = self.__assignExampleTupTypes(catName, atName, self.__dApi.getExampleListPdbx(catName, atName), pType)
            fD["EXAMPLES"].extend(self.__assignExampleTupTypes(catName, atName, self.__dApi.getExampleList(catName, atName), pType))
            # -----
            scL = []
            for scTup in self.__dApi.getItemSubCategoryList(catName, atName):
                if scTup[1] is not None:
                    qD = {"id": scTup[0], "label": scTup[1]}
                else:
                    qD = {"id": scTup[0]}
                scL.append(qD)
            fD["SUB_CATEGORIES"] = scL
            if len(scL) > 1:
                logger.debug("Multiple subcategories for %r %r %r", catName, atName, scL)
            #
            # bList = self.__dApi.getBoundaryListAlt(catName, atName, fallBack=True)
            bdList = self.__dApi.getBoundaryList(catName, atName)
            if bdList:
                minD = {}
                maxD = {}
                for (minV, maxV) in bdList:
                    if minV == maxV:
                        continue
                    if minV not in [".", "?"]:
                        minD[minV] = False
                    if maxV not in [".", "?"]:
                        maxD[maxV] = False
                for (minV, maxV) in bdList:
                    if minV == maxV and minV in minD:
                        minD[minV] = True
                    if minV == maxV and maxV in maxD:
                        maxD[maxV] = True
                for ky in minD:
                    if "." in ky:
                        kyV = float(ky)
                    else:
                        kyV = int(ky)
                    if minD[ky]:
                        fD["MIN_VALUE"] = kyV
                    else:
                        fD["MIN_VALUE_EXCLUSIVE"] = kyV
                for ky in maxD:
                    if "." in ky:
                        kyV = float(ky)
                    else:
                        kyV = int(ky)
                    if maxD[ky]:
                        fD["MAX_VALUE"] = kyV
                    else:
                        fD["MAX_VALUE_EXCLUSIVE"] = kyV
            #
            aD[atName] = fD
        #
        return aD

    def __getCategoryFeatures(self, catName, unitCardinalityList, subCategoryD):
        cD = {"KEY_ATTRIBUTES": []}
        # cD['KEY_ATTRIBUTES'] = [CifName.attributePart(keyItem) for keyItem in self.__dApi.getCategoryKeyList(catName)]
        cD["KEY_ATTRIBUTES"] = [CifName.attributePart(keyItem) for keyItem in self.__getCategoryKeysWithReplacement(catName)]
        cD["UNIT_CARDINALITY"] = catName in unitCardinalityList
        cD["CONTENT_CLASSES"] = self.__getContentClasses(catName)
        #
        # Exclude all categories beginning with "ma_" from being mandatory
        # (temporarily hardcoded here until new configuration file section added to achieve same effect)
        cD["IS_MANDATORY"] = True if str(self.__dApi.getCategoryMandatoryCode(catName)).lower() == "yes" and not catName.startswith("ma_") else False
        # cD["IS_MANDATORY"] = True if str(self.__dApi.getCategoryMandatoryCode(catName)).lower() == "yes" else False
        #
        cD["SUB_CATEGORIES"] = subCategoryD[catName] if catName in subCategoryD else []
        #
        return cD

    def __getUnitCardinalityCategories(self, parentDList):
        """Assign categories with unit cardinality relative to the input list of parent key items.

        parentDList (dict):  [{'CATEGORY_NAME':xxx 'ATTRIBUTE_NAME': xxxx}]

        Return: category name list
        """
        numParents = len(parentDList)
        logger.debug("Parent slice count %d def %r", numParents, parentDList)
        ucL = []
        #
        #  Find the common set of child categories for the input parent items
        comCatList = []
        for pD in parentDList:
            catList = [pD["CATEGORY_NAME"]]
            for childItem in self.__dApi.getFullChildList(pD["CATEGORY_NAME"], pD["ATTRIBUTE_NAME"]):
                childCategoryName = CifName.categoryPart(childItem)
                primaryKeyItemList = self.__dApi.getCategoryKeyList(childCategoryName)
                logger.debug("child category %r primary key items  %r", childCategoryName, primaryKeyItemList)
                # child must be part of the primary key to be a candidate
                if childItem in primaryKeyItemList:
                    catList.append(childCategoryName)
            if comCatList:
                comCatList = list(set(catList) & set(comCatList))
            else:
                comCatList.extend(catList)
        logger.debug("Common category list %r", comCatList)
        for cat in comCatList:
            primaryKeyItemList = self.__dApi.getCategoryKeyList(cat)
            if len(primaryKeyItemList) == numParents:
                ucL.append(cat)
        #
        logger.debug("Slice unit cardinality categories from parent-child relationships %r", ucL)
        return sorted(ucL)

    #

    def __getCategoryKeysWithReplacement(self, categoryName):
        if categoryName in self.__keyReplaceCategoryD:
            keyItems = [CifName.itemName(categoryName, atName) for atName in self.__keyReplaceCategoryD[categoryName]]
        else:
            keyItems = self.__dApi.getCategoryKeyList(categoryName)
        return sorted(keyItems)

    def __indexContexts(self, dictSchema):
        """Extract the category an item level dictionary contexts."""
        catIndex = {}
        atIndex = {}
        for catName in dictSchema:
            for cT in self.__dApi.getCategoryContextList(catName):
                catIndex.setdefault(cT, []).append(catName)

            for atName in dictSchema[catName]:
                for cT in self.__dApi.getContextList(catName, atName):
                    atIndex.setdefault(cT, []).append((catName, atName))
        return catIndex, atIndex


import gym
from gym import spaces
from gym.utils import seeding
import numpy as np
import itertools

class TradingEnv(gym.Env):

	def __init__(self, train_data, init_invest=20000):
		self.stock_price_history = np.around(train_data)
		self.n_stock, self.n_step = self.stock_price_history.shape

		self.init_invest = init_invest
		self.cur_step = None
		self.stock_owned = None
		self.stock_price = None
		self.cash_in_hand = None

		self.action_space = spaces.Discrete(3**self.n_stock)

		stock_max_price = self.stock_price_history.max(axis=1)
		stock_range = [[0, init_invest * 2 // mx] for mx in stock_max_price]
		price_range = [[0, mx] for mx in stock_max_price]
		cash_in_hand_range = [[0, init_invest*2]]
		self.observation_space = spaces.MultiDiscrete(stock_range + price_range + cash_in_hand_range)

		self._seed()
		self._reset()

	def _seed(self, seed=None):
		self.np_random, seed = seeding.np_random(seed)
		return [seed]

	def _reset(self):
		self.cur_step = 0
		self. stock_owned = [0] * self.n_stock
		self.stock_price = self.stock_price_history[:, self.cur_step]
		self.cash_in_hand = self.init_invest
		return self._get_obs()

	def _step(self, action):
		assert self.action_space.contains(action)
		prev_val = self._get_val()
		self.cur_step += 1
		self.stock_price = self.stock_price_history[:, self.cur_step]
		self._trade()
		cur_val = self._get_val()
		reward = cur_val - prev_val
		done = self.cur_step == self.n_step - 1
		info = {'cur_val': cur_val}
		return self._get_obs(), reward, done, info

	def _get_obs(self):
		obs = []
		obs.extend(self.stock_owned)
		obs.extend(list(self.stock_price))
		obs.append(self.cash_in_hand)
		return obs

	def _get_val(self):
		return np.sum(self.stock_owned*self.stock_price) + self.cash_in_hand

	def _trade(self, action):
		action_combo = map(list, itertools.product([0,1,2], repeat=self.n_stock))
		action_vec = action_combo[action]

		sell_index = []
		buy_index = []
		for i, a in enumerate(action_vec):
			if a == 0:
				sell_index.append(i)
			elif a == 2:
				buy_index.append(i)

		if sell_index:
			for i in sell_index:
				self.cash_in_hand += self.stock_price[i]*self.stock_owned[i]
				self.stock_owned[i] = 0
			if buy_index:
				can_buy = True
				while can_buy:
					for i in buy_index:
						if self.cash_in_hand > self.stock_price[i]:
							self.stock_owned[i] += 1
							self.cash_in_hand -= self.stock_price[i]
						else:
							can_buy = False


# import h5py
import numpy
import sys

from volumina.api import Viewer
# from PyQt4.QtGui import QApplication, QColor

from lazyflow.graph import Graph
# from lazyflow.operators.ioOperators.opStreamingHdf5Reader import OpStreamingHdf5Reader
# from lazyflow.operators import OpCompressedCache

from PyQt4 import QtCore, uic
from PyQt4.QtGui import QApplication, QPainter, QPainterPath
from PyQt4 import QtGui
from PyQt4 import Qt

# class testWidget(QtGui.QWidget):
#
#     def __init__(self):
#         super(testWidget, self).__init__()
#
#
# app = QtGui.QApplication(sys.argv)
# main_window = QtGui.QMainWindow
# main_window.show()
# sys.exit(app.exec_())

# w = testWidget()

class testWidget(QtGui.QWidget):

    def __init__(self):
        super(testWidget, self).__init__()
        self.initUI()

    def initUI(self):
        self.setGeometry(300, 300, 250, 150)
        self.setWindowTitle('Test Widget')
        self.show()

    def paintEvent(self, QPaintEvent):

        back = (numpy.ones((100, 200, 300)) * 0).astype(numpy.uint8)
        back[0:60, 0:60, 0:60] = 255
        ol = (numpy.zeros((100, 200, 300))).astype(numpy.uint8)
        ol[0:99, 0:99, 0:99] = 120
        ol[0:99, 120:140, 0:99] = 255

        path = QPainterPath()
        path.addRect(20, 20, 60, 60)
        path.moveTo(0, 0)
        path.cubicTo(99, 0,  50, 50,  99, 99)
        path.cubicTo(0, 99,  50, 50,  0, 0)
        painter = QPainter(self)
        painter.fillRect(0, 0, 100, 100, Qt.Qt.red)

def main():

    app = QtGui.QApplication(sys.argv)

    w = testWidget()

    sys.exit(app.exec_())

if __name__ == '__main__':
    main()
else:
    main()


from game import *

class player():
	pos = geoma.point(0, 0)
	def __init__(self):
		self.pos.set(0, 0)
	def set(self, p: geoma.point):
		self.pos = p

class game():
	UNINIT = 0
	TURN = 1
	WIN = 2

	dx = 100
	dy = 100
	tek_map: map = map()
	players = []
	players_color = [pygame.Color("Brown"), pygame.Color("Magenta"), pygame.Color("White")]
	tek_player = 0
	state = UNINIT
	surface: pygame.Surface
	screen: pygame.display

	tek_pt: geoma.point

	def __init__(self, screen_init: pygame.display, sur: pygame.Surface):
		self.state = self.UNINIT
		self.screen = screen_init
		self.surface = sur

	def load_map(self):
		self.tek_map.load()

	def turn(self):
		self.state = self.TURN
		while (self.state != self.WIN or self.state != self.LOSE):
			if (self.state != self.TURN):
				break
			self.turn()
			# break

	def draw1(self):
		self.surface.fill(pygame.Color("WHITE"))
		self.tek_map.draw(self.surface, GREEN, RED, BLUE, 3, 3, 3)
		print(len(self.players))
		for i in range(len(self.players)):
			# print(self.players[i])
			pygame.draw.circle(self.surface, self.players_color[i], self.players[i].pos.to_arr(), 5)
		if (len(self.players) < 2):	pygame.draw.circle(self.surface, self.players_color[len(self.players)], self.tek_pt.to_arr(), 5)
		# if self.mode == 1 and self.sel_reb != []:
		# 	self.draw_sel_reb()
		# if self.mode == 1:
		# 	tek = pygame.mouse.get_pos()
		# 	num = self.find_nearest_reb(geoma.point(tek[0] - self.dx, tek[1] - self.dy))
		# 	self.draw_reb(num, self.sel_time_color)
		# if (self.mode == 1 and self.is_obl):
		# 	a = min(self.point1.x, self.point2.x)
		# 	b = min(self.point1.y, self.point2.y)
		# 	w = max(self.point1.x, self.point2.x) - a
		# 	h = max(self.point1.y, self.point2.y) - b
		# 	pygame.draw.rect(self.surface, self.rect_color, pygame.Rect(a, b, w, h))
		# if self.mode == 0 and self.is_first:
		# 	self.draw_new_reb()
		# if (self.mode == 0 and not(self.is_first)):
		# 	tek = pygame.mouse.get_pos()
		# 	pygame.draw.circle(self.surface, self.sel_time_color, self.get_point(geoma.point(tek[0], tek[1])).to_arr(), 5)
		self.screen.blit(self.surface, (self.dx, self.dy))
		pygame.display.update()

	def get_nearest(self, pt: geoma.point):
		dist = float("Inf")
		ans = -1
		for i in range(len(self.tek_map.start)):
			tek_dist = geoma.dist_pt_otr(pt, self.tek_map.start[i])
			if (dist - tek_dist >= geoma.eps):
				ans = i
				dist = tek_dist
		if (ans == -1):	return pt
		l = geoma.line(self.tek_map.start[ans].a, self.tek_map.start[ans].b)
		perp_line = geoma.perp_line_pt(l, pt)
		# return pt
		pt_inter = geoma.line_inter(perp_line, l)
		val = geoma.inside_pt_otr(pt, self.tek_map.start[ans])
		if (val == 0):
			return pt_inter
		elif (val == 1):
			return self.tek_map.start[ans].a
		elif (val == 2):
			return self.tek_map.start[ans].b

	def get_mouse_pos(self):
		a = pygame.mouse.get_pos()
		tek_pt = geoma.point(a[0] - self.dx, a[1] - self.dy)
		tek_pt = self.get_nearest(tek_pt)
		return tek_pt

	def get_point(self):
		run = True
		while run:
			self.tek_pt = self.get_nearest(self.get_mouse_pos())
			self.draw1()
			for event in pygame.event.get():
				if event.type == pygame.QUIT:
					exit()
				if event.type == pygame.KEYDOWN:
					if (event.key == pygame.K_ESCAPE):
						exit()
				if event.type == pygame.MOUSEBUTTONUP:
					return self.tek_pt

	def start_new_game(self):
		# load map
		self.load_map()

		# set start point for players
		self.players.clear()
		for i in range(2):
			a = player()
			a.set(self.get_point())
			self.players.append(a)
			self.draw1()

		# game_process
		while True:
			for event in pygame.event.get():
				if event.type == pygame.QUIT:
					exit()
				if event.type == pygame.KEYDOWN:
					if (event.key == pygame.K_ESCAPE):
						exit()
		# end

	def start(self):
		self.start_new_game()

		# exit()
		# check for exit
		# while True:
		# 	for event in pygame.event.get():
		# 		if i.type == pygame.QUIT:
		# 			exit()
		# 		if event.type == pygame.KEYDOWN:
		# 			if (event.key == pygame.K_ESCAPE):
		# 				exit()

if __name__ == "__main__":
	pygame.init()
	screen = pygame.display.set_mode([1000, 1000])
	screen.fill(pygame.Color("lightgray"))
	pygame.display.update()

	tek_surf = pygame.Surface((800, 800))
	tek_game = game(screen, tek_surf)
	tek_game.start()
	pygame.quit()


def DATA2JSON(X, Y, filename, str_format = '05d', digits = 3):
    n = len(X)
    data = dict()
    for i in range(n):
        data[format(i, str_format)] = {'X':round(X[i], digits), 'Y':round(Y[i], digits)}
    with open(filename, "w") as outfile:
        json.dump(data, outfile, indent=4)

def JSON2DATA(filename, str_format = '05d', digits = 3):
    with open(filename) as infile:
        data = load(infile)
    N = len(data)
    X = np.zeros(N)
    Y = np.zeros(N)

    for i in range(N):
        X[i] = data[format(i, str_format)]['X']
        Y[i] = data[format(i, str_format)]['Y']
    return X,Y