using ForwardDiff
using Distances
using LinearAlgebra
using ReverseDiff
using AutoGrad
#using CuArrays
using Flux
using Flux.Tracker
using Flux: @epochs


LR = 0.1 # learning rate
EPOCHS = 100 # number of epochs
DIM = 2 # dimensionality of embeddings

function radius(x)
    return x[end]
end

function centerpoint(x)
    return x[1:end-1]
end

function loss1(c, d) # loss for normal form 1
    return max(0, euclidean(centerpoint(c), centerpoint(d)) + abs(radius(c))-abs(radius(d)))
end

function loss2(c1, c2, d) # loss for c1 and c2 SubClassOf: d
    dist = euclidean(centerpoint(c1), centerpoint(c2))
    if dist > abs(radius(c1) + radius(c2))
        return abs(radius(c2) + radius(c2)) # no solution, circles are separate
    elseif dist < abs(radius(c1)-radius(c2)) # no solution, one circle contained in other
        if radius(c1) < radius(c2)
            return loss1(c1, d)
        else
            return loss1(c2, d)
        end
    elseif dist == 0 && radius(c1) == radius(c2) ## circles are coincident
        return loss1(c1, d)
    else
        a = (radius(c1)^2 - radius(c2)^2 + dist^2)/(2 * dist)
        rad = sqrt(radius(c1)^2 - a^2)
        cent = centerpoint(c1) + a * (centerpoint(c2) - centerpoint(c1)) / dist
        return loss1(vcat(cent, rad),d)
    end
end

function loss3(c, d, r) # normal form 3
    return max(0, euclidean(centerpoint(c)-r, centerpoint(d)) + abs(radius(d)) - abs(radius(c)))
end

function loss4(c, d, r) # normal form 4
    return max(0, euclidean(centerpoint(c) + r, centerpoint(d)) + abs(radius(d)) - abs(radius(c)))
end

classes = Dict()
relations = Dict()
ccount = 1
rcount = 1

function filldict(x)
    global ccount
    if haskey(classes, x)
        return classes[x]
    else
        classes[x] = ccount
        ccount += 1
        return ccount - 1
    end
end

function rfilldict(x)
    global rcount
    if haskey(relations, x)
        return relations[x]
    else
        relations[x] = rcount
        rcount += 1
        return rcount - 1
    end
end

nf1 = Set()
nf2 = Set()
nf3 = Set()
nf4 = Set()


# Reading the file of normalized axioms (generated by Normalize.groovy)
open(ARGS[1]) do file
    for line in eachline(file)
        if (startswith(line, "SubClassOf")) # ignore subproperty axioms
            if occursin(r"ObjectIntersectionOf", line) # normal form 2
                m = match(r"(<.*>).*(<.*>).*(<.*>)", line) # 3 captures: C and D SubClassOf: E
                if m != nothing
                    c, d, e = map(filldict, m.captures)
                    push!(nf2, (c,d,e))
                end
            elseif occursin(r"SubClassOf.<[^\s]*>[\s]*<[^\s]*>.$", line) # normal form 1
                m = match(r"(<.*>).*(<.*>)", line) # 2 captures: C SubClassOf: D
                c,d = map(filldict, m.captures)
                push!(nf1, (c,d))
            elseif occursin(r"SubClassOf.ObjectSomeValuesFrom", line)  # normal form 4
                m = match(r"(<.*>).*(<.*>).*(<.*>)", line) # 3 captures: R some C SubClassOf: D
                r,c,d = m.captures
                c = filldict(c)
                d = filldict(d)
                r = rfilldict(r)
                push!(nf4, (c,d,r))
            elseif occursin(r"SubClassOf.<[^\s]*>[\s]*ObjectSomeValuesFrom", line)  # normal form 3
                m = match(r"(<.*>).*(<.*>).*(<.*>)", line) # 3 captures: C SubClassOf: R some D
                c,r,d = m.captures
                c = filldict(c)
                d = filldict(d)
                r = rfilldict(r)
                push!(nf3, (c,d,r))
            end
        end
    end
end

# there are now rcount-1 relations and ccount-1 classes; generate random embeddings; offload to GPU; make them trainable
#cvec = param(cu(rand(ccount-1, DIM + 1)))
#rvec = param(cu(rand(rcount-1, DIM)))
cvec = param((rand(ccount-1, DIM + 1)))
rvec = param((rand(rcount-1, DIM)))

# setting up reverse maps for faster lookups
rclasses = Dict(value => key for (key, value) in classes)
rrelations = Dict(value => key for (key, value) in relations)

# set to array, just for convenience
nf1arr =  collect(nf1)
nf2arr =  collect(nf2)
nf3arr =  collect(nf3)
nf4arr =  collect(nf4)

# loss functions that work with arrays
function loss1(x::Int, y::Int)
    return loss1(cvec[x,:], cvec[y,:])
end
function loss2(x::Int, y::Int, z::Int)
    return loss2(cvec[x,:], cvec[y,:], cvec[z,:])
end
function loss3(x::Int, y::Int, r::Int)
    return loss3(cvec[x,:], cvec[y,:], rvec[r,:])
end
function loss4(x::Int, y::Int, r::Int)
    return loss4(cvec[x,:], cvec[y,:], rvec[r,:])
end

opt = Momentum([cvec, rvec], LR, decay=0.0)

for k in 1:EPOCHS
    println("Epoch: $k")
    Flux.train!(loss1, nf1, opt)
    Flux.train!(loss2, nf2, opt)
    Flux.train!(loss3, nf3, opt)
    Flux.train!(loss4, nf4, opt)
    if k % 1 == 0
        println("Loss 1: ",  sum([loss1(x,y) for (x,y) in nf1arr]))
        println("Loss 2: ",  sum([loss2(x,y,z) for (x,y,z) in nf2arr]))
        println("Loss 3: ",  sum([loss3(x,y,r) for (x,y,r) in nf3arr]))
        println("Loss 4: ",  sum([loss4(x,y,r) for (x,y,r) in nf4arr]))
    end
end
