# Fit to multiple relative binding free energy edges
import sys
from argparse import ArgumentParser
import jax
from jax import numpy as jnp
import numpy as np
import datetime
import timemachine

# forcefield handlers
from ff import Forcefield
from ff.handlers.serialize import serialize_handlers
from ff.handlers.deserialize import deserialize_handlers
from ff.handlers.nonbonded import AM1CCCHandler, LennardJonesHandler

# free energy classes
from fe.free_energy import RelativeFreeEnergy, construct_lambda_schedule
from fe.estimator import SimulationResult
from fe.model import RBFEModel
from fe.loss import pseudo_huber_loss  # , l1_loss, flat_bottom_loss

# MD initialization
from md import builders

# parallelization across multiple GPUs
from parallel.client import CUDAPoolClient, GRPCClient
from parallel.utils import get_gpu_count

from collections import namedtuple

from pickle import load

from optimize.step import truncated_step

from typing import Tuple, Dict, List, Union

from pathlib import Path
from time import time

array = Union[np.array, jnp.array]

Handler = Union[AM1CCCHandler, LennardJonesHandler]  # TODO: relax this assumption

NUM_GPUS = get_gpu_count()

# how much MD to run, on how many GPUs
Configuration = namedtuple(
    'Configuration',
    ['num_complex_windows', 'num_solvent_windows', 'num_equil_steps', 'num_prod_steps'])

# define a couple configurations: one for quick tests, and one for production
production_configuration = Configuration(
    num_complex_windows=60,
    num_solvent_windows=60,
    num_equil_steps=10000,
    num_prod_steps=100000,
)

intermediate_configuration = Configuration(
    num_complex_windows=60,
    num_solvent_windows=60,
    num_equil_steps=10000,
    num_prod_steps=10000,
)

testing_configuration = Configuration(
    num_complex_windows=10,
    num_solvent_windows=10,
    num_equil_steps=1000,
    num_prod_steps=1000,
)

# TODO: rename this to something more descriptive than "Configuration"...
#   want to distinguish later between an "RBFE configuration"
#       (which describes the computation of a single edge)
#   and a "training configuration"
#       (which describes the overall training loop)

# locations relative to project root
root = Path(timemachine.__file__).parent
path_to_protein = str(root.joinpath('tests/data/hif2a_nowater_min.pdb'))

class ParameterUpdate:
    def __init__(self, before, after, gradient, update):
        self.before = before
        self.after = after
        self.gradient = gradient
        self.update = update

    # TODO: def __str__

    def save(self, fname='parameter_update.npz'):
        """save numpy arrays to fname"""
        print(f'saving parameter updates to {fname}')
        np.savez(
            file=fname,
            before=self.before,
            after=self.after,
            gradient=self.gradient,
            update=self.update,
        )

def _save_forcefield(fname, ff_params):
    with open(fname, 'w') as fh:
        fh.write(ff_params)

if __name__ == "__main__":
    default_output_path = f"results_{str(datetime.datetime.now())}"
    parser = ArgumentParser(description="Fit Forcefield parameters to hif2a")
    parser.add_argument("--num-gpus", default=None, type=int,
                        help=f"Number of GPUs to run against, defaults to {NUM_GPUS} if no hosts provided")
    parser.add_argument("--hosts", nargs="*", default=None, help="Hosts running GRPC worker to use for compute")
    parser.add_argument("--param-updates", default=1000, type=int, help="Number of updates for parameters")
    parser.add_argument("--seed", default=2021, type=int, help="Seed for shuffling ordering of transformations")
    parser.add_argument("--config", default="intermediate", choices=["intermediate", "production", "test"])

    parser.add_argument("--path_to_ff", default=str(root.joinpath('ff/params/smirnoff_1_1_0_ccc.py')))
    parser.add_argument("--path_to_edges", default="relative_transformations.pkl",
                        help="Path to pickle file containing list of RelativeFreeEnergy objects")
    parser.add_argument("--output_path", default=default_output_path, help="Path to output directory")
    # TODO: also make configurable: forces_to_refit, optimizer params, path_to_protein, path_to_protein_ff, ...
    args = parser.parse_args()

    # create path if it doesn't exist
    output_path = Path(args.output_path)
    output_path.mkdir(parents=True, exist_ok=True)
    print(f'output path: {output_path}')

    # xor num_gpus and hosts args
    args = parser.parse_args()
    if args.num_gpus is not None and args.hosts is not None:
        print("Unable to provide --num-gpus and --hosts together")
        sys.exit(1)

    # which force field components we'll refit
    forces_to_refit = [AM1CCCHandler, LennardJonesHandler]

    # how much computation to spend per refitting step
    configuration = None
    if args.config == "intermediate":  # goldilocks
        configuration = intermediate_configuration
    elif args.config == "test":
        configuration = testing_configuration  # a little
    elif args.config == "production":
        configuration = production_configuration  # a lot
    assert configuration is not None, "No configuration provided"

    if not args.hosts:
        num_gpus = args.num_gpus
        if num_gpus is None:
            num_gpus = NUM_GPUS
        # set up multi-GPU client
        client = CUDAPoolClient(max_workers=num_gpus)
    else:
        # Setup GRPC client
        client = GRPCClient(hosts=args.hosts)
    client.verify()

    # load and construct forcefield
    with open(args.path_to_ff) as f:
        ff_handlers = deserialize_handlers(f.read())

    forcefield = Forcefield(ff_handlers)

    # load pre-defined collection of relative transformations
    with open(args.path_to_edges, 'rb') as f:
        relative_transformations: List[RelativeFreeEnergy] = load(f)

    # update the forcefield parameters for a few steps, each step informed by a single free energy calculation

    # compute and save the sequence of relative_transformation indices
    num_epochs = int(np.ceil(args.param_updates / len(relative_transformations)))
    np.random.seed(args.seed)
    step_inds = []
    for epoch in range(num_epochs):
        inds = np.arange(len(relative_transformations))
        np.random.shuffle(inds)
        step_inds.append(inds)
    step_inds = np.hstack(step_inds)[:args.param_updates]

    np.save(output_path.joinpath('step_indices.npy'), step_inds)

    # build the complex system
    complex_system, complex_coords, _, _, complex_box, _ = builders.build_protein_system(
        path_to_protein)
    # TODO: optimize box
    complex_box += np.eye(3) * 0.1  # BFGS this later

    # build the water system
    solvent_system, solvent_coords, solvent_box, _ = builders.build_water_system(4.0)
    # TODO: optimize box
    solvent_box += np.eye(3) * 0.1  # BFGS this later

    # note: "complex" means "protein + solvent"
    binding_model = RBFEModel(
        client=client,
        ff=forcefield,
        complex_system=complex_system,
        complex_coords=complex_coords,
        complex_box=complex_box,
        complex_schedule=construct_lambda_schedule(configuration.num_complex_windows),
        solvent_system=solvent_system,
        solvent_coords=solvent_coords,
        solvent_box=solvent_box,
        solvent_schedule=construct_lambda_schedule(configuration.num_solvent_windows),
        equil_steps=configuration.num_equil_steps,
        prod_steps=configuration.num_prod_steps,
    )

    def loss_fxn(ff_params, mol_a, mol_b, core, label_ddG, callback=None):
        pred_ddG = binding_model.predict(ff_params, mol_a, mol_b, core, callback)
        return pseudo_huber_loss(pred_ddG - label_ddG)

    # TODO: how to get intermediate results from the computational pipeline encapsulated in binding_model.loss ?
    #   e.g. stage_results, and further diagnostic information
    #   * x trajectories,
    #   * d U / d parameters trajectories,
    #   * matrix of U(x; lambda) for all x, lambda
    #   * the deltaG pred
    #   (proper way is probably something like has_aux=True https://jax.readthedocs.io/en/latest/jax.html#jax.value_and_grad)

    ordered_params = forcefield.get_ordered_params()
    ordered_handles = forcefield.get_ordered_handles()

    handle_types_being_optimized = [AM1CCCHandler, LennardJonesHandler]

    # TODO: move flatten into optimize.utils
    def flatten(params) -> Tuple[np.array, callable]:
        """Turn params dict into flat array, with an accompanying unflatten function

        TODO: note that the result is going to be in the order given by ordered_handles (filtered by presence in handle_types)
            rather than in the order they appear in handle_types_being_optimized

        TODO: maybe leave out the reference to handle_types_being optimized altogether

        TODO: does Jax have a pytree-based flatten / unflatten utility?
        """

        theta_list = []
        _shapes = dict()
        _handle_types = []

        for param, handle in zip(params, ordered_handles):
            assert handle.params.shape == param.shape
            key = type(handle)

            if key in handle_types_being_optimized:
                theta_list.append(param.flatten())
                _shapes[key] = param.shape
                _handle_types.append(key)

        theta = np.hstack(theta_list)

        def unflatten(theta: array) -> Dict[Handler, array]:
            params = dict()
            i = 0
            for key in _handle_types:
                shape = _shapes[key]
                num_params = int(np.prod(shape))
                params[key] = np.array(theta[i: i + num_params]).reshape(shape)
                i += num_params
            return params

        return theta, unflatten

    relative_improvement_bound = 0.8

    def _compute_step_lower_bound(loss, blown_up):
        """problem this addresses: on a small fraction of steps, the free energy estimate may be grossly unreliable
        away from target, typically indicating an instability was encountered.
        detect if this occurs, and don't allow a step.

        """
        if not blown_up:
            return loss * relative_improvement_bound
        else:
            return loss  # don't move!

    def _results_to_arrays(results: List[SimulationResult]):
        """each result object was constructed by SimulationResult(xs=xs, du_dls=full_du_dls, du_dps=grads)

        for each field, concatenate into an array
        """

        xs = np.array([r.xs for r in results])
        du_dls = np.array([r.du_dls for r in results])
        du_dps = np.array([r.du_dps for r in results])

        return xs, du_dls, du_dps

    def _blew_up(results: List[SimulationResult]):
        """if stddev(du_dls) for any window exceeded 1000 kJ/mol, don't trust result enough to take a step
        if du_dls contains any nans, don't trust result enough to take a step"""
        du_dls = _results_to_arrays(results)[1]

        # TODO: adjust this threshold a bit, move reliability calculations into fe/estimator.py or fe/model.py
        return np.isnan(du_dls).any() or (du_dls.std(1).max() > 1000)

    results_this_step = dict()  # {stage : result} pairs # TODO: proper type hint

    def save_in_memory_callback(results, stage):
        global results_this_step
        results_this_step[stage] = results
        print(f'collected {stage} results!')

    # in each optimizer step, look at one transformation from relative_transformations
    for step, rfe_ind in enumerate(step_inds):
        rfe = relative_transformations[rfe_ind]

        # compute a step, measuring total wall-time
        t0 = time()

        # TODO: perhaps update this to accept an rfe argument, instead of all of rfe's attributes as arguments
        loss, loss_grads = jax.value_and_grad(loss_fxn, argnums=0)(ordered_params, rfe.mol_a, rfe.mol_b, rfe.core,
                                                                   rfe.label, callback=save_in_memory_callback)
        print(f"at optimizer step {step}, loss={loss:.3f}")

        # check if it's probably okay to take an optimizer step on the basis of this result
        # TODO: move responsibility for returning error flags / simulation uncertainty estimates further upstream
        blown_up = _blew_up(results_this_step['complex']) or _blew_up(results_this_step['solvent'])

        # note: unflatten_grad and unflatten_theta have identical definitions for now
        flat_loss_grad, unflatten_grad = flatten(loss_grads)
        flat_theta, unflatten_theta = flatten(ordered_params)

        # based on current estimate of (loss, grad, and simulation stability), return a conservative step to take in parameter space
        theta_increment = truncated_step(flat_theta, loss, flat_loss_grad,
                                         step_lower_bound=_compute_step_lower_bound(loss, blown_up))
        param_increments = unflatten_theta(theta_increment)

        # for any parameter handler types being updated, update in place
        for handle in ordered_handles:
            handle_type = type(handle)
            if handle_type in param_increments:
                print(f'updating {handle_type.__name__}')

                # TODO: careful -- this must be a "+=" or "-=" not an "="!
                handle.params += param_increments[handle_type]

                increment = param_increments[handle_type]
                update_mask = increment != 0

                # TODO: replace with a function that knows what to report about each handle type
                print(
                    f'updated {int(np.sum(update_mask))} params by between {np.min(increment[update_mask]):.4f} and {np.max(increment[update_mask])}')

        t1 = time()
        elapsed = t1 - t0

        print(f'completed forcefield-updating step {step} in {elapsed:.3f} s !')

        # save du_dls snapshot
        path_to_du_dls = output_path.joinpath(f'du_dls_snapshot_{step}.npz')
        print(f'saving du_dl trajs to {path_to_du_dls}')
        du_dls_dict = dict()  # keywords here must be strings
        for stage, results in results_this_step.items():
            du_dls_dict[stage] = _results_to_arrays(results)[1]
        np.savez(path_to_du_dls, **du_dls_dict)

        # also save information about this step's parameter gradient and parameter update
        # results to npz
        path_to_npz = output_path.joinpath(f'theta_grad_loss_snapshot_{step}.npz')
        print(f'saving theta, grad, loss snapshot to {path_to_npz}')
        np.savez(
            path_to_npz,
            theta=np.array(flat_theta),
            grad=np.array(flat_loss_grad),
            loss=loss
        )

        # TODO: same for xs and du_dps snapshots

        # save updated forcefield .py files after every gradient step
        step_params = serialize_handlers(ff_handlers)
        # TODO: consider if there's a more modular way to keep track of ff updates
        _save_forcefield(output_path.joinpath("forcefield_checkpoint_{step}.py"), ff_params=step_params)


#!/usr/bin/env python
from __future__ import division
import numpy as np
import pandas as pd
import warnings
from .helpers import *

def analyze_chunk(data, subjgroup=None, subjname='Subject', listgroup=None, listname='List', analysis=None, analysis_type=None, pass_features=False, **kwargs):
    """
    Private function that groups data by subject/list number and performs analysis for a chunk of data.

    Parameters
    ----------
    data : Egg data object
        The data to be analyzed

    subjgroup : list of strings or ints
        String/int variables indicating how to group over subjects.  Must be
        the length of the number of subjects

    subjname : string
        Name of the subject grouping variable

    listgroup : list of strings or ints
        String/int variables indicating how to group over list.  Must be
        the length of the number of lists

    listname : string
        Name of the list grouping variable

    analysis : function
        This function analyzes data and returns it.

    pass_features : bool
        Logical indicating whether the analyses uses the features field of the Egg

    Returns
    ----------
    analyzed_data : Pandas DataFrame
        DataFrame containing the analysis results

    """
    # if no grouping, set default to iterate over each list independently
    subjgroup = subjgroup if subjgroup else data.pres.index.levels[0].values
    listgroup = listgroup if listgroup else data.pres.index.levels[1].values

    # create a dictionary for grouping
    subjdict = {subj : data.pres.index.levels[0].values[subj==np.array(subjgroup)] for subj in set(subjgroup)}
    # listdict = {lst : data.pres.index.levels[1].values[lst==np.array(listgroup)] for lst in set(listgroup)}

    # allow for lists of listgroup arguments
    if all(isinstance(el, list) for el in listgroup):
        listdict = [{lst : data.pres.index.levels[1].values[lst==np.array(listgrpsub)] for lst in set(listgrpsub)} for listgrpsub in listgroup]
    else:
        listdict = [{lst : data.pres.index.levels[1].values[lst==np.array(listgroup)] for lst in set(listgroup)} for subj in subjdict]

    # perform the analysis
    def perform_analysis(subj, lst):

        # get data slice for presentation and recall
        pres_slice = data.pres.loc[[(s,l) for s in subjdict[subj] for l in listdict[subj][lst] if all(~pd.isnull(data.pres.loc[(s,l)]))]]
        pres_slice.list_length = data.list_length

        rec_slice = data.rec.loc[[(s,l) for s in subjdict[subj] for l in listdict[subj][lst] if all(~pd.isnull(data.pres.loc[(s,l)]))]]

        # if features are need for analysis, get the features for this slice of data
        if pass_features:
            feature_slice = data.features.loc[[(s,l) for s in subjdict[subj] for l in listdict[subj][lst] if all(~pd.isnull(data.pres.loc[(s,l)]))]]

        # generate indices
        index = pd.MultiIndex.from_arrays([[subj],[lst]], names=[subjname, listname])

        # perform analysis for each data chunk
        if pass_features:
            return pd.DataFrame([analysis(pres_slice, rec_slice, feature_slice, data.dist_funcs, **kwargs)], index=index, columns=[feature for feature in feature_slice[0].as_matrix()[0].keys()])
        else:
            return pd.DataFrame([analysis(pres_slice, rec_slice, **kwargs)], index=index)

    # create list of chunks to process
    a=[]
    b=[]
    for subj in subjdict:
        for lst in listdict[0]:
            a.append(subj)
            b.append(lst)

    # handle parellel kwarg
    parallel=kwargs['parallel']
    del kwargs['parallel']

    # if we're running permutation tests, use multiprocessing
    if parallel==True:
        import multiprocessing
        from pathos.multiprocessing import ProcessingPool as Pool
        p = Pool(multiprocessing.cpu_count())
        analyzed_data = p.map(perform_analysis, a, b)
    else:
        analyzed_data = [perform_analysis(ai, bi) for ai,bi in zip(a,b)]

    # concatenate slices
    analyzed_data = pd.concat(analyzed_data)

    analyzed_data.attrs = {
        'analysis_type' : analysis_type,
        'list_length' : data.list_length
    }

    for key in kwargs:
        analyzed_data.attrs[key] = kwargs[key]

    return analyzed_data

# recall matrix
def recall_matrix(presented, recalled):
    """
    Computes recall matrix given list of presented and list of recalled words

    Parameters
    ----------
    presented : list of list of strings
      presentedWords are the words presented in the experiment, in order, grouped by list

    recalled : list of list of strings
      recalledWords are the words recalled by the subject, in order, grouped by list

    Returns
    ----------
    recall_matrix : list of lists of ints
      each integer represents the presentation position of the recalled word in a given list in order of recall
      0s represent recalled words not presented
      negative ints represent words recalled from previous lists

    """

    def recall_pos(pres_list,rec_list):
        pres_list = list(pres_list)
        rec_list = list(rec_list)
        result = np.zeros(len(pres_list)) if len(pres_list)>=len(rec_list) else np.zeros(len(rec_list))
        result.fill(np.nan)
        for idx,rec_word in enumerate(rec_list):
            if rec_word in pres_list:
                if type(rec_word) is str:
                    result[idx]=int(pres_list.index(rec_word)+1)
        return result

    result = []
    for pres_list, rec_list in zip(presented.values, recalled.values):
        result.append(recall_pos(pres_list, rec_list))
    return result

def compute_distances(pres_list, feature_list, dist_funcs):
    """
    Compute distances between list words along n feature dimensions

    Parameters
    ----------
    pres_list : list
        list of presented words
    feature_list : list
        list of feature dicts for presented words
    dist_funcs : dict
        dict of distance functions for each feature

    Returns
    ----------
    distances : dict
        dict of distance matrices for each feature
    """

    # initialize dist dict
    distances = {}

    # for each feature in dist_funcs
    for feature in dist_funcs:

        # initialize dist matrix
        dists = np.zeros((len(pres_list), len(pres_list)))

        # for each word in the list
        for idx1, item1 in enumerate(pres_list):

            # for each word in the list
            for idx2, item2 in enumerate(pres_list):

                # compute the distance between word 1 and word 2 along some feature dimension
                dists[idx1,idx2] = dist_funcs[feature](feature_list[idx1][feature],feature_list[idx2][feature])

        # set that distance matrix to the value of a dict where the feature name is the key
        distances[feature] = dists

    return distances

def compute_feature_weights(pres_list, rec_list, feature_list, distances):
    """
    Compute clustering scores along a set of feature dimensions

    Parameters
    ----------
    pres_list : list
        list of presented words
    rec_list : list
        list of recalled words
    feature_list : list
        list of feature dicts for presented words
    distances : dict
        dict of distance matrices for each feature

    Returns
    ----------
    weights : list
        list of clustering scores for each feature dimension
    """

    # initialize the weights object for just this list
    weights = {}
    for feature in feature_list[0]:
        weights[feature] = []

    # return default list if there is not enough data to compute the fingerprint
    if len(rec_list) <= 2:
        print('Not enough recalls to compute fingerprint, returning default fingerprint.. (everything is .5)')
        for feature in feature_list[0]:
            weights[feature] = .5
        return [weights[key] for key in weights]

    # initialize past word list
    past_words = []
    past_idxs = []

    # loop over words
    for i in range(len(rec_list)-1):

        # grab current word
        c = rec_list[i]

        # grab the next word
        n = rec_list[i + 1]

        # if both recalled words are in the encoding list and haven't been recalled before
        if (c in pres_list and n in pres_list) and (c not in past_words and n not in past_words):

            # for each feature
            for feature in feature_list[0]:

                # get the distance vector for the current word
                dists = distances[feature][pres_list.index(c),:]

                # distance between current and next word
                cdist = dists[pres_list.index(n)]

                # filter dists removing the words that have already been recalled
                dists_filt = np.array([dist for idx, dist in enumerate(dists) if idx not in past_idxs])

                # get indices
                avg_rank = np.mean(np.where(np.sort(dists_filt)[::-1] == cdist)[0]+1)

                # compute the weight
                weights[feature].append(avg_rank / len(dists_filt))

            # keep track of what has been recalled already
            past_idxs.append(pres_list.index(c))
            past_words.append(c)

    # average over the cluster scores for a particular dimension
    for feature in weights:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", category=RuntimeWarning)
            weights[feature] = np.nanmean(weights[feature])

    return [weights[key] for key in weights]

# def single_perm(p, r, f, distances):
#     r_real = compute_feature_weights(p, r, f, distances)
#     perm = list(np.random.permutation(r))
#     r_perm = compute_feature_weights(p, perm, f, distances)
#     return [feature_perm < r_real[idx] for idx, feature_perm in enumerate(r_perm)]
#
# def permute_fingerprint_parallel(p, r, f, distances, n_perms=100):
#
#     executor = concurrent.futures.ThreadPoolExecutor(10)
#     futures = [executor.submit(single_perm, p, r, f, distances) for perm in range(n_perms)]
#     concurrent.futures.wait(futures)
#
#     results = [perm.result() for perm in futures]
#
#     print(np.sum(np.array(results), axis=0) / n_perms)
#
#     return np.sum(np.array(results), axis=0) / n_perms

def permute_fingerprint_serial(p, r, f, distances, n_perms=100):

    r_perms = []
    r_real = compute_feature_weights(p, r, f, distances)

    for iperm in range(n_perms):
        r_perm = list(np.random.permutation(r))
        r_perms.append(compute_feature_weights(p, r_perm, f, distances))

    r_perms_bool = []
    for perm in r_perms:
        r_perm_bool = []
        for idx, feature_perm in enumerate(perm):
            r_perm_bool.append(feature_perm < r_real[idx])
        r_perms_bool.append(r_perm_bool)

    return np.sum(np.array(r_perms_bool), axis=0) / n_perms

# accuracy analysis
def accuracy_helper(pres_slice, rec_slice):
    """
    Computes proportion of words recalled

    Parameters
    ----------
    pres_slice : Pandas Dataframe
        chunk of presentation data to be analyzed
    rec_slice : Pandas Dataframe
        chunk of recall data to be analyzed

    Returns
    ----------
    prop_recalled : numpy array
      proportion of words recalled

    """

    # compute recall_matrix for data slice
    recall = recall_matrix(pres_slice, rec_slice)

    # simple function that returns 1 if item encoded in position n is in recall list
    def compute_acc(lst):
        return len([i for i in np.unique(lst) if i>0])/(pres_slice.list_length)

    # get spc for each row in recall matrix
    acc_matrix = [compute_acc(lst) for lst in recall]

    # average over rows
    prop_recalled = np.mean(acc_matrix,axis=0)

    return prop_recalled

# serial position curve
def spc_helper(pres_slice, rec_slice):
    """
    Computes probability of a word being recalled (in the appropriate recall list), given its presentation position

    Parameters
    ----------
    pres_slice : Pandas Dataframe
        chunk of presentation data to be analyzed
    rec_slice : Pandas Dataframe
        chunk of recall data to be analyzed

    Returns
    ----------
    prop_recalled : numpy array
      each number represents the probability of recall for a word presented in given position/index

    """

    # compute recall_matrix for data slice
    recall = recall_matrix(pres_slice, rec_slice)

    # get spc for each row in recall matrix
    spc_matrix = [[1 if pos in lst else 0 for pos in range(1,len(lst)+1)] for lst in recall]

    # average over rows
    prop_recalled = np.mean(spc_matrix, axis=0)

    return prop_recalled

# probability of nth recall
def pnr_helper(pres_slice, rec_slice, n):

    """
    Computes probability of a word being recalled nth (in the appropriate recall
    list), given its presentation position.  Note: zero indexed

    Parameters
    ----------
    pres_slice : Pandas Dataframe
        chunk of presentation data to be analyzed
    rec_slice : Pandas Dataframe
        chunk of recall data to be analyzed

    Returns
    ----------
    prob_recalled : numpy array
      each number represents the probability of nth recall for a word presented in given position/index

    """

    # compute recall_matrix for data slice
    recall = recall_matrix(pres_slice, rec_slice)

    # simple function that returns 1 if item encoded in position n is recalled first
    def pos_recalled_first(pos,lst,n):
        return 1 if pos==lst[n] else 0

    # get pfr for each row in recall matrix
    pnr_matrix = [[pos_recalled_first(pos,lst,n) for pos in range(1,len(lst)+1)] for lst in recall]

    # average over rows
    prob_recalled = np.mean(pnr_matrix,axis=0)

    return prob_recalled

# lag-crp
def lagcrp_helper(pres_slice, rec_slice):
    """
    Computes probabilities for each transition distance (probability that a word recalled will be a given distance--in presentation order--from the previous recalled word)

    Parameters
    ----------
    pres_slice : Pandas Dataframe
        chunk of presentation data to be analyzed
    rec_slice : Pandas Dataframe
        chunk of recall data to be analyzed

    Returns
    ----------
    prob_recalled : numpy array
      each float is the probability of transition distance (distnaces indexed by position, from -(n-1) to (n-1), excluding zero

    """

    # compute recall_matrix for data slice
    recall = recall_matrix(pres_slice, rec_slice)

    def check_pair(a, b):
        if (a>0 and b>0) and (a!=b):
            return True
        else:
            return False

    def compute_actual(recall_list, list_length):
        arr=pd.Series(data=np.zeros((list_length)*2), index=list(range(-list_length,0))+list(range(1,list_length+1)))
        recalled=[]
        for trial in range(0,list_length-1):
            a=recall_list[trial]
            b=recall_list[trial+1]
            if check_pair(a, b) and (a not in recalled) and (b not in recalled):
                arr[b-a]+=1
            recalled.append(a)
        return arr

    def compute_possible(recall_list, list_length):
        arr=pd.Series(data=np.zeros((list_length)*2), index=list(range(-list_length,0))+list(range(1,list_length+1)))
        recalled=[]
        for trial in recall_list:

            if np.isnan(trial):
                pass
            else:

                low_bound=int(1-trial)
                up_bound=int(list_length-trial)

                chances=list(range(low_bound,0))+list(range(1,up_bound+1))
                #ALL transitions

                #remove transitions not possible
                for each in recalled:
                    if each-trial in chances:
                        chances.remove(each-trial)

                #update array with possible transitions
                arr[chances]+=1

                recalled.append(trial)

        return arr

    ########

    list_crp = []
    for n_list in recall:
        actual = compute_actual(n_list, pres_slice.list_length)
        possible = compute_possible(n_list, pres_slice.list_length)
        crp = [0.0 if j==0 else i/j for i,j in zip(actual,possible)]
        crp.insert(int(len(crp)/2),np.nan)
        list_crp.append(crp)

    prob_recalled = np.mean(list_crp, axis=0)

    return prob_recalled

# temporal clustering analysis
def temporal_helper(pres_slice, rec_slice, permute=False, n_perms=1000):
    """
    Computes temporal clustering score

    Parameters
    ----------
    pres_slice : Pandas Dataframe
        chunk of presentation data to be analyzed
    rec_slice : Pandas Dataframe
        chunk of recall data to be analyzed

    Returns
    ----------
    score : float
        a number representing temporal clustering

    """

    # initialize temporal clustering list
    temporal_clustering = []

    # define distance function for temporal clustering
    dist_funcs = {
        'temporal' : lambda a, b : np.abs(a-b)
    }

    # define features (just positions for temporal clustering)
    f = [{'temporal' : i} for i in range(pres_slice.list_length+1)]

    # loop over lists
    for p, r in zip(pres_slice.as_matrix(), rec_slice.as_matrix()):

        # turn arrays into lists
        p = list(p)
        r = list(filter(lambda ri: isinstance(ri, str), list(r)))

        if len(r)>1:

            # compute distances
            distances = compute_distances(p, f, dist_funcs)

            # add optional bootstrapping
            if permute:
                temporal_clustering.append(permute_fingerprint_serial(p, r, f, distances, n_perms=n_perms))
            else:
                temporal_clustering.append(compute_feature_weights(p, r, f, distances))
        else:
            temporal_clustering.append([np.nan]*len(f[0].keys()))

    # return average over rows
    return np.nanmean(temporal_clustering, axis=0)

# fingerprint analysis
def fingerprint_helper(pres_slice, rec_slice, feature_slice, dist_funcs, permute=False, n_perms=1000):
    """
    Computes clustering along a set of feature dimensions

    Parameters
    ----------
    pres_slice : Pandas Dataframe
        chunk of presentation data to be analyzed
    rec_slice : Pandas Dataframe
        chunk of recall data to be analyzed
    feature_slice : Pandas Dataframe
        chunk of features data to be analyzed
    dist_funcs : dict
        Dictionary of distance functions for feature clustering analyses

    Returns
    ----------
    probabilities : numpy array
      each number represents clustering along a different feature dimension

    """
    import time

    # compute fingerprint for each list within a chunk
    fingerprint_matrix = []

    for p, r, f in zip(pres_slice.as_matrix(), rec_slice.as_matrix(), feature_slice.as_matrix()):

        # turn arrays into lists
        p = list(p)
        f = list(f)
        r = list(filter(lambda ri: isinstance(ri, str), list(r)))

        if len(r)>1:

            # compute distances
            distances = compute_distances(p, f, dist_funcs)

            # add optional bootstrapping
            if permute:
                fingerprint_matrix.append(permute_fingerprint_serial(p, r, f, distances, n_perms=n_perms))
            else:
                fingerprint_matrix.append(compute_feature_weights(p, r, f, distances))
        else:
            fingerprint_matrix.append([np.nan]*len(f[0].keys()))

    # return average over rows
    return np.mean(fingerprint_matrix, axis=0)

# fingerprint + temporal clustering analysis
def fingerprint_temporal_helper(pres_slice, rec_slice, feature_slice, dist_funcs, permute=True, n_perms=1000):
    """
    Computes clustering along a set of feature dimensions

    Parameters
    ----------
    pres_slice : Pandas Dataframe
        chunk of presentation data to be analyzed
    rec_slice : Pandas Dataframe
        chunk of recall data to be analyzed
    feature_slice : Pandas Dataframe
        chunk of features data to be analyzed
    dist_funcs : dict
        Dictionary of distance functions for feature clustering analyses

    Returns
    ----------
    probabilities : numpy array
      each number represents clustering along a different feature dimension

    """
    # compute fingerprint for each list within a chunk
    fingerprint_matrix = []

    for p, r, f in zip(pres_slice.as_matrix(), rec_slice.as_matrix(), feature_slice.as_matrix()):

        # turn arrays into lists
        p = list(p)
        f = list(f)
        r = list(filter(lambda ri: isinstance(ri, str), list(r)))

        # add in temporal clustering
        nf = []
        for idx, fi in enumerate(f):
            fi['temporal'] = idx
            nf.append(fi)

        dist_funcs_copy = dist_funcs.copy()
        dist_funcs_copy['temporal'] = lambda a, b : np.abs(a-b)

        # if there is at least 1 transition
        if len(r)>1:

            # compute distances
            distances = compute_distances(p, nf, dist_funcs_copy)

            # add optional bootstrapping
            if permute:
                fingerprint_matrix.append(permute_fingerprint_serial(p, r, nf, distances, n_perms=n_perms))
            else:
                fingerprint_matrix.append(compute_feature_weights(p, r, nf, distances))
        else:
            fingerprint_matrix.append([np.nan]*len(nf[0].keys()))

    return np.nanmean(fingerprint_matrix, axis=0)

# main analysis function
def analyze(data, subjgroup=None, listgroup=None, subjname='Subject',
            listname='List', analysis=None, n=0, permute=False, n_perms=1000,
            parallel=False):
    """
    General analysis function that groups data by subject/list number and performs analysis.

    Parameters
    ----------
    data : Egg data object
        The data to be analyzed

    subjgroup : list of strings or ints
        String/int variables indicating how to group over subjects.  Must be
        the length of the number of subjects

    subjname : string
        Name of the subject grouping variable

    listgroup : list of strings or ints
        String/int variables indicating how to group over list.  Must be
        the length of the number of lists

    listname : string
        Name of the list grouping variable

    analysis : string
        This is the analysis you want to run.  Can be accuracy, spc, pfr,
        temporal or fingerprint

    n : int
        Optional argument for pnr analysis.  Defines encoding position of item
        to run pnr.  Default is 0, and it is zero indexed

    permute : bool
        Optional argument for fingerprint/temporal cluster analyses. Determines
        whether to correct clustering scores by shuffling recall order for each list
        to create a distribution of clustering scores (for each feature). The
        "corrected" clustering score is the proportion of clustering scores in
        that random distribution that were lower than the clustering score for
        the observed recall sequence. Default is False.

    n_perms : int
        Optional argument for fingerprint/temporal cluster analyses. Number of
        permutations to run for "corrected" clustering scores. Default is 1000 (
        per recall list).

    parallel : bool
        Option to use multiprocessing (this can help speed up the permutations
        tests in the clustering calculations)

    Returns
    ----------
    analyzed_data : Pandas DataFrame
        DataFrame containing the analysis results

    """

    # make sure an analysis is specified
    if analysis is None:
        raise ValueError('You must pass an analysis type.')

    # check if subject/list grouping variables exist on the egg
    if hasattr(data, 'subjgroup'):
        if data.subjgroup is not None:
            subjgroup = data.subjgroup

    if hasattr(data, 'subjname'):
        if data.subjname is not None:
            subjname = data.subjname

    if hasattr(data, 'listgroup'):
        if data.listgroup is not None:
            listgroup = data.listgroup

    if hasattr(data, 'listname'):
        if data.listname is not None:
            listname = data.listname

    if type(data) != list:
        data = [data]

    if type(analysis) != list:
        analysis = [analysis]

    result = [[] for d in range(len(data))]

    for idx,d in enumerate(data):
        for a in analysis:

            if a is 'accuracy':
                r = analyze_chunk(d, subjgroup=subjgroup,
                                  listgroup=listgroup,
                                  subjname=subjname,
                                  listname=listname,
                                  analysis=accuracy_helper,
                                  analysis_type='accuracy',
                                  pass_features=False,
                                  parallel=parallel)
            elif a is 'spc':
                r = analyze_chunk(d, subjgroup=subjgroup,
                                  listgroup=listgroup,
                                  subjname=subjname,
                                  listname=listname,
                                  analysis=spc_helper,
                                  analysis_type='spc',
                                  pass_features=False,
                                  parallel=parallel)
            elif a is 'pfr':
                r = analyze_chunk(d, subjgroup=subjgroup,
                                  listgroup=listgroup,
                                  subjname=subjname,
                                  listname=listname,
                                  analysis=pnr_helper,
                                  analysis_type='pfr',
                                  pass_features=False,
                                  n=0,
                                  parallel=parallel)
            elif a is 'pnr':
                r = analyze_chunk(d, subjgroup=subjgroup,
                                  listgroup=listgroup,
                                  subjname=subjname,
                                  listname=listname,
                                  analysis=pnr_helper,
                                  analysis_type='pnr',
                                  pass_features=False,
                                  n=n,
                                  parallel=parallel)
            elif a is 'lagcrp':
                r = analyze_chunk(d, subjgroup=subjgroup,
                                  listgroup=listgroup,
                                  subjname=subjname,
                                  listname=listname,
                                  analysis=lagcrp_helper,
                                  analysis_type='lagcrp',
                                  pass_features=False,
                                  parallel=parallel)
                # set indices for lagcrp
                r.columns=range(-int((len(r.columns)-1)/2),int((len(r.columns)-1)/2)+1)
            elif a is 'fingerprint':
                r = analyze_chunk(d, subjgroup=subjgroup,
                                  listgroup=listgroup,
                                  subjname=subjname,
                                  listname=listname,
                                  analysis=fingerprint_helper,
                                  analysis_type='fingerprint',
                                  pass_features=True,
                                  permute=permute,
                                  n_perms=n_perms,
                                  parallel=parallel)
            elif a is 'temporal':
                r = analyze_chunk(d, subjgroup=subjgroup,
                                  listgroup=listgroup,
                                  subjname=subjname,
                                  listname=listname,
                                  analysis=temporal_helper,
                                  analysis_type='temporal',
                                  permute=permute,
                                  n_perms=n_perms,
                                  parallel=parallel)
            elif a is 'fingerprint_temporal':
                r = analyze_chunk(d, subjgroup=subjgroup,
                                  listgroup=listgroup,
                                  subjname=subjname,
                                  listname=listname,
                                  analysis=fingerprint_temporal_helper,
                                  analysis_type='fingerprint_temporal',
                                  pass_features=True,
                                  permute=permute,
                                  n_perms=n_perms,
                                  parallel=parallel)

            result[idx].append(r)

    # return analysis result
    if len(data)>1 and len(analysis)>1:
        return result
    elif len(data)>1 and len(analysis)==1:
        return [item[0] for item in result]
    elif len(data)==1 and len(analysis)>1:
        return [item for item in result[0]]
    else:
        return result[0][0]


from copy import deepcopy
import os
ROOT = os.path.dirname(os.path.abspath(__file__)) + '/'

from astropy import log as logger
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

from .filter import Filter

from .utils.one_filter import OneFilter
from .utils.plot import MakePlots
from .utils.plot_filters import PlotFilters
from .utils.tools import properties, grid_units, get_slices, average_collapse, central_wav
from .utils.units import ConvertUnits

class SyntheticSED(object):

    '''
    SyntheticSED is part of the FluxCompensator. It converts
    input_arrays (e. g. HYPERION ModelOutput) to "realistic" synthetic observations.
    It contains attributes like ModelOutput (see Notes).
    If input_array is already a SyntheticSED object, the attributes are
    passed. If input_array is not a SyntheticSED object, SyntheticSED
    specific attributes are defined and then passed.

    Parameters
    ----------

    input_array : SyntheticSED, ModelOutput (get_sed), optional
        input_array also reads arrays with ModelOutput like properties.

    unit_out : str, optional
        The output units for SyntheticSED val. Valid options are:

            * ``'ergs/cm^2/s'``
            * ``'ergs/cm^2/s/Hz'``
            * ``'Jy'``
            * ``'mJy'``
            * ``'MJy/sr'``

        The default is ``'ergs/cm^2/s'``.

    name : str
        The name of the FluxCompensator object until another
        input_array is called. The default is ``None``.

    Attributes
    ----------

    wav : numpy.ndarray
        The wavelengths of the val vector entries in microns.

    val : numpy.ndarray
        The 1D vector with shape like wav.

    units : str
        Current units of the val vector.

    distance : str
        Distance to the observed object in cm.

    x_min : float
        Physical offset from axis origin in FOV in cm.

    x_max : float
        Physical offset from axis origin in FOV in cm.

    y_min : float
        Physical offset from axis origin in FOV in cm.

    y_max : float
        Physical offset from axis origin in FOV in cm.

    lon_min : float
        Minimal longitudinal angle.

    lon_max : float
        Maximal longitudinal angle.

    lat_min : float
        Minimal latitudinal angle.

    lat_max : float
        Maximal latitudinal angle.

    pix_area_sr : float
        Pixel area per sr.

    ap_min : float
        Minimal aperture.

    ap_max : float
        Maximal aperture.

    Notes
    -----
    unit_in : str
        Unit of val in input_array. Valid options are:

            * ``'ergs/cm^2/s'``
            * ``'ergs/cm^2/s/Hz'``
            * ``'Jy'``
            * ``'mJy'``
            * ``'MJy/sr'``

    grid_unit : float
        Physical unit of FOV axis in cm. Valid options are:

            * ``au`` in cm
            * ``pc`` in cm
            * ``kpc`` in cm

    grid_unit_name
        Astronomical unit of FOV axis. Valid options are:

            * ``'au'``
            * ``'pc'``
            * ``'kpc'``

    FOV : tuple
        Tuple ``FOV(x,y)`` of Field of View pixel entries.

            * pixel in x direction: ``FOV[0]``
            * pixel in y direction: ``FOV[1]``

    name : str
        The name of the FluxCompensator object until another
        input_array is called. The default is ``None``.

    stage : str
        Gives current operation stage of SyntheticSED.
        E. g. ``'SyntheticSED: convolve_filter'``

    log : list
         List of strings of the previous and current stages.

    filter : dict
        Dictionary ``filter = {name, waf_0, waf_min, waf_max}``
        of the applied filter.

            * name of filter:     ``filter['name']``
            * central wavelength: ``filter['waf_0']``
            * minimal wavelength:  ``filter['waf_min']``
            * maximal wavelength: ``filter['waf_max']``

    Returns
    -------

    sed : SyntheticSED
        1D val array (collapsed rough SED) with SyntheticSED properties.

    flux : SyntheticFlux
        0D val array (scalar) with SyntheticFlux properties.
    '''

    def __init__(self, input_array, unit_out='ergs/cm^2/s', name=None):

        # Hyperion ModelOutput attributes (image and sed)
        if input_array.val.ndim == 2 and input_array.val[:, 0].ndim == 1:
            self.val = np.array(input_array.val[0, :])
        else:
            self.val = np.array(deepcopy(input_array.val))

        self.wav = np.array(deepcopy(input_array.wav))
        self.units = input_array.units
        self.distance = input_array.distance

        # Hyperion Image
        try:
            self.x_max = input_array.x_max
            self.x_max = input_array.x_max
            self.x_min = input_array.x_min
            self.y_max = input_array.y_max
            self.y_min = input_array.y_min
            self.lon_min = input_array.lon_min
            self.lon_max = input_array.lon_max
            self.lat_min = input_array.lat_min
            self.lat_max = input_array.lat_max
            self.pix_area_sr = input_array.pix_area_sr
            # switch
            self.hyperion_cube = True

        except AttributeError:
            self.x_max = None
            self.x_max = None
            self.x_min = None
            self.y_max = None
            self.y_min = None
            self.lon_min = None
            self.lon_max = None
            self.lat_min = None
            self.lat_max = None
            self.pix_area_sr = None
            # switch
            self.hyperion_cube = None

        # Hyperion SED
        try:
            self.ap_min = input_array.ap_min
            self.ap_max = input_array.ap_max
            # switch
            self.hyperion_sed = True

        except AttributeError:
            self.ap_min = None
            self.ap_max = None
            # switch
            self.hyperion_sed = None

        ##################
        # new attributes #
        ##################

        from .cube import SyntheticCube

        if isinstance(input_array, SyntheticSED) or isinstance(input_array, SyntheticCube):
            # attributes with are passed, since input_array is SyntheticSED or SyntheticCube

            # physical values
            self.unit_in = input_array.unit_in
            self.unit_out = input_array.unit_out

            if self.x_max is not None:
                self.grid_unit = grid_units(self.x_max - self.x_min)['grid_unit']
                self.grid_unit_name = grid_units(self.x_max - self.x_min)['grid_unit_name']
            else:
                self.grid_unit = None
                self.grid_unit_name = None

            # properties of cube
            self.FOV = deepcopy(input_array.FOV)

            # name
            self.name = input_array.name
            self.stage = input_array.stage
            self.log = deepcopy(input_array.log)

            # filter
            self.filter = deepcopy(input_array.filter)

        elif not isinstance(input_array, SyntheticSED) and not isinstance(input_array, SyntheticCube) and self.hyperion_sed is True:
            # attributes with are defined, since input_array is NOT SyntheticSED or SyntheticCube but HyperionSED
            self.unit_in = input_array.units
            self.unit_out = unit_out

            self.grid_unit = None
            self.grid_unit_name = None

            self.FOV = None

            # name
            self.name = name
            self.stage = 'SyntheticSED:  initial'
            self.log = [self.stage]

            # filter
            self.filter = {'name': None, 'waf_0': None, 'waf_min': None, 'waf_max': None}

            # convert into val units into unit_out
            if self.unit_in == 'MJy/sr' or self.unit_out == 'MJy/sr' or self.unit_in == 'Jy/arcsec^2' or self.unit_out == 'Jy/arcsec^2':
                raise Exception('WARNING: Input or Output units needs to differ from MJy/sr if Input_array is not SyntheticCube or SyntheticSED or HyperionCube')

            s = ConvertUnits(wav=self.wav, val=self.val)
            self.val = s.get_unit(in_units=self.unit_in, out_units=self.unit_out)

            self.units = self.unit_out

        else:   # attributes with are defined, since input_array is NOT SyntheticSED or SyntheticCube or HyperionOutput
            # physical values
            self.unit_in = input_array.units
            self.unit_out = unit_out

            self.grid_unit = grid_units(self.x_max)['grid_unit']
            self.grid_unit_name = grid_units(self.x_max)['grid_unit_name']

            self.FOV = (self.x_max - self.x_min, self.y_max - self.y_min)

            # name
            self.name = name
            self.stage = 'SyntheticSED:  initial'
            self.log = [self.stage]

            # filter
            self.filter = {'name': None, 'waf_0': None, 'waf_min': None, 'waf_max': None}

            # convert into val units into unit_out
            s = ConvertUnits(wav=self.wav, val=self.val)
            self.val = s.get_unit(in_units=self.unit_in, out_units=self.unit_out, input_resolution=self.resolution['arcsec'])

            self.units = self.unit_out

    def extinction(self, A_v, input_opacities=None):
        '''
        Accounts for reddening.

        Parameters
        ----------

        A_v : Value of the visible extinction.

        input_opacities : ``None``, str
            If ``None`` standard extinction law is used.
            Otherwise a e. g. input_opacities.txt file can be passed
            as a str to read an opacity file with column #1 wav in microns
            and column #2 in cm^2/g.
            Default is ``None``.

        Returns
        -------

        sed : SyntheticSED
        '''

        stage = 'SyntheticSED: extinction'

        # read own extinction law
        if input_opacities is None:
            t = np.loadtxt(ROOT + 'database/extinction/extinction_law.txt')

        else:
            t = np.loadtxt(input_opacities)

        wav_ext = t[:, 0]
        k_lam = t[:, 1]

        # wav_ext monotonically increasing
        if wav_ext[0] > wav_ext[1]:
            wav_ext = wav_ext[::-1]
            k_lam = k_lam[::-1]

        k_v = np.interp(0.550, wav_ext, k_lam)

        # interpolate to get A_int for a certain wavelength
        k = np.interp(self.wav, wav_ext, k_lam)
        A_int_lam = A_v * (k / k_v)

        # apply extinction law
        val_ext = np.zeros(shape=np.shape(self.val))
        val_ext[:len(self.wav)] = self.val[:len(self.wav)] * 10 ** (-0.4 * A_int_lam[:len(self.wav)])

        # return SyntheticSED
        s = SyntheticSED(self)
        s.val = val_ext
        s.stage = stage
        s.log.append(s.stage)

        return s

    def convolve_filter(self, filter_input, plot_rebin=None, plot_rebin_dpi=None):
        '''
        Convolves vector val entries within filter limits in a 0D val array.

        Parameters
        ----------

        filter_input : object

            * database : if filter ``name`` from FluxCompensator database is used.
            * Filter : if own filter is used.

        plot_rebin : ``True``, ``None``
            Switch to plot the rebined filter and the original filter in one plot.

        plot_rebin_dpi  :  ``None``, scalar > 0
            The resolution in dots per inch.
            ``None`` is default and will use the val savefig.dpi
            in the matplotlibrc file.

        Returns
        -------

        flux : SyntheticFlux
        '''

        stage = 'SyntheticSED: convolve_filter'

        # debugging comment
        logger.debug('-' * 70)
        logger.debug(stage)
        logger.debug('-' * 70)

        weight = filter_input.rebin(self.wav, self.val)

        # returns weight{'wav_short' 'val_short' 'Response_new' 'filter_index' 'wavf_0' 'waf_min' 'waf_max' 'filter_name'}
        wav_short = weight['wav_short']
        val_short = weight['val_short']
        filter_index = weight['filter_index']
        Response_new = weight['Response_new']
        waf_0 = weight['waf_0']
        waf_min = weight['waf_min']
        waf_max = weight['waf_max']
        filter_name = weight['filter_name']

        if plot_rebin is not None:
            plot = filter_input.plot(val_name=self.name, dpi=plot_rebin_dpi)

        # weight val_short with rebined response
        val = val_short.copy()
        val[:len(wav_short)] = val_short[:len(wav_short)] * Response_new[:len(wav_short)]

        # collapse remaining vector to val scalar
        val_tot = np.sum(val)

        # return SyntheticFlux
        from .flux import SyntheticFlux
        f = SyntheticFlux(self)
        f.log.append(stage)
        f.stage = 'SyntheticFlux: initial'
        f.log.append(f.stage)
        f.val = val_tot
        f.wav = np.array(waf_0)
        f.filter = {'name': filter_name, 'waf_0': waf_0, 'waf_min': waf_min, 'waf_max': waf_max}

        return f

    def get_total_val(self, wav_1, wav_2):
        '''
        Collapses the val entries in the vector within the boundaries wav_1
        and wav_2 into a 0D val array.

        WARNING: This tool cannot replace convolve_filter!
                 But it can be used to produce rough estimates
                 in-between the processes.

        Parameters
        ----------

        wav_1, wav_2 : float
            Boundaries in microns.

        Returns
        -------

        flux : SyntheticFlux
        '''

        stage = 'SyntheticSED: get_total_val'

        # for MJy/sr convert first, add and then convert back
        if self.unit_out == 'MJy/sr' or self.unit_out == 'Jy/arcsec^2':
            s = ConvertUnits(wav=self.wav, val=self.val)
            self.val = s.get_unit(in_units=self.units, out_units='Jy', input_resolution=self.resolution['arcsec'])

        # slices within boundaries are extracted, averaged collapsed to a single scalar val
        vec = get_slices(wav=self.wav, val=self.val, wav_1=wav_1, wav_2=wav_2)
        f_total = average_collapse(val=vec['val_short'])

        # real limits within collapse
        wav_max = 10 ** (np.log10(self.wav[vec['filter_index'][0]]) + self.spacing_wav / 2.)
        wav_min = 10 ** (np.log10(self.wav[vec['filter_index'][-1]]) - self.spacing_wav / 2.)
        wav_total = central_wav(wav=[wav_min, wav_max])

        # for MJy/sr convert first, add and then convert back
        if self.unit_out == 'MJy/sr' or self.unit_out == 'Jy/arcsec^2':
            s = ConvertUnits(wav=wav_total, val=f_total)
            f_total = s.get_unit(in_units='Jy', out_units=self.unit_out, input_resolution=self.resolution['arcsec'] * self.pixel[0])

        # return SyntheticFlux
        from .flux import SyntheticFlux
        f = SyntheticFlux(self)
        f.log.append(stage)
        f.stage = 'SyntheticFlux: initial'
        f.log.append(f.stage)
        f.val = np.array(f_total)
        f.wav = np.array(wav_total)
        f.filter = {'name': 'val_tot', 'waf_0': wav_total, 'waf_min': wav_min, 'waf_max': wav_max}

        return f

    def plot_sed_multi_filter(self, multi_filter_val, multi_filter_wav, names, filter_label_size=None, ymin=10. ** (-5), dpi=None):
        '''
        Reads in array of filtered val and plots it with the current passed val of
        SyntheticSED in a log-log diagram.
        That way the quality of multiple filtered val can be checked, if filters are
        in database.
        If the used filters are not part of the filter database, define a OneFilter object.

        Parameters
        ----------

        multi_filter_val : np.ndarray
            1D vector with val entries from several filters.

        multi_filter_wav : np.ndarray
            1D vector of central wavelengths from several filters.

        names : np.ndarray
            Original name of several filters, if found in database.

        filter_label_size : ``None``, ``True``
            Switch wether to set print labels above the filters.

                * ``None``: No labels plotted.
                * ``True``: Plots names[index] as label.

        ymin : float
            Minimal vertical limit for SED plot.
            Default is 10.**(-10) of the maximum val.

        dpi  :  ``None``, scalar > 0
            The resolution in dots per inch.
            ``None`` is default and will use the val savefig.dpi
            in the matplotlibrc file.

        Returns
        -------

        flux : SyntheticFlux
        '''

        stage = 'SyntheticSED: plot_sed_multi_filter'

        # plot sed and val within boundaries of sed
        fig = plt.figure()
        font = {'size': 8}

        mpl.rc('font', **font)
        suptitle = fig.suptitle('output SED with filtered val', fontsize=12.)

        x_max = self.wav[0]
        x_min = self.wav[-1]
        y_max = 10 * max(self.val)
        y_min = max(self.val) * ymin

        # subplot with SED and collapsed fluval
        gs = gridspec.GridSpec(2, 1, height_ratios=[2, 1])
        ax = plt.subplot(gs[0])
        plt.subplots_adjust(wspace=0, hspace=0)
        ax.set_xlim([x_min, x_max])
        ax.set_ylim([y_min, y_max])

        #ax.set_xlabel(r'$\lambda$ [$\mu$m]')
        ax.set_ylabel('val ' + '[' + str(self.unit_out) + ']')
        plt.loglog(self.wav, self.val, 'o-k', label='rough SED', linewidth=2)

        # sort c.fluval.wav, filt after c.wav
        # get index of c.wav in right order
        # [::-1] largest cwav comes first
        sort_order = np.argsort(multi_filter_wav)
        multi_filter_val = np.take(multi_filter_val, sort_order)[::-1]
        multi_filter_wav = np.take(multi_filter_wav, sort_order)[::-1]
        names = np.take(names, sort_order)[::-1]

        # color
        f = np.linspace(0, 200, len(names))

        for i in range(len(names)):
            current_waf_0 = multi_filter_wav[i]
            current_val = multi_filter_val[i]
            current_name = names[i]

            color = plt.cm.RdYlBu(int(f[i]))
            #color = plt.cm.autumn(int(f[i]))

            # filter
            text = 'filtered val by ' + str(current_name)
            x = PlotFilters(style='loglog', normalized=True, unit='unit energy')
            plot_filter = x.collect_filters(current_name)
            plt.loglog(current_waf_0, current_val, 'o', color=color)

        # subplot with filter responses
        ax2 = plt.subplot(gs[1])
        ax2.set_xlim([x_min, x_max])
        ax2.set_ylim([10. ** (-4), 10.])
        ax2.set_ylabel('filter response [unit energy]', color='k')

        for i in range(len(names)):
            current_waf_0 = multi_filter_wav[i]
            current_val = multi_filter_val[i]
            current_name = names[i]

            color = plt.cm.RdYlBu(int(f[i]))
            #color = plt.cm.autumn(int(f[i]))
            fancy_line = {'color': color, 'linestyle': '-', 'linewidth': 2}

            x = PlotFilters(style='loglog', normalized=True, unit='unit energy')
            plot_filter = x.collect_filters(current_name)
            x.plot(plot_filter, line=fancy_line)

            # text above filters
            if filter_label_size is not None:
                plt.text(current_waf_0, 3, current_name, horizontalalignment='center', verticalalignment='center', fontsize=filter_label_size)

        ax2.set_xlabel(r'$\lambda$ [$\mu$m]')

        # legend merged
        h1, l1 = ax.get_legend_handles_labels()
        h2, l2 = ax2.get_legend_handles_labels()

        leg = ax.legend(h1 + h2, l1 + l2, loc=(1.03, -0.5))
        for t in leg.get_texts():
            t.set_fontsize('small')    # the legend text fontsize

        fig.savefig(str(self.name) + '_' + 'process-output_SS-multi-filter.png', dpi=dpi, bbox_inches='tight')

        # return SyntheticSED
        s = SyntheticSED(self)
        s.log.append(stage)
        s.stage = 'SyntheticSED: initial'
        s.log.append(s.stage)

        return s

    @property
    def spacing_wav(self):
        '''
        The property spacing_wav estimates the width of the logarithmic
        spaced wav entries.
        '''

        if self.wav.ndim != 0:
            spacing_wav = np.log10(self.wav[0] / self.wav[-1]) / (len(self.wav) - 1)
        else:
            spacing_wav = None
        return spacing_wav

    @property
    def pixel(self):
        '''
        The property pixel is a tuple which resembles the current pixel in a
        val slice. ``pixel(x,y)`` are calls as follows:

        ``x = pixel[0]``
        ``y = pixel[1]``
        '''
        if self.val.ndim in (0, 1):
            pixel = (None, None)
        if self.val.ndim in (2, 3):
            pixel = (self.val.shape[0], self.val.shape[1])
        return pixel

    @property
    def shape(self):
        '''
        The property shape is a string, which resembles the current shape of
        the val array.

        scalar: ``'()'``
        1D:     ``'(wav)'``
        2D:     ``'(x, y)'``
        3D:     ``'(x, y , wav)'``
        '''
        if self.val.ndim == 0:
            shape = '()'
        if self.val.ndim == 1:
            shape = '(wav)'
        if self.val.ndim == 2:
            shape = '(x, y)'
        if self.val.ndim == 3:
            shape = '(x, y, wav)'
        return shape

    @property
    def resolution(self):
        '''
        The property resolution tells you the current resolution. If we are already
        in the SED or val dimension everything is considered as one large pixel.

            resolution in arcsec per pixel : ``resolution['arcsec']``
            resolution in rad per pixel    : ``resolution['rad']``

        '''
        resolution = {}
        if self.pixel[0] is None:
            resolution['rad'] = self.FOV[0] / 1. / self.distance
        else:
            resolution['rad'] = self.FOV[0] / self.pixel[0] / self.distance
        resolution['arcsec'] = np.degrees(resolution['rad']) * 3600
        return resolution


import logging

import numba
import numpy as np
import numpy.random as npr

import second.core.box_np_ops as box_np_ops

logger = logging.getLogger(__name__)

def unmap(data, count, inds, fill=0):
    """Unmap a subset of item (data) back to the original set of items (of
    size count)"""
    if count == len(inds):
        return data

    if len(data.shape) == 1:
        ret = np.empty((count, ), dtype=data.dtype)
        ret.fill(fill)
        ret[inds] = data
    else:
        ret = np.empty((count, ) + data.shape[1:], dtype=data.dtype)
        ret.fill(fill)
        ret[inds, :] = data
    return ret

def create_target_np(all_anchors,
                     gt_boxes,
                     similarity_fn,
                     box_encoding_fn,
                     prune_anchor_fn=None,
                     gt_classes=None,
                     matched_threshold=0.6,
                     unmatched_threshold=0.45,
                     bbox_inside_weight=None,
                     positive_fraction=None,
                     rpn_batch_size=300,
                     norm_by_num_examples=False,
                     gt_importance=None,
                     box_code_size=7):
    """Modified from FAIR detectron.
    Args:
        all_anchors: [num_of_anchors, box_ndim] float tensor.
        gt_boxes: [num_gt_boxes, box_ndim] float tensor.
        similarity_fn: a function, accept anchors and gt_boxes, return
            similarity matrix(such as IoU).
        box_encoding_fn: a function, accept gt_boxes and anchors, return
            box encodings(offsets).
        prune_anchor_fn: a function, accept anchors, return indices that
            indicate valid anchors.
        gt_classes: [num_gt_boxes] int tensor. indicate gt classes, must
            start with 1.
        matched_threshold: float, iou greater than matched_threshold will
            be treated as positives.
        unmatched_threshold: float, iou smaller than unmatched_threshold will
            be treated as negatives.
        bbox_inside_weight: unused
        positive_fraction: [0-1] float or None. if not None, we will try to
            keep ratio of pos/neg equal to positive_fraction when sample.
            if there is not enough positives, it fills the rest with negatives
        rpn_batch_size: int. sample size
        norm_by_num_examples: bool. norm box_weight by number of examples, but
            I recommend to do this outside.
        gt_importance: 1d array. loss weight per gt.
    Returns:
        labels, bbox_targets, bbox_outside_weights
    """
    total_anchors = all_anchors.shape[0]
    if prune_anchor_fn is not None:
        inds_inside = prune_anchor_fn(all_anchors)
        anchors = all_anchors[inds_inside, :]
        if not isinstance(matched_threshold, float):
            matched_threshold = matched_threshold[inds_inside]
        if not isinstance(unmatched_threshold, float):
            unmatched_threshold = unmatched_threshold[inds_inside]
    else:
        anchors = all_anchors
        inds_inside = None
    num_inside = len(inds_inside) if inds_inside is not None else total_anchors
    box_ndim = all_anchors.shape[1]
    logger.debug('total_anchors: {}'.format(total_anchors))
    logger.debug('inds_inside: {}'.format(num_inside))
    logger.debug('anchors.shape: {}'.format(anchors.shape))
    if gt_classes is None:
        gt_classes = np.ones([gt_boxes.shape[0]], dtype=np.int32)
    if gt_importance is None:
        gt_importance = np.ones([gt_boxes.shape[0]], dtype=np.float32)

    # Compute anchor labels:
    # label=1 is positive, 0 is negative, -1 is don't care (ignore)
    labels = np.empty((num_inside, ), dtype=np.int32)
    gt_ids = np.empty((num_inside, ), dtype=np.int32)
    labels.fill(-1)
    gt_ids.fill(-1)
    importance = np.empty((num_inside, ), dtype=np.float32)
    importance.fill(1)
    if len(gt_boxes) > 0:
        # Compute overlaps between the anchors and the gt boxes overlaps
        anchor_by_gt_overlap = similarity_fn(anchors, gt_boxes)
        # Map from anchor to gt box that has highest overlap
        anchor_to_gt_argmax = anchor_by_gt_overlap.argmax(axis=1)
        # For each anchor, amount of overlap with most overlapping gt box
        anchor_to_gt_max = anchor_by_gt_overlap[np.arange(num_inside),
                                                anchor_to_gt_argmax]  #
        # Map from gt box to an anchor that has highest overlap
        gt_to_anchor_argmax = anchor_by_gt_overlap.argmax(axis=0)
        # For each gt box, amount of overlap with most overlapping anchor
        gt_to_anchor_max = anchor_by_gt_overlap[gt_to_anchor_argmax,
                                                np.arange(anchor_by_gt_overlap.
                                                          shape[1])]
        # must remove gt which doesn't match any anchor.
        empty_gt_mask = gt_to_anchor_max == 0
        gt_to_anchor_max[empty_gt_mask] = -1
        """
        if not np.all(empty_gt_mask):
            gt_to_anchor_max = gt_to_anchor_max[empty_gt_mask]
            anchor_by_gt_overlap = anchor_by_gt_overlap[:, empty_gt_mask]
            gt_classes = gt_classes[empty_gt_mask]
            gt_boxes = gt_boxes[empty_gt_mask]
        """
        # Find all anchors that share the max overlap amount
        # (this includes many ties)
        anchors_with_max_overlap = np.where(
            anchor_by_gt_overlap == gt_to_anchor_max)[0]
        # Fg label: for each gt use anchors with highest overlap
        # (including ties)
        gt_inds_force = anchor_to_gt_argmax[anchors_with_max_overlap]
        labels[anchors_with_max_overlap] = gt_classes[gt_inds_force]
        gt_ids[anchors_with_max_overlap] = gt_inds_force
        # Fg label: above threshold IOU
        pos_inds = anchor_to_gt_max >= matched_threshold
        gt_inds = anchor_to_gt_argmax[pos_inds]
        labels[pos_inds] = gt_classes[gt_inds]
        gt_ids[pos_inds] = gt_inds
        bg_inds = np.where(anchor_to_gt_max < unmatched_threshold)[0]
        importance[pos_inds] = gt_importance[gt_inds]
    else:
        # labels[:] = 0
        bg_inds = np.arange(num_inside)
    fg_inds = np.where(labels > 0)[0]
    fg_max_overlap = None
    if len(gt_boxes) > 0:
        fg_max_overlap = anchor_to_gt_max[fg_inds]
    gt_pos_ids = gt_ids[fg_inds]
    # bg_inds = np.where(anchor_to_gt_max < unmatched_threshold)[0]
    # bg_inds = np.where(labels == 0)[0]
    # subsample positive labels if we have too many
    if positive_fraction is not None:
        num_fg = int(positive_fraction * rpn_batch_size)
        if len(fg_inds) > num_fg:
            disable_inds = npr.choice(
                fg_inds, size=(len(fg_inds) - num_fg), replace=False)
            labels[disable_inds] = -1
            fg_inds = np.where(labels > 0)[0]

        # subsample negative labels if we have too many
        # (samples with replacement, but since the set of bg inds is large most
        # samples will not have repeats)
        num_bg = rpn_batch_size - np.sum(labels > 0)
        # print(num_fg, num_bg, len(bg_inds) )
        if len(bg_inds) > num_bg:
            enable_inds = bg_inds[npr.randint(len(bg_inds), size=num_bg)]
            labels[enable_inds] = 0
        bg_inds = np.where(labels == 0)[0]
    else:
        if len(gt_boxes) == 0:
            labels[:] = 0
        else:
            labels[bg_inds] = 0
            # re-enable anchors_with_max_overlap
            labels[anchors_with_max_overlap] = gt_classes[gt_inds_force]
    bbox_targets = np.zeros((num_inside, box_code_size),
                            dtype=all_anchors.dtype)
    if len(gt_boxes) > 0:
        # print(anchors[fg_inds, :].shape, gt_boxes[anchor_to_gt_argmax[fg_inds], :].shape)
        # bbox_targets[fg_inds, :] = box_encoding_fn(
        #     anchors[fg_inds, :], gt_boxes[anchor_to_gt_argmax[fg_inds], :])
        bbox_targets[fg_inds, :] = box_encoding_fn(
            gt_boxes[anchor_to_gt_argmax[fg_inds], :], anchors[fg_inds, :])
    # Bbox regression loss has the form:
    #   loss(x) = weight_outside * L(weight_inside * x)
    # Inside weights allow us to set zero loss on an element-wise basis
    # Bbox regression is only trained on positive examples so we set their
    # weights to 1.0 (or otherwise if config is different) and 0 otherwise
    # NOTE: we don't need bbox_inside_weights, remove it.
    # bbox_inside_weights = np.zeros((num_inside, box_ndim), dtype=np.float32)
    # bbox_inside_weights[labels == 1, :] = [1.0] * box_ndim

    # The bbox regression loss only averages by the number of images in the
    # mini-batch, whereas we need to average by the total number of example
    # anchors selected
    # Outside weights are used to scale each element-wise loss so the final
    # average over the mini-batch is correct
    # bbox_outside_weights = np.zeros((num_inside, box_ndim), dtype=np.float32)
    bbox_outside_weights = np.zeros((num_inside, ), dtype=all_anchors.dtype)
    # uniform weighting of examples (given non-uniform sampling)
    if norm_by_num_examples:
        num_examples = np.sum(labels >= 0)  # neg + pos
        num_examples = np.maximum(1.0, num_examples)
        bbox_outside_weights[labels > 0] = 1.0 / num_examples
    else:
        bbox_outside_weights[labels > 0] = 1.0
    # bbox_outside_weights[labels == 0, :] = 1.0 / num_examples

    # Map up to original set of anchors
    if inds_inside is not None:
        labels = unmap(labels, total_anchors, inds_inside, fill=-1)
        bbox_targets = unmap(bbox_targets, total_anchors, inds_inside, fill=0)
        # bbox_inside_weights = unmap(
        #     bbox_inside_weights, total_anchors, inds_inside, fill=0)
        bbox_outside_weights = unmap(
            bbox_outside_weights, total_anchors, inds_inside, fill=0)
        importance = unmap(importance, total_anchors, inds_inside, fill=0)
    # return labels, bbox_targets, bbox_outside_weights
    ret = {
        "labels": labels,
        "bbox_targets": bbox_targets,
        "bbox_outside_weights": bbox_outside_weights,
        "assigned_anchors_overlap": fg_max_overlap,
        "positive_gt_id": gt_pos_ids,
        "importance": importance,
    }
    if inds_inside is not None:
        ret["assigned_anchors_inds"] = inds_inside[fg_inds]
    else:
        ret["assigned_anchors_inds"] = fg_inds
    return ret


#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# 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.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Generalized Kohn-Sham
'''

import time
import numpy
import scipy.linalg
from pyscf import lib
from pyscf.lib import logger
from pyscf.scf import ghf
from pyscf.dft import rks

def get_veff(ks, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
    '''Coulomb + XC functional for GKS.
    '''
    if mol is None: mol = self.mol
    if dm is None: dm = ks.make_rdm1()
    t0 = (time.clock(), time.time())

    ground_state = (isinstance(dm, numpy.ndarray) and dm.ndim == 2)

    assert(hermi == 1)
    dm = numpy.asarray(dm)
    nso = dm.shape[-1]
    nao = nso // 2
    dm_a = dm[...,:nao,:nao].real
    dm_b = dm[...,nao:,nao:].real

    if ks.grids.coords is None:
        ks.grids.build(with_non0tab=True)
        if ks.small_rho_cutoff > 1e-20 and ground_state:
            ks.grids = rks.prune_small_rho_grids_(ks, mol, dm_a+dm_b, ks.grids)
        t0 = logger.timer(ks, 'setting up grids', *t0)
    if ks.nlc != '':
        if ks.nlcgrids.coords is None:
            ks.nlcgrids.build(with_non0tab=True)
            if ks.small_rho_cutoff > 1e-20 and ground_state:
                ks.nlcgrids = rks.prune_small_rho_grids_(ks, mol, dm_a+dm_b, ks.nlcgrids)
            t0 = logger.timer(ks, 'setting up nlc grids', *t0)

    max_memory = ks.max_memory - lib.current_memory()[0]
    ni = ks._numint
    n, exc, vxc = ni.nr_uks(mol, ks.grids, ks.xc, (dm_a,dm_b), max_memory=max_memory)
    if ks.nlc != '':
        assert('VV10' in ks.nlc.upper())
        _, enlc, vnlc = ni.nr_rks(mol, ks.nlcgrids, ks.xc+'__'+ks.nlc, dm_a+dm_b,
                                  max_memory=max_memory)
        exc += enlc
        vxc += vnlc
    logger.debug(ks, 'nelec by numeric integration = %s', n)
    t0 = logger.timer(ks, 'vxc', *t0)
    if vxc.ndim == 4:
        raise NotImplementedError
    vxc = numpy.asarray(scipy.linalg.block_diag(*vxc), dtype=dm.dtype)

    #enabling range-separated hybrids
    omega, alpha, hyb = ni.rsh_and_hybrid_coeff(ks.xc, spin=mol.spin)

    if abs(hyb) < 1e-10 and abs(alpha) < 1e-10:
        vk = None
        if (ks._eri is None and ks.direct_scf and
            getattr(vhf_last, 'vj', None) is not None):
            ddm = numpy.asarray(dm) - numpy.asarray(dm_last)
            vj = ks.get_j(mol, ddm, hermi)
            vj += vhf_last.vj
        else:
            vj = ks.get_j(mol, dm, hermi)
        vxc += vj
    else:
        if (ks._eri is None and ks.direct_scf and
            getattr(vhf_last, 'vk', None) is not None):
            ddm = numpy.asarray(dm) - numpy.asarray(dm_last)
            vj, vk = ks.get_jk(mol, ddm, hermi)
            vk *= hyb
            if abs(omega) > 1e-10:
                vklr = _get_k_lr(mol, ddm, omega, hermi)
                vklr *= (alpha - hyb)
                vk += vklr
            vj += vhf_last.vj
            vk += vhf_last.vk
        else:
            vj, vk = ks.get_jk(mol, dm, hermi)
            vk *= hyb
            if abs(omega) > 1e-10:
                vklr = _get_k_lr(mol, dm, omega, hermi)
                vklr *= (alpha - hyb)
                vk += vklr
        vxc += vj - vk

        if ground_state:
            exc -= numpy.einsum('ij,ji', dm, vk).real * .5
    if ground_state:
        ecoul = numpy.einsum('ij,ji', dm, vj).real * .5
    else:
        ecoul = None

    vxc = lib.tag_array(vxc, ecoul=ecoul, exc=exc, vj=vj, vk=vk)
    return vxc

def _get_k_lr(mol, dm, omega=0, hermi=0):
    nso = dm.shape[-1]
    nao = nso // 2
    dms = dm.reshape(-1,nso,nso)
    n_dm = dms.shape[0]

    dmaa = dms[:,:nao,:nao]
    dmab = dms[:,nao:,:nao]
    dmbb = dms[:,nao:,nao:]
    dms = numpy.vstack((dmaa, dmbb, dmab))
    if dm.dtype == numpy.complex128:
        dms = numpy.vstack((dms.real, dms.imag))
        hermi = 0

    k1 = rks._get_k_lr(mol, dms, omega, hermi)
    k1 = k1.reshape(-1,n_dm,nao,nao)

    if dm.dtype == numpy.complex128:
        k1 = k1[:3] + k1[3:] * 1j

    vk = numpy.zeros((n_dm,nso,nso), dm.dtype)
    vk[:,:nao,:nao] = k1[0]
    vk[:,nao:,nao:] = k1[1]
    vk[:,:nao,nao:] = k1[2]
    vk[:,nao:,:nao] = k1[2].transpose(0,2,1).conj()
    vk = vk.reshape(dm.shape)
    return vk

class GKS(ghf.GHF):
    '''Generalized Kohn-Sham'''
    def __init__(self, mol):
        ghf.GHF.__init__(self, mol)
        rks._dft_common_init_(self)

    def dump_flags(self, verbose=None):
        ghf.GHF.dump_flags(self, verbose)
        logger.info(self, 'XC functionals = %s', self.xc)
        if self.nlc!='':
            logger.info(self, 'NLC functional = %s', self.nlc)
        logger.info(self, 'small_rho_cutoff = %g', self.small_rho_cutoff)
        self.grids.dump_flags(verbose)

    get_veff = get_veff
    energy_elec = rks.energy_elec
    define_xc_ = rks.define_xc_

    def nuc_grad_method(self):
        raise NotImplementedError

if __name__ == '__main__':
    from pyscf import gto
    mol = gto.Mole()
    mol.verbose = 3
    mol.atom = 'H 0 0 0; H 0 0 1; O .5 .6 .2'
    mol.basis = 'ccpvdz'
    mol.build()

    mf = GKS(mol)
    mf.xc = 'b3lyp'
    mf.kernel()

    dm = mf.init_guess_by_1e(mol)
    dm = dm + 0j
    nao = mol.nao_nr()
    numpy.random.seed(12)
    dm[:nao,nao:] = numpy.random.random((nao,nao)) * .1j
    dm[nao:,:nao] = dm[:nao,nao:].T.conj()
    mf.kernel(dm)
    mf.canonicalize(mf.mo_coeff, mf.mo_occ)
    mf.analyze()
    print(mf.spin_square())
    print(mf.e_tot - -76.2760115704274)


#!/usr/bin/env python
# ------------------------------------------------------------------------------------------------------%
# Created by "Thieu Nguyen" at 14:14, 20/09/2020                                                        %
#
#       Email:      nguyenthieu2102@gmail.com                                                           %
#       Homepage:   https://www.researchgate.net/profile/Thieu_Nguyen6                                  %
#       Github:     https://github.com/thieu1995                                                  %
#-------------------------------------------------------------------------------------------------------%

from numpy import ones, zeros, concatenate

D = [9, 11, 7, 6, 9, 38, 48, 2, 3, 3, 7, 7, 5, 10, 7, 14, 3, 4, 4, 2, 5, 9, 5, 7, 4, 22, 10, 10, 4, 3, 4, 5,
     30, 118, 153, 158, 126, 126, 126, 76, 74, 86, 86, 30, 25, 25, 25, 30, 30, 30, 59, 59, 59, 59, 64, 64, 64]
gn = [0, 0, 14, 1, 2, 0, 0, 2, 1, 3, 4, 9, 3, 10, 11, 15, 4, 4, 5, 3, 8, 10, 8, 7, 7, 86, 3, 9, 1, 8, 1, 6, 30, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 105, 24, 24, 24,
      29, 29, 29, 14, 14, 14, 14, 0, 0, 0]
hn = [8, 9, 0, 4, 4, 32, 38, 0, 1, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 3, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 108, 148, 148, 116, 116, 116, 76, 74, 76, 76, 0, 1,
      1, 1, 1, 1, 1, 1, 1, 1, 1, 6, 6, 6]

# bound constraint definitions for all 53 test functions
xmin1 = [0, 0, 0, 0, 1000, 0, 100, 100, 100]
xmax1 = [10, 200, 100, 200, 2000000, 600, 600, 600, 900]
xmin2 = [10 ** 4, 10 ** 4, 10 ** 4, 0, 0, 0, 100, 100, 100, 100, 100]
xmax2 = [0.819 * 10 ** 6, 1.131 * 10 ** 6, 2.05 * 10 ** 6, 0.05074, 0.05074, 0.05074, 200, 300, 300, 300, 400]
xmin3 = [1000, 0, 2000, 0, 0, 0, 0]
xmax3 = [2000, 100, 4000, 100, 100, 20, 200]
xmin4 = [0, 0, 0, 0, 1e-5, 1e-5]
xmax4 = [1, 1, 1, 1, 16, 16]
xmin5 = zeros(D[4])
xmax5 = [100, 200, 100, 100, 100, 100, 200, 100, 200]
xmin6 = zeros(D[5])
xmax6 = [90, 150, 90, 150, 90, 90, 150, 90, 90, 90, 150, 150, 90, 90, 150, 90, 150, 90, 150, 90, 1, 1.2,
         1, 1, 1, 0.5, 1, 1, 0.5, 0.5, 0.5, 1.2, 0.5, 1.2, 1.2, 0.5, 1.2, 1.2]
xmin7 = zeros(D[6])
xmin7[23] = xmin7[25] = xmin7[27] = xmin7[30] = 0.849999
xmax7 = ones(D[6])
xmax7[3] = 140
xmax7[24] = xmax7[26] = xmax7[31] = xmax7[34] = xmax7[36] = xmax7[28] = 30
xmax7[1] = xmax7[2] = xmax7[4] = xmax7[12:15] = 90
xmax7[0] = xmax7[5:12] = xmax7[15:20] = 35
xmin8 = [0, -0.51]
xmax8 = [1.6, 1.49]
xmin9 = [0.5, 0.5, -0.51]
xmax9 = [1.4, 1.4, 1.49]
xmin10 = [0.2, -2.22554, -0.51]
xmax10 = [1, -1, 1.49]
xmin11 = [0, 0, 0, 0, -0.51, -0.51, 0]
xmax11 = [20, 20, 10, 10, 1.49, 1.49, 40]
xmin12 = [0, 0, 0, -0.51, -0.51, -0.51, -0.51]
xmax12 = [100, 100, 100, 1.49, 1.49, 1.49, 1.49]
xmin13 = [27, 27, 27, 77.51, 32.51]
xmax13 = [45, 45, 45, 102.49, 45.49]
xmin14 = [0.51, 0.51, 0.51, 250, 250, 250, 6, 4, 40, 10]
xmax14 = [3.49, 3.49, 3.49, 2500, 2500, 2500, 20, 16, 700, 450]
xmin15 = [2.6, 0.7, 17, 7.3, 7.3, 2.9, 5]
xmax15 = [3.6, 0.8, 28, 8.3, 8.3, 3.9, 5.5]
xmin16 = 0.001 * ones(D[15])
xmax16 = +5 * ones(D[15])
xmin17 = [0.05, 0.25, 2.00]
xmax17 = [2, 1.3, 15.0]
xmin18 = [0.51, 0.51, 10, 10]
xmax18 = [99.49, 99.49, 200, 200]
xmin19 = [0.125, 0.1, 0.1, 0.1]
xmax19 = [2, 10, 10, 2]
xmin20 = zeros(D[19])
xmax20 = 1 * ones(D[19])
xmin21 = [60, 90, 1, 0, 2]
xmax21 = [80, 110, 3, 1000, 9]
xmin22 = [16.51, 13.51, 13.51, 16.51, 13.51, 47.51, 0.51, 0.51, 0.51]
xmax22 = [96.49, 54.49, 51.49, 46.49, 51.49, 124.49, 3.49, 6.49, 6.49]
xmin23 = [0, 0, 0, 0, 0]
xmax23 = [60, 60, 90, 90, 90]
xmin24 = [10, 10, 100, 0, 10, 100, 1]
xmax24 = [150, 150, 200, 50, 150, 300, 3.14]
xmin25 = [1, 1, 1e-6, 1]
xmax25 = [16, 16, 16 * 1e-6, 16]
xmin26 = concatenate((6.51 * ones(8), 0.51 * ones(14)), axis=0)
xmax26 = concatenate((76.49 * ones(8), 4.49 * ones(4), 9.49 * ones(10)), axis=0)
xmin27 = 0.645e-4 * ones(D[26])
xmax27 = 50e-4 * ones(D[26])
xmin28 = [125, 10.5, 4.51, 0.515, 0.515, 0.4, 0.6, 0.3, 0.02, 0.6]
xmax28 = [150, 31.5, 50.49, 0.6, 0.6, 0.5, 0.7, 0.4, 0.1, 0.85]
xmin29 = [20, 1, 20, 0.1]
xmax29 = [50, 10, 50, 60]
xmin30 = [0.51, 0.6, 0.51]
xmax30 = [70.49, 3, 42.49]
xmin31 = 12. * ones(4)
xmax31 = 60. * ones(4)
xmin32 = [78, 33, 27, 27, 27]
xmax32 = [102, 45, 45, 45, 45]
xmin33 = 0.001 * ones(D[32])
xmax33 = ones(D[32])
xmin34 = -1 * ones(D[33])
xmax34 = +1 * ones(D[33])
xmin35 = -1 * ones(D[34])
xmax35 = +1 * ones(D[34])
xmin36 = -1 * ones(D[35])
xmax36 = +1 * ones(D[35])
xmin37 = -1 * ones(D[36])
xmin37[117: 126] = 0
xmax37 = +1 * ones(D[36])
xmin38 = -1 * ones(D[37])
xmin38[117: 126] = 0
xmax38 = +1 * ones(D[37])
xmin39 = -1 * ones(D[38])
xmin39[117: 126] = 0
xmax39 = +1 * ones(D[38])
xmin40 = -1 * ones(D[39])
xmin40[75: 76] = 0
xmax40 = +1 * ones(D[39])
xmax40[75: 76] = 2
xmin41 = -1 * ones(D[40])
xmax41 = +1 * ones(D[40])
xmin42 = -1 * ones(D[41])
xmin42[75: 76] = 0
xmin42[77: 86] = 0
xmax42 = +1 * ones(D[41])
xmax42[75: 76] = 2
xmax42[77: 86] = 500
xmin43 = -1 * ones(D[42])
xmin43[75: 76] = 0
xmin43[77: 86] = 0
xmax43 = +1 * ones(D[42])
xmax43[75: 76] = 2
xmax43[77: 86] = 500
xmin44 = 40 * ones(D[43])
xmax44 = 1960 * ones(D[43])
xmin45 = zeros(D[44])
xmax45 = +90 * ones(D[44])
xmin46 = zeros(D[45])
xmax46 = +90 * ones(D[45])
xmin47 = zeros(D[46])
xmax47 = +90 * ones(D[46])
xmin48 = zeros(D[47])
xmax48 = +90 * ones(D[47])
xmin49 = zeros(D[48])
xmax49 = +90 * ones(D[48])
xmin50 = zeros(D[49])
xmax50 = +90 * ones(D[49])
xmin51 = zeros(D[50])
xmax51 = 10. * ones(D[50])
xmin52 = zeros(D[51])
xmax52 = 10. * ones(D[51])
xmin53 = zeros(D[52])
xmax53 = 10. * ones(D[52])
xmin54 = zeros(D[53])
xmax54 = 10. * ones(D[53])
xmin55 = zeros(D[54])
xmax55 = 10. * ones(D[54])
xmin56 = zeros(D[55])
xmax56 = 10. * ones(D[55])
xmin57 = zeros(D[56])
xmax57 = 10. * ones(D[56])

def benchmark_function(idx):
    # % prob_k -> Index of problem.
    # % D[idx  -> Dimension of the problem.
    # % par.g  -> Number of inequility constraints.
    # % par.h  -> Number of equality constraints.
    # % par.xmin -> lower bound of decision variables.
    # % par.xmax -> upper bound of decision variables.
    xmin = globals()['xmin' + str(idx)]
    xmax = globals()['xmax' + str(idx)]
    return {"D": D[idx-1], "g": gn[idx-1], "h": hn[idx-1], "xmin": xmin, "xmax": xmax}


import os
# os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   # see issue #152
# os.environ["CUDA_VISIBLE_DEVICES"]="2"

import torch as T
from torch.autograd import Variable
import numpy as np
import pdb
import copy
import time

cuda = True
if cuda:
    floatX = T.cuda.FloatTensor
    intX = T.cuda.IntTensor
    byteX = T.cuda.ByteTensor
    longX = T.cuda.LongTensor
else:
    floatX = T.FloatTensor
    intX = T.IntTensor
    byteX = T.ByteTensor
    longX = T.LongTensor

def m_eye(n, k=0):
    assert k < n and k >= 0
    if k == 0:
        return T.eye(n).type(floatX)
    else:
        return T.cat((T.cat((T.zeros(n-k, k), T.eye(n-k)), dim=1), T.zeros(k, n)), dim=0).type(floatX)

def ctc_loss(pred, pred_len, token, token_len, blank=0):
    '''
    :param pred: (Time, batch, voca_size+1)
    :param pred_len: (batch)
    :param token: (batch, U)
    :param token_len: (batch)
    '''
    Time, batch = pred.size(0), pred.size(1)
    U = token.size(1)
    eps = 0

    # token_with_blank
    token_with_blank = T.cat((T.zeros(batch, U, 1).type(longX), token[:, :, None]), dim=2).view(batch, -1)    # (batch, 2U)
    token_with_blank = T.cat((token_with_blank, T.zeros(batch, 1).type(longX)), dim=1)  # (batch, 2U+1)
    length = token_with_blank.size(1)

    pred = pred[T.arange(0, Time).type(longX)[:, None, None], T.arange(0, batch).type(longX)[None, :, None], token_with_blank[None, :]]  # (T, batch, 2U+1)

    # recurrence relation
    sec_diag = T.cat((T.zeros((batch, 2)).type(floatX), T.ne(token_with_blank[:, :-2], token_with_blank[:, 2:]).type(floatX)), dim=1) * T.ne(token_with_blank, blank).type(floatX)	# (batch, 2U+1)
    recurrence_relation = (m_eye(length) + m_eye(length, k=1)).repeat(batch, 1, 1) + m_eye(length, k=2).repeat(batch, 1, 1) * sec_diag[:, None, :]	# (batch, 2U+1, 2U+1)

    # alpha
    alpha_t = T.cat((pred[0, :, :2], T.zeros(batch, 2*U-1).type(floatX)), dim=1) # (batch, 2U+1)
    probability = alpha_t[None] # (1, batch, 2U+1)

    # dynamic programming
    # (T, batch, 2U+1)
    for t in T.arange(1, Time).type(longX):
        alpha_t = T.bmm(alpha_t[:, None], recurrence_relation)[:, 0] * pred[t]
        probability = T.cat((probability, alpha_t[None]), dim=0)

    labels_2 = probability[pred_len-1, T.arange(batch).type(longX), 2*token_len-1]
    labels_1 = probability[pred_len-1, T.arange(batch).type(longX), 2*token_len]
    labels_prob = labels_2 + labels_1

    cost = -T.log(labels_prob+eps)
    return cost

def log_batch_dot(alpha_t, rec):
    '''
    alpha_t: (batch, 2U+1)
    rec: (batch, 2U+1, 2U+1)
    '''
    eps_nan = -1e8
    # a+b
    _sum = alpha_t[:, :, None] + rec
    _max_sum = T.max(_sum, dim=1)[0]
    nz_mask1 = T.gt(_max_sum, eps_nan) # max > eps_nan
    nz_mask2 = T.gt(_sum, eps_nan)     # item > eps_nan

    # a+b-max
    _sum = _sum - _max_sum[:, None]

    # exp
    _exp = T.zeros_like(_sum).type(floatX)
    _exp[nz_mask2] = T.exp(_sum[nz_mask2])

    # sum exp
    _sum_exp = T.sum(_exp, dim=1)

    out = T.ones_like(_max_sum).type(floatX) * eps_nan
    out[nz_mask1] = T.log(_sum_exp[nz_mask1]) + _max_sum[nz_mask1]
    return out

def log_sum_exp_axis(a, uniform_mask=None, dim=0):
    assert dim == 0
    eps_nan = -1e8
    eps = 1e-26
    _max = T.max(a, dim=dim)[0]

    if not uniform_mask is None:
        nz_mask2 = T.gt(a, eps_nan) * uniform_mask
        nz_mask1 = T.gt(_max, eps_nan) * T.ge(T.max(uniform_mask, dim=dim)[0], 1)
    else:
        nz_mask2 = T.gt(a, eps_nan)
        nz_mask1 = T.gt(_max, eps_nan)

    # a-max
    a = a - _max[None]

    # exp
    _exp_a = T.zeros_like(a).type(floatX)
    _exp_a[nz_mask2] = T.exp(a[nz_mask2])

    # sum exp
    _sum_exp_a = T.sum(_exp_a, dim=dim)

    out = T.ones_like(_max).type(floatX) * eps_nan
    out[nz_mask1] = T.log(_sum_exp_a[nz_mask1] + eps) + _max[nz_mask1]
    return out

def log_sum_exp(*arrs):
#    return T.max(a.clone(), b.clone()) + T.log1p(T.exp(-T.abs(a.clone()-b.clone())))
    c = T.cat(list(map(lambda x:x[None], arrs)), dim=0)
    return log_sum_exp_axis(c, dim=0)

def ctc_loss_log(pred, pred_len, token, token_len, blank=0):
    '''
    :param pred: (Time, batch, voca_size+1)
    :param pred_len: (batch)
    :param token: (batch, U)
    :param token_len: (batch)
    '''
    Time, batch = pred.size(0), pred.size(1)
    U = token.size(1)
    eps_nan = -1e8

    # token_with_blank
    token_with_blank = T.cat((T.zeros(batch, U, 1).type(longX), token[:, :, None]), dim=2).view(batch, -1)    # (batch, 2U)
    token_with_blank = T.cat((token_with_blank, T.zeros(batch, 1).type(longX)), dim=1)  # (batch, 2U+1)
    length = token_with_blank.size(1)

    pred = pred[T.arange(0, Time).type(longX)[:, None, None], T.arange(0, batch).type(longX)[None, :, None], token_with_blank[None, :]]  # (T, batch, 2U+1)

    # recurrence relation
    sec_diag = T.cat((T.zeros((batch, 2)).type(floatX), T.ne(token_with_blank[:, :-2], token_with_blank[:, 2:]).type(floatX)), dim=1) * T.ne(token_with_blank, blank).type(floatX)	# (batch, 2U+1)
    recurrence_relation = (m_eye(length) + m_eye(length, k=1)).repeat(batch, 1, 1) + m_eye(length, k=2).repeat(batch, 1, 1) * sec_diag[:, None, :]	# (batch, 2U+1, 2U+1)
    recurrence_relation = eps_nan * (T.ones_like(recurrence_relation) - recurrence_relation)

    # alpha
    alpha_t = T.cat((pred[0, :, :2], T.ones(batch, 2*U-1).type(floatX)*eps_nan), dim=1) # (batch, 2U+1)
    probability = alpha_t[None] # (1, batch, 2U+1)

    # dynamic programming
    # (T, batch, 2U+1)
    for t in T.arange(1, Time).type(longX):
        alpha_t = log_batch_dot(alpha_t, recurrence_relation) + pred[t]
        probability = T.cat((probability, alpha_t[None]), dim=0)

    labels_2 = probability[pred_len-1, T.arange(batch).type(longX), 2*token_len-1]
    labels_1 = probability[pred_len-1, T.arange(batch).type(longX), 2*token_len]
    labels_prob = log_sum_exp(labels_2, labels_1)
#     pdb.set_trace()

    cost = -labels_prob
    return cost

def ctc_cost(out, targets, sizes, target_sizes):
#    A batched version for uni_alpha_cost
#    param out: (Time, batch, voca_size+1)
#    param targets: targets without splited
#    param sizes: size for out (N)
#    param target_sizes: size for targets (N)

    Time = out.size(0)
    pred = T.nn.functional.log_softmax(out, dim=-1)

    offset = 0
    batch = target_sizes.size(0)
    target_max = target_sizes.max().item()
    target = T.zeros(batch, target_max).type(longX)

    for index, (target_size, size) in enumerate(zip(target_sizes, sizes)):
        target[index, :target_size.item()] = targets[offset: offset+target_size.item()].data
        offset += target_size.item()

    if not cuda:
        costs = ctc_loss_log(pred.cpu(), sizes.data.type(longX), target, target_sizes.data.type(longX))
    else:
        costs = ctc_loss_log(pred, sizes.data.type(longX), target, target_sizes.data.type(longX))
    return costs.sum()

def test_seg_ctc(use_mine=True):
    # (T, voca_size+1)
    pred_np = np.array([[0.5, 0.4, 0.1], [0.3, 0.1, 0.6], [0.7, 0.2, 0.1], [0.3, 0.5, 0.2]])[:, None]
    pred_np = np.log(np.tile(pred_np, (1,2,1)))
    # (U)
    token_np = np.array([2, 2, 1, 2])

    pred = Variable(floatX(pred_np), requires_grad=True)
    token = Variable(T.IntTensor(token_np))
    sizes = Variable(T.IntTensor(np.array([4, 4])))
    target_sizes = Variable(T.IntTensor(np.array([2, 2])))

    for i in range(40):
        if use_mine:
            cost = ctc_cost(pred, token, sizes, target_sizes)
            print(cost.data.item())
        else:
            from warpctc_pytorch import CTCLoss
            criterion = CTCLoss().cuda()
            cost = criterion(pred, token, sizes, target_sizes)
            print(cost.data.item())

        optimizer = T.optim.SGD([pred], lr=3e-2, momentum=0.9, nesterov=True)
        optimizer.zero_grad()
        cost.backward()
        optimizer.step()

if __name__ == '__main__':
    print('_________')
    test_seg_ctc(use_mine=True)