/****************************************************************************
 * pySurprise - Python bindings for the Surprise metric                     *
 *                                                                          *
 * Copyright (C) 2012 Rodrigo Aldecoa and Ignacio Marín                     *
 * Python bindings by pySurprise contributors                               *
 *                                                                          *
 *  This program is free software: you can redistribute it and/or modify    *
 *  it under the terms of the GNU General Public License as published by    *
 *  the Free Software Foundation, either version 3 of the License, or       *
 *  (at your option) any later version.                                     *
 ****************************************************************************/

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <vector>
#include <map>
#include <set>
#include <stdexcept>
#include <cmath>

#include "surprise.h"

namespace py = pybind11;


// High-level: compute Surprise from edges and partition assignment
// edges: list of (int, int) pairs
// partition: list of int, where partition[i] = community of node i
double surprise_from_partition(
    const std::vector<std::pair<int, int>>& edges,
    const std::vector<int>& partition)
{
    int num_nodes = static_cast<int>(partition.size());
    if (num_nodes < 2) {
        throw std::invalid_argument("Partition must have at least 2 nodes");
    }

    // F = number of possible pairs = n*(n-1)/2
    double F = static_cast<double>(num_nodes) * (num_nodes - 1) / 2.0;

    // n = number of edges
    double n = static_cast<double>(edges.size());

    // Compute community sizes for M
    std::map<int, int> community_sizes;
    for (int i = 0; i < num_nodes; ++i) {
        community_sizes[partition[i]]++;
    }

    // M = sum of s_i*(s_i-1)/2 for each community
    double M = 0.0;
    for (const auto& kv : community_sizes) {
        double s = static_cast<double>(kv.second);
        M += s * (s - 1.0) / 2.0;
    }

    // p = intra-community edges
    double p = 0.0;
    for (const auto& edge : edges) {
        int u = edge.first;
        int v = edge.second;
        if (u < 0 || u >= num_nodes || v < 0 || v >= num_nodes) {
            throw std::invalid_argument(
                "Edge (" + std::to_string(u) + ", " + std::to_string(v) +
                ") references node outside partition range [0, " +
                std::to_string(num_nodes - 1) + "]");
        }
        if (partition[u] == partition[v]) {
            p++;
        }
    }

    return computeSurprise(F, M, n, p);
}


// High-level: compute Surprise from edges as dict-labelled partition
// edges: list of (str, str) pairs
// partition: dict mapping node label -> community id
double surprise_from_partition_dict(
    const std::vector<std::pair<std::string, std::string>>& edges,
    const std::map<std::string, int>& partition)
{
    int num_nodes = static_cast<int>(partition.size());
    if (num_nodes < 2) {
        throw std::invalid_argument("Partition must have at least 2 nodes");
    }

    double F = static_cast<double>(num_nodes) * (num_nodes - 1) / 2.0;
    double n = static_cast<double>(edges.size());

    // Community sizes
    std::map<int, int> community_sizes;
    for (const auto& kv : partition) {
        community_sizes[kv.second]++;
    }

    double M = 0.0;
    for (const auto& kv : community_sizes) {
        double s = static_cast<double>(kv.second);
        M += s * (s - 1.0) / 2.0;
    }

    // Intra-community edges
    double p = 0.0;
    for (const auto& edge : edges) {
        auto it_u = partition.find(edge.first);
        auto it_v = partition.find(edge.second);
        if (it_u == partition.end() || it_v == partition.end()) {
            throw std::invalid_argument(
                "Edge node not found in partition");
        }
        if (it_u->second == it_v->second) {
            p++;
        }
    }

    return computeSurprise(F, M, n, p);
}


PYBIND11_MODULE(_core, m) {
    m.doc() = "pySurprise: Python bindings for the Surprise network community metric";

    // Low-level function: direct access to the C++ computeSurprise
    m.def("compute_surprise", &computeSurprise,
          py::arg("F"), py::arg("M"), py::arg("n"), py::arg("p"),
          R"pbdoc(
          Compute the Surprise value from the four parameters.

          Parameters
          ----------
          F : float
              Total number of possible node pairs: N*(N-1)/2
          M : float
              Sum of intra-community possible pairs: sum(s_i*(s_i-1)/2)
          n : float
              Total number of edges in the network.
          p : float
              Number of intra-community edges.

          Returns
          -------
          float
              The Surprise value (in -log10 scale).
          )pbdoc");

    m.def("log_hyper_probability", &logHyperProbability,
          py::arg("F"), py::arg("M"), py::arg("n"), py::arg("j"),
          "Compute a single log10 hypergeometric probability term.");

    // High-level: from integer-indexed partition
    m.def("surprise_from_partition", &surprise_from_partition,
          py::arg("edges"), py::arg("partition"),
          R"pbdoc(
          Compute Surprise from a list of edges and a partition vector.

          Parameters
          ----------
          edges : list of (int, int)
              Edge list. Each tuple (u, v) is an undirected edge.
              Node indices must be in [0, len(partition)-1].
          partition : list of int
              Community assignment for each node.
              partition[i] = community ID of node i.

          Returns
          -------
          float
              The Surprise value.
          )pbdoc");

    // High-level: from string-labelled partition
    m.def("surprise_from_partition_dict", &surprise_from_partition_dict,
          py::arg("edges"), py::arg("partition"),
          R"pbdoc(
          Compute Surprise from labelled edges and a partition dict.

          Parameters
          ----------
          edges : list of (str, str)
              Edge list with string node labels.
          partition : dict of {str: int}
              Maps each node label to its community ID.

          Returns
          -------
          float
              The Surprise value.
          )pbdoc");
}
