###############
# Computes cartesian product of polynomial rings by overloading the product to drop terms 
###############

using Oscar
using Unicode


export product_polynomial_ring, ProdPolyRing, ProdPolyIdeal

# export print_with_subscripts, display_with_subscripts, latex_representation

struct ProdPolyRing{T, R<:Ring, S<:MPolyRing} <: Ring
    R::R
    N::Int
    S::S
    Sgr::MPolyDecRing

    function ProdPolyRing(R::Ring, N::Int, vars::Vector{String})
        S, _ = polynomial_ring(R, vars)
        T = Oscar.elem_type(S)
        Sgr = grade(S,vcat([[1,0] for i in 1:N],[[0,1] for i in 1:N]))[1]
        new{T, typeof(R), typeof(S)}(R, N, S, Sgr)
    end
end

# Constructor
function ProdPolyRing(F::Ring, N::Int; z::Union{AbstractString, Char, Symbol} = "z")
    vars = vcat([string(z, "_", i) for i in 0:N-1],[string(z,"bar_",j) for j in 0:N-1])
    ProdPolyRing(F, N, vars)
end


product_polynomial_ring(F::Ring, N::Int; z::Union{AbstractString, Char} = "z") = ProdPolyRing(F, N, z=z)
product_polynomial_ring(F::Ring, N::Int, z::Symbol) = ProdPolyRing(F, N, z=z)

# Display method
function Base.show(io::IO, R::ProdPolyRing)
    z = split(string(R.S.S[1]),'_')[1]
    print(io, "Polynomial ring in $(R.N)+$(R.N) variables $(z)_⋅, $(z)bar_⋅ over ", R.R)
end




# Implement necessary methods
Oscar.base_ring(R::ProdPolyRing) = R.R
Oscar.ngens(R::ProdPolyRing) = 2*R.N
Oscar.gens(R::ProdPolyRing) = gens(R.S)
Oscar.gen(R::ProdPolyRing, i::Int) = gen(R.S, i)

# Add matrix-specific methods
#function gen(R::ProdPolyRing)
#    matrix(R.S, R.N, R.N, gens(R.S))
#end

#function Oscar.gen(R::ProdPolyRing, i::Int, j::Int)
#    gen(R.S, (i % R.N)*R.N + (j % R.N) + 1)
#end

# Element type
Oscar.elem_type(::Type{<:ProdPolyRing{T}}) where {T} = T

# Define zero and one for the ring
function Oscar.zero(R::ProdPolyRing)
    return matrix(R.S, R.N, R.N, [zero(R.S) for _ in 1:R.N^2])
end

function Oscar.one(R::ProdPolyRing)
    return matrix(R.S, R.N, R.N, [i == j ? one(R.S) : zero(R.S) for i in 1:R.N for j in 1:R.N])
end

# Conversion method
function (R::ProdPolyRing)(a)
    return R.S(a)
end





##############
# ProdPolyIdeal interface
##############
#  https://docs.oscar-system.org/stable/AlgebraicGeometry/FrameWorks/ArchitectureOfAffineSchemes/

# ProdPolyIdeal definition
struct ProdPolyIdeal{T, R<:ProdPolyRing} <: Ideal{T}
    R::R
    gens::Vector{T}
    I::MPolyIdeal

    function ProdPolyIdeal(R::ProdPolyRing, g::Union{RingElem, RingElement, Integer}...)
        T = Oscar.elem_type(R.S)
        oscar_gens = [R.S(R(f)) for f in g]
        I = ideal(R.S, oscar_gens)
        new{T, typeof(R)}(R, [R(f) for f in g], I)
    end

    function ProdPolyIdeal(R::ProdPolyRing, g::Vector)
        T = Oscar.elem_type(R.S)
        oscar_gens = [R.S(R(f)) for f in g]
        I = ideal(R.S, oscar_gens)
        new{T, typeof(R)}(R, [R(f) for f in g], I)
    end
end

# Ideal constructors
function Oscar.ideal(R::ProdPolyRing, g::Union{RingElem, RingElement, Integer}...)
    ProdPolyIdeal(R, g...)
end

function Oscar.ideal(R::ProdPolyRing, g::Vector)
    ProdPolyIdeal(R, g)
end

# Quotient ring constructor
function Oscar.quo(R::ProdPolyRing, I::ProdPolyIdeal)
    Q, f = quo(R.S, I.I)
    return Q, f
end

# Ideal membership test
function Base.in(f::RingElem, I::ProdPolyIdeal)
    return I.R.S(f) in I.I
end

# Generators
Oscar.gens(I::ProdPolyIdeal) = I.gens

# Coordinates
function Oscar.coordinates(f::RingElem, I::ProdPolyIdeal)
    return Oscar.coordinates(I.R.S(f), I.I)
end

# Show method for ProdPolyIdeal
function Base.show(io::IO, I::ProdPolyIdeal)
    print(io, "Ideal in ", I.R, " generated by ", I.gens)
end

# Implement necessary methods for Ideal interface
Oscar.base_ring(I::ProdPolyIdeal) = I.R
Oscar.ngens(I::ProdPolyIdeal) = length(I.gens)
Oscar.gen(I::ProdPolyIdeal, i::Int) = I.gens[i]


# Sum of ideals
function Base.:+(I::ProdPolyIdeal, J::ProdPolyIdeal)
    @assert I.R == J.R "Ideals must be in the same ring"
    K = I.I + J.I  # Using Oscar's implementation
    return ProdPolyIdeal(I.R, gens(K))
end

# Product of ideals
function Base.:*(I::ProdPolyIdeal, J::ProdPolyIdeal)
    @assert I.R == J.R "Ideals must be in the same ring"
    K = I.I * J.I  # Using Oscar's implementation
    return ProdPolyIdeal(I.R, gens(K))
end

# Intersection of ideals
function Oscar.intersect(I::ProdPolyIdeal, J::ProdPolyIdeal)
    @assert I.R == J.R "Ideals must be in the same ring"
    K = intersect(I.I, J.I)  # Using Oscar's implementation
    return ProdPolyIdeal(I.R, gens(K))
end

# Quotient of ideals
function Oscar.quotient(I::ProdPolyIdeal, J::ProdPolyIdeal)
    @assert I.R == J.R "Ideals must be in the same ring"
    K = quotient(I.I, J.I)  # Using Oscar's implementation
    return ProdPolyIdeal(I.R, gens(K))
end

# Saturation of an ideal
function Oscar.saturation(I::ProdPolyIdeal, J::ProdPolyIdeal)
    @assert I.R == J.R "Ideals must be in the same ring"
    K = saturation(I.I, J.I)  # Using Oscar's implementation
    return ProdPolyIdeal(I.R, gens(K))
end

# Radical of an ideal
function Oscar.radical(I::ProdPolyIdeal)
    K = radical(I.I)  # Using Oscar's implementation
    return ProdPolyIdeal(I.R, gens(K))
end

# Primary decomposition of an ideal
function Oscar.primary_decomposition(I::ProdPolyIdeal)
    decomp = primary_decomposition(I.I)  # Using Oscar's implementation
    return [ProdPolyIdeal(I.R, gens(J)) for J in decomp]
end

# Equality test for ideals
function Base.:(==)(I::ProdPolyIdeal, J::ProdPolyIdeal)
    return I.R == J.R && I.I == J.I
end

# Subset test for ideals
function Base.issubset(I::ProdPolyIdeal, J::ProdPolyIdeal)
    @assert I.R == J.R "Ideals must be in the same ring"
    return issubset(I.I, J.I)
end




#######################################################################
#######################################################################
# Useful polynomials 
#######################################################################
#######################################################################


function fij(S::ProdPolyRing,j1::Union{Int,zzModRingElem},j2::Union{Int,zzModRingElem}) 
    X(i,j) = gen(S,i,j)
    sum([X(a,a+j1)*X(a+j1+j2,a+j2) for a=0:S.N-1])
end

function fij(S::ProdPolyRing,j::zzModMatrix) 
    fij(S,j[1],j[2])
end

function fplus(S::ProdPolyRing,j1::Union{Int,zzModRingElem},j2::Union{Int,zzModRingElem}) 
    (1//2)*(fij(S,j1,j2) + fij(S,j2,j1))
end

function fplus(S::ProdPolyRing,j::zzModMatrix) 
    fplus(S,j[1],j[2])
end

function fminus(S::ProdPolyRing,j1::Union{Int,zzModRingElem},j2::Union{Int,zzModRingElem}) 
    (1//2)*(fij(S,j1,j2) - fij(S,j2,j1))
end

function fminus(S::ProdPolyRing,j::zzModMatrix) 
    fminus(S,j[1],j[2])
end

## Projecting onto harmonic subspace

function hij(S::ProdPolyRing,j1::Union{Int,zzModRingElem},j2::Union{Int,zzModRingElem}) 
    X = gen(S) 
    c = (1//2)*(tr(X^2) + tr(X)^2)*(1//(S.N+1))
    fij(S,j1,j2) - S.Sgr((Int(j1) == 0 ? 1 : 0) + (Int(j2) == 0 ? 0 : 1))*c
end

function hij(S::ProdPolyRing,j::zzModMatrix) 
    hij(S,j[1],j[2])
end

function hplus(S::ProdPolyRing,j1::Union{Int,zzModRingElem},j2::Union{Int,zzModRingElem})
    X = gen(S) 
    c = (1//2)*(tr(X^2) + tr(X)^2)*(1//(S.N+1))
    fplus(S,j1,j2) - ((Int(j1) == 0 ? 1 : 0) + (Int(j2) == 0 ? 1 : 0))*c
end

function hplus(S::ProdPolyRing,j::zzModMatrix) 
    hplus(S,j[1],j[2])
end

function hminus(S::ProdPolyRing,j1::Union{Int,zzModRingElem},j2::Union{Int,zzModRingElem}) 
    X = gen(S) 
    c = (1//2)*(tr(X^2) - tr(X)^2)*(1//(S.N-1))
    fminus(S,j1,j2) - ((Int(j1) == 0) || (Int(j2) == 0) ? 1 : 0)*c
end

function hminus(S::ProdPolyRing,j::zzModMatrix) 
    hminus(S,j[1],j[2])
end


## Laplacian 
# Should eventually replace Laplacian and laplacian from Polys.jl
function laplacian(S::ProdPolyRing,f)
    sum([derivative(f,gen(S,a,a)) for a in 0:S.N-1])
end

## Ideals

function Ihplus(S::ProdPolyRing)
    ideal(S.Sgr,[hplus(S,j1,j2) for j1 in 0:S.N-1 for j2 in 0:j1])
end

Ihminus(S::ProdPolyRing) = ideal(S.Sgr,[hminus(S,i,j) for i in 0:S.N-1 for j in 0:i if i != j])

function Iminors(S::ProdPolyRing,k) 
    ideal(S.Sgr,minors(matrix(S),k))
end

Ireal(S::ProdPolyRing) = ideal(S.Sgr,[gen(S,i,j) - gen(S,j,i) for i in 0:S.N-1 for j in 0:i if i != j])

function Itorus(S::ProdPolyRing,t::zzModMatrix) 
    N = S.N
    ideal(S.Sgr,[gen(S,ZN(N)[j1 j2]) - gen(S,ZN(N)[j1 j2]*t) for j1 in 0:S.N-1 for j2 in 0:S.N-1])
end

# Ideal generated by action of cyclic group <t> on generators.  
function Icyclic(S::ProdPolyRing,t::zzModMatrix) 
    N = S.N
    ideal(S.Sgr,[gen(S,ZN(N)[j1 j2]) - gen(S,ZN(N)[j1 j2]*t) for j1 in 0:S.N-1 for j2 in 0:S.N-1])
end

#function Oscar.irrelevant_ideal(R::MPolyRing)
#    ideal(gens(R))
#end

function irrelevant_ideal(S::ProdPolyRing)
    irrelevant_ideal(S.Sgr)
end


function dwork_polynomial(P::ProdPolyRing,mu,mubar)
    X = gen(P)
    N = P.N
    
    ss = sum(X[i,j]^N for i in 1:N for j in 1:N) # (sum z_i^N)(sum w_i^N)
    pp = prod(X[i,i] for i in 1:N) # (prod z_i^N)(prod w_i^N)
    sp = sum(prod(X[i,j] for j in 1:N) for i in 1:N) # (sum z_i^N)(prod w_i^N)
    ps = sum(prod(X[j,i] for j in 1:N) for i in 1:N) # (prod z_i^N)(sum w_i^N)
    
    ss - N*(ps*mu + sp*mubar) + N^2*mu*mubar*pp
end