using DimensionalityReduction
using Gadfly

# TODO: approximate (generalized) principal components rather than principal components could be more efficient for N > T
type FactorModel
    # models X_t = lambda*factors_t + e_t  # Here lambda is not equal to lambda(L) --> static factors
    x::Array{Float64, 2}  # variables to calculate factors from
    number_of_factors::Int64  # columns of factors we use (which e.g. capture a certain percentage of the variation or from Bai Ng info criteria)
    number_of_factors_criterion::String  # Bai Ng info criterion to use
    factor_type::String  # "principal components" later maybe also "maximum likelihood"
    factors::Array{Float64, 2}  # matrix of static factors
    loadings::Array{Float64, 2}
    factor_residuals::Array{Float64, 2}  # residuals from the factor estimation  x = factors * lambda

    function FactorModel(x::Matrix{Float64}, number_of_factors::Int64=int(ceil(minimum(size(x))/2)), number_of_factors_criterion::String="", factor_type::String="principal components")
        # (static) factor equation:
        factors, loadings, number_of_factors = calculate_factors(x, factor_type, number_of_factors)
        # TODO: check if correct factors are used (p 1136 on Bai Ng 2006)
        factor_residuals = x .- factors[:, 1:number_of_factors] * loadings[:, 1:number_of_factors]'
        #factor_residuals = x - x*loadings[:, 1:number_of_factors]*inv(loadings[:, 1:number_of_factors]'loadings[:, 1:number_of_factors])*loadings[:, 1:number_of_factors]' # rescaled version from french dude
        # TODO: for lags we have to regress residuals on its lags and use info criterio to choose lag length p. this should be offloaded into another method
        # TODO: solve chicken and egg problem between number_of_factors and number_of_factor_lags. Both require residuals. --> maybe select one after the other repeatedly until convergence?
        # TODO: do the above e.g. factors, loadings, number_of_factors = augment_factors_by_lags(factors, number_of_factor_lags)  # TODO: do we have to

        return(new(x, number_of_factors, number_of_factors_criterion, factor_type, factors, loadings, factor_residuals))
    end

    function FactorModel(x::Matrix{Float64}, number_of_factors_criterion::String, factor_type::String="principal components"; max_factors = int(ceil(minimum(size(x))/2)))
        #println("finding optimal number of factors with maximum number of factors=", max_factors, " and information criterion=", number_of_factors_criterion)
        # number of factors to include is not given but a criterion --> we use
        # the criterion to determine the number of factors using the Bai, Ng
        # information criteria.
        models = [apply(FactorModel, (x, number_of_factors, number_of_factors_criterion, factor_type)) for number_of_factors in 1:max_factors]
        # we keep all the models in memory which can be a problem depending on the dimensions of x. TODO: will refactor later when debugging and testing is done
        criteria = [calculate_criterion(model) for model in models]
        return models[indmin(criteria[1:max_factors])]  # keep the model with the best information criterion
    end

end

function principal_components(x)  # calculate factors and loadings for a given x matrix
    T, N = size(x)  # T: time dimension, N cross-sectional dimension
    # see Stock, Watson (1998)
    if T >= N
        eigen_values, eigen_vectors = eig(x'x)  # x'x is NxN
        eigen_values, eigen_vectors = reverse(eigen_values), eigen_vectors[:, size(eigen_vectors, 2):-1:1]  # reverse the order from highest to lowest eigenvalue
        loadings = sqrt(N) * eigen_vectors  # we may rescale the eigenvectors say Bai, Ng 2002
        factors = x*loadings/N  # this is from Bai, Ng 2002 except that for me the inverse of the loadings matrix primed is not the same as the loadings matrix
        #factors = factors*(factors'factors/T)^(1/2)  # TODO: not sure it is correct to rescale like this (see Bai, Ng 2002 p.198)
        #factors = 1/N*x*loadings'  # this is from Stock, Watson 2010
        #loadings = (x'x)*loadings/(N*T)  # and this is from C. Hurlin from University of Orléans... rescales factor loadings to have the same variance as x
    end
    if N > T
        eigen_values, eigen_vectors = eig(x*x')  # x*x' is TxT
        eigen_values, eigen_vectors = reverse(eigen_values), eigen_vectors[:, size(eigen_vectors, 2):-1:1]  # reverse the order from highest to lowest eigenvalue
        factors = sqrt(T) * eigen_vectors  # sqrt(T) comes from normalization F'F/T = eye(r) where r is number of factors (see Bai 2003), factors are Txr
        loadings = x'factors/T  # see e.g. Bai 2003 p.6 or Breitung, Eickmeier 2011 p. 80. Dimension of loadings is Nxr where r here is also N
        # betahat = loadings * chol(loadings'loadings/N)  # C. Hurlin from University of Orléans rescales like this
    end
        # resultingly factors*loadings' estimates x
    return factors, loadings, eigen_values, eigen_vectors
end

# calculate the factors and the rotation matrix to transform data into the space spanned by the factors
function calculate_factors(x::Matrix, factor_type::String="principal components", number_of_factors=ceil(minimum(size(x))/2))
    if factor_type == "principal components"
        factors, loadings = principal_components(x)
        max_factor_number = int(ceil(minimum(size(x))/2))
    elseif factor_type == "squared principal components"  # include squares of X
        factors, loadings = principal_components([x x.^2])
        max_factor_number = minimum([size(x, 1), size(x,2)*2])
    elseif factor_type == "quadratic principal components"  # include squares of X and interaction terms - better only use in combination with targeted_predictors
        pca_cols = x  # columns to use for principal components
        for i in 1:size(x, 2)
            for j in 1:size(x, 2)
                pca_cols = hcat(pca_cols, x[:, i].*x[:, j])  # not very efficient!
            end
        end
        max_factor_number = int(ceil(minimum(size(pca_cols))/2))
        factors, loadings = principal_components(pca_cols)
    end
    if number_of_factors > max_factor_number
        number_of_factors = max_factor_number
        warn("can not estimate more than `minimum(size(x))` factors with $factor_type. Number of factors set to $number_of_factors")
    end
    return factors, loadings, number_of_factors
end

function scree_plot(x::Array{Float64, 2}, max_factors::Int64=size(x, 2); file_name::String="")
    eigen_values = principal_components(x)[3][1:max_factors]
    plt = plot(
        layer(x=[1:length(eigen_values)], y=eigen_values, Geom.point, Theme(default_color=color("green"))),
        layer(x=[1:length(eigen_values)], y=eigen_values, Geom.line),
        Guide.XLabel("Number of Factors"), Guide.YLabel("Eigenvalue"),
        Scale.y_continuous(format=:plain),
        Scale.x_discrete
    )
    if length(file_name) > 0
        draw(PNG(file_name, 20cm, 15cm), plt)
    end
    plt
end

function Base.show(io::IO, fm::FactorModel)
    @printf io "Static Factor Model\n"
    @printf io "Dimensions of X: %s\n" size(fm.x)
    @printf io "Number of factors used: %s\n" fm.number_of_factors
    @printf io "Factors calculated by: %s\n" fm.factor_type
    @printf io "Factor model residual variance: %s\n" sum(fm.factor_residuals.^2)/apply(*, size(fm.x))
end

function calculate_criterion(fm::FactorModel)
    if fm.number_of_factors_criterion != ""
        number_of_factors_criterion = fm.number_of_factors_criterion
        number_of_factors_criterion_value::Float64 = eval(symbol("criterion_$number_of_factors_criterion"))(fm)
    end
    return(number_of_factors_criterion_value)
end

function predict(fm::FactorModel, y::Array{Float64, 1}, h::Int64=1, number_of_lags::Int64=5, number_of_factors::Int64=0)
    if number_of_factors == 0  # number of factors can be given to set them to a different value in the forecasting equation than in the factor equation
        #println("using the same number of factors for forecasting as in the factor equation: ", fm.number_of_factors)
        number_of_factors = fm.number_of_factors
    end
    # makes a h step ahead forecast of y (which can also be in the factor model)
    w = lagged_matrix(y, [0, [h:number_of_lags+h-1]])
    # estimate y_{t+h} = alpha'w_t + Gamma'x_t

    y, x = w[1:end-1, 1], hcat(w[:, 2:end], fm.factors[end-size(w,1)+1:end, 1:number_of_factors])  # last observation of y is reserved for prediction
    x = hcat(ones(size(x, 1)), x)  # add a constant term to the regression
    new_x, design_matrix = x[end,:], x[1:end-1, :]  # we reserve the last observation for prediction and dont use it for learning
    coefficients = inv(design_matrix'design_matrix)*design_matrix'y
    prediction = new_x*coefficients
    residuals = y - design_matrix*coefficients
    #hat_matrix = design_matrix*inv(design_matrix'design_matrix)*design_matrix'
    #residual_variance = (residuals.^2)./(1.-diag(hat_matrix))  # HC 2
    #coefficient_covariance = inv(design_matrix'design_matrix)*(design_matrix'diagm(residual_variance)design_matrix)*inv(design_matrix'design_matrix)

    return(residuals, prediction[1])
end


using Turing

#Returns a categorical distribution.
@sf fdirichlet(alpha) = begin
    ps = dirichlet(alpha)
    () -> categorical(ps)
end

#Makes 10 draws from the categorical distribution returned by fdirichlet.
@sf fdp() = begin
    dist = fdirichlet(ones(10)/5)
    Int[dist() for i = 1:30]
end

#Because we haven't conditioned on anything, we only need to draw one sample.
println(sample(fdp, 1)[end])


# test function kurt
x = [-1,0,2,-1]
@test kurt(x) == 18/4 - 3*(6/4)^2

x = [0]
@test kurt(x) == 0

x = [2,1,3,-2,-1,-3]
@test kurt(x) == mean(x.^4) - 3*mean(x.^2)^2

# test function grad_neg_sum_kurt()
@testset "Test grad_neg_sum_kurt" begin
    for idx = 1: 10
        Z = [randn(100000) for i = 1: 3]
        W = randn(3,3)
        @test rel_error(ForwardDiff.gradient(x -> neg_abs_sum_kurt(x'*Z), W), grad_neg_abs_sum_kurt(W, Z)) < 1e-8
    end
end

@testset "Test kurtosis-based icf $opt" for opt in ["orth", "sphe"]
    for idx = 1: 10
        # set up
        T = 100000
        x = [rand(T) + 1*ones(T), rand(T) - 1*ones(T)]
        A = randn(2,2)
        z = A*x;

        # apply icf
        Aest = icf(z; opt = opt)[1]
        @test err_pd(pinv(Aest)*A) < 5e-2
    end
end

# test function grad_ent_sum()
@testset "Test grad_ent_sum" begin
    for idx = 1: 10
        Z = [randn(100000) for i = 1: 3]
        W = randn(3,3)
        @test rel_error(ForwardDiff.gradient(x -> ent_sum(x'*Z), W), grad_ent_sum(W, Z)) < 1e-8
    end
end

@testset "Test entropy-based icf" begin
    for idx = 1: 10
        # set up
        T = 100000
        x = [rand(T) + 1*ones(T), rand(T) - 1*ones(T)]
        A = randn(2,2)
        z = A*x;

        # apply icf
        Aest = icf(z; obj = "ent")[1]
        @test err_pd(pinv(Aest)*A) < 5e-2
    end
end


using FastTransforms
import Base: *

#=
N = 2.^(7:10)

for n in N
    A = sphones(Float64, n, n)
    B = deepcopy(A)

    @time PC = c_plan_sph2fourier(n)
    @time c_sph2fourier(PC, B)

    @time PJ = SlowSphericalHarmonicPlan(A)
    @time A = PJ*A;

    @show vecnorm(A-B)/vecnorm(A)

    println()
    println()
end
=#

#=
α, β, γ = 0.0, -0.5, -0.5
#α, β, γ = -0.5, -0.5, -0.5
#α, β, γ = 0.0, 0.0, 0.0

for n in N
    A = trirand(Float64, n, n)
    B = deepcopy(A)

    @time PC = c_plan_tri2cheb(n, α, β, γ)
    @time c_tri2cheb(PC, B)

    @time PJ = SlowTriangularHarmonicPlan(A, α, β, γ);
    @time A = PJ*A;

    @show vecnorm(A-B)/vecnorm(B)

    println()
    println()
end

CF = zero(F); CF[1,2] = 1.0; CF[2,2] = 1.0; CF[3,2] = 1.0;
c_tri2cheb(cp, CF)
CF
c_cheb2tri(cp, CF)
CF

cprot = c_plan_rottriangle(size(F, 1), α, β, γ)

CF = zero(F); CF[1,2] = 1.0; CF[2,2] = 1.0; CF[3,2] = 1.0;
c_execute_tri_hi2lo(cprot, CF)
CF
c_execute_tri_lo2hi(cprot, CF)
CF
=#

using ApproxFun, FastTransforms

import ApproxFun: jacobip

jacobinorm(n,a,b) = if n ≠ 0
    sqrt((2n+a+b+1))*exp((lgamma(n+a+b+1)+lgamma(n+1)-log(2)*(a+b+1)-lgamma(n+a+1)-lgamma(n+b+1))/2)
else
    sqrt(exp(lgamma(a+b+2)-log(2)*(a+b+1)-lgamma(a+1)-lgamma(b+1)))
end

njacobip(n,a,b,x) = jacobinorm(n,a,b) * jacobip(n,a,b,x)

α = 0.0#0.25
β = 0.0#-0.25
γ = 0.0#-0.125

P = (ℓ,m,x,y) -> (2*(1-x))^m*njacobip(ℓ-m,2m+β+γ+1,α,2x-1)*njacobip(m,γ,β,2y/(1-x)-1)

p_T = chebyshevpoints(10)

f = (x,y) -> P(0,0,x,y) + P(1,1,x,y) + P(2,2,x,y) + P(3,3,x,y) + P(4,4,x,y) +
             P(1,0,x,y) + P(2,1,x,y) + P(3,2,x,y) +
             P(2,0,x,y) + P(3,1,x,y) +
             P(3,0,x,y)

f̃ = (s,t) -> f((s+1)/2, (t+1)/2*(1-(s+1)/2))

F = f̃.(p_T, p_T')

for j = 1:size(F,2)
    F[:,j] = chebyshevtransform(F[:,j])
end

for k = 1:size(F,1)
    F[k,:] = chebyshevtransform(F[k,:])
end

P = Tri2Cheb(

cp = c_plan_tri2cheb(size(F, 1), α, β, γ)

CF̌ = copy(F)

c_cheb2tri(cp, CF̌)

CF̌


# Generated functions for tensor-product polynomials where nodes are in an array of arrays

@generated function chebyshev_weights(f::AbstractArray{T,N},nodes::Array{Array{T,1},1},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T,N,S}

  chebyshev_polynomials = :( poly = Array{Array{T,2},1}(undef,N);
                             for k = 1:N;
                               orderk = order[k];
                               xk = nodes[k];

                               # normalize nodes

                               if domain[1,k] == domain[2,k];
                                 fill!(xk,(domain[1,k].+domain[2,k])/2);
                               else;
                                 xk = 2*(xk.-domain[2,k])/(domain[1,k].-domain[2,k]).-one(T);
                               end;

                               polynomial = Array{T}(undef,length(xk),orderk+1);
                               for i = 1:orderk+1;
                                 for j = 1:length(xk);
                                   if i == 1;
                                     polynomial[j,i] = one(T);
                                   elseif i == 2;
                                     polynomial[j,i] = xk[j];
                                   else;
                                     polynomial[j,i] = 2*xk[j]*polynomial[j,i-1]-polynomial[j,i-2];
                                   end;
                                 end;
                               end;
                               poly[i] = polynomial;#push!(poly,polynomial);
                             end )

  i_vars = Array{Symbol}(undef,N)
  s_vars = Array{Symbol}(undef,N)
  for i = 1:N
    i_vars[i] = Symbol("i$i")
    s_vars[i] = Symbol("s$i")
  end

  # Construct denominator term

  inner_prod_denominator = :( poly[1][s1,i1] )
  for i = 2:N
    inner_prod_denominator = :( poly[$i][$(s_vars[i]),$(i_vars[i])]*$inner_prod_denominator )
  end
  denominator_term = :( $inner_prod_denominator*$inner_prod_denominator )

  # Construct numerator term

  inner_prod_numerator = :( poly[1][s1,i1]*f[$(s_vars...)] )
  for i = 2:N
      inner_prod_numerator = :( poly[$i][$(s_vars[i]),$(i_vars[i])]*$inner_prod_numerator )
  end
  numerator_term = inner_prod_numerator

  # Construct weights term

  weights_term = :( weights[$(i_vars...)] = numerator/denominator )

  # Construct the inner loops that compute numerator and denominator

  inner = :( numerator += $numerator_term;
             denominator += $denominator_term )
  outer  = inner

  for i = N:-1:1#1:N
    outer = :(
      for $(s_vars[i]) = 1:size(f,$i)
        $outer
      end
    )
  end

  # Construct the outer loops that compute the weights

  new_inner = :( numerator = zero(T);
                 denominator = zero(T);
                 $outer;
                 $weights_term )
  new_outer = new_inner

  for i = N:-1:1#1:N
    new_outer = :(
      for $(i_vars[i]) = 1:(order[$i]+1)
        $new_outer
      end
    )
  end

  # Put it all together to compute the weights array

  final = :( $chebyshev_polynomials;
             weights = Array{T}(undef,(order+1)...);
             $new_outer;
             return weights )

  return final

end

@generated function chebyshev_weights(f::AbstractArray{T,N},poly::Array{Array{T,2},1},order::Array{S,1}) where {T,N,S}

  i_vars = Array{Symbol}(undef,N)
  s_vars = Array{Symbol}(undef,N)
  for i = 1:N
    i_vars[i] = Symbol("i$i")
    s_vars[i] = Symbol("s$i")
  end

  # Construct denominator term

  inner_prod_denominator = :( poly[1][s1,i1] )
  for i = 2:N
    inner_prod_denominator = :( poly[$i][$(s_vars[i]),$(i_vars[i])]*$inner_prod_denominator )
  end
  denominator_term = :( $inner_prod_denominator*$inner_prod_denominator )

  # Construct numerator term

  inner_prod_numerator = :( poly[1][s1,i1]*f[$(s_vars...)] )
  for i = 2:N
      inner_prod_numerator = :( poly[$i][$(s_vars[i]),$(i_vars[i])]*$inner_prod_numerator )
  end
  numerator_term = inner_prod_numerator

  # Construct weights term

  weights_term = :( weights[$(i_vars...)] = numerator/denominator )

  # Construct the inner loops that compute numerator and denominator

  inner = :( numerator += $numerator_term;
             denominator += $denominator_term )
  outer  = inner

  for i = N:-1:1#1:N
    outer = :(
      for $(s_vars[i]) = 1:size(f,$i)
        $outer
      end
    )
  end

  # Construct the outer loops that compute the weights

  new_inner = :( numerator = zero(T);
                 denominator = zero(T);
                 $outer;
                 $weights_term )
  new_outer = new_inner

  for i = N:-1:1#1:N
    new_outer = :(
      for $(i_vars[i]) = 1:(order[$i]+1)
        $new_outer
      end
    )
  end

    # Put it all together to compute the weights array

  final = :( weights = Array{T}(undef,(order.+1)...);
             $new_outer;
             return weights )

  return final

end

# Generated functions for tensor-product polynomials where nodes are in a tuple

@generated function chebyshev_weights(f::AbstractArray{T,N},nodes::NTuple{N,Array{T,1}},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T,N,S}

  chebyshev_polynomials = :( poly = Array{Array{T,2},1}(undef,N);
                             for k = 1:N;
                               orderk = order[k];
                               xk = nodes[k];

                               # normalize nodes

                               if domain[1,k] == domain[2,k];
                                 fill!(xk,(domain[1,k].+domain[2,k])/2);
                               else;
                                 xk = 2*(xk.-domain[2,k])/(domain[1,k].-domain[2,k]).-one(T);
                               end;

                               polynomial = Array{T}(undef,length(xk),orderk+1);
                               for i = 1:orderk+1;
                                 for j = 1:length(xk);
                                   if i == 1;
                                     polynomial[j,i] = one(T);
                                   elseif i == 2;
                                     polynomial[j,i] = xk[j];
                                   else;
                                     polynomial[j,i] = 2*xk[j]*polynomial[j,i-1]-polynomial[j,i-2];
                                   end;
                                 end;
                               end;
                               poly[k] = polynomial;#push!(poly,polynomial);
                             end )

  i_vars = Array{Symbol}(undef,N)
  s_vars = Array{Symbol}(undef,N)
  for i = 1:N
    i_vars[i] = Symbol("i$i")
    s_vars[i] = Symbol("s$i")
  end

  # Construct denominator term

  inner_prod_denominator = :( poly[1][s1,i1] )
  for i = 2:N
    inner_prod_denominator = :( poly[$i][$(s_vars[i]),$(i_vars[i])]*$inner_prod_denominator )
  end
  denominator_term = :( $inner_prod_denominator*$inner_prod_denominator )

  # Construct numerator term

  inner_prod_numerator = :( poly[1][s1,i1]*f[$(s_vars...)] )
  for i = 2:N
      inner_prod_numerator = :( poly[$i][$(s_vars[i]),$(i_vars[i])]*$inner_prod_numerator )
  end
  numerator_term = inner_prod_numerator

  # Construct weights term

  weights_term = :( weights[$(i_vars...)] = numerator/denominator )

  # Construct the inner loops that compute numerator and denominator

  inner = :( numerator += $numerator_term;
             denominator += $denominator_term )
  outer  = inner

  for i = N:-1:1#1:N
    outer = :(
      for $(s_vars[i]) = 1:size(f,$i)
        $outer
      end
    )
  end

  # Construct the outer loops that compute the weights

  new_inner = :( numerator = zero(T);
                 denominator = zero(T);
                 $outer;
                 $weights_term )
  new_outer = new_inner

  for i = N:-1:1#1:N
    new_outer = :(
      for $(i_vars[i]) = 1:(order[$i]+1)
        $new_outer
      end
    )
  end

  # Put it all together to compute the weights array

  final = :( $chebyshev_polynomials;
             weights = Array{T}(undef,(order.+1)...);
             $new_outer;
             return weights )

  return final

end

@generated function chebyshev_weights(f::AbstractArray{T,N},poly::NTuple{N,Array{T,2}},order::Array{S,1}) where {T,N,S}

  i_vars = Array{Symbol}(undef,N)
  s_vars = Array{Symbol}(undef,N)
  for i = 1:N
    i_vars[i] = Symbol("i$i")
    s_vars[i] = Symbol("s$i")
  end

  # Construct denominator term

  inner_prod_denominator = :( poly[1][s1,i1] )
  for i = 2:N
    inner_prod_denominator = :( poly[$i][$(s_vars[i]),$(i_vars[i])]*$inner_prod_denominator )
  end
  denominator_term = :( $inner_prod_denominator*$inner_prod_denominator )

  # Construct numerator term

  inner_prod_numerator = :( poly[1][s1,i1]*f[$(s_vars...)] )
  for i = 2:N
      inner_prod_numerator = :( poly[$i][$(s_vars[i]),$(i_vars[i])]*$inner_prod_numerator )
  end
  numerator_term = inner_prod_numerator

  # Construct weights term

  weights_term = :( weights[$(i_vars...)] = numerator/denominator )

  # Construct the inner loops that compute numerator and denominator

  inner = :( numerator += $numerator_term;
             denominator += $denominator_term )
  outer  = inner

  for i = N:-1:1#1:N
    outer = :(
      for $(s_vars[i]) = 1:size(f,$i)
        $outer
      end
    )
  end

  # Construct the outer loops that compute the weights

  new_inner = :( numerator = zero(T);
                 denominator = zero(T);
                 $outer;
                 $weights_term )
  new_outer = new_inner

  for i = N:-1:1#1:N
    new_outer = :(
      for $(i_vars[i]) = 1:(order[$i]+1)
        $new_outer
      end
    )
  end

  # Put it all together to compute the weights array

  final = :( weights = Array{T}(undef,(order.+1)...);
             $new_outer;
             return weights )

  return final

end

# Generated functions for complete polynomials where nodes are in an array of arrays

@generated function chebyshev_weights(f::AbstractArray{T,N},nodes::Array{Array{T,1},1},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T,N,S}

  chebyshev_polynomials = :( poly = Array{Array{T,2},1}(undef,N);
                             for k = 1:N;
                               xk = nodes[k];

                               # normalize nodes

                               if domain[1,k] == domain[2,k];
                                 fill!(xk,(domain[1,k].+domain[2,k])/2);
                               else;
                                 xk = 2*(xk.-domain[2,k])/(domain[1,k].-domain[2,k]).-one(T);
                               end;

                               polynomial = Array{T}(undef,length(xk),order+1);
                               for i = 1:order+1;
                                 for j = 1:length(xk);
                                   if i == 1;
                                     polynomial[j,i] = one(T);
                                   elseif i == 2;
                                     polynomial[j,i] = xk[j];
                                   else;
                                     polynomial[j,i] = 2*xk[j]*polynomial[j,i-1]-polynomial[j,i-2];
                                   end;
                                 end;
                               end;
                               poly[k] = polynomial;#push!(poly,polynomial);
                             end )

  i_vars = Array{Symbol}(undef,N)
  s_vars = Array{Symbol}(undef,N)
  for i = 1:N
    i_vars[i] = Symbol("i$i")
    s_vars[i] = Symbol("s$i")
  end

  # Construct denominator term

  inner_prod_denominator = :( poly[1][s1,i1] )
  for i = 2:N
    inner_prod_denominator = :( poly[$i][$(s_vars[i]),$(i_vars[i])]*$inner_prod_denominator )
  end
  denominator_term = :( $inner_prod_denominator*$inner_prod_denominator )

  # Construct numerator term

  inner_prod_numerator = :( poly[1][s1,i1]*f[$(s_vars...)] )
  for i = 2:N
      inner_prod_numerator = :( poly[$i][$(s_vars[i]),$(i_vars[i])]*$inner_prod_numerator )
  end
  numerator_term = inner_prod_numerator

  # Construct weights term

  weights_term = :( weights[$(i_vars...)] = numerator/denominator )

  # Construct the inner loops that compute numerator and denominator

  inner = :( numerator += $numerator_term;
             denominator += $denominator_term )
  outer  = inner

  for i = N:-1:1#1:N
    outer = :(
      for $(s_vars[i]) = 1:size(f,$i)
        $outer
      end
    )
  end

  # Construct the outer loops that compute the weights

  new_inner = :( numerator = zero(T);
                 denominator = zero(T);
                 if sum(tuple($(i_vars...))) <= order+N;
                   $outer;
                   $weights_term;
                 end )
  new_outer = new_inner

  for i = N:-1:1#1:N
    new_outer = :(
      for $(i_vars[i]) = 1:(order+1)
        $new_outer
      end
    )
  end

  # Initialize the weight Array

#=

  initial_weight = string("weights = zeros(order+1,",)
  for i = 2:N
    initial_weight = string(initial_weight,"order+1,")
  end
  initial_weight = parse(string(initial_weight,")"))

=#

  # Put it all together to compute the weights array

  final = :( $chebyshev_polynomials;
             weights = Array{T}(undef,(order.+1)...);
#             $initial_weight;
             $new_outer;
             return weights )

  return final

end

@generated function chebyshev_weights(f::AbstractArray{T,N},poly::Array{Array{T,2},1},order::S) where {T,N,S}

  i_vars = Array{Symbol}(undef,N)
  s_vars = Array{Symbol}(undef,N)
  for i = 1:N
    i_vars[i] = Symbol("i$i")
    s_vars[i] = Symbol("s$i")
  end

  # Construct denominator term

  inner_prod_denominator = :( poly[1][s1,i1] )
  for i = 2:N
    inner_prod_denominator = :( poly[$i][$(s_vars[i]),$(i_vars[i])]*$inner_prod_denominator )
  end
  denominator_term = :( $inner_prod_denominator*$inner_prod_denominator )

  # Construct numerator term

  inner_prod_numerator = :( poly[1][s1,i1]*f[$(s_vars...)] )
  for i = 2:N
      inner_prod_numerator = :( poly[$i][$(s_vars[i]),$(i_vars[i])]*$inner_prod_numerator )
  end
  numerator_term = inner_prod_numerator

  # Construct weights term

  weights_term = :( weights[$(i_vars...)] = numerator/denominator )

  # Construct the inner loops that compute numerator and denominator

  inner = :( numerator += $numerator_term;
             denominator += $denominator_term )
  outer  = inner

  for i = N:-1:1#1:N
    outer = :(
      for $(s_vars[i]) = 1:size(f,$i)
        $outer
      end
    )
  end

  # Construct the outer loops that compute the weights

  new_inner = :( numerator = zero(T);
                 denominator = zero(T);
                 if sum(tuple($(i_vars...))) <= order+N;
                   $outer;
                   $weights_term;
                 end )
  new_outer = new_inner

  for i = N:-1:1#1:N
    new_outer = :(
      for $(i_vars[i]) = 1:(order+1)
        $new_outer
      end
    )
  end

  # Initialize the weight Array

#=

  initial_weight = string("weights = zeros(order+1,",)
  for i = 2:N
    initial_weight = string(initial_weight,"order+1,")
  end
  initial_weight = parse(string(initial_weight,")"))

=#

  # Put it all together to compute the weights array

  final = :( weights = Array{T}(undef,(order.+1)...);
#             $initial_weight;
             $new_outer;
             return weights )

  return final

end

# Generated functions for complete polynomials where nodes are in a tuple

@generated function chebyshev_weights(f::AbstractArray{T,N},nodes::NTuple{N,Array{T,1}},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T,N,S}

  chebyshev_polynomials = :( poly = Array{Array{T,2},1}(undef,N);
                             for k = 1:N;
                               xk = nodes[k];

                               # normalize nodes

                               if domain[1,k] == domain[2,k];
                                 fill!(xk,(domain[1,k].+domain[2,k])/2);
                               else;
                                 xk = 2*(xk.-domain[2,k])/(domain[1,k].-domain[2,k]).-one(T);
                               end;

                               polynomial = Array{T}(undef,length(xk),order+1);
                               for i = 1:order+1;
                                 for j = 1:length(xk);
                                   if i == 1;
                                     polynomial[j,i] = one(T);
                                   elseif i == 2;
                                     polynomial[j,i] = xk[j];
                                   else;
                                     polynomial[j,i] = 2*xk[j]*polynomial[j,i-1]-polynomial[j,i-2];
                                   end;
                                 end;
                               end;
                               poly[k] = polynomial;#push!(poly,polynomial);
                             end )

  i_vars = Array{Symbol}(undef,N)
  s_vars = Array{Symbol}(undef,N)
  for i = 1:N
    i_vars[i] = Symbol("i$i")
    s_vars[i] = Symbol("s$i")
  end

  # Construct denominator term

  inner_prod_denominator = :( poly[1][s1,i1] )
  for i = 2:N
    inner_prod_denominator = :( poly[$i][$(s_vars[i]),$(i_vars[i])]*$inner_prod_denominator )
  end
  denominator_term = :( $inner_prod_denominator*$inner_prod_denominator )

  # Construct numerator term

  inner_prod_numerator = :( poly[1][s1,i1]*f[$(s_vars...)] )
  for i = 2:N
      inner_prod_numerator = :( poly[$i][$(s_vars[i]),$(i_vars[i])]*$inner_prod_numerator )
  end
  numerator_term = inner_prod_numerator

  # Construct weights term

  weights_term = :( weights[$(i_vars...)] = numerator/denominator )

  # Construct the inner loops that compute numerator and denominator

  inner = :( numerator += $numerator_term;
             denominator += $denominator_term )
  outer  = inner

  for i = N:-1:1#1:N
    outer = :(
      for $(s_vars[i]) = 1:size(f,$i)
        $outer
      end
    )
  end

  # Construct the outer loops that compute the weights

  new_inner = :( numerator = zero(T);
                 denominator = zero(T);
                 if sum(tuple($(i_vars...))) <= order+N;
                   $outer;
                   $weights_term;
                 end )
  new_outer = new_inner

  for i = N:-1:1#1:N
    new_outer = :(
      for $(i_vars[i]) = 1:(order+1)
        $new_outer
      end
    )
  end

  # Initialize the weight Array

  initial_weight = string("weights = zeros(order+1,",)
  for i = 2:N
    initial_weight = string(initial_weight,"order+1,")
  end
  initial_weight = Meta.parse(string(initial_weight,")"))

  # Put it all together to compute the weights array

  final = :( $chebyshev_polynomials;
#             weights = Array{T}(undef,(order.+1)...);
             $initial_weight;
             $new_outer;
             return weights )

  return final

end

@generated function chebyshev_weights(f::AbstractArray{T,N},poly::NTuple{N,Array{T,2}},order::S) where {T,N,S}

  i_vars = Array{Symbol}(undef,N)
  s_vars = Array{Symbol}(undef,N)
  for i = 1:N
    i_vars[i] = Symbol("i$i")
    s_vars[i] = Symbol("s$i")
  end

  # Construct denominator term

  inner_prod_denominator = :( poly[1][s1,i1] )
  for i = 2:N
    inner_prod_denominator = :( poly[$i][$(s_vars[i]),$(i_vars[i])]*$inner_prod_denominator )
  end
  denominator_term = :( $inner_prod_denominator*$inner_prod_denominator )

  # Construct numerator term

  inner_prod_numerator = :( poly[1][s1,i1]*f[$(s_vars...)] )
  for i = 2:N
      inner_prod_numerator = :( poly[$i][$(s_vars[i]),$(i_vars[i])]*$inner_prod_numerator )
  end
  numerator_term = inner_prod_numerator

  # Construct weights term

  weights_term = :( weights[$(i_vars...)] = numerator/denominator )

  # Construct the inner loops that compute numerator and denominator

  inner = :( numerator += $numerator_term;
             denominator += $denominator_term )
  outer  = inner

  for i = N:-1:1#1:N
    outer = :(
      for $(s_vars[i]) = 1:size(f,$i)
        $outer
      end
    )
  end

  # Construct the outer loops that compute the weights

  new_inner = :( numerator = zero(T);
                 denominator = zero(T);
                 if sum(tuple($(i_vars...))) <= order+N;
                   $outer;
                   $weights_term;
                 end )
  new_outer = new_inner

  for i = N:-1:1#1:N
    new_outer = :(
      for $(i_vars[i]) = 1:(order+1)
        $new_outer
      end
    )
  end

  # Initialize the weight Array

  initial_weight = string("weights = zeros(order+1,",)
  for i = 2:N
    initial_weight = string(initial_weight,"order+1,")
  end
  initial_weight = Meta.parse(string(initial_weight,")"))

  # Put it all together to compute the weights array

  final = :( #weights = Array{T}(undef,(order.+1)...);
             $initial_weight;
             $new_outer;
             return weights )

  return final

end

# Functions to ensure backward compatibility with an older API

# Tensor product polynomials

function chebyshev_weights(f::AbstractArray{T,1},nodes_1::Array{T,1},order::Array{S,1},domain=[ones(1);-ones(1)]) where {T<:AbstractFloat,S<:Integer}

  weights = chebyshev_weights(f,(nodes_1,),order,reshape(domain,2,1))

  return weights

end

function chebyshev_weights(f::AbstractArray{T,2},nodes_1::Array{T,1},nodes_2::Array{T,1},order::Array{S,1},domain=[ones(1);-ones(1)]) where {T<:AbstractFloat,S<:Integer}

  weights = chebyshev_weights(f,(nodes_1,nodes_2),order,domain)

  return weights

end

function chebyshev_weights(f::AbstractArray{T,3},nodes_1::Array{T,1},nodes_2::Array{T,1},nodes_3::Array{T,1},order::Array{S,1},domain=[ones(1);-ones(1)]) where {T<:AbstractFloat,S<:Integer}

  weights = chebyshev_weights(f,(nodes_1,nodes_2,nodes_3),order,domain)

  return weights

end

function chebyshev_weights(f::AbstractArray{T,4},nodes_1::Array{T,1},nodes_2::Array{T,1},nodes_3::Array{T,1},nodes_4::Array{T,1},order::Array{S,1},domain=[ones(1);-ones(1)]) where {T<:AbstractFloat,S<:Integer}

  weights = chebyshev_weights(f,(nodes_1,nodes_2,nodes_3,nodes_4),order,domain)

  return weights

end

function chebyshev_weights(f::AbstractArray{T,5},nodes_1::Array{T,1},nodes_2::Array{T,1},nodes_3::Array{T,1},nodes_4::Array{T,1},nodes_5::Array{T,1},order::Array{S,1},domain=[ones(1);-ones(1)]) where {T<:AbstractFloat,S<:Integer}

  weights = chebyshev_weights(f,(nodes_1,nodes_2,nodes_3,nodes_4,nodes_5),order,domain)

  return weights

end

function chebyshev_weights(f::AbstractArray{T,6},nodes_1::Array{T,1},nodes_2::Array{T,1},nodes_3::Array{T,1},nodes_4::Array{T,1},nodes_5::Array{T,1},nodes_6::Array{T,1},order::Array{S,1},domain=[ones(1);-ones(1)]) where {T<:AbstractFloat,S<:Integer}

  weights = chebyshev_weights(f,(nodes_1,nodes_2,nodes_3,nodes_4,nodes_5,nodes_6),order,domain)

  return weights

end

# Complete polynomials

function chebyshev_weights(f::AbstractArray{T,1},nodes_1::Array{T,1},order::S,domain=[ones(1);-ones(1)]) where {T<:AbstractFloat,S<:Integer}

  weights = chebyshev_weights(f,(nodes_1,),order,reshape(domain,2,1))

  return weights

end

function chebyshev_weights(f::AbstractArray{T,2},nodes_1::Array{T,1},nodes_2::Array{T,1},order::S,domain=[ones(1);-ones(1)]) where {T<:AbstractFloat,S<:Integer}

  weights = chebyshev_weights(f,(nodes_1,nodes_2),order,domain)

  return weights

end

function chebyshev_weights(f::AbstractArray{T,3},nodes_1::Array{T,1},nodes_2::Array{T,1},nodes_3::Array{T,1},order::S,domain=[ones(1);-ones(1)]) where {T<:AbstractFloat,S<:Integer}

  weights = chebyshev_weights(f,(nodes_1,nodes_2,nodes_3),order,domain)

  return weights

end

function chebyshev_weights(f::AbstractArray{T,4},nodes_1::Array{T,1},nodes_2::Array{T,1},nodes_3::Array{T,1},nodes_4::Array{T,1},order::S,domain=[ones(1);-ones(1)]) where {T<:AbstractFloat,S<:Integer}

  weights = chebyshev_weights(f,(nodes_1,nodes_2,nodes_3,nodes_4),order,domain)

  return weights

end

function chebyshev_weights(f::AbstractArray{T,5},nodes_1::Array{T,1},nodes_2::Array{T,1},nodes_3::Array{T,1},nodes_4::Array{T,1},nodes_5::Array{T,1},order::S,domain=[ones(1);-ones(1)]) where {T<:AbstractFloat,S<:Integer}

  weights = chebyshev_weights(f,(nodes_1,nodes_2,nodes_3,nodes_4,nodes_5),order,domain)

  return weights

end

function chebyshev_weights(f::AbstractArray{T,6},nodes_1::Array{T,1},nodes_2::Array{T,1},nodes_3::Array{T,1},nodes_4::Array{T,1},nodes_5::Array{T,1},nodes_6::Array{T,1},order::S,domain=[ones(1);-ones(1)]) where {T<:AbstractFloat,S<:Integer}

  weights = chebyshev_weights(f,(nodes_1,nodes_2,nodes_3,nodes_4,nodes_5,nodes_6),order,domain)

  return weights

end


f = open("2019/01/input.txt", "r")
inlines = [parse(Int, x) for x in readlines(f)]
close(f)

outlines = [[(x÷3)-2] for x in inlines]

while true
    going = false
    for l in outlines
        fuel = (last(l)÷3)-2
        if fuel > 0
            #println(fuel)
            push!(l, fuel)
            going = true
        end
    end
    if !going
        break
    end
end

println(sum([sum(l) for l in outlines]))


function test_point_sampling()
    radius = 1.0f0
    origin = Point3f0(0, 0, 0)
    mesh = icosphere(4, radius)
    verts = coordinates(mesh)
    triangles = faces(mesh)
    areas_array = area.(Ref(verts), triangles)
    cumsum_area = similar(areas_array)
    cumsum_area = cumsum!(cumsum_area, areas_array)
    total_area = cumsum_area[end]
    bins = cumsum_area / total_area
    point1 = point_sampling(verts, triangles, bins)
    d_point1 = norm(point1 - origin)
    @test isapprox(d_point1, radius, atol = 0.001)
    points1 = points_sampling(mesh, 10000)
    points_dist = norm.(points1)
    mean_dist = mean(points_dist)
    sd_dist = std(points_dist)
    @test isapprox(mean_dist, 1.0f0, atol = 0.001)
    @test isapprox(sd_dist, 0.0f0, atol = 0.001)
end


using Pkg
using CSV
using DataFrames
using Distributions
using Random
using StatsBase
using NamedArrays
using Plots
using PyPlot

prices = [20, 25, 18, 20, 15]
score = [0.0, 0, 0, 0]
<<<<<<< HEAD

##Rozkłady
#Ilość samochodów na minutę tworzymy rozkład zeby uzyc sample w funkcji
cars = [1, 2, 3, 4, 5]
prob = [0.15, 0.25, 0.2, 0.2, 0.1]
w = Weights(prob)

##Główna funkcja kroku
#jeden krok = jedna minuta

# function OneStep()
#     #wylosować z rozkładu samochdów na minutę ile samochdów przyjedzie w tym kroku
#     cars_num = sample(cars, w)
#     #wylosować ilość osób w kazdym samochodzie, 1 samochod = jedna zmienna?
#     num_ppl = Dict(i => rand(1:cars_num) for i in 1:cars_num)
#     # wylosować zamówienia dla kazdego samochodu, zapisujemy je jako data DataFrame
#     #gdzie kolumny to produkty zamowione a rzedy to nr samochodu
#     orders = DataFrame(BigMac = zeros(cars_num), Drwal = zeros(cars_num),
#     McNuggets = zeros(cars_num), McChicken = zeros(cars_num), Wege = zeros(cars_num))

#     #losujemy zamowienia dla kazdego samochodu
#     for i in 1:nrow(orders)
#         for j in 1:length(num_ppl)
#             prod = rand(1:ncol(orders))
#             orders[i, prod] += 1
#         end
#     end

function OneStep(start_inv)
    cars_num = sample(cars, w)
=======

function OneStep(start_inv, start, num_cars, weights, start_val, val)
    cars_num = sample(num_cars, weights)
>>>>>>> 0c7659ff753bbe001737e16c6caf863e93042415
    #wylosować ilość osób w kazdym samochodzie, 1 samochod = jedna zmienna?
    num_ppl = Dict(i => rand(1:cars_num) for i in 1:cars_num)
    # wylosować zamówienia dla kazdego samochodu, zapisujemy je jako data DataFrame
    #gdzie kolumny to produkty zamowione a rzedy to nr samochodu
    products = ["BigMac", "Drwal", "McNuggets", "McChicken", "Wege"]
    orders = DataFrame(zeros(Int64, cars_num, length(products)), products)

    #losujemy zamowienia dla kazdego samochodu
    for i in 1:nrow(orders)
        for j in 1:length(num_ppl)
                prod = rand(1:ncol(orders)) # można zmienić na jakiś sensowniejszy rozkład
                orders[i, prod] += 1 # każda osoba wybiera po jednym produkcie?
        end
    end

    #sumujemy zamówienia po kolumnach z wszystkich samochodów
    sum_order = sum.(eachcol(orders))
    #odejmujemy od inventory zamówione składniki
    end_inv = start - sum_order
    #vektor ile zamówień nie udalo sie wypelnic na razie 0
    unfulfilled = zeros(5)

    #jesli inventory zeszlo ponizej zera to tyle ile jest na minusie tyle zamowien nie uda sie spelnic
    for i in 1:length(end_inv)
        if(end_inv[i] < 0)
            unfulfilled[i] -= end_inv[i]
            end_inv[i] = 0
        end
    end
    #spelnione zamowienia to wszystkie zamowienia - niespelnione
    fulfilled = sum_order - unfulfilled
    #przychod
    gain = fulfilled .* prices
    #kara za niespelnienie zamowien
    penalty = unfulfilled .* (0.25*prices)
    #zmniejszamy o jedna minute waznosc produktow
    val .-= 1
    #to co musielismy wyrzucic na razie zera
    thrown_out = zeros(5)
    #jesli waznosc spadnie do 0 to musimy wyrzucic to co zostalo i zrobic nowe oraz zaktualizowac date wwaznosci
    for i in 1:length(val)
        if(val[i] <= 0)
            thrown_out[i] = end_inv[i]
            end_inv[i] += start_inv[i]
            val[i] = start_val[i]
        end
    end
    #strata policzona na podstawie wyrzuconych
    #mozna zmienic te 0.25 i powiedziec ze co sie stanie jak produkty zdrozeja
    loss = thrown_out.* (0.25*prices)
    # zysk
    profit = gain - penalty - loss
    #wyniki
    score[1] = sum(profit) #zysk
    score[2] = sum(thrown_out)#ile wyrzucilismy w tej minucie
    score[3] = abs(sum(unfulfilled)) #ile zamowien zostalo niespelnionych w  tej minucie
    score[4] = sum(fulfilled) # ile zamowien spelnionych w tej minucie

    #nadpisujemy startowe inventory
    start = end_inv
    #zwracamy wynik koncowe inventory oraz obecna date waznosci do nastepnej iteracji
    return score, end_inv, val
end

#horizon to ilosc minut jaka chcemy miec w symulacji najlepiej 60
#val to daty waznosci
#start to inventory na poczatku minuty
#mean_val to ile razy odpalamy symulacje zeby policzyc srednie wartosci
function SimulateOneRun(horizon, val, start, mean_val, prob)
    #Ceny produktów
    prices = [20, 30, 23, 25, 18]
    #pusty wektor do wynikow
    summed_out = zeros(4)
    #tutaj jest petla zeby wyciagac srednia z wynikow, poniewaz losujemy wiele rzeczy
    #chcemy zeby wyniki byly srednia z wielu przebiegow
    #mean_val to ilosc iteracji do liczenia sredniej (im wiecej tym lepiej)
    for i in 1:mean_val
        #obecne daty waznosci
        _val = val
        #stałe poczatkowe daty
        start_val = deepcopy(_val)
        score = [0.0, 0.0, 0.0, 0.0]
        #obecny stan inventory
        _start = start
        #staly poczatkowy stan inventory
        start_inventory = deepcopy(_start)
        #Ilość samochodów na minutę tworzymy rozkład zeby uzyc sample w funkcji
        cars = [1, 2, 3, 4, 5]
        #prob = [0.15, 0.25, 0.2, 0.2, 0.1]
        w = Weights(prob)
        #Wynik symulacji jako dataframe
        df = DataFrame(Profit = 0.0, Thrown_out = 0.0, Unfulfilled = 0.0, Fulfilled = 0.0)
        #symulacja, horizon powinno byc 60 bo symulujemy godzine, mozna zwiekszac
        for i in 1:horizon
            #out to wyniki kroku, _start to stan inventory na koniec kroku/minuty,
            #_val to daty waznosci na koniec kroku/minuty
            out, _start, _val = OneStep(start_inventory, _start, cars, w, start_val, _val)
            #pushujemy do data frame
            push!(df, out[1:4])
        end
        #bierzemy wszystko od drugiego wyniku bo pierwszy to zera,
        #bo tworzylismy data frame z rzedem zer wczesniej
        df_out = df[2:end,:]
        #tutaj musimy skleic rzedami wyniki z kazdej z tych symulacji zeby potem policzyc srednia po rzedach
        summed_out = hcat(summed_out, sum.(eachcol(df_out)))
    end

    #println(start_val)
    #liczymy srednia po rzedach od drugiej do ostatniej kolumny, bo pierwsza to zera
    return mean(summed_out[:, 2:end], dims = 2)
end

#mean_val to ilosc razy jaka odpalimy simulateOneRun zeby policzyc srednia, tam wyzej bylo opisane
function SimulatemanyRuns(mean_val, prob)
    #daty waznosci
    validity = [10, 10, 10, 10, 10]
    #data frame z wynikami
    temp = DataFrame(Profit = 0.0, Thrown_out = 0.0, Unfulfilled = 0.0, Fulfilled = 0.0)
    #przykladowe stany inventory od 1 sztuki do 40 sztuk
    x = [fill(i, 5) for i in 1:40]
    #petla dla kazdego stanu inventory
    for i in x
       push!(temp, SimulateOneRun(60, validity, i, mean_val, prob))
    end
    return temp
end

<<<<<<< HEAD

final = SimulateOneRun(60, [20, 20, 20, 20, 20])

=======
#uruchomienie symulacji
@time xd = SimulatemanyRuns(100, [0.15, 0.25, 0.2, 0.2, 0.1])[2:end, :]
#to wyswietla caly data frame bez ucinania rzedow
show(stdout, "text/plain", xd)
#---------------------------------------------------------
###sprawdzmy dla jakich wartosci jest najwiekszy profit
#wylaczamy kolumne profit
profit = xd[!, "Profit"]
#wykres slupkowy
plt.clf()
plt.bar([1:1:40;], profit, label = "Profit")
#splotujmy jeszcze raz od 10 do 30
profit_2 = profit[10:30]
plt.bar([10:1:30;], profit_2)
profit_3 = profit[18:22]
plt.bar([18:1:22;], profit_3)
xlabel("Number of products")
ylabel("Profit")
gcf()
#---------------------------------------------
#sprawdzmy ile było wyrzucane jedzenia
thrown = xd[!, "Thrown_out"]
plt.clf()
plt.bar([1:1:40;], thrown, color = "red", alpha = 0.7)
xlabel("Number of products")
ylabel("Loss from wasted products")
gcf()

#nakladamy na profit
plt.clf()
plt.bar([1:1:40;], profit, label = "Profit")
#splotujmy jeszcze raz od 10 do 30
profit_2 = profit[10:30]
plt.bar([10:1:30;], profit_2)
profit_3 = profit[18:22]
plt.bar([18:1:22;], profit_3)
plt.bar([1:1:40;], thrown, alpha=0.5)
gcf()

#ilosc spelnionych i niespelnionych zamowien
fulfilled = xd[!, "Fulfilled"]
unfulfilled = xd[!, "Unfulfilled"]
plt.clf()
plt.bar([1:1:40;], fulfilled, color = "green")
plt.bar([1:1:40;], unfulfilled, color = "red")
legend(["Fulfilled", "Unfulfilled"])
xlabel("Number of products")
ylabel("Orders")
gcf()

#zmiana ruchu samochodów na bardziej intensywny i mniej intensywny
@time xd_2 = SimulatemanyRuns(100, [0.1, 0.2, 0.25, 0.25, 0.2])[2:end, :]
@time xd_3 = SimulatemanyRuns(100, [0.25, 0.3, 0.2, 0.15, 0.1])[2:end, :]
profit = xd[!, "Profit"]
profit_22 = xd_2[!, "Profit"]
profit_33 = xd_3[!, "Profit"]
#wykres slupkowy
plt.clf()

plt.bar([1:1:40;], profit_22, alpha = 0.5, color = "red")
plt.bar([1:1:40;], profit, color = "green")
plt.bar([1:1:40;], profit_33, alpha = 0.8, color = "blue")
xlabel("Number of products")
ylabel("Profit")
legend(["Low Traffic", "High Traffic", "Medium traffic"])
gcf()
>>>>>>> 0c7659ff753bbe001737e16c6caf863e93042415


module WindowTrans

using DSP

# The windowed transforms work by sliding a window over the signal and
# performing the Fourier transform. The input signal is assumed to be
# apodized (zero at the ends of the interval).

function wind_transf(t, x, w)
    f = fftshift(fftfreq(length(t),1.0/(t[2]-t[1])))

    Cₓ = hcat([x.*circshift(w,i) for i = 1:length(t)]...)
    Wₓ = fftshift(fft(Cₓ, 1), 1)
    f,Wₓ
end

#=
The Gabor transform
$$
G_x(t,f) = \int_{-\infty}^\infty \mathrm{d}\tau
\exp[-\pi(\tau-t)^2]\exp(-j\pi f\tau)x(\tau)
$$

Here, we use a version with variable width.
=#

gabor(t, x, σ) = wind_transf(t, x, fftshift(exp.(-(t-mean(t)).^2/2σ^2)))

#=
The Wigner distribution function (WDF) is defined as [[1]]

$$
W_x(t,f) =
\int_{-\infty}^\infty d\tau
x\left(t+\frac{\tau}{2}\right)
x^*\left(t-\frac{\tau}{2}\right)
\exp(-j2\pi f\tau)
$$

Thus, for every $t$, we need to form

$$C_x\left(t+\frac{\tau}{2}, t-\frac{\tau}{2}\right)
\equiv x\left(t+\frac{\tau}{2}\right)
x^*\left(t-\frac{\tau}{2}\right),$$

but since the integration limits are $\pm\infty$, we instead form

$$x\left(t\right)
x^*\left(t-\tau\right).$$

The function is assumed to be windowed on input.

The good thing about the WDF is that the resolution is very high. The
bad thing is the cross term that shows up as soon as there are more
than one frequency present in the signal.

[1]: http://dx.doi.org/10.1109/TSP.2007.896271
=#
wigner(t, x) = wind_transf(t, x, fftshift(conj(reverse(x))))

#=
The "NMR" transform uses a window function that is derivating--apodizing:

$$
W(t) = t^n\exp(-\lambda_w t)
$$
=#
nmr(t, x, λ, n = 2) = wind_transf(t, x, t.^n.*exp(-λ*t))

export wind_transf, gabor, wigner, nmr

end # module


mutable struct MSetLPSolver <: MSetSolver
  solver
  m # Int
  n_arms # Int

  # Optimisation model (reused by solve_linear).
  model # ::Model
  x # ::Vector{Variable}

  function MSetLPSolver(solver)
    return new(solver, nothing, nothing, nothing, nothing)
  end
end

function build!(solver::MSetLPSolver, m::Int, n_arms::Int)
  solver.m = m
  solver.n_arms = n_arms
  solver.model = Model(solver.solver)
  solver.x = @variable(solver.model, [1:n_arms], binary=true)

  for i in 1:n_arms
    set_name(solver.x[i], "x_$i")
  end

  @constraint(solver.model, sum(solver.x) <= m)
end

has_lp_formulation(::MSetLPSolver) = true

function get_lp_formulation(solver::MSetLPSolver, rewards::Dict{T, Float64}) where T
  return solver.model,
    sum(rewards[i] * solver.x[i] for i in keys(rewards)),
    Dict{Int, JuMP.VariableRef}(i => solver.x[i] for i in keys(rewards))
end

function solve_linear(solver::MSetLPSolver, rewards::Dict{Int, Float64})
  m, obj, vars = get_lp_formulation(solver, rewards)
  @objective(m, Max, obj)

  # Maybe solve_budgeted_linear has already been called on this model.
  if :MSetLP in keys(m.ext)
    budget_constraint = m.ext[:MSetLP][:budget_constraint]
    set_normalized_rhs(budget_constraint, 0)
  end

  set_silent(m)
  optimize!(m)

  if termination_status(m) != MOI.OPTIMAL
    return Int[]
  end
  return Int[i for i in keys(rewards) if value(vars[i]) > 0.5]
end

function solve_budgeted_linear(solver::MSetLPSolver, rewards::Dict{Int, Float64}, weights::Dict{Int, Int}, budget::Int)
  m, obj, vars = get_lp_formulation(solver, rewards)
  @objective(m, Max, obj)

  # Add the budget constraint (or change the existing constraint).
  if :MSetLP in keys(m.ext)
    budget_constraint = m.ext[:MSetLP][:budget_constraint]
    set_normalized_rhs(budget_constraint, budget)
  else
    budget_constraint = @constraint(m, sum(weights[i] * solver.x[i] for i in keys(rewards)) >= budget)
    m.ext[:MSetLP] = Dict(:budget_constraint => budget_constraint)
  end

  set_silent(m)
  optimize!(m)

  if termination_status(m) != MOI.OPTIMAL
    return Int[]
  end
  return Int[i for i in keys(rewards) if value(vars[i]) > 0.5]
end


"""
This script contains all functions used for TSP
"""

"Construct TSP pre-model, without cycle elimination mechanism"
function prebuild_tsp(dist::Matrix)

    n = size(dist, 1)
    # Definition of model
    model = Model(GLPK.Optimizer)
    set_optimizer_attribute(model, "tm_lim", 300 * 1_000)
    #model = Model(Gurobi.Optimizer)
    #set_optimizer_attributes(model, "OutputFlag" => 0)

    # Main variable: x_ij=1 if the tour visits i and j in that order, 0 otherwise
    @variable(model, x[1:n, 1:n], Bin)
    # Objective: minimizing the total cost (distance) of the tour
    @objective(model, Min, sum(dist[i, j] * x[i, j] for i = 1:n, j = 1:n))
    # SHARED CONSTRAINTS
    @constraint(
        model, no_self_edges[i = 1:n], x[i,i] == 0
    )
    @constraint(
        model, exactly_one_successor[i = 1:n], sum(x[i, j] for j = 1:n) == 1
    )
    @constraint(
        model, exactly_one_predecessor[j = 1:n], sum(x[i, j] for i = 1:n) == 1
    )
    return model, x
end

"""
    Given the induced graph as an adjacency list (i.e., next[i] is the next node to visit after node i),
        compute all subtours.
    Return them as a list of lists of nodes in the same component
"""
function find_subtours(next::Vector{Int})
    n = length(next)
    g = DiGraph(n)
    for i = 1:n
        add_edge!(g, i, next[i])
    end
    components = strongly_connected_components(g)
    return sort(components, by=length)
end

"Solve the TSP using an iterative approach"
function solve_iterative(dist_matrix::Dict, zone_list::Array; lower_bound=0, time_limit_seconds::Real = 600,
                         eliminate_length_2::Bool=false,
                         verbose::Bool = true)
    # We first solve the model without any subtour elimination consideration

    dist = zeros(length(zone_list), length(zone_list))
    for i in 1:length(zone_list), j in 1:length(zone_list)
        if i != j
            dist[i, j] = dist_matrix[zone_list[i]][zone_list[j]]
        end
    end

    model, x = prebuild_tsp(dist)
    n = size(dist,1)
    @constraint(model, sum(dist[i, j] * x[i, j] for i = 1:n, j = 1:n) >= lower_bound)
    if eliminate_length_2
        @constraint(model, no_length_2[i = 1:n, j = 1:n], x[i, j] + x[j, i] <= 1)
    end
    verbose || set_optimizer_attribute(model, "OutputFlag", 0)
    start=time()
    optimize!(model)

    while true
        # We store the incumbent solution
        next = [findfirst(x -> x > 0.5, value.(x[i, :])) for i = 1:n]
        # Note: checking for >0.5 is conservative (x is binary!) but it avoids numerical errors
        subtours = find_subtours(next)
        #println("Found $(length(subtours)) subtours after $(time() - start) seconds")
        if length(subtours) == 1 # only one cycle, the TSP solution
            solvetime = time() - start
            return solvetime, objective_value(model), value.(x)
        else
            # eliminate subtours
            for subtour in subtours
                @constraint(model, sum(x[i, j] for i=subtour, j=setdiff(1:n, subtour)) >= 1)
                @constraint(model, sum(x[i, j] for i=setdiff(1:n, subtour), j=subtour) >= 1)
            end
        end
        optimize!(model)
        time() - start > time_limit_seconds && return solvetime, objective_value(model), value.(x)
    end
end

function calculate_tsp_seq(route_df, zone_distance_matrix; first_zone = "INIT")

    zone_list = unique(skipmissing(route_df.zone_id))
    opt_zone_seq = zone_list
    compute_time, obj, opt_x = solve_iterative(zone_distance_matrix, zone_list)
    opt_ind = find_subtours([findfirst(x -> x > 0.5, opt_x[i, :]) for i = 1:size(opt_x,1)])[1]
    init_output_seq = [zone_list[i] for i in opt_ind]
    init_ind = findfirst(x -> x == first_zone,init_output_seq)
    zone_output_seq = vcat(init_output_seq[init_ind:end],init_output_seq[1:init_ind-1])

    return zone_output_seq
end

"Construct path TSP pre-model, without cycle elimination mechanism"
function prebuild_path_tsp(dist::Matrix, s, t)

    n = size(dist, 1)
    other_node_list = setdiff(collect(1:n), [s, t])
    # Definition of model
    model = Model(GLPK.Optimizer)
    set_optimizer_attribute(model, "tm_lim", 300 * 1_000)
    #model = Model(Gurobi.Optimizer)
    #set_optimizer_attributes(model, "OutputFlag" => 0)

    # Main variable: x_ij=1 if the tour visits i and j in that order, 0 otherwise
    @variable(model, x[1:n, 1:n], Bin)
    # Objective: minimizing the total cost (distance) of the tour
    @objective(model, Min, sum(dist[i, j] * x[i, j] for i = 1:n, j = 1:n))
    # SHARED CONSTRAINTS
    @constraint(
        model, no_self_edges[i = 1:n], x[i,i] == 0
    )
    @constraint(
        model, exactly_one_successor[i in other_node_list], sum(x[i, j] for j = 1:n) == 1
    )
    @constraint(model, sum(x[s, j] for j in 1:n) == 1)
    @constraint(model, sum(x[j, s] for j in 1:n) == 0)
    @constraint(model, sum(x[i, t] for i in 1:n) == 1)
    @constraint(model, sum(x[t, i] for i in 1:n) == 0)
    @constraint(
        model, exactly_one_predecessor[j in other_node_list], sum(x[i, j] for i = 1:n) == 1
    )
    return model, x
end

"""
    Given the induced graph as an adjacency list (i.e., next[i] is the next node to visit after node i),
        compute all subtours.
    Return them as a list of lists of nodes in the same component
"""
function find_path_subtours(next)
    n = length(next)
    g = DiGraph(n)
    for i = 1:n
        if next[i] != nothing
            add_edge!(g, i, next[i])
        end
    end
    cycles = simplecycles(g)
    #components = connected_components(g)
    return cycles #sort(components, by=length)
end

"Solve the TSP using an iterative approach"
function solve_path_iterative(dist_matrix::Dict, zone_list::Array, first_stop, last_stop; time_limit_seconds::Real = 600,
                         eliminate_length_2::Bool=false)
    # We first solve the model without any subtour elimination consideration

    dist = zeros(length(zone_list), length(zone_list))
    for i in 1:length(zone_list), j in 1:length(zone_list)
        if i != j
            dist[i, j] = dist_matrix[zone_list[i]][zone_list[j]]
        end
    end

    s = findfirst(x -> x == first_stop, zone_list)
    t = findfirst(x -> x == last_stop, zone_list)

    model, x = prebuild_path_tsp(dist, s, t)
    n = size(dist,1)
    if eliminate_length_2
        @constraint(model, no_length_2[i = 1:n, j = 1:n], x[i, j] + x[j, i] <= 1)
    end
    start=time()
    optimize!(model)

    while true
        # We store the incumbent solution
        next = [findfirst(x -> x > 0.5, value.(x[i, :])) for i = 1:n]
        # Note: checking for >0.5 is conservative (x is binary!) but it avoids numerical errors
        subtours = find_path_subtours(next)
        #println("Found $(length(subtours)) subtours after $(time() - start) seconds")
        if length(subtours) == 0 # no cycles
            solvetime = time() - start
            return solvetime, objective_value(model), value.(x), s
        else
            # eliminate subtours
            for subtour in subtours
                @constraint(model, sum(x[i, j] for i=subtour, j=setdiff(1:n, subtour)) >= 1)
                @constraint(model, sum(x[i, j] for i=setdiff(1:n, subtour), j=subtour) >= 1)
            end
        end
        optimize!(model)
        time() - start > time_limit_seconds && return solvetime, objective_value(model), value.(x), s
    end
end

function calculate_path_tsp_seq(route_df, first_zone, last_zone, zone_distance_matrix)

    zone_list = unique(skipmissing(route_df.zone_id))
    if first_zone != "INIT" && last_zone != "INIT"
        zone_list = zone_list[2:end]
    end

    adj_zone_distance_matrix = Dict()
    for i in zone_list
        adj_zone_distance_matrix[i] = Dict()
        for j in zone_list
            adj_zone_distance_matrix[i][j] = 0
        end
    end

    for i in zone_list
        for j in zone_list
            adj_zone_distance_matrix[i][j] = zone_distance_matrix[i][j]
        end
    end

    compute_time, obj, opt_x = solve_path_iterative(adj_zone_distance_matrix, zone_list, first_zone, last_zone)

    next = [findfirst(x -> x > 0.5, opt_x[i, :]) for i = 1:size(opt_x,1)]

    output_seq = [first_zone]
    while length(output_seq) < length(zone_list)
        next_zone_ind = next[findfirst(x -> x == output_seq[end], zone_list)]
        next_zone = zone_list[next_zone_ind]
        push!(output_seq, next_zone)
    end

    @assert output_seq[1] == first_zone
    @assert output_seq[end] == last_zone

    return output_seq
end

function find_first_and_last_stop_list(prev_node, zone, next_zone, zone_to_node_dict, stop_distance_matrix)

    zone_node = zone_to_node_dict[zone]
    next_zone_node = zone_to_node_dict[next_zone]

    prev_dist_list = []
    next_dist_list = []
    for node in zone_node

        prev_dist = stop_distance_matrix[prev_node][node]
        push!(prev_dist_list, (prev_dist, node))

        next_dist = 0
        for next_node in next_zone_node
            next_dist += stop_distance_matrix[node][next_node]
        end
        push!(next_dist_list, (next_dist, node))
    end

    first_node_list = [i[2] for i in sort(prev_dist_list)[1:min(length(zone_node), 3)]]
    last_node_list = [i[2] for i in sort(next_dist_list)[1:min(length(zone_node), 3)]]

    return first_node_list, last_node_list
end

function generate_complete_seq_multiple(zone_seq, route_df, stop_distance_matrix)
    zone_output_seq = zone_seq # opt_zone_seq

    # Generate zone-node correspondence dict
    zone_to_node_dict = Dict(i => [] for i in zone_seq)
    selected_df = route_df[:, ["stops", "zone_id"]]
    for i in 1:size(selected_df, 1)
        stop_id = selected_df[i,"stops"]
        zone_id = selected_df[i,"zone_id"]
        push!(zone_to_node_dict[zone_id], stop_id)
    end

    init_stop = zone_to_node_dict["INIT"][1]
    output_complete_sequence = [init_stop]
    prev_node = init_stop
    for i in 2:length(zone_output_seq)
        curr_zone = zone_output_seq[i]
        # Zone only contains one stop
        if length(zone_to_node_dict[curr_zone]) <= 1
            push!(output_complete_sequence, zone_to_node_dict[curr_zone][1])
            continue
        end
        if i == length(zone_output_seq)
            next_zone = zone_output_seq[1]
        else
            next_zone = zone_output_seq[i+1]
        end
        first_node_list, last_node_list = find_first_and_last_stop_list(prev_node, curr_zone, next_zone, zone_to_node_dict, stop_distance_matrix)
        stop_list = zone_to_node_dict[curr_zone]

        inner_option_list = []
        for first_node in first_node_list
            for last_node in last_node_list
                if first_node == last_node
                    compute_time, obj, opt_x = solve_iterative(stop_distance_matrix, stop_list)
                    opt_ind = find_subtours([findfirst(x -> x > 0.5, opt_x[i, :]) for i = 1:size(opt_x,1)])[1]
                    init_output_seq = [stop_list[i] for i in opt_ind]
                    init_ind = findfirst(x -> x == first_node,init_output_seq)
                    output_path = vcat(init_output_seq[init_ind:end],init_output_seq[1:init_ind-1])
                    reduced_obj = obj - stop_distance_matrix[output_path[end]][output_path[1]]

                    push!(inner_option_list, (reduced_obj, first_node, last_node, output_path))
                else
                    compute_time, obj, opt_x, s = solve_path_iterative(stop_distance_matrix, stop_list, first_node, last_node);
                    next = [findfirst(x -> x > 0.5, opt_x[i, :]) for i = 1:size(opt_x,1)]
                    output_seq = [s]
                    while length(output_seq) < length(stop_list)
                        next_visit_node = next[output_seq[end]]
                        push!(output_seq, next_visit_node)
                    end
                    output_path = [stop_list[i] for i in output_seq]

                    @assert first_node == output_path[1]
                    @assert last_node == output_path[end]

                    push!(inner_option_list, (obj, first_node, last_node, output_path)) # output_path
                end
            end
        end

        inner_zone_tsp_path = sort(inner_option_list)[1][4]
        first_node = inner_zone_tsp_path[1]
        last_node = inner_zone_tsp_path[end]

        output_complete_sequence = vcat(output_complete_sequence, inner_zone_tsp_path)

        prev_node = last_node
    end

    return output_complete_sequence
end

function adjust_zone_dist_matrix(zone_distance_matrix, zone_list, route_min_tt_nearest, route_ecu_tt_nearest, β, γ, α)

    adj_zone_dist_matrix = Dict(i => Dict() for i in zone_list)

    for i in zone_list, j in zone_list

        try
            global big_zone_i = split(i, ".")[1]
            global big_zone_j = split(j, ".")[1]
        catch
            global big_zone_i = i
            global big_zone_j = j
        end

        if big_zone_i != "INIT" && big_zone_j != "INIT" && big_zone_i != big_zone_j
            adj_zone_dist_matrix[i][j] = zone_distance_matrix[i][j] * β
        else
            adj_zone_dist_matrix[i][j] = zone_distance_matrix[i][j]
        end

        if i == "INIT"
            if j ∉ route_ecu_tt_nearest[route_ecu_tt_nearest.from_zone .== i, :].to_zone && j ∉ route_min_tt_nearest[route_min_tt_nearest.from_zone .== i, :].to_zone
                adj_zone_dist_matrix[i][j] = zone_distance_matrix[i][j] * α
            end
        end

        # if i == "INIT" && j ∉ route_zone_min_tt_df[route_zone_min_tt_df.from_zone .== i, :].to_zone
        #     adj_zone_dist_matrix[i][j] = zone_distance_matrix[i][j] * α
        # end

        if i == "INIT" || j == "INIT"
            continue
        end

        try
            global after_dot_i1 = split(i, ".")[2][1]
            global after_dot_i2 = split(i, ".")[2][2]
            global after_dot_j1 = split(j, ".")[2][1]
            global after_dot_j2 = split(j, ".")[2][2]
            global first_diff = abs(Int(after_dot_i1) - Int(after_dot_j1))
            global second_diff = abs(Int(after_dot_i2) - Int(after_dot_j2))
            global sum_diff = first_diff + second_diff
        catch
            global sum_diff = 0
        end

        if big_zone_i == big_zone_j && sum_diff != 1
            adj_zone_dist_matrix[i][j] = adj_zone_dist_matrix[i][j] * γ
        else
            adj_zone_dist_matrix[i][j] = adj_zone_dist_matrix[i][j]
        end
    end

    return adj_zone_dist_matrix
end


# function eval_metric(::Val{:mse}, pred::AbstractMatrix{T}, Y::AbstractVector{T}, α=0.0) where T <: AbstractFloat
#     eval = mean((pred .- Y) .^ 2)
#     return eval
# end
function eval_metric(::Val{:mse}, pred::Vector{SVector{1,T}}, Y::AbstractVector{T}, α) where T <: AbstractFloat
    eval = zero(T)
    @inbounds for i in 1:length(pred)
        eval += (pred[i][1] - Y[i]) ^ 2
    end
    eval /= length(pred)
    return eval
end

function eval_metric(::Val{:rmse}, pred::Vector{SVector{1,T}}, Y::AbstractVector{T}, α=0.0) where T <: AbstractFloat
    eval = zero(T)
    @inbounds for i in 1:length(pred)
        eval += (pred[i][1] - Y[i]) ^ 2
    end
    eval = sqrt(eval/length(pred))
    return eval
end

function eval_metric(::Val{:mae}, pred::Vector{SVector{1,T}}, Y::AbstractVector{T}, α=0.0) where T <: AbstractFloat
    eval = zero(T)
    @inbounds for i in 1:length(pred)
        eval += abs(pred[i][1] - Y[i])
    end
    eval /= length(pred)
    return eval
end

function eval_metric(::Val{:logloss}, pred::Vector{SVector{1,T}}, Y::AbstractVector{T}, α=0.0) where T <: AbstractFloat
    eval = zero(T)
    ϵ = T(2e-7)
    @inbounds for i in 1:length(pred)
        eval -= Y[i] * log(max(ϵ, sigmoid(pred[i][1]))) + (1 - Y[i]) * log(max(ϵ, 1 - sigmoid(pred[i][1])))
    end
    eval /= length(pred)
    return eval
end

function eval_metric(::Val{:mlogloss}, pred::Vector{SVector{L,T}}, Y::Vector{S}, α=0.0) where {T <: AbstractFloat, L, S <: Integer}
    eval = zero(T)
    ϵ = T(2e-7)
    @inbounds for i in 1:length(pred)
        # pred[i] = pred[i] .- maximum(pred[i])
        soft_pred = min.(1-ϵ, max.(ϵ, exp.(pred[i]) / sum(exp.(pred[i]))))
        p = pred[i] .- maximum(pred[i])
        soft_pred = min.(1-ϵ, max.(ϵ, exp.(p) / sum(exp.(p))))
        eval -= log(soft_pred[Y[i]])
    end
    eval /= length(Y)
    return eval
end

function eval_metric(::Val{:poisson}, pred::Vector{SVector{1,T}}, Y::AbstractVector{T}, α=0.0) where T <: AbstractFloat
    eval = zero(T)
    @inbounds for i in 1:length(pred)
        eval += exp(pred[i][1]) * (1 - Y[i]) + log(factorial(Y[i]))
    end
    eval /= length(pred)
    return eval
end

# gaussian
# pred[i][1] = μ
# pred[i][2] = log(σ)
function eval_metric(::Val{:gaussian}, pred::Vector{SVector{L,T}}, Y::AbstractVector{T}, α=0.0) where {L, T <: AbstractFloat}
    eval = zero(T)
    ϵ = T(2e-7)
    @inbounds for i in 1:length(pred)
        eval += pred[i][2] + (Y[i] - pred[i][1])^2 / (2*max(ϵ, exp(2*pred[i][2])))
    end
    eval /= length(Y)
    return eval
end

function eval_metric(::Val{:quantile}, pred::Vector{SVector{1,T}}, Y::AbstractVector{T}, α=0.0) where T <: AbstractFloat
    eval = zero(T)
    for i in 1:length(pred)
        eval += α * max(Y[i] - pred[i][1], zero(T)) + (1-α) * max(pred[i][1] - Y[i], zero(T))
    end
    eval /= length(pred)
    return eval
end

function gini_raw(labels::T, preds::S) where {T,S}
    if length(preds) < 2
        return 0.0
    end
    random = (1:length(preds)) ./ length(preds)
    l_sort = labels[sortperm(preds)]
    l_cum_w = cumsum(l_sort) ./ sum(labels)
    gini = sum(l_cum_w .- random)
    return gini
end

function gini_norm(labels::T, preds::S) where {T,S}
    if length(labels) < 2
        return 0.0
    end
    return gini_raw(labels, preds) / gini_raw(labels, labels)
end

function eval_metric(::Val{:gini}, pred::Vector{SVector{1,T}}, Y::AbstractVector{T}, α=0.0) where T <: AbstractFloat
    return -gini_norm(Y, pred)
end

function eval_metric(::Val{:gini}, pred::Vector{SVector{2,T}}, Y::AbstractVector{T}, α=0.0) where T <: AbstractFloat
    return -gini_norm(Y, pred)
end


using RNLA, Test, LinearAlgebra

@testset "rcur_test1" begin
    m = 100
    n = 50
    k = 4

    C = rand(m, k)
    U = rand(k, k)
    R = rand(k, n)
    A = C * U * R

    C1, U1, R1 = rcur(A)

    @test norm(C - C1) <= sqrt(eps())
    @test norm(U - U1) <= sqrt(eps())
    @test norm(R - R1) <= sqrt(eps())
end

@testset "rcur_test2" begin
end


@doc """
```math
LFST = \\frac{\\sum^q_{j=1} X^2_{j,T}}{\\sum^q_{j=1}}
```
""" ->
function lfst_test(X, g)
  q = size(X, 1)
  tmp = (collect(1:1:q)*π).^2
  cov_rw = 1./tmp
  cov_ha = ones(q).+(g*g*cov_rw)
  cov_ha_inv = 1./cov_ha
  chol_ha = sqrt(cov_ha_inv)
  Xa = X.*chol_ha
  lfst = sum(X.^2)/sum(Xa.^2)
  nsim = 50000
  Xsim = randn(q, nsim)
  Xasim = Xsim.*repmat(chol_ha, 1, nsim)
  nx = sum(Xsim.^2, 1)
  dx = sum(Xasim.^2, 1)
  lfst_sim = (nx./dx)'
  lfst_pvalue = mean(lfst_sim.>=lfst)
  return lfst, lfst_pvalue
end
function lfst_pv(X, g)
  q = size(X, 1)
  tmp = (collect(1:1:q)*π).^2
  cov_rw = 1./tmp
  cov_ha = ones(q).+(g*g*cov_rw)
  cov_ha_inv = 1./cov_ha
  chol_ha = sqrt(cov_ha_inv)
  Xa = X.*chol_ha
  lfst = sum(X.^2)/sum(Xa.^2)
  nsim = 50000
  Xsim = randn(q, nsim)
  Xasim = Xsim.*repmat(chol_ha, 1, nsim)
  nx = sum(Xsim.^2, 1)
  dx = sum(Xasim.^2, 1)
  lfst_sim = (nx./dx)'
  lfst_pvalue = mean(lfst_sim.>=lfst)
  return lfst, lfst_sim
end

function  mw_lfst_pv(x_t, itrend, param_m, param_t, nsim; stat=false)
# mwlfst_pv - rewrite of Muller and Watson (2015) GAUSS code. Calculates LFST
#   test pvalues.
# Syntax:  pv::Float64 = mwlfst_pv(x_t, itrend, param_m, param_t, nsim)
#
# Inputs:
#   x_t::Vector{Float64} - Discrete cosine transforms of data.
#   itrend::Int64 - 1 if x_t is demeaned, 2 if x_t is detrended.
#   param_m::Float64 - "g" parameter for demeaned alternative.
#   param_t::Float64 - "g" parameter for detrended alternative.
#   nsim::Int64 - Number of simulations to compute p-value
# Outputs:
#   pv::Float64 - p-value
#------------- BEGIN CODE --------------
t = 1000 # Sample Size Used for Asymptotic Approximation
q = size(x_t, 1)
# Compute Null and Alternative Covariance Matrices
if itrend == 1
  psi = psi_compute(t, q)
  cov_h0 = psi'psi
  g_t = param_m/t
  gt2 = g_t.^2
  tmp = mw_rw_cov(t)
  tmp = gt2*tmp
  tmp1 = diag(tmp) + ones(t, 1)
  cov_u = diagrv(tmp, tmp1)
  cov_ha = psi'cov_u*psi
end
# Compute Test Statistic @
  chol_cov_h0 = chol(cov_h0)
  a = inv(chol_cov_h0')
  x_t_trans = a*x_t # @ X transformed under Ho @
  cov_ha_trans = a*cov_ha*a' # @ Covariance under Ha Transformed @
#   @ Note ... LF under Ho is unity @
  lr = mw_marg_lf(x_t_trans, cov_ha_trans)
# @ Simulations for Pvalue @
  xsim = rand(Normal(), size(x_t,1), nsim)
  lr_sim = mw_marg_lf(xsim, cov_ha_trans)
  pv = mean(lr_sim .>(lr - 0.0001))
return lr, lr_sim
end

function diagrv(x, v)
  y = x - diagm(diag(x)) + diagm(v[:])
  return y
end

function mw_marg_lf(x::Vector{Float64}, cov::Array{Float64,2})
# This is a scaled version of the density in King and Hillier (1985) Equation (2) @
# -- Note x is ncov x n ... compute n marginal likelihood values @
# Compute Marginal Likelihood value for x with covariance matrix cov
  x = x./sqrt(sum(x.^2, 1))'
  m = size(x, 1)
  d = det(cov)
  # (gauss code) x = chol(invpd(cov))
  c = chol(Hermitian(inv(x'x)))
  #x = c*x
  tmp = sum(x, 1)
  tmp = d*(tmp.^m)
  mlf = sqrt(1.0./tmp)
  return mlf
end
function mw_marg_lf(x::Array{Float64,2}, cov::Array{Float64,2})
  nsim = size(x, 2)
  mlf_vec = Vector{Float64}(nsim)
  for isim ∈ 1:nsim
    mlf_vec[isim, :] = mw_marg_lf(x[:, isim], cov)
  end
return mlf_vec
end

function mw_rw_cov(t)
  sum_mat = zeros(t,t)
  for i ∈ 1:t
    sum_mat[i, 1:i] = ones(1, i)
  end
  cov = sum_mat*sum_mat'
  return cov
end


@diffusion_process FitzHughDiffusionAux begin
    :dimensions
    process --> 2
    wiener --> 1

    :parameters
    (ϵ, s, γ, β, σ, t, T) --> Float64
    (u, v) --> ℝ{2}

    :additional
    constdiff --> true
    linear --> true
end

function B(t, P::FitzHughDiffusionAux)
    @SMatrix [
        1/P.ϵ-3*P.v^2/P.ϵ   -1/P.ϵ;
        P.γ                 -1.0
    ]
end

function β(t, P::FitzHughDiffusionAux)
    ℝ{2}(
        P.s/P.ϵ+2*P.v^3/P.ϵ,
        P.β
    )
end

σ(t, P::FitzHughDiffusionAux) = ℝ{2}(0.0, P.σ)


@doc raw"""
    truncated_conjugate_gradient_descent(M, F, ∇F, x, η, H, Δ)

solve the trust-region subproblem

```math
\operatorname*{arg\,min}_{\eta  ∈  T_{x}M} m_{x}(\eta) = F(x) + \langle \nabla F(x), \eta \rangle_{x} + \frac{1}{2} \langle \operatorname{Hess}[F](\eta)_ {x}, \eta \rangle_{x}
```
```math
\text{s.t.} \; \langle \eta, \eta \rangle_{x} \leqq {\Delta}^2
```

with the Steihaug-Toint truncated conjugate-gradient method.
For a description of the algorithm and theorems offering convergence guarantees,
see the reference:

* P.-A. Absil, C.G. Baker, K.A. Gallivan,
    Trust-region methods on Riemannian manifolds, FoCM, 2007.
    doi: [10.1007/s10208-005-0179-9](https://doi.org/10.1007/s10208-005-0179-9)
* A. R. Conn, N. I. M. Gould, P. L. Toint, Trust-region methods, SIAM,
    MPS, 2000. doi: [10.1137/1.9780898719857](https://doi.org/10.1137/1.9780898719857)

# Input
* `M` – a manifold $\mathcal M$
* `F` – a cost function $F\colon\mathcal M\to\mathbb R$ to minimize
* `∇F` – the gradient $\nabla F\colon\mathcal M\to T\mathcal M$ of F
* `x` – a point on the manifold $x ∈ \mathcal M$
* `η` – an update tangential vector $\eta ∈ \mathcal{T_{x}M}$
* `H` – the hessian $H( \mathcal M, x, \xi)$ of F
* `Δ` – a trust-region radius

# Optional
* `preconditioner` – a preconditioner for the hessian H
* `θ` – 1+θ is the superlinear convergence target rate. The algorithm will
    terminate early if the residual was reduced by a power of 1+theta.
* `κ` – the linear convergence target rate: algorithm will terminate
    early if the residual was reduced by a factor of kappa.
* `useRandom` – set to true if the trust-region solve is to be initiated with a
    random tangent vector. If set to true, no preconditioner will be
    used. This option is set to true in some scenarios to escape saddle
    points, but is otherwise seldom activated.
* `stopping_criterion` – ([`StopWhenAny`](@ref), [`StopAfterIteration`](@ref),
    [`StopIfResidualIsReducedByFactor`](@ref), [`StopIfResidualIsReducedByPower`](@ref),
    [`StopWhenCurvatureIsNegative`](@ref), [`StopWhenTrustRegionIsExceeded`](@ref) )
    a functor inheriting from [`StoppingCriterion`](@ref) indicating when to stop,
    where for the default, the maximal number of iterations is set to the dimension of the
    manifold, the power factor is `θ`, the reduction factor is `κ`.
    .
* `return_options` – (`false`) – if actiavated, the extended result, i.e. the
    complete [`Options`](@ref) re returned. This can be used to access recorded values.
    If set to false (default) just the optimal value `xOpt` is returned

and the ones that are passed to [`decorate_options`](@ref) for decorators.

# Output
* `η` – an approximate solution of the trust-region subproblem in
    $\mathcal{T_{x}M}$.
OR
* `options` - the options returned by the solver (see `return_options`)

# see also
[`trust_regions`](@ref)
"""
function truncated_conjugate_gradient_descent(
    M::mT,
    F::TF,
    ∇F::TdF,
    x,
    η,
    H::Union{Function,Missing},
    Δ::Float64;
    preconditioner::Tprec=(M, x, ξ) -> ξ,
    θ::Float64=1.0,
    κ::Float64=0.1,
    useRandom::Bool=false,
    stopping_criterion::StoppingCriterion=StopWhenAny(
        StopAfterIteration(manifold_dimension(M)),
        StopIfResidualIsReducedByPower(
            sqrt(inner(
                M,
                x,
                ∇F(M, x) + (useRandom ? H(M, x, η) : zero_tangent_vector(M, x)),
                ∇F(M, x) + (useRandom ? H(M, x, η) : zero_tangent_vector(M, x)),
            )),
            θ,
        ),
        StopIfResidualIsReducedByFactor(
            sqrt(inner(
                M,
                x,
                ∇F(M, x) + (useRandom ? H(M, x, η) : zero_tangent_vector(M, x)),
                ∇F(M, x) + (useRandom ? H(M, x, η) : zero_tangent_vector(M, x)),
            )),
            κ,
        ),
        StopWhenTrustRegionIsExceeded(),
        StopWhenCurvatureIsNegative(),
    ),
    return_options=false,
    kwargs..., #collect rest
) where {mT<:Manifold,TF,TdF,Tprec}
    p = HessianProblem(M, F, ∇F, H, preconditioner)
    o = TruncatedConjugateGradientOptions(
        x,
        stopping_criterion,
        η,
        zero_tangent_vector(M, x),
        Δ,
        zero_tangent_vector(M, x),
        useRandom,
    )
    o = decorate_options(o; kwargs...)
    resultO = solve(p, o)
    if return_options
        resultO
    else
        return get_solver_result(resultO)
    end
end
function initialize_solver!(
    p::P, o::O
) where {P<:HessianProblem,O<:TruncatedConjugateGradientOptions}
    o.η = o.useRand ? o.η : zero_tangent_vector(p.M, o.x)
    Hη = o.useRand ? getHessian(p, o.x, o.η) : zero_tangent_vector(p.M, o.x)
    o.residual = get_gradient(p, o.x) + Hη
    # Initial search direction (we maintain -delta in memory, called mdelta, to
    # avoid a change of sign of the tangent vector.)
    return o.δ = o.useRand ? o.residual : get_preconditioner(p, o.x, o.residual)
    # If the Hessian or a linear Hessian approximation is in use, it is
    # theoretically guaranteed that the model value decreases strictly
    # with each iteration of tCG. Hence, there is no need to monitor the model
    # value. But, when a nonlinear Hessian approximation is used (such as the
    # built-in finite-difference approximation for example), the model may
    # increase. It is then important to terminate the tCG iterations and return
    # the previous (the best-so-far) iterate. The variable below will hold the
    # model value.
    # o.model_value = o.useRand ? 0 : inner(p.M,o.x,o.η,get_gradient(p,o.x)) + 0.5 * inner(p.M,o.x,o.η,Hη)
end
function step_solver!(
    p::P, o::O, iter
) where {P<:HessianProblem,O<:TruncatedConjugateGradientOptions}
    ηOld = o.η
    δOld = o.δ
    z = o.useRand ? o.residual : get_preconditioner(p, o.x, o.residual)
    # this is not correct, it needs to be the inverse of the preconditioner!
    zrOld = inner(p.M, o.x, z, o.residual)
    HηOld = getHessian(p, o.x, ηOld)
    # This call is the computationally expensive step.
    Hδ = getHessian(p, o.x, δOld)
    # Compute curvature (often called kappa).
    δHδ = inner(p.M, o.x, δOld, Hδ)
    # Note that if d_Hd == 0, we will exit at the next "if" anyway.
    α = zrOld / δHδ
    # <neweta,neweta>_P =
    # <eta,eta>_P + 2*alpha*<eta,delta>_P + alpha*alpha*<delta,delta>_P
    e_Pd = -inner(p.M, o.x, ηOld, o.useRand ? δOld : get_preconditioner(p, o.x, δOld)) # It must be clarified if it's negative or not
    d_Pd = inner(p.M, o.x, δOld, o.useRand ? δOld : get_preconditioner(p, o.x, δOld))
    e_Pe = inner(p.M, o.x, ηOld, o.useRand ? ηOld : get_preconditioner(p, o.x, ηOld))
    e_Pe_new = e_Pe + 2α * e_Pd + α^2 * d_Pd # vielleicht müssen doch die weiteren Optionen gespeichert werden
    # Check against negative curvature and trust-region radius violation.
    # If either condition triggers, we bail out.
    if δHδ <= 0 || e_Pe_new >= o.Δ^2
        tau = (-e_Pd + sqrt(e_Pd^2 + d_Pd * (o.Δ^2 - e_Pe))) / d_Pd
        o.η = ηOld - tau * (δOld)
    else
        # No negative curvature and eta_prop inside TR: accept it.
        o.η = ηOld - α * (δOld)
        # Verify that the model cost decreased in going from eta to new_eta. If
        # it did not (which can only occur if the Hessian approximation is
        # nonlinear or because of numerical errors), then we return the
        # previous eta (which necessarily is the best reached so far, according
        # to the model cost). Otherwise, we accept the new eta and go on.
        # -> Stopping Criterion
        old_model_value =
            inner(p.M, o.x, ηOld, get_gradient(p, o.x)) + 0.5 * inner(p.M, o.x, ηOld, HηOld)
        new_model_value =
            inner(p.M, o.x, o.η, get_gradient(p, o.x)) +
            0.5 * inner(p.M, o.x, o.η, getHessian(p, o.x, o.η))
        if new_model_value >= old_model_value
            o.η = ηOld
        end
    end
    # Update the residual.
    o.residual = o.residual - α * Hδ
    # Precondition the residual.
    # It's actually the inverse of the preconditioner in o.residual
    z = o.useRand ? o.residual : get_preconditioner(p, o.x, o.residual)
    # this is not correct, it needs to be the inverse of the preconditioner!
    # Save the old z'*r.
    # Compute new z'*r.
    zr = inner(p.M, o.x, z, o.residual)
    # Compute new search direction.
    β = zr / zrOld
    return o.δ = project(p.M, o.x, z + β * o.δ)
end
get_solver_result(o::O) where {O<:TruncatedConjugateGradientOptions} = o.η


module NewtonMethod

function newton_solver!(R, f!, ∂f∂p!, p, num_iters = 3, α = 1.0)
    # Array to store errors
    ε = zeros(num_iters);

    # Newton iteration loop
    for i in 1:num_iters
        # (Generalize to Jac-vec matrix-free method such as Krylov subspace for large problems.)

        # Compute residuals
        R    = f!(R, p)

        # Compute Jacobian
        ∂R∂p = ∂f∂p!(R, p)

        # Compute Newton step
        Δp   = ∂R∂p \ -R

        # Update state with relaxation factor
        p    = newton_update!(p, reshape(Δp, size(p)), α)

        # Error processing
        ε[i] = maximum(abs.(Δp))
        println("Newton step error L² norm: $(ε[i])")
    end
    R, p, ε
end


"""
    T_transform(Yc, Xc, Xp, x) (stationnary)

"""
function T_transform(Yc::ContinuousUnivariateDistribution,  # Distribution des données observées pour la période de calibration
                     Xc::ContinuousUnivariateDistribution,  # Distribution des données du modèle pour la période de calibration
                     Xp::ContinuousUnivariateDistribution,  # Distribution des données du modèle pour la période projetée
                     x::Real)

    F_Xp = cdf(Xp, x)  # Évaluation en x de la CDF des données du modèle pour la période projeté

    Fm1_Xc = quantile(Xc, F_Xp)  # Quantile correspondant de la distribution des données du modèle pour la période de calibration

    F_Yp = cdf(Yc,  Fm1_Xc)  # Évaluation de la CDF des données observées pour la période de calibration au point correspondant

    return F_Yp  # Valeur de la CDF évaluée au point x
end

"""
    T_transform(Yx, Xc, Xc_cov, Xp, Xp_cov, x)  (non-stationnary)

"""
function T_transform(Yc::EGPpower,  # Distribution des données observées pour la période de calibration
                     Xc::nonstatEGPpower,  # Distribution des données du modèle pour la période de calibration
                     Xc_cov::Real,  # covariable pour la période de calibration (ex : [CO₂] de l'année 1980)
                     Xp::nonstatEGPpower,  # Distribution des données du modèle pour la période projetée
                     Xp_cov::Real,  # covariable pour la période projetée (ex : [CO₂] de l'année 2050)
                     x::Real)

    F_Xp = cdf(Xp, x, Xp_cov)  # Évaluation en x de la CDF des données du modèle pour la période projeté

    Fm1_Xc = quantile(Xc, F_Xp, Xc_cov)  # Quantile correspondant de la distribution des données du modèle pour la période de calibration

    F_Yp = cdf(Yc,  Fm1_Xc)  # Évaluation de la CDF des données observées pour la période de calibration au point correspondant

    return F_Yp  # Valeur de la CDF évaluée au point x
end

## Bias correction

"""
    bias_correct(obs, ref, covariate_ref; fut, covariate_fut, censoring_obs, censoring_sim, wet_thresh)

"""
function bias_correct(obs::Array{Float64, 1},
                ref::Array{Float64, 1},
                covariate_ref::Array{Float64, 1};
                fut::Array{Float64, 1}=Float64[],
                covariate_fut::Array{Float64, 1}=Float64[],
                censoring_obs::Real=5.0,
                censoring_sim::Real=5.0,
                wet_thresh::Float64=1.0)

    # Transformation des valeurs en fonction du seuil
    obs_t = obs[obs .> wet_thresh] .- wet_thresh

    ref_t = ref[ref .> wet_thresh] .- wet_thresh
    covariate_ref_t = covariate_ref[ref .> wet_thresh]

    fut_t = fut[fut .> wet_thresh] .- wet_thresh
    covariate_fut_t = covariate_fut[fut .> wet_thresh]

    # Ajustement du modèle sur les observations
    fm_obs = EGPpowerfit(obs_t, censoring=censoring_obs)

    # Fusion des prédictions et des projections pour l'ajustement
    sim = vcat(ref_t, fut_t)
    covariate_sim = vcat(covariate_ref_t, covariate_fut_t)

    # Ajustement du modèle sur les simulations
    fm_sim, fms = EGPpowerfit(sim, censoring=censoring_sim, covariate=covariate_sim, scalecov=false, uppertailcov=true, lowertailcov=false)

    # Correction des prédictions
    output_ref = similar(ref)
    output_ref[ref .<= wet_thresh] .= 0.0
    output_ref[ref .> wet_thresh] = quantile.(fm_obs, cdf.(EGPpower.(fm_sim, covariate_ref_t), ref_t)) .+ wet_thresh

    if !isempty(fut)
        # Récupération de la CDF des obs pour le futur (Kallache)
        Yp = CDFT.(Ref(fm_obs), Ref(EGPpower(fm_sim, last(covariate_ref))), EGPpower.(fm_sim, covariate_fut_t))
        p = cdf.(getfield.(Yp, :Xp), fut_t)

        # Correction des projections
        output_fut = similar(fut)
        output_fut[fut .<= wet_thresh] .= 0.0
        output_fut[fut .> wet_thresh] = quantile.(Yp, p) .+ wet_thresh
        append!(output_ref, output_fut)

    end
    return output_ref
end

function bias_correct(obs::Array{Float64, 1},
                ref::Array{Float64, 1},
                covariate_ref::Array{Float64, 1},
                fut::Array{Float64, 1},
                covariate_fut::Array{Float64, 1};
                censoring_obs::Real=5.0,
                censoring_sim::Real=5.0,
                wet_thresh::Float64=1.0)

    return bias_correct(obs, ref, covariate_ref, fut=fut, covariate_fut=covariate_fut, censoring_obs=censoring_obs, censoring_sim=censoring_sim, wet_thresh=wet_thresh)
end

function bias_correct(obs::TimeArray{Float64,1,T,Array{Float64,1}},
                ref::TimeArray{Float64,1,T,Array{Float64,1}},
                covariate_ref::TimeArray{Float64,1,T,Array{Float64,1}},
                fut::TimeArray{Float64,1,T,Array{Float64,1}},
                covariate_fut::TimeArray{Float64,1,T,Array{Float64,1}};
                censoring_obs::Real=5.0,
                censoring_sim::Real=5.0,
                wet_thresh::Float64=1.0) where {T <: TimeType}

    x = bias_correct(values(obs), values(ref), values(covariate_ref), fut=values(fut), covariate_fut=values(covariate_fut), censoring_obs=censoring_obs,
      censoring_sim=censoring_sim, wet_thresh=wet_thresh)
    t = vcat(timestamp(ref), timestamp(fut))

    return TimeArray(t, x)
end

function bias_correct(obs::TimeArray{Float64,1,T,Array{Float64,1}},
                ref::TimeArray{Float64,1,T,Array{Float64,1}},
                covariate_ref::TimeArray{Float64,1,T,Array{Float64,1}};
                censoring_obs::Real=5.0,
                censoring_sim::Real=5.0,
                wet_thresh::Float64=1.0) where {T <: TimeType}

    x = bias_correct(values(obs), values(ref), values(covariate_ref), censoring_obs=censoring_obs, censoring_sim=censoring_sim, wet_thresh=wet_thresh)
    t = timestamp(ref)

    return TimeArray(t, x)
end


@testset verbose=true "Math" begin
    components = [:x, :y, :z]

    basic_functions = [conj, abs2, abs, inv, log, exp, sqrt]

    ### Test equivalence with complex math
    ℍ_to_ℂ(q::Real, component) = q + 0im
    function ℍ_to_ℂ(q::AbstractQuaternion, component)
        c = q.w + im * getproperty(q, component)
    end
    function ℂ_to_ℍ(c::Complex, component)
        Quaternion(c.re, [comp_i == component ? c.im : zero(real(c)) for comp_i in components]...)
    end
    @testset "Complex equivalence $T" for T in FloatTypes
        ϵ = 10eps(T)
        scalars = [zero(T), one(T), -one(T)]
        for c in [a+b*im for a in scalars for b in scalars]
            # The following constructs a quaternion from the complex number `c` by equating `im`
            # with one of the three quaternionic imaginaries, looping over each option.  We then
            # apply the quaternionic math function to that quaternion, and map back to ℂ; the result
            # should be identical to just applying the complex math function directly.  The only
            # exceptions are for `log` and `sqrt` when the complex number is a negative real number,
            # for which we have to arbitrarily choose the "vector" part of the quaternion, which is
            # chosen to be `imz` in the code.
            for component in components
                q = ℂ_to_ℍ(c, component)
                for unary_function in basic_functions
                    if unary_function ∈ [log, sqrt] && component!=:z && c.re < zero(T) && c.im == zero(T)
                        continue  # We arbitrarily chose the z component for these return values
                    end
                    cval = unary_function(c)
                    qval = ℍ_to_ℂ(unary_function(q), component)
                    @test cval ≈ qval rtol=ϵ nans=true
                    # Repeat the test on QuatVec for `exp` for pure-imaginary input
                    if unary_function ∈ [exp] && iszero(c.re)
                        qval = ℍ_to_ℂ(unary_function(QuatVec(q)), component)
                        @test cval ≈ qval rtol=ϵ nans=true
                    end
                    # Repeat the test on Rotor for `log` and `sqrt` for nonzero input
                    if unary_function ∈ [log, sqrt] && !iszero(c)
                        cval = unary_function(c/abs(c))
                        qval = ℍ_to_ℂ(unary_function(Rotor(q)), component)
                        @test cval ≈ qval rtol=ϵ nans=true
                    end
                end
                for unary_function in [angle]  # Not identical, but related
                    cval = unary_function(c)
                    qval = ℍ_to_ℂ(unary_function(q), component)
                    @test 2*abs(cval) ≈ qval rtol=ϵ nans=true
                    if !iszero(c)
                        cval = unary_function(c/abs(c))
                        qval = ℍ_to_ℂ(unary_function(Rotor(q)), component)
                        @test 2*abs(cval) ≈ qval rtol=ϵ nans=true
                    end
                end
                for binary_function in [(^)]
                    if binary_function ∈ [(^)] && ((component!=:z && c.re < zero(T) && c.im == zero(T)) || iszero(c))
                        continue  # We arbitrarily chose the z component for these return values
                    end
                    if T !== BigFloat
                        S = 10randn(T, 100)
                        for s ∈ S
                            cval = binary_function(c, float(s))
                            qval = ℍ_to_ℂ(binary_function(q, s), component)
                            @test cval ≈ qval rtol=6ϵ nans=true
                            if !iszero(c)
                                cval = binary_function(c/abs(c), s)
                                qval = ℍ_to_ℂ(binary_function(Rotor(q), s), component)
                                @test cval ≈ qval rtol=6ϵ nans=true
                            end
                        end
                    end
                    S = [-3, -2, -1, 0, 1, 2, 3]
                    for s ∈ S
                        cval = binary_function(c, s)
                        qval = ℍ_to_ℂ(binary_function(q, s), component)
                        @test cval ≈ qval rtol=ϵ nans=true
                        if iszero(c.re)
                            qval = ℍ_to_ℂ(binary_function(QuatVec(q), s), component)
                            @test cval ≈ qval rtol=ϵ nans=true
                        end
                        if !iszero(c)
                            cval = binary_function(c/abs(c), s)
                            qval = ℍ_to_ℂ(binary_function(Rotor(q), s), component)
                            @test cval ≈ qval rtol=ϵ nans=true
                        end
                    end
                end

                @test exp(log(q)) ≈ q rtol=ϵ nans=true
                @test exp(zero(q)) ≈ one(q) rtol=ϵ nans=true
                @test exp(log(Rotor(q))) ≈ Rotor(q) rtol=ϵ nans=true
                @test exp(zero(QuatVec(q))) ≈ one(q) rtol=ϵ nans=true

            end
        end
    end

end


using CSV
using Plots

include("mlp.jl")

using .MLP

data = CSV.File("data/breast-cancer-wisconsin.data", header=false)
data = broadcast(collect.(data)) do x
    if x[7] == "?"
        x[7] = 0
    else
        x[7] = parse(Int, x[7])
    end
    x
end
data = collect(hcat(convert(Array{Array{Int64, 1}, 1}, data)...)')

d_tra = Dict(
    2 => [0, 1],
    4 => [1, 0],
)

data_X = data[:, 2:end-1] / 10
data_Y = collect(hcat(broadcast(data[:, end]) do x
    d_tra[x]
end...)')

data_X = Float64.(data_X)
data_Y = Float64.(data_Y)

let quz = "question-2", n_hidden_layers=5, n_epochs=800, η=0.9
    mlp = SimpleMLP{Float64}(data_X, data_Y, n_hidden_layers=n_hidden_layers, n_epochs=n_epochs, η=η)
    fit!(mlp)
    save_figures(mlp, "figs/$(quz)_")
    println("$(quz) done.")
end

let quz = "question-3", n_hidden_layers=5, η=0.9
    mlp = SimpleMLP{Float64}(data_X, data_Y, n_hidden_layers=n_hidden_layers, n_epochs=800, η=η)
    fit!(mlp)
    mlp_10 = SimpleMLP{Float64}(data_X, data_Y, n_hidden_layers=n_hidden_layers, n_epochs=10, η=η)
    fit!(mlp_10)
    mlp_100_000 = SimpleMLP{Float64}(data_X, data_Y, n_hidden_layers=n_hidden_layers, n_epochs=100_000, η=η)
    fit!(mlp_100_000)

    plt_1 = plot(
        mlp.loss_history,
        label="n_epochs=800, last=$(mlp.loss_history[end])",
        title="loss",
        legend=:outertopright,
        size=(1024, 512),
    )
    plot!(plt_1, mlp_10.loss_history, label="n_epochs=10, last=$(mlp_10.loss_history[end])")
    plot!(plt_1, mlp_100_000.loss_history, label="n_epochs=10,000, last=$(mlp_100_000.loss_history[end])")
    savefig(plt_1, "figs/$(quz)_loss_fig.png")

    plt_2 = plot(
        [mlp.train_accuracy_history mlp.test_accuracy_history],
        label=["trainset, n_epochs=800, last=$(mlp.train_accuracy_history[end])" "testset, n_epochs=800, last=$(mlp.test_accuracy_history[end])"],
        title="accuracy",
        legend=:outertopright,
        size=(1024, 512),
    )
    plot!(
        plt_2,
        [mlp_10.train_accuracy_history mlp_10.test_accuracy_history],
        label=["trainset, n_epochs=10, last=$(mlp_10.train_accuracy_history[end])" "testset, n_epochs=10, last=$(mlp_10.test_accuracy_history[end])"],
    )
    plot!(
        plt_2,
        [mlp_100_000.train_accuracy_history mlp_100_000.test_accuracy_history],
        label=["trainset, n_epochs=100,000, last=$(mlp_100_000.train_accuracy_history[end])" "testset, n_epochs=100,000, last=$(mlp_100_000.test_accuracy_history[end])"],
    )
    savefig(plt_2, "figs/$(quz)_accuracy_fig.png")
    println("$(quz) done.")
end

let quz = "question-4", n_epochs=800, η=0.9
    mlp = SimpleMLP{Float64}(data_X, data_Y, n_hidden_layers=5, n_epochs=n_epochs, η=η)
    fit!(mlp)
    mlp_2 = SimpleMLP{Float64}(data_X, data_Y, n_hidden_layers=2, n_epochs=n_epochs, η=η)
    fit!(mlp_2)
    mlp_30 = SimpleMLP{Float64}(data_X, data_Y, n_hidden_layers=30, n_epochs=n_epochs, η=η)
    fit!(mlp_30)

    plt_1 = plot(
        mlp.loss_history,
        label="n_hidden_layers=5, last=$(mlp.loss_history[end])",
        title="loss",
        legend=:outertopright,
        size=(1024, 512),
    )
    plot!(plt_1, mlp_2.loss_history, label="n_hidden_layers=2, last=$(mlp_2.loss_history[end])")
    plot!(plt_1, mlp_30.loss_history, label="n_hidden_layers=30, last=$(mlp_30.loss_history[end])")
    savefig(plt_1, "figs/$(quz)_loss_fig.png")

    plt_2 = plot(
        [mlp.train_accuracy_history mlp.test_accuracy_history],
        label=["trainset, n_hidden_layers=5, last=$(mlp.train_accuracy_history[end])" "testset, n_hidden_layers=5, last=$(mlp.test_accuracy_history[end])"],
        title="accuracy",
        legend=:outertopright,
        size=(1024, 512),
    )
    plot!(
        plt_2,
        [mlp_2.train_accuracy_history mlp_2.test_accuracy_history],
        label=["trainset, n_hidden_layers=2, last=$(mlp_2.train_accuracy_history[end])" "testset, n_hidden_layers=2, last=$(mlp_2.test_accuracy_history[end])"],
    )
    plot!(
        plt_2,
        [mlp_30.train_accuracy_history mlp_30.test_accuracy_history],
        label=["trainset, n_hidden_layers=30, last=$(mlp_30.train_accuracy_history[end])" "testset, n_hidden_layers=30, last=$(mlp_30.test_accuracy_history[end])"],
    )
    savefig(plt_2, "figs/$(quz)_accuracy_fig.png")
    println("$(quz) done.")
end

let quz = "question-5", n_tries=100, n_hidden_layers=5, n_epochs=800, η=0.9
    mlp_l = SimpleMLP{Float64}[]
    for t ∈ 1:n_tries
        mlp = SimpleMLP{Float64}(data_X, data_Y, n_hidden_layers=n_hidden_layers, n_epochs=n_epochs, η=η)
        fit!(mlp)
        push!(mlp_l, mlp)
    end
    w_mlp = mlp_l[argmax(map(x -> x.loss_history[end], mlp_l))]
    save_figures(w_mlp, "figs/$(quz)_")
    println("$(quz) done.")
end

let quz = "question-6", n_hidden_layers=5, n_epochs=800, η=0.9
    mlp = SimpleMLP{Float64}(data_X, data_Y, n_hidden_layers=n_hidden_layers, n_epochs=n_epochs, η=η)
    fit!(mlp)
    mlp_a = SimpleMLP{Float64}(data_X, data_Y, n_hidden_layers=n_hidden_layers, n_epochs=n_epochs, η=η)
    adaptive_fit!(mlp_a)

    plt_1 = plot(
        mlp.loss_history,
        label="Gradient Descent, last=$(mlp.loss_history[end])",
        title="loss",
        legend=:outertopright,
        size=(1024, 512),
    )
    plot!(plt_1, mlp_a.loss_history, label="AdaGrad, last=$(mlp_a.loss_history[end])")
    savefig(plt_1, "figs/$(quz)_loss_fig.png")

    plt_2 = plot(
        [mlp.train_accuracy_history mlp.test_accuracy_history],
        label=["trainset, Gradient Descent, last=$(mlp.train_accuracy_history[end])" "testset, Gradient Descent, last=$(mlp.test_accuracy_history[end])"],
        title="accuracy",
        legend=:outertopright,
        size=(1024, 512),
    )
    plot!(
        plt_2,
        [mlp_a.train_accuracy_history mlp_a.test_accuracy_history],
        label=["trainset, AdaGrad, last=$(mlp_a.train_accuracy_history[end])" "testset, AdaGrad, last=$(mlp_a.test_accuracy_history[end])"],
    )
    savefig(plt_2, "figs/$(quz)_accuracy_fig.png")
    println("$(quz) done.")
end