# This is the main output file for PS3 Question 1
# We write down matrices T, Z, Q, H of each process
# α_{t} = T * α_{t-1} + η_{t} ~ N(0, Q)
# y_{t} = Z * α_{t} + ϵ_{t} ~ N(0, H)

## Packages
using Random
using Distributions
using Plots
using Optim
using ForwardDiff
using LinearAlgebra

## Files
include("KalmanFilter.jl")

rng = Random.seed!(1234);
n = 100; # number of periods

############################## Part 1 (AR(1)) #############################
## Question 1.1: AR(1) process (x_{t} = ρ * x_{t-1} + ϵ_{t})
# Parameters
ρ_AR1 = 0.6;
σ_AR1 = 0.5;

# Simulate AR(1) process
y_AR1 = repeat([reshape([0.0], 1, 1)], n); # [0.0] is a Vector{Float64}, 'reshape' changes it to Matrix{Float64}
[y_AR1[i] = [ρ_AR1] * y_AR1[i - 1] + σ_AR1 * randn(rng, 1) for i in 2:n];

## Getting Matrices to apply kalman filter
T_AR1 = reshape([ρ_AR1], 1, 1);
Z_AR1 = reshape([1.0], 1, 1);
Q_AR1 = reshape([σ_AR1^2], 1, 1);
H_AR1 = reshape([0.0], 1, 1);
α0_AR1 = reshape(mean(y_AR1), 1, 1);

## Applying kalman_filter
αhat_AR1, v_AR1, yhat_AR1, P_AR1, F_AR1, G_AR1, K_AR1  = KalmanFilter(T_AR1, Z_AR1, Q_AR1, H_AR1, n, α0_AR1, y_AR1)

## Plots
vec_y_AR1 = [y_AR1[i][1, 1] for i in 1:n]
vec_yhat_AR1 = [yhat_AR1[i][1, 1] for i in 2:n]

plot(vec_y_AR1, xlabel = " Time ", label =" Simulated AR(1) ", legend=:bottomright)
plot!(vec_yhat_AR1, label = "Estimate from Kalman Filter")
savefig("figs/AR1.png")

## Writing the likelhood function
# For the AR1 process we need to estimate two Parameters ρ and σ:
function log_L_AR1(θ::Vector)  # θ = [ρ, log_σ]
    ρ1 = θ[1]
    σ1 = exp(θ[2])
    T1 = reshape([ρ1], 1, 1)
    Q1 = reshape([σ1^2], 1, 1)

    v1 = [KalmanFilter(T1, Z_AR1, Q1, H_AR1, n, α0_AR1, y_AR1)[2][i][1] for i in 1:n]
    F1 = [KalmanFilter(T1, Z_AR1, Q1, H_AR1, n, α0_AR1, y_AR1)[5][i][1] for i in 1:n]
    Ft = abs.(det.(vec(F1)))

    return sum([-0.5 * log(Ft[i]) - 0.5 * v1[i]' * inv(Ft[i]) * v1[i] for i in 1:n])
end

## Optimization
res_AR1 = optimize(θ->-log_L_AR1(θ), [-0.6, -0.69])
ρ_hat_AR1, σ_hat_AR1 = Optim.minimizer(res_AR1)[1], exp(Optim.minimizer(res_AR1)[2]) # (0.61, 0.48)

############################## Part 2 (AR(2)) #############################
## Question 1.2: AR(2) process (x_{t} = ρ1 * x_{t-1} + ρ2 * x_{t-2} + ϵ_{t})
# Parameters
ρ1_AR2 = 0.3
ρ2_AR2 = 0.4
σ_AR2 = 0.5

# Simulate AR(2) process
y_AR2 = repeat([reshape([0.0], 1, 1)], n);
[y_AR2[i] = [ρ1_AR2] * y_AR2[i - 1] + [ρ2_AR2] * y_AR2[i - 2] + σ_AR2 * randn(rng, 1) for i in 3:n];

## Getting Matrices to apply kalman filter
T_AR2 = [ρ1_AR2 ρ2_AR2; 1.0 0.0];
Z_AR2 = [1.0 0.0];
Q_AR2 = [σ_AR2^2 0.0; 0.0 0.0];
H_AR2 = reshape([0.0], 1, 1);
α0_AR2 = [mean(y_AR2); mean(y_AR2)];

αhat_AR2, v_AR2, yhat_AR2, P_AR2, F_AR2, G_AR2, K_AR2 = KalmanFilter(T_AR2, Z_AR2, Q_AR2, H_AR2, n, α0_AR2, y_AR2)

vec_y_AR2 = [y_AR2[i][1, 1] for i in 3:n]
vec_y_hat_AR2 = [yhat_AR2[i][1, 1] for i in 3:n]

plot(vec_y_AR2, xlabel = " Time ", label =" Simulated AR(2) ", legend=:bottomright)
plot!(vec_y_hat_AR2, label = "Estimate from Kalman Filter")
savefig("figs/AR2.png")

## Writing the likelhood function
# For the AR2 process we need to estimate three Parameters ρ1, ρ2, and σ:
function log_L_AR2(θ::Vector)  # θ = [ρ1, ρ2, log_σ]
    ρ1 = θ[1]
    ρ2 = θ[2]
    σ1 = exp(θ[3])
    T1 = [ρ1 ρ2; 1.0 0.0]
    Q1 = [σ1^2 0.0; 0.0 0.0]

    v1 = [KalmanFilter(T1, Z_AR2, Q1, H_AR2, n, α0_AR2, y_AR2)[2][i][1] for i in 1:n]
    F1 = [KalmanFilter(T1, Z_AR2, Q1, H_AR2, n, α0_AR2, y_AR2)[5][i][1] for i in 1:n]
    Ft = abs.(det.(vec(F1)))

    return sum([-0.5 * log(Ft[i]) - 0.5 * v1[i]' * inv(Ft[i]) * v1[i] for i in 3:n])
end

## Optimization
res_AR2 = optimize(θ->-log_L_AR2(θ), [0.3, 0.7, -0.6])
ρ1_hat_AR2,ρ2_hat_AR2,σ_hat_AR2 = Optim.minimizer(res_AR2)[1],Optim.minimizer(res_AR2)[2],exp(Optim.minimizer(res_AR2)[3])
# (0.26, 0.48, 0.54)

############################## Part 3 (MA(1)) #############################
## Question 1.3: MA(1) process (x_{t} = ϵ_{t} + ρ * ϵ_{t-1})
# Parameters
ρ_MA1 = 0.7
σ_MA1 = 0.8

# Simulate MA(1) process
ϵ_MA1 = rand(Normal(0, σ_MA1), n)
y_MA1 = repeat([reshape([0.0], 1, 1)], n);
[y_MA1[i] = reshape([ϵ_MA1[i]], 1,1) + [ρ_MA1]*ϵ_MA1[i-1] for i in 2:n]

## Getting Matrices to apply kalman filter
T_MA1 = [0.0 0.0; 1.0 0.0]
Z_MA1 = [1.0 ρ_MA1]
Q_MA1 = [σ_MA1^2 0.0; 0.0 0.0]
H_MA1 = reshape([0.0], 1, 1)

α0_MA1 = reshape([0.0; 0.0],2,1);
αhat_MA1, v_MA1, yhat_MA1, P_MA1, F_MA1, G_MA1, K_MA1 = KalmanFilter(T_MA1, Z_MA1, Q_MA1, H_MA1, n, α0_MA1, y_MA1)

vec_y_MA1 = [y_MA1[i][1, 1] for i in 2:n]
vec_y_hat_MA1= [yhat_MA1[i][1, 1] for i in 2:n]

plot(vec_y_MA1, xlabel = " Time ", label =" Simulated MA(1) ", legend=:bottomright)
plot!(vec_y_hat_MA1, label = "Estimate from Kalman Filter")
savefig("figs/MA1.png")

## Writing the likelhood function
# For the MA1 process we need to estimate two Parameters ρ and σ:
function log_L_MA1(θ::Vector)  # θ = [ρ, log_σ]
    ρ1 = θ[1]
    σ1 = exp(θ[2])
    Z1 = [1.0 ρ1]
    Q1 = [σ1^2 0.0; 0.0 0.0]

    v1 = [KalmanFilter(T_MA1, Z1, Q1, H_MA1, n, α0_MA1, y_MA1)[2][i][1] for i in 1:n]
    F1 = [KalmanFilter(T_MA1, Z1, Q1, H_MA1, n, α0_MA1, y_MA1)[5][i][1] for i in 1:n]
    Ft = abs.(det.(vec(F1)))

    return sum([-0.5 * log(Ft[i]) - 0.5 * v1[i]' * inv(Ft[i]) * v1[i] for i in 2:n])
end

## Optimization
res_MA1 = optimize(θ->-log_L_MA1(θ), [0.6, -0.6])
ρ_hat_MA1,σ_hat_MA1 = Optim.minimizer(res_MA1)[1],exp(Optim.minimizer(res_MA1)[2]) # (0.61, 0.73)

############################## Part 4 (Random Walk) #############################
## Question 1.4: Random Walk process
# Parameters
σ_ϵ = 0.8
σ_η = 0.6

# Simulate MA(1) process
ϵ_RW = rand(Normal(0, σ_ϵ), n)
η_RW = rand(Normal(0, σ_η), n)

μ_RW = repeat([reshape([0.0], 1, 1)], n)
[μ_RW[i] = μ_RW[i-1] + reshape([η_RW[i]], 1,1) for i in 2:n]

y_RW = repeat([reshape([0.0], 1, 1)], n)
[y_RW[i] = μ_RW[i] + reshape([ϵ_RW[i]], 1,1) for i in 1:n]

## Getting Matrices to apply kalman filter
T_RW = reshape([1.0], 1, 1)
Z_RW = copy(T_RW)
Q_RW = reshape([σ_η^2], 1, 1)
H_RW = reshape([σ_ϵ^2], 1, 1)

α0_RW = reshape([0.0],1,1);
αhat_RW, v_RW, yhat_RW, P_RW, F_RW, G_RW, K_RW = KalmanFilter(T_RW, Z_RW, Q_RW, H_RW, n, α0_RW, y_RW)

vec_y_RW = [y_RW[i][1, 1] for i in 2:n]
vec_y_hat_RW= [yhat_RW[i][1, 1] for i in 2:n]

plot(vec_y_RW, xlabel = " Time ", label =" Simulated Random Walk ", legend=:bottomright)
plot!(vec_y_hat_RW, label = "Estimate from Kalman Filter")
savefig("figs/RW.png")

## Writing the likelhood function
# For the Random Walk process we need to estimate two Parameters σ_ϵ and σ_η:
function log_L_RW(θ::Vector)  # θ = [log_σ_ϵ, log_σ_η]
    σ_ϵ = exp(θ[1])
    σ_η = exp(θ[2])
    Q1 = reshape([σ_η^2], 1, 1)
    H1 = reshape([σ_ϵ^2], 1, 1)

    v1 = [KalmanFilter(T_RW, Z_RW, Q1, H1, n, α0_RW, y_RW)[2][i][1] for i in 1:n]
    F1 = [KalmanFilter(T_RW, Z_RW, Q1, H1, n, α0_RW, y_RW)[5][i][1] for i in 1:n]
    Ft = abs.(det.(vec(F1)))

    return sum([-0.5 * log(Ft[i]) - 0.5 * v1[i]' * inv(Ft[i]) * v1[i] for i in 2:n])
end

## Optimization
res_RW = optimize(θ->-log_L_RW(θ), [-0.3, -0.6])
σ_ϵ_hat, σ_η_hat= exp(Optim.minimizer(res_RW)[1]), exp(Optim.minimizer(res_RW)[2]) # (0.70, 0.65)


export mean_motion
"""
Compute the mean motion given a semi-major axis.

Arguments:
- `a::Real`: Semi-major axis. [m]
- `use_degrees:Bool`: If `true` returns result in units of degrees
- `GM::Real`: Gravitational constant of central body. Defaults to `SatelliteDynamics.GM_EARTH` if none is provided.

Returns:
- `n::Real`: Orbital mean motion. [rad/s] or [deg/s]
"""
function mean_motion(a::Real; use_degrees::Bool=false, GM::Real=GM_EARTH)
    n = sqrt(GM/a^3)

    if use_degrees
        n *= 180.0/pi
    end

    return n
end

export semimajor_axis
"""
Calculate semi-major axis given mean-motion.

Arguments:
- `n::Real`: Orbital mean motion. [rad/s] or [deg/s]
- `use_degrees:Bool`: If `true` interpret input as being in degrees.
- `GM::Real`: Gravitational constant of central body. Defaults to `SatelliteDynamics.GM_EARTH` if none is provided.

Returns:
- `a::Real`: Semi-major axis. [m]
"""
function semimajor_axis(n::Real; use_degrees::Bool=false, GM::Real=GM_EARTH)
    if use_degrees
        n *= pi/180.0
    end

    a = (GM/n^2)^(1.0/3.0)

    return a
end

export orbit_period
"""
Compute the satellite orbital period given the semi-major axis.

Arguments:
- `a::Real`: Semi-major axis. [m]
- `GM::Real`: Gravitational constant of central body. Defaults to `SatelliteDynamics.GM_EARTH` if none is provided.

Returns:
- `T::Real`: Orbital period. [s]
"""
function orbit_period(a::Real; GM::Real=GM_EARTH)
    return 2.0*pi*sqrt(a^3/GM)
end

export sun_sync_inclination
"""
Compute the required inclination for a Sun-synchronous Earth orbit.

Algorithm assumes the nodal precession is entirely due to the J2 perturbation, and no other perturbations are considered.

The inclination is computed using a first-order, non-iterative approximation.

Arguments:
- `a::Real`: Semi-major axis. [m]
- `e::Real`: Eccentricity. [dimensionless]
- `use_degrees:Bool`: If `true` interpret output will be returned in degrees.

Returns:
- `iss::Real`: Requierd inclination for a sun-synchronous orbit. [rad] or [deg]
"""
function sun_sync_inclination(a::Real, e::Real; use_degrees::Bool=false)
    # Compute the required RAAN precession of a sun-synchronous orbit
    OMEGA_DOT_SS = 2*pi/365.2421897/86400.0

    # Inclination required for sun-synchronous orbits
    i_ss = acos(-2*a^(7/2)*OMEGA_DOT_SS*(1-e^2)^2/(3*(R_EARTH^2)*J2_EARTH*sqrt(GM_EARTH)))

    if use_degrees == true
        i_ss *= 180.0/pi
    end

    return i_ss
end

export anomaly_eccentric_to_mean
"""
Convert eccentric anomaly into mean anomaly.

Arguments:
- `E::Real`: Eccentric anomaly. [rad] or [deg]
- `e::Real`: Eccentricity. [dimensionless]
- `use_degrees:Bool`: If `true` interpret input will be interpreted as being in degrees, and output will be returned in degrees.

Returns:
- `M::Real`: Mean anomaly. [rad] or [deg]
"""
function anomaly_eccentric_to_mean(E::Real, e::Real; use_degrees::Bool=false)
    # Convert degree input
    if use_degrees == true
        E *= pi/180.0
    end

    # Convert eccentric to mean
    M = E - e*sin(E)

    # Convert degree output
    if use_degrees == true
        M *= 180.0/pi
    end

    return M
end

export anomaly_mean_to_eccentric
"""
Convert mean anomaly into eccentric anomaly.

Arguments:
- `M::Real`: Mean anomaly. [deg] or [deg]
- `e::Real`: Eccentricity. [dimensionless]
- `use_degrees:Bool`: If `true` interpret input will be interpreted as being in degrees, and output will be returned in degrees.

Returns:
- `E::Real`: Eccentric anomaly. [rad] or [deg]
"""
function anomaly_mean_to_eccentric(M::Real, e::Real; use_degrees::Bool=false)
    # Convert degree input
    if use_degrees == true
        M *= pi/180.0
    end

    # Convert mean to eccentric
    max_iter = 15
    epsilson = eps(Float64)*100.0

    # Initialize starting values
    M = mod(M, 2.0*pi)
    if e < 0.8
        E = M
    else
        E = pi
    end

    # Initialize working variable
    f = E - e*sin(E) - M
    i = 0

    # Iterate until convergence
    while abs(f) > epsilson
        f = E - e*sin(E) - M
        E = E - f / (1.0 - e*cos(E))

        # Increase iteration counter
        i += 1
        if i == max_iter
            error("Maximum number of iterations reached before convergence.")
        end
    end

    # Convert degree output
    if use_degrees == true
        E *= 180.0/pi
    end

    return E
end

export sOSCtoCART
"""
Given an orbital state expressed in osculating orbital elements compute the equivalent Cartesean position and velocity of the inertial state.

The osculating elements are assumed to be (in order):
1. _a_, Semi-major axis [m]
2. _e_, Eccentricity [dimensionless]
3. _i_, Inclination [rad]
4. _Ω_, Right Ascension of the Ascending Node (RAAN) [rad]
5. _ω_, Argument of Perigee [ramd]
6. _M_, Mean anomaly [rad]

Arguments:
- x_oe `x::Array{<:Real, 1}`: Osculating orbital elements. See above for desription of the elements and their required order.
- `use_degrees:Bool`: If `true` interpret input will be interpreted as being in degrees, and output will be returned in degrees.
- `GM::Real`: Gravitational constant of central body. Defaults to `SatelliteDynamics.GM_EARTH` if none is provided.

# Returns
- x `x::Array{<:Real, 1}`: Cartesean inertial state. Returns position and velocity. [m; m/s]
"""
function sOSCtoCART(x_oe::Array{<:Real, 1}; use_degrees::Bool=false, GM::Real=GM_EARTH)

    if use_degrees == true
        # Copy and convert input from degrees to radians if necessary
        oe = deepcopy(x_oe)
        oe[3:6] = oe[3:6]*pi/180.0
    else
        oe = x_oe
    end

    # Unpack input
    a, e, i, RAAN, omega, M = oe

    E = anomaly_mean_to_eccentric(M, e)

    # Create perifocal coordinate vectors
    P    = zeros(Float64, 3)
    P[1] = cos(omega)*cos(RAAN) - sin(omega)*cos(i)*sin(RAAN)
    P[2] = cos(omega)*sin(RAAN) + sin(omega)*cos(i)*cos(RAAN)
    P[3] = sin(omega)*sin(i)

    Q    = zeros(Float64, 3)
    Q[1] = -sin(omega)*cos(RAAN) - cos(omega)*cos(i)*sin(RAAN)
    Q[2] = -sin(omega)*sin(RAAN) + cos(omega)*cos(i)*cos(RAAN)
    Q[3] =  cos(omega)*sin(i)

    # Find 3-Dimensional Position
    x = zeros(Float64, 6)
    x[1:3] = a*(cos(E)-e)*P + a*sqrt(1-e*e)*sin(E)*Q
    x[4:6] = sqrt(GM*a)/norm(x[1:3])*(-sin(E)*P + sqrt(1-e*e)*cos(E)*Q)

    return x
end

export sCARTtoOSC
"""
Given a Cartesean position and velocity in the inertial frame, return the
state expressed in terms of  osculating orbital elements.

The osculating elements are assumed to be (in order):
1. _a_, Semi-major axis [m]
2. _e_, Eccentricity [dimensionless]
3. _i_, Inclination [rad]
4. _Ω_, Right Ascension of the Ascending Node (RAAN) [rad]
5. _ω_, Argument of Perigee [ramd]
6. _M_, Mean anomaly [rad]

Arguments:
- x `x::Array{<:Real, 1}`: Cartesean inertial state. Returns position and velocity. [m; m/s]
- `use_degrees:Bool`: If `true` interpret input will be interpreted as being in degrees, and output will be returned in degrees.
- `GM::Real`: Gravitational constant of central body. Defaults to `SatelliteDynamics.GM_EARTH` if none is provided.

# Returns
- x_oe `x::Array{<:Real, 1}`: Osculating orbital elements. See above for desription of the elements and their required order.
"""
function sCARTtoOSC(x::Array{<:Real, 1}; use_degrees::Bool=false, GM::Real=GM_EARTH)

    # Initialize Cartesian Polistion and Velocity
    r = x[1:3]
    v = x[4:6]

    h = cross(r, v)  # Angular momentum vector
    W = h/norm(h)    # Unit vector along angular momentum vector

    i     = atan(sqrt(W[1]*W[1] + W[2]*W[2]), W[3])    # Compute inclination
    OMEGA = atan(W[1], -W[2])                          # Right ascension of ascending node                     # Compute RAAN
    p     = norm(h)*norm(h)/GM                         # Semi-latus rectum
    a     = 1.0/(2.0/norm(r) - norm(v)*norm(v)/GM)     # Semi-major axis
    n     = sqrt(GM/(a^3))                             # Mean motion

    # Numerical stability hack for circular and near-circular orbits
    # Ensures that (1-p/a) is always positive
    if isapprox(a, p, atol=1e-9, rtol=1e-8)
        p = a
    end

    e     = sqrt(1 - p/a)                              # Eccentricity
    E     = atan(dot(r, v)/(n*a^2), (1-norm(r)/a))     # Eccentric Anomaly
    M     = anomaly_eccentric_to_mean(E, e)            # Mean Anomaly
    u     = atan(r[3], -r[1]*W[2] + r[2]*W[1])         # Mean longiude
    nu    = atan(sqrt(1-e*e)*sin(E), cos(E)-e)         # True Anomaly
    omega = u - nu                                     # Argument of perigee

    # Correct angles to run from 0 to 2PI
    OMEGA = OMEGA + 2.0*pi
    omega = omega + 2.0*pi
    M     = M     + 2.0*pi

    OMEGA = mod(OMEGA, 2.0*pi)
    omega = mod(omega, 2.0*pi)
    M     = mod(M, 2.0*pi)

    # Create Orbital Element Vector
    x_oe = [a, e, i, OMEGA, omega, M]

    # Convert output to degrees if necessary
    if use_degrees == true
        x_oe[3:6] = x_oe[3:6]*180.0/pi
    end

    return x_oe
end


## Packages
using LinearAlgebra
using Plots

## Parameters
α = 0.333;
δ = 0.1;
β = 0.975;
A = 1.05;
nk = 1000; # number of capital grids
tol = 1e-5;
maxDiff = 10.0;
iter = 0;
maxiter = 5000;

## Steady state of capital
kss = ((1/β+δ-1)/(α*A))^(1/(α-1));

## Capital grids
kmin = 0.5*kss;
kmax = 1.5*kss;
k_grid = LinRange(kmin, kmax,nk);
kp_grid = LinRange(kmin, kmax,nk);

# VFI
V_old = zeros(nk);
V_new = zeros(nk);
kpol = zeros(nk);

while maxDiff >= tol && iter < maxiter
    iter = iter+1;
    println("Current iteration is $iter")

    temp = zeros(nk, nk);
    V_old = V_new;

    # Take care of negative consumption
    for i = 1:nk
        for j = 1:nk
            c = A*(k_grid[i]^α)+(1-δ)*k_grid[i]-kp_grid[j];
            if c < 0
                temp[i,j] = -10.0;
            else
                temp[i,j] = log(c)+β*V_old[j];
            end
        end
    end

    # Find max and argmax
    V_new,argmax = findmax(temp,dims=2);

    for i = 1:nk
        kpol[i] = k_grid[argmax[i][2]];
    end

    maxDiff = maximum(abs.(V_new-V_old));
    println("maxDiff = $maxDiff")
end

plot(k_grid,k_grid,label = "", color = :black, linestyle = :dash)
plot!(k_grid,kpol,label = "Capital policy")

plot(k_grid,V_new,label = "Value Function")

# This is a check when δ = 1
# plot!(k_grid,α*β*A*k_grid.^α,label = "true")


using NumericalMethodsforEngineers, DataFrames, Plots
pyplot(size=(700,700))

ProjDir = dirname(@__FILE__)
cd(ProjDir) #do

  x = [1.0, 3.0, 6.0, 5.0]
  y = [1.0, 5.0, 10.0, 9.0]
  xi = [2.0, 4.5]
  (dfin, dfxi) = lagrangianpolynomial(length(x), x, y, xi)

  xint = 1:0.1:5
  (dfin, dfxint) = lagrangianpolynomial(length(x), x, y, collect(xint))

  dfin |> display
  println()

  dfxi |> display
  println()

  dfxint |> display
  println()

  p = plot(dfxint[:xi], dfxint[:yi], line=(:path, 1), label="interpolated curve")
  scatter!(p, dfin[:x], dfin[:y], marker=(:circle, 4), label="input points", color=:blue)
  scatter!(p, dfxi[:xi], dfxi[:yi], marker=(:star, 8), color=:red, label="interpolated points")

  plot(p)
  savefig("ex.5.1.png")
  gui()

#end


eval_for(p::Poly, x₀) = p.a₀ + (p.n > 0 ? sum([coef*x₀^k for (k,coef) in p.terms]) : 0)
eval_for(p::Poly, x, x₀) = isequal(p.x, x) ? eval_for(p, x₀) : eval_for(sym(p), x, x₀)
eval_for(p, x, x₀) = substitute(p, Dict(x => x₀))

(p::Poly)(x₀) = eval_for(p, x₀)

"""
    solve_newton is a symbolic Newton-Ralphson solver
    f is a symbolic equation to be solved (f ~ 0)
    x is the variable to solve
    x₀ is the initial guess
"""
function solve_newton(f, x, x₀; abstol=1e-8, maxiter=50)
    xₙ = Complex(x₀)
    ∂f = value(derivative(f, x))

    for i = 1:maxiter
        xₙ₊₁ = xₙ - eval_for(f, x, xₙ) / eval_for(∂f, x, xₙ)

        if abs(xₙ₊₁ - xₙ) < abstol
            return xₙ₊₁
        else
            xₙ = xₙ₊₁
        end
    end
    return nothing
end

"""
    find_roots returns all the real and complex roots of a polynomial (poly)
    for variable x

    the output is rs, sz, where rs is a list of real roots and zs is a list of
    complex roots

    if poly is not a polynomial, the number of desired roots (n) needs to be
    provided. For example, find_roots(sin(x), x, 10) returns 10 different, but
    not necessarily sequential, multiples of π
"""
function find_roots(p, x, n=-1; abstol=1e-8)
    n = (n == -1 ? degree(p, x) : n)

    rs = Float64[]
    zs = Complex[]

    while n > 0
        z = solve_newton(p, x, Complex(rand(),rand()))
        if z != nothing
            if abs(imag(z)) < abstol
                r = real(z)
                push!(rs, r)
                if p isa Poly
                    p = p / gen_poly(x, -r, 1)
                else
                    p = p / (x - r)
                end
                n -= 1
            else
                if abs(real(z)) < abstol
                    z = Complex(0, imag(z))
                end
                append!(zs, [z, conj(z)])
                #poly = poly / ((x-z)*(x-conj(z)))
                if p isa Poly
                    p = p / gen_poly(x, abs2(z), -2*real(z), 1)
                else
                    p = p / (x^2 - 2x*real(z) + abs2(z))
                end
                n -= 2
            end
        end
    end
    sort(rs), zs
end


# okay, this is for my autodiff package which I'm going to do briefly just to try to understand
#how this actually works out, as it's my next julia project. It shuold be very very educatoinal to be honest as it'll teach us more abuot the type system and the maths behind autodiff as well as using macros and a bit about compilation if we are going to get it to work as we should

# the aims are as follows: we implement autodiff types and the autodiff dual number functions for all the common mathemtical operations
# we implement a macro which can then take a function written with thos standard and automagically compile it into an autodiffed one
# should be fairl straightforward, but requires a deep understanding. now let's try to gain that understanding

# first our dual number type

import Base: *

type Dual{T} <: Real
	r::T # the real part
	d::T # the derivative part

	#function Dual(r::T, d::T)
	#	return new(r,d)
	#end

	#function Dual(r::T)
	#	return new(r,0)	#this sets the default derivative as zero, which seems reasonable atm
	#end
end

#some simple helper functoins
function getReal(num::Dual)
	return num.r
end

function getDerivative(num::Dual)
	return num.d
end

# okay, now we're going to implement finite differences so I can check my answers

function finite_differences(f::Function, start, step)
	return (f(start+step) - f(step))/step
end

# okay, let's implement our AD functions using overloaded operators

#plus function
function +(u::Dual, v::Dual)
	return Dual((u.r + v.r), (u.d+v.d))
end

#multiply function
function *(u::Dual, v::Dual)
	return Dual((u.r*v.r), (u.d *v.r + u.r * v.d))
end

function -(u::Dual, v::Dual)
	return Dual((u.r-v.r), (u.d-r.d))
end

function negate(u::Dual)
	return Dual(-u.r, -u.d)
end

function signum(u::Dual)
	return Dual(signum(u.r),0)
end

function abs(u::Dual)
	return dual(abs(u.r), u.d*(signum(u.r)))
end

function /(u::Dual, v::Dual)
	return Dual((u.r/v.r), ((u.d*v.r - u.r*v.d)/(v.r^2)))
end

function exp(u::Dual)
	return Dual(exp(u.r), (u.d*exp(u.r)))
end

function sqrt(u::Dual)
	return Dual(sqrt(u.r), (u.d/(2*sqrt(u.r))))
end

function log(u::Dual)
	return Dual(log(u.r), (u.r/u.d))
end

function sin(u::Dual)
	return Dual(sin(u.r), (u.d*cos(u.r)))
end

function cos(u::Dual)
	return Dual(cos(u.r), (-u.d*sin(u.r)))
end

function tan(u::Dual)
	return Dual(tan(u.r), (1/(cos(u.r)^2)))
end

function asin(u::Dual)
	return Dual(asin(u.r), (u.d/(sqrt(1-u.r)^2)))
end

function acos(u::Dual)
	return Dual(acos(u.r), (-u.d/sqrt(1-u.r)^2))
end

function atan(u::Dual)
	return Dual(atan(u.r), (u.d/(1+u.r)^2))
end

function sinh(u::Dual)
	return Dual(sinh(u.r), (u.d*cosh(u.r)))
end

function cosh(u::Dual)
	return Dual(cosh(u.r), (u.d*sinh(u.r)))
end

function tanh(u::Dual)
	return Dual(tanh(u.r), (u.d * (1- (tanh(u.r))^2)))
end

function asinh(u::Dual)
	return Dual(asinh(u.r), (u.d/(sqrt(1 + u.r^2))))
end

function acosh(u::Dual)
	return Dual(acosh(u.r), (u.d/sqrt(u.r^2 -1)))
end

function atanh(u::Dual)
	return Dual(atanh(u.r), (u.d/(1-u.r^2)))
end

# okay, let's test this

function f(a,b)
	return a + b
end

a = Dual(4,1)
b = Dual(6,2)

f(a,b)


include("rootfinding.jl")

f(x) = sin(x)
df(x) = cos(x)

x, it = bisect(f, 0.1, 4.0)

x, it = newton(f, df, 4.0)

g(x) = cos(x)
#anim = cobweb(g, 0, 1, 0.5, 50)
#gif(anim, "anim_fps15.gif", fps = 1)

g1(x) = x - x^3 - 4*x^2 + 10
g2(x) = sqrt(complex(10/x - 4*x))
g3(x) = sqrt(10 - x^3)/2
g4(x) = sqrt(10/(4+x))
g5(x) = x - (x^3 + 4*x^2 - 10)/(3*x^2 + 8*x)
#=
x1, it1 = fixedpoint(g1, 1.5)
x2, it2 = fixedpoint(g2, 1.5)
x3, it3 = fixedpoint(g3, 1.5)
x4, it4 = fixedpoint(g4, 1.5)
x5, it5 = fixedpoint(g5, 1.5)
println((x1,it1))
println((x2,it2))
println((x3,it3))
println((x4,it4))
println((x5,it5))

f(x) = cos(x)-x
df(x) = -sin(x) - 1
#anim = viznewton(f, df, -3, 10, 6, 100)
#gif(anim, "newton_fps1.gif", fps = 1)

x, it = secant(f, 6, 3)
=#

x = brent(x->cos(x)-x, 0.0, 4.0)
println(x)
x1 = bisect(x->cos(x)-x, 0.0, 4.0)
println(x1)


#############################
### 2D Coordinate systems ###
#############################
"""
`Polar{T,A}(r::T, θ::A)` - 2D polar coordinates
"""
struct Polar{T,A}
    r::T
    θ::A

    Polar{T, A}(r, θ) where {T, A} = new(r, θ)
end

function Polar(r, θ)
    r2, θ2 = promote(r, θ)

    return Polar{typeof(r2), typeof(θ2)}(r2, θ2)
end

Base.show(io::IO, x::Polar) = print(io, "Polar(r=$(x.r), θ=$(x.θ) rad)")
Base.isapprox(p1::Polar, p2::Polar; kwargs...) = isapprox(p1.r, p2.r; kwargs...) && isapprox(p1.θ, p2.θ; kwargs...)

"`PolarFromCartesian()` - transformation from `AbstractVector` of length 2 to `Polar` type"
struct PolarFromCartesian <: Transformation; end
"`CartesianFromPolar()` - transformation from `Polar` type to `SVector{2}` type"
struct CartesianFromPolar <: Transformation; end

Base.show(io::IO, trans::PolarFromCartesian) = print(io, "PolarFromCartesian()")
Base.show(io::IO, trans::CartesianFromPolar) = print(io, "CartesianFromPolar()")

function (::PolarFromCartesian)(x::AbstractVector)
    length(x) == 2 || error("Polar transform takes a 2D coordinate")

    Polar(hypot(x[1], x[2]), atan(x[2], x[1]))
end

function transform_deriv(::PolarFromCartesian, x::AbstractVector)
    length(x) == 2 || error("Polar transform takes a 2D coordinate")

    r = hypot(x[1], x[2])
    f = x[2] / x[1]
    c = one(eltype(x))/(x[1]*(one(eltype(x)) + f*f))
    @SMatrix [ x[1]/r    x[2]/r ;
              -f*c       c      ]
end
transform_deriv_params(::PolarFromCartesian, x::AbstractVector) = error("PolarFromCartesian has no parameters")

function (::CartesianFromPolar)(x::Polar)
    s,c = sincos(x.θ)
    SVector(x.r * c, x.r * s)
end
function transform_deriv(::CartesianFromPolar, x::Polar)
    sθ, cθ = sincos(x.θ)
    @SMatrix [cθ  -x.r*sθ ;
              sθ   x.r*cθ ]
end
transform_deriv_params(::CartesianFromPolar, x::Polar) = error("CartesianFromPolar has no parameters")

Base.inv(::PolarFromCartesian) = CartesianFromPolar()
Base.inv(::CartesianFromPolar) = PolarFromCartesian()

compose(::PolarFromCartesian, ::CartesianFromPolar) = IdentityTransformation()
compose(::CartesianFromPolar, ::PolarFromCartesian) = IdentityTransformation()

# For convenience
Base.convert(::Type{Polar}, v::AbstractVector) = PolarFromCartesian()(v)
@inline Base.convert(::Type{V}, p::Polar) where {V <: AbstractVector} = convert(V, CartesianFromPolar()(p))
@inline Base.convert(::Type{V}, p::Polar) where {V <: StaticVector} = convert(V, CartesianFromPolar()(p))

#############################
### 3D Coordinate Systems ###
#############################
"""
Spherical(r, θ, ϕ) - 3D spherical coordinates

There are many Spherical coordinate conventions and this library uses a somewhat exotic one.
Given a vector `v` with Cartesian coordinates `xyz`, let `v_xy = [x,y,0]` be the orthogonal projection of `v` on the `xy` plane.

* `r` is the radius. It is given by `norm(v, 2)`.
* `θ` is the azimuth. It is the angle from the x-axis to `v_xy`
* `ϕ` is the latitude. It is the angle from `v_xy` to `v`.

```jldoctest
julia> using CoordinateTransformations

julia> v = randn(3);

julia> sph = SphericalFromCartesian()(v);

julia> r = sph.r; θ=sph.θ; ϕ=sph.ϕ;

julia> v ≈ [r * cos(θ) * cos(ϕ), r * sin(θ) * cos(ϕ), r*sin(ϕ)]
true
"""
struct Spherical{T,A}
    r::T
    θ::A
    ϕ::A

    Spherical{T, A}(r, θ, ϕ) where {T, A} = new(r, θ, ϕ)
end

function Spherical(r, θ, ϕ)
    r2, θ2, ϕ2 = promote(r, θ, ϕ)

    return Spherical{typeof(r2), typeof(θ2)}(r2, θ2, ϕ2)
end

Base.show(io::IO, x::Spherical) = print(io, "Spherical(r=$(x.r), θ=$(x.θ) rad, ϕ=$(x.ϕ) rad)")
Base.isapprox(p1::Spherical, p2::Spherical; kwargs...) = isapprox(p1.r, p2.r; kwargs...) && isapprox(p1.θ, p2.θ; kwargs...) && isapprox(p1.ϕ, p2.ϕ; kwargs...)

"""
Cylindrical(r, θ, z) - 3D cylindrical coordinates
"""
struct Cylindrical{T,A}
    r::T
    θ::A
    z::T

    Cylindrical{T, A}(r, θ, z) where {T, A} = new(r, θ, z)
end

function Cylindrical(r, θ, z)
    r2, θ2, z2 = promote(r, θ, z)

    return Cylindrical{typeof(r2), typeof(θ2)}(r2, θ2, z2)
end

Base.show(io::IO, x::Cylindrical) = print(io, "Cylindrical(r=$(x.r), θ=$(x.θ) rad, z=$(x.z))")
Base.isapprox(p1::Cylindrical, p2::Cylindrical; kwargs...) = isapprox(p1.r, p2.r; kwargs...) && isapprox(p1.θ, p2.θ; kwargs...) && isapprox(p1.z, p2.z; kwargs...)

"`SphericalFromCartesian()` - transformation from 3D point to `Spherical` type"
struct SphericalFromCartesian <: Transformation; end
"`CartesianFromSpherical()` - transformation from `Spherical` type to `SVector{3}` type"
struct CartesianFromSpherical <: Transformation; end
"`CylindricalFromCartesian()` - transformation from 3D point to `Cylindrical` type"
struct CylindricalFromCartesian <: Transformation; end
"`CartesianFromCylindrical()` - transformation from `Cylindrical` type to `SVector{3}` type"
struct CartesianFromCylindrical <: Transformation; end
"`CylindricalFromSpherical()` - transformation from `Spherical` type to `Cylindrical` type"
struct CylindricalFromSpherical <: Transformation; end
"`SphericalFromCylindrical()` - transformation from `Cylindrical` type to `Spherical` type"
struct SphericalFromCylindrical <: Transformation; end

Base.show(io::IO, trans::SphericalFromCartesian) = print(io, "SphericalFromCartesian()")
Base.show(io::IO, trans::CartesianFromSpherical) = print(io, "CartesianFromSpherical()")
Base.show(io::IO, trans::CylindricalFromCartesian) = print(io, "CylindricalFromCartesian()")
Base.show(io::IO, trans::CartesianFromCylindrical) = print(io, "CartesianFromCylindrical()")
Base.show(io::IO, trans::CylindricalFromSpherical) = print(io, "CylindricalFromSpherical()")
Base.show(io::IO, trans::SphericalFromCylindrical) = print(io, "SphericalFromCylindrical()")

# Cartesian <-> Spherical
function (::SphericalFromCartesian)(x::AbstractVector)
    length(x) == 3 || error("Spherical transform takes a 3D coordinate")

    Spherical(hypot(x[1], x[2], x[3]), atan(x[2], x[1]), atan(x[3], hypot(x[1], x[2])))
end
function transform_deriv(::SphericalFromCartesian, x::AbstractVector)
    length(x) == 3 || error("Spherical transform takes a 3D coordinate")
    T = eltype(x)

    r = hypot(x[1], x[2], x[3])
    rxy = hypot(x[1], x[2])
    fxy = x[2] / x[1]
    cxy = one(T)/(x[1]*(one(T) + fxy*fxy))
    f = -x[3]/(rxy*r*r)

    @SMatrix [ x[1]/r   x[2]/r  x[3]/r;
          -fxy*cxy  cxy     zero(T);
           f*x[1]   f*x[2]  rxy/(r*r) ]
end
transform_deriv_params(::SphericalFromCartesian, x::AbstractVector) = error("SphericalFromCartesian has no parameters")

function (::CartesianFromSpherical)(x::Spherical)
    sθ, cθ = sincos(x.θ)
    sϕ, cϕ = sincos(x.ϕ)
    SVector(x.r * cθ * cϕ, x.r * sθ * cϕ, x.r * sϕ)
end
function transform_deriv(::CartesianFromSpherical, x::Spherical{T}) where T
    sθ, cθ = sincos(x.θ)
    sϕ, cϕ = sincos(x.ϕ)
    @SMatrix [cθ*cϕ -x.r*sθ*cϕ -x.r*cθ*sϕ ;
              sθ*cϕ  x.r*cθ*cϕ -x.r*sθ*sϕ ;
              sϕ     zero(T)    x.r * cϕ  ]
end
transform_deriv_params(::CartesianFromSpherical, x::Spherical) = error("CartesianFromSpherical has no parameters")

# Cartesian <-> Cylindrical
function (::CylindricalFromCartesian)(x::AbstractVector)
    length(x) == 3 || error("Cylindrical transform takes a 3D coordinate")

    Cylindrical(hypot(x[1], x[2]), atan(x[2], x[1]), x[3])
end

function transform_deriv(::CylindricalFromCartesian, x::AbstractVector)
    length(x) == 3 || error("Cylindrical transform takes a 3D coordinate")
    T = eltype(x)

    r = hypot(x[1], x[2])
    f = x[2] / x[1]
    c = one(T)/(x[1]*(one(T) + f*f))
    @SMatrix [ x[1]/r   x[2]/r   zero(T) ;
              -f*c      c        zero(T) ;
               zero(T)  zero(T)  one(T)  ]
end
transform_deriv_params(::CylindricalFromCartesian, x::AbstractVector) = error("CylindricalFromCartesian has no parameters")

function (::CartesianFromCylindrical)(x::Cylindrical)
    sθ, cθ = sincos(x.θ)
    SVector(x.r * cθ, x.r * sθ, x.z)
end
function transform_deriv(::CartesianFromCylindrical, x::Cylindrical{T}) where {T}
    sθ, cθ = sincos(x.θ)
    @SMatrix [cθ      -x.r*sθ  zero(T) ;
              sθ       x.r*cθ  zero(T) ;
              zero(T)  zero(T) one(T)  ]
end
transform_deriv_params(::CartesianFromPolar, x::Cylindrical) = error("CartesianFromCylindrical has no parameters")

function (::CylindricalFromSpherical)(x::Spherical)
    sϕ, cϕ = sincos(x.ϕ)
    Cylindrical(x.r*cϕ,x.θ,x.r*sϕ)
end
function transform_deriv(::CylindricalFromSpherical, x::Spherical)
    M1 = transform_deriv(CylindricalFromCartesian(), CartesianFromSpherical()(x))
    M2 = transform_deriv(CartesianFromSpherical(), x)
    return M1*M2
end
transform_deriv_params(::CylindricalFromSpherical, x::Spherical) = error("CylindricalFromSpherical has no parameters")

function (::SphericalFromCylindrical)(x::Cylindrical)
    Spherical(hypot(x.r,x.z),x.θ,atan(x.z,x.r))
end
function transform_deriv(::SphericalFromCylindrical, x::Cylindrical)
    M1 = transform_deriv(SphericalFromCartesian(), CartesianFromCylindrical()(x))
    M2 = transform_deriv(CartesianFromCylindrical(), x)
    return M1*M2
end
transform_deriv_params(::SphericalFromCylindrical, x::Cylindrical) = error("SphericalFromCylindrical has no parameters")

Base.inv(::SphericalFromCartesian)   = CartesianFromSpherical()
Base.inv(::CartesianFromSpherical)   = SphericalFromCartesian()
Base.inv(::CylindricalFromCartesian) = CartesianFromCylindrical()
Base.inv(::CartesianFromCylindrical) = CylindricalFromCartesian()
Base.inv(::CylindricalFromSpherical) = SphericalFromCylindrical()
Base.inv(::SphericalFromCylindrical) = CylindricalFromSpherical()

# Inverse composition
compose(::SphericalFromCartesian,   ::CartesianFromSpherical)   = IdentityTransformation()
compose(::CartesianFromSpherical,   ::SphericalFromCartesian)   = IdentityTransformation()
compose(::CylindricalFromCartesian, ::CartesianFromCylindrical) = IdentityTransformation()
compose(::CartesianFromCylindrical, ::CylindricalFromCartesian) = IdentityTransformation()
compose(::CylindricalFromSpherical, ::SphericalFromCylindrical) = IdentityTransformation()
compose(::SphericalFromCylindrical, ::CylindricalFromSpherical) = IdentityTransformation()

# Cyclic compositions
compose(::SphericalFromCartesian,   ::CartesianFromCylindrical) = SphericalFromCylindrical()
compose(::CartesianFromSpherical,   ::SphericalFromCylindrical) = CartesianFromCylindrical()
compose(::CylindricalFromCartesian, ::CartesianFromSpherical)   = CylindricalFromSpherical()
compose(::CartesianFromCylindrical, ::CylindricalFromSpherical) = CartesianFromSpherical()
compose(::CylindricalFromSpherical, ::SphericalFromCartesian)   = CylindricalFromCartesian()
compose(::SphericalFromCylindrical, ::CylindricalFromCartesian) = SphericalFromCartesian()

# For convenience
Base.convert(::Type{Spherical}, v::AbstractVector) = SphericalFromCartesian()(v)
Base.convert(::Type{Cylindrical}, v::AbstractVector) = CylindricalFromCartesian()(v)

Base.convert(::Type{V}, s::Spherical) where {V <: AbstractVector} = convert(V, CartesianFromSpherical()(s))
Base.convert(::Type{V}, c::Cylindrical) where {V <: AbstractVector} = convert(V, CartesianFromCylindrical()(c))
Base.convert(::Type{V}, s::Spherical) where {V <: StaticVector} = convert(V, CartesianFromSpherical()(s))
Base.convert(::Type{V}, c::Cylindrical) where {V <: StaticVector} = convert(V, CartesianFromCylindrical()(c))

Base.convert(::Type{Spherical}, c::Cylindrical) = SphericalFromCylindrical()(c)
Base.convert(::Type{Cylindrical}, s::Spherical) = CylindricalFromSpherical()(s)


# -*- coding: utf-8 -*-
# # Poisson Equation
#
# $$
# \frac{\partial^2 u}{\partial x^2} = b  \qquad x \in [0,1]
# $$
#
# $$
# u(0) = u(1) = 0, \qquad b = \sin(2\pi x)
# $$

using Plots, BenchmarkTools

function plot_solution(x, u)
    plot([0;x;1],[0;u;0], label="computed")
    scatter!([0;x;1],-sin.(2π*[0;x;1])/(4π^2),label="exact")
end

Δx = 0.05
x = Δx:Δx:1-Δx ## Solve only interior points: the endpoints are set to zero.
N = length(x)

A = zeros(N,N)
for i in 1:N, j in 1:N
    abs(i-j) <= 1 && (A[i,j] +=1)
    i==j          && (A[i,j] -=3)
end

B = sin.(2π*x) * Δx^2
u = A \ B

plot_solution(x, u)

# # SparseArrays

using SparseArrays

Δx = 0.05
x = Δx:Δx:1-Δx ## Solve only interior points: the endpoints are set to zero.
N = length(x)
B = sin.(2π*x) * Δx^2

P = spdiagm( -1 =>    ones(Float64,N-1),
              0 => -2*ones(Float64,N),
              1 =>    ones(Float64,N-1))

u = P \ B

plot_solution(x, u)

# # LinearAlgebra

using LinearAlgebra

Δx = 0.05
x = Δx:Δx:1-Δx ## Solve only interior points: the endpoints are set to zero.
N = length(x)
B = sin.(2π*x) * Δx^2

DU =  ones(Float64, N-1)
D = -2ones(Float64, N)
DL =  ones(Float64, N-1)

T = Tridiagonal(DL, D, DU)

u = T \ B

plot_solution(x, u)
#----------------------------------------------------------------------------

# # LAPACK

using LinearAlgebra

Δx = 0.05
x = Δx:Δx:1-Δx ## Solve only interior points: the endpoints are set to zero.
N = length(x)
B = sin.(2π*x) * Δx^2

DU =   ones(Float64, N-1)
D  = -2ones(Float64, N)
DL =   ones(Float64, N-1)

LAPACK.gtsv!(DL, D, DU, B)

plot_solution(x, B)
#----------------------------------------------------------------------------

# # DIffEqOperators.jl
#
# [docs](http://docs.juliadiffeq.org/latest/features/diffeq_operator.html)

using DiffEqOperators
#----------------------------------------------------------------------------

Δx = 0.05
x = Δx:Δx:1-Δx ## Solve only interior points: the endpoints are set to zero.
N = length(x)
B = sin.(2π*x)

## Second order approximation to the second derivative
order = 2
deriv = 2

A = DerivativeOperator{Float64}(order, deriv, Δx, N, :Dirichlet0,:Dirichlet0)

u = A \ B

plot_solution(x, u)


"""
    expon_rng(λ, shape=1; seed=nothing)

Generate a `shape` element array of random variables from a Exponential(`λ`) distribution. Optionally you can set a specific seed.

# Notes

The pdf of an Exponential(λ) distribution is given as:

``f(x, λ) = λe^{-λx} \\quad x ≥ 0``

!!! note
    Some Exponential distributions are parameterized by ``β = \\frac{1}{λ}``, where λ is the number of events in an interval.
    In such cases, β represents the mean interarrival time; here we use λ to represent mean arrival rate per unit of time.

# Examples

```julia-repl
julia> expon_rng(3)
1-element Vector{Float64}:
 0.07033135663980515

julia> expon_rng(1.2, seed=42)
1-element Vector{Float64}:
 0.3296112244200808

julia> expon_rng(1.2, (2, 2))
2×2 Matrix{Float64}:
 1.9327    0.134739
 0.746861  0.155614
```

# References

D.P. Kroese, T. Taimre, Z.I. Botev. Handbook of Monte Carlo Methods.
  Wiley Series in Probability and Statistics, John Wiley & Sons, New York, 2011.
"""
function expon_rng(λ::Real, shape::Union{Int, Tuple{Vararg{Int}}}=1; seed::Union{Int, Nothing}=nothing)
    U = get_std_uniform(shape, seed=seed)
    X = zeros(shape)
    X = (-1/λ) .* log.(1 .- U)  # could also use just U
    return X
end


"""Julia program to check if a number is a Neon Number or not
A neon number is a number where the sum of digits of square of the number is equal to the number.
"""

function check_neon_num(n)
    sum = 0
    temp = n ^ 2
    while(temp != 0)
        sum = sum + (temp % 10)
        temp = temp  ÷ 10
    end
    if(sum == n)
        return true
    else
        return false
    end
end

print("Enter the number: ")
n = readline()
n = parse(Int, n)
res = check_neon_num(n)
if res
    println("The given number $n is a Neon Number.")
else
    println("The given number $n is not a Neon Number.")
end

"""
Time Complexity: O(log(num)), where 'num' is the given number
Space Complexity: O(1)

SAMPLE INPUT AND OUTPUT

SAMPLE 1
Enter the number: 0
The given number 0 is a Neon Number..

SAMPLE 2
Enter the number: 14
The given number 14 is not a Neon Number.

"""


using Distributions, LinearAlgebra, LaTeXStrings, Random, Plots; pyplot()
Random.seed!(1)

N = 10^5

SigY = [ 6 4 ;
         4 9]
muY = [15 ;
       20]
A = cholesky(SigY).L

rngGens = [()->rand(Normal()),
           ()->rand(Uniform(-sqrt(3),sqrt(3))),
           ()->rand(Exponential())-1]

rv(rg) = A*[rg(),rg()] + muY

data = [[rv(r) for _ in 1:N] for r in rngGens]

stats(data) = begin
    data1, data2 = first.(data),last.(data)
    println(round(mean(data1),digits=2), "\t",round(mean(data2),digits=2),"\t",
            round(var(data1),digits=2), "\t", round(var(data2),digits=2), "\t",
            round(cov(data1,data2),digits=2))
end

println("Mean1\tMean2\tVar1\tVar2\tCov")
for d in data
    stats(d)
end

scatter(first.(data[1]), last.(data[1]), c=:blue, ms=1, msw=0, label="Normal")
scatter!(first.(data[2]), last.(data[2]), c=:red, ms=1, msw=0, label="Uniform")
scatter!(first.(data[3]),last.(data[3]),c=:green, ms=1,msw=0,label="Exponential",
	xlims=(0,40), ylims=(0,40), legend=:bottomright, ratio=:equal,
    xlabel=L"X_1", ylabel=L"X_2")


export factorize_qr,
       factorize_lq,
       factorize_svd_right,
       factorize_svd_left

"""
    factorize_qr(M::Array{<:Number, 3})

QR factorization of M tensor.
"""
function factorize_qr(M::Array{<:Number, 3})
    m1, m2, m3 = size(M)
    M = reshape(M, m1*m2, m3)
    Q, R = qr(M)
    Q = Matrix(Q)
    Q = reshape(Q, m1, m2, size(Q, 2))
    return Q, R
end

"""
    factorize_lq(M::Array{<:Number, 3})

LQ factorization of M tensor.
"""
function factorize_lq(M::Array{<:Number, 3})
    m1, m2, m3 = size(M)
    M = reshape(M, m1, m2*m3)
    L, Q = lq(M)
    Q = Matrix(Q)
    Q = reshape(Q, size(Q, 1), m2, m3)
    return L, Q
end

"""
    do_svd_with_options(M::AbstractMatrix{<:Number},
                        cutoff::Float64,
                        dimcutoff::Int,
                        normalize_S::Bool)

Do SVD with cutoff truncation and normalization of singular values.
"""
function do_svd_with_options(M::AbstractMatrix{<:Number};
                             cutoff::Float64=1e-8,
                             dimcutoff::Int=1000,
                             normalize_S::Bool=true)
    F = svd(M)
    S = F.S
    ix_cutoff = min(findlast(S .> cutoff), dimcutoff)

    S = Diagonal(S[1:ix_cutoff])
    if normalize_S
        S ./= norm(S)
    end
    U = F.U[:, 1:ix_cutoff]
    Vt = F.Vt[1:ix_cutoff, :]
    return U, S, Vt
end

"""
    factorize_svd_right(M::AbstractMatrix{<:Number}; kwargs...)

SVD factorization of M matrix (U unitary).

 -- M -- = -- U -- SVt --
"""
function factorize_svd_right(M::AbstractMatrix{<:Number}; kwargs...)
    m1, m2 = size(M)
    U, S, Vt = do_svd_with_options(M; kwargs...)
    SVt = S*Vt
    return U, SVt
end

"""
    factorize_svd_left(M::AbstractMatrix{<:Number}; kwargs...)

SVD factorization of M matrix (Vt unitary).

 -- M -- = -- US -- Vt --
"""
function factorize_svd_left(M::AbstractMatrix{<:Number}; kwargs...)
    m1, m2 = size(M)
    U, S, Vt = do_svd_with_options(M; kwargs...)
    US = U*S
    return US, Vt
end

"""
    factorize_svd_right(M::Array{<:Number, 3}; kwargs...)

SVD factorization of M tensor (U unitary).

 -- M -- = -- U -- SVt --
    |         |
"""
function factorize_svd_right(M::Array{<:Number, 3}; kwargs...)
    m1, m2, m3 = size(M)
    U, SVt = factorize_svd_right(reshape(M, m1*m2, m3); kwargs...)
    U = reshape(U, m1, m2, size(U, 2))
    return U, SVt
end

"""
    factorize_svd_left(M::Array{<:Number, 3}; kwargs...)

SVD factorization of M tensor (Vt unitary).

 -- M -- = -- US -- Vt --
    |               |
"""
function factorize_svd_left(M::Array{<:Number, 3}; kwargs...)
    m1, m2, m3 = size(M)
    US, Vt = factorize_svd_left(reshape(M, m1, m2*m3); kwargs...)
    Vt = reshape(Vt, size(Vt, 1), m2, m3)
    return US, Vt
end

"""
    factorize_svd_right(M::Array{<:Number, 4}; kwargs...)

SVD factorization of M tensor (U unitary).

 -- M -- = -- U -- SVt --
   | |        |     |
"""
function factorize_svd_right(M::Array{<:Number, 4}; kwargs...)
    m1, m2, m3, m4 = size(M)
    U, SVt = factorize_svd_right(reshape(M, m1*m2, m3*m4); kwargs...)
    U = reshape(U, m1, m2, size(U, 2))
    SVt = reshape(SVt, size(SVt, 1), m3, m4)
    return U, SVt
end

"""
    factorize_svd_left(M::Array{<:Number, 4}; kwargs...)

SVD factorization of M tensor (Vt unitary).

 -- M -- = -- US -- Vt --
   | |        |     |
"""
function factorize_svd_left(M::Array{<:Number, 4}; kwargs...)
    m1, m2, m3, m4 = size(M)
    US, Vt = factorize_svd_left(reshape(M, m1*m2, m3*m4); kwargs...)
    US = reshape(US, m1, m2, size(US, 2))
    Vt = reshape(Vt, size(Vt, 1), m3, m4)
    return US, Vt
end


# # Convex Optimization in Julia
#
# ## Madeleine Udell | ISMP 2015

#-

# ## Convex.jl team
#
# * [Convex.jl](https://github.com/cvxgrp/Convex.jl): Madeleine Udell, Karanveer Mohan, David Zeng, Jenny Hong

#-

# ## Collaborators/Inspiration:
#
# * [CVX](http://www.cvxr.com): Michael Grant, Stephen Boyd
# * [CVXPY](https://github.com/cvxgrp/cvxpy): Steven Diamond, Eric Chu, Stephen Boyd
# * [JuliaOpt](https://github.com/JuliaOpt): Miles Lubin, Iain Dunning, Joey Huchette

## initial package installation

#-

## Make the Convex.jl module available
using Convex, SparseArrays, LinearAlgebra
using SCS # first order splitting conic solver [O'Donoghue et al., 2014]

## Generate random problem data
m = 50;  n = 100
A = randn(m, n)
x♮ = sprand(n, 1, .5) # true (sparse nonnegative) parameter vector
noise = .1*randn(m)    # gaussian noise
b = A*x♮ + noise      # noisy linear observations

## Create a (column vector) variable of size n.
x = Variable(n)

## nonnegative elastic net with regularization
λ = 1
μ = 1
problem = minimize(square(norm(A * x - b)) + λ*square(norm(x)) + μ*norm(x, 1),
                   x >= 0)

## Solve the problem by calling solve!
solve!(problem, SCS.Optimizer(verbose=0))

println("problem status is ", problem.status) # :Optimal, :Infeasible, :Unbounded etc.
println("optimal value is ", problem.optval)

#-

using Interact, Plots
## Interact.WebIO.install_jupyter_nbextension() # might be helpful if you see `WebIO` warnings in Jupyter
@manipulate throttle=.1 for λ=0:.1:5, μ=0:.1:5
    global A
    problem = minimize(square(norm(A * x - b)) + λ*square(norm(x)) + μ*norm(x, 1),
                   x >= 0)
    solve!(problem, SCS.Optimizer(verbose=0))
    histogram(evaluate(x), xlims=(0,3.5), label="x")
end

# # Quick convex prototyping

#-

# ## Variables

## Scalar variable
x = Variable()

#-

## (Column) vector variable
y = Variable(4)

#-

## Matrix variable
Z = Variable(4, 4)

# # Expressions

#-

# Convex.jl allows you to use a [wide variety of functions](http://convexjl.readthedocs.org/en/latest/operations.html) on variables and on expressions to form new expressions.

x + 2x

#-

e = y[1] + logdet(Z) + sqrt(x) + minimum(y)

# ### Examine the expression tree

e.children[2]

# # Constraints
#
# A constraint is convex if convex combinations of feasible points are also feasible. Equivalently, feasible sets are convex sets.
#
# In other words, convex constraints are of the form
#
# * `convexExpr <= 0`
# * `concaveExpr >= 0`
# * `affineExpr == 0`

x <= 0

#-

square(x) <= sum(y)

#-

M = Z
for i = 1:length(y)
    global M += rand(size(Z)...)*y[i]
end
M ⪰ 0

# # Problems

x = Variable()
y = Variable(4)
objective = 2*x + 1 - sqrt(sum(y))
constraint = x >= maximum(y)
p = minimize(objective, constraint)

#-

## solve the problem
solve!(p, SCS.Optimizer(verbose=0))
p.status

#-

x.value

#-

## can evaluate expressions directly
evaluate(objective)

# ## Pass to solver
#
# call a `MathProgBase` solver suited for your problem class
#
# * see the [list of Convex.jl operations](http://convexjl.readthedocs.org/en/latest/operations.html) to find which cones you're using
# * see the [list of solvers](http://www.juliaopt.org/) for an up-to-date list of solvers and which cones they support

#-

# to solve problem using a different solver, just import the solver package and pass the solver to the `solve!` method: eg
#
#     using Mosek
#     solve!(p, Mosek.Optimizer())

#-

# ## Warmstart

## Generate random problem data
m = 50;  n = 100
A = randn(m, n)
x♮ = sprand(n, 1, .5) # true (sparse nonnegative) parameter vector
noise = .1*randn(m)    # gaussian noise
b = A*x♮ + noise      # noisy linear observations

## Create a (column vector) variable of size n.
x = Variable(n)

## nonnegative elastic net with regularization
λ = 1
μ = 1
problem = minimize(square(norm(A * x - b)) + λ*square(norm(x)) + μ*norm(x, 1),
                   x >= 0)
@time solve!(problem, SCS.Optimizer(verbose=0))
λ = 1.5
@time solve!(problem, SCS.Optimizer(verbose=0), warmstart = true)

# # DCP examples

## affine
x = Variable(4)
y = Variable(2)
sum(x) + y[2]

#-

2*maximum(x) + 4*sum(y) - sqrt(y[1] + x[1]) - 7 * minimum(x[2:4])

#-

## not dcp compliant
log(x) + square(x)

#-

## $f$ is convex increasing and $g$ is convex
square(pos(x))

#-

## $f$ is convex decreasing and $g$ is concave
invpos(sqrt(x))

#-

## $f$ is concave increasing and $g$ is concave
sqrt(sqrt(x))


#=
Challenge #263 [Easy] Calculating Shannon Entropy of a String

Description

Shannon entropy was introduced by Claude E. Shannon in his 1948 paper "A
Mathematical Theory of Communication". Somewhat related to the physical and
chemical concept entropy, the Shannon entropy measures the uncertainty associated
with a random variable, i.e. the expected value of the information in the message
(in classical informatics it is measured in bits). This is a key concept in
information theory and has consequences for things like compression, cryptography
and privacy, and more.
The Shannon entropy H of input sequence X is calculated as -1 times the sum of
the frequency of the symbol i times the log base 2 of the frequency:

            n
            _   count(i)          count(i)
H(X) = -1 * >   --------- * log  (--------)
            -       N          2      N
            i=1

(That funny thing is the summation for i=1 to n. I didn't see a good way to do
this in Reddit's markup so I did some crude ASCII art.)
For more, see Wikipedia for Entropy in information theory).

Input Description

You'll be given a string, one per line, for which you should calculate the
Shannon entropy. Examples:

1223334444
Hello, world!

Output Description

Your program should emit the calculated entropy values for the strings to at least
five decimal places. Examples:

1.84644
3.18083

Challenge Input

122333444455555666666777777788888888
563881467447538846567288767728553786
https://www.reddit.com/r/dailyprogrammer
int main(int argc, char *argv[])

Challenge Output

2.794208683
2.794208683
4.056198332
3.866729296

(source: https://redd.it/4fc896)
=#

function __shannon_entropy( string::AbstractString )
    local frequency::Dict{Char, Integer} = Dict{Char, Integer}()

    for i in 1:length( string )
        if haskey( frequency, string[ i ] )
            frequency[ string[ i ] ] += 1
        else
            frequency[ string[ i ] ] = 1
        end
    end

    return frequency
end

function Shannon_entropy( string::AbstractString )
    local entropy::AbstractFloat = 0
    local frequency::Dict{Char, Integer} = __shannon_entropy( string )
    local N::Integer = length( string )

    for i in values( frequency )
        entropy += ( i/N ) * log2( i/N )
    end
    entropy *= -1

    return @sprintf "%.9f" entropy
end

println( Shannon_entropy( "1223334444" ) )
println( Shannon_entropy( "Hello, world!") )
println( Shannon_entropy( "122333444455555666666777777788888888") )
println( Shannon_entropy( "563881467447538846567288767728553786") )
println( Shannon_entropy( "https://www.reddit.com/r/dailyprogrammer") )
println( Shannon_entropy( "int main(int argc, char *argv[])") )


using LinearAlgebra
using SparseArrays
using Plots

"This code solves the 2D Laplace's equation with nonzero Dirichlet and Neumann
boundary conditions using a finite difference method."

m = 100 # number of points
x = LinRange(-1,1,m+2)
y = LinRange(-1,1,m+2)
xint = x[2:end-1]
yint = y[2:end-1]
h = x[2]-x[1]

# ordering of unknowns
# u14 u24 u34 u44
# u13 u23 u33 u43
# u12 u22 u32 u42
# u11 u21 u31 u41

id(i,j,m) = i + (j-1)*m # return global index from i,j

function FD_matrix_2D(h,m)
    A = spzeros(m*m,m*m)
    for i = 1:m, j = 1:m # loop thru x,y indices
        A[id(i,j,m),id(i,j,m)] = 4/h^2
        if i > 1 # avoids leftmost line of nodes
            A[id(i,j,m),id(i-1,j,m)] = -1/h^2 # x-derivative
        end
        if i < m # avoids rightmost line of nodes
            A[id(i,j,m),id(i+1,j,m)] = -1/h^2
        end
        if j > 1
            A[id(i,j,m),id(i,j-1,m)] = -1/h^2 # y-derivative
        end
        if j < m
            A[id(i,j,m),id(i,j+1,m)] = -1/h^2
        end
    end
    return A
end

A = FD_matrix_2D(h,m)

f(x,y) = exp(-10*((x-.25)^2 + (y-.75)^2))
f(x,y) = 0.0
b = vec(f.(xint,yint'))

# impose non-zero Dirichlet BCs
u_left(y) = 0.0
u_right(y) = 1.0 # 0 if y in [-1,0], 1 if y in [0,1]
u_bottom(x) = 0.0
u_top(x) = -1.0

for i = 1:m, j = 1:m
    if i==1
        b[id(i,j,m)] += u_left(yint[j])/h^2
    end
    if i==m
        b[id(i,j,m)] += u_right(yint[j])/h^2
    end
    if j==1
        b[id(i,j,m)] += u_bottom(xint[i])/h^2
    end
    if j==m
        b[id(i,j,m)] += u_top(xint[i])/h^2
    end
end

# solve Au = f(xi,yj)
u = A\b
contourf(xint,yint,u,clims=(-1,1)) #vec(reshape(u,m,m)'))


using LinearAlgebra
using SparseArrays
using Plots

"This code solves the 2D Laplace's equation with nonzero Dirichlet and Neumann
boundary conditions using a finite difference method."

m = 100 # number of points
x = LinRange(-1,1,m+2)
y = LinRange(-1,1,m+2)
xint = x[1:end]
yint = y[2:end-1]
h = x[2]-x[1]

id(i,j,m) = i + (j-1)*(m+2) # return global index from i,j

function FD_matrix_2D(h,m)
    A = spzeros((m+2)*m,(m+2)*m)
    for i = 1:m+2, j = 1:m # loop thru x,y indices
        A[id(i,j,m),id(i,j,m)] = 4/h^2

        # stencil at interior nodes
        if i > 1 # avoids leftmost line of nodes
            A[id(i,j,m),id(i-1,j,m)] = -1/h^2 # x-derivative
        end
        if i < m+2 # avoids rightmost line of nodes
            A[id(i,j,m),id(i+1,j,m)] = -1/h^2
        end
        if j > 1
            A[id(i,j,m),id(i,j-1,m)] = -1/h^2 # y-derivative
        end
        if j < m
            A[id(i,j,m),id(i,j+1,m)] = -1/h^2
        end

        # modify for Neumann BCs on left/right
        # 2*(u_{m+1,j} - u_{m,j})/h^2 at Neumann boundaries in the x coordinate
        if i==1 # on left
            A[id(i,j,m),id(i+1,j,m)] = -2/h^2
        end
        if i==m+2 # on right
            A[id(i,j,m),id(i-1,j,m)] = -2/h^2
        end
    end
    return A
end

A = FD_matrix_2D(h,m)

# f(x,y) = exp(-10*((x-.25)^2 + (y-.75)^2))
f(x,y) = 0.0
b = vec(f.(xint,yint'))

# impose non-zero Dirichlet BCs
u_top(x) = 1.0
u_bottom(x) = -1.0
dudx_left(y) = y^2
dudx_right(y) = 20*y

for i = 1:m+2, j = 1:m
    # modify Neumann BCs first, then add Dirichlet contributions next
    if i==1 # left
        b[id(i,j,m)] += 2*dudx_left(yint[j])/h
    end
    if i==m+2  # right
        b[id(i,j,m)] += 2*dudx_right(yint[j])/h
    end

    if j==1 # bottom
        uij = u_bottom(xint[i])
        b[id(i,j,m)] += uij/h^2
    end
    if j==m # top
        uij = u_top(xint[i])
        b[id(i,j,m)] += uij/h^2
    end
end

# solve Au = f(xi,yj)
u = A\vec(b)
contourf(xint,yint,u) #vec(reshape(u,m,m)'))


import RowEchelon

#=
Homework 2
Section 1.4
Problem 37
=#

# Create the Matrix to be evaluated
A = [ 7 2 -5 8 ;
	 -5 -3 4 -9 ;
	 6 10 -2 7 ;
	 -7 9 2 15 ]

# Store the columns for future reference
columns = 4

# Row reduce the Matrix
A_rref = RowEchelon.rref(A)

println("Original Matrix, A")
println(A)
println("Row Reduced Echelon Matrix:")
println(A_rref)

# Determine whether or not matrix has a pivot position in each row
hasAllZeroRow = false
index = 1
currentZeros = 0
while index <= columns
	global index
	global hasAllZeroRow
	global currentZeros

	# Parse through each entry of the rows, searching for an all 0 row.
	for entry in A_rref[index, :]
		global hasAllZeroRow
		global currentZeros

		if entry == 0
			# If here then:
			# Current entry is a zero, increment zero counter
			currentZeros += 1
		end
	end
	if currentZeros == columns
		# If here then:
		# A row with all zeroes was found
		hasAllZeroRow = true
	end
	currentZeros = 0
	index += 1
end

if hasAllZeroRow == true
	# If here then:
	# The matrix does not span R4
	println("The Matrix does not span R4")
else
	# If here then:
	# The matrix does span all elements in R4
	println("The Matrix does span R4")
end

println()
println()
exit(0)


#
# Jerzy Wroczyński (nr. indeksu 250075)
# Obliczenia Naukowe, Laboratorium
# Lista 3, Zadania 1.–3.
#

module IterativeMethods

FL = Float64
MAX_ITERATIONS = 10^9
NEAR_ZERO = 2^(-2.0^32)

"""
### Finds the solution to the equation `f(x) = 0` using the bisection method.

### Arguments
- `f` — the function to find a root of
- `a` — beginning of the interval
- `b` — ending of the interval
- `delta` — the precision of the solution in terms of the X axis
- `epsilon` — the precision of the solution in terms of the Y axis

### Return value: a tuple of four values
- `[0]` — the approximation of the root of the `f` function
- `[1]` — the actual value of the `f` function in `[0]`
- `[2]` — number of iterations performed
- `[3]` — error indicator, possible values:
    - `0` — the method was convergent in this case
    - `1` — the `f` function doesn’t change its sign inside the given interval (doesn’t cut through the X axis)

"""
function mbisekcji(f, a::FL, b::FL, delta::FL, epsilon::FL)

    # begin with the starting values
    u = f(a)
    v = f(b)
    # the length of the section
    e = b - a

    if sign(u) == sign(v)
        # `f` doesn’t change its sign between `a` and `b`
        return (a, u, 0, 1);
    end

    # predefine some variables so they are accessible outside the loop
    k = 0
    c = 0
    w = 0
    # iterate
    for k in 0:MAX_ITERATIONS

        # divide into two sections — divide the length of the section by two
        e = e/2
        # pick the middle point
        c = a + e
        w = f(c)

        # if the solution is sufficient with respect to the given precision parameters terminate this method and return the output values
        if abs(e) < delta || abs(w) < epsilon
            return (c, w, k, 0)
        end

        # choose that section that contains the solution
        if sign(w) != sign(u)
            b = c
            v = w
        else
            a = c
            u = w
        end

    end

    # somehow the function exceeded MAX_ITERATIONS, however it seems so unlikely that a separate error code for this edge case would be an excess
    return (c, w, k, 0)

end

"""
### Finds the solution to the equation `f(x) = 0` using the Newton’s method.

### Arguments
- `f` — the function to find a root of
- `pf` — the derivative of the `f` function
- `x0` — first approximation
- `delta` — the precision of the solution in terms of the X axis
- `epsilon` — the precision of the solution in terms of the Y axis
- `maxint` — maximum iterations this function can perform

### Return value: a tuple of four values
- `[0]` — the approximation of the root of the `f` function
- `[1]` — the actual value of the `f` function in `[0]`
- `[2]` — number of iterations performed
- `[3]` — error indicator, possible values:
    - `0` — the method was convergent in this case
    - `1` — given solution is not within the given precision parameters because of not enough iterations
    - `2` — the derivative function was near zero

"""
function mstycznych(f, pf, x0::FL, delta::FL, epsilon::FL, maxint::Int)

    v = f(x0)

    if abs(v) < epsilon
        # given x₀ is already within given precision
        return (x0, v, 0, 0)
    end

    # iterate until the MAX_ITERATIONS is reached or the precision is satisfied
    xn = x0
    xn1 = x0

    # predefine the iterator so it would be accessible outside the loop
    k = 0
    # iterate
    for k in 1:maxint

        # perform the recursive step
        xn1 = xn - (v/pf(xn))
        v = f(xn1)

        # check the derivative
        if abs(pf(xn)) < NEAR_ZERO || isinf(abs(pf(xn)))
            # the derivative function is near zero which could lead to a not very precise solution
            return (xn, f(xn), k, 2)
        end

        # if the solution is within given precision return the output values
        if abs(xn1 - xn) < delta || abs(v) < epsilon
            return (xn1, v, k, 0)
        end

        # otherwise continue iterating
        xn = xn1

    end

    # not enough iterations to satisfy given precision
    return (xn1, v, k, 1)

end

"""
### Finds the solution to the equation `f(x) = 0` using the secant method.

### Arguments
- `f` — the function to find a root of
- `x0` — the first approximation
- `x1` — the second approximation
- `delta` — the precision of the solution in terms of the X axis
- `epsilon` — the precision of the solution in terms of the Y axis
- `maxint` — maximum iterations this function can perform

### Return value: a tuple of four values
- `[0]` — the approximation of the root of the `f` function
- `[1]` — the actual value of the `f` function in `[0]`
- `[2]` — number of iterations performed
- `[3]` — error indicator, possible values:
    - `0` — the method was convergent in this case
    - `1` — given solution is not within the given precision parameters because of not enough iterations

"""
function msiecznych(f, x0::FL, x1::FL, delta::FL, epsilon::FL, maxint::Int)

    # calc the first values
    v = f(x0)
    w = f(x1)
    a = x0
    b = x1

    # predefine the iterator so it would be accessible outside the loop
    k = 0
    #iterate
    for k in 0:maxint

        if abs(v) > abs(w)
            # swap the values, if necessary to maintain the order such that the latter value (`f(b)`) is equal or bigger than former (`f(a)`)
            # because of that the absolute value of the function will not rise — and that’s what we want since we’re calculating the root of the `f` function
            a,b = b,a
            v,w = w,v
        end

        # calculate the pseudo-derivative
        d = (b-a)/(w-v)
        # shift the values
        b = a
        w = v
        # calc the next value
        a = a - d * v
        v = f(a)

        if abs(v) < epsilon || abs(b-a) < delta
            # given solution is within given precision
            return (a, v, k, 0)
        end

    end

    # not enough iterations to give an accurate solution (not within given precision)
    return (a, v, k, 1)

end

end


export cheb_diff, cheb_mesh

"""
    cheb_mesh(T::Type, N::Integer)
    cheb_mesh(N::Integer)

Return the Chebyshev nodes of degree `N` on the interval [0, 1].
"""
cheb_mesh(T::Type, N::Integer) = (-cos.((T(π) * (0:N)) ./ N) .+ 1) ./ 2
cheb_mesh(N::Integer) = cheb_mesh(Float64, N)

"""
    cheb_diff(T::Type, N::Integer)
    cheb_diff(N::Integer)

Return a tuple `(x, D)` where `x` contains the Chebyshev nodes of degree `N` on the interval
[0, 1] using the numerical type `T` and `D` is the Chebyshev differentiation matrix on the
same interval.

## Example - differentiate a sinusoid

```julia
julia> (x, D) = cheb_diff(10) ;

julia> D * sinpi.(x) ≈ π * cospi.(x)
true
```
"""
function cheb_diff(T::Type, N::Integer)
    if N == 0
        return (ones(T, 1), zeros(T, 1, 1))
    else
        x = cheb_mesh(T, N)
        c = [2; ones(T, N - 1, 1); 2] .* (-1) .^ (0:N)
        dx = x .- x'
        D = (c * (1 ./ c)') ./ (dx + I)
        D[diagind(D)] .-= vec(sum(D; dims=2))
        return (x, D)
    end
end
cheb_diff(N::Integer) = cheb_diff(Float64, N)


export Sigmoid

"""
`Sigmoid([domainType=Float64::Type,] dim_in::Tuple, γ = 1.)`

Creates the sigmoid non-linear operator with input dimensions `dim_in`.
```math
\\sigma(\\mathbf{x}) = \\frac{1}{1+e^{-\\gamma \\mathbf{x} } }
```

"""
struct Sigmoid{T,N,G<:Real} <: NonLinearOperator
	dim::NTuple{N,Int}
	gamma::G
end

function Sigmoid(DomainType::Type, DomainDim::NTuple{N,Int}, gamma::G=1.) where {N, G <: Real}
	Sigmoid{DomainType,N,G}(DomainDim,gamma)
end

Sigmoid(DomainDim::NTuple{N,Int}, gamma::G=1.) where {N,G} = Sigmoid{Float64,N,G}(DomainDim,gamma)

function mul!(y::AbstractArray{T,N}, L::Sigmoid{T,N,G}, x::AbstractArray{T,N}) where {T,N,G}
	y .= (1 .+exp.(-L.gamma.*x)).^(-1)
end

function mul!(y::AbstractArray,
              J::AdjointOperator{Jacobian{A,TT}},
              b::AbstractArray) where {T,N,G, A<: Sigmoid{T,N,G}, TT <: AbstractArray{T,N}}
    L = J.A
	y .= exp.(-L.A.gamma.*L.x)
	y ./= (1 .+y).^2
	y .= conj.(L.A.gamma.*y)
	y .*= b
end

fun_name(L::Sigmoid) = "σ"

size(L::Sigmoid) = (L.dim, L.dim)

domainType(L::Sigmoid{T,N,D}) where {T,N,D} = T
codomainType(L::Sigmoid{T,N,D}) where {T,N,D} = T


module NewtonsMethod

using LinearAlgebra, Statistics, Compat, ForwardDiff

function newtonroot(f, f′; x₀, tolerance = 1E-7, maxiter = 1000)
    # setup the algorithm
    x_old = x₀
    normdiff = Inf
    iter = 1
    while normdiff > tolerance && iter <= maxiter
        x_new = x_old - f(x_old)/f′(x_old) # use the passed in map
        normdiff = norm(x_new - x_old)
        x_old = x_new
        iter = iter + 1
    end

    if iter < maxiter+1
        return (root = x_old, normdiff = normdiff, iter = iter) # A named tuple
    else
        return nothing
    end
end

# for auto-differentiation, can just forward to the existing one
D(f) = x -> ForwardDiff.derivative(f, x)

# same name! different parameters
newtonroot(f; x₀, tolerance = 1E-7, maxiter = 1000) = newtonroot(f, D(f), x₀=x₀, tolerance = tolerance, maxiter = maxiter)

export newtonroot
end # module


using Random, Statistics, Plots, Printf
Random.seed!(430)
N = 10000;

# generate 10000 realizations of a Gaussian random variable Xn with mean 1 and variance 4
X_n = 2 * randn(N) .+ 1;
@printf("%d Gaussian random variable Xn with mean: %.5f and variance: %.5f\n", N, mean(X_n), var(X_n))

# Calculate S_n
S_n = 1/length(X_n[1:N]) * sum(X_n[1:N])
@printf("Obtained Sn: %.5f from Xn samples\n", S_n)

# Define the running average
function running_avg(n)
    values = []
    for i in 1:n
        S_n = 1/length(X_n[1:i]) * sum(X_n[1:i])
        push!(values, S_n)
    end
    return values
end

# Samples for each n values from 1 to 10,000
X_nSamples = running_avg(N);

# Plot |Sn − 1| as a function of n
S_nMinus1 = abs.(X_nSamples .- 1);
display(plot(1:N, S_nMinus1, labels="|Sn − 1|", xlabel="N", ylabel="|Sn − 1|"))

# Define the time correlation
function time_correlation(X_n, n)
    Kn_list = []
    for n in 1:N-1
        Kn=0
        for k in 1:n
            Kn += X_n[k]*X_n[k+1]
        end
        val = 1/n * Kn;
        push!(Kn_list, val)
    end
    return Kn_list
end

# Plot Kn
T = time_correlation(X_n, N);
display(plot(1:length(T), T, labels="Kn", xlabel="N", ylabel="Kn"))