# A collection of date and time functions.
#
# Some of these have/should be changed to use the built-in Dates module.
#
# The functions which use Julian Day Numbers are valid only for positive values,
# i.e., for dates after -4712 (4713BC).
#
# Unless otherwise specified, Julian Day Numbers are fractional values.
#
# Numeric years use the astronomical convention of a year 0: 0 = 1BC, -1 = 2BC,
# etc.
#
# Numeric months are 1-based: Jan = 1 Dec = 12.
#
# Numeric days are the same as the calendar value.
#
# Reference: Jean Meeus, _Astronomical Algorithms_, second edition, 1998, Willmann-Bell, Inc.

export
    is_leap_year,
    cal_to_jd,
    date_to_jd,
    jd_to_date,
    jd_to_cal,
    cal_to_jd_m,
    cal_to_day_of_year,
    day_of_year_to_cal,
    pesach,
    jewish_new_year,
    moslem_to_christian,
    christian_to_moslem,
    easter,
    jd_to_day_of_week,
    is_dst,
    jd_to_jcent,
    lt_to_str,
    ut_to_lt

"""
    is_leap_year(yr, gregorian=true)

Parameters:
yr        : year
gregorian : If True, use Gregorian calendar, else use Julian calendar (default: True)

Return:
True is this is a leap year, else False.

Return True if this is a leap year in the Julian or Gregorian calendars
"""
function is_leap_year(yr, gregorian=true)
    yr =  floor(Integer, yr)
    if gregorian
        return (yr % 4 == 0) && ((yr % 100 != 0) || (yr % 400 == 0))
    else
        return yr % 4 == 0
    end
end

"""
    cal_to_jd(yr, mo, day, gregorian=true)

Convert a date in the Julian or Gregorian calendars to the Julian Day Number

    Parameters:
        yr        : year
        mo        : month
        day       : day, may be fractional
        gregorian : if True, use Gregorian calendar, else use Julian calendar (default: True)

    Return:
        jd        : julian day number

presumably returns midday time, since Julian days start at noon?

"""
function cal_to_jd(yr, mo, day, gregorian=true)
    a =  floor(Integer, (14 - mo)/12)

    y = yr + 4800 - a
    m = mo + 12 * a - 3

    jdn = day +  floor(Integer, ((153 * m + 2) / 5)) + 365 * y +  floor(Integer, (y / 4)) - 32083.5
    if !gregorian
    	return jdn
    else
    	return jdn -  floor(Integer, (y/100)) +  floor(Integer, y/400) + 38
    end
end

"""
    date_to_jd(yr, mo, d, h, m, s, gregorian=true)

be careful as they don't take into account the current Daylight Savings Time.
"""
function date_to_jd(yr, mo, d, h, m, s, gregorian=true)
   return(cal_to_jd(yr, mo, d, gregorian) + hms_to_fday(h,m,s))
end

"""
    jd_to_date(jd, gregorian=true)

be careful as they don't take into account the current Daylight Savings Time.
"""
function jd_to_date(jd, gregorian=true)
    (y, month, df) = jd_to_cal(jd, gregorian)
    (f, d) = modf(df)
    (h, m, s) = fday_to_hms(f)
    return (y, month, floor(Integer, d), h, m, s)
end

"""
    jd_to_cal(jd, gregorian=true)

Convert a Julian day number to a date in the Julian or Gregorian calendars.

Parameters:
    jd        : Julian Day number
    gregorian : if True, use Gregorian calendar, else use Julian calendar (default: True)

Return:
    year
    month
    day (may be fractional)

Return a tuple (year, month, day).
"""
function jd_to_cal(jd, gregorian=true)
    (F, Z) = modf(jd + 0.5)
    if gregorian
        alpha =  floor(Integer, ((Z - 1867216.25) / 36524.25))
        A = Z + 1 + alpha -  floor(Integer, alpha / 4)
    else
        A = Z
    end

    B = A + 1524
    C =  floor(Integer, (B - 122.1) / 365.25)
    D =  floor(Integer, 365.25 * C)
    E =  floor(Integer, (B - D) / 30.6001)
    day = B - D -  floor(Integer, 30.6001 * E) + F
    if E < 14
        mo = E - 1
    else
        mo = E - 13
    end

    if mo > 2
        yr = C - 4716
    else
        yr = C - 4715
    end

    return (yr, mo, day)
end

"""
    cal_to_jd_m(yr, mo, day, gregorian=true)

Meeus' implementation of cal_to_jd - which gives the same results
"""
function cal_to_jd_m(yr, mo, day, gregorian=true)
    if mo <= 2
        yr = yr-1
        mo = mo+12
    end
    if gregorian
        A = floor(Integer, yr / 100)
        B = 2 - A + floor(Integer, A / 4)
    else
        B = 0
    end
    return  floor(Integer, 365.25 * (yr + 4716)) +  floor(Integer, 30.6001 * (mo + 1)) + day + B - 1524.5
end

"""
    cal_to_day_of_year(yr, mo, dy, gregorian=true)

Convert a date in the Julian or Gregorian calendars to day of the year (Meeus 7.1).

Parameters:
    yr        : year
    mo        : month
    day       : day
    gregorian : if True, use Gregorian calendar, else use Julian calendar (default: True)

Return:
    day number : 1 = Jan 1...365 (or 366 for leap years) = Dec 31.
"""
function cal_to_day_of_year(yr, mo, dy, gregorian=true)
    if gregorian == false
        gregorian = true
    end
    if is_leap_year(yr, gregorian)
        K = 1
    else
        K = 2
    end
    dy =  floor(Integer, dy)
    return  floor(Integer, 275 * mo / 9.0) - (K *  floor(Integer, (mo + 9) / 12.0)) + dy - 30
end

"""
    day_of_year_to_cal(yr, N, gregorian=true)

Convert a day of year number to a month and day in the Julian or Gregorian calendars.

Parameters:
    yr        : year
    N        : day of year, 1..365 (or 366 for leap years)
    gregorian : if True, use Gregorian calendar, else use Julian calendar (default: True)

Return:
    month
    day
"""
function day_of_year_to_cal(yr, N, gregorian=true)
    if is_leap_year(yr, gregorian)
        K = 1
    else
        K = 2
    end
    if (N < 32)
        mo = 1
    else
        mo =  floor(Integer, 9 * (K+N) / 275.0 + 0.98)
    end
    dy =  floor(Integer, N -  floor(Integer, 275 * mo / 9.0) + K *  floor(Integer, (mo + 9) / 12.0) + 30)
    return (mo, dy)
end

"""
    pesach(yr, gregorian=true)

Return the date of Jewish Pesach  for a year in the Julian or Gregorian calendars.

Parameters:
    yr        : year (in the Jewish calendar)
    gregorian : if True, use Gregorian calendar, else use Julian calendar (default: True)

Return:
    month
    day

Meeus - chapter 9
"""
function pesach(yr, gregorian=true)
    C =  floor(Integer, yr/100)
    S = 0
    if gregorian
        S =  floor(Integer, (3*C-5)/4)
    end
    A = yr - 3760
    a = (12 * yr + 12) % 19
    b = yr % 4
    Q = -1.904412361576 + 1.554241796621 * a + 0.25 * b -0.003177794022 * yr + S
    j = ( floor(Integer, Q) + 3 * yr + 5 * b + 2 - S) % 7
    r = Q- floor(Integer, Q)
    if j==2 || j==4 || j==6
        D =  floor(Integer, Q) + 23
    elseif j==1 && a > 6 && r >= 0.63287037
        D =  floor(Integer, Q)+24
    elseif j==0 && a > 11 && r >= 0.897723765
        D =  floor(Integer, Q) + 23
    else
        D =  floor(Integer, Q) + 22
    end

    if D > 31
        return (4, D-31)
    else
        return (3, D)
    end
end

"""
    jewish_new_year(yr)

"""
function jewish_new_year(yr)
    gr_yr = yr-1-3760
    (m, d) = pesach(yr-1)
    jd = cal_to_jd(gr_yr, m, d) + 316
    return jd_to_cal(Float64)
end

"""
    moslem_to_christian(h, m, d)

Converts a date in the moslem calendar to a date in the Christian calendar
Meeus - chapter 9

Input: h - moslem yearhe christian year
       m - moslem month number
       d - moslem day
Returns: yr, month, day in the christian calendar

This function will return meaningless results if called with a date earlier than
622, July 16th
"""
function moslem_to_christian(h, m, d)
    N = d +  floor(Integer, 29.5001*(m-1)+0.99)
    Q  =  floor(Integer, h/30)
    R  = h%30
    A =  floor(Integer, (11 * R + 3)/30)
    W = 404*Q + 354*R+208 + A
    Q1 =  floor(Integer, W/1461)
    Q2 = W % 1461
    G  = 621 + 4 *  floor(Integer, 7 * Q + Q1)
    K =  floor(Integer, Q2/365.2422)
    E =  floor(Integer, K*365.2422)
    J = Q2-E+N-1
    X = G+K

    if J > 366 && X%4 == 0
        X = X+1
        J = J - 366
    elseif J > 365 && X%4 > 0
        X = X+1
        J = J - 365
    end

    # X is the year, J is the day in the julian calendar
    # Convert to MJD
    jd =  floor(Integer, 365.25*(X-1)) + 1721423 + J
    alpha =  floor(Integer, (jd-1867216.25)/36524.25)
    if jd < 2299161
        # Before 1582, Oct 15th
        bbeta = jd
    else
        bbeta = jd+ 1 + alpha -  floor(Integer, (alpha/4))
    end
    b = bbeta+1524
    c =  floor(Integer, (b-122.1)/365.25)
    d =  floor(Integer, 365.25*c)
    e1 =  floor(Integer, (b-d)/30.6001)
    D = b-d- floor(Integer, 30.6001*e1)
    if e1 < 14
        M = e1 - 1 else M = e1 - 13
    end
    if M > 2
        X = c - 4716 else X = c - 4715
    end
    return (X, M, D)
end

"""
    christian_to_moslem(X, M, D, gregorian=true)

"""
function christian_to_moslem(X, M, D, gregorian=true)
    # first convert the gregorian date to a julian date
    if gregorian
        if M < 3
            X = X-1
            M = M+12
        end
        alpha =  floor(Integer, X/100)
        bbeta  = 2-alpha+ floor(Integer, alpha/4)
        b =  floor(Integer, 365.25 * X) +  floor(Integer, 30.6001*(M+1)) + D + 1722519 + bbeta
        c =  floor(Integer, (b-122.1)/365.25)
        d =  floor(Integer, 365.25 * c)
        e1 =  floor(Integer, (b - d)/30.6001)
        D = b-d- floor(Integer, 30.6001 * e1)
        if e1 < 14
            M = e1 - 1 else M = e1 - 13
        end
        if M > 2
            X = c-4716
        else
            X = c-4715
        end
    end
    if X % 4 == 0
        W = 1
    else
        W = 2
    end
    N =  floor(Integer, 275 * M / 9) - W *  floor(Integer, (M + 9) / 12) + D - 30
    A = X - 623
    B =  floor(Integer, A/4)
    C = A % 4
    C1 = 365.2501 * C
    C2 =  floor(Integer, C1)
    if C1 - C2 > 0.5
        C2 = C2 + 1
    end
    Dp = 1461 * B + 170 + C2
    Q =  floor(Integer, Dp/10631)
    R = Dp % 10631
    J =  floor(Integer, R/354)
    K = R % 354
    O =  floor(Integer, (11 * J + 14)/30)
    H = 30 * Q + J + 1
    JJ = K - O + N - 1

    if JJ > 354
        CL = H % 30
        DL = (11 * CL + 3) % 30
        if DL < 19
           JJ = JJ - 354
           H=H + 1
        else
           JJ = JJ - 355
           H=H + 1
        end
        if JJ == 0
           JJ = 355
           H = H - 1
        end
    end
    S =  floor(Integer, (JJ - 1)/29.5)
    m = 1 + S
    d =  floor(Integer, JJ - 29.5 * S)
    if JJ == 355
        m = 12
        d = 30
    end
    return (H, m, d)
end

"""
    easter(yr, gregorian=true)

Return the date of Western ecclesiastical Easter for a year in the Julian or Gregorian calendars.

Parameters:
    yr        : year
    gregorian : if True, use Gregorian calendar, else use Julian calendar (default: True)

Return:
    month
    day

Julian formula found in Meeus.

Gregorian formula taken from Wikipedia (for the formula published in Nature, April 20th, 1876)
"""
function easter(yr, gregorian=true)
    yr =  floor(Integer, yr)
    if gregorian
        a = yr % 19
        b =  floor(Integer, yr / 100)
        c = yr % 100
        d =  floor(Integer, b / 4)
        e1 = b % 4
        f =  floor(Integer, (b + 8) / 25)
        g =  floor(Integer, (b - f + 1) / 3)
        h = (19 * a + b - d - g + 15) % 30
        i =  floor(Integer, c / 4)
        k = c % 4
        l = (32 + 2 * e1 + 2 * i - h - k) % 7
        m =  floor(Integer, (a + 11 * h + 22 * l) / 451)
        tmp = h + l - 7 * m + 114
    else
        a = yr % 4
        b = yr % 7
        c = yr % 19
        d = (19 * c + 15) % 30
        e1 = (2 * a + 4 * b - d + 34) % 7
        tmp = d + e1 + 114
    end
    mo =  floor(Integer, tmp / 31)
    dy = (tmp % 31) + 1
    return (mo, dy)
end

"""
    jd_to_day_of_week(jd::Float64)

Return the day of week for a Julian Day Number.

The Julian Day Number must be for 0h UT.

Parameters:
    jd : Julian Day number

Return:
    day of week: 0 = Sunday...6 = Saturday.

"""
function jd_to_day_of_week(jd)
    i = jd + 1.5
    return  floor(Integer, i) % 7
end

"""
    is_dst(jd)

Is this instant within the Daylight Savings Time period

# for Europe: last sunday of march 1:00UTC to last sunday of october 1:00UTC

if daylight_timezone_name is false, the function always returns false.

Parameters:
    jd : Julian Day number representing an instant in Universal Time

Return:
    True if Daylight Savings Time is in effect, False otherwise.
"""
function is_dst(jd::Float64)
    global daylight_timezone_name
    if daylight_timezone_name == ""
        return false
    end
    #  What year is this?
    (yr, mon, day) = jd_to_cal(jd)
    # Last day of march
    starting = cal_to_jd(yr, 3, 31)
    # Back to last sunday
    dow = jd_to_day_of_week(starting)
    if dow >=0
        starting = starting - dow
    end
    # Advance to 1AM
    starting = starting + 1.0 / 24
    # Before the beginning of the period ?
    if jd < starting
        return false
    end
    # Last day in October
    stop = cal_to_jd(yr, 10, 31)
    # Backup to the last Sunday
    dow = jd_to_day_of_week(stop)
    stop = stop - dow
    #  Advance to 1AM
    stop = stop + 1.0 / 24
    # Before the end of the period ?
    if jd < stop
        return true
    end
    # After the end of the period
    return false
end

"""
    jd_to_jcent(jd)

Return the number of Julian centuries since J2000.0

Parameters:
    jd : Julian Day number

Return:
    Julian centuries
"""
function jd_to_jcent(jd)
    return (jd - 2451545.0) / 36525.0
end

"""
    lt_to_str(jd, zone="", level="second")

This now uses Julia Dates module. Timezones are currently not done - everything's in UTC.

Convert time in Julian Days to a formatted string.

    The general format is:

        YYYY-MMM-DD HH:MM:SS ZZZ

    Truncate the time value to seconds, minutes, hours or days as
    indicated. if level = "day", don't print the time zone string.

    Pass an empty string ("", the default) for zone if you want to do
    your own zone formatting in the calling module.

    Parameters:
        jd    : Julian Day number
        zone  : Time zone string (default = "")
        level : "day" or "hour" or "minute" or "second" (default = "second")

    Return:
        formatted date/time string

"""
function lt_to_str(jd, zone="", level="second")
    (yr, mon, day) = jd_to_cal(jd)
    (fday, iday) = modf(day)
    iday =  floor(Integer, iday)
    (hr, mn, second) = fday_to_hms(fday)
    second =  floor(Integer, second)

    if level == "second"
        return Dates.format(DateTime(yr, mon, iday, hr, mn, second), Dates.DateFormat("e, dd u yyyy HH:MM:SS"))
    end

    if level == "minute"
        return Dates.format(DateTime(yr, mon, iday, hr, mn), Dates.DateFormat("e, dd u yyyy HH:MM"))
    end

    if level == "day"
        return Dates.format(DateTime(yr, mon, iday), Dates.DateFormat("e, dd u yyyy"))
    end

    error("Unknown time level = $level")
end

"""
    ut_to_lt(jd)

Convert universal time in Julian Days to a time.

Include Daylight Savings Time offset, if any.

Parameters:
    jd : Julian Day number, universal time

Return:
    Julian Day number, time
    zone string of the zone used for the conversion

not yet doing time zones... :(

"""
function ut_to_lt(jd)
    global daylight_timezone_name, daylight_timezone_offset, standard_timezone_name, standard_timezone_offset
    if is_dst(jd)
        zone   = daylight_timezone_name
        offset = daylight_timezone_offset
    else
        zone   = standard_timezone_name
        offset = standard_timezone_offset
    end
    return (jd - offset/24, zone)
end


module TestGroupGenerics

using Test
using RLEVectors

# Test function for broadcasting
count = 0
function simple_add(x)
    global count += 1
    x + 100
end

@testset begin

    # compare group
    vec = [1, 1, 2, 2, 4, 4]
    x = RLEVector(vec)
    @test (x .< 3) == RLEVector([true, false], [4, 6])

    # math group
    vec = [1, 1, 2, 2, 4, 4]
    x = RLEVector(vec)
    @test x .+ 5 == 5 .+ x
    @test x .+ 4 == RLEVector([5, 5, 6, 6, 8, 8])
    @test div.(vec, 2) == collect(div.(x, 2))
    @test median(RLEVector([1, 2, 3, 2, 1, 5, 4])) == median([1, 2, 3, 2, 1, 5, 4])
    @test median(RLEVector([1, 2, 3, 2, 1, 5])) == median([1, 2, 3, 2, 1, 5])
    @test median(RLEVector([3, 2, 1])) == 2.0
    @test median(RLEVector([2, 2, 2, 3, 3, 3])) == 2.5 # median is average of end of run and next value
    @test sum(RLEVector([4, 4, 5, 5, 6, 6])) == 30
    @test mean(RLEVector([4, 4, 5, 5, 6, 6])) == 5.0
    @test x .^ 2 == RLEVector([1, 1, 4, 4, 16, 16])
    @test x .^ 3 == RLEVector([1, 1, 8, 8, 64, 64])
    @test x + x == RLEVector([2, 2, 4, 4, 8, 8])
    @test x - x == RLEVector([0, 0, 0, 0, 0, 0])
    rle = RLEVector([4, 4, 9, 9, 16, 16])
    @test sqrt.(rle) == RLEVector([2, 2, 3, 3, 4, 4])

    # math on bools
    vec = [1, 1, 2, 2, 4, 4]
    x = RLEVector(vec)
    @test x .+ true == x .+ 1
    @test x .+ false == x
    vec = [1, 1, 2, 2, 4, 4]
    x = RLEVector(vec)
    @test x .+ true == x .+ 1
    @test x .+ false == x

    # broadcast with different sizes
    x = RLEVector([1, 1, 3])
    @test x .+ 4 == RLEVector([5, 5, 7])
    @test x .+ [4 5 6] == [[5 6 7]; [5 6 7]; [7 8 9]]
    @test x .+ [4, 5, 6] == RLEVector([5, 6, 9])
    @test x .+ [[2 2 2]; [4 4 4]]' == [[3 3 5]; [5 5 7]]'

    # test that broadcast only happens once per element
    x = RLEVector([1,1,1,2,2,3]) # 6 elements, 3 runs
    y = simple_add.(x)
    @test count == 6

    # findmax, findmin
    @test findmin(RLEVector([1, 2, 3, 4, 1, 1])) == findmin([1, 2, 3, 4, 1, 1])
    @test findmax(RLEVector([1, 2, 3, 4, 1, 1])) == findmax([1, 2, 3, 4, 1, 1])

    # indexin
    foo = IntegerRle(collect(1:1000), collect(5:5:5000))
    x = RLEVector([2, 2, 4, 4, 3, 3])
    y = RLEVector([0, 0, 0, 3, 3, 3, 4, 4])
    @test indexin(x, y) == indexin(collect(x), collect(y))
    @test indexin(x, collect(3:11)) == indexin(collect(x), collect(3:11))
    @test indexin([200, 200, 1, 1, 5, 5], foo) == indexin([200, 200, 1, 1, 5, 5], foo)

    @test findall(in(RLEVector(collect(3:10))), RLEVector([1, 1, 2, 2, 3, 3])) == collect(5:6)
    @test findall(in(3:10), RLEVector([1, 1, 2, 2, 3, 3])) == collect(5:6)
    @test findall(in(collect(3:10)), RLEVector([1, 1, 2, 2, 3, 3])) == collect(5:6)
    @test findall(in(RLEVector(collect(1:4))), [3, 4, 5]) == [1, 2]

    # in
    @test in(3, RLEVector([1, 2, 2, 3])) == true
    @test in(4, RLEVector([1, 2, 2, 3])) == false
    @test setdiff(Set([1, 2, 3, 4, 5]), RLEVector([1, 1, 2, 2, 4, 4, 5, 5])) == Set([3])

end # testset

end # module


function polyfit(x, y, n)
  A = [ float(x[i])^p for i = 1:length(x), p = 0:n ]
  A \ y
end

function polyval(p,x::Number)
	y = zero(promote_type(eltype(p), eltype(x)))
	for i=1:length(p)
		y+=p[i]*x^(i-1)
	end
	y
end
function polyval(p,x::AbstractArray)
	y = zeros(promote_type(eltype(p), eltype(x)), length(x))
	for i = 1:length(x)
		y[i] = polyval(p,x[i])
	end
	y
end
function polyresiduals(p,x,y)
	yp = polyval(p,x)
	residuals = y-yp
end

### CALIBRATION
function peakassignlog(locations_arb, locations_true)
# locations_arb is an array of possible peak locations in arbitrary units, assumed to have an unknown,
# nonlinear (but monotonic and smooth) relationship to locations_true
# locations_arb may contain more peaks than locations_true, including some spurious peaks
# the goal is to find the best possible set of assignments for entries in locations_arb to
# correspond to locations_true by examining all possible assignments and calcualting some figure of merit
@assert(issorted(locations_arb))
@assert(issorted(locations_true))
@assert(length(locations_arb)>=length(locations_true))
combos = combinations(1:length(locations_arb), length(locations_true))
length(combos)==1 && return (locations_arb[first(combos)], 0.0)

assert(!(0 in locations_true)) # don't pass zero and force zero
assert(!(0 in locations_arb))

log_locations_arb = log(locations_arb)
log_locations_true = log(locations_true)

best_combo = first(combos)
best_fom = realmax(Float64)
for c in combos
	cu = log_locations_arb[c]
	fom = peakassign_fig_of_merit2(cu, log_locations_true)
	if fom < best_fom
		best_fom = fom
		best_combo = c
	end
end
locations_arb[best_combo], best_fom
end

function peakassign2(locations_arb, locations_true)
# locations_arb is an array of possible peak locations in arbitrary units, assumed to have an unknown,
# nonlinear (but monotonic and smooth) relationship to locations_true
# locations_arb may contain more peaks than locations_true, including some spurious peaks
# the goal is to find the best possible set of assignments for entries in locations_arb to
# correspond to locations_true by examining all possible assignments and calcualting some figure of merit
@assert(issorted(locations_arb))
@assert(issorted(locations_true))
@assert(length(locations_arb)>=length(locations_true))
combos = combinations(1:length(locations_arb), length(locations_true))
length(combos)==1 && return (locations_arb[first(combos)], 0.0)

assert(!(0 in locations_true)) # don't pass zero and force zero
assert(!(0 in locations_arb))

best_combo = first(combos)
best_fom = realmax(Float64)
for c in combos
	cu = locations_arb[c]
	fom = peakassign_fig_of_merit2(cu, locations_true)
	if fom < best_fom
		best_fom = fom
		best_combo = c
	end
end
locations_arb[best_combo], best_fom
end

function peakassign2zeros(locations_arb, locations_true)
# locations_arb is an array of possible peak locations in arbitrary units, assumed to have an unknown,
# nonlinear (but monotonic and smooth) relationship to locations_true
# locations_arb may contain more peaks than locations_true, including some spurious peaks
# the goal is to find the best possible set of assignments for entries in locations_arb to
# correspond to locations_true by examining all possible assignments and calcualting some figure of merit
@assert(issorted(locations_arb))
@assert(issorted(locations_true))
@assert(length(locations_arb)>=length(locations_true))
combos = combinations(1:length(locations_arb), length(locations_true))
length(combos)==1 && return (locations_arb[first(combos)], 0.0)

assert(!(0 in locations_true)) # don't pass zero and force zero
assert(!(0 in locations_arb))
locations_true = unshift!(copy(locations_true), 0)

best_combo = first(combos)
best_fom = realmax(Float64)
for c in combos
	cu = locations_arb[c]
	unshift!(cu,0)
	fom = peakassign_fig_of_merit2(cu, locations_true)
	if fom < best_fom
		best_fom = fom
		best_combo = c
	end
end
locations_arb[best_combo], best_fom
end

function peakassign2(locations_arb, locations_true)
# locations_arb is an array of possible peak locations in arbitrary units, assumed to have an unknown,
# nonlinear (but monotonic and smooth) relationship to locations_true
# locations_arb may contain more peaks than locations_true, including some spurious peaks
# the goal is to find the best possible set of assignments for entries in locations_arb to
# correspond to locations_true by examining all possible assignments and calcualting some figure of merit
@assert(issorted(locations_arb))
@assert(issorted(locations_true))
@assert(length(locations_arb)>=length(locations_true))
combos = combinations(1:length(locations_arb), length(locations_true))
length(combos)==1 && return (locations_arb[first(combos)], 0.0)

assert(!(0 in locations_true)) # don't pass zero and force zero
assert(!(0 in locations_arb))

best_combo = first(combos)
best_fom = realmax(Float64)
for c in combos
	cu = locations_arb[c]
	fom = peakassign_fig_of_merit(cu, locations_true)
	if fom < best_fom
		best_fom = fom
		best_combo = c
	end
end
locations_arb[best_combo], best_fom
end

function peakassign_fig_of_merit(x, y)
	# average fractional error trying to predict the y value of the middle of 3 points using the slope
	# from the outer two points and the x value of the middle point
	err = 0.0
	for i = 1:length(y)-2
		yip1_prediction = y[i]+(y[i+2]-y[i])*(x[i+1]-x[i])/(x[i+2]-x[i])
		err += abs(1-y[i+1]/yip1_prediction)
	end
	err/length(y-2)
end

function peakassign_fig_of_merit2(x, y)
	n = length(x) > 2 ? 2 : 1
	p = polyfit(x,y,n)
	std(polyresiduals(p,x,y)./y)
end

# test peakassign

known_energies = [3691, 4090, 4510, 4952, 5414, 5898, 6403, 6929, 7477, 8040, 8637] # ka energies Ca-Zn

filt_values = known_energies.^0.9 # caluclate filt values with some amount of non-linearity

tested = 0
peakassignlogwrong=0
peakassign2wrong=0
peakassign2zerowrong=0
peakassignwrong=0

nc = 7
	tested+=length(combinations(1:length(known_energies),nc))
	println("testing peakassignlog $nc")
	for c in combinations(1:length(known_energies),nc)
		locations_arb = filt_values
		locations_true = known_energies[c]
		best_combo, best_fom = peakassignlog(locations_arb, locations_true)
		if !all(best_combo .== filt_values[c])
			peakassignlogwrong+=1
			# println(locations_true)
			# println(best_combo)
			# println(filt_values[c])
			# @test best_fom < 1
			# @test all(best_combo .== filt_values[c])
		end
	end

	println("testing peakassign2 $nc")
	for c in combinations(1:length(known_energies),nc)
		locations_arb = filt_values
		locations_true = known_energies[c]
		best_combo, best_fom = peakassign2(locations_arb, locations_true)
		if !all(best_combo .== filt_values[c])
			peakassign2wrong+=1

		end
	end

	println("testing peakassign2zeros $nc")
	for c in combinations(1:length(known_energies),nc)
		locations_arb = filt_values
		locations_true = known_energies[c]
		best_combo, best_fom = peakassign2zeros(locations_arb, locations_true)
		if !all(best_combo .== filt_values[c])
			peakassign2zerowrong+=1

		end
	end
	println("testing peakassign $nc")
	for c in combinations(1:length(known_energies),nc)
		locations_arb = filt_values
		locations_true = known_energies[c]
		best_combo, best_fom = peakassign2(locations_arb, locations_true)
		if !all(best_combo .== filt_values[c])
			peakassignwrong+=1

		end
	end

println("peakassignlog $(peakassignlogwrong/tested)")
println("peakassign2 $(peakassign2wrong/tested)")
println("peakassign2zeros $(peakassign2zerowrong/tested)")
println("peakassign $(peakassignwrong/tested)")


using QuasiArrays, DomainSets, StaticArrays, Test

@testset "Kron" begin
    @testset "InclusionKron" begin
        a = Inclusion(0:0.5:2)
        b = Inclusion([1,3,4])
        c = quasikron(a,b)
        @test eltype(c) == SVector{2,Float64}
        @test c[SVector(0.5,3)] == SVector(0.5,3)
        @test_throws BoundsError c[SVector(0.6,3)]
        @test_throws BoundsError c[SVector(0.5,2)]
    end

    @testset "QuasiKron" begin
        a = QuasiVector(randn(5), 0:0.5:2)
        b = QuasiVector([5,6,8], [1,3,4])
        K = quasikron(a, b)
        @test axes(K) == (quasikron(axes(a,1),axes(b,1)),)
        @test K[SVector(0.5,3)] == a[0.5]b[3]
        @test_throws BoundsError K[SVector(0.6,3)]
        @test_throws BoundsError K[SVector(0.5,2)]
    end
end


module Blobs

export Blob

using ..Elements
import ..Motions: induce_velocity, mutually_induce_velocity!, self_induce_velocity!, advect,
                  allocate_velocity

#== Type definition ==#

"""
    Blob <: Elements.Element

An immutable structure representing a regularized point source/vortex

## Fields
- `z`: position
- `S`: strength/circulation
- `δ`: regularization radius

## Constructors

"""
struct Blob{S <: Number,R <: Real} <: Element
    z::Complex{R}
    S::Union{R,Complex{R}}
    δ::Float64
    Blob{S}(z::Complex{R}, s::T, δ)           where {S <: Complex, T <: Real, R <: Real} = (U = promote_type(R,T); new{Complex{U},U}(z, im*convert(U,s),δ))
    Blob{S}(z::Complex{R}, s::Complex{T}, δ)  where {S <: Complex, T <: Real, R <: Real} = (U = promote_type(R,T); new{Complex{U},U}(z, convert(Complex{U},s),δ))
    Blob{S}(z::Complex{R}, s::T, δ)           where {S <: Real, T <: Real, R <: Real}    = (U = promote_type(R,T); new{U,U}(z, convert(U,s),δ))
end
# struct Blob{T <: Number, R <: Number} <: Element
#     z::Complex{R}
#     S::T
#     δ::Float64
#     Blob{T,R}(z, s::Real, δ) where {T <: Complex, R <: Number} = new(z, im*s, δ)
#     Blob{T,R}(z, s, δ)       where {T <: Complex, R <: Number} = new(z, s, δ)
#     Blob{T,R}(z, s::Real, δ) where {T <: Real, R <: Number}    = new(z, s, δ)
# end

Elements.kind(::Blob) = Singleton
Elements.kind(::Type{Blob{T,R}}) where {T,R} = Singleton

Elements.property_type(::Blob{T,R}) where {T,R} = R
Elements.property_type(::Type{Blob{T,R}}) where {T,R} = R

#Elements.promote_property_type(::Blob{T,R}) where {T,R} = promote_type(T,R)
#Elements.promote_property_type(::Type{Blob{T,R}}) where {T,R} = promote_type(T,R)

#(b::Blob{T,R})(; z = b.z, S = b.S, δ = b.δ) where {T,R} = Blob{T,R}(z, S, δ)

#== Methods to be extended ==#

Elements.position(b::Blob) = b.z
Elements.blobradius(b::Blob) = b.δ
Elements.streamfunction(z::Complex{T}, b::Blob) where {T} = real(-0.5b.S*log(z - b.z)/π)

Elements.complexpotential(z::Complex{T}, b::Blob) where {T} = -0.5im*b.S*log(z - b.z)/π

blob_kernel(z, δ) = 0.5im*z/(π*(abs2(z) + δ^2))

# ensures that the velocity is of same type as position
allocate_velocity(v::Vector{Blob{T,R}}) where {T,R} = zeros(Complex{R},length(v))

function induce_velocity(z::Complex{T}, b::Blob, t) where {T}
    b.S'*blob_kernel(z - b.z, b.δ)
end

function induce_velocity(target::Blob, source::Blob, t)
    δ = √(0.5(target.δ^2 + source.δ^2))
    source.S'*blob_kernel(target.z - source.z, δ)
end

function mutually_induce_velocity!(ws₁, ws₂,
                                   blobs₁::Vector{Blob{T₁}},
                                   blobs₂::Vector{Blob{T₂}}, t) where {T₁, T₂}
    for (s, source) in enumerate(blobs₁)
        for (t, target) in enumerate(blobs₂)
            δ = √(0.5(target.δ^2 + source.δ^2))
            K = blob_kernel(target.z - source.z, δ)
            ws₂[t] += source.S'*K
            ws₁[s] -= target.S'*K
        end
    end
    nothing
end

function self_induce_velocity!(ws, blobs::Vector{Blob{T,R}}, t) where {T,R}
    N = length(blobs)

    for s in 1:N, t in s+1:N
        δ = √(0.5(blobs[t].δ^2 + blobs[s].δ^2))
        K = blob_kernel(blobs[t].z -blobs[s].z, δ)
        ws[t] += blobs[s].S'*K
        ws[s] -= blobs[t].S'*K
    end
    ws
end

function advect(b::Blob{T}, w::ComplexF64, Δt::Float64) where {T}
    Blob{T}(b.z + w*Δt, b.S, b.δ)
end
#
#function Base.show(io::IO, b::Blob)
#    print(io, "Vortex Blob: z = $(round(b.z, digits=3)), Γ = $(round(b.Γ, digits=3)), δ = $(round(b.δ, digits=3))")
#end

end


using DiffEqFlux, Flux, Optim, OrdinaryDiffEq

n = 5  # number of ODEs
tspan = (0.0, 1.0)

d = 5  # number of data pairs
x_data = rand(n, 5)
y_data = rand(n, 5)
training_data = Iterators.cycle([(x_data[:, i], y_data[:, i]) for i in 1:d])

NN = Chain(Dense(n, 10n, tanh),
           Dense(10n, n))

@info "ROCK4"
nODE = NeuralODE(NN, tspan, ROCK4(), reltol=1e-4, saveat=[tspan[end]])

loss_function(θ, x, y) = Flux.mse(y, nODE(x, θ))
res = DiffEqFlux.sciml_train(loss_function, nODE.p, LBFGS(), training_data)
res = DiffEqFlux.sciml_train(loss_function, nODE.p, NewtonTrustRegion(), training_data)

NN = SlowChain(SlowDense(n, 10n, tanh),
               SlowDense(10n, n))

@info "ROCK2"
nODE = NeuralODE(NN, tspan, ROCK2(), reltol=1e-4, saveat=[tspan[end]])

loss_function(θ, x, y) = Flux.mse(y, nODE(x, θ))
res = DiffEqFlux.sciml_train(loss_function, nODE.p, LBFGS(), training_data)
res = DiffEqFlux.sciml_train(loss_function, nODE.p, NewtonTrustRegion(), training_data)

@info "ROCK4"
nODE = NeuralODE(NN, tspan, ROCK4(), reltol=1e-4, saveat=[tspan[end]])

loss_function(θ, x, y) = Flux.mse(y, nODE(x, θ))
res = DiffEqFlux.sciml_train(loss_function, nODE.p, LBFGS(), training_data)
res = DiffEqFlux.sciml_train(loss_function, nODE.p, NewtonTrustRegion(), training_data)


"""
    Lagrange{N, S, R, T}

Represent a polynomial in Lagrange form using nodes `xs`, weights `ws`, and coefficients `coeffs`.
The Lagrange form does polynomial interpolation between `xs` and `ys` through `p(x) = Σ_{0..n} ℓ_i(x) y_i`,
where if `ℓ(x) = prod(x-x_i)`, `w_i = 1/prod_{j≠i}(x_i - x_j)`, then `ℓ_i(x) = ℓ(x) w_i/(x-x_i)`. The `ℓ_i`
satisfy `ℓ_i(x_j) = δ_{ij}`, so the coefficients are just the `ys`.

```jldoctest Lagrange
julia> using Polynomials, SpecialPolynomials

julia> p =  Lagrange([1,2,3], [1,2,3])
Lagrange(1⋅ℓ^2_0(x) + 2⋅ℓ^2_1(x) + 3⋅ℓ^2_2(x))

julia> p.([1,2,3]) # the coefficients
3-element Array{Int64,1}:
 1
 2
 3

julia> convert(Polynomial,  p)
Polynomial(1.0*x)
```

The instances hold the nodes and weights, which are necessary for
representation, so the type alone can not be used for functions such
as `variable` or `convert(Lagrange, ...)`. For the former we can  use an instance, for the latter we can use `fit`:

```jldoctest Lagrange
julia> p =  Lagrange([1,2,3], [1,2,3])
Lagrange(1⋅ℓ^2_0(x) + 2⋅ℓ^2_1(x) + 3⋅ℓ^2_2(x))

julia> variable(p)
Lagrange(1⋅ℓ^2_0(x) + 2⋅ℓ^2_1(x) + 3⋅ℓ^2_2(x))

julia> q = Polynomial([0,0,1])
Polynomial(x^2)

julia> qq = fit(Lagrange, p.xs, q)
Lagrange(1⋅ℓ^2_0(x) + 4⋅ℓ^2_1(x) + 9⋅ℓ^2_2(x))

julia> convert(Polynomial, qq)
Polynomial(1.0*x^2)
```

Interpolating polynomials suffer from the Runge phenomenon unless the
nodes are well chosen. For `P=Chebyshvev` and `P=ChebyshevU`, the
function `SpecialPolynomials.lagrange_barycentric_nodes_weights(P, n)`
will return a good choice of `n+1` points over `[-1,1]` along with
precomputed weights.

```jldoctest Lagrange
julia> xs, ws = SpecialPolynomials.lagrange_barycentric_nodes_weights(Chebyshev, 64);

julia> f(x) = exp(-x)*sinpi(x)
f (generic function with 1 method)

julia> p = fit(Lagrange, xs, f.(xs));

julia> degree(p)
64

julia> maximum(abs.(f(x) - p(x) for x in range(-1, 1, length=20))) <= 1e-14
true
```

!!! note
    The above example  is  more directly  done through `fit(Chebyshev, f, 64)`, though  the resulting
    polynomial will reference a different  basis.

"""
struct Lagrange{N, S<:Number, R <: Number, T <: Number} <: AbstractInterpolatingPolynomial{T}
    xs::Vector{S}
    ws::Vector{R}
    coeffs::Vector{T}
    var::Symbol
    function  Lagrange(xs::Vector{S}, ws::Vector{R}, coeffs::Vector{T}, var::Symbol=:x) where {S,R,T}
        xs = unique(xs)
        N = length(xs)
        N == length(coeffs) || throw(ArgumentError("the unique xs and the coeffs must have the same length"))
        new{N, S,R,T}(xs, ws, coeffs, var)
    end
    function  Lagrange(xs::Vector{S}, coeffs::Vector{T}, var::Symbol=:x) where {S,T}
        xs = unique(xs)
        N = length(xs)
        ws = lagrange_barycentric_weights(xs)
        R = eltype(ws)
        length(coeffs) == length(xs) || throw(ArgumentError("the unique xs and the coeffs must have the same length"))
        new{N, S, R, T}(xs, ws, coeffs, :x)
    end
end

export Lagrange

basis_symbol(::Type{Lagrange{N,S,R,T}}) where {N,S,R,T} = "ℓ^$(N-1)"

## Boilerplate code reproduced here, as there are three type parameters
Base.convert(::Type{P}, p::P) where {P <: Lagrange} =  p
Base.convert(::Type{Lagrange{N,S,R,T}}, p::Lagrange) where {N,S,R,T} = Lagrange{N,S,R,T}(p.xs, p.ws, coeffs(p), p.var)
Base.promote_rule(::Type{Lagrange{N, S, R, T}}, ::Type{Lagrange{N, S, R, Q}}) where {N, S, R, T, Q <: Number} = Lagrange{N, S, R, promote_type(T,Q)}
Base.promote_rule(::Type{Lagrange{N, S, R, T}}, ::Type{Q}) where {N, S, R, T, Q <: Number} = Lagrange{N, S, R, promote_type(T,Q)}

Polynomials.domain(::Type{<:Lagrange}) = Polynomials.Interval(-Inf, Inf)
function Polynomials.variable(p::Lagrange{S, R, T}, var::Polynomials.SymbolLike=:x) where {S, R, T}
    _fit(Lagrange, p.xs, p.ws, p.xs, var)
end
function Polynomials.one(p::Lagrange)
    _fit(Lagrange, p.xs, p.ws, ones(eltype(p.xs),length(p.xs)), p.var)
end
Polynomials.zero(p::Lagrange{N, S, R, T}) where {N,S,R,T} = 0*p

##  Evaluation
##  From https://people.maths.ox.ac.uk/trefethen/barycentric.pdf
function (p::Lagrange)(x::Number)
    ws, xs, cs = p.ws, p.xs, coeffs(p)
    a = b = zero(x/1)
    for j in eachindex(xs)
        Δ = (x - xs[j])
        iszero(Δ) && return cs[j] # xj ∈ xs, so sum c_j l_i(xj) = sum c_j δ_{ij} = c_j
        l = ws[j]/Δ
        a += l * cs[j]
        b += l
   end
   a/b
end

## convert to poly
## avoids division in barycentric form, which isn't available for Polynomial  arguments, such as `x-xs[j]`
## This has numeric issues, unlike the barycentric form used for evaluation
function  Base.convert(Q::Type{<:Polynomial},  p::Lagrange{N, S,R,T})  where  {N, S, R, T}
    q = zero(Q)/1
    x =  variable(q)
    xs = p.xs
    cs = coeffs(p)
    ws = p.ws

    if length(xs) == 1
       return cs[1] + 0*x #Polynomial(cs[1]*ones(T,1))
    else
        l = fromroots(Polynomial, xs)
        for i in  eachindex(xs)
            li = prod(x - xs[j] for j in eachindex(xs) if j != i);
            q += li * ws[i] * cs[i]
        end
    end
    q
end

## return w⁻¹s = {1/w_j, 0 <= j <= n} as 1 based ws[j+1] = w_j
"""
   lagrange_barycentric_weights(xs)

return [1/w_i for 0 <i <= n] where wi = ∏_{j≠i} (xi - xj)
"""
function lagrange_barycentric_weights(xs)
    n =  length(xs)
    ws =  ones(eltype(xs), n)
    for j in 2:n
        for k in 1:j
            ws[k] =  (xs[k] - xs[j]) * ws[k]
        end
        ws[j] = prod(xs[j]-xs[i] for i in 1:j-1)
    end
    1 ./ ws
end

# https://people.maths.ox.ac.uk/trefethen/barycentric.pdf (5.1)
function lagrange_barycentric_weights(xs::Union{AbstractRange{T},AbstractUnitRange{T}}) where {T}
    n = length(xs)-1
    ws = zeros(float(T), n+1)
    o = one(T)
    for j = 0:n
        ws[1+j] = o * binomial(n, j)
        o *= -1
    end
    ws
end

##  add a new point
## check for  xk ∈ xs?
"""
    update_lagrange_barycentric_weights(ws, xs, xn1)

Adds `xn1` to the nodes `xs` and updates the weights `ws` accordingly.
"""
function update_lagrange_barycentric_weights(ws, xs, xn1)
    ys = [wj/(xj - xn1) for (wj,xj) in zip(ws, xs)]
    wn1 = one(eltype(ws))
    for xj in xs
        wn1 *= (xn1 - xj)
    end
    push!(ys, 1/wn1)
    ys
end

"""
    lagrange_barycentric_nodes_weights(::Type{<:SpecialPolynomial}, n::Int)

Return a collection of `n+1` nodes and weights the given family of
polynomials. Defined for `ChebyshevT` and `ChebyshevU` familes, as there
are `O(n)` algorithms available. Otherwise, `lagrange_barycentric_weights`
is needed.

"""
lagrange_barycentric_nodes_weights(::Type{<:AbstractSpecialPolynomial}, n::Int) = throw(MethodError())

## can add easily if the nodes are shared
## otherwise, this needs to interpolate
function Base.:+(p1::Lagrange{N,S,R,T}, p2::Lagrange{M,S1,R1,T1}) where {N,S,R,T,M,S1,R1,T1}
    p1.var == p2.var || throw(ArgumentError("p1 and p2 must share the same variable"))

    # can just add if using the same set of nodes
    if N == M && p1.xs == p2.xs
        cs =  coeffs(p1) +  coeffs(p2)
        return Lagrange(p1.xs, p1.ws, cs, p1.var)
    else
        p,q = N >= M ? (p1, p2) : (p2, p1)
        xs, ws, cs = p.xs, p.ws, copy(coeffs(p))
        for i in eachindex(xs)
            cs[i] += q(xs[i])
        end
        return _fit(Lagrange, xs,  ws, cs, p1.var)
    end
end

## XXX is there a better way?
## This refits using  updated weights,
## but this can be wildly inaccurate
function Base.:*(p1::Lagrange{N,S,R,T}, p2::Lagrange{M,S1,R1,T1}) where {N,S,R,T,M,S1,R1,T1}

    p1.var == p2.var || throw(ArgumentError("p1 and p2 must share the same variable"))

    p, q =  N >= M ? (p1, p2)  : (p2, p1)
    S2 = typeof(one(S) * one(S1)/1)
    xs = convert(Vector{S2}, copy(p.xs))
    ws = copy(p.ws)

    # new_xs ## XXX this is where the work is
    new_xs = _new_nodes(xs, q.xs)
    for y in new_xs
        ws = update_lagrange_barycentric_weights(ws, xs, y)
        push!(xs, y)
    end
    ys = p.(xs) .* q.(xs)
    _fit(Lagrange, xs,  ws, ys, p1.var)

end

# handle scalar case
function Base.:+(p::Lagrange, c::Number)
    q = Lagrange(p.xs[1:1], [c], p.var)
    p + q
end

Base.:*(c::Number, p::Lagrange) = p*c
function Base.:*(p::P, c::Number) where {P <: Lagrange}
    Lagrange(p.xs, c*coeffs(p), p.var)
end

Base.:-(p::P) where {P<:Lagrange} = (-1)*p

## short cut if weights are known.
_fit(P::Type{Lagrange}, xs, ws, ys, var=:x) = Lagrange(xs, ws, ys, var)

function Polynomials.fit(P::Type{<:Lagrange},
                         xs::AbstractVector{S},
                         ys::AbstractVector{T};
                         var = :x) where {S, T}
    ws = lagrange_barycentric_weights(xs)
    _fit(P, xs, ws, ys, var)
end

function Polynomials.fit(P::Type{<:Lagrange}, xs::AbstractVector{S}, f; var=:x) where {S}
     ws = lagrange_barycentric_weights(xs)
    _fit(P, xs, ws, f.(xs), var)
end


using CoordinateTransformations
using LinearAlgebra
using Optim

const XYPosition = SArray{Tuple{2},Float64,1,2}

"""
    translate(oldPts::AbstractVector{XYPosition}, newPts::AbstractVector{XYPosition})::AffineMap

Determines the optimal AffineMap to transform the X-Y coordinates in `oldPts` into the X-Y coordinates in `newPts`.
The coordinates in `oldPts` and `newPts` are assumed to represent the same features in the same order and thus
must contain the same number of X-Y coordinates.
"""
function translate(oldPts::AbstractVector{XYPosition}, newPts::AbstractVector{XYPosition})::AffineMap
    @assert length(oldPts)==length(newPts) "The number of old and new points must match."
    if length(oldPts)==1
        offset, scale, rotatio = SA[ newPts[1] .- oldPts[1] ], SA[1.0, 1.0], 0.0
    else # Compute a first estimate from points 1 and 2
        offset = 0.5*(newPts[2] .+ newPts[1]) .- 0.5*(oldPts[2] .+ oldPts[1])
        dold, dnew = oldPts[2] .- oldPts[1], newPts[2] .- newPts[1]
        scale = SA[ norm(dnew)/norm(dold), norm(dnew)/norm(dold) ]
        rotation = acos((dold⋅dnew)/(norm(dold)*norm(dnew)))
    end
    res = AffineMap(SA[ cos(rotation) -sin(rotation) ; sin(rotation) cos(rotation) ] * SA[ scale[1] 0.0; 0.0 scale[2]], offset)
    if length(oldPts)>2
        # Simplex optimize to discover the minimum offset, scale and rotation over all point pairs
        x0 = [offset..., scale..., rotation]
        function f(x)
            res = AffineMap(SA[ cos(x[5]) -sin(x[5]) ; sin(x[5]) cos(x[5]) ] * SA[ x[3] 0.0; 0.0 x[4] ], SA[x[1], x[2]])
            return sum(norm(res(oldPt) .- newPt) for (oldPt, newPt) in zip(oldPts, newPts))
        end
        xp = Optim.minimizer(optimize(f, x0, Optim.NelderMead()))
        res = AffineMap(SA[ cos(xp[5]) -sin(xp[5]) ; sin(xp[5]) cos(xp[5]) ] * SA[ xp[3] 0.0; 0.0 xp[4] ], SA[xp[1], xp[2]])
    end
    return res
end


export
    # cudnnSoftmaxMode_t
    CUDNN_SOFTMAX_MODE_INSTANCE,
    CUDNN_SOFTMAX_MODE_CHANNEL,

    # cudnnSoftmaxAlgorithm_t
    CUDNN_SOFTMAX_FAST,
    CUDNN_SOFTMAX_ACCURATE,
    CUDNN_SOFTMAX_LOG

function softmax{T}(x::CuArray{T}; algo=CUDNN_SOFTMAX_ACCURATE, mode=CUDNN_SOFTMAX_MODE_CHANNEL)
    h = handle(x)
    dims1 = ntuple(_ -> 1, 4-ndims(x))
    x4d = reshape(x, dims1..., size(x)...)
    xdesc = TensorDesc(x4d)
    y = similar(x)
    cudnnSoftmaxForward(h, algo, mode, T[1], xdesc, x, T[0], xdesc, y)

    function backward!(dy, dx)
        isvoid(dx) && return
        cudnnSoftmaxBackward(h, algo, mode, T[1], xdesc, y, xdesc, dy, T[1], xdesc, dx)
    end
    y, backward!
end

logsoftmax(x) = softmax(x, algo=CUDNN_SOFTMAX_LOG)


using SafeTestsets
using ThreeBodyDecay

# ThreeBodyDecay structure
@safetestset "Test of the tbs structure" begin
    include("test_threebodymasses.jl")
    include("test_threebodyspins.jl")
    include("test_threebodyparities.jl")
    include("test_invariants.jl")
end

include("wignerrotationdispatch.jl")

@safetestset "four vectors" begin include("test_generation.jl") end

# angular functions
@safetestset "Angular Functions block" begin
    include("circle_relations.jl")
    include("test_representaiton_property.jl")
    include("check_lineshape_functions.jl")
end

@safetestset "Permutations" begin include("test_wigner_permutations.jl") end

@safetestset "Sum rule" begin include("wigner_angle_sumrules.jl") end

# Lineshapes
@safetestset "Lineshapes" begin
    # include("breit-wigner.jl")
    include("scattering_length_X3872.jl")
end

# JPC functions
@safetestset "Couplings" begin
    include("test_couplings.jl")
    include("test_recoupling.jl")
    include("test_ls_amplitude.jl")
    include("test_sum_over_polarization.jl")
end

# Integrals
@safetestset "Integrals" begin
    include("test_three_body_phase_space_integral.jl")
end
# examples
@safetestset "Complete example" begin include("test_Xic_decay.jl") end

@safetestset "square dalitz" begin
    include("test_transformations_of_squaredalitz.jl")
end

include("test_inference.jl")


# Run ttv computations:
#include("stepsize_precision.jl")

# Now run "high-resolution" simulations:
#dq = ttv_elements!(n,t0,h/16.,tmax,elements,tt1,count,0.0,0,0,rstar;fout="test_output_highres.txt",iout=1)
#dqbig = ttv_elements!(n,t0big,hbig/16,tmaxbig,elements_big,tt1big,count,big(0.0),0,0,rstarbig;fout="test_output_highres_big.txt",iout=1)

# Fits a 2D Keplerian to a set of data
using PyPlot
using Optim
using DelimitedFiles
using Statistics

# Kepler solver:
include("../../src/kepler.jl")

function return_xy(t,elements)
  P,ecc,tp,semi,omega = elements
  if ecc >= 1.0
    return 0.0,0.0
  end
  # First, compute M:
  M = (t-tp)*2pi/P
  # Now, compute eccentric anomaly:
  ekep = ekepler(M,ecc)
  # Now compute xp & yp (the planar values aligned with pericenter):
  xp = semi*(cos(ekep)-ecc)
  yp = semi*sqrt(1-ecc^2)*sin(ekep)
  # Now rotate:
  x = xp*cos(omega)-yp*sin(omega)
  y = xp*sin(omega)+yp*cos(omega)
return x,y
end

function fit_keplerian(t,x,y,elements0)

  # elements are given by: [P,ecc,tp,semi,omega]
  # Chi-square function:
  function chi_square(elements)
    # Call model:
    chisq = 0.0
    for i=1:length(t)
      xmod,ymod = return_xy(t[i],elements)
      chisq += (x[i]-xmod)^2+(y[i]-ymod)^2
    end
    #println("elements: ",elements," chi^2: ",chisq)
  return chisq
  end

  # Optimize the fit:
  result =  optimize(chi_square,elements0,LBFGS(),autodiff=:forward)
#  result =  optimize(chi_square,elements0,LBFGS())
#  result =  optimize(chi_square,elements0)
  elements_best = result.minimizer
  nt = length(t)
  xmod = zeros(nt)
  ymod = zeros(nt)
  for i=1:nt
     xmod[i],ymod[i] = return_xy(t[i],elements_best)
  end
return xmod,ymod, result
end

data_hr  = readdlm("test_output_highres.txt")
elements_all = readdlm("elements.txt",',')
semi = median(sqrt.(data_hr[:,5].^2 .+data_hr[:,7].^2))
period = elements_all[2,2]
elements_best = zeros(5)
chi_best = Inf
chisq = zeros(20,20)
tobs = data_hr[:,1]; xobs  = data_hr[:,5]-data_hr[:,2]; yobs = -(data_hr[:,7]-data_hr[:,4])
nobs = length(tobs)
xmod = zeros(nobs)
ymod = zeros(nobs)
k=1
for j=1:20
  tp = elements_all[2,3]+(j-0.5)/20*elements_all[2,2]
  M = (tobs[1]-tp)*2*pi/period
#  omega = atan2(-(yobs[1]*cos(M)-xobs[1]*sin(M))/semi,(yobs[1]*sin(M)+xobs[1]*cos(M))/semi)
  omega = atan(-(yobs[1]*cos(M)-xobs[1]*sin(M))/semi,(yobs[1]*sin(M)+xobs[1]*cos(M))/semi)
  for k=1:20
    omega = (k-0.5)/20*2*pi
    elements = [period,sqrt(elements_all[2,4]^2+elements_all[2,5]^2),tp,semi,omega]
  # P,ecc,tp,semi,omega = elements
    for i=1:length(tobs)
      xmod[i],ymod[i] = return_xy(tobs[i],elements)
    end
    chisq[j,k] =sum((xobs .- xmod).^2+(yobs .- ymod).^2)
    if chisq[j,k] < chi_best
      global chi_best = chisq[j,k]
      ibest = j
      global elements_best = elements
      println(j," chi_best: ",chi_best," elements_best: ",elements_best)
      clf()
      plot(xobs,xmod)
      plot(yobs,ymod)
    end
  end
end

xmod, ymod , res = fit_keplerian(tobs,xobs,yobs,elements_best)

# Now fit big-float data:
data_hr_big  = readdlm("test_output_highres_big.txt")
tobs_big = data_hr_big[:,1]; xobs_big  = data_hr_big[:,5]-data_hr_big[:,2]; yobs_big = -(data_hr_big[:,7]-data_hr_big[:,4])
elements_best_big = copy(res.minimizer)
xmod_big, ymod_big , res_big = fit_keplerian(tobs_big,xobs_big,yobs_big,elements_best_big)

clf()
#plot(tobs,xobs-xmod)
#plot(tobs,yobs-ymod)

# Plot the difference in angle between the n-body and Keplerian angles:
#dtheta = atan2(yobs.*xmod .- ymod.*xobs,xobs.*xmod+yobs.*ymod)
dtheta = atan.(yobs.*xmod .- ymod.*xobs,xobs.*xmod+yobs.*ymod)
xdot = data_hr[:,14]-data_hr[:,11]; ydot = -(data_hr[:,16]-data_hr[:,13])
# Compute the change in theta as a function of time:
dthetadt = (ydot.*xobs .-yobs.*xdot)./(xobs.^2 .+yobs.^2)
plot(tobs,-dtheta./dthetadt*24*3600)
#dtheta_big = atan2(yobs_big.*xmod_big .- ymod_big.*xobs_big,xobs_big.*xmod_big .+yobs_big.*ymod_big)
dtheta_big = atan.(yobs_big.*xmod_big .- ymod_big.*xobs_big,xobs_big.*xmod_big .+yobs_big.*ymod_big)
plot(tobs,-dtheta_big./dthetadt*24*3600)


using EffectiveWaves, Plots

pyplot()

Maxtime=100.
T=Float64

"Derivative of Hankel function of the first kind"
function diffhankelh1(n,z)
  if n!=0
    0.5*(hankelh1(-1 + n, z) - hankelh1(1 + n, z))
  else
    - hankelh1(1, z)
  end
end

"Derivative of Bessel function of first kind"
function diffbesselj(n,z)
  if n!=0
    0.5*(besselj(-1 + n, z) - besselj(1 + n, z))
  else
    - besselj(1, z)
  end
end

# background medium
medium = Medium(1.0,1.0+0.0im)
radius_multiplier = 1.005
max_hankel_order=15

incident_medium = medium
θin = 0.2
ω = 15.
k = ω/incident_medium.c

## Strong scatterers
species = [
    Specie(ρ=0.8,r=0.1, c=0.2, volfrac=0.1),
    Specie(ρ=0.2, r=0.2, c=0.1, volfrac=0.1)
]

S = length(species)
# @memoize Z_l_n(l,n) = Zn(ω,species[l],medium,n)

as = radius_multiplier*[(s1.r + s2.r) for s1 in species, s2 in species]
function M(keff,j,l,m,n)
    (n==m ? 1.0+im*0.0:0.0+im*0.0)*(j==l ? 1.0+im*0.0:0.0+im*0.0) + 2.0pi*species[l].num_density*Z_l_n[l,n]*
        kernelN(n-m,k*as[j,l],keff*as[j,l])/(k^2.0-keff^2.0)
end

ho = -1 + sum([ tol .< norm([M(0.9*k + 0.1im,j,l,1,n) for j = 1:S, l = 1:S]) for n=0:max_hankel_order ])

# this matrix is needed to calculate the eigenvectors
MM(keff::Complex{T}) = reshape(
    [M(keff,j,l,m,n) for m in -ho:ho, j = 1:S, n in -ho:ho, l = 1:S]
, ((2ho+1)*S, (2ho+1)*S))
detMM2(keff_vec::Array{T}) = map(x -> real(x*conj(x)), det(MM(keff_vec[1]+im*keff_vec[2])))

function detMM!(F,x)
    F[1] = abs(det(MM(x[1]+im*x[2])))
end

k0 = real(k)
x = k0 .* LinRange(0.0,1.8,100)/10.
y = k0 .* LinRange(0.,0.22,100)

X = repmat(x',length(y),1)
Y = repmat(y,1,length(x))

Z0 = map( (x,y) -> detMM2([x,y]),X,Y)
Z = map( z -> (abs(z)> 2.) ? NaN : z, Z0)
contour(x,y,Z,fill=true, xlab = "Re k*", ylab = "Im k*", title="Roots of secular det M - high ω - low φ")

x = k0 .* LinRange(0.18,0.23,50)
y = k0 .* LinRange(0.01,0.17,50)

x = kx; y = ky;

X = repmat(x',length(y),1)
Y = repmat(y,1,length(x))
Z = map( (x,y) -> (z = detMM2([x,y]); (abs(z)> 0.5) ? NaN: z),X,Y)

contour(x,y,Z,fill=true, xlab = "Re k*", ylab = "Im k*", title="Roots of secular det M")

x = LinRange(3.,6.,40)
y = LinRange(-1.5,14.,40)

X = repmat(x',length(y),1)
Y = repmat(y,1,length(x))
Z = map( (x,y) -> (z = detMM2([x,y]); (abs(z)> 1.) ? NaN: z),X,Y)
# Z = map( (x,y) -> detMM2([x,y]),X,Y)

# contour(x,y,Z,fill=true, xlab = "Re k*", ylab = "Im k*", title="Roots of secular det M")
heatmap(x,y,Z, xlab = "Re k*", ylab = "Im k*", title="det M - high ω and almost dirichlet")


using Revise
using Plots
using Basins
using DynamicalSystems
using RegionTrees

function newton_map(dz,z, p, n)
    f(x) = x^p[1]-1
    df(x)= p[1]*x^(p[1]-1)
    z1 = z[1] + im*z[2]
    dz1 = f(z1)/df(z1)
    z1 = z1 - dz1
    dz[1]=real(z1)
    dz[2]=imag(z1)
    return
end

# dummy function to keep the initializator happy
function newton_map_J(J,z0, p, n)
   return
end

ds = DiscreteDynamicalSystem(newton_map,[0.1, 0.2], [3] , newton_map_J)
integ  = integrator(ds)
res=10
xg=range(-1.,1.,length=res)
yg=range(-1.,1.,length=res)

@time bsn=Basins.basins_map2D_tree(xg, yg, integ; r_init=0.125, r_max=0.125/4)

function evaluate(basin,v)
    return findleaf(basin,v).data
end

function compute_frac_dim(bsn)

    root_n(x,n) = foldl((y,_) -> parent(y), 1:n; init=x)

    N = zeros(Int, 20)
    wd=1e3
    k_mx = 0
    for lf in allcells(bsn.basin)

        k = 1
        while !isnothing(root_n(lf,k))
            k += 1
        end
        N[k] += 1
        wd = min(wd,lf.boundary.widths[1])
        k_mx = max(k_mx,k)
    end
    @show N
    epsi = reverse(wd * 2 .^ range(0,k_mx-1,step=1))
    return epsi,N[1:k_mx]
end

epsi,N = compute_frac_dim(bsn)

elog = log10.(1 ./ epsi)
Nlog = log10.(N)
D = linear_region(elog[4:end],Nlog[4:end])
@show D

plt = plot(xlim=(xg[1], xg[end]), ylim=(yg[1], yg[end]), legend=nothing)

x = range(xg[1], xg[end], length=1000)
y = range(yg[1], yg[end], length=1000)
heatmap!(plt, x, y, (x, y) -> evaluate(bsn.basin, SVector(x, y)), fill=true)

# plt = plot(xlim=(-1.5, 1.5), ylim=(-1.5, 1.5), legend=nothing)
 for leaf in allleaves(bsn.basin)
     v = hcat(collect(vertices(leaf.boundary))...)
     #plot!(plt, v[1,[1,2,4,3,1]], v[2,[1,2,4,3,1]])
     #@show leaf.data
     plot!(plt, v[1,[1,2,4,3,1]], v[2,[1,2,4,3,1]], linecolor=:black, lw=0.4)

 end
 plt


###
  Name     : c5_31_CAPM.jl
  Book     : Hands-on Data Science with Anaconda )
  Publisher: Packt Publishing Ltd.
  Author   : Yuxing Yan and James Yan
  Date     : 5/30/2018
  email    : yany@canisius.edu
             paulyxy@hotmail.com

Note #1: go to http://finance.yahoo.com to dowload the last 5-year monthly
          data for IBM and S&P500. The ticker for S&P500 is ^GSPC
      #2: delete the first observation which contains 'null'
      #3: check if the last line is a blank line. If it is, delete it.
###

using DataFrames, CSV
using GLM, StatsModels
#
function f(x)
  p=x[Symbol("Adj Close")]
  n=length(p)
  a=p[2:n]./p[1:(n-1)]-1.0
  return convert(Array{Float64,1}, a)
end
#
x = CSV.read("c:/temp/ibmMonthly5years.txt")
stockRet=f(x)
y = CSV.read("c:/temp/sp500Monthly5years.txt")
mktRet=f(y)
n3=min(length(stockRet),length(mktRet))
data = DataFrame(X=mktRet[1:n3], Y=stockRet[1:n3])
OLS = glm(@formula(Y ~ X), data, Normal(), IdentityLink())


using FluxReconstruction, Test

ps = FRPSpace1D(0, 1, 100, 2)

#--- values ---#
V = vandermonde_matrix(ps.deg, ps.xpl)
ψf = vandermonde_matrix(ps.deg, [-1.0, 1.0])

lf = zeros(2, ps.deg + 1)
for i in axes(lf, 1)
    lf[i, :] .= V' \ ψf[i, :]
end

@test lf[1, :] ≈ ps.ll
@test lf[2, :] ≈ ps.lr

#--- derivatives ---#
Vr = ∂vandermonde_matrix(ps.deg, ps.xpl)

∂l = zeros(ps.deg + 1, ps.deg + 1)
for i = 1:ps.deg+1
    ∂l[i, :] .= V' \ Vr[i, :]
end

@test ∂l ≈ ps.dl

dVf = ∂vandermonde_matrix(ps.deg, [-1.0, 1.0])
∂lf = zeros(2, ps.deg + 1)
for i = 1:2
    ∂lf[i, :] .= V' \ dVf[i, :]
end


using Base.Test
using ForwardNets

net = ForwardNet{Float32}()
push!(net, Variable(:A, Float32[0]))
push!(net, Variable(:B, Float32, 2))
push!(net, Variable, :C, Float32[1])
push!(net, Variable, :D, 3, :LAST)
@test lastindex(net) == 4

A = net[:A]
B = net[2]

@test name(A) == :A
@test name(B) == :B
@test output(A) == [0.0]
@test size(output(B)) == (2,)

push!(net, Affine, :affine, :A, 1)
push!(net, Affine, :A, 1)
affine = net[:affine]
@test name(affine) == :affine
@test size(output(affine)) == (1,)

A.tensor[1] = 1.0f0
@test affine.input[1] == 1.0f0
affine.W[1] = 2.0f0
affine.b[1] = 3.0f0
affine.output[1] = NaN
forward!(affine)
@test affine.output[1] == 5.0f0

push!(net, Concatenator, :concat, [:A, :C])
concat = net[:concat]
@test size(output(concat)) == (2,)
output(A)[1] = 1.0f0
output(net[:C])[1] = 2.0f0
forward!(concat)
@test concat.output == [1.0f0, 2.0f0]

push!(net, Reshaper, :reshaper, :concat, (1,2))
reshaper = net[:reshaper]
@test size(output(reshaper)) == (1,2)
forward!(reshaper)
@test output(reshaper) == [1.0f0, 2.0f0]'

push!(net, TanH, :tanh, :A)
push!(net, ReLU, :relu, :A)
push!(net, SoftPlus, :softplus, :A)
_tanh = net[:tanh]
_relu = net[:relu]
_softplus = net[:softplus]

output(A)[1] = -2.0f0
forward!(_tanh)
@test isapprox(output(_tanh)[1], tanh(-2.0))

output(A)[1] = -2.0f0
forward!(_relu)
@test isapprox(output(_relu)[1], relu(-2.0))

output(A)[1] = -2.0f0
forward!(_softplus)
@test isapprox(output(_softplus)[1], softplus(-2.0))

push!(net, LSTM, :lstm, :D, 3)
lstm = net[:lstm]
@test name(lstm) == :lstm
lstm.state[1] = NaN
zero!(lstm)
@test lstm.state[1] == 0.0f0
forward!(lstm)

push!(net, Variable, :output, :affine, :LAST)

pass = calc_forwardpass(net, [:A], [:output])
output(A)[1] = 2.0f0
forward!(pass)
@test output(net[:output])[1] == convert(Float32, 2*2 + 3)


@testset "evaluate" begin

    @test evaluate(GPNode(Const(1)), []) == 1
    @test evaluate(GPNode(Const(1.25)), []) == 1.25

    @test evaluate(GPNode(Variable(1)), [3.14]) == 3.14
    @test evaluate(GPNode(Variable(1)), [3.14]) == 3.14

    add = PrimitiveFunction(+, 2)
    mul = PrimitiveFunction(*, 2)
    sub = PrimitiveFunction(-, 2)

    t = GPNode(mul, GPNode[
        GPNode(sub, GPNode[
            GPNode(Const(10)),
            GPNode(Variable(1))
        ]),
        GPNode(add, GPNode[
            GPNode(Variable(1))
            GPNode(Variable(2))
        ])
    ])
    t_func(vars) = (10 - vars[1]) * (vars[1] + vars[2])

    @test evaluate(t, [0, 1]) == t_func([0, 1])
    @test evaluate(t, [123, 456]) == t_func([123, 456])

    @test evaluate(t, [0., 1.]) == t_func([0., 1.])

end


module CovarianceEstimation

using Statistics
using StatsBase
using LinearAlgebra
import StatsBase: cov

export cov
export CovarianceEstimator, SimpleCovariance,
    LinearShrinkage,
    # Targets for linear shrinkage
    DiagonalUnitVariance, DiagonalCommonVariance, DiagonalUnequalVariance,
    CommonCovariance, PerfectPositiveCorrelation, ConstantCorrelation,
    # Eigendecomposition-based methods
    AnalyticalNonlinearShrinkage,
    # Biweight midcovariance
    BiweightMidcovariance

include("utils.jl")
include("biweight.jl")
include("linearshrinkage.jl")
include("nonlinearshrinkage.jl")

end # module


# Unit Test for by default supported convex sets and their functions
using COSMO, Test, Random, LinearAlgebra

rng = Random.MersenneTwister(13131)

@testset "Convex Sets" begin
tol = 1e-4

    @testset "Create and project" begin

    # Zero Cone
    zset = COSMO.ZeroSet(10)
    x = randn(rng,10)
    COSMO.project!(view(x, 1:length(x)), zset)
    @test norm(x,Inf) == 0.

    # Positive Orthant R+
    nonnegatives = COSMO.Nonnegatives(10)
    x = randn(rng, 10)
    COSMO.project!(view(x, 1:length(x)), nonnegatives)
    @test minimum(x) >= 0.

    # Box
    l = -1 * ones(10)
    u = 1 * ones(10)
    box = COSMO.Box(l, u)
    box = COSMO.Box{Float64}(10)
    box.l .= l
    box.u .= u
    x = 100 * randn(rng, 10)
    COSMO.project!(view(x, 1:length(x)), box)
    @test minimum(x) >= -1. && maximum(x) <= 1.

    # Second Order (Lorentz) cones
    soc = COSMO.SecondOrderCone(10)
    x = 10 * randn(rng, 9)
    t = norm(x, 2) - 0.5
    x = [t; x]

    COSMO.project!(view(x, 1:length(x)), soc)
    @test norm(x[2:10], 2) <= x[1]

    # Positive Semidefinite cones
    psd = COSMO.PsdCone(16)
    X = Symmetric(randn(rng, 4, 4))
    X = X - 4 * Matrix(1.0I, 4, 4)
    x = vec(X)
    COSMO.project!(view(x, 1:length(x)), psd)
    @test minimum(eigen(reshape(x, 4, 4)).values) >= -1e-9

    C = COSMO.CompositeConvexSet([COSMO.ZeroSet(10), COSMO.Nonnegatives(10)])
    x = -rand(20)
    xs = COSMO.SplitVector(x, C)
    COSMO.project!(xs, C)
    @test norm(x, Inf) == 0.

    end

    @testset "in_dual Functions" begin

    # Dual of zero cone
    x = randn(rng, 10)
    convex_set = COSMO.ZeroSet(10)
    @test COSMO.in_dual(view(x, 1:length(x)), convex_set, tol)

    # Dual of Positive Orthant R+ (self-dual)
    xpos = rand(rng, 10)
    xneg = -rand(rng, 10)
    xzeros = zeros(10)
    convex_set = COSMO.Nonnegatives(10)
    @test COSMO.in_dual(view(xpos, 1:length(xpos)), convex_set, tol)
    @test !COSMO.in_dual(view(xneg,1:length(xneg)), convex_set, tol)
    @test COSMO.in_dual(view(xzeros,1:length(xzeros)), convex_set,tol)

    #TODO: Dual of Box [important!]

    # Dual of Second Order Cone (self-dual)
    tol = 1e-4
    x = randn(rng, 9)
    t = norm(x, 2)
    xpos = [t + 0.5; x]
    xneg = [t - 0.5; x]
    convex_set = COSMO.SecondOrderCone(10)
    @test COSMO.in_dual(view(xpos, 1:length(xpos)), convex_set, tol)
    @test !COSMO.in_dual(view(xneg, 1:length(xneg)), convex_set, tol)

    # Dual of Positive Semidefinite Cone (Square) (self-dual)
    tol = 1e-4
    X = Symmetric(randn(rng, 4, 4))
    Xpos = X + 4 * Matrix(1.0I, 4, 4)
    Xneg = X - 4 * Matrix(1.0I, 4, 4)
    xpos = vec(Xpos)
    xneg = vec(Xneg)
    convex_set = COSMO.PsdCone(16)
    @test COSMO.in_dual(view(xpos, 1:length(xpos)), convex_set, tol)
    @test !COSMO.in_dual(view(xneg, 1:length(xneg)), convex_set, tol)

    # Dual of Positive Semidefinite Cone (Triangle)
    convex_set = COSMO.PsdConeTriangle(10)
    xpos = zeros(10)
    xneg = zeros(10)
    COSMO.extract_upper_triangle!(Xpos, xpos, sqrt(2))
    COSMO.extract_upper_triangle!(Xneg, xneg, sqrt(2))
    @test COSMO.in_dual(view(xpos, 1:length(xpos)), convex_set, tol)
    @test !COSMO.in_dual(view(xneg, 1:length(xneg)), convex_set, tol)

    end

    @testset "inPolRec Functions" begin

    # Polar Recession cone of zero cone
    xpos = zeros(10)
    xneg = randn(rng, 10)
    convex_set = COSMO.ZeroSet(10)
    @test COSMO.in_pol_recc(view(xpos, 1:length(xpos)), convex_set, tol)
    @test !COSMO.in_pol_recc(view(xneg, 1:length(xneg)), convex_set, tol)

    # Polar Recession cone of Positive Orthant R+
    xpos = -rand(rng, 10)
    xneg = rand(rng, 10)
    convex_set = COSMO.Nonnegatives(10)
    @test COSMO.in_pol_recc(view(xpos,1:length(xpos)), convex_set, tol)
    @test !COSMO.in_pol_recc(view(xneg,1:length(xneg)), convex_set, tol)

    #TODO: Polar Recc of Box [important!]

    # Polar Recc of Second Order Cone
    tol = 1e-4
    x = randn(rng, 9)
    t = norm(x, 2)
    xpos = [-t - 0.5; x]
    xneg = [-t + 0.5; x]
    convex_set = COSMO.SecondOrderCone(10)
    @test COSMO.in_pol_recc(view(xpos, 1:length(xpos)), convex_set, tol)
    @test !COSMO.in_pol_recc(view(xneg, 1:length(xneg)), convex_set, tol)

    # Polar Recc of Positive Semidefinite Cone (Square)
    tol = 1e-4
    X = Symmetric(randn(rng, 4, 4))
    Xpos = X - 20 * Matrix(1.0I, 4, 4)
    Xneg = X + 4 * Matrix(1.0I, 4, 4)
    xpos = vec(Xpos)
    xneg = vec(Xneg)
    convex_set = COSMO.PsdCone(16)
    @test COSMO.in_pol_recc(view(xpos, 1:length(xpos)), convex_set, tol)
    @test !COSMO.in_pol_recc(view(xneg, 1:length(xneg)), convex_set, tol)

    # Polar Recc of Positive Semidefinite Cone (Triangle)
    COSMO.extract_upper_triangle!(Xpos, xpos, sqrt(2))
    COSMO.extract_upper_triangle!(Xneg, xneg, sqrt(2))
    convex_set = COSMO.PsdConeTriangle(10)
    @test COSMO.in_pol_recc(view(xpos, 1:length(xpos)), convex_set, tol)
    @test !COSMO.in_pol_recc(view(xneg, 1:length(xneg)), convex_set, tol)
   end

   @testset "Set Utilities" begin
    convex_set = COSMO.PsdConeTriangle(3)
    composite_set = COSMO.CompositeConvexSet([COSMO.ZeroSet(2), COSMO.Nonnegatives(1)])
    @test COSMO.get_subset(convex_set, 1) == convex_set
    @test_throws DimensionMismatch COSMO.get_subset(convex_set, 2)
    @test COSMO.get_subset(composite_set, 1) == COSMO.ZeroSet(2)
    @test_throws ErrorException COSMO.CompositeConvexSet([convex_set; composite_set])
   end

end

nothing


using Plots.PlotMeasures
using Plots

export
    plot_ϵ,
    plot_e_field

function plot_ϵ(s::Simulator; figsize=(600, 750), left_margin=-100px)
    ϵ = s.permittivity.ϵ

    return heatmap(
		axes(s.grid, 1), axes(s.grid, 2), ϵ',
		color=:algae,
		size=figsize, left_margin=left_margin, aspect_ratio=:equal
	)
end

function plot_e_field(s::Simulator; figsize=(600, 750), left_margin=-100px)
    ez = s.ez
    ϵ = s.permittivity.ϵ

    lim = maximum(abs.(ez))
    lim_ϵ = maximum(abs.(ϵ))

    p = plot(
        clim=(-lim, lim),  colorbar=false,
		size=figsize, left_margin=left_margin, aspect_ratio=:equal
    )
    p = heatmap!(
        p,
		axes(s.grid, 1), axes(s.grid, 2), ez',
		color=:coolwarm
	)
    p = contour!(
        p,
        axes(s.grid, 1), axes(s.grid, 2), lim .* ϵ' ./ lim_ϵ,
        color=:algae
    )

    return p
end


# ============================================================
# Thermodynamics
# ============================================================

"""
    heat_capacity_ratio(K, D)
    heat_capacity_ratio(K, Nr, D)

Calculate heat capacity ratio

"""
function heat_capacity_ratio(K, D::T) where {T<:Integer}
    γ = begin
        if D == 1
            (K + 3.0) / (K + 1.0)
        elseif D == 2
            (K + 4.0) / (K + 2.0)
        elseif D == 3
            (K + 5.0) / (K + 3.0)
        end
    end

    return γ
end

#--- diatomic ---#
function heat_capacity_ratio(K, Nr, D::T) where {T<:Integer}
    γ = begin
        if D == 1
            (K + 3.0 + Nr) / (K + 1.0 + Nr)
        elseif D == 2
            (K + 4.0 + Nr) / (K + 2.0 + Nr)
        elseif D == 3
            (K + 5.0 + Nr) / (K + 3.0 + Nr)
        end
    end

    return γ
end

"""
    sound_speed(λ, γ)
    sound_speed(prim, γ)

Calculate speed of sound

"""
sound_speed(λ::Real, γ::Real) = (0.5 * γ / λ)^0.5

sound_speed(prim::AbstractVector{T}, γ) where {T<:Real} = sound_speed(prim[end], γ)

#--- mixture ---#
function sound_speed(prim::AbstractMatrix{T}, γ) where {T<:Real}
    c = similar(prim, axes(prim, 2))
    for j in eachindex(c)
        c[j] = sound_speed(prim[end, j], γ)
    end

    return maximum(c)
end


# This script looks at the coefficient plot shown in the first example of [Robust causal interference
# using direct acyclic graphs: the R package `dagitty`](https://academic.oup.com/ije/article/45/6/1887/2907796)
# by J Textor, B van der Zanden, M Gilthorpe, M Lisliewicz and G Ellison.
# [R Dagitty](https://doi.org/10.1093/ije/dyw341)

using StatisticalRethinking, StatsPlots

ProjDir = @__DIR__

for i in 1:5
  include(scm_path("..", "examples", "sports", "m$i.jl"))
end

models=[m1s, m2s, m3s, m4s, m5s];
pars=[:bC, :bFL, :bPGP, :bTM]

df = DataFrame!(CSV.File(scm_path("..", "examples", "sports", "sports.csv"), delim=','));

scatter(df[:, :TM], df[:, :WUE])
savefig("$(ProjDir)/sports_01.png")

r1 = plotcoef(models, pars, "$(ProjDir)/sports_02.png", "Normal (Particles) estimates")
display(r1)

println()
r2 = plotcoef(models, pars, "$(ProjDir)/sports_03.png", "Quap estimates", quap)
display(r2)


using BlackBoxOptim, Gadfly

# run Borg MOAE
schaffer1(x) = (sumabs2(x), sumabs2(x .- 2.0))
res = bboptimize(schaffer1; Method=:borg_moea,
                FitnessScheme=ParetoFitnessScheme{2}(is_minimizing=true),
                SearchRange=(-10.0, 10.0), NumDimensions=3, ϵ=0.1,
                MaxSteps=15000, TraceInterval=1.0, TraceMode=:verbose);

# the parameterized version of the exact Pareto frontier
# N is the number of the problem (not fitness) dimensions
pareto_curve_func(t, ::Type{Val{N}}) where {N} = (N*t[1]^2, N*(2-t[1])^2)
pareto_curve = BlackBoxOptim.Hypersurface(pareto_curve_func,
                                          symmetric_search_space(1, (0.0, 2.0)))

# generate the set of ϵ-indexed points on the exact Pareto frontier
pareto_pts = BlackBoxOptim.generate(pareto_curve,
                                   fitness_scheme(res), Val{numdims(res)});
# calculate the distance between the solution and the exact frontier
BlackBoxOptim.IGD(pareto_curve, pareto_frontier(res), fitness_scheme(res), Val{numdims(res)})

# draw the results
plot(layer(x = [x.orig[1] for x in values(pareto_pts)], # Draw the exact Pareto frontier
           y = [x.orig[2] for x in values(pareto_pts)],
           Geom.line),
    layer(x= [fitness(x).orig[1] for x in pareto_frontier(res)], # the result of the method
          y= [fitness(x).orig[2] for x in pareto_frontier(res)],
          color= [haskey(pareto_pts, fitness(x).index) for x in pareto_frontier(res)],
          Geom.point))


using Random

@testset "Profiles" begin

σ = rand()
g = Gaussian(σ,1)
@test g(σ) ≈ exp(-1)/sqrt(π*σ^2)

@test radius(g) == σ
@test strength(g) == 1

σ = 0.2
A = 1
dg = DGaussian(σ,A)
x = 0.1
@test dg(x) == -2*A*x/sqrt(π)/σ^3*exp(-x^2/σ^2)

end

@testset "SpatialFields" begin

σ = 0.2
A = 1

g = EmptySpatialField()
@test g(rand(2)...) == 0.0

g = SpatialGaussian(σ,σ,0.0,0.5,1)
@test g(0,0.5+σ) ≈ g(σ,0.5) ≈ g(-σ,0.5) ≈ g(0,0.5-σ) ≈ exp(-1)/(π*σ^2)

u = 1
v = 0
gc = SpatialGaussian(σ,σ,0.0,0.5,1,u,v)
t = 1
@test gc(u*t,0.5+σ+v*t,t) ≈ gc(σ+u*t,0.5+v*t,t) ≈ gc(-σ+u*t,0.5+v*t,t) ≈ gc(0+u*t,0.5-σ+v*t,t) ≈ exp(-1)/(π*σ^2)

g = EmptySpatialField()
for x in [-0.5,0,0.5], y in [-0.5,0,0.5]
  g += SpatialGaussian(0.2,0.5,x,y,1)
end
@test g(0.5,0) ≈ 5.5357580598289235

g2 = -g
@test g2(0.5,0) ≈ -5.5357580598289235

gr = PhysicalGrid((-3.0,3.0),(-2.0,2.0),0.03)
w = Nodes(Dual,size(gr))
xg, yg = coordinates(w,gr)

gfield = GeneratedField(w,g,gr)

@test gfield()[104,70] == g(xg[104],yg[70])
@test datatype(gfield) == typeof(w)

wfield = PulseField(gfield,0.5,0.1)

@test maximum(wfield(0.5)) ≈ maximum(gfield()) ≈ 5.532960088678624
@test maximum(wfield(2)) ≈ 0.0
@test datatype(wfield) == datatype(gfield) <: ScalarGridData

q = Edges(Primal,size(gr))
gauss = SpatialGaussian(0.5,0,0,1)
gaussfield = GeneratedField(q,gauss,EmptySpatialField(),gr)
q .= gaussfield()

ffield = PulseField(gaussfield,0.5,0.1)

@test maximum(ffield(0.5).u) ≈ maximum(q.u) ≈ 1.272094144652253
@test maximum(ffield(2).u) ≈ 0.0
@test maximum(ffield(0.5).v) ≈ maximum(q.v) ≈ 0.0
@test maximum(ffield(2).v) ≈ maximum(q.v) ≈ 0.0

@test datatype(ffield) == datatype(gaussfield) <: VectorGridData

gq = EdgeGradient(Dual,size(gr))
tfield = GeneratedField(gq,[g,gauss,EmptySpatialField(),EmptySpatialField()],gr)
gq .= tfield()
@test maximum(gq.dudx) ≈ 5.532960088678624
@test maximum(gq.dudy) ≈ 1.2732395447351625
@test maximum(gq.dvdx) == maximum(gq.dvdy) == 0.0

ffield = PulseField(tfield,0.5,0.1)
@test datatype(ffield) <: TensorGridData

end

@testset "Derivatives of gaussians" begin

  σx = 0.5
  σy = 0.5
  x0 = 0
  y0 = 0
  A = 1
  dgaussx = SpatialGaussian(σx,σy,x0,y0,A,deriv=1)
  dgaussy = SpatialGaussian(σx,σy,x0,y0,A,deriv=2)
  x = 0.1
  y = 0.2
  @test dgaussx(0.1,0.2) == dgaussy(0.2,0.1) ≈ -2*A*x/π/σx^3/σy*exp(-x^2/σx^2)*exp(-y^2/σy^2)

end


###########################################################################
#       Basic functionality for contracting and manipulating tensors      #
###########################################################################

using LinearAlgebra

"""
    contract_tensors(T1,cT1::Vector{Int64},T2,cT2::Vector{Int64},ordering::Vector{Int}=Vector{Int64}())

Function to contract to tensors T1 and T2, the contraction is performed as matrix multiplication by reordering the indices and reshaping the tensors
num_of_indices_T1: number of indices of tensor T1
num_of_indices_T2: number of indices of tensor T2
iT1: vector of indices of T1 which are contracted
iT2: vector of indices of T2 which are contracted
"""
function contract_tensors(T1, cT1::Vector{Int}, T2, cT2::Vector{Int}, ordering::Vector{Int}=Vector{Int64}())
    # Make sure the number of indices to contract is the same
    @assert(length(cT1) == length(cT2))

    # Determine size of the tensors
    dimT1 = size(T1)
    dimT2 = size(T2)

    # Make sure the dimensions along the contracted indices are equal
    if (dimT1[cT1] != dimT2[cT2])
        error("Dimensions along the contracted indices must be the same, failed to contract the indices")
    end

    # Get the uncontracted indices of T1 and T2
    uT1 = setdiff(collect(1:ndims(T1)), cT1)
    uT2 = setdiff(collect(1:ndims(T2)), cT2)

    # Bring the contracted indices of T1 to the end
    T1 = permutedims(T1, (uT1..., cT1...))
    # Bring the contracted indices of T2 to the front
    T2 = permutedims(T2, (cT2..., uT2...))

    # Build single index out of all contracted and uncontracted indices to
    # obtain a matrix with two indices
    il = prod(dimT1[uT1])
    ir = prod(dimT2[uT2])
    T1 = reshape(T1, (il, prod(dimT1[cT1])))
    T2 = reshape(T2, (prod(dimT2[cT2]), ir))
    # Actual contraction as matrix multiplication
    res = T1 * T2

    # Expand the single indices
    res = reshape(res, (dimT1[uT1]..., dimT2[uT2]...))

    # Permute order if necessary
    if (!isempty(ordering))
        res = permutedims(res, (ordering...,))
    end

    return res
end


to_int(s) = parse(Int, s)

diff(x, y) = abs(x-y)

parse_file(filename) = to_int.(split(read(filename, String), ","))

function solve(filename, cost_fn)
    xs = parse_file(filename)
    costs = [sum([cost_fn(i, x) for x in xs]) for i in min(xs...):max(xs...)]
    min(costs...)
end

# Part 1
solve("7eg.txt", diff) == 37
solve("7.txt", diff) == 336701

cost(x,y) = sum(1:diff(x,y))

# Part 2
solve("7eg.txt", cost)
solve("7.txt", cost)


using CairoMakie
using DifferentialEquations
using Random: MersenneTwister

mrng = MersenneTwister(1234)

##################################
####### Plot Setup ###############
##################################

noto_sans = "./resources/NotoSans-Regular.ttf"

tickfontsize    = 28
labelfontsize   = 32

#f = Figure(font=noto_sans, resolution=(800,400))
f = Figure(resolution = (2000, 400))

################################
########## Morse Function ######
################################

# (x^2 - 1)(x^2 - 4)
# f(x) = (x-1) * (x+1) * (x-2) * (x+2)
g(x) = 4 * x^3 - 10*x

function morseFlow!(du, u, μ, t)
    g(x) = 4.0 * x^3 - 10.0 * x

    du .= -0.5 .* g.(u)
end

curve_count = 50
u0 = 3 * randn(mrng, curve_count)

solution_curves = zeros(101, curve_count)
timings = 0:0.01:1

for i in 1:curve_count
    tspan = (0.0, 1.0)
    prob = ODEProblem(morseFlow!, [u0[i]], tspan)

    sol = solve(prob, saveat=0.01)
    solution_curves[:,i] = sol[1,:]
end

ax1 = Axis(  f[1,1],
            xticklabelsize=tickfontsize,
            yticklabelsize=tickfontsize,
            yticklabelpad=2,
            xlabel="τ", ylabel = "x",
            xlabelsize=labelfontsize,
            ylabelsize=labelfontsize,
            ylims = (-5.0, 5.0))

for i in 1:curve_count
    lines!(ax1, timings, solution_curves[:,i], color=:black)
end
xlims!(ax1,(0,1))

################################
##### Van Der Pol Oscillator ###
################################

function vanderpol!(du,u,μ,t)
    x,y = u
    du[1] = μ * (x - x^3/3 - y)
    du[2] = x/μ
end

curve_count = 20
u0 = 5 * rand(mrng, 2,curve_count) .- 2.5

solution_curves = zeros(201, 2, curve_count)

for i in 1:curve_count
    tspan = (0.0, 10.0)
    prob = ODEProblem(vanderpol!, u0[:,i], tspan, 1)

    sol = solve(prob, saveat=0.05)
    solution_curves[:, 1, i] = sol[1,:]
    solution_curves[:, 2, i] = sol[2,:]
end

ax2 = Axis(  f[1,2],
            xticklabelsize=tickfontsize,
            yticklabelsize=tickfontsize,
            yticklabelpad=2,
            xlabel="x", ylabel = "y",
            xlabelsize=labelfontsize,
            ylabelsize=labelfontsize,
            xlims = (-2.5,2.5),
            ylims = (-2.5,2.5))

for i in 1:curve_count
    lines!(ax2, solution_curves[:,1,i], solution_curves[:,2,i], color=:black)
end

################################
###### Hopf Bifurcaton #########
################################

function hopf!(du,u,μ,t)
    x,y = u
    du[1] = x * (1.0 - (x^2 + y^2)) - y
    du[2] = y * (1.0 - (x^2 + y^2)) + x
end

curve_count = 20
u0 = 4 * rand(mrng,2,curve_count) .- 2.0

solution_curves = zeros(201, 2, curve_count)

for i in 1:curve_count
    tspan = (0.0, 10.0)
    prob = ODEProblem(hopf!, u0[:,i], tspan, 1)

    sol = solve(prob, saveat=0.05)
    solution_curves[:, 1, i] = sol[1,:]
    solution_curves[:, 2, i] = sol[2,:]
end

ax3 = Axis(  f[1,3],
            xticklabelsize=tickfontsize,
            yticklabelsize=tickfontsize,
            yticklabelpad=2,
            xlabel="x", ylabel = "y",
            xlabelsize=labelfontsize,
            ylabelsize=labelfontsize,
            xlims = (-2.0,2.0),
            ylims = (-2.0,2.0))

for i in 1:curve_count
    lines!(ax3, solution_curves[:,1,i], solution_curves[:,2,i], color=:black)
end

###############################
####### Lorenz System #########
###############################

function lorenz!(du,u,μ,t)
    x, y, z = u
    du[1] = 10.0 * (y - x)
    du[2] = x*(28.0 - z) - y
    du[3] = x * y  - 8.0/3.0 * z
end

curve_count = 3
u0 = 100 * rand(mrng, 3,curve_count) .- 50.0

solution_curves = zeros(2001, 3, curve_count)

for i in 1:curve_count
    tspan = (0.0, 10.0)
    prob = ODEProblem(lorenz!, u0[:,i], tspan, 1)

    sol = solve(prob, saveat=0.005)
    solution_curves[:, 1, i] = sol[1,:]
    solution_curves[:, 2, i] = sol[2,:]
    solution_curves[:, 3, i] = sol[3,:]
end

ax4 = Axis( f[1,4],
            xticklabelsize=tickfontsize,
            yticklabelsize=tickfontsize,
            yticklabelpad=2,
            xlabel="x", ylabel = "z",
            xlabelsize=labelfontsize,
            ylabelsize=labelfontsize)

for i in 1:curve_count
    lines!(ax4, solution_curves[:,1,i], solution_curves[:,3,i], color=:black)
end

##############################
######## Heat Equation #######
############################3#

heat = zeros(100,200)
heat[:,1] = randn(mrng, 100) .+ 1
heat[50:60,1] .+= 4
heat[90:100,1] .+= 4
α = 0.2
for i in 2:200
    heat[:,i] += (1.0 - 2.0 * α) * heat[:,i-1]
    heat[2:end,i] += α .* heat[1:end-1,i-1]
    heat[1:end-1,i] += α .* heat[2:end,i-1]
    heat[1,i] += α * heat[end,i-1]
    heat[end,i] += α * heat[1,i-1]
end

ax5 = Axis( f[1,5],
            xticklabelsize=tickfontsize,
            yticklabelsize=tickfontsize,
            yticklabelpad=2,
            xlabel="τ", ylabel = "x",
            xlabelsize=labelfontsize,
            ylabelsize=labelfontsize)

hm = heatmap!(ax5, (1:200)./100.0, (1:100)./100.0, heat')
Colorbar(f[1,6],hm, ticklabelsize=tickfontsize)

save("tablePlot.pdf", f)


using Test
using UnitTestDesign
#### Set-based

### tuples_in_trials
tt_cases = [
    [[[1,2,3], [1,2,1], [2,3,4]], 2, 8],
    [[[1,1,1], [1,1,1], [1,1,1]], 2, 3],
    [[[1,1,1], [1,1,1], [1,1,1]], 3, 1],
]
for tt_case in tt_cases
    res = UnitTestDesign.tuples_in_trials(tt_case[1], tt_case[2])
    total_covered = sum([length(aset) for aset in values(res)])
    @test total_covered == tt_case[3]
end

sc = UnitTestDesign.SetCoverage([2,3,2], 2)
@test UnitTestDesign.remaining(sc) == 0
@test eltype(sc) == Int64
UnitTestDesign.build_all_combinations!(sc, 2)
@test UnitTestDesign.remaining(sc) == 16
@test UnitTestDesign.add_coverage!(sc, [1, 1, 1]) == 3
@test UnitTestDesign.remaining(sc) == 13
@test UnitTestDesign.add_coverage!(sc, [1, 2, 1]) == 2
@test UnitTestDesign.remaining(sc) == 11

UnitTestDesign.build_all_combinations!(sc, 2)
multiple = [2 2; 2 3; 1 2]
@test UnitTestDesign.add_coverage!(sc, multiple) == 6
@test UnitTestDesign.remaining(sc) == 10

# Add two types of coverage
sc = UnitTestDesign.SetCoverage([2,3,2], 2)
UnitTestDesign.build_all_combinations!(sc, 2)
UnitTestDesign.build_all_combinations!(sc, 3)
@test UnitTestDesign.remaining(sc) == 28


include("../utils.jl")

@testset "Probability simplex" begin
    M = ProbabilitySimplex(2)
    @test M^4 == MultinomialMatrices(3, 4)
    p = [0.1, 0.7, 0.2]
    q = [0.3, 0.6, 0.1]
    X = zeros(3)
    Y = [-0.1, 0.05, 0.05]
    @test is_point(M, p)
    @test_throws DomainError is_point(M, p .+ 1, true)
    @test_throws DomainError is_point(M, [0], true)
    @test_throws DomainError is_point(M, -ones(3), true)
    @test manifold_dimension(M) == 2
    @test is_vector(M, p, X)
    @test is_vector(M, p, Y)
    @test_throws DomainError is_vector(M, p .+ 1, X, true)
    @test_throws DomainError is_vector(M, p, zeros(4), true)
    @test_throws DomainError is_vector(M, p, Y .+ 1, true)

    @test Manifolds.default_metric_dispatch(M, Manifolds.FisherRaoMetric()) === Val{true}()

    @test injectivity_radius(M, p) == injectivity_radius(M, p, ExponentialRetraction())
    @test injectivity_radius(M, p, SoftmaxRetraction()) == injectivity_radius(M, p)
    @test injectivity_radius(M, ExponentialRetraction()) == 0
    @test injectivity_radius(M) == 0
    @test injectivity_radius(M, SoftmaxRetraction()) == 0

    pE = similar(p)
    embed!(M, pE, p)
    @test pE == p
    XE = similar(X)
    embed!(M, XE, p, X)
    @test XE == X

    @test mean(M, [p, q]) == shortest_geodesic(M, p, q)(0.5)

    types = [Vector{Float64}]
    TEST_FLOAT32 && push!(types, Vector{Float32})
    TEST_STATIC_SIZED && push!(types, MVector{3,Float64})

    basis_types = (DefaultOrthonormalBasis(), ProjectedOrthonormalBasis(:svd))
    for T in types
        @testset "Type $T" begin
            pts = [
                convert(T, [0.5, 0.3, 0.2]),
                convert(T, [0.4, 0.4, 0.2]),
                convert(T, [0.3, 0.5, 0.2]),
            ]
            test_manifold(
                M,
                pts,
                test_injectivity_radius=false,
                test_project_tangent=true,
                test_musical_isomorphisms=true,
                test_vee_hat=false,
                test_forward_diff=false,
                test_reverse_diff=false,
                is_tangent_atol_multiplier=5.0,
                inverse_retraction_methods=[SoftmaxInverseRetraction()],
                retraction_methods=[SoftmaxRetraction()],
                test_inplace=true,
            )
        end
    end

    @testset "Projection testing" begin
        p = [1 / 3, 1 / 3, 1 / 3]
        q = [0.2, 0.3, 0.5]
        X = log(M, p, q)
        X2 = X .+ 10
        Y = project(M, p, X2)
        @test isapprox(M, p, X, Y)

        # Check adaption of metric and representer
        Y1 = change_metric(M, EuclideanMetric(), p, X)
        @test Y1 == X .* sqrt.(p)
        Y2 = change_representer(M, EuclideanMetric(), p, X)
        @test Y2 == X .* p

        X = log(M, q, p)
        X2 = X + [1, 2, 3]
        Y = project(M, q, X2)
        @test is_vector(M, q, Y; atol=1e-15)

        @test_throws DomainError project(M, [1, -1, 2])
        @test isapprox(M, [0.6, 0.2, 0.2], project(M, [0.3, 0.1, 0.1]))
    end
end


#=
  https://www.youtube.com/watch?v=e1Ykk_CqKTY&t=458s

  Probabilistic Programming: What It Is and How It Works - Noel Welsh

  We can see either 1, 2, or 3 cats.
  There are 3 different enticements:

  - Milkshake
  - Fish
  - Nothing

  And there are different probabilities how many cats there are given
  an enticement, see below.

  Now: We see 3 cats, what is the probability that it's a milkshake?

  The video got the following (for 3 cats):
   - milkshake: 0.42
   - fish: 0.04
   - nothing: 0.03

  Normalized to percentage (from the video):

  0.42/(0.42 + 0.04 + 0.03) milkshake
      0.85714285714285714286
  0.04/(0.42 + 0.04 + 0.03) fish
      0.081632653061224489796
  0.03/(0.42 + 0.04 + 0.03)  nothing
      0.061224489795918367347

  Here are two models.
  Model2 seems to be the correct one (i.e. corresponds to the video),
  i.e. where the enticements are mutual exclusive.
  In model1 we accept that fish and milkshake can coexist.

  See ~/cplint/cat.pl
      ~/blog/cat.blog
      ~/psi/cat.psi
=#

using Turing # , StatsPlots, DataFrames
include("jl_utils.jl")
# using ReverseDiff, Zygote, Tracker
# Turing.setadbackend(:reversediff)
# Turing.setadbackend(:zygote)
# Turing.setadbackend(:tracker)

# This is not correct
@model function cat1()
    NoCat = 0
    OneCat = 1
    TwoCats = 2
    ThreeCats = 3

    # probability of an enticement in the garden
    milkshake      ~ Bernoulli(0.6)
    fish           ~ Bernoulli(0.1)
    no_enticement  ~ Bernoulli(0.3)
    println("milkshake:$milkshake fish:$fish no_enticement:$no_enticement")

    # Number of cats per enticement
    vs = [NoCat,OneCat,TwoCats,ThreeCats]
    vs_cat ~ DiscreteUniform(NoCat,ThreeCats)
    if !no_enticement && milkshake
        cat ~ Categorical(simplex([0.0,0.1,0.2,0.7]))
    elseif !no_enticement && fish
        cat ~ Categorical(simplex([0.0,0.2,0.3,0.4]))
    elseif no_enticement && !fish && !milkshake
        cat ~ Categorical(simplex([0.0,0.6,0.3,0.1]))
    else
        cat ~ Categorical(simplex([1.0,0.0,0.0,0.0]))
    end
    println("cat before:$cat  (vs[cat]: $(vs[cat]))")
    true ~ Dirac(vs_cat == vs[cat]) # || begin Turing.@addlogprob! -Inf;  end
    true ~ Dirac(vs_cat == ThreeCats) #  || begin Turing.@addlogprob! -Inf;  end
    println("cat after:$cat")

end

#=
model = cat1()
num_chains = 4
# chains = sample(model, Prior(), MCMCThreads(), 1000, num_chains)
chains = sample(model, MH(), MCMCThreads(), 10_000, num_chains)
# chains = sample(model, MH(), 1000)

display(chains)
=#

# This is correct (according to the video)
#=
Summary Statistics
            parameters      mean       std   naive_se      mcse         ess      rhat
                Symbol   Float64   Float64    Float64   Float64     Float64   Float64

                  cats    3.0000    0.0000     0.0000    0.0000         NaN       NaN
            enticement    1.2009    0.5322     0.0013    0.0090   2218.9994    1.0026
       enticement_fish    0.0783    0.2686     0.0007    0.0043   2703.4937    1.0026
  enticement_milkshake    0.8600    0.3470     0.0009    0.0058   2345.7527    1.0017
    enticement_nothing    0.0617    0.2407     0.0006    0.0041   2176.3555    1.0048
=#
enticements = []
enticements_dict = Dict(1=>"milkshake",2=>"fish",3=>"no_enticement")
@model cat2() = begin

    milkshake = 1
    fish = 2
    no_enticement = 3

    # Priors
    # probability of an enticement in the garden
    enticement ~ Categorical(simplex([0.6,0.1,0.3]))
    # Note: This seems a little too complex...
    enticement_milkshake ~ Bernoulli(0.5)
    enticement_fish      ~ Bernoulli(0.5)
    enticement_nothing   ~ Bernoulli(0.5)
    cats ~ DiscreteUniform(1,3)
    if enticement == milkshake
        cats ~ Categorical(simplex([0.1,0.2,0.7]))
        enticement_milkshake ~ Bernoulli(1.0)
        enticement_fish      ~ Bernoulli(0.0)
        enticement_nothing   ~ Bernoulli(0.0)
    elseif enticement == fish
        cats ~ Categorical(simplex([0.2,0.4,0.4]))
        enticement_milkshake ~ Bernoulli(0.0)
        enticement_fish      ~ Bernoulli(1.0)
        enticement_nothing   ~ Bernoulli(0.0)
    else
        cats ~ Categorical(simplex([0.6,0.3,0.1]))
        enticement_milkshake ~ Bernoulli(0.0)
        enticement_fish      ~ Bernoulli(0.0)
        enticement_nothing   ~ Bernoulli(1.0)
    end

    # println("enticement:$(enticements_dict[$enticement]) cats:$cats")

    # We observe 3 cats
    true ~ Dirac(cats == 3) # || begin Turing.@addlogprob! -Inf; return; end
    # push!(enticements,[enticements_dict[enticement],cats])

    return enticement
end

model = cat2()
num_chains = 4
# chains = sample(model, Prior(), MCMCThreads(), 1000, num_chains)
chains = sample(model, MH(), MCMCThreads(), 40_000, num_chains)
# chains = sample(model, MH(), 40_000)

# chains = sample(model, PG(20), MCMCThreads(), 1000, num_chains)

display(chains)

println("\nEnticement: milkshake:1 fish:2 no_enticement:3")
gen = generated_quantities(model, chains)
show_var_dist_pct(gen, 120)


# ------------------------------------------------------------------------------------------
# # Loops
#
# Topics:
# 1. `while` loops
# 2. `for` loops
#

#
# ## while loops
#
# The syntax for a `while` is
#
# ```julia
# while *condition*
#     *loop body*
# end
# ```
#
# For example, we could use `while` to count or to iterate over an array.
# ------------------------------------------------------------------------------------------

n = 0
while n < 10
    n += 1
    println(n)
end
n

myfriends = ["Ted", "Robyn", "Barney", "Lily", "Marshall"]

i = 1
while i <= length(myfriends)
    friend = myfriends[i]
    println("Hi $friend, it's great to see you!")
    i += 1
end

# ------------------------------------------------------------------------------------------
# ## for loops
#
# The syntax for a `for` loop is
#
# ```julia
# for *var* in *loop iterable*
#     *loop body*
# end
# ```
#
# We could use a for loop to generate the same results as either of the examples above:
# ------------------------------------------------------------------------------------------

for n in 1:10
    println(n)
end

myfriends = ["Ted", "Robyn", "Barney", "Lily", "Marshall"]

for friend in myfriends
    println("Hi $friend, it's great to see you!")
end

# ------------------------------------------------------------------------------------------
# Now let's use `for` loops to create some addition tables, where the value of every entry
# is the sum of its row and column indices.

#
# Note that we iterate over this array via column-major loops in order to get the best
# performance. More information about fast indexing of multidimensional arrays inside nested
# loops can be found at https://docs.julialang.org/en/v1/manual/performance-tips/#Access-
# arrays-in-memory-order,-along-columns-1
#
# First, we initialize an array with zeros.
# ------------------------------------------------------------------------------------------

m, n = 5, 5
A = fill(0, (m, n))

for j in 1:n
    for i in 1:m
        A[i, j] = i + j
    end
end
A

# ------------------------------------------------------------------------------------------
# Here's some syntactic sugar for the same nested `for` loop
# ------------------------------------------------------------------------------------------

B = fill(0, (m, n))

for j in 1:n, i in 1:m
    B[i, j] = i + j
end
B

# ------------------------------------------------------------------------------------------
# The more "Julia" way to create this addition table would have been with an *array
# comprehension*.
# ------------------------------------------------------------------------------------------

C = [i + j for i in 1:m, j in 1:n]

# ------------------------------------------------------------------------------------------
# ### Exercises
#
# #### 4.1
# Loop over integers between 1 and 100 and print their squares.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### 4.2
# Add to the code above a bit to create a dictionary, `squares` that holds integers and
# their squares as key, value pairs such that
#
# ```julia
# squares[10] == 100
# ```
# ------------------------------------------------------------------------------------------

@assert squares[10] == 100
@assert squares[11] == 121

# ------------------------------------------------------------------------------------------
# #### 4.3
# Use an array comprehension to create an an array `squares_arr` that stores the squares for
# all integers between 1 and 100.
# ------------------------------------------------------------------------------------------

@assert length(squares_arr) == 100
@assert sum(squares_arr) == 338350

# ------------------------------------------------------------------------------------------
# Please click on `Validate` on the top, once you are done with the exercises.
# ------------------------------------------------------------------------------------------


### mcvar and mcse stand for Monte Carlo variance and Monte Carlo error respectively

## Monte Carlo variance assuming IID samples

mcvar_iid(v::AbstractArray) = var(v)/length(v)

mcvar_iid(v::AbstractArray, region) = mapslices(mcvar_iid, v, region)

mcvar_iid(s::VariableNState{Univariate}) = mcvar_iid(s.value)

mcvar_iid(s::VariableNState{Multivariate}, i::Integer) = mcvar_iid(s.value[i, :])

mcvar_iid(s::VariableNState{Multivariate}, r::AbstractVector=1:s.size) = eltype(s)[mcvar_iid(s, i) for i in r]

## Monte Carlo standard error assuming IID samples

mcse_iid(v::AbstractArray) = sqrt(mcvar_iid(v))

mcse_iid(v::AbstractArray, region) = mapslices(mcse_iid, v, region)

mcse_iid(s::VariableNState{Univariate}) = mcse_iid(s.value)

mcse_iid(s::VariableNState{Multivariate}, i::Integer) = mcse_iid(s.value[i, :])

mcse_iid(s::VariableNState{Multivariate}, r::AbstractVector=1:s.size) = eltype(s)[mcse_iid(s, i) for i in r]

## Monte Carlo variance using batch means
## Reference:
## Flegal J.M and Jones G.L
## Batch Means and Spectral Variance Estimators in Markov chain Monte Carlo
## Annals of Statistics, 2010, 38 (2), pp 1034-1070

function mcvar_bm{T}(v::AbstractArray{T}; batchlen::Integer=100)
  nbatches = div(length(v), batchlen)
  @assert nbatches > 1 "Choose batch size such that the number of batches is greather than one"
  nbsamples = nbatches*batchlen
  batchmeans = T[mean(v[((j-1)*batchlen+1):(j*batchlen)]) for j = 1:nbatches]
  return batchlen*var(batchmeans)/nbsamples
end

mcvar_bm(v::AbstractArray, region; batchlen::Integer=100) = mapslices(x -> mcvar_bm(x, batchlen=batchlen), v, region)

mcvar_bm(s::VariableNState{Univariate}; batchlen::Integer=100) = mcvar_bm(s.value, batchlen=batchlen)

mcvar_bm(s::VariableNState{Multivariate}, i::Integer; batchlen::Integer=100) = mcvar_bm(s.value[i, :], batchlen=batchlen)

mcvar_bm(s::VariableNState{Multivariate}, r::AbstractVector=1:s.size; batchlen::Integer=100) =
  eltype(s)[mcvar_bm(s, i, batchlen=batchlen) for i in r]

## Monte Carlo standard error using batch means

mcse_bm(v::AbstractArray; batchlen::Integer=100) = sqrt(mcvar_bm(v, batchlen=batchlen))

mcse_bm(v::AbstractArray, region; batchlen::Integer=100) = mapslices(x -> mcse_bm(x, batchlen=batchlen), v, region)

mcse_bm(s::VariableNState{Univariate}; batchlen::Integer=100) = mcse_bm(s.value, batchlen=batchlen)

mcse_bm(s::VariableNState{Multivariate}, i::Integer; batchlen::Integer=100) = mcse_bm(s.value[i, :], batchlen=batchlen)

mcse_bm(s::VariableNState{Multivariate}, r::AbstractVector=1:s.size; batchlen::Integer=100) =
  eltype(s)[mcse_bm(s, i, batchlen=batchlen) for i in r]

## Initial monotone sequence estimator (IMSE) of Monte Carlo variance
## Reference:
## Geyer C.J
## Practical Markov Chain Monte Carlo
## Statistical Science, 1992, 7 (4), pp 473-483

function mcvar_imse{T}(v::AbstractVector{T}; maxlag::Integer=length(v)-1)
  k = convert(Integer, floor((maxlag-1)/2))
  m = k+1

  # Preallocate memory
  g = Array(T, k+1)

  # Calculate empirical autocovariance
  acv = autocov(v, 0:maxlag)

  # Calculate \hat{G}_{n, m} from the autocovariances, see pp 477 in Geyer
  for j = 0:k
    g[j+1] = acv[2*j+1]+acv[2*j+2]
    if g[j+1] <= 0
      m = j
      break
    end
  end

  # Create the monotone sequence of g
  if m > 1
    for j = 2:m
      if g[j] > g[j-1]
        g[j] = g[j-1]
      end
    end
  end

  # Calculate the initial monotone sequence estimator
  return (-acv[1]+2*sum(g[1:m]))/length(v)
end

mcvar_imse(v::AbstractArray; maxlag::Integer=length(v)-1) = mcvar_imse(vec(v), maxlag=maxlag)

mcvar_imse(v::AbstractArray, region; maxlag::Integer=length(v)-1) = mapslices(x -> mcvar_imse(x, maxlag=maxlag), v, region)

mcvar_imse(v::AbstractArray, region) = mapslices(x -> mcvar_imse(x, maxlag=length(x)-1), v, region)

mcvar_imse(s::VariableNState{Univariate}; maxlag::Integer=s.n-1) = mcvar_imse(s.value, maxlag=maxlag)

mcvar_imse(s::VariableNState{Multivariate}, i::Integer; maxlag::Integer=s.n-1) = mcvar_imse(s.value[i, :], maxlag=maxlag)

mcvar_imse(s::VariableNState{Multivariate}, r::AbstractVector=1:s.size; maxlag::Integer=s.n-1) =
  eltype(s)[mcvar_imse(s, i, maxlag=maxlag) for i in r]

## Initial monotone sequence estimator (IMSE) of Monte Carlo standard error

mcse_imse(v::AbstractArray; maxlag::Integer=length(v)-1) = sqrt(mcvar_imse(v, maxlag=maxlag))

mcse_imse(v::AbstractArray, region; maxlag::Integer=length(v)-1) = mapslices(x -> mcse_imse(x, maxlag=maxlag), v, region)

mcse_imse(v::AbstractArray, region) = mapslices(x -> mcse_imse(x, maxlag=length(x)-1), v, region)

mcse_imse(s::VariableNState{Univariate}; maxlag::Integer=s.n-1) = mcse_imse(s.value, maxlag=maxlag)

mcse_imse(s::VariableNState{Multivariate}, i::Integer; maxlag::Integer=s.n-1) = mcse_imse(s.value[i, :], maxlag=maxlag)

mcse_imse(s::VariableNState{Multivariate}, r::AbstractVector=1:s.size; maxlag::Integer=s.n-1) =
  eltype(s)[mcse_imse(s, i, maxlag=maxlag) for i in r]

## Initial positive sequence estimator (IPSE) of Monte Carlo variance
## Reference:
## Geyer C.J
## Practical Markov Chain Monte Carlo
## Statistical Science, 1992, 7 (4), pp 473-483

function mcvar_ipse{T}(v::AbstractVector{T}; maxlag::Integer=length(v)-1)
  k = convert(Integer, floor((maxlag-1)/2))

  # Preallocate memory
  g = Array(T, k+1)
  m = k+1

  # Calculate empirical autocovariance
  acv = autocov(v, 0:maxlag)

  # Calculate \hat{G}_{n, m} from the autocovariances, see pp. 477 in Geyer
  for j = 0:k
    g[j+1] = acv[2*j+1]+acv[2*j+2]
    if g[j+1] <= 0
      m = j
      break
    end
  end

  # Calculate the initial positive sequence estimator
  return (-acv[1]+2*sum(g[1:m]))/length(v)
end

mcvar_ipse(v::AbstractArray; maxlag::Integer=length(v)-1) = mcvar_ipse(vec(v), maxlag=maxlag)

mcvar_ipse(v::AbstractArray, region; maxlag::Integer=length(v)-1) = mapslices(x -> mcvar_ipse(x, maxlag=maxlag), v, region)

mcvar_ipse(v::AbstractArray, region) = mapslices(x -> mcvar_ipse(x, maxlag=length(x)-1), v, region)

mcvar_ipse(s::VariableNState{Univariate}; maxlag::Integer=s.n-1) = mcvar_ipse(s.value, maxlag=maxlag)

mcvar_ipse(s::VariableNState{Multivariate}, i::Integer; maxlag::Integer=s.n-1) = mcvar_ipse(s.value[i, :], maxlag=maxlag)

mcvar_ipse(s::VariableNState{Multivariate}, r::AbstractVector=1:s.size; maxlag::Integer=s.n-1) =
  eltype(s)[mcvar_ipse(s, i, maxlag=maxlag) for i in r]

## Initial positive sequence estimator (IPSE) of Monte Carlo standard error

mcse_ipse(v::AbstractArray; maxlag::Integer=length(v)-1) = sqrt(mcvar_ipse(v, maxlag=maxlag))

mcse_ipse(v::AbstractArray, region; maxlag::Integer=length(v)-1) = mapslices(x -> mcse_ipse(x, maxlag=maxlag), v, region)

mcse_ipse(v::AbstractArray, region) = mapslices(x -> mcse_ipse(x, maxlag=length(x)-1), v, region)

mcse_ipse(s::VariableNState{Univariate}; maxlag::Integer=s.n-1) = mcse_ipse(s.value, maxlag=maxlag)

mcse_ipse(s::VariableNState{Multivariate}, i::Integer; maxlag::Integer=s.n-1) = mcse_ipse(s.value[i, :], maxlag=maxlag)

mcse_ipse(s::VariableNState{Multivariate}, r::AbstractVector=1:s.size; maxlag::Integer=s.n-1) =
  eltype(s)[mcse_ipse(s, i, maxlag=maxlag) for i in r]

## Wrapper function for estimating Monte Carlo variance using various approaches

mcvar_types = (:iid, :bm, :imse, :ipse)
mcvar_functions = Dict{Symbol, Function}(zip(mcvar_types, (mcvar_iid, mcvar_bm, mcvar_imse, mcvar_ipse)))
mcse_functions = Dict{Symbol, Function}(zip(mcvar_types, (mcse_iid, mcse_bm, mcse_imse, mcse_ipse)))

mcvar(v::AbstractArray; vtype::Symbol=:imse, args...) = mcvar_functions[vtype](v; args...)

mcvar(v::AbstractArray, region; vtype::Symbol=:imse, args...) = mcvar_functions[vtype](v, region; args...)

mcvar(s::VariableNState{Univariate}; vtype::Symbol=:imse, args...) = mcvar_functions[vtype](s; args...)

mcvar(s::VariableNState{Multivariate}, i::Integer; vtype::Symbol=:imse, args...) = mcvar_functions[vtype](s, i; args...)

mcvar(s::VariableNState{Multivariate}, r::AbstractVector=1:s.size; vtype::Symbol=:imse, args...) =
  mcvar_functions[vtype](s, r; args...)

## Wrapper function for estimating Monte Carlo standard error using various approaches

mcse(v::AbstractArray; vtype::Symbol=:imse, args...) = mcse_functions[vtype](v; args...)

mcse(v::AbstractArray, region; vtype::Symbol=:imse, args...) = mcse_functions[vtype](v, region; args...)

mcse(s::VariableNState{Univariate}; vtype::Symbol=:imse, args...) = mcse_functions[vtype](s; args...)

mcse(s::VariableNState{Multivariate}, i::Integer; vtype::Symbol=:imse, args...) = mcse_functions[vtype](s, i; args...)

mcse(s::VariableNState{Multivariate}, r::AbstractVector=1:s.size; vtype::Symbol=:imse, args...) =
  mcse_functions[vtype](s, r; args...)