export Cora


"""
    Cora

The Cora citation network dataset from Ref. [1].
Nodes represent documents and edges represent citation links.
Each node has a predefined feature with 1433 dimensions. 
The dataset is designed for the node classification task. 
The task is to predict the category of certain paper.
The dataset is retrieved from Ref. [2].

## Statistics 

- Nodes: 2708
- Edges: 10556
- Number of Classes: 7
- Label split:
    - Train:  140
    - Val:    500
    - Test:  1000

The split is the one used in the original paper [1] and 
doesn't consider all nodes.

## Interface

- [`Cora.dataset`](@ref)

## References

[1]: [Deep Gaussian Embedding of Graphs: Unsupervised Inductive Learning via Ranking](https://arxiv.org/abs/1707.03815)
[2]: [Planetoid](https://github.com/kimiyoung/planetoid
"""
module Cora

using DataDeps
using ..MLDatasets: datafile, read_planetoid_data
using DelimitedFiles: readdlm

using PyCall

const DEPNAME = "Cora"
# LINK = "https://github.com/shchur/gnn-benchmark/raw/master/data/npz"
# LINK = "https://github.com/abojchevski/graph2gauss/raw/master/data/"
const LINK = "https://github.com/kimiyoung/planetoid/raw/master/data"
const DOCS = "https://github.com/kimiyoung/planetoid"
const DATA = "ind.cora." .* ["x", "y", "tx", "allx", "ty", "ally", "graph", "test.index"]

function __init__()
    register(DataDep(
        DEPNAME,
        """
        Dataset: The $DEPNAME dataset.
        Website: $DOCS
        """,
        map(x -> "$LINK/$x", DATA),
        "81de017067dc045ebdb8ffd5c0e69a209973ffdb1fe2d5b434e94d3614f3f5c7",  # if checksum omitted, will be generated by DataDeps
        # post_fetch_method = unpack
    ))
end

"""
    dataset(; dir=nothing, reverse_edges=true)

Retrieve the Cora dataset. The output is a named tuple with fields
```juliarepl
julia> keys(Cora.dataset())
(:node_features, :node_labels, :adjacency_list, :train_indices, :val_indices, :test_indices, :num_classes, :num_nodes, :num_edges, :directed)
```
In particular, `adjacency_list` is a vector of vector, 
where `adjacency_list[i]` will contain the neighbors of node `i`
through outgoing edges.

If `reverse_edges=true`, the graph will contain
the reverse of each edge and the graph will be undirected.

See also [`Cora`](@ref).

## Usage Examples

```juliarepl
using MLDatasets: Cora

data = Cora.dataset()
train_labels = data.node_labels[data.train_indices]
```
"""
dataset(; dir=nothing, reverse_edges=true) =
    read_planetoid_data(DEPNAME, dir=dir, reverse_edges=reverse_edges)


end #module 

