# Copyright (C) 2012 Victor Semionov
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#  * Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys
import os

import math

import numpy as np

import model

import params
import output
import unitconv

class ComputationError(Exception):
    pass

class PulseTrainMock(object):

    def __init__(self, pulse, count, period):
        self.pulse = pulse
        self.count = count
        self.period = period

def create_medium(ref_inversion):
    initial_inversion = None
    if ref_inversion is not None:
        initial_inversion = initial_inversion = model.inversion.UniformInversion(ref_inversion)
    doping_agent = model.dopant.DopingAgent(params.lasing_wavelen, params.dopant_xsection, params.dopant_upper_lifetime, params.dopant_lower_lifetime, params.dopant_branching_ratio, params.dopant_concentration)
    medium = model.medium.ActiveMedium(initial_inversion, doping_agent, params.medium_radius, params.medium_length, params.medium_refr_idx)
    return medium

def create_beam():
    pulse_photon_count = model.energy.photon_count(params.lasing_wavelen, params.pulse_energy)
    ref_fluence = params.beam_class.ref_fluence(params.beam_radius, pulse_photon_count)
    beam = params.beam_class(params.beam_radius, ref_fluence)
    return beam

def create_pulse(active_medium, beam, rho, phi, ret_time_trunc_rel_error=False):
    fluence = beam.fluence(rho, phi)
    ref_density = params.pulse_class.ref_density(active_medium.light_speed, params.pulse_duration, fluence)
    pulse = params.pulse_class(-params.pulse_duration/2.0, params.pulse_duration, ref_density)
    scale, time_trunc_rel_error = model.error.pulse_scale(pulse, params.time_trunc_rtol)
    pulse = model.pulse.ExtendedPulse(pulse, scale)
    pulse = model.pulse.TruncatedPulse(pulse)
    return (pulse, time_trunc_rel_error) if ret_time_trunc_rel_error else pulse

def create_train(pulse):
    train = PulseTrainMock(pulse, params.train_pulse_count, params.train_pulse_period)
    return train

def create_depop_model(active_medium, depop_model_class):
    is_numerical = issubclass(depop_model_class, model.depop.NumericalDepopulationModel)
    if is_numerical:
        depop_model_kwargs = dict(rtol=params.depop_rate_rtol, min_samples=params.depop_rate_min_samples)
    else:
        depop_model_kwargs = {}
    if depop_model_class is params.depop_model_class:
        depop_model_kwargs.update(params.depop_model_extra_args)
    depop_model = depop_model_class(active_medium, **depop_model_kwargs)
    return depop_model

def compute_inversion(dirname):
    print output.div_line
    print "computing population inversion"

    active_medium = create_medium(None)
    pump_system = model.pump.PumpSystem(params.pump_wavelen, params.pump_duration, params.pump_power, params.pump_efficiency)
    depop_model = create_depop_model(active_medium, params.depop_model_class)
    inv = params.inverter_class(active_medium, pump_system, depop_model)

    ref_inversion = inv.invert(params.inversion_rtol, params.inversion_min_count_t)
    rate_evals = (len(inv.inversion) - 1) * inv.evals_per_step
    pump_energy = params.pump_duration * params.pump_power
    stored_energy = model.energy.energy(params.lasing_wavelen, ref_inversion * active_medium.volume)

    if params.verbose:
        print "count_t:", len(inv.T)
        print "depopulation rate evaluation count:", rate_evals

    if params.inversion_validate:
        print "validating uniform ASE-induced depopulation rate approximation"
        ross_num_model = depop_model if isinstance(depop_model, model.depop.RossNumericalASEModel) else model.depop.RossNumericalASEModel(active_medium, params.depop_rate_rtol, params.depop_rate_min_samples)
        rate_rel_stddev = ross_num_model.rate_rel_stddev(ref_inversion)
        unitconv.print_result("depopulation rate rel. std. deviation [{}]: {}", ("%",), (rate_rel_stddev,))
        if rate_rel_stddev > 10.0e-2:
            output.warn("uniform ASE-induced depopulation rate approximation is invalid")

    if isinstance(depop_model, model.depop.NumericalDepopulationModel):
        print "perturbing population inversion"
        perturb_depop_model = model.depop.PerturbedDepopulationModel(depop_model)
        perturb_inv = params.inverter_class(active_medium, pump_system, perturb_depop_model)
        perturb_ref_inversion = perturb_inv.invert(params.inversion_rtol, params.inversion_min_count_t)
        rel_error = model.error.perturbed_inversion_rel_error(ref_inversion, perturb_ref_inversion, params.inversion_rtol)
    else:
        rel_error = params.inversion_rtol

    gain_coef = ref_inversion * active_medium.doping_agent.xsection
    gain = math.exp(gain_coef * active_medium.length)

    ref_inversion_atol = ref_inversion * rel_error
    gain_atol = gain * (math.exp(ref_inversion_atol * active_medium.doping_agent.xsection * active_medium.length) - 1.0)
    stored_energy_atol = stored_energy * rel_error

    unitconv.print_result("pump energy [{}]: {}", ("mJ",), (pump_energy,))
    unitconv.print_result("population inversion [{}]: {} ~ {}", ("cm^-3",), (ref_inversion, ref_inversion_atol))
    unitconv.print_result("small signal gain: {} ~ {}", (), (gain, gain_atol))
    unitconv.print_result("stored energy [{}]: {} ~ {}", ("mJ",), (stored_energy, stored_energy_atol))

    if params.graphs:
        print output.status_writing
        dirname = output.init_dir(dirname)
        output.plot_inversion(dirname, inv)

    return ref_inversion, rel_error

def most_efficient_methods((int_types, amp_types), active_medium, input_beam, ref_pulse, quiet=False):
    if not quiet:
        print output.div_line
    if not quiet or params.verbose:
        print "determining most efficient method combination"

    min_xverse = (params.min_count_rho, params.min_count_phi)
    min_evo = (min_count_z, min_count_t) = (params.min_count_z, params.min_count_t)

    pulse_train = create_train(ref_pulse)

    best_method = None
    limit = 0

    for int_type in int_types:
        if params.verbose:
            print int_type.__name__

        integrator = model.integrator.DomainIntegrator(int_type)
        min_count_int = integrator.num_integrator.min_count
        min_count_amp_z = model.discrete.steps(model.discrete.divs(max(min_count_z, min_count_int)))
        min_count_amp_t = model.discrete.steps(model.discrete.divs(max(min_count_t, min_count_int)))
        min_count_amp = min_count_amp_z * min_count_amp_t
        int_limit = limit // min_count_amp
        try:
            (count_rho, count_phi), int_rel_error = model.error.min_integration_steps(integrator, min_xverse, params.int_rtol, int_limit, active_medium, input_beam)
        except model.exc.SoftLimitError:
            sys.exc_clear()
            continue
        except (ValueError, MemoryError):
            output.print_exception()
            print >>sys.stderr, "attempting to recover"
            sys.exc_clear()
            continue

        for amp_type in amp_types:
            if params.verbose:
                print amp_type.__name__

            count_xverse = count_rho * count_phi
            amp_limit = limit // count_xverse
            try:
                (count_z, count_t), amp_rel_error = model.error.min_amplification_steps(amp_type, min_evo, params.amp_rtol, amp_limit, active_medium, pulse_train, None, integrator)
            except model.exc.SoftLimitError:
                sys.exc_clear()
                continue
            except (ValueError, MemoryError):
                output.print_exception()
                print >>sys.stderr, "attempting to recover"
                sys.exc_clear()
                continue

            count = count_rho * count_phi * count_z * count_t
            is_best = False
            if best_method is None:
                is_best = True
            else:
                _, _, best_count, (best_amp_rel_error, best_int_rel_error) = best_method
                if count < best_count:
                    is_best = True
                elif count == best_count:
                    if (amp_rel_error + int_rel_error) < (best_amp_rel_error + best_int_rel_error):
                        is_best = True
            if is_best:
                best_method = (int_type, amp_type), (count_rho, count_phi, count_z, count_t), count, (amp_rel_error, int_rel_error)
                limit = count

    if best_method is None:
        raise ComputationError("no suitable numerical method combination found")

    (int_type, amp_type), (count_rho, count_phi, count_z, count_t), _, (amp_rel_error, int_rel_error) = best_method
    return (int_type, amp_type), (count_rho, count_phi, count_z, count_t), (amp_rel_error, int_rel_error)

def select_methods((int_types, amp_types), ref_inversion, quiet=False):
    active_medium = create_medium(ref_inversion)
    input_beam = create_beam()
    ref_pulse, time_trunc_rel_error = create_pulse(active_medium, input_beam, input_beam.rho_ref, input_beam.phi_ref, ret_time_trunc_rel_error=True)

    methods = most_efficient_methods((int_types, amp_types), active_medium, input_beam, ref_pulse, quiet)
    (int_type, amp_type), (count_rho, count_phi, count_z, count_t), (amp_rel_error, int_rel_error) = methods

    if params.verbose:
        print "int_type: %s; amp_type: %s" % (int_type.__name__, amp_type.__name__, )
        print "count_rho: %d; count_phi: %d" % (count_rho, count_phi)
        print "count_z: %d; count_t: %d" % (count_z, count_t)

    numerics = (int_type, amp_type), (count_rho, count_phi, count_z, count_t)
    rel_errors = time_trunc_rel_error, amp_rel_error, int_rel_error
    return numerics, rel_errors

def amplify_ref_pulse(dirname, num_types, counts, ref_inversion):
    print output.div_line
    print "amplifying ref. pulse"

    dirname = os.path.join(dirname, output.ref_pulse_rel_path)

    (int_type, amp_type), (_, _, count_z, count_t) = num_types, counts

    active_medium = create_medium(ref_inversion)
    input_beam = create_beam()
    rho, phi = input_beam.rho_ref, input_beam.phi_ref
    ref_pulse = create_pulse(active_medium, input_beam, rho, phi)

    integrator = model.integrator.DomainIntegrator(int_type)
    amp = amp_type(active_medium, count_z)

    num_density_out, _ = amp.amplify(rho, phi, ref_pulse, count_t)

    if active_medium.doping_agent.lower_lifetime in model.amplifier.ExactAmplifier.analytical_lower_lifetimes:
        exact_amp = model.amplifier.ExactOutputAmplifier(active_medium, count_z)
        exact_density_out, exact_population_final = exact_amp.amplify(rho, phi, ref_pulse, count_t)
    else:
        exact_density_out, exact_population_final = None, None

    fluence_out = integrator.integrate(amp.T, num_density_out) * active_medium.light_speed
    fluence_gain = fluence_out / input_beam.ref_fluence
    unitconv.print_result("fluence gain: {}", (), (fluence_gain,))

    if params.graphs:
        count_z = len(amp.Z)
        fluences = np.empty(count_z)
        for l in range(count_z):
            fluences[l] = integrator.integrate(amp.T, amp.density[l]) * active_medium.light_speed
        print output.status_writing
        dirname = output.init_dir(dirname)
        output.plot_output(dirname, input_beam, ref_pulse, params.pulse_duration, amp, fluences, exact_density_out, exact_population_final)

def amplify_train(dirname, num_types, counts, ref_inversion, quiet=False):
    if not quiet:
        print output.div_line
    if not quiet or params.verbose:
        print "amplifying pulse train"

    active_medium = create_medium(ref_inversion)
    input_beam = create_beam()
    ref_pulse = create_pulse(active_medium, input_beam, input_beam.rho_ref, input_beam.phi_ref)
    pulse_train = create_train(ref_pulse)

    int_type, amp_type = num_types
    count_rho, count_phi, count_z, count_t = counts
    integrator = model.integrator.DomainIntegrator(int_type)
    amp = amp_type(active_medium, count_z)

    radius = min(active_medium.radius, input_beam.rho_trunc)
    Rho = np.linspace(0.0, radius, count_rho)
    Phi = np.linspace(0.0, 2.0*math.pi, count_phi)

    output_fluence = np.empty((count_rho, count_phi))

    max_fluences = np.empty(params.train_pulse_count)
    output_photon_counts = np.empty(params.train_pulse_count)

    populations = [[None] * count_phi for _ in range(count_rho)]
    for m, rho in enumerate(Rho):
        for n, phi in enumerate(Phi):
            upper = np.vectorize(active_medium.initial_inversion.inversion)(rho, phi, amp.Z)
            lower = np.zeros(count_z)
            populations[m][n] = (upper, lower)

    norm_beam = params.beam_class(params.beam_radius, 1.0)
    norm_pulse = create_pulse(active_medium, norm_beam, norm_beam.rho_ref, norm_beam.phi_ref)
    amp._init_time(norm_pulse, count_t)
    norm_input_density = np.vectorize(norm_pulse.density)(amp.T)

    lower_decay = model.amplifier.lower_state_decay(active_medium, pulse_train)

    pulse_num_stride = params.pulse_num_stride
    for pnum in range(params.train_pulse_count):
        if not quiet or params.verbose:
            output.show_status((pnum, None), (pulse_num_stride, None), False)

        for m, rho in enumerate(Rho):
            for n, phi in enumerate(Phi):
                input_density = norm_input_density * input_beam.fluence(rho, phi)

                density_out, population_final = amp.amplify(rho, phi, None, None, T=amp.T, initial_population=populations[m][n], input_density=input_density)

                upper = np.copy(population_final[0])
                lower = population_final[1] * lower_decay
                populations[m][n] = (upper, lower)

                fluence_out = integrator.integrate(amp.T, density_out) * active_medium.light_speed
                output_fluence[m, n] = fluence_out

        if pnum == 0:
            ref_idx = np.unravel_index(output_fluence.argmax(), output_fluence.shape)
            ref_output_fluence = np.copy(output_fluence)
        max_fluences[pnum] = output_fluence[ref_idx]
        output_photon_counts[pnum] = integrator.integrate_base(active_medium, input_beam, Rho, Phi, output_fluence)

    del input_density, density_out, population_final, upper, lower
    del amp, output_fluence, populations, norm_input_density

    if not quiet or params.verbose:
        output.show_status((pnum+1, None), (pulse_num_stride, None), True)

    if not quiet or params.verbose:
        print "processing results"

    max_output_fluence = max_fluences[::-1].sum()
    max_output_fluence = model.energy.energy(params.lasing_wavelen, max_output_fluence)
    del max_fluences

    train_output_photon_count = output_photon_counts[::-1].sum()
    train_output_energy = model.energy.energy(params.lasing_wavelen, train_output_photon_count)

    rel_gain_decrease = 1.0 - output_photon_counts[-1] / output_photon_counts[0]

    if not quiet:
        if params.graphs:
            print output.status_writing
            ref_pulse_dir = os.path.join(dirname, output.ref_pulse_rel_path)
            dirname = output.init_dir(dirname)
            ref_pulse_dir = output.init_dir(ref_pulse_dir)
            output.plot_beam(ref_pulse_dir, input_beam, Rho, Phi, ref_output_fluence)
            output.plot_train(dirname, input_beam, active_medium, output_photon_counts)

    return max_output_fluence, output_photon_counts, train_output_energy, rel_gain_decrease

def report_results(ref_inversion, max_output_fluence, output_photon_counts, output_energy, rel_gain_decrease, inversion_rel_error, rel_errors):
    print output.div_line
    print "results:"

    active_medium = create_medium(ref_inversion)

    energy_rel_error = model.error.energy_rel_error(active_medium, inversion_rel_error, rel_errors)

    pump_energy = params.pump_duration * params.pump_power
    stored_energy = model.energy.energy(params.lasing_wavelen, ref_inversion * active_medium.volume)

    input_beam = create_beam()

    input_photon_count = input_beam.fluence_integral(active_medium.radius)
    input_energy = model.energy.energy(params.lasing_wavelen, input_photon_count)
    input_energy *= params.train_pulse_count

    energy_gain = output_energy / input_energy
    added_energy = output_energy - input_energy
    extraction_eff = added_energy / stored_energy
    total_eff = added_energy / pump_energy

    stored_energy_abs_error = inversion_rel_error * stored_energy
    output_energy_abs_error = energy_rel_error * output_energy
    energy_gain_abs_error = energy_rel_error * energy_gain

    extraction_eff_abs_error = (added_energy + output_energy_abs_error) / max(stored_energy - stored_energy_abs_error, 0.0) - extraction_eff
    total_eff_abs_error = output_energy_abs_error / pump_energy

    max_output_fluence_abs_error = max_output_fluence * energy_rel_error

    photon_count_first, photon_count_last = output_photon_counts[0], output_photon_counts[-1]
    photon_count_first_abs_error, photon_count_last_abs_error = photon_count_first * energy_rel_error, photon_count_last * energy_rel_error
    rel_gain_decrease_abs_error = 0.0
    if params.train_pulse_count > 1:
        rel_gain_decrease_abs_error = (photon_count_last + photon_count_last_abs_error) / max(photon_count_first - photon_count_first_abs_error, 0.0) - photon_count_last / photon_count_first

    unitconv.print_result("input energy [{}]: {}", ("mJ",), (input_energy,))
    unitconv.print_result("output energy [{}]: {} ~ {}", ("mJ",), (output_energy, output_energy_abs_error))
    unitconv.print_result("energy gain: {} ~ {}", (), (energy_gain, energy_gain_abs_error))
    unitconv.print_result("extraction efficiency [{}]: {} ~ {}", ("%",), (extraction_eff, extraction_eff_abs_error))
    unitconv.print_result("opt.-opt. efficiency [{}]: {} ~ {}", ("%",), (total_eff, total_eff_abs_error))
    unitconv.print_result("max. output fluence [{}]: {} ~ {}", ("J/cm^2",), (max_output_fluence, max_output_fluence_abs_error))
    unitconv.print_result("rel. gain decrease [{}]: {} ~ {}", ("%",), (rel_gain_decrease, rel_gain_decrease_abs_error))


"""
@author: mkowalska
"""
import os
from os.path import expanduser
import numpy as np
import matplotlib.pyplot as plt
import datetime
import time

from kcsd import csd_profile as CSD
from kcsd import ValidateKCSD1D, SpectralStructure

__abs_file__ = os.path.abspath(__file__)

def makemydir(directory):
    """
    Creates a new folder if it doesn't exist

    Parameters
    ----------
    directory: string
        directory

    Returns
    -------
    None
    """
    try:
        os.makedirs(directory)
    except OSError:
        pass
    os.chdir(directory)

def save_source_code(save_path, timestr):
    """
    Saves the source code.

    Parameters
    ----------
    save_path: string
        directory
    timestr: float

    Returns
    -------
    None
    """
    with open(save_path + '/source_code_' + str(timestr), 'w') as sf:
        sf.write(open(__file__).read())

def stability_M(csd_profile, csd_seed, n_src, ele_lims, true_csd_xlims,
                total_ele, noise=0, method='cross-validation', Rs=None,
                lambdas=None):
    """
    Investigates stability of reconstruction for different number of basis
    sources

    Parameters
    ----------
    csd_profile: function
        Function to produce csd profile.
    csd_seed: int
        Seed for random generator to choose random CSD profile.
    n_src: int
        Number of basis sources.
    ele_lims: list
        Boundaries for electrodes placement.
    true_csd_xlims: list
        Boundaries for ground truth space.
    total_ele: int
        Number of electrodes.
    noise: float
        Determines the level of noise in the data.
        Default: 0.
    method: string
        Determines the method of regularization.
        Default: cross-validation.
    Rs: numpy 1D array
        Basis source parameter for crossvalidation.
        Default: None.
    lambdas: numpy 1D array
        Regularization parameter for crossvalidation.
        Default: None.

    Returns
    -------
    obj_all: class object
    rms: float
        Normalized error of reconstruction.
    point_error_all: numpy array
        Normalized error of reconstruction calculated separetly at every point
        point of estimation space.
    eigenvalues: numpy array
        Eigenvalues of k_pot matrix.
    eigenvectors: numpy array
        Eigen vectors of k_pot matrix.
    """
    obj_all = []
    rms = np.zeros((len(n_src)))
    point_error_all = []
    eigenvectors = np.zeros((len(n_src), total_ele, total_ele))
    eigenvalues = np.zeros((len(n_src), total_ele))
    for i, value in enumerate(n_src):
        KK = ValidateKCSD1D(csd_seed, n_src_init=value, R_init=0.23,
                            ele_lims=ele_lims, true_csd_xlims=true_csd_xlims,
                            sigma=0.3, h=0.25, src_type='gauss', est_xres=0.01)
        obj, rms[i], point_error = KK.make_reconstruction(csd_profile,
                                                          csd_seed,
                                                          total_ele=total_ele,
                                                          noise=noise,
                                                          Rs=Rs,
                                                          lambdas=lambdas)
        ss = SpectralStructure(obj)
        eigenvectors[i], eigenvalues[i] = ss.evd()
        point_error_all.append(point_error)
        obj_all.append(obj)
    return obj_all, rms, point_error_all, eigenvalues, eigenvectors

def plot_M(n_src_init, rms, save_path):
    """
    Creates plot of relationship between RMS error and different number of
    basis sources

    Parameters
    ----------
    n_src_init: list
        List of number of basis sources.
    rms: numpy array
        Error of reconstruction.
    save_path: string
        Directory.

    Returns
    -------
    None
    """
    fig = plt.figure()
    plt.plot(n_src_init, rms, '--', marker='.')
    plt.xscale('log')
    plt.title('Stability of reconstruction for different number of basis '
              'sources')
    plt.xlabel('Number of basis sources')
    plt.ylabel('RMS')
    plt.show()
    save_as = (save_path + '/RMS_for_different_M')
    fig.savefig(os.path.join(save_path, save_as+'.png'))
    plt.close()

def plot_eigenvalues(eigenvalues, save_path, n_src):
    """
    Creates plot of eigenvalues of kernel matrix (k_pot) for different number
    of basis sources

    Parameters
    ----------
    eigenvalues: numpy array
        Eigenvalues of k_pot matrix.
    save_path: string
        Directory.
    n_src: list
        List of number of basis sources.

    Returns
    -------
    None
    """
    fig = plt.figure()
    for indx, i in enumerate(n_src):
        plt.plot(eigenvalues[indx], '--', marker='.', label='M='+str(i))
    plt.legend()
#    plt.title('Eigenvalue decomposition of kernel matrix for different number '
#              'of basis sources')
    plt.xlabel('Number of components')
    plt.ylabel('Eigenvalues')
    plt.yscale('log')
    save_as = (save_path + '/eigenvalues_for_different_M')
    fig.savefig(os.path.join(save_path, save_as+'.png'))
    plt.close()

def plot_max_eigenvalue_M(eigenvalues, save_path, n_src):
    """
    Creates plot of eigenvalues of kernel matrix (k_pot) for different number
    of basis sources

    Parameters
    ----------
    eigenvalues: numpy array
        Eigenvalues of k_pot matrix.
    save_path: string
        Directory.
    n_src: list
        List of number of basis sources.

    Returns
    -------
    None
    """
    fig = plt.figure(figsize=(8, 6))
    plt.plot(n_src, eigenvalues[:, 0], '--', marker='.', label=r'$\mu_1$')
    plt.legend()
    plt.title('First eigenvalue in the function of different number of basis'
              'sources')
    plt.xlabel('Number of basis sources')
    plt.xscale('log')
    plt.ylabel('Eigenvalues')
    plt.yscale('log')
    save_as = (save_path + '/max_eigenvalue_different_M')
    fig.savefig(os.path.join(save_path, save_as+'.png'))
    plt.close()

def plot_eigenvectors(eigenvectors, save_path, n_src):
    """
    Creates plot of eigenvectors of kernel matrix (k_pot) for different number
    of basis sources

    Parameters
    ----------
    eigenvectors: numpy array
        Eigenvectors of k_pot matrix.
    save_path: string
        Directory.
    n_src: list
        List of number of basis sources.

    Returns
    -------
    None
    """
    fig = plt.figure(figsize=(15, 15))
#    plt.suptitle('Eigenvalue decomposition of kernel matrix for different '
#                 'number of basis sources')
    for i in range(eigenvectors.shape[2]):
        plt.subplot(int(eigenvectors.shape[2]/2) + 1, 2, i + 1)
        for idx, j in enumerate(n_src):
            plt.plot(eigenvectors[idx, :, i].T, '--', marker='.',
                     label='M='+str(j))
        plt.ylabel('Eigenvectors')
        plt.title(r'$v_' + str(i + 1) + '$')
    plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
    plt.xlabel('Number of components')
    plt.tight_layout()
    plt.show()
    save_as = (save_path + '/eigenvectors_for_different_M')
    fig.savefig(os.path.join(save_path, save_as+'.png'))
    plt.close()

def plot_k_interp_cross(k_icross, save_path, n_src):
    """
    Creates plot of vectors of cross kernel matrix (k_interp_cross) for
    different number of basis sources

    Parameters
    ----------
    k_icross: numpy array
        List of cross kernel matrixes for different number of basis sources.
    save_path: string
        Directory.
    n_src: list
        List of number of basis sources.

    Returns
    -------
    None
    """
    fig = plt.figure(figsize=(k_icross[0].shape[1] + 5,
                              k_icross[0].shape[1] + 5))
#    plt.suptitle('Vectors of cross kernel matrix for different number '
#                 'of basis sources')
    for i in range(k_icross[0].shape[1]):
        plt.subplot(int(k_icross[0].shape[1]/2) + 1, 2, i + 1)
        for idx, j in enumerate(n_src):
            plt.plot(k_icross[idx][:, i], '--', marker='.',
                     label='M='+str(j))
            plt.title(r'$\tilde{K}_' + str(i + 1) + '$')
        plt.ylabel('Cross kernel')
    plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
    plt.xlabel('Number of estimation points')
    plt.tight_layout()
    plt.show()
    save_path = save_path + '/cross_kernel'
    makemydir(save_path)
    save_as = (save_path + '/cross_kernel_for_different_M')
    fig.savefig(os.path.join(save_path, save_as+'.png'))
    plt.close()

def plot_eigenvalue_lambda(eigenvalues, lambd, save_path, n_src):
    """
    Creates plot of eigenvalues of kernel matrix (k_pot) with lambda for
    different number of basis sources

    Parameters
    ----------
    eigenvalues: numpy array
        Eigenvalues of k_pot matrix.
    lambd: list
        Regularization parameter.
    save_path: string
        Directory.
    n_src: list
        List of number of basis sources.

    Returns
    -------
    None
    """
    fig = plt.figure(figsize=(7, 7))
    x = np.arange(1, eigenvalues.shape[1] + 1)
    for indx, i in enumerate(n_src):
        plt.plot(x, 1/(eigenvalues[indx] + lambd[indx]), '--', marker='.',
                 label='M='+str(i))
    plt.legend()
    plt.title(r'$\frac{1}{(\mu_j + \lambda)}$')
    plt.xlabel('Components number j')
    plt.ylabel(r'1/($\mu_j + \lambda)$')
    plt.yscale('log')
    plt.tight_layout()
    plt.show()
    save_path = save_path + '/cross_kernel'
    makemydir(save_path)
    save_as = (save_path + '/eigenvalues_coefficients_for_different_M')
    fig.savefig(os.path.join(save_path, save_as+'.png'))
    plt.close()

def plot_k_interp_cross_v(k_icross, eigenvectors, save_path, n_src):
    """
    Creates plot of product of cross kernel vectors and eigenvectors for
    different number of basis sources

    Parameters
    ----------
    k_icross: numpy array
        List of cross kernel matrixes for different number of basis sources.
    eigenvectors: numpy array
        Eigenvectors of k_pot matrix.
    save_path: string
        Directory.
    n_src: list
        List of number of basis sources.

    Returns
    -------
    None
    """
    fig = plt.figure(figsize=(k_icross[0].shape[1] + 5,
                              k_icross[0].shape[1] + 5))
#    plt.suptitle('Vectors of cross kernel and eigenvectors product for '
#                 'different number of basis sources')
    for i in range(k_icross[0].shape[1]):
        plt.subplot(int(k_icross[0].shape[1]/2) + 1, 2, i + 1)
        for idx, j in enumerate(n_src):
            plt.plot(np.dot(k_icross[idx], eigenvectors[idx, :, i]), '--',
                     marker='.', label='M='+str(j))
            plt.title(r'$\tilde{K}*v_' + str(i) + '$')
#        plt.ylabel(r'$\tilde{K}*v_' + str(i) + '$')
    plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
    plt.xlabel('Number of estimation points')
    plt.tight_layout()
    plt.show()
    save_path = save_path + '/cross_kernel'
    makemydir(save_path)
    save_as = (save_path + '/cross_kernel_eigenvector_product_for_different_M')
    fig.savefig(os.path.join(save_path, save_as+'.png'))
    plt.close()

def plot_k_pot(k_pot, save_path, n_src):
    """
    Creates plot of vectors of kernel matrix (k_pot) for
    different number of basis sources

    Parameters
    ----------
    k_pot: numpy array
        List of kernel matrixes for different number of basis sources.
    save_path: string
        Directory.
    n_src: list
        List of number of basis sources.

    Returns
    -------
    None
    """
    fig = plt.figure(figsize=(k_pot[0].shape[1] + 5,
                              k_pot[0].shape[1] + 5))
#    plt.suptitle('Vectors of kernel matrix for different number '
#                 'of basis sources')
    for i in range(k_pot[0].shape[1]):
        plt.subplot(int(k_pot[0].shape[1]/2) + 1, 2, i + 1)
        for idx, j in enumerate(n_src):
            plt.plot(k_pot[idx][:, i], '--', marker='.',
                     label='M='+str(j))
            plt.title(r'$K_' + str(i + 1) + '$')
        plt.ylabel('Kernel')
    plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
    plt.xlabel('Number of components')
    plt.tight_layout()
    plt.show()
    save_path = save_path + '/kernel'
    makemydir(save_path)
    save_as = (save_path + '/kernel_for_different_M')
    fig.savefig(os.path.join(save_path, save_as+'.png'))
    plt.close()

if __name__ == '__main__':
    HOME = expanduser('~')
    DAY = datetime.datetime.now()
    DAY = DAY.strftime('%Y%m%d')
    TIMESTR = time.strftime("%H%M%S")
    SAVE_PATH = HOME + "/kCSD_results/" + DAY + '/' + TIMESTR
    makemydir(SAVE_PATH)
    save_source_code(SAVE_PATH, time.strftime("%Y%m%d-%H%M%S"))

    CSD_PROFILE = CSD.gauss_1d_mono
    CSD_SEED = 15
#    N_SRC = [2, 4, 8, 16, 32, 64, 128, 256, 512]
    N_SRC = [2, 8, 16, 512]
    ELE_LIMS = [0.1, 0.9]  # range of electrodes space
    TRUE_CSD_XLIMS = [0., 1.]
    TOTAL_ELE = 10
    noise = None
    Rs = np.arange(0.1, 0.5, 0.1)
    lambdas = None
    method = 'cross-validation'
    OBJ, RMS, POINT_ERROR, eigenval, eigenvec = stability_M(CSD_PROFILE,
                                                            CSD_SEED,
                                                            N_SRC, ELE_LIMS,
                                                            TRUE_CSD_XLIMS,
                                                            TOTAL_ELE, Rs=Rs,
                                                            noise=noise,
                                                            lambdas=lambdas,
                                                            method=method)
    k_pot_list = []
    k_interp_cross_list = []
    lambdas = []
    for index in range(len(OBJ)):
        k_pot_list.append(OBJ[index].k_pot)
        k_interp_cross_list.append(OBJ[index].k_interp_cross)
        lambdas.append(OBJ[index].lambd)

    plot_M(N_SRC, RMS, SAVE_PATH)
    plot_eigenvalues(eigenval, SAVE_PATH, N_SRC)
    plot_max_eigenvalue_M(eigenval, SAVE_PATH, N_SRC)
    plot_eigenvectors(eigenvec, SAVE_PATH, N_SRC)
    plot_k_interp_cross(k_interp_cross_list, SAVE_PATH, N_SRC)
    plot_k_interp_cross_v(k_interp_cross_list, eigenvec, SAVE_PATH, N_SRC)
    plot_eigenvalue_lambda(eigenval, lambdas, SAVE_PATH, N_SRC)
    plot_k_pot(k_pot_list, SAVE_PATH, N_SRC)


"""Profiles to reduce data to 1-dimension."""

# Heavily inspired by pynbody (https://pynbody.github.io/).

from __future__ import annotations

from bisect import bisect
from copy import copy
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    List,
    Optional,
    Sequence,
    Tuple,
    Union,
)

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from numpy import ndarray
from pandas import DataFrame
from scipy.interpolate import interp1d

from .._logging import logger
from .._units import Quantity
from .._units import units as plonk_units
from ..utils.math import average
from ..utils.snap import dust_array_names, vector_array_names
from ..utils.strings import is_documented_by, pretty_array_name
from .extra import extra_profiles

if TYPE_CHECKING:
    from ..snap.snap import SnapLike

_aggregations = ('average', 'mean', 'median', 'std', 'sum')

class Profile:
    """Profiles.

    A profile is a binning of particles in Cartesian slices, or
    cylindrical or spherical shells around the origin, i.e. (0, 0, 0).
    For cylindrical profiles the cylindrical cells are perpendicular to
    the xy-plane, i.e. the bins are averaged azimuthally and in the
    z-direction. Cartesian profiles can be in the x-, y-, and
    z-directions.

    Parameters
    ----------
    snap
        The Snap object.
    ndim : optional
        The dimension of the profile. For ndim == 2, the radial binning
        is cylindrical in the xy-plane. For ndim == 3, the radial
        binning is spherical. For ndim == 1, the radial binning is
        Cartesian along the x-axis. Default is 2.
    cmin : optional
        The minimum coordinate for binning. Can be a string, e.g.
        '10 au', or a quantity with units, e.g. plonk.units('10 au').
        Defaults to minimum on the particles.
    cmax : optional
        The maximum coordinate for binning. Can be a string, e.g.
        '10 au', or a quantity with units, e.g. plonk.units('10 au').
        Defaults to the 99 percentile distance.
    n_bins : optional
        The number of radial bins. Default is 100.
    aggregation : optional
        The method to aggregate particle quantities in bins by. Options
        are 'average', 'mean', or 'median'. Here 'average' is a
        mass-weighted average. Default is 'average'.
    spacing : optional
        The spacing of radial bins. Can be 'linear' or 'log'. Default is
        'linear'.
    coordinate : optional
        The coordinate ('x', 'y', or 'z') for Cartesian profiles only,
        i.e. when ndim==1. Default is 'x'. For cylindrical and spherical
        profiles the coordinate is 'radius'.
    ignore_accreted : optional
        Ignore particles accreted onto sinks. Default is True.

    Examples
    --------
    Generate profile from snapshot.

    >>> prof = plonk.load_profile(snap=snap)
    >>> prof = plonk.load_profile(snap=snap, n_bins=300)
    >>> prof = plonk.load_profile(snap=snap, cmin='10 au', cmax='300 au')
    >>> prof = plonk.load_profile(snap=snap, spacing='log')

    To access a profile.

    >>> prof['surface_density']
    >>> prof['scale_height']

    To set a new profile.

    >>> prof['aspect_ratio'] = prof['scale_height'] / prof['radius']

    Alternatively use the add_profile decorator.

    >>> @prof.add_profile
    ... def mass(prof):
    ...     M = prof.snap['mass']
    ...     return prof.particles_to_binned_quantity('sum', M)

    Plot one or many quantities on the profile.

    >>> prof.plot('radius', 'density')
    >>> prof.plot('radius', ['angular_momentum_x', 'angular_momentum_y'])

    Plot a quantity on the profile with units.

    >>> units = {'position': 'au', 'surface_density'='g/cm^2'}
    >>> prof.plot('radius', 'surface_density', units=units)
    """

    _profile_aliases: Dict[str, str] = {}

    def __init__(
        self,
        snap: SnapLike,
        ndim: int = 2,
        cmin: Any = None,
        cmax: Any = None,
        n_bins: int = 100,
        aggregation: str = 'average',
        spacing: str = 'linear',
        coordinate: str = 'x',
        ignore_accreted: bool = True,
    ):

        self.snap = snap
        self.ndim = ndim
        self.aggregation = _check_aggregation(aggregation)
        self.spacing = _check_spacing(spacing)
        self.properties: Dict[str, Any] = {}

        self._profiles: Dict[str, Quantity] = {}
        self._profile_functions: Dict[str, Callable] = {}
        self._default_units = copy(self.snap._default_units)

        self._weights = self.snap['mass']
        self._mask = _setup_particle_mask(snap, ignore_accreted)
        self._x = _calculate_x(snap, self._mask, ndim, coordinate)
        self.range = _set_range(self._x, cmin, cmax)
        self.n_bins = n_bins

        self.bin_edges, self['size'] = _setup_bins(ndim, spacing, self.range, n_bins)
        self.bin_centers = 0.5 * (self.bin_edges[:-1] + self.bin_edges[1:])
        self._x = self._x.to(self.bin_edges.units)
        self._particle_bin = np.digitize(self._x.magnitude, self.bin_edges.magnitude)
        self.bin_indices = _set_particle_bin_indices(self._particle_bin, n_bins)

        if ndim == 1:
            self._coordinate = coordinate
        else:
            self._coordinate = 'radius'
        self._profiles[self._coordinate] = self.bin_centers
        self._profiles['number'] = np.histogram(
            self._x.magnitude, self.bin_edges.magnitude
        )[0] * plonk_units('dimensionless')

        # Add pre-defined profiles
        try:
            num_separate_dust = len(self.snap.num_particles_of_type['dust'])
        except KeyError:
            num_separate_dust = 0
        num_mixture_dust = self.snap.num_dust_species - num_separate_dust
        extra_profiles(self, num_separate_dust, num_mixture_dust)

    def add_profile(self, fn: Callable) -> Callable:
        """Decorate function to add profile to Profile.

        Parameters
        ----------
        fn
            A function that returns the profile as an array. The name of
            the function is the string with which to reference the array.

        Returns
        -------
        Callable
            The function which returns the array.
        """
        self._profile_functions[fn.__name__] = fn
        return fn

    def add_alias(self, name: str, alias: str) -> None:
        """Add alias to array.

        Parameters
        ----------
        name
            The name of the array.
        alias
            The alias to reference the array.
        """
        self._profile_aliases[alias] = name

    def loaded_profiles(self) -> List[str]:
        """Return a listing of loaded profiles."""
        return sorted(self._profiles.keys())

    def available_profiles(self) -> List[str]:
        """Return a listing of available profiles."""
        loaded = list(self.loaded_profiles())
        available = list(self._profile_functions.keys())
        snap_arrays = _1d_arrays(list(self.snap.available_arrays(verbose=True)))
        return sorted(set(loaded + available + snap_arrays))

    @property
    def default_units(self) -> Dict[str, Any]:
        """Profile default units."""
        return {
            key: self._default_units[key] for key in sorted(self._default_units.keys())
        }

    def set_units(self, **kwargs) -> Profile:
        """Set default unit for profiles.

        Parameters
        ----------
        kwargs
            Keyword arguments with keys as the profile name, e.g.
            'pressure', and with values as the unit as a string, e.g.
            'pascal'.

        Examples
        --------
        Set multiple default units.

        >>> profile.set_units(pressure='pascal', density='g/cm^3')
        """
        for key, val in kwargs.items():
            defaults = list(self.default_units) + list(self.snap.default_units)
            if key not in defaults:
                logger.info(f'adding profile {key} to default_units dict')
            self._default_units[key] = val

        return self

    def base_profile_name(self, name: str) -> str:
        """Get the base profile name from a string.

        For example, 'velocity_x' returns 'velocity', 'density' returns
        'density', 'dust_fraction_001' returns 'dust_fraction', 'x'
        returns 'position'.

        Parameters
        ----------
        name
            The name as a string

        Returns
        -------
        str
            The base name.
        """
        try:
            return self.snap.base_array_name(name)
        except ValueError:
            if name == self._coordinate:
                return 'position'
            name_root = '_'.join(name.split('_')[:-1])
            name_suffix = name.split('_')[-1]
            if name_root == '' and name_suffix in ('x', 'y', 'z'):
                return 'position'
            if name_root in self.snap._array_aliases:
                return self.snap._array_aliases[name_root]
            if name_suffix in ('x', 'y', 'z', 'mag'):
                return name_root
            if _str_is_int(name_suffix):
                return name_root
            return name

    def plot(
        self,
        x: str,
        y: Union[str, List[str]],
        units: Dict[str, Union[str, List[str]]] = None,
        std: str = None,
        label: Union[str, List[str]] = None,
        ax: Any = None,
        ax_kwargs={},
        **kwargs,
    ) -> Any:
        """Plot profile.

        Parameters
        ----------
        x
            The x axis to plot as a string.
        y
            The y axis to plot. Can be string or a list of strings.
        units
            The units of the plot as a dictionary. The keys correspond
            to quantities such as 'position', 'density', 'velocity', and
            so on. The values are strings representing units, e.g.
            'g/cm^3' for density.
        std : optional
            Add standard deviation on profile. Can be 'shading' or
            'errorbar'.
        label : optional
            A label for the plot. Can be a string or a list of strings,
            one per y.
        ax : optional
            A matplotlib Axes object to plot to.
        ax_kwargs
            Keyword arguments to pass to matplotlib Axes.
        **kwargs
            Keyword arguments to pass to Axes plot method.

        Returns
        -------
        ax
            The matplotlib Axes object.
        """
        if std is not None and std not in ('shading', 'errorbar'):
            raise ValueError('std must be "shading" or "errorbar"')

        ynames = _yname_from_yinput(y, self)
        labels = _labels(label, ynames)

        xunit = _get_unit(self, x, units)
        yunits = [_get_unit(self, y, units) for y in ynames]

        xdata = self[x].to(xunit)

        if ax is None:
            _, ax = plt.subplots()

        for yname, yunit, label in zip(ynames, yunits, labels):
            ydata = self[yname].to(yunit)
            if label is None:
                label = f'{pretty_array_name(yname)}'
                if ydata.units != plonk_units.dimensionless:
                    label = f'{pretty_array_name(yname)}'
                    label += f' [{ydata.units:~P}]'
            [line] = ax.plot(xdata.magnitude, ydata.magnitude, label=label, **kwargs)
            if std:
                color = line.get_color()
                _std_plot(self, xdata, ydata, yname, yunit, std, color, ax)

        ax.set_xlabel(f'{pretty_array_name(x)} [{xdata.units:~P}]')
        ax.legend()
        ax.set(**ax_kwargs)

        return ax

    def to_function(self, profile: str, **kwargs) -> Callable:
        """Create function via interpolation.

        The function is of the coordinate of the profile, e.g.
        'radius', and returns values of the selected profile, e.g.
        'scale_height'. The function is generated from the
        scipy.interpolate function interp1d.

        Parameters
        ----------
        profile
            The profile function to create as a string, e.g.
            'scale_height'.

        Returns
        -------
        Callable
            The function.

        Examples
        --------
        Select all particles within a scale height in a disc.

        >>> scale_height = prof.to_function('scale_height')
        >>> subsnap = snap[np.abs(snap['z']) < scale_height(snap['R'])]
        """
        coord = self.bin_centers
        prof = self[profile]

        def fn(x):
            nonlocal coord, prof
            _coord = coord.to(x.units).magnitude
            _prof = prof.magnitude
            y = interp1d(_coord, _prof, fill_value='extrapolate', **kwargs)(x.magnitude)
            return y * self[profile].units

        return fn

    def to_dataframe(
        self, columns: List[str] = None, units: List[str] = None
    ) -> DataFrame:
        """Convert Profile to DataFrame.

        Parameters
        ----------
        columns : optional
            A list of columns to add to the data frame. If None, add all
            loaded columns. Default is None.
        units : optional
            A list of units corresponding to columns add to the data
            frame. Units must be strings, and must be base units. I.e.
            'cm' not '10 cm'. If None, use default, i.e. cgs. Default is
            None.

        Returns
        -------
        DataFrame
        """
        data = dict()
        if columns is None:
            columns = self.loaded_profiles()
        if units is None:
            _units = list()
            for column in columns:
                try:
                    _units.append(self[column].units)
                except AttributeError:
                    _units.append(plonk_units('dimensionless'))
        else:
            _units = list()
            for unit in units:
                u = plonk_units(unit)
                if np.allclose(u.m, 1.0):
                    _units.append(u.units)
                else:
                    raise ValueError(
                        'Units must be strings, and must be base units. '
                        'I.e. "cm" not "10 cm".'
                    )
        if len(_units) != len(columns):
            raise ValueError('units and columns must have same length')
        for column, unit in zip(columns, _units):
            try:
                name = column + f' [{unit:~}]'
                array = self[column].to(unit).magnitude
            except AttributeError:
                name = column
                array = self[column]
            data[name] = array
        return pd.DataFrame(data)

    def particles_to_binned_quantity(
        self, aggregation: str, array: Quantity
    ) -> Quantity:
        """Calculate binned quantities from particles.

        This takes care of the bin indices and ignoring accreted
        particles (if requested in instantiating the profile).

        Parameters
        ----------
        aggregation
            The aggregation function that acts on particles in the
            radial bin.
        array
            The particle array.
        """
        if aggregation not in _aggregations:
            raise ValueError('Cannot determine aggregation method')

        _array = array[self._mask]
        binned_quantity = np.zeros(self.n_bins) * _array.units
        for idx, bin_ind in enumerate(self.bin_indices):
            if bin_ind.size == 0:
                continue
            if aggregation == 'average':
                val = average(
                    _array[bin_ind], weights=self._weights[self._mask][bin_ind],
                )
            elif aggregation == 'mean':
                val = np.mean(_array[bin_ind])
            elif aggregation == 'median':
                val = np.median(_array[bin_ind])
            elif aggregation == 'std':
                val = np.std(_array[bin_ind])
            elif aggregation == 'sum':
                val = np.sum(_array[bin_ind])

            binned_quantity[idx] = val

        return binned_quantity

    def __getitem__(self, name: str) -> Quantity:
        """Return the profile of a given kind."""
        name_root = '_'.join(name.split('_')[:-1])
        name_suffix = name.split('_')[-1]

        if name in self._profile_aliases:
            name = self._profile_aliases[name]
        if name in self._profiles:
            return self._profiles[name]
        if name in self._profile_functions:
            self._profiles[name] = self._profile_functions[name](self)
            return self._profiles[name]

        if name_suffix in _aggregations:
            aggregation = name_suffix
            array_name = name_root
        else:
            aggregation = self.aggregation
            array_name = name
        try:
            array: Quantity = self.snap[array_name]
        except ValueError as e:
            logger.error(e)
            raise ValueError(f'array "{array_name}" not available on snap')
        if array.ndim == 1:
            self._profiles[name] = self.particles_to_binned_quantity(aggregation, array)
            return self._profiles[name]
        raise ValueError(
            'Requested profile has array dimension > 1.\nTo access x-, y-, or '
            'z-components, or magnitude of vector quantities,\ntry, for '
            'example, prof["velocity_x"] or prof["momentum_mag"].\nTo '
            'access dust profiles, try, for example, prof["stopping_time_001"]'
        )

    def __setitem__(self, name: str, item: Quantity):
        """Set the profile directly."""
        if not isinstance(item, Quantity):
            raise ValueError('"item" must be pint Quantity')
        if item.shape[0] != self.n_bins:
            raise ValueError('Length of array does not match number of bins')
        if name in self.loaded_profiles():
            raise ValueError(
                'Attempting to overwrite existing profile. To do so, first delete the '
                'profile\nwith del prof["profile"], then try again.'
            )
        if name in self.available_profiles():
            raise ValueError(
                'Attempting to set profile already available. '
                'See prof.available_profiles().'
            )
        self._profiles[name] = item

    def __delitem__(self, name):
        """Delete a profile from memory."""
        del self._profiles[name]

    def __len__(self):
        """Length as number of bins."""
        return self.n_bins

    def __repr__(self):
        """Dunder repr method."""
        return self.__str__()

    def __str__(self):
        """Dunder str method."""
        return f'<plonk.Profile "{self.snap.file_path.name}">'

    def _ipython_key_completions_(self):
        """Tab completion for IPython __getitem__ method."""
        return self.available_profiles()

@is_documented_by(Profile)
def load_profile(
    snap: SnapLike,
    ndim: int = 2,
    cmin: Any = None,
    cmax: Any = None,
    n_bins: int = 100,
    aggregation: str = 'average',
    spacing: str = 'linear',
    coordinate: str = 'x',
    ignore_accreted: bool = True,
) -> Profile:
    logger.debug(f'Loading profile: {snap.file_path.name}')
    return Profile(
        snap=snap,
        ndim=ndim,
        cmin=cmin,
        cmax=cmax,
        n_bins=n_bins,
        aggregation=aggregation,
        spacing=spacing,
        coordinate=coordinate,
        ignore_accreted=ignore_accreted,
    )

def _setup_particle_mask(snap: SnapLike, ignore_accreted: bool) -> ndarray:
    if ignore_accreted is False:
        return np.ones(len(snap), dtype=bool)
    h: Quantity = snap['h']
    return h > 0

def _calculate_x(snap: SnapLike, mask: ndarray, ndim: int, coordinate: str) -> Quantity:
    pos: Quantity = snap['position']
    pos = pos[mask]
    if ndim == 1:
        if coordinate == 'x':
            return pos[:, 0]
        if coordinate == 'y':
            return pos[:, 1]
        if coordinate == 'z':
            return pos[:, 2]
        raise ValueError('coordinate must be "x", "y", or "z" for ndim==1')
    if ndim == 2:
        return np.sqrt(pos[:, 0] ** 2 + pos[:, 1] ** 2)
    if ndim == 3:
        return np.sqrt(pos[:, 0] ** 2 + pos[:, 1] ** 2 + pos[:, 2] ** 2)
    raise ValueError('Unknown ndim: cannot calculate x array')

def _set_range(x: Quantity, cmin: Any, cmax: Any) -> Tuple[float, float]:
    if cmin is None:
        rmin = x.min()
    else:
        rmin = Quantity(cmin)
        if not rmin.dimensionality == Quantity('cm').dimensionality:
            raise ValueError('must specify cmin units, e.g. cmin="10 au"')
    if cmax is None:
        rmax = np.percentile(x.magnitude, 99, axis=0) * x.units
    else:
        rmax = Quantity(cmax)
        if not rmax.dimensionality == Quantity('cm').dimensionality:
            raise ValueError('must specify cmin units, e.g. cmax="100 au"')

    return rmin, rmax

def _setup_bins(ndim: int, spacing: str, xrange: Quantity, n_bins: int) -> Quantity:
    bin_edges = _bin_edges(spacing, xrange, n_bins)
    if ndim == 1:
        bin_sizes = bin_edges[1:] - bin_edges[:-1]
    elif ndim == 2:
        bin_sizes = np.pi * (bin_edges[1:] ** 2 - bin_edges[:-1] ** 2)
    elif ndim == 3:
        bin_sizes = 4 / 3 * np.pi * (bin_edges[1:] ** 3 - bin_edges[:-1] ** 3)
    return bin_edges, bin_sizes

def _bin_edges(spacing: str, xrange: Quantity, n_bins: int) -> Quantity:
    if spacing == 'linear':
        bin_edges = (
            np.linspace(xrange[0].magnitude, xrange[1].magnitude, n_bins + 1)
            * xrange[0].units
        )
    elif spacing == 'log':
        bin_edges = (
            np.logspace(
                np.log10(xrange[0].magnitude),
                np.log10(xrange[1].magnitude),
                n_bins + 1,
            )
            * xrange[0].units
        )
    else:
        raise ValueError('Cannot determine spacing to setup bins')
    return bin_edges

def _set_particle_bin_indices(particle_bin: ndarray, n_bins: int) -> List[ndarray]:
    sortind = particle_bin.argsort()
    sort_pind = particle_bin[sortind]
    binind = list()
    prev_index = bisect(sort_pind, 0)
    for i in range(n_bins):
        new_index = bisect(sort_pind, i + 1)
        binind.append(np.sort(sortind[prev_index:new_index]))
        prev_index = new_index
    return binind

def _check_aggregation(method: str) -> str:
    if method in _aggregations:
        return method
    raise ValueError(
        f'Cannot determine aggregation method: choose from {_aggregations}'
    )

def _check_spacing(spacing: str) -> str:
    if spacing.lower() in ('lin', 'linear'):
        return 'linear'
    if spacing.lower() in ('log', 'logarithm', 'logarithmic'):
        return 'log'
    raise ValueError('Cannot determine spacing')

def _1d_arrays(arrays: list) -> list:
    _arrays = list()
    for array in arrays:
        if (array + '_x' in arrays) or (array + '_001' in arrays):
            pass
        else:
            _arrays.append(array)
    return _arrays

def _get_unit(profile, name, units):
    if name is None:
        return None
    base_name = profile.base_profile_name(name)
    if units is not None:
        if name in units:
            return 1 * plonk_units(units[name])
        if base_name in units:
            return 1 * plonk_units(units[base_name])
    if name in profile.default_units:
        return 1 * plonk_units(profile.default_units[name])
    if base_name in profile.default_units:
        return 1 * plonk_units(profile.default_units[base_name])
    return 1 * profile[name].units

def _yname_from_yinput(y, profile):
    if isinstance(y, str):
        if y not in profile.available_profiles():
            if y + '_001' in profile.available_profiles():
                return dust_array_names(
                    name=y, num_dust_species=profile.snap.num_dust_species
                )
            elif y + '_x' in profile.available_profiles():
                return vector_array_names(name=y)
            return [y]
        return [y]
    return y

def _labels(label, ynames):
    labels: Sequence[Optional[str]]
    if label is not None:
        if isinstance(label, str):
            return [label for _ in ynames]
        return label
    else:
        return [None for _ in ynames]

def _std_plot(profile, xdata, ydata, yname, yunit, std, color, ax):
    if yname.split('_')[-1] in _aggregations:
        _yname = '_'.join(yname.split('_')[:-1])
    else:
        _yname = yname
    try:
        y_std = profile[_yname + '_std']
    except ValueError:
        logger.warning('Cannot calculate standard deviation')
        return
    if profile.aggregation in ('std', 'sum'):
        y_mean = profile[_yname + '_mean']
    else:
        y_mean = ydata
    y_std = y_std.to(yunit).magnitude
    y_mean = y_mean.to(yunit).magnitude
    xdata = xdata.magnitude
    if std == 'shading':
        ax.fill_between(
            xdata, y_mean - y_std, y_mean + y_std, color=color, alpha=0.2,
        )
    elif std == 'errorbar':
        ax.errorbar(xdata, y_mean, yerr=y_std, linestyle='', color=color, alpha=0.5)

def _str_is_int(string: str) -> bool:
    try:
        int(string)
        return True
    except ValueError:
        return False


"""
Constant molecular properties to be used in other modules.

Arrays are ordered by atomic number for convenience. Atomic symbols are case
sensitive. This module should not depend on other modules.

Unit types
----------
Length: Angstrom (ang), Bohr (bohr), picometre (pm), nanometre (nm)

Angle: radian (rad), degree (deg)

Time: femtosecond (fs), picosecond (ps), atomic unit (au)

Mass: atomic mass unit (amu), electron mass (me), proton mass (mp),
kilogram (kg)

Energy: electron volt (ev), Hartree (har), kilocalorie per mole (kcm),
kilojoule per mole (kjm), reciprocal centimetre (cm), joule (j),
terahertz (thz)

For all types, 'auto' will give the default unit.

Attributes
----------
sym : ndarray
    List of atomic symbols up to Krypton. The ordering (with the
    exception of deuterium) yields the correct atomic number from
    ``sym.index(elem)``.
mass : ndarray
    List of atomic masses corresponding to the elements in `sym`.
covrad : ndarray
    List of covalent radii corresponding to the elements in `sym`.
lenunits : dict
    Dictionary of units of length and their conversions from the
    default (angstroms).
angunits : dict
    Dictionary of units of angle and their conversions from the
    default (radians).
timunits : dict
    Dictionary of units of time and their conversions from the
    default (femtoseconds).
masunits : dict
    Dictionary of units of mass and their conversions from the
    default (atomic mass units).
eneunits : dict
    Dictionary of units of energy and their conversion from the
    default (electron volts).
"""
import numpy as np

# Global constants
sym = np.array(['X', 'H', 'He', 'Li', 'Be', 'B', 'C', 'N', 'O', 'F', 'Ne',
                'Na', 'Mg', 'Al', 'Si', 'P', 'S', 'Cl', 'Ar', 'K', 'Ca',
                'Sc', 'Ti', 'V', 'Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn',
                'Ga', 'Ge', 'As', 'Se', 'Br', 'Kr'])
mass = np.array([0.00000000, 1.00782504, 4.00260325, 7.01600450, 9.01218250,
                 11.00930530, 12.00000000, 14.00307401, 15.99491464,
                 18.99840325, 19.99243910, 22.98976970, 23.98504500,
                 26.98154130, 27.97692840, 30.97376340, 31.97207180,
                 34.96885273, 39.96238310, 38.96370790, 39.0983,
                 40.078, 44.955908, 47.867, 50.9415, 51.9961, 54.938044,
                 55.845, 58.933194, 58.6934, 63.546, 65.38, 69.723, 72.630,
                 74.921595, 78.971, 79.904, 83.798])
covrad = np.array([0.000, 0.320, 1.600, 0.680, 0.352, 0.832, 0.720, 0.680,
                   0.680, 0.640, 1.120, 0.972, 1.100, 1.352, 1.200, 1.036,
                   1.020, 1.000, 1.568, 1.328, 0.992, 1.440, 1.472, 1.328,
                   1.352, 1.352, 1.340, 1.328, 1.620, 1.520, 1.448, 1.220,
                   1.168, 1.208, 1.220, 1.208, 1.600])
lenunits = dict(auto=1., ang=1., bohr=1./0.52917721, pm=100., nm=0.1)
angunits = dict(auto=1., rad=1., deg=180./np.pi)
timunits = dict(auto=1., fs=1., ps=0.001, au=1./0.024188843)
masunits = dict(auto=1., amu=1., me=1822.888486209, mp=1./1.00727647,
                kg=1.66053904e-27)
eneunits = dict(auto=1., ev=1., har=1./27.21138505, kcm=23.061, kjm=96.485,
                cm=8065.5, j=1.602176634e-19, thz=241.8)

def get_num(elem):
    """Returns atomic number from atomic symbol.

    Takes advantage of the fact that sym indices match atomic numbers.

    Parameters
    ----------
    elem : str or array_like
        The atomic symbol(s) to be parsed.

    Returns
    -------
    int or ndarray
        The atomic numbers corresponding to each symbol.
    """
    if isinstance(elem, str):
        return _find_index(elem)
    else:
        for atm in elem:
            if atm not in sym and atm[0] not in ['X', 'D']:
                raise ValueError('Unrecognized atomic symbol \'' + atm +
                                 '\'. Use X prefix for dummy atoms.')
        return np.array([_find_index(atm) for atm in elem])

def get_mass(elem):
    """Returns atomic mass from atomic symbol.

    Parameters
    ----------
    elem : str of array_like
        The atomic symbol(s) to be parsed.

    Returns
    -------
    float or ndarray
        The atomic masses corresponding to each symbol.
    """
    return mass[get_num(elem)]

def get_covrad(elem):
    """Returns covalent radius from atomic symbol.

    Parameters
    ----------
    elem : str of array_like
        The atomic symbol(s) to be parsed.

    Returns
    -------
    float or ndarray
        The atomic covalent radii corresponding to each symbol.
    """
    return covrad[get_num(elem)]

def unit_vec(v):
    """Returns a unit vector aligned with a given vector.

    Parameters
    ---------
    v : array_like
        The input, un-normalized vector.

    Returns
    -------
    ndarray
        The normalized (unit) vector."""
    vlen = np.linalg.norm(v)
    if np.isclose(vlen, 0):
        raise ValueError('Cannot make unit vector from zero vector.')
    else:
        return v / vlen

def arccos(val):
    """Returns the arccosine of an angle allowing for numerical errors.

    NumPy's arccos function is defined for the range [-1, 1], but
    returns NaN for :math:`|x| = 1 + \delta`, where :math:`\delta` is
    small.  This can be avoided by checking for limiting cases with
    ``numpy.isclose``.

    Parameters
    ----------
    val : float
        The x-coordinate on the unit circle.

    Returns
    -------
    float
        The angle intersecting the unit circle at x = val.
    """
    if np.isclose(val, -1):
        return np.pi
    elif np.isclose(val, 1):
        return 0.
    else:
        return np.arccos(val)

def conv(old='auto', new='auto'):
    """Returns conversion factor from old units to new units.

    Parameters
    ----------
    old : str, optional
        The units to be converted from. See different units types for
        defaults.
    new : str, optional
        The units to be converted to. See different units types for
        defaults.

    Returns
    -------
    float
        The conversion factor, new_units / old_units.
    """
    if old == new:
        return 1.
    for unittype in [lenunits, angunits, timunits, masunits, eneunits]:
        if old in unittype and new in unittype:
            return unittype[new] / unittype[old]

    raise ValueError('Units \'{}\' and \'{}\' unrecognized or '
                     'not of same unit type'.format(old, new))

def _find_index(string):
    """Determines if dummy or regular atom and returns index.

    Parameters
    ----------
    string : str
        The atomic symbol.

    Returns
    -------
    int
        The atomic number of the given atomic symbol.
    """
    if string[0] == 'X':
        return 0
    elif string  == 'D':
        return 1
    else:
        return np.where(sym == string)[0][0]


"""
This module generates waveforms for compact binary coalescences.

One first generates a high level specification JSON file that only specifies
meta information such as the number of waveforms and from which ranges to
sample the parameters. Then we actually sample the parameters for the waveforms
and saves those in a big JSON file. Next, we generate all these specified
waveforms and save them together with the configuration file in an hdf file.
"""

from __future__ import absolute_import, print_function

import os
import json
import pylab
import numpy as np
import h5py
from tqdm import tqdm

from pycbc.waveform import get_td_waveform

def data_generation_pipeline(meta_config_file='../data/config.json',
                             sample_config_file='../data/sample_config.json',
                             sample_file='../data/samples.h5'):
    """
    Run through the whole data generation pipeline.

    Arguments:

        meta_config_file: Path of a JSON file created by `generate_meta_config`

        sample_config_file: Path where to store the sample config JSON file

        sample_file: Path where to store the samples hdf data file

    Given a meta config file, sample the specified waveform parameters, compute
    the waveforms and store them.
    """
    WaveformConfigGenerator(meta_config_file,
                            sample_config_file).generate_and_save_config()
    WaveformGenerator(sample_config_file,
                      sample_file).generate_and_save_waveforms()

class WaveformGenerator:
    """
    Compute waveforms.
    """

    def __init__(self, config_file, output_file):
        """
        Initialize a WaveformGenerator.

        Arguments:

            config_file: A JSON configuration file generated by
                         `WaveformConfigGenerator`

            output_file: File path where to store the created hdf file
        """
        self.config_file = config_file
        self.output_file = output_file
        print("Reading the configuration {}...".format(config_file), end=' ')
        with open(self.config_file, 'r') as f:
            self.config = json.loads(f.read())
        print("DONE")

    def generate_and_save_waveforms(self):
        """
        Generate and save all specified waveforms.
        """
        # How many waveforms do we have to generate
        n_samples = len(self.config['injections'])
        duration = self.config['meta']['duration']
        sample_rate = self.config['meta']['sample_rate']

        # Collect indices of failures during waveform generation
        # Sometimes the model just fail to compute the requested waveform.
        # This indicates that one should choose different parameter regions.
        failed = []
        # Compute the waveforms
        print("Generate the waveforms...", end=' ')
        # If the duration is smaller 0, all waveforms are kept in full length
        if duration > 0:
            N = int(duration * sample_rate)
            waveforms = np.zeros((n_samples, N))

            for i, conf in enumerate(tqdm(self.config['injections'])):
                x = np.zeros(N)
                try:
                    # Compute the current waveform
                    hp, cp = tuple(map(np.array, get_td_waveform(**conf)))
                    # Get the correct indices to inject it into the noise
                    inj_time = conf['injection_time']
                    Iembed, Isignal = self._get_embedding_indices(hp,
                                                                  inj_time,
                                                                  N)
                    x[Iembed] += hp[Isignal]
                except Exception as e:
                    # If waveform fails, remember the index
                    # and keep the waveform as all zeros
                    failed.append(i)
                    err = type(e).__name__
                    print("Failure {}\n during waveform {}: {}".format(err,
                                                                       i,
                                                                       conf))
                waveforms[i, :] = x
        else:
            waveforms_raw = []
            for i, conf in enumerate(tqdm(self.config['injections'])):
                try:
                    # Compute the current waveform
                    hp, cp = tuple(map(np.array, get_td_waveform(**conf)))
                    waveforms_raw.append(hp)

                    # Get the correct indices to inject it into the noise
                    inj_time = conf['injection_time']
                except Exception as e:
                    # If waveform fails, remember the index
                    # and keep the waveform as all zeros
                    waveforms_raw.append(np.zeros(1))
                    failed.append(i)
                    err = type(e).__name__
                    print("Failure {}\n during waveform {}: {}".format(err,
                                                                       i,
                                                                       conf))
            N = max(map(len, waveforms_raw))
            waveforms = np.zeros((n_samples, N))
            for i, x in enumerate(waveforms_raw):
                waveforms[i, :len(x)] = x
                inj_time = np.argmax(x) / float(N)
                self.config['injections'][i]['injection_time'] = inj_time
        print("DONE")

        print("Save everything to file...", end=' ')
        with h5py.File(self.output_file, 'w') as f:
            f['waveforms'] = waveforms
            f['config'] = np.string_(json.dumps(self.config))
            f['failed'] = failed
        print("DONE")

    def plot_waveform(self, index=0):
        """
        Plot a waveform.

        This is mostly just for testing and visualization.
        """
        fig = pylab.figure()
        with h5py.File(self.output_file, 'r') as f:
            x = f['X'][index, :]
        pylab.plot(x)
        pylab.ylabel('Strain')
        pylab.xlabel('Time')
        return fig

    def _get_embedding_indices(self, f, frac, N):
        """
        Return matching indices of background and waveform for superposition.

        The waveform generating functions generate waveforms of different
        length, depending on the input parameters. Together with a variable
        injection time with in the sample, the waveform could be too
        long or too short towards both sides. Hence computing the injection
        indices is not completely straight forward.

        This is highly non-trivial magic and after extensive testing I never
        want to have to open this box again.
        """
        ni = int(frac * N)
        nn = N - 1
        sn = len(f) - 1
        si = np.argmax(f)
        start = ni - si
        nl, sl = abs(start), 0
        if start < 0:
            nl, sl = sl, nl
        end = (nn - ni) - (sn - si)
        nr, sr = ni + sn - si, sn
        if end < 0:
            nr, sr = nn, si + nn - ni
        Iembed = range(nl, nr + 1)
        Isignal = range(sl, sr + 1)
        if len(Iembed) == 0 or len(Isignal) == 0:
            print('Could not find appropriate ranges:')
            print('noise length: {}, signal length: {}'.format(nn, sn))
            print('noise: {} : {}, signal: {} : {}'.format(nl, nr, sl, sr))
        return Iembed, Isignal

# -----------------------------------------------------------------------------
# ------------------------------  Generate a Waveform Configuration JSON File
# -----------------------------------------------------------------------------

class WaveformConfigGenerator:
    """
    Sample and save waveform parameters in a waveform configuration JSON file.
    """

    waveform_config = {}

    def __init__(self, meta_config_file, output_file):
        """
        Initialize the WaveformConfigGenerator.

        Arguments:

            meta_config_file: A meta configuration file as generated by
            `generate_meta_config`. This meta config file can be hand tuned.

            output_file: The sample_configuration file holding information
            about parameters of each individual waveform.
        """
        self.output_file = output_file
        self.meta_config_file = meta_config_file
        print("Load meta configuration...", end=' ')
        with open(self.meta_config_file, 'r') as f:
            self.meta_config = json.loads(f.read())
        print("DONE")
        self.n_samples = self.meta_config['n_samples']
        self.default_parameters = self.meta_config['default_parameters']
        self.update_list = self.meta_config['update_list']

    def generate_and_save_config(self):
        """
        Generate and save the configuration JSON file for the injections.
        """
        self.waveform_config.update(self._injections_spec())
        self.waveform_config.update({'meta': self.meta_config})
        print("Write configuration to {}...".format(self.output_file), end=' ')
        with open(self.output_file, 'w') as f:
            json.dump(self.waveform_config, f, sort_keys=True, indent=2)
        print("DONE")

    def _injections_spec(self):
        """
        Generate and return specification of the injection waveforms.
        """
        injections = []
        print("Sampling the waveform parameters...", end=' ')
        for i in range(self.n_samples):
            cp = self.default_parameters.copy()
            # Set the id first as counter
            cp.update({'id': i})
            # Update other values in the update list with new samples
            # Some magic to figure out how to sample the samples
            # Very bad style with eval and everything. Works. Wouldn't touch.
            for update, pars in self.update_list.iteritems():
                args = '("' + update + '", ' + str(pars) + ')'
                fname = 'self._sample_' + update + args
                cp.update(eval(fname))
            injections.append(cp)
        print("DONE")
        return {'injections': injections}

    def _sample_distance(self, key, pars):
        """
        Draw a random source distance.
        """
        return {key: self._uniform_in_range(pars[0], pars[1])}

    def _sample_masses(self, key, pars):
        """
        Draw random compact object masses.
        """
        mass1 = self._uniform_in_range(pars[0], pars[1])
        mass2 = self._uniform_in_range(pars[0], pars[1])
        return {'mass1': mass1, 'mass2': mass2}

    def _sample_injection_time(self, key, pars):
        """
        Draw a random injection time.
        """
        return {key: self._uniform_in_range(pars[0], pars[1])}

    def _uniform_in_range(self, lower_bound, upper_bound):
        """
        Draw uniformly distributed number in a given range.
        """
        return np.random.rand() * (upper_bound - lower_bound) + lower_bound

def generate_meta_config(output_file='../data/meta_config.json'):
    """
    Generate a meta configuration JSON file for the data generation pipeline.

    This spits out a valid template of a meta configuration file as the one
    required by the `WaveformConfigGenerator`. Just run this once to have a
    valid template and then you can go ahead an tweak the individual numbers in
    that template. All the defaults are considered sane choices.
    """
    print("Generate meta config file {}...".format(output_file), end=' ')
    n_samples = 1024
    # If the duration is < 0.0 the whole waveform will be used.
    # All waveforms will be padded to the right with zeros to the length of the
    # longest generated waveform.
    # Then the injection times are overriden(!) in the config file coming with
    # the output h5 file with the fraction where the maximum of the waveform
    # lies within the whole sample.
    duration = -1.
    sample_rate = 4096
    # Default parameters passed to simulator in case a parameter is not sampled

    # Possible approximants: ['TaylorEt', 'SEOBNRv3_opt', 'IMRPhenomA',
    # 'IMRPhenomC', 'IMRPhenomB', 'EOBNRv2', 'SEOBNRv4_opt', 'PhenSpinTaylor',
    # 'PhenSpinTaylorRD', 'NR_hdf5', 'SEOBNRv3_pert', 'EOBNRv2HM',
    # 'SpinTaylorT4', 'TaylorT1', 'TaylorT3', 'TaylorT2', 'HGimri', 'TaylorT4',
    # 'IMRPhenomD', 'IMRPhenomPv2', 'SEOBNRv1', 'SpinDominatedWf', 'SEOBNRv3',
    # 'SEOBNRv2', 'SpinTaylorT1', 'SEOBNRv4', 'SpinTaylorT2', 'EccentricTD',
    # 'SEOBNRv2_opt', 'SEOBNRv3_opt_rk4']
    default_parameters = {
        'mass1': 1.0,
        'mass2': 1.0,
        'spin1z': 0.0,
        'spin2z': 0.0,
        'lambda1': 0.0,
        'lambda2': 0.0,
        'distance': 1000.0,
        'coa_phase': 0.0,
        'inclination': 0.0,
        'delta_t': 1.0 / sample_rate,
        'f_lower': 15.0,
        'approximant': 'SEOBNRv4',
        'injection_time': 0.95
        # id does not have a default because it must be unique
    }

    # Specify the parameters that should be varied and the ranges (uniform)
    # For the chosen range consult
    # https://www.lsc-group.phys.uwm.edu/ligovirgo/cbcnote/O2/OfflineTuningVerificationInjections
    update_list = {
        'masses': [2., 50.],
        # 'injection_time': [0.5, 0.9],
        'distance': [400., 800.]
    }

    # Collect everything in a big meta config dictionary
    meta_config = {
        'n_samples': n_samples,
        'duration': duration,
        'sample_rate': sample_rate,
        'default_parameters': default_parameters,
        'update_list': update_list,
    }

    with open(output_file, 'w') as f:
        json.dump(meta_config, f, sort_keys=True, indent=2)
    print("DONE -- Saved as {}".format(output_file))

if __name__ == '__main__':
    import datetime
    timestamp = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
    new_dir = os.path.join(os.getcwd(), timestamp)
    print("Create new folder ", new_dir)
    os.makedirs(new_dir)

    meta_config_file = os.path.join(new_dir, 'config.json')
    sample_config_file = os.path.join(new_dir, 'sample_config.json')
    sample_file = os.path.join(new_dir, 'samples.h5')

    generate_meta_config(meta_config_file)

    data_generation_pipeline(meta_config_file,
                             sample_config_file,
                             sample_file)


import numpy as np
import math

##################################
#
#    Global variables
#
#################################

PERIODIC_TABLE = {
    1:  "H",
    2:  "He",
    3:  "Li",
    4:  "Be",
    5:  "B",
    6:  "C",
    7:  "N",
    8:  "O",
    9:  "F",
    10:	"Ne",
    11:	"Na",
    12:	"Mg",
    13:	"Al",
    14:	"Si",
    15: "P",
    16: "S",
    17:	"Cl",
    18:	"Ar",
    19: "K",
    20:	"Ca",
    21:	"Sc",
    22:	"Ti",
    23: "V",
    24:	"Cr",
    25:	"Mn",
    26:	"Fe",
    27:	"Co",
    28:	"Ni",
    29:	"Cu",
    30:	"Zn",
    31:	"Ga",
    32:	"Ge",
    33:	"As",
    34:	"Se",
    35:	"Br",
    36:	"Kr",
    37:	"Rb",
    38:	"Sr",
    39: "Y",
    40:	"Zr",
    41:	"Nb",
    42:	"Mo",
    43:	"Tc",
    44:	"Ru",
    45:	"Rh",
    46:	"Pd",
    47:	"Ag",
    48:	"Cd",
    49:	"In",
    50:	"Sn",
    51:	"Sb",
    52:	"Te",
    53: "I",
    54:	"Xe",
    55:	"Cs",
    56:	"Ba",
    57:	"La",
    58:	"Ce",
    59:	"Pr",
    60:	"Nd",
    61:	"Pm",
    62:	"Sm",
    63:	"Eu",
    64:	"Gd",
    65:	"Tb",
    66:	"Dy",
    67:	"Ho",
    68:	"Er",
    69:	"Tm",
    70:	"Yb",
    71:	"Lu",
    72:	"Hf",
    73:	"Ta",
    74: "W",
    75:	"Re",
    76:	"Os",
    77:	"Ir",
    78:	"Pt",
    79:	"Au",
    80:	"Hg",
    81:	"Tl",
    82:	"Pb",
    83:	"Bi",
    84:	"Po",
    85:	"At",
    86:	"Rn",
    87:	"Fr",
    88:	"Ra",
    89:	"Ac",
    90:	"Th",
    91:	"Pa",
    92: "U",
    93:	"Np",
    94:	"Pu",
    95:	"Am",
    96:	"Cm",
    97:	"Bk",
    98:	"Cf",
    99:	"Es",
    100: "Fm",
    101: "Md",
    102: "No",
    103: "Lr",
    104: "Rf",
    105: "Db",
    106: "Sg",
    107: "Bh",
    108: "Hs",
    109: "Mt",
    110: "Ds",
    111: "Rg",
    112: "Cn",
    113: "Nh",
    114: "Fl",
    115: "Mc",
    116: "Lv",
    117: "Ts",
    118: "Og"
}

####################################
#
#        Classes
#
####################################

class atom(object):

    def __init__(self, num, cartesian=None, unit="angstrom"):

        if cartesian is None:
            cartesian = [0, 0, 0]
        self.__atom_num = num
        self.__label = PERIODIC_TABLE[num]
        self.__cartesian = cartesian
        self.__unit = unit

    def get_atom_num(self):
        return self.__atom_num

    def get_label(self):
        return self.__label

    def get_cartesian(self):
        return self.__cartesian

    def get_unit(self):
        return self.__unit

    def set_atom_num(self, num):
        self.__atom_num = num
        self.__label = PERIODIC_TABLE[num]

    def set_label(self, label):
        self.__label = label
        if dict_inv_search(label, PERIODIC_TABLE):
            self.__atom_num = dict_inv_enquiry(label, PERIODIC_TABLE)
        else:
            raise Exception("Element {} is not found in periodic table.".format(label))

    def set_cartesian(self, cartesian):
        self.__cartesian = cartesian

    def unit_convert_to_bohr(self):
        if self.__unit == 'angstrom':
            self.__cartesian = [i / 0.529177210903 for i in self.__cartesian]

        if self.__unit == 'SI':
            self.__cartesian = [i / 5.29177210903e-11 for i in self.__cartesian]

        self.__unit = 'bohr'

    def unit_convert_to_angstrom(self):
        if self.__unit == 'bohr':
            self.__cartesian = [i * 0.529177210903 for i in self.__cartesian]

        if self.__unit == 'SI':
            self.__cartesian = [i * 1e10 for i in self.__cartesian]

        self.__unit = 'angstrom'

#########################################
#
#        Functions
#
########################################

def is_member(value, iterable):
    for i in iterable:
        if value == i:
            return True

    return False

def dict_inv_enquiry(value, dictionary):
    for i in dictionary:
        if dictionary[i] == value:
            return i

    return value

def dict_inv_search(value, dictionary):
    for i in dictionary:
        if dictionary[i] == value:
            return True

    return False

def rot(a_vec, b_vec):
    x = np.linalg.det([[a_vec[1], a_vec[2]], [b_vec[1], b_vec[2]]])
    y = - np.linalg.det([[a_vec[0], a_vec[2]], [b_vec[0], b_vec[2]]])
    z = np.linalg.det([[a_vec[0], a_vec[1]], [b_vec[0], b_vec[1]]])

    return [x, y, z]

def norm(vec):
    return np.dot(np.array(vec), np.array(vec))

def ortho_mat(vec):
    z = np.array(vec) / math.sqrt(norm(vec))
    if z[0] == 0:
        x = np.array([1, 0, 0])
    else:
        x = np.array([- (z[1] + z[2]) / z[0], 1, 1])

    x = x / math.sqrt(norm(x))

    y = np.array(rot(z, x))

    y = y / math.sqrt(norm(y))

    return np.array([x, y, z])

def rotation(coordinate, ref_point, ref_vec, deg):

    reference_point = np.array(ref_point)
    translated_cartesian = np.array(coordinate) - reference_point

    tr_basis_mat = ortho_mat(ref_vec).transpose()

    coefficient = np.linalg.solve(tr_basis_mat, translated_cartesian)

    coefficient = np.array([coefficient[0] * math.cos(-deg) - coefficient[1] * math.sin(-deg),
                            coefficient[0] * math.sin(-deg) + coefficient[1] * math.cos(-deg),
                            coefficient[2]])

    return list(tr_basis_mat.dot(coefficient) + reference_point)

def remove_axis_vec(coordinate, ref_point, ref_vec):

    reference_point = np.array(ref_point)
    translated_cartesian = np.array(coordinate) - reference_point

    tr_basis_mat = ortho_mat(ref_vec).transpose()

    coefficient = np.linalg.solve(tr_basis_mat, translated_cartesian)

    coefficient = np.array([coefficient[0], coefficient[1], 0])

    return list(tr_basis_mat.dot(coefficient))

def bond_length(a_atom, b_atom):
    ab = np.array(a_atom.get_cartesian()) - np.array(b_atom.get_cartesian())

    return math.sqrt(np.dot(ab, ab))

def bond_angle(a_atom, center_atom, b_atom):
    ac = np.array(a_atom.get_cartesian()) - np.array(center_atom.get_cartesian())
    bc = np.array(b_atom.get_cartesian()) - np.array(center_atom.get_cartesian())

    return math.acos(np.dot(ac, bc) / math.sqrt(norm(ac)) / math.sqrt(norm(bc)))

def bond_dihedral_angle(a_atom, b_atom, c_atom, d_atom):
    axis = np.array(b_atom.get_cartesian()) - np.array(c_atom.get_cartesian())

    axis_removed_ab = remove_axis_vec(a_atom.get_cartesian(), b_atom.get_cartesian(), axis)
    axis_removed_dc = remove_axis_vec(d_atom.get_cartesian(), c_atom.get_cartesian(), axis)

    return math.acos(np.dot(axis_removed_ab, axis_removed_dc)
                     / math.sqrt(norm(axis_removed_ab))
                     / math.sqrt(norm(axis_removed_dc)))

# all atoms in `group' should be bound to a_atom
def set_bond_length(a_atom, b_atom, length, group=None):
    vec = np.array(b_atom.get_cartesian()) - np.array(a_atom.get_cartesian())
    change = (math.sqrt(norm(vec)) - length) / math.sqrt(norm(vec)) * vec

    if group is not None:
        for i in group:
            i.set_cartesian(list(np.array(i.get_cartesian()) + change))
    else:
        a_atom.set_cartesian(list(np.array(a_atom.get_cartesian()) + change))

# all atoms in `group' should be bound to a_atom
def set_bond_angle(a_atom, center_atom, b_atom, angle, group=None):
    original_angle = bond_angle(a_atom, center_atom, b_atom)
    ac = np.array(a_atom.get_cartesian()) - np.array(center_atom.get_cartesian())
    bc = np.array(b_atom.get_cartesian()) - np.array(center_atom.get_cartesian())
    axis = np.array(rot(ac, bc))

    change = angle - original_angle

    if group is not None:
        for i in group:
            i.set_cartesian(rotation(i.get_cartesian(), center_atom.get_cartesian(), axis, change))

    else:
        a_atom.set_cartesian(rotation(a_atom.get_cartesian(), center_atom.get_cartesian(), axis, change))

# all atoms in `group' should be bound to a_atom
def set_bond_dihedral_angle(a_atom, b_atom, c_atom, d_atom, angle, group=None):
    axis = np.array(b_atom.get_cartesian()) - np.array(c_atom.get_cartesian())
    original_angle = bond_dihedral_angle(a_atom, b_atom, c_atom, d_atom)

    change = angle - original_angle

    if group is not None:
        for i in group:
            i.set_cartesian(rotation(i.get_cartesian(), b_atom, axis, change))

    else:
        a_atom.set_cartesian(rotation(a_atom.get_cartesian(), b_atom, axis, change))

def atoms_to_string_template(atoms_list):
    atoms_string = ""

    for i in atoms_list:
        atoms_string += ' ' + i.get_label().ljust(28)
        for j in i.get_cartesian():
            atoms_string += str(float(j)).ljust(28)

        atoms_string += '\n'

    return atoms_string