"""
    Perceptron.jl file

Implement the BetaML.Perceptron module

`?BetaML.Perceptron` for documentation

- [Importable source code (most up-to-date version)](https://github.com/sylvaticus/BetaML.jl/blob/master/src/Perceptron.jl) - [Julia Package](https://github.com/sylvaticus/BetaML.jl)
- [Demonstrative static notebook](https://github.com/sylvaticus/lmlj.jl/blob/master/notebooks/Perceptron.ipynb)
- [Demonstrative live notebook](https://mybinder.org/v2/gh/sylvaticus/BetaML.jl/master?filepath=notebooks%2FPerceptron.ipynb) (temporary personal online computational environment on myBinder) - it can takes minutes to start with!
- Theory based on [MITx 6.86x - Machine Learning with Python: from Linear Models to Deep Learning](https://github.com/sylvaticus/MITx_6.86x) ([Unit 3](https://github.com/sylvaticus/MITx_6.86x/blob/master/Unit%2003%20-%20Neural%20networks/Unit%2003%20-%20Neural%20networks.md))
- New to Julia? [A concise Julia tutorial](https://github.com/sylvaticus/juliatutorial) - [Julia Quick Syntax Reference book](https://julia-book.com)

"""

"""
    Perceptron module

Provide linear and kernel classifiers.

See a [runnable example on myBinder](https://mybinder.org/v2/gh/sylvaticus/BetaML.jl/master?filepath=notebooks%2FPerceptron.ipynb)

- [`perceptron`](@ref): Train data using the classical perceptron
- [`kernelPerceptron`](@ref): Train data using the kernel perceptron
- [`pegasos`](@ref): Train data using the pegasos algorithm
- [`predict`](@ref): Predict data using parameters from one of the above algorithms

All algorithms are multiclass, with `perceptron` and `pegasos` employing a one-vs-all strategy, while `kernelPerceptron` employs a _one-vs-one_ approach, and return a "probability" for each class in term of a dictionary for each record. Use `mode(ŷ)` to return a single class prediction per record.

The binary equivalent algorithms, accepting only `{-1,+1}` labels, are available as `peceptronBinary`, `kernelPerceptronBinary` and `pegasosBinary`. They are slighly faster as they don't need to be wrapped in the multi-class equivalent and return a more informative output.

The multi-class versions are available in the MLJ framework as `PerceptronClassifier`,`KernelPerceptronClassifier` and `PegasosClassifier` respectivly.
"""
module Perceptron

using LinearAlgebra, Random, ProgressMeter, Reexport, CategoricalArrays

using ForceImport
@force using ..Api
@force using ..Utils

export perceptron, perceptronBinary, kernelPerceptron, kernelPerceptronBinary, pegasos, pegasosBinary, predict

# Todo (breaking): change the API so that a struct kernelPerceptronModel or linearModel is reported
# instead of a named tuple
# linearModel <: betaSupervisedModel <: betaModel

"""
    perceptron(x,y;θ,θ₀,T,nMsgs,shuffle,forceOrigin,returnMeanHyperplane)

Train the multiclass classifier "perceptron" algorithm  based on x and y (labels).

The perceptron is a _linear_ classifier. Multiclass is supported using a one-vs-all approach.

# Parameters:
* `x`:           Feature matrix of the training data (n × d)
* `y`:           Associated labels of the training data, can be in any format (string, integers..)
* `θ`:           Initial value of the weights (parameter) [def: `zeros(d)`]
* `θ₀`:          Initial value of the weight (parameter) associated to the constant
                 term [def: `0`]
* `T`:           Maximum number of iterations across the whole set (if the set
                 is not fully classified earlier) [def: 1000]
* `nMsg`:        Maximum number of messages to show if all iterations are done [def: `0`]
* `shuffle`:     Whether to randomly shuffle the data at each iteration [def: `false`]
* `forceOrigin`: Whether to force `θ₀` to remain zero [def: `false`]
* `returnMeanHyperplane`: Whether to return the average hyperplane coefficients instead of the final ones  [def: `false`]
* `rng`:         Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]

# Return a named tuple with:
* `θ`:          The weights of the classifier
* `θ₀`:         The weight of the classifier associated to the constant term
* `classes`:    The classes (unique values) of y

# Notes:
* The trained parameters can then be used to make predictions using the function `predict()`.
* This model is available in the MLJ framework as the `PerceptronClassifier`

# Example:
```jldoctest
julia> model = perceptron([1.1 2.1; 5.3 4.2; 1.8 1.7], [-1,1,-1])
julia> ŷ     = predict([2.1 3.1; 7.3 5.2], model.θ, model.θ₀, model.classes)
```
"""
function perceptron(x::AbstractMatrix, y::AbstractVector; θ=zeros(size(x,2)),θ₀=0.0, T=1000, nMsgs=0, shuffle=false, forceOrigin=false, returnMeanHyperplane=false, rng = Random.GLOBAL_RNG)
    yclasses = unique(y)
    nCl      = length(yclasses)
    if nCl == 2
        outθ        = Array{Vector{Float64},1}(undef,1)
        outθ₀       = Array{Float64,1}(undef,1)
    else
        outθ        = Array{Vector{Float64},1}(undef,nCl)
        outθ₀       = Array{Float64,1}(undef,nCl)
    end
    for (i,c) in enumerate(yclasses)
        ybin = ((y .== c) .*2 .-1)  # conversion to -1/+1
        outBinary = perceptronBinary(x, ybin; θ=θ,θ₀=θ₀, T=T, nMsgs=nMsgs, shuffle=shuffle, forceOrigin=forceOrigin, rng=rng)
        if returnMeanHyperplane
            outθ[i]  = outBinary.avgθ
            outθ₀[i] = outBinary.avgθ₀
        else
            outθ[i]  = outBinary.θ
            outθ₀[i] = outBinary.θ₀
        end
        if nCl == 2
            break    # if there are only two classes we do compute only one passage, as A vs B would be the same as B vs A
        end
    end
    return (θ=outθ,θ₀=outθ₀,classes=yclasses)
end

"""
    perceptronBinary(x,y;θ,θ₀,T,nMsgs,shuffle,forceOrigin)

Train the binary classifier "perceptron" algorithm based on x and y (labels)

# Parameters:
* `x`:           Feature matrix of the training data (n × d)
* `y`:           Associated labels of the training data, in the format of ⨦ 1
* `θ`:           Initial value of the weights (parameter) [def: `zeros(d)`]
* `θ₀`:          Initial value of the weight (parameter) associated to the constant
                 term [def: `0`]
* `T`:           Maximum number of iterations across the whole set (if the set
                 is not fully classified earlier) [def: 1000]
* `nMsg`:        Maximum number of messages to show if all iterations are done
* `shuffle`:     Whether to randomly shuffle the data at each iteration [def: `false`]
* `forceOrigin`: Whether to force `θ₀` to remain zero [def: `false`]
* `rng`:         Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]

# Return a named tuple with:
* `θ`:          The final weights of the classifier
* `θ₀`:         The final weight of the classifier associated to the constant term
* `avgθ`:       The average weights of the classifier
* `avgθ₀`:      The average weight of the classifier associated to the constant term
* `errors`:     The number of errors in the last iteration
* `besterrors`: The minimum number of errors in classifying the data ever reached
* `iterations`: The actual number of iterations performed
* `separated`:  Weather the data has been successfully separated

# Notes:
* The trained parameters can then be used to make predictions using the function `predict()`.

# Example:
```jldoctest
julia> model = perceptronBinary([1.1 2.1; 5.3 4.2; 1.8 1.7], [-1,1,-1])
```
"""
function perceptronBinary(x, y; θ=zeros(size(x,2)),θ₀=0.0, T=1000, nMsgs=10, shuffle=false, forceOrigin=false, rng = Random.GLOBAL_RNG)
   if nMsgs != 0
       @codeLocation
       println("***\n*** Training perceptron for maximum $T iterations. Random shuffle: $shuffle")
   end
   x = makeMatrix(x)
   (n,d) = size(x)
   bestϵ = Inf
   lastϵ = Inf
   if forceOrigin θ₀ = 0.0; end
   sumθ = θ; sumθ₀ = θ₀
   @showprogress 1 "Training Perceptron..." for t in 1:T
       ϵ = 0
       if shuffle
          # random shuffle x and y
          ridx = Random.shuffle(rng, 1:size(x)[1])
          x = x[ridx, :]
          y = y[ridx]
       end
       for i in 1:n
           if y[i]*(θ' * x[i,:] + θ₀) <= eps()
               θ  = θ + y[i] * x[i,:]
               θ₀ = forceOrigin ? 0.0 : θ₀ + y[i]
               sumθ += θ; sumθ₀ += θ₀
               ϵ += 1
           end
       end
       if (ϵ == 0)
           if nMsgs != 0
               println("*** Avg. error after epoch $t : $(ϵ/size(x)[1]) (all elements of the set has been correctly classified")
           end
           return (θ=θ,θ₀=θ₀,avgθ=sumθ/(n*T),avgθ₀=sumθ₀/(n*T),errors=0,besterrors=0,iterations=t,separated=true)
       elseif ϵ < bestϵ
           bestϵ = ϵ
       end
       lastϵ = ϵ
       if nMsgs != 0 && (t % ceil(T/nMsgs) == 0 || t == 1 || t == T)
         println("Avg. error after iteration $t : $(ϵ/size(x)[1])")
       end
   end
   return  (θ=θ,θ₀=θ₀,avgθ=sumθ/(n*T),avgθ₀=sumθ₀/(n*T),errors=lastϵ,besterrors=bestϵ,iterations=T,separated=false)
end

"""
   kernelPerceptron(x,y;K,T,α,nMsgs,shuffle)

Train a multiclass kernel classifier "perceptron" algorithm based on x and y.

`kernelPerceptron` is a (potentially) non-linear perceptron-style classifier employing user-defined kernel funcions. Multiclass is supported using a one-vs-one approach.

# Parameters:
* `x`:        Feature matrix of the training data (n × d)
* `y`:        Associated labels of the training data, in the format of ⨦ 1
* `K`:        Kernel function to employ. See `?radialKernel` or `?polynomialKernel`for details or check `?BetaML.Utils` to verify if other kernels are defined (you can alsways define your own kernel) [def: [`radialKernel`](@ref)]
* `T`:        Maximum number of iterations (aka "epochs") across the whole set (if the set is not fully classified earlier) [def: 100]
* `α`:        Initial distribution of the errors [def: `zeros(length(y))`]
* `nMsg`:     Maximum number of messages to show if all iterations are done [def: `0`]
* `shuffle`:  Whether to randomly shuffle the data at each iteration [def: `false`]
* `rng`:      Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]

# Return a named tuple with:
* `x`: The x data (eventually shuffled if `shuffle=true`)
* `y`: The label
* `α`: The errors associated to each record
* `classes`: The labels classes encountered in the training

# Notes:
* The trained model can then be used to make predictions using the function `predict()`.
* This model is available in the MLJ framework as the `KernelPerceptronClassifier`

# Example:
```jldoctest
julia> model  = kernelPerceptron([1.1 2.1; 5.3 4.2; 1.8 1.7], [-1,1,-1])
julia> ŷtrain = Perceptron.predict(xtrain,model.x,model.y,model.α, model.classes,K=model.K)
julia> ϵtrain = error(ytrain, mode(ŷtrain))
```
"""
function kernelPerceptron(x, y; K=radialKernel, T=100, α=zeros(Int64,length(y)), nMsgs=0, shuffle=false, rng = Random.GLOBAL_RNG)
    x         = makeMatrix(x)
    yclasses  = unique(y)
    nCl       = length(yclasses)
    nModels   = Int((nCl  * (nCl - 1)) / 2)
    (n,d) = size(x)
    outX = Array{typeof(x),1}(undef,nModels)
    outY = Array{Array{Int64,1},1}(undef,nModels)
    outα = Array{Array{Int64,1},1}(undef,nModels)

    modelCounter = 1
    for (i,c) in enumerate(yclasses)
        for (i2,c2) in enumerate(yclasses)
            if i2 <= i continue end
            ids = ( (y .== c) .| (y .== c2) )
            thisx = x[ids,:]
            thisy = y[ids]
            thisα = α[ids]
            ybin = ((thisy .== c) .*2 .-1)  # conversion to +1 (if c) or -1 (if c2)
            outBinary = kernelPerceptronBinary(thisx, ybin; K=K, T=T, α=thisα, nMsgs=nMsgs, shuffle=shuffle, rng = rng)
            outX[modelCounter] = outBinary.x
            outY[modelCounter] = outBinary.y
            outα[modelCounter] = outBinary.α
            modelCounter += 1
        end
    end
    return (x=outX,y=outY,α=outα,classes=yclasses,K=K)
end

"""
   kernelPerceptronBinary(x,y;K,T,α,nMsgs,shuffle)

Train a multiclass kernel classifier "perceptron" algorithm based on x and y

# Parameters:
* `x`:        Feature matrix of the training data (n × d)
* `y`:        Associated labels of the training data, in the format of ⨦ 1
* `K`:        Kernel function to employ. See `?radialKernel` or `?polynomialKernel`for details or check `?BetaML.Utils` to verify if other kernels are defined (you can alsways define your own kernel) [def: [`radialKernel`](@ref)]
* `T`:        Maximum number of iterations across the whole set (if the set is not fully classified earlier) [def: 1000]
* `α`:        Initial distribution of the errors [def: `zeros(length(y))`]
* `nMsg`:     Maximum number of messages to show if all iterations are done
* `shuffle`:  Whether to randomly shuffle the data at each iteration [def: `false`]
* `rng`:      Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]

# Return a named tuple with:
* `x`: the x data (eventually shuffled if `shuffle=true`)
* `y`: the label
* `α`: the errors associated to each record
* `errors`: the number of errors in the last iteration
* `besterrors`: the minimum number of errors in classifying the data ever reached
* `iterations`: the actual number of iterations performed
* `separated`: a flag if the data has been successfully separated

# Notes:
* The trained data can then be used to make predictions using the function `predict()`. **If the option `shuffle` has been used, it is important to use there the returned (x,y,α) as these would have been shuffle compared with the original (x,y)**.
* Please see @kernelPerceptron for a multi-class version
# Example:
```jldoctest
julia> model = kernelPerceptronBinary([1.1 2.1; 5.3 4.2; 1.8 1.7], [-1,1,-1])
```
"""
function kernelPerceptronBinary(x, y; K=radialKernel, T=1000, α=zeros(Int64,length(y)), nMsgs=10, shuffle=false, rng = Random.GLOBAL_RNG)
    if nMsgs != 0
        @codeLocation
        println("***\n*** Training kernel perceptron for maximum $T iterations. Random shuffle: $shuffle")
    end
    x = makeMatrix(x)
    (n,d) = size(x)
    bestϵ = Inf
    lastϵ = Inf
    @showprogress 1 "Training Kernel Perceptron..." for t in 1:T
        ϵ = 0
        if shuffle
           # random shuffle x, y and alpha
           ridx = Random.shuffle(rng, 1:size(x)[1])
           x = x[ridx, :]
           y = y[ridx]
           α = α[ridx]
        end
        for i in 1:n
            if y[i]*sum([α[j]*y[j]*K(x[j,:],x[i,:]) for j in 1:n]) <= 0 + eps()
                α[i] += 1
                ϵ += 1
            end
        end
        if (ϵ == 0)
            if nMsgs != 0
                println("*** Avg. error after epoch $t : $(ϵ/size(x)[1]) (all elements of the set has been correctly classified")
            end
            return (x=x,y=y,α=α,errors=0,besterrors=0,iterations=t,separated=true,K=K)
        elseif ϵ < bestϵ
            bestϵ = ϵ
        end
        lastϵ = ϵ
        if nMsgs != 0 && (t % ceil(T/nMsgs) == 0 || t == 1 || t == T)
          println("Avg. error after iteration $t : $(ϵ/size(x)[1])")
        end
    end
    return  (x=x,y=y,α=α,errors=lastϵ,besterrors=bestϵ,iterations=T,separated=false)
end

"""
    pegasos(x,y;θ,θ₀,λ,η,T,nMsgs,shuffle,forceOrigin,returnMeanHyperplane)

Train the multiclass classifier "pegasos" algorithm according to x (features) and y (labels)

Pegasos is a _linear_, gradient-based classifier. Multiclass is supported using a one-vs-all approach.

# Parameters:
* `x`:           Feature matrix of the training data (n × d)
* `y`:           Associated labels of the training data, can be in any format (string, integers..)
* `θ`:           Initial value of the weights (parameter) [def: `zeros(d)`]
* `θ₀`:          Initial value of the weight (parameter) associated to the constant term [def: `0`]
* `λ`:           Multiplicative term of the learning rate
* `η`:           Learning rate [def: (t -> 1/sqrt(t))]
* `T`:           Maximum number of iterations across the whole set (if the set is not fully classified earlier) [def: 1000]
* `nMsg`:        Maximum number of messages to show if all iterations are done
* `shuffle`:     Whether to randomly shuffle the data at each iteration [def: `false`]
* `forceOrigin`: Whehter to force `θ₀` to remain zero [def: `false`]
* `returnMeanHyperplane`: Whether to return the average hyperplane coefficients instead of the average ones  [def: `false`]
* `rng`:         Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]

# Return a named tuple with:
* `θ`:          The weights of the classifier
* `θ₀`:         The weight of the classifier associated to the constant term
* `classes`:    The classes (unique values) of y

# Notes:
* The trained parameters can then be used to make predictions using the function `predict()`.
* This model is available in the MLJ framework as the `PegasosClassifier`

# Example:
```jldoctest
julia> model = pegasos([1.1 2.1; 5.3 4.2; 1.8 1.7], [-1,1,-1])
julia> ŷ     = predict([2.1 3.1; 7.3 5.2], model.θ, model.θ₀, model.classes)
```
"""
function pegasos(x, y; θ=zeros(size(x,2)),θ₀=0.0, λ=0.5,η= (t -> 1/sqrt(t)), T=1000, nMsgs=0, shuffle=false, forceOrigin=false,returnMeanHyperplane=false, rng = Random.GLOBAL_RNG)
    yclasses = unique(y)
    nCl      = length(yclasses)
    if nCl == 2
        outθ        = Array{Vector{Float64},1}(undef,1)
        outθ₀       = Array{Float64,1}(undef,1)
    else
        outθ        = Array{Vector{Float64},1}(undef,nCl)
        outθ₀       = Array{Float64,1}(undef,nCl)
    end
    for (i,c) in enumerate(yclasses)
        ybin = ((y .== c) .*2 .-1)  # conversion to -1/+1
        outBinary = pegasosBinary(x, ybin; θ=θ,θ₀=θ₀, λ=λ,η=η, T=T, nMsgs=nMsgs, shuffle=shuffle, forceOrigin=forceOrigin, rng=rng)
        if returnMeanHyperplane
            outθ[i]  = outBinary.avgθ
            outθ₀[i] = outBinary.avgθ₀
        else
            outθ[i]  = outBinary.θ
            outθ₀[i] = outBinary.θ₀
        end
        if nCl == 2
            break    # if there are only two classes we do compute only one passage, as A vs B would be the same as B vs A
        end
    end
    return (θ=outθ,θ₀=outθ₀,classes=yclasses)
end

"""
    pegasosBinary(x,y;θ,θ₀,λ,η,T,nMsgs,shuffle,forceOrigin)

Train the peagasos algorithm based on x and y (labels)

# Parameters:
* `x`:           Feature matrix of the training data (n × d)
* `y`:           Associated labels of the training data, in the format of ⨦ 1
* `θ`:           Initial value of the weights (parameter) [def: `zeros(d)`]
* `θ₀`:          Initial value of the weight (parameter) associated to the constant term [def: `0`]
* `λ`:           Multiplicative term of the learning rate
* `η`:           Learning rate [def: (t -> 1/sqrt(t))]
* `T`:           Maximum number of iterations across the whole set (if the set is not fully classified earlier) [def: 1000]
* `nMsg`:        Maximum number of messages to show if all iterations are done
* `shuffle`:    Whether to randomly shuffle the data at each iteration [def: `false`]
* `forceOrigin`: Whether to force `θ₀` to remain zero [def: `false`]

# Return a named tuple with:
* `θ`:          The final weights of the classifier
* `θ₀`:         The final weight of the classifier associated to the constant term
* `avgθ`:       The average weights of the classifier
* `avgθ₀`:      The average weight of the classifier associated to the constant term
* `errors`:     The number of errors in the last iteration
* `besterrors`: The minimum number of errors in classifying the data ever reached
* `iterations`: The actual number of iterations performed
* `separated`:  Weather the data has been successfully separated

# Notes:
* The trained parameters can then be used to make predictions using the function `predict()`.

# Example:
```jldoctest
julia> pegasos([1.1 2.1; 5.3 4.2; 1.8 1.7], [-1,1,-1])
```
"""
function pegasosBinary(x, y; θ=zeros(size(x,2)),θ₀=0.0, λ=0.5,η= (t -> 1/sqrt(t)), T=1000, nMsgs=10, shuffle=false, forceOrigin=false, rng = Random.GLOBAL_RNG)
  if nMsgs != 0
      @codeLocation
      println("***\n*** Training pegasos for maximum $T iterations. Random shuffle: $shuffle")
  end
  x = makeMatrix(x)
  (n,d) = size(x)
  bestϵ = Inf
  lastϵ = Inf
  if forceOrigin θ₀ = 0.0; end
  sumθ = θ; sumθ₀ = θ₀
  @showprogress 1 "Training Pegasos..." for t in 1:T
      ϵ = 0
      ηₜ = η(t)
      if shuffle
         # random shuffle x and y
         ridx = Base.shuffle(rng, 1:size(x)[1])
         x = x[ridx, :]
         y = y[ridx]
      end
      for i in 1:n
          if y[i]*(θ' * x[i,:] + θ₀) <= eps()
              θ  = (1-ηₜ*λ) * θ + ηₜ * y[i] * x[i,:]
              θ₀ = forceOrigin ? 0.0 : θ₀ + ηₜ * y[i]
              sumθ += θ; sumθ₀ += θ₀
              ϵ += 1
          else
              θ  = (1-ηₜ*λ) * θ
          end
      end
      if (ϵ == 0)
          if nMsgs != 0
              println("*** Avg. error after epoch $t : $(ϵ/size(x)[1]) (all elements of the set has been correctly classified")
          end
          return (θ=θ,θ₀=θ₀,avgθ=sumθ/(n*T),avgθ₀=sumθ₀/(n*T),errors=0,besterrors=0,iterations=t,separated=true)
      elseif ϵ < bestϵ
          bestϵ = ϵ
      end
      lastϵ = ϵ
      if nMsgs != 0 && (t % ceil(T/nMsgs) == 0 || t == 1 || t == T)
        println("Avg. error after iteration $t : $(ϵ/size(x)[1])")
      end
  end
  return  (θ=θ,θ₀=θ₀,avgθ=sumθ/(n*T),avgθ₀=sumθ₀/(n*T),errors=lastϵ,besterrors=bestϵ,iterations=T,separated=false)
end

# ------------------------------------------------------------------------------
# Other functions

"""
  predict(x,θ,θ₀)

Predict a binary label {-1,1} given the feature vector and the linear coefficients

# Parameters:
* `x`:        Feature matrix of the training data (n × d)
* `θ`:        The trained parameters
* `θ₀`:       The trained bias barameter [def: `0`]

# Return :
* `y`: Vector of the predicted labels

# Example:
```julia
julia> predict([1.1 2.1; 5.3 4.2; 1.8 1.7], [3.2,1.2])
```
"""
function predict(x,θ,θ₀=0.0)
    x = makeMatrix(x)
    θ = makeColVector(θ)
    (n,d) = size(x)
    d2 = length(θ)
    if (d2 != d) error("x and θ must have the same dimensions."); end
    y = zeros(Int64,n)
    for i in 1:n
        y[i] = (θ' * x[i,:] + θ₀) > eps() ? 1 : -1  # no need to divide by the norm to get the sign!
    end
    return y
end

"""
  predict(x,θ,θ₀,classes)

Predict a multiclass label given the feature vector, the linear coefficients and the classes vector

# Parameters:
* `x`:       Feature matrix of the training data (n × d)
* `θ`:       Vector of the trained parameters for each one-vs-all model (i.e. `model.θ`)
* `θ₀`:      Vector of the trained bias barameter for each one-vs-all model (i.e. `model.θ₀`)
* `classes`: The overal classes encountered in training (i.e. `model.classes`)

# Return :
* `ŷ`: Vector of dictionaries `label=>probability`

 # Notes:
 * Use `mode(ŷ)` if you want a single predicted label per record

# Example:
```julia
julia> model  = perceptron([1.1 2.1; 5.3 4.2; 1.8 1.7], [-1,1,-1])
julia> ŷtrain = predict([10 10; 2.5 2.5],model.θ,model.θ₀, model.classes)
"""
function predict(x,θ::AbstractVector{T},θ₀::AbstractVector{Float64},classes::Vector{Tcl}) where {T<: AbstractVector{Float64},Tcl}
    (n,d) = size(x)
    nCl   = length(classes)
    y     = Array{Dict{Tcl,Float64},1}(undef,n)
    for i in 1:n
        probRaw = Array{Float64,1}(undef,nCl)
        for (c,cl) in enumerate(classes)
            if nCl == 2 && c ==2
                probRaw[2] = - probRaw[1]
            else
                probRaw[c] = (θ[c]' * x[i,:] + θ₀[c])
            end
        end
        prob = softmax(probRaw)
        y[i] = Dict(zip(classes,prob))
    end
    return y
end

"""
  predict(x,xtrain,ytrain,α;K)

Predict a binary label {-1,1} given the feature vector and the training data together with their errors (as trained by a kernel perceptron algorithm)

# Parameters:
* `x`:      Feature matrix of the training data (n × d)
* `xtrain`: The feature vectors used for the training
* `ytrain`: The labels of the training set
* `α`:      The errors associated to each record
* `K`:      The kernel function used for the training and to be used for the prediction [def: [`radialKernel`](@ref)]

# Return :
* `y`: Vector of the predicted labels

# Example:
```julia
julia> predict([1.1 2.1; 5.3 4.2; 1.8 1.7], [3.2,1.2])
```
"""
function predict(x,xtrain,ytrain,α;K=radialKernel)
    x = makeMatrix(x)
    xtrain = makeMatrix(xtrain)
    (n,d) = size(x)
    (ntrain,d2) = size(xtrain)
    if (d2 != d) error("xtrain and x must have the same dimensions."); end
    if ( length(ytrain) != ntrain || length(α) != ntrain) error("xtrain, ytrain and α must all have the same length."); end
    y = zeros(Int64,n)
    for i in 1:n
        y[i] = sum([ α[j] * ytrain[j] * K(x[i,:],xtrain[j,:]) for j in 1:ntrain]) > eps() ? 1 : -1
    end
    return y
 end

 """
   predict(x,xtrain,ytrain,α,classes;K)

 Predict a multiclass label given the new feature vector and a trained kernel perceptron model.

 # Parameters:
 * `x`:      Feature matrix of the training data (n × d)
 * `xtrain`: A vector of the feature matrix used for training each of the one-vs-one class matches (i.e. `model.x`)
 * `ytrain`: A vector of the label vector used for training each of the one-vs-one class matches (i.e. `model.y`)
 * `α`:      A vector of the errors associated to each record (i.e. `model.α`)
 * `classes`: The overal classes encountered in training (i.e. `model.classes`)
 * `K`:      The kernel function used for the training and to be used for the prediction [def: [`radialKernel`](@ref)]

 # Return :
 * `ŷ`: Vector of dictionaries `label=>probability` (warning: it isn't really a probability, it is just the standardized number of matches "won" by this class compared with the other classes)

  # Notes:
  * Use `mode(ŷ)` if you want a single predicted label per record

 # Example:
 ```julia
 julia> model  = kernelPerceptron([1.1 2.1; 5.3 4.2; 1.8 1.7], [-1,1,-1])
 julia> ŷtrain = Perceptron.predict([10 10; 2.2 2.5],model.x,model.y,model.α, model.classes,K=model.K)
 ```
 """
 function predict(x,xtrain,ytrain,α,classes::AbstractVector{Tcl};K=radialKernel) where {Tcl}
     (n,d)   = size(x)
     nCl     = length(classes)
     y       = Array{Dict{Tcl,Float64},1}(undef,n)
     nModels = Int((nCl  * (nCl - 1)) / 2)
     if !(nModels == length(xtrain) == length(ytrain) == length(α)) error("xtrain, ytrain or α have a length not compatible with the number of classes in this model."); end
     x = makeMatrix(x)
     d2 = size(xtrain[1],2)
     if (d2 != d) error("xtrain and x must have the same dimensions."); end
     for i in 1:n
         #countByClass = zeros(Float64,nCl)
         countByClass = zeros(Int64,nCl)
         mCounter = 1
          for (ic,c) in enumerate(classes)
             for (ic2,c2) in enumerate(classes)
                 if ic2 <= ic
                     continue
                 end
                 nThisModel = size(xtrain[mCounter],1)
                 if ( length(ytrain[mCounter]) != nThisModel || length(α[mCounter]) != nThisModel) error("xtrain, ytrain and α must all have the same length."); end
                 # note that we assign "winning" scores between pair of classes matches only based on who win, not by how much he did
                 # todo check
                 #if sum([ α[mCounter][j] * ((ytrain[mCounter][j] .== c) .*2 .-1) * K(x[i,:],xtrain[mCounter][j,:]) for j in 1:nThisModel]) > eps()
                 score = sum([ α[mCounter][j] * (ytrain[mCounter][j]) * K(x[i,:],xtrain[mCounter][j,:]) for j in 1:nThisModel])
                 #println(score)
                 if score > eps()
                     countByClass[ic] += 1
                     #countByClass[ic] += score
                 else
                     countByClass[ic2] += 1
                     #countByClass[ic2] += -score
                 end
                 mCounter += 1
             end
         end
         #println(countByClass)
         prob = softmax(countByClass)
         y[i] = Dict(zip(classes,prob))
    end
    return y
end

# MLJ interface
include("Perceptron_MLJ.jl")

end


using Base.Test
using PowerSeries

# Represent the series 1.0 + 1.0*x - 2.0*x^2
a = series(1.0, 1.0, -2.0)

# Represent the series 1.0 + 0.0*x + 1.0*x^2
b = series(1.0, 0.0, 1.0)

# Series add linearly
@test a+b == series(2.0, 1.0, -1.0)

# The output of series operations is truncated to match the input size.
# Represents (1+x-2x^2)(1 + x^2) = 1+x-x^2+o(x^3)
@test a*b == series(1.0,1.0,-1.0)

# Functions with known derivatives can easily be overloaded to operate on
# power series.
# You can generate the taylor series of a function about a point x up to
# e.g. 6th order by computing f(Series(x, 1.0, 0.0, 0.0, 0.0, 0.0))
x = series(0.0, 1.0, 0.0, 0.0, 0.0, 0.0)
@test sin(x) == series(0.0, 1.0, 0.0, -1.0/6.0, 0.0, 1.0/120)
@test log(1 + x) == series(0.0, 1.0, -1.0/2, 1.0/3, -1.0/4, 1.0/5)
@test 1/(1 - x) == series(1.0, 1.0, 1.0, 1.0, 1.0, 1.0)

@test polyval(a, 0.1) == 1.08
@test polyval(a, -0.1) == 0.88
@test polyder(a) == series(1.0, -4.0)
@test polyint(a) == series(0.0, 1.0, 1.0/2, -2.0/3)
@test restrict(a) == series(1.0, 1.0)

# Power series are one of the best ways to take higher order derivatives
# of generic functions.
#
# This example tests promotion rules in the constructor because of the
# 1 and 0, instead of 1.0 and 0.0
f(x) = exp(-x^2)
f2(x) = polyder(polyder(f(series(x, 1, 0))))
@test f2(2.0) == let x = 2.0; -2exp(-x^2)+4x^2*exp(-x^2); end

# macro-expansion-time conditional.
# This is probably not very good or general, but gets around the fact that
# the signture for @test_throws changed between versions 0.2 and 0.3
macro when(cond, ex)
  eval(cond) ? esc(ex) : nothing
end

# PowerSeries comes with types defined for series up to order 7. By default,
# trying to construct a higher order series is a type error.
@when(
  VERSION >= VersionNumber(0,2,typemax(Int64)),
  @test_throws MethodError series(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)
)

# If you want to work with higher order series, you can generate types up
# to a given order with PowerSeries.generate(order)
PowerSeries.generate(9)
@test series(0, 1, 2, 3, 4, 5, 6, 7, 8, 9) != false

@test series(1, 2) - 1 == series(0, 2)

# Regression test for https://github.com/jwmerrill/PowerSeries.jl/issues/6
let x=series(0.0, 1.0, 0.0, 0.0)
  @test x*x == x^2
end

@test series(1.0, 2.0, 3.0) < 2.0
@test 2.0 > series(1.0, 2.0, 3.0)
@test 0.5 < series(1.0, 2.0, 3.0)
@test series(1.0, 2.0, 3.0) > 0.5
@test !(series(1.0, 2.0, 3.0) < 1.0)
@test !(series(1.0, 2.0, 3.0) > 1.0)
@test series(1.0, 1.0) < series(2.0, 0.0)
@test series(2.0, 0.0) > series(1.0, 1.0)
@test !(series(1.0, 2.0) < series(1.0, 0.0))
@test !(series(1.0, 0.0) > series(1.0, 2.0))

# Only 1 argument functions
fns = [
  sqrt,
  exp,
  log,
  sin,
  cos,
  tan,
  asin,
  acos,
  atan,
  sinh,
  cosh,
  tanh,
  asinh,
  acosh,
  atanh,
  csc,
  sec,
  cot,
  acsc,
  asec,
  acot,
  csch,
  sech,
  coth,
  acsch,
  asech,
  acoth,
  gamma,
  floor,
  ceil,
  round,
  sign,
  abs
]

for val in [-2.3, -1.2, -0.1, 0.7, 3.6]
  for fn in fns
    if (val < 1 && fn === acosh) continue end
    if (val < 0 && fn === sqrt) continue end
    if (val < 0 && fn === log) continue end
    if (abs(val) > 1 && fn === asin) continue end
    if (abs(val) > 1 && fn === acos) continue end
    if (abs(val) > 1 && fn === atanh) continue end
    if ((val < 0 || val > 1) && fn === asech) continue end
    if (abs(val) < 1 && fn === acsc) continue end
    if (abs(val) < 1 && fn === asec) continue end
    if (abs(val) < 1 && fn === acoth) continue end
    @test_approx_eq_eps constant(polyder(fn(series(val, 1.0)))) 5000*(fn(val + .0001) - fn(val- .0001)) .01
  end
end


using HyperRectangles
using HyperRectangles.Relations
using HyperRectangles.Operations
using Base.Test

# test constructors and containment
let
    a= HyperRectangle{Float64,4}([Inf,Inf,Inf,Inf],[-Inf,-Inf,-Inf,-Inf])

    update!(a, [1,2,3,4])

    @test a == HyperRectangle{Float64,4}([1.0,2.0,3.0,4.0],[1.0,2.0,3.0,4.0])
    @test a != HyperRectangle{Float64,4}([3.0,2.0,3.0,4.0],[1.0,2.0,3.0,4.0])

    update!(a, [5,6,7,8])

    b = HyperRectangle{Float64,4}([1.0,2.0,3.0,4.0],[5.0,6.0,7.0,8.0])
    @test a == b
    @test isequal(a,b)

    @test max(a) == [5.0,6.0,7.0,8.0]
    @test min(a) == [1.0,2.0,3.0,4.0]

    @test_throws ErrorException HyperRectangle([1.0,2.0,3.0],[1.0,2.0,3.0,4.0])

    @test in(a,b) && in(b,a) && contains(a,b) && contains(b,a)

    c = HyperRectangle([1.1,2.1,3.1,4.1],[4.0,5.0,6.0,7.0])

    @test !in(a,c) && in(c,a) && contains(a,c) && !contains(c,a)
end

# Testing split function
let
    d = HyperRectangle{Float64,4}([1.0,2.0,3.0,4.0],[2.0,3.0,4.0,5.0])
    d1, d2 = split(d, 3, 3.5)

    @test d1.max[3] == 3.5 && d1.min[3] == 3.0
    @test d2.max[3] == 4.0 && d2.min[3] == 3.5
end

#test points function
let
    a = HyperRectangle([0,0],[1,1])
    pt_expa = Vector[[0,0],[0,1],[1,0],[1,1]]
    @test points(a) == pt_expa
    b = HyperRectangle([0,0,0],[1,1,1])
    pt_expb = Vector[[0,0,0],[0,0,1],[0,1,0],[0,1,1],[1,0,0],[1,0,1],[1,1,0],[1,1,1]]
    @test points(b) == pt_expb
end

# test empty constructor on 0.4
if VERSION >= v"0.4.0-"
    let
        a = HyperRectangle{Float64, 4}()
        @test a == HyperRectangle{Float64,4}([Inf,Inf,Inf,Inf],[-Inf,-Inf,-Inf,-Inf])
    end
end


import DomainSets: boundary

"""
    abstract type AbstractSubGrid{T,N} <: AbstractGrid{T,N} end

A subgrid of an underlying grid.
"""
abstract type AbstractSubGrid{T,N} <: AbstractGrid{T,N} end

"""
    supergrid(g::AbstractSubGrid)

The underlying grid of the subgrid.
"""
supergrid(g::AbstractSubGrid) = g.supergrid

include("indexsubgrid.jl")
include("maskedsubgrid.jl")
include("product.jl")
include("boundary.jl")

subgrid(grid::AbstractGrid, domain::Domain) = MaskedGrid(grid, domain)

function subgrid(grid::AbstractEquispacedGrid, domain::AbstractInterval)
    a = infimum(domain)
    b = supremum(domain)
    h = step(grid)
    idx_a = convert(Int, ceil( (a-grid[1])/step(grid))+1 )
    idx_b = convert(Int, floor( (b-grid[1])/step(grid))+1 )
    idx_a = max(idx_a, 1)
    idx_b = min(idx_b, length(grid))
    IndexSubGrid(grid, idx_a:idx_b, domain)
end

function subgrid(grid::ScatteredGrid, domain::Domain)
    mask = in.(grid, Ref(domain))
    points = grid.points[mask]
    ScatteredGrid(points, domain)
end

subgrid(grid::ProductGrid, domain::ProductDomain) =
    ProductGrid(map(subgrid, elements(grid), elements(domain))...)

# supergrid(grid::ProductGrid) = grid
# supergrid(grid::ProductGrid{Tuple{Vararg{G}}}) where G<:Union{MaskedGrid,IndexSubGrid} = ProductGrid(map(supergrid, elements(grid))...)


# buidling connection matrices

# Utility functions
# for connectivity

"""
        lognorm_reparametrize(m,std) -> distr::LogNormal
# parameters
  + `m`   sample mean
  + `std` sample std
"""
function lognorm_reparametrize(m::Real,std::Real)
    vm2= (std/m)^2
    μ = log(m / sqrt(1. + vm2))
    σ = sqrt(log( 1. + vm2))
    return LogNormal(μ,σ)
end

function no_autapses!(mat::SparseMatrixCSC)
  for i in 1: min(size(mat)...)
    mat[i,i]=0.0
  end
  return dropzeros!(mat)
end

# scale rows so that the sum of the nonzero elements is exactly x
# for simplicity ignores the sign of x
function scale_rows!(x::Real,mat::SparseMatrixCSC)
  x=abs(x) # there are better ways to change sign!
  n = size(mat,1)
  vals=nonzeros(mat)
  rows=rowvals(mat)
  for i in 1:n
    idxs = findall(r->r==i,rows)
    if !isempty(idxs) # is it control useful? or less efficient?
      s = sum(vals[idxs])
      vals[idxs] .*= x/s
    end
  end
  return mat
end

function _sparse_wmat(npost::Integer,npre::Integer,p::Real,make_weights!::Function ;
    scal::Float64=1.0, # final scaling (e.g. -1 for inhibitory neurons)
    noself::Bool=true, # no autapses ?
    rowsum::Union{Nothing,Float64}=nothing) #mean of each row (before scal)
  wmat = sprand(Float64,npost,npre,p)
  if noself
    no_autapses!(wmat)
  end
  make_weights!(wmat)
  if !isnothing(rowsum) # scale so that each row has mean exactly rowsum mean
    scale_rows!(rowsum,wmat)
  end
  if scal != 1.0
    lmul!(scal,wmat)
  end
  return wmat
end

function sparse_wmat_rowfun(make_a_row::Function,
    npost::Integer,npre::Integer,p::Real;no_autapses::Bool=false)
  wmat = sprand(Float64,npost,npre,p)
  if no_autapses
    no_autapses!(wmat)
  end
  rw = rowvals(wmat)
  wvec = nonzeros(wmat)
  for r in unique(rw)
    ridx = r .== rw
    n = count(ridx)
    newrow = make_a_row(n)
    wvec[ridx] .= newrow
  end
  return wmat
end

function sparse_constant_wmat(npost::Integer,npre::Integer,p::Real,j_val::Real ;
    scal::Float64=1.0, # final scaling
    no_autapses::Bool=true, # no autapses ?
    rowsum::Union{Nothing,Float64}=nothing)
  make_weights=function(n::Integer)
    ret=fill(j_val,n)
    if !isnothing(rowsum)
      ret .*= rowsum / sum(ret)
    end
    if scal != 1
      lmul!(scal,ret)
    end
    return ret
  end
  return sparse_wmat_rowfun(make_weights,npost,npre,p;no_autapses=no_autapses)
end

function sparse_wmat_lognorm(npost::Integer,npre::Integer,
    p::Real,μ::Real,σ::Real;
    noself::Bool=true, # no autapses ?
    rowsum::Union{Nothing,Float64}=nothing)
  μp = abs(μ)
  distr = lognorm_reparametrize(μp,σ)
  make_weights=function (smat)
    vals=nonzeros(smat)
    Distributions.rand!(distr,vals)
  end
  scal = μ < 0 ? -1.0 : 1.0
  return _sparse_wmat(npost,npre,p,make_weights;
    scal=scal,noself=noself,rowsum=rowsum)
end

function make_sparse_weights(npost::Integer,npre::Integer,p::Real,weigth_value::Real ;
    scal::Float64=1.0, # final scaling
    noself::Bool=true, # no autapses ?
    rowsum::Union{Nothing,Float64}=nothing)
  make_weights=function (smat)
    vals=nonzeros(smat)
    fill!(vals,weigth_value)
  end
  return _sparse_wmat(npost,npre,p,make_weights;
    scal=scal,noself=noself,rowsum=rowsum)
end

# not needed with new Poisson input types!
# """
#     wmat_train_assemblies_protocol(Nneus::Integer,p::Real; scal::Float64=1.0)
#       -> (weights::SparseMatrixCSC,assembly_dictionary::Dict)

#   # Arguments
#   + Npost::Integer : number of postsynaptic neurons
#   + Nassemblies::Integer : number of assemblies to train
#   + p_assembly::Real : probability for a neuron to be in any assembly
#   + scal::Real=1.0 : connection weights

#   # Outputs
#   + weights : sparse weight matrix
#   + assembly_list::Vector{NamedTuple{(:neupost,:neupre),Vector{Int64,Vector{Int64}}}}
#     element `i` contains the pre and post neurons associated to assembly `i`
# """
# function wmat_train_assemblies_protocol(Npost::Integer,Nassemblies::Integer,
#     p_assembly::Real;scal::Float64=1.0)
#   # unoptimized... but called only once
#   post_neus = [findall(rand(Npost) .< p_assembly) for _ in 1:Nassemblies]
#   c = 0
#   pre_neus = map(post_neus) do _po
#     _ln = length(_po)
#     ret = c .+ collect(1:_ln)
#     c += _ln
#     return ret
#   end
#   _pre_neus_all = vcat(pre_neus...)
#   wmat = sparse(vcat(post_neus...),_pre_neus_all,scal,
#     Npost,length(_pre_neus_all)) # include total matrix size
#   assembly_list = map(zip(post_neus,pre_neus)) do (_post,_pre)
#     (neupost=_post,neupre=_pre)
#   end
#   return wmat,assembly_list
# end

# to order assemblies (no idea where to place this function)
function order_by_pattern_idxs(pattern_idxs::Vector{Vector{Int64}},
    nneus::Int64;rev::Bool=false)
  _tosort = map(1:nneus) do neu
    pa = findall(patt->neu in patt,pattern_idxs)
    isnothing(pa) && return 10_000_000
    (length(pa) == 1) && return pa[1]
    (length(pa)>1) && return sample(pa)
  end
  return sortperm(_tosort;rev=rev)
end


function consolidate{T}(a::Array{Set{T},1})
    1 < length(a) || return a
    b = copy(a)
    c = Set{T}[]
    while 1 < length(b)
        x = shift!(b)
        cme = true
        for (i, y) in enumerate(b)
            !isempty(intersect(x, y)) || continue
            cme = false
            b[i] = union(x, y)
            break
        end
        !cme || push!(c, x)
    end
    push!(c, b[1])
    return c
end


#
# Uniform scores test statistic
#
# Mardia, K. V. (1972). A Multi-Sample Uniform Scores Test on a Circle
# and Its Parametric Competitor. Journal of the Royal Statistical
# Society. Series B (Methodological), 34(1), 102–113.
#
using StatsBase
immutable UniformScores <: PairwiseStatistic
    groups::Vector{Int32}
    ningroup::Vector{Int32}

    function UniformScores{T<:Integer}(groups::Vector{T})
        ext = extrema(groups)
        ext[1] == 1 || throw(ArgumentError("group indices must start at 1"))
        ningroup = zeros(Int32, ext[2])
        for i = 1:length(groups)
            ningroup[groups[i]] += 1
        end
        new(groups, ningroup)
    end
end
Base.eltype{T<:Real}(::UniformScores, X::AbstractArray{Complex{T}}) = T

allocwork{T<:Real}(t::UniformScores, X::AbstractVecOrMat{Complex{T}}) =
    (zeros(T, size(X)), zeros(T, ntrials(X)), zeros(Int, ntrials(X)), [cis(2π*i/size(X, 1)) for i = 1:size(X, 1)], zeros(Complex{T}, length(t.ningroup)))

# Single input matrix
using Base.Sort: DEFAULT_UNSTABLE, Perm, Forward
function computestat!{T<:Real}(t::UniformScores, out::AbstractMatrix{T},
                               work::Tuple{Matrix{T}, Vector{T}, Vector{Int}, Vector{Complex{T}}, Vector{Complex{T}}}, X::AbstractVecOrMat{Complex{T}})
    chkinput(out, X)
    angles, anglediff, ranks, circranks, d = work
    (size(X, 1) == size(angles, 1) == length(anglediff) == length(ranks) == length(circranks) == length(t.groups) &&
        size(X, 2) == size(angles, 2)) && length(d) == length(t.ningroup) || throw(ArgumentError("invalid work object"))
    for i = 1:length(X)
        @inbounds angles[i] = angle(X[i])
    end

    for k = 1:size(X, 2), j = 1:k-1
        @inbounds @simd for i = 1:size(X, 1)
            anglediff[i] = mod2pi(angles[i, j] - angles[i, k])
            ranks[i] = i
        end
        # sortperm!, but don't allocate memory
        sort!(ranks, DEFAULT_UNSTABLE, Perm(Forward, anglediff))
        fill!(d, zero(T))
        @simd for i = 1:size(X, 1)
            @inbounds d[t.groups[ranks[i]]] += circranks[i]
        end
        Wr = zero(T)
        @simd for i = 1:length(d)
            @inbounds Wr += abs2(d[i])/t.ningroup[i]
        end
        out[j, k] = 2*Wr
    end
    out
end


using Optim
using HTTP
using GLM
using LinearAlgebra
using Random
using Statistics
using DataFrames
using DataFramesMeta
using CSV
using MultivariateStats
using Distributions

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 1
#:::::::::::::::::::::::::::::::::::::::::::::::::::
url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2021/master/ProblemSets/PS8-factor/nlsy.csv"
df = CSV.read(HTTP.get(url).body,DataFrame)
show(df, allcols=true)

ols = lm(@formula(logwage ~ black + hispanic + female + schoolt + grad4yr + gradHS),df)
println(ols)

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 2
#:::::::::::::::::::::::::::::::::::::::::::::::::::
println(cor(Matrix(@select(df, :asvabAR, :asvabCS, :asvabMK, :asvabNO, :asvabPC, :asvabWK))))

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 3
#:::::::::::::::::::::::::::::::::::::::::::::::::::
ols = lm(@formula(logwage ~ black + hispanic + female + schoolt + grad4yr + gradHS + asvabAR + asvabCS + asvabMK + asvabNO + asvabPC + asvabWK),df)
println(ols)
#it will be problematic to directly include these in the regression since the correlation among six asvabs is too high

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 4
#:::::::::::::::::::::::::::::::::::::::::::::::::::
asvabMat = Matrix(@select(df, :asvabAR, :asvabCS, :asvabMK, :asvabNO, :asvabPC, :asvabWK))'
M = fit(PCA, asvabMat; maxoutdim=1)
asvabPCA = MultivariateStats.transform(M, asvabMat)
insertcols!(df, 16, :asvabFac => vec(asvabPCA'))
ols = lm(@formula(logwage ~ black + hispanic + female + schoolt + grad4yr + gradHS + asvabFac),df)
println(ols)

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 5
#:::::::::::::::::::::::::::::::::::::::::::::::::::
asvabMat = Matrix(@select(df, :asvabAR, :asvabCS, :asvabMK, :asvabNO, :asvabPC, :asvabWK))'
M = fit(FactorAnalysis, asvabMat; maxoutdim=1)
asvabFA= MultivariateStats.transform(M, asvabMat)
insertcols!(df, 17, :asvabFacFA => vec(asvabFA'))
ols = lm(@formula(logwage ~ black + hispanic + female + schoolt + grad4yr + gradHS + asvabFacFA),df)
println(ols)

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 6
#:::::::::::::::::::::::::::::::::::::::::::::::::::
include("E://oklohoma//ECON6343//fall-2021-master//ProblemSets//PS4-mixture//lgwt.jl") # make sure the function gets read in

function ms_mle(x,df,N_grid)
    N = size(df,1)
    J = 6
    alpha = reshape(x[1:4*J],(4,J))
    beta = reshape(x[4*J+1:4*J+7],(7,1))
    gamma = x[4*J+8:4*J+13]
    sigmaj = x[4*J+14:4*J+19]
    delta = x[end-1]
    sigmaw = x[end]
    y = Matrix(@select(df, :logwage))
    M = Matrix(@select(df, :asvabAR, :asvabCS, :asvabMK, :asvabNO, :asvabPC, :asvabWK))
    Xm = hcat(ones(N),Matrix(@select(df, :black, :hispanic, :female)))
    X = hcat(ones(N),Matrix(@select(df, :black, :hispanic, :female,:schoolt, :grad4yr, :gradHS)))
    d = Normal(0,1)
    nodes, weights = lgwt(N_grid,-4,4)
    Random.seed!(1234)
    xi = randn(N_grid)
    Li = ones(N)
    lik = zeros(N)
    for i = 1:N_grid
        for j = 1:J
            Li .*= pdf.(d, (M[:,j].-Xm*alpha[:,j].-gamma[j]*xi[i])./sigmaj[j])./sigmaj[j]
        end
        Li = Li .*pdf.(d, (y.-X*beta.-delta*xi[i])./sigmaw)./sigmaw
        lik =  lik .+ Li .*weights[i].*pdf.(d,xi[i])
    end
    loglik = -sum(log.(lik))
    return loglik
end

x_start = rand(45)
lower = vcat(repeat([-Inf],37),zeros(6),[-Inf],[0])
upper = repeat([Inf],45)
N_grid = 1000
x_optim = optimize(x -> ms_mle(x, df, N_grid),lower, upper, x_start, Fminbox(NelderMead()), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
x_mle = x_optim.minimizer
println(x_mle)


#---------------------------------------------------------
# # [Rectangle](@id 03-rect)
#---------------------------------------------------------

#=
This page illustrates the `Rect` shape in the Julia package
[`ImagePhantoms`](https://github.com/JuliaImageRecon/ImagePhantoms.jl).

This page was generated from a single Julia file:
[03-rect.jl](@__REPO_ROOT_URL__/03-rect.jl).
=#

#md # In any such Julia documentation,
#md # you can access the source code
#md # using the "Edit on GitHub" link in the top right.

#md # The corresponding notebook can be viewed in
#md # [nbviewer](http://nbviewer.jupyter.org/) here:
#md # [`03-rect.ipynb`](@__NBVIEWER_ROOT_URL__/03-rect.ipynb),
#md # and opened in [binder](https://mybinder.org/) here:
#md # [`03-rect.ipynb`](@__BINDER_ROOT_URL__/03-rect.ipynb).

# ### Setup

# Packages needed here.

using ImagePhantoms: Rect, phantom, radon, spectrum
using ImageGeoms: ImageGeom, axesf
using MIRTjim: jim, prompt
using FFTW: fft, fftshift
using Unitful: mm, unit, °
using UnitfulRecipes
using Plots: plot, plot!, scatter!, default; default(markerstrokecolor=:auto)

# The following line is helpful when running this example.jl file as a script;
# this way it will prompt user to hit a key after each figure is displayed.

isinteractive() ? jim(:prompt, true) : prompt(:draw);

# ### Overview

#=
A basic shape used in constructing 2D digital image phantoms
is the rectangle, specified by its center, widths, angle and value.
All of the methods in `ImagePhantoms` support physical units,
so we use such units throughout this example.
(Using units is recommended but not required.)

Define an rectangle object, using physical units.
=#

width = (2mm, 8mm)
ob = Rect((4mm, 3mm), width, π/6, 1.0f0)

# ### Phantom image using `phantom`

# Make a digital image of it using `phantom` and display it.
dx = 0.02mm
dy = 0.025mm
(M,N) = (2^10,2^10+2)
x = (-M÷2:M÷2-1) * dx
y = (-N÷2:N÷2-1) * dy
img = phantom(x, y, [ob])
jim(x, y, img)

# Hereafter we use `ImageGeoms` to simplify the indexing.

ig = ImageGeom(dims=(M,N), deltas=(dx,dy), offsets=(0.5,0.5))
@assert all(axes(ig) .≈ (x,y))
p1 = jim(axes(ig)..., img, "Rect phantom", xlabel="x", ylabel="y")

# ### Spectrum using `spectrum`

#=
There are two ways to examine the spectrum of this image:
* using the analytical Fourier transform of the rect via `spectrum`
* applying the DFT via FFT to the digital image.
Because the shape has units `mm`, the spectra axes have units cycles/mm.
=#

zscale = 1 / prod(width) # normalize spectra by area
spectrum_exact = spectrum(axesf(ig)..., [ob]) * zscale
sp = z -> max(log10(abs(z)/oneunit(abs(z))), -6) # log-scale for display
clim = (-6, 0) # colorbar limit for display
(xlabel, ylabel) = ("ν₁", "ν₂")
p2 = jim(axesf(ig)..., sp.(spectrum_exact), "log10|Spectrum|"; clim, xlabel, ylabel)

# Sadly `fft` cannot handle units currently, so this function is a work-around:
function myfft(x)
    u = unit(eltype(x))
    return fftshift(fft(fftshift(x) / u)) * u
end

#src fx = (-M÷2:M÷2-1) / M / dx # appropriate frequency axes for DFT,
#src fy = (-N÷2:N÷2-1) / N / dy # that are provided by axesf(ig)
spectrum_fft = myfft(img) * dx * dy * zscale
p3 = jim(axesf(ig)..., sp.(spectrum_fft), "log10|DFT|"; clim, xlabel, ylabel)

# Compare the DFT and analytical spectra to validate the code
@assert maximum(abs, spectrum_exact - spectrum_fft) /
        maximum(abs, spectrum_exact) < 2e-2
p4 = jim(axesf(ig)..., abs.(spectrum_fft - spectrum_exact), "Difference"; xlabel, ylabel)
jim(p1, p4, p2, p3)

# ### Radon transform using `radon`

# Examine the Radon transform of the rect using `radon`,
# and validate it using the projection-slice theorem aka Fourier-slice theorem.

dr = 0.02mm # radial sample spacing
nr = 2^10 # radial sinogram bins
r = (-nr÷2:nr÷2-1) * dr # radial samples
fr = (-nr÷2:nr÷2-1) / nr / dr # corresponding spectral axis
ϕ = deg2rad.(0:180) # * Unitful.rad # todo round unitful Unitful.°
sino = radon(ob).(r, ϕ') # sample Radon transform of a single shape object
p5 = jim(r, rad2deg.(ϕ), sino; aspect_ratio=:none, title="sinogram", yflip=false, xlabel="r", ylabel="ϕ")

#src clim=(0, maximum(width)*ob.value), # todo

#=
The maximum sinogram value is about sqrt(8^2+2^2) = 8.2mm,
which makes sense for a 8mm × 2mm rect.

The above sampling generated a parallel-beam sinogram,
but one could make a fan-beam sinogram by sampling `(r, ϕ)` appropriately.
=#

# ### Fourier-slice theorem illustration

# Pick one particular view angle (55°) and look at its slice and spectra.
ia = argmin(abs.(ϕ .- 55°))
slice = sino[:,ia]
slice_fft = myfft(slice) * dr
angle = round(rad2deg(ϕ[ia]), digits=1)

kx, ky = (fr * cos(ϕ[ia]), fr * sin(ϕ[ia])) # Fourier-slice theorem
slice_ft = spectrum(ob).(kx, ky)
@assert maximum(abs, slice_ft - slice_fft) / maximum(abs, slice_ft) < 2e-4

p3 = plot(r, slice, title="profile at ϕ = $angle", label="")
p4 = plot(title="1D spectra")
scatter!(fr, abs.(slice_fft), label="abs fft", color=:blue)
scatter!(fr, real(slice_fft), label="real fft", color=:green)
scatter!(fr, imag(slice_fft), label="imag fft", color=:red, xlims=(-1,1).*(1.0/mm))

plot!(fr, abs.(slice_ft), label="abs", color=:blue)
plot!(fr, real(slice_ft), label="real", color=:green)
plot!(fr, imag(slice_ft), label="imag", color=:red)
plot(p1, p5, p3, p4)

#=
The good agreement between the analytical spectra (solid lines)
and the DFT samples (disks)
validates that `phantom`, `radon`, and `spectrum`
are all self consistent for this `Rect` object.
=#


using BandedMatrices, LinearAlgebra
export ProductInversion

"""
    ProductInversion(r, μ; verbose=false) <: Decomposition
Decompose a timeseries `s` into a **product** `x * r`, given that you have a good
estimate of the second factor `r` (the "input") and you need `x` (the "multiplier")
but you can't do simply `x =  s ./ r` because `r` contains zeros.

This method works well when the characteristic timescales of `x` are comparable,
or larger than those of `r` but not much smaller.

The second argument `μ` is a regularization parameter. In short, we estimate `r`
by minimizing a cost with two components: that `x` is close to `s/r` and that `x`
is smooth. `μ` is the multiplier of the smoothness cost.

You can give a vector as `μ`. The process will be repeated for all `μ` and the
[`rmse`](@ref) between `s` and the estimated `x * r` will be computed. The
`x` that gives the least error will be returned finally.
If `verbose = true`, the method also prints the pairs `(μ, err)` for each `μ`.

Use the low level `SignalDecomposition.matrix_invert(s, r, μ::Real) → x, err`
to get the error values.
"""
struct ProductInversion{R, M} <: Decomposition
    r::R
    μ::M
    x::R
    xdummy::R
    verbose::Bool
end
ProductInversion(r, μ; verbose=false) = ProductInversion(r, μ, copy(r), copy(r), verbose)

# TODO: performant method `update!(method, s, μ)`

function decompose(t, s, method::ProductInversion{R, <: Real}) where R
    x, err = matrix_invert(s, method.r, method.μ)
    return x, method.r
end

function decompose(t, s, method::ProductInversion{R, <: AbstractArray}) where R
    isequispaced(t) || error("Input time axis must be equispaced for method ProductInversion.")
    leasterr = Inf
    bestx = copy(s)
    bestμ = method.μ[1]
    for μ in method.μ
        x, err = matrix_invert(s, method.r, μ, method.x, method.xdummy)
        method.verbose && @show (μ, err)
        if err < leasterr
            bestx .= x
            leasterr = err
            bestμ = μ
        end
    end
    method.verbose && @show (bestμ, leasterr)
    return bestx, method.r
end

function matrix_invert(s::AbstractVector{T}, r, μ::Real, x = copy(r), xdummy = copy(r)) where {T}
    # TODO: This can be optimized by pre-allocating the diags dictinary
    # and simply updating the values for each μ
    # TODO : add @inbounds
    L = length(s)
    diags = Dict(
        0 => T[6μ + r[i]^2 for i in 1:L],
        1 => fill(T(-4μ), L-1),
        -1 => fill(T(-4μ), L-1),
        2 => fill(T(μ), L-2),
        -2 => fill(T(μ), L-2),
    )
    # take care of border elements
    diags[1][1] = -8μ
    diags[2][1:2] .= 2μ
    diags[-1][end] = -8μ
    diags[-2][end-1:end] .= 2μ
    B = BandedMatrix(diags...)
    @. xdummy = s * r
    ldiv!(x, B, xdummy)
    @. xdummy = x * r
    err = rmse(s, xdummy)
    return x, err
end


export ProbabilityDistribution, ConditionalDistribution, JointProbabilityDistribution
export is_probability_distribution, is_conditional_distribution
export Probabilities, Conditionals, JointProbabilities
export outcome_probabilities

"""
    abstract type ProbabilityDistribution <: AbstractVector{Real} end

The abstract type representing a discrete probability distribution.
"""
abstract type ProbabilityDistribution{T} <: AbstractVector{T} end
Base.size(probs::ProbabilityDistribution) = size(probs.distribution)
Base.getindex(probs::ProbabilityDistribution, id::Int) = getindex(probs.distribution, id...)
Base.setindex!(probs::ProbabilityDistribution, v, id::Int) = (probs.distribution[id...] = v)

"""
    abstract type ConditionalDistribution <: AbstractMatrix{Real} end

The abstract type representing a discrete conditional probability distribution.
"""
abstract type ConditionalDistribution{T} <: AbstractMatrix{T} end
Base.size(probs::ConditionalDistribution) = size(probs.distribution)
Base.getindex(probs::ConditionalDistribution, id::Vararg{Int,2}) = getindex(probs.distribution, id...)
Base.setindex!(probs::ConditionalDistribution, v, id::Vararg{Int,2}) = (probs.distribution[id...] = v)

"""
    abstract type JointProabilityDistribution <: AbstractMatrix{Real} end

The abstract type representing a discrete joint probability distribution.
"""
abstract type JointProbabilityDistribution{T} <: AbstractMatrix{T} end
Base.size(probs::JointProbabilityDistribution) = size(probs.distribution)
Base.getindex(probs::JointProbabilityDistribution, id::Vararg{Int,2}) = getindex(probs.distribution, id...)
Base.setindex!(probs::JointProbabilityDistribution, v, id::Vararg{Int,2}) = (probs.distribution[id...] = v)

"""
    is_probability_distribution(
        probabilities :: Vector{<:Real};
        atol=ATOL :: Float64
    ) :: Bool

Returns `true` if the provided vector is a valid probability distribution:
* `sum(probabilities) ≈ 1`
* `probabilities[i] ≥ 0 ∀ i`
"""
function is_probability_distribution(probabilities :: AbstractVector{<:Real}; atol=ATOL :: Float64) :: Bool
    _is_prob(prob :: Real) :: Bool = begin
        (prob ≥ 0) || isapprox(prob, 0, atol=atol)
    end

    if findfirst(!_is_prob, probabilities) !== nothing
        return false
    end

    if !isapprox(sum(probabilities), 1, atol=atol)
        return false
    end

    return true
end
is_probability_distribution(::ProbabilityDistribution) = true
is_probability_distribution(::JointProbabilityDistribution) = true

"""
    is_conditional_distribution( probabilities :: AbstractMatrix{<:Real}; atol=ATOL :: Flaot64 ) :: Bool

Returns `true` if the provided matrix is column stochastic. That is, each column
is a valid probability distribution.
"""
function is_conditional_distribution(conditionals :: Matrix{<:Real}; atol=ATOL :: Float64) :: Bool
    all(col -> is_probability_distribution(col, atol=atol), eachcol(conditionals))
end
is_conditional_distribution(::ConditionalDistribution) = true

"""
    Probabilities( distribution :: Vector{<:Real}; atol=ATOL :: Float64 ) <: ProbabilityDistribution

A struct representing a discrete probability distribution. All elements
in the marginal distribution must be positive and their sum must be one.
"""
struct Probabilities{T} <: ProbabilityDistribution{T}
    distribution :: Vector{T}
    atol :: Float64
    Probabilities(
        distribution :: Vector{<:Real};
        atol=ATOL :: Float64
    ) = is_probability_distribution(distribution, atol=atol) ? new{eltype(distribution)}(distribution,atol) : throw(
        DomainError(distribution, "is not a valid probability distribution.")
    )
end

"""
    Conditionals( distribution :: Matrix{<:Real}; atol=ATOL :: Float64 ) <: ConditionalDistribution

A struct representing a conditional probability distribution function. `Conditionals`
are organized in a matrix with rows correpsonding to outputs and columns corresponding
to inputs. For example if there are ``M`` inputs and ``N`` outputs, the corresponding
conditionals matrix takes the form:

```math
\\begin{bmatrix}
p(1|1) & \\dots & p(1|M) \\\\
\\vdots & \\ddots & \\vdots \\\\
p(N|1) & \\dots & p(N|M) \\\\
\\end{bmatrix}
```
"""
struct Conditionals{T} <: ConditionalDistribution{T}
    distribution :: Matrix{T}
    atol :: Float64
    Conditionals(
        distribution :: Matrix{<:Real};
        atol=ATOL :: Float64
    ) = is_conditional_distribution(
            distribution, atol=atol
        ) ? new{eltype(distribution)}(distribution, atol) : throw(
            DomainError(distribution, "is not a valid conditional probability distribution.")
        )
end

"""
    JointProbabilities(
        distribution :: Matrix{<:Real};
        atol=ATOL :: Float64
    ) <: JointProbabilityDistribution

A struct representing a discrete probability distribution. All elements
in the marginal distribution must be positive and their sum must be one.
For convenience, the `JointProbabilities` can also be constructed by passing in
a `ConditionalDistribution` and a `ProbabilityDistribution`,

    JointProbabilities(
        priors :: AbstractVector{<:Real},
        conditionals :: AbstractMatrix{<:Real};
        atol=ATOL :: Float64
    )

Or, two `ProbabilityDistribution`s can be provided

    JointProbabilities(
        priors1 :: AbstractVector{<:Real},
        priors2 :: AbstractVector{<:Real};
        atol=ATOL :: Float64
    )

A `DomainError` is thrown if the joint probability distribution is invalid.
"""
struct JointProbabilities{T} <: JointProbabilityDistribution{T}
    distribution :: Matrix{T}
    atol :: Float64
    JointProbabilities(
        distribution :: Matrix{<:Real};
        atol=ATOL :: Float64
    ) = is_probability_distribution(
            distribution[:], atol=atol
        ) ? new{eltype(distribution)}(distribution, atol) : throw(
            DomainError(distribution, "is not a valid joint probability distribution.")
        )
    JointProbabilities(
        priors :: AbstractVector{<:Real},
        conditionals :: AbstractMatrix{<:Real};
        atol=ATOL :: Float64
    ) = JointProbabilities(priors' .* conditionals, atol=atol)
    JointProbabilities(
        conditionals :: AbstractMatrix{<:Real},
        priors :: AbstractVector{<:Real};
        atol=ATOL :: Float64
    ) = JointProbabilities(priors' .* conditionals, atol=atol)
    JointProbabilities(
        priors1 :: AbstractVector{<:Real},
        priors2 :: AbstractVector{<:Real};
        atol=ATOL :: Float64
    ) = JointProbabilities(priors2 * priors1', atol=atol)
end

"""
    *(
        C :: ConditionalDistribution,
        P :: ProbabilityDistribution;
        atol=ATOL :: Float64
    ) :: ProbabilityDistribution

Multiplication of a conditional with a prior distribution yields the outcome probabilites.
"""
*(
    conditional :: ConditionalDistribution,
    prior :: ProbabilityDistribution;
    atol=ATOL :: Float64
) :: ProbabilityDistribution = Probabilities(conditional.distribution * prior.distribution, atol=atol)

"""
    outcome_probabilities(
        conditionals::ConditionalDistribution,
        priors :: ProbabilityDistribution;
        atol=ATOL :: Float64
    ) :: ProbabilityDistribution

Returns the probability of each outcome given priors and conditional probabilities.
For convenience, this method can be called with an abstract vector/matrix

    outcome_probabilities(
        conditionals :: AbstractMatrix{<:Real},
        priors :: AbstractVector{<:Real};
        atol=ATOL :: Float64
    )

Furthermore, the ordering of the arguments can be reversed.
A `DomainError` is thrown if the constructed vector is not a valid probability
distribution.
"""
function outcome_probabilities(
    conditionals :: ConditionalDistribution,
    priors :: ProbabilityDistribution;
    atol=ATOL :: Float64
) :: ProbabilityDistribution
    *(conditionals,priors,atol=atol)
end
function outcome_probabilities(
    priors :: ProbabilityDistribution,
    conditionals :: ConditionalDistribution;
    atol=ATOL :: Float64
) :: ProbabilityDistribution
    *(conditionals,priors, atol=atol)
end
function outcome_probabilities(
    conditionals :: AbstractMatrix{<:Real},
    priors :: AbstractVector{<:Real};
    atol=ATOL :: Float64
) :: ProbabilityDistribution
    Probabilities(conditionals*priors, atol=atol)
end
function outcome_probabilities(
    priors :: AbstractVector{<:Real},
    conditionals :: AbstractMatrix{<:Real};
    atol=ATOL :: Float64
) :: ProbabilityDistribution
    Probabilities(conditionals*priors, atol=atol)
end


include("matrixgen.jl")
include("reptileutils.jl")

## Tools: functions for angles and sides
function findconstraints(A1, B1, n, digits = 12)
  #OBSOLTETE
  #returns a vector with ones in the locations where additional constraints would  fully constrain the system

  l = size(A)[2] #number of columns in A
  output = zeros(n)
  i = 0 #loop counter
  r = rank([A B])
  while r < n && i < n #underconstrained and not out of variables
    i += 1 #increment
    c = zeros(1, l) #new row
    c[i] = 1 #add constraint x_i = 0
    r1 = rank([A B; c 0]) #rank with new constraint x_i = 0
    if r == r1 #nothing changed
      continue #try next variable
    else #rank has increased. rank decrease should be impossible
      if size(A)[1] == l #sqaure. could be singular
        X = [A; zeros(1, l)] \ [B; 0.0] #add an empty row to trick the algorithm out of square matrix, and solve
      else
        X = A \ B #solve
      end
      if floateq(A * X, B, digits) == false #solution failed
        continue #try next variable
      end
      #new constraint is valid
      A = [A; c] #add new constraint
      B = [B; 0]
      r = rank([A B])
      output[i] = 1.0 #save this as a good constraint
    end
  end
  return output #return good constraints, or all zeros if it cannot be constrained
end

function buildanglesys(matint)
  #returns A and B of the equation AX=B for the given angle matrix
  mat = convert(Array{Angle, 2}, matint) #convert
  n = size(mat)[1]
  a = mat[:,2:end] - eye(Angle, n, n) #move lhs over
  A = [a; ones(Angle, 1, n)] # sum(theta) = pi*(n-2) constraint
  B = [-mat[:,1] * Angle(0//1, 1//1); Angle(0//1, 1//1) * (n-2)//1] #straight angle conditions; pi*(n-2 constraint)
  return A, B, n
end

function buildsidesys(matint::Array{Int, 2}, m::Int)
  #returns A and B of the equation AX=B for the given side matrix with s1 = 1
  mat = convert(Array{Side, 2}, matint, m)
  n = size(mat)[1]
  a = mat - sideeye(Side, n, m) * Side(0//1, 1//1, m)
  A = [one(Side, m) zeros(Side, (1, n-1), mvalue = m)
    zeros(Side, (n, 1), mvalue = m) a[:, 2:end]]::Array{Side, 2} #with fake first row and col forcing s1=1
  B = [one(Side, m); -a[:, 1]] #with fake first row forcing s1=1
  return A, B, n
end

function isvalidsol(angles::Array{Angle,1})
    #checks if the `angles` vector of polygon angles are valid
    #only checking the p-component since that should be all that is possible

  for i in angles
    if i.p <= 0//1 #negative angles (can't be changed because the system has no d.o.f.)
        return false
    end
    if i.p >= 2//1 #angle over 2pi (can't be changed because the system has no d.o.f.)
        return false
    end
    if i.p == 1//1 #there is a straight angle
        return false
    end
  end

  return true #all tests pass
end

function isvalidsol(sides::Array{Side,1})
    #tests whether the vector of side lengths `s` are valid
    m = findm(sides)

    if minimum(sides) <= zero(Side, m) #minimum individual side lenth
        return false
    end
    #longest side is more than the sum of the other sides (s1 is length 1)
    if 2//1 * maximum(sides) >= sum(sides)  #max side >= all sides - max side
        return false
    end
    return true #all test pass
end

function solvemat(mat, m::Int, isangle::Bool)
  #solves the given angle matrix if possible
  #returns rref form of the solution matrix
  #returns marix of zeros and rank = n+1 if there is no solution
  #returns rank as 2nd argument

  if isangle #angle matrix
    A, B, n = buildanglesys(mat) #matrices for A*X=B
  else #side matrix
    A, B, n = buildsidesys(mat, m) #matrices for A*X=B
  end

  rrefmat = rref([A B]) #row reduce

  if isangle
    rank = countnz(sum(rrefmat, 2)) #rank of matrix
  else #side
    #ifferent because of special side zero functions
    rank = countnz(sum(rrefmat .!= zero(Side, m), 2)) #rank of matrix
  end

  if rank < n #underconstrained.
    return rrefmat[1:n, :], rank
  end

  if rank == n #fully constrained
    if isvalidsol(rrefmat[1:n,end]) #last column is solutions
      return rrefmat[1:n, :], rank
    end
  end

  return zeros(rrefmat[1:n, :]), n+1 #overconstrained or invalid solutions
end


using PyPlot
Tmax = 5
npoints = 10
g(n) = range(1/Tmax^n,1., length=npoints)
h(n) = range(0, n*log10(Tmax), length=npoints)
j(n) = range(0, 1-exp(-n*(Tmax-1)), length=npoints)

##
figure()
for n in range(0, 5, length=11)
    #plot(g(n).^(-1/n), g(n), label="T^$(-n)", ".-")
    plot(1 .- 1/n*log.(1 .- j(n)), j(n), "x-", label="1-exp(-$n(T -1))")
end
n = 1
plot(10 .^(1/n*h(n)), h(n), label="log T", "v-k")
xlabel("T")
ylabel("equally spaced f(T)")
legend(bbox_to_anchor=(1.04,1), borderaxespad=0)
plt.tight_layout()


#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   Compute the satellite position.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

export satellite_position_i

"""
    satellite_position_i(a::Number, e::Number, i::Number, RAAN::Number, w::Number, f::Number)

Compute the satellite position in the Earth-Centered Inertial (ECI) reference
frame given the orbital elements `a`, `e`, `i`, `RAAN`, `w`, and `f`.

Notice that the ECI frame used will be the same as the frame of the orbital
elements.

# Args

* `a`: Semi-major axis.
* `e`: Eccentricity.
* `i`: Inclination [rad].
* `RAAN`: Right ascension of the ascending node [rad].
* `w`: Argument of perigee [rad].
* `f`: True anomaly [rad].

# Returns

* The satellite position vector represented in the ECI reference frame.
* The unit vector perpendicular to the satellite position vector that lies on
  the orbit plane represented in the ECI reference frame.

# Remarks

The satellite position vector will have the same unit of the semi-major axis.

"""
function satellite_position_i(a::Number, e::Number, i::Number, RAAN::Number,
                              w::Number, f::Number)
    # Compute the radius from the focus.
    norm_r = a*(1-e^2)/(1+e*cos(f))

    # Let s be the coordinate system in which:
    #     - The X axis points towards the satellite;
    #     - The Z axis is normal to the orbit plane (right-hand direction);
    #     - The Y axis completes a right-hand coordinate system.
    #
    # Thus, the satellite vector represented in the s coordinate frame is:
    r_s = SVector{3}(1,0,0)*norm_r

    # rt is the versor perpendicular to the r vector that lies on the orbit
    # plane.
    rt_s = SVector{3}(0,1,0)

    # Compute the matrix that rotates from the S coordinate frame to the
    # Inertial coordinate Frame.
    Dis = angle_to_dcm(RAAN, i, w+f, :ZXZ)'

    # Compute the satellite vector represented in the Inertial coordinate
    # frame.
    r_i = Dis*r_s

    # Compute unit vector `rt` represented in the Inertial coordinate frame.
    rt_i = Dis*rt_s

    # Return.
    return r_i, rt_i
end


using GNSSTools
using Statistics
using ProgressMeter
using JLD
using PyPlot
pygui(true)

h_parms = [1.9e-6, 2e-6, 9.7e-8, 9.1e-11, 9.9e-15]./16.3676e6^2  # TCXO
CN0 = 45.  # dB⋅Hz
f_d = 0
nADC = 16
code_start_idx = 1000
iterations = 100
Ts = Array(1e-3:1e-3:100e-3)
# t_length = maximum(Ts)

# L1 C/A
f_s = 2.048e6
sigtype = define_l1ca_code_type(; sig_freq=1*L1_freq)
# generatesignal!(signal)
l1ca_SNRs = zeros(length(Ts))
p = Progress(length(Ts)*iterations, 1, "Processing L1 C/A...")
for i in 1:length(Ts)
    SNR = 0
    T = Ts[i]
    signal = definesignal(sigtype, f_s, T;
                          prn=1, code_start_idx=code_start_idx, f_d=f_d, fd_rate=0,
                          CN0=CN0, include_phase_noise=false, f_if=1.25e6,
                          include_thermal_noise=true, include_adc=false, nADC=nADC,
                          include_carrier=true, receiver_h_parms=h_parms,
                          skip_noise_generation=false)
    replica = definesignal(sigtype, f_s, T; f_d=f_d, allocate_noise_vectors=false,
                           skip_noise_generation=true)
    generatereplica!(replica)
    for j in 1:iterations
        # definesignal!(signal; new_thermal_noise=true, new_phase_noise=false)
        # generatesignal!(signal)
        generatereplica!(signal)
        A = sqrt(2*GNSSTools.k*535*10^(CN0/10))
        σₙ = sqrt(GNSSTools.k*2.046e6*535)
        signal.data .= A.*signal.data .+ σₙ.*randn(signal.sample_num)
        corr_result = fft_correlate(signal.data[1:replica.sample_num], replica.data)
        corr_result_fft = fft(corr_result)
        PS = 2*maximum(abs2.(corr_result_fft))
        # PS = abs2(corr_result[code_start_idx])
        PN = (sum(abs2.(corr_result_fft)) - PS) / (length(corr_result_fft) - 2)
        SNR += PS/PN
        next!(p)
    end
    l1ca_SNRs[i] = 10log10(SNR/iterations)
end

fig = figure()
plot(Ts.*1000, l1ca_SNRs, label="L1 C/A")
xlabel("Integration Times (ms)")
ylabel("Acquired SNR (dB)")
legend()
xscale("log")

save("ch2_snr_vs_integration_time.jld", "l1ca_SNRs", l1ca_SNRs)
savefig("figures/ch2_snr_vs_integration_time.pdf", dpi=300)


module LVAChannels

export Stim, run_lva, LVAParam

using Sundials
using SIUnits
using SIUnits.ShortUnits

const mS = Milli*Siemens

immutable LVAParam
    c_m       # membrane spacific capacitance
    gbar      # Ca channel cross-membrane conductivity
    eca       # Ca channel reversal potential
    gl        # leak conductivity
    el        # leak reversal potential
    q10_1     # rate scaling for 'm' activation gate
    q10_2     # rate scaling for 'dsh' inactivation gate
    vrest     # (as hoc) resting potential

    # constructor with default values, with q10 values
    # corresponding to the body temperature current-clamp experiments
    LVAParam(;
        c_m    = 0.01F*m^-2,
        gbar   = 0.2mS*cm^-2,
        eca    = 120mV,
        gl     = .1mS*cm^-2,
        el     = -65mV,
        q10_1  = 5,
        q10_2  = 3,
#        vrest  = -61.47mV
        vrest  = -63mV
    ) = new(c_m, gbar, eca, gl, el, q10_1, q10_2, vrest)

end

immutable Stim
    t0        # start time of stimulus
    t1        # stop time of stimulus
    i_e       # stimulus current density

    Stim() = new(0s, 0s, 0A/m^2)
    Stim(t0, t1, i_e) = new(t0, t1, i_e)
end

# 'm' activation gate
function m_lims(v, q10)
    quotient = 1+exp(-(v+63mV)/7.8mV)
    mtau = 1ms*(1.7+exp(-(v+28.8mV)/13.5mV))/(q10*quotient)
    minf = 1/quotient
    return mtau, minf
end

# 'dsh' inactivation gate:
# d <-> s (alpha2, beta2)
# s <-> h (alpha1, beta1)
# subject to s=1-h-d
function dsh_lims(v, q10)
    k = sqrt(0.25+exp((v+83.5mV)/6.3mV))-0.5
    alpha1 = q10*exp(-(v+160.3mV)/17.8mV)/1ms
    beta1 = alpha1*k

    tau2 = 240ms/(1+exp((v+37.4mV)/30mV))/q10
    alpha2 = 1/tau2/(1+k)
    beta2 = alpha2*k

    hinf = 1/(1+k+k^2)
    dinf = k^2*hinf

    return alpha1, beta1, alpha2, beta2, dinf, hinf
end

# Choose initial conditions for the system such that the gating variables
# are at steady state for the user-specified voltage v.
function initial_conditions(v, q10_1, q10_2)
    mtau, minf = m_lims(v, q10_1)
    alpha1, beta1, alpha2, beta2, dinf, hinf = dsh_lims(v, q10_2)

    return (v, minf, dinf, hinf)
end

# Given time t and state (v, m, d, h),
# return (vdot, mdot, ddot, hdot).
function f(t, state; p=LVAParam(), istim=0A/m^2)
    v, m, d, h = state

    ica = p.gbar*m^3*h*(v - p.eca)
    il = p.gl*(v - p.el)
    itot = ica + il - istim

    # Calculate the voltage dependent rates for the gating variables.
    mtau, minf = m_lims(v, p.q10_1)
    alpha1, beta1, alpha2, beta2, dinf, hinf = dsh_lims(v, p.q10_2)

    mdot = (minf-m)/mtau
    hdot = alpha1*(1-h-d)-beta1*h
    ddot = beta2*(1-h-d)-alpha2*d

    return (-itot/p.c_m, mdot, ddot, hdot)
end

function make_range(t_start, dt, t_end)
    r = collect(t_start: dt: t_end)
    if length(r)>0 && r[length(r)]<t_end
        push!(r, t_end)
    end
    return r
end

function run_lva(t_end; stim=Stim(), param=LVAParam(), sample_dt=0.01ms)
    v_scale = 1V
    t_scale = 1s

    v0, m0, h0, d0 = initial_conditions(param.vrest, param.q10_1, param.q10_2)
    y0 = [ v0/v_scale, m0, h0, d0 ]

    fbis(t, y, ydot, istim) = begin
        vdot, mdot, hdot, ddot =
            f(t*t_scale, (y[1]*v_scale, y[2], y[3], y[4]), istim=istim, p=param)

        ydot[1], ydot[2], ydot[3], ydot[4] =
            vdot*t_scale/v_scale, mdot*t_scale, hdot*t_scale, ddot*t_scale

        return Sundials.CV_SUCCESS
    end

    fbis_nostim(t, y, ydot) = fbis(t, y, ydot, 0A/m^2)
    fbis_stim(t, y, ydot) = fbis(t, y, ydot, stim.i_e)

    # Ideally would run with vector absolute tolerance to account for v_scale,
    # but this would prevent us using the nice cvode wrapper.

    res = []
    samples = []

    t1 = clamp(stim.t0, 0s, t_end)
    if t1>0s
        ts = make_range(0s, sample_dt, t1)
        r = Sundials.cvode(fbis_nostim, y0, map(t->t/t_scale, ts), abstol=1e-6, reltol=5e-10)
        y0 = vec(r[size(r)[1], :])
        push!(res, r)
        push!(samples, ts)
    end
    t2 = clamp(stim.t1, t1, t_end)
    if t2>t1
        ts = make_range(t1, sample_dt, t2)
        r = Sundials.cvode(fbis_stim, y0, map(t->t/t_scale, ts), abstol=1e-6, reltol=5e-10)
        y0 = vec(r[size(r)[1], :])
        push!(res, r)
        push!(samples, ts)
    end
    if t_end>t2
        ts = make_range(t2, sample_dt, t_end)
        r = Sundials.cvode(fbis_nostim, y0, map(t->t/t_scale, ts), abstol=1e-6, reltol=5e-10)
        y0 = vec(r[size(r)[1], :])
        push!(res, r)
        push!(samples, ts)
    end

    res = vcat(res...)
    samples = vcat(samples...)

    # Use map here because of issues with type deduction with arrays and SIUnits.
    return samples, map(v->v*v_scale, res[:, 1]), res[:, 2], res[:, 3], res[:, 4]
end

end # module LVAChannels


using ModelingToolkit, LinearAlgebra, SparseArrays

# Define the constants for the PDE
const α₂ = 1.0
const α₃ = 1.0
const β₁ = 1.0
const β₂ = 1.0
const β₃ = 1.0
const r₁ = 1.0
const r₂ = 1.0
const _DD = 100.0
const γ₁ = 0.1
const γ₂ = 0.1
const γ₃ = 0.1
const N = 8
const X = reshape([i for i in 1:N for j in 1:N],N,N)
const Y = reshape([j for i in 1:N for j in 1:N],N,N)
const α₁ = 1.0.*(X.>=4*N/5)

const Mx = Tridiagonal([1.0 for i in 1:N-1],[-2.0 for i in 1:N],[1.0 for i in 1:N-1])
const My = copy(Mx)
Mx[2,1] = 2.0
Mx[end-1,end] = 2.0
My[1,2] = 2.0
My[end,end-1] = 2.0

# Define the initial condition as normal arrays
@variables du[1:N,1:N,1:3] u[1:N,1:N,1:3] MyA[1:N,1:N] AMx[1:N,1:N] DA[1:N,1:N]

# Define the discretized PDE as an ODE function
function f(du,u,p,t)
   A = @view  u[:,:,1]
   B = @view  u[:,:,2]
   C = @view  u[:,:,3]
  dA = @view du[:,:,1]
  dB = @view du[:,:,2]
  dC = @view du[:,:,3]
  mul!(MyA,My,A)
  mul!(AMx,A,Mx)
  @. DA = _DD*(MyA + AMx)
  @. dA = DA + α₁ - β₁*A - r₁*A*B + r₂*C
  @. dB = α₂ - β₂*B - r₁*A*B + r₂*C
  @. dC = α₃ - β₃*C + r₁*A*B - r₂*C
end

f(du,u,nothing,0.0)

multithreadedf = eval(ModelingToolkit.build_function(du,u,parallel=ModelingToolkit.MultithreadedForm())[2])

MyA = zeros(N,N);
AMx = zeros(N,N);
DA = zeros(N,N);
# Loop to catch syncronization issues
for i in 1:100
   _du = rand(N,N,3)
   _u = rand(N,N,3)
   multithreadedf(_du,_u)
   _du2 = copy(_du)
   f(_du2,_u,nothing,0.0)
   @test _du ≈ _du2
end

#=
jac = sparse(ModelingToolkit.jacobian(vec(du),vec(u)))
fjac = eval(ModelingToolkit.build_function(jac,u,parallel=ModelingToolkit.SerialForm())[2])
multithreadedfjac = eval(ModelingToolkit.build_function(jac,u,parallel=ModelingToolkit.MultithreadedForm())[2])

u = rand(N,N,3)
J = similar(jac,Float64)
fjac(J,u)

J2 = similar(jac,Float64)
multithreadedfjac(J2,u)
@test J ≈ J2

using FiniteDiff
J3 = Array(similar(jac,Float64))
FiniteDiff.finite_difference_jacobian!(J2,(du,u)->f!(du,u,nothing,nothing),u)
maximum(J2 .- Array(J)) < 1e-5
=#

using Distributed
addprocs(4)
distributedf = eval(ModelingToolkit.build_function(du,u,parallel=ModelingToolkit.DistributedForm())[2])

using Dagger
daggerf = eval(ModelingToolkit.build_function(du,u,parallel=ModelingToolkit.DaggerForm())[2])

jac = sparse(ModelingToolkit.jacobian(vec(du),vec(u),simplify=false))
serialjac = eval(ModelingToolkit.build_function(vec(jac),u)[2])
multithreadedjac = eval(ModelingToolkit.build_function(vec(jac),u,parallel=ModelingToolkit.MultithreadedForm())[2])
distributedjac = eval(ModelingToolkit.build_function(vec(jac),u,parallel=ModelingToolkit.DistributedForm())[2])
daggerjac = eval(ModelingToolkit.build_function(vec(jac),u,parallel=ModelingToolkit.DaggerForm())[2])

MyA = zeros(N,N)
AMx = zeros(N,N)
DA  = zeros(N,N)
_du = rand(N,N,3)
_u = rand(N,N,3)

f(_du,_u,nothing,0.0)
multithreadedf(_du,_u)
#distributedf(_du,_u)
#daggerf(_du,_u)

#=
using BenchmarkTools
@btime f(_du,_u,nothing,0.0)
@btime multithreadedf(_du,_u)
@btime distributedf(_du,_u)
@btime daggerf(_du,_u)

_jac = similar(jac,Float64)
@btime serialjac(_jac,_u)
@btime multithreadedjac(_jac,_u)
@btime distributedjac(_jac,_u)
@btime daggerjac(_jac,_u)
=#


#==============================================================================
    Code for solving the Hamiltonian Jacobi Bellman for
	   a Ramsey model with a Diffusion process

	Orignially based on Matlab code from Ben Moll:
        http://www.princeton.edu/~moll/HACTproject.htm

        Updated to Julia 1.0.0
==============================================================================#

using Distributions, Plots, SparseArrays, LinearAlgebra, DifferentialEquations

using Random

Random.seed!(12676673)

include("B_Switch.jl")

γ= 2.0 #gamma parameter for CRRA utilit
ρ = 0.05 #the discount rate
α = 0.3 # the curvature of the production function (cobb-douglas)
δ = 0.05 # the depreciation rate

# Create a continuous time OrnsteinUhlenbeckProcess
var = 0.07
T = 10001 # Forecasting periods/ length of process
T_obs= 100# set the number of inital observations
global dt_inv = 5# Does this dt matter?
global dt= 1/dt_inv
OU_process=zeros((T+T_obs)*dt_inv)
OU_process[1]=exp(var/2)

# Z our state variable follows a stochastic process with the given parameters
corr = 0.9
θ = -log(corr)
σ = 1.0
ε_OU = rand(Normal(0,1), (T+T_obs)*dt_inv).*sqrt(dt)

for i in 1:((T+T_obs)*dt_inv-1)
	OU_process[i+1]=(1-θ*dt)*OU_process[i].+σ*ε_OU[i]
end

# Now plot the process
plot(OU_process[:], grid=false,
		xlabel="Time", ylabel="z",
		legend=false, color="red", title="\$ \\textrm{Ornstein-Uhlenbeck Process for } \\log(z_t) \$")
png("OU_Process")

# Step up grid space
var = mean(OU_process)
μ_z = exp(var/2)

#=============================================================================
 	k our capital follows a process that depends on z,

	using the regular formula for capital accumulation
	we would have:
		(1+ρ)k_{t+1} = k_{t}⋅f'(k_{t}) + (1-δ)k_{t}
	where:
		f(k_{t}) = z⋅k^{α} so f'(k_{t}) = (α)⋅z⋅k^{α-1}
	so in steady state where k_{t+1} = k_{t}
		(1+ρ)k = α⋅z⋅k^{α} + (1-δ)k
		k = [(α⋅z)/(ρ+δ)]^{1/(1-α)

=============================================================================#
#K_starting point, for mean of z process

k_st = ((α*exp(μ_z))/(ρ+δ))^(1/(1-α))

# create the grid for k
H = 100 #number of points on grid
k_min = 0.3*k_st # min value
k_max = 9.0*k_st # max value
k = LinRange(k_min, k_max, H)
dk = (k_max-k_min)/(H-1)

# create the grid for z
J = 40
z_min = μ_z*0.8
z_max = μ_z.*1.2
z = LinRange(z_min, z_max, J)
dz = (z_max-z_min)/(J-1)
dz_sq = dz^2

# Check the pdf to make sure our grid isn't cutting off the tails of
	#= our distribution
y = pdf.(LogNormal(0, var), z)
plot(z,y, grid=false,
		xlabel="z", ylabel="Probability",
		legend=false, color="purple", title="PDF of z")
png("PDF_of_z")=#

#create matrices for k and z
z= convert(Array, z)'
kk = k*ones(1,J)
zz = ones(H,1)*z

# use Ito's lemma to find the drift and variance of our optimization equation

μ = (-θ*log.(z).+σ.^2/2).*z # the drift from Ito's lemma
Σ_sq = σ.^2 .*z.^2#the variance from Ito's lemma
global B_switch_SS = B_switch(μ, Σ_sq, dz, H)

max_it = 100
ε = 0.1^(6)
Δ = 1000

# set up all of these empty matrices
Vaf, Vab, Vzf, Vzb, Vzz = [zeros(H,J) for i in 1:6]

# Now it's time to solve the model, first put in a guess for the value function
v0 = (zz.*kk.^α).^(1-γ)/(1-γ)/ρ
v=v0

maxit= 30 #set number of iterations (only need 6 to converge)
dist = [] # set up empty array for the convergence criteria

for n = 1:maxit
    V=v

    #Now set up the forward difference

    Vaf[1:H-1,:] = (V[2:H, :] - V[1:H-1,:])/dk
    Vaf[H,:] = (z.*k_max.^α .- δ.*k_max).^(-γ) # imposes a constraint

    #backward difference
    Vab[2:H,:] = (V[2:H, :] - V[1:H-1,:])/dk
    Vab[1,:] = (z.*k_min.^α .- δ.*k_min).^(-γ)

    #I_concave = Vab .> Vaf # indicator for whether the value function is concave

    # Consumption and savings functions
    cf = Vaf.^(-1/γ)
    sf = zz .* kk.^α - δ.*kk - cf

    # consumption and saving backwards difference

    cb = Vab.^(-1.0/γ)
    sb = zz .* kk.^α - δ.*kk - cb
    #println(sb)
    #consumption and derivative of the value function at the steady state

    c0 = zz.*kk.^α - δ.*kk
    Va0 = c0.^(-γ)

    # df chooses between the forward or backward difference

    If = sf.>0 # positive drift will ⇒ forward difference
    Ib = sb.<0 # negative drift ⇒ backward difference
    I0=(1.0.-If-Ib) # at steady state

    Va_upwind = Vaf.*If + Vab.*Ib + Va0.*I0 # need to include SS term

    global c = Va_upwind.^(-1/γ)
    u = (c.^(1-γ))/(1-γ)

    # Now to constuct the A matrix
    X = -min.(sb, zeros(H,1))/dk
    #println(X)
    Y = -max.(sf, zeros(H,1))/dk + min.(sb, zeros(H,1))/dk
    Z = max.(sf, zeros(H,1))/dk

    updiag = 0
       for j = 1:J
           updiag =[updiag; Z[1:H-1,j]; 0]
       end
    updiag =(updiag[:])

    centerdiag=reshape(Y, H*J, 1)
    centerdiag = (centerdiag[:]) # for tuples

   lowdiag = X[2:H, 1]
       for j = 2:J
           lowdiag = [lowdiag; 0; X[2:H,j]]
       end
   lowdiag=(lowdiag)

   # spdiags in Matlab allows for automatic trimming/adding of zeros
       # spdiagm does not do this
  AA = sparse(Diagonal(centerdiag))+ [zeros(1, H*J); sparse(Diagonal(lowdiag)) zeros(H*J-1,1)] + sparse(Diagonal(updiag))[2:end, 1:(H*J)] # trim first element

  A = AA + B_switch_SS
  B = (1/Δ + ρ)*sparse(I, H*J, H*J) - A

  u_stacked= reshape(u, H*J, 1)
  V_stacked = reshape(V,H*J, 1)

  b = u_stacked + (V_stacked./Δ)

  V_stacked = B\b

  V = reshape(V_stacked, H, J)

  V_change = V-v

  global v= V

  # need push function to add to an already existing array
  push!(dist, findmax(abs.(V_change))[1])
  if dist[n].< ε
      println("Value Function Converged Iteration=")
      println(n)
      break
  end

end

v1=v

# calculate the savings for kk
ss = zz.*kk.^α - δ.*kk - c

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

		Now, we have found the steady state,
				we can move on to misspecification and forecasting

		Our Agent will now misspecify θ the drift parameter from our process
		They will then forecast an AR(1) and use this to update their
		forecast of the process, since ϕ = 1-θ.

		Therefore θ = 1-ϕ

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

# Our Agent will now misspecify η the drift parameter from our process
# They will then forecast an AR(1) and use this to update

# Use the initai points to form a OLS estimate of our coefficients

X = [ones(T_obs,1) OU_process[1:T_obs]]
Y = OU_process[1+dt_inv:T_obs+dt_inv]
ϵ_OU = ε_OU[1:T_obs]*sqrt(dt)

η_g = (X'X)^(-1) * X'Y
θ_g = (1-η_g[2])
const_g =η_g[1]

R_g = X'*X*dt
#R_g = [1 0 ; 0 1]#Use identity matrix for R_g?
guesses_θ, guesses_const =[zeros(1,T-1) for i in 1:2]

Value_functions=[]

v =v0

for t = 1:T-1
	global v=v0
	global θ_g = θ_g
	global const_g = const_g
	global R_g = R_g
	guesses_θ[1,t] = θ_g
	guesses_const[1,t]=const_g

	μ = (-θ_g*log.(z).+σ.^2/2).*z # the drift from Ito's lemma
	Σ_sq = σ.^2 .*z.^2#the variance from Ito's lemma

	global B_switch_G = B_switch(μ, Σ_sq, dz, H)

	for n = 1:maxit
	    V=v
	    #Now set up the forward difference
	    Vaf[1:H-1,:] = (V[2:H, :] - V[1:H-1,:])/dk
	    Vaf[H,:] = (z.*k_max.^α .- δ.*k_max).^(-γ) # imposes a constraint

	    #backward difference
	    Vab[2:H,:] = (V[2:H, :] - V[1:H-1,:])/dk
	    Vab[1,:] = (z.*k_min.^α .- δ.*k_min).^(-γ)

	    #I_concave = Vab .> Vaf # indicator for whether the value function is concave

	    # Consumption and savings functions
	    cf = Vaf.^(-1/γ)
	    sf = zz .* kk.^α - δ.*kk - cf

	    # consumption and saving backwards difference

	    cb = Vab.^(-1.0/γ)
	    sb = zz .* kk.^α - δ.*kk - cb
	    #println(sb)
	    #consumption and derivative of the value function at the steady state

	    c0 = zz.*kk.^α - δ.*kk
	    Va0 = c0.^(-γ)

	    # df chooses between the forward or backward difference

	    If = sf.>0 # positive drift will ⇒ forward difference
	    Ib = sb.<0 # negative drift ⇒ backward difference
	    I0=(1.0.-If-Ib) # at steady state

	    Va_upwind = Vaf.*If + Vab.*Ib + Va0.*I0 # need to include SS term

	    global c = Va_upwind.^(-1/γ)
	    u = (c.^(1-γ))/(1-γ)

	    # Now to constuct the A matrix
	    X = -min.(sb, zeros(H,1))/dk
	    #println(X)
	    Y = -max.(sf, zeros(H,1))/dk + min.(sb, zeros(H,1))/dk
	    Z = max.(sf, zeros(H,1))/dk

	    updiag = 0
	       for j = 1:J
	           updiag =[updiag; Z[1:H-1,j]; 0]
	       end
	    updiag =(updiag[:])

	    centerdiag=reshape(Y, H*J, 1)
	    centerdiag = (centerdiag[:]) # for tuples

	   lowdiag = X[2:H, 1]
	       for j = 2:J
	           lowdiag = [lowdiag; 0; X[2:H,j]]
	       end
	   lowdiag=(lowdiag)

	   # spdiags in Matlab allows for automatic trimming/adding of zeros
	       # spdiagm does not do this
	  AA = sparse(Diagonal(centerdiag))+ [zeros(1, H*J); sparse(Diagonal(lowdiag)) zeros(H*J-1,1)] + sparse(Diagonal(updiag))[2:end, 1:(H*J)] # trim first element

	  A = AA + B_switch_G
	  B = (1/Δ + ρ)*sparse(I, H*J, H*J) - A

	  u_stacked= reshape(u, H*J, 1)
	  V_stacked = reshape(V,H*J, 1)

	  b = u_stacked + (V_stacked./Δ)

	  V_stacked = B\b

	  V = reshape(V_stacked, H, J)

	  V_change = V-v

	  Vaf[1:H-1,:] = (V[2:H, :] - V[1:H-1,:])/dk
	  Vaf[H,:] = (z.*k_max.^α .- δ.*k_max).^(-γ) # imposes a constraint

	  #throw out bad Value functions, good for debugging not for results
	  #if findmin(Vaf)[1]>0
	  	global v= V
	#end

	  # need push function to add to an already existing array
	  push!(dist, findmax(abs.(V_change))[1])
	  if dist[n].< ε
	      println("Value Function Converged Iteration=")
	      println(n)
	      break
	  end
	end

	push!(Value_functions, v)
	x = OU_process[T_obs+(t-1)*dt_inv+1:T_obs+(t)*dt_inv]
	y= OU_process[T_obs+(t-1)*dt_inv+2:T_obs+(t)*dt_inv+1]
	#ϵ_OU = ε_OU[T_obs+(t-1)*dt_inv+1:T_obs+(t)*dt_inv]

	global ϕ_g = [const_g; 1-θ_g*dt]# prediction
	# trues values are ϕ = [0.0; 1-θ; σ]
	W = [ones(dt_inv,1) x]

	#Better to loop or treat as a vector

	global R_g = R_g + (1/t).*(W'*W*dt - R_g)
	global ϕ_g = ϕ_g + (1/t).*R_g^(-1)*W'*(y-W*ϕ_g)*dt

	#=for j in 1:dt_inv
		global R_g = R_g + (1/t) .*(W[j,:]*W[j,:]' - R_g)
		global ϕ_g = ϕ_g + (1/t) .*R_g^(-1)*W[j,:]*(y[j,:][1]-ϕ_g'*W[j,:])
	end=#

	global const_g = ϕ_g[1]
	global θ_g = (1-ϕ_g[2])*dt_inv

	println("loop:$(t)")
end

plot(guesses_θ[:], label="Estimates",
title="\$ \\textrm{Estimate of } \\theta \\textrm{ over time}\$", legend=:bottomright)
plot!(θ.*ones(T,1), label="True value", legend=:topright, ylims=(0,.5))
png("Theta_estimates")

plot(guesses_const[:], label="Estimates",
title="\$ \\textrm{Estimate of the constant}\\textrm{ over time}\$", legend=:bottomright)
plot!(zeros(T,1), label="True value", legend=:topright)
png("const_estimates")

all_drifts=(-guesses_θ[:].*log.(z) .+(σ.^2)/2 ).*z

plot(all_drifts[:,20], label="Estimates",title="Drift for Median Z", legend=:bottomright)
plot!(μ[20]*ones(T,1), label="True value")
#png("drift_estimates_median")

plot(all_drifts[:,1], label="Estimates",title="Drift for Lowest Z", legend=:bottomright)
plot!(μ[1]*ones(T,1), label="True value")
#png("drift_estimates_low")

plot(all_drifts[:,end], label="Estimates",title="Drift for Highest Z", legend=:bottomright)
plot!(μ[end]*ones(T,1), label="True value")
#png("drift_estimates_high")

plot(Value_functions[1][:,20], label="Period 1",
	title="Value Functions For median Z", grid=false,
	legend=:bottomright, xlabel="k", line=:dashdot)
plot!(Value_functions[100][:,20], label="Period 100", line=:dot)
plot!(Value_functions[500][:,20], label="Period 500",line=:dashdotdot)
plot!(Value_functions[end][:,20], label="Period 10,000", color=:orange)
plot!(v1[:,20], label="True Value", line=:dash, color=:black)
png("Value_med_z")

plot(Value_functions[1][50,:], label="Period 1",
title="Value Functions For median K", grid=false,
 legend=:topleft, xlabel="z",line=:dashdot)
plot!(Value_functions[100][50,:], label="Period 100",line=:dot)
plot!(Value_functions[1000][50,:], label="Period 1,000", color=:purple, line=:dashdotdot)
plot!(Value_functions[end][50,:], label="Period 10,000", color=:orange,linewidth = 1.25)
plot!(v1[50,:], label="True Value", line=:dash, color=:black)
png("Value_med_k")


function matview(A::AbstractArray{T,N}) where {T,N}
    n = size(A,1)
    m = size(A,2)
    L = n*m
    reinterpret(SArray{Tuple{n,m},T,2,L}, reshape(A, (L, map(i->size(A,i),3:N)...,)))
end

extractmat(A::AbstractArray{SMatrix{N,M,T}}) where {N,M,T} =
    reshape(reinterpret(T, A), (N,M, size(A)...,))

# Array of Vectors
vecview(A::AbstractArray{T,N}) where {T,N} =
    reinterpret(SVector{size(A,1),T},A)
extractvec(A::AbstractArray{SVector{L,T}}) where {L,T} =
    reshape(reinterpret(T, A), (L, size(A)...,))


# From https://github.com/unkcpz/LibSymspg.jl/blob/f342e72/test/runtests.jl#L83-L85
@testset "Niggli reduce" begin
    lattice = [
        4.0 0.0 0.0
        20.0 2.0 0.0
        0.0 0.0 12.0
    ]
    @test niggli_reduce(lattice, 1e-3) ≈ [
        0.0 4.0 0.0
        -2.0 0.0 0.0
        0.0 0.0 12.0
    ]
end
# From https://github.com/unkcpz/LibSymspg.jl/blob/f342e72/test/runtests.jl#L87-89
@testset "Delaunay reduce" begin
    lattice = [
        4.0 0.0 0.0
        20.0 2.0 0.0
        0.0 0.0 12.0
    ]
    @test delaunay_reduce(lattice, 1e-3) ≈ [
        0.0 -4.0 0.0
        2.0 0.0 0.0
        0.0 0.0 12.0
    ]
end


export Tropical, TropicalF64, TropicalF32, TropicalF16, content

neginf(::Type{T}) where T = typemin(T)
neginf(::Type{T}) where T<:AbstractFloat = typemin(T)
neginf(::Type{T}) where T<:Rational = typemin(T)
neginf(::Type{T}) where T<:Integer = T(-999999)
neginf(::Type{Int16}) = Int16(-16384)
neginf(::Type{Int8}) = Int8(-64)

struct Tropical{T} <: Number
    n::T
    Tropical{T}(x) where T = new{T}(T(x))
    function Tropical(x::T) where T
        new{T}(x)
    end
    function Tropical{T}(x::Tropical{T}) where T
        x
    end
    function Tropical{T1}(x::Tropical{T2}) where {T1,T2}
        new{T1}(T2(x.n))
    end
end

Base.show(io::IO, t::Tropical) = Base.print(io, "$(t.n)ₜ")

Base.:*(a::Tropical, b::Tropical) = Tropical(a.n + b.n)
function Base.:*(a::Tropical{<:Rational}, b::Tropical{<:Rational})
    if a.n.den == 0
        a
    elseif b.n.den == 0
        b
    else
        Tropical(a.n + b.n)
    end
end
Base.:+(a::Tropical, b::Tropical) = Tropical(max(a.n, b.n))
Base.zero(::Type{Tropical{T}}) where T = Tropical(neginf(T))
Base.zero(::Tropical{T}) where T = zero(Tropical{T})

Base.one(::Type{Tropical{T}}) where T = Tropical(zero(T))
Base.one(::Tropical{T}) where T = one(Tropical{T})

# inverse and division
Base.inv(x::Tropical) = Tropical(-x.n)
Base.:/(x::Tropical, y::Tropical) = Tropical(x.n - y.n)
Base.div(x::Tropical, y::Tropical) = Tropical(x.n - y.n)

Base.isapprox(x::Tropical, y::Tropical; kwargs...) = isapprox(x.n, y.n; kwargs...)

# promotion rules
Base.promote_type(::Type{Tropical{T1}}, b::Type{Tropical{T2}}) where {T1, T2} = Tropical{promote_type(T1,T2)}


# # [Cyber Physical Systems Example](@id cyberphysical_example)
#
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/examples/Cyber-Physical.ipynb)

# Authors: Georgios Bakirtzis [https://bakirtzis.net/](https://bakirtzis.net) and Raul Gonzalez Garcia (raulg@iastate.edu)
#
# The following example is a mechanization from
# 1. Compositional Cyber-Physical Systems Modeling - [http://dx.doi.org/10.4204/EPTCS.333.9](http://dx.doi.org/10.4204/EPTCS.333.9)
# 2. Categorical Semantics of Cyber-Physical Systems Theory - [https://doi.org/10.1145/3461669](https://doi.org/10.1145/3461669)

using AlgebraicDynamics
using AlgebraicDynamics.DWDDynam

using Catlab.WiringDiagrams
using Catlab.Graphics
using Catlab.Graphics.Graphviz

using LabelledArrays
using DifferentialEquations
using Plots

# We use  functorial semantics to model a cyper-physical system, namely an unmanned aerial vehicle (UAV).
# We define a diagram of systems (the composition syntax) that is the architecture of the composition.
# Then, we apply behaviors of the individual parts of the system to the architecture. This composition produces
# a complete UAV model.

#-

# We first have to define our boxes and specify what the inports and outports are.
# For example, the sensor box has two inports `:e` and `:s` and one outport `s_prime`.

s = Box(:sensor, [:s, :e], [:s′])
c = Box(:controller, [:d, :s′], [:c])
d = Box(:dynamics, [:c], [:s]);

# A wiring diagram has outer inports and outports which define the interface of target system.
# Then we add the boxes and wires to the diagram and visualize the result.

UAV = WiringDiagram([:e,:d], [:s])

sensor     = add_box!(UAV, s)
controller = add_box!(UAV, c)
dynamics   = add_box!(UAV, d)

add_wires!(UAV, [
    ## net inputs
    (input_id(UAV), 1) => (sensor, 2),
    (input_id(UAV), 2) => (controller, 2),

    ## connections
    (sensor, 1) => (controller, 1),
    (controller, 1) => (dynamics, 1),
    (dynamics, 1) => (sensor, 1),

    ## net output
    (dynamics, 1) => (output_id(UAV), 1)
])

to_graphviz(UAV)

# Then we assign behaviors to inhabit the boxes.
function 𝗟(𝐖)
    𝐿(u, x, p, t) = LVector( sc = -p.𝓐l * (u[1] - x[1] - x[2]) );
    𝐶(u, x, p, t) = LVector( sl = -p.𝓐c * (u[1] + p.𝓑c*x[1] - x[2]) );
    𝐷(u, x, p, t) = LVector( α = -0.313*u[1] +  56.7*u[2] +  0.232*x[1],
                             q = -0.013*u[1] - 0.426*u[2] + 0.0203*x[1],
                             θ =  56.7*u[2]              );

    u_𝐿(u,p,t) = [ u[1] ];  # outputs sl
    u_𝐶(u,p,t) = [ u[1] ];  # outputs sc
    u_𝐷(u,p,t) = [ u[3] ];  # outputs θ

    return oapply(𝐖,
                  Dict(:sensor     => ContinuousMachine{Float64}(2, 1, 1, 𝐿, u_𝐿),
                       :controller => ContinuousMachine{Float64}(2, 1, 1, 𝐶, u_𝐶),
                       :dynamics   => ContinuousMachine{Float64}(1, 3, 1, 𝐷, u_𝐷)));
end

𝑢ᵤₐᵥ = 𝗟(UAV)

# Lastly, we compute and plot the solution.

## initial values
xₒ = LVector( e = 0.01,  # [e, d] -> [θ offset, 𝛿 control input]
              d = 0.05);

uₒ = [0.0, 0, 0, 0, 0]
tspan = (0, 20.0)

params = (𝓐l = 100,  # decay constant of sensor
          𝓐c = 100,  # decay constant of controller
          𝓑c = 0)    # ratio of velocity to reference velocity

prob = ODEProblem(𝑢ᵤₐᵥ, uₒ, xₒ, tspan, params)
sol = solve(prob, alg_hints=[:stiff]);

#-

plot(sol, vars = [1,2, ((t,y) -> (t, y*1e2), 0, 4), 3, 5],
    label = ["sl" "sc" "α" "q" "θ"],
    lw = 2, title  = "Aircraft pitch behaviour",
    xlabel = "time", ylabel = "response"
)


# test EuclidDistance

using IncrementalInference
using Test

##

@testset "test EuclidDistance on 1 dim" begin

##

fg = initfg()

addVariable!(fg, :x0, ContinuousScalar)
addFactor!(fg, [:x0], Prior(Normal()))

addVariable!(fg, :x1, ContinuousScalar)

eud = EuclidDistance(Normal(10,1))
addFactor!(fg, [:x0;:x1], eud)

##

tree, _, = solveTree!(fg)

##

@test isapprox(getPPE(fg, :x0).suggested[1], 0, atol=1)

pts = getBelief(fg, :x1) |> getPoints
N = size(pts, 2)

@test 0.3*N < sum( 5 .< pts )
@test 0.3*N < sum( pts .< -5 )
@test sum( -5 .< pts .< 5 ) < 0.1*N

##

end

@testset "test EuclidDistance on 2 dim" begin

##

fg = initfg()

addVariable!(fg, :x0, ContinuousEuclid{2})
addFactor!(fg, [:x0], Prior(MvNormal(zeros(2),diagm([1;1.0]))))

addVariable!(fg, :x1, ContinuousEuclid{2})

eud = EuclidDistance(Normal(10,1))
addFactor!(fg, [:x0;:x1], eud)

tree, _, = solveTree!(fg)

@test isapprox(getPPE(fg, :x0).suggested[1], 0, atol=1)
@test isapprox(getPPE(fg, :x0).suggested[1], 0, atol=1)

pts = getBelief(fg, :x1) |> getPoints
N = size(pts, 2)

pts .^= 2
@test 0.5*N < sum( 7 .< sqrt.(sum(pts, dims=1)) .< 13 )

##

end