#TODO: move shock sensor implementation to own file

#------------------------------------------------------------------------------
# Method from: Persson and Peraire, "Sub-Cell Shock Capturing for DG Methods"
#              AIAA 2006
# Some additional details are from Barter's Ph.D Thesis, "Shock Capturing
# with PDE-Based Artificial Viscosity for an Adaptive, Higher-Order DG
# Finite Element Method", MIT 2008.

"""
  Computes the shock sensor and the numerical viscoscity for the Persson
  and Perairi paper
"""
function getShockSensor(params::ParamType{Tdim}, sbp::AbstractOperator,
                          sensor::ShockSensorPP,
                          q::AbstractMatrix{Tsol}, elnum::Integer,
                          coords::AbstractMatrix,
                          dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                          Se_mat::AbstractMatrix, ee_mat::AbstractMatrix
                         ) where {Tsol, Tmsh, Tdim}

  Tres = promote_type(Tsol, Tmsh)
  numNodesPerElement = size(q, 2)

  @unpack sensor up up_tilde up1_tilde s0 kappa e0
  #fill!(up_tilde, 0); fill!(up1_tilde, 0)

  # use density as the key variable
  for i=1:numNodesPerElement
    up[i] = q[1, i]
  end

  getFilteredSolutions(params, sensor.Vp, up, up_tilde, up1_tilde)
  #getFilteredSolution(params, sensor.Vp, up, up_tilde)
  #getFilteredSolution(params, sensor.Vp1, up, up1_tilde)

  # compute the inner product
  num = zero(Tres)
  den = zero(Tres)

  # use the filtered variables for (u, u) and ||u - u_tilde||.
  # This is a bit different than
  # finite element methods, where the original solution has a basis, and the
  # norm in any basis should be the same.  Here we use the filtered u rather
  # than the original because it is probably smoother.
  if sensor.use_filtered
    for i=1:numNodesPerElement
      fac = sbp.w[i]/jac[i]
      delta_u = up_tilde[i] - up1_tilde[i]

      num += delta_u*fac*delta_u
      den += up_tilde[i]*fac*up_tilde[i]
    end
  else
    # use the original u for (u, u) and ||u - u_tilde||.  For the diagonal E
    # operators, this means u - u_tilde contains the highest 2 degree
    # polynomials (because diagonal E operators have enough nodes for one
    # degree higher solution)
    for i=1:numNodesPerElement
      fac = sbp.w[i]/jac[i]
      delta_u = up[i] - up1_tilde[i]

      num += delta_u*fac*delta_u
      den += up[i]*fac*up[i]
    end
  end

  Se = num/den
  se = log10(Se)

  # should this be a separate function from computing Se?
  if se < s0 - kappa
    ee = Tres(0.0)
    is_nonzero = false
  elseif se > s0 - kappa && se < s0 + kappa
    ee = 0.5*e0*(1 + sinpi( (se - s0)/(2*kappa)))
    is_nonzero = true
  else
    ee = Tres(e0)
    is_nonzero = true
  end

  # scale by lambda_max * h/p to get subcell resolution
  lambda_max = zero(Tsol)
  h_avg = zero(Tmsh)
  for i=1:numNodesPerElement
    q_i = sview(q, :, i)
    lambda_max += getLambdaMax(params, q_i)
    h_avg += sbp.w[i]/jac[i]
  end

  lambda_max /= numNodesPerElement
  h_avg = h_avg^(1/Tdim)

  ee *= lambda_max*h_avg/sbp.degree

  # use elementwise constant
  fill!(Se_mat, Se)
  fill!(ee_mat, ee)

  return is_nonzero
end

"""
  Filter the solution for the shock sensor

  **Inputs**

   * params
   * vand: the [`VandermondeData`](@ref) to use
   * u_nodal: vector, length `numNodesPerElement` containing the solution at the
              nodes of the element

  **Inputs/Outputs**

   * u_filt: vector to be overwritten with the filtered solution, same length
             as `u_nodal`
"""
function getFilteredSolution(params::ParamType, vand::VandermondeData,
                             u_nodal::AbstractVector, u_filt::AbstractVector)

  # u_modal = Vpinv * u_nodal
  # u_filt = Vp * u_modal, thus this can be done in one step by:
  # u_filt = (Vp*Vpinv)*u_nodal
  smallmatvec!(vand.filt, u_nodal, u_filt)

  return nothing
end

"""
  Computed the two filtered solutions required by the shock sensor (well, only
  one of them is required, the other one may or may not be useful)

  **Inputs**

   * params
   * vand: `VandermondeData`
   * u_nodal: `numNodesPerElement` nodal solution

  **Inputs/Outputs**

   * u_filt1: The result of `u = Vp*pinv(Vp)*u_nodal`  This projects the
              solution onto the orthogonal basis and then projects it back.
              This may or may not be useful
   * u_filt2: The result of projecting `u` onto the orthgonal basis and then
              projecting all but the highest mode(s) back to a nodal back to
              the nodal basis
"""
function getFilteredSolutions(params::ParamType, vand::VandermondeData,
                              u_nodal::AbstractVector,
                              u_filt1::AbstractVector, u_filt2::AbstractVector)

  smallmatvec!(vand.filt, u_nodal, u_filt1)
  smallmatvec!(vand.filt1, u_nodal, u_filt2)

  return nothing
end

#------------------------------------------------------------------------------
# ShockSensorNone

function getShockSensor(params::ParamType, sbp::AbstractOperator,
                          sensor::ShockSensorNone,
                          q::AbstractMatrix{Tsol}, elnum::Integer,
                          coords::AbstractMatrix,
                          dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                          Se_mat::AbstractMatrix, ee_mat::AbstractMatrix
                         ) where {Tsol, Tmsh}

  error("getShockSensor called for ShockSensorNone: did you forget to specify the shock capturing scheme?")
end

function isShockElement(params::ParamType, sbp::AbstractOperator,
                          sensor::ShockSensorNone,
                          q::AbstractMatrix{Tsol}, elnum::Integer,
                          coords::AbstractMatrix,
                          dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                         ) where {Tsol, Tmsh}

  error("isShockElement called for ShockSensorNone: did you forget to specify the shock capturing scheme?")
end

#------------------------------------------------------------------------------
# ShockSensorEverywhere

function getShockSensor(params::ParamType, sbp::AbstractOperator,
                        sensor::ShockSensorEverywhere{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        Se_mat::AbstractMatrix, ee_mat::AbstractMatrix
                        ) where {Tsol, Tres, Tmsh}

  fill!(Se_mat, sensor.alpha)
  fill!(ee_mat, sensor.alpha)

  return true
end

function isShockElement(params::ParamType, sbp::AbstractOperator,
                        sensor::ShockSensorEverywhere{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        ) where {Tsol, Tres, Tmsh}

  return true
end

#------------------------------------------------------------------------------
# ShockSensorVelocity

function getShockSensor(params::ParamType, sbp::AbstractOperator,
                        sensor::ShockSensorVelocity{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        Se_mat::AbstractMatrix, ee_mat::AbstractMatrix
                        ) where {Tsol, Tres, Tmsh}

  numNodesPerElement = size(q, 2)
  dim = size(coords, 1)

  alpha = sensor.alpha
  for i=1:numNodesPerElement
    for d=1:dim
      Se_mat[d, i] = alpha*d*q[d+1, i]
      ee_mat[d, i] = alpha*d*q[d+1, i]
    end
  end

  return true
end

function isShockElement(params::ParamType, sbp::AbstractOperator,
                        sensor::ShockSensorVelocity{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        ) where {Tsol, Tres, Tmsh}

  return true
end

#------------------------------------------------------------------------------
# ShockSensorHIso

function getShockSensor(params::ParamType{Tdim}, sbp::AbstractOperator,
                        sensor::ShockSensorHIso{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        Se_mat::AbstractMatrix, ee_mat::AbstractMatrix
                         ) where {Tsol, Tres, Tmsh, Tdim}

  # compute | div(F) |
  # Do this in Cartesian coordinates because it makes the differentiation easier
  numDofPerNode, numNodesPerElement = size(q)
  res = sensor.res

  computeStrongResidual(params, sbp, sensor.strongdata, q, dxidx, jac, res)
  # Note: Hartmann's paper referrs to Jaffre's paper which says to use the
  #       maximum value of res over the element, but that is not differentiable,
  #       so take the L2 norm instead
  val = computeL2Norm(params, sbp, jac, res)
  h_avg = computeElementVolume(params, sbp, jac)

  h_fac = h_avg^((2 - sensor.beta)/Tdim)

  ee = Tres(sensor.C_eps*h_fac*val)
#=
  if ee < 1e-12  # make this zero so the shockmesh machinery knows to exclude
                 # this element, saving cost
    ee = Tres(0)
  end
=#

  fill!(Se_mat, val)
  fill!(ee_mat, ee)

  # calling | div(f) | as the shock sensor, which is somewhat arbitrary
  return true
end

function isShockElement(params::ParamType{Tdim}, sbp::AbstractOperator,
                        sensor::ShockSensorHIso{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                       ) where {Tsol, Tres, Tmsh, Tdim}

  return true
end

"""
  Computes del * F.  where del is the diverage in the xyz directions,
  for a single element.

  **Inputs**

   * params
   * sbp
   * data: `StrongFormData` to hold temporary arrays
   * q: solution on the entire element, `numDofPerNode` x `numNodesPerElement`
   * dxidx: scaled mapping jacobian for entire element, `dim` x `dim` x
            `numNodesPerElement`
   * jac: mapping jacobian determinant for entire element, `numNodesPerElement`

  **Inputs/Outputs**

   * res: array to overwrite with residual, same size as `q`.
"""
function computeStrongResidual(params::ParamType{Tdim}, sbp::AbstractOperator,
                               data::StrongFormData{Tsol, Tres},
                               q::AbstractMatrix,
                               dxidx::Abstract3DArray,
                               jac::AbstractVector,
                               res::AbstractMatrix) where {Tdim, Tsol, Tres}

  @unpack data flux nrm aux_vars work
  fill!(res, 0); fill!(nrm, 0)

  numDofPerNode, numNodesPerElement = size(q)

  for i=1:numNodesPerElement
    q_i = sview(q, :, i)
    for d=1:Tdim
      nrm[d] = 1
      flux_d = sview(flux, :, i, d)

      calcEulerFlux(params, q_i, aux_vars, nrm, flux_d)

      nrm[d] = 0
    end
  end

  applyDx(sbp, flux, dxidx, jac, work, res)

  return nothing
end

"""
  Computes the integral L2 norm for a single element

  **Inputs**

   * params
   * sbp
   * jac: mapping jacobian determinant for the element, `numNodesPerElement`
   * res: `numDofPerNode` x `numNodesPerElement` array to compute the norm
          of.
"""
function computeL2Norm(params::ParamType{Tdim}, sbp::AbstractOperator,
                       jac::AbstractVector, res::AbstractMatrix{Tres}
                      ) where {Tres, Tdim}

  numDofPerNode, numNodesPerElement = size(res)

  # compute norm
  val = zero(Tres)
  for i=1:numNodesPerElement
    fac = sbp.w[i]/jac[i]
    for j=1:numDofPerNode
      val += res[j, i]*fac*res[j, i]
    end
  end

  val = sqrt(val)

  return val
end

"""
  Computes volume of element by integrating the mapping jacobian determinant
"""
function computeElementVolume(params::ParamType{Tdim}, sbp::AbstractOperator,
                         jac::AbstractVector{Tmsh}) where {Tdim, Tmsh}

  # compute norm and mesh size h
  vol = zero(Tmsh)
  for i=1:length(jac)
    vol += sbp.w[i]/jac[i]
  end

  return vol
end

#------------------------------------------------------------------------------
# ShockSensorBO

function getShockSensor(params::ParamType{Tdim}, sbp::AbstractOperator,
                        sensor::ShockSensorBO{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        Se_mat::AbstractMatrix, ee_mat::AbstractMatrix
                         ) where {Tsol, Tres, Tmsh, Tdim}

  numDofPerNode, numNodesPerElement = size(q)

  lambda_max = zero(Tsol)
  h_avg = zero(Tmsh)
  for i=1:numNodesPerElement
    q_i = sview(q, :, i)
    lambda_max += getLambdaMax(params, q_i)
    h_avg += sbp.w[i]/jac[i]
  end

  lambda_max /= numNodesPerElement
  h_avg = h_avg^(1/Tdim)

  fill!(Se_mat, lambda_max)
  fill!(ee_mat, sensor.alpha*lambda_max*h_avg/sbp.degree)
#  println("element ", elnum, " SensorBO value = ", real(ee_mat[1, 1]))

  return true
end

function isShockElement(params::ParamType{Tdim}, sbp::AbstractOperator,
                        sensor::ShockSensorBO{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                       ) where {Tsol, Tres, Tmsh, Tdim}

  return true
end

#------------------------------------------------------------------------------
# ShockSensorHHO

function getShockSensor(params::ParamType{Tdim}, sbp::AbstractOperator,
                        sensor::ShockSensorHHO{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        Se_mat::AbstractMatrix, ee_mat::AbstractMatrix
                       ) where {Tsol, Tres, Tmsh, Tdim}

  numDofPerNode, numNodesPerElement = size(q)

  @unpack sensor p_dot press_el press_dx work res Rp fp
  fill!(press_dx, 0); fill!(res, 0); fill!(Rp, 0)

  # compute Rm
  computeStrongResidual(params, sbp, sensor.strongdata, q, dxidx, jac, res)

  # compute Rp
  @simd for i=1:numNodesPerElement
    q_i = ro_sview(q, :, i)
    press = calcPressure_diff(params, q_i, p_dot)
    press_el[1, i] = press

    @simd for j=1:numDofPerNode
      Rp[i] += p_dot[j]*res[j, i]/press
    end
  end

  # compute pressure switch fp (excluding the h_k factor)
  applyDx(sbp, press_el, dxidx, jac, work, press_dx)
  @simd for i=1:numNodesPerElement
    val = zero(Tres)
    @simd for d=1:Tdim
      val += press_dx[1, i, d]*press_dx[1, i, d]
    end
    fp[i] = sqrt(val)/(press_el[1, i] + 1e-12)
  end

  @simd for i=1:numNodesPerElement
    @simd for d=1:Tdim
      h_fac = sensor.h_k_tilde[d, elnum]
      Se_mat[d, i] = h_fac*fp[i]*absvalue(Rp[i])
      ee_mat[d, i] = Se_mat[d, i]*h_fac*h_fac*sensor.C_eps
    end
  end

  return true
end

#------------------------------------------------------------------------------
# ShockSensorHHOConst

function getShockSensor(params::ParamType{Tdim}, sbp::AbstractOperator,
                        sensor::ShockSensorHHOConst{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        Se_mat::AbstractMatrix, ee_mat::AbstractMatrix
                       ) where {Tsol, Tres, Tmsh, Tdim}

  numDofPerNode, numNodesPerElement = size(q)

  @unpack sensor p_dot press_el press_dx work res Rp fp epsilon
  fill!(press_dx, 0); fill!(res, 0); fill!(Rp, 0)

  # compute Rm
  computeStrongResidual(params, sbp, sensor.strongdata, q, dxidx, jac, res)

  # compute Rp
  @simd for i=1:numNodesPerElement
    q_i = ro_sview(q, :, i)
    press = calcPressure_diff(params, q_i, p_dot)
    press_el[1, i] = press

    @simd for j=1:numDofPerNode
      Rp[i] += p_dot[j]*res[j, i]/press
    end
  end

  # compute pressure switch fp (excluding the h_k factor)
  applyDx(sbp, press_el, dxidx, jac, work, press_dx)
  @simd for i=1:numNodesPerElement
    val = zero(Tres)
    @simd for d=1:Tdim
      val += press_dx[1, i, d]*press_dx[1, i, d]
    end
    fp[i] = sqrt(val)/(press_el[1, i] + 1e-12)
  end

  # compute L2 norm of fp * Rp
  @simd for i=1:numNodesPerElement
    epsilon[1, i] = fp[i]*Rp[i]
  end

  #val = computeL2Norm(params, sbp, jac, epsilon)
  val = zero(Tres)
  @simd for i=1:numNodesPerElement
    val += sbp.w[i]*absvalue(epsilon[1, i])
  end

  @simd for i=1:numNodesPerElement
    @simd for d=1:Tdim
      h_fac = sensor.h_k_tilde[d, elnum]
      Se_mat[d, i] = h_fac*val
      ee_mat[d, i] = Se_mat[d, i]*h_fac*h_fac*sensor.C_eps
    end
  end

  return true
end

function isShockElement(params::ParamType{Tdim}, sbp::AbstractOperator,
                        sensor::HHOSensors{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                       ) where {Tsol, Tres, Tmsh, Tdim}
  return true
end

"""
  Compute the h_k tilde factors from the paper
"""
function calcAnisoFactors(mesh::AbstractMesh, sbp, opts,
                          hk_all::AbstractMatrix{T}) where {T}

  # compute h_k tilde = h_k/(p + 1), where h_k takes into account the
  # anisotropy of the element in directory k
  # compute the p_k first, then compute h_k, where p_k is the integral of the
  # face normal

  @assert size(hk_all) == (mesh.dim, mesh.numEl)

  fill!(hk_all, 0)

  # do all faces (interior, boundary, shared)
  @simd for i=1:mesh.numInterfaces
    iface_i = mesh.interfaces[i]
    nrm_i = ro_sview(mesh.nrm_face, :, :, i)

    @simd for j=1:mesh.numNodesPerFace
      @simd for k=1:mesh.dim
        fac = mesh.sbpface.wface[j]*absvalue2(nrm_i[k, j])
        hk_all[k, iface_i.elementL] += fac
        hk_all[k, iface_i.elementR] += fac
      end
    end
  end

  @simd for i=1:mesh.numBoundaryFaces
    bndry_i = mesh.bndryfaces[i]
    nrm_i = ro_sview(mesh.nrm_bndry, :, :, i)

    @simd for j=1:mesh.numNodesPerFace
      @simd for k=1:mesh.dim
        fac = mesh.sbpface.wface[j]*absvalue2(nrm_i[k, j])
        hk_all[k, bndry_i.element] += fac
      end
    end
  end

  @simd for peer=1:mesh.npeers
    @simd for i=1:length(mesh.shared_interfaces[peer])
      iface_i = mesh.shared_interfaces[peer]
      nrm_i = ro_sview(mesh.nrm_sharedface[peer], :, :, i)

      @simd for j=1:mesh.numNodesPerFace
        @simd for k=1:mesh.dim
          fac = mesh.sbpface.wface[j]*absvalue2(nrm_i[k, j])
          hk_all[k, iface_i.elementL] += fac
        end
      end
    end
  end

  # hk_all now contains the p_i in each direction
  # Now compute h_k tilde from p_i
  @simd for i=1:mesh.numEl

    jac_i = ro_sview(mesh.jac, :, i)
    vol = zero(T)
    @simd for j=1:mesh.numNodesPerElement
      vol += sbp.w[j]/jac_i[j]
    end

    h_k = vol^(1/mesh.dim)
    pk_prod = one(T)
    @simd for d=1:mesh.dim
      pk_prod *= hk_all[d, i]
    end
    pk_prod = pk_prod^(1/mesh.dim)

    @simd for d=1:mesh.dim
      hk_all[d, i] = h_k*pk_prod/(hk_all[d, i]*(sbp.degree + 1))
    end
  end

  return nothing
end

#------------------------------------------------------------------------------
# SensorHApprox

# even though this shared the name with all the other getShockSensor methods,
# it signature and outputs are different.  This is sort-of ok because this
# is a special shock sensor and should never be used as the main shock sensor.
function getShockSensor(params::ParamType{Tdim}, sbp::AbstractOperator,
                        sensor::ShockSensorHApprox{Tsol, Tres},
                        elnum::Integer, jac::AbstractVector{Tmsh},
                       ) where {Tsol, Tres, Tmsh, Tdim}

  numNodesPerElement = length(jac)

  shared_idx = elnum - first(sensor.shared_els) + 1  # index of shared element
  if (elnum in sensor.local_els) ||
     (elnum in sensor.shared_els && sensor.shared_isShocked[shared_idx])

    # compute element size h for both elements
    h = zero(Tmsh)
    @simd for i=1:numNodesPerElement
      h += sbp.w[i]/jac[i]
    end

    # compute estimates viscosity
    lambda_max = sensor.lambda_max
    eps_L = sensor.alpha*lambda_max*(h^(1/Tdim))/sbp.degree

#    println("element ", elnum, " HApprox value = ", real(eps_L), ", with h = ")
    return Tres(eps_L)
  else
    return zero(Tres)
  end

end

function isShockElement(params::ParamType, sbp::AbstractOperator,
                          sensor::ShockSensorHApprox,
                          q::AbstractMatrix{Tsol}, elnum::Integer,
                          coords::AbstractMatrix,
                          dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                         ) where {Tsol, Tmsh}

  error("ShockSensorHApprox does not know which elements are shocked or not, don't ask it that.")
end

#------------------------------------------------------------------------------
# ShockSensorOddBO

function getShockSensor(params::ParamType, sbp::AbstractOperator,
                        sensor::ShockSensorOddBO{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        Se_mat::AbstractMatrix, ee_mat::AbstractMatrix
                        ) where {Tsol, Tres, Tmsh}

  if elnum % 2 == 1
    getShockSensor(params, sbp, sensor.sensor, q, elnum, coords, dxidx, jac,
                   Se_mat, ee_mat)
  else
    fill!(Se_mat, 0)
    fill!(ee_mat, 0)
  end

  return true
end

function isShockElement(params::ParamType, sbp::AbstractOperator,
                        sensor::ShockSensorOddBO{Tsol, Tres},
                        q::AbstractMatrix, elnum::Integer,
                        coords::AbstractMatrix,
                        dxidx::Abstract3DArray, jac::AbstractVector{Tmsh},
                        ) where {Tsol, Tres, Tmsh}

  return (elnum % 2 ) == 1
end


"""
```
B, flags = pmf(A; ωₘ, slope, dhₘ, dh₀, cellsize)
```
Applies the progressive morphological filter by *Zhang et al. (2003)* [^zhang2003] to `A`.

# Output
- `B::Array{T,2}` Maximum allowable values
- `flags::Array{Float64,2}` A sized array with window sizes if filtered, zero if not filtered.

Afterwards, one can retrieve the resulting mask for `A` by `A .<= B` or `flags .== 0.`.

# Arguments
- `A::Array{T,2}` Input Array
- `ωₘ::Float64=20.` Maximum window size [m]
- `slope::Float64=0.01` Terrain slope [m/m]
- `dhₘ::Float64=2.5` Maximum elevation threshold [m]
- `dh₀::Float64=0.2` Initial elevation threshold [m]
- `cellsize::Float64=1.` Cellsize in [m]

[^zhang2003]: Zhang, Keqi, Shu-Ching Chen, Dean Whitman, Mei-Ling Shyu, Jianhua Yan, and Chengcui Zhang. “A Progressive Morphological Filter for Removing Nonground Measurements from Airborne LIDAR Data.” IEEE Transactions on Geoscience and Remote Sensing 41, no. 4 (2003): 872–82. <https://doi.org/10.1109/TGRS.2003.810682>.
"""
function pmf(A::AbstractMatrix{T};
    ωₘ::Float64 = 20.0,
    slope::Float64 = 0.01,
    dhₘ::Float64 = 2.5,
    dh₀::Float64 = 0.2,
    cellsize::Float64 = 1.0,
    circular = false) where {T<:Real}

    # Compute windowsizes and thresholds
    ωₘ = round(Int, ωₘ / cellsize)
    κ_max = floor(Int, log2(ωₘ - 1))  # determine # iterations based on exp growth
    windowsizes = Int.(exp2.(1:κ_max)) .+ 1

    # Compute tresholds
    dwindows = vcat(windowsizes[1], windowsizes)  # prepend first element so we get 0 as diff
    window_diffs = [dwindows[i] - dwindows[i-1] for i = 2:length(dwindows)]
    height_tresholds = [min(dhₘ, slope * window_diff * cellsize + dh₀) for window_diff in window_diffs]
    @info "Using the following thresholds: $height_tresholds for the following windows: $windowsizes"

    # Set up arrays
    Af = copy(A)  # array to be morphed
    nan_mask = isnan.(Af)
    Af[nan_mask] .= Inf  # Replace NaN with Inf, as to always filter these

    B = copy(A)  # max_elevation raster
    out = copy(A)  # max_elevation raster

    flags = zeros(size(A))  # 0 = ground, other values indicate window size
    flags[nan_mask] .= NaN

    mask = falses(size(A))

    # Iterate over window sizes and height tresholds
    for (ωₖ, dhₜ) in zip(windowsizes, height_tresholds)
        if circular
            opening_circ!(Af, ωₖ, out)
        else
            opening!(Af, ωₖ, out)
        end
        mask .= (A .- Af) .> dhₜ
        for I in eachindex(flags)
            if mask[I] && flags[I] == 0
                flags[I] = ωₖ
            end
        end
        B .= min.(B, Af .+ dhₜ)
    end

    B, flags
end


function cable_transform(y, z)
    v1 = [0.0, 0.0, 1.0]
    # if norm(z) > norm(y)
    #     v2 = y[1:3,1] - z[1:3,1]
    # else
    #     v2 = z[1:3,1] - y[1:3,1]
    # end
    v2 = y[1:3,1] - z[1:3,1]
    normalize!(v2)
    ax = cross(v1, v2)
    ang = acos(v1'*v2)
    R = AngleAxis(ang, ax...)

    if any(isnan.(R))
        R = I
    else
        nothing
    end

    compose(Translation(z), LinearMap(R))
end

function cast3d(v::AbstractVector{T}) where {T}
	l = length(v)
	if l == 3
		return v
	elseif l == 2
		return [v[1], 0.0, v[2]]
	else
		error("Wrong vector length")
	end
end

function bivector_rotation(a::AbstractVector{T}, b::AbstractVector{T}) where {T}
	na = norm(a)
	nb = norm(b)
	v = cross(a,b)/(na*nb)
	s = norm(v)
	c = (a'*b)/(na*nb)
	vx = skew(v)
	α = (1 - c)/s^2
	R = I + vx + vx*vx*α
	return R
end

function default_background!(vis; grid::Bool=false, axes::Bool=false, zoom=1.0)
    setvisible!(vis["/Background"], true)
    setprop!(vis["/Background"], "top_color", RGBA(1.0, 1.0, 1.0, 1.0))
    setprop!(vis["/Background"], "bottom_color", RGBA(1.0, 1.0, 1.0, 1.0))
	setvisible!(vis["/Axes"], axes)
    # setvisible!(vis["/Grid"], grid)
	# setprop!(vis["/Cameras/default/rotated/<object>"], "zoom", zoom)
end

function get_line_material(size::Real; α=1.0)
    orange_col = [1,153/255,51/255]
    blue_col = [51/255,1,1]
    black_col = [0,0,0]
    orange_mat = LineBasicMaterial(color=color=RGBA(orange_col...,α), linewidth=size)
    blue_mat = LineBasicMaterial(color=color=RGBA(blue_col...,α), linewidth=size)
    black_mat = LineBasicMaterial(color=color=RGBA(black_col...,α), linewidth=size)
    return orange_mat, blue_mat, black_mat
end

function get_material(;α=1.0)
    orange_col = [1,153/255,51/255]
    blue_col = [51/255,1,1]
    black_col = [0,0,0]
    orange_mat = MeshPhongMaterial(color=RGBA(orange_col...,α))
    blue_mat = MeshPhongMaterial(color=RGBA(blue_col...,α))
    black_mat = MeshPhongMaterial(color=RGBA(black_col...,α))
    return orange_mat, blue_mat, black_mat
end


println(sum(filter(x -> x%3 == 0 || x%5 == 0, 1:999)))


type Biconnected_components_output
    bcc_edge_list::Vector{Int64}  #list of bcc edges
    components::Vector{Int64} #Indices of bcc corresponding to the edges in List
    number::Int64 # Total number of Biconnected components
    A::MatrixNetwork #MatrixNetwork
end

function BiCC(A::MatrixNetwork)
    n=A.n
    rp=A.rp
    ci=A.ci
    bcc_edges=zeros(Int64,2*length(ci))
    bcc_edges_count=0
    cn=1
    low=zeros(Int64,n)
    dt=zeros(Int64,n)
    pred=zeros(Int64,n)
    components=-1 * ones(Int64,length(ci))
    t=0
    cs=zeros(Int64,2*length(ci))
    css=0
    rs=zeros(Int64,2*n)
    rss=0 #recursion stack (v,ri)
    children = zeros(Int64, n)
    #start dfs at 1
    for sv=1:n
        v=sv
        if dt[v]>0
            continue
        end
        rss=rss+1
        rs[2*rss-1]=v
        rs[2*rss]=rp[v] #add v,rp[v] to stack
        low[v]=t
        dt[v]=t;
        t=t+1;

        while rss>0
            v=rs[2*rss-1]
            ri=rs[2*rss]
            rss=rss-1
            while ri<rp[v+1]
                w=ci[ri]
                ri=ri+1
                if dt[w] == 0
                    low[w]=w; dt[w]=t; t=t+1;
                    css=css+1
                    cs[2*css-1]=v
                    cs[2*css]=w
                    children[v]=children[v]+1
                    pred[w]=v
                    rss=rss+1
                    rs[2*rss-1]=v
                    rs[2*rss]=ri
                    v=w
                    ri=rp[w]
                    continue
                elseif(w != pred[v] && dt[w] < low[v])
 		    if low[v] > dt[w]
			low[v] = dt[w]
		    end
                    css=css+1
                    cs[2*css-1]=v
                    cs[2*css]=w
                end
            end
             for ri=rp[v]:rp[v+1]-1
                w=ci[ri]
                if components[w] == -1
                    if low[v] > low[w]
                        low[v]=low[w]
                    end
                end
                if (( (dt[v] == 1 ) && ( children[v] > 1 ) ) || ( dt[v] > 1 && low[w] >= dt[v]))
                    while css > 0
                        bcc_edges_count = bcc_edges_count + 1
                        bcc_edges[2*bcc_edges_count-1] = cs[2*css-1]
                        bcc_edges[2*bcc_edges_count] = cs[2*css]
                        components[bcc_edges_count] = cn
			if (cs[2*css-1] == v && cs[2*css] ==w)
				css=css-1
				break
			end
                        css = css-1
                    end

                    cn=cn+1
                end
            end

          end
	end
	print(components)
          return bcc_edges
end


"""
Julia expression wrapper
"""
struct Expression
    expr::Expr
    syms::Dict{Symbol,Int}
end
function Expression(ex::Expr)
    syms = Dict( s=>i for (i,s) in pairs(sort!(symbols(ex))) )
    Expression(ex, syms)
end
show(io::IO, e::Expression) = infix(io, e.expr)
(e::Expression)(vals::T...) where {T} = evaluate(e.expr, e.syms, vals...)

function symbols(ex::Expr)
    syms = Symbol[]
    for e in ex.args
        isa(e, Symbol) && push!(syms, e)
        isa(e, Expr) && append!(syms, symbols(e))
    end
    unique!(syms)
end

function compile(ex::Expr, params::Vector{Symbol})
    tprm = Expr(:tuple, params...)
    Expr(:->, tprm, ex)
end

height(ex) = isa(ex, Expr) ? maximum( height(e) for e in ex.args )+1 : 0
nodes(ex) = !isa(ex, Expr) ? 1 : length(ex.args) > 0 ? sum( nodes(e) for e in ex.args ) : 0
length(ex) = nodes(ex)

function depth(root, ex; d=0)
    return if root == ex
        d
    elseif isa(root, Expr)
        maximum( depth(e, ex, d=d+1) for e in root.args )
    else
        -1
    end
end

function copyto!(dst::Expr, src::Expr)
    dst.head = src.head
    dst.args = deepcopy(src.args)
    dst
end

function randsubexpr(ex)
    !isa(ex, Expr) && return ex
    csize = map(nodes, ex.args[2:end])
    cidx = rand(1:sum(csize)+1)-1
    if cidx == 0
        ex
    else
        nidx = findfirst(i->cidx<=i, cumsum(csize))
        randsubexpr(ex.args[nidx+1])
    end
end

function getindex(ex::Expr, idx::Int)
    if idx == 0
        ex
    else
        csize = cumsum(map(nodes, ex.args[2:end]))
        nidx = findfirst(i->i>=idx, csize)
        nex = ex.args[nidx+1]
        !isa(nex, Expr) ? nex : nex[csize[nidx]-idx]
    end
end

function setindex!(ex::Expr, subex, idx::Int)
    csize = cumsum(map(nodes, ex.args[2:end]))
    nidx = findfirst(i->i>=idx, csize)
    nex = ex.args[nidx+1]
    if !isa(nex, Expr) || csize[nidx]-idx == 0
        ex.args[nidx+1] = subex
    else
        nex[csize[nidx]-idx] = subex
    end
end

isnum(ex) = isa(ex, Number)
iszeronum(root) = isnum(root) && iszero(root)
isonenum(root) = isnum(root) && isone(root)
isdiv(ex) = ex in [/, div, pdiv, aq]
isexpr(ex) = isa(ex, Expr)
issym(ex) = isa(ex, Symbol)
isexprsym(ex) = isexpr(ex) || issym(ex)
isbinexpr(ex) = isexpr(ex) && length(ex.args) == 3

function evaluate(ex::Expr, psyms::Dict{Symbol,Int}, vals::T...)::T where {T}
    exprm = ex.args
    exvals = (isexpr(nex) || issym(nex) ? evaluate(nex, psyms, vals...) : nex for nex in exprm[2:end])
    exprm[1](exvals...)
end

function evaluate(ex::Symbol, psyms::Dict{Symbol,Int}, vals::T...)::T where {T}
    pidx = get(psyms, ex, 0)
    if pidx == 0
        @error "Undefined symbol: $ex"
    end
    return T(vals[pidx])
end

function simplifyunary!(root)
    fn, op = root.args
    if isexpr(op)
        fn2, op2 = op.args
        if fn == log && fn2 == exp
            return op2
        elseif fn == exp && fn2 == log
            return op2
        end
    end
    return root
end

function simplifybinary!(root)
    fn, op1, op2 = root.args

    # some elementary simplification rules
    if fn == (-) && op1 == op2 # additive inverse: x-x = 0
        root = 0
    elseif isdiv(fn) && (op1 == op2 || iszeronum(op2)) # look for multiplicative inverse x/x = 1
        root = 1
    elseif (fn == (*) || isdiv(fn)) &&  (iszeronum(op1) || iszeronum(op2)) # look for 0: x*0 = 0*x = 0
        root = 0
    elseif (fn == (+) || fn == (-)) && iszeronum(op2) # look for 0: x ± 0 = x
        root = op1
    elseif fn == (+) && iszeronum(op1) # look for 0: 0 + x = x
        root = op2
    elseif (fn == (*) || isdiv(fn)) && isonenum(op2) # x*1 = x || x/1 = x
        root = op1
    elseif fn == (*) && isonenum(op1) # 1*x = x
        root = op2
    elseif (fn == (+)) && (op1 == op2) # x+x = 2x
        root.args[1] = (*)
        root.args[2] = 2
    elseif (fn == (+) || fn == (-)) && (isbinexpr(op1) && isnum(op2))
        fn2, op11, op12 = op1.args
        if fn == (+) && fn2 == (-) && isnum(op12)
            # (x-m)+n = x+(n-m)
            root.args[2] = op11
            root.args[3] = op2-op12
        elseif fn == (-) && fn2 == (-) && isnum(op12)
            # (x-m)-n = x-(n+m)
            root.args[2] = op11
            root.args[3] = op12+op2
        elseif fn2 == (+) && (isnum(op11) || isnum(op12))
            # (m+x)±n = (x+m)±n = x+(n±m)
            var, n2 = isnum(op11) ? (op12, op11) : (op11, op12)
            n3 = fn(n2, op2)
            root.args[1] = fn2
            root.args[2] = var
            root.args[3] = n3
        elseif fn2 == (-) && isnum(op11)
            # (m-x)±n = (n±m)-x
            n3 = fn(op11, op2)
            root.args[1] = fn2
            root.args[2] = n3
            root.args[3] = op12
        end
    elseif fn == (+) && isnum(op1) && isbinexpr(op2)
        fn2, op21, op22 = op2.args
        if fn2 == (+) || fn2 == (-)
            if isnum(op21)
                # n+(m±x) = (n+m)±x
                root.args[1] = fn2
                root.args[2] = op1 + op21
                root.args[3] = op22
            elseif isnum(op22)
                # n+(x±m) = (n±m)+x
                root.args[1] = fn
                root.args[2] = fn2(op1, op22)
                root.args[3] = op21
            end
        end
    elseif fn == (-) && isnum(op1) && isbinexpr(op2)
        fn2, op21, op22 = op2.args
        if fn2 == (+) && (isnum(op21) || isnum(op22))
            # n-(x+m) = n-(m+x) = (n-m)-x
            var, n2 = isnum(op21) ? (op22, op21) : (op21, op22)
            root.args[2] = op1 - n2
            root.args[3] = var
        elseif fn2 == (-) && (isnum(op21) || isnum(op22))
            # n-(m-x) = (n-m)+x
            # n-(x-m) = (n+m)-x
            var, n2, f1p = isnum(op21) ? (op22, op21, true) : (op21, op22, false)
            root.args[1] = f1p ? (+) : (-)
            root.args[2] = f1p ? op1 - n2 : op1 + n2
            root.args[3] = var
        end
    elseif fn == (+) && (isexpr(op1) || isexpr(op2))
        if isbinexpr(op2)
            # x + (n - x) = n
            fn2, op21, op22 = op2.args
            if fn2 == (-) && op1 == op22
                root = op21
            end
        elseif isbinexpr(op1)
            # (n - x) + x = n
            fn2, op11, op12 = op1.args
            if fn2 == (-) && op12 == op2
                root = op11
            end
        end
    elseif fn == (-) && (isexpr(op1) || isexpr(op2))
        if isbinexpr(op2)
            fn2, op21, op22 = op2.args
            if op1 == op21
                # x - (x ± n) = -±n
                if fn2 == (-)
                    root = op22
                else
                    pop!(root.args)
                    root.args[end] = op22
                end
            elseif op1 == op22 && fn2 == (+)
                # x - (n + x) = -n
                pop!(root.args)
                root.args[end] = op21
            end
        elseif isbinexpr(op1)
            # (x ± n) - x = ±n
            fn2, op11, op12 = op1.args
            if op11 == op2
                if fn2 == (+)
                    root = op12
                else
                    root.args[1] = fn2
                    root.args[2] = op12
                    pop!(root.args)
                end
            end
        end
    end

    # evaluate numerical arithmetic
    isnum(op1) && isnum(op2) && return fn(op1, op2)

    return root
end

"""
    simplify!(expr)

Simplify an algebraic expression
"""
function simplify!(root)
    # return if not an expression
    !isa(root, Expr) && return root
    # return if not a function
    root.head != :call && return root

    n = length(root.args)
    # recursively simplify arguments
    for i in 2:n
        !isexpr(root.args[i]) && continue
        root.args[i] = simplify!(root.args[i])
    end

    n == 2 && return simplifyunary!(root)
    n == 3 && return simplifybinary!(root)
    return root
end

function infix(io::IO, root; digits=3)
    if isa(root, Number)
        print(io, round(root, digits=digits))
    elseif isa(root, Expr) && root.head == :call
        if root.args[1] ∈ [+, -, *, /, ^]
            print(io, "(")
            infix(io, root.args[2])
            infix(io, root.args[1])
            length(root.args)>2 && infix(io, root.args[3])
            print(io, ")")
        else
            infix(io, root.args[1])
            print(io, "(")
            for (i, ex) in enumerate(root.args[2:end])
                i > 1 && print(io, ", ")
                infix(io, ex)
            end
            print(io, ")")
        end
    else
        print(io, root)
    end
end
show(io::IO, ::MIME"text/html", e::Expression) = infix(io, e.expr)

latex(io::IO, root::Number; digits=3) = print(io, round(root, digits=digits))
latex(io::IO, root::Symbol; kwargs...) = print(io, root)
function latex(io::IO, root::Expr; digits=3)
    root.head != :call && print(io, expr)
    if root.args[1] == (/)
        print(io, "\\frac{")
        latex(io, root.args[2], digits=digits)
        print(io, "}{")
        latex(io, root.args[3], digits=digits)
        print(io, "}")
    elseif root.args[1] ∈ [+, -, *]
        print(io, "\\left(")
        latex(io, root.args[2], digits=digits)
        print(io, root.args[1])
        latex(io, root.args[3], digits=digits)
        print(io, "\\right)")
    elseif root.args[1] == (^)
        print(io, "{")
        latex(io, root.args[2], digits=digits)
        print(io, "}^{")
        latex(io, root.args[3], digits=digits)
        print(io, "}")
    else
        print(io, "\\")
        print(io, root.args[1])
        print(io, "\\left(")
        for ex in root.args[2:end]
            latex(io, ex, digits=digits)
        end
        print(io, "\\right)")
    end
end
show(io::IO, ::MIME"text/latex", e::Expression) = latex(io, e.expr)


"""
Methods for the query generator: specifically, to

1. generate sparsity coefficients b and subsampling matrices M
2. get the indices of a signal subsample
3. compute a subsampled and delayed Walsh-Hadamard transform.
"""

# potential TODO figure out how to load in parts of the signal bit by bit
# see how Amirali did it

include("input_signal.jl")
using StatsBase
using Random

"""
A semi-arbitrary fixed choice of the sparsity coefficient.
See get_b for full signature.
"""
function get_b_simple(signal::Signal)
    # I don't know why, but changing this to signal.n \div 3 makes it break
    if signal.n == 4
        return 2
    elseif (signal.n == 8) || (signal.n == 16)
        return 4
    elseif (signal.n < 16)
        return signal.n ÷ 2
    else
        return signal.n ÷ 4
    end
end

get_b_lookup = Dict(
    :simple => get_b_simple,
    :random => get_b_simple
)

"""
Get the sparsity coefficient for the signal.

Arguments
---------
signal : TestSignal
The signal whose WHT we want.

method : Symbol
The method to use. All methods referenced must use this signature (minus "method".)

Returns
-------
b : Int64
The sparsity coefficient.
"""
function get_b(signal::Signal; method=:simple)::Int64
    return get_b_lookup[method](signal)
end

function get_partial_I(n, b, k)
    M = falses(n, b)
    for j in 1:b
        M[k+j, j] = 1
    end
    return M
end

"""
A semi-arbitrary fixed choice of the subsampling matrices. See get_Ms for full signature.
"""
function get_Ms_simple(n::Int64, b::Int64)
    # @assert n % b == 0 "n must be exactly divisible by b"
    num_to_get = Int64(ceil(n / b))

    Ms = BitArray[]
    push!(Ms, get_partial_I(n, b, n-b))
    for i in num_to_get-2:-1:0
        push!(Ms, get_partial_I(n, b, b*i))
    end
    return Ms
end

function get_Ms_random(n::Int64, b::Int64)
    num_to_get = n ÷ b
    Ms = BitArray[]

    for i in 1:num_to_get
        push!(Ms, bitrand(n, b))
    end
    return Ms
end

get_Ms_lookup = Dict(
    :simple => get_Ms_simple,
    :random => get_Ms_random
)

function get_Ms(n::Int64, b::Int64; method=:simple)::Array{BitArray,1}
    return get_Ms_lookup[method](n, b)
end

"""
Gets the delays matrix [0; I], of dimension (n+1, n). See get_D for full signature.
"""
function get_D_identity_like(n::Int64; kwargs...)
    return BitArray(vcat(zeros(1, n), Matrix{Bool}(I, n, n)))
end

"""
Gets a random delays matrix of dimension (num_delays, n). See get_D for full signature.
"""
function get_D_random(n::Int64; kwargs...)
    num_delays = kwargs[:num_delays]
    choices = sample(0:2^n-1, num_delays, replace=false)
    return @pipe dec_to_bin.(choices, n) |> hcat(_...) |> transpose |> BitArray
end

"""
Get a repetition code based (NSO-SPRIGHT) delays matrix. See get_D for full signature.
Not sure of correctness based on the paper: this is not (num_delays, n).
"""
function get_D_nso(n::Int64; kwargs...)::BitArray{2}
    num_delays = kwargs[:num_delays]
    p1 = num_delays ÷ n
    random_offsets = get_D_random(n; num_delays=p1)
    D = BitArray(undef, 0, n)
    identity_like = get_D_identity_like(n)
    for row in eachrow(random_offsets)
        modulated_offsets = @pipe [(row .⊻ r) for r in eachrow(identity_like)] |> hcat(_...) |> transpose
        D = vcat(D, modulated_offsets)
    end
    return D
end

function get_generator_matrix(codelen::Int64, n::Int64, code::Symbol)
    return zeros((codelen, n))
end

"""
Get a channel code based (SO-SPRIGHT) delays matrix. See get_D for full signature.
Returns a matrix of dimension (p1+p2+p3, n).
"""
function get_D_so(n::Int64; kwargs...)::BitArray{2}
    p1, p2, p3 = kwargs[:num_delays]
    code = kwargs[:code]
    D1 = get_D_random(n, num_delays=p1)
    D2 = zeros((p2, n))
    D3 = get_generator_matrix(p3, n, code)
    return vcat(D1, D2, D3)
end

get_D_lookup = Dict(
    :identity_like => get_D_identity_like,
    :random => get_D_random,
    :nso => get_D_nso,
    :so => get_D_so
)

"""
Delay generator: gets a delays matrix.

Arguments
---------
n : Int64

number of bits: log2 of the signal length.

Returns
-------
D : num_delays×n Array{Bool,2}

The delays matrix; if num_delays is not specified in kwargs, see the relevant sub-function for a default.
"""
function get_D(n::Int64; method=:identity_like, kwargs...)
    return get_D_lookup[method](n; kwargs...)
end

"""
Query generator: creates indices for signal subsamples.
Arguments
---------
M : n×b Array{Bool,2}
The subsampling matrix.

d : n-element Array{Bool,1}
The subsampling offset.

Returns
-------
indices : B-element Array{Int64,1}
The (decimal) subsample indices.
"""
function subsample_indices(M, d)
    L = binary_ints(size(M, 2))
    inds_binary = Bool.((M*L .⊻ d) .% 2)
    return @pipe inds_binary |> eachcol |> collect |> bin_to_dec.(_) .+ 1
end

"""
Creates random delays, subsamples according to M and the random delays,
and returns the subsample WHT along with the delays.

Arguments
---------
signal : Signal
The signal to subsample, delay, and compute the WHT of.

M : n×b Array{Bool,2}
The subsampling matrix.

num_delays : Int64
The number of delays to apply; or, the number of rows in the delays matrix.

force_identity_like : Bool
Whether to make D = [0; I] like in the noiseless case; for debugging.
"""
function compute_delayed_wht(signal::Signal, M, D)
    return compute_delayed_subtransform(signal, M, D, fwht)
end

function compute_delayed_fft(signal::Signal, M, D)
    return compute_delayed_subtransform(signal, M, D, fft)
end

function compute_delayed_subtransform(signal::Signal, M, D, transform::Function)
    inds = map(d -> subsample_indices(M, d), D |> eachrow |> collect)
    used_inds = reduce(union, inds)
    samples_to_transform = map(x -> signal[x], inds)
    return hcat(transform.(samples_to_transform)...), used_inds
end


module Mod1DynamicDefinition
using ValidatedNumerics
using ..DynamicDefinition, ..Contractors

using ..DynamicDefinition: derivative

export Mod1Dynamic, preim, nbranches, plottable

"""
Defines a Dynamic on [0,1] as the Mod-1 quotient of a given map.

An alternative newer implementation relying on piecewise-defined functions is in `mod1_dynamic`
"""
struct Mod1Dynamic{FT} <: MarkovDynamic
	T::FT
	nbranches::Int
	orientation::Float64
	domain::Interval{Float64}
	is_full_branch::Bool
end

Mod1Dynamic(T::FT, nbranches = undef, domain = Interval{Float64}(0,1)) where {FT} = Mod1Dynamic{FT}(T, nbranches, domain)

function Mod1Dynamic{FT}(T, nbranches = undef, domain = Interval{Float64}(0,1)) where {FT}
	@assert domain == 0..1 # TODO: this only works for domain == 0..1, for now
	range_diff = T(@interval(1.))-T(@interval(0.))
	orientation = unique_sign(range_diff)
	nbranches = ceil(orientation * range_diff).hi
	is_full_branch = isinteger(T(0..0)) & isinteger(T(1..1))
	return Mod1Dynamic{FT}(T, nbranches, orientation, domain, is_full_branch)
end

DynamicDefinition.domain(S::Mod1Dynamic{FT}) where {FT} = S.domain

DynamicDefinition.nbranches(S::Mod1Dynamic{FT}) where {FT} =S.nbranches

DynamicDefinition.is_full_branch(S::Mod1Dynamic{FT}) where {FT} = S.is_full_branch

# TODO: serious doubts that this works if T(0) is not an integer...
function DynamicDefinition.preim(D::Mod1Dynamic{FT}, k, y, ϵ) where {FT}
	# we need to treat the case with the other orientation, 0 not fixed point...
	@assert 1 <= k <= D.nbranches
	f(x) = D.T(x)-D.T(0)-(y-D.T(0)+(k-1)*D.orientation)
	root(f, D.domain, ϵ)
end

DynamicDefinition.derivative(n, D::Mod1Dynamic{FT}, x) where {FT} = derivative(n, D.T, x)
DynamicDefinition.distorsion(D::Mod1Dynamic{FT}, x) where {FT} = distorsion(D.T, x)

DynamicDefinition.max_distorsion(D::Mod1Dynamic{FT}, tol=1e-3) where {FT} = maximise(x -> distorsion(D.T, x), domain(D), tol=tol)[1]
DynamicDefinition.expansivity(D::Mod1Dynamic{FT}, tol=1e-3) where {FT} = maximise(x -> abs(1/derivative(D, x)), domain(D), tol=tol)[1]

function DynamicDefinition.plottable(D::Mod1Dynamic{FT}, x) where {FT}
	@assert 0 <= x <= 1
	return mod(D.T(x), 1.)
end

using RecipesBase
@recipe f(::Type{Mod1Dynamic{FT}}, D::Mod1Dynamic{FT}) where {FT} = x -> plottable(D, x)

orientation(D::Mod1Dynamic, k) = D.orientation

end


@doc doc"""
`Q` state clock model with energy $E = -\sum_{ij} J_{ij} \cos(\theta_i - \theta_j)$,
where $\theta_i = 2\pi \sigma_i/Q$ and $\sigma_i$ takes an integer value from $1$ to $Q$.
"""
mutable struct Clock <: Model
    lat :: Lattice
    Q :: Int
    spins :: Matrix{Int}
    cosines :: Vector{Float64}
    sines :: Vector{Float64}
    sines_sw :: Vector{Float64}
    rng :: Random.MersenneTwister

    function Clock(lat::Lattice, Q::Integer, rng::Random.AbstractRNG)
        if Q < 2
            error("Q should be 2 or more: Q = $Q")
        end
        spins = rand(rng, 1:Q, 1, numsites(lat))
        cosines = [cospi(2s/Q) for s in 1:Q]
        sines = [sinpi(2s/Q) for s in 1:Q]
        sines_sw = [sinpi(2(s-0.5)/Q) for s in 1:Q]
        return new(lat, Q, spins, cosines, sines, sines_sw, rng)
    end
end
Clock(lat::Lattice, Q::Integer) = Clock(lat, Q, Random.seed!(Random.MersenneTwister(0)))
Clock(lat::Lattice, Q::Integer, seed) = Clock(lat, Q, Random.seed!(Random.MersenneTwister(0), seed...))

@doc doc"""
    Clock(param)

Generates `Clock` using `param["Lattice"]`, `param["Q"]`,  and `param["Seed"]` (if defined).
Each spin $\sigma_i$ will be initialized randomly and independently.
"""
function Clock(param::Parameter)
    lat = generatelattice(param)
    Q = param["Q"]
    if "Seed" in keys(param)
        return Clock(lat, Q, param["Seed"])
    else
        return Clock(lat, Q)
    end
end

include("update.jl")
include("estimator.jl")
include("postproc.jl")


using Statistics
using Printf

function compute_simd(x::Array{Float64, 1}, y::Array{Float64, 1},
                 grads::Array{Float64, 3}, out::Array{Float64, 1})

    det = -(x[2] - x[3])*y[1] + (x[1] - x[3])*y[2] - (x[1] - x[2])*y[3]
    dv = 1.0 / (6.0 * det)
    ii = 1
    for i=1:6
        for j=1:i
            s = 0.0
            @simd for k=1:3
                s += grads[1, k, i] * grads[1, k, j] + grads[2, k, i] * grads[2, k, j]
            end
            out[ii] = dv * s
            ii += 1
        end
    end

    return out
end

function compute(x::Array{Float64, 1}, y::Array{Float64, 1},
                 grads::Array{Float64, 3}, out::Array{Float64, 1})

    det = -(x[2] - x[3])*y[1] + (x[1] - x[3])*y[2] - (x[1] - x[2])*y[3]
    dv = 1.0 / (6.0 * det)
    ii = 1
    for i=1:6
        for j=1:i
            s = 0.0
            for k=1:3
                s += grads[1, k, i] * grads[1, k, j] + grads[2, k, i] * grads[2, k, j]
            end
            out[ii] = dv * s
            ii += 1
        end
    end

    return out
end

const x = Float64[0., 1., 0.]
const y = Float64[0., 0., 1.]

a11 = -y[1] + y[3]
a12 =  y[1] - y[2]
a21 =  x[1] - x[3]
a22 = -x[1] + x[2]

gq = reshape(
    [-1.0,-1.0,1.0,1.0,-1.0,-1.0,1.0,0.0,1.0,0.0,-1.0,0.0,0.0,
     -1.0,0.0,1.0,0.0,1.0,0.0,-2.0,-2.0,-2.0,2.0,0.0,0.0,2.0,2.0,
     2.0,2.0,0.0,0.0,2.0,-2.0,-2.0,-2.0,0], 2, 3, 6)

grads = zeros(2, 3, 6)

for f=1:6
    for p=1:3
        grads[1, p, f] = a11 * gq[1, p, f] + a12 * gq[2, p, f]
        grads[2, p, f] = a21 * gq[1, p, f] + a22 * gq[2, p, f]
    end
end

result = zeros(21)

nb_runs = 1000

times = zeros(nb_runs)

for irun in 1:nb_runs
    times[irun] = @elapsed compute(x, y, grads, result)
end

time_without_simd = median(times)

@printf("Julia without simd %.3f µs\n", time_without_simd * 1e6)

for irun in 1:nb_runs
    times[irun] = @elapsed compute_simd(x, y, grads, result)
end

time_with_simd = median(times)

@printf("Julia with simd    %.3f µs\n", time_with_simd * 1e6)

@printf("ratio (without simd)/(with simd): %.2f\n", time_without_simd / time_with_simd)


using GLMakie
function flower(n; npetals = 8)
    n = div(n, npetals)
    x = mapreduce(hcat, (1:npetals) .* (2π/npetals)) do θ
        ct = cos(θ)
        st = sin(θ)

        x0 = tanh.(randn(1, n) .- 1) .+ 4.0 .+ 0.05.* randn(1, n)
        y0 = randn(1, n) .* 0.3

        x₁ = x0 * cos(θ) .- y0 * sin(θ)
        x₂ = x0 * sin(θ) .+ y0 * cos(θ)
        vcat(x₁, x₂)
    end
    _y = mapreduce(i -> fill(i, n), vcat, 1:npetals)
    y = zeros(npetals, length(_y))
    foreach(i -> y[_y[i], i] = 1, 1:length(_y))
    Float32.(x), Float32.(y)
end

x, y = flower(900)
scatter(x[1,:], x[2,:], color = mapslices(argmax, y, dims = 1)[:])

#######
#   Define a Tracked Array for operator overloading AD
#######
struct TrackedArray{T,N,V<:AbstractArray{T,N}} <: AbstractArray{T,N}
    value::V
    grad::Union{Nothing,V}
    tape::Vector{Any}
end

TrackedArray(a::AbstractArray) = TrackedArray(a, similar(a) .= 0, [])
TrackedMatrix{T,V} = TrackedArray{T,2,V} where {T,V<:AbstractMatrix{T}}
TrackedVector{T,V} = TrackedArray{T,1,V} where {T,V<:AbstractVector{T}}
Base.size(a::TrackedArray) = size(a.value)
Base.show(io::IO, ::MIME"text/plain", a::TrackedArray) = show(io, a)
Base.show(io::IO, a::TrackedArray) = print(io, "TrackedArray($(size(a.value)))")
value(A::TrackedArray) = A.value
resetgrad!(A::TrackedArray) = (A.grad .= 0; empty!(A.tape))
value(A) = A
track(A) = TrackedArray(A)
track(a::Number) = TrackedArray(reshape([a], 1, 1))

function accum!(A::TrackedArray)
    isempty(A.tape) && return(A.grad)
    A.grad .= sum(g(accum!(r)) for (r, g) in A.tape)
    empty!(A.tape)
    A.grad
end

#######
#   Define AD rules for few operations appearing in FFNN
#######
import Base: +, *
import Base.Broadcast: broadcasted
function *(A::TrackedMatrix, B::TrackedMatrix)
    a, b = value.((A, B))
    C = track(a * b)
    push!(A.tape, (C, Δ -> Δ * b'))
    push!(B.tape, (C, Δ -> a' * Δ))
    C
end

function *(A::TrackedMatrix, B::AbstractMatrix)
    a, b = value.((A, B))
    C = track(a * b)
    push!(A.tape, (C, Δ -> Δ * b'))
    C
end

function broadcasted(::typeof(+), A::TrackedMatrix, B::TrackedVector)
    C = track(value(A) .+ value(B))
    push!(A.tape, (C, Δ -> Δ))
    push!(B.tape, (C, Δ -> sum(Δ, dims = 2)[:]))
    C
end

function σ(x::Real)
    t = @fastmath exp(-abs(x))
    y = ifelse(x ≥ 0, inv(1 + t), t / (1 + t))
    ifelse(x > 40, one(y), ifelse(x < -80, zero(y), y))
end

broadcasted(::typeof(identity), A::TrackedArray) = A

function broadcasted(::typeof(σ), A::TrackedArray)
    Ω = σ.(value(A))
    C = track(Ω)
    push!(A.tape, (C, Δ -> Δ .* Ω .* (1 .- Ω)))
    C
end

function mse(A::TrackedMatrix, B::AbstractMatrix)
    n = size(A, 2)
    a = value(A)
    c = similar(a, 1, 1)
    c .= sum((a .- B).^2)/2n
    C = track(c)
    push!(A.tape, (C, Δ -> Δ .* (a .- B) ./ n))
    C
end

mse(x::AbstractMatrix, y::AbstractMatrix) = sum((x - y).^2) / (2*size(x,2))

#######
#   Define a Dense layer
#######
struct Dense{F,W,B}
    σ::F
    w::W
    b::B
end

Base.show(io::IO, m::Dense) = print(io, "Dense($(size(m.w,2)) → $(size(m.w,1)))")
Dense(i::Int, o::Int, σ = identity) = Dense(σ, randn(Float32, o, i), randn(Float32, o))
track(m::Dense) = Dense(m.σ, track(m.w), track(m.b))
track(m::ComposedFunction) = track(m.outer) ∘ track(m.inner)
(m::Dense)(x) = m.σ.(m.w * x .+ m.b)
params(m::ComposedFunction) = vcat(params(m.outer), params(m.inner))
params(m::Dense) = [m.w, m.b]

#######
#   Let's try to actually train a model
#######
x, y = flower(900)
function initmodel()
    m₁ = track(Dense(2, 20, σ))
    m₂ = track(Dense(20, 20, σ))
    m₃ = track(Dense(20, size(y,1)))
    m = m₃ ∘ m₂ ∘ m₁
end
m = initmodel()
m(x) |> value

######
#   Let's try to learn the parameters
######
α = 0.01
ps = params(m)
@elapsed for i in 1:10000
    foreach(resetgrad!, ps)
    loss = mse(m(x), y)
    fill!(loss.grad, 1)
    foreach(accum!, ps)
    foreach(x -> x.value .-= α .* x.grad, ps)
    mod(i,250) == 0 && println("loss after $(i) iterations = ", sum(value(loss)))
end

all(mapslices(argmax, value(m(x)), dims = 1)[:] .== mapslices(argmax, y, dims = 1)[:])
scatter(x[1,:], x[2,:], color = mapslices(argmax, value(m(x)), dims = 1)[:])

######
#   Let's try to move the computation to GPU
######
using CUDA
gpu(x::AbstractArray) = CuArray(x)
gpu(x::TrackedArray) = TrackedArray(CuArray(value(x)))
gpu(m::Dense) = Dense(m.σ, gpu(m.w), gpu(m.b))
gpu(m::ComposedFunction) = gpu(m.outer) ∘ gpu(m.inner)

gx, gy = gpu(x), gpu(y)
m = gpu(m)
ps = params(m)
@elapsed for i in 1:10000
    foreach(resetgrad!, ps)
    loss = mse(m(gx), gy)
    fill!(loss.grad, 1)
    foreach(accum!, ps)
    foreach(x -> x.value .-= α .* x.grad, ps)
    mod(i,250) == 0 && println("loss after $(i) iterations = ", sum(value(loss)))
end

#######
#   Why we see a small speed-up? The problem is small
#######
using BenchmarkTools
p = randn(Float32, 20, 2)
@benchmark $(p) * $(x)
gp = gpu(p)
@benchmark $(gp) * $(gx)

######
#   Let's verify the gradients
######
using FiniteDifferences
ps = [m₃.w, m₃.b, m₂.w, m₂.b, m₁.w, m₁.b]
map(ps) do p
    foreach(resetgrad!, ps)
    loss = mse(m(x), y)
    fill!(loss.grad, 1)
    foreach(accum!, ps)
    accum!(p)
    θ = deepcopy(value(p))
    Δθ = deepcopy(p.grad)
    f = θ -> begin
        p.value .= θ
        value(mse(m(x), y))
    end
    sum(abs2.(grad(central_fdm(5, 1), f, θ)[1] - Δθ))
end


using Plots
gr()
x = LinRange(0,2π,100)
plot(x, sin.(x), c = :black, legend =:bottomleft, label ="sin(x)",
    fg_legend = :red, bg_legend = nothing,xlab ="x",ylab ="f(x)",
    title = "Sine", guidefont=(18,"Times"),tickfont=(14, "Times"),
    legendfont=(13, "Times"), titlefont=(18, "Times"))
savefig("figSin1.pdf")


##############################################################################
#
# Symmetric triangular distribution from Distributions Handbook
#
##############################################################################

immutable Triangular <: ContinuousUnivariateDistribution
    location::Float64
    scale::Float64
    function Triangular(l::Real, s::Real)
        if s > 0.0
            new(float64(l), float64(s))
        else
            error("scale must be positive")
        end
    end
end

Triangular(location::Real) = Triangular(location, 1.0)
Triangular() = Triangular(0.0, 1.0)

function insupport(d::Triangular, x::Number)
    return isreal(x) && isfinite(x) &&
           d.location - d.scale <= x <= d.location + d.scale
end

kurtosis(d::Triangular) = d.scale^4 / 15.0

mean(d::Triangular) = d.location

median(d::Triangular) = d.location

modes(d::Triangular) = [d.location]

function pdf(d::Triangular, x::Real)
    if insupport(d, x)
        return -abs(x - d.location) / (d.scale^2) + 1.0 / d.scale
    else
        return 0.0
    end
end

function rand(d::Triangular)
    xi1, xi2 = rand(), rand()
    return d.location + (xi1 - xi2) * d.scale
end

skewness(d::Triangular) = 0.0

var(d::Triangular) = d.scale^2 / 6.0


using LightGraphs
using LightGraphs.SimpleGraphs

include("SI\\metaplex.jl")

N = 100
M = 10

g = random_configuration_model(N*M, fill(6,N*M))
g = complete_graph(N*M)
h = random_configuration_model(M,fill(3,M))
h = erdos_renyi(M, 0.5)
h = path_graph(M)
h = complete_graph(M)

using GraphPlot
gplot(h)

mpx = Metaplex(g,h)

Xi0 = falses(N*M)
Xi0[1:10] .= true

Xμ0 = [div(i-1,N)+1 for i in 1:N*M]
Xμ0 = [rand(1:M) for i in 1:N*M]

β = 0.2
D = [0.1, 0.1]
tmax = 3.0

@time t, pcs = metaplex_gillespie(mpx, Xi0, Xμ0, β, D, nmax = 500000, tmax=tmax)

using Plots
gr()
pyplot()

plot(t, pcs[:,2,:], label="", line=:steppre)

plot(t, sum(pcs[:,2,:], dims=2)/(N*M),
    label="",
    line=:steppre,
    xlabel="Time",
    ylabel="Fraction of infected individuals"
    )

plot(t, vec(sum(pcs[:,:,:], dims=[2,3])), label="")

t_mf, s_mf, i_mf, sμ_mf, iμ_mf, sol = metaplex_ode(mpx, Xi0, Xμ0, β, D, tmax = tmax)

plot(t_mf, reshape(sum(i_mf, dims=3), length(sol.t), N*M), label="")

plot(t_mf, iμ_mf, label="")

sum(iμ_mf, dims=2)/(N*M)

plot(t, sum(pcs[:,2,:], dims=2)/(N*M),
    label="Gillespie",
    line=:steppre,
    xlabel="Time",
    ylabel="Fraction of infected individuals")
plot!(t_mf, sum(iμ_mf, dims=2)/(N*M), label="Mean Field")

t_mc, s_mc, i_mc, sd_s_mc, sd_i_mc = metaplex_montecarlo(mpx, Xi0, Xμ0, β, D, tmax=tmax, nmax=500000)

plot(t_mc, i_mc, label="")
plot!(t_mf, iμ_mf, label="", linestyle=:dash)

plot(t_mc, sum(i_mc, dims=2)/(N*M), label="Average", legend=:bottomright)
plot!(t_mf, sum(iμ_mf, dims=2)/(N*M), label="Mean Field")

sum(i_mf, dims=2)/(N*M)
plot(t_mc[1:20], sum(i_mc[1:20,:], dims=2)/(N*M), label="Average", legend=:topleft)
plot!(t_mf[1:50], sum(iμ_mf, dims=2)[1:50]/(N*M), label="Mean Field")


@testset "General Manifold functions and fallbacks" begin
M = Euclidean(3)
x = RnPoint([1.,0.,0.])
y = RnPoint([0.,1.,0.])
ξ = RnTVector([1.,0.,1.])
η = RnTVector([0.,1.,1.])
@test copy(x) == x
@test copy(ξ) == ξ
@test getValue(2*ξ) == 2*getValue(ξ)
@test getValue(ξ*2) == getValue(ξ)*2
@test getValue(2/RnTVector([1.,1.])) == 2 ./ [1.,1.]
@test getValue(ξ/2) == getValue(ξ)/2
@test getValue(ξ+η) == getValue(ξ) + getValue(η)
@test getValue(ξ-η) == getValue(ξ) - getValue(η)
@test getValue(+η) == getValue(η)
@test getValue(-η) == -getValue(η)
@test x==x
@test ξ==ξ
@test geodesic(M,x,y,3) == RnPoint.([[1.,0.,0.],[.5,.5,0.],[0.,1.,0.]])
@test geodesic(M,x,y,[0.,.5,1.]) == RnPoint.([[1.,0.,0.],[.5,.5,0.],[0.,1.,0.]])
@test reflection(M,x,y) == RnPoint(2*getValue(x)-getValue(y))
#
# Test Fallbacks – check that errors are thrown if the manifold does not fit
struct TestManifold <: Manifold end
struct TestMPoint <: MPoint end
struct TestTVector <: TVector end
M = TestManifold()
x = TestMPoint()
y = TestMPoint()
ξ = TestTVector()
η = TestTVector()
@test_throws DomainError distance(M, x, y)
@test_throws DomainError dot(M, x, ξ, η)
@test_throws ErrorException dot(Euclidean(3), ξ, η)
@test_throws DomainError exp(M, x, ξ)
@test_throws DomainError log(M, x, y)
@test_throws DomainError parallelTransport(M, x, y, ξ)
@test_throws DomainError tangentONB(M, x, ξ)
@test_throws DomainError tangentONB(M, x, y)
@test_logs (:warn,"""No valitadion for a TestMPoint on TestManifold available. Continuing without
  validation. To turn this warning off, either deactivate the validate flag
  in (one of) your extended MPoints or implement a corresponding validation.""") validateMPoint(M,x)
@test_logs (:warn,"""No valitadion for a TestMPoint and a TestTVector on TestManifold available.
  Continuing without validation. To turn this warning off, either deactivate
  the validate flag in (one of) your extended TVectors or MPoints or
  implement a corresponding validation""") validateTVector(M,x,ξ)
@test_throws DomainError zeroTVector(M,x)
#
#
@test_throws DomainError manifoldDimension(M)
@test_throws DomainError manifoldDimension(x)
@test_throws DomainError getValue(x)
@test_throws DomainError getValue(ξ)
#
@test_throws ErrorException typeofMPoint(ξ)
@test_throws ErrorException typeofTVector(x)
@test_throws ErrorException typicalDistance(M)
end


### A Pluto.jl notebook ###
# v0.12.11

using Markdown
using InteractiveUtils

# ╔═╡ f3a59724-187d-11eb-2f7c-2df2f3339a6e
using Pkg, DrWatson

# ╔═╡ f3a5da02-187d-11eb-24ed-a7c460c0219a
begin
	@quickactivate "StatisticalRethinkingTuring"
	using Turing
	using StatisticalRethinking
end

# ╔═╡ f3b3b5aa-187c-11eb-2e75-07bc5598cd16
md"## Turing intro clip 3_quap.jl"

# ╔═╡ f3a66514-187d-11eb-22eb-47f74d2724cd
begin
	df = CSV.read(sr_datadir("Howell1.csv"), DataFrame)
	df = df[df.age .>= 18, :]
	x̄ = mean(df.weight)
	x = range(minimum(df.weight), maximum(df.weight), length = 100)     # either
end;

# ╔═╡ 2f0ce4de-187e-11eb-296e-8b4e16814350
@model function ppl4_3(weights, heights)
    a ~ Normal(178, 20)
    b ~ LogNormal(0, 1)
    σ ~ Uniform(0, 50)
    μ = a .+ b .* weights
    for i in eachindex(heights)
        heights[i] ~ Normal(μ[i], σ)
    end
end

# ╔═╡ 33b27276-187e-11eb-26a8-1bb7ad288539
begin
	m4_3t = ppl4_3(df.weight, df.height)
    sampler=NUTS(0.65)
	nsamples=2000
	nchains=4
    chns4_3t = mapreduce(c -> sample(m4_3t, sampler, nsamples), chainscat, 1:nchains)
	Particles(Chains(chns4_3t, [:parameters]))
end

# ╔═╡ 7779262e-187e-11eb-0441-37f564373c1e
q4_3t = quap(m4_3t)

# ╔═╡ 32db0b98-2c08-11eb-145d-a1ea44641e75
rand(q4_3t.distr, 4_000)'

# ╔═╡ 3efadbfe-1882-11eb-0900-b17b00bf8f31

begin
	scatter(df.weight, df.height, leg=false)
	quap4_3t = sample(q4_3t)
	a_map = mean(quap4_3t.a)
	b_map = mean(quap4_3t.b)
	plot!(x, a_map .+ b_map .* x)
end

# ╔═╡ Cell order:
# ╟─f3b3b5aa-187c-11eb-2e75-07bc5598cd16
# ╠═f3a59724-187d-11eb-2f7c-2df2f3339a6e
# ╠═f3a5da02-187d-11eb-24ed-a7c460c0219a
# ╠═f3a66514-187d-11eb-22eb-47f74d2724cd
# ╠═2f0ce4de-187e-11eb-296e-8b4e16814350
# ╠═33b27276-187e-11eb-26a8-1bb7ad288539
# ╠═7779262e-187e-11eb-0441-37f564373c1e
# ╠═32db0b98-2c08-11eb-145d-a1ea44641e75
# ╠═3efadbfe-1882-11eb-0900-b17b00bf8f31


commands = readlines("day14/input.txt")

function decode_addresses(mask, value)
    value_bits = bitstring(value)
    out_bits = []
    for i in 0:35
        if mask[end-i] == '0'
            pushfirst!(out_bits, value_bits[end-i])
        else
            pushfirst!(out_bits, mask[end-i])
        end
    end
    out_bitstring = join(out_bits)
    addresses_out = [out_bitstring]
    while contains(addresses_out[1], 'X')
        address = popfirst!(addresses_out)
        idx = findfirst('X', address)
        push!(addresses_out, address[1:idx-1]*"0"*address[idx+1:end], address[1:idx-1]*"1"*address[idx+1:end])
    end
    return addresses_out
end

current_mask = ""
memory = Dict{Int, Int}()
for instruction in commands
    if contains(instruction, "[")
        address, value = match(r"mem\[(\d+)\] = (\d+)", instruction).captures
        for address in decode_addresses(current_mask, parse(Int, address))
            global memory[parse(Int, address; base=2)] = parse(Int, value)
        end
    else
        global current_mask = match(r"mask = (\w+)", instruction).captures[1]
    end
end

println(sum(v for v in values(memory)))


#=
These examples use ArbNumerics v0.4.4, no earlier version supports `exp(matrix)`.
So, either first do `pkg> up` or do `pkg> rm ArbNumerics` then `pkg> add ArbNumerics`.
Sometimes, rebuilding with `pkg> build ArbNumerics` cleans up after itself.
All the matrix handling is new code, the examples here are working as of now.
If you find anything unclear or something awry, post an issue and
please use slack's direct message to draw my attention.
=#

#=
Using Float64 matrices with entries from `rand` in 0.0..1.0:
with a   250x250 matrix of Float64s, working with 350 bits used  10 seconds, trace ~ 1.0e+53
with a   500x500 matrix of Float64s, working with 350 bits used  50 seconds, trace ~ 1.0e+108
with a   750x750 matrix of Float64s, working with 350 bits used 140 seconds, trace ~ 1.0e+162
with a 1000x1000 matrix of Float64s, working with 350 bits used 380 seconds, trace ~ 1.0e+215
with a   250x250 matrix of Float64s, working with 700 bits used  25 seconds
with a   500x500 matrix of Float64s, working with 350 bits used  25 seconds
with a   750x750 matrix of Float64s, working with 175 bits used  25 seconds
=#

using ArbNumerics # GenericLinearAlgebra, LinearAlgebra, Readables

#=
ArbNumerics' interactive users see the "viewable precision"; computations utilize "working precision".
By definition, _viewable precision <= working_precision_. They may share the same number; then the
results shown are precisely the computed numerical results. Should they differ, there are "extra bits".
Julia should be restarted just before these precisions are altered.  When the precision is being altered
in the course of iterative refinement, it is the precision of successive values which moves, not their
bedrock relationship.  Individual values may be constructed more or less precisely than their sources.
Arithmetic over values of differing precision yields a result that adopts the lesser precision.
When the precision is lower than that which permits crisp numerical results, the interval known
to contain the actual result eclipses the magnitude of that value.  It appears so: `[+/- 12.1]`.
=#

setextrabits(0)  #  obtain the raw results
setextrabits(24) #  more accurately, less precision is observed
#  This is not a means of protection from numerical misadventure.
#  This may be adaptive accomodation to the partially [un]known.

                                    # Use one of these perspectives on bits that count.
setworkingprecision(ArbFloat, 125)  # This focus is on available computational precision.
setprecision(ArbReal, 100)          # Or prefer less precision with more accuracy.

                                    # the most recent `set_` gives the operative precision
                                    # for constructing new Arb values using type name only
                                    # the setting is shared by all of the Arb types
                                    # and that makes it easier to characterize accuracy

#=
#   During the course of computation, ArbFloat, ArbReal, ArbComplex values
#   may be constructed at various precisions.  Once given, precision stays.
#   Change of precision occurs through construction of that changed value.
=#

a₂  = ArbFloat(1, digits=100, base=2);
a₁₀ = ArbFloat(1, digits=29, base=10);
workingprecision.((a₂, a₁₀))
(100, 98)

b₂  = ArbFloat(1, digits=3000, base=2);
b₁₀ = ArbFloat(1, digits=902, base=10);
workingprecision.((b₂, b₁₀))
(3000, 2999)

workingprecision(ArbReal(1)) == 128
workingprecision(ArbFloat(1, bits=1024) + ArbFloat(1, bits=256)) == 256

#>  ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ...  <#

#=
    To regard a matrix of extended precision ArbFloats or ArbReals,
    where the elements are in the Float64 domain: `Float64.(matrix)`.
    To characterize or classify a large matrix of the same nature:
    `approx = Float64.(matrix); result = describe(process(approx))`
#
const Arb{P} = Union{ArbFloat{P}, ArbReal{P}} where {P}
#=
#   Obtaining the exponential of a matrix
#
#     (a) provision a square matrix
#     (b) generate a similar matrix of eltype ArbReal
#     (c) receive each of (b) carried from (a)
#     (d) request exp(c) of ArbNumerics                (faster than exp(Matrix{ArbFloat}))
#     (e) generate a similar matrix given (a)
#     (f) receive each of (e) carried from (d)
=#
function expmatrix(m::Array{T,2}) where {T<:Real}
    a = Array{ArbReal,2}(undef, size(m))
    a[:] = m[:]
    z = exp(a)
    result = similar(m)
    result[:] = z[:]
    return result
end


n1 = 1. # bimolecular 1
n2 = 2. # bimolecular 2
n3 = 3. # trimolecular
function cooperativity_index(n1,n2,n3)
    Σ_bi = n1 + n2
    Σ_tri = n3
    Δn = Σ_tri - Σ_bi
    # cooperativity index formulav (relative change)
    Coop_indx = Δn / Σ_bi
end

cooperativity_index(n1,n2,n3)

n1 = [1. ± 0.1, 1. ± 0.1]  # bimolecular 1
n2 = [2. ± 0.1, 2. ± 0.1]# bimolecular 2
n3 = [3. ± 0.1, 4. ± 0.1]

cooperativity_index.(n1,n2,n3)

function solve_with_sds(rates_sds)
    prob = ODEProblem(f,u₀,tspan,10 .^ rates_sds)
    # prob = remake(prob, p=10 .^ rates)
    sol = solve(prob,Vern7(),saveat=0.01, abs_tol=1e-10,rel_tol=1e-10)
    n_tot = reduce(vcat,sum(sol,dims=1))
    return n_tot
end
function solve_model(rates_sds)
    prob = ODEProblem(f,u₀,tspan,10 .^ rates_sds)
    # prob = remake(prob, p=10 .^ rates)
    sol = solve(prob,Vern7(),saveat=0.01, abs_tol=1e-10,rel_tol=1e-10)
    return sol
end

#-------------------------------------------------------------------------------

densities_tcr = [25.0, 38.0, 1.0] # TCR, pMHC, Ac
densities_tc1 = [15.0, 38.0, 10.0, 1.0] # TCR, pMHC, CD4, Ac
densities_tc2 = [8.0,  38.0, 17.0, 1.0]
densities_tc3 = [3.0,  38.0, 22.0, 1.0]

# rates_T = [-3.37638358011457 ± 0.0483562400944489
#             -0.302013187183543 ± 0.0607189115063286]
# rates_TC1 = [-1.72689171836867 ± 0.469108114325548
#             -0.646695212785203 ± 0.685834531823618]
# rates_TC2 = [-0.644293816490510 ± 0.408735939559103
#             0.617796010276307 ± 0.467542299161307]
# rates_TC3 = [0.891017383250648 ± 0.559684851912230
#             2.16934160248305 ± 0.568940078989975]

rates_T = [-3.37638358011457
            -0.302013187183543]
rates_TC1 = [-1.72689171836867
            -0.646695212785203]
rates_TC2 = [-0.644293816490510
            0.617796010276307]
rates_TC3 = [0.891017383250648
            2.16934160248305]

sd_data = CSV.read("C:/Users/stravaglino3/Downloads/julia_code/n_stds_for_julia.csv")
nT_sds = sd_data[:,1]
nTC1_sds = Array(sd_data[:,2:3])
nTC2_sds = Array(sd_data[:,4:5])
nTC3_sds = Array(sd_data[:,6:7])

tspan = (0.0,16.0)
t = collect(0.01:0.01:16)

gr()
# bimolecular TCR-pMHC bonds
u₀ = [0.0]
f = (du,u,p,t) -> bimolecular!(du,u,p,t,densities_tcr,[])
n_tcr = Array(solve_model(rates_T))'[2:end]
n_tcr_sd = n_tcr .± nT_sds

# TCR-pMHC-CD4 (1)
u₀ = [0.0, 0.0]
f = (du,u,p,t) -> trimolecular!(du,u,p,t,densities_tc1,10 .^ rates_tcr)
n_tc1 = Array(solve_model(rates_TC1))'[2:end,:]
n_tc1_sd = sum(n_tc1 .± nTC1_sds, dims=2)

# TCR-pMHC-CD4 (2)
f = (du,u,p,t) -> trimolecular!(du,u,p,t,densities_tc2,10 .^ rates_tcr)
n_tc2 = Array(solve_model(rates_TC2))'[2:end,:]
n_tc2_sd = sum(n_tc2 .± nTC2_sds, dims=2)
plot!(t,n_tc2_sd, errorbars=:ribbon, alpha=.1,c=:red)

# TCR-pMHC-CD4 (3)
f = (du,u,p,t) -> trimolecular!(du,u,p,t,densities_tc3,10 .^ rates_tcr)
n_tc3 = Array(solve_model(rates_TC3))'[2:end,:]
n_tc3_sd = sum(n_tc3 .± nTC3_sds, dims=2)
plot!(t,n_tc3_sd, errorbars=:ribbon, alpha=0.1,ribboncolor=:cyan)

plot(t, n_tcr_sd, errorbars=:ribbon, alpha=0.1,msc=:black)
plot!(t, n_tcr, lw=3, lc=:black)
plot!(t,n_tc1_sd, errorbars=:ribbon, alpha=0.1, msc=:blue)
plot!(t, sum(n_tc1,dims=2), lw=3, lc=:blue)
plot!(t,n_tc2_sd, errorbars=:ribbon, alpha=0.1, msc=:red)
plot!(t, sum(n_tc2,dims=2), lw=3, lc=:red)
plot!(t,n_tc3_sd, errorbars=:ribbon, alpha=0.1, msc=:cyan)
plot!(t, sum(n_tc3,dims=2), lw=3, lc=:cyan)

scatter!(n_se_tcr[:,1],n_se_tcr[:,2],yerror=n_se_tcr[:,3],markercolor=:black,markersize=5)
scatter!(n_se_1[:,1],n_se_1[:,2],yerror=n_se_1[:,3],markercolor=:blue,markersize=5)
scatter!(n_se_2[:,1],n_se_2[:,2],yerror=n_se_2[:,3],markercolor=:red,markersize=5)
h = scatter!(n_se_3[:,1],n_se_3[:,2],yerror=n_se_3[:,3],markercolor=:cyan,markersize=5,legend=false)

xlabel!("Contact Time tc")
ylabel!("<n>")
# savefig("all_n_with_error_bands")

#-------------------------------------------------------------------------------
# FIT CD4 only data

cd4_1 = CSV.read("cd4_data_from_muaz.csv")
t_cd41, Pa_cd41, n_se_cd41 = clean_data(cd4_1)
n_cd41 = Pa2n(Pa_cd41)

f = (du,u,p,t) -> bimolecular!(du,u,p,t,dens,cons) # enclose constants

# u0 = [0.0;0.0]
u₀ = [0.0]          # Initial condition
tspan = (0.0,8.0)  # Simulation interval
pinit = [1e-4,0.5]      # equation parameter. p = [k₁, k₋₁]
densities_cd41 = [1700., 850., 1.] #CD4, pMHC, A
cons = []

rates_cd41,hes_cd41,l_cd41,boot_cd41 = optimization(bimolecular!,n_cd41,densities_cd41',[],model_opt2,t_cd41)

# Setup the ODE problem, then solve
plotly()
t = collect(0.01:0.01:16)
tspan = (0.0,16.0)
f = (du,u,p,t) -> bimolecular!(du,u,p,t,densities_cd41,cons)
prob = ODEProblem(f,u₀,tspan,10 .^ rates_cd41)
sol_cd41 = Array(solve(prob,Vern7(),saveat=t,abstol=1e-10,reltol=1e-10))'
plot(t,sol_cd41,lc=:orange,lw=3)
scatter!(n_se_cd41[:,1],n_se_cd41[:,2],yerror=n_se_cd41[:,3],markercolor=:orange)

densities_tcr = [25.0, 38.0, 1.0] # TCR, pMHC, Ac
densities_tc1 = [15.0, 38.0, 10.0, 1.0] # TCR, pMHC, CD4, Ac
densities_tc2 = [8.0,  38.0, 17.0, 1.0]
densities_tc3 = [3.0,  38.0, 22.0, 1.0]
densities_cd41 = [1700., 850., 1.] #CD4, pMHC, A
cooperativity_index(n1,n2,n3)
#---------------------
#normalize everything to tc2

n_tcr_norm = n_tcr./ 25 .* 8
n_tc1_norm = sum(n_tc1,dims=2) ./15 .* 8 ./10 .*17
n_tc2_norm = sum(n_tc2,dims=2)
n_tc3_norm = sum(n_tc3,dims=2) ./3 .* 8 ./22 .* 17
n_cd41_norm = sol_cd41 ./1700 .*17 ./850 .*38
n_tcr_norm = n_tcr ./25. .*15.
n_tc1_norm = sum(n_tc1,dims=2)
n_tc2_norm = sum(n_tc2,dims=2) ./8. .*15. ./17. .*10.
n_tc3_norm = sum(n_tc3,dims=2) ./3. .*15. ./22. .*10.
n_cd41_norm = sol_cd41 ./1700. .*10. ./850. .*38

plotly()
t = collect(0.01:0.01:16)
plot(t,n_tcr_norm,lw=3,lc=:black)
plot!(t,n_tc1_norm,lw=3,lc=:blue)
plot!(t,n_tc2_norm,lw=3,lc=:red)
plot!(t,n_tc3_norm,lw=3,lc=:cyan)
plot!(t,n_cd41_norm,lw=3,lc=:orange,yscale=:linear,legend=false)
yaxis!("log<n>")
xlabel!("Contact Time tc")
savefig("n_log_scaled")

coop_indx_tc1 = cooperativity_index.(n_tcr_norm,n_cd41_norm,n_tc1_norm)
coop_indx_tc2 = cooperativity_index.(n_tcr_norm,n_cd41_norm,n_tc2_norm)
coop_indx_tc3 = cooperativity_index.(n_tcr_norm,n_cd41_norm,n_tc3_norm)
plot(t,coop_indx_tc1 .* 100,lw=3,lc=:blue)
plot!(t,coop_indx_tc2 .* 100,lw=3,lc=:red)
plot!(t,coop_indx_tc3 .* 100,lw=3,lc=:cyan,legend=false)
xlabel!("Contact Time tc")
ylabel!("Cooperative Index (%)")
# savefig("Cooperative Index.png")


using DataStructures
using Utils

function load()
    re = Int[]
    open(joinpath(@__DIR__, "01_expense-report.txt"), "r") do f
        for i in eachline(f)
            push!(re, parse(Int, i))
        end
    end
    re
end

# Original solution. My first Julia code ever.
# See below for a generalized, count-independent version.
function find_product2(r)
    for (o, a) in enumerate(r)
        for i in o+1:length(r)
            if a + r[i] == 2020
                return a * r[i]
            end
        end
    end
end

# Original solution. My first Julia code ever.
# See below for a generalized, count-independent version.
function find_product3(r)
    for (o, a) in enumerate(r)
        for i in o+1:length(r)
            for j in i+1:length(r)
                if a + r[i] + r[j] == 2020
                    return a * r[i] * r[j]
                end
            end
        end
    end
end

using DataStructures

function find(r, n, start = 1, which = nil())::Union{Int,Nothing}
    if start <= length(r)
        current = r[start]
        extended = cons(current, which)
        with_current = sum(extended)
        if with_current == 2020 && n == 1
            return prod(extended)
        else
            return @something_nothing find(r, n, start + 1, which) find(
                r,
                n - 1,
                start + 1,
                extended,
            )
        end
    end
end

const expenses = load()
println("Product of 2: ", find_product2(expenses))
println("Product of 3: ", find_product3(expenses))

using Test
using OffsetArrays

@testset begin
    a = [11, 222, 1010, 3333, 1010]
    b = [20, 11, 222, 1010, 3333, 2000]
    c = [11, 222, 1010, 3333, 1009, 55555, 1]
    d = [20, 11, 222, 2, 1010, 3333, 1998]

    @testset "2 loops" begin
        @test find_product2(a) == 1010 * 1010
        @test find_product2(b) == 20 * 2000
    end

    @testset "3 loops" begin
        @test find_product3(c) == 1009 * 1010
        @test find_product3(d) == 2 * 20 * 1998
    end

    finds = OffsetArray([find_product2, find_product3], 2:3)
    @testset "recursive equivalence" begin
        for i in 2:3
            @test find(a, i) == finds[i](a)
            @test find(b, i) == finds[i](b)
            @test find(c, i) == finds[i](c)
            @test find(d, i) == finds[i](d)
        end
    end

    @testset "solutions" begin
        @test find_product2(expenses) == 692916
        @test find_product3(expenses) == 289270976
    end
end


using HypothesisTests, Base.Test

@test SWCoeffs(3).A == [sqrt(2.0)/2.0]
@test SWCoeffs(3).N == 3

S = SWCoeffs(10)
@test S.N == 10
@test length(S) == 10
@test endof(S) == 10
@test length(S.A) == 5

@test S.A[4] == S[4] == -S[7]
@test S.A[2] == S[2] == -S[9]

S = SWCoeffs(11)
@test S.N == 11
@test length(S) == 11
@test endof(S) == 11
@test length(S.A) == 5

@test S[6] == 0.0
@test S.A[5] == S[5] == -S[7]
@test S.A[3] == S[3] == -S[9]
@test S.A[1] == S[1] == -S[11]

#anti-symmetry
S = SWCoeffs(20)
@test all([S[i] == -S[end-i+1] for i in 1:length(S)])

# Values obtained by calling `_swilk` fortran subroutine directly.

SWc10 = SWCoeffs(10)
a = SWc10.A .- [0.573715, 0.32897, 0.214349, 0.122791, 0.0400887]
@test norm(a,1) ≈ 0.0 atol=S.N*eps(Float32)

SWc20 = SWCoeffs(20)
b = SWc20.A .- [0.473371, 0.32174, 0.255663, 0.208297, 0.16864, 0.133584, 0.101474, 0.0712893, 0.0423232, 0.0140351]
@test norm(b,1) ≈ 0.0 atol=S.N*eps(Float32)

srand(1)
X = sort(rand(Float64, 10))
@test HypothesisTests.swstat(X, SWc10) ≈ 0.8453033192812377 atol=1.2e-8
@test pvalue(HypothesisTests.swstat(X, SWc10), SWc10) ≈ 0.05106360837817192 atol=4.9e-8

# **Worked Example** (Section 4) from
# PATRICK ROYSTON Approximating the Shapiro-Wilk W-test for non-normality
# *Statistics and Computing* (1992) **2**, 117-119

X = [48.4, 49.0, 59.5, 59.6, 60.7, 88.8, 98.2, 109.4, 169.1, 227.1]
SWc = SWCoeffs(length(X))
@test all(abs.(SWc.A .- [0.5737, 0.3290, 0.2143, 0.1228, 0.0401]) .< 5.0e-5)
W = HypothesisTests.swstat(X, SWc)
@test W ≈ 0.8078 atol=2.9e-5
pval = pvalue(W, SWc)
@test pval ≈ 0.018 atol=4.7e-5

t = ShapiroWilkTest(X)
@test t.W == W
@test pvalue(t) == pval
@test t.N1 == length(X)
@test t.SWc.N == length(X)
@test length(t.SWc) == length(X)

@test pvalue(t) == pval


#=
Tools for solving the standard optimal savings / income fluctuation
problem for an infinitely lived consumer facing an exogenous income
process that evolves according to a Markov chain.

@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date: 2014-08-18

References
----------

http://quant-econ.net/jl/ifp.html

=#
using Interpolations
using Optim

# utility and marginal utility functions
u(x) = log(x)
du(x) = 1 ./ x

"""
Income fluctuation problem

##### Fields

- `r::Float64` : Strictly positive interest rate
- `R::Float64` : The interest rate plus 1 (strictly greater than 1)
- `bet::Float64` : Discount rate in (0, 1)
- `b::Float64` :  The borrowing constraint
- `Pi::Matrix{Floa64}` : Transition matrix for `z`
- `z_vals::Vector{Float64}` : Levels of productivity
- `asset_grid::LinSpace{Float64}` : Grid of asset values

"""
type ConsumerProblem
    r::Float64
    R::Float64
    bet::Float64
    b::Float64
    Pi::Matrix{Float64}
    z_vals::Vector{Float64}
    asset_grid::LinSpace{Float64}
end

function ConsumerProblem(;r=0.01, bet=0.96, Pi=[0.6 0.4; 0.05 0.95],
                         z_vals=[0.5, 1.0], b=0.0, grid_max=16, grid_size=50)
    R = 1 + r
    asset_grid = linspace(-b, grid_max, grid_size)

    ConsumerProblem(r, R, bet, b, Pi, z_vals, asset_grid)
end

"""
Given a matrix of size `(length(cp.asset_grid), length(cp.z_vals))`, construct
an interpolation object that does linear interpolation in the asset dimension
and has a lookup table in the z dimension
"""
function Interpolations.interpolate(cp::ConsumerProblem, x::AbstractMatrix)
    sz = (length(cp.asset_grid), length(cp.z_vals))
    if size(x) != sz
        msg = "x must have dimensions $(sz)"
        throw(DimensionMismatch(msg))
    end

    itp = interpolate(x, (BSpline(Linear()), NoInterp()), OnGrid())
    scale(itp, cp.asset_grid, 1:sz[2])
end

"""
Apply the Bellman operator for a given model and initial value.

##### Arguments

- `cp::ConsumerProblem` : Instance of `ConsumerProblem`
- `v::Matrix`: Current guess for the value function
- `out::Matrix` : Storage for output
- `;ret_policy::Bool(false)`: Toggles return of value or policy functions

##### Returns

None, `out` is updated in place. If `ret_policy == true` out is filled with the
policy function, otherwise the value function is stored in `out`.

"""
function bellman_operator!(cp::ConsumerProblem, V::Matrix, out::Matrix;
                           ret_policy::Bool=false)
    # simplify names, set up arrays
    R, Pi, bet, b = cp.R, cp.Pi, cp.bet, cp.b
    asset_grid, z_vals = cp.asset_grid, cp.z_vals

    z_idx = 1:length(z_vals)
    vf = interpolate(cp, V)

    # compute lower_bound for optimization
    opt_lb = 1e-8

    # solve for RHS of Bellman equation
    for (i_z, z) in enumerate(z_vals)
        for (i_a, a) in enumerate(asset_grid)

            function obj(c)
                y = 0.0
                for j in z_idx
                    y += vf[R*a+z-c, j] * Pi[i_z, j]
                end
                return -u(c)  - bet * y
            end

            res = optimize(obj, opt_lb, R.*a.+z.+b)
            c_star = res.minimum

            if ret_policy
                out[i_a, i_z] = c_star
            else
               out[i_a, i_z] = - obj(c_star)
            end

        end
    end
    out
end

bellman_operator(cp::ConsumerProblem, V::Matrix; ret_policy=false) =
    bellman_operator!(cp, V, similar(V); ret_policy=ret_policy)

"""
Extract the greedy policy (policy function) of the model.

##### Arguments

- `cp::CareerWorkerProblem` : Instance of `CareerWorkerProblem`
- `v::Matrix`: Current guess for the value function
- `out::Matrix` : Storage for output

##### Returns

None, `out` is updated in place to hold the policy function

"""
get_greedy!(cp::ConsumerProblem, V::Matrix, out::Matrix) =
    bellman_operator!(cp, V, out, ret_policy=true)

get_greedy(cp::ConsumerProblem, V::Matrix) =
    bellman_operator(cp, V, ret_policy=true)

"""
The approximate Coleman operator.

Iteration with this operator corresponds to policy function
iteration. Computes and returns the updated consumption policy
c.  The array c is replaced with a function cf that implements
univariate linear interpolation over the asset grid for each
possible value of z.

##### Arguments

- `cp::CareerWorkerProblem` : Instance of `CareerWorkerProblem`
- `c::Matrix`: Current guess for the policy function
- `out::Matrix` : Storage for output

##### Returns

None, `out` is updated in place to hold the policy function

"""
function coleman_operator!(cp::ConsumerProblem, c::Matrix, out::Matrix)
    # simplify names, set up arrays
    R, Pi, bet, b = cp.R, cp.Pi, cp.bet, cp.b
    asset_grid, z_vals = cp.asset_grid, cp.z_vals
    z_size = length(z_vals)
    gam = R * bet
    vals = Array(Float64, z_size)

    cf = interpolate(cp, c)

    # linear interpolation to get consumption function. Updates vals inplace
    cf!(a, vals) = map!(i->cf[a, i], vals, 1:z_size)

    # compute lower_bound for optimization
    opt_lb = 1e-8

    for (i_z, z) in enumerate(z_vals)
        for (i_a, a) in enumerate(asset_grid)
            function h(t)
                cf!(R*a+z-t, vals)  # update vals
                expectation = dot(du(vals), vec(Pi[i_z, :]))
                return abs(du(t) - max(gam * expectation, du(R*a+z+b)))
            end
            opt_ub = R*a + z + b  # addresses issue #8 on github
            res = optimize(h, min(opt_lb, opt_ub - 1e-2), opt_ub,
                           method=Optim.Brent())
            out[i_a, i_z] = res.minimum
        end
    end
    out
end

"""
Apply the Coleman operator for a given model and initial value

See the specific methods of the mutating version of this function for more
details on arguments
"""
coleman_operator(cp::ConsumerProblem, c::Matrix) =
    coleman_operator!(cp, c, similar(c))

function init_values(cp::ConsumerProblem)
    # simplify names, set up arrays
    R, bet, b = cp.R, cp.bet, cp.b
    asset_grid, z_vals = cp.asset_grid, cp.z_vals
    shape = length(asset_grid), length(z_vals)
    V, c = Array(Float64, shape...), Array(Float64, shape...)

    # Populate V and c
    for (i_z, z) in enumerate(z_vals)
        for (i_a, a) in enumerate(asset_grid)
            c_max = R*a + z + b
            c[i_a, i_z] = c_max
            V[i_a, i_z] = u(c_max) ./ (1 - bet)
        end
    end

    return V, c
end


# ---
# title: 741. Cherry Pickup
# id: problem741
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/cherry-pickup/description/>
# hidden: true
# ---
#
# In a N x N `grid` representing a field of cherries, each cell is one of three
# possible integers.
#
#
#
#   * 0 means the cell is empty, so you can pass through;
#   * 1 means the cell contains a cherry, that you can pick up and pass through;
#   * -1 means the cell contains a thorn that blocks your way.
#
#
#
# Your task is to collect maximum number of cherries possible by following the
# rules below:
#
#
#
#   * Starting at the position (0, 0) and reaching (N-1, N-1) by moving right or down through valid path cells (cells with value 0 or 1);
#   * After reaching (N-1, N-1), returning to (0, 0) by moving left or up through valid path cells;
#   * When passing through a path cell containing a cherry, you pick it up and the cell becomes an empty cell (0);
#   * If there is no valid path between (0, 0) and (N-1, N-1), then no cherries can be collected.
#
#
#
#
#
# **Example 1:**
#
#
#
#     Input: grid =
#     [[0, 1, -1],
#      [1, 0, -1],
#      [1, 1,  1]]
#     Output: 5
#     Explanation:
#     The player started at (0, 0) and went down, down, right right to reach (2, 2).
#     4 cherries were picked up during this single trip, and the matrix becomes [[0,1,-1],[0,0,-1],[0,0,0]].
#     Then, the player went left, up, up, left to return home, picking up one more cherry.
#     The total number of cherries picked up is 5, and this is the maximum possible.
#
#
#
#
# **Note:**
#
#   * `grid` is an `N` by `N` 2D array, with `1 <= N <= 50`.
#   * Each `grid[i][j]` is an integer in the set `{-1, 0, 1}`.
#   * It is guaranteed that grid[0][0] and grid[N-1][N-1] are not -1.
#   *
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


# Randomization option slows down processing
# but improves quality of solution considerably
# Would be better to do randomization in place
function cddual{T<:Real}(X::AbstractMatrix{T},
                         Y::AbstractVector{T};
                         C::Real = 1.0,
                         norm::Integer = 2,
                         randomized::Bool = true,
                         maxpasses::Integer = 2)
    # l: # of samples
    # n: # of features
    n, l = size(X)
    alpha = zeros(Float64, l)
    w = zeros(Float64, n)

    # Set U and D
    #  * L1-SVM: U = C, D[i] = 0
    #  * L2-SVM: U = Inf, D[i] = 1 / (2C)
    U = 0.0
    if norm == 1
        U = C
        D = zeros(Float64, l)
    elseif norm == 2
        U = Inf
        D = fill(1.0 / (2.0 * C), l)
    else
        throw(ArgumentError("Only L1-SVM and L2-SVM are supported"))
    end

    # Set Qbar
    Qbar = Array{Float64,1}(l)
    for i in 1:l
        Qbar[i] = D[i] + dot(X[:, i], X[:, i])
    end

    # Loop over examples
    converged = false
    pass = 0

    while !converged
        # Assess convergence
        pass += 1
        if pass == maxpasses
            converged = true
        end

        # Choose order of observations to process
        if randomized
            indices = randperm(l)
        else
            indices = 1:l
        end

        # Process all observations
        for i in indices
            g = Y[i] * dot(w, X[:, i]) - 1.0 + D[i] * alpha[i]

            if alpha[i] == 0.0
                pg = min(g, 0.0)
            elseif alpha[i] == U
                pg = max(g, 0.0)
            else
                pg = g
            end

            if abs(pg) > 0.0
                alphabar = alpha[i]
                alpha[i] = min(max(alpha[i] - g / Qbar[i], 0.0), U)
                for j in 1:n
                    w[j] = w[j] + (alpha[i] - alphabar) * Y[i] * X[j, i]
                end
            end
        end
    end

    return SVMFit(w, pass, converged)
end


using Bokeh
using Compat
using Base.Dates

start = now()
x = map(d -> start + Dates.Hour(d), 1:(30 * 24))
println(typeof(x))
# x: Array{DateTime,1}
day = 1000 * 3600 * 24
y = 0.5sin.(Dates.value.(x) * pi / day) .+ cos.(Dates.value.(x) * pi / (day * 7))
# note x needs to be an array of type DateTime or Date
plot(x, y)
showplot("dates_xy.html")

start = Date(1980, 1, 1)
x = map(d -> start + Dates.Day(d), 1:(365 * 30))
println(typeof(x))
# x: Array{Date,1}
y = 0.5sin.(Dates.value.(x) * pi / 365) .+ cos.(Dates.value.(x) * pi / (365 * 7))

plot(y, x)
showplot("dates_yx.html")


using Test

using DesignAutomation

@testset "bracket" begin
    fn(x) = 2 - 4 * x + exp(x)
    res = [round(i, digits=3) for i in bracket(fn, .5)]
    @test res == [.5, 1.309, 2.618]
end


#-----------------------------------------------------------------------
# JuMPeR  --  JuMP Extension for Robust Optimization
# http://github.com/IainNZ/JuMPeR.jl
#-----------------------------------------------------------------------
# Copyright (c) 2016: Iain Dunning
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
#-----------------------------------------------------------------------
# example/inventory.jl
# Implements the inventory management example from
#   A. Ben-Tal, A. Goryashko, E. Guslitzer, A. Nemirovski
#   "Adjustable Robust Solutions of Uncertain Linear Programs"
# Requires a linear optimization solver.
#-----------------------------------------------------------------------

using JuMP, JuMPeR, Gurobi

# Define problem parameters
I = 3           # Number of factories
T = 24          # Number of time periods
d_nom = 1000*[1 + 0.5*sin(π*(t-1)/12)  for t in 1:T]  # Nominal demand
θ = 0.20        # Uncertainty level
α = [1.0, 1.5, 2.0]  # Production costs
c = [α[i] * (1 + 0.5*sin(π*(t-1)/12))  for i in 1:I, t in 1:T]
P = 567         # Maximimum production per period
Q = 13600       # Maximumum production over all
Vmin = 500      # Minimum inventory at warehouse
Vmax = 2000     # Maximum inventory at warehouse
v1 = Vmin       # Initial inventory (not provided in paper)

# Setup robust model
invmgmt = RobustModel()

# Uncertain parameter: demand at each time stage lies in a interval
@uncertain(invmgmt, d_nom[t]*(1-θ) <= d[t=1:T] <= d_nom[t]*(1+θ))

# Decision: how much to produce at each factory at each time
# As this decision can be updated as demand is realized, we will use adaptive
# policy - in particular, an affine policy where production at time t is an
# affine function of the demand realized previously.
@adaptive(invmgmt, p[i=1:I,t=1:T], policy=Affine, depends_on=d[1:t-1])

# Objective: minimize total cost of production
@variable(invmgmt, F)  # Overall cost
@objective(invmgmt, Min, F)
@constraint(invmgmt, F >= sum{c[i,t] * p[i,t], i=1:I, t=1:T})

# Constraint: cannot exceed production limits
for i in 1:I, t in 1:T
    @constraint(invmgmt, p[i,t] >= 0)
    @constraint(invmgmt, p[i,t] <= P)
end
for i in 1:I
    @constraint(invmgmt, sum{p[i,t], t=1:T} <= Q)
end

# Constraint: cannot exceed inventory limits
for t in 1:T
    @constraint(invmgmt,
        v1 + sum{p[i,s], i=1:I, s=1:t} - sum{d[s],s=1:t} >= Vmin)
    @constraint(invmgmt,
        v1 + sum{p[i,s], i=1:I, s=1:t} - sum{d[s],s=1:t} <= Vmax)
end

# Solve
setsolver(invmgmt, GurobiSolver())
status = solve(invmgmt)

println(getobjectivevalue(invmgmt))