#! /usr/bin/env python
# -*- coding: utf-8 -*-
#
# graph_tool -- a general graph manipulation python module
#
# Copyright (C) 2006-2021 Tiago de Paula Peixoto <tiago@skewed.de>
#
# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU Lesser General Public License as published by the Free
# Software Foundation; either version 3 of the License, or (at your option) any
# later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
# details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.

import numpy
from . util import *
from . mcmc import *
from . bisection import *
from . blockmodel import *
from . overlap_blockmodel import *
from . layered_blockmodel import *
from . nested_blockmodel import *

def default_args(mcmc_args={}, anneal_args={}, mcmc_equilibrate_args={},
                 shrink_args={}, mcmc_multilevel_args={}, overlap=False):
    mcmc_args = dict(dict(beta=numpy.inf, c=0, niter=5, allow_vacate=False,
                          entropy_args=dict(dl=True)), **mcmc_args)
    if overlap:
        mcmc_args = dict(mcmc_args, bundled=True)
    mcmc_equilibrate_args = dict(dict(wait=1, nbreaks=1, epsilon=1e-4,
                                      mcmc_args=mcmc_args),
                                 **dmask(mcmc_equilibrate_args,
                                         ["mcmc_args"]))
    shrink_entropy_args = dict(dict(mcmc_args["entropy_args"]),
                               **shrink_args.get("entropy_args", {}))
    if not shrink_entropy_args.get("dense", False):
         shrink_entropy_args["multigraph"] = False
    shrink_args = dict(dict(entropy_args=shrink_entropy_args, niter=20),
                       **dmask(shrink_args, ["entropy_args"]))
    mcmc_multilevel_args = \
            dict(dict(r=1.3, anneal=False,
                      shrink_args=shrink_args,
                      mcmc_equilibrate_args=mcmc_equilibrate_args,
                      anneal_args=anneal_args),
                 **dmask(mcmc_multilevel_args,
                         ["shrink_args", "mcmc_equilibrate_args",
                          "anneal_args"]))
    return mcmc_multilevel_args


def get_states(g, B_min=None, B_max=None, b_min=None, b_max=None, deg_corr=True,
               overlap=False, nonoverlap_init=True, layers=False, clabel=None,
               state_args={}, mcmc_multilevel_args={}):

    if B_min is None:
        if clabel is None:
            B_min = 1
        else:
            B_min = len(set(clabel.fa))
            if b_min is None:
                b_min = clabel

    _B_max = g.num_vertices()
    if overlap and not nonoverlap_init:
        _B_max = 2 * g.num_edges()

    if B_max is None:
        B_max = _B_max
        if overlap and nonoverlap_init and b_max is None:
            b_max = g.vertex_index.copy("int")

    if layers:
        State = LayeredBlockState
        state_args = dict(state_args, overlap=overlap)
    elif overlap:
        State = OverlapBlockState
    else:
        State = BlockState

    if b_max is not None:
        max_state = State(g, b=b_max, deg_corr=deg_corr, clabel=clabel,
                          **dmask(state_args, ["B", "b", "deg_corr", "clabel"]))
    else:
        max_state = State(g, B=_B_max, deg_corr=deg_corr, clabel=clabel,
                          **dmask(state_args,["B", "b", "deg_corr", "clabel"]))

    if max_state.get_nonempty_B() > B_max:
        max_state = mcmc_multilevel(max_state, B=B_max, **mcmc_multilevel_args)

    if b_min is not None:
        min_state = State(g, B=B_min, b=b_min, deg_corr=deg_corr, clabel=clabel,
                          **dmask(state_args, ["B", "b", "deg_corr", "clabel"]))
    elif B_min == 1:
        min_state = State(g, B=1, deg_corr=deg_corr, clabel=clabel,
                          **dmask(state_args, ["B", "b", "deg_corr", "clabel"]))
    else:
        min_state = max_state.copy()

    if min_state.get_nonempty_B() > B_min:
        min_state = mcmc_multilevel(min_state, B=B_min, **mcmc_multilevel_args)

    return min_state, max_state


def minimize_blockmodel_dl(g, B_min=None, B_max=None, b_min=None, b_max=None,
                           deg_corr=True, overlap=False, nonoverlap_init=True,
                           layers=False, state_args={}, bisection_args={},
                           mcmc_args={}, anneal_args={},
                           mcmc_equilibrate_args={}, shrink_args={},
                           mcmc_multilevel_args={}, verbose=False):
    r"""Fit the stochastic block model, by minimizing its description length using an
    agglomerative heuristic.

    Parameters
    ----------
    g : :class:`~graph_tool.Graph`
        The graph.
    B_min : ``int`` (optional, default: ``None``)
        The minimum number of blocks.
    B_max : ``int`` (optional, default: ``None``)
        The maximum number of blocks.
    b_min : :class:`~graph_tool.VertexPropertyMap` (optional, default: ``None``)
        The partition to be used with the minimum number of blocks.
    b_max : :class:`~graph_tool.VertexPropertyMap` (optional, default: ``None``)
        The partition to be used with the maximum number of blocks.
    deg_corr : ``bool`` (optional, default: ``True``)
        If ``True``, the degree-corrected version of the model will be used.
    overlap : ``bool`` (optional, default: ``False``)
        If ``True``, the overlapping version of the model will be used.
    nonoverlap_init : ``bool`` (optional, default: ``True``)
        If ``True``, and ``overlap == True`` a non-overlapping initial state
        will be used.
    layers : ``bool`` (optional, default: ``False``)
        If ``True``, the layered version of the model will be used.
    state_args : ``dict`` (optional, default: ``{}``)
        Arguments to be passed to appropriate state constructor (e.g.
        :class:`~graph_tool.inference.blockmodel.BlockState`,
        :class:`~graph_tool.inference.overlap_blockmodel.OverlapBlockState` or
        :class:`~graph_tool.inference.layered_blockmodel.LayeredBlockState`)
    bisection_args : ``dict`` (optional, default: ``{}``)
        Arguments to be passed to :func:`~graph_tool.inference.bisection.bisection_minimize`.
    mcmc_args : ``dict`` (optional, default: ``{}``)
        Arguments to be passed to :meth:`graph_tool.inference.blockmodel.BlockState.mcmc_sweep`,
        :meth:`graph_tool.inference.overlap_blockmodel.OverlapBlockState.mcmc_sweep` or
        :meth:`graph_tool.inference.layered_blockmodel.LayeredBlockState.mcmc_sweep`.
    mcmc_equilibrate_args : ``dict`` (optional, default: ``{}``)
        Arguments to be passed to :func:`~graph_tool.inference.mcmc.mcmc_equilibrate`.
    shrink_args : ``dict`` (optional, default: ``{}``)
        Arguments to be passed to :meth:`graph_tool.inference.blockmodel.BlockState.shrink`,
        :meth:`graph_tool.inference.overlap_blockmodel.OverlapBlockState.shrink` or
        :meth:`graph_tool.inference.layered_blockmodel.LayeredBlockState.shrink`.
    mcmc_multilevel_args : ``dict`` (optional, default: ``{}``)
        Arguments to be passed to :func:`~graph_tool.inference.mcmc.mcmc_multilevel`.
    verbose : ``bool`` or ``tuple`` (optional, default: ``False``)
        If ``True``, progress information will be shown. Optionally, this
        accepts arguments of the type ``tuple`` of the form ``(level, prefix)``
        where ``level`` is a positive integer that specifies the level of
        detail, and ``prefix`` is a string that is prepended to the all output
        messages.

    Returns
    -------
    min_state : :class:`~graph_tool.inference.blockmodel.BlockState` or  :class:`~graph_tool.inference.overlap_blockmodel.OverlapBlockState` or  :class:`~graph_tool.inference.layered_blockmodel.LayeredBlockState`
        State with minimal description length.

    Notes
    -----

    This function is a convenience wrapper around
    :func:`~graph_tool.inference.bisection.bisection_minimize`.

    See [peixoto-efficient-2014]_ for details on the algorithm.

    This algorithm has a complexity of :math:`O(V \ln^2 V)`, where :math:`V` is
    the number of nodes in the network.

    Examples
    --------

    .. testsetup:: mdl

       gt.seed_rng(43)
       np.random.seed(43)

    .. doctest:: mdl

       >>> g = gt.collection.data["polbooks"]
       >>> state = gt.minimize_blockmodel_dl(g)
       >>> state.draw(pos=g.vp["pos"], vertex_shape=state.get_blocks(),
       ...            output="polbooks_blocks_mdl.svg")
       <...>

    .. figure:: polbooks_blocks_mdl.*
       :align: center

       Block partition of a political books network, which minimizes the
       description length of the network according to the degree-corrected
       stochastic blockmodel.


    .. testsetup:: mdl_overlap

       gt.seed_rng(42)
       np.random.seed(42)

    .. doctest:: mdl_overlap

       >>> g = gt.collection.data["polbooks"]
       >>> state = gt.minimize_blockmodel_dl(g, overlap=True)
       >>> state.draw(pos=g.vp["pos"], output="polbooks_overlap_blocks_mdl.svg")
       <...>

    .. figure:: polbooks_overlap_blocks_mdl.*
       :align: center

       Overlapping partition of a political books network, which minimizes the
       description length of the network according to the overlapping
       degree-corrected stochastic blockmodel.


    References
    ----------
    .. [peixoto-efficient-2014] Tiago P. Peixoto, "Efficient Monte Carlo and greedy
       heuristic for the inference of stochastic block models", Phys. Rev. E 89,
       012804 (2014), :doi:`10.1103/PhysRevE.89.012804`, :arxiv:`1310.4378`.
    """

    b_cache = {} # keep a global cache

    mcmc_multilevel_args = \
        default_args(mcmc_args=mcmc_args,
                     anneal_args=anneal_args,
                     mcmc_equilibrate_args=mcmc_equilibrate_args,
                     shrink_args=shrink_args,
                     mcmc_multilevel_args=dict(mcmc_multilevel_args,
                                               b_cache=b_cache),
                     overlap=overlap)

    bisection_args = dict(dict(mcmc_multilevel_args=mcmc_multilevel_args,
                               random_bisection=False),
                          **bisection_args)

    clabel = state_args.get("clabel", None)
    if clabel is None:
        clabel = state_args.get("pclabel", None)

    min_state, max_state = get_states(g, B_min=B_min, B_max=B_max, b_min=b_min,
                                      b_max=b_max, deg_corr=deg_corr,
                                      overlap=overlap,
                                      nonoverlap_init=nonoverlap_init,
                                      layers=layers, clabel=clabel,
                                      state_args=state_args,
                                      mcmc_multilevel_args=mcmc_multilevel_args)

    if B_min is None:
        B_min = 1
    if B_max is None:
        B_max = numpy.inf

    Bs = list(b_cache.keys())
    for B in Bs:
        if B > B_max or B < B_min:
            del b_cache[B]

    state = bisection_minimize([min_state, max_state], verbose=verbose,
                               **bisection_args)

    return state

def minimize_nested_blockmodel_dl(g, B_min=None, B_max=None, b_min=None,
                                  b_max=None, Bs=None, bs=None, deg_corr=True,
                                  overlap=False, nonoverlap_init=True,
                                  layers=False, hierarchy_minimize_args={},
                                  state_args={}, bisection_args={},
                                  mcmc_args={}, anneal_args={},
                                  mcmc_equilibrate_args={}, shrink_args={},
                                  mcmc_multilevel_args={}, verbose=False):
    r"""Fit the nested stochastic block model, by minimizing its description length
    using an agglomerative heuristic.

    Parameters
    ----------
    g : :class:`~graph_tool.Graph`
        The graph.
    B_min : ``int`` (optional, default: ``None``)
        The minimum number of blocks.
    B_max : ``int`` (optional, default: ``None``)
        The maximum number of blocks.
    b_min : :class:`~graph_tool.VertexPropertyMap` (optional, default: ``None``)
        The partition to be used with the minimum number of blocks.
    b_max : :class:`~graph_tool.VertexPropertyMap` (optional, default: ``None``)
        The partition to be used with the maximum number of blocks.
    Bs : ``list`` of ints (optional, default: ``None``)
        If provided, it will correspond to the sizes of the initial hierarchy.
    bs : ``list`` of integer-valued :class:`numpy.ndarray` objects (optional, default: ``None``)
        If provided, it will correspond to the initial hierarchical partition.
    deg_corr : ``bool`` (optional, default: ``True``)
        If ``True``, the degree-corrected version of the model will be used.
    overlap : ``bool`` (optional, default: ``False``)
        If ``True``, the overlapping version of the model will be used.
    nonoverlap_init : ``bool`` (optional, default: ``True``)
        If ``True``, and ``overlap == True`` a non-overlapping initial state
        will be used.
    layers : ``bool`` (optional, default: ``False``)
        If ``True``, the layered version of the model will be used.
    hierarchy_minimize_args : ``dict`` (optional, default: ``{}``)
        Arguments to be passed to :func:`~graph_tool.inference.nested_blockmodel.hierarchy_minimize`.
    state_args : ``dict`` (optional, default: ``{}``)
        Arguments to be passed to appropriate state constructor (e.g.
        :class:`~graph_tool.inference.blockmodel.BlockState`,
        :class:`~graph_tool.inference.overlap_blockmodel.OverlapBlockState` or
        :class:`~graph_tool.inference.layered_blockmodel.LayeredBlockState`)
    bisection_args : ``dict`` (optional, default: ``{}``)
        Arguments to be passed to :func:`~graph_tool.inference.bisection.bisection_minimize`.
    mcmc_args : ``dict`` (optional, default: ``{}``)
        Arguments to be passed to :meth:`graph_tool.inference.blockmodel.BlockState.mcmc_sweep`,
        :meth:`graph_tool.inference.overlap_blockmodel.OverlapBlockState.mcmc_sweep` or
        :meth:`graph_tool.inference.layered_blockmodel.LayeredBlockState.mcmc_sweep`.
    mcmc_equilibrate_args : ``dict`` (optional, default: ``{}``)
        Arguments to be passed to :func:`~graph_tool.inference.mcmc.mcmc_equilibrate`.
    shrink_args : ``dict`` (optional, default: ``{}``)
        Arguments to be passed to :meth:`graph_tool.inference.blockmodel.BlockState.shrink`,
        :meth:`graph_tool.inference.overlap_blockmodel.OverlapBlockState.shrink` or
        :meth:`graph_tool.inference.layered_blockmodel.LayeredBlockState.shrink`.
    mcmc_multilevel_args : ``dict`` (optional, default: ``{}``)
        Arguments to be passed to :func:`~graph_tool.inference.mcmc.mcmc_multilevel`.
    verbose : ``bool`` or ``tuple`` (optional, default: ``False``)
        If ``True``, progress information will be shown. Optionally, this
        accepts arguments of the type ``tuple`` of the form ``(level, prefix)``
        where ``level`` is a positive integer that specifies the level of
        detail, and ``prefix`` is a string that is prepended to the all output
        messages.

    Returns
    -------
    min_state : :class:`~graph_tool.inference.nested_blockmodel.NestedBlockState`
        Nested state with minimal description length.

    Notes
    -----

    This function is a convenience wrapper around
    :func:`~graph_tool.inference.nested_blockmodel.hierarchy_minimize`.

    See [peixoto-hierarchical-2014]_ for details on the algorithm.

    This algorithm has a complexity of :math:`O(V \ln^2 V)`, where :math:`V` is
    the number of nodes in the network.

    Examples
    --------
    .. testsetup:: nested_mdl

       gt.seed_rng(43)
       np.random.seed(43)

    .. doctest:: nested_mdl

       >>> g = gt.collection.data["power"]
       >>> state = gt.minimize_nested_blockmodel_dl(g, deg_corr=True)
       >>> state.draw(output="power_nested_mdl.pdf")
       (...)

    .. testcleanup:: nested_mdl

       conv_png("power_nested_mdl.pdf")

    .. figure:: power_nested_mdl.png
       :align: center
       :width: 60%

       Hierarchical Block partition of a power-grid network, which minimizes
       the description length of the network according to the nested
       (degree-corrected) stochastic blockmodel.


    .. doctest:: nested_mdl_overlap

       >>> g = gt.collection.data["celegansneural"]
       >>> state = gt.minimize_nested_blockmodel_dl(g, deg_corr=True, overlap=True)
       >>> state.draw(output="celegans_nested_mdl_overlap.pdf")
       (...)

    .. testcleanup:: nested_mdl_overlap

       conv_png("celegans_nested_mdl_overlap.pdf")

    .. figure:: celegans_nested_mdl_overlap.png
       :align: center
       :width: 60%

       Overlapping block partition of the *C. elegans* neural network, which
       minimizes the description length of the network according to the nested
       overlapping degree-corrected stochastic blockmodel.

    References
    ----------
    .. [peixoto-hierarchical-2014] Tiago P. Peixoto, "Hierarchical block
       structures and high-resolution model selection in large networks ",
       Phys. Rev. X 4, 011047 (2014), :doi:`10.1103/PhysRevX.4.011047`,
       :arxiv:`1310.4377`.
    """

    mcmc_multilevel_args = \
            default_args(mcmc_args=mcmc_args,
                         anneal_args=anneal_args,
                         mcmc_equilibrate_args=mcmc_equilibrate_args,
                         shrink_args=shrink_args,
                         mcmc_multilevel_args=mcmc_multilevel_args,
                         overlap=overlap)

    if bs is None:
        clabel = state_args.get("clabel", None)
        if clabel is None:
            clabel = state_args.get("pclabel", None)
        min_state, max_state = get_states(g, B_min=B_min, B_max=B_max,
                                          b_min=b_min, b_max=b_max,
                                          deg_corr=deg_corr, overlap=overlap,
                                          nonoverlap_init=nonoverlap_init,
                                          layers=layers, clabel=clabel,
                                          state_args=dmask(state_args,
                                                           ["hstate_args",
                                                            "hentropy_args"]),
                                          mcmc_multilevel_args=mcmc_multilevel_args)
        if b_max is None:
            b_max = max_state.b.fa
        if b_min is None:
            b_min = min_state.b.fa
        if Bs is None:
            bs = [min_state.b.fa, zeros(min_state.b.fa.max() + 1, dtype="int")]
        else:
            bs = []
            bstate = max_state
            for B in Bs:
                bstate = mcmc_multilevel(bstate, B=B,
                                         verbose=verbose_push(verbose,
                                                              ("l = %d " %
                                                               len(bs))),
                                         **mcmc_multilevel_args)
                bs.append(bstate.b.a)
                bstate = bstate.get_block_state()
        State = type(min_state)
    else:
        if layers:
            State = LayeredBlockState
            state_args = dict(state_args, overlap=overlap)
        elif overlap:
            State = OverlapBlockState
        else:
            State = BlockState

    if layers:
        state_args = dict(state_args, overlap=overlap)

    state = NestedBlockState(g, bs=bs,
                             base_type=State,
                             deg_corr=deg_corr,
                             sampling=False,
                             **dmask(state_args, ["deg_corr"]))

    bisection_args = dict(dict(mcmc_multilevel_args=mcmc_multilevel_args,
                               random_bisection=False),
                          **bisection_args)

    hierarchy_minimize(state, B_max=B_max, B_min=B_min, b_max=b_max,
                       b_min=b_min, bisection_args=bisection_args,
                       verbose=verbose,
                       **dmask(hierarchy_minimize_args,
                               ["B_max", "B_min", "bisection_args", "verbose"]))

    return state
