#ifndef DRESS_H
#define DRESS_H

#include <stdint.h>

#ifdef __cplusplus
extern "C" {
#endif

// dress graph variant: determines how adjacency lists are constructed
// from the input edge list.
//
//   UNDIRECTED: N[u] = N(u)           (symmetric neighborhood)
//   DIRECTED:   N[u] = in(u) + out(u) (union of in- and out-neighbors)
//   FORWARD:    N[u] = out(u)         (only outgoing neighbors)
//   BACKWARD:   N[u] = in(u)          (only incoming neighbors)
typedef enum __dress_variant_t {
    DRESS_VARIANT_UNDIRECTED = 0,
    DRESS_VARIANT_DIRECTED   = 1,
    DRESS_VARIANT_FORWARD    = 2,
    DRESS_VARIANT_BACKWARD   = 3
} dress_variant_t;



// dress graph structure.
//
// Stores the augmented variant adjacency in CSR (Compressed Sparse Row)
// format, with per-edge weight and similarity arrays. An optional set of
// precomputed neighborhood intercepts accelerates iterative fitting from
// O(deg_u + deg_v) to O(|N[u] ∩ N[v]|) per edge.
//
// Memory layout (S = total half-edges, T = total intercept entries):
//   Permanent:  12N + 52E + 8T  bytes  (with intercepts)
//               12N + 48E       bytes  (without intercepts)
typedef struct __dress_graph_t {
    dress_variant_t variant;       // graph variant (determines adjacency)
    int      N;                    // number of vertices
    int      E;                    // number of input edges

    int     *U;                    // [E]   input edge sources (owned)
    int     *V;                    // [E]   input edge targets (owned)

    // CSR variant adjacency — sorted per-node neighbor lists stored flat.
    // Neighbors of node u are adj_target[adj_offset[u] .. adj_offset[u+1]).
    int     *adj_offset;           // [N+1] CSR row offsets
    int     *adj_target;           // [S]   neighbor vertex ids
    int     *adj_edge_idx;         // [S]   maps half-edge to input edge index

    // Per-edge arrays — indexed by edge id 0..E-1.
    double  *W;                    // [E]   raw input weight (NULL when unweighted)
    double  *edge_weight;          // [E]   variant-specific edge weight
    double  *edge_dress;           // [E]   current dress values
    double  *edge_dress_next;      // [E]   next-iteration dress values (double-buffer)

    // Per-node arrays — indexed by vertex id 0..N-1.
    double  *node_dress;           // [N]   sqrt of weighted dress sum for node

    // Precomputed intercepts (allocated only when precompute_intercepts == 1).
    // For edge e = (u,v), common neighbors are stored at
    // intercept_edge_ux/vx[intercept_offset[e] .. intercept_offset[e+1]).
    int      precompute_intercepts;
    int     *intercept_offset;     // [E+1] CSR-style offsets into intercept arrays
    int     *intercept_edge_ux;    // [T]   edge index for (u,x) per common neighbor x
    int     *intercept_edge_vx;    // [T]   edge index for (v,x) per common neighbor x
} dress_graph_t, *p_dress_graph_t;

// Construct a dress graph from an edge list.
// Takes ownership of U, V, W (all freed by free_dress_graph).
// W may be NULL (unweighted).
p_dress_graph_t init_dress_graph(int N, int E, int *U, int *V,
                                 double *W, dress_variant_t variant,
                                 int precompute_intercepts);

// Run iterative dress fitting for at most max_iterations, stopping early
// when the maximum per-edge change falls below epsilon.
// On return, *iterations and *delta (if non-NULL) hold the iteration
// count and final max delta.
void   dress_fit(p_dress_graph_t g, int max_iterations, double epsilon,
                int *iterations, double *delta);

// Free all memory associated with the dress graph (including U, V).
void   free_dress_graph(p_dress_graph_t g);

// Query the DRESS value for any vertex pair (u, v).
//
// If edge (u,v) exists in the graph, returns its converged edge_dress value.
// If the edge does not exist (virtual edge), estimates what the DRESS value
// *would be* using a local fixed-point iteration against the frozen
// steady-state node/edge values.  Useful for link prediction.
//
// The graph must have been fitted (dress_fit) before calling this.
// edge_weight is the hypothetical weight of the virtual edge (ignored
// when the edge already exists).  Pass 1.0 for unweighted graphs.
double dress_get(const p_dress_graph_t g, int u, int v,
                 int max_iterations, double epsilon,
                 double edge_weight);

// Backward compatibility aliases (C only — disabled for C++ to avoid
// collisions with member function names).
#ifndef __cplusplus
#define fit(g, max_iterations, epsilon, iterations, delta) \
        dress_fit((g), (max_iterations), (epsilon), (iterations), (delta))
#endif

// ---- Δ^k-DRESS ----------------------------------------------------------

// Compute the Δ^k-DRESS histogram of a graph.
//
// Enumerates all C(N, k) subsets of k vertices, removes each subset
// from the graph, runs DRESS on the resulting subgraph, and accumulates
// every converged edge value into a single histogram.
//
// The deletion subsets are generated by an iterative (stack-based) DFS
// over k-combinations, avoiding recursion.
//
// Parameters:
//   g          - input dress graph (used for topology only: N, E, U, V,
//                variant).  Not modified.
//   k          - deletion depth: number of vertices removed per subset.
//                k = 0 runs DRESS on the original graph (Δ^0).
//   iterations - maximum DRESS iterations per subgraph.
//   epsilon    - convergence tolerance for DRESS and histogram bin width.
//                The histogram has floor(dmax / epsilon) + 1 bins
//                (dmax = 2 for unweighted graphs; larger for weighted).
//   hist_size  - [out] if non-NULL, set to floor(dmax / epsilon) + 1 on return.
//   keep_multisets - if non-zero, allocate and fill a flat matrix of
//                per-subgraph edge DRESS values.
//   multisets  - [out] on return, *multisets points to a heap-allocated
//                flat array of size C(N,k) * E.  Row r, edge e is at
//                (*multisets)[r * E + e].  Removed edges are NAN.
//                The caller must free(*multisets).  Ignored (and set to
//                NULL) when keep_multisets is 0.  May be NULL itself.
//   num_subgraphs - [out] if non-NULL, set to C(N,k) on return.
//
// Returns:
//   A heap-allocated int64_t array of length floor(dmax / epsilon) + 1.
//   hist[i] counts the number of converged edge values falling in bin
//   i = floor(d / epsilon).  The top bin holds the maximum possible
//   value.  The caller must free() the returned pointer.
//
// Complexity:
//   O( C(N,k) * iterations * E * d_max )
//   where d_max is the maximum degree.  Each subgraph is independent;
//   the outer loop is embarrassingly parallel (not parallelised here).
int64_t *delta_dress_fit(p_dress_graph_t g, int k, int iterations,
                        double epsilon, int *hist_size,
                        int keep_multisets, double **multisets,
                        int64_t *num_subgraphs);

// Strided variant of delta_dress_fit for distributed computation.
//
// Processes only the subgraphs whose sequential index satisfies
// index % stride == offset.  With offset=0, stride=1 this is
// identical to delta_dress_fit (all subgraphs).
//
// Intended for MPI distribution: each rank calls with
// offset=rank, stride=nprocs, then the per-rank histograms are
// summed via MPI_Allreduce.
int64_t *delta_dress_fit_strided(p_dress_graph_t g, int k, int iterations,
                                 double epsilon, int *hist_size,
                                 int keep_multisets, double **multisets,
                                 int64_t *num_subgraphs,
                                 int offset, int stride);

// Backward compatibility alias (C only).
#ifndef __cplusplus
#define delta_fit(g, k, iterations, epsilon, hist_size, keep_multisets, multisets, num_subgraphs) \
        delta_dress_fit((g), (k), (iterations), (epsilon), (hist_size), (keep_multisets), (multisets), (num_subgraphs))
#endif

#ifdef __cplusplus
}
#endif

#endif // DRESS_H