import logging

import numpy as np
from nengo.dists import Choice
from nengo.exceptions import BuildError
from nengo.neurons import (
    LIF,
    LIFRate,
    RectifiedLinear,
    RegularSpiking,
    SpikingRectifiedLinear,
)

from nengo_loihi.compat import HAS_DL, HAS_TF, nengo_dl, tf

try:
    from nengo_extras.loihi_training import (
        LoihiLIFBuilder,
        LoihiSpikingRectifiedLinearBuilder,
    )

except ImportError:  # pragma: no cover

    class ErrorBuilder:
        def __init__(self, ops):
            raise BuildError(
                "Building Loihi neuron types in nengo-dl requires nengo-extras>=0.5. "
                "Please install or upgrade nengo-extras."
            )

    LoihiLIFBuilder = LoihiSpikingRectifiedLinearBuilder = ErrorBuilder

logger = logging.getLogger(__name__)

class Installer:
    def __init__(self):
        self.installed = False

    def __call__(self):
        if self.installed or not HAS_DL:
            pass
        else:
            logger.info("Installing NengoDL neuron builders")
            nengo_dl.neuron_builders.SimNeuronsBuilder.TF_NEURON_IMPL[
                LoihiLIF
            ] = LoihiLIFBuilder
            nengo_dl.neuron_builders.SimNeuronsBuilder.TF_NEURON_IMPL[
                LoihiSpikingRectifiedLinear
            ] = LoihiSpikingRectifiedLinearBuilder
            self.installed = True

install_dl_builders = Installer()

def discretize_tau_rc(tau_rc, dt):
    """Discretize tau_rc as per discretize_compartment.

    Parameters
    ----------
    tau_rc : float
        The neuron membrane time constant.
    dt : float
        The simulator time step.
    """
    lib = tf.math if HAS_TF and isinstance(tau_rc, tf.Tensor) else np

    decay_rc = -(lib.expm1(-dt / tau_rc))
    decay_rc = lib.round(decay_rc * (2**12 - 1)) / (2**12 - 1)
    return -dt / lib.log1p(-decay_rc)

def discretize_tau_ref(tau_ref, dt):
    """Discretize tau_ref as per Compartment.configure_lif.

    Parameters
    ----------
    tau_rc : float
        The neuron membrane time constant.
    dt : float
        The simulator time step.
    """
    lib = tf.math if HAS_TF and isinstance(tau_ref, tf.Tensor) else np

    return dt * lib.round(tau_ref / dt)

def loihi_lif_rates(neuron_type, x, gain, bias, dt, amplitude=None):
    tau_ref = discretize_tau_ref(neuron_type.tau_ref, dt)
    tau_rc = discretize_tau_rc(neuron_type.tau_rc, dt)
    amplitude = neuron_type.amplitude if amplitude is None else amplitude

    j = neuron_type.current(x, gain, bias) - 1
    out = np.zeros_like(j)
    period = tau_ref + tau_rc * np.log1p(1.0 / j[j > 0])
    out[j > 0] = (amplitude / dt) / np.ceil(period / dt)
    return out

def loihi_spikingrectifiedlinear_rates(neuron_type, x, gain, bias, dt, amplitude=None):
    amplitude = neuron_type.amplitude if amplitude is None else amplitude

    j = neuron_type.current(x, gain, bias)
    out = np.zeros_like(j)
    period = 1.0 / j[j > 0]
    out[j > 0] = (amplitude / dt) / np.ceil(period / dt)
    return out

def loihi_regularspiking_rates(neuron_type, x, gain, bias, dt):
    base_type = neuron_type.base_type
    if type(base_type) is LIFRate:
        return loihi_lif_rates(
            base_type, x, gain, bias, dt, amplitude=neuron_type.amplitude
        )
    elif type(base_type) is RectifiedLinear:
        return loihi_spikingrectifiedlinear_rates(
            base_type, x, gain, bias, dt, amplitude=neuron_type.amplitude
        )
    else:
        return neuron_type.rates(x, gain, bias)

def _broadcast_rates_inputs(x, gain, bias):
    x = np.array(x, dtype=float, copy=False, ndmin=1)
    gain = np.array(gain, dtype=float, copy=False, ndmin=1)
    bias = np.array(bias, dtype=float, copy=False, ndmin=1)
    if x.ndim == 1:
        x = x[:, np.newaxis] * np.ones(gain.shape[-1])
    return x, gain, bias

def loihi_rates(neuron_type, x, gain, bias, dt):
    x, gain, bias = _broadcast_rates_inputs(x, gain, bias)
    for cls in type(neuron_type).__mro__:
        if cls in loihi_rate_functions:
            return loihi_rate_functions[cls](neuron_type, x, gain, bias, dt)
    return neuron_type.rates(x, gain, bias)

def nengo_rates(neuron_type, x, gain, bias):
    """Call NeuronType.rates with Nengo 3.0 broadcasting rules"""
    x, gain, bias = _broadcast_rates_inputs(x, gain, bias)
    return neuron_type.rates(x, gain, bias)

loihi_rate_functions = {
    LIF: loihi_lif_rates,
    SpikingRectifiedLinear: loihi_spikingrectifiedlinear_rates,
    RegularSpiking: loihi_regularspiking_rates,
}

class LoihiLIF(LIF):
    """Simulate LIF neurons as done by Loihi.

    On Loihi, the inter-spike interval has to be an integer. This causes
    aliasing the firing rates where a wide variety of inputs can produce the
    same output firing rate. This class reproduces this effect, as well as
    the discretization of some of the neuron parameters. It can be used in
    e.g. ``nengo`` or ``nengo_dl`` to reproduce these unique Loihi effects.

    Parameters
    ----------
    nengo_dl_noise : `nengo_extras.loihi_training.NeuronOutputNoise`
        Noise added to the rate-neuron output when training with this neuron
        type in ``nengo_dl``.
    """

    state = {
        "voltage": Choice([0]),
        "refractory_time": Choice([0]),
    }

    def __init__(
        self,
        tau_rc=0.02,
        tau_ref=0.002,
        min_voltage=0,
        amplitude=1,
        nengo_dl_noise=None,
        **kwargs,
    ):
        super().__init__(
            tau_rc=tau_rc,
            tau_ref=tau_ref,
            min_voltage=min_voltage,
            amplitude=amplitude,
            **kwargs,
        )
        self.nengo_dl_noise = nengo_dl_noise
        install_dl_builders()

    @property
    def _argreprs(self):
        args = super()._argreprs
        if self.nengo_dl_noise is not None:
            args.append("nengo_dl_noise=%s" % self.nengo_dl_noise)
        return args

    def rates(self, x, gain, bias, dt=0.001):
        return loihi_lif_rates(self, x, gain, bias, dt)

    def step(self, dt, J, output, voltage, refractory_time):
        tau_ref = discretize_tau_ref(self.tau_ref, dt)
        tau_rc = discretize_tau_rc(self.tau_rc, dt)

        refractory_time -= dt
        delta_t = (dt - refractory_time).clip(0, dt)
        voltage -= (J - voltage) * np.expm1(-delta_t / tau_rc)

        spikes_mask = voltage > 1
        output[:] = spikes_mask * (self.amplitude / dt)

        voltage[voltage < self.min_voltage] = self.min_voltage
        voltage[spikes_mask] = 0
        refractory_time[spikes_mask] = tau_ref + dt

class LoihiSpikingRectifiedLinear(SpikingRectifiedLinear):
    """Simulate spiking rectified linear neurons as done by Loihi.

    On Loihi, the inter-spike interval has to be an integer. This causes
    aliasing in the firing rates such that a wide variety of inputs produce the
    same output firing rate. This class reproduces this effect. It can be used
    in e.g. ``nengo`` or ``nengo_dl`` to reproduce these unique Loihi effects.
    """

    state = {
        "voltage": Choice([0]),
    }

    def __init__(self, amplitude=1, **kwargs):
        super().__init__(amplitude=amplitude, **kwargs)
        install_dl_builders()

    def rates(self, x, gain, bias, dt=0.001):
        return loihi_spikingrectifiedlinear_rates(self, x, gain, bias, dt)

    def step(self, dt, J, output, voltage):
        voltage += J * dt

        spikes_mask = voltage > 1
        output[:] = spikes_mask * (self.amplitude / dt)

        voltage[voltage < 0] = 0
        voltage[spikes_mask] = 0


#! /usr/bin/python

""" Walking gait task.
"""

### IMPORTS ###
import random

# Libraries
import pymunk
import numpy as np

# Local
from ..networks.rnn import NeuralNetwork

# Shortcuts
pi = np.pi

### FUNCTIONS ###

def angle_fix(theta):
    """ Fixes an angle to a value between -pi and pi.

        >>> angle_fix(-2*pi)
        0.0
    """
    return ((theta + pi) % (2*pi)) - pi

### CLASSES ###

class Joint(object):
    """ Joint object, contains pivot, motor and limit"""

    def __init__(self, a, b, position, range=(-pi, pi), max_rate=2.0):
        self.pivot = pymunk.PivotJoint(a, b, position)
        self.motor = pymunk.SimpleMotor(a, b, 0)
        self.motor.max_force = 1e7
        self.limit = pymunk.RotaryLimitJoint(a, b, range[0], range[1])
        self.limit.max_force = 1e1
        self.limit.max_bias = 0.5
        self.max_rate = max_rate

    def angle(self):
        return angle_fix(self.pivot.a.angle - self.pivot.b.angle)

    def set_target(self, target):
        """ Target is between 0 and 1, representing min and max angle. """
        cur = self.angle()
        tgt = angle_fix(target * (self.limit.max - self.limit.min) + self.limit.min)
        if tgt > cur + 0.1:
            self.motor.rate = self.max_rate
        elif tgt < cur - 0.1:
            self.motor.rate = -self.max_rate
        else:
            self.motor.rate = 0

class Leg(object):
    """ Leg object, contains joints and shapes """

    def __init__(self, parent, position, walking_task):
        self.walking_task = walking_task
        (w, l) = walking_task.leg_length / 5.0, walking_task.leg_length
        mass = w * l * 0.2
        # Upper leg
        upperleg = pymunk.Body(mass, pymunk.moment_for_box(mass, w, l))
        upperleg.position = pymunk.Vec2d(parent.position) + pymunk.Vec2d(position) + pymunk.Vec2d(0, l/2.0 - w/2.0)
        shape = pymunk.Poly.create_box(upperleg, (w,l))
        shape.group = 1
        shape.friction = 2.0
        # Joints
        pos = pymunk.Vec2d(parent.position) + pymunk.Vec2d(position)
        hip = Joint(parent, upperleg, pos, (-0.1*pi, 0.9*pi), self.walking_task.max_rate)
        walking_task.space.add(hip.pivot, hip.motor, hip.limit, upperleg, shape)

        # Lower leg
        lowerleg = pymunk.Body(mass, pymunk.moment_for_box(mass, w, l * 1.2))
        lowerleg.position = pymunk.Vec2d(upperleg.position) + pymunk.Vec2d(0, l - w/2.0)
        shape = pymunk.Poly.create_box(lowerleg, (w, l * 1.2))
        shape.group = 1
        shape.friction = 2.0
        # Joints
        pos =  pymunk.Vec2d(upperleg.position) + pymunk.Vec2d(0, l/2.0 - w/2.0)
        knee = Joint(upperleg, lowerleg, pos, (-0.9*pi, 0.1*pi), self.walking_task.max_rate)
        walking_task.space.add(knee.pivot, knee.motor, knee.limit, lowerleg, shape)

        self.upperleg = upperleg
        self.lowerleg = lowerleg
        self.hip = hip
        self.knee = knee

class WalkingTask(object):
    """ Walking gait task.
    """

    def __init__(self, max_steps=1000,
                       track_length=1000,
                       max_rate=2.0,
                       torso_height=40,
                       torso_density=0.2,
                       leg_spacing=30,
                       leg_length=30,
                       num_legs=4):
        # Settings
        self.max_steps = max_steps
        self.track_length = track_length
        self.max_rate = max_rate
        self.torso_height = torso_height
        self.torso_density = torso_density
        self.leg_spacing = leg_spacing
        self.leg_length = leg_length
        self.num_legs = num_legs

    def evaluate(self, network, draw=False):
        """ Evaluate the efficiency of the given network. Returns the
            distance that the walker ran in the given time (max_steps).
        """
        if not isinstance(network, NeuralNetwork):
            network = NeuralNetwork(network)

        if draw:
            import pygame
            pygame.init()
            screen = pygame.display.set_mode((self.track_length, 200))
            pygame.display.set_caption("Simulation")
            clock = pygame.time.Clock()
            running = True
            font = pygame.font.Font(pygame.font.get_default_font(), 8)

        # Initialize pymunk
        self.space = space = pymunk.Space()
        space.gravity = (0.0, 900.0)
        space.damping = 0.7
        self.touching_floor = False
        # Create objects
        # Floor

        floor = pymunk.Body()
        floor.position = pymunk.Vec2d(self.track_length/2.0 , 210)
        sfloor = pymunk.Poly.create_box(floor, (self.track_length, 40))
        sfloor.friction = 1.0
        sfloor.collision_type = 1
        space.add_static(sfloor)

        # Torso
        torsolength = 20 + (self.num_legs // 2 - 1) * self.leg_spacing
        mass = torsolength * self.torso_height * self.torso_density
        torso = pymunk.Body(mass, pymunk.moment_for_box(mass, torsolength, self.torso_height))
        torso.position = pymunk.Vec2d(200, 200 - self.leg_length * 2 - self.torso_height)
        storso = pymunk.Poly.create_box(torso, (torsolength, self.torso_height))
        storso.group = 1
        storso.collision_type = 1
        storso.friction = 2.0
        # space.add_static(storso)
        space.add(torso, storso)

        # Legs
        legs = []
        for i in range(self.num_legs // 2):
            x = 10 - torsolength / 2.0 + i * self.leg_spacing
            y = self.torso_height / 2.0 - 10
            legs.append( Leg(torso, (x,y), self) )
            legs.append( Leg(torso, (x,y), self) )

        # Collision callback
        def oncollide(space, arb):
            self.touching_floor = True
        space.add_collision_handler(1, 1, post_solve=oncollide)

        for step in xrange(self.max_steps):

            # Query network
            input_width = max(len(legs), 4)
            net_input = np.zeros((3, input_width))
            torso_y = torso.position.y
            torso_a = torso.angle
            sine = np.sin(step / 10.0)
            hip_angles = [leg.hip.angle() for leg in legs]
            knee_angles = [leg.knee.angle() for leg in legs]
            other = [torso_y, torso_a, sine, 1.0]
            # Build a 2d input grid,
            # as in Clune 2009 Evolving Quadruped Gaits, p4
            net_input[0, :len(legs)] = hip_angles
            net_input[1, :len(legs)] = knee_angles
            net_input[2, :4] = other
            act = network.feed(net_input, add_bias=False)

            output = np.clip(act[-self.num_legs*2:] * self.max_rate, -1.0, 1.0) / 2.0 + 0.5

            for i, leg in enumerate(legs):
                leg.hip.set_target( output[i * 2] )
                leg.knee.set_target( output[i * 2 + 1] )

            # Advance simulation
            space.step(1/50.0)
            # Check for success/failure
            if torso.position.x < 0:
                break
            if torso.position.x > self.track_length - 50:
                break
            if self.touching_floor:
                break

            # Draw
            if draw:
                print act
                # Clear
                screen.fill((255, 255, 255))
                # Do all drawing
                txt = font.render('%d' % step, False, (0,0,0) )
                screen.blit(txt, (0,0))
                # Draw objects
                for o in space.shapes + space.static_shapes:
                    if isinstance(o, pymunk.Circle):
                        pygame.draw.circle(screen, (0,0,0), (int(o.body.position.x), int(o.body.position.y)), int(o.radius))
                    else:
                        pygame.draw.lines(screen, (0,0,0), True, [(int(p.x), int(p.y)) for p in o.get_points()])
                # Flip buffers
                pygame.display.flip()
                clock.tick(50)

        if draw:
            pygame.quit()

        distance = torso.position.x
        # print "Travelled %.2f in %d steps." % (distance, step)
        return {'fitness':distance}

    def solve(self, network):
        return False

    def visualize(self, network, filename=None):
        """ Visualize a solution strategy by the given individual. """
        self.evaluate(network, draw=True)


from typing import Tuple, List, Dict, Union
from threading import Thread
from multiprocessing import cpu_count

from numba import njit
from numpy import int16 as nshort, inf, ndarray, float32 as nfloat
import numpy as np
from pandas import DataFrame, Index, MultiIndex

_NULL_INDEX = -1
_NEG_INF = -inf

def _get_breaks(n_items: int, n_workers: int) -> List[Tuple[int, int]]:
    """Gets starts and stops for breaks similar to numpy.array_slice()"""
    div, remainder = divmod(n_items, n_workers)
    slices = []
    bottom = 0
    for i in range(n_workers):
        top = bottom + div
        top += int(i < remainder)
        slices.append((bottom, top))
        bottom = top
    return slices

@njit
def _update_heap(utilities: ndarray, zones: ndarray, new_u: float, new_zone: int):
    """Inserts a value into a sorted array, maintaining sort order, and the number of items in that array. The array
    is sorted lowest-to-highest"""
    i = 0
    top = len(utilities)
    while i < top:
        current_u = utilities[i]
        if new_u < current_u:
            break
        i += 1
    if i <= 0:
        return
    for j in range(i - 1):
        utilities[j] = utilities[j + 1]
        zones[j] = zones[j + 1]
    utilities[i - 1] = new_u
    zones[i - 1] = new_zone

@njit(nogil=True)
def _nbf_twopart_worker(access_utils: ndarray, egress_utils: ndarray, result_utils: ndarray, result_stations: ndarray,
                        start: int, stop: int, k: int):
    """Performance-tuned NBF to operate on its own thread"""
    n_origins, n_intermediate = access_utils.shape
    n_destinations = egress_utils.shape[1]

    # Allocate the sorted heap of utilities (and associated indices) once per thread
    util_heap = np.full(shape=k, fill_value=-inf)
    zones_heap = np.full(shape=k, fill_value=_NULL_INDEX)

    for offset in range(start, stop):
        origin_zone, destination_zone = divmod(offset, n_destinations)

        # Reset the heap for this OD
        util_heap[:] = _NEG_INF
        zones_heap[:] = _NULL_INDEX

        for interim_zone in range(n_intermediate):

            if interim_zone >= n_intermediate:
                print("ERR", offset, origin_zone, destination_zone, interim_zone)
                raise AssertionError()

            interim_util = access_utils[origin_zone, interim_zone] + egress_utils[interim_zone, destination_zone]

            # In general, for problems where (n_origins * n_destinations) >> k, most values will not be in the top k.
            # So quickly check against the lowest utility in the heap to avoid calling the updater func
            if interim_util < util_heap[0] or interim_util == _NEG_INF:
                continue
            _update_heap(util_heap, zones_heap, interim_util, interim_zone)

        result_utils[origin_zone, destination_zone, :] = util_heap
        result_stations[origin_zone, destination_zone, :] = zones_heap

def _validate_access_egress_tables(access_table: DataFrame, egress_table: DataFrame) -> Tuple[Index, Index, Index]:

    assert access_table.index.nlevels == 2, "Access table index must have two levels"
    assert egress_table.index.nlevels == 2, "Egress table index must have two levels"

    # Take the unique index of each level, as Pandas can return more items than is present, if the frame is a slice
    origin_zones: Index = access_table.index.unique(level=0)
    intermediate_zones: Index = access_table.index.unique(level=1)
    destination_zones: Index = egress_table.index.unique(level=1)

    # Check that the access and egress tables have compatible indices
    assert intermediate_zones.equals(egress_table.index.unique(level=0)), \
        "Access index level 2 and egress index level 1 must be the same"

    return origin_zones, intermediate_zones, destination_zones

def best_intermediate_zones(access_table: DataFrame, egress_table: DataFrame, cost_column: str, k: int = 1,
                            n_threads: int = None, squeeze=True, other_columns=True,
                            intermediate_name: str = "intermediate_zone", maximize=True,
                            availability_column: str = "available", null_index=0
                            ) -> Union[DataFrame, Dict[int, DataFrame]]:
    """Numba-accelerated.

    Triple-index operation for two matrices, finding the most- or least-cost intermediate zones. Takes an access matrix
    of the shape (O, I) and an egress matrix of the shape (I, D) to produce a combined matrix of the shape (O, D), with
    the best intermediate I. Also works to construct multiple (O, D) matrices - for the top _k_ intermediate zones in
    _I_.

    There is no restriction on the label dtypes, as long as the access and egress tables share the same _I_ index.

    Both the input matrices must be provided in "tall" format - as Pandas Series with a 2-level MultiIndex.
    Essentially, the access and egress tables are DataFrames with multiple matrices defined within. The output table(s)
    are also returned in a tall format.

    When constructing the result tables, columns in the access and egress tables are "carried forward" such that the
    results columns will be the union of columns in the input tables. Columns in one table only will be carried forward
    unmodified and retain their data type. Columns in both tables will be added together, and thus MUST be numeric.

    In the specified cost column, a value of `-inf` (or `inf` when minimizing) is respected as the sentinel value for
    unavailable. (O, I) or (I, D) interchanges with this sentinel value will not be considered.

    Args:
        access_table: DataFrame with 2-level MultiIndex of the shape ((O, I), A). Must include the specified cost column
        egress_table: DataFrame with 2-level MultiIndex of the shape ((I, D), E). Must include the specified cost column
        cost_column: Name of the column in the access and egress table to use as the cost to minimize/maximize. Values
            of `+/- inf` are respected to indicate unavailable choices
        k: The number of ranks to return (e.g., find the _k_ best intermediate zones). If k <= 0, it will be corrected
            to 1.
        n_threads: Number of threads to use. Defaults to cpu_count()
        squeeze: If k == 1 and squeeze=True, a single DataFrame is returned. Otherwise, a Dictionary of DataFrames will
            be returned.
        other_columns: If True, the result DataFrame will include all columns in the access and egress tables. The
            result table will be of the shape ((O, D), A | E + 3)
        intermediate_name: Name of the column in the result table containing the selected intermediate zone.
        maximize: If True, this function maximize the result. If False, it minimizes it.
        availability_column: Name of the column in the result table containing a flag whether ANY intermediate zone
            was found to be available.
        null_index: Fill value used if NO intermediate zone is available.

    Returns:
        DataFrame: If k == 1 and squeeze=True. A DataFrame of the shape ((O, D), A | E + 3), containing the intermediate
            zone selected, the associated max/min cost, and a flag indicating its availability. Additional columns from
            the access and egress tables, indexed for the appropriately chosen intermediate zone, will also be included
            if other_columns=True.
        Dict[int, DataFrame]: If k > 1. The keys represent the ranks, so result[1] is the best intermediate zone,
            result[2] is the second-best, etc. The value DataFrames are in the same format as if k == 1, just with
            different intermediate zones chosen.
    """

    # Check inputs
    k = max(1, k)
    if n_threads is None:
        n_threads = cpu_count()
    origins, intermediates, destinations = _validate_access_egress_tables(access_table, egress_table)
    n_origins, n_intermediate, n_destinations = len(origins), len(intermediates), len(destinations)

    # Compute best path(s)
    access_cost = access_table[cost_column].values.reshape([n_origins, n_intermediate]).astype(nfloat)
    egress_cost = egress_table[cost_column].values.reshape([n_intermediate, n_destinations]).astype(nfloat)

    if not maximize:  # Invert the cost to use code set to maximize
        access_cost *= -1
        egress_cost *= -1

    result_cost = np.zeros(shape=(n_origins, n_destinations, k), dtype=nfloat)
    result_indices = np.zeros(dtype=nshort, shape=(n_origins, n_destinations, k))

    # Setup the workers (1 per thread) to select best paths for a subset of ODs
    breaks = _get_breaks(n_origins * n_destinations, n_threads)
    threads = [
        Thread(target=_nbf_twopart_worker, args=[
            access_cost, egress_cost, result_cost, result_indices,
            start, stop, k
        ])
        for start, stop in breaks
    ]
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    # Construct composite result tables
    if other_columns:
        remaining_columns = set(access_table.columns | egress_table.columns) - {cost_column, intermediate_name,
                                                                                availability_column}
    else:
        remaining_columns = set()

    row_index = MultiIndex.from_product([origins, destinations])        # Labels for the rows
    access_indexer = np.repeat(np.arange(n_origins), n_destinations)    # Indexer for the access table
    egress_indexer = np.tile(np.arange(n_destinations), n_origins)      # Indexer for the egress table
    tables = {}                                                         # Results

    for i in range(k):
        table = DataFrame(index=row_index)

        offsets_i = result_indices[:, :, i]  # 2D indexer for this (i ∈ k) path
        flat_offsets_i = offsets_i.flatten()  # Convert to 1D indexer
        availability_i = flat_offsets_i != _NULL_INDEX

        intermediate_result_i = intermediates.take(flat_offsets_i)
        intermediate_result_i[~availability_i] = null_index
        table[intermediate_name] = intermediate_result_i

        table[availability_column] = availability_i
        table[cost_column] = result_cost[offsets_i]

        # If there are any columns left, add them to the composite table
        for column in remaining_columns:
            in_access = column in access_table
            in_egress = column in egress_table

            if in_access:
                access_matrix = access_table[column].values.reshape([n_origins, n_intermediate])
                access_component = access_matrix[access_indexer, flat_offsets_i]
            if in_egress:
                egress_matrix = egress_table[column].values.reshape([n_intermediate, n_destinations])
                egress_component = egress_matrix[flat_offsets_i, egress_indexer]

            if in_access and in_egress:
                composite_data = access_component + egress_component
            elif in_access:
                composite_data = access_component
            elif in_egress:
                composite_data = egress_component
            else:
                raise RuntimeError("This shouldn't happen")

            table[column] = composite_data

        tables[k - i] = table

    if k == 1 and squeeze:
        return tables[1]
    return tables


#!/usr/bin/env python3
import numpy as np

def rank_properties(prop_name, obs):
    """Select ranking method based on property name.

    Parameters
    ----------
    prop_name : str
        The property to rank.
    obs : dict
        The observations including metadata

    Returns
    -------
    list of bool
        Entry "preferred" in propertyCollection, True if preferred, else False
    """
    if prop_name in ["taxonomy", "taxonomies"]:
        return select_taxonomy(obs)
    elif prop_name in ["mass", "masses"]:
        return select_numeric_property(obs, "mass")
    elif prop_name in ["albedo", "albedos"]:
        return select_numeric_property(obs, "albedo")
    elif prop_name in ["diameter", "diameters"]:
        return select_numeric_property(obs, "diameter")
    else:
        return [True for i in obs]

def select_taxonomy(taxonomies):
    """Select a single taxonomic classification from multiple choices.

    Evaluates the wavelength ranges, methods, schemes, and recency of
    classification.

    Parameters
    ----------
    taxonomies: dict
        Taxonomic classifications retrieved from SsODNet:datacloud.

    Returns
    -------
    list of bool
        True if preferred, else False
    """
    POINTS = {
        "scheme": {"bus-demeo": 3, "bus": 2, "smass": 2, "tholen": 1, "sdss": 1},
        "waverange": {"vis": 1, "nir": 3, "visnir": 6, "mix": 4},
        "method": {"spec": 7, "phot": 3, "mix": 4},
    }

    taxonomies = [dict(zip(taxonomies, t)) for t in zip(*taxonomies.values())]

    # Compute points of each classification
    points = []

    for row in taxonomies:
        points.append(
            sum(
                [
                    POINTS[crit][row[crit].lower()]
                    for crit in ["scheme", "waverange", "method"]
                ]
            )
        )

    # Convert points to boolean
    preferred = [True if p == max(points) else False for p in points]
    return preferred

def select_numeric_property(obs, prop_name):
    """Select preferred observations ranking methods.

    Parameters
    ----------
    obs : dict
        Property measurements and metadata retrieved from SsODNet:datacloud.
    prop_name : str
        Name of the asteroid property.

    Returns
    -------
    list of bool
        True if selected, else False.

    Notes
    -----
    The method ranking depends on the observable.
    """
    RANKING = PROPERTIES[prop_name]["ranking"]
    methods = set(obs["method"])

    for method in RANKING:

        if set(method) & methods:  # method used at least once

            # Ensure that rows do not contain all 0 values, as can be the case
            # for albedo in diamalbedo
            if all(
                [
                    obs[prop_name][i] == 0
                    for i, m in enumerate(obs["method"])
                    if m in method
                ]
            ):
                continue

            # All entries using this method are preferred
            preferred = [True if m in method else False for m in obs["method"]]
            return preferred

    # No property entry is preferred (likely all 0), return list of False
    return [False for o in obs["number"]]

# In alphabetic order
PROPERTIES = {
    # TEMPLATE
    # property_name: Rock instance attribute name
    #   attribute: attribute key in datacloud
    #   collection: Rock instance attribute name for collection of parameters
    #               Use plural form
    #   extra_columns: names of additional columns to output on CLI
    #   ssodnet_path: json path to asteroid property
    #                 (to be replaced by ssoCard
    #   type: asteroid property type, float or str
    "albedo": {
        "ranking": [
            ["SPACE"],
            ["ADAM", "KOALA", "SAGE", "Radar"],
            ["LC+TPM", "TPM", "LC+AO", "LC+Occ", "TE-IM", "TE-Occ"],
            ["AO", "Occ", "IM"],
            ["NEATM"],
            ["STM"],
        ],
    },
    "diameter": {
        "ranking": [
            ["SPACE"],
            ["ADAM", "KOALA", "SAGE", "Radar"],
            ["LC+TPM", "TPM", "LC+AO", "LC+Occ", "TE-IM", "TE-Occ"],
            ["AO", "Occ", "IM"],
            ["NEATM"],
            ["STM"],
        ],
    },
    "mass": {
        "ranking": [
            ["SPACE"],
            ["Bin-Genoid"],
            ["Bin-IM", "Bin-Radar", "Bin-PheMu"],
            ["EPHEM", "DEFLECT"],
        ],
    },
}

# Classes to complexes mapping
CLASS_TO_COMPLEX = {
    "A": "A",
    "Ad": "U",
    "AQ": "U",
    "AS": "U",
    "AU": "U",
    "AV": "U",
    "B": "B",
    "BC": "C",
    "BCF": "C",
    "BCU": "U",
    "BFC": "C",
    "BFU": "U",
    "BFX": "U",
    "Bk": "U",
    "BU": "B",
    "C": "C",
    "Caa": "C",
    "Cb": "C",
    "CB": "C",
    "CBU": "C",
    "CD": "U",
    "CDX": "U",
    "CF": "C",
    "CFB": "C",
    "CFU": "U",
    "CFXU": "U",
    "Cg": "C",
    "CG": "C",
    "CGSU": "U",
    "CGTP": "U",
    "CGU": "U",
    "Cgx": "C",
    "CL": "U",
    "CO": "U",
    "CP": "U",
    "CPF": "U",
    "CPU": "U",
    "CQ": "U",
    "CS": "U",
    "CSGU": "U",
    "CSU": "U",
    "CTGU": "U",
    "CU": "U",
    "CX": "U",
    "CXF": "U",
    "Cgh": "Ch",
    "Ch": "Ch",
    "D": "D",
    "DCX": "U",
    "DL": "D",
    "DP": "D",
    "DU": "D",
    "Ds": "D",
    "DS": "D",
    "DSU": "D",
    "DT": "D",
    "DTU": "D",
    "DU": "D",
    "DX": "D",
    "DXCU": "D",
    "E": "E",
    "EM": "X",
    "EU": "U",
    "F": "C",
    "FBCU": "U",
    "FC": "C",
    "FCB": "C",
    "FCU": "U",
    "FCX": "U",
    "FP": "U",
    "FU": "U",
    "FX": "U",
    "FXU": "U",
    "G": "C",
    "GC": "C",
    "GS": "U",
    "J": "V",
    "K": "K",
    "Kl": "U",
    "L": "L",
    "LA": "U",
    "Ld": "L",
    "LQ": "U",
    "LS": "U",
    "M": "M",
    "MU": "U",
    "O": "O",
    "OV": "U",
    "P": "P",
    "PC": "U",
    "PCD": "U",
    "PD": "U",
    "PDC": "U",
    "PF": "U",
    "PU": "U",
    "Q": "Q",
    "QO": "Q",
    "QRS": "U",
    "QSV": "U",
    "QV": "U",
    "Qw": "Q",
    "R": "R",
    "S": "S",
    "SA": "S",
    "Sa": "S",
    "SC": "U",
    "SCTU": "U",
    "SD": "U",
    "SDU": "U",
    "SG": "U",
    "Sk": "S",
    "Sl": "S",
    "SMU": "U",
    "SO": "S",
    "Sq": "S",
    "SQ": "S",
    "Sqw": "S",
    "Sr": "S",
    "SR": "S",
    "Srw": "S",
    "ST": "U",
    "STD": "U",
    "STGD": "U",
    "STU": "U",
    "SU": "U",
    "SV": "S",
    "Sv": "S",
    "Svw": "S",
    "Sw": "S",
    "SX": "U",
    "T": "T",
    "TCG": "U",
    "TD": "U",
    "TDG": "U",
    "TDS": "U",
    "TS": "U",
    "TSD": "U",
    "TX": "U",
    "V": "V",
    "Vw": "V",
    "X": "X",
    "XB": "U",
    "Xc": "X",
    "XC": "U",
    "XCU": "U",
    "XD": "U",
    "XDC": "U",
    "Xe": "X",
    "XF": "U",
    "XFC": "U",
    "XFCU": "U",
    "XFU": "U",
    "Xk": "X",
    "XL": "U",
    "Xn": "X",
    "XS": "U",
    "XSC": "U",
    "XSCU": "U",
    "XT": "U",
    "Xt": "X",
    "XU": "U",
    "Z": "U",
    np.nan: None,
    "": None,
    None: None,
    float("nan"): None,
}


import os
import sys
import collections.abc
import multiprocessing
import queue
import signal

import numpy as np
import scipy.interpolate

from cosmosis.datablock import option_section, names

from cosmosis.datablock.cosmosis_py import errors

import bias_module

def setup(options):
    config = {}

    here = os.path.dirname(__file__)
    window_file = options[option_section, "window_file"]
    bands_file = options[option_section, "bands_file"]
    config["window"] = np.loadtxt(window_file)
    config["bands"] = np.loadtxt(bands_file)

    config["verbose"] = options.get_bool(option_section, "verbose", False)
    config["timeout"] = options.get_double(option_section, "timeout", default=-1.0)
    if config["timeout"] <= 0:
        config["timeout"] = None

    try:
        bands_range = options[option_section, "bands_range"]
    except errors.BlockNameNotFound:
        bands_range = [20, 160]
    if len(bands_range) != 2 or not all(np.issubdtype(type(b), np.integer) for b in bands_range):
        raise ValueError(f"bands_range needs to be two integers, got {bands_range}, {[type(b) for b in bands_range]}")

    try:
        points_range = options[option_section, "points_range"]
    except errors.BlockNameNotFound:
        points_range = [4, 32]
    if len(points_range) != 2 or not all(np.issubdtype(type(b), np.integer) for b in points_range):
        raise ValueError(f"points_range needs to be two integers, got {points_range}")

    n_points, n_bands = config["window"].shape

    cut_points_idx = list(range(points_range[0])) + list(range(points_range[1], n_points))
    cut_bands_idx = list(range(bands_range[0])) + list(range(bands_range[1], n_bands))

    if config["verbose"]:
        print("Cutting bands", cut_bands_idx)
        print("Cutting points", cut_points_idx)
    if len(cut_points_idx) > 0:
        config["window"] = np.delete(config["window"], cut_points_idx, axis=0)
    if len(cut_bands_idx) > 0:
        config["bands"] = np.delete(config["bands"], cut_bands_idx)
        config["window"] = np.delete(config["window"], cut_bands_idx, axis=1)

    config["output_section_wedges"] = options.get_string(option_section, "output_section_wedges", "xi_wedges")
    config["output_section_pk_mm"] = options.get_string(option_section, "output_section_pk_mm", "matter_matter_power_spectrum_pt")
    config["output_section_pk_gm"] = options.get_string(option_section, "output_section_pk_gm", "galaxy_matter_power_spectrum_pt")
    config["output_section_pk_gg"] = options.get_string(option_section, "output_section_pk_gg", "galaxy_galaxy_power_spectrum_pt")

    config["compute_lss_parameters"] = options.get_bool(option_section, "compute_lss_parameters", True)

    config["twopt_type"] = 4
    config["num_ell"] = 3
    config["num_points_use"] = config["window"].shape[0]
    config["num_bands_use"] = config["window"].shape[1]
    config["z_index"] = 4
    config["zm"] = options[option_section, "z_eff"]
    if not isinstance(config["zm"], collections.abc.Iterable):
        config["zm"] = [config["zm"]]
    config["om_fid"] = 0.31
    config["h0_fid"] = 0.7

    config["use_growth"] = options.get_bool(option_section, "use_growth", False)
    config["local_lag_g2"] = options.get_bool(option_section, "local_lag_g2", True)
    config["local_lag_g3"] = options.get_bool(option_section, "local_lag_g3", False)

    config["no_interpolation"] = options.get_bool(option_section, "no_interpolation", False)

    # config["derived_parameters"] = np.loadtxt(os.path.join(here, "../output/derived_params.txt"))
    # config["data_parameters"] = np.loadtxt(os.path.join(here, "../output/data_params.txt"))

    # config["z_pk"] = np.loadtxt(os.path.join(here, "../output/z_p_k.txt"))
    # config["log_k_pk"] = np.loadtxt(os.path.join(here, "../output/log_k_h.txt"))
    # config["log_pk"] = np.loadtxt(os.path.join(here, "../output/log_p_k.txt"))

    # if config["verbose"]:
    #     print(config)
    return config

def run_wedges(q, config,
               h, omdm, omb, omv, omk, omnuh2, nnu, w, wa,
               z_Pk, log_k_h, log_Pk,
               growth, sigma8,
               params):

    module = bias_module.BiasModule()
    wedges_config = []

    for z in config["zm"]:
        wedges_config.append(module.initialize_wedges(config["twopt_type"], config["num_ell"],
                                                         config["num_points_use"], config["num_bands_use"],
                                                         z, config["om_fid"], config["h0_fid"],
                                                         config["window"], verbose=config["verbose"]))

    # Setup cosmology
    module.setup_cosmology(h, omdm, omb, omv, omk, omnuh2, nnu, w, wa,
                                     config["use_growth"], config["local_lag_g2"], config["local_lag_g3"],
                                     z_Pk, log_k_h, log_Pk.T,
                                     growth, sigma8, verbose=config["verbose"])

    results = []
    #print("Running compute_wedges", flush=True)
    for i, zm in enumerate(config["zm"]):
        b = i + 1
        b1, b2, gamma2, gamma3, a_vir, gamma, z_index, H_z, DA_z = params[i]

        vtheo, vtheo_convolved, Pk_mm, Pk_gm, Pk_gg = module.compute_wedges(
                                                        wedges_config[i],
                                                        b1, b2, gamma2, gamma3, a_vir, gamma,
                                                        z_index+2,
                                                        H_z, DA_z,
                                                        config["bands"],
                                                        verbose=config["verbose"])
        results.append((vtheo, vtheo_convolved, Pk_mm, Pk_gm, Pk_gg))

    for c in wedges_config:
        module.cleanup_wedges(c)

    module.cleanup_cosmology()

    q.put(results)

def execute(block, config):
    h = block[names.cosmological_parameters, "h0"]
    omdm = block[names.cosmological_parameters, "omega_c"]
    omb = block[names.cosmological_parameters, "omega_b"]
    omv = block[names.cosmological_parameters, "omega_lambda"]
    omk = block[names.cosmological_parameters, "omega_k"]
    omnuh2 = block[names.cosmological_parameters, "omnuh2"]
    nnu = block.get_double(names.cosmological_parameters, "nnu", default=3.046)
    w = block[names.cosmological_parameters, "w"]
    wa = block[names.cosmological_parameters, "wa"]

    if config["use_growth"]:
        gamma = block["bias_parameters", f"gamma"]
    else:
        gamma = 1.0

    log_Pk = np.log(block[names.matter_power_lin, "p_k"])
    log_k_h = np.log(block[names.matter_power_lin, "k_h"])
    z_Pk = block[names.matter_power_lin, "z"]

    z_growth = block[names.growth_parameters, "z"]
    if not np.allclose(z_Pk, z_growth):
        raise ValueError("Redshifts of power spectrum and growth do not match.")

    sigma8 = block[names.growth_parameters, "SIGMA_8"]
    if block.has_value(names.growth_parameters, "fsigma_8"):
        growth = block[names.growth_parameters, "fsigma_8"]
    else:
        sigma2_vdelta_8 = block[names.growth_parameters, "SIGMA2_VDELTA_8"]
        growth = sigma2_vdelta_8/sigma8

    Pk_mm_pt = np.zeros((len(config["zm"]), len(log_k_h)))
    Pk_gm_pt = np.zeros((len(config["zm"]), len(log_k_h)))
    Pk_gg_pt = np.zeros((len(config["zm"]), len(log_k_h)))

    params = []
    for i, zm in enumerate(config["zm"]):
        b = i + 1

        z = block[names.distances, "z"]
        H_z = block[names.distances, "h"]*2.99792458e8/1e3
        DA_z = block[names.distances, "d_a"]

        if config["no_interpolation"]:
            # Original implementation taylored to the way the CosmoMC module reads
            # these parameters. Sensitive to z resolution
            z_index = np.argmin(np.abs(z-zm))

            if z_index >= (200-13)//4:
                # z_index too large to fit H_z and DA_z into derived_parameters array
                # so we take a range around z[z_index]
                s = slice(max(0, z_index-4), min(len(z), z_index+4))
                z = z[s]
                H_z = H_z[s]
                DA_z = DA_z[s]
                z_index = np.argmin(np.abs(z-zm))
        else:
            # More stable and sane.
            H_z_intp = scipy.interpolate.InterpolatedUnivariateSpline(z, H_z, ext=2)
            DA_z_intp = scipy.interpolate.InterpolatedUnivariateSpline(z, DA_z, ext=2)

            H_z = np.atleast_1d(H_z_intp(zm))
            DA_z = np.atleast_1d(DA_z_intp(zm))
            z_index = 0

        # Bias parameters
        b1 = block["bias_parameters", f"b1_bin_{b}"]
        b2 = block["bias_parameters", f"b2_bin_{b}"]
        if not config["local_lag_g2"]:
            gamma2 = block["bias_parameters", f"gamma2_bin_{b}"]
        else:
            gamma2 = 1.0
        if not config["local_lag_g3"]:
            gamma3 = block["bias_parameters", f"gamma3_bin_{b}"]
        else:
            gamma3 = 1.0
        a_vir = block["bias_parameters", f"a_vir_bin_{b}"]

        params.append((b1, b2, gamma2, gamma3, a_vir, gamma, z_index, H_z, DA_z))

    # Run wedges
    # Need fork. Using spawn launches the whole cosmosis process
    mp_context = multiprocessing.get_context("fork")

    # Queue to get results back from the process
    q = multiprocessing.Queue()

    # Create the process
    proc = mp_context.Process(target=run_wedges,
                              args=(q, config, h, omdm, omb, omv, omk, omnuh2, nnu, w, wa,
                                    z_Pk, log_k_h, log_Pk,
                                    growth, sigma8,
                                    params))
    proc.start()
    try:
        # Wait for results to show up in the queue
        result = q.get(block=True, timeout=config["timeout"])
    except queue.Empty:
        print(f"wedges module timed out after {config['timeout']} s. Attempting to tell process to stop.", file=sys.stderr, flush=True)
        # os.kill(proc.pid, signal.SIGTERM)
        proc.join(0.5)
        if proc.is_alive():
            print("Wedges process is not cooperating. Terminating it.", file=sys.stderr, flush=True)
            proc.terminate()
            proc.join(0.5)
            if proc.exitcode is None:
                print("So you have choosen death.", file=sys.stderr, flush=True)
                proc.kill()
        return 1

    proc.join(0.5)
    if proc.is_alive():
        print("Wedges process is not cooperating. Terminating it.", file=sys.stderr, flush=True)
        proc.terminate()
        proc.join(0.5)
        if proc.exitcode is None:
            print("So you have choosen death.", file=sys.stderr, flush=True)
            proc.kill()

    proc.close()

    # Put the results into the datablock
    for i, zm in enumerate(config["zm"]):
        b = i + 1

        vtheo, vtheo_convolved, Pk_mm, Pk_gm, Pk_gg = result[i]

        n = len(vtheo)//config["num_ell"]
        vtheo = np.array([vtheo[i*n:(i+1)*n] for i in range(config["num_ell"])])
        n = len(vtheo_convolved)//config["num_ell"]
        vtheo_convolved = np.array([vtheo_convolved[i*n:(i+1)*n] for i in range(config["num_ell"])])

        block[config["output_section_wedges"], f"vtheo_bin_{b}"] = vtheo
        block[config["output_section_wedges"], f"bin_{b}"] = vtheo_convolved
        Pk_mm_pt[i] = Pk_mm
        Pk_gm_pt[i] = Pk_gm
        Pk_gg_pt[i] = Pk_gg

    block[config["output_section_wedges"], "n_wedge"] = config["num_ell"]
    block[config["output_section_wedges"], "bands"] = config["bands"]
    block[config["output_section_wedges"], "z"] = config["zm"]

    block[config["output_section_pk_mm"], "z"] = config["zm"]
    block[config["output_section_pk_gm"], "z"] = config["zm"]
    block[config["output_section_pk_gg"], "z"] = config["zm"]

    block[config["output_section_pk_mm"], "k_h"] = np.exp(log_k_h)
    block[config["output_section_pk_gm"], "k_h"] = np.exp(log_k_h)
    block[config["output_section_pk_gg"], "k_h"] = np.exp(log_k_h)

    block[config["output_section_pk_mm"], "p_k"] = Pk_mm_pt
    block[config["output_section_pk_gm"], "p_k"] = Pk_gm_pt
    block[config["output_section_pk_gg"], "p_k"] = Pk_gg_pt

    if config["compute_lss_parameters"]:
        z = block[names.growth_parameters, "z"]
        fsigma_8 = scipy.interpolate.InterpolatedUnivariateSpline(z, block[names.growth_parameters, "fsigma_8"])

        z_background = block[names.distances, "z"]
        F_AP = scipy.interpolate.InterpolatedUnivariateSpline(z_background, block[names.distances, "F_AP"])
        rs_DV = scipy.interpolate.InterpolatedUnivariateSpline(z_background[1:], block[names.distances, "rs_DV"][1:])

        for i, zm in enumerate(config["zm"]):
            b = i + 1

            block["lss_parameters", f"rs_DV_bin_{b}"] = float(rs_DV(zm))
            block["lss_parameters", f"F_AP_bin_{b}"] = float(F_AP(zm))
            block["lss_parameters", f"fsigma_8_bin_{b}"] = float(fsigma_8(zm))

    return 0

def cleanup(config):
    pass

# def execute(block, config):
#     h = block[names.cosmological_parameters, "h0"]
#     omdm = block[names.cosmological_parameters, "omega_c"]
#     omb = block[names.cosmological_parameters, "omega_b"]
#     omv = block[names.cosmological_parameters, "omega_lambda"]
#     omk = block[names.cosmological_parameters, "omega_k"]
#     omnuh2 = block[names.cosmological_parameters, "omnuh2"]
#     nnu = block.get_double(names.cosmological_parameters, "nnu", default=3.046)
#     w = block[names.cosmological_parameters, "w"]
#     wa = block[names.cosmological_parameters, "wa"]

#     if config["use_growth"]:
#         gamma = block["bias_parameters", f"gamma"]
#     else:
#         gamma = 1.0

#     log_Pk = np.log(block[names.matter_power_lin, "p_k"])
#     log_k_h = np.log(block[names.matter_power_lin, "k_h"])
#     z_Pk = block[names.matter_power_lin, "z"]

#     z_growth = block[names.growth_parameters, "z"]
#     if not np.allclose(z_Pk, z_growth):
#         raise ValueError("Redshifts of power spectrum and growth do not match.")

#     sigma8 = block[names.growth_parameters, "SIGMA_8"]
#     if block.has_value(names.growth_parameters, "fsigma_8"):
#         growth = block[names.growth_parameters, "fsigma_8"]
#     else:
#         sigma2_vdelta_8 = block[names.growth_parameters, "SIGMA2_VDELTA_8"]
#         growth = sigma2_vdelta_8/sigma8

#     # Setup cosmology
#     config["module"].setup_cosmology(h, omdm, omb, omv, omk, omnuh2, nnu, w, wa,
#                                      config["use_growth"], config["local_lag_g2"], config["local_lag_g3"],
#                                      z_Pk, log_k_h, log_Pk.T,
#                                      growth, sigma8, verbose=config["verbose"])

#     Pk_mm_pt = np.zeros((len(config["zm"]), len(log_k_h)))
#     Pk_gm_pt = np.zeros((len(config["zm"]), len(log_k_h)))
#     Pk_gg_pt = np.zeros((len(config["zm"]), len(log_k_h)))

#     for i, zm in enumerate(config["zm"]):
#         b = i + 1

#         z = block[names.distances, "z"]
#         H_z = block[names.distances, "h"]*2.99792458e8/1e3
#         DA_z = block[names.distances, "d_a"]
#         z_index = np.argmin(np.abs(z-zm))

#         if z_index >= (200-13)//4:
#             # z_index too large to fit H_z and DA_z into derived_parameters array
#             # so we take a range around z[z_index]
#             s = slice(max(0, z_index-4), min(len(z), z_index+4))
#             z = z[s]
#             H_z = H_z[s]
#             DA_z = DA_z[s]
#             z_index = np.argmin(np.abs(z-zm))

#         # Bias parameters
#         b1 = block["bias_parameters", f"b1_bin_{b}"]
#         b2 = block["bias_parameters", f"b2_bin_{b}"]
#         if not config["local_lag_g2"]:
#             gamma2 = block["bias_parameters", f"gamma2_bin_{b}"]
#         else:
#             gamma2 = 1.0
#         if not config["local_lag_g3"]:
#             gamma3 = block["bias_parameters", f"gamma3_bin_{b}"]
#         else:
#             gamma3 = 1.0
#         a_vir = block["bias_parameters", f"a_vir_bin_{b}"]

#         vtheo, vtheo_convolved, Pk_mm, Pk_gm, Pk_gg = config["module"].compute_wedges(
#                                                         config["wedges_config"][i],
#                                                         b1, b2, gamma2, gamma3, a_vir, gamma,
#                                                         z_index+2,
#                                                         H_z, DA_z,
#                                                         config["bands"],
#                                                         verbose=config["verbose"])

#         n = len(vtheo)//config["num_ell"]
#         vtheo = np.array([vtheo[i*n:(i+1)*n] for i in range(config["num_ell"])])
#         n = len(vtheo_convolved)//config["num_ell"]
#         vtheo_convolved = np.array([vtheo_convolved[i*n:(i+1)*n] for i in range(config["num_ell"])])

#         block[config["output_section_wedges"], f"vtheo_bin_{b}"] = vtheo
#         block[config["output_section_wedges"], f"bin_{b}"] = vtheo_convolved
#         Pk_mm_pt[i] = Pk_mm
#         Pk_gm_pt[i] = Pk_gm
#         Pk_gg_pt[i] = Pk_gg

#     block[config["output_section_wedges"], "n_wedge"] = config["num_ell"]
#     block[config["output_section_wedges"], "bands"] = config["bands"]
#     block[config["output_section_wedges"], "z"] = config["zm"]

#     block[config["output_section_pk_mm"], "z"] = config["zm"]
#     block[config["output_section_pk_gm"], "z"] = config["zm"]
#     block[config["output_section_pk_gg"], "z"] = config["zm"]

#     block[config["output_section_pk_mm"], "k_h"] = np.exp(log_k_h)
#     block[config["output_section_pk_gm"], "k_h"] = np.exp(log_k_h)
#     block[config["output_section_pk_gg"], "k_h"] = np.exp(log_k_h)

#     block[config["output_section_pk_mm"], "p_k"] = Pk_mm_pt
#     block[config["output_section_pk_gm"], "p_k"] = Pk_gm_pt
#     block[config["output_section_pk_gg"], "p_k"] = Pk_gg_pt

#     if config["compute_lss_parameters"]:
#         z = block[names.growth_parameters, "z"]
#         fsigma_8 = scipy.interpolate.InterpolatedUnivariateSpline(z, block[names.growth_parameters, "fsigma_8"])

#         z_background = block[names.distances, "z"]
#         F_AP = scipy.interpolate.InterpolatedUnivariateSpline(z_background, block[names.distances, "F_AP"])
#         rs_DV = scipy.interpolate.InterpolatedUnivariateSpline(z_background[1:], block[names.distances, "rs_DV"][1:])

#         for i, zm in enumerate(config["zm"]):
#             b = i + 1

#             block["lss_parameters", f"rs_DV_bin_{b}"] = float(rs_DV(zm))
#             block["lss_parameters", f"F_AP_bin_{b}"] = float(F_AP(zm))
#             block["lss_parameters", f"fsigma_8_bin_{b}"] = float(fsigma_8(zm))

#     return 0

# def cleanup(config):
#     for c in config["wedges_config"]:
#         config["module"].cleanup_wedges(c)

#     config["module"].cleanup_cosmology()


# (C) Crown Copyright, Met Office. All rights reserved.
#
# This file is part of ocean_error_covs and is released under the BSD 3-Clause license.
# See LICENSE in the root of the repository for full licensing details.
#######################################################################
import numpy as np
import arrays

class ObsProfiles():
      """Class of functions to deal with profile observations"""

      def __init__(self):
         self.undef = 1000

      def random_subsample_profiles(self, depths, depth_range, fdbk_var_array):
          """ Thin profile observations vertically by choosing a single observation
              at random within a pre-defined depth range.

              **** PARAMETERS ****
              1. depths: Array of depths of the profile observations
              2. depth_range: Tuple defining boundaries over which to sub-sample
                              profile observations
              3. fdbk_var_arrays: object containing the arrays describing the
                                  innovations

              ***** RETURNS *****
              1. subsampled_fdbk_var_array: object containing the sub-sampled
                                            arrays describing the innovations
                                            with a single obs for each depth level
          """

          # Define new tuples
          lat_dep=[] ; lon_dep =[] ; qc_dep =[] ; mod_dep =[] ; obs_dep = []

          # Loop over obs number
          for n in range(0, len(fdbk_var_array.lats)):
              dep_mask = np.logical_and(depths[n,:] >= depth_range[0],
                                        depths[n,:] <  depth_range[1])
              num_in_mask=np.sum(dep_mask)
              if num_in_mask > 0:
                 lat_dep += [fdbk_var_array.lats[n]]
                 lon_dep += [fdbk_var_array.lons[n]]
                 qc_dep += [fdbk_var_array.obs_qc[n]]
                 random_pick = np.random.random_integers(0, num_in_mask-1)       # NOTE: Results will not be
                 obs_dep += [fdbk_var_array.obs_vals[n,dep_mask][random_pick]]   # reproduceable unless a
                 mod_dep += [fdbk_var_array.mod_vals[n,dep_mask][random_pick]]   # random seed is pre-selected

          # define a new fdbk array object
          subsampled_fdbk_var_array = arrays.FdbkVarArrays()
          subsampled_fdbk_var_array.lats = np.array(lat_dep)
          subsampled_fdbk_var_array.lons = np.array(lon_dep)
          subsampled_fdbk_var_array.obs_qc = np.array(qc_dep)
          subsampled_fdbk_var_array.obs_vals = np.array(obs_dep)
          subsampled_fdbk_var_array.mod_vals = np.array(mod_dep)

          # additional masking for profiles
          subsampled_fdbk_var_array.obs_qc[np.abs(subsampled_fdbk_var_array.obs_vals)>self.undef] = self.undef
          return subsampled_fdbk_var_array


from __future__ import print_function
from coordinate_systems import Distance, Angle, Dihedral, OutOfPlane
from utilities import nifty, options, block_matrix
from wrappers import Molecule
from utilities.manage_xyz import write_molden_geoms
# standard library imports
import sys
import os
from os import path

# third party
import numpy as np
import multiprocessing as mp
from collections import Counter
from copy import copy
from itertools import chain

# local application imports
sys.path.append(path.dirname(path.dirname(path.abspath(__file__))))

def worker(arg):
   obj, methname = arg[:2]
   return getattr(obj, methname)(*arg[2:])

#######################################################################################
#### This class contains the main constructor, object properties and staticmethods ####
#######################################################################################

# TODO interpolate is still sloppy. It shouldn't create a new molecule node itself
# but should create the xyz. GSM should create the new molecule based off that xyz.
# TODO nconstraints in ic_reparam and write_iters is irrelevant

class GSM(object):

    from utilities import units

    @staticmethod
    def default_options():
        if hasattr(GSM, '_default_options'):
            return GSM._default_options.copy()

        opt = options.Options()

        opt.add_option(
            key='reactant',
            required=True,
            # allowed_types=[Molecule,wrappers.Molecule],
            doc='Molecule object as the initial reactant structure')

        opt.add_option(
            key='product',
            required=False,
            # allowed_types=[Molecule,wrappers.Molecule],
            doc='Molecule object for the product structure (not required for single-ended methods.')

        opt.add_option(
            key='nnodes',
            required=False,
            value=1,
            allowed_types=[int],
            doc="number of string nodes"
        )

        opt.add_option(
            key='optimizer',
            required=True,
            doc='Optimzer object  to use e.g. eigenvector_follow, conjugate_gradient,etc. \
                        most of the default options are okay for here since GSM will change them anyway',
        )

        opt.add_option(
            key='driving_coords',
            required=False,
            value=[],
            allowed_types=[list],
            doc='Provide a list of tuples to select coordinates to modify atoms\
                 indexed at 1')

        opt.add_option(
            key='CONV_TOL',
            value=0.0005,
            required=False,
            allowed_types=[float],
            doc='Convergence threshold'
        )

        opt.add_option(
            key='CONV_gmax',
            value=0.001,
            required=False,
            allowed_types=[float],
            doc='Convergence threshold'
        )

        opt.add_option(
            key='CONV_Ediff',
            value=0.1,
            required=False,
            allowed_types=[float],
            doc='Convergence threshold'
        )

        opt.add_option(
            key='CONV_dE',
            value=0.5,
            required=False,
            allowed_types=[float],
            doc='Convergence threshold'
        )

        opt.add_option(
            key='ADD_NODE_TOL',
            value=0.1,
            required=False,
            allowed_types=[float],
            doc='Convergence threshold')

        opt.add_option(
            key="growth_direction",
            value=0,
            required=False,
            doc="how to grow string,0=Normal,1=from reactant"
        )

        opt.add_option(
            key="DQMAG_MAX",
            value=0.8,
            required=False,
            doc="max step along tangent direction for SSM"
        )
        opt.add_option(
            key="DQMAG_MIN",
            value=0.2,
            required=False,
            doc=""
        )

        opt.add_option(
            key='print_level',
            value=1,
            required=False
        )

        opt.add_option(
            key='xyz_writer',
            value=write_molden_geoms,
            required=False,
            doc='Function to be used to format and write XYZ files',
        )

        opt.add_option(
            key='mp_cores',
            value=1,
            doc='multiprocessing cores for parallel programming. Use this with caution.',
        )

        opt.add_option(
            key="BDIST_RATIO",
            value=0.5,
            required=False,
            doc="SE-Crossing uses this \
                        bdist must be less than 1-BDIST_RATIO of initial bdist in order to be \
                        to be considered grown.",
        )

        opt.add_option(
            key='ID',
            value=0,
            required=False,
            doc='A number for identification of Strings'
        )

        opt.add_option(
            key='interp_method',
            value='DLC',
                allowed_values=['Geodesic', 'DLC'],
            required=False,
            doc='Which reparameterization method to use',
        )

        opt.add_option(
            key='noise',
            value=100.0,
            allowed_types=[float],
            required=False,
            doc='Noise to check for intermediate',
        )

        GSM._default_options = opt
        return GSM._default_options.copy()

    @classmethod
    def from_options(cls, **kwargs):
        return cls(cls.default_options().set_values(kwargs))

    @classmethod
    def copy_from_options(cls, gsm_obj, reactant, product):
        new_gsm = cls.from_options(gsm_obj.options.copy().set_values({'reactant': reactant, 'product': product}))
        return new_gsm

    def __init__(
            self,
            options,
    ):
        """ Constructor """
        self.options = options

        os.system('mkdir -p scratch')

        # Cache attributes
        self.nnodes = self.options['nnodes']
        self.nodes = [None]*self.nnodes
        self.nodes[0] = self.options['reactant']
        self.nodes[-1] = self.options['product']
        self.driving_coords = self.options['driving_coords']
        self.growth_direction = self.options['growth_direction']
        self.isRestarted = False
        self.DQMAG_MAX = self.options['DQMAG_MAX']
        self.DQMAG_MIN = self.options['DQMAG_MIN']
        self.BDIST_RATIO = self.options['BDIST_RATIO']
        self.ID = self.options['ID']
        self.optimizer = []
        self.interp_method = self.options['interp_method']
        self.CONV_TOL = self.options['CONV_TOL']
        self.noise = self.options['noise']
        self.mp_cores = self.options['mp_cores']
        self.xyz_writer = self.options['xyz_writer']

        optimizer = options['optimizer']
        for count in range(self.nnodes):
            self.optimizer.append(optimizer.__class__(optimizer.options.copy()))
        self.print_level = options['print_level']

        # Set initial values
        self.current_nnodes = 2
        self.nR = 1
        self.nP = 1
        self.climb = False
        self.find = False
        self.ts_exsteps = 3  # multiplier for ts node
        self.n0 = 1  # something to do with added nodes? "first node along current block"
        self.end_early = False
        self.tscontinue = True  # whether to continue with TS opt or not
        self.found_ts = False
        self.rn3m6 = np.sqrt(3.*self.nodes[0].natoms-6.)
        self.gaddmax = self.options['ADD_NODE_TOL']  # self.options['ADD_NODE_TOL']/self.rn3m6;
        print(" gaddmax:", self.gaddmax)
        self.ictan = [None]*self.nnodes
        self.active = [False] * self.nnodes
        self.climber = False  # is this string a climber?
        self.finder = False   # is this string a finder?
        self.done_growing = False
        self.nclimb = 0
        self.nhessreset = 10  # are these used??? TODO
        self.hessrcount = 0   # are these used?!  TODO
        self.hess_counter = 0   # it is probably good to reset the hessian
        self.newclimbscale = 2.
        self.TS_E_0 = None
        self.dE_iter = 100.  # change in max TS node

        self.nopt_intermediate = 0     # might be a duplicate of endearly_counter
        self.flag_intermediate = False
        self.endearly_counter = 0  # Find the intermediate x time
        self.pot_min = []
        self.ran_out = False   # if it ran out of iterations

        self.newic = Molecule.copy_from_options(self.nodes[0])  # newic object is used for coordinate transformations

    @property
    def TSnode(self):
        '''
        The current node with maximum energy
        '''
        # Treat GSM with penalty a little different since penalty will increase energy based on energy
        # differences, which might not be great for Climbing Image
        if self.__class__.__name__ != "SE_Cross" and self.nodes[0].PES.__class__.__name__ == "Penalty_PES":
            energies = np.asarray([0.]*self.nnodes)
            for i, node in enumerate(self.nodes):
                if node is not None:
                    energies[i] = (node.PES.PES1.energy + node.PES.PES2.energy)/2.
            return int(np.argmax(energies))
        else:
            # make sure TS is not zero or last node
            return int(np.argmax(self.energies[1:self.nnodes-1])+1)

    @property
    def emax(self):
        return self.energies[self.TSnode]

    @property
    def npeaks(self):
        '''
        '''
        minnodes = []
        maxnodes = []
        energies = self.energies
        if energies[1] > energies[0]:
            minnodes.append(0)
        if energies[self.nnodes-1] < energies[self.nnodes-2]:
            minnodes.append(self.nnodes-1)
        for n in range(self.n0, self.nnodes-1):
            if energies[n+1] > energies[n]:
                if energies[n] < energies[n-1]:
                    minnodes.append(n)
            if energies[n+1] < energies[n]:
                if energies[n] > energies[n-1]:
                    maxnodes.append(n)

        return len(maxnodes)

    @property
    def energies(self):
        '''
        Energies of string
        '''
        E = []
        for ico in self.nodes:
            if ico is not None:
                E.append(ico.energy - self.nodes[0].energy)
        return E

    @energies.setter
    def energies(self, list_of_E):
        '''
        setter for energies
        '''
        self.E = list_of_E

    @property
    def geometries(self):
        geoms = []
        for ico in self.nodes:
            if ico is not None:
                geoms.append(ico.geometry)
        return geoms

    @property
    def gradrmss(self):
        self._gradrmss = []
        for ico in self.nodes:
            if ico is not None:
                self._gradrmss.append(ico.gradrms)
        return self._gradrmss

    @property
    def dEs(self):
        self._dEs = []
        for ico in self.nodes:
            if ico is not None:
                self._dEs.append(ico.difference_energy)
        return self._dEs

    @property
    def ictan(self):
        return self._ictan

    @ictan.setter
    def ictan(self, value):
        self._ictan = value

    @property
    def dqmaga(self):
        return self._dqmaga

    @dqmaga.setter
    def dqmaga(self, value):
        self._dqmaga = value

    @staticmethod
    def add_xyz_along_tangent(
            xyz1,
            constraints,
            step,
            coord_obj,
    ):
        dq0 = step*constraints
        new_xyz = coord_obj.newCartesian(xyz1, dq0)

        return new_xyz

    @staticmethod
    def add_node(
            nodeR,
            nodeP,
            stepsize,
            node_id,
            **kwargs
    ):
        '''
        Add a node between  nodeR and nodeP or if nodeP is none use driving coordinate to add new node
        '''

        # get driving coord
        driving_coords = kwargs.get('driving_coords', None)
        DQMAG_MAX = kwargs.get('DQMAG_MAX', 0.8)
        DQMAG_MIN = kwargs.get('DQMAG_MIN', 0.2)

        if nodeP is None:

            if driving_coords is None:
                raise RuntimeError("You didn't supply a driving coordinate and product node is None!")

            BDISTMIN = 0.05
            ictan, bdist = GSM.get_tangent(nodeR, None, driving_coords=driving_coords)

            if bdist < BDISTMIN:
                print("bdist too small %.3f" % bdist)
                return None
            new_node = Molecule.copy_from_options(nodeR, new_node_id=node_id)
            new_node.update_coordinate_basis(constraints=ictan)
            constraint = new_node.constraints[:, 0]
            sign = -1.

            dqmag_scale = 1.5
            minmax = DQMAG_MAX - DQMAG_MIN
            a = bdist/dqmag_scale
            if a > 1.:
                a = 1.
            dqmag = sign*(DQMAG_MIN+minmax*a)
            if dqmag > DQMAG_MAX:
                dqmag = DQMAG_MAX
            print(" dqmag: %4.3f from bdist: %4.3f" % (dqmag, bdist))

            dq0 = dqmag*constraint
            print(" dq0[constraint]: %1.3f" % dqmag)

            new_node.update_xyz(dq0)
            new_node.bdist = bdist

        else:
            ictan, _ = GSM.get_tangent(nodeR, nodeP)
            nodeR.update_coordinate_basis(constraints=ictan)
            constraint = nodeR.constraints[:, 0]
            dqmag = np.linalg.norm(ictan)
            print(" dqmag: %1.3f" % dqmag)
            # sign=-1
            sign = 1.
            dqmag *= (sign*stepsize)
            print(" scaled dqmag: %1.3f" % dqmag)

            dq0 = dqmag*constraint
            old_xyz = nodeR.xyz.copy()
            new_xyz = nodeR.coord_obj.newCartesian(old_xyz, dq0)
            new_node = Molecule.copy_from_options(MoleculeA=nodeR, xyz=new_xyz, new_node_id=node_id)

        return new_node

    @staticmethod
    def interpolate_xyz(nodeR, nodeP, stepsize):
        '''
        Interpolate between two nodes
        '''
        ictan, _ = GSM.get_tangent(nodeR, nodeP)
        Vecs = nodeR.update_coordinate_basis(constraints=ictan)
        constraint = nodeR.constraints[:, 0]
        prim_constraint = block_matrix.dot(Vecs, constraint)
        dqmag = np.dot(prim_constraint.T, ictan)
        print(" dqmag: %1.3f" % dqmag)
        # sign=-1
        sign = 1.
        dqmag *= (sign*stepsize)
        print(" scaled dqmag: %1.3f" % dqmag)

        dq0 = dqmag*constraint
        old_xyz = nodeR.xyz.copy()
        new_xyz = nodeR.coord_obj.newCartesian(old_xyz, dq0)

        return new_xyz

    @staticmethod
    def interpolate(start_node, end_node, num_interp):
        '''

        '''
        nifty.printcool(" interpolate")

        num_nodes = num_interp + 2
        nodes = [None]*(num_nodes)
        nodes[0] = start_node
        nodes[-1] = end_node
        sign = 1
        nR = 1
        nP = 1
        nn = nR + nP

        for n in range(num_interp):
            if num_nodes - nn > 1:
                stepsize = 1./float(num_nodes - nn)
            else:
                stepsize = 0.5
            if sign == 1:
                iR = nR-1
                iP = num_nodes - nP
                iN = nR
                nodes[nR] = GSM.add_node(nodes[iR], nodes[iP], stepsize, iN)
                if nodes[nR] is None:
                    raise RuntimeError

                # print(" Energy of node {} is {:5.4}".format(nR,nodes[nR].energy-E0))
                nR += 1
                nn += 1

            else:
                n1 = num_nodes - nP
                n2 = n1 - 1
                n3 = nR - 1
                nodes[n2] = GSM.add_node(nodes[n1], nodes[n3], stepsize, n2)
                if nodes[n2] is None:
                    raise RuntimeError
                # print(" Energy of node {} is {:5.4}".format(nR,nodes[nR].energy-E0))
                nP += 1
                nn += 1
            sign *= -1

        return nodes

    @staticmethod
    def get_tangent_xyz(xyz1, xyz2, prim_coords):
        PMDiff = np.zeros(len(prim_coords))
        for k, prim in enumerate(prim_coords):
            if type(prim) is Distance:
                PMDiff[k] = 2.5 * prim.calcDiff(xyz2, xyz1)
            else:
                PMDiff[k] = prim.calcDiff(xyz2, xyz1)
        return np.reshape(PMDiff, (-1, 1))

    @staticmethod
    def get_tangent(node1, node2, print_level=1, **kwargs):
        '''
        Get internal coordinate tangent between two nodes, assumes they have unique IDs
        '''

        if node2 is not None and node1.node_id != node2.node_id:
            print(" getting tangent from between %i %i pointing towards %i" % (node2.node_id, node1.node_id, node2.node_id))
            assert node2 != None, 'node n2 is None'

            PMDiff = np.zeros(node2.num_primitives)
            for k, prim in enumerate(node2.primitive_internal_coordinates):
                if type(prim) is Distance:
                    PMDiff[k] = 2.5 * prim.calcDiff(node2.xyz, node1.xyz)
                else:
                    PMDiff[k] = prim.calcDiff(node2.xyz, node1.xyz)

            return np.reshape(PMDiff, (-1, 1)), None
        else:
            print(" getting tangent from node ", node1.node_id)

            driving_coords = kwargs.get('driving_coords', None)
            assert driving_coords is not None, " Driving coord is None!"

            c = Counter(elem[0] for elem in driving_coords)
            nadds = c['ADD']
            nbreaks = c['BREAK']
            nangles = c['nangles']
            ntorsions = c['ntorsions']

            ictan = np.zeros((node1.num_primitives, 1), dtype=float)
            # breakdq = 0.3
            bdist = 0.0
            atoms = node1.atoms
            xyz = node1.xyz.copy()

            for i in driving_coords:
                if "ADD" in i:

                    # order indices to avoid duplicate bonds
                    if i[1] < i[2]:
                        index = [i[1]-1, i[2]-1]
                    else:
                        index = [i[2]-1, i[1]-1]

                    bond = Distance(index[0], index[1])
                    prim_idx = node1.coord_obj.Prims.dof_index(index, 'Distance')
                    if len(i) == 3:
                        # TODO why not just use the covalent radii?
                        d0 = (atoms[index[0]].vdw_radius + atoms[index[1]].vdw_radius)/2.8
                    elif len(i) == 4:
                        d0 = i[3]
                    current_d = bond.value(xyz)

                    # TODO don't set tangent if value is too small
                    ictan[prim_idx] = -1*(d0-current_d)
                    # if nbreaks>0:
                    #    ictan[prim_idx] *= 2
                    # => calc bdist <=
                    if current_d > d0:
                        bdist += np.dot(ictan[prim_idx], ictan[prim_idx])
                    if print_level > 0:
                        print(" bond %s target (less than): %4.3f current d: %4.3f diff: %4.3f " % ((i[1], i[2]), d0, current_d, ictan[prim_idx]))

                elif "BREAK" in i:
                    # order indices to avoid duplicate bonds
                    if i[1] < i[2]:
                        index = [i[1]-1, i[2]-1]
                    else:
                        index = [i[2]-1, i[1]-1]
                    bond = Distance(index[0], index[1])
                    prim_idx = node1.coord_obj.Prims.dof_index(index, 'Distance')
                    if len(i) == 3:
                        d0 = (atoms[index[0]].vdw_radius + atoms[index[1]].vdw_radius)
                    elif len(i) == 4:
                        d0 = i[3]

                    current_d = bond.value(xyz)
                    ictan[prim_idx] = -1*(d0-current_d)

                    # => calc bdist <=
                    if current_d < d0:
                        bdist += np.dot(ictan[prim_idx], ictan[prim_idx])

                    if print_level > 0:
                        print(" bond %s target (greater than): %4.3f, current d: %4.3f diff: %4.3f " % ((i[1], i[2]), d0, current_d, ictan[prim_idx]))
                elif "ANGLE" in i:

                    if i[1] < i[3]:
                        index = [i[1]-1, i[2]-1, i[3]-1]
                    else:
                        index = [i[3]-1, i[2]-1, i[1]-1]
                    angle = Angle(index[0], index[1], index[2])
                    prim_idx = node1.coord_obj.Prims.dof_index(index, 'Angle')
                    anglet = i[4]
                    ang_value = angle.value(xyz)
                    ang_diff = anglet*np.pi/180. - ang_value
                    # print(" angle: %s is index %i " %(angle,ang_idx))
                    if print_level > 0:
                        print((" anglev: %4.3f align to %4.3f diff(rad): %4.3f" % (ang_value, anglet, ang_diff)))
                    ictan[prim_idx] = -ang_diff
                    # TODO need to come up with an adist
                    # if abs(ang_diff)>0.1:
                    #    bdist+=ictan[ICoord1.BObj.nbonds+ang_idx]*ictan[ICoord1.BObj.nbonds+ang_idx]
                elif "TORSION" in i:

                    if i[1] < i[4]:
                        index = [i[1]-1, i[2]-1, i[3]-1, i[4]-1]
                    else:
                        index = [i[4]-1, i[3]-1, i[2]-1, i[1]-1]
                    torsion = Dihedral(index[0], index[1], index[2], index[3])
                    prim_idx = node1.coord_obj.Prims.dof_index(index, 'Dihedral')
                    tort = i[5]
                    torv = torsion.value(xyz)
                    tor_diff = tort - torv*180./np.pi
                    if tor_diff > 180.:
                        tor_diff -= 360.
                    elif tor_diff < -180.:
                        tor_diff += 360.
                    ictan[prim_idx] = -tor_diff*np.pi/180.

                    if tor_diff*np.pi/180. > 0.1 or tor_diff*np.pi/180. < 0.1:
                        bdist += np.dot(ictan[prim_idx], ictan[prim_idx])
                    if print_level > 0:
                        print((" current torv: %4.3f align to %4.3f diff(deg): %4.3f" % (torv*180./np.pi, tort, tor_diff)))

                elif "OOP" in i:
                    index = [i[1]-1, i[2]-1, i[3]-1, i[4]-1]
                    oop = OutOfPlane(index[0], index[1], index[2], index[3])
                    prim_idx = node1.coord_obj.Prims.dof_index(index, 'OutOfPlane')
                    oopt = i[5]
                    oopv = oop.value(xyz)
                    oop_diff = oopt - oopv*180./np.pi
                    if oop_diff > 180.:
                        oop_diff -= 360.
                    elif oop_diff < -180.:
                        oop_diff += 360.
                    ictan[prim_idx] = -oop_diff*np.pi/180.

                    if oop_diff*np.pi/180. > 0.1 or oop_diff*np.pi/180. < 0.1:
                        bdist += np.dot(ictan[prim_idx], ictan[prim_idx])
                    if print_level > 0:
                        print((" current oopv: %4.3f align to %4.3f diff(deg): %4.3f" % (oopv*180./np.pi, oopt, oop_diff)))

            bdist = np.sqrt(bdist)
            if np.all(ictan == 0.0):
                raise RuntimeError(" All elements are zero")
            return ictan, bdist

    @staticmethod
    def get_tangents(nodes, n0=0, print_level=0):
        '''
        Get the normalized internal coordinate tangents and magnitudes between all nodes
        '''
        nnodes = len(nodes)
        dqmaga = [0.]*nnodes
        ictan = [[]]*nnodes

        for n in range(n0+1, nnodes):
            # print "getting tangent between %i %i" % (n,n-1)
            assert nodes[n] is not None, "n is bad"
            assert nodes[n-1] is not None, "n-1 is bad"
            ictan[n] = GSM.get_tangent_xyz(nodes[n-1].xyz, nodes[n].xyz, nodes[0].primitive_internal_coordinates)

            dqmaga[n] = 0.
            # ictan0= np.copy(ictan[n])
            dqmaga[n] = np.linalg.norm(ictan[n])

            ictan[n] /= dqmaga[n]

            # NOTE:
            # vanilla GSM has a strange metric for distance
            # no longer following 7/1/2020
            # constraint = self.newic.constraints[:,0]
            # just a fancy way to get the normalized tangent vector
            # prim_constraint = block_matrix.dot(Vecs,constraint)
            # for prim in self.newic.primitive_internal_coordinates:
            #    if type(prim) is Distance:
            #        index = self.newic.coord_obj.Prims.dof_index(prim)
            #        prim_constraint[index] *= 2.5
            # dqmaga[n] = float(np.dot(prim_constraint.T,ictan0))
            # dqmaga[n] = float(np.sqrt(dqmaga[n]))
            if dqmaga[n] < 0.:
                raise RuntimeError

        # TEMPORORARY parallel idea
        # ictan = [0.]
        # ictan += [ Process(target=get_tangent,args=(n,)) for n in range(n0+1,self.nnodes)]
        # dqmaga = [ Process(target=get_dqmag,args=(n,ictan[n])) for n in range(n0+1,self.nnodes)]

        if print_level > 1:
            print('------------printing ictan[:]-------------')
            for n in range(n0+1, nnodes):
                print("ictan[%i]" % n)
                print(ictan[n].T)
        if print_level > 0:
            print('------------printing dqmaga---------------')
            for n in range(n0+1, nnodes):
                print(" {:5.4}".format(dqmaga[n]), end='')
                if (n) % 5 == 0:
                    print()
            print()
        return ictan, dqmaga

    @staticmethod
    def get_three_way_tangents(nodes, energies, find=True, n0=0, print_level=0):
        '''
        Calculates internal coordinate tangent with a three-way tangent at TS node
        '''
        nnodes = len(nodes)
        ictan = [[]]*nnodes
        dqmaga = [0.]*nnodes
        # TSnode = np.argmax(energies[1:nnodes-1])+1
        TSnode = np.argmax(energies)   # allow for the possibility of TS node to be endpoints?

        last_node_max = (TSnode == nnodes-1)
        first_node_max = (TSnode == 0)
        if first_node_max or last_node_max:
            print("*********** This will cause a range error in the following for loop *********")
            print("** Setting the middle of the string to be TS node to get proper directions **")
            TSnode = nnodes//2

        for n in range(n0, nnodes):
            do3 = False
            print('getting tan[{' + str(n) + '}]')
            if n < TSnode:
                # The order is very important here
                # the way it should be ;(
                intic_n = n+1
                newic_n = n

                # old way
                # intic_n = n
                # newic_n = n+1

            elif n > TSnode:
                # The order is very important here
                intic_n = n
                newic_n = n-1
            else:
                do3 = True
                newic_n = n
                intic_n = n+1
                int2ic_n = n-1

            if do3:
                if first_node_max or last_node_max:
                    t1, _ = GSM.get_tangent(nodes[intic_n], nodes[newic_n])
                    t2, _ = GSM.get_tangent(nodes[newic_n], nodes[int2ic_n])
                    print(" done 3 way tangent")
                    ictan0 = t1 + t2
                else:
                    f1 = 0.
                    dE1 = abs(energies[n+1]-energies[n])
                    dE2 = abs(energies[n] - energies[n-1])
                    dEmax = max(dE1, dE2)
                    dEmin = min(dE1, dE2)
                    if energies[n+1] > energies[n-1]:
                        f1 = dEmax/(dEmax+dEmin+0.00000001)
                    else:
                        f1 = 1 - dEmax/(dEmax+dEmin+0.00000001)

                    print(' 3 way tangent ({}): f1:{:3.2}'.format(n, f1))

                    t1, _ = GSM.get_tangent(nodes[intic_n], nodes[newic_n])
                    t2, _ = GSM.get_tangent(nodes[newic_n], nodes[int2ic_n])
                    print(" done 3 way tangent")
                    ictan0 = f1*t1 + (1.-f1)*t2
            else:
                ictan0, _ = GSM.get_tangent(nodes[newic_n], nodes[intic_n])

            ictan[n] = ictan0/np.linalg.norm(ictan0)
            dqmaga[n] = np.linalg.norm(ictan0)

        return ictan, dqmaga

    @staticmethod
    def ic_reparam(nodes, energies, climbing=False, ic_reparam_steps=8, print_level=1, NUM_CORE=1, MAXRE=0.25):
        '''
        Reparameterizes the string using Delocalizedin internal coordinatesusing three-way tangents at the TS node
        Only pushes nodes outwards during reparameterization because otherwise too many things change.
            Be careful, however, if the path is allup or alldown then this can cause
        Parameters
        ----------
        nodes : list of molecule objects
        energies : list of energies in kcal/mol
        ic_reparam_steps : int max number of reparameterization steps
        print_level : int verbosity
        '''
        nifty.printcool("reparametrizing string nodes")

        nnodes = len(nodes)
        rpart = np.zeros(nnodes)
        for n in range(1, nnodes):
            rpart[n] = 1./(nnodes-1)
        deltadqs = np.zeros(nnodes)
        TSnode = np.argmax(energies)
        disprms = 100
        if ((TSnode == nnodes-1) or (TSnode == 0)) and climbing:
            raise RuntimeError(" TS node shouldn't be the first or last node")

        ideal_progress_gained = np.zeros(nnodes)
        if climbing:
            for n in range(1, TSnode):
                ideal_progress_gained[n] = 1./(TSnode)
            for n in range(TSnode+1, nnodes):
                ideal_progress_gained[n] = 1./(nnodes-TSnode-1)
            ideal_progress_gained[TSnode] = 0.
        else:
            for n in range(1, nnodes):
                ideal_progress_gained[n] = 1./(nnodes-1)

        for i in range(ic_reparam_steps):

            ictan, dqmaga = GSM.get_tangents(nodes)
            totaldqmag = np.sum(dqmaga)

            if climbing:
                progress = np.zeros(nnodes)
                progress_gained = np.zeros(nnodes)
                h1dqmag = np.sum(dqmaga[:TSnode+1])
                h2dqmag = np.sum(dqmaga[TSnode+1:nnodes])
                if print_level > 0:
                    print(" h1dqmag, h2dqmag: %3.2f %3.2f" % (h1dqmag, h2dqmag))
                progress_gained[:TSnode] = dqmaga[:TSnode]/h1dqmag
                progress_gained[TSnode+1:] = dqmaga[TSnode+1:]/h2dqmag
                progress[:TSnode] = np.cumsum(progress_gained[:TSnode])
                progress[TSnode:] = np.cumsum(progress_gained[TSnode:])
            else:
                progress = np.cumsum(dqmaga)/totaldqmag
                progress_gained = dqmaga/totaldqmag

            if i == 0:
                orig_dqmaga = copy(dqmaga)
                orig_progress_gained = copy(progress_gained)

            if climbing:
                difference = np.zeros(nnodes)
                for n in range(TSnode):
                    difference[n] = ideal_progress_gained[n] - progress_gained[n]
                    deltadqs[n] = difference[n]*h1dqmag
                for n in range(TSnode+1, nnodes):
                    difference[n] = ideal_progress_gained[n] - progress_gained[n]
                    deltadqs[n] = difference[n]*h2dqmag
            else:
                difference = ideal_progress_gained - progress_gained
                deltadqs = difference*totaldqmag

            if print_level > 1:
                print(" ideal progress gained per step", end=' ')
                for n in range(nnodes):
                    print(" step [{}]: {:1.3f}".format(n, ideal_progress_gained[n]), end=' ')
                print()
                print(" path progress                 ", end=' ')
                for n in range(nnodes):
                    print(" step [{}]: {:1.3f}".format(n, progress_gained[n]), end=' ')
                print()
                print(" difference                    ", end=' ')
                for n in range(nnodes):
                    print(" step [{}]: {:1.3f}".format(n, difference[n]), end=' ')
                print()
                print(" deltadqs                      ", end=' ')
                for n in range(nnodes):
                    print(" step [{}]: {:1.3f}".format(n, deltadqs[n]), end=' ')
                print()

            # disprms = np.linalg.norm(deltadqs)/np.sqrt(nnodes-1)
            disprms = np.linalg.norm(deltadqs)/np.sqrt(nnodes-1)
            print(" disprms: {:1.3}\n".format(disprms))

            if disprms < 0.02:
                break

            # Move nodes
            if climbing:
                deltadqs[TSnode-2] -= deltadqs[TSnode-1]
                deltadqs[nnodes-2] -= deltadqs[nnodes-1]
                for n in range(1, nnodes-1):
                    if abs(deltadqs[n]) > MAXRE:
                        deltadqs[n] = np.sign(deltadqs[n])*MAXRE
                for n in range(TSnode-1):
                    deltadqs[n+1] += deltadqs[n]
                for n in range(TSnode+1, nnodes-2):
                    deltadqs[n+1] += deltadqs[n]
                for n in range(nnodes):
                    if abs(deltadqs[n]) > MAXRE:
                        deltadqs[n] = np.sign(deltadqs[n])*MAXRE

                if NUM_CORE > 1:

                    # 5/14/2021 TS node fucks this up?!
                    tans = [ictan[n] if deltadqs[n] < 0 else ictan[n+1] for n in chain(range(1, TSnode), range(TSnode+1, nnodes-1))]  # + [ ictan[n] if deltadqs[n]<0 else ictan[n+1] for n in range(TSnode+1,nnodes-1)]
                    pool = mp.Pool(NUM_CORE)
                    Vecs = pool.map(worker, ((nodes[0].coord_obj, "build_dlc", node.xyz, tan) for node, tan in zip(nodes[1:TSnode] + nodes[TSnode+1:nnodes-1], tans)))
                    pool.close()
                    pool.join()
                    for n, node in enumerate(nodes[1:TSnode] + nodes[TSnode+1:nnodes-1]):
                        node.coord_basis = Vecs[n]

                    # move the positions
                    dqs = [deltadqs[n]*nodes[n].constraints[:, 0] for n in chain(range(1, TSnode), range(TSnode+1, nnodes-1))]
                    pool = mp.Pool(NUM_CORE)
                    newXyzs = pool.map(worker, ((node.coord_obj, "newCartesian", node.xyz, dq) for node, dq in zip(nodes[1:TSnode] + nodes[TSnode+1:nnodes-1], dqs)))
                    pool.close()
                    pool.join()
                    for n, node in enumerate(nodes[1:TSnode] + nodes[TSnode+1:nnodes-1]):
                        node.xyz = newXyzs[n]
                else:
                    for n in chain(range(1, TSnode), range(TSnode+1, nnodes-1)):
                        if deltadqs[n] < 0:
                            # print(f" Moving node {n} along tan[{n}] this much {deltadqs[n]}")
                            print(" Moving node {} along tan[{}] this much {}".format(n, n, deltadqs[n]))
                            nodes[n].update_coordinate_basis(ictan[n])
                            constraint = nodes[n].constraints[:, 0]
                            dq = deltadqs[n]*constraint
                            nodes[n].update_xyz(dq, verbose=(print_level > 1))
                        elif deltadqs[n] > 0:
                            print(" Moving node {} along tan[{}] this much {}".format(n, n+1, deltadqs[n]))
                            nodes[n].update_coordinate_basis(ictan[n+1])
                            constraint = nodes[n].constraints[:, 0]
                            dq = deltadqs[n]*constraint
                            nodes[n].update_xyz(dq, verbose=(print_level > 1))
            else:
                # e.g 11-2 = 9, deltadq[9] -= deltadqs[10]
                deltadqs[nnodes-2] -= deltadqs[nnodes-1]
                for n in range(1, nnodes-1):
                    if abs(deltadqs[n]) > MAXRE:
                        deltadqs[n] = np.sign(deltadqs[n])*MAXRE
                for n in range(1, nnodes-2):
                    deltadqs[n+1] += deltadqs[n]
                for n in range(1, nnodes-1):
                    if abs(deltadqs[n]) > MAXRE:
                        deltadqs[n] = np.sign(deltadqs[n])*MAXRE

                if NUM_CORE > 1:
                    # Update the coordinate basis
                    tans = [ictan[n] if deltadqs[n] < 0 else ictan[n+1] for n in range(1, nnodes-1)]
                    pool = mp.Pool(NUM_CORE)
                    Vecs = pool.map(worker, ((nodes[0].coord_obj, "build_dlc", node.xyz, tan) for node, tan in zip(nodes[1:nnodes-1], tans)))
                    pool.close()
                    pool.join()
                    for n, node in enumerate(nodes[1:nnodes-1]):
                        node.coord_basis = Vecs[n]
                    # move the positions
                    dqs = [deltadqs[n]*nodes[n].constraints[:, 0] for n in range(1, nnodes-1)]
                    pool = mp.Pool(NUM_CORE)
                    newXyzs = pool.map(worker, ((node.coord_obj, "newCartesian", node.xyz, dq) for node, dq in zip(nodes[1:nnodes-1], dqs)))
                    pool.close()
                    pool.join()
                    for n, node in enumerate(nodes[1:nnodes-1]):
                        node.xyz = newXyzs[n]
                else:
                    for n in range(1, nnodes-1):
                        if deltadqs[n] < 0:
                            # print(f" Moving node {n} along tan[{n}] this much {deltadqs[n]}")
                            print(" Moving node {} along tan[{}] this much {}".format(n, n, deltadqs[n]))
                            nodes[n].update_coordinate_basis(ictan[n])
                            constraint = nodes[n].constraints[:, 0]
                            dq = deltadqs[n]*constraint
                            nodes[n].update_xyz(dq, verbose=(print_level > 1))
                        elif deltadqs[n] > 0:
                            print(" Moving node {} along tan[{}] this much {}".format(n, n+1, deltadqs[n]))
                            nodes[n].update_coordinate_basis(ictan[n+1])
                            constraint = nodes[n].constraints[:, 0]
                            dq = deltadqs[n]*constraint
                            nodes[n].update_xyz(dq, verbose=(print_level > 1))

        if climbing:
            ictan, dqmaga = GSM.get_tangents(nodes)
            h1dqmag = np.sum(dqmaga[:TSnode+1])
            h2dqmag = np.sum(dqmaga[TSnode+1:nnodes])
            if print_level > 0:
                print(" h1dqmag, h2dqmag: %3.2f %3.2f" % (h1dqmag, h2dqmag))
            progress_gained[:TSnode] = dqmaga[:TSnode]/h1dqmag
            progress_gained[TSnode+1:] = dqmaga[TSnode+1:]/h2dqmag
            progress[:TSnode] = np.cumsum(progress_gained[:TSnode])
            progress[TSnode:] = np.cumsum(progress_gained[TSnode:])
        else:
            ictan, dqmaga = GSM.get_tangents(nodes)
            totaldqmag = np.sum(dqmaga)
            progress = np.cumsum(dqmaga)/totaldqmag
            progress_gained = dqmaga/totaldqmag
        print()
        if print_level > 0:
            print(" ideal progress gained per step", end=' ')
            for n in range(nnodes):
                print(" step [{}]: {:1.3f}".format(n, ideal_progress_gained[n]), end=' ')
            print()
            print(" original path progress        ", end=' ')
            for n in range(nnodes):
                print(" step [{}]: {:1.3f}".format(n, orig_progress_gained[n]), end=' ')
            print()
            print(" reparameterized path progress ", end=' ')
            for n in range(nnodes):
                print(" step [{}]: {:1.3f}".format(n, progress_gained[n]), end=' ')
            print()

        print(" spacings (begin ic_reparam, steps", end=' ')
        for n in range(nnodes):
            print(" {:1.2}".format(orig_dqmaga[n]), end=' ')
        print()
        print(" spacings (end ic_reparam, steps: {}/{}):".format(i+1, ic_reparam_steps), end=' ')
        for n in range(nnodes):
            print(" {:1.2}".format(dqmaga[n]), end=' ')
        print("\n  disprms: {:1.3}".format(disprms))

        return

    # TODO move to string utils or delete altogether
    #def get_current_rotation(self,frag,a1,a2):
    #    '''
    #    calculate current rotation for single-ended nodes
    #    '''
    #
    #    # Get the information on fragment to rotate
    #    sa,ea,sp,ep = self.nodes[0].coord_obj.Prims.prim_only_block_info[frag]
    #
    #    theta = 0.
    #    # Haven't added any nodes yet
    #    if self.nR==1:
    #        return theta

    #    for n in range(1,self.nR):
    #        xyz_frag = self.nodes[n].xyz[sa:ea].copy()
    #        axis = self.nodes[n].xyz[a2] - self.nodes[n].xyz[a1]
    #        axis /= np.linalg.norm(axis)
    #
    #        # only want the fragment of interest
    #        reference_xyz = self.nodes[n-1].xyz.copy()

    #        # Turn off
    #        ref_axis = reference_xyz[a2] - reference_xyz[a1]
    #        ref_axis /= np.linalg.norm(ref_axis)

    #        # ALIGN previous and current node to get rotation around axis of rotation
    #        #print(' Rotating reference axis to current axis')
    #        I = np.eye(3)
    #        v = np.cross(ref_axis,axis)
    #        if v.all()==0.:
    #            print('Rotation is identity')
    #            R=I
    #        else:
    #            vx = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
    #            c = np.dot(ref_axis,axis)
    #            s = np.linalg.norm(v)
    #            R = I + vx + np.dot(vx,vx) * (1. - c)/(s**2)
    #        new_ref_axis = np.dot(ref_axis,R.T)
    #        #print(' overlap of ref-axis and axis (should be 1.) %1.2f' % np.dot(new_ref_axis,axis))
    #        new_ref_xyz = np.dot(reference_xyz,R.T)

    #
    #        # Calculate dtheta
    #        ca = self.nodes[n].primitive_internal_coordinates[sp+3]
    #        cb = self.nodes[n].primitive_internal_coordinates[sp+4]
    #        cc = self.nodes[n].primitive_internal_coordinates[sp+5]
    #        dv12_a = ca.calcDiff(self.nodes[n].xyz,new_ref_xyz)
    #        dv12_b = cb.calcDiff(self.nodes[n].xyz,new_ref_xyz)
    #        dv12_c = cc.calcDiff(self.nodes[n].xyz,new_ref_xyz)
    #        dv12 = np.array([dv12_a,dv12_b,dv12_c])
    #        #print(dv12)
    #        dtheta = np.linalg.norm(dv12)  #?
    #
    #        dtheta = dtheta + np.pi % (2*np.pi) - np.pi
    #        theta += dtheta

    #    theta = theta/ca.w
    #    angle = theta * 180./np.pi
    #    print(angle)

    #    return theta

    @staticmethod
    def calc_optimization_metrics(nodes):
        '''
        '''

        nnodes = len(nodes)
        rn3m6 = np.sqrt(3*nodes[0].natoms-6)
        totalgrad = 0.0
        gradrms = 0.0
        sum_gradrms = 0.0
        for i, ico in enumerate(nodes[1:nnodes-1]):
            if ico != None:
                print(" node: {:02d} gradrms: {:.6f}".format(i, float(ico.gradrms)), end='')
                if i % 5 == 0:
                    print()
                totalgrad += ico.gradrms*rn3m6
                gradrms += ico.gradrms*ico.gradrms
                sum_gradrms += ico.gradrms
        print('')
        # TODO wrong for growth
        gradrms = np.sqrt(gradrms/(nnodes-2))
        return totalgrad, gradrms, sum_gradrms