from abc import ABCMeta, abstractmethod
from collections import Iterable
from numbers import Integral, Real
from warnings import warn

from six import add_metaclass
import numpy as np

from .function import Tabulated1D, INTERPOLATION_SCHEME
from openmc.stats.univariate import Univariate, Tabular, Discrete, Mixture
import openmc.checkvalue as cv
from openmc.mixin import EqualityMixin
from .data import EV_PER_MEV
from .endf import get_tab1_record, get_tab2_record

@add_metaclass(ABCMeta)
class EnergyDistribution(EqualityMixin):
    """Abstract superclass for all energy distributions."""
    def __init__(self):
        pass

    @abstractmethod
    def to_hdf5(self, group):
        pass

    @staticmethod
    def from_hdf5(group):
        """Generate energy distribution from HDF5 data

        Parameters
        ----------
        group : h5py.Group
            HDF5 group to read from

        Returns
        -------
        openmc.data.EnergyDistribution
            Energy distribution

        """
        energy_type = group.attrs['type'].decode()
        if energy_type == 'maxwell':
            return MaxwellEnergy.from_hdf5(group)
        elif energy_type == 'evaporation':
            return Evaporation.from_hdf5(group)
        elif energy_type == 'watt':
            return WattEnergy.from_hdf5(group)
        elif energy_type == 'madland-nix':
            return MadlandNix.from_hdf5(group)
        elif energy_type == 'discrete_photon':
            return DiscretePhoton.from_hdf5(group)
        elif energy_type == 'level':
            return LevelInelastic.from_hdf5(group)
        elif energy_type == 'continuous':
            return ContinuousTabular.from_hdf5(group)
        else:
            raise ValueError("Unknown energy distribution type: {}"
                             .format(energy_type))

    @staticmethod
    def from_endf(file_obj, params):
        """Generate energy distribution from an ENDF evaluation

        Parameters
        ----------
        file_obj : file-like object
            ENDF file positioned at the start of a section for an energy
            distribution.
        params : list
            List of parameters at the start of the energy distribution that
            includes the LF value indicating what type of energy distribution is
            present.

        Returns
        -------
        openmc.data.EnergyDistribution
            A sub-class of :class:`openmc.data.EnergyDistribution`

        """
        lf = params[3]
        if lf == 1:
            return ArbitraryTabulated.from_endf(file_obj, params)
        elif lf == 5:
            return GeneralEvaporation.from_endf(file_obj, params)
        elif lf == 7:
            return MaxwellEnergy.from_endf(file_obj, params)
        elif lf == 9:
            return Evaporation.from_endf(file_obj, params)
        elif lf == 11:
            return WattEnergy.from_endf(file_obj, params)
        elif lf == 12:
            return MadlandNix.from_endf(file_obj, params)

class ArbitraryTabulated(EnergyDistribution):
    r"""Arbitrary tabulated function given in ENDF MF=5, LF=1 represented as

    .. math::
         f(E \rightarrow E') = g(E \rightarrow E')

    Parameters
    ----------
    energy : numpy.ndarray
        Array of incident neutron energies
    pdf : list of openmc.data.Tabulated1D
        Tabulated outgoing energy distribution probability density functions

    Attributes
    ----------
    energy : numpy.ndarray
        Array of incident neutron energies
    pdf : list of openmc.data.Tabulated1D
        Tabulated outgoing energy distribution probability density functions

    """

    def __init__(self, energy, pdf):
        super(ArbitraryTabulated, self).__init__()
        self.energy = energy
        self.pdf = pdf

    def to_hdf5(self, group):
        raise NotImplementedError

    @classmethod
    def from_endf(cls, file_obj, params):
        """Generate arbitrary tabulated distribution from an ENDF evaluation

        Parameters
        ----------
        file_obj : file-like object
            ENDF file positioned at the start of a section for an energy
            distribution.
        params : list
            List of parameters at the start of the energy distribution that
            includes the LF value indicating what type of energy distribution is
            present.

        Returns
        -------
        openmc.data.ArbitraryTabulated
            Arbitrary tabulated distribution

        """
        params, tab2 = get_tab2_record(file_obj)
        n_energies = params[5]

        energy = np.zeros(n_energies)
        pdf = []
        for j in range(n_energies):
            params, func = get_tab1_record(file_obj)
            energy[j] = params[1]
            pdf.append(func)
        return cls(energy, pdf)

class GeneralEvaporation(EnergyDistribution):
    r"""General evaporation spectrum given in ENDF MF=5, LF=5 represented as

    .. math::
        f(E \rightarrow E') = g(E'/\theta(E))

    Parameters
    ----------
    theta : openmc.data.Tabulated1D
        Tabulated function of incident neutron energy :math:`E`
    g : openmc.data.Tabulated1D
        Tabulated function of :math:`x = E'/\theta(E)`
    u : float
        Constant introduced to define the proper upper limit for the final
        particle energy such that :math:`0 \le E' \le E - U`

    Attributes
    ----------
    theta : openmc.data.Tabulated1D
        Tabulated function of incident neutron energy :math:`E`
    g : openmc.data.Tabulated1D
        Tabulated function of :math:`x = E'/\theta(E)`
    u : float
        Constant introduced to define the proper upper limit for the final
        particle energy such that :math:`0 \le E' \le E - U`

    """

    def __init__(self, theta, g, u):
        super(GeneralEvaporation, self).__init__()
        self.theta = theta
        self.g = g
        self.u = u

    def to_hdf5(self, group):
        raise NotImplementedError

    @classmethod
    def from_ace(cls, ace, idx=0):
        raise NotImplementedError

    @classmethod
    def from_endf(cls, file_obj, params):
        """Generate general evaporation spectrum from an ENDF evaluation

        Parameters
        ----------
        file_obj : file-like object
            ENDF file positioned at the start of a section for an energy
            distribution.
        params : list
            List of parameters at the start of the energy distribution that
            includes the LF value indicating what type of energy distribution is
            present.

        Returns
        -------
        openmc.data.GeneralEvaporation
            General evaporation spectrum

        """
        u = params[0]
        params, theta = get_tab1_record(file_obj)
        params, g = get_tab1_record(file_obj)
        return cls(theta, g, u)

class MaxwellEnergy(EnergyDistribution):
    r"""Simple Maxwellian fission spectrum represented as

    .. math::
        f(E \rightarrow E') = \frac{\sqrt{E'}}{I} e^{-E'/\theta(E)}

    Parameters
    ----------
    theta : openmc.data.Tabulated1D
        Tabulated function of incident neutron energy
    u : float
        Constant introduced to define the proper upper limit for the final
        particle energy such that :math:`0 \le E' \le E - U`

    Attributes
    ----------
    theta : openmc.data.Tabulated1D
        Tabulated function of incident neutron energy
    u : float
        Constant introduced to define the proper upper limit for the final
        particle energy such that :math:`0 \le E' \le E - U`

    """

    def __init__(self, theta, u):
        super(MaxwellEnergy, self).__init__()
        self.theta = theta
        self.u = u

    @property
    def theta(self):
        return self._theta

    @property
    def u(self):
        return self._u

    @theta.setter
    def theta(self, theta):
        cv.check_type('Maxwell theta', theta, Tabulated1D)
        self._theta = theta

    @u.setter
    def u(self, u):
        cv.check_type('Maxwell restriction energy', u, Real)
        self._u = u

    def to_hdf5(self, group):
        """Write distribution to an HDF5 group

        Parameters
        ----------
        group : h5py.Group
            HDF5 group to write to

        """

        group.attrs['type'] = np.string_('maxwell')
        group.attrs['u'] = self.u
        self.theta.to_hdf5(group, 'theta')

    @classmethod
    def from_hdf5(cls, group):
        """Generate Maxwell distribution from HDF5 data

        Parameters
        ----------
        group : h5py.Group
            HDF5 group to read from

        Returns
        -------
        openmc.data.MaxwellEnergy
            Maxwell distribution

        """
        theta = Tabulated1D.from_hdf5(group['theta'])
        u = group.attrs['u']
        return cls(theta, u)

    @classmethod
    def from_ace(cls, ace, idx=0):
        """Create a Maxwell distribution from an ACE table

        Parameters
        ----------
        ace : openmc.data.ace.Table
            An ACE table
        idx : int
            Offset to read from in XSS array (default of zero)

        Returns
        -------
        openmc.data.MaxwellEnergy
            Maxwell distribution

        """
        # Read nuclear temperature -- since units are MeV, convert to eV
        theta = Tabulated1D.from_ace(ace, idx)
        theta.y *= EV_PER_MEV

        # Restriction energy
        nr = int(ace.xss[idx])
        ne = int(ace.xss[idx + 1 + 2*nr])
        u = ace.xss[idx + 2 + 2*nr + 2*ne]*EV_PER_MEV

        return cls(theta, u)

    @classmethod
    def from_endf(cls, file_obj, params):
        """Generate Maxwell distribution from an ENDF evaluation

        Parameters
        ----------
        file_obj : file-like object
            ENDF file positioned at the start of a section for an energy
            distribution.
        params : list
            List of parameters at the start of the energy distribution that
            includes the LF value indicating what type of energy distribution is
            present.

        Returns
        -------
        openmc.data.MaxwellEnergy
            Maxwell distribution

        """
        u = params[0]
        params, theta = get_tab1_record(file_obj)
        return cls(theta, u)

class Evaporation(EnergyDistribution):
    r"""Evaporation spectrum represented as

    .. math::
        f(E \rightarrow E') = \frac{E'}{I} e^{-E'/\theta(E)}

    Parameters
    ----------
    theta : openmc.data.Tabulated1D
        Tabulated function of incident neutron energy
    u : float
        Constant introduced to define the proper upper limit for the final
        particle energy such that :math:`0 \le E' \le E - U`

    Attributes
    ----------
    theta : openmc.data.Tabulated1D
        Tabulated function of incident neutron energy
    u : float
        Constant introduced to define the proper upper limit for the final
        particle energy such that :math:`0 \le E' \le E - U`

    """

    def __init__(self, theta, u):
        super(Evaporation, self).__init__()
        self.theta = theta
        self.u = u

    @property
    def theta(self):
        return self._theta

    @property
    def u(self):
        return self._u

    @theta.setter
    def theta(self, theta):
        cv.check_type('Evaporation theta', theta, Tabulated1D)
        self._theta = theta

    @u.setter
    def u(self, u):
        cv.check_type('Evaporation restriction energy', u, Real)
        self._u = u

    def to_hdf5(self, group):
        """Write distribution to an HDF5 group

        Parameters
        ----------
        group : h5py.Group
            HDF5 group to write to

        """

        group.attrs['type'] = np.string_('evaporation')
        group.attrs['u'] = self.u
        self.theta.to_hdf5(group, 'theta')

    @classmethod
    def from_hdf5(cls, group):
        """Generate evaporation spectrum from HDF5 data

        Parameters
        ----------
        group : h5py.Group
            HDF5 group to read from

        Returns
        -------
        openmc.data.Evaporation
            Evaporation spectrum

        """
        theta = Tabulated1D.from_hdf5(group['theta'])
        u = group.attrs['u']
        return cls(theta, u)

    @classmethod
    def from_ace(cls, ace, idx=0):
        """Create an evaporation spectrum from an ACE table

        Parameters
        ----------
        ace : openmc.data.ace.Table
            An ACE table
        idx : int
            Offset to read from in XSS array (default of zero)

        Returns
        -------
        openmc.data.Evaporation
            Evaporation spectrum

        """
        # Read nuclear temperature -- since units are MeV, convert to eV
        theta = Tabulated1D.from_ace(ace, idx)
        theta.y *= EV_PER_MEV

        # Restriction energy
        nr = int(ace.xss[idx])
        ne = int(ace.xss[idx + 1 + 2*nr])
        u = ace.xss[idx + 2 + 2*nr + 2*ne]*EV_PER_MEV

        return cls(theta, u)

    @classmethod
    def from_endf(cls, file_obj, params):
        """Generate evaporation spectrum from an ENDF evaluation

        Parameters
        ----------
        file_obj : file-like object
            ENDF file positioned at the start of a section for an energy
            distribution.
        params : list
            List of parameters at the start of the energy distribution that
            includes the LF value indicating what type of energy distribution is
            present.

        Returns
        -------
        openmc.data.Evaporation
            Evaporation spectrum

        """
        u = params[0]
        params, theta = get_tab1_record(file_obj)
        return cls(theta, u)

class WattEnergy(EnergyDistribution):
    r"""Energy-dependent Watt spectrum represented as

    .. math::
        f(E \rightarrow E') = \frac{e^{-E'/a}}{I} \sinh \left ( \sqrt{bE'}
        \right )

    Parameters
    ----------
    a, b : openmc.data.Tabulated1D
        Energy-dependent parameters tabulated as function of incident neutron
        energy
    u : float
        Constant introduced to define the proper upper limit for the final
        particle energy such that :math:`0 \le E' \le E - U`

    Attributes
    ----------
    a, b : openmc.data.Tabulated1D
        Energy-dependent parameters tabulated as function of incident neutron
        energy
    u : float
        Constant introduced to define the proper upper limit for the final
        particle energy such that :math:`0 \le E' \le E - U`

    """

    def __init__(self, a, b, u):
        super(WattEnergy, self).__init__()
        self.a = a
        self.b = b
        self.u = u

    @property
    def a(self):
        return self._a

    @property
    def b(self):
        return self._b

    @property
    def u(self):
        return self._u

    @a.setter
    def a(self, a):
        cv.check_type('Watt a', a, Tabulated1D)
        self._a = a

    @b.setter
    def b(self, b):
        cv.check_type('Watt b', b, Tabulated1D)
        self._b = b

    @u.setter
    def u(self, u):
        cv.check_type('Watt restriction energy', u, Real)
        self._u = u

    def to_hdf5(self, group):
        """Write distribution to an HDF5 group

        Parameters
        ----------
        group : h5py.Group
            HDF5 group to write to

        """

        group.attrs['type'] = np.string_('watt')
        group.attrs['u'] = self.u
        self.a.to_hdf5(group, 'a')
        self.b.to_hdf5(group, 'b')

    @classmethod
    def from_hdf5(cls, group):
        """Generate Watt fission spectrum from HDF5 data

        Parameters
        ----------
        group : h5py.Group
            HDF5 group to read from

        Returns
        -------
        openmc.data.WattEnergy
            Watt fission spectrum

        """
        a = Tabulated1D.from_hdf5(group['a'])
        b = Tabulated1D.from_hdf5(group['b'])
        u = group.attrs['u']
        return cls(a, b, u)

    @classmethod
    def from_ace(cls, ace, idx):
        """Create a Watt fission spectrum from an ACE table

        Parameters
        ----------
        ace : openmc.data.ace.Table
            An ACE table
        idx : int
            Offset to read from in XSS array (default of zero)

        Returns
        -------
        openmc.data.WattEnergy
            Watt fission spectrum

        """
        # Energy-dependent a parameter -- units are MeV, convert to eV
        a = Tabulated1D.from_ace(ace, idx)
        a.y *= EV_PER_MEV

        # Advance index
        nr = int(ace.xss[idx])
        ne = int(ace.xss[idx + 1 + 2*nr])
        idx += 2 + 2*nr + 2*ne

        # Energy-dependent b parameter -- units are MeV^-1
        b = Tabulated1D.from_ace(ace, idx)
        b.y /= EV_PER_MEV

        # Advance index
        nr = int(ace.xss[idx])
        ne = int(ace.xss[idx + 1 + 2*nr])
        idx += 2 + 2*nr + 2*ne

        # Restriction energy
        u = ace.xss[idx]*EV_PER_MEV

        return cls(a, b, u)

    @classmethod
    def from_endf(cls, file_obj, params):
        """Generate Watt fission spectrum from an ENDF evaluation

        Parameters
        ----------
        file_obj : file-like object
            ENDF file positioned at the start of a section for an energy
            distribution.
        params : list
            List of parameters at the start of the energy distribution that
            includes the LF value indicating what type of energy distribution is
            present.

        Returns
        -------
        openmc.data.WattEnergy
            Watt fission spectrum

        """
        u = params[0]
        params, a = get_tab1_record(file_obj)
        params, b = get_tab1_record(file_obj)
        return cls(a, b, u)

class MadlandNix(EnergyDistribution):
    r"""Energy-dependent fission neutron spectrum (Madland and Nix) given in
    ENDF MF=5, LF=12 represented as

    .. math::
        f(E \rightarrow E') = \frac{1}{2} [ g(E', E_F(L)) + g(E', E_F(H))]

    where

    .. math::
        g(E',E_F) = \frac{1}{3\sqrt{E_F T_M}} \left [ u_2^{3/2} E_1 (u_2) -
        u_1^{3/2} E_1 (u_1) + \gamma \left ( \frac{3}{2}, u_2 \right ) - \gamma
        \left ( \frac{3}{2}, u_1 \right ) \right ] \\ u_1 = \left ( \sqrt{E'} -
        \sqrt{E_F} \right )^2 / T_M \\ u_2 = \left ( \sqrt{E'} + \sqrt{E_F}
        \right )^2 / T_M.

    Parameters
    ----------
    efl, efh : float
        Constants which represent the average kinetic energy per nucleon of the
        fission fragment (efl = light, efh = heavy)
    tm : openmc.data.Tabulated1D
        Parameter tabulated as a function of incident neutron energy

    Attributes
    ----------
    efl, efh : float
        Constants which represent the average kinetic energy per nucleon of the
        fission fragment (efl = light, efh = heavy)
    tm : openmc.data.Tabulated1D
        Parameter tabulated as a function of incident neutron energy

    """

    def __init__(self, efl, efh, tm):
        super(MadlandNix, self).__init__()
        self.efl = efl
        self.efh = efh
        self.tm = tm

    @property
    def efl(self):
        return self._efl

    @property
    def efh(self):
        return self._efh

    @property
    def tm(self):
        return self._tm

    @efl.setter
    def efl(self, efl):
        name = 'Madland-Nix light fragment energy'
        cv.check_type(name, efl, Real)
        cv.check_greater_than(name, efl, 0.)
        self._efl = efl

    @efh.setter
    def efh(self, efh):
        name = 'Madland-Nix heavy fragment energy'
        cv.check_type(name, efh, Real)
        cv.check_greater_than(name, efh, 0.)
        self._efh = efh

    @tm.setter
    def tm(self, tm):
        cv.check_type('Madland-Nix maximum temperature', tm, Tabulated1D)
        self._tm = tm

    def to_hdf5(self, group):
        """Write distribution to an HDF5 group

        Parameters
        ----------
        group : h5py.Group
            HDF5 group to write to

        """

        group.attrs['type'] = np.string_('madland-nix')
        group.attrs['efl'] = self.efl
        group.attrs['efh'] = self.efh
        self.tm.to_hdf5(group)

    @classmethod
    def from_hdf5(cls, group):
        """Generate Madland-Nix fission spectrum from HDF5 data

        Parameters
        ----------
        group : h5py.Group
            HDF5 group to read from

        Returns
        -------
        openmc.data.MadlandNix
            Madland-Nix fission spectrum

        """
        efl = group.attrs['efl']
        efh = group.attrs['efh']
        tm = Tabulated1D.from_hdf5(group['tm'])
        return cls(efl, efh, tm)

    @classmethod
    def from_endf(cls, file_obj, params):
        """Generate Madland-Nix fission spectrum from an ENDF evaluation

        Parameters
        ----------
        file_obj : file-like object
            ENDF file positioned at the start of a section for an energy
            distribution.
        params : list
            List of parameters at the start of the energy distribution that
            includes the LF value indicating what type of energy distribution is
            present.

        Returns
        -------
        openmc.data.MadlandNix
            Madland-Nix fission spectrum

        """
        params, tm = get_tab1_record(file_obj)
        efl, efh = params[0:2]
        return cls(efl, efh, tm)

class DiscretePhoton(EnergyDistribution):
    """Discrete photon energy distribution

    Parameters
    ----------
    primary_flag : int
        Indicator of whether the photon is a primary or non-primary photon.
    energy : float
        Photon energy (if lp==0 or lp==1) or binding energy (if lp==2).
    atomic_weight_ratio : float
        Atomic weight ratio of the target nuclide responsible for the emitted
        particle

    Attributes
    ----------
    primary_flag : int
        Indicator of whether the photon is a primary or non-primary photon.
    energy : float
        Photon energy (if lp==0 or lp==1) or binding energy (if lp==2).
    atomic_weight_ratio : float
        Atomic weight ratio of the target nuclide responsible for the emitted
        particle

    """

    def __init__(self, primary_flag, energy, atomic_weight_ratio):
        super(DiscretePhoton, self).__init__()
        self.primary_flag = primary_flag
        self.energy = energy
        self.atomic_weight_ratio = atomic_weight_ratio

    @property
    def primary_flag(self):
        return self._primary_flag

    @property
    def energy(self):
        return self._energy

    @property
    def atomic_weight_ratio(self):
        return self._atomic_weight_ratio

    @primary_flag.setter
    def primary_flag(self, primary_flag):
        cv.check_type('discrete photon primary_flag', primary_flag, Integral)
        self._primary_flag = primary_flag

    @energy.setter
    def energy(self, energy):
        cv.check_type('discrete photon energy', energy, Real)
        self._energy = energy

    @atomic_weight_ratio.setter
    def atomic_weight_ratio(self, atomic_weight_ratio):
        cv.check_type('atomic weight ratio', atomic_weight_ratio, Real)
        self._atomic_weight_ratio = atomic_weight_ratio

    def to_hdf5(self, group):
        """Write distribution to an HDF5 group

        Parameters
        ----------
        group : h5py.Group
            HDF5 group to write to

        """

        group.attrs['type'] = np.string_('discrete_photon')
        group.attrs['primary_flag'] = self.primary_flag
        group.attrs['energy'] = self.energy
        group.attrs['atomic_weight_ratio'] = self.atomic_weight_ratio

    @classmethod
    def from_hdf5(cls, group):
        """Generate discrete photon energy distribution from HDF5 data

        Parameters
        ----------
        group : h5py.Group
            HDF5 group to read from

        Returns
        -------
        openmc.data.DiscretePhoton
            Discrete photon energy distribution

        """
        primary_flag = group.attrs['primary_flag']
        energy = group.attrs['energy']
        awr = group.attrs['atomic_weight_ratio']
        return cls(primary_flag, energy, awr)

    @classmethod
    def from_ace(cls, ace, idx):
        """Generate discrete photon energy distribution from an ACE table

        Parameters
        ----------
        ace : openmc.data.ace.Table
            An ACE table
        idx : int
            Offset to read from in XSS array (default of zero)

        Returns
        -------
        openmc.data.DiscretePhoton
            Discrete photon energy distribution

        """
        primary_flag = int(ace.xss[idx])
        energy = ace.xss[idx + 1]*EV_PER_MEV
        return cls(primary_flag, energy, ace.atomic_weight_ratio)

class LevelInelastic(EnergyDistribution):
    r"""Level inelastic scattering

    Parameters
    ----------
    threshold : float
        Energy threshold in the laboratory system, :math:`(A + 1)/A * |Q|`
    mass_ratio : float
        :math:`(A/(A + 1))^2`

    Attributes
    ----------
    threshold : float
        Energy threshold in the laboratory system, :math:`(A + 1)/A * |Q|`
    mass_ratio : float
        :math:`(A/(A + 1))^2`

    """

    def __init__(self, threshold, mass_ratio):
        super(LevelInelastic, self).__init__()
        self.threshold = threshold
        self.mass_ratio = mass_ratio

    @property
    def threshold(self):
        return self._threshold

    @property
    def mass_ratio(self):
        return self._mass_ratio

    @threshold.setter
    def threshold(self, threshold):
        cv.check_type('level inelastic threhsold', threshold, Real)
        self._threshold = threshold

    @mass_ratio.setter
    def mass_ratio(self, mass_ratio):
        cv.check_type('level inelastic mass ratio', mass_ratio, Real)
        self._mass_ratio = mass_ratio

    def to_hdf5(self, group):
        """Write distribution to an HDF5 group

        Parameters
        ----------
        group : h5py.Group
            HDF5 group to write to

        """

        group.attrs['type'] = np.string_('level')
        group.attrs['threshold'] = self.threshold
        group.attrs['mass_ratio'] = self.mass_ratio

    @classmethod
    def from_hdf5(cls, group):
        """Generate level inelastic distribution from HDF5 data

        Parameters
        ----------
        group : h5py.Group
            HDF5 group to read from

        Returns
        -------
        openmc.data.LevelInelastic
            Level inelastic scattering distribution

        """
        threshold = group.attrs['threshold']
        mass_ratio = group.attrs['mass_ratio']
        return cls(threshold, mass_ratio)

    @classmethod
    def from_ace(cls, ace, idx):
        """Generate level inelastic distribution from an ACE table

        Parameters
        ----------
        ace : openmc.data.ace.Table
            An ACE table
        idx : int
            Offset to read from in XSS array (default of zero)

        Returns
        -------
        openmc.data.LevelInelastic
            Level inelastic scattering distribution

        """
        threshold = ace.xss[idx]*EV_PER_MEV
        mass_ratio = ace.xss[idx + 1]
        return cls(threshold, mass_ratio)

class ContinuousTabular(EnergyDistribution):
    """Continuous tabular distribution

    Parameters
    ----------
    breakpoints : Iterable of int
        Breakpoints defining interpolation regions
    interpolation : Iterable of int
        Interpolation codes
    energy : Iterable of float
        Incoming energies at which distributions exist
    energy_out : Iterable of openmc.stats.Univariate
        Distribution of outgoing energies corresponding to each incoming energy

    Attributes
    ----------
    breakpoints : Iterable of int
        Breakpoints defining interpolation regions
    interpolation : Iterable of int
        Interpolation codes
    energy : Iterable of float
        Incoming energies at which distributions exist
    energy_out : Iterable of openmc.stats.Univariate
        Distribution of outgoing energies corresponding to each incoming energy

    """

    def __init__(self, breakpoints, interpolation, energy, energy_out):
        super(ContinuousTabular, self).__init__()
        self.breakpoints = breakpoints
        self.interpolation = interpolation
        self.energy = energy
        self.energy_out = energy_out

    @property
    def breakpoints(self):
        return self._breakpoints

    @property
    def interpolation(self):
        return self._interpolation

    @property
    def energy(self):
        return self._energy

    @property
    def energy_out(self):
        return self._energy_out

    @breakpoints.setter
    def breakpoints(self, breakpoints):
        cv.check_type('continuous tabular breakpoints', breakpoints,
                      Iterable, Integral)
        self._breakpoints = breakpoints

    @interpolation.setter
    def interpolation(self, interpolation):
        cv.check_type('continuous tabular interpolation', interpolation,
                      Iterable, Integral)
        self._interpolation = interpolation

    @energy.setter
    def energy(self, energy):
        cv.check_type('continuous tabular incoming energy', energy,
                      Iterable, Real)
        self._energy = energy

    @energy_out.setter
    def energy_out(self, energy_out):
        cv.check_type('continuous tabular outgoing energy', energy_out,
                      Iterable, Univariate)
        self._energy_out = energy_out

    def to_hdf5(self, group):
        """Write distribution to an HDF5 group

        Parameters
        ----------
        group : h5py.Group
            HDF5 group to write to

        """

        group.attrs['type'] = np.string_('continuous')

        dset = group.create_dataset('energy', data=self.energy)
        dset.attrs['interpolation'] = np.vstack((self.breakpoints,
                                                 self.interpolation))

        # Determine total number of (E,p) pairs and create array
        n_pairs = sum(len(d) for d in self.energy_out)
        pairs = np.empty((3, n_pairs))

        # Create array for offsets
        offsets = np.empty(len(self.energy_out), dtype=int)
        interpolation = np.empty(len(self.energy_out), dtype=int)
        n_discrete_lines = np.empty(len(self.energy_out), dtype=int)
        j = 0

        # Populate offsets and pairs array
        for i, eout in enumerate(self.energy_out):
            n = len(eout)
            offsets[i] = j

            if isinstance(eout, Mixture):
                discrete, continuous = eout.distribution
                n_discrete_lines[i] = m = len(discrete)
                interpolation[i] = 1 if continuous.interpolation == 'histogram' else 2
                pairs[0, j:j+m] = discrete.x
                pairs[1, j:j+m] = discrete.p
                pairs[2, j:j+m] = discrete.c
                pairs[0, j+m:j+n] = continuous.x
                pairs[1, j+m:j+n] = continuous.p
                pairs[2, j+m:j+n] = continuous.c
            else:
                if isinstance(eout, Tabular):
                    n_discrete_lines[i] = 0
                    interpolation[i] = 1 if eout.interpolation == 'histogram' else 2
                elif isinstance(eout, Discrete):
                    n_discrete_lines[i] = n
                    interpolation[i] = 1
                pairs[0, j:j+n] = eout.x
                pairs[1, j:j+n] = eout.p
                pairs[2, j:j+n] = eout.c
            j += n

        # Create dataset for distributions
        dset = group.create_dataset('distribution', data=pairs)

        # Write interpolation as attribute
        dset.attrs['offsets'] = offsets
        dset.attrs['interpolation'] = interpolation
        dset.attrs['n_discrete_lines'] = n_discrete_lines

    @classmethod
    def from_hdf5(cls, group):
        """Generate continuous tabular distribution from HDF5 data

        Parameters
        ----------
        group : h5py.Group
            HDF5 group to read from

        Returns
        -------
        openmc.data.ContinuousTabular
            Continuous tabular energy distribution

        """
        interp_data = group['energy'].attrs['interpolation']
        energy_breakpoints = interp_data[0, :]
        energy_interpolation = interp_data[1, :]
        energy = group['energy'].value

        data = group['distribution']
        offsets = data.attrs['offsets']
        interpolation = data.attrs['interpolation']
        n_discrete_lines = data.attrs['n_discrete_lines']

        energy_out = []
        n_energy = len(energy)
        for i in range(n_energy):
            # Determine length of outgoing energy distribution and number of
            # discrete lines
            j = offsets[i]
            if i < n_energy - 1:
                n = offsets[i+1] - j
            else:
                n = data.shape[1] - j
            m = n_discrete_lines[i]

            # Create discrete distribution if lines are present
            if m > 0:
                eout_discrete = Discrete(data[0, j:j+m], data[1, j:j+m])
                eout_discrete.c = data[2, j:j+m]
                p_discrete = eout_discrete.c[-1]

            # Create continuous distribution
            if m < n:
                interp = INTERPOLATION_SCHEME[interpolation[i]]
                eout_continuous = Tabular(data[0, j+m:j+n], data[1, j+m:j+n], interp)
                eout_continuous.c = data[2, j+m:j+n]

            # If both continuous and discrete are present, create a mixture
            # distribution
            if m == 0:
                eout_i = eout_continuous
            elif m == n:
                eout_i = eout_discrete
            else:
                eout_i = Mixture([p_discrete, 1. - p_discrete],
                                 [eout_discrete, eout_continuous])
            energy_out.append(eout_i)

        return cls(energy_breakpoints, energy_interpolation,
                   energy, energy_out)

    @classmethod
    def from_ace(cls, ace, idx, ldis):
        """Generate continuous tabular energy distribution from ACE data

        Parameters
        ----------
        ace : openmc.data.ace.Table
            ACE table to read from
        idx : int
            Index in XSS array of the start of the energy distribution data
            (LDIS + LOCC - 1)
        ldis : int
            Index in XSS array of the start of the energy distribution block
            (e.g. JXS[11])

        Returns
        -------
        openmc.data.ContinuousTabular
            Continuous tabular energy distribution

        """
        # Read number of interpolation regions and incoming energies
        n_regions = int(ace.xss[idx])
        n_energy_in = int(ace.xss[idx + 1 + 2*n_regions])

        # Get interpolation information
        idx += 1
        if n_regions > 0:
            breakpoints = ace.xss[idx:idx + n_regions].astype(int)
            interpolation = ace.xss[idx + n_regions:idx + 2*n_regions].astype(int)
        else:
            breakpoints = np.array([n_energy_in])
            interpolation = np.array([2])

        # Incoming energies at which distributions exist
        idx += 2*n_regions + 1
        energy = ace.xss[idx:idx + n_energy_in]*EV_PER_MEV

        # Location of distributions
        idx += n_energy_in
        loc_dist = ace.xss[idx:idx + n_energy_in].astype(int)

        # Initialize variables
        energy_out = []

        # Read each outgoing energy distribution
        for i in range(n_energy_in):
            idx = ldis + loc_dist[i] - 1

            # intt = interpolation scheme (1=hist, 2=lin-lin)
            INTTp = int(ace.xss[idx])
            intt = INTTp % 10
            n_discrete_lines = (INTTp - intt)//10
            if intt not in (1, 2):
                warn("Interpolation scheme for continuous tabular distribution "
                     "is not histogram or linear-linear.")
                intt = 2

            n_energy_out = int(ace.xss[idx + 1])
            data = ace.xss[idx + 2:idx + 2 + 3*n_energy_out].copy()
            data.shape = (3, n_energy_out)
            data[0,:] *= EV_PER_MEV

            # Create continuous distribution
            eout_continuous = Tabular(data[0][n_discrete_lines:],
                                      data[1][n_discrete_lines:]/EV_PER_MEV,
                                      INTERPOLATION_SCHEME[intt])
            eout_continuous.c = data[2][n_discrete_lines:]

            # If discrete lines are present, create a mixture distribution
            if n_discrete_lines > 0:
                eout_discrete = Discrete(data[0][:n_discrete_lines],
                                         data[1][:n_discrete_lines])
                eout_discrete.c = data[2][:n_discrete_lines]
                if n_discrete_lines == n_energy_out:
                    eout_i = eout_discrete
                else:
                    p_discrete = min(sum(eout_discrete.p), 1.0)
                    eout_i = Mixture([p_discrete, 1. - p_discrete],
                                     [eout_discrete, eout_continuous])
            else:
                eout_i = eout_continuous

            energy_out.append(eout_i)

        return cls(breakpoints, interpolation, energy, energy_out)


# -*- coding: utf-8 -*-
"""
This module allows the analysis of surface-based molecular cavities in various
volumes. To do so, a volume and a list of atom positions and their cutoff radii
is required. The following example shows reading the first frame of a ``xyz``
file (using the ``pybel`` module), defining a hexagonal volume and moving all
atoms into this volume.

.. code-block:: python
   :emphasize-lines: 7,10

   import pybel

   atoms = pybel.readfile("xyz", "hexagonal.xyz").next().atoms
   num_atoms = len(atoms)
   atom_positions = [atom.coords for atom in atoms]

   volume = volumes.HexagonalVolume(17.68943, 22.61158)
   for atom_index in range(num_atoms):
       atom_positions[atom_index] = volume.get_equivalent_point(atom_positions[atom_index])
   atoms = Atoms(atom_positions, [2.8]*num_atoms)

After this, a discretization of the volume is needed. This module supports
caching of these with the ``DiscretizationCache`` class.

.. code-block:: python

   discretization_cache = DiscretizationCache('cache.hdf5')
   discretization = discretization_cache.get_discretization(volume, 192)

Using this ``discretization`` and the ``Atoms`` object created above, the atom
positions and their cutoff radii are also discretized.

.. code-block:: python

   atom_discretization = AtomDiscretization(atoms, discretization)

With these objects created, all preparations are complete and the domain and
cavity calculation can be done:

.. code-block:: python

   domain_calculation = DomainCalculation(discretization, atom_discretization)
   cavity_calculation = CavityCalculation(domain_calculation)

Additionally, calculation of center-based cavities is possible by passing an
additional parameter to the CavityCalculation:

.. code-block:: python

   cavity_calculation = CavityCalculation(domain_calculation,
                                          use_surface_points=False)

The FakeDomainCalculation class provides a drop-in replacement for the
domain_calculation object in case the results of a previous calculation need to
be used (this is possible as only those attributes which are stored are
actually used during center-based cavity calculation, which is not the case for
surface-based cavity calculations, which at least require the surface point
lists).

The CalculationResults class provides a container for the results and allows
storage to and retrieval from HDF5 files. These files have several groups which
contain the relevant information.

Author: Florian Rhiem <f.rhiem@fz-juelich.de>
"""
from math import pi as PI
import sys
import numpy as np

from computation.split_and_merge.pipeline import start_split_and_merge_pipeline
from computation.split_and_merge.algorithm import ObjectType
from core.calculation.gyrationtensor import calculate_gyration_tensor_parameters
from util.message import print_message
from extension import atomstogrid, mark_cavities, cavity_triangles, cavity_intersections

dimension = 3
dimensions = range(dimension)

class DomainCalculation:
    """
    Cavity domain calulation is performed by the following steps:
     1. A grid is created with the resolution defined in the volume
        discretization and filled with zeros.
     2. For each atom, all points in the grid closer to this atom than the
        discrete cavity cutoff radius are set to a point indicating the atom
        index (atom_index+1).
     3. At this point, every point in the grid which is inside of the volume
        and still has a value of zero is part of a cavity domain. To find these
        domains, an optimized split and merge algorithm is applied to the whole
        grid. It returns the center and surface points of each cavity domain
        (points with a neighbor outside of the cavity domain) stored in lists.
        Points inside of a domain are marked with a negative value indicating
        which domain they are part of.
    """

    def __init__(self, discretization, atom_discretization):
        # step 1
        self.discretization = discretization
        self.atom_discretization = atom_discretization
        self.grid = np.zeros(self.discretization.d, dtype=np.int64)

        # step 2
        atomstogrid(self.grid,
                    self.atom_discretization.discrete_positions,
                    self.atom_discretization.atoms.radii_as_indices,
                    self.atom_discretization.sorted_discrete_radii,
                    [(0, 0, 0)] + self.discretization.combined_translation_vectors,
                    self.discretization.grid)
        # step 3
        result = start_split_and_merge_pipeline(self.grid,
                                                self.discretization.grid,
                                                self.atom_discretization.discrete_positions,
                                                self.discretization.combined_translation_vectors,
                                                self.discretization.get_translation_vector,
                                                ObjectType.DOMAIN)
        self.centers, translated_areas, non_translated_areas, self.surface_point_list, self.cyclic_area_indices = result
        print_message("Number of domains:", len(self.centers))

        self.domain_volumes = []
        self.critical_domains = []   # count of very small domains -> domains that can disappear on cutoff radius changes
        for domain_index in range(len(self.centers)):
            current_cell_sum = (self.grid == -(domain_index + 1)).sum()
            if current_cell_sum == 1:
                self.critical_domains.append(domain_index)
            domain_volume = current_cell_sum * (self.discretization.s_step ** 3)
            self.domain_volumes.append(domain_volume)

        self.characteristic_radii = [(0.75 * volume / PI)**(1.0/3.0) for volume in self.domain_volumes]

        if translated_areas:
            gyration_tensor_parameters = tuple(calculate_gyration_tensor_parameters(area) for area in translated_areas)
            (self.mass_centers, self.squared_gyration_radii, self.asphericities,
             self.acylindricities, self.anisotropies) = zip(*gyration_tensor_parameters)
            self.mass_centers = [self.discretization.discrete_to_continuous(point, result_inside_volume=True)
                                 for point in self.mass_centers]
            self.squared_gyration_radii = [self.discretization.discrete_to_continuous(value, unit_exponent=2)
                                           for value in self.squared_gyration_radii]
        else:
            (self.mass_centers, self.squared_gyration_radii, self.asphericities,
             self.acylindricities, self.anisotropies) = 5*([], )

        self.triangles()

    def triangles(self):
        if hasattr(self, "domain_triangles"):
            return self.domain_triangles
        number_of_domains = len(self.centers)
        print_message("Number of domains:", number_of_domains)
        triangles = []
        surface_areas = []
        step = (self.discretization.s_step,) * 3
        offset = self.discretization.discrete_to_continuous((0, 0, 0))
        for domain_index in range(number_of_domains):
            print_message("Calculating triangles for domain", domain_index)
            vertices, normals, surface_area = cavity_triangles(
                    self.grid,
                    [domain_index],
                    1, step, offset,
                    self.discretization.grid)
            triangles.append((vertices, normals))
            surface_areas.append(surface_area)

        self.domain_triangles = triangles
        self.domain_surface_areas = surface_areas
        return triangles

class CavityCalculation:
    """
    Cavity domain calulation is performed by the following steps:
     1. The discrete volume grid is divided into subgrid cells with a side
        length based on the maximum discrete cavity cutoff radius. For each
        subgrid, a tuple of three lists is stored (in self.sg).
     2. The first list for each subgrid cell is filled with the atoms inside
        the cell (their 'real' positions, which might be outside of the
        volume).
     3. The second and third lists are filled with surface points and their
        domain index. (These might also be moved with the translation vectors
        and might thereby also be outside of the volume.)
     4. A new grid is created (grid3) and each point in this grid is set to
        zero if it is outside of the volume or part inside of the cavity cutoff
        radius of an atom, or a negative value if it is part of a cavity domain
        (see the domain calculation step 2 for this).
     5. For each point which is inside the cavity cutoff radius of an
        atom, the nearest atom and the nearest domain surface point are found
        by using the neighbor subgrid cells. If a domain surface point is
        nearer than the nearest atom, than the point belongs to the cavity
        which this surface point belonged to and is marked with a negative
        value.
     6. At this point, two cavities constructed from two cavity domains might
        actually be one multicavity. In this step, these are found and a list
        of multicavities is created.

    About the subgrid cells:
    If a point inside a subgrid cell was marked as 'near an atom' during the
    domain calculation, then the position of this atom must be either in the
    same cell or in one of the cell's neighbors. This is guaranteed, because
    the atom must be at most its on cavity cutoff radius away, and the subgrid
    cell size is the maximum cavity cutoff radius.

    To calculate center-based cavities, use a grid filled with zeros instead of
    resuing some values from the domain calculation grid and then iterate over
    the domain centers instead of the domain surface points.
    """

    def __init__(self, domain_calculation, use_surface_points=True, gyration_tensor_parameters=False):
        self.domain_calculation = domain_calculation
        if use_surface_points:
            self.grid = self.domain_calculation.grid
            num_surface_points = sum(map(len, self.domain_calculation.surface_point_list))
            print_message("Number of surface points:", num_surface_points)
        else:
            self.grid = None

        self.sg_cube_size = self.domain_calculation.atom_discretization.sorted_discrete_radii[0]
        if use_surface_points:
            domain_seed_point_lists = self.domain_calculation.surface_point_list
        else:
            domain_seed_point_lists = [[center] for center in self.domain_calculation.centers]

        discretization = self.domain_calculation.discretization
        atom_discretization = self.domain_calculation.atom_discretization

        # steps 1 to 5
        self.grid3 = mark_cavities(self.grid,
                                   discretization.grid,
                                   discretization.d,
                                   self.sg_cube_size,
                                   atom_discretization.discrete_positions,
                                   [(0, 0, 0)] + discretization.combined_translation_vectors,
                                   domain_seed_point_lists,
                                   use_surface_points)

        if gyration_tensor_parameters:
            result = start_split_and_merge_pipeline(self.grid3,
                                                    discretization.grid,
                                                    atom_discretization.discrete_positions,
                                                    discretization.combined_translation_vectors,
                                                    discretization.get_translation_vector,
                                                    ObjectType.CAVITY)
            translated_areas, non_translated_areas, cyclic_area_indices = result

        num_domains = len(self.domain_calculation.centers)
        grid_volume = (discretization.grid == 0).sum()
        self.cavity_volumes = []
        for domain_index in range(num_domains):
            self.cavity_volumes.append(1.0 * (self.grid3 == -(domain_index + 1)).sum() * (discretization.s_step ** 3))
        self.characteristic_radii = [(0.75 * volume / PI)**(1.0/3.0) for volume in self.cavity_volumes]

        # step 6
        intersection_table = cavity_intersections(self.grid3, num_domains)
        multicavities = []
        cavity_to_neighbors = num_domains * [None]
        for domain in range(num_domains):
            current_neighbors = set([domain])
            for neighbor in range(num_domains):
                if intersection_table[domain][neighbor] == 1:
                    current_neighbors.add(neighbor)
            for multicavity in multicavities[:]:
                if any([neighbor in multicavity for neighbor in current_neighbors]):
                    current_neighbors = current_neighbors | multicavity
                    multicavities.remove(multicavity)
            multicavities.append(current_neighbors)
            for neighbor in current_neighbors:
                cavity_to_neighbors[neighbor] = current_neighbors
        self.multicavities = multicavities
        self.multicavity_volumes = []
        for multicavity in multicavities:
            self.multicavity_volumes.append(sum(self.cavity_volumes[cavity_index] for cavity_index in multicavity))
        print_message("Multicavity volumes:", self.multicavity_volumes)

        if gyration_tensor_parameters:
            def key_func(cavity_index):
                cavity_area = non_translated_areas[cavity_index]
                a_single_cavity_index = -self.grid3[cavity_area[0]] - 1
                max_neighbor_index = max(cavity_to_neighbors[a_single_cavity_index])
                return max_neighbor_index
            sorted_area_indices = sorted(range(len(self.multicavities)), key=key_func)
            sorted_translated_areas = [translated_areas[i] for i in sorted_area_indices]
            sorted_cyclic_area_indices = [i for i, index in enumerate(sorted_area_indices)
                                        if index in cyclic_area_indices]
            self.cyclic_area_indices = sorted_cyclic_area_indices

            if sorted_translated_areas:
                gyration_tensor_parameters = tuple(calculate_gyration_tensor_parameters(area)
                                                for area in sorted_translated_areas)
                (self.mass_centers, self.squared_gyration_radii, self.asphericities,
                self.acylindricities, self.anisotropies) = zip(*gyration_tensor_parameters)
                self.mass_centers = [discretization.discrete_to_continuous(point, result_inside_volume=True)
                                    for point in self.mass_centers]
                self.squared_gyration_radii = [discretization.discrete_to_continuous(value, unit_exponent=2)
                                            for value in self.squared_gyration_radii]
            else:
                (self.mass_centers, self.squared_gyration_radii, self.asphericities,
                self.acylindricities, self.anisotropies) = 5*([], )

        self.triangles()

    def squared_distance(self, a, b):
        '''
        Calculates the squared distance between two points while taking the
        translation vectors into account.
        '''
        sqd = sys.maxint
        for v in self.domain_calculation.discretization.combined_translation_vectors + [(0, 0, 0)]:
            sqd = min(sqd, sum([(a[i] - b[i] + v[i]) * (a[i] - b[i] + v[i]) for i in dimensions]))
        return sqd

    def triangles(self):
        if hasattr(self, "cavity_triangles"):
            return self.cavity_triangles
        step = (self.domain_calculation.discretization.s_step,) * 3
        offset = self.domain_calculation.discretization.discrete_to_continuous((0, 0, 0))
        triangles = []
        surface_areas = []
        for i, multicavity in enumerate(self.multicavities):
            print_message("Generating triangles for multicavity:", i+1)
            vertices, normals, surface_area = cavity_triangles(
                    self.grid3,
                    multicavity,
                    4, step, offset,
                    self.domain_calculation.discretization.grid)
            triangles.append((vertices, normals))
            surface_areas.append(surface_area)

        self.cavity_triangles = triangles
        self.cavity_surface_areas = surface_areas
        return cavity_triangles

    def __getattr__(self, attr):
        optional_attributes = ('mass_centers', 'squared_gyration_radii', 'asphericities', 'acylindricities',
                               'anisotropies', 'characteristic_radii', 'cyclic_area_indices')
        if attr in optional_attributes:
            return None
        else:
            return super(CavityCalculation, self).__getattr__(attr)

class FakeDomainCalculation(object):
    '''
    When calculating center-based cavities, a DomainCalculation object is
    required. This has to provide the domain central points, but not the surface
    points (as those are only needed for surface-based cavity calculation).
    Objects of this class can be used as a drop-in for 'real'
    DomainCalculations, e.g. when the required data is loaded from a file.
    '''

    def __init__(self, discretization, atom_discretization, results):
        self.centers = results.domains.centers
        self.discretization = discretization
        self.atom_discretization = atom_discretization


#!/usr/bin/env python

import datetime
import json
import math
import random
import sys
from ipaddress import ip_address, ip_network

import numpy as np
import pytricia
import yaml

import networkx

ip_prefixes = pytricia.PyTricia()

seed = hash(datetime.datetime.now())
global_policy = {}

def read_topo(filename, local_policy=None):
    data = yaml.load(open(filename))
    G = networkx.Graph()
    nodes = data["nodes"]
    links = data["links"]
    for n in nodes:
        nid = nodes[n]['id']
        G.add_node(nid, name=n, **nodes[n])
        G.node[nid]['type'] = nodes[n]['type']
        G.node[nid]['ip-prefix'] = nodes[n].get('ip-prefix', [])
        G.node[nid]['routing'] = pytricia.PyTricia()
        G.node[nid]['fine_grained'] = pytricia.PyTricia()
        for prefix in nodes[n].get('ip-prefix', []):
            ip_prefixes[prefix] = nid
    for l in links:
        G.add_edge(*l)
    return G

bgp_general_policy = networkx.shortest_path

# FIXME: The generation algorithm does not make sense
# Refer to the evaluation part of SDX and SIDR
def generate_local_policy(G, **args):
    """
    Generate the local policy for each node in G.

    Args:
        G: graph.
        args: additional arguments to config the generation algorithm.

    Returns:
        The mapping from node id to a local_policy table.
        local_policy_table ::= Trie<prefix, Map<port, next_hop>>
    """
    global global_policy
    random.seed(seed)
    policies = dict()
    max_ports = 30
    for node in G.nodes():
        if node not in policies:
            policies[node] = pytricia.PyTricia()
        # FIXME: select prefix by following a distribution
        for prefix in G.node[node]["ip-prefix"]:
            # FIXME: select ports by following a distribution
            ports = set([
                random.randint(10000, 60000)
                for _ in range(random.randint(1, max_ports))
            ])
            for port in ports:
                policies[node][port] = random.choice(
                    [i for i in G.neighbors(node)] + [None])
    global_policy = policies

DEFAULT_SERVICE_TYPES = {21: 0.1, 80: 0.1, 2801: 0.2, 8444: 0.3, 8445: 0.3}

# TODO: Some key points
# - What's the distribution for policy type (blackhole/deflection) selection
# - What's the distribution for node selection
# - How many network we need to fwd
# - What's the distribution for the tcp port selection
def generate_random_policy(G,
                           service_types=DEFAULT_SERVICE_TYPES,
                           network_ratio=0.5,
                           prefix_ratio=0.2,
                           policy_place='transit',
                           policy_type='both',
                           **args):
    """
    Randomly setup black-hole or deflection policies in transit network

    Args:
        G: Topology
        service_types: supported service port distribution.
        network_ratio: how many networks is selected. default 0.5
        prefix_ratio: how many prefix is seleted. default: 0.2
        policy_place: transit, edge or both. default: 'transit'
        policy_type: blackhole, deflection or both. default: 'both'
        args: additional arguments.

    Returns:
        global local policy table.
    """
    edge_networks = [
        n for n in G.nodes() if G.node[n].get('type', '') == 'edge'
    ]
    transit_networks = [
        n for n in G.nodes() if G.node[n].get('type', '') == 'transit'
    ]
    policies = {}
    if policy_place == 'transit':
        policy_networks = transit_networks
    elif policy_place == 'edge':
        policy_networks = edge_networks
    else:
        policy_networks = G.nodes()
    for d in policy_networks:
        if d not in policies:
            policies[d] = pytricia.PyTricia()
        for dest in random.sample([n for n in edge_networks if n != d],
                                  math.ceil(
                                      len(edge_networks) * float(network_ratio))):
            node_prefixes = G.node[dest]['ip-prefix']
            for prefix in random.sample(
                    node_prefixes, math.ceil(
                        len(node_prefixes) * float(prefix_ratio))):
                policies[d][prefix] = {
                    port: gen_single_policy(G, d, dest, policy_type)
                    for port in random.sample(service_types.keys(),
                                              random.randint(1, 4))
                }
    return policies

def gen_single_policy(G, d, dest, policy_type):
    if policy_type == 'both':
        return random.choice([None, random_deflection(G, d, dest)])
    elif policy_type == 'deflection':
        return random_deflection(G, d, dest)
    return None

def random_deflection(G, d, dest):
    next_hop = networkx.shortest_path(G, d, dest)[1]
    subG = G.copy()
    subG.remove_edge(d, next_hop)
    if dest not in networkx.descendants(subG, d):
        return None
    else:
        return networkx.shortest_path(subG, d, dest)[1]

def get_local_policy(node):
    """
    local_policy is not a simple prefix based routing.

    For simplicity, assume the format of local_policy for each node is:

    local_policy ::= Trie<prefix, Map<port, next_hop>>
    """
    global global_policy
    return global_policy.get(node, pytricia.PyTricia())

def default_routing_policy(node,
                           dst_ip,
                           dst_port=None,
                           src_ip=None,
                           src_port=None,
                           protocol='tcp',
                           **args):
    """
    Default routing policy for networks.

    Args:
        node: node id for the network.
        dst_ip: destination ip address.
        dst_port: optional.
        src_ip: optional.
        src_port: optional.
        protocol: optional.
        args: additional flow spec.

    Returns:
        The next hop of the give flow spec from this node.
    """
    for prefix in node.get('ip-prefix', []):
        if ip_address(dst_ip) in ip_network(prefix):
            return None
    local_policy = get_local_policy(node['id'])
    if dst_ip in local_policy:
        local_policy_for_ip = local_policy[dst_ip]
        if dst_port in local_policy_for_ip:
            return local_policy_for_ip[dst_port]
    fg_routing = node['fine_grained']
    if dst_ip in fg_routing:
        fg_routing_for_ip = fg_routing[dst_ip]
        if dst_port in fg_routing_for_ip:
            return fg_routing_for_ip[dst_port]
    if dst_ip in node['routing']:
        return node['routing'][dst_ip]

def fp_bgp_convergence(G):
    """
    False-positive BGP Convergence. node["routing"] is the table of  {ip-prefix -> next hop node}
    """
    paths = networkx.shortest_path(G)
    for src in G.nodes:
        for dst in G.nodes:
            if src != dst:
                try:
                    path = paths[src].get(dst)
                except KeyError:
                    path = None
                if path:
                    prefixes = G.node[dst]['ip-prefix']
                    for hop, next_hop in zip(path[:-1], path[1:]):
                        for prefix in prefixes:
                            G.node[hop]["routing"][prefix] = next_hop

def all_fg_nodes(G, prefix):
    """
    Let's assume the local_policy is generated from the prefixes of nodes first.
    """
    # FIXME: If the granularity of prefixes in local_policy is different from ones in node, it will conduct error.
    fg_nodes = []
    for node in G.nodes:
        local = get_local_policy(node)
        if local.get(prefix):
            fg_nodes.append(node)
    return fg_nodes

def correct_bgp_convergence(G):
    """
    Correct BGP Convergence
    """
    for dst in G.nodes:
        prefixes = G.node[dst]['ip-prefix']
        for prefix in prefixes:
            H = G.copy()
            for node in all_fg_nodes(G, prefix):
                H.remove_node(node)
            paths = networkx.shortest_path(H)
            for src in H.nodes:
                if src != dst:
                    # path = paths[src][dst]
                    path = paths[src].get(dst, [])
                    for hop, next_hop in zip(path[:-1], path[1:]):
                        G.node[hop]["routing"][prefix] = next_hop

def find_fine_grained_routes(G, prefix_port):
    for prefix, ports in prefix_port.items():
        for port in ports:
            delete_nodes = set()
            delete_links = set()
            H = G.copy()  # A copy of directed graph to modify links and nodes
            # Step 1: Remove unused links for <prefix, port>
            for node in H.nodes:
                action = get_local_policy(node).get(prefix)
                if action and (port in action):
                    action = action[port]
                    if action:
                        for neig in H.neighbors(node):
                            if action != neig:
                                delete_links.add((node, neig))
                    else:
                        delete_nodes.add(node)
            for edge in delete_links:
                H.remove_edge(*edge)
            for node in delete_nodes:
                H.remove_node(node)
            # Step 2: Find shortest path (Is it possible to be not found?)
            dst = ip_prefixes[prefix]
            paths = networkx.shortest_path(H, target=dst)
            # Step 3: Traverse the shortest path
            for src in H.nodes:
                if src != dst:
                    path = paths.get(src)
                    if path:
                        for hop, next_hop in zip(path[:-1], path[1:]):
                            if prefix not in G.node[hop]["fine_grained"]:
                                G.node[hop]["fine_grained"][prefix] = dict()
                            G.node[hop]["fine_grained"][prefix][port] = next_hop

    # Step 4: BGP shortest path in all nodes
    fp_bgp_convergence(G)
    return G

def fine_grained_announcement(G):
    for node in G.nodes:
        del G.node[node]["routing"]
        del G.node[node]["fine_grained"]
    G = G.to_directed()
    for node in G.nodes:
        G.node[node]["routing"] = pytricia.PyTricia()
        G.node[node]["fine_grained"] = pytricia.PyTricia()
    prefix_port = dict()  # type: dict{str, set{int}}
    for node in G.nodes:
        local = get_local_policy(node)
        for prefix in local:
            ports_actions = local[prefix]
            if prefix not in prefix_port:
                prefix_port[prefix] = set()
            for port in ports_actions:
                prefix_port[prefix].add(port)
    return find_fine_grained_routes(G, prefix_port)

def read_flows(filename, port_dist=DEFAULT_SERVICE_TYPES):
    """
    Examples:
        [{
            "src_ip": "10.0.0.1",
            "src_port": 22,
            "dst_ip": "10.0.10.1",
            "dst_port": 80,
            "protocol": "tcp",
            "start_time": 1516292713,
            "end_time": 1516313885,
            "volume": 4089456904
        }]

        required: src_ip, dst_ip, start_time, end_time, volume
        optional: src_port, dst_port, protocol
    """
    data = json.load(open(filename))
    for d in data:
        if not d.get('dst_port', None):
            d['dst_port'] = int(np.random.choice(
                list(port_dist.keys()), p=list(port_dist.values())))
        # (d['src_ip'], d['dst_ip']) = (d['dst_ip'], d['src_ip'])
    return data

statistic_as_length = {}

def coarse_grained_correct_bgp(G, F):
    """
    The principle of Correct BGP is very simple:

    For a specific IP p, compute the subgraph G' of G in which every node with local_policy covering p will not show.
    The final path of p is the shortest_path of subgraph G'.

    The Correct BGP MUST guarantee the local_policy NEVER be triggered.
    """
    for flow in F:
        node = G.node[ip_prefixes[flow["src_ip"]]]
        if node is not None:
            as_length = [node]
            prefixes = [ip_network(p) for p in node['ip-prefix']]
            while ip_address(flow["dst_ip"]) not in prefixes:
                if len(node["fine_grained"]) > 0:
                    try:
                        next_hop = node["fine_grained"][flow["dst_ip"]][
                            "dst_port"]
                        in_fg = True  # In fine grained policy
                    except KeyError:
                        in_fg = False
                if not in_fg:
                    next_hop = node["routing"][flow["dst_ip"]]
                as_length.append(next_hop)
        print("AS Length: %d" % len(as_length))
        if len(as_length) not in statistic_as_length:
            statistic_as_length[len(as_length)] = 0
        statistic_as_length[len(as_length)] += 1

# def match(ip, prefix):
#     return ip in prefix

# TODO: No need, merge it with routing_policy
# def compute_next_hop(flow, rib):
#     """
#     Compute the next hop of the flow from  RIB.
#
#     Args:
#         flow: flow spec.
#         rib: the RIB table from the Graph node.
#
#     Returns:
#         the next hop. (None if no route.)
#     """
#     for rule in rib:
#         if match(flow['dst-ip'], rule):
#             return rib[rule]
#     return None

def check_path(flow, G, routing_policy=default_routing_policy, debug=False):
    """
    Check the path of a given flow in the topology.

    Args:
        flow: The flow spec to check.
        G: The topology object.

    Returns:
        the AS-PATH length of the route.
        nan - no route
        inf - there is a loop
    """
    if debug:
        path = []
    loop_remover = {}
    src = ip_prefixes[flow['src_ip']]
    dst = ip_prefixes[flow['dst_ip']]
    if src == dst:
        if debug:
            return 1, [src]
        return 1
    d = src
    if debug:
        path = [src]
    dn = routing_policy(G.node[src], **flow)
    path_len = 1
    while dn:
        loop_remover[d] = loop_remover.get(d, 0) + 1
        # print d, p, loop_remover
        if loop_remover[d] > 1:
            if debug:
                return math.inf, []
            return math.inf
        d = dn
        if debug:
            path.append(dn)
        dn = routing_policy(G.node[d], **flow)
        path_len += 1
    if d != dst:
        if debug:
            return math.nan, []
        return math.nan
    if debug:
        return path_len, path
    return path_len

debug_dict = {}

def check_reachability(G, F, max_len=10, debug=False, debug_num=None):
    as_length_dist = {}
    success_volume = 0
    unsuccess_volume = 0
    R_F = []
    for f in F:
        if debug:
            global debug_dict
            result, path = check_path(f, G, debug=True)
            if debug_num not in debug_dict:
                debug_dict[debug_num] = dict()
            debug_dict[debug_num][(f["src_ip"], f["dst_ip"], f["start_time"],
                                   f["end_time"], f["volume"])] = (result, path)
        else:
            result = check_path(f, G)
        as_length_dist[result] = as_length_dist.get(result, 0) + 1
        if type(result) == float:
            unsuccess_volume += f['volume']
        elif result > 1:
            success_volume += f['volume']
            R_F.append(f)
    # print 'block_policies', 'deflection_policies'
    # print '%d\t%d' % policy_summary(G)
    if debug:
        as_len_nan = as_length_dist.pop(math.nan) if math.nan in as_length_dist else 0
        as_lens = sorted(as_length_dist.keys())
        as_lens.append(math.nan)
        as_length_dist[math.nan] = as_len_nan
        print('\t'.join([str(l) for l in as_lens]))
        print('\t'.join([str(as_length_dist[l]) for l in as_lens]))
    as_len_pdf = []
    for al in range(max_len - 1):
        as_len_pdf.append(as_length_dist.get(al + 2, 0))
    as_len_pdf.append(as_length_dist.get(math.inf, 0))
    as_len_pdf.append(as_length_dist.get(math.nan, 0))
    as_len_pdf.append(success_volume)
    as_len_pdf.append(unsuccess_volume)
    # Format: flow_num from 2 to max_len, inf, nan, success_volume, unsuccess_volume
    print('\t'.join([str(a) for a in as_len_pdf]))
    return R_F

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("%s topo-filename flow-filename mode" % sys.argv[0])
    topo_filename = sys.argv[1]
    flow_filename = sys.argv[2]
    mode = sys.argv[3]
    args = dict([tuple(k.split('=')) for k in sys.argv[4:]])

    G = read_topo(topo_filename)
    F = read_flows(flow_filename)
    # generate_local_policy(G)
    # global_policy = generate_random_policy(G, network_ratio=0.1, prefix_ratio=0.1, policy_type='blackhole')
    global_policy = generate_random_policy(G, **args)

    if '1' in mode:
        H = G.copy()
        fp_bgp_convergence(H)
        # print("FP_BGP:")
        check_reachability(H, F)
    if '2' in mode:
        H = G.copy()
        correct_bgp_convergence(H)
        # print("C_BGP:")
        R_F = check_reachability(H, F, debug=False, debug_num=2)
    if '3' in mode:
        H = G.copy()
        H = fine_grained_announcement(H)
        # print("SFP:")
        check_reachability(H, F, debug=False, debug_num=3)
    if '4' in mode:
        H = G.copy()
        H = fine_grained_announcement(H)
        # print("Reachable-SFP:")
        check_reachability(H, R_F, debug=False, debug_num=4)

    if len(debug_dict) > 0:
        flow_diff = list()
        for flow in debug_dict[3]:
            if debug_dict[3][flow][0] == math.nan and debug_dict[2][flow][0] != math.nan:
                print(str(flow))
        with open("debug.json", 'w') as f:
            f.write(debug_dict.__str__())

    # coarse_grained_correct_bgp(G, F)
    # print("AS Length statistics: %d" % statistic_as_length)


#!/usr/bin/env python3
# coding: utf-8

__author__ = 'cleardusk'

"""
The pipeline of 3DDFA prediction: given one image, predict the 3d face vertices, 68 landmarks and visualization.

[todo]
1. CPU optimization: https://pmchojnacki.wordpress.com/2018/10/07/slow-pytorch-cpu-performance
"""

import torch
import torchvision.transforms as transforms
import mobilenet_v1
import numpy as np
import cv2
import os
import math
from tqdm import tqdm
import time
# import face_alignment
# import single_align
from utils.ddfa import ToTensorGjz, NormalizeGjz, str2bool
import scipy.io as sio
from utils.inference import get_suffix, parse_roi_box_from_landmark, crop_img, predict_68pts, dump_to_ply, dump_vertex, \
    draw_landmarks, predict_dense, parse_roi_box_from_bbox, get_colors, write_obj_with_colors, get_aligned_param, parse_quality_list_part
from utils.cv_plot import plot_pose_box
from utils.estimate_pose import parse_pose
from utils.params import param_mean, param_std
from utils.render import get_depths_image, cget_depths_image, cpncc, crender_colors
from utils.paf import gen_img_paf
import argparse
import torch.backends.cudnn as cudnn
from simple_dataset import McDataset
from torch.utils.data import DataLoader

STD_SIZE = 120

def main(args):
    # 1. load pre-tained model
    checkpoint_fp = 'models/phase1_wpdc_vdc.pth.tar'
    arch = 'mobilenet_1'

    checkpoint = torch.load(checkpoint_fp, map_location=lambda storage, loc: storage)['state_dict']
    model = getattr(mobilenet_v1, arch)(num_classes=62)  # 62 = 12(pose) + 40(shape) +10(expression)

    model_dict = model.state_dict()
    # because the model is trained by multiple gpus, prefix module should be removed
    for k in checkpoint.keys():
        model_dict[k.replace('module.', '')] = checkpoint[k]
    model.load_state_dict(model_dict)
    if args.mode == 'gpu':
        cudnn.benchmark = True
        model = model.cuda()
    model.eval()

    tri = sio.loadmat('visualize/tri.mat')['tri']
    transform = transforms.Compose([ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])

    if not os.path.exists(args.save_dir):
        os.mkdir(args.save_dir)

    # 2. parse images list and landmark
    lmk_file = args.lmk_file
    ts = time.time()
    rank_land, rank_img_list, start, end = parse_quality_list_part(lmk_file, args.world_size, args.rank, args.resume_idx)
    print('parse land file in {:.3f} seconds'.format(time.time() - ts))

    # for batch processing
    print('World size {}, rank {}, start from {}, end with {}'.format(args.world_size, args.rank, start, end))
    dataset = McDataset(rank_img_list, rank_land, transform=transform, std_size=STD_SIZE)
    dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=2, pin_memory=True)

    for img_idx, (inputs, ori_imgs, img_fps, roi_boxes) in enumerate(tqdm(dataloader)):

        # forward: one step
        with torch.no_grad():
            if args.mode == 'gpu':
                inputs = inputs.cuda()
            params = model(inputs)
            params = params.cpu().numpy()

        roi_boxes = roi_boxes.numpy()
        outputs_roi_boxes = roi_boxes
        if args.bbox_init == 'two':
            step_two_ori_imgs = []
            step_two_roi_boxes = []
            ori_imgs = ori_imgs.numpy()
            for ii in range(params.shape[0]):
                # 68 pts
                pts68 = predict_68pts(params[ii], roi_boxes[ii])

                # two-step for more accurate bbox to crop face
                roi_box = parse_roi_box_from_landmark(pts68)
                img_step2 = crop_img(ori_imgs[ii], roi_box)
                img_step2 = cv2.resize(img_step2, dsize=(STD_SIZE, STD_SIZE), interpolation=cv2.INTER_LINEAR)
                # input = transform(img_step2).unsqueeze(0)
                step_two_ori_imgs.append(transform(img_step2))
                step_two_roi_boxes.append(roi_box)
            with torch.no_grad():
                step_two_ori_imgs = torch.stack(step_two_ori_imgs, dim=0)
                inputs = step_two_ori_imgs
                if args.mode == 'gpu':
                    inputs = inputs.cuda()
                params = model(inputs)
                params = params.cpu().numpy()
            outputs_roi_boxes = step_two_roi_boxes

        # dump results
        if args.dump_param:
            for img_fp, param, roi_box in zip(img_fps, params, outputs_roi_boxes):
                split = img_fp.split('/')
                save_name = os.path.join(args.save_dir, '{}.txt'.format(os.path.splitext(split[-1])[0]))
                this_param = param * param_std + param_mean
                this_param = np.concatenate((this_param, roi_box))
                this_param.tofile(save_name, sep=' ')

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='3DDFA inference pipeline')
    parser.add_argument('-m', '--mode', default='gpu', type=str, help='gpu or cpu mode')
    parser.add_argument('--bbox_init', default='two', type=str,
                        help='one|two: one-step bbox initialization or two-step')
    parser.add_argument('--dump_2d_img', default='false', type=str2bool, help='whether to save 3d rendered image')
    parser.add_argument('--dump_param', default='true', type=str2bool, help='whether to save param')
    parser.add_argument('--save_dir', default='results', type=str, help='dir to save result')
    parser.add_argument('--lmk_file', default='quality_list', type=str, help='landmarks file')
    parser.add_argument('--rank', default=0, type=int, help='used when parallel run')
    parser.add_argument('--world_size', default=1, type=int, help='used when parallel run')
    parser.add_argument('--resume_idx', default=0, type=int)
    parser.add_argument('--batch_size', default=80, type=int, help='batch size')

    args = parser.parse_args()
    main(args)