using AlphaGo
using AlphaGo: from_kgs, from_sgf, BLACK, WHITE, EMPTY, MISSING_GROUP_ID,
              GoPosition, is_koish, is_eyeish, PlayerMove, is_move_legal,
              is_move_suicidal, score, from_board, flip_playerturn!, add_stone!,
              all_legal_moves, IllegalMove, from_flat, pass_move!, play_move!
using Base.Test

include("test_utils.jl")

pc_set(string, env::GoEnv) = Set(map(from_kgs, split(string),
                                        repmat([env], length(split(string)))))

@testset "go" begin
  env = GoEnv(9)
  EMPTY_ROW = repeat(".", env.N) * "\n"
  TEST_BOARD = load_board("""
                          .X.....OO
                          X........
                          """ * repeat(EMPTY_ROW, 7), env)
    @testset "load_board" begin
    @test assertEqualArray(env.EMPTY_BOARD, zeros(env.N, env.N))
    @test assertEqualArray(env.EMPTY_BOARD, load_board(repeat(". \n", env.N ^ 2), env))
  end

  @testset "parsing" begin
    @test from_kgs("A9", env) == (1, 1)
    @test from_sgf("aa") == (1, 1)
    @test from_kgs("A3", env) == (7, 1)
    @test from_sgf("ac") == (3, 1)
    @test from_kgs("D4", env) == from_sgf("df")
  end

  @testset "neighbors" begin
    corner = from_kgs("A1", env)
    neighbors = [env.EMPTY_BOARD[c...] for c in env.NEIGHBORS[corner...]]
    @test length(neighbors) == 2

    side = from_kgs("A2", env)
    side_neighbors = [env.EMPTY_BOARD[c...] for c in env.NEIGHBORS[side...]]
    @test length(side_neighbors) == 3
  end

  @testset "is_koish" begin
    @test is_koish(TEST_BOARD, from_kgs("A9", env), env) == BLACK
    @test is_koish(TEST_BOARD, from_kgs("B8", env), env) == nothing
    @test is_koish(TEST_BOARD, from_kgs("B9", env), env) == nothing
    @test is_koish(TEST_BOARD, from_kgs("E5", env), env) == nothing
  end

  @testset "is_eyeish" begin
    board = load_board("""
              .XX...XXX
              X.X...X.X
              XX.....X.
              ........X
              XXXX.....
              OOOX....O
              X.OXX.OO.
              .XO.X.O.O
              XXO.X.OO.
          """, env)
    B_eyes = pc_set("A2 A9 B8 J7 H8", env)
    W_eyes = pc_set("H2 J1 J3", env)
    not_eyes = pc_set("B3 E5", env)
    for be in B_eyes
      @test is_eyeish(board, be, env) == BLACK ? true : throw(AssertionError(string(be)))
    end
    for we in W_eyes
      @test is_eyeish(board, we, env) == WHITE ? true : throw(AssertionError(string(we)))
    end
    for ne in not_eyes
      @test is_eyeish(board, ne, env) == nothing ? true : throw(AssertionError(string(ne)))
    end
  end
  @testset "lib_tracker_init" begin
    board = load_board("X........" * repeat(EMPTY_ROW, 8), env)

    lib_tracker = from_board(board, env)
    @test length(lib_tracker.groups) == 1
    @test lib_tracker.group_index[from_kgs("A9", env)...] != MISSING_GROUP_ID
    @test lib_tracker.liberty_cache[from_kgs("A9", env)...] == 2
    sole_group = lib_tracker.groups[lib_tracker.group_index[from_kgs("A9", env)...]]
    @test sole_group.stones == pc_set("A9", env)
    @test sole_group.liberties == pc_set("B9 A8", env)
    @test sole_group.color == BLACK
  end

  @testset "place_stone" begin
    board = load_board("X........" * repeat(EMPTY_ROW, 8), env)
    lib_tracker = from_board(board, env)
    add_stone!(lib_tracker, BLACK, from_kgs("B9", env), env)
    @test length(lib_tracker.groups) == 1
    @test lib_tracker.group_index[from_kgs("A9", env)...] != MISSING_GROUP_ID
    @test lib_tracker.liberty_cache[from_kgs("A9", env)...] == 3
    @test lib_tracker.liberty_cache[from_kgs("B9", env)...] == 3
    sole_group = lib_tracker.groups[lib_tracker.group_index[from_kgs("A9", env)...]]
    @test sole_group.stones == pc_set("A9 B9", env)
    @test sole_group.liberties == pc_set("C9 A8 B8", env)
    @test sole_group.color == BLACK
  end

  @testset "place_stone_opposite_color" begin
    board = load_board("X........" * repeat(EMPTY_ROW, 8), env)
    lib_tracker = from_board(board, env)
    add_stone!(lib_tracker, WHITE, from_kgs("B9", env), env)
    @test length(lib_tracker.groups) == 2
    @test lib_tracker.group_index[from_kgs("A9", env)...] != MISSING_GROUP_ID
    @test lib_tracker.group_index[from_kgs("B9", env)...] != MISSING_GROUP_ID
    @test lib_tracker.liberty_cache[from_kgs("A9", env)...] == 1
    @test lib_tracker.liberty_cache[from_kgs("B9", env)...] == 2
    black_group = lib_tracker.groups[lib_tracker.group_index[from_kgs("A9", env)...]]
    white_group = lib_tracker.groups[lib_tracker.group_index[from_kgs("B9", env)...]]
    @test black_group.stones == pc_set("A9", env)
    @test black_group.liberties == pc_set("A8", env)
    @test black_group.color == BLACK
    @test white_group.stones == pc_set("B9", env)
    @test white_group.liberties == pc_set("C9 B8", env)
    @test white_group.color == WHITE
  end

  @testset "merge_multiple_groups" begin
    board = load_board("""
        .X.......
        X.X......
        .X.......
    """ * repeat(EMPTY_ROW, 6), env)
    lib_tracker = from_board(board,env)
    add_stone!(lib_tracker, BLACK, from_kgs("B8", env), env)
    @test length(lib_tracker.groups) == 1
    @test lib_tracker.group_index[from_kgs("B8", env)...] != MISSING_GROUP_ID
    sole_group = lib_tracker.groups[lib_tracker.group_index[from_kgs("B8", env)...]]
    @test sole_group.stones == pc_set("B9 A8 B8 C8 B7", env)
    @test sole_group.liberties == pc_set("A9 C9 D8 A7 C7 B6", env)
    @test sole_group.color == BLACK

    liberty_cache = lib_tracker.liberty_cache
    for stone in sole_group.stones
      @test liberty_cache[stone...] == 6 ? true : throw(AssertionError(string(stone)))
    end
  end

  @testset "capture_stone" begin
    board = load_board("""
        .X.......
        XO.......
        .X.......
    """ * repeat(EMPTY_ROW, 6), env)
    lib_tracker = from_board(board, env)
    captured = add_stone!(lib_tracker, BLACK, from_kgs("C8", env), env)
    @test length(lib_tracker.groups) == 4
    @test lib_tracker.group_index[from_kgs("B8", env)...] == MISSING_GROUP_ID
    @test captured == pc_set("B8", env)
  end

  @testset "capture_many" begin
    board = load_board("""
        .XX......
        XOO......
        .XX......
    """ * repeat(EMPTY_ROW, 6), env)
    lib_tracker = from_board(board, env)
    captured = add_stone!(lib_tracker, BLACK, from_kgs("D8", env), env)
    @test length(lib_tracker.groups) == 4
    @test lib_tracker.group_index[from_kgs("B8", env)...] == MISSING_GROUP_ID
    @test captured == pc_set("B8 C8", env)

    left_group = lib_tracker.groups[lib_tracker.group_index[from_kgs("A8", env)...]]
    @test left_group.stones == pc_set("A8", env)
    @test left_group.liberties == pc_set("A9 B8 A7", env)

    right_group = lib_tracker.groups[lib_tracker.group_index[from_kgs("D8", env)...]]
    @test right_group.stones == pc_set("D8", env)
    @test right_group.liberties == pc_set("D9 C8 E8 D7", env)

    top_group = lib_tracker.groups[lib_tracker.group_index[from_kgs("B9", env)...]]
    @test top_group.stones == pc_set("B9 C9", env)
    @test top_group.liberties == pc_set("A9 D9 B8 C8", env)

    bottom_group = lib_tracker.groups[lib_tracker.group_index[from_kgs("B7", env)...]]
    @test bottom_group.stones == pc_set("B7 C7", env)
    @test bottom_group.liberties == pc_set("B8 C8 A7 D7 B6 C6", env)

    liberty_cache = lib_tracker.liberty_cache
    for stone in top_group.stones
      @test liberty_cache[stone...] == 4 ? true : throw(AssertionError(string(stone)))
    end
    for stone in left_group.stones
      @test liberty_cache[stone...] == 3 ? true: throw(AssertionError(string(stone)))
    end
    for stone in right_group.stones
      @test liberty_cache[stone...] == 4 ? true : throw(AssertionError(string(stone)))
    end
    for stone in bottom_group.stones
      @test liberty_cache[stone...] == 6 ? true : throw(AssertionError(string(stone)))
    end
    for stone in captured
      @test liberty_cache[stone...] == 0 ? true : throw(AssertionError(string(stone)))
    end
  end

  @testset "capture_multiple_groups" begin
    board = load_board("""
        .OX......
        OXX......
        XX.......
    """ * repeat(EMPTY_ROW, 6), env)
    lib_tracker = from_board(board, env)
    captured = add_stone!(lib_tracker, BLACK, from_kgs("A9", env), env)
    @test length(lib_tracker.groups) == 2
    @test captured == pc_set("B9 A8", env)

    corner_stone = lib_tracker.groups[lib_tracker.group_index[from_kgs("A9", env)...]]
    @test corner_stone.stones == pc_set("A9", env)
    @test corner_stone.liberties == pc_set("B9 A8", env)

    surrounding_stones = lib_tracker.groups[lib_tracker.group_index[from_kgs("C9", env)...]]
    @test surrounding_stones.stones == pc_set("C9 B8 C8 A7 B7", env)
    @test surrounding_stones.liberties == pc_set("B9 D9 A8 D8 C7 A6 B6", env)

    liberty_cache = lib_tracker.liberty_cache
    for stone in corner_stone.stones
      @test liberty_cache[stone...] == 2 ? true : throw(AssertionError(string(stone)))
    end
    for stone in surrounding_stones.stones
      @test liberty_cache[stone...] == 7 ? true : throw(AssertionError(string(stone)))
    end
  end

  @testset "same_friendly_group_neighboring_twice" begin
    board = load_board("""
        XX.......
        X........
    """ * repeat(EMPTY_ROW, 7), env)

    lib_tracker = from_board(board, env)
    captured = add_stone!(lib_tracker, BLACK, from_kgs("B8", env), env)
    @test length(lib_tracker.groups) == 1
    sole_group_id = lib_tracker.group_index[from_kgs("A9", env)...]
    sole_group = lib_tracker.groups[sole_group_id]
    @test sole_group.stones == pc_set("A9 B9 A8 B8", env)
    @test sole_group.liberties == pc_set("C9 C8 A7 B7", env)
    @test captured == Set()
  end

  @testset "same_opponent_group_neighboring_twice" begin
    board = load_board("""
        XX.......
        X........
    """ * repeat(EMPTY_ROW, 7), env)

    lib_tracker = from_board(board, env)
    captured = add_stone!(lib_tracker, WHITE, from_kgs("B8", env), env)
    @test length(lib_tracker.groups) == 2
    black_group = lib_tracker.groups[lib_tracker.group_index[from_kgs("A9", env)...]]
    @test black_group.stones == pc_set("A9 B9 A8", env)
    @test black_group.liberties == pc_set("C9 A7", env)

    white_group = lib_tracker.groups[lib_tracker.group_index[from_kgs("B8", env)...]]
    @test white_group.stones == pc_set("B8", env)
    @test white_group.liberties == pc_set("C8 B7", env)

    @test captured == Set()
  end

  @testset "passing" begin
    start_position = GoPosition(env,
        board = TEST_BOARD,
        n = 0,
        komi = 6.5,
        caps = (1, 2),
        ko = PlayerMove(from_kgs("A1", env)...),
        recent = Vector{PlayerMove}(),
        to_play = BLACK
    )
    expected_position = GoPosition(env,
        board = TEST_BOARD,
        n = 1,
        komi = 6.5,
        caps = (1, 2),
        ko = nothing,
        recent = [PlayerMove(BLACK, nothing)],
        to_play = WHITE
    )
    pass_position = pass_move!(start_position)
    @test assertEqualPositions(pass_position, expected_position)
  end

  @testset "flipturn" begin
    start_position = GoPosition(env,
        board = TEST_BOARD,
        n = 0,
        komi = 6.5,
        caps = (1, 2),
        ko = PlayerMove(from_kgs("A1", env)...),
        recent = Vector{PlayerMove}(),
        to_play = BLACK
    )
    expected_position = GoPosition(env,
        board = TEST_BOARD,
        n = 0,
        komi = 6.5,
        caps = (1, 2),
        ko = nothing,
        recent = Vector{PlayerMove}(),
        to_play = WHITE
    )
    flip_position = flip_playerturn!(start_position)
    @test assertEqualPositions(flip_position, expected_position)
  end

  @testset "is_move_suicidal" begin
    board = load_board("""
        ...O.O...
        ....O....
        XO.....O.
        OXO...OXO
        O.XO.OX.O
        OXO...OOX
        XO.......
        ......XXO
        .....XOO.
    """, env)
    position = GoPosition(env,
        board = board,
        to_play = BLACK
    )
    suicidal_moves = pc_set("E9 H5", env)
    nonsuicidal_moves = pc_set("B5 J1 A9", env)
    for move in suicidal_moves
      @test position.board[move...] == EMPTY #sanity check my coordinate input
      @test is_move_suicidal(position, move) ? true : throw(AssertionError(string(move)))
    end
    for move in nonsuicidal_moves
      @test position.board[move...] == EMPTY # sanity check my coordinate input
      @test !is_move_suicidal(position, move) ? true : throw(AssertionError(string(move)))
    end
  end

  @testset "legal_moves" begin
    board = load_board("""
        .O.O.XOX.
        O..OOOOOX
        ......O.O
        OO.....OX
        XO.....X.
        .O.......
        OX.....OO
        XX...OOOX
        .....O.X.
    """, env)
    position = GoPosition(env,
        board = board,
        to_play = BLACK)
    illegal_moves = pc_set("A9 E9 J9", env)
    legal_moves = pc_set("A4 G1 J1 H7", env)
    for move in illegal_moves
      @test !is_move_legal(position, move)
    end
    for move in legal_moves
      @test is_move_legal(position, move)
    end
    # check that the bulk legal test agrees with move-by-move illegal test.
    bulk_legality = all_legal_moves(position)
    for (i, bulk_legal) in enumerate(bulk_legality)
      @test is_move_legal(position, from_flat(i, position)) == bulk_legal
    end
    # flip the colors and check that everything is still (il)legal
    position = GoPosition(env,board = -board, to_play = WHITE)
    for move in illegal_moves
      @test !is_move_legal(position, move)
    end
    for move in legal_moves
      @test is_move_legal(position, move)
    end
    bulk_legality = all_legal_moves(position)
    for (i, bulk_legal) in enumerate(bulk_legality)
      @test is_move_legal(position, from_flat(i, position)) == bulk_legal
    end
  end

  @testset "move" begin
    start_position = GoPosition(env,
        board=TEST_BOARD,
        n = 0,
        komi = 6.5,
        caps =(1, 2),
        ko = nothing,
        recent = Vector{PlayerMove}(),
        to_play = BLACK
    )
    expected_board = load_board("""
        .XX....OO
        X........
    """ * repeat(EMPTY_ROW, 7), env)
    expected_position = GoPosition(env,
        board = expected_board,
        n = 1,
        komi = 6.5,
        caps = (1, 2),
        ko = nothing,
        recent=[PlayerMove(BLACK, from_kgs("C9", env))],
        to_play = WHITE
    )
    actual_position = play_move!(start_position, from_kgs("C9", env))
    @test assertEqualPositions(actual_position, expected_position)

    expected_board2 = load_board("""
        .XX....OO
        X.......O
    """ * repeat(EMPTY_ROW, 7), env)
    expected_position2 = GoPosition(env,
        board = expected_board2,
        n = 2,
        komi = 6.5,
        caps = (1, 2),
        ko = nothing,
        recent = [PlayerMove(BLACK, from_kgs("C9", env)), PlayerMove(WHITE, from_kgs("J8", env))],
        to_play = BLACK
    )
    actual_position2 = play_move!(actual_position, from_kgs("J8", env))
    @test assertEqualPositions(actual_position2, expected_position2)
  end

  @testset "move_with_capture" begin
    start_board = load_board(repeat(EMPTY_ROW, 5) * """
        XXXX.....
        XOOX.....
        O.OX.....
        OOXX.....
    """, env)
    start_position = GoPosition(env,
        board = start_board,
        n = 0,
        komi = 6.5,
        caps = (1, 2),
        ko = nothing,
        recent = Vector{PlayerMove}(),
        to_play = BLACK
    )
    expected_board = load_board(repeat(EMPTY_ROW, 5) * """
        XXXX.....
        X..X.....
        .X.X.....
        ..XX.....
    """, env)
    expected_position = GoPosition(env,
        board = expected_board,
        n = 1,
        komi = 6.5,
        caps = (7, 2),
        ko = nothing,
        recent = Vector{PlayerMove}([PlayerMove(BLACK, from_kgs("B2", env))]),
        to_play = WHITE
    )
    actual_position = play_move!(start_position, from_kgs("B2", env))
    @test assertEqualPositions(actual_position, expected_position)
  end

  @testset "ko_move" begin
    start_board = load_board("""
        .OX......
        OX.......
    """ * repeat(EMPTY_ROW, 7), env)
    start_position = GoPosition(env,
        board = start_board,
        n = 0,
        komi = 6.5,
        caps = (1, 2),
        ko = nothing,
        recent = Vector{PlayerMove}(),
        to_play = BLACK
    )
    expected_board = load_board("""
        X.X......
        OX.......
    """ * repeat(EMPTY_ROW, 7), env)
    expected_position = GoPosition(env,
        board = expected_board,
        n = 1,
        komi = 6.5,
        caps = (2, 2),
        ko = from_kgs("B9", env),
        recent = Vector{PlayerMove}([PlayerMove(BLACK, from_kgs("A9", env))]),
        to_play = WHITE
    )
    actual_position = play_move!(start_position, from_kgs("A9", env))
    @test assertEqualPositions(actual_position, expected_position)

    # Check that retaking ko is illegal until two intervening moves
    @test_throws IllegalMove play_move!(actual_position, from_kgs("B9", env))
    pass_twice = pass_move!(pass_move!(actual_position))
    ko_delayed_retake = play_move!(pass_twice, from_kgs("B9", env))
    expected_position = GoPosition(env,
        board = start_board,
        n = 4,
        komi = 6.5,
        caps = (2, 3),
        ko = from_kgs("A9", env),
        recent = Vector{PlayerMove}([
            PlayerMove(BLACK, from_kgs("A9", env)),
            PlayerMove(WHITE, nothing),
            PlayerMove(BLACK, nothing),
            PlayerMove(WHITE, from_kgs("B9", env))]),
        to_play = BLACK
        )

    @test assertEqualPositions(ko_delayed_retake, expected_position)
  end

  @testset "is_game_over" begin
    root = GoPosition(env)
    @test !root.done
    first_pass = play_move!(root, nothing)
    @test !first_pass.done
    second_pass = play_move!(first_pass, nothing)
    @test second_pass.done
  end

  @testset "scoring" begin
    board = load_board("""
        .XX......
        OOXX.....
        OOOX...X.
        OXX......
        OOXXXXXX.
        OOOXOXOXX
        .O.OOXOOX
        .O.O.OOXX
        ......OOO
    """, env)
    position = GoPosition(env,
        board = board,
        n = 54,
        komi = 6.5,
        caps = (2, 5),
        ko = nothing,
        recent = Vector{PlayerMove}(),
        to_play = BLACK
    )
    expected_score = 1.5
    @test score(position) == expected_score

    board = load_board("""
        XXX......
        OOXX.....
        OOOX...X.
        OXX......
        OOXXXXXX.
        OOOXOXOXX
        .O.OOXOOX
        .O.O.OOXX
        ......OOO
      """, env)
    position = GoPosition(env,
        board = board,
        n = 55,
        komi = 6.5,
        caps = (2, 5),
        ko = nothing,
        recent = Vector{PlayerMove}(),
        to_play = WHITE
    )
    expected_score = 2.5
    @test score(position) == expected_score
  end
end


function extract_nodes(file, ::Type{TF}=Float64, ::Type{TI}=Int) where {TF,TI}
    line = readline(file)

    pattern = r"(\d+)\s*,\s*(-?\d+\.?\d*(e[-\+]?\d*)?)\s*,\s*(-?\d+\.?\d*(e[-\+]?\d*)?)\s*(,\s*(-?\d+\.?\d*(e[-\+]?\d*)?))?"
    m = match(pattern, line)

    first_node_idx = parse(TI, m[1])
    first_node_idx == TI(1) || throw("First node index is not 1.")

    if m[6] isa Nothing
        node_coords = [(parse(TF, m[2]), parse(TF, m[4]))]
        nextline = _extract_nodes!(node_coords, file, TI(1))
    else
        node_coords = [(parse(TF, m[2]), parse(TF, m[4]), parse(TF, m[7]))]
        nextline = _extract_nodes!(node_coords, file, TI(1))
    end
    return node_coords, nextline
end

function _extract_nodes!(
    node_coords::AbstractVector{NTuple{dim,TF}}, file, prev_node_idx::TI
) where {dim,TF,TI}
    if dim === 2
        pattern = r"(\d+)\s*,\s*(-?\d+\.?\d*(e[-\+]?\d*)?)\s*,\s*(-?\d+\.?\d*(e[-\+]?\d*)?)"
    elseif dim === 3
        pattern = r"(\d+)\s*,\s*(-?\d+\.?\d*(e[-\+]?\d*)?)\s*,\s*(-?\d+\.?\d*(e[-\+]?\d*)?)\s*,\s*(-?\d+\.?\d*(e[-\+]?\d*)?)"
    else
        error("Dimension is not supported.")
    end

    line = readline(file)
    m = match(stopping_pattern, line)
    while m isa Nothing
        m = match(pattern, line)
        if m != nothing
            node_idx = parse(Int, m[1])
            node_idx == prev_node_idx + TI(1) || throw("Node indices are not consecutive.")
            if dim === 2
                push!(node_coords, (parse(TF, m[2]), parse(TF, m[4])))
            else
                push!(node_coords, (parse(TF, m[2]), parse(TF, m[4]), parse(TF, m[6])))
            end
            prev_node_idx = node_idx
        end
        line = readline(file)
        m = match(stopping_pattern, line)
    end
    return line
end


module CPUTime

export
    CPUtime_us,
    CPUtic,
    CPUtoq,
    CPUtoc,
    @CPUtime,
    @CPUelapsed

function CPUtime_us()
    rusage = Libc.malloc(4*sizeof(Clong) + 14*sizeof(UInt64))  # sizeof(uv_rusage_t); this is different from sizeof(rusage)
    ccall(:uv_getrusage, Cint, (Ptr{Nothing},), rusage)
    utime = UInt64(1000000)*unsafe_load(convert(Ptr{Clong}, rusage + 0*sizeof(Clong))) +    # user CPU time
                            unsafe_load(convert(Ptr{Clong}, rusage + 1*sizeof(Clong)))
    stime = UInt64(1000000)*unsafe_load(convert(Ptr{Clong}, rusage + 2*sizeof(Clong))) +    # system CPU time
                            unsafe_load(convert(Ptr{Clong}, rusage + 3*sizeof(Clong)))
    ttime = utime + stime  # total CPU time
    Libc.free(rusage)
    return ttime
end

function CPUtic()
    t0 = CPUtime_us()
    task_local_storage(:CPUTIMERS, (t0, get(task_local_storage(), :CPUTIMERS, ())))
    return t0
end

function CPUtoq()
    t1 = CPUtime_us()
    timers = get(task_local_storage(), :CPUTIMERS, ())
    if timers === ()
        error("CPUtoc() without CPUtic()")
    end
    t0 = timers[1]::UInt64
    task_local_storage(:CPUTIMERS, timers[2])
    (t1-t0)/1e6
end

function CPUtoc()
    t = CPUtoq()
    println("elapsed CPU time: ", t, " seconds")
    return t
end

# print elapsed CPU time, return expression value
macro CPUtime(ex)
    quote
        local t0 = CPUtime_us()
        local val = $(esc(ex))
        local t1 = CPUtime_us()
        println("elapsed CPU time: ", (t1-t0)/1e6, " seconds")
        val
    end
end

# print nothing, return elapsed CPU time
macro CPUelapsed(ex)
    quote
        local t0 = CPUtime_us()
        local val = $(esc(ex))
        (CPUtime_us()-t0)/1e6
    end
end

end # module


using ArgParse

include("./src/raft.jl")
using Raft

## I'm just a follower to start with
ps = Raft.PersistentState(0, 0, [])
vs = Raft.VolatileState(0, 0)

s = ArgParseSettings("Runs a Raft server")

@add_arg_table s begin
  "--server_id", "-s"
      help = "id of this server"
      arg_type = UInt64
      required = true
  "--port", "-p"
      help = "port to listen on"
      required = true
      arg_type = UInt64
  "--buddies", "-b"
      nargs = '+'
      help = "ports of buddies to talk to"
      arg_type = UInt64
      required = true
  "--amleader", "-l"
      action = :store_true
      help = "This server should be the leader"
end

## Read config:
parsed_args = parse_args(ARGS, s)
for key in parsed_args
  @show key
end

buddies = Dict{Raft.Server, Raft.Address}()
i = UInt64(1)
for b in parsed_args["buddies"]
  buddies[i] = Raft.Address(ip"127.0.0.1", b)
  i+=1
end

my_config = Raft.ServerConfig(parsed_args["server_id"],
  parsed_args["port"], buddies)
am_leader = parsed_args["amleader"]

## Wait for leader to contact me
if !am_leader
  Raft.run_follower(my_config)
else
  Raft.run_leader(my_config)
end


"""
Package: ForecastPlots

    candel(x, t, args...; colors = (:green, :red), kw...)

Plot a candlestick for stock prices

# Arguments
- `x`: Matrix with four columns for Open, High, Low and Close values.
- `t`: Vector with time units, it defaults to 1:size(x,1)
- `args...`: `plot` arguments
- `colors`: Tuple with the up and down color symbols, it defaults to (:green,:red)
- `kw...`: `plot` keyword arguments

# Returns
Candle plot

# Examples
```julia-repl
using MarketData

x = yahoo(:INTC)
last_month_intc = values(x)[end-30:end,1:4]
candle(last_month_intc)
candle(last_month_intc,timestamp(x)[end-30:end])
candle(last_month_intc,timestamp(x)[end-30:end]; colors=(:white,:black), title = "INTC")
```
"""
function candle(x::Matrix{<:Real},t=1:size(x,1); colors=(:green,:red), kw...)

    @assert length(colors) == 2 "Two colors required"
    @assert eltype(colors) == Symbol "'colors' must be symbols e.g. (:white,:black)"
    @assert colors[1] != colors[2] "'colors' must be different"

    plot(t,repeat([missing],size(x,1));labels=nothing,kw...)

    ttype = eltype(t) <: TimeType
    t = ttype ? Dates.value.(t) : t
    woc = ttype ? (t[2]-t[1])/2  : 1/2
    whl = woc/10

    rect(w, h, x, y) = Shape(x .+ [0, w, w, 0, 0], y .+ [0, 0, h, h, 0])

    for (i,ohlc) in zip(t,eachrow(x))

        # High / Low
        plot!(rect(whl,ohlc[2]-ohlc[3],i,ohlc[3]);
              seriescolor = :black,
              linewidth = 0,
              label = nothing)

        # Open / Close
        plot!(rect(woc,ohlc[4]-ohlc[1],i-woc/2,ohlc[1]);
              seriescolor = ohlc[1]<ohlc[4] ? colors[1] : colors[2],
              linewidth = 0,
              label = nothing)

    end
    current()
end


__precompile__()

module ERFA

const depsfile = joinpath(dirname(dirname(@__FILE__)), "deps", "deps.jl")
if isfile(depsfile)
    include(depsfile)
else
    error("ERFA is not properly installed. Please run Pkg.build(\"ERFA\")")
end

include("erfa_common.jl")
include("deprecated.jl")

include("a.jl")
include("b.jl")
include("c.jl")
include("d.jl")
include("e.jl")
include("f.jl")
include("g.jl")
include("h.jl")
include("i.jl")
include("j.jl")
include("l.jl")
include("n.jl")
include("o.jl")
include("p.jl")
include("r.jl")
include("s.jl")
include("t.jl")
include("u.jl")
include("x.jl")

function ASTROM(pmt, eb::AbstractArray, eh::AbstractArray, em, v::AbstractArray, bm1, bpn::AbstractArray, along, phi, xpl, ypl, sphi, cphi, diurab, eral, refa, refb)
    ASTROM(pmt,
              (eb[1], eb[2], eb[3]),
              (eh[1], eh[2], eh[3]),
              em,
              (v[1], v[2], v[3]),
              bm1,
              (bpn[1], bpn[2], bpn[3], bpn[4], bpn[5], bpn[6], bpn[7], bpn[8], bpn[9]),
              along,
              phi,
              xpl,
              ypl,
              sphi,
              cphi,
              diurab,
              eral,
              refa,
              refb)
end

function LDBODY(bm, dl, pv::AbstractArray)
    LDBODY(bm, dl, (pv...,))
end

end # module


# Tests data_import.jl

using Test
using GeophysicalModelGenerator

# test the import of a csv file with depth given as positive values and units in km

# test loaded structure
# NOTE: currently deactivated these tests -----------------------------
#LoadedData   =   ReadCSV_LatLon("TestData.csv", "positive");
#@test LoadedData.lat.name  ==  "lat"
#@test LoadedData.lat.unit  ==  "deg"
#@test LoadedData.lat.values  ==  [36.424, 36.194, 36.144, 35.486, 36.106, 35.767, 36.24, 36.301, 36.401, 36.321]

#@test LoadedData.lon.name  ==  "lon"
#@test LoadedData.lon.unit  ==  "deg"
#@test LoadedData.lon.values  ==  [68.74, 71.111, 70.757, 70.38, 71.051, 70.55, 68.95, 71.303, 71.567, 70.829]

#@test LoadedData.depth.name  ==  "depth"
#@test LoadedData.depth.unit  ==  "km"
#@test LoadedData.depth.values  ==  [-52.04, -60.51, -60.74, -61.0, -75.5, -81.52, -82.55, -84.36, -90.9, -110.41]

#@test LoadedData.values.varnames == ["local_magnitude", "absError(km)"]
#@test LoadedData.values.vals == [5.5   9.7;5.6  10.0;5.1  10.7;5.5  20.4;5.1   6.5;5.9  10.1;5.2   5.5;5.1   3.2;5.3   3.6;5.3   3.6]
# ---------------------------------------------------------------------

# test loading images (profiles & mapviews)

# Extact & save profile in GeoData format
filename            =   "test.png";             # fake png
Corner_LowerLeft    =   (18.0, 51.0, -590.0)
Corner_UpperRight   =   (9.0, 42.0,    0.0)
data_Image          =   Screenshot_To_GeoData(filename,Corner_LowerLeft, Corner_UpperRight)
@test data_Image.lon[1000]==17.592964824120603
@test data_Image.lat[1000]==50.59296482412061
@test data_Image.depth[1000]==-590km
@test Write_Paraview(data_Image, "Profile_1")==nothing

# Test in CartData
data_Image          =   Screenshot_To_GeoData(filename,Corner_LowerLeft, Corner_UpperRight, Cartesian=true)
@test data_Image.x.val[22] == 18.0km
@test data_Image.y.val[22] == 51.0km
@test data_Image.z.val[22] ≈ -125.15151515151516km

# Test in UTM zone [note that depth should be in m]
data_Image          =   Screenshot_To_GeoData(filename,Corner_LowerLeft, Corner_UpperRight, UTM=true, UTMzone=33, isnorth=true)
@test data_Image.EW.val[22] == 18.0
@test data_Image.NS.val[22] == 51.0
@test data_Image.depth.val[22] ≈ -125.15151515151516m

# Mapview (distorted) in GeoData format
filename            =   "test.png";             # fake png
Corner_LowerLeft    =   (2.0,  40.0, -15.0)
Corner_UpperRight   =   (22.0, 51.0, -15.0)
Corner_LowerRight   =   (20.0, 40.0, -15.0)
Corner_UpperLeft    =   (0.0,  51.0, -15.0)
data_Image          =   Screenshot_To_GeoData(filename,Corner_LowerLeft, Corner_UpperRight, Corner_LowerRight=Corner_LowerRight, Corner_UpperLeft=Corner_UpperLeft)
@test data_Image.lon[1000]==2.814070351758794
@test data_Image.lat[1000]==40.00000000000001
@test data_Image.depth[1000]==-15km
@test Write_Paraview(data_Image, "MapView_1") == nothing

# MapView in CartData
data_Image          =   Screenshot_To_CartData(filename,Corner_LowerLeft, Corner_UpperRight, Corner_LowerRight=Corner_LowerRight, Corner_UpperLeft=Corner_UpperLeft)
@test data_Image.x.val[22] ≈ 0.42424242424242425km
@test data_Image.y.val[22] ≈ 48.666666666666664km
@test data_Image.z.val[22] ≈ -15km

# MapView in UTMData
data_Image          =   Screenshot_To_UTMData(filename,Corner_LowerLeft, Corner_UpperRight, Corner_LowerRight=Corner_LowerRight, Corner_UpperLeft=Corner_UpperLeft, UTMzone=33, isnorth=true)
@test data_Image.EW.val[22] ≈ 0.42424242424242425
@test data_Image.NS.val[22] ≈ 48.666666666666664
@test data_Image.depth.val[22] ≈ -15.0m


#
# Protein - TMAO (compare new and old implementations)
#

include("../../src/MDDF.jl")
using PDBTools
using Plots

dir="/home/leandro/Drive/Alunos/Ivan/PCCP_revision"

# Here we use the PDBTools package to read the pdb file (from http://github.com/m3g/PDBTools)
atoms = PDBTools.readPDB("$dir/6Mnative.pdb")

# Set solute and solvent

protein = PDBTools.select(atoms,"protein")
solute = MDDF.Selection(protein,nmols=1 )

water = PDBTools.select(atoms,"water")
solvent = MDDF.Selection(water,natomspermol=3)

options = MDDF.Options()

# Initialize trajectroy data structure and open input stream
#R = MDDF.mddf(trajectory,options)
#MDDF.save(R,"ivan.json")

trajectory = MDDF.NamdDCD("$dir/6Mnative.dcd",solute,solvent)
#lc = MDDF.mddf_linkedcells(trajectory,options)

#trajectory = MDDF.Trajectory("$dir/6Mnative.dcd",solute,solvent)
lcP = MDDF.mddf(trajectory,options)
#lcP = MDDF.load("./ivan.json")

plot(layout=(5,1))

sp=1
plot!(ylabel="mddf",sp=sp)
plot!(lc.d,lc.mddf,linewidth=2,label="lc",subplot=sp)
scatter!(lcP.d,lcP.mddf,linewidth=2,label="lcP",subplot=sp)

sp=2
plot!(ylabel="md_count",sp=sp)
plot!(lc.d,lc.md_count,linewidth=2,label="lc",subplot=sp)
scatter!(lcP.d,lcP.md_count,linewidth=2,label="lcP",subplot=sp)

sp=3
plot!(ylabel="md_count_random",sp=sp)
plot!(lc.d,lc.md_count_random,linewidth=2,label="lc",subplot=sp)
scatter!(lcP.d,lcP.md_count_random,linewidth=2,label="lcP",subplot=sp)

sp=4
plot!(ylabel="rdf_count",sp=sp)
plot!(lc.d,lc.rdf_count,linewidth=2,label="lc",subplot=sp)
scatter!(lcP.d,lcP.rdf_count,linewidth=2,label="lcP",subplot=sp)

sp=5
plot!(ylabel="rdf_count_random",sp=sp)
plot!(lc.d,lc.rdf_count_random,linewidth=2,label="lc",subplot=sp)
scatter!(lcP.d,lcP.rdf_count_random,linewidth=2,label="lcP",subplot=sp)

plot!(size=(400,800))
savefig("./ivan.pdf")
savefig("./plots.pdf")


using AtomicLevels: CSF
using HalfIntegers

angularmomentum(o::RelativisticOrbital) = AtomicLevels.κ2j(o.κ)
angularmomentum(csf::CSF{O,<:HalfInteger}) where O = last(csf.terms)

# Relativistic CSLs and parsing of GRASP CSL files

function parse_rcsf(filename)
    open(filename, "r") do io
        # First line should be "Core subshells:"
        line = readline(io); @assert strip(line) == "Core subshells:"
        line_cores = strip(readline(io), ['\n', '\r'])
        line = readline(io); @assert strip(line) == "Peel subshells:"
        line_peels = readline(io)
        line = readline(io); @assert strip(line) == "CSF(s):"

        core_orbitals = parse_cores(line_cores)
        core_couplings = zeros(HalfInt, length(core_orbitals))
        core_occupations = map(degeneracy, core_orbitals)

        blockid, csfid = 1, 1
        csfs = CSF{RelativisticOrbital{Int},HalfInt,Seniority}[]
        while ! eof(io)
            line1 = readline(io)
            if startswith(line1, " *")
                blockid += 1
                csfid = 1
                line1 = readline(io)
            end
            line1 = strip(line1, ['\n', '\r'])
            line2 = strip(readline(io), ['\n', '\r'])
            line3 = strip(readline(io), ['\n', '\r'])

            orbitals, noccupations, orbcouplings, csfcouplings = parse_csflines(line1, line2, line3)

            @assert !any(isequal(0), noccupations) # we should never get 0-electron orbitals

            # Fix orbcouplings that are not explicitly written in the CSL file (represented
            # with nothings in the array). It is assumed that this is the case only if the
            # orbital is fully occupied.
            #
            # NOTE: we could, though, also omit values if there is only 1 electron or
            # maxelectrons(orb) - 1 electrons, this which case the angular momentum can
            # only have only one value too.
            for i = 1:length(orbitals)
                orbital, nelec = orbitals[i], noccupations[i]
                if orbcouplings[i] === nothing
                    nelec == degeneracy(orbital) || error("""
                    Unable to fix missing orbital coupling.
                      Orbital $i, orbital=$(repr(orbital)), nelec=$nelec
                    1: $(line1)
                    2: $(line2)
                    3: $(line3)
                    """)
                    orbcouplings[i] = 0
                elseif (nelec == 1 || nelec == degeneracy(orbital) - 1) && orbcouplings[i] != angularmomentum(orbital)
                    @warn "Bad orbital coupling" orbcouplings[i] nelec angularmomentum(orbital)
                elseif orbcouplings[i] > nelec * angularmomentum(orbital)
                    # If the orbital coupling is larger than (# particles) * (l of orbital),
                    # then that has to be an error.
                    @warn "Bad orbital coupling" orbcouplings[i] nelec angularmomentum(orbital)
                end
            end

            # Fix csfcouplings which are not explicitly written to the CSL file. This appears
            # to be the case for "trivial" zero couplings, if the previous CSF layer and
            # current orbital are both zeros.
            for i = 1:length(orbitals)
                oj = orbcouplings[i]
                cj = (i > 1) ? csfcouplings[i-1] : zero(HalfInt)
                Δupper, Δlower = oj+cj, abs(oj-cj)
                if csfcouplings[i] === nothing
                    Δupper == Δlower || error("""
                    Unable to fix missing CSF coupling.
                      Orbital $i, orbital=$(repr(orbitals[i])), oj=$(repr(oj)), cj=$(repr(cj))
                    1: $(line1)
                    2: $(line2)
                    3: $(line3)
                    """)
                    csfcouplings[i] = Δlower
                elseif !(Δlower <= csfcouplings[i] <= Δupper)
                    @warn "Invalid csfcoupling value?" csfcouplings[i] Δupper Δlower
                end
            end

            config = Configuration(vcat(core_orbitals, orbitals), vcat(core_occupations, noccupations), sorted=true)
            # I have no idea how to generate correct seniority numbers, even in the case where
            # it does not matter (no degeneracy). So instead, I'll try to match the couplings
            # to the intermediate terms generated by AtomicLevels.
            subshell_terms = let its = intermediate_terms.(vcat(core_orbitals, orbitals), vcat(core_occupations, noccupations))
                map(its, vcat(core_couplings, Vector{HalfInt}(orbcouplings))) do its, grasp_term
                    idxs = findall(it -> it.term == grasp_term, its)
                    @assert length(idxs) == 1 # make sure there is only one intermediate term with the coupling we're interested in
                    its[first(idxs)]
                end
            end
            terms = map(x -> convert(Rational{Int}, x),
                vcat(core_couplings, Vector{HalfInt}(csfcouplings)))
            csf = CSF(config, subshell_terms, terms)
            push!(csfs, csf)
        end
        return csfs
    end
end

function parse_csflines(line1, line2, line3)
    # Assuming that the CSF line consists of NNNLL(EE) blocks, each 9 characters long.
    @assert length(line1) % 9 == 0

    orbs = RelativisticOrbital{Int}[]
    orbs_nelectrons = Int[]
    orbs_orbcouplings = Union{HalfInt,Nothing}[]
    orbs_csfcouplings = Union{HalfInt,Nothing}[]
    norbitals = div(length(line1), 9) # number of orbitals in this group
    for i = 1:norbitals
        orb = line1[9*(i-1)+1:9*i]
        @assert orb[6]=='(' && orb[9]==')'
        orbital = parse(RelativisticOrbital, strip(orb[1:5]))
        # n = parse(Int, orb[1:3])
        # kappa = parse_j(orb[4:5])
        nelec = parse(Int, orb[7:8])

        # pick out coupled angular momentum from the second line:
        angmom =  if length(line2) < 9*i
            nothing
        else
            parse_angmom_string(line2[9*(i-1)+1:9*i])
        end

        # Pick the J-coupling from between the orbitals (line3).
        # The items in that line are shifted by three characters to the right for
        # some reason.. except for the last one, which defines the J^P values of
        # the CSF. That one is shifted by 1 character.. and then one after that
        # is the parity +/- symbol.
        c2J_idx_first, c2J_idx_last = if i < norbitals
            # So, for the non-last ones we assume a 9 character block that has been
            # shifted by 3 characters to the right.
            # +1 to go from 0-based to 1-based
            9*(i-1)+3+1, 9*i+3
        else # i == norbitals -- the last non-regular coupling
            9*(i-1)+3+1, 9*i+1
        end
        c2J_string = line3[c2J_idx_first:c2J_idx_last]
        coupled_angmom = try
            parse_angmom_string(c2J_string)
        catch e
            error("""
            Error parsing 2J value on line 3 (i=$i)
              range $(c2J_idx_first):$(c2J_idx_last) -> '$(c2J_string)'
            1: $(line1)
            2: $(line2)
            3: $(line3)
            $(' '^(c2J_idx_first+2))^$('-'^(c2J_idx_last-c2J_idx_first-1))^
            """)
        end
        push!(orbs_csfcouplings, coupled_angmom)

        push!(orbs, orbital)
        push!(orbs_nelectrons, nelec)
        push!(orbs_orbcouplings, angmom)
    end
    @assert length(orbs_csfcouplings) == norbitals

    return orbs, orbs_nelectrons, orbs_orbcouplings, orbs_csfcouplings
end

function parse_angmom_string(s)
    length(strip(s)) == 0 && return nothing
    parse(HalfInt, s)
end

function parse_cores(line)
    orbstrings = split(line)
    orbs = RelativisticOrbital{Int}[]
    for os in orbstrings
        push!(orbs, parse(RelativisticOrbital, strip(os)))
    end
    orbs
end


abstract type AbstractFlowDynamics end

function flow!(balls::Vector{<:MovableBall}, t, flow_type::AbstractFlowDynamics)
	for ball in balls
		flow!(ball, t, flow_type)
	end
end


# Code originally by Kristoffer Carlsson (2016-ish)
# Code modified by Jesse Chan (2020)
struct SparseMatrixBSC{Tv, Ti <: Integer}
    R::Int                 # Block size in rows
    C::Int                 # Block size in columns
    m::Int                 # Number of rows
    n::Int                 # Number of columns
    colptr::Vector{Ti}     # Column i is in colptr[i]:(colptr[i+1]-1)
    rowval::Vector{Ti}     # Row values of blocks
    nzval::Array{Tv, 3}    # Nonzero values, one "matrix" per block, nzval[i, j, block]

    function SparseMatrixBSC(R::Integer, C::Integer,
                             m::Integer, n::Integer,
                             colptr::Vector{Ti}, rowval::Vector{Ti},  nzval::Array{Tv, 3}) where {Ti, Tv}
      m < 0 && throw(ArgumentError("number of rows (m) must be ≥ 0, got $m"))
      n < 0 && throw(ArgumentError("number of columns (n) must be ≥ 0, got $n"))

      R < 0 && throw(ArgumentError("block size y must be ≥ 0, got $x_block_size"))
      C < 0 && throw(ArgumentError("block size x must be ≥ 0, got $y_block_size"))

      m % R != 0 && throw(ArgumentError("row block size: $(R) must evenly divide number of rows: $m"))
      n % C != 0 && throw(ArgumentError("column block size: $(C) must evenly divide number of rows: $n"))
      new{Tv, Ti}(Int(R), Int(C), Int(m), Int(n), colptr, rowval, nzval)
  end
end

function SparseMatrixBSC(m::Integer, n::Integer, colptr::Vector{Ti}, rowval::Vector{Ti}, nzval::Array{Tv, 3}) where {Tv, Ti<:Integer}
    R, C = size(nzval, 1), size(nzval, 2)
    SparseMatrixBSC(R, C, m, n, colptr, rowval, nzval)
end

###########################
# Conversions CSC <-> BSC #
###########################

nblocks(A::SparseMatrixBSC) = (length(A.colptr) - 1, A.n ÷ A.C)
nblocks(A::SparseMatrixBSC, i::Int) = nblocks(A)[i]
blocksize(A::SparseMatrixBSC) = A.R, A.C
blocksize(A::SparseMatrixBSC, i::Int) = blocksize(A)[i]

# Column i is in colptr[i]:(colptr[i+1]-1)
nzblockrange(A::SparseMatrixBSC, col::Integer) =  Int(A.colptr[col]):Int(A.colptr[col + 1] - 1) # returns range for blocks in column "col"

Base.size(A::SparseMatrixBSC) = (A.m, A.n)
Base.size(A::SparseMatrixBSC, i) = size(A)[i]
Base.eltype(A::SparseMatrixBSC) = eltype(A.nzval)
SparseArrays.nnz(A::SparseMatrixBSC) = length(A.nzval)

function SparseArrays.SparseMatrixCSC(A::SparseMatrixBSC{Tv, Ti}) where {Tv, Ti <: Integer}
    if blocksize(A) == (1,1)
        return SparseMatrixCSC(A.m, A.n, A.colptr, A.rowval, vec(A.nzval))
    end
    rowval = zeros(Ti, nnz(A))
    colptr = zeros(Ti, size(A, 2) + 1)
    nzval = zeros(Tv, length(A.nzval))

    count_row = 1
    count_col = 2
    colptr[1] = 1
    @inbounds for col in 1:nblocks(A, 2)
        blockrange = nzblockrange(A, col)
        n_blocks_col = length(blockrange)
        nnz_values_col = n_blocks_col * blocksize(A, 1)
        for j_blk in 1:blocksize(A, 2)
            # The new colptr is the previous one plus the number of nonzero elements in this column.
            colptr[count_col] = colptr[count_col - 1] + nnz_values_col
            count_col += 1
            for block in blockrange
                i_offset = (A.rowval[block] - 1) * blocksize(A, 1)
                for i_blk in 1:blocksize(A, 1)
                    nzval[count_row] = A.nzval[i_blk, j_blk, block]
                    rowval[count_row] = i_blk + i_offset
                    count_row += 1
                end
            end
        end
    end
    return SparseMatrixCSC(A.m, A.n, colptr, rowval, nzval)
end

function SparseMatrixBSC(A::SparseMatrixCSC{Tv, Ti}, R::Integer, C::Integer) where {Tv, Ti}
    if (R, C) == (1,1)
        return SparseMatrixBSC(1, 1, A.m, A.n, A.colptr, A.rowval, reshape(A.nzval, length(A.nzval), 1, 1))
    end

    A.m % R != 0 && throw(ArgumentError("row block size: $(R) must evenly divide number of rows: $(A.m)"))
    A.n % C != 0 && throw(ArgumentError("column block size: $(C) must evenly divide number of rows: $(A.n)"))

    Anzval = A.nzval
    Arowval = A.rowval
    Acolptr = A.colptr

    # Upper bound of number of nonzero blocks is nnz(A).
    rowval = zeros(Ti, nnz(A))
    colptr = zeros(Ti, A.n ÷ C + 1)

    n_colblocks = div(A.n, C)

    colptr[1] = 1
    row_counter = 1
    rows = Int[]

    # Strategy to compute rowval and colptr:
    # For each column block accumulate all the rowvalues in the CSC matrix for that block.
    # Convert these to what rowblock they represent
    # Each unique rowblock should now be entered in order into rowval.
    # Colptr for this column block is incremented by the number of unique rowblock values.
    @inbounds for colblock in 1:n_colblocks
        j_offset = (colblock - 1) * C
        row_block_counter = 1

        # Count the number of non zero values for the columns in this column block
        nzvals_block = Acolptr[j_offset + C + 1] - Acolptr[j_offset + 1]
        if nzvals_block == 0 # No nz values in this column block, exit early
            colptr[colblock + 1] = row_counter
            continue
        end

        # Accumulate rowvals for this block
        resize!(rows, nzvals_block)
        for j_blk in 1:C
            col = j_offset + j_blk
            nz_range = Acolptr[col]:Acolptr[col + 1] - 1
            for r in nz_range
                rows[row_block_counter] = Arowval[r]
                row_block_counter += 1
            end
        end

        # Convert from row values -> block rows
        @simd for i in 1:length(rows)
            rows[i] = (rows[i] - 1) ÷ R + 1
        end

        # Pick out the unique block rows and put them into rowval

        #######################
        sort!(rows) # <- A bit of a bottle enck, takes about 30% of tot time.
        #######################
        rowval[row_counter] = rows[1] # We have at least one value in rows so this is ok
        row_counter += 1
        for i in 2:length(rows)
            if rows[i] > rows[i-1]
                rowval[row_counter] = rows[i]
                row_counter += 1
            end
        end
        colptr[colblock + 1] = row_counter
    end

    # We now know the true number of non zero blocks so we reshape rowval
    # and allocate the exact space we need for nzval
    deleteat!(rowval, row_counter:length(rowval))
    nzval = zeros(Tv, R, C, length(rowval))

    @inbounds for colblock in 1:n_colblocks
        j_offset = (colblock - 1) * C
        for j_blk in 1:C
            current_block = colptr[colblock]
            col = j_offset + j_blk
            for r in Acolptr[col]:Acolptr[col + 1] - 1
                row = Arowval[r]
                # Looking for the correct block for this column
                while row > rowval[current_block] * R
                    current_block += 1
                end
                i_blk = row - (rowval[current_block] - 1) * R
                nzval[i_blk, j_blk, current_block] = Anzval[r]
            end
        end
    end

    SparseMatrixBSC(A.m, A.n, colptr, rowval, nzval)
end

# for fast CSC conversion: convert to SparseMatrixCSC nzval ordering
function getCSCordering(A::SparseMatrixBSC)

    # loop through all active columns
    R,C = blocksize(A)
    ntotalblocks = size(A.nzval,3)
    flattened_indices = reshape(1:R*C*ntotalblocks,R,C,ntotalblocks) # UnitRange avoids allocation
    CSC_permuted_indices = zeros(Int,R,C*ntotalblocks) # [col1, col2, ..., col_blocksize*nblocks]

    sk = 1
    for i = 1:length(A.colptr)-1
        # find i,j indices of each block in a col
        blocks_in_column = nzblockrange(A,i)
        for col = 1:C
            for block_id in blocks_in_column
                CSC_permuted_indices[:,sk] .= flattened_indices[:,col,block_id]
                sk += 1
            end
        end
    end

    return CSC_permuted_indices[:]
end

# fast conversion from BSC to CSC matrix
function SparseMatrixCSC!(B::SparseMatrixCSC,A::SparseMatrixBSC{Tv, Ti}, CSCpermuted_indices) where {Tv, Ti <: Integer}
    B.nzval .= A.nzval[:][CSCpermuted_indices] # need to add permutation, but this should be super fast
end

############
# Base interfaces #
############

function Base.show(io::IO, A::SparseMatrixBSC;
                   header::Bool=true, repr=false)
    print(io, A.m, "×", A.n, " sparse block matrix with ", size(A.nzval,3), " ",
        eltype(A), " blocks of size ", blocksize(A, 1),"×",blocksize(A, 2))
end

# convert from block index to CSC index
function getCSCindex(A::SparseMatrixBSC{Tv,Ti}, index::Block{2,Ti}) where {Tv,Ti}
    row,col  = index.n
    colrange = nzblockrange(A,col)
    rows     = A.rowval[colrange]
    index    = findfirst(rows.==row)
    if isnothing(index)
        return nothing
    else
        return colrange[index]
    end
end

function Base.getindex(A::SparseMatrixBSC{Tv,Ti}, blockindex::Block{2,Ti}) where {Tv,Ti}
    CSC_id = getCSCindex(A,blockindex)
    if isnothing(CSC_id)
        return zeros(blocksize(A)...)
    else
        return A.nzval[:,:,CSC_id]
    end
end

function Base.setindex!(A::SparseMatrixBSC{Tv,Ti}, val::Array{Tv,2}, blockindex::Block{2,Ti}) where {Tv,Ti}
    CSC_id = getCSCindex(A,blockindex)
    if isnothing(CSC_id)
        error("block index not in original sparsity pattern. adding new sparse blocks not currently supported")
    end
    if size(val)!=blocksize(A)
        error("assigned block is size ", size(val), ", block size is ", blocksize(A))
    end
    A.nzval[:,:,CSC_id] .= val
end

# ##########
# # LinAlg #
# ##########
#
# # We do dynamic dispatch here so that the size of the blocks are known at compile time
# function LinearAlgebra.mul!(b::Vector{Tv}, A::SparseMatrixBSC{Tv}, x::Vector{Tv}) where {Tv}
#     if blocksize(A) == (1,1)
#         mul!(b, SparseMatrixCSC(A), x)
#     else
#         _mul!(b, A, x, Val(A.C), Val(A.R))
#     end
# end
#
# function _mul!(b::Vector{Tv}, A::SparseMatrixBSC{Tv}, x::Vector{Tv},
#                              ::Val{C}, ::Val{R}) where {Tv, C, R}
#     fill!(b, 0.0)
#     n_cb, n_rb = nblocks(A)
#     for J in 1:n_cb
#         j_offset = (J - 1)  * blocksize(A, 2)
#         for r in nzblockrange(A, J)
#             @inbounds i_offset = (A.rowval[r] - 1) * blocksize(A, 1)
#             matvec_kernel!(b, A, x, r, i_offset, j_offset, Val(C), Val(R))
#         end
#     end
#     return b
# end
#
# # TODO: Possibly SIMD.jl and do a bunch of simd magic coolness for small mat * vec
# @inline function matvec_kernel!(b::Vector{T}, A::SparseMatrixBSC{T}, x::Vector{T}, r,
#                                 i_offset, j_offset, ::Val{C}, ::Val{R}) where {C, R, T}
#     @inbounds for j in 1:C
#         for i in 1:R
#             b[i_offset + i] += A.nzval[i, j, r] * x[j_offset + j]
#         end
#     end
# end
#
# function Base.:*(A::SparseMatrixBSC{Tv}, x::Vector{Tv}) where {Tv}
#     mul!(similar(x, A.m), A, x)
# end


@testset "action interface" begin
    roadway = get_test_roadway()
    trajdata = get_test_trajdata(roadway)
    veh = trajdata[1][1]
    s = VehicleState()
    @test VehicleState() == propagate(veh, s, roadway, NaN)
end

@testset "AccelTurnrate" begin
    roadway = get_test_roadway()
    trajdata = get_test_trajdata(roadway)
    veh = trajdata[1][1]
    a = AccelTurnrate(0.1,0.2)
    io = IOBuffer()
    show(io, a)
    close(io)
    @test a == convert(AccelTurnrate, [0.1,0.2])
    @test copyto!([NaN, NaN], AccelTurnrate(0.1,0.2)) == [0.1,0.2]
    @test length(AccelTurnrate) == 2

    s = propagate(veh, AccelTurnrate(0.0,0.0), roadway, 1.0)
    @test isapprox(s.posG.x, veh.state.v)
    @test isapprox(s.posG.y, 0.0)
    @test isapprox(s.posG.θ, 0.0)

    # TODO Test the get method
end

@testset "AccelDesang" begin
    roadway = get_test_roadway()
    trajdata = get_test_trajdata(roadway)
    veh = trajdata[1][1]
    a = AccelDesang(0.1,0.2)
    @test a == convert(AccelDesang, [0.1,0.2])
    @test copyto!([NaN, NaN], AccelDesang(0.1,0.2)) == [0.1,0.2]
    @test length(AccelDesang) == 2

    io = IOBuffer()
    show(io, a)

    s = propagate(veh, AccelDesang(0.0,0.0), roadway, 1.0)
    @test isapprox(s.posG.x, veh.state.v*1.0)
    @test isapprox(s.posG.y, 0.0)
    @test isapprox(s.posG.θ, 0.0)
    @test isapprox(s.v, veh.state.v)
end

@testset "AccelSteeringAngle" begin
    roadway = get_test_roadway()
    trajdata = get_test_trajdata(roadway)
    veh = trajdata[1][1]
    a = AccelSteeringAngle(0.1,0.2)
    io = IOBuffer()
    show(io, a)
    close(io)
    @test a == convert(AccelSteeringAngle, [0.1,0.2])
    @test copyto!([NaN, NaN], AccelSteeringAngle(0.1,0.2)) == [0.1,0.2]
    @test length(AccelSteeringAngle) == 2

    # Check that propagate
    # won't work with the veh type loaded from test_tra
    @test_throws ErrorException propagate(veh, AccelSteeringAngle(0.0,0.0), roadway, 1.0)

    veh = Entity(veh.state, BicycleModel(veh.def), veh.id)
    s = propagate(veh, AccelSteeringAngle(0.0,0.0), roadway, 1.0)
    @test isapprox(s.posG.x, veh.state.v*1.0)
    @test isapprox(s.posG.y, 0.0)
    @test isapprox(s.posG.θ, 0.0)
    @test isapprox(s.v, veh.state.v)
    # Test the branch condition within propagate
    s = propagate(veh, AccelSteeringAngle(0.0,1.0), roadway, 1.0)
    @test isapprox(s.v, veh.state.v)
end

@testset "LaneFollowingAccel" begin
    a = LaneFollowingAccel(1.0)
    roadway = gen_straight_roadway(3, 100.0)
    s = VehicleState(VecSE2(0.0, 0.0, 0.0), roadway, 0.0)
    veh = Entity(s, VehicleDef(), 1)
    vehp = propagate(veh, a, roadway, 1.0)
    @test posf(vehp).s == 0.5
end

@testset "LatLonAccel" begin
    roadway = get_test_roadway()
    trajdata = get_test_trajdata(roadway)
    veh = trajdata[1][1]
    a = LatLonAccel(0.1,0.2)
    io = IOBuffer()
    show(io, a)
    close(io)
    @test a == convert(LatLonAccel, [0.1,0.2])
    @test copyto!([NaN, NaN], LatLonAccel(0.1,0.2)) == [0.1,0.2]
    @test length(LatLonAccel) == 2

    Δt = 0.1
    s = propagate(veh, LatLonAccel(0.0,0.0), roadway, Δt)
    @test isapprox(s.posG.x, veh.state.v*Δt)
    @test isapprox(s.posG.y, 0.0)
    @test isapprox(s.posG.θ, 0.0)
    @test isapprox(s.v, veh.state.v)
end

@testset "Pedestrian LatLon" begin
    roadway = get_test_roadway()
    trajdata = get_test_trajdata(roadway)
    veh = trajdata[1][1]
    a = PedestrianLatLonAccel(0.5,1.0, roadway[LaneTag(2,1)])
    Δt = 1.0
    s = propagate(veh, a, roadway, Δt)
    @test s.posF.roadind.tag == LaneTag(2,1)
    @test isapprox(s.posG.x, 4.0)
    @test isapprox(s.posG.y, 0.25)
end

@testset "Action Mapping" begin
    a1 = LaneFollowingAccel(1.2)
    a2 = LaneFollowingAccel(.4)
    a3 = LaneFollowingAccel(0.)
    actions = [a1, a2, a3]
    ids = [:vehA, :vehB, :vehC]
    action_mappings = Scene([EntityAction(a, id) for (a,id) in zip(actions, ids)])
    for (action,id) in zip(actions, ids)
        @test get_by_id(action_mappings, id).action.a == action.a
    end
end


"""
    couple_models(aero::LiftingLine, stru::GEBT, flap::LiftingLineFlaps)

Create a coupled model using a lifting line aerodynamic model, a
geometrically exact beam theory model, and a lifting line control surface model.
This model introduces additional parameters corresponding to the freestream
air density ``\\rho_\\infty``, followed by the external loads ``F_{x,i}, F_{y,i},
F_{z,i}, M_{x,i}, M_{y,i}, M_{z,i}`` or displacements ``u_{x,i}, u_{y,i}, u_{z,i},
\\theta_{x,i}, \\theta_{y,i}, \\theta_{z,i}`` for each node, followed by the constant
distributed loads ``f_{x,i}, f_{y,i}, f_{z,i}, m_{x,i}, m_{y,i}, m_{z,i}``
applied on each beam element (excluding aerodynamic loads), followed by the body
frame linear and angular velocities ``u, v, w, p, q, r`` and the control surface
deflections ``\\delta_1, \\delta_2, \\dots, \\delta_N``.

**NOTE: When using this model, the local frame for each beam element should be
oriented with the x-axis along the beam's axis, the y-axis forward, and the
z-axis normal to the surface**
"""
function couple_models(aero::LiftingLine, stru::GEBT, flap::LiftingLineFlaps)
    return (aero, stru, flap)
end

# --- traits --- #

function number_of_additional_parameters(aero::LiftingLine, stru::GEBT, flap::LiftingLineFlaps)

    return 1 + 6*length(stru.icol_point) + 6*length(stru.icol_elem) + 6 + length(flap.gains)
end

function coupling_inplaceness(::Type{<:LiftingLine}, ::Type{<:GEBT}, ::Type{<:LiftingLineFlaps})
    return InPlace()
end

function coupling_rate_jacobian_type(::Type{LiftingLine{NA,TA}}, ::Type{<:GEBT},
    ::Type{LiftingLineFlaps{NF,NG,TF,TG}}) where {NA,NF,NG,TA,TF,TG}

    aero_model_types = TA.parameters
    flap_model_types = TF.parameters
    if all(islinear.(coupling_rate_jacobian_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Linear()
    else
        return Nonlinear()
    end
end

function coupling_state_jacobian_type(::Type{LiftingLine{NA,TA}}, ::Type{<:GEBT},
    ::Type{LiftingLineFlaps{NF,NG,TF,TG}}) where {NA,NF,NG,TA,TF,TG}

    aero_model_types = TA.parameters
    flap_model_types = TF.parameters
    if all(islinear.(coupling_state_jacobian_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Linear()
    else
        return Nonlinear()
    end
end

coupling_parameter_jacobian_type(::Type{<:LiftingLine}, ::Type{<:GEBT},
    ::Type{LiftingLineFlaps}) = Nonlinear()

function coupling_time_gradient_type(::Type{LiftingLine{NA,TA}}, ::Type{<:GEBT},
    ::Type{LiftingLineFlaps{NF,NG,TF,TG}}) where {NA,NF,NG,TA,TF,TG}

    aero_model_types = TA.parameters
    flap_model_types = TF.parameters
    if all(isempty.(coupling_time_gradient_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Empty()
    elseif all(iszero.(coupling_time_gradient_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Zeros()
    elseif all(isconstant.(coupling_time_gradient_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Constant()
    elseif all(islinear.(coupling_time_gradient_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Linear()
    else
        return Nonlinear()
    end
end

# --- methods --- #

function get_coupling_inputs!(y, aero::LiftingLine{NA,TA}, stru::GEBT,
    flap::LiftingLineFlaps{NF,NG,TF,TG}, dx, x, p, t) where {NA,NF,NG,TA,TF,TG}

    # extract number of state variables, inputs, and parameters
    nua = number_of_states(aero) # number of aerodynamic states
    nus = number_of_states(stru) # number of structural states
    nuf = number_of_states(flap) # number of control surface states
    nya = number_of_inputs(aero) # number of aerodynamic inputs
    nys = number_of_inputs(stru) # number of structural inputs
    nyf = number_of_inputs(flap) # number of control surface inputs
    npa = number_of_parameters(aero) # number of aerodynamic parameters
    nps = number_of_parameters(stru) # number of structural parameters
    npf = number_of_parameters(flap) # number of control surface parameters
    npadd = number_of_additional_parameters(aero, stru, flap) # number of additional parameters

    # get indices for state variables, inputs, and parameters
    iua = 1:nua # indices of aerodynamic states
    iya = 1:nya # indices of aerodynamic inputs
    ipa = 1:npa # indices of aerodynamic parameters
    ius = nua + 1 : nua + nus # indices of structural states
    iys = nya + 1 : nya + nys # indices of structural inputs
    ips = npa + 1 : npa + nps # indices of structural parameters
    iuf = nua + nus + 1 : nua + nus + nuf # indices of control surface states
    iyf = nya + nys + 1 : nya + nys + nyf # indices of control surface inputs
    ipf = npa + nps + 1 : npa + nps + npf # indices of control surface parameters
    ipadd = npa + nps + npf + 1 : npa + nps + npf + npadd # indices of additional parameters
    iuas = state_indices(aero.models) # indices of aerodynamic states for each section
    iyas = input_indices(aero.models) # indices of aerodynamic inputs for each section
    ipas = parameter_indices(aero.models) # indices of aerodynamic parameters for each section
    iufs = state_indices(flap.models) # indices of control surface states for each section
    iyfs = input_indices(flap.models) # indices of control surface inputs for each section
    ipfs = parameter_indices(flap.models) # indices of control surface parameters for each section

    # separate state variables, inputs, and parameters
    dua = view(dx, iua) # aerodynamic rates
    ua = view(x, iua) # aerodynamic states
    ya = view(y, iya) # aerodynamic inputs
    pa = view(p, ipa) # aerodynamic parameters
    dus = view(dx, ius) # structural rates
    us = view(x, ius) # structural states
    ys = view(y, iys) # structural inputs
    ps = view(p, ips) # structural parameters
    duf = view(dx, iuf) # control surface rates
    uf = view(x, iuf) # control surface states
    yf = view(y, iyf) # control surface inputs
    pf = view(p, ipf) # control surface parameters
    padd = view(p, ipadd) # additional parameters for coupled model
    duas = view.(Ref(dua), iuas) # aerodynamic rates for each section
    uas = view.(Ref(ua), iuas) # aerodynamic states for each section
    yas = view.(Ref(ya), iyas) # aerodynamic inputs for each section
    pas = view.(Ref(pa), ipas) # aerodynamic parameters for each section
    dufs = view.(Ref(duf), iufs) # control surface rates for each section
    ufs = view.(Ref(uf), iufs) # control surface states for each section
    yfs = view.(Ref(yf), iyfs) # control surface inputs for each section
    pfs = view.(Ref(pf), ipfs) # control surface parameters for each section

    # number of points and elements
    npoint = length(stru.icol_point)
    nelem = length(stru.icol_elem)

    # global parameters
    ρ = padd[1]
    ur = padd[1 + 6*npoint + 6*nelem + 1]
    vr = padd[1 + 6*npoint + 6*nelem + 2]
    wr = padd[1 + 6*npoint + 6*nelem + 3]
    pr = padd[1 + 6*npoint + 6*nelem + 4]
    qr = padd[1 + 6*npoint + 6*nelem + 5]
    rr = padd[1 + 6*npoint + 6*nelem + 6]
    δ = padd[SVector{NG}(1 + 6*npoint + 6*nelem + 6 + 1 : 1 + 6*npoint + 6*nelem + 6 + NG)]
    dδ = zero(δ)

    # body linear and angular velocity
    V = SVector(ur, vr, wr)
    Ω = SVector(pr, qr, rr)

    # construct assembly from structural parameters
    assembly = gxbeam_assembly(ps, npoint, nelem, stru.start, stru.stop)

    # save prescribed point loads/displacements
    for ip = 1:npoint
        yoff = 6*(ip-1)
        poff = 1 + 6*(ip-1)
        y[iys[yoff+1]] = padd[poff+1]
        y[iys[yoff+2]] = padd[poff+2]
        y[iys[yoff+3]] = padd[poff+3]
        y[iys[yoff+4]] = padd[poff+4]
        y[iys[yoff+5]] = padd[poff+5]
        y[iys[yoff+6]] = padd[poff+6]
    end

    # save aerodynamic inputs and distributed loads
    for i = 1:NA
        # models for this section
        section_aero = aero.models[i]
        section_stru = LiftingLineSection()
        section_flap = flap.models[i]
        section_ctrl = LiftingLineSectionControl()
        section_models = (section_aero, section_stru, section_flap, section_ctrl)

        # model dimensions for this section
        Nuai = number_of_states(section_aero)
        Nyai = number_of_inputs(section_aero)
        Npai = number_of_parameters(section_aero)
        Nusi = number_of_states(section_stru)
        Nysi = number_of_inputs(section_stru)
        Npsi = number_of_parameters(section_stru)
        Nufi = number_of_states(section_flap)
        Nyfi = number_of_inputs(section_flap)
        Npfi = number_of_parameters(section_flap)
        Nuci = number_of_states(section_ctrl)
        Nyci = number_of_inputs(section_ctrl)
        Npci = number_of_parameters(section_ctrl)

        # local section properties
        icol = stru.icol_elem[i]
        element = assembly.elements[i]
        u_elem = SVector(us[icol], us[icol+1], us[icol+2]) # linear displacement
        θ_elem = SVector(us[icol+3], us[icol+4], us[icol+5]) # angular displacement
        F_elem = SVector(us[icol+6], us[icol+7], us[icol+8]) .* stru.force_scaling # internal force
        M_elem = SVector(us[icol+9], us[icol+10], us[icol+11]) .* stru.force_scaling # internal moment
        P_elem = SVector(us[icol+12], us[icol+13], us[icol+14]) .* stru.mass_scaling # linear momentum
        H_elem = SVector(us[icol+15], us[icol+16], us[icol+17]) .* stru.mass_scaling # angular momentum

        dF_elem = SVector(dus[icol+6], dus[icol+7], dus[icol+8]) .* stru.force_scaling # internal force
        dM_elem = SVector(dus[icol+9], dus[icol+10], dus[icol+11]) .* stru.force_scaling # internal moment
        dP_elem = SVector(dus[icol+12], dus[icol+13], dus[icol+14]) .* stru.mass_scaling # linear momentum
        dH_elem = SVector(dus[icol+15], dus[icol+16], dus[icol+17]) .* stru.mass_scaling # angular momentum

        # convert rotation parameter to Wiener-Milenkovic parameters
        scaling = GXBeam.rotation_parameter_scaling(θ_elem)
        θ_elem *= scaling # angular displacement (Wiener-Milenkovic parameters)

        # local to body transformation
        CtCab = GXBeam.get_C(θ_elem)'*element.Cab

        # local structural to aerodynamic transformation
        R = @SMatrix [0 -1 0; 1 0 0; 0 0 1]

        # local freestream linear and angular velocities
        vi = -R*(CtCab'*V + GXBeam.element_linear_velocity(element, P_elem, H_elem)) # local linear freestream velocity
        ωi = R*GXBeam.element_angular_velocity(element, P_elem, H_elem) # local angular freestream velocity

        # local freestream linear and angular accelerations
        dvi = -R*GXBeam.element_linear_velocity(element, dP_elem, dH_elem)
        dωi = R*GXBeam.element_angular_velocity(element, dP_elem, dH_elem)

        # NOTE: We assume that the local beam frame y-axis is oriented towards
        # the leading edge and the z-axis is oriented up
        # NOTE: We assume the local aerodynamic frame is oriented with the x-axis
        # in the chordwise direction and the z-axis in the (airfoil) normal
        # direction.

        # section state variable rates
        duai = SVector{Nuai}(duas[i]) # aerodynamic state variables
        dusi = vcat(dvi, dωi) # structural state variables
        dufi = SVector{Nufi}(dufs[i]) # control surface state variables
        duci = SMatrix{1,NG}(getindex.(flap.gains, i))* dδ # controller state variables
        dui = vcat(duai, dusi, dufi, duci)

        # section state variables
        uai = SVector{Nuai}(uas[i]) # aerodynamic state variables
        usi = vcat(vi, ωi) # structural state variables
        ufi = SVector{Nufi}(ufs[i]) # control surface state variables
        uci = SMatrix{1,NG}(getindex.(flap.gains, i)) * δ # controller state variables
        ui = vcat(uai, usi, ufi, uci)

        # section parameters
        pai = SVector{Npai}(pas[i]) # aerodynamic parameters
        psi = SVector(ρ) # structural parameters
        pfi = SVector{Npfi}(pfs[i]) # control surface parameters
        pci = SVector{0,Float64}() # controller parameters
        pi = vcat(pai, psi, pfi, pci)

        # section inputs
        yi = get_coupling_inputs(section_models, dui, ui, pi, t)

        # separate inputs
        yai = view(yi, 1:Nyai)
        ysi = view(yi, Nyai + 1 : Nyai + Nysi)
        yfi = view(yi, Nyai + Nysi + 1 : Nyai + Nysi + Nyfi)
        yci = view(yi, Nyai + Nysi + Nyfi + 1 : Nyai + Nysi + Nyfi + Nyci)

        # save local inputs
        y[iya[iyas[i]]] .= yai # aerodynamic inputs
        y[iyf[iyfs[i]]] .= yfi # control surface inputs

        # section aerodynamic loads (in body frame)
        fi = CtCab*R'*SVector(yi[Nyai+1], yi[Nyai+2], yi[Nyai+3])
        mi = CtCab*R'*SVector(yi[Nyai+4], yi[Nyai+5], yi[Nyai+6])

        # add constant distributed loads (in body frame)
        poff = 1 + 6*npoint + 6*(i-1)
        fi += SVector(padd[poff+1], padd[poff+2], padd[poff+3])
        mi += SVector(padd[poff+4], padd[poff+5], padd[poff+6])

        # save distributed loads for this element (in body frame)
        yoff = 6*npoint + 6*(i-1)
        y[iys[yoff+1 : yoff+3]] = fi
        y[iys[yoff+4 : yoff+6]] = mi
    end

    # save body frame linear/angular velocities
    yoff = 6*npoint + 6*nelem
    y[iys[yoff+1]] = ur
    y[iys[yoff+2]] = vr
    y[iys[yoff+3]] = wr
    y[iys[yoff+4]] = pr
    y[iys[yoff+5]] = qr
    y[iys[yoff+6]] = rr

    return y
end

# --- performance overloads --- #

# TODO

# --- convenience methods --- #

function set_additional_parameters!(padd, model::LiftingLine, stru::GEBT,
    flap::LiftingLineFlaps{NF,NG,TF,TG}; rho, point_conditions,
    element_loads, u, v, w, p, q, r, delta) where {NF,NG,TF,TG}

    np = length(stru.icol_point)
    ne = length(stru.icol_elem)

    padd[1] = rho

    for ip = 1:np
        padd[1+6*(ip-1)+1] = point_conditions[6*(ip-1)+1]
        padd[1+6*(ip-1)+2] = point_conditions[6*(ip-1)+2]
        padd[1+6*(ip-1)+3] = point_conditions[6*(ip-1)+3]
        padd[1+6*(ip-1)+4] = point_conditions[6*(ip-1)+4]
        padd[1+6*(ip-1)+5] = point_conditions[6*(ip-1)+5]
        padd[1+6*(ip-1)+6] = point_conditions[6*(ip-1)+6]
    end

    for ie = 1:ne
        padd[1+6*np+6*(ie-1)+1] = element_loads[6*(ie-1)+1]
        padd[1+6*np+6*(ie-1)+2] = element_loads[6*(ie-1)+2]
        padd[1+6*np+6*(ie-1)+3] = element_loads[6*(ie-1)+3]
        padd[1+6*np+6*(ie-1)+4] = element_loads[6*(ie-1)+4]
        padd[1+6*np+6*(ie-1)+5] = element_loads[6*(ie-1)+5]
        padd[1+6*np+6*(ie-1)+6] = element_loads[6*(ie-1)+6]
    end

    padd[1 + 6*np + 6*ne + 1] = u
    padd[1 + 6*np + 6*ne + 2] = v
    padd[1 + 6*np + 6*ne + 3] = w
    padd[1 + 6*np + 6*ne + 4] = p
    padd[1 + 6*np + 6*ne + 5] = q
    padd[1 + 6*np + 6*ne + 6] = r
    padd[1 + 6*np + 6*ne + 6 + 1 : 1 + 6*np + 6*ne + 6 + NG] = delta

    return padd
end

function separate_additional_parameters(model::LiftingLine, stru::GEBT,
    flap::LiftingLineFlaps{NF,NG,TF,TG}, padd) where {NF,NG,TF,TG}

    np = length(stru.icol_point)
    ne = length(stru.icol_elem)

    rho = padd[1]

    point_conditions = reshape(view(padd, 1 + 1 : 1 + 6*np), 6, np)

    element_loads = reshape(view(padd, 1 + 6*np + 1 : 1 + 6*np + 6*ne), 6, ne)

    u = padd[1 + 6*np + 6*ne + 1]
    v = padd[1 + 6*np + 6*ne + 2]
    w = padd[1 + 6*np + 6*ne + 3]
    p = padd[1 + 6*np + 6*ne + 4]
    q = padd[1 + 6*np + 6*ne + 5]
    r = padd[1 + 6*np + 6*ne + 6]
    delta = view(padd, 1 + 6*np + 6*ne + 6 + 1 : 1 + 6*np + 6*ne + 6 + NG)

    return (rho = rho, point_conditions = point_conditions, element_loads = element_loads,
        u = u, v = v, w = w, p = p, q = q, r = r, delta = delta)
end


################################################################################
#
#  Two element generated ideals
#
################################################################################

@doc Markdown.doc"""
    has_2_elem(A::NfAbsOrdIdl) -> Bool

Returns whether $A$ is generated by two elements.
"""
function has_2_elem(A::NfAbsOrdIdl)
  return isdefined(A, :gen_two)
end

@doc Markdown.doc"""
    has_weakly_normal(A::NfAbsOrdIdl) -> Bool

Returns whether $A$ has weakly normal two element generators.
"""
function has_weakly_normal(A::NfAbsOrdIdl)
  return (isdefined(A, :gens_weakly_normal) &&
        A.gens_weakly_normal == true) || has_2_elem_normal(A)
end

@doc Markdown.doc"""
    has_2_elem_normal(A::NfAbsOrdIdl) -> Bool

Returns whether $A$ has normal two element generators.
"""
function has_2_elem_normal(A::NfAbsOrdIdl)
  #the one ideal <1, ?> is automatomatically normal>
  return isdefined(A, :gens_normal) && (A.gen_one == 1 || A.gens_normal > 1)
end

################################################################################
#
#  Predicates
#
################################################################################

function isconsistent(A::NfAbsOrdIdl)
  if has_2_elem_normal(A) && !iszero(A.gen_two)
    if !defines_2_normal(A)
      return false
    end
  end
  if has_norm(A)
    b = basis_mat(A, copy = false)
    if det(b) != A.norm
      return false
    end
  end
  return true
end

# check if gen_one,gen_two is a P(gen_one)-normal presentation
# see Pohst-Zassenhaus p. 404
function defines_2_normal(A::NfAbsOrdIdl)
  m = A.gen_one
  gen = A.gen_two
  mg = denominator(inv(gen.elem_in_nf), order(A))
  # the minimum of ideal generated by g
  g = gcd(m,mg)
  return gcd(m, div(mg, g)) == 1
end

###########################################################################################
#
#  2-element normal presentation
#
###########################################################################################

# The following makes sure that A has a weakly normal presentation
# Recall that (x,y) are a weakly normal presentation for A
# if and only if norm(A) = gcd(norm(x), norm(y))
#
# Maybe we should allow an optional paramter (an fmpz),
# which should be the first generator.
# So far, the algorithm just samples (lifts of) random elements of A/m^2,
# where m is the minimum of A.

function _assure_weakly_normal_presentation(A::NfAbsOrdIdl)
  if has_2_elem(A) && has_weakly_normal(A)
    return nothing
  end

  if isdefined(A, :princ_gen)
    x = A.princ_gen
    b = x.elem_in_nf

    bi = inv(b)

    A.gen_one = denominator(bi, order(A))
    A.minimum = A.gen_one
    A.gen_two = x
    if !isdefined(A, :norm)
      A.norm = abs(_normmod(A.gen_one^degree(order(A)), b))
    end
    @hassert :NfOrd 1 gcd(A.gen_one^degree(order(A)),
                    FlintZZ(norm(A.gen_two))) == A.norm

    if A.gen_one == 1
      A.gens_normal = 2*A.gen_one
    else
      A.gens_normal = A.gen_one
    end
    A.gens_weakly_normal = 1
    return nothing
  end

  O = order(A)
  if minimum(A) == 0
    A.gen_one = minimum(A)
    A.gen_two = zero(O)
    A.gens_weakly_normal = true
    return nothing
  end

  if minimum(A) == 1
    A.gen_one = minimum(A)
    A.gen_two = one(O)
    A.gens_weakly_normal = true
    A.gens_normal = fmpz(2)
    return nothing
  end

  if !issimple(nf(order(A))) || !isdefining_polynomial_nice(nf(order(A)))
    B = Array{fmpz}(undef, degree(O))
    Amin2 = minimum(A)^2
    Amind = minimum(A)^degree(O)
    BB = basis(A, copy = false)
    r = -Amin2:Amin2

    while true
      rand!(B, r)
      gen = dot(B, BB)
      if norm(A) == gcd(Amind, norm(gen))
        A.gen_one = minimum(A)
        A.gen_two = O(gen, false)
        A.gens_weakly_normal = true
        return nothing
      end
    end
  end

  Amin2 = minimum(A)^2
  Amind = gcd(minimum(A)^degree(O), minimum(A)*norm(A))

  B = Array{fmpz}(undef, degree(O))

  gen = O()

  r = -Amin2:Amin2

  m = zero_matrix(FlintZZ, 1, degree(O))

  cnt = 0
  while true
    cnt += 1

    if cnt > 100 && is2_normal_difficult(A)
      assure_2_normal_difficult(A)
      return nothing
    end

    #if cnt > 1000
    #  println("Having a hard time find weak generators for $A")
    #end

    rand!(B, r)

    # Put the entries of B into the (1 x d)-Matrix m
    for i in 1:degree(O)
      s = ccall((:fmpz_mat_entry, :libflint), Ptr{fmpz}, (Ref{fmpz_mat}, Int, Int), m, 0, i - 1)
      ccall((:fmpz_set, :libflint), Nothing, (Ptr{fmpz}, Ref{fmpz}), s, B[i])
    end

    if iszero(m)
      continue
    end

    mul!(m, m, basis_mat(A, copy = false))
    d = denominator(basis_mat(O, copy = false))
    mul!(m, m, basis_mat(O, copy = false).num)
    gen = elem_from_mat_row(nf(O), m, 1, d)
    d = denominator(gen)
    f, e = ppio(d, minimum(A, copy = false))
    gen = mod(numerator(gen), f*minimum(A)^2)//f
    if iszero(gen)
      continue
    end

    # the following should be done inplace
    #gen = dot(reshape(Array(mm), degree(O)), basis(O))
    if norm(A) == _normmod(Amind, O(gen, false))#gcd(Amind, numerator(norm(gen)))
      A.gen_one = minimum(A)
      A.gen_two = O(gen, false)
      A.gens_weakly_normal = true
      return nothing
    end
  end
end
#
# Here some random guesses for the difficult 2-element generators
# degree  | d
#   < 7   | 1
#  8 - 12 | 2 * 3
# 13 - 20 | 2 * 3 * 5
#  >= 21  | 2 * 3 * 5 * 7

function is2_normal_difficult(A::NfAbsOrdIdl)
  d = 2
  m = minimum(A)
  ZK = order(A)
  n = degree(ZK)

  if n < 7
    return false
  end

  if n < 12 && isone(gcd(m, 2 * 3))
    return false
  end

  if n < 20 && isone(gcd(m, 2 * 3 * 5))
    return false
  end

  if isone(gcd(m, 2 * 3 * 5 * 7))
    return false
  end

  return true
end

function valuation_of_ideal_difficult(A::NfAbsOrdIdl, P::NfAbsOrdIdl)
  B = basis(A, copy = false)
  res = valuation(B[1], P)
  for i = 2:length(B)
    if iszero(res)
      break
    end
    res = min(res, valuation(B[i], P))
  end
  return res
end

function assure_2_normal_difficult(A::NfAbsOrdIdl)

  m = minimum(A)
  ZK = order(A)
  n = degree(ZK)

  if !is2_normal_difficult(A)
    assure_2_normal(A)
    return nothing
  end

  if n < 12
    d = Int[2, 3]
  elseif n < 20
    d = Int[2, 3, 5]
  else
    d = Int[2, 3, 5, 7]
  end

  I = ideal(ZK, 1)
  for i = 1:length(d)
    m1, m = ppio(m, fmpz(d[i]))
    if isone(m1)
      continue
    end
    A1 = gcd(A, m1)
    lp = prime_decomposition(ZK, d[i])
    v = Int[valuation_of_ideal_difficult(A1, p[1]) for p = lp]
    for i = 1:length(v)
      if v[i] > 0
        I *= lp[i][1]^v[i]
      end
    end
  end
  #Now the ideal I has a 2 normal presentation
  if isone(m)
    A.gen_one = I.gen_one
    A.gen_two = I.gen_two
    A.gens_normal = I.gens_normal
    A.gens_weakly_normal = I.gens_weakly_normal
    A.gens_short = I.gens_short
    return nothing
  end
  J = gcd(A, m)
  assure_2_normal(J)
  C = I * J
  A.gen_one = C.gen_one
  A.gen_two = C.gen_two
  A.gens_normal = C.gens_normal
  A.gens_weakly_normal = C.gens_weakly_normal
  A.gens_short = C.gens_short
  return nothing
end

################################################################################
#
#  Assure 2 normal
#
################################################################################

function assure_2_normal(A::NfAbsOrdIdl)
  if has_2_elem(A) && has_2_elem_normal(A)
    @hassert :NfOrd 1 isconsistent(A)
    return nothing
  end
  O = order(A)
  K = nf(O)
  n = degree(K)

  if norm(A) == 1
    A.gen_one = fmpz(1)
    A.gen_two = one(O)
    A.gens_normal = fmpz(1)
    return
  end

  if has_2_elem(A)
    m = minimum(A)
    bas = basis(O)
    # Magic constants
    if m > 1000
      r = -500:500
    else
      r = -div(Int(m)+1,2):div(Int(m)+1,2)
    end
    #gen = K()
    #s = K()
    gen = zero(O)
    s = O()
    cnt = 0
    while true
      cnt += 1
      if cnt > 100 && is2_normal_difficult(A)
        assure_2_normal_difficult(A)
        @hassert :NfOrd 1 isconsistent(A)
        return
      end
      if cnt > 1000
        error("Having a hard time making generators normal for $A")
      end
      rand!(s, O, r)
      mul!(s, s, A.gen_two)
      add!(gen, rand(r)*A.gen_one, gen)
      add!(gen, s, gen)
      gen = mod(gen, m^2)

      if iszero(gen)
        continue
      end

      mg = _minmod(m^2, gen)
      g = gcd(m, mg)
      if gcd(m, div(mg, g)) == 1
        if gcd(m^n, norm(gen)) != norm(A)
          @vprint :NfOrd 2 "\n\noffending ideal $A \ngen is $gen\nWrong ideal\n"
          cnt += 10
          continue
        end
        break
      end
    end
    @vprint :NfOrd 2 "used $cnt attempts\n"
    A.gen_one = m
    A.gen_two = gen
    A.gens_normal = m
    @hassert :NfOrd 1 isconsistent(A)
    return
  end

  m = minimum(A)
  bas = nf_elem[ elem_in_nf(b) for b in basis(A, copy = false)]
  # Magic constants
  if m > 1000
    r = -500:500
  else
    r = -div(Int(m)+1,2):div(Int(m)+1,2)
  end
  cnt = 0
  while true
    cnt += 1
    if cnt > 100 && is2_normal_difficult(A)
      assure_2_normal_difficult(A)
      @hassert :NfOrd 1 isconsistent(A)
      return
    end
    if cnt > 1000
      error("Having a hard time making generators normal for $A")
    end
    gen = O(rand(bas, r))
    gen = mod(gen, m^2)

    if iszero(gen)
      continue
    end

    mg = _minmod(m^2, gen)
    if gcd(m, div(mg, gcd(mg, m))) == 1
      if gcd(m^n, norm(gen)) != norm(A)
        @vprint :NfOrd 1 "\n\noffending ideal $A \ngen is $gen\nWrong ideal\n"
        cnt += 10
        continue
      end
      break
    end
  end
  @vprint :NfOrd 2 "used $cnt attempts\n"
  A.gen_one = m
  A.gen_two = gen
  A.gens_normal = m
  @hassert :NfOrd 1 isconsistent(A)
  return nothing
end


##########################################################################################
#
#    MCMC specific derivation rules
#
##########################################################################################

####### creates multiple rules at once for logpdf(Distrib, x)
# FIXME : using undocumented dprefix function of ReverseDiffSource (should be replaced)
macro dlogpdfd(dist::Symbol, rule)
	sig = :( logpdf($(Expr(:(::), :d, dist)), x::Real) )
	ReverseDiffSource.deriv_rule( sig, :d, rule )

	sig = :( logpdf($(Expr(:(::), :d, dist)), x::AbstractArray) )
	rule2 = ReverseDiffSource.substSymbols(rule, {:x => :(x[i]), :ds => :(ds[i])})
	ReverseDiffSource.deriv_rule( sig, :d, :(for i in 1:length(x) ; $rule2 ; end))

	sig = :( logpdf($(Expr(:(::), :d, Expr(:curly, :Array, dist))), x::AbstractArray) )
	rule2 = ReverseDiffSource.substSymbols(rule, {:dd1 => :(dd1[i]), :dd2 => :(dd2[i]), :dd3 => :(dd3[i]),
		:x => :(x[i]), :ds => :(ds[i]), :d => :(d[i]) })
	ReverseDiffSource.deriv_rule(sig, :d, :(for i in 1:length(x) ; $rule2 ; end))
end

macro dlogpdfx(dist::Symbol, rule)
	sig = :( logpdf($(Expr(:(::), :d, dist)), x::Real) )
	ReverseDiffSource.deriv_rule( sig, :x, rule )

	sig = :( logpdf($(Expr(:(::), :d, dist)), x::AbstractArray) )
	rule2 = ReverseDiffSource.substSymbols(rule, {:dx => :(dx[i]), :x => :(x[i]), :ds => :(ds[i])})
	ReverseDiffSource.deriv_rule( sig, :x, :(for i in 1:length(x) ; $rule2 ; end))

	sig = :( logpdf($(Expr(:(::), :d, Expr(:curly, :Array, dist))), x::AbstractArray) )
	rule3 = ReverseDiffSource.substSymbols(rule2, {:d => :(d[i])})
	ReverseDiffSource.deriv_rule( sig, :x, :(for i in 1:length(x) ; $rule3 ; end))
end

####### derivation for Distribution types constructors
ReverseDiffSource.declareType(Distribution, :Distribution)

for d in [:Bernoulli, :TDist, :Exponential, :Poisson]
	ReverseDiffSource.declareType(eval(d), d)

	ReverseDiffSource.deriv_rule(:( ($d)(p::Real) ), :p, :( dp = ds1 ))
	ReverseDiffSource.deriv_rule(:( ($d)(p::AbstractArray) ), :p, :( copy!(dp, ds1) ))
end

for d in [ :Normal, :Uniform, :Weibull, :Gamma, :Cauchy, :LogNormal, :Binomial, :Beta, :Laplace]
	ReverseDiffSource.declareType(eval(d), d)

	ReverseDiffSource.deriv_rule(:( ($d)(p1::Real, p2::Real) ),   :p1, :( dp1 = ds1 ) )
	ReverseDiffSource.deriv_rule(:( ($d)(p1::Real, p2::Real) ),   :p2, :( dp2 = ds2 ) )
	ReverseDiffSource.deriv_rule(:( ($d)(p1::AbstractArray, p2::AbstractArray) ), :p1, :( copy!(dp1, ds1) ) )
	ReverseDiffSource.deriv_rule(:( ($d)(p1::AbstractArray, p2::AbstractArray) ), :p2, :( copy!(dp2, ds2) ) )
end

#######   Normal distribution
@dlogpdfx Normal dx += (d.μ - x) / (d.σ * d.σ) * ds
@dlogpdfd Normal ( 	dd1 += (x - d.μ) / (d.σ*d.σ) * ds;
					dd2 += ((x - d.μ)*(x - d.μ) / (d.σ*d.σ) - 1.) / d.σ * ds )

## Uniform distribution
@dlogpdfx Uniform dx += 0.
@dlogpdfd Uniform ( dd1 += (d.a <= x <= d.b) / (d.b - d.a) * ds ;
					dd2 += (d.a <= x <= d.b) / (d.a - d.b) * ds )

## Weibull distribution
@dlogpdfd Weibull   ( 	dd1 += ((1. - (x/d.scale)^d.shape) * log(x/d.scale) + 1./d.shape) * ds ;
						dd2 += ((x/d.scale)^d.shape - 1.) * d.shape/d.scale * ds )
@dlogpdfx Weibull   dx += ((1. - (x/d.scale)^d.shape) * d.shape - 1.) / x * ds

## Beta distribution
@dlogpdfd Beta   ( dd1 += (digamma(d.alpha+d.beta) - digamma(d.alpha) + log(x)) * ds ;
				    dd2 += (digamma(d.alpha+d.beta) - digamma(d.beta) + log(1-x)) * ds )
@dlogpdfx Beta   dx += ((d.alpha-1) / x - (d.beta-1)/(1-x)) * ds

## TDist distribution
@dlogpdfd TDist   dd1 += ((x*x-1)/(x*x + d.df)+log(d.df/(x*x+d.df))+digamma((d.df+1)/2)-digamma(d.df/2))/2 * ds
@dlogpdfx TDist   dx += (-(d.df+1)*x / (d.df+x*x)) * ds

## Exponential distribution
@dlogpdfd Exponential   dd1 += (x-d.scale) / (d.scale*d.scale) * ds
@dlogpdfx Exponential   dx -= ds / d.scale

## Gamma distribution
@dlogpdfd Gamma   ( dd1 += (log(x) - log(d.scale) - digamma(d.shape)) * ds ;
					dd2 += ((x - d.scale*d.shape) / (d.scale*d.scale)) * ds )
@dlogpdfx Gamma   dx += (-( d.scale + x - d.shape*d.scale)/(d.scale*x)) * ds

## Cauchy distribution
@dlogpdfd Cauchy   ( dd1 += (2(x-d.location) / (d.scale*d.scale + (x-d.location)*(x-d.location))) * ds ;
					 dd2 += (((x-d.location)*(x-d.location) - d.scale*d.scale) / (d.scale*(d.scale*d.scale + (x-d.location)*(x-d.location)))) * ds )
@dlogpdfx Cauchy   dx += (2(d.location-x) / (d.scale*d.scale + (x-d.location)*(x-d.location))) * ds

## Log-normal distribution
@dlogpdfd LogNormal   ( dd1 += (log(x) - d.meanlog) / (d.sdlog*d.sdlog) * ds ;
					 	dd2 += (d.meanlog*d.meanlog - d.sdlog*d.sdlog - log(x)*(2d.meanlog-log(x))) / (d.sdlog*d.sdlog*d.sdlog) * ds )
@dlogpdfx LogNormal   dx += (d.meanlog - d.sdlog*d.sdlog - log(x)) / (d.sdlog*d.sdlog*x) * ds

## Laplace distribution
@dlogpdfd Laplace   ( dd1 += (x > d.location ? 1 : -1) / d.scale * ds ;  # location
	                  dd2 += ( abs(x-d.location)/d.scale - 2. ) /d.scale ) # scale
@dlogpdfx Laplace   dx += (x > d.location ? -1 : 1) / d.scale * ds

# Note : vectorization will not be easily possible for multi variate distribs (Dirichlet, Categorical)
# TODO : add other continuous distribs ? : Pareto, Rayleigh, Logistic, Levy, Dirichlet, FDist
# TODO : add other discrete distribs ? : NegativeBinomial, DiscreteUniform, HyperGeometric, Geometric, Categorical

## Bernoulli distribution (Note : no derivation on x parameter as it is an integer)
@dlogpdfd Bernoulli     dd1 += 1. / (d.p1 - 1. + x) * ds

## Binomial distribution (Note : no derivation on x and n parameters as they are integers)
@dlogpdfd Binomial      dd2 += (x / d.prob - (d.size-x) / (1 - d.prob)) * ds

## Poisson distribution (Note : no derivation on x parameter as it is an integer)
@dlogpdfd Poisson       dd1 += (x / d.lambda - 1) * ds


using Conda, PyCall, Logging
try
	Conda.add("numpy")
	Conda.add("scipy")
	Conda.add("cython")
	Conda.add("pip")
	run(`$(PyCall.pyprogramname) -m pip install --user ttpy`)
	@info "Python TT Toolbox installed"
catch;
	@warn "Python TT Toolbox installation failed; the TTPy submodule will not be usable"
end


# *********************************************************************************
# REopt, Copyright (c) 2019-2020, Alliance for Sustainable Energy, LLC.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
# Redistributions of source code must retain the above copyright notice, this list
# of conditions and the following disclaimer.
#
# Redistributions in binary form must reproduce the above copyright notice, this
# list of conditions and the following disclaimer in the documentation and/or other
# materials provided with the distribution.
#
# Neither the name of the copyright holder nor the names of its contributors may be
# used to endorse or promote products derived from this software without specific
# prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
# OF THE POSSIBILITY OF SUCH DAMAGE.
# *********************************************************************************
"""
data for electric tariff in reopt model
    can be defined using custom rates or URDB rate
"""
# TODO function for creating BAU inputs: don't need for tariff b/c Tech repeats no longer needed?
struct ElectricTariff
    energy_rates::Array{Float64,1}

    monthly_demand_rates::Array{Float64,1}
    time_steps_monthly::Array{Array{Int64,1},1}  # length = 0 or 12

    tou_demand_rates::Array{Float64,1}
    tou_demand_ratchet_timesteps::Array{Array{Int64,1},1}  # length = n_tou_demand_ratchets

    fixed_monthly_charge::Float64
    annual_min_charge::Float64
    min_monthly_charge::Float64

    export_rates::DenseAxisArray{Array{Float64,1}}
    export_bins::Array{Symbol,1}
end

"""
    ElectricTariff

    function ElectricTariff(;
        urdb_label::String="",
        urdb_response::Dict=Dict(),
        year::Int=2020,
        time_steps_per_hour::Int=1,
        NEM::Bool=false,
        wholesale_rate::T=nothing,
        monthly_energy_rates::Array=[],
        monthly_demand_rates::Array=[],
    ) where {T <: Union{Nothing, Int, Float64, Array}, S <: Union{Nothing, Int, Float64, Array}}

"""
function ElectricTariff(;
    urdb_label::String="",
    urdb_response::Dict=Dict(),
    year::Int=2020,
    time_steps_per_hour::Int=1,
    NEM::Bool=false,
    wholesale_rate::T=nothing,
    monthly_energy_rates::Array=[],
    monthly_demand_rates::Array=[],
    ) where {T <: Union{Nothing, Int, Float64, Array}, S <: Union{Nothing, Int, Float64, Array}}

    nem_rate = [100.0 for _ in 1:8760*time_steps_per_hour]

    u = nothing
    if !isempty(urdb_label)

        u = URDBrate(urdb_label, year, time_steps_per_hour=time_steps_per_hour)

    elseif !isempty(urdb_response)

        u = URDBrate(urdb_response, year, time_steps_per_hour=time_steps_per_hour)

    elseif !isempty(monthly_energy_rates) && !isempty(monthly_demand_rates)

        invalid_args = String[]
        if !(length(monthly_energy_rates) == 12)
            push!(invalid_args, "length(monthly_energy_rates) must equal, got length $(length(monthly_energy_rates))")
        end
        if !(length(monthly_demand_rates) == 12)
            push!(invalid_args, "length(monthly_demand_rates) must equal, got length $(length(monthly_demand_rates))")
        end
        if length(invalid_args) > 0
            error("Invalid argument values: $(invalid_args)")
        end

        tou_demand_rates = Float64[]
        tou_demand_ratchet_timesteps = []
        demand_rates_monthly = monthly_demand_rates
        time_steps_monthly = get_monthly_timesteps(year, time_steps_per_hour=time_steps_per_hour)
        energy_rates = Real[]
        for m in 1:12
            append!(energy_rates, [monthly_energy_rates[m] for ts in time_steps_monthly[m]])
        end

        fixed_monthly_charge = 0.0
        annual_min_charge = 0.0
        min_monthly_charge = 0.0

        if NEM
            nem_rate = [-0.999 * x for x in energy_rates]
        end
    else
        error("Creating ElectricTariff requires at least urdb_label or monthly rates.")
    end

    if !isnothing(u)

        if NEM
            t = get_tier_with_lowest_energy_rate(u)
            nem_rate = [-0.999 * x for x in u.energy_rates[t,:]]
        end

        energy_rates, monthly_demand_rates, tou_demand_rates = remove_tiers_from_urdb_rate(u)
        time_steps_monthly = Array[]
        if !isempty(u.monthly_demand_rates)
            time_steps_monthly =
                get_monthly_timesteps(year, time_steps_per_hour=time_steps_per_hour)
        end

        tou_demand_ratchet_timesteps = u.tou_demand_ratchet_timesteps
        fixed_monthly_charge = u.fixed_monthly_charge
        annual_min_charge = u.annual_min_charge
        min_monthly_charge = u.min_monthly_charge
    end

    #= export_rates
    3 "tiers": 1. NEM (Net Energy Metering), 2. WHL (Wholesale), 3. CUR (Curtail)
    - if NEM then set ExportRate[:Nem, :] to energy_rate[tier_with_lowest_energy_rate, :]
        - otherwise set to 100 dollars/kWh
    - user can provide either scalar wholesale rate or vector of timesteps,
        - otherwise set to 100 dollars/kWh
    - curtail cost set to zero by default, but can be specified same as wholesale rate
    =#
    whl_rate = create_export_rate(wholesale_rate, length(energy_rates), time_steps_per_hour)

    if !NEM & (sum(whl_rate) >= 0)
        export_rates = DenseAxisArray{Array{Float64,1}}(undef, [])
        export_bins = Symbol[]
    elseif !NEM
        export_bins = [:WHL]
        export_rates = DenseAxisArray([whl_rate], export_bins)
    elseif (sum(whl_rate) >= 0)
        export_bins = [:NEM]
        export_rates = DenseAxisArray([nem_rate], export_bins)
    else
        export_bins = [:NEM, :WHL]  # TODO add :EXC to align with REopt Lite API
        export_rates = DenseAxisArray([nem_rate, whl_rate], export_bins)
    end

    ElectricTariff(
        energy_rates,
        monthly_demand_rates,
        time_steps_monthly,
        tou_demand_rates,
        tou_demand_ratchet_timesteps,
        fixed_monthly_charge,
        annual_min_charge,
        min_monthly_charge,
        export_rates,
        export_bins
    )
end

function get_tier_with_lowest_energy_rate(u::URDBrate)
    """
    ExportRate should be lowest energy cost for tiered rates.
    Otherwise, ExportRate can be > FuelRate, which leads REopt to export all PV energy produced.
    """
    tier_with_lowest_energy_cost = 1
    if length(u.energy_tier_limits) > 1
        annual_energy_charge_sums = Float64[]
        for etier in u.energy_rates
            push!(annual_energy_charge_sums, sum(etier))
        end
        tier_with_lowest_energy_cost =
            findall(annual_energy_charge_sums .== minimum(annual_energy_charge_sums))[1]
    end
    return tier_with_lowest_energy_cost
end

# TODO: dispatch custom rates based on options: TOU, monthly, etc.
function CustomRate(flat_energy::Real, flat_demand::Real=0) end

"""
    function create_export_rate(e::Nothing, N::Int, ts_per_hour::Int=1)
No export rate provided by user: set to 100 dollars/kWh for all time
"""
function create_export_rate(e::Nothing, N::Int, ts_per_hour::Int=1)
    [100 for _ in range(1, stop=N) for ts in 1:ts_per_hour]
end

"""
    function create_export_rate(e::T, N::Int, ts_per_hour::Int=1) where T<:Real
Case for scaler export rate provided -> convert to array of timesteps
"""
function create_export_rate(e::T, N::Int, ts_per_hour::Int=1) where T<:Real
    [float(-1*e) for ts in range(1, stop=N) for ts_per_hour::Int=1]
end

"""
    function create_export_rate(e::Array{<:Real, 1}, N::Int, ts_per_hour::Int=1)

Check length of e and upsample if length(e) != N
"""
function create_export_rate(e::Array{<:Real, 1}, N::Int, ts_per_hour::Int=1)
    Ne = length(e)
    if Ne != Int(N/ts_per_hour) || Ne != N
        @error "Export rates do not have correct number of entries. Must be $(N) or $(Int(N/ts_per_hour))."
    end
    if Ne != N  # upsample
        export_rates = [-1*x for x in e for ts in 1:ts_per_hour]
    else
        export_rates = -1*e
    end
    return export_rates
end

"""
    get_monthly_timesteps(year::Int; time_steps_per_hour=1)

return Array{Array{Int64,1},1}, size = (12,)
"""
function get_monthly_timesteps(year::Int; time_steps_per_hour=1)
    a = Array[]
    i = 1
    for m in range(1, stop=12)
        n_days = daysinmonth(Date(string(year) * "-" * string(m)))
        stop = n_days * 24 * time_steps_per_hour + i - 1
        if m == 2 && isleapyear(year)
            stop -= 24 * time_steps_per_hour  # TODO support extra day in leap years?
        end
        steps = [step for step in range(i, stop=stop)]
        append!(a, [steps])
        i = stop + 1
    end
    return a
end

# TODO use this function only for URDBrate
function remove_tiers_from_urdb_rate(u::URDBrate)
    # tariff args: have to validate that there are no tiers
    if length(u.energy_tier_limits) > 1
        @warn "Energy rate contains tiers. Using the first tier!"
    end
    elec_rates = vec(u.energy_rates[1,:])

    if u.n_monthly_demand_tiers > 1
        @warn "Monthly demand rate contains tiers. Using the last tier!"
    end
    if u.n_monthly_demand_tiers > 0
        demand_rates_monthly = vec(u.monthly_demand_rates[:,u.n_monthly_demand_tiers])
    else
        demand_rates_monthly = vec(u.monthly_demand_rates)  # 0×0 Array{Float64,2}
    end

    if u.n_tou_demand_tiers > 1
        @warn "TOU demand rate contains tiers. Using the last tier!"
    end
    if u.n_tou_demand_tiers > 0
        demand_rates = vec(u.tou_demand_rates[:,u.n_tou_demand_tiers])
    else
        demand_rates = vec(u.tou_demand_rates)
    end

    return elec_rates, demand_rates_monthly, demand_rates
end