text:
# This file is part of Kpax3. License is MIT.

function test_data_exceptions()
  f = "../build/test"

  @test_throws Kpax3.KFASTAError Kpax3.NucleotideData(Kpax3.KSettings("data/read_empty_file.fasta", f))
  @test_throws Kpax3.KFASTAError Kpax3.NucleotideData(Kpax3.KSettings("data/read_no_id_char.fasta", f))
  @test_throws Kpax3.KFASTAError Kpax3.NucleotideData(Kpax3.KSettings("data/read_no_1st_seq.fasta", f))
  @test_throws Kpax3.KFASTAError Kpax3.NucleotideData(Kpax3.KSettings("data/read_no_nth_seq.fasta", f))

  @test_throws Kpax3.KFASTAError Kpax3.NucleotideData(Kpax3.KSettings("data/read_utf8_id.fasta", f))
  @test_throws Kpax3.KFASTAError Kpax3.NucleotideData(Kpax3.KSettings("data/read_utf8_seq.fasta", f))

  @test_throws Kpax3.KFASTAError Kpax3.AminoAcidData(Kpax3.KSettings("data/read_empty_file.fasta", f))
  @test_throws Kpax3.KFASTAError Kpax3.AminoAcidData(Kpax3.KSettings("data/read_no_id_char.fasta", f))
  @test_throws Kpax3.KFASTAError Kpax3.AminoAcidData(Kpax3.KSettings("data/read_no_1st_seq.fasta", f))
  @test_throws Kpax3.KFASTAError Kpax3.AminoAcidData(Kpax3.KSettings("data/read_no_nth_seq.fasta", f))

  @test_throws Kpax3.KFASTAError Kpax3.AminoAcidData(Kpax3.KSettings("data/read_utf8_id.fasta", f))
  @test_throws Kpax3.KFASTAError Kpax3.AminoAcidData(Kpax3.KSettings("data/read_utf8_seq.fasta", f))

  @test_throws Kpax3.KCSVError Kpax3.CategoricalData(Kpax3.KSettings("data/read_empty_file.csv", f))
  @test_throws Kpax3.KCSVError Kpax3.CategoricalData(Kpax3.KSettings("data/read_no_id_char.csv", f))
  @test_throws Kpax3.KCSVError Kpax3.CategoricalData(Kpax3.KSettings("data/read_no_1st_seq.csv", f))
  @test_throws Kpax3.KCSVError Kpax3.CategoricalData(Kpax3.KSettings("data/read_no_nth_seq.csv", f))

  nothing
end

test_data_exceptions()

function test_data_blanks()
  f = "../build/test"

  nt = Kpax3.NucleotideData(Kpax3.KSettings("data/read_blanks.fasta", f))
  @test nt.data == UInt8[0 1;
                         1 0;
                         1 0;
                         0 1;
                         0 1;
                         1 0;
                         1 0;
                         0 1;
                         1 0;
                         0 1;
                         0 1;
                         1 0]
  @test nt.id == ["ID1", "ID5"]
  @test nt.ref == UInt8['a', 't', 'g', '.', '.', '.', 'g', '.', '.', 'a', '.', 'a']
  @test nt.val == UInt8['a', 'g', 'a', 'g', 'c', 'g', 'c', 'g', 'a', 'g', 'c', 'g']
  @test nt.key == [1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6]

  nt = Kpax3.NucleotideData(Kpax3.KSettings("data/read_blanks.fasta", f, l=1))
  @test nt.data == UInt8[0 1;
                         1 0;
                         1 0;
                         0 1;
                         0 1;
                         1 0;
                         1 0;
                         0 1;
                         1 0;
                         0 1;
                         0 1;
                         1 0]
  @test nt.id == ["ID1", "ID5"]
  @test nt.ref == UInt8['a', 't', 'g', '.', '.', '.', 'g', '.', '.', 'a', '.', 'a']
  @test nt.val == UInt8['a', 'g', 'a', 'g', 'c', 'g', 'c', 'g', 'a', 'g', 'c', 'g']
  @test nt.key == [1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6]

  aa = Kpax3.AminoAcidData(Kpax3.KSettings("data/read_blanks.fasta", f))
  @test aa.data == UInt8[0 1;
                         1 0;
                         1 0;
                         0 1;
                         0 1;
                         1 0;
                         1 0;
                         0 1;
                         1 0;
                         0 1;
                         0 1;
                         1 0]
  @test aa.id == ["ID1", "ID5"]
  @test aa.ref == UInt8['a', 't', 'g', '.', '.', '.', 'g', '.', '.', 'a', '.', 'a']
  @test aa.val == UInt8['a', 'g', 'a', 'g', 'c', 'g', 'c', 'g', 'a', 'g', 'c', 'g']
  @test aa.key == [1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6]

  aa = Kpax3.AminoAcidData(Kpax3.KSettings("data/read_blanks.fasta", f, l=1))
  @test aa.data == UInt8[0 1;
                         1 0;
                         1 0;
                         0 1;
                         0 1;
                         1 0;
                         1 0;
                         0 1;
                         1 0;
                         0 1;
                         0 1;
                         1 0]
  @test aa.id == ["ID1", "ID5"]
  @test aa.ref == UInt8['a', 't', 'g', '.', '.', '.', 'g', '.', '.', 'a', '.', 'a']
  @test aa.val == UInt8['a', 'g', 'a', 'g', 'c', 'g', 'c', 'g', 'a', 'g', 'c', 'g']
  @test aa.key == [1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6]

  cg = Kpax3.CategoricalData(Kpax3.KSettings("data/read_blanks.csv", f, l=1))

  nothing
end

test_data_blanks()

function test_data_proper_dna_file()
  f = "../build/test"

  nt = Kpax3.NucleotideData(Kpax3.KSettings("data/read_proper_nt.fasta", f))
  @test nt.data == UInt8[0 0 0 0 0 1;
                         1 1 1 1 1 0;
                         0 0 1 0 1 1;
                         1 1 0 1 0 0;
                         1 1 0 0 0 0;
                         0 0 0 1 1 0;
                         0 0 0 1 1 1;
                         1 1 1 0 0 0;
                         0 0 1 0 0 0;
                         1 0 0 1 0 1;
                         0 1 0 0 1 0;
                         1 1 1 0 0 0;
                         0 0 0 1 1 0;
                         0 0 0 0 0 1;
                         1 1 0 0 1 1;
                         0 0 1 1 0 0;
                         0 1 0 0 1 0;
                         1 0 1 1 0 1]
  @test nt.id == String["ID1", "ID2", "ID3", "ID4", "ID5", "ID6"]
  @test nt.ref == UInt8['a', '.', 'g', '.', '.', '.', 'g', '.', '.', '.', '.', 'a']
  @test nt.val == UInt8['c', 't', 'a', 'g', 'a', 'g', 'c', 'g', 'a', 'c', 'g', 'a', 'g', 't', 'a', 't', 'c', 'g']
  @test nt.key == [1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7, 8, 8]

  nt = Kpax3.NucleotideData(Kpax3.KSettings("data/read_proper_nt.fasta", f, l=1))
  @test nt.data == UInt8[0 0 0 0 0 1;
                         1 1 1 1 1 0;
                         0 0 1 0 1 1;
                         1 1 0 1 0 0;
                         1 1 0 0 0 0;
                         0 0 0 1 1 0;
                         0 0 0 1 1 1;
                         1 1 1 0 0 0;
                         0 0 1 0 0 0;
                         1 0 0 1 0 1;
                         0 1 0 0 1 0;
                         1 1 1 0 0 0;
                         0 0 0 1 1 0;
                         0 0 0 0 0 1;
                         1 1 0 0 1 1;
                         0 0 1 1 0 0;
                         0 1 0 0 1 0;
                         1 0 1 1 0 1]
  @test nt.id == String["ID1", "ID2", "ID3", "ID4", "ID5", "ID6"]
  @test nt.ref == UInt8['a', '.', 'g', '.', '.', '.', 'g', '.', '.', '.', '.', 'a']
  @test nt.val == UInt8['c', 't', 'a', 'g', 'a', 'g', 'c', 'g', 'a', 'c', 'g', 'a', 'g', 't', 'a', 't', 'c', 'g']
  @test nt.key == [1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7, 8, 8]

  # consider all characters
  nt = Kpax3.NucleotideData(Kpax3.KSettings("data/read_proper_nt.fasta", f,
                                miss=zeros(UInt8, 1)))
  @test nt.data == UInt8[0 0 0 0 0 1;
                         1 1 1 1 1 0;
                         0 0 1 0 1 1;
                         1 1 0 1 0 0;
                         1 1 0 0 0 0;
                         0 0 0 1 1 0;
                         0 0 1 0 0 1;
                         0 0 0 1 1 1;
                         1 1 1 0 0 0;
                         0 0 1 0 0 0;
                         1 0 0 1 0 1;
                         0 1 0 0 1 0;
                         1 1 1 0 0 0;
                         0 0 0 1 1 0;
                         0 0 0 0 0 1;
                         1 1 0 0 1 1;
                         0 0 1 1 0 0;
                         0 1 0 0 1 0;
                         1 0 1 1 0 1;
                         1 0 1 1 1 1;
                         0 1 0 0 0 0]
  @test nt.id == String["ID1", "ID2", "ID3", "ID4", "ID5", "ID6"]
  @test nt.ref == UInt8['a', '.', 'g', '.', '.', '.', 'g', '.', '.', '.', '.', '.']
  @test nt.val == UInt8['c', 't', 'a', 'g', 'a', 'g', 'x', 'c', 'g', 'a', 'c', 'g', 'a', 'g', 't', 'a', 't', 'c', 'g', 'a', 'x']
  @test nt.key == [1, 1, 2, 2, 3, 3, 3, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7, 8, 8, 9, 9]

  nothing
end

test_data_proper_dna_file()

function test_data_proper_protein_file()
  f = "../build/test"

  aa = Kpax3.AminoAcidData(Kpax3.KSettings("data/read_proper_aa.fasta", f))
  @test aa.data == UInt8[0 0 0 0 0 1;
                         1 1 1 1 1 0;
                         0 0 1 0 1 1;
                         1 1 0 1 0 0;
                         1 1 0 0 0 0;
                         0 0 0 1 1 0;
                         1 1 1 0 0 0;
                         0 0 0 1 1 1;
                         0 0 1 0 0 0;
                         0 1 0 0 1 0;
                         1 0 0 1 0 1;
                         0 0 0 0 0 1;
                         1 1 1 0 0 0;
                         0 0 0 1 1 0;
                         1 1 0 0 1 1;
                         0 0 1 1 0 0;
                         1 1 0 1 0 0;
                         0 0 1 0 1 1]
  @test aa.id == String["ID1", "ID2", "ID3", "ID4", "ID5", "ID6"]
  @test aa.ref == UInt8['m', '.', 'a', '.', '.', '.', 'v', '.', '.', '.', '.', 'f']
  @test aa.val == UInt8['e', 'k', 'i', 'k', 'l', 'v', 'l', 'v', 'c', 'l', 't', 'e', 'l', 'm', 'c', 'y', 'a', 't']
  @test aa.key == [1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7, 8, 8]

  aa = Kpax3.AminoAcidData(Kpax3.KSettings("data/read_proper_aa.fasta", f, l=1))
  @test aa.data == UInt8[0 0 0 0 0 1;
                         1 1 1 1 1 0;
                         0 0 1 0 1 1;
                         1 1 0 1 0 0;
                         1 1 0 0 0 0;
                         0 0 0 1 1 0;
                         1 1 1 0 0 0;
                         0 0 0 1 1 1;
                         0 0 1 0 0 0;
                         0 1 0 0 1 0;
                         1 0 0 1 0 1;
                         0 0 0 0 0 1;
                         1 1 1 0 0 0;
                         0 0 0 1 1 0;
                         1 1 0 0 1 1;
                         0 0 1 1 0 0;
                         1 1 0 1 0 0;
                         0 0 1 0 1 1]
  @test aa.id == String["ID1", "ID2", "ID3", "ID4", "ID5", "ID6"]
  @test aa.ref == UInt8['m', '.', 'a', '.', '.', '.', 'v', '.', '.', '.', '.', 'f']
  @test aa.val == UInt8['e', 'k', 'i', 'k', 'l', 'v', 'l', 'v', 'c', 'l', 't', 'e', 'l', 'm', 'c', 'y', 'a', 't']
  @test aa.key == [1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7, 8, 8]

  # consider all characters
  aa = Kpax3.AminoAcidData(Kpax3.KSettings("data/read_proper_aa.fasta", f, miss=zeros(UInt8,1)))
  @test aa.data == UInt8[0 0 0 0 0 1;
                         1 1 1 1 1 0;
                         0 0 1 0 1 1;
                         1 1 0 1 0 0;
                         1 1 0 0 0 0;
                         0 0 0 1 1 0;
                         0 0 1 0 0 1;
                         1 1 1 0 0 0;
                         0 0 0 1 1 1;
                         0 0 1 0 0 0;
                         0 1 0 0 1 0;
                         1 0 0 1 0 1;
                         0 0 0 0 0 1;
                         1 1 1 0 0 0;
                         0 0 0 1 1 0;
                         1 1 0 0 1 1;
                         0 0 1 1 0 0;
                         1 1 0 1 0 0;
                         0 0 1 0 1 1;
                         1 0 1 1 1 1;
                         0 1 0 0 0 0]
  @test aa.id == String["ID1", "ID2", "ID3", "ID4", "ID5", "ID6"]
  @test aa.ref == UInt8['m', '.', 'a', '.', '.', '.', 'v', '.', '.', '.', '.', '.']
  @test aa.val == UInt8['e', 'k', 'i', 'k', 'l', 'v', 'x', 'l', 'v', 'c', 'l', 't', 'e', 'l', 'm', 'c', 'y', 'a', 't', 'f', 'x']
  @test aa.key == [1, 1, 2, 2, 3, 3, 3, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7, 8, 8, 9, 9]

  nothing
end

test_data_proper_protein_file()

function test_data_input_output()
  f = "../build/test"

  nt = Kpax3.NucleotideData(Kpax3.KSettings("data/read_proper_nt.fasta", f))

  # TODO: Test exception when saving to a location without writing permissions
  Kpax3.save("../build/nt.jld", nt)
  @test isfile("../build/nt.jld")

  @test_throws SystemError Kpax3.loadnt("../build/non_existent.file")

  nt = Kpax3.loadnt("../build/nt.jld")

  @test isa(nt, Kpax3.NucleotideData)
  @test nt.data == UInt8[0 0 0 0 0 1;
                         1 1 1 1 1 0;
                         0 0 1 0 1 1;
                         1 1 0 1 0 0;
                         1 1 0 0 0 0;
                         0 0 0 1 1 0;
                         0 0 0 1 1 1;
                         1 1 1 0 0 0;
                         0 0 1 0 0 0;
                         1 0 0 1 0 1;
                         0 1 0 0 1 0;
                         1 1 1 0 0 0;
                         0 0 0 1 1 0;
                         0 0 0 0 0 1;
                         1 1 0 0 1 1;
                         0 0 1 1 0 0;
                         0 1 0 0 1 0;
                         1 0 1 1 0 1]
  @test nt.id == String["ID1", "ID2", "ID3", "ID4", "ID5", "ID6"]
  @test nt.ref == UInt8['a', '.', 'g', '.', '.', '.', 'g', '.', '.', '.', '.', 'a']
  @test nt.val == UInt8['c', 't', 'a', 'g', 'a', 'g', 'c', 'g', 'a', 'c', 'g', 'a', 'g', 't', 'a', 't', 'c', 'g']
  @test nt.key == [1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7, 8, 8]

  aa = Kpax3.AminoAcidData(Kpax3.KSettings("data/read_proper_aa.fasta", f))

  # TODO: Test exception when saving to a location without writing permissions
  Kpax3.save("../build/aa.jld", aa)
  @test isfile("../build/aa.jld")

  @test_throws SystemError Kpax3.loadaa("../build/non_existent.file")

  aa = Kpax3.loadaa("../build/aa.jld")
  @test isa(aa, Kpax3.AminoAcidData)
  @test aa.data == UInt8[0 0 0 0 0 1;
                         1 1 1 1 1 0;
                         0 0 1 0 1 1;
                         1 1 0 1 0 0;
                         1 1 0 0 0 0;
                         0 0 0 1 1 0;
                         1 1 1 0 0 0;
                         0 0 0 1 1 1;
                         0 0 1 0 0 0;
                         0 1 0 0 1 0;
                         1 0 0 1 0 1;
                         0 0 0 0 0 1;
                         1 1 1 0 0 0;
                         0 0 0 1 1 0;
                         1 1 0 0 1 1;
                         0 0 1 1 0 0;
                         1 1 0 1 0 0;
                         0 0 1 0 1 1]
  @test aa.id == String["ID1", "ID2", "ID3", "ID4", "ID5", "ID6"]
  @test aa.ref == UInt8['m', '.', 'a', '.', '.', '.', 'v', '.', '.', '.', '.', 'f']
  @test aa.val == UInt8['e', 'k', 'i', 'k', 'l', 'v', 'l', 'v', 'c', 'l', 't', 'e', 'l', 'm', 'c', 'y', 'a', 't']
  @test aa.key == [1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7, 8, 8]

  nothing
end

test_data_input_output()


text:
#
# utils.jl -
#
# General purpose functions.
#

"""

```julia
circularmask(R, [rmin = 0,] rmax) -> M
```

yields an array `M` of same dimensions (and element type) as `R` which is
set with a circular mask at radii `R`.  Argument `rmax` is the outer radius
of the mask while optional argument `rmin` is the inner radius of the mask
(0 by default).

The value of the result is 1 inside the mask, 0 ouside the mask and in the
range `[0,1]` at the mask borders.  Keyword `rev` can be set true to
reverse the mask, that is to set the result to 0 inside the mask and 1
outside the mask.

Keyword `shape` can be set with a non-zero value to use simple
anti-aliasing rules for the edges.  The thickness (in same units as `R`) of
the mask border is given by the absolute value of `shape`.  Thus, the
support of the mask has a radius: `rmax + abs(shape)/2`.  With `shape > 0`,
a linear ramp is used; with `shape = 0` (the default), a rectangular
profile is used; with `shape < 0`, a sine profile is used.

The destination mask may be specified:

```julia
circularmask!(M, R, [rmin = 0,] rmax) -> M
```

"""
circularmask(R::AbstractArray{<:AbstractFloat}, args...; kwds...) =
     circularmask!(similar(R), R, args...; kwds...)

function circularmask!(M::AbstractArray{Tm,N},
                       R::AbstractArray{Tr,N},
                       r1::Real, r2::Real=0;
                       shape::Real = 0,
                       rev::Bool = false) where {Tm<:AbstractFloat,
                                                 Tr<:AbstractFloat, N}
    @assert size(M) == size(R)
    const rmin = convert(Tr, min(r1, r2))
    const rmax = convert(Tr, max(r1, r2))
    const inside = (rev ? zero(Tm) : one(Tm))  # value inside mask
    const outside = (rev ? one(Tm) : zero(Tm)) # value ouside mask
    const edge = (inside + outside)/2          # value at middle of edge

    if shape == 0
        # Points exactly at the edges are set to 1/2.
        if rmin > 0
            # Account for a central obscuration.
            @inbounds for i in eachindex(M, R)
                r = R[i]
                M[i] = (r < rmin || r > rmax ? outside :
                        rmin < r < rmax ? inside : edge)
            end
        else
            # No central obscuration.
            @inbounds for i in eachindex(M, R)
                r = R[i]
                M[i] = (r > rmax ? outside :
                        r < rmax ? inside  : edge)
            end
        end
    else
        # Points near the edges are set to a value between 0 and 1
        # following a given profile.  Note that the expressions at the
        # edges are centered to minimize rounding errors.
        const h = abs(convert(Tr, shape))/2 # half border thickness
        const lim1 = rmin - h
        const lim2 = (rmin + h ≤ rmax - h ? rmin + h : (rmin + rmax)/2)
        const lim3 = (rmin + h ≤ rmax - h ? rmax - h : (rmin + rmax)/2)
        const lim4 = rmax + h
        if shape > 0
            # Use a linear B-spline for borders.
            const a = (inside - outside)/(h + h)
            if lim2 > 0
                # Account for a central obscuration.
                @inbounds for i in eachindex(M, R)
                    r = R[i]
                    M[i] = (r ≥ lim4 ? outside             :
                            r > lim3 ? (rmax - r)*a + edge :
                            r ≥ lim2 ? inside              :
                            r > lim1 ? (r - rmin)*a + edge : outside)
                end
            else
                # No central obscuration.
                @inbounds for i in eachindex(M, R)
                    r = R[i]
                    M[i] = (r ≥ lim4 ? outside             :
                            r > lim3 ? (rmax - r)*a + edge : inside)
                end
            end
        else # shape < 0
            # Use a sine window for borders.
            const a = convert(Tr, π)*(inside - outside)/(h + h)
            if lim2 > 0
                # Account for a central obscuration.
                @inbounds for i in eachindex(M, R)
                    r = R[i]
                    M[i] = (r ≥ lim4 ? outside                       :
                            r > lim3 ? sin((rmax - r)*a)*edge + edge :
                            r ≥ lim2 ? inside                        :
                            r > lim1 ? sin((r - rmin)*a)*edge + edge : outside)
                end
            else
                # No central obscuration.
                @inbounds for i in eachindex(M, R)
                    r = R[i]
                    M[i] = (r ≥ lim4 ? outside                       :
                            r > lim3 ? sin((rmax - r)*a)*edge + edge : inside)
                end
            end
        end
    end
    return M
end

@doc @doc(circularmask) circularmask!

"""
```julia
fftshiftphasor([T,] s, n)
```

yields a vector of `n` complexes whose elements have values `exp(2*i*π*k*s/n)`
for `k ∈ fftfreq(n)`.  Optional argument `T <: Complex{<:AbstractFloat}` is the
type of the elements of the result.  By default, `T =
Complex{float(typeof(s))}`.

"""
fftshiftphasor(s::T, n::Integer) where {T<:Real} =
    fftshiftphasor(Complex{float(T)}, s, n)

function fftshiftphasor(::Type{Complex{T}}, s::Real,
                        n::Integer) where {T<:AbstractFloat}
    fftshiftphasor!(Array{Complex{T}}(n), s)
end

function fftshiftphasor!(dst::AbstractVector{Complex{T}},
                         s::Real) where {T<:AbstractFloat}
    # For fractional shifts, it is important to use the signed discrete
    # frequency index.
    len = length(dst)
    n = div(len, 2)
    α = convert(T, 2π*s/len)
    @inbounds @simd for k in 1:len-n
        ϕ = α*(k - 1)
        dst[k] = complex(cos(ϕ), sin(ϕ))
    end
    l = len + 1
    @inbounds @simd for k in len-n+1:len
        ϕ = α*(k - l)
        dst[k] = complex(cos(ϕ), sin(ϕ))
    end
    if iseven(len)
        # Treat the Nyquist frequency specifically for a slight gain in
        # precision.
        k = n + 1
        dst[k] = real(dst[k])
    end
    return dst
end

phasor(phi::AbstractArray{T,N}) where {T<:Real,N} =
    phasor(Complex{float(T)}, phi)

function phasor(::Type{Complex{T}},
                phi::AbstractArray{<:Real,N}) where {T<:AbstractFloat,N}
    phasor!(Array{Complex{T}}(size(phi)), phi)
end

function phasor!(dst::AbstractArray{Complex{T},N},
                 phi::AbstractArray{<:Real,N}) where {T<:AbstractFloat,N}
    @assert size(dst) == size(phi)
    @inbounds @simd for k in eachindex(dst, phi)
        ϕ = convert(T, phi[k])
        dst[k] = complex(cos(ϕ), sin(ϕ))
    end
    return dst
end

phasor(alpha::Ta, X::AbstractArray{Tx,N}) where {Ta<:Real,Tx<:Real,N} =
    phasor(Complex{float(promote_type(Ta, Tx))}, alpha, X)

function phasor(::Type{Complex{T}},
                alpha::Real,
                X::AbstractArray{<:Real,N}) where {T<:AbstractFloat,N}
    phasor!(Array{Complex{T}}(size(X)), alpha, X)
end

function phasor!(dst::AbstractArray{Complex{T},N},
                 alpha::Real,
                 X::AbstractArray{<:Real,N}) where {T<:AbstractFloat,N}
    @assert size(dst) == size(X)
    α = convert(T, alpha)
    @inbounds @simd for k in eachindex(dst, X)
        ϕ = α*convert(T, X[k])
        dst[k] = complex(cos(ϕ), sin(ϕ))
    end
    return dst
end

ext: jl


text:
using Yao
using BitBasis
using Test, LinearAlgebra
using QuAlgorithmZoo

"""
random phase estimation problem setup.
"""
function rand_phaseest_setup(N::Int)
    U = rand_unitary(1<<N)
    b = randn(ComplexF64, 1<<N); b=b/norm(b)
    phases, evec = eigen(U)
    ϕs = @. mod(angle(phases)/2/π, 1)
    return U, b, ϕs, evec
end

@testset "phaseest" begin
    # Generate a random matrix.
    N = 3
    V = rand_unitary(1<<N)

    # Initial Set-up.
    phases = rand(1<<N)
    ϕ = Int(0b111101)/(1<<6)
    phases[3] = ϕ
    signs = exp.(2π*im.*phases)
    U = V*Diagonal(signs)*V'
    b = V[:,3]

    # Define ArrayReg and U operator.
    M = 6
    reg1 = zero_state(M)
    reg2 = ArrayReg(b)
    UG = matblock(U)

    # circuit
    circuit = PEBlock(UG, M, N)

    # run
    reg = apply!(join(reg2, reg1), circuit)

    # measure
    res = breflect(measure(focus!(copy(reg), 1:M); nshots=10)[1]; nbits=M) / (1<<M)

    @test res ≈ ϕ
    @test apply!(reg, circuit |> adjoint) ≈ join(reg2, reg1)
end

@testset "phaseest, non-eigen" begin
    # Generate a random matrix.
    N = 3
    U, b, ϕs, evec = rand_phaseest_setup(N)

    # Define ArrayReg and U operator.
    M = 6
    reg1 = zero_state(M)
    reg2 = ArrayReg(b)
    UG = matblock(U);

    # run circuit
    reg= join(reg2, reg1)
    pe = PEBlock(UG, M, N)
    apply!(reg, pe)

    # measure
    bs, proj, amp_relative = projection_analysis(evec, focus!(reg, M+1:M+N))
    @test isapprox(ϕs, bfloat.(bs, nbits=M), atol=0.05)
end


text:
module Fronts

include("_Diff.jl")
using ._Diff: value_and_derivative, derivative

include("_Rootfinding.jl")
using ._Rootfinding: bracket_bisect

include("PorousModels/PorousModels.jl")

using OrdinaryDiffEq.DiffEqBase: NullParameters
using OrdinaryDiffEq: ODEFunction, ODEProblem, ODESolution
using OrdinaryDiffEq: ContinuousCallback, DiscreteCallback, terminate!
using OrdinaryDiffEq: RadauIIA5
import OrdinaryDiffEq

using ArgCheck: @argcheck
using StaticArrays: @SVector
using PCHIPInterpolation: Interpolator, integrate
using RecipesBase

"""
    solve(prob::DirichletProblem[; itol, maxiter, d_dϕb_hint]) -> Solution
    solve(prob::FlowrateProblem[; itol, ϕbtol, maxiter, b_hint]) -> Solution
    solve(prob::CauchyProblem) -> Solution

Solve the problem `prob`.

# Keyword arguments
- `itol=1e-3`: absolute tolerance for the initial condition.
- `ϕbtol=1e-6`: maximum tolerance for `ϕb`. Allows solving `FlowrateProblem`s with boundaries at `r=0`.
- `maxiter=100`: maximum number of iterations.
- `d_dϕb_hint`, `b_hint`: optional hints for the algorithms.

Type `\\phi<tab>` to obtain the `ϕ` symbol.

# Exceptions
This function throws an `SolvingError` if an acceptable solution is not found (within the
maximum number of iterations, if applicable). However, in situations where `solve` can determine
that the problem is "unsolvable" before the attempt to solve it, it will signal this by throwing a
`DomainError` instead. Other invalid argument values will raise `ArgumentError`s.

See also: [`Solution`](@ref), [`SolvingError`](@ref)
"""
function solve end

include("equations.jl")
include("boltzmann.jl")
include("odes.jl")
include("solution.jl")
include("problems.jl")
include("integration.jl")
include("shooting.jl")
include("exceptions.jl")
include("inverse.jl")

export Equation, DiffusionEquation, RichardsEquation, flux, isindomain
export TransformedFunction, d_dϕ, ∂_∂r, ∂_∂t, transform
export DirichletProblem, FlowrateProblem, CauchyProblem, monotonicity
export solve
export Solution, rb
export SolvingError
export inverse

end


text:
using Plots
using ImplicitEquations
using LaTeXStrings

f(x, y) = x^4 - x^2 * y + y^3
g(u, w) = f(u, u * w)

p1 = plot(Eq(f, 0), title = L"x^4 - x^2y + y^3 = 0")
p2 = plot(Eq(g, 0), title = L"u^3(u - w + w^3) = 0")
plot(
    p1,
    p2,
    xlabel = "x",
    ylabel = "y",
    xlim = [-1, 1],
    ylim = [-1, 1],
    aspect_ratio = :equal
)

max_stars_repo_name: terasakisatoshi/juliaExer

max_issues_repo_name: terasakisatoshi/juliaExer

max_forks_repo_name: terasakisatoshi/juliaExer


text:
using SpinMonteCarlo
using Printf

const ups = [0,1,2]
const Ls = [8,12,16]
const MCS = 8192
const Therm = MCS >> 3
const Ts = range(0.5, stop=2.0, length=16)

params = Dict{String, Any}[]
for update in ups
    for L in Ls
        for T in Ts
            push!(params,
                  Parameter("Model"=>XY, "Lattice"=>"square lattice",
                            "L"=>L, "T"=>T, "J"=>1.0,
                            "MCS"=>MCS, "Thermalization"=>Therm,
                            "Update Method"=> (update==0 ? local_update! :
                                               update==1 ? SW_update! : Wolff_update!),
                            "update"=>update,
                            "Verbose"=>true,
                           ))
        end
    end
end

obs = map(runMC, params)

const pnames = ["update", "L", "T"]
const onames = ["|Magnetization|",
                "|Magnetization|^2",
                "|Magnetization|^4",
                "Binder Ratio",
                "Magnetization x",
                "|Magnetization x|",
                "Magnetization x^2",
                "Magnetization x^4",
                "Binder Ratio x",
                "Magnetization y",
                "|Magnetization y|",
                "Magnetization y^2",
                "Magnetization y^4",
                "Binder Ratio y",
                "Susceptibility",
                "Connected Susceptibility",
                "Susceptibility x",
                "Connected Susceptibility x",
                "Susceptibility y",
                "Connected Susceptibility y",
                "Helicity Modulus x",
                "Helicity Modulus y",
                "Energy",
                "Energy^2",
                "Specific Heat",
                "MCS per Second",
                "Time per MCS",
               ]

const io = open("res-xy.dat", "w")
i=1
for pname in pnames
    println(io, "# \$$i : $pname")
    global i+=1
end
for oname in onames
    println(io, "# \$$i, $(i+1): $oname")
    global i+=2
end

for (p,o) in zip(params, obs)
    for pname in pnames
        print(io, p[pname], " ")
    end
    for oname in onames
        @printf(io, "%.15f %.15f ", mean(o[oname]), stderror(o[oname]))
    end
    println(io)
end


text:
module TestNewick

using Phylo
using Compat.Test

@testset "A few simple binary trees" begin
    @test length(nodeiter(parsenewick("((,),(,));"))) == 7
    @test ["where", "when it's good", "Not mine", "MyLeaf"] ⊆
        nodenameiter(parsenewick("""((MyLeaf,"when it's good"),
                                     ('Not mine',where));"""))
    tree = parsenewick("((MyLeaf[&Real=23,'Not real'={5,4}]:4.0,)Parent,(,));")
    branches = branchfilter(tree) do branch
        return src(branch) == "Parent" && dst(branch) == "MyLeaf"
    end
    @test length(branches) == 1
    @test getlength(first(branches)) ≈ 4.0
    @test getnoderecord(tree, "MyLeaf")["Real"] == 23
    @test 5 ∈ getnoderecord(tree, "MyLeaf")["Not real"]
    @test 4 ∈ getnoderecord(tree, "MyLeaf")["Not real"]
    if VERSION < v"0.7.0-"
        @test_warn "Tree ended but not finished" parsenewick("((,),(,));a")
    end
    @test_throws ErrorException parsenewick("((,),(,)));")
    @test_throws ErrorException parsenewick("((,),(,))")
    @test_throws ErrorException parsenewick("((,),(,);")
    @test_throws ErrorException parsenewick("((MyLeaf:-4.0,)Parent,(,));")
    tree = open(parsenewick, Phylo.path("H1N1.newick"))
    @test nleaves(tree) == 507
    @test ntrees(tree) == 1
    treex = open(parsenexus, "H1N1.trees")
    @test nleaves(tree) == nleaves(treex) ==
          nleaves(treex["TREE1"]) == nleaves(treex["TREE2"])
    @test ntrees(treex) == 2
end

@testset "A few simple polytomous trees" begin
    @test length(nodeiter(parsenewick("((,),(,,));", NamedPolytomousTree))) == 8
    @test ["where", "when it's good", "Not mine", "MyLeaf", "next"] ⊆
        nodenameiter(parsenewick("""((MyLeaf,"when it's good",next),
                                     ('Not mine',where));""", NamedPolytomousTree))
    tree = parsenewick("((MyLeaf:4.0,)Parent,());", NamedPolytomousTree)
    branches = branchfilter(tree) do branch
        return src(branch) == "Parent" && dst(branch) == "MyLeaf"
    end
    @test length(branches) == 1
    @test getlength(first(branches)) ≈ 4.0
    @test_throws ErrorException parsenewick("((,),(,)));", NamedPolytomousTree)
    @test_throws ErrorException parsenewick("((,),(,))", NamedPolytomousTree)
    @test_throws ErrorException parsenewick("((,),(,);", NamedPolytomousTree)
end

end


text:
using PooksoftAlphaVantageDataStore

# ok, so lets build by user object -
my_current_dir = pwd()   # where am I?
path_to_config_file = my_current_dir*"/test/configuration/Configuration.json"

# build the api user model -
user_model_result = build_api_user_model(path_to_config_file)
user_model = checkresult(user_model_result)

# now that we have the userModel - lets populate the request dictionary -
required_api_keys = [
        "symbol"        ;   # ticker symbol -
        "function"      ;   # what function are we doing?
        "apikey"        ;   # the API key
]
requestDictionary = Dict{String,Any}()
requestDictionary["apikey"] = user_model.alphavantage_api_key
requestDictionary["symbol"] = "ALLY"
requestDictionary["function"] = "EARNINGS"
api_call_result = execute_company_earnings_api_call(requestDictionary)
df = checkresult(api_call_result)


text:
function (sqrt){G<:Grasp,R<:Real}(a::Rvl{G,R})
    lo = (sqrt)(a.lo, RoundDown)
    hi = (sqrt)(a.hi, RoundUp)
    Rvl{G,R}(lo, hi)
end


text:
# julia> include("/Users/ian/jd/10-19 Code/12 Julia/12.03 PE/pe2.jl")

fibs = [1,2]

while (fibs[end-1] + fibs[end]) < 4*10^6
    push!(fibs, fibs[end-1] + fibs[end])
end

println("ans = $(sum([i for i in fibs if i%2==0]))")

max_stars_repo_path: learning exercises/pe/pe2.jl

max_stars_repo_name: ian-double-u/NNfromScratch

max_issues_repo_path: learning exercises/pe/pe2.jl

max_issues_repo_name: ian-double-u/NNfromScratch

max_forks_repo_path: learning exercises/pe/pe2.jl

max_forks_repo_name: ian-double-u/NNfromScratch


text:
struct NameDataset
    firstyear::Int32
    lastyear::Int32
    namesdict::Dict

    NameDataset(fy, ly) = new(Int32(fy), Int32(ly), Dict())
end

firstyear(nds::NameDataset) = nds.firstyear
lastyear(nds::NameDataset) = nds.lastyear

_years(nds::NameDataset) = (firstyear(nds), lastyear(nds))

getindex(nds::NameDataset, key::AbstractString) = Base.getindex(nds.namesdict, lowercase(key))
haskey(nds::NameDataset, name::AbstractString) = Base.haskey(nds.namesdict, lowercase(name))

# check to prevent overwriting
function addname(nds::NameDataset, name::AbstractString)
    haskey(nds, name) && return true
    nds.namesdict[lowercase(name)] = (female=spzeros(Int32, nds.lastyear), male=spzeros(Int32, nds.lastyear))
end

function _update_genders!(nds::NameDataset, name, year, gender, n)
    firstyear, lastyear = _years(nds)
    firstyear <= year <= lastyear || throw(ArgumentError("This dataset only has years $firstyear to $lastyear"))
    name = lowercase(name)
    addname(nds, name)
    nds[name][gender][year] > 0 && error("duplicate $gender entry for $name $year")

    nds[name][gender][year] = n
end

struct RawDataSet{T} end

RawDataSet(s::String) = RawDataSet{Symbol(s)}()
RawDataSet(s::Symbol) = RawDataSet{s}()

parsedataset(datfolder, ::RawDataSet) = throw(ArgumentError("Unknown Dataset"))

function parsedataset(datfolder, ::RawDataSet{:USCensus})
    nds = NameDataset(1880, 2018)

    for y in filter(f-> occursin(r"^yob\d{4}", f), readdir(datfolder))
        year = match(r"yob(\d{4})\.txt", y).captures[1] |> x -> parse(Int32, x)

        for line in eachline(joinpath(datfolder, y))
            (name, gender, n) = split(line, ',')

            gender in ("M", "F") || error("Gender $gender not recognized")
            gender == "F" ? gender = :female : gender = :male

            n = parse(Int32, n)

            _update_genders!(nds, name, year, gender, n)
        end

    end
    return nds
end


text:
import PowerModels
import Ipopt
import JuMP
import JSON
# using .PP2PM

function run_pf(json_net, method, print_level=5, max_iter=Nothing)
    # pm = PP2PM.load_pm_from_json(json_path)
    if max_iter == Nothing
        solver = JuMP.with_optimizer(Ipopt.Optimizer, print_level=print_level)
    else
        solver = JuMP.with_optimizer(Ipopt.Optimizer, print_level=print_level)
    end
    # solver = JuMP.with_optimizer(GLPK.Optimizer, tm_lim = 60.0, msg_lev = GLPK.OFF)
    data = JSON.parse(json_net)
    if method == "ac"
        result = PowerModels.run_ac_pf(data, solver, setting = Dict("output" => Dict("branch_flows" => true)))
    elseif method == "dc"
        result = PowerModels.run_dc_pf(data, solver, setting = Dict("output" => Dict("branch_flows" => true)))
    else
        throw(DomainError(method), "method not recognized")
    end
    return JSON.json(result), JSON.json(data)
end


text:
# in place adjoint of square matrices

const _adjointblocksize = 32
function _adjoint!(A::AbstractMatrix, r = axes(A,1), c = axes(A,2))
    @assert r == c
    if length(r) > _adjointblocksize
        r1, r2 = _split(r)
        _adjoint!(A, r1, r1)
        _adjoint!(A, r2, r2)
        _swapadjoint!(A, r1, r2)
    else
        @inbounds for j in c
            @simd for i = first(r):(j-1)
                A[i,j], A[j,i] = conj(A[j,i]), conj(A[i,j])
            end
            A[j,j] = conj(A[j,j])
        end
    end
    return A
end

function _swapadjoint!(A::AbstractMatrix, r, c)
    if length(r) > length(c) && length(r) > _adjointblocksize
        r1, r2 = _split(r)
        _swapadjoint!(A, r1, c)
        _swapadjoint!(A, r2, c)
    elseif length(c) > _adjointblocksize
        c1, c2 = _split(c)
        _swapadjoint!(A, r, c1)
        _swapadjoint!(A, r, c2)
    else
        @inbounds for j in c
            @simd for i in r
                A[i,j], A[j,i] = conj(A[j,i]), conj(A[i,j])
            end
        end
    end
    return A
end

function _split(r::AbstractUnitRange{Int})
    i = first(r)
    f = last(r)
    m = (i+f) >> 1
    return (i:m, m+1:f)
end


text:
# convert_a3d

# Method that allows federation by setting the `output_format`  in the Stanmodel().

"""

# convert_a3d

Convert the output file created by cmdstan to the shape of choice.

### Method
```julia
convert_a3d(a3d_array, cnames, ::Val{Symbol})
```
### Required arguments
```julia
* `a3d_array::Array{Float64}(n_draws, n_variables, n_chains`      : Read in from output files created by cmdstan
* `cnames::Vector{AbstractString}`                                                 : Monitored variable names
```

### Optional arguments
```julia
* `::Val{Symbol} = :array`                                                                             : Output format
```
Method called is based on the output_format defined in the stanmodel, e.g.:

   stanmodel = Stanmodel(num/_samples=1200, thin=2, name="bernoulli",
   model=bernoullimodel, output_format=:array);

Current formats supported are:

1. :array (a3d_array format, the default for CmdStan)
2. :namedarray (NamedArrays object)
3. :dataframe (DataFrames object)
4. :mambachains (Mamba.Chains object)
5. :mcmcchain (TuringLang/MCMCChain.Chains object, will be deprecated)
6. :mcmcchain (TuringLang/MCMCChains.Chains object)

Options 3 through 6 are respectively provided by the packages StanDataFrames,
StanMamba, StanMCMCChain and StanMCMCChains.
```

### Return values
```julia
* `res`                       : Draws converted to the specified format.
```
"""
convert_a3d(a3d_array, cnames, ::Val{:array}) = a3d_array

convert_a3d(a3d_array, cnames, ::Val{:namedarray}) =
  [NamedArray(a3d_array[:,:,i], (collect(1:size(a3d_array, 1)), Symbol.(cnames)))
    for i in 1:size(a3d_array, 3)]


text:
@testset "Datums and ellipsoids" begin
    # Check the ellipsoids
    @test ellipsoid(WGS84) == wgs84_ellipsoid
    @test ellipsoid(NAD27) == clarke1866
    @test ellipsoid(NAD83) == grs80
    @test ellipsoid(OSGB36) == airy1830
    @test ellipsoid(GDA94) == grs80

    # Check transverse-Mercator pre-calculations
    @test Geodesy.TransverseMercator(WGS84) == Geodesy.wgs84_tm
    @test Geodesy.TransverseMercator(NAD27) == Geodesy.clarke1866_tm
    @test Geodesy.TransverseMercator(OSGB36) == Geodesy.airy1830_tm
    @test Geodesy.TransverseMercator(GDA94) == Geodesy.grs80_tm

    # The LLAfromECEF aren't pre-cached, as yet. Probably should just create a
    # large ellispoid thingie for all of them to share...
    @test (LLAfromECEF(wgs84); true)
    @test (LLAfromECEF(nad27); true)
    @test (LLAfromECEF(nad83); true)
    @test (LLAfromECEF(osgb36); true)
    @test (LLAfromECEF(gda94); true)
end


text:
"""
    two_sum(a, b)

Computes `hi = fl(a+b)` and `lo = err(a+b)`.
"""
@inline function two_sum(a::T, b::T) where {N, F<:Base.IEEEFloat, T<:NTuple{N,F}}
    hi = a .+ b
    v  = hi .- a
    lo = (a .- (hi .- v)) .+ (b .- v)
    return hi, lo
end

"""
    two_sum(a, b, c)

Computes `hi = fl(a+b+c)` and `lo = err(a+b+c)`.
"""
function two_sum(a::T,b::T,c::T) where {N, F<:Base.IEEEFloat, T<:NTuple{N,F}}
    s, t   = two_sum(b, c)
    hi, u  = two_sum(a, s)
    lo     = u .+ t
    hi, lo = two_hilo_sum(hi, lo)
    return hi, lo
end

"""
    three_sum(a, b, c, d)

Computes `hi = fl(a+b+c+d)` and `md = err(a+b+c+d), lo = err(md)`.
"""
function three_sum(a::T,b::T,c::T,d::T) where {N, F<:Base.IEEEFloat, T<:NTuple{N,F}}
    t0, t1 = two_sum(a ,  b)
    t0, t2 = two_sum(t0,  c)
    hi, t3 = two_sum(t0,  d)
    t0, t1 = two_sum(t1, t2)
    hm, t2 = two_sum(t0, t3) # here, t0 >= t3
    ml     = t1 .+ t2
    return hi, hm, ml
end

"""
    two_sum(a, b, c, d)

Computes `hi = fl(a+b+c+d)` and `lo = err(a+b+c+d)`.
"""
function two_sum(a::T,b::T,c::T,d::T) where {N, F<:Base.IEEEFloat, T<:NTuple{N,F}}
    t0, t1 = two_sum(a ,  b)
    t0, t2 = two_sum(t0,  c)
    hi, t3 = two_sum(t0,  d)
    t0, t1 = two_sum(t1, t2)
    lo     = t0 .+ t3
    return hi, lo
end

"""
    two_diff(a, b)

Computes `s = fl(a-b)` and `e = err(a-b)`.
"""
@inline function two_diff(a::T, b::T) where {N, F<:Base.IEEEFloat, T<:NTuple{N,F}}
    hi = a .- b
    a1 = hi .+ b
    b1 = hi .- a1
    lo = (a .- a1) .- (b .+ b1)
    return hi, lo
end
"""
    three_diff(a, b, c)

Computes `s = fl(a-b-c)` and `e1 = err(a-b-c), e2 = err(e1)`.
"""
function three_diff(a::T,b::T,c::T) where {N, F<:Base.IEEEFloat, T<:NTuple{N,F}}
    negb = (-).(b)
    s, t = two_diff(negb, c)
    x, u = two_sum(a, s)
    y, z = two_sum(u, t)
    x, y = two_hilo_sum(x, y)
    return x, y, z
end

"""
    two_square(a)

Computes `s = fl(a*a)` and `e = err(a*a)`.
"""
@inline function two_square(a::T) where {N, F<:Base.IEEEFloat, T<:NTuple{N,F}}
    hi = a .* a
    neghi = (-).(hi)
    lo = (fma).(a, a, neghi)
    hi, lo
end

"""
    two_prod(a, b)

Computes `s = fl(a*b)` and `e = err(a*b)`.
"""
@inline function two_prod(a::T, b::T) where {N, F<:Base.IEEEFloat, T<:NTuple{N,F}}
    hi = a .* b
    neghi = (-).(hi)
    lo = (fma).(a, b, neghi)
    hi, lo
end

#=
   three_fma algorithm from
   Sylvie Boldo and Jean-Michel Muller
   Some Functions Computable with a Fused-mac
=#

"""
   three_fma(a, b, c)

Computes `s = fl(fma(a,b,c))` and `e1 = err(fma(a,b,c)), e2 = err(e1)`.
"""
function three_fma(a::T, b::T, c::T) where {N, F<:Base.IEEEFloat, T<:NTuple{N,F}}
     x = (fma).(a, b, c)
     y, z = two_prod(a, b)
     t, z = two_sum(c, z)
     t, u = two_sum(y, t)
     y = ((t .- x) .+ u)
     y, z = two_hilo_sum(y, z)
     infs = isinf.(x)
     if any(infs)
        ys = [y...]
        zs = [z...]
        ys[infs] .= zero(F)
        zs[infs] .= zero(F)
        y = (ys...,)
        z = (zs...,)
     end
     return x, y, z
end

# with arguments sorted by magnitude

"""
    two_hilo_sum(a, b)

*unchecked* requirement `|a| ≥ |b|`

Computes `hi = fl(a+b)` and `lo = err(a+b)`.
"""
@inline function two_hilo_sum(a::T, b::T) where {N, F<:Base.IEEEFloat, T<:NTuple{N,F}}
    hi = a .+ b
    lo = b .- (hi .- a)
    return hi, lo
end

"""
    two_lohi_sum(a, b)

*unchecked* requirement `|b| ≥ |a|`

Computes `hi = fl(a+b)` and `lo = err(a+b)`.
"""
@inline function two_lohi_sum(a::T, b::T) where {N, F<:Base.IEEEFloat, T<:NTuple{N,F}}
    hi = b .+ a
    lo = a .- (hi .- b)
    return hi, lo
end

"""
    two_hilo_diff(a, b)

*unchecked* requirement `|a| ≥ |b|`

Computes `hi = fl(a-b)` and `lo = err(a-b)`.
"""
@inline function two_hilo_diff(a::T, b::T) where {N, F<:Base.IEEEFloat, T<:NTuple{N,F}}
    hi = a .- b
    lo = (a .- hi) .- b
    hi, lo
end

"""
    two_lohi_diff(a, b)

*unchecked* requirement `|b| ≥ |a|`

Computes `hi = fl(a-b)` and `lo = err(a-b)`.
"""
@inline function two_lohi_diff(a::T, b::T) where {N, F<:Base.IEEEFloat, T<:NTuple{N,F}}
    hi = b .- a
    lo = (b .- hi) .- a
    hi, lo
end

"""
    three_hilo_diff(a, b, c)

*unchecked* requirement `|a| ≥ |b| ≥ |c|`

Computes `x = fl(a-b-c)` and `y = err(a-b-c), z = err(y)`.
"""
function three_hilo_diff(a::T,b::T,c::T) where {N, F<:Base.IEEEFloat, T<:NTuple{N,F}}
    negc = (-).(c)
    s, t = two_hilo_diff(b, negc)
    x, u = two_hilo_sum(a, s)
    y, z = two_hilo_sum(u, t)
    x, y = two_hilo_sum(x, y)
    return x, y, z
end

"""
    three_lohi_diff(a, b, c)

*unchecked* requirement `|c| ≥ |b| ≥ |a|`

Computes `x = fl(a-b-c)` and `y = err(a-b-c), z = err(y)`.
"""
function three_lohi_diff(c::T,b::T,a::T) where {N, F<:Base.IEEEFloat, T<:NTuple{N,F}}
    negc = (-).(c)
    s, t = two_hilo_diff(b, negc)
    x, u = two_hilo_sum(a, s)
    y, z = two_hilo_sum(u, t)
    x, y = two_hilo_sum(x, y)
    return x, y, z
end


text:
function boxplot(data, outputfile="", model_names=0:length(data)-1)
    # data must be an iterble object of lists of numbers

    @assert length(model_names) == length(data)
    n = length(model_names)
    for vals in data
        @assert length(data[1]) == length(vals)
    end
    k = length(data[1])

    df = DataFrame(
        ids = collect(Base.flatten([1:length(vals) for vals in data])),
        value = collect(Base.flatten(data)),
        model = collect(Base.flatten([[name for i in 1:k] for name in model_names]))
    )
    if outputfile == ""
        plot(df, x="model", y="value", Geom.boxplot)
    else
        draw(SVG(outputfile*".svg", n*5cm, 15cm), plot(df, x="model", y="value", Geom.boxplot))
    end
end


text:
function bathroom(p::Player)
    if p.diarrhea == true
        diarrhea(p)
        return
    end
    print_dict(p, "spacer")
    print_dict(p, "bathroom_0")
    s = readline(p.socket)
    if s == "w"
        print_dict(p, "bathroom_1")
    elseif s == "l"
        print_dict(p, "bathroom_2")
    else
        print_dict(p, "bathroom_3")
        return
    end

    print_dict(p, "bathroom_4")

    s = readline(p.socket)
    if s == "v"
        print_dict(p, "spacer")
        use_cryptomat(p)
    elseif s == "r"
        p.status = reception
    end

    print_dict(p, "bathroom_5")
end

function code(p::Player)
    path = string("data/.cheesecode_", p.dimension)
    if !(isfile(path))
        code = "CASINO_ROYALE"
    else
        f = open_file_try(path, "r")
        code = read(f, String)
        close(f)
    end
    return md5(code * string(p.dimension))
end

function diarrhea(p::Player)
    print_dict(p, "bathroom_diarrhea")
    while true
        user_input = readline(p.socket)
        if user_input == "r"
            print_dict(p, "bathroom_stall_walls")
            write(p.socket, JSON.json(code(p)) * "\n")
            break
        elseif user_input == "l"
            break
        end
    end
    print_dict(p, "bathroom_diarrhea_leave")
    while true
        user_input = readline(p.socket)
        if user_input == "c"
            use_cryptomat(p)
            break
        elseif user_input == "r"
            break
        end
    end

end


text:
printstyled("\n\n================================================================\n", color=:yellow)
printstyled("
   #    ######  ######  ####### #     #          ####### #######
  # #   #     # #     # #     # #  #  #             #    #
 #   #  #     # #     # #     # #  #  #             #    #
#     # ######  ######  #     # #  #  #   #####     #    #####
####### #   #   #   #   #     # #  #  #             #    #
#     # #    #  #    #  #     # #  #  #             #    #
#     # #     # #     # #######  ## ##              #    #######\n\n", color=:yellow)
printstyled("ARROW: Agile RestoRation of Optical Wavelengths\n\n", bold=true, color=:blue)
printstyled("Traffic Engineering Simulator for ARROW\n\n", bold=true, color=:blue)
printstyled("[Paper]: Z. Zhong, M. Ghobadi, A. Khaddaj, J. Leach, Y. Xia, Y. Zhang, ARROW: Restoration-Aware Traffic Engineering, ACM SIGCOMM, 2021.\n", color=:blue)
printstyled("[Website]: http://arrow.csail.mit.edu\n", color=:blue)
printstyled("[Code Contributors]: zhizhenz@mit.edu, alaakh@mit.edu\n", color=:blue)
printstyled("[Code Release Date]: October, 2021\n", color=:blue)
printstyled("==================================================================\n\n", color=:yellow)

max_stars_repo_name: hipersys-team/arrow

max_issues_repo_name: arrow-bow/arrow

max_forks_repo_name: arrow-bow/arrow


text:
header("TOOLSPECIFIC")
@testset "parse tools" begin
    str1 = (tool="JARP", version="1.2", str = """
 <toolspecific tool="JARP" version="1.2">
        <FrameColor>
          <value>java.awt.Color[r=0,g=0,b=0]</value>
        </FrameColor>
        <FillColor>
          <value>java.awt.Color[r=255,g=255,b=255]</value>
        </FillColor>
 </toolspecific>
""", contentparse = (c) -> begin end)

    str2 = (tool="de.uni-freiburg.telematik.editor", version="1.0", str = """
 <toolspecific tool="de.uni-freiburg.telematik.editor" version="1.0">
     <visible>true</visible>
 </toolspecific>
""", contentparse = (c) -> begin end)

    str3 = (tool="petrinet3", version="1.0", str = """
 <toolspecific tool="petrinet3" version="1.0">
     <placeCapacity capacity="0"/>
 </toolspecific>
""", contentparse = (c) -> begin end)

str4 = (tool="org.pnml.tool", version="1.0", str = """
 <toolspecific tool="org.pnml.tool" version="1.0">
    <tokengraphics>
         <tokenposition x="-9" y="-2"/>
         <tokenposition x="2"  y="3"/>
     </tokengraphics>
 </toolspecific>
""", contentparse = (c) -> begin end)

    str5 = (tool="org.pnml.tool", version="1.0", str = """
 <toolspecific tool="org.pnml.tool" version="1.0">
    <tokengraphics>
         <tokenposition x="-9" y="-2"/>
         <tokenposition x="2"  y="3"/>
     </tokengraphics>
     <visible>true</visible>
 </toolspecific>
""", contentparse = (c) -> begin end)

    @testset for s in [str1, str2, str3, str4, str5]
        @show s
        n = parse_node(root(EzXML.parsexml(s.str)); reg=PNML.IDRegistry())
        printnode(n)

        @test typeof(n) <: ToolInfo
        @test xmlnode(n) isa Maybe{EzXML.Node}
        @test_call xmlnode(n)
        @test n.toolname == s.tool
        @test name(n) == s.tool
        @test n.version == s.version
        @test version(n) == s.version

        @test get_toolinfo(n, s.tool, s.version) isa ToolInfo
        @test get_toolinfo(n, s.tool, s.version) == n # Is identity on scalar
        @test get_toolinfo(n, s.tool) == n
        @test get_toolinfo(n, s.tool, r"^.*$") == n
        @test get_toolinfo(n, Regex(s.tool), r"^.*$") == n
        @test get_toolinfo(n, Regex(s.tool)) == n

        @test_call get_toolinfo(n, s.tool, s.version)

        #s.contentparse(n.infos) #TODO
        # contentparse should handle a vector or scalar of well-formed xml.

        @test n.infos isa Vector{AnyElement}
        @test PNML.infos(n) isa Vector{AnyElement}
        foreach(n.infos) do toolinfo
            @test toolinfo isa AnyElement
            # Content may optionally attach its xml.
            @test !PNML.has_xml(toolinfo) || xmlnode(toolinfo) isa Maybe{EzXML.Node}
        end
        println()
    end
    @testset "combined" begin
        println("combined tools")
        str = """
<?xml version="1.0"?>
<pnml xmlns="http://www.pnml.org/version-2009/grammar/pnml">
 <net id="net0" type="pnmlcore">
 <page id="page0">
 <place id="place0">
    $(str1.str)
    $(str2.str)
    $(str3.str)
    $(str4.str)
    $(str5.str)
 </place>
 </page>
 </net>
</pnml>
"""
        model = parse_str(str)
        @show model

        page = firstpage(first_net(model))
        @test_call firstpage(first_net(model))

        @test !has_tools(page)
        @test_call has_tools(page)

        place = first(page.places)
        @test has_tools(place)
        @test_call has_tools(place)
        t = tools(place)
        @test_call tools(place)
        @test t isa Vector{ToolInfo}
        @test length(t) == 5

        for ti in t
            @test ti isa ToolInfo
        end
        for (i,s) in enumerate([str1, str2, str3, str4, str5])
            ti = get_toolinfo(t, s.tool, s.version)
            @test ti isa ToolInfo
            @test PNML.name(t[i]) == PNML.name(ti)
            @test PNML.version(t[i]) == PNML.version(ti)
            @test PNML.name(t[i]) == s.tool
            @test PNML.version(t[i]) == s.version

            @test_call PNML.name(t[i])
            @test_call PNML.version(t[i])

            @test typeof(t[i].infos) == typeof(ti.infos)

            for y in zip(t[i].infos, ti.infos)
                @test tag(y[1]) == tag(y[2])
                @test y[1].dict == y[2].dict
            end

            x = parse_node(root(EzXML.parsexml(s.str)); reg=PNML.IDRegistry())

            @test typeof(t[i].infos) == typeof(x.infos)
            for y in zip(t[i].infos, x.infos)
                @test tag(y[1]) == tag(y[2])
                @test y[1].dict == y[2].dict
            end
        end
    end
end


text:
module Gadgetron
import Sockets

using PartialFunctions

include("MRD.jl")
include("Types.jl")
include("StreamAlgorithms.jl")

export listen, connect,register_type, MRD, close, Stream

@Base.enum fixed_message_ids::UInt16 FILENAME = 1 CONFIG = 2 HEADER = 3 CLOSE = 4 TEXT = 5 QUERY =6 RESPONSE = 7 ERROR = 8
const message_ids = Dict{UInt16,Type}([(1008,MRD.Acquisition),(1022,MRD.Image),(1026, MRD.Waveform,1026),(1050,Types.Bucket),(1051, Types.Bundle)])
const message_types = Dict(value => key for (key,value ) in message_ids)

"""
	register_type!(type::Type, message_id::UInt16)
Registers a type with the Gadgetron system.  Types must implement MRD.write and MRD.read
"""
function register_type!(type::Type, message_id::UInt16)
	message_ids[message_id] = type
	message_types[type] = message_id
end

"""
	MRDChannel
An Abstract channel of MRD data.

# Fields
- config Configuration of the channel. Currently a string but may change in the future
- header The dataset header
- channel The channel containing the message
"""
mutable struct MRDChannel <: AbstractChannel{Any}
	config
	header::MRD.MRDHeader
	channel::AbstractChannel
	function MRDChannel(config, header, channel)
		conn = new(config, header,channel)
		finalizer(close,conn)
	end
end

Base.iterate(m::MRDChannel, state...) = iterate(m.channel, state...)

"""
	MRDRemoteConnection
	An connection to an external endpoint using the MRD protocol
"""
mutable struct MRDRemoteConnection <: AbstractChannel{Any}
	socket::IO
	inputchannel::Channel
	outputchannel::Channel
	outputtask::Ref{Task}
	"Creates a remote connection from a newly opened socket "
	function MRDRemoteConnection(socket::IO; buffer::Real=Inf, spawn=true)

		input = Channel(buffer, spawn=spawn) do c
			message_id = read_id(socket)
			while message_id != UInt16(CLOSE)
				message_type = message_ids[message_id]
				push!(c, MRD.read(socket, message_type))
				message_id = read_id(socket)
			end
		end

		outputtask = Ref{Task}()
		output = Channel(buffer, spawn=spawn, taskref=outputtask) do c
			for msg in c
				message_id = message_types[typeof(msg)]
				write(socket,message_id)
				MRD.write(socket, msg)
			end
		end

		new(socket, input, output, outputtask )
	end

end

Base.take!(m::MRDRemoteConnection) = take!(m.inputchannel)
Base.fetch(m::MRDRemoteConnection) = fetch(m.inputchannel)
Base.isready(m::MRDRemoteConnection) = isread(m.inputchannel)
Base.wait(m::MRDRemoteConnection) = wait(m.inputchannel)
Base.put!(m::MRDRemoteConnection,v ) = put!(m.outputchannel,v)

"Closes the connection by sending a close message and closing the network socket"
function Base.close(m::MRDRemoteConnection)
	close(m.outputchannel)
	wait(m.outputtask[])
	if typeof(m.socket) == Sockets.TCPSocket
		if iswritable(m.socket)
			write(m.socket,CLOSE)
			Sockets.close(m.socket)
		end
	end
end

Base.iterate(m::MRDRemoteConnection, state...) = iterate(m.inputchannel, state...)

"Creates an MRD channel from a newly opened socket"
MRDChannel(socket; kwargs...) = MRDChannel(read_config(socket),read_header(socket), MRDRemoteConnection(socket; kwargs...))

"Waits for a connection on the given address and port "
listen(addr, port::Integer) = Sockets.listen(Sockets.IPv6(0),addr,port) |> Sockets.accept |> MRDChannel
function listen(port::Integer; kwargs...)
	socket = Sockets.listen(Sockets.IPv6(0),port) |> Sockets.accept
	return MRDChannel(socket; kwargs...)
end

"Waits for a connection on the given adress and port"
function listen(recon_function::Function, port::Integer; kwargs... )
	connection = listen(port; kwargs...)
	try
		recon_function(connection)
	finally
		close(connection)
	end
end

"Similar to listen, but will keep serving the function"
function serve_forever(recon_function::Function, port::Integer; kwargs...)
	server =Sockets.listen(Sockets.IPv6(0),port)
	while true
		socket = Sockets.accept(server)
		channel = MRDChannel(socket; kwargs...)
		try
			recon_function(channel)
		finally
			close(channel)
		end
	end
end

"Connect to an open MRD server "
connect(addr, port::Integer) = Sockets.connect(addr,port) |> MRDChannel

read_id(x) = Sockets.read(x,UInt16)

Base.put!(c::MRDChannel,v) = put!(c.channel,v)
Base.take!(c::MRDChannel) = take!(c.channel)
Base.close(c::MRDChannel) = close(c.channel)
Base.fetch(c::MRDChannel) = fetch(c.channel)
Base.isready(c::MRDChannel) = isread(c.channel)
Base.wait(c::MRDChannel) = wait(c.channel)

"Reads a config message from a socket "
function read_config(socket::IO)
	message_id = read_id(socket)
	@assert message_id == UInt16(CONFIG) "Expected CONFIG message id, received $message_id"
	return MRD.read_string(socket)
end

"Reads a header message form a socket "
function read_header(socket::IO)
	message_id = read_id(socket)
	@assert message_id == UInt16(HEADER) "Expected HEADER message id, received $message_id"
	return MRD.read_string(socket) |> MRD.MRDHeader
end

include("Main.jl")
include("Default.jl")

end


text:
function run(s)
    lines = split(s, "\n")
    constraints = [(idx, parse(Int, el)) for (idx, el) in enumerate(split(lines[2], ",")) if el != "x"]
    ts::Int = 0
    jump::Int = 1
    current_idx::Int = 1
    max_idx = length(constraints)
    while true
        offset, mul = constraints[current_idx]
        if (ts + (offset - 1)) % mul == 0
            current_idx += 1
            jump *= mul
        end
        if current_idx > max_idx
            return ts
        end
        ts += jump
    end
    return 0
end

#########################################

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

main()


text:
#=
# 1. Basic potential flow problem

This part introduces how to solve the basic potential flow problem in this package. We consider here a staggered, Cartesian grid with uniform cell size and of infinite extent. The basic (unbounded) potential flow problem is expressed as

$\mathsf{Ls} = -\mathsf{w},$

where $\mathsf{L}$ is the discrete 5-point Laplacian operator, $\mathsf{s}$ is the discrete streamfunction, and $\mathsf{w}$ is the discrete vorticity. Following the vortex-in-cell approach, the discrete vorticity is obtained by regularizing the vorticity from the $N_v$ vortex elements onto the cell vertices,

$\mathsf{w}_{i, j}=\sum_{q=1}^{N_{v}} \frac{1}{\Delta x} \Gamma_{v, q} d\left(\frac{\mathsf{x}_{i}-X_{q}}{\Delta x}\right) d\left(\frac{\mathsf{y}_{j}-Y_{q}}{\Delta x}\right)$

where $d$ is the $M_{4}'$ interpolation kernel, and $(X_{q},Y_{q})$ and $\Gamma_{v,q}$ are the position and strength of the $q$th vortex element.
=#

#=
## Flow around a point vortex
Now we will solve this discrete potential flow problem using `GridPotentialFlow.jl` to obtain the streamfunction field and velocity field around a single point vortex.
=#

# The first step in creating *any* `GridPotentialFlow` model is to create a `PhysicalGrid` to discretize the domain.
using GridPotentialFlow
Δx = 0.01
Lx = 2.0
xlim = (-Lx/2,Lx/2)
ylim = (-Lx/2,Lx/2)
g = PhysicalGrid(xlim,ylim,Δx);

# The second step is to create our point vortex using `Vortex`.
v = Vortex(0.0,0.0,1.0);

# Now we can create a `VortexModel` using the grid and an array containing the point vortex.
model = VortexModel(g,vortices=[v]);

# The discrete streamfunction `s` is then obtained using `streamfunction`.
s = streamfunction(model);
using Plots
plot(s,g,xlabel="x",ylabel="y")

# The function `streamfunction` returns a `Nodes` array. If we want to perform differential calculus operations this data, we can use the methods of the  `CartesianGrids` package. For example, we can easily obtain the velocity field from the streamfunction field using the `curl` operation.
q = curl(s);
plot(q,g,xlabel="x",ylabel="y")

#=
## Accuracy of the discretized Poisson equation
To verify that the discretization technique of the Poisson equation is second-order accurate, we perform a mesh refinement analysis for a flow consisting of point vortices of random strenght that are randomly positioned on the lower-left quadrant of a square domain.
=#

# The error is calculated by comparing the numerically calculated streamfunction field with the exact solution for the streamfunction field. We therefore need a function to calculate the exact streamfunction.
function ψ_vortex!(ψ::Nodes{Dual,nx,ny},vortex::Vortex,g::PhysicalGrid) where {nx,ny}
    x,y = coordinates(ψ,g)
    for i in 2:nx-1, j in 2:ny-1
        r = sqrt((x[i]-vortex.x)^2+(y[j]-vortex.y)^2)
        ψ[i,j] = ψ_vortex(r,vortex.Γ)
    end
end

function ψ_vortex(r::Real,Γ::Real)
    return -Γ/(2π)*log(r)
end

# We create four vortices with random strenghts, randomly positioned in the lower-left quadrant of the domain.
nv = 4;
vl = Vortex.(-Lx/4 .+ 0.4*Lx*(rand(nv).-0.5),-Lx/4 .+ 0.4*Lx*(rand(nv).-0.5),0.5*rand(nv).+0.5);

# Next, we create a series of grids, with each grid doubling the number of grid points of the previous grid in each direction.
grids = [PhysicalGrid(xlim,ylim,Lx/(nx-2)) for nx in [2^p for p in 5:9]];

# The error is calculated as $\epsilon = \Vert \psi(\mathsf{x},\mathsf{y})/\Delta x - \mathsf{s} \Vert_2 / \Vert \mathsf{s} \Vert_2$, for which we use the `norm` function.
using LinearAlgebra: norm

# We now loop over the grids and calculate the error.
errors = []
gridspacings = []
for g in grids
    ## LGF
    local model = VortexModel(g,vortices=vl)
    s_lgf = streamfunction(model);

    ## Exact solution
    s_exact = Nodes(Dual,size(g))
    s_temp = Nodes(Dual,size(g))
    for v in vl
        ψ_vortex!(s_temp,v,g)
        s_exact += s_temp
    end

    ## Bring to same reference level. A constant value can be added/subtracted to any potential flow solution
    s_lgf .-= s_lgf[g.I0[1],2];
    s_exact .-= s_exact[g.I0[1],2];

    error = s_lgf-s_exact
    idx = g.I0[1] + 1 : g.N[1] - 1 # Only look at top right corner
    push!(errors,norm(error[idx,idx])/norm(s_lgf[idx,idx]))
    push!(gridspacings,g.Δx)
end

# And finally, we create a log-log plot of the error versus the grid spacing.
firstorder=(errors[end]/gridspacings[end].*gridspacings).^1
secondorder=(sqrt(errors[end])/gridspacings[end].*gridspacings).^2
p=plot(gridspacings,errors,xaxis=:log,yaxis=:log,marker=:circle,lab="error",xlabel="dx",legend=:bottomright,title="Error")
plot!(gridspacings,firstorder,lab="1st order",linestyle=:dot,linecolor=:black)
plot!(gridspacings,secondorder,lab="2nd order",linestyle=:dash,linecolor=:black)

#=
## Corotating point vortices
Now we will try advancing a vortex model in time. The simplest unsteady vortex model consists of two point vortices. When both point vortices have the same strength, they will rotate around each other on a trajectory that is easy to describe analytically. In this example, we will compare the analytical and simulated trajectories during one revolution.
=#

# We create two vortices at a distance $d$ from each other and give them a strength $\Gamma$
d = Lx/2
Γ = 1
v1 = Vortex(d/2,0.0,Γ);
v2 = Vortex(-d/2,0.0,Γ);

# We can analytically determine the time $T$ it takes for these vortices to complete one revolution around their centroid.
Vθ = Γ/(2*π*d); # Analytical tangential velocity
T = π*d/Vθ # Analytical period

# Let's now create a vortex model with the two point vortices.
model = VortexModel(g,vortices=[v1,v2]);

# If we update the positions of the point vortices repeatedly to simulate the model advancing in time from $t=0$ to $t=T$, we can check if they end up again at their original positions.
tspan = (0.0,T);

# To step in time, we can simply update the position of the $q$th vortex as $X^{n+1}_q = X^{n}_q + Δt Ẋ^{n}_q$ (forward Euler) in a for-loop. Alternatively, we can make use of the `OrdinaryDiffEq.jl` package and use, for example, their fourth-order Runge-Kutta time stepping scheme.

import OrdinaryDiffEq

#= To construct the `ODEProblem` from this package, we will have to provide a right-hand side function that computes the velocities $Ẋ$ of the vortices, which is the local flow velocity at their positions. These velocities are obtained using `vortexvelocities`, which regularizes the vorticity to the grid, solves the potential flow system, and interpolates the velocities from the grid to the vortex locations as

$\left(U_{q}, V_{q}\right)=\sum_{i, j} \mathsf{v}_{i j} d\left(\frac{\mathsf{x}_{i}-X_{q}}{\Delta x}\right)\left(\frac{\mathsf{y}_{j}-Y_{q}}{\Delta x}\right),$

where $v$ is the velocity field on the nodes.
=#

function rhs(X,model,t)
    setvortexpositions!(model,X)
    Ẋ = vortexvelocities!(model)
    return Ẋ
end

# We finally create the problem and solve it with a fourth-order Runge-Kutta time stepping scheme.
X = getvortexpositions(model)
prob = OrdinaryDiffEq.ODEProblem(rhs,X,tspan,model);
sol = OrdinaryDiffEq.solve(prob,dt=0.1,OrdinaryDiffEq.RK4(),dense=false,adaptive=false);

# When we plot the x- and y-position versus time of the first vortex, we see it completed one revolution.
plot(sol.t,map(s->s.u[1],sol.u),xlabel='t',label="x₁")
plot!(sol.t,map(s->s.v[1],sol.u),label="y₁")

#jl @testset "Corotating point vortices" begin
#jl     @test isapprox(sol.u[end].u[1], 0.5; atol = 1e-2)
#jl     @test isapprox(sol.u[end].u[2], -0.5; atol = 1e-2)
#jl     @test isapprox(sol.u[end].v[1], 0.0; atol = 1e-2)
#jl     @test isapprox(sol.u[end].v[2], 0.0; atol = 1e-2)
#jl end

#=
## Corotating vortex patches
A more complex example is the evolution of two circular regions of spatially uniform vorticity that have an equal radius $r_0$ and circulation $\Gamma$, and whose centers are separated by a distance $d_0$. If $r_0/d_0$ is small enough, these two vortex patches can be regarded as two point vortices with the same circulation and separation. We therefore assume the period of this system is approximately the same as the the period for the point vortices case.
=#

#md # ```@setup 1.-Basic-potential-flow-problem
#md # xlim = (-2,2);
#md # ylim = (-2,2);
#md # Δx = 0.05;
#md # g = PhysicalGrid(xlim,ylim,Δx);
#md # ```
#!md xlim = (-2,2);
#!md ylim = (-2,2);
#!md Δx = 0.05;
#!md g = PhysicalGrid(xlim,ylim,Δx);

# We discretize the vortex patches with point vortices arranged on concentric rings using the following function.

function vortexpatch!(vort,xc,yc,Γ,radius,nring)
    Δr = radius/(nring-1/2)
    dΓ = Γ/(1+8*nring*(nring-1)/2)
    push!(vort,Vortex(xc,yc,dΓ))
    for ir in 1:nring-1
        nθ = 8*ir
        for j = 0:nθ-1
            push!(vort,Vortex(xc + ir*Δr*cos(2π*j/nθ),yc + ir*Δr*sin(2π*j/nθ),dΓ))
        end
    end
    return vort
end

vortexpatch(xc,yc,Γ,radius,nring) = vortexpatch!(Vortex[],xc,yc,Γ,radius,nring)

# We will simulate two cases with different values for $r_0/d_0$, the ratio of the vortex patch radius to the distance between them.
d0 = d
anims = []
for r0 in [0.2*d0, 0.4*d0]
    vortices = vcat(vortexpatch(0.0,0.0+d0/2,Γ,r0,10),vortexpatch(0.0,0.0-d0/2,Γ,r0,10));
    vortexcolors = vcat(fill(:blue,length(vortices)÷2),fill(:red,length(vortices)÷2));
    local model = VortexModel(g,vortices=vortices);
    local X = getvortexpositions(model)
    local prob = OrdinaryDiffEq.ODEProblem(rhs,X,tspan,model);
    local sol = OrdinaryDiffEq.solve(prob,dt=0.1,OrdinaryDiffEq.RK4(),dense=false,adaptive=false);
    anim = @animate for i=1:length(sol.t)-1
        plot(xlims=xlim,ylims=ylim,ratio=:equal,legend=:none,title="r0/d0 = $(r0/d0)")
        scatter!(sol.u[i].u,sol.u[i].v,markerstrokewidth=0,markersize=3,color=vortexcolors)
    end
    push!(anims,anim)
end

gif(anims[1])

# If the ratio $r_0/d_0$ is big enough, the two vortex patches merge with eachother, a result that has been widely reported in literature [^1].

gif(anims[2])

# [^1]: Eldredge J. D. (2019) "Mathematical  Modeling  of  Unsteady  Inviscid  Flows, Interdisciplinary Applied Mathematics" *Springer*, vol. 50.


text:
"""
`module DB`

## Overview

This module implements a "database" for storing a precomputed LSQ system. This
is useful for fitting e.g. NBodyIPs since it allows one to precompute the <basis,
data> inner products which are very expensive, and then quickly construct LSQ
systems from them using e.g., many variants of weighting and regularisation.

A db called (e.g.) 'lsqdb' consists of two files:
 * `lsqdb_info.json` : this stores a list of basis functions and a list
 of configurations with attached observations (normally DFT data)
 * `lsqdb_kron.h5` : this stores all the inner products <basis, data>
 as a single large matrix. The observations and basis functions in
 `lsqdb_info` store the corresponding column and row indices.

## More details and remarks

### INFO

* BASIS: The basis is (for the time being) simply represented as a list (Vector) of
JuLIP.AbstractCalculator. These are stored as `INFO["basis"]`.

* DATA: Each "datum" is an atomistic configuration, which must have a
"configtype", stored as a `Dat`.
Informally two pieces of data with the same configtype should represent
"similar" kinds of configurations. `data = INFO["data"]` is a Dictionary
where the keys are the configtypes.

### (De-)Vectorising Data

Forces in `JuLIP` are represented as `Vector{SVector{...}}`, which is the
same memory layout as a 3 x N matrix (`JuLIP.mat`), which can can then be
vectorised (`[:]`), and this vectoriation is readily undone again.
Analogously, any data must be stored in such a vectorised format. This is
achieved (e.g. for forces) via
* `vec_obs(::Val{:F}, F) -> Vector`
* `devec_obs(::Val{:F}, Fvec) -> Vector{SVector{...}}`
or equivalently
* `vec_obs("F", F) -> Vector`
* `devec_obs("F", Fvec) -> Vector{SVector{...}}`
(the `Val` versions are used for performance optimisation)
"""
module DB

using Base.Threads:          SpinLock, nthreads
using StaticArrays:          SVector
using JuLIP:                 AbstractCalculator, AbstractAtoms, Atoms,
                             save_json, load_json, decode_dict
using IPFitting:        Dat, LsqDB, basis, eval_obs, observations,
                             observation, vec_obs, devec_obs,
                             tfor_observations
using IPFitting.Data:   configtype
using HDF5:                  h5open, read

import Base: flush, append!, union

export LsqDB, info, configtypes

const KRONFILE = "_kron.h5"
const INFOFILE = "_info.json"

"""
`dbpath(db::LsqDB)` : return the absolute path to the database files, not
including the 'kron.h5' and 'info.json' endings.
"""
dbpath(db::LsqDB) = db.dbpath

infofile(dbpath::AbstractString) = dbpath * INFOFILE

kronfile(dbpath::AbstractString) = dbpath * KRONFILE

# ------------ Save and load the info file

"if the file exists, append a random string to avoid overwriting"
function _backupfile(fname)
   if isfile(fname)
      fnew = fname * "." * String(rand('a':'z', 5))
      @warn("The file $fname already exists. It will be renamed to $fnew to avoid overwriting.")
      run(`mv $fname $fnew`)
   end
   return nothing
end

function load_info(dbpath::String)
   dbinfo = load_json(infofile(dbpath))
   basis = decode_dict.(dbinfo["basis"])
   configs = Dat.(dbinfo["configs"])   # here we already know the type
   return basis, configs
end

function save_info(dbpath::String, db)
   _backupfile(infofile(dbpath))
   save_json(infofile(dbpath),
             Dict("basis" => Dict.(db.basis),
                  "configs" => Dict.(db.configs))
            )
   return nothing
end

save_info(db) = save_info(dbpath(db), db)

# ----------- Save and load the KRON file

"save a single matrix to HDF5"
_savemath5(A, fname) =
   h5open(fname, "w") do fid
      fid["A"] = A
      nothing
   end

"load a single matrix from HDF5"
_loadmath5(fname) =
   h5open(fname, "r") do fid
      read(fid["A"])
   end

function save_kron(dbpath, db)
   _backupfile(kronfile(dbpath))
   _savemath5(db.Ψ, kronfile(dbpath))
end

save_kron(db) = save_kron(dbpath(db), db)

load_kron(dbpath::String; mmap=false) = _loadmath5(kronfile(dbpath))

load_kron(db::LsqDB; mmap=false) = load_kron(dbpath(db); mmap=mmap)

function LsqDB(dbpath::AbstractString; mmap=false)
   basis, configs = load_info(dbpath)
   Ψ = load_kron(dbpath; mmap=mmap)
   return LsqDB(basis, configs, Ψ, dbpath)
end

"""
for the time being, this just checks that the db director and db name are
admissible. In the future, this should initialise the DB files.
"""
function initdb(basedir, dbname)
   @assert !('/' in dbname)
   @assert isdir(basedir)
   dbpath = joinpath(basedir, dbname)
   initdb(dbpath)
   return nothing
end

function initdb(dbpath)
   # TODO: seems this fails to detect an existing database?
   @assert !isfile(infofile(dbpath))
   @assert !isfile(kronfile(dbpath))
   # check that we can actually create and delete this file
   save_json(infofile(dbpath), Dict("basis" => [], "data" => []))
   rm(infofile(dbpath))
   _savemath5(rand(5,5), kronfile(dbpath))
   rm(kronfile(dbpath))
   return nothing
end

function flush(db::LsqDB)
   save_info(db)
   save_kron(db)
   return nothing
end

function LsqDB(dbpath::AbstractString,
               basis::AbstractVector{<: AbstractCalculator},
               configs::AbstractVector{Dat};
               verbose=true,
               maxnthreads=nthreads())
   # assign indices, count observations and allocate a matrix
   Ψ = _alloc_lsq_matrix(configs, basis)
   # create the struct where everything is stored
   db = LsqDB(basis, configs, Ψ, dbpath)
   # parallel assembly of the LSQ matrix
   tfor_observations( configs,
      (n, okey, cfg, lck) -> safe_append!(db, lck, cfg, okey),
      msg = "Assemble LSQ blocks",
      verbose=verbose,
      maxnthreads=maxnthreads )
   # save to file
   if dbpath != ""
      verbose && @info("Writing db to disk...")
      try
         flush(db)
      catch
         @warn("""something went wrong trying to save the db to disk, but the data
               should be ok; if it is crucial to keep it, try to save manually.""")
      end
      verbose && @info("... done")
   else
      verbose && @info("db is not written to disk since `dbpath` is empty.")
   end
   return db
end

function set_matrows!(d::Dat, okey::String, irows::Vector{Int})
   d.rows[okey] = irows
   return d
end

matrows(d::Dat, okey::String) = d.rows[okey]

function _alloc_lsq_matrix(configs, basis)
   # the indices associated with the basis are simply the indices within
   # the array - there is nothing else to do here
   #
   # loop through all observations and assign indices
   nrows = 0
   for (okey, d, _) in observations(configs)
      len = length(observation(d, okey))
      set_matrows!(d, okey, collect(nrows .+ (1:len)))
      nrows += len
   end
   # allocate and return the matrix
   return zeros(Float64, nrows, length(basis))
end

function safe_append!(db::LsqDB, db_lock, cfg, okey)
   # computing the lsq blocks ("rows") can be done in parallel,
   lsqrow = evallsq(cfg, basis(db), okey)
   irows = matrows(cfg, okey)
   # but writing them into the DB must be done in a threadsafe way
   lock(db_lock)
   db.Ψ[irows, :] .= lsqrow
   unlock(db_lock)
   return nothing
end

# ------------------- Evaluating LSQ Blocks -----------------

"""
Take a basis and split it into individual basis groups.
"""
function split_basis(basis; splitfun = b -> hash(Val{:basis}(), b))
   # get some basis hashs of the individual basis functions
   tps = splitfun.(basis)
   Iord = Vector{Int}[]
   Bord = Any[]
   for tp in unique(tps)
      # find which elements of basis have type `tp`
      I = findall( [tp == t  for t in tps] )
      push!(Iord, I)
      push!(Bord, [b for b in basis[I]])
   end
   return Bord, Iord
end

# fill the LSQ system, i.e. evaluate basis at data points
function evallsq(d::Dat, B::AbstractVector{TB}, okey) where {TB <: AbstractCalculator}
   B1 = [b for b in B]
   if !(isconcretetype(eltype(B1)))
      return evallsq_split(d, B1, okey)
   end
   # TB is a leaf-type so we can use "evaluate_many"
   return _evallsq(Val(Symbol(okey)), B1, Atoms(d))
end

"""
evaluate one specific kind of datum, such as energy, forces, etc
for one Dat across a large section of the basis; implicitly this will
normally be done via `evaluate_many` or similar. The line
```
vals = eval_obs(vDT, B, at)
```
should likely be the bottleneck of `evallsq`
"""
function _evallsq(vDT::Val,
                  B::AbstractVector{<:AbstractCalculator},
                  at::AbstractAtoms)
   # vals will be a vector containing multiple evaluations
   vals = eval_obs(vDT, B, at)
   # vectorise the first so we know the length of the observation
   vec1 = vec_obs(vDT, vals[1])
   # create a multi-D array to reshape these into
   A = Array{Float64}(undef, length(vec1), length(B))
   A[:, 1] = vec1
   for n = 2:length(B)
      A[:, n] = vec_obs(vDT, vals[n])
   end
   return A
end

"""
split the Basis `B` into subsets of identical types and evaluate
those independently (fast).
"""
function evallsq_split(d, basis, okey)
   # TB is not a leaf-type so we should split the basis to be able to
   # evaluate_many & co
   Bord, Iord = split_basis(basis)
   lenobs = length(d.D[okey])
   lsqrow = zeros(lenobs, length(basis))
   for (B, IB) in zip(Bord, Iord)
      lsqrow[:, IB] = evallsq(d, B, okey)
   end
   return lsqrow
end

# =========================== Convenience =================================

# TODO: THIS NEEDS TO BE REWRITTEN! => ALSO WRITE TESTS!
# function union(db1::LsqDB,db2::LsqDB; dbpath = (db1.dbpath * "_u"))
#    @info("Warning: the union implies that the data in the two databases are the same")
#    configtypes = collect(keys(db1.data_groups))
#    basis = cat(1, db1.basis, db2.basis)
#    data_groups = db1.data_groups
#
#    kron_groups1 = db1.kron_groups
#    kron_groups2 = db2.kron_groups
#
#    kron_groups = Dict{String, KronGroup}()
#    for k in configtypes
#       kron_groups[k] = KronGroup()
#       @assert haskey(db2.data_groups, k)
#       for ot in ["E", "F", "V"]
#          if !haskey(data_groups[k][1].D, ot)
#             continue
#          end
#          kron_groups[k][ot] = cat(3,kron_groups1[k][ot],kron_groups2[k][ot])
#       end
#    end
#    db = LsqDB(basis, data_groups, kron_groups, dbpath)
#    flush(db)
#    return db
# end

# function union(db1::LsqDB,db2::LsqDB, dbpath = (db1.dbpath * "_u"))
#    basis = [basis(db1), basis(db2)]
#    data_groups::Dict{String, DataGroup}
#       kron_groups::Dict{String, KronGroup}
#       dbpath::String
#    end
#
#    basis(db::LsqDB) = db.basis
#    basis(db::LsqDB, i::Integer) = db.basis[i]
#    data(db::LsqDB) = db.data
# end

configtypes(db::LsqDB) = unique(configtype.(db.configs))

_nconfigs(db::LsqDB, ct::AbstractString) =
   length(find(configtype.(db.configs) .== ct))

function info(db::LsqDB)
   # config names, how many
   all_cts = configtype.(db.configs)
   unq_cts = unique(all_cts)
   configs_info = Dict{String, Int}()
   for ct in unq_cts
      configs_info[ct] = sum(all_cts .== ct)
   end

   println("======================================================")
   println("       LsqDB Summary ")
   println("------------------------------------------------------")
   println(" Datagroup: configname  =>  number of configs ")
   for (i, (key, val)) in enumerate(configs_info)
      println("        $i : $key  =>  $val ")
   end
   # basis groups, how many, max degree
   B = db.basis
   Bord, Iord = split_basis(B)
   println("------------------------------------------------------")
   println(" Basis Group  | length | description ")
   for (i, B1) in enumerate(Bord)
      lenstr = replace(string(length(B1), pad=5), '0' => ' ')
      idxstr = replace(string(i, pad=2), '0' => ' ')
      desc = string(B1)
      desc = desc[1:min(50, length(desc))]
      println("          $idxstr  | $lenstr  | $desc" )
      #
   end
   println("======================================================")
end

end


text:
using JLD
using Plots
using LaTeXStrings
using Formatting
using TexTables
#reset past font scaling
Plots.scalefontsizes()
#scale fonts
Plots.scalefontsizes(1.75)

lightblue = colorant"rgb(63%,74%,78%)"
orange = colorant"rgb(85%,55%,13%)"
silver = colorant"rgb(69%,67%,66%)"
rust = colorant"rgb(72%,26%,6%)"

#Extracting stored varibles
ld = load("./out/jld/TimingVaryRho.jld")
for key in keys(ld)
  symb = Symbol(key)
  data = ld[key]
  @eval $(symb) = $(data)
end
# ----------------------------------------
# Creating plot of RMSE
# ----------------------------------------
plotTiming = plot(ylabel=L"$\mathrm{Time} \ \mathrm{in} \ \mathrm{s}$",
                  xlabel=L"\rho",
                  legend=:topleft
                  )

plotTimingVnonz = plot(ylabel=L"$\mathrm{Time} \ \mathrm{in} \ \mathrm{s}$",
                  xlabel=L"$\mathrm{nonzeros}$",
                  legend=false,
                  )

plotTimingBessel = plot(ylabel=L"$\mathrm{Time} \ \mathrm{in} \ \mathrm{s}$",
                        xlabel=L"\rho",
                        legend=:topleft,
                        )

plotTimingBesselVnonz = plot(ylabel=L"$\mathrm{Time} \ \mathrm{in} \ \mathrm{s}$",
                        xlabel=L"$\mathrm{nonzeros}$",
                        legend=false,
                        )

plot!(plotTiming,
      ρList,
      timings,
      markershape=:circle,
      color=orange,
      markersize=15,
      linewidth=5,
      linestyle=:dash,
      label="Exponential, no aggretation ")

plot!(plotTiming,
      ρList,
      timingsAgg,
      markershape=:diamond,
      color=lightblue,
      markersize=15,
      linewidth=5,
      linestyle=:dot,
      label="Exponential, with aggregation")

plot!(plotTimingVnonz,
      nonz,
      timings,
      markershape=:circle,
      color=orange,
      markersize=15,
      linewidth=5,
      linestyle=:dash,
      label="Exponential, no aggretation ")

plot!(plotTimingVnonz,
      nonzAgg,
      timingsAgg,
      markershape=:diamond,
      color=lightblue,
      markersize=15,
      linewidth=5,
      linestyle=:dot,
      label="Exponential, with aggregation")

plot!(plotTimingBessel,
      ρList,
      timingsBessel,
      markershape=:circle,
      color=rust,
      markersize=15,
      linewidth=5,
      linestyle=:dash,
      label="Bessel, no aggregation")

plot!(plotTimingBessel,
      ρList,
      timingsBesselAgg,
      markershape=:diamond,
      color=silver,
      markersize=15,
      linewidth=5,
      linestyle=:dot,
      label="Bessel, aggregation")

plot!(plotTimingBesselVnonz,
      nonzBessel,
      timingsBessel,
      markershape=:circle,
      color=rust,
      markersize=15,
      linewidth=5,
      linestyle=:dash,
      label="Bessel, no aggregation")

plot!(plotTimingBesselVnonz,
      nonzBesselAgg,
      timingsBesselAgg,
      markershape=:diamond,
      color=silver,
      markersize=15,
      linewidth=5,
      linestyle=:dot,
      label="Bessel, aggregation")

savefig(plotTiming, "./out/plots/timing.pdf")
savefig(plotTimingBessel, "./out/plots/timingBessel.pdf")

savefig(plotTimingVnonz, "./out/plots/timingVnonz.pdf")
savefig(plotTimingBesselVnonz, "./out/plots/timingBesselVnonz.pdf")


text:
function initialize!(integrator, cache::LinearImplicitEulerConstantCache)
  integrator.kshortsize = 2
  integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)
  integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal

  # Avoid undefined entries if k is an array of arrays
  integrator.fsallast = zero(integrator.fsalfirst)
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
end

@muladd function perform_step!(integrator, cache::LinearImplicitEulerConstantCache, repeat_step=false)
  @unpack t,dt,uprev,u,k = integrator

  L = update_coefficients(integrator.f,u,p,t+dt)

  if typeof(uprev) <: AbstractArray
    W = I - dt*L
  else
    W = 1 - dt*L
  end

  # if Bs
  k = uprev

  u = W\k

  if integrator.opts.adaptive && integrator.success_iter > 0
    # local truncation error (LTE) bound by dt^2/2*max|y''(t)|
    # use 2nd divided differences (DD) a la SPICE and Shampine

    # TODO: check numerical stability
    uprev2 = integrator.uprev2
    tprev = integrator.tprev

    dt1 = dt*(t+dt-tprev)
    dt2 = (t-tprev)*(t+dt-tprev)
    c = 7/12 # default correction factor in SPICE (LTE overestimated by DD)
    r = c*dt^2 # by mean value theorem 2nd DD equals y''(s)/2 for some s

    tmp = @. r*abs((u - uprev)/dt1 - (uprev - uprev2)/dt2)
    atmp = calculate_residuals(tmp, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm)
    integrator.EEst = integrator.opts.internalnorm(atmp)
  else
    integrator.EEst = 1
  end

  integrator.fsallast = f(u, p, t+dt)
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
  integrator.u = u
end

function initialize!(integrator, cache::LinearImplicitEulerCache)
  integrator.kshortsize = 2
  integrator.fsalfirst = cache.fsalfirst
  integrator.fsallast = cache.k
  resize!(integrator.k, integrator.kshortsize)
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
  integrator.f(integrator.fsalfirst, integrator.uprev, integrator.p, integrator.t) # For the interpolation, needs k at the updated point
end

@muladd function perform_step!(integrator, cache::LinearImplicitEulerCache, repeat_step=false)
  @unpack t,dt,uprev,u = integrator
  @unpack W,k,tmp,atmp = cache
  mass_matrix = integrator.f.mass_matrix

  L = integrator.f
  update_coefficients!(L,u,p,t+dt)

  if typeof(L) <: AbstractDiffEqLinearOperator

      # Of the form u' = A(t)u

      # Check is_constant before redoing
      for j in 1:length(u), i in 1:length(u)
          @inbounds W[i,j] = @muladd mass_matrix[i,j]-dt*L[i,j]
      end
      k .= uprev # + B
  else # Must be a DiffEqAffineOperator

      # Of the form u' = A(t)u + B(t)
      # Generalize later!
      A = L.As[1]
      B = L.Bs[1]

      for j in 1:length(u), i in 1:length(u)
          @inbounds W[i,j] = @muladd mass_matrix[i,j]-dt*A[i,j]
      end
      @. k = uprev + dt*B
  end

  cache.linsolve(vec(u), W, vec(k), true)

  if integrator.opts.adaptive && integrator.success_iter > 0
    # local truncation error (LTE) bound by dt^2/2*max|y''(t)|
    # use 2nd divided differences (DD) a la SPICE and Shampine

    # TODO: check numerical stability
    uprev2 = integrator.uprev2
    tprev = integrator.tprev

    dt1 = (dt)*(t+dt-tprev)
    dt2 = (t-tprev)*(t+dt-tprev)
    c = 7/12 # default correction factor in SPICE (LTE overestimated by DD)
    r = c*dt^2 # by mean value theorem 2nd DD equals y''(s)/2 for some s

    @. tmp = r*abs((u - uprev)/dt1 - (uprev - uprev2)/dt2)
    calculate_residuals!(atmp, tmp, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm)
    integrator.EEst = integrator.opts.internalnorm(atmp)
  else
    integrator.EEst = 1
  end

  f(integrator.fsallast,u,p,t+dt)
end

function initialize!(integrator, cache::MidpointSplittingCache)
  integrator.kshortsize = 2
  integrator.fsalfirst = cache.fsalfirst
  integrator.fsallast = cache.k
  resize!(integrator.k, integrator.kshortsize)
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
  integrator.f(integrator.fsalfirst, integrator.uprev, integrator.p, integrator.t) # For the interpolation, needs k at the updated point
end

function perform_step!(integrator, cache::MidpointSplittingCache, repeat_step=false)
  @unpack t,dt,uprev,u = integrator
  @unpack W,k,tmp = cache
  mass_matrix = integrator.f.mass_matrix

  L = integrator.f
  update_coefficients!(L,u,p,t+dt/2)

  A = L.As[1]
  Bs = L.As[2:end]

  copyto!(tmp, uprev)
  for B in reverse(Bs)
    u .= exp((dt/2)*B)*tmp
    @swap!(tmp,u)
  end

  u .= exp(dt*A)*tmp

  for B in Bs
    tmp .= exp((dt/2)*B)*u
    @swap!(u,tmp)
  end

  f(integrator.fsallast,u,p,t+dt)
end


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

function get_theo_Pd(f_s, T, CN0, B, p_fa, Tsys=535)
    N = f_s * T
    v_i = N * sqrt(2*GNSSTools.k*535)*10^(CN0/20)
    σ_n_theoretical = sqrt(GNSSTools.k*B*Tsys)
    SNR_theo = CN0 - 10*log10(1/T)
    v_t_theoretical = p_fa2v_t(p_fa, sqrt(N)*σ_n_theoretical)
    P_d_theo = v_t2p_d(v_t_theoretical,
                       sqrt(N)*σ_n_theoretical,
                       v_i=v_i)
    # P_d_theo = v_t2p_d(v_t_theoretical, sqrt(N)*σ_n_theoretical,
                    #    SNR=SNR_theo)
    return P_d_theo
end

"""
sim_pd(CN0s, state_num; iterations=100, T=2e-3)

Simulates signals and processes them to calculate the propability of detection,
P_d.
"""
function sim_pd(CN0s, f_s; iterations=100, T=1e-3, M=1,
                return_theoretical=true, f_d=0,
                sigtype=define_l1ca_code_type(),
                doppler_bin_num=5)
    t_length = T*M
    prn = 1
    fd_range = 5e3
    # f_if = 1.25e6
    f_if = 0
    h_parms = [1.9e-6, 2e-6, 9.7e-8, 9.1e-11, 9.9e-15] ./ 16.3676e6^2  # TCXO
    # h_parms = [2.2e-9, 1.6e-10, 1.6e-10, 7.1e-13, 7.2e-16] ./ 10e6^2  # OCXO
    # sigtype = define_l1ca_code_type()
    # sigtype = define_l5_code_type()
    if sigtype.include_I && sigtype.include_Q
        B = max(sigtype.B_I, sigtype.B_Q)
    elseif sigtype.include_I && !sigtype.include_Q
        B = sigtype.B_I
    elseif !sigtype.include_I && sigtype.include_Q
        B = sigtype.B_Q
    else
        error("No CodeType defined in SigType struct.")
    end
    d = ceil(Int, f_s/l1ca_chipping_rate/2)
    replica = definereplica(sigtype, f_s, T)
    signal = definesignal(sigtype, f_s, t_length;
                          prn=prn, fd_rate=0, f_if=f_if,
                          include_phase_noise=false, include_thermal_noise=true,
                          include_adc=true, include_carrier=true,
                          receiver_h_parms=h_parms)
    if return_theoretical
        P_ds_theoretical = Array{Float64}(undef, length(CN0s))
        SNR_theoretical = Array{Float64}(undef, length(CN0s))
    end
    P_ds_experimental = Array{Float64}(undef, length(CN0s))
    P_ds_exp_from_dopler_code = Array{Float64}(undef, length(CN0s))
    SNR_experimental = Array{Float64}(undef, length(CN0s))
    p = Progress(iterations*length(CN0s), 1, "Processing...")
    for j in 1:length(CN0s)
        σ²_ns = zeros(iterations)
        powers = zeros(iterations)
        above_threshold_num_1 = 0
        above_threshold_num_2 = 0
        SNR_avg = 0.
        for iteration in 1:iterations
            # ϕ₀ = rand(0:0.0001:2π)
            ϕ₀ = 0
            # f_d = rand(-fd_range:0.1:fd_range)
            # code_start_idx = rand(1:1:replica.sample_num)
            code_start_idx = 1
            # ϕ₀ = 0
            # code_start_idx = 1000
            definesignal!(signal; CN0=CN0s[j], new_phase_noise=false,
                          new_thermal_noise=true, phi=ϕ₀, f_d=f_d,
                          code_start_idx=code_start_idx,
                          nADC=16)
            generatesignal!(signal)
            fd_range = (doppler_bin_num - 1)/(2*T)
            fd_course, n0_est, SNR_est, P_d, above_threshold,
            corr_result = courseacquisition(signal,
                                            replica,
                                            prn;
                                            fd_center=0,
                                            fd_range=fd_range,
                                            fd_rate=0,
                                            return_corrresult=true,
                                            M=M)
            rows, cols = size(corr_result)
            idx = argmax(corr_result)
            row_idx = idx[1]
            col_idx = idx[2]
            noise_row_idx = (row_idx + 3) % rows + 1
            # Peak power should be peak + the two adjacent bins
            # lower_row_idx = row_idx - 1
            # if lower_row_idx < 1
            #     lower_row_idx = rows + lower_row_idx
            # end
            # upper_row_idx = row_idx + 1
            # if upper_row_idx > rows
            #     upper_row_idx = upper_row_idx % rows
            # end
            # lower_col_idx = col_idx - 1
            # if lower_col_idx < 1
            #     lower_col_idx = replica.sample_num + lower_col_idx
            # end
            # upper_col_idx = col_idx + 1
            # if upper_col_idx > replica.sample_num
            #     upper_col_idx = upper_col_idx % replica.sample_num
            # end
            # power = sum(sqrt.([corr_result[row_idx,lower_col_idx],
            #                   corr_result[row_idx,col_idx],
            #                   corr_result[row_idx,upper_col_idx]]))
            # power = sum(sqrt.([corr_result[lower_row_idx,col_idx],
            #                   corr_result[row_idx,col_idx],
            #                   corr_result[upper_row_idx,col_idx]]))
            power = sqrt(corr_result[argmax(corr_result)])
            σ²_n = sum(corr_result[noise_row_idx,:]) / (cols - 1)
            σ_n = sqrt(σ²_n)
            # σ²_n = (sum(corr_result) - maximum(corr_result))/(rows*cols - 1)
            # v_t = p_fa2v_t(p_fa, sqrt(σ²_n))
            v_t = p_fa2v_t(p_fa, σ_n)
            # SNR_avg += power/σ_n
            powers[iteration] = power
            σ²_ns[iteration] = σ_n
            if power >= v_t
                above_threshold_num_2 += 1
            end
            correct_doppler = fd_course == (round(f_d*T, digits=0)/T)
            # correct_n0 = (n0_est >= code_start_idx-d) && (n0_est <= code_start_idx+d)
            correct_n0 = n0_est == code_start_idx
            # correct_n0 = n0_est == code_start_idx
            if correct_doppler && correct_n0
                above_threshold_num_1 += 1
            end
            # println([n0_est, code_start_idx, f_d, fd_course])
            next!(p)
        end
        # σ_n_exp = mean(σ²_ns)
        # # println([minimum(σ²_ns), maximum(σ²_ns), maximum(σ²_ns)-minimum(σ²_ns)])
        # v_t_exp = p_fa2v_t(p_fa, σ_n_exp)
        # for iteration in 1:iterations
        #     if powers[iteration] >= v_t_exp
        #         above_threshold_num_2 += 1
        #     end
        # end

        SNR_experimental[j] = 10*log10(SNR_avg / iterations)
        P_ds_experimental[j] = above_threshold_num_2 / iterations
        P_ds_exp_from_dopler_code[j] = above_threshold_num_1 nbins = 1000
        fig = figure(figsize=(7.5, 4.5))
        ax1 = fig.add_subplot(2,2,1, aspect="auto")
        hist2D(gps_hist.dopplers./1000, gps_hist.doppler_rates, bins=nbins)
        xlabel("Doppler (kHz)")
        ylabel("Doppler Rate (Hz/s)")
        title("(a)")

        ax2 = fig.add_subplot(2,2,2, aspect="auto")
        hist2D(iridium_hist.dopplers./1000, iridium_hist.doppler_rates, bins=nbins)
        xlabel("Doppler (kHz)")
        ylabel("Doppler Rate (Hz/s)")
        title("(b)")

        ax3 = fig.add_subplot(2,2,3, aspect="auto")
        hist2D(starlink_hist.dopplers./1000, starlink_hist.doppler_rates, bins=nbins)
        xlabel("Doppler (kHz)")
        ylabel("Doppler Rate (Hz/s)")
        title("(c)")

        ax4 = fig.add_subplot(2,2,4, aspect="auto")
        hist2D(oneweb_hist.dopplers./1000, oneweb_hist.doppler_rates, bins=nbins)
        xlabel("Doppler (kHz)")
        ylabel("Doppler Rate (Hz/s)")
        title("(d)")/ iterations
        if return_theoretical
            P_d_theo = get_theo_Pd(f_s, T, CN0s[j], B, p_fa, signal.Tsys)
            # σ_n_theoretical = sqrt(GNSSTools.k*B*signal.Tsys)
            # # v_i = f_s * T * (sqrt(2*GNSSTools.k*535)*10^(CN0s[j]/20))^2
            # N = f_s*T
            # v_i = N * sqrt(GNSSTools.k*535)*10^(CN0s[j]/20)
            # SNR_theo = CN0s[j] - 10*log10(1/T)
            # SNR_theoretical[j] = SNR_theo
            # v_t_theoretical = p_fa2v_t(p_fa, sqrt(N)*σ_n_theoretical)
            # P_ds_theoretical[j] = v_t2p_d(v_t_theoretical,
            #                               sqrt(N)*σ_n_theoretical,
            #                               v_i=v_i)
            # # P_ds_theoretical[j] = v_t2p_d(v_t_theoretical, σ_n_theoretical,
            #                             #   SNR=SNR_theo)
            P_ds_theoretical[j] = P_d_theo
        end
    end
    if return_theoretical
        return (P_ds_theoretical, P_ds_experimental, P_ds_exp_from_dopler_code)
    else
        return (P_ds_experimental, P_ds_exp_from_dopler_code)
    end
end

CN0s = Array(range(25, 50, step=0.5))
iterations = 100
p_fa = 1e-7
f_d = 0

# f_s = 2.5e6
# sigtype = define_l1ca_code_type()
# Ts = [1e-3 2e-3 5e-3 10e-3]
# name = "l1ca"

f_s = 25e6
sigtype = define_l5_code_type(;channel="Q")
Ts = [1e-3, 20e-3]  # seconds
name = "l5"

M = 1
P_ds_theoretical = []
P_ds_experimental = []
P_ds_exp_from_dopler_code = []

println("Working on $(sigtype.name) signal...")

directory = "/home/sjbilardi/Dropbox/Apps/Overleaf/MastersThesis/"
file_name = "ch4_compare_theoretical_empirical_Pd_$(name)_fd=$(f_d)Hz"
file = string(directory, file_name, ".jld")

for i in 1:length(Ts)
        P_ds_theo, P_ds_exp, P_ds_exp2 = sim_pd(CN0s, f_s; iterations=iterations,
                                                T=Ts[i], M=M,
                                                return_theoretical=true,
                                                sigtype=sigtype,
                                                f_d=f_d,
                                                doppler_bin_num=5)
        push!(P_ds_theoretical, P_ds_theo)
        push!(P_ds_experimental, P_ds_exp)
        push!(P_ds_exp_from_dopler_code, P_ds_exp2)
        save(file,
             "P_ds_theoretical", P_ds_theoretical,
             "P_ds_experimental", P_ds_experimental,
             "P_ds_exp_from_dopler_code", P_ds_exp_from_dopler_code,
             "CN0s", CN0s,
             "M", M,
             "p_fa", p_fa,
             "iterations", iterations,
             "Ts", Ts,
             "f_s", f_s)
end

P_ds_theoretical, P_ds_experimental, CN0s, f_s = load(file,
                                                 "P_ds_theoretical",
                                                 "P_ds_experimental",
                                                 "CN0s",
                                                 "f_s")

directory = "/home/sjbilardi/Dropbox/Apps/Overleaf/MastersThesis/"
file_names = [string(directory, "ch4_compare_theoretical_empirical_Pd_l1ca_fd=0Hz.jld"),
              string(directory, "ch4_compare_theoretical_empirical_Pd_l5_fd=0Hz.jld")]
cmap = get_cmap("viridis")
fig = figure(figsize=(7.5, 3.5))
for f in 1:length(file_names)
    file_name = file_names[f]
    P_ds_theoretical, P_ds_experimental, CN0s, f_s, Ts = load(file_name,
                                                          "P_ds_theoretical",
                                                          "P_ds_experimental",
                                                          "CN0s",
                                                          "f_s",
                                                          "Ts")
    Pds_to_plot = length(P_ds_theoretical)
    Pd_to_plot = Array(1:Pds_to_plot)
    ax = fig.add_subplot(1, 2, f)
    for i in 1:Pds_to_plot
        if i in Pd_to_plot
            color = cmap(float(i)/Pds_to_plot)
            label = "T = $(floor(Int, Ts[i]*1e3))ms"
            # label = "T = $(floor(Int, Ts[i]*1e3))ms, f_s = $(round(f_s*1e-6, digits=1))MHz"
            ax.plot(CN0s, P_ds_theoretical[i], "-", color=color,
                    label=label)
            ax.plot(CN0s, P_ds_experimental[i], ":", color=color)
            ax.plot(CN0s, P_ds_experimental[i], ".", color=color)
        end
    end
    legend()
    xlabel("C/N₀ (dB⋅Hz)")
    ylabel(L"P_d")
    letter = ('a' + (f-1))
    title("($(letter))")
end
subplots_adjust(wspace=0.3, bottom=0.15, left=0.08, right=0.92, top=0.93)
savefig(string(directory, "figures/ch4_compare_theoretical_empirical_Pd.svg"), dpi=300)

# """

# The error between the truth and observation is near zero when f_d is an integer
# multiple of 1/T. If f_d is not a multiple of 1/T and is 1/2T from either
# close multiples of 1/T, the error will be the maximum (~3.5dB).

# Also, if f_s is a multiple of the code chipping rate, then it will also
# result in the worst acquisition performance and greater error between the
# the truth and observation.

# """