include("HittingSet.jl")

module Engine

export Construction, mprod, codimension, dimension

import Subscripts
using LinearAlgebra
using AbstractAlgebra
using Groebner
using HomotopyContinuation: Variable, Expression, System
using ..HittingSet

# --- commutative algebra ---

# as of version 0.36.6, AbstractAlgebra only supports ideals in multivariate
# polynomial rings when the coefficients are integers. we use Groebner to extend
# support to rationals and to finite fields of prime order
Generic.reduce_gens(I::Generic.Ideal{U}) where {T <: FieldElement, U <: MPolyRingElem{T}} =
  Generic.Ideal{U}(base_ring(I), groebner(gens(I)))

function codimension(I::Generic.Ideal{U}, maxdepth = Inf) where {T <: RingElement, U <: MPolyRingElem{T}}
  leading = [exponent_vector(f, 1) for f in gens(I)]
  targets = [Set(findall(.!iszero.(exp_vec))) for exp_vec in leading]
  length(HittingSet.solve(HittingSetProblem(targets), maxdepth))
end

dimension(I::Generic.Ideal{U}, maxdepth = Inf) where {T <: RingElement, U <: MPolyRingElem{T}} =
  length(gens(base_ring(I))) - codimension(I, maxdepth)

# hat tip Sascha Timme
#   https://github.com/JuliaHomotopyContinuation/HomotopyContinuation.jl/issues/520#issuecomment-1317681521
function Base.convert(::Type{Expression}, f::MPolyRingElem)
  variables = Variable.(symbols(parent(f)))
  f_data = zip(coefficients(f), exponent_vectors(f))
  sum(cf * prod(variables .^ exp_vec) for (cf, exp_vec) in f_data)
end

# create a ModelKit.System from an ideal in a multivariate polynomial ring. the
# variable ordering is taken from the polynomial ring
function System(I::Generic.Ideal)
  eqns = Expression.(gens(I))
  variables = Variable.(symbols(base_ring(I)))
  System(eqns, variables = variables)
end

## [to do] not needed right now
# create a ModelKit.System from a list of elements of a multivariate polynomial
# ring. the variable ordering is taken from the polynomial ring
##function System(eqns::AbstractVector{MPolyRingElem})
##  if isempty(eqns)
##    return System([])
##  else
##    variables = Variable.(symbols(parent(f)))
##    return System(Expression.(eqns), variables = variables)
##  end
##end

# --- primitve elements ---

abstract type Element{T} end

mutable struct Point{T} <: Element{T}
  coords::Vector{MPolyRingElem{T}}
  vec::Union{Vector{MPolyRingElem{T}}, Nothing}
  rel::Nothing
  
  ## [to do] constructor argument never needed?
  Point{T}(
    coords::Vector{MPolyRingElem{T}} = MPolyRingElem{T}[],
    vec::Union{Vector{MPolyRingElem{T}}, Nothing} = nothing
  ) where T = new(coords, vec, nothing)
end

function buildvec!(pt::Point)
  coordring = parent(pt.coords[1])
  pt.vec = [one(coordring), dot(pt.coords, pt.coords), pt.coords...]
end

mutable struct Sphere{T} <: Element{T}
  coords::Vector{MPolyRingElem{T}}
  vec::Union{Vector{MPolyRingElem{T}}, Nothing}
  rel::Union{MPolyRingElem{T}, Nothing}
  
  ## [to do] constructor argument never needed?
  Sphere{T}(
    coords::Vector{MPolyRingElem{T}} = MPolyRingElem{T}[],
    vec::Union{Vector{MPolyRingElem{T}}, Nothing} = nothing,
    rel::Union{MPolyRingElem{T}, Nothing} = nothing
  ) where T = new(coords, vec, rel)
end

function buildvec!(sph::Sphere)
  coordring = parent(sph.coords[1])
  sph.vec = sph.coords
  sph.rel = mprod(sph.coords, sph.coords) + one(coordring)
end

const coordnames = IdDict{Symbol, Vector{Union{Symbol, Nothing}}}(
  nameof(Point) => [nothing, nothing, :xₚ, :yₚ, :zₚ],
  nameof(Sphere) => [:rₛ, :sₛ, :xₛ, :yₛ, :zₛ]
)

coordname(elt::Element, index) = coordnames[nameof(typeof(elt))][index]

function pushcoordname!(coordnamelist, indexed_elt::Tuple{Any, Element}, coordindex)
  eltindex, elt = indexed_elt
  name = coordname(elt, coordindex)
  if !isnothing(name)
    subscript = Subscripts.sub(string(eltindex))
    push!(coordnamelist, Symbol(name, subscript))
  end
end

function takecoord!(coordlist, indexed_elt::Tuple{Any, Element}, coordindex)
  elt = indexed_elt[2]
  if !isnothing(coordname(elt, coordindex))
    push!(elt.coords, popfirst!(coordlist))
  end
end

# --- primitive relations ---

abstract type Relation{T} end

mprod(v, w) = (v[1]*w[2] + w[1]*v[2]) / 2 - dot(v[3:end], w[3:end])

# elements: point, sphere
struct LiesOn{T} <: Relation{T}
  elements::Vector{Element{T}}
  
  LiesOn{T}(pt::Point{T}, sph::Sphere{T}) where T = new{T}([pt, sph])
end

equation(rel::LiesOn) = mprod(rel.elements[1].vec, rel.elements[2].vec)

# elements: sphere, sphere
struct AlignsWithBy{T} <: Relation{T}
  elements::Vector{Element{T}}
  cos_angle::T
  
  AlignsWithBy{T}(sph1::Sphere{T}, sph2::Sphere{T}, cos_angle::T) where T = new{T}([sph1, sph2], cos_angle)
end

equation(rel::AlignsWithBy) = mprod(rel.elements[1].vec, rel.elements[2].vec) - rel.cos_angle

# --- constructions ---

mutable struct Construction{T}
  points::Set{Point{T}}
  spheres::Set{Sphere{T}}
  relations::Set{Relation{T}}
  
  function Construction{T}(; elements = Set{Element{T}}(), relations = Set{Relation{T}}()) where T
    allelements = union(elements, (rel.elements for rel in relations)...)
    new{T}(
      filter(elt -> isa(elt, Point), allelements),
      filter(elt -> isa(elt, Sphere), allelements),
      relations
    )
  end
end

function Base.push!(ctx::Construction{T}, elt::Point{T}) where T
  push!(ctx.points, elt)
end

function Base.push!(ctx::Construction{T}, elt::Sphere{T}) where T
  push!(ctx.spheres, elt)
end

function Base.push!(ctx::Construction{T}, rel::Relation{T}) where T
  push!(ctx.relations, rel)
  for elt in rel.elements
    push!(ctx, elt)
  end
end

function realize(ctx::Construction{T}) where T
  # collect coordinate names
  coordnamelist = Symbol[]
  eltenum = enumerate(Iterators.flatten((ctx.spheres, ctx.points)))
  for coordindex in 1:5
    for indexed_elt in eltenum
      pushcoordname!(coordnamelist, indexed_elt, coordindex)
    end
  end
  
  # construct coordinate ring
  coordring, coordqueue = polynomial_ring(parent_type(T)(), coordnamelist, ordering = :degrevlex)
  
  # retrieve coordinates
  for (_, elt) in eltenum
    empty!(elt.coords)
  end
  for coordindex in 1:5
    for indexed_elt in eltenum
      takecoord!(coordqueue, indexed_elt, coordindex)
    end
  end
  
  # construct coordinate vectors
  for (_, elt) in eltenum
    buildvec!(elt)
  end
  
  # turn relations into equations
  eqns = vcat(
    equation.(ctx.relations),
    [elt.rel for (_, elt) in eltenum if !isnothing(elt.rel)]
  )
  
  # add relations to center and orient the construction
  if !isempty(ctx.points)
    append!(eqns, [sum(pt.coords[k] for pt in ctx.points) for k in 1:3])
  end
  if !isempty(ctx.spheres)
    append!(eqns, [sum(sph.coords[k] for sph in ctx.spheres) for k in 3:4])
  end
  
  (Generic.Ideal(coordring, eqns), eqns)
end

end

# ~~~ sandbox setup ~~~

using Random
using Distributions
using AbstractAlgebra
using HomotopyContinuation

CoeffType = Rational{Int64}

a = Engine.Point{CoeffType}()
s = Engine.Sphere{CoeffType}()
a_on_s = Engine.LiesOn{CoeffType}(a, s)
ctx = Engine.Construction{CoeffType}(elements = Set([a]), relations= Set([a_on_s]))
##ideal_a_s = Engine.realize(ctx)
##println("A point on a sphere: ", Engine.dimension(ideal_a_s), " degrees of freedom")

b = Engine.Point{CoeffType}()
b_on_s = Engine.LiesOn{CoeffType}(b, s)
Engine.push!(ctx, b)
Engine.push!(ctx, s)
Engine.push!(ctx, b_on_s)
ideal_ab_s, eqns_ab_s = Engine.realize(ctx)
freedom = Engine.dimension(ideal_ab_s)
println("Two points on a sphere: ", freedom, " degrees of freedom")

##spheres = [Engine.Sphere{CoeffType}() for _ in 1:3]
##tangencies = [
##  Engine.AlignsWithBy{CoeffType}(
##    spheres[n],
##    spheres[mod1(n+1, length(spheres))],
##    CoeffType(-1//1)
##  )
##  for n in 1:3
##]
##ctx_tan_sph = Engine.Construction{CoeffType}(elements = Set(spheres), relations = Set(tangencies))
##ideal_tan_sph = Engine.realize(ctx_tan_sph)
##println("Three mutually tangent spheres: ", Engine.dimension(ideal_tan_sph), " degrees of freedom")

# --- test rational cut ---

coordring = base_ring(ideal_ab_s)
vbls = Variable.(symbols(coordring))
##cut_system = CompiledSystem(System([eqns_ab_s; cut], variables = vbls))
##cut_result = HomotopyContinuation.solve(cut_system)
##println("non-singular solutions:")
##for soln in solutions(cut_result)
##  display(soln)
##end
##println("singular solutions:")
##for sing in singular(cut_result)
##  display(sing.solution)
##end

# test a random witness set
system = CompiledSystem(System(eqns_ab_s, variables = vbls))
max_slope = 2
binom = Binomial(2max_slope, 1/2)
Random.seed!(6071)
samples = []
for _ in 1:3
  cut_matrix = rand(binom, freedom, length(gens(coordring))) .- max_slope
  ##cut_matrix = [
  ##  1 1 1 1 0 1 1 0 1 1 0;
  ##  1 2 1 2 0 1 1 0 1 1 0;
  ##  1 1 0 1 0 1 2 0 2 0 0
  ##]
  sph_z_ind = indexin([sph.coords[5] for sph in ctx.spheres], gens(coordring))
  cut_offset = [sum(cf[sph_z_ind]) for cf in eachrow(cut_matrix)]
  println("sphere z variables: ", vbls[sph_z_ind])
  display(cut_matrix)
  display(cut_offset)
  cut_subspace = LinearSubspace(cut_matrix, cut_offset)
  wtns = witness_set(system, cut_subspace)
  append!(samples, solution.(filter(isreal, results(wtns))))
end
println("$(length(samples)) sample solutions:")
for soln in samples
  display([vbls round.(soln, digits = 6)])
  k_sq = abs2(soln[1])
  if abs2(soln[end-2]) > 1e-12
    if k_sq < 1e-12
      println("center at infinity: z coordinates $(round(soln[end], digits = 6)) and $(round(soln[end-1], digits = 6))}")
    else
      sum_sq = soln[4]^2 + soln[7]^2 + soln[end-2]^2 / k_sq
      println("center on z axis: r² = $(round(1/k_sq, digits = 6)), x² + y² + h² = $(round(sum_sq, digits = 6))")
    end
  else
    sum_sq = sum(soln[[4, 7, 10]] .^ 2)
    println("center at origin: r² = $(round(1/k_sq, digits = 6)); x² + y² + z² = $(round(sum_sq, digits = 6))")
  end
end