function __init__tudataset()
    DEPNAME = "TUDataset"
    LINK = "https://www.chrsmrrs.com/graphkerneldatasets"
    DOCS = "https://chrsmrrs.github.io/datasets/docs/home/"
    DATA = "PROTEINS.zip"

    register(DataDep(
        DEPNAME,
        """
        Dataset: The $DEPNAME dataset.
        Website: $LINK)
        """,
        "$LINK/$DATA",
        "2da8de15284b88edabca2888ce5444d62f364ed41159260977088c4e53d4d848",  # if checksum omitted, will be generated by DataDeps
        post_fetch_method = unpack
    ))
end

"""
    TUDataset(name; dir=nothing)

A variety of graph benchmark datasets, *.e.g.* "QM9", "IMDB-BINARY",
"REDDIT-BINARY" or "PROTEINS", collected from the [TU Dortmund University](https://chrsmrrs.github.io/datasets/).
Retrieve from TUDataset collection the dataset `name`, where `name`
is any of the datasets available [here](https://chrsmrrs.github.io/datasets/docs/datasets/). 

A `TUDataset` object can be indexed to retrieve a specific graph or a subset of graphs.

See [here](https://chrsmrrs.github.io/datasets/docs/format/) for an in-depth 
description of the format. 

# Usage Example

```julia
using MLDatasets: TUDataset
using LightGraphs: SimpleGraph, add_edge!

data = TUDataset("PROTEINS")

# Access first graph
d1 = data[1] 

# Create a LightGraphs' graph
g = SimpleGraph(d1.num_nodes)
for (s, t) in zip(d1.source, d1.target)
    add_edge!(g, s, t)
end

# Node features
X = d1.node_attributes # (nfeatures x nnodes) matrix
```
"""
struct TUDataset <: AbstractDataset
    name::String
    metadata::Dict{String, Any}
    graphs::Vector{Graph}
    graph_data::Union{Nothing,NamedTuple}
    num_nodes::Int
    num_edges::Int
    num_graphs::Int
end

function TUDataset(name; dir=nothing)
    d = tudataset_datadir(name, dir)
    # See here for the file format: https://chrsmrrs.github.io/datasets/docs/format/
    
    st = readdlm(joinpath(d, "$(name)_A.txt"), ',', Int)
    # Check that the first node is labeled 1.
    # TODO this will fail if the first node is isolated
    @assert minimum(st) == 1
    source, target = st[:,1], st[:,2]

    graph_indicator = readdlm(joinpath(d, "$(name)_graph_indicator.txt"), Int) |> vec      
    @assert all(sort(unique(graph_indicator)) .== 1:length(unique(graph_indicator)))

    num_nodes = length(graph_indicator)
    num_edges = length(source)
    num_graphs = length(unique(graph_indicator))

    # LOAD OPTIONAL FILES IF EXIST
    
    node_labels = isfile(joinpath(d, "$(name)_node_labels.txt")) ?
                    readdlm(joinpath(d, "$(name)_node_labels.txt"), Int) |> vec :
                    nothing
    edge_labels = isfile(joinpath(d, "$(name)_edge_labels.txt")) ?
                    readdlm(joinpath(d, "$(name)_edge_labels.txt"), Int) |> vec :
                    nothing
    graph_labels = isfile(joinpath(d, "$(name)_graph_labels.txt")) ?
                    readdlm(joinpath(d, "$(name)_graph_labels.txt"), Int) |> vec :
                    nothing

    node_attributes = isfile(joinpath(d, "$(name)_node_attributes.txt")) ?
                        readdlm(joinpath(d, "$(name)_node_attributes.txt"), ',', Float32)' |> collect :
                        nothing
    edge_attributes = isfile(joinpath(d, "$(name)_edge_attributes.txt")) ?
                        readdlm(joinpath(d, "$(name)_edge_attributes.txt"), ',', Float32)' |> collect :
                        nothing
    graph_attributes = isfile(joinpath(d, "$(name)_graph_attributes.txt")) ?
                        readdlm(joinpath(d, "$(name)_graph_attributes.txt"), ',', Float32)' |> collect :
                        nothing

    # We need this two vectors sorted for efficiency in tudataset_getgraph(full_dataset, i)
    @assert issorted(graph_indicator)
    if !issorted(source)
        p = sortperm(source)
        source, target[p] = source[p], target[p]
        if edge_labels !== nothing
            edge_labels = edge_labels[p]
        end
        if edge_attributes !== nothing
            edge_attributes = edge_attributes[:,p]
        end
    end


    full_dataset = (; num_nodes, num_edges, num_graphs,
                        source, target, 
                        graph_indicator,
                        node_labels,
                        edge_labels,            
                        graph_labels,
                        node_attributes, 
                        edge_attributes,
                        graph_attributes)

    graphs = [tudataset_getgraph(full_dataset, i) for i in 1:num_graphs]
    graph_data = (; features = graph_attributes, targets = graph_labels) |> clean_nt
    metadata = Dict{String, Any}("name" => name)
    return TUDataset(name, metadata, graphs, graph_data, num_nodes, num_edges, num_graphs)
end

function tudataset_datadir(name, dir = nothing)
    dir = isnothing(dir) ? datadep"TUDataset" : dir
    LINK = "https://www.chrsmrrs.com/graphkerneldatasets/$name.zip"
    d  = joinpath(dir, name)
    if !isdir(d)
        DataDeps.fetch_default(LINK, dir)
        currdir = pwd()
        cd(dir) # Needed since `unpack` extracts in working dir
        DataDeps.unpack(joinpath(dir, "$name.zip"))
        cd(currdir)
    end
    @assert isdir(d)
    return d
end


function tudataset_getgraph(data::NamedTuple, i::Int)
    vmin = searchsortedfirst(data.graph_indicator, i)
    vmax = searchsortedlast(data.graph_indicator, i)
    nodes = vmin:vmax
    node_labels = isnothing(data.node_labels) ? nothing : getobs(data.node_labels, nodes)
    node_attributes = isnothing(data.node_attributes) ? nothing : getobs(data.node_attributes, nodes)
    
    emin = searchsortedfirst(data.source, vmin)
    emax = searchsortedlast(data.source, vmax)
    edges = emin:emax
    source = data.source[edges] .- vmin .+ 1
    target = data.target[edges] .- vmin .+ 1
    edge_labels = isnothing(data.edge_labels) ? nothing : getobs(data.edge_labels, edges)
    edge_attributes = isnothing(data.edge_attributes) ? nothing : getobs(data.edge_attributes, edges)

    num_nodes = length(nodes)
    num_edges = length(source)
    node_data = (features = node_attributes, targets = node_labels)
    edge_data = (features = edge_attributes, targets = edge_labels)
    
    return Graph(; num_nodes, num_edges, 
                edge_index = (source, target), 
                node_data = node_data |> clean_nt,
                edge_data = edge_data |> clean_nt,
                )
end


Base.length(data::TUDataset) = length(data.graphs)

function Base.getindex(data::TUDataset) 
    if data.graph_data === nothing
        return getobs(data.graphs)
    else
        return getobs((; data.graphs, data.graph_data...))    
    end
end

function Base.getindex(data::TUDataset, i) 
    if data.graph_data === nothing
        return getobs(data.graphs, i)
    else
        return getobs((; data.graphs, data.graph_data...), i)    
    end
end
