#!/usr/bin/env python
import argparse
import logging
import os
import sys
import numpy as np
from shutil import copyfile, rmtree
from ase import units
import torch
import torch.nn as nn
from torch.optim import Adam
from torch.utils.data.sampler import RandomSampler
from schnetpack.utils.script_utils.settings import get_environment_provider
from schnetpack.utils import get_dataset
import schnetpack as spk
import schnet_ev
from schnet_ev.data import Eigenvalue_properties, Delta_EV_properties, Charges
logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))

def get_parser():
    """ Setup parser for command line arguments """
    main_parser = argparse.ArgumentParser()

    # command-specific
    cmd_parser = argparse.ArgumentParser(add_help=False)
    cmd_parser.add_argument('--cuda', help='Set flag to use GPU(s)', action='store_true')
    cmd_parser.add_argument('--parallel',
                            help='Run data-parallel on all available GPUs (specify with environment variable'
                                 + ' CUDA_VISIBLE_DEVICES)', action='store_true')
    cmd_parser.add_argument('--batch_size', type=int,
                            help='Mini-batch size for training and prediction (default: %(default)s)',
                            default=100)
    cmd_parser.add_argument("--aggregation_mode",type=str,default="build",choices=["avg","build"])
    cmd_parser.add_argument('--environment_provider',type=str, default="simple", choices=["simple","ase","torch"],help="Environment provider for dataset. (default: %(default)s)",)
    cmd_parser.add_argument('--force_mask_index',help='Atom number of excluded atom for force training', type=int)
    cmd_parser.add_argument('--force_mask', help= 'Only train on forces for specific atoms', action="store_true")
    # training
    train_parser = argparse.ArgumentParser(add_help=False, parents=[cmd_parser])
    train_parser.add_argument('datapath', help='Path / destination of MD17 dataset')
    train_parser.add_argument('modelpath', help='Destination for models and logs')
    train_parser.add_argument('--seed', type=int, default=None, help='Set random seed for torch and numpy.')
    train_parser.add_argument('--overwrite', action='store_true', help='Overwrite old directories')

    # data split
    train_parser.add_argument('--split_path', help='Path / destination of npz with data splits',
                              default=None)
    train_parser.add_argument('--split', help='Give sizes of train and validation splits and use remaining for testing',
                              type=int, nargs=2, default=[None, None])
    train_parser.add_argument('--max_epochs', type=int, help='Maximum number of training epochs (default: %(default)s)',
                              default=5000)
    train_parser.add_argument('--lr', type=float, help='Initial learning rate (default: %(default)s)',
                              default=1e-4)
    train_parser.add_argument('--lr_patience', type=int,
                              help='Epochs without improvement before reducing the learning rate (default: %(default)s)',
                              default=25)
    train_parser.add_argument('--lr_decay', type=float, help='Learning rate decay (default: %(default)s)',
                              default=0.8)
    train_parser.add_argument('--lr_min', type=float, help='Minimal learning rate (default: %(default)s)',
                              default=1e-6)
    train_parser.add_argument('--logger', help='Choose logger for training process (default: %(default)s)',
                              choices=['csv', 'tensorboard'], default='csv')
    train_parser.add_argument('--log_every_n_epochs', type=int,
                              help='Log metrics every given number of epochs (default: %(default)s)',
                              default=1)
    train_parser.add_argument('--tradeoffs', type=str, help='Property tradeoffs for training', default=None)
    train_parser.add_argument('--verbose', action='store_true', help='Print property error magnitudes')
    train_parser.add_argument('--real_socs', action='store_true',
                              help='If spin-orbit couplings are predicted, information should be given whether they are real or not.')
    train_parser.add_argument('--hamiltonian', action='store_true',
                              help='If True, approximate diabatic Hamiltonian using partial charges without diagonalization to train on eigenvalues.')
    train_parser.add_argument('--diabatic', action='store_true',
                              help='If True, approximate diabatic Hamiltonian and diagonalize to train on eigenvalues.',default=True)
    train_parser.add_argument('--phase_loss', action='store_true', help='Use special loss, which ignores phase.')
    train_parser.add_argument('--inverse_nacs', action='store_true', help='Weight NACs with inverse energies.')
    train_parser.add_argument('--smooth_loss', action='store_true', help='Use phase independent min loss.')
    train_parser.add_argument('--mixed_loss', action='store_true', help='Use phase independent min loss.')
    train_parser.add_argument('--diagonal', action='store_true', help='Train SOCs via diagonal elements. Must be included in the training data base', default=None)
    train_parser.add_argument('--L1', action='store_true', help='Use L1 norm')
    train_parser.add_argument('--Huber', action='store_true', help='Use L1 norm')

    ## evaluation
    eval_parser = argparse.ArgumentParser(add_help=False, parents=[cmd_parser])
    eval_parser.add_argument('datapath', help='Path / destination of MD17 dataset directory')
    eval_parser.add_argument('modelpath', help='Path of stored model')
    eval_parser.add_argument('--split', help='Evaluate on trained model on given split',
                             choices=['train', 'validation', 'test', 'all'], default=['all'], nargs='+')
    eval_parser.add_argument('--hessian', action='store_true', help='Gives back the hessian of the PES.')
    eval_parser.add_argument('--hamiltonian', action='store_true',
                              help='If True, approximate diabatic Hamiltonian using partial charges without diagonalization to train on eigenvalues.')
    eval_parser.add_argument('--diabatic', action='store_true',
                              help='If True, approximate diabatic Hamiltonian and diagonalize to train on eigenvalues.',default=True)

    pred_parser = argparse.ArgumentParser(add_help=False, parents=[cmd_parser])
    pred_parser.add_argument('datapath', help='Path / destination of MD17 dataset directory')
    pred_parser.add_argument('--print_uncertainty',action='store_true', help='Print uncertainty of each property',default=False)
    pred_parser.add_argument('--diabatic', action='store_true',
                              help='If True, approximate diabatic Hamiltonian and diagonalize to train on eigenvalues.',default=True)
    pred_parser.add_argument('--thresholds',type=float,nargs=5, help='Percentage of mean predicted by NNs taken as thresholds for adaptive sampling - first value for energy, second for forces, third for dipoles, fourth for nacs, fifth for socs', default=None)
    pred_parser.add_argument('--modelpaths',type=str, help='Path of stored models')
    pred_parser.add_argument('modelpath', help='Path of stored model')
    pred_parser.add_argument('--hessian', action='store_true', help='Gives back the hessian of the PES.')
    pred_parser.add_argument('--adaptive',type=float, nargs=1,help='Adaptive Sampling initializer + float for reducing the step size')
    pred_parser.add_argument('--nac_approx',type=float, nargs=3, default=[1,0.018,0.036],help='Type of NAC approximation as first value and threshold for energy gap in Hartree as second value.')
    # model-specific parsers
    opt_parser = argparse.ArgumentParser(add_help=False, parents=[cmd_parser])
    opt_parser.add_argument('--nac_approx',type=float, nargs=3, default=[1,0.018,0.036],help='Type of NAC approximation as first value and threshold for energy gap in Hartree as second value.')
    model_parser = argparse.ArgumentParser(add_help=False)

    #######  SchNet  #######
    schnet_parser = argparse.ArgumentParser(add_help=False, parents=[model_parser])
    schnet_parser.add_argument('--features', type=int, help='Size of atom-wise representation (default: %(default)s)',
                               default=256)
    schnet_parser.add_argument("--cutoff_function",help="Functional form of the cutoff", choices=["hard","cosine","mollifier"],default="cosine")
    schnet_parser.add_argument('--interactions', type=int, help='Number of interaction blocks (default: %(default)s)',
                               default=6)
    schnet_parser.add_argument('--cutoff', type=float, default=10.,
                               help='Cutoff radius of local environment (default: %(default)s)')
    schnet_parser.add_argument('--num_gaussians', type=int, default=50,
                               help='Number of Gaussians to expand distances (default: %(default)s)')
    schnet_parser.add_argument('--n_layers', type=int, default=3,
                               help='Number of layers in output networks (default: %(default)s)')
    schnet_parser.add_argument('--n_neurons', type=int, default=100,
                              help='Number of nodes in atomic networks (default: %(default)s)')

    #######  invD  ########
    invD_parser = argparse.ArgumentParser(add_help=False, parents=[model_parser])
    invD_parser.add_argument('--n_atoms', type=int, help='Number of atoms in the molecule',
                             default="ADD NUMBER OF ATOMS WITH ARGUENT --n_atoms")
    invD_parser.add_argument('--n_layers', type=int, default=3,
                               help='Number of layers in output networks (default: %(default)s)')
    invD_parser.add_argument('--n_neurons', type=int, default=100,
                              help='Number of nodes in atomic networks (default: %(default)s)')
     #######  wACSF  ########
    wacsf_parser = argparse.ArgumentParser(add_help=False, parents=[model_parser])
    # wACSF parameters
    wacsf_parser.add_argument('--radial', type=int, default=22,
                              help='Number of radial symmetry functions (default: %(default)s)')
    wacsf_parser.add_argument('--angular', type=int, default=5,
                              help='Number of angular symmetry functions (default: %(default)s)')
    wacsf_parser.add_argument('--zetas', type=int, nargs='+', default=[1],
                              help='List of zeta exponents used for angle resolution (default: %(default)s)')
    wacsf_parser.add_argument('--standardize', action='store_true',
                              help='Standardize wACSF before atomistic network.')
    wacsf_parser.add_argument('--cutoff', type=float, default=5.0,
                              help='Cutoff radius of local environment (default: %(default)s)')
    # Atomistic network parameters
    wacsf_parser.add_argument('--n_nodes', type=int, default=100,
                              help='Number of nodes in atomic networks (default: %(default)s)')
    wacsf_parser.add_argument('--n_layers', type=int, default=2,
                              help='Number of layers in atomic networks (default: %(default)s)')
    # Advances wACSF settings
    wacsf_parser.add_argument('--centered', action='store_true', help='Use centered Gaussians for radial functions')
    wacsf_parser.add_argument('--crossterms', action='store_true', help='Use crossterms in angular functions')
    wacsf_parser.add_argument('--behler', action='store_true', help='Switch to conventional ACSF')
    wacsf_parser.add_argument('--elements', default=['H', 'C', 'O'], nargs='+',
                              help='List of elements to be used for symmetry functions (default: %(default)s).')

    # setup subparser structure
    cmd_subparsers = main_parser.add_subparsers(dest='mode', help='Command-specific arguments')
    cmd_subparsers.required = True
    subparser_train = cmd_subparsers.add_parser('train', help='Training help')
    subparser_eval = cmd_subparsers.add_parser('eval', help='Eval help')
    subparser_opt = cmd_subparsers.add_parser('opt', help='Eval help',parents=[opt_parser])
    subparser_pred = cmd_subparsers.add_parser('pred', help='Eval help', parents=[pred_parser])

    train_subparsers = subparser_train.add_subparsers(dest='model', help='Model-specific arguments')
    train_subparsers.required = True
    train_subparsers.add_parser('invD', help='invD help', parents=[train_parser, invD_parser])
    train_subparsers.add_parser('schnet', help='SchNet help', parents=[train_parser, schnet_parser])
    train_subparsers.add_parser('wacsf', help='wACSF help', parents=[train_parser, wacsf_parser])

    eval_subparsers = subparser_eval.add_subparsers(dest='model', help='Model-specific arguments')
    eval_subparsers.required = True
    eval_subparsers.add_parser('invD', help='invD help', parents=[eval_parser, invD_parser])
    eval_subparsers.add_parser('schnet', help='SchNet help', parents=[eval_parser, schnet_parser])
    eval_subparsers.add_parser('wacsf', help='wACSF help', parents=[eval_parser, wacsf_parser])

    opt_subparsers = subparser_opt.add_subparsers(dest="mdel",help="Model-specific arguments")
    opt_subparsers.required = True
    opt_subparsers.add_parser('schnet', help='SchNet help', parents=[eval_parser,schnet_parser])
    return main_parser

def train(args, model, tradeoffs, train_loader, val_loader, device, n_orb,mean,stddev,offset):
    # setup hook and logging
    hooks = [
        spk.train.MaxEpochHook(args.max_epochs)
    ]

    # setup optimizer for training
    to_opt = filter(lambda p: p.requires_grad, model.parameters())
    optimizer = Adam(to_opt, lr=args.lr)

    schedule = spk.train.ReduceLROnPlateauHook(optimizer, patience=args.lr_patience, factor=args.lr_decay,
                                               min_lr=args.lr_min,
                                               window_length=1, stop_after_min=True)
    hooks.append(schedule)

    # Build metrics based on the properties in the model
    metrics = []
    for prop in tradeoffs:
      if prop == "eigenvalues_active_forces" or prop == "eigenvalues_forces":
          pass
      else:
          metrics += [
              schnet_ev.metrics.MeanAbsoluteError(prop, prop),
              schnet_ev.metrics.RootMeanSquaredError(prop, prop)
          ]
    if args.logger == 'csv':
        logger = spk.train.CSVHook(os.path.join(args.modelpath, 'log'),
                                   metrics, every_n_epochs=args.log_every_n_epochs)
        hooks.append(logger)
    elif args.logger == 'tensorboard':
        logger = spk.train.TensorboardHook(os.path.join(args.modelpath, 'log'),
                                           metrics, every_n_epochs=args.log_every_n_epochs)
        hooks.append(logger)
    n_act = n_orb['n_act']
    # Automatically construct loss function based on tradoffs
    def loss(batch, result):
        err_sq = 0.0
        if args.verbose:
            print('===================')
        for prop in tradeoffs:
            if prop == 'energy' or prop == 'forces' or prop == 'homo_lumo':
                prop_diff = batch[prop] - result[prop]
                prop_err = torch.mean(prop_diff.view(-1) ** 2)
                err_sq = err_sq + float(tradeoffs[prop]) * prop_err
            if prop in Eigenvalue_properties.properties or prop in Delta_EV_properties.properties :
                if "eigenvalues_active_forces" in tradeoffs or "eigenvalues_forces" in tradeoffs:
                    result['ev_all']= result[prop]
                    result[prop] = result[prop][prop]
                #results are ordered in descending order
                if args.smooth_loss == True:
                     prop_diff =torch.normal(batch[prop],0.01)-torch.normal(result[prop],0.01)
                     prop_err = torch.mean(prop_diff.view(-1)**2)
                     err_sq = err_sq+float(tradeoffs[prop])*prop_err
                elif args.hamiltonian == True and prop == "eigenvalues":
                     #every vector has a different size - masked for every batch different the mask however is max. length and diff. for every molecule
                     prop_diff = (batch[prop][:,:result[prop].shape[1]]- result[prop]) * batch["eigenvalues_all_mask"][:,:result[prop].shape[1]]
                     prop_err = torch.mean(prop_diff.view(-1)**2)
                     err_sq = err_sq+float(tradeoffs[prop])*prop_err
                elif args.hamiltonian == True and prop != "eigenvalues":
                     prop_diff_ev = (batch[prop][:,:result[prop].shape[1]] - result[prop])*batch["eigenvalues_active_mask"][:,:result[prop].shape[1]]
                     prop_err_ev = torch.mean(prop_diff_ev.view(-1)**2)
                     prop_err_charges=0
                     #get charges
                     if "hirshfeld_pbe" in tradeoffs:
                         q_r=torch.sum((result['coeff']**2)[:,:batch['n_elec_atoms'].shape[1],:]*batch["occupation"][:,None,:result[prop].shape[1]],dim=2)
                         result["hirshfeld_pbe"] = q_r -batch['n_elec_atoms']
                         prop_diff_charges = batch["hirshfeld_pbe"] - result["hirshfeld_pbe"]
                         prop_err_charges += torch.mean(prop_diff_charges.view(-1)**2)*tradeoffs["hirshfeld_pbe"]
                     if "hirshfeld_pbe0" in tradeoffs:
                         q_r=torch.sum((result['coeff']**2)[:,:batch['n_elec_atoms'].shape[1],:]*batch["occupation"][:,None,:result[prop].shape[1]],dim=2)
                         result["hirshfeld_pbe0"]= q_r -batch['n_elec_atoms']
                         prop_diff_charges = batch["hirshfeld_pbe0"] - result["hirshfeld_pbe0"]
                         prop_err_charges += torch.mean(prop_diff_charges.view(-1)**2)*tradeoffs["hirshfeld_pbe0"]
                     err_sq = err_sq+float(tradeoffs[prop])*(prop_err_ev)+prop_err_charges
                elif prop in Delta_EV_properties.properties:
                    prop_diff = torch.abs(batch[prop] - result[prop])*batch["eigenvalues_active_mask"]
                    prop_err = torch.mean(prop_diff.view(-1)**2)
                    err_sq = err_sq  + float(tradeoffs[prop])*prop_err
                else:
                    prop_diff = torch.abs(batch[prop] - result[prop])#*batch["eigenvalues_active_mask"]
                    prop_err = torch.mean(prop_diff.view(-1)**2)
                    err_sq = err_sq  + float(tradeoffs[prop])*prop_err
                if args.mixed_loss == True:
                    #outer subtraction
                    deltaE = torch.abs(batch[prop][...,None] - batch[prop][...,None,:])
                    #make diabatic couplings prop to 1/delta E Delta function
                    mask= batch["%s_mask"%prop][...,None]*  batch["%s_mask"%prop][...,None,:]
                    prop_diff=(torch.abs(torch.triu(torch.normal((1/deltaE),0.1),1))*0.01-torch.abs(torch.triu((result['ev_all']['ev_diabatic']),1)))*mask
                    #print(mask,torch.triu(torch.normal((1/deltaE),0.1),1)*0.01,torch.triu((result['ev_all']['ev_diabatic'])))
                    prop_err2 = torch.mean(prop_diff.view(-1) **2 )
                    #diag forces
                    #print(result['ev_all']['ev_diab_forces'][0].shape)
                    #import numpy as np
                    #a=result['ev_all']['ev_diab_forces'][0].to('cpu').detach().numpy()
                    #a=a.reshape((3,3,7,7))
                    #u,e=np.linalg.eigh(a)
                    #print(e.shape,u.shape)
                    #multiply with properr the tradeoff so that the error gets smaller when the iegenvalues get more accurate - was not good
                    err_sq = err_sq + float(tradeoffs['eigenvalues_active_forces'])  *prop_err2
            if prop == "occupation":
                prop_diff = batch[prop] - result[prop]
                prop_err = torch.mean(prop_diff.view(-1)**2)
                err_sq = err_sq + float(tradeoffs[prop])*prop_err
            if prop == "eigenvalues_active_forces" or prop == "eigenvalues_forces" or prop in Charges.properties:
                #batch[prop]=result['ev_all'][prop]
                pass #continue
            if prop == 'occ_orb':
                #if args.mixed_loss==True:
                #    #error on occupied orbitals
                #    prop_diff_occ_orb = batch[prop] - result[prop]
                #    prop_err_occ_orb = torch.mean(prop_diff_occ_orb.view(-1)**2)
                #    err_sq = err_sq + float(tradeoffs[prop])*prop_err_occ_orb/2
                #    #additionally include the sum of the errors
                #    sum_occ_orb_NN = torch.sum(result[prop],1)*2
                #    prop_diff_sum_orbE = batch['sum_occ_orb'] - sum_occ_orb_NN
                #    prop_err_sum_orbE = torch.mean(prop_diff_sum_orbE.view(-1)**2)
                #    err_sq = err_sq + float(tradeoffs[prop])*prop_err_sum_orbE/2
                #    #additionally try to fit the total energy and force by taking the sum of occ.orb.energies and delta E
                #    total_E_NN = sum_occ_orb_NN + result['delta_E']
                #    total_F_NN = result['delta_E_forces']
                #    for batch_value in range(result['energy'].shape[0]):
                #        total_F_NN[batch_value][0] += torch.sum(result['occ_orb_forces'][batch_value],0)
                #    prop_diff_totE = batch['energy'] - total_E_NN
                #    prop_diff_totF = batch['forces'] - total_F_NN
                #    prop_err_totE = torch.mean(prop_diff_totE.view(-1)**2)
                #    prop_err_totF = torch.mean(prop_diff_totF.view(-1)**2)
                #    err_sq = err_sq + float(tradeoffs['energy']) * prop_err_totE + float(tradeoffs['forces']) * prop_err_totF
                if args.mixed_loss == True:
                    #error of eigenvalues based on distribution function
                    prop_diff_occ_orb =torch.normal(batch[prop],stddev[prop])-torch.normal(result[prop],stddev[prop])
                    prop_err_occ_orb = torch.mean(prop_diff_occ_orb.view(-1)**2)
                    err_sq = err_sq+float(tradeoffs[prop])*prop_err_occ_orb
                else:
                    prop_diff_occ_orb = batch[prop] - result[prop]
                    prop_err_occ_orb = torch.mean(prop_diff_occ_orb.view(-1)**2)
                    err_sq = err_sq + float(tradeoffs[prop])*prop_err_occ_orb
                    #print("err occ", prop_err_occ_orb)
            if prop == "unocc_orb":
                n_unocc=n_orb['n_unocc']
                prop_diff = batch[prop][:2] - result[prop][:2]
                prop_err = torch.mean(prop_diff.view(-1)**2)
                err_sq = err_sq + float(tradeoffs[prop])*prop_err
                #print(prop_err)
            if args.verbose:
                print(prop, prop_err)

        return err_sq

    trainer = spk.train.Trainer(args.modelpath, model, loss, optimizer,
                                train_loader, val_loader, hooks=hooks)
    trainer.train(device)

def evaluate(args, model, train_loader, val_loader, test_loader, device):
    # Get property names from model
    if args.parallel:
        model=model.module.to(device)
        properties=model.output_modules[0].properties
    else:
        properties = model.output_modules[0].properties
    header = ['Subset']
    metrics = []
    for prop in properties:
      if prop == "eigenvalues_active_forces" or prop== "eigenvalues_forces":
        pass
      else:
        header += [f'{prop}_MAE', f'{prop}_RMSE']
        metrics += [
                spk.metrics.MeanAbsoluteError(prop, prop),
                spk.metrics.RootMeanSquaredError(prop, prop)
            ]

    results = []
    if ('train' in args.split) or ('all' in args.split):
        logging.info('Training split...')
        results.append(['training'] + ['%.7f' % i for i in evaluate_dataset(metrics, model, train_loader, device,properties)])

    if ('validation' in args.split) or ('all' in args.split):
        logging.info('Validation split...')
        results.append(['validation'] + ['%.7f' % i for i in evaluate_dataset(metrics, model, val_loader, device,properties)])
    if ('test' in args.split) or ('all' in args.split):
        logging.info('Testing split...')
        results.append(['test'] + ['%.7f' % i for i in evaluate_dataset(metrics, model, test_loader, device,properties)])
    header = ','.join(header)
    results = np.array(results)

    np.savetxt(os.path.join(args.modelpath, 'evaluation.csv'), results, header=header, fmt='%s', delimiter=',')

def evaluate_dataset(metrics, model, loader, device,properties):
    # TODO: Adapt for SCHNARC, copy old
    for metric in metrics:
        metric.reset()

    qm_values={}
    predicted={}
    batchs={}
    header=[]
    results2={}
    """for batch in loader:
        batch = {
            k: v.to(device)
            for k, v in batch.items()
        }
        result = model(batch)
        for prop in result:
            if prop =="eigenvalues_active":
              if prop in result:
                result['ev_all'] += result[prop]
                result[prop] += result[prop][prop]
              else:
                result['ev_all'] = result[prop]
                result[prop]=result[prop][prop]"""
    #do only 1000 batches for testing
    #index_=-1
    for batch in loader:
    #    index_ +=1
    #    if index_==int(1000):
    #       break
        batch = {
            k: v.to(device)
            for k, v in batch.items()
        }
        result = model(batch)
        for prop in result:
          if prop == "eigenvalues_active":
              if "eigenvalues_active_forces" in properties:
                if prop in predicted:
                    predicted[prop] += [(result[prop][prop]*batch["%s_mask"%prop]).cpu().detach().numpy()]
                else:
                    predicted[prop] = [(result[prop][prop]*batch["%s_mask"%prop]).cpu().detach().numpy()]
              else:
                if prop in predicted:
                    predicted[prop] += [(result[prop]*batch["%s_mask"%prop]).cpu().detach().numpy()]
                else:
                    predicted[prop] = [(result[prop]*batch["%s_mask"%prop]).cpu().detach().numpy()]
          if prop =="eigenvalues" or prop in Eigenvalue_properties.properties:
                if prop in predicted:
                    predicted[prop] += [(result[prop]*batch["%s_mask"%prop]).cpu().detach().numpy()]
                else:
                    predicted[prop] = [(result[prop]*batch["%s_mask"%prop]).cpu().detach().numpy()]
          if prop == "delta_eigenvalues_pbe0_gbw" or prop=="delta_eigenvalues_pbe0_h2o":
            if prop in predicted:
                predicted[prop] += [(result[prop]*batch["eigenvalues_gbw_mask"]).cpu().detach().numpy()]
            else:
                predicted[prop] = [(result[prop]*batch["eigenvalues_gbw_mask"]).cpu().detach().numpy()]
          if prop == "energy":
            if prop in predicted:
                predicted[prop] += [(result[prop]).cpu().detach().numpy()]
            else:
                predicted[prop] = [(result[prop]).cpu().detach().numpy()]

        for prop in result:
          if prop =="eigenvalues_active":
                if prop in qm_values:
                    qm_values[prop] += [(batch[prop]*batch["%s_mask"%prop]).cpu().detach().numpy()]
                    qm_values["%s_mask"%prop] +=[batch["%s_mask"%prop].cpu().detach().numpy()]
                else:
                    qm_values[prop] = [(batch[prop]*batch["%s_mask"%prop]).cpu().detach().numpy()]
                    qm_values["%s_mask"%prop] = [batch["%s_mask"%prop].cpu().detach().numpy()]
          if prop == "delta_eigenvalues_pbe0_gbw" or prop =="delta_eigenvalues_pbe0_h2o":
                if prop in qm_values:
                    qm_values[prop] += [(batch[prop]*batch["eigenvalues_gbw_mask"]).cpu().detach().numpy()]
                    qm_values["eigenvalues_gbw_mask"] +=[batch["eigenvalues_gbw_mask"].cpu().detach().numpy()]
                else:
                    qm_values[prop] = [(batch[prop]*batch["eigenvalues_gbw_mask"]).cpu().detach().numpy()]
                    qm_values["eigenvalues_gbw_mask"] = [batch["eigenvalues_gbw_mask"].cpu().detach().numpy()]

          if prop =="eigenvalues" or prop in Eigenvalue_properties.properties:
                if prop in qm_values:
                    qm_values[prop] += [(batch[prop]).cpu().detach().numpy()]
                else:
                    qm_values[prop] = [(batch[prop]).cpu().detach().numpy()]
          if prop == "energy":
            if prop in qm_values:
                qm_values[prop] += [(batch[prop]).cpu().detach().numpy()]
            else:
                qm_values[prop] = [(batch[prop]).cpu().detach().numpy()]
        for metric in metrics:
          if "eigenvalues_active" in result:
              prop="eigenvalues_active"
              batch[prop]=batch[prop]*batch["%s_mask"%prop]
              if "eigenvalues_active_forces" in properties:
                  result[prop][prop] = result[prop][prop]*batch["%s_mask"%prop]
                  metric.add_batch(batch,result[prop])
              else:
                  result[prop]=  result[prop]#*batch["%s_mask"%prop][:len(result[prop][0])]
                  metric.add_batch(batch,result)
          else:
              metric.add_batch(batch,result)
    results = [
    metric.aggregate() for metric in metrics
    ]
    for prop in result:
        if prop == "eigenvalues_active" or prop=="eigenvalues":
            if "eigenvalues_active_forces" in properties:
                predicted['diabatic_H'] = [result[prop]['ev_diabatic'].cpu().detach().numpy()]
                #predicted['eigenvalues_active_forces'] = [result[prop]['eigenvalues_active_forces'].cpu().detach().numpy()]
    for p in predicted.keys():
        predicted[p]=np.vstack(predicted[p])
        if p in qm_values.keys():
            qm_values[p]=np.vstack(qm_values[p])
    if "eigenvalues_active_mask" in qm_values.keys():
        qm_values['eigenvalues_active_mask'] = np.array(qm_values["eigenvalues_active_mask"])
    if "eigenvalues_gbw_mask" in qm_values.keys():
        qm_values["eigenvalues_gbw_mask"] = np.array(qm_values["eigenvalues_gbw_mask"])
    if "eigenvalues_mask" in qm_values.keys():
        qm_values["eigenvalues_mask"] = np.array(qm_values["eigenvalues_mask"])
    prediction_path = os.path.join(args.modelpath,"evaluation_values.npz")
    prediction_path_qm = os.path.join(args.modelpath,"evaluation_qmvalues.npz")
    np.savez(prediction_path,**predicted)
    np.savez(prediction_path_qm,**qm_values)
    logging.info('Stored model predictions in {:s} ...'.format(prediction_path))

    return results

def run_prediction(model, loader, device, args):
    from tqdm import tqdm
    import numpy as np

    predicted = {}
    qm_values = {}
    for batch in tqdm(loader, ncols=120):
        batch = {
            k: v.to(device)
            for k, v in batch.items()
        }
        if device != "cuda" and args.parallel==True:
	        result = model.module(batch)
        else:
                result = model(batch)
        for prop in result:
            if prop in predicted:
                predicted[prop] += [result[prop].cpu().detach().numpy()]
            else:
                predicted[prop] = [result[prop].cpu().detach().numpy()]
        for prop in batch:
            if prop in qm_values:
                qm_values[prop] += [batch[prop].cpu().detach().numpy()]
            else:
                qm_values[prop] = [batch[prop].cpu().detach().numpy()]

    for p in predicted.keys():
        predicted[p] = np.vstack(predicted[p])
    #for p in qm_values.keys():
    #    qm_values[p]=np.vstack(qm_values[p])
    #prediction_path_qm = os.path.join(args.modelpath,"evaluation_qmvalues.npz")
    #np.savez(prediction_path_qm,**qm_values)
    prediction_path = os.path.join(args.modelpath, 'predictions.npz')
    np.savez(prediction_path, **predicted)
    logging.info('Stored model predictions in {:s}...'.format(prediction_path))

def get_model(args, n_orb, properties, atomref=None, mean=None, stddev=None, train_loader=None, parallelize=False,
              mode='train'):
    if args.model == 'schnet':
        representation = spk.representation.SchNet(args.features, args.features, args.interactions,
                                                   args.cutoff / units.Bohr, args.num_gaussians)

        property_output = schnet_ev.model.MultiStatePropertyModel(args.features, n_orb,n_neurons=args.n_neurons, properties=properties,
                                                                mean=mean, stddev=stddev, atomref=atomref,
                                                                n_layers=args.n_layers, diabatic=args.diabatic,hamiltonian = args.hamiltonian,
                                                                inverse_energy=args.inverse_nacs, force_mask = args.force_mask, force_mask_index = args.force_mask_index)

        model = spk.atomistic.AtomisticModel(representation, property_output)

    elif args.model == 'invD':
        representation = spk.representation.schnet.invD()
        n_in = (args.n_atoms*args.n_atoms-args.n_atoms)
        #neurons = number of neurons in each layer of the network; if None, it is divided by two in each layer
        property_output = schnet_ev.model.MultiStatePropertyModel(n_in, n_orb,n_neurons=args.n_neurons, properties=properties,
                                                                mean=mean, stddev=stddev, atomref=atomref,
                                                                n_layers=args.n_layers, diabatic=args.diabatic,hamiltonian = args.hamiltonian,
                                                                inverse_energy=args.inverse_nacs, force_mask=args.force_mask, force_mask_index = args.force_mask_index)

        model = spk.atomistic.AtomisticModel(representation, property_output)

    else:
        raise ValueError('Unknown model class:', args.model)

    if parallelize:
        model = nn.DataParallel(model)

    logging.info("The model you built has: %d parameters" % spk.utils.count_params(model))

    return model

def orca_opt(args, model, device):
    # reads QM.in and writes QM.out for optimizations with ORCA

    #read QMin
    QMin=read_QMin("QM.in")
    # get input for schnet
    schnet_inputs = QMin2schnet(QMin,device)
    #perform prediction
    schnet_outputs = model(schnet_inputs)
    for key,value in schnet_outputs.items():
        schnet_outputs[key] = value.cpu().detach().numpy()
    #transform to QM.out
    get_QMout(schnet_outputs, QMin,args)
    return

def get_QMout(schnet_outputs,QMin,args):
    from transform_prediction_QMout import QMout
    QMout(schnet_outputs,args.modelpath,args.nac_approx,QMin)
    return

def QMin2schnet(QMin,args):
    from schnetpack import Properties
    from ase import Atoms
    from schnetpack.environment import SimpleEnvironmentProvider, collect_atom_triples

    #environment_provider = SimpleEnvironmentProvider()
    environment_provider = get_environment_provider(args,device=device)
    schnet_inputs = dict()
    molecule = Atoms(QMin['atypes'],QMin['geom'])
    schnet_inputs[Properties.Z] = torch.LongTensor(molecule.numbers.astype(np.int))
    schnet_inputs[Properties.atom_mask] = torch.ones_like(schnet_inputs[Properties.Z]).float()
    #set positions
    positions = molecule.positions.astype(np.float32)
    schnet_inputs[Properties.R] = torch.FloatTensor(positions)
    #get atom environment
    nbh_idx, offsets = environment_provider.get_environment(molecule)
    #neighbours and neighour mask
    mask = torch.FloatTensor(nbh_idx) >= 0
    schnet_inputs[Properties.neighbor_mask] = mask.float()
    schnet_inputs[Properties.neighbors] = torch.LongTensor(nbh_idx.astype(np.int)) * mask.long()
    #get cells
    schnet_inputs[Properties.cell] = torch.FloatTensor(molecule.cell.astype(np.float32))
    schnet_inputs[Properties.cell_offset]= torch.FloatTensor(offsets.astype(np.float))
    collect_triples = False
    if collect_triples is not None:
        nbh_idx_j, nbh_idx_k, offset_idx_j, offset_idx_k = collect_atom_triples(nbh_idx)
        schnet_inputs[Properties.neighbor_pairs_j] = torch.LongTensor(nbh_idx_j.astype(np.int))
        schnet_inputs[Properties.neighbor_pairs_k] = torch.LongTensor(nbh_idx_k.astype(np.int))
        schnet_inputs[Properties.neighbor_pairs_mask] = torch.ones_like(schnet_inputs[Properties.neighbor_pairs_j]).float()                                                                 # Add batch dimension and move to CPU/GPU
    for key, value in schnet_inputs.items():
        schnet_inputs[key] = value.unsqueeze(0).to(device)

    return schnet_inputs

def read_QMin(filename):
    A2Bohr = 1/0.529177211
    data = open(filename, 'r').readlines()
    QMin = {}
    natoms = int(data[0].split()[0])
    QMin['NAtoms'] = natoms
    #unit
    for line in data:
        if "angstrom" in line or "Angstrom" in line:
            angstrom = True
            break
        else:
            angstrom = False
    #states
    for line in data:
        if "states" in line:
            s=line.split()[1:]
            break
    states = [ int(i) for i in s]
    n_doublets = 0
    n_unocc = 0
    n_quartets = 0
    for index,istate in enumerate(states):
        if index == int(0):
            n_occ = istate
        elif index == int(1):
            n_doublets = istate
        elif index == int(2):
            n_unocc = istate
        elif index == int(3):
            n_quartets = istate
    nmstates = n_occ + 2*n_doublets + 3*n_unocc + 4*n_quartets
    QMin['n_occ'] = n_occ
    QMin['n_doublets'] = n_doublets
    QMin['n_unocc'] = n_unocc
    QMin['n_quartets'] = n_quartets
    QMin['n_orb'] = nmstates

    #geometries
    atypes = []
    geom = []
    for curr_atom in range(natoms):
        atom_data = data[curr_atom+2].strip().split()
        curr_atype = atom_data[0]
        if angstrom == True:
            curr_geom = [float(coordinate) for coordinate in atom_data[1:4]]
            for xyz in range(3):
                curr_geom[xyz] = curr_geom[xyz] * A2Bohr
        else:
            curr_geom = [float(coordinate) for coordinate in atom_data[1:4]]
        atypes.append(curr_atype)
        geom.append(curr_geom)
    QMin['geom'] = np.array(geom)
    QMin['atypes'] = np.array(atypes)

    return QMin

def get_deltaE(model):
    return

if __name__ == '__main__':
    # Parse command line arguments
    parser = get_parser()
    args = parser.parse_args()

    # Determine the device
    device = torch.device("cuda" if args.cuda else "cpu")
    if args.mode != 'train':
        model = torch.load(os.path.join(args.modelpath, 'best_model'), map_location='cpu').to(device)
        if args.hessian == True:
            model.output_modules[0].output_dict['energy'].return_hessian = [True,1,1,1,1]
        else:
            if args.parallel==False:
                for p in  model.output_modules[0].output_dict.keys():
                    model.output_modules[0].output_dict[p].return_hessian = [False,1,1]
            else:
                for p in model.module.output_modules[0].output_dict.keys():
                    model.module.output_modules[0].output_dict[p].return_hessian = [False,1,1]

    if args.mode == 'opt':
        orca_opt(args, model, device)
        sys.exit()

    if args.mode == 'pred':
        pred_data = spk.data.AtomsData(args.datapath)
        pred_loader = spk.data.AtomsLoader(pred_data, batch_size=args.batch_size, num_workers=2, pin_memory=True)
        run_prediction(model, pred_loader, device, args)
        sys.exit(0)

    # Load settings
    argparse_dict = vars(args)
    jsonpath = os.path.join(args.modelpath, 'args.json')

    # Setup directories and files for training or load settings for restart.
    if args.mode == 'train':
        if args.overwrite and os.path.exists(args.modelpath):
            logging.info('Existing model will be overwritten...')
            rmtree(args.modelpath)

        if not os.path.exists(args.modelpath):
            os.makedirs(args.modelpath)

        spk.utils.to_json(jsonpath, argparse_dict)

        spk.utils.set_random_seed(args.seed)
        train_args = args

        # Read the tradeoff file or generate a default file
        tradeoff_file = os.path.join(args.modelpath, 'tradeoffs.yaml')

        if train_args.tradeoffs is None:
            schnet_ev.utils.generate_default_tradeoffs(tradeoff_file)
            tradeoffs = schnet_ev.utils.read_tradeoffs(tradeoff_file)
        else:
            tradeoffs = schnet_ev.utils.read_tradeoffs(args.tradeoffs)
            schnet_ev.utils.save_tradeoffs(tradeoffs, tradeoff_file)
    else:
        train_args = spk.utils.read_from_json(jsonpath)
        tradeoff_file = os.path.join(args.modelpath, 'tradeoffs.yaml')
        tradeoffs = schnet_ev.utils.read_tradeoffs(tradeoff_file)

    # Determine the properties to load based on the tradeoffs
    properties = [p for p in tradeoffs]
    if "eigenvalues_forces" in properties:
        tradeoffs['eigenvalues_active_forces']=tradeoffs['eigenvalues_forces']
    # Read and process the data using the properties found in the tradeoffs.
    logging.info('Loading {:s}...'.format(args.datapath))
    environment_provider = get_environment_provider(train_args,device=torch.device("cuda" if args.cuda else "cpu"))
    #previous function:
    dataset = spk.data.AtomsData(args.datapath, environment_provider=environment_provider, collect_triples=args.model == 'wacsf')
    # Determine the number of states based on the metadata
    n_orb = {}
    if "mean_energy" in dataset.get_metadata():
        mean_energy = dataset.get_metadata('mean_energy')
    else:
        mean_energy = int(0)
    n_orb['mean_energy']=mean_energy
    n_orb['n_act'] = int(1)
    offset=0
    for prop in properties:
        if prop in Eigenvalue_properties.properties or prop in Delta_EV_properties.properties:

        # if "eigenvalues" in properties or "eigenvalues_active" in properties or "eigenvalues_pbe0" in properties or "eigenvalues_pbe" in properties or "delta_eigenvalues_gbw" in properties or "eigenvalues_gbw" in properties:
            n_orb['n_act'] = dataset.get_metadata('active')
            if dataset.get_metadata("active_all"):
                n_orb["active_all"]= dataset.get_metadata("active_all")
            if dataset.get_metadata('offset'):
                offset = dataset.get_metadata('offset')
            else:
                offset = float(0.0)
    ##activate if only one state is learned or not all
    # only needed if unocc and occ orbitals are trained separately
    n_orb['n_occ']=int(1)
    n_orb['n_unocc']=int(0)
    n_orb['n_orb']=n_orb['n_occ']+n_orb['n_unocc']
    n_orb['offset']=offset
    n_orb["aggregation_mode"]=args.aggregation_mode
    if args.mode == 'eval':
        split_path = os.path.join(args.modelpath, 'split.npz')
        data_train, data_val, data_test = dataset.create_splits(*args.split, split_file=split_path)

        train_loader = spk.data.AtomsLoader(data_train, batch_size=args.batch_size, sampler=RandomSampler(data_train),
                                        num_workers=4, pin_memory=True)
        val_loader = spk.data.AtomsLoader(data_val, batch_size=args.batch_size, num_workers=2, pin_memory=True)

        logging.info("evaluating...")
        test_loader = spk.data.AtomsLoader(data_test, batch_size=args.batch_size,
                                           num_workers=2, pin_memory=True)
        evaluate(args, model, train_loader, val_loader, test_loader, device)
        logging.info("... done!")
        exit()
     # Splits the dataset in test, val, train sets
    split_path = os.path.join(args.modelpath, 'split.npz')

    if args.mode == 'train':
        if args.split_path is not None:
            copyfile(args.split_path, split_path)

    logging.info('Creating splits...')
    data_train, data_val, data_test = dataset.create_splits(*train_args.split, split_file=split_path)

    # Generate loaders for training
    train_loader = spk.data.AtomsLoader(data_train, batch_size=args.batch_size, sampler=RandomSampler(data_train),
                                        num_workers=4, pin_memory=True)
    val_loader = spk.data.AtomsLoader(data_val, batch_size=args.batch_size, num_workers=2, pin_memory=True)

    # Determine statistics for scaling
    if args.mode == 'train':
        # Set the random seed
        spk.utils.set_random_seed(args.seed)

        logging.info('Calculate statistics...')
        mean = {}
        stddev = {}
        # Compute mean and stddev for relevant properties
        for p in properties:
            if p in schnet_ev.data.Properties.normalize:
                mean_p, stddev_p = train_loader.get_statistics(p, True)
                mean_p = mean_p[p]
                stddev_p = stddev_p[p]
                mean[p] = mean_p
                stddev[p] = stddev_p
                logging.info('{:s} MEAN: {:20.11f} STDDEV: {:20.11f}'.format(p, mean_p.detach().cpu().numpy()[0],
                                                                             stddev_p.detach().cpu().numpy()[0]))
            else:
                mean[p], stddev[p] = None, None

    # Construct the model.
    model = get_model(train_args,
                      n_orb,
                      properties,
                      mean=mean,
                      stddev=stddev,
                      train_loader=train_loader,
                      parallelize=args.parallel,
                      mode=args.mode
                      ).to(device)

    if args.mode == 'train':
        logging.info("training...")
        batch_size = args.batch_size
        train(args, model, tradeoffs, train_loader, val_loader, device, n_orb,mean,stddev,offset)
        logging.info("...training done!")


from PyQt5 import QtCore, QtGui
from PyQt5.QtWidgets import QWidget, QDialog, QTableWidgetItem
from ui_mainscreen import Ui_MainScreen
from ui_wearabledialog import Ui_WearableDialog
from ui_zonedialog import Ui_ZoneDialog
from ui_devicedialog import Ui_DeviceDialog
from ui_configdialog import Ui_ConfigDialog
from ui_pausechangedialog import Ui_PauseChangeDialog
from ui_confirmdialog import Ui_ConfirmDialog
from drew_util import *
from constants import *

class UiControl:
    def __init__(self, systemState):
        self.systemState = systemState
        self.currXmlId = [-1, -1, -1, -1]
        self.afterPauseTask = TASK_NONE

        self.iconPause = QtGui.QIcon("./images/pause.png")
        self.iconActivate = QtGui.QIcon("./images/play.png")
        self.iconWait = QtGui.QIcon("./images/wait.png")

        self.mainWindow = QWidget()
        self.mainUi = Ui_MainScreen()
        self.mainUi.setupUi(self.mainWindow)

        self.pauseChangeDialog = QDialog()
        self.pauseChangeDialog.setWindowFlags(QtCore.Qt.CustomizeWindowHint|QtCore.Qt.WindowTitleHint|QtCore.Qt.WindowStaysOnTopHint)
        self.pauseChangeDialogUi = Ui_PauseChangeDialog()
        self.pauseChangeDialogUi.setupUi(self.pauseChangeDialog)
        self.pauseChangeGif = QtGui.QMovie("./images/waiting.gif")
        self.pauseChangeGif.setScaledSize(QtCore.QSize(140, 140))
        self.pauseChangeDialogUi.labelGIF.setMovie(self.pauseChangeGif)
        self.pauseChangeThread = PauseChangeThread(self.systemState)
        self.showWhenPauseDone = None

        self.dropdownRefreshThread = DropdownRefreshThread(self.systemState)
        self.doRefresh = -1

        self.statusRefreshThread1 = StatusRefreshThread()
        self.statusRefreshThread2 = StatusRefreshThread()
        self.statuRefreshThreadSelect = False

        self.delTypeId = -1
        self.confirmDialog = QDialog()
        self.confirmDialog.setWindowFlags(QtCore.Qt.CustomizeWindowHint|QtCore.Qt.WindowTitleHint)
        self.confirmDialogUi = Ui_ConfirmDialog()
        self.confirmDialogUi.setupUi(self.confirmDialog)

        self.dialogs = []
        for i in range(NUM_DIALOGS):
            dialog = QDialog()
            dialog.setWindowFlags(QtCore.Qt.CustomizeWindowHint|QtCore.Qt.WindowTitleHint)
            self.dialogs.append(dialog)

        self.dialogUis = []
        self.dialogUis.append(Ui_WearableDialog())
        self.dialogUis.append(Ui_ZoneDialog())
        self.dialogUis.append(Ui_DeviceDialog())
        self.dialogUis.append(Ui_ConfigDialog())
        for i in range(NUM_DIALOGS):
            self.dialogUis[i].setupUi(self.dialogs[i])

        self.tables = [self.mainUi.tableWearable, self.mainUi.tableZone, self.mainUi.tableDevice, self.mainUi.tableConfig]
        self.zoTable = self.mainUi.tableZoneOccupation
        self.dsTable = self.mainUi.tableDeviceState

        self.buttonGroups = []
        self.buttonGroups.append([self.mainUi.buttonEditWearable, self.mainUi.buttonDeleteWearable, self.mainUi.buttonNewWearable])
        self.buttonGroups.append([self.mainUi.buttonEditZone, self.mainUi.buttonDeleteZone, self.mainUi.buttonNewZone])
        self.buttonGroups.append([self.mainUi.buttonEditDevice, self.mainUi.buttonDeleteDevice, self.mainUi.buttonNewDevice])
        self.buttonGroups.append([self.mainUi.buttonEditConfig, self.mainUi.buttonClearConfig])

        self.connectUiElements()

        # populate tables with current system info
        self.populateHardwareTables()
        self.createZoneOccupationTable()
        self.createDeviceStateTable()

        # populate device type dropdown
        dropdown = self.dialogUis[TID_D].dropdownType
        for deviceType in DEVICE_TYPES.keys():
            dropdown.addItem(DEVICE_TYPES[deviceType], deviceType)

        # set initial button enabled/disabled state, etc
        for i in range(4):
            self.selectionUpdate(i)

        # initially, system should be unpaused
        self.mainUi.labelPauseStatus.setText("<b>System Active<b>")
        self.mainUi.buttonPause.setIcon(self.iconPause)
        self.manuallyPaused = False

        # now that everything is set up, let the other threads start doing their thing
        self.systemState.setSystemPause(RESUME)

        # since we start in the status tab, we want to make sure it's updating when we begin
        self.statusRefresh()

        # let's go!
        self.mainWindow.show()

    def connectUiElements(self):
        # "New" buttons
        self.mainUi.buttonNewWearable.clicked.connect(lambda: self.editWearable(True))
        self.mainUi.buttonNewZone.clicked.connect(lambda: self.editZone(True))
        self.mainUi.buttonNewDevice.clicked.connect(lambda: self.editDevice(True))

        # "Edit" buttons
        self.mainUi.buttonEditWearable.clicked.connect(lambda: self.editWearable(False))
        self.mainUi.buttonEditZone.clicked.connect(lambda: self.editZone(False))
        self.mainUi.buttonEditDevice.clicked.connect(lambda: self.editDevice(False))
        self.mainUi.buttonEditConfig.clicked.connect(lambda: self.editConfig())

        # "Delete" buttons
        self.mainUi.buttonDeleteWearable.clicked.connect(lambda: self.confirmDeletion(TID_W))
        self.mainUi.buttonDeleteZone.clicked.connect(lambda: self.confirmDeletion(TID_Z))
        self.mainUi.buttonDeleteDevice.clicked.connect(lambda: self.confirmDeletion(TID_D))
        self.mainUi.buttonClearConfig.clicked.connect(lambda: self.beginTaskWithPause(TASK_CLEAR_CONFIG))

        # deletion confirmation buttons
        self.confirmDialogUi.buttonNo.clicked.connect(lambda: self.confirmDialog.hide())
        self.confirmDialogUi.buttonYes.clicked.connect(lambda: self.beginTaskWithPause(TASK_DELETE_TABLE_ENTRY))

        # dialog "Cancel" buttons
        self.dialogUis[TID_W].buttonCancel.clicked.connect(lambda: self.cancelHardwareDialog(TID_W))
        self.dialogUis[TID_Z].buttonCancel.clicked.connect(lambda: self.cancelHardwareDialog(TID_Z))
        self.dialogUis[TID_D].buttonCancel.clicked.connect(lambda: self.cancelHardwareDialog(TID_D))
        self.dialogUis[TID_C].buttonCancel.clicked.connect(lambda: self.cancelHardwareDialog(TID_C))

        # dialog "Save" buttons
        self.dialogUis[TID_W].buttonSave.clicked.connect(lambda: self.saveWearable())
        self.dialogUis[TID_Z].buttonSave.clicked.connect(lambda: self.saveZone())
        self.dialogUis[TID_D].buttonSave.clicked.connect(lambda: self.saveDevice())
        self.dialogUis[TID_C].buttonSave.clicked.connect(lambda: self.saveConfig())

        # dialog "Refresh List" buttons
        self.dialogUis[TID_W].buttonRefresh.clicked.connect(lambda: self.refreshWearableDropdown())
        self.dialogUis[TID_Z].buttonRefresh.clicked.connect(lambda: self.initiateDropdownRefresh(TID_Z))
        self.dialogUis[TID_D].buttonRefresh.clicked.connect(lambda: self.initiateDropdownRefresh(TID_D))

        # enable/disable buttons when table selection changes, etc
        self.tables[TID_W].itemSelectionChanged.connect(lambda: self.selectionUpdate(TID_W))
        self.tables[TID_Z].itemSelectionChanged.connect(lambda: self.selectionUpdate(TID_Z))
        self.tables[TID_D].itemSelectionChanged.connect(lambda: self.selectionUpdate(TID_D))
        self.tables[TID_C].itemSelectionChanged.connect(lambda: self.selectionUpdate(TID_C))

        # manual system pause/resume button
        self.mainUi.buttonPause.clicked.connect(lambda: self.beginPauseChange((not self.systemState.systemIsPaused), True))
        self.pauseChangeThread.doneSignal.connect(lambda: self.finishPauseChange())

        self.dropdownRefreshThread.doneSignal.connect(lambda: self.finishDropdownRefresh())

        self.statusRefreshThread1.refreshSignal.connect(lambda: self.statusRefresh())
        self.statusRefreshThread2.refreshSignal.connect(lambda: self.statusRefresh())

        # link values of zone threshold slider & spinner
        spinner = self.dialogUis[TID_Z].spinnerThreshold
        slider = self.dialogUis[TID_Z].sliderThreshold
        spinner.valueChanged.connect(lambda: slider.setValue(spinner.value()))
        slider.valueChanged.connect(lambda: spinner.setValue(slider.value()))

        #
        tabs = self.mainUi.tabWidget
        tabs.currentChanged.connect(lambda: self.tabChange(tabs.currentIndex()))

    def tabChange(self, tabNum):
        if (tabNum == TAB_STAT):
            self.statusRefresh()
        elif (tabNum == TAB_SYS):
            self.createConfigTable()
        # other cases -> do nothing

    def beginPauseChange(self, desiredAction, isManualChange):
        if (desiredAction == self.systemState.systemIsPaused):
            return False # system is already in the desired state

        # manual changes are initiated by pressing the pause/resume button
        # automatic changes are initiated when editing any part of the system
        if (isManualChange):
            self.manuallyPaused = (desiredAction == PAUSE)

        if (desiredAction == PAUSE):
            waitText = "Pausing System..."
        else:
            waitText = "Resuming System... "
        self.mainUi.labelPauseStatus.setText(waitText)
        self.mainUi.buttonPause.setIcon(self.iconWait)

        self.pauseChangeDialog.setWindowTitle(waitText)
        self.pauseChangeGif.start()
        self.pauseChangeDialog.show()

        self.pauseChangeThread.setAction(desiredAction)
        self.pauseChangeThread.start()

        return True

    def finishPauseChange(self):
        if (self.systemState.systemIsPaused):
            doneText = "<b>System Paused</b>"
            newIcon = self.iconActivate
        else:
            doneText = "<b>System Active<b>"
            newIcon = self.iconPause

        self.mainUi.buttonPause.setIcon(newIcon)
        self.mainUi.labelPauseStatus.setText(doneText)

        self.pauseChangeGif.stop()
        self.pauseChangeDialog.hide()

        if (self.showWhenPauseDone != None):
            self.showWhenPauseDone.show()
            self.showWhenPauseDone = None

        if (self.doRefresh != -1):
            self.initiateDropdownRefresh(self.doRefresh)

        self.doRefresh = -1

        self.executeTask()

    def statusRefresh(self):
        # update each row (zone) in zone occupation table
        for row in range(self.zoTable.rowCount()):
            items = [self.zoTable.item(row, 0), self.zoTable.item(row, 1)]
            zone = self.systemState.getHardwareObjectByXmlId(TID_Z, items[0].data(5))
            self.updateZoneOccupationTableRow(zone, items)

        # update each row (zone) in zone occupation table
        for row in range(self.dsTable.rowCount()):
            items = [self.dsTable.item(row, 0), self.dsTable.item(row, 1)]
            device = self.systemState.getHardwareObjectByXmlId(TID_D, items[0].data(5))
            self.updateDeviceStateTableRow(device, items)

        # if still in status tab, need to keep refreshing
        if (self.mainUi.tabWidget.currentIndex() == TAB_STAT):
            # alternate between two refresh threads so signal emits don't "overlap"
            if (self.statuRefreshThreadSelect):
                self.statusRefreshThread1.start()
            else:
                self.statusRefreshThread2.start()
            self.statuRefreshThreadSelect = (not self.statuRefreshThreadSelect)

    def editWearable(self, isNew, isAfterPause=False):
        self.showWhenPauseDone = self.dialogs[TID_W]
        self.beginPauseChange(PAUSE, False)
        self.dialogUis[TID_W].labelInvalidName.setVisible(False)
        self.dialogUis[TID_W].labelInvalidWearable.setVisible(False)
        self.newFlag = isNew
        if (isNew):
            wearable = self.systemState.newWearable()
            self.currXmlId[TID_W] = wearable.xmlId
        else:
            wearable = self.systemState.getHardwareObjectByXmlId(TID_W, self.currXmlId[TID_W])
        self.dialogUis[TID_W].inputName.setText(wearable.name)
        self.refreshWearableDropdown()

    def editZone(self, isNew):
        self.showWhenPauseDone = self.dialogs[TID_Z]
        self.doRefresh = TID_Z
        pauseNeeded = self.beginPauseChange(PAUSE, False)
        if (not pauseNeeded):
            self.doRefresh = -1
            self.initiateDropdownRefresh(TID_Z)
        self.dialogUis[TID_Z].labelInvalidName.setVisible(False)
        self.dialogUis[TID_Z].labelInvalidModule.setVisible(False)
        self.newFlag = isNew
        if (isNew):
            zone = self.systemState.newZone()
            self.currXmlId[TID_Z] = zone.xmlId
        else:
            zone = self.systemState.getHardwareObjectByXmlId(TID_Z, self.currXmlId[TID_Z])
        self.dialogUis[TID_Z].inputName.setText(zone.name)
        self.dialogUis[TID_Z].spinnerThreshold.setValue(zone.threshold)
        if (not pauseNeeded):
            self.dialogs[TID_Z].show()

    def editDevice(self, isNew):
        self.showWhenPauseDone = self.dialogs[TID_D]
        self.doRefresh = TID_D
        pauseNeeded = self.beginPauseChange(PAUSE, False)
        if (not pauseNeeded):
            self.doRefresh = -1
            self.initiateDropdownRefresh(TID_D)
        self.dialogUis[TID_D].labelInvalidName.setVisible(False)
        self.dialogUis[TID_D].labelInvalidDevice.setVisible(False)
        self.newFlag = isNew
        if (isNew):
            device = self.systemState.newDevice()
            self.currXmlId[TID_D] = device.xmlId
        else:
            device = self.systemState.getHardwareObjectByXmlId(TID_D, self.currXmlId[TID_D])
        self.dialogUis[TID_D].inputName.setText(device.name)
        if (not pauseNeeded):
            self.dialogs[TID_D].show()

    def editConfig(self):
        self.showWhenPauseDone = self.dialogs[TID_C]
        self.beginPauseChange(PAUSE, False)
        self.dialogUis[TID_C].labelInvalidZone.setVisible(False)
        device = self.systemState.getHardwareObjectByXmlId(TID_D, self.currXmlId[TID_C])
        self.dialogUis[TID_C].labelConfig.setText("Configuration for \"" + device.name + "\"")
        self.populateConfigDropdowns(device)

    def cancelHardwareDialog(self, typeId):
        if (typeId != TID_C):
            if (self.newFlag):
                self.systemState.deleteHardwareObject(typeId, self.currXmlId[typeId])
            self.selectionUpdate(typeId)

        if (not self.manuallyPaused):
            self.beginPauseChange(RESUME, False)
        self.dialogs[typeId].hide()

    def saveWearable(self):
        givenName = (self.dialogUis[TID_W].inputName.text()).strip()
        if (self.systemState.nameInUse(TID_W, givenName, self.currXmlId[TID_W])):
            self.dialogUis[TID_W].labelInvalidName.setVisible(True)
            return
        else:
            self.dialogUis[TID_W].labelInvalidName.setVisible(False)

        selectedHwId = self.dialogUis[TID_W].dropdownWearable.currentData()
        if (selectedHwId == -1):
            self.dialogUis[TID_W].labelInvalidWearable.setVisible(True)
            return
        else:
            self.dialogUis[TID_W].labelInvalidWearable.setVisible(False)

        wearable = self.systemState.getHardwareObjectByXmlId(TID_W, self.currXmlId[TID_W])
        wearable.name = givenName
        wearable.hwId = selectedHwId
        self.updateWearableTable(wearable, self.newFlag)
        self.selectionUpdate(TID_W)
        if (not self.manuallyPaused):
            self.beginPauseChange(RESUME, False)
        self.dialogs[TID_W].hide()

    def saveZone(self):
        # disallow duplicate names
        givenName = (self.dialogUis[TID_Z].inputName.text()).strip()
        if (self.systemState.nameInUse(TID_Z, givenName, self.currXmlId[TID_Z])):
            self.dialogUis[TID_Z].labelInvalidName.setVisible(True)
            return
        else:
            self.dialogUis[TID_Z].labelInvalidName.setVisible(False)

        # disallow invalid hardware IDs
        selectedHwId = self.dialogUis[TID_Z].dropdownModule.currentData()
        if (selectedHwId == -1):
            self.dialogUis[TID_Z].labelInvalidModule.setVisible(True)
            return
        else:
            self.dialogUis[TID_Z].labelInvalidModule.setVisible(False)

        # input validated -- do the save operation
        zone = self.systemState.getHardwareObjectByXmlId(TID_Z, self.currXmlId[TID_Z])
        zone.name = givenName
        zone.hwId = selectedHwId
        zone.threshold = self.dialogUis[TID_Z].spinnerThreshold.value()
        self.updateZoneTable(zone, self.newFlag)
        self.selectionUpdate(TID_Z)

        # update zone occupation table in status tab
        if (self.newFlag):
            self.updateZoneOccupationTableRow(zone)
        else:
            for row in range(self.zoTable.rowCount()):
                if (self.zoTable.item(row, 0).data(5) == zone.xmlId):
                    items = [self.zoTable.item(row, 0), self.zoTable.item(row, 1)]
                    self.updateZoneOccupationTableRow(zone, items)
                    break

        self.updateDeviceStateTableRow
        if (not self.manuallyPaused):
            self.beginPauseChange(RESUME, False)
        self.dialogs[TID_Z].hide()

    def saveDevice(self):
        # disallow duplicate names
        givenName = (self.dialogUis[TID_D].inputName.text()).strip()
        if (self.systemState.nameInUse(TID_D, givenName, self.currXmlId[TID_D])):
            self.dialogUis[TID_D].labelInvalidName.setVisible(True)
            return
        else:
            self.dialogUis[TID_D].labelInvalidName.setVisible(False)

        # disallow invalid hardware IDs
        (selectedHwId, selectedListName) = self.dialogUis[TID_D].dropdownDevice.currentData()
        if (selectedHwId == 'INVALID'):
            self.dialogUis[TID_D].labelInvalidDevice.setVisible(True)
            return
        else:
            self.dialogUis[TID_D].labelInvalidDevice.setVisible(False)

        # input validated -- do the save operation
        device = self.systemState.getHardwareObjectByXmlId(TID_D, self.currXmlId[TID_D])
        device.name = givenName
        device.hwId = selectedHwId
        device.listName = selectedListName
        device.devType = self.dialogUis[TID_D].dropdownType.currentData()
        self.updateDeviceTable(device, self.newFlag)
        self.selectionUpdate(TID_D)

        # update device state table in status tab
        if (self.newFlag):
            self.updateDeviceStateTableRow(device)
        else:
            for row in range(self.dsTable.rowCount()):
                if (self.dsTable.item(row, 0).data(5) == device.xmlId):
                    items = [self.dsTable.item(row, 0), self.dsTable.item(row, 1)]
                    self.updateDeviceStateTableRow(device, items)
                    break

        if (not self.manuallyPaused):
            self.beginPauseChange(RESUME, False)
        self.dialogs[TID_D].hide()

    def saveConfig(self):
        configUi = self.dialogUis[TID_C]
        selectedZone = configUi.dropdownZone.currentData()
        if (selectedZone == -1):
            self.dialogUis[TID_C].labelInvalidZone.setVisible(True)
            return
        else:
            self.dialogUis[TID_C].labelInvalidZone.setVisible(False)

        device = self.systemState.getHardwareObjectByXmlId(TID_D, self.currXmlId[TID_C])
        device.zone = selectedZone
        device.enter = configUi.dropdownEntryAction.currentData()
        device.exit = configUi.dropdownExitAction.currentData()
        self.updateConfigTableEntry(device)
        self.selectionUpdate(TID_C)
        if (not self.manuallyPaused):
            self.beginPauseChange(RESUME, False)
        self.dialogs[TID_C].hide()

    def confirmDeletion(self, typeId):
        self.delTypeId = typeId
        hwObj = self.systemState.getHardwareObjectByXmlId(typeId, self.currXmlId[typeId])
        text = DEL_STR + HW_STR_DICT[typeId] + " \"" + hwObj.name + "\"?"
        self.confirmDialogUi.labelAreYouSure.setText(text)
        self.confirmDialogUi.labelWarning.setVisible(typeId == TID_Z)
        self.confirmDialog.show()

    def beginTaskWithPause(self, taskId):
        self.afterPauseTask = taskId
        ret = self.beginPauseChange(PAUSE, False)
        if (not ret):
            self.executeTask()

    def executeTask(self):
        task = self.afterPauseTask
        self.afterPauseTask = TASK_NONE

        if (task == TASK_NONE):
            return
        elif (task == TASK_CLEAR_CONFIG):
            self.clearConfig()
        elif (task == TASK_DELETE_TABLE_ENTRY):
            self.deleteTableEntry(self.delTypeId)

    def deleteTableEntry(self, typeId):
        xmlId = self.currXmlId[typeId]
        self.systemState.deleteHardwareObject(typeId, xmlId)
        self.tables[typeId].removeRow(self.tables[typeId].currentRow())
        self.removeFromStatusTable(typeId, xmlId)
        if (not self.manuallyPaused):
            self.beginPauseChange(RESUME, False)
        self.confirmDialog.hide()

    def clearConfig(self):
        device = self.systemState.getHardwareObjectByXmlId(TID_D, self.currXmlId[TID_C])
        device.zone = -1
        device.enter = 0
        device.exit = 0
        self.updateConfigTableEntry(device)
        if (not self.manuallyPaused):
            self.beginPauseChange(RESUME, False)

    def selectionUpdate(self, tableIdx):
        items = self.tables[tableIdx].selectedItems()
        if (len(items) > 0):
            self.currXmlId[tableIdx] = items[0].data(5)
        else:
            self.currXmlId[tableIdx] = -1
        shouldEnable = ( len(self.tables[tableIdx].selectedIndexes()) != 0)
        buttons = self.buttonGroups[tableIdx]
        for i in range(2): # only update edit & delete buttons
            buttons[i].setEnabled(shouldEnable)

    def refreshWearableDropdown(self):
        hwIdList = self.systemState.discoverHardware(TID_W)
        dropdown = self.dialogUis[TID_W].dropdownWearable
        dropdown.clear()

        for hwId in hwIdList:
            dropdown.addItem(str(hwId), hwId)

        if (not self.newFlag): # add and select current hwId if editing existing wearable
            wearable = self.systemState.getHardwareObjectByXmlId(TID_W, self.currXmlId[TID_W])
            dropdown.insertItem(0, str(wearable.hwId), int(wearable.hwId))
            dropdown.setCurrentIndex(0)
        elif (len(hwIdList) == 0):
            dropdown.addItem("[no unassigned wearables detected]", -1)

    def initiateDropdownRefresh(self, typeId):
        if (typeId == TID_Z):
            titleText = "Searching for Zone Modules..."
        else: # typeId = TID_D
            titleText = "Searching for Devices..."
        self.pauseChangeDialog.setWindowTitle(titleText)
        self.dropdownRefreshThread.setTypeId(typeId)
        self.dropdownRefreshThread.start()
        self.pauseChangeGif.start()
        self.pauseChangeDialog.show()

    def finishDropdownRefresh(self):
        if (self.dropdownRefreshThread.typeId == TID_Z):
            self.refreshZoneModuleDropdown()
        else: # typeId == TID_D
            self.refreshDeviceDropdown()

        self.pauseChangeGif.stop()
        self.pauseChangeDialog.hide()

    def refreshZoneModuleDropdown(self):
        hwIdList = self.dropdownRefreshThread.hwIdList
        dropdown = self.dialogUis[TID_Z].dropdownModule
        dropdown.clear()

        for hwId in hwIdList:
            dropdown.addItem(str(hwId), hwId)

        if (not self.newFlag): # add and select current hwId if editing existing zone
            zone = self.systemState.getHardwareObjectByXmlId(TID_Z, self.currXmlId[TID_Z])
            dropdown.insertItem(0, str(zone.hwId), int(zone.hwId))
            dropdown.setCurrentIndex(0)
        elif (len(hwIdList) == 0):
            dropdown.addItem("[no unassigned modules detected]", -1)

    def refreshDeviceDropdown(self):
        hwTupleList = self.dropdownRefreshThread.hwIdList
        dropdown = self.dialogUis[TID_D].dropdownDevice
        dropdown.clear()

        for (hwId, listName) in hwTupleList:
            dropdown.addItem(hwId + " - " + listName, (hwId, listName))

        # add and select current hwId if editing existing zone
        if (not self.newFlag):
            device = self.systemState.getHardwareObjectByXmlId(TID_D, self.currXmlId[TID_D])
            dropdown.insertItem(0, device.hwId + " - " + device.listName, (device.hwId, device.listName))
            dropdown.setCurrentIndex(0)
        elif (len(hwTupleList) == 0):
            dropdown.addItem("[no unassigned devices detected]", 'INVALID')

    def populateConfigDropdowns(self, device):
        configUi = self.dialogUis[TID_C]
        dropdown = configUi.dropdownZone
        dropdown.clear()

        for zone in self.systemState.dicts[TID_Z].values():
            dropdown.addItem(zone.name, zone.xmlId)
            if (zone.xmlId == device.zone):
                dropdown.setCurrentIndex(dropdown.count()-1)

        if (dropdown.count() == 0):
            dropdown.addItem("[no zones exist]", -1)

        actions = DEVICE_ACTIONS[device.devType]
        dropdown = configUi.dropdownEntryAction
        dropdown.clear()
        dropdown2 = configUi.dropdownExitAction
        dropdown2.clear()

        for actionId in DEVICE_ACTIONS[device.devType].keys():
            dropdown.addItem(actions[actionId], actionId)
            if (actionId == device.enter):
                dropdown.setCurrentIndex(dropdown.count()-1)
            dropdown2.addItem(actions[actionId], actionId)
            if (actionId == device.exit):
                dropdown2.setCurrentIndex(dropdown2.count()-1)

    def populateHardwareTables(self):
        for wearable in self.systemState.dicts[TID_W].values():
            self.updateWearableTable(wearable, True)

        for zone in self.systemState.dicts[TID_Z].values():
            self.updateZoneTable(zone, True)

        for device in self.systemState.dicts[TID_D].values():
            self.updateDeviceTable(device, True)

    def updateWearableTable(self, wearable, isNew):
        if (isNew):
            self.tables[TID_W].insertRow(0)
            # "Name" item with xmlId data
            item = QTableWidgetItem()
            item.setTextAlignment(ITEM_ALIGN_FLAGS)
            item.setFlags(ITEM_INTERACT_FLAGS)
            item.setText(wearable.name)
            item.setData(5, wearable.xmlId)
            self.tables[TID_W].setItem(0, 0, item)
            # "Wearable ID" item
            item = QTableWidgetItem()
            item.setTextAlignment(ITEM_ALIGN_FLAGS)
            item.setFlags(ITEM_INTERACT_FLAGS)
            item.setText(str(wearable.hwId))
            self.tables[TID_W].setItem(0, 1, item)
        else:
            items = self.tables[TID_W].selectedItems()
            items[0].setText(wearable.name)
            items[1].setText(str(wearable.hwId))

    def updateZoneTable(self, zone, isNew):
        if (isNew):
            self.tables[TID_Z].insertRow(0)
            # "Name" item with xmlId data
            item = QTableWidgetItem()
            item.setTextAlignment(ITEM_ALIGN_FLAGS)
            item.setFlags(ITEM_INTERACT_FLAGS)
            item.setText(zone.name)
            item.setData(5, zone.xmlId)
            self.tables[TID_Z].setItem(0, 0, item)
            # "Module ID" item
            item = QTableWidgetItem()
            item.setTextAlignment(ITEM_ALIGN_FLAGS)
            item.setFlags(ITEM_INTERACT_FLAGS)
            item.setText(str(zone.hwId))
            self.tables[TID_Z].setItem(0, 1, item)
            # "Threshold" item
            item = QTableWidgetItem()
            item.setTextAlignment(ITEM_ALIGN_FLAGS)
            item.setFlags(ITEM_INTERACT_FLAGS)
            item.setText(str(zone.threshold))
            self.tables[TID_Z].setItem(0, 2, item)
        else:
            items = self.tables[TID_Z].selectedItems()
            items[0].setText(zone.name)
            items[1].setText(str(zone.hwId))
            items[2].setText(str(zone.threshold))

    def updateDeviceTable(self, device, isNew):
        if (isNew):
            self.tables[TID_D].insertRow(0)
            # "Name" item with xmlId data
            item = QTableWidgetItem()
            item.setTextAlignment(ITEM_ALIGN_FLAGS)
            item.setFlags(ITEM_INTERACT_FLAGS)
            item.setText(device.name)
            item.setData(5, device.xmlId)
            self.tables[TID_D].setItem(0, 0, item)
            # "Bluetooth Address" item
            item = QTableWidgetItem()
            item.setTextAlignment(ITEM_ALIGN_FLAGS)
            item.setFlags(ITEM_INTERACT_FLAGS)
            item.setText(device.hwId)
            self.tables[TID_D].setItem(0, 1, item)
            # "Type" item
            item = QTableWidgetItem()
            item.setTextAlignment(ITEM_ALIGN_FLAGS)
            item.setFlags(ITEM_INTERACT_FLAGS)
            item.setText(DEVICE_TYPES[device.devType])
            self.tables[TID_D].setItem(0, 2, item)
        else:
            items = self.tables[TID_D].selectedItems()
            items[0].setText(device.name)
            items[2].setText(device.hwId)
            items[2].setText(DEVICE_TYPES[device.devType])

    def createConfigTable(self):
        configTable = self.tables[TID_C]
        configTable.clearContents()
        for i in range(configTable.rowCount()):
            configTable.removeRow(0)
        for device in self.systemState.dicts[TID_D].values():
            configTable.insertRow(0)
            # "Device" item with xmlId data
            item = QTableWidgetItem()
            item.setTextAlignment(ITEM_ALIGN_FLAGS)
            item.setFlags(ITEM_INTERACT_FLAGS)
            item.setText(device.name)
            item.setData(5, device.xmlId)
            configTable.setItem(0, 0, item)
            # "Containing Zone" item
            item = QTableWidgetItem()
            item.setTextAlignment(ITEM_ALIGN_FLAGS)
            item.setFlags(ITEM_INTERACT_FLAGS)
            if (device.zone == -1):
                text = "[unassigned]"
            else:
                text = self.systemState.dicts[TID_Z][device.zone].name
            item.setText(text)
            item.setData(5, device.zone)
            configTable.setItem(0, 1, item)
            # "Entry Action" item
            item = QTableWidgetItem()
            item.setTextAlignment(ITEM_ALIGN_FLAGS)
            item.setFlags(ITEM_INTERACT_FLAGS)
            item.setText(DEVICE_ACTIONS[device.devType][device.enter])
            configTable.setItem(0, 2, item)
            # "Exit Action" item
            item = QTableWidgetItem()
            item.setTextAlignment(ITEM_ALIGN_FLAGS)
            item.setFlags(ITEM_INTERACT_FLAGS)
            item.setText(DEVICE_ACTIONS[device.devType][device.exit])
            configTable.setItem(0, 3, item)

    def updateConfigTableEntry(self, device):
        items = self.tables[TID_C].selectedItems()
        if (device.zone == -1):
            text = "[unassigned]"
        else:
            text = self.systemState.dicts[TID_Z][device.zone].name
        items[1].setText(text)
        items[1].setData(5, device.zone)
        items[2].setText(DEVICE_ACTIONS[device.devType][device.enter])
        items[3].setText(DEVICE_ACTIONS[device.devType][device.exit])

    def createZoneOccupationTable(self):
        for zone in self.systemState.dicts[TID_Z].values():
            self.updateZoneOccupationTableRow(zone)

    def updateZoneOccupationTableRow(self, zone, tableItems=None):
        if (tableItems == None):
            tableItems = []
            for i in range(2):
                item = QTableWidgetItem()
                item.setTextAlignment(ITEM_ALIGN_FLAGS)
                item.setFlags(ITEM_INTERACT_FLAGS)
                tableItems.append(item)
            self.zoTable.insertRow(0)
            self.zoTable.setItem(0, 0, tableItems[0])
            self.zoTable.setItem(0, 1, tableItems[1])

        tableItems[0].setText(zone.name)
        tableItems[0].setData(5, zone.xmlId)

        wearablesPresent = ""
        firstOneFlag = True
        signalDataList = zone.wearablesInZone.items()

        counter = 0
        for wearbleId, signalData in signalDataList:
            wearable = self.systemState.getHardwareObjectByHwId(TID_W, wearbleId)
            if (not signalData.isInZone):
                continue
            counter += 1
            if (firstOneFlag):
                firstOneFlag = False
            else:
                wearablesPresent = wearablesPresent + ", "
            wearablesPresent = wearablesPresent + wearable.name

        if (counter == 0):
            wearablesPresent = "[none]"

        tableItems[1].setText(wearablesPresent)

    def createDeviceStateTable(self):
        for device in self.systemState.dicts[TID_D].values():
            self.updateDeviceStateTableRow(device)

    def updateDeviceStateTableRow(self, device, tableItems=None):
        if (tableItems == None):
            tableItems = []
            for i in range(2):
                item = QTableWidgetItem()
                item.setTextAlignment(ITEM_ALIGN_FLAGS)
                item.setFlags(ITEM_INTERACT_FLAGS)
                tableItems.append(item)
            self.dsTable.insertRow(0)
            self.dsTable.setItem(0, 0, tableItems[0])
            self.dsTable.setItem(0, 1, tableItems[1])

        tableItems[0].setText(device.name)
        tableItems[0].setData(5, device.xmlId)
        tableItems[1].setText(DEVICE_STATES[device.devType][device.state])

    def removeFromStatusTable(self, typeId, xmlId):
        if (typeId != TID_Z and typeId != TID_D):
            return
        else:
            if (typeId == TID_Z):
                table = self.zoTable
            else: # typeId == TID_D
                table = self.dsTable

            for row in range(table.rowCount()):
                if ((table.item(row, 0)).data(5) == xmlId):
                    table.removeRow(row)
                    return

class PauseChangeThread(QtCore.QThread):
    doneSignal = QtCore.pyqtSignal()
    action = RESUME
    isDone = False

    def __init__(self, systemState):
        super().__init__()
        self.systemState = systemState

    def setAction(self, action):
        self.action = action

    def run(self):
        self.isDone = False
        self.systemState.setSystemPause(self.action)
        self.doneSignal.emit()
        self.isDone = True

class DropdownRefreshThread(QtCore.QThread):
    doneSignal = QtCore.pyqtSignal()
    typeId = -1
    hwIdList = None

    def __init__(self, systemState):
        super().__init__()
        self.systemState = systemState

    def setTypeId(self, typeId):
        self.typeId = typeId

    def run(self):
        self.hwIdList = self.systemState.discoverHardware(self.typeId)
        self.doneSignal.emit()

class StatusRefreshThread(QtCore.QThread):
    refreshSignal = QtCore.pyqtSignal()

    def run(self):
        time.sleep(0.25)
        self.refreshSignal.emit()


import os
import json
from tracemalloc import start
import discord
from discord.ext import commands
import asyncio

INITIAL_EXTENSIONS = [
    'cogs.talk_and_question',
    'cogs.get_image',
    ]

try:
    intents = discord.Intents.all()
    bot = commands.Bot(command_prefix='??',intents=intents)
except Exception as  e:
    print(e)

if os.path.exists('token.json') :
    pass
else :
    check_token = {'token': 'plz replace tihs text'}
    check_1PSID = {'1PSID': 'plz replace tihs text'}
    with open('token.json', 'w') as f:
     jsondata = [check_token, check_1PSID]
     json.dump(jsondata, f, indent=2)
     print('token.json created plz edit it file. file location is RootDirectory/token.json\r\n')
     exit()

@bot.event
async def on_ready():
    print(f'complete logging in! {bot.user}')
    try:
         await bot.tree.sync()
    except Exception as e:
        print(e)

@bot.event
async def on_message(message):
    if message.author == bot.user:
        return

    if message.content.startswith('$こんにちは'):
        await message.channel.send('こんにちは！')

with open('token.json', 'r') as f:
     data = json.load(f)
bot_token = data[0]['token']
try:
 async def load_extension():
  for cog in INITIAL_EXTENSIONS:
   await bot.load_extension(cog)
except Exception as  e:
    print(e)

async def main():
 async with bot:
  await load_extension()
  await bot.start(bot_token)
try:
   asyncio.run(main())
except Exception as e:
    print(e)
    exit()


from typing import List
class DSU:
    def __init__(self, N):
        self.pre = [i for i in range(N)]

    def find(self, root):
        if self.pre[root] != root: self.pre[root] = self.find(self.pre[root])
        return self.pre[root]

    def union(self, x, y):
        xroot = self.find(x)
        yroot = self.find(y)
        if xroot != yroot:
            self.pre[x] = yroot
            self.pre[xroot] = yroot

class Solution:
    def removeStones(self, stones: List[List[int]]) -> int:
        N = len(stones)
        dsu = DSU(20000)
        for x, y in stones: dsu.union(x, y + 10000)

        return N - len(set(dsu.find(x) for x, y in stones))

if __name__ == '__main__':
    solution = Solution()
    print(solution.removeStones([[0,0],[0,1],[1,0],[1,2],[2,1],[2,2]]))


from flask import make_response
from app.utils.http_code import HTTP_404_NOT_FOUND,HTTP_500_INTERNAL_SERVER_ERROR
from flask import jsonify

def handle_404_error(error):
    """ Return 404 error in json for api"""

    return make_response(
        jsonify({
            'message': {
                'error' :'Not found'
            }
        }),
        HTTP_404_NOT_FOUND
    )

def handle_500_error(error):
    """ Return 500 error in json for api"""

    return make_response(
        jsonify({
            'message': {
                'error' :'Internal Server Error'
            }
        }),
        HTTP_500_INTERNAL_SERVER_ERROR
    )


# -*- coding: utf-8 -*-
"""Project Euler: Problem 61

Find the sum of the only ordered set of six cyclic 4-digit numbers
for which each polygonal type: triangle, square, pentagonal, hexagonal,
heptagonal, and octagonal, is represented by a different number
in the set.

https://projecteuler.net/problem=61
"""
from typing import List
import itertools

def generate_4digit_figurate_number(sides: int) -> List[int]:
    """Generates all 4-digit figurate number for a given number of sides

    Args:
        sides: the number of "sides" for the figurate number;
               e.g. 3 for triangular, 8 for octagonal, etc.

    Returns:
        List of all 4-digit figurate numbers
    """
    if not 2 < sides < 9:
        return []
    figurate_list = []
    figurate_number = 0
    n = 1
    while True:
        if sides == 3:
            figurate_number = n * (n + 1) // 2
        if sides == 4:
            figurate_number = n ** 2
        if sides == 5:
            figurate_number = n * (3 * n - 1) // 2
        if sides == 6:
            figurate_number = n * (2 * n - 1)
        if sides == 7:
            figurate_number = n * (5 * n - 3) // 2
        if sides == 8:
            figurate_number = n * (3 * n - 2)

        if 10000 > figurate_number > 999:
            figurate_list.append(figurate_number)
        elif figurate_number > 10000:
            return figurate_list
        n += 1

def is_cyclical(n: int, m: int) -> bool:
    """Finds if n and m are cyclical numbers

    Args:
        n: first number to check (looks at last two digits)
        m: second number to check (looks at first two digits)

    Returns:
        True if first two digits of m equal last two digits of n
    """
    if n % 100 == m // 100:
        return True
    else:
        return False

def find_cyclicals(numbers: List[int], comparator: int) -> List[int]:
    """Finds all numbers in a list that are cyclicals of a comparator.

    Args:
        numbers: list to search through
        comparator: number of which cyclicals to find

    Returns:
        List of all cyclicals of the number
    """
    return [n for n in numbers if is_cyclical(comparator, n)]

def main():
    numbers = [generate_4digit_figurate_number(n) for n in range(3, 9)]
    possible_permutations = list(itertools.permutations(range(0, 6)))
    possibilities = []
    for order in possible_permutations:
        cyclical_groups = []
        # Find initial list
        for item in numbers[order[0]]:
            for n in find_cyclicals(numbers[order[1]], item):
                cyclical_groups.append([item, n])
        # Generate further steps
        for cycle in cyclical_groups:
            for i in range(2, 6):
                for n in find_cyclicals(numbers[order[i]], cycle[-1]):
                    cycle.append(n)
        # Trim non-6 element lists
        cyclical_groups = [c for c in cyclical_groups if (len(c) == 6)]
        # See if all elements are cyclical (also checks first and last elements' cycle)
        # For some reason the find_cyclicals function returns non-cyclical combinations
        cyclical_groups = [c for c in cyclical_groups if all([is_cyclical(c[i], c[i + 1]) for i in range(-1, 5)])]
        for c in cyclical_groups:
            possibilities.append(c)
    # Sum possibilities and return the set (only unique value)
    print(set([sum(p) for p in possibilities]))

if __name__ == '__main__':
    main()


#================================================================
#
#   File name   : indicators.py
#   Author      : PyLessons
#   Created date: 2021-01-20
#   Website     : https://pylessons.com/
#   GitHub      : https://github.com/pythonlessons/RL-Bitcoin-trading-bot
#   Description : Used to plot 5 indicators with OHCL bars
#
#================================================================
import pandas as pd
from ta.trend import SMAIndicator, macd, PSARIndicator
from ta.volatility import BollingerBands
from ta.momentum import rsi
from utils import Plot_OHCL

def AddIndicators(df):
    # Add Simple Moving Average (SMA) indicators
    df["sma7"] = SMAIndicator(close=df["Close"], window=7, fillna=True).sma_indicator()
    df["sma25"] = SMAIndicator(close=df["Close"], window=25, fillna=True).sma_indicator()
    df["sma99"] = SMAIndicator(close=df["Close"], window=99, fillna=True).sma_indicator()

    # Add Bollinger Bands indicator
    indicator_bb = BollingerBands(close=df["Close"], window=20, window_dev=2)
    df['bb_bbm'] = indicator_bb.bollinger_mavg()
    df['bb_bbh'] = indicator_bb.bollinger_hband()
    df['bb_bbl'] = indicator_bb.bollinger_lband()

    # Add Parabolic Stop and Reverse (Parabolic SAR) indicator
    indicator_psar = PSARIndicator(high=df["High"], low=df["Low"], close=df["Close"], step=0.02, max_step=2, fillna=True)
    df['psar'] = indicator_psar.psar()

    # Add Moving Average Convergence Divergence (MACD) indicator
    df["MACD"] = macd(close=df["Close"], window_slow=26, window_fast=12, fillna=True) # mazas

    # Add Relative Strength Index (RSI) indicator
    df["RSI"] = rsi(close=df["Close"], window=14, fillna=True) # mazas

    return df

if __name__ == "__main__":
    df = pd.read_csv('./pricedata.csv')
    df = df.sort_values('Date')
    df = AddIndicators(df)

    test_df = df[-400:]

    Plot_OHCL(df)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

""""
@author:Yoann Poupart

File containing an example of conditional GAN training.

"""

# Libraries imports
from dcgan import FLEXI_C_DCGAN

import scipy.stats as sc
import numpy as np
import matplotlib.pyplot as plt
import tikzplotlib

from keras.models import Model, load_model
from keras.layers import Input
from keras.optimizers import Adam
from sklearn.utils import shuffle

# Class definition

# Functions definition

# Constants definition
BLOCK_SIZE = 90
SAMPLING_SIZE = 165
DATA_SET_SIZE = 2 ** 5 * 3 ** 5 * 5 ** 2
FILE_X = './Samples/real_ds_cond_trend_temp_glob.npy'
FILE_Y = './Samples/real_lab_cond_trend_temp_glob.npy'
FILE_GEN_STEP = "./Networks/gan_cond_trend_temp_step_gen"
FILE_DISC_STEP = "./Networks/gan_cond_trend_temp_step_disc"

NUM_EPOCHS = 1  # 50_000
BATCH_SIZE = 2 ** 3 * 3 ** 5
INIT_LR = 4e-4
LAST_ACTI = "linear"
H = 10
FACT_LR = 1
GRAD_NORM = 1
SUB_SET_SIZE = 1000

TYPE = 'student'
SHAPE = 50
SIGMA_NOISE_DATA = 1e-3
W_COND = 2

SIMU_GAN = True
SAVE_STEPS = False
SAVE_FIG = False
SAVE_TEX = False
FS = 30

if __name__ == '__main__':
    if SIMU_GAN:
        X = np.load(FILE_X)
        Y = np.load(FILE_Y)
        if SAMPLING_SIZE == 165:
            print("[INFO] building generator...")
            gen = FLEXI_C_DCGAN.build_generator(w_end=SAMPLING_SIZE, w_start=SAMPLING_SIZE,
                                                w_1=110, s_1=1, w_2=150, final_act=LAST_ACTI, h=H, w_cond=W_COND)
            print("[INFO] building discriminator...")
            disc = FLEXI_C_DCGAN.build_discriminator(w_start=SAMPLING_SIZE,
                                                     w_1=150, w_2=90, w_3=15, h=H, w_cond=W_COND)
        elif SAMPLING_SIZE == 156:
            print("[INFO] building generator...")
            gen = FLEXI_C_DCGAN.build_generator(w_end=SAMPLING_SIZE, w_start=SAMPLING_SIZE,
                                                w_1=100, s_1=1, w_2=140, final_act=LAST_ACTI, h=H, w_cond=W_COND)
            print("[INFO] building discriminator...")
            disc = FLEXI_C_DCGAN.build_discriminator(w_start=SAMPLING_SIZE,
                                                     w_1=140, w_2=80, w_3=15, h=H, w_cond=W_COND)
        else:
            raise NotImplementedError
        try:
            gen = load_model(FILE_GEN_STEP, compile=False)
            disc = load_model(FILE_DISC_STEP, compile=False)
        except:
            pass
        print('[INFO] test...')
        discOpt = Adam(learning_rate=INIT_LR, beta_1=0.999, decay=INIT_LR / NUM_EPOCHS, clipnorm=GRAD_NORM)
        disc.compile(loss="binary_crossentropy", optimizer=discOpt)

        zeros = np.zeros((BATCH_SIZE, SAMPLING_SIZE))
        zeros_coord = np.random.uniform(-90, 90, size=(BATCH_SIZE, 2)) * np.array([2, 1])
        _ = gen.predict((zeros, zeros_coord), verbose=1)
        _ = disc.predict((zeros, zeros_coord), verbose=1)
        print(f"[INFO] discriminator summary : ")
        disc.summary()
        print(f"[INFO] generator summary : ")
        gen.summary()

        print("[INFO] building GAN...")
        disc.trainable = False
        data_input = Input(shape=(SAMPLING_SIZE,))
        coord_input = Input(shape=(2,))
        ganOutput = disc((gen((data_input, coord_input)), coord_input))
        gan = Model((data_input, coord_input), ganOutput)
        ganOpt = Adam(learning_rate=INIT_LR / 4, beta_1=0.5, decay=INIT_LR / 4 / NUM_EPOCHS, clipnorm=GRAD_NORM)
        gan.compile(loss="binary_crossentropy", optimizer=discOpt)

        print("[INFO] starting training...")
        benchmark_noise = np.random.normal(size=(1, SAMPLING_SIZE))
        benchmark_coord = np.random.uniform(-90, 90, size=(1, 2)) * np.array([2, 1])
        disc_losses = []
        gen_losses = []

        for epoch in range(0, NUM_EPOCHS):
            print("[INFO] starting epoch {} of {}...".format(epoch + 1,
                                                             NUM_EPOCHS))
            X, Y = shuffle(X, Y)
            batchesPerEpoch = int(X.shape[0] / BATCH_SIZE)
            for i in range(0, batchesPerEpoch):
                batch_data = X[i * BATCH_SIZE:(i + 1) * BATCH_SIZE, :] + np.random.normal(loc=SIGMA_NOISE_DATA, size=(
                    BATCH_SIZE, SAMPLING_SIZE))
                batch_coord = Y[i * BATCH_SIZE:(i + 1) * BATCH_SIZE]

                noise = np.random.normal(size=(BATCH_SIZE, SAMPLING_SIZE))
                gen_coord = np.random.uniform(-90, 90, size=(BATCH_SIZE, 2)) * np.array([2, 1])
                gen_data = gen.predict((noise, gen_coord), verbose=0)

                X_tr = np.concatenate((batch_data, gen_data))
                y_tr = np.array(([1] * BATCH_SIZE) + ([0] * BATCH_SIZE))
                y_lab_tr = np.concatenate((batch_coord, gen_coord))
                (X_tr, y_tr, y_lab_tr) = shuffle(X_tr, y_tr, y_lab_tr)
                disc_loss = disc.train_on_batch((X_tr, y_lab_tr), y_tr)
                disc_losses.append(disc_loss)

                noise = np.random.normal(size=(BATCH_SIZE, SAMPLING_SIZE))
                gen_coord = np.random.uniform(-90, 90, size=(BATCH_SIZE, 2)) * np.array([2, 1])
                fake_labels = np.array([1] * BATCH_SIZE)
                gen_loss = gan.train_on_batch((noise, gen_coord), fake_labels)
                gen_losses.append(gen_loss)

            if NUM_EPOCHS >= 10:
                if epoch % (NUM_EPOCHS // 10) == 0:
                    gen_sample = gen.predict((benchmark_noise, benchmark_coord), verbose=0)
                    plt.subplot(2, 5, epoch // (NUM_EPOCHS // 10) + 1)
                    plt.plot(gen_sample.mean(axis=0))
            if SAVE_STEPS:
                gen.save(FILE_GEN_STEP, include_optimizer=True)
                disc.save(FILE_DISC_STEP, include_optimizer=True)

        if SAVE_FIG:
            if SAVE_TEX:
                tikzplotlib.save("./Images/states.tex")
            else:
                plt.savefig("./Images/States.png")
            plt.close()
        else:
            plt.show()

        plt.figure(figsize=(15, 15))
        steps = 1 + np.arange(len(gen_losses))
        plt.plot(steps, gen_losses, label="Loss of the generator")
        plt.plot(steps, disc_losses, label="Loss of the discriminator")
        plt.xticks(fontsize=FS)
        plt.yticks(fontsize=FS)
        plt.xlabel("Step", fontsize=FS)
        plt.ylabel(r"Loss", fontsize=FS)
        plt.legend(loc='best', fontsize=FS)
        if SAVE_FIG:
            if SAVE_TEX:
                tikzplotlib.save("Images/Loss.tex")
            else:
                plt.savefig("./Images/Loss.png")
            plt.close()
        else:
            plt.show()


import requests
import selenium
from selenium import webdriver
from bs4 import BeautifulSoup

from pymongo import MongoClient
import datetime

client = MongoClient('localhost', 27017)
db = client.dbsparta
collection = db.collection
docs = collection.find({})
weekno = datetime.datetime.today().weekday()

def update():
    players = list(db.myfantasystar.find({}, {'_id': False}, {'position': False}, {'match': False}, {'goal': False},
                                         {'assi': False}, {'img_url': False}).sort('name', ))
    if weekno == 2:
        players['like'] = '0'
        players['pre_like'] = '0'
    else:
        players['pre_like'] = players['like']
        players['like'] = '0'

    print(players)

    # like_info = {
    #     'like': players['pre_like'],
    #     'pre_like': players['like'],
    # }
    # db.likedb.insert_one(like_info)

#
# def like_reset():
#     like = list(db.myfantasystar.find({}, {'_id': False}).sort('name', ))
#     pre_like =
#     if weekno == 2:
#         like = '0'
#         pre_like = '0'
#
# #오잉?! like를 0으로 선언했기 때문에 like를 바로 넣으면 안될듯?
#
# def like_update():
#     if weekno != 2:
#         pre_like = 'like'
#         like = '0'


import pandas as pd
import requests
import zipfile
import os
import re
import datetime
import numpy as np

def unzip(arg):
    zf = zipfile.ZipFile(arg)
    return [zf.read(name) for name in zf.namelist()]

def extract_date(arg):
    arg = re.findall(string=str(arg), pattern='[0-9]{6}')[0] + '01'
    dt = pd.to_datetime(arg)

    i = 32
    next_dt = dt + datetime.timedelta(i)
    while next_dt.month != dt.month:
        next_dt = dt + datetime.timedelta(i)
        i -= 1
    return dt + datetime.timedelta(i + 1)

def get_line_for_portfolio_group(arg, text_to_find=''):
    for elem in range(0, len(arg)):
        if text_to_find in arg[elem]:
            return elem

    return elem

def industry_portfolio_csv(N_portfolios=49, cleanup=True, weighting='equal'):
    """

    :param N_portfolios:
    :param cleanup:
    :return:
    """

    if weighting=='value':
        text_to_find = 'Average Value Weighted Returns'
    elif weighting=='equal':
        text_to_find = 'Average Equal Weighted Returns'
    url = 'https://mba.tuck.dartmouth.edu/pages/faculty/ken.french/ftp/{N}_Industry_Portfolios_CSV.zip'.format(
        N=N_portfolios)

    resp = requests.get(url=url)

    fdir = os.path.expanduser('~') + '/temp/'
    fname = os.path.split(url)[1]
    fpath = fdir + fname

    f = open(fpath, 'wb')
    f.write(resp.content)
    f.close()

    uz = unzip(open(fpath, 'rb'))

    if cleanup:
        os.unlink(fpath)

    lines = str(uz[0]).split('\\r\\n')
    line_for_group = get_line_for_portfolio_group(lines,text_to_find)

    lines = lines[line_for_group+1:]

    lines = [l.split(',') for l in lines]

    df = pd.DataFrame(lines[1:], columns=lines[0])

    first_line_break = df[df.iloc[:, 0] == ''].index[0]

    df = df.iloc[0:first_line_break, :]

    df.index = df.iloc[:, 0].apply(extract_date)
    del df['']
    df = df.replace(' -99.99', np.nan)
    df = df.astype(float)

    return df / 100.

if __name__ == '__main__':
    industry_portfolio_csv()


import sys
from collections import defaultdict
from enum import Enum

class Color(Enum):
    WHITE = 0,
    GRAY = 1,
    BLACK = 2

def top_sort(edges, s, entry, result):
    entry[s] = Color.GRAY
    for v in edges[s]:
        if entry[v] == Color.WHITE:
            top_sort(edges, v, entry, result)
    entry[s] = Color.BLACK
    result.append(s)

if __name__ == '__main__':
    line = sys.stdin.readline().split()
    n = int(line[0])
    m = int(line[1])
    edges = defaultdict(list)
    for _ in range(m):
        line = sys.stdin.readline().split()
        u = int(line[0])
        v = int(line[1])
        edges[u].append(v)
    for key in edges.keys():
        edges[key].sort(reverse=True)
    entry = [Color.WHITE] * (n + 1)
    order = []
    for i in range(1, n + 1):
        if entry[i] is Color.WHITE:
            top_sort(edges, i, entry, order)
    order.reverse()
    print(" ".join(map(str, order)))