export @p_str, StaticPolynomial, differentiate, integrate, degree, coeffs, variable, isconstant

abstract type AbstractStaticPolynomial{T} <: Function end

struct StaticPolynomial{T<:Number, N} <: AbstractStaticPolynomial{T}
    coeffs::NTuple{N, T}
    var::Symbol
    function StaticPolynomial(coeffs::Tuple, var::Symbol = :x)
        T = supertype_all(coeffs...)
        isempty(coeffs) && return StaticPolynomial{T, 0}()
        N = length(coeffs)
        c = promote(coeffs...) |> NTuple{N, T}
        var = N > 1 ? var : :CONSTANT
        var = isconstant(coeffs) ? :CONSTANT : var
        return new{T, N}(c, var)
    end
    StaticPolynomial{T, N}(coeffs::NTuple{N, T}, var::Symbol) where {T<:Number, N}= new{T, N}(coeffs, var)
end

coeffs(p::AbstractStaticPolynomial) = p.coeffs
variable(p::AbstractStaticPolynomial) = p.var

(p::StaticPolynomial)(x) = evalpoly(x, coeffs(p))

StaticPolynomial(coeffs::AbstractVector, var::Symbol = :x) = StaticPolynomial(Tuple(coeffs), var)

# Constructors
function StaticPolynomial(spcoeffs::AbstractVector{T}, exponents::AbstractVector{<:Integer}, var::Symbol = :x) where T<:Number
    coeffs = zeros(T, maximum(exponents) + 1)
    for (c, e) in zip(spcoeffs, exponents)
        coeffs[e + 1] += c
    end
    return StaticPolynomial(coeffs, var)
end

function StaticPolynomial(spcoeffs::Tuple, exponents::Tuple, var::Symbol = :x)
    coeffs = zeros(supertype_all(spcoeffs...), maximum(exponents) + 1)
    for (c, e) in zip(spcoeffs, exponents)
        coeffs[e + 1] += c
    end
    return StaticPolynomial(coeffs, var)
end

StaticPolynomial(coeffs...) = StaticPolynomial(coeffs)

# Indexing and utility functions
degree(p::StaticPolynomial{T, N}) where {T<:Number, N} = N - 1
Base.length(p::StaticPolynomial{T, N}) where {T<:Number, N} = N

Base.getindex(p::StaticPolynomial{T, N}, inds...) where {T<:Number, N} = try
        coeffs(p)[inds...]
    catch BoundsError
        zero(T)
    end

# Conversion and promotion rules
Base.convert(::Type{T}, x::StaticPolynomial) where T<:StaticPolynomial = T(x)
Base.convert(::Type{T}, x::StaticPolynomial) where T<:Number = T(x)

StaticPolynomial(p::StaticPolynomial) = p

StaticPolynomial{T}(p::StaticPolynomial{S}) where {T<:Number, S<:Number} = StaticPolynomial(broadcast(T, p.coeffs))

function StaticPolynomial{T, N}(p::StaticPolynomial{S, M}) where {T<:Number, S<:Number, M, N}

    N < M && error("Cannot convert a polynomial of degree $(M-1) to degree $(N-1)")
    coeffs = zeros(T, N)
    for (i, c) in enumerate(p.coeffs)
        coeffs[i] = c
    end
    return StaticPolynomial(coeffs, p.var)
end

function StaticPolynomial{T, N}(p::StaticPolynomial{T, M}) where {T<:Number, N, M}
    N < M && error("Cannot convert a polynomial of degree $(M-1) to degree $(N-1)")
    p1 = StaticPolynomial(zeros(T, N))::StaticPolynomial{T, N}
    return p + p1
end

(::Type{T})(p::StaticPolynomial) where {T<:Number} =
    degree(p) |> iszero ? T(coeffs(p)[1])::T : throw(InexactError(nameof(T), T, p))

Base.promote_rule(::Type{StaticPolynomial{T, N}}, ::Type{S}) where {T<:Number, S<:Number, N} =
    StaticPolynomial{promote_type(T, S), N}

Base.promote_rule(::Type{StaticPolynomial{T, N}}, ::Type{StaticPolynomial{S, N}}) where {T<:Number, S<:Number, N} =
    StaticPolynomial{promote_type(T, S), N}

Base.promote_rule(::Type{StaticPolynomial{T, N}}, ::Type{StaticPolynomial{S, M}}) where {T<:Number, S<:Number, M, N} =
    StaticPolynomial{promote_type(T, S), max(N, M)}

# define basic arithmetic operations
Base.isapprox(p1::StaticPolynomial{T, N}, p2::StaticPolynomial{S, M}) where {T<:Number, S<:Number, M, N} =
    p1.var == p2.var && M == N && all([p1[i] == p2[i] for i in 1:M])

Base.:(==)(p1::StaticPolynomial, p2::StaticPolynomial) = (p1.coeffs == p2.coeffs && p1.var == p2.var) || (iszero(p1) && iszero(p2)) || (isone(p1) && isone(p2))

isconstant(p::StaticPolynomial) = isconstant(p.coeffs) && p.var == :CONSTANT
isconstant(coeffs) = all(iszero.(coeffs[2:end]))

function polynomial_arithmetic(operation::Function, p1::StaticPolynomial{T, N}, p2::StaticPolynomial{S, M}) where {T<:Number, S<:Number, M, N}
    R = promote_type(T, S)
    #@show p1
    #@show p2
    if p1.var == p2.var || isconstant(p1) || isconstant(p2)
        coeffs = (R(operation(p1[i], p2[i])) for i in 1:max(M, N)) |> collect
        if isconstant(p1)
            return StaticPolynomial(coeffs, p2.var)
        else
            return StaticPolynomial(coeffs, p1.var)
        end
    else
        error("Polynomials must be in terms of the same variable.")
    end
end

Base.:(+)(p1::StaticPolynomial, p2::StaticPolynomial) = polynomial_arithmetic(+, p1, p2)
Base.:(+)(p::StaticPolynomial, n::S) where {S<:Number} = p + StaticPolynomial(n)
Base.:(+)(n::S, p::StaticPolynomial) where {S<:Number} = p + n

Base.:(-)(p1::StaticPolynomial, p2::StaticPolynomial) = polynomial_arithmetic(-, p1, p2)
Base.:(-)(p::StaticPolynomial, n::S) where {S<:Number} = p - StaticPolynomial(n)
Base.:(-)(n::S, p::StaticPolynomial) where {S<:Number} = StaticPolynomial(n) - p

Base.:(/)(p::StaticPolynomial, n::S) where {S<:Number} = StaticPolynomial(coeffs(p)./n)

Base.:(//)(p::StaticPolynomial, n::S) where {S<:Number} = StaticPolynomial(coeffs(p).//n)

function Base.:(^)(p::StaticPolynomial, n::S) where {S<:Integer}
    if n == 1
        return p
    elseif n == 0
        return one(p)
    else
        return p * p^(n-1)
    end
end

Base.zero(p::StaticPolynomial{T, N}) where {T<:Number, N} = zero(StaticPolynomial{T, N}, p.var)
Base.zero(::Type{StaticPolynomial{T, N}}, var=:x) where {T, N} = StaticPolynomial(zeros(T, N), :CONSTANT)
Base.zero(::Type{StaticPolynomial{T}}) where T = zero(StaticPolynomial{T, 1})
Base.zero(::Type{StaticPolynomial}) = zero(StaticPolynomial{Int})

Base.one(p::StaticPolynomial{T, N}) where {T<:Number, N} = one(StaticPolynomial{T, N}, p.var)
Base.one(::Type{StaticPolynomial{T, N}}, var=:x) where {T, N} = zero(StaticPolynomial{T, N}, var) + 1
Base.one(::Type{StaticPolynomial{T}}) where T = one(StaticPolynomial{T, 1})
Base.one(::Type{StaticPolynomial}) = one(StaticPolynomial{Int})

function Base.:(*)(p1::StaticPolynomial{T, N}, p2::StaticPolynomial{S, M}) where {T<:Number, S<:Number, N, M}
    !(p1.var == p2.var || p1.var == :CONSTANT || p2.var == :CONSTANT) && error("Polynomials must be in terms of the same variable.")

    if M > 1 && N > 1
        degree = M + N - 2
        R = promote_type(T, S)
        cs = zeros(R, degree + 1)
        for i in 1:N
            for j in 1:M
                cs[i + j - 1] += p1[i] * p2[j]
            end
        end
        return StaticPolynomial(cs, p1.var)
    elseif N > 1
        return p1 * coeffs(p2)[1]
    else
        return p2 * coeffs(p1)[1]
    end
end

function Base.:(*)(p::StaticPolynomial{T}, n::S) where {T<:Number, S<:Number}
    if iszero(n)
        return StaticPolynomial(n)
    elseif isone(n)
        R = promote_type(T, S)
        return StaticPolynomial{R}(p)
    else
        return StaticPolynomial(n .* p.coeffs, p.var)
    end
end

Base.:(*)(n::S, p::StaticPolynomial) where {S<:Number} = p * n

Base.iszero(p::StaticPolynomial) = map(iszero, p.coeffs) |> all
Base.isone(p::StaticPolynomial) = isone(p.coeffs[1]) && map(iszero, p.coeffs[2:end]) |> all

# Differentiation and integration
function differentiate(p::StaticPolynomial, order = 1)
    if order == 0 || (iszero(p) && length(p) <= 1)
        return p
    elseif length(p) == 1 && !iszero(p)
        return zero(p)
    else
        coeffs = (i * p[i + 1] for i in 1:length(p)-1) |> collect
        var = length(coeffs) <= 1 ? :CONSTANT : p.var
        p = StaticPolynomial(coeffs, var)
        return differentiate(p, order - 1)
    end
end

function integrate(p::StaticPolynomial{T}; C = 0) where T<:AbstractFloat
    n = length(p)
    cs = zeros(T, n + 1)
    cs[1] = C
    cs[2:end] = [p[i] / i for i in 1:n]
    var = p.var == :CONSTANT ? :x : p.var
    return StaticPolynomial(cs, var)
end

function integrate(p::StaticPolynomial; C = 0)
    n = length(p)
    cs = zeros(Rational, n + 1)
    cs[1] = C
    cs[2:end] = [p[i] // i for i in 1:n]
    var = p.var == :CONSTANT ? :x : p.var
    return StaticPolynomial(cs, var)
end

#StaticPolynomial IO stuff
function exponent_string(exponent::Integer, mimetype=MIME"text/plain"())
    if exponent < 0
        return "⁻" * exponent_string(-exponent, mimetype)
    elseif exponent > 9
        return exponent_string(exponent ÷ 10, mimetype) * exponent_string(exponent % 10, mimetype)
    else
        if mimetype === MIME"text/latex"()
            return "^{" * string(exponent) * "}"
        else
            return exponent_dict[exponent]
        end
    end
end

const exponent_dict = Dict{Int, String}(
    0 => "⁰",
    1 => "¹",
    2 => "²",
    3 => "³",
    4 => "⁴",
    5 => "⁵",
    6 => "⁶",
    7 => "⁷",
    8 => "⁸",
    9 => "⁹",
)

poly_coeff_string(coeff, exponent) = isone(coeff) ? "" : (exponent == 0 ? string(coeff) : _poly_coeff_string(coeff))

_poly_coeff_string(coeff::(BigFloat | Float16 | Float32 | Float64 | Integer)) = string(coeff)
_poly_coeff_string(coeff) = string(coeff) |> wrap_string

wrap_string(str; left='(', right=')') = left * str * right

function poly_term_string(coeff, exponent, var, mimetype)
    if iszero(exponent)
        return isone(coeff) ? string(coeff) : poly_coeff_string(coeff, exponent)
    elseif isone(exponent)
        return poly_coeff_string(coeff, exponent) * string(var)
    else
        return poly_coeff_string(coeff, exponent) * string(var) * exponent_string(exponent, mimetype)
    end
end

signstring(n::Number) = signbit(n) ? "-" : "+"

function Base.show(io::IO, p::StaticPolynomial{T, N}) where {T<:Number, N}
    print(io, polyname(p), "(")
    print(io, printpoly(p))
    print(io, ")")
end

polyname(p) = string(typeof(p))

function printpoly(p::StaticPolynomial{T, N}, mimetype=MIME"text/plain"()) where {T<:Number, N}
    str = ""
    if iszero(p)
        str = string(zero(T))
    else
        if T <: Real
            for (i, c) in enumerate(coeffs(p))
                if c != 0
                    str *= isempty(str) ? poly_term_string(c, i - 1, p.var, mimetype) : " " * signstring(c) * " " * poly_term_string(abs(c), i-1, p.var, mimetype)
                end
            end
        else
            for (i, c) in enumerate(coeffs(p))
                if c != 0
                    str *= isempty(str) ? "" : " + "
                    str *= poly_term_string(c, i - 1, p.var, mimetype)
                end
            end
        end
    end
    return str
end

Base.string(p::StaticPolynomial) = printpoly(p)

Base.show(io::IO, mimetype::MIME"text/plain",  p::StaticPolynomial{T, N}) where {T<:Number, N} = begin
    Base.show(io, p)
end

function Base.show(io::IO, ::MIME"text/latex",  p::StaticPolynomial{T, N}) where {T, N}
    print(io, "\$")
    print(io, printpoly(p, MIME"text/latex"()))
    print(io, "\$")
end


"""
    @test_formany [ntests = 100] [maxtests = 1000] [logto = ""] [argument_data, proposition]

Takes proposition depending on some variables, generates sample of their values
and tests the proposition. Returns `Pass` if proposition holds for all values
in the sample, `Fail` if proposition does not hold for some value, and `Error`
if a part of code could not be evaluated.

# Arguments

Macro accepts a tuple of expressions that describe either statement of a test
or specify optional keyword arguments.

Statement of a test must be a sequence of comma-separated expressions, where
all but last expressions are read as `argument_data` and last statement is
`proposition` depending on these arguments.

* `argument_data` consists of arbitrarily many expressions containing
information about variables, their types and conditions on them. Expressions
must have one of the following forms:

  * `a<x::T<b` or `a>x::T>b`, where middle variable `x` must be a symbol;
  `T`  must be a type; bounds `a` & `b` can be arbitrary expressions;
  * `a<<x::T<b`, `a<x::T<<b`, `a>>x::T>b` or `a>x::T>>b`,
  similar to the one above;
  * shorthand `x` instead of `x::T` may be used, e.g. `a<x<b`,
  then type of variable defaults to `Float64`.

* `proposition`: is the expression to be tested. Two types of propositions are allowed:

  * `proposition` of form `condition-->statement`, where `condition` and
  `statement` are boolean-valued expressions
  * `proposition` is boolean-valued expression itself

Optional keyword arguments include the following:

* `ntests = 100`: number of tests that must be performed.

* `maxtests = 1000`: maximal number of attempts to generate arguments
satisfying `condition`

* `logto = "path_to_file"`: if provided then parsed as a path to file,
where log of tests will be recorded.

# Description
Macro @test_formany is a test that attempts to generate sample of argument values
of size `ntests` satisfying the `condition` if it is given. Then it checks whether
the proposition holds for all elements in the sample. If number of attempts
exceeds `maxtests` parameter before at least `ntests` values for arguments
are generated, macro gives output `Base.Test.Error`.

If `proposition` is of form `condition-->statement` (see arguments list),
then test checks whether implication holds. Namely, the test passes if whenever
`condition` evaluates to `true` for generated argument values, then `statement`
also evaluates to `true`.

If `proposition` is just a boolean-valued expression, code simply checks whether
`prop` holds for generated values.

Some caution is required to avoid vacuous conditions. If test is unable to
generate values satisfying the `condition` after `maxtests` attempts, output
of the test is Base.Test.Error as mentioned before.

Default for optional keyword arguments is the following: `ntests = 100, maxtests = 1000`.
If only `ntests` parameter is specified, `maxtests = 10*ntests`. If
`logto = "path_to_file"` is given, log of test will be written to corresponding
file,for example:

* `@test_formany ntests = 10 maxtests = 1000  logto =  "./test_log.csv" 1<x::Float32<10, proposition`
* `@test_formany ntests = 50 logto = "./test_log.csv" 1<x<10, proposition` ## maxtests = 500
* `@test_formany logto = "./test_log.csv" 1<x<<Inf, proposition` ## ntests = 100; maxtests = 1000
* `@test_formany ntests = 50 maxtests = 1000 1<x<10, proposition` ## not logged
* `@test_formany 1<x::Int32<Inf, proposition logto = "./test_log.csv" maxtests = 1000 ntests = 10` ## order is not important

# Examples
```julia
julia> @test_formany 100>x::Float64>10, x+10<y::Float64<<1000, y-5<z::Float64<Inf, z>x+5
Test Passed
  Expression: (:((100 > x::Float64 > 10,x + 10 < y::Float64 << 1000,y - 5 < z::Float64 < Inf,z > x + 5)),)

julia> @test_formany ntests = 1000 100>x::Float64>10, x+10<y::Float64<<1000, y-5<z::Float64<Inf, z>x+6
Test Failed
  Expression: (:(ntests = 1000),:((100 > x::Float64 > 10,x + 10 < y::Float64 << 1000,y - 5 < z::Float64 < Inf,z > x + 6)))

## example of bad design for generators:
julia> @test_formany -1000<x::Float64<1000, x>999-->x+1>1000
Error During Test
  Test threw an exception of type ErrorException
  Expression: (:((-1000 < x::Float64 < 1000,$(Expr(:-->, :(x > 999), :(x + 1 > 1000))))),)
  Found only 1/100 values satisfying given condition.

## log is increasing
julia> @test_formany Inf>x::Float64>0,Inf>y::Float64>0, x<y-->log(x)<log(y) ntests = 1000
Test Passed
  Expression: (:((Inf > x::Float64 > 0,Inf > y::Float64 > 0,$(Expr(:-->, :(x < y), :(log(x) < log(y)))))),:(ntests = 1000))

## f(x) = x^3 is not convex on [-100,100]
julia> @test_formany 0<a<1,-100<x<100,-100<y<100,(a*x+(1-a)*y)^3<a*x^3+(1-a)*y^3
Test Failed
  Expression: (:((0 < a < 1,-100 < x < 100,-100 < y < 100,(a * x + (1 - a) * y) ^ 3 < a * x ^ 3 + (1 - a) * y ^ 3)),)

## f(x) = x^3 is convex on [-100,100]
julia> @test_formany 0<a<1,0<x<100,0<y<100, (a*x+(1-a)*y)^3<=a*x^3+(1-a)*y^3
Test Passed
  Expression: (:((0 < a < 1,0 < x < 100,0 < y < 100,(a * x + (1 - a) * y) ^ 3 <= a * x ^ 3 + (1 - a) * y ^ 3)),)

## Cobb-Douglas utility function is concave
julia> function respects_concavity(f,t,x1,x2)
                  f(t*x1+(1-t)*x2) > t * f(x1) + (1-t) * f(x2)
              end

julia> function respects_weak_concavity(f,t,x1,x2)
                  f(t*x1+(1-t)*x2) >= t * f(x1) + (1-t) * f(x2)
              end

julia> function cobb_douglas(a,x)
                  prod(x.^a)
              end

# This passes the @test_formany, since probability of discovering a counter-example with a = .5 is zero.
julia> @test_formany 0<t<1, 0<a<1, 1<x1<10, 1<x2<10, 1<y1<10, 1<y2<10,respects_concavity(x -> cobb_douglas([a,1-a],x),t,[x1,y1],[x2,y2])
Test Passed
  Expression: (:((0 < t < 1,0 < a < 1,1 < x1 < 10,1 < x2 < 10,1 < y1 < 10,1 < y2 < 10,respects_concavity((x->begin  # REPL[5], line 1:
                    cobb_douglas([a,1 - a],x)
                end),t,[x1,y1],[x2,y2]))),)

julia> @test_formany 0<t<1, 0<a<1, 1<x1<10, 1<x2<10, 1<y1<10, 1<y2<10,respects_weak_concavity(x -> cobb_douglas([a,1-a],x),t,[x1,y1],[x2,y2])
Test Passed
  Expression: (:((0 < t < 1,0 < a < 1,1 < x1 < 10,1 < x2 < 10,1 < y1 < 10,1 < y2 < 10,respects_weak_concavity((x->begin  # REPL[6], line 1:
                    cobb_douglas([a,1 - a],x)
                end),t,[x1,y1],[x2,y2]))),)

## handling errors
julia> @test_formany -1<x<1, abs(log(x))>0
Error During Test
  Test threw an exception of type DomainError
  Expression: (:((-1 < x < 1,abs(log(x)) > 0)),)
  DomainError
```
"""
macro test_formany(exprs...)
   inner_ex = try
	maxtests = 0; ntests = 0
	prop = :(); cond = :()
	logging  = ""
	var_data = Array(Any,1,6)
	for ex in exprs
		if isa(ex, Expr) && ex.head == :tuple
			var_data = parse_argument_data(ex)
			if ex.args[end].head == :-->
				prop = ex.args[end].args[2]
				cond = ex.args[end].args[1]
			else
				prop = ex.args[end]
				cond = true
			end
		elseif isa(ex, Expr) && ex.head == :(=)
			if ex.args[1] == :ntests
				ntests = ex.args[2]
			elseif ex.args[1] == :maxtests
				maxtests = ex.args[2]
			elseif ex.args[1]== :logto
				logging = ex.args[2]
			else
				error("Invalid macro input $ex.")
			end
		else
			error("Invalid macro input $ex.")
		end
	end
	if ntests == 0
		ntests = 100
	end
	if maxtests == 0
		maxtests = 10*ntests
	end
	num_of_vars = size(var_data,1)
	generate_values = :()
	values = :([])
	for i in 1:num_of_vars
		next_expr = :($(esc(var_data[i,1]))=
			   custom_generator($(esc(var_data[i,2])),$(esc(var_data[i,3])),
			      $(esc(var_data[i,4])),$(esc(var_data[i,5])),$(esc(var_data[i,6])))(div(n,2)+3))
		generate_values = Expr(:block,generate_values.args...,next_expr)
		values = Expr(:vect, values.args..., :($(esc(var_data[i,1]))))
	end
#####
	if logging == ""    ## no logging
		inner_ex = quote try
			for n in 1:$maxtests
				$generate_values
				if $(esc(cond))
					num_good_args += 1
					res = res && $(esc(prop))
					if !res
						fail_data = $(esc(prop))
						break
					end
					if num_good_args >= $ntests
						break
					end
				end
			end
		        # no counter-example but insufficient amount of tests
			if res && num_good_args<$ntests
				nt = $ntests
 				error("Found only $num_good_args/$nt values satisfying given condition.")
			end
			res ? Pass(:test,$(Expr(:quote, exprs)), nothing, nothing) :
		   	   Fail(:test,$(Expr(:quote, exprs)), fail_data, nothing)
		catch err
			Error(:test_error,$(Expr(:quote, exprs)), err, catch_backtrace())
		end end #quote #try
	else 		## logging is a path to file where the log is written to
		inner_ex = quote try
			log_file = open($logging,"a")
			writedlm(log_file,reshape([string(s) for s in $(var_data[:,1])],(1,$num_of_vars)),",")
			for n in 1:$maxtests
				$generate_values
				if $(esc(cond))
					num_good_args += 1
					res = res && $(esc(prop))
					writedlm(log_file,transpose(push!(convert(Vector{Any},$values),res)),",")
					if !res
						fail_data = $(esc(prop))
						break
					end
					if num_good_args >= $ntests
						break
					end
				end
			end
			close(log_file)
		        # no counter-example but insufficient amount of tests
			if res && num_good_args<$ntests
				nt = $ntests
 				error("Found only $num_good_args/$nt values satisfying given condition.")
			end
			res ? Pass(:test,$(Expr(:quote, exprs)), nothing, nothing) :
		   	   Fail(:test,$(Expr(:quote, exprs)), fail_data, nothing)
		catch err
			Error(:test_error,$(Expr(:quote, exprs)), err, catch_backtrace())
		end end #quote #try
	end
	inner_ex
#####
   catch err
	inner_ex = quote
		Error(:test_error,$(Expr(:quote, exprs)), $err, catch_backtrace())
	end
   end # try defining inner_ex
	return quote
		num_good_args = 0
		res = true
		fail_data = false
		result = $inner_ex
		record(get_testset(), result)
	end
end


using Distributions
using Statistics
import LinearModel

function assert_same_rows(X, Y)
    if size(X,1) == size(Y,1)
        return true
    else
        throw(ArgumentError("X and Y must have same number of rows"))
    end
end

function assert_vector(Y)
    if size(Y,2) == 1
        return true
    else
        # more readable error message than struct type-checking error message
        throw(ArgumentError("Y must be a vector"))
    end
end

struct Sample
    X::Array
    Y::Vector
    Sample(X, Y) = (assert_same_rows(X, Y) & assert_vector(Y)
        ? new(X, Y)
        : throw(ArgumentError)
    )
end

# Create function that will return normal observations
function make_sample_generator_normal(; a, b, β_0, β, sd_e)
    k = size(β, 1)
    make_x_sample(n) = rand(Uniform(a, b), n, k)
    y_process(X) = β_0 .+ X * β
    noise_process(X) = rand(Normal(0, sd_e), size(X, 1), 1)
    function normal_sample(n)::Sample
        X = make_x_sample(n)
        Y = vec(y_process(X) + noise_process(X)) # vec reshapes Array{T, 2} to Vector
        return Sample(X, Y)
    end
    return normal_sample
end

function add_intercept(X::Array)
    nrows = size(X, 1)
    X0 = ones(nrows, 1)
    return [X0 X]
end

function add_intercept(sample::Sample)
    return Sample(add_intercept(sample.X), sample.Y)
end

normal_sampler = make_sample_generator_normal(a=0, b=20, β_0=5, β=[1,2], sd_e=1.1)
sam1 = normal_sampler(100)
lm1 = evaluate_ols(sam1, intercept=true)

println(show(lm1))
# TODO: add below this to show
println("R-squared: ", round(r2(lm1), digits=4))

# TODO: import data
# TODO: replicate https://itl.nist.gov/div898/strd/lls/data/LINKS/v-Norris.shtml

# the difference between two ols methods is really small
#@assert all(ols(obs1.X, obs1.Y)-ols_(obs1.X,obs1.Y) .< 1e-10) # was 14
#@assert add_intercept(collect(2:4)) == [1 1 1; 2 3 4]'
#@assert r2(lm1) > .8
using Test
@test_throws ArgumentError Sample([1, 2, 3], [1, 2])


# These algorithms are from
#   GARSTKA M., CANNON M., GOULART P.:
#   A clique graph based merging strategy for decomposable sdps.
#   IFAC-PapersOnLine 53, 2 (2020), 7355–7361.

import Base.length
length(cg::CliqueGraph) = size(cg.edge_mat, 1)

# Input:
#   cliques : vector of cliques
#   n       : num nodes of graph
# Output:
#   CliqueGraph : rep of relationship bt cliques
"""
    generate_clique_graph(cliques, n::Integer)

Generates CliqueGraph datastructure for an undirected graph from `cliques` and
the number of nodes `n`.

# Reference
* [A clique graph based merging strategy for decomposable SDPs](https://arxiv.org/pdf/1911.05615) by Michael Garstka, Mark Cannon, Paul Goulart
"""
function generate_clique_graph(cliques, n::Integer)
    m = length(cliques)
    membership_mat = spzeros(Bool, n, m)

    # TODO: edge mat as matrix or as list of tuples?
    edge_mat = zeros(Float64, m, m)
    active_cliques = Set(1:m)

    for (idx, clique) in enumerate(cliques)
        for node in clique
            membership_mat[node, idx] = true
        end
    end
    for i in 1:m, j in i+1:m
        ci = @view(membership_mat[:, i])
        cj = @view(membership_mat[:, j])
        l_int = sum(ci .& cj)
        if l_int > 0
            edge_mat[i,j] = edge_mat[j,i] = weight_function(ci, cj)
        end
    end

    return CliqueGraph(membership_mat, edge_mat, active_cliques)
end

# TODO: this should be a user-defined function
"""
    weight_function(ci, cj)

Defines the weight function used to determine if cliques `ci` and `cj` should be
merged. Cliques are merged if this is positive.

# Reference
* [A clique graph based merging strategy for decomposable SDPs](https://arxiv.org/pdf/1911.05615) by Michael Garstka, Mark Cannon, Paul Goulart
"""
function weight_function(ci, cj)
    l_ci = sum(ci)
    l_cj = sum(cj)
    l_union = sum(ci[k] | cj[k] for k in 1:length(ci))
    # l_union = 0.0
    # for k in 1:length(ci)
    #     # l_ci += ci[k]
    #     # l_cj += cj[k]
    #     l_union += ci[k] .| cj[k]
    # end
    return l_ci^3 + l_cj^3 - l_union^3
end

function _merge_cliques!(cg::CliqueGraph, i::Integer, j::Integer)
    # nodes X cliques
    n, m = size(cg.membership_mat)

    ci = @view(cg.membership_mat[:, i])
    cj = @view(cg.membership_mat[:, j])

    n_rem = sum(ci) < sum(cj) ? i : j
    n_keep = sum(ci) < sum(cj) ? j : i
    c_keep = @view(cg.membership_mat[:,n_keep])
    c_rem =  @view(cg.membership_mat[:,n_rem])

    # Update membership matrix
    c_keep .= ci .| cj
    c_rem .= false
    cg.edge_mat[n_keep,n_rem] = cg.edge_mat[n_rem,n_keep] = 0.0

    # TODO: change to only cliques with nonzero intersection
    for k in 1:m
        k == n_keep && continue
        # Remove edges from n_rem
        cg.edge_mat[k,n_rem] = cg.edge_mat[n_rem,k] = 0.0

        # Update edges to n_keep
        wij = weight_function(c_keep, @view(cg.membership_mat[:,k]))
        cg.edge_mat[k,n_keep] = cg.edge_mat[n_keep,k] = wij
    end

    pop!(cg.active_cliques, n_rem)

end

"""
    merge_cliques!(cg::CliqueGraph; verbose=false)

Merges cliques in `cg` in a greedy fashion starting with the pair with the largest
`weight_function`. Stops when all pairs of cliques have a non-positive weight.

# Reference
* [A clique graph based merging strategy for decomposable SDPs](https://arxiv.org/pdf/1911.05615) by Michael Garstka, Mark Cannon, Paul Goulart
"""
function merge_cliques!(cg::CliqueGraph; verbose=false)
    # nodes X cliques
    n, m = size(cg.membership_mat)
    if verbose
        max_clique_size = maximum(nnz, @view(cg.membership_mat[:,i]) for i in 1:m)
        @info "Starting merging; num cliques = $m. Max size is $max_clique_size."
    end

    max_val, max_ind = findmax(cg.edge_mat)
    while max_val > 0
        #TODO: verbose printing
        i = max_ind[1]
        j = max_ind[2]
        _merge_cliques!(cg, i, j)
        max_val, max_ind = findmax(cg.edge_mat)
    end
    dropzeros!(cg.membership_mat)
    if verbose
        max_clique_size = maximum(nnz, @view(cg.membership_mat[:,i]) for i in cg.active_cliques)
        @info "Finished merging; num cliques = $(length(cg.active_cliques)). Max size is $max_clique_size."
    end
end

"""
    get_cliques(cg::CliqueGraph)

Returns a list of the cliques in CliqueGraph `cg`.
"""
function get_cliques(cg::CliqueGraph)
    return [findnz(cg.membership_mat[:,i])[1] for i in sort(collect(cg.active_cliques))]
end


#=
Package adapted from https://github.com/luca-aki/TraceEstimation.jl
which has an MIT license as follows.

MIT License

Copyright (c) 2021 Akshay Jain and Contributors

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
=#

module TraceEstimators

include("common.jl")
include("hutchinson.jl")
include("diagapp.jl")
include("slq.jl")
include("chebyhutch.jl")
include("diagonalapprox.jl")

end


# You reach another volcanically active part of the cave. It would be nice if
# you could do some kind of thermal imaging so you could tell ahead of time
# which caves are too hot to safely enter.
#
# Fortunately, the submarine seems to be equipped with a thermal camera! When
# you activate it, you are greeted with:
#
# Congratulations on your purchase! To activate this infrared thermal imaging
# camera system, please enter the code found on page 1 of the manual.
#
# Apparently, the Elves have never used this feature. To your surprise, you
# manage to find the manual; as you go to open it, page 1 falls out. It's a
# large sheet of transparent paper! The transparent paper is marked with random
# dots and includes instructions on how to fold it up (your puzzle input). For
# example:
#
# 6,10
# 0,14
# 9,10
# 0,3
# 10,4
# 4,11
# 6,0
# 6,12
# 4,1
# 0,13
# 10,12
# 3,4
# 3,0
# 8,4
# 1,10
# 2,14
# 8,10
# 9,0
#
# fold along y=7
# fold along x=5
#
# The first section is a list of dots on the transparent paper. 0,0 represents
# the top-left coordinate. The first value, x, increases to the right. The
# second value, y, increases downward. So, the coordinate 3,0 is to the right
# of 0,0, and the coordinate 0,7 is below 0,0. The coordinates in this example
# form the following pattern, where # is a dot on the paper and . is an empty,
# unmarked position:
#
# ...#..#..#.
# ....#......
# ...........
# #..........
# ...#....#.#
# ...........
# ...........
# ...........
# ...........
# ...........
# .#....#.##.
# ....#......
# ......#...#
# #..........
# #.#........
#
# Then, there is a list of fold instructions. Each instruction indicates a line
# on the transparent paper and wants you to fold the paper up (for horizontal
# y=... lines) or left (for vertical x=... lines). In this example, the first
# fold instruction is fold along y=7, which designates the line formed by all
# of the positions where y is 7 (marked here with -):
#
# ...#..#..#.
# ....#......
# ...........
# #..........
# ...#....#.#
# ...........
# ...........
# -----------
# ...........
# ...........
# .#....#.##.
# ....#......
# ......#...#
# #..........
# #.#........
#
# Because this is a horizontal line, fold the bottom half up. Some of the dots
# might end up overlapping after the fold is complete, but dots will never
# appear exactly on a fold line. The result of doing this fold looks like this:
#
# #.##..#..#.
# #...#......
# ......#...#
# #...#......
# .#.#..#.###
# ...........
# ...........
#
# Now, only 17 dots are visible.
#
# Notice, for example, the two dots in the bottom left corner before the
# transparent paper is folded; after the fold is complete, those dots appear in
# the top left corner (at 0,0 and 0,1). Because the paper is transparent, the
# dot just below them in the result (at 0,3) remains visible, as it can be seen
# through the transparent paper.
#
# Also notice that some dots can end up overlapping; in this case, the dots
# merge together and become a single dot.
#
# The second fold instruction is fold along x=5, which indicates this line:
#
# #.##.|#..#.
# #...#|.....
# .....|#...#
# #...#|.....
# .#.#.|#.###
# .....|.....
# .....|.....
#
# Because this is a vertical line, fold left:
#
# #####
# #...#
# #...#
# #...#
# #####
# .....
# .....
#
# The instructions made a square!
#
# The transparent paper is pretty big, so for now, focus on just completing the
# first fold. After the first fold in the example above, 17 dots are visible -
# dots that end up overlapping after the fold is completed count as a single
# dot.
#
# How many dots are visible after completing just the first fold instruction on
# your transparent paper?

# read data
iter = eachline("input.txt")
values = map(l -> parse.(Int, split(l, ',')), Iterators.takewhile(l -> l != "", iter))
folds = map(iter) do l
  (f, i) = split(l, '=')
  (f[end], parse(Int, i))
end

# build sheet - I really hate 1-based indexing
width = maximum(first, values) + 1
height = maximum(last, values) + 1
sheet = falses(width, height)
foreach(((x,y),) -> sheet[x+1,y+1] = true, values)

# fold
fold_along = folds[1][2]
if folds[1][1] == 'x'
  sheet = sheet[1:fold_along,:] .| sheet[end:-1:(fold_along+2),:]
else
  sheet = sheet[:,1:fold_along] .| sheet[:,end:-1:(fold_along+2)]
end

println("Count: ", count(sheet))


### A Pluto.jl notebook ###
# v0.14.6

using Markdown
using InteractiveUtils

# ╔═╡ 0279facc-3f77-401b-ab39-352a6ac9bc9f
begin
    import Pkg
    Pkg.activate(mktempdir())
    Pkg.add([
        Pkg.PackageSpec(name="Plots", version="1"),
        Pkg.PackageSpec(name="PlutoUI", version="0.7"),
	    Pkg.PackageSpec(name="DifferentialEquations", version="6"),
	    Pkg.PackageSpec(name="Parameters", version="0.12"),
        Pkg.PackageSpec(name="LabelledArrays", version="1"),
        Pkg.PackageSpec(name="Setfield", version="0.7"),
    ])
    using Plots, PlutoUI, LinearAlgebra, DifferentialEquations, Parameters, LabelledArrays, Setfield

	Plots.gr(fmt=:png, lw=2)

	PlutoUI.TableOfContents()
end

# ╔═╡ f0ab21d0-bdcb-11eb-39e8-59f7ff34b88c
md"""
# Fig 1.09 Hodgkin-Huxley model
"""

# ╔═╡ 51592b89-56fc-40b3-be86-8ae48241a80b
# Parameters for Hodgkin-Huxley model
    p = (E_N = 55.0,       # Reversal potential of Na (mV)
		 E_K = -72.0,      # Reversal potential of K (mV)
		 E_LEAK = -49.0,   # Reversal potential of leaky channels (mV)
		 G_N_BAR = 120.0,  # Max. Na channel conductance (mS/cm^2)
		 G_K_BAR = 36.0,   # Max. K channel conductance (mS/cm^2)
		 G_LEAK = 0.30,    # Max. leak channel conductance (mS/cm^2)
		 C_M = 1.0,        # membrane capacitance (uF/cm^2)
		 I_STIM = 0.0     # Stimulus current (uA/cm^2)
    )

# ╔═╡ 8018875c-258e-4516-a36c-2413fee97ccd
begin
	# Events and callbacks
	on_i1!(integrator) = integrator.p = @set p.I_STIM = -6.80
    off_i1!(integrator) = integrator.p = @set p.I_STIM = 0.0
    on_i2!(integrator) = integrator.p = @set p.I_STIM = -6.90

    cbs = CallbackSet(PresetTimeCallback([20.0], on_i1!),
                  PresetTimeCallback([21.0, 61.0], off_i1!),
                  PresetTimeCallback([60.0], on_i2!))
end

# ╔═╡ 999aae28-fe75-4a2f-a55d-12360b647f05
u0 = LVector(v=-59.8977, m=0.0536, h=0.5925, n=0.3192)

# ╔═╡ 23ec33a8-b6f3-4a89-bc24-95474ed5a9cb
tend = 100.0

# ╔═╡ 8e7633bb-1d79-43cc-81ec-8fc740ba13ff
md"""
## Appendix
Running environment and some auxillary functions.

**You can use [this helper tool](https://fonsp.com/article-test-3/pkghelper.html) to generate these commands!**
"""

# ╔═╡ 94adbc56-8f68-4048-bfbc-326acb359fde
begin
	# Convenience functions
    hill(x, k) = x / (x + k)
    hill(x, k, n) = hill(x^n, k^n)
    exprel(x) = ifelse(x≈zero(x), one(x), x / expm1(x))
end

# ╔═╡ 31549c4b-2aec-4c09-9888-a15b2b0ebbd6
"ODE system of HH model"
	function hh!(du, u, p, t)
		@unpack E_N, E_K, E_LEAK, G_N_BAR, G_K_BAR, G_LEAK, C_M, I_STIM = p
		@unpack v, m, h, n = u

		mαV = -0.10 * (v + 35)
		mα = exprel(mαV)
		mβ = 4.0 * exp(-(v + 60) / 18.0)

		hα = 0.07 * exp(-(v+60)/20)
		hβ = 1 / (exp(-(v+30)/10) + 1)

		nαV = -0.1 * (v+50)
		nα = 0.1 * exprel(nαV)
		nβ = 0.125 * exp( -(v+60) / 80)

		iNa = G_N_BAR * (v - E_N) * (m^3) * h
		iK  = G_K_BAR * (v - E_K) * (n^4)
		iLeak = G_LEAK * (v - E_LEAK)
		du.v = -(iNa + iK + iLeak + I_STIM) / C_M
		du.m = -(mα + mβ) * m + mα
		du.h = -(hα + hβ) * h + hα
		du.n = -(nα + nβ) * n + nα
		return du
	end

# ╔═╡ 2e369b99-0f70-4ff2-be75-d3eec6c46777
sol = solve(ODEProblem(hh!, u0, tend, p), callback = cbs)

# ╔═╡ 2cec94d6-fcc8-49b6-970c-b84f8a736f61
plot(sol, vars=(0, 1), ylabel="Membrane potential (mV)", xlabel="Time (ms)", legend=false, title="Figure 1.9")

# ╔═╡ Cell order:
# ╠═f0ab21d0-bdcb-11eb-39e8-59f7ff34b88c
# ╠═31549c4b-2aec-4c09-9888-a15b2b0ebbd6
# ╠═8018875c-258e-4516-a36c-2413fee97ccd
# ╠═51592b89-56fc-40b3-be86-8ae48241a80b
# ╠═999aae28-fe75-4a2f-a55d-12360b647f05
# ╠═23ec33a8-b6f3-4a89-bc24-95474ed5a9cb
# ╠═2e369b99-0f70-4ff2-be75-d3eec6c46777
# ╠═2cec94d6-fcc8-49b6-970c-b84f8a736f61
# ╠═8e7633bb-1d79-43cc-81ec-8fc740ba13ff
# ╠═94adbc56-8f68-4048-bfbc-326acb359fde
# ╠═0279facc-3f77-401b-ab39-352a6ac9bc9f


# Bogumił Kamiński, 2019-03-25

using LinearAlgebra

function uppersum(M)
    n = size(M, 1)
    s = 0
    for i in 2:n, j in 1:i
        s += M[j, i]
    end
    s
end

@time m = randn(10_000, 10_000);
@time m = randn(10_000, 10_000);

m

@time sum(m)
@time sum(m)

@time sum(UpperTriangular(m))
@time sum(UpperTriangular(m))

@time uppersum(m)
@time uppersum(m)


macro test_elements(name, μ, ele, pos, vel)
    quote
        @testset $name begin
            μ = $μ
            a, e, i, Ω, ω, ν = $ele
            pos = $pos
            vel = $vel

            pos₁, vel₁ = cartesian(a, e, i, Ω, ω, ν, μ)
            a₁, e₁, i₁, Ω₁, ω₁, ν₁ = keplerian(pos, vel, μ)

            @testset "Cartesian" begin
                @testset for i = 1:3
                    @test pos₁[i] ≈ pos[i]
                    @test vel₁[i] ≈ vel[i]
                end
            end

            @testset "Keplerian" begin
                @test a ≈ a₁
                @test e ≈ e₁ atol=sqrt(eps())
                @test i ≈ i₁
                @test Ω ≈ Ω₁
                @test ω ≈ ω₁
                @test ν ≈ ν₁
            end
        end
    end
end

macro test_elements(name, μ, ele)
    quote
        @testset $name begin
            μ = $μ
            a, e, i, Ω, ω, ν = $ele

            pos₁, vel₁ = cartesian(a, e, i, Ω, ω, ν, μ)
            a₁, e₁, i₁, Ω₁, ω₁, ν₁ = keplerian(pos₁, vel₁, μ)

            @test a ≈ a₁
            @test e ≈ e₁ atol=sqrt(eps())
            @test i ≈ i₁
            @test Ω ≈ Ω₁
            @test ω ≈ ω₁
            @test ν ≈ ν₁
        end
    end
end

@testset "Two Body" begin
    @testset "Orbital Elements" begin
        @testset "Perifocal" begin
            p = 1.13880762905224e7
            ecc = 0.7311
            ν = 0.44369564302687126
            μ = 3.9860047e14

            pos = [6194863.12535486, 2944437.90016286, 0.0]
            vel = [-2539.71254827, 9668.69568539, 0.0]
            pos₁, vel₁ = perifocal(p, ecc, ν, μ)
            @testset for i = 1:3
                @test pos₁[i] ≈ pos[i]
                @test vel₁[i] ≈ vel[i]
            end
        end
        # Reference values from Orekit
        @test_elements("Elliptic",
                       3.9860047e14,
                       [24464560.0, 0.7311, 0.122138, 1.00681, 3.10686, 0.44369564302687126],
                       [-0.107622532467967e+07, -0.676589636432773e+07, -0.332308783350379e+06],
                       [0.935685775154103e+04, -0.331234775037644e+04, -0.118801577532701e+04])
        # Reference values from poliastro
        @test_elements("Circular",
                       3.986004418e14,
                       [6778136.6, 0.0, deg2rad(15), deg2rad(20), 0.0, deg2rad(30)],
                       [4396398.60746266, 5083838.45333733,  877155.42119322],
                       [-5797.06004014,  4716.60916063,  1718.86034246])
        @test_elements("Circular Orekit",
                       3.9860047e14,
                       # Orekit uses an argument of pericenter of 3.10686 here
                       [24464560.0, 0.0, 0.122138, 1.00681, 0.0, 0.048363])
        @test_elements("Hyperbolic Orekit",
                       3.9860047e14,
                       [-24464560.0, 1.7311, 0.122138, 1.00681,  3.10686, 0.12741601769795755])
        @test_elements("Equatorial",
                       3.9860047e14,
                       [24464560.0, 0.7311, 0.0, 0.0, 3.10686, 0.44369564302687126])
        @test_elements("Circular-Equatorial",
                       3.9860047e14,
                       [24464560.0, 0.0, 0.0, 0.0, 0.0, 0.44369564302687126])
    end
    @testset "Anomalies" begin
        anomalies = -π : 0.1π : π
        anomalies2π = 0.0 : 0.1π : 2π
        @testset "Elliptic" begin
            eccentricities = 0.1:0.1:0.9
            @testset "E <-> ν" begin
                @testset for E = anomalies, ecc = eccentricities
                    ν = transform(elliptic, eccentric_anomaly, true_anomaly, E, ecc)
                    E₁ = transform(elliptic, true_anomaly, eccentric_anomaly, ν, ecc)
                    @test E₁ ≈ E
                end
                @testset for E = anomalies2π, ecc = eccentricities
                    ν = transform(elliptic, eccentric_anomaly, true_anomaly, E, ecc)
                    E₁ = transform(elliptic, true_anomaly, eccentric_anomaly, ν, ecc)
                    @test E₁ ≈ E
                end
            end
            @testset "E <-> M" begin
                @testset for E = anomalies, ecc = eccentricities
                    M = transform(elliptic, eccentric_anomaly, mean_anomaly, E, ecc)
                    E₁ = transform(elliptic, mean_anomaly, eccentric_anomaly, M, ecc)
                    @test E₁ ≈ E
                end
                @testset for E = anomalies2π, ecc = eccentricities
                    M = transform(elliptic, eccentric_anomaly, mean_anomaly, E, ecc)
                    E₁ = transform(elliptic, mean_anomaly, eccentric_anomaly, M, ecc)
                    @test E₁ ≈ E
                end
            end
            @testset "M <-> ν" begin
                @testset for M = anomalies, ecc = eccentricities
                    ν = transform(elliptic, mean_anomaly, true_anomaly, M, ecc)
                    M₁ = transform(elliptic, true_anomaly, mean_anomaly, ν, ecc)
                    @test M₁ ≈ M
                end
                @testset for M = anomalies2π, ecc = eccentricities
                    ν = transform(elliptic, mean_anomaly, true_anomaly, M, ecc)
                    M₁ = transform(elliptic, true_anomaly, mean_anomaly, ν, ecc)
                    @test M₁ ≈ M
                end
            end
        end
        @testset "Hyperbolic" begin
            eccentricities = 1.1:0.1:10.0
            @testset "E <-> ν" begin
                @testset for E = anomalies, ecc = eccentricities
                    ν = transform(hyperbolic, eccentric_anomaly, true_anomaly, E, ecc)
                    E₁ = transform(hyperbolic, true_anomaly, eccentric_anomaly, ν, ecc)
                    @test E₁ ≈ E
                end
                @testset for E = anomalies2π, ecc = eccentricities
                    ν = transform(hyperbolic, eccentric_anomaly, true_anomaly, E, ecc)
                    E₁ = transform(hyperbolic, true_anomaly, eccentric_anomaly, ν, ecc)
                    @test E₁ ≈ E
                end
            end
            @testset "E <-> M" begin
                @testset for E = anomalies, ecc = eccentricities
                    M = transform(hyperbolic, eccentric_anomaly, mean_anomaly, E, ecc)
                    E₁ = transform(hyperbolic, mean_anomaly, eccentric_anomaly, M, ecc)
                    @test E₁ ≈ E
                end
                @testset for E = anomalies2π, ecc = eccentricities
                    M = transform(hyperbolic, eccentric_anomaly, mean_anomaly, E, ecc)
                    E₁ = transform(hyperbolic, mean_anomaly, eccentric_anomaly, M, ecc)
                    @test E₁ ≈ E
                end
            end
            @testset "M <-> ν" begin
                @testset for M = anomalies, ecc = eccentricities
                    ν = transform(hyperbolic, mean_anomaly, true_anomaly, M, ecc)
                    M₁ = transform(hyperbolic, true_anomaly, mean_anomaly, ν, ecc)
                    @test M₁ ≈ M
                end
                @testset for M = anomalies2π, ecc = eccentricities
                    ν = transform(hyperbolic, mean_anomaly, true_anomaly, M, ecc)
                    M₁ = transform(hyperbolic, true_anomaly, mean_anomaly, ν, ecc)
                    @test M₁ ≈ M
                end
            end
        end
        @testset "Parabolic" begin
            @testset "E <-> ν" begin
                @testset for E = anomalies
                    ν = transform(parabolic, eccentric_anomaly, true_anomaly, E, 1.0)
                    E₁ = transform(parabolic, true_anomaly, eccentric_anomaly, ν, 1.0)
                    @test E₁ ≈ E
                end
                @testset for E = anomalies2π
                    ν = transform(parabolic, eccentric_anomaly, true_anomaly, E, 1.0)
                    E₁ = transform(parabolic, true_anomaly, eccentric_anomaly, ν, 1.0)
                    @test E₁ ≈ E
                end
            end
            @testset "E <-> M" begin
                @testset for E = anomalies
                    M = transform(parabolic, eccentric_anomaly, mean_anomaly, E, 1.0)
                    E₁ = transform(parabolic, mean_anomaly, eccentric_anomaly, M, 1.0)
                    @test E₁ ≈ E
                end
                @testset for E = anomalies2π
                    M = transform(parabolic, eccentric_anomaly, mean_anomaly, E, 1.0)
                    E₁ = transform(parabolic, mean_anomaly, eccentric_anomaly, M, 1.0)
                    @test E₁ ≈ E
                end
            end
            @testset "M <-> ν" begin
                @testset for M = anomalies
                    ν = transform(parabolic, mean_anomaly, true_anomaly, M, 1.0)
                    M₁ = transform(parabolic, true_anomaly, mean_anomaly, ν, 1.0)
                    @test M₁ ≈ M
                end
                @testset for M = anomalies2π
                    ν = transform(parabolic, mean_anomaly, true_anomaly, M, 1.0)
                    M₁ = transform(parabolic, true_anomaly, mean_anomaly, ν, 1.0)
                    @test M₁ ≈ M
                end
            end
        end
    end
    @testset "Kepler" begin
        # Source: Vallado, Fundamentals of Astrodynamics and Applications, 4th edition, p. 94-95
        μ = 3.986004418e5
        r₀exp = [1131.340, -2282.343, 6672.423]
        v₀exp = [-5.64305, 4.30333, 2.42879]
        Δt = 40*60
        r₁exp = [-4219.7527, 4363.0292, -3958.7666]
        v₁exp = [3.689866, -1.916735, -6.112511]
        r₁, v₁ = kepler(μ, r₀exp, v₀exp, Δt)
        @test r₁ ≈ r₁exp rtol=1e-4
        @test v₁ ≈ v₁exp rtol=1e-5
        r₀, v₀ = kepler(μ, r₁, v₁, -Δt)
        @test r₀ ≈ r₀exp
        @test v₀ ≈ v₀exp
        r₁, v₁ = kepler(μ, r₀exp, v₀exp, eps())
        @test r₁ == r₀exp
        @test v₁ == v₀exp
        @test_throws ErrorException kepler(μ, r₀exp, v₀exp, Δt, 1)
    end
end


"""
$(TYPEDEF)

Represents a differential operator.

# Fields
$(FIELDS)

# Examples

```jldoctest
julia> using ModelingToolkit

julia> @variables x y;

julia> D = Differential(x)
(D'~x())

julia> D(y)  # Differentiate y wrt. x
(D'~x())(y())
```
"""
struct Differential <: Function
    """The variable or expression to differentiate with respect to."""
    x::Expression
end
(D::Differential)(x) = Operation(D, Expression[x])

Base.show(io::IO, D::Differential) = print(io, "(D'~", D.x, ")")
Base.convert(::Type{Expr}, D::Differential) = D

Base.:(==)(D1::Differential, D2::Differential) = isequal(D1.x, D2.x)

"""
$(SIGNATURES)

TODO
"""
function expand_derivatives(O::Operation,simplify=true)
    @. O.args = expand_derivatives(O.args,false)

    if isa(O.op, Differential)
        (D, o) = (O.op, O.args[1])

        isequal(o, D.x)     && return Constant(1)
        occursin(D.x, o)    || return Constant(0)
        isa(o, Operation)   || return O
        isa(o.op, Variable) && return O

        x = sum(1:length(o.args)) do i
            derivative(o, i) * expand_derivatives(D(o.args[i]),false)
        end

        return simplify ? ModelingToolkit.simplify(x) : x
    end

    return simplify ? ModelingToolkit.simplify(O) : O
end
expand_derivatives(x,args...) = x

# Don't specialize on the function here
"""
$(SIGNATURES)

Calculate the derivative of the op `O` with respect to its argument with index
`idx`.

# Examples

```jldoctest label1
julia> using ModelingToolkit

julia> @variables x y;

julia> ModelingToolkit.derivative(sin(x), 1)
cos(x())
```

Note that the function does not recurse into the operation's arguments, i.e., the
chain rule is not applied:

```jldoctest label1
julia> myop = sin(x) * y^2
sin(x()) * y() ^ 2

julia> typeof(myop.op)  # Op is multiplication function
typeof(*)

julia> ModelingToolkit.derivative(myop, 1)  # wrt. sin(x)
y() ^ 2

julia> ModelingToolkit.derivative(myop, 2)  # wrt. y^2
sin(x())
```
"""
derivative(O::Operation, idx) = derivative(O.op, (O.args...,), Val(idx))
derivative(O::Constant, ::Any) = Constant(0)

# Pre-defined derivatives
import DiffRules, SpecialFunctions, NaNMath
for (modu, fun, arity) ∈ DiffRules.diffrules()
    fun in [:*, :+] && continue # special
    for i ∈ 1:arity
        @eval function derivative(::typeof($modu.$fun), args::NTuple{$arity,Any}, ::Val{$i})
            M, f = $(modu, fun)
            partials = DiffRules.diffrule(M, f, args...)
            dx = @static $arity == 1 ? partials : partials[$i]
            convert(Expression, dx)
        end
    end
end

derivative(::typeof(+), args::NTuple{N,Any}, ::Val) where {N} = 1
derivative(::typeof(*), args::NTuple{N,Any}, ::Val{i}) where {N,i} = Operation(*, deleteat!(collect(args), i))
derivative(::typeof(one), args::Tuple{<:Any}, ::Val) = 0

function count_order(x)
    @assert !(x isa Symbol) "The variable $x must have an order of differentiation that is greater or equal to 1!"
    n = 1
    while !(x.args[1] isa Symbol)
        n = n+1
        x = x.args[1]
    end
    n, x.args[1]
end

_repeat_apply(f, n) = n == 1 ? f : f ∘ _repeat_apply(f, n-1)
function _differential_macro(x)
    ex = Expr(:block)
    lhss = Symbol[]
    x = x isa Tuple && first(x).head == :tuple ? first(x).args : x # tuple handling
    x = flatten_expr!(x)
    for di in x
        @assert di isa Expr && di.args[1] == :~ "@derivatives expects a form that looks like `@derivatives D''~t E'~t` or `@derivatives (D''~t), (E'~t)`"
        lhs = di.args[2]
        rhs = di.args[3]
        order, lhs = count_order(lhs)
        push!(lhss, lhs)
        expr = :($lhs = $_repeat_apply(Differential($rhs), $order))
        push!(ex.args,  expr)
    end
    push!(ex.args, Expr(:tuple, lhss...))
    ex
end

"""
$(SIGNATURES)

Define one or more differentials.

# Examples

```jldoctest
julia> using ModelingToolkit

julia> @variables x y z;

julia> @derivatives Dx'~x Dy'~y  # Create differentials wrt. x and y
((D'~x()), (D'~y()))

julia> Dx(z)  # Differentiate z wrt. x
(D'~x())(z())

julia> Dy(z)  # Differentiate z wrt. y
(D'~y())(z())
```
"""
macro derivatives(x...)
    esc(_differential_macro(x))
end

function calculate_jacobian(eqs, dvs)
    Expression[Differential(dv)(eq) for eq ∈ eqs, dv ∈ dvs]
end


include("../src/geom_module.jl")

module pointtest
using ..geom
using PyPlot
PyPlot.pygui(true)

function test_rot_mat()

    phi = 0.0
    rot_mat = geom.RotationMatrix(phi)

    phi = pi
    rot_mat = geom.RotationMatrix(phi)

    phi = pi/2
    rot_mat = geom.RotationMatrix(phi)

    println(rot_mat)

end

function test_point()

    p1 = geom.Point{Float64}(0, 1)
    p2 = geom.Point{Float64}(1, 0)
    p3 = geom.Point{Float64}(1, 1)
    p4 = geom.Point{Float64}(0, 0)

    println("at first index: $(p1[1])" )
    println("at second index: $(p1[2])" )

    println("at first index: $(p1[1] == p1.x)" )
    println("at first index: $(p1[2] == p1.y)" )

end

function test_point2()

    p1 = geom.Point{Float64}(2.0, 1)
    p2 = geom.Point{Float64}(1, -1.0)

    p_sum = p1 + p2
    p_neg = p1 - p2
    p_prod = p1 * 10

    println(" p sum = $(p_sum) ")
    println(" p sum = $(p_sum - geom.Point(p1.x + p2.x, p1.y + p2.y) ) ")
    println(" p prod = $(p_prod) ")
    println(" p_neg = $(p_neg) ")

    rot_mat = geom.RotationMatrix(pi/2)

    p_rot = geom.rotate(p1, rot_mat)
    println(" p_rot = $(p_rot) ")

end

function test_point3()

    p1 = geom.Point(10., 3.)

    p_shift = geom.shift(p1, 10., 5.)
    p_scale = geom.scale(p1, 3)

    println(" p_shift = $(p_shift) ")
    println(" p_scale = $(p_scale) ")

end

end # module


# include("./utils.jl")

abstract type Smoother end
mutable struct GaussianProcessSmoother <: Smoother
    smoother_function :: Function
    kernel
    mean_function
    GaussianProcessSmoother(x,y,z) = GaussianProcessSmoother(x,y,z)
    GaussianProcessSmoother() = new()
    function GaussianProcessSmoother(mean_function,kernel)
        gp = new()
        gp.smoother_function = GP
        gp.mean_function = mean_function
        gp.kernel = kernel
        return gp
    end
end

function fit(smoother::GaussianProcessSmoother,x,y)
    model = smoother.smoother_function(
                        x,
                        y,
                        smoother.mean_function,
                        smoother.kernel
                        )
    return model
end

function predict(model::GPBase,smoothed_t)
    smoothed_x,_ = predict_y(model, smoothed_t)
    return smoothed_x
end

function smoothing(x::AbstractVector,smoother::Smoother;step=0.01)
    t = 1.0:length(x)
    standardized_x = standardize(x)
    model = fit(smoother,t,standardized_x)
    lower_bound, upper_bound = minimum(t), maximum(t)
    smoothed_t = range(lower_bound, stop = upper_bound, step = step)
    n_points = length(smoothed_t)
    smoothed_x = predict(model, smoothed_t)
    smoothed_x = reverse_standardization(smoothed_x,x)
    return smoothed_x
end
function smoothing(x::AbstractVector;step=0.01)
    smoother =  GaussianProcessSmoother(MeanZero(),SE(0.0,0.0))
    smoothed_x = smoothing(x,smoother,step=step)
    return smoothed_x
end

function smoothing(x::Matrix{T},smoothers::Vector{Smoother};step=0.01) where {T<:Number}
    n_row, n_col = size(x,1),size(x,2)
    @assert n_col == length(smoothers)
    t = 1.0:n_row
    lower_bound, upper_bound = minimum(t), maximum(t)
    smoothed_t = range(lower_bound, stop = upper_bound, step = step)
    n_points = length(smoothed_t)
    new_x = zeros(eltype(x),n_points,n_col)

    for i in 1:n_col
        smoothed_x = smoothing(x[:,i],smoothers[i],step=step)
        new_x[:, i] = smoothed_x
    end
    return new_x
end
function smoothing(x::Matrix{T},smoother::Smoother;step=0.01) where {T<:Number}
    smoother_x = smoother
    smoother_y = deepcopy(smoother)
    smoother_z = deepcopy(smoother)
    smoothers = [smoother_x,smoother_y,smoother_z]
    return smoothing(x,smoothers;step=step)
end
function smoothing(x::Matrix{T};step=0.01) where {T<:Number}
    n_row, n_col = size(x,1),size(x,2)
    t = 1.0:n_row
    lower_bound, upper_bound = minimum(t), maximum(t)
    smoothed_t = range(lower_bound, stop = upper_bound, step = step)
    n_points = length(smoothed_t)
    new_x = zeros(eltype(x),n_points,n_col)
    for i in 1:n_col
        smoothed_x = smoothing(x[:,i],step=step)
        new_x[:, i] = smoothed_x
    end
    return new_x
end


using Test
using jInvVisPyPlot
using jInv.Mesh
using PyPlot

# tests for regular mesh
domain = [1 2 0 .1 3 4]
n      = [81 70 30]
M2D    = getRegularMesh(domain[1:4],n[1:2])
m2D    = ones(tuple(M2D.n...))
M3D    = getRegularMesh(domain,n)
m3D    = ones(tuple(M3D.n...))

println("=== test plotModel (regular mesh) ===")
figure()
plotModel(m2D,true,M2D,2,[0.8,1.2]);

figure()
plotModel(m3D,true,M3D,2,[0.8,1.2]);

println("==== all passed === ")


using RobustAndOptimalControl
G1 = ssrand(2,3,4, proper=true)
G2 = ssrand(2,3,4, proper=true)

@test h2norm(G1-G1) < 1e-6
@test hinfnorm2(G1-G1)[1] < 1e-6
@test nugap(G1, G1)[1] < 1e-6
@test hankelnorm(G1-G1)[1] < 1e-6
@test hinfnorm2(baltrunc2(G1, n=4)[1]-G1)[1] < 1e-10

@test count(1:100) do _
    G1 = ssrand(1,1,3, proper=true)
    G2 = ssrand(1,1,3, proper=true)
    try
        νgap(G1\G2, ss(inv(tf(G1))*tf(G2)))[1] < 1e-8
    catch
        false
    end
end >= 95


### Model
m = SmetsWoutersOrig()
TTT, RRR, CCC = solve(m)
meas = measurement(m, TTT, RRR, CCC)
sys = compute_system(m)

nothing


using Combinatorics: with_replacement_combinations
using StatsBase: countmap

function compute(n::Int, m::Int)::Int
    factorials = map(factorial, 0:9)
    chains = Dict{Int, Int}()
    result = 0
    for combination ∈ with_replacement_combinations(9:-1:0, n)
        for _ ∈ 1:count(isequal(0), combination) + 1
            number = parse(Int, join(combination))
            chain = [number]
            term = sum(factorials[parse.(Int, collect(string(number))) .+ 1])
            while term ∉ chain
                push!(chain, term)
                if haskey(chains, term) break end
                term = sum(factorials[parse.(Int, collect(string(term))) .+ 1])
            end
            index = findfirst(isequal(term), chain)
            loop_len = haskey(chains, term) ? chains[term] : length(chain) - index + 1
            for (i, j) ∈ enumerate(chain) chains[j] = i ≥ index ? loop_len : index - i + loop_len end
            if chains[number] == m
                digits = length(combination)
                result += (digits - count(isequal(0), combination)) * factorials[digits] ÷ prod(factorials[values(countmap(combination)) .+ 1])
            end
            if 0 ∈ combination
                deleteat!(combination, findfirst(isequal(0), combination))
                if length(combination) == 0 break end
            end
        end
    end
    return result
end


function explicit_crank_nicolson(mu, beta, p0, h, t0, t1, a0, a1)
    a = a0:h:a1
    t = t0:h:t1
    m = length(a)
    n = length(t)

    u = zeros(h, m, n)
    u[:, 1] = p0(a, t0)

    for i in 1:(n - 1)
        h2 = 0.5 * (t[i] + t[i + 1])
        u[:, i + 1] = (1 .- 0.5 * h * mu(a, h2)) ./ (1 .+ 0.5 * h * mu(a, h2)) * u[:, i]
        u[1, i + 1] = h * sum(beta(a, t[i + 1]) .* u[1:m, i + 1]) + 0.5 * h * u[1, i]
    end
    u
end


# using SymPy

# x, y = @vars x y
# u = x*(1-x)*y*(1-y)
# v = x*(1-x)*y*(1-y)^2

# p = x*(1-x)*y*(1-y)

# ux = diff(u,x)
# uy = diff(u,y)
# vx = diff(v,x)
# vy = diff(v,y)
# px = diff(p,x)
# py = diff(p,y)
# f = u*ux + v*uy - (diff(ux,x)+diff(uy,y)) + px
# g = u*vx + v*vy - (diff(vx,x)+diff(vy,y)) + py
# h = diff(u, x) + diff(v, y)
# println(replace(replace(sympy.julia_code(simplify(f)), ".^"=>"^"), ".*"=>"*"))
# println(replace(replace(sympy.julia_code(simplify(g)), ".^"=>"^"), ".*"=>"*"))
# println(replace(replace(sympy.julia_code(simplify(h)), ".^"=>"^"), ".*"=>"*"))

function u_exact(x,y)
    x*(1-x)*y*(1-y)
end

function v_exact(x,y)
    x*(1-x)*y*(1-y)
end

function p_exact(x,y)
    x*(1-x)*y*(1-y)
end

function t_exact(x,y)
    x*(1-x)*y*(1-y)
end

function k_exact(x,y)
    1 + x^2 + x / (1 + y^2)
end

function ffunc_(x, y)
    x*y*(1 - x)*(1 - y)*(-x*y*(1 - x) + x*(1 - x)*(1 - y)) + x*y*(1 - x)*(1 - y)*(-x*y*(1 - y) + y*(1 - x)*(1 - y)) - x*y*(1 - y) + 0.02*x*(1 - x) + y*(1 - x)*(1 - y) + 0.02*y*(1 - y)
end

function gfunc_(x, y)
    x*y*(1 - x)*(1 - y)*(-x*y*(1 - x) + x*(1 - x)*(1 - y)) + x*y*(1 - x)*(1 - y)*(-x*y*(1 - y) + y*(1 - x)*(1 - y)) - x*y*(1 - x) + x*(1 - x)*(1 - y) + 0.02*x*(1 - x) + 0.02*y*(1 - y)
end

function hfunc_(x,y)
    -x*y*(1 - x) - x*y*(1 - y) + x*(1 - x)*(1 - y) + y*(1 - x)*(1 - y)
end

function heat_source_func(x, y)
    # x^2*y*(x - 1)^2*(y - 1)*(2*y - 1) + x*y^2*(x - 1)*(2*x - 1)*(y - 1)^2 - 2.0*x*(x - 1) - 2.0*y*(y - 1) # for k=1.0
    (2*x^2*y*(x - 1)*(2*y - 1) + x*y*(x - 1)*(y - 1)*(y^2 + 1)^2*(x*(x - 1)*(2*y - 1) + y*(2*x - 1)*(y - 1)) - (y^2 + 1)*(2*x*(x - 1)*(x + (x^2 + 1)*(y^2 + 1)) + 2*y*(x + (x^2 + 1)*(y^2 + 1))*(y - 1) + y*(2*x - 1)*(y - 1)*(2*x*(y^2 + 1) + 1)))/(y^2 + 1)^2
end

using LinearAlgebra
using MAT
using AdFem
using PyPlot
using SparseArrays

m = 20
n = 20
h = 1/n
nu = 0.01

F1 = compute_fem_source_term1(eval_f_on_gauss_pts(ffunc_, m, n, h), m, n, h)
F2 = compute_fem_source_term1(eval_f_on_gauss_pts(gfunc_, m, n, h), m, n, h)
H = h^2*eval_f_on_fvm_pts(hfunc_, m, n, h)
B = constant(compute_interaction_matrix(m, n, h))

# compute F
Laplace = nu * constant(compute_fem_laplace_matrix1(m, n, h))
heat_source = eval_f_on_gauss_pts(heat_source_func, m, n, h)
heat_source = constant(compute_fem_source_term1(heat_source, m, n, h))

kgauss = eval_f_on_gauss_pts(k_exact, m, n, h)
LaplaceK = constant(compute_fem_laplace_matrix1(kgauss, m, n, h))

# The temperature term is
# T_xx + T_yy + q = u T_x + v T_y
#
function compute_residual(S)
    u, v, p, T = S[1:(m+1)*(n+1)],
        S[(m+1)*(n+1)+1:2(m+1)*(n+1)],
        S[2(m+1)*(n+1)+1:2(m+1)*(n+1)+m*n],
        S[2(m+1)*(n+1)+m*n+1:end]
    G = eval_grad_on_gauss_pts([u;v], m, n, h)
    ugauss = fem_to_gauss_points(u, m, n, h)
    vgauss = fem_to_gauss_points(v, m, n, h)
    ux, uy, vx, vy = G[:,1,1], G[:,1,2], G[:,2,1], G[:,2,2]

    interaction = compute_interaction_term(p, m, n, h) # julia kernel needed
    f1 = compute_fem_source_term1(ugauss.*ux, m, n, h)
    f2 = compute_fem_source_term1(vgauss.*uy, m, n, h)
    f3 = -interaction[1:(m+1)*(n+1)]
    f4 = Laplace*u
    f5 = -F1
    F = f1 + f2 + f3 + f4 + f5

    g1 = compute_fem_source_term1(ugauss.*vx, m, n, h)
    g2 = compute_fem_source_term1(vgauss.*vy, m, n, h)
    g3 = -interaction[(m+1)*(n+1)+1:end]
    g4 = Laplace*v
    g5 = -F2
    G = g1 + g2 + g3 + g4 + g5

    H0 = -B * [u;v] + H

    # T_xx + T_yy = u T_x + v T_y - heat_source
    T0 = LaplaceK * T + compute_fem_advection_matrix1(ugauss,vgauss, m, n, h) * T - heat_source

    R = [F;G;H0;T0]
    return R
end

function compute_jacobian(S)
    u, v, p, T = S[1:(m+1)*(n+1)],
        S[(m+1)*(n+1)+1:2(m+1)*(n+1)],
        S[2(m+1)*(n+1)+1:2(m+1)*(n+1)+m*n],
        S[2(m+1)*(n+1)+m*n+1:end]

    G = eval_grad_on_gauss_pts([u;v], m, n, h)
    ugauss = fem_to_gauss_points(u, m, n, h)
    vgauss = fem_to_gauss_points(v, m, n, h)
    ux, uy, vx, vy = G[:,1,1], G[:,1,2], G[:,2,1], G[:,2,2]

    M1 = constant(compute_fem_mass_matrix1(ux, m, n, h))
    M2 = constant(compute_fem_advection_matrix1(constant(ugauss), constant(vgauss), m, n, h)) # a julia kernel needed
    M3 = Laplace
    Fu = M1 + M2 + M3

    Fv = constant(compute_fem_mass_matrix1(uy, m, n, h))

    N1 = constant(compute_fem_mass_matrix1(vy, m, n, h))
    N2 = constant(compute_fem_advection_matrix1(constant(ugauss), constant(vgauss), m, n, h))
    N3 = Laplace
    Gv = N1 + N2 + N3

    Gu = constant(compute_fem_mass_matrix1(vx, m, n, h))

    # M = constant(compute_fem_laplace_matrix1(m, n, h)) +
        #  compute_fem_advection_matrix1(ugauss,vgauss, m, n, h)

    M = LaplaceK + constant(compute_fem_advection_matrix1(ugauss,vgauss, m, n, h))

    gradT = eval_grad_on_gauss_pts1(T, m, n, h)
    Tx, Ty = gradT[:,1], gradT[:,2]
    DU_TX = constant(compute_fem_mass_matrix1(Tx, m, n, h))       # (m+1)*(n+1), (m+1)*(n+1)
    DV_TY = constant(compute_fem_mass_matrix1(Ty, m, n, h))       # (m+1)*(n+1), (m+1)*(n+1)

    J0 = [Fu Fv                         # Jacobian of [u; v]
          Gu Gv]

    J1 = [J0 -B'                        # Jacobian of [u; v; p]
          -B spdiag(zeros(size(B,1)))]

    N = 2*(m+1)*(n+1) + m*n                                     # length of [u; v; p]
    J = [J1 SparseTensor(spzeros(N,(m+1)*(n+1)))                # Jacobian of [u; v; p; T]
        DU_TX DV_TY SparseTensor(spzeros((m+1)*(n+1), m*n)) M]
end

NT = 5    # number of iterations for Newton's method, 10 => 5

bd = bcnode("all", m, n, h)
bd = [bd; bd .+ (m+1)*(n+1);
    2*(m+1)*(n+1)+(n-1)*m+1:2*(m+1)*(n+1)+(n-1)*m+2;
    (2*(m+1)*(n+1)+m*n) .+ bd]
# only apply Dirichlet to velocity; set left bottom two points to zero to fix rank deficient problem for pressure

function solve_steady_cavityflow_one_step(S)
    residual = compute_residual(S)
    J = compute_jacobian(S)

    J, _ = fem_impose_Dirichlet_boundary_condition1(J, bd, m, n, h)
    residual = scatter_update(residual, bd, zeros(length(bd)))    # residual[bd] .= 0.0 in Tensorflow syntax

    d = J\residual
    residual_norm = norm(residual)
    op = tf.print("residual norm", residual_norm)
    d = bind(d, op)
    S_new = S - d
    return S_new
end

function condition(i, S_arr)
    i <= NT + 1
end

function body(i, S_arr)
    S = read(S_arr, i-1)
    op = tf.print("i=",i)
    i = bind(i, op)
    S_new = solve_steady_cavityflow_one_step(S)
    S_arr = write(S_arr, i, S_new)
    return i+1, S_arr
end

xy = fem_nodes(m, n, h)
x, y = xy[:,1], xy[:,2]
u0 = @. u_exact(x,y)
v0 = @. v_exact(x,y)
t0 = @. t_exact(x,y)

xy = fvm_nodes(m, n, h)
x, y = xy[:,1], xy[:,2]
p0 = @. p_exact(x,y)

S_arr = TensorArray(NT+1)
S_arr = write(S_arr, 1, zeros(m*n+3*(m+1)*(n+1)))

i = constant(2, dtype=Int32)

_, S = while_loop(condition, body, [i, S_arr])
S = set_shape(stack(S), (NT+1, 2*(m+1)*(n+1)+m*n+(m+1)*(n+1)))

sess = Session(); init(sess)
output = run(sess, S)

matwrite("CHTcoupled_data.mat",
    Dict(
        "V"=>output[end, :]
    ))

figure(figsize=(25,10))
subplot(241)
title("u velocity")
visualize_scalar_on_fem_points(output[NT+1, 1:(m+1)*(n+1)], m, n, h)
subplot(245)
visualize_scalar_on_fem_points(u0, m, n, h)

subplot(242)
title("v velocity")
visualize_scalar_on_fem_points(output[NT+1, (m+1)*(n+1)+1:2*(m+1)*(n+1)], m, n, h)
subplot(246)
visualize_scalar_on_fem_points(v0, m, n, h)

subplot(243)
visualize_scalar_on_fvm_points(output[NT+1, 2*(m+1)*(n+1)+1:2*(m+1)*(n+1)+m*n], m, n, h)
title("pressure")
subplot(247)
visualize_scalar_on_fvm_points(p0, m, n, h)
title("")

subplot(244)
title("temperature")
visualize_scalar_on_fem_points(output[NT+1, 2*(m+1)*(n+1)+m*n+1:end], m, n, h)
subplot(248)
visualize_scalar_on_fem_points(t0, m, n, h)

tight_layout()

# final_u=output[NT+1, 1:(1+m)*(1+n)]
# final_v=output[NT+1, (1+m)*(1+n)+1:2*(m+1)*(n+1)]
# final_p=output[NT+1, 2*(m+1)*(n+1)+1:end]

# u1 = final_u[50*101+1: 50*101+101]
# u2 = final_u[51:101:end]

# v1 = final_v[50*101+1: 50*101+101]
# v2 = final_v[51:101:end]
# xx = 0:0.01:1

# figure();plot(xx, u1);
# savefig("u_horizontal.png")

# figure();plot(xx, u2);
# savefig("u_vertical.png")

# figure();plot(xx, v1);
# savefig("v_horizontal.png")

# figure();plot(xx, v2);
# savefig("v_vertical.png")

# p1 = final_p[49*100+1: 49*100+100]
# p2 = final_p[50*100+1: 50*100+100]
# p3 = 0.5 * (p1 .+ p2)

# p4 = final_p[50:100:end]
# p5 = final_p[51:100:end]
# p6 = 0.5 * (p4 .+ p5)

# xx = 0.005:0.01:1

# figure();plot(xx, p3);
# savefig("p_horizontal.png")

# figure();plot(xx, p6);
# savefig("p_vertical.png")


# Copyright (c) 2017: Miles Lubin and contributors
# Copyright (c) 2017: Google Inc.
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

"""
    RSOCtoSOCBridge{T,F,G} <: Bridges.Constraint.AbstractBridge

`RSOCtoSOCBridge` implements the following reformulation:

  * ``||x||_2^2 \\le 2tu`` into ``||v|| \\le w ``, with the substitution rules
    ``t = \\frac{w}{\\sqrt 2} + \\frac{v_1}{\\sqrt 2}``,
    ``u = \\frac{w}{\\sqrt 2} - \\frac{v_1}{\\sqrt 2}``, and
    ``x = (v_2,\\ldots,v_N)``.

## Source node

`RSOCtoSOCBridge` supports:

  * `G` in [`MOI.RotatedSecondOrderCone`](@ref)

## Target node

`RSOCtoSOCBridge` creates:

  * `F` in [`MOI.SecondOrderCone`](@ref)
"""
struct RSOCtoSOCBridge{T,F,G} <:
       SetMapBridge{T,MOI.SecondOrderCone,MOI.RotatedSecondOrderCone,F,G}
    constraint::MOI.ConstraintIndex{F,MOI.SecondOrderCone}
end

const RSOC{T,OT<:MOI.ModelLike} = SingleBridgeOptimizer{RSOCtoSOCBridge{T},OT}

function concrete_bridge_type(
    ::Type{<:RSOCtoSOCBridge{T}},
    G::Type{<:MOI.AbstractVectorFunction},
    ::Type{MOI.RotatedSecondOrderCone},
) where {T}
    S = MOI.Utilities.promote_operation(/, T, MOI.Utilities.scalar_type(G), T)
    Y = MOI.Utilities.promote_operation(-, T, S, S)
    Z = MOI.Utilities.promote_operation(+, T, S, S)
    F = MOI.Utilities.promote_operation(vcat, T, Z, Y, G)
    return RSOCtoSOCBridge{T,F,G}
end

function MOI.Bridges.map_set(
    ::Type{<:RSOCtoSOCBridge},
    set::MOI.RotatedSecondOrderCone,
)
    return MOI.SecondOrderCone(MOI.dimension(set))
end

function MOI.Bridges.inverse_map_set(
    ::Type{<:RSOCtoSOCBridge},
    set::MOI.SecondOrderCone,
)
    return MOI.RotatedSecondOrderCone(MOI.dimension(set))
end

function MOI.Bridges.map_function(::Type{<:RSOCtoSOCBridge{T}}, func) where {T}
    scalars = MOI.Utilities.eachscalar(func)
    t, u, x = scalars[1], scalars[2], scalars[3:end]
    ts = MOI.Utilities.operate!(/, T, t, sqrt(T(2)))
    us = MOI.Utilities.operate!(/, T, u, sqrt(T(2)))
    return MOI.Utilities.operate(vcat, T, ts + us, ts - us, x)
end

# The map is an involution
function MOI.Bridges.inverse_map_function(BT::Type{<:RSOCtoSOCBridge}, func)
    return MOI.Bridges.map_function(BT, func)
end

# The map is symmetric
function MOI.Bridges.adjoint_map_function(BT::Type{<:RSOCtoSOCBridge}, func)
    return MOI.Bridges.map_function(BT, func)
end

# The map is a symmetric involution
function MOI.Bridges.inverse_adjoint_map_function(
    BT::Type{<:RSOCtoSOCBridge},
    func,
)
    return MOI.Bridges.map_function(BT, func)
end


using StaticArrays
using Rotations
using LinearAlgebra
using Test

import Rotations: RotationError, rotation_error, add_error, params

@testset "Rotation Error" begin
q1 = rand(UnitQuaternion)
dq = expm(0.1 * normalize(@SVector rand(3)))
q2 = q1*dq
@test q1\q2 ≈ dq
dg = RodriguesParam(dq)

emap = CayleyMap()
e1 = rotation_error(q2, q1, emap)
@test e1 ≈ q2 ⊖ q1
@test e1 ≈ params(dg)
@test add_error(q1, e1) ≈ q2
@test q1 ⊕ e1 isa UnitQuaternion
@test rotation_angle(UnitQuaternion(e1)) ≈ 0.1

aa1 = AngleAxis(q1)
@test add_error(aa1, e1) ≈ q2

e1 = rotation_error(q2, aa1, emap)
@test add_error(aa1, e1) ≈ AngleAxis(q2)
@test aa1 ⊕ e1 ≈ q2

# Test IdentityMap
g1 = rand(RodriguesParam)
g2 = RodriguesParam(g1*dg)
@test rotation_angle(g1\g2) ≈ 0.1

e1 = rotation_error(g2, g1, IdentityMap())
@test e1 ≈ g2 ⊖ g1
@test e1 ≈ params(dg)

p1 = rand(MRP)
dp = MRP(dg)
p2 = p1*dp
@test rotation_angle(p1\p2) ≈ 0.1

e1 = rotation_error(p2, p1, IdentityMap())
@test e1 ≈ p2 ⊖ p1
@test e1 ≈ params(dp)

@test_throws ArgumentError rotation_error(q2, q1, IdentityMap())
@test_throws ArgumentError q1 ⊕ e1

# Test inv
emap = CayleyMap()
imap = Rotations.InvCayleyMap()
g = inv(emap)(q1)
@test g ≈ imap(q1)
@test emap(g) ≈ q1
@test inv(imap)(g) ≈ q1
@test inv(inv(emap))(g) ≈ q1
@test inv(inv(imap))(q1) ≈ g

g = RodriguesParam(g)
@test UnitQuaternion(g) ≈ q1
@test imap(g) ≈ params(g)
end


using Turing
using Test
using Random
Random.seed!(0)

@model gdemo() = begin
  s ~ InverseGamma(2,3)
  m ~ Normal(0,sqrt(s))
  1.5 ~ Normal(m, sqrt(s))
  2.0 ~ Normal(m, sqrt(s))
  return s, m
end

GaussianKernel(var) = (x) -> Normal(x, sqrt(var))
N = 2000
s1 = MH(N, (:s, GaussianKernel(3.0)), (:m, GaussianKernel(3.0)))
s2 = MH(N, :s, :m)
s3 = MH(N)
s4 = Gibbs(N, MH(5, :m), MH(5, :s))

c1 = sample(gdemo(), s1)
c2 = sample(gdemo(), s2)
c3 = sample(gdemo(), s3)
c4 = sample(gdemo(), s4)

# Very loose bound, only for testing constructor.
for c in [c1, c2, c3, c4]
  @test mean(c[:s]) ≈ 49/24 atol=1.0
  @test mean(c[:m]) ≈ 7/6 atol=1.0
end


"""
simulation & machine learning of diffusion advection PDE
"""

using IterTools
using Plots
using Random
using Flux
using LinearAlgebra
Random.seed!(1)
using Main.EquivariantOperators # end users should omit Main.

# make grid
n = 2
dx = 0.02
cell = dx * Matrix(I, n, n)
rmax = 1.0
grid = Grid(cell, rmax)
sz = size(grid)
▽ = Op(:▽, cell)

# put small blob as IC
u0 = zeros(sz..., 1)
put!(u0, grid, [0.0, 0.0], [1.0])
g = Op(:Gaussian, cell; σ = dx)
u0 = g(u0)

# diffusion advection
p1 = [0.1]
p2 = [0.5, 0.4]
vf = cat(p2[1] * ones(sz), p2[2] * ones(sz), dims = 3)
f(u, p, t) = p1[1] * (▽ ⋅ ▽(u)) - vf ⊗ ▽(u)

# simulate PDE
using DifferentialEquations
tspan = (0.0, 1.0)
prob = ODEProblem(f, u0, tspan)
sol = solve(prob, Tsit5(), reltol = 1e-3, abstol = 1e-3)

# plot
using Plots
gr()
anim = Animation()

t = 0:0.02:1
for t in t
heatmap(sol(t)[:, :, 1], clim=(0,10))
frame(anim)
end
gif(anim, "f.gif", fps = 10)

##
data = [(sol(t), f(sol(t), 0, 0)) for t in t]
# op = Op(Radfunc(),-1e-6, 2dx, cell)
# ps=Flux.params(op)
p1_ = ones(1)
p2_ = ones(2)
ps = Flux.params(p1_, p2_)

function loss(u, du)
    vf_ = cat(p2_[1] * ones(sz), p2_[2] * ones(sz), dims = 3)
    duhat = p1_[1] .* (▽ ⋅ ▽(u)) - vf_ ⊗ ▽(u)
    @show l = nae(duhat, du)
end

loss(data[1]...)
opt = ADAM(0.1)

Flux.@epochs 20 Flux.train!(loss, ps, data, opt)

@show p1_, p2_
# heatmap(op.kernel[:, :, 1])


# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENCE in the project root.
# ------------------------------------------------------------------

module LocallyWeightedRegression

using GeoStatsBase

using LinearAlgebra
using NearestNeighbors
using StaticArrays
using Distances
using Variography

import GeoStatsBase: solve

export LocalWeightRegress

"""
    LocalWeightRegress(var₁=>param₁, var₂=>param₂, ...)

Locally weighted regression (LOESS) estimation solver.

## Parameters

* `neighbors` - Number of neighbors (default to all data locations)
* `variogram` - A variogram defined in Variography.jl (default to `ExponentialVariogram()`)
* `distance`  - A distance defined in Distances.jl (default to `Euclidean()`)

### References

Cleveland 1979. *Robust Locally Weighted Regression and Smoothing Scatterplots*
"""
@estimsolver LocalWeightRegress begin
  @param neighbors = nothing
  @param variogram = ExponentialVariogram()
  @param distance = Euclidean()
end

function solve(problem::EstimationProblem, solver::LocalWeightRegress)
  # retrieve problem info
  pdata = data(problem)
  pdomain = domain(problem)

  # result for each variable
  μs = []; σs = []

  for (var,V) in variables(problem)
    # get user parameters
    if var ∈ keys(solver.params)
      varparams = solver.params[var]
    else
      varparams = LocalWeightRegressParam()
    end

    # get valid data for variable
    X, z = valid(pdata, var)

    # number of data points for variable
    ndata = length(z)

    @assert ndata > 0 "estimation requires data"

    # allocate memory
    varμ = Vector{V}(undef, npoints(pdomain))
    varσ = Vector{V}(undef, npoints(pdomain))

    # fit search tree
    kdtree = KDTree(X, varparams.distance)

    # determine number of nearest neighbors to use
    k = varparams.neighbors == nothing ? ndata : varparams.neighbors

    @assert k ≤ ndata "number of neighbors must be smaller or equal to number of data points"

    # determine kernel (or weight) function
    γ = varparams.variogram
    kern(x, y) = sill(γ) - γ(x, y)

    # pre-allocate memory for coordinates
    x = MVector{ndims(pdomain),coordtype(pdomain)}(undef)

    # estimation loop
    for location in LinearPath(pdomain)
      coordinates!(x, pdomain, location)

      inds, dists = knn(kdtree, x, k)

      Xₗ = [ones(eltype(X), k) X[:,inds]']
      zₗ = view(z, inds)

      Wₗ = Diagonal([kern(x, view(X,:,j)) for j in inds])

      # weighted least-squares
      θₗ = Xₗ'*Wₗ*Xₗ \ Xₗ'*Wₗ*zₗ

      # add intercept term to estimation location
      xₗ = [one(eltype(x)), x...]

      # linear combination of response values
      rₗ = Wₗ*Xₗ*(Xₗ'*Wₗ*Xₗ\xₗ)

      varμ[location] = θₗ ⋅ xₗ
      varσ[location] = norm(rₗ)
    end

    push!(μs, var => varμ)
    push!(σs, var => varσ)
  end

  EstimationSolution(pdomain, Dict(μs), Dict(σs))
end

end