# ------------------------------------------------------------------------
# Copyright (C) 2017 Nobuya Sato, Hiori Kino, and Takashi Miyake
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ------------------------------------------------------------------------

""" Soiap potential."""

# [note] set environmental variable SOIAP_COMMAND for soiap command

import os
import numpy as np

import ase
from ase.calculators.calculator import Calculator, all_changes

#---- calculation class for soiap
class Soiap(Calculator):

    implemented_properties = ["energy", "forces", "stress"]

    default_parameters = dict(
        number_max_relax=10,
        md_mode=4,
        number_max_relax_cell=0,
        external_stress_v=[0.0, 0.0, 0.0],
        unit_time=1,
        time_step=1.0,
        mass_scale=0,
        mass_cell=5.e+4,
        md_mode_cell=2,
        th_force=5.e-5,
        th_stress=5.e-7,
        force_field=1
    )

    #---- calculate enregy, forces, and stress
    def calculate(self, atoms, properties, system_changes):
        Calculator.calculate(self,atoms,properties,system_changes)

        self.write_input()

        self.run()

        self.read_results()
    # end def calculate

    #---- write input file for the solve
    def write_input(self, atoms=None, properties=None, system_changes=None):
        # make convinience local variables
        cell      = self.atoms.cell
        symbols   = self.atoms.get_chemical_symbols()
        numbers   = self.atoms.get_atomic_numbers()
        positions = self.atoms.get_scaled_positions()
        constraint_indices = []
        for c in self.atoms.constraints:
            constraint_indices.extend(c.index)
        # end for

        numbers_uniq = np.unique(numbers)
        numbers_index = {}
        for e in range(len(numbers_uniq)):
            numbers_index[numbers_uniq[e]] = e+1
        # end for

        # write data to the input file
        fd = open( "input", "w" )
        fd.write("number_atom    %d\n" % len(self.atoms) )
        fd.write("number_element %d\n" % len(numbers_uniq) )
        fd.write("\n")

        fd.write("unit_vec\n")
        fd.write("( %f, %f, %f )\n" % ( cell[0][0], cell[0][1], cell[0][2]) )
        fd.write("( %f, %f, %f )\n" % ( cell[1][0], cell[1][1], cell[1][2]) )
        fd.write("( %f, %f, %f )\n" % ( cell[2][0], cell[2][1], cell[2][2]) )
        fd.write("\n")

        fd.write("unit_vec_column 1 ! column\n")
        fd.write("unit_length     1 ! Angstrom\n")
        fd.write("atom_pos        1 ! relative\n")
        fd.write("\n")

        fd.write("atom_list\n")
        for a in range(len(self.atoms)):
            fd.write("%2d %f %f %f %d %d ! %s\n" %
              (numbers[a], positions[a][0], positions[a][1], positions[a][2],
               numbers_index[numbers[a]],
               a not in constraint_indices, symbols[a] ) )
        # end for
        fd.write("\n")

        keys = sorted(self.parameters)
        for key in keys:
            if isinstance(self.parameters[key],list) :
                fd.write(
                  "%s " % key
                + "".join(" %r" %
                    self.parameters[key][i]
                      for i in range(len(self.parameters[key])))
                + "\n")
            else:
                fd.write("%s %r\n" % (key,self.parameters[key]) )
            # end if
        # end for

        fd.close()
    # end def write_input

    #---- run the solver
    def run(self):
        # exec the solver
        if "SOIAP_COMMAND" in os.environ:
            soiap = os.environ["SOIAP_COMMAND"]
            exitcode = os.system("%s input > log.stdout" % soiap )
        else:
            raise RuntimeError("Please set SOIAP_COMMAND environment variable")
        # end if

        if exitcode != 0:
            raise RuntimeError("Soiap exited with exit code: %d.  " % exitcode)
        # end if
    # end def run

    #---- read struct, enregy, forces, and stress
    def read_results(self):
        self.read_struct()
        self.read_energy()
        self.read_forces()
        self.read_stress()
    # end def read_results

    #---- read cell and positions
    def read_struct(self):
        # init local variable
        cell = [[0,0,0],[0,0,0],[0,0,0]]
        positions = []

        # get the latest values
        fd = open("log.struc","r")
        line = fd.readline()
        while line:
            if line.find("*** unit vectors [Bohr]:") != -1:
                positions = []

                line = fd.readline()
                terms = line.split()
                cell[0][0] = float(terms[0])
                cell[1][0] = float(terms[1])
                cell[2][0] = float(terms[2])

                line = fd.readline()
                terms = line.split()
                cell[0][1] = float(terms[0])
                cell[1][1] = float(terms[1])
                cell[2][1] = float(terms[2])

                line = fd.readline()
                terms = line.split()
                cell[0][2] = float(terms[0])
                cell[1][2] = float(terms[1])
                cell[2][2] = float(terms[2])

                line = fd.readline() # skip messeges ...
                for a in range(len(self.atoms)):
                    line = fd.readline()
                    terms = line.split()
                    position = [ float(terms[0]), float(terms[1]), float(terms[2]) ]
                    positions.append(position)
                # end for
            # end if
            line = fd.readline()
        # end while

        fd.close()

        # conversion of unit from bohr to angstrom
        for i in range(3):
           for j in range(3):
              self.atoms.cell[i][j] = cell[i][j]*ase.units.Bohr
           # end for
        # end for

        self.atoms.set_scaled_positions(positions)
    # end read_struct

    #---- read energy
    def read_energy(self):
        # init local variable
        energy = 0.0

        # get the latest values
        fd = open("log.tote","r")
        line = fd.readline()
        while line:
            terms = line.split()
            if len(terms) == 7:
                energy = float(terms[2])
            # end if
            line = fd.readline()
        # end while

        fd.close()

        # no conversion of unit from eV to eV
        self.results["energy"] = energy
    # end def read_energy

    #---- read forces
    def read_forces(self):
        # init local variable
        forces = []

        # get the latest values
        fd = open("log.frc","r")
        line = fd.readline()
        while line:
            if line.find("*** atomic forces: QMD%loopc") != -1:
                forces = []
                for a in range(len(self.atoms)):
                    line = fd.readline()
                    terms = line.split()
                    force = [ float(terms[0]), float(terms[1]), float(terms[2]) ]
                    forces.append( np.array(force) )
                # end for
            # end if
            line = fd.readline()
        # end while

        fd.close()

        # conversion of unit from eV/bohr to eV/angstrom
        for a in range(len(forces)):
           for j in range(3):
              forces[a][j] = forces[a][j]*(ase.units.eV/ase.units.Bohr)
           # end for
        # end for

        self.results["forces"] = np.array(forces)
    # end def read_forces

    #---- read stress
    def read_stress(self):
        # init local variable
        stress = np.zeros(6)

        # get the latest values
        fd = open("log.strs","r")
        line = fd.readline()
        while line:
            if line.find("QMD%loopc, QMD%smax") != -1:
                line = fd.readline()
                terms = line.split()
                stress[0] = float(terms[0]) # xx
                stress[5] = float(terms[1]) # xy
                stress[4] = float(terms[2]) # xz

                line = fd.readline()
                terms = line.split()
                stress[1] = float(terms[1]) # yy
                stress[3] = float(terms[2]) # yz

                line = fd.readline()
                terms = line.split()
                stress[2] = float(terms[2]) # zz
            # end if
            line = fd.readline()
        # end while

        fd.close()

        # conversion of unit from eV/bohr^3 to eV/angstrom^3
        for i in range(6):
            stress[i] = stress[i]*((-1.0)*ase.units.eV/ase.units.Bohr**3)
        # end for

        self.results["stress"] = stress
    # end def read_stress

# end class Soiap


import os
import time
import logging, csv
import numpy as np
import tensorflow as tf
from collections import deque

from utils import find_trainable_variables, make_session, set_global_seeds
from utils import cat_entropy, mse, explained_variance, normalize_obs, discount_with_dones
from utils import add_to_collection_rnn_state

# TODO nsteps == batch_size
class Model(object):
    def __init__(self, policy, ob_space, ac_space, lr, cliprange,
                 max_grad_norm, units_per_hlayer, activ_fcn, log_interval, logdir,
                 nenvs, nbatch_train, nsteps, ent_coef, vf_coef, keep_model, meta=False):
        self.logger = logging.getLogger(self.__module__ + "." + self.__class__.__name__)
        self.logger.info('Set up PPO learning agent')
        self.num_steps_trained = 0
        self.log_interval = log_interval

        sess = make_session()
        nact = ac_space.n
        nbatch = nenvs * nsteps

        self.global_step = tf.get_variable('global_step',
                                           [],
                                           tf.int32,
                                           tf.constant_initializer(0, tf.int32),
                                           trainable=False)

        eval_model = policy(sess, ob_space, ac_space, 1, 1, units_per_hlayer, reuse=False, activ_fcn=activ_fcn)
        act_model = policy(sess, ob_space, ac_space, nenvs, 1, units_per_hlayer, reuse=True, activ_fcn=activ_fcn)
        train_model = policy(sess, ob_space, ac_space, nenvs, nbatch_train//nenvs, units_per_hlayer, reuse=True, activ_fcn=activ_fcn)

        A = tf.placeholder(tf.int32, [None], name='A')
        ADV = tf.placeholder(tf.float32, [None], name='ADV')
        R = tf.placeholder(tf.float32, [None], name='R')
        OLDNEGLOGPAC = tf.placeholder(tf.float32, [None])
        OLDVPRED = tf.placeholder(tf.float32, [None])
        LR = lr
        CLIPRANGE = cliprange

        def get_loss(model, placeholder_dict):
            a = placeholder_dict["A"]
            adv = placeholder_dict["ADV"]
            r = placeholder_dict["R"]
            oldneglocpac = placeholder_dict["OLDNEGLOGPAC"]
            oldvpred =placeholder_dict["OLDVPRED"]
            clip_range = placeholder_dict["CLIPRANGE"]

            neglogpac = model.pd.neglogp(a)
            entropy = tf.reduce_mean(cat_entropy(model.pi_logit))
            vpred = model.vf
            vpredclipped = oldvpred + tf.clip_by_value(model.vf - oldvpred, - clip_range, clip_range)
            vf_losses1 = tf.square(vpred - r)
            vf_losses2 = tf.square(vpredclipped - r)
            vf_loss = .5 * tf.reduce_mean(tf.maximum(vf_losses1, vf_losses2))
            ratio = tf.exp(oldneglocpac - neglogpac)
            pg_losses = -adv * ratio
            pg_losses2 = -adv * tf.clip_by_value(ratio, 1.0 - clip_range, 1.0 + clip_range)
            pg_loss = tf.reduce_mean(tf.maximum(pg_losses, pg_losses2))
            approxkl = .5 * tf.reduce_mean(tf.square(neglogpac - oldneglocpac))
            clipfrac = tf.reduce_mean(tf.to_float(tf.greater(tf.abs(ratio - 1.0), clip_range)))
            return pg_loss, entropy, vf_loss, vpred, neglogpac, approxkl, clipfrac

        self.input_plchld = {'A':A, 'ADV':ADV, 'R':R, 'OLDNEGLOGPAC':OLDNEGLOGPAC, 'OLDVPRED':OLDVPRED, 'CLIPRANGE':CLIPRANGE}
        pg_loss, entropy, vf_loss, vpred, neglogpac, approxkl, clipfrac = \
            get_loss(train_model, self.input_plchld)
        loss = pg_loss - entropy * ent_coef + vf_loss * vf_coef
        with tf.variable_scope('model'):
            params = tf.trainable_variables()
        grads = tf.gradients(loss, params)
        if max_grad_norm is not None:
            grads, _grad_norm = tf.clip_by_global_norm(grads, max_grad_norm)
        grads = list(zip(grads, params))
        trainer = tf.train.AdamOptimizer(learning_rate=LR, epsilon=1e-5)
        _train = [trainer.apply_gradients(grads), self.global_step.assign_add(nsteps)]

        if log_interval > 0:
            for g, v in grads:
                if g is not None:
                    tf.summary.histogram("%s-grad" % v.name.replace(':', '_'), g)
            for p in params:
                if p is not None:
                    tf.summary.histogram("train/%s" % p.name.replace(':', '_'), p.value())
            tf.summary.scalar("train/pg_loss", pg_loss)
            tf.summary.scalar("train/vf_loss", vf_loss)
            tf.summary.scalar("train/entropy", entropy)
            tf.summary.histogram("others/ADV", ADV)
            tf.summary.histogram("others/neglocpac", neglogpac)
            tf.summary.histogram("others/vf", vpred)
            self.summary_step = tf.summary.merge_all()

        # Adding these to collection so we can restore them again
        tf.add_to_collection('inputs', eval_model.X)
        # tf.add_to_collection('pi', eval_model.pi)
        # tf.add_to_collection('pi_logit', eval_model.pi_logit)
        tf.add_to_collection('val', eval_model.vf)
        tf.add_to_collection('step', eval_model.ac)
        if eval_model.initial_state is not None:
            add_to_collection_rnn_state('state_in', eval_model.rnn_state_in)
            add_to_collection_rnn_state('state_out', eval_model.rnn_state_out)

        tf.global_variables_initializer().run(session=sess)

        def train(obs, returns, masks, actions, values, neglogpacs, states=None):
            advs = returns - values
            advs = (advs - advs.mean()) / (advs.std() + 1e-8)  # TODO why?
            td_map = {train_model.X:obs, A:actions, ADV:advs, R:returns, # LR:lr, CLIPRANGE:cliprange,
                    OLDNEGLOGPAC:neglogpacs, OLDVPRED:values}
            if states is not None:
                td_map[train_model.rnn_state_in] = states
                # td_map[train_model.M] = masks  # TODO masks are required to reset rnn state if episode terminated

            policy_loss, value_loss, policy_entropy, app_kl, clipratio, gs = \
                sess.run([pg_loss, vf_loss, entropy, approxkl, clipfrac, self.global_step, _train], td_map)[:-1]
            # TF summary logging
            if log_interval > 0 and (self.num_steps_trained % self.log_interval == 0):
                self.logger.info('Save summary of network weights, grads and losses.')
                summary_str = sess.run(self.summary_step, td_map)
                self.summary_writer.add_summary(tf.Summary.FromString(summary_str), gs)
            self.num_steps_trained += 1
            return policy_loss, value_loss, policy_entropy, app_kl, clipratio

        # saver = tf.train.Saver(max_to_keep=2)
        saver = tf.train.Saver(max_to_keep=keep_model)
        self.loss_names = ['policy_loss', 'value_loss', 'policy_entropy', 'approxkl', 'clipfrac']

        def save(f_name):
            gs = sess.run(self.global_step)
            self.logger.info('Save network parameters of model at global step %s' % gs)
            saver.save(sess, os.path.join(logdir, f_name), global_step=gs)

        def load(load_path):
            saver.restore(sess, load_path)

        def test_run(env, n_eps, n_pipes):
            self.logger.info('Evaluating current agent')
            ep_return = []
            ep_length = []
            for i in range(0, n_eps):
                obs = env.reset()
                obs = normalize_obs(obs)
                done = False
                if eval_model.initial_state is not None:
                    if len(eval_model.initial_state) > 1:
                        rnn_s_in = (np.zeros(eval_model.initial_state[0].shape), np.zeros(eval_model.initial_state[1].shape))  # init lstm cell vector
                    else:
                        rnn_s_in = np.zeros(eval_model.initial_state.shape)  # init gru cell vector
                total_return = 0
                total_length = 0

                while not done and (total_return < n_pipes):
                    if eval_model.initial_state is not None:
                        pi, pi_log, act, rnn_s_out = sess.run([eval_model.pi, eval_model.pi_logit, eval_model.ac, eval_model.rnn_state_out], feed_dict={eval_model.X: [obs], eval_model.rnn_state_in: rnn_s_in})
                    else:
                        pi, pi_log, act = sess.run([eval_model.pi, eval_model.pi_logit, eval_model.ac], feed_dict={eval_model.X: [obs]})

                    ac = np.argmax(pi_log)
                    obs, reward, done, _ = env.step(ac)
                    obs = normalize_obs(obs)
                    total_length += 1
                    total_return += reward
                    if eval_model.initial_state is not None:
                        rnn_s_in = rnn_s_out
                self.logger.info('Episode %s: %s, %s' % (i, total_return, total_length))
                ep_length.append(total_length)
                ep_return.append(total_return)
            return ep_return

        self.get_loss = get_loss
        self.loss = loss
        self.trainer = trainer
        self.train = train
        self.train_model = train_model
        self.eval_model = eval_model
        self.act_model = act_model
        self.step = act_model.step
        self.value = act_model.value
        self.initial_state = act_model.initial_state
        self.save = save
        self.load = load
        self.test_run = test_run
        # tf.global_variables_initializer().run(session=sess) #pylint: disable=E1101

        # Set the summary writer to write to the given logdir if logging is enabled
        if log_interval > 0:
            self.summary_writer = tf.summary.FileWriter(logdir, graph_def=sess.graph_def)
        else:
            self.summary_writer = None

        self.sess = sess

    def get_summary_writer(self):
        return self.summary_writer

class Runner(object):

    def __init__(self, *, env, model, nsteps, gamma, lam, horizon=100, show_interval=0, summary_writer=None):
        self.logger = logging.getLogger(self.__module__ + "." + self.__class__.__name__)
        self.logger.info('Set up %s-step Runner' % nsteps)
        self.env = env
        self.model = model
        nd, = env.observation_space.shape
        self.nenv = nenv = env.num_envs
        self.obs = np.zeros((nenv,) + env.observation_space.shape, dtype=model.train_model.X.dtype.name)
        self.obs[:] = env.reset()
        self.gamma = gamma
        self.lam = lam
        self.nsteps = nsteps
        self.states = model.initial_state
        self.dones = [False for _ in range(nenv)]
        self.reward_window = [deque(maxlen=horizon) for _ in range(nenv)]

        # episode stats - ONLY required in episodic setting
        self.eplength = [0 for _ in range(nenv)]
        self.epreturn = [0 for _ in range(nenv)]
        self.return_threshold = -2.  # 0.0  # 40  # threshold below which no model will be saved.

        # rendering
        self.show_interval = show_interval
        self.ep_idx = [0 for _ in range(nenv)]

        self.summary_writer = summary_writer

    def run(self):
        mb_obs, mb_rewards, mb_actions, mb_values, mb_dones, mb_neglogpacs = [],[],[],[],[],[]
        mb_states = self.states
        for _ in range(self.nsteps):
            actions, pi, values, self.states, neglogpacs = self.model.step(self.obs, self.states, self.dones)
            mb_obs.append(np.copy(self.obs))
            mb_actions.append(actions)
            mb_values.append(values)
            mb_neglogpacs.append(neglogpacs)
            mb_dones.append(self.dones)
            obs, rewards, self.dones, _ = self.env.step(actions)
            self.obs[:] = normalize_obs(obs)
            mb_rewards.append(rewards)

            self.logger.debug('Observations: %s' % self.obs)

            # render only every i-th episode
            if self.show_interval != 0:
                if (self.ep_idx[0] % self.show_interval) == 0:
                    self.env.render()

            self.eplength = [self.eplength[i] + 1 for i in
                             range(self.nenv)]  # Todo use already implemented functions in run_ple_utils!!!
            self.epreturn = [self.epreturn[i] + rewards[i] for i in range(self.nenv)]
            [self.reward_window[i].append(rewards[i]) for i in range(self.nenv)]

            # Check for terminal states in every env - this is only used in terminating version of FlappyBird
            for i, done in enumerate(self.dones):  # i -> environment ID
                if done:
                    self.ep_idx[i] += 1
                    self.obs[i] = self.obs[i] * 0

                    # update tensorboard summary
                    if self.summary_writer is not None:
                        summary = tf.Summary()
                        summary.value.add(tag='envs/environment%s/episode_length' % i,
                                          simple_value=self.eplength[i])
                        summary.value.add(tag='envs/environment%s/episode_reward' % i,
                                          simple_value=self.epreturn[i])
                        self.summary_writer.add_summary(summary, self.ep_idx[i])  # self.global_step.eval())
                        self.summary_writer.flush()
                    # self.retbuffer.append(self.epreturn[i])
                    if self.epreturn[i] > self.return_threshold:
                        self.return_threshold = self.epreturn[i]
                        self.logger.info('Save model at max reward %s' % self.return_threshold)
                        self.model.save('inter_model')
                    self.eplength[i] = 0
                    self.epreturn[i] = 0

        # batch of steps to batch of rollouts
        mb_obs = np.asarray(mb_obs, dtype=self.obs.dtype)
        mb_rewards = np.asarray(mb_rewards, dtype=np.float32)
        mb_actions = np.asarray(mb_actions)
        mb_values = np.asarray(mb_values, dtype=np.float32)
        mb_neglogpacs = np.asarray(mb_neglogpacs, dtype=np.float32)  # TODO this is an array of tensors, output values instead!!
        mb_dones = np.asarray(mb_dones, dtype=np.bool)
        last_values = self.model.value(self.obs, self.states, self.dones)
        # discount/bootstrap off value fn
        mb_returns = np.zeros_like(mb_rewards)
        mb_advs = np.zeros_like(mb_rewards)
        lastgaelam = 0
        for t in reversed(range(self.nsteps)):
            if t == self.nsteps - 1:
                nextnonterminal = 1.0 - self.dones
                nextvalues = last_values
            else:
                nextnonterminal = 1.0 - mb_dones[t+1]
                nextvalues = mb_values[t+1]
            delta = mb_rewards[t] + self.gamma * nextvalues * nextnonterminal - mb_values[t]  # 1-step td-error
            mb_advs[t] = lastgaelam = delta + self.gamma * self.lam * nextnonterminal * lastgaelam
        mb_returns = mb_advs + mb_values

        self.logger.debug('Actions: %s' % mb_actions)
        self.logger.debug('Q values: %s' % mb_values)
        # self.logger.debug('Done mask: %s' % mb_masks)  # ?
        self.logger.debug('Observations: %s' % mb_obs)

        return (*map(sf01, (mb_obs, mb_returns, mb_dones, mb_actions, mb_values, mb_neglogpacs)),
            mb_states, self.reward_window, mb_rewards)

def sf01(arr):
    """
    swap and then flatten axes 0 and 1
    """
    s = arr.shape
    return arr.swapaxes(0, 1).reshape(s[0] * s[1], *s[2:])

def constfn(val):
    def f(_):
        return val
    return f

# nsteps == batchsize
def learn(policy, env, test_env, seed, total_timesteps,
          log_interval, test_interval, show_interval, logdir,
          lr, max_grad_norm, units_per_hlayer, activ_fcn,
          gamma=0.99, vf_coef=0.5, ent_coef=0.01, nsteps=5,
          lam=0.95, nminibatches=4, noptepochs=4, cliprange=0.2,
          early_stop=False, keep_model=2,
          save_model=True, restore_model=False, save_traj=False):

    total_timesteps = int(total_timesteps)

    logger = logging.getLogger(__name__)
    tf.reset_default_graph()
    set_global_seeds(seed)

    nenvs = env.num_envs
    ob_space = env.observation_space
    ac_space = env.action_space
    nbatch = nenvs * nsteps
    nbatch_train = nbatch // nminibatches  # TODO number of samples per minibatch in an optimization episode

    make_model = lambda : Model(policy=policy,
                                ob_space=ob_space,
                                ac_space=ac_space,
                                nenvs=nenvs,
                                nbatch_train=nbatch_train,
                                nsteps=nsteps,
                                ent_coef=ent_coef,
                                vf_coef=vf_coef,
                                max_grad_norm=max_grad_norm,
                                activ_fcn=activ_fcn,
                                units_per_hlayer=units_per_hlayer,
                                log_interval=log_interval,
                                logdir=logdir,
                                keep_model=keep_model,
                                lr=lr,
                                cliprange=cliprange)
    model = make_model()

    sum_write = model.get_summary_writer()
    result_path = os.path.join(logdir, 'train_results.csv')
    if save_traj:
        rew_traj = []
        rew_results_path = os.path.join(logdir, ('lr'+str(lr)+'_tracking_results.csv'))
    else:
        rew_results_path = None

    i_sample, i_train = 0, 0
    return_threshold = -2.
    horizon = 100
    avg_rm = deque(maxlen=30)

    runner = Runner(env=env, model=model, nsteps=nsteps, gamma=gamma, lam=lam, horizon=horizon, show_interval=show_interval, summary_writer=sum_write)

    if restore_model:
        for el in os.listdir(logdir):
            if 'final' in el and '.meta' in el:
                # Load pre trained model and set network parameters
                model.load(os.path.join(logdir, el[:-5]))
                # Reset global step parameter.
                model.sess.run(model.global_step.assign(0))

    logger.info('Start Training')
    breaked = False

    nupdates = total_timesteps//nbatch
    for update in range(1, nupdates+1):
        assert nbatch % nminibatches == 0  # nbatch should be a multiple of nminibatches
        obs, returns, masks, actions, values, neglogpacs, states, reward_window, rewards = \
            runner.run()  #pylint: disable=E0632  # returns are estimates of the discounted reward

        if rew_results_path is not None:
            rew_traj.append(rewards)

        nbatch_train = nbatch // nminibatches  # number of samples per minibatch
        tstart = time.time()
        # frac = 1.0 - (update - 1.0) / nupdates  # converges to 0
        # lrnow = lr(frac)  #
        # cliprangenow = cliprange(frac)  # cliprange converges to 0

        # Update step
        mblossvals = []
        if states is None: # nonrecurrent version
            inds = np.arange(nbatch)
            for _ in range(noptepochs):
                np.random.shuffle(inds)  #
                for start in range(0, nbatch, nbatch_train):
                    end = start + nbatch_train
                    mbinds = inds[start:end]
                    slices = (arr[mbinds] for arr in (obs, returns, masks, actions, values, neglogpacs))
                    # mblossvals.append(model.train(lrnow, cliprangenow, *slices))
                    mblossvals.append(model.train(*slices))
        else:  # recurrent version
            assert nenvs % nminibatches == 0
            envsperbatch = nenvs // nminibatches  # minibatch contains batch data from several envs.
            envinds = np.arange(nenvs, dtype=np.int32)
            flatinds = np.arange(nenvs * nsteps).reshape(nenvs, nsteps)
            envsperbatch = nbatch_train // nsteps
            for _ in range(noptepochs):
                np.random.shuffle(envinds)
                for start in range(0, nenvs, envsperbatch):
                    end = start + envsperbatch
                    mbenvinds = np.array(envinds[start:end])  # TODO int() does not work here. ensure that indices are integers beforehand
                    mbflatinds = flatinds[mbenvinds].ravel()
                    slices = (arr[mbflatinds] for arr in (obs, returns, masks, actions, values, neglogpacs))
                    if nenvs == 1:
                        mbstates = states[:]
                    else:
                        if type(states) == tuple or type(states) == tf.contrib.rnn.LSTMStateTuple:  # LSTM state
                            mbstates = [el[mbenvinds] for el in states]
                        else:  # GRU state
                            mbstates = states[mbenvinds]
                    # mblossvals.append(model.train(lrnow, cliprangenow, *slices, mbstates))
                    mblossvals.append(model.train(*slices, mbstates))

        if test_interval > 0 and i_train > 0 and (update % test_interval == 0):
            ep_return = model.test_run(test_env, n_eps=10, n_pipes=2000)  # TODO test, whether results.csv is saved properly
            with open(result_path, "a") as csvfile:
                writer = csv.writer(csvfile)
                ep_return[0:0] = [i_sample, i_train]
                writer.writerow(ep_return)

        if ('ContFlappyBird' in env.env_id):
            saved = False
            for i, rw in enumerate(reward_window):
                rm = sum(rw) / horizon
                if sum_write is not None:
                    s_summary = tf.Summary()
                    s_summary.value.add(tag='envs/environment%s/isample_return' % i,
                                        simple_value=rm)
                    sum_write.add_summary(s_summary, i_sample)

                    t_summary = tf.Summary()
                    t_summary.value.add(tag='envs/environment%s/itrain_return' % i,
                                        simple_value=rm)
                    sum_write.add_summary(t_summary, i_train)
                    sum_write.flush()
                if save_model and not saved and rm > return_threshold:
                    return_threshold = rm
                    logger.info('Save model at max rolling mean %s' % return_threshold)
                    model.save('inter_model')
                    saved = True
                avg_rm.append(rm)

        if early_stop:
            if (i_sample > 500000) and (i_sample <= 500000 + nbatch):  # TODO how to determine early-stopping criteria non-heuristically, but automatically? - BOHB algorithm?
                if (sum(avg_rm)/30) <= -0.88:
                    print('breaked')
                    breaked = True
                    break
        i_sample += nbatch
        i_train += 1

    if save_model:
        model.save('final_model')
        logger.info('Finished Training. Saving Final model.')

    if rew_results_path is not None:
        with open(rew_results_path, "a") as csvfile:
            writer = csv.writer(csvfile)
            traj = np.asanyarray(rew_traj).reshape(-1).tolist()
            traj[0:0] = [np.mean(traj)]  # i_train, i_sample
            writer.writerow(traj)

    logger.info('*******************************************************')
    logger.info('Total number of interactions with the environment: %s' % i_sample)
    logger.info('Total number of finished episodes during training: sum(%s) = %s' % (runner.ep_idx, sum(runner.ep_idx)))
    logger.info('Total number of parameter updates during training: %s' % i_train)
    logger.info('*******************************************************\n')

    return breaked

def safemean(xs):
    return np.nan if len(xs) == 0 else np.mean(xs)

from run_ple_utils import make_ple_envs, make_ple_env
from models import MLPPolicy, LSTMPolicy, GRUPolicy
if __name__ == '__main__':
    seed = 1
    env = make_ple_envs('ContFlappyBird-hNS-nrf0-train-v0', num_env=1, seed=seed)
    test_env = make_ple_env('ContFlappyBird-v3', seed=seed)

    logger = logging.getLogger()
    ch = logging.StreamHandler()  # Handler which writes to stderr (in red)
    ch.setLevel(logging.INFO)
    ch.setFormatter(logging.Formatter('%(levelname)s:%(name)s: %(message)s'))
    logger.addHandler(ch)
    logger.setLevel(logging.INFO)

    logdir = '/home/mara/Desktop/logs/LSTM_PPO'
    logdir = os.path.join(logdir, 'ppo_output1')

    learn(policy=LSTMPolicy,
          env=env,
          test_env=test_env,
          seed=seed,
          total_timesteps=10000,  # int(1e7),
          log_interval=0,
          test_interval=0,
          show_interval=0,
          logdir=logdir,
          lr=5e-4,
          # lrschedule='constant',
          max_grad_norm=0.01,
          units_per_hlayer=(24, 24, 24),
          activ_fcn='relu6',
          gamma=0.9,
          vf_coef=0.3,
          ent_coef=0.00005,
          nsteps=64,
          lam=0.95,
          nminibatches=1,
          noptepochs=1,
          cliprange=0.2,
          early_stop=False,
          keep_model=2,
          restore_model=False)
    env.close()


#! /usr/bin/env python
# coding=utf-8
""" Find best PSF stars using GA to minimize mean error.

    .. seealso:: :ref:`gapick`
"""
from __future__ import absolute_import, division, print_function
__metaclass__ = type

# For now:
# TODO: --neigbours option: once for generation neightbour removal
# TODO: --timestamp: add current timestamp to directory name (+)
# TODO: conf option for
# TODO: warnings about: daophot.opt, allstar.opt
# TODO: single parameter --aperture, instead of photo.opt
# TODO: write down text based statistics from logbook.pkl
# TODO: some docs, and link to them in --help
# TODO: allstar.opt missing in result dir!
# For later, on request:
# TODO: Option for provide own PSF stars set, used for comaprision, including in candidates, put in initial population
# TODO: Checkpoints - continue calculation

import os
import logging
import random
import pickle
import time
from datetime import timedelta
from copy import deepcopy

try:
    from itertools import izip_longest as zip_longest
except ImportError:
    from itertools import zip_longest

import numpy
from bitarray import bitarray
from scipy.stats import sigmaclip
from deap import base
from deap import creator
from deap import tools

import astwro.tools.__commons as commons
import astwro.starlist as sl
import astwro.pydaophot as dao
import astwro.tools
import astwro.utils as utils
from astwro.phot import PhotError

_time_format = '%a %H:%M:%S'

# Definitions for genetic algorithm:
# 'individual' - subset of candidates competing with other subsets to be the best in fitness
# 'genome' - the individual's definition with subset of candidates stars, is represented by binary string
#     of length equal to number of candidate stars. '1' means that star is in subset.
# 'fitness' - minimized function on individual, mean of errors form allstar if individual's stars are
#     taken as PSF stars
# 'population' - set of all individuals in some iteration

# 2.1 Definition of routines used by algorithms: initializations, scoring
def select_stars(starlist, genome):
    # type: (sl.StarList, bitarray) -> sl.StarList
    # Select stars present in genome
    return starlist[genome.tolist()]

def random_genome(ind_class, len, prob_limits):
    # type: (type, int, [float, float]) -> bitarray
    # Create random genome of length `len` in which probability of '1' on any position is from `prob_limits` range.
    prob = random.uniform(prob_limits[0], prob_limits[1])
    return ind_class([random.random() <= prob for _ in range(len)])

def clone_individual(individual):
    # type: (bitarray) -> bitarray
    # Make a deepcopy-clone - deepcopy genome bitarray and associated fitness
    n = deepcopy(individual)
    n.fitness = deepcopy(individual.fitness)
    return n

def calc_spectrum(pop):
    # calculates 'spectrogram' of population
    # which is a statistic of stars occurrences in population individuals
    spec = numpy.zeros(len(pop[0]))
    for ind in pop:
        spec += ind.tolist()
    return spec

# === astwro version <= 0.7.0
# def fitness_for_als(als):
#     # type: (sl.StarList) -> (float,)
#     # Calucalates fitness from allstar result
#     clipped = sigmaclip(als.chi)[0]
#     return numpy.sqrt((clipped * clipped).sum()),  # fitness is tuple (val,)

# === astwro version > 0.7.0
# def fitness_for_als(als):
#     # type: (sl.StarList) -> (float,)
#     # Calucalates fitness from allstar result
#     w = 100*(-als.mag/5)
#
#     return (als.chi*w).sum()/w.sum(),  # fitness is tuple (val,)

# === astwro version > 0.7.2
def fitness_for_als(als):
    # type: (sl.StarList) -> (float,)
    # Calucalates fitness from allstar result
    w = 100*(-als.mag/5)

    return (als.mag_err*w).sum()/w.sum(),  # fitness is tuple (val,)

def eval_population(population, candidates, workers, show_progress, fine_tune):
    if fine_tune:
        return eval_population_fine_psf(population, candidates, workers, show_progress)
    else:
        return eval_population_simple(population, candidates, workers, show_progress)

def eval_population_fine_psf(population, candidates, workers, show_progress):
    # type: (list(bitarray), sl.StarList, list(dict), bool) -> list
    # Evaluates fitness for all individual in population.

    # This version uses sofisticated process from daophot_bialkow
    # Uses daophots and allstars processes form `workers` running them in parallel.
    # :return: list fitnesses (1-element couples as `deap` lib likes)

    progress = None
    if show_progress:
        progress = utils.progressbar(total=len(population), step=len(workers))
        progress.print_progress(0)
    fitnesses = []
    f_max = None
    # https://docs.python.org/3/library/itertools.html#itertools-recipes grouper()
    # grouping population into chunks of size `parallel`
    for chunk in zip_longest(*([iter(population)] * len(workers))):
        active = [g is not None for g in chunk]  # all active, on errors some could be deactivated
        # PSF
        for individual, worker in zip(chunk, workers):
            if individual:
                pdf_s = select_stars(candidates, individual)
                worker['daophot'].write_starlist(pdf_s, 'i.lst')
                worker['daophot'].PSf(psf_stars='i.lst')  # add to queue only - batch mode
                worker['daophot'].run(wait=False)  # parallel start all workers without waiting
        # wait for PSF and start ALLSTAR nei
        for i, worker in enumerate(workers):
            if active[i]:
                worker['daophot'].wait_for_results()  # now wait before using results
                active[i] = worker['daophot'].PSf_result.converged  # PSF is not always successful
                if active[i]:
                    worker['allstar'].ALlstar(stars='i.nei')  # enqueue calculation
                    worker['allstar'].run(wait=False)  # asynchronous / parallel
        # second PSF
        for i, worker in enumerate(workers):
            if active[i]:
                worker['allstar'].wait_for_results()
                if worker['allstar'].ALlstars_result.success:
                    worker['daophot'].SUbstar(subtract='i.als', leave_in='i.lst')
                    worker['daophot'].run()  # quick run
                    worker['daophot'].ATtach('is')
                    worker['daophot'].PSf(photometry='i.als', psf_stars='i.lst')
                    worker['daophot'].run(wait=False)
                else:
                    active[i] = False
        # second allstar
        for i, worker in enumerate(workers):
            if active[i]:
                worker['daophot'].wait_for_results()
                if worker['daophot'].PSf_result.success:
                    worker['allstar'].ALlstar(stars='i.nei')  # enqueue calculation
                    worker['allstar'].run(wait=False)  # asynchronous / parallel
                else:
                    active[i] = False
        # third PSF
        for i, worker in enumerate(workers):
            if active[i]:
                worker['allstar'].wait_for_results()
                if worker['allstar'].ALlstars_result.success:
                    worker['daophot'].SUbstar(subtract='i.als', leave_in='i.lst')
                    worker['daophot'].run()  # quick run
                    worker['daophot'].ATtach('is')
                    worker['daophot'].PSf(photometry='i.als', psf_stars='i.lst')
                    worker['daophot'].run(wait=False)
                else:
                    active[i] = False
        # final allstar
        for i, worker in enumerate(workers):
            if active[i]:
                worker['daophot'].wait_for_results()
                if worker['daophot'].PSf_result.success:
                    worker['allstar'].ALlstar(stars='als.ap')  # enqueue calculation
                    worker['allstar'].run(wait=False)  # asynchronous / parallel
                else:
                    active[i] = False

        # wait for allstar for workers and process results
        for i, worker in enumerate(workers):
            if active[i]:
                worker['allstar'].wait_for_results()
                all_s = worker['allstar'].ALlstars_result.als_stars
                f = fitness_for_als(all_s)
                fitnesses.append(f)  # fitness is tuple (val,)
                f_max = f if f_max is None or f_max[0] < f[0] else f_max
            #                f_max = f[0] if f_max is None else max(f_max, f[0])
            else:
                fitnesses.append(None)
        # cut fitnesses if longer than population
        # (finesses mod parallel == 0, some None at the end can appear if population mod parallel != 0)
        fitnesses = fitnesses[:len(population)]
        # fill gaps in fitnesses with maximum of rest of population
        for i, f in enumerate(fitnesses):
            if f is None:
                fitnesses[i] = f_max
        if progress:
            progress.print_progress()

    return fitnesses

def eval_population_simple(population, candidates, workers, show_progress):
    # type: (list(bitarray), sl.StarList, list(dict), bool) -> list
    # Evaluates fitness for all individual in population.
    # Uses daophots and allstars processes form `workers` running them in parallel.
    # :return: list fitnesses (1-element couples as `deap` lib likes)

    progress = None
    if show_progress:
        progress = utils.progressbar(total=len(population), step=len(workers))
        progress.print_progress(0)
    fitnesses = []
    f_max = None
    # https://docs.python.org/3/library/itertools.html#itertools-recipes grouper()
    # grouping population into chunks of size `parallel`
    for chunk in zip_longest(*([iter(population)] * len(workers))):
        active = [g is not None for g in chunk]  # all active, on errors some could be deactivated
        # PSF
        for individual, worker in zip(chunk, workers):
            if individual:
                pdf_s = select_stars(candidates, individual)
                worker['daophot'].PSf(psf_stars=pdf_s)  # add to queue only - batch mode
                worker['daophot'].run(wait=False)  # parallel start all workers without waiting
        # wait for PSF and start ALLSTAR
        for i, worker in enumerate(workers):
            if active[i]:
                worker['daophot'].wait_for_results()  # now wait before using results
                active[i] = worker['daophot'].PSf_result.converged  # PSF is not always successful
                if active[i]:
                    worker['allstar'].ALlstar(stars='als.ap')  # enqueue calculation
                    worker['allstar'].run(wait=False)  # asynchronous / parallel
        # wait for allstar for workers and process results
        for i, worker in enumerate(workers):
            if active[i]:
                worker['allstar'].wait_for_results()
                all_s = worker['allstar'].ALlstars_result.als_stars
                f = fitness_for_als(all_s)
                fitnesses.append(f)  # fitness is tuple (val,)
                f_max = f if f_max is None or f_max[0] < f[0] else f_max
            #                f_max = max(f_max, f)
            else:
                fitnesses.append(None)
        # cut fitnesses if longer than population
        # (finesses mod parallel == 0, some None at the end can appear if population mod parallel != 0)
        fitnesses = fitnesses[:len(population)]
        # fill gaps in fitnesses by maximum of rest of population
        for i, f in enumerate(fitnesses):
            if f is None:
                fitnesses[i] = f_max
        if progress:
            progress.print_progress()

    return fitnesses

def _prepare_output_dir(outdir, overwrite, srcdir, arg):
    # type: (str, bool, str, object) -> (utils.CycleFile, utils.CycleFile, utils.CycleFile, str)
    # Prepare output directory for results
    if outdir:
        from shutil import copytree, rmtree

        outdir = os.path.abspath(os.path.expanduser(outdir))
        if os.path.exists(outdir):
            if overwrite:
                rmtree(outdir)
            else:
                logging.error('--out-dir:{} already exists and no --overwrite requested.'.format(outdir))
                raise Exception('Output directory {} exists.'.format(outdir))
        copytree(srcdir, outdir)
        logging.info('Results dir created: {}'.format(outdir))
        # todo allstar.opt missing in result dir!
        basename = 'gen'
        lst_file = utils.cyclefile(outdir, basename, '.lst')
        reg_file = utils.cyclefile(outdir, basename, '.reg')
        gen_file = utils.cyclefile(outdir, basename, '.gen')
        # write down parameters
        with open(os.path.join(outdir, 'about.txt'), 'w') as f:
            print("astwro pydaophot/tools version: {}/{}\n".format(dao.__version__, astwro.tools.__version__), file=f)
            for k, v in arg.__dict__.items():
                print('{}\t= {}'.format(k, v), file=f)
        return lst_file, reg_file, gen_file, outdir
    else:
        return None, None, None, None

def _plot_err(pe, pe_psf, psf, fig_file):
    # type: (PhotError) -> None
    try:
        import matplotlib
        matplotlib.use('pdf')
        import matplotlib.pyplot as plt
    except ImportError:
        logging.warning("'matplotlib' cannot be imported. No control diagrams will be generated")
    except:
        logging.warning('matplotlib PDF output cannot be initialized. No control diagrams will be generated')
        return
    else:
        try:
            plt.rc('text', usetex=True)
        except:
            pass

        ax = plt.subplot()
        ls = numpy.linspace(pe.mag.min(), pe.mag.max())
        ax.plot(pe.mag[~pe.outlayers_mask], pe.err[~pe.outlayers_mask], '.', alpha=0.7, ms=2,
                label='training set')
        ax.errorbar(pe.mag_means, pe.err_means, yerr=pe.err_means_stderr * 1.0, fmt='D', markerfacecolor='none', alpha=0.7, lw=1,
                    label='error bins')
        ax.plot(ls, pe.evaluate_fit(ls), lw=1,
                label='error fit')
        ax.plot(ls, pe.evaluate_fit_plus_sigmas(ls),  '--', lw=1,
                label='error$+{:.1f}\\sigma$ fit (training limit)'.format(pe.fitplussigma))
        ax.plot(ls, pe_psf.evaluate_fit_plus_sigmas(ls),  '--', lw=1,
                label='error$+{:.1f}\\sigma$ fit (PSF limit)'.format(pe_psf.fitplussigma))
        ax.plot(pe.mag[pe.outlayers_mask], pe.err[pe.outlayers_mask], '.', alpha=0.7, ms=3,
                label='outlayers')
        ax.plot(psf.mag, psf.mag_err, '.', alpha=0.7, ms=3,
                label='PSF candidates')
        ax.set_xlabel('mag'); ax.set_ylabel('mag error')
        ax.legend()
        ax.figure.savefig(fig_file, bbox_inches='tight')
        ax.set_ylim(0,0.2)
        ax.figure.savefig('{}_zoomed{}'.format(*os.path.splitext(fig_file)), bbox_inches='tight')
        logging.info('Control plot mag/err {} generated'.format(fig_file))

def __do(arg):
    # Main routine, common for command line, and python scripts call
    # :type arg: Namespace

    start_time = time.time()

    # Configure progress logging
    clogger = logging.getLogger('clean logger')  # create another logger for stats without prefixes (clean)
    chandler = logging.StreamHandler()
    chandler.setFormatter(logging.Formatter('%(message)s'))
    clogger.propagate = False
    clogger.addHandler(chandler)

    check_arguments(arg)

    # image_file
    if arg.image is None:
        from astwro.sampledata import fits_image
        arg.image = fits_image()  # sample image

        logging.warning('No image file argument provided (-h for help), using demonstration image: %s', arg.image)

    # expand patches of file parameters
    arg.all_stars_file = dao.Daophot.expand_path(arg.all_stars_file)
    arg.psf_stars_file = dao.Daophot.expand_path(arg.psf_stars_file)
    arg.photo_opt = dao.Daophot.expand_path(arg.photo_opt)
    arg.daophot_opt = dao.Daophot.expand_path(arg.daophot_opt)

    # get single daophot and ATtach file
    dp = dao.Daophot(image=arg.image, daophotopt=arg.daophot_opt)
    al = dao.Allstar(image=arg.image, dir=dp.dir)

    # all stars file
    if not arg.all_stars_file:
        logging.warning('No all-stars-file provided, Stars will be found by daophot FIND (frames av/sum: %d/%d)',
                        arg.frames_av, arg.frames_sum)
        find = dp.FInd(arg.frames_av, arg.frames_sum)
        logging.info('{} stars found by FIND, sky estimation: {}, err/dev: {}/{}'.format(find.stars, find.sky, find.err,
                                                                                find.skydev))
        arg.all_stars_file = find.starlist_file

    # photometry
    if not arg.photo_opt:  # no file no problem, but recreate default radius if not provided
        if arg.photo_is == 0:
            arg.photo_is = 35
        if arg.photo_os == 0:
            arg.photo_os = 50
        if not arg.photo_ap:
            arg.photo_ap = [8]
    logging.info('Performing preliminary {} pixels aperture  photometry'.format(arg.photo_ap[0]))
    photometry = dp.PHotometry(photoopt=arg.photo_opt,
                               IS=arg.photo_is,
                               OS=arg.photo_os,
                               apertures=arg.photo_ap,
                               stars=arg.all_stars_file)

    # pick PFS stars candidates
    PICKcandidates = not arg.psf_stars_file
    if PICKcandidates:
        pick = dp.PIck(arg.stars_to_pick, arg.faintest_to_pick)
        arg.psf_stars_file = pick.picked_stars_file

    # psf (for errors collection)
    PSf_result = dp.PSf(psf_stars=arg.psf_stars_file)

    # all stars (filtering by photometry errors and magnitudes)
    stars = photometry.photometry_starlist
    count0 = stars.shape[0]
    stars = stars[stars.mag < arg.max_ph_mag]
    count1 = stars.shape[0]
    stars = stars[stars.mag_err < arg.max_ph_err]
    count2 = stars.shape[0]
    logging.info('{} (of {}) stars left after filtering against {:.2f} photometry magnitude threshold, '
                 .format(count1, count0, arg.max_ph_mag))
    logging.info('{}  left after {:.2f} aperture photometry error threshold'
                 .format(count2, arg.max_ph_err))

    logging.info('Performing preliminary PSF profile photometry')
    ALlstar_result = al.ALlstar(stars=stars)
    logging.info('{} stars converged in PSF profile photometry'
                 .format(ALlstar_result.stars_no[0]))
    logging.info('Fitting PSF profile photometry errors')
    pe = PhotError(ALlstar_result.als_stars.mag, ALlstar_result.als_stars.mag_err, fitplussigma=arg.clip_fit_sigmas)
    logging.info('{} stars left, after excluding profile photometry error {} sigma outliers'
                 .format((~pe.outlayers_mask).sum(), arg.clip_fit_sigmas))
    stars = ALlstar_result.als_stars[~pe.outlayers_mask]

    # for the sake of optimisation we write starlist to file to avoid multiple serialization of this list
    # (instead of more obvious providing starlist as the allstar argument)
    dp.write_starlist(stars, 'als.ap', sl.DAO.AP_FILE)

    # candidates (filter out big psf errors)
    pe2 = None
    if PICKcandidates: # do it again
        if arg.clip_fit_sigmas > arg.clip_psf_fit_sigmas: # tighter filtering for PSF candids
            logging.info('Excluding {} sigma outlayers for PICK'.format(arg.clip_psf_fit_sigmas))
            pe2 = PhotError(ALlstar_result.als_stars.mag, ALlstar_result.als_stars.mag_err,
                           fitplussigma=arg.clip_psf_fit_sigmas)
            cands_for_cands = ALlstar_result.als_stars[~pe2.outlayers_mask]
        else:
            cands_for_cands = stars
        pick = dp.PIck(arg.stars_to_pick, arg.faintest_to_pick, photometry=cands_for_cands)
        logging.info('{} PSF stars candidates selected by PICK'.format(pick.stars))
        arg.psf_stars_file = pick.picked_stars_file
        PSf_result = dp.PSf(psf_stars=arg.psf_stars_file) # collect PSF errors
    candidates = dp.read_starlist(arg.psf_stars_file, add_psf_errors=True)
    org_cand_no = candidates.stars_number()
    err = PSf_result.errors
    averr = err.psf_err.mean()
    err = err[(err.psf_err < arg.max_psf_err_mult * averr) & (
                err.flag == ' ')]  # filter out big errors and * or ? marked stars
    candidates = candidates.loc[err.index]

    # candidates = candidates[candidates.psf_err < arg.max_psf_err]  # old filter, new above
    logging.info(
        "{} good candidates ({} rejected: * or ? or psf error exceeded max-psf-err-mult*averge-error "
        "= {:.2f}*{:.2f} = {:.2f})"
            .format(
            candidates.stars_number(),
            org_cand_no - candidates.stars_number(),
            arg.max_psf_err_mult,
            averr,
            arg.max_psf_err_mult * averr)
    )
    if arg.include_psf:
        logging.info('PSF star candidates stays in training set (--include-psf)')
    else:
        overlapping = stars.index.isin(candidates.index)
        stars = stars[~overlapping]
        logging.info('{} stars left in training set after exclusion PSF star candidates (no --include-psf)'
                     .format(stars.stars_number()))

    if candidates.stars_number() < 15:
        logging.error("Number of candidates lass than 15. GA needs more. Sorry...")

    if stars.stars_number() < 15:
        logging.error("Number stars in training set lass than 15. GA needs more. Sorry...")

    # Prepare output directory for results
    lst_file, reg_file, gen_file, result_dir = _prepare_output_dir(arg.out_dir, arg.overwrite, str(dp.dir), arg)
    _plot_err(pe, pe2, candidates, os.path.join(arg.out_dir, 'prel_psf_filtering.pdf'))

    #  From all candidates find best subset, where best means minimizing mean of errors form allstar
    #       using genetic algorithm
    #  For details about setting up genetic algorithm with DEAP visit:
    #       http://deap.gel.ulaval.ca/doc/default/examples/ga_onemax.html

    creator.create("FitnessMax", base.Fitness, weights=(-1.0,))
    creator.create("Individual", bitarray, fitness=creator.FitnessMax)

    toolbox = base.Toolbox()
    # Structure initializes
    toolbox.register("individual", random_genome, creator.Individual, candidates.stars_number(), arg.ga_init_prob)
    toolbox.register("population", tools.initRepeat, list, toolbox.individual)
    toolbox.register('clone', clone_individual)

    # The Genetic Operators

    # set min_stars to all_cand_no*ga_init_prob/2
    toolbox.register('mate', tools.cxTwoPoint)
    toolbox.register('mutate', tools.mutFlipBit, indpb=arg.ga_mut_str)
    toolbox.register('select', tools.selTournament, tournsize=3)

    # setup stats
    stats_fits = tools.Statistics(key=lambda ind: ind.fitness.values)
    stats_star = tools.Statistics(key=sum)  # number of stars is an sum of bitarray: [001101010001] has 5 stars
    stats = tools.MultiStatistics(fitness=stats_fits, size=stats_star)
    stats.register('avg', numpy.mean)
    stats.register('std', numpy.std)
    stats.register('min', numpy.min)
    stats.register('max', numpy.max)

    # Initiate workers. Each worker has Daophot and Allstar objects sharing runner directory,
    # working in batch mode.
    workers = []
    workers_logger = logging.getLogger('worker')
    workers_logger.setLevel('ERROR')  # prevent workers flood output with logrecords
    for i in range(arg.parallel):
        d = dp.clone()  # clone previously used daophot
        d.batch_mode = True
        a = dao.Allstar(dir=d.dir, image=d.image, batch=True, options={'MA': 100})
        d.logger = workers_logger
        a.logger = workers_logger
        workers.append({'daophot': d, 'allstar': a})

    # Setup initial population, HoF and logbook and  or load it from checkpoint when continuing previous calculation
    start_gen = 0

    #    if arg.checkpoint:   ## Not implemented
    if False:
        with open(os.path.expanduser(arg.checkpoint), "rb") as f:
            checkpoint = pickle.load(f)
        pop = checkpoint['population']
        start_gen = checkpoint['generation']
        hof = checkpoint['halloffame']
        logbook = checkpoint['logbook']
        logging.info('Restoring genetic algorithm on {} of {} generations'.format(start_gen, arg.ga_max_iter))
    else:
        hof = tools.HallOfFame(maxsize=10)

        logbook = tools.Logbook()
        logbook.header = 'gen', 'fitness', 'size'
        logbook.chapters['fitness'].header = 'min', 'avg', 'max', 'std'
        logbook.chapters['size'].header = 'min', 'avg', 'max'

        pop = toolbox.population(n=arg.ga_pop)
        logging.info('Starting genetic algorithm for {} generations at {}'.format(
            arg.ga_max_iter,
            time.strftime(_time_format, time.localtime())
        ))
        logging.info('{} parallel threads, ETA will be calculated after generation 1'.format(arg.parallel))

    # Calculate fitnesses of initial population
    fitnesses = eval_population(pop, candidates, workers, show_progress=not arg.no_progress, fine_tune=arg.fine)
    for ind, fit in zip(pop, fitnesses):
        ind.fitness.values = fit

    record = stats.compile(pop)
    logbook.record(gen=0, spectrum=calc_spectrum(pop), **record)
    clogger.info('{}\t ETA: [... to be determined]'.format(logbook.stream))

    evolution_start_time = time.time()

    # Begin the evolution
    for g in range(start_gen + 1, arg.ga_max_iter):
        #  New Generation
        #  select the next generation individuals
        offspring = toolbox.select(pop, len(pop))
        # Clone the selected individuals
        offspring = list(map(toolbox.clone, offspring))
        # Apply crossover and mutation on the offspring
        for child1, child2 in zip(offspring[::2], offspring[1::2]):
            if random.random() < arg.ga_cross_prob:
                toolbox.mate(child1, child2)
                del child1.fitness.values
                del child2.fitness.values
        for mutant in offspring:
            if random.random() < arg.ga_mut_prob:
                toolbox.mutate(mutant)
                del mutant.fitness.values

        # calculate fitnesses of new individuals
        invalid_ind = [ind for ind in offspring if not ind.fitness.valid]
        fitnesses = eval_population(invalid_ind, candidates, workers, show_progress=not arg.no_progress,
                                    fine_tune=arg.fine)
        for ind, fit in zip(invalid_ind, fitnesses):
            ind.fitness.values = fit
        # New population from offspring
        pop[:] = offspring

        # Stats
        # hof.update(pop)  # not implemented yet, __deapcopy__ of the Individual should work first
        ETA = time.strftime(_time_format, time.localtime(
            evolution_start_time + (time.time() - evolution_start_time) * arg.ga_max_iter / g))
        record = stats.compile(pop)
        logbook.record(gen=g, spectrum=calc_spectrum(pop), **record)
        clogger.info('{}\t ETA: {}'.format(logbook.stream, ETA))

        # For every generation create lst file and ds9 reg file of best and point symlinks to last generation
        if result_dir:
            best_ind = tools.selBest(pop, 1)[0]
            best_stars = select_stars(candidates, best_ind)
            lst_file.next_file(g)
            reg_file.next_file(g)
            gen_file.next_file(g)
            sl.write_dao_file(best_stars, lst_file.file, sl.DAO.LST_FILE)
            sl.write_ds9_regions(best_stars, reg_file.file)
            for ind in pop:
                gen_file.file.write(ind.to01() + '\n')
            with open(os.path.join(result_dir, 'logbook.pkl'), 'wb') as f:
                pickle.dump(logbook, f)
            checkpoint = dict(population=pop, generation=g, halloffame=hof, logbook=logbook)
            with open(os.path.join(result_dir, 'checkpoint.chk'), 'wb') as f:
                pickle.dump(checkpoint, f)

        # end of evolution loop

    if result_dir:
        lst_file.close()
        reg_file.close()
        gen_file.close()

    logging.info('Successful evolution finished at {} (elapsed time: {:s})'.format(
        time.strftime(_time_format, time.localtime()),
        str(timedelta(seconds=time.time() - start_time))
    ))

    best_ind = tools.selBest(pop, 1)[0]
    logging.info('Best individual is {}, {}'.format(best_ind, best_ind.fitness.values))

    best_stars = select_stars(candidates, best_ind)
    return best_stars

__arg_parser_singlethon = None

def __arg_parser():
    global __arg_parser_singlethon
    if __arg_parser_singlethon:
        return __arg_parser_singlethon
    import argparse
    parser = argparse.ArgumentParser(
        description='Find best PSF stars using genetic algorithm running the daophot/allstar'
                    ' workflow for number of a PSF stars sets.'
                    ' Minimized function'
                    ' is the mean of allstar\'s chi statistic, weighted by flux, calculated on '
                    ' training subset of of all stars. Results'
                    ' are stored in --dir directory if provided. List of star ids'
                    ' of the best set found, is printed to stdout (until suppressed by -no_stdout)',
        epilog=commons.version_string(),
    )
    parser.add_argument('image', default=None, nargs='?',
                        help='FITS image file (obligatory if no --demo)')
    g_dao = parser.add_argument_group('DAOPHOT/ALLSTARS related parameters',
                                      'Those parameters are used by DAOPHOT/ALLSTARS workflow. gapick optimizes PSF '
                                      'fitting errors of that workflow applied to training set stars as a function'
                                      'of \'PSF stars\' sets - stars used to derive PSF function')
    g_dao.add_argument('-c', '--all-stars-file', metavar='FILE', default=None,
                       help='all stars input file in one of daophot\'s formats (default: obtained by daophot FIND)')
    g_dao.add_argument('-D', '--daophot-opt', metavar='FILE', default=None,
                       help='daophot.opt file for daophot (default: internal)')
    g_dao.add_argument('--frames-av', metavar='n', type=int, default=1,
                       help='frames ave - parameter of daophot FIND when --all-stars-file not provided  (default: 1)')
    g_dao.add_argument('--frames-sum', metavar='n', type=int, default=1,
                       help='frames summed - parameter of daophot FIND when --all-stars-file not provided (default: 1)')
    g_dao.add_argument('-O', '--photo-opt', metavar='FILE', default=None,
                       help='photo.opt file for aperture photometry (default: none)')
    g_dao.add_argument('--photo-is', metavar='r', type=int, default=0,
                       help='PHOTOMETRY inner sky radius, overwrites photo.opt, (default: from --photo-opt or 35)')
    g_dao.add_argument('--photo-os', metavar='r', type=int, default=0,
                       help='PHOTOMETRY outher sky radius, overwrites photo.opt, (default: from --photo-opt or 50)')
    g_dao.add_argument('--photo-ap', metavar='r', type=int, default=[], nargs='+',
                       help='PHOTOMETRY apertures radius (up to 12), overwrites photo.opt, '
                            '(default: from --photo-opt or 8)')
    g_dao.add_argument('-l', '--psf-stars-file', metavar='FILE', default=None,
                       help='PSF candidates input file in one of daophot\'s formats, the result of algorithm is '
                            'a subset of those stars (default: obtained by daophot PICK)')
    g_dao.add_argument('-P', '--stars-to-pick', metavar='n', type=int, default=100,
                       help='number of stars to PICK as candidates when --stars-to-pick not provided (default: 100)')
    g_dao.add_argument('--faintest-to-pick', metavar='MAG', type=int, default=20,
                       help='faintest magnitude to PICK as candidates when --stars-to-pick not provided (default: 20)')
    g_dao.add_argument('-f', '--fine', action='store_true',
                       help='fine tuned PSF calculation (3 iter) for crowded fields, without this option no neighbours'
                            'subtraction will be performed (slower, usually not needed)')
    g_flt = parser.add_argument_group('Filtering training set stars',
                                      'Options controls which stars will be included in a training set -- whose '
                                      'residuals after the PSF fitting will be minimized by genetic algorithm')
    g_flt.add_argument('-x', '--include-psf', action='store_true',
                       help='NEW | include candidates of PSF stars in training set; by default PSF candidates are'
                            'excluded from calculation of minimized PSF error; may be included if there is low number'
                            'of stars')
    g_flt.add_argument('--max-ph-err', metavar='x', type=float, default=0.1,
                       help='threshold of aperture photometry error for training set stars; '
                            'stars for which aperture photometry (daophot PHOTO) error is greater than x '
                            'will be excluded and have no effect on minimization '
                            '(default 0.1)')
    g_flt.add_argument('--max-ph-mag', metavar='m', type=float, default=20,
                       help='threshold of aperture photometry magnitude of stars for for training set stars; '
                            '(analogous to --max-ph-err)stars for which aperture photometry (daophot PHOTO) magnitude '
                            'is greater than m (fainter than m) will be excluded form allstar run and have no effect '
                            'on quality measurement (default 20)')
    g_flt.add_argument('--clip-fit-sigmas', metavar='x', type=float, default=3.0,
                       help='NEW | threshold of errors from the preliminary PSF fitting (allstars) in sigmas above '
                            'mean error level for specific magnitude (a mean error level for the magnitude is fitted) '
                            'stars with preliminary PSF fitting error greater than mean_err(mag)+x*sigma will '
                            'be rejected from training set (default 3.0)')
    g_psf = parser.add_argument_group('Filtering candidates for PSF  stars',
                                      'Options controls which stars can be considered as PSF candidates. '
                                      'Filtering is applied to stars provided in --psf-stars-fil or found by '
                                      'daophot/PICK')
    g_psf.add_argument('--max-psf-err-mult', metavar='x', type=float, default=3.0,
                       help='threshold of PSF errors of PSF candidates - as multiplier of average PSF error; '
                            'candidates with PSF error (from daophot PSF command) greater than x*av_err will '
                            'be rejected (as well as other marked ? or * by daophot/PSF command)'
                            '(default 3.0)')
    g_psf.add_argument('--clip-psf-fit-sigmas', metavar='x', type=float, default=2.0,
                       help='NEW | threshold of errors from the preliminary PSF fitting (allstars) in sigmas above '
                            'mean error level for specific magnitude (a mean error level for the magnitude is fitted) '
                            'candidates with preliminary PSF fitting error greater than mean_err(mag)+x*sigma will '
                            'be rejected (default 2.0)')
    g_ga = parser.add_argument_group('Genetic algorithm parameters')
    g_ga.add_argument('-I', '--ga-init-prob', metavar='x', default=[0.3, 0.8], type=float, nargs='+',
                      help='what portion of candidates is used to initialize GA individuals;'
                           ' e.g. if there is 100 candidates, each of them will be '
                           ' chosen to initialize individual genome with probability x; '
                           ' in other words if x=0.3 first population in GA will contain'
                           ' individuals with around 30 stars each; '
                           ' two numbers (e.g. -I 0.3 0.8) indicates range: for each initial genome'
                           ' random number from that range will be taken (spreading number of initial stars)'
                           ' (default: 0.3 0.8)')
    g_ga.add_argument('-i', '--ga-max-iter', metavar='n', default=50, type=int,
                      help='maximum number of iterations of generic algorithm - generations (default: 50)')
    g_ga.add_argument('-n', '--ga-pop', metavar='n', default=80, type=int,
                      help='population size of GA (default: 80)')
    g_ga.add_argument('--ga-cross-prob', metavar='x', default=0.5, type=float,
                      help='crossover probability of GA (default: 0.5)')
    g_ga.add_argument('--ga-mut-prob', metavar='x', default=0.2, type=float,
                      help='mutation probability of GA - probability to became a mutant (default: 0.2)')
    g_ga.add_argument('--ga-mut-str', metavar='x', default=0.05, type=float,
                      help='mutation strength of GA - probability of every bit flip in mutant (default: 0.05)')
    g_cnt = parser.add_argument_group('gapick run control parameters')
    g_cnt.add_argument('-p', '--parallel', metavar='n', type=int, default=8,
                       help='how many parallel processes can be forked; '
                            'n=1 avoids parallelism (default: 8)')
    g_cnt.add_argument('-d', '--out-dir', metavar='output_dir', type=str, default='RESULTS',
                       help='output directory; directory will be created and result files will be stored there;'
                            ' directory should not exist or --overwrite flag should be set'
                            ' (default: RESULTS)')
    g_cnt.add_argument('-o', '--overwrite', action='store_true',
                       help='if directory specified by --out_dir parameter exists, '
                            'then ALL its content WILL BE DELETED')
    # g_cnt.add_argument('-C', '--checkpoint', metavar='file.chk', type=str, default=None,
    #                     help='restore evaluation from checkpoint; algorithm saves checkpoint.chk file every generation,'
    #                          ' which allows resuming evolution, even with another parameters')
    g_cnt.add_argument('-L', '--loglevel', metavar='level', default='info',
                       help='logging level: debug, info, warning, error, critical (default: info)')
    g_cnt.add_argument('-t', '--no-stdout', action='store_true',
                       help='suppress printing result (list of best choice of PSF stars) to stdout at finish')
    g_cnt.add_argument('-b', '--no-progress', action='store_true',
                       help='suppress showing progress bar')
    g_cnt.add_argument('-v', '--version', action='store_true',
                       help='show version and exit')
    g_cnt.add_argument('--demo', action='store_true',
                       help='use internal demo FITS file')
    __arg_parser_singlethon = parser
    return parser

def main(**kwargs):
    """Entry point for python script calls. Parameters identical to command line"""

    args = commons.bunch_kwargs(__arg_parser(), **kwargs)
    # call main routine - common form command line and python calls
    return __do(args)

def check_arguments(arg, parser=None):
    if len(arg.ga_init_prob) > 2:
        exit_error('--ga-init-prob parameter be one or two numbers', parser)
    if len(arg.ga_init_prob) == 1:
        arg.ga_init_prob = [arg.ga_init_prob[0], arg.ga_init_prob[0]]
    if max(arg.ga_init_prob) > 0.99 or min(arg.ga_init_prob) < 0.01:
        exit_error('--ga-init-prob must be in range [0.01 : 0.99]', parser)
    if arg.ga_init_prob[0] > arg.ga_init_prob[1]:
        exit_error('--ga-init-prob a b : a <= b must hold', parser)
    if not arg.demo and not arg.image:
        exit_error('FITS image parameter not provided (and no --version, --help or --demo)', parser)

def exit_error(message, parser):
    logging.error(message)
    if parser:
        parser.error()
        parser.exit(1)
    else:
        raise ValueError(message)

def info():
    """Prints commandline help message"""
    commons.info(__arg_parser())

def commandline_entry():
    # Entry point for command line
    __args = __arg_parser().parse_args()  # parse command line arguments
    if __args.version:
        print('astwro.tools ' + astwro.tools.__version__)
        exit()

    # Configure logging
    #    logging.basicConfig(format='[%(levelname)s] %(module)s: %(message)s', level=__args.loglevel.upper())
    logging.basicConfig(format='%(levelname)7s | %(message)s', level=__args.loglevel.upper())
    check_arguments(__args, __arg_parser())
    __stars = __do(__args)  # call main routine - common form command line and python calls
    if __stars is None:
        return 1
    if not __args.no_stdout:
        print('\n'.join(map(str, __stars.index)))
    return 0

if __name__ == '__main__':
    code = commandline_entry()
    exit(code)


# -*- coding: utf-8 -*-
"""
Created on Mon Dec 17 19:11:22 2018

@author: Jian
"""

import numpy as np

from . import utils

from . import BSSA14
from . import CB14
from . import CY14
from . import ASK14
from . import helper

#%%----------------------------------------------------------------------------
def GMPE(model_name, Mw, Rjb, Vs30, period, epislon=0, NGAs=None,
         rake=None, Mech=3, Ftype=None, Fnm=None, Frv=None,
         dip=None, W=None, Ztor=None, Zhypo=None, Fas=0,
         Rrup=None, Rx=None, Fhw=None, azimuth=None,
         VsFlag=0, Z25=None, Z15=None, Z10=None,
         ArbCB=0, SJ=0,
         country='California', region='CA',
         Dregion='GlobalCATW',
         CRjb=15, Ry0=None,
         D_DPP=0):

    model_name = helper.check_model_name(model_name)

    if NGAs == None:
        NGAs={'CB':{'NewCoefs':None,'terms':(1,1,1,1,1,1,1,1,1)},
              'BSSA':{'NewCoefs':None,'terms':(1,1,1)},
              'CY':{'NewCoefs':None,'terms':(1,1,1,1,1,1,1)},
              'ASK':{'NewCoefs':None,'terms':(1,1,1,1,1,1,1)}}
    else:
        if 'BSSA' not in list(NGAs.keys()):
            NGAs['BSSA'] = NGAs['BA']
        if 'ASK' not in list(NGAs.keys()):
            NGAs['ASK'] = NGAs['AS']

    dict1 = NGAs

    # check the input period
    # Note: this function is better used at a given period with a set of other parameters (not with a set of periods)
    if period > 10.0 or 0 < period < 0.01:
        raise ValueError('Positive period value should be within [0.01,10] for SA at corresponding periods')
    if period < 0 and period not in [-1,-2]:
        raise ValueError('negative period should be -1,-2 for PGA and PGV')

    if model_name == 'BSSA':
        ngaM = BSSA14.BSSA14_nga()
        kwds = {'Mech': Mech, 'Ftype': Ftype, 'Z10': Z10, 'Dregion': Dregion,
                'country': country, 'CoefTerms': dict1[model_name]}
    elif model_name == 'CB':
        ngaM = CB14.CB14_nga()
        if Fhw is None:
            Fhw = 0
        kwds = {'Ftype': Ftype, 'Rrup': Rrup, 'Ztor': Ztor, 'dip': dip,
                'Z25': Z25, 'W': W, 'Zhypo': Zhypo, 'azimuth': azimuth,
                'Fhw': Fhw, 'Z10': Z10, 'Z15': Z15, 'Arb': ArbCB, 'SJ': SJ,
                'region': region, 'CoefTerms': dict1[model_name]}
    elif model_name == 'CY':
        ngaM = CY14.CY14_nga()
        kwds = {'Ftype': Ftype, 'Rrup': Rrup, 'Rx': Rx, 'Ztor': Ztor,
                'dip': dip, 'W': W, 'Zhypo': Zhypo, 'azimuth': azimuth,
                'Fhw': Fhw, 'Z10': Z10, 'AS': Fas, 'VsFlag': VsFlag,
                'country': country, 'D_DPP': D_DPP,
                'CoefTerms': dict1[model_name]}
        # the new CY model treat PGA = SA(0.01)
        if period == -1:
            period = 0.01
    elif model_name == 'ASK':
        ngaM = ASK14.ASK14_nga()
        kwds = {'Ftype': Ftype, 'Rrup': Rrup, 'Rx': Rx, 'Ztor': Ztor,
                'dip': dip, 'W': W, 'Zhypo': Zhypo, 'azimuth': azimuth,
                'Fhw': Fhw, 'Z10': Z10, 'Fas': Fas, 'CRjb': CRjb, 'Ry0': Ry0,
                'region': region, 'country': country, 'VsFlag': VsFlag,
                'CoefTerms': dict1[model_name]}
    else:
        raise ValueError("`model_name` must be one of {'ASK', 'BSSA', 'CB', 'CY'}")

    # common interpolate for all models
    periods = np.array(ngaM.periods)
    if period in periods:
        # compute median, std directly for the existing period in the period list of the NGA model
        values = utils.mapfunc( ngaM, Mw, Rjb, Vs30, period, rake, **kwds )
        values = np.array( values )
    else:
        #print 'do the interpolation for periods that is not in the period list of the NGA model'
        ind_low = (periods <= period*1.0).nonzero()[0]
        ind_high = (periods >= period*1.0).nonzero()[0]

        period_low = max( periods[ind_low] )
        period_high = min( periods[ind_high] )
        values_low = np.array(utils.mapfunc(ngaM, Mw, Rjb, Vs30, period_low,
                                            rake, **kwds))
        values_high = np.array(utils.mapfunc(ngaM, Mw, Rjb, Vs30, period_high,
                                             rake, **kwds))
        N1,N2 = np.array(values_low).shape
        values = np.zeros( (N1,N2) )
        for icmp in range( N2 ):
            if icmp != 0:
                # stardand values are in ln (g)
                values[:,icmp] = utils.logline(np.log(period_low),
                                               np.log(period_high),
                                               values_low[:,icmp],
                                               values_high[:,icmp],
                                               np.log(period) )
            else:
                # median value is in g
                values[:,icmp] = utils.logline(np.log(period_low),
                                               np.log(period_high),
                                               np.log(values_low[:,icmp]),
                                               np.log(values_high[:,icmp]),
                                               np.log(period))
                values[:,icmp] = np.exp( values[:,icmp] )    # change the median into g unit (logline gives the result in ln(g))

    # outputs
    NGAsigmaT = values[:,1]
    NGAtau = values[:,2]
    NGAsigma = values[:,3]

    if epislon:
        NGAmedian = np.exp( np.log(values[:,0]) + epislon * NGAsigmaT )
    else:
        NGAmedian = values[:,0]

    # returned quantities are all in g, not in log(g), event for the standard deviations
    return NGAmedian[0], \
           np.exp(NGAsigmaT)[0], \
           np.exp(NGAtau)[0], \
           np.exp(NGAsigma)[0]      # all in g, include the standard deviation

#%%----------------------------------------------------------------------------
def GMPE_array(model_name, Mw, Rjb, Vs30, period_array, epislon=0, NGAs=None,
         rake=None, Mech=3, Ftype=None, Fnm=None, Frv=None,
         dip=None, W=None, Ztor=None, Zhypo=None, Fas=0,
         Rrup=None, Rx=None, Fhw=None, azimuth=None,
         VsFlag=0, Z25=None, Z15=None, Z10=None,
         ArbCB=0, SJ=0,
         country='California', region='CA',
         Dregion='GlobalCATW',
         CRjb=15, Ry0=None,
         D_DPP=0):

    import collections
    if not isinstance(period_array, collections.Iterable):
        raise TypeError('`period_array` must be a array-like object.')

    model_name = helper.check_model_name(model_name)

    results = []
    for period in period_array:
        if period < 0:
            raise ValueError('Period value must be positive')
        result_ = GMPE(model_name, Mw, Rjb, Vs30, period, epislon=epislon,
               NGAs=NGAs, rake=rake, Mech=Mech, Ftype=Ftype, Fnm=Fnm,
               Frv=Frv, dip=dip, W=W, Ztor=Ztor, Zhypo=Zhypo, Fas=Fas,
               Rrup=Rrup, Rx=Rx, Fhw=Fhw, azimuth=azimuth,
               VsFlag=VsFlag, Z25=Z25, Z15=Z15, Z10=Z10,
               ArbCB=ArbCB, SJ=SJ,country=country, region=region,
               Dregion=Dregion, CRjb=CRjb, Ry0=Ry0, D_DPP=D_DPP)

        results.append(result_)

    results_T = list(zip(*results))

    return tuple([np.array(_) for _ in results_T])


import cv2
import os
from keras.models import load_model
import numpy as np
from pygame import mixer
import time

mixer.init()
sound = mixer.Sound('alarm.wav')

face = cv2.CascadeClassifier('haar cascade files\haarcascade_frontalface_alt.xml')
leye = cv2.CascadeClassifier('haar cascade files\haarcascade_lefteye_2splits.xml')
reye = cv2.CascadeClassifier('haar cascade files\haarcascade_righteye_2splits.xml')

lbl=['Close','Open']

model = load_model('models/cnncat2.h5')
path = os.getcwd()
cap = cv2.VideoCapture(0)
font = cv2.FONT_HERSHEY_COMPLEX_SMALL
count=0
score=0
thicc=2
rpred=[99]
lpred=[99]

while(True):
    ret, frame = cap.read()
    height,width = frame.shape[:2]

    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    faces = face.detectMultiScale(gray,minNeighbors=5,scaleFactor=1.1,minSize=(25,25))
    left_eye = leye.detectMultiScale(gray)
    right_eye =  reye.detectMultiScale(gray)

    cv2.rectangle(frame, (0,height-50) , (200,height) , (0,0,0) , thickness=cv2.FILLED )

    for (x,y,w,h) in faces:
        cv2.rectangle(frame, (x,y) , (x+w,y+h) , (100,100,100) , 1 )

    for (x,y,w,h) in right_eye:
        r_eye=frame[y:y+h,x:x+w]
        count=count+1
        r_eye = cv2.cvtColor(r_eye,cv2.COLOR_BGR2GRAY)
        r_eye = cv2.resize(r_eye,(24,24))
        r_eye= r_eye/255
        r_eye=  r_eye.reshape(24,24,-1)
        r_eye = np.expand_dims(r_eye,axis=0)
        rpred = model.predict_classes(r_eye)
        if(rpred[0]==1):
            lbl='Alert'
        if(rpred[0]==0):
            lbl='Sleepy!'
        break

    for (x,y,w,h) in left_eye:
        l_eye=frame[y:y+h,x:x+w]
        count=count+1
        l_eye = cv2.cvtColor(l_eye,cv2.COLOR_BGR2GRAY)
        l_eye = cv2.resize(l_eye,(24,24))
        l_eye= l_eye/255
        l_eye=l_eye.reshape(24,24,-1)
        l_eye = np.expand_dims(l_eye,axis=0)
        lpred = model.predict_classes(l_eye)
        if(lpred[0]==1):
            lbl='Alert'
        if(lpred[0]==0):
            lbl='Sleepy!'
        break

    if(rpred[0]==0 and lpred[0]==0):
        score=score+1
        cv2.putText(frame,"Sleepy!",(10,height-20), font, 1,(255,255,255),1,cv2.LINE_AA)
    # if(rpred[0]==1 or lpred[0]==1):
    else:
        score=score-1
        cv2.putText(frame,"Alert",(10,height-20), font, 1,(255,255,255),1,cv2.LINE_AA)

    if(score<0):
        score=0
    cv2.putText(frame,'Score:'+str(score),(100,height-20), font, 1,(255,255,255),1,cv2.LINE_AA)
    if(score>15):
        #person is feeling sleepy so we beep the alarm
        cv2.imwrite(os.path.join(path,'image.jpg'),frame)
        try:
            sound.play()

        except:  # isplaying = False
            pass
        if(thicc<16):
            thicc= thicc+2
        else:
            thicc=thicc-2
            if(thicc<2):
                thicc=2
        cv2.rectangle(frame,(0,0),(width,height),(0,0,255),thicc)
    cv2.imshow('frame',frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break
cap.release()
cv2.destroyAllWindows()

SEM VII/Data Science and Analytics/Mini Project/Driver Drowsiness Detection/DDDS_CNN/main_capture.py


#!/usr/bin/env python

"""
Inference  with frozen_inference_graph.pb
Author: @developmentseed

Use:
   python frozen_pred.py \
    --frozen_inference_graph=data/frozen_inference_graph.pb \
    --images_path=data/chips/*.jpg \
    --threshold=0.5 \
    --batch_size=5
"""
import os
from os import makedirs, path as op
import tensorflow as tf
import tensorflow.compat.v1 as tf1
import numpy as np
from joblib import Parallel, delayed
import glob
import json
from PIL import Image
from itertools import zip_longest
from numpyencoder import NumpyEncoder
import click
import time
from tqdm import tqdm

start_time = time.time()
# Just disables the warning, doesn't enable AVX/FMA
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
print("Tensorflow version " + tf.__version__)

inference_graph="models/livestock/frozen_inference_graph.pb"
detection_graph = tf1.Graph()
with detection_graph.as_default():
    od_graph_def = tf1.GraphDef()
    with tf1.gfile.GFile(inference_graph, "rb") as fid:
        serialized_graph = fid.read()
        od_graph_def.ParseFromString(serialized_graph)
        tf1.import_graph_def(od_graph_def, name="")

def _grouper(iterable, n, fillvalue=None):
    "Itertool recipe to collect data into fixed-length chunks or blocks"
    # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"
    args = [iter(iterable)] * n
    return zip_longest(*args, fillvalue=fillvalue)

def tf1_od_pred(list_images, conf_threshold, frozen_inference_graph):
    """
    funtion to make predictions
    """
    detections = ([])  # contains img name as key and detection info as value (boxes, scores, classes, num)
    with detection_graph.as_default():
        with tf1.Session(graph=detection_graph) as sess:
            # all_tensor_names= [tensor.name for tensor in tf1.get_default_graph().as_graph_def().node]
            # print(all_tensor_names)
            image_tensor = detection_graph.get_tensor_by_name("image_tensor:0")
            # Each box represents a part of the image where a particular object was detected.
            detection_boxes = detection_graph.get_tensor_by_name("detection_boxes:0")
            # Each score represent how level of confidence for each of the objects.
            # Score is shown on the result image, together with the class label.
            detection_scores = detection_graph.get_tensor_by_name("detection_scores:0")
            detection_classes = detection_graph.get_tensor_by_name("detection_classes:0")
            num_detections = detection_graph.get_tensor_by_name("num_detections:0")
            # the array based representation of the image will be used later in order to prepare the
            # result image with boxes and labels on it.
            # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
            # image_np_expanded = np.expand_dims(item['img_arr'], axis=0)
            # Actual detection.
            for item_img in list_images:
                (preds_boxes, preds_scores, preds_classes, preds_num,) = sess.run(
                    [
                        detection_boxes,
                        detection_scores,
                        detection_classes,
                        num_detections,
                    ],
                    feed_dict={image_tensor: item_img["img_arr"]},
                )

                #####################################################
                # Filter predictions
                #####################################################
                preds_num = int(preds_num)
                if preds_num > 0:

                    filtered_boxes = []
                    filtered_score = []
                    filtered_classes = []

                    for idx, scores in enumerate(preds_scores):
                        for i, score in enumerate(scores):
                            if score > conf_threshold:
                                filtered_boxes.append(preds_boxes[idx][i])
                                filtered_score.append(score)
                                filtered_classes.append(preds_classes[idx][i])

                    detections.append(
                        {
                            "image_id": item_img["image_id"],
                            "boxes": filtered_boxes,
                            "scores": filtered_score,
                            "classes": filtered_classes,
                            "num": preds_num,
                        }
                    )
                print(f'Predictions for {item_img["image_id"]}: {str(time.time() - start_time)}')

        return detections

def load_image_into_numpy_array(image):
    (im_width, im_height) = image.size
    return np.array(image.getdata()).reshape((im_height, im_width, 3)).astype(np.uint8)

@click.command(short_help="Testing local tfserving api")
@click.option(
    "-m",
    "--frozen_inference_graph",
    help="frozen inference graph path",
    required=True,
    type=str,
    default="data/ryan_sprint_rcnn_resnet101_serengeti_wildlife_v3_tfs/frozen_inference_graph.pb",
)
@click.option(
    "-s",
    "--images_path",
    help="Path to images",
    required=True,
    type=str,
    default="data/chips/*.jpg",
)
@click.option(
    "-t",
    "--threshold",
    help="threshold to filter the predictions",
    required=True,
    type=float,
    default=0.5,
)
@click.option(
    "-b",
    "--batch_size",
    help="batch size for running inference",
    required=True,
    type=float,
    default=5,
)

def main(frozen_inference_graph, images_path, threshold, batch_size):
    ##########################################
    ### Read images and convert to numpy array
    ##########################################
    images = glob.glob(images_path)
    # images_tensors = []

    for i, img_group in enumerate(_grouper(images, int(batch_size))):
        images_tensors = []
        for batch_img_fname in img_group:
            if batch_img_fname is not None:
                image = Image.open(batch_img_fname)
                image_np = load_image_into_numpy_array(image)
                # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
                image_np_expanded = np.expand_dims(image_np, axis=0)
                images_tensors.append({"image_id": batch_img_fname, "img_arr": image_np_expanded})
        print(f"Convert done {len(images)} images to images_tensor: {str(time.time() - start_time)}")

        ##########################################
        ### Inference with frozen_inference_graph.pb in parallel
        ##########################################

        detections = tf1_od_pred(images_tensors, threshold, frozen_inference_graph)
        print(f"Predictions and filters done: {str(time.time() - start_time)}")
        with open(f"result_{i}.json", "w") as f:
            f.write(json.dumps(detections, cls=NumpyEncoder))

if __name__ == "__main__":
    main()