struct CodePoint
    expr
    line::Union{LineNumberNode, Nothing}
end


struct CalculationSequenceError <: Exception
    var::String
end
CalculationSequenceError() = begin
    errmessage = join(split(
        """
        The `@T` macro was unable to order the given formula(s) in a way that 
        resulted in a correct calculation flow. Note that `@T` cannot take 
        indexing into account with (1) a mix of forwards `A[t+1]` and backwards 
        `A[t-1]` referencing, (2) with runtime variables like `A[c+t]`, (3) with 
        nonlinear t indexing like `A[t^2-3t]`, and (4) with non-t indexing 
        like A[34].
        """, "\n"
    ))

    CalculationSequenceError(errmessage)
end


"
Function to make testing between macros easier
"
function striplines(ex)
    ex′ = Expr(:block, deepcopy(ex))

    # Define the mutation
    function recurse(ex)
        if ex isa Expr
            for i ∈ 1:length(ex.args)
                if ex.args[i] isa LineNumberNode
                    ex.args[i] = nothing
                else 
                    recurse(ex.args[i])
    
                end 
            end 
        end 
    end

    recurse(ex′)

    ex = ex′.args[1]
    return ex

end

@testset "has_symbol" begin
    @test has_symbol(:(a + b + c + foo(a + b + c + foo(d))), :d)
    @test !has_symbol(:(a + b + c + foo(a + b + c + foo(d))), :q)
end

"
Test if an expression contains the specefic variable
"
has_symbol(ex::Expr, s::Symbol) = any(has_symbol.(ex.args, s))
has_symbol(ex::Symbol, s::Symbol) = ex == s
has_symbol(ex, s::Symbol) = false


@testset "get_variables" begin
    @test get_variables(:(a = foo(b + c + bar(d) + e[f]))) == [:a, :b, :c, :d, :e, :f]
end

"
Get all the variable names from an expression, this excludes macro names and function names
"
function get_variables(ex)
    varnames = Array{Symbol, 1}()

    # behaviour for each type
    recurse(ex) = nothing
    recurse(ex::Symbol) = push!(varnames, ex)
    recurse(ex::Expr) = begin
        i₁ = ex.head ∈ [:call, :macocall] ? 2 : 1 
        for i ∈ i₁:length(ex.args)
            recurse(ex.args[i])
        end
    end 

    # Apply and return
    recurse(ex) 
    return varnames
end


"
Assign a value to a variable within an expression
"
function assign_variable(ex, var::Symbol, value)
    ex′ = Expr(:block, deepcopy(ex))

    # behaviour for each type
    recurse(ex) = nothing
    recurse(ex::Expr) = begin
        i₁ = ex.head ∈ [:call, :macocall] ? 2 : 1 
        for i ∈ i₁:length(ex.args)
            if ex.args[i] == var
                ex.args[i] = value
            else
                recurse(ex.args[i])
            end
        end
    end 

    # Apply and return
    recurse(ex′)
    ex = ex′.args[1]
    return ex
end


@testset "get_t_refs" begin
    @test( 
        get_t_refs(:(a[5] + a[8] + b[t+40])) == 
        [:b => :(t + 40)]
    )
    @test(
        get_t_refs(:(a[5] + a[8] + b[t+40] + b[t-30] + c[t])) == 
        [:b => :(t + 40), :b => :(t - 30), :c => :t]
    )
end
"
Find references involving `t` in an equation. Note: nested referencing is currently not supported.
"
function get_t_refs(ex)
    references = Vector()

    recurse(ex) = nothing
    recurse(ex::Expr) = begin
        if ex.head == :ref && length(ex.args) == 2 && ex.args[1] isa Symbol
            (ex₁, ex₂) = ex.args 

            if has_symbol(ex₂, :t)
                push!(references, ex₁=>ex₂)
            end
        end
        for i in ex.args
            recurse(i)
        end
    end

    recurse(ex)
    return references
end


@testset "vars_without_t_index" begin
    @test vars_without_t_index(:(a+b[t])) == Set([:a])
    @test vars_without_t_index(:a) == Set([:a])
    @test vars_without_t_index(:(a+b[t]+c+t)) == Set([:a,:c])
end
"
Collect all variables in an expression that isn't indexed by `t`
"
function vars_without_t_index(ex)
    vars = Set()
    ex′ = Expr(:begin, ex)

    # behaviour for each type
    recurse(ex::Expr) = begin
        i₁ = ex.head ∈ [:call, :macocall] ? 2 : 1 
        for i ∈ i₁:length(ex.args)
            if i==i₁ && ex.head == :ref 
                # nothing
            elseif ex.args[i] isa Symbol && ex.args[i] != :t
                push!(vars, ex.args[i])
            elseif ex.args[i] isa Expr
                recurse(ex.args[i])
            end
        end
    end 

    recurse(ex′)
    return vars
end


@testset "t_linearcombinations" begin
    @test t_linearcombinations(:(t*t)) == false
    @test t_linearcombinations(:(t^2)) == false
    @test t_linearcombinations(:((5+67)*foo(t))) == false
    @test t_linearcombinations(:(t == 0 ? 1 : 2)) == false
    
    @test t_linearcombinations(:(t+1+2/3*5))
    @test t_linearcombinations(:(t*2+1+2/3*5))
    @test t_linearcombinations(:((5+67)*t))

    @test t_linearcombinations(:(1+(a+t))) == true
    @test t_linearcombinations(:(1+(a+2t))) == true
    @test t_linearcombinations(:(1+(a+(t*t)))) == false    
    @test t_linearcombinations(:(1+(a+(q*q)))) === nothing # no t
    @test t_linearcombinations(:(1+(a+(t)+t-t+5))) == false
end
"
Test if expression contains an linear combination of t.
    if t is linear: true
    if t is non-linear: false
    if not explicit: nothing
"
function t_linearcombinations(ex)

    # Can we skip the investigation?
    vars = get_variables(ex)
    if !(:t in vars)
        return nothing
    elseif [i for i in vars if i == :t] != [:t]
        return false
    end

    # behaviour for each type
    recurse(args...) = nothing
    recurse(ex::Symbol, linear_parents) = ex == :t ? linear_parents : nothing 
    recurse(ex::Expr, linear_parents) = begin
        # Test if only + - and multiplication by a constant in t
        if ex.head == :call && (ex.args[1] == :+ || ex.args[1] == :- || ex.args[1] == :*)
            calls = [recurse(i, linear_parents) for i in ex.args]

            if false in calls 
                false
            elseif true in calls
                true
            else
                nothing
            end
        else
            calls = [recurse(i, false) for i in ex.args]
            if false in calls
                false
            else 
                nothing
            end
        end
    end
    

    # call the recursion
    return recurse(ex, true)
end


@testset "get_t" begin
    t = 2
    @test get_t([1,2,3], t) == 2
    @test get_t(37246, 2) == 37246
end
function get_t(var, t)
    var isa Vector ? var[t] : var
end


@testset "wrap_t_getter" begin
    @test wrap_t_getter(:(a+b[t])) == :($get_t(a, t) + b[t])
    @test wrap_t_getter(:a) == :($get_t(a, t))
    @test wrap_t_getter(:(a+b[t]+c+t)) == :($get_t(a, t) + b[t] + $get_t(c, t) + t)
end
function wrap_t_getter(ex)
    ex′ = Expr(:begin, deepcopy(ex))

    # behaviour for each type
    recurse(ex::Expr) = begin
        i₁ = ex.head ∈ [:call, :macocall] ? 2 : 1 
        for i ∈ i₁:length(ex.args)
            if i==i₁ && ex.head == :ref 
                # nothing
            elseif ex.args[i] isa Symbol && ex.args[i] != :t
                ex.args[i] = Expr(:call, get_t, ex.args[i], :t)
            elseif ex.args[i] isa Expr
                recurse(ex.args[i])
            end
        end
    end 

    recurse(ex′)

    ex = ex′.args[1]
    return ex
end


expr_to_assignment_dict(expr; line::Union{Nothing, LineNumberNode}=nothing) = begin
    # Convert to convenient form
    if !(expr isa Expr) || (expr.head != :(block))
        expr = Expr(:block, expr)
    end

    # Collect all definitions into dict (e.g. Dict(:A=>:(B+C), :B=>:(C+D))
    formulas = OrderedDict{Symbol, CodePoint}()
    lastline = line
    for e ∈ expr.args
        if e isa LineNumberNode || e === nothing
            lastline = e
            continue
        end

        if !(expr isa Expr) || (e.head != :(=))
            throw(ErrorException("Only assignment expressions are valid, like `A = 5`, got: $(e)"))
        end
        formulas[e.args[1]] = CodePoint(e.args[2], lastline)
        lastline = line
    end

    return formulas
end


"
Last item in an OrderedDict
"
last(d::OrderedDict) = begin
    for (i, item) ∈ enumerate(d)
        i == length(d) && return item
    end
end