# pylint: disable=too-many-lines
"""Nuclear potential module.

Module containing representations of 3D potentials for use in nuclear theory.
Also contains logic to read them from and save them to files with a standard
naming convention.

class Channel
-------------
A container for the channel information for a potential. It has the following
method::

    channel = Channel(spin, orb_ang_mom_1, orb_ang_mom_2, tot_ang_mom, isospin)

These are also commonly read as S, L, L, J, and T.

class CoupledChannel
--------------------
A container to handle coupled channels. It has the following method::

    channel = CoupledChannel(list_of_channels)

All channels in coupled channel should have same S, J, and T.

class PotentialType
-------------------
A container class to hold all the physical information about the potential. It
has the following method::

    potential_type = PotentialType(n_body, order, name, channel, particles)

class Potential
---------------
Abstraction for the representation of a potential. Handles the logic of adding
and removing weights. Can generate corresponding kinetic energy. It has the
following methods::

    potential = Potential(potential_type, nodes, weights, potential, lam=50.0,
                          has_weights=False)
    kinetic_energy = potential.kinetic_energy()
    potential_data_wo_weights = potential.without_weights()
    potential_data_w_weights = potential.with_weights()
    new_potential = potential.copy(potential_data, lam)
    reduced_potential = potential.reduce_dim(dim)

class CoupledPotential
----------------------
Abstraction for representation for potential of coupled channel. Handles logic
of adding and removing weights. Can generate kinetic energy. It has the
following methods::

    potential = CoupledPotential([potential1, potential2, potential3,
                                  potential4])
    kinetic_energy = potential.kinetic_energy()
    potential_data_wo_weights = potential.without_weights()
    potential_data_w_weights = potential.with_weights()
    new_potential = potential.copy(potential_data, lam)
    reduced_potential = potential.reduce_dim(dim)
    channel_potential = potential.extract_channel_potential(
        potential1.potential_type.channel
    )

Methods
-------
potential = load_from_file(file_str)

Method to load a potential from a file. Requires that standard file-naming
conventions have been followed.

potential = load(n_body, order, name, channel, lambda, particles,
                 num_points='*')

Method to load potential from a standard directory. Requires that potential was
saved there earlier.

save(potential, directory=None)

Method to save potential with correct naming convention either to a standard
folder or to a user-specified directory.

Changelog:

2018.11.14
    Added:
        CoupledChannel for coupled channels
        CoupledPotential for potentials in coupled channels

2018.11.09
    Added:
        load_from_file method
    Changed:
        Make load take parameters and use load_from_file for loading from a
        specific file
        Save now has different parameter ordering with the dir_str param being
        optional

2018.11.06
    Added:
        Initial creation of module

"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import glob
from math import pi
from math import sqrt
import os
import re

import numpy as np
import matplotlib.pyplot as plt

NBODY_DICT = {
    'NN': 2,
    '3N': 3,
}

STANDARD_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                             'potentials')

INV_NBODY_DICT = {v: k for k, v in NBODY_DICT.items()}

ORDER_DICT = {
    'LO': 0,
    'NLO': 1,
    'N2LO': 2,
    'N3LO': 3,
}

INV_ORDER_DICT = {v: k for k, v in ORDER_DICT.items()}

class Channel:
    """Container for information on channel for potential."""

    # pylint: disable=too-many-arguments
    def __init__(self, spin, orb_ang_mom_1, orb_ang_mom_2, tot_ang_mom,
                 isospin):
        """Create Channel object.

        Parameters
        ----------
        spin : int
            Spin quantum number.
        orb_ang_mom_1 : int
            First angular momentum quantum number.
        orb_ang_mom_2 : int
            Second angular momentum quantum number.
        tot_ang_mom : int
            Total angular momentum.
        isospin : int
            2-body isospin quantum number.

        """
        self._spin = spin
        self._l1 = orb_ang_mom_1
        self._l2 = orb_ang_mom_2
        self._j = tot_ang_mom
        self._isospin = isospin

    def as_5tuple(self):
        """Return 5-tuple representation of channel.

        Returns
        -------
        (int, int, int, int, int)
            5-tuple with channel quantum numbers.

        """
        return (self._spin, self._l1, self._l2, self._j, self._isospin)

    def __str__(self):
        """Return string representation of channel.

        Returns
        -------
        str
            String of 5 integers with channel information which are SLLJT.

        """
        return '{}{}{}{}{}'.format(self._spin, self._l1, self._l2, self._j,
                                   self._isospin)

    def __eq__(self, other):
        """Return whether channel is same as another channel object.

        Returns
        -------
        bool
            True if self and other are the same, False otherwise.

        """
        return self.as_5tuple() == other.as_5tuple()

    def __ne__(self, other):
        """Return whether channel is different from another channel object.

        Returns
        -------
        bool
            False if self and other are the same, False otherwise.

        """
        return self.as_5tuple() != other.as_5tuple()

class CoupledChannel(Channel):
    """Container for information about coupled channel."""

    def __init__(self, list_of_channels):
        """Create coupled channel container.

        Parameters
        ----------
        list_of_channels : list of Channel objects
            List of channels in coupled channel.

        """
        spins = {x.as_5tuple()[0] for x in list_of_channels}
        tot_ang_moms = {x.as_5tuple()[3] for x in list_of_channels}
        isospins = {x.as_5tuple()[4] for x in list_of_channels}
        if len(spins) * len(isospins) * len(tot_ang_moms) != 1:
            raise ValueError('Given channels cannot be coupled.')
        super(CoupledChannel, self).__init__(spins.pop(), '*', '*',
                                             tot_ang_moms.pop(),
                                             isospins.pop())
        self._channels = list_of_channels

    @property
    def channels(self):
        """Return list of channels in coupled channel.

        Returns
        -------
        list of Channel objects

        """
        return self._channels

    def __eq__(self, other):
        """Return whether coupled channel object is same as another.

        Returns
        -------
        bool
            True if coupled channels are equal, False otherwise.

        """
        return False not in {x == y for x, y in zip(self.channels,
                                                    other.channels)}

    def __ne__(self, other):
        """Return whether coupled channel object is not same as another.

        Returns
        -------
        bool
            True if coupled channels are not equal, False otherwise.

        """
        return False in {x == y for x, y in zip(self.channels,
                                                other.channels)}

class PotentialType:
    """Container for information related to potential."""

    # pylint: disable=too-many-arguments
    def __init__(self, n_body, order, name, channel, particles):
        """Construct potential type.

        Parameters
        ----------
        n_body : int
            Number of particles interacting in potential.
        order : int
            Order to which potential was calculated.
        name : str
            Name for potential, may reflect something about origin.
        channel: Channel
            Object representing the partial wave channel for the potential.
        particles: str
            String representing constituent particles in the interaction.

        """
        self._n_body = n_body
        self._order = order
        self._name = name
        self._channel = channel
        self._particles = particles

    @property
    def n_body(self):
        """Return number of particles in potential.

        Returns
        -------
        int
            Number of particles.

        """
        return self._n_body

    @property
    def order(self):
        """Return order to which potential was calculated.

        Returns
        -------
        int
            Order of potential.

        """
        return self._order

    @property
    def name(self):
        """Return name of potential.

        Returns
        -------
        str
            Name of potential.

        """
        return self._name

    @property
    def channel(self):
        """Return partial wave channel of potential.

        Returns
        -------
        Channel
            Channel object representing partial wave channel.

        """
        return self._channel

    @property
    def particles(self):
        """Return particles in interaction to which potential applies.

        Returns
        -------
        str
            String with particles in interaction.

        """
        return self._particles

    def __eq__(self, other):
        """Return whether potential type is same as other potential type.

        Returns
        -------
        bool
            True if same, False otherwise.

        """
        return ((self.n_body == other.n_body)
                and (self.order == other.order)
                and (self.name == other.name)
                and (self.channel == other.channel)
                and (self.particles == other.particles))

    def __ne__(self, other):
        """Return whether potential type is not same as other potential type.

        Returns
        -------
        bool
            False if same, True otherwise.

        """
        return not ((self.n_body == other.n_body)
                    and (self.order == other.order)
                    and (self.name == other.name)
                    and (self.channel == other.channel)
                    and (self.particles == other.particles))

class Potential:
    """Class encapsulating all relevant information about a potential."""

    # pylint: disable=too-many-arguments
    def __init__(self, potential_type, nodes, weights, potential, lam=50.0,
                 has_weights=False):
        """Create potential from parameters.

        Parameters
        ----------
        potential_type : PotentialType
            PotentialType instance with information about the potential.
        nodes : list of floats
            List of momenta at which the potential is defined.
        weights : list of floats
            List of integration weights corresponding to nodes.
        potential : matrix of floats
            Value of potential at incoming and outgoing momenta in nodes.
        lam : float, optional
            Value of lambda (SRG flow parameter) for potential. For unevolved
            potentials, a value of 50 is the default.
        has_weights : bool, optional
            Specifies whether potential given has weights factored in already.

        """
        self._potential_type = potential_type
        self._nodes = nodes
        self._weights = weights
        self._lam = lam
        if has_weights:
            potential = _rem_w(potential, self._weights, self._nodes)
        self._potential = potential

    def copy(self, potential, lam):
        """Create potential from current potential with new data and lam.

        Parameters
        ----------
        potential : matrix of floats
            Potential data.
        lam : float
            Value of lambda

        Returns
        -------
        Potential
            New potential with new data.

        """
        return Potential(self._potential_type, self._nodes, self._weights,
                         potential, lam)

    def with_weights(self):
        """Return potential with weights factored in (for calculations).

        Returns
        -------
        matrix of floats
            Potential with integration weights.

        """
        return _add_w(self._potential, self._weights, self._nodes)

    def without_weights(self):
        """Return potential without weights (for visualization).

        Returns
        -------
        matrix of floats
            Potential without integration weights.

        """
        return np.array(self._potential)

    def reduce_dim(self, dim):
        """Return new potential with only `dim` lowest energy states.

        Parameters
        ----------
        dim : int
            Dimension to which potential is to be reduced.

        Returns
        -------
        Potential
            New reduced dimension potential.

        Raises
        ------
        ValueError
            When value for new dim is too small or too large.

        """
        if dim >= len(self.nodes):
            raise ValueError('Value of dim is not smaller than current dim.')
        if dim <= 0:
            raise ValueError('Zero or negative dim is not allowed.')

        new_data = self._potential[np.ix_(list(range(dim)), list(range(dim)))]
        new_nodes = self._nodes[:dim]
        new_weights = self._weights[:dim]

        return Potential(self._potential_type, new_nodes, new_weights,
                         new_data, self._lam)

    def kinetic_energy(self):
        """Return kinetic energy for potential (for calculations).

        Returns
        -------
        matrix of floats
            Kinetic energy matrix.

        """
        return np.diag(np.array([p**2 for p in self._nodes]))

    def __eq__(self, other):
        """Return whether two potentials are equal to with numerical error.

        Returns
        -------
        bool
            True when potential type, nodes, weights, potential, and lam are
            all equal within epsilon, False otherwise.

        """
        # Numerical errors smaller than this are acceptable
        # If there is something wrong with the physics, it should produce
        # errors larger than this.
        eps = 10**(-4)

        if self.potential_type != other.potential_type:
            return False
        if self.dim != other.dim:
            return False
        if abs(self.lam - other.lam) > eps:
            return False
        for p_self, p_other, w_self, w_other in zip(self.nodes, other.nodes,
                                                    self.weights,
                                                    other.weights):
            if abs(p_self - p_other) > eps or abs(w_self - w_other) > eps:
                return False
        for i in range(self.dim):
            for j in range(self.dim):
                diff = abs(self.without_weights()[i][j] -
                           other.without_weights()[i][j])
                if diff > eps:
                    return False
        return True

    def __ne__(self, other):
        """Return whether two potentials are not equal to with numerical error.

        Returns
        -------
        bool
            False when potential type, nodes, weights, potential, and lam are
            all equal within epsilon, True otherwise.

        """
        # Numerical errors smaller than this are acceptable
        # If there is something wrong with the physics, it should produce
        # errors larger than this.
        eps = 10**(-4)

        if self.potential_type != other.potential_type:
            return True
        if self.dim != other.dim:
            return True
        if abs(self.lam - other.lam) > eps:
            return True
        for p_self, p_other, w_self, w_other in zip(self.nodes, other.nodes,
                                                    self.weights,
                                                    other.weights):
            if abs(p_self - p_other) > eps or abs(w_self - w_other) > eps:
                return True
        for i in range(self.dim):
            for j in range(self.dim):
                diff = abs(self.without_weights()[i][j] -
                           other.without_weights()[i][j])
                if diff > eps:
                    return True
        return False

    @property
    def dim(self):
        """Return the dimension of the potential matrix.

        Returns
        -------
        int
            The dimension of the (square) potential matrix.

        """
        return len(self._nodes)

    @property
    def potential_type(self):
        """Return `PotentialType` object for potential.

        Returns
        -------
        PotentialType
            Object with all physics related information for the potential.

        """
        return self._potential_type

    @property
    def nodes(self):
        """Return the nodes for the potential.

        Returns
        -------
        list of floats
            List of momenta at which potential is defined.

        """
        return self._nodes

    @property
    def weights(self):
        """Return weights for the potential.

        Returns
        -------
        list of floats
            Integration weights corresponding to nodes for potential.

        """
        return self._weights

    @property
    def lam(self):
        """Return lambda for potential.

        Returns
        -------
        float
            Value of lambda, the SRG flow parameter, for potential.

        """
        return self._lam

class CoupledPotential(Potential):
    """Representation of potential of coupled channel."""

    def __init__(self, list_of_potentials):  # pylint: disable=too-many-locals
        """Create potential from list of potentials in a coupled channel.

        Parameters
        ----------
        list_of_potentials : list of Potential objects
            List of potentials to form coupled channel.

        Returns
        -------
        Potential
            New potential with full coupled channel.

        """
        self._construction = list_of_potentials
        channels = [x.potential_type.channel for x in list_of_potentials]
        n_body = {x.potential_type.n_body for x in list_of_potentials}
        order = {x.potential_type.order for x in list_of_potentials}
        name = {x.potential_type.name for x in list_of_potentials}
        particles = {x.potential_type.particles for x in list_of_potentials}
        if len(n_body) * len(order) * len(name) * len(particles) != 1:
            raise ValueError('Given potentials cannot be coupled.')
        coupled_channel = CoupledChannel(channels)
        potential_type = PotentialType(n_body.pop(), order.pop(), name.pop(),
                                       coupled_channel, particles.pop())
        lam = {x.lam for x in list_of_potentials}
        if len(lam) != 1:
            raise ValueError('Not all given potentials are at the same lam.')
        lam = lam.pop()
        dim = {x.dim for x in list_of_potentials}
        if len(dim) != 1:
            raise ValueError('Not all given potentials have same dim.')
        dim = dim.pop()
        c_dim = int(sqrt(len(list_of_potentials)))
        if c_dim**2 != len(list_of_potentials):
            raise ValueError('Non-square number of potentials given.')
        nodes = []
        weights = []
        for pot in list_of_potentials[:c_dim]:
            nodes += pot.nodes
            weights += pot.weights
        nodes = np.array(nodes)
        weights = np.array(weights)
        potential_data = np.zeros((c_dim * dim, c_dim * dim))
        self._channel_indexes = []
        for i in range(c_dim):
            for j in range(c_dim):
                r_s = i * dim
                r_e = (i + 1) * dim
                c_s = j * dim
                c_e = (j + 1) * dim
                data = list_of_potentials[i * c_dim + j].without_weights()
                potential_data[r_s:r_e, c_s:c_e] = data
                self._channel_indexes.append((r_s, r_e, c_s, c_e))
        super(CoupledPotential, self).__init__(potential_type, nodes, weights,
                                               potential_data, lam)
        self._c_dim = c_dim
        self._w_dim = dim
        self._channels = channels

    def copy(self, potential, lam):
        """Create potential from current potential with new data and lam.

        Parameters
        ----------
        potential : matrix of floats
            Potential data.
        lam : float
            Value of lambda

        Returns
        -------
        Potential
            New potential with new data.

        """
        new_potentials = []
        for pot, ranges in zip(self._construction, self._channel_indexes):
            sub_matrix = _submatrix(potential, ranges)
            new_potentials.append(pot.copy(sub_matrix, lam))
        return CoupledPotential(new_potentials)

    def reduce_dim(self, dim):
        """Return new potential with only `dim` lowest energy states.

        Parameters
        ----------
        dim : int
            Dimension to which potential is to be reduced.

        Returns
        -------
        Potential
            New reduced dimension potential.

        Raises
        ------
        ValueError
            When value for new dim is too small or too large.

        """
        if dim >= self._w_dim:
            raise ValueError('Value of dim is not smaller than current dim.')
        if dim <= 0:
            raise ValueError('Zero or negative dim is not allowed.')
        new_potentials = []
        for pot, ranges in zip(self._construction, self._channel_indexes):
            sub_matrix = _submatrix(self._potential, ranges)
            new_potentials.append(pot.copy(sub_matrix,
                                           self._lam).reduce_dim(dim))
        return CoupledPotential(new_potentials)

    def extract_channel_potential(self, channel):
        """Return potential corresponding to channel.

        Parameters
        ----------
        channel : Channel
            Channel to extract.

        Returns
        -------
        Potential
            Potential corresponding to channel.

        """
        for chan, potential, ranges in zip(self._channels, self._construction,
                                           self._channel_indexes):
            if channel == chan:
                sub_matrix = _submatrix(self._potential, ranges)
                return potential.copy(sub_matrix, self._lam)
        raise ValueError('Channel not found.')

    @property
    def dim(self):
        """Return the dimension of single channel in the potential matrix.

        Returns
        -------
        int
            The dimension of a single channel in the (square) potential matrix.

        """
        return self._w_dim

# pylint: disable=too-many-locals
def load_from_file(file_str):
    """Load potential from file.

    Parameters
    ----------
    file_str : str
        String path to file with potential data.

    Returns
    -------
    Potential
        Potential created from extracted information and data from file.

    """
    # Parse info about potential from filename
    # Strip directory structure
    end = file_str.split('/')[-1]

    # Match regular expression
    regex_str = r'V(.*)_(.*)_(.*)_SLLJT_(.*)_lambda_(.*)_Np_(.*)_(.*)\.dat'
    result = re.search(regex_str, end)

    # Extract values from matches
    n_body_str = result.group(1)
    order_str = result.group(2)
    name = result.group(3)
    channel_str = result.group(4)
    lam = float(result.group(5))
    particles = result.group(7)

    # Convert string values to integer values
    n_body = NBODY_DICT[n_body_str]
    order = ORDER_DICT[order_str]

    # Convert channel to 5-tuple, then Channel object
    channel = Channel(*tuple([int(n) for n in channel_str]))

    # Get number of points
    num_points = int(result.group(6))

    # Read potential
    with open(file_str) as file:
        nodes = []
        weights = []
        for _ in range(num_points):
            vals = file.readline().split()
            weights.append(float(vals[0]))
            nodes.append(float(vals[1]))
        potential = np.array([[float(file.readline().split()[-1]) for _ in
                               range(num_points)] for _ in range(num_points)])

    # Create potential_type
    potential_type = PotentialType(n_body, order, name, channel, particles)

    # Return potential
    return Potential(potential_type, nodes, weights, potential, lam)

# pylint: disable=too-many-arguments
def load(n_body, order, name, channel, lam, particles, num_points='*'):
    """Load potential based on parameters.

    Parameters
    ----------
     n_body : int
        Number of particles interacting in potential.
    order : int
        Order to which potential was calculated.
    name : str
        Name for potential, may reflect something about origin.
    channel: Channel or (int, int, int, int, int) or str
        Object representing the partial wave channel for the potential.
    lam : float
        Value of SRG flow parameter for potential.
    particles: str
        String representing constituent particles in the interaction.
    num_points : int, optional
        Number of points in potential. Should only be specified if multiple
        versions of same potential are saved and you need a specific one.
        Otherwise, will match the first one in lexicographical ordering.

    Returns
    -------
    Potential
        Potential created from extracted information and data from file.

    Raises
    ------
    FileNotFoundError
        If globbing doesn't match any files.

    """
    # Set up format string
    file_format_str = '{}/V{}_{}_{}_SLLJT_{}_lambda_{:.2f}_Np_{}_{}.dat'

    # Get values for format string
    n_body_str = INV_NBODY_DICT[n_body]
    order_str = INV_ORDER_DICT[order]

    # Handle non-string formats
    if isinstance(channel, Channel):
        channel = str(channel)
    elif isinstance(channel, tuple):
        channel = ''.join(channel)

    dir_str = os.path.join(STANDARD_PATH, n_body_str,
                           'SLLJT_{}'.format(channel))

    # Create full file path string
    file_path = file_format_str.format(dir_str, n_body_str, order_str, name,
                                       channel, lam, num_points, particles)

    # Handle globbing
    if num_points == '*':
        try:
            file_path = glob.glob(file_path)[0]
        except IndexError:
            raise FileNotFoundError('No potential with those params found.')

    return load_from_file(file_path)

def save(potential, dir_str=None):
    """Save potential with correct file-naming.

    Parameters
    ----------
    potential : Potential
        Potential to be saved.
    dir_str : str, optional
        String corresponding to directory where file should be saved. May have
        trailing `/`.

    """
    # Set up format strings
    file_format_str = '{}/V{}_{}_{}_SLLJT_{}_lambda_{:.2f}_Np_{}_{}.dat'
    nodes_format_str = '{:.5e} {:.5e}\n'
    potential_format_str = '{:.5e} {:.5e} {:.5e}\n'

    # Get values for format string
    potential_type = potential.potential_type
    n_body = potential_type.n_body
    n_body_str = INV_NBODY_DICT[n_body]
    order = potential_type.order
    order_str = INV_ORDER_DICT[order]
    name = potential_type.name
    channel_str = str(potential_type.channel)
    lam = potential.lam
    num_points = len(potential.nodes)
    particles = potential_type.particles

    # Handle optional argument
    if dir_str is None:
        dir_str = os.path.join(STANDARD_PATH, n_body_str,
                               'SLLJT_{}'.format(channel_str))

    # Strip potential trailing '/'
    if dir_str[-1] == '/':
        dir_str = dir_str[:-1]

    # Create full file path string
    file_path = file_format_str.format(dir_str, n_body_str, order_str, name,
                                       channel_str, lam, num_points, particles)

    # Create directory if it doesnt exist
    _ensure_dir_for_file(file_path)

    # Output potential
    with open(file_path, 'w+') as file:
        for weight, node in zip(potential.weights, potential.nodes):
            file.write(nodes_format_str.format(weight, node))
        for i in range(num_points):
            for j in range(num_points):
                file.write(potential_format_str.format(
                    potential.nodes[i], potential.nodes[j],
                    potential.without_weights()[i][j]))

def plot(potential, v_min=None, v_max=None):
    """Plot potential with colorbar.

    Parameters
    ----------
    potential : Potential
        Potential to be plotted.
    v_min : int, optional
        Minimum value to be reflected on the colorbar scale.
    v_max : int, optional
        Maximum value to be reflected on the colorbar scale.

    """
    if v_min is None or v_max is None:
        plt.matshow(potential.without_weights())
    else:
        plt.matshow(potential.without_weights(), vmin=v_min, vmax=v_max)
    plt.colorbar()
    plt.show()
    plt.close()

# ------------------- Internal Methods ------------------------------------- #

def _add_w(matrix, weights, nodes):
    factor_vector = [sqrt(w) * p for w, p in zip(weights, nodes)]
    weighted_matrix = np.dot(np.dot(np.diag(factor_vector), matrix),
                             np.diag(factor_vector))
    return 2 / pi * weighted_matrix

def _rem_w(matrix, weights, nodes):
    factor_vector = [1/(sqrt(w) * p) for w, p in zip(weights, nodes)]
    unweighted_matrix = np.dot(np.dot(np.diag(factor_vector), pi / 2 * matrix),
                               np.diag(factor_vector))
    return unweighted_matrix

def _ensure_dir_for_file(file):
    directory = os.path.dirname(file)
    if not os.path.exists(directory):
        os.makedirs(directory)

def _submatrix(potential, ranges):
    return potential[np.ix_(list(range(ranges[0], ranges[1])),
                            list(range(ranges[2], ranges[3])))]


"""
Set of programs and tools to read the outputs from RH (Han's version)
"""
import os
import sys
import io
import xdrlib
import numpy as np

class Rhout:
    """
    Reads outputs from RH.

    Currently the reading the following output files is supported:
     - input.out
     - geometry.out
     - atmos.out
     - spectrum.out (no Stokes)
     - spectrum_XX (no Stokes, from solveray)
     - brs.out
     - J.dat
     - opacity.out (no Stokes)

    These output files are NOT supported:
      - Atom (atom, collrate, damping, pops, radrate)
      - Flux
      - metals
      - molecule

    Parameters
    ----------
    fdir : str, optional
        Directory with output files.
    verbose : str, optional
        If True, will print more details.

    Notes
    -----
    In general,  the way to read all the XDR files should be:

     Modify read_xdr_file so that it returns only xdata.
     Then, on each read_xxx, read the necessary header variables,
     rewind (xdata.set_position(0)), then read the variables in order.
     This allows the flexibility of derived datatypes, and appending to dictionary
     (e.g. as in readatmos for all the elements and etc.). It also allows one to
     read directly into attribute of the class (with setattr(self,'aa',<data>))
    """
    def __init__(self, fdir='.', verbose=True):
        ''' Reads all the output data from a RH run.'''
        self.verbose = verbose
        self.fdir = fdir
        self.read_input('{0}/input.out'.format(fdir))
        self.read_geometry('{0}/geometry.out'.format(fdir))
        self.read_atmosphere('{0}/atmos.out'.format(fdir))
        self.read_spectrum('{0}/spectrum.out'.format(fdir))
        if os.path.isfile('{0}/spectrum_1.00'.format(fdir)):
            self.read_ray('{0}/spectrum_1.00'.format(fdir))

    def read_input(self, infile='input.out'):
        ''' Reads RH input.out file. '''
        data = read_xdr_file(infile)
        self.input = {}
        input_vars = [('magneto_optical', 'i'), ('PRD_angle_dep', 'i'),
                      ('XRD', 'i'), ('start_solution', 'i'),
                      ('stokes_mode', 'i'), ('metallicity', 'd'),
                      ('backgr_pol', 'i'), ('big_endian', 'i')]
        for v in input_vars:
            self.input[v[0]] = read_xdr_var(data, v[1:])
        close_xdr(data, infile, verbose=self.verbose)

    def read_geometry(self, infile='geometry.out'):
        ''' Reads RH geometry.out file. '''
        data = read_xdr_file(infile)
        self.geometry = {}
        geom_type = ['ONE_D_PLANE', 'TWO_D_PLANE',
                     'SPHERICAL_SYMMETRIC', 'THREE_D_PLANE']
        type = read_xdr_var(data, ('i',))
        if type not in list(range(4)):
            raise ValueError('read_geometry: invalid geometry type {0} in {1}'.
                             format(type, infile))
        nrays = read_xdr_var(data, ('i',))
        self.nrays = nrays
        self.geometry_type = geom_type[type]
        # read some parameters and define structure to be read
        if self.geometry_type == 'ONE_D_PLANE':
            ndep = read_xdr_var(data, ('i',))
            self.ndep = ndep
            geom_vars = [('xmu', 'd', (nrays,)), ('wmu', 'd', (nrays,)),
                         ('height', 'd', (ndep,)), ('cmass', 'd', (ndep,)),
                         ('tau500', 'd', (ndep,)), ('vz', 'd', (ndep,))]
        elif self.geometry_type == 'TWO_D_PLANE':
            nx = read_xdr_var(data, ('i',))
            nz = read_xdr_var(data, ('i',))
            self.nx = nx
            self.nz = nz
            geom_vars = [('angleSet', 'i'), ('xmu', 'd', (nrays,)),
                         ('ymu', 'd', (nrays,)), ('wmu', 'd', (nrays,)),
                         ('x', 'd', (nx,)), ('z', 'd', (nz,)),
                         ('vx', 'd', (nx, nz)), ('vz', 'd', (nx, nz))]
        elif self.geometry_type == 'THREE_D_PLANE':
            nx = read_xdr_var(data, ('i',))
            ny = read_xdr_var(data, ('i',))
            nz = read_xdr_var(data, ('i',))
            self.nx = nx
            self.ny = ny
            self.nz = nz
            geom_vars = [('angleSet', 'i'), ('xmu', 'd', (nrays,)),
                         ('ymu', 'd', (nrays,)), ('wmu', 'd', (nrays,)),
                         ('dx', 'd'), ('dy', 'd'),
                         ('z', 'd', (nz,)), ('vx', 'd', (nx, ny, nz)),
                         ('vy', 'd', (nx, ny, nz)), ('vz', 'd', (nx, ny, nz))]
        elif self.geometry_type == 'SPHERICAL_SYMMETRIC':
            nradius = read_xdr_var(data, ('i',))
            ncore = read_xdr_var(data, ('i',))
            self.nradius = nradius
            self.ncore = ncore
            geom_vars = [('radius', 'd'), ('xmu', 'd', (nrays,)),
                         ('wmu', 'd', (nrays,)), ('r', 'd', (nradius,)),
                         ('cmass', 'd', (nradius,)), ('tau500', 'd', (nradius,)),
                         ('vr', 'd', (nradius,))]
        # read data
        for v in geom_vars:
            self.geometry[v[0]] = read_xdr_var(data, v[1:])
        close_xdr(data, infile, verbose=self.verbose)

    def read_atmosphere(self, infile='atmos.out'):
        ''' Reads RH atmos.out file '''
        if not hasattr(self, 'geometry'):
            em = ('read_atmosphere: geometry data not loaded, '
                  'call read_geometry() first!')
            raise ValueError(em)
        data = read_xdr_file(infile)
        self.atmos = {}
        nhydr = read_xdr_var(data, ('i',))
        nelem = read_xdr_var(data, ('i',))
        self.atmos['nhydr'] = nhydr
        self.atmos['nelem'] = nelem
        # read some parameters and define structure to be read
        if self.geometry_type == 'ONE_D_PLANE':
            ndep = self.ndep
            atmos_vars = [('moving', 'i'), ('T', 'd', (ndep,)),
                          ('n_elec', 'd', (ndep,)), ('vturb', 'd', (ndep,)),
                          ('nh', 'd', (ndep, nhydr)), ('id', 's')]
        elif self.geometry_type == 'TWO_D_PLANE':
            nx, nz = self.nx, self.nz
            atmos_vars = [('moving', 'i'), ('T', 'd', (nx, nz)),
                          ('n_elec', 'd', (nx, nz)), ('vturb', 'd', (nx, nz)),
                          ('nh', 'd', (nx, nz, nhydr)), ('id', 's')]
        elif self.geometry_type == 'THREE_D_PLANE':
            nx, ny, nz = self.nx, self.ny, self.nz
            atmos_vars = [('moving', 'i'), ('T', 'd', (nx, ny, nz)),
                          ('n_elec', 'd', (nx, ny, nz)
                           ), ('vturb', 'd', (nx, ny, nz)),
                          ('nh', 'd', (nx, ny, nz, nhydr)), ('id', 's')]
        elif self.geometry_type == 'SPHERICAL_SYMMETRIC':
            nradius = self.nradius
            atmos_vars = [('moving', 'i'), ('T', 'd', (nradius,)),
                          ('n_elec', 'd', (nradius,)), ('vturb', 'd', (nradius,)),
                          ('nh', 'd', (nradius, nhydr)), ('id', 's')]
        # read data
        for v in atmos_vars:
            self.atmos[v[0]] = read_xdr_var(data, v[1:])
        # read elements into nested dictionaries
        self.elements = {}
        for v in range(nelem):
            el = read_xdr_var(data, ('s',)).strip()
            weight = read_xdr_var(data, ('d',))
            abund = read_xdr_var(data, ('d',))
            self.elements[el] = {'weight': weight, 'abund': abund}
        # read stokes data, if present
        self.stokes = False
        if self.geometry_type != 'SPHERICAL_SYMMETRIC':
            try:
                stokes = read_xdr_var(data, ('i',))
            except EOFError or IOError:
                if self.verbose:
                    print('(WWW) read_atmos: no Stokes data in atmos.out,'
                          ' skipping.')
                return
            self.stokes = True
            ss = self.atmos['T'].shape
            stokes_vars = [('B', 'd', ss), ('gamma_B', 'd', ss),
                           ('chi_B', 'd', ss)]
            for v in stokes_vars:
                self.atmos[v[0]] = read_xdr_var(data, v[1:])
        close_xdr(data, infile, verbose=self.verbose)

    def read_spectrum(self, infile='spectrum.out'):
        ''' Reads RH spectrum.out file '''
        if not hasattr(self, 'geometry'):
            em = ('read_spectrum: geometry data not loaded, '
                  'call read_geometry() first!')
            raise ValueError(em)
        if not hasattr(self, 'atmos'):
            em = ('read_spectrum: atmos data not loaded, '
                  'call read_atmos() first!')
            raise ValueError(em)
        data = read_xdr_file(infile)
        profs = {}
        self.spec = {}
        nspect = read_xdr_var(data, ('i',))
        self.spec['nspect'] = nspect
        nrays = self.nrays
        self.wave = read_xdr_var(data, ('d', (nspect,)))
        if self.geometry_type == 'ONE_D_PLANE':
            ishape = (nrays, nspect)
        elif self.geometry_type == 'TWO_D_PLANE':
            ishape = (self.nx, nrays, nspect)
        elif self.geometry_type == 'THREE_D_PLANE':
            ishape = (self.nx, self.ny, nrays, nspect)
        elif self.geometry_type == 'SPHERICAL_SYMMETRIC':
            ishape = (nrays, nspect)
        self.imu = read_xdr_var(data, ('d', ishape))
        self.spec['vacuum_to_air'] = read_xdr_var(data, ('i',))
        self.spec['air_limit'] = read_xdr_var(data, ('d',))
        if self.stokes:
            self.stokes_Q = read_xdr_var(data, ('d', ishape))
            self.stokes_U = read_xdr_var(data, ('d', ishape))
            self.stokes_V = read_xdr_var(data, ('d', ishape))
        close_xdr(data, infile, verbose=self.verbose)
        # read as_rn, if it exists
        if os.path.isfile('asrs.out'):
            data = read_xdr_file('asrs.out')
            if self.atmos['moving'] or self.stokes or self.input['PRD_angle_dep']:
                self.spec['as_rn'] = read_xdr_var(data, ('i', (nrays, nspect)))
            else:
                self.spec['as_rn'] = read_xdr_var(data, ('i', (nspect,)))
            close_xdr(data, 'asrs.out', verbose=self.verbose)

    def read_ray(self, infile='spectrum_1.00'):
        ''' Reads spectra for single ray files (e.g. mu=1). '''
        if not hasattr(self, 'geometry'):
            em = ('read_spectrum: geometry data not loaded,'
                  ' call read_geometry() first!')
            raise ValueError(em)
        if not hasattr(self, 'spec'):
            em = ('read_spectrum: spectral data not loaded, '
                  'call read_spectrum() first!')
            raise ValueError(em)
        data = read_xdr_file(infile)
        nspect = self.spec['nspect']
        self.ray = {}
        if self.geometry_type == 'ONE_D_PLANE':
            self.muz = read_xdr_var(data, ('d',))
            ishape = (nspect,)
            sshape = (self.ndep,)
        elif self.geometry_type == 'TWO_D_PLANE':
            self.mux = read_xdr_var(data, ('d',))
            self.muz = read_xdr_var(data, ('d',))
            ishape = (self.nx, nspect)
            sshape = (self.nx, self.nz)
        elif self.geometry_type == 'THREE_D_PLANE':
            self.mux = read_xdr_var(data, ('d',))
            self.muy = read_xdr_var(data, ('d',))
            ishape = (self.nx, self.ny, nspect)
            sshape = (self.nx, self.ny, self.nz)
        elif self.geometry_type == 'SPHERICAL_SYMMETRIC':
            self.muz = read_xdr_var(data, ('d',))
            ishape = (nspect,)
            sshape = (self.nradius,)
        # read intensity
        self.int = read_xdr_var(data, ('d', ishape))
        # read absorption and source function if written
        ns = read_xdr_var(data, ('i',))
        if ns > 0:
            nshape = (ns,) + sshape
            self.ray['chi'] = np.zeros(nshape, dtype='d')
            self.ray['S'] = np.zeros(nshape, dtype='d')
            self.ray['wave_idx'] = np.zeros(ns, dtype='l')
            for i in range(ns):
                self.ray['wave_idx'][i] = read_xdr_var(data, ('i',))
                self.ray['chi'][i] = read_xdr_var(data, ('d', sshape))
                self.ray['S'][i] = read_xdr_var(data, ('d', sshape))
        if self.stokes:
            self.ray_stokes_Q = read_xdr_var(data, ('d', ishape))
            self.ray_stokes_U = read_xdr_var(data, ('d', ishape))
            self.ray_stokes_V = read_xdr_var(data, ('d', ishape))
        close_xdr(data, infile, verbose=self.verbose)

    def read_brs(self, infile='brs.out'):
        ''' Reads the file with the background opacity record settings,
            in the old (xdr) format. '''
        if not hasattr(self, 'geometry'):
            em = ('read_brs: geometry data not loaded, call read_geometry()'
                  ' first!')
            raise ValueError(em)
        if not hasattr(self, 'spec'):
            em = ('read_brs: spectrum data not loaded, call read_spectrum()'
                  ' first!')
            raise ValueError(em)
        data = read_xdr_file(infile)
        atmosID = read_xdr_var(data, ('s',)).strip()
        nspace = read_xdr_var(data, ('i',))
        nspect = read_xdr_var(data, ('i',))
        if nspect != self.spec['nspect']:
            em = ('(EEE) read_brs: nspect in file different from atmos. '
                  'Aborting.')
            raise ValueError(em)
        self.brs = {}
        if self.atmos['moving'] or self.stokes:
            ishape = (2, self.nrays, nspect)
        else:
            ishape = (nspect,)
        self.brs['hasline'] = read_xdr_var(
            data, ('i', (nspect,))).astype('Bool')
        self.brs['ispolarized'] = read_xdr_var(
            data, ('i', (nspect,))).astype('Bool')
        self.brs['backgrrecno'] = read_xdr_var(data, ('i', ishape))
        close_xdr(data, infile, verbose=self.verbose)

    def read_j(self, infile='J.dat'):
        ''' Reads the mean radiation field, for all wavelengths. '''
        if not hasattr(self, 'geometry'):
            em = 'read_j: geometry data not loaded, call read_geometry() first!'
            raise ValueError(em)
        if not hasattr(self, 'spec'):
            em = 'read_j: spectrum data not loaded, call read_spec() first!'
            raise ValueError(em)
        data_file = open(infile, 'r')
        nspect = self.spec['nspect']
        if self.geometry_type == 'ONE_D_PLANE':
            rec_len = self.ndep * 8
            ishape = (nspect, self.ndep)
        elif self.geometry_type == 'TWO_D_PLANE':
            rec_len = (self.nx * self.nz) * 8
            ishape = (nspect, self.nx, self.nz)
        elif self.geometry_type == 'THREE_D_PLANE':
            rec_len = (self.nx * self.ny * self.nz) * 8
            ishape = (nspect, self.nx, self.ny, self.nz)
        elif self.geometry_type == 'SPHERICAL_SYMMETRIC':
            rec_len = self.nradius * 8
            ishape = (nspect, self.nradius)
        self.J = np.zeros(ishape)
        for i in range(nspect):
            # point background file to position and read
            data_file.seek(i * rec_len)
            self.J[i] = read_file_var(data_file, ('d', ishape[1:]))
        data_file.close()

    def read_opacity(self, infile_line='opacity.out', infile_bg='background.dat',
                     imu=0):
        ''' Reads RH atmos.out file '''
        if not hasattr(self, 'geometry'):
            em = ('read_opacity: geometry data not loaded,'
                  ' call read_geometry() first!')
            raise ValueError(em)
        if not hasattr(self, 'spec'):
            em = ('read_opacity: spectrum data not loaded,'
                  ' call read_spec() first!')
            raise ValueError(em)
        if not hasattr(self.atmos, 'brs'):
            self.read_brs()
        data_line = read_xdr_file(infile_line)
        file_bg = open(infile_bg, 'r')
        nspect = self.spec['nspect']
        if self.geometry_type == 'ONE_D_PLANE':
            as_rec_len = 2 * self.ndep * 8
            bg_rec_len = self.ndep * 8
            ishape = (nspect, self.ndep)
        elif self.geometry_type == 'TWO_D_PLANE':
            as_rec_len = 2 * (self.nx * self.nz) * 8
            bg_rec_len = (self.nx * self.nz) * 8
            ishape = (nspect, self.nx, self.nz)
        elif self.geometry_type == 'THREE_D_PLANE':
            as_rec_len = 2 * (self.nx * self.ny * self.nz) * 8
            bg_rec_len = (self.nx * self.ny * self.nz) * 8
            ishape = (nspect, self.nx, self.ny, self.nz)
        elif self.geometry_type == 'SPHERICAL_SYMMETRIC':
            as_rec_len = 2 * self.nradius * 8
            bg_rec_len = self.nradius * 8
            ishape = (nspect, self.nradius)
        # create arrays
        chi_as = np.zeros(ishape)
        eta_as = np.zeros(ishape)
        chi_c = np.zeros(ishape)
        eta_c = np.zeros(ishape)
        scatt = np.zeros(ishape)
        # NOTE: this will not work when a line is polarised.
        #       For those cases these arrays must be read per wavelength, and will
        #       have different sizes for different wavelengths.
        if np.sum(self.brs['ispolarized']):
            em = ('read_opacity: Polarized line(s) detected, cannot continue'
                  ' with opacity extraction')
            raise ValueError(em)
        # get record numbers
        if self.atmos['moving'] or self.stokes or self.input['PRD_angle_dep']:
            as_index = self.spec['as_rn'][imu] * as_rec_len
            bg_index = self.brs['backgrrecno'][1, imu] * bg_rec_len
        else:
            as_index = self.spec['as_rn'] * as_rec_len
            bg_index = self.brs['backgrrecno'] * bg_rec_len
        # Read arrays
        for i in range(nspect):
            if as_index[i] >= 0:  # avoid non-active set lines
                # point xdr buffer to position and read
                data_line.set_position(as_index[i])
                chi_as[i] = read_xdr_var(data_line, ('d', ishape[1:]))
                eta_as[i] = read_xdr_var(data_line, ('d', ishape[1:]))
            # point background file to position and read
            file_bg.seek(bg_index[i])
            chi_c[i] = read_file_var(file_bg, ('d', ishape[1:]))
            eta_c[i] = read_file_var(file_bg, ('d', ishape[1:]))
            scatt[i] = read_file_var(file_bg, ('d', ishape[1:]))
        self.chi_as = chi_as
        self.eta_as = eta_as
        self.chi_c = chi_c
        self.eta_c = eta_c
        self.scatt = scatt
        close_xdr(data_line, infile_line, verbose=False)
        file_bg.close()

    def get_contrib_imu(self, imu, type='total', op_file='opacity.out',
                        bg_file='background.dat', j_file='J.dat'):
        ''' Calculates the contribution function for intensity, for a
            particular ray, defined by imu.

            type can be: \'total\', \'line, or \'continuum\'

            The units of self.contribi are J m^-2 s^-1 Hz^-1 sr^-1 km^-1

            NOTE: This only calculates the contribution function for
                  the quadrature rays (ie, often not for disk-centre)
                  For rays calculated with solve ray, one must use
                  get_contrib_ray

        '''
        type = type.lower()
        if not hasattr(self, 'geometry'):
            em = ('get_contrib_imu: geometry data not loaded,'
                  ' call read_geometry() first!')
            raise ValueError(em)
        if not hasattr(self, 'spec'):
            em = ('get_contrib_imu: spectrum data not loaded,'
                  ' call read_spec() first!')
            raise ValueError(em)
        self.read_opacity(infile_line=op_file, infile_bg=bg_file, imu=imu)
        self.read_j(infile=j_file)
        mu = self.geometry['xmu'][imu]
        # Calculate optical depth
        ab = (self.chi_c + self.chi_as)
        self.tau = get_tau(self.geometry['height'], mu, ab)
        # Calculate source function
        if type == 'total':
            self.S = (self.eta_as + self.eta_c + self.J * self.scatt) / ab
        elif type == 'line':
            self.S = self.eta_as / ab
        elif type == 'continuum':
            self.S = (self.eta_c + self.J * self.scatt) / ab
        else:
            raise ValueError('get_contrib_imu: invalid type!')
        # Calculate contribution function
        self.contribi = get_contrib(
            self.geometry['height'], mu, self.tau, self.S)
        return

    def get_contrib_ray(self, inray='ray.input', rayfile='spectrum_1.00'):
        ''' Calculates the contribution function for intensity, for a
            particular ray

            The units of self.contrib are J m^-2 s^-1 Hz^-1 sr^-1 km^-1
        '''
        inray = self.fdir + '/' + inray
        rayfile = self.fdir + '/' + rayfile
        if not hasattr(self, 'ray'):
            self.read_ray(infile=rayfile)
        if 'wave_idx' not in list(self.ray.keys()):
            em = ('get_contrib_ray: no chi/source function written to '
                  'ray file, aborting.')
            raise ValueError(em)
        # read mu from ray.input file
        mu = np.loadtxt(inray, dtype='f')[0]
        if not (0 <= mu <= 1.):
            em = 'get_contrib_ray: invalid mu read: %f' % mu
            raise ValueError(em)
        idx = self.ray['wave_idx']
        # Calculate optical depth
        self.tau = get_tau(self.geometry['height'], mu, self.ray['chi'])
        # Calculate contribution function
        self.contrib = get_contrib(self.geometry['height'], mu, self.tau,
                                   self.ray['S'])
        return

class RhAtmos:
    """
    Reads input atmosphere from RH. Currently only 2D format supported.

    Parameters
    ----------
    format : str, optional
        Atmosphere format. Currently only '2D' (default) supported.
    filename : str, optional
        File to read.
    verbose : str, optional
        If True, will print more details.
    """
    def __init__(self, format="2D", filename=None, verbose=True):
        ''' Reads RH input atmospheres. '''
        self.verbose = verbose
        if format.lower() == "2d":
            if filename is not None:
                self.read_atmos2d(filename)
        else:
            raise NotImplementedError("Format %s not yet supported" % format)

    def read_atmos2d(self, filename):
        """
        Reads input 2D atmosphere
        """
        data = read_xdr_file(filename)
        self.nx = read_xdr_var(data, ('i',))
        self.nz = read_xdr_var(data, ('i',))
        self.nhydr = read_xdr_var(data, ('i',))
        self.hboundary = read_xdr_var(data, ('i',))
        self.bvalue = read_xdr_var(data, ('i', (2, )))
        nx, nz, nhydr = self.nx, self.nz, self.nhydr
        atmos_vars = [('dx', 'd', (nx,)), ('z', 'd', (nz,)),
                      ('T', 'd', (nx, nz)), ('ne', 'd', (nx, nz)),
                      ('vturb', 'd', (nx, nz)), ('vx', 'd', (nx, nz)),
                      ('vz', 'd', (nx, nz)), ('nh', 'd', (nx, nz, nhydr))
                      ]
        for v in atmos_vars:
            setattr(self, v[0], read_xdr_var(data, v[1:]))

    def write_atmos2d(self, filename, dx, z, T, ne, vturb, vx, vz, nh,
                      hboundary, bvalue):
        nx, nz = T.shape
        nhydr = nh.shape[-1]
        assert T.shape == ne.shape
        assert ne.shape == vturb.shape
        assert vturb.shape == nh.shape[:-1]
        assert dx.shape[0] == nx
        assert z.shape[0] == nz
        # Pack as double
        p = xdrlib.Packer()
        p.pack_int(nx)
        p.pack_int(nz)
        p.pack_int(nhydr)
        p.pack_int(hboundary)
        p.pack_int(bvalue[0])
        p.pack_int(bvalue[1])
        p.pack_farray(nx, dx.ravel().astype('d'), p.pack_double)
        p.pack_farray(nz, z.ravel().astype('d'), p.pack_double)
        p.pack_farray(nx * nz, T.ravel().astype('d'), p.pack_double)
        p.pack_farray(nx * nz, ne.ravel().astype('d'), p.pack_double)
        p.pack_farray(nx * nz, vturb.ravel().astype('d'), p.pack_double)
        p.pack_farray(nx * nz, vx.ravel().astype('d'), p.pack_double)
        p.pack_farray(nx * nz, vz.ravel().astype('d'), p.pack_double)
        p.pack_farray(nx * nz * nhydr, nh.T.ravel().astype('d'), p.pack_double)
        # Write to file
        f = open(filename, 'wb')
        f.write(p.get_buffer())
        f.close()

#############################################################################
# TOOLS
#############################################################################
class EmptyData:
    def __init__(self):
        pass

def read_xdr_file(filename):  # ,var,cl=None,verbose=False):
    """
    Reads data from XDR file.

    Because of the way xdrlib works, this reads the whole file to
    memory at once. Avoid with very large files.

    Parameters
    ----------
    filename : string
        File to read.

    Returns
    -------
    result   : xdrlib.Unpacker object
    """
    try:
        f = io.open(filename, 'rb')
        data = f.read()
        f.close()
    except IOError as e:
        raise IOError(
            'read_xdr_file: problem reading {0}: {1}'.format(filename, e))
    # return XDR data
    return xdrlib.Unpacker(data)

def close_xdr(buf, ofile='', verbose=False):
    """
    Closes the xdrlib.Unpacker object, gives warning if not all data read.

    Parameters
    ----------
    buf : xdrlib.Unpacker object
        data object.
    ofile : string, optional
        Original file from which data was read.
    verbose : bool, optional
        Whether to print warning or not.
    """
    try:
        buf.done()
    except:  # .done() will raise error if data remaining
        if verbose:
            print(('(WWW) close_xdr: {0} not all data read!'.format(ofile)))

def read_xdr_var(buf, var):
    """
    Reads a single variable/array from a xdrlib.Unpack buffer.

    Parameters
    ----------

    buf:  xdrlib.Unpack object
        Data buffer.
    var: tuple with (type[,shape]), where type is 'f', 'd', 'i', 'ui',
             or 's'. Shape is optional, and if true is shape of array.
        Type and shape of variable to read

    Returns
    -------
    out :  int/float or array
        Resulting variable.
    """
    assert len(var) > 0
    if var[0] not in ['f', 'd', 'i', 'ui', 's']:
        raise ValueError('read_xdr_var: data type'
                         ' {0} not currently supported'.format(var[0]))
    fdict = {'f': buf.unpack_float,
             'd': buf.unpack_double,
             'i': buf.unpack_int,
             'ui': buf.unpack_uint,
             's': buf.unpack_string}
    func = fdict[var[0]]
    # Single or array?
    if len(var) == 1:
        # this is because RH seems to write the size of the string twice
        if var[0] == 's':
            buf.unpack_int()
        out = func()
    else:
        nitems = np.prod(var[1])
        out = np.array(buf.unpack_farray(nitems, func)).reshape(var[1][::-1])
        # invert order of indices, to match IDL's
        out = np.transpose(out, list(range(len(var[1])))[::-1])
    return out

def read_file_var(buf, var):
    ''' Reads a single variable/array from a file buffer.

    IN:
       buf:  open file object
       var:  tuple with (type[,shape]), where type is 'f', 'd', 'i', 'ui',
             or 's'. Shape is optional, and if true is shape of array.
    OUT:
       variable/array

    '''
    assert len(var) > 0
    if len(var) == 1:
        out = np.fromfile(buf, dtype=var, count=1)
    elif len(var) == 2:
        out = np.fromfile(buf, dtype=var[0], count=var[1][0])
    else:
        nitems = np.prod(var[1])
        out = np.array(np.fromfile(buf, dtype=var[0], count=nitems)).\
            reshape(var[1][::-1])
        out = np.transpose(out, list(range(len(var[1])))[::-1])
    return out

def get_tau(x, mu, chi):
    ''' Calculates the optical depth, given x (height), mu (cos[theta]) and
        chi, absorption coefficient. Chi can be n-dimensional, as long as
        last index is depth.
    '''
    # With scipy, this could be done in one line with
    # scipy.integrate.quadrature.cumtrapz, but we are avoiding scipy to keep
    # these tools more independent
    if len(x) != chi.shape[-1]:
        raise ValueError('get_tau: x and chi have different sizes!')
    path = x / mu
    npts = len(x)
    # bring depth to first index, to allow n-d algebra
    chi_t = np.transpose(chi)
    tau = np.zeros(chi_t.shape)
    for i in range(1, npts):
        tau[i] = tau[i - 1] + 0.5 * \
            (chi_t[i - 1] + chi_t[i]) * (path[i - 1] - path[i])
    return tau.T

def get_contrib(z, mu, tau_in, S):
    ''' Calculates contribution function using x, mu, tau, and the source
        function. '''
    # Tau truncated at 100 (large enough to be useless)
    tau = tau_in.copy()
    tau[tau_in > 100.] = 100.
    # Calculate dtau (transpose to keep n-D generic form), and dx
    dtau = np.zeros(tau_in.shape[::-1])
    tt = np.transpose(tau_in)
    dtau[1:] = tt[1:] - tt[:-1]
    dtau = np.transpose(dtau)
    dx = np.zeros(z.shape)
    dx[1:] = (z[1:] - z[:-1]) / mu
    dx[0] = dx[1]
    # Calculate contribution function
    contrib = S * np.exp(-tau) * (- dtau / dx) / mu
    # convert from m^-1 to km^-1, units are now: J m^-2 s^-1 Hz^-1 sr^-1 km^-1
    contrib *= 1.e3
    return contrib

def write_B(outfile, Bx, By, Bz):
    ''' Writes a RH magnetic field file. Input B arrays can be any rank, as
        they will be flattened before write. Bx, By, Bz units should be T.'''
    if (Bx.shape != By.shape) or (By.shape != Bz.shape):
        raise TypeError('writeB: B arrays have different shapes!')
    n = np.prod(Bx.shape)
    # Convert into spherical coordinates
    B = np.sqrt(Bx**2 + By**2 + Bz**2)
    gamma_B = np.arccos(Bz / B)
    chi_B = np.arctan(By / Bx)
    # Pack as double
    p = xdrlib.Packer()
    p.pack_farray(n, B.ravel().astype('d'), p.pack_double)
    p.pack_farray(n, gamma_B.ravel().astype('d'), p.pack_double)
    p.pack_farray(n, chi_B.ravel().astype('d'), p.pack_double)
    # Write to file
    f = open(outfile, 'wb')
    f.write(p.get_buffer())
    f.close()
    return


#!/usr/bin/env python
#
#  A Python package for temporal consistency and scheduling.
#
#  Copyright (c) 2015 MIT. All rights reserved.
#
#   author: Pedro Santana
#   e-mail: psantana@mit.edu
#   website: people.csail.mit.edu/psantana
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions
#  are met:
#
#  1. Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.
#  2. Redistributions in binary form must reproduce the above copyright
#     notice, this list of conditions and the following disclaimer in
#     the documentation and/or other materials provided with the
#     distribution.
#  3. Neither the name(s) of the copyright holders nor the names of its
#     contributors or of the Massachusetts Institute of Technology may be
#     used to endorse or promote products derived from this software
#     without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
#  FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
#  COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
#  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
#  BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
#  OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
#  AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
#  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
#  ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
#  POSSIBILITY OF SUCH DAMAGE.
"""
pytemporal: A Python package for temporal consistency and scheduling.

Chance-constrained, strong controllability checker implemented using Gurobi

@author: Pedro Santana (psantana@mit.edu).
"""
from gurobipy import *
import numpy as np
import scipy.stats as st
from .gaussian_piecewise import get_gaussian_partition

class PARIS(object):
    """
    Class encapsulating the different functionality of PARIS, the Polynomial-time,
    RIsk-sensitive Scheduler.
    """
    def __init__(self,gaussian_div=5,gaussian_optimize_partition=False,
                 gaussian_lr=0.05,gaussian_tol=1e-5,gaussian_max_iter=10000,
                 gaussian_lb_sigma=None,gaussian_ub_sigma=None,
                 random_part_init=False,verbose=False):
        self.gaussian_div = gaussian_div
        self.gaussian_optimize_partition = gaussian_optimize_partition
        self.gaussian_lr = gaussian_lr
        self.gaussian_tol = gaussian_tol
        self.gaussian_max_iter = gaussian_max_iter
        self.gaussian_lb_sigma = gaussian_lb_sigma
        self.gaussian_ub_sigma = gaussian_ub_sigma
        self.random_part_init = random_part_init
        self.verbose = verbose

    def init_model(self,):
        """
        Initializes an empty optimization model
        """
        self.model = Model('paris_model')
        self.model.setParam('OutputFlag',1 if self.verbose else 0)
        self.modelSense = GRB.MINIMIZE #Minimize risk of STNU reformulation

    def schedule(self,tcs,makespan=False,cc=-1.0,lin_sched_obj=None,reset=True):
        """
        Automates the process of calling the different forms of the PARIS module.
        """
        if reset:
            self.init_model()

        if self.set_minimum_risk_strong_controllability(tcs):
            #Objective function is makespan
            if makespan:
                self.set_makespan_optimization()
            else:
                #Objective function has been specified by the user
                if lin_sched_obj != None:
                    self.set_linear_schedule_objective(lin_sched_obj)
            #Add a chance constraint
            if cc>=0.0:
                self.set_chance_constraint(cc)
            return self.solve()
        else:
            None,None,None

    def stnu_reformulation(self,pstnu_or_prog=None,makespan=False,cc=-1.0,digits=3):
        """
        Takes a PSTNU or RMPyL program and performs the appropriate STNU reformulation.
        It modifies the original object!
        """
        #The input is a list of temporal constraints (PSTNU)
        if isinstance(pstnu_or_prog,list):
            tcs = pstnu_or_prog
        #The input is an RMPyL program
        else:
            if len(pstnu_or_prog.choices)>0:
                print('\nWARNING: PARIS should not be used for conditional scheduling!')
            tcs = pstnu_or_prog.temporal_constraints

        squeeze_dict,risk_bound,sc_schedule = self.schedule(tcs,makespan,cc)
        if squeeze_dict != None:
            for tc,tc_dict in squeeze_dict.items():
                tc.set_stcu(round(tc_dict['lb'],digits),round(tc_dict['ub'],digits))
            return risk_bound,sc_schedule
        else:
            print('The %s is not strongly controllable'%('PSTNU' if isinstance(pstnu_or_prog,list) else 'RMPyL program'))
            return None,None

    def set_minimum_risk_strong_controllability(self,tcs):
        """
        Sets up the strong controllability problem for risk minimization.
        """
        #Extracts the temporal reference structure from the pSTPN, while adding
        #the required squeezing variables for probabilistic durations.
        self.sc_reform={'_baseline_risk_':0.0}
        self.tc_squeeze_var_dict={}
        requirements=[]
        for tc in tcs:
            if tc.type in ['uncontrollable_probabilistic','uncontrollable_bounded']:
                if tc.end.name in self.sc_reform:
                    #NOTE:some Zipcar tests had uncontrollable durations sharing
                    #end points, so I had to convert the exception into a failure in
                    #order to be able to run all the tests.
                    print('An event cannot be the end point of two uncontrollable durations.')
                    return False
                else:
                    self.sc_reform[tc.end.name]={'ref':tc.start.name,'lb':tc.lb,'ub':tc.ub}

                if tc.type=='uncontrollable_probabilistic':
                    dist_type=tc.distribution['type']
                    if dist_type=='uniform':
                        lb,ub=tc.distribution['lb'],tc.distribution['ub']
                        uniform_sc_reformulation(tc,lb,ub,self.model,self.sc_reform,
                                                 self.tc_squeeze_var_dict)
                    elif dist_type=='gaussian':
                        #FIXME:some Zipcar tests had distributions with variance equal to 0.0,
                        #which cause errors in the piecewise approximation. This is clearly
                        #inconsistent.
                        if tc.distribution['variance']>0.0:
                            smart_piecewise_gaussian_sc_reformulation(tc,self.model,self.sc_reform,
                                                                      self.tc_squeeze_var_dict,
                                                                      num_div=self.gaussian_div,
                                                                      optimize_partition=self.gaussian_optimize_partition,
                                                                      lr=self.gaussian_lr,
                                                                      tol=self.gaussian_tol,
                                                                      max_iter=self.gaussian_max_iter,
                                                                      gaussian_lb_sigma=self.gaussian_lb_sigma,
                                                                      gaussian_ub_sigma=self.gaussian_ub_sigma,
                                                                      random_part_init=self.random_part_init)
                    else:
                        raise TypeError('Gurobi only supports uniform and Gaussian distributions at the moment.')

            else:#Controllable duration
                requirements.append(tc)

        #Adds baseline risk as a constant to the objective (allows correct chance constraints)
        self.model.setObjective(self.model.getObjective()+self.sc_reform['_baseline_risk_'],GRB.MINIMIZE)
        self.model.update() # Integrate new variables

        #For strong controllability, only controllable temporal events are variables
        #in the optimization model. All uncontrollable time points should be replaced
        #by their maximum and minimum values.
        self.var_map={}; path_cache={}
        for tc in requirements:
            #Minimal and maximal values of temporal distances
            dv1 = ctg_path(tc.start.name,self.sc_reform,path_cache,self.var_map,
                           self.model,set())
            dv2 = ctg_path(tc.end.name,self.sc_reform,path_cache,self.var_map,
                           self.model,set())
            df_min,df_max = diff_min_max(dv2,dv1,self.sc_reform,self.var_map)

            #Strongly controllable reformulation
            self.model.addConstr(df_min>=tc.lb)
            self.model.addConstr(df_max<=tc.ub)

        self.risk_bound = self.model.getObjective() #Gets the risk bound expression
        self.model.update()

        return True

    def set_makespan_optimization(self,tight_risk=True,mult=1000.0):
        """
        Adds makespan to the objective and constraints.
        """
        #Surrogate objective that ensures a tight risk bound with approximate
        #makespan
        if tight_risk:
            M = len(self.sc_reform)*mult
            t = self.model.addVar(vtype=GRB.CONTINUOUS,lb=0.0,obj=M) #Makespan
            self.model.update()
        else: #Optimizes makespan directly
            t = self.model.addVar(vtype=GRB.CONTINUOUS,lb=0.0) #Makespan
            self.model.update()
            self.model.setObjective(t,GRB.MINIMIZE)

        for ec in self.var_map.values():
            self.model.addConstr(t-ec>=0.0) #All controllable events should be less

        self.model.update()
        return t

    def set_linear_schedule_objective(self,lin_sched_obj):
        """
        Sets a generic linear objective over the controllable schedule variables.
        """
        valid_objective=True
        if isinstance(lin_sched_obj,dict):
            for field in ['events','coefficients','maximize']:
                if not field in lin_sched_obj:
                    valid_objective=False
                    print('\nERROR: missing field '+field+'\n')
                    break
        else:
            valid_objective=False

        if valid_objective:
            #Extracts the names of the events, whether they are given as event
            #objects or their names directly.
            ec_names = [e if isinstance(e,str) else e.name for e in lin_sched_obj['events']]

            #Constructs the objective
            obj = quicksum([lin_sched_obj['coefficients'][i]*self.var_map[n] for i,n in enumerate(ec_names)])

            #Updates the model
            self.model.setObjective(obj,GRB.MAXIMIZE if lin_sched_obj['maximize'] else GRB.MINIMIZE)
            self.model.update()
        else:
            print('\nERROR: Invalid specification of linear schedule objective as dictionary.')
            print('\t"events": list of event objects or names involved in the objective')
            print('\t"coefficients": list of corresponding coefficients for the events in the objective.')
            print('\t"maximize": whether to maximize (TRUE) or minimize (FALSE).')
            obj = None

        return obj

    def set_chance_constraint(self,theta):
        """
        Adds a chance constraint to the model.
        """
        cc = self.model.addConstr(self.risk_bound<=theta)
        return cc

    def solve(self):
        """
        Solves the internal optimization problem and returns the solution.
        """
        self.model.optimize()#Solves LP with SC constraints

        if self.model.status == GRB.status.OPTIMAL: #Found optimal squeezing
            #Returns amount by which temporal constraints should be squeezed,
            #but only for the ones with non-zero (or very close to zero)
            #squeezing
            squeeze_dict={}
            for tc,tc_dict in self.tc_squeeze_var_dict.items():
                new_lb_ub_bounds=[]
                for bound in [tc_dict['lb'],tc_dict['ub']]:
                    if isinstance(bound,gurobipy.LinExpr):
                        new_lb_ub_bounds.append(bound.getValue())
                    elif isinstance(bound,gurobipy.Var):
                        new_lb_ub_bounds.append(bound.X)
                    else:
                        new_lb_ub_bounds.append(bound)
                risk= tc_dict['risk_function'](tc,new_lb_ub_bounds[0],new_lb_ub_bounds[1])
                squeeze_dict[tc]={'lb':new_lb_ub_bounds[0],'ub':new_lb_ub_bounds[1],'risk':risk}

            #One possible strongly controllable schedule
            sc_schedule={event_id:var.X for event_id,var in self.var_map.items()}

            # Returns the objective with the constant _baseline_risk_ added to it already
            return squeeze_dict,self.risk_bound.getValue(),sc_schedule
            # return squeeze_dict,sc_model.ObjVal+sc_reform['_baseline_risk_'],sc_schedule

        else: #No relaxation found
            return None,None,None

def ctg_path(event_name,sc_reform,path_cache,var_map,model,prev_events):
    """
    Recursively computes the controllable and contigent events that influence
    the schedule of a given event.
    """
    if event_name in path_cache:#If solution has been already computed, use it
        return path_cache[event_name]
    else:
        if event_name in sc_reform: #End point of uncontrollable duration
            if event_name in prev_events:
                raise RuntimeError('Contigent duration loop detected!')
            else:
                prev_events.add(event_name)
            path_ref = ctg_path(sc_reform[event_name]['ref'],sc_reform,path_cache,var_map,model,prev_events)
            path = [event_name]+path_ref
        else: #Controllable event
            if not event_name in var_map:#1-to-1 mapping between events and variables
                var_map[event_name]=model.addVar(vtype=GRB.CONTINUOUS,lb=0.0)
                model.update()
            path = [event_name]

        path_cache[event_name]=path #Caches solution for future use
        return path

def diff_min_max(dv2,dv1,sc_reform,var_map):
    """
    Computes the maximum and the minimum differences between two sequences of events.
    """
    dvs = [dv1[:-1],dv2[:-1]] #All contingent durations
    mins = [0.0,0.0]
    maxs = [0.0,0.0]
    for i in range(2):
        for ev in [e for e in dvs[i] if not e in dvs[1-i]]:
            mins[i] += sc_reform[ev]['lb']
            maxs[i] += sc_reform[ev]['ub']
    ec1 = var_map[dv1[-1]]; ec2 = var_map[dv2[-1]]
    diff_min = ec2-ec1+mins[1]-maxs[0]
    diff_max = ec2-ec1+maxs[1]-mins[0]
    return diff_min, diff_max

def uniform_sc_reformulation(tc,lb,ub,model,sc_reform,tc_squeeze_var_dict):
    """
    Creates the necessary variables and constraints to treat strong controllability
    of probabilistic durations with uniform distribution.
    """
    pdf = 1.0/(ub-lb) #PDF associated with squeezing this variable
    lb_var = model.addVar(vtype=GRB.CONTINUOUS,lb=0.0,ub=ub-lb,obj=pdf)
    ub_var = model.addVar(vtype=GRB.CONTINUOUS,lb=0.0,ub=ub-lb,obj=pdf)
    model.update()
    #The relaxation cannot make the upper bound bigger than the
    #lower bound.
    model.addConstr(lb_var+ub_var<=ub-lb)
    #Updates the strong controllability reformulation dictionary
    sc_reform[tc.end.name]['lb']+=lb_var
    sc_reform[tc.end.name]['ub']+=-ub_var
    #Updates the correspondence between uncontrollable temporal constraints
    #and squeezing variables
    tc_squeeze_var_dict[tc]={'lb':sc_reform[tc.end.name]['lb'],
                             'ub':sc_reform[tc.end.name]['ub'],
                             'risk_function':uniform_squeeze_risk}

def smart_piecewise_gaussian_sc_reformulation(tc,model,sc_reform,tc_squeeze_var_dict,
                                               num_div=5,optimize_partition=False,
                                               lr=0.05,tol=1e-5,max_iter=10000,
                                               gaussian_lb_sigma=None,
                                               gaussian_ub_sigma=None,
                                               random_part_init=False):
    """
    Generates an upper bound for a Gaussian CDF by performing a piecewise constraint
    approximation of its PDF (therefore approximating the CDF by straight line
    segments).
    """
    mean,var = tc.distribution['mean'],tc.distribution['variance']
    stdev = np.sqrt(var)
    lb_breaks, ub_breaks = get_gaussian_partition(mean=mean,var=var,num_div=num_div,
                                                  lb_sigma=gaussian_lb_sigma,
                                                  ub_sigma=gaussian_ub_sigma,
                                                  optimize_partition=optimize_partition,
                                                  random_part_init=random_part_init,
                                                  lr=lr,tol=tol,max_iter=max_iter)

    lb_vars=[]
    for i in range(1,len(lb_breaks)):
        lb_vars.append(model.addVar(vtype=GRB.CONTINUOUS,lb=0.0,ub=lb_breaks[i]-lb_breaks[i-1],
                                    obj=st.norm.pdf(lb_breaks[i],loc=mean,scale=stdev)))
    ub_vars=[]
    for i in range(1,len(ub_breaks)):
        ub_vars.append(model.addVar(vtype=GRB.CONTINUOUS,lb=0.0,ub=ub_breaks[i]-ub_breaks[i-1],
                                    obj=st.norm.pdf(ub_breaks[i-1],loc=mean,scale=stdev)))

    #Adds the variables to the model
    model.update()

    lower_bound = lb_breaks[0]+quicksum(lb_vars) #Total lower bound
    upper_bound = ub_breaks[-1]-quicksum(ub_vars)#Total upper bound

    #model.addConstr(lower_bound-upper_bound<=0.0)#Lower bound cannot be bigger than upper bound

    #Updates the strong controllability reformulation dictionary
    sc_reform[tc.end.name]['lb']=lower_bound
    sc_reform[tc.end.name]['ub']=upper_bound

    #Baseline risk of cutting a Gaussian
    sc_reform['_baseline_risk_']+=st.norm.cdf(lb_breaks[0],loc=mean,scale=stdev)+1.0-st.norm.cdf(ub_breaks[-1],loc=mean,scale=stdev)

    #Updates the correspondence between uncontrollable temporal constraints
    #and squeezing variables
    tc_squeeze_var_dict[tc]={'lb':sc_reform[tc.end.name]['lb'],
                             'ub':sc_reform[tc.end.name]['ub'],
                             'risk_function':gaussian_squeeze_risk}

def uniform_squeeze_risk(tc,new_lb,new_ub):
    """
    Risk for squeezing uniform distributions.
    """
    ub_squeeze=tc.ub-new_ub; lb_squeeze=new_lb-tc.lb
    if ub_squeeze>=0.0 and lb_squeeze>=0.0:
        return (1.0/(tc.ub-tc.lb))*(ub_squeeze+lb_squeeze)
    else:
        raise ValueError('Invalid squeezing of uniform distribution:[%f,%f]'%(lb_squeeze,ub_squeeze))

def gaussian_squeeze_risk(tc,new_lb,new_ub):
    """
    Risk for squeezing Gaussian distributions.
    """
    mean = tc.distribution['mean']; stdev=np.sqrt(tc.distribution['variance'])
    return st.norm.cdf(new_lb,loc=mean,scale=stdev)+1.0-st.norm.cdf(new_ub,loc=mean,scale=stdev)