import argparse
import os
import sys
import math
import cv2
import numpy as np
import multiprocessing
from sklearn.metrics import confusion_matrix
import matplotlib.pyplot as plt
caffe_root = '/home/zhujiagang/temporal-segment-networks/lib/caffe-action/'
sys.path.insert(0, caffe_root + 'python')
from pyActionRecog import parse_directory
from pyActionRecog import parse_split_file
from pyActionRecog.utils.video_funcs import default_aggregation_func

dataset ='ucf101'
split =1
modality ='flow'
frame_path = '/home/zjg/zjg/tsncaffe/UCF-101-result/'

net_proto = "/home/zhujiagang/temporal-segment-networks/models/ucf101/tsn_bn_inception_flow_deploy.prototxt"
net_weights = "/home/zhujiagang/temporal-segment-networks/models/ucf101_split_1_tsn_flow_reference_bn_inception.caffemodel"
rgb_prefix ='img_'
flow_x_prefix ='flow_x_'
flow_y_prefix ='flow_y_'
num_frame_per_video = 1
save_scores = None
num_worker = 1
num_id = 1
gpus = [1]
from pyActionRecog.action_caffe import CaffeNet

# build neccessary information
print dataset
split_tp = parse_split_file(dataset)
f_info = parse_directory(frame_path, rgb_prefix, flow_x_prefix, flow_y_prefix)

gpu_list = None
eval_video_list = split_tp[split - 1][1]
score_name = 'prob'

def build_net():
    global net
    my_id = multiprocessing.current_process()._identity[0] \
        if num_worker > 1 else 1

    if gpu_list is None:
        net = CaffeNet(net_proto, net_weights, my_id-1)
    else:
        net = CaffeNet(net_proto, net_weights, gpu_list[my_id - 1])

def vis_square(data):
    """Take an array of shape (n, height, width) or (n, height, width, 3)
       and visualize each (height, width) thing in a grid of size approx. sqrt(n) by sqrt(n)"""

    # normalize data for display
    data = (data - data.min()) / (data.max() - data.min())

    # force the number of filters to be square
    n = int(np.ceil(np.sqrt(data.shape[0])))
    padding = (((0, n ** 2 - data.shape[0]),
                (0, 1), (0, 1))  # add some space between filters
               + ((0, 0),) * (data.ndim - 3))  # don't pad the last dimension (if there is one)
    data = np.pad(data, padding, mode='constant', constant_values=1)  # pad with ones (white)

    # tile the filters into an image
    data = data.reshape((n, n) + data.shape[1:]).transpose((0, 2, 1, 3) + tuple(range(4, data.ndim + 1)))
    data = data.reshape((n * data.shape[1], n * data.shape[3]) + data.shape[4:])

    plt.imshow(data);
    plt.axis('off')

def eval_video(video):
    global net
    label = video[1]
    vid = video[0]
    video_frame_path = f_info[0][vid]

    cnt_indexer = 2
    frame_cnt = f_info[cnt_indexer][vid]
    stack_depth = 5

    if num_frame_per_video >= 2:
        step = (frame_cnt - stack_depth) / (num_frame_per_video-1)
        if step > 0:
            frame_ticks = range(1, min((2 + step * (num_frame_per_video-1)), frame_cnt+1), step)
        else:
            frame_ticks = [1] * num_frame_per_video
    else:
        if frame_cnt > stack_depth:
            frame_ticks = [(frame_cnt - stack_depth) / 2]
        else:
            frame_ticks = [frame_cnt / 2]

    assert(len(frame_ticks) == num_frame_per_video)
    frame_scores = []

    for tick in frame_ticks:
        ii = 0
        if modality == 'rgb':
            name = '{}{:05d}.jpg'.format(rgb_prefix, tick)
            frame = cv2.imread(os.path.join(video_frame_path, name), cv2.IMREAD_COLOR)
            scores = net.predict_single_frame([frame, ], score_name, frame_size=(340, 256))
            frame_scores.append(scores)
        if modality == 'flow':
            frame_idx = [min(frame_cnt, tick + offset) for offset in xrange(stack_depth)]
            flow_stack = []
            for idx in frame_idx:
                ii = ii + 1
                x_name = '{}{:05d}.jpg'.format(flow_x_prefix, idx)
                y_name = '{}{:05d}.jpg'.format(flow_y_prefix, idx)
                flow_stack.append(cv2.imread(os.path.join(video_frame_path, x_name), cv2.IMREAD_GRAYSCALE))
                flow_stack.append(cv2.imread(os.path.join(video_frame_path, y_name), cv2.IMREAD_GRAYSCALE))

            scores, feat = net.predict_single_flow_stack_feature_map(flow_stack, score_name, frame_size=(224, 224))
            vis_square(feat)
            frame_scores.append(scores)

    print 'video {} done'.format(vid)
    sys.stdin.flush()
    return np.array(frame_scores), label

if num_worker > 1:
    pool = multiprocessing.Pool(num_worker, initializer=build_net)
    video_scores = pool.map(eval_video, eval_video_list)
else:
    build_net()
    video_scores = map(eval_video, eval_video_list)

video_pred = [np.argmax(default_aggregation_func(x[0])) for x in video_scores]
video_labels = [x[1] for x in video_scores]

cf = confusion_matrix(video_labels, video_pred).astype(float)

cls_cnt = cf.sum(axis=1)
cls_hit = np.diag(cf)

cls_acc = cls_hit/cls_cnt

print cls_acc

print 'Accuracy {:.02f}%'.format(np.mean(cls_acc)*100)

if save_scores is not None:
    np.savez(save_scores, scores=video_scores, labels=video_labels)


import codecs
import math
import numpy as np
import tensorflow as tf
import configparser
from bilm import Batcher, BidirectionalLanguageModel, weight_layers
import pyemblib
import configlogger
from drgriffis.common import log
from lib import wsd_parser
from lib.elmo_runner import ELMoParams, ELMoRunner, prepVocabulary

def getELMoRepresentations(sentences_words, sentences_instances, semcor_labels,
        unique_sense_IDs, bilm_params):

    sense_embeddings = {}
    for sense_ID in unique_sense_IDs:
        sense_embeddings[sense_ID] = []

    with tf.Session() as sess:
        log.writeln('  (1) Setting up ELMo')
        elmo = ELMoRunner(sess, bilm_params)

        # batch up the data
        sentence_ids = elmo.preprocess(sentences_words)

        batch_size = 25
        num_batches = math.ceil(sentence_ids.shape[0] / batch_size)
        batch_start = 0
        log.writeln('  (2) Extracting sense embeddings from sentences')
        log.track(message='    >> Processed {0}/{1:,} batches'.format('{0:,}',num_batches), writeInterval=5)
        while batch_start < sentence_ids.shape[0]:
            batch_sentence_ids = sentence_ids[batch_start:batch_start + batch_size]
            elmo_sentence_input_ = elmo(batch_sentence_ids)

            for i in range(elmo_sentence_input_.shape[0]):
                sentence_indices = sentences_instances[batch_start+i]
                for (instance_ID, ix) in sentence_indices:
                    senses = semcor_labels[instance_ID]
                    for sense in senses:
                        sense_embeddings[sense].append(
                            elmo_sentence_input_[i][ix]
                        )

            log.tick()
            batch_start += batch_size
    log.flushTracker()

    log.writeln('  (3) Calculating mean per-sense embeddings')
    mean_sense_embeddings = pyemblib.Embeddings()
    for (sense_ID, embedding_list) in sense_embeddings.items():
        if len(embedding_list) > 0:
            mean_sense_embeddings[sense_ID] = np.mean(embedding_list, axis=0)
        else:
            log.writeln('[WARNING] Sense ID "%s" found no embeddings' % sense_ID)

    return mean_sense_embeddings

if __name__ == '__main__':
    def _cli():
        import optparse
        parser = optparse.OptionParser(usage='Usage: %prog CONFIG')
        parser.add_option('-l', '--logfile', dest='logfile',
                help='name of file to write log contents to (empty for stdout)',
                default=None)
        (options, args) = parser.parse_args()
        if len(args) != 1:
            parser.print_help()
            exit()
        return args, options
    (configf,), options = _cli()
    config = configparser.ConfigParser()
    config.read(configf)

    log.start(logfile=options.logfile)
    configlogger.writeConfig(log, [
        ('SemCor', [
            ('XML', config['SemCor']['XML']),
            ('Labels', config['SemCor']['Labels']),
            ('Vocab', config['SemCor']['Vocab']),
        ]),
        ('ELMo', [
            ('Weights', config['ELMo']['Weights']),
            ('Options', config['ELMo']['Options']),
        ]),
        ('Output file', config['SemCor']['Embeddings']),
    ])

    t_sub = log.startTimer('Reading SemCor labels from %s...' % config['SemCor']['Labels'])
    semcor_labels, unique_sense_IDs = wsd_parser.readLabels(config['SemCor']['Labels'])
    log.stopTimer(t_sub, message='Read {0:,} labels ({1:,} unique senses) in {2}s.\n'.format(
        len(semcor_labels), len(unique_sense_IDs), '{0:.2f}'
    ))

    t_sub = log.startTimer('Pre-processing SemCor text from %s...' % config['SemCor']['XML'])
    (sentences_words, sentences_instances) = wsd_parser.processSentences(config['SemCor']['XML'])
    log.stopTimer(t_sub, message='Read {0:,} sentences in {1}s.\n'.format(
        len(sentences_words), '{0:.2f}'
    ))

    log.writeln('Pre-processing SemCor vocabulary...')
    max_char_len = prepVocabulary(sentences_words, config['SemCor']['Vocab'])
    log.writeln('Wrote vocabulary to {0}.\nMax character length: {1:,}\n'.format(
        config['SemCor']['Vocab'], max_char_len
    ))

    log.writeln('OVERRIDING max_char_len to 50!\n')
    max_char_len = 50

    bilm_params = ELMoParams()
    bilm_params.options_file = config['ELMo']['Options']
    bilm_params.weights_file = config['ELMo']['Weights']
    bilm_params.vocab_file = config['SemCor']['Vocab']
    bilm_params.max_char_len = max_char_len

    t_sub = log.startTimer('Getting ELMo representations for SemCor senses...')
    sense_embeddings = getELMoRepresentations(
        sentences_words,
        sentences_instances,
        semcor_labels,
        unique_sense_IDs,
        bilm_params
    )
    log.stopTimer(t_sub, message='Calculated embeddings for {0:,} senses in {1}s.\n'.format(
        len(sense_embeddings), '{0:.2f}'
    ))

    t_sub = log.startTimer('Writing sense embeddings to %s...' % config['SemCor']['Embeddings'])
    pyemblib.write(sense_embeddings, config['SemCor']['Embeddings'])
    log.stopTimer(t_sub, message='Completed writing embeddings in {0:.2f}s.\n')

    log.stop()


import matplotlib
matplotlib.use('Agg')
# import pickle
import copy
import argparse
import time
import datetime
import numpy as np
from io import StringIO
import itertools
import os
import sys

from ff.handlers.serialize import serialize_handlers
from ff.handlers.deserialize import deserialize

from multiprocessing import Process, Pipe

from jax.config import config as jax_config
# this always needs to be set
jax_config.update("jax_enable_x64", True)

import jax.numpy as jnp

import rdkit
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem import rdFMCS

from fe import dataset

from fe import loss, bar
from fe.pdb_writer import PDBWriter

import configparser
import grpc

from training import trainer
from training import service_pb2_grpc

from fe.utils import convert_uIC50_to_kJ_per_mole

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='Absolute Binding Free Energy Script')
    parser.add_argument('--config_file', type=str, required=True, help='Location of config file.')

    args = parser.parse_args()

    print("Launch Time:", datetime.datetime.now())

    config = configparser.ConfigParser()
    config.read(args.config_file)
    print("Config Settings:")
    config.write(sys.stdout)

    general_cfg = config['general']

    if not os.path.exists(general_cfg['out_dir']):
        os.makedirs(general_cfg['out_dir'])

    suppl = Chem.SDMolSupplier(general_cfg['ligand_sdf'], removeHs=False)

    all_guest_mols = []

    data = []

    for guest_idx, mol in enumerate(suppl):
        mol_dG = -1*convert_uIC50_to_kJ_per_mole(float(mol.GetProp(general_cfg['bind_prop'])))
        data.append((mol, mol_dG))

    full_dataset = dataset.Dataset(data)
    train_frac = float(general_cfg['train_frac'])
    train_dataset, test_dataset = full_dataset.split(train_frac)

    # process the host first
    host_pdbfile = general_cfg['protein_pdb']

    stage_dGs = []

    ff_raw = open(general_cfg['forcefield'], "r").read()
    ff_handlers = deserialize(ff_raw)

    worker_address_list = []
    for address in config['workers']['hosts'].split(','):
        worker_address_list.append(address)

    stubs = []

    for address in worker_address_list:
        print("connecting to", address)
        channel = grpc.insecure_channel(address,
            options = [
                ('grpc.max_send_message_length', 500 * 1024 * 1024),
                ('grpc.max_receive_message_length', 500 * 1024 * 1024)
            ]
        )

        stub = service_pb2_grpc.WorkerStub(channel)
        stubs.append(stub)

    intg_cfg = config['integrator']
    lr_config = config['learning_rates']
    restr_config = config['restraints']

    lambda_schedule = {}

    for stage_str, v in config['lambda_schedule'].items():

        stage = int(stage_str)
        stage_schedule = np.array([float(x) for x in v.split(',')])

        assert stage not in lambda_schedule

        if stage == 0 or stage == 1:
            # stage 0 must be monotonically decreasing
            assert np.all(np.diff(stage_schedule) > 0)
        else:
            raise Exception("unknown stage")
            # stage 1 and 2 must be monotonically increasing
            # assert np.all(np.diff(stage_schedule) > 0)
        lambda_schedule[stage] = stage_schedule

    learning_rates = {}
    for k, v in config['learning_rates'].items():
        learning_rates[k] = np.array([float(x) for x in v.split(',')])

    engine = trainer.Trainer(
        host_pdbfile,
        stubs,
        worker_address_list,
        ff_handlers,
        lambda_schedule,
        int(general_cfg['du_dl_cutoff']),
        float(restr_config['search_radius']),
        float(restr_config['force_constant']),
        int(general_cfg['n_frames']),
        int(intg_cfg['steps']),
        float(intg_cfg['dt']),
        float(intg_cfg['temperature']),
        float(intg_cfg['friction']),
        learning_rates,
        general_cfg['precision']
    )

    for epoch in range(100):

        print("Starting Epoch", epoch, datetime.datetime.now().strftime("%d/%m/%Y %H:%M:%S"))
        epoch_dir = os.path.join(general_cfg['out_dir'], "epoch_"+str(epoch))

        if not os.path.exists(epoch_dir):
            os.makedirs(epoch_dir)

        epoch_params = serialize_handlers(ff_handlers)
        with open(os.path.join(epoch_dir, "start_epoch_params.py"), 'w') as fh:
            fh.write(epoch_params)

        for mol, experiment_dG in test_dataset.data:
            print("test mol", mol.GetProp("_Name"), "Smiles:", Chem.MolToSmiles(mol))
            mol_dir = os.path.join(epoch_dir, "test_mol_"+mol.GetProp("_Name"))
            start_time = time.time()
            dG, ci, loss = engine.run_mol(mol, inference=True, run_dir=mol_dir, experiment_dG=experiment_dG)
            print(mol.GetProp("_Name"), "test loss", loss, "pred_dG", dG, "exp_dG", experiment_dG, "time", time.time() - start_time, "ci 95% (mean, lower, upper)", ci.value, ci.lower_bound, ci.upper_bound)

        train_dataset.shuffle()

        for mol, experiment_dG in train_dataset.data:
            print("train mol", mol.GetProp("_Name"), "Smiles:", Chem.MolToSmiles(mol))
            mol_dir = os.path.join(epoch_dir, "train_mol_"+mol.GetProp("_Name"))
            start_time = time.time()
            dG, ci, loss = engine.run_mol(mol, inference=False, run_dir=mol_dir, experiment_dG=experiment_dG)
            print(mol.GetProp("_Name"), "train loss", loss, "pred_dG", dG, "exp_dG", experiment_dG, "time", time.time() - start_time, "ci 95% (mean, lower, upper)", ci.value, ci.lower_bound, ci.upper_bound)

        epoch_params = serialize_handlers(ff_handlers)
        with open(os.path.join(epoch_dir, "end_epoch_params.py"), 'w') as fh:
            fh.write(epoch_params)


import os, sys
import time
import numpy as np

# Classes and Functions for preparing and continuing expanded-ensemble simulations

class expanded_ensemble_mdpfile(object):

    def __init__(self,  couple_moltype = '1MQ',
			pull_group1_name = 'a1-Protein',
                        pull_group2_name = 'a2-Ligand',
                        pull_coord1_k    = 200.0,
                        pull_coord1_init  = 0.4,
                        fep_lambdas      = np.arange(0.00, 1.05, 0.05),
                        init_lambda_weights = np.zeros(21),
                        init_lambda_state   = 20,
                        wl_increment_in_kT  = 3.0 ):

        """Initialize the class."""

        self.couple_moltype = couple_moltype
        self.pull_group1_name = pull_group1_name
        self.pull_group2_name = pull_group2_name
        self.pull_coord1_k    = pull_coord1_k     # the spring constant in kJ/nm^2
        self.pull_coord1_init = pull_coord1_init  # the equilibrium distance between the pull group 1 and 2
        self.fep_lambdas      = fep_lambdas
        self.nlambdas         = len(self.fep_lambdas)
        self.fep_lambdas_string   = ' '.join(['%2.3f'%fep_lambdas[i] for i in range(self.nlambdas)])

        self.init_lambda_weights  = init_lambda_weights
        self.init_lambda_weights_string =  ' '.join(['%2.3f'%init_lambda_weights[i] for i in range(self.nlambdas)])

        self.init_lambda_state    = init_lambda_state
        self.wl_increment_in_kT   = wl_increment_in_kT

    def mdpfile_text(self):
        """Returns a string corresponding to the mdpfile contents."""

        self.header_desc = ";       generated from expanded_ensemble_mdpfile() on %s \n;\n;\n"%time.asctime()

        self.mdp_text = """; VARIOUS PREPROCESSING OPTIONS
; Preprocessor information: use cpp syntax.
; e.g.: -I/home/joe/doe -I/home/mary/roe
include                  = -I../top
; e.g.: -DPOSRES -DFLEXIBLE (note these variable names are case sensitive)
define                   =

; RUN CONTROL PARAMETERS
integrator               = md-vv   ;  only md-vv works with ee;  sd
; Start time and timestep in ps
tinit                    = 0
dt                       = 0.002
nsteps                   = 500000   ; 1000 ps = 1 ns
; For exact run continuation or redoing part of a run
init_step                = 0
; Part index is updated automatically on checkpointing (keeps files separate)
simulation_part          = 1
; mode for center of mass motion removal
comm-mode                = Linear
; number of steps for center of mass motion removal
nstcomm                  = 10
; group(s) for center of mass motion removal
comm_grps                = System

; LANGEVIN DYNAMICS OPTIONS
; Friction coefficient (amu/ps) and random seed
bd-fric                  = 0
ld-seed                  = -1

; TEST PARTICLE INSERTION OPTIONS
rtpi                     = 0.05

; OUTPUT CONTROL OPTIONS
; Output frequency for coords (x), velocities (v) and forces (f)
nstxout                  = 500000  ; 1 ns
nstvout                  = 500000  ; 1 ns
nstfout                  = 0
; Output frequency for energies to log file and energy file
nstlog                   = 500     ; 1 ps
nstcalcenergy            = 500     ; 1 ps
nstenergy                = 500     ; 1 ps
; Output frequency and precision for .xtc file
nstxtcout                = 500     ; every 1 ps
xtc-precision            = 1000
; This selects the subset of atoms for the .xtc file. You can
; select multiple groups. By default all atoms will be written.
xtc_grps                 = 1MQ
; Selection of energy groups
energygrps               = WaterProtein Water_and_ions 1MQ

; NEIGHBORSEARCHING PARAMETERS
; nblist update frequency
nstlist                  = 10
; ns algorithm (simple or grid)
ns_type                  = grid
; Periodic boundary conditions: xyz, no, xy
pbc                      = xyz
periodic_molecules       = no
; nblist cut-off
rlist                    = 0.9
; long-range cut-off for switched potentials
rlistlong                = -1

; OPTIONS FOR ELECTROSTATICS AND VDW
; Method for doing electrostatics
coulombtype              = PME
rcoulomb-switch          = 0
rcoulomb                 = 0.9
; Relative dielectric constant for the medium and the reaction field
epsilon_r                = 1
epsilon_rf               = 1
; Method for doing Van der Waals
vdw-type                 = Cut-off
; cut-off lengths
rvdw-switch              = 0
rvdw                     = 0.9
; Apply long range dispersion corrections for Energy and Pressure
DispCorr                 = No
; Extension of the potential lookup tables beyond the cut-off
table-extension          = 1
; Seperate tables between energy group pairs
energygrp_table          =
; Spacing for the PME/PPPM FFT grid
fourierspacing           = 0.12
; FFT grid size, when a value is 0 fourierspacing will be used
fourier_nx               = 0
fourier_ny               = 0
fourier_nz               = 0
; EWALD/PME/PPPM parameters
pme_order                = 4
ewald_rtol               = 1e-5
ewald_geometry           = 3d
epsilon_surface          = 0
optimize_fft             = no

; OPTIONS FOR WEAK COUPLING ALGORITHMS
; Temperature coupling
tcoupl                   = Berendsen
nsttcouple               = -1
nh-chain-length          = 10
; Groups to couple separately
tc-grps                  = Protein non-Protein
; Time constant (ps) and reference temperature (K)
tau_t                    = 1.0  1.0
ref_t                    = 300  300
; Pressure coupling
Pcoupl                   = no
Pcoupltype               = Isotropic
nstpcouple               = -1
; Time constant (ps), compressibility (1/bar) and reference P (bar)
tau_p                    = 1.0
compressibility          = 4.5e-5
ref_p                    = 1.0
; Scaling of reference coordinates, No, All or COM
refcoord_scaling         = No
; Random seed for Andersen thermostat
andersen_seed            = 815131

; GENERATE VELOCITIES FOR STARTUP RUN
gen_vel                  = yes
gen_temp                 = 300
;gen_seed		 = -1 <<<--- FAHWorkServer complains that there are multiple lines; it adds its own line!

; OPTIONS FOR BONDS
constraints              = hbonds
; Type of constraint algorithm
constraint-algorithm     = Lincs
; Do not constrain the start configuration
continuation             = no
; Use successive overrelaxation to reduce the number of shake iterations
Shake-SOR                = no
; Relative tolerance of shake
shake-tol                = 0.0001
; Highest order in the expansion of the constraint coupling matrix
lincs-order              = 4
; Number of iterations in the final step of LINCS. 1 is fine for
; normal simulations, but use 2 to conserve energy in NVE runs.
; For energy minimization with constraints it should be 4 to 8.
lincs-iter               = 1
; Lincs will write a warning to the stderr if in one step a bond
; rotates over more degrees than
lincs-warnangle          = 30
; Convert harmonic bonds to morse potentials
morse                    = no

; ENERGY GROUP EXCLUSIONS
; Pairs of energy groups for which all non-bonded interactions are excluded
energygrp_excl           =

; WALLS
; Number of walls, type, atom types, densities and box-z scale factor for Ewald
nwall                    = 0
wall_type                = 9-3
wall_r_linpot            = -1
wall_atomtype            =
wall_density             =
wall_ewald_zfac          = 3

; Free energy control stuff ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
free-energy              = expanded
init-lambda-state        = {init_lambda_state}
;  delta-lambda             = 0
fep-lambdas              = {fep_lambdas_string}
calc-lambda-neighbors    = -1
sc-alpha                 = 0
sc-power                 = 0
sc-sigma                 = 0.3
nstdhdl                  = 500
separate-dhdl-file       = yes
dhdl-derivatives         = yes
;   dh_hist_size             = 0
;   dh_hist_spacing          = 0.1
couple-moltype           = {couple_moltype}
couple-lambda0           = none ; off
couple-lambda1           = vdw-q ; on
couple-intramol          = no

; expanded ensemble variables
nstexpanded              = 500     ; 1 ps
lmc-stats                = wang-landau
lmc-move                 = metropolized-gibbs
lmc-weights-equil        = no
weight-equil-number-all-lambda = -1
weight-equil-number-samples = -1
weight-equil-number-steps = -1
weight-equil-wl-delta    = -1
weight-equil-count-ratio = -1

; Seed for Monte Carlo in lambda space
lmc-seed                 = -1
mc-temperature           = -1
lmc-repeats              = 1
lmc-gibbsdelta           = -1
lmc-forced-nstart        = 0
symmetrized-transition-matrix = no
nst-transition-matrix    = -1
mininum-var-min          = 100
init-lambda-weights      = {init_lambda_weights}
weight-c-range           = 0
wl-scale                 = 0.8
wl-ratio                 = 0.8
init-wl-delta            = {wl_increment_in_kT}  ; kT -- default is 1.0
wl-oneovert              = no

; pulling parameters
pull                     = umbrella
pull_ngroups             = 2
pull_ncoords             = 1
pull_group1_name         = {pull_group1_name}
pull_group2_name         = {pull_group2_name}
pull-geometry            = direction-periodic
pull_coord1_groups       = 1 2
pull-dim                 = Y Y Y
pull_coord1_rate         = 0.00
pull_coord1_k            = {pull_coord1_k}
pull-start               = no
pull-coord1-init         = {pull_coord1_init}
pull-nstxout             = 500   ; 1 ps
pull-nstfout             = 500   ; 1 ps
""".format(couple_moltype      = self.couple_moltype,
           pull_group1_name    = self.pull_group1_name,
           pull_group2_name    = self.pull_group2_name,
           pull_coord1_k       = self.pull_coord1_k,
           pull_coord1_init    = self.pull_coord1_init,
           fep_lambdas_string  = self.fep_lambdas_string,
           init_lambda_state   = self.init_lambda_state,
           init_lambda_weights = self.init_lambda_weights_string,
           wl_increment_in_kT  = self.wl_increment_in_kT )

        return self.header_desc + self.mdp_text

    def write_to_filename(self, filename):
        """Writes the ee mdpfile to specified filename."""

        fout = open(filename, 'w')
        fout.write( self.mdpfile_text() )
        fout.close()

### A derived class for an mdpfile wth NO RESTRAINT

class expanded_ensemble_mdpfile_NOREST(expanded_ensemble_mdpfile):

    def __init__(self,  couple_moltype = '1MQ',
                        fep_lambdas      = np.arange(0.00, 1.05, 0.05),
                        init_lambda_weights = np.zeros(21),
                        init_lambda_state   = 20,
                        wl_increment_in_kT  = 3.0 ):

        """Initialize the class."""

        self.couple_moltype = couple_moltype
        self.fep_lambdas      = fep_lambdas
        self.nlambdas         = len(self.fep_lambdas)
        self.fep_lambdas_string   = ' '.join(['%2.3f'%fep_lambdas[i] for i in range(self.nlambdas)])

        self.init_lambda_weights  = init_lambda_weights
        self.init_lambda_weights_string =  ' '.join(['%2.3f'%init_lambda_weights[i] for i in range(self.nlambdas)])

        self.init_lambda_state    = init_lambda_state
        self.wl_increment_in_kT   = wl_increment_in_kT

    def mdpfile_text(self):
        """Returns a string corresponding to the mdpfile contents.

        NOTE -- these will be used for ligand-only simulations in small boxes of water, so total time is 10 ns"""

        self.header_desc = ";       generated from expanded_ensemble_mdpfile() on %s \n;\n;\n"%time.asctime()

        self.mdp_text = """; VARIOUS PREPROCESSING OPTIONS
; Preprocessor information: use cpp syntax.
; e.g.: -I/home/joe/doe -I/home/mary/roe
include                  = -I../top
; e.g.: -DPOSRES -DFLEXIBLE (note these variable names are case sensitive)
define                   =

; RUN CONTROL PARAMETERS
integrator               = md-vv   ;  only md-vv works with ee;  sd
; Start time and timestep in ps
tinit                    = 0
dt                       = 0.002
nsteps                   = 5000000   ; 10 000 ps = 10 ns
; For exact run continuation or redoing part of a run
init_step                = 0
; Part index is updated automatically on checkpointing (keeps files separate)
simulation_part          = 1
; mode for center of mass motion removal
comm-mode                = Linear
; number of steps for center of mass motion removal
nstcomm                  = 10
; group(s) for center of mass motion removal
comm_grps                = System

; LANGEVIN DYNAMICS OPTIONS
; Friction coefficient (amu/ps) and random seed
bd-fric                  = 0
ld-seed                  = -1

; TEST PARTICLE INSERTION OPTIONS
rtpi                     = 0.05

; OUTPUT CONTROL OPTIONS
; Output frequency for coords (x), velocities (v) and forces (f)
nstxout                  = 500000  ; 1 ns
nstvout                  = 500000  ; 1 ns
nstfout                  = 0
; Output frequency for energies to log file and energy file
nstlog                   = 5000     ; 10 ps
nstcalcenergy            = 5000     ; 10 ps
nstenergy                = 5000     ; 10 ps
; Output frequency and precision for .xtc file
nstxtcout                = 50000  ; every 100 ps
xtc-precision            = 1000
; This selects the subset of atoms for the .xtc file. You can
; select multiple groups. By default all atoms will be written.
xtc_grps                 = 1MQ
; Selection of energy groups
energygrps               = Water non-Water

; NEIGHBORSEARCHING PARAMETERS
; nblist update frequency
nstlist                  = 10
; ns algorithm (simple or grid)
ns_type                  = grid
; Periodic boundary conditions: xyz, no, xy
pbc                      = xyz
periodic_molecules       = no
; nblist cut-off
rlist                    = 0.9
; long-range cut-off for switched potentials
rlistlong                = -1

; OPTIONS FOR ELECTROSTATICS AND VDW
; Method for doing electrostatics
coulombtype              = PME
rcoulomb-switch          = 0
rcoulomb                 = 0.9
; Relative dielectric constant for the medium and the reaction field
epsilon_r                = 1
epsilon_rf               = 1
; Method for doing Van der Waals
vdw-type                 = Cut-off
; cut-off lengths
rvdw-switch              = 0
rvdw                     = 0.9
; Apply long range dispersion corrections for Energy and Pressure
DispCorr                 = No
; Extension of the potential lookup tables beyond the cut-off
table-extension          = 1
; Seperate tables between energy group pairs
energygrp_table          =
; Spacing for the PME/PPPM FFT grid
fourierspacing           = 0.12
; FFT grid size, when a value is 0 fourierspacing will be used
fourier_nx               = 0
fourier_ny               = 0
fourier_nz               = 0
; EWALD/PME/PPPM parameters
pme_order                = 4
ewald_rtol               = 1e-5
ewald_geometry           = 3d
epsilon_surface          = 0
optimize_fft             = no

; OPTIONS FOR WEAK COUPLING ALGORITHMS
; Temperature coupling
tcoupl                   = Berendsen
nsttcouple               = -1
nh-chain-length          = 10
; Groups to couple separately
tc-grps                  = Water   non-Water
; Time constant (ps) and reference temperature (K)
tau_t                    = 1.0  1.0
ref_t                    = 300  300
; Pressure coupling
Pcoupl                   = no
Pcoupltype               = Isotropic
nstpcouple               = -1
; Time constant (ps), compressibility (1/bar) and reference P (bar)
tau_p                    = 1.0
compressibility          = 4.5e-5
ref_p                    = 1.0
; Scaling of reference coordinates, No, All or COM
refcoord_scaling         = No
; Random seed for Andersen thermostat
andersen_seed            = 815131

; GENERATE VELOCITIES FOR STARTUP RUN
gen_vel                  = yes
gen_temp                 = 300
;gen_seed		 = -1 <<<--- FAHWorkServer complains that there are multiple lines; it adds its own line!

; OPTIONS FOR BONDS
constraints              = hbonds
; Type of constraint algorithm
constraint-algorithm     = Lincs
; Do not constrain the start configuration
continuation             = no
; Use successive overrelaxation to reduce the number of shake iterations
Shake-SOR                = no
; Relative tolerance of shake
shake-tol                = 0.0001
; Highest order in the expansion of the constraint coupling matrix
lincs-order              = 4
; Number of iterations in the final step of LINCS. 1 is fine for
; normal simulations, but use 2 to conserve energy in NVE runs.
; For energy minimization with constraints it should be 4 to 8.
lincs-iter               = 1
; Lincs will write a warning to the stderr if in one step a bond
; rotates over more degrees than
lincs-warnangle          = 30
; Convert harmonic bonds to morse potentials
morse                    = no

; ENERGY GROUP EXCLUSIONS
; Pairs of energy groups for which all non-bonded interactions are excluded
energygrp_excl           =

; WALLS
; Number of walls, type, atom types, densities and box-z scale factor for Ewald
nwall                    = 0
wall_type                = 9-3
wall_r_linpot            = -1
wall_atomtype            =
wall_density             =
wall_ewald_zfac          = 3

; Free energy control stuff ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
free-energy              = expanded
init-lambda-state        = {init_lambda_state}
;  delta-lambda             = 0
fep-lambdas              = {fep_lambdas_string}
calc-lambda-neighbors    = -1
sc-alpha                 = 0
sc-power                 = 0
sc-sigma                 = 0.3
nstdhdl                  = 5000
separate-dhdl-file       = yes
dhdl-derivatives         = yes
;   dh_hist_size             = 0
;   dh_hist_spacing          = 0.1
couple-moltype           = {couple_moltype}
couple-lambda0           = none ; off
couple-lambda1           = vdw-q ; on
couple-intramol          = no

; expanded ensemble variables
nstexpanded              = 5000     ; 10 ps
lmc-stats                = wang-landau
lmc-move                 = metropolized-gibbs
lmc-weights-equil        = no
weight-equil-number-all-lambda = -1
weight-equil-number-samples = -1
weight-equil-number-steps = -1
weight-equil-wl-delta    = -1
weight-equil-count-ratio = -1

; Seed for Monte Carlo in lambda space
lmc-seed                 = -1
mc-temperature           = -1
lmc-repeats              = 1
lmc-gibbsdelta           = -1
lmc-forced-nstart        = 0
symmetrized-transition-matrix = no
nst-transition-matrix    = -1
mininum-var-min          = 100
init-lambda-weights      = {init_lambda_weights}
weight-c-range           = 0
wl-scale                 = 0.8
wl-ratio                 = 0.8
init-wl-delta            = {wl_increment_in_kT}  ; kT -- default is 1.0
wl-oneovert              = no

; pulling parameters
pull                     = no
""".format(couple_moltype      = self.couple_moltype,
           fep_lambdas_string  = self.fep_lambdas_string,
           init_lambda_state   = self.init_lambda_state,
           init_lambda_weights = self.init_lambda_weights_string,
           wl_increment_in_kT  = self.wl_increment_in_kT )

        return self.header_desc + self.mdp_text

    def write_to_filename(self, filename):
        """Writes the ee mdpfile to specified filename."""

        fout = open(filename, 'w')
        fout.write( self.mdpfile_text() )
        fout.close()

if __name__ == "__main__":
    e = expanded_ensemble_mdpfile()
    print(e.mdpfile_text())
    e.write_to_filename('test.mdp')

    e2 = expanded_ensemble_mdpfile_NOREST()
    print(e2.mdpfile_text())
    e2.write_to_filename('test-NOREST.mdp')


import numpy as np

class Py2C02:

    def __init__(self):
        # Variables for the screen and sprites
        self.palScreen = [0]*0x40
        self.sprScreen = np.zeros((256, 240, 3), dtype=np.uint8)
        self.sprNameTable = [np.zeros((256, 240, 3), dtype=np.uint8),np.zeros((256, 240, 3), dtype=np.uint8)]
        self.sprPatternTable = [np.zeros((128,128, 3), dtype=np.uint8),np.zeros((128,128, 3), dtype=np.uint8)]

        self.tblName = np.zeros((2,1024), dtype=np.uint8)
        self.tblPattern = np.zeros((2,1024), dtype=np.uint8)
        self.tblPalette = np.zeros(32, dtype=np.uint8)

        #Variables for the scanning of the screen
        self.scanline = 0
        self.cycle = 0
        self.frame_complete = False

        #Variable for holding the cartridge
        self.cart = None

        # Makes palScreen
        self.setpalScreen()

    """
    Function for filling the array holding the different pixel values
    """
    def setpalScreen(self):
        self.palScreen[0x00] = np.array([84, 84, 84])
        self.palScreen[0x01] = np.array([0, 30, 116])
        self.palScreen[0x02] = np.array([8, 16, 144])
        self.palScreen[0x03] = np.array([48, 0, 136])
        self.palScreen[0x04] = np.array([68, 0, 100])
        self.palScreen[0x05] = np.array([92, 0, 48])
        self.palScreen[0x06] = np.array([84, 4, 0])
        self.palScreen[0x07] = np.array([60, 24, 0])
        self.palScreen[0x08] = np.array([32, 42, 0])
        self.palScreen[0x09] = np.array([8, 58, 0])
        self.palScreen[0x0A] = np.array([0, 64, 0])
        self.palScreen[0x0B] = np.array([0, 60, 0])
        self.palScreen[0x0C] = np.array([0, 50, 60])
        self.palScreen[0x0D] = np.array([0, 0, 0])
        self.palScreen[0x0E] = np.array([0, 0, 0])
        self.palScreen[0x0F] = np.array([0, 0, 0])

        self.palScreen[0x10] = np.array([152, 150, 152])
        self.palScreen[0x11] = np.array([8, 76, 196])
        self.palScreen[0x12] = np.array([48, 50, 236])
        self.palScreen[0x13] = np.array([92, 30, 228])
        self.palScreen[0x14] = np.array([136, 20, 176])
        self.palScreen[0x15] = np.array([160, 20, 100])
        self.palScreen[0x16] = np.array([152, 34, 32])
        self.palScreen[0x17] = np.array([120, 60, 0])
        self.palScreen[0x18] = np.array([84, 90, 0])
        self.palScreen[0x19] = np.array([40, 114, 0])
        self.palScreen[0x1A] = np.array([8, 124, 0])
        self.palScreen[0x1B] = np.array([0, 118, 40])
        self.palScreen[0x1C] = np.array([0, 102, 120])
        self.palScreen[0x1D] = np.array([0, 0, 0])
        self.palScreen[0x1E] = np.array([0, 0, 0])
        self.palScreen[0x1F] = np.array([0, 0, 0])

        self.palScreen[0x20] = np.array([236, 238, 236])
        self.palScreen[0x21] = np.array([76, 154, 236])
        self.palScreen[0x22] = np.array([120, 124, 236])
        self.palScreen[0x23] = np.array([176, 98, 236])
        self.palScreen[0x24] = np.array([228, 84, 236])
        self.palScreen[0x25] = np.array([236, 88, 180])
        self.palScreen[0x26] = np.array([236, 106, 100])
        self.palScreen[0x27] = np.array([212, 136, 32])
        self.palScreen[0x28] = np.array([160, 170, 0])
        self.palScreen[0x29] = np.array([116, 196, 0])
        self.palScreen[0x2A] = np.array([76, 208, 32])
        self.palScreen[0x2B] = np.array([56, 204, 108])
        self.palScreen[0x2C] = np.array([56, 180, 204])
        self.palScreen[0x2D] = np.array([60, 60, 60])
        self.palScreen[0x2E] = np.array([0, 0, 0])
        self.palScreen[0x2F] = np.array([0, 0, 0])

        self.palScreen[0x30] = np.array([236, 238, 236])
        self.palScreen[0x31] = np.array([168, 204, 236])
        self.palScreen[0x32] = np.array([188, 188, 236])
        self.palScreen[0x33] = np.array([212, 178, 236])
        self.palScreen[0x34] = np.array([236, 174, 236])
        self.palScreen[0x35] = np.array([236, 174, 212])
        self.palScreen[0x36] = np.array([236, 180, 176])
        self.palScreen[0x37] = np.array([228, 196, 144])
        self.palScreen[0x38] = np.array([204, 210, 120])
        self.palScreen[0x39] = np.array([180, 222, 120])
        self.palScreen[0x3A] = np.array([168, 226, 144])
        self.palScreen[0x3B] = np.array([152, 226, 180])
        self.palScreen[0x3C] = np.array([160, 214, 228])
        self.palScreen[0x3D] = np.array([160, 162, 160])
        self.palScreen[0x3E] = np.array([0, 0, 0])
        self.palScreen[0x3F] = np.array([0, 0, 0])

    """
    Functions for returning screen properties
    """

    def GetScreen(self) -> np.zeros((256, 240, 3), dtype=np.uint8):
        return self.sprScreen

    def GetNameTable(self, i: np.uint8) -> np.zeros((256, 240, 3), dtype=np.uint8):
        return self.sprNameTable[i]

    def GetPatternTable(self, i: np.uint8) -> np.zeros((256, 240, 3), dtype=np.uint8):
        return self.sprPatternTable[i]

    """
    Functions for reading and writing
    """

    def cpuRead(self, addr: np.uint16) -> np.uint8:
        data = 0x00
        if addr == 0x0000: #Control
            pass
        elif addr == 0x0001: # Mask
            pass
        elif addr == 0x0002: # Status
            pass
        elif addr == 0x0003: # OAM Address
            pass
        elif addr == 0x0004: # OAM Data
            pass
        elif addr == 0x0005: # Scroll
            pass
        elif addr == 0x0006: # PPU Address
            pass
        elif addr == 0x0007: # PPU Data
            pass

        return data

    def cpuWrite(self, addr: np.uint16, data: np.uint8):

        if addr == 0x0000: #Control
            pass
        elif addr == 0x0001: # Mask
            pass
        elif addr == 0x0002: # Status
            pass
        elif addr == 0x0003: # OAM Address
            pass
        elif addr == 0x0004: # OAM Data
            pass
        elif addr == 0x0005: # Scroll
            pass
        elif addr == 0x0006: # PPU Address
            pass
        elif addr == 0x0007: # PPU Data
            pass

    def ppuRead(self, addr: np.uint16) -> np.uint8:
        data = 0x00

        addr &= 0x3FFF

        if self.cart.ppuRead(addr, data):
            pass

        return data

    def ppuWrite(self, addr: np.uint16, data: np.uint8):
        addr &= 0x3FFF

        if self.cart.ppuWrite(addr, data):
            pass

    """
    Function for connecting the cartridge
    """
    def connectCartridge(self, cart):
        self.cart = cart

    """
    Clock Function
    """
    def clock(self):
        try:
            self.sprScreen[self.cycle-1, self.scanline, :] = self.palScreen[np.random.choice([0x3F,0x30])]
        except:
            pass

        self.cycle += 1

        if self.cycle >= 341:

            self.cycle = 0

            self.scanline += 1

            if self.scanline >= 261:
                self.scanline = -1
                self.frame_complete = True


# -*- coding: utf-8 -*-

import os
import time
import shutil
import numpy as np
import tensorflow as tf
import yolov4.utils as utils
from tqdm import tqdm
from yolov4.dataset import Dataset
from yolov4.model import YOLOV4
from yolov4.config import cfg

os.environ['CUDA_VISIBLE_DEVICES'] = '0'

class YoloTrain(object):
    def __init__(self):
        self.anchor_per_scale = cfg.YOLO.ANCHOR_PER_SCALE
        self.classes = utils.read_class_names(cfg.YOLO.CLASSES)
        self.num_classes = len(self.classes)
        self.learn_rate_init = cfg.TRAIN.LEARN_RATE_INIT
        self.learn_rate_end = cfg.TRAIN.LEARN_RATE_END
        self.first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS
        self.second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS
        self.warmup_periods = cfg.TRAIN.WARMUP_EPOCHS
        self.initial_weight = cfg.TRAIN.INITIAL_WEIGHT
        self.time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))
        self.moving_ave_decay = cfg.YOLO.MOVING_AVE_DECAY
        self.max_bbox_per_scale = 150
        self.train_logdir = "./data/log/train"
        self.trainset = Dataset('train')
        self.valset = Dataset('test')
        self.steps_per_period = len(self.trainset)
        self.sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))

        with tf.name_scope('input'):
            self.input_data = tf.placeholder(dtype=tf.float32, name='input_data')
            self.label_sbbox = tf.placeholder(dtype=tf.float32, name='label_sbbox')
            self.label_mbbox = tf.placeholder(dtype=tf.float32, name='label_mbbox')
            self.label_lbbox = tf.placeholder(dtype=tf.float32, name='label_lbbox')
            self.true_sbboxes = tf.placeholder(dtype=tf.float32, name='sbboxes')
            self.true_mbboxes = tf.placeholder(dtype=tf.float32, name='mbboxes')
            self.true_lbboxes = tf.placeholder(dtype=tf.float32, name='lbboxes')
            self.trainable = tf.placeholder(dtype=tf.bool, name='training')

        with tf.name_scope("define_loss"):
            iou_use = 1  # (0, 1, 2) ==> (giou_loss, diou_loss, ciou_loss)
            focal_use = False  # (False, True) ==> (normal, focal_loss)
            label_smoothing = 0  # 0.001
            self.model = YOLOV4(self.input_data, self.trainable)
            self.net_var = tf.global_variables()
            self.iou_loss, self.conf_loss, self.prob_loss = self.model.compute_loss(
                                                self.label_sbbox,  self.label_mbbox,  self.label_lbbox,
                                                self.true_sbboxes, self.true_mbboxes, self.true_lbboxes,
                                                iou_use, focal_use, label_smoothing)
            self.loss = self.iou_loss + self.conf_loss + self.prob_loss
            # self.loss = tf.Print(self.loss, [self.iou_loss, self.conf_loss, self.prob_loss], message='loss: ')

        with tf.name_scope('learn_rate'):
            self.global_step = tf.Variable(1.0, dtype=tf.float64, trainable=False, name='global_step')
            warmup_steps = tf.constant(self.warmup_periods * self.steps_per_period,
                                       dtype=tf.float64, name='warmup_steps')
            train_steps = tf.constant((self.first_stage_epochs + self.second_stage_epochs) * self.steps_per_period,
                                      dtype=tf.float64, name='train_steps')
            self.learn_rate = tf.cond(
                pred=self.global_step < warmup_steps,
                true_fn=lambda: self.global_step / warmup_steps * self.learn_rate_init,
                false_fn=lambda: self.learn_rate_end + 0.5 * (self.learn_rate_init - self.learn_rate_end) *
                                    (1 + tf.cos(
                                        (self.global_step - warmup_steps) / (train_steps - warmup_steps) * np.pi))
            )
            global_step_update = tf.assign_add(self.global_step, 1.0)

        with tf.name_scope("define_weight_decay"):
            moving_ave = tf.train.ExponentialMovingAverage(self.moving_ave_decay).apply(tf.trainable_variables())

        with tf.name_scope("define_first_stage_train"):
            self.first_stage_trainable_var_list = []
            for var in tf.trainable_variables():
                var_name = var.op.name
                var_name_mess = str(var_name).split('/')
                if var_name_mess[0] in ['conv_sbbox', 'conv_mbbox', 'conv_lbbox']:
                    self.first_stage_trainable_var_list.append(var)

            first_stage_optimizer = tf.train.AdamOptimizer(self.learn_rate).minimize(self.loss,
                                                    var_list=self.first_stage_trainable_var_list)
            with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
                with tf.control_dependencies([first_stage_optimizer, global_step_update]):
                    with tf.control_dependencies([moving_ave]):
                        self.train_op_with_frozen_variables = tf.no_op()

        with tf.name_scope("define_second_stage_train"):
            second_stage_trainable_var_list = tf.trainable_variables()
            second_stage_optimizer = tf.train.AdamOptimizer(self.learn_rate).minimize(self.loss,
                                                    var_list=second_stage_trainable_var_list)

            with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
                with tf.control_dependencies([second_stage_optimizer, global_step_update]):
                    with tf.control_dependencies([moving_ave]):
                        self.train_op_with_all_variables = tf.no_op()

        with tf.name_scope('loader_and_saver'):
            self.loader = tf.train.Saver(self.net_var)
            self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=10)
            # self.saver  = tf.train.Saver(tf.trainable_variables(), max_to_keep=10)

        with tf.name_scope('summary'):
            tf.summary.scalar("learn_rate", self.learn_rate)
            tf.summary.scalar("iou_loss", self.iou_loss)
            tf.summary.scalar("conf_loss", self.conf_loss)
            tf.summary.scalar("prob_loss", self.prob_loss)
            tf.summary.scalar("total_loss", self.loss)

            logdir = "./data/log/"
            if os.path.exists(logdir): shutil.rmtree(logdir)
            os.mkdir(logdir)
            self.write_op = tf.summary.merge_all()
            self.summary_writer = tf.summary.FileWriter(logdir, graph=self.sess.graph)

    def train(self):
        self.sess.run(tf.global_variables_initializer())
        frozen_epoch = 0
        try:
            print('=> Restoring weights from: %s ... ' % self.initial_weight)
            self.loader.restore(self.sess, self.initial_weight)
            if '-' in self.initial_weight:
                frozen_epoch = self.initial_weight.split('-')[-1]
                frozen_epoch = int(frozen_epoch)
        except:
            print('=> %s does not exist !!!' % self.initial_weight)
            print('=> Now it starts to train YOLOV4 from scratch ...')
            self.first_stage_epochs = 0

        saving = 0.0
        early_stop = 0
        for epoch in range(1+frozen_epoch, 1+self.first_stage_epochs+self.second_stage_epochs):
            if epoch <= self.first_stage_epochs:
                train_op = self.train_op_with_frozen_variables
            else:
                train_op = self.train_op_with_all_variables

            pbar = tqdm(self.trainset)
            train_epoch_loss, val_epoch_loss = [], []

            for train_data in pbar:
                _, summary, train_step_loss, global_step_val = self.sess.run(
                    [train_op, self.write_op, self.loss, self.global_step], feed_dict={
                                                self.input_data:   train_data[0],
                                                self.label_sbbox:  train_data[1],
                                                self.label_mbbox:  train_data[2],
                                                self.label_lbbox:  train_data[3],
                                                self.true_sbboxes: train_data[4],
                                                self.true_mbboxes: train_data[5],
                                                self.true_lbboxes: train_data[6],
                                                self.trainable:    True,
                })

                train_epoch_loss.append(train_step_loss)
                self.summary_writer.add_summary(summary, global_step_val)
                pbar.set_description("train loss: %.2f" % train_step_loss)

            for val_data in self.valset:
                val_step_loss = self.sess.run(self.loss, feed_dict={
                                                self.input_data:   val_data[0],
                                                self.label_sbbox:  val_data[1],
                                                self.label_mbbox:  val_data[2],
                                                self.label_lbbox:  val_data[3],
                                                self.true_sbboxes: val_data[4],
                                                self.true_mbboxes: val_data[5],
                                                self.true_lbboxes: val_data[6],
                                                self.trainable:    False,
                })

                val_epoch_loss.append(val_step_loss)

            train_epoch_loss, val_epoch_loss = np.mean(train_epoch_loss), np.mean(val_epoch_loss)
            train_epoch_loss = np.mean(train_epoch_loss)
            ckpt_file = "./checkpoint/yolov4_val_loss=%.4f.ckpt" % val_epoch_loss
            log_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
            if saving == 0.0:
                if frozen_epoch > 0:
                    org_val_loss = float('.'.join(self.initial_weight.split('=')[-1].split('.')[:2]))
                    saving = org_val_loss
                else:
                    saving = val_epoch_loss
                print("=> Epoch: %2d Time: %s Val loss: %.2f" % (epoch, log_time, val_epoch_loss))
            elif saving > val_epoch_loss:
                print("=> Epoch: %2d Time: %s Train loss: %.2f Val loss: %.2f Saving %s ..."
                      % (epoch, log_time, train_epoch_loss, val_epoch_loss, ckpt_file))
                self.saver.save(self.sess, ckpt_file, global_step=epoch)
                saving = val_epoch_loss
                early_stop = 0
            else:
                if early_stop > 10:
                    self.saver.save(self.sess, ckpt_file, global_step=epoch)
                    break
                print("=> Epoch: %2d Time: %s Val loss: %.2f" % (epoch, log_time, val_epoch_loss))
                early_stop += 1
        self.saver.save(self.sess, ckpt_file, global_step=epoch)

if __name__ == '__main__':
    YoloTrain().train()


import os
import time
from datetime import datetime
from glob import glob
from multiprocessing import Pool
from contextlib import closing
from natsort import natsorted
import numpy as np
# from numba import guvectorize, float32
from scipy.stats.stats import pearsonr
import pandas as pd
import seaborn as sns
import argparse
import pickle
import mdtraj as md
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
from matplotlib.widgets import Button
import matplotlib
matplotlib.use('TkAgg')  # MUST BE CALLED BEFORE IMPORTING plt
import matplotlib.pyplot as plt
from pyvol.spheres import Spheres

# @guvectorize([(float32[:], float32[:], float32)], '(n),(n)->()')#, nopython=True)
# def dist_v(pid, gp, res):
#     res = np.linalg.norm(pid - gp)
#     print(res)
#     return

def dist(pid, gp):
    return np.linalg.norm(pid - gp)

class ResultsSet:
    def __init__(self, *args, **kwargs):
        self.cutoff = 3.0
        try:
            # self.wrkdr = os.path.dirname(kwargs["pattern"])
            self.wrkdr = os.getcwd()
            self.folders = natsorted(glob(kwargs["pattern"]))
        except KeyError:
            self.wrkdr = os.curdir
            self.folders = []
        self.n = len(self.folders)
        self.pocket_names = [[] for _ in range(self.n)]
        self.volumes = np.zeros(shape=self.n, dtype=np.float_)
        self.xval = None
        self.frames = []
        self.results = []
        self.pocket_legend = [] # marked for deprecation
        self.pocket_IDs = None
        self.all_volumes = None
        self.ref_sel = None
        self.ref_poc = None
        self.n_clusters = None
        return

    def save(self, fname="resultset.p"):
        fout = open(fname, "wb")
        pickle.dump(self, fout)
        fout.close()
        return

    def load(self, fname="resultset.p"):
        fout = open(fname, "rb")
        loaded = pickle.load(fout)
        fout.close()
        return loaded

    def parse(self):
        for f in self.folders:
            try:
                self.frames.append(glob(os.path.join(f, "*pdb"))[0])
            except IndexError:
                print(f, "does not contain a pdb file")
            try:
                summary = glob(os.path.join(f, "*rept"))[0]
                self.results.append(pd.read_csv(summary))
            except IndexError:
                self.results.append(None)
            try:
                self.xval = np.loadtxt(os.path.join(self.wrkdr, "xval.dat"))
            except IOError:
                pass
        return

    def parse_reference(self, fname):
        self.ref_poc = pd.read_csv(fname, sep=';', header=None).to_numpy()
        return

    def sum_pockets(self):
        self.volumes = np.zeros(shape=self.n, dtype=np.float_)
        self.pocket_names = [[] for _ in range(self.n)]
        for j in range(self.n):
            if self.results[j] is None:
                continue
            for i in range(self.results[j].shape[0]):
                if self.results[j].distance[i] < self.cutoff:
                    self.volumes[j] += self.results[j].volume[i]
                    self.pocket_names[j].append(self.results[j].name[i])
        for i in range(len(self.volumes)):
            if self.volumes[i] == 0:
                self.volumes[i] = float("NaN")
        return

    def update(self, event):
        if event.button == 1:
            self.cutoff += 0.5
            self.sum_pockets()
        plt.clf()
        plt.plot(np.arange(self.n), self.volumes, "c-")
        plt.title(self.cutoff)
        plt.draw()
        return

    def inc(self, event):
        if event.button == 1:
            self.cutoff += 0.5
            self.sum_pockets()
        self.figdata.set_ydata(self.volumes)
        self.ax.set_title(self.cutoff)
        self.ax.set_ylim((0, 1.05*np.nanmax(self.volumes)))
        plt.draw()
        return

    def dcr(self, event):
        if event.button == 1:
            self.cutoff -= 0.5
            self.sum_pockets()
        self.figdata.set_ydata(self.volumes)
        self.ax.set_title(self.cutoff)
        self.ax.set_ylim((0, 1.05 * np.nanmax(self.volumes)))
        plt.draw()
        return

    def opt_cutoff(self):
        self.sum_pockets()
        plt.rcParams.update({'font.size': 14})
        f, self.ax = plt.subplots()
        self.ax.set_title(self.cutoff)
        self.ax.set_ylim((0, 1.05 * np.nanmax(self.volumes)))
        self.ax.set_ylabel("Volume / A$^3$")
        self.figdata, = self.ax.plot(np.arange(self.n), self.volumes, "o-")
        if self.xval is None:
            self.ax.xaxis.set_visible(False)
        else:
            self.ax.set_xticks([0, 24, 49, 74, 99])
            self.ax.set_xticklabels(["{:.1f}".format(self.xval[0]), "{:.1f}".format(self.xval[24]),
                                     "{:.1f}".format(self.xval[49]), "{:.1f}".format(self.xval[74]),
                                     "{:.1f}".format(self.xval[99])])
        axi = plt.axes([0.60, 0.0, 0.18, 0.07])
        axd = plt.axes([0.78, 0.0, 0.18, 0.07])
        bi = Button(axi, 'Increase')
        bi.on_clicked(self.inc)
        bd = Button(axd, 'Decrease')
        bd.on_clicked(self.dcr)
        plt.tight_layout()
        plt.show()
        delattr(self, "figdata")
        delattr(self, "ax")
        return

    def get_pocketID(self, index):
        """
        Generates distances to pocket for identification.
        """
        frame = md.load(self.frames[index])
        print(self.frames[index])
        refs = np.empty(shape=(len(self.ref_sel), 3), dtype=np.float_)
        for i in range(len(self.ref_sel)):
            if "CA" in self.ref_sel[i]:
                refs[i] = frame.xyz[0][frame.top.select(self.ref_sel[i])[0]]
            else:
                refs[i] = md.compute_center_of_mass(frame.atom_slice(frame.top.select(self.ref_sel[i])))
        refs = 10 * refs
        try:
            pocketID = np.empty(shape=(self.results[index].shape[0], refs.shape[0]), dtype=np.float32)
            for i in range(self.results[index].shape[0]):
                p = Spheres(spheres_file="{0:s}.obj".format(os.path.join(self.folders[index], self.results[index].name[i])))
                for j in range(refs.shape[0]):
                    nearest = p.nearest_coord_to_external(refs[j])
                    pocketID[i][j] = np.linalg.norm(refs[j] - nearest)
            return pocketID
        except AttributeError:
            return []

    def prepare_for_clustering(self):
        all = []
        self.n_clusters = 2
        for i in range(self.n):
            for pid in self.get_pocketID(i):
                all.append(pid)
            pid = self.get_pocketID(i)
            for j in range(pid.shape[0]):
                if j == self.n_clusters:
                    self.n_clusters += 1
                all.append(pid[j])
                self.pocket_legend.append((i, j))
        self.pocket_IDs = np.array(all)
        return

    def opt_cluster(self, krange=range(2, 16), plot=False):
        """
        Optimises the number of clusters.
        """
        # km = KMeans(n_clusters=8).fit(self.pocket_IDs)
        # https://blog.cambridgespark.com/how-to-determine-the-optimal-number-of-clusters-for-k-means-clustering-14f27070048f
        Sum_of_squared_distances = []
        for k in krange:
            km = KMeans(n_clusters=k)
            km = km.fit(self.pocket_IDs)
            Sum_of_squared_distances.append(km.inertia_)
        if plot:
            plt.plot(krange, Sum_of_squared_distances, 'bx-')
            plt.xlabel('Number of clusters')
            plt.ylabel('Sum of squared distances')
            plt.show()
        # https://realpython.com/k-means-clustering-python/#choosing-the-appropriate-number-of-clusters
        # A list holds the silhouette coefficients for each k
        silhouette_coefficients = []
        # Notice you start at 2 clusters for silhouette coefficient
        for k in krange:
            kmeans = KMeans(n_clusters=k)
            kmeans.fit(self.pocket_IDs)
            score = silhouette_score(self.pocket_IDs, kmeans.labels_)
            silhouette_coefficients.append(score)
        if plot:
            plt.plot(krange, silhouette_coefficients)
            plt.xticks(krange)
            plt.xlabel("Number of Clusters")
            plt.ylabel("Silhouette Coefficient")
            plt.show()
        self.n_clusters = krange[silhouette_coefficients.index(max(silhouette_coefficients))]
        return

    def plot_cluster(self):
        """
        Plots clustered ID vectors in the first 3 dimensions.
        """
        km = KMeans(n_clusters=self.n_clusters)
        km = km.fit(self.pocket_IDs)
        centroids = km.cluster_centers_
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        ax.scatter(self.pocket_IDs[:, 0], self.pocket_IDs[:, 1], self.pocket_IDs[:, 2],
                   c=km.labels_.astype(float), s=50, alpha=0.5)
        ax.scatter(centroids[:, 0], centroids[:, 1], centroids[:, 2], c='red', s=50)
        plt.show()
        return

    def identify_clusters(self, plot=False):
        self.all_volumes = np.zeros(shape=(self.n, self.n_clusters), dtype=np.float_)
        km = KMeans(n_clusters=self.n_clusters)
        km = km.fit(self.pocket_IDs)
        for i in range(self.pocket_IDs.shape[0]):
            self.all_volumes[self.pocket_legend[i][0]][km.labels_[i]] +=\
                self.results[self.pocket_legend[i][0]].volume[self.pocket_legend[i][1]]
        self.all_volumes[self.all_volumes == 0] = 'nan'
        if plot:
            f, ax = plt.subplots()
            x = range(self.n)
            for i in range(self.n_clusters):
                ax.plot(x, self.all_volumes[:, i], "o-", label="Cluster {:d}".format(i + 1))
            plt.xlabel("Frame")
            plt.ylabel("Volume / A$^3$")
            plt.legend(loc="best")
            plt.show()
        return

    def process(self):
        """
        Processed all frames for pocket IDs.
        """
        # how to flatten:
        # v = p.reshape((1020, 3))
        # w = v[~np.isnan(v).any(axis=1)]
        # all = []
        # for i in range(self.n):
        #     all.append(self.get_pocketID(i))
        if self.ref_sel is None:
            frame = md.load(self.frames[0])
            # self.ref_sel = ["resid {0:d} to {1:d}".format(i, i + 9) for i in range(0, frame.n_residues, 10)]
            self.ref_sel = ["resid {0:d} and name CA".format(i) for i in range(2, frame.n_residues, 10)]
        if int(self.n/args.threads) == 0:
            chunk = 1
        else:
            chunk = int(self.n/args.threads)
        # for i in range(self.n):
        #     rv = self.get_pocketID(i)
        with closing(Pool(processes=int(args.threads))) as pool:
            all = pool.map(self.get_pocketID, range(self.n), chunk)
            pool.terminate()
        b = np.zeros(shape=(len(all), len(max(all, key=lambda x: len(x))), len(self.ref_sel)), dtype=np.float32)
        for i, j in enumerate(all):
            try:
                b[i][0:len(j)] = j
            except ValueError:
                pass
        b[b == 0] = 'nan'
        self.pocket_IDs = b
        return

    def grid(self, bins=10, plot=False, cutoff_factor=0.5, keep_bin_ratio=0.25):
        all_volumes = []
        dimensions = self.pocket_IDs.shape[2]
        boundaries = [np.nanmin(self.pocket_IDs, axis=(0, 1)), np.nanmax(self.pocket_IDs, axis=(0, 1))]
        lspace = [np.linspace(boundaries[0][i], boundaries[1][i], bins) for i in range(len(boundaries[0]))]
        mesh = np.array(np.meshgrid(*lspace), dtype=np.float32)
        cutoff = dist(mesh[:, 0, 0, 0], mesh[:, 1, 1, 1]) * cutoff_factor
        mesh = mesh.swapaxes(0, dimensions).reshape(bins**dimensions, 3)
        pockets = []
        for i in range(mesh.shape[0]):
            dists = np.empty(shape=(self.n, self.pocket_IDs.shape[1]), dtype=np.float32)
            for j in range(self.n):
                for k in range(self.pocket_IDs.shape[1]):
                    dists[j][k] = dist(self.pocket_IDs[j][k], mesh[i])
            min_dist = np.nanmin(dists, axis=1)
            if np.nanmin(min_dist) < cutoff:
                if len(np.where(min_dist < cutoff)[0]) > keep_bin_ratio*self.n:
                    print(mesh[i])
                    idx = np.where(dists < cutoff)
                    pockets.append(idx)
                    vol = []
                    for j in range(self.n):
                        if j in idx[0]:
                            k = np.where(idx[0] == j)[0][0]
                            vol.append(self.results[j].volume[idx[1][k]])
                        else:
                            vol.append('nan')
                    all_volumes.append(vol)
        all_volumes = np.array(all_volumes, dtype=np.float32)
        cov = np.zeros(shape=(all_volumes.shape[0], all_volumes.shape[0]))
        for i in range(all_volumes.shape[0]):
            for j in range(i):
                x = all_volumes[i]
                y = all_volumes[j]
                bad = ~np.logical_or(np.isnan(x), np.isnan(y))
                cov[i][j] = pearsonr(np.compress(bad, x), np.compress(bad, y))[0]
        groups = [0]
        for i in range(1, all_volumes.shape[0]):
            for j in range(i):
                FOUND = False
                if cov[i][j] > 0.8:
                    groups.append(groups[j])
                    FOUND = True
                    break
            if not FOUND:
                groups.append(np.max(groups) + 1)
        groups = np.array(groups)
        filtered_volumes = []
        for i in range(np.max(groups) + 1):
            filtered_volumes.append(np.nanmean(all_volumes[np.where(groups == i)[0], :], axis=0))
        filtered_volumes = np.array(filtered_volumes)
        if plot:
            f, ax = plt.subplots()
            x = range(self.n)
            for i in range(filtered_volumes.shape[0]):
                ax.plot(x, filtered_volumes[i], "o-", label="Cluster {:d}".format(i + 1))
            plt.xlabel("Frame")
            plt.ylabel("Volume / A$^3$")
            plt.legend(loc="best")
            plt.show()
        return

    def write_pml(self, fname="pyvol", stride=1, spec="helicase"):
        """
        This function was specific to the SARS-CoV-2 helicase visualization.
        The generalization is underway.
        """
        if spec == "helicase":
            tmpl = "/mnt/data/covid/pyvol/PCA/visualization_template.pse"
        with open(os.path.join(self.wrkdr, "{0:s}_{1:%Y%m%d%H%M}.pml".format(fname, datetime.now())), "w") as fout:
            if spec == "helicase":
                fout.write("load {:s}\n".format(tmpl))
            for i in range(0, self.n, stride):
                fout.write("load {0:s}, frame{1:03d}\n".format(self.frames[i], i))
                if spec == "helicase":
                    fout.write("super frame{:03d}, struct\n".format(i))
                else:
                    fout.write("super frame{:03d}, frame001\n".format(i))
                for j in range(len(self.pocket_names[i])):
                    fout.write("load {0:s}.xyz, frame{1:d}p{2:d}\n".format(os.path.join(self.folders[i],
                                                                                self.pocket_names[i][j]), i, j))
                    fout.write(
                        "color cyan, frame{0:d}p{1:d}; hide everything, frame{0:d}p{1:d}; show surface, frame{0:d}p{1:d}\n".format(
                            i, j))
                    fout.write("matrix_copy frame{0:03d}, frame{0:d}p{1:d};\n".format(i, j))
            if spec == "helicase":
                fout.write("select domain1, frame* and i. 1-240\n")
                fout.write("select RecA1, frame* and i. 241-442\n")
                fout.write("select RecA2, frame* and i. 443-601\n")
                fout.write("color deepteal, domain1 and elem C\n")
                fout.write("color yellow, RecA1 and elem C\n")
                fout.write("color hotpink, RecA2 and elem C\n")
            fout.write("zoom frame*;\n")
            fout.write("mset; rewind\n")
            fout.write("mset 1x{:d};\n".format(int(self.n/stride)))
            k = 0
            for i in range(0, self.n, stride):
                enable_pockets = ""
                for j in range(len(self.pocket_names[i])):
                    enable_pockets += "enable frame{0:d}p{1:d};".format(i, j)
                fout.write("disable all; enable struct; enable frame{0:03d}; {1:s}scene s{0:03d}, store;\n".format(
                    i, enable_pockets))
                fout.write("mview store, {1:d}, scene=s{0:03d};\n".format(i, k-1))
                k += 1
        return

    def write_vis_pml(self, fname="pyvol_movie", stride=1):
        """
        This function is specific to the SARS-CoV-2 helicase visualization. Resaving movie frames after adjustment.
        """
        with open(os.path.join(self.wrkdr, "{:s}.pml".format(fname)), "w") as fout:
            fout.write("mset; rewind\n")
            fout.write("mset 1x{:d};\n".format(int(self.n/stride)))
            k = 0
            for i in range(0, self.n, stride):
                enable_pockets = ""
                for j in range(len(self.pocket_names[i])):
                    enable_pockets += "enable frame{0:d}p{1:d};".format(i, j)
                fout.write("disable all; enable struct; enable frame{0:03d}; {1:s}scene s{0:03d}, store;\n".format(
                    i, enable_pockets))
                fout.write("mview store, {1:d}, scene=s{0:03d};\n".format(i, k))
                k += 1
        return

    def write_RNA_pml(self):
        """
        This function is specific to the SARS-CoV-2 helicase visualization, with the extra assumption, that the
        largest pocket is the RNA
        """
        with open(os.path.join(self.wrkdr, "pyvol_largest_{:%Y%m%d%H%M}.pml".format(datetime.now())), "w") as fout:
            for i in range(self.n):
                if self.results[i] is not None:
                    fout.write("load {0:s}.xyz, largest{1:03d}\n".format(
                        os.path.join(self.folders[i], self.results[i].name[0]), i))
            fout.write("hide everything, largest*; show surface, largest*\n")
            fout.write("zoom struct;\n")
            for i in range(self.n):
                fout.write("scene s{0:03d}, recall; enable largest{0:03d}; scene s{0:03d}, update;\n".format(i))
                fout.write("mview store, {0:d}, scene=s{0:03d};\n".format(i))
        return

    def plot_selected(self, fout=None, label="ATP"):
        plt.rcParams.update({'font.size': 14})
        f, a = plt.subplots()
        plt.ylabel("Volume / A$^3$")
        if self.xval is None:
            a.xaxis.set_visible(False)
        else:
            a.set_xticks([0, 24, 49, 74, 99])
            a.set_xticklabels(["{:.1f}".format(self.xval[0]), "{:.1f}".format(self.xval[24]),
                               "{:.1f}".format(self.xval[49]), "{:.1f}".format(self.xval[74]),
                               "{:.1f}".format(self.xval[99])])
        a.plot(np.arange(self.n), self.volumes, "o-", label=label)
        plt.legend(loc="best")
        plt.tight_layout()
        if fout is None:
            plt.show()
        else:
            plt.savefig(fout, dpi=300)
            np.savetxt(fout.replace(".png", ".txt"), self.volumes)
        return

    def plot_selected_and_largest(self, fout=None):
        plt.rcParams.update({'font.size': 14})
        f, a = plt.subplots()
        plt.ylabel("Volume / A$^3$")
        if self.xval is None:
            a.xaxis.set_visible(False)
        else:
            a.set_xticks([0, 24, 49, 74, 99])
            a.set_xticklabels(["{:.1f}".format(self.xval[0]), "{:.1f}".format(self.xval[24]),
                               "{:.1f}".format(self.xval[49]), "{:.1f}".format(self.xval[74]),
                               "{:.1f}".format(self.xval[99])])
        a.plot(np.arange(self.n), self.volumes, "o-", label="ATP")
        largest = [r.volume[0] if r is not None else float("NaN") for r in self.results]
        a.plot(np.arange(self.n), largest, "o-", label="largest")
        a.set_ylim([0, np.nanmin([np.nanmax(largest), 3800])+200])
        plt.legend(loc="best")
        plt.tight_layout()
        if fout is None:
            plt.show()
        else:
            plt.savefig(fout, dpi=300)
        return

    def follow_pocket(self, spec, identifiers=6, mode="single", stride=1):
        if self.ref_poc is None:
            raise ValueError("No pocket is selected to be tracked.")
        if self.pocket_IDs is None:
            self.process()
        for r in self.ref_poc:
            if identifiers == "all":
                PIDs = self.pocket_IDs
            else:
                PIDs = self.select_ID(r, identifiers, plot=False)
            ref = PIDs[r[0]][r[1]]
            dists = np.empty(shape=(self.n, PIDs.shape[1]), dtype=np.float32)
            for j in range(self.n):
                for k in range(PIDs.shape[1]):
                    dists[j][k] = dist(PIDs[j][k], ref)
            if mode == "single":
                min_dist = np.nanmin(dists, axis=1)
                sns.displot(min_dist[~np.isnan(min_dist)], kde=True, binwidth=0.3)
                plt.show()
            elif mode == "multiple":
                flat_dist = dists.flatten()
                sns.displot(flat_dist[~np.isnan(flat_dist)], kde=True, binwidth=0.3)
                plt.show()
            while True:
                ans = input("Please enter a distance cutoff:")
                try:
                    cutoff = float(ans)
                    break
                except ValueError:
                    pass
            # reinitialize selected pockets
            self.pocket_names = [[] for _ in range(self.n)]
            self.volumes = np.zeros(shape=self.n, dtype=np.float_)
            if mode == "single":
                for i in range(self.n):
                    if min_dist[i] < cutoff:
                        self.pocket_names[i].append(self.results[i].
                                                    name[np.where(dists[i] == np.nanmin(dists[i]))[0][0]])
                        self.volumes[i] = self.results[i].volume[np.where(dists[i] == np.nanmin(dists[i]))[0][0]]
            elif mode == "multiple":
                for i in range(self.n):
                    for j in range(dists.shape[1]):
                        if dists[i][j] < cutoff:
                            self.pocket_names[i].append(self.results[i].name[j])
                            self.volumes[i] += self.results[i].volume[j]
            # self.volumes = np.where(self.volumes == 0, 199, self.volumes)
            self.volumes = np.where(self.volumes > 10000, np.nan, self.volumes)
            self.write_pml(fname="pocket_tracking_{0:d}_{1:d}".format(r[0], r[1]), stride=stride, spec=spec)
            self.write_vis_pml(fname="pocket_vis_{0:d}_{1:d}".format(r[0], r[1]), stride=stride)
            self.plot_selected(fout="pocket_{0:d}_{1:d}_{2:.1f}.png".format(r[0], r[1], cutoff),
                               label="{0:d}_{1:d}_{2:.1f}".format(r[0], r[1], cutoff))
        return

    def select_ID(self, r, dim, plot=False):
        refID = self.pocket_IDs[r[0]][r[1]]
        idx = refID.argsort()[:dim]
        selected = self.pocket_IDs[:, :, idx]
        if plot:
            frame = md.load(self.frames[r[0]])
            refs = np.empty(shape=(len(self.ref_sel), 3), dtype=np.float32)
            for i in range(len(self.ref_sel)):
                if "CA" in self.ref_sel[i]:
                    refs[i] = frame.xyz[0][frame.top.select(self.ref_sel[i])[0]]
                else:
                    refs[i] = md.compute_center_of_mass(frame.atom_slice(frame.top.select(self.ref_sel[i])))
            refs = 10 * refs
            fig = plt.figure()
            ax = fig.add_subplot(111, projection='3d')
            ax.scatter(refs[:, 0], refs[:, 1], refs[:, 2], s=50, alpha=0.5)
            ax.scatter(refs[idx, 0], refs[idx, 1], refs[idx, 2], c='red', s=50)
            plt.show()
        return selected

parser = argparse.ArgumentParser()
parser.add_argument("-p", "--pattern")
parser.add_argument("--fig")
parser.add_argument("-t", "--threads", type=int, default=4)
parser.add_argument("--ps", type=int, default=100, help="Stride for pml printing.")
parser.add_argument("-s", "--selected_pockets",
                    help="Zero based index pairs <WINDOW>;<POCKET> for selecting reference in a file")
parser.add_argument("-S", "--specific", default=None, help="Name a predefined topic specification.")
parser.add_argument("-m", "--mode", default="single", help="Pocket tracking mode {single|multiple}")
parser.add_argument("-l", "--load", help="Processed result-set file, i.e. 'ref_ca.p'")
args = parser.parse_args()
start = time.time()
# ref_selection = ["resSeq <= 240", "240 < resSeq <= 445", "445 < resSeq"]
# ref_selection = ["resSeq <= 100", "100 < resSeq <= 200", "200 < resSeq <= 300",
#                  "300 < resSeq <= 400", "400 < resSeq <= 500", "500 < resSeq"]
if args.pattern is not None:
    rs = ResultsSet(pattern=args.pattern)
else:
    rs = ResultsSet()
# rs.ref_sel = ref_selection
if args.load is None:
    rs.parse()
    rs.process()
else:
    rs = rs.load(fname=args.load)
rs.save(fname="ref_ca.p")
rs.parse_reference(args.selected_pockets)
rs.follow_pocket(args.specific, mode=args.mode, identifiers='all', stride=args.ps)

# rs.opt_cluster()
# rs.n_clusters = 10
# rs.plot_cluster()
# rs.identify_clusters(plot=True)

# rs.grid(plot=True)
# rs.opt_cutoff()
# rs.write_pml()
# rs.write_RNA_pml()
# rs.write_vis_pml()
# rs.plot_selected(args.fig)
end =time.time()
print(end - start)


#!/usr/bin/env python

import rospy
from nav_msgs.msg import OccupancyGrid
from std_msgs.msg import Int16
from geometry_msgs.msg import Twist
from geometry_msgs.msg import PointStamped
from geometry_msgs.msg import *
from compressed_image_transport import *
from nav_msgs.msg import Odometry
from sensor_msgs.msg import LaserScan
from tf2_msgs.msg import TFMessage
from std_srvs.srv import Empty
from visualization_msgs.msg import Marker
from visualization_msgs.msg import MarkerArray
from  mazeSolver import Solver
import tf
import math
import cv2
import numpy as np
import threading
from std_msgs.msg import String
import heapq
import timeit
from time import sleep
import tensorflow as tf

from collections import deque
import dqn
from typing import List
import random

pi = np.pi
class Maze_pathfinder():
    def __init__(self):

        self._sub = rospy.Subscriber('/map', OccupancyGrid, self.callback, queue_size=1)
        self._sub = rospy.Subscriber('/odom', Odometry, self.callback2, queue_size=1)

        self._sub = rospy.Subscriber('/clicked_point', PointStamped, self.callback4, queue_size=1)

        self._pub = rospy.Publisher('/cmd_vel', Twist, queue_size=1)
        self.count = 0
        #self._sub = rospy.Subscriber('/scan', LaserScan, self.getState, queue_size=1)
        # variables used in maze solve

        self.reset_proxy = rospy.ServiceProxy('gazebo/reset_simulation', Empty) # reset service

        self.action_size = 11
        self.input_size = 37

    def callback(self, map_data):
        return

    def callback2(self, odometry):
        ## converting odometry to x position, y position and theta
        return
    def callback4(self, clicked_point):
        return

    def publishing_vel(self, angular_x, angular_y, angular_z, linear_x, linear_y, linear_z):
        vel = Twist()
        vel.angular.x = angular_x
        vel.angular.y = angular_y
        vel.angular.z = angular_z
        vel.linear.x = linear_x
        vel.linear.y = linear_y
        vel.linear.z = linear_z
        self._pub.publish(vel)

    def getState(self, scan):
        scan_range = []
        min_range = 0.13
        done = False

        scan_idx = []
        scan_idx.append(0)
        for i in range(1,19):
            scan_idx.append(i * 5)
            scan_idx.append(359 - i*5)

        for i in scan_idx:
            if scan.ranges[i] == float('Inf'):
                scan_range.append(3.5)
            elif np.isnan(scan.ranges[i]):
                scan_range.append(0)
            else:
                scan_range.append(scan.ranges[i])

        if min_range > min(scan_range) > 0:
            done = True

        return scan_range, done

    def setReward(self, state, done, action):
        right = state[35]
        left = state[36]
        reward = 1/abs(right - left)
        #print "reward : ", reward
        if done:
            #rospy.loginfo("Collision!!")
            reward = -200
            self.move(0, 0)
        return reward

    def step(self, action):
        max_angular_vel = 0.3
        #print "action : ", action
        ang_vel = (float(self.action_size - 1)/2 - action) / (self.action_size/2) * max_angular_vel

        speed = 0.22
        self.move(speed, ang_vel)

        data = None
        while data is None:
            try:
                data = rospy.wait_for_message('/scan', LaserScan, timeout=5)
            except:
                pass

        state, done = self.getState(data)
        reward = self.setReward(state, done, action)

        return np.asarray(state), reward, done
    def reset(self):

        try:
            self.reset_proxy()
            #print "reset"
        except (rospy.ServiceException) as e:
            print "gazebo/reset_simulation service call failed"

        data = None
        while data is None:
            try:
                #print "wait"
                data = rospy.wait_for_message('/scan', LaserScan, timeout=5)
            except:
                pass
        state, done = self.getState(data)
        return state
    def move(self, speed, direction):
        #print "speed : ", speed, " direction : ", direction
        self.publishing_vel(0, 0, direction, speed, 0, 0)

    def main(self):
        rospy.spin()

if __name__ == '__main__':
    rospy.init_node('maze_pathfinder')
    mazesolver = Maze_pathfinder()
    mazesolver.main()