// graph-tool -- a general graph modification and manipulation thingy
//
// Copyright (C) 2006-2026 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/>.

#ifndef GRAPH_SBM_SAMPLE_EDGE_HH
#define GRAPH_SBM_SAMPLE_EDGE_HH

#include <tuple>

#include "graph.hh"
#include "graph_filtering.hh"
#include "graph_util.hh"
#include "dynamic_sampler.hh"

#include "random.hh"

namespace graph_tool
{
using namespace std;
using namespace boost;

template <class State>
class SBMEdgeSampler
{
public:
    typedef typename State::vertex_t vertex_t;

    SBMEdgeSampler(State& state, double esample=.25, double usample=.25)
        : _state(state),
          _vertices(vertices(state._g).first,
                    vertices(state._g).second),
          _N(_vertices.size()),
          _v_in_sampler(is_directed(state._g) ?
                        __v_in_sampler : _v_out_sampler),
          _esample(esample),
          _usample(usample)
    {
        if (_esample == 0 && _usample == 1)
            return;

        for (auto e : edges_range(state._g))
        {
            auto u = source(e, state._g);
            auto v = target(e, state._g);
            if (_edge_pos.find(get_edge(u, v)) != _edge_pos.end() || _state._eweight[e] == 0)
                continue;
            _edges.push_back(get_edge(u, v));
            _edge_pos[_edges.back()] = _edges.size() - 1;
            _E += _state._eweight[e];
        }

        if (_esample == 1 || _usample == 1)
            return;

        for (auto me : edges_range(state._bg))
        {
            auto r = source(me, state._bg);
            auto s = target(me, state._bg);

            std::tie(r, s) = get_edge(r, s);

            size_t mrs = state._mrs[me];
            if (mrs == 0)
                continue;
            _rs_pos[me] = _rs_sampler.insert({r, s}, mrs);

            if (!is_directed(_state._g) && r == s)
                mrs *= 2;

            if (r >= _r_out_sampler.size())
                _r_out_sampler.resize(r + 1);
            _r_out_pos[me] = _r_out_sampler[r].insert(s, mrs);

            if (is_directed(_state._g) || r != s)
            {
                auto& r_in_sampler = is_directed(_state._g) ?
                    _r_in_sampler : _r_out_sampler;
                if (s >= r_in_sampler.size())
                    r_in_sampler.resize(s + 1);
                _r_in_pos[me] = r_in_sampler[s].insert(r, mrs);
            }
        }

        for (auto v : _vertices)
        {
            size_t r = state._b[v];
            if (r >= _v_out_sampler.size())
            {
                if (is_directed(_state._g))
                    _v_in_sampler.resize(r+1);
                _v_out_sampler.resize(r+1);
            }

            size_t kin = 0;
            size_t kout = 0;
            if (state._deg_corr)
                std::tie(kin, kout) = state._degs[v];

            if (is_directed(_state._g))
                _v_in_pos[v] = _v_in_sampler[r].insert(v, kin + 1);
            _v_out_pos[v] = _v_out_sampler[r].insert(v, kout + 1);
        }
    }

    SBMEdgeSampler(const SBMEdgeSampler&) = delete;

    template <class T>
    std::tuple<T, T> get_edge(T u, T v)
    {
        if (!is_directed(_state._g) && u > v)
            return {v, u};
        return {u, v};
    }

    void update_edge(vertex_t u, vertex_t v, size_t m, int64_t delta)
    {
        if (_esample == 0 && _usample == 1)
            return;

        if (m == 0 && delta > 0)
        {
            _edges.push_back(get_edge(u, v));
            _edge_pos[_edges.back()] = _edges.size() - 1;
        }

        _E += delta;

        if (m > 0 && (m + delta) == 0)
        {
            auto iter = _edge_pos.find(get_edge(u, v));
            size_t pos = iter->second;
            _edge_pos.erase(iter);
            if (pos < _edges.size() - 1)
            {
                _edges[pos] = _edges.back();
                _edge_pos[_edges.back()] = pos;
            }
            _edges.pop_back();
        }

        if (_esample == 1 || _usample == 1)
            return;

        size_t r = _state._b[u];
        size_t s = _state._b[v];

        std::tie(r, s) = get_edge(r, s);

        auto me = _state._emat.get_me(r, s);
        assert(me != _state._emat.get_null_edge());

        int64_t ers = _state._mrs[me] + ((delta < 0) ? delta : 0);
        if (ers == 0)
        {
            _rs_sampler.remove(_rs_pos[me]);
            _rs_pos[me] = std::numeric_limits<size_t>::max();

            _r_out_sampler[r].remove(_r_out_pos[me]);
            _r_out_pos[me] = std::numeric_limits<size_t>::max();

            if (is_directed(_state._g) || r != s)
            {
                auto& r_in_sampler = is_directed(_state._g) ?
                    _r_in_sampler : _r_out_sampler;

                r_in_sampler[s].remove(_r_in_pos[me]);
                _r_in_pos[me] = std::numeric_limits<size_t>::max();
            }
        }
        else if (delta == ers)
        {
            _rs_pos[me] = _rs_sampler.insert({r, s}, ers);

            if (!is_directed(_state._g) && r == s)
                ers *= 2;

            if (r >= _r_out_sampler.size())
                _r_out_sampler.resize(r + 1);
            _r_out_pos[me] = _r_out_sampler[r].insert(s, ers);

            if (is_directed(_state._g) || r != s)
            {
                auto& r_in_sampler = is_directed(_state._g) ?
                    _r_in_sampler : _r_out_sampler;

                if (s >= r_in_sampler.size())
                    r_in_sampler.resize(s + 1);
                _r_in_pos[me] = r_in_sampler[s].insert(r, ers);
            }
        }
        else
        {
            _rs_sampler.update(_rs_pos[me], delta, true);

            int64_t dm = delta;
            if (!is_directed(_state._g) && r == s)
                dm *= 2;

            _r_out_sampler[r].update(_r_out_pos[me], dm, true);

            if (is_directed(_state._g) || r != s)
            {
                auto& r_in_sampler = is_directed(_state._g) ?
                    _r_in_sampler : _r_out_sampler;

                r_in_sampler[s].update(_r_in_pos[me], dm, true);
            }
        }

        if (_state._deg_corr)
        {
            size_t ku = get<1>(_state._degs[u]);
            size_t kv = (is_directed(_state._g)) ?
                get<0>(_state._degs[v]) :
                get<1>(_state._degs[v]);

            if (delta < 0)
            {
                if (u != v || is_directed(_state._g))
                {
                    ku += delta;
                    kv += delta;
                }
                else
                {
                    ku += 2 * delta;
                }
            }

            r = _state._b[u];
            _v_out_sampler[r].remove(_v_out_pos[u]);
            _v_out_pos[u] = _v_out_sampler[r].insert(u, ku + 1);

            if (u != v || is_directed(_state._g))
            {
                s = _state._b[v];
                if (is_directed(_state._g))
                {
                    _v_in_sampler[s].remove(_v_in_pos[v]);
                    _v_in_pos[v] = _v_in_sampler[s].insert(v, kv + 1);
                }
                else
                {
                    _v_out_sampler[s].remove(_v_out_pos[v]);
                    _v_out_pos[v] = _v_out_sampler[s].insert(v, kv + 1);
                }
            }
        }
    }

    template <class RNG>
    std::tuple<vertex_t, vertex_t> sample(RNG& rng, bool edges_only = false)
    {
        // if (_edges.empty())
        //     return _null_edge;
        // return uniform_sample(_edges, rng);
        if (_esample == 0 && _usample == 1)
        {
            return get_edge(uniform_sample(_vertices, rng),
                            uniform_sample(_vertices, rng));
        }

        if (_esample == 1 || edges_only)
        {
            if (_edges.empty())
                return _null_edge;

            std::bernoulli_distribution coin(_E / double(_E + _N));
            if (edges_only || coin(rng))
            {
                return uniform_sample(_edges, rng);
            }
            else
            {
                auto v = uniform_sample(_vertices, rng);
                return {v, v};
            }
        }

        std::bernoulli_distribution esample(_esample);

        if (!_edges.empty() && esample(rng))
            return uniform_sample(_edges, rng);

        std::bernoulli_distribution usample(_usample);

        if (_edges.empty() || usample(rng))
            return get_edge(uniform_sample(_vertices, rng),
                            uniform_sample(_vertices, rng));

        auto [r, s] = _rs_sampler.sample(rng);
        auto& r_sampler = _v_out_sampler[r];
        auto& s_sampler = _v_in_sampler[s];
        return get_edge(r_sampler.sample(rng),
                        s_sampler.sample(rng));
    }

    double log_prob(vertex_t u, vertex_t v, size_t m, int64_t delta,
                    bool edges_only = false)
    {
        // if (_edges.size() + delta == 0)
        //     return -numeric_limits<double>::infinity();
        // return -log(_edges.size() + delta);
        if (_esample == 0 && _usample == 1)
            return 0;

        if (_esample == 1 || edges_only)
            return -log(_edges.size() + delta);

        auto& g = _state._g;
        auto E = is_directed(g) ? (_E + delta) :
            2 * (_E + delta);

        double lu = -safelog_fast(_N) - safelog_fast(_vertices.size());
        if (!is_directed(g) && u != v)
            lu += log(2);

        if (_usample == 1 || E == 0)
            return lu;

        size_t r = _state._b[u];
        size_t s = _state._b[v];

        size_t ku = 0, kv = 0;
        if (_state._deg_corr)
        {
            ku = get<1>(_state._degs[u]);
            kv = (is_directed(_state._g)) ?
                get<0>(_state._degs[v]) :
                get<1>(_state._degs[v]);
        }

        auto&& me = _state._emat.get_me(r, s);
        size_t ers = (me == _state._emat.get_null_edge()) ? 0 : size_t(_state._mrs[me]);
        ers += delta;

        if (!is_directed(g) && r == s)
            ers *= 2;

        size_t nr = _state._wr[r];
        size_t ns = _state._wr[s];
        size_t er = _state._mrp[r];
        size_t es = is_directed(g) ? _state._mrm[s] : _state._mrp[s];

        if (_state._deg_corr)
        {
            if (r != s || is_directed(g))
            {
                er += delta;
                es += delta;
            }
            else
            {
                er += 2 * delta;
                es += 2 * delta;
            }

            if (u != v || is_directed(g))
            {
                ku += delta;
                kv += delta;
            }
            else
            {
                ku += 2 * delta;
                kv += 2 * delta;
            }
        }
        else
        {
            er = es = 0;
        }

        double lp = 0;

        if (E > 0)
        {
            if (ers > 0)
            {
                lp = ((safelog_fast(ers) - safelog_fast(E)) +
                      (safelog_fast(ku + 1) - safelog_fast(er + nr)) +
                      (safelog_fast(kv + 1) - safelog_fast(es + ns)));

                if (!is_directed(g) && u != v)
                    lp += log(2);

                lp = log_sum_exp(lp + log1p(-_usample),
                                 lu + log(_usample));
            }
            else
            {
                lp = lu + log(_usample);
            }
        }
        else
        {
            lp = lu;
        }

        if (_esample > 0)
        {
            if (m + delta > 0)
            {
                double rp;
                if (m == 0)
                    rp = -safelog_fast(_edges.size() + 1);
                else
                    rp = -safelog_fast(_edges.size());
                return log_sum_exp(rp + log(_esample),
                                   lp + log1p(-_esample));
            }
            else
            {
                return lp + log1p(-_esample);
            }
        }

        return lp;
    }

    template <class RNG>
    vertex_t sample_out_neighbor(vertex_t v, RNG& rng)
    {
        auto r = _state._b[v];
        std::bernoulli_distribution usample(_usample);
        if (_state._mrp[r] == 0 || usample(rng))
            return uniform_sample(_vertices, rng);
        auto s = _r_out_sampler[r].sample(rng);
        return _v_in_sampler[s].sample(rng);
    }

    template <class RNG>
    vertex_t sample_in_neighbor(vertex_t v, RNG& rng)
    {
        auto r = _state._b[v];
        std::bernoulli_distribution usample(_usample);
        if (_state._mrm[r] == 0 || usample(rng))
            return uniform_sample(_vertices, rng);
        auto s = _r_in_sampler[r].sample(rng);
        return _v_out_sampler[s].sample(rng);
    }

    double log_prob_out(vertex_t v, vertex_t u)
    {
        if (_usample == 1)
            return -safelog_fast(_N);

        auto& g = _state._g;
        size_t r = _state._b[v];
        size_t s = _state._b[u];

        size_t ku = 0;
        if (_state._deg_corr)
            ku = (is_directed(_state._g)) ?
                get<0>(_state._degs[u]) :
                get<1>(_state._degs[u]);

        auto&& me = _state._emat.get_me(r, s);
        size_t ers = (me == _state._emat.get_null_edge()) ? 0 : size_t(_state._mrs[me]);

        if (!is_directed(g) && r == s)
            ers *= 2;

        size_t ns = _state._wr[s];
        size_t er = _state._mrp[r];
        size_t es = (_state._deg_corr) ?
            (is_directed(g) ? _state._mrm[s] : _state._mrp[s]) : 0;

        double lN = safelog_fast(_N);

        if (_state._mrp[r] == 0)
            return -lN;

        if (ers > 0)
        {
            double lp = (safelog_fast(ers) - safelog_fast(er) +
                         safelog_fast(ku + 1) - safelog_fast(es + ns));

            return log_sum_exp(lp + log1p(-_usample), -lN + log(_usample));
        }
        else
        {
            return -lN + log(_usample);
        }
    }

    double log_prob_in(vertex_t v, vertex_t u)
    {
        if (_usample == 1)
            return -safelog_fast(_N);

        size_t r = _state._b[v];
        size_t s = _state._b[u];

        size_t ku = 0;
        if (_state._deg_corr)
            ku = get<1>(_state._degs[u]);

        auto&& me = _state._emat.get_me(s, r);
        size_t esr = (me == _state._emat.get_null_edge()) ? 0 : size_t(_state._mrs[me]);

        size_t ns = _state._wr[s];
        size_t er = _state._mrm[r];
        size_t es = (_state._deg_corr) ? _state._mrp[s] : 0;

        double lN = safelog_fast(_N);

        if (_state._mrp[r] == 0)
            return -lN;

        if (esr > 0)
        {
            double lp = (safelog_fast(esr) - safelog_fast(er) +
                         safelog_fast(ku + 1) - safelog_fast(es + ns));

            return log_sum_exp(lp + log1p(-_usample), -lN + log(_usample));
        }
        else
        {
            return -lN + log(_usample);
        }
    }

    void check_counts()
    {
        for (auto me : edges_range(_state._bg))
        {
            size_t mrs = _state._mrs[me];
            assert(_rs_sampler.get_prob(_rs_pos[me]) == mrs);

            size_t r = source(me, _state._bg);
            size_t s = target(me, _state._bg);
            std::tie(r, s) = get_edge(r, s);

            if (!is_directed(_state._g) && r == s)
                mrs *= 2;

            assert(_r_out_sampler[r].get_prob(_r_out_pos[me]) == mrs);
            if (is_directed(_state._g) || r != s)
            {
                auto& r_in_sampler = is_directed(_state._g) ?
                    _r_in_sampler : _r_out_sampler;
                assert(r_in_sampler[s].get_prob(_r_in_pos[me]) == mrs);
            }
        }

        for (size_t i = 0; i < _rs_sampler.size(); ++i)
        {
            auto& [r,s] = _rs_sampler[i];
            auto ers = _rs_sampler.get_prob(i);
            auto&& me = _state._emat.get_me(r, s);
            size_t mrs = 0;
            if (me != _state._emat.get_null_edge())
                mrs = _state._mrs[me];
            assert(mrs == ers);
        }

        for (size_t r = 0; r < _r_out_sampler.size(); ++r)
        {
            auto& out_sampler = _r_out_sampler[r];
            for (size_t i = 0; i < out_sampler.size(); ++i)
            {
                if (!out_sampler.is_valid(i))
                    continue;
                auto s = out_sampler[i];
                auto ers = out_sampler.get_prob(i);
                auto&& me = _state._emat.get_me(r, s);
                size_t mrs = 0;
                if (me != _state._emat.get_null_edge())
                    mrs = _state._mrs[me];
                if (!is_directed(_state._g) && r == s)
                    mrs *= 2;
                assert(mrs == ers);
            }
        }

        auto& r_in_sampler = is_directed(_state._g) ?
            _r_in_sampler : _r_out_sampler;

        for (size_t s = 0; s < r_in_sampler.size(); ++s)
        {
            auto& in_sampler = r_in_sampler[s];
            for (size_t i = 0; i < in_sampler.size(); ++i)
            {
                if (!in_sampler.is_valid(i))
                    continue;
                auto r = in_sampler[i];
                auto ers = in_sampler.get_prob(i);
                auto&& me = _state._emat.get_me(r, s);
                size_t mrs = 0;
                if (me != _state._emat.get_null_edge())
                    mrs = _state._mrs[me];
                if (!is_directed(_state._g) && r == s)
                    mrs *= 2;
                assert(mrs == ers);
            }
        }
    }

    std::tuple<vertex_t, vertex_t> get_null_edge()
    {
        return _null_edge;
    }

    size_t num_edges()
    {
        return _edges.size();
    }

//private:
    State& _state;

    std::vector<vertex_t> _vertices;
    size_t _N;

    typedef DynamicSampler<std::tuple<size_t, size_t>> rs_sampler_t;
    rs_sampler_t _rs_sampler;
    eprop_map_t<size_t> _rs_pos;

    typedef DynamicSampler<size_t> r_sampler_t;
    std::vector<r_sampler_t> _r_out_sampler;
    eprop_map_t<size_t> _r_out_pos;

    std::vector<r_sampler_t> _r_in_sampler;
    eprop_map_t<size_t> _r_in_pos;

    typedef DynamicSampler<vertex_t> vsampler_t;
    vector<vsampler_t> __v_in_sampler, _v_out_sampler;
    vector<vsampler_t>& _v_in_sampler;

    std::conditional_t<std::is_integral_v<vertex_t>,
                       vprop_map_t<size_t>,
                       gt_hash_map<vertex_t, size_t>>
        _v_in_pos, _v_out_pos;

    std::vector<std::tuple<vertex_t, vertex_t>> _edges;
    gt_hash_map<std::tuple<vertex_t, vertex_t>, size_t> _edge_pos;
    size_t _E = 0;

    double _esample;
    double _usample;

    std::tuple<vertex_t, vertex_t> _null_edge = \
            { boost::graph_traits<typename State::g_t>::null_vertex(),
              boost::graph_traits<typename State::g_t>::null_vertex() };

};


} // graph_tool namespace

#endif // GRAPH_SBM_SAMPLE_EDGE_HH
