module SymbolicAnalysis

using DomainSets
using LinearAlgebra
using LogExpFunctions
using PrecompileTools
using StatsBase
using Distributions
using DSP, DataStructures

using Symbolics
import Symbolics: issym, Term
using SymbolicUtils: iscall, BasicSymbolic
using SymbolicUtils.Rewriters
SymbolicUtils.inspect_metadata[] = true

# Symbolics v7 / SymbolicUtils v4 removed the `Symbolic` abstract type: every
# symbolic — scalar or array — is now a `BasicSymbolic{SymReal}`.
const Symbolic = BasicSymbolic

# The scalar-symbolic type Symbolics uses when dispatching `in(::symbolic, ::Domain)`.
# Matching it exactly lets the `in(::_, ::CustomDomain)` disambiguators below stay
# strictly more specific than Symbolics' `IntervalSets.Domain` method (which is
# keyed on `BasicSymbolic{SymReal}`).
const InDomainSymbolic = BasicSymbolic{SymbolicUtils.SymReal}

struct VarDomain end

include("rules.jl")
include("atoms.jl")
include("gdcp/gdcp_rules.jl")
include("gdcp/spd.jl")
include("gdcp/lorentz.jl")
include("canon.jl")

struct AnalysisResult
    curvature::SymbolicAnalysis.Curvature
    sign::SymbolicAnalysis.Sign
    gcurvature::Union{SymbolicAnalysis.GCurvature, Nothing}
end

"""
    analyze(ex)
    analyze(ex, M)

Analyze the expression `ex` and return the curvature and sign of the expression. If a manifold `M` from [Manifolds.jl](https://juliamanifolds.github.io/Manifolds.jl/stable/) is provided, also return the geodesic curvature of the expression.
Currently supports the `SymmetricPositiveDefinite` and `Lorentz` manifolds.

The returned `AnalysisResult` contains the following fields:

  - `curvature::SymbolicAnalysis.Curvature`: The curvature of the expression.
  - `sign::SymbolicAnalysis.Sign`: The sign of the expression.
  - `gcurvature::Union{SymbolicAnalysis.GCurvature,Nothing}`: The geodesic curvature of the expression if `M` is provided. Otherwise, `nothing`.
"""
function analyze(ex, M::Union{AbstractManifold, Nothing} = nothing)
    ex = unwrap(ex)
    ex = canonize(ex)
    ex = propagate_sign(ex)
    ex = propagate_curvature(ex)
    if isnothing(M)
        return AnalysisResult(getcurvature(ex), getsign(ex), nothing)
    else
        @assert M isa SymmetricPositiveDefinite || M isa Lorentz "Only SymmetricPositiveDefinite and Lorentz manifolds are currently supported"
        ex = propagate_gcurvature(ex, M)
        return AnalysisResult(getcurvature(ex), getsign(ex), getgcurvature(ex))
    end
end

export analyze

@setup_workload begin
    @compile_workload begin
        @variables x y
        y_with_domain = setmetadata(
            y, VarDomain, DomainSets.HalfLine{Number, :open}()
        )

        ex1 = exp(y_with_domain) - log(y_with_domain) |> unwrap
        analyze(ex1)

        ex2 = abs(x)^2 + abs(x)^3 |> unwrap
        analyze(ex2)

        ex3 = 2 * abs(x) - 1 |> unwrap
        analyze(ex3)

        @variables z[1:3]
        ex4 = exp.(z) |> unwrap
        analyze(ex4)
    end
end

end
