text:
# Testing:
#
#  - computation of sufficient statistics
#  - distribution fitting (i.e. estimation)
#

using Distributions
using Test, Random, LinearAlgebra

n0 = 100
N = 10^5

rng = MersenneTwister(123)

const funcs = ([rand,rand], [dist -> rand(rng, dist), (dist, n) -> rand(rng, dist, n)])

@testset "Testing fit for DiscreteUniform" begin
    for func in funcs
        w = func[1](n0)

        x = func[2](DiscreteUniform(10, 15), n0)
        d = fit(DiscreteUniform, x)
        @test isa(d, DiscreteUniform)
        @test minimum(d) == minimum(x)
        @test maximum(d) == maximum(x)

        d = fit(DiscreteUniform, func[2](DiscreteUniform(10, 15), N))
        @test minimum(d) == 10
        @test maximum(d) == 15
    end
end

@testset "Testing fit for Bernoulli" begin
    for func in funcs, dist in (Bernoulli, Bernoulli{Float64})
        w = func[1](n0)
        x = func[2](dist(0.7), n0)

        ss = suffstats(dist, x)
        @test isa(ss, Distributions.BernoulliStats)
        @test ss.cnt0 == n0 - count(t->t != 0, x)
        @test ss.cnt1 == count(t->t != 0, x)

        ss = suffstats(dist, x, w)
        @test isa(ss, Distributions.BernoulliStats)
        @test ss.cnt0 ≈ sum(w[x .== 0])
        @test ss.cnt1 ≈ sum(w[x .== 1])

        d = fit(dist, x)
        p = count(t->t != 0, x) / n0
        @test isa(d, dist)
        @test mean(d) ≈ p

        d = fit(dist, x, w)
        p = sum(w[x .== 1]) / sum(w)
        @test isa(d, dist)
        @test mean(d) ≈ p

        d = fit(dist, func[2](dist(0.7), N))
        @test isa(d, dist)
        @test isapprox(mean(d), 0.7, atol=0.01)
    end
end

@testset "Testing fit for Beta" begin
    for func in funcs, dist in (Beta, Beta{Float64})
        d = fit(dist, func[2](dist(1.3, 3.7), N))
        @test isa(d, dist)
        @test isapprox(d.α, 1.3, atol=0.1)
        @test isapprox(d.β, 3.7, atol=0.1)

        d = fit_mle(dist, func[2](dist(1.3, 3.7), N))
        @test isa(d, dist)
        @test isapprox(d.α, 1.3, atol=0.1)
        @test isapprox(d.β, 3.7, atol=0.1)

    end
end

@testset "Testing fit for Binomial" begin
    for func in funcs, dist in (Binomial, Binomial{Float64})
        w = func[1](n0)

        x = func[2](dist(100, 0.3), n0)

        ss = suffstats(dist, (100, x))
        @test isa(ss, Distributions.BinomialStats)
        @test ss.ns ≈ sum(x)
        @test ss.ne == n0
        @test ss.n == 100

        ss = suffstats(dist, (100, x), w)
        @test isa(ss, Distributions.BinomialStats)
        @test ss.ns ≈ dot(Float64[xx for xx in x], w)
        @test ss.ne ≈ sum(w)
        @test ss.n == 100

        d = fit(dist, (100, x))
        @test isa(d, dist)
        @test ntrials(d) == 100
        @test succprob(d) ≈ sum(x) / (n0 * 100)

        d = fit(dist, (100, x), w)
        @test isa(d, dist)
        @test ntrials(d) == 100
        @test succprob(d) ≈ dot(x, w) / (sum(w) * 100)

        d = fit(dist, 100, func[2](dist(100, 0.3), N))
        @test isa(d, dist)
        @test ntrials(d) == 100
        @test isapprox(succprob(d), 0.3, atol=0.01)
    end
end

# Categorical

@testset "Testing fit for Categorical" begin
    for func in funcs
        p = [0.2, 0.5, 0.3]
        x = func[2](Categorical(p), n0)
        w = func[1](n0)

        ss = suffstats(Categorical, (3, x))
        h = Float64[count(v->v == i, x) for i = 1 : 3]
        @test isa(ss, Distributions.CategoricalStats)
        @test ss.h ≈ h

        d = fit(Categorical, (3, x))
        @test isa(d, Categorical)
        @test ncategories(d) == 3
        @test probs(d) ≈ h / sum(h)

        d2 = fit(Categorical, x)
        @test isa(d2, Categorical)
        @test probs(d2) == probs(d)

        ss = suffstats(Categorical, (3, x), w)
        h = Float64[sum(w[x .== i]) for i = 1 : 3]
        @test isa(ss, Distributions.CategoricalStats)
        @test ss.h ≈ h

        d = fit(Categorical, (3, x), w)
        @test isa(d, Categorical)
        @test probs(d) ≈ h / sum(h)

        d = fit(Categorical, suffstats(Categorical, 3, x, w))
        @test isa(d, Categorical)
        @test probs(d) ≈ (h / sum(h))

        d = fit(Categorical, func[2](Categorical(p), N))
        @test isa(d, Categorical)
        @test isapprox(probs(d), p, atol=0.01)
    end
end

@testset "Testing fit for Cauchy" begin
    @test fit(Cauchy, collect(-4.0:4.0)) === Cauchy(0.0, 2.0)
    @test fit(Cauchy{Float64}, collect(-4.0:4.0)) === Cauchy(0.0, 2.0)
end

@testset "Testing fit for Exponential" begin
    for func in funcs, dist in (Exponential, Exponential{Float64})
        w = func[1](n0)
        x = func[2](dist(0.5), n0)

        ss = suffstats(dist, x)
        @test isa(ss, Distributions.ExponentialStats)
        @test ss.sx ≈ sum(x)
        @test ss.sw == n0

        ss = suffstats(dist, x, w)
        @test isa(ss, Distributions.ExponentialStats)
        @test ss.sx ≈ dot(x, w)
        @test ss.sw == sum(w)

        d = fit(dist, x)
        @test isa(d, dist)
        @test scale(d) ≈ mean(x)

        d = fit(dist, x, w)
        @test isa(d, dist)
        @test scale(d) ≈ dot(x, w) / sum(w)

        d = fit(dist, func[2](dist(0.5), N))
        @test isa(d, dist)
        @test isapprox(scale(d), 0.5, atol=0.01)
    end
end

@testset "Testing fit for Normal" begin
    for func in funcs, dist in (Normal, Normal{Float64})
        μ = 11.3
        σ = 3.2
        w = func[1](n0)

        x = func[2](dist(μ, σ), n0)

        ss = suffstats(dist, x)
        @test isa(ss, Distributions.NormalStats)
        @test ss.s  ≈ sum(x)
        @test ss.m  ≈ mean(x)
        @test ss.s2 ≈ sum((x .- ss.m).^2)
        @test ss.tw ≈ n0

        ss = suffstats(dist, x, w)
        @test isa(ss, Distributions.NormalStats)
        @test ss.s  ≈ dot(x, w)
        @test ss.m  ≈ dot(x, w) / sum(w)
        @test ss.s2 ≈ dot((x .- ss.m).^2, w)
        @test ss.tw ≈ sum(w)

        d = fit(dist, x)
        @test isa(d, dist)
        @test d.μ ≈ mean(x)
        @test d.σ ≈ sqrt(mean((x .- d.μ).^2))

        d = fit(dist, x, w)
        @test isa(d, dist)
        @test d.μ ≈ dot(x, w) / sum(w)
        @test d.σ ≈ sqrt(dot((x .- d.μ).^2, w) / sum(w))

        d = fit(dist, func[2](dist(μ, σ), N))
        @test isa(d, dist)
        @test isapprox(d.μ, μ, atol=0.1)
        @test isapprox(d.σ, σ, atol=0.1)
    end
end

@testset "Testing fit for Normal with known moments" begin
    import Distributions.NormalKnownMu, Distributions.NormalKnownSigma
    μ = 11.3
    σ = 3.2

    for func in funcs

        w = func[1](n0)
        x = func[2](Normal(μ, σ), n0)

        ss = suffstats(NormalKnownMu(μ), x)
        @test isa(ss, Distributions.NormalKnownMuStats)
        @test ss.μ == μ
        @test ss.s2 ≈ sum(abs2.(x .- μ))
        @test ss.tw ≈ n0

        ss = suffstats(NormalKnownMu(μ), x, w)
        @test isa(ss, Distributions.NormalKnownMuStats)
        @test ss.μ == μ
        @test ss.s2 ≈ dot((x .- μ).^2, w)
        @test ss.tw ≈ sum(w)

        d = fit_mle(Normal, x; mu=μ)
        @test isa(d, Normal)
        @test d.μ == μ
        @test d.σ ≈ sqrt(mean((x .- d.μ).^2))

        d = fit_mle(Normal, x, w; mu=μ)
        @test isa(d, Normal)
        @test d.μ == μ
        @test d.σ ≈ sqrt(dot((x .- d.μ).^2, w) / sum(w))

        ss = suffstats(NormalKnownSigma(σ), x)
        @test isa(ss, Distributions.NormalKnownSigmaStats)
        @test ss.σ == σ
        @test ss.sx ≈ sum(x)
        @test ss.tw ≈ n0

        ss = suffstats(NormalKnownSigma(σ), x, w)
        @test isa(ss, Distributions.NormalKnownSigmaStats)
        @test ss.σ == σ
        @test ss.sx ≈ dot(x, w)
        @test ss.tw ≈ sum(w)

        d = fit_mle(Normal, x; sigma=σ)
        @test isa(d, Normal)
        @test d.σ == σ
        @test d.μ ≈ mean(x)

        d = fit_mle(Normal, x, w; sigma=σ)
        @test isa(d, Normal)
        @test d.σ == σ
        @test d.μ ≈ dot(x, w) / sum(w)
    end
end

@testset "Testing fit for Uniform" begin
    for func in funcs, dist in (Uniform, Uniform{Float64})
        x = func[2](dist(1.2, 5.8), n0)
        d = fit(dist, x)
        @test isa(d, dist)
        @test 1.2 <= minimum(d) <= maximum(d) <= 5.8
        @test minimum(d) == minimum(x)
        @test maximum(d) == maximum(x)

        d = fit(dist, func[2](dist(1.2, 5.8), N))
        @test 1.2 <= minimum(d) <= maximum(d) <= 5.8
        @test isapprox(minimum(d), 1.2, atol=0.02)
        @test isapprox(maximum(d), 5.8, atol=0.02)
    end
end

@testset "Testing fit for Gamma" begin
    for func in funcs, dist in (Gamma, Gamma{Float64})
        x = func[2](dist(3.9, 2.1), n0)
        w = func[1](n0)

        ss = suffstats(dist, x)
        @test isa(ss, Distributions.GammaStats)
        @test ss.sx    ≈ sum(x)
        @test ss.slogx ≈ sum(log.(x))
        @test ss.tw    ≈ n0

        ss = suffstats(dist, x, w)
        @test isa(ss, Distributions.GammaStats)
        @test ss.sx    ≈ dot(x, w)
        @test ss.slogx ≈ dot(log.(x), w)
        @test ss.tw    ≈ sum(w)

        d = fit(dist, func[2](dist(3.9, 2.1), N))
        @test isa(d, dist)
        @test isapprox(shape(d), 3.9, atol=0.1)
        @test isapprox(scale(d), 2.1, atol=0.2)
    end
end

@testset "Testing fit for Geometric" begin
    for func in funcs, dist in (Geometric, Geometric{Float64})
        x = func[2](dist(0.3), n0)
        w = func[1](n0)

        ss = suffstats(dist, x)
        @test isa(ss, Distributions.GeometricStats)
        @test ss.sx ≈ sum(x)
        @test ss.tw ≈ n0

        ss = suffstats(dist, x, w)
        @test isa(ss, Distributions.GeometricStats)
        @test ss.sx ≈ dot(x, w)
        @test ss.tw ≈ sum(w)

        d = fit(dist, x)
        @test isa(d, dist)
        @test succprob(d) ≈ inv(1. + mean(x))

        d = fit(dist, x, w)
        @test isa(d, dist)
        @test succprob(d) ≈ inv(1. + dot(x, w) / sum(w))

        d = fit(dist, func[2](dist(0.3), N))
        @test isa(d, dist)
        @test isapprox(succprob(d), 0.3, atol=0.01)
    end
end

@testset "Testing fit for Laplace" begin
    for func in funcs, dist in (Laplace, Laplace{Float64})
        d = fit(dist, func[2](dist(5.0, 3.0), N + 1))
        @test isa(d, dist)
        @test isapprox(location(d), 5.0, atol=0.02)
        @test isapprox(scale(d)   , 3.0, atol=0.02)
    end
end

@testset "Testing fit for Pareto" begin
    for func in funcs, dist in (Pareto, Pareto{Float64})
        x = func[2](dist(3., 7.), N)
        d = fit(dist, x)

        @test isa(d, dist)
        @test isapprox(shape(d), 3., atol=0.1)
        @test isapprox(scale(d), 7., atol=0.1)
    end
end

@testset "Testing fit for Poisson" begin
    for func in funcs, dist in (Poisson, Poisson{Float64})
        x = func[2](dist(8.2), n0)
        w = func[1](n0)

        ss = suffstats(dist, x)
        @test isa(ss, Distributions.PoissonStats)
        @test ss.sx ≈ sum(x)
        @test ss.tw ≈ n0

        ss = suffstats(dist, x, w)
        @test isa(ss, Distributions.PoissonStats)
        @test ss.sx ≈ dot(x, w)
        @test ss.tw ≈ sum(w)

        d = fit(dist, x)
        @test isa(d, dist)
        @test mean(d) ≈ mean(x)

        d = fit(dist, x, w)
        @test isa(d, dist)
        @test mean(d) ≈ dot(Float64[xx for xx in x], w) / sum(w)

        d = fit(dist, func[2](dist(8.2), N))
        @test isa(d, dist)
        @test isapprox(mean(d), 8.2, atol=0.2)
    end
end

@testset "Testing fit for InverseGaussian" begin
    for func in funcs, dist in (InverseGaussian, InverseGaussian{Float64})
        x = rand(dist(3.9, 2.1), n0)
        w = func[1](n0)

        ss = suffstats(dist, x)
        @test isa(ss, Distributions.InverseGaussianStats)
        @test ss.sx    ≈ sum(x)
        @test ss.sinvx ≈ sum(1 ./ x)
        @test ss.sw    ≈ n0

        ss = suffstats(dist, x, w)
        @test isa(ss, Distributions.InverseGaussianStats)
        @test ss.sx    ≈ dot(x, w)
        @test ss.sinvx ≈ dot(1 ./ x, w)
        @test ss.sw    ≈ sum(w)

        d = fit(dist, rand(dist(3.9, 2.1), N))
        @test isa(d, dist)
        @test isapprox(mean(d), 3.9, atol=0.1)
        @test isapprox(shape(d), 2.1, atol=0.1)

        d = fit_mle(dist, rand(dist(3.9, 2.1), N))
        @test isapprox(mean(d), 3.9, atol=0.1)
        @test isapprox(shape(d), 2.1, atol=0.1)
    end
end

@testset "Testing fit for Rayleigh" begin
    for func in funcs, dist in (Rayleigh, Rayleigh{Float64})
        x = func[2](dist(3.6), N)
        d = fit(dist, x)

        @test isa(d, dist)
        @test isapprox(mode(d), 3.6, atol=0.1)

        # Test automatic differentiation
        f(x) = mean(fit(Rayleigh, x))
        @test all(ForwardDiff.gradient(f, x) .>= 0)
    end
end

@testset "Testing fit for Weibull" begin
    for func in funcs, dist in (Weibull, Weibull{Float64})
        d = fit(dist, func[2](dist(8.1, 4.3), N))
        @test isa(d, dist)
        @test isapprox(d.α, 8.1, atol = 0.1)
        @test isapprox(d.θ, 4.3, atol = 0.1)

    end
end


text:
# Scott P Jones' solution as reinterpreted by Colin Caine

using Random: shuffle!

# Use Int32 to save some bits because this is quite a big vector.
const names = shuffle!(Int32[0:26^2 * 10^3 - 1;])

mutable struct Robot
    # Robots are comfortable with integer names ;)
    id::Int32

    """
        Robot()

    Construct a robot with a unique name within its class.
    """
    Robot() = new(mint_id())
end

"""
    mint_id()

Return a unique robot id. If there are none left, error.
"""
function mint_id()
    isempty(names) ? error("No unique identifiers left!") : pop!(names)
end

"""
    id2name(id)

Convert an integer name to a human-friendly name.
"""
function id2name(id)
    id, c1 = divrem(id, 26)
    id, c2 = divrem(id, 26)
    id, d1 = divrem(id, 10)
    id, d2 = divrem(id, 10)
    d3 = id
    return join((Char('A' + c1),
                 Char('A' + c2),
                 string.((d1, d2, d3))...))
end

"""
    reset!(instance::Robot)

Factory-reset this robot, assigning it a new unique name.
"""
function reset!(instance::Robot)
    instance.id = mint_id()
    instance
end

"""
    name(instance::Robot)

Return this robot's name.
"""
function name(instance::Robot)
    id2name(instance.id)
end

ext: jl


text:
module Utils
using Gadfly

export plot_response

"""
Plot a comparison between the expected and calculated apparent resistivity and
phase response as well as the accompanying model.

Arguments
=========

- `expected::Matrix`: Expected response, ρa in column 1 and Φ in column 2.
- `calculated::Matrix`: Calculated response, ""
- `model::Matrix`: Model used to find response. Depth in column 1, ρ in column 2.
- `fs::Vector`: Frequencies used in forward modelling calculations

Returns
=======

- `::Plot`: Stacked plots of resistivity, phase and the model.
"""
function plot_response(expected::Matrix, calculated::Matrix, model::Matrix, fs::Vector)
    expected_theme = Theme(default_color=colorant"orange")
    lres_expected = layer(x=fs,
                          y=expected[:,1],
                          Geom.point,
                          expected_theme)
    lres_calculated = layer(x=fs,
                            y=calculated[:,1],
                            Geom.line)

    res_plot = plot(lres_expected,
                    lres_calculated,
                    Guide.xlabel("Frequency (Hz)"),
                    Guide.ylabel("ρa (Ωm)"),
                    Scale.x_log10,
                    Scale.y_log10,
                    Coord.cartesian(xflip=true))

    lphase_expected = layer(x=fs,
                            y=expected[:,2],
                            Geom.point,
                            expected_theme)
    lphase_calculated = layer(x=fs,
                              y=calculated[:,2],
                              Geom.line)
    phase_plot = plot(lphase_expected,
                      lphase_calculated,
                      Guide.xlabel("Frequency (Hz)"),
                      Guide.ylabel("Degree (°)"),
                      Scale.x_log10,
                      Scale.y_continuous(minvalue=0, maxvalue=90),
                      Coord.cartesian(xflip=true))

    model_plot = plot(y=[model[:,2]; model[end,2]],
                      x=[model[:,1] / 1000; model[end,1] / 1000 + 0.5],
                      Stat.step(direction=:hv),
                      Geom.line,
                      Guide.ylabel("Resistivity (Ωm)"),
                      Guide.xlabel("Depth (km)"),
                      Guide.title("Model"),
                      Scale.y_log10,
                      Guide.manual_color_key("Legend", ["Expected";"Calculated"], ["orange";"deepskyblue"]))

    return hstack(vstack(res_plot, phase_plot), model_plot)
end

end


text:
using ArgParse
using YAML
using Distributed

@everywhere include("evo/robo.jl")
@everywhere r(i::Individual) = robo_eval(i, 10, length(RoboGrid.ACTIONS))

s = ArgParse.ArgParseSettings()

ArgParse.@add_arg_table(
    s,
    "--seed", arg_type=Int, default=0,
    "--id", arg_type=String, default="test",
    "--log", arg_type=String, default="evolution.log",
    "--cfg", arg_type=String, default="cfg/darwin.yaml",
)

args = ArgParse.parse_args(s)
cfg = YAML.load_file(args["cfg"])
cfg["seed"] = args["seed"]
cfg["n_fitness"] = 1
cfg["nin"] = 10
cfg["nout"] = length(RoboGrid.ACTIONS)

e = Evolution(CGPInd, cfg; id=args["id"], logfile=args["log"])
e.evaluation = r
e.generation = generation
Darwin.run!(e)


text:
"""
This module defines the following data types:
* `AGeom` : acquisition ageometry, i.e., positions of supersources, sources and receivers
* `Src` : source related acquisition parameters, e.g., source wavelet
It also provides methods that either does operations on these data type or
help their construction.
"""
module Acquisition

import GeoPhyInv: Medium, χ
import GeoPhyInv.Utils
using Distributions
using DataFrames
using Random
using LinearAlgebra
using SparseArrays
using CSV
"""
Return some derived fields of `AGeom`

# Arguments

* `acq::Vector{AGeom}` : a vector of `AGeom`
* `attrib::Symbol` : attribute to determine the return object
  * `=:nus` number of unique source positions in acquisition
  * `=:nur` number of unique receiver positions in acquisition
  * `=:uspos` a tuple of x and z positions of all the unique sources
  * `=:urpos` a tuple of x and z position of all the unique receivers
  * `=:ageomurpos` a `AGeom` vector as if all the unique receiver positions are used for each supersource
  * `=:ageomuspos` a `AGeom` vector as if all the unique source positions are used for each supersource
"""
function AGeom_get(acq::Vector{AGeom}, attrib::Symbol)
	nageom = length(acq)
	sposx = AGeom_getvec(acq,:sx);	sposz = AGeom_getvec(acq,:sz)
	isequal(length(sposx), length(sposz)) ?
		spos = [[sposz[is],sposx[is]] for is in 1:length(sposx)] : error("input acq corrupt")

	rposx = AGeom_getvec(acq,:rx);	rposz = AGeom_getvec(acq,:rz)
	isequal(length(rposx), length(rposz)) ?
		rpos = [[rposz[ir],rposx[ir]] for ir in 1:length(rposx)] : error("input acq corrupt")

	uspos = unique(spos); nus=length(uspos)
	urpos = unique(rpos); nur=length(urpos)
	uspost = ([uspos[iu][1] for iu in 1:nus], [uspos[iu][2] for iu in 1:nus]);
        urpost = ([urpos[iu][1] for iu in 1:nur], [urpos[iu][2] for iu in 1:nur]);

	if(attrib == :nus)
		return nus
	elseif(attrib == :nur)
		return nur
	elseif(attrib == :uspos)
		return uspost
	elseif(attrib == :urpos)
		return urpost
	elseif(attrib == :ageomurpos)
		return [AGeom(acq[iageom].sx, acq[iageom].sz,
	       fill(urpost[2],acq[iageom].nss),
	       fill(urpost[1],acq[iageom].nss), acq[iageom].nss, acq[iageom].ns,
	       fill(nur,acq[iageom].nss)) for iageom=1:nageom]
	elseif(attrib == :ageomuspos)
		return [AGeom(fill(uspost[2],acq[iageom].nss),
	       fill(uspost[1],acq[iageom].nss),
	       acq[iageom].rx, acq[iageom].rz,
	       acq[iageom].nss, fill(nus,acq[iageom].nss),
	       acq[iageom].nr) for iageom=1:nageom]
	else
		error("invalid attrib")
	end
end

"""
Given receiver positions `rpos` and `rpos0`.
Returns an array Int indices of the dimension of number of supersources
with `true` at indices, if the waves due to that particular source are
recorded.
"""
function AGeom_find(acq::AGeom; rpos::Array{Float64,1}=nothing, rpos0::Array{Float64,1}=nothing)
	rpos==nothing ? error("need rpos") : nothing
	sson = Array{Vector{Int64}}(undef,acq.nss);
	for iss = 1:acq.nss
		rvec = [[acq.rz[iss][ir],acq.rx[iss][ir]] for ir=1:acq.nr[iss]]
		ir=findfirst(rvec, rpos)
		if(rpos0==nothing)
			ir0=1;
		else
			ir0=findfirst(rvec, rpos0);
		end
		sson[iss] = ((ir != 0) && (ir0 != 0)) ? [ir,ir0] : [0]
	end
	return sson
end

#=
"""
Modify input `AGeom` such that the output `AGeom` has
either sources or receivers on the boundary of
`mgrid`.

# Arguments

* `ageom::AGeom` : input ageometry
* `mgridi` : grid to determine the boundary
* `attrib::Symbol` : decide return
  * `=:srcborder` sources on boundary (useful for back propagation)
  * `=:recborder` receivers on boundary
"""
function AGeom_boundary(ageom::AGeom,
	      mgrid
	      attrib::Symbol
	     )

	if((attrib == :recborder) | (attrib == :srcborder))
		bz, bx, nb = border(mgrid, 3, :inner)
	end
	"change the position of receivers to the boundary"
	if(attrib == :recborder)
		nr = nb; 		sx = ageom.sx;
		sz = ageom.sz;		ns = ageom.ns;
		nss = ageom.nss;
		return AGeom(sx, sz, fill(bx,nss), fill(bz,nss), nss, ns, fill(nr,nss))
	"change the position of source (not supersources) to the boundary for back propagation"
	elseif(attrib == :srcborder)
		ns = nb;		rx = ageom.rx;
		rz = ageom.rz;		nr = ageom.nr;
		nss = ageom.nss;
		return AGeom(fill(bx,nss), fill(bz,nss), rx, rz, nss, fill(ns,nss), nr)
	else
		error("invalid attrib")
	end
end
=#

"""
Check if the input acquisition ageometry is fixed spread.
"""
function AGeom_isfixed(ageom::AGeom)
	isfixed=false
	# can have different source positions, but also different number of sources?
	for field in [:rx, :rz, :nr, :ns]
		f=getfield(ageom, field)
		if(all(f .== f[1:1]))
			isfixed=true
		else
			isfixed=false
			return isfixed
		end
	end
	return isfixed
end

"""
A fixed spread acquisition has same set of sources and
receivers for each supersource.
This method constructs a
fixed spread acquisition ageometry using either a
horizontal or vertical array of supersources/ receivers.
Current implementation has only one source for every supersource.

# Arguments

* `ssmin::Float64` : minimum coordinate for sources
* `ssmax::Float64` : maximum coordinate for sources
* `ss0::Float64` : consant coordinate for sources
* `rmin::Float64` : minimum coordinate for receivers
* `rmax::Float64` : maximum coordinate for receivers
* `r0::Float64` : consant coordinate for receivers
* `nss::Int64` : number of supersources
* `nr::Int64` : number of receivers
* `ssattrib::Symbol=:horizontal` : supersource array kind
  `=:vertical` : vertical array of supersources
  `=:horizontal` horizontal array of supersources
* `rattrib::Symbol=:horizontal` : receiver array kind
  `=:vertical` : vertical array of receivers
  `=:horizontal` horizontal array of receivers
* `rand_flags::Vector{Bool}=[false, false]` : decide placement of supersources and receivers
  `=[true, false]` : randomly place supersources for regularly spaced receivers
  `=[true, true]` : randomly place supersources and receivers
  `=[false, false]` : regularly spaced supersources and receivers
  `=[false, true]` : randomly place receivers for regularly spaced supersources

# Return
* a fixed spread acquisition ageometry `AGeom`
"""
function AGeom_fixed(
	      ssmin::Real,
	      ssmax::Real,
	      ss0::Real,
	      rmin::Real,
	      rmax::Real,
	      r0::Real,
	      nss::Int64,
	      nr::Int64,
	      ssattrib::Symbol=:horizontal,
	      rattrib::Symbol=:horizontal,
 	      rand_flags::Vector{Bool}=[false, false];
	      ssα::Real=0.0,
	      rα::Real=0.0,
	      ns::Vector{Int64}=ones(Int,nss),
	      srad::Real=0.0
	     )

	ssmin=Float64(ssmin); rmin=Float64(rmin)
	ssmax=Float64(ssmax); rmax=Float64(rmax)
	ss0=Float64(ss0); r0=Float64(r0)
	ssα=Float64(ssα)*pi/180.
	rα=Float64(rα)*pi/180.

	ssarray = isequal(nss,1) ? fill(ssmin,1) : (rand_flags[1] ?
					     Random.rand(Uniform(minimum([ssmin,ssmax]),maximum([ssmin,ssmax])),nss) : range(ssmin,stop=ssmax,length=nss))
	if(ssattrib==:horizontal)
		ssz = ss0.+(ssarray.-minimum(ssarray)).*sin(ssα)/cos(ssα); ssx=ssarray
	elseif(ssattrib==:vertical)
		ssx = ss0.+(ssarray.-minimum(ssarray)).*sin(ssα)/cos(ssα); ssz=ssarray
	else
		error("invalid ssattrib")
	end

	rarray = isequal(nr,1) ? fill(rmin,1) : (rand_flags[2] ?
					  Random.rand(Uniform(minimum([rmin,rmax]),maximum([rmin,rmax])),nr) : range(rmin,stop=rmax,length=nr))
	if(rattrib==:horizontal)
		rz = r0.+(rarray.-minimum(rarray)).*sin(rα)/cos(rα); rx = rarray
	elseif(rattrib==:vertical)
		rx = r0.+(rarray.-minimum(rarray)).*sin(rα)/cos(rα); rz = rarray
	else
		error("invalid rattrib")
	end
	rxall = [rx for iss=1:nss];
	rzall = [rz for iss=1:nss];
	ssxall=[zeros(ns[iss]) for iss=1:nss];
	sszall=[zeros(ns[iss]) for iss=1:nss];
	for iss in 1:nss
		for is in 1:ns[iss]
			θ=Random.rand(Uniform(-Float64(pi),Float64(pi)))
			r = iszero(srad) ? 0.0 : Random.rand(Uniform(0,srad))
			x=r*cos(θ)
			z=r*sin(θ)
			ssxall[iss][is]=ssx[iss]+x
			sszall[iss][is]=ssz[iss]+z
		end
	end
	return AGeom(ssxall, sszall, rxall, rzall, nss, ns, fill(nr,nss))
end

"""
Circular acquisition. The sources and receivers can also be placed on a circle of radius
`rad`. The origin of the circle is at `loc`.
This ageometry is unrealistic, but useful for testing.
Receivers are placed such that the limits
of the angular offset are given by `θlim`

# Arguments

* `nss::Int64=10` : number of supersources
* `nr::Int64=10` : number receivers for each super source
* `loc::Vector{Float64}=[0.,0.]` : location of origin
* `rad::Vector{Float64}=[100.,100.]` : radius for source and receiver circles, for example,
  * `=[0.,100.]` for sources at the center of circle
* `θlim::Vector{Float64}=[0.,π]` : acquisition is limited to these angular offsets between 0 and π

# Return

* a circular acquisition ageometry `AGeom`
"""
function AGeom_circ(;
		   nss::Int64=10,
		   nr::Int64=10,
		   loc::Vector=[0.,0.],
		   rad::Vector=[100.,100.],
		   θlim::Vector=[0.,π]
	       )

	# modify nr such that approximately nr receivers are present in θlim
	nra = nr * convert(Int64, floor(π/(abs(diff(θlim)[1]))))

	(nra==0) && error("θlim should be from 0 to π")

	sxx, szz = circ_coord(loc..., nss, rad[1])
	rxxa, rzza = circ_coord(loc..., nra, rad[2])
	sx=Array{Array{Float64}}(undef,nss)
	sz=Array{Array{Float64}}(undef,nss)
	rx=Array{Array{Float64}}(undef,nss)
	rz=Array{Array{Float64}}(undef,nss)
	ns=Array{Int64}(undef,nss)
	nr=Array{Int64}(undef,nss)

	for iss=1:nss
		rangles = [angle(complex(rxxa[ira]-loc[2], rzza[ira]-loc[1])) for ira in 1:nra]
		sangles = angle(complex(sxx[iss]-loc[2], szz[iss]-loc[1]))
		diff = abs.(rangles.-sangles)
		diff = [min(diff[ira], (2π - diff[ira])) for ira in 1:nra]
		angles = findall(minimum(θlim) .<= diff .<= maximum(θlim))
		sx[iss] = [sxx[iss]]
		sz[iss] = [szz[iss]]
		ns[iss] = 1
		nr[iss] = length(angles)
		rx[iss] = rxxa[angles]
		rz[iss] = rzza[angles]
	end
	ageom =  AGeom(sx, sz, rx, rz, nss, ns, nr)

	print(ageom, "circular acquisition ageometry")
	return ageom
end

function circ_coord(z, x, n, rad)
	θ = (n==1) ? [0.0] : collect(range(0, stop=2π, length=n+1))
	xx = (rad .* cos.(θ) .+ x)[1:n]
	zz = (rad .* sin.(θ) .+ z)[1:n]
	return xx, zz
end

end # module


text:
using Revise
using Pkg
using WAdag
using Test
#=
@testset "WAdag.jl" begin
    @test cleanaddress(;wb="us-500-input.xlsx",ws="Sheet1",rng="A1:B501",teststate=true) == DataFrame(XLSX.readtable("us-500-input.xlsx","Sheet1")...)
end
=#

@testset "basicfunc.jl" begin
    @test bfunc(5) == 25
end


text:
# julia --project=. /scripts/nsys.jl
# nsys launch julia /path/to/nsys.jl
using GPUAcceleratedTracking, CUDA, Tracking, GNSSSignals, StructArrays
import Tracking: Hz, ms
system = GPSL1(use_gpu = Val(true))
enable_gpu = Val(true)
num_samples = 2500
num_ants = 1
num_correlators = 3
system = GPSL1(use_gpu = enable_gpu)
codes = system.codes
#convert to text_mem
codes = CuTexture(
    CuTextureArray(codes),
    address_mode = CUDA.ADDRESS_MODE_WRAP,
    interpolation = CUDA.NearestNeighbour(),
    normalized_coordinates = true
)
code_frequency = get_code_frequency(system)
code_length = get_code_length(system)
start_code_phase = 0.0f0
carrier_phase = 0.0f0
carrier_frequency = 1500Hz
prn = 1
signal, sampling_frequency = gen_signal(system, prn, carrier_frequency, num_samples, num_ants = NumAnts(num_ants), start_code_phase = start_code_phase, start_carrier_phase = carrier_phase)
correlator = EarlyPromptLateCorrelator(NumAnts(num_ants), NumAccumulators(num_correlators))
correlator_sample_shifts = get_correlator_sample_shifts(system, correlator, sampling_frequency, 0.5)
num_of_shifts = correlator_sample_shifts[end] - correlator_sample_shifts[1]
code_replica = CUDA.zeros(Float32, num_samples + num_of_shifts)
carrier_replica = StructArray{ComplexF32}((CUDA.zeros(Float32, num_samples), CUDA.zeros(Float32, num_samples)))
downconverted_signal = StructArray{ComplexF32}((CUDA.zeros(Float32, num_samples, num_ants), CUDA.zeros(Float32, num_samples, num_ants)))
threads_per_block = zeros(Int, 3)
blocks_per_grid = zeros(Int, 3)
accum = StructArray{ComplexF32}((CUDA.zeros(Float32, (num_samples, num_ants, length(correlator_sample_shifts))),CUDA.zeros(Float32, (num_samples, num_ants, length(correlator_sample_shifts)))))
threads_per_block[3] = 1024
blocks_per_grid[3] = cld(num_samples, threads_per_block[3])
phi = StructArray{ComplexF32}((CUDA.zeros(Float32, (blocks_per_grid[3], num_ants, length(correlator_sample_shifts))),CUDA.zeros(Float32, (blocks_per_grid[3], num_ants, length(correlator_sample_shifts)))))
shmem_size = sizeof(ComplexF32) * 1024 * num_correlators * num_ants
code_replica_kernel = @cuda launch=false gen_code_replica_texture_mem_strided_kernel!(
        code_replica,
        codes,
        code_frequency,
        sampling_frequency,
        start_code_phase,
        prn,
        num_samples,
        num_of_shifts,
        code_length
    )
blocks_per_grid[1], threads_per_block[1] = launch_configuration(code_replica_kernel.fun)
downconvert_and_accumulate_kernel = @cuda launch=false downconvert_and_accumulate_strided_kernel!(
    accum.re,
    accum.im,
    code_replica,
    carrier_replica.re,
    carrier_replica.im,
    downconverted_signal.re,
    downconverted_signal.im,
    signal.re,
    signal.im,
    carrier_frequency,
    sampling_frequency,
    carrier_phase,
    num_samples,
    NumAnts(num_ants),
    correlator_sample_shifts
)
blocks_per_grid[2], threads_per_block[2] = launch_configuration(downconvert_and_accumulate_kernel.fun)
algorithm = KernelAlgorithm(2431)
kernel_algorithm(
    threads_per_block,
    blocks_per_grid,
    shmem_size,
    code_replica,
    codes,
    code_frequency,
    sampling_frequency,
    start_code_phase,
    prn,
    num_samples,
    num_of_shifts,
    code_length,
    accum.re,
    accum.im,
    phi.re,
    phi.im,
    carrier_replica.re,
    carrier_replica.im,
    downconverted_signal.re,
    downconverted_signal.im,
    signal.re,
    signal.im,
    correlator_sample_shifts,
    carrier_frequency,
    carrier_phase,
    NumAnts(num_ants),
    nothing,
    algorithm
)
CUDA.@profile NVTX.@range "kernel_algorithm" begin
    kernel_algorithm(
    threads_per_block,
    blocks_per_grid,
    shmem_size,
    code_replica,
    codes,
    code_frequency,
    sampling_frequency,
    start_code_phase,
    prn,
    num_samples,
    num_of_shifts,
    code_length,
    accum.re,
    accum.im,
    phi.re,
    phi.im,
    carrier_replica.re,
    carrier_replica.im,
    downconverted_signal.re,
    downconverted_signal.im,
    signal.re,
    signal.im,
    correlator_sample_shifts,
    carrier_frequency,
    carrier_phase,
    NumAnts(num_ants),
    nothing,
    algorithm
)
end

max_stars_repo_name: mfkiwl/GPUAcceleratedTracking

max_issues_repo_name: mfkiwl/GPUAcceleratedTracking

max_forks_repo_name: mfkiwl/GPUAcceleratedTracking


text:
module UtilitiesTest

export test_distance_between

using ProtoBuf
using Test

using dataset_pb2

using SmuUtilities
#using TopologyFromGeometry

function test_bohr_angstroms()
  @test isapprox(angstroms_to_bohr(SmuUtilities.bohr_to_angstroms(1.0)), 1.0)
end

function test_add_atom()
  bond_topology = BondTopology()
  add_atom_by_atomic_number!(bond_topology, 6)
  @test length(bond_topology.atoms) == 1
  @test bond_topology.atoms[1] == BondTopology_AtomType.ATOM_C

  add_atom_by_atomic_number!(bond_topology, 7)
  @test length(bond_topology.atoms) == 2
  @test bond_topology.atoms[2] == BondTopology_AtomType.ATOM_N

  add_atom_by_atomic_number!(bond_topology, 7, 1)
  @test length(bond_topology.atoms) == 3
  @test bond_topology.atoms[3] == BondTopology_AtomType.ATOM_NPOS

  # Do it systematically
  atype = [[1, 0], [6, 0], [7, 0], [7, 1], [8, 0], [8, 1], [9, 0]]
  expected = [
    BondTopology_AtomType.ATOM_H,
    BondTopology_AtomType.ATOM_C,
    BondTopology_AtomType.ATOM_N,
    BondTopology_AtomType.ATOM_NPOS,
    BondTopology_AtomType.ATOM_O,
    BondTopology_AtomType.ATOM_ONEG,
    BondTopology_AtomType.ATOM_F,
  ]
  @test length(atype) == length(expected)

  bt = BondTopology()
  for i in zip(atype, expected)
    add_atom_by_atomic_number!(bt, i[1][1], i[1][2])
    @test bt.atoms[end] == i[2]
  end
  @test length(bt.atoms) == length(expected)
  true
end

function test_add_single_bond()
  bond_topology = BondTopology()
  atoms = Vector{Int32}()
  push!(atoms, BondTopology_AtomType.ATOM_C)
  push!(atoms, BondTopology_AtomType.ATOM_C)
  setproperty!(bond_topology, :atoms, atoms)

  add_bond!(0, 1, 1, bond_topology)
  @test length(bond_topology.bonds) == 1
  @test bond_topology.bonds[1].bond_type == BondTopology_BondType.BOND_SINGLE
  true
end

function test_add_double_bond()
  bond_topology = BondTopology()
  atoms = Vector{Int32}()
  push!(atoms, BondTopology_AtomType.ATOM_C)
  push!(atoms, BondTopology_AtomType.ATOM_C)
  setproperty!(bond_topology, :atoms, atoms)

  add_bond!(0, 1, 2, bond_topology)
  @test length(bond_topology.bonds) == 1
  @test bond_topology.bonds[1].bond_type == BondTopology_BondType.BOND_DOUBLE
  true
end

function test_add_triple_bond()
  bond_topology = BondTopology()
  atoms = Vector{Int32}()
  push!(atoms, BondTopology_AtomType.ATOM_C)
  push!(atoms, BondTopology_AtomType.ATOM_C)
  setproperty!(bond_topology, :atoms, atoms)

  add_bond!(0, 1, 3, bond_topology)
  @test length(bond_topology.bonds) == 1
  @test bond_topology.bonds[1].bond_type == BondTopology_BondType.BOND_TRIPLE
  true
end

function test_add_multiple_bonds()
  bond_topology = BondTopology()
  atoms = Vector{Int32}()
  push!(atoms, BondTopology_AtomType.ATOM_C)
  push!(atoms, BondTopology_AtomType.ATOM_C)
  push!(atoms, BondTopology_AtomType.ATOM_C)
  setproperty!(bond_topology, :atoms, atoms)

  add_bond!(0, 1, 1, bond_topology)
  add_bond!(1, 2, 1, bond_topology)
  @test length(bond_topology.bonds) == 2
  @test bond_topology.bonds[1].bond_type == BondTopology_BondType.BOND_SINGLE
  @test bond_topology.bonds[2].bond_type == BondTopology_BondType.BOND_SINGLE
  true
end

function test_distance_between()
  geom = Geometry()
  atoms = Vector{Geometry_AtomPos}()
  a1 = Geometry_AtomPos(x=0.0, y=0.0, z=0.0)
  push!(atoms, a1)
  a2 = Geometry_AtomPos(x=0.0, y=0.0, z=angstroms_to_bohr(1.0))
  push!(atoms, a2)
  setproperty!(geom, :atom_positions, atoms)
  @test isapprox(distance_between_atoms(geom, 1, 2), 1.0)
  setproperty!(geom.atom_positions[2], :x, angstroms_to_bohr(1.0))
  setproperty!(geom.atom_positions[2], :y, angstroms_to_bohr(1.0))
  setproperty!(geom.atom_positions[2], :z, angstroms_to_bohr(1.0))
  @test isapprox(distance_between_atoms(geom, 1, 2), sqrt(3.0))

  distances = SmuUtilities.distances(geom)
  @test distances[1,1] == 0.0
  @test distances[2,2] == 0.0
  @test isapprox(distances[1,2], sqrt(3.0))
  @test isapprox(distances[2,1], sqrt(3.0))
  true
end

function test_bonded()
  bond_topology = BondTopology()
  atoms = Vector{Int32}()
  bonds = Vector{BondTopology_Bond}()
  push!(atoms, BondTopology_AtomType.ATOM_C)
  push!(atoms, BondTopology_AtomType.ATOM_C)
  push!(atoms, BondTopology_AtomType.ATOM_C)
  push!(bonds, BondTopology_Bond(atom_a=0, atom_b=1, bond_type=BondTopology_BondType.BOND_SINGLE))
  push!(bonds, BondTopology_Bond(atom_a=1, atom_b=2, bond_type=BondTopology_BondType.BOND_SINGLE))
  setproperty!(bond_topology, :atoms, atoms)
  setproperty!(bond_topology, :bonds, bonds)

  connection_matrix = SmuUtilities.bonded(bond_topology)
  @test length(atoms) == size(connection_matrix, 1)
  @test connection_matrix[1,2] > 0
  @test connection_matrix[2,3] > 0
  @test count(!iszero, connection_matrix) == 4  # Bonds are placed twice.
  push!(bond_topology.bonds, BondTopology_Bond(atom_a=0, atom_b=2, bond_type=BondTopology_BondType.BOND_DOUBLE))
  connection_matrix = SmuUtilities.bonded(bond_topology)
  @test count(!iszero, connection_matrix) == 6  # Bonds are placed twice.
  @test connection_matrix[1,2] == connection_matrix[2,3]
  @test connection_matrix[1,2] != connection_matrix[3,1]
  true
end

function test_connections()
  bond_topology = BondTopology()
  atoms = Vector{Int32}()
  bonds = Vector{BondTopology_Bond}()
  push!(atoms, BondTopology_AtomType.ATOM_C)
  push!(atoms, BondTopology_AtomType.ATOM_C)
  push!(atoms, BondTopology_AtomType.ATOM_C)
  push!(bonds, BondTopology_Bond(atom_a=0, atom_b=1, bond_type=BondTopology_BondType.BOND_SINGLE))
  push!(bonds, BondTopology_Bond(atom_a=1, atom_b=2, bond_type=BondTopology_BondType.BOND_SINGLE))
  setproperty!(bond_topology, :atoms, atoms)
  setproperty!(bond_topology, :bonds, bonds)

  connections = SmuUtilities.connections(SmuUtilities.bonded(bond_topology))
  @test length(connections) == length(atoms)
  @test length(connections[1]) == 1
  @test length(connections[2]) == 2
  @test length(connections[3]) == 1
  @test connections[1][1] == 2
  @test connections[2][1] == 1
  @test connections[2][2] == 3
  @test connections[3][1] == 2
  true
end

function test_canonical_bond_topology()
  bond_topology = BondTopology()
  atoms = Vector{Int32}()
  bonds = Vector{BondTopology_Bond}()
  push!(atoms, BondTopology_AtomType.ATOM_C)
  push!(atoms, BondTopology_AtomType.ATOM_C)
  push!(atoms, BondTopology_AtomType.ATOM_C)
  push!(bonds, BondTopology_Bond(atom_a=2, atom_b=1, bond_type=BondTopology_BondType.BOND_SINGLE))
  push!(bonds, BondTopology_Bond(atom_a=1, atom_b=0, bond_type=BondTopology_BondType.BOND_SINGLE))
  setproperty!(bond_topology, :atoms, atoms)
  setproperty!(bond_topology, :bonds, bonds)

  SmuUtilities.canonical_bond_topology!(bond_topology)

  expected = BondTopology()
  atoms = Vector{Int32}()
  bonds = Vector{BondTopology_Bond}()
  push!(atoms, BondTopology_AtomType.ATOM_C)
  push!(atoms, BondTopology_AtomType.ATOM_C)
  push!(atoms, BondTopology_AtomType.ATOM_C)
  push!(bonds, BondTopology_Bond(atom_a=0, atom_b=1, bond_type=BondTopology_BondType.BOND_SINGLE))
  push!(bonds, BondTopology_Bond(atom_a=1, atom_b=2, bond_type=BondTopology_BondType.BOND_SINGLE))
  setproperty!(expected, :atoms, atoms)
  setproperty!(expected, :bonds, bonds)
  @test bond_topology == expected
  @test SmuUtilities.same_bond_topology(bond_topology, expected)

  true
end

function test_single_fragment1()
  bond_topology = BondTopology()
  atoms = Vector{Int32}()
  bonds = Vector{BondTopology_Bond}()
  push!(atoms, BondTopology_AtomType.ATOM_C)  # atom 0
  push!(atoms, BondTopology_AtomType.ATOM_C)  # atom 1
  push!(atoms, BondTopology_AtomType.ATOM_C)  # atom 2
  push!(bonds, BondTopology_Bond(atom_a=0, atom_b=1, bond_type=BondTopology_BondType.BOND_SINGLE))
  setproperty!(bond_topology, :atoms, atoms)
  setproperty!(bond_topology, :bonds, bonds)

  @test !SmuUtilities.is_single_fragment(bond_topology)

  push!(bond_topology.bonds, BondTopology_Bond(atom_a=1, atom_b=2, bond_type=BondTopology_BondType.BOND_SINGLE))
  @test SmuUtilities.is_single_fragment(bond_topology)

  push!(bond_topology.atoms, BondTopology_AtomType.ATOM_C)  # atom 3
  @test !SmuUtilities.is_single_fragment(bond_topology)

  push!(bond_topology.atoms, BondTopology_AtomType.ATOM_C)  # atom 4
  @test !SmuUtilities.is_single_fragment(bond_topology)
  push!(bond_topology.bonds, BondTopology_Bond(atom_a=3, atom_b=4, bond_type=BondTopology_BondType.BOND_SINGLE))
  @test !SmuUtilities.is_single_fragment(bond_topology)

  push!(bond_topology.bonds, BondTopology_Bond(atom_a=2, atom_b=3, bond_type=BondTopology_BondType.BOND_SINGLE))
  @test SmuUtilities.is_single_fragment(bond_topology)
  push!(bond_topology.bonds, BondTopology_Bond(atom_a=0, atom_b=4, bond_type=BondTopology_BondType.BOND_SINGLE))
  @test SmuUtilities.is_single_fragment(bond_topology)

  true
end

function all_tests()
  @test test_add_atom()
  @test test_add_single_bond()
  @test test_add_double_bond()
  @test test_add_triple_bond()
  @test test_distance_between()
  @test test_bonded()
  @test test_connections()
  @test test_canonical_bond_topology()
  @test test_single_fragment1()
  print("All tests complete\n")
end

end  # module UtilitiesTest

max_stars_repo_name: IanAWatson/smu

max_issues_repo_name: IanAWatson/smu

max_forks_repo_name: IanAWatson/smu


text:
using Graphene
using GeometricalPredicates: Point2D, Line2D, Line
using Test

export test_xy
export test_indexed_atoms
export test_atom_group_collection

atom1 = Point2D(4., 6.)
atom2 = Point2D(1., 2.)
atom3 = Point2D(7., 2.)
atom4 = Point2D(8., 3.)
bond12 = Line(atom1, atom2)
@test isa(bond12, Line2D)

test_indexed_atoms = pairs(IndexLinear(), [atom1, atom2, atom3, atom4])
test_atom_group = [1, 2, 3, 4]
@test make_bonds(test_atom_group) == [(1,2), (1,3), (1,4)]
test_atom_group = [1, 2, 3]
@test make_bonds(test_atom_group) == [(1,2), (1,3)]
test_atom_group = [1, 2]
@test make_bonds(test_atom_group) == [(1,2)]

test_xy = [ 0.   5.    5.  15.   15.  20.  20.  30.  30.  35.  45.   55.   55.  64.   100.;
           17.4  8.7  26.1  8.7  26.1  0.  17.4  0.  17.4  8.7  8.7   8.7  18.7 10.4  100.]
test_indexed_atoms = make_atoms(test_xy)

test_atom_group = [4, 2, 6, 7]
@test make_paths(test_atom_group, test_indexed_atoms) == [(7, 4, 6), (6, 4, 2), (2, 4, 7)]
test_atom_group = [12, 11, 13, 14]
@test make_paths(test_atom_group, test_indexed_atoms) == [(11, 12, 13), (13, 12, 14), (14, 12, 11)]
test_atom_group = [11, 10, 12]
@test bond2path(test_atom_group, test_indexed_atoms) == ([(11, 10), (11, 12)], [(10, 11, 12), (12, 11, 10)])
test_atom_group = [15]
# @test bond2path(test_atom_group, test_indexed_atoms) == (nothing, nothing)
@test bond2path(test_atom_group, test_indexed_atoms) == ([], [])

test_atom_group_collection = [[1, 2, 3],
                              [2, 1, 4],
                              [3, 1, 5],
                              [4, 2, 6, 7],
                              [5, 3, 7],
                              [6, 4, 8],
                              [7, 4, 5, 9],
                              [8, 6, 10],
                              [9, 7, 10],
                              [10, 8, 9, 11],
                              [11, 10, 12],
                              [12, 11, 13, 14],
                              [13, 12],
                              [14, 12]]

@test collect_bonds_paths(test_atom_group_collection[9:14], test_indexed_atoms) ==
    ([(9, 7), (9, 10), (10, 8), (10, 9), (10, 11), (11, 10), (11, 12), (12, 11),
      (12, 13), (12, 14), (13, 12), (14, 12)],
     [(7, 9, 10), (10, 9, 7), (11, 10, 8), (8, 10, 9), (9, 10, 11),
     (10, 11, 12), (12, 11, 10), (11, 12, 13), (13, 12, 14), (14, 12, 11),
     (12, 13, 12), (12, 14, 12)])


text:
using StructuredProximalOperators
using Random
using LinearAlgebra
using FiniteDifferences
using ForwardDiff

function main()
    m, n = 5, 6
    k = 4
    α = 2.0
    g = regularizer_lnuclear(1.0)

    M = FixedRankMatrices(m, n, k, ℝ)
    Random.seed!(1234)
    x = randomMPoint(M)
    Random.seed!(1456)
    η = randomTVector(M, x)
    x_emb = embed(M, x)
    η_emb = embed(M, x, η)

    y, N = prox_αg(g, x, α)
    y_emb = embed(M, y)
    @show M
    @show N

    # Input curve, image by prox.
    c(t) = retract(M, x, t * η)

    ## Check the prox image lies in the same manifold for considered t range
    # for t in 10 .^ range(-8, stop = 0, length = 20)
    #     y_t, M_t = prox_αg(g, c(t), α)
    #     @show M_t
    # end

    # 1. Approximate DT(prox(x))[η] by finite differences:
    T(x) = prox_αg(g, x, α)[1]
    DT_x = FiniteDifferences.jacobian(central_fdm(5, 1), T, x_emb)[1]
    function DT_x_at(η)
        η_vec, vectomat = to_vec(η)
        return vectomat(DT_x * η_vec)
    end

    DT_x_η_emb = DT_x_at(η_emb)
    DT_x_η = zero_tangent_vector(N, y)
    project!(N, DT_x_η, y, DT_x_η_emb)
    println("dist(DT(x)[η],  T_y M):\t\t", norm(DT_x_η_emb - embed(N, y, DT_x_η)))

    finitediff_prox(t) = norm(T(x_emb + t * η_emb) - y_emb - t * DT_x_η_emb)

    # 2. Approximate projection derivative by finite differences:
    function proj(x, v)
        return project(N, x, v)
    end
    proj_v(x) = proj(x, y_emb - x_emb)

    println("")
    println("")

    Dproj_Tc = FiniteDifferences.jacobian(central_fdm(4, 1), t -> proj_v(T(c(t))), 0.0)[1]
    _, vectomat = to_vec(rand(5, 6))
    Dproj_Tc = vectomat(Dproj_Tc)

    finitediff_proj(t) = norm(proj_v(T(c(t))) - proj_v(T(c(0))) - t * Dproj_Tc)

    # @show proj(y, y_emb-x_emb)

    # 2. Check DT(prox(x))[η] theoretical formula:
    Hess_DT = zero_tangent_vector(N, y)
    ∇²M_g_ξ!(g, N, Hess_DT, y, DT_x_η)
    Hess_DT_emb = embed(M, y, Hess_DT)

    η_proj = project(N, y, η_emb)
    η_proj_emb = project(N, y, η_proj)
    @show norm(DT_x_η_emb + α * Hess_DT_emb + Dproj_Tc - project(N, y, η_emb))
    @show norm(DT_x_η_emb + α * Hess_DT_emb + Dproj_Tc)

    comparison = compare_curves(finitediff_prox, finitediff_proj)
    return display_curvescomparison(comparison)
end

main()


text:
# Copyright (c) 2018: Matthew Wilhelm & Matthew Stuber.
# This code is licensed under MIT license (see LICENSE.md for full details)
#############################################################################
# EAGO
# A development environment for robust and global optimization
# See https://github.com/PSORLab/EAGO.jl
#############################################################################
# src/eago_optimizer/functions/nonlinear/cone.jl
# Defines buffered structures to store quadratic functions:
# BufferedSOC as well as the  lower_interval_bound, interval_bound, and
# eliminate_fixed_variables! functions associated with each structure.
#############################################################################

###
### Structure definitions
###
"""
$(TYPEDEF)

Stores a second-order cone with a buffer.
"""
mutable struct BufferedSOC <: AbstractEAGOConstraint
    variables::VECOFVAR
    buffer::Dict{Int, Float64}
    saf::SAF
    len::Int
end

###
### Constructor definitions
###
function BufferedSOC(func::VECOFVAR, set::SOC)
    len = length(func.variables)
    buffer = Dict{Int, Float64}([(variable.value, 0.0) for variable in func.variables])
    saf = SAF(fill(SAT(0.0, VI(1)), len), 0.0)
    return BufferedSOC(copy(func), buffer, saf, len)
end

###
### Parsing definitions
###


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

module TOML
    using Dates

    const TYPE = Union{AbstractDict,AbstractArray,AbstractString,DateTime}

    include("parser.jl")
    include("print.jl")

    "Convert `TOML.Table` to `Dict{String,Any}`"
    function table2dict(tbl::Union{Nothing,Table})
        tbl == nothing && return Dict{String,Any}()
        return table2dict(get(tbl))
    end

    function table2dict(tbl::Table)
        ret = Dict{String,Any}()
        for (k,v) in tbl.values
            if isa(v, Table)
                ret[k] = table2dict(v)
            elseif isa(v, Array) && length(v)>0 && isa(v[1], Table)
                ret[k] = [table2dict(e) for e in v]
            else
                ret[k] = v
            end
        end
        return ret
    end

    "Parse IO input and return result as dictionary."
    function parse(io::IO)
        parser = Parser(io)
        res = parse(parser)
        length(parser.errors)>0 && throw(CompositeException(parser.errors))
        return table2dict(res)
    end

    "Parse string"
    function parse(str::AbstractString)
        io = IOBuffer(str)
        res = parse(io)
        close(io)
        return res
    end

    "Parse file"
    parsefile(filename::AbstractString) = open(parse, filename, "r")

end


text:
using Test, NiLang, NiLang.AD

include("vars.jl")
include("gradfunc.jl")

include("instructs.jl")
include("ulog.jl")
include("complex.jl")
include("manual.jl")
include("jacobian.jl")
include("hessian_backback.jl")


text:
include("StructEquality.jl")
include("ImportACounter.jl")

b_counter = CounterX(1)

println(a_counter.id)
println(b_counter.id)
println(a_counter == b_counter) # False

max_stars_repo_name: Sohva/oolong

max_issues_repo_name: Sohva/oolong

max_forks_repo_name: oolongtensor/oolong


text:
"""
    subtreeMutation(genotype::Union{CGPGenotype, STGPGenotype},
                    gpExperimentInfo::Union{CGPInfo, STGPInfo},
                    rng::Random.AbstractRNG,
                    method::Function = fullGenerator)

Performs a subtree mutation for the given genotype. The subtree mutation performs
a crossover operation between the given genotype and a random generated one.

`Self-provided Arguments` are provided by the library, so only `User Arguments` must be provided.

# Self-provided Arguments
- `genotype::Union{CGPGenotype, STGPGenotype}`: genotype of the individual that is
    going to be mutated.
- `gpExperimentInfo::Union{CGPInfo, STGPInfo}`: information about the GP experiment.
- `rng::Random.AbstractRNG`: random number generator for random number consistency
    along an experiment.
- `method::Function = fullGenerator`: generation method for the random generated
    tree that is going to be crossed with the given one.

# User arguments
None.

# Returns
The mutated genotype.
"""
function subtreeMutation(genotype::Union{CGPGenotype, STGPGenotype},
                         gpExperimentInfo::Union{CGPInfo, STGPInfo},
                         rng::Random.AbstractRNG,
                         method::Function = fullGenerator)

    randomGenotype = method(gpExperimentInfo, rng)
    mutatedGenotype = oneChildSubtreeCross(genotype, randomGenotype, gpExperimentInfo, rng)[1]

    return mutatedGenotype
end # function

max_stars_repo_name: SergioGmezM/GeneticJulia

max_issues_repo_name: SergioGmezM/GeneticJulia

max_forks_repo_name: SergioGmezM/GeneticJulia


text:
function hash_struct(struc, h::UInt, fields=fieldnames(typeof(struc)))
    for f in fields
        h = Base.hash(getfield(struc,f), h)
    end
    return h
end

const EarthTime =  Dict("year" =>   365,    # days per year
                        "day"  =>   24 ,    # hours per day
                        "hour" =>   60 ,    # minutes per hour
                        "minute" => 60 ,    # seconds per minute
)

const KerbalTime = Dict("year" =>   424,    # days per year
                        "day"  =>   6,      # hours per day
                        "hour" =>   60,     # minutes per hour
                        "minute" => 60,     # seconds per minute
)

function seconds_to_datetime(nseconds, timedefs=KerbalTime)
    year, day, hour, minute = timedefs["year"], timedefs["day"], timedefs["hour"], timedefs["minute"]

    y = floor(nseconds/(minute*hour*day*year))+1       # starting from year 1
    nseconds -= minute*hour*day*year*(y-1)
    d = floor(nseconds/(minute*hour*day))+1            # starting from day 1
    nseconds -= minute*hour*day*(d-1)
    h = floor(nseconds/(minute*hour))                  # starting from hour 0
    nseconds -= minute*hour*(h)
    m = floor(nseconds/minute)                         # starting from minute 0
    nseconds -= minute*(m)
    s = round(nseconds)
    return NTuple{5,Int}((y,d,h,m,s))
end

function seconds_to_datetime_string(nseconds, timedefs=KerbalTime)
    ny, nd, nh, nm, ns = seconds_to_datetime(nseconds, timedefs)
    return "Year $ny, Day $nd, $(lpad(nh, 2, '0')):$(lpad(nm, 2, '0')):$(lpad(ns, 2, '0'))"
end

function wrap_angle(θ::Real, min::Real = 0)
    if min <= θ < (min + 2π)
        return θ
    else
        return θ - 2π*floor((θ-min)/2π)
    end
end

function angledist(θ₁, θ₂)
    dist = wrap_angle(θ₁ - θ₂)
    return dist > π ? dist - 2π : dist
end

function isapproxangle(θ₁,θ₂; atol=1e-12, rtol=1e-12, kwargs...)
    return isapprox(angledist(θ₁, θ₂),0;atol=atol,rtol=rtol,kwargs...) || isapprox(rem,1;atol=atol,rtol=rtol,kwargs...)
end

function wrap_acos(x::Real)
    if abs(x) > 1
        x = copysign(1,x)
    end
    return acos(x)
end

function wrap_asin(x::Real)
    if abs(x) > 1
        x = copysign(1,x)
    end
    return asin(x)
end


text:
using QuantEcon
using Statistics
using LinearAlgebra

include(string(pwd(),"/src/gensys.jl"))

function log_like_dsge(par,data;kalman_tol = 1e-10)
    #order to par
    #alfa
    #beta
    #epsilon
    #theta
    #sig
    #sigma: this is the std dev of the innovation
    #phi
    #phi_pi
    #phi_y
    #rho_v

    #data will have t x p dimension: lines are periods p are variables

    alfa = par[1]
    bet = par[2]
    epsilon = par[3]
    theta = par[4]
    sig = par[5]
    #par 6 is coded bellow see line 56
    phi = par[7]
    phi_pi = par[8]
    phi_y = par[9]
    rho_v = par[10]

    THETA = (1-alfa)/(1-alfa+alfa*epsilon)
    lamb = (1-theta)*(1-bet*theta)/theta*THETA
    kappa = lamb*(sig+(phi+alfa)/(1-alfa))

    nobs = size(data,1)

    GAMMA_0 = [bet  0    0  0;
               1    sig  0  0;
               0    0    0  0;
               0    0    0  1]

    GAMMA_1 = [ 1       -kappa  0   0;
                0        sig    1   0;
               -phi_pi  -phi_y  1  -1;
                0        0      0   rho_v]

    PSI = [0; 0; 0; 1]

    PI = [bet  0;
          1    sig;
          0    0;
          0    0]

    p = size(GAMMA_1,1) #number of endogenous vars

    sol = gensys(GAMMA_0,GAMMA_1,PSI,PI; verbose = false)
    if sum(sol.eu) != 2
        return -99999999999
    end

    #Sig = zeros(p,p)
    #Sig[4,4] = par[6]

    G = zeros(1,p)
    G[1,2] = 1

    A = sol.Theta1
    R = [0] .+ 1e-8
    Q = par[6]^2*sol.Theta2*sol.Theta2'

    kalman_res = Kalman(A,G,Q,R) #create a Kalman filter instance

    y_mean = mean(data,dims=1)
    y_var = var(data,dims=1)
    y_var = reshape(y_var,size(y_var,2))
    #y_var = repeat([1],p)

    x_hat = repeat(y_mean,p) #initial mean of the state
    x_var = diagm(repeat(y_var,p))#variance initial of state
    #x_var = x_var*x_var'
    set_state!(kalman_res,x_hat,x_var)

    fit = zeros(nobs,4)

    llh = zeros(nobs)

    for j in 1:(nobs-1)
        med = kalman_res.cur_x_hat
        fit[j,:] = med
        varian = kalman_res.cur_sigma
        #println(det(varian))
        eta = data[j+1,:] - kalman_res.G*med #mean loglike
        P = kalman_res.G*varian*kalman_res.G' + kalman_res.R
        teste_cond = 1/cond(P)
        if teste_cond < kalman_tol
            llh[j] = -500
        #elseif det(P) <= 0
        #    llh[j] = -500
        else
            llh[j] = -(p*log(2*pi) + logdet(P) .+ eta'*inv(P)*eta)/2
            QuantEcon.update!(kalman_res,data[j+1,:]) #updating the kalman estimates
        end
        #println(kalman_res.cur_sigma)
        #println(det(varian))
        #println(P)
    end
    llh = llh[10:length(llh)]
    return sum(llh)
end

max_stars_repo_name: danmrc/DSGE-HMC

max_issues_repo_name: danmrc/DSGE-HMC

max_forks_repo_name: danmrc/DSGE-HMC


text:
mutable struct Descriptor <: AbstractDescriptor
    name::String
    p::libgb.GrB_Descriptor
    function Descriptor(name, p::libgb.GrB_Descriptor)
        d = new(name, p)
        function f(descriptor)
            libgb.GrB_Descriptor_free(Ref(descriptor.p))
        end
        return finalizer(f, d)
    end
end

function Descriptor(;kwargs...)
    desc = Descriptor("", libgb.GrB_Descriptor_new())
    for (s, x) in kwargs
        setproperty!(desc, s, x)
    end
    return desc
end

Base.unsafe_convert(::Type{libgb.GrB_Descriptor}, d::Descriptor) = d.p

function Base.getproperty(d::Descriptor, s::Symbol)
    if s === :p
        return getfield(d, s)
    elseif s === :replace_output
        x = libgb.GxB_Desc_get(d, libgb.GrB_OUTP)
        if x == libgb.GrB_REPLACE
            return true
        else
            return false
        end
    elseif s === :complement_mask
        x = libgb.GxB_Desc_get(d, libgb.libgb.GrB_MASK)
        if x == libgb.GrB_COMP || x == libgb.GrB_STRUCT_COMP
            return true
        else
            return false
        end
    elseif s === :structural_mask
        x = libgb.GxB_Desc_get(d, libgb.GrB_MASK)
        if x == libgb.GrB_STRUCTURE || x == libgb.GrB_STRUCT_COMP
            return true
        else
            return false
        end
    elseif s === :transpose_input1
        x = libgb.GxB_Desc_get(d, libgb.GrB_INP0)
        if x == libgb.GrB_TRAN
            return true
        else
            return false
        end
    elseif s === :transpose_input2
        x = libgb.GxB_Desc_get(d, libgb.GrB_INP1)
        if x == libgb.GrB_TRAN
            return true
        else
            return false
        end
    elseif s === :nthreads
        return libgb.GxB_Desc_get(d, libgb.GxB_DESCRIPTOR_NTHREADS)
    elseif s === :chunk
        return libgb.GxB_Desc_get(d, libgb.GxB_DESCRIPTOR_CHUNK)
    elseif s === :sort
        if libgb.GxB_Desc_get(d, libgb.GxB_SORT) == libgb.GxB_DEFAULT
            return false
        else
            return true
        end
    elseif s === :axb_method
        x = libgb.GxB_Desc_get(d, libgb.GxB_AxB_METHOD)
        if x == libgb.GxB_AxB_GUSTAVSON
            return :gustavson
        elseif x == libgb.GxB_AxB_DOT
            return :dot
        elseif x == libgb.AxB_HASH
            return :hash
        elseif x == libgb.GxB_AxB_SAXPY
            return :saxpy
        else
            return :default
        end
    else
        return getfield(d, s)
    end
end

function Base.setproperty!(d::Descriptor, s::Symbol, x)
    if s === :p
        setfield!(d, s, x)
        return nothing
    elseif s === :replace_output
        x ? (y = libgb.GrB_REPLACE) : (y = libgb.GxB_DEFAULT)
        libgb.GxB_Desc_set(d, libgb.GrB_OUTP, y)
    elseif s === :complement_mask
        if x == false
            if d.structural_mask
                libgb.GxB_Desc_set(d, libgb.GrB_MASK, libgb.GrB_STRUCTURE)
            else
                libgb.GxB_Desc_set(d, libgb.GrB_MASK, libgb.GxB_DEFAULT)
            end
        else
            libgb.GxB_Desc_set(d, libgb.GrB_MASK, libgb.GrB_COMP)
        end
    elseif s === :structural_mask
        if x == false
            if d.complement_mask
                libgb.GxB_Desc_set(d, libgb.GrB_MASK, libgb.GrB_COMP)
            else
                libgb.GxB_Desc_set(d, libgb.GrB_MASK, libgb.GxB_DEFAULT)
            end
        else
            libgb.GxB_Desc_set(d, libgb.GrB_MASK, libgb.GrB_STRUCTURE)
        end
    elseif s === :transpose_input1
        libgb.GxB_Desc_set(d, libgb.GrB_INP0, x ? libgb.GrB_TRAN : libgb.GxB_DEFAULT)
    elseif s === :transpose_input2
        libgb.GxB_Desc_set(d, libgb.GrB_INP1, x ? libgb.GrB_TRAN : libgb.GxB_DEFAULT)
    elseif s === :nthreads
        libgb.GxB_Desc_set(d, libgb.GxB_DESCRIPTOR_NTHREADS, x)
    elseif s === :chunk
        libgb.GxB_Desc_set(d, libgb.GxB_DESCRIPTOR_CHUNK, x)
    elseif s === :sort
        libgb.GxB_Desc_set(d, libgb.GxB_SORT, x ? 3 : 0)
    end
end
function Base.propertynames(::Descriptor)
    return (
    :replace_output,
    :complement_mask,
    :structural_mask,
    :transpose_input1,
    :transpose_input2,
    :nthreads,
    :chunk,
    :sort,
    )
end

Base.show(io::IO, ::MIME"text/plain", d::Descriptor) = gxbprint(io, d)
Base.print(io::IO, d::Descriptor) = gxbprint(io, d)

function _handledescriptor(desc; out=nothing, in1 = nothing, in2 = nothing)
    if out === nothing && in1 === nothing && in2 === nothing
        if !(desc isa Descriptor)
            return C_NULL
        else
            return desc
        end
    end
    if desc == C_NULL || desc === nothing
        desc = Descriptor()
    end
    in1 isa Transpose && (desc.transpose_input1 = true)
    in2 isa Transpose && (desc.transpose_input2 = true)
    return desc
end


text:
"""
    CH.jl

This is a Julia port of a MATLAB implementation of batch and incremental
Calinski-Harabasz (CH) Cluster Validity Index.

Authors:
MATLAB implementation: Leonardo Enzo Brito da Silva
Julia port: Sasha Petrenko <sap625@mst.edu>

REFERENCES
[1] L. E. Brito da Silva, N. M. Melton, and D. C. Wunsch II, "Incremental
Cluster Validity Indices for Hard Partitions: Extensions  and  Comparative
Study," ArXiv  e-prints, Feb 2019, arXiv:1902.06711v1 [cs.LG].
[2] T. Calinski and J. Harabasz, "A dendrite method for cluster analysis,"
Communications in Statistics, vol. 3, no. 1, pp. 1–27, 1974.
[3] M. Moshtaghi, J. C. Bezdek, S. M. Erfani, C. Leckie, and J. Bailey,
"Online Cluster Validity Indices for Streaming Data," ArXiv e-prints, 2018,
arXiv:1801.02937v1 [stat.ML]. [Online].
[4] M. Moshtaghi, J. C. Bezdek, S. M. Erfani, C. Leckie, J. Bailey, "Online
cluster validity indices for performance monitoring of streaming data clustering,"
Int. J. Intell. Syst., pp. 1–23, 2018.
"""

const common_string = """

# Arguments
- `cvi::CVI`: the CVI object.
- `sample::RealVector`: a vector of features used in clustering the sample.
- `label::Integer`: the cluster label prescribed to the sample by the clustering algorithm.
"""

"""
    CH

The stateful information of the Calinski-Harabasz (CH) Cluster Validity Index

# References
1. L. E. Brito da Silva, N. M. Melton, and D. C. Wunsch II, "Incremental Cluster Validity Indices for Hard Partitions: Extensions  and  Comparative Study," ArXiv  e-prints, Feb 2019, arXiv:1902.06711v1 [cs.LG].
2. T. Calinski and J. Harabasz, "A dendrite method for cluster analysis," Communications in Statistics, vol. 3, no. 1, pp. 1–27, 1974.
3. M. Moshtaghi, J. C. Bezdek, S. M. Erfani, C. Leckie, and J. Bailey, "Online Cluster Validity Indices for Streaming Data," ArXiv e-prints, 2018, arXiv:1801.02937v1 [stat.ML]. [Online].
4. M. Moshtaghi, J. C. Bezdek, S. M. Erfani, C. Leckie, J. Bailey, "Online cluster validity indices for performance monitoring of streaming data clustering," Int. J. Intell. Syst., pp. 1–23, 2018.
"""
mutable struct CH <: CVI
    label_map::LabelMap
    dim::Integer
    n_samples::Integer
    mu::RealVector      # dim
    n::IntegerVector    # dim
    v::RealMatrix       # dim x n_clusters
    CP::RealVector      # dim
    SEP::RealVector     # dim
    G::RealMatrix       # dim x n_clusters
    BGSS::Float
    WGSS::Float
    n_clusters::Integer
    criterion_value::Float
end # CH <: CVI

"""
    CH()

Default constructor for the Calinski-Harabasz (CH) Cluster Validity Index.
"""
function CH()
    CH(
        LabelMap(),                     # label_map
        0,                              # dim
        0,                              # n_samples
        Array{Float, 1}(undef, 0),      # mu
        Array{Integer, 1}(undef, 0),    # n
        Array{Float, 2}(undef, 0, 0),   # v
        Array{Float, 1}(undef, 0),      # CP
        Array{Float, 1}(undef, 0),      # SEP
        Array{Float, 2}(undef, 0, 0),   # G
        0.0,                            # BGSS
        0.0,                            # WGSS
        0,                              # n_clusters
        0.0                             # criterion_value
    )
end # CH()

"""
    setup!(cvi::CH, sample::Vector{T}) where {T<:RealFP}
"""
function setup!(cvi::CH, sample::Vector{T}) where {T<:RealFP}
    # Get the feature dimension
    cvi.dim = length(sample)
    # Initialize the augmenting 2-D arrays with the correct feature dimension
    # NOTE: R is emptied and calculated in evaluate!, so it is not defined here
    cvi.v = Array{T, 2}(undef, cvi.dim, 0)
    cvi.G = Array{T, 2}(undef, cvi.dim, 0)
end # setup!(cvi::CH, sample::Vector{T}) where {T<:Real}

function param_inc!(cvi::CH, sample::RealVector, label::Integer)
    # Get the internal label
    i_label = get_internal_label!(cvi.label_map, label)

    n_samples_new = cvi.n_samples + 1
    if isempty(cvi.mu)
        mu_new = sample
        setup!(cvi, sample)
    else
        mu_new = (1 - 1/n_samples_new) .* cvi.mu + (1/n_samples_new) .* sample
    end

    if i_label > cvi.n_clusters
        n_new = 1
        v_new = sample
        CP_new = 0.0
        G_new = zeros(cvi.dim)
        # Update 1-D parameters with a push
        cvi.n_clusters += 1
        push!(cvi.CP, CP_new)
        push!(cvi.n, n_new)
        # Update 2-D parameters with appending and reassignment
        cvi.v = [cvi.v v_new]
        cvi.G = [cvi.G G_new]
    else
        n_new = cvi.n[i_label] + 1
        v_new = (1 - 1/n_new) .* cvi.v[:, i_label] + (1/n_new) .* sample
        delta_v = cvi.v[:, i_label] - v_new
        diff_x_v = sample .- v_new
        CP_new = cvi.CP[i_label] + transpose(diff_x_v)*diff_x_v + cvi.n[i_label]*transpose(delta_v)*delta_v + 2*transpose(delta_v)*cvi.G[:, i_label]
        G_new = cvi.G[:, i_label] + diff_x_v + cvi.n[i_label].*delta_v
        # Update parameters
        cvi.n[i_label] = n_new
        cvi.v[:, i_label] = v_new
        cvi.CP[i_label] = CP_new
        cvi.G[:, i_label] = G_new
    end
    cvi.n_samples = n_samples_new
    cvi.mu = mu_new
    cvi.SEP = [cvi.n[ix] * sum((cvi.v[:, ix] - cvi.mu).^2) for ix=1:cvi.n_clusters]
end # param_inc!(cvi::CH, sample::RealVector, label::Integer)

function param_batch!(cvi::CH, data::RealMatrix, labels::IntegerVector)
    cvi.dim, cvi.n_samples = size(data)
    # Take the average across all samples, but cast to 1-D vector
    cvi.mu = mean(data, dims=2)[:]
    # u = findfirst.(isequal.(unique(labels)), [labels])
    u = unique(labels)
    cvi.n_clusters = length(u)
    cvi.n = zeros(Integer, cvi.n_clusters)
    cvi.v = zeros(cvi.dim, cvi.n_clusters)
    cvi.CP = zeros(cvi.n_clusters)
    cvi.SEP = zeros(cvi.n_clusters)
    for ix = 1:cvi.n_clusters
        subset = data[:, findall(x->x==u[ix], labels)]
        cvi.n[ix] = size(subset, 2)
        cvi.v[1:cvi.dim, ix] = mean(subset, dims=2)
        diff_x_v = subset - cvi.v[:, ix] * ones(1, cvi.n[ix])
        cvi.CP[ix] = sum(diff_x_v.^2)
        cvi.SEP[ix] = cvi.n[ix] * sum((cvi.v[:, ix] - cvi.mu).^2);
    end
end # param_batch!(cvi::CH, data::RealMatrix, labels::IntegerVector)

function evaluate!(cvi::CH)
    # Within group sum of scatters
    cvi.WGSS = sum(cvi.CP)
    if cvi.n_clusters > 2
        # Between groups sum of scatters
        cvi.BGSS = sum(cvi.SEP)
        # CH index value
        cvi.criterion_value = (cvi.BGSS / cvi.WGSS) * ((cvi.n_samples - cvi.n_clusters)/(cvi.n_clusters - 1))
    else
        cvi.BGSS = 0.0
        cvi.criterion_value = 0.0
    end
end # evaluate(cvi::CH)


text:
#! /usr/bin/julia
#! /mnt/d/p/Julia/juliaLinux/bin/julia
#PATH=/mnt/d/p/Julia/juliaLinux/bin:$PATH

include("RawFileofflineProc/interface.jl") # 如果需要适配不同的算法来分析相应文件，需要修改 intrface.jl 中包含的算法库

##########
startMeters =100
endmeters  =31100
ProtectMeters =30
chnid =1#  1 or  2
maxTimeCnt1 =-1## -1：time loop forever ,until data end

println( "startMeters:$startMeters,\nendmeters:$endmeters,
 \nProtectMeters:$ProtectMeters,
 \nchnid :$chnid ,\nmaxTimeCnt1:$maxTimeCnt1\n")

###################### Linux  READ RAW  LC  Files ，必须指定  cfgfilename 和 dataFilename
basePath1 = "g:/other/fiber/fiberdata" #/home/wang/Documents/wangzht/fiber/fiberdata
basePath1 = "E:/DATAs/xxxx号机告警分析/20190212GY01BigWind"
# cfgfilename  =""# joinpath(basePath1, "20190212T131313.CONF"  )
# dataFilename = joinpath( basePath1,"20190212T131313.lc") #  处理单个文件时使用
# @time calInternal(chnid,startMeters ,endmeters,maxTimeCnt1, dataFilename ,cfgfilename, ProtectMeters  )

######################  READ RAW1~RAW4 Files

dataFilename1="20190212T131313.RAW3"
dataFilename = joinpath(basePath1 ,dataFilename1)
# dataFilename="G:\\Bd003\\data\\userdata1\\Data^Ver00^BD-003^20181107045148^ID00.RAW4"
calPT(chnid,startMeters ,endmeters,maxTimeCnt1, dataFilename , ProtectMeters )

exit(1)

max_stars_repo_name: wangzhongtian/fiberAlarmAnalyseTool

max_issues_repo_name: wangzhongtian/fiberAlarmAnalyseTool

max_forks_repo_name: wangzhongtian/fiberAlarmAnalyseTool


text:
function is_leap_year(year::Int)
    year % 4 == 0 && (year % 100 != 0 || year % 400 == 0);
end

max_stars_repo_name: cmccandless/ExercismSolutions-julia

max_issues_repo_name: cmccandless/ExercismSolutions-julia

max_forks_repo_name: cmccandless/ExercismSolutions-julia


text:
# linear softmax classifier (with stochastic gradient descent)
type LinearClassifier
    k::Int64 # number of outputs
    n::Int64 # number of inputs
    weights::Array{Float64, 2} # n * k weight matrix

    outputs :: Vector{Float64}
end

function LinearClassifier(k, n)
    weights = rand(n, k) * 2 - 1; # range is [-1, 1]
    LinearClassifier(k, n, weights, zeros(k))
end

function predict(c::LinearClassifier, x::Array{Float64})
    # the softmax() function from the NumericExtension package is more numeric stable
    return vec(softmax(x * c.weights))
end

function predict!(c::LinearClassifier, x::Array{Float64})
    # c.outputs = vec(softmax(x * c.weights))
    s = 0.0
    for i in 1:c.k
        s = 0.0
        for j in 1:c.n
            s += x[j] * c.weights[j, i]
        end
        c.outputs[i] = s
    end

    softmax!(c.outputs, c.outputs);
end

function train_one(c::LinearClassifier, x::Array{Float64}, y::Int64, α::Float64=0.025)
    # if !in(y, 1 : c.k)
    #     msg = @sprintf "A sample is discarded because the label y = %d is not in range of 1 to %d" y c.k
    #     warn(msg)
    #     return
    # end

    predict!(c, x)
    c.outputs[y] -= 1

    # c.weights -= α * x' * outputs;
    # BLAS.ger!(-α, vec(x), c.outputs, c.weights)
    m = 0.0
    j = 0
    limit = c.n - 4
    for i in 1:c.k
        m = α * c.outputs[i]
        j = 1
        while j <= limit
            @nexprs 4 (idx->c.weights[j + idx - 1, i] -= m * x[j + idx - 1])
            j+=4
        end
        while j <= c.n
            c.weights[j, i] -= m * x[j]
            j+=1
        end
    end
end

function train_one(c::LinearClassifier, x::Array{Float64}, y::Int64, input_gradient::Array{Float64}, α::Float64=0.025)
    predict!(c, x)
    c.outputs[y] -= 1

    # input_gradient = ( c.weights * outputs' )'
    # BLAS.gemv!('N', α, c.weights, c.outputs, 1.0, input_gradient)
    m = 0.0
    j = 0
    limit = c.n - 4
    for i in 1:c.k
        m = α * c.outputs[i]
        j = 1
        while j <= limit
            @nexprs 4 (idx->input_gradient[j+idx-1] += m * c.weights[j+idx-1, i])
            j+=4
        end
        while j <= c.n
            input_gradient[j] += m * c.weights[j, i]
            j+=1
        end
    end

    # c.weights -= α * x' * outputs;
    # BLAS.ger!(-α, vec(x), c.outputs, c.weights)
    for i in 1:c.k
        m = α * c.outputs[i]
        j = 1
        while j <= limit
            @nexprs 4 (idx->c.weights[j + idx - 1, i] -= m * x[j + idx - 1])
            j+=4
        end
        while j <= c.n
            c.weights[j, i] -= m * x[j]
            j+=1
        end
    end
end

# calculate the overall log likelihood. Mainly used for debugging
function log_likelihood(c, X, y)
    n = size(X, 1)
    l = 0
    for i in 1:n
        l += log(predict(c, X[i, :])[y[i]])
    end
    return l
end

# calculate the accuracy on the testing dataset
function accuracy(c::LinearClassifier, X::Array{Float64}, y::Array{Int64})
    n = size(X, 1)
    succ = 0
    for i in 1 : n
        output = predict(c, X[i, :])
        if maximum(output) == output[y[i]]
            succ += 1
        end
    end
    return succ / n
end


text:
module BasicTests

using Photon, Test
using Photon: batchfirst, batchlast

function KorATest()
    a = randn(10,10)
    b = KorA(a)
    @test size(a) == size(b)
    @test isapprox(sum(a), sum(b))
end

function utils()
    a = randn(32,10)
    b = batchlast(a)
    @test  a[1,:] == b[:,1]

    a = randn(10, 32)
    b = batchfirst(a)
    @test  a[:,1] == b[1,:]

    a = randn(32, 10)
    @test batchfirst(batchlast(a)) == a
end

@testset "Basic" begin
    KorATest()
    utils()
end

end


text:
module Cylindrical

    using Plots
    using Plots.PlotMeasures
    using RecipesBase
    using LaTeXStrings

    using GeometricIntegrators.Solutions: TimeSeries, DataSeries, Solution
    using GeometricProblems.PlotRecipes: subscript

    include("common.jl")

    @userplot PlotTrajectory
    @recipe function f(p::PlotTrajectory; plottype=nothing, nplot=1, xlims=:auto, ylims=:auto, latex=true)
        if length(p.args) != 1 || !(typeof(p.args[1]) <: Solution)
            error("Guiding center trajectory plots should be given a solution. Got: $(typeof(p.args))")
        end

        sol = p.args[1]

        linewidth=2
        if sol.nt ≤ 1000
            markersize := 5
        else
            markersize  := 1
            markercolor := 1
            linecolor   := 1
            markerstrokewidth := 1
            markerstrokecolor := 1
        end

        seriestype := :scatter
        guidefont  := font(18)
        tickfont   := font(12)
        legend := :none

        if plottype == nothing
            size := (600,800)

            layout := @layout [Rplot
                               Zplot
                               ϕplot
                               Uplot]

            if latex
                ylabels = (L"R", L"Z", L"\varphi", L"u")
            else
                ylabels = ("R", "Z", "φ", "u")
            end

            for i in 1:4
                @series begin
                    subplot := i
                    if i == 4
                        if latex
                            xlabel := L"t"
                        else
                            xlabel := "t"
                        end
                    else
                        xaxis := false
                    end
                    ylabel := ylabels[i]
                    sol.t[0:nplot:end], sol.q[i,0:nplot:end]
                end
            end

        elseif plottype==:trajectoryRZ
            if latex
                xlabel := L"R"
                ylabel := L"Z"
            else
                xlabel := "R"
                ylabel := "Z"
            end
            xlims  := xlims
            ylims  := ylims
            size   := (500,500)

            sol.q[1,0:nplot:end], sol.q[2,0:nplot:end]

        elseif plottype==:trajectoryXY
            if latex
                xlabel := L"X"
                ylabel := L"Y"
            else
                xlabel := "X"
                ylabel := "Y"
            end
            xlims  := xlims
            ylims  := ylims
            size   := (500,500)

            sol.q[1,0:nplot:end] .* cos.(sol.q[3,0:nplot:end]), sol.q[1,0:nplot:end] .* sin.(sol.q[3,0:nplot:end])

        elseif plottype==:trajectoryYZ
            if latex
                xlabel := L"Y"
                ylabel := L"Z"
            else
                xlabel := "Y"
                ylabel := "Z"
            end
            xlims  := xlims
            ylims  := ylims
            size   := (500,500)

            sol.q[1,0:nplot:end] .* sin.(sol.q[3,0:nplot:end]), sol.q[2,0:nplot:end]

        elseif plottype==:trajectoryXZ
            if latex
                xlabel := L"X"
                ylabel := L"Z"
            else
                xlabel := "X"
                ylabel := "Z"
            end
            xlims  := xlims
            ylims  := ylims
            size   := (500,500)

            sol.q[1,0:nplot:end] .* cos.(sol.q[3,0:nplot:end]), sol.q[2,0:nplot:end]

        elseif plottype==:trajectory3d
            if latex
                xlabel := L"X"
                ylabel := L"Y"
                zlabel := L"Z"
            else
                xlabel := "X"
                ylabel := "Y"
                zlabel := "Z"
            end
            xlims  := xlims
            ylims  := ylims
            size --> (500,500)

            sol.q[1,0:nplot:end] .* cos.(sol.q[3,0:nplot:end]), sol.q[1,0:nplot:end] .* sin.(sol.q[3,0:nplot:end]), sol.q[2,0:nplot:end]
        end
    end

end