# This file is a part of SolidStateDetectors.jl, licensed under the MIT License (MIT).

"""
    add_fano_noise(E_dep::RealQuantity, E_ionisation::RealQuantity, f_fano::Real)::RealQuantity

Adds Fano noise to an energy deposition `E_dep`, assuming a detector material
ionisation energy `E_ionisation` and a Fano factor `f_fano`.

## Arguments
* `E_dep::RealQuantity`: Energy deposited in a semiconductor material.
* `E_ionisation`: Energy needed to create one electron-hole-pair in the semiconductor material.
* `f_fano`: Fano factor of the material.

## Example
```julia
add_fano_noise(100u"keV", 2.95u"eV", 0.129)
```

Some material properties are stored in `SolidStateDetectors.material_properties` and can be used here:
```julia
material = SolidStateDetectors.material_properties[:HPGe]
add_fano_noise(100u"keV", material.E_ionisation, material.f_fano)
```

!!! note
    Using values with units for `E_dep` or `E_ionisation` requires the package [Unitful.jl](https://github.com/PainterQubits/Unitful.jl).
"""
function add_fano_noise end
export add_fano_noise

function add_fano_noise(rng::AbstractRNG, E_dep::RealQuantity, E_ionisation::RealQuantity, f_fano::Real)
    target_unit = unit(E_dep)
    n_expected = uconvert(Unitful.NoUnits, E_dep/E_ionisation)

    # dist = Distributions.Poisson(n_expected)
    # n_observed = n_expected + f_fano * (rand(dist) - n_expected)

    # Use normal distribution instead of Poisson, makes no differenc for
    # typical numbers of charges:
    σ = sqrt(n_expected)
    n_observed = n_expected + f_fano * σ * randn(rng)
    n_observed_pos = max(0, n_observed)

    uconvert(target_unit, E_ionisation * n_observed_pos)
end

function add_fano_noise(rng::AbstractRNG, evts::DetectorHitEvents, E_ionisation::RealQuantity, f_fano::Real)
    noisy_edep = deepmap(x -> add_fano_noise(x, E_ionisation, f_fano), evts.edep)

    TypedTables.Table(merge(
        Tables.columns(evts),
        (edep = noisy_edep,)
    ))
end

add_fano_noise(x, E_ionisation::RealQuantity, f_fano::Real) =
    add_fano_noise(Random.GLOBAL_RNG, x, E_ionisation, f_fano)


"""
    module Partition

This submodule contains the following partition functions to be used with circular
coordinates.

* `Partition.linear(r, d) = max(r - d, 0.0)`
* `Partition.quadratic(r, d) = (max(r - d, 0.0))^2`
* `Partition.exponential(r, d) = r - d > 0 ? exp(r^2/(d^2 - r^2)) : 0.0`

"""
module Partition
"""
    Partition.linear(r, d) = max(r - d, 0.0)
"""
linear(r, d) = max(r - d, 0.0)
"""
    Partition.quadratic(r, d) = (max(r - d, 0.0))^2
"""
quadratic(r, d) = (max(r - d, 0.0))^2
"""
    Partition.exponential(r, d) = r - d > 0 ? exp(r^2/(d^2 - r^2)) : 0.0
"""
exponential(r, d) = r - d > 0 ? exp(r^2 / (d^2 - r^2)) : 0.0
end

###
### Landmark selection
###
function _maxmin_sample(points, n, metric=Euclidean())
    l = rand(1:length(points))
    landmarks = zeros(Int, n)
    i = 1
    closest_landmark = fill(Inf, length(points))
    while i ≤ n
        landmarks[i] = l
        for (j, p) in enumerate(points)
            dist = metric(SVector(points[l]), SVector(p))
            closest_landmark[j] = min(closest_landmark[j], dist)
        end
        i += 1
        l = argmax(closest_landmark)
    end
    return SVector.(points[landmarks]), maximum(closest_landmark)
end

function _landmarks_and_radius(points, landmarks::Int, metric)
    return _maxmin_sample(points, landmarks, metric)
end
function _landmarks_and_radius(points, landmarks::AbstractVector{Int}, metric)
    return _landmarks_and_radius(points, points[landmarks], metric)
end
function _landmarks_and_radius(points, landmarks, metric)
    radius = 0.0
    for p in points
        nearest = minimum(metric(SVector(l), SVector(p)) for l in landmarks)
        radius = max(nearest, radius)
    end
    return SVector.(landmarks), radius
end

###
### Cocycle massage
###
function _to_integer_coefficients(chain::Chain{Mod{M},S}) where {M,S}
    result = Chain{Int,S}()
    resize!(result, length(chain))
    for (i, (sx, coef)) in enumerate(chain)
        new_coef = Int(coef) > (M - 1) ÷ 2 ? Int(coef) - M : Int(coef)
        @inbounds result[i] = (sx, Int(new_coef))
    end
    return result
end

function _zero_coboundary_matrix(filtration::AbstractFiltration, time)
    is = Int[]
    js = Int[]
    vs = Float64[]
    for e in edges(filtration)
        if birth(e) < time
            i = index(e)
            u, v = e
            append!(is, (i, i))
            append!(js, (u, v))
            append!(vs, (-1.0, 1.0))
        end
    end
    n = nv(filtration)
    return sparse(is, js, vs, n * (n - 1) ÷ 2, n)
end

function _to_vector(filtration::AbstractFiltration, cocycle::Chain{Int}, time)
    n = nv(filtration)
    vector = zeros(n * (n - 1) ÷ 2)
    for (sx, c) in cocycle
        if birth(sx) < time
            vector[index(sx)] = Float64(c)
        end
    end
    return vector
end

function _mod_z(t)
    t = rem(t, 1.0)
    return t < 0.0 ? 1.0 + t : t
end

function _harmonic_smoothing(filtration, chain; time)
    integral_cocycle = _to_integer_coefficients(chain)
    if !is_cocycle(filtration, integral_cocycle, time)
        error(
            "the cocycle cannot be converted to `Int` coefficients. ",
            "Try using changing the landmarks, or using a different `modulus`.",
        )
    end
    coboundary = _zero_coboundary_matrix(filtration, time)
    real_cocycle = _to_vector(filtration, integral_cocycle, time)

    minimizer = coboundary \ real_cocycle
    real_cocycle .-= coboundary * minimizer
    return minimizer, real_cocycle
end

const PRIMES = [7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79]

struct CircularCoordinateData
    radius::Float64
    coordinate::Vector{Float64}
    cocycle::Vector{Float64}
end

"""
    CircularCoordinates

This `struct` implements Perea's sparse circular coordinates (see reference below).

The idea behind this method is that you pick a subset of your data set (also called
landmarks), compute persistent cohomology on that subset, and use the result to construct
circular coordinates on the whole data set.

For this to work correctly, you need the following.

* A set of landmarks that cover the points well. The default minmax sampling method is
usually a good choice here.
* The persistence diagram on the landmarks must have an interval that is persistent enough.

The circular coordinates returned are on the interval [0, 1). If you are looking for angles,
make sure to multiply them by `2π`.

If you try to compute circular coordinates for a point that is not near any landmark, the
result will be `missing`.

# Constructor

`CircularCoordinates([::AbstractFiltration, ]points, landmarks; kwargs...)`

# Arguments

* `points`: a vector of points. The `eltype` of this vector can be `Tuple`s, `SVector`s, or
  similar.

* `landmarks`: can be an integer, a vector of indices, or a vector of points. If set to an
  integer, it sets the number of landmarks to choose with maxmin sampling. If you are
  looking for non-sparse circular coordinates, use `landmarks=eachindex(points)`.

* `out_dim`: number of most persistent persistence intervals to use to compute coordinates.
  If less than `out_dim` suitable intervals are found, the construction will show a warning
  and construct `CircularCoordinates` with a lower `out_dim`. This warning can be suppressed
  by passing `warn=false`.

* `partition`: a function that defines the partition of unity used when determining the
  coordinates. Should take two arguments, `r`, the radius of the balls, and `d`, the
  distance from the landmark. The partition function should only have support on the ball
  around the landmark (in other words, it should evaluate to 0 when d ≥ r). If that is not
  the case, the circular coordinates are not well-defined and the results may not always
  make sense. See [`Partition`](@ref) for a small collection of predefined functions.

* `metric`: a metric from [Distances.jl](https://github.com/JuliaStats/Distances.jl). Can
  only be set if the filtration argument is `Rips`.

* `kwargs...`: additional keyword arguments passed to `ripserer`. Note that `modulus` is set
  to a random prime between 7 and 79 by default.

# Example

```jldoctest
julia> data = [(sin(t), cos(t)) for t in range(0, 2π, length=101)[1:end-1]];

julia> cc = CircularCoordinates(data, 1:10:100)
CircularCoordinates(
  out_dim=1,
  radius=(0.9510565162951535,),
  n_landmarks=10,
  partition=linear,
  metric=Euclidean(0.0),
)

julia> summary(cc(data))
"100×1 Matrix{Union{Missing, Float64}}"

julia> summary(cc(data, 1))
"100-element Vector{Union{Missing, Float64}}"

```

# Reference

Perea, J. A. (2020, June). Sparse circular coordinates via principal Z-bundles. In
Topological Data Analysis: The Abel Symposium 2018 (Vol. 15, p. 435). Springer Nature.

"""
struct CircularCoordinates{F,P<:SVector,M}
    out_dim::Int
    landmarks::Vector{P}
    partition::F
    metric::M

    coordinate_data::Vector{CircularCoordinateData}
    meta::NamedTuple

    partition_buffer::Vector{Float64}
end

function CircularCoordinates(args...; kwargs...)
    return CircularCoordinates(Rips, args...; kwargs...)
end
function CircularCoordinates(
    ::Type{F},
    points::AbstractVector,
    landmarks;
    out_dim=1,
    partition=Partition.linear,
    metric=Euclidean(),
    coverage=1,
    modulus=rand(PRIMES),
    threshold=nothing,
    verbose=false,
    warn=true,
    kwargs...,
) where {F<:AbstractFiltration}
    @prog_print verbose "Determining radius...   "
    landmarks, min_radius = _landmarks_and_radius(points, landmarks, metric)
    @prog_println verbose "done."

    # TODO with new ripserer interface, change this.
    flt_kwargs = metric == Euclidean() ? NamedTuple() : (metric = Euclidean())
    flt_kwargs = isnothing(threshold) ? flt_kwargs : (; threshold=threshold, flt_kwargs...)

    # compute cohomology
    @prog_print verbose "Computing cohomology... "
    filtration = F(landmarks; flt_kwargs...)
    diagram = ripserer(filtration; modulus=modulus, reps=true, kwargs...)[2]
    @prog_println verbose "done."

    @prog_print verbose "Smoothing cocycles...   "
    coord_data = CircularCoordinateData[]
    for d in 1:out_dim
        if length(diagram) < d
            break
        end
        interval = diagram[end + 1 - d]
        b, d = interval
        if !isfinite(death(interval))
            d = threshold(filtration)
        end
        if max(b, min_radius) ≥ d / 2
            break
        else
            radius = (1 - coverage) * max(b, min_radius) + coverage * d / 2
        end
        coords, cocycle = _harmonic_smoothing(
            filtration, interval.representative; time=2 * radius
        )
        push!(coord_data, CircularCoordinateData(radius, coords, cocycle))
    end
    @prog_println verbose "done."
    n_success = length(coord_data)
    if n_success == 0
        error("no interval is persistent enough to cover the data")
    elseif n_success < out_dim
        s = n_success == 1 ? " was" : "s were"
        warn && @warn "only $n_success interval$s persistent enough to cover the data"
    end

    meta = (
        filtration=filtration,
        diagram=diagram,
        intervals=diagram[end:-1:(end + 1 - n_success)],
        modulus=modulus,
    )
    return CircularCoordinates(
        n_success, landmarks, partition, metric, coord_data, meta, zeros(length(landmarks))
    )
end

function Base.show(io::IO, cc::CircularCoordinates)
    return print(
        io,
        "CircularCoordinates(\n",
        "  out_dim=$(cc.out_dim),\n",
        "  radius=$(tuple((d.radius for d in cc.coordinate_data)...)),\n",
        "  n_landmarks=$(length(cc.landmarks)),\n",
        "  partition=$(cc.partition),\n",
        "  metric=$(cc.metric),\n)",
    )
end

# Transform single
function _transform(cc::CircularCoordinates, point, dim)
    cd = cc.coordinate_data[dim]

    selected_ball = 0
    min_distance = Inf
    for (j, l) in enumerate(cc.landmarks)
        dist = cc.metric(point, l)
        value = cc.partition(cd.radius, dist)
        if isinf(value)
            cc.partition_buffer .= 0.0
            cc.partition_buffer[j] = 1.0
            selected_ball = j
            break
        else
            cc.partition_buffer[j] = value
        end
        if dist < min_distance
            selected_ball = j
            min_distance = dist
        end
    end
    norm = sum(cc.partition_buffer)
    if norm > 0
        cc.partition_buffer ./= norm
    else
        return missing
    end

    j = selected_ball
    coord = cd.coordinate[j]
    for k in eachindex(cc.landmarks)
        j == k && continue
        cocycle_value = j > k ? -cd.cocycle[index((j, k))] : +cd.cocycle[index((k, j))]
        coord += cc.partition_buffer[k] * cocycle_value
    end
    return _mod_z(coord)
end

function (cc::CircularCoordinates)(points, dims=1:(cc.out_dim))
    if dims isa Integer
        result = zeros(Union{Missing,Float64}, length(points))
    else
        result = zeros(Union{Missing,Float64}, length(points), cc.out_dim)
    end
    for (j, dim) in enumerate(dims)
        for (i, p) in enumerate(points)
            result[i, j] = _transform(cc, SVector(p), dim)
        end
    end
    return result
end


# Demo using DiscreteHMM and Baum-Welch algorithm
# to estimate a simple model of natural language.
include("models.jl")

module demo_discrete
using DiscreteHMM
can_save = (Pkg.installed("JLD")!=nothing)  # check if JLD is installed
if can_save; using JLD; end

# Settings
m = 100  # number of hidden states to use in estimated model
n_max = 100000  # maximum number of words/symbols to use
from_file = false  # load estimated parameters from file
run_profiler = false  # run profiler to assess performance
tolerance = 1e-2  # convergence tolerance for Baum-Welch
n_gen = 1000  # length of sequence to generate using estimated parameters
filename = "simple.txt"  # file containing a sample of simple natural language

# Read and preprocess data
function preprocess(filename,n_max)
	# Read file
	f = open(filename)
	s = readall(f)
	close(f)
	# Remove whitespace and remove some punctuation
	s = replace(s,r"\s+"," ")
	s = replace(s,r"[\"]","")
	s = replace(s,r"' | '"," ")
	s = replace(s,r"(\W) | (\W)",s" \1 ")
	s = strip(s)
	words = split(s,' ')
	# Consider only the first n_max words/symbols
	words = words[1:min(n_max,length(words))]
	# Extract vocabulary
	vocabulary = sort(unique(words))
	V = length(vocabulary)
	code = Dict(zip(vocabulary,1:V))  # mapping from words to indices
	return words,vocabulary,code
end

# Produce the numerical sequence corresponding to a sequence of words.
function encode(words,code)
	return Int[code[w] for w in words]
end

# Produce the sequence of text corresponding to a numerical sequence.
function compose(x,vocabulary)
	s = join(vocabulary[x]," ")
	s = replace(s,r" (\W) ",s"\1 ")
	return s
end

# Preprocess and encode the text in the given file
words,vocabulary,code = preprocess(filename,n_max)
V = length(vocabulary)
x = encode(words,code)
n = length(x)

# Save encoding to file
#writedlm("x.txt",x)
#writedlm("code.txt",[(1:V) vocabulary],' ')

# Sanity check of forward algorithm implementation
if true
	log_pi = log(ones(m)/m)
	log_T = log(ones(m,m)/m)
	phi = [ones(V)/V for r = 1:m]
	G,log_m = DiscreteHMM.forward(x,log_pi,log_T,phi)
	log_m_exact = -n*log(V)
	@assert(abs(log_m - log_m_exact) < 1e-10, "Log marginal does not match")
end

# Estimate parameters using Baum-Welch
if from_file
	(log_pi,log_T,phi,log_m) = load("estimated_params.jld","params")
else
    if run_profiler
        @profile log_pi,log_T,phi,log_m = DiscreteHMM.estimate(x,m,tolerance)
        Profile.print() #format = :flat)
        Profile.clear()
    else
		log_pi,log_T,phi,log_m = DiscreteHMM.estimate(x,m,tolerance)
	end
	if can_save
		save("estimated_params.jld","params",(log_pi,log_T,phi,log_m))
	end
end

# Generate a new sequence from the estimated model
println("Generating a sample of text using the estimated model:")
y,zy = DiscreteHMM.generate(n_gen,log_pi,log_T,phi)
println(compose(y,vocabulary))

end # module


include("utils.jl")

@testset "Circle" begin
    M = Circle()
    @testset "Real Circle Basics" begin
        @test representation_size(M) == ()
        @test manifold_dimension(M) == 1
        @test !is_manifold_point(M, 9.)
        @test_throws DomainError is_manifold_point(M, 9., true)
        @test !is_tangent_vector(M, 9., 0.)
        @test_throws DomainError is_tangent_vector(M, 9., 0., true)
        @test flat(M,0.0, FVector(TangentSpace,1.0)) == FVector(CotangentSpace,1.0)
        @test sharp(M,0.0, FVector(CotangentSpace,1.0)) == FVector(TangentSpace,1.0)
        @test vector_transport_to(M,0.0,1.0,1.0, ParallelTransport()) == 1.0
        @test retract(M,0.0,1.0) == exp(M,0.0,1.0)
        @test injectivity_radius(M) ≈ π
        @test mean(M, [-π/2,0.,π]) ≈ π/2
        @test mean(M, [-π/2,0.,π],[1., 1., 1.]) == π/2
        v = MVector(0.0)
        x = SVector(0.0)
        log!(M,v,x,SVector(π/4))
        @test norm(M,x,v) ≈ π/4
        @test is_tangent_vector(M,x,v)
        @test project_point(M,1.0) == 1.0
        x = MVector(0.0)
        project_point!(M,x)
        @test x == MVector(0.0)
        x .+= 2*π
        project_point!(M,x)
        @test x == MVector(0.0)
    end
    types = [Float64, Float32]

    basis_types = (ArbitraryOrthonormalBasis(),)
    basis_types_real = (ArbitraryOrthonormalBasis(),
        DiagonalizingOrthonormalBasis([-1]),
        DiagonalizingOrthonormalBasis([1])
    )
    for T in types
        @testset "Type $T" begin
            pts = convert.(Ref(T), [-π/4,0.,π/4])
            test_manifold(
                M,
                pts,
                test_forward_diff = false,
                test_reverse_diff = false,
                test_vector_spaces = false,
                test_project_tangent = true,
                test_musical_isomorphisms = true,
                test_vector_transport = true,
                is_mutating = false
            )
            ptsS = SVector.(pts)
            test_manifold(
                M,
                ptsS,
                test_forward_diff = false,
                test_reverse_diff = false,
                test_project_tangent = true,
                test_musical_isomorphisms = true,
                test_vector_transport = true,
                basis_types_vecs = basis_types_real,
                basis_types_to_from = basis_types_real
            )
        end
    end
    Mc = Circle(ℂ)
    @testset "Complex Circle Basics" begin
        @test representation_size(Mc) == ()
        @test manifold_dimension(Mc) == 1
        @test is_tangent_vector(Mc, 1im, 0.)
        @test is_manifold_point(Mc, 1im)
        @test !is_manifold_point(Mc, 1+1im)
        @test_throws DomainError is_manifold_point(Mc, 1+1im, true)
        @test !is_tangent_vector(Mc, 1+1im, 0.)
        @test_throws DomainError is_tangent_vector(Mc, 1+1im, 0., true)
        @test !is_tangent_vector(Mc, 1im, 2im)
        @test_throws DomainError is_tangent_vector(Mc, 1im, 2im, true)
        @test flat(Mc,0.0+0.0im, FVector(TangentSpace,1.0im)) == FVector(CotangentSpace,1.0im)
        @test sharp(Mc,0.0+0.0im, FVector(CotangentSpace,1.0im)) == FVector(TangentSpace,1.0im)
        @test norm(Mc,1.0,log(Mc, 1.0,-1.0)) ≈ π
        @test is_tangent_vector(Mc,1.0,log(Mc,1.0,-1.0))
        v = MVector(0.0+0.0im)
        x = SVector(1.0+0.0im)
        log!(Mc,v,x,SVector(-1.0+0.0im))
        @test norm(Mc,SVector(1.0),v) ≈ π
        @test is_tangent_vector(Mc,x,v)
        @test project_point(Mc,1.0) == 1.0
        project_point(Mc,1/sqrt(2.0) + 1/sqrt(2.0) * im) == 1/sqrt(2.0) + 1/sqrt(2.0) * im
        x = MVector(1.0+0.0im)
        project_point!(Mc,x)
        @test x == MVector(1.0+0.0im)
        x .*= 2
        project_point!(Mc,x)
        @test x == MVector(1.0+0.0im)

    end
    types = [Complex{Float64}, Complex{Float32}]
    for T in types
        @testset "Type $T" begin
            a = 1/sqrt(2.0)
            pts = convert.(Ref(T), [a-a*im, 1+0im, a+a*im])
            test_manifold(
                Mc,
                pts,
                test_forward_diff = false,
                test_reverse_diff = false,
                test_vector_spaces = false,
                test_project_tangent = true,
                test_musical_isomorphisms = true,
                test_vector_transport = true,
                is_mutating=false,
                exp_log_atol_multiplier = 2.0,
                is_tangent_atol_multiplier = 2.0
            )
            ptsS = SVector.(pts)
            test_manifold(
                Mc,
                ptsS,
                test_forward_diff = false,
                test_reverse_diff = false,
                test_project_tangent = true,
                test_musical_isomorphisms = true,
                test_vector_transport = true,
                exp_log_atol_multiplier = 2.0,
                is_tangent_atol_multiplier = 2.0,
                basis_types_vecs = basis_types,
                basis_types_to_from = basis_types
            )
        end
    end
end


export UpdateEloF

"""
    UpdateEloF <: UpdateRule

 Update using standard Elo where
* K is constant, but
* performance model distribution can be set (with constant parameters)
* extra factors are added to ratings, with a scale parameter

## Parameters
* `r0::Float64`: default rating
* `K::Float64`: the value of the gain parameter (use EloF2 for variable K)
* `dist::ContinuousUnivariateDistribution`: the model to use for performance as a function of rating, e.g, Normal or Logistic
* `factor_scale::Float64`
* the distribution can have its own parameters, e.g., `β::Float64`, the spread parameter for the Logistic

```
"""
struct UpdateEloF <: UpdateRule
    r0::Float64
    K::Float64
    dist::ContinuousUnivariateDistribution
    factor_scale::Float64
    function UpdateEloF(r0, K, dist, factor_scale)
        @check_args(UpdateEloF, K >= zero(K))
        @check_args(UpdateEloF, r0 >= zero(r0))
        @check_args(UpdateEloF, factor_scale >= zero(factor_scale))
        new(r0, K, dist, factor_scale)
    end
end
# defaults imply the distribution is  Logistic(μ=0.0, θ=100.0)
#   https://en.wikipedia.org/wiki/Logistic_distribution
#   https://juliastats.github.io/Distributions.jl/latest/univariate.html#Continuous-Distributions-1
#      mean=mu=0.0, sdev=π θ / √3
default_θ = 400 / log(10) # this results in the standard values used for Elo in Chess and many other cases
UpdateEloF(;r0=1500.0, K=32.0, θ=default_θ, factor_scale=default_θ) = UpdateEloF( r0, K, Logistic(0.0, θ), factor_scale )
# default factor_scale is one "standard deviation"

function update_info( rule::UpdateEloF )
    info = Dict(
                :name => "EloF",
                :reference => "\"Whos's #1\", Langville and Meyer, p.?? and ??",
                :computation => "simultaneous",
                :state_model => "recursive",
                :input => "outcome,factors",
                :output => "probabilistic",
                :model => "single",
                :ties => true,
                :factors => true,
                :parameters => ["r0, default rating", "K, gain", "dist(), performance model", "factor_scale"],
                :record => false
                )
    return info
end
update_info( ::Type{UpdateEloF} ) =  update_info( UpdateEloF() )

function update_ratings( rule::UpdateEloF,
                         input_ratings::RatingsList,
                         input_competitions::DataFrame)
    n = size(input_competitions,1)
    m = length( input_ratings.players )
    I = player_indexes( input_ratings.players )

    # set initial ratings if they aren't already
    old_r = copy( input_ratings.ratings )
    for p in input_ratings.players
        if !haskey( old_r, p )
            old_r[p] = rule.r0
        end
    end
    new_r = copy(old_r) # start with the old ratings, in case one isn't modified at all

    # perform update for each outcome
    d = input_competitions # just an abbreviation
    for i=1:n
        pA = d[i,PlayerA]
        pB = d[i,PlayerB]
        old_rA = old_r[ pA ]
        old_rB = old_r[ pB ]
        (ΔrA, ΔrB) = update( rule, old_rA, old_rB, d[i,Outcome], rule.factor_scale*d[i,FactorA], rule.factor_scale*d[i,FactorB])
        new_r[ pA ] += ΔrA
        new_r[ pB ] += ΔrB
    end

    # output ratings list
    output_ratings = RatingsList(input_ratings.players, new_r )
    return output_ratings
end

# see Elo.jl for definition of `update`


import AngularMomentumAlgebra: @linearly_combinable

struct MySym
    s::Symbol
end
@linearly_combinable MySym

@testset "Linear combinations" begin
    x = MySym(:x)
    y = MySym(:y)

    w = x + y
    @test length(w) == 2
    @test eltype(w) == (MySym,Int)
    @test w.Ts == [x,y]
    @test w.coeffs == [1,1]

    @test (-x).coeffs == [-1]
    @test length((-x)) == 1

    z = 4x - 5y
    z′ = -z
    z′′ = 2z

    @test z.Ts == [x,y]
    @test z.coeffs == [4, -5]

    @test z′.Ts == [x,y]
    @test z′.coeffs == [-4, 5]

    @test z′′.Ts == [x,y]
    @test z′′.coeffs == [8, -10]

    @test all((z′′/2).coeffs .== z.coeffs)

    @test [(t,c) for (t,c) in z] == [(x,4), (y,-5)]
end


import Base
mutable struct Transpose <: DiffableFunction
    in_is_vector
    grad_field::GradField
    Transpose(in_shape, grad_field) = new(length(in_shape) == 1, grad_field)
end

function forward(::Transpose, x)
    return transpose(x)
end

function backward(f::Transpose, gy)
    gx = transpose(gy)
    (f.in_is_vector) && (return flatten(gx))
    return gx
end

"""
    transpose(x::Variable)
# Examples
```julia-repl
julia> x = Variable([1 2; 3 4])
name: nothing
values: [1 2; 3 4]
creator: User-Defined(nothing)

julia> transpose(x)
name: nothing
values: [1 3; 2 4]
creator: JITrench.Transpose
```
"""
Base.transpose(x::Variable) = Transpose(size(x.values), GradField())(x)
is_support(::typeof(Base.transpose)) = true


function compute()
N=4000000
ans=0
x=1
y=2
while x<=N
	if(x%2==0)
		 ans+=x
	end
	 x,y=y,x+y

end
ans
end
print(compute())


using SparseArrays: sparse
using LinearAlgebra: cholesky, ldiv!
# import Convex
# import ECOS

"""Fit Ax = b using Huber loss in ADMM (alternatinve directions method of multiplers)
Adapted from https://web.stanford.edu/~boyd/papers/admm/huber/huber_fit.html
"""
# TODO: differentiate this from the regularization alpha
function huber_fit(
    A,
    b,
    rho = 1.0,
    alpha = 1.0;
    lu_tuple = nothing,
    quiet = true,
    max_iter = 1000,
    abstol = 1e-4,
    reltol = 1e-2,
)
    m, n = size(A)
    Atb = A' * b

    x = zeros(n)
    z = zeros(m)
    zold = zeros(m)
    u = zeros(m)
    # If you prefactor the A matrix
    if !isnothing(lu_tuple)
        L, U = lu_tuple
    else
        L, U = factor(A)
    end

    idx = 0

    kappa = 1 + 1 / rho  # For shrinkage
    r_factor = rho / (1 + rho)
    s_factor = 1 / (1 + rho)
    while idx < max_iter
        # x update
        zold .= z
        q = Atb .+ A' * (z - u)
        x .= U \ (L \ q)

        # x update with relaxation
        Ax_hat = alpha .* A * x + (1 - alpha) * (z .+ b)
        tmp = Ax_hat - b + u
        z .= (r_factor .* tmp) + (s_factor .* shrinkage(tmp, kappa))
        u += (Ax_hat - z - b)

        # Stopping check
        r_norm = norm(A * x - z - b)

        s_norm = norm(BLAS.gemv('T', -rho, A, (z - zold)))
        # ^^ equivalent to this below,  but faster
        # s_norm  = norm(-rho * A' * (z - zold))

        eps_pri = sqrt(n) * abstol + reltol * maximum([norm(A * x), norm(-z), norm(b)])
        eps_dual = sqrt(n) * abstol + reltol * norm(rho * u)

        if !quiet
            @show r_norm, eps_pri, s_norm, eps_dual, objective(z)
        end

        if r_norm < eps_pri && s_norm < eps_dual
            if !quiet
                println("Number of iterations to converge: $idx")
            end
            return x
        else
            idx += 1
        end

    end
    println("Caution: problem did not converge within $max_iter iterations to within $abstol, $reltol (alpha: $alpha, rho: $rho, sizes: $(size(A)), $(size(b)))")
    return x
end

"""Evaluate the huber loss with a bandwidth of `M`"""
function huber(x; M = 1)
    y = max.(x, 0)
    z = min.(y, M)
    return z .* (2 .* y .- z)
end

objective(z) = sum(huber(z)) / 2

# Sparse factoring
function factor(A)
    ch = cholesky(A' * A)
    return sparse(ch.L), sparse(ch.U)
end

pos(x) = max.(x, 0)

# Soft threshold operator (section 4.4.3)
shrinkage(x, kappa) = pos(1 .- kappa ./ abs.(x)) .* x

# ### Alternate: Solve the L1 fitting problem using Convex
# # Note: this had a memory leak earlier when run on many pixels,
# # possibly from the ECOS library
# function l1_fit(A, b, x::Union{Convex.Variable, Nothing}=nothing, quiet=true)
#     solver = ECOS.ECOSSolver(verbose=0)
#     if isnothing(x)
#         x = Convex.Variable(size(A, 2))
#     end
#
#     problem = Convex.minimize(norm(A*x - b, 1))
#     Convex.solve!(problem, solver)
#     return x.value
# end

#### IRLS algorithm functions: ####
l1_objective(A, x, b) = sum(abs.(A * x - b))

"""Iteratively reweighted least squares (IRLS), used to solve L1 minimization
Source: https://en.wikipedia.org/wiki/Iteratively_reweighted_least_squares"""
function irls(
    A::AbstractArray{<:AbstractFloat,2},
    b::AbstractArray{<:AbstractFloat,1};
    p::Int = 1,
    iters = 50,
)
    M, N = size(A)

    # Use Float64 to avoid roundoff NaNs
    x = Array{Float64,1}(undef, N)
    W = zeros(Float64, (size(A, 1), size(A, 1)))

    ep = sqrt(eps(eltype(x)))
    W .= diagm(0 => (abs.(b - A * x) .+ ep).^(p - 2))

    for ii = 1:iters
        x .= (A' * W * A) \ (A' * W * b)
        W .= diagm(0 => (abs.(b - A * x) .+ ep).^(p - 2))
    end
    # println("objective: ", l1_objective(A, x, b))
    return Float32.(x)
end

# # Function compatible for run_sbas
# function invert_pixel(pixel::AbstractArray{T, 1}, B::AbstractArray{T,2}; iters=50, p=1) where {T<:AbstractFloat}
#     return irls(B, pixel, iters=iters, p=p)
# end


mutable struct CholeskyFactorization{Tv, Ti} <: AbstractLUFactorization{Tv,Ti}
    A::Union{ExtendableSparseMatrix{Tv,Ti},Nothing}
    fact::Union{SuiteSparse.CHOLMOD.Factor{Tv},Nothing}
    phash::UInt64
    A64
end

"""
CholeskyFactorization(;valuetype=Float64, indextype=Int64)
CholeskyFactorization(matrix)

Default Cholesky factorization via cholmod.
"""
CholeskyFactorization(;valuetype::Type=Float64, indextype::Type=Int64)=CholeskyFactorization{valuetype,indextype}(nothing,nothing,0,nothing)

function update!(cholfact::CholeskyFactorization)
    A=cholfact.A
    flush!(A)
    if A.phash!=cholfact.phash
        cholfact.A64=Symmetric(SparseMatrixCSC{Float64,Int64}(A.cscmatrix))
        cholfact.fact=cholesky(cholfact.A64)
        cholfact.phash=A.phash
    else
        cholfact.A64.data.nzval.=A.cscmatrix.nzval
        cholfact.fact=cholesky!(cholfact.fact,cholfact.A64)
    end
    cholfact
end

LinearAlgebra.ldiv!(fact::CholeskyFactorization, v)=fact.fact\v
LinearAlgebra.ldiv!(u,fact::CholeskyFactorization, v)=u.=fact.fact\v


struct Problem
    C::Vector{Int64} # Vecteur des pofits de chaque variables
    A::Union{Vector{Vector{Int64}}, Array{Int64, 2}} # Matrice creuse/Liste d'adgacence des contraintes

    #Donées annexes

	nbVar::Int64
	nbConst::Int64
	equalWeights::Bool

    listOccur::Array{Int64, 2}
    occurVar :: Vector{Int64} # Nombre de variable que chaque variable penalise
    occurConst::Vector{Int64} #Nombre de fois qu'une variable apparait  dans une contraintes
end

function Problem(fname::String = "Data/test.txt")::Problem

	    f=open(fname)

	    # lecture du nbre de contraintes (m) et de variables (n)
	    m, n = parse.(Int, split(readline(f)) )

	    # lecture des n coefficients de la fonction economique et cree le vecteur d'entiers c
	    C = parse.(Int, split(readline(f)) )

		# Savoir si tout les profits sont égaux ou non
		equalWeights::Bool = true
		for var = 1:n
			equalWeights = equalWeights && (C[var] == C[1])
		end

	    # lecture des m contraintes et reconstruction de la matrice binaire A
	    A = zeros(Int64, m, n)

		# Matrice décrivant la corrélation entre les variables
	    listOccur = zeros(Int64, n, n)

		# Liste du nombre de variable correlée à une variable
	    occurVar = zeros(Int64, n)

		# Liste du nombre de contraintes où apparaissent une variable
	    occurConst = zeros(Int64, n)

	    for i=1:m
	        # lecture du nombre d'elements non nuls sur la contrainte i (non utilise)
	        nbVarConst = parse.(Int, split(readline(f)))

	        # lecture des indices des elements non nuls sur la contrainte i
	        line = parse.(Int, split(readline(f)))

	        for var in line
	            A[i, var] = 1
	            occurConst[var] += 1
	            for var2 in line
	                listOccur[var, var2] = 1
	                occurVar[var] += 1
	            end
	        end
	    end
	    close(f)

	    return Problem(C, A, n, m, equalWeights, listOccur, occurVar, occurConst)
end

struct Solution
    x::Vector{Int64}
    z::Int64

	# Données annexes

	var1::Vector{Int64} # Indice des variables misent à 1
	var0::Vector{Int64} # Indice des variables misent à 0

	nbVarLibre::Int64
	estLibre::BitArray
end

function Solution(x::Vector{Int64}, z:: Int64, estLibre::BitArray)::Solution
	nbVar::Int64 = length(x)

	num1 = sum(x)
	var1 = Vector{Int64}(undef, num1)
	var0 = Vector{Int64}(undef, nbVar-num1)

	num0 = 0
	num1 = 0
	for ind = 1:length(x)
		if x[ind] == 0
			num0 += 1
			var0[num0] = ind
		else
			num1 += 1
			var1[num1] = ind
		end
	end

	nbVarLibre::Int64 = 0
	for ind = 1:nbVar
		if estLibre[ind] == 1
			nbVarLibre += 1
		end
	end
	return Solution(x, z, var1, var0, nbVarLibre, estLibre)
end

import Base.copy

function Base.copy(sol::Solution)::Solution
	return Solution(
		copy(x),
		z,
		copy(var1),
		copy(var0),
		nbVarLibre,
		copy(estLibre)
	)
end


# -*- coding: utf-8 -*-
# %%
# Copyright 2020-2021 John T. Foster
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# %%
include("bspline.jl")
using LinearAlgebra, IterativeSolvers
export construct_spline_matrix, reconstruct_trajectory, construct_helix, L2error

function construct_helix(n::Integer = 100)
    """
    Helper function which outputs an nx3 array of form [cos.(x)' sin.(x)'  linspace(0,1,n+1)]
    """
    t = LinRange(0, 4*π, n+1)
    arr = zeros(Float64, (length(t), 3))
    arr[:, 1] = cos.(t)
    arr[:, 2] = sin.(t)
    arr[:, 3] = LinRange(0, 1, n+1)
    arr
end

function create_knot_vector(Qk::Vector{<:Float64}, p::Integer = 3)
    ū = create_ūk(Qk)
    n = length(ū)
    m = length(ū) + p + 1
    kv = zeros(m)
    for j = 2:(n - p)
        kv[j+p] = sum(ū[j:(j+p-1)]) / float(p)
    end
    kv[(end - p):end] .= 1
    kv
end

#Knots = m
#From definition, m = n + p + 1
#set m = length(tangents) so the linear algebra works out
#then set n (num samples) = m - p - 1
function reconstruct_trajectory(tangents::Matrix{<:Float64}, p::Integer=3)
    control_points = zeros(size(tangents))
    control_points[:,1] = reconstruct_trajectory_1d(tangents[:,1], p)
    control_points[:,2] = reconstruct_trajectory_1d(tangents[:,2], p)
    control_points[:,3] = tangents[:,3] #Don't reconstruct the domain
    kv = create_knot_vector(tangents[:,1], p)
    basis = BSplineBasis(kv, p, k=2)
    curve = BSplineCurve(basis, control_points)
    return curve
end

function reconstruct_trajectory_1d(tangents::Vector{<:Float64}, p::Integer=3)
    kv = create_knot_vector(tangents, p)
    ū = create_ūk(tangents) # n = m - p - 1
    basis = BSplineBasis(kv, p, k=2)
    N, Nprime = construct_spline_matrix(basis, ū, length(kv), p)
    control_points = lsmr(N, tangents)
    return control_points
end

function construct_spline_matrix(basis::BSplineBasis, samples::Vector{<:Float64}, num_knots::Integer, p::Integer=3)
    rows, cols = length(samples), num_knots - p - 1
    N, Nprime = zeros(Float64, (rows, cols)), zeros(Float64, (rows, cols))
    N[1, 1] = 1
    N[end, end] = 1
    Nprime[1, 1] = 1
    Nprime[end, end] = 1
    for i in p:rows
        evals = basis(samples[i])
        column = find_knot_span(basis, samples[i])
        N[i,column-p:column] = evals[1,:]#evals[1, :][evals[1, :] .> 0]
        Nprime[i,column-p:column] = evals[2,:]#evals[2, :][evals[2, :] .> 0]
    end
    N, Nprime
end

function L2error(C::BSplineCurve, tangents::Matrix{<:Float64})
    """
    Returns a vector with the cumulative sum of the L2 error for each dimension.
    """
    evals = evaluate(C, length(tangents[:,1]))
    error = zeros(length(tangents[:,1]))
    for (i, u) in enumerate(evals[:,3]')
        index = findall(x-> x>=u, tangents[:,3])[1] #Find index of tangents vec closest to z axis of our control point
        error[i] = abs(tangents[index,2] - evals[i,2]) + abs(tangents[index,3] - evals[i,3])
    end
    return error
end

function main(n::Integer=100, p::Integer=3)
    Q = construct_helix(n)
    Curve = reconstruct_trajectory(Q, p)
    Q, Curve
end


using Statistics
using Plots
using FFTW
using Statistics
using Optim
using Images, FileIO, ImageIO
using Printf
using Revise
using Profile
using LinearAlgebra
using JLD2
using Random
using Distributions
using FITSIO
using LineSearches
using SparseArrays

push!(LOAD_PATH, pwd()*"/../../../main")
using DHC_2DUtils
push!(LOAD_PATH, pwd()*"/../../../scratch_NM")
using Deriv_Utils_New
using Data_Utils
using Visualization
using ReconFuncs

numfile = 1000

logbool = false
apdbool = true
isobool = false
invcovtype = "Full+Eps"

direc = "scratch_NM/StandardizedExp/Nx64/" #Change
datfile = direc * "Data_" * string(numfile) * ".jld2" #Replace w SLURM array
loaddf = load(datfile)
true_img = loaddf["true_img"]
init = loaddf["init"]

fname_save = direc * "SFDTargSFDCov/" * string(logbool) * "_" * string(apdbool) * "_" * string(isobool) * "/LambdaVary/" * string(numfile) * "rand_init"  #Change
Nx=size(true_img)[1]
filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(S1iso, Nf) = size(filter_hash["S1_iso_mat"])
(S2iso, Nfsq) = size(filter_hash["S2_iso_mat"])
dhc_args = Dict(:doS2=>false, :doS20=>true, :apodize=>apdbool, :iso=>isobool) #Iso #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
if dhc_args[:iso]
    coeff_mask = falses(2+S1iso+S2iso)
    coeff_mask[S1iso+3:end] .= true
else #Not iso
    coeff_mask = falses(2+Nf+Nf^2)
    coeff_mask[Nf+3:end] .= triu(trues(Nf, Nf))[:]
end

white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :norm=>false, :smooth=>false, :smoothval=>0.8) #Iso #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", ConjugateGradient(linesearch=LineSearches.BackTracking()))])
recon_settings = Dict([("target_type", "ground_truth"), ("covar_type", "sfd_dbn"), ("log", logbool), ("GaussianLoss", false), ("iteration_init", randn(Nx, Nx)), ("TransformedGaussianLoss", false), ("RandomInit", true), ("Invcov_matrix", invcovtype), ("epsvalue", 1e-8),
  ("optim_settings", optim_settings), ("white_noise_args", white_noise_args)]) #Add constraints

recon_settings["datafile"] = datfile

regs_true = DHC_compute_wrapper(log.(true_img), filter_hash; norm=false, dhc_args...)[coeff_mask]
s2mean, s2icov = meancov_generator(true_img, filter_hash, dhc_args, coeff_mask, recon_settings, safety=nothing)
#Check
#s2mean - s_true

if recon_settings["log"] & dhc_args[:apodize]
    lval = std(apodizer(log.(true_img))).^(-2)
elseif !recon_settings["log"] & dhc_args[:apodize]
    s_init = DHC_compute_wrapper(apodizer(init), filter_hash, norm=false; dhc_args...)[coeff_mask]
    l1init = ( 0.5 .* (s_init - s2mean)' * s2icov * (s_init - s2mean))[1]
    println("L1init", l1init)
    est1 = std(apodizer(true_img)).^(-2)
    wn_exp = mean(wind_2d(Nx).^2)*Nx*Nx*(loaddf["std"])^2
    lval = minimum([est1, 0.01*(l1init - 1.0)/(wn_exp)])
elseif recon_settings["log"] & !dhc_args[:apodize]
    lval = std(log.(true_img)).^(-2)
else
    s_init = DHC_compute_wrapper(init, filter_hash, norm=false; dhc_args...)[coeff_mask]
    l1init = ( 0.5 .* (s_init - s2mean)' * s2icov * (s_init - s2mean))[1]
    println("L1init", l1init)
    est1 = std(true_img).^(-2)
    wn_exp = Nx*Nx*(loaddf["std"])^2
    lval = minimum([est1, 0.01*(l1init - 1.0)/(wn_exp)])
end

recon_settings["lambda"] = lval
println("Regularizer Lambda=", round(lval, sigdigits=3))

recon_settings["s_targ_mean"] = s2mean
recon_settings["s_invcov"] = s2icov
recon_settings["safemode"] = false
recon_settings["fname_save"] = fname_save * ".jld2"
recon_settings["optim_settings"] = optim_settings
resobj, recon_img = reconstruction_wrapper(true_img, imfilter(init, Kernel.gaussian(0.8)), filter_hash, dhc_args, coeff_mask, recon_settings)

plot(heatmap(true_img), heatmap(recon_settings["iteration_init"]), heatmap(recon_img))

DHC_compute_wrapper(recon_img, filter_hash, norm=false; dhc_args...) .- DHC_compute_wrapper(true_img, filter_hash, norm=false; dhc_args...)

p= plot([t.value for t in Optim.trace(resobj)])
plot!(title="Loss: S20 | Targ = S(SFD) | Cov = SFD Covariance", xlabel = "No. Iterations", ylabel = "Loss")
savefig(p, fname_save * "_trace.png")

Visualization.plot_diffscales([apodizer(true_img), apodizer(init), apodizer(recon_img), apodizer(recon_img) - apodizer(true_img)], ["GT", "Init", "Reconstruction", "Residual"], fname=fname_save*".png")
Visualization.plot_synth_QA(apodizer(true_img), apodizer(init), apodizer(recon_img), fname=fname_save*"_6p.png")

println("Input Data", datfile)
println("Output File", fname_save)
println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))
println("Mean Abs Res: Init-True = ", mean((init - true_img).^2).^0.5, " Recon-True = ", mean((recon_img - true_img).^2).^0.5)


const threshold = 1e-8

myeqzero{T<:Real}(x::T) = x == zero(T)
myeqzero{T<:AbstractFloat}(x::T) = -threshold < x < threshold
myeqzero{T<:Real}(x::Vector{T}) = myeqzero(sum(abs(x)))

myeq{T<:Real}(x::T, y::T) = x == y
# I check with zero because isapprox(0, 1e-100) is false...
# but isapprox(1e-100, 2e-100) should be false
myeq{T<:AbstractFloat}(x::T, y::T) = (x == zero(T) ? myeqzero(y) : (y == zero(T) ? myeqzero(x) : isapprox(x, y)))
myeq{T<:Real, S<:Real}(x::T, y::S) = myeq(promote(x, y)...)

mylt{T<:Real}(x::T, y::T) = x < y
mylt{T<:AbstractFloat}(x::T, y::T) = x < y && !myeq(x, y)
mylt{S<:Real,T<:Real}(x::S, y::T) = mylt(promote(x, y)...)
mygt{S<:Real, T<:Real}(x::S, y::T) = mylt(y, x)
myleq{T<:Real}(x::T, y::T) = x <= y
myleq{T<:AbstractFloat}(x::T, y::T) = x <= y || myeq(x, y)
myleq{S<:Real,T<:Real}(x::S, y::T) = myleq(promote(x, y)...)
mygeq{T<:Real}(x::T, y::T) = myleq(y, x)
mypos{T<:Real}(x::T) = mygt(x, zero(T))
myneg{T<:Real}(x::T) = mylt(x, zero(T))
mynonneg{T<:Real}(x::T) = mygeq(x, zero(T))
mynonpos{T<:Real}(x::T) = myleq(x, zero(T))

function mypromote_type{T1,T2}(::Type{T1}, ::Type{T2})
  if T1 <: AbstractFloat || T2 <: AbstractFloat
    Float64
  else
    Rational{BigInt}
  end
end


using Statistics

# Quick and dirty Anderson implementation ... lacks important things like
# control of the condition number of the anderson matrix.
# Not particularly optimised. Also should be moved to NLSolve ...
#
# Accelerates the iterative solution of f(x) = 0 according to a
# damped preconditioned scheme
#    xₙ₊₁ = xₙ + αₙ P⁻¹ f(xₙ)
# Where f(x) computes the residual (e.g. SCF(x) - x)
# Further define
#    preconditioned residual  Pf(x) = P⁻¹ f(x)
#    fixed-point map          g(x)  = V + α Pf(x)
# where the α may vary between steps.
#
# Finds the linear combination xₙ₊₁ = g(xₙ) + ∑ᵢ βᵢ (g(xᵢ) - g(xₙ))
# such that |Pf(xₙ) + ∑ᵢ βᵢ (Pf(xᵢ) - Pf(xₙ))|² is minimal
struct AndersonAcceleration
    m::Int                  # maximal history size
    iterates::Vector{Any}   # xₙ
    residuals::Vector{Any}  # Pf(xₙ)
    maxcond::Real           # Maximal condition number for Anderson matrix
end
AndersonAcceleration(;m=10, maxcond=1e6) = AndersonAcceleration(m, [], [], maxcond)

function Base.push!(anderson::AndersonAcceleration, xₙ, αₙ, Pfxₙ)
    push!(anderson.iterates,  vec(xₙ))
    push!(anderson.residuals, vec(Pfxₙ))
    if length(anderson.iterates) > anderson.m
        popfirst!(anderson.iterates)
        popfirst!(anderson.residuals)
    end
    @assert length(anderson.iterates) <= anderson.m
    @assert length(anderson.iterates) == length(anderson.residuals)
    anderson
end

# Gets the current xₙ, Pf(xₙ) and damping αₙ
function (anderson::AndersonAcceleration)(xₙ, αₙ, Pfxₙ)
    xs   = anderson.iterates
    Pfxs = anderson.residuals

    # Special cases with fast exit
    anderson.m == 0 && return xₙ .+ αₙ .* Pfxₙ
    if isempty(xs)
        push!(anderson, xₙ, αₙ, Pfxₙ)
        return xₙ .+ αₙ .* Pfxₙ
    end

    M = hcat(Pfxs...) .- vec(Pfxₙ)  # Mᵢⱼ = (Pfxⱼ)ᵢ - (Pfxₙ)ᵢ
    # We need to solve 0 = M' Pfxₙ + M'M βs <=> βs = - (M'M)⁻¹ M' Pfxₙ

    # Ensure the condition number of M stays below maxcond, else prune the history
    Mfac = qr(M)
    while size(M, 2) > 1 && cond(Mfac.R) > anderson.maxcond
        M = M[:, 2:end]  # Drop oldest entry in history
        popfirst!(anderson.iterates)
        popfirst!(anderson.residuals)
        Mfac = qr(M)
    end

    xₙ₊₁ = vec(xₙ) .+ αₙ .* vec(Pfxₙ)
    βs   = -(Mfac \ vec(Pfxₙ))
    for (iβ, β) in enumerate(βs)
        xₙ₊₁ .+= β .* (xs[iβ] .- vec(xₙ) .+ αₙ .* (Pfxs[iβ] .- vec(Pfxₙ)))
    end

    push!(anderson, xₙ, αₙ, Pfxₙ)
    reshape(xₙ₊₁, size(xₙ))
end

function ScfAcceptStepAll()
    accept_step(info, info_next) = true
end

"""
Accept a step if the energy is at most increasing by `max_energy` and the residual
is at most `max_relative_residual` times the residual in the previous step.
"""
function ScfAcceptImprovingStep(;max_energy_change=1e-12, max_relative_residual=1.0)
    function accept_step(info, info_next)
        energy_change = info_next.energies.total - info.energies.total
        relative_residual = norm(info_next.Pinv_δV) / norm(info.Pinv_δV)

        # Accept if energy goes down or residual decreases
        accept = energy_change < max_energy_change || relative_residual < max_relative_residual
        mpi_master() && @debug "Step $(accept ? "accepted" : "discarded")" energy_change relative_residual
        accept
    end
end

"""
Use the two iteration states `info` and `info_next` to find a damping value
from a quadratic model for the SCF energy. Returns `nothing` if the constructed
model is not considered trustworthy, else returns the suggested damping.
"""
function scf_damping_quadratic_model(info, info_next; modeltol=0.1)
    T     = eltype(info.Vin)
    dvol  = info.basis.dvol

    Vin   = info.Vin
    ρin   = info.ρout      # = ρ(Vin)
    Vout  = info.Vout      # = step(Vin), where step(V) = (Vext + Vhxc(ρ(V)))
    α0    = info_next.α
    Vnext = info_next.Vin  # = Vin + α0 * (Anderson(Vin, P⁻¹( Vout - Vin )) - Vin)
    ρnext = info_next.ρout # = ρ(Vnext)
    δρ    = ρnext - ρin
    # α0 * δV = α0 * (Vnext - Vin) = α0 * (Anderson(Vin, P⁻¹( Vout - Vin )) - Vin)

    # We build a quadratic model for
    #   ϕ(α) = E(Vin  + α δV)  at α = 0
    #        = E(Vin) + α ∇E|_(V=Vin) ⋅ δV + ½ α^2 <δV | ∇²E|_(V=Vin) | δV>
    #        = E(Vin) + (α/α0) ∇E|_(V=Vin) ⋅ (α0 * δV) + ½ (α/α0)² <α0 * δV | ∇²E|_(V=Vin) | α0 * δV>
    #
    # Now
    #      ∇E|_(V=Vin)  = - χ₀(Vout - Vin)
    #      ∇²E|_(V=Vin) ≃ - χ₀ (1 - K χ₀)        (only true if Vin is an SCF minimum, K taken at Vin)
    # and therefore using the self-adjointness of χ₀
    #      ∇E|_(V=Vin) ⋅ δV         = -(Vout - Vin) ⋅ χ₀(δV) = - (Vout - Vin) ⋅ δρ
    #      <δV | ∇²E|_(V=Vin) | δV> = - δV ⋅ δρ + δρ ⋅ K(δρ)
    #
    slope = dot(Vout .- Vin, δρ) / α0 * dvol
    Kδρ   = apply_kernel(info.basis, δρ; ρ=ρin)
    curv  = dvol * (-dot(Vnext .- Vin, δρ) + dot(δρ, Kδρ)) / α0^2
    Emodel(α) = info.energies.total + slope * α + curv * α^2 / 2

    # Relative error of the model at α0 (the damping we used to get info_next)
    Etotal_next = info_next.energies.total

    # TODO Is this a good error measure ... for larger problems it seems to be over-demanding
    model_relerror = abs(Etotal_next - Emodel(α0)) / abs(Etotal_next - info.energies.total)

    minimum_exists = curv > eps(T)  # Does the model predict a minimum along the search direction
    trusted_model  = model_relerror < modeltol      # Model fits observation
    tight_model    = model_relerror < modeltol / 5  # Model fits observation very well

    # Accept model if it leads to minimum and is either tight or shows a negative slope
    α_model = -slope / curv
    if minimum_exists && (tight_model || (slope < -eps(T) && trusted_model))
        mpi_master() && @debug "Quadratic model accepted" model_relerror slope curv α_model
        (α=α_model, relerror=model_relerror)
    else
        mpi_master() && @debug "Quadratic model discarded" model_relerror slope curv α_model
        (α=nothing, relerror=model_relerror) # Model not trustworthy ...
    end
end

# Adaptive damping using a quadratic model
@kwdef struct AdaptiveDamping
    α_min = 0.05        # Minimal damping
    α_max = 1.0         # Maximal damping
    α_trial_init = 0.8  # Initial trial damping used (i.e. in the first SCF step)
    α_trial_min = 0.2   # Minimal trial damping used in a step
    α_trial_enhancement = 1.1  # Enhancement factor to α_trial in case a step is immediately successful
    modeltol = 0.1      # Maximum relative error on the predicted energy for model
    #                     to be considered trustworthy
end
function AdaptiveDamping(α_trial_min; kwargs...)
    # Select some reasonable defaults.
    # The free tweaking parameter here should be increased a bit for cases,
    # where the Anderson does weird stuff in case of too small damping.
    AdaptiveDamping(;α_min=α_trial_min / 4,
                     α_max=max(1.25α_trial_min, 1.0),
                     α_trial_init=max(α_trial_min, 0.8),
                     α_trial_min=α_trial_min,
                     kwargs...)
end

function ensure_damping_within_range(damping::AdaptiveDamping, α, α_next)
    α_sign = sign(α_next)
    abs(α_next) ≤ damping.α_min / 5 && (α_sign = +1.0)
    if α_sign > 0.0
        # Avoid getting stuck
        α_next = min(0.95abs(α), abs(α_next))
    else
        # Don't move too far backwards (where model validity cannot be ensured)
        α_next = min(0.50abs(α), abs(α_next))
    end
    α_next = clamp(α_next, damping.α_min, damping.α_max)
    α_sign * α_next
end

function propose_backtrack_damping(damping::AdaptiveDamping, info, info_next)
    if abs(info_next.α) < 1.75damping.α_min
        # Too close to α_min to be worth making another step ... just give up
        return info_next.α
    end

    α_next, relerror = scf_damping_quadratic_model(info, info_next; modeltol=damping.modeltol)
    if isnothing(α_next)
        # Model failed ... use heuristics: Half for small model error, else use a quarter
        α_next = info_next.α / (relerror < 10 ? 2 : 4)
    end
    ensure_damping_within_range(damping, info_next.α, α_next)
end

trial_damping(damping::AdaptiveDamping) = damping.α_trial_init
function trial_damping(damping::AdaptiveDamping, info, info_next, step_successful)
    n_backtrack = length(info_next.diagonalization)

    α_trial = abs(info_next.α)  # By default use the α that worked in this step
    if step_successful && n_backtrack == 1  # First step was good => speed things up
        α_trial ≥ damping.α_max && return damping.α_max  # No need to compute model
        α_model = scf_damping_quadratic_model(info, info_next; modeltol=damping.modeltol).α
        if !isnothing(α_model)  # Model is meaningful
            α_trial = max(damping.α_trial_enhancement * abs(α_model), α_trial)
        end
    end

    # Ensure returned damping is in valid range
    clamp(α_trial, damping.α_trial_min, damping.α_max)
end

struct FixedDamping
    α
end
FixedDamping() = FixedDamping(0.8)
propose_backtrack_damping(damping::FixedDamping) = damping.α
trial_damping(damping::FixedDamping, args...) = damping.α

# Notice: For adaptive damping to run smoothly, multiple defaults need to be changed.
#         See the function scf_potential_mixing_adaptive for that use case.
@timing function scf_potential_mixing(
    basis::PlaneWaveBasis;
    damping=FixedDamping(0.8),
    n_bands=default_n_bands(basis.model),
    ρ=guess_density(basis),
    V=nothing,
    ψ=nothing,
    tol=1e-6,
    maxiter=100,
    eigensolver=lobpcg_hyper,
    n_ep_extra=3,
    diag_miniter=1,
    determine_diagtol=ScfDiagtol(),
    mixing=SimpleMixing(),
    is_converged=ScfConvergenceEnergy(tol),
    callback=ScfDefaultCallback(),
    acceleration=AndersonAcceleration(;m=10),
    accept_step=ScfAcceptStepAll(),
    max_backtracks=3,  # Maximal number of backtracking line searches
)
    # TODO Test other mixings and lift this
    @assert (   mixing isa SimpleMixing
             || mixing isa KerkerMixing
             || mixing isa KerkerDosMixing)
    damping isa Number && (damping = FixedDamping(damping))

    if ψ !== nothing
        @assert length(ψ) == length(basis.kpoints)
        for ik in 1:length(basis.kpoints)
            @assert size(ψ[ik], 2) == n_bands + n_ep_extra
        end
    end

    # Initial guess for V (if none given)
    energies, ham = energy_hamiltonian(basis, nothing, nothing; ρ=ρ)
    isnothing(V) && (V = total_local_potential(ham))

    function EVρ(Vin; diagtol=tol / 10, ψ=nothing)
        ham_V = hamiltonian_with_total_potential(ham, Vin)
        res_V = next_density(ham_V; n_bands=n_bands, ψ=ψ, n_ep_extra=n_ep_extra,
                             miniter=diag_miniter, tol=diagtol, eigensolver=eigensolver)
        new_E, new_ham = energy_hamiltonian(basis, res_V.ψ, res_V.occupation;
                                            ρ=res_V.ρout, eigenvalues=res_V.eigenvalues,
                                            εF=res_V.εF)
        (basis=basis, ham=new_ham, energies=new_E, Vin=Vin,
         Vout=total_local_potential(new_ham), res_V...)
    end

    n_iter    = 1
    converged = false
    α_trial   = trial_damping(damping)
    diagtol   = determine_diagtol((ρin=ρ, Vin=V, n_iter=n_iter))
    info      = EVρ(V; diagtol=diagtol, ψ=ψ)
    Pinv_δV   = mix_potential(mixing, basis, info.Vout - info.Vin; n_iter, info...)
    info      = merge(info, (α=NaN, diagonalization=[info.diagonalization], ρin=ρ,
                             n_iter=n_iter, Pinv_δV=Pinv_δV))
    ΔEdown    = 0.0

    while n_iter < maxiter
        info = merge(info, (stage=:iterate, algorithm="SCF", converged=converged))
        callback(info)
        if MPI.bcast(is_converged(info), 0, MPI.COMM_WORLD)
            # TODO Debug why these MPI broadcasts are needed
            converged = true
            break
        end
        n_iter += 1
        info = merge(info, (n_iter=n_iter, ))

        # Ensure same α on all processors
        α_trial = MPI.bcast(α_trial, 0, MPI.COMM_WORLD)
        δV = (acceleration(info.Vin, α_trial, info.Pinv_δV) - info.Vin) / α_trial

        # Determine damping and take next step
        guess   = ψ
        α       = α_trial
        successful  = false  # Successful line search (final step is considered good)
        n_backtrack = 1
        diagonalization = empty(info.diagonalization)
        info_next = info
        while n_backtrack ≤ max_backtracks
            diagtol = determine_diagtol(info_next)
            mpi_master() && @debug "Iteration $n_iter linesearch step $n_backtrack   α=$α diagtol=$diagtol"
            Vnext = info.Vin .+ α .* δV

            info_next    = EVρ(Vnext; ψ=guess, diagtol=diagtol)
            Pinv_δV_next = mix_potential(mixing, basis, info_next.Vout - info_next.Vin;
                                         n_iter, info_next...)
            push!(diagonalization, info_next.diagonalization)
            info_next = merge(info_next, (α=α, diagonalization=diagonalization,
                                          ρin=info.ρout, n_iter=n_iter,
                                          Pinv_δV=Pinv_δV_next))

            successful = accept_step(info, info_next)
            successful = MPI.bcast(successful, 0, MPI.COMM_WORLD)  # Ensure same successful
            if successful || n_backtrack ≥ max_backtracks
                break
            end
            n_backtrack += 1

            # Adjust α to try again ...
            α_next = propose_backtrack_damping(damping, info, info_next)
            α_next = MPI.bcast(α_next, 0, MPI.COMM_WORLD)  # Ensure same α on all processors
            if α_next == α  # Backtracking further not useful ...
                break
            end

            # Adjust to guess fitting α best:
            guess = α_next > α / 2 ? info_next.ψ : ψ
            α = α_next
        end

        # Switch off acceleration in case of very bad steps
        ΔE = info_next.energies.total - info.energies.total
        ΔE < 0 && (ΔEdown = -max(abs(ΔE), tol))

        # Update α_trial and commit the next state
        α_trial = trial_damping(damping, info, info_next, successful)
        info = info_next
    end

    ham  = hamiltonian_with_total_potential(ham, info.Vout)
    info = (ham=ham, basis=basis, energies=info.energies, converged=converged,
            ρ=info.ρout, eigenvalues=info.eigenvalues, occupation=info.occupation,
            εF=info.εF, n_iter=n_iter, n_ep_extra=n_ep_extra, ψ=info.ψ,
            diagonalization=info.diagonalization, stage=:finalize, algorithm="SCF")
    callback(info)
    info
end

# Wrapper function setting a few good defaults for adaptive damping
function scf_potential_mixing_adaptive(basis; tol=1e-6, damping=AdaptiveDamping(), kwargs...)
    @assert damping isa AdaptiveDamping
    scf_potential_mixing(basis; tol=tol, diag_miniter=2,
                         accept_step=ScfAcceptImprovingStep(max_energy_change=tol),
                         determine_diagtol=ScfDiagtol(ratio_ρdiff=0.03, diagtol_max=5e-3),
                         damping=damping, kwargs...)
end


using Test

function parseInput(input)
    chars = split(input, ",")
    chars = parse.(Int, chars)
    return chars
end

function getVal(val, codes, mode)
    if mode == 0
        return codes[val + 1]
    elseif mode == 1
        return val
    end
end

function part1(codes)
    i = 1
    outputs = []
    while codes[i] != 99
        operation = codes[i]
        numbers = digits(operation, pad=5)
        instruction = numbers[1] + 10*numbers[2]
        mode1 = numbers[3]
        mode2 = numbers[4]
        mode3 = numbers[5]
        if instruction == 1
            val1, val2, dest = codes[i+1:i+3]
            codes[dest+1] = getVal(val1, codes, mode1) + getVal(val2, codes, mode2) #codes[val2]
            i += 4
        elseif instruction == 2
            val1, val2, dest = codes[i+1:i+3]
            codes[dest+1] = getVal(val1, codes, mode1) * getVal(val2, codes, mode2) #codes[val2]
            i += 4
        elseif instruction == 3
            dest = codes[i+1] + 1
            codes[dest] = 1
            i += 2
        elseif instruction == 4
            dest = codes[i+1] + 1
            append!(outputs, codes[dest])
            i += 2
        end
    end
    return codes, outputs
end

function run(s)
    input = parseInput(split(s, "\n")[1])
    # input[2] = 12
    # input[3] = 2
    return part1(input)[2][end]
end

@test part1(parseInput("1002,4,3,4,33"))[1] == [1002, 4, 3, 4, 99]

function main()
    res, time, memory = @timed run(ARGS[1])
    println("_duration:$(time*1000)")
    println(res)
end

main()