#=============================================

Defines a RealInterval type, used to represent intervals of the real line.
RealIntervals are primarily used for representing the output ranges of similarity
functions.

=============================================#

#====================
RealInterval definition and constructors
====================#

@doc """
    struct RealInterval{T<:Real}

Encodes an interval of the real line, such as `[-1,1]` or `[0,Inf)`.

# Fields
- `lower::T`: lower bound on the interval.
- `upper::T`: upper bound on the interval.
- `closed_below::Bool`: whether or not the interval is closed below.
- `closed_above::Bool`: whether or not the interval is closed above.

# Examples
The following snippet constructs `RealInterval` represeting the interval [0,1)

```jldoctest; setup = :(using LSHFunctions)
julia> interval = LSHFunctions.RealInterval(0, 1, true, false);
```

It's generally easier to construct an interval using the `@interval` macro. Check out the documentation for [`@interval`](@ref) for more information.

See also: [`@interval`](@ref)
"""
struct RealInterval{T<:Real}
    # Lower bound on the interval
    lower :: T

    # Upper bound on the interval
    upper :: T

    # Whether or not the lower end of the interval is closed
    closed_below :: Bool

    # Whether or not the upper end of the interval is closed
    closed_above :: Bool

    ### Internal RealInterval constructors
    RealInterval{T}(L, U, CB, CA) where T = RealInterval{T}(T(L), T(U), CB, CA)
    RealInterval{T}(L::T, U::T, CB::Bool, CA::Bool) where T =
        new{T}(L, U, CB, CA)
end

### External RealInterval constructors
@generated function RealInterval(L::Real, U::Real, CB::Bool, CA::Bool)
    bounds_type = begin
        if L <: Integer && U <: AbstractFloat
            U
        else
            L
        end
    end
    return :(RealInterval{$bounds_type}(L, U, CB, CA))
end

### String conversion and displaying RealInterval structs
function Base.string(interval::RealInterval)
    lower_bracket = (interval.closed_below) ? "[" : "("
    upper_bracket = (interval.closed_above) ? "]" : ")"
    lower_bracket * string(interval.lower) * "," * string(interval.upper) * upper_bracket
end

show(io::IO, interval::RealInterval) = print(io, string(interval))

#====================
Operator definitions for RealInterval
====================#
# Interval membership
function Base.:(∈)(x::Real, interval::RealInterval)
    if interval.lower < x < interval.upper
        true
    elseif interval.lower == x && interval.closed_below
        true
    elseif interval.upper == x && interval.closed_above
        true
    else
        false
    end
end

# Interval intersection
function Base.:(∩)(int1::RealInterval, int2::RealInterval)
    lower = max(int1.lower, int2.lower)
    upper = min(int1.upper, int2.upper)
    closed_below = (lower ∈ int1) && (lower ∈ int2)
    closed_above = (upper ∈ int1) && (upper ∈ int2)
    RealInterval(lower, upper, closed_below, closed_above)
end

intersect(int1, int2) = int1 ∩ int2

# Equality
# There are two cases in which equality can hold:
# - both intervals are empty; or
# - both intervals are non-empty, in which case the intervals can only be
#   equal if all of their fields are equal.
function Base.:(==)(int1::RealInterval, int2::RealInterval)
    # Case 1
    if isempty(int1)
        isempty(int2)
    # Case 2
    else
        int1 === int2
    end
end

#====================
Definitions of other methods for RealInterval
====================#

@doc """
    isempty(interval::RealInterval)

Returns `true` if `interval` is empty (i.e. there doesn't exist an `x` for which `x ∈ interval` is `true`), and `false` otherwise.
"""
Base.isempty(interval::RealInterval) =
    interval.lower > interval.upper ||
    (interval.lower == interval.upper && !(interval.closed_below && interval.closed_above))

@doc """
    width(interval::RealInterval)

Return the width of a `RealInterval` (i.e. the difference between its upper and lower bounds.
"""
width(interval::RealInterval{T}) where T =
    isempty(interval) ? T(0) : (interval.upper - interval.lower)

#====================
The @interval macro, which provides the primary interface for users to construct
new RealIntervals.
====================#

@doc """
    @interval(expr)

Construct a new `LSHFunctions.RealInterval` representing an interval on the real line from an expression such as

    0 ≤ x < 1

The returned expression constructs an `LSHFunctions.RealInterval` encoding the lower and upper bounds on the interval, as well as whether the ends are opened or closed.

# Examples
You can construct an interval using the following syntax:

```jldoctest; setup = :(using LSHFunctions)
julia> interval = @interval(0 ≤ x < 1);
```

There are usually multiple ways of constructing the same interval. For instance, each of the expressions below are equivalent ways of constructing the interval `[-1,1]`.

```jldoctest; setup = :(using LSHFunctions)
julia> @interval(-1 ≤  x ≤  1) ==
       @interval(-1 <= x <= 1) ==
       @interval(-1 ≤  y ≤  1) ==
       @interval( 1 ≥  x ≥ -1)
true
```

You can even create intervals with `Inf` or `-Inf` at the endpoints, e.g. `@interval(-Inf < x < Inf)`.

There are two primary operations you can run on an interval: testing for membership and intersection. You can test whether or not `x` is in an interval using `x ∈ interval`, as shown below.

```jldoctest; setup = :(using LSHFunctions)
julia> interval = @interval(0 ≤ x < 1);

julia> 0 ∈ interval && 1 ∉ interval
true

julia> 0 in interval    # This syntax also works
true
```

You can also intersect two intervals using the `∩` operator (or by using `intersect(interval_1, interval_2)`).

```
julia> @interval(0 ≤ x < 1) ∩ @interval(1/2 < x ≤ 1) == @interval(1/2 < x < 1)
true
```

See also: [`RealInterval`](@ref)
"""
macro interval(expr)
    invalid_expr_msg = """
    Expression is invalid. The expression passed to @interval must take the form

        a {<,≤,>,≥} x {<,≤,>,≥} b

    where a and b evaluate to real numbers, and x is a valid variable identifier.
    See the documentation for @interval for more information.
    """

    if !isa(expr, Expr)
        return :($invalid_expr_msg |> ErrorException |> throw)
    end

    args = expr.args

    # Convert >= and <= to ≥ and ≤ to reduce the number of symbols we have to deal with
    map!(sym -> (sym == :>=) ? :≥ : sym, args, args)
    map!(sym -> (sym == :<=) ? :≤ : sym, args, args)

    valid_expr = (
        length(args) == 5          &&
        expr.head == :comparison   &&

        # 3rd argument must be a valid variable identifier
        isa(args[3], Symbol)       &&
        Base.isidentifier(args[3]) &&

        # 2nd and 4th arguments must be inequalities
        args[2] ∈ (:<, :≤, :>, :≥) &&
        args[4] ∈ (:<, :≤, :>, :≥)
    )

    if !valid_expr
        return :($invalid_expr_msg |> ErrorException |> throw)
    end

    # Extract a tuple (a, op1, b, op2) from the expression $a $op1 x $op2 $b. E.g.
    # if expr == :(0 ≤ x < b), return (0, :≤, :b, :<)
    interval_tuple = begin
        # Flip inequalities if expression is of the form b {>,≥} x {>,≥} a
        if args[2] in (:>, :≥) && args[4] in (:>, :≥)
            map!(x -> (x == :>) ? :< : x, args, args)
            map!(x -> (x == :≥) ? :≤ : x, args, args)
            reverse!(args)
        elseif !(args[2] in (:<, :≤) && args[4] in (:<, :≤))
            return quote
                "Inequalities point in opposite directions" |>
                ErrorException |>
                throw
            end
        end
        (args[1], args[2], args[5], args[4])
    end

    # Construct a new RealInterval based on the inequalities that were provided
    a, b = interval_tuple[1], interval_tuple[3]
    closed_below = (interval_tuple[2] == :≤)
    closed_above = (interval_tuple[4] == :≤)
    quote
        RealInterval($(esc(a)), $(esc(b)), $(esc(closed_below)), $(esc(closed_above)))
    end
end


# ---
# title: 391. Perfect Rectangle
# id: problem391
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Line Sweep
# link: <https://leetcode.com/problems/perfect-rectangle/description/>
# hidden: true
# ---
#
# Given N axis-aligned rectangles where N > 0, determine if they all together
# form an exact cover of a rectangular region.
#
# Each rectangle is represented as a bottom-left point and a top-right point.
# For example, a unit square is represented as [1,1,2,2]. (coordinate of bottom-
# left point is (1, 1) and top-right point is (2, 2)).
#
# ![](https://assets.leetcode.com/uploads/2018/10/22/rectangle_perfect.gif)
#
# **Example 1:**
#
#
#
#     rectangles = [
#       [1,1,3,3],
#       [3,1,4,2],
#       [3,2,4,4],
#       [1,3,2,4],
#       [2,3,3,4]
#     ]
#
#     Return true. All 5 rectangles together form an exact cover of a rectangular region.
#
#
#
#
#
#
# ![](https://assets.leetcode.com/uploads/2018/10/22/rectangle_separated.gif)
#
# **Example 2:**
#
#
#
#     rectangles = [
#       [1,1,2,3],
#       [1,3,2,4],
#       [3,1,4,2],
#       [3,2,4,4]
#     ]
#
#     Return false. Because there is a gap between the two rectangular regions.
#
#
#
#
#
#
# ![](https://assets.leetcode.com/uploads/2018/10/22/rectangle_hole.gif)
#
# **Example 3:**
#
#
#
#     rectangles = [
#       [1,1,3,3],
#       [3,1,4,2],
#       [1,3,2,4],
#       [3,2,4,4]
#     ]
#
#     Return false. Because there is a gap in the top center.
#
#
#
#
#
#
# ![](https://assets.leetcode.com/uploads/2018/10/22/rectangle_intersect.gif)
#
# **Example 4:**
#
#
#
#     rectangles = [
#       [1,1,3,3],
#       [3,1,4,2],
#       [1,3,2,4],
#       [2,2,4,4]
#     ]
#
#     Return false. Because two of the rectangles overlap with each other.
#
#
#
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


function DFSVisit(G::DictGraph, dfsDict::SortedDict, vertex::Int, count::Int)
    count  += 1;
    dfsDict[vertex][2] = t;
    dfsDict[vertex][1] = 0;

    for v in G[vertex]
        if dfsDict[v[1]][1] == -1
            dfsDict[v[1]][3] = vertex
            DFSVisit(G, dfsDict, v[1], count);
        end
    end

    dfsDict[vertex][1] = 1
    count += 1
    dfsDict[vertex][4] += count

    return;
end

function DFS(G::DictGraph; vertex::Int = -1)
    dfsDict = SortedDict(i => [-1 0 -1 0] for i in keys(G.relation))
    #                         [c  d  π f]
    count = -1

    if vertex > -1
        vertex in keys(G) || error("$vertex is not a vertex of this graph.")
        DFSVisit(G, dfsDict, vertex, count)
    end

    for v in keys(G.relation)
        dfsDict[v][1] == -1 && DFSVisit(G, dfsDict, v, count)
    end

    return dfsDict
end


###############
# Composition #
###############

"""
    Composed(ts::A)

    ∘(b1::Bijector{N}, b2::Bijector{N})::Composed{<:Tuple}
    composel(ts::Bijector{N}...)::Composed{<:Tuple}
    composer(ts::Bijector{N}...)::Composed{<:Tuple}

where `A` refers to either
- `Tuple{Vararg{<:Bijector{N}}}`: a tuple of bijectors of dimensionality `N`
- `AbstractArray{<:Bijector{N}}`: an array of bijectors of dimensionality `N`

A `Bijector` representing composition of bijectors. `composel` and `composer` results in a
`Composed` for which application occurs from left-to-right and right-to-left, respectively.

Note that all the alternative ways of constructing a `Composed` returns a `Tuple` of bijectors.
This ensures type-stability of implementations of all relating methdos, e.g. `inv`.

If you want to use an `Array` as the container instead you can do

    Composed([b1, b2, ...])

In general this is not advised since you lose type-stability, but there might be cases
where this is desired, e.g. if you have a insanely large number of bijectors to compose.

# Examples
## Simple example
Let's consider a simple example of `Exp`:
```julia-repl
julia> using Bijectors: Exp

julia> b = Exp()
Exp{0}()

julia> b ∘ b
Composed{Tuple{Exp{0},Exp{0}},0}((Exp{0}(), Exp{0}()))

julia> (b ∘ b)(1.0) == exp(exp(1.0))    # evaluation
true

julia> inv(b ∘ b)(exp(exp(1.0))) == 1.0 # inversion
true

julia> logabsdetjac(b ∘ b, 1.0)         # determinant of jacobian
3.718281828459045
```

# Notes
## Order
It's important to note that `∘` does what is expected mathematically, which means that the
bijectors are applied to the input right-to-left, e.g. first applying `b2` and then `b1`:
```julia
(b1 ∘ b2)(x) == b1(b2(x))     # => true
```
But in the `Composed` struct itself, we store the bijectors left-to-right, so that
```julia
cb1 = b1 ∘ b2                  # => Composed.ts == (b2, b1)
cb2 = composel(b2, b1)         # => Composed.ts == (b2, b1)
cb1(x) == cb2(x) == b1(b2(x))  # => true
```

## Structure
`∘` will result in "flatten" the composition structure while `composel` and
`composer` preserve the compositional structure. This is most easily seen by an example:
```julia-repl
julia> b = Exp()
Exp{0}()

julia> cb1 = b ∘ b; cb2 = b ∘ b;

julia> (cb1 ∘ cb2).ts # <= different
(Exp{0}(), Exp{0}(), Exp{0}(), Exp{0}())

julia> (cb1 ∘ cb2).ts isa NTuple{4, Exp{0}}
true

julia> Bijectors.composer(cb1, cb2).ts
(Composed{Tuple{Exp{0},Exp{0}},0}((Exp{0}(), Exp{0}())), Composed{Tuple{Exp{0},Exp{0}},0}((Exp{0}(), Exp{0}())))

julia> Bijectors.composer(cb1, cb2).ts isa Tuple{Composed, Composed}
true
```

"""
struct Composed{A, N} <: Bijector{N}
    ts::A

    Composed(bs::C) where {N, C<:Tuple{Vararg{<:Bijector{N}}}} = new{C, N}(bs)
    Composed(bs::A) where {N, A<:AbstractArray{<:Bijector{N}}} = new{A, N}(bs)
end

isclosedform(b::Composed) = all(isclosedform, b.ts)

"""
    composel(ts::Bijector...)::Composed{<:Tuple}

Constructs `Composed` such that `ts` are applied left-to-right.
"""
composel(ts::Bijector{N}...) where {N} = Composed(ts)

"""
    composer(ts::Bijector...)::Composed{<:Tuple}

Constructs `Composed` such that `ts` are applied right-to-left.
"""
composer(ts::Bijector{N}...) where {N} = Composed(reverse(ts))

# The transformation of `Composed` applies functions left-to-right
# but in mathematics we usually go from right-to-left; this reversal ensures that
# when we use the mathematical composition ∘ we get the expected behavior.
# TODO: change behavior of `transform` of `Composed`?
@generated function ∘(b1::Bijector{N1}, b2::Bijector{N2}) where {N1, N2}
    if N1 == N2
        return :(composel(b2, b1))
    else
        return :(throw(DimensionMismatch("$(typeof(b1)) expects $(N1)-dim but $(typeof(b2)) expects $(N2)-dim")))
    end
end

# type-stable composition rules
∘(b1::Composed{<:Tuple}, b2::Bijector) = composel(b2, b1.ts...)
∘(b1::Bijector, b2::Composed{<:Tuple}) = composel(b2.ts..., b1)
∘(b1::Composed{<:Tuple}, b2::Composed{<:Tuple}) = composel(b2.ts..., b1.ts...)

# type-unstable composition rules
∘(b1::Composed{<:AbstractArray}, b2::Bijector) = Composed(pushfirst!(copy(b1.ts), b2))
∘(b1::Bijector, b2::Composed{<:AbstractArray}) = Composed(push!(copy(b2.ts), b1))
function ∘(b1::Composed{<:AbstractArray}, b2::Composed{<:AbstractArray})
    return Composed(append!(copy(b2.ts), copy(b1.ts)))
end

# if combining type-unstable and type-stable, return type-unstable
function ∘(b1::T1, b2::T2) where {T1<:Composed{<:Tuple}, T2<:Composed{<:AbstractArray}}
    error("Cannot compose compositions of different container-types; ($T1, $T2)")
end
function ∘(b1::T1, b2::T2) where {T1<:Composed{<:AbstractArray}, T2<:Composed{<:Tuple}}
    error("Cannot compose compositions of different container-types; ($T1, $T2)")
end

∘(::Identity{N}, ::Identity{N}) where {N} = Identity{N}()
∘(::Identity{N}, b::Bijector{N}) where {N} = b
∘(b::Bijector{N}, ::Identity{N}) where {N} = b

inv(ct::Composed) = Composed(reverse(map(inv, ct.ts)))

# # TODO: should arrays also be using recursive implementation instead?
function (cb::Composed{<:AbstractArray{<:Bijector}})(x)
    @assert length(cb.ts) > 0
    res = cb.ts[1](x)
    for b ∈ Base.Iterators.drop(cb.ts, 1)
        res = b(res)
    end

    return res
end

@generated function (cb::Composed{T})(x) where {T<:Tuple}
    @assert length(T.parameters) > 0
    expr = :(x)
    for i in 1:length(T.parameters)
        expr = :(cb.ts[$i]($expr))
    end
    return expr
end

function logabsdetjac(cb::Composed, x)
    y, logjac = forward(cb.ts[1], x)
    for i = 2:length(cb.ts)
        res = forward(cb.ts[i], y)
        y = res.rv
        logjac += res.logabsdetjac
    end

    return (rv = y, logabsdetjac = logjac)
end

@generated function logabsdetjac(cb::Composed{T}, x) where {T<:Tuple}
    N = length(T.parameters)

    expr = Expr(:block)
    push!(expr.args, :((y, logjac) = forward(cb.ts[1], x)))

    for i = 2:N - 1
        temp = gensym(:res)
        push!(expr.args, :($temp = forward(cb.ts[$i], y)))
        push!(expr.args, :(y = $temp.rv))
        push!(expr.args, :(logjac += $temp.logabsdetjac))
    end
    # don't need to evaluate the last bijector, only it's `logabsdetjac`
    push!(expr.args, :(logjac += logabsdetjac(cb.ts[$N], y)))

    push!(expr.args, :(return logjac))

    return expr
end

function forward(cb::Composed, x)
    rv, logjac = forward(cb.ts[1], x)

    for t in cb.ts[2:end]
        res = forward(t, rv)
        rv = res.rv
        logjac = res.logabsdetjac + logjac
    end
    return (rv=rv, logabsdetjac=logjac)
end

@generated function forward(cb::Composed{T}, x) where {T<:Tuple}
    expr = Expr(:block)
    push!(expr.args, :((y, logjac) = forward(cb.ts[1], x)))
    for i = 2:length(T.parameters)
        temp = gensym(:temp)
        push!(expr.args, :($temp = forward(cb.ts[$i], y)))
        push!(expr.args, :(y = $temp.rv))
        push!(expr.args, :(logjac += $temp.logabsdetjac))
    end
    push!(expr.args, :(return (rv = y, logabsdetjac = logjac)))

    return expr
end


"""Length of Longest Common Substring problem in Julia

Given two strings, we have to find the longest common substring in Julia

The problem can be solved using Dynamic Programming
"""

function lengthoflongestcommonsubstring(S1, S2)
    # Build a 2D matrix in Bottom-Up manner. Every element
    # dp[i][j] contains the length of the longest common substring
    # of S1[0 .. i] and S2[0 .. j].
    n = length(S1)
    m = length(S2)
    ans = 0
    # Fill every element of the 2D matrix with 0 as value
    dp = zeros(Int, n+1, m+1)

    for i in 1:(n+1)
        for j in 1:(m+1)

            # length of LCS when one of the string is empty is zero
            if i == 1 || j == 1
                dp[i,j] = 0

            # If the elements of both the strings match, then
            # the value of LCS increments by one.
            # Also Check if the value of lcs has reached the max till now.
            elseif S1[i-1] == S2[j-1]
                dp[i,j] = dp[(i - 1),(j - 1)] + 1
                ans = max(ans, dp[i,j])
            else
                dp[i,j] = 0
            end
        end
    end
    return ans
end

print("Enter the first string: ")
s1 = readline()
print("Enter the second string: ")
s2 = readline()
res = lengthoflongestcommonsubstring(s1, s2)
print("The length of the longest common substring is ")
print(res)

# Time Complexity: O(n * m), where n, m are the size of the given strings
# Space Complexity: O(n * m)

# SAMPLE INPUT AND OUTPUT

# Enter the first string: runner
# Enter the second string: run
# The length of the longest common substring is 3


using Test
using ConvOptDL.Utils

@testset "Utils Test (batched_outer)" begin
    A = hcat([ones(4)*i for i in 1:5]...)
    B = hcat([ones(3)*i for i in 1:5]...)
    refC = ones(4, 3, 5)
    for i in 1:5
        refC[:, :, i] .*= i^2
    end
    @test batched_outer(A, B) == refC
end

@testset "Utils Test (batched_kronecker)" begin
    A = [
        1 2;
        3 4;
    ]
    B = ones(3, 2)
    ans_ = repeat(repeat(A, inner=(3, 2)), outer=(1, 1, 5))
    @test batched_kronecker(repeat(A, outer=(1, 1, 5)), repeat(B, outer=(1, 1, 5))) == ans_
end


@testset "Sampling Density" begin

# create a 10x10 grid of unit spaced sampling points
N = 10
g = (0:(N-1)) ./ N .- 0.5
x = vec(ones(N) * g')
y = vec(g * ones(N)')
nodes = cat(x',y', dims=1)

# approximate the density weights
p = plan_nfft(nodes, (N,N), m = 5, σ = 2.0);
weights = sdc(p, iters = 10)

# test if they approximate the true weights (1/(N*N))
@test all( (≈).(vec(weights), 1/(N*N), rtol=1e-7) )

end


using JuMP, EAGO, Gurobi

m = Model(with_optimizer(EAGO.Optimizer, relaxed_optimizer = Gurobi.Optimizer(OutputFlag=0)))

# ----- Variables ----- #
x_Idx = Any[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22]
@variable(m, x[x_Idx])
JuMP.set_lower_bound(x[21], 0.0)
JuMP.set_lower_bound(x[1], -0.5)
JuMP.set_upper_bound(x[1], 0.5)
JuMP.set_lower_bound(x[2], 5.4)
JuMP.set_upper_bound(x[2], 6.4)
JuMP.set_lower_bound(x[3], 0.4)
JuMP.set_upper_bound(x[3], 1.4)
JuMP.set_lower_bound(x[4], 4.9)
JuMP.set_upper_bound(x[4], 5.9)
JuMP.set_lower_bound(x[5], 1.3)
JuMP.set_upper_bound(x[5], 2.3)
JuMP.set_lower_bound(x[6], 3.9)
JuMP.set_upper_bound(x[6], 4.9)
JuMP.set_lower_bound(x[7], 2.1)
JuMP.set_upper_bound(x[7], 3.1)
JuMP.set_lower_bound(x[8], 4.1)
JuMP.set_upper_bound(x[8], 5.1)
JuMP.set_lower_bound(x[9], 2.8)
JuMP.set_upper_bound(x[9], 3.8)
JuMP.set_lower_bound(x[10], 3.0)
JuMP.set_upper_bound(x[10], 4.0)
JuMP.set_lower_bound(x[11], 3.9)
JuMP.set_upper_bound(x[11], 4.9)
JuMP.set_lower_bound(x[12], 3.2)
JuMP.set_upper_bound(x[12], 4.2)
JuMP.set_lower_bound(x[13], 4.7)
JuMP.set_upper_bound(x[13], 5.7)
JuMP.set_lower_bound(x[14], 2.3)
JuMP.set_upper_bound(x[14], 3.3)
JuMP.set_lower_bound(x[15], 5.6)
JuMP.set_upper_bound(x[15], 6.6)
JuMP.set_lower_bound(x[16], 2.3)
JuMP.set_upper_bound(x[16], 3.3)
JuMP.set_lower_bound(x[17], 6.0)
JuMP.set_upper_bound(x[17], 7.0)
JuMP.set_lower_bound(x[18], 1.9)
JuMP.set_upper_bound(x[18], 2.9)
JuMP.set_lower_bound(x[19], 6.9)
JuMP.set_upper_bound(x[19], 7.9)
JuMP.set_lower_bound(x[20], 1.0)
JuMP.set_upper_bound(x[20], 2.0)
JuMP.set_upper_bound(x[21], 10.0)
JuMP.set_lower_bound(x[22], -2.0)
JuMP.set_upper_bound(x[22], 2.0)

# ----- Constraints ----- #
@NLconstraint(m, e2, x[22]*x[1]-x[2]+x[21] == 0.0)
@NLconstraint(m, e3, x[22]*x[3]-x[4]+x[21] == 0.0)
@NLconstraint(m, e4, x[22]*x[5]-x[6]+x[21] == 0.0)
@NLconstraint(m, e5, x[22]*x[7]-x[8]+x[21] == 0.0)
@NLconstraint(m, e6, x[22]*x[9]-x[10]+x[21] == 0.0)
@NLconstraint(m, e7, x[22]*x[11]-x[12]+x[21] == 0.0)
@NLconstraint(m, e8, x[22]*x[13]-x[14]+x[21] == 0.0)
@NLconstraint(m, e9, x[22]*x[15]-x[16]+x[21] == 0.0)
@NLconstraint(m, e10, x[22]*x[17]-x[18]+x[21] == 0.0)
@NLconstraint(m, e11, x[22]*x[19]-x[20]+x[21] == 0.0)

# ----- Objective ----- #
@NLobjective(m, Min, ( (x[1])^2+ (x[2]-5.9)^2+ (x[3]-0.9)^2+ (x[4]-5.4)^2+ (x[5]-1.8)^2+ (x[6]-4.4)^2+ (x[7]-2.6)^2+ (x[8]-4.6)^2+ (x[9]-3.3)^2+ (x[10]-3.5)^2+ (x[11]-4.4)^2+ (x[12]-3.7)^2+ (x[13]-5.2)^2+ (x[14]-2.8)^2+ (x[15]-6.1)^2+ (x[16]-2.8)^2+ (x[17]-6.5)^2+ (x[18]-2.4)^2+ (x[19]-7.4)^2+ (x[20]-1.5)^2))

JuMP.optimize!(m)

backend_opt = backend(m).optimizer.model.optimizer
last_relaxed_opt = backend_opt.relaxed_optimizer
run_time = backend_opt._run_time
println("run time: $run_time")


include("model_search/mountain_car_model_search_utils.jl")

struct MountainCarModelSearchAgent
    mountaincar::MountainCar
    model::MountainCar
    data::Array{MountainCarContTransition}
    optimization_params::MountainCarOptimizationParameters
    horizon::Int64
end

function return_based_model_search(
    mountaincar::MountainCar,
    data::Array{MountainCarContTransition},
    optimization_params::MountainCarOptimizationParameters,
    horizon::Int64;
    hardcoded::Bool = false,
    num_episodes_eval::Int64 = 1,
    debug::Bool = false,
    eval_distance::Bool = false,
    optimistic::Bool = false,
    model_evaluation::Bool = false,
)
    params = true_params
    least_squares_params = get_least_squares_fit(mountaincar, params, data)
    x_matrices_array, _, x_next_array, _, cost_array = preprocess_data(data)
    function eval_fn(p)
        policy, values, converged = value_iteration(mountaincar, p)
        if converged
            if optimistic
                return mfmc_optimistic_evaluation(
                    mountaincar,
                    policy,
                    p,
                    horizon,
                    x_matrices_array,
                    x_next_array,
                    cost_array,
                    num_episodes_eval,
                    0.015;
                    debug = debug,
                )
            elseif model_evaluation
                return mfmc_model_evaluation(
                    mountaincar,
                    policy,
                    p,
                    horizon,
                    x_matrices_array,
                    x_next_array,
                    cost_array,
                    num_episodes_eval,
                    0.01;
                    debug = debug,
                )
            else
                return mfmc_evaluation(
                    mountaincar,
                    policy,
                    values,
                    horizon,
                    x_matrices_array,
                    x_next_array,
                    cost_array,
                    num_episodes_eval;
                    hardcoded = hardcoded,
                    debug = debug,
                )
            end
        else
            if debug
                println("Value Iteration did not converge. Skipping parameters ", p)
            end
            return Inf
        end
    end

    params = hill_climb(
        eval_fn,
        optimization_params;
        least_squares_params = least_squares_params,
        debug = debug,
    )
    if eval_distance
        policy, values, _ = value_iteration(mountaincar, params)
        estimated_return, eval_distances = mfmc_evaluation(
            mountaincar, policy, values, horizon,
            x_matrices_array, x_next_array, cost_array,
            num_episodes_eval;
            eval_distance = eval_distance,
        )
        println("Mean distance ", mean(eval_distances),
                " Std distance ", std(eval_distances),
                " Max distance ", maximum(eval_distances),
                " Min distance ", minimum(eval_distances))
        println("Estimated return is ", estimated_return)
    end
    params
end

function bellman_based_model_search(
    mountaincar::MountainCar,
    data::Array{MountainCarContTransition},
    optimization_params::MountainCarOptimizationParameters,
    horizon::Int64;
    num_episodes_eval::Int64 = 1,
    debug::Bool = false,
)
    params = true_params
    least_squares_params = get_least_squares_fit(mountaincar, params, data)
    x_matrices_array, _, xnext_array, _, _ = preprocess_data(data)
    function eval_fn(p)
        policy, values, converged = value_iteration(mountaincar, p)
        if converged
            return bellman_evaluation(
                mountaincar,
                p,
                policy,
                values,
                horizon,
                x_matrices_array,
                xnext_array,
                num_episodes_eval;
                debug = debug,
            )
        else
            if debug
                println("Value Iteration did not converge. Skipping parameters ", p)
            end
            return Inf
        end
    end
    params = hill_climb(
        eval_fn,
        optimization_params;
        least_squares_params = least_squares_params,
        debug = debug,
    )
    params
end

function MountainCarModelSearchAgent(
    mountaincar::MountainCar,
    model::MountainCar,
    horizon::Int64,
    num_episodes_offline::Int64,
)
    data = generate_batch_data(mountaincar, true_params, num_episodes_offline, horizon)
    # println("Generated ", length(data), " transitions offline")
    MountainCarModelSearchAgent(mountaincar, model, horizon, data)
end

function MountainCarModelSearchAgent(
    mountaincar::MountainCar,
    model::MountainCar,
    horizon::Int64,
    data::Array{MountainCarContTransition},
)
    maximum_num_evaluations = 90
    optimization_params =
        MountainCarOptimizationParameters([0.0012, 0.5], [-0.0025, 3], maximum_num_evaluations, false)
    MountainCarModelSearchAgent(mountaincar, model, data, optimization_params, horizon)
end

function run_return_based_model_search(
    agent::MountainCarModelSearchAgent;
    ensemble = false,
    hardcoded = false,
    max_steps = 500,
    debug = false,
    eval_distance = false,
    optimistic = false,
    model_evaluation = false,
)
    params = return_based_model_search(
        agent.model,
        agent.data,
        agent.optimization_params,
        agent.horizon;
        hardcoded = hardcoded,
        debug = debug,
        eval_distance = eval_distance,
        optimistic = optimistic,
        model_evaluation = model_evaluation,
    )
    println("Params found is ", params)
    run(agent, params; max_steps = max_steps)
end

function run_bellman_based_model_search(
    agent::MountainCarModelSearchAgent;
    max_steps = 500,
    debug = false,
)
    params = bellman_based_model_search(
        agent.model,
        agent.data,
        agent.optimization_params,
        agent.horizon;
        debug = debug,
    )
    run(agent, params, max_steps = max_steps)
end

function run(
    agent::MountainCarModelSearchAgent,
    params::Array{Float64};
    max_steps = 500,
    debug = false,
    num_eval = 10,
)
    policy, _, _ = value_iteration(agent.model, params)
    rng = MersenneTwister(0)
    num_steps = 0.0
    for _ in 1:num_eval
        num_steps += length(
            simulate_episode(
                agent.mountaincar,
                true_params,
                agent.horizon;
                policy = policy,
                rng = rng,
            )
        )
    end
    num_steps/num_eval
end


# Stiff 1
using ParameterizedFunctions, OrdinaryDiffEq, Plots
plotly()
rober = @ode_def Rober begin
  dy₁ = -k₁*y₁+k₃*y₂*y₃
  dy₂ =  k₁*y₁-k₂*y₂^2-k₃*y₂*y₃
  dy₃ =  k₂*y₂^2
end k₁ k₂ k₃
prob = ODEProblem(rober,[1.0;0.0;0.0],(0.0,1e11),(0.04,3e7,1e4))
sol = solve(prob, Rosenbrock23())
plot(sol, xscale=:log10, tspan=(1e-6, 1e5), layout=(3,1))

#######################################
#######################################
#######################################
# Stiff 2

using OrdinaryDiffEq, Plots
function orego(du,u,p,t)
  s,q,w = p
  y1,y2,y3 = u
  du[1] = s*(y2+y1*(1-q*y1-y2))
  du[2] = (y3-(1+y1)*y2)/s
  du[3] = w*(y1-y3)
end
p = [77.27,8.375e-6,0.161]
prob = ODEProblem(orego,[1.0,2.0,3.0],(0.0,360.0),p)
sol = solve(prob, Rodas5())
plot(sol)
# phase space

#######################################
#######################################
#######################################
# Stiff SDE

using StochasticDiffEq
function orego(du,u,p,t)
  s,q,w = p
  y1,y2,y3 = u
  du[1] = s*(y2+y1*(1-q*y1-y2))
  du[2] = (y3-(1+y1)*y2)/s
  du[3] = w*(y1-y3)
end
function g(du,u,p,t)
  du[1] = 0.1u[1]
  du[2] = 0.1u[2]
  du[3] = 0.1u[3]
end
p = [77.27,8.375e-6,0.161]
prob = SDEProblem(orego,g,[1.0,2.0,3.0],(0.0,30.0),p)
sol = solve(prob,SOSRI())
plot(sol)

#######################################
#######################################
#######################################
# DAE

function f(out, da, a, p, t)
   L1, m1, L2, m2, g = p

   u1, v1, x1, y1, T1,
   u2, v2, x2, y2, T2 = a

   du1, dv1, dx1, dy1, dT1,
   du2, dv2, dx2, dy2, dT2 = da

   out[1]  = x2*T2/(m2*L2) - du2
   out[2]  = y2*T2/(m2*L2) - g - dv2
   out[3]  = u2 - dx2
   out[4]  = v2 - dy2
   out[5]  = u2^2 + v2^2 -y2*g + T2/m2

   out[6]  = x1*T1/(m1*L1) - x2*T2/(m2*L2) - du1
   out[7]  = y1*T1/(m1*L1) - g - y2*T2/(m2*L2) - dv1
   out[8]  = u1 - dx1
   out[9]  = v1 - dy1
   out[10] = u1^2 + v1^2 + T1/m1 +
                (-x1*x2 - y1*y2)/(m1*L2)*T2 - y1*g
   nothing
end

# Release pendulum from top right
u0 = zeros(10)
u0[3] = 1.0
u0[8] = 1.0
du0 = zeros(10)
du0[2] = 9.8
du0[7] = 9.8

p = [1,1,1,1,9.8]
tspan = (0.,100.)

differential_vars = [true, true, true, true, false,
                     true, true, true, true, false]
prob = DAEProblem(f, du0, u0, tspan, p, differential_vars = differential_vars)
using Sundials
sol = solve(prob, IDA())
plot(sol, vars=((x1,y1,x2,y2)->(x1, y1),3,4,8,9), lab="")
plot!(sol, vars=((x1,y1,x2,y2)->(x1+x2, y1+y2),3,4,8,9), xlims=(-2.2, 2.2), aspect_ratio=1, lab="")

#######################################
#######################################
#######################################
# DDE
using DiffEqProblemLibrary.DDEProblemLibrary, LinearAlgebra
using Plots
using DelayDiffEq
DDEProblemLibrary.importddeproblems()
prob = DDEProblemLibrary.prob_dde_RADAR5_waltman_5
sol = solve(prob, MethodOfSteps(Rosenbrock23(linsolve=LinSolveFactorize(lu))); reltol = 1e-6, abstol = [1e-18, 1e-18, 1e-18, 1e-18, 1e-6, 1e-6])
plot(sol)

#######################################
#######################################
#######################################
# Stiff PDE

using OrdinaryDiffEq, Sundials, Plots

# initial condition
function init_brusselator_2d(xyd)
    N = length(xyd)
    u = zeros(N, N, 2)
    for I in CartesianIndices((N, N))
        x = xyd[I[1]]
        y = xyd[I[2]]
        u[I,1] = 22*(y*(1-y))^(3/2)
        u[I,2] = 27*(x*(1-x))^(3/2)
    end
    u
end

N = 32

xyd_brusselator = range(0,stop=1,length=N)

u0 = vec(init_brusselator_2d(xyd_brusselator))

tspan = (0, 22.)

p = (3.4, 1., 10., xyd_brusselator)

brusselator_f(x, y, t) = ifelse((((x-0.3)^2 + (y-0.6)^2) <= 0.1^2) &&
                                (t >= 1.1), 5., 0.)
function brusselator_2d_loop(du, u, p, t)
    A, B, α, xyd = p
    dx = step(xyd)
    N = length(xyd)
    α = α/dx^2
    limit = a -> let N=N
        a == N+1 ? 1 :
        a == 0 ? N :
        a
    end
    II = LinearIndices((N, N, 2))

    @inbounds begin
        for I in CartesianIndices((N, N))
            x = xyd[I[1]]
            y = xyd[I[2]]
            i = I[1]
            j = I[2]
            ip1 = limit(i+1)
            im1 = limit(i-1)
            jp1 = limit(j+1)
            jm1 = limit(j-1)

            ii1 = II[i,j,1]
            ii2 = II[i,j,2]

            du[II[i,j,1]] = α*(u[II[im1,j,1]] + u[II[ip1,j,1]] + u[II[i,jp1,1]] + u[II[i,jm1,1]] - 4u[ii1]) +
            B + u[ii1]^2*u[ii2] - (A + 1)*u[ii1] + brusselator_f(x, y, t)

            du[II[i,j,2]] = α*(u[II[im1,j,2]] + u[II[ip1,j,2]] + u[II[i,jp1,2]] + u[II[i,jm1,2]] - 4u[II[i,j,2]]) +
            A*u[ii1] - u[ii1]^2*u[ii2]
        end
    end
    nothing
end
gr()
function plot_sol(sol)
    off = N^2
    for t in sol.t[1]:0.1:sol.t[end]
        solt = sol(t)
        plt1 = surface(reshape(solt[1:off], N, N), zlims=(0, 5), leg=false)
        surface!(plt1, reshape(solt[off+1:end], N, N), zlims=(0, 5), leg=false)
        display(plt1)
        sleep(0.05)
    end
    nothing
end

prob2 = ODEProblem(brusselator_2d_loop, u0, tspan, p)

using Sundials
sol2 = @time solve(prob2, CVODE_BDF(linear_solver=:GMRES))

plot_sol(sol2)


### observation simulation and beam fitting

## baseline coordinates
## arr: DataFrames with at least :X, :Y, :Z
## check at https://web.njit.edu/~gary/728/Lecture6.html

function calc_baselines(arr::AbstractDataFrame , lat=-23.0262015)

    coordinates= vcat(hcat(convert(Vector{Float64},arr[:X]))' ,
        hcat(convert(Vector{Float64},arr[:Y]))' ,
        hcat(convert(Vector{Float64},arr[:Z]))')
    rot= [0 -sind(lat) cosd(lat) ; 1 0 0 ; 0  cosd(lat)  sind(lat)]
    xyz= *(rot,coordinates)

    Nc= size(xyz)
    Nbl= convert(Int64,(Nc[2]*(Nc[2]-1))/2)
    bl= zeros(Float64,3,2Nbl)   ## symmetry of the bl

    ibl= 1
    for i in 1:Nc[2]
        for j in i+1:Nc[2]
            bl[1,ibl]= xyz[1,i]-xyz[1,j] ; bl[1,ibl+Nbl]= xyz[1,j]-xyz[1,i]
            bl[2,ibl]= xyz[2,i]-xyz[2,j] ; bl[2,ibl+Nbl]= xyz[2,j]-xyz[2,i]
            bl[3,ibl]= xyz[3,i]-xyz[3,j] ; bl[3,ibl+Nbl]= xyz[3,j]-xyz[3,i]
            ibl += 1
        end
    end
    return(bl)
end

####
### bl: baselines computed in calc_baselines
### h0: hour angle in degrees
### d0: source declination in degrees
### only snapshot still...
function calc_uv(bl, h0=0, d0=-50)

    h0 *= 15   #hours to degree
    proj= [sind(h0) cosd(h0) 0 ; -sind(d0)cosd(h0) sind(d0)sind(h0) cosd(d0) ; cosd(d0)cosd(h0) -cosd(d0)sind(h0) sind(d0)]
    uv= *(proj, bl)
    return(uv)
end

function calc_uv_coverage
    ## will compute the uv_coverage with an observing and sampling time ..
    ## for now only a snapshot.
    ## ...
end

################################
## setup the npix,fov, freq, etc for the beam extraction from the uv
## coverage. It is based on the uniform weighting 2/FOV for the uv-cell
##
function calc_dirtybeam(uv, npix=511, sizefit=255 ; robust=2)
    c= 299792458.0  ; ν= 100*1e9 ; λ= c / ν
    D= 12 ; PB= 1.13 * λ / D
    duv= 2/(PB*npix) ; ds= duv/λ
    dr= 180*3600/(π*ds*npix)
    lmax= maximum(abs.(uv))

    if (duv*npix/2) < lmax
        println("##UV-gridding error, npix too small...")
    end

    binuv= range(-duv*npix/2, stop=duv*npix/2, length=npix)
    Hh= fit(Histogram, (uv[1,:] , uv[2,:]) , (binuv , binuv))
    H= Hh.weights

    ## Robust weighting R (Briggs)
    ## R=-2 : uniform
    ## R=2  : natural
    ## R=0.5: "briggs"
    Hsum= sum(H)
    H2sum= sum(1 ./ H[H .> 0].^2)/Hsum
    f2= (5 * 10.0^(-robust))^2 * H2sum
    H2 = H ./ (1 .+ f2 .* H)

    h= fft(H2)
    hshift= fftshift(abs.(real.(h)))

    center= npix/2
    coord_a= convert(Int64,center-sizefit/2)
    coord_b= convert(Int64,center+sizefit/2)

    hshiftcut= hshift'[coord_a:coord_b, coord_a:coord_b]
    return(hshiftcut , dr)
end

## fitting of the Gaussian 2D for the beam
## assuming offset=0 for the beam

function gaussian2D(x,y, amplitude, xo, yo, σx, σy, θ)

    a = (cosd(θ)^2)/(2σx^2) + (sind(θ)^2)/(2*σy^2)
    b = -(sind(2θ))/(2σx^2) - (sind(2θ))/(2σy^2)
    c = (sind(θ)^2)/(2σx^2) + (cosd(θ)^2)/(2σy^2)
    gauss2d= amplitude*exp( -( a*((x-xo)^2) + 2b*(x-xo)*(y-yo) + c*((y-yo)^2)) )
    return(gauss2d)
end

#################################
function fit_beam(beam , dr)
    bs= size(beam)
    h= reshape(beam,(bs[1]*bs[2],))

    xy= zeros(Float64,bs[1]*bs[2], 2)

    for i in 1:bs[1]
        for j in 1:bs[2]
            xy[(i-1)*bs[1]+j,1]= i
        end
    end

    for i in 1:bs[1]
        for j in 1:bs[2]
            xy[(i-1)*bs[1]+j,2]= bs[2]+1-j
        end
    end

    @. multimodel(x, p) = gaussian2D(x[:,1], x[:,2], p[1], p[2] , p[3] , p[4] , p[5], p[6])
    p0 = [maximum(beam) , bs[1]/2 , bs[2]/2 , 1, 1, 0]
    f= curve_fit(multimodel, xy, h, p0)

    gaussian_sigma_to_fwhm=2.3548200450309493

    ## some bad fitting turn to be negative....
    bx= abs(f.param[4]*dr*gaussian_sigma_to_fwhm)
    by= abs(f.param[5]*dr*gaussian_sigma_to_fwhm)

    e= max(bx/by, by/bx)
    ar= sqrt(bx*by)
    sidelobe= maximum(abs.(f.resid * 100 / f.param[1]))

    b = synthbeam(bx, by,ar,e, sidelobe)

    return(b)
end

#################
### MRS in arcsec at 100 GHz
function calc_mrs(uv)
    minbl = minimum(uv[:,1].^2 .+ uv[:,2].^2)
    lmin= sqrt(minbl)
    ν= 100

    mrs= 37100 / (lmin*ν)
    return(mrs)
end


"""

```
icell3D
```

Returns the i,j,k coordinates of the cell from the coordinates of an atom and box properties

"""
function icell3D(x::AbstractVector, box::Box)
    i = trunc(Int, (x[1] + box.sides[1] / 2) / box.l[1]) + 1
    j = trunc(Int, (x[2] + box.sides[2] / 2) / box.l[2]) + 1
    k = trunc(Int, (x[3] + box.sides[3] / 2) / box.l[3]) + 1
    return i, j, k
end

#
# given the 1D representation of the index of a 3D cell, returns the 3 indexes of the
# coordinates in 3 dimensions
#

# if the number of cells in each dimension is different

function icell3D(n::AbstractVector{Int}, i1D)
    jk = i1D % (n[2] * n[3])
    if jk == 0
        jk = n[2] * n[3]
    end
    j, k = icell2D(n[2:3], jk)
    i = round(Int, (i1D - jk) / (n[2] * n[3])) + 1
    return i, j, k
end

# If the number of cells in each dimension is the same

function icell3D(n::Int, i1D::Int)
    n2 = n^2
    jk = i1D % n2
    if jk == 0
        jk = n2
    end
    j, k = icell2D(n, jk)
    i = round(Int, (i1D - jk) / (n2)) + 1
    return i, j, k
end

# Given the 1D represetation of the index of a 2D cell, returns the 2
# indexes of the 2D representation

# If the two dimensions are different

function icell2D(n::AbstractVector{Int}, i1D::Int)
    j = i1D % n[2]
    if j == 0
        j = n[2]
    end
    i = round(Int, (i1D - j) / n[2]) + 1
    return i, j
end

# If the two dimensions are the same

function icell2D(n::Int, i1D::Int)
    j = i1D % n
    if j == 0
        j = n
    end
    i = round(Int, (i1D - j) / n) + 1
    return i, j
end


using Decimals
using Test

@testset "Arithmetic" begin

@testset "Addition" begin
    @test Decimal(0.1) + 0.2 == 0.1 + Decimal(0.2) == Decimal(0.1) + Decimal(0.2) == Decimal(0.3)
    @test Decimal.([0.1 0.2]) .+ [0.3 0.1] == Decimal.([0.4 0.3])
    @test Decimal(2147483646) + Decimal(1) == Decimal(2147483647)
    @test Decimal(1,3,-2) + parse(Decimal, "0.2523410412138103") == Decimal(0,2223410412138103,-16)
end

@testset "Subtraction" begin
    @test Decimal(0.3) - 0.1 == 0.3 - Decimal(0.1)
    @test 0.3 - Decimal(0.1) == Decimal(0.3) - Decimal(0.1)
    @test Decimal(0.3) - Decimal(0.1) == Decimal(0.2)
    @test Decimal.([0.3 0.1]) .- [0.1 0.5] == Decimal.([0.2 -0.4])
end

@testset "Negation" begin
    @test -Decimal.([0.3 0.2]) == [-Decimal(0.3) -Decimal(0.2)]
    @test -Decimal(0.3) == zero(Decimal) - Decimal(0.3)
    @test iszero(decimal(12.1) - decimal(12.1))
end

@testset "Multiplication" begin
    @test Decimal(12.21) * Decimal(2.12) == Decimal(0,258852,-4)
    @test Decimal(12.2112543) * Decimal(2.121352) == Decimal(0,259043687318136,-13)
    @test Decimal(0.2) * 0.1 == 0.2 * Decimal(0.1)
    @test 0.2 * Decimal(0.1) == Decimal(0.02)
    @test Decimal(12.34) * 0.1234 == 12.34 * Decimal(0.1234)
    @test 12.34 * Decimal(0.1234) == Decimal(1.522756)
    @test Decimal(0.21084210) * -2 == -2 * Decimal(0.21084210)
    @test -2 * Decimal(0.21084210) == Decimal(-0.4216842)
    @test Decimal(0, 2, -1) * 0.0 == zero(Decimal)
    @test Decimal.([0.3, 0.6]) .* 5 == [Decimal(0.3)*5, Decimal(0.6)*5]
    @test one(Decimal) * 1 == Decimal(0, 1, 0)
end

@testset "Inversion" begin
    @test inv(Decimal(0, 1, -1)) == Decimal(0, 1, 1)
    @test inv(Decimal(0, 1, 1)) == Decimal(0, 1, -1)
    @test inv(Decimal(1, 2, -1)) == Decimal(1, 5, 0)
    @test inv(Decimal(1, 5, 0)) == Decimal(1, 2, -1)
    @test inv(Decimal(0, 2, -2)) == Decimal(0, 5, 1)
    @test inv(Decimal(0, 5, 1)) == Decimal(0, 2, -2)
    @test inv(Decimal(1, 4, -1)) == Decimal(1, 25, -1)
    @test inv(Decimal(1, 25, -1)) == Decimal(1, 4, -1)
end

@testset "Division" begin
    @test Decimal(0.2) / Decimal(0.1) == Decimal(2)
    @test Decimal(0.3) / Decimal(0.1) == Decimal(0,3,0)
    @test [Decimal(0.3) / Decimal(0.1), Decimal(0.6) / Decimal(0.1)] == [Decimal(0.3), Decimal(0.6)] ./ Decimal(0.1)
    @test [Decimal(0.3) / 0.1, Decimal(0.6) / 0.1] == [Decimal(0.3), Decimal(0.6)] ./ 0.1
end

end


# Check whether the test is being run as standalone
if ~@isdefined Test
    using Test
end

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

@testset "hierarchy matrix" begin
    # Makes sure the hierarchy matrix generator is working as expected for a select set of orbital systems.
    # Fully nested Keplerians
    ϵ_nested::Matrix{Float64} = [-1  1  0  0;
                                 -1 -1  1  0;
                                 -1 -1 -1  1;
                                 -1 -1 -1 -1;]
    @test hierarchy([4,1,1,1]) == ϵ_nested

    # Double Binary
    ϵ_double::Matrix{Float64} = [-1  1  0  0;
                                  0  0 -1  1;
                                 -1 -1  1  1;
                                 -1 -1 -1 -1;]
    @test hierarchy([4,2,1]) == ϵ_double

    # Double binary + top level body
    ϵ_asym::Matrix{Float64} = [-1  1  0  0  0;
                                0  0 -1  1  0;
                               -1 -1  1  1  0;
                               -1 -1 -1 -1  1;
                               -1 -1 -1 -1 -1;]
    @test hierarchy([5,2,1,1]) == ϵ_asym

    # Test of specific ambiguous initialization case of a 'symmetric fully-nested' (top) and the double binary + tertiary binary (bottom)
    ϵ_sym_full::Matrix{Float64} = [-1  1  0  0  0  0;
                                    0  0  0 -1  1  0;
                                   -1 -1  1  0  0  0;
                                    0  0  0 -1 -1  1;
                                   -1 -1 -1  1  1  1;
                                   -1 -1 -1 -1 -1 -1;]

    ϵ_asym_multi::Matrix{Float64}=[-1  1  0  0  0  0;
                                    0  0 -1  1  0  0;
                                   -1 -1  1  1  0  0;
                                    0  0  0  0 -1  1;
                                   -1 -1 -1 -1  1  1;
                                   -1 -1 -1 -1 -1 -1;]
    h_sym_full,h_asym_multi = hierarchy.([[6,2,2,1,2],[6,2,2,1]])

    @test ϵ_sym_full == h_sym_full
    @test ϵ_asym_multi == h_asym_multi
end


# Utilities
using LinearAlgebra: checksquare, opnorm

include("SymplecticForm.jl")

# calculate the 'nonSymplecticity".
nonSymplecticity(A::AbstractMatrix) = opnorm(A'*Ω*A - Ω) / opnorm(A)

# check if a matrix is symplectic, returning its non-symplecticity
function checkSymplectic(A::AbstractMatrix)
    A <: Real || throw(ArgumentError("matrix is not real-valued"))
    checksquare(A) % 2 == 0 || throw(DimensionMismatch("matrix is not even-dimensional"))
    return nonSymplecticity(A)
end
function checkSymplecticMatrix(A...)
    nonSymplecticities = []
    for a in A
        a <: Real || throw(ArgumentError("matrix is not real-valued"))
        checksquare(a) % 2 == 0 || throw(DimensionMismatch("matrix is not even-dimensional"))
        push!(nonSymplecticities, nonSymplecticity(a))
    end
    return nonSymplecticities
end

using LinearAlgebra

# Define the AbstractSymplecticMatrix data type
abstract type AbstractSymplecticMatrix{T} <: AbstractArray{T, 2} end
Base.size(A::AbstractSymplecticMatrix) = error("size not implemented")
Base.getindex(A::AbstractSymplecticMatrix, i::Int) = error("getindex not implemented")
Base.getindex(A::AbstractSymplecticMatrix, I::Vararg{Int, N}) where {N} = error("getindex not implemented")
Base.setindex!(A::AbstractSymplecticMatrix, v, i::Int) = error("setindex! not implemented")
Base.setindex!(A::AbstractSymplecticMatrix, v, I::Vararg{Int, N}) where {N} = error("setindex! not implemented")
# get a Q-quadrature from columns
getQCol(A::AbstractSymplecticMatrix, i::Int) = error("getQCol not implemented")
# get a P-quadrature from columns
getPCol(A::AbstractSymplecticMatrix, i::Int) = error("getPCol not implemented")
# get a Q-quadrature from rows
getQRow(A::AbstractSymplecticMatrix, i::Int) = error("getQCol not implemented")
# get a P-quadrature from rows
getPRow(A::AbstractSymplecticMatrix, i::Int) = error("getPCol not implemented")

Base.transpose(S::AbstractSymplecticMatrix) = error("transpose not implemented")
Base.adjoint(S::AbstractSymplecticMatrix) = error("adjoin not implemented")

Base.:-(S::AbstractSymplecticMatrix) = error("- not implemented")
Base.:+(S::AbstractSymplecticMatrix) = error("+ not implemented")
Base.:*(S1::AbstractSymplecticMatrix, S2::AbstractSymplecticMatrix, Ss::Vararg{AbstractSymplecticMatrix}) = error("* not implemented")
Base.inv(S::AbstractSymplecticMatrix) = error("inv not implemented")
Base.:^(S::AbstractSymplecticMatrix, p::Int) = error("^ not implemented")

# Define the SymplecticMatrix data type
struct SymplecticMatrix{T} <: AbstractSymplecticMatrix{T}
    matrix:: AbstractArray{T, 2}
    function SymplecticMatrix{T}(matrix:: AbstractArray{T, 2}) where T
        if !(T<:Real) return error("not real") end
        if !(LinearAlgebraUtilities.isSymplectic(matrix)) return error("not symplectic") end
        return new{T}(AbstractMatrix{T}(matrix))
    end
end

Base.size(S::SymplecticMatrix) = size(S.matrix)
Base.getindex(S::SymplecticMatrix, i::Int) = getindex(S.matrix, i)
Base.getindex(S::SymplecticMatrix, I::Vararg{Int, N}) where {N} = getindex(S.matrix, I)
getQCol(S::SymplecticMatrix, i::Int) = S.matrix[:, 2*i-1]
getPCol(S::SymplecticMatrix, i::Int) = S.matrix[:, 2*i]
getQRow(S::SymplecticMatrix, i::Int) = S.matrix[2*i-1, :]
getPRow(S::SymplecticMatrix, i::Int) = S.matrix[2*i, :]

SymplecticMatrix(M::AbstractArray{T, 2}) where T = SymplecticMatrix{T}(M)
SymplecticMatrix(S::SymplecticMatrix) = S

Base.transpose(S::SymplecticMatrix)::SymplecticMatrix = transpose( S.matrix )
Base.adjoint(S::SymplecticMatrix)::SymplecticMatrix = adjoint( S.matrix )

Base.:-(S::SymplecticMatrix)::SymplecticMatrix = Base.:-( S.matrix )
Base.:+(S::SymplecticMatrix)::SymplecticMatrix = Base.:+( S.matrix )
Base.:*(S1::SymplecticMatrix, S2::SymplecticMatrix, Ss::Vararg{SymplecticMatrix})::SymplecticMatrix = *(S1.matrix, S2.matrix,  [s.matrix for s in Ss]... )
Base.inv(S::SymplecticMatrix)::SymplecticMatrix = - Ω(S) * S' * Ω(S)
Base.:^(S::SymplecticMatrix, p::Int)::SymplecticMatrix = p < 0 ? inv(S).matrix ^ (-p)  : S.matrix^p

# Unary Minus and Unary Plus
⊗ = Base.kron

# Direct Sum
dsum(M1::AbstractMatrix, M2::AbstractMatrix)::AbstractMatrix = LinearAlgebraUtilities.dsum(M1, M2)
dsum(S1::SymplecticMatrix, S2::SymplecticMatrix, Ss::Vararg{SymplecticMatrix})::SymplecticMatrix = dsum(S1.matrix, S2.matrix,  [s.matrix for s in Ss]... )
⊕ = dsum

Base.Matrix(S::SymplecticMatrix)::Array{T, 2} = S.matrix
Base.convert(::Type{SymplecticMatrix}, S::SymplecticMatrix) = S
Base.convert(::Type{SymplecticMatrix}, M::AbstractArray{T, 2}) where T = SymplecticMatrix(M)
Base.convert(::Type{AbstractArray}, S::SymplecticMatrix) = Matrix(S)

# Omega(n::Int)::SymplecticMatrix = LinearAlgebraUtilities.Omega(n)
# Omega(S::SymplecticMatrix)::SymplecticMatrix = LinearAlgebraUtilities.Omega(S.matrix)
# Ω = Omega

isGeneric(S::SymplecticMatrix)::Bool = LinearAlgebraUtilities.isGeneric(S.matrix)

# Override functions in Base
# Base.Matrix(S::SymplecticMatrix)::AbstractMatrix = S.matrix
# Base.convert(::Type{SymplecticMatrix}, matrix::AbstractMatrix) = SymplecticMatrix(matrix)
# Base.convert(::Type{SymplecticMatrix}, S::SymplecticMatrix) = S
# Base.convert(::Type{AbstractMatrix}, S::SymplecticMatrix) = Matrix(S)
# Base.size(S::SymplecticMatrix, dim::Integer) = size(S.matrix, dim)
# Base.size(S::SymplecticMatrix) = size(S.matrix)
## Import !!!
# function Base.getindex(S::SymplecticMatrix{T}, i::Integer, j::Integer)::T where {T}
#     return S.matrix[i, j]
# end


"""
Staggered Bernoulli scheme is a generalization of Bernoulli scheme were each positive outcome is possible
for different number of steps.

The constructor takes probabilities of succeses p_1, ... p_N-1 and
n_1, ... n_N-1, number of steps outcome i is possible.
"""
struct StaggeredBernoulliScheme{N}<:DiscreteCountingProcess
    n::SVector{N, Int64}
    p_suc::SVector{N, Float64}

    function StaggeredBernoulliScheme(n::AbstractArray{<:Number}, p_suc::AbstractArray{<:Number})
        @assert 0<sum(p_suc)<=1
        # @assert issorted(n)
        new{length(p_suc)}(SA[n...], SA[p_suc...])
    end
end

# # Distributions.jl interface (https://juliastats.org/Distributions.jl/stable/extends/)
# TODO: implment also for staggered case
# Base.length(bs::BernoulliScheme) = 2

# """
# Get step-shift of the next success and the next succes type.
# Support for step-shift is 1,2, ....
# Support for success type is 1, 2, ..., N  (N is the number of success results)
# """
# function Distributions._rand!(rng::AbstractRNG, bs::BernoulliScheme, x::AbstractVector{T}) where T <: Real
#     x[2] = Distributions.rand(rng, Categorical(bs.p_suc./sum(bs.p_suc)))
#     x[1] = Distributions.rand(rng, Geometric(sum(bs.p_suc)))+1
#     return x
# end

"""
Get parameters of the process
"""
Distributions.params(bs::StaggeredBernoulliScheme) = (bs.n, bs.p_suc)

"""
Get outcome distribution for given index.
Support for the outcome distribution is 1...N+1. N+1 means failure.
Index can be in the range 1, ..., n (n is the number of random variables in the scheme).
"""
function Base.getindex(bs::StaggeredBernoulliScheme, i::Integer)
    if 1<=i<=maximum(bs.n)
        active = bs.n .>= i
        active_p = bs.p_suc .* active
        return Categorical([active_p..., 1-sum(active_p)])
    else
        throw(BoundsError(bs, i))
    end
end


#
# This example assumes below packages are available in your environment
#

using Turing, MCMCChains, Distributions, ParetoSmooth, Random
using StatsFuns

Random.seed!(5574)

data = rand(Normal(0, 1), 30)

@model function model(y)
    μ ~ Normal(0, 1)
    σ ~ truncated(Cauchy(0, 1), 0.0, Inf)

    y .~ Normal(μ, σ)
end

chains = sample(model(data), NUTS(1000, .65), MCMCThreads(), 1000, 3)

# method for MCMCChains
function pointwise_loglikes(chain::Chains, data, ll_fun)
    samples = Array(Chains(chain, :parameters).value)
    pointwise_loglikes(samples, data, ll_fun)
end

# generic method for arrays
function pointwise_loglikes(samples::Array{Float64,3}, data, ll_fun)
    n_data = length(data)
    n_samples, n_chains = size(samples)[[1,3]]
    pointwise_lls = fill(0.0, n_data, n_samples, n_chains)
    for c in 1:n_chains
        for s in 1:n_samples
            for d in 1:n_data
                pointwise_lls[d,s,c] = ll_fun(samples[s,:,c], data[d])
            end
        end
    end
    return pointwise_lls
end

function compute_loo(psis_output, pointwise_lls)
    dims = size(pointwise_lls)
    lwp = deepcopy(pointwise_lls)
    lwp += psis_output.weights;
    lwpt = reshape(lwp, dims[1], dims[2] * dims[3])';
    loos = reshape(logsumexp(lwpt; dims=1), size(lwpt, 2));
    return sum(loos)
end

# compute the pointwise log likelihoods where indices correspond to [data, sample, chain]
pointwise_lls = pointwise_loglikes(chains, data, (p,d)->logpdf(Normal(p...), d))

# compute the psis object
psis_output = psis(pointwise_lls)

# return loo based on Rob's example
loo = compute_loo(psis_output, pointwise_lls)


# matrix and transformations

"""
    getmatrix()

Get the current matrix. Returns an array of six float64 numbers:

- xx component of the affine transformation

- yx component of the affine transformation

- xy component of the affine transformation

- yy component of the affine transformation

- x0 translation component of the affine transformation

- y0 translation component of the affine transformation

Some basic matrix transforms:

- translate

`transform([1, 0, 0, 1, dx, dy])` shifts by `dx`, `dy`

- scale

`transform([fx 0 0 fy 0 0])` scales by `fx` and `fy`

- rotate

`transform([cos(a), -sin(a), sin(a), cos(a), 0, 0])` rotates around to `a` radians

    rotate around O: [c -s s c 0 0]

- shear

`transform([1 0 a 1 0 0])` shears in x direction by `a`
-  shear in y: [1  B 0  1 0 0]

- x-skew

`transform([1, 0, tan(a), 1, 0, 0])` skews in x by `a`

- y-skew

`transform([1, tan(a), 0, 1, 0, 0])` skews in y by `a`

- flip

`transform([fx, 0, 0, fy, centerx * (1 - fx), centery * (fy-1)])` flips with center at `centerx`/`centery`

- reflect

`transform([1 0 0 -1 0 0])` reflects in xaxis

`transform([-1 0 0 1 0 0])` reflects in yaxis

When a drawing is first created, the matrix looks like this:

    getmatrix() = [1.0, 0.0, 0.0, 1.0, 0.0, 0.0]

When the origin is moved to 400/400, it looks like this:

    getmatrix() = [1.0, 0.0, 0.0, 1.0, 400.0, 400.0]

To reset the matrix to the original:

    setmatrix([1.0, 0.0, 0.0, 1.0, 0.0, 0.0])

"""
function getmatrix()
    gm = Cairo.get_matrix(get_current_cr())
    return([gm.xx, gm.yx, gm.xy, gm.yy, gm.x0, gm.y0])
end

"""
    setmatrix(m::Array)

Change the current matrix to matrix `m`. Use `getmatrix()` to get the current matrix.
"""
function setmatrix(m::Array)
    if eltype(m) != Float64
        m = map(Float64, m)
    end
    # some matrices make Cairo freak out and need reset. Not sure what the rules are yet…
    if length(m) < 6
        throw("didn't like that matrix $m: not enough values")
    elseif count(!iszero, m) == 0
        throw("didn't like that matrix $m: too many zeroes")
    else
        cm = Cairo.CairoMatrix(m[1], m[2], m[3], m[4], m[5], m[6])
        Cairo.set_matrix(get_current_cr(), cm)
    end
end

"""
    transform(a::Array)

Modify the current matrix by multiplying it by matrix `a`.

For example, to skew the current state by 45 degrees in x and move by 20 in y direction:

    transform([1, 0, tand(45), 1, 0, 20])

Use `getmatrix()` to get the current matrix.
"""
function transform(a::Array)
    b = Cairo.get_matrix(get_current_cr())
    setmatrix([
        (a[1] * b.xx)  + a[2]  * b.xy,             # xx
        (a[1] * b.yx)  + a[2]  * b.yy,             # yx
        (a[3] * b.xx)  + a[4]  * b.xy,             # xy
        (a[3] * b.yx)  + a[4]  * b.yy,             # yy
        (a[5] * b.xx)  + (a[6] * b.xy) + b.x0,     # x0
        (a[5] * b.yx)  + (a[6] * b.yy) + b.y0      # y0
    ])
end

"""
    rotationmatrix(a)

Return a 3x3 Julia matrix that will apply a rotation through `a` radians.
"""
function rotationmatrix(a)
    return ([cos(a)  -sin(a)  0.0 ;
             sin(a)   cos(a)  0.0 ;
             0.0         0.0  1.0 ])
end

"""
    translationmatrix(x, y)

Return a 3x3 Julia matrix that will apply a translation in `x` and `y`.
"""
function translationmatrix(x, y)
    return ([1.0   0.0   x ;
             0.0   1.0   y ;
             0.0   0.0   1.0 ])
end

"""
    scalingmatrix(sx, sy)

Return a 3x3 Julia matrix that will apply a scaling by `sx` and `sy`.
"""
function scalingmatrix(sx, sy)
    return ([sx 0.0  0.0 ;
           0.0   sy  0.0 ;
           0.0  0.0  1.0])
end

"""
    cairotojuliamatrix(c)

Return a 3x3 Julia matrix that's the equivalent of the six-element matrix in `c`.
"""
function cairotojuliamatrix(c::Array)
    return [c[1] c[3] c[5] ; c[2] c[4] c[6] ; 0.0 0.0 1.0]
end

"""
    juliatocairomatrix(c)

Return a six-element matrix that's the equivalent of the 3x3 Julia matrix in `c`.
"""
function juliatocairomatrix(c::Matrix)
    return [c[1] c[2] c[4] c[5] c[7] c[8]]
end

"""
    getrotation(R::Matrix)
    getrotation()

Get the rotation of a Julia 3x3 matrix, or the current Luxor rotation.

```math
\\begin{bmatrix}
a & b & tx \\\\
c & d & ty \\\\
0 & 0 & 1  \\\\
\\end{bmatrix}
```

The rotation angle is `atan(-b, a)` or `atan(c, d)`.
"""
function getrotation(R::Matrix)
    # t = atan(-R[4], R[1]) # should be the same as:
    t = atan(R[2], R[5])
    return mod2pi(t)
end

function getrotation()
    getrotation(cairotojuliamatrix(getmatrix()))
end

"""
    getscale(R::Matrix)
    getscale()

Get the current scale of a Julia 3x3 matrix, or the current Luxor scale.

Returns a tuple of x and y values.
"""
function getscale(R::Matrix)
    sx = hypot(R[1], R[2])
    sy = hypot(R[4], R[5])
    return (sx, sy)
end

function getscale()
    getscale(cairotojuliamatrix(getmatrix()))
end

"""
    gettranslation(R::Matrix)
    gettranslation()

Get the current translation of a Julia 3x3 matrix, or the current Luxor translation.

Returns a tuple of x and y values.
"""
function gettranslation(R::Matrix)
    return (R[7], R[8])
end

function gettranslation()
    gettranslation(cairotojuliamatrix(getmatrix()))
end


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 1c3ff5c4-90d5-11eb-1459-7d7d43335a80
using Yao, YaoPlots, YaoExtensions, PlutoUI, Flux, Plots

# ╔═╡ e9342104-91cf-11eb-01d7-ffd747975c7c
using LinearAlgebra

# ╔═╡ aa09eef8-8eac-11eb-00e6-6794d5c9e0b5
md"""
# Quantum Computing Basics

In case you are not familar with some basic concept of quantum computing or quantum algorithms, we will have a quick and short introduction for quantum computing.
"""

# ╔═╡ 22f46d6e-90d5-11eb-3094-e7015012995c
md"""

## Qubits

The computation states are often represented as qubits in quantum computing model. A qubit is a quantum state that has two basis states ``|0⟩`` and ``|1⟩``

## Quantum Registers

A quantum register is an abstraction of the qubit memory, it usualy contains many qubits.

## Quantum Gates

A quantum gate is a unitary operator that can be applied on the (quantum) register. we usually visualize it as a square block with two wires, e.g the following is a Pauli X gate.
"""

# ╔═╡ 920da102-90d5-11eb-0437-495807d8b345
YaoPlots.plot(X)

# ╔═╡ a5d5d56a-90d5-11eb-1e4a-8bf893d8534a
md"""
## Parameterized Quantum Gates

There are some quantum gates that can take classical parameters, e.g the rotation gates on X, Y, Z or other directions.
"""

# ╔═╡ be1f51c6-90d5-11eb-3db1-99708f112467
YaoPlots.plot(Rx(0.1))

# ╔═╡ 4f4a700e-90d6-11eb-3058-b1809082853a
md"""
## Controlled Gate

The gate operation can be controlled by another qubits, we usually represent them as the following diagram
"""

# ╔═╡ 6b2079e2-90d6-11eb-330c-0194f2906114
YaoPlots.plot(control(2, 1, 2=>Rx(0.1)))

# ╔═╡ 7854398c-90d6-11eb-2aad-b75727ab7b79
md"for controlled Pauli X gate, we also call it the CNOT gate, it can be plotted in a simpler way"

# ╔═╡ 9003bb5c-90d6-11eb-102e-f1954fe71ec6
YaoPlots.plot(cnot(2, 1, 2))

# ╔═╡ debe78d2-90d5-11eb-1526-87f514b030bf
md"""
## Quantum Circuits & Quantum Algorithms

Now with the concept of register and gates, we can define quantum circuits. The quantum circuits are composition of gates and measurements.

e.g the following is a quantum fourier transformation circuit, it is defined using hadamard gate ``H`` and control phase shift gate.
"""

# ╔═╡ 0a379552-90d6-11eb-1a07-130ebf3c1264
begin
	A(i, j) = control(i, j=>shift(2π/(1<<(i-j+1))))
	B(n, k) = chain(n, j==k ? put(k=>H) : A(j, k) for j in k:n)
	qft(n) = chain(B(n, k) for k in 1:n)
	YaoPlots.plot(qft(3))
end

# ╔═╡ e4f1f268-91ca-11eb-231a-712cb8cf4a55
md"""
## Program with Yao

checkout the website: [yaoquantum.org](https://yaoquantum.org/)
"""

# ╔═╡ 0478efb2-91cb-11eb-034e-0b0c3be90bb7
html"""
<iframe src="https://yaoquantum.org/" width=100% height=1000></iframe>
"""

# ╔═╡ 2a546ef0-91cb-11eb-17fd-4b8c74878e6a
md"""
### the quantum blocks

Yao represents the quantum circuits as blocks

![](https://github.com/Roger-luo/PIML/blob/master/notebooks/assets/qblock.png?raw=true)
"""

# ╔═╡ 76b6bf64-91cb-11eb-3c39-576a66e70cf4
qft(3)

# ╔═╡ c75b0244-8eac-11eb-2528-b9f59f61634d
md"""
## Variational Quantum Algorithms

*in some context, this is also called quantum machine learning.*

Traditionally, a quantum algorithm usually requires a well designed circuits that provides an exponential speed up comparing to its classical counterparts, e.g the well known Shor algorithm.

However, designing and finding such algorithms can be extremely hard, and may not be possible implemented on near term quantum devices that is noisy.
"""

# ╔═╡ e1345ee8-8eac-11eb-3c22-0db185142833
md"""
## Differentiating Quantum Circuits: the forward mode

In literature, this is usually refered as the **Parameter Shift Rule**, it was first discovered in [Quantum Circuit Learning, arXiv: 1803.00745](https://arxiv.org/pdf/1803.00745.pdf)

we will now derive the rules, assume the unitary of our parameterized quantum circuit is ``U(\mathbf{θ}) = \prod_{i=1}^n U(θ_i)``, and for convenience, we will denote ``U_{j:k} = U_j⋯U_k``

then for an observable ``B``, the expectation value is ``⟨B(\mathbf{θ})⟩ = Tr(B U_{l:1}ρU_{l:1}^{†})``

now if we assume our unitary is generated by Pauli strings (which is not always true), we have ``U_j(θ) = exp(-iθ_j P_j/2)``.

The gradient of ``⟨B(\mathbf{θ})⟩`` can be calculated as following

```math
\frac{\partial ⟨B⟩}{\partial θ_j} = -\frac{i}{2} (B U_{l:j} [P_j, U_{j-1:1}ρU_{j-1:1}^{†}] U_{l:j}^{†})
```

and because we have the following relationship for arbitrary ``ρ``

```math
[P_j, ρ] = i[U_j(\frac{π}{2})ρ U_j^{†}(\frac{π}{2}) - U_j(-\frac{π}{2})ρ U_j(-\frac{π}{2})]
```

thus the gradient can be evaluated by

```math
\begin{aligned}
\frac{\partial B}{\partial θ_j} =& \frac{1}{2} (Tr[B U_{l:j+1}U(θ_j+\frac{π}{2})U_{j-1:1}ρU_{j-1:1}^{†}U_j^{†}(θ_j+\frac{π}{2})U_{l:j+1}^{†}] -\\
&Tr[B U_{l:j+1}U(θ_j-\frac{π}{2})U_{j-1:1}ρU_{j-1:1}^{†}U_j^{†}(θ_j-\frac{π}{2})U_{l:j+1}^{†}])
\end{aligned}
```

which means

```math
\frac{\partial B}{\partial θ_j} = \frac{1}{2} (⟨B(\theta_j + \frac{π}{2})⟩ - ⟨B(\theta_j - \frac{π}{2})⟩)
```
"""

# ╔═╡ f8270b02-8eac-11eb-2ab3-3323e01f4f62
md"""
## Differentiating Quantum Circuits: the reverse mode

is it possible to differentiate quantum circuits using reverse mode?
"""

# ╔═╡ c97fffa8-91d0-11eb-0a0c-092545e368f8
md"""
Only for simulation

![](https://github.com/Roger-luo/PIML/blob/master/notebooks/assets/yao-reverse-ad.png?raw=true)
"""

# ╔═╡ 3c919a50-8ead-11eb-058f-27eceb1deeaa
md"""
## Implementing A Simple Variational Quantum Eigensolver (VQE)

You have learned how variational Monte Carlo works, now let's:

1. use a quantum circuit as our ansatz
2. directly calculate the energy by measure the Hamiltonian
3. use gradient based optimization to optimize the energy
"""

# ╔═╡ 68e75bba-91ca-11eb-128f-1d6b0b7cbb78
md"first let's create a variational circuit using the `variational_circuit` function"

# ╔═╡ 4422a0fa-91ca-11eb-3bda-f37b2dab68b3
YaoPlots.plot(variational_circuit(5, 3))

# ╔═╡ 82f2d962-91ca-11eb-2e66-597a90047e87
md"then create a quantum register at state ``|00⋯00⟩``"

# ╔═╡ 895600fc-91ca-11eb-3d03-e5aedfa44576
r = zero_state(5)

# ╔═╡ 940657d8-91ca-11eb-33fd-63a019b4ebfc
md"we can evaluate this circuit on the emulator"

# ╔═╡ a0527026-91ca-11eb-0604-21a298350c8d
r |> variational_circuit(5, 3)

# ╔═╡ b8a71bf4-91ca-11eb-2f81-eb67cfeeada0
md"""
we can now create a 1-D heisenberg hamiltonian

```math
\sum_i \mathbf{S}_i \mathbf{S}_{i+1}
```
"""

# ╔═╡ da9ba108-91ca-11eb-28a4-755cf76ea605
h = heisenberg(5)

# ╔═╡ b3cc56a2-91cb-11eb-3736-87df27d04545
md"we can calculate the expectation value of this hamiltonian"

# ╔═╡ 848d146c-91cb-11eb-17cc-87a1aa3c0d76
expect(heisenberg(5), r)

# ╔═╡ bd3c4bca-91cb-11eb-29cf-657eb95db73d
@doc expect

# ╔═╡ 9fe3db9c-91cb-11eb-1888-cf7dd2aa3aa1
expect(heisenberg(5), r=>variational_circuit(5, 3))

# ╔═╡ a6c13932-91cb-11eb-3822-7da569e520e0
reg, ∇θ = expect'(heisenberg(5), zero_state(5)=>variational_circuit(5, 3))

# ╔═╡ cbacc1f8-91cb-11eb-316c-37e3da59855c
md"we can use the optimizers from the ML framework `Flux`"

# ╔═╡ dddcbf20-91cd-11eb-2d73-2daa1643074d
opt = ADAM()

# ╔═╡ 18d83a62-91d0-11eb-2d8d-2bc2e69785e8
md"depth=$(@bind depth Slider(1:10; show_value=true))"

# ╔═╡ 12aff314-91ce-11eb-36d3-772d32cb8f3e
circuit = variational_circuit(5, depth)

# ╔═╡ e172b6f0-91cd-11eb-35be-2d44f6f26a80
begin
	history = []
	θ = rand(nparameters(circuit))
	for _ in 1:1000
		reg, ∇θ = expect'(heisenberg(5), zero_state(5)=>dispatch!(circuit, θ))
		Flux.Optimise.update!(opt, θ, ∇θ)
		current = expect(heisenberg(5), zero_state(5)=>dispatch!(circuit, θ))
		push!(history, real(current))
	end
	Plots.plot(history)
end

# ╔═╡ ec0935e0-91cf-11eb-1837-334e0e92edee
minimum(eigvals(Matrix(heisenberg(5))))

# ╔═╡ ad0071da-8ead-11eb-31e8-b94b8bff7635
md"""
## Placement Ads: Google Summer of Code 2022

https://summerofcode.withgoogle.com/
"""

# ╔═╡ Cell order:
# ╠═1c3ff5c4-90d5-11eb-1459-7d7d43335a80
# ╟─aa09eef8-8eac-11eb-00e6-6794d5c9e0b5
# ╟─22f46d6e-90d5-11eb-3094-e7015012995c
# ╟─920da102-90d5-11eb-0437-495807d8b345
# ╟─a5d5d56a-90d5-11eb-1e4a-8bf893d8534a
# ╟─be1f51c6-90d5-11eb-3db1-99708f112467
# ╟─4f4a700e-90d6-11eb-3058-b1809082853a
# ╟─6b2079e2-90d6-11eb-330c-0194f2906114
# ╟─7854398c-90d6-11eb-2aad-b75727ab7b79
# ╟─9003bb5c-90d6-11eb-102e-f1954fe71ec6
# ╟─debe78d2-90d5-11eb-1526-87f514b030bf
# ╟─0a379552-90d6-11eb-1a07-130ebf3c1264
# ╟─e4f1f268-91ca-11eb-231a-712cb8cf4a55
# ╟─0478efb2-91cb-11eb-034e-0b0c3be90bb7
# ╟─2a546ef0-91cb-11eb-17fd-4b8c74878e6a
# ╠═76b6bf64-91cb-11eb-3c39-576a66e70cf4
# ╟─c75b0244-8eac-11eb-2528-b9f59f61634d
# ╟─e1345ee8-8eac-11eb-3c22-0db185142833
# ╟─f8270b02-8eac-11eb-2ab3-3323e01f4f62
# ╟─c97fffa8-91d0-11eb-0a0c-092545e368f8
# ╟─3c919a50-8ead-11eb-058f-27eceb1deeaa
# ╟─68e75bba-91ca-11eb-128f-1d6b0b7cbb78
# ╠═4422a0fa-91ca-11eb-3bda-f37b2dab68b3
# ╟─82f2d962-91ca-11eb-2e66-597a90047e87
# ╠═895600fc-91ca-11eb-3d03-e5aedfa44576
# ╟─940657d8-91ca-11eb-33fd-63a019b4ebfc
# ╠═a0527026-91ca-11eb-0604-21a298350c8d
# ╟─b8a71bf4-91ca-11eb-2f81-eb67cfeeada0
# ╠═da9ba108-91ca-11eb-28a4-755cf76ea605
# ╟─b3cc56a2-91cb-11eb-3736-87df27d04545
# ╠═848d146c-91cb-11eb-17cc-87a1aa3c0d76
# ╠═bd3c4bca-91cb-11eb-29cf-657eb95db73d
# ╠═9fe3db9c-91cb-11eb-1888-cf7dd2aa3aa1
# ╠═a6c13932-91cb-11eb-3822-7da569e520e0
# ╟─cbacc1f8-91cb-11eb-316c-37e3da59855c
# ╠═dddcbf20-91cd-11eb-2d73-2daa1643074d
# ╠═18d83a62-91d0-11eb-2d8d-2bc2e69785e8
# ╠═12aff314-91ce-11eb-36d3-772d32cb8f3e
# ╠═e172b6f0-91cd-11eb-35be-2d44f6f26a80
# ╠═e9342104-91cf-11eb-01d7-ffd747975c7c
# ╠═ec0935e0-91cf-11eb-1837-334e0e92edee
# ╟─ad0071da-8ead-11eb-31e8-b94b8bff7635


# # Car
#
# function car_dynamics!(ẋ::AbstractVector,x::AbstractVector,u::AbstractVector)
#     ẋ[1] = u[1]*cos(x[3])
#     ẋ[2] = u[1]*sin(x[3])
#     ẋ[3] = u[2]
#     return nothing
# end
# n,m = 3,2
#
# car = Model(car_dynamics!,n,m)

@with_kw struct DubinsCar <: AbstractModel
    radius::Float64 = 0.175  # [m] radius of a Roomba
end
Base.size(::DubinsCar) = 3,2

function dynamics(::DubinsCar,x,u)
    ẋ = @SVector [u[1]*cos(x[3]),
                  u[1]*sin(x[3]),
                  u[2]]
end

Base.position(::DubinsCar, x::SVector) = @SVector [x[1], x[2], 0.0]
orientation(::DubinsCar, x::SVector) = expm(x[3]*@SVector [0,0,1.])

# using Plots
# function Plots.plot(model::DubinsCar, Z::Traj; kwargs...)
#     x = [z.z[1] for z in Z]
#     y = [z.z[2] for z in Z]
#     Plots.plot(x, y, xlabel="x position", ylabel="y position", markershape=:circle;
#         kwargs...)
# end
# function Plots.plot!(model::DubinsCar, Z::Traj; kwargs...)
#     x = [z.z[1] for z in Z]
#     y = [z.z[2] for z in Z]
#     Plots.plot!(x, y, xlabel="x position", ylabel="y position", markershape=:circle;
#         kwargs...)
# end


struct CYBERTRUCK{T} <: Model
    n::Int
    m::Int
    d::Int

    mass::T
    J::T
    μ::T
    g::T

    nq
    nu
    nc
    nf
    nb
	ns

    idx_u
    idx_λ
    idx_b
    idx_ψ
    idx_η
    idx_s
end

# Dimensions
nq = 4 # configuration dim
nu = 2 # control dim
nc = 1 # number of contact points
nf = 4 # number of faces for friction cone pyramid
nb = nc * nf
ns = 1

# Parameters
g = 9.81   # gravity
μ = 0.5    # coefficient of friction
mass = 1.0 # body mass
J = 1.0    # body inertia

n = 2 * nq
m = nu + nc + nb + nc + nb + ns
d = 0

idx_u = (1:nu)
idx_λ = nu .+ (1:nc)
idx_b = nu + nc .+ (1:nb)
idx_ψ = nu + nc + nb .+ (1:nc)
idx_η = nu + nc + nb + nc .+ (1:nb)
idx_s = nu + nc + nb + nc + nb .+ (1:ns)

# Methods
M_func(model::CYBERTRUCK, q) = Diagonal(@SVector [model.mass, model.mass, model.mass, model.J])
G_func(model::CYBERTRUCK, q) = @SVector [0.0, 0.0, model.mass * model.g, 0.0]

function ϕ_func(::CYBERTRUCK, q)
    return @SVector [q[3]]
end

N_func(::CYBERTRUCK, q) = @SMatrix [0.0 0.0 1.0 0.0]

P_func(::CYBERTRUCK, q) = @SMatrix [1.0 0.0 0.0 0.0;
                                     -1.0 0.0 0.0 0.0;
                                     0.0 1.0 0.0 0.0;
                                     0.0 -1.0 0.0 0.0]

B_func(::CYBERTRUCK, q) = @SMatrix [cos(q[4]) sin(q[4]) 0.0 0.0;
                                    0.0 0.0 0.0 1.0]

function friction_cone(model::CYBERTRUCK, u)
	λ = u[model.idx_λ]
	b = u[model.idx_b]
	return @SVector [model.μ * λ[1] - sum(b)]
end

function maximum_dissipation(model::CYBERTRUCK, x⁺, u, h)
	q3 = x⁺[model.nq .+ (1:model.nq)]
	q2 = x⁺[1:model.nq]
	ψ = u[model.idx_ψ]
	ψ_stack = ψ[1] * ones(4)
	η = u[model.idx_η]
	return P_func(model, q3) * (q3 - q2) / h + ψ_stack - η
end

function fd(model::CYBERTRUCK, x⁺, x, u, w, h, t)
	q3 = view(x⁺, model.nq .+ (1:model.nq))
	q2⁺ = view(x⁺, 1:model.nq)
	q2⁻ = view(x, model.nq .+ (1:model.nq))
	q1 = view(x, 1:model.nq)
	u_ctrl = view(u, model.idx_u)
	λ = view(u, model.idx_λ)
	b = view(u, model.idx_b)

    [q2⁺ - q2⁻;
    ((1.0 / h) * (M_func(model, q1) * (SVector{4}(q2⁺) - SVector{4}(q1))
    - M_func(model, q2⁺) * (SVector{4}(q3) - SVector{4}(q2⁺)))
    + transpose(B_func(model, q3)) * SVector{2}(u_ctrl)
    + transpose(N_func(model, q3)) * SVector{1}(λ)
    + transpose(P_func(model, q3)) * SVector{4}(b)
    - h * G_func(model, q2⁺))]
end

model = CYBERTRUCK(n, m, d,
				   mass, J, μ, g,
				   nq,
  				   nu,
  				   nc,
  				   nf,
  				   nb,
  				   ns,
  				   idx_u,
  				   idx_λ,
  				   idx_b,
  				   idx_ψ,
  				   idx_η,
  				   idx_s)

function visualize!(vis, model::CYBERTRUCK, q;
		Δt = 0.1)

    obj_path = joinpath(pwd(), "src/models/cybertruck/cybertruck.obj")
    mtl_path = joinpath(pwd(), "src/models/cybertruck/cybertruck.mtl")

    ctm = ModifiedMeshFileObject(obj_path, mtl_path, scale=0.1)
    setobject!(vis["cybertruck"], ctm)
    settransform!(vis["cybertruck"], LinearMap(RotZ(pi) * RotX(pi / 2.0)))

    anim = MeshCat.Animation(convert(Int,floor(1.0 / Δt)))

    for t = 1:length(q)
        MeshCat.atframe(anim, t) do
            settransform!(vis["cybertruck"], compose(Translation(q[t][1:3]...),
				LinearMap(RotZ(q[t][4] + pi) * RotY(0.0) * RotX(pi / 2.0))))
        end
    end
    # settransform!(vis["/Cameras/default"], compose(Translation(-1, -1, 0),LinearMap(RotZ(pi/2))))
    MeshCat.setanimation!(vis, anim)
end


using CellListMap
using StaticArrays
using Test

@testset "CellListMap.jl" begin

  if Threads.nthreads() == 1
    println("""

       WARNING: Ideally, run a multi-threaded test to check the parallel versions.

    """)
  end

  # Number of particles, sides and cutoff
  sides = @SVector [250,250,250]
  cutoff = 10.
  box = Box(sides,cutoff)

  # Particle positions
  N = 2000
  x = [ sides .* rand(SVector{3,Float64}) for i in 1:N ]
  # Add some pathological coordinates
  x[1] = -sides/2
  x[10] = sides
  x[100] = @SVector [sides[1]/2, -sides[2]/2, 2*sides[3]]
  y = [ sides .* rand(SVector{3,Float64}) for i in 1:N ]

  # Initialize auxiliary linked lists
  cl = CellList(x,box)

  # Function to be evalulated for each pair: sum of displacements on x
  f(x,y,avg_dx) = avg_dx + abs(x[1] - y[1])

  naive = CellListMap.map_naive!((x,y,i,j,d2,avg_dx) -> f(x,y,avg_dx),0.,x,box)
  @test map_pairwise!((x,y,i,j,d2,avg_dx) -> f(x,y,avg_dx),0.,box,cl,parallel=true) ≈ naive
  @test map_pairwise!((x,y,i,j,d2,avg_dx) -> f(x,y,avg_dx),0.,box,cl,parallel=false) ≈ naive

  # Check if changing lcell breaks something
  box = Box(sides,cutoff,lcell=1); cl = CellList(x,box)
  @test map_pairwise!((x,y,i,j,d2,avg_dx) -> f(x,y,avg_dx),0.,box,cl,parallel=true) ≈ naive
  box = Box(sides,cutoff,lcell=2); cl = CellList(x,box)
  @test map_pairwise!((x,y,i,j,d2,avg_dx) -> f(x,y,avg_dx),0.,box,cl,parallel=true) ≈ naive
  box = Box(sides,cutoff,lcell=3); cl = CellList(x,box)
  @test map_pairwise!((x,y,i,j,d2,avg_dx) -> f(x,y,avg_dx),0.,box,cl,parallel=true) ≈ naive
  box = Box(sides,cutoff,lcell=5); cl = CellList(x,box)
  @test map_pairwise!((x,y,i,j,d2,avg_dx) -> f(x,y,avg_dx),0.,box,cl,parallel=true) ≈ naive

  # Function to be evalulated for each pair: build distance histogram
  function build_histogram!(d2,hist)
    d = sqrt(d2)
    ibin = floor(Int,d) + 1
    hist[ibin] += 1
    return hist
  end

  naive = CellListMap.map_naive!(
    (x,y,i,j,d2,hist) -> build_histogram!(d2,hist),
    zeros(Int,10),x,box
  )
  @test map_pairwise!(
    (x,y,i,j,d2,hist) -> build_histogram!(d2,hist),
    zeros(Int,10),box,cl,parallel=true
  ) ≈ naive
  @test map_pairwise!(
    (x,y,i,j,d2,hist) -> build_histogram!(d2,hist),
    zeros(Int,10),box,cl,parallel=false
  ) ≈ naive

  # Function to be evalulated for each pair: build distance histogram
  function potential(i,j,d2,u,mass)
    d = sqrt(d2)
    u = u - 9.8*mass[i]*mass[j]/d
    return u
  end

  # Run pairwise computation
  mass = rand(N)
  naive = CellListMap.map_naive!((x,y,i,j,d2,u) -> potential(i,j,d2,u,mass),0.0,x,box)
  @test map_pairwise!((x,y,i,j,d2,u) -> potential(i,j,d2,u,mass),0.0,box,cl,parallel=true) ≈ naive
  @test map_pairwise!((x,y,i,j,d2,u) -> potential(i,j,d2,u,mass),0.0,box,cl,parallel=false) ≈ naive

  # Function to be evalulated for each pair: build distance histogram
  function calc_forces!(x,y,i,j,d2,mass,forces)
    G = 9.8*mass[i]*mass[j]/d2
    d = sqrt(d2)
    df = (G/d)*(x - y)
    forces[i] = forces[i] - df
    forces[j] = forces[j] + df
    return forces
  end

  # forces
  forces = [ zeros(SVector{3,Float64}) for i in 1:N ]

  # Run pairwise computation
  naive = CellListMap.map_naive!(
    (x,y,i,j,d2,forces) -> calc_forces!(x,y,i,j,d2,mass,forces),
    copy(forces),x,box
  )
  @test map_pairwise!(
    (x,y,i,j,d2,forces) -> calc_forces!(x,y,i,j,d2,mass,forces),
    copy(forces),box,cl,parallel=true
  ) ≈ naive
  @test map_pairwise!(
    (x,y,i,j,d2,forces) -> calc_forces!(x,y,i,j,d2,mass,forces),
    copy(forces),box,cl,parallel=false
  ) ≈ naive

  # Compute some properties of disjoint sets
  box = Box(sides,cutoff,lcell=1)
  cl = CellList(x,y,box)

  naive = CellListMap.map_naive!((x,y,i,j,d2,u) -> potential(i,j,d2,u,mass),0.0,x,y,box)
  @test map_pairwise!(
    (x,y,i,j,d2,u) -> potential(i,j,d2,u,mass),
    0.0,box,cl,parallel=false
  ) ≈ naive
  @test map_pairwise!(
    (x,y,i,j,d2,u) -> potential(i,j,d2,u,mass),
    0.0,box,cl,parallel=true
  ) ≈ naive

  # Check different lcell
  box = Box(sides,cutoff,lcell=3)
  cl = CellList(x,y,box)
  @test map_pairwise!(
    (x,y,i,j,d2,u) -> potential(i,j,d2,u,mass),
    0.0,box,cl,parallel=false
  ) ≈ naive
  @test map_pairwise!(
    (x,y,i,j,d2,u) -> potential(i,j,d2,u,mass),
    0.0,box,cl,parallel=true
  ) ≈ naive

  # Test the examples, to check further if the parallelization didn't break something
  N = 100_000
  x = [ sides .* rand(SVector{3,Float64}) for i in 1:N ]
  y = [ sides .* rand(SVector{3,Float64}) for i in 1:N ]
  @test CellListMap.test1(parallel=true,x=x)[2] ≈ CellListMap.test1(parallel=false,x=x)[2]
  @test CellListMap.test2(parallel=true,x=x) ≈ CellListMap.test2(parallel=false,x=x)
  @test CellListMap.test3(parallel=true,x=x) ≈ CellListMap.test3(parallel=false,x=x)
  @test CellListMap.test4(parallel=true,x=x) ≈ CellListMap.test4(parallel=false,x=x)
  @test count(CellListMap.test5(parallel=true,x=x,y=y) .≈ CellListMap.test5(parallel=false,x=x,y=y)) == 3

  pairs1 = sort!(CellListMap.test7(parallel=true,x=x),by=x->x[3])
  pairs2 = sort!(CellListMap.test7(parallel=false,x=x),by=x->x[3])
  @test count([ count(pairs1[i] .≈ pairs2[i]) == 3 for i in 1:length(pairs1) ]) == length(pairs1)

  x = [ sides .* rand(SVector{3,Float64}) for i in 1:1_500 ]
  y = [ sides .* rand(SVector{3,Float64}) for i in 1:1_500_000 ]
  @test count(CellListMap.test6(parallel=true,x=x,y=y) .≈ CellListMap.test6(parallel=false,x=x,y=y)) == 3

  # invert x and y to test swap
  ixy = CellListMap.test6(parallel=false,x=x,y=y)
  iyx = CellListMap.test6(parallel=false,x=y,y=x)
  @test ( ixy[1] == iyx[2] && ixy[2] == iyx[1] && ixy[3] ≈ iyx[3] )
  ixy = CellListMap.test6(parallel=true,x=x,y=y)
  iyx = CellListMap.test6(parallel=true,x=y,y=x)
  @test ( ixy[1] == iyx[2] && ixy[2] == iyx[1] && ixy[3] ≈ iyx[3] )

#  # Test resizing of the cell lists
#  x = [ rand(SVector{3,Float64}) for i in 1:1000 ]
#  box = Box([0.83,0.41,0.97],0.1)
#  cl = CellList(x,box)
#  @test length(cl.cwp) == 924
#
#  box = Box([0.33,0.41,0.97],0.1)
#  cl = UpdateCellList!(x,box,cl)
#  @test length(cl.cwp) == 924
#
#  box = Box([0.83,0.81,0.97],0.1)
#  cl = UpdateCellList!(x,box,cl)
#  @test length(cl.cwp) == 1598
#
#  x .= 0.9*x
#  cl = UpdateCellList!(x,box,cl)
#  @test length(cl.cwp) == 1598
#
#  x .= 1.2*x
#  cl = UpdateCellList!(x,box,cl)
#  @test length(cl.cwp) == 1598
#
#  box = Box([0.83,0.81,0.97],0.2)
#  cl = UpdateCellList!(x,box,cl)
#  @test length(cl.cwp) == 1598
#
#  box = Box([0.83,0.81,0.97],0.05)
#  cl = UpdateCellList!(x,box,cl)
#  @test length(cl.cwp) == 8737

end


# Test calling Knitro with license manager

using KNITRO, JuMP
using Test
using MathOptInterface
const MOI = MathOptInterface

# hs071
# Polynomial objective and constraints
# min x1 * x4 * (x1 + x2 + x3) + x3
# st  x1 * x2 * x3 * x4 >= 25
#     x1^2 + x2^2 + x3^2 + x4^2 = 40
#     1 <= x1, x2, x3, x4 <= 5
# Start at (1,5,5,1)
# End at (1.000..., 4.743..., 3.821..., 1.379...)

# Instantiate license manager
lm = KNITRO.LMcontext()

nruns = 10

for i in 1:nruns
    model = Model(with_optimizer(KNITRO.Optimizer, license_manager=lm, outlev=1))

    initval = [1, 5, 5, 1]

    @variable(model, 1 <= x[i=1:4] <= 5, start=initval[i])
    @NLobjective(model, Min, x[1] * x[4] * (x[1] + x[2] + x[3]) + x[3])
    c1 = @NLconstraint(model, x[1] * x[2] * x[3] * x[4] >= 25)
    c2 = @NLconstraint(model, sum(x[i]^2 for i=1:4) == 40)

    JuMP.optimize!(model)
    @test JuMP.termination_status(model) == MOI.LOCALLY_SOLVED

    # Warning! Free Knitro model before freeing license manager!
    MOI.empty!(backend(model))
end

# Free license manager
KNITRO.KN_release_license(lm)


using Test

using Petri
using AlgebraicPetri
using Catlab.Theories
using Catlab.CategoricalAlgebra.ShapeDiagrams
using Catlab.CategoricalAlgebra.FinSets

function compare_petricospan(p₁::PetriCospan, p₂::PetriCospan)
    @test dom(p₁) == dom(p₂)
    @test codom(p₁) == codom(p₂)
    @test force(left(p₁)) == force(left(p₂))
    @test force(right(p₁)) == force(right(p₂))
    @test decoration(p₁) == decoration(p₂)
end

@testset "Core" begin
    include("core.jl")
end

@testset "Petri" begin
    include("petri.jl")
end


#=
learn a bit more about MCMC. This explores using
a proposal density that takes into account the
correlations of the posterior, to get better
mixing

* play with scaling to see effects on acceptance rate
* examine chains with different acceptance rates
* look at scatter plots to see posterior correlations
=#
using Plots, LinearAlgebra, Statistics, Econometrics
include("QIVmodel.jl")
function main()
LNW, X, Z = getdata()
n = size(LNW,1)
mcmcreps = 100000
burnin = 100000
# use st. errs. from ordinary QR as quide to setting tuning.
basetuning = [0.1, 0.05, 0.03, 0.005, 0.3, 0.2, 0.2]
# adjust this up or down to achieve decent acceptance rate
scale = 0.1 # higher for higher taus
tuning = scale*basetuning
τ = 0.5
Σ = τ*(1.0-τ)*Z'Z/n
Σinv = inv(Σ)
θ = X\LNW  # OLS start values
# to do ordinary QR via MCMC, set Z=X (just to verify that MCMC works!)
lnL = θ -> likelihood(θ, LNW, X, Z, τ, Σinv)
Proposal = θ -> proposal(θ, tuning)
Prior = θ -> prior(θ)
chain = mcmc(θ, mcmcreps, burnin, Prior, lnL, Proposal)
d = dstats(chain)
# explore things:
p1 = plot(chain[:,1], legend=false, title="simple proposal: theta1") # chain is highly autocorrelated!
p2 = scatter(chain[:,1],chain[:,2], legend=false, title="simple proposal: theta1 vs. theta2") # there is significant posterior correlation
#= these results suggest that the proposal could be improved:
    * compute the covariance of the chain, and make draws from random walk MVN with that covariance
    * the following implements this idea
=#

# the following uses a different proposal, with correlations
# across the parameters, based on the previous chain.
V = cov(chain[:,1:end-1])
cholV = cholesky(V).U
tuning = 0.5
Proposal = θ -> proposal2(θ, tuning*cholV)
chain = mcmc(θ, mcmcreps, burnin, Prior, lnL, Proposal)
p3 = plot(chain[:,1], legend=false, title="correlated proposal: theta1") # chain is highly autocorrelated!
p4 = scatter(chain[:,1],chain[:,2], legend=false, title="correlated proposal: theta1 vs. theta2") # there is significant posterior correlation
plot(p1,p2,p3,p4,layout=(2,2))
gui()
return
end
main()


@block AnthonyWang [tutorials] begin

    @cell "Tutorial simple scatter" [tutorial, scatter, "2d"] begin
        x = rand(10)
        y = rand(10)
        colors = rand(10)
        scene = scatter(x, y, color = colors)
    end

    @cell "Tutorial markersize" [tutorial, scatter, markersize, "2d"] begin
        x = 1:10
        y = 1:10
        sizevec = x ./ 10
        scene = scatter(x, y, markersize = sizevec)
    end

    @cell "Tutorial simple line" [tutorial, line, "2d"] begin
        x = range(0, stop = 2pi, length = 40)
        f(x) = sin.(x)
        y = f(x)
        scene = lines(x, y, color = :blue)
    end

    @cell "Tutorial adding to a scene" [tutorial, line, scene, markers, "2d"] begin
        x = range(0, stop = 2pi, length = 80)
        f1(x) = sin.(x)
        f2(x) = exp.(-x) .* cos.(2pi*x)
        y1 = f1(x)
        y2 = f2(x)

        scene = lines(x, y1, color = :blue)
        scatter!(scene, x, y1, color = :red, markersize = 0.1)

        lines!(scene, x, y2, color = :black)
        scatter!(scene, x, y2, color = :green, marker = :utriangle, markersize = 0.1)
    end

    @cell "Tutorial removing from a scene" [tutorial, line, scene, markers, "2d"] begin
        x = range(0, stop = 2pi, length = 80)
        y1 = sin.(x)
        y2 = exp.(-x) .* cos.(2pi * x)

        scene = lines(x, y1, color = :blue)
        scatter!(scene, x, y1, color = :red, markersize = 0.1)

        lines!(scene, x, y2, color = :black)
        scatter!(scene, x, y2, color = :green, marker = :utriangle, markersize = 0.1)
        # initialize the stepper and give it an output destination
        st = Stepper(scene, @replace_with_a_path)
        step!(st)

        # remove last plot in scene
        delete!(scene, scene[end])
        step!(st)
        # remove second plot in scene
        delete!(scene, scene[2])
        step!(st)
        st
    end

    @cell "Tutorial adjusting scene limits" [tutorial, scene, limits, "2d"] begin
        x = range(0, stop = 10, length = 40)
        y = x
        #= specify the scene limits, note that the arguments for FRect are
            x_min, y_min, x_dist, y_dist,
            therefore, the maximum x and y limits are then x_min + x_dist and y_min + y_dist
        =#
        limits = FRect(-5, -10, 20, 30)

        scene = lines(x, y, color = :blue, limits = limits)

    end

    @cell "Tutorial basic theming" [tutorial, scene, limits, "2d"] begin
        x = range(0, stop = 2pi, length = 40)
        y = cos.(x)
        scene = lines(x, y, color = :blue)

        axis = scene[Axis] # get the axis object from the scene
        axis.grid.linecolor = ((:red, 0.5), (:blue, 0.5))
        axis.names.textcolor = ((:red, 1.0), (:blue, 1.0))
        axis.names.axisnames = ("x", "y = cos(x)")
        scene
    end

    @cell "Tutorial heatmap" [tutorial, heatmap, "2d"] begin
        data = rand(50, 50)
        scene = heatmap(data)
    end

    @cell "Tutorial linesegments" [tutorial, linesegments, "2d"] begin
        points = [
            Point2f0(0, 0) => Point2f0(5, 5);
            Point2f0(15, 15) => Point2f0(25, 25);
            Point2f0(0, 15) => Point2f0(35, 5);
        ]
        linesegments(points, color = :red, linewidth = 2)
    end

    @cell "Tutorial barplot" [tutorial, barplot, "2d"] begin
        data = sort(randn(100))
        barplot(data)
    end

    @cell "Tutorial title" [tutorial, title, "2d"] begin
        scene = lines(rand(10), axis=(names=(title="Random lines",),))
    end

    @cell "Tutorial title function" [tutorial, title, "2d"] begin
        scene = lines(rand(10))
        sc_t = title(scene, "Random lines")
        sc_t
    end

    @cell "Tutorial plot transformation" [tutorial, transformation, "2d"] begin
        data = rand(10)

        scene = Scene()
        st = Stepper(scene, @replace_with_a_path)
        lineplot = lines!(scene, data)[end]    # gets the last defined plot for the Scene
        scatplot = scatter!(scene, data)[end]  # same thing but the last defined plot is scatter
        step!(st)

        rotate!(lineplot, 0.025π) # only the lines are rotated, not the scatter
        step!(st)
        st
    end

    # @cell "Tutorial markersizes" [tutorial, markersize, scatter] begin
    #
    #     scene = scatter(rand(10); markersize = 10px)
    #
    # end

end


using MPI, MPIArrays

MPI.Init()
rank = MPI.Comm_rank(MPI.COMM_WORLD)
N = 8 # size of the matrix

# Create an uninitialized matrix and vector
x = MPIArray{Float64}(N)
A = MPIArray{Float64}(N,N)

# Set random values
forlocalpart!(m -> fill!(m,1+rank),x)
forlocalpart!(m -> fill!(m,1+rank),A)

# Make sure every process finished initializing the coefficients
sync(A, x)

b = A*x

# This will print on the first process, using slow element-by-element communication, but that's OK to print to screen
rank == 0 && println("Matvec result: $b")

# Clean up
free(A)
free(x)
MPI.Finalize()


# Generic newton method

function newton(F_DF, x₀; tol::Real=1e-12, maxiter::Int=15, verbose::Bool=true)
    verbose && return _newton_verbose(F_DF, x₀, tol, maxiter)
    return _newton_silent(F_DF, x₀, tol, maxiter)
end

function _newton_silent(F_DF, x₀::Number, tol, maxiter)
    F, DF = F_DF(x₀)
    nF = abs(F)
    nF ≤ tol && return x₀, true
    x = x₀
    i = 1
    while i ≤ maxiter && isfinite(nF)
        x -= DF \ F
        F, DF = F_DF(x)
        nF = abs(F)
        nF ≤ tol && return x, true
        i += 1
    end
    return x, false
end

function _newton_silent(F_DF, x₀, tol, maxiter)
    F, DF = F_DF(x₀)
    nF = norm(F)
    nF ≤ tol && return x₀, true
    x = copy(x₀)
    i = 1
    while i ≤ maxiter && isfinite(nF)
        x .-= ldiv!(lu!(DF), F)
        F, DF = F_DF(x)
        nF = norm(F)
        nF ≤ tol && return x, true
        i += 1
    end
    return x, false
end

function _newton_verbose(F_DF, x₀::Number, tol, maxiter)
    _display_newton_infos(tol, maxiter)
    F, DF = F_DF(x₀)
    nF = abs(F)
    AF = DF \ F
    nAF = abs(AF)
    _display_newton_iteration(0, nF, nAF)
    nF ≤ tol && return x₀, true
    x = x₀
    i = 1
    while i ≤ maxiter && isfinite(nF)
        x -= AF
        F, DF = F_DF(x)
        nF = abs(F)
        AF = DF \ F
        nAF = abs(AF)
        _display_newton_iteration(i, nF, nAF)
        nF ≤ tol && return x, true
        i += 1
    end
    return x, false
end

function _newton_verbose(F_DF, x₀, tol, maxiter)
    _display_newton_infos(tol, maxiter)
    F, DF = F_DF(x₀)
    nF = norm(F)
    ldiv!(lu!(DF), F)
    nAF = norm(F)
    _display_newton_iteration(0, nF, nAF)
    nF ≤ tol && return x₀, true
    x = copy(x₀)
    i = 1
    while i ≤ maxiter && isfinite(nF)
        x .-= F
        F, DF = F_DF(x)
        nF = norm(F)
        ldiv!(lu!(DF), F)
        nAF = norm(F)
        _display_newton_iteration(i, nF, nAF)
        nF ≤ tol && return x, true
        i += 1
    end
    return x, false
end

# in-place version

function newton!(F_DF!, x₀, F, DF; tol::Real=1e-12, maxiter::Int=15, verbose::Bool=true)
    verbose && return _newton_verbose!(F_DF!, x₀, F, DF, tol, maxiter)
    return _newton_silent!(F_DF!, x₀, F, DF, tol, maxiter)
end

function _newton_silent!(F_DF!, x₀, F, DF, tol, maxiter)
    F_DF!(F, DF, x₀)
    nF = norm(F)
    nF ≤ tol && return x₀, true
    i = 1
    while i ≤ maxiter && isfinite(nF)
        x₀ .-= ldiv!(lu!(DF), F)
        F_DF!(F, DF, x₀)
        nF = norm(F)
        nF ≤ tol && return x₀, true
        i += 1
    end
    return x₀, false
end

function _newton_verbose!(F_DF!, x₀, F, DF, tol, maxiter)
    _display_newton_infos(tol, maxiter)
    F_DF!(F, DF, x₀)
    nF = norm(F)
    ldiv!(lu!(DF), F)
    nAF = norm(F)
    _display_newton_iteration(0, nF, nAF)
    nF ≤ tol && return x₀, true
    i = 1
    while i ≤ maxiter && isfinite(nF)
        x₀ .-= F
        F_DF!(F, DF, x₀)
        nF = norm(F)
        ldiv!(lu!(DF), F)
        nAF = norm(F)
        _display_newton_iteration(i, nF, nAF)
        nF ≤ tol && return x₀, true
        i += 1
    end
    return x₀, false
end

#

function _display_newton_infos(tol, maxiter)
    println(string("\n Newton's method: tol = ", tol, ", maxiter = ", maxiter))
    println("      iteration        |F(x)|             |AF(x)|")
    println("-------------------------------------------------")
end

_display_newton_iteration(i, nF, nAF) = @printf("%11d %19.4e %19.4e\n", i, nF, nAF)


using MinimaxCalibratedEBayes
using Test
using Plots
pgfplotsx()

xs = range(-3,3;step=0.5)
ws = Float64.(collect(1:14))
tmp_hist = MCEBHistogram(xs)

@test StatsBase.midpoints(tmp_hist.grid) == StatsBase.midpoints(tmp_hist)

calib_normal = DiscretizedAffineEstimator(tmp_hist, Normal(0, 0.5))
calib_sin = DiscretizedAffineEstimator(tmp_hist, sin)

calib_sinc = DiscretizedAffineEstimator(tmp_hist, SincKernel(m=10_000))

calib_array = [calib_sin, calib_normal, calib_sinc]
@test isa(calib_array, AbstractVector{<:DiscretizedAffineEstimator})

plot(calib_normal)
plot(calib_sin)

plot(calib_array, label=["bla" "blabla" "sinc"])

z = DiscretizedStandardNormalSample(0.5, tmp_hist)
zs = DiscretizedStandardNormalSample.(randn(100), tmp_hist)

@test var(marginalize(Normal(0,0), StandardNormalSample(1.0))) == 1.0
@test var(marginalize(Normal(0,1), StandardNormalSample(1.0))) ≈ 2.0

tmp_marginalize = marginalize(Normal(), z)
marginalize_probs = pdf(tmp_marginalize)
@test sum(pdf(tmp_marginalize)) == 1

midpoints(marginalize_probs)

plot(marginalize_probs)

first(tmp_hist)

zs[1].mhist
tmp_calib = DiscretizedAffineEstimator(tmp_hist, ws, 1.0)

@test tmp_calib(-4) == 2.00
@test tmp_calib.([-4, 4]) == [2.0, 15.0]
@test tmp_calib.([-4, -3, -2.9, 4]) == [2.0, 2.0, 3.0, 15.0]

@test step(tmp_hist) == step(xs)
@test first(tmp_hist) == first(xs)
@test last(tmp_hist) == last(xs)


using CSV
using DataFrames
using StatsBase
using StatsModels
using Statistics
using Lathe
using GLM

data = DataFrame(CSV.File("motor_cena.csv"))

#display(describe(data))
#display(countmap(data[!,:Boja]))
#display(countmap(data[!,:Stanje]))
#display(countmap(data[!,:Ostecenje]))
#display(describe(data))

select!(data, Not(:Boja))

data[ismissing.(data[!,:Stanje]), :Stanje] .= mode(skipmissing(data[!,:Stanje]))
data[ismissing.(data[!,:Ostecenje]), :Ostecenje] .= mode(skipmissing(data[!,:Ostecenje]))
data[ismissing.(data[!,:Kilometraza]), :Kilometraza] .= trunc(Int64, mode(skipmissing(data[!,:Kilometraza])))
dropmissing!(data,[:Cena])
filter!(row -> row.Cena != 0, data)

filter!(row -> row.Kilometraza <= 500000, data)
filter!(row -> row.Godiste <= 2021 && row.Godiste >= 1900, data)
filter!(row -> row.BrojCilindara > 0 && row.BrojCilindara < 200, data)

covKwKs = cov(data.kW, data.KS)

if covKwKs > 0.6
    select!(data, Not(:kW))
end

fm = @formula(Cena ~ Stanje+Tip+Godiste+Kilometraza+Kubikaza+KS+BrojCilindara+Ostecenje)
dataTrain, dataTest = Lathe.preprocess.TrainTestSplit(data, .75)
linearRegressor = lm(fm, dataTrain)

predictTest = predict(linearRegressor, dataTest)
errorTest = dataTest.Cena - predictTest
rmseTest = sqrt(mean(errorTest.*errorTest))

predictTrain = predict(linearRegressor, dataTrain)
errorTrain = dataTrain.Cena - predictTrain
rmseTrain = sqrt(mean(errorTrain.*errorTrain))

println("RMSE trening skupa je $rmseTrain")
println("RMSE test skupa je: $rmseTest")

if r2(linearRegressor) > 0.5
    println("Model je dobar za predvidjanje!")
end

#plotlyjs()

#plotGodisteCena = scatter(data.Godiste, data.Cena, title = "Godiste - Cena", xlabel="Godiste", ylabel = "Cena")

#display(plotGodisteCena)
#savefig(plotGodisteCena, "godisteCena.html")

#plotKilometrazaCena = scatter(data.Kilometraza, data.Cena, title = "Kilometraza - Cena", xlabel="Kilometraza", ylabel = "Cena")

#display(plotKilometrazaCena)
#savefig(plotKilometrazaCena, "kilometrazaCena.html")


export pwelch,
    windowproperty,
    hanning,
    hamming,
    rect

using DSP

# When x is a vector, it is treated as a single channel.
# When x is a matrix, the PSD is computed independently for each column
# and stored in the corresponding column of pxx
# This function is comparable with pwelch Matlab function for calcualting
# PSD

"""
    pwelch(x, window, noverlap, nfft,fs)

Return the power spectral density (PSD) estimate, pxx, of the input signal, x,
found using Welch's method.

### Fields:
* `x`    - a vector (single channel) or array (multi-channels).
* `window` - a window function used to divided x into segments
* `noverlap`- number of samples overlaping between segments
* `nfft` - number of fft point that is >= length of segment
* `fs`- Hz, sampling frequency
### Available Constructors:
* `pxx, f = pwelch(signal,window,noverlap,nfft,fs)`
* `pxx, f = pwelch(signal,window,noverlap,nfft,fs,"power")`
"""
function pwelch(x::AbstractVector,window,noverlap,nfft,fs)

    n = length(window)  # length of segment
    ob = welch_pgram(x, n, noverlap;
        onesided=true,  # one sided spectrum
        nfft=nfft,      # number of fft points
        fs=fs,          # sampling frequency
        window=window)  # window function

    return ob.power, ob.freq
end

# Using when x is multiple channels
function pwelch(x::AbstractArray,window,noverlap,nfft,fs)

    nrow, ncol = size(x)
    pxx = Array{Float64}(undef, floor(Int,nfft/2)+1, ncol) # assume 1 sided
    f = Array{Float64}(undef, floor(Int,nfft/2)+1,1)

        for i=1:ncol
            n = length(window)  # length of segment
            ob = welch_pgram(x[:,i], n, noverlap;
                onesided=true,  # one sided spectrum
                nfft=nfft,      # number of fft points
                fs=fs,          # sampling frequency
                window=window)  # window function
            pxx[:,i] = ob.power
            f[:,1] = ob.freq
        end

    return pxx, f

end

# Equivalent noise bandwidth, amplitude and energy correction factor
"""
    enbw(window)

    returs the two-sided equivalent noise bandwidth, amplitude correction factor
    and energy correction factor.
"""
function windowproperty(window)

    # equivalent noise bandwith
    ENBW = length(window)*sum(window .^2) / sum(window)^2

    # amplitude correction factor
    ACF = 1/( sum(window)/ length(window) )

    # energy correction factor
    ECF = ACF/sqrt(ENBW)

    return ENBW, ACF, ECF
end

# example
# windowproperty(hamming(1000))

function pwelch(x::AbstractVector,window,noverlap,nfft,fs,type::AbstractString)

    n = length(window)  # length of segment
    ob = welch_pgram(x, n, noverlap;
        onesided=true,  # one sided spectrum
        nfft=nfft,      # number of fft points
        fs=fs,          # sampling frequency
        window=window)  # window function
    if type=="power"
        ENBW = windowproperty(window)[1] # equivalent bandwith
        res = ob.freq[2] # frequency resolution
        pxx = ob.power * (ENBW*res) # estimate Power spectrum
        return pxx, ob.freq
    else
        return ob.power, ob.freq
    end

end

function pwelch(x::AbstractArray,window,noverlap,nfft,fs,type::AbstractString)

    nrow, ncol = size(x)
    pxx = Array{Float64}(undef, floor(Int,nfft/2)+1, ncol) # assume 1 sided
    f = Array{Float64}(undef, floor(Int,nfft/2)+1,1)

        for i=1:ncol
            n = length(window)  # length of segment
            ob = welch_pgram(x[:,i], n, noverlap;
                onesided=true,  # one sided spectrum
                nfft=nfft,      # number of fft points
                fs=fs,          # sampling frequency
                window=window)  # window function

            if type == "power"
                ENBW = windowproperty(window)[1] # equivalent bandwith
                res = ob.freq[2] # frequency resolution
                pxx[:,i] = ob.power * (ENBW*res) # estimate Power spectrum
            else
                pxx[:,i] = ob.power
            end

            f[:,1] = ob.freq
        end

    return pxx, f

end

# example
#=
filename = "R:\\CMPH-Windfarm Field Study\\Hallett\\set1\\Recording-3.1.pti"
signal, Fs, Date, Time = ptiread(filename)

res = 0.1
nfft = floor(Int,Fs/res)
ovlap = Int(0.5*nfft)

 pxx, f =  pwelch(signal,hanning(nfft),ovlap,nfft,Fs)
 pxx, f =  pwelch(signal,hanning(nfft),ovlap,nfft,Fs,"power")
=#