"""

:mod:`meshes` -- Discretization
===============================

Everything related to meshes appropriate for the multigrid solver.

"""
# Copyright 2018-2020 The emg3d Developers.
#
# This file is part of emg3d.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License.  You may obtain a copy
# of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
# License for the specific language governing permissions and limitations under
# the License.

import numpy as np
from copy import deepcopy
from scipy import optimize

__all__ = ['TensorMesh', 'get_hx_h0', 'get_cell_numbers', 'get_stretched_h',
           'get_domain', 'get_hx']

class TensorMesh:
    """Rudimentary mesh for multigrid calculation.

    The tensor-mesh :class:`discretize.TensorMesh` is a powerful tool,
    including sophisticated mesh-generation possibilities in 1D, 2D, and 3D,
    plotting routines, and much more. However, in the multigrid solver we have
    to generate a mesh at each level, many times over and over again, and we
    only need a very limited set of attributes. This tensor-mesh class provides
    all required attributes. All attributes here are the same as their
    counterparts in :class:`discretize.TensorMesh` (both in name and value).

    .. warning::
        This is a slimmed-down version of :class:`discretize.TensorMesh`, meant
        principally for internal use by the multigrid modeller. It is highly
        recommended to use :class:`discretize.TensorMesh` to create the input
        meshes instead of this class. There are no input-checks carried out
        here, and there is only one accepted input format for `h` and `x0`.

    Parameters
    ----------
    h : list of three ndarrays
        Cell widths in [x, y, z] directions.

    x0 : ndarray of dimension (3, )
        Origin (x, y, z).

    """

    def __init__(self, h, x0):
        """Initialize the mesh."""
        self.x0 = x0

        # Width of cells.
        self.hx = h[0]
        self.hy = h[1]
        self.hz = h[2]

        # Cell related properties.
        self.nCx = int(self.hx.size)
        self.nCy = int(self.hy.size)
        self.nCz = int(self.hz.size)
        self.vnC = np.array([self.hx.size, self.hy.size, self.hz.size])
        self.nC = int(self.vnC.prod())
        self.vectorCCx = np.r_[0, self.hx[:-1].cumsum()]+self.hx*0.5+self.x0[0]
        self.vectorCCy = np.r_[0, self.hy[:-1].cumsum()]+self.hy*0.5+self.x0[1]
        self.vectorCCz = np.r_[0, self.hz[:-1].cumsum()]+self.hz*0.5+self.x0[2]

        # Node related properties.
        self.nNx = self.nCx + 1
        self.nNy = self.nCy + 1
        self.nNz = self.nCz + 1
        self.vnN = np.array([self.nNx, self.nNy, self.nNz], dtype=int)
        self.nN = int(self.vnN.prod())
        self.vectorNx = np.r_[0., self.hx.cumsum()] + self.x0[0]
        self.vectorNy = np.r_[0., self.hy.cumsum()] + self.x0[1]
        self.vectorNz = np.r_[0., self.hz.cumsum()] + self.x0[2]

        # Edge related properties.
        self.vnEx = np.array([self.nCx, self.nNy, self.nNz], dtype=int)
        self.vnEy = np.array([self.nNx, self.nCy, self.nNz], dtype=int)
        self.vnEz = np.array([self.nNx, self.nNy, self.nCz], dtype=int)
        self.nEx = int(self.vnEx.prod())
        self.nEy = int(self.vnEy.prod())
        self.nEz = int(self.vnEz.prod())
        self.vnE = np.array([self.nEx, self.nEy, self.nEz], dtype=int)
        self.nE = int(self.vnE.sum())

    def __repr__(self):
        """Simple representation."""
        return (f"TensorMesh: {self.nCx} x {self.nCy} x {self.nCz} "
                f"({self.nC:,})")

    def copy(self):
        """Return a copy of the TensorMesh."""
        return TensorMesh.from_dict(self.to_dict(True))

    def to_dict(self, copy=False):
        """Store the necessary information of the TensorMesh in a dict."""
        out = {'hx': self.hx, 'hy': self.hy, 'hz': self.hz, 'x0': self.x0,
               '__class__': self.__class__.__name__}
        if copy:
            return deepcopy(out)
        else:
            return out

    @classmethod
    def from_dict(cls, inp):
        """Convert dictionary into :class:`TensorMesh` instance.

        Parameters
        ----------
        inp : dict
            Dictionary as obtained from :func:`TensorMesh.to_dict`.
            The dictionary needs the keys `hx`, `hy`, `hz`, and `x0`.

        Returns
        -------
        obj : :class:`TensorMesh` instance

        """
        try:
            return cls(h=[inp['hx'], inp['hy'], inp['hz']], x0=inp['x0'])
        except KeyError as e:
            print(f"* ERROR   :: Variable {e} missing in `inp`.")
            raise

    @property
    def vol(self):
        """Construct cell volumes of the 3D model as 1D array."""
        if getattr(self, '_vol', None) is None:
            self._vol = (self.hx[None, None, :]*self.hy[None, :, None] *
                         self.hz[:, None, None]).ravel()
        return self._vol

def get_hx_h0(freq, res, domain, fixed=0., possible_nx=None, min_width=None,
              pps=3, alpha=None, max_domain=100000., raise_error=True, verb=1,
              return_info=False):
    r"""Return cell widths and origin for given parameters.

    Returns cell widths for the provided frequency, resistivity, domain extent,
    and other parameters using a flexible amount of cells. See input parameters
    for more details. A maximum of three hard/fixed boundaries can be provided
    (one of which is the grid center).

    The minimum cell width is calculated through :math:`\delta/\rm{pps}`, where
    the skin depth is given by :math:`\delta = 503.3 \sqrt{\rho/f}`, and the
    parameter `pps` stands for 'points-per-skindepth'. The minimum cell width
    can be restricted with the parameter `min_width`.

    The actual calculation domain adds a buffer zone around the (survey)
    domain. The thickness of the buffer is six times the skin depth. The field
    is basically zero after two wavelengths. A wavelength is
    :math:`2\pi\delta`, hence roughly 6 times the skin depth. Taking a factor 6
    gives therefore almost two wavelengths, as the field travels to the
    boundary and back. The actual buffer thickness can be steered with the
    `res` parameter.

    One has to take into account that the air is very resistive, which has to
    be considered not just in the vertical direction, but also in the
    horizontal directions, as the airwave will bounce back from the sides
    otherwise. In the marine case this issue reduces with increasing water
    depth.

    See Also
    --------
    get_stretched_h : Get `hx` for a fixed number `nx` and within a fixed
                      domain.

    Parameters
    ----------

    freq : float
        Frequency (Hz) to calculate the skin depth. The skin depth is a concept
        defined in the frequency domain. If a negative frequency is provided,
        it is assumed that the calculation is carried out in the Laplace
        domain. To calculate the skin depth, the value of `freq` is then
        multiplied by :math:`-2\pi`, to simulate the closest
        frequency-equivalent.

    res : float or list
        Resistivity (Ohm m) to calculate the skin depth. The skin depth is
        used to calculate the minimum cell width and the boundary thicknesses.
        Up to three resistivities can be provided:

        - float: Same resistivity for everything;
        - [min_width, boundaries];
        - [min_width, left boundary, right boundary].

    domain : list
        Contains the survey-domain limits [min, max]. The actual calculation
        domain consists of this domain plus a buffer zone around it, which
        depends on frequency and resistivity.

    fixed : list, optional
        Fixed boundaries, one, two, or maximum three values. The grid is
        centered around the first value. Hence it is the center location with
        the smallest cell. Two more fixed boundaries can be added, at most one
        on each side of the first one.
        Default is 0.

    possible_nx : list, optional
        List of possible numbers of cells. See :func:`get_cell_numbers`.
        Default is ``get_cell_numbers(500, 5, 3)``, which corresponds to
        [16, 24, 32, 40, 48, 64, 80, 96, 128, 160, 192, 256, 320, 384].

    min_width : float, list or None, optional
        Minimum cell width restriction:

        - None : No restriction;
        - float : Fixed to this value, ignoring skin depth and `pps`.
        - list [min, max] : Lower and upper bounds.

        Default is None.

    pps : int, optional
        Points per skindepth; minimum cell width is calculated via
        `dmin = skindepth/pps`.
        Default = 3.

    alpha : list, optional
        Maximum alpha and step size to find a good alpha. The first value is
        the maximum alpha of the survey domain, the second value is the maximum
        alpha for the buffer zone, and the third value is the step size.
        Default = [1, 1.5, .01], hence no stretching within the survey domain
        and a maximum stretching of 1.5 in the buffer zone; step size is 0.01.

    max_domain : float, optional
        Maximum calculation domain from fixed[0] (usually source position).
        Default is 100,000.

    raise_error : bool, optional
        If True, an error is raised if no suitable grid is found. Otherwise it
        just prints a message and returns None's.
        Default is True.

    verb : int, optional
        Verbosity, 0 or 1.
        Default = 1.

    return_info : bool
        If True, a dictionary is returned with some grid info (min and max
        cell width and alpha).

    Returns
    -------
    hx : ndarray
        Cell widths of mesh.

    x0 : float
        Origin of the mesh.

    info : dict
        Dictionary with mesh info; only if ``return_info=True``.

        Keys:

        - `dmin`: Minimum cell width;
        - `dmax`: Maximum cell width;
        - `amin`: Minimum alpha;
        - `amax`: Maximum alpha.

    """
    # Get variables with default lists:
    if alpha is None:
        alpha = [1, 1.5, 0.01]
    if possible_nx is None:
        possible_nx = get_cell_numbers(500, 5, 3)

    # Cast resistivity value(s).
    res = np.array(res, ndmin=1)
    if res.size == 1:
        res_arr = np.array([res[0], res[0], res[0]])
    elif res.size == 2:
        res_arr = np.array([res[0], res[1], res[1]])
    else:
        res_arr = np.array([res[0], res[1], res[2]])

    # Cast and check fixed.
    fixed = np.array(fixed, ndmin=1)
    if fixed.size > 2:

        # Check length.
        if fixed.size > 3:
            print("\n* ERROR   :: Maximum three fixed boundaries permitted.\n"
                  f"             Provided: {fixed.size}.")
            raise ValueError("Wrong input for fixed")

        # Sort second and third, so it doesn't matter how it was provided.
        fixed = np.array([fixed[0], max(fixed[1:]), min(fixed[1:])])

        # Check side.
        if np.sign(np.diff(fixed[:2])) == np.sign(np.diff(fixed[::2])):
            print("\n* ERROR   :: 2nd and 3rd fixed boundaries have to be "
                  "left and right of the first one.\n             "
                  f"Provided: [{fixed[0]}, {fixed[1]}, {fixed[2]}]")
            raise ValueError("Wrong input for fixed")

    # Calculate skin depth.
    skind = 503.3*np.sqrt(res_arr/abs(freq))
    if freq < 0:  # For Laplace-domain calculations.
        skind /= np.sqrt(2*np.pi)

    # Minimum cell width.
    dmin = skind[0]/pps
    if min_width is not None:  # Respect user input.
        min_width = np.array(min_width, ndmin=1)
        if min_width.size == 1:
            dmin = min_width
        else:
            dmin = np.clip(dmin, *min_width)

    # Survey domain; contains all sources and receivers.
    domain = np.array(domain, dtype=float)

    # Calculation domain; big enough to avoid boundary effects.
    # To avoid boundary effects we want the signal to travel two wavelengths
    # from the source to the boundary and back to the receiver.
    # => 2*pi*sd ~ 6.3*sd = one wavelength => signal is ~ 0.2 %.
    # Two wavelengths we can safely assume it is zero.
    #
    # The air does not follow the concept of skin depth, as it is a wave rather
    # than diffusion. For this is the factor `max_domain`, which restricts
    # the domain in each direction to this value from the center.

    # (a) Source to edges of domain.
    dist_in_domain = abs(domain - fixed[0])

    # (b) Two wavelengths.
    two_lambda = skind[1:]*4*np.pi

    # (c) Required buffer, additional to domain.
    dist_buff = np.max([np.zeros(2), (two_lambda - dist_in_domain)/2], axis=0)

    # (d) Add buffer to domain.
    calc_domain = np.array([domain[0]-dist_buff[0], domain[1]+dist_buff[1]])

    # (e) Restrict total domain to max_domain.
    calc_domain[0] = max(calc_domain[0], fixed[0]-max_domain)
    calc_domain[1] = min(calc_domain[1], fixed[0]+max_domain)

    # Initiate flag if terminated.
    finished = False

    # Initiate alpha variables for survey and calculation domains.
    sa, ca = 1.0, 1.0

    # Loop over possible cell numbers from small to big.
    for nx in np.unique(possible_nx):

        # Loop over possible alphas for domain.
        for sa in np.arange(1.0, alpha[0]+alpha[2]/2, alpha[2]):

            # Get current stretched grid cell sizes.
            thxl = dmin*sa**np.arange(nx)  # Left of origin.
            thxr = dmin*sa**np.arange(nx)  # Right of origin.

            # 0. Adjust stretching for fixed boundaries.
            if fixed.size > 1:  # Move mesh to first fixed boundary.
                t_nx = np.r_[fixed[0], fixed[0]+np.cumsum(thxr)]
                ii = np.argmin(abs(t_nx-fixed[1]))
                thxr *= abs(fixed[1]-fixed[0])/np.sum(thxr[:ii])

            if fixed.size > 2:  # Move mesh to second fixed boundary.
                t_nx = np.r_[fixed[0], fixed[0]-np.cumsum(thxl)]
                ii = np.argmin(abs(t_nx-fixed[2]))
                thxl *= abs(fixed[2]-fixed[0])/np.sum(thxl[:ii])

            # 1. Fill from center to left domain.
            nl = np.sum((fixed[0]-np.cumsum(thxl)) > domain[0])+1

            # 2. Fill from center to right domain.
            nr = np.sum((fixed[0]+np.cumsum(thxr)) < domain[1])+1

            # 3. Get remaining number of cells and check termination criteria.
            nsdc = nl+nr  # Number of domain cells.
            nx_remain = nx-nsdc

            # Not good, try next.
            if nx_remain <= 0:
                continue

            # Create the current hx-array.
            hx = np.r_[thxl[:nl][::-1], thxr[:nr]]
            hxo = np.r_[thxl[:nl][::-1], thxr[:nr]]

            # Get actual domain:
            asurv_domain = [fixed[0]-np.sum(thxl[:nl]),
                            fixed[0]+np.sum(thxr[:nr])]
            x0 = float(fixed[0]-np.sum(thxl[:nl]))

            # Get actual stretching (differs in case of fixed layers).
            sa_adj = np.max([hx[1:]/hx[:-1], hx[:-1]/hx[1:]])

            # Loop over possible alphas for calc_domain.
            for ca in np.arange(sa, alpha[1]+alpha[2]/2, alpha[2]):

                # 4. Fill to left calc_domain.
                thxl = hx[0]*ca**np.arange(1, nx_remain+1)
                nl = np.sum((asurv_domain[0]-np.cumsum(thxl)) >
                            calc_domain[0])+1

                # 5. Fill to right calc_domain.
                thxr = hx[-1]*ca**np.arange(1, nx_remain+1)
                nr = np.sum((asurv_domain[1]+np.cumsum(thxr)) <
                            calc_domain[1])+1

                # 6. Get remaining number of cells and check termination
                # criteria.
                ncdc = nl+nr  # Number of calc_domain cells.
                nx_remain2 = nx-nsdc-ncdc

                if nx_remain2 < 0:  # Not good, try next.
                    continue

                # Create hx-array.
                nl += int(np.floor(nx_remain2/2))  # If uneven, add one cell
                nr += int(np.ceil(nx_remain2/2))   # more on the right.
                hx = np.r_[thxl[:nl][::-1], hx, thxr[:nr]]

                # Calculate origin.
                x0 = float(asurv_domain[0]-np.sum(thxl[:nl]))

                # Mark it as finished and break out of the loop.
                finished = True
                break

            if finished:
                break

        if finished:
            break

    # Check finished and print info about found grid.
    if not finished:
        # Throw message if no solution was found.
        print("\n* ERROR   :: No suitable grid found; relax your criteria.\n")
        if raise_error:
            raise ArithmeticError("No grid found!")
        else:
            hx, x0 = None, None

    elif verb > 0:
        print(f"   Skin depth ", end="")
        if res.size == 1:
            print(f"         [m] : {skind[0]:.0f}")
        elif res.size == 2:
            print(f"(m/l-r)  [m] : {skind[0]:.0f} / {skind[1]:.0f}")
        else:
            print(f"(m/l/r)  [m] : {skind[0]:.0f} / {skind[1]:.0f} / "
                  f"{skind[2]:.0f}")
        print(f"   Survey domain       [m] : {domain[0]:.0f} - "
              f"{domain[1]:.0f}")
        print(f"   Calculation domain  [m] : {calc_domain[0]:.0f} - "
              f"{calc_domain[1]:.0f}")
        print(f"   Final extent        [m] : {x0:.0f} - "
              f"{x0+np.sum(hx):.0f}")
        extstr = f"   Min/max cell width  [m] : {min(hx):.0f} / "
        alstr = f"   Alpha survey"
        nrstr = "   Number of cells "
        if not np.isclose(sa, sa_adj):
            sastr = f"{sa:.3f} ({sa_adj:.3f})"
        else:
            sastr = f"{sa:.3f}"
        print(extstr+f"{max(hxo):.0f} / {max(hx):.0f}")
        print(alstr+f"/calc       : {sastr} / {ca:.3f}")
        print(nrstr+f"(s/c/r) : {nx} ({nsdc}/{ncdc}/{nx_remain2})")
        print()

    if return_info:
        if not fixed.size > 1:
            sa_adj = sa

        info = {'dmin': dmin,
                'dmax': np.nanmax(hx),
                'amin': np.nanmin([ca, sa, sa_adj]),
                'amax': np.nanmax([ca, sa, sa_adj])}

        return hx, x0, info
    else:
        return hx, x0

def get_cell_numbers(max_nr, max_prime=5, min_div=3):
    r"""Returns 'good' cell numbers for the multigrid method.

    'Good' cell numbers are numbers which can be divided by 2 as many times as
    possible. At the end there will be a low prime number.

    The function adds all numbers :math:`p 2^n \leq M` for :math:`p={2, 3, ...,
    p_\text{max}}` and :math:`n={n_\text{min}, n_\text{min}+1, ..., \infty}`;
    :math:`M, p_\text{max}, n_\text{min}` correspond to `max_nr`, `max_prime`,
    and `min_div`, respectively.

    Parameters
    ----------
    max_nr : int
        Maximum number of cells.

    max_prime : int
        Highest permitted prime number p for p*2^n. {2, 3, 5, 7} are good upper
        limits in order to avoid too big lowest grids in the multigrid method.
        Default is 5.

    min_div : int
        Minimum times the number can be divided by two.
        Default is 3.

    Returns
    -------
    numbers : array
        Array containing all possible cell numbers from lowest to highest.

    """
    # Primes till 20.
    primes = np.array([2, 3, 5, 7, 11, 13, 17, 19])

    # Sanity check; 19 is already ridiculously high.
    if max_prime > primes[-1]:
        print(f"* ERROR   :: Highest prime is {max_prime}, "
              "please use a value < 20.")
        raise ValueError("Highest prime too high")

    # Restrict to max_prime.
    primes = primes[primes <= max_prime]

    # Get possible values.
    # Currently restricted to prime*2**30 (for prime=2 => 1,073,741,824 cells).
    numbers = primes[:, None]*2**np.arange(min_div, 30)

    # Get unique values.
    numbers = np.unique(numbers)

    # Restrict to max_nr and return.
    return numbers[numbers <= max_nr]

def get_stretched_h(min_width, domain, nx, x0=0, x1=None, resp_domain=False):
    """Return cell widths for a stretched grid within the domain.

    Returns `nx` cell widths within `domain`, where the minimum cell width is
    `min_width`. The cells are not stretched within `x0` and `x1`, and outside
    uses a power-law stretching. The actual stretching factor and the number of
    cells left and right of `x0` and `x1` are find in a minimization process.

    The domain is not completely respected. The starting point of the domain
    is, but the endpoint of the domain might slightly shift (this is more
    likely the case for small `nx`, for big `nx` the shift should be small).
    The new endpoint can be obtained with ``domain[0]+np.sum(hx)``. If you want
    the domain to be respected absolutely, set ``resp_domain=True``. However,
    be aware that this will introduce one stretch-factor which is different
    from the other stretch factors, to accommodate the restriction. This
    one-off factor is between the left- and right-side of `x0`, or, if `x1` is
    provided, just after `x1`.

    See Also
    --------
    get_hx_x0 : Get `hx` and `x0` for a flexible number of `nx` with
                given bounds.

    Parameters
    ----------

    min_width : float
        Minimum cell width. If x1 is provided, the actual minimum cell width
        might be smaller than min_width.

    domain : list
        [start, end] of model domain.

    nx : int
        Number of cells.

    x0 : float
        Center of the grid. `x0` is restricted to `domain`.
        Default is 0.

    x1 : float
        If provided, then no stretching is applied between `x0` and `x1`. The
        non-stretched part starts at `x0` and stops at the first possible
        location at or after `x1`. `x1` is restricted to `domain`. This will
        min_width so that an integer number of cells fit within x0 and x1.

    resp_domain : bool
        If False (default), then the domain-end might shift slightly to assure
        that the same stretching factor is applied throughout. If set to True,
        however, the domain is respected absolutely. This will introduce one
        stretch-factor which is different from the other stretch factors, to
        accommodate the restriction. This one-off factor is between the left-
        and right-side of `x0`, or, if `x1` is provided, just after `x1`.

    Returns
    -------
    hx : ndarray
        Cell widths of mesh.

    """

    # Cast to arrays
    domain = np.array(domain, dtype=float)
    x0 = np.array(x0, dtype=float)
    x0 = np.clip(x0, *domain)  # Restrict to model domain
    min_width = np.array(min_width, dtype=float)
    if x1 is not None:
        x1 = np.array(x1, dtype=float)
        x1 = np.clip(x1, *domain)  # Restrict to model domain

    # If x1 is provided (a part is not stretched)
    if x1 is not None:

        # Store original values
        xlim_orig = domain.copy()
        nx_orig = int(nx)
        x0_orig = x0.copy()
        h_min_orig = min_width.copy()

        # Get number of non-stretched cells
        n_nos = int(np.ceil((x1-x0)/min_width))

        # Re-calculate min_width to fit with x0-x1-limits:
        min_width = (x1-x0)/n_nos

        # Subtract one cell, because the standard scheme provides one
        # min_width-cell.
        n_nos -= 1

        # Reset x0, because the first min_width comes from normal scheme
        x0 += min_width

        # Reset xmax for normal scheme
        domain[1] -= n_nos*min_width

        # Reset nx for normal scheme
        nx -= n_nos

        # If there are not enough points reset to standard procedure. The limit
        # of five is arbitrary. However, nx should be much bigger than five
        # anyways, otherwise stretched grid doesn't make sense.
        if nx <= 5:
            print("Warning :: Not enough points for non-stretched part,"
                  "ignoring therefore `x1`.")
            domain = xlim_orig
            nx = nx_orig
            x0 = x0_orig
            x1 = None
            min_width = h_min_orig

    # Get stretching factor (a = 1+alpha).
    if min_width == 0 or min_width > np.diff(domain)/nx:
        # If min_width is bigger than the domain-extent divided by nx, no
        # stretching is required at all.
        alpha = 0
    else:

        # Wrap _get_dx into a minimization function to call with fsolve.
        def find_alpha(alpha, min_width, args):
            """Find alpha such that min(hx) = min_width."""
            return min(get_hx(alpha, *args))/min_width-1

        # Search for best alpha, must be at least 0
        args = (domain, nx, x0)
        alpha = max(0, optimize.fsolve(find_alpha, 0.02, (min_width, args)))

    # With alpha get actual cell spacing with `resp_domain` to respect the
    # users decision.
    hx = get_hx(alpha, domain, nx, x0, resp_domain)

    # Add the non-stretched center if x1 is provided
    if x1 is not None:
        hx = np.r_[hx[: np.argmin(hx)], np.ones(n_nos)*min_width,
                   hx[np.argmin(hx):]]

    # Print warning min_width could not be respected.
    if abs(hx.min() - min_width) > 0.1:
        print(f"Warning :: Minimum cell width ({np.round(hx.min(), 2)} m) is "
              "below `min_width`, because `nx` is too big for `domain`.")

    return hx

def get_domain(x0=0, freq=1, res=0.3, limits=None, min_width=None,
               fact_min=0.2, fact_neg=5, fact_pos=None):
    r"""Get domain extent and minimum cell width as a function of skin depth.

    Returns the extent of the calculation domain and the minimum cell width as
    a multiple of the skin depth, with possible user restrictions on minimum
    calculation domain and range of possible minimum cell widths.

    .. math::

            \delta &= 503.3 \sqrt{\frac{\rho}{f}} , \\
            x_\text{start} &= x_0-k_\text{neg}\delta , \\
            x_\text{end} &= x_0+k_\text{pos}\delta , \\
            h_\text{min} &= k_\text{min} \delta .

    Parameters
    ----------

    x0 : float
        Center of the calculation domain. Normally the source location.
        Default is 0.

    freq : float
        Frequency (Hz) to calculate the skin depth. The skin depth is a concept
        defined in the frequency domain. If a negative frequency is provided,
        it is assumed that the calculation is carried out in the Laplace
        domain. To calculate the skin depth, the value of `freq` is then
        multiplied by :math:`-2\pi`, to simulate the closest
        frequency-equivalent.

        Default is 1 Hz.

    res : float, optional
        Resistivity (Ohm m) to calculate skin depth.
        Default is 0.3 Ohm m (sea water).

    limits : None or list
        [start, end] of model domain. This extent represents the minimum extent
        of the domain. The domain is therefore only adjusted if it has to reach
        outside of [start, end].
        Default is None.

    min_width : None, float, or list of two floats
        Minimum cell width is calculated as a function of skin depth:
        fact_min*sd. If `min_width` is a float, this is used. If a list of
        two values [min, max] are provided, they are used to restrain
        min_width. Default is None.

    fact_min, fact_neg, fact_pos : floats
        The skin depth is multiplied with these factors to estimate:

            - Minimum cell width (`fact_min`, default 0.2)
            - Domain-start (`fact_neg`, default 5), and
            - Domain-end (`fact_pos`, defaults to `fact_neg`).

    Returns
    -------

    h_min : float
        Minimum cell width.

    domain : list
        Start- and end-points of calculation domain.

    """

    # Set fact_pos to fact_neg if not provided.
    if fact_pos is None:
        fact_pos = fact_neg

    # Calculate the skin depth.
    skind = 503.3*np.sqrt(res/abs(freq))
    if freq < 0:  # For Laplace-domain calculations.
        skind /= np.sqrt(2*np.pi)

    # Estimate minimum cell width.
    h_min = fact_min*skind
    if min_width is not None:  # Respect user input.
        if np.array(min_width).size == 1:
            h_min = min_width
        else:
            h_min = np.clip(h_min, *min_width)

    # Estimate calculation domain.
    domain = [x0-fact_neg*skind, x0+fact_pos*skind]
    if limits is not None:  # Respect user input.
        domain = [min(limits[0], domain[0]), max(limits[1], domain[1])]

    return h_min, domain

def get_hx(alpha, domain, nx, x0, resp_domain=True):
    r"""Return cell widths for given input.

    Find the number of cells left and right of `x0`, `nl` and `nr`
    respectively, for the provided alpha. For this, we solve

    .. math::   \frac{x_\text{max}-x_0}{x_0-x_\text{min}} =
                \frac{a^{nr}-1}{a^{nl}-1}

    where :math:`a = 1+\alpha`.

    Parameters
    ----------

    alpha : float
        Stretching factor `a` is given by ``a=1+alpha``.

    domain : list
        [start, end] of model domain.

    nx : int
        Number of cells.

    x0 : float
        Center of the grid. `x0` is restricted to `domain`.

    resp_domain : bool
        If False (default), then the domain-end might shift slightly to assure
        that the same stretching factor is applied throughout. If set to True,
        however, the domain is respected absolutely. This will introduce one
        stretch-factor which is different from the other stretch factors, to
        accommodate the restriction. This one-off factor is between the left-
        and right-side of `x0`, or, if `x1` is provided, just after `x1`.

    Returns
    -------
    hx : ndarray
        Cell widths of mesh.

    """
    if alpha <= 0.:  # If alpha <= 0: equal spacing (no stretching at all)
        hx = np.ones(nx)*np.diff(np.squeeze(domain))/nx

    else:            # Get stretched hx
        a = alpha+1

        # Get hx depending if x0 is on the domain boundary or not.
        if np.isclose(x0, domain[0]) or np.isclose(x0, domain[1]):
            # Get al a's
            alr = np.diff(domain)*alpha/(a**nx-1)*a**np.arange(nx)
            if x0 == domain[1]:
                alr = alr[::-1]

            # Calculate differences
            hx = alr*np.diff(domain)/sum(alr)

        else:
            # Find number of elements left and right by solving:
            #     (xmax-x0)/(x0-xmin) = a**nr-1/(a**nl-1)
            nr = np.arange(2, nx+1)
            er = (domain[1]-x0)/(x0-domain[0]) - (a**nr[::-1]-1)/(a**nr-1)
            nl = np.argmin(abs(np.floor(er)))+1
            nr = nx-nl

            # Get all a's
            al = a**np.arange(nl-1, -1, -1)
            ar = a**np.arange(1, nr+1)

            # Calculate differences
            if resp_domain:
                # This version honours domain[0] and domain[1], but to achieve
                # this it introduces one stretch-factor which is different from
                # all the others between al to ar.
                hx = np.r_[al*(x0-domain[0])/sum(al),
                           ar*(domain[1]-x0)/sum(ar)]
            else:
                # This version moves domain[1], but each stretch-factor is
                # exactly the same.
                fact = (x0-domain[0])/sum(al)  # Take distance from al.
                hx = np.r_[al, ar]*fact

                # Note: this hx is equivalent as providing the following h
                # to TensorMesh:
                # h = [(min_width, nl-1, -a), (min_width, n_nos+1),
                #      (min_width, nr, a)]

    return hx


from __future__ import absolute_import, division, print_function, unicode_literals
from typing import Optional

import numpy as np
import os

from astropy.io import fits
from astropy.table import Table

from .. import this_project as P
from .. import utils as ut

__all__ = ["Planck14Data", "Planck14DataAlt", "Planck14Model", "Maniyar18Model"]

class CIBxCIB:

    _l = None  # multipole
    _Cl = None  # angular power
    _dCl = None  # uncertainty on the angular power
    _Dl = None  # l(l+1)/2pi * Cl
    _dDl = None  # uncertainty on the angular power
    _l3Cl = None  # l**3 * Cl

    _S = None  # shot noise level
    _dS = None  # uncertainty on the shot noise

    _raw_table = None  # raw table, taken from publications, emails, etc

    def __init__(self, freq1, freq2=None, unit="Jy^2/sr"):
        if unit not in ["Jy^2/sr", "MJy^2/sr", "K^2.sr", "uK^2.sr"]:
            raise ValueError(
                'Unit must be either "Jy^2/sr", ' '"MJy^2/sr", "uK^2.sr" or "K^2.sr"'
            )
        self.unit = unit

        self.freq1 = self.freq2int(freq1)
        if freq2 is None:
            self.freq2 = freq1
        else:
            self.freq2 = self.freq2int(freq2)

    # Methods
    #########
    def freq2int(self, freq):
        if isinstance(freq, str):
            return int(freq)
        else:
            if isinstance(freq, int):
                return freq
            else:
                raise TypeError("freq must be int or str")

    # Properties
    ############
    @property
    def freqstr(self):
        self._freqstr = "x".join(
            (
                str(max([int(self.freq1), int(self.freq2)])),
                str(min([int(self.freq1), int(self.freq2)])),
            )
        )
        return self._freqstr

    @property
    def l(self):
        if self._l is None:
            self._l = self.raw_table[:, 0]
        return self._l

    @property
    def Cl(self):
        return None

    @property
    def dCl(self):
        return None

    @property
    def Dl(self):
        if self._Dl is None:
            self._Dl = self.l * (self.l + 1.0) / 2.0 / np.pi * self.Cl
        return self._Dl

    @property
    def l3Cl(self):
        if self._l3Cl is None:
            self._l3Cl = self.l ** 3 * self.Cl
        return self._l3Cl

    @property
    def Jy2K(self):
        self._Jy2K = P.Jy2K
        return self._Jy2K

    @property
    def K2Jy(self):
        self._K2Jy = P.K2Jy
        return self._K2Jy

class Planck14DataAlt(CIBxCIB):
    def __init__(self, freq1, freq2=None, unit="Jy^2/sr"):
        super(Planck14Data, self).__init__(freq1, freq2=freq2, unit=unit)

        self.Cl_contains_SN = True

    # Properties
    ############
    @property
    def raw_table(self):
        if self._raw_table is None:
            self._raw_table = np.loadtxt(
                os.path.join(P.PACKAGE_DIR, "resources/cibxcib/Planck14_data.txt")
            )
        return self._raw_table

    @property
    def Cl(self):
        if self._Cl is None:
            # Native unit is Jy^2/sr
            self._Cl = self.raw_table[:, self._freq2col(self.freqstr)].copy()

            # We add the shot noise, which is also in Jy^2/sr
            self._Cl += self.S

            if self.unit in ["K^2.sr", "uK^2.sr"]:
                self._Cl *= self.Jy2K[str(self.freq1)] * self.Jy2K[str(self.freq2)]

                if self.unit == "uK^2.sr":
                    self._Cl *= 1.0e12

            if self.unit == "MJy^2/sr":
                self._Cl /= 1.0e12

        return self._Cl

    @property
    def S(self):
        """
        Shot noise, units are Jy^2/sr. Taken from
        Planck (2014 XXX)
        """
        if self._S is None:
            self._S = {
                "857x857": 5364,
                "545x545": 1690,
                "353x353": 262,
                "217x217": 21,
                "3000x3000": 9585,
                "857x545": 2702,
                "857x353": 953,
                "857x217": 181,
                "545x353": 626,
                "545x217": 121,
                "353x217": 54,
                "3000x857": 4158,
                "3000x545": 1449,
                "3000x353": 411,
                "3000x217": 95,
            }

            if self.unit in ["K^2.sr", "uK^2.sr"]:
                self._S *= self.Jy2K[str(self.freq1)] * self.Jy2K[str(self.freq2)]

                if self.unit == "uK^2.sr":
                    self._S *= 1.0e12

            if self.unit == "MJy^2/sr":
                self._S /= 1.0e12

        return self._S[self.freqstr]

    @property
    def dS(self):
        """
        Shot noise
        """
        if self._dS is None:
            self._dS = {
                "857x857": 343,
                "545x545": 45,
                "353x353": 8,
                "217x217": 2,
                "3000x3000": 1090,
                "857x545": 124,
                "857x353": 54,
                "857x217": 6,
                "545x353": 19,
                "545x217": 6,
                "353x217": 3,
                "3000x857": 443,
                "3000x545": 176,
                "3000x353": 48,
                "3000x217": 11,
            }
            if self.unit in ["K^2.sr", "uK^2.sr"]:
                self._dS *= self.Jy2K[str(self.freq1)] * self.Jy2K[str(self.freq2)]
                if self.unit == "uK^2.sr":
                    self._dS *= 1.0e12

            if self.unit == "MJy^2/sr":
                self._dS /= 1.0e12

        return self._dS[self.freqstr]

    # Methods
    #########
    def _freq2col(self, freqstr):
        mapping = {
            "857x857": 1,
            "545x545": 2,
            "353x353": 3,
            "217x217": 4,
            "3000x3000": 5,
            "857x545": 6,
            "857x353": 7,
            "857x217": 8,
            "545x353": 9,
            "545x217": 10,
            "353x217": 11,
            "3000x857": 12,
            "3000x545": 13,
            "3000x353": 14,
            "3000x217": 15,
            "100x100": 16,
            "857x100": 17,
            "545x100": 18,
            "353x100": 19,
            "217x100": 20,
            "143x143": 21,
            "857x143": 22,
            "545x143": 23,
            "353x143": 24,
            "217x143": 25,
            "143x100": 26,
        }

        return mapping[self.freqstr]

class Planck14Model(CIBxCIB):
    def __init__(self, freq1, freq2=None, unit="Jy^2/sr"):
        super(Planck14Model, self).__init__(freq1, freq2=freq2, unit=unit)

        self.Cl_contains_SN = True

    # Properties
    ############
    @property
    def raw_table(self):
        if self._raw_table is None:
            self._raw_table = np.loadtxt(
                os.path.join(P.PACKAGE_DIR, "resources/cibxcib/Planck14_model.txt")
            )
        return self._raw_table

    @property
    def Cl(self):
        if self._Cl is None:
            # native unit is be Jy^2/sr
            self._Cl = self.raw_table[:, self._freq2col(self.freqstr)].copy()

            # Possibly convert the units
            if self.unit in ["K^2.sr", "uK^2.sr"]:
                self._Cl *= self.Jy2K[str(self.freq1)] * self.Jy2K[str(self.freq2)]
                if self.unit == "uK^2.sr":
                    self._Cl *= 1.0e12

            if self.unit == "MJy^2/sr":
                self._Cl /= 1.0e12

            # Apply the correction factor to the PR1 calibration
            # self._Cl /= (
            #     ut.PLANCK_PR1PR3_CALCORR[str(self.freq1)]
            #     * ut.PLANCK_PR1PR3_CALCORR[str(self.freq2)]
            # )
        return self._Cl

    # Methods
    #########
    def _freq2col(self, freqstr):
        # l, 353x353, 353x545, 353x857, 545x545, 545x857, 857x857
        mapping = {
            "353x353": 1,
            "545x353": 2,
            "857x353": 3,
            "545x545": 4,
            "857x545": 5,
            "857x857": 6,
        }

        return mapping[self.freqstr]

class Maniyar18Model(CIBxCIB):
    def __init__(self, freq1, freq2=None, unit="Jy^2/sr"):
        super(Maniyar18Model, self).__init__(freq1, freq2=freq2, unit=unit)

        self.Cl_contains_SN = True
        self.Cl_contains_1halo = True

    # Properties
    ############
    @property
    def raw_table(self):
        if self._raw_table is None:
            self._raw_table = fits.open(
                os.path.join(
                    P.PACKAGE_DIR, "resources/cibxcib/Maniyar18_model_crosspowers.dat"
                )
            )
        return self._raw_table

    @property
    def l(self):
        if self._l is None:
            self._l = self.raw_table[1].data
        return self._l

    @property
    def Cl(self):
        if self._Cl is None:
            # native unit is Jy^2/sr
            self._Cl = self.raw_table[0].data[
                self._freq2col(str(self.freq1)), self._freq2col(str(self.freq2))
            ]

            # Possibly convert the units
            if self.unit in ["K^2.sr", "uK^2.sr"]:
                self._Cl *= self.Jy2K[str(self.freq1)] * self.Jy2K[str(self.freq2)]
                if self.unit == "uK^2.sr":
                    self._Cl *= 1.0e12

            if self.unit == "MJy^2/sr":
                self._Cl /= 1.0e12

        return self._Cl

    # Methods
    #########
    def _freq2col(self, freq):
        # 217x217, 353x353, 545x545, 857x857, 3000x3000
        mapping = {
            "100": 0,
            "143": 1,
            "217": 2,
            "353": 3,
            "545": 4,
            "857": 5,
            "3000": 6,
        }

        return mapping[freq]

class Planck14Data(CIBxCIB):
    def __init__(self, freq1, freq2=None, unit="Jy^2/sr"):
        super(Planck14Data, self).__init__(freq1, freq2=freq2, unit=unit)
        self.Cl_contains_SN = True

    # Properties
    ############
    @property
    def raw_table(self):
        if self._raw_table is None:
            self._raw_table = Table.read(
                os.path.join(
                    P.PACKAGE_DIR, "resources/cibxcib/Planck14_data_frompaper.txt"
                ),
                format="csv",
                delimiter=";",
            )
        return self._raw_table

    @property
    def l(self):
        if self._l is None:
            self._l = self.raw_table["ell"].data.data

        return self._l

    @property
    def Cl(self):
        if self._Cl is None:
            # Native unit is Jy^2/sr
            self._Cl = self.raw_table[f"{self.freq1}x{self.freq2}"].data.data

            # We add the shot noise, which is also in Jy^2/sr
            # self._Cl += self.S

            if self.unit in ["K^2.sr", "uK^2.sr"]:
                self._Cl *= self.Jy2K[str(self.freq1)] * self.Jy2K[str(self.freq2)]

                if self.unit == "uK^2.sr":
                    self._Cl *= 1.0e12

            if self.unit == "MJy^2/sr":
                self._Cl /= 1.0e12

            # Apply the correction factor to the PR1 calibration
            # self._Cl /= (
            #     ut.PLANCK_PR1PR3_CALCORR[str(self.freq1)]
            #     * ut.PLANCK_PR1PR3_CALCORR[str(self.freq2)]
            # )

        return self._Cl

    @property
    def dCl(self):
        if self._dCl is None:
            # Native unit is Jy^2/sr
            self._dCl = self.raw_table[f"d{self.freq1}x{self.freq2}"].data.data

            if self.unit in ["K^2.sr", "uK^2.sr"]:
                self._dCl *= self.Jy2K[str(self.freq1)] * self.Jy2K[str(self.freq2)]

                if self.unit == "uK^2.sr":
                    self._dCl *= 1.0e12

            if self.unit == "MJy^2/sr":
                self._dCl /= 1.0e12

            # Apply the correction factor to the PR1 calibration
            # self._dCl /= (
            #     ut.PLANCK_PR1PR3_CALCORR[str(self.freq1)]
            #     * ut.PLANCK_PR1PR3_CALCORR[str(self.freq2)]
            # )

        return self._dCl

    @property
    def dDl(self):
        if self._dDl is None:
            self._dDl = self.l * (self.l + 1.0) / 2.0 / np.pi * self.dCl
        return self._dDl

class Mak17Model(CIBxCIB):
    def __init__(
        self,
        freq1: str,
        freq2: Optional[str] = None,
        lmax: Optional[int] = None,
        unit: str = "uK2.sr",
        mask="mask40",
    ):
        super().__init__(freq1, freq2=freq2, unit=unit)

        if lmax is None:
            self.lmax = 3000
        else:
            self.lmax = lmax

        #         self.freq1 = freq1
        #         if freq2 is None:
        #             self.freq2 = freq1
        #         else:
        #             self.freq2 = freq2
        #
        self.mask = mask

    @property
    def l(self):
        if self._l is None:
            self._l = np.arange(self.lmax)
        return self._l

    @staticmethod
    def model(ells, A_cib, A_ps, gamma):
        """
        Returns Dl^total = Dl^CIB + Dl^PS
        The Dl^PS are given as value at ell=2000, hence we need to convert this back to Cls
        and then apply the appropriate scaling.
        """
        return (
            A_cib * np.power(ells / 2000.0, gamma)
            + A_ps / 2000 / 2001 * ells * ells
            + 1.0
        )

    @property
    def mask(self):
        return self._mask

    @mask.setter
    def mask(self, val):
        if not hasattr(self, "_mask"):
            allowed = ["mask30", "mask40", "mask50"]
            if val not in allowed:
                raise ValueError(f"mask must be in {allowed}")
            self._mask = val
        else:
            raise RuntimeError("mask can only be set once during initialization.")

    @staticmethod
    def get_rho_cib(freq1, freq2):
        d = {
            "353x353": 1.0,
            "545x545": 1.0,
            "857x857": 1.0,
            "353x545": 0.975,
            "353x857": 0.892,
            "545x857": 0.949,
        }

        return d[f"{freq1}x{freq2}"]

    @staticmethod
    def get_rho_ps(freq1, freq2):
        d = {
            "353x353": 1.0,
            "545x545": 1.0,
            "857x857": 1.0,
            "353x545": 0.98,
            "353x857": 0.86,
            "545x857": 0.97,
        }

        return d[f"{freq1}x{freq2}"]

    @property
    def _raw_data(self):
        # Define mask30
        mask30 = {
            "gamma": 0.51,
            "cib": {"353": 2.5e3, "545": 4.5e5, "857": 1.09e9},
            "ps": {"353": 2.15e3, "545": 3.54e5, "857": 7.2e8},
            "cal": {"353": 1.0, "545": 1.05, "857": 1.01},
        }

        # Define mask40
        mask40 = {
            "gamma": 0.53,
            "cib": {"353": 2.56e3, "545": 4.47e5, "857": 1.09e9},
            "ps": {"353": 2.1e3, "545": 3.42e5, "857": 7.34e8},
            "cal": {"353": 1.0, "545": 1.03, "857": 1.01},
        }

        d = dict(
            mask30=mask30,
            mask40=mask40,
            #                mask50=mask50,
        )

        return d

    def calibration(self, freq1: int, freq2: int):
        raw_data = self._raw_data[self.mask]
        cal1, cal2 = raw_data["cal"][str(freq1)], raw_data["cal"][str(freq2)]

        return cal1 * cal2

    def get_model_parameters(self):
        raw_data = self._raw_data[self.mask]

        # Correlation coefficients
        rho_cib = self.get_rho_cib(self.freq1, self.freq2)
        rho_ps = self.get_rho_ps(self.freq1, self.freq2)

        # Amplitudes
        gamma = raw_data["gamma"]

        A_cib = rho_cib * np.sqrt(
            raw_data["cib"][str(self.freq1)] * raw_data["cib"][str(self.freq2)]
        )

        A_ps = rho_ps * np.sqrt(
            raw_data["ps"][str(self.freq1)] * raw_data["ps"][str(self.freq2)]
        )

        return gamma, A_cib, A_ps

    @property
    def Cl(self):
        """
        We copy the dipole for the monopole, because it would naturally be inf due to
        the division by zero.
        """
        _Cl = self.Dl / (self.l * (self.l + 1)) * 2.0 * np.pi
        _Cl = np.concatenate([_Cl[1:2], _Cl[1:]])

        return _Cl

    @property
    def Dl(self):
        if self._Dl is None:
            gamma, A_cib, A_ps = self.get_model_parameters()
            self._Dl = self.model(self.l, A_cib, A_ps, gamma)
            self._Dl /= self.calibration(self.freq1, self.freq2)

            if self.unit == "K^2.sr":
                self._Dl /= 1.0e12
            if self.unit == "Jy^2/sr":
                self._Dl *= 1.0e-12 * (
                    self.K2Jy[str(self.freq1)] * self.K2Jy[str(self.freq2)]
                )
            if self.unit == "MJy^2/sr":
                self._Dl *= self.K2Jy[str(self.freq1)] * self.K2Jy[str(self.freq2)]
        return self._Dl

class Mak17(Mak17Model):
    """We keep this copy of the Mak17Model to ensure downward compatibility."""
    pass


#!/bin/python3.6

import subprocess
import os
import networkx as nx
import random
from typing import Optional, Callable, Dict, List
import re
import signal
from glob import glob
from time import sleep

from utils import filled_in, TreeDecomposition, pairs, stream_bn, get_bn_stats, \
    filter_read_bn, compute_complexity_width, read_jkl, write_jkl, get_domain_sizes, \
    CWDecomposition

import matplotlib.pyplot as plt
from networkx.drawing.nx_agraph import pygraphviz_layout

SOLVER_DIR = "../solvers"
SCORE_PATN = re.compile(r"New improvement! (?P<score>[\d.-]+) \(after (?P<time>[\d.-]+) s.\)")
CHECKPOINT_MINUTES = 30  # in minutes
CHECKPOINT_INTERVAL = int(CHECKPOINT_MINUTES*60)  # in seconds
CHECKPOINT_RETRY_INTERVAL = 5  # in seconds

# not that crucial, if you consider the stdout updates of the heuristic as mere
# triggers for reading new (not necessarily next) bn from .res file
PARSE_RETRY_INTERVAL = 0.01  # in seconds

class BayesianNetwork(object):
    def __init__(self, input_file, data=None, parents=None):
        self.parents = dict() if parents is None else parents
        self.input_file = input_file
        self.data = data
        self.sum_scores, self.best_score, self.offsets = get_bn_stats(self.input_file)
        self.best_norm_score = self.best_score - sum(self.offsets.values())
        self._score = None
        self._dag = None

    def _clear_cached(self):
        self._score = self._dag = None

    def compute_all_scores(self, subset: set = None) -> Dict[int, float]:
        if self.data is None:
            if subset is None:
                data = stream_bn(self.input_file, normalize=True)
            else:
                data = filter_read_bn(self.input_file, subset, normalize=True).items()
        else:
            data = self.data.items()  # data is assumed to be normalized
        scores = dict()
        for node, psets in data:
            if node in self.parents:  # could be a bn with a subset of nodes
                if subset is None or node in subset:
                    scores[node] = psets[self.parents[node]] + self.offsets[node]
        return scores

    def compute_score(self, subset: set=None) -> float:
        """recompute score based on parents sets"""
        scores = self.compute_all_scores(subset)
        return sum(scores.values())

    @property
    def score(self) -> float:
        if self._score is None:
            self._score = self.compute_score()
        return self._score

    def add(self, node, parents):
        self.parents[node] = frozenset(parents)
        self._clear_cached()

    def recompute_dag(self):
        dag = nx.DiGraph()
        for node in self.parents.keys():
            dag.add_node(node)
            for parent in self.parents[node]:
                assert not dag.has_edge(node, parent), f"cyclic parent set {node}<->{parent}"
                dag.add_edge(parent, node)
        self._dag = dag

    @property
    def dag(self) -> nx.DiGraph:
        if self._dag is None:
            self.recompute_dag()
        return self._dag

    def verify(self):
        assert nx.is_directed_acyclic_graph(self.dag), "network is not acyclic"

    def get_moralized(self):
        moral = self.dag.to_undirected()
        for node in self.parents.keys():
            for p1, p2 in pairs(self.parents[node]):
                moral.add_edge(p1, p2)
        return moral

    def draw(self, subset=None):
        if subset is None:
            dag = self.dag
        else:
            dag = self.dag.subgraph(subset)
        pos = pygraphviz_layout(dag, prog='dot')
        nx.draw(dag, pos, with_labels=True)
        plt.show()

    def replace(self, newbn: 'BayesianNetwork'):
        for node, new_parents in newbn.parents.items():
            self.parents[node] = new_parents
        self._clear_cached()

class TWBayesianNetwork(BayesianNetwork):
    def __init__(self, input_file, tw=0, elim_order=None, td=None, *args, **kwargs):
        super().__init__(input_file, *args, **kwargs)
        self.tw = tw
        self.elim_order: List[int] = elim_order
        self._td: Optional[TreeDecomposition] = td

    @property
    def td(self) -> TreeDecomposition:
        if self._td is None:
            raise RuntimeError("td requested before finalizing (call `.done`)")
        return self._td

    def done(self):
        """
        compute and store tree decomp and width based on
        elim_order(default: reverse topological order of the dag)
        """
        if self.elim_order is None:
            self.elim_order = list(nx.topological_sort(self.dag))[::-1]
        self._td = TreeDecomposition(self.get_moralized(), self.elim_order, self.tw)
        if self.tw <= 0:
            self.tw = self.td.width

    def verify(self, verify_treewidth=True):
        super().verify()
        if verify_treewidth:
            self.td.verify(graph=self.get_moralized())

    def get_triangulated(self, elim_order=None):
        if elim_order is None: elim_order = self.elim_order
        assert elim_order is not None, "elim order not specified"
        moral = self.get_moralized().subgraph(elim_order)
        triangulated, max_degree = filled_in(moral, elim_order)
        assert max_degree <= self.tw, f"tw {self.tw} < {max_degree} for elim order {elim_order}"
        return triangulated

class CWBayesianNetwork(TWBayesianNetwork):
    def __init__(self, input_file, cwidth, datfile, elim_order=None, td=None, *args, **kwargs):
        super().__init__(input_file, *args, **kwargs)
        self.tw = cwidth
        self.elim_order: List[int] = elim_order
        self._td: Optional[TreeDecomposition] = td
        self.domain_sizes = get_domain_sizes(datfile)

    def done(self):
        if self.elim_order is None:
            self.elim_order = list(nx.topological_sort(self.dag))[::-1]
        self._td = CWDecomposition(self.get_moralized(), self.elim_order, self.tw, self.domain_sizes)

def inject_tuples(basejkl, extra_tuples, destname, strong_injection=False):
    src = read_jkl(basejkl, normalize=False)
    inject_count = 0
    for node, pset in extra_tuples.items():
        parents, score = pset
        if strong_injection or parents not in src[node]:
            src[node][parents] = score
            inject_count += 1
    write_jkl(src, destname)

def parse_res(filename: str, treewidth: int, outfile: str, cwidth=-1,
              add_extra_tuples=False, augfile: str = "augmented.jkl",
              datfile=None, retry=True, debug=False) -> TWBayesianNetwork:
    """
    Parse a .res file containing a solution BN. Optionally merge the parent set
    tuples from the jkl file `filename` and the the res file `outfile` and
    save as a temporary file `augfile` (only when `add_extra_tuples` is True)

    :param filename: input jkl file
    :param treewidth: treewidth bound
    :param outfile: res file containing BN
    :param cwidth: cwidth bound, use -1 to ignore and simply parse a TWBN
    :param add_extra_tuples: whether to inject tuples from outfile into filename
    :param augfile: name of temporary file containing merged set of tuples,
                    (only applicable if add_extra_tuples is True)
    :param datfile: path to data file (with header rows)
    :param retry: keep retrying until file is non-empty
    :param debug: debugging
    :return: parsed BN as a TWBayesianNetwork
    """
    elim_order = None
    tuples = []
    extra_tuples = dict()
    score = None
    # keep retrying until file is non-empty
    while retry and os.path.getsize(outfile) == 0:
        sleep(PARSE_RETRY_INTERVAL)
    with open(outfile) as out:
        for line in out:
            if not line.strip():
                continue
            elif line.startswith("Score:"):
                score = float(line.split()[1].strip())
                break
            elif line.startswith("elim-order:"):
                elim_order = line.split()[1].strip("()").split(",")
                elim_order = list(map(int, elim_order))
                if debug: print("elim-order:", elim_order)
            else:
                vertex, rest = line.split(":")
                vertex = int(vertex)
                rest = rest.strip().split()
                if len(rest) == 1:
                    local_score = float(rest[0])
                    parents = frozenset()
                else:
                    local_score, parents = rest
                    local_score = float(local_score)
                    parents = frozenset(map(int, parents.strip("()").split(",")))
                # bn.add(int(vertex), map(int, parents))
                tuples.append((vertex, parents))
                if add_extra_tuples:
                    extra_tuples[vertex] = (parents, local_score)
                if debug: print(f"v:{vertex}, sc:{local_score}, par:{parents})")
    if add_extra_tuples:
        inject_tuples(filename, extra_tuples, augfile, True)
        if debug: print("temporary merged file saved as", augfile)
        input_file = augfile
    else:
        input_file = filename
    if cwidth > 0:
        assert datfile, "datfile needed for cwidth"
        bn = CWBayesianNetwork(cwidth=cwidth, input_file=input_file, datfile=datfile)
    else:
        bn = TWBayesianNetwork(tw=treewidth, input_file=input_file)
    if elim_order is not None: bn.elim_order = elim_order
    for node, parents in tuples: bn.add(node, parents)
    bn.done()
    if score is not None:
        bn._score = score
    else:
        print(f"warning: score not found in {outfile}")
    return bn

def parse_res_score(outfile: str):
    with open(outfile) as out:
        for line in out:
            if not line.strip():
                continue
            elif line.startswith("Score:"):
                score = float(line.split()[1].strip())
                return score
    return None

def write_res(bn: BayesianNetwork, outfile, write_elim_order=False, debug=False):
    nodes = sorted(bn.dag.nodes())
    scores = bn.compute_all_scores()
    with open(outfile, "w") as out:
        if write_elim_order and isinstance(bn, TWBayesianNetwork):
            elim_order = bn.td.recompute_elim_order()
            out.write(f"elim-order: ({','.join(map(str, elim_order))})\n")
        for node in nodes:
            parents = bn.parents[node]
            pstr = ""
            if parents:
                pstr = " (" + ",".join(map(str, parents)) + ")"
            out.write(f"{node}: {scores[node]:.2f} {pstr}\n")
        out.write(f"\nScore: {bn.compute_score():.3f}\n")

def write_net(bn: BayesianNetwork, datfile, tempprefix="lltemp",
              use_res="", debug=False):
    netfile = f"{tempprefix}.net"
    if use_res:
        resfile = use_res
    else:
        resfile = f"{tempprefix}.res"
        write_res(bn, resfile)
    basecmd = ["java", "-jar", os.path.join(SOLVER_DIR, "blip.jar"), "parle"]
    args = ["-d", datfile, "-r", resfile, "-n", netfile]
    cmd = basecmd + args
    if debug: print("running parle, cmd:", cmd)
    proc = subprocess.run(cmd, stdout=subprocess.PIPE)
    if proc.returncode == 0:  # success
        if debug: print("file generated:", netfile)
    else:  # error
        print("error encountered during 'parle', returncode:", proc.returncode)
        return None

def run_blip(filename, treewidth, outfile="temp.res", timeout=10, seed=0,
             solver="kg", logfile="temp.log", debug=False) -> TWBayesianNetwork:
    basecmd = ["java", "-jar", os.path.join(SOLVER_DIR, "blip.jar"),
               f"solver.{solver}", "-v", "1"]
    args = ["-j", filename, "-w", str(treewidth), "-r", outfile,
            "-t", str(timeout), "-seed", str(seed), "-l", logfile]
    cmd = basecmd + args
    if debug: print("running blip, cmd:", cmd)
    proc = subprocess.run(cmd, stdout=subprocess.PIPE)
    if proc.returncode == 0:  # success
        return parse_res(filename, treewidth, outfile, debug)
    else:  # error
        print("error encountered, returncode:", proc.returncode)
        return None

def activate_checkpoints(bnprovider, save_as):
    def alarm_handler(signalnum, frame):
        print("alarm triggered")
        try:
            bn = bnprovider()
        except IndexError:
            print("bn res file invalid (probably got overwritten)")
            signal.alarm(CHECKPOINT_RETRY_INTERVAL)  # snooze alarm
        else:
            patn = save_as.replace(".res", "*.res")
            prefix, ext = os.path.splitext(save_as)
            prev_checkpoint = 0
            for check_file in glob(patn):
                try:
                    saved_at = int(check_file.replace(prefix, "").replace(ext, ""))
                except ValueError:
                    continue
                else:
                    prev_checkpoint = max(prev_checkpoint, saved_at)
            new_checkpoint = prev_checkpoint + CHECKPOINT_MINUTES
            fname = save_as.replace(".res", f"{new_checkpoint}.res")
            print("saving checkpoint to", fname)
            write_res(bn, fname, write_elim_order=True)
            signal.alarm(CHECKPOINT_INTERVAL)  # reset alarm
    signal.signal(signal.SIGALRM, alarm_handler)  # register handler
    signal.alarm(CHECKPOINT_INTERVAL)  # set first alarm
    print(f"checkpointing for {CHECKPOINT_MINUTES}m activated")

def monitor_blip(filename, treewidth, logger: Callable, outfile="temp.res",
                 timeout=10, seed=0, solver="kg", datfile=None,
                 cwidth=0, onlyfilter=False, save_as="", debug=False):
    """
    Run BLIP in monitoring mode, where each new score update is logged

    :param filename: path to jkl file
    :param treewidth: treewidth bound (ignored if in CWIDTH_MODE)
    :param logger: logging function to be used
    :param outfile: path to .res file containing learned network (volatile)
    :param timeout: total time limit on blip computation
    :param seed: random seed passed on to blip
    :param solver: blip sub-algo to use (for tw: [kg, ka, kmax], cwidth: [old, greedy, max])
    :param datfile: path to data file (with header rows, for domain sizes)
    :param cwidth: cwidth bound (if positive, activates CWIDTH_MODE)
    :param onlyfilter: only use pset filtering algo (ignored if not in CWIDTH_MODE)
    :param save_as: filepath prefix to use for saving checkpoint solutions
    :param debug: enable debug mode
    """
    CWIDTH_MODE = cwidth > 0
    if CWIDTH_MODE:
        assert solver in ("old", "greedy", "max"), \
            f"invalid solver({solver}) for monitor_blip in CWIDTH_MODE"
        basecmd = ["java", "-jar", os.path.join(SOLVER_DIR, "blip-cw.jar"),
                   f"solver.kg.adv", "-v", "1", "-src", f"cwidth-{solver}"]
        args = ["-j", filename, "-d", datfile, "-w", "0", "-cw", str(cwidth),
                "-r", outfile, "-t", str(timeout), "-seed", str(seed)]
        if onlyfilter: args.append("-filter")
    else:
        basecmd = ["java", "-jar", os.path.join(SOLVER_DIR, "blip.jar"),
                   f"solver.{solver}", "-v", "1"]
        args = ["-j", filename, "-w", str(treewidth), "-r", outfile,
                "-t", str(timeout), "-seed", str(seed)]
    cmd = basecmd + args
    if debug: print("monitoring blip, cmd:", " ".join(cmd))
    if save_as:
        if CWIDTH_MODE:
            bnprovider = lambda: parse_res(filename, 0, outfile, cwidth=cwidth,
                                           datfile=datfile)
        else:
            bnprovider = lambda: parse_res(filename, treewidth, outfile)
        activate_checkpoints(bnprovider, save_as)
    domain_sizes = None if datfile is None else get_domain_sizes(datfile)
    with subprocess.Popen(cmd, stdout=subprocess.PIPE, bufsize=1,
                          universal_newlines=True) as proc:
        for line in proc.stdout:
            if debug: print("got line:", line, end='')
            match = SCORE_PATN.match(line)
            if match:
                score = float(match['score'])
                logdata = {"score": score}
                if domain_sizes is not None:
                    try:
                        bn = bnprovider()
                    except IndexError:  # todo: not reached anymore, retries, rethink
                        print("bn res file invalid (probably got overwritten)")
                        cw = acw = -1
                    else:
                        tw = bn.td.compute_width()
                        cw = compute_complexity_width(bn.td, domain_sizes)
                        logdata["tw"] = tw
                        logdata["cw"] = cw
                logger(logdata)
    print(f"done returncode: {proc.returncode}")

def start_blip_proc(filename, treewidth, outfile="temp.res",
                    timeout=10, seed=0, solver="kg", debug=False) -> subprocess.Popen:
    basecmd = ["java", "-jar", os.path.join(SOLVER_DIR, "blip.jar"),
               f"solver.{solver}", "-v", "1"]
    args = ["-j", filename, "-w", str(treewidth), "-r", outfile,
            "-t", str(timeout), "-seed", str(seed)]
    cmd = basecmd + args
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, bufsize=1,
                            universal_newlines=True)
    os.set_blocking(proc.stdout.fileno(), False)  # set stdout to be non-blocking
    if debug: print(f"starting blip proc, pid: {proc.pid}, \ncmd: {cmd}")
    return proc

def start_blip_proc_cw(filename, datfile, cwidth, outfile="temp.res", timeout=10,
                       seed=0, searcher="greedy", onlyfilter=False, debug=False) -> subprocess.Popen:
    basecmd = ["java", "-jar", os.path.join(SOLVER_DIR, "blip-cw.jar"),
               f"solver.kg.adv", "-v", "1", "-src", f"cwidth-{searcher}"]
    args = ["-j", filename, "-d", datfile, "-w", "0", "-cw", str(cwidth),
            "-r", outfile, "-t", str(timeout), "-seed", str(seed)]
    if onlyfilter: args.append(["-filter"])
    cmd = basecmd + args
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, bufsize=1,
                            universal_newlines=True)
    os.set_blocking(proc.stdout.fileno(), False)  # set stdout to be non-blocking
    if debug: print(f"starting blip-cw proc, pid: {proc.pid}, \ncmd: {cmd}")
    return proc

def check_blip_proc(proc: subprocess.Popen, debug=False) -> float:
    score = float("-infinity")
    rc = proc.poll()
    if rc is not None:
        if debug: print(f"blip proc already terminated with rc:", rc)
        return score
    while True:
        line = proc.stdout.readline()
        if line:
            if debug: print("got line:", line, end='')
            match = SCORE_PATN.match(line)
            if match:
                score = float(match['score'])
        else:
            rc = proc.poll()
            if rc is not None:
                if debug: print(f"no output and proc is completed (rc: {rc})")
            else:
                if debug: print(f"no output and proc is still running")
            break
    return score

def stop_blip_proc(proc: subprocess.Popen):
    proc.terminate()
    proc.stdout.close()
    if proc.stderr is not None: proc.stderr.close()
    if proc.stdin is not None: proc.stdin.close()
    proc.wait()

if __name__ == '__main__':
    # res = run_blip("../past-work/blip-publish/data/child-5000.jkl", 10, timeout=2, seed=4)
    # dag = res.get_dag()
    # print("acyclic", nx.is_directed_acyclic_graph(dag))
    # fig, axes = plt.subplots(1, 3)
    # nx.draw(dag, with_labels=True, ax=axes[0])
    # moral = res.get_moralized()
    # nx.draw(moral, with_labels=True, ax=axes[1])
    # tri = res.get_triangulated()
    # nx.draw(tri, with_labels=True, ax=axes[2])
    # fig.show()
    for seed in range(1,100):
        print(seed)
        res = run_blip("../past-work/blip-publish/data/child-5000.jkl", 10, timeout=2, seed=seed, debug=False)
        tri = res.get_triangulated()
    print("done")