import matplotlib
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_style("dark")
plt.rcParams['figure.figsize'] = 16, 12
import pandas as pd
from tqdm import tqdm_notebook, tqdm
import io
from PIL import Image
from glob import glob
from collections import defaultdict
import os
import pickle
from optparse import OptionParser
from datetime import datetime
import json
import sys
import time
from shutil import copyfile
import cv2
cv2.ocl.setUseOpenCL(True)
import random

import imgaug as ia
from imgaug import augmenters as iaa

import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
import torchvision.models as models

from kaggle_camera_model_id_lib.utils import PechkaBot, ImageList, NpzFolder, NCrops, MultiDataset
from kaggle_camera_model_id_lib.models import VggHead, StyleVggHead, IEEEfcn, ResNetFC, ResNetX, FatNet1
from kaggle_camera_model_id_lib.models import InceptionResNetV2fc, InceptionResNetV2fcSmall, InceptionResNetV2
from kaggle_camera_model_id_lib.models import ResNetDense, ResNetDenseFC
from kaggle_camera_model_id_lib.utils import jpg_compress, equalize_v_hist, hsv_convert
from kaggle_camera_model_id_lib.utils import scale_crop_pad, gamma_correction
from kaggle_camera_model_id_lib.utils import patch_quality_dich, n_random_crops, n_pseudorandom_crops

_bot = PechkaBot()

def log(txt):
    print(txt)
    _bot.send_message(txt)

def train_pass(train_loader, model, criterion, optimizer):
    loss_train_batch = 0
    acc_train_batch = 0

    for ix_batch, (X, Y) in tqdm(
        enumerate(train_loader),
        total=int(len(train_loader.dataset.imgs)/batch_size_train),
        desc='Train #%i' % ix_epoch):

        bs, ncrops, c, h, w = X.shape
        X = X.view(-1, c, h, w)
        Y = Y.view(ncrops*bs)

        X_var = Variable(X.cuda())
        Y_var = Variable(Y.cuda())

        log_p = model(X_var)
        loss = criterion(log_p, Y_var)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        loss_train_batch += loss.data[0]
        acc_train_batch += ((log_p.max(1)[1] == Y_var).float().sum()/Y_var.shape[0]).data[0]

        if options.is_debug and ix_batch > 50:
            break

    X_var = X_var.cpu()
    del(X_var)
    Y_var = Y_var.cpu()
    del(Y_var)

    loss_train_batch /= ix_batch + 1
    acc_train_batch /= ix_batch + 1

    return loss_train_batch, acc_train_batch

def val_pass(val_loader, model, criterion):
    loss_val_batch = 0
    acc_val_batch = 0

    for ix_batch, (X, Y) in tqdm(
        enumerate(val_loader),
        total=int(len(val_loader.dataset.imgs)/batch_size_val),
        desc='Val #%i' % ix_epoch):

        bs, ncrops, c, h, w = X.shape
        X = X.view(-1, c, h, w)
        Y = Y.view(ncrops*bs)

        X_var = Variable(X.cuda(), volatile=True)
        Y_var = Variable(Y.cuda(), volatile=True)

        log_p = model(X_var)
        loss = criterion(log_p, Y_var)

        loss_val_batch += loss.data[0]
        acc_val_batch += ((log_p.max(1)[1] == Y_var).float().sum()/Y_var.shape[0]).data[0]

        if options.is_debug and ix_batch > 50:
            break

    X_var = X_var.cpu()
    del(X_var)
    Y_var = Y_var.cpu()
    del(Y_var)

    loss_val_batch /= ix_batch + 1
    acc_val_batch /= ix_batch + 1

    return loss_val_batch, acc_val_batch

model_factory = {
    'Vgg19Head_E_2b_bn': lambda n_classes: VggHead(num_classes=n_classes, vgg_key='E_2b', load_vgg_bn=True, batch_norm=True),
    'Vgg19Head_E_3b_bn': lambda n_classes: VggHead(num_classes=n_classes, vgg_key='E_3b', load_vgg_bn=True, batch_norm=True),
    'Vgg19Head_E_bn': lambda n_classes: VggHead(num_classes=n_classes, load_vgg_bn=True, vgg_key='E', batch_norm=True),
    'Vgg11Head_A_bn': lambda n_classes: VggHead(num_classes=n_classes, load_vgg_bn=True, vgg_key='A', batch_norm=True),
    'Vgg11Head_A': lambda n_classes: VggHead(num_classes=n_classes, load_vgg_bn=True, vgg_key='A', batch_norm=False),
    'StyleVggHead_bn': lambda n_classes: StyleVggHead(num_classes=n_classes, load_vgg_bn=True),
    'IEEEfcn': lambda n_classes: IEEEfcn(n_classes),
    'resnet18fc_pretrained': lambda n_classes: ResNetFC(
        models.resnet.BasicBlock, [2, 2, 2, 2], num_classes=n_classes, load_resnet='resnet18'),
    'resnet18fc': lambda n_classes: ResNetFC(
        models.resnet.BasicBlock, [2, 2, 2, 2], num_classes=n_classes, load_resnet=None),
    'resnet18X_pretrained': lambda n_classes: ResNetX(
        models.resnet.BasicBlock, [2, 2, 2, 2], num_classes=n_classes, load_resnet='resnet18'),
    'InceptionResNetV2fc_5_10_4': lambda n_classes: InceptionResNetV2fc(
        num_classes=n_classes, nun_block35=5, num_block17=10, num_block8=4),
    'InceptionResNetV2fcSmall_5_10': lambda n_classes: InceptionResNetV2fcSmall(
        num_classes=n_classes, nun_block35=5, num_block17=10),
    'resnet34fc_pretrained': lambda n_classes: ResNetFC(
        models.resnet.BasicBlock, [3, 4, 6, 3], num_classes=n_classes, load_resnet='resnet34'),
    'resnet34fc_pretrained_maxpool': lambda n_classes: ResNetFC(
        models.resnet.BasicBlock, [3, 4, 6, 3], num_classes=n_classes, load_resnet='resnet34', pool_type='max'),
    'resnet50fc_pretrained': lambda n_classes: ResNetFC(
        models.resnet.Bottleneck, [3, 4, 6, 3], num_classes=n_classes, load_resnet='resnet50'),
    'FatNet1': lambda n_classes: FatNet1(n_classes),
    'resnet34X_pretrained_maxpool': lambda n_classes: ResNetX(
        models.resnet.BasicBlock, [3, 4, 6, 3], num_classes=n_classes, load_resnet='resnet34', pool_type='max'),
    'resnet50X_pretrained_maxpool': lambda n_classes: ResNetX(
        models.resnet.Bottleneck, [3, 4, 6, 3], num_classes=n_classes, load_resnet='resnet50', pool_type='max'),
    'InceptionResNetV2': lambda n_classes: InceptionResNetV2(num_classes=n_classes),
    'ResNetDense34_pretrained': lambda n_classes: ResNetDense(
        models.resnet.BasicBlock, [3, 4, 6, 3], num_classes=n_classes, load_resnet='resnet34'),
    'ResNetDenseFC34_pretrained': lambda n_classes: ResNetDenseFC(
        models.resnet.BasicBlock, [3, 4, 6, 3], num_classes=n_classes, load_resnet='resnet34',
        zero_first_center=False),
    'ResNetDenseFC34_pretrained_zfc': lambda n_classes: ResNetDenseFC(
        models.resnet.BasicBlock, [3, 4, 6, 3], num_classes=n_classes, load_resnet='resnet34',
        zero_first_center=True)
}

def create_CELoss(prms):
    if prms is None:
        return nn.CrossEntropyLoss()
    if 'weight' in prms:
        prms['weight'] = torch.FloatTensor(prms['weight'])
    return nn.CrossEntropyLoss(**prms)

criterion_factory = {
    'CrossEntropyLoss': lambda prms: create_CELoss(prms),
    'MultiMarginLoss': lambda prms: nn.MultiMarginLoss(**prms)
}

if __name__ == '__main__':
    parser = OptionParser()
    parser.add_option('-c',
                      '--config',
                      dest='cfg_path',
                      help='config path')
    parser.add_option('-d',
                      '--debug',
                      action="store_true",
                      dest="is_debug")
    (options, args) = parser.parse_args()
    if options.cfg_path is None:
        sys.exit('cfg_path is not provided')
    if options.is_debug:
        log('DEBUG MODE ON')
    log('-----------\n\nStarting training process: \n  %s\n  %s' % (str(datetime.now()), __file__))
    log('config: %s' % options.cfg_path)

    with open(options.cfg_path) as f:
        cfg = json.load(f)

    log('Config:')
    for k, v in cfg.items():
        log('  %s = %s' % (k, v))

    train_list_path = cfg['train_list_path']
    val_path = cfg['val_path']
    out_dir = cfg['out_dir']
    model_path = cfg['model_path']
    crop_size = cfg['crop_size']
    step_crop_val = cfg['step_crop_val']
    n_crops_train = cfg['n_crops_train']
    batch_size_train = cfg['batch_size_train']
    batch_size_val = cfg['batch_size_val']
    workers = cfg['workers']
    n_epoches = cfg['n_epoches']
    model_type = cfg['model_type']
    n_classes = cfg['n_classes']
    learning_rate = cfg['learning_rate']
    momentum = cfg['momentum']
    lr_scheduler_step_size = cfg['lr_scheduler_step_size']
    lr_scheduler_gamma = cfg['lr_scheduler_gamma']
    weight_decay = cfg['weight_decay']
    optim_type = cfg['optim_type']
    crop_center_size = cfg['crop_center_size']
    do_random_aug_kaggle = cfg['do_random_aug_kaggle']
    p_random_aug_kaggle_train = cfg['p_random_aug_kaggle_train']
    p_random_aug_kaggle_val = cfg['p_random_aug_kaggle_val']
    do_hard_aug = cfg['do_hard_aug']
    p_hard_aug_train = cfg['p_hard_aug_train']
    p_hard_aug_val = cfg['p_hard_aug_val']
    criterion_type = cfg['criterion_type']
    criterion_params = cfg['criterion_params']
    n_crops_search_train = cfg['n_crops_search_train']
    train_list_pseudo_npz = cfg['train_list_pseudo_npz']
    train_list_flickr = cfg['train_list_flickr']

    to_tensor = transforms.ToTensor()
    normalize = transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
    random_crop = transforms.RandomCrop(crop_size)
    center_crop = transforms.CenterCrop(crop_center_size)
    rvf = transforms.RandomVerticalFlip()
    rhf = transforms.RandomHorizontalFlip()
    random_flip = lambda img: rvf(rhf(img))

    scale_05 = lambda img: scale_crop_pad(img, 0.5)
    scale_08 = lambda img: scale_crop_pad(img, 0.8)
    scale_15 = lambda img: scale_crop_pad(img, 1.5)
    scale_20 = lambda img: scale_crop_pad(img, 2.0)
    gamma_08 = lambda img: gamma_correction(img, 0.8)
    gamma_12 = lambda img: gamma_correction(img, 1.2)
    jpg_70 = lambda img: jpg_compress(img, (70, 71))
    jpg_90 = lambda img: jpg_compress(img, (90, 91))
    augs = [scale_05, scale_08, scale_15, scale_20, gamma_08, gamma_12, jpg_70, jpg_90]

    def random_aug_kaggle(img, p=0.5):
        if np.random.rand() < p:
            return random.choice(augs)(img)
        return img

    blur = iaa.GaussianBlur(sigma=(0, 2))
    sharpen = iaa.Sharpen(alpha=(0, 1), lightness=(0.5, 2))
    emboss = iaa.Emboss(alpha=(0, 1), strength=(0, 2))
    contrast_normalization = iaa.ContrastNormalization(alpha=(0.7, 1.3))
    hard_aug = iaa.OneOf([blur, sharpen, emboss, contrast_normalization])
    sometimes_train = iaa.Sometimes(p_hard_aug_train, hard_aug)
    sometimes_val = iaa.Sometimes(p_hard_aug_val, hard_aug)

    def aug_train(img):
        #if min(img.size) > crop_center_size:
        #    return random_flip(random_crop(center_crop(img)))
        #img_np = np.array(img)
        #if img_np.shape[0] < crop_center_size and img_np.shape[1] > crop_center_size:
        #    n = np.random.randint(img_np.shape[1] - crop_center_size)
        #    return random_flip(random_crop(Image.fromarray(img_np[:, n:(n + crop_center_size), :])))
        #if img_np.shape[1] < crop_center_size and img_np.shape[0] > crop_center_size:
        #    n = np.random.randint(img_np.shape[0] - crop_center_size)
        #    return random_flip(random_crop(Image.fromarray(img_np[n:(n + crop_center_size), :, :])))
        return random_flip(random_crop(img))

    def aug_train_fscore(img):
        if min(img.size) > crop_center_size:
            img_np = np.array(center_crop(img))
        else:
            img_np = np.array(img)
            if img_np.shape[0] < crop_center_size and img_np.shape[1] > crop_center_size:
                n = np.random.randint(img_np.shape[1] - crop_center_size)
                img_np = img_np[:, n:(n + crop_center_size), :]
            if img_np.shape[1] < crop_center_size and img_np.shape[0] > crop_center_size:
                n = np.random.randint(img_np.shape[0] - crop_center_size)
                img_np = img_np[n:(n + crop_center_size), :, :]

        crops = n_pseudorandom_crops(img_np, crop_size, n_crops_train, n_crops_search_train, patch_quality_dich)
        for img in crops:
            yield random_flip(random_crop(Image.fromarray(img)))

    def aug_optional_train(img):
        if do_hard_aug:
            img = Image.fromarray(sometimes_train.augment_image(np.array(img)))

        if do_random_aug_kaggle:
            img = random_aug_kaggle(img, p_random_aug_kaggle_train)
        return img

    def aug_optional_val(img):
        if do_hard_aug:
            img = Image.fromarray(sometimes_val.augment_image(np.array(img)))

        if do_random_aug_kaggle:
            img = random_aug_kaggle(img, p_random_aug_kaggle_val)
        return img

    if n_crops_search_train is None:
        log(' -> default transform_train is selected')
        transform_train = transforms.Compose([
            transforms.Lambda(lambda img: [
                aug_optional_train(aug_train(img))
                for i in range(n_crops_train)
            ]),
            transforms.Lambda(lambda crops: torch.stack([normalize(to_tensor(crop)) for crop in crops]))
        ])
    else:
        log(' -> dich fscore transform_train is selected')
        transform_train = transforms.Compose([
            transforms.Lambda(lambda img: [
                aug_optional_train(img) for img in aug_train_fscore(img)
            ]),
            transforms.Lambda(lambda crops: torch.stack([normalize(to_tensor(crop)) for crop in crops]))
        ])

    ds_train = ImageList(
        train_list_path,
        transform=transform_train,
        target_transform=transforms.Compose([
            transforms.Lambda(lambda y: [y]*n_crops_train),
            transforms.Lambda(lambda ylist: torch.LongTensor(ylist))
        ]))

    train_ds_list = []
    if train_list_pseudo_npz is not None:
        ds_train_pseudo = NpzFolder(
            train_list_pseudo_npz,
            transform=transforms.Compose([
                transforms.Lambda(lambda img: [
                    aug_train(Image.fromarray(img))
                    for i in range(n_crops_train)
                ]),
                transforms.Lambda(lambda crops: torch.stack([normalize(to_tensor(crop)) for crop in crops]))
            ]),
            target_transform=transforms.Compose([
                transforms.Lambda(lambda y: [y]*n_crops_train),
                transforms.Lambda(lambda ylist: torch.LongTensor(ylist))
            ]))
        train_ds_list.append(ds_train_pseudo)
        log(' -> pseudo dataset is loaded')

    if train_list_flickr is not None:
        ds_train_flickr = ImageList(
            train_list_flickr,
            transform=transform_train,
            target_transform=transforms.Compose([
                transforms.Lambda(lambda y: [y]*n_crops_train),
                transforms.Lambda(lambda ylist: torch.LongTensor(ylist))
            ]))
        train_ds_list.append(ds_train_flickr)
        log(' -> flickr dataset is loaded')

    if len(train_ds_list) > 0:
        train_ds_list = [ds_train] + train_ds_list
        ds_train = MultiDataset(train_ds_list)
        log(' -> MultiDataset is created: %i' % len(train_ds_list))

        #for ds in train_ds_list:
        #    print('; '.join(['%s: %i' % (k, v) for (k, v) in sorted(ds.class_to_idx.items(), key=lambda t: t[1])]))
        #sys.exit('DEBUG EXIT')

    train_loader = torch.utils.data.DataLoader(
        ds_train,
        batch_size=batch_size_train,
        shuffle=True,
        num_workers=workers,
        pin_memory=True)
    log('train_loader.size: %i' % len(train_loader.dataset.imgs))

    if val_path is not None:
        ds_val = NpzFolder(
            val_path,
            transform=transforms.Compose([
                transforms.Lambda(lambda img: NCrops(img, crop_size=crop_size, step=step_crop_val)),
                transforms.Lambda(lambda crops: torch.stack([normalize(to_tensor(aug_optional_val(Image.fromarray(crop))))
                                                             for crop in crops]))
            ]),
            target_transform=transforms.Compose([
                transforms.Lambda(lambda y: [y]*int(np.floor(1 + (512 - crop_size)/step_crop_val))**2),
                transforms.Lambda(lambda ylist: torch.LongTensor(ylist))
            ]))
        val_loader = torch.utils.data.DataLoader(
            ds_val,
            batch_size=batch_size_val,
            shuffle=False,
            num_workers=workers,
            pin_memory=True)
        log('val_loader.size: %i' % len(val_loader.dataset.imgs))

    model = model_factory[model_type](n_classes)
    if model_path is not None:
        checkpoint = torch.load(model_path)
        model.load_state_dict(checkpoint['model'])
        loss_train = checkpoint['loss_train']
        acc_train = checkpoint['acc_train']
        loss_val = checkpoint['loss_val']
        acc_val = checkpoint['acc_val']
        log('Last state:\n  TLoss: %0.6f\n  TAcc:  %0.4f\n  VLoss: %0.6f\n  VAcc:  %0.4f' %
            (loss_train[-1], acc_train[-1], loss_val[-1], acc_val[-1]))
        del(checkpoint)
        log('model loaded: %s' % model_path)
    model = model.cuda()

    criterion = criterion_factory[criterion_type](criterion_params)
    criterion = criterion.cuda()

    if optim_type == 'sgd':
        optimizer = optim.SGD(
            model.parameters(),
            lr=learning_rate,
            momentum=momentum,
            dampening=0,
            weight_decay=weight_decay,
            nesterov=True)
    elif optim_type == 'adam':
        optimizer = optim.Adam(
            model.parameters(),
            lr=learning_rate,
            betas=(0.9, 0.999),
            weight_decay=weight_decay)

    log('optimizer %s:\n' % str(type(optimizer)) +
        '\n'.join([('%s: %s' % (k, str(v))) for (k, v) in
                   optimizer.param_groups[0].items() if k != 'params']))
    lr_scheduler = torch.optim.lr_scheduler.StepLR(
        optimizer,
        step_size=lr_scheduler_step_size,
        gamma=lr_scheduler_gamma)

    loss_train = []
    acc_train = []
    loss_val = []
    acc_val = []
    epoch_time = []

    for ix_epoch in range(n_epoches):
        start_time = time.time()
        lr_scheduler.step()

        model.train()
        loss_train_batch, acc_train_batch = train_pass(train_loader, model, criterion, optimizer)

        if val_path is not None:
            model.eval()
            loss_val_batch, acc_val_batch = val_pass(val_loader, model, criterion)
        else:
            loss_val_batch, acc_val_batch = 0.0, 0.0

        loss_train.append(loss_train_batch)
        acc_train.append(acc_train_batch)
        loss_val.append(loss_val_batch)
        acc_val.append(acc_val_batch)
        epoch_time.append(time.time() - start_time)

        log('ix_epoch: %i' % ix_epoch)
        log('  Time:  %0.2f\n  TLoss: %0.6f\n  TAcc:  %0.4f\n  VLoss: %0.6f\n  VAcc:  %0.4f' %
            (epoch_time[-1], loss_train[-1], acc_train[-1], loss_val[-1], acc_val[-1]))

        torch.save({
            'model': model.state_dict(),
            'opt': optimizer.state_dict(),
            'epoch': ix_epoch + 1,
            'loss_train': loss_train,
            'acc_train': acc_train,
            'loss_val': loss_val,
            'acc_val': acc_val,
            'class_to_idx': train_loader.dataset.class_to_idx,
            'cfg': cfg
        }, os.path.join(out_dir, 'checkpoint.tar'))

        if val_path is not None and acc_val[-1] == np.max(acc_val):
            log('Best found!')
            copyfile(
                os.path.join(out_dir, 'checkpoint.tar'),
                os.path.join(out_dir, 'best_model.tar'))

        if options.is_debug and ix_epoch == 1:
            log('END OF DEBUG')
            break


#1시간 #이분탐색 #매개변수탐색
"""
https://www.acmicpc.net/problem/2805
[2805] : 나무 자르기
"""

import sys
input = sys.stdin.readline

N, M = map(int, input().split())
lst = list(map(int, input().split()))

start, end = 0, max(lst)

while start <= end:
    mid = (start + end) // 2
    cost = 0
    for i in lst:
        if i >= mid:
            cost += (i - mid)

    if cost >= M:
        start = mid + 1

    else:
        end = mid - 1

print(end)


from sklearn import datasets
from sklearn.model_selection import train_test_split
import numpy as np
import pandas as pd
from collections import Counter
from sklearn import metrics
import matplotlib.pyplot as plt

def KNN(Test_input, X_train, y_train, K):
    #Calculate dist. due to Test_input
    dist_train = np.zeros(X_train.shape[0], dtype=float)
    for i in range(0, X_train.shape[0]):
        dist_train[i] = np.linalg.norm(Test_input - X_train[i])

    Train_set_with_dist = np.concatenate((X_train, y_train[:, None], dist_train[:, None]), axis = 1)
    Sorting_Train_set = Train_set_with_dist[Train_set_with_dist[:,5].argsort()]

    Train_set_with_sort = Sorting_Train_set[:, 0 : -1]
    #Take out Top K entry as neighbor
    Neighbor_set = Train_set_with_sort[0 : K]
    A = Neighbor_set[:,-1].tolist()
    c = Counter(Neighbor_set[:,-1].tolist())
    value, count = c.most_common()[0]
    return value

#Load data
iris_dataset = datasets.load_iris()
X_train, X_test, y_train, y_test = train_test_split(iris_dataset['data'], iris_dataset['target'], test_size = 0.5, random_state = 0)

predicted_y_train = np.zeros((X_train.shape[0], 1), dtype = float)
predicted_y_test  = np.zeros((X_test.shape[0], 1), dtype = float)

train_error_array = np.zeros((20,1), dtype = float)
test_error_array = np.zeros((20, 1), dtype = float)
for i in range(0 , 20):
    train_error_point = 0.
    test_error_point = 0.

    for j in range(0, X_train.shape[0]):
        predicted_y_train[j] = KNN(X_train[j], X_train, y_train, i + 1)
        predicted_y_test[j] =  KNN(X_test[j], X_train, y_train, i + 1)

    train_error = 1 - metrics.accuracy_score(y_train, predicted_y_train)
    test_error = 1 - metrics.accuracy_score(y_test, predicted_y_test)
    print("K : {0:d} Train_Error: {1:.2f} Test_Error: {2:.2f}".format(i + 1, train_error, test_error))
    train_error_array[i] = train_error
    test_error_array[i] = test_error

y =  np.concatenate((train_error_array, test_error_array), axis = 1)
# Draw the outcome
plots = plt.plot(range(1, 21), y)
plt.legend(plots, ('Train_Error', 'Test_Error'), loc='best', framealpha=0.5, prop={'size': 'large', 'family': 'monospace'})
plt.show()


'''
Created on Nov 17, 2018

@author: ameier
'''
import math
import re

def get_variables_and_names(ntllayers):
    import tensorflow as tf
    '''
    Returns three different types of variables and their names:
        - the pre-trained weights and bias of transfer learning layers
        - the variables of the new layers after the transfer learning layers
          (recurrent and dense output layers)
        - all trainable variables both in tl and prediction layers

    TODO funktioniert das auch ohne Transfer Learning?
    '''

    allow_print = False
    # variable names for weights/bias have the following pattern, e.g:
    # rnn/multi_rnn_cell/cell_0/lstm_cell/bias:0
    # rnn/multi_rnn_cell/cell_0/lstm_cell/kernel:0
    # https://stackoverflow.com/questions/45617026/tensorflow-transfer-learning-how-to-load-part-of-layers-from-one-checkpoint-file
    # (7.11.18)
    pattern = re.compile(
        "/multi_rnn_cell/cell_[0-" + str(ntllayers - 1) + "]")

    # pre-trained variables
    tl_variables = [var for var in tf.global_variables()
                    if pattern.search(var.name) and (var.name.endswith('/bias:0')
                                                     or var.name.endswith('/kernel:0'))]
    tl_vars_names = [var.name for var in tl_variables]
    if allow_print:
        print("\n\n", flush=True)
        print("variables to restore: ", flush=True)
        for name in tl_vars_names:
            print(name, flush=True)

    # all trainable variables
    all_variables = tf.trainable_variables()
    all_vars_names = [var.name for var in all_variables]
    if allow_print:
        print("all variables: ", flush=True)
        for name in all_vars_names:
            print(name, flush=True)

    # the variables that should be trained (the new prediction layers, but
    # not the pre-trained tl layers)
    predlayer_variables = []
    predlayer_vars_names = []
    for var, name in zip(all_variables, all_vars_names):
        if name not in tl_vars_names:
            predlayer_variables.append(var)
            predlayer_vars_names.append(name)
    if allow_print:
        print("predlayer_vars_names: ", flush=True)
        for name in predlayer_vars_names:
            print(name, flush=True)

    return tl_variables, tl_vars_names, predlayer_variables, \
        predlayer_vars_names,  all_variables, all_vars_names

def build_tl_model(rnn_type, ntllayers, ntimesteps,
                   epochs, dims, returnseq, test_b_size, with_dense_first):
    from predictors.tensorflow_rnnclass_withoutstate import TFRNNWithoutState
    # build transfer learning network
    return TFRNNWithoutState(dims, n_time_steps_to_use=ntimesteps,
                             test_b_size=test_b_size, n_epochs=epochs, has_time_outputs=returnseq,
                             custom_reset=False, n_rnn_layers=ntllayers, n_neurons_per_layer=None,
                             rnn_type=rnn_type, with_dense_first=with_dense_first)

def build_tl_rnn_predictor(rnn_type, ntllayers, n_overall_layers,
                           ntimesteps, epochs, dims,
                           returnseq, test_b_size, apply_tl, with_dense_first,
                           tl_learn_rate):
    '''
    Builds the prediction model.

    In case transfer learning is applied, the weights and biases for the tl layers
    (except those for the dense output layer) are reused. An additional
    rnn (or lstm) layer is added after the tl layers and a new dense output
    layer is created. The overall number of recurrent layers: ntllayers+npredlayers

    In case transfer learning is not applied, a new untrained network is created.
    Overall number of recurrent layers: npredlayers

    @return: prediction model (a tensorflow model)
    '''
    import tensorflow as tf
    # build completely new network with all layers
    pred_model = build_tl_model(rnn_type, n_overall_layers, ntimesteps,
                                epochs, dims, returnseq, test_b_size, with_dense_first)

    if apply_tl:
        # set learning rates for tl and prediction layers separately (8.11.18)
        # https://stackoverflow.com/questions/34945554/how-to-set-layer-wise-learning-rate-in-tensorflow

        # default learning rate for AdamOptimizer in Tensorflow: 0.001
        # https://www.tensorflow.org/api_docs/python/tf/train/AdamOptimizer
        tl_variables, _, predlayer_variables, _, _, _ = get_variables_and_names(
            ntllayers)
        var_list1 = tl_variables  # first layers
        var_list2 = predlayer_variables  # remaining layers
        opt1 = tf.train.AdamOptimizer(tl_learn_rate)
        opt2 = tf.train.AdamOptimizer(0.001)
        grads = tf.gradients(pred_model.loss, var_list1 + var_list2)
        grads1 = grads[:len(var_list1)]
        grads2 = grads[len(var_list1):]
        train_op1 = opt1.apply_gradients(zip(grads1, var_list1))
        train_op2 = opt2.apply_gradients(zip(grads2, var_list2))
        pred_model.train_op = tf.group(train_op1, train_op2)  # new train_op
    return pred_model


import os

import engine.core_game as pt

width, height = 700, 650
pos_x, pos_y = 350, 40
x_cells, y_cells = 50, 50
os.environ['SDL_VIDEO_WINDOW_POS'] = "%d,%d" % (pos_x, pos_y)

background_color = (0, 0, 0)
grid_color = (25, 170, 190)
tetramino_color = (225, 19, 19)

pt.init_game(width, height, x_cells, y_cells, background_color, grid_color, tetramino_color)


from .QtCore import QObject, QTimer, Qt, Signal
from .QtGui import QApplication

class QDeferredCaller(QObject):
    """ A QObject subclass which facilitates executing callbacks on the
    main application thread.

    """
    _posted = Signal(object)

    def __init__(self):
        """ Initialize a QDeferredCaller.

        """
        super(QDeferredCaller, self).__init__()
        app = QApplication.instance()
        if app is not None:
            self.moveToThread(app.thread())
        self._posted.connect(self._onPosted, Qt.QueuedConnection)

    #--------------------------------------------------------------------------
    # Private API
    #--------------------------------------------------------------------------
    def _onPosted(self, callback):
        """ A private signal handler for the '_callbackPosted' signal.

        This handler simply executes the callback.

        """
        callback()

    #--------------------------------------------------------------------------
    # Public API
    #--------------------------------------------------------------------------
    def deferredCall(self, callback, *args, **kwargs):
        """ Execute the callback on the main gui thread.

        Parameters
        ----------
        callback : callable
            The callable object to execute on the main thread.

        *args, **kwargs
            Any additional positional and keyword arguments to pass to
            the callback.

        """
        f = lambda: callback(*args, **kwargs)
        self._posted.emit(f)

    def timedCall(self, ms, callback, *args, **kwargs):
        """ Execute a callback on a timer in the main gui thread.

        Parameters
        ----------
        ms : int
            The time to delay, in milliseconds, before executing the
            callable.

        callback : callable
            The callable object to execute at on the timer.

        *args, **kwargs
            Any additional positional and keyword arguments to pass to
            the callback.

        """
        f = lambda: callback(*args, **kwargs)
        f2 = lambda: QTimer.singleShot(ms, f)
        self._posted.emit(f2)

#: A globally available caller instance. This will be created on demand
#: by the globally available caller functions.
_caller = None

def deferredCall(callback, *args, **kwargs):
    """ Execute the callback on the main gui thread.

    This is a convenience wrapper around QDeferredCaller.deferredCall.
    This should only be called after the QApplication is created.

    """
    global _caller
    c = _caller
    if c is None:
        c = _caller = QDeferredCaller()
    c.deferredCall(callback, *args, **kwargs)

def timedCall(ms, callback, *args, **kwargs):
    """ Execute a callback on a timer in the main gui thread.

    This is a convenience wrapper around QDeferredCaller.timedCall.
    This should only be called after the QApplication is created.

    """
    global _caller
    c = _caller
    if c is None:
        c = _caller = QDeferredCaller()
    c.timedCall(ms, callback, *args, **kwargs)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#----------------------------------------------------------------------------
# Created By  : Mohammad Zarei
# Created Date: 30 Dec 2019
# ---------------------------------------------------------------------------
"""Metrics used to evaluate performance of NB models (SPFs)"""
# ---------------------------------------------------------------------------
# Imports

import pandas as pd
import numpy as np

def MAPE(actual, predicted):
    '''
    Returns Mean Absolute Percentage Error.

    actual, predicted: numpy arrays
    '''
    res = np.empty(len(actual))
    for j in range(len(actual)):
        if actual[j] != 0:
            res[j] = (actual[j] - predicted[j]) / actual[j]
        else:
            res[j] = predicted[j] / np.mean(actual)
    return np.mean(np.abs(res))

def FI(results):
    '''
    Returns average of False Identification Test for top n% hotspots where n is 2,4,6,8,10.
    results dataframe must have 'true_rank' and 'rank' columns
    '''

    FI_test = 0
    for HS_level in [0.02,0.04,0.06,0.08,0.1]:
        FI_test += 1-results[(results['true_rank'] > 1- HS_level) & (results['rank'] > 1- HS_level)].count()[0]/results[(results['rank'] > 1- HS_level)].count()[0]
    return round(FI_test/5, 3)

def PMD(results):
    '''
    Returns average of Poison Mean Different Test for top n% hotspots where n is 2,4,6,8,10.
    results dataframe must have 'true_rank', 'rank', 'lambda' columns
    '''
    PMD_test = 0
    for HS_level in [0.02,0.04,0.06,0.08,0.1]:
        PMD_test += (results[results['true_rank'] > 1-HS_level]['lambda'].mean() - results[results['rank'] > 1-HS_level]['lambda'].mean())/results[results['true_rank'] > 1-HS_level]['lambda'].mean()
    return round(PMD_test/5,3)

def CURE(df, ax, x_label, plot_label, ls='-', y_name='Obs', x='Ftot', y_pred='y_pred',boundry=True):
    '''
    Plot Cumulative Residual plot based on predictions ('y_pred') and observed crash ('Obs') counts in 'df'
    '''
    df.sort_values(by=x, ascending=True, inplace=True)
    df.reset_index(drop=True, inplace=True)
    df['i'] = range(1, len(df[y_name])+1)
    df['res'] = df[y_pred] - df[y_name]
    df['res_sq'] = (df[y_pred] - df[y_name])**2

    df['e1'] = df['res'].cumsum()

    df.iloc[-1, df.columns.get_loc('e1')] = 0

    df['e2'] = df['res_sq'].cumsum()
    df['e3'] = (df['e2']*((1 - df['e2']/df['e2'].iloc[-1])))**0.5
    area_out_of_bound = len(df[(df.e1 > 2*df.e3) | (df.e1 < -2*df.e3)])/len(df)
    print(area_out_of_bound)
    if boundry == True:
        ax.plot(df[x],  2.5*df['e3'], linestyle = 'dotted',label="2$\sigma$")
        ax.plot(df[x], -2.5*df['e3'], linestyle = 'dotted',label="-2$\sigma$")
    ax.plot(df[x], df['e1'], linestyle = ls, label = plot_label)
    ax.legend(loc='upper right',fontsize=10)
    ax.set_xlabel(x_label)
    ax.set_ylabel('CURE')
    return df


import jax
from jax import numpy as jnp
import torch
from torch import nn as nn
from torch.nn import functional as F

import numpy as np

from einops import rearrange

import jax_from_pt as jpt

class SelfAttention(jpt.JaxModule):
    def setup(config):
        module = jpt.Module()
        module.config = config

        dim = config.embedding_dim

        module.key_matrix = jpt.Linear.init(dim, dim)
        module.query_matrix = jpt.Identity.init()
        module.value_matrix = jpt.Linear.init(dim, dim)

        t_mask = torch.tril(torch.ones(config.context_length, config.context_length)).view(1, 1, config.context_length, config.context_length)
        j_mask = jnp.array(t_mask.numpy())

        module.register_buffer("mask", j_mask, t_mask)

        return module, T_SelfAttention(module)

    def apply(module, state, x):
        module = module(state)

        config = module.config

        *_, T, C = x.shape

        split_heads_shape = x.shape[:-1] + (config.n_heads, config.embedding_dim // config.n_heads)

        key = rearrange(module.key_matrix(x).reshape(split_heads_shape), '... T nh hs -> ... nh T hs') #.transpose(-2, -3)
        query = rearrange(module.query_matrix(x).reshape(split_heads_shape), '... T nh hs -> ... nh hs T')
        value = rearrange(module.value_matrix(x).reshape(split_heads_shape), '... T nh hs -> ... nh T hs') #.transpose(-2, -3)

        logits = jnp.matmul(key, query) # [..., nh, T, hs] x [..., nh, hs, T] -> [..., nh, T, T]

        broadcast_mask = jnp.broadcast_to((module.mask[:, :, :T, :T] == 0), logits.shape)
        masked_logits = jnp.where(broadcast_mask, -jnp.inf, logits)
        # logits.masked_fill(module.mask[:,:,:T,:T] == 0, float('-inf'))

        att_weights = jax.nn.softmax(masked_logits, axis=-1)

        y = jnp.matmul(att_weights, value) #  [..., nh, T, T]  x [..., nh, T, hs] -> [..., nh, T, hs]
        y = rearrange(y, '... nh T hs -> ... T nh hs').reshape(x.shape) # [..., nh, T, hs] -> [..., T, nh, hs] -> [..., T, C]

        return y

class T_SelfAttention(jpt.TModule):

    def forward(self, x):

        # I don't totally understand why heads are a good idea, but apparently they are...
        *_, T, C = x.shape

        config = self.config

        # assert C == self.dim, "specified axis does not have correct dimension: was {}, expected {}".format(C, self.dim)

        split_heads_shape = x.shape[:-1] + (config.n_heads, config.embedding_dim // config.n_heads)
        key = self.key_matrix(x).reshape(split_heads_shape).transpose(-2, -3) # [..., T, C] -> [..., T, nh, hs] -> [..., nh, T, hs]
        query = self.query_matrix(x).reshape(split_heads_shape).transpose(-2, -3) # [..., T, C] -> [..., T, nh, hs] -> [..., nh, T, hs]
        value = self.value_matrix(x).reshape(split_heads_shape).transpose(-2, -3) # [..., T, C] -> [..., T, nh, hs] -> [..., nh, T, hs]

        assert key.shape[-2] == T, "shape mismatch: {}".format(key.shape)

        # print(key.shape)
        # print(query.shape)

        logits = torch.matmul(key, query.transpose(-1, -2)) # [..., nh, T, hs] x [..., nh, hs, T] -> [..., nh, T, T]
        masked_logits = logits.masked_fill(self.mask[:,:,:T,:T] == 0, float('-inf'))

        att_weights = F.softmax(masked_logits, dim=-1)

        y = torch.matmul(att_weights, value) #  [..., nh, T, T]  x [..., nh, T, hs] -> [..., nh, T, hs]
        y = y.transpose(-2, -3).reshape(x.shape) # [..., nh, T, hs] -> [..., T, nh, hs] -> [..., T, C]

        return y

# class SelfAttention(torch.nn.Module):
#     def __init__(self, config):
#         super().__init__()
#         self.config = config
#         self.dim = config.embedding_dim
#         self.key_matrix = torch.nn.Linear(self.dim, self.dim) # maybe we want to mess with initialization schemes later?
#         self.query_matrix = torch.nn.Identity()#torch.nn.Linear(self.dim, self.dim)

#         self.value_matrix = torch.nn.Linear(self.dim, self.dim)

#         self.n_heads = config.n_heads
#         self.context_length = config.context_length
#         self.register_buffer("mask", torch.tril(torch.ones(config.context_length, config.context_length)).view(1, 1, config.context_length, config.context_length))

#         assert self.dim % self.n_heads == 0, "number of heads ({}) does not evenly divide embedding dim ({})".format(self.n_heads, self.dim)

#     def forward(self, x):

#         # I don't totally understand why heads are a good idea, but apparently they are...
#         *_, T, C = x.shape

#         assert C == self.dim, "specified axis does not have correct dimension: was {}, expected {}".format(C, self.dim)

#         split_heads_shape = x.shape[:-1] + (self.n_heads, self.dim // self.n_heads)
#         key = self.key_matrix(x).reshape(split_heads_shape).transpose(-2, -3) # [..., T, C] -> [..., T, nh, hs] -> [..., nh, T, hs]
#         query = self.query_matrix(x).reshape(split_heads_shape).transpose(-2, -3) # [..., T, C] -> [..., T, nh, hs] -> [..., nh, T, hs]
#         value = self.value_matrix(x).reshape(split_heads_shape).transpose(-2, -3) # [..., T, C] -> [..., T, nh, hs] -> [..., nh, T, hs]

#         assert key.shape[-2] == T, "shape mismatch: {}".format(key.shape)

#         # print(key.shape)
#         # print(query.shape)

#         logits = torch.matmul(key, query.transpose(-1, -2)) # [..., nh, T, hs] x [..., nh, hs, T] -> [..., nh, T, T]
#         masked_logits = logits.masked_fill(self.mask[:,:,:T,:T] == 0, float('-inf'))

#         att_weights = F.softmax(masked_logits, dim=-1)

#         y = torch.matmul(att_weights, value) #  [..., nh, T, T]  x [..., nh, T, hs] -> [..., nh, T, hs]
#         y = y.transpose(-2, -3).reshape(x.shape) # [..., nh, T, hs] -> [..., T, nh, hs] -> [..., T, C]

#         return y

class ResidualSelfAttention(jpt.JaxModule):

    def setup(config):
        module = jpt.Module()
        module.config = config

        module.selfattention = SelfAttention.init(config)
        module.ln = jpt.LayerNorm.init(config.embedding_dim)

        module.fc1 = jpt.Linear.init(config.embedding_dim, 2*config.embedding_dim)
        module.fc2 = jpt.Linear.init(2*config.embedding_dim, config.embedding_dim)

        return module, T_ResidualSelfAttention(module)

    def apply(module, state, x):
        module = module(state)

        y = module.ln(x)
        y = module.selfattention(y)

        if module.config.use_fc:
            y = module.fc1(y)
            y = jax.nn.gelu(y)
            y = module.fc2(y)

        y = x + y
        return y

class T_ResidualSelfAttention(jpt.TModule):
    def forward(self, x):
        y = self.ln(x)
        y = self.selfattention(y)

        if self.config.use_fc:
            y = self.fc1(y)
            y = F.gelu(y)
            y = self.fc2(y)

        y = x + y
        return y

# from omegaconf import OmegaConf
# config = OmegaConf.load('config/jax_test/test.yaml').model
# config.vocab_size = 1000

# state, apply, t_module, t_state = StackedAttention.init(config)

# x = np.random.normal(np.ones((10, config.context_length, config.embedding_dim)))

# idx = np.ones((10, config.context_length), dtype=int)
# mask = idx
# labels = np.zeros((10, config.context_length), dtype=int)

# t_idx = torch.tensor(idx, dtype=int)
# t_labels = torch.tensor(labels)
# t_mask = torch.tensor(mask)

# j_idx =

# t_x = torch.tensor(x, dtype=torch.float32)
# j_x = jnp.array(x)

# t_y = t_module(t_x)
# j_y = apply(state, j_x)

# print(f"torch: {t_y}")
# print(f"jax: {j_y}")

# print(jnp.sum((t_y.detach().numpy() - j_y)**2))

class StackedAttention(jpt.JaxModule):
    def setup(config):
        module = jpt.Module()
        module.config = config

        module.features = jpt.Sequential.init(*[ResidualSelfAttention.init(config) for _ in range(config.n_layers)])
        module.tok_embeddings = jpt.Embedding.init(config.vocab_size, config.embedding_dim)
        t_pos_embeddings = torch.zeros(1, config.context_length, config.embedding_dim, requires_grad=True)
        j_pos_embeddings = jnp.array(t_pos_embeddings.detach().numpy())
        module.register_parameter("pos_embeddings", j_pos_embeddings, t_pos_embeddings)

        module.head = jpt.Linear.init(config.embedding_dim, config.vocab_size)

        module.head[0]['params']['weight'] = jnp.zeros_like(module.head[0]['params']['weight'])
        module.head[0]['params']['bias'] = jnp.zeros_like(module.head[0]['params']['bias'])

        with torch.no_grad():
            module.head[2].weight.mul_(0.0)
            module.head[2].bias.mul_(0.0)

        return module, T_StackedAttention(module)

    def apply(module, state, idx, mask, labels, compute_loss=True):
        module = module(state)

        B, T = idx.shape

        T = min(T-1, module.config.context_length)

        tok_embd = module.tok_embeddings(idx[:, :T])
        pos_embd = module.pos_embeddings[:, :T, :]

        x = jax.nn.softmax(tok_embd + pos_embd, axis=-1) # input for attention layers: shape [B, T, C]

        features = module.features(x)

        if not compute_loss:
            return features

        logits = module.head(features) # shape [B, T, V]

        targets = labels[:, :T]#
        # targets = StackedAttention.get_targets(mask, idx, T)

        # cross entropy loss doesn't know about T, so we flatten the time dimension:
        # print("logits: ", logits.shape)
        # print("targets: ", targets.shape)

        logits_for_CE = logits.reshape(-1, logits.shape[-1]) # shape [BT, V]
        # print(f"logits for ce shape: {logits_for_CE.size()}, original logits shape: {logits.size()}, targets shape: {targets.size()}")
        targets_for_CE = targets.reshape(-1) # shape [BT]

        # logits_for_CE = logits_for_CE[targets_for_CE != -100]
        # targets_for_CE = targets_for_CE[targets_for_CE != -100]

        stopped_logits = jax.lax.stop_gradient(logits_for_CE)
        stopped_targets = jax.lax.stop_gradient(targets_for_CE)

        predictions = jnp.argmax(stopped_logits, 1)
        num_targets = jnp.sum(stopped_targets != -100)
        num_correct = jnp.sum(stopped_targets == predictions)
        accuracy = num_correct/num_targets

    #         predictions = torch.argmax(logits_for_CE, 1)
    #         num_targets = torch.sum(targets_for_CE != -100)
    #         num_correct = torch.sum(targets_for_CE == predictions)
    #         accuracy = num_correct/num_targets

        # loss = jnp.average(optax.softmax_cross_entropy_with_integer_labels(logits_for_CE, targets_for_CE))
        loss = jpt.softmax_cross_entropy(logits_for_CE, targets_for_CE)

        return features, loss, accuracy

class T_StackedAttention(jpt.TModule):
    def forward(self, idx, mask, labels, compute_loss=True):

        # x is 1-hot encoding
        B, T = idx.size()

        T = min(T-1, self.config.context_length)

        tok_embd = self.tok_embeddings(idx[:, :T])

        pos_embd = self.pos_embeddings[:, :T, :]

        x = F.softmax(tok_embd + pos_embd, dim=-1) # input for attention layers: shape [B, T, C]

        features = self.features(x)

        if not compute_loss:
            return features

        logits = self.head(features) # shape [B, T, V]

        targets = labels[:, :T]#
        # targets = StackedAttention.get_targets(mask, idx, T)

        # cross entropy loss doesn't know about T, so we flatten the time dimension:
        # print("logits: ", logits.shape)
        # print("targets: ", targets.shape)

        logits_for_CE = logits.reshape(-1, logits.size(-1)) # shape [BT, V]
        # print(f"logits for ce shape: {logits_for_CE.size()}, original logits shape: {logits.size()}, targets shape: {targets.size()}")
        targets_for_CE = targets.reshape(-1) # shape [BT]

        with torch.no_grad():
            predictions = torch.argmax(logits_for_CE, 1)
            num_targets = torch.sum(targets_for_CE != -100)
            num_correct = torch.sum(targets_for_CE == predictions)
            accuracy = num_correct/num_targets

        loss = F.cross_entropy(logits_for_CE, targets_for_CE)

        return features, loss, accuracy

# from omegaconf import OmegaConf
# config = OmegaConf.load('config/jax_test/test.yaml').model
# config.vocab_size = 1000

# state, apply, t_module, t_state = StackedAttention.init(config)

# x = np.random.normal(np.ones((10, config.context_length, config.embedding_dim)))

# idx = np.ones((10, config.context_length), dtype=int)
# mask = idx
# labels = np.zeros((10, config.context_length), dtype=int)

# t_idx = torch.tensor(idx, dtype=int)
# t_labels = torch.tensor(labels)
# t_mask = torch.tensor(mask)

# j_idx = jnp.array(idx)
# j_mask = jnp.array(mask)
# j_labels = jnp.array(labels)

# lr = 0.001

# t_adamw = torch.optim.AdamW(t_module.parameters(), lr, betas=(0.9, 0.99), eps=1e-08, weight_decay=0.001)

# t_f, t_l, t_a = t_module(t_idx, t_mask, t_labels)
# j_f, j_l, j_a = apply(state, j_idx, j_mask, j_labels)

# t_l.backward()
# t_adamw.step()

# import optax

# j_adamw = optax.adamw(lr, b1=0.9, b2=0.99, eps=1e-08, eps_root=0.0, weight_decay=0.001)
# j_adamw_state = j_adamw.init(state['params'])

# def loss_fn(params, constants, idx, mask, labels):
#     return apply({'params': params, 'constants': constants}, idx, mask, labels)[1]

# grad_fn = jax.grad(loss_fn)

# params = state['params']
# constants = state['constants']

# grads = grad_fn(params, constants, j_idx, j_mask, j_labels)

# updates, j_adamw_state = j_adamw.update(grads, j_adamw_state, params)
# params = optax.apply_updates(params, updates)

# state = {
#     'params': params,
#     'constants': constants
# }

# n_t_f, n_t_l, n_t_a = t_module(t_idx, t_mask, t_labels)
# n_j_f, n_j_l, n_j_a = apply(state, j_idx, j_mask, j_labels)

# print(f"torch: {t_l}")
# print(f"jax: {j_l}")

# print(f"new torch: {n_t_l}")
# print(f"new jax: {n_j_l}")

# t_f = t_f.detach().numpy()

# n_t_f = n_t_f.detach().numpy()

# print(jnp.sum((t_f - j_f)**2)/ jnp.sum(t_f**2 + j_f**2))

# print(jnp.sum((n_t_f - n_j_f)**2)/ jnp.sum(n_t_f**2 + n_j_f**2))

# class StackedAttention(torch.nn.Module):
#     def __init__(self, config):
#         super().__init__()
#         self.config = config
#         self.features = torch.nn.Sequential(*[ResidualSelfAttention(config) for _ in range(config.n_layers)])
#         self.tok_embeddings = torch.nn.Embedding(config.vocab_size, config.embedding_dim)
#         self.pos_embeddings = torch.nn.Parameter(torch.zeros(1, config.context_length, config.embedding_dim))
#         self.head = torch.nn.Linear(config.embedding_dim, config.vocab_size)
#         with torch.no_grad():
#             self.head.weight.mul_(0.0)
#             self.head.bias.mul_(0.0)

#     def get_targets(mask, idx, T):
#         targets = idx[:,1:T+1]
#         targets = targets.masked_fill(mask[:,1:T+1] == 0, -100)
#         # print(targets)
#         return targets

#     def forward(self, idx, mask, labels, compute_loss=True):
#         """
#         idx is 1-hot encoding integer tensor shape [B, T] entries are indices into vocab
#         targets is 1-hot encoding integer tensor shape [B, T], entries are indices into vocab for labels.
#             ith entry of bth row of targets is label for ith prefix of idx in bth example in the batch.
#         """

#         # x is 1-hot encoding
#         B, T = idx.size()

#         T = min(T-1, self.config.context_length)

#         tok_embd = self.tok_embeddings(idx[:, :T])

#         pos_embd = self.pos_embeddings[:, :T, :]

#         x = F.softmax(tok_embd + pos_embd, dim=-1) # input for attention layers: shape [B, T, C]

#         features = self.features(x)

#         if not compute_loss:
#             return features

#         logits = self.head(features) # shape [B, T, V]

#         targets = labels[:, :T]#
#         # targets = StackedAttention.get_targets(mask, idx, T)

#         # cross entropy loss doesn't know about T, so we flatten the time dimension:
#         # print("logits: ", logits.shape)
#         # print("targets: ", targets.shape)

#         logits_for_CE = logits.reshape(-1, logits.size(-1)) # shape [BT, V]
#         # print(f"logits for ce shape: {logits_for_CE.size()}, original logits shape: {logits.size()}, targets shape: {targets.size()}")
#         targets_for_CE = targets.reshape(-1) # shape [BT]

#         with torch.no_grad():
#             predictions = torch.argmax(logits_for_CE, 1)
#             num_targets = torch.sum(targets_for_CE != -100)
#             num_correct = torch.sum(targets_for_CE == predictions)
#             accuracy = num_correct/num_targets

#         loss = F.cross_entropy(logits_for_CE, targets_for_CE)

#         return features, loss, accuracy

# config = ModelConfig(vocab_size=4, context_length=3, num_layers=2, embedding_dim=4, n_heads=2)

# l = StackedAttention(config)

# idx = torch.tensor([ [2, 3, 0, 1],
#                      [0, 1, 2, 3],
#                      [3, 0, 1, 2],
#                      [1, 2, 3, 0] ])

# mask = 0*torch.tensor([[1, 1, 1, 1],
#                      [1, 1, 1, 1],
#                      [1, 1, 1, 1],
#                      [1, 1, 1, 1]])

# print(l(idx, mask)[1])

# x = torch.tensor([ [[1.0,  0.0,  0.0, 0.0],
#                     [0.0,  0.0,  0.0, 1.0],
#                     [0.0,  0.5,  0.5, 0.0]],

#                    [[0.5,  0.2,  0.1, 0.2],
#                     [0.0,  1.0,  0.0, 0.0],
#                     [0.0,  0.0,  1.0, 0.0]]])
# print(x.shape)

# y = l(x)
# print(y)


import traceback
import os.path
from struct import unpack

import eidolon
from eidolon import (
    vec3, Vec3Matrix, IndexMatrix, RealMatrix, PyDataSet, GeomType, ElemType, StdProps, splitPathExt, MeshScenePlugin,
    MeshSceneObject,SceneObject,SceneObjectRepr, Future, QtWidgets, ScenePlugin, Qt, BasisGenFuncs, toIterable,
    taskroutine, first, copyfileSafe, getValidFilename, fillList, fillEnumTable
)

from ..ui import Ui_ObjDialog, Ui_DataDialog, Ui_TDDialog

CheartKeywords=eidolon.enum(
    ('loadargs','Argument values for loading mesh data','(.X filename list,.T filename,ElemType of topology,load sequential or not,# of copies for timesteps,filename of base .X file)'),
    ('datafieldargs','List of argument value for loading fields','(.D filename list,.T file,.D file dimensions,Elemtype of topology,name of spatial topology,load sequential or not)')
)

def readDataSet(xfile,tfile,typename, isBinary=False):
    assert typename in ElemType, 'Unknown typename '+typename

    xdata,xheader=readXFile(open(xfile),splitPathExt(xfile)[1],isBinary)
    tdata,theader=readTFile(open(tfile),typename,splitPathExt(tfile)[1],isBinary)

    return PyDataSet(splitPathExt(xfile)[1],xdata,[tdata])

def readXFile(data,buffname, isBinary=False,buff=None,task=None):
    return readCheartBuffer(data,vec3, 1,2,isBinary,buff,buffname,task)

def readTFile(data, typename,buffname,isBinary=False,buff=None,task=None):
    tdata,theader= readCheartBuffer(data,int, ElemType[typename].numNodes(),2,isBinary,buff,buffname,task,True)
    tdata.setType(typename)
    return tdata,theader

def readDFile(data, dim,buffname,isBinary=False,buff=None,task=None):
    return readCheartBuffer(data,float, dim,2,isBinary,buff,buffname,task)

readXFileTask=taskroutine(taskLabel="Read Node File")(readXFile)
readTFileTask=taskroutine(taskLabel="Read Topology File")(readTFile)
readDFileTask=taskroutine(taskLabel="Read Data File")(readDFile)

def readCheartBuffer(data, pydatatype, dim,numHeaders,isBinary=False,buff=None,buffname=None,task=None,isIndex=False):
    if not buff: # create an appropriate matrix if there isn't one provided
        if not buffname:
            buffname='CheartBuffer'

        if pydatatype==vec3:
            buff=Vec3Matrix(buffname,0,dim)
        elif pydatatype==int:
            buff=IndexMatrix(buffname,0,dim)
        elif pydatatype==float:
            buff=RealMatrix(buffname,0,dim)

    if pydatatype==vec3: # replace the data type for vec3 with a pseudo-constructor that takes 1-3 numbers and returns a vec3
        def readVec3(*vals):
            if len(vals)==2:
                return vec3(vals[0],vals[1],0)
            elif len(vals)==1:
                return vec3(vals[0],0,0)

            return vec3(vals[0],vals[1],vals[2])

        pydatatype=readVec3

    elemsize=3 if pydatatype==vec3 else 1
    doClose=False

    if isinstance(data,str):
        doClose=True
        data=open(data,'rb' if isBinary else 'r')

    isFileSrc=hasattr(data,'name') and os.path.isfile(data.name)

    if isFileSrc:
        buff.meta(StdProps._filename,os.path.abspath(data.name))

    try:
        if isBinary:
            header=[unpack('i',data.read(4))[0] for i in range(numHeaders)]
        else:
            header=list(map(int,data.readline().replace(',',' ').split()))

        buff.meta(StdProps._header,' '.join(map(str,header)))
        buff.addRows(header[0]-buff.n())

        if header[1]==2 and elemsize==3:
            elemsize=2

        if task:
            task.setMaxProgress(header[0])

        if isFileSrc: # read from file
            if isBinary:
                buff.readBinaryFile(os.path.abspath(data.name),numHeaders*4)
            else:
                buff.readTextFile(os.path.abspath(data.name),numHeaders)

            if task:
                task.setProgress(header[0])
        elif isBinary: # read binary stream data
            done=False
            numlines=0

            while not done:
                for j in range(dim):
                    comps=[]
                    for i in range(elemsize):
                        d=data.read(4 if pydatatype==int else 8)
                        if d!='':
                            comps.append(unpack('i' if pydatatype==int else 'd',d)[0])
                        else:
                            done=True
                            break

                    if not done:
                        buff.setAt(pydatatype(*comps),numlines,j)

                if task:
                    task.setProgress(numlines)

                numlines+=1
        else: # read text stream data
            i=0
            line=data.readline()
            mapfunc=int if pydatatype==int else float

            while line!='':
                if line.strip()!='':
                    if i==buff.n():
                        buff.addRows(1)

                    comps=list(map(mapfunc, line.replace(',',' ').split()))

                    if task:
                        task.setProgress(i)

                    for j in range(dim):
                        buff.setAt(pydatatype(*comps[j*elemsize:(j+1)*elemsize]),i,j)

                i+=1
                line=data.readline()

        return buff,tuple(header)
    finally:
        if doClose:
            data.close()

def writeCheartBuffer(outfile,header,data,addVal=None):
    if isinstance(data,Vec3Matrix):
        getRow=lambda n:tuple(data.getAt(n))
    else:
        getRow=lambda n:data.getRow(n)

    if addVal!=None:
        getRow1=getRow
        getRow=lambda n:[v+addVal for v in getRow1(n)]

    with open(outfile,'w') as o:
        o.write(' '.join(map(str,header))+'\n')
        for n in range(data.n()):
            line=getRow(n)
            o.write(' '.join(map(str,line))+'\n')

def guessTopologyType(tfile):
    '''Returns a guess of the (geom,order,basis) values for a given text topology file.'''
    geom,order,basis=GeomType._Tet,1,eidolon.BasisGenFuncs._NL # default if `tfile' is not text, has 4 values on the first line, or is otherwise bogus

    if os.path.isfile(tfile) and splitPathExt(tfile)[2]=='.T':
        with open(tfile) as o:
            o.readline() # skip header
            line=first(l for l in o if l.strip()) # skip empty lines

        linelen=len(line.split())
        if linelen==2:
            geom,order=GeomType._Line,1
        elif linelen==3:
            geom,order=GeomType._Tri,1
        elif linelen==6:
            geom,order=GeomType._Tri,2
        elif linelen==8:
            geom,order=GeomType._Hex,1
        elif linelen==9:
            geom,order=GeomType._Quad,2
        elif linelen==10:
            geom,order=GeomType._Tet,2
        elif linelen==27:
            geom,order=GeomType._Hex,2

    return geom,order,basis

@eidolon.concurrent
def loadFileSequenceRange(process,files,typename,dim):
    start=process.startval
    end=process.endval

    results=[]
    count=0

    for f in files[start:end]:
        buffname=splitPathExt(f)[1]
        isBinary=f[-1].lower()=='b'

        try:
            if typename!=None:
                buff,header=readTFile(f,typename,buffname,isBinary)
            elif dim!=None:
                buff,header=readDFile(f,dim,buffname,isBinary)
            else:
                buff,header=readXFile(f,buffname,isBinary)

            buff.setShared(True)
            results.append((buff,header))
        except Exception as e:
            eidolon.printFlush(e)
            traceback.print_exc()
            results.append((None,None))

        count+=1
        process.setProgress(count)

    return results

def loadFileSequence(files,typename=None,dim=None,task=None):
    files=list(map(os.path.abspath,files))

    for f in files:
        assert os.path.isfile(f),'File not found:'+f

    proccount=eidolon.chooseProcCount(len(files),0,10)
    result=loadFileSequenceRange(len(files),proccount,task,files,typename,dim)

    filemap={}

    for buff,header in eidolon.listSum(result.values()):
        filename=buff.meta(StdProps._filename)
        filemap[filename]=(buff,header)

    return [filemap.get(f) for f in files]

class BaseLoadDialog(QtWidgets.QDialog):
    def __init__(self,mgr,parent=None):
        QtWidgets.QDialog.__init__(self,parent)
        self.setupUi(self)
        self.mgr=mgr
        self.lastdir=self.tr('.')

    def accept(self):
        self.done(1)

    def reject(self):
        self.done(0)

    def _chooseFile(self,title,filterstr,filebox,chooseMultiple=False):
        result = self.mgr.win.chooseFileDialog(title,self.lastdir,filterstr,parent=self,chooseMultiple=chooseMultiple)
        if not chooseMultiple and filebox!=None and result.strip()!='':
            filebox.setText(result)
            self.lastdir=os.path.dirname(os.path.abspath(result))

        return result

    def _updateTable(self,table,newitems=[]):
        itemlist=[str(table.item(n,1).text()) for n in range(table.rowCount())]+newitems
        itemlist=[i for i in itemlist if i.strip()!='']
        fillEnumTable(itemlist,table)
        return itemlist

    def _guessTopology(self,tfile,geomBox,basisBox,orderBox):
        geom,order,basis=guessTopologyType(tfile)
        geomBox.setCurrentIndex(geomBox.findText(GeomType[geom][0]))
        basisBox.setCurrentIndex(basisBox.findText(BasisGenFuncs[basis][0]))
        orderBox.setValue(order)

    def _fillGeomBasisBoxes(self,geomBox,basisBox):
        fillList(geomBox,[nf for n,nf,dim,simp in GeomType])
        fillList(basisBox,[desc for n,desc,gen in BasisGenFuncs])

    def _getElemType(self,geomBox,basisBox,orderBox):
        geomtype=GeomType[geomBox.currentIndex()][0]
        basistype=BasisGenFuncs[basisBox.currentIndex()][0]
        return ElemType.getTypeName(geomtype,basistype,orderBox.value())

    def getParams(self):
        pass

class LoadDialog(Ui_ObjDialog,BaseLoadDialog):
    def __init__(self,mgr,parent=None):
        BaseLoadDialog.__init__(self,mgr,parent)
        self._fillGeomBasisBoxes(self.geomBox,self.basisBox)

        self.xchoose.clicked.connect(self.choose_x)
        self.tchoose.clicked.connect(lambda:self._chooseFile("Choose Topology .T File",".T Files (*.T *.Tb)",self.tfile))
        self.tfile.textChanged.connect(lambda:self._guessTopology(str(self.tfile.text()),self.geomBox,self.basisBox,self.orderBox))

    def choose_x(self):
        fname = self._chooseFile("Choose Node .X File",".X .D Files (*.X *.Xb *.D *.Db)",self.xfile)
        if fname:
            fnameext=os.path.splitext(fname)
            if os.path.isfile(fnameext[0]+'.T') and fnameext[1] in ('.X','.D'):
                self.tfile.setText(fnameext[0]+'.T')
            elif os.path.isfile(fnameext[0]+'.Tb') and fnameext[1] in ('.Xb','Db'):
                self.tfile.setText(fnameext[0]+'.Tb')

    def getParams(self):
        result=self.exec_()
        if result!=1:
            return None

        xfile=str(self.xfile.text()).strip()
        tfile=str(self.tfile.text()).strip()
        elemtype=self._getElemType(self.geomBox,self.basisBox,self.orderBox)

        return xfile,tfile,elemtype

class LoadDataDialog(Ui_DataDialog,BaseLoadDialog):
    def __init__(self,mgr,parent=None):
        BaseLoadDialog.__init__(self,mgr,parent)
        self._fillGeomBasisBoxes(self.geomBox,self.basisBox)
        self.dfiles=[]

        self.tchoose.clicked.connect(lambda:self._chooseFile("Choose Topology .T File",".T Files (*.T *.Tb)",self.tfile))
        self.dchoose.clicked.connect(self.addFields)
        self.tfile.textChanged.connect(lambda:self._guessTopology(str(self.tfile.text()),self.geomBox,self.basisBox,self.orderBox))

    def _updateDTable(self,newitems=[]):
        self.dfiles=self._updateTable(self.dfileTable,newitems)

    def addFields(self):
        fnames = self._chooseFile("Choose Field Files",".D Files (*.D *.Db)",None,True)
        self._updateDTable(fnames)

    def getParams(self):
        result=self.exec_()
        if result!=1:
            return None

        elemtype=self._getElemType(self.geomBox,self.basisBox,self.orderBox)
        tfile=str(self.tfile.text()).strip()
        return self.dfiles,self.dimBox.value(), tfile,elemtype

class LoadTDDialog(Ui_TDDialog,BaseLoadDialog):
    def __init__(self,mgr,parent=None):
        BaseLoadDialog.__init__(self,mgr,parent)
        self._fillGeomBasisBoxes(self.geomBox,self.basisBox)
        self._fillGeomBasisBoxes(self.fgeomBox,self.fbasisBox)

        self.xchoose.clicked.connect(lambda:self._chooseFile("Choose Node .X .D File",".X .D Files (*.X *.Xb *.D *.Db)",self.xfile))
        self.tchoose.clicked.connect(lambda:self._chooseFile("Choose Topology .T File",".T Files (*.T *.Tb)",self.tfile))
        self.ftchoose.clicked.connect(lambda:self._chooseFile("Choose Topology .T File",".T Files (*.T *.Tb)",self.ftfile))
        self.tfile.textChanged.connect(lambda:self._guessTopology(str(self.tfile.text()),self.geomBox,self.basisBox,self.orderBox))
        self.ftfile.textChanged.connect(lambda:self._guessTopology(str(self.tfile.text()),self.fgeomBox,self.fbasisBox,self.forderBox))

        self.absoluteRadio.clicked.connect(self._absoluteCheck)
        self.displaceRadio.clicked.connect(self._absoluteCheck)

        self.addNodesButton.clicked.connect(self.addNodes)
        self.addFieldsButton.clicked.connect(self.addFields)

        self.xkeyframes=[]
        self.dkeyframes=[]

        #self.xfileTable.verticalHeader().setMovable(True)
        #self.dfileTable.verticalHeader().setMovable(True)

        self.xfileTable.setDragDropMode(QtWidgets.QAbstractItemView.InternalMove)
        self.dfileTable.setDragDropMode(QtWidgets.QAbstractItemView.InternalMove)

        self.xfileTable.itemChanged.connect(lambda i:self._moveRow(i.row(),self.xkeyframes,self.xfileTable))
        setattr(self.xfileTable,'keyPressEvent',lambda e:self._delRow(e,self.xkeyframes,self.xfileTable))

        self.dfileTable.itemChanged.connect(lambda i:self._moveRow(i.row(),self.dkeyframes,self.dfileTable))
        setattr(self.dfileTable,'keyPressEvent',lambda e:self._delRow(e,self.dkeyframes,self.dfileTable))

    def _absoluteCheck(self):
        isabs=self.absoluteRadio.isChecked()
        if isabs:
            self.xfile.setText('')
        self.xchoose.setEnabled(not isabs)
        self.xfile.setEnabled(not isabs)

    def _delRow(self,event,itemlist,table):
        index=first(i.row() for i in table.selectedIndexes()) or -1
        if event.key() in (Qt.Key_Delete,Qt.Key_Backspace) and 0<=index<len(itemlist):
            itemlist.pop(index)
            fillEnumTable(itemlist,table)

        QtWidgets.QTableWidget.keyPressEvent(table,event)

    def _moveRow(self,j,itemlist,table):
        index=first(i.row() for i in table.selectedIndexes()) or -1
        if j>index:
            j+=1
        if 0<=index<len(itemlist) and j>=0:
            itemlist.insert(j,itemlist.pop(index))
            fillEnumTable(itemlist,table)

    def addNodes(self):
        fnames = self._chooseFile("Choose Keyframe Node Files",".X .D Files (*.X *.Xb *.D *.Db)",None,True)
        if fnames:
            self.xkeyframes=self._updateTable(self.xfileTable,fnames)

    def addFields(self):
        fnames = self._chooseFile("Choose Keyframe Field Files",".D Files (*.D *.Db)",None,True)
        if fnames:
            self.dkeyframes=self._updateTable(self.dfileTable,fnames)

    def getParams(self):
        result=self.exec_()
        if result!=1:
            return None

        xfile=str(self.xfile.text()).strip()
        tfile=str(self.tfile.text()).strip()
        ftfile=str(self.ftfile.text()).strip()
        isAbsolute=self.absoluteRadio.isChecked()
        datadim=self.dimBox.value()
        starttime=self.startTimeBox.value()
        interval=self.intervalBox.value()

        elemtype=self._getElemType(self.geomBox,self.basisBox,self.orderBox)
        felemtype=self._getElemType(self.fgeomBox,self.fbasisBox,self.forderBox)

        return xfile,tfile,ftfile,isAbsolute,datadim,elemtype,felemtype,starttime,interval,self.xkeyframes,self.dkeyframes

class CheartPlugin(MeshScenePlugin):
    def __init__(self):
        ScenePlugin.__init__(self,'CHeart')

    def init(self,plugid,win,mgr):
        ScenePlugin.init(self,plugid,win,mgr)
        if win:
            win.addMenuItem('Import','CHeartObjLoad'+str(plugid),'&CHeart Object',self._openLoadDialog)
            win.addMenuItem('Import','CHeartDataLoad'+str(plugid),'CHeart &Field',self._openLoadDataDialog)
            win.addMenuItem('Import','CHeartDataLoad'+str(plugid),'CHeart &Time Series',self._openLoadTDDialog)
            win.addMenuItem('Export','CHeartExport'+str(plugid),'&CHeart File(s)',self._openSaveDialog)

        conf=mgr.conf

        cheartload=conf.get('args','--cheartload').split(',')
        cheartfield=conf.get('args','--cheartfield').split(',')
        reprload=conf.get('args','--cheartrepr').split(',')

        # If the arguments are present and in the right format, load an object, a field, and/or produce a simple representation
        if len(cheartload)==3:
            obj=self.loadSceneObject(cheartload[0], cheartload[1],cheartload[2])
            self.mgr.addFuncTask(lambda:self.mgr.addSceneObject(obj()))

            # if the field argument is present and correct, load the field with the optional field topology
            if conf.hasValue('args','--cheartfield') and len(cheartfield)>1:
                dfile=cheartfield[0]
                dim=int(cheartfield[1])
                tfile=cheartfield[2] if len(cheartfield)==4 else None
                typename=cheartfield[3] if len(cheartfield)==4 else None
                self.loadDataField(obj,dfile,dim,tfile,typename)

            #if the representation argument is present and correct, generate a representation using the default material
            if conf.hasValue('args','--cheartrepr') and len(reprload)>0:
                reprtype=reprload[0]
                refine=int(reprload[1]) if len(reprload)>1 else 0
                externalOnly=bool(reprload[2]) if len(reprload)>2 else True
                rep=Future()
                self.mgr.addFuncTask(lambda:rep.setObject(obj().createRepr(reprtype,refine,externalOnly=externalOnly)))
                self.mgr.addFuncTask(lambda:self.mgr.addSceneObjectRepr(rep()))
                self.mgr.addFuncTask(mgr.setCameraSeeAll)

    def getHelp(self):
        return '''
Usage: --cheartload=XFILE,TFILE,BASIS [--cheartfield=DFILE,DIM[,TFILE,BASIS]]
                   [--cheartrepr=TYPE[,REFINE[,EXTERNONLY]]]

  Where:
    XFILE, TFILE, DFILE  Data file paths
    BASIS                Basis type identifier (eg Tet1NL)
    DIM, REFINE          Positive integers
    EXTERNONLY           'true' or 'false'.'''

    def getObjFiles(self,obj):
        xfiles,tfile,_,_,_,initXfile=obj.kwargs.get(CheartKeywords._loadargs,([],None,None,None,None,None))

        files=list(toIterable(xfiles))

        if tfile:
            files.append(tfile)

        if initXfile:
            files.append(initXfile)

        for dfiles,tfile,_,_,_,_ in obj.kwargs.get(CheartKeywords._datafieldargs,[]):
            files+=list(toIterable(dfiles))

            if tfile:
                files.append(tfile)

        return list(map(os.path.abspath,files))

    def checkFileOverwrite(self,obj,dirpath,name=None):
        oldname=obj.getName()
        newname=name or oldname
        prefixpath=os.path.join(dirpath,newname)

        inds=max(filter(eidolon.isSpatialIndex,obj.datasets[0].indices.values()),key=len)
        indsname=inds.getName()

        result=[prefixpath+'.T']

        if len(obj.datasets)==1:
            result.append(prefixpath+'.X')
        else:
            result+=['%s%.4i.X'%(prefixpath,i) for i in range(len(obj.datasets))]

        for fieldname in obj.getFieldNames():
            dfs=toIterable(obj.getDataField(fieldname))
#            spatial=dfs[0].meta(StdProps._spatial)

            # write out topology if needed
            topo=dfs[0].meta(StdProps._topology)
            if topo and topo!=indsname:
                topo=obj.datasets[0].getIndexSet(topo)
                result.append(prefixpath+getValidFilename(os.path.splitext(topo.getName())[0])+'_topo.T')

            isTimeCopy=dfs[0].meta(StdProps._timecopy)=='True' # true if the field has been copied to each timestep, so only write it out once
            if isTimeCopy: # if this is a time copy, reduce the list to one field and ensure the name doesn't get numbered
                dfs=[dfs[0]]

            for i,df in enumerate(dfs):
                dfname=getValidFilename(os.path.splitext(df.getName())[0])
                if dfname.startswith(newname):
                    dfprefix=os.path.join(dirpath,dfname)
                else:
                    dfprefix=prefixpath+'_'+dfname

                if not isTimeCopy and len(dfs)>1: # number the filenames to ensure uniqueness and order
                    dfprefix='%s%.4i'%(dfprefix,i)

                result.append(dfprefix+'.D')

        return list(filter(os.path.exists,result))

    def renameObjFiles(self,obj,oldname,overwrite=False):
        if not isinstance(obj,SceneObject) or obj.getName()==oldname:
            return

        newname=obj.getName()
        indexnames={} # maps the old names of topologies to their new names
        xfiles,tfile,typename,loadSeq,loadTS,initXfile=obj.kwargs[CheartKeywords._loadargs]
        fields=list(obj.kwargs.get(CheartKeywords._datafieldargs,[]))

        def _moveFile(f):
            if f==None or not os.path.isfile(f):
                return None # if there's no filename or actual file, return None which indicates a file isn't used for a particular purpose
            else:
                newbasename=splitPathExt(f)[1]
                if oldname in newbasename:
                    newbasename=newbasename.replace(oldname,newname,1)
                else:
                    newbasename='%s_%s'%(newname,newbasename)

                return eidolon.renameFile(f,newbasename,overwriteFile=overwrite)

        # if a topology's name has `oldname' in it, store the mapping between names in indexnames
        for ds in obj.datasets:
            for ind in ds.enumIndexSets():
                oldindname=ind.getName()
                if oldname in oldindname:
                    newindname=oldindname.replace(oldname,newname,1)
                    indexnames[oldindname]=newindname

            ds.renameIndexSet(oldname,newname)
            ds.renameDataField(oldname,newname)

        # move .X files
        if isinstance(xfiles,str):
            xfiles=_moveFile(xfiles) # single file for static mesh
        else:
            xfiles=list(map(_moveFile,xfiles)) # multiple files for time-dependent mesh

        # rename .T file, initial .X file (for displacement series), and change the kwargs for the object
        tfile=_moveFile(tfile)
        initXfile=_moveFile(initXfile)
        obj.kwargs[CheartKeywords._loadargs]=(xfiles,tfile,typename,loadSeq,loadTS,initXfile)

        # rename fields and their topologies as necessary
        for i in range(len(fields)):
            dfiles,tfile,dim,typename,spatialName,loadSeq=fields[i]

            if isinstance(dfiles,str):
                dfiles=_moveFile(dfiles)
            else:
                dfiles=list(map(_moveFile,dfiles))

            spatialName=indexnames.get(spatialName,spatialName) # associate this field with the changed spatial name

            tfile=_moveFile(tfile)
            fields[i]=(dfiles,tfile,dim,typename,spatialName,loadSeq)

        # apply renaming to the field info stored in the object's kwargs
        if fields:
            obj.kwargs[CheartKeywords._datafieldargs]=tuple(fields)

    def copyObjFiles(self,obj,sdir,overwrite=False):
        xfiles,tfile,typename,loadSeq,loadTS,initXfile=obj.kwargs[CheartKeywords._loadargs]

        if isinstance(xfiles,str):
            filename=os.path.join(sdir,os.path.basename(xfiles))
            copyfileSafe(xfiles,filename,overwrite)
            xfiles=filename
        else:
            xfiles=list(xfiles)
            for i,x in enumerate(xfiles):
                xfiles[i]=os.path.join(sdir,os.path.basename(x))
                copyfileSafe(x,xfiles[i],overwrite)

        if tfile:
            filename=os.path.join(sdir,os.path.basename(tfile))
            copyfileSafe(tfile,filename,overwrite)
            tfile=filename

        if initXfile:
            filename=os.path.join(sdir,os.path.basename(initXfile))
            copyfileSafe(initXfile,filename,overwrite)
            initXfile=filename

        obj.kwargs[CheartKeywords._loadargs]=(xfiles,tfile,typename,loadSeq,loadTS,initXfile)

        if CheartKeywords._datafieldargs in obj.kwargs:
            newfields=[]
            for i,args in enumerate(obj.kwargs[CheartKeywords._datafieldargs]):
                dfiles,tfile,dim,typename,spatialName,loadSeq=args
                if isinstance(dfiles,str):
                    filename=os.path.join(sdir,os.path.basename(dfiles))
                    copyfileSafe(dfiles,filename,overwrite)
                    dfiles=filename
                else:
                    dfiles=list(dfiles)
                    for i,x in enumerate(dfiles):
                        dfiles[i]=os.path.join(sdir,os.path.basename(x))
                        copyfileSafe(x,dfiles[i],overwrite)

                if tfile:
                    filename=os.path.join(sdir,os.path.basename(tfile))
                    copyfileSafe(tfile,filename,overwrite)
                    tfile=filename

                newfields.append((dfiles,tfile,dim,typename,spatialName,loadSeq))

            obj.kwargs[CheartKeywords._datafieldargs]=newfields

    def getScriptCode(self,obj,**kwargs):
        configSection=kwargs.get('configSection',False)
        namemap=kwargs.get('namemap',{})
        varname=namemap[obj]
#        scriptdir=kwargs['scriptdir']
        convertpath=kwargs['convertPath']

        if isinstance(obj,MeshSceneObject):
            script=''
            args={}

            if not configSection:
                xfiles,tfile,typename,loadSeq,loadTS,initXfile=obj.kwargs[CheartKeywords._loadargs]

                if isinstance(xfiles,str):
                    xfiles=[xfiles]

                tfile=convertpath(tfile)
                initXfile=convertpath(initXfile)

                if len(xfiles)>1:
                    xfiles=('['+','.join(map(convertpath,xfiles))+']')
                else:
                    xfiles=convertpath(xfiles[0])

                args={
                    'varname':varname,
                    'xfiles':xfiles,
                    'tfile':tfile,
                    'elemtype':('ElemType._'+typename if typename and tfile else '""'),
                    'loadSeq':loadSeq,
                    'loadTS':loadTS,
                    'initXfile':initXfile,
                    'objname':obj.getName()
                }
                script+='%(varname)s=CHeart.loadSceneObject(%(xfiles)s,%(tfile)s,%(elemtype)s,%(loadSeq)r,%(loadTS)r,%(initXfile)s,%(objname)r)'
            else:
                for args in obj.kwargs.get(CheartKeywords._datafieldargs,[]):
                    dfiles,tfile,dim,typename,spatialName,loadSeq=args
                    if isinstance(dfiles,str):
                        dfiles=[dfiles]

                    tfile=convertpath(tfile)

                    if len(dfiles)>1:
                        dfiles=('['+','.join(map(convertpath,dfiles))+']')
                    else:
                        dfiles=convertpath(dfiles[0])

                    args={
                        'varname':varname,
                        'dfiles':dfiles,
                        'tfile':tfile,
                        'dim':dim,
                        'spatialName':spatialName,
                        'loadSeq':loadSeq,
                        'elemtype':('ElemType._'+typename if typename and tfile else '""')
                    }

                    script+='%(varname)s.loadDataField(%(dfiles)s,%(dim)r,%(tfile)s,%(elemtype)s,%(spatialName)r,%(loadSeq)r)\n'%args

                args={'varname':varname,'timesteps':obj.getTimestepList()}
                script+='%(varname)s.setTimestepList(%(timesteps)r)\n'

            return eidolon.setStrIndent(script % args).strip()+'\n'
        else:
            return MeshScenePlugin.getScriptCode(self,obj,**kwargs)

    def _openLoadDialog(self):
        d=LoadDialog(self)
        params=d.getParams()
        if params:
            xfile,tfile,elemtype=params
            if xfile=='':
                self.mgr.showMsg("Invalid X file '"+xfile+"'")
            elif tfile=='':
                self.mgr.showMsg("Invalid T file '"+tfile+"'")
            else:
                obj=self.loadSceneObject(xfile,tfile,elemtype)
                self.mgr.addFuncTask(lambda:self.mgr.addSceneObject(obj()))

    def _openLoadDataDialog(self):
        obj=self.win.getSelectedObject()

        if not obj or not isinstance(obj,MeshSceneObject):
            self.mgr.showMsg('A mesh object must be selected before data fields can be loaded for it','No scene object')
        else:
            d=LoadDataDialog(self)
            params=d.getParams()
            if params:
                dfiles,dim,tfile,elemtype=params
                self.loadDataField(obj,dfiles,dim,tfile,elemtype)

    def _openLoadTDDialog(self):
        d=LoadTDDialog(self)
        params=d.getParams()
        if params:
            xfile,tfile,ftfile,isAbsolute,datadim,elemtype,felemtype,starttime,interval,xkeyframes,dkeyframes=params

            lx=len(xkeyframes)
            ld=len(dkeyframes)

            # if one x file is specified copy this multiple times to produce a time series to match the time field
            if lx==1 and lx<ld:
                xkeyframes=[xkeyframes[0]]*ld
                lx=ld

            assert lx!=0 or ld!=0, 'Spatial or field keyframes must be provided.'
            assert lx==0 or ld==0 or lx==ld, 'Spatial and field keyframes must have the same number of both provided.'

            assert (lx>0 and isAbsolute) or len(xfile)>0

            timesteps=max(lx,ld)

            initialXFile=None if isAbsolute else xfile

            #if lx==0:
            #    xkeyframes=xfile

            obj=self.loadSceneObject(xkeyframes, tfile,elemtype,loadTimesteps=timesteps,initialXFile=initialXFile)

            if len(dkeyframes)>0:
                self.loadDataField(obj,dkeyframes,datadim,ftfile if ftfile!='' else None,felemtype)

            @self.mgr.addFuncTask
            def _add():
                o=Future.get(obj)
                o.setTimestepScheme(starttime,interval)
                self.mgr.addSceneObject(o)

    def _openSaveDialog(self):
        obj=self.win.getSelectedObject()
        if isinstance(obj,SceneObjectRepr):
            obj=obj.parent

        if not isinstance(obj,MeshSceneObject):
            self.mgr.showMsg('Error: Must select mesh data object to export','Cheart Export')
        else:
            filename=self.mgr.win.chooseFileDialog('Choose Cheart .T filename',filterstr='.T Files (*.T)',isOpen=False)
            if filename!='':
                self.saveSceneObject(filename,obj)

    def loadObject(self,filename,name=None,**kwargs):
        xfile=eidolon.ensureExt(filename,'.X',True) # guess X filename based on T filename
        if not os.path.isfile(xfile):
            raise IOError('Cannot find X file to match %r'%filename)

        ttype=guessTopologyType(filename)
        return self.loadSceneObject(xfile,filename,'%s%i%s'%ttype)

    def saveObject(self,obj,path,overwrite=False,setFilenames=False,**kwargs):
        return self.saveSceneObject(path,obj,kwargs.get('indsname',None),setFilenames,overwrite)

    def loadSceneObject(self,xfiles,tfile,typename,loadSequential=False,loadTimesteps=0,initialXFile=None,objname=None):
        f=Future()
        @taskroutine('Loading CHeart Mesh')
        def _load(xfiles,tfile,typename,loadSequential,loadTimesteps,initialXFile,objname,task):
            with f:
                loadargs=(xfiles,tfile,typename,loadSequential,loadTimesteps,initialXFile)
                spatialName=splitPathExt(tfile)[1] if tfile else None
                elems=[]
                datasets=[]

                if isinstance(xfiles,str): # single node file, make into a list
                    xfiles=[xfiles]
                elif initialXFile: # initial node position file, all subsequent node files are displacement values
                    xfiles=[initialXFile]+xfiles

                if not objname:
                    basenames=list(map(os.path.basename,xfiles))
                    if tfile:
                        basenames.append(os.path.basename(tfile))
                    common=eidolon.getStrListCommonality(basenames)
                    if common>3:
                        if basenames[0][common-1]=='.':
                            common-=1

                        objname=basenames[0][:common]
                    else:
                        objname='CheartObj'

                xfiles=loadFileSequence(xfiles,task=task)

                if not all(xfiles):
                    raise IOError('Unable to read CHeart data from files')

                initialXMat,header=xfiles.pop(0) if initialXFile else (None,None)

                if tfile:
                    topo,header=readTFile(tfile,typename,spatialName,tfile[-1].lower()=='b')
                    topo.sub(1) # 0-based indexing, the only sensible way
                    topo.meta(StdProps._isspatial,'True') # indicate that this is a spatial topology
                    elems.append(topo)

                for nodes,nheader in xfiles:
                    if initialXMat: # add initial position to displacement values
                        nodes.add(initialXMat)

                    name=splitPathExt(nodes.meta(StdProps._filename))[1]
                    datasets.append(PyDataSet(name,nodes,[e.clone() for e in elems]))

                # if a scene object with multiple timesteps is requested but the node data is static, clone the first dataset
                if loadTimesteps>0 and len(datasets)==1:
                    datasets=[datasets[0].clone('%s %i'%(datasets[0].getName(),i+1)) for i in range(loadTimesteps)]

                f.setObject(MeshSceneObject(objname,datasets,self,loadargs=loadargs))

        return self.mgr.runTasks(_load(xfiles,tfile,typename,loadSequential,loadTimesteps,initialXFile,objname),f,loadSequential)

    def saveSceneObject(self,prefixpath,obj,indsname=None,setObjArgs=False,overwrite=False):
        f=Future()
        @taskroutine('Saving CHeart Mesh')
        def _save(prefixpath,obj,indsname,setObjArgs,overwrite,task):
            with f:
                # these will be stored in the object's kwargs member to reflect the arguments to be passed to loadSceneObject() and loadDataField()
                loadargs=[[],None,None,False,0,None] # new arguments for loading the object
                datafieldargs=[] # new arguments for loading the fields

                prefixpath=os.path.splitext(prefixpath)[0]
                if os.path.isdir(prefixpath):
                    prefixdir=prefixpath
                    prefixname=obj.getName()
                else:
                    prefixdir,prefixname=os.path.split(prefixpath)

                dds=obj.datasets
                inds=None

                if len(dds)==1:
                    dslist=[(prefixpath,dds[0])]
                else:
                    dslist=[('%s%.4i'%(prefixpath,i),ds) for i,ds in enumerate(dds)]

                # get the spatial topology, this either the topology named by `indsname' or the largest one stored in dds[0]
                if indsname:
                    inds=obj.getIndexSet(indsname) # choose the named index

                if not inds:
                    # choose the largest spatial index, this is a kludge and obviously can cause problems
                    inds=max(filter(eidolon.isSpatialIndex,dds[0].indices.values()),key=len)
                    indsname=inds.getName()

                # names of topologies and the files they get stored in must match so change the topology names in each of dds
#               for ds in dds:
#                   if not any(n.startswith(prefixname) for n in ds.getIndexNames()):
#                       ds.renameIndexSet(indsname,prefixname)

                indsname=inds.getName()
                loadargs[1]=prefixpath+'.T'
                loadargs[2]=inds.getType()

                # write out topology
                writeCheartBuffer(prefixpath+'.T',[inds.n(),dds[0].getNodes().n()],inds,1) # 1-based

                # write out .X files
                for prefix,ds in dslist:
                    nodes=ds.getNodes()
                    writeCheartBuffer(prefix+'.X',[nodes.n(),3],nodes)
                    loadargs[0].append(prefix+'.X')

                # write out fields with their topologies if needed
                for fieldname in obj.getFieldNames():
                    dfs=toIterable(obj.getDataField(fieldname))
                    spatial=dfs[0].meta(StdProps._spatial)
                    if spatial!=inds.getName():
                        continue

                    fieldarg=[[],None,1,None,splitPathExt(loadargs[1])[1],False]
                    datafieldargs.append(fieldarg)

                    # write out topology if needed
                    topo=dfs[0].meta(StdProps._topology)
                    if topo and topo!=indsname:
                        topo=dds[0].getIndexSet(topo)
                        toponame=prefixpath+getValidFilename(os.path.splitext(topo.getName())[0])+'_topo.T'
                        if not os.path.exists(toponame):
                            writeCheartBuffer(toponame,[topo.n(),dds[0].getNodes().n()],inds,1) # 1-based

                        fieldarg[1]=toponame
                        fieldarg[3]=topo.getType()

                    isTimeCopy=dfs[0].meta(StdProps._timecopy)=='True' # true if the field has been copied to each timestep, so only write it out once
                    if isTimeCopy: # if this is a time copy, reduce the list to one field and ensure the name doesn't get numbered
                        dfs=[dfs[0]]

                    # write out field matrices for this field
                    for i,df in enumerate(dfs):
                        dfname=getValidFilename(os.path.splitext(df.getName())[0])
                        if prefixname and dfname.startswith(prefixname):
                            dfprefix=os.path.join(prefixdir,dfname)
                        else:
                            dfprefix=prefixpath+'_'+dfname

                        if not isTimeCopy and len(dfs)>1: # number the filenames to ensure uniqueness and order
                            dfprefix='%s%.4i'%(dfprefix,i)

                        writeCheartBuffer(dfprefix+'.D',[df.n(),df.m()],df)

                        fieldarg[0].append(dfprefix+'.D')
                        fieldarg[2]=df.m()

                loadargs=tuple(loadargs)
                datafieldargs=tuple(map(tuple,datafieldargs))
                if setObjArgs:
                    if obj.plugin:
                        obj.plugin.removeObject(obj)
                    obj.plugin=self
                    obj.kwargs[CheartKeywords._loadargs]=loadargs
                    obj.kwargs[CheartKeywords._datafieldargs]=datafieldargs

                f.setObject((loadargs,datafieldargs))

        return self.mgr.runTasks(_save(prefixpath,obj,indsname,setObjArgs,overwrite),f)

    def loadDataField(self,obj,dfiles,dim,tfile=None,typename=None,spatialName=None,loadSequential=False,*args,**kwargs):
        f=Future()
        @taskroutine('Loading CHeart Field')
        def _load(obj,dfiles,dim,tfile,typename,spatialName,loadSequential,task):
            with f:
                dfiles=toIterable(dfiles)
                obj=Future.get(obj)
                datafieldargs=(dfiles,tfile,dim,typename,spatialName,loadSequential)

                if isinstance(obj,MeshSceneObject):
                    objname=obj.getName()
                    ds=obj.datasets[0]
                    obj.kwargs.setdefault(CheartKeywords._datafieldargs,[]).append(datafieldargs)
                else:
                    ds=obj
                    objname=''

                # if no data topology name given, choose the first spatial topology for this field
                spatialName=spatialName or first(i.getName() for i in ds.enumIndexSets() if eidolon.isSpatialIndex(i))
                toponame=splitPathExt(tfile)[1] if tfile else spatialName
                ind=ds.getIndexSet(toponame)

                results=loadFileSequence(dfiles,None,dim,task)

                if not all(results):
                    raise IOError('Unable to read CHeart data field from files')

                if not ind and tfile: # if there's no topology by the given name and theres a tfile, load the tfile
                    ind,header=readTFile(tfile,typename,toponame,tfile[-1].lower()=='b')
                    ind.sub(1) # sensible 0-based indexing
                    ind.meta(StdProps._spatial,spatialName)

                datas=[]
                for buff,header in results:
                    bname=buff.getName()
                    if objname and bname.startswith(objname): # if the field names are prefixed with the object name, remove this prefixing
                        newname=bname[len(objname)+1:]
                        if newname[0] in ('_','-',' '):
                            newname=newname[1:]

                        buff.setName(newname)
                    buff.meta(StdProps._topology,toponame)
                    buff.meta(StdProps._spatial,spatialName)

                    if buff.n()!=ds.getNodes().n() and ind and buff.n()==ind.n():
                        buff.meta(StdProps._elemdata,'True')

                    datas.append(buff)

                # determine a field name from the commonality amongst the loaded field matrix names, used for time-dependent fields only
                names=[d.getName() for d in datas]
                fieldname=min(names,key=len)
                common=eidolon.getStrListCommonality(names) or len(fieldname)
                fieldname=fieldname[:common].rstrip('0')

                for d in datas: # set each field matrix to have the same name
                    d.meta('oldname',d.getName())
                    d.setName(fieldname)

                if isinstance(obj,MeshSceneObject) and len(obj.datasets)>1:
                    # if there's only one matrix, make a list of the one matrix for each timestep, this shares the matrix between datasets
                    if len(datas)==1:
                        datas[0].meta(StdProps._timecopy,'True')
                        datas=datas*len(obj.datasets)

                    for d,ds1 in zip(datas,obj.datasets):
                        if not d.meta('oldname'): # recall the previous name of the matrix just in case
                            d.meta('oldname',d.getName())
                        d.setName(fieldname) # set the matrix name to the field name
                        ds1.setDataField(d)
                        if ind:
                            ds1.setIndexSet(ind)
                else:
                    for d in datas:
                        ds.setDataField(d)

                    if ind:
                        ds.setIndexSet(ind)

                f.setObject(datas if len(datas)>1 else datas[0])

        return self.mgr.runTasks(_load(obj,dfiles,dim,tfile,typename,spatialName,loadSequential),f)

eidolon.addPlugin(CheartPlugin())


from rest_framework.serializers import ModelSerializer

from events.models import EventItem
from events.serializers import (EventItemGroupSerializer, TransportationPickupLocationSerializer,
                                AccommodationRoomSerializer)
from items.serializers import ItemMasterSerializer

class EventItemSerializer(ModelSerializer):
    class Meta:
        model = EventItem

        exclude = ('uuid', 'event', 'items_booked', 'deleted_at')

    def to_representation(self, instance):
        response = super().to_representation(instance)
        data = {
            'uuid': instance.uuid,
            'early_bird_rate': instance.get_early_bird_rate(),
            'late_bird_rate': instance.get_late_bird_rate(),
            'current_rate': instance.get_current_rate(),
            'items_booked': instance.items_booked,

            'item_master': ItemMasterSerializer(instance.item_master).data,
            'group': EventItemGroupSerializer(instance.group).data,

            'transportation_pickup_locations': TransportationPickupLocationSerializer(
                instance.transportation_pickup_locations, many=True).data,
            'accommodation_rooms': AccommodationRoomSerializer(
                instance.accommodation_rooms, many=True).data
        }
        response.update(data)
        return response