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

#include "graph_filtering.hh"
#include "graph.hh"
#include "graph_properties.hh"
#include "graph_exceptions.hh"

#include "graph_sfdp.hh"
#include "random.hh"
#include "hash_map_wrap.hh"
#include "numpy_bind.hh"

#include <boost/multi_array/multi_array_ref.hpp>

using namespace std;
using namespace boost;
using namespace graph_tool;

void sfdp_layout(GraphInterface& g, std::any pos, std::any vweight,
                 std::any eweight, std::any pin, python::object spring_parms,
                 double theta, double init_step, double step_schedule,
                 size_t max_level, double epsilon, size_t max_iter,
                 bool adaptive, bool verbose, rng_t& rng)
{
    typedef UnityPropertyMap<int,GraphInterface::vertex_t> vweight_map_t;
    typedef UnityPropertyMap<int,GraphInterface::edge_t> eweight_map_t;
    auto vertex_props_t = hana::append(vertex_scalar_properties,
                                       hana::type<vweight_map_t>());
    auto edge_props_t = hana::append(edge_scalar_properties,
                                     hana::type<eweight_map_t>());

    double C = python::extract<double>(spring_parms[0]);
    double K = python::extract<double>(spring_parms[1]);
    double p = python::extract<double>(spring_parms[2]);
    std::vector<double> gamma;
    for (int i = 0; i < python::len(spring_parms[3]); ++i)
        gamma.push_back(python::extract<double>(spring_parms[3][i]));

    typedef multi_array_ref<int32_t,1> group_map_t;
    std::vector<group_map_t> groups;
    for (int i = 0; i < python::len(spring_parms[4]); ++i)
        groups.push_back(get_array<int32_t,1>(spring_parms[4][i]));

    double r = python::extract<double>(spring_parms[5]);
    size_t kc = python::extract<double>(spring_parms[6]);
    typedef vprop_map_t<int32_t> c_map_t;
    c_map_t c_map = std::any_cast<c_map_t>(python::extract<any&>(spring_parms[7])());

    double R = python::extract<double>(spring_parms[8]);
    typedef vprop_map_t<double> o_map_t;
    o_map_t o_map = std::any_cast<o_map_t>(python::extract<any&>(spring_parms[9])());

    if (!vweight.has_value())
        vweight = vweight_map_t();
    if (!eweight.has_value())
        eweight = eweight_map_t();

    typedef vprop_map_t<uint8_t> pin_map_t;
    pin_map_t pin_map = std::any_cast<pin_map_t>(pin);

    run_action<decltype(graph_tool::never_directed)>()
        (g,
         [&](auto&& graph, auto&& pos, auto&& vweight, auto&& eweight)
         {
             return get_sfdp_layout(graph, pos, vweight, eweight,
                                    pin_map.get_unchecked(num_vertices(g.get_graph())),
                                    groups, C, K, p, theta, gamma, r, kc,
                                    c_map.get_unchecked(num_vertices(g.get_graph())),
                                    R, o_map.get_unchecked(num_vertices(g.get_graph())),
                                    init_step,
                                    step_schedule, max_level, epsilon, max_iter,
                                    adaptive, verbose, rng);
         },
         vertex_floating_vector_properties, vertex_props_t,
         edge_props_t)(pos, vweight, eweight);
}

struct do_propagate_pos
{
    template <class Graph, class CoarseGraph, class VertexMap, class PosMap,
              class RNG>
    void operator()(Graph& g, CoarseGraph& cg, VertexMap vmap,
                    std::any acvmap, PosMap pos, std::any acpos,
                    double delta, RNG& rng) const
    {
        typename PosMap::checked_t cpos =
            std::any_cast<typename PosMap::checked_t>(acpos);
        typename VertexMap::checked_t cvmap =
            std::any_cast<typename VertexMap::checked_t>(acvmap);
        typedef typename property_traits<VertexMap>::value_type c_t;
        typedef typename property_traits<PosMap>::value_type pos_t;
        typedef typename pos_t::value_type val_t;

        uniform_real_distribution<val_t> noise(-delta, delta);
        gt_hash_map<c_t, pos_t> cmap;

        for (auto v : vertices_range(cg))
            cmap[cvmap[v]] = cpos[v];

        for (auto v : vertices_range(g))
        {
            pos[v] = cmap[vmap[v]];

            if (delta > 0)
            {
                for (size_t j = 0; j < pos[v].size(); ++j)
                    pos[v][j] += noise(rng);
            }
        }
    }
};

void propagate_pos(GraphInterface& gi, GraphInterface& cgi, std::any vmap,
                   std::any cvmap, std::any pos, std::any cpos,
                   double delta, rng_t& rng)
{
    auto vmaps_t = hana::tuple_t<vprop_map_t<int32_t>>;

    gt_dispatch<>()
        (
            [&](auto&& graph, auto&& a2, auto&& a3, auto&& a4)
            {
                return do_propagate_pos()
                    (std::forward<decltype(graph)>(graph),
                     std::forward<decltype(a2)>(a2),
                     std::forward<decltype(a3)>(a3), cvmap,
                     std::forward<decltype(a4)>(a4), cpos, delta, rng);
            },
            all_graph_views, all_graph_views, vmaps_t,
            vertex_floating_vector_properties)(
            gi.get_graph_view(), cgi.get_graph_view(), vmap, pos);
}

struct do_propagate_pos_mivs
{
    template <class Graph, class MIVSMap, class PosMap,
              class RNG>
    void operator()(Graph& g, MIVSMap mivs, PosMap pos, double delta, RNG& rng) const
    {
        typedef typename property_traits<PosMap>::value_type pos_t;
        typedef typename pos_t::value_type val_t;

        uniform_real_distribution<val_t> noise(-delta, delta);

        for (auto v : vertices_range(g))
        {
            if (mivs[v])
                continue;
            size_t count = 0;
            for (auto a : adjacent_vertices_range(v, g))
            {
                if (!mivs[a])
                    continue;
                pos[v].resize(pos[a].size(), 0);
                for (size_t j = 0; j < pos[a].size(); ++j)
                    pos[v][j] += pos[a][j];
                ++count;
            }

            if (count == 0)
                throw ValueException("invalid MIVS! Vertex has no neighbors "
                                     "belonging to the set!");

            if (count == 1)
            {
                if (delta > 0)
                {
                    for (size_t j = 0; j < pos[v].size(); ++j)
                        pos[v][j] += noise(rng);
                }
            }
            else
            {
                for (size_t j = 0; j < pos[v].size(); ++j)
                    pos[v][j] /= count;
            }
        }
    }
};


void propagate_pos_mivs(GraphInterface& gi, std::any mivs, std::any pos,
                        double delta, rng_t& rng)
{
    run_action<>()
        (gi,
         [&](auto&& graph, auto&& a2, auto&& a3)
         {
             return do_propagate_pos_mivs()
                 (std::forward<decltype(graph)>(graph),
                  std::forward<decltype(a2)>(a2),
                  std::forward<decltype(a3)>(a3), delta, rng);
         },
         vertex_scalar_properties,
         vertex_floating_vector_properties)(mivs, pos);
}


struct do_avg_dist
{
    template <class Graph, class PosMap>
    void operator()(Graph& g, PosMap pos, double& ad) const
    {
        size_t count = 0;
        double d = 0;
        #pragma omp parallel if (num_vertices(g) > get_openmp_min_thresh()) \
            reduction(+: d, count)
        parallel_vertex_loop_no_spawn
                (g,
                 [&](auto v)
                 {
                     for (auto a : adjacent_vertices_range(v, g))
                     {
                         d += dist(pos[v], pos[a]);
                         count++;
                     }
                 });
        if (count > 0)
            d /= count;
        ad = d;
    }
};


double avg_dist(GraphInterface& gi, std::any pos)
{
    double d;
    run_action<>()
        (gi,
         [&](auto&& graph, auto&& a2)
         {
             return do_avg_dist()
                 (std::forward<decltype(graph)>(graph),
                  std::forward<decltype(a2)>(a2), d);
         },
         vertex_scalar_vector_properties)(pos);
    return d;
}

void sanitize_pos(GraphInterface& gi, std::any apos)
{
    run_action<>()
        (gi,
         [&](auto&& g, auto&& pos)
         {
             parallel_vertex_loop
                 (g,
                  [&](auto v)
                  {
                      pos[v].resize(2);
                  });
         },
         vertex_scalar_vector_properties)(apos);
}

#include <boost/python.hpp>

#define __MOD__ layout
#include "module_registry.hh"
REGISTER_MOD
([]
 {
     python::def("sfdp_layout", &sfdp_layout);
     python::def("propagate_pos", &propagate_pos);
     python::def("propagate_pos_mivs", &propagate_pos_mivs);
     python::def("avg_dist", &avg_dist);
     python::def("sanitize_pos", &sanitize_pos);
 });
