using Unitful, UnitfulAstro

import PhysicalConstants.CODATA2018: g_n

import ..SfUnits: Angle, to_angle, Acceleration

export kinetic_energy, distance, duration, acceleration, projectile_displacement, projectile_velocity,
	projectile_flight_time, projectile_apex, projectile_range, projectile_angle, projectile_angle_planetary

# TODO: rename these when the underlying compilation issues have been resolved
export projectile_range_planetary
# TODO: better names?
export projectile_peak_displacement_planetary

"""
	kinetic_energy(m::Unitful.Mass, v::Unitful.Velocity)

Compute the kinetic energy of a body with mass `m` travelling at velocity `v`.

No relativistic corrections are applied. Use `relativistic_kinetic_energy` if those are required.
"""
kinetic_energy(m::Unitful.Mass, v::Unitful.Velocity) = 0.5m*v^2 |> u"J"

"""
    distance(a::Unitful.Acceleration, t::Unitful.Time, initial_v::Unitful.Velocity = 0u"m/s")

Compute the distance travelled in time `t` by a body with initial velocity `initial_v` and uniform acceleration `a`.
"""
distance(a::Unitful.Acceleration, t::Unitful.Time, initial_v::Unitful.Velocity = 0u"m/s") = t*initial_v + a*t^2 / 2

"""
    duration(a::Unitful.Acceleration, d::Unitful.Length, initial_v::Unitful.Velocity = 0u"m/s")

Compute the time taken for a body with initial velocity `initial_v` and uniform acceleration `a` to travel a distance of `d`.
"""
function duration(a::Unitful.Acceleration, d::Unitful.Length, initial_v::Unitful.Velocity = 0u"m/s")
    if d < 0u"m"
		throw(DomainError(d, "Negative distances not allowed."))
	end

	if initial_v == 0u"m/s"
		return sqrt(2d / a)
	elseif a == 0u"m/s^2"
		if initial_v < 0u"m/s"
			throw(DomainError(d, "Negative initial velocities not allowed with zero acceleration."))
		end

		return d / initial_v
	end

    # at^2 + 2v_0t - 2d = 0
    # ax^2 + bx + c

    b = 2 * initial_v
    c = -2d

    discr = b^2 - 4*a*c

    t1 = (-b + sqrt(discr))/(2a)
    t2 = (-b - sqrt(discr))/(2a)

    # the negative results aren't wrong, per se, but as the parameters are
    # vector quantities the negative results aren't particularly interesting.

    if t1 < 0u"s"
        return t2
    elseif t2 < 0u"s"
        return t1
    else
        return min(t1, t2)
    end
end

"""
    acceleration(d::Unitful.Length, v_final::Unitful.Velocity, v_initial::Unitful.Velocity = 0u"m/s")

Given a distance `d`, initial velocity `v_initial and desired velocity `v_final` compute the required uniform acceleration
"""
acceleration(d::Unitful.Length, v_final::Unitful.Velocity, v_initial::Unitful.Velocity = 0u"m/s") = (v_final^2 - v_initial^2) / 2d

"""
    projectile_displacement(v_0::Unitful.Velocity, θ::Angle, t::Unitful.Time; g::Acceleration = g_n)

Horizontal and vertical displacement at time `t` of a projectile with initial velocity `v_0` and launch angle `θ` in a unitform gravitational acceleration `g`.
"""
projectile_displacement(v_0::Unitful.Velocity, θ::Angle, t::Unitful.Time; g::Acceleration = g_n) =
    [ v_0 * t * cos(θ), v_0 * t * sin(θ) - 0.5g * t^2 ]

"""
    projectile_displacement(v_0::Unitful.Velocity, t::Unitful.Time; g::Acceleration = g_n)

Vertical displacement at time `t` of a projectile launched vertically with initial velocity `v_0` in a unitform gravitational acceleration `g`.
"""
projectile_displacement(v_0::Unitful.Velocity, t::Unitful.Time; g::Acceleration = g_n) = projectile_displacement(v_0, 90u"°", t, g = g)[2]

"""
    projectile_velocity(v_0::Unitful.Velocity, θ::Angle, t::Unitful.Time; g::Acceleration = g_n)

Horizontal and vertical velocity at time `t` of a projectile with initial velocity `v_0` and launch angle `θ` in a unitform gravitational acceleration `g`.
"""
projectile_velocity(v_0::Unitful.Velocity, θ::Angle, t::Unitful.Time; g::Acceleration = g_n) =
    [ v_0 * cos(θ), v_0 * sin(θ) - g * t ]

"""
    projectile_velocity(v_0::Unitful.Velocity, t::Unitful.Time; g::Acceleration = g_n)

Velocity at time `t` of a projectile launched vertically with initial velocity `v_0` in a unitform gravitational acceleration `g`.
"""
projectile_velocity(v_0::Unitful.Velocity, t::Unitful.Time; g::Acceleration = g_n) = projectile_velocity(v_0, 90u"°", g, t)[2]

"""
    projectile_flight_time(v_0::Unitful.Velocity, θ::Angle = 90u"°"; g::Acceleration = g_n)

Compute total flight time for a projectile launched at angle `θ` with initial velocity `v_0` in a uniform gravitation acceleration `g`.

Flight is considered to be complete when vertical displacement returns to zero.
"""
projectile_flight_time(v_0::Unitful.Velocity, θ::Angle = 90u"°"; g::Acceleration = g_n) = 2v_0 * sin(θ) / g |> u"s"

"""
    projectile_apex(v_0::Unitful.Velocity, θ::Angle = 90u"°"; g::Acceleration = g_n)

Horizontal and vertical displacement of a projectile launched at angle `θ` with initial velocity `v_0` in a uniform gravitation acceleration `g` at the apex of its trajectory.

Default launch angle is 90° for maximum height.
"""
projectile_apex(v_0::Unitful.Velocity, θ::Angle = 90u"°"; g::Acceleration = g_n) = [ v_0^2 * sin(2θ) / 2g,  v_0^2 * sin(θ)^2 / 2g ]

"""
    projectile_range(v_0::Unitful.Velocity, θ::Angle = 45u"°"; g::Acceleration = g_n)

Compute range of a projectile launched at angle `θ` with initial velocity `v_0` in a uniform gravitation acceleration `g`.

Default launch angle is 45° for maximum range.
"""
projectile_range(v_0::Unitful.Velocity, θ::Angle = 45u"°"; g::Acceleration = g_n) = (v_0^2 * sin(2θ)) / g

"""
    projectile_range(v_0::Unitful.Velocity, y_0::Unitful.Length, θ::Angle = 45u"°"; g::Acceleration = g_n)

Compute range of a projectile launched from altitude `y_0` at angle `θ` with initial velocity `v_0` in a uniform gravitation acceleration `g`.

Range is reached when vertical displacement is zero. Default launch angle is 45° for maximum range.
"""
projectile_range(v_0::Unitful.Velocity, y_0::Unitful.Length, θ::Angle = 45u"°"; g::Acceleration = g_n) =
	(v_0 * cos(θ) / g) * (v_0 * sin(θ) + sqrt((v_0 * sin(θ))^2 + 2g * y_0))

"""
    projectile_range_planetary(v_0::Unitful.Velocity, r_planet::Unitful.Length = 1u"Rearth", θ::Angle = 45u"°"; g::Acceleration = g_n)

Compute range of a projectile launched from altitude `y_0` at angle `θ` with initial velocity `v_0` on a spherical planet with radius `r_planet` and surface gravity `g`.

NOTE: this should be called `projectile_range`, but weird internal compiler errors prevent that with this version of Julia (1.5.3).
"""
function projectile_range_planetary(v_0::Unitful.Velocity, θ::Angle = 45u"°", r_planet::Unitful.Length = 1u"Rearth"; g::Acceleration = g_n)
	v_rat2 = (v_0 / sqrt(r_planet * g))^2

	if v_rat2 > 1
		throw(DomainError("Initial velocity exceeds orbital velocity; Range undefined"))
	end

	a = v_0^2 * sin(2θ) / g
	b = sqrt(1 - (2 - v_rat2) * v_rat2 * cos(θ)^2)

	return a / b |> u"m"
end

function projectile_peak_displacement_planetary(v_0::Unitful.Velocity, θ::Angle = 45u"°", r_planet::Unitful.Length = 1u"Rearth"; g::Acceleration = g_n)
	v_rat2 = (v_0 / sqrt(r_planet * g))^2

	if v_rat2 > 1
		throw(DomainError("Initial velocity exceeds orbital velocity; Range undefined"))
	end

	a = v_0^2 * sin(θ) / g
	b = 1 - v_rat2 + sqrt(1 - (2 - v_rat2) * v_rat2 * cos(θ)^2)

	return a / b |> u"m"
end

"""
    projectile_angle(v_0::Unitful.Velocity, d::Unitful.Length; g::Acceleration = g_n)

Compute the possible launch angles for a projectile to reach a horizontal range of `d` given initial velocity `v_0` in a uniform gravitation acceleration `g`.

The shallow angle is the first result, the steep angle the second. A domain error is raised if the target is out of range.
"""
function projectile_angle(v_0::Unitful.Velocity, d::Unitful.Length; g::Acceleration = g_n)
	k = (g * d) / v_0^2

	if (k > 1)
		throw(DomainError("Target out of range; no solutions"))
	end

	return ( 0.5asin(k) * 1u"rad" |> u"°", 45u"°" + (0.5acos(k) * 1u"rad" |> u"°") )
end

"""
    projectile_angle(v_0::Unitful.Velocity, d::Unitful.Length, y::Unitful.Length; g::Acceleration = g_n)

Compute the possible launch angles for a projectile to reach a horizontal range of `d` and altitude of `y` relative to the starting point given initial velocity `v_0` in a uniform gravitation acceleration `g`.

The shallow angle is the first result, the steep angle the second. A domain error is raised if the target is out of range.
"""
function projectile_angle(v_0::Unitful.Velocity, d::Unitful.Length, y::Unitful.Length; g::Acceleration = g_n)
	det = v_0^4 - g * (g * d^2 + 2y * v_0^2)

	if (ustrip(det) < 0)
		throw(DomainError("Target out of range; no solutions"))
	end

	a = atan((v_0^2 + sqrt(det)) / (g * d)) * 1u"rad" |> u"°"
	b = atan((v_0^2 - sqrt(det)) / (g * d)) * 1u"rad" |> u"°"

	if a < b
		return (a, b)
	else
		return (b, a)
	end
end

"""
    projectile_optimum_angle(v_0::Unitful.Velocity, r_planet::Unitful.Length = 1u"Rearth"; g::Acceleration = g_n)
"""
function projectile_optimum_angle(v_0::Unitful.Velocity, r_planet::Unitful.Length = 1u"Rearth"; g::Acceleration = g_n)
	v_rat2 = (v_0 / sqrt(r_planet * g))^2

	if v_rat2 > 0
		throw(DomainError("Initial velocity exceeds orbital velocity; optimum angle undefined"))
	end

	return 0.5acos(v_rat2 / (2 - v_rat2)) * u"rad" |> u"°"
end


function eller(n)
    result = 1
    answer,hw  = factor(n)
    for i in 1:length(answer)
        if hw[i] == 1
            result *= answer[i]-1
        else
            result *= (answer[i]^hw[i] - answer[i]^(hw[i]-1))
        end
    end
    return result
end

function factor(n)
    answer = []
    hw = []
    for i in eratosphen(n)
        if n%i==0
            push!(answer,i)
            push!(hw,1)
            n=n/i
            while n%i==0
                hw[end]+=1
                n=n/i
            end
        end
    end
    return answer,hw
end

function eratosphen(n)
    ser = fill(true,n)
    ser[1] = false
    k = 2
    while k !== nothing && k<n
        ser[k^2:k:end] .= false
        k=findnext(ser, k+1)
    end
    return findall(ser)
end


using Optim, Plots
# method for contour
rosenbrock(x1,x2) = (1.0 - x1)^2 + 100.0 * (x2 - x1^2)^2
# method for optim
rosenbrock(θ) = (1.0 - θ[1])^2 + 100.0 * (θ[2] - θ[1]^2)^2
θ = zeros(2)
b1 = range(0.0,stop=1.5,length=100)
b2 = range(0.0,stop=1.5,length=100)
contour(b1,b2,(b1,b2)->rosenbrock(b1,b2),fill=true, c=:viridis)
θs = zeros(100,2)
for i = 1:100
    #global θ = optimize(rosenbrock, θ, GradientDescent(), Optim.Options(iterations=1)).minimizer
    global θ = optimize(rosenbrock, θ, Newton(), Optim.Options(iterations=1)).minimizer
    θs[i,:] = θ'
    display(scatter!(θs[1:i,1], θs[1:i,2], legend=false))
    sleep(0.1)
end
θs


# Magnetic field topology in fusion research.

using Elliptic, SpecialFunctions

"""
    getB_bottle(x, y, z, distance, a, b, I1, I2) -> Vector{Float}

Get magnetic field from a magnetic bottle.
Reference: https://en.wikipedia.org/wiki/Magnetic_mirror#Magnetic_bottles

# Arguments
- `x,y,z::Float`: location in [m].
- `distance::Float`: distance between solenoids in [m].
- `a::Float`: radius of each side coil in [m].
- `b::Float`: radius of central coil in [m].
- `I1::Float`: current in the solenoid times number of windings in side coils.
- `I2::Float`: current in the central solenoid times number of windings in the
central loop.
"""
function getB_bottle(x, y, z, distance, a, b, I1, I2)

   r = √(x^2 + y^2) # distance from z-axis

   # 1st loop
   z₁ = z + 0.5*distance
   k = √(4*r*a / (z₁^2 + (a + r)^2) )
   K, E = ellipke(k)
   Bz1 = μ₀*I1 / (2π*√(z₁^2+(a+r)^2)) * ((a^2-z₁^2-r^2)/(z₁^2+(r-a)^2)*E + K)
   Br1 = μ₀*z₁*I1/(2π*r*√(z₁^2+(a+r)^2))*((z₁^2+r^2+a^2)/(z₁^2+(r-a)^2)*E - K)
   Bx1 = Br1 * x / r
   By1 = Br1 * y / r

   # 2nd loop
   z₂ = z - 0.5*distance
   k = √(4*r*a / (z₂^2 + (a + r)^2) )
   K, E = ellipke(k)
   Bz2 = μ₀*I1 / (2π*√(z₂^2+(a+r)^2)) * ((a^2-z₂^2-r^2)/(z₂^2+(r-a)^2)*E + K)
   Br2 = μ₀*z₂*I1/(2π*r*√(z₂^2+(a+r)^2))*((z₂^2+r^2+a^2)/(z₂^2+(r-a)^2)*E - K)
   Bx2 = Br2 * x / r
   By2 = Br2 * y / r

   # central loop
   z₃ = z
   k = √(4*r*b / (z₃^2 + (b+r)^2) )
   K, E = ellipke(k)
   Bz3 = μ₀*I2 / (2π*√(z₃^2+(b+r)^2)) * ((b^2-z₃^2-r^2)/(z₃^2+(r-b)^2)*E + K)
   Br3 = μ₀*z₃*I2/(2π*r*√(z₃^2+(b+r)^2))*((z₃^2+r^2+b^2)/(z₃^2+(r-b)^2)*E - K)
   Bx3 = Br3 * x / r
   By3 = Br3 * y / r

   # total magnetic field
   if x == 0.0 && y == 0.0
      Bx = 0.0
      By = 0.0
      Bz = Bz1 + Bz2 + Bz3
   else
      Bx = Bx1 + Bx2 + Bx3
      By = By1 + By2 + By3
      Bz = Bz1 + Bz2 + Bz3
   end
   [Bx, By, Bz]
end

"""
    getB_tokamak(x, y, z, a, b, ICoils, IPlasma)

Get the magnetic field from a Tokamak topology consists of 16 coils.
Original: https://github.com/BoschSamuel/Simulation-of-a-Tokamak-Fusion-Reactor/blob/master/Simulation2.m
# Arguments
- `x,y,z::Float`: location in [m].
- `a::Float`: radius of each coil in [m].
- `b::Float`: radius of central region in [m].
- `ICoil::Float`: current in the coil times number of windings.
- `IPlasma::Float`: current of the plasma?
"""
function getB_tokamak(x, y, z, a, b, ICoils, IPlasma)

   a *= 2

   Bx, By, Bz = 0.0, 0.0, 0.0

   # magnetic field of the coils
   for i = 0:15
      θ = π/16 + i*π/8 # angle between the i-th coil and the x-axis

      if abs(sin(θ)) > 0.01
         r1_ = x/(cos(θ)-sin(θ)*tan(atan(y, x) - θ)) - a - b
         r1 = √(r1_^2 + z^2)
         z1 = ((b + a)*cos(θ) + r1_*cos(θ) - x) / sin(θ)
      else
         r1_ = x - b - a
         r1 = √(r1_^2 + z^2)
         z1 = y
      end

      k = √(4r1*a / (z1^2 + (a + r1)^2))

      K, E = ellipke(k)
      # Bz1_ is the magnetic field in the coil frame
      Bz1_ = μ₀*ICoils/(2π*√(z1^2+(a+r1)^2))*((a^2-z1^2-r1^2)/(z1^2+(r1-a)^2)*E+K)
      # Br1_ is the magnetic field in the coil frame
      Br1_ = μ₀*z1*ICoils/(2π*r1*√(z1^2+(a+r1)^2))*((z1^2+r1^2+a^2)/(z1^2+(r1-a)^2)*E-K)
      # normal coordinates
      Bx1 = -sin(θ)*Bz1_ + Br1_*r1_/r1*cos(θ)
      By1 = cos(θ)*Bz1_ + sin(θ)*Br1_*r1_/r1
      Bz1 = Br1_ * z / r1

      # add the field of a single coil to the total field
      Bx += Bx1
      By += By1
      Bz += Bz1

      if abs(Bx) < 5e-12
         Bx = 0.0
      end
      if abs(By) < 5e-12
         By = 0.0
      end
      if abs(Bz) < 5e-12
         Bz = 0.0
      end
   end

   # magnetic field of the plasma current
   σ = a/3 # parameter of the Gauss curve
   ϕ = atan(y, x)
   # distance to centre of plasma ring
   distance = √( z^2 + (x - (a + b)*cos(ϕ))^2 + (y - (a + b)*sin(ϕ))^2 )
   I2_r_plasma = IPlasma * erf(distance/(σ*√2))

   r = hypot(x, y)
   k = √(4r*(a+b)/(z^2+((a+b)+r)^2))
   K, E = ellipke(k)
   Bz_plasma = μ₀*I2_r_plasma/(2π*√(z^2+((a+b)+r)^2))*(((a+b)^2-z^2-r^2)/(z^2+(r-(a+b))^2)*E+K)
   Br_plasma = μ₀*z*I2_r_plasma/(2π*r*√(z^2+(b+r)^2))*((z^2+r^2+(a+b)^2)/(z^2+(r-(a+b))^2)*E-K)
   Bx_plasma = Br_plasma*x/r
   By_plasma = Br_plasma*y/r

   if distance > 0.0001
      Bx += Bx_plasma
      By += By_plasma
      Bz += Bz_plasma
   end
   [Bx, By, Bz]
end


function rouwenhorst(n, mu, sigma, rho)

    # grid
    width = sqrt((n-1) * sigma^2 / ( 1 - rho^2))
    grid = vcat(range( mu-width, mu + width, length=n))

    # transition matrix
    p0 = (1 + rho) / 2
    trans = [p0 1-p0; 1-p0 p0]

    if n > 2
        for i = 1:n-2
            cstr_temp = zeros(length(trans[:,1]), 1)
            trans = p0 .* [trans cstr_temp; cstr_temp' 0] + (1 - p0 ) .* [cstr_temp trans; cstr_temp' 0]  + (1 - p0 ) .*  [ cstr_temp' 0; trans cstr_temp] + p0 .* [ cstr_temp' 0; cstr_temp trans]
        end
        trans = trans ./ sum(trans,dims=2)
    end

    # ergodic distribution
    dist = ones(1,n)./n
    for i = 1:100
        dist = dist*(trans^i)
    end
    dist = vcat(dist...)

    return grid, trans, dist
end


"""
    UCB1 Implementation
    Based on: Figure-1, Auer, P., Bianchi, N. C., & Fischer, P. (2002). Finite time analysis of the multiarmed bandit problem. Machine Learning, 47, 235–256.
"""

type UCB1 <: BanditAlgorithmBase
    noOfArms::Int64
    noOfSteps::Int64
    lastPlayedArm::Int64

    c::Float64

    cummReward::Vector{Float64}
    count::Vector{Int64}
    ucbIndices::Vector{Float64}

    function UCB1( noOfArms::Int, c::Real = √2 )
        new( noOfArms,
             0,
             0,
             c,
             zeros(Float64,noOfArms),
             zeros(Int64,noOfArms),
             zeros(Float64,noOfArms)
        )
    end
end

function get_arm_index( agent::UCB1 )
    if any(agent.count.==0)
        agent.lastPlayedArm =  rand( find(agent.count.==0) )
    else
        agent.lastPlayedArm = findmax(agent.ucbIndices)[2]
    end

    return agent.lastPlayedArm
end

function update_reward!( agent::UCB1, r::Real )

    # Update cummulative reward
    agent.cummReward[agent.lastPlayedArm] += r

    # Update play count for arm
    agent.count[agent.lastPlayedArm] += 1

    # Update number of steps played
    agent.noOfSteps += 1

    # Update UCB indices
    agent.ucbIndices = agent.cummReward./agent.count +
                        agent.c * sqrt.(log(agent.noOfSteps)./agent.count)

    nothing
end

function reset!( agent::UCB1 )
    agent.noOfSteps     = 0
    agent.lastPlayedArm = 0

    agent.cummReward    = zeros( Float64, agent.noOfArms )
    agent.count         = zeros( Int64, agent.noOfArms )
    agent.ucbIndices    = zeros( Float64, agent.noOfArms )

    nothing
end

function info_str( agent::UCB1, latex::Bool )
    if agent.c == √2
        return @sprintf( "UCB1" )
    else
        return @sprintf( "UCB1 (c=%3.2f)", agent.c )
    end
end

# """
#     UCB2 Implementation
#     Based on: Figure-2, Auer, P., Bianchi, N. C., & Fischer, P. (2002). Finite time analysis of the multiarmed bandit problem. Machine Learning, 47, 235–256.
# """
# type UCB2 <: BanditAlgorithmBase
#     noOfArms::Int64
#     noOfSteps::Int64
#     lastPlayedArm::Int64
#     α::Float64
#
#     cummReward::Vector{Float64}
#     count::Vector{Int64}
#
#     function UCB2( noOfArms::Int, α::Float64 )
#         new( noOfArms,
#              0,
#              0,
#              α,
#              zeros(Float64,noOfArms),
#              zeros(Int64,noOfArms)
#         )
#     end
# end

# function getArmIndex( agent::UCB2 )
#
# end
#
# function updateReward( agent::UCB2, r::Float64 )
#
# end

"""
    UCB-Normal Implementation
    Based on: Figure-4, Auer, P., Bianchi, N. C., & Fischer, P. (2002). Finite time analysis of the multiarmed bandit problem. Machine Learning, 47, 235–256.
"""

type UCBNormal <: BanditAlgorithmBase
    noOfArms::Int64
    noOfSteps::Int64
    lastPlayedArm::Int64

    cummReward::Vector{Float64}
    count::Vector{Int64}
    cummSqReward::Vector{Float64}
    ucbIndices::Vector{Float64}

    function UCBNormal( noOfArms::Int )
        new( noOfArms,
             0,
             0,
             zeros(Float64,noOfArms),
             zeros(Int64,noOfArms),
             zeros(Float64,noOfArms),
             zeros(Float64,noOfArms)
        )
    end
end

function get_arm_index( agent::UCBNormal )
    tmp_lb      = ceil( 8 * log(agent.noOfSteps) )  # Quantity for comparison
    tmp_lbArms  = find( agent.count .< tmp_lb )     # Find index of under played arms
    if( size(tmp_lbArms,1) > 0 )    # if any machine is played less than ceil(8 log n), play that arm
        agent.lastPlayedArm = rand( tmp_lbArms )    # Randomly play one of those arms
    else    # else play arm with highest UCB index
        agent.lastPlayedArm = findmax(agent.ucbIndices)[2]
    end
    return agent.lastPlayedArm
end

function update_reward!( agent::UCBNormal, r::Real )
    # Update cummulative reward
    agent.cummReward[agent.lastPlayedArm] += r
    # Update squared cummulative reward
    agent.cummSqReward[agent.lastPlayedArm] += (r^2)
    # Update count for last played arm
    agent.count[agent.lastPlayedArm] += 1
    # Update number of steps
    agent.noOfSteps += 1
    # Update UCB indices
    agent.ucbIndices    = agent.cummReward ./ agent.count +
                            sqrt.(16 *
                                (agent.cummSqReward-((agent.cummReward).^2)./agent.count) ./ (agent.count-1) *
                                log(agent.noOfSteps-1)./agent.count )

    nothing
end

function reset!( agent::UCBNormal )
    agent.noOfSteps     = 0
    agent.lastPlayedArm = 0

    agent.cummReward    = zeros( Float64, agent.noOfArms )
    agent.count         = zeros( Int64, agent.noOfArms )
    agent.cummSqReward  = zeros( Float64, agent.noOfArms )
    agent.ucbIndices    = zeros( Float64, agent.noOfArms )

    nothing
end

function info_str( agent::UCBNormal, latex::Bool )
    return @sprintf( "UCB Normal" )
end

"""
    Discounted UCB
    Based on: Moulines, E., & Paristech, T. (2008). On Upper-Confidence Bound Policies for Non-Stationary Bandit Problems. arXiv Preprint, (2008), 1–24.
"""

type DUCB <: BanditAlgorithmBase
    noOfArms::Int64
    noOfSteps::Int64
    lastPlayedArm::Int64

    γ::Float64                      # Discouting factor
    ξ::Float64                      # Confidence Scale Parameter
    cummReward::Vector{Float64}     # Cummulattive Reward from each arm
    discCumReward::Vector{Float64}  # Discounted reward
    count::Vector{Int64}            # Count of how many time an arm is pulled
    discCount::Vector{Float64}      # Discounted Count
    indices::Vector{Float64}        # Calculated Indices

    function DUCB( noOfArms::Int, γ::Float64, ξ::Float64 = 0.5 )
        new( noOfArms,
             0,
             0,
             γ,
             ξ,
             zeros(Float64,noOfArms),
             zeros(Float64,noOfArms),
             zeros(Int64,noOfArms),
             zeros(Float64,noOfArms),
             zeros(Float64,noOfArms)
        )
    end
end

function get_arm_index( agent::DUCB )
    if any(agent.count.==0)
        agent.lastPlayedArm =  rand( find(agent.count.==0) )
    else
        agent.lastPlayedArm = findmax(agent.indices)[2]
    end

    return agent.lastPlayedArm
end

function update_reward!( agent::DUCB, r::Real )
    # Update cummulative reward
    agent.cummReward[agent.lastPlayedArm] += r
    # Update play count for arm
    agent.count[agent.lastPlayedArm] += 1
    # Update number of steps played
    agent.noOfSteps += 1
    # Update discounted cummulative reward
    agent.discCumReward *= agent.γ                  # Discount for all arms
    agent.discCumReward[agent.lastPlayedArm] += r   # Update reward to last played arm
    # Update discounted count
    agent.discCount *= agent.γ
    agent.discCount[agent.lastPlayedArm] += 1
    # Update UCB indices
    agent.indices   = agent.discCumReward./agent.discCount +
                        2*sqrt.(agent.ξ*log(sum(agent.discCount))./agent.discCount)

    nothing
end

function reset!( agent::DUCB )
    agent.noOfSteps     = 0;
    agent.lastPlayedArm = 0;

    agent.cummReward    = zeros( Float64, agent.noOfArms );
    agent.discCumReward = zeros( Float64, agent.noOfArms );
    agent.count         = zeros( Int64, agent.noOfArms );
    agent.discCount     = zeros( Float64, agent.noOfArms );
    agent.indices       = zeros( Float64, agent.noOfArms );

    nothing
end

function info_str( agent::DUCB, latex::Bool )
    if latex
        return @sprintf( "Discounted-UCB\$(\\gamma = %4.3f,\\xi = %4.3f)\$", agent.γ, agent.ξ );
    else
        return @sprintf( "Discounted-UCB(γ = %4.3f, ξ = %4.3f)", agent.γ, agent.ξ );
    end
end

"""
    Sliding Window UCB
    Based on: Moulines, E., & Paristech, T. (2008). On Upper-Confidence Bound Policies for Non-Stationary Bandit Problems. arXiv Preprint, (2008), 1–24.
"""

type SWUCB <: BanditAlgorithmBase
    noOfArms::Int64
    noOfSteps::Int64
    lastPlayedArm::Int64

    τ::Int64                        # Window Length
    ξ::Float64                      # Confidence Scale Parameter
    cummReward::Vector{Float64}     # Cummulative Reward from each arm
    count::Vector{Int64}            # Count of how many time an arm is pulled
    swCount::Vector{Int64}          # Count of how many times an arm is picked in this sliding Window
    armsInWindow::Vector{Int64}     # Vector of Arms picked in window
    rewardsInWindow::Vector{Float64}# Vector of Rewards corresponding to arms picked in this window
    swCumRew::Vector{Float64}       # Cummulative reward of arms in this window
    indices::Vector{Float64}        # Calculated Indices

    function SWUCB( noOfArms::Int, τ::Int64, ξ::Float64 = 2.0 )
        new( noOfArms,
             0,
             0,
             τ,
             ξ,
             zeros(Float64,noOfArms),
             zeros(Int64,noOfArms),
             zeros(Int64,noOfArms),
             Vector{Int64}(),
             Vector{Float64}(),
             zeros(Float64,noOfArms),
             zeros(Float64,noOfArms)
        )
    end
end

function get_arm_index( agent::SWUCB )
    if any(agent.swCount.==0)
        agent.lastPlayedArm =  rand( find(agent.swCount.==0) )
    else
        agent.lastPlayedArm = findmax(agent.indices)[2]
    end
    return agent.lastPlayedArm
end

function update_reward!( agent::SWUCB, r::Real )
    # Update cummulative reward
    agent.cummReward[agent.lastPlayedArm] += r
    # Update play count for arm
    agent.count[agent.lastPlayedArm] += 1
    # Update number of steps played
    agent.noOfSteps += 1

    # For sliding window
    push!( agent.armsInWindow, agent.lastPlayedArm )
    push!( agent.rewardsInWindow, r )
    agent.swCount[agent.lastPlayedArm]  += 1
    agent.swCumRew[agent.lastPlayedArm] += r
    if length(agent.armsInWindow) > agent.τ
        _arm    = shift!( agent.armsInWindow )
        _r      = shift!( agent.rewardsInWindow )
        agent.swCount[_arm]     -= 1
        agent.swCumRew[_arm]    -= _r
    end

    # Sanity check
    assert( length(agent.armsInWindow) == length(agent.rewardsInWindow) <= agent.τ )

    # Update UCB indices
    agent.indices = agent.swCumRew./agent.swCount +
                        sqrt.(agent.ξ*log(min(agent.noOfSteps,agent.τ))./agent.swCount)

    nothing
end

function reset!( agent::SWUCB )
    agent.noOfSteps         = 0
    agent.lastPlayedArm     = 0
    agent.cummReward        = zeros(Float64,agent.noOfArms)
    agent.count             = zeros(Int64,agent.noOfArms)
    agent.swCount           = zeros(Int64,agent.noOfArms)
    agent.armsInWindow      = Vector{Int64}()
    agent.rewardsInWindow   = Vector{Int64}()
    agent.indices           = zeros(Float64,agent.noOfArms)

    nothing
end

function info_str( agent::SWUCB, latex::Bool )
    if latex
        return @sprintf( "SW-UCB\$(\\tau = %3.2f,\\xi = %3.2f)\$", agent.τ, agent.ξ );
    else
        return @sprintf( "SW-UCB(τ = %3.2f, ξ = %3.2f)", agent.τ, agent.ξ );
    end
end

"""
    UCB-V: Variance Aware UCB Implementation
    Based on: Audibert, J. Y., Munos, R., & Szepesvári, C. (2009). Exploration-exploitation tradeoff using variance estimates in multi-armed bandits. Theoretical Computer Science, 410(19), 1876–1902. http://doi.org/10.1016/j.tcs.2009.01.016
"""

type UCBV <: BanditAlgorithmBase
    noOfArms::Int64
    noOfSteps::Int64
    lastPlayedArm::Int64

    # Parameters
    ζ::Float64
    c::Float64
    b::Float64

    empiricalMean::Vector{Float64}
    empiricalVariance::Vector{Float64}
    count::Vector{Int64}
    ucbIndices::Vector{Float64}

    function UCBV( noOfArms::Int; ζ::Real = 1.0, c::Real = 1/3, b::Real = 1.0 )
        new( noOfArms,
             0,
             0,
             ζ,
             c,
             b,
             zeros(Float64,noOfArms),
             zeros(Float64,noOfArms),
             zeros(Int64,noOfArms),
             zeros(Float64,noOfArms)
        )
    end
end

function get_arm_index( agent::UCBV )
    if any(agent.count.==0)
        agent.lastPlayedArm =  rand( find(agent.count.==0) )
    else
        agent.lastPlayedArm = findmax(agent.ucbIndices)[2]
    end

    return agent.lastPlayedArm
end

function update_reward!( agent::UCBV, r::Real )

    # Save old mean
    μ = agent.empiricalMean[agent.lastPlayedArm]

    # Update cummulative reward
    agent.empiricalMean[agent.lastPlayedArm] = (agent.empiricalMean[agent.lastPlayedArm]*agent.count[agent.lastPlayedArm]+r) /
                                                    (agent.count[agent.lastPlayedArm]+1)
    # Update empirical Variance : Based on (30) in http://mathworld.wolfram.com/SampleVarianceComputation.html
    if agent.count[agent.lastPlayedArm] != 0
        agent.empiricalVariance[agent.lastPlayedArm] = (agent.count[agent.lastPlayedArm]-1)/agent.count[agent.lastPlayedArm] *
                                                            agent.empiricalVariance[agent.lastPlayedArm] +
                                                        1/(agent.count[agent.lastPlayedArm]+1) * (r-μ)^2
    end

    # Update play count for arm
    agent.count[agent.lastPlayedArm] += 1

    # Update number of steps played
    agent.noOfSteps += 1

    # Update UCB indices
    agent.ucbIndices = agent.empiricalMean +
                        sqrt.(2*agent.ζ*log(agent.noOfSteps)*agent.empiricalVariance./agent.count) +
                            3*agent.c*agent.b*log(agent.noOfSteps)./agent.count
    nothing
end

function reset!( agent::UCBV )
    agent.noOfSteps     = 0
    agent.lastPlayedArm = 0

    agent.empiricalMean     = zeros( Float64, agent.noOfArms )
    agent.empiricalVariance = zeros( Float64, agent.noOfArms )
    agent.count             = zeros( Int64, agent.noOfArms )
    agent.ucbIndices        = zeros( Float64, agent.noOfArms )

    nothing
end

function info_str( agent::UCBV, latex::Bool )
    if latex
        return @sprintf( "UCB-V(\$\\zeta=%3.2f,c=%3.2f,b=%3.2f\$)", agent.ζ, agent.c, agent.b )
    else
        return @sprintf( "UCB-V(ζ=%3.2f,c=%3.2f,b=%3.2f)", agent.ζ, agent.c, agent.b )
    end
end

"""
    Bayes-UCB Implementation
    Based on: Kaufmann, E., Cappé, O., & Garivier, A. (2012). On Bayesian upper confidence bounds for bandit problems. International Conference on Artificial Intelligence and Statistics, 592–600.
    Using 1 - 1/t as quantile
"""
type BayesUCB <: BanditAlgorithmBase
    noOfArms::Int64
    noOfSteps::Int64
    lastPlayedArm::Int64

    count::Vector{Int64}
    α0::Vector{Int64}
    β0::Vector{Int64}
    cummSuccess::Vector{Int64}
    cummFailure::Vector{Int64}

    samplingDist::Vector{Distributions.Beta}

    function BayesUCB( noOfArms::Integer )
        new( noOfArms,
             0,
             0,
             zeros(Int64,noOfArms),
             ones(Int64,noOfArms),
             ones(Int64,noOfArms),
             zeros(Float64,noOfArms),
             zeros(Float64,noOfArms),
             fill(Distributions.Beta(1,1),noOfArms)
        )
    end

    function BayesUCB( armParams::Array{Tuple{Int64,Int64},1} )
        _noOfArms   = length( armParams )
        _priorDist  = [ Distributions.Beta(armParams[idx][1],armParams[idx][2]) for idx=1:_noOfArms ]
        new( _noOfArms,
             0,
             0,
             zeros(Int64,noOfArms),
             [ armParams[idx][1] for idx=1:_noOfArms ],
             [ armParams[idx][2] for idx=1:_noOfArms ],
             zeros(Float64,_noOfArms),
             zeros(Float64,_noOfArms),
             _priorDist
        )
    end
end

function get_arm_index( agent::BayesUCB )
    if any(agent.count.==0)
        agent.lastPlayedArm =  rand( find(agent.count.==0) )
    else
        agent.lastPlayedArm = findmax(  map(dist->quantile(dist,1-1/agent.noOfSteps),agent.samplingDist) )[2]
    end
    return agent.lastPlayedArm
end

function update_reward!( agent::BayesUCB, r::Int64 )
    # Update S and F
    agent.cummSuccess[agent.lastPlayedArm] += (r==0?0:1)
    agent.cummFailure[agent.lastPlayedArm] += (r==0?1:0)

    # Update Distributions
    agent.samplingDist[agent.lastPlayedArm] = Distributions.Beta(
                                                agent.cummSuccess[agent.lastPlayedArm]+agent.α0[agent.lastPlayedArm],
                                                agent.cummFailure[agent.lastPlayedArm]+agent.β0[agent.lastPlayedArm]
                                            )

    # Update play count for arm
    agent.count[agent.lastPlayedArm] += 1

    # Update time steps
    agent.noOfSteps += 1

    nothing
end

function reset!( agent::BayesUCB )
    agent.noOfSteps     = 0
    agent.lastPlayedArm = 0

    agent.count         = zeros( Int64, agent.noOfArms )
    agent.cummSuccess   = zeros( Float64, agent.noOfArms )
    agent.cummFailure   = zeros( Float64, agent.noOfArms )
    agent.samplingDist  = fill( Distributions.Beta(1,1), agent.noOfArms )

    nothing
end

function info_str( agent::BayesUCB, latex::Bool )
    return @sprintf( "Bayes UCB" )
end

"""
    KL-UCB Implementation
    Based on: Garivier, A., & Cappe, O. (2011). The KL-UCB Algorithm for Bounded Stochastic Bandits and Beyond. In CoLT (Vol. 24).
"""
type KLUCB <: BanditAlgorithmBase
    noOfArms::Int64
    noOfSteps::Int64
    lastPlayedArm::Int64

    cummReward::Vector{Float64}
    count::Vector{Int64}
    ucbIndices::Vector{Float64}

    function KLUCB( noOfArms::Int )
        new( noOfArms,
             0,
             0,
             zeros(Float64,noOfArms),
             zeros(Int64,noOfArms),
             zeros(Float64,noOfArms)
        )
    end
end

function get_arm_index( agent::KLUCB )
    if any(agent.count.==0)
        agent.lastPlayedArm =  rand( find(agent.count.==0) )
    else
        agent.lastPlayedArm = findmax(agent.ucbIndices)[2]
    end

    return agent.lastPlayedArm
end

function update_reward!( agent::KLUCB, r::Real )

    # Update cummulative reward
    agent.cummReward[agent.lastPlayedArm] += r

    # Update play count for arm
    agent.count[agent.lastPlayedArm] += 1

    # Update number of steps played
    agent.noOfSteps += 1

    # Update UCB indices
    agent.ucbIndices = find_q.( agent.cummReward, agent.count, agent.noOfSteps )

    nothing
end

function reset!( agent::KLUCB )
    agent.noOfSteps     = 0
    agent.lastPlayedArm = 0

    agent.cummReward    = zeros( Float64, agent.noOfArms )
    agent.count         = zeros( Int64, agent.noOfArms )
    agent.ucbIndices    = zeros( Float64, agent.noOfArms )

    nothing
end

function info_str( agent::KLUCB, latex::Bool )
    return @sprintf( "KL-UCB" )
end

# Utility functions
function kl_div( p::Float64, q::Float64 )
    if p==q
        return 0.00
    elseif p == 0.0
        return -log(1-q)
    elseif p == 1.0
        return -log(q)
    else
        return p*log(p/q) + (1-p)*log((1-p)/(1-q))
    end
end

function kl_div( X1::Distributions.Bernoulli, X2::Distributions.Bernoulli )
   kl_div( X1.p, X2.p )
end

function find_q( sa, na, t; ϵ = 1e-5 )
    lQ = sa/na
    uQ = min( 1, sa/na+√(log(t)/na/2))      # Applying Pinkser's Inequality
    q̄ = 1.0
    q  = (lQ+uQ)/2
    while abs(q̄-q) > ϵ
        q̄  = q
        # print( @sprintf( "lQ = %f, q = %f, uQ = %f", lQ, q, uQ ) )
        klDiv = kl_div( sa/na, q )
        # println( @sprintf( "    kl = %f, d = %f", klDiv, log(t)/na ))
        if klDiv < (log(t)/na)
            lQ = q
        else
            uQ = q
        end
        q  = (lQ+uQ)/2
    end
    return lQ
end

"""
    M-UCB
    Based on: Cao, Y., Zheng, W., Kveton, B., & Xie, Y. (2018). Nearly Optimal Adaptive Procedure for Piecewise-Stationary Bandit: a Change-Point Detection Approach. Retrieved from http://arxiv.org/abs/1802.03692
"""

type MUCB <: BanditAlgorithmBase
    noOfArms::Int64
    noOfSteps::Int64
    lastPlayedArm::Int64

    # Algorithm params
    w::Int              # Test window length - Should be a even number
    b::Float64          # Test threshold
    γ::Float64          # Expoloration parameter

    # Internal
    τ::Int              # Last detected change point
    _ucb::UCB1          # UCB agent to run
    rew::Vector{Vector{Float64}}    # Past rewards from each arm

    function MUCB( noOfArms::Int, w::Int, b::Float64, γ::Float64 )
        new(
            noOfArms,
            0,
            0,
            w,
            b,
            γ,
            0,
            UCB1(noOfArms),
            [ Vector{Float64}() for k = 1:noOfArms ]
        )
    end
end

function get_arm_index( agent::MUCB )
    A = (agent.noOfSteps-agent.τ) % floor(agent.noOfArms/agent.γ)
    if A < agent.noOfArms   # Uniform Exploration
        agent.lastPlayedArm = A+1
        agent._ucb.lastPlayedArm = agent.lastPlayedArm
    else
        agent.lastPlayedArm = get_arm_index( agent._ucb )
    end
    return agent.lastPlayedArm
end

function update_reward!( agent::MUCB, r::Real )
    agent.noOfSteps += 1
    push!( agent.rew[agent.lastPlayedArm], r )
    update_reward!( agent._ucb, r )
    if length(agent.rew[agent.lastPlayedArm]) > agent.w # If the window length is reached
        # Then discard the oldest reward
        shift!( agent.rew[agent.lastPlayedArm] )
        # Also do the test
        if mucb_cd( agent.rew[agent.lastPlayedArm], agent.b )   # If test detects a change
            # Reset UCB1
            reset!( agent._ucb )
            # Update last detection point
            agent.τ = agent.noOfSteps
            agent.rew = [ Vector{Float64}() for k = 1:agent.noOfArms ]
        end
    end
    nothing
end

function reset!( agent::MUCB )
    agent.noOfSteps     = 0
    agent.lastPlayedArm = 0
    agent.τ             = 0
    agent.rew           = [ Vector{Float64}() for k = 1:agent.noOfArms ]

    reset!( agent._ucb )    # Reset internal UCB1 agent also
end

function info_str( agent::MUCB, latex::Bool )
    if latex
        return @sprintf( "M-UCB( w = %d, b = %4.3f, \$\\gamma\$ = %4.3f )", agent.w, agent.b, agent.γ )
    else
        return @sprintf( "M-UCB( w = %d, b = %4.3f, γ = %4.3f )", agent.w, agent.b, agent.γ )
    end
end

#=
    Change Detection algorithm for M-UCB
        b - detection threshold
        rew - list of rewards; must be of length 'w' an even number
=#
function mucb_cd( rew::Vector{Float64}, b::Float64 )
    l = length( rew )
    assert( l%2 == 0 )
    return abs( sum(rew[1:Int(l/2)])-sum(rew[Int(l/2)+1:end]) ) > b
end


## Bisection methods to find a root.

"""
    bisection(f,l,r)

Find the root of a function via a bisection approach. This is a
  pedagogical implementation that can be improved.
  ## Add link to "julia bisection floating point"
"""
function bisection(f::Function, l::Real, r::Real)
  if r < l
    l,r = r,l
  end
  fl = f(l)
  fr = f(r)
  if fl*fr == 0
    al = abs(fl)
    ar = abs(fr)
    if al < ar
      return l
    else
      return r
    end
  end

  @assert(fl*fr < 0) # make sure the signs are different
  while r-l >= max(min(eps(l),eps(r)),eps(1.0)) # while there is still floating point space left.
    m = l/2 + r/2
    fm = f(m)
        @show l, m, r, fl, fm, fr
    if abs(fm) <= eps(1.0)
      return m
    elseif sign(fl) == sign(fm)
      l = m
      fl = fm
    else
      @assert(fm*fl <= 0)
      r = m
      frm = fm
    end
  end
  return l
end
l = bisection(x -> x-5, 0.0, 12.0)
##
l = bisection(x -> x-5, 0.0, 10.0)


# complexspace.jl
#
# Defines the immutable ComplexSpace for a standard complex vector space C^n and is
# characterized by its dimension and whether or not it is the dual space. Tensors with
# ComplexSpace as index spaces make a distinction between covariant and contravariant indices.

# ComplexSpace:
#---------------
immutable ComplexSpace <: EuclideanSpace{ℂ}
  d::Int
  dual::Bool
end
ComplexSpace(d::Int) = ComplexSpace(d, false)
^(::Type{ℂ},d::Int) = ComplexSpace(d)

# Corresponding methods:
dim(V::ComplexSpace) = V.d
dual(V::ComplexSpace) = ComplexSpace(V.d, !V.dual)
cnumber(V::ComplexSpace) = ComplexSpace(1, V.dual)
cnumber(::Type{ComplexSpace}) = ComplexSpace(1)
iscnumber(V::ComplexSpace) = dim(V)==1

# Show methods
Base.show(io::IO, V::ComplexSpace) = print(io, V.dual ? "ℂ^$(V.d)*" : "ℂ^$(V.d)")

# direct sum of ComplexSpaces
directsum(V1::ComplexSpace, V2::ComplexSpace) = (V1.dual==V2.dual ? ComplexSpace(V1.d+V2.d, V1.dual) : throw(SpaceError("Direct sum of a vector space and its dual do not exist")))

# fusing and splitting ComplexSpaces
fuse(V1::ComplexSpace,V2::ComplexSpace,V::ComplexSpace) = dim(V1)*dim(V2)==dim(V)

# # basis and basisvector
# typealias ComplexBasisVector BasisVector{ComplexSpace,Int} # use integer from 1 to dim as identifier
# typealias ComplexBasis Basis{ComplexSpace}
#
# Base.length(B::ComplexBasis) = dim(space(B))
# Base.start(B::ComplexBasis) = 1
# Base.next(B::ComplexBasis, state::Int) = (EuclideanBasisVector(space(B),state),state+1)
# Base.done(B::ComplexBasis, state::Int) = state>length(B)
#
# Base.to_index(b::ComplexBasisVector) = b.identifier