#############################
#           pHM.jl          #
#          pH-Mouse         #
#     @author Tanner Lee    #
#      @date 9/27/2018      #
#############################

# Start by loading the initializing data in memory
using DelimitedFiles
print("Which input number would you like to load? \t")
input = readline(stdin)
print("How many iterations would you like to run? \t")
iterations = readline(stdin)
iterations = parse(Int64, iterations)
print("Which room would you like to start in? \t\t")
startingLocation = readline(stdin)
startingLocation = parse(Int64, startingLocation)
print("Enter an intensity for the heatmap {1-10}: \t")
intensity = readline(stdin)
intensity = parse(Int64, intensity)
input = "input" * input * ".phm"
println("!CLEARING TMP FOLDER!")
rm("./tmp", force=true, recursive=true)
mkdir("./tmp")
println("!TMP SUCCESSFULLY CLEARED!")
println("!IMPORTING DATA FROM $input !")
A = readdlm(input, String)
println("!DATA IMPORTED SUCCESSFULLY!")

Arows = size(A)[1]
Acols = size(A)[2]

# Create the transpose manually
# Julia cannot transpose a string matrix
Atrans = Array{String}(undef, Arows, Acols)
for i=1:Arows
    for j=1:Acols
        Atrans[j,i] = A[i,j]
    end
end

println(A)
println(Atrans)

# Create an empty matrix we can concat with
# The columns we want are the pivot columns
probMatrix = zeros(Arows*Acols,Arows*Acols)

# Parse through A to create probability matrix
for i=1:Arows       # 1 < i < rows
    for j=1:Acols   # 1 < j < cols
        global A
        global probMatrix
        up = false
        down = false
        right = false
        left = false
        upDigit = 0
        downDigit = 0
        leftDigit = 0
        rightDigit = 0
        probCounter = 0
        tempString = ""
        tempRoomNumber = 0
        currentRoomString = A[i,j]
        currentRoomNumber = chop(currentRoomString)
        currentRoomNumber = parse(Int64,currentRoomNumber)
        # Check if the current room is a wall, or cheese
        if (string(currentRoomString[length(currentRoomString)]) == "W")
            # Current Room is a wall
            continue
        elseif (string(currentRoomString[length(currentRoomString)]) == "C")
            # Current Room is cheese
            probMatrix[currentRoomNumber,currentRoomNumber] = 1.0
            continue
        end
        # Check how many neighbors A[i,j] has
        # A[i,j] has neighbors at --
        # A[i+1,j], A[i-1,j], A[i,j+1], A[i,j-1]
        # However all of these may not exist
        # Boundary conditions must be considered
        if i != 1
            # We are not on the top row, decrease row
            tempString = A[i-1,j]
            # Get the room number from this neighbor
            tempRoomNumber = chop(tempString)
            # Ensure this room isnt a wall
            if (string(tempString[length(tempString)]) != "W")
                # We can enter this room
                probCounter += 1
                up = true
                upDigit = parse(Int64,tempRoomNumber)
            end
        end
        if i != Arows
            # We are not on the last row, increase row
            tempString = A[i+1,j]
            # Get the room number from this neighbor
            tempRoomNumber = chop(tempString)
            # Ensure this room isnt a wall
            if (string(tempString[length(tempString)]) != "W")
                # We can enter this room
                probCounter += 1
                down = true
                downDigit = parse(Int64,tempRoomNumber)
            end
        end
        if j != 1
            # We are not on the first col, decrease col
            tempString = A[i,j-1]
            # Get the room number from this neighbor
            tempRoomNumber = chop(tempString)
            # Ensure this room isnt a wall
            if (string(tempString[length(tempString)]) != "W")
                # We can enter this room
                probCounter += 1
                left = true
                leftDigit = parse(Int64,tempRoomNumber)
            end
        end
        if j != Acols
            # We are not on the last col, increase col
            tempString = A[i,j+1]
            # Get the room number from this neighbor
            tempRoomNumber = chop(tempString)
            # Ensure this room isnt a wall
            if (string(tempString[length(tempString)]) != "W")
                # We can enter this room
                probCounter += 1
                right = true
                rightDigit = parse(Int64,tempRoomNumber)
            end
        end
        # Create this rooms column in the probability matrix
        if (up)
            # We can go up
            probMatrix[upDigit,currentRoomNumber] = 1/probCounter
        end
        if (down)
            probMatrix[downDigit,currentRoomNumber] = 1/probCounter
        end
        if (left)
            probMatrix[leftDigit,currentRoomNumber] = 1/probCounter
        end
        if (right)
            probMatrix[rightDigit,currentRoomNumber] = 1/probCounter
        end
    end
end

# We now have a probability matrix for our maze
# Create some an input for it
x = zeros(Arows*Acols,1)
x[startingLocation] = intensity

# Use the following code segment to print the probMatrix
# For debugging purposes only
#=
sizeProbMatrix = size(probMatrix)
for i=1:sizeProbMatrix[1]
    print("Row $i:\t [")
    for j=1:15
        print(probMatrix[i,j])
        print(", ")
    end
    println("]")
end
=#

# Complete Luxor Initialization
# Initialize Room Size in pixels
global rSize = 100
# Initialize Canvas Size
global cSize = rSize*Arows
println("!INITIALIZING LUXOR!")
using Luxor
animation = Movie(cSize, cSize, "pHM", 1:iterations)
println("!LUXOR INITIALIZED!")
println("!INITIALIZING COLORSCHEMES!")
using Colors, ColorSchemes
println("!COLORSCHEMES INITIALIZED!")

# Run the simulation
println("!INITIAL INPUT!")
println(x)
println("!STARTING SIMULATION!")

function backdrop(scene, framenumber)
    background("white")
    # Initialize Table and digits
    global table = Table(Arows, Acols, rSize, rSize)
    sethue("black")
    # Loop through each room of the maze and print it to the image
    for i in 1:length(table)
        if (string(Atrans[i][length(Atrans[i])]) == "W")
            # If the current room is a wall make it grey
            box.(table[i], rSize, rSize, :fill)
        end
        # Put a black border on the room and print its name in the middle
        text(Atrans[i], table[i], halign=:center, valign=:middle)
        box.(table[i], rSize, rSize, :stroke)
    end
end

function frame(scene, framenumber)
    global table
    global rSize
    global x
    global probMatrix
    for i=1:length(x)
        sethue(get(ColorSchemes.YlOrRd_7, x[i]))
        setmode("darken")   # Set mode to darken to ensure text isnt covered
        box.(table[i], rSize, rSize, :fill)
    end
    x = probMatrix*x
end

animate(animation, [
    Scene(animation, backdrop, 0:iterations),
    Scene(animation, frame, 0:iterations)],
    creategif=true,
    framerate=20,
    pathname="./images/pHM.gif",
    tempdirectory="./tmp")

println("!SIMULATION COMPLETE!")
println("!FINAL OUTPUT!")
println(x)
println()


#
# Rhapsodie.jl
#
# Package for the Reconstruction of High-contrAst Polarized
# SOurces and Deconvolution for cIrcumstellar Environments (Rhapsodie)
#
#----------------------------------------------------------
#
#
# Copiyright (c) 2017-2021 Laurence Denneulin (see LICENCE.md)
#

module Rhapsodie

    export
        PolarimetricPixel,
        PolarimetricMap,
        write,
        read,
        convert,
        load_parameters,
        get_par,
        load_data,
        fg!,
        set_fft_op,
        TwoDimensionalTransformInterpolator,
        FieldTransformOperator,
        data_simulator,
        generate_model,
        data_generator,
        generate_parameters,
        Double_Difference,
        Double_Ratio,
        Linear_Method,
        NonLinear_Method,
        apply_rhapsodie

    import Base: +, -, *, /, ==, getindex, setindex!, read, write, convert

    using OptimPackNextGen
    import OptimPackNextGen: BraDi #va devenir BraDi avec un D majuscule
    using SpecialFunctions
    using TwoDimensional
    using FFTW
    using LinearInterpolators
    using Statistics
    using LinearAlgebra
    using LazyAlgebra
    import LazyAlgebra: Mapping, vcreate, vcopy, apply!
    using StaticArrays
    using FITSIO
    using EasyFITS
    using DelimitedFiles
    using TiPi
    using Random

    include("Polarimetric_Parameters.jl")
    include("grad_tools.jl")
    include("separable_methods.jl")
    include("rhapsodie_methods.jl")
    include("datasimul_tools.jl")
end


struct TimesteppingParams{T <: Real, N <: Integer, TO, C, P}
            niter0 :: N      #starting iteration number
                dt :: T      #timestep
          end_time :: T      #end time of this simulation
                t0 :: T      #start time of this simulation
        chkpt_freq :: T      #temporary checkpoint frequency
       pchkpt_freq :: T      #permanent checkpoint frequency
       chkpt_path  :: String #path to location of permanent and tempoary checkpoint output
      n_iter_total :: TO     #total number of timesteps counting from zero
      n_iter_chkpt :: C      #number of iterations per temporary checkpoint
     n_iter_pchkpt :: P      #number of iterations per permanent checkpoint
    step_thickness :: Bool   #toggle whether to step the thickness at each timestep or not (coupling control)
end

"""
TimesteppingParams(;
                    niter0 = 0,
                    dt = 1.0,
                    end_time = 1.0,
                    t0 = nothing,
                    chkpt_freq = Inf,
                    pchkpt_freq = Inf,
                    chkpt_path = './',
                    step_thickness = true)

Construct a WAVI.jl TimesteppingParams object.
TimesteppingParams stores information relating to timestepping.

Keyword arguments
=================
- 'niter0': Iteration number of the first timestep. niter0 = 0 corresponds to a new simulation, while niter0 > 0 (positive integer) corresponds to a pickup.
- 'dt': Model timestep
- 'end_time': Simulation termination time
- 't0': Starting time of the simulation
- 'chkpt_freq': Frequency of outputting temporary checkpoints
- 'pchkpt_freq': Frequecy with which permanent checkpoints are pass
- 'chkpt_path' : Path to location checkpoint output
- 'step_thickness': Toggle whether to update the ice thickness (true) or not (false) at each timestep
"""
function TimesteppingParams(;
                        niter0 = 0,
                        dt = 1.0,
                        end_time = nothing,
                        n_iter_total = nothing,
                        t0 = nothing,
                        chkpt_freq = Inf,
                        pchkpt_freq = Inf,
                        chkpt_path = "./",
                        step_thickness = true)

    #initialize t0 (really you should read start time from pickup file)
    t0 = niter0 > 0 ? niter0 * dt : 0
    t0 = map(typeof(dt), t0)

    #check compatibility of n_iter_total and end_time, and compute them
    end_time, n_iter_total = compute_iterations_and_end_time(end_time, n_iter_total,dt)

    #compute number of timesteps checkpoint number of timesteps
    chkpt_freq == Inf ? n_iter_chkpt = Inf : n_iter_chkpt  = round(Int, chkpt_freq/dt)
    pchkpt_freq == Inf ? n_iter_pchkpt = Inf : n_iter_pchkpt = round(Int, pchkpt_freq/dt)

    #check the output path ends in '/' and exists
    endswith(chkpt_path, "/") || (chkpt_path = string(chkpt_path, "/"))
    if ~isdir(chkpt_path)
        @warn string("Did not find output path ", chkpt_path, ". Any outputs will go to the working directory", pwd())
        chkpt_path = "./"
    end

    return TimesteppingParams(niter0, dt, end_time, t0, chkpt_freq, pchkpt_freq,
                            chkpt_path,n_iter_total, n_iter_chkpt, n_iter_pchkpt, step_thickness)
end

"""
    compute_iterations_and_end_time(end_time, n_iter_total, dt)

Check input end time and total iterations, and compute those not passed.
"""
function compute_iterations_and_end_time(end_time, n_iter_total, dt)
    (isa(dt, Real) && (dt > 0))  || throw(ArgumentError("timestep dt must be a positive number"))

    if (~(n_iter_total === nothing) && ~(end_time === nothing)) #if both passed, throw error if incompatible
        (end_time ≈ (n_iter_total * dt)) ||  throw(ArgumentError("You have specified both end time (end_time) and total iterations (n_iter_total), but their values are incompatible: end time mustequal n_iter_total * dt"))
    elseif ((n_iter_total === nothing) && ~(end_time === nothing)) #if only end time passed, n_iter_total is the nearest integer
        end_time == Inf ? n_iter_total = Inf : n_iter_total  = round(Int, end_time/dt)
    elseif (~(n_iter_total === nothing) && (end_time === nothing)) #if only number of iterations
        n_iter_total == Inf ? end_time = Inf : end_time  = dt*n_iter_total
    elseif (n_iter_total === nothing) && (end_time === nothing) #if neither is passed
        throw(ArgumentError("You must pass at least one of end_time or n_iter_total"))
    end

    #check both are positive numbers, now that both end_time and n_iter_total assigned
    (isa(n_iter_total, Integer) && (n_iter_total > 0)) || throw(ArgumentError("n_iter_total must be a positive integer"))
    (isa(end_time, Real) && (end_time > 0))|| throw(ArgumentError("end_time must be a positive number"))
    return end_time, n_iter_total
end


using DataDeps

register(DataDep(
  "elevators",
  "https://www.openml.org/d/216",
  "https://www.openml.org/data/get_csv/3653/dataset_2202_elevators.arff",
  "6c80dc6d2e78b42693ecff38ac42062774a5294cfa1f22e509601ae72f7a6af4",
  post_fetch_method=(path -> begin
    UCIData.process_dataset(path,
        target_index=19,
        feature_indices=1:18,
        header=true,
    )
  end),
))


@testset "583.delete-operation-for-two-strings.jl" begin
    @test min_distance583("sea", "eat") == 2
end


using HackerNewsAnalysis

top_story_id = first ∘ fetch_top_stories
#=
julia> top_story_id = first ∘ fetch_top_stories
#58 (generic function with 1 method)
=#

top_story_id()
#=
julia> top_story_id()
21676252
=#

top_story = fetch_story ∘ first ∘ fetch_top_stories
top_story()
#=
julia> top_story = fetch_story ∘ first ∘ fetch_top_stories
#58 (generic function with 1 method)

julia> top_story()
Story("kkm", 384, 538, 1575218518, 21676252, "The world needs more search engines", [21677682, 21676507, 21676549, 21676457, 21676352, 21677053, 21676790, 21676616, 21676315, 21676385  …  21676442, 21676505, 21676858, 21676598, 21676736, 21677029, 21676919, 21676853, 21676397, 21676703], "https://www.0x65.dev/blog/2019-12-01/the-world-needs-cliqz-the-world-needs-more-search-engines.html")=#

title(s::Story) = s.title
top_story_title = title ∘ fetch_story ∘ first ∘ fetch_top_stories
top_story_title()
#=
julia> title(s::Story) = s.title
title (generic function with 1 method)

julia> top_story_title = title ∘ fetch_story ∘ first ∘ fetch_top_stories
#58 (generic function with 1 method)

julia> top_story_title()
"The world needs more search engines"
=#


module Day6

export get_inputs, get_solution1, get_solution2

## Input getting
function get_inputs()
    test_input1 = test_input2 = read_inputs("test_input1.txt")
    test_output1 = 11
    test_output2 = 6
    data = read_inputs("input.txt")
    return (; test_input1, test_input2, test_output1, test_output2, data)
end

read_inputs(f_name) = parse_declaration.(split(read(joinpath(@__DIR__, f_name), String), r"\n\n|\r\n\r\n"))

parse_declaration(str) = split(str, r"\n|\r\n", keepempty=false)

## Solution functions
# Part 1
get_solution1(data) = sum(length.(reduce.(union, data)))

# Part 2
get_solution2(data) = sum(length.(reduce.(intersect, data)))

end


function Ellipsoid(arg0::Frame, arg1::jdouble, arg2::jdouble, arg3::jdouble)
    return Ellipsoid((Frame, jdouble, jdouble, jdouble), arg0, arg1, arg2, arg3)
end

function get_a(obj::Ellipsoid)
    return jcall(obj, "getA", jdouble, ())
end

function get_b(obj::Ellipsoid)
    return jcall(obj, "getB", jdouble, ())
end

function get_c(obj::Ellipsoid)
    return jcall(obj, "getC", jdouble, ())
end

function get_frame(obj::Ellipsoid)
    return jcall(obj, "getFrame", Frame, ())
end

function get_plane_section(obj::Ellipsoid, arg0::Vector3D, arg1::Vector3D)
    return jcall(obj, "getPlaneSection", Ellipse, (Vector3D, Vector3D), arg0, arg1)
end

function is_inside(obj::Ellipsoid, arg0::Vector3D)
    return jcall(obj, "isInside", jboolean, (Vector3D,), arg0)
end

function point_on_limb(obj::Ellipsoid, arg0::Vector3D, arg1::Vector3D)
    return jcall(obj, "pointOnLimb", Vector3D, (Vector3D, Vector3D), arg0, arg1)
end


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

using SHA
using Test

const VERBOSE = false

# Define some data we will run our tests on
lorem = "Lorem ipsum dolor sit amet, consectetur adipisicing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum."
so_many_as_array = repeat([0x61], 1000000)
so_many_as_tuple = ntuple((i) -> 0x61, 1000000)
tempdir = mktempdir()
file = joinpath(tempdir, ".sha")
fIO = open(file, "w")
write(fIO, '\0')
close(fIO)
data = Any["", "test", lorem, file, so_many_as_array, so_many_as_tuple]

# Descriptions of the data, the SHA functions we'll run on the data, etc...
data_desc = ["the empty string", "the string \"test\"", "lorem ipsum",
             "0 file", "one million a's Array", "one million a's Tuple"]
sha_types = Dict(sha1 => SHA.SHA1_CTX,
            sha2_224 => SHA.SHA2_224_CTX, sha2_256 => SHA.SHA2_256_CTX, sha2_384 => SHA.SHA2_384_CTX, sha2_512 => SHA.SHA2_512_CTX,
            sha3_224 => SHA.SHA3_224_CTX, sha3_256 => SHA.SHA3_256_CTX, sha3_384 => SHA.SHA3_384_CTX, sha3_512 => SHA.SHA3_512_CTX)
sha_funcs = [sha1,
             sha2_224, sha2_256, sha2_384, sha2_512,
             sha3_224, sha3_256, sha3_384, sha3_512]
ctxs = [SHA1_CTX,
        SHA2_224_CTX, SHA2_256_CTX, SHA2_384_CTX, SHA2_512_CTX,
        SHA3_224_CTX, SHA3_256_CTX, SHA3_384_CTX, SHA3_512_CTX]
shws = ["SHA1 hash state",
        "SHA2 224-bit hash state", "SHA2 256-bit hash state", "SHA2 384-bit hash state", "SHA2 512-bit hash state",
        "SHA3 224-bit hash state", "SHA3 256-bit hash state", "SHA3 384-bit hash state", "SHA3 512-bit hash state"]

answers = Dict(
sha1 => [
"da39a3ee5e6b4b0d3255bfef95601890afd80709",
"a94a8fe5ccb19ba61c4c0873d391e987982fbbd3",
"19afa2a4a37462c7b940a6c4c61363d49c3a35f4",
"5ba93c9db0cff93f52b521d7420e43f6eda2784f",
"34aa973cd4c4daa4f61eeb2bdbad27316534016f",
"34aa973cd4c4daa4f61eeb2bdbad27316534016f"
],
sha2_224 => [
"d14a028c2a3a2bc9476102bb288234c415a2b01f828ea62ac5b3e42f",
"90a3ed9e32b2aaf4c61c410eb925426119e1a9dc53d4286ade99a809",
"6a0644abcf1e2cecbec2814443dab5f24b7ad8ebb66c75667ab67959",
"fff9292b4201617bdc4d3053fce02734166a683d7d858a7f5f59b073",
"20794655980c91d8bbb4c1ea97618a4bf03f42581948b2ee4ee7ad67",
"20794655980c91d8bbb4c1ea97618a4bf03f42581948b2ee4ee7ad67"
],
sha2_256 => [
"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
"9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08",
"2c7c3d5f244f1a40069a32224215e0cf9b42485c99d80f357d76f006359c7a18",
"6e340b9cffb37a989ca544e6bb780a2c78901d3fb33738768511a30617afa01d",
"cdc76e5c9914fb9281a1c7e284d73e67f1809a48a497200e046d39ccc7112cd0",
"cdc76e5c9914fb9281a1c7e284d73e67f1809a48a497200e046d39ccc7112cd0"
],
sha2_384 => [
"38b060a751ac96384cd9327eb1b1e36a21fdb71114be07434c0cc7bf63f6e1da274edebfe76f65fbd51ad2f14898b95b",
"768412320f7b0aa5812fce428dc4706b3cae50e02a64caa16a782249bfe8efc4b7ef1ccb126255d196047dfedf17a0a9",
"63980fd0425cd2c3d8a400ee0f2671ef135db03b947ec1af21b6e28f19c16ca272036469541f4d8e336ac6d1da50580f",
"bec021b4f368e3069134e012c2b4307083d3a9bdd206e24e5f0d86e13d6636655933ec2b413465966817a9c208a11717",
"9d0e1809716474cb086e834e310a4a1ced149e9c00f248527972cec5704c2a5b07b8b3dc38ecc4ebae97ddd87f3d8985",
"9d0e1809716474cb086e834e310a4a1ced149e9c00f248527972cec5704c2a5b07b8b3dc38ecc4ebae97ddd87f3d8985"
],
sha2_512 => [
"cf83e1357eefb8bdf1542850d66d8007d620e4050b5715dc83f4a921d36ce9ce47d0d13c5d85f2b0ff8318d2877eec2f63b931bd47417a81a538327af927da3e",
"ee26b0dd4af7e749aa1a8ee3c10ae9923f618980772e473f8819a5d4940e0db27ac185f8a0e1d5f84f88bc887fd67b143732c304cc5fa9ad8e6f57f50028a8ff",
"f41d92bc9fc1157a0d1387e67f3d0893b70f7039d3d46d8115b5079d45ad601159398c79c281681e2da09bf7d9f8c23b41d1a0a3c5b528a7f2735933a4353194",
"b8244d028981d693af7b456af8efa4cad63d282e19ff14942c246e50d9351d22704a802a71c3580b6370de4ceb293c324a8423342557d4e5c38438f0e36910ee",
"e718483d0ce769644e2e42c7bc15b4638e1f98b13b2044285632a803afa973ebde0ff244877ea60a4cb0432ce577c31beb009c5c2c49aa2e4eadb217ad8cc09b",
"e718483d0ce769644e2e42c7bc15b4638e1f98b13b2044285632a803afa973ebde0ff244877ea60a4cb0432ce577c31beb009c5c2c49aa2e4eadb217ad8cc09b"
],
sha3_224 => [
"6b4e03423667dbb73b6e15454f0eb1abd4597f9a1b078e3f5b5a6bc7",
"3797bf0afbbfca4a7bbba7602a2b552746876517a7f9b7ce2db0ae7b",
"ea5395370949ad8c7d2ca3e7c045ef3306fe3a3f4740de452ef87a28",
"bdd5167212d2dc69665f5a8875ab87f23d5ce7849132f56371a19096",
"d69335b93325192e516a912e6d19a15cb51c6ed5c15243e7a7fd653c",
"d69335b93325192e516a912e6d19a15cb51c6ed5c15243e7a7fd653c"
],
sha3_256 => [
"a7ffc6f8bf1ed76651c14756a061d662f580ff4de43b49fa82d80a4b80f8434a",
"36f028580bb02cc8272a9a020f4200e346e276ae664e45ee80745574e2f5ab80",
"8c8142d2ca964ab307ace567ddd5764f17ebb76eb8ff25543ab54c14fe2ab139",
"5d53469f20fef4f8eab52b88044ede69c77a6a68a60728609fc4a65ff531e7d0",
"5c8875ae474a3634ba4fd55ec85bffd661f32aca75c6d699d0cdcb6c115891c1",
"5c8875ae474a3634ba4fd55ec85bffd661f32aca75c6d699d0cdcb6c115891c1",
],
sha3_384 => [
"0c63a75b845e4f7d01107d852e4c2485c51a50aaaa94fc61995e71bbee983a2ac3713831264adb47fb6bd1e058d5f004",
"e516dabb23b6e30026863543282780a3ae0dccf05551cf0295178d7ff0f1b41eecb9db3ff219007c4e097260d58621bd",
"eb9fbba3eb916a4efe384b3125f5d03ceb9c5c1b94431ac30fa86c54408b92701ca5d2628cd7113aa5541177ec3ccd1d",
"127677f8b66725bbcb7c3eae9698351ca41e0eb6d66c784bd28dcdb3b5fb12d0c8e840342db03ad1ae180b92e3504933",
"eee9e24d78c1855337983451df97c8ad9eedf256c6334f8e948d252d5e0e76847aa0774ddb90a842190d2c558b4b8340",
"eee9e24d78c1855337983451df97c8ad9eedf256c6334f8e948d252d5e0e76847aa0774ddb90a842190d2c558b4b8340",
],
sha3_512 => [
"a69f73cca23a9ac5c8b567dc185a756e97c982164fe25859e0d1dcc1475c80a615b2123af1f5f94c11e3e9402c3ac558f500199d95b6d3e301758586281dcd26",
"9ece086e9bac491fac5c1d1046ca11d737b92a2b2ebd93f005d7b710110c0a678288166e7fbe796883a4f2e9b3ca9f484f521d0ce464345cc1aec96779149c14",
"3a4318353396a12dfd20442cfce1d8ad4d7e732e85cc56b01b4cf9057a41c8827c0a03c70812e76ace68d776759225c213b4f581aac0dba5dd43b785b1a33fe5",
"7127aab211f82a18d06cf7578ff49d5089017944139aa60d8bee057811a15fb55a53887600a3eceba004de51105139f32506fe5b53e1913bfa6b32e716fe97da",
"3c3a876da14034ab60627c077bb98f7e120a2a5370212dffb3385a18d4f38859ed311d0a9d5141ce9cc5c66ee689b266a8aa18ace8282a0e0db596c90b0a7b87",
"3c3a876da14034ab60627c077bb98f7e120a2a5370212dffb3385a18d4f38859ed311d0a9d5141ce9cc5c66ee689b266a8aa18ace8282a0e0db596c90b0a7b87",
]
)

function describe_hash(T::Type{S}) where {S <: SHA.SHA_CTX}
    if T <: SHA.SHA1_CTX return "SHA1" end
    if T <: SHA.SHA2_CTX return "SHA2-$(SHA.digestlen(T)*8)" end
    if T <: SHA.SHA3_CTX return "SHA3-$(SHA.digestlen(T)*8)" end
end

VERBOSE && println("Loaded hash types: $(join(sort([describe_hash(t[2]) for t in sha_types]), ", ", " and "))")

# First, test processing the data in one go
nerrors = 0
for idx in 1:length(data)
    global nerrors

    desc = data_desc[idx]
    VERBOSE && print("Testing on $desc$(join(["." for z in 1:(34-length(desc))]))")
    nerrors_old = nerrors
    for sha_idx in 1:length(sha_funcs)
        sha_func = sha_funcs[sha_idx]

        if idx == 4
            open(data[idx]) do f
                hash = bytes2hex(sha_func(f))
            end
        else
            hash = bytes2hex(sha_func(data[idx]))
        end

        if hash != answers[sha_func][idx]
            print("\n")
            @warn(
            """
            For $(describe_hash(sha_types[sha_func])) expected:
                $(answers[sha_func][idx])
            Calculated:
                $(hash)
            """)
            nerrors += 1
        else
            VERBOSE && print(".")
        end
    end
    VERBOSE && println("Done! [$(nerrors - nerrors_old) errors]")
end

# Do another test on the "so many a's" data where we chunk up the data into
# two chunks, (sized appropriately to AVOID overflow from one update to another)
# in order to test multiple update!() calls
VERBOSE && print("Testing on one million a's (chunked properly)")
nerrors_old = nerrors
for sha_idx in 1:length(sha_funcs)
    global nerrors

    ctx = sha_types[sha_funcs[sha_idx]]()
    SHA.update!(ctx, so_many_as_array[1:2*SHA.blocklen(typeof(ctx))])
    SHA.update!(ctx, so_many_as_array[2*SHA.blocklen(typeof(ctx))+1:end])
    hash = bytes2hex(SHA.digest!(ctx))
    if hash != answers[sha_funcs[sha_idx]][end]
        print("\n")
        @warn(
        """
        For $(describe_hash(sha_types[sha_funcs[sha_idx]])) expected:
            $(answers[sha_funcs[sha_idx]][end-1])
        Calculated:
            $(hash)
        """)
        nerrors += 1
    else
        VERBOSE && print(".")
    end
end
VERBOSE && println("Done! [$(nerrors - nerrors_old) errors]")

# Do another test on the "so many a's" data where we chunk up the data into
# three chunks, (sized appropriately to CAUSE overflow from one update to another)
# in order to test multiple update!() calls as well as the overflow codepaths
VERBOSE && print("Testing on one million a's (chunked clumsily)")
nerrors_old = nerrors
for sha_idx in 1:length(sha_funcs)
    global nerrors
    ctx = sha_types[sha_funcs[sha_idx]]()

    # Get indices awkwardly placed for the blocklength of this hash type
    idx0 = round(Int, 0.3*SHA.blocklen(typeof(ctx)))
    idx1 = round(Int, 1.7*SHA.blocklen(typeof(ctx)))
    idx2 = round(Int, 2.6*SHA.blocklen(typeof(ctx)))

    # Feed data in according to our dastardly blocking scheme
    SHA.update!(ctx, so_many_as_array[0      + 1:1*idx0])
    SHA.update!(ctx, so_many_as_array[1*idx0 + 1:2*idx0])
    SHA.update!(ctx, so_many_as_array[2*idx0 + 1:3*idx0])
    SHA.update!(ctx, so_many_as_array[3*idx0 + 1:4*idx0])
    SHA.update!(ctx, so_many_as_array[4*idx0 + 1:idx1])
    SHA.update!(ctx, so_many_as_array[idx1 + 1:idx2])
    SHA.update!(ctx, so_many_as_array[idx2 + 1:end])

    # Ensure the hash is the appropriate one
    hash = bytes2hex(SHA.digest!(ctx))
    if hash != answers[sha_funcs[sha_idx]][end]
        print("\n")
        @warn(
        """
        For $(describe_hash(sha_types[sha_funcs[sha_idx]])) expected:
            $(answers[sha_funcs[sha_idx]][end-1])
        Calculated:
            $(hash)
        """)
        nerrors += 1
    else
        VERBOSE && print(".")
    end
end
VERBOSE && println("Done! [$(nerrors - nerrors_old) errors]")

# test hmac correctness using the examples on [wiki](https://en.wikipedia.org/wiki/Hash-based_message_authentication_code#Examples)
VERBOSE && print("Testing on the hmac functions")
nerrors_old = nerrors
for (key, msg, fun, hash) in (
    (b"", b"", hmac_sha1, "fbdb1d1b18aa6c08324b7d64b71fb76370690e1d"),
    (b"", b"", hmac_sha256, "b613679a0814d9ec772f95d778c35fc5ff1697c493715653c6c712144292c5ad"),
    (b"key", b"The quick brown fox jumps over the lazy dog", hmac_sha1, "de7c9b85b8b78aa6bc8a7a36f70a90701c9db4d9"),
    (b"key", b"The quick brown fox jumps over the lazy dog", hmac_sha256, "f7bc83f430538424b13298e6aa6fb143ef4d59a14946175997479dbc2d1a3cd8"),
)
    global nerrors
    digest1 = bytes2hex(fun(Vector(key), Vector(msg)))
    digest2 = bytes2hex(fun(Vector(key), IOBuffer(Vector(msg))))
    if digest1 != hash || digest2 != hash
        print("\n")
        @warn(
        """
        For $fun($(String(key)), $(String(msg))) expected:
            $hash
        Calculated:
            $digest1
            $digest2
        """)
        nerrors += 1
    else
        VERBOSE && print(".")
    end
end
VERBOSE && println("Done! [$(nerrors - nerrors_old) errors]")

replstr(x) = sprint((io, x) -> show(IOContext(io, :limit => true), MIME("text/plain"), x), x)

for idx in 1:length(ctxs)
    global nerrors
    # Part #1: copy
    VERBOSE && print("Testing copy function @ $(ctxs[idx]) ...")
    try
        copy(ctxs[idx]())
    catch
        print("\n")
        @warn("Some weird copy error happened with $(ctxs[idx])")
        nerrors += 1
    end
    VERBOSE && println("Done! [$(nerrors - nerrors_old) errors]")

    # Part #2: show
    VERBOSE && print("Testing show function @ $(ctxs[idx]) ...")
    if replstr(ctxs[idx]()) != shws[idx]
        print("\n")
        @warn("Some weird show error happened with $(ctxs[idx])")
        nerrors += 1
    end
    VERBOSE && println("Done! [$(nerrors - nerrors_old) errors]")
end

# test error if eltype of input is not UInt8
for f in sha_funcs
    global nerrors
    local data = UInt32[0x23467, 0x324775]
    try
        f(data)
    catch ex
        if ex isa MethodError &&
            ex.f === f &&
            ex.args === (data,)
            continue
        end
        rethrow()
    end
    @warn("Non-UInt8 Arrays should fail")
    nerrors += 1
end

# Clean up the I/O mess
rm(file)
rm(tempdir)

if nerrors == 0
    VERBOSE && println("ALL OK")
else
    println("Failed with $nerrors failures")
end
@test nerrors == 0


module TestCaretWrapper

include(joinpath("..", "fixture_learners.jl"))
using .FixtureLearners
nfcp = NumericFeatureClassification()

using FactCheck

using MLBase
importall Orchestra.Types
importall Orchestra.Transformers.CaretWrapper
CW = CaretWrapper
using PyCall
@pyimport rpy2.robjects as RO
@pyimport rpy2.robjects.packages as RP
@pyimport rpy2.robjects.numpy2ri as N2R
N2R.activate()
RP.importr("caret")

function behavior_check(caret_learner::String, impl_options=Dict())
  # Predict with Orchestra learner
  srand(1)
  pycall(RO.r["set.seed"], PyObject, 1)
  learner = CRTLearner({
    :learner => caret_learner,
    :impl_options => impl_options
  })
  orchestra_predictions = fit_and_transform!(learner, nfcp)

  # Predict with backend learner
  srand(1)
  pycall(RO.r["set.seed"], PyObject, 1)
  (r_dataset_df, label_factors) = CW.dataset_to_r_dataframe(
    nfcp.train_instances, nfcp.train_labels
  )
  label_factors = collect(label_factors)
  caret_formula = RO.Formula("Y ~ .")
  r_fit_control = pycall(RO.r[:trainControl], PyObject,
    method = "none"
  )
  if isempty(impl_options)
    r_model = pycall(RO.r[:train], PyObject,
      caret_formula,
      method = caret_learner,
      data = r_dataset_df,
      trControl = r_fit_control,
      tuneLength = 1
    )
  else
    r_model = pycall(RO.r[:train], PyObject,
      caret_formula,
      method = caret_learner,
      data = r_dataset_df,
      trControl = r_fit_control,
      tuneGrid = RO.DataFrame(impl_options)
    )
  end
  (r_instance_df, _) = CW.dataset_to_r_dataframe(nfcp.test_instances)
  original_predictions = collect(pycall(RO.r[:predict], PyObject,
    r_model,
    newdata = r_instance_df
  ))
  original_predictions = map(x -> label_factors[x], original_predictions)

  # Verify same predictions
  @fact orchestra_predictions => original_predictions
end

facts("CARET learners") do
  context("CRTLearner gives same results as its backend") do
    caret_learners = ["svmLinear", "nnet", "earth"]
    for caret_learner in caret_learners
      behavior_check(caret_learner)
    end
  end
  context("CRTLearner with options gives same results as its backend") do
    behavior_check("svmLinear", {:C => 5.0})
  end

  context("CRTLearner throws on incompatible feature") do
    instances = {
      1 "a";
      2 3;
    }
    labels = [
      "a";
      "b";
    ]

    learner = CRTLearner()
    @fact_throws fit!(learner, instances, labels)
  end
end

end # module


using Oceananigans
using Oceananigans.Units
using OceanTurbulenceParameterEstimation
using OceanTurbulenceParameterEstimation: Transformation
using LinearAlgebra, CairoMakie, DataDeps, Distributions

using ElectronDisplay

using Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivities:
    CATKEVerticalDiffusivity,
    SurfaceTKEFlux,
    MixingLength

#####
##### Compile LESbrary
#####

case_path(case) = @datadep_str("four_day_suite_1m/$(case)_instantaneous_statistics.jld2")

Δz = 2
#times = [96hours - 20minutes, 96hours]
times = [12hours, 96hours]
field_names = (:b, :e, :u, :v)
regrid_size = (1, 1, Int(256/Δz))

k₁ = floor(Int, 120 / Δz)
k₂ = ceil(Int, 150 / Δz)
space = SpaceIndices(z=k₁:k₂)
mult = length(space.z) / regrid_size[3]

transformation = (
                  b = Transformation(; normalization=ZScore()),
                  u = Transformation(; normalization=ZScore()),
                  v = Transformation(; normalization=ZScore()),
                  e = Transformation(; space, normalization=RescaledZScore(mult)),
                 )

observation_library = Dict()

# Don't optimize u, v for free_convection
observation_library["free_convection"] =
    SyntheticObservations(case_path("free_convection"); transformation, times, regrid_size,
                          field_names = (:b, :e))

# Don't optimize v for non-rotating cases
observation_library["strong_wind_no_rotation"] =
    SyntheticObservations(case_path("strong_wind_no_rotation"); transformation, times, regrid_size,
                          field_names = (:b, :e, :u))

# The rest are standard
for case in ["strong_wind", "strong_wind_weak_cooling", "weak_wind_strong_cooling"]
    observation_library[case] = SyntheticObservations(case_path(case); field_names, transformation, times, regrid_size)
end

@show znodes(Center, observation_library["free_convection"].grid)[space.z]

cases = [
         "free_convection",
         "weak_wind_strong_cooling",
         "strong_wind_weak_cooling",
         "strong_wind",
         "strong_wind_no_rotation",
        ]

observations = [observation_library[case] for case in cases]

#####
##### Simulation
#####

# Constant Ri, no convection
mixing_length = MixingLength(Cᴬu   = 0.0,
                             Cᴬc   = 0.0,
                             Cᴬe   = 0.0,
                             Cᴸᵇ   = 1.36,
                             Cᴷu⁻  = 0.101,
                             Cᴷc⁻  = 0.0574,
                             Cᴷe⁻  = 3.32,
                             Cᵟu   = 0.5,
                             Cᵟc   = 0.5,
                             Cᵟe   = 0.5,
                             CᴷRiᶜ = 2.0,
                             Cᴷuʳ  = 0.0,
                             Cᴷcʳ  = 0.0,
                             Cᴷeʳ  = 0.0)

surface_TKE_flux = SurfaceTKEFlux(CᵂwΔ=4.74, Cᵂu★=2.76)
catke = CATKEVerticalDiffusivity(; Cᴰ=1.78, mixing_length)

#####
##### Calibration
#####

mass = 0.6
prior_library = Dict()
prior_library[:CᵂwΔ]  = ScaledLogitNormal(; bounds=(4, 10)) #, interval=(2, 5), mass)
prior_library[:Cᵂu★]  = ScaledLogitNormal(; bounds=(2, 10)) #, interval=(3, 5), mass)
prior_library[:Cᴸᵇ]   = ScaledLogitNormal(; bounds=(0, 0.5)) #, interval=(0.1, 2), mass)
prior_library[:Cᴰ]    = ScaledLogitNormal(; bounds=(0, 1)) #, interval=(0.5, 2), mass)

prior_library[:Cᴷu⁻]  = ScaledLogitNormal(; bounds=(0, 0.1)) #, interval=(0.01, 0.1), mass)
prior_library[:Cᴷc⁻]  = ScaledLogitNormal(; bounds=(0, 1)) #, interval=(0.5, 1.0), mass)
prior_library[:Cᴷe⁻]  = ScaledLogitNormal(; bounds=(0, 3)) #, interval=(1.5, 3), mass)

prior_library[:Cᴷuʳ]  = ScaledLogitNormal(; bounds=(-1, 1))
prior_library[:Cᴷcʳ]  = ScaledLogitNormal(; bounds=(-1, 1))
prior_library[:Cᴷeʳ]  = ScaledLogitNormal(; bounds=(-1, 1))

prior_library[:CᴷRiᶜ] = ScaledLogitNormal(; bounds=(1, 3))
prior_library[:CᴷRiʷ] = ScaledLogitNormal(; bounds=(0, 0.5))

prior_library[:Cᴬu]   = ScaledLogitNormal(; bounds=(0, 0.1))
prior_library[:Cᴬc]   = ScaledLogitNormal(; bounds=(0, 2))
prior_library[:Cᴬe]   = ScaledLogitNormal(; bounds=(0, 0.1)) #, interval=(1, 3), mass)

prior_library[:Cᵟu]  = ScaledLogitNormal(; bounds=(0, 10))
prior_library[:Cᵟc]  = ScaledLogitNormal(; bounds=(0, 10))
prior_library[:Cᵟe]  = ScaledLogitNormal(; bounds=(0, 10))

# No convective adjustment:
constant_Ri_parameters = (:Cᴰ, :CᵂwΔ, :Cᵂu★, :Cᴸᵇ, :Cᴷu⁻, :Cᴷc⁻, :Cᴷe⁻, :Cᵟu, :Cᵟc, :Cᵟe)
variable_Ri_parameters = (:Cᴷuʳ, :Cᴷcʳ, :Cᴷeʳ, :CᴷRiʷ, :CᴷRiᶜ, :Cᴰ, :Cᴸᵇ, :CᵂwΔ, :Cᵂu★)
convective_adjustment_parameters = (:Cᴬc, :Cᴬe) # Cᴬu

# :Cᴸˢ
parameter_names = (:CᵂwΔ, :Cᵂu★, :Cᴷe⁻, :Cᴸᵇ, :Cᴰ, :Cᴷc⁻, :Cᴷu⁻, :Cᴷuʳ, :Cᴷcʳ, :Cᴷeʳ, :CᴷRiᶜ, :CᴷRiʷ, :Cᴬc, :Cᴬe)
free_parameters = FreeParameters(prior_library, names=parameter_names)

Nensemble = 4000
Δt = 20.0

function build_simulation()
    simulation = ensemble_column_model_simulation(observations;
                                                  Nensemble,
                                                  architecture = GPU(),
                                                  tracers = (:b, :e),
                                                  closure = catke)

    simulation.Δt = Δt

    Qᵘ = simulation.model.velocities.u.boundary_conditions.top.condition
    Qᵇ = simulation.model.tracers.b.boundary_conditions.top.condition
    N² = simulation.model.tracers.b.boundary_conditions.bottom.condition

    for (case, obs) in enumerate(observations)
        f = obs.metadata.parameters.coriolis_parameter
        view(Qᵘ, :, case) .= obs.metadata.parameters.momentum_flux
        view(Qᵇ, :, case) .= obs.metadata.parameters.buoyancy_flux
        view(N², :, case) .= obs.metadata.parameters.N²_deep
        view(simulation.model.coriolis, :, case) .= Ref(FPlane(f=f))
    end

    return simulation
end

simulation = build_simulation()
calibration = InverseProblem(observations, simulation, free_parameters)
resampler = Resampler(resample_failure_fraction=0.5, acceptable_failure_fraction=1.0)
eki = EnsembleKalmanInversion(calibration; resampler, convergence_rate=0.8)

#####
##### Plot utils
#####

Nt = length(times)
observed_data = []

for observation in observations
    time_serieses = observation.field_time_serieses
    names = keys(time_serieses)
    case_data = NamedTuple(n => interior(getproperty(time_serieses, n)[Nt])[1, 1, :] for n in names)
    push!(observed_data, case_data)
end

function get_modeled_case(icase, name, k=1)
    model_time_serieses = calibration.time_series_collector.field_time_serieses
    field = getproperty(model_time_serieses, name)[Nt]
    return view(interior(field), k, icase, :)
end

colorcycle =  [:black, :royalblue1, :darkgreen, :lightsalmon, :seagreen, :magenta2, :red4, :khaki1,   :darkgreen, :bisque4,
               :silver, :lightsalmon, :lightseagreen, :teal, :royalblue1, :darkorchid4]

markercycle = [:rect, :utriangle, :star5, :circle, :cross, :+, :pentagon, :ltriangle, :airplane, :diamond, :star4]
markercycle = repeat(markercycle, inner=2)

function make_axes(fig, row=1, label=nothing)
    ax_b = Axis(fig[row, 1], xlabel = "Buoyancy \n[cm s⁻²]", ylabel = "z [m]")
    ax_u = Axis(fig[row, 2], xlabel = "x-velocity \n[cm s⁻¹]")
    ax_v = Axis(fig[row, 3], xlabel = "y-velocity \n[cm s⁻¹]")
    ax_e = Axis(fig[row, 4], xlabel = "Turbulent kinetic energy \n[cm² s⁻²]")
    if !isnothing(label)
        ax_t = Axis(fig[row, 5])
        xlims!(0, 1)
        ylims!(0, 1)
        hidespines!(ax_t)
        hidedecorations!(ax_t)
        text!(ax_t, label, justification=:left, align=(:left, :center), position=(0, 0.5))
    end
    return (ax_b, ax_u, ax_v, ax_e)
end

function plot_fields!(axs, label, color, b, e, u=zeros(size(b)), v=zeros(size(b)); linewidth=2, linestyle=:solid)
    grid = first(values(observation_library)).grid
    z = znodes(Center, grid)
    b, u, v, e = Tuple(Array(f) for f in (b, u, v, e))

    for (q, name) in zip((b, u, v, e), ("b", "u", "v", "e"))
        any(isnan.(q)) && @warn("NaNs found in $label $(name)!")
    end

    ## Note unit conversions below, eg m s⁻¹ -> cm s⁻¹:cyan
    lines!(axs[1], 1e2 * b, z; color, linestyle, label, linewidth)
    lines!(axs[2], 1e2 * u, z; color, linestyle, label, linewidth)
    lines!(axs[3], 1e2 * v, z; color, linestyle, label, linewidth)
    lines!(axs[4], 1e4 * e, z; color, linestyle, label, linewidth)

    return nothing
end

function min_max_parameters(summary)
    names = keys(summary.ensemble_mean)
    Nens = length(summary.parameters)
    parameter_matrix = [summary.parameters[k][name] for name in names, k = 1:Nens]
    θ_min = minimum(parameter_matrix, dims=2)
    θ_max = maximum(parameter_matrix, dims=2)
    return θ_min, θ_max
end

function finitefind(a, val, find)
    b = deepcopy(a)
    b[.!isfinite.(a)] .= val
    return find(b)
end

finitefindmin(a) = finitefind(a, Inf, findmin)
finitefindmax(a) = finitefind(a, -Inf, findmax)

function visualize_parameter_evolution(eki)
    summaries = eki.iteration_summaries
    Niters = length(summaries)
    names = eki.inverse_problem.free_parameters.names
    θ_mean = NamedTuple(name => map(s -> s.ensemble_mean[name], summaries) for name in names)

    k_best(s) = finitefindmin(s.mean_square_errors)[2]
    θ_best = NamedTuple(name => map(s -> s.parameters[k_best(s)][name], summaries) for name in names)

    θ_min_max = [min_max_parameters(s) for s in summaries]
    θ_min = [[θn[1][i] for θn in θ_min_max] for i in 1:length(names)]
    θ_max = [[θn[2][i] for θn in θ_min_max] for i in 1:length(names)]

    θᵢ = NamedTuple(name => first(θ_mean[name]) for name in names)
    Δθ = NamedTuple(name => (θ_mean[name] .- θᵢ[name]) ./ θᵢ[name] for name in names)
    iterations = 0:length(summaries)-1

    fig = Figure(resolution=(1200, 1200))
    ax1 = Axis(fig[1:3, 1], xlabel = "Iteration", ylabel = "Δθ")
    for (i, name) in enumerate(names)
        label = string(name)
        marker = markercycle[i]
        color = colorcycle[i]
        scatterlines!(ax1, iterations, parent(Δθ[name]); marker, color=(color, 0.8), label, linewidth=4)
    end

    fig[1:3, 2] = Legend(fig, ax1)

    Nparts = 3
    Nθpart = floor(Int, length(names) / Nparts)

    for p in 1:Nparts
        axp = Axis(fig[p+3, 1], xlabel = "Iteration", ylabel = "θ")

        if p == Nparts
            np = UnitRange((p-1) * Nθpart + 1, length(names))
        else
            np = UnitRange((p-1) * Nθpart + 1, p * Nθpart)
        end

        partnames = names[np]

        for (n, name) in enumerate(partnames)
            i = (p - 1) * Nθpart + n
            label = string(name)
            marker = markercycle[i]
            color = colorcycle[i]
            scatterlines!(axp, iterations, parent(θ_mean[name]); marker, color=(color, 0.6), label, linewidth=4)
            lines!(axp, iterations, parent(θ_best[name]); color, linewidth=2)
            band!(axp, iterations, parent(θ_min[i]), parent(θ_max[i]), color=(color, 0.3))
        end

        fig[p+3, 2] = Legend(fig, axp)
    end

    display(fig)

    return nothing
end

function plot_latest(eki)
    latest_summary = eki.iteration_summaries[end]
    min_error, k_min = finitefindmin(latest_summary.mean_square_errors)
    max_error, k_max = finitefindmax(latest_summary.mean_square_errors)

    fig = Figure(resolution=(1200, 1200))

    for (c, case) in enumerate(cases)
        label = replace(case, "_" => "\n")
        axs = make_axes(fig, c, label)
        observed = observed_data[c]
        obs = observations[c]

        min_error_data = NamedTuple(n => get_modeled_case(c, n, k_min) for n in keys(obs.field_time_serieses))
        max_error_data = NamedTuple(n => get_modeled_case(c, n, k_max) for n in keys(obs.field_time_serieses))

        plot_fields!(axs, "observed at t = " * prettytime(times[end]), (:gray23, 0.6), observed...; linewidth=4)
        plot_fields!(axs, "min", :navy, min_error_data...)
        plot_fields!(axs, "max", :orangered3, max_error_data...)

        fig[1, 6] = Legend(fig, axs[1])
    end

    display(fig)

    return nothing
end

#####
##### Visualize observations
#####

fig = Figure(resolution=(1200, 1200))
linestyles = [:solid, :dash, :dot, :dashdot, :dashdotdot]
all_axs = []

for (o, observation) in enumerate(observations)
    axs_label = replace(cases[o], "_" => "\n")
    axs = make_axes(fig, o, axs_label)
    append!(all_axs, axs)
    for (n, t) in enumerate(times)
        linestyle = linestyles[o]
        label = "t = " * prettytime(t)
        names = keys(observation.field_time_serieses)
        data = map(name -> interior(observation.field_time_serieses[name][n])[1, 1, :], names)
        plot_fields!(axs, label, colorcycle[n], data...)
    end
end

display(fig)

#####
##### Calibrate
#####

# Initial state after 0 iterations
plot_latest(eki)

# Continuously update
for i = 1:20
    @info "Iterating..."
    start_time = time_ns()
    iterate!(eki)
    elapsed = 1e-9 * (time_ns() - start_time)
    @info string("   done. (", prettytime(elapsed), ")")
    @show eki.iteration_summaries[end]
    visualize_parameter_evolution(eki)
    plot_latest(eki)
end


"Parameters for right-hand side function of Variational Partitioned Runge-Kutta methods."
const ParametersVPRKpVariational = AbstractParametersVPRK{:vprk_pvariational}

"Variational partitioned Runge-Kutta integrator."
struct IntegratorVPRKpVariational{DT, TT, D, S,
                PT  <: ParametersVPRK{DT,TT},
                PPT <: ParametersVPRKpVariational{DT,TT},
                ST  <: NonlinearSolver{DT},
                PST <: NonlinearSolver{DT},
                IT <: InitialGuessIODE{TT}} <: AbstractIntegratorVPRKwProjection{DT,TT,D,S}

    params::PT
    pparams::PPT
    solver::ST
    projector::PST
    iguess::IT
    caches::CacheDict{PPT}

    function IntegratorVPRKpVariational(params::ParametersVPRK{DT,TT,D,S},
                    pparams::ParametersVPRKpVariational{DT,TT,D,S},
                    solver::ST, projector::PST, iguess::IT, caches) where {DT,TT,D,S,ST,PST,IT}
        new{DT, TT, D, S, typeof(params), typeof(pparams), ST, PST, IT}(params, pparams, solver, projector, iguess, caches)
    end

    function IntegratorVPRKpVariational{DT,D}(equations::NamedTuple, tableau::TableauVPRK{TT}, Δt::TT; R=[1,1]) where {DT, TT, D}
        # get number of stages
        S = tableau.s

        # create params
        params = ParametersVPRK{DT,D}(equations, tableau, Δt)

        # create projector params
        R  = convert(Vector{TT}, R)
        R1 = [R[1], zero(TT)]
        R2 = [zero(TT), R[2]]
        pparams = ParametersVPRKpVariational{DT,D}(equations, tableau, Δt, NamedTuple{(:R,:R1,:R2)}((R,R1,R2)))

        # create cache dict
        caches = CacheDict(pparams)

        # create nonlinear solver
        solver = create_nonlinear_solver(DT, D*S, params, caches)

        # create projector
        projector = create_nonlinear_solver(DT, 2*D, pparams, caches)

        # create initial guess
        iguess = InitialGuessIODE(get_config(:ig_extrapolation), equations[:v̄], equations[:f̄], Δt)

        # create integrator
        IntegratorVPRKpVariational(params, pparams, solver, projector, iguess, caches)
    end

    function IntegratorVPRKpVariational(equation::Union{IODE{DT},LODE{DT}}, tableau, Δt; kwargs...) where {DT}
        IntegratorVPRKpVariational{DT, ndims(equation)}(get_functions(equation), tableau, Δt; kwargs...)
    end
end

function Integrators.initialize!(int::IntegratorVPRKpVariational{DT}, sol::AtomicSolutionPODE{DT},
                                 cache::IntegratorCacheVPRK{DT}=int.caches[DT]) where {DT}
    sol.t̄ = sol.t - timestep(int)

    equation(int, :v̄)(sol.t, sol.q, sol.v)
    equation(int, :f̄)(sol.t, sol.q, sol.v, sol.f)

    initialize!(int.iguess, sol.t, sol.q, sol.p, sol.v, sol.f,
                            sol.t̄, sol.q̄, sol.p̄, sol.v̄, sol.f̄)

    # initialise projector
    equation(int, :g)(sol.t, sol.q, cache.λ, cache.G[1])
    # cache.U[1] .= cache.λ
end

function initial_guess!(int::IntegratorVPRKpVariational{DT}, sol::AtomicSolutionPODE{DT},
                        cache::IntegratorCacheVPRK{DT}=int.caches[DT]) where {DT}
    for i in eachstage(int)
        evaluate!(int.iguess, sol.q̄, sol.p̄, sol.v̄, sol.f̄,
                              sol.q, sol.p, sol.v, sol.f,
                              cache.q̃, cache.ṽ,
                              tableau(int).q.c[i])

        for k in eachdim(int)
            int.solver.x[ndims(int)*(i-1)+k] = cache.ṽ[k]
        end
    end
end

function initial_guess_projection!(int::IntegratorVPRKpVariational, sol::AtomicSolutionPODE)
    offset_q = 0
    offset_λ = ndims(int)

    for k in eachdim(int)
        int.projector.x[offset_q+k] = sol.q[k]
        int.projector.x[offset_λ+k] = 0
    end
end

function compute_projection!(
                x::Vector{ST}, q::SolutionVector{ST}, p::SolutionVector{ST}, λ::SolutionVector{ST},
                U::Vector{Vector{ST}}, G::Vector{Vector{ST}},
                params::ParametersVPRKpVariational{DT,TT,D,S}) where {ST,DT,TT,D,S}

    @assert D == length(q) == length(p) == length(λ)
    @assert D == length(U[1]) == length(U[2])
    @assert D == length(G[1]) == length(G[2])

    # copy x to q, λ
    for k in 1:D
        q[k] = x[0*D+k]
        λ[k] = x[1*D+k]
    end

    # compute u=λ and g=∇ϑ(q)⋅λ
    U[1] .= 0
    U[2] .= λ
    # U[1] .= λ
    # U[2] .= 0

    params.equ[:g](params.t̄, q, λ, G[1])
    G[2] .= 0
    # G[1] .= 0
    # params.equ.g(params.t̄, q, λ, G[2])

    # compute p=ϑ(q)
    params.equ[:ϑ](params.t̄, q, λ, p)
end

"Compute stages of projected variational partitioned Runge-Kutta methods."
function Integrators.function_stages!(x::Vector{ST}, b::Vector{ST},
                params::ParametersVPRKpVariational{DT,TT,D,S},
                caches::CacheDict) where {ST,DT,TT,D,S}

    @assert length(x) == length(b)

    # get cache for internal stages
    cache = caches[ST]

    compute_projection!(x, cache.q̃, cache.p̃, cache.λ, cache.U, cache.G, params)

    # # compute b = - [q̄-q-U]
    for k in 1:D
        b[0*D+k] = - (cache.q̃[k] - params.q̄[k]) + params.Δt * params.pparams[:R][2] * cache.U[2][k]
        # b[0*D+k] = - (cache.q[k] - params.q̄[k])
    end

    # compute b = - [p̄-p-G]
    for k in 1:D
        b[1*D+k] = - (cache.p̃[k] - params.p̄[k])
        # b[1*D+k] = - (cache.p[k] - params.p̄[k]) + params.Δt * params.R[2] * cache.G[2][k]
    end
end

function Integrators.integrate_step!(int::IntegratorVPRKpVariational{DT,TT}, sol::AtomicSolutionPODE{DT,TT},
                                     cache::IntegratorCacheVPRK{DT}=int.caches[DT]) where {DT,TT}
    # add perturbation for next time step to solution
    # (same vector field as previous time step)
    project_solution!(int, sol, int.pparams.pparams[:R1], cache)

    # update nonlinear solver parameters from cache
    update_params!(int.params, sol)

    # compute initial guess
    initial_guess!(int, sol, cache)

    # reset solution
    reset!(sol, timestep(int))

    # call nonlinear solver
    solve!(int.solver)

    # print solver status
    print_solver_status(int.solver.status, int.solver.params)

    # check if solution contains NaNs or error bounds are violated
    check_solver_status(int.solver.status, int.solver.params)

    # compute vector fields at internal stages
    compute_stages!(int.solver.x, cache.Q, cache.V, cache.P, cache.F, int.params)

    # compute unprojected solution
    update_solution!(int, sol, cache)

    # set time and solution for projection solver
    update_params!(int.pparams, sol)

    # set initial guess for projection
    initial_guess_projection!(int, sol)

    # call projection solver
    solve!(int.projector)

    # print solver status
    print_solver_status(int.projector.status, int.projector.params)

    # check if solution contains NaNs or error bounds are violated
    check_solver_status(int.projector.status, int.projector.params)

    # compute projection vector fields
    compute_projection!(int.projector.x, cache.q̃, cache.p̃, cache.λ,
                                         cache.U, cache.G, int.pparams)

    # add projection to solution
    project_solution!(int, sol, int.pparams.pparams[:R2], cache)

    # copy solution to initial guess
    update_vector_fields!(int.iguess, sol.t, sol.q, sol.p, sol.v, sol.f)
end


module IncomeTaxCalculations

#
# This module calculates income tax for an individual (and also, at the end a married couple).
# It's based mainly on the descriptions and examples in the excellent Melville's Taxation:
# Melville, Alan. 2019. Melville’s Taxation: Finance Act 2019. 25th ed. London. (and 2020 edition).
#

using Dates
using Dates: Date, now, TimeType, Year
using Parameters: @with_kw

using ScottishTaxBenefitModel

using .Definitions

using .ModelHousehold:
    Person

using .STBParameters:
    IncomeTaxSys

using .GeneralTaxComponents:
    RateBands,
    TaxResult,
    calctaxdue,
    delete_thresholds_up_to
    # , *

using .Utils:
    get_if_set

using .STBIncomes

using .Results:
    BenefitUnitResult,
    IndividualResult,
    ITResult

export
    apply_allowance,
    calc_income_tax!,
    calculate_company_car_charge,
    old_enough_for_mca

## FIXME just use the dict..
function guess_car_percentage_2020_21( sys :: IncomeTaxSys, company_car_fuel_type :: Fuel_Type )
    return sys.company_car_charge_by_CO2_emissions[company_car_fuel_type]
end

function calculate_company_car_charge(
    pers   :: Person,
    sys    :: IncomeTaxSys,
    calculator :: Function = guess_car_percentage_2020_21 ) :: Real
    value = max(0.0, pers.company_car_value-
        pers.company_car_contribution )
    if pers.fuel_supplied > 0.0
        value += sys.fuel_imputation
    end
    prop = calculator( sys, pers.company_car_fuel_type )
    value * prop
end

"""
Very rough approximation to MCA age - ignores all months since we don't have that in a typical dataset
TODO maybe overload this with age as a Date?
"""
function old_enough_for_mca(
    sys            :: IncomeTaxSys,
    age            :: Integer,
    model_run_date :: TimeType = now() ) :: Bool
    (model_run_date - Year(age)) < sys.mca_date
end

function calculate_allowance( pers::Person, sys :: IncomeTaxSys ) :: Real
    allowance = sys.personal_allowance
    if pers.registered_blind
        allowance += sys.blind_persons_allowance
    end
    allowance
end

function apply_allowance( allowance::Real, income::Real )::Tuple where RT<:Real
    r = max( 0.0, income - allowance )
    allowance = max(0.0, allowance-income)
    allowance,r
end

"""
  from Melville, ch13.
  Changes are in itres: add to pension fields and extend bands
  notes:
  Melville talks of "earned income" - using non-savings income
  FIXME: check: does the personal_allowance_income_limit get bumped up?
"""
function calculate_pension_taxation!(
    pres   :: IndividualResult,
    sys    :: IncomeTaxSys,
    pers   ::Person,
    total_income::Real,
    earned_income:: Real )

    pres.it.savings_thresholds = copy( sys.savings_thresholds )
    pres.it.dividend_thresholds = copy( sys.dividend_thresholds )
    pres.it.non_savings_thresholds = copy( sys.non_savings_thresholds )
    # fixme check avs here
    # FIXME not right?
    avc = pres.income[AVCS]
    pen = pres.income[PENSION_CONTRIBUTIONS_EMPLOYEE] +
        pres.income[PENSION_CONTRIBUTIONS_EMPLOYER]
    eligible_contribs = avc + pen
    if eligible_contribs <= 0.0
        return
    end

    max_relief = sys.pension_contrib_annual_allowance
    if total_income < sys.pension_contrib_basic_amount
        max_relief = sys.pension_contrib_basic_amount
    end
    if total_income > sys.pension_contrib_threshold_income
        excess = total_income - sys.pension_contrib_threshold_income
        max_relief = max( sys.pension_contrib_annual_minimum,
            sys.pension_contrib_annual_allowance - excess*sys.pension_contrib_withdrawal_rate )
    end
    # println( "total_income $total_income max_relief $max_relief eligible_contribs $eligible_contribs" )
    eligible_contribs = min( eligible_contribs, max_relief );
    pres.it.pension_eligible_for_relief = eligible_contribs
    basic_rate = sys.non_savings_rates[ sys.non_savings_basic_rate ]
    # println("total_income=$total_income max_relief=$max_relief eligible_contribs=$eligible_contribs basic_rate=$basic_rate sys.pension_contrib_withdrawal_rate=$(sys.pension_contrib_withdrawal_rate)")
    gross_contribs = eligible_contribs/(1-basic_rate)
    pres.it.pension_relief_at_source = gross_contribs - eligible_contribs
    pres.it.non_savings_thresholds .+= gross_contribs
    pres.it.savings_thresholds .+= gross_contribs
    pres.it.dividend_thresholds .+= gross_contribs
end

"""

Complete(??) income tax calculation, based on the Scottish/UK 2019 system.
Mostly taken from Melville (2019) chs 2-4.

FIXME this is too long and needs broken up.

problems:

1. we do this in strict non-savings, savings, dividends order; see 2(11) for examples where it's now advantageous to use a different order
2.

"""
function calc_income_tax!(
    pres   :: IndividualResult,
    pers   :: Person,
    sys    :: IncomeTaxSys,
    spouse_transfer :: Real = 0.0 )
    if sys.abolished
        return;
    end
    total_income = isum( pres.income, sys.all_taxable )
    non_savings_income = isum( pres.income, sys.non_savings_income )

    savings_income = isum( pres.income, sys.savings_income )
    dividends_income = isum( pres.income, sys.dividend_income )
    allowance = calculate_allowance( pers, sys )
    # allowance reductions goes here

    adjusted_net_income = total_income

    calculate_pension_taxation!( pres, sys, pers, total_income, non_savings_income )

    # adjusted_net_income -= pres.it.pension_eligible_for_relief

    adjusted_net_income += calculate_company_car_charge(pers, sys)
    # ...

    non_savings_tax = TaxResult(0.0, 0)
    savings_tax = TaxResult(0.0, 0)
    dividend_tax = TaxResult(0.0, 0)

    if adjusted_net_income > sys.personal_allowance_income_limit
        allowance =
            max(0.0,
                allowance -
                    sys.personal_allowance_withdrawal_rate*(
                        adjusted_net_income - sys.personal_allowance_income_limit ))
    end
    taxable_income = adjusted_net_income-allowance
    # note: we copy from the expanded versions from pension_contributions
    savings_thresholds = deepcopy( pres.it.savings_thresholds )
    savings_rates = deepcopy( sys.savings_rates )
    # FIXME model all this with parameters
    toprate = size( savings_thresholds )[1]
    non_savings_taxable = 0.0
    savings_taxable = 0.0
    dividends_taxable = 0.0

    if taxable_income > 0
        allowance,non_savings_taxable = apply_allowance( allowance, non_savings_income )
        non_savings_tax = calctaxdue(
            taxable=non_savings_taxable,
            rates=sys.non_savings_rates,
            thresholds=pres.it.non_savings_thresholds )

        # horrific savings calculation see Melville Ch2 "Savings Income" & examples 2-3
        # FIXME Move to separate function
        # delete the starting bands up to non_savings taxabke icome
        savings_rates, savings_thresholds = delete_thresholds_up_to(
            rates=savings_rates,
            thresholds=savings_thresholds,
            upto=non_savings_taxable );
        if sys.personal_savings_allowance > 0
            psa = sys.personal_savings_allowance
            # println( "taxable income $taxable_income sys.savings_thresholds[2] $(sys.savings_thresholds[2])")
            if taxable_income > sys.savings_thresholds[toprate]
                psa = 0.0
            elseif taxable_income > sys.savings_thresholds[2] # above the basic rate FIXME parameterise '2' here
                psa *= 0.5 # FIXME parameterise 0.5 here
            end
            if psa > 0.0 ## if we haven't deleted the zero band already, just widen it
                if savings_rates[1] == 0.0
                    savings_thresholds[1] += psa;
                else ## otherwise, insert a  new one.
                    savings_thresholds = vcat([psa], savings_thresholds )
                    savings_rates = vcat([0.0], savings_rates )
                end
            end
            pres.it.personal_savings_allowance = psa
        end # we have a personal_savings_allowance
        pres.it.savings_rates = savings_rates
        pres.it.savings_thresholds= savings_thresholds
        allowance,savings_taxable = apply_allowance( allowance, savings_income )
        savings_tax = calctaxdue(
            taxable=savings_taxable,
            rates=savings_rates,
            thresholds=savings_thresholds )

        # Dividends
        # see around example 8-9 ch2
        allowance,dividends_taxable =
            apply_allowance( allowance, dividends_income )
        dividend_rates=deepcopy(sys.dividend_rates)
        dividend_thresholds=deepcopy(pres.it.dividend_thresholds )
        # always preserve any bottom zero rate
        add_back_zero_band = false
        zero_band = 0.0
        used_thresholds = non_savings_taxable+savings_taxable
        copy_start = 1
        # handle the zero rate
        if dividend_rates[1] == 0.0
            add_back_zero_band = true
            zero_band = dividend_thresholds[1]
            used_thresholds += min( zero_band, dividends_taxable )
            copy_start = 2
        end
        dividend_rates, dividend_thresholds =
            delete_thresholds_up_to(
                rates=dividend_rates[copy_start:end],
                thresholds=dividend_thresholds[copy_start:end],
                upto=used_thresholds );
        if add_back_zero_band
            dividend_rates = vcat( [0.0], dividend_rates )
            dividend_thresholds .+= zero_band # push all up
            dividend_thresholds = vcat( zero_band, dividend_thresholds )
        end
        pres.it.dividend_rates = dividend_rates
        pres.it.dividend_thresholds= dividend_thresholds

        dividend_tax = calctaxdue(
            taxable=dividends_taxable,
            rates=dividend_rates,
            thresholds=dividend_thresholds )
    else # some allowance left
        allowance = -taxable_income # e.g. allowance - taxable_income
    end

    #
    # tax reducers
    #
    @assert non_savings_tax.due >= 0
    @assert savings_tax.due >= 0
    @assert dividend_tax.due >= 0
    total_tax = non_savings_tax.due+savings_tax.due+dividend_tax.due
    nst = non_savings_tax.due
    st = savings_tax.due
    dt = dividend_tax.due
    if spouse_transfer > 0
        # note that if scotland sets a high basic rate the sp_reduction here
        # can wipe out all savings taxes which should be remitted to uk gov
        sp_reduction =
            sys.non_savings_rates[sys.non_savings_basic_rate]*spouse_transfer
        total_tax = max( 0.0, total_tax - sp_reduction )
        # assign spouse reduction to non-savings - needed for Scottish income tax
        nst -= sp_reduction
        # println( "nst $nst sp_reduction=$sp_reduction total_tax=$total_tax")
        if nst < 0.0
            # assign any leftover to savings tax
            st += nst # subtract, really
            nst = 0.0
        end
        if st < 0.0
            # assign any leftover to savings tax
            dt = max(0.0, dt+st) # subtract, really
            st = 0.0
        end
    end

    # total_tax -= pres.it.pension_relief_at_source
    pres.income[INCOME_TAX] = total_tax

    pres.it.taxable_income = taxable_income
    pres.it.allowance = allowance
    pres.it.total_income = total_income
    pres.it.adjusted_net_income = adjusted_net_income

    # memo items, but we need nst for Scottish Income Tax
    pres.it.non_savings_tax = nst # non_savings_tax.due, less transferrable allowances
    pres.it.non_savings_income = non_savings_income
    pres.it.non_savings_band = non_savings_tax.end_band
    pres.it.non_savings_taxable = non_savings_taxable

    pres.it.savings_tax = st
    pres.it.savings_band = savings_tax.end_band
    pres.it.savings_income = savings_income
    pres.it.savings_taxable = savings_taxable

    pres.it.dividends_tax = dt
    pres.it.dividend_band = dividend_tax.end_band
    pres.it.dividends_income = dividends_income
    pres.it.dividends_taxable = dividends_taxable

    pres.it.unused_allowance = allowance
end

function allowed_to_transfer_allowance(
    sys  :: IncomeTaxSys;
    from :: ITResult,
    to   :: ITResult ) :: Bool

    can_transfer :: Bool = true
    if ! (from.unused_allowance > 0.0 &&
        to.unused_allowance <= 0.0)
         # nothing to transfer - this is actually wrong since
         # you can opt to transfer some allowance even if you
         # can technically use it.
        can_transfer = false
    elseif to.savings_band > sys.savings_basic_rate ||
        to.non_savings_band > sys.non_savings_basic_rate ||
        to.dividend_band > sys.dividend_basic_rate
        can_transfer = false
    end
    ## TODO disallow if mca claimed
    return can_transfer
end # can_transfer

function calculate_mca(
    pers :: Person,
    tax :: ITResult,
    sys :: IncomeTaxSys)::Real
    ## FIXME parameterise this
    mca = sys.married_couples_allowance
    if tax.adjusted_net_income > sys.mca_income_maximum
        mca = max( sys.mca_minimum, mca -
           (tax.adjusted_net_income-sys.mca_income_maximum)*sys.mca_withdrawal_rate)
    end
    mca * sys.mca_credit_rate
end

"""
 FIXME maybe send an actual benefit unit in here? With kids..
"""
function calc_income_tax!(
    bres   :: BenefitUnitResult,
    head   :: Person,
    spouse :: Union{Nothing, Person},
    sys    :: IncomeTaxSys )
    hdres = bres.pers[head.pid]
    calc_income_tax!( hdres, head, sys )
    # FIXME the transferable stuff here
    # is not quite right as you can elect to transfer more than
    # the surplus allowance in some cases.
    # also - add in restrictions on transferring to
    # higher rate payers.
    if spouse !== nothing
        spres = bres.pers[spouse.pid]
        calc_income_tax!( spres, spouse, sys )
        # This is not quite right - you can't claim the
        # MCA AND transfer an allowance. We're assuming
        # always MCA first (I think it's always more valuable?)
        if old_enough_for_mca( sys, head.age ) || old_enough_for_mca( sys, spouse.age )
            # shoud usually just go to the head but.. some stuff about partner
            # with greater income if married after 2005 and you can elect to do this if
            # married before, so:
            if hdres.it.adjusted_net_income > spres.it.adjusted_net_income
                hdres.it.mca = calculate_mca( head, hdres.it, sys )
                hdres.income[INCOME_TAX] = max( 0.0, hdres.income[INCOME_TAX]- hdres.it.mca )
                # FIXME not quite right - trickle down to savings/divs
                if hdres.it.mca > 0
                    hdres.it.non_savings_tax = max( 0.0, hdres.it.non_savings_tax - hdres.it.mca )
                end
            else
                spres.it.mca = calculate_mca( spouse, spres.it, sys )
                spres.income[INCOME_TAX] = max( 0.0, spres.income[INCOME_TAX] - spres.it.mca )
                if spres.it.mca > 0
                    spres.it.non_savings_tax = max( 0.0, spres.it.non_savings_tax - spres.it.mca )
                end
            end
        end
        if spres.it.mca == 0.0 == hdres.it.mca
            if allowed_to_transfer_allowance( sys, from=spres.it, to=hdres.it )
                transferable_allow = min( spres.it.unused_allowance, sys.marriage_allowance )
                calc_income_tax!( hdres, head, sys, transferable_allow )
                hdres.it.transferred_allowance = transferable_allow
            elseif allowed_to_transfer_allowance( sys, from=hdres.it, to=spres.it )
                transferable_allow = min( hdres.it.unused_allowance, sys.marriage_allowance )
                calc_income_tax!( spres, spouse, sys, transferable_allow )
                spres.it.transferred_allowance = transferable_allow
            end
        end
        if spres.income[INCOME_TAX] < 0.0
            println( spouse.income )
            println( spres.it )
            println( spres.income )
            @assert false "income tax can't be negative; spouse; pid=$(spouse.pid)"
        end
        if spres.it.non_savings_tax > spres.income[INCOME_TAX]
            println( spouse.income )
            println( spres.it )
            println( spres.income )
            @assert false "non savings part can't be greater than total income tax; spouse ; pid=$(spouse.pid)"
        end
    end
    if hdres.income[INCOME_TAX] < 0.0
        println( head.income )
        println( hdres.it )
        println( hdres.income )

        @assert false "income tax can't be negative; head; pid=$(head.pid)"
    end
    if hdres.it.non_savings_tax > hdres.income[INCOME_TAX]
        println( head.income )
        println( hdres.it )
        println( hdres.income )
        @assert false "non savings part can't be greater than total income tax; head ; pid=($head.pid)"
    end
end # calc_income_tax

end # module


using SubHunt
using Test
using Random
using POMDPs
using POMDPPolicies
using POMDPSimulators
using DiscreteValueIteration
using ParticleFilters
using POMDPModelTools
using QMDP

@testset "VI" begin
    rng = MersenneTwister(6)
    pomdp = SubHuntPOMDP()
    # show(STDOUT, MIME("text/plain"), SubVis(pomdp))

    solver = ValueIterationSolver(verbose=true)
    vipolicy = solve(solver, UnderlyingMDP(pomdp))
    #
    s = rand(rng, initialstate(pomdp))
    @show value(vipolicy, s)
    @show value(vipolicy, SubState(s.own, s.target, s.goal, true))
end

@testset "QMDP and PF" begin
    rng = MersenneTwister(6)
    pomdp = SubHuntPOMDP()
    policy = solve(QMDPSolver(verbose=true), pomdp)

    # policy = RandomPolicy(pomdp, rng=rng)

    filter = BootstrapFilter(pomdp, 10000, rng)

    for (s, a, r, sp) in stepthrough(pomdp, policy, filter, "s,a,r,sp", max_steps=200, rng=rng)
        v = SubVis(pomdp, s=s, a=a, r=r)
        show(stdout, MIME("text/plain"), v)
    end
end

@testset "DPOMDP and PF" begin
    rng = MersenneTwister(6)
    dpomdp = DSubHuntPOMDP(SubHuntPOMDP(), 1.0)
    policy = RandomPolicy(dpomdp, rng=rng)
    filter = BootstrapFilter(dpomdp, 10000, rng)

    for (s, a, r, sp) in stepthrough(dpomdp, policy, filter, "s,a,r,sp", max_steps=200, rng=rng)
        v = SubVis(dpomdp.cp, s=s, a=a, r=r)
        show(stdout, MIME("text/plain"), v)
    end
end

@testset "Visualization" begin
    rng = MersenneTwister(6)
    pomdp = SubHuntPOMDP()
    policy = solve(QMDPSolver(verbose=true), pomdp)

    # policy = RandomPolicy(pomdp, rng=rng)

    filter = BootstrapFilter(pomdp, 10000, rng)

    for step in stepthrough(pomdp, policy, filter, "s,a,r,sp", max_steps=200, rng=rng)
        show(stdout, MIME("text/plain"), render(pomdp, step))
    end
end


"""
Notes:
 - `rand` will store randomness only when trace type matches TraceR.
 - `randc` never stores randomness.
 - `randr` will store and replay randomness regardless trace type (N.B. Particle Gibbs uses `randr`).
 - `fork1` will perform replaying immediately and fix the particle weight to 1.
 - `fork2` will perform lazy replaying and accumulate likelihoods like a normal particle.
"""

module Traces
using Distributions

include("taskcopy.jl")
include("tarray.jl")

export Trace, TraceR, TraceC, current_trace, fork, fork2, randr, TArray, tzeros

type Trace{T}
  task :: Task
  randomness :: Array{Any, 1}    # elem t is the randomness created by the t’th assume call.
  index :: Int                 # index of current randomness
  num_produce :: Int           # num of produce calls from trace, each produce corresponds to an observe.
  Trace() = (res = new(); res.randomness = Array{Any,1}(); res.index = 0; res.num_produce = 0; res)
end

function call{T}(::Type{Trace{T}}, f::Function)
  res = Trace{T}();
  res.task = Task(()->f());
  if isa(res.task.storage, Void)
    res.task.storage = ObjectIdDict()
  end
  res.task.storage[:turing_trace] = res # create a backward reference in task_local_storage
  res
end

typealias TraceR Trace{:R} # Task Copy
typealias TraceC Trace{:C} # Replay

# generate a new random variable, replay if t.counter < length(t.randomness)
function randr( t::Trace, distr :: Distribution )
  t.index += 1
  if t.index < length(t.randomness)
    res = t.randomness[t.index]
  else # sample, record
    res = Distributions.rand(distr)
    push!(t.randomness, res)
  end
  return res
end

# generate a new random variable, no replay
randc(t::Trace, distr :: Distribution) = Distributions.rand(distr)

Distributions.rand(t::TraceR, distr :: Distribution) = randr(t, distr)
Distributions.rand(t::TraceC, distr :: Distribution) = randc(t, distr)

# step to the next observe statement, return log likelihood
Base.consume(t::Trace) = (t.num_produce += 1; Base.consume(t.task))

# Task copying version of fork for both TraceR and TraceC.
function forkc(trace :: Trace)
  newtrace = typeof(trace)()
  newtrace.task = Base.copy(trace.task)
  n_rand = min(trace.index, length(trace.randomness))
  newtrace.randomness = trace.randomness[1:n_rand]
  newtrace.index = trace.index
  newtrace.num_produce = trace.num_produce
  newtrace.task.storage[:turing_trace] = newtrace
  newtrace
end

# fork s and replay until observation t; drop randomness between y_t:T if keep == false
#  N.B.: PG requires keeping all randomness even we only replay up to observation y_t
function forkr(trace :: TraceR, t :: Int, keep :: Bool)
  # Step 0: create new task and copy randomness
  newtrace = TraceR(trace.task.code)
  newtrace.randomness = deepcopy(trace.randomness)

  # Step 1: Call consume t times to replay randomness
  map((i)->consume(newtrace), 1:t)

  # Step 2: Remove remaining randomness if keep==false
  if !keep
    newtrace.randomness = newtrace.randomness[1:newtrace.index]
  end
  newtrace
end

# Default fork implementation, replay immediately.
fork(s :: TraceR) = forkr(s, s.num_produce, false)
fork(s :: TraceC) = forkc(s)

# Lazily replay on demand, note that:
#  - lazy replay is only possible for TraceR
#  - lazy replay accumulates likelihoods
#  - lazy replay is useful for implementing PG (i.e. ref particle)
fork2(s :: TraceR) = forkr(s, 0, true)

current_trace() = current_task().storage[:turing_trace]

end


"""
`Basics.estimate()`
    ... estimates various quantities, typically, due to semi-empirical formulas and parametrizations.

  + `("ionization potential: inner-shell", shell::Shell, Z::Int64)`
    ... to estimate the ionization potential = mean binding energy of an electron in the given shell; an energy value::Float64
        is returned. These ionization potentials are taken from the tabulation ....
"""
function Basics.estimate(sa::String, shell::Shell, Z::Int64)
    if      sa == "ionization potential: inner-shell"
        shellIndex = Dict( Subshell("1s_1/2") => 1, Subshell("2s_1/2") => 2, Subshell("2p_1/2") => 3, Subshell("2p_3/2") => 4,
                           Subshell("3s_1/2") => 5, Subshell("3p_1/2") => 6, Subshell("3p_3/2") => 7,
                           Subshell("3d_3/2") => 8, Subshell("3d_3/2") => 9  )
        idx = shellIndex[shell]
        return( JAC.store("binding energies: Williams (2000)", Z)[idx] )
    else
        error("Unsupported keystring = $sa")
    end
end


using GeoParams: AbstractMaterialParam, AbstractMaterialParamsStruct
using ..Units
using Parameters, Unitful

# Computational routines needed for computations with the MaterialParams structure

# with tuple & vector - apply for all phases in MatParam
function compute_param!(fn::F, rho::Vector{_T}, MatParam::NTuple{N,AbstractMaterialParamsStruct}, P::_T=zero(_T),T::_T=zero(_T)) where {F,N,_T}
    rho .= map(x->fn(x,P,T), MatParam)
end

# each individual calcuation
function compute_param(fn::F, MatParam::NTuple{N,AbstractMaterialParamsStruct}, P::_T=zero(_T),T::_T=zero(_T)) where {F,N,_T}
    map(x->fn(x,P,T), MatParam)
end

#---------------------------------------------------------------------------------------------------------------------------#
#Computational routines for Phases

#performs computation given a single Phase
function compute_param(fn::F, MatParam::NTuple{N,AbstractMaterialParamsStruct}, Phase::Int64, P::_T=zero(_T),T::_T=zero(_T)) where {F,N,_T}
    Phase_tup = ntuple(i->MatParam[i].Phase, Val(N))
    ind = find_ind(Phase_tup, Phase)
    return fn(MatParam, P, T)[ind]
end

function compute_param!(fn::F, rho::AbstractArray{_T, ndim}, MatParam::NTuple{N,AbstractMaterialParamsStruct}, Phases::AbstractArray{_I, ndim}, P=nothing, T=nothing) where {F,ndim,N,_T,_I<:Integer}
    Phase_tup = ntuple(i->MatParam[i].Phase, Val(N))

    # check if we compute with P/T or use default values
    isnothing(T) ? compute_T = false : compute_T = true
    isnothing(P) ? compute_P = false : compute_P = true

    Tval = zero(_T)
    Pval = zero(_T)

    @inbounds for I in eachindex(Phases)
        phase   = find_ind(Phase_tup, Phases[I])

        # Extract relevant value if requested
        if compute_T; Tval = T[I]; end
        if compute_P; Pval = P[I]; end

        # this computes density for ALL phases, which is a bit of an overkill as we only need density for a single phase
        rho_tup = fn(MatParam, Pval, Tval)
        rho[I]  = rho_tup[phase]
    end
end

#Phase PhaseRatios

function compute_param!(fn::F, rho::AbstractArray{_T, N}, MatParam::NTuple{K,AbstractMaterialParamsStruct}, PhaseRatios::AbstractArray{_T, M}, P=nothing, T=nothing) where {F,_T<:AbstractFloat, N,M, K}
    if M!=(N+1)
        error("The PhaseRatios array should have one dimension more than the other arrays")
    end

    # check if we compute with P/T or use default values
    isnothing(T) ? compute_T = false : compute_T = true
    isnothing(P) ? compute_P = false : compute_P = true

    Tval = zero(_T)
    Pval = zero(_T)
    if compute_P
        Rindex = CartesianIndices(P)
    else
        Rindex = CartesianIndices(T)
    end

    @inbounds for I in Rindex
        frac    = view(PhaseRatios, Tuple(I)..., 1:K)    # fraction of each phase @ point I

        # Extract relevant value if requested
        if compute_T; Tval = T[I]; end
        if compute_P; Pval = P[I]; end

        # compute point-wise density:
        rho[I] = compute_param_times_frac(fn, frac, MatParam, Pval, Tval)
    end

    return
end

#Multiplies parameter with the fraction of a phase
function compute_param_times_frac(fn::F, PhaseRatios::AbstractArray{_T, 1}, MatParam::NTuple{N,AbstractMaterialParamsStruct}, P::_T, T::_T) where {F,_T, N}

    value_tup = fn(MatParam, P, T)

    # Sum & multiply density with fraction
    val = zero(_T)
    @inbounds for j=1:N
        val += PhaseRatios[j]*value_tup[j]
    end

    return val
end