# pylint: disable=dangerous-default-value

"""
  Various routines for dealing with generating input and processing output
  for abstracted records of experimental tests.

  These routines do things like:
    * Convert simplified descriptions of tests into strain/strain time
      histories that the models can use
    * Generate random experiments and the corresponding inputs

  The actual deterministic and statistical fitting routines expect data
  in a common format.  There are a few conditions on the input test data:

  1. You have to sample all tests with the same number of time points (ntime).
  2. You have to provide the full, correct history of the input conditions.
     For strain controlled experiments this is the full history of test
     time, strain, and temperature.  For stress control it is the full
     history of time, stress, and temperature.
  3. You do not have to provide the full *results* history.  That is,
     the fit routines can accommodate some types of abstracted test data,
     like maximum stress as a function of cycles for strain-controlled cyclic
     tests.  The convert_results function can then convert a full (simulated)
     results history to the right abstracted format for comparison to the
     test data.

  The fitting routines rely on five tensors with the following format:

  * data :code:`(3, ntime, nexp)`: the input experimental conditions.
    The first
    index contains the experimental time, temperature, and strain history
    for strain controlled tests and the experimental time, temperature
    and stress history for stress controlled tests
  * results :code:`(ntime, nexp)`: the (potentially abstracted) experimental
    results.
    For example, this could be the stress history for a tensile test or
    the strain history for a creep test.  The specific requirements for
    certain tests types are described below.
  * cycles :code:`(ntime, nexp)`: the cycle counts for cyclic  tests or an
    integer
    used to indicate features of the experiment for other test types.
    Specific requirements are given below.
  * types :code:`(nexp,)`: an integer giving the test type (see list below)
  * control :code:`(nexp,)`: 0 for strain control, 1 for stress control

  The current test types are:

  0. "tension" -- uniaxial, monotonic tension or compression test
  1. "relaxation" -- stress relaxation test
  2. "strain_cyclic" -- strain controlled cyclic (i.e. creep or creep-fatigue)
  3. "creep" -- a creep test
  4. "stress_cyclic" -- a stress controlled cyclic test
  5. "abstract_tensile" -- tension tests where only the yield strength
     and ultimate tensile strength are known
  6. "direct_data" -- the full results history is known and provided

  The :func:`pyoptmat.experiments.load_results`
  function provides a way to load data for this type
  from an xarray format.  The same tensors are stored in the xarray
  structure, however the descriptive string names for the tests are used
  and "strain" or "stress" is used instead of 0 and 1 to indicate the
  control type.  Additional metadata can be stored in the xarray format with
  source information, heat ID, etc.
"""

import numpy as np
import numpy.random as ra

import torch

def load_results(xdata, device=torch.device("cpu")):
    """
    Load experimental data from xarray into torch tensors

    Args:
      xdata (xarray.DataArray): xarray data structure

    Keyword Args:
      device (torch.device):   the device to dump the resulting arrays on

    Returns:
      tuple: see below

    This function returns a tuple of five tensors

      data :code:`(3, ntime, nexperiment)`

        The initial index are the experimental (times, temperatures, idata)
        For strain controlled tests idata is strain
        For stress controlled tests idata is stress

      results :code:`(ntime, nexperiment)`

        For strain controlled tests this is stress
        For stress controlled tests this is strain

      cycles :code:`(ntime, nexperiment)`

        Cycle count for all tests

      types :code:`(nexperiment,)`

        Experiment types, converted to integers per the dicts above

      control :code:`(nexperiment,)`

        Maps the string control type ("strain" or "stress") to an integer
        using the dict above
    """
    time = torch.tensor(xdata["time"].values, device=device)
    temp = torch.tensor(xdata["temperature"].values, device=device)
    strain = torch.tensor(xdata["strain"].values, device=device)
    stress = torch.tensor(xdata["stress"].values, device=device)

    cycle = torch.tensor(xdata["cycle"].values, device=device)
    types = torch.tensor([exp_map[t] for t in xdata["type"].values], device=device)
    control = torch.tensor(
        [control_map[t] for t in xdata["control"].values], device=device
    )

    data = torch.empty((4,) + time.shape, device=device)

    data[0] = time
    data[1] = temp
    data[2, :, control == 0] = strain[:, control == 0]
    data[2, :, control == 1] = stress[:, control == 1]

    results = torch.empty_like(time)
    results[:, control == 0] = stress[:, control == 0]
    results[:, control == 1] = strain[:, control == 1]

    return data, results, cycle, types, control

def convert_results(results, cycles, types):
    """
    Process a raw results vector to our common format based on test type

    Args:
      results (torch.tensor):   raw results data :code:`(ntime, nexperiment)`
      cycles (torch.tensor):    cycle counts :code:`(ntime, nexperiment)`
      types (torch.tensor):     test types :code:`(nexperiment,)`

    Returns:
      torch.tensor: tensor of processed results
    """
    processed = torch.empty_like(results)

    for i, func in exp_fns_num.items():
        current = types == i
        if torch.sum(current) > 0:
            processed[:, current] = func(cycles[:, current], results[:, current])

    return processed

def format_direct_data(cycles, predictions):
    """
    Format direct stress/strain results

    For this test type cycles just equals 0 for all time steps

    This function does nothing, it just returns :code:`predictions`

    Args:
      cycles (torch.tensor):        cycle count
      predictions (torch.tensor):   input to format

    Returns:
      torch.tensor:                 tensor of processed results

    """
    return predictions

def format_abstract_tensile(cycles, predictions):
    """
    Format abstracted tensile test data, where only the yield strength
    and ultimate tensile strength are available.

    This method relies on the input cycles being:
    * 0: in the elastic regime
    * 1: for the first 1/2 of the remaining time points outside of the elastic regime
    * 2: for the second 1/2 of the remaining time points.

    The inputs are the full stress history.  This function:

    * Cycle 0: overwrite with zeros
    * Cycle 1: overwrite with the first value in cycle 1
    * Cycle 2: overwrite with the maximum stress

    Args:
      cycles (torch.tensor):        cycle count
      predictions (torch.tensor):   input to format

    Returns:
      torch.tensor:                 tensor of processed results
    """
    result = torch.zeros_like(predictions)

    # cycles == 0 -> 0
    result += torch.where(cycles == 0, 0.0, 0.0)

    # cycles == 1 -> first value in that region
    _, index = torch.max(cycles == 1, 0)
    result += torch.where(cycles == 1, predictions.gather(0, index.view(1, -1))[0], 0.0)

    # cycles == 2 -> max value overall
    mvalues, _ = torch.max(predictions, 0)
    result += torch.where(cycles == 2, mvalues, 0.0)

    return result

def format_tensile(cycles, predictions):
    """
    Format tension test data to our "post-processed" form for comparison

    Input data are stresses for this test type

    Cycles are all 0

    This function doesn't do anything, just returns :code:`predictions`

    Args:
      cycles (torch.tensor):        cycle count/listing
      predictions (torch.tensor):   input data

    Returns:
      torch.tensor:                 processed results
    """
    return predictions

def format_relaxation(cycles, predictions):
    """
    Format stress relaxation test data to our "post-processed" form for
    comparison. This works for both creep and stress relaxation tests.

    Input data are stresses for stress relaxation and strains for creep

    Cycle 0 indicates the loading part of the test, cycle 1 is the hold

    Zero out the loading results, replace the relaxation results with the
    normalized (subtract t=0) curve

    Args:
      cycles (torch.tensor):        cycle count/listing
      predictions (torch.tensor):   input data

    Returns:
      torch.tensor:                 processed results
    """
    result = torch.zeros_like(predictions)
    rcurve = cycles[:, 0] == 1  # This is right, but dangerous in the future

    curve = predictions[rcurve]
    curve = curve - curve[0]

    result[rcurve] = curve

    return result

def format_cyclic(cycles, predictions):
    """
    Format a generic cyclic test -- works for both stress and strain control

    Input data are stresses for strain control and strains for stress control.

    We format this as a "block" -- the values for each cycle are replaced
    by the maximum value within the cycle

    Args:
      cycles (torch.tensor):        cycle count/listing
      predictions (torch.tensor):   input data

    Returns:
      torch.tensor:                 processed results
    """
    # If this is slow we can probably remove the for loop
    result = torch.zeros_like(predictions)
    uc = cycles[:, 0]  # Correct but dangerous for future expansion
    for i in range(uc[-1] + 1):
        curr = uc == i
        vals, _ = torch.max(predictions[curr], axis=0)
        result[curr] = vals

    return result

def make_tension_tests(rates, temperatures, elimits, nsteps):
    """
    Produce tension test (time,strain,temperature) history blocks
    given tensor inputs for the strain rates, temperatures, and
    maximum strain of each test

    Args:
      rates (torch.tensor):         1D tensor giving the strain rate of each test
      temperaturess (torch.tensor): 1D tensor giving the constant temperature of each test
      elimits (torch.tensor):       1D tensor giving the maximum strain of each test
      nsteps (torch.tensor):        integer number of steps

    Returns:
      tuple:                        tuple of
                                    :code:`(times, strains, temperatures, cycles)`
    """
    nbatch = temperatures.shape[0]
    times = torch.zeros(nsteps, nbatch)
    strains = torch.zeros_like(times)
    temps = torch.zeros_like(strains)

    for i in range(nbatch):
        times[:, i] = torch.linspace(0, elimits[i] / rates[i], nsteps)
        strains[:, i] = torch.linspace(0, elimits[i], nsteps)
        temps[:, i] = temperatures[i]

    return times, strains, temps, torch.zeros_like(times, dtype=int)

def make_creep_tests(
    stress, temperature, rate, hold_times, nsteps_load, nsteps_hold, logspace=False
):
    """
    Produce creep test input (time,stress,temperature) given tensor
    inputs for the target stress, target temperature, loading rate

    Args:
      stress (torch.tensor):        1D tensor of target stresses
      temperature (torch.tensor):   1D tensor of target temperature
      rate (torch.tensor):          1D tensor of target rates
      hold_times (torch.tensor):    1D tensor of hold times
      nsteps_load (torch.tensor):   number of time steps to load up the sample
      nsteps_hold (torch.tensor):   number of time steps to hold the sample

    Keyword Args:
      logspace (bool):              log-space time increments during holds

    Returns:
      tuple:                        tuple of
                                    :code:`(times, strains, temperatures, cycles)`
    """
    nbatch = stress.shape[0]
    nsteps = nsteps_load + nsteps_hold

    stresses = torch.zeros(nsteps, nbatch)
    times = torch.zeros_like(stresses)
    temperatures = torch.zeros_like(stresses)

    for i, (s, t, lr, T) in enumerate(zip(stress, hold_times, rate, temperature)):
        stresses[:nsteps_load, i] = torch.linspace(0, s, nsteps_load)
        stresses[nsteps_load:, i] = s

        times[:nsteps_load, i] = torch.linspace(0, s / lr, nsteps_load)
        temperatures[:, i] = T
        if logspace:
            times[nsteps_load:, i] = torch.logspace(
                torch.log10(times[nsteps_load - 1, i]), torch.log10(t), nsteps_hold + 1
            )[1:]
        else:
            times[nsteps_load:, i] = torch.linspace(
                times[nsteps_load - 1, i], t, nsteps_hold + 1
            )[1:]

    cycles = torch.ones_like(times, dtype=int)
    cycles[:nsteps_load] = 0

    return times, stresses, temperatures, cycles

def generate_random_tension(strain_rate=[1.0e-6, 1.0e-2], max_strain=0.2):
    """
    Generate a random tension test condition in the provided ranges

    Keyword Args:
      strain_rate (list): Range of strain rates
      max_strain (float): Maximum strain to simulate

    Returns:
      dict:               dictionary with :code:`"max_strain"`
                          and :code:`"strain_rate"`
    """
    return {
        "max_strain": max_strain,
        "strain_rate": 10.0
        ** ra.uniform(np.log10(strain_rate[0]), np.log10(strain_rate[1])),
    }

def sample_tension(test, nsteps=50):
    """
    Generate the times and strains for a tensile test

    Args:
      test (dict):    Dictionary defining the test case

    Keyword Args:
      nsteps (int):   Number of steps to sample

    Returns:
      tuple:          tuple of :code:`(times, strains)`
    """
    tmax = test["max_strain"] / test["strain_rate"]
    times = np.linspace(0, tmax, nsteps)
    strains = times * test["strain_rate"]

    return times, strains

def generate_random_cycle(
    max_strain=[0, 0.02],
    R=[-1, 1],
    strain_rate=[1.0e-3, 1.0e-5],
    tension_hold=[0, 1 * 3600.0],
    compression_hold=[0, 600],
):
    """
    Generate a random cycle in the provided ranges

    Keyword Args:
      max_strain (list):        range of the maximum strains
      R (list):                 range of R ratios
      strain_rate (list):       range of loading strain rates
      tension_hold (list):      range of tension hold times
      compression_hold (list):  range of compression hold times

    Returns:
      dict:                     dictionary describing cycle, described below

    * :code:`"max_strain"` -- maximum strain value
    * :code:`"R"` -- R ratio :math:`\\frac{max}{min}`
    * :code:`"strain_rate"` -- strain rate during load/unload
    * :code:`"tension_hold"` -- hold on the tension end of the cycle
    * :code:`"compression_hold"` -- hold on the compressive end of the cycle
    """
    return {
        "max_strain": ra.uniform(*max_strain),
        "R": ra.uniform(*R),
        "strain_rate": 10.0
        ** ra.uniform(np.log10(strain_rate[0]), np.log10(strain_rate[1])),
        "tension_hold": ra.uniform(*tension_hold),
        "compression_hold": ra.uniform(*compression_hold),
    }

def sample_cycle_normalized_times(cycle, N, nload=10, nhold=10):
    # pylint: disable=too-many-locals
    """
    Sample a cyclic test at a normalized series of times

    Take a random cycle dictionary and expand into discrete
    times, strains samples where times are the actual, physical
    times, given over the fixed phases

      * :math:`0 \\rightarrow  t_{phase}` -- tension load
      * :math:`t_{phase} \\rightarrow 2 t_{phase}` -- tension hold
      * :math:`2 t_{phase} \\rightarrow 3 t_{phase}` --   unload
      * :math:`3 t_{phase} \\rightarrow 4 t_{phase}` -- compression load
      * :math:`4 t_{phase} \\rightarrow 5 t_{phase}` -- compression hold
      * :math:`5 t_{phase} \\rightarrow 6 t_{phase}` -- unload

    This pattern repeats for N cycles

    Args:
      cycle (dict): dictionary defining the load cycle
      N (int):      number of repeats to include in the history

    Keyword Args:
      nload (int):  number of steps to use for the load time
      nhold (int):  number of steps to use for the hold time
    """
    emax = cycle["max_strain"]
    emin = cycle["R"] * cycle["max_strain"]
    erate = cycle["strain_rate"]

    # Segments:
    t1 = np.abs(emax) / erate
    t2 = cycle["tension_hold"]
    t3 = np.abs(emax - emin) / erate
    t4 = cycle["compression_hold"]
    t5 = np.abs(emin) / erate
    divisions = [t1, t2, t3, t4, t5]
    timesteps = [nload, nhold, 2 * nload, nhold, nload]
    cdivisions = np.cumsum(divisions)
    period = cdivisions[-1]

    Ntotal = nload * 4 + nhold * 2

    times = np.zeros((1 + Ntotal * N,))
    cycles = np.zeros(times.shape, dtype=int)

    n = 1
    tc = 0
    for k in range(N):
        for ti, ni in zip(divisions, timesteps):
            times[n : n + ni] = np.linspace(tc, tc + ti, ni + 1)[1:]
            cycles[n : n + ni] = k
            n += ni
            tc += ti

    tp = times % period
    strains = np.piecewise(
        tp,
        [
            np.logical_and(tp >= 0, tp < cdivisions[0]),
            np.logical_and(tp >= cdivisions[0], tp < cdivisions[1]),
            np.logical_and(tp >= cdivisions[1], tp < cdivisions[2]),
            np.logical_and(tp >= cdivisions[2], tp < cdivisions[3]),
            np.logical_and(tp >= cdivisions[3], tp < cdivisions[4]),
        ],
        [
            lambda tt: tt / t1 * emax,
            lambda tt: tt * 0 + emax,
            lambda tt: emax - (tt - cdivisions[1]) / t3 * (emax - emin),
            lambda tt: tt * 0 + emin,
            lambda tt: emin - (tt - cdivisions[3]) / t5 * emin,
        ],
    )

    return times, strains, cycles

# Numerical codes for each test type
exp_map = {
    "tensile": 0,
    "relaxation": 1,
    "strain_cyclic": 2,
    "creep": 3,
    "stress_cyclic": 4,
    "abstract_tensile": 5,
    "direct_data": 6,
}
# Function to use to process each test type
exp_fns = {
    "tensile": format_tensile,
    "relaxation": format_relaxation,
    "strain_cyclic": format_cyclic,
    "creep": format_relaxation,
    "stress_cyclic": format_cyclic,
    "abstract_tensile": format_abstract_tensile,
    "direct_data": format_direct_data,
}
# Map to numbers instead
exp_fns_num = {exp_map[k]: v for k, v in exp_fns.items()}

# Map control type to number
control_map = {"strain": 0, "stress": 1}


"""Main module."""

from scipy.interpolate import interp1d
from pathlib import Path
import pandas as pd
import numpy as np
from typing import Union, Any, Sequence
import json

class QINCM:
    # Quick Inland Navigation Cost Model

    # Under keel clearance
    ukc = 0.20

    def __init__(self,
                 route_depth_costs_file: Union[str, Path] = None,
                 knelpunt_discharge_depth_file: Union[str, Path] = None,
                 reference: str = None,
                 ):
        """
        Initialise

        :param route_depth_costs_file:
        :param knelpunt_discharge_depth_file:
        :param reference_point_mode: set True if the discharge in the earlier point links to a single discharge
        """

        # Initialise model
        self._read_routes_depth_costs(route_depth_costs_file)
        self._read_knelpunt_discharge_depth(knelpunt_discharge_depth_file, reference=reference)

    def _compute_knelpunt_depth(self, discharges):
        depths = {}
        for k in self.knelpunt_names:
            depths[k] = self.knelpunt_discharge_depth[k](discharges[k])
        depths = pd.DataFrame(data=depths, index=discharges.index)

        # if isinstance(discharges, pd.DataFrame):
        #     depths = {}
        #     for k in self.knelpunt_names:
        #         depths[k] = self.knelpunt_discharge_depth[k](discharges[k])
        #     depths = pd.DataFrame(data=depths, index=discharges.index)
        # else:
        #     depths = {}
        #     for k in self.knelpunt_names:
        #         depths[k] = self.knelpunt_discharge_depth[k](discharges)
        #     depths = pd.DataFrame(data=depths, index=discharges)

        return depths

    def _costs_at_routes_at_discharge(self, discharges: pd.DataFrame) -> pd.DataFrame:
        """
        Compute costs per route per discharge

        param discharges: list of unique discharges. dim1: unique discharges, dim2: knelpunten

        returns: DataFrame (index: routes, columns: discharges)
        """

        depths = self._compute_knelpunt_depth(discharges)

        # For each r (=FrozenList of knelpunten)
        r_costs = {}
        for r in self.routes:

            # Get the limiting depth on the route
            r_depth = depths.reindex(r, axis=1).min(axis=1).fillna(999999)

            # Depth to draught/draft
            r_draughts = r_depth - self.ukc

            # Get costs
            r_costs[r] = self.routes_depth_costs[r](r_draughts)

        costs = pd.DataFrame(r_costs, index=discharges.index)

        return costs

    def _compute_local_discharge(self, discharges):
        if np.ndim(discharges) == 1:
            # Only reference discharge is given, compute local discharge from Q-Q-relation
            Q_ref = discharges

            if isinstance(discharges, pd.Series):
                Q_local = pd.DataFrame(index=discharges.index)
            else:
                Q_local = pd.DataFrame(index=Q_ref)  # The index is only for convenience.

            for k, QQ in self.knelpunt_discharge_distribution.items():
                Q_local[k] = QQ(Q_ref)
        else:
            k_names = self.knelpunt_names

            if isinstance(discharges, pd.DataFrame):
                Q_local = pd.DataFrame(data=discharges, columns=k_names, index=discharges.index)
            else:
                Q_local = pd.DataFrame(data=discharges, columns=k_names)

                Q_local.index = Q_local[self.knelpunt_reference]  # The index is only for convenience.
        return Q_local

    def costs_per_discharge(self, discharges: Union) -> pd.DataFrame:
        """
        Compute total costs per discharge

        param discharges: list of unique discharges. ALso supports timeseries

        returns: Series (index=discharges)
        """
        Q_local = self._compute_local_discharge(discharges)

        costs = self._costs_at_routes_at_discharge(Q_local)
        return costs

    def costs_for_scenario(self, discharges, occurance=None, delta: bool = True):
        """
        Compute total costs in scenario

        param discharge: list of unique discharges
        param occurance: float, or list with for each discharge the number of days. If none, it assumes every discharges occured one day
        """

        # TODO: Discharge may also be a pandas
        if occurance is not None:
            # Validate input
            # if isinstance(occurance, float):
            #     occurance = np.ones(np.shape(discharges)) * occurance
            # assert len(discharges) == len(occurance), 'Input should have same length'

            # Compute total costs
            costs = self.costs_per_discharge(discharges)

            # Compute delta costs
            if delta:
                costs_no_problems = self.costs_per_discharge([99999999])  #TODO: This way is not good for depths that are included as constants...
                costs = costs.subtract(costs_no_problems.values, axis=1)

            costs_occurance = costs.multiply(occurance, axis=0)

        else:
            costs = self.costs_per_discharge(discharges)

            if delta:
                costs_no_problems = self.costs_per_discharge([99999999])
                costs = costs.subtract(costs_no_problems.values, axis=1)

            costs_occurance = costs

        total_costs_per_route = costs_occurance.sum(axis=0)
        # total_costs_per_discharge = costs_occurance.sum(axis=1) # Or per day if that's the index
        # total_costs = total_costs_per_route.sum()

        return total_costs_per_route

    def _read_routes_depth_costs(self, routes_depth_costs_file: Union[str, Path]):
        """
        # Set for each route (combination of knelpunten) the function of [draught]-[response].

        Read json file with following format:

        {
          [list_of_knelpunten_for_route1]: {depth: costs_per_route_at_depth_per_day},
          [list_of_knelpunten_for_route2]: {depth: costs_per_route_at_depth_per_day},
          ...
        }

        e.g. (truncated notation)

        {
          [A, B]: {0.0: 1000, 1.0: 2000, 2.0: 1500},
          [B, C]: {0.0: 2000, 1.0: 3000, 2.0: 1800}
        }

        """
        # Read output
        routes_depth_costs = pd.read_json(routes_depth_costs_file, convert_dates=False, convert_axes=False)
        routes_depth_costs.columns = [frozenset(i[1:-1].split(', ')) for i in routes_depth_costs.columns]
        routes_depth_costs = routes_depth_costs.rename({frozenset({''}): frozenset()}, axis=1)
        routes_depth_costs.index = [float(c) for c in routes_depth_costs.index]

        self.routes = routes_depth_costs.columns

        # Convert into interpolation function
        depth_costs_functions = {}
        for r in self.routes:
            depth_costs = routes_depth_costs.xs(r, axis=1)
            # Test this also: depth_costs = routes_depth_costs.xs(r)

            depth_costs_function = interp1d(
                x=depth_costs.index,
                y=depth_costs.values,
                kind='linear',
                bounds_error=False,
                fill_value=tuple(depth_costs.values[[0, -1]]),
            )
            depth_costs_functions[r] = depth_costs_function

        self.routes_depth_costs = depth_costs_functions

    def _read_knelpunt_discharge_depth(self, knelpunt_discharge_depth_file: Union[str, Path], reference=None):
        """
        Read json file with for each knelpunt the discharge-depth relation. The discharges for all knelpunten (probably)
        need to be identical

        {
               knelpunt1: {discharge: depth}
               knelpunt2: {discharge: depth}
        }

        e.g.
        {
          'k1': {500: 2.5, 1000: 3, 2000: 4},
          'k2': {200, 5, 800: 6, 1500: 7}
        }

        :param knelpunt_discharge_depth_file: path to json file

        """
        if not isinstance(knelpunt_discharge_depth_file, dict):
            with open(knelpunt_discharge_depth_file) as fin:
                discharge_depth = json.load(fin)
        else:
            discharge_depth = knelpunt_discharge_depth_file

        if reference is None:
            # If not set, take the first
            self.knelpunt_reference = list(discharge_depth.keys())[0]
        else:
            self.knelpunt_reference = reference

        self.knelpunt_names = list(discharge_depth.keys())

        self.knelpunt_discharge_depth = {}
        knelpunt_discharge = {}
        for k, QD in discharge_depth.items():
            Q, D = zip(*QD.items())
            Q = [float(q) for q in Q]
            D = [float(d) for d in D]

            # Make this an interpolation function, incl. extrapolation
            discharge_depth_function = interp1d(
                x=Q,
                y=D,
                kind='linear',
                bounds_error=False,
                fill_value='extrapolate',
            )
            self.knelpunt_discharge_depth[k] = discharge_depth_function
            knelpunt_discharge[k] = Q

        # Make lookup function of discharge at reference, to local discharge
        Q_ref = knelpunt_discharge[self.knelpunt_reference]

        self.knelpunt_discharge_distribution = {}
        for k in discharge_depth:
            Q = knelpunt_discharge[k]

            # Make this an interpolation function, extrapolation=constant
            Qref_Q = interp1d(
                x=Q_ref,
                y=Q,
                kind='linear',
                bounds_error=False,
                fill_value='extrapolate',
            )

            self.knelpunt_discharge_distribution[k] = Qref_Q

    def stats_knelpunten(self, Qmin=500, Qmax=2000):
        """
        for each discharge determine the number of trips that is influenced by the knelpunt (alltrips)
        for each knelpunt show only the trips that are limited by that specific point

        return alltrips, mintrips
        """

        alltrips = {}
        mintrips = {}
        mintrips_increase = {}

        routes = self.routes_depth_costs.keys()

        discharge_series = np.linspace(Qmin, Qmax, 100)

        # Get discharge per knelpunt
        Q_local = self._compute_local_discharge(discharge_series)

        # Get depth per knelpunt and convert to draught by using ukc
        QH = self._compute_knelpunt_depth(Q_local) - self.ukc

        for r in routes:
            # Depth on route
            r_QH = QH[r]

            # Get total costs on this route
            depth_for_route = r_QH.min(axis=1).fillna(999)
            costs_for_route = pd.Series(
                data = self.routes_depth_costs[r](depth_for_route),
                index = discharge_series
            )
            costs_for_route_increase = costs_for_route - costs_for_route.iloc[-1]

            for k in r:

                depth = r_QH[k].values

                # Get number of affected trips (per waterdepth)
                costs_for_all_passing_trips = self.routes_depth_costs[r](depth)
                alltrips[(r, k)] = pd.Series(
                    data=costs_for_all_passing_trips,
                    index=discharge_series
                )

                # Only select those levels where k is the minimum depth on the route
                k_minimal = r_QH.idxmin(axis=1) == k
                mintrips[(r, k)] = costs_for_route.multiply(k_minimal, axis=0)
                mintrips_increase[(r, k)] = costs_for_route_increase.multiply(k_minimal, axis=0)

        alltrips = pd.concat(alltrips)
        mintrips = pd.concat(mintrips)
        mintrips_increase = pd.concat(mintrips_increase)

        alltrips_sum = alltrips.unstack(level=1).sum(axis=0, level=1)
        mintrips_sum = mintrips.unstack(level=1).sum(axis=0, level=1)
        mintrips_increase_sum = mintrips_increase.unstack(level=1).sum(axis=0, level=1)
        return alltrips_sum, mintrips_sum, mintrips_increase_sum

if __name__ == '__main__':

    inputdir = Path('../data')

    M = QINCM(inputdir / 'route_depth_costs.json',
              inputdir / 'knelpunt_discharge_waterdepth.json',
              reference='WA_Nijmegen'
              )

    discharges = np.linspace(500, 3000, 26)

    # print('Costs per route per discharge')
    # a1 = M.costs_per_route_per_discharge(discharges)
    # print(a1.head())

    print('Costs per discharge')
    a2 = M.costs_per_discharge(discharges)
    print(a2.head())

    # Create a random occurance so the total sum is a full year
    occurance = np.random.rand(*discharges.shape)
    occurance *= 365 / occurance.sum()

    print('Costs per year')
    a3 = M.costs_for_scenario(discharges, occurance)
    print(a3)


# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

from __future__ import division, unicode_literals

"""
This module contains the class describing the coordination geometries that can exist in a given structure. These
"model" coordination geometries are described in the following articles :
 - Pure Appl. Chem., Vol. 79, No. 10, pp. 1779--1799, 2007.
 - Acta Cryst. A, Vol. 46, No. 1, pp. 1--11, 1990.
The module also contains descriptors of part of these geometries (plane of separation, ...) that are used in the
identification algorithms.
"""

__author__ = "David Waroquiers"
__copyright__ = "Copyright 2012, The Materials Project"
__credits__ = "Geoffroy Hautier"
__version__ = "2.0"
__maintainer__ = "David Waroquiers"
__email__ = "david.waroquiers@gmail.com"
__date__ = "Feb 20, 2016"

import numpy as np
from scipy.misc import factorial
import itertools
import abc
from monty.json import MSONable, MontyDecoder
import json
import os
from six import with_metaclass

module_dir = os.path.dirname(os.path.abspath(__file__))

UNKNOWN_ENVIRONMENT_SYMBOL = 'UNKNOWN'
UNCLEAR_ENVIRONMENT_SYMBOL = 'UNCLEAR'
EXPLICIT_PERMUTATIONS = 'EXPLICIT_PERMUTATIONS'
SEPARATION_PLANE = 'SEPARATION_PLANE'

class AbstractChemenvAlgorithm(with_metaclass(abc.ABCMeta, MSONable)):
    """
    Class used to define a Chemenv strategy for the neighbors and coordination environment to be applied to a
    StructureEnvironments object
    """

    def __init__(self, algorithm_type):
        self._algorithm_type = algorithm_type

    @abc.abstractmethod
    def as_dict(self):
        """
        A JSON serializable dict representation of the algorithm
        """
        pass

    @property
    def algorithm_type(self):
        return self._algorithm_type

    @abc.abstractmethod
    def __str__(self):
        return

class ExplicitPermutationsAlgorithm(AbstractChemenvAlgorithm):
    def __init__(self, permutations):
        """
            Initializes a separation plane for a given perfect coordination geometry
        """
        super(ExplicitPermutationsAlgorithm, self).__init__(
            algorithm_type=EXPLICIT_PERMUTATIONS)
        self._permutations = permutations

    def __str__(self):
        return self.algorithm_type

    @property
    def permutations(self):
        return self._permutations

    @property
    def as_dict(self):
        return {"@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "permutations": self._permutations}

    @classmethod
    def from_dict(cls, dd):
        return cls(dd['permutations'])

class SeparationPlane(AbstractChemenvAlgorithm):
    def __init__(self, plane_points, mirror_plane=False, ordered_plane=False,
                 point_groups=None,
                 ordered_point_groups=None,  # include_inverted_plane=False,
                 point_groups_permutations=None,
                 # do_inverse_pt_gp_permutations=False, plane_type='MIRROR',
                 explicit_permutations=None, minimum_number_of_points=None,
                 explicit_optimized_permutations=None,
                 multiplicity=None,
                 other_plane_points=None):  # , plane_safe_permutations=False):
        """
            Initializes a separation plane for a given perfect coordination geometry

            :param mirror_plane: True if the separation plane is a mirror plane, in which case there is a correspondence
            of the points in each point_group (can reduce the number of permutations)
            :param ordered_plane : True if the order of the points in the plane can be taken into account to reduce the
            number of permutations
            :param plane_points: Indices of the points that are in the plane in the perfect structure (and should be
            found in the defective one as well)
            :param point_groups: The two groups of points separated by the plane
            :param plane_type: can be "MIRROR", if the plane is a mirror plane going through the central site,
             'BASAL_THROUGH_CENTER', if the plane is a basal plane (no point on the "left" side) going through the central
             site, 'BASAL', if the is a basal plane not going through the central site, 'UNEQUILIBRATED_THROUGH_CENTER', if
             the plane cuts the geometry in two groups of points with different numbers of points on each side, and is going
             through the centre, 'UNEQUILIBRATED', if the plane cuts the geometry in two groups of points with different
             numbers of points on each side, and is not going through the centre, 'EQUILIBRATED_THROUGH_CENTER', if the
             plane cuts the geometry in two groups of points of the same size, is going through the centre but is not a
             mirror plane, 'EQUILIBRATED', if the plane cuts the geometry in two groups of points of the same size, is not
             going through the centre but is not a mirror plane.
            """
        super(SeparationPlane, self).__init__(algorithm_type=SEPARATION_PLANE)
        self.mirror_plane = mirror_plane
        self.plane_points = plane_points
        self.point_groups = point_groups
        if len(point_groups[0]) > len(point_groups[1]):
            raise RuntimeError(
                "The number of points in the first group should be\n"
                "less than or equal to the number of points in the second group")
        self._hash = 10000 * len(plane_points) + 100 * len(
            point_groups[0]) + len(point_groups[1])
        self.ordered_plane = ordered_plane
        self.ordered_point_groups = [False,
                                     False] if ordered_point_groups is None else ordered_point_groups
        self._ordered_indices = list(point_groups[0])
        self._ordered_indices.extend(plane_points)
        self._ordered_indices.extend(point_groups[1])
        self._inv_ordered_indices = np.argsort(self._ordered_indices)
        self._point_groups_permutations = point_groups_permutations
        self.explicit_permutations = explicit_permutations
        self.explicit_optimized_permutations = explicit_optimized_permutations
        self._safe_permutations = None
        if self.explicit_optimized_permutations is not None:
            self._permutations = self.explicit_optimized_permutations
        elif self.explicit_permutations is not None:
            self._permutations = self.explicit_permutations
        self.multiplicity = multiplicity
        self.other_plane_points = other_plane_points
        self.minimum_number_of_points = minimum_number_of_points
        self.maximum_number_of_points = len(self.plane_points)
        self._ref_separation_perm = list(self.point_groups[0])
        self._ref_separation_perm.extend(list(self.plane_points))
        self._ref_separation_perm.extend(list(self.point_groups[1]))
        self._argsorted_ref_separation_perm = list(
            np.argsort(self._ref_separation_perm))

    @property
    def ordered_indices(self):
        return self._ordered_indices

    @property
    def inv_ordered_indices(self):
        return self._inv_ordered_indices

    @property
    def permutations(self):
        return self._permutations

    @property
    def ref_separation_perm(self):
        return self._ref_separation_perm

    @property
    def argsorted_ref_separation_perm(self):
        return self._argsorted_ref_separation_perm

    def safe_plane_permutations(self, ordered_plane=False,
                                ordered_point_groups=None):
        ordered_point_groups = [False,
                                False] if ordered_point_groups is None else ordered_point_groups
        rotate = lambda s, n: s[-n:] + s[:-n]
        if ordered_plane and self.ordered_plane:
            plane_perms = [rotate(self.plane_points, ii) for ii in
                           range(len(self.plane_points))]
            invplanepoints = self.plane_points[::-1]
            plane_perms.extend([rotate(invplanepoints, ii) for ii in
                                range(len(self.plane_points) - 1, -1, -1)])
        else:
            plane_perms = list(itertools.permutations(self.plane_points))
        if ordered_point_groups[0] and self.ordered_point_groups[0]:
            s0_perms = [rotate(self.point_groups[0], ii) for ii in
                        range(len(self.point_groups[0]))]
            invpg0 = self.point_groups[0][::-1]
            s0_perms.extend([rotate(invpg0, ii) for ii in range(len(invpg0))])
        else:
            s0_perms = list(itertools.permutations(self.point_groups[0]))
        if ordered_point_groups[1] and self.ordered_point_groups[1]:
            s2_perms = [rotate(self.point_groups[1], ii) for ii in
                        range(len(self.point_groups[1]))]
            invpg2 = self.point_groups[1][::-1]
            s2_perms.extend([rotate(invpg2, ii) for ii in range(len(invpg2))])
        else:
            s2_perms = list(itertools.permutations(self.point_groups[1]))
        add_opposite = False
        if self._safe_permutations is None:
            self._safe_permutations = []
            for perm_side1 in s0_perms:
                for perm_sep_plane in plane_perms:
                    for perm_side2 in s2_perms:
                        perm = list(perm_side1)
                        perm.extend(list(perm_sep_plane))
                        perm.extend(list(perm_side2))
                        self._safe_permutations.append(perm)
                        if add_opposite:
                            perm = list(perm_side2)
                            perm.extend(list(perm_sep_plane))
                            perm.extend(list(perm_side1))
                            self._safe_permutations.append(perm)
        return self._safe_permutations

    def safe_separation_permutations(self, ordered_plane=False,
                                     ordered_point_groups=None,
                                     add_opposite=False):
        s0 = range(len(self.point_groups[0]))
        plane = range(len(self.point_groups[0]),
                      len(self.point_groups[0]) + len(self.plane_points))
        s1 = range(len(self.point_groups[0]) + len(self.plane_points),
                   len(self.point_groups[0]) + len(self.plane_points) + len(
                       self.point_groups[1]))
        ordered_point_groups = [False,
                                False] if ordered_point_groups is None else ordered_point_groups
        rotate = lambda s, n: s[-n:] + s[:-n]
        if ordered_plane and self.ordered_plane:
            plane_perms = [rotate(plane, ii) for ii in range(len(plane))]
            inv_plane = plane[::-1]
            plane_perms.extend(
                [rotate(inv_plane, ii) for ii in range(len(inv_plane))])
        else:
            plane_perms = list(itertools.permutations(plane))
        if ordered_point_groups[0] and self.ordered_point_groups[0]:
            s0_perms = [rotate(s0, ii) for ii in range(len(s0))]
            inv_s0 = s0[::-1]
            s0_perms.extend([rotate(inv_s0, ii) for ii in range(len(inv_s0))])
        else:
            s0_perms = list(itertools.permutations(s0))
        if ordered_point_groups[1] and self.ordered_point_groups[1]:
            s1_perms = [rotate(s1, ii) for ii in range(len(s1))]
            inv_s1 = s1[::-1]
            s1_perms.extend([rotate(inv_s1, ii) for ii in range(len(inv_s1))])
        else:
            s1_perms = list(itertools.permutations(s1))
        if self._safe_permutations is None:
            self._safe_permutations = []
            for perm_side1 in s0_perms:
                for perm_sep_plane in plane_perms:
                    for perm_side2 in s1_perms:
                        perm = list(perm_side1)
                        perm.extend(list(perm_sep_plane))
                        perm.extend(list(perm_side2))
                        self._safe_permutations.append(perm)
                        if add_opposite:
                            perm = list(perm_side2)
                            perm.extend(list(perm_sep_plane))
                            perm.extend(list(perm_side1))
                            self._safe_permutations.append(perm)
        return self._safe_permutations

    @property
    def as_dict(self):
        return {"@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "plane_points": self.plane_points,
                "mirror_plane": self.mirror_plane,
                "ordered_plane": self.ordered_plane,
                "point_groups": self.point_groups,
                "ordered_point_groups": self.ordered_point_groups,
                "point_groups_permutations": self._point_groups_permutations,
                "explicit_permutations": self.explicit_permutations,
                "explicit_optimized_permutations": self.explicit_optimized_permutations,
                "multiplicity": self.multiplicity,
                "other_plane_points": self.other_plane_points,
                "minimum_number_of_points": self.minimum_number_of_points}

    @classmethod
    def from_dict(cls, dd):
        eop = dd[
            'explicit_optimized_permutations'] if 'explicit_optimized_permutations' in dd else None
        return cls(plane_points=dd['plane_points'],
                   mirror_plane=dd['mirror_plane'],
                   ordered_plane=dd['ordered_plane'],
                   point_groups=dd['point_groups'],
                   ordered_point_groups=dd['ordered_point_groups'],
                   point_groups_permutations=dd['point_groups_permutations'],
                   explicit_permutations=dd['explicit_permutations'],
                   explicit_optimized_permutations=eop,
                   multiplicity=dd[
                       'multiplicity'] if 'multiplicity' in dd else None,
                   other_plane_points=dd[
                       'other_plane_points'] if 'other_plane_points' in dd else None,
                   minimum_number_of_points=dd['minimum_number_of_points'])

    def __str__(self):
        out = 'Separation plane algorithm with the following reference separation :\n'
        out += '[{}] | [{}] | [{}]'.format(
            '-'.join(str(pp) for pp in [self.point_groups[0]]),
            '-'.join(str(pp) for pp in [self.plane_points]),
            '-'.join(str(pp) for pp in [self.point_groups[1]]),
        )
        return out

class CoordinationGeometry(object):
    """
    Class used to store the ideal representation of a chemical environment or "coordination geometry"
    """

    class NeighborsSetsHints(object):

        ALLOWED_HINTS_TYPES = ['single_cap', 'double_cap', 'triple_cap']

        def __init__(self, hints_type, options):
            if hints_type not in self.ALLOWED_HINTS_TYPES:
                raise ValueError('Type "{}" for NeighborsSetsHints is not allowed'.format(type))
            self.hints_type = hints_type
            self.options = options

        def hints(self, hints_info):
            if hints_info['csm'] > self.options['csm_max']:
                return []
            return object.__getattribute__(self, '{}_hints'.format(self.hints_type))(hints_info)

        def single_cap_hints(self, hints_info):
            cap_index_perfect = self.options['cap_index']
            nb_set = hints_info['nb_set']
            permutation = hints_info['permutation']
            nb_set_voronoi_indices_perfect_aligned = nb_set.get_neighb_voronoi_indices(permutation=permutation)
            cap_voronoi_index = nb_set_voronoi_indices_perfect_aligned[cap_index_perfect]
            new_site_voronoi_indices = list(nb_set.site_voronoi_indices)
            new_site_voronoi_indices.remove(cap_voronoi_index)
            return [new_site_voronoi_indices]

        def double_cap_hints(self, hints_info):
            first_cap_index_perfect = self.options['first_cap_index']
            second_cap_index_perfect = self.options['second_cap_index']
            nb_set = hints_info['nb_set']
            permutation = hints_info['permutation']
            nb_set_voronoi_indices_perfect_aligned = nb_set.get_neighb_voronoi_indices(permutation=permutation)
            first_cap_voronoi_index = nb_set_voronoi_indices_perfect_aligned[first_cap_index_perfect]
            second_cap_voronoi_index = nb_set_voronoi_indices_perfect_aligned[second_cap_index_perfect]
            new_site_voronoi_indices1 = list(nb_set.site_voronoi_indices)
            new_site_voronoi_indices2 = list(nb_set.site_voronoi_indices)
            new_site_voronoi_indices3 = list(nb_set.site_voronoi_indices)
            new_site_voronoi_indices1.remove(first_cap_voronoi_index)
            new_site_voronoi_indices2.remove(second_cap_voronoi_index)
            new_site_voronoi_indices3.remove(first_cap_voronoi_index)
            new_site_voronoi_indices3.remove(second_cap_voronoi_index)
            return [new_site_voronoi_indices1, new_site_voronoi_indices2, new_site_voronoi_indices3]

        def triple_cap_hints(self, hints_info):
            first_cap_index_perfect = self.options['first_cap_index']
            second_cap_index_perfect = self.options['second_cap_index']
            third_cap_index_perfect = self.options['third_cap_index']
            nb_set = hints_info['nb_set']
            permutation = hints_info['permutation']
            nb_set_voronoi_indices_perfect_aligned = nb_set.get_neighb_voronoi_indices(permutation=permutation)
            first_cap_voronoi_index = nb_set_voronoi_indices_perfect_aligned[first_cap_index_perfect]
            second_cap_voronoi_index = nb_set_voronoi_indices_perfect_aligned[second_cap_index_perfect]
            third_cap_voronoi_index = nb_set_voronoi_indices_perfect_aligned[third_cap_index_perfect]
            new_site_voronoi_indices1 = list(nb_set.site_voronoi_indices)
            new_site_voronoi_indices2 = list(nb_set.site_voronoi_indices)
            new_site_voronoi_indices3 = list(nb_set.site_voronoi_indices)
            new_site_voronoi_indices4 = list(nb_set.site_voronoi_indices)
            new_site_voronoi_indices5 = list(nb_set.site_voronoi_indices)
            new_site_voronoi_indices6 = list(nb_set.site_voronoi_indices)
            new_site_voronoi_indices7 = list(nb_set.site_voronoi_indices)
            new_site_voronoi_indices1.remove(first_cap_voronoi_index)
            new_site_voronoi_indices2.remove(second_cap_voronoi_index)
            new_site_voronoi_indices3.remove(third_cap_voronoi_index)
            new_site_voronoi_indices4.remove(second_cap_voronoi_index)
            new_site_voronoi_indices4.remove(third_cap_voronoi_index)
            new_site_voronoi_indices5.remove(first_cap_voronoi_index)
            new_site_voronoi_indices5.remove(third_cap_voronoi_index)
            new_site_voronoi_indices6.remove(first_cap_voronoi_index)
            new_site_voronoi_indices6.remove(second_cap_voronoi_index)
            new_site_voronoi_indices7.remove(first_cap_voronoi_index)
            new_site_voronoi_indices7.remove(second_cap_voronoi_index)
            new_site_voronoi_indices7.remove(third_cap_voronoi_index)
            return [new_site_voronoi_indices1, new_site_voronoi_indices2, new_site_voronoi_indices3,
                    new_site_voronoi_indices4, new_site_voronoi_indices5, new_site_voronoi_indices6,
                    new_site_voronoi_indices7]

        def as_dict(self):
            return {'hints_type': self.hints_type,
                    'options': self.options}

        @classmethod
        def from_dict(cls, dd):
            return cls(hints_type=dd['hints_type'],
                       options=dd['options'])

    def __init__(self, mp_symbol, name, alternative_names=None,
                 IUPAC_symbol=None, IUCr_symbol=None, coordination=None,
                 central_site=np.zeros(3), points=None, solid_angles=None,
                 permutations_safe_override=False,
                 plane_ordering_override=True, deactivate=False, faces=None,
                 edges=None,
                 plane_safe_permutations=False, algorithms=None,
                 equivalent_indices=None,
                 neighbors_sets_hints=None):
        """
        Initializes one "coordination geometry" according to [Pure Appl. Chem., Vol. 79, No. 10, pp. 1779--1799, 2007]
        and [Acta Cryst. A, Vol. 46, No. 1, pp. 1--11, 1990].
        :param mp_symbol: Symbol used internally for the coordination geometry.
        :param name: Name of the coordination geometry.
        :param alternative_names: Alternative names for this coordination geometry.
        :param IUPAC_symbol: The IUPAC symbol of this coordination geometry.
        :param IUCr_symbol: The IUCr symbol of this coordination geometry.
        :param coordination: The coordination number of this coordination geometry (number of neighboring atoms).
        :param central_site: The coordinates of the central site of this coordination geometry.
        :param points: The list of the coordinates of all the points of this coordination geometry.
        :param separation_planes: List of separation facets to help set up the permutations
        :param permutation_safe_override: Computes all the permutations if set to True (overrides the plane separation
        algorithms or any other algorithm, for testing purposes)
        :param plane_ordering_override: Computes all the permutations of the plane separation algorithm if set to False
        otherwise, uses the anticlockwise ordering of the separation facets (for testing purposes)
        :param deactivate: deactivates this coordination geometry in the search
        :param faces : list of the faces with their vertices given in a clockwise or anticlockwise order, for drawing
        purposes
        :param : list of edges, for drawing purposes
        """
        self._mp_symbol = mp_symbol
        self.name = name
        self.alternative_names = alternative_names if alternative_names is not None else []
        self.IUPACsymbol = IUPAC_symbol
        self.IUCrsymbol = IUCr_symbol
        self.coordination = coordination
        self.central_site = np.array(central_site)
        self.points = points
        self._solid_angles = solid_angles
        self.permutations_safe_override = permutations_safe_override
        self.plane_ordering_override = plane_ordering_override
        self.plane_safe_permutations = plane_safe_permutations
        # self.setup_permutations(permutations)
        self.deactivate = deactivate
        self._faces = faces
        self._edges = edges
        self._algorithms = algorithms
        if points is not None:
            self.centroid = np.mean(np.array(points), axis=0)
        else:
            self.centroid = None
        self.equivalent_indices = equivalent_indices
        self.neighbors_sets_hints = neighbors_sets_hints

    def as_dict(self):
        return {'mp_symbol': self._mp_symbol,
                'name': self.name,
                'alternative_names': self.alternative_names,
                'IUPAC_symbol': self.IUPACsymbol,
                'IUCr_symbol': self.IUCrsymbol,
                'coordination': self.coordination,
                'central_site': [float(xx) for xx in self.central_site],
                'points': [[float(xx) for xx in pp] for pp in
                           self.points] if self.points is not None else None,
                'solid_angles': [float(ang) for ang in
                                 self._solid_angles] if self._solid_angles is not None else None,
                'deactivate': self.deactivate,
                '_faces': self._faces,
                '_edges': self._edges,
                '_algorithms': [algo.as_dict for algo in
                                self._algorithms] if self._algorithms is not None else None,
                'equivalent_indices': self.equivalent_indices,
                'neighbors_sets_hints': [nbsh.as_dict() for nbsh in self.neighbors_sets_hints]
                if self.neighbors_sets_hints is not None else None}

    @classmethod
    def from_dict(cls, dd):
        dec = MontyDecoder()
        return cls(mp_symbol=dd['mp_symbol'],
                   name=dd['name'],
                   alternative_names=dd['alternative_names'],
                   IUPAC_symbol=dd['IUPAC_symbol'],
                   IUCr_symbol=dd['IUCr_symbol'],
                   coordination=dd['coordination'],
                   central_site=dd['central_site'],
                   points=dd['points'],
                   solid_angles=(dd['solid_angles'] if 'solid_angles' in dd
                                 else [4.0 * np.pi / dd['coordination']] * dd[
                       'coordination']),
                   deactivate=dd['deactivate'],
                   faces=dd['_faces'],
                   edges=dd['_edges'],
                   algorithms=[dec.process_decoded(algo_d)
                               for algo_d in dd['_algorithms']] if dd[
                                                                       '_algorithms'] is not None else None,
                   equivalent_indices=dd[
                       'equivalent_indices'] if 'equivalent_indices' in dd else None,
                   neighbors_sets_hints=[cls.NeighborsSetsHints.from_dict(nbshd)
                                         for nbshd in dd['neighbors_sets_hints']]
                   if 'neighbors_sets_hints' in dd else None)

    def __str__(self):
        symbol = ''
        if self.IUPAC_symbol is not None:
            symbol += ' (IUPAC: {s}'.format(s=self.IUPAC_symbol)
            if self.IUCr_symbol is not None:
                symbol += ' || IUCr: {s})'.format(s=self.IUCr_symbol)
            else:
                symbol += ')'
        elif self.IUCr_symbol is not None:
            symbol += ' (IUCr: {s})'.format(s=self.IUCr_symbol)
        outs = ['Coordination geometry type : {n}{s}\n'.format(n=self.name,
                                                               s=symbol),
                '  - coordination number : {c}'.format(c=self.coordination)]
        if self.points is None:
            outs.append('... not yet implemented')
        else:
            outs.append('  - list of points :')
            for pp in self.points:
                outs.append('    - {p}'.format(p=pp))
        outs.append(
            '------------------------------------------------------------')
        outs.append('')

        return '\n'.join(outs)

    def __repr__(self):
        symbol = ''
        if self.IUPAC_symbol is not None:
            symbol += ' (IUPAC: {s}'.format(s=self.IUPAC_symbol)
            if self.IUCr_symbol is not None:
                symbol += ' || IUCr: {s})'.format(s=self.IUCr_symbol)
            else:
                symbol += ')'
        elif self.IUCr_symbol is not None:
            symbol += ' (IUCr: {s})'.format(s=self.IUCr_symbol)
        outs = ['Coordination geometry type : {n}{s}\n'.format(n=self.name,
                                                               s=symbol),
                '  - coordination number : {c}'.format(c=self.coordination)]
        outs.append(
            '------------------------------------------------------------')
        outs.append('')
        return '\n'.join(outs)

    def __len__(self):
        return self.coordination

    def set_permutations_safe_override(self, permutations_safe_override):
        self.permutations_safe_override = permutations_safe_override
        # self.setup_permutations()

    @property
    def distfactor_max(self):
        dists = [np.linalg.norm(pp - self.central_site) for pp in self.points]
        return np.max(dists) / np.min(dists)

    @property
    def coordination_number(self):
        """
        Returns the coordination number of this coordination geometry.
        """
        return self.coordination

    @property
    def mp_symbol(self):
        """
        Returns the MP symbol of this coordination geometry.
        """
        return self._mp_symbol

    @property
    def ce_symbol(self):
        """
        Returns the symbol of this coordination geometry.
        """
        return self._mp_symbol

    def get_coordination_number(self):
        """
        Returns the coordination number of this coordination geometry.
        """
        return self.coordination

    def is_implemented(self):
        """
        Returns True if this coordination geometry is implemented.
        """
        return bool(self.points)

    def get_name(self):
        """
        Returns the name of this coordination geometry.
        """
        return self.name

    @property
    def IUPAC_symbol(self):
        """
        Returns the IUPAC symbol of this coordination geometry.
        """
        return self.IUPACsymbol

    @property
    def IUPAC_symbol_str(self):
        """
        Returns a string representation of the IUPAC symbol of this coordination geometry.
        """
        return str(self.IUPACsymbol)

    @property
    def IUCr_symbol(self):
        """
        Returns the IUCr symbol of this coordination geometry.
        """
        return self.IUCrsymbol

    @property
    def IUCr_symbol_str(self):
        """
        Returns a string representation of the IUCr symbol of this coordination geometry.
        """
        return str(self.IUCrsymbol)

    @property
    def number_of_permutations(self):
        """
        Returns the number of permutations of this coordination geometry.
        """
        if self.permutations_safe_override:
            return factorial(self.coordination)
        elif self.permutations is None:
            return factorial(self.coordination)
        return len(self.permutations)

    def ref_permutation(self, permutation):
        perms = []
        for eqv_indices in self.equivalent_indices:
            perms.append(tuple([permutation[ii] for ii in eqv_indices]))
        perms.sort()
        return perms[0]

    @property
    def algorithms(self):
        """
        Returns the list of algorithms that are used to identify this coordination geometry.
        """
        return self._algorithms

    def get_central_site(self):
        """
        Returns the central site of this coordination geometry.
        """
        return self.central_site

    def faces(self, sites, permutation=None):
        """
        Returns the list of faces of this coordination geometry. Each face is given as a
        list of its vertices coordinates.
        """
        if permutation is None:
            coords = [site.coords for site in sites]
        else:
            coords = [sites[ii].coords for ii in permutation]
        return [[coords[ii] for ii in f] for f in self._faces]

    def edges(self, sites, permutation=None, input='sites'):
        """
        Returns the list of edges of this coordination geometry. Each edge is given as a
        list of its end vertices coordinates.
        """
        if input == 'sites':
            coords = [site.coords for site in sites]
        elif input == 'coords':
            coords = sites
        # if permutation is None:
        #     coords = [site.coords for site in sites]
        # else:
        #     coords = [sites[ii].coords for ii in permutation]
        if permutation is not None:
            coords = [coords[ii] for ii in permutation]
        return [[coords[ii] for ii in e] for e in self._edges]

    def solid_angles(self, permutation=None):
        """
        Returns the list of "perfect" solid angles Each edge is given as a
        list of its end vertices coordinates.
        """
        if permutation is None:
            return self._solid_angles
        else:
            return [self._solid_angles[ii] for ii in permutation]

    def get_pmeshes(self, sites, permutation=None):
        """
        Returns the pmesh strings used for jmol to show this geometry.
        """
        pmeshes = []
        # _vertices = [site.coords for site in sites]
        if permutation is None:
            _vertices = [site.coords for site in sites]
        else:
            _vertices = [sites[ii].coords for ii in permutation]
        _face_centers = []
        number_of_faces = 0
        for face in self._faces:
            if len(face) in [3, 4]:
                number_of_faces += 1
            else:
                number_of_faces += len(face)

            _face_centers.append(np.array([np.mean([_vertices[face_vertex][ii]
                                                    for face_vertex in face])
                                           for ii in range(3)]))

        out = '{}\n'.format(len(_vertices) + len(_face_centers))
        for vv in _vertices:
            out += '{:15.8f} {:15.8f} {:15.8f}\n'.format(vv[0], vv[1], vv[2])
        for fc in _face_centers:
            out += '{:15.8f} {:15.8f} {:15.8f}\n'.format(fc[0], fc[1], fc[2])
        out += '{:d}\n'.format(number_of_faces)
        for iface, face in enumerate(self._faces):
            if len(face) == 3:
                out += '4\n'
            elif len(face) == 4:
                out += '5\n'
            else:
                for ii in range(len(face)):
                    out += '4\n'
                    out += '{:d}\n'.format(len(_vertices) + iface)
                    out += '{:d}\n'.format(face[ii])
                    out += '{:d}\n'.format(face[np.mod(ii + 1, len(face))])
                    out += '{:d}\n'.format(len(_vertices) + iface)
            if len(face) in [3, 4]:
                for face_vertex in face:
                    out += '{:d}\n'.format(face_vertex)
                out += '{:d}\n'.format(face[0])
        pmeshes.append({"pmesh_string": out})
        return pmeshes

    def get_pmeshes_test(self, sites, permutation=None):
        """
        Returns the pmesh strings used for jmol to show this geometry.
        """
        pmeshes = []
        _vertices = [site.coords for site in sites]
        # if permutation is None:
        #    _vertices = [site.coords for site in sites]
        # else:
        #    _vertices = [sites[ii].coords for ii in permutation]
        _face_centers = []
        number_of_faces = 0
        for face in self._faces:
            if len(face) in [3, 4]:
                number_of_faces += 1
            else:
                number_of_faces += len(face)

            _face_centers.append(np.array([np.mean([_vertices[face_vertex][ii]
                                                    for face_vertex in face])
                                           for ii in range(3)]))

        out = '{}\n'.format(len(_vertices) + len(_face_centers))
        for vv in _vertices:
            out += '{:15.8f} {:15.8f} {:15.8f}\n'.format(vv[0], vv[1], vv[2])
        for fc in _face_centers:
            out += '{:15.8f} {:15.8f} {:15.8f}\n'.format(fc[0], fc[1], fc[2])
        out += '{:d}\n'.format(number_of_faces)
        for iface, face in enumerate(self._faces):
            if len(face) == 3:
                out += '4\n'
            elif len(face) == 4:
                out += '5\n'
            else:
                for ii in range(len(face)):
                    out += '4\n'
                    out += '{:d}\n'.format(len(_vertices) + iface)
                    out += '{:d}\n'.format(permutation[face[ii]])
                    out += '{:d}\n'.format(
                        permutation[face[np.mod(ii + 1, len(face))]])
                    out += '{:d}\n'.format(len(_vertices) + iface)
            if len(face) in [3, 4]:
                for face_vertex in face:
                    out += '{:d}\n'.format(permutation[face_vertex])
                out += '{:d}\n'.format(permutation[face[0]])
        pmeshes.append({"pmesh_string": out})
        return pmeshes

class AllCoordinationGeometries(dict):
    """
    Class used to store all the reference "coordination geometries" (list with instances of the CoordinationGeometry
    classes)
    """

    def __init__(self, permutations_safe_override=False, only_symbols=None):
        """
            Initializes the list of Coordination Geometries
            :param permutations_safe_override:
            :param only_symbols:
            """
        dict.__init__(self)
        self.cg_list = list()
        if only_symbols is None:
            f = open(
                '{}/coordination_geometries_files/allcg.txt'.format(module_dir),
                'r')
            data = f.readlines()
            f.close()
            for line in data:
                cg_file = '{}/{}'.format(module_dir, line.strip())
                f = open(cg_file, 'r')
                dd = json.load(f)
                f.close()
                self.cg_list.append(CoordinationGeometry.from_dict(dd))
        else:
            for symbol in only_symbols:
                fsymbol = symbol.replace(':', '#')
                cg_file = '{}/coordination_geometries_files/{}.json'.format(
                    module_dir, fsymbol)
                f = open(cg_file, 'r')
                dd = json.load(f)
                f.close()
                self.cg_list.append(CoordinationGeometry.from_dict(dd))

        self.cg_list.append(CoordinationGeometry(UNKNOWN_ENVIRONMENT_SYMBOL,
                                                 "Unknown environment",
                                                 deactivate=True))
        self.cg_list.append(CoordinationGeometry(UNCLEAR_ENVIRONMENT_SYMBOL,
                                                 "Unclear environment",
                                                 deactivate=True))
        if permutations_safe_override:
            for cg in self.cg_list:
                cg.set_permutations_safe_override(True)

    def __getitem__(self, key):
        return self.get_geometry_from_mp_symbol(key)

    def __repr__(self):
        """
        Returns a string with the list of coordination geometries.
        """
        outs = ['', '#=================================#',
                '# List of coordination geometries #',
                '#=================================#', '']
        for cg in self.cg_list:
            outs.append(repr(cg))

        return '\n'.join(outs)

    def __str__(self):
        """
        Returns a string with the list of coordination geometries that are implemented.
        """
        outs = ['', '#=======================================================#',
                '# List of coordination geometries currently implemented #',
                '#=======================================================#', '']
        for cg in self.cg_list:
            if cg.is_implemented():
                outs.append(str(cg))

        return '\n'.join(outs)

    def get_geometries(self, coordination=None, returned='cg'):
        """
        Returns a list of coordination geometries with the given coordination number.
        :param coordination: The coordination number of which the list of coordination geometries are returned.
        """
        geom = list()
        if coordination is None:
            for gg in self.cg_list:
                if returned == 'cg':
                    geom.append(gg)
                elif returned == 'mp_symbol':
                    geom.append(gg.mp_symbol)
        else:
            for gg in self.cg_list:
                if gg.get_coordination_number() == coordination:
                    if returned == 'cg':
                        geom.append(gg)
                    elif returned == 'mp_symbol':
                        geom.append(gg.mp_symbol)
        return geom

    def get_symbol_name_mapping(self, coordination=None):
        geom = {}
        if coordination is None:
            for gg in self.cg_list:
                geom[gg.mp_symbol] = gg.name
        else:
            for gg in self.cg_list:
                if gg.get_coordination_number() == coordination:
                    geom[gg.mp_symbol] = gg.name
        return geom

    def get_symbol_cn_mapping(self, coordination=None):
        geom = {}
        if coordination is None:
            for gg in self.cg_list:
                geom[gg.mp_symbol] = gg.coordination_number
        else:
            for gg in self.cg_list:
                if gg.get_coordination_number() == coordination:
                    geom[gg.mp_symbol] = gg.coordination_number
        return geom

    def get_implemented_geometries(self, coordination=None, returned='cg',
                                   include_deactivated=False):
        """
        Returns a list of the implemented coordination geometries with the given coordination number.
        :param coordination: The coordination number of which the list of implemented coordination geometries
        are returned.
        """
        geom = list()
        if coordination is None:
            for gg in self.cg_list:
                if gg.points is not None and (
                            (not gg.deactivate) or include_deactivated):
                    if returned == 'cg':
                        geom.append(gg)
                    elif returned == 'mp_symbol':
                        geom.append(gg.mp_symbol)
        else:
            for gg in self.cg_list:
                if gg.get_coordination_number() == coordination and gg.points is not None and \
                        ((not gg.deactivate) or include_deactivated):
                    if returned == 'cg':
                        geom.append(gg)
                    elif returned == 'mp_symbol':
                        geom.append(gg.mp_symbol)
        return geom

    def get_not_implemented_geometries(self, coordination=None,
                                       returned='mp_symbol'):
        """
        Returns a list of the implemented coordination geometries with the given coordination number.
        :param coordination: The coordination number of which the list of implemented coordination geometries
        are returned.
        """
        geom = list()
        if coordination is None:
            for gg in self.cg_list:
                if gg.points is None:
                    if returned == 'cg':
                        geom.append(gg)
                    elif returned == 'mp_symbol':
                        geom.append(gg.mp_symbol)
        else:
            for gg in self.cg_list:
                if gg.get_coordination_number() == coordination and gg.points is None:
                    if returned == 'cg':
                        geom.append(gg)
                    elif returned == 'mp_symbol':
                        geom.append(gg.mp_symbol)
        return geom

    def get_geometry_from_name(self, name):
        """
        Returns the coordination geometry of the given name.
        :param name: The name of the coordination geometry.
        """
        for gg in self.cg_list:
            if gg.name == name or name in gg.alternative_names:
                return gg
        raise LookupError(
            'No coordination geometry found with name "{name}"'.format(
                name=name))

    def get_geometry_from_IUPAC_symbol(self, IUPAC_symbol):
        """
        Returns the coordination geometry of the given IUPAC symbol.
        :param IUPAC_symbol: The IUPAC symbol of the coordination geometry.
        """
        for gg in self.cg_list:
            if gg.IUPAC_symbol == IUPAC_symbol:
                return gg
        raise LookupError(
            'No coordination geometry found with IUPAC symbol "{symbol}"'.format(
                symbol=IUPAC_symbol))

    def get_geometry_from_IUCr_symbol(self, IUCr_symbol):
        """
        Returns the coordination geometry of the given IUCr symbol.
        :param IUCr_symbol: The IUCr symbol of the coordination geometry.
        """
        for gg in self.cg_list:
            if gg.IUCr_symbol == IUCr_symbol:
                return gg
        raise LookupError(
            'No coordination geometry found with IUCr symbol "{symbol}"'.format(
                symbol=IUCr_symbol))

    def get_geometry_from_mp_symbol(self, mp_symbol):
        """
        Returns the coordination geometry of the given mp_symbol.
        :param mp_symbol: The mp_symbol of the coordination geometry.
        """
        for gg in self.cg_list:
            if gg.mp_symbol == mp_symbol:
                return gg
        raise LookupError(
            'No coordination geometry found with mp_symbol "{symbol}"'.format(
                symbol=mp_symbol))

    def is_a_valid_coordination_geometry(self, mp_symbol=None,
                                         IUPAC_symbol=None, IUCr_symbol=None,
                                         name=None, cn=None):
        """
        Checks whether a given coordination geometry is valid (exists) and whether the parameters are coherent with
        each other.
        :param IUPAC_symbol:
        :param IUCr_symbol:
        :param name:
        :param cn:
        :param mp_symbol: The mp_symbol of the coordination geometry.
        """
        if name is not None:
            raise NotImplementedError(
                'is_a_valid_coordination_geometry not implemented for the name')
        if mp_symbol is None and IUPAC_symbol is None and IUCr_symbol is None:
            raise SyntaxError(
                'missing argument for is_a_valid_coordination_geometry : at least one of mp_symbol, '
                'IUPAC_symbol and IUCr_symbol must be passed to the function')
        if mp_symbol is not None:
            try:
                cg = self.get_geometry_from_mp_symbol(mp_symbol)
                if IUPAC_symbol is not None:
                    if IUPAC_symbol != cg.IUPAC_symbol:
                        return False
                if IUCr_symbol is not None:
                    if IUCr_symbol != cg.IUCr_symbol:
                        return False
                if cn is not None:
                    if int(cn) != int(cg.coordination_number):
                        return False
                return True
            except LookupError:
                return False
        elif IUPAC_symbol is not None:
            try:
                cg = self.get_geometry_from_IUPAC_symbol(IUPAC_symbol)
                if IUCr_symbol is not None:
                    if IUCr_symbol != cg.IUCr_symbol:
                        return False
                if cn is not None:
                    if cn != cg.coordination_number:
                        return False
                return True
            except LookupError:
                return False
        elif IUCr_symbol is not None:
            try:
                cg = self.get_geometry_from_IUCr_symbol(IUCr_symbol)
                if cn is not None:
                    if cn != cg.coordination_number:
                        return False
                return True
            except LookupError:
                return True
        raise Exception('Should not be here !')

    def pretty_print(self, type='implemented_geometries', maxcn=8, additional_info=None):
        if type == 'all_geometries_latex_images':
            mystring = ''
            for cn in range(1, maxcn + 1):
                mystring += '\\section*{{Coordination {cn}}}\n\n'.format(cn=cn)
                for cg in self.get_implemented_geometries(coordination=cn,
                                                          returned='cg'):
                    mystring += '\\subsubsection*{{{mp} : {name}}}\n\n'.format(
                        mp=cg.mp_symbol, name=cg.get_name())
                    mystring += 'IUPAC : {iupac}\n\nIUCr : {iucr}\n\n'.format(
                        iupac=cg.IUPAC_symbol, iucr=cg.IUCr_symbol)
                    mystring += '\\begin{center}\n'
                    mystring += '\\includegraphics[scale=0.15]{{images/{let}_{cif}.png}}\n'.format(
                        let=cg.mp_symbol.split(':')[0],
                        cif=cg.mp_symbol.split(':')[1])
                    mystring += '\\end{center}\n\n'
                for cg in self.get_not_implemented_geometries(cn,
                                                              returned='cg'):
                    mystring += '\\subsubsection*{{{mp} : {name}}}\n\n'.format(
                        mp=cg.mp_symbol, name=cg.get_name())
                    mystring += 'IUPAC : {iupac}\n\nIUCr : {iucr}\n\n'.format(
                        iupac=cg.IUPAC_symbol, iucr=cg.IUCr_symbol)
        elif type == 'all_geometries_latex':
            mystring = ''
            for cn in range(1, maxcn + 1):
                mystring += '\\subsection*{{Coordination {cn}}}\n\n'.format(
                    cn=cn)
                mystring += '\\begin{itemize}\n'
                for cg in self.get_implemented_geometries(coordination=cn,
                                                          returned='cg'):
                    mystring += '\\item {mp} $\\rightarrow$ {name} '.format(
                        mp=cg.mp_symbol.replace('_',
                                                '\\_'),
                        name=cg.get_name())
                    mystring += '(IUPAC : {iupac} - IUCr : {iucr})\n'.format(
                        iupac=cg.IUPAC_symbol_str,
                        iucr=cg.IUCr_symbol_str.replace('[', '$[$').replace(']',
                                                                            '$]$'))
                for cg in self.get_not_implemented_geometries(cn,
                                                              returned='cg'):
                    mystring += '\\item {mp} $\\rightarrow$ {name} '.format(
                        mp=cg.mp_symbol.replace('_',
                                                '\\_'),
                        name=cg.get_name())
                    mystring += '(IUPAC : {iupac} - IUCr : {iucr})\n'.format(
                        iupac=cg.IUPAC_symbol_str,
                        iucr=cg.IUCr_symbol_str.replace('[', '$[$').replace(']',
                                                                            '$]$'))
                mystring += '\\end{itemize}\n\n'
        else:
            mystring = '+-------------------------+\n| Coordination geometries |\n+-------------------------+\n\n'
            for cn in range(1, maxcn + 1):
                mystring += '==>> CN = {cn} <<==\n'.format(cn=cn)
                if type == 'implemented_geometries':
                    for cg in self.get_implemented_geometries(coordination=cn):
                        if additional_info is not None:
                            if 'nb_hints' in additional_info:
                                if cg.neighbors_sets_hints is not None:
                                    addinfo = ' *'
                                else:
                                    addinfo = ''
                            else:
                                addinfo = ''
                        else:
                            addinfo = ''
                        mystring += ' - {mp} : {name}{addinfo}\n'.format(mp=cg.mp_symbol,
                                                                  name=cg.get_name(),
                                                                  addinfo=addinfo)
                elif type == 'all_geometries':
                    for cg in self.get_geometries(coordination=cn):
                        mystring += ' - {mp} : {name}\n'.format(mp=cg.mp_symbol,
                                                                name=cg.get_name())
                mystring += '\n'
        return mystring