text:
module CTReg

#NOTE: This file contains all of Clinton Tepper's regression code

if pwd() ∉ LOAD_PATH
  push!(LOAD_PATH,pwd())
end

using  DataFrames, Distributions, StatsBase, GLM, CategoricalArrays#, Base.BLAS, Base.LinAlg, RData

const DEBUG_CTMOD = false
const DEFAULT_STAR_LEGEND = "*p < 0.1, **p < 0.05, ***p < 0.01 s.t. p=Pr(>|T|)"

export CTLM, CT2SLS, project!, getResid!, pullModelMatrix, getCoeff!,
  getHomoskedΣ!, getModWhiteΣ!,  getWhiteΣ!,getNeweyWest!,
  getModWhiteΣSlow, getHomoskedΣSlow, getWhiteΣSlow, getNeweyWestSlow,
  dropNullsFromDF, CTExpr, CTSym, CTQR, get1stStage, getTerm, getR,getNeweyWestFunc

#this alias is useful when parsing inputs to formula objects
const CTExpr = Union{Symbol,Expr,Void}
const CTSym = Union{Symbol,Void}

abstract type CTModel end

##################Utility Functions ###############################

#helper function  to get the model matrix from a dataframe given a formula
#IN: A dataframe and formula object
#OUT: the model matrix
function getModelMatrix(df::T, f::Formula)::Matrix{Float64} where
  T <: AbstractDataFrame
  return ModelMatrix(ModelFrame(f, df)).m
end

#Same as above but allows for an expression
function getModelMatrix(df::T, exp::V)::Matrix{Float64} where
  {T <: AbstractDataFrame, V <: CTExpr}

  #special case which crashes Formula
  if exp == Symbol("")
    return ones(Float64,size(df,1),1)
  end

  return getModelMatrix(df, get1SidedFormula(exp))
end

#helper function  to get the list of symbols in an expression
#IN: A dataframe and a string object, typically representing a formula
#OUT: a vector of symbols in the string object
function getSymbolsFromExpr(exp::V)::Vector{Symbol} where {V <: CTExpr}
  symStrings::Vector{String} = split(string(exp),
    ['|','=','~',' ','+','*','&',')','(','-'], keep=false)

  filter!(s::String->typeof(parse(s))≠Int,symStrings)

  #convert the strings to symbols
  return Symbol.(symStrings)
end

#helper function to create a one-sided formula given an expression
#IN: an expression and dataframe
#OUT: A one-sided formula object
function get1SidedFormula(RHS::T)::Formula where
  {T <: CTExpr}

  return Formula(nothing, RHS)
end

#drops nulls  from dataframe given a list of symbols
#IN: a data frame and symbols
#OUT: dataframe less the null entires as dictated by the symbol list
dropNullsFromDF(df::DataFrame, syms::Vector{Symbol})::DataFrame =
  df[completecases(df[:,syms]),:]

#drops nulls  from dataframe given a list of symbols
#IN: a sub-dataframe and symbols, #OUT: subdataframe less the null entries
dropNullsFromDF(dfSub::SubDataFrame, syms::Vector{Symbol})::SubDataFrame =
    view(dfSub,completecases(dfSub[:,syms]))

#helper function for the above, takes in a single symbol
#IN: a subdataframe and symbol #OUT: subdataframe less the null entries
function dropNullsFromDF(df::T, s::Symbol) where {T<:AbstractDataFrame}

  return dropNullsFromDF(df, collect([s]))
end

#helper function for the above, de-nulls entire dataframe
#IN: a subdataframe #OUT: subdataframe less the null entries
function dropNullsFromDF(df::T) where {T<:AbstractDataFrame}
  return dropNullsFromDF(df, names(df))
end

#####################CTQR Type###########################
#=This is a storage object specific to the QR optimizations
meant to hold the QR factorization and a handle for X
Values: X matrix that is decomposed, Q matrix, and R matrix
The X matrix has dimensions of NxK. R is the inverse, stored
for computational efficiency.=#

struct CTQR
  X::Matrix{Float64}
  Q::Matrix{Float64}
  R::Matrix{Float64}

  N::Int
  K::Int
  RInv::Matrix{Float64}
end

#helper method which takes the inverse of the R matrix (which is used in a lot)
#IN: X matrix and its Q R decomposition in matrix form
#OUT: CTQR Object
CTQR(X::Matrix{Float64},Q::Matrix{Float64},R::Matrix{Float64};keepX=true)::CTQR =
  CTQR(keepX?X:Matrix{Float64}(),Q,R,size(X,1),size(X,2),(R)\eye(size(X,2)))

#main constructor for the CTQR object
#IN: X Matrix
#OUT: CTQR object
function CTQR(X::Matrix{Float64}; keepX = true)::CTQR
  Q::Matrix{Float64},R::Matrix{Float64} = qr(X)
  return CTQR(X,Q,R)
end

#default constructor to clear memory
CTQR() = CTQR(Matrix{Float64}(),Matrix{Float64}(),Matrix{Float64}(),
    0,0,Matrix{Float64}())::CTQR

#####################CTLM Model Type#########################
#This is designed to hold the minimal info for a linear model
#Plus the associated CTQR object
 mutable struct CTLM <: CTModel
  xqr::CTQR #handle to the xqr object
  X::Matrix{Float64} #the data and x matrix (handle)
  Y::Vector{Float64} #Y data

  N::Int #Number of data points (rows)
  K::Int #Number of data columns
  β::Vector{Float64} #beta coefficient
  ε::Vector{Float64} #residuals

  XNames::Vector{Symbol} #names of the X variables
  YName::Symbol #names of the Y variables
  #rowLabels::Vector{Date}

end

function CTLM(X::Matrix{Float64}, Y::Vector{Float64};
    XNames::Vector{Symbol} = Symbol.(:X, 1:size(X,2)),
    YName::Symbol = :Y)::CTLM

    #println("flag5")
    #showall(X[1:100,1:2])
    xqr::CTQR = CTQR(X, keepX=false)
    #println("flag6")
    β::Vector{Float64} = Vector{Float64}(xqr.K)
    getCoeff!(xqr,Y,β)
    ε::Vector{Float64} = Vector{Float64}(xqr.N)
    getResid!(X,Y,β,ε)

    return CTLM(xqr, X, Y, xqr.N, xqr.K, β, ε, XNames, YName)
end

#=Constructor and helper method which gets required info from DataFrame
IN: The source dataframe, a formula expression for the RHS,
    the dependent variable as a symbol, optionally names for X columns,
    Y and a switch to elliminate null values
OUT: A CTLM object
NOTE: RHS expression must be ordered with factors last, otherwise
    the factor names will be incorrect =#
function CTLM(df::DataFrame,  XExpr::T, YSym::Symbol;
    XNames::Vector{Symbol}=[Symbol(:Intercept)], YName::Symbol=:Y,
    eliminateNulls=true, fixedEffectsSym::V = nothing)::CTLM  where {T <: CTExpr, V<:Union{Symbol,Void}}

    #println("flag1")
    #get the list of symbols from the expression for X
    XSym::Vector{Symbol} = getSymbolsFromExpr(XExpr)

    #make a view so if we drop nulls it doesn't affect the original
    if fixedEffectsSym == nothing || fixedEffectsSym ∈ [XSym; YSym]
        dfSub::SubDataFrame = view(df[:,[XSym; YSym]],1:size(df,1))
    else
        dfSub = view(df[:,[XSym; YSym; fixedEffectsSym]],1:size(df,1))
    end

    if eliminateNulls     # if we are going to drop nulls
        dfSub = dropNullsFromDF(dfSub)
    end

    #get the factor expanded model matrix (Adds the dummies)
  #  println("XExpr: $XExpr")
    XModelMatrix::Matrix{Float64} = getModelMatrix(dfSub, XExpr)
    #println("got here 41")
    #showcols(dfSub)
    #this is a check to make sure the factor expansion follows other columns
    if length(XSym) > 1
        for i ∈ 2:length(XSym) #all types
            if typeof(dfSub[:,XSym[i-1]])<:CategoricalArray &&
                !(typeof(dfSub[:,XSym[i]])<:CategoricalArray)
                warn("Type of $(typeof(dfSub[:,i])) is preceeded by type of
                    factor. Column names are likely MISALIGNED. Put factors
                    after numerical variables in RHS expressions.")
            end
        end
    end

    #println("flag2")
    #Now assign the names as appropriate
    XNamesFull::Vector{Symbol} = Vector{Symbol}(size(XModelMatrix,2))
    for i ∈ 1:length(XNamesFull)
        XNamesFull[i] = i≤length(XNames)?XNames[i]:Symbol(:X,i)
    end

    #get the appropriate matrices and construct the CTLM object

    if fixedEffectsSym != nothing
      #println("flag3a")
            #println(typeof(Vector(dfSub[:,fixedEffectsSym])))
        return CTLM(XModelMatrix, Vector{Float64}(dfSub[:,YSym]),
          Vector(dfSub[:,fixedEffectsSym]),
            XNames=XNamesFull, YName=YName)
    else
      #  println("flag3b")
        return CTLM(XModelMatrix,
            Vector{Float64}(dfSub[:,YSym]),
            XNames=XNamesFull, YName=YName)
    end

end

#One-way clustered SEs
function CTLM(X::Matrix{Float64}, Y::Vector{Float64}, fixedEffects::Vector{V};
    XNames::Vector{Symbol} = Symbol.(:X, 1:size(X,2)),
    YName::Symbol = :Y)::CTLM where V<:Union{
      Float64,Int,Date,Symbol, CategoricalValue, Missing}

    #println("flag4")
    #we get the unique values and make a two way table
    #tVars::Vector{T} = unique(tFI)
    iVars::Vector{V} = unique(fixedEffects)
    iN::Int = length(iVars)
    iTable::Dict = Dict(iVars[i] => i for i::Int ∈ 1:iN)

    K::Int = size(X,2)
    N::Int = size(X,1)

    #now create the codified versions
    #Xi::Matrix{Float64} = Matrix{Float64}(length(iVars),size(X,2))
    fixedEffectsCode::Vector{Int} = ((x::V)->iTable[x]).(fixedEffects)

    meanXY::Matrix{Float64} = zeros(iN,K+1)
    NXY::Vector{Int} = zeros(iN)

    #println("$N,$K,$iN,$(size(NXY))")
    @fastmath @inbounds for r ∈ 1:N
        meanXY[fixedEffectsCode[r],K+1] += Y[r]
        NXY[fixedEffectsCode[r]] += 1.0
    end

    @fastmath @inbounds for c ∈2:K, r ∈1:N
            meanXY[fixedEffectsCode[r],c] += X[r,c]
    end

    meanXY ./= NXY

    @fastmath @inbounds for c ∈2:K, r ∈1:N
            X[r,c] -= meanXY[fixedEffectsCode[r],c]
    end

    @fastmath @inbounds for r ∈ 1:N
         Y[r] -= meanXY[fixedEffectsCode[r],K+1]
    end

    #=@fastmath Threads.@threads for i ∈ 1:iN
        cluster::Vector{Int} = find((x::Int)->x==i,fixedEffectsCode)

        meanY::Float64 = mean(Y[cluster])
        #print(i<10?"$meanY\n":"")
        Y[cluster] .-= meanY

        #We get a dramatic benefit from the parallelization
        for j::Int ∈ 1:K
            if XNames[j] != :intercept
                meanX::Float64 = mean(X[cluster,j])
                X[cluster, j] .-= meanX
            end
        end
    end=#

    return CTLM(X, Y, XNames=XNames, YName=YName)
end

#this function calculates the Pearson correlation coefficient of the predicted values
#relative to the realized values
getR(lin::CTLM)::Float64 = cor(lin.Y, lin.Y .- lin.ε)

#blank default constructor for memory management
CTLM() = CTLM(CTQR(), Matrix{Float64}(), Vector{Float64}(), 0, 0,
        Vector{Float64}(), Vector{Float64}(), Vector{Symbol}(), Symbol())::CTLM

#fucntion for getting regression coefficients
getTerm(lm::CTLM, s::Symbol) = (lm.β[findfirst(lm.XNames, s)])::Float64

######2SLS constructor
#Holds the modelling componenets for a 2SLS model (except errors)
#Variable descriptions are in the struct
# Runs standard 2SLS using QR optimization
mutable struct CT2SLS <: CTModel
    zaqr::CTQR #QR decomposition of Za (Z|W)
    xaqr::CTQR #QR decomposition of fitted 1st stage values (X̃|W)

    X::Matrix{Float64} #Endogeneous covariates
    W::Matrix{Float64} #Exogeneous covariates
    Y::Vector{Float64} #Dependent var
    Z::Matrix{Float64} #IV data

    N::Int #number of data points
    K::Int # of endogeneous covariates
    L::Int # of instrumental variables
    KW::Int # of exogeneous covariates

    Π1::Matrix{Float64} #first stage results
    Ξ1::Matrix{Float64} #first stage residuals

    δ2::Vector{Float64} #2SLS coefficients
    ξ2::Vector{Float64} #2SLS residuals

    XNames::Vector{Symbol} #names of the X variables
    WNames::Vector{Symbol} #names of the W variables
    YName::Symbol #names of the Y variables
    ZNames::Vector{Symbol} #names of the Z variables

    #CT2SLS()
end

#= Convenience constructor for CT2SLS that allocates a 0 element array for the
case of no exogeneous covariates. NOTE: no exog covariates implies no intercept,
so there probably should be  some exogeneous covariates
#IN: X Matrix, Y Vector, and the Z matrix of Instrumental variables
#OUT: CT2SLS object=#
CT2SLS(X::Matrix{Float64},Y::Vector{Float64},Z::Matrix{Float64};
    XNames::Vector{Symbol} = Symbol.(:X, 1:size(X,2)),
    YName::Symbol = :Y,
    ZNames::Vector{Symbol} = Symbol.(:Z, 1:size(Z,2)))::CT2SLS =
        CT2SLS(X,Matrix{Float64}(Size(X,1),0),Y,Z)

#=Main constructor for CT2SLS object
IN: Endogeneous matrix X, Exogeneous matrix W, Dependent vector Y,
instrumental zariable matrix Z
OUT: A 2SLS object
=#
function CT2SLS(X::Matrix{Float64}, W::Matrix{Float64},
  Y::Vector{Float64},Z::Matrix{Float64};
  XNames::Vector{Symbol} = Symbol.(:X, 1:size(X,2)),
  WNames::Vector{Symbol} = Symbol.(:W, 1:size(W,2)),
  YName::Symbol = :Y,
  ZNames::Vector{Symbol} = Symbol.(:Z, 1:size(Z,2)))::CT2SLS

  #memory pre-allocaiton
  N::Int = size(X,1) #number of data points
  K::Int = size(X,2) # of endogeneous variables
  L::Int = size(Z,2) # of instrumental variables
  KW::Int = size(W,2)# of exogeneous covariates

  Π1::Matrix{Float64} = Matrix{Float64}(L+KW,K)
  X̂::Matrix{Float64} = Matrix{Float64}(N,K) #holds fitted first stage values
  Ξ1::Matrix{Float64} = similar(X̂) #holds first stage residuals
  δ2::Vector{Float64} = Vector{Float64}(K+KW)
  ξ2::Vector{Float64} = similar(Y)
  #Xa::Matrix{Float64} = Matrix{Float64}(N,K+KW) #fitted 1st stage plus exog covariates

  #start with setting up and running the first stage
  ZW::Matrix{Float64} = [Z W]
  zaqr::CTQR = CTQR(ZW, keepX=false) #augment the matrices to ZW and get the QR
  getCoeff!(zaqr, X, Π1) #get the first stage coefficients
  BLAS.gemm!('N','N',1.0,ZW,Π1,0.0,X̂) #get the fitted first stage values
  getResid!(ZW, X, Π1, Ξ1) #get the residuals

  #now run the second stage
  xaqr::CTQR = CTQR([X̂ W], keepX = false) #augment X̂W and get the QR
  getCoeff!(xaqr, Y, δ2) #get the 2SLS coefficients
  getResid!([X W], Y, δ2, ξ2) #get the 2SLS residuals

  #Finally, return the constructed model
  return CT2SLS(zaqr, xaqr, X, W, Y, Z, N, K, L, KW, Π1, Ξ1, δ2, ξ2,
    XNames, WNames, YName, ZNames)
end

#Constructor and helper method which gets required info from DataFrame
#IN: The source dataframe, a formula expression/symbol for X, W, Y, and Z
# the dependent variable as a symbol
#OUT: A CTLM object
function CT2SLS(df::DataFrame, XExpr::T, WExpr::U, YSym::Symbol, ZExpr::V;
    XNames::Vector{Symbol} = Vector{Symbol}(),
    WNames::Vector{Symbol} = Vector{Symbol}([:Intercept]),
    YName::Symbol = :Y,
    ZNames::Vector{Symbol} = Vector{Symbol}(),
    eliminateNulls=true)::CT2SLS where {T <: CTExpr, U <: CTExpr, V <: CTExpr}

    #get the list of symbols for X, W and Z
    XSym::Vector{Symbol} = getSymbolsFromExpr(XExpr)
    WSym::Vector{Symbol} = getSymbolsFromExpr(WExpr)
    ZSym::Vector{Symbol} = getSymbolsFromExpr(ZExpr)

    #make a view so if we drop nulls it doesn't affect the original
    dfSub::SubDataFrame = view(df[:, [XSym; WSym; YSym; ZSym]],1:size(df,1))

    # if we are going to drop nulls
    if eliminateNulls
        dfSub = dropNullsFromDF(dfSub)
    end

    #get the model matrices, which will expand all factors
    XModelMatrix::Matrix{Float64} = getModelMatrix(dfSub, XExpr)
    WModelMatrix::Matrix{Float64} = getModelMatrix(dfSub, WExpr)
    ZModelMatrix::Matrix{Float64} = getModelMatrix(dfSub, ZExpr)

    #this is a check to make sure the factor expansion follows other columns
    #We must repeat the check on each variable matrix
    for symVec::Vector{Symbol} ∈ [XSym, WSym, ZSym]
        if length(symVec) > 1
            for i ∈ 2:length(symVec) #all types
                if typeof(dfSub[:,symVec[i-1]])<:CategoricalArray &&
                    !(typeof(dfSub[:,symVec[i]])<:CategoricalArray)
                    warn("Type of $(typeof(dfSub[:,i])) is preceeded by type of
                        factor. Column names are likely MISALIGNED. Put factors
                        after numerical variables in RHS expressions.")
                end
            end
        end
    end

    #now create the namesfor the expanded matrix
    XNamesFull::Vector{Symbol} = Vector{Symbol}(size(XModelMatrix,2))
    for i ∈ 1:length(XNamesFull)
        XNamesFull[i] = i≤length(XNames)?XNames[i]:Symbol(:X,i)
    end

    WNamesFull::Vector{Symbol} = Vector{Symbol}(size(WModelMatrix,2))
    for i ∈ 1:length(WNamesFull)
        WNamesFull[i] = i≤length(WNames)?WNames[i]:Symbol(:W,i)
    end

    ZNamesFull::Vector{Symbol} = Vector{Symbol}(size(ZModelMatrix,2))
    for i ∈ 1:length(ZNamesFull)
        ZNamesFull[i] = i≤length(ZNames)?ZNames[i]:Symbol(:Z,i)
    end

    #get the appropriate matrices and construct the CT2SLS object
    return CT2SLS(XModelMatrix, WModelMatrix, Vector{Float64}(dfSub[:,YSym]),
        ZModelMatrix, XNames=XNamesFull, WNames=WNamesFull,
        YName=YName, ZNames=ZNamesFull)

end

#default constructor to help with memory management
CT2SLS() = CT2SLS(CTQR(), CTQR(), Matrix{Float64}(), Matrix{Float64}(),
    Vector{Float64}(), Matrix{Float64}(), 0, 0, 0, 0, Matrix{Float64}(),
    Matrix{Float64}(), Vector{Float64}(), Vector{Float64}(),
    Vector{Symbol}(), Vector{Symbol}(), Symbol(), Vector{Symbol}())::CT2SLS

#=function get1stStage
#The purpose is to get the first stage of the IV as a linear model
#Unfortunately, we need to copy some of the data, so this can be expensive,
#although Z and W should only be copied once
#IN: an IV object
#OUT: a vector of CTLM objects, one for each focal (X) variable in the
#initial regression=#
function get1stStage(iv::CT2SLS)::Vector{CTLM}
    ZW::Matrix{Float64} =  [iv.Z iv.W]
    iv1st::Vector = Vector{CTLM}()

    for i ∈ 1:length(iv.XNames)
        push!(iv1st, CTLM(iv.zaqr, [iv.Z iv.W], iv.X[:,i],
            iv.zaqr.N, iv.zaqr.K, iv.Π1[:,i], iv.Ξ1[:,i],
            [iv.ZNames; iv.WNames], iv.XNames[i]))
    end

    return iv1st
end

#########################Project!##############################
#=gets the projection matrix

IN: THe CTQR decomposition and the memory for  the projection matrix
OUT: Writes and returns the projection matrix
NOTE: This method will allocate the projection matrix if no second
argument is supplied. =#
function project!(xqr::CTQR, P::Matrix{Float64}=Matrix{Float64}(xqr.N, xqr.N)
  )::Matrix{Float64}

  #use the BLAS library for fast matrix algebra
  return BLAS.gemm!('N','T',1.0,xqr.Q,xqr.Q,0.0,P)

end

#=This version gets only the diagonal of the projection matrix
IN: THe CTQR decomposition and the memory for  the projection diagonal
OUT: Writes and returns the projection matrix=#
function project!(xqr::CTQR, P::Vector{Float64})::Vector{Float64}
  #get the diagonal quickly
  #equivlenet to diag(Q*Q')
  P .= 0.0
  @fastmath for j::Int ∈ 1:xqr.K, i::Int ∈ 1:xqr.N
    P[i] += xqr.Q[i,j] * xqr.Q[i,j]
  end

  return P

end

#=This version runs either of the above methods after generating
the xqr object given an independent variable matrix
IN: Independent variable X matrix and the memory for  the projection matrix
or vector for the diagonal
OUT: Writes and returns the projection matrix=#
function project!(X::Matrix{Float64},
  P::Array{Float64}=Matrix{Float64}(size(X,1),size(X,1)))::Array{Float64}

  return project!(CTQR(X),P)
end

#=simple test function to verify the projection optimizations
output should be identical to the matrix function above
IN: Independent variable X matrix, memory for the projection matrix
OUT: Writes and returns teh projection matrix=#
function projectSlow!(X::Matrix{Float64},P::Matrix{Float64})
  P .= (X * ((X' * X)\eye(size(X,2))) * X')
  return P
end

#########################getResid!#########################
#=gets the residuals in the 2SLS framework

Note this is a wrapper function which can accept multiple independent
and multiple dependent variables. NOTE the change in the naming convention
relative to other verisons of this method (due to the 2SLS context)
IN: QR decomposition of the RHS variables, LHS dependent variable matrix,
a coefficient matrix, optionally the memory for the residual matrix
OUT: Writes and returns the residual matrix=#
function getResid!(ZW::Matrix{Float64}, X::Matrix{Float64}, Π1::Matrix{Float64},
  Ξ1::Matrix{Float64}=similar(X))::Matrix{Float64}

  #get the residuals by reference
  for i ∈ 1:size(X,2) #get the vector of reisduals multiple times
    getResid!(ZW, view(X,:,i), view(Π1,:,i), view(Ξ1,:,i))
  end
  return Ξ1
end

#=Main method to get residuals
IN: X matrix the RHS variables, LHS dependent variable,
a coefficient vector, optionally the memory for the residual vector
OUT: Writes and returns the residual vector=#
function getResid!(X::Matrix{Float64}, Y::T,
  β::T, ε::T=similar(Y))::T where T<:StridedVector{Float64}

  return BLAS.axpy!(1.0, Y, BLAS.gemv!('N',-1.0, X, β,0.0,ε)) #ε=Y-Xβ
end

#########################getCoef!###########################
# Gets the coefficients in a regression

#=This is a wrapper function which gets a matrix of coefficients for use in
# the first stage of 2SLS. NOTE the change in the naming convention
relative to other verisons of this method (due to the 2SLS context)
IN: QR decomposition of the RHS variables, LHS dependent variable matrix,
optionally the memory for the coefficient matrix
OUT: Writes and returns the coefficient matrix=#
function getCoeff!(zaqr::CTQR, X::Matrix{Float64},
  Π1::Matrix{Float64} = Matrix{Float64}(zaqr.K, size(X,2)))::Matrix{Float64}

  for i ∈ 1:size(X,2) # for each dependent variable vector
    getCoeff!(zaqr, view(X,:,i), view(Π1,:,i)) #run the regression as needed
  end

  return Π1
end

#=Main method to get the regression coefficients
IN: QR decomposition of X matrix the RHS variables, LHS dependent variable,
optionally the memory for the coefficient vector
OUT: Writes and returns the coefficient vector=#
function getCoeff!(xqr::CTQR, Y::T,
  β::T = Vector{Float64}(xqr.K))::T where  T<:StridedVector{Float64}

  return  BLAS.gemv!('N',1.0,BLAS.gemm('N', 'T', xqr.RInv, xqr.Q),Y,0.0,β)
end

#=Convenience method to get the regression coefficients
IN: X matrix the RHS variables, LHS dependent variable,
optionally the memory for the coefficient vector
OUT: Writes and returns the coefficeint vector=#
getCoeff!(X::Vector{Float64}, Y::Vector{Float64}) =
  getCoeff!(CTQR(X), Y, Vector{Float64}(size(X,2)))::Vector{Float64}

###################getHomosked!#######################
#Gets the covariance matrix under Homoskedastic assumptions σ^2[X'X]^-1

#main method to get the covariance matrix
#IN: The QR decomposition from the regression, a vector of residuals,
#optionally memory for the covariance matrix
#OUT: Writes and returns the covariance matrix
function getHomoskedΣ!(xqr::CTQR, ε::Vector{Float64},
  Σ::Matrix{Float64} = Matrix{Float64}(lin.K, lin.K))::Matrix{Float64}

  return BLAS.gemm!('N','T',(ε⋅ε)/(xqr.N-xqr.K),xqr.RInv,xqr.RInv,0.0,Σ) #[X'X]^-1*σ2
end

#=helper method for the above function which extracts the QR decomposition
from a linear model
IN: A linear model, optionally memory for the covariance matrix
OUT: WRites and returns the covariance matrix=#
function getHomoskedΣ!(lin::CTLM,
  Σ::Matrix{Float64}=Matrix{Float64}(lin.K, lin.K))::Matrix{Float64}

  return getHomoskedΣ!(lin.xqr, lin.ε, Σ)
end

#=helper method for the above function which extracts the QR decomposition
from a CT2SLS object
IN: A linear model, optionally memory for the covariance matrix
OUT: WRites and returns the covariance matrix=#
function getHomoskedΣ!(iv::CT2SLS,
  Σ::Matrix{Float64}=Matrix{Float64}(iv.xaqr.K, iv.xaqr.K))::Matrix{Float64}

  return getHomoskedΣ!(iv.xaqr, iv.ξ2, Σ)
end

#=Simple test function to verify the QR-related optimizations
Output should be identical to the standard methods
IN: Independent variable X matrix
OUT: Returns the covariance matrix=#
getHomoskedΣSlow(X::Matrix{Float64}, ε::Vector{Float64})::Matrix{Float64} =
  (X.'*X)\eye(size(X,2))*(ε⋅ε)/(size(X,1)-size(X,2))

#Convenience method for above: IN: A linear model #OUT: A covariance matrix
getHomoskedΣSlow(lin::CTLM)::Matrix{Float64} =
  getHomoskedΣSlow(lin.X, lin.Y.-lin.X * lin.β)

#Convenience method for above: IN: A CT2SLS object #OUT: A covariance matrix
getHomoskedΣSlow(iv::CT2SLS)::Matrix{Float64} =
  getHomoskedΣSlow([[iv.Z iv.W]*iv.Π1 iv.W], iv.Y.-[iv.X iv.W] * iv.δ2)

#########################getNeweyWest ##########OLS Only
function getNeweyWest!(X::Matrix{Float64}, xqr::CTQR, ε::Vector{Float64}, lag::Int,
  Σ::Matrix{Float64} = Matrix{Float64}(xqr.K, xqr.K))::Matrix{Float64}

  #pre-allocate for the spectral matrix

  Rv::Matrix{Float64} = Matrix{Float64}(xqr.K, xqr.K) #pre-allocate working matrix
  RRInv::Matrix{Float64} = BLAS.gemm('N', 'T', xqr.RInv, xqr.RInv) #this is equivelent to [X'X]^-1

  #need to multiply through by the error
  Xe::Matrix{Float64} = X .* ε
  ST::Matrix{Float64} = BLAS.gemm('T','N',1.0/xqr.N, Xe, Xe)
  for v::Int ∈ 1:lag
    #overwrites Rv with (1/N)R'R
    BLAS.gemm!('T', 'N', 1.0/xqr.N, view(Xe, (v+1):(xqr.N), :),view(Xe, 1:(xqr.N-v), :), 0.0, Rv)
    #Rv .= view(Xe, (v+1):(xqr.N), :)' * view(Xe, 1:(xqr.N-v), :) .* (1.0/xqr.N)
    ST .+= (lag+1.-v)/(lag+1.) .* (Rv .+ Rv')
  end

  #this is [X'X]^-1S=[R'R]^-1S
  RRInvS::Matrix{Float64} = BLAS.gemm('N', 'N', RRInv, ST)

  #finally we have T[X'X]^-1S[X'X]^-1
  BLAS.gemm!('N','N',Float64(xqr.N), RRInvS, RRInv, 0.0, Σ)
  return Σ
end

function getNeweyWest!(lin::CTLM, lag::Int,
  Σ::Matrix{Float64} = Matrix{Float64}(lin.K, lin.K))::Matrix{Float64}

  return getNeweyWest!(lin.X, lin.xqr, lin.ε, lag, Σ)
end

#helper function in case the input requires a single argument function
getNeweyWestFunc(lag::Int) = (lin::CTLM)-> getNeweyWest!(lin, lag)

#testing for NeweyWest
function getNeweyWestSlow(X::Matrix{Float64}, ε::Vector{Float64}, lag::Int)
  N::Int, K::Int = size(X)

  #first form the spectral matrix

  #initial value for spectral matrix
  xxt::Matrix{Float64} = Matrix{Float64}(K,K)
  xxt .= zeros(K,K)
  for t::Int ∈ 1:N
    xxt .+= (X[t, :] * X[t, :]') .* ε[t]^2
  end
  ST::Matrix{Float64} =  xxt ./ N

  #make the rest of the spectral matrix
  for v::Int ∈ 1:lag
    xxt .= zeros(K,K)
    for t::Int ∈ (1+v):(N)
      xxt .+= (X[t, :] * X[t-v, :]') .* ε[t]*ε[t-v]
    end
    xxt ./= N
    ST .+= (lag+1.-v)/(lag+1.) .* (xxt .+ xxt')
  end

  XXInv::Matrix{Float64} = (X' * X) \ eye(K)

  return N*XXInv*ST* XXInv
end

getNeweyWestSlow(lin::CTLM, lag::Int) = getNeweyWestSlow(lin.X, lin.Y .- lin.X * lin.β, lag)

getNeweyWestFuncSlow(lag::Int)::Function = (lin::CTLM)-> getNeweyWestSlow(lin, lag)
###################getWhiteERRORs ###OLS only

function getWhiteΣ!(xqr::CTQR, ε::Vector{Float64},
  Σ::Matrix{Float64} = Matrix{Float64}(xqr.K, xqr.K))::Matrix{Float64}

  Λ::Vector{Float64} = ε.^2.0

  QRtInv::Matrix{Float64} = BLAS.gemm('N','T',xqr.Q,xqr.RInv) #Q(R')^-1
  RInvQΛ::Matrix{Float64} = QRtInv.'

  #loop to scale the matrix by Λ (modified part of modified white)
  @fastmath for j ∈ 1:xqr.N, i∈1:xqr.K
    RInvQΛ[i,j] *= Λ[j]
  end

  #final multiplicaiton and assignment
  return BLAS.gemm!('N','N',1.0,RInvQΛ, QRtInv,0.0,Σ) #R^-1Q'ΛQ(R')^-1
end

function getWhiteΣ!(lin::CTLM,
  Σ::Matrix{Float64} = Matrix{Float64}(lin.K, lin.K))::Matrix{Float64}

  return getWhiteΣ!(lin.xqr, lin.ε, Σ)
end

#for testing purposes only
function getWhiteΣSlow(X::Matrix{Float64}, ε::Vector{Float64})
  XXInv::Matrix{Float64} = (X.' * X)\eye(size(X,2))

  return XXInv * X.' * diagm(ε) .^ 2.0 * X * XXInv
end

#Convenience method for above: IN: A linear model #OUT: A covariance matrix
getWhiteΣSlow(lin::CTLM) = getWhiteΣSlow(lin.X, lin.Y.-lin.X * lin.β)

###################getMWErrors!##############################
#Main method to get the modified white SEs Does #[X'X]^-1X'ΛX[X'X]^-1

#=Main method to calculate the modified white SEs
IN: A CTQR decomposition object, a residuals vector,
optional memory for a covariance matrix
OUT: Writes and returns the covariance matrix=#
function getModWhiteΣ!(xqr::CTQR, ε::Vector{Float64},
  Σ::Matrix{Float64} = Matrix{Float64}(lin.K, lin.K))::Matrix{Float64}

  Λ::Vector{Float64} = similar(ε)

  project!(xqr,Λ) #get the projection diagonal
  Λ .= (ε./(1.0.-Λ)).^2.0

  QRtInv::Matrix{Float64} = BLAS.gemm('N','T',xqr.Q,xqr.RInv) #Q(R')^-1
  RInvQΛ::Matrix{Float64} = QRtInv.'

  #loop to scale the matrix by Λ (modified part of modified white)
  @fastmath for j ∈ 1:xqr.N, i∈1:xqr.K
    RInvQΛ[i,j] *= Λ[j]
  end

  #final multiplicaiton and assignment
  return BLAS.gemm!('N','N',1.0,RInvQΛ, QRtInv,0.0,Σ) #R^-1Q'ΛQ(R')^-1
end

#=helper method which extracts the required components from a linear model
IN: A linear model, optionally memory for the covariance matrix
OUT: Writes and allocates the covariance matrix=#
function getModWhiteΣ!(lin::CTLM,
  Σ::Matrix{Float64} = Matrix{Float64}(lin.K, lin.K))::Matrix{Float64}

  return getModWhiteΣ!(lin.xqr, lin.ε, Σ)
end

#=helper method which extracts the required components from a 2SLS model
IN: A 2SLS model, optionally memory for the covariance matrix
OUT: Writes and allocates the covariance matrix=#
function getModWhiteΣ!(iv::CT2SLS,
  Σ::Matrix{Float64}=Matrix{Float64}(iv.xaqr.K, iv.xaqr.K))::Matrix{Float64}

  return getModWhiteΣ!(iv.xaqr, iv.ξ2, Σ)
end

#=Simple test function to verify the QR-related optimizations
Output should be identical to the standard methods
IN: Independent variable X matrix
OUT: Returns the covariance matrix=#
function getModWhiteΣSlow(X::Matrix{Float64}, ε::Vector{Float64})
  XXInv::Matrix{Float64} = (X.' * X)\eye(size(X,2))

  P::Matrix{Float64} = X * (XXInv) * X.'
  return XXInv * X.' * diagm(ε  ./ (1.0 .- diag(P))) .^ 2.0 * X * XXInv

end

#Convenience method for above: IN: A linear model #OUT: A covariance matrix
getModWhiteΣSlow(lin::CTLM) = getModWhiteΣSlow(lin.X, lin.Y.-lin.X * lin.β)

#Convenience method for above: IN: A 2SLS object #OUT: A covariance matrix
getModWhiteΣSlow(iv::CT2SLS)::Matrix{Float64} =
  getModWhiteΣSlow([[iv.Z iv.W]*iv.Π1 iv.W], iv.Y.-[iv.X iv.W] * iv.δ2)

  #################testing methods for this module
function LMtest()

  #Number of Observations
  N::Int = 200
  K::Int = 2

  #Allocate
  X::Matrix{Float64} = Matrix{Float64}(N,K)
  Y::Vector{Float64} = Vector{Float64}(N)
  ε::Vector{Float64} = similar(Y)

  #parameters for the simulation
  σ2::Vector{Float64} = [2.0^2.0 for i::Int ∈ 1:N]
  σ2[1:ceil(Int, N/2)] .= 0.5^2.0
  β::Vector{Float64} = [(1.0 * i) for i::Int ∈ 1:K]
  #println("β: ", β)
  #println("σ2: $σ2")

  #this will hold the sampled beta
  e::Vector{Float64} = similar(ε)

  #run the simulation
  X[:,1] .= 1.0 #intercept
  for i ∈ 2:K
    X[:,i] .= rand(Uniform(),N)
  end

  ε .= map((s2::Float64)->rand(Normal(0.0,s2^0.5)),σ2)
  Y .=  X*β .+ ε

  #get the linear model
  lin::CTLM = CTLM(X, Y)

  #get the homoskedastic SEs
  ΣHomosked::Matrix{Float64} = getHomoskedΣ!(lin)
  ΣHomoskedSlow::Matrix{Float64} = getHomoskedΣSlow(lin)

  #print the coefficients
  println("Coefficients: ",lin.β)
  println("Homoskedastic Errors: ", diag(ΣHomosked).^.5)
  println("Check: ", diag(ΣHomoskedSlow).^.5)

  #get the modified white SEs
  ΣWhite::Matrix{Float64} = similar(ΣHomosked)
  getModWhiteΣ!(lin, ΣWhite)
  ΣWhiteSlow::Matrix{Float64} = getModWhiteΣSlow(lin)

  #print the coefficients
  println("Modified White Errors: ",diag(ΣWhite).^.5)
  println("Check: ", diag(ΣWhiteSlow).^.5)

    #=Π1::Matrix{Float64} = Matrix{Float64}(K,2)
  getCoeff!(lin.xqr, [Y 2.*Y], Π1)
  println("1st stage Z: $(lin.X)")
  println("1st Stage X: $([Y 2.*Y])" )
  println("1st Stage Test Coef: $Π1" )
  Ξ::Matrix{Float64} = getResid!(lin.xqr, [Y 2.*Y], Π1)
  println("1st Stage Resid: $Ξ" )=#

  #test the project routines
  Pa::Vector{Float64} = Vector{Float64}(N)
  PM::Matrix{Float64} = Matrix{Float64}(N,N)
  PS::Matrix{Float64} = Matrix{Float64}(N,N)

  project!(X,PM)
  project!(X,Pa)

  projectSlow!(X,PS)
  println("P: ", Pa[1:5])
  println("P (from full matrix): ", PM[1:3,1:3])
  println("P Slow: ", diag(PS)[1:10])
end

function IVTest()

  srand(11)

  N::Int = 200 #Number of Observations
  K::Int = 3 # of endogneous covariates
  L::Int = 4 # of instruments
  KW::Int = 5# of exog covariates

  #Allocate
  X::Matrix{Float64} = Matrix{Float64}(N,K)
  Z::Matrix{Float64} = Matrix{Float64}(N,L)
  W::Matrix{Float64} = Matrix{Float64}(N,KW)
  Y::Vector{Float64} = Vector{Float64}(N)
  A::Vector{Float64} = Vector{Float64}(N)

  Π1::Matrix{Float64} = Matrix{Float64}(L+KW,K)
  π2::Vector{Float64} = Vector{Float64}(L+KW)
  Ξ1::Matrix{Float64} = similar(X)
  e2::Vector{Float64} = similar(Y)

  #parameters for the simulation
  σ2::Vector{Float64} = [2.0^2 for i::Int ∈ 1:N]
  σ2[1:ceil(Int, N/2)] .= 0.5^2.0

  #generate π2 coefficients (pattern is 1 times the index, ie 1 2 3 ...)
  π2 .= [(10*i) for i::Int ∈ 1:(L+KW)]
  π2[L+1] .= 1.0 #intercept term
  #println("π2: ",π2)

  #generate Π1, pattern is δ2 times 1/10 the index, ie ie .1δ2 .2δ2 .3δ2 ...)
  for i ∈ 1:(L+KW), j∈1:K
    Π1[i,j] .= i + j^2
  end
  Π1[L+1,:] .= 1.0 #set the intercept term
  #display(Π1)

  #generate the exogeneous variables
  for i ∈ 2:KW
    W[:,i] .= rand(Uniform(-10,10),N)
  end
  W[:,1] .= 1.0 #intercept

  #generate the instrumental variables
  for i ∈ 1:L
    Z[:,i] .= rand(Uniform(-10,10),N)
  end

  #map((s2::Float64)->rand(Normal(0.0,s2^0.5)),σ2)
  #generate the errors
  for i ∈ 1:N
    for j ∈ 1:K
      Ξ1[i,j] = rand(Normal(0.0,σ2[i]^0.5))
    end
    e2[i] = rand(Normal(0.0,σ2[i]^0.5))
  end

  #println(Ξ1)
  #populate Y and X
  X .= [Z W]*Π1 .+ Ξ1
  Y .= [Z W]*π2 .+ e2

  #add some endogeneity
  for i ∈ 1:K
    A .= rand(Uniform(-i,i),N)
    X[:,i] .+= A
    Y .+= A .* 10
  end

  #probably is a more graceful way to do this
  XNames::Vector{Symbol} = [Symbol("X$i") for i = 1:K]
  WNames::Vector{Symbol} = [Symbol("W$i") for i = 1:KW]
  ZNames::Vector{Symbol} = [Symbol("Z$i") for i = 1:L]

  iv::CT2SLS = CT2SLS(X,W,Y,Z, XNames=XNames, YName=:Y, WNames=WNames, ZNames=ZNames)
  println("coefficients: \n", [XNames; WNames] ,"\n", iv.δ2)

  #get the homoskedastic SEs
  ΣHomosked::Matrix{Float64} = getHomoskedΣ!(iv)
  ΣHomoskedSlow::Matrix{Float64} = getHomoskedΣSlow(iv)

  #print the coefficients
  println("Homoskedastic Errors: ", diag(ΣHomosked).^.5)
  println("Check: ", diag(ΣHomoskedSlow).^.5)

  #get the modified white SEs
  ΣWhite::Matrix{Float64} = getModWhiteΣ!(iv)
  ΣWhiteSlow::Matrix{Float64} = getModWhiteΣSlow(iv)

  #print the coefficients
  println("Modified White Errors: ",diag(ΣWhite).^.5)
  println("Check: ", diag(ΣWhiteSlow).^.5)

  iv1st::Vector{CTLM} = get1stStage(iv)
  println("1st Stage results: \n $(iv1st[1].XNames)")
  for iv1 ∈ iv1st
      println("coefficients $(iv1.YName):  $(iv1.β )")
      println("Error (Homosked $(iv1.YName)): $(diag(getHomoskedΣ!(iv1)).^.5)")
  end

  oStream::IOStream = open(
    "C:\\Users\\Clinton\\Dropbox\\Projects\\Summer17RTester\\Summer17RTester\\IVTestOut.csv","w+")

  write(oStream, [string.(iv.XNames); string.(iv.WNames)]::Vector{String})
  writecsv(oStream, [X W Y Z])
  close(oStream)
end

if DEBUG_CTMOD
  #LMtest()
  IVTest()
  #IOTest()
  #print(perfTestStr(10^6))
end

end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
"""
    SGHMC(n_iters::Int, learning_rate::Float64, momentum_decay::Float64)

Stochastic Gradient Hamiltonian Monte Carlo sampler.

Usage:

```julia
SGHMC(1000, 0.01, 0.1)
```

Example:

```julia
@model example begin
  ...
end

sample(example, SGHMC(1000, 0.01, 0.1))
```
"""
mutable struct SGHMC{T} <: Hamiltonian
  n_iters        ::  Int       # number of samples
  learning_rate  ::  Float64   # learning rate
  momentum_decay ::  Float64   # momentum decay
  space          ::  Set{T}    # sampling space, emtpy means all
  gid            ::  Int
end
function SGHMC(learning_rate::Float64, momentum_decay::Float64, space...)
  SGHMC(1, learning_rate, momentum_decay, space..., 0)
end
function SGHMC(n_iters, learning_rate, momentum_decay)
  SGHMC(n_iters, learning_rate, momentum_decay, Set(), 0)
end
function SGHMC(n_iters, learning_rate, momentum_decay, space...)
  _space = isa(space, Symbol) ? Set([space]) : Set(space)
  SGHMC(n_iters, learning_rate, momentum_decay, _space, 0)
end
SGHMC(alg::SGHMC, new_gid::Int) = SGHMC(alg.n_iters, alg.learning_rate, alg.momentum_decay, alg.space, new_gid)

function step(model, spl::Sampler{<:SGHMC}, vi::VarInfo, is_first::Bool)
  if is_first
    if ~haskey(spl.info, :wum)
      if spl.alg.gid != 0 link!(vi, spl) end    # X -> R

      #wum = WarmUpManager(1, 1, Dict())

      D = length(vi[spl])
      ve = VarEstimator{Float64}(0, zeros(D), zeros(D))
      wum = WarmUpManager(1, Dict(), ve)
      wum[:ϵ] = [spl.alg.learning_rate]
      wum[:stds] = ones(D)
      spl.info[:wum] = wum

      oldθ = realpart(vi[spl])
      vi[spl] = oldθ

      # Initialize velocity
      v = zeros(Float64, size(oldθ))
      spl.info[:v] = v

      if spl.alg.gid != 0 invlink!(vi, spl) end # R -> X
    end

    push!(spl.info[:accept_his], true)

    vi
  else
    # Set parameters
    η, α = spl.alg.learning_rate, spl.alg.momentum_decay

    @debug "X-> R..."
    if spl.alg.gid != 0
      link!(vi, spl)
      runmodel(model, vi, spl)
    end

    @debug "recording old variables..."
    old_θ = realpart(vi[spl]);
    θ = deepcopy(old_θ)
    _, grad = gradient_forward(old_θ, vi, model, spl)
    old_v = deepcopy(spl.info[:v])
    v = deepcopy(old_v)

    if verifygrad(grad)
      @debug "update latent variables and velocity..."
      # Implements the update equations from (15) of Chen et al. (2014).
      for k in 1:size(old_θ, 1)
        θ[k,:] = old_θ[k,:] + old_v[k,:]
        noise = rand(MvNormal(zeros(length(old_θ[k,:])), sqrt.(2 * η * α)*ones(length(old_θ[k,:]))))
        v[k,:] = (1. - α) * old_v[k,:] - η * grad[k,:] + noise # NOTE: divide η by batch size
      end
    end

    @debug "saving new latent variables and velocity..."
    spl.info[:v] = v
    vi[spl] = θ

    @debug "always accept..."
    push!(spl.info[:accept_his], true)

    @debug "R -> X..."
    if spl.alg.gid != 0 invlink!(vi, spl); cleandual!(vi) end

    vi
  end
end

ext: jl


text:
using JLD
using PyPlot
using FileIO
using Statistics

a = load("Results/Videos/09_24_segway/Actor_Critic_2018-09-24iteartion_930.jld")
a = load("Results/Videos/10_01_segway/Actor_Critic_2018-10-01iteartion_1260.jld")
res = a["res"]
res = res[1:1259] #./ 500

function f_filt(x)
    l = length(x)
    tau = .1
    y = zeros(l)
    y[1] = x[1]

    for i = 2:l
        y[i] = (1-tau)*y[i-1] + tau *x[i]
    end

    return y
end

clf()
figure(2)
grid(alpha = 0.25)
plot(f_filt(res))
xlabel("Iterations", fontsize = 16)
ylabel("Cost", fontsize = 16)

max_stars_repo_name: NikonPic/REINFORCEMENT-JULIA

max_issues_repo_name: NikonPic/REINFORCEMENT-JULIA

max_forks_repo_name: NikonPic/REINFORCEMENT-JULIA


text:
using URIParser, Requests

eacct = "EX123456"
uname = "malcolm@amisllp.com"
pword = "passwd"

dispatcher = "api.esendex.com/v1.0/messagedispatcher2
uri = URI(string(https://,uname,'@',pword,'/',dispatcher)

message = """?xml version'1.0' encoding='UTF-8' ?>
<messages>
<accountreference>EX123456</accountreference>
<message>
<to>447777555666</to>
<body>Every message matters</body>
</message>
</messages>"""

post(uri; data = message)

max_stars_repo_path: Module 3/Chapter09/esendex.jl

max_issues_repo_path: Module 3/Chapter09/esendex.jl

max_forks_repo_path: Module 3/Chapter09/esendex.jl


text:
# Full-scale parsing tests of JuliaSyntax itself, Julia Base, etc.

@testset "Parse JuliaSyntax" begin
    pkgdir = joinpath(@__DIR__, "..")
    test_parse_all_in_path(joinpath(pkgdir, "src"))
    test_parse_all_in_path(joinpath(pkgdir, "test"))
end

base_path = joinpath(Sys.BINDIR, Base.DATAROOTDIR, "julia", "base")
@testset "Parse Base at $base_path" begin
    test_parse_all_in_path(base_path)
end

if haskey(ENV, "PARSE_BASE_TEST")
# TODO: Turn on by default

base_tests_path = joinpath(Sys.BINDIR, Base.DATAROOTDIR, "julia", "test")
@testset "Parse Base tests at $base_tests_path" begin
    test_parse_all_in_path(base_tests_path)
end

end

@testset "Parse Julia stdlib at $(Sys.STDLIB)" begin
    for stdlib in readdir(Sys.STDLIB)
        fulldir = joinpath(Sys.STDLIB, stdlib)
        if isdir(fulldir)
            @testset "Parse $stdlib" begin
                test_parse_all_in_path(joinpath(Sys.STDLIB, fulldir))
            end
        end
    end
end


text:
module TestVisualization

using Test
using Documenter
using Trixi
using Plots

include("test_trixi.jl")

# pathof(Trixi) returns /path/to/Trixi/src/Trixi.jl, dirname gives the parent directory
EXAMPLES_DIR = joinpath(pathof(Trixi) |> dirname |> dirname, "examples", "2d")

# Start with a clean environment: remove Trixi output directory if it exists
outdir = "out"
isdir(outdir) && rm(outdir, recursive=true)

# Run various visualization tests
@testset "Visualization tests" begin
  # Run Trixi
  @test_nowarn_debug trixi_include(@__MODULE__, joinpath(examples_dir(), "2d", "elixir_euler_blast_wave_amr.jl"),
                                   tspan=(0,0.1))

  @testset "PlotData2D, PlotDataSeries2D, PlotMesh2D" begin
    # Constructor
    @test PlotData2D(sol) isa PlotData2D
    @test PlotData2D(sol; nvisnodes=0, grid_lines=false, solution_variables=cons2cons) isa PlotData2D
    pd = PlotData2D(sol)

    # show
    @test_nowarn_debug show(stdout, pd)
    println(stdout)

    # getindex
    @test pd["rho"] == Trixi.PlotDataSeries2D(pd, 1)
    @test pd["v1"] == Trixi.PlotDataSeries2D(pd, 2)
    @test pd["v2"] == Trixi.PlotDataSeries2D(pd, 3)
    @test pd["p"] == Trixi.PlotDataSeries2D(pd, 4)
    @test_throws KeyError pd["does not exist"]

    # convenience methods for mimicking a dictionary
    @test pd[begin] == Trixi.PlotDataSeries2D(pd, 1)
    @test pd[end] == Trixi.PlotDataSeries2D(pd, 4)
    @test length(pd) == 4
    @test size(pd) == (4,)
    @test keys(pd) == ("rho", "v1", "v2", "p")
    @test eltype(pd) == Pair{String, Trixi.PlotDataSeries2D}
    @test [v for v in pd] == ["rho" => Trixi.PlotDataSeries2D(pd, 1),
                              "v1" => Trixi.PlotDataSeries2D(pd, 2),
                              "v2" => Trixi.PlotDataSeries2D(pd, 3),
                              "p" => Trixi.PlotDataSeries2D(pd, 4)]

    # PlotDataSeries2D
    pds = pd["p"]
    @test pds.plot_data == pd
    @test pds.variable_id == 4
    @test_nowarn_debug show(stdout, pds)
    println(stdout)

    # getmesh/PlotMesh2D
    @test getmesh(pd) == Trixi.PlotMesh2D(pd)
    @test getmesh(pd).plot_data == pd
    @test_nowarn_debug show(stdout, getmesh(pd))
    println(stdout)
  end

  @testset "2D plot recipes" begin
    pd = PlotData2D(sol)

    @test_nowarn_debug plot(sol)
    @test_nowarn_debug plot(pd)
    @test_nowarn_debug plot(pd["p"])
    @test_nowarn_debug plot(getmesh(pd))
  end

  @testset "PlotData1D, PlotDataSeries1D, PlotMesh1D" begin
    # Run Trixi
    @test_nowarn_debug trixi_include(@__MODULE__, joinpath(examples_dir(), "1d", "elixir_euler_blast_wave.jl"),
                                     tspan=(0,0.1))

    # Constructor
    @test PlotData1D(sol) isa PlotData1D
    pd = PlotData1D(sol)

    # show
    @test_nowarn_debug show(stdout, pd)
    println(stdout)

    # getindex
    @test pd["rho"] == Trixi.PlotDataSeries1D(pd, 1)
    @test pd["v1"] == Trixi.PlotDataSeries1D(pd, 2)
    @test pd["p"] == Trixi.PlotDataSeries1D(pd, 3)
    @test_throws KeyError pd["does not exist"]

    # convenience methods for mimicking a dictionary
    @test pd[begin] == Trixi.PlotDataSeries1D(pd, 1)
    @test pd[end] == Trixi.PlotDataSeries1D(pd, 3)
    @test length(pd) == 3
    @test size(pd) == (3,)
    @test keys(pd) == ("rho", "v1", "p")
    @test eltype(pd) == Pair{String, Trixi.PlotDataSeries1D}
    @test [v for v in pd] == ["rho" => Trixi.PlotDataSeries1D(pd, 1),
                              "v1" => Trixi.PlotDataSeries1D(pd, 2),
                              "p" => Trixi.PlotDataSeries1D(pd, 3)]

    # PlotDataSeries1D
    pds = pd["p"]
    @test pds.plot_data == pd
    @test pds.variable_id == 3
    @test_nowarn_debug show(stdout, pds)
    println(stdout)

    # getmesh/PlotMesh1D
    @test getmesh(pd) == Trixi.PlotMesh1D(pd)
    @test getmesh(pd).plot_data == pd
    @test_nowarn_debug show(stdout, getmesh(pd))
    println(stdout)

    @testset "1D plot recipes" begin
      pd = PlotData1D(sol)

      @test_nowarn_debug plot(sol)
      @test_nowarn_debug plot(pd)
      @test_nowarn_debug plot(pd["p"])
      @test_nowarn_debug plot(getmesh(pd))
    end

    # Fake a PlotDataXD objects to test code for plotting multiple variables on at least two rows
    # with at least one plot remaining empty
    @testset "plotting multiple variables" begin
      x = collect(0.0:0.1:1.0)
      data1d = rand(5, 11)
      variable_names = string.('a':'e')
      mesh_vertices_x1d = [x[begin], x[end]]
      fake1d = PlotData1D(x, data1d, variable_names, mesh_vertices_x1d)
      @test_nowarn_debug plot(fake1d)

      y = x
      data2d = [rand(11,11) for _ in 1:5]
      mesh_vertices_x2d = [0.0, 1.0, 1.0, 0.0]
      mesh_vertices_y2d = [0.0, 0.0, 1.0, 1.0]
      fake2d = PlotData2D(x, y, data2d, variable_names, mesh_vertices_x2d, mesh_vertices_y2d)
      @test_nowarn_debug plot(fake2d)
    end
  end

  @testset "plot 3D" begin
    @test_nowarn_debug trixi_include(@__MODULE__, joinpath(examples_dir(), "3d", "elixir_advection_basic.jl"),
                                     tspan=(0,0.1))
    @test PlotData2D(sol) isa PlotData2D
  end

  @testset "plotting TimeIntegratorSolution" begin
    @test_nowarn_debug trixi_include(@__MODULE__, joinpath(examples_dir(), "2d", "elixir_hypdiff_lax_friedrichs.jl"))
    @test_nowarn_debug plot(sol)
  end

  @testset "VisualizationCallback" begin
    # To make CI tests work, disable showing a plot window with the GR backend of the Plots package
    # Xref: https://github.com/jheinen/GR.jl/issues/278
    # Xref: https://github.com/JuliaPlots/Plots.jl/blob/8cc6d9d48755ba452a2835f9b89d3880e9945377/test/runtests.jl#L103
    if !isinteractive()
      restore = get(ENV, "GKSwstype", nothing)
      ENV["GKSwstype"] = "100"
    end

    @test_nowarn_debug trixi_include(@__MODULE__,
                               joinpath(examples_dir(), "2d", "elixir_advection_amr_visualization.jl"),
                               visualization = VisualizationCallback(interval=20,
                                               clims=(0,1),
                                               plot_creator=Trixi.save_plot),
                               tspan=(0.0, 2.0))

    @testset "elixir_advection_amr_visualization.jl with save_plot" begin
      @test isfile(joinpath(outdir, "solution_000000.png"))
      @test isfile(joinpath(outdir, "solution_000020.png"))
      @test isfile(joinpath(outdir, "solution_000024.png"))
    end

    @testset "show" begin
      @test_nowarn_debug show(stdout, visualization)
      println(stdout)

      @test_nowarn_debug show(stdout, "text/plain", visualization)
      println(stdout)
    end

    # Restore GKSwstype to previous value (if it was set)
    if !isinteractive()
      if isnothing(restore)
        delete!(ENV, "GKSwstype")
      else
        ENV["GKSwstype"] = restore
      end
    end
  end
end

end #module


text:
using Combinatorics
using StaticArrays

include("common.jl")

struct Digit
    # Record whether each segment is active.
    segments::SVector{7,Bool}
end

function Base.parse(::Type{Digit}, code::AbstractString)
    result = MVector{7,Bool}(undef)
    fill!(result, false)
    for char in code
        char_index = Int64(char) - Int64('a') + 1
        result[char_index] = true
    end
    return Digit(SVector(result))
end

struct Display
    digits::Vector{Digit}
end

function Base.parse(::Type{Display}, input::AbstractString)
    digits = Digit[]
    for digit in split(strip(input), ('|', ' '))
        isempty(digit) && continue
        push!(digits, parse(Digit, digit))
    end
    return Display(digits)
end

real_digits = [
    parse(Digit, "abcefg"),  # 0
    parse(Digit, "cf"),
    parse(Digit, "acdeg"),
    parse(Digit, "acdfg"),
    parse(Digit, "bcdf"),
    parse(Digit, "abdfg"),
    parse(Digit, "abdefg"),
    parse(Digit, "acf"),
    parse(Digit, "abcdefg"),
    parse(Digit, "abcdfg"),  # 9
]

permute(digit::Digit, perm) = Digit(digit.segments[perm])
permute(display::Display, perm) = Display(map(x -> permute(x, perm), display.digits))

"""Get all the digits that can be parsed. Terminates as soon as invalid digit is found."""
function get_all_numbers(display::Display)
    result = Int64[]
    for digit in display.digits
        index = findfirst(==(digit), real_digits)
        isnothing(index) && return result  # invalid digit!
        push!(result, index - 1)
    end
    return result
end

example = """
acedgfb cdfbe gcdfa fbcad dab cefabd cdfgeb eafb cagedb ab | cdfeb fcadb cdfeb cdbaf
"""

function unscramble_display(display::Display)
    for perm in permutations(1:7)  # Only 5040 cases... not too bad to bash.
        this_display = permute(display, perm)
        numbers = get_all_numbers(this_display)
        length(numbers) == 14 && return this_display
    end
    error("Couldn't find correct permutation :(")
end

function get_test_numbers(display::Display)
    all_numbers = get_all_numbers(unscramble_display(display))
    return all_numbers[11:14]
end

example_2 = """
be cfbegad cbdgef fgaecd cgeb fdcge agebfd fecdb fabcd edb | fdgacbe cefdb cefbgd gcbe
edbfga begcd cbg gc gcadebf fbgde acbgfd abcde gfcbed gfec | fcgedb cgb dgebacf gc
fgaebd cg bdaec gdafb agbcfd gdcbef bgcad gfac gcb cdgabef | cg cg fdcagb cbg
fbegcd cbd adcefb dageb afcb bc aefdc ecdab fgdeca fcdbega | efabcd cedba gadfec cb
aecbfdg fbg gf bafeg dbefa fcge gcbea fcaegb dgceab fcbdga | gecf egdcabf bgf bfgea
fgeab ca afcebg bdacfeg cfaedg gcfdb baec bfadeg bafgc acf | gebdcfa ecba ca fadegcb
dbcfg fgd bdegcaf fgec aegbdf ecdfab fbedc dacgb gdcebf gf | cefg dcbef fcge gbcadfe
bdfegc cbegaf gecbf dfcage bdacg ed bedf ced adcbefg gebcd | ed bcgafe cdgba cbgef
egadfb cdbfeg cegd fecab cgb gbdefca cg fgcdab egfdb bfceg | gbdfcae bgc cg cgb
gcafb gcf dcaebfg ecagb gf abcdeg gaef cafbge fdbac fegbdc | fgae cfgab fg bagce
"""

function get_ans(input::AbstractString)
    num_interesting_digits = 0
    for line in split(input, "\n")
        isempty(line) && continue
        display = parse(Display, line)
        for n in get_test_numbers(display)
            if n in (1, 4, 7, 8)
                num_interesting_digits += 1
            end
        end
    end
    return num_interesting_digits
end

@assert get_ans(example_2) == 26
get_ans(load_input(8))

# Part 2

function make_num(digits::Vector{Int64})
    result = 0
    multiplier = 1
    for digit in reverse(digits)
        result += multiplier * digit
        multiplier *= 10
    end
    return result
end

function get_ans2(input::AbstractString)
    total = 0
    for line in split(input, "\n")
        isempty(line) && continue
        display = parse(Display, line)
        number = make_num(get_test_numbers(display))
        total += number
    end
    return total
end

@assert get_ans2(example_2) == 61229
get_ans2(load_input(8))


text:
using KernelStat
using ARIMA
using DependentBootstrap

#-------------------------------
# SIMULATION FUNCTIONS
#-------------------------------
#Simulate kernel functions
const numObsBandwidthSim = 500::Int
const maxCorLag = 6::Int
const kernelList = [KernelUniform(), KernelEpanechnikov(), KernelGaussian(), KernelPR1993FlatTop(), KernelP2003FlatTop(), KernelPP2002Trap(), KernelPP2002Smooth(), KernelPR1994SB(numObsBandwidthSim)]::Vector{KernelFunction}
const bandwidthList = [BandwidthWhiteNoise(), BandwidthBartlett(), BandwidthP2003(numObsBandwidthSim)]
const hacList = [HACVarianceBasic(KernelUniform(0, maxCorLag), BandwidthP2003(numObsBandwidthSim)), HACVarianceBasic(KernelGaussian(1), BandwidthP2003(numObsBandwidthSim)), HACVarianceBasic(KernelEpanechnikov(1), BandwidthP2003(numObsBandwidthSim)), HACVarianceBasic(KernelBartlett(1), BandwidthP2003(numObsBandwidthSim)), HACVarianceBasic(KernelPR1994SB(numObsBandwidthSim, 0.1), BandwidthMax())]

#------------------------------
# TEST FUNCTIONS
#------------------------------
function testevaluate()
	for k = 1:length(kernelList)
		kF = kernelList[k]
		aD = activedomain(kF)
		if aD.start == -Inf || aD.stop == Inf
			aD.start == -Inf ? (tempStart = -10) : (tempStart = aD.start)
			aD.stop == Inf ? (tempStop = 10) : (tempStop = aD.stop)
			aD = UnitRange(tempStart, tempStop)
		end
		xIn = [aD.start:(aD.stop - aD.start)/10:aD.stop]
		xOut = [ evaluate(xIn[n], kF) for n = 1:length(xIn) ]
		println("Kernel func = " * string(kF) * ". evaluate of active domain = " * string(xOut))
	end
	println("")
	println("")
	return(true)
end

function testbandwidth()
	for k = 0:maxCorLag
		k == 0 ? maCoef = Array(Float64, 0) : maCoef = [ 1 / 2^p for p = 1:k ]
		for j = 1:length(bandwidthList)
			bC = bandwidthList[j]
			mVec = Array(Int, 10)
			for i = 1:10
				x = simulate(numObsBandwidthSim, ARIMAModel(maCoef=maCoef))
				(mVec[i], xVar, xCovVec) = bandwidth(x, bC)
			end
			println("Bandwidth method = " * string(bC))
			println("    MA Order = " * string(k))
			println("    Bandwidth estimates = " * string(mVec))
			println("")
		end
	end
	println("")
	println("")
	return(true)
end

function testhacvariance()
	for k = 0:maxCorLag
		k == 0 ? maCoef = Array(Float64, 0) : maCoef = [ 1 / 2^p for p = 1:k ]
		x = simulate(numObsBandwidthSim, ARIMAModel(maCoef=maCoef))
		meanBoot = dbootstrapstatistic(x, BootstrapParam(x, statistic=mean))
		bootTrueVar = var(sqrt(numObsBandwidthSim) * meanBoot)
		xMean = [ mean(simulate(numObsBandwidthSim, ARIMAModel(maCoef=maCoef))) for i = 1:2000 ]
		simTrueVar = var(sqrt(numObsBandwidthSim) * xMean)
		for j = 1:length(hacList)
			hC = hacList[j]
			vVec = Array(Float64, 10)
			bwVec = Array(Int, 10)
			for i = 1:10
				x = simulate(numObsBandwidthSim, ARIMAModel(maCoef=maCoef))
				(vVec[i], bwVec[i]) = hacvariance(x, hC)
			end
			show(hC)
			println("True MA Order = " * string(k))
			println("Bandwidth estimates = " * string(bwVec))
			println("Simulated true variance estimate = " * string(simTrueVar))
			println("Bootstrapped true variance estimate = " * string(bootTrueVar))
			println("HAC estimates = " * string(vVec))
			println("")
		end
	end
	println("")
	println("")
	return(true)
end

testhacvariance()

function testall(KK::Int)
	testevaluate()
	testbandwidth()
	testhacvariance()
	println("Tests complete.")
end


text:
using DataFrames, CSV, Gadfly, Compose, TableView
filterminlog2foldchange = 0.50
filterminTPM = 2.5

dst = CSV.read("/Users/emsanford/Dropbox (RajLab)/Shared_Eric/Signal_Integration/Analysis_SI2-SI4/extractedData/DeSeqOutputAllConds.tsv")
metadata = CSV.read("/Users/emsanford/Dropbox (RajLab)/Shared_Eric/Signal_Integration/Analysis_SI2-SI4/extractedData/sampleMetadata_SI2-SI4.txt")

######## create a series of "and" filters from singleton conditions ############
dstboolhigh = dst[:, r".*high_isDeGene"]
boolselectionvec1 = dstboolhigh[:, Symbol("RA-high_isDeGene")] .&
                    dstboolhigh[:, Symbol("TGFb-high_isDeGene")]

dstlogfoldchangehigh = dst[:, r".*high_log2fc"]
boolselectionvec2 = (dstlogfoldchangehigh[:, Symbol("RA-high_log2fc")] .> filterminlog2foldchange) .&
                    (dstlogfoldchangehigh[:, Symbol("TGFb-high_log2fc")] .> filterminlog2foldchange)

dstavgtpmhigh = dst[:, r".*high_avgTPM"]
boolselectionvec3 = (dstavgtpmhigh[:, Symbol("RA-high_avgTPM")] .> filterminTPM) .&
                    (dstavgtpmhigh[:, Symbol("TGFb-high_avgTPM")] .> filterminTPM)

finalselectionvec = boolselectionvec1 .& boolselectionvec2 .& boolselectionvec3
println(sum(finalselectionvec))
finalselectionbool = convert(Array{Bool, 1}, finalselectionvec)
################################################################################

dstfiltered = dst[finalselectionbool, :]
dstfiltselect = dstfiltered[:, r"hgnc.symbol|.*_log2fc"]
dststacked = stack(dstfiltered, r".*_log2fc", variable_name=:condition, value_name=:log2fc)
dststacked[:, :foldchange] = 2 .^ dststacked[:, :log2fc]

# now build long table with condition column
dst_tpms = dst[:, r".*_tpm|ensg|hgnc.*"]
dst_tpms_filt = dst_tpms[finalselectionbool, :]
dst_tpms_stacked = stack(dst_tpms_filt, r".*tpm", variable_name=:sampleID, value_name=:tpm)
conditions = [split(join(split(string(s), "-")[2:end], "-"), "_")[1] for s in dst_tpms_stacked.sampleID]
dst_tpms_stacked[:, :condition] = conditions

# add replicate info
dst_tpms_stacked[:, :sampleID] = map(s -> split(string(s), "_")[1], dst_tpms_stacked.sampleID)
dst_tpms_stacked = join(dst_tpms_stacked, metadata[:, [:sampleID; :replicate]], on = :sampleID, kind = :left)

function lessthanfx(left1, right1)
    left = string(left1)
    right = string(right1)
    if (left[1:3] != right[1:3]) | (occursin("and", left) ⊻ occursin("and", right))
        if (occursin("and", left))
            return false
        elseif (occursin("and", right))
            return true
        else
            return isless(left1, right1)
        end
    elseif occursin("low", left)
        return true
    elseif occursin("high", left)
        return false
    elseif occursin("med", left)
        return occursin("high", right)
    end
    return false
end

sort!(dststacked, :condition, lt=lessthanfx)
sort!(dst_tpms_stacked, :condition, lt=lessthanfx)
dstfiltered_avgTPM = stack(dstfiltered[:, r".*_avgTPM|ensg"], r".*_avgTPM", variable_name = :condition_tpm, value_name = :avg_tpm)
sort!(dstfiltered_avgTPM, :condition_tpm, lt=lessthanfx)

a = []
b = []
c = []
for (i,j) in enumerate(unique(dststacked.ensg))
    tempfiltereddst = dst_tpms_stacked[dst_tpms_stacked.ensg .== j, :]
    hgnc = tempfiltereddst[1, Symbol("hgnc.symbol")]
    println(hgnc)

    tempFilteredDstAvgTpm = dstfiltered_avgTPM[dstfiltered_avgTPM.ensg .== j, :]

    # calculate average variance of the "both" conditions
    variances = []
    for cond in ["TGFb-and-RA-low"; "TGFb-and-RA-med"; "TGFb-and-RA-high"]
        tpms = tempfiltereddst[tempfiltereddst.condition .== cond, :tpm]

    end

    # calculate predictions for additive and multiplicative models
    control_tpm_value = tempFilteredDstAvgTpm[tempFilteredDstAvgTpm.condition_tpm .== Symbol("EtOH-nlDensity_avgTPM"), Symbol("avg_tpm")]
    control_tpm_value = control_tpm_value[1]
    prediction_df = DataFrame(ensg = String[],
                              dosage = String[],
                              additive_pred_TPM = Float64[],
                              multiplicative_pred_TPM = Float64[])
    for dosage in ["low"; "med"; "high"]
        ra_tpm   = tempFilteredDstAvgTpm[tempFilteredDstAvgTpm.condition_tpm .== Symbol("RA-"*dosage*"_avgTPM"), Symbol("avg_tpm")][1]
        tgfb_tpm = tempFilteredDstAvgTpm[tempFilteredDstAvgTpm.condition_tpm .== Symbol("TGFb-"*dosage*"_avgTPM"), Symbol("avg_tpm")][1]
        ra_effect = ra_tpm - control_tpm_value
        tgfb_effect = tgfb_tpm - control_tpm_value
        addpred  = control_tpm_value + ra_effect + tgfb_effect
        multpred = (ra_tpm/control_tpm_value) * (tgfb_tpm/control_tpm_value) * control_tpm_value
        push!(prediction_df, (j, dosage, addpred, multpred))

        #for beeswarm plot without layers working
        push!(tempfiltereddst, ("NA", addpred, j, hgnc, "TGFb-and-RA-"*dosage, "addpred"))
        push!(tempfiltereddst, ("NA", multpred, j, hgnc, "TGFb-and-RA-"*dosage, "multpred"))
    end

    print(prediction_df)
    ## attempt to overlay lines on beeswarm plot seems to keep ignoring "x axis" values...
    beeswarmplot = layer(x = tempfiltereddst.condition,
                         y = tempfiltereddst.tpm,
                         color = tempfiltereddst.replicate,
                         Geom.beeswarm)
    lineplot_add_pred_low = layer(x = [1;2;3;0], y = [prediction_df[:, 3]; prediction_df[3, 3]], Geom.line)
    p = plot(beeswarmplot, lineplot_add_pred_low, Guide.title(hgnc))
    if hgnc == "NA"
        hgnc = j
    end

    push!(b, tempfiltereddst)
    push!(c, tempFilteredDstAvgTpm)

    # p = plot(tempfiltereddst, x = :condition, y = :tpm, color = :replicate, Geom.beeswarm, Guide.title(hgnc))
    draw(SVG(joinpath("/Users/emsanford/Dropbox (RajLab)/Shared_Eric/Signal_Integration/Analysis_SI2-SI4/plots/tpmBeeSwarmsAll", hgnc * ".svg")), p)
    push!(a, p)
    break
end
a[1]
# showtable(b[1])

# plot(dststacked, x="condition", y="log2fc", Geom.beeswarm)
# gridstack([a[1] a[2]; a[3] a[4]])


text:
"""
Case 8:
This case study a three bus system with 1 machine (One d- One q-: 4th order model), a VSM of 19 states and an infinite source. All lines are modeled as a static lines.
The perturbation trips two of the three circuits of line between buses 1 and 2, triplicating its impedance.
"""

##################################################
############### LOAD DATA ########################
##################################################

############### Data Network ########################

nodes_case8   = [Bus(1 , "Bus 1"  , "REF" , 0 , 1.02  , (min=0.94, max=1.06), 138),
                    Bus(2 , "Bus 2"  , "PV" , 0 , 1.00 , (min=0.94, max=1.06), 138),
                    Bus(3 , "Bus 3"  , "PQ" , 0 , 1.00 , (min=0.94, max=1.06), 138)]

branch_case8  =  [Line("Line1", true, 0.0, 0.0, Arc(from=nodes_case8[1], to=nodes_case8[3]), 0.01, 0.12, (from=0.1, to=0.1), 100, 1.04),
                    Line("Line2", true, 0.0, 0.0, Arc(from=nodes_case8[1], to=nodes_case8[2]), 0.01, 0.12, (from=0.1, to=0.1), 100, 1.04),
                    Line("Line3", true, 0.0, 0.0, Arc(from=nodes_case8[2], to=nodes_case8[3]), 0.02, 0.9, (from=0.5, to=0.5), 100, 1.04)]

#Trip of Line 1.
branch_case8_fault = [Line("Line1", true, 0.0, 0.0, Arc(from=nodes_case8[1], to=nodes_case8[3]), 0.03, 0.36, (from=0.03, to=0.03), 100, 1.04),
    Line("Line2", true, 0.0, 0.0, Arc(from=nodes_case8[1], to=nodes_case8[2]), 0.01, 0.12, (from=0.1, to=0.1), 100, 1.04),
    Line("Line3", true, 0.0, 0.0, Arc(from=nodes_case8[2], to=nodes_case8[3]), 0.02, 0.9, (from=0.5, to=0.5), 100, 1.04)]

loads_case8 =   [PowerLoad("Bus1", true, nodes_case8[1], PowerSystems.ConstantPower, 0.5, 0.1, 1.5, 0.8),
                   PowerLoad("Bus2", true, nodes_case8[2], PowerSystems.ConstantPower, 1.0, 0.3, 1.5, 0.8),
                   PowerLoad("Bus3", true, nodes_case8[3], PowerSystems.ConstantPower, 0.3, 0.1, 0.5, 0.3)]

############### Data devices ########################

inf_gen_case8 = StaticSource(1, #number
                    :InfBus, #name
                    nodes_case8[1],#bus
                    1.00, #VR
                    0.0, #VI
                    0.000005) #Xth

######## Machine Data #########

### Case 2: 4th Order Model with AVR (3-bus case) ###
case8_machine =  OneDOneQMachine(0.0, #R
                                  1.3125, #Xd
                                  1.2578, #Xq
                                  0.1813, #Xd_p
                                  0.25, #Xq_p
                                  5.89, #Td0_p
                                  0.6, #Tq0_p
                                  100.0)   #MVABase

######## Shaft Data #########

### Shafts for Gen ###
case8_shaft = SingleMass(3.01, #H (M = 6.02 -> H = M/2)
                            0.0) #D

######## PSS Data #########
cases_no_pss = PSSFixed(0.0)

######## TG Data #########

### No TG for Cases 1, 2, 3, 4 ###
case8_no_tg = TGFixed(1.0) #eff

########  AVR Data #########
### AVRs for Case 2, 3, 4 and 5 ###
case8_avr = AVRTypeI(20.0, #Ka - Gain
                        0.01, #Ke
                        0.063, #Kf
                        0.2, #Ta
                        0.314, #Te
                        0.35, #Tf
                        0.001, #Tr
                        5.0, #Vrmax
                        -5.0, #Vrmin
                        0.0039, #Ae - 1st ceiling coefficient
                        1.555) #Be - 2nd ceiling coefficient

### Case 7 Generators ###
case8_gen = DynGenerator(1, #Number
                         :Case8Gen,
                         nodes_case8[2], #bus
                         1.0, # ω_ref,
                         1.0124, #V_ref
                         0.6, #P_ref
                         case8_machine, #machine
                         case8_shaft, #shaft
                         case8_avr, #avr
                         case8_no_tg, #tg
                         cases_no_pss) #pss

############### Inverter Data ########################

converter = AvgCnvFixedDC(138.0, #Rated Voltage
                          100.0) #Rated MVA

dc_source = FixedDCSource(1500.0) #Not in the original data, guessed.

filt = LCLFilter(0.08, #Series inductance lf in pu
                   0.003, #Series resitance rf in pu
                   0.074, #Shunt capacitance cf in pu
                   0.2, #Series reactance rg to grid connection (#Step up transformer or similar)
                   0.01) #Series resistance lg to grid connection (#Step up transformer or similar)

pll = PLL(500.0, #ω_lp: Cut-off frequency for LowPass filter of PLL filter.
          0.084, #k_p: PLL proportional gain
          4.69) #k_i: PLL integral gain

virtual_H = VirtualInertia(2.0, #Ta:: VSM inertia constant
                           400.0, #kd:: VSM damping coefficient
                           20.0, #kω:: Frequency droop gain in pu
                           2*pi*50.0) #ωb:: Rated angular frequency

Q_control = ReactivePowerDroop(0.2, #kq:: Reactive power droop gain in pu
                              1000.0) #ωf:: Reactive power cut-off low pass filter frequency

outer_control = VirtualInertiaQdroop(virtual_H, Q_control)

vsc = CombinedVIwithVZ(0.59, #kpv:: Voltage controller proportional gain
                       736.0, #kiv:: Voltage controller integral gain
                       0.0, #kffv:: Binary variable enabling the voltage feed-forward in output of current controllers
                       0.0, #rv:: Virtual resistance in pu
                       0.2, #lv: Virtual inductance in pu
                       1.27, #kpc:: Current controller proportional gain
                       14.3, #kiv:: Current controller integral gain
                       0.0, #kffi:: Binary variable enabling the current feed-forward in output of current controllers
                       50.0, #ωad:: Active damping low pass filter cut-off frequency
                       0.2) #kad:: Active damping gain

case8_inv = DynInverter(2, #number
                             :DARCO, #name
                             nodes_case8[3], #bus location
                             1.0, #ω_ref
                             0.8, #V_ref
                             0.5, #P_ref
                             -0.3, #Q_ref
                             100.0, #MVABase
                             converter, #Converter
                             outer_control, #OuterControl
                             vsc, #Voltage Source Controller
                             dc_source, #DC Source
                             pll, #Frequency Estimator
                             filt) #Output Filter

######################### Dynamical System ########################

case8_DynSystem = DynamicSystem(nodes_case8, branch_case8,  [case8_inv, case8_gen], vcat(inf_gen_case8, loads_case8),  100.0, 50.0);

##################################################
############### SOLVE PROBLEM ####################
##################################################

dx0 = zeros(LITS.get_total_rows(case8_DynSystem))
x0 = [1.00, #V1_R
          1.00, #V2_R
          1.00, #V3_R
          0.0, #V1_I
         -0.01, #V2_I
         -0.01, #V3_I
         0.0, #δω_vsm
          0.2, #δθ_vsm
          0.025, #qm
          0.0015, #ξ_d
          -0.07, #ξ_q
          0.05, #γ_d
        -0.001, #γ_q
         0.95, #ϕ_d
         -0.10, #ϕ_q
         1.004, #vpll_d
         0.0, #vpll_q
         0.0, #ε_pll
         0.1, #δθ_pll
         0.5, #id_cv
         0.0, #iq_cv
         0.95, #vod
         -0.1, #voq
         0.49, #iod
        -0.1, #ioq
          1.0, #eq_p
          0.47, #ed_p
          0.6, #δ
          1.0, #ω
          2.1, #Vf
          0.28, #Vr1
          -0.39, #Vr2,
          1.0] #Vm

case8_inv.inner_vars[13] = 0.95 #Vd_cnv var
case8_inv.inner_vars[14] = -0.1 #Vq_cnv var
Ybus_fault = Ybus(branch_case8_fault, nodes_case8)[:,:]
u0 = [0.2]
tspan = (0.0, 30.0);

#Find initial condition
inif! = (out,x) -> LITS.system_model!(out, dx0 ,x, (Ybus_fault,case8_DynSystem), 0.0)
sys_solve = nlsolve(inif!, x0, xtol=:1e-8,ftol=:1e-8,method=:trust_region)
x0_init = sys_solve.zero

#Define Fault using Callbacks
cb = DiffEqBase.DiscreteCallback(LITS.change_t_one, LITS.Y_change!)

#Define Simulation Problem
sim = DynamicSimulation(case8_DynSystem, tspan, Ybus_fault, cb, x0_init)

#Solve problem in equilibrium
run_simulation!(sim, IDA());

#Obtain data for voltages
series = get_voltagemag_series(sim, 2)
zoom = [(series[1][ix], series[2][ix]) for (ix, s) in enumerate(series[1]) if (s > 0.90 && s < 1.6)]

@test sim.solution.retcode == :Success


text:
using Pkg
Pkg.activate(joinpath(@__DIR__, "../.."))
Pkg.instantiate()

# ## visualizer
vis = Visualizer()
open(vis)

# ## setup
using Dojo
using IterativeLQR
using LinearAlgebra
using FiniteDiff
using DojoEnvironments
using JLD2

# ## scripts
include(joinpath(module_dir(), "examples/policy/methods/continuation.jl"))
include(joinpath(module_dir(), "examples/policy/methods/tvlqr.jl"))
include(joinpath(module_dir(), "DojoEnvironments/src",
    "quadruped/methods/template.jl"))

################################################################################
# ## system
################################################################################
gravity = -9.81
timestep = 0.02
friction_coefficient = 0.8
damper = 0.5
spring = 1.0
env = get_environment(:quadruped,
    representation=:minimal,
    timestep=timestep,
    contact_body=false,
    gravity=gravity,
    friction_coefficient=friction_coefficient,
    damper=damper,
    spring=spring,
    infeasible_control=true,
    vis=vis)

# ## dimensions
n = env.num_states
m = env.num_inputs
nu_infeasible = 6

################################################################################
# ## simulation test
################################################################################
mech = get_mechanism(:quadruped,
    contact_body=false,
    timestep=timestep,
    gravity=gravity,
    friction_coefficient=friction_coefficient,
    damper=damper,
    spring=spring)

initialize!(mech, :quadruped, body_position=[0,0,0.0])
u_hover = [0.02;0;2.0; 0;0.03;0; zeros(12)]
function ctrl!(m, k; u=u_hover)
    nu = input_dimension(m)
    set_input!(m, SVector{nu}(u))
end

Main.@elapsed storage = simulate!(mech, 3.0, ctrl!,
    opts=SolverOptions(rtol=1e-5, btol=1e-4, undercut=5.0),
    )
Dojo.visualize(mech, storage, vis=env.vis)

################################################################################
# ## reference trajectory
################################################################################
velocity = 0.13
width_scale = 0.6
radius = 0.08
width_scale = ((2Tm-1) * timestep) * velocity /(2 * radius)
Tm = 19
xref = quadruped_trajectory(env.mechanism,
    r=radius,
    z=0.29;
    Δx=-0.02,
    Δfront=0.10,
    width_scale=width_scale,
    height_scale=1.0,
    N=Tm,
    Ncycles=1)
DojoEnvironments.visualize(env, xref)

# ## horizon
T = length(xref)

################################################################################
# ## ILQR problem
################################################################################
# ## model
dyn = IterativeLQR.Dynamics(
    (y, x, u, w) -> dynamics(y, env, x, u, w),
    (dx, x, u, w) -> dynamics_jacobian_state(dx, env, x, u, w),
    (du, x, u, w) -> dynamics_jacobian_input(du, env, x, u, w),
    n, n, m)

model = [dyn for t = 1:T-1]

# ## rollout
x1 = deepcopy(xref[1])
ū = [u_hover for t = 1:T-1]

x̄ = IterativeLQR.rollout(model, x1, ū)
DojoEnvironments.visualize(env, x̄)

# ## objective
############################################################################
qt = [0.3; 0.05; 0.05;
    5e-1 * ones(3);
    1e-3 * ones(3);
    1e-3 * ones(3);
    fill([2, 1e-3], 12)...]
ots = [(x, u, w) -> transpose(x - xref[t]) * Diagonal(timestep * qt) * (x - xref[t]) +
    transpose(u) * Diagonal(timestep * 0.3 * ones(m)) * u for t = 1:T-1]
oT = (x, u, w) -> transpose(x - xref[end]) * Diagonal(timestep * qt) * (x - xref[end])

cts = [IterativeLQR.Cost(ot, n, m) for ot in ots]
cT = IterativeLQR.Cost(oT, n, 0)
obj = [cts..., cT]

# ## constraints
############################################################################
ul = -1.0 * 1e-3*ones(nu_infeasible)
uu = +1.0 * 1e-3*ones(nu_infeasible)

function contt(x, u, w)
    [
        1e-1 * (ul - u[1:nu_infeasible]);
        1e-1 * (u[1:nu_infeasible] - uu);
    ]
end

function goal(x, u, w)
    Δ = 1e-1 * (x - xref[end])[[1:6;13:2:36]]
    return Δ
end

con_policyt = IterativeLQR.Constraint(contt, n, m, indices_inequality=collect(1:2nu_infeasible))
con_policyT = IterativeLQR.Constraint(goal, n, 0)

cons = [[con_policyt for t = 1:T-1]..., con_policyT]

# ## solver
options = Options(line_search=:armijo,
        max_iterations=50,
        max_dual_updates=12,
        min_step_size=1e-2,
        objective_tolerance=1e-3,
        lagrangian_gradient_tolerance=1e-3,
        constraint_tolerance=1e-3,
        initial_constraint_penalty=1e-1,
        scaling_penalty=10.0,
        max_penalty=1e4,
        verbose=true)

s = IterativeLQR.Solver(model, obj, cons, options=options)

IterativeLQR.initialize_controls!(s, ū)
IterativeLQR.initialize_states!(s, x̄)

# ## solve
local_callback!(solver::IterativeLQR.Solver) = continuation_callback!(solver, env, build=false)
reset!(env)
@time IterativeLQR.constrained_ilqr_solve!(s, augmented_lagrangian_callback! = local_callback!)

# ## solution
x_sol, u_sol = IterativeLQR.get_trajectory(s)

# ## visualize
x_view = [[x_sol[1] for t = 1:15]..., x_sol..., [x_sol[end] for t = 1:15]...]
DojoEnvironments.visualize(env, x_view)

################################################################################
# Save
################################################################################
JLD2.jldsave(joinpath(@__DIR__, "../data/trotting_forward.jld2"), x=x_sol, u=u_sol)
file = JLD2.jldopen(joinpath(@__DIR__, "../data/trotting_forward.jld2"))
file["x"]
file["u"]
JLD2.close(file)


text:
using Yelp
using Base.Test


text:
module FurutaPendulums

export FurutaPendulum, SimulatedFurutaPendulum, AbstractFurutaPendulum

using Random, UnPack

using Reexport                           # Use reexport to avoid writing all exports
@reexport using AbstractControlProcesses # Import and export all exported names from ACP
const ACP = AbstractControlProcesses     # Less to write, extending function with new

include("physical.jl")
include("simulated.jl")

const AbstractFurutaPendulum = Union{FurutaPendulum, SimulatedFurutaPendulum}

ACP.num_outputs(p::AbstractFurutaPendulum) = 4
ACP.num_inputs(p::AbstractFurutaPendulum) = 1
ACP.outputrange(p::AbstractFurutaPendulum) = [(0, 2pi), (-1e10, 1e10), (0, 2pi), (-1e10, 1e10)]
ACP.inputrange(p::AbstractFurutaPendulum) = [(-1, 1)]
ACP.isstable(p::AbstractFurutaPendulum) = false
ACP.isasstable(p::AbstractFurutaPendulum) = false
ACP.sampletime(p::FurutaPendulum) = 0.001 # TODO find the correct one
ACP.sampletime(p::SimulatedFurutaPendulum) = p.params.h
ACP.bias(p::AbstractFurutaPendulum) = 0

ACP.control(p::FurutaPendulum, u::AbstractVector) = write(p.voltage, only(u))
ACP.control(p::SimulatedFurutaPendulum, u::AbstractVector) = p.u = only(u)
ACP.measure(p::FurutaPendulum) = [
    read(p.base_angle),
    read(p.base_velocity),
    mod2pi(read(p.arm_angle) + π), # This package has convention that θ=0 means up, the physical hardware has down.
    read(p.arm_velocity)
]
ACP.measure(p::SimulatedFurutaPendulum) = p.x .+ p.params.noise * randn(p.rng, 4)

function ACP.periodic_wait(p::SimulatedFurutaPendulum, last_time, dt)
    step!(p, dt)
    last_time + dt
end

# ACP.initialize(p::FurutaPendulum) = write(p.calibrate, true)
ACP.initialize(p::SimulatedFurutaPendulum) = fill!(p.x, 0) # Used as a reset

end


text:
abstract type AbstractParameter{T} end
const AP{T} = AbstractParameter{T}

# interface
value(::AP) = throw("Not implemented yet")
Base.eltype(::AP{T}) where {T} = T
Base.length(p::AP) = length(value(p))
Base.size(p::AP) = size(value(p))

# Sometime when setup our models, we require it's parameters to follow
# certain constrains, e.g. the support for scaling factor `σ` & length scale `l`
# of a GP kernel should be ℜ⁺, however, adding those constrains make the optimization
# process less convenient. `Parameter` is used to handle this, it's bijector is used
# to map constrained variable to unconstrained one, which is then used during the optimization,
# the inverse of the bijector does the reverse, which will provide us with the constrained
# variables we need.
struct Parameter{T, N, fT} <: AP{T}
    x::AbstractArray{T}
    f::fT

    # parameters with single constraint
    function Parameter{T}(y, f::fT) where {T<:Real, N, fT<:Bijector{N}}
        x = inv(f)(y)
        if N == 0
            return new{T, 0, fT}(T[x], f)
        else
            return new{T, N, fT}(x, f)
        end
    end

    # a set of scalar parameter with multiple constraints
    function Parameter{T}(y::Vector{T}, fs::Bijector{0}...) where {T<:Real}
        length(y) != length(fs) && throw(error("Number of parameters and Bijectors must match !"))
        x = map((f,z)->inv(f)(z), fs, y)
        return new{T, 0, typeof(fs)}(x, fs)
    end
end

Base.:(==)(p1::Parameter{T, N, fT}, p2::Parameter{T, N, fT}) where {T, N, fT} = p1.x==p2.x
Base.:(==)(p1::Parameter{T1, N1, fT1}, p2::Parameter{T2, N2, fT2}) where {T1, N1, fT1, T2, N2, fT2} = false

# functor is used to flatten the `Parameter` into a namedtuple, also with a function to reconstruct from that
# namedtuple to `Parameter`
functor(p::Parameter) = (x=p.x,), x->Parameter(p.f(x...), p.f)
functor(p::Parameter{T, 0}) where {T} = (x=p.x,), x->Parameter(p.f(first(x...)), p.f)
functor(p::Parameter{T, 0, S}) where {T, S<:Tuple} = (x=p.x,), x->Parameter(map((f,x)->f(x), p.f, x...), p.f...)

# outer constructor for a set of scalar parameter with multiple constraints
Parameter(y, fs...) = Parameter{eltype(y)}(y, fs...)

# interface
value(y::Parameter) = y.f(y.x)
value(y::Parameter{T, 0}) where {T} = y.f(first(y.x))
value(y::Parameter{T, 0, S}) where {T, S<:Tuple} = map((f, x)->f(x), y.f, y.x)

# if no constrain is added, it will use `Identity` bijector by default
Parameter(y) = Parameter(y, Identity{ndims(y)}())

# if we want positivity constrain, then use `Exp` bijector
Parameter(y, ::Val{:pos}) = Parameter(y, Bijectors.Exp{ndims(y)}())

# `Constant` is used to hold data that we are not intended to optimize
# Note: in QuantumLang, data is either `Parameter` or `Constant`.
struct Constant{T, N} <: AP{T}
    x::AbstractArray{T}

    function Constant{T}(c) where {T}
        N = ndims(c)
        if N == 0
            return new{T, 0}(T[c])
        else
            return new{T, N}(c)
        end
    end
end

Base.:(==)(c1::Constant{T, N}, c2::Constant{T, N}) where {T, N} = c1.x==c2.x
Base.:(==)(c1::Constant{T1, N1}, c2::Constant{T2, N2}) where {T1, N1, T2, N2} = false

functor(c::Constant) = (), _ -> c

Constant(c) = Constant{eltype(c)}(c)

# interface
value(c::Constant{T, 0}) where {T} = first(c.x)
value(c::Constant) = c.x