"""
Procedures needed for Common support estimation.

Created on Thu Dec  8 15:48:57 2020.

@author: MLechner

# -*- coding: utf-8 -*-
"""
import copy
import numpy as np
import pandas as pd
from mcf import mcf_data_functions as mcf_data
from mcf import general_purpose as gp
from mcf import general_purpose_estimation as gp_est

def common_support(predict_file, tree_file, fill_y_file, fs_file, var_x_type,
                   v_dict, c_dict, cs_list=None, prime_values_dict=None,
                   pred_tr_np=None, d_tr_np=None):
    """
    Remove observations from data files that are off-support.

    Parameters
    ----------
    predict_file : String of csv-file. Data to predict the RF.
    train_file : String of csv-file. Data to train the RF.
    fill_y_file : String of csv-file. Data with y to be used by RF.
    fs_file : String of csv-file. Data with y to be used by RF.
    var_x_type : Dict. Features.
    v_dict : Dict. Variables.
    c_dict : Dict. Parameters.
    cs_list: Tuple. Contains the information from estimated propensity score
                    needed to predict for other data. Default is None.
    prime_values_dict: Dict. List of unique values for variables to dummy.
                    Default is None.
    pred_t: Numpy array. Predicted treatment probabilities in training data.
                         Needed to define cut-offs.
    d_train: Numpy series. Observed treatment in training data (tree_file).

    Returns
    -------
    predict_file_new : String of csv-file. Adjusted data.
    cs_list: Tuple. Contains the information from estimated propensity score
                    needed to predict for other data.
    pred_t: Numpy array. Predicted treatment probabilities in training data.
    d_train_tree: estimated tree by sklearn.

    """
    def r2_obb(c_dict, idx, oob_best):
        if c_dict['with_output']:
            print('\n')
            print('-' * 80)
            print('Treatment: {:2}'.format(c_dict['d_values'][idx]),
                  'OOB Score (R2 in %): {:6.3f}'.format(oob_best * 100))
            print('-' * 80)

    def get_data(file_name, x_name):
        data = pd.read_csv(file_name)
        x_all = data[x_name]    # deep copies
        obs = len(x_all.index)
        return data, x_all, obs

    def check_cols(x_1, x_2, name1, name2):
        var1 = set(x_1.columns)
        var2 = set(x_2.columns)
        if var1 != var2:
            if len(var1-var2) > 0:
                print('Variables in ', name1, 'not contained in ', name2,
                      *(var1-var2))
            if len(var2-var1) > 0:
                print('Variables in ', name2, 'not contained in ', name1,
                      *(var2-var1))
            raise Exception(name1 + ' data and ' + name2 + ' data contain' +
                            ' differnt variables. Programm stopped.')

    def mean_by_treatment(treat_pd, data_pd):
        treat_pd = treat_pd.squeeze()
        treat_vals = pd.unique(treat_pd)
        print('--------------- Mean by treatment status ------------------')
        if len(treat_vals) > 0:
            mean = data_pd.groupby(treat_pd).mean()
            print(mean.transpose())
        else:
            print('All obs have same treatment:', treat_vals)

    def on_support_data_and_stats(obs_to_del_np, data_pd, x_data_pd, out_file,
                                  upper_l, lower_l, c_dict, header=False,
                                  d_name=None):
        obs_to_keep = np.invert(obs_to_del_np)
        data_keep = data_pd[obs_to_keep]
        gp.delete_file_if_exists(out_file)
        data_keep.to_csv(out_file, index=False)
        if c_dict['with_output']:
            x_keep = x_data_pd[obs_to_keep]
            x_delete = x_data_pd[obs_to_del_np]
            if header:
                print('\n')
                print('=' * 80)
                print('Common support check')
                print('-' * 80)
                print('Upper limits on treatment probabilities: ', upper_l)
                print('Lower limits on treatment probabilities: ', lower_l)
            print('-' * 80)
            print('Data investigated and saved:', out_file)
            print('-' * 80)
            print('Observations deleted: {:4}'.format(np.sum(obs_to_del_np)),
                  ' ({:6.3f}%)'.format(np.mean(obs_to_del_np)*100))
            with pd.option_context(
                    'display.max_rows', 500,
                    'display.max_columns', 500,
                    'display.expand_frame_repr', True,
                    'display.width', 150,
                    'chop_threshold', 1e-13):
                all_var_names = [name.upper() for name in data_pd.columns]
                if d_name[0].upper() in all_var_names:
                    d_keep = data_keep[d_name]
                    d_delete = data_pd[d_name]
                    d_delete = d_delete[obs_to_del_np]
                    d_keep_count = d_keep.value_counts(sort=False)
                    d_delete_count = d_delete.value_counts(sort=False)
                    d_keep_count = pd.concat(
                        [d_keep_count,
                         d_keep_count / np.sum(obs_to_keep) * 100], axis=1)
                    d_delete_count = pd.concat(
                        [d_delete_count,
                         d_delete_count / np.sum(obs_to_del_np) * 100], axis=1)
                    d_keep_count.columns = ['Obs.', 'Share in %']
                    d_delete_count.columns = ['Obs.', 'Share in %']
                    if c_dict['panel_data']:
                        cluster_id = data_pd[v_dict['cluster_name']].squeeze()
                        cluster_keep = cluster_id[obs_to_keep].squeeze()
                        cluster_delete = cluster_id[obs_to_del_np].squeeze()
                    print('-' * 80)
                    print('Observations kept by treatment')
                    print(d_keep_count)
                    print('-   ' * 20)
                    print('Observations deleted by treatment')
                    print(d_delete_count)
                    if c_dict['panel_data']:
                        print('-   ' * 20)
                        print('Total number of panel unit:',
                              len(cluster_id.unique()))
                        print('Observations belonging to ',
                              len(cluster_keep.unique()),
                              'panel units are ON support')
                        print('Observations belonging to ',
                              len(cluster_delete.unique()),
                              'panel units are OFF support')
                if d_name[0].upper() in all_var_names:
                    print()
                    print('Full sample (ON and OFF support observations)')
                    mean_by_treatment(data_pd[d_name], x_data_pd)
                print('-' * 80)
                print('Data ON support')
                print('-' * 80)
                print(x_keep.describe().transpose())
                if d_name[0].upper() in all_var_names:
                    print()
                    mean_by_treatment(d_keep, x_keep)
                print('-' * 80)
                print('Data OFF support')
                print('-' * 80)
                print(x_delete.describe().transpose())
                if d_name[0].upper() in all_var_names:
                    print()
                    if np.sum(obs_to_del_np) > 1:
                        mean_by_treatment(d_delete, x_delete)
                    else:
                        print('Only single observation deleted.')
            if np.mean(obs_to_del_np) > c_dict['support_max_del_train']:
                raise Exception(
                    'Less than {:3}%'.format(
                        100-c_dict['support_max_del_train']*100)
                    + ' observations left after common support check of'
                    + ' training data. Programme terminated. Improve'
                    + ' balance of input data for forest building.')

    x_name, x_type = gp.get_key_values_in_list(var_x_type)
    names_unordered = []  # Split ordered variables into dummies
    for j, val in enumerate(x_type):
        if val > 0:
            names_unordered.append(x_name[j])
    fs_adjust = False
    obs_fs = 0
    if c_dict['train_mcf']:
        data_tr, x_tr, obs_tr = get_data(tree_file, x_name)  # train,adj.
        data_fy, x_fy, obs_fy = get_data(fill_y_file, x_name)  # adj.
        if c_dict['fs_yes']:
            # if not ((fs_file == tree_file) or (fs_file == fill_y_file)):
            if fs_file not in (tree_file, fill_y_file):
                data_fs, x_fs, obs_fs = get_data(fs_file, x_name)  # adj.
                fs_adjust = True
    if c_dict['pred_mcf']:
        data_pr, x_pr, obs_pr = get_data(predict_file, x_name)
    else:
        obs_pr = 0
    if names_unordered:  # List is not empty
        if c_dict['train_mcf'] and c_dict['pred_mcf']:
            x_total = pd.concat([x_tr, x_fy, x_pr], axis=0)
            if fs_adjust:
                x_total = pd.concat([x_total, x_fs], axis=0)
            x_dummies = pd.get_dummies(x_total[names_unordered],
                                       columns=names_unordered)
            x_total = pd.concat([x_total, x_dummies], axis=1)
            x_tr = x_total[:obs_tr]
            x_fy = x_total[obs_tr:obs_tr+obs_fy]
            x_pr = x_total[obs_tr+obs_fy:obs_tr+obs_fy+obs_pr]
            if fs_adjust:
                x_fs = x_total[obs_tr+obs_fy+obs_pr:]
        elif c_dict['train_mcf'] and not c_dict['pred_mcf']:
            x_total = pd.concat([x_tr, x_fy], axis=0)
            if fs_adjust:
                x_total = pd.concat([x_total, x_fs], axis=0)
            x_dummies = pd.get_dummies(x_total[names_unordered],
                                       columns=names_unordered)
            x_total = pd.concat([x_total, x_dummies], axis=1)
            x_tr = x_total[:obs_tr]
            x_fy = x_total[obs_tr:obs_tr+obs_fy]
            if fs_adjust:
                x_fs = x_total[obs_tr+obs_fy:]
        else:
            x_add_tmp = check_if_obs_needed(names_unordered, x_pr,
                                            prime_values_dict)
            if x_add_tmp is not None:
                x_total = pd.concat([x_pr, x_add_tmp], axis=0)
            else:
                x_total = x_pr
            x_dummies = pd.get_dummies(x_total[names_unordered],
                                       columns=names_unordered)
            x_pr = pd.concat([x_total, x_dummies], axis=1)
            if x_add_tmp is not None:  # remove add_temp
                x_pr = x_pr[:obs_pr]
    if c_dict['train_mcf']:
        x_name_all = x_tr.columns.values.tolist()
    else:
        x_name_all = x_pr.columns.values.tolist()
    if c_dict['train_mcf']:
        x_tr_np = x_tr.to_numpy(copy=True)
        d_all_in = pd.get_dummies(data_tr[v_dict['d_name']],
                                  columns=v_dict['d_name'])
        d_tr_np = d_all_in.to_numpy(copy=True)
        pred_tr_np = np.empty((np.shape(x_tr_np)[0], c_dict['no_of_treat']))
    if c_dict['train_mcf']:
        x_pred_all = x_fy.copy()
        if c_dict['pred_mcf']:
            x_pred_all = pd.concat([x_pred_all, x_pr], axis=0)
        if fs_adjust:
            x_pred_all = pd.concat([x_pred_all, x_fs], axis=0)
    else:
        obs_fy = 0
        x_pred_all = x_pr.copy()
    x_pred_all_np = x_pred_all.to_numpy(copy=True)
    pred_all_np = np.empty((obs_fy+obs_fs+obs_pr, c_dict['no_of_treat']))
    if c_dict['no_parallel'] > 1:
        workers_mp = copy.copy(c_dict['no_parallel'])
    else:
        workers_mp = None
    if c_dict['train_mcf']:
        check_cols(x_tr, x_fy, 'Tree', 'Fill_y')
        if fs_adjust:
            check_cols(x_tr, x_fs, 'Tree', 'Feature selection')
    if c_dict['train_mcf'] and c_dict['pred_mcf']:
        check_cols(x_tr, x_pr, 'Tree', 'Prediction')
    if c_dict['train_mcf']:
        if c_dict['with_output'] and c_dict['verbose']:
            print('\n')
            print('-' * 80)
            print('Computing random forest based common support')
    if c_dict['train_mcf']:
        cs_list = []
        c_dict_new = mcf_data.m_n_grid(copy.deepcopy(c_dict),    # dict only
                                       len(x_pred_all.columns))  # used here
        for idx in range(c_dict['no_of_treat']):
            return_forest = bool(c_dict['save_forest'])
            ret_rf = gp_est.RandomForest_scikit(
                x_tr_np, d_tr_np[:, idx], x_pred_all_np, boot=c_dict['boot'],
                n_min=c_dict_new['grid_n_min'],
                no_features=c_dict_new['grid_m'], workers=workers_mp,
                pred_p_flag=True, pred_t_flag=True, pred_oob_flag=True,
                with_output=False, variable_importance=True, x_name=x_name_all,
                var_im_with_output=c_dict['with_output'],
                return_forest_object=return_forest)
            pred_all_np[:, idx] = np.copy(ret_rf[0])
            pred_tr_np[:, idx] = np.copy(ret_rf[1])
            oob_best = np.copy(ret_rf[2])
            if c_dict['save_forest']:
                cs_list.append(ret_rf[6])
            r2_obb(c_dict, idx, oob_best)
            if c_dict['no_of_treat'] == 2:
                pred_all_np[:, idx+1] = 1 - pred_all_np[:, idx]
                pred_tr_np[:, idx+1] = 1 - pred_tr_np[:, idx]
                break
    else:
        for idx in range(c_dict['no_of_treat']):
            pred_all_np[:, idx] = cs_list[idx].predict(x_pred_all_np)
            if c_dict['no_of_treat'] == 2:
                pred_all_np[:, idx+1] = 1 - pred_all_np[:, idx]
                break
    obs_to_del_all, obs_to_del_tr, upper_l, lower_l = indicate_off_support(
        pred_tr_np, pred_all_np, d_tr_np, c_dict)
    # split obs_to_del_all into its parts
    obs_to_del_fs, obs_to_del_fy, obs_to_del_pr = False, False, False
    predict_file_new = None
    if c_dict['train_mcf']:
        obs_to_del_fy = obs_to_del_all[:obs_fy]
        if c_dict['pred_mcf']:
            obs_to_del_pr = obs_to_del_all[obs_fy:obs_fy+obs_pr]
        if fs_adjust:
            obs_to_del_fs = obs_to_del_all[obs_fy+obs_pr:]
    else:
        obs_to_del_pr = obs_to_del_all
    if c_dict['train_mcf']:
        if np.any(obs_to_del_tr):
            on_support_data_and_stats(obs_to_del_tr, data_tr, x_tr, tree_file,
                                      upper_l, lower_l, c_dict, header=True,
                                      d_name=v_dict['d_name'])
        if np.any(obs_to_del_fs):
            on_support_data_and_stats(obs_to_del_fs, data_fs, x_fs, fs_file,
                                      upper_l, lower_l, c_dict,
                                      d_name=v_dict['d_name'])
        if np.any(obs_to_del_fy):
            on_support_data_and_stats(obs_to_del_fy, data_fy, x_fy,
                                      fill_y_file, upper_l, lower_l, c_dict,
                                      d_name=v_dict['d_name'])
    if c_dict['pred_mcf']:
        if np.any(obs_to_del_pr):
            on_support_data_and_stats(
                obs_to_del_pr, data_pr, x_pr, c_dict['preddata3_temp'],
                upper_l, lower_l, c_dict, d_name=v_dict['d_name'])
            predict_file_new = c_dict['preddata3_temp']
        else:
            predict_file_new = predict_file
    else:
        predict_file_new = None
    return predict_file_new, cs_list, pred_tr_np, d_tr_np

def check_if_obs_needed(names_unordered, x_all_p, prime_values_dict):
    """Generate new rows -> all values of unordered variables are in data."""
    no_change = True
    max_length = 1
    for name in names_unordered:
        length = len(prime_values_dict[name])
        if length > max_length:
            max_length = length
    x_add_tmp = x_all_p[:max_length].copy()
    for name in names_unordered:
        unique_vals_p = np.sort(x_all_p[name].unique())
        unique_vals_t = np.sort(prime_values_dict[name])
        if len(unique_vals_p) > len(unique_vals_t) or (
               (len(unique_vals_p) == len(unique_vals_t))
               and not np.all(unique_vals_p == unique_vals_t)):
            print(name, 'Training values: ', unique_vals_t,
                  'Prediction values:', unique_vals_p)
            raise Exception('Common support variable value error')
        add_vals_in_train = np.setdiff1d(unique_vals_t, unique_vals_p)
        if add_vals_in_train.size > 0:
            for i, val in enumerate(add_vals_in_train):
                x_add_tmp[i, name] = val
            no_change = False
    if no_change:
        return None
    return x_add_tmp

def indicate_off_support(pred_t, pred_p, d_t, c_dict):
    """
    Indicate which observations are off support.

    Parameters
    ----------
    pred_t : N x no of treat Numpy array. Predictions of treat probs in train.
    pred_p : N x no of treat Numpy array. Predictions of treat probs in pred.
    d_t: N x no of treat Numpy array. Treatment dummies.
    c_dict : Dict. Parameters.

    Returns
    -------
    off_support_p : N x 1 Numpy array of boolean. True if obs is off support.
    off_support_t : N x 1 Numpy array of boolean. True if obs is off support.
    upper : No of treatment x 1 Numpy array of float. Upper limits.
    lower : No of treatment x 1 Numpy array of float. Lower limits.

    """
    # Normalize such that probabilities add up to 1
    pred_t = pred_t / pred_t.sum(axis=1, keepdims=True)
    pred_p = pred_p / pred_p.sum(axis=1, keepdims=True)
    n_p = np.shape(pred_p)[0]
    n_t = np.shape(pred_t)[0]
    q_s = c_dict['support_quantil']
    if c_dict['common_support'] == 1:
        upper_limit = np.empty((c_dict['no_of_treat'], c_dict['no_of_treat']))
        lower_limit = np.empty_like(upper_limit)
        for idx in range(c_dict['no_of_treat']):
            if q_s == 1:
                upper_limit[idx, :] = np.max(pred_t[d_t[:, idx] == 1], axis=0)
                lower_limit[idx, :] = np.min(pred_t[d_t[:, idx] == 1], axis=0)
            else:
                upper_limit[idx, :] = np.quantile(pred_t[d_t[:, idx] == 1],
                                                  q_s, axis=0)
                lower_limit[idx, :] = np.quantile(pred_t[d_t[:, idx] == 1],
                                                  1-q_s, axis=0)
        if c_dict['with_output']:
            print('Treatment sample     Treatment probabilities in %')
            print('--------------------- Upper limits ----------------')
            for idx, ival in enumerate(c_dict['d_values']):
                print('D = {:2}'.format(ival), end='              ')
                for jdx in range(c_dict['no_of_treat']):
                    print('{:7.4f} '.format(upper_limit[idx, jdx]), end=' ')
                print(' ')
            print('--------------------- Lower limits ----------------')
            for idx, ival in enumerate(c_dict['d_values']):
                print('D = {:2}'.format(ival), end='              ')
                for jdx in range(c_dict['no_of_treat']):
                    print('{:7.4f} '.format(lower_limit[idx, jdx]), end=' ')
                print(' ')
        upper = np.min(upper_limit, axis=0)
        lower = np.max(lower_limit, axis=0)
    else:
        # Normalize such that probabilities add up to 1
        upper = np.ones(
            c_dict['no_of_treat']) * (1 - c_dict['support_min_p'])
        lower = np.ones(c_dict['no_of_treat']) * c_dict['support_min_p']
    off_support_p = np.empty(n_p, dtype=bool)
    off_support_t = np.empty(n_t, dtype=bool)
    off_upper = np.empty(c_dict['no_of_treat'], dtype=bool)
    off_lower = np.empty_like(off_upper)
    for i in range(n_p):
        off_upper = np.any(pred_p[i, :] > upper)
        off_lower = np.any(pred_p[i, :] < lower)
        off_support_p[i] = off_upper or off_lower
    for i in range(n_t):
        off_upper = np.any(pred_t[i, :] > upper)
        off_lower = np.any(pred_t[i, :] < lower)
        off_support_t[i] = off_upper or off_lower
    return off_support_p, off_support_t, upper, lower


import os
import time
from datetime import datetime

import numpy as np
import tensorflow as tf

from ops import *

"""
cppgan-vae

compositional pattern-producing generative adversarial network combined with variational autoencoder

UPDATE: modified to get rid of likelihood function, and use pure GAN to draw to correct class.

I learned a lot from studying the below pages:

https://github.com/carpedm20/DCGAN-tensorflow
https://jmetzen.github.io/2015-11-27/vae.html

it wouldn't have been possible without referencing those two guy's code!

Description of CPPNs:

https://en.wikipedia.org/wiki/Compositional_pattern-producing_network

"""

class CPPNResnet:
    def __init__(
        self,
        batch_size=1,
        z_dim=32,
        x_dim=26,
        y_dim=26,
        c_dim=1,
        scale=8.0,
        learning_rate_g=0.01,
        learning_rate_d=0.001,
        learning_rate_vae=0.0001,
        beta1=0.9,
        net_size_g=6,
        net_depth_g=24,
        subnet_depth_g=4,
        net_size_q=512,
        keep_prob=1.0,
        df_dim=32,
        model_name="cppn_resnet",
        grad_clip=5.0,
        logdir="save",
    ):
        """

        Args:
        z_dim               dimensionality of the latent vector
        x_dim, y_dim        default resolution of generated images for training
        c_dim               1 for monotone, 3 for colour
        learning_rate_g     learning rate for the generator
                     _d    learning rate for the discriminiator
                     _vae  learning rate for the variational autoencoder
        net_size_g          number of activations per layer for cppn generator function
        net_depth_g         depth of generator
        net_size_q          number of activations per layer for decoder (real image -> z). 2 layers.
        df_dim              discriminiator is a convnet.  higher -> more activtions -> smarter.
        keep_prob           dropout probability

        when training, use I used dropout on training the decoder, batch norm on discriminator, nothing on cppn
        choose training parameters so that over the long run, decoder and encoder log errors hover around 0.7 each (so they are at the same skill level)
        while the error for vae should slowly move lower over time with D and G balanced.

        """

        self.batch_size = batch_size
        self.learning_rate_g = learning_rate_g
        self.learning_rate_d = learning_rate_d
        self.learning_rate_vae = learning_rate_vae
        self.beta1 = beta1
        self.net_size_g = net_size_g
        self.net_size_q = net_size_q
        self.x_dim = x_dim
        self.y_dim = y_dim
        self.scale = scale
        self.c_dim = c_dim
        self.z_dim = z_dim
        self.net_depth_g = net_depth_g
        self.subnet_depth_g = subnet_depth_g
        self.model_name = model_name
        self.keep_prob = keep_prob
        self.df_dim = df_dim
        self.num_class = 11  # 0->9 are MNIST classes, 10 are fake digits.
        self.grad_clip = grad_clip
        self.logdir = logdir

        # tf Graph batch of image (batch_size, height, width, depth)
        self.batch = tf.placeholder(tf.float32, [batch_size, x_dim, y_dim, c_dim])
        self.batch_flatten = tf.reshape(self.batch, [batch_size, -1])
        self.batch_label = tf.placeholder(
            tf.float32, [batch_size, self.num_class]
        )  # mnist labels for the batch (one-hot)
        self.fake_label = np.array(
            self.batch_size * [10], dtype=np.int32
        )  # label of fake batches
        self.fake_label_one_hot = self.to_one_hot(self.fake_label)

        n_points = x_dim * y_dim
        self.n_points = n_points

        self.x_vec, self.y_vec, self.r_vec = self.coordinates(x_dim, y_dim, scale)

        # latent vector
        # self.z = tf.placeholder(tf.float32, [self.batch_size, self.z_dim])
        # inputs to cppn, like coordinates and radius from centre
        self.x = tf.placeholder(tf.float32, [self.batch_size, None, 1])
        self.y = tf.placeholder(tf.float32, [self.batch_size, None, 1])
        self.r = tf.placeholder(tf.float32, [self.batch_size, None, 1])

        # batch normalization : deals with poor initialization helps gradient flow
        self.d_bn1 = batch_norm(batch_size, name=self.model_name + "_d_bn1")
        self.d_bn2 = batch_norm(batch_size, name=self.model_name + "_d_bn2")

        # Use recognition network to determine mean and
        # (log) variance of Gaussian distribution in latent
        # space
        self.z_mean, self.z_log_sigma_sq = self.encoder()

        # Draw one sample z from Gaussian distribution
        eps = tf.random_normal((self.batch_size, self.z_dim), 0, 1, dtype=tf.float32)
        # z = mu + sigma*epsilon
        self.z = tf.add(
            self.z_mean, tf.multiply(tf.sqrt(tf.exp(self.z_log_sigma_sq)), eps)
        )

        # Use generator to determine mean of
        # Bernoulli distribution of reconstructed input
        self.G = self.generator()
        # self.batch_reconstruct_flatten = tf.reshape(self.G, [batch_size, -1]) # not needed

        self.predict_real_samples = self.discriminator(
            self.batch
        )  # discriminiator on correct examples
        self.predict_fake_samples = self.discriminator(
            self.G, reuse=True
        )  # feed generated images into D

        self.create_vae_loss_terms()
        self.create_gan_loss_terms()

        self.balanced_loss = (
            1.0 * self.g_loss + 10.0 * self.vae_loss
        )  # can try to weight these.

        self.t_vars = tf.trainable_variables()

        self.q_vars = [
            var for var in self.t_vars if (self.model_name + "_q_") in var.name
        ]
        self.g_vars = [
            var for var in self.t_vars if (self.model_name + "_g_") in var.name
        ]
        self.d_vars = [
            var for var in self.t_vars if (self.model_name + "_d_") in var.name
        ]
        self.both_vars = self.q_vars + self.g_vars
        # self.vae_vars = self.q_vars # in this version, g_vars don't concern vae_loss

        # clip gradients
        d_opt_real_grads, _ = tf.clip_by_global_norm(
            tf.gradients(self.d_loss_real, self.d_vars), self.grad_clip
        )
        d_opt_grads, _ = tf.clip_by_global_norm(
            tf.gradients(self.d_loss, self.d_vars), self.grad_clip
        )
        g_opt_grads, _ = tf.clip_by_global_norm(
            tf.gradients(self.balanced_loss, self.both_vars), self.grad_clip
        )
        vae_opt_grads, _ = tf.clip_by_global_norm(
            tf.gradients(self.vae_loss, self.q_vars), self.grad_clip
        )

        # Use ADAM optimizer
        with tf.variable_scope(self.model_name + "_opt", reuse=tf.AUTO_REUSE) as scope:
            d_real_optimizer = tf.train.AdamOptimizer(
                self.learning_rate_d, beta1=self.beta1
            )
            d_optimizer = tf.train.AdamOptimizer(self.learning_rate_d, beta1=self.beta1)
            g_optimizer = tf.train.AdamOptimizer(self.learning_rate_g, beta1=self.beta1)
            vae_optimizer = tf.train.AdamOptimizer(
                self.learning_rate_vae, beta1=self.beta1
            )

            self.d_opt_real = d_real_optimizer.apply_gradients(
                zip(d_opt_real_grads, self.d_vars)
            )
            self.d_opt = d_optimizer.apply_gradients(zip(d_opt_grads, self.d_vars))
            self.g_opt = g_optimizer.apply_gradients(zip(g_opt_grads, self.both_vars))
            self.vae_opt = vae_optimizer.apply_gradients(
                zip(vae_opt_grads, self.q_vars)
            )

        """
    self.d_opt_real = tf.train.AdamOptimizer(self.learning_rate_d, beta1=self.beta1) \
                      .minimize(self.d_loss_real, var_list=self.d_vars)
    self.d_opt = tf.train.AdamOptimizer(self.learning_rate_d, beta1=self.beta1) \
                      .minimize(self.d_loss, var_list=self.d_vars)
    self.g_opt = tf.train.AdamOptimizer(self.learning_rate_g, beta1=self.beta1) \
                      .minimize(self.g_loss, var_list=self.both_vars)
    self.vae_opt = tf.train.AdamOptimizer(self.learning_rate_vae, beta1=self.beta1) \
                      .minimize(self.vae_loss, var_list=self.q_vars)
    """

        """
    tvars = tf.trainable_variables()
    grads, _ = tf.clip_by_global_norm(tf.gradients(self.cost, tvars), args.grad_clip)
    optimizer = tf.train.AdamOptimizer(self.lr, epsilon=0.001)
    self.train_op = optimizer.apply_gradients(zip(grads, tvars))
    """

        self.init()
        # self.saver = tf.train.Saver(tf.all_variables())

    def init(self):
        # Launch the session
        self.sess = tf.InteractiveSession()

        # init all vars
        self.all_vars = tf.get_collection_ref(tf.GraphKeys.GLOBAL_VARIABLES)
        self.sess.run(tf.variables_initializer(self.all_vars))

        # filter to trainable vars, and include only those in the Saver
        self.trainable_vars = [
            v
            for v in self.all_vars
            if "beta1_power" not in v.name and "beta2_power" not in v.name
        ]
        self.saver = tf.train.Saver(var_list=self.trainable_vars, max_to_keep=50)

        # initialize writer for tensorboard logs
        self.writer = tf.summary.FileWriter(self.logdir)

    def reinit(self):
        self.all_vars = tf.get_collection_ref(tf.GraphKeys.GLOBAL_VARIABLES)
        self.sess.run(tf.variables_initializer(self.all_vars))

    def to_one_hot(self, label):
        # convert labels, a numpy list of labels (of size batch_size) to the one hot equivalent
        return np.eye(self.num_class)[label]

    def create_vae_loss_terms(self):
        # The loss is composed of two terms:
        # 1.) The reconstruction loss (the negative log probability
        #     of the input under the reconstructed Bernoulli distribution
        #     induced by the decoder in the data space).
        #     This can be interpreted as the number of "nats" required
        #     for reconstructing the input when the activation in latent
        #     is given.
        # Adding 1e-10 to avoid evaluatio of log(0.0)

        # stop using likelihood function for similarity
        # reconstr_loss = \
        #     -tf.reduce_sum(self.batch_flatten * tf.log(1e-10 + self.batch_reconstruct_flatten)
        #                    + (1-self.batch_flatten) * tf.log(1e-10 + 1 - self.batch_reconstruct_flatten), 1)

        # 2.) The latent loss, which is defined as the Kullback Leibler divergence
        ##    between the distribution in latent space induced by the encoder on
        #     the data and some prior. This acts as a kind of regularizer.
        #     This can be interpreted as the number of "nats" required
        #     for transmitting the the latent space distribution given
        #     the prior.

        latent_loss = -0.5 * tf.reduce_sum(
            1
            + self.z_log_sigma_sq
            - tf.square(self.z_mean)
            - tf.exp(self.z_log_sigma_sq),
            1,
        )

        # self.vae_loss = tf.reduce_mean(reconstr_loss + latent_loss) / self.n_points # average over batch and pixel

        # vae loss is now purely kl divergence loss term.  let GAN take care of mnist class accuracy.
        self.vae_loss = (
            tf.reduce_mean(latent_loss) / self.n_points
        )  # average over batch and pixel
        self.summ_vae_loss = tf.summary.scalar("vae_loss", self.vae_loss)

    def create_gan_loss_terms(self):
        # Define loss function and optimiser
        """replace below with class-based disriminiator
        self.d_loss_real = binary_cross_entropy_with_logits(tf.ones_like(self.D_right), self.D_right)
        self.d_loss_fake = binary_cross_entropy_with_logits(tf.zeros_like(self.D_wrong), self.D_wrong)
        self.d_loss = 1.0*(self.d_loss_real + self.d_loss_fake)/ 2.0
        self.g_loss = 1.0*binary_cross_entropy_with_logits(tf.ones_like(self.D_wrong), self.D_wrong)
        """

        # cross entropy loss of predicting real mnist to real classes
        self.d_loss_real = tf.reduce_mean(
            -tf.reduce_sum(
                self.batch_label * tf.log(self.predict_real_samples),
                reduction_indices=[1],
            )
        )
        # accuracy of using discriminiator as a normal mnist classifier
        self.d_loss_real_accuracy = tf.reduce_mean(
            tf.cast(
                tf.equal(
                    tf.argmax(self.predict_real_samples, 1),
                    tf.argmax(self.batch_label, 1),
                ),
                tf.float32,
            )
        )
        # cross entropy loss of predicting that fake generated mnist are in fact fake
        self.d_loss_fake = tf.reduce_mean(
            -tf.reduce_sum(
                self.fake_label_one_hot * tf.log(self.predict_fake_samples),
                reduction_indices=[1],
            )
        )
        # accuracy of discriminator predicting a fake mnist digit
        self.d_loss_fake_accuracy = tf.reduce_mean(
            tf.cast(
                tf.equal(
                    tf.argmax(self.predict_fake_samples, 1),
                    tf.argmax(self.fake_label_one_hot, 1),
                ),
                tf.float32,
            )
        )
        # take the average of two d_loss to be the defacto d_loss
        self.d_loss = (
            10.0 * self.d_loss_real + self.d_loss_fake
        ) / 11.0  # balance out the classes
        self.summ_d_loss = tf.summary.scalar("d_loss", self.d_loss)
        self.summ_d_loss_fake = tf.summary.scalar("d_loss_fake", self.d_loss_fake)
        self.summ_d_loss_real = tf.summary.scalar("d_loss_real", self.d_loss_real)
        self.summ_d_loss_fake_accuracy = tf.summary.scalar(
            "d_loss_fake_accuracy", self.d_loss_fake_accuracy
        )
        self.summ_d_loss_real_accuracy = tf.summary.scalar(
            "d_loss_real_accuracy", self.d_loss_real_accuracy
        )

        # cross entropy of generator fooling discriminiator that its shit is real.
        self.g_loss = tf.reduce_mean(
            -tf.reduce_sum(
                self.batch_label * tf.log(self.predict_fake_samples),
                reduction_indices=[1],
            )
        )
        # accuracy of generated samples being fooled to be classified as their supposed ground truth labels
        self.g_loss_accuracy = tf.reduce_mean(
            tf.cast(
                tf.equal(
                    tf.argmax(self.predict_fake_samples, 1),
                    tf.argmax(self.batch_label, 1),
                ),
                tf.float32,
            )
        )
        self.summ_g_loss = tf.summary.scalar("g_loss", self.g_loss)
        self.summ_g_loss_accuracy = tf.summary.scalar(
            "g_loss_accuracy", self.g_loss_accuracy
        )

    def coordinates(self, x_dim=32, y_dim=32, scale=1.0):
        n_pixel = x_dim * y_dim
        x_range = scale * (np.arange(x_dim) - (x_dim - 1) / 2.0) / (x_dim - 1) / 0.5
        y_range = scale * (np.arange(y_dim) - (y_dim - 1) / 2.0) / (y_dim - 1) / 0.5
        x_mat = np.matmul(np.ones((y_dim, 1)), x_range.reshape((1, x_dim)))
        y_mat = np.matmul(y_range.reshape((y_dim, 1)), np.ones((1, x_dim)))
        r_mat = np.sqrt(x_mat * x_mat + y_mat * y_mat)
        x_mat = np.tile(x_mat.flatten(), self.batch_size).reshape(
            self.batch_size, n_pixel, 1
        )
        y_mat = np.tile(y_mat.flatten(), self.batch_size).reshape(
            self.batch_size, n_pixel, 1
        )
        r_mat = np.tile(r_mat.flatten(), self.batch_size).reshape(
            self.batch_size, n_pixel, 1
        )
        return x_mat, y_mat, r_mat

    def show_image(self, image):
        """
        image is in [height width depth]
        """
        plt.subplot(1, 1, 1)
        y_dim = image.shape[0]
        x_dim = image.shape[1]
        if self.c_dim > 1:
            plt.imshow(image, interpolation="nearest")
        else:
            plt.imshow(
                image.reshape(y_dim, x_dim), cmap="Greys", interpolation="nearest"
            )
        plt.axis("off")
        plt.show()

    def encoder(self):
        # Generate probabilistic encoder (recognition network), which
        # maps inputs onto a normal distribution in latent space.
        # The transformation is parametrized and can be learned.
        H1 = tf.nn.dropout(
            tf.nn.softplus(
                linear(self.batch_flatten, self.net_size_q, self.model_name + "_q_lin1")
            ),
            self.keep_prob,
        )
        H2 = tf.nn.dropout(
            tf.nn.softplus(linear(H1, self.net_size_q, self.model_name + "_q_lin2")),
            self.keep_prob,
        )
        z_mean = linear(H2, self.z_dim, self.model_name + "_q_lin3_mean")
        z_log_sigma_sq = linear(
            H2, self.z_dim, self.model_name + "_q_lin3_log_sigma_sq"
        )
        return (z_mean, z_log_sigma_sq)

    def discriminator(self, image, reuse=False):

        if reuse:
            tf.get_variable_scope().reuse_variables()

        h0 = lrelu(conv2d(image, self.df_dim, name=self.model_name + "_d_h0_conv"))
        h1 = lrelu(
            self.d_bn1(conv2d(h0, self.df_dim * 2, name=self.model_name + "_d_h1_conv"))
        )
        h2 = lrelu(
            self.d_bn2(conv2d(h1, self.df_dim * 4, name=self.model_name + "_d_h2_conv"))
        )
        h3 = linear(
            tf.reshape(h2, [self.batch_size, -1]),
            self.num_class,
            self.model_name + "_d_h2_lin",
        )

        return tf.nn.softmax(h3)

    def generator(self, gen_x_dim=26, gen_y_dim=26, reuse=False):

        if reuse:
            tf.get_variable_scope().reuse_variables()

        n_network = self.net_size_g
        gen_n_points = gen_x_dim * gen_y_dim

        z_scaled = (
            tf.reshape(self.z, [self.batch_size, 1, self.z_dim])
            * tf.ones([gen_n_points, 1], dtype=tf.float32)
            * self.scale
        )
        z_unroll = tf.reshape(z_scaled, [self.batch_size * gen_n_points, self.z_dim])
        x_unroll = tf.reshape(self.x, [self.batch_size * gen_n_points, 1])
        y_unroll = tf.reshape(self.y, [self.batch_size * gen_n_points, 1])
        r_unroll = tf.reshape(self.r, [self.batch_size * gen_n_points, 1])

        U = (
            fully_connected(z_unroll, n_network, self.model_name + "_g_0_z")
            + fully_connected(
                x_unroll, n_network, self.model_name + "_g_0_x", with_bias=False
            )
            + fully_connected(
                y_unroll, n_network, self.model_name + "_g_0_y", with_bias=False
            )
            + fully_connected(
                r_unroll, n_network, self.model_name + "_g_0_r", with_bias=False
            )
        )

        # H = tf.nn.relu(U)
        H = tf.nn.tanh(U)

        for i in range(0, self.net_depth_g):
            H0 = H
            for j in range(0, self.subnet_depth_g):
                H0 = tf.nn.relu(
                    fully_connected(
                        H0,
                        n_network,
                        self.model_name + "_g_relu_skip_" + str(i) + "_" + str(j),
                        stddev=1.0,
                    )
                )
            H0 = tf.nn.tanh(
                fully_connected(
                    H0,
                    n_network,
                    self.model_name + "_g_tanh_skip_" + str(i),
                    stddev=0.001,
                )
            )
            H = H + H0

        output = tf.sigmoid(
            fully_connected(
                H, self.c_dim, self.model_name + "_g_" + str(self.net_depth_g)
            )
        )

        result = tf.reshape(output, [self.batch_size, gen_y_dim, gen_x_dim, self.c_dim])

        return result

    def partial_train(self, batch, label):
        """Train model based on mini-batch of input data.

        Return cost of mini-batch.

        I should really seperate the below tricks into parameters, like number of times/pass
        and also the regulator threshold levels.
        """

        counter = 0

        label_one_hot = self.to_one_hot(label)

        """
    for i in range(1):
      counter += 1
      _, vae_loss = self.sess.run((self.vae_opt, self.vae_loss),
                              feed_dict={self.batch: batch, self.x: self.x_vec, self.y: self.y_vec, self.r: self.r_vec, self.batch_label: label_one_hot})
    """

        for i in range(16):
            counter += 1
            _, g_loss, vae_loss, g_accuracy = self.sess.run(
                (self.g_opt, self.g_loss, self.vae_loss, self.g_loss_accuracy),
                feed_dict={
                    self.batch: batch,
                    self.x: self.x_vec,
                    self.y: self.y_vec,
                    self.r: self.r_vec,
                    self.batch_label: label_one_hot,
                },
            )
            if g_accuracy > 0.98:
                break

        # train classifier on only real mnist digits
        # _ = self.sess.run((self.d_opt_real), feed_dict={self.batch: batch, self.x: self.x_vec, self.y: self.y_vec, self.r: self.r_vec, self.batch_label: label_one_hot})

        # calculate accuracy before deciding whether to train discriminator
        (
            d_loss,
            d_loss_real,
            d_loss_fake,
            d_real_accuracy,
            d_fake_accuracy,
        ) = self.sess.run(
            (
                self.d_loss,
                self.d_loss_real,
                self.d_loss_fake,
                self.d_loss_real_accuracy,
                self.d_loss_fake_accuracy,
            ),
            feed_dict={
                self.batch: batch,
                self.x: self.x_vec,
                self.y: self.y_vec,
                self.r: self.r_vec,
                self.batch_label: label_one_hot,
            },
        )

        if (
            d_fake_accuracy < 0.7 and g_accuracy > 0.6
        ):  # only train discriminiator if generator is good and d is behind.
            for i in range(8):
                counter += 1
                (
                    _,
                    d_loss,
                    d_loss_real,
                    d_loss_fake,
                    d_real_accuracy,
                    d_fake_accuracy,
                ) = self.sess.run(
                    (
                        self.d_opt,
                        self.d_loss,
                        self.d_loss_real,
                        self.d_loss_fake,
                        self.d_loss_real_accuracy,
                        self.d_loss_fake_accuracy,
                    ),
                    feed_dict={
                        self.batch: batch,
                        self.x: self.x_vec,
                        self.y: self.y_vec,
                        self.r: self.r_vec,
                        self.batch_label: label_one_hot,
                    },
                )
                if d_fake_accuracy > 0.75:
                    break
        elif d_real_accuracy < 0.6:
            for i in range(8):
                counter += 1
                _, d_real_accuracy = self.sess.run(
                    (self.d_opt_real, self.d_loss_real_accuracy),
                    feed_dict={
                        self.batch: batch,
                        self.x: self.x_vec,
                        self.y: self.y_vec,
                        self.r: self.r_vec,
                        self.batch_label: label_one_hot,
                    },
                )
                if d_real_accuracy > 0.7:
                    break

        return (
            d_loss,
            g_loss,
            vae_loss,
            counter,
            d_real_accuracy,
            d_fake_accuracy,
            g_accuracy,
            d_loss_real,
            d_loss_fake,
        )

    def encode(self, X):
        """Transform data by mapping it into the latent space."""
        # Note: This maps to mean of distribution, we could alternatively
        # sample from Gaussian distribution
        return self.sess.run(self.z_mean, feed_dict={self.batch: X})

    def generate(self, z=None, x_dim=26, y_dim=26, scale=5.0):
        """Generate data by sampling from latent space.

        If z is not None, data for this point in latent space is
        generated. Otherwise, z is drawn from prior in latent
        space.
        """
        if z is None:
            z = np.random.normal(size=self.z_dim).astype(np.float32)
        # Note: This maps to mean of distribution, we could alternatively
        # sample from Gaussian distribution

        z = np.reshape(z, (self.batch_size, self.z_dim))

        G = self.generator(gen_x_dim=x_dim, gen_y_dim=y_dim, reuse=True)
        gen_x_vec, gen_y_vec, gen_r_vec = self.coordinates(x_dim, y_dim, scale=scale)
        image = self.sess.run(
            G,
            feed_dict={
                self.z: z,
                self.x: gen_x_vec,
                self.y: gen_y_vec,
                self.r: gen_r_vec,
            },
        )
        return image

    def save_model(self, checkpoint_path, epoch):
        """ saves the model to a file """
        self.saver.save(self.sess, checkpoint_path, global_step=epoch)

    def load_model(self, checkpoint_path):

        ckpt = tf.train.get_checkpoint_state(checkpoint_path)
        print("loading model: ", ckpt.model_checkpoint_path)

        # self.saver.restore(self.sess, checkpoint_path+'/'+ckpt.model_checkpoint_path)
        # use the below line for tensorflow 0.7
        self.saver.restore(self.sess, ckpt.model_checkpoint_path)

    def close(self):
        self.sess.close()


"""
optimize.py: Driver and core functions for geometry optimization

Copyright 2016-2020 Regents of the University of California and the Authors

Authors: Lee-Ping Wang, Chenchen Song

Contributors: Yudong Qiu, Daniel G. A. Smith, Alberto Gobbi, Josh Horton

Redistribution and use in source and binary forms, with or without modification,
are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice,
this list of conditions and the following disclaimer.

2. Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.

3. Neither the name of the copyright holder nor the names of its contributors
may be used to endorse or promote products derived from this software
without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
POSSIBILITY OF SUCH DAMAGE.
"""

from __future__ import print_function, division

import os
import sys
import time
import traceback
import pkg_resources
from copy import deepcopy

import numpy as np
from numpy.linalg import multi_dot

import geometric
from .info import print_logo, print_citation
from .internal import CartesianCoordinates, PrimitiveInternalCoordinates, DelocalizedInternalCoordinates
from .ic_tools import check_internal_grad, check_internal_hess, write_displacements
from .normal_modes import calc_cartesian_hessian, frequency_analysis
from .step import brent_wiki, Froot, calc_drms_dmax, get_cartesian_norm, rebuild_hessian, get_delta_prime, trust_step, force_positive_definite
from .prepare import get_molecule_engine, parse_constraints
from .params import OptParams, parse_optimizer_args
from .nifty import row, col, flat, bohr2ang, ang2bohr, logger, bak, createWorkQueue
from .errors import InputError, HessianExit, EngineError, GeomOptNotConvergedError, GeomOptStructureError, LinearTorsionError

class Optimizer(object):
    def __init__(self, coords, molecule, IC, engine, dirname, params):
        """
        Object representing the geometry optimization of a molecular system.

        Parameters
        ----------
        coords : np.ndarray
            Nx3 array of Cartesian coordinates in atomic units
        molecule : Molecule
            Molecule object (Units Angstrom)
        IC : InternalCoordinates
            Object describing the internal coordinate system
        engine : Engine
            Object containing methods for calculating energy and gradient
        dirname : str
            Directory name for files to be written
        params : OptParams object
            Contains optimization parameters (really just a struct)
            Includes xyzout and qdata output file names (written if not None)
        """
        # Copies of data passed into constructor
        self.coords = coords
        self.molecule = deepcopy(molecule)
        self.IC = IC
        self.engine = engine
        self.dirname = dirname
        self.params = params
        # Set initial value of the trust radius.
        self.trust = self.params.trust
        # Copies of molecule object for preserving the optimization trajectory and the last frame
        self.progress = deepcopy(self.molecule)
        self.progress.xyzs = []
        self.progress.qm_energies = []
        self.progress.qm_grads = []
        self.progress.comms = []
        # Cartesian coordinates
        self.X = self.coords.copy()
        # Loop of optimization
        self.Iteration = 0
        # Counts how many steps it has been since checking the coordinate system
        self.CoordCounter = 0
        # Current state, used to control logic of optimization loop.
        self.state = OPT_STATE.NEEDS_EVALUATION
        # Some more variables to be updated throughout the course of the optimization
        self.trustprint = "="
        self.ForceRebuild = False
        # Sanity check - if there's only one atom, it will probably crash
        if self.molecule.na < 2:
            raise InputError("Geometry optimizer assumes there are at least two atoms in the system")

    def get_cartesian_norm(self, dy, verbose=None):
        if not verbose: verbose = self.params.verbose
        return get_cartesian_norm(self.X, dy, self.IC, self.params.enforce, self.params.verbose)

    def get_delta_prime(self, v0, verbose=None):
        # This method can be called at a different verbose level than the master
        # because it can occur inside a nested loop
        if not verbose: verbose = self.params.verbose
        return get_delta_prime(v0, self.X, self.G, self.H, self.IC, self.params.transition, verbose)

    def createFroot(self, v0):
        return Froot(self.trust, v0, self.X, self.G, self.H, self.IC, self.params)

    def refreshCoordinates(self):
        """
        Refresh the Cartesian coordinates used to define parts of the internal coordinate system.
        These include definitions of delocalized internal coordinates and reference coordinates for rotators.
        """
        logger.info("Refreshing coordinate system and resetting rotations\n")
        # Resetting of rotations
        self.IC.resetRotations(self.X)
        if isinstance(self.IC, DelocalizedInternalCoordinates):
            self.IC.build_dlc(self.X)
        # With redefined internal coordinates, the Hessian needs to be rebuilt
        self.H0 = self.IC.guess_hessian(self.coords)
        self.rebuild_hessian()
        # Current values of internal coordinates and IC gradient are recalculated
        self.Y = self.IC.calculate(self.X)
        self.G = self.IC.calcGrad(self.X, self.gradx)

    def checkCoordinateSystem(self, recover=False, cartesian=False):
        """
        Build a new internal coordinate system from current Cartesians and replace the current one if different.
        """
        # Reset the check counter
        self.CoordCounter = 0
        # Build a new molecule object and connectivity graph
        newmol = deepcopy(self.molecule)
        newmol.xyzs[0] = self.X.reshape(-1,3) * bohr2ang
        newmol.build_topology()
        # Build the new internal coordinate system
        if cartesian:
            if self.IC.haveConstraints():
                raise ValueError("Cannot continue a constrained optimization; please implement constrained optimization in Cartesian coordinates")
            IC1 = CartesianCoordinates(newmol)
        else:
            IC1 = self.IC.__class__(newmol, connect=self.IC.connect, addcart=self.IC.addcart, build=False, conmethod=self.IC.conmethod)
            if self.IC.haveConstraints(): IC1.getConstraints_from(self.IC)
        # Check for differences
        changed = (IC1 != self.IC)
        if changed:
            logger.info("\x1b[1;94mInternal coordinate system may have changed\x1b[0m\n")
            if self.IC.repr_diff(IC1) != "":
                logger.info(self.IC.repr_diff(IC1)+'\n')
        # Set current ICs to the new one
        if changed or recover or cartesian:
            self.IC = IC1
            self.refreshCoordinates()
            return True
        else: return False

    def trust_step(self, iopt, v0, verbose=None):
        # This method can be called at a different verbose level than the master
        # because it can occur inside a nested loop
        if not verbose: verbose = self.params.verbose
        return trust_step(iopt, v0, self.X, self.G, self.H, self.IC, self.params.transition, verbose)

    def newCartesian(self, dy):
        if self.IC.haveConstraints() and self.params.enforce:
            self.X = self.IC.newCartesian_withConstraint(self.X, dy, thre=self.params.enforce, verbose=self.params.verbose)
        else:
            self.X = self.IC.newCartesian(self.X, dy, self.params.verbose)

    def calcGradNorm(self):
        gradxc = self.IC.calcGradProj(self.X, self.gradx) if self.IC.haveConstraints() else self.gradx.copy()
        atomgrad = np.sqrt(np.sum((gradxc.reshape(-1,3))**2, axis=1))
        rms_gradient = np.sqrt(np.mean(atomgrad**2))
        max_gradient = np.max(atomgrad)
        return rms_gradient, max_gradient

    def rebuild_hessian(self):
        self.H = rebuild_hessian(self.IC, self.H0, self.X_hist, self.Gx_hist, self.params)

    def frequency_analysis(self, hessian, suffix, afterOpt):
        do_wigner = False
        if self.params.wigner:
            # Wigner sampling should only be performed on the final Hessian calculation of a run
            if self.params.hessian in ['last', 'first+last', 'each'] and afterOpt:
                do_wigner = True
            elif self.params.hessian in ['first', 'stop']:
                do_wigner = True
        if do_wigner:
            logger.info("Requesting %i samples from Wigner distribution.\n" % self.params.wigner)
        prefix = self.params.xyzout.replace("_optim.xyz", "").replace(".xyz", "")
        # Call the frequency analysis function with an input Hessian, with most arguments populated from self.params
        frequency_analysis(self.X, hessian, self.molecule.elem, energy=self.E, temperature=self.params.temperature, pressure=self.params.pressure, verbose=self.params.verbose,
                           outfnm='%s.vdata_%s' % (prefix, suffix), note='Iteration %i Energy % .8f%s' % (self.Iteration, self.E, ' (Optimized Structure)' if afterOpt else ''),
                           wigner=((self.params.wigner, os.path.join(self.dirname, 'wigner')) if do_wigner else None))

    def calcEnergyForce(self):
        """
        Calculate the energy and Cartesian gradients of the current structure.
        """
        # Check to confirm that the structure has nothing that would cause a cryptic error
        self.checkStructure()
        ### Calculate Energy and Gradient ###
        # Dictionary containing single point properties (energy, gradient)
        # For frequency calculations and multi-step jobs, the gradient from an existing
        # output file may be read in.
        spcalc = self.engine.calc(self.X, self.dirname, read_data=(self.Iteration==0))
        self.E = spcalc['energy']
        self.gradx = spcalc['gradient']
        # Calculate Hessian at the first step, or at each step if desired
        if self.params.hessian == 'each':
            # Hx is assumed to be the Cartesian Hessian at the current step.
            # Otherwise we use the variable name Hx0 to avoid almost certain confusion.
            self.Hx = calc_cartesian_hessian(self.X, self.molecule, self.engine, self.dirname, read_data=True, verbose=self.params.verbose)
            if self.params.frequency:
                self.frequency_analysis(self.Hx, 'iter%03i' % self.Iteration, False)
        elif self.Iteration == 0:
            if self.params.hessian in ['first', 'stop', 'first+last']:
                self.Hx0 = calc_cartesian_hessian(self.X, self.molecule, self.engine, self.dirname, read_data=True, verbose=self.params.verbose)
                if self.params.frequency:
                    self.frequency_analysis(self.Hx0, 'first', False)
                if self.params.hessian == 'stop':
                    logger.info("Exiting as requested after Hessian calculation.\n")
                    logger.info("Cartesian Hessian is stored in %s/hessian/hessian.txt.\n" % self.dirname)
                    raise HessianExit
                    # sys.exit(0)
            elif hasattr(self.params, 'hess_data') and self.Iteration == 0:
                self.Hx0 = self.params.hess_data.copy()
                if self.params.frequency:
                    self.frequency_analysis(self.Hx0, 'first', False)
                if self.Hx0.shape != (self.X.shape[0], self.X.shape[0]):
                    raise IOError('hess_data passed in via OptParams does not have the right shape')
            # self.Hx = self.Hx0.copy()
        # Add new Cartesian coordinates, energies, and gradients to history
        self.progress.xyzs.append(self.X.reshape(-1,3) * bohr2ang)
        self.progress.qm_energies.append(self.E)
        self.progress.qm_grads.append(self.gradx.copy())
        self.progress.comms.append('Iteration %i Energy % .8f' % (self.Iteration, self.E))

    def prepareFirstStep(self):
        """
        After computing the initial set of energies and forces, carry out some preparatory tasks
        prior to entering the optimization loop.
        """
        # Initial internal coordinates (optimization variables) and internal gradient
        self.Y = self.IC.calculate(self.coords)
        self.G = self.IC.calcGrad(self.X, self.gradx).flatten()
        # Print initial iteration
        rms_gradient, max_gradient = self.calcGradNorm()
        msg = "Step %4i :" % self.Iteration
        logger.info(msg + " Gradient = %.3e/%.3e (rms/max) Energy = % .10f\n" % (rms_gradient, max_gradient, self.E))
        # Initial history
        self.X_hist = [self.X]
        self.Gx_hist = [self.gradx]
        # Initial Hessian
        if hasattr(self, 'Hx'):
            # Compute IC Hessian from Cartesian Hessian at the current step
            self.H0 = self.IC.calcHess(self.X, self.gradx, self.Hx)
        elif hasattr(self, 'Hx0'):
            # Compute IC Hessian from input Cartesian Hessian
            self.H0 = self.IC.calcHess(self.X, self.gradx, self.Hx0)
        else:
            # Form guess Hessian if initial Hessian is not provided
            self.H0 = self.IC.guess_hessian(self.coords)
        self.H = self.H0.copy()

    def SortedEigenvalues(self):
        Eig = sorted(np.linalg.eigh(self.H)[0])
        if self.params.transition and len(Eig) >= 12:
            # logger.info("Hessian Eigenvalues:  %.3e %.3e %.3e %.3e %.3e %.3e %.3e %.3e %.3e ... %.3e %.3e %.3e\n" %
            #             (Eig[0],Eig[1],Eig[2],Eig[3],Eig[4],Eig[5],Eig[6],Eig[7],Eig[8],Eig[-3],Eig[-2],Eig[-1]))
            logger.info("Hessian Eigenvalues:  % .3e % .3e % .3e % .3e % .3e % .3e % .3e\n" % (Eig[0],Eig[1],Eig[2],Eig[3],Eig[4],Eig[5],Eig[6])),
            logger.info("% .3e % .3e % .3e % .3e % .3e    .....   % .3e % .3e % .3e\n" % (Eig[7],Eig[8],Eig[9],Eig[10],Eig[11],Eig[-3],Eig[-2],Eig[-1])),
        elif len(Eig) >= 6:
            logger.info("Hessian Eigenvalues: %.5e %.5e %.5e ... %.5e %.5e %.5e\n" % (Eig[0],Eig[1],Eig[2],Eig[-3],Eig[-2],Eig[-1]))
        else:
            logger.info("Hessian Eigenvalues: " + ' '.join("%.5e" % i for i in Eig) + '\n')
        return Eig

    def step(self):
        """
        Perform one step of the optimization.
        """
        params = self.params
        if np.isnan(self.G).any():
            raise RuntimeError("Gradient contains nan - check output and temp-files for possible errors")
        if np.isnan(self.H).any():
            raise RuntimeError("Hessian contains nan - check output and temp-files for possible errors")
        self.Iteration += 1
        if (self.Iteration%5) == 0:
            self.engine.clearCalcs()
            self.IC.clearCache()

        # At the start of the loop, the optimization variables, function value, gradient and Hessian are known.
        # (i.e. self.Y, self.E, self.G, self.H)
        if params.verbose: self.IC.printRotations(self.X)
        Eig = self.SortedEigenvalues()
        Emin = Eig[0].real
        if params.transition:
            v0 = 1.0
        elif Emin < params.epsilon:
            v0 = params.epsilon-Emin
        else:
            v0 = 0.0
        # Are we far from constraint satisfaction?
        self.farConstraints = self.IC.haveConstraints() and self.IC.maxConstraintViolation(self.X) > 1e-1
        ### OBTAIN AN OPTIMIZATION STEP ###
        # The trust radius is to be computed in Cartesian coordinates.
        # First take a full-size optimization step
        if params.verbose: logger.info("  Optimizer.step : Attempting full-size optimization step\n")
        dy, _, __ = self.get_delta_prime(v0, verbose=self.params.verbose)
        # Internal coordinate step size
        inorm = np.linalg.norm(dy)
        # Cartesian coordinate step size
        self.cnorm = self.get_cartesian_norm(dy)
        # If the full-size step is within the trust radius, then call get_delta_prime again with diagnostic messages if needed
        if (self.params.verbose >= 2 and self.params.verbose < 4 and self.cnorm <= 1.1*self.trust):
            self.get_delta_prime(v0, verbose=self.params.verbose+2)
        if params.verbose: logger.info("  Optimizer.step : Internal-step: %.4f Cartesian-step: %.4f Trust-radius: %.4f\n" % (inorm, self.cnorm, self.trust))
        # If the step is above the trust radius in Cartesian coordinates, then
        # do the following to reduce the step length:
        if self.cnorm > 1.1 * self.trust:
            # This is the function f(inorm) = cnorm-target that we find a root
            # for obtaining a step with the desired Cartesian step size.
            froot = self.createFroot(v0)
            froot.stores[inorm] = self.cnorm
            ### Find the internal coordinate norm that matches the desired Cartesian coordinate norm
            if params.verbose: logger.info("  Optimizer.step : Using Brent algorithm to target Cartesian trust radius\n")
            iopt = brent_wiki(froot.evaluate, 0.0, inorm, self.trust, cvg=0.1, obj=froot, verbose=params.verbose)
            if froot.brentFailed and froot.stored_arg is not None:
                # If Brent fails but we obtained an IC step that is smaller than the Cartesian trust radius, use it
                if params.verbose: logger.info("  Optimizer.step : \x1b[93mUsing stored solution at %.3e\x1b[0m\n" % froot.stored_val)
                iopt = froot.stored_arg
            elif self.IC.bork:
                # Decrease the target Cartesian step size and try again
                for i in range(3):
                    froot.target /= 2
                    if params.verbose: logger.info("  Optimizer.step : \x1b[93mReducing target to %.3e\x1b[0m\n" % froot.target)
                    froot.above_flag = True # Stop at any valid step between current target step size and trust radius
                    iopt = brent_wiki(froot.evaluate, 0.0, iopt, froot.target, cvg=0.1, verbose=params.verbose)
                    if not self.IC.bork: break
            LastForce = self.ForceRebuild
            self.ForceRebuild = False
            if self.IC.bork:
                logger.info("\x1b[91mInverse iteration for Cartesians failed\x1b[0m\n")
                # This variable is added because IC.bork is unset later.
                self.ForceRebuild = True
            else:
                if params.verbose: logger.info("  Optimizer.step : \x1b[93mBrent algorithm requires %i evaluations\x1b[0m\n" % froot.counter)
            ##### If IC failed to produce valid Cartesian step, it is "borked" and we need to rebuild it.
            if self.ForceRebuild:
                # Force a rebuild of the coordinate system and skip the energy / gradient and evaluation steps.
                if LastForce:
                    logger.warning("\x1b[1;91mFailed twice in a row to rebuild the coordinate system; continuing in Cartesian coordinates\x1b[0m\n")
                self.checkCoordinateSystem(recover=True, cartesian=LastForce)
                logger.info("\x1b[1;93mSkipping optimization step\x1b[0m\n")
                self.Iteration -= 1
                self.state = OPT_STATE.SKIP_EVALUATION
                return
            ##### End Rebuild
            # Finally, take an internal coordinate step of the desired length.
            dy, _ = self.trust_step(iopt, v0, verbose=(self.params.verbose+1 if self.params.verbose >= 2 else 0))
            self.cnorm = self.get_cartesian_norm(dy)
        ### DONE OBTAINING THE STEP ###
        if isinstance(self.IC, PrimitiveInternalCoordinates):
            idx = np.argmax(np.abs(dy))
            iunit = np.zeros_like(dy)
            iunit[idx] = 1.0
            self.prim_msg = "Along %s %.3f" % (self.IC.Internals[idx], np.dot(dy/np.linalg.norm(dy), iunit))
        ### These quantities, computed previously, are no longer used.
        # Dot product of the gradient with the step direction
        # Dot = -np.dot(dy/np.linalg.norm(dy), self.G/np.linalg.norm(self.G))
        # Whether the Cartesian norm comes close to the trust radius
        # bump = cnorm > 0.8 * self.trust
        ### Before updating any of our variables, copy current variables to "previous"
        self.Yprev = self.Y.copy()
        self.Xprev = self.X.copy()
        self.Gxprev = self.gradx.copy()
        self.Gprev = self.G.copy()
        self.Eprev = self.E
        ### Update the Internal Coordinates ###
        X0 = self.X.copy()
        self.newCartesian(dy)
        ## The "actual" dy may be different from the one passed to newCartesian(),
        ## for example if we enforce constraints or don't get the step we expect.
        dy = self.IC.calcDiff(self.X, X0)
        # dyp = self.IC.Prims.calcDiff(self.X, X0)
        # print("Actual dy:", dy)
        self.Y += dy
        self.expect = flat(0.5*multi_dot([row(dy),self.H,col(dy)]))[0] + np.dot(dy,self.G)
        self.state = OPT_STATE.NEEDS_EVALUATION

    def evaluateStep(self):
        ### At this point, the state should be NEEDS_EVALUATION
        assert self.state == OPT_STATE.NEEDS_EVALUATION
        # Shorthand for self.params
        params = self.params
        # Write current optimization trajectory to file
        if self.params.xyzout is not None: self.progress.write(self.params.xyzout)
        if self.params.qdata is not None: self.progress.write(self.params.qdata, ftype='qdata')
        # Project out the degrees of freedom that are constrained
        rms_gradient, max_gradient = self.calcGradNorm()
        rms_displacement, max_displacement = calc_drms_dmax(self.X, self.Xprev)
        # The ratio of the actual energy change to the expected change
        Quality = (self.E-self.Eprev)/self.expect
        colors = {}
        colors['quality'] = "\x1b[0m"
        # 2020-03-10: Step quality thresholds are hard-coded here.
        # At the moment, no need to set them as variables.
        if params.transition:
            if Quality > 0.8 and Quality < 1.2: step_state = StepState.Good
            elif Quality > 0.5 and Quality < 1.5: step_state = StepState.Okay
            elif Quality > 0.0 and Quality < 2.0: step_state = StepState.Poor
            else:
                colors['energy'] = "\x1b[91m"
                colors['quality'] = "\x1b[91m"
                step_state = StepState.Reject
        else:
            if Quality > 0.75: step_state = StepState.Good
            elif Quality > 0.25: step_state = StepState.Okay
            elif Quality > 0.0: step_state = StepState.Poor
            else:
                colors['energy'] = "\x1b[91m"
                colors['quality'] = "\x1b[91m"
                step_state = StepState.Poor if Quality > -1.0 else StepState.Reject
        # Check convergence criteria
        Converged_energy = np.abs(self.E-self.Eprev) < params.Convergence_energy
        Converged_grms = rms_gradient < params.Convergence_grms
        Converged_gmax = max_gradient < params.Convergence_gmax
        Converged_drms = rms_displacement < params.Convergence_drms
        Converged_dmax = max_displacement < params.Convergence_dmax
        if 'energy' not in colors: colors['energy'] = "\x1b[92m" if Converged_energy else "\x1b[0m"
        colors['grms'] = "\x1b[92m" if Converged_grms else "\x1b[0m"
        colors['gmax'] = "\x1b[92m" if Converged_gmax else "\x1b[0m"
        colors['drms'] = "\x1b[92m" if Converged_drms else "\x1b[0m"
        colors['dmax'] = "\x1b[92m" if Converged_dmax else "\x1b[0m"
        # Molpro defaults for convergence
        Converged_molpro_gmax = max_gradient < params.Convergence_molpro_gmax
        Converged_molpro_dmax = max_displacement < params.Convergence_molpro_dmax
        self.conSatisfied = not self.IC.haveConstraints() or self.IC.maxConstraintViolation(self.X) < 1e-2
        # Print status
        msg = "Step %4i :" % self.Iteration
        msg += " Displace = %s%.3e\x1b[0m/%s%.3e\x1b[0m (rms/max)" % (colors['drms'], rms_displacement, colors['dmax'], max_displacement)
        msg += " Trust = %.3e (%s)" % (self.trust, self.trustprint)
        msg += " Grad%s = %s%.3e\x1b[0m/%s%.3e\x1b[0m (rms/max)" % ("_T" if self.IC.haveConstraints() else "", colors['grms'], rms_gradient, colors['gmax'], max_gradient)
        logger.info(msg + " E (change) = % .10f (%s%+.3e\x1b[0m) Quality = %s%.3f\x1b[0m" % (self.E, colors['energy'], self.E-self.Eprev, colors['quality'], Quality) + "\n")

        if self.IC is not None and self.IC.haveConstraints():
            self.IC.printConstraints(self.X, thre=1e-3)
        if isinstance(self.IC, PrimitiveInternalCoordinates):
            logger.info(self.prim_msg + '\n')

        ### Check convergence criteria ###
        if Converged_energy and Converged_grms and Converged_drms and Converged_gmax and Converged_dmax and self.conSatisfied:
            self.SortedEigenvalues()
            logger.info("Converged! =D\n")
            self.state = OPT_STATE.CONVERGED
            return

        if self.Iteration > params.maxiter:
            self.SortedEigenvalues()
            logger.info("Maximum iterations reached (%i); increase --maxiter for more\n" % params.maxiter)
            self.state = OPT_STATE.FAILED
            return

        if params.qccnv and Converged_grms and (Converged_drms or Converged_energy) and self.conSatisfied:
            self.SortedEigenvalues()
            logger.info("Converged! (Q-Chem style criteria requires grms and either drms or energy)\n")
            self.state = OPT_STATE.CONVERGED
            return

        if params.molcnv and Converged_molpro_gmax and (Converged_molpro_dmax or Converged_energy) and self.conSatisfied:
            self.SortedEigenvalues()
            logger.info("Converged! (Molpro style criteria requires gmax and either dmax or energy)\nThis is approximate since convergence checks are done in cartesian coordinates.\n")
            self.state = OPT_STATE.CONVERGED
            return

        assert self.state == OPT_STATE.NEEDS_EVALUATION

        ### Adjust Trust Radius and/or Reject Step ###
        prev_trust = self.trust
        if step_state in (StepState.Poor, StepState.Reject):
            new_trust = max(params.tmin, min(self.trust, self.cnorm)/2)
            self.trustprint = "\x1b[91m-\x1b[0m" if new_trust < self.trust else "="
            self.trust = new_trust
        elif step_state == StepState.Good:
            new_trust = min(params.tmax, np.sqrt(2)*self.trust)
            self.trustprint = "\x1b[92m+\x1b[0m" if new_trust > self.trust else "="
            self.trust = new_trust
        elif step_state == StepState.Okay:
            self.trustprint = "="

        if step_state == StepState.Reject:
            if prev_trust <= params.thre_rj:
                logger.info("\x1b[93mNot rejecting step - trust below %.3e\x1b[0m\n" % params.thre_rj)
            elif (not params.transition) and self.E < self.Eprev:
                logger.info("\x1b[93mNot rejecting step - energy decreases during minimization\x1b[0m\n")
            elif self.farConstraints:
                logger.info("\x1b[93mNot rejecting step - far from constraint satisfaction\x1b[0m\n")
            else:
                logger.info("\x1b[93mStep Is Rejected\x1b[0m\n")
                self.trustprint = "\x1b[1;91mx\x1b[0m"
                self.Y = self.Yprev.copy()
                self.X = self.Xprev.copy()
                self.gradx = self.Gxprev.copy()
                self.G = self.Gprev.copy()
                self.E = self.Eprev
                return

        # Append steps to history (for rebuilding Hessian)
        self.X_hist.append(self.X)
        self.Gx_hist.append(self.gradx)

        ### Rebuild Coordinate System if Necessary ###
        UpdateHessian = (not self.params.hessian == 'each')
        if self.IC.bork:
            logger.info("Failed inverse iteration - checking coordinate system\n")
            self.checkCoordinateSystem(recover=True)
            UpdateHessian = False
        elif self.CoordCounter == (params.check - 1):
            logger.info("Checking coordinate system as requested every %i cycles\n" % params.check)
            if self.checkCoordinateSystem(): UpdateHessian = False
        else:
            self.CoordCounter += 1
        # Check for large rotations (debugging purposes)
        if self.params.verbose >= 1: self.IC.largeRots()
        # Check for large rotations in linear molecules
        if self.IC.linearRotCheck():
            logger.info("Large rotations in linear molecules - refreshing Rotator reference points and DLC vectors\n")
            self.refreshCoordinates()
            UpdateHessian = False
        self.G = self.IC.calcGrad(self.X, self.gradx).flatten()

        ### Update the Hessian ###
        if UpdateHessian:
            self.UpdateHessian()
        if hasattr(self, 'Hx'):
            self.H = self.IC.calcHess(self.X, self.gradx, self.Hx)
        # Then it's on to the next loop iteration!
        return

    def UpdateHessian(self):
        params = self.params

        if params.transition:
            ts_bfgs = False
            if ts_bfgs: # pragma: no cover
                logger.info("TS-BFGS Hessian update\n")
                # yk = Dg; dk = Dy
                dk = col(self.Y - self.Yprev)
                yk = col(self.G - self.Gprev)
                jk = yk - np.dot(self.H, dk)
                B = force_positive_definite(self.H)
                # Scalar 1: dk^T |Bk| dk
                s1 = multi_dot([dk.T, B, dk])
                # Scalar 2: (yk^T dk)^2 + (dk^T |Bk| dk)^2
                s2 = np.dot(yk.T, dk)**2 + s1**2
                # Vector quantities
                v2 = np.dot(yk.T, dk)*yk + s1*np.dot(B, dk)
                uk = v2/s2
                Ek = np.dot(jk, uk.T) + np.dot(uk, jk.T) + np.dot(jk.T, dk) * np.dot(uk, uk.T)
                self.H += Ek
            else:
                Dy   = col(self.Y - self.Yprev)
                Dg   = col(self.G - self.Gprev)
                # Murtagh-Sargent-Powell update
                Xi = Dg - np.dot(self.H,Dy)
                # ndy2 = np.dot(Dy.T,Dy)
                dH_MS = np.dot(Xi, Xi.T)/np.dot(Dy.T, Xi)
                dH_P = np.dot(Xi, Dy.T) + np.dot(Dy, Xi.T) - np.dot(Dy, Dy.T)*np.dot(Xi.T, Dy)/np.dot(Dy.T, Dy)
                dH_P /= np.dot(Dy.T, Dy)
                phi = 1.0 - np.dot(Dy.T,Xi)**2/(np.dot(Dy.T,Dy)*np.dot(Xi.T,Xi))
                # phi = 1.0
                self.H += (1.0-phi)*dH_MS + phi*dH_P
                if params.verbose:
                    logger.info("Hessian update: %.5f Powell + %.5f Murtagh-Sargent\n" % (phi, 1.0-phi))
        else:
            Dy   = col(self.Y - self.Yprev)
            Dg   = col(self.G - self.Gprev)
            # Catch some abnormal cases of extremely small changes.
            if np.linalg.norm(Dg) < 1e-6: return
            if np.linalg.norm(Dy) < 1e-6: return
            # BFGS Hessian update
            Mat1 = np.dot(Dg,Dg.T)/np.dot(Dg.T,Dy)[0,0]
            Mat2 = np.dot(np.dot(self.H,Dy), np.dot(self.H,Dy).T)/multi_dot([Dy.T,self.H,Dy])[0,0]
            Eig = np.linalg.eigh(self.H)[0]
            Eig.sort()
            ndy = np.array(Dy).flatten()/np.linalg.norm(np.array(Dy))
            ndg = np.array(Dg).flatten()/np.linalg.norm(np.array(Dg))
            nhdy = np.dot(self.H,Dy).flatten()/np.linalg.norm(np.dot(self.H,Dy))
            if params.verbose:
                msg = "Denoms: %.3e %.3e" % (np.dot(Dg.T,Dy)[0,0], multi_dot((Dy.T,self.H,Dy))[0,0])
                msg +=" Dots: %.3e %.3e" % (np.dot(ndg, ndy), np.dot(ndy, nhdy))
            #H1 = H.copy()
            self.H += Mat1-Mat2
            Eig1 = np.linalg.eigh(self.H)[0]
            Eig1.sort()
            if params.verbose:
                msg += " Eig-ratios: %.5e ... %.5e" % (np.min(Eig1)/np.min(Eig), np.max(Eig1)/np.max(Eig))
                logger.info(msg+'\n')
            if np.min(Eig1) <= params.epsilon and params.reset:
                logger.info("Eigenvalues below %.4e (%.4e) - returning guess\n" % (params.epsilon, np.min(Eig1)))
                self.H = self.IC.guess_hessian(self.coords)

    def optimizeGeometry(self):
        """
        High-level optimization loop.
        This allows calcEnergyForce() to be separated from the rest of the codes
        """
        self.calcEnergyForce()
        self.prepareFirstStep()
        while self.state not in [OPT_STATE.CONVERGED, OPT_STATE.FAILED]:
            self.step()
            if self.state == OPT_STATE.NEEDS_EVALUATION:
                self.calcEnergyForce()
                self.evaluateStep()
        if self.state == OPT_STATE.FAILED:
            raise GeomOptNotConvergedError("Optimizer.optimizeGeometry() failed to converge.")
        # If we want to save the Hessian used by the optimizer (in Cartesian coordinates)
        if self.params.write_cart_hess:
            # One last Hessian update before writing it out
            self.UpdateHessian()
            logger.info("Saving current approximate Hessian (Cartesian coordinates) to %s" % self.params.write_cart_hess)
            Hx = self.IC.calcHessCart(self.X, self.G, self.H)
            np.savetxt(self.params.write_cart_hess, Hx, fmt='% 14.10f')
        if self.params.hessian in ['last', 'first+last', 'each']:
            Hx = calc_cartesian_hessian(self.X, self.molecule, self.engine, self.dirname, read_data=False, verbose=self.params.verbose)
            if self.params.frequency:
                self.frequency_analysis(Hx, 'last', True)
        return self.progress

    def checkStructure(self):
        """
        A function that checks for problematic structures and throws an error before
        calling any QC method.
        """
        # Check for three consecutive atoms in torsion angle becoming linear
        torsion_constraint_linear_angles = self.IC.torsionConstraintLinearAngles(self.X)
        if torsion_constraint_linear_angles:
            errorStr = "> Atoms Angle\n"
            for key, val in torsion_constraint_linear_angles.items():
                errorStr += "> %i-%i-%i %6.2f\n" % (key[0]+1, key[1]+1, key[2]+1, val)
            raise LinearTorsionError("A constrained torsion has three consecutive atoms\n"
                                     "forming a nearly linear angle, making the torsion angle poorly defined.\n"+errorStr)

class OPT_STATE(object):
    """ This describes the state of an OptObject during the optimization process
    """
    NEEDS_EVALUATION = 0  # convergence has not been evaluated -> calcualte Energy, Forces
    SKIP_EVALUATION  = 1  # We know this is not yet converged -> skip Energy
    CONVERGED        = 2
    FAILED           = 3  # optimization failed with no recovery option

class StepState(object):
    """ This describes the state of an OptObject during the optimization process
    """
    Reject  = 0 # Reject the step
    Poor    = 1 # Poor step; decrease the trust radius down to the lower limit.
    Okay    = 2 # Okay step; do not change the trust radius.
    Good    = 3 # Good step; increase the trust radius up to the limit.

def Optimize(coords, molecule, IC, engine, dirname, params):
    """
    Optimize the geometry of a molecule. This function used to contain the whole
    optimization loop, which has since been moved to the Optimizer() class;
    now a wrapper and kept for compatibility.

    Parameters
    ----------
    coords : np.ndarray
        Nx3 array of Cartesian coordinates in atomic units
    molecule : Molecule
        Molecule object
    IC : InternalCoordinates
        Object describing the internal coordinate system
    engine : Engine
        Object containing methods for calculating energy and gradient
    dirname : str
        Directory name for files to be written
    params : OptParams object
        Contains optimization parameters (really just a struct)
    hessian : np.ndarray, optional
        3Nx3N array of Cartesian Hessian of initial structure

    Returns
    -------
    progress: Molecule
        A molecule object for opt trajectory and energies
    """
    optimizer = Optimizer(coords, molecule, IC, engine, dirname, params)
    return optimizer.optimizeGeometry()

def run_optimizer(**kwargs):
    """
    Run geometry optimization, constrained optimization, or
    constrained scan job given arguments from command line.
    """
    #==============================#
    #|   Log file configuration   |#
    #==============================#
    # By default, output should be written to <args.prefix>.log and also printed to the terminal.
    # This behavior may be changed by editing the log.ini file.
    # Output will only be written to log files after the 'logConfig' line is called!
    if kwargs.get('logIni') is None:
        import geometric.optimize
        logIni = pkg_resources.resource_filename(geometric.optimize.__name__, 'config/log.ini')
    else:
        logIni = kwargs.get('logIni')
    logfilename = kwargs.get('prefix')
    # Input file for optimization; QC input file or OpenMM .xml file
    inputf = kwargs.get('input')
    verbose = kwargs.get('verbose', False)
    # Get calculation prefix and temporary directory name
    arg_prefix = kwargs.get('prefix', None) #prefix for output file and temporary directory
    prefix = arg_prefix if arg_prefix is not None else os.path.splitext(inputf)[0]
    logfilename = prefix + ".log"
    # Create a backup if the log file already exists
    backed_up = bak(logfilename)
    import logging.config
    logging.config.fileConfig(logIni,defaults={'logfilename': logfilename},disable_existing_loggers=False)
    #==============================#
    #| End log file configuration |#
    #==============================#

    import geometric
    logger.info('geometric-optimize called with the following command line:\n')
    logger.info(' '.join(sys.argv)+'\n')
    print_logo(logger)
    logger.info('-=# \x1b[1;94m geomeTRIC started. Version: %s \x1b[0m #=-\n' % geometric.__version__)
    if backed_up:
        logger.info('Backed up existing log file: %s -> %s\n' % (logfilename, os.path.basename(backed_up)))

    t0 = time.time()

    # Create the params object, containing data to be passed into the optimizer
    params = OptParams(**kwargs)
    params.printInfo()

    # Create "dirname" folder for writing
    dirname = prefix+".tmp"
    if not os.path.exists(dirname):
        os.makedirs(dirname)
    kwargs['dirname'] = dirname

    # Get the Molecule and engine objects needed for optimization
    M, engine = get_molecule_engine(**kwargs)

    # Create Work Queue object
    if kwargs.get('port', 0):
        logger.info("Creating Work Queue object for distributed Hessian calculation\n")
        createWorkQueue(kwargs['port'], debug=verbose>1)

    # Get initial coordinates in bohr
    coords = M.xyzs[0].flatten() * ang2bohr

    # Read in the constraints
    constraints = kwargs.get('constraints', None) #Constraint input file (optional)

    if constraints is not None:
        Cons, CVals = parse_constraints(M, open(constraints).read())
    else:
        Cons = None
        CVals = None

    #=========================================#
    #| Set up the internal coordinate system |#
    #=========================================#
    # First item in tuple: The class to be initialized
    # Second item in tuple: Whether to connect nonbonded fragments
    # Third item in tuple: Whether to throw in all Cartesians (no effect if second item is True)
    CoordSysDict = {'cart':(CartesianCoordinates, False, False),
                    'prim':(PrimitiveInternalCoordinates, True, False),
                    'dlc':(DelocalizedInternalCoordinates, True, False),
                    'hdlc':(DelocalizedInternalCoordinates, False, True),
                    'tric-p':(PrimitiveInternalCoordinates, False, False),
                    'tric':(DelocalizedInternalCoordinates, False, False)}
    coordsys = kwargs.get('coordsys', 'tric')
    CoordClass, connect, addcart = CoordSysDict[coordsys.lower()]

    IC = CoordClass(M, build=True, connect=connect, addcart=addcart, constraints=Cons, cvals=CVals[0] if CVals is not None else None,
                    conmethod=params.conmethod)
    #========================================#
    #| End internal coordinate system setup |#
    #========================================#

    # Auxiliary functions (will not do optimization):
    displace = kwargs.get('displace', False) # Write out the displacements of the coordinates.
    if displace:
        write_displacements(coords, M, IC, dirname, verbose)
        return

    fdcheck = kwargs.get('fdcheck', False) # Check internal coordinate gradients using finite difference..
    if fdcheck:
        IC.Prims.checkFiniteDifferenceGrad(coords)
        IC.Prims.checkFiniteDifferenceHess(coords)
        check_internal_grad(coords, M, IC.Prims, engine, dirname, verbose)
        check_internal_hess(coords, M, IC.Prims, engine, dirname, verbose)
        return

    # Print out information about the coordinate system
    if isinstance(IC, CartesianCoordinates):
        logger.info("%i Cartesian coordinates being used\n" % (3*M.na))
    else:
        logger.info("%i internal coordinates being used (instead of %i Cartesians)\n" % (len(IC.Internals), 3*M.na))
    logger.info(IC)
    logger.info("\n")

    if Cons is None:
        # Run a standard geometry optimization
        params.xyzout = prefix+"_optim.xyz"
        progress = Optimize(coords, M, IC, engine, dirname, params)
    else:
        # Run a single constrained geometry optimization or scan over a grid of values
        if isinstance(IC, (CartesianCoordinates, PrimitiveInternalCoordinates)):
            raise RuntimeError("Constraints only work with delocalized internal coordinates")
        Mfinal = None
        for ic, CVal in enumerate(CVals):
            if len(CVals) > 1:
                logger.info("---=== Scan %i/%i : Constrained Optimization ===---\n" % (ic+1, len(CVals)))
            IC = CoordClass(M, build=True, connect=connect, addcart=addcart, constraints=Cons, cvals=CVal, conmethod=params.conmethod)
            IC.printConstraints(coords, thre=-1)
            if len(CVals) > 1:
                params.xyzout = prefix+"_scan-%03i.xyz" % (ic+1)
                # In the special case of a constraint scan, we write out multiple qdata.txt files
                if params.qdata is not None: params.qdata = 'qdata_scan-%03i.txt' % (ic+1)
            else:
                params.xyzout = prefix+"_optim.xyz"
            if ic == 0:
                progress = Optimize(coords, M, IC, engine, dirname, params)
            else:
                progress += Optimize(coords, M, IC, engine, dirname, params)
            # update the structure for next optimization in SCAN (by CNH)
            M.xyzs[0] = progress.xyzs[-1]
            coords = progress.xyzs[-1].flatten() * ang2bohr
            if Mfinal:
                Mfinal += progress[-1]
            else:
                Mfinal = progress[-1]
            cNames = IC.getConstraintNames()
            cVals = IC.getConstraintTargetVals()
            comment = ', '.join(["%s = %.2f" % (cName, cVal) for cName, cVal in zip(cNames, cVals)])
            Mfinal.comms[-1] = "Scan Cycle %i/%i ; %s ; %s" % (ic+1, len(CVals), comment, progress.comms[-1])
            #print
        if len(CVals) > 1:
            Mfinal.write('scan-final.xyz')
            if params.qdata is not None: Mfinal.write('qdata-final.txt')
    print_citation(logger)
    logger.info("Time elapsed since start of run_optimizer: %.3f seconds\n" % (time.time()-t0))
    return progress

def main(): # pragma: no cover
    # Read user input (look in params.py for full list of options).
    # args is a dictionary containing only user-specified arguments
    # (i.e. keys without provided values are removed.)
    args = parse_optimizer_args(sys.argv[1:])

    # Run the optimizer.
    try:
        run_optimizer(**args)
    except EngineError:
        logger.info("EngineError:\n" + traceback.format_exc())
        sys.exit(51)
    except GeomOptNotConvergedError:
        logger.info("Geometry Converge Failed Error:\n" + traceback.format_exc())
        sys.exit(50)
    except GeomOptStructureError:
        logger.info("Structure Error:\n" + traceback.format_exc())
        sys.exit(50)
    except HessianExit:
        logger.info("Exiting normally.\n")
        sys.exit(0)
    except:
        logger.info("Unknown Error:\n" + traceback.format_exc())
        raise

if __name__ == "__main__":
    main()