import RayCaster as RC
import Test

Test.@testset "RayCaster.jl" begin
    obstacle_tile_map = BitArray([
                                  1 1 1 1 1
                                  1 0 0 0 1
                                  1 0 0 0 1
                                  1 0 0 0 1
                                  1 1 1 1 1
                                 ])

    height_obstacle_tile_map = size(obstacle_tile_map, 1)
    width_obstacle_tile_map = size(obstacle_tile_map, 2)

    Test.@testset "Continuous world" begin
        T1 = Float32
        T2 = Float32

        x_start = convert(T1, 5 / 2)
        y_start = convert(T1, 5 / 2)

        Test.@testset "theta = 0" begin
            theta = convert(T2, 0)
            cos_theta = cos(theta)
            sin_theta = sin(theta)
            i_stop, j_stop, hit_dimension, delta_x_world_units_to_exit_start_tile, delta_y_world_units_to_exit_start_tile = RC.cast_ray_continous_world(obstacle_tile_map, x_start, y_start, cos_theta, sin_theta)
            Test.@test i_stop == 5
            Test.@test j_stop == 3
            Test.@test hit_dimension == 1
            # Test.@test total_euclidean ≈ 1.5
            Test.@test delta_x_world_units_to_exit_start_tile ≈ 0.5
            Test.@test delta_y_world_units_to_exit_start_tile ≈ 0.5
        end

        Test.@testset "theta = atan(2 / 3)" begin
            theta = convert(T2, atan(2 / 3))
            cos_theta = cos(theta)
            sin_theta = sin(theta)
            i_stop, j_stop, hit_dimension, delta_x_world_units_to_exit_start_tile, delta_y_world_units_to_exit_start_tile = RC.cast_ray_continous_world(obstacle_tile_map, x_start, y_start, cos_theta, sin_theta)
            Test.@test i_stop == 5
            Test.@test j_stop == 4
            Test.@test hit_dimension == 1
            # Test.@test total_euclidean ≈ sqrt(1.5 ^ 2 + 1 ^ 2)
            Test.@test delta_x_world_units_to_exit_start_tile ≈ 0.5
            Test.@test delta_y_world_units_to_exit_start_tile ≈ 0.5
        end

        Test.@testset "theta = (pi / 2) - atan(2 / 3)" begin
            theta = convert(T2, (pi / 2) - atan(2 / 3))
            cos_theta = cos(theta)
            sin_theta = sin(theta)
            i_stop, j_stop, hit_dimension, delta_x_world_units_to_exit_start_tile, delta_y_world_units_to_exit_start_tile = RC.cast_ray_continous_world(obstacle_tile_map, x_start, y_start, cos_theta, sin_theta)
            Test.@test i_stop == 4
            Test.@test j_stop == 5
            Test.@test hit_dimension == 2
            # Test.@test total_euclidean ≈ sqrt(1.5 ^ 2 + 1 ^ 2)
            Test.@test delta_x_world_units_to_exit_start_tile ≈ 0.5
            Test.@test delta_y_world_units_to_exit_start_tile ≈ 0.5
        end

        Test.@testset "theta = pi / 2" begin
            theta = convert(T2, pi / 2)
            cos_theta = cos(theta)
            sin_theta = sin(theta)
            i_stop, j_stop, hit_dimension, delta_x_world_units_to_exit_start_tile, delta_y_world_units_to_exit_start_tile = RC.cast_ray_continous_world(obstacle_tile_map, x_start, y_start, cos_theta, sin_theta)
            Test.@test i_stop == 3
            Test.@test j_stop == 5
            Test.@test hit_dimension == 2
            # Test.@test total_euclidean ≈ 1.5
            Test.@test delta_x_world_units_to_exit_start_tile ≈ 0.5
            Test.@test delta_y_world_units_to_exit_start_tile ≈ 0.5
        end

        Test.@testset "theta = (pi / 2) + atan(2 / 3)" begin
            theta = convert(T2, (pi / 2) + atan(2 / 3))
            cos_theta = cos(theta)
            sin_theta = sin(theta)
            i_stop, j_stop, hit_dimension, delta_x_world_units_to_exit_start_tile, delta_y_world_units_to_exit_start_tile = RC.cast_ray_continous_world(obstacle_tile_map, x_start, y_start, cos_theta, sin_theta)
            Test.@test i_stop == 2
            Test.@test j_stop == 5
            Test.@test hit_dimension == 2
            # Test.@test total_euclidean ≈ sqrt(1.5 ^ 2 + 1 ^ 2)
            Test.@test delta_x_world_units_to_exit_start_tile ≈ 0.5
            Test.@test delta_y_world_units_to_exit_start_tile ≈ 0.5
        end

        Test.@testset "theta = pi - atan(2 / 3)" begin
            theta = convert(T2, pi - atan(2 / 3))
            cos_theta = cos(theta)
            sin_theta = sin(theta)
            i_stop, j_stop, hit_dimension, delta_x_world_units_to_exit_start_tile, delta_y_world_units_to_exit_start_tile = RC.cast_ray_continous_world(obstacle_tile_map, x_start, y_start, cos_theta, sin_theta)
            Test.@test i_stop == 1
            Test.@test j_stop == 4
            Test.@test hit_dimension == 1
            # Test.@test total_euclidean ≈ sqrt(1.5 ^ 2 + 1 ^ 2)
            Test.@test delta_x_world_units_to_exit_start_tile ≈ 0.5
            Test.@test delta_y_world_units_to_exit_start_tile ≈ 0.5
        end

        Test.@testset "theta = pi" begin
            theta = convert(T2, pi)
            cos_theta = cos(theta)
            sin_theta = sin(theta)
            i_stop, j_stop, hit_dimension, delta_x_world_units_to_exit_start_tile, delta_y_world_units_to_exit_start_tile = RC.cast_ray_continous_world(obstacle_tile_map, x_start, y_start, cos_theta, sin_theta)
            Test.@test i_stop == 1
            Test.@test j_stop == 3
            Test.@test hit_dimension == 1
            # Test.@test total_euclidean ≈ 1.5
            Test.@test delta_x_world_units_to_exit_start_tile ≈ 0.5
            Test.@test delta_y_world_units_to_exit_start_tile ≈ 0.5
        end

        Test.@testset "theta = pi + atan(2 / 3)" begin
            theta = convert(T2, pi + atan(2 / 3))
            cos_theta = cos(theta)
            sin_theta = sin(theta)
            i_stop, j_stop, hit_dimension, delta_x_world_units_to_exit_start_tile, delta_y_world_units_to_exit_start_tile = RC.cast_ray_continous_world(obstacle_tile_map, x_start, y_start, cos_theta, sin_theta)
            Test.@test i_stop == 1
            Test.@test j_stop == 2
            Test.@test hit_dimension == 1
            # Test.@test total_euclidean ≈ sqrt(1.5 ^ 2 + 1 ^ 2)
            Test.@test delta_x_world_units_to_exit_start_tile ≈ 0.5
            Test.@test delta_y_world_units_to_exit_start_tile ≈ 0.5
        end

        Test.@testset "theta = (3 * pi / 2) - atan(2 / 3)" begin
            theta = convert(T2, (3 * pi / 2) - atan(2 / 3))
            cos_theta = cos(theta)
            sin_theta = sin(theta)
            i_stop, j_stop, hit_dimension, delta_x_world_units_to_exit_start_tile, delta_y_world_units_to_exit_start_tile = RC.cast_ray_continous_world(obstacle_tile_map, x_start, y_start, cos_theta, sin_theta)
            Test.@test i_stop == 2
            Test.@test j_stop == 1
            Test.@test hit_dimension == 2
            # Test.@test total_euclidean ≈ sqrt(1.5 ^ 2 + 1 ^ 2)
            Test.@test delta_x_world_units_to_exit_start_tile ≈ 0.5
            Test.@test delta_y_world_units_to_exit_start_tile ≈ 0.5
        end

        Test.@testset "theta = 3 * pi / 2" begin
            theta = convert(T2, 3 * pi / 2)
            cos_theta = cos(theta)
            sin_theta = sin(theta)
            i_stop, j_stop, hit_dimension, delta_x_world_units_to_exit_start_tile, delta_y_world_units_to_exit_start_tile = RC.cast_ray_continous_world(obstacle_tile_map, x_start, y_start, cos_theta, sin_theta)
            Test.@test i_stop == 3
            Test.@test j_stop == 1
            Test.@test hit_dimension == 2
            # Test.@test total_euclidean ≈ 1.5
            Test.@test delta_x_world_units_to_exit_start_tile ≈ 0.5
            Test.@test delta_y_world_units_to_exit_start_tile ≈ 0.5
        end

        Test.@testset "theta = (3 * pi / 2) + atan(2 / 3)" begin
            theta = convert(T2, (3 * pi / 2) + atan(2 / 3))
            cos_theta = cos(theta)
            sin_theta = sin(theta)
            i_stop, j_stop, hit_dimension, delta_x_world_units_to_exit_start_tile, delta_y_world_units_to_exit_start_tile = RC.cast_ray_continous_world(obstacle_tile_map, x_start, y_start, cos_theta, sin_theta)
            Test.@test i_stop == 4
            Test.@test j_stop == 1
            Test.@test hit_dimension == 2
            # Test.@test total_euclidean ≈ sqrt(1.5 ^ 2 + 1 ^ 2)
            Test.@test delta_x_world_units_to_exit_start_tile ≈ 0.5
            Test.@test delta_y_world_units_to_exit_start_tile ≈ 0.5
        end

        Test.@testset "theta = 2 * pi - atan(2 / 3)" begin
            theta = convert(T2, 2 * pi - atan(2 / 3))
            cos_theta = cos(theta)
            sin_theta = sin(theta)
            i_stop, j_stop, hit_dimension, delta_x_world_units_to_exit_start_tile, delta_y_world_units_to_exit_start_tile = RC.cast_ray_continous_world(obstacle_tile_map, x_start, y_start, cos_theta, sin_theta)
            Test.@test i_stop == 5
            Test.@test j_stop == 2
            Test.@test hit_dimension == 1
            # Test.@test total_euclidean ≈ sqrt(1.5 ^ 2 + 1 ^ 2)
            Test.@test delta_x_world_units_to_exit_start_tile ≈ 0.5
            Test.@test delta_y_world_units_to_exit_start_tile ≈ 0.5
        end
    end

    Test.@testset "Discrete world" begin
        I = Int
        world_units_per_tile_unit = convert(I, 256)

        i_start_world_units = convert(I, height_obstacle_tile_map * world_units_per_tile_unit ÷ 2 + 1)
        j_start_world_units = convert(I, width_obstacle_tile_map * world_units_per_tile_unit ÷ 2 + 1)

        Test.@testset "delta_i = 1, delta_j = 0" begin
            delta_i_world_units = convert(I, 1)
            delta_j_world_units = convert(I, 0)
            i_stop_tile_units, j_stop_tile_units, hit_dimension, delta_i_world_units_to_exit_start_tile, delta_j_world_units_to_exit_start_tile = RC.cast_ray_discrete_world(obstacle_tile_map, i_start_world_units, j_start_world_units, delta_i_world_units, delta_j_world_units, world_units_per_tile_unit)
            Test.@test i_stop_tile_units == convert(I, 5)
            Test.@test j_stop_tile_units == convert(I, 3)
            Test.@test hit_dimension == 1
            Test.@test delta_i_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2
            Test.@test delta_j_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2
        end

        Test.@testset "delta_i = 2, delta_j = 1" begin
            delta_i_world_units = convert(I, 2)
            delta_j_world_units = convert(I, 1)
            i_stop_tile_units, j_stop_tile_units, hit_dimension, delta_i_world_units_to_exit_start_tile, delta_j_world_units_to_exit_start_tile = RC.cast_ray_discrete_world(obstacle_tile_map, i_start_world_units, j_start_world_units, delta_i_world_units, delta_j_world_units, world_units_per_tile_unit)
            Test.@test i_stop_tile_units == convert(I, 5)
            Test.@test j_stop_tile_units == convert(I, 4)
            Test.@test hit_dimension == 1
            Test.@test delta_i_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2
            Test.@test delta_j_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2
        end

        Test.@testset "delta_i = 1, delta_j = 1" begin
            delta_i_world_units = convert(I, 1)
            delta_j_world_units = convert(I, 1)
            i_stop_tile_units, j_stop_tile_units, hit_dimension, delta_i_world_units_to_exit_start_tile, delta_j_world_units_to_exit_start_tile = RC.cast_ray_discrete_world(obstacle_tile_map, i_start_world_units, j_start_world_units, delta_i_world_units, delta_j_world_units, world_units_per_tile_unit)
            Test.@test i_stop_tile_units == convert(I, 5)
            Test.@test j_stop_tile_units == convert(I, 4)
            Test.@test hit_dimension == 1
            Test.@test delta_i_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2
            Test.@test delta_j_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2
        end

        Test.@testset "delta_i = 1, delta_j = 2" begin
            delta_i_world_units = convert(I, 1)
            delta_j_world_units = convert(I, 2)
            i_stop_tile_units, j_stop_tile_units, hit_dimension, delta_i_world_units_to_exit_start_tile, delta_j_world_units_to_exit_start_tile = RC.cast_ray_discrete_world(obstacle_tile_map, i_start_world_units, j_start_world_units, delta_i_world_units, delta_j_world_units, world_units_per_tile_unit)
            Test.@test i_stop_tile_units == convert(I, 4)
            Test.@test j_stop_tile_units == convert(I, 5)
            Test.@test hit_dimension == 2
            Test.@test delta_i_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2
            Test.@test delta_j_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2
        end

        Test.@testset "delta_i = 0, delta_j = 1" begin
            delta_i_world_units = convert(I, 0)
            delta_j_world_units = convert(I, 1)
            i_stop_tile_units, j_stop_tile_units, hit_dimension, delta_i_world_units_to_exit_start_tile, delta_j_world_units_to_exit_start_tile = RC.cast_ray_discrete_world(obstacle_tile_map, i_start_world_units, j_start_world_units, delta_i_world_units, delta_j_world_units, world_units_per_tile_unit)
            Test.@test i_stop_tile_units == convert(I, 3)
            Test.@test j_stop_tile_units == convert(I, 5)
            Test.@test hit_dimension == 2
            Test.@test delta_i_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2
            Test.@test delta_j_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2
        end

        Test.@testset "delta_i = -1, delta_j = 2" begin
            delta_i_world_units = convert(I, -1)
            delta_j_world_units = convert(I, 2)
            i_stop_tile_units, j_stop_tile_units, hit_dimension, delta_i_world_units_to_exit_start_tile, delta_j_world_units_to_exit_start_tile = RC.cast_ray_discrete_world(obstacle_tile_map, i_start_world_units, j_start_world_units, delta_i_world_units, delta_j_world_units, world_units_per_tile_unit)
            Test.@test i_stop_tile_units == convert(I, 2)
            Test.@test j_stop_tile_units == convert(I, 5)
            Test.@test hit_dimension == 2
            Test.@test delta_i_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2 + 1
            Test.@test delta_j_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2
        end

        Test.@testset "delta_i = -1, delta_j = 1" begin
            delta_i_world_units = convert(I, -1)
            delta_j_world_units = convert(I, 1)
            i_stop_tile_units, j_stop_tile_units, hit_dimension, delta_i_world_units_to_exit_start_tile, delta_j_world_units_to_exit_start_tile = RC.cast_ray_discrete_world(obstacle_tile_map, i_start_world_units, j_start_world_units, delta_i_world_units, delta_j_world_units, world_units_per_tile_unit)
            Test.@test i_stop_tile_units == convert(I, 2)
            Test.@test j_stop_tile_units == convert(I, 5)
            Test.@test hit_dimension == 2
            Test.@test delta_i_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2 + 1
            Test.@test delta_j_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2
        end

        Test.@testset "delta_i = -2, delta_j = 1" begin
            delta_i_world_units = convert(I, -2)
            delta_j_world_units = convert(I, 1)
            i_stop_tile_units, j_stop_tile_units, hit_dimension, delta_i_world_units_to_exit_start_tile, delta_j_world_units_to_exit_start_tile = RC.cast_ray_discrete_world(obstacle_tile_map, i_start_world_units, j_start_world_units, delta_i_world_units, delta_j_world_units, world_units_per_tile_unit)
            Test.@test i_stop_tile_units == convert(I, 1)
            Test.@test j_stop_tile_units == convert(I, 4)
            Test.@test hit_dimension == 1
            Test.@test delta_i_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2 + 1
            Test.@test delta_j_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2
        end

        Test.@testset "delta_i = -1, delta_j = 0" begin
            delta_i_world_units = convert(I, -1)
            delta_j_world_units = convert(I, 0)
            i_stop_tile_units, j_stop_tile_units, hit_dimension, delta_i_world_units_to_exit_start_tile, delta_j_world_units_to_exit_start_tile = RC.cast_ray_discrete_world(obstacle_tile_map, i_start_world_units, j_start_world_units, delta_i_world_units, delta_j_world_units, world_units_per_tile_unit)
            Test.@test i_stop_tile_units == convert(I, 1)
            Test.@test j_stop_tile_units == convert(I, 3)
            Test.@test hit_dimension == 1
            Test.@test delta_i_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2 + 1
            Test.@test delta_j_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2
        end

        Test.@testset "delta_i = -2, delta_j = -1" begin
            delta_i_world_units = convert(I, -2)
            delta_j_world_units = convert(I, -1)
            i_stop_tile_units, j_stop_tile_units, hit_dimension, delta_i_world_units_to_exit_start_tile, delta_j_world_units_to_exit_start_tile = RC.cast_ray_discrete_world(obstacle_tile_map, i_start_world_units, j_start_world_units, delta_i_world_units, delta_j_world_units, world_units_per_tile_unit)
            Test.@test i_stop_tile_units == convert(I, 1)
            Test.@test j_stop_tile_units == convert(I, 2)
            Test.@test hit_dimension == 1
            Test.@test delta_i_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2 + 1
            Test.@test delta_j_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2 + 1
        end

        Test.@testset "delta_i = -1, delta_j = -1" begin
            delta_i_world_units = convert(I, -1)
            delta_j_world_units = convert(I, -1)
            i_stop_tile_units, j_stop_tile_units, hit_dimension, delta_i_world_units_to_exit_start_tile, delta_j_world_units_to_exit_start_tile = RC.cast_ray_discrete_world(obstacle_tile_map, i_start_world_units, j_start_world_units, delta_i_world_units, delta_j_world_units, world_units_per_tile_unit)
            Test.@test i_stop_tile_units == convert(I, 1)
            Test.@test j_stop_tile_units == convert(I, 2)
            Test.@test hit_dimension == 1
            Test.@test delta_i_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2 + 1
            Test.@test delta_j_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2 + 1
        end

        Test.@testset "delta_i = -1, delta_j = -2" begin
            delta_i_world_units = convert(I, -1)
            delta_j_world_units = convert(I, -2)
            i_stop_tile_units, j_stop_tile_units, hit_dimension, delta_i_world_units_to_exit_start_tile, delta_j_world_units_to_exit_start_tile = RC.cast_ray_discrete_world(obstacle_tile_map, i_start_world_units, j_start_world_units, delta_i_world_units, delta_j_world_units, world_units_per_tile_unit)
            Test.@test i_stop_tile_units == convert(I, 2)
            Test.@test j_stop_tile_units == convert(I, 1)
            Test.@test hit_dimension == 2
            Test.@test delta_i_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2 + 1
            Test.@test delta_j_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2 + 1
        end

        Test.@testset "delta_i = 0, delta_j = -1" begin
            delta_i_world_units = convert(I, 0)
            delta_j_world_units = convert(I, -1)
            i_stop_tile_units, j_stop_tile_units, hit_dimension, delta_i_world_units_to_exit_start_tile, delta_j_world_units_to_exit_start_tile = RC.cast_ray_discrete_world(obstacle_tile_map, i_start_world_units, j_start_world_units, delta_i_world_units, delta_j_world_units, world_units_per_tile_unit)
            Test.@test i_stop_tile_units == convert(I, 3)
            Test.@test j_stop_tile_units == convert(I, 1)
            Test.@test hit_dimension == 2
            Test.@test delta_i_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2
            Test.@test delta_j_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2 + 1
        end

        Test.@testset "delta_i = 1, delta_j = -2" begin
            delta_i_world_units = convert(I, 1)
            delta_j_world_units = convert(I, -2)
            i_stop_tile_units, j_stop_tile_units, hit_dimension, delta_i_world_units_to_exit_start_tile, delta_j_world_units_to_exit_start_tile = RC.cast_ray_discrete_world(obstacle_tile_map, i_start_world_units, j_start_world_units, delta_i_world_units, delta_j_world_units, world_units_per_tile_unit)
            Test.@test i_stop_tile_units == convert(I, 4)
            Test.@test j_stop_tile_units == convert(I, 1)
            Test.@test hit_dimension == 2
            Test.@test delta_i_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2
            Test.@test delta_j_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2 + 1
        end

        Test.@testset "delta_i = 1, delta_j = -1" begin
            delta_i_world_units = convert(I, 1)
            delta_j_world_units = convert(I, -1)
            i_stop_tile_units, j_stop_tile_units, hit_dimension, delta_i_world_units_to_exit_start_tile, delta_j_world_units_to_exit_start_tile = RC.cast_ray_discrete_world(obstacle_tile_map, i_start_world_units, j_start_world_units, delta_i_world_units, delta_j_world_units, world_units_per_tile_unit)
            Test.@test i_stop_tile_units == convert(I, 5)
            Test.@test j_stop_tile_units == convert(I, 2)
            Test.@test hit_dimension == 1
            Test.@test delta_i_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2
            Test.@test delta_j_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2 + 1
        end

        Test.@testset "delta_i = 2, delta_j = -1" begin
            delta_i_world_units = convert(I, 2)
            delta_j_world_units = convert(I, -1)
            i_stop_tile_units, j_stop_tile_units, hit_dimension, delta_i_world_units_to_exit_start_tile, delta_j_world_units_to_exit_start_tile = RC.cast_ray_discrete_world(obstacle_tile_map, i_start_world_units, j_start_world_units, delta_i_world_units, delta_j_world_units, world_units_per_tile_unit)
            Test.@test i_stop_tile_units == convert(I, 5)
            Test.@test j_stop_tile_units == convert(I, 2)
            Test.@test hit_dimension == 1
            Test.@test delta_i_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2
            Test.@test delta_j_world_units_to_exit_start_tile == world_units_per_tile_unit ÷ 2 + 1
        end
    end
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.3'
#       jupytext_version: 0.8.6
#   kernelspec:
#     display_name: Julia 1.0.3
#     language: julia
#     name: julia-1.0
# ---

module ScalingModel
using DifferentialEquations

function micro_1(du, u, parms, time)
    # PARAMETER DEFS
    # β transmition rate
    # r net population growth rate
    # μ hosts' natural mortality rate
    # Κ population size
    # α disease induced mortality rate

    β, r, μ, K, α = parms
    S, I = u
    dS = r*(1-S/K)*S - β*S*I
    dI = β*S*I-(μ+α)*I
    du = [dS,dI]
end

# +
# PARAMETER DEFS
# w and m are used to define the other parameters allometrically

function main()
    w = 1;
    m = 10;
    β = 0.0247*m*w^0.44;
    r = 0.6*w^-0.27;
    μ = 0.4*w^-0.26;
    K = 16.2*w^-0.7;
    α = (m-1)*μ;
    # -

    parms = [β,r,μ,K,α];
    init = [K,1.];
    tspan = (0.0,10.0);

    sir_prob = ODEProblem(micro_1,init,tspan,parms)

    sir_sol = solve(sir_prob);
    return sir_sol
end

# using Plots

# plot(sir_sol,xlabel="Time",ylabel="Number")

# m = [5,10,20,40]
# ws = 10 .^collect(range(-3,length = 601,3))
# βs = zeros(601,4)
# for i = 1:4
#     βs[:,i] = 0.0247*m[i]*ws.^0.44
# end
# plot(ws,βs,xlabel="Weight",ylabel="\\beta_min", xscale=:log10,yscale=:log10, label=["m = 5" "m = 10" "m = 20" "m = 40"],lw=3)

end


"""
    natural_from_canonical(y::AbstractVector{<:Real}, σ²::AbstractVector{<:Real})

Compute natural parameters `η1` and `η2` from canonical paramters `y` and `σ²`.
Inverse of [`canonical_from_natural`](@ref).
"""
function natural_from_canonical(y::AbstractVector{<:Real}, σ²::AbstractVector{<:Real})
    η1, η2 = natural_from_canonical(y, Diagonal(σ²))
    return η1, η2.diag
end

function natural_from_canonical(y::AbstractVector{<:Real}, Σ::Diagonal{<:Real})
    return Σ \ y, -inv(Σ) / 2
end

"""
    canonical_from_natural(η1::AbstractVector{<:Real}, η2::AbstractVector{<:Real})

Compute canonical parameters `y` and `σ²` from natural parameters `η1` and `η2`.
Inverse of [`natural_from_canonical`](@ref).
"""
function canonical_from_natural(η1::AbstractVector{<:Real}, η2::AbstractVector{<:Real})
    y, Σ = canonical_from_natural(η1, Diagonal(η2))
    return y, diag(Σ)
end

function canonical_from_natural(η1::AbstractVector{<:Real}, η2::Diagonal{<:Real})
    Σ = -inv(η2) / 2
    y = Σ * η1
    return y, Σ
end

"""
    expectation_from_canonical(m::AbstractVector{<:Real}, σ²::AbstractVector{<:Real})

Compute expectation parameters `m1` and `m2` from canonical parameters `m` and `σ²`.
Inverse of [`expectation_from_canonical`](@ref).
"""
function expectation_from_canonical(m::AbstractVector{<:Real}, σ²::AbstractVector{<:Real})
    m, S = expectation_from_canonical(m, Diagonal(σ²))
    return m, diag(S)
end

function expectation_from_canonical(m::AbstractVector{<:Real}, Σ::Diagonal{<:Real})
    return m, Diagonal(m.^2 + diag(Σ))
end

"""
    canonical_from_expectation(m1::AbstractVector{<:Real}, m2::AbstractVector{<:Real})

Compute canonical parameters `m` and `σ²` from expectation parameters `m1` and `m2`.
Inverse of [`expectation_from_canonical`](@ref).
"""
function canonical_from_expectation(m1::AbstractVector{<:Real}, m2::AbstractVector{<:Real})
    m1, M2 = canonical_from_expectation(m1, Diagonal(m2))
    return m1, diag(M2)
end

function canonical_from_expectation(m1::AbstractVector{<:Real}, M2::Diagonal{<:Real})
    return m1, Diagonal(diag(M2) .- m1.^2)
end


function displacement(alpha, x)
    u1 = alpha * x[2] * sin(pi * x[1])
    u2 = alpha * (x[1]^3 + cos(pi * x[2]))
    return [u1, u2]
end

function stress_field(lambda, mu, alpha, x)
    s11 =
        (lambda + 2mu) * alpha * pi * x[2] * cos(pi * x[1]) -
        lambda * alpha * pi * sin(pi * x[2])
    s22 =
        -(lambda + 2mu) * alpha * pi * sin(pi * x[2]) +
        lambda * alpha * pi * x[2] * cos(pi * x[1])
    s12 = alpha * mu * (3x[1]^2 + sin(pi * x[1]))
    return [s11, s22, s12]
end

function body_force(lambda, mu, alpha, x)
    b1 = alpha * (lambda + 2mu) * pi^2 * x[2] * sin(pi * x[1])
    b2 =
        -alpha * (6mu * x[1] + (lambda + mu) * pi * cos(pi * x[1])) +
        alpha * (lambda + 2mu) * pi^2 * cos(pi * x[2])
    return [b1, b2]
end

function onleftboundary(x, L, W)
    return x[1] ≈ 0.0
end

function onbottomboundary(x, L, W)
    return x[2] ≈ 0.0
end

function onrightboundary(x, L, W)
    return x[1] ≈ L
end

function ontopboundary(x, L, W)
    return x[2] ≈ W
end

function onboundary(x, L, W)
    return onbottomboundary(x, L, W) ||
           onrightboundary(x, L, W) ||
           ontopboundary(x, L, W) ||
           onleftboundary(x, L, W)
end

function construct_mesh_and_quadratures(
    meshwidth,
    nelmts,
    elasticitybasis,
    distancefunction,
    numqp,
)
    levelsetbasis = HermiteTensorProductBasis(2)

    dim, nf = size(interpolation_points(levelsetbasis))
    refpoints = interpolation_points(elasticitybasis)

    mesh = CutCellDG.DGMesh([0.0, 0.0], meshwidth, [nelmts, nelmts], refpoints)
    cgmesh = CutCellDG.CGMesh([0.0, 0.0], meshwidth, [nelmts, nelmts], nf)

    levelset = CutCellDG.LevelSet(distancefunction, cgmesh, levelsetbasis)

    cutmesh = CutCellDG.CutMesh(mesh, levelset)
    cellquads = CutCellDG.CellQuadratures(cutmesh, levelset, numqp)
    interfacequads = CutCellDG.InterfaceQuadratures(cutmesh, levelset, numqp)
    facequads = CutCellDG.FaceQuadratures(cutmesh, levelset, numqp)

    mergedmesh =
        CutCellDG.MergedMesh!(cutmesh, cellquads, facequads, interfacequads)

    return mergedmesh, cellquads, facequads, interfacequads, levelset
end


function finite_diff(f::Function, args::Tuple, i::Int, dx::Float64)
    pos_args = Any[args...]
    pos_args[i] += dx
    neg_args = Any[args...]
    neg_args[i] -= dx
    return (f(pos_args...) - f(neg_args...)) / (2. * dx)
end

const dx = 1e-6

@testset "bernoulli" begin
    f = (x::Bool, prob::Float64) -> logpdf(Gen.Bernoulli(), x, prob)
    args = (false, 0.3,)
    actual = logpdf_grad(Gen.Bernoulli(), args...)
    @test isapprox(actual[2], finite_diff(f, args, 2, dx))
    args = (true, 0.3,)
    actual = logpdf_grad(Gen.Bernoulli(), args...)
    @test isapprox(actual[2], finite_diff(f, args, 2, dx))
end

@testset "beta" begin
    f = (x, alpha, beta) -> logpdf(Gen.Beta(), x, alpha, beta)
    args = (0.4, 0.2, 0.3)
    actual = logpdf_grad(Gen.Beta(), args...)
    @test isapprox(actual[1], finite_diff(f, args, 1, dx))
    @test isapprox(actual[2], finite_diff(f, args, 2, dx))
    @test isapprox(actual[3], finite_diff(f, args, 3, dx))
end


using FastGaussQuadrature

@testset "Gauss–Hermite" begin
    tol = 1e-14

    @testset "Golub-Welsch" begin
        n = 18;
        x,w = gausshermite( n )
        @test isa(x,Vector{Float64})
        @test isa(w,Vector{Float64})
        @test (length(x) == n && length(w) == n)
        @test (dot(w,x) < tol && abs(dot(w,x.^2) - sqrt(pi)/2) < tol)
    end

    @testset "Recurrence" begin
        n = 42
        x,w = gausshermite( n )
        @test isa(x,Vector{Float64})
        @test isa(w,Vector{Float64})
        @test (length(x) == n && length(w) == n)
        @test (dot(w,x) < tol && abs(dot(w,x.^2) - sqrt(pi)/2) < tol)
        @test abs(x[37] - 5.660357581283058) < tol
        @test abs(w[17] - 0.032202101288908) < tol

        f = x -> 1 + x + x^2 + x^3
        @test w'f.(x) ≈ 2.6586807763789717
    end

    @testset "Asymptotics"  begin
        n = 251
        x,w = gausshermite( n )
        @test isa(x,Vector{Float64})
        @test isa(w,Vector{Float64})
        @test (length(x) == n && length(w) == n)
        @test (dot(w,x) < tol && abs(dot(w,x.^2) - sqrt(pi)/2) < 300*tol)
        @test abs(x[37] - -13.292221459334638) < tol
        @test abs(w[123] - 0.117419270715955) < 2*tol

        n = 500
        x,w = gausshermite( n )
        @test isa(x,Vector{Float64})
        @test isa(w,Vector{Float64})
        @test (length(x) == n && length(w) == n)
        @test (dot(w,x) < tol && abs(dot(w,x.^2) - sqrt(pi)/2) < 300*tol)
    end

    @testset "Recurrence" begin
        x = 0.1
        He0 = 1 # Probabilists He_n(x) = H_n(x/sqrt(2))/2^(n/2)
        He1 = x; He2 = x^2-1; He3 = x*(x^2-3); He4 = 3 - 6x^2 + x^4
        @test FastGaussQuadrature.hermpoly_rec(1,x)[1] ≈ He1*exp(-x^2/4)
        @test FastGaussQuadrature.hermpoly_rec(2,x)[1] ≈ He2*exp(-x^2/4)/sqrt(2)
        @test FastGaussQuadrature.hermpoly_rec(3,x)[1] ≈ He3*exp(-x^2/4)/sqrt(2*3)
        @test FastGaussQuadrature.hermpoly_rec(4,x)[1] ≈ He4*exp(-x^2/4)/sqrt(2*3*4)

        @test FastGaussQuadrature.hermpoly_rec(1,60)[1] ≈ 0.0
        @test FastGaussQuadrature.hermpoly_rec(20^2,20)[1] ≈ 0.20019391063012504
        @test FastGaussQuadrature.hermpoly_rec(60^2,60)[1] ≈ 0.07918022667865038
        @test FastGaussQuadrature.hermpoly_rec(60^2,100)[1] ≈ -0.14191347555044895
        @test FastGaussQuadrature.hermpoly_rec(60^2,130)[1] ≈ 1.7334093012299562E-73
        @test FastGaussQuadrature.hermpoly_rec(60^2,200)[1] ≈ 0.0

        @test FastGaussQuadrature.hermpoly_rec(1:9,x) ≈ [FastGaussQuadrature.hermpoly_rec(n,x)[1] for n=1:9]
        @test FastGaussQuadrature.hermpoly_rec(0:20^2,20)[end] ≈ FastGaussQuadrature.hermpoly_rec(20^2,20)[1]
    end

    @testset "Transform" begin
        n = 500
        x,w = FastGaussQuadrature.unweightedgausshermite( n )
        @test w[1] ≠ 0

        V = Array{Float64}(undef,n,n)
        for k=1:n
            V[k,:] = FastGaussQuadrature.hermpoly_rec(0:n-1, sqrt(2)*x[k])
        end

        f = x -> first(FastGaussQuadrature.hermpoly_rec(1, sqrt(2)*x))
        @test V' * (w.* f.(x)) ≈ [0; sqrt(π); zeros(n-2)]
        f = x -> first(FastGaussQuadrature.hermpoly_rec(2, sqrt(2)*x))
        @test V' * (w.* f.(x)) ≈ [0; 0; sqrt(π); zeros(n-3)]
        f = x -> first(FastGaussQuadrature.hermpoly_rec(3, sqrt(2)*x))
        @test V' * (w.* f.(x)) ≈ [0; 0; 0; sqrt(π); zeros(n-4)]
    end
end


using CSV, Distributions, HypothesisTests

data = CSV.read("../data/machine1.csv", header=false)[:,1]
xBar, n = mean(data), length(data)
s = std(data)
alpha = 0.1
t = quantile(TDist(n-1),1-alpha/2)

println("Calculating formula: ", (xBar - t*s/sqrt(n), xBar + t*s/sqrt(n)))
println("Using confint() function: ", confint(OneSampleTTest(xBar,s,n),alpha))


import DualNumbers

export ∇n

function gradient(f::Function,Γ::Domain)
    fx = Fun(x->f(DualNumbers.Dual(real(x),1),imag(x)),Γ)
    fy = Fun(x->f(real(x),DualNumbers.Dual(imag(x),1)),Γ)
    Fun(Γ,DualNumbers.epsilon.(fx.coefficients)),Fun(Γ,DualNumbers.epsilon.(fy.coefficients))
end

∇n(f::Function,d::Domain) = normalderivative(f,d)

normalderivative(f::Function,Γ::UnionDomain) = Fun(map(d->normalderivative(f,d),Γ),PiecewiseSpace)

function normalderivative(f::Function,Γ::Domain)
    fx,fy = gradient(f,Γ)
    rp = Fun(x->fromcanonical(Γ,x),Γ)'
    n = -im*rp/abs(rp)
    fx*real(n)+fy*imag(n)
end


fanprojection(e::Ellipse, βᵢ, γᵢ, D) = projection(e, βᵢ + γᵢ, D * sin(γᵢ))

function fanprojection(
    ellipses::AbstractVector{<:Ellipse},
    β::AbstractVector,
    γ::AbstractVector,
    D::AbstractFloat,
)
    [sum(fanprojection(e, βᵢ, γᵢ, D) for e in ellipses) for γᵢ in γ, βᵢ in β]
end

function g(N::Integer, α::AbstractFloat)
    # p. 83 eq. 96
    g = zeros(N)
    N2 = div(N, 2)
    for (i, gi) in enumerate(g)
        n = i - N2 - 1
        if mod(n, 2) != 0
            g[i] = 0.5 * (1 / (π * sin(n * α)))^2
        elseif n == 0
            g[i] = 1 / (8 * α^2)
        end
    end
    g
end

function fanWeight!(R::AbstractMatrix, γ::AbstractVector, D::AbstractFloat)
    views = size(R, 2)
    for βᵢ = 1:views
        @inbounds @. R[:, βᵢ] *= D * cos(γ)
    end
end

function fanFilter!(Q::AbstractMatrix, α::AbstractFloat)
    samples, views = size(Q)
    N = nextpow(2, 2 * samples - 1)
    ZP(x) = vcat(x, zeros(N - samples))
    w = ZP(g(samples, α))
    # indices to slice out unpadded projection?
    i = samples ÷ 2 + 1
    j = i + samples - 1
    for βᵢ = 1:views
        # p. 83 eq. 94
        @inbounds Q[:, βᵢ] .= α * real.(ifft(fft(ZP(Q[:, βᵢ])) .* fftshift(fft(w))))[i:j]
    end
end

function fanReconstruction(
    R::AbstractMatrix,
    β::AbstractVector,
    γ::AbstractVector,
    D::AbstractFloat,
)
    Q = copy(R)
    α = step(γ)
    fanWeight!(Q, γ, D)
    fanFilter!(Q, α)

    Δβ = 2π / length(β)
    pixels = 128
    f = zeros(pixels, pixels)

    Q = LinearInterpolation((γ, β), Q, extrapolation_bc = Line())

    for (i, xᵢ) in enumerate(range(-1, 1, length = pixels)),
        (j, yⱼ) in enumerate(range(-1, 1, length = pixels))

        @inbounds f[j, i] = sum(β) do βᵢ
            r = √(xᵢ^2 + yⱼ^2)
            ϕ = atan(yⱼ, xᵢ)
            γ′ = atan((r * cos(βᵢ - ϕ)) / (D + r * sin(βᵢ - ϕ)))
            L² = (xᵢ - D * cos(βᵢ))^2 + (yⱼ - D * sin(βᵢ))^2
            Q(γ′, βᵢ) / L²
        end
    end
    Δβ .* f
end


using Church
using Distributions

mem(f::Function) = begin
    d = Dict()
    nf(args...) = begin
        if !haskey(d, args)
            d[args] = f(args...)
        end
        d[args]
    end
    nf
end

dp(concentration::Real, base_measure::Function) = begin
    sticks = mem(i::Int -> sample(Beta(1., concentration)))  #Dict{Int, Sample}
    atoms  = mem(i::Int -> base_measure())                   #Dict{Int, Sample}
    loop(i::Int) =
        @branch(sample(Bernoulli[sticks(i)]), atoms(i), loop(i+1))
    () -> loop(0)
end

dp_mixture(concentration::Real, base_measure::Function, parameter::Function) = begin
    dp_ = dp(concentration, base_measure)
    () -> parameter(dp_())
end

d = dp(1., () -> sample(Normal()))
ds = [d() for i = 1:20]

f = dp_mixture(1., () -> sample(Normal()), x -> sample(Normal[x, 0.1]))
fs = [f() for i = 1:20]


struct Rocket{T} <: Model{T}
      nq::Int
      nu::Int
      nw::Int
      nc::Int

      mass::T                               # mass
      inertia::Diagonal{T,Vector{T}}        # inertia matrix
      inertia_inv::Diagonal{T,Vector{T}}    # inertia matrix inverse
      gravity::Vector{T}                    # gravity
      length::T                             # length (com to thruster)
end

function f(model::Rocket, z, u, w)
      # states
      x = z[1:3]
      r = z[4:6]
      v = z[7:9]
      ω = z[10:12]

      # force in body frame
      F = u[1:3]

      # torque in body frame
      τ = [model.length * u[2],
           -model.length * u[1],
           0.0]

      [v;
       0.25 * ((1.0 - r' * r) * ω - 2.0 * cross(ω, r) + 2.0*(ω' * r) * r);
       model.gravity + (1.0 / model.mass) * MRP(r[1], r[2], r[3]) * F;
       model.inertia_inv * (τ - cross(ω, model.inertia * ω))]
end

nq, nu, nw = 12, 3, 0

mass = 1.0
len = 1.0

inertia = Diagonal([1.0 / 12.0 * mass * len^2.0, 1.0 / 12.0 * mass * len^2.0, 1.0e-5])
inertia_inv = Diagonal([1.0 / (1.0 / 12.0 * mass * len^2.0), 1.0 / (1.0 / 12.0 * mass * len^2.0), 1.0 / (1.0e-5)])

rocket = Rocket(nq, nu, nw, 0,
                  mass,
                  inertia,
                  inertia_inv,
                  [0.0, 0.0, -9.81],
                  len)


@testset "Vec" begin
    for T ∈ Floats
        a = Vec{3,T}(1,2,3)
        b = Vec{3,T}(2,4,6)
        @test a ⊙ b ≈ [2, 8, 18]
        @test b ⊘ a ≈ [2, 2, 2]
        @test inv(a)*a ≈ 1
        @test norm²(a) ≈ 1 + 4 + 9
    end
end


"""
Methods to find a low-rank subset
"""

"""
Return (error, indices) for the subset of size n, whose k-rank approximation has lowest error, as measured by the 2-norm.
Checks all subsets of size n. Has complexity Ω(|data_set|^n).
"""
function find_low_rank_subset_checkall(data_set,n,k)
    f(S) = let M = hcat([data_set[i] for i in S]...);
        truncated_svd(M, k)[1], S end
    minimum(f , powerset(1:length(data_set), n, n) )
end

find_low_rank_subset_checkall(A :: Array{T, 2}, n, k) where T <: Number = find_low_rank_subset_checkall(collect(eachcol(A)), n, k)

"""
Return (error, indices) for the subset of size n, whose k-rank approximation with vectors from data set has lowest error, as measured by the 2-norm.
Checks all subsets of size k. Has complexity Ω(|data_set|^k).
"""
function find_low_rank_subset_sample_rep(data_set,n,k)
    f(E) = let M = hcat([data_set[i] for i in E]...);
    M_times_pinv = M * pinv(M);
    distances_squared = [let d = v - M_times_pinv*v; dot(d,d) end for v in data_set]
    indices, distances = findmin(distances_squared, 1:n);
    sum(distances), indices
    end
    minimum(f, powerset(1:length(data_set), k, k))
end

find_low_rank_subset_sample_rep(A :: Array{T, 2}, n, k) where T <: Number = find_low_rank_subset_sample_rep(collect(eachcol(A)), n, k)

"""
Return an array of current_trajectories for each iteration step.
"""
function FastLowRank(data_set, n, k; initial_indices = [], step = L2_step, num_trajectories = k*trunc(Int(length(data_set)/n)), num_iterations = 50, verbose = false)
    current_trajectories = Vector{Trajectory}(undef, num_trajectories)
    # all_trajectories = Vector{Vector{Trajectory}}(undef, num_iterations)

    # initialize
    for t in 1:length(initial_indices)
        current_trajectories[t] = LowRank.Trajectory(data_set, n, k, initial_indices[t])
    end
    for t in length(initial_indices)+1:num_trajectories
        current_trajectories[t] = LowRank.Trajectory(data_set, n, k)
    end

    # do iterations
    is_converged = [false for i in 1:num_trajectories]
    best_converged_error = Inf
    best_error = Inf
    for i in 1:num_iterations
        for (ind, traj) in enumerate(current_trajectories)
            is_converged[ind] && continue
            new_traj = step(traj)
            best_error = min(best_error, new_traj.error)
            if Set(traj.indices) == Set(new_traj.indices)
                is_converged[ind] = true
                best_converged_error = min(best_converged_error, new_traj.error)
            end
            current_trajectories[ind] = new_traj
        end

        # @show best_error, best_converged_error
        if best_converged_error - best_error < eps(best_error)
            verbose && println("converged in $i steps")
            break
        elseif i == num_iterations
            verbose && println("did not converge in $i steps")
        end
    end
    best_error, best_trajectory_index = findmin((x -> x.error).(current_trajectories))
    best_trajectory = current_trajectories[best_trajectory_index]
    verbose && println("error = $best_error")
    return best_trajectory.error, best_trajectory.indices
end

FastLowRank(A :: Array{T, 2}, n, k; kwargs...) where T <: Number = FastLowRank(collect(eachcol(A)), n, k; kwargs...)


"""
Functions to deal with various types of norms and seminorms
"""

using FastRounding
using IntervalOptimisation
using TaylorSeries: Taylor1
using SparseArrays: getcolptr
using .DynamicDefinition

using .DynamicDefinition: derivative

"""
'Absolute value' definition that returns mag(I) for an interval and abs(x) for a real
"""
abs_or_mag(x::Number) = abs(x)
abs_or_mag(x::Interval) = mag(x)

"""
Computes a rigorous upper bound for z*z'
"""
z_times_conjz(z::Complex) = square_round(abs_or_mag(real(z)), RoundUp) ⊕₊ square_round(abs_or_mag(imag(z)), RoundUp)
abs_or_mag(z::Complex) = sqrt_round(z_times_conjz(z), RoundUp)

"""
Certified upper bound to ||A|| (of specified NormKind)
"""
function opnormbound(::Type{L1}, A::AbstractVecOrMat{T}) where {T}
    # partly taken from JuliaLang's LinearAlgebra/src/generic.jl
    Tnorm = typeof(abs_or_mag(float(real(zero(T)))))
    Tsum = promote_type(Float64, Tnorm)
    nrm::Tsum = 0
    @inbounds begin
        for j = 1:size(A, 2)
            nrmj::Tsum = 0
            for i = 1:size(A, 1)
                nrmj = nrmj ⊕₊ abs_or_mag(A[i,j])
            end
            nrm = max(nrm,nrmj)
        end
    end
    return convert(Tnorm, nrm)
end

function opnormbound(::Type{Linf}, A::AbstractVecOrMat{T}) where {T}
    # partly taken from JuliaLang's LinearAlgebra/src/generic.jl
    Tnorm = typeof(abs_or_mag(float(real(zero(T)))))
    Tsum = promote_type(Float64, Tnorm)
    nrm::Tsum = 0
    @inbounds begin
        for i = 1:size(A, 1)
            nrmi::Tsum = 0
            for j = 1:size(A, 2)
                nrmi = nrmi ⊕₊ abs_or_mag(A[i,j])
            end
            nrm = max(nrm,nrmi)
        end
    end
    return convert(Tnorm, nrm)
end

"""
These functions compute a rigorous upper bound for the 2-norm of a vector;
we have a specialized version for complex numbers to avoid taking
the sqrt root and squaring again
"""
function opnormbound(::Type{L2}, v::Vector{T}) where {T<:Real}
    # partly taken from JuliaLang's LinearAlgebra/src/generic.jl
    Tnorm = typeof(abs_or_mag(float(real(zero(T)))))
    Tsum = promote_type(Float64, Tnorm)
    nrm::Tsum = 0
    @inbounds begin
        for j = 1:length(v)
            nrm = nrm ⊕₊ square_round(abs_or_mag(v[j]), RoundUp)
        end
    end
    return convert(Tnorm, sqrt_round(nrm, RoundUp))
end

function opnormbound(::Type{L2}, v::Vector{T}) where {T<:Complex}
    # partly taken from JuliaLang's LinearAlgebra/src/generic.jl
    Tnorm = typeof(abs_or_mag(float(real(zero(T)))))
    Tsum = promote_type(Float64, Tnorm)
    nrm::Tsum = 0
    @inbounds begin
        for j = 1:length(v)
            nrm = nrm ⊕₊ z_times_conjz(v[j])
        end
    end
    return convert(Tnorm, sqrt_round(nrm, RoundUp))
end

function opnormbound(::Type{L1}, A::SparseMatrixCSC) where {T}
    # partly taken from JuliaLang's Sparsearray/src/linalg.jl
    m, n = size(A)
    Tnorm = typeof(abs_or_mag(float(real(zero(eltype(A))))))
    Tsum = promote_type(Float64,Tnorm)
    nA::Tsum = 0
    @inbounds begin
        for j=1:n
            colSum::Tsum = 0
            for i = getcolptr(A)[j]:getcolptr(A)[j+1]-1
                colSum = colSum ⊕₊ abs_or_mag(nonzeros(A)[i])
            end
            nA = max(nA, colSum)
        end
    end
    return convert(Tnorm, nA)
end

function opnormbound(::Type{Linf}, A::SparseMatrixCSC) where {T}
    # partly taken from JuliaLang's Sparsearray/src/linalg.jl
    m, n = size(A)
    Tnorm = typeof(abs_or_mag(float(real(zero(eltype(A))))))
    Tsum = promote_type(Float64,Tnorm)
    rowSum = zeros(Tsum,m)
    @inbounds begin
        for i=1:length(nonzeros(A))
            rowSum[rowvals(A)[i]] = rowSum[rowvals(A)[i]] ⊕₊ abs_or_mag(nonzeros(A)[i])
        end
    end
    return convert(Tnorm, maximum(rowSum))
end

"""
Rigorous upper bound on a vector norm. Note that Linf, L1 are the "analyst's" norms
"""
normbound(N::Type{L1}, v::AbstractVector) = opnormbound(L1, v) ⊘₊ Float64(length(v), RoundDown)
normbound(N::Type{Linf}, v::AbstractVector) = opnormbound(Linf, v)

"""
Types to compute norms iteratively by "adding a column at a time".
"""
abstract type NormCacher{T} end

mutable struct NormCacherL1 <: NormCacher{L1}
    C::Float64
    function NormCacherL1(n)
        new(0)
    end
end
"""
Create a new NormCacher to compute the normbound of the empty matrix with n rows
"""
NormCacher{L1}(n) = NormCacherL1(n)

mutable struct NormCacherLinf <: NormCacher{Linf}
    C::Vector{Float64}
    function NormCacherLinf(n)
        new(zeros(n))
    end
end
NormCacher{Linf}(n) = NormCacherLinf(n)

"""
Update a NormCacher to add one column to the matrix it is computing a norm of.
This column may be affected by an error ε (in the same norm).
"""
function add_column!(Cacher::NormCacherL1, v::AbstractVector, ε::Float64)
    Cacher.C = max(Cacher.C, opnormbound(L1, v) ⊕₊ ε)
end

function add_column!(Cacher::NormCacherLinf, v::AbstractVector, ε::Float64)
    Cacher.C = Cacher.C .⊕₊ abs.(v) .⊕₊ ε
end

"""
Return the norm of the matrix the NormCacher is working on.
"""
function get_norm(Cacher::NormCacherL1)
    return Cacher.C
end

function get_norm(Cacher::NormCacherLinf)
    return maximum(Cacher.C)
end

"""
(A, B) = dfly(strongnorm, auxnorm, dynamic)

Constants (A, B) such that ||Lf||_s ≦ A||f||_s + B||f||_aux
"""
dfly(::Type{<:NormKind}, ::Type{<:NormKind}, ::Dynamic) = @error "Not implemented"

# I don't think this is used in production anymore
function dfly(::Type{TotalVariation}, ::Type{L1}, D::Dynamic)
    dist = max_distorsion(D)
    lam = expansivity(D)

    if !(abs(lam) < 1) # these are intervals, so this is *not* equal to abs(lam) >= 1.
        @error "The function is not expanding"
    end

    if is_full_branch(D)
        return lam.hi, dist.hi
    else
        if !(abs(lam) < 0.5)
            @error "Expansivity is insufficient to prove a DFLY. Try with an iterate."
        end
        # We need a way to estimate the branch widths
        @error "Not implemented"
    end
end

function dfly(::Type{TotalVariation}, ::Type{L1}, D::PwMap)
    dist = max_distorsion(D)
    lam = expansivity(D)
    vec = endpoints(D)
    disc = maximum(2/abs(vec[i]-vec[i+1]) for i in nbranches(D))

    if is_full_branch(D)
        if !(abs(lam) < 1) # these are intervals, so this is *not* equal to abs(lam) >= 1.
            @error "The function is not expanding"
        end
        return lam.hi, dist.hi
    else
        if !(abs(2*lam) < 1)
            @error "Expansivity is insufficient to prove a DFLY. Try with an iterate."
        end
        return (2*lam).hi, (dist + disc).hi
    end
end

function dfly(::Type{Lipschitz}, ::Type{L1}, D::Dynamic)
    # TODO: should assert that D is globally C2 instead, but we don't have that kind of infrastructure yet.
    @assert is_full_branch(D)

    dist = max_distorsion(D)
    lam = expansivity(D)

    return ((lam*(2*dist+1)).hi, (dist*(dist+1)).hi)
end


using ReachabilityAnalysis, Symbolics, Plots

LazySets.set_ztol(Float64, 1e-15)

function multistable_oscillator(; X0 = Interval(0.0, 0.05),
                                  V₊ = +13.5, V₋ = -13.5,
                                  R = 20.E3, C = 5.5556E-8,
                                  R1 = 20.E3, R2 = 20.E3)

    @variables x
    τ = 1/(R*C)
    α = R2/(R1+R2)

    A = -τ
    b = (τ/α) * V₊
    I₊ = HalfSpace(x <= α*V₊)
    m1 = @system(x' = Ax + b, x ∈ I₊)

    b = (τ/α) * V₋
    I₋ = HalfSpace(x >= α*V₋)
    m2 = @system(x' = Ax + b, x ∈ I₋)

    automaton = LightAutomaton(2)
    add_transition!(automaton, 1, 2, 1)
    g1 = Hyperplane(x == α*V₊)
    r1 = ConstrainedIdentityMap(1, g1)

    add_transition!(automaton, 2, 1, 2)
    g2 = Hyperplane(x == α*V₋)
    r2 = ConstrainedIdentityMap(1, g2)

    modes = [m1, m2]
    resetmaps = [r1, r2]
    H = HybridSystem(automaton, modes, resetmaps)

    # initial condition in mode 1
    X0e = [(1, X0)]
    return IVP(H, X0e)
end

prob = multistable_oscillator()

sol = solve(prob, T=100e-4, alg=INT(δ=1.E-6), fixpoint_check=false);

plot(sol, vars=(0, 1), xlab="t", ylab="v-")

tspan.(sol)

location.(sol)

plot(sol[1][end-10:end], vars=(0, 1), xlab="t", ylab="v-")
plot!(x -> 6.75, xlims=(3.1e-4, 3.3e-4), lab="Guard", lw=2.0, color=:red)

Xc = cluster(sol[1], [318, 319, 320], BoxClustering(1))

plot(sol[1][end-10:end], vars=(0, 1))
plot!(sol[2][1:10], vars=(0, 1))
plot!(x -> 6.75, xlims=(3.1e-4, 3.3e-4), lab="Guard", lw=2.0, color=:red)
plot!(Xc[1], vars=(0, 1), c=:grey)

sol = solve(prob, T=100e-4, alg=INT(δ=1.E-6), fixpoint_check=true);

plot(sol, vars=(0, 1), xlab="t", ylab="v-")

tspan(sol)


"""
```
  genlasso(
      X    :: AbstractMatrix{T},
      y    :: AbstractVector{T};
      path :: Bool = true,
      ρ    :: Union{AbstractVector, Number} = zero(T),
      D    :: AbstractMatrix{T} = eye(size(X, 2)),
      solver = ECOSSolver(maxit=10e8, verbose=0)
    )
```
Solve generalized lasso problem by reformulating it as constrained lasso problem.
Note genralized lasso minimizes
    `0.5sumabs2(√obswt .* (y - X * β)) + ρ * sumabs(D * β)`.

### Arguments
- `X`       : predictor matrix.
- `y`       : response vector.

### Optional arguments
- `path`    : set `path=false` if user wishes to supply parameter value(s).
              Default is true.
- `ρ`       : tuning parameter value(s). Default is 0.
- `D`       : penalty matrix. Default is identity matrix.
- `solver`  : a solver Convex.jl supports. Default is ECOS.
              Note that Mosek and Gurobi are more robust than ECOS. Unlike ECOS or
              SCS, both Mosek and Gurobi require a license (free for academic
              use). <http://convexjl.readthedocs.io/en/latest/solvers.html>

### Returns
- `β`       : estimated coefficents.
- `objval`  : optimal objective value.
- `problem` : Convex.jl problem.
"""

function genlasso(
    X::AbstractMatrix{T},
    y::AbstractArray{T};
    path::Bool = true,
    ρ::Union{AbstractVector, Number} = zero(T),
    D::AbstractMatrix{T} = eye(size(X, 2)),
    solver = ECOSSolver(maxit=10e8, verbose=0)
    ) where T

  # singular value decomposition on D
  m = size(D, 1)
  F = svdfact!(D, thin = false)
  # extract singular values and submatrices
  singvals = F[:S]
  rankD = countnz(F[:S] .> abs(F[:S][1]) * eps(F[:S][1]) * maximum(size(D)))
  V1 = F[:V][:, 1:rankD]
  V2 = F[:V][:, rankD+1:end]
  U1 = F[:U][:, 1:rankD]
  U2 = F[:U][:, rankD+1:end]
  # calculate the MP-inverse of D
  Dplus = V1 * broadcast(*, U1', 1./F[:S])
  # transform the design matrix
  XDplus = X * Dplus
  XV2 = X * V2
  # projection matrix onto C(XV2)
  Pxv2 = (1 / dot(XV2, XV2)) * A_mul_Bt(XV2, XV2)
  # orthogonal projection matrix
  Mxv2 = eye(size(XV2, 1)) - Pxv2
  # obtain the new design matrix and response vector
  ỹ = vec(Mxv2 * y)
  X̃ = Mxv2 * XDplus
  # solve the constrained lasso problem
  if path
    α̂path, ρpath, = lsq_classopath(X̃, ỹ; solver=solver, Aeq = U2',
          beq = zeros(size(U2, 2)))
  else
    α̂path, = lsq_constrsparsereg(X̃, ỹ, ρ; solver=solver, Aeq = U2',
          beq = zeros(size(U, 2)))
    ρpath = ρ
  end
  # transform back to beta
  β̂path = Base.LinAlg.BLAS.ger!(1.0, vec(V2 * ((1 / dot(XV2, XV2)) *
    		At_mul_B(XV2, y))), ones(size(ρpath)), (eye(size(V2, 1)) -
    		V2 * ((1 / dot(XV2, XV2)) * At_mul_B(XV2, X))) * Dplus * α̂path)

  return β̂path, ρpath

end # end of function


import QuantumOpticsBase:
    tracenorm, tracenorm_h, tracenorm_nh, tracedistance, tracedistance_h, tracedistance_nh
import LinearAlgebra: ishermitian

"""
    tracenorm(rho)
Trace norm of `rho`.
It is defined as
```math
T(ρ) = \\operatorname{tr}\\{\\sqrt{ρ^† ρ}\\}.
```
Depending if `rho` is hermitian either [`tracenorm_h`](@ref) or
[`tracenorm_nh`](@ref) is called.
"""
function tracenorm(rho::DenseSuperOpType)
    ishermitian(rho) ? tracenorm_h(rho) : tracenorm_nh(rho)
end

#=
function tracenorm!(rho::DenseSuperOpType)
    if ishermitian(rho)
        rho.data ./= tracenorm_h(rho)
    else
        rho.data ./= tracenorm_nh(rho)
    end
end
=#

"""
    tracenorm_h(rho)
Trace norm of `rho`.
It uses the identity
```math
T(ρ) = \\operatorname{tr}\\{\\sqrt{ρ^† ρ}\\} = \\sum_i |λ_i|
```
where ``λ_i`` are the eigenvalues of `rho`.
"""
function tracenorm_h(rho::DenseSuperOpType{B,B,T}) where {B<:Tuple{Basis,Basis},T}
    s = eigvals(Hermitian(rho.data))
    sum(abs.(s))
end

#=
function tracenorm_h(rho::T) where {T<:AbstractOperator}
    throw(ArgumentError("tracenorm_h not implemented for $(typeof(rho)). Use dense operators instead."))
end
=#

"""
    tracenorm_nh(rho)
Trace norm of `rho`.
Note that in this case `rho` doesn't have to be represented by a square
matrix (i.e. it can have different left-hand and right-hand bases).
It uses the identity
```math
    T(ρ) = \\operatorname{tr}\\{\\sqrt{ρ^† ρ}\\} = \\sum_i σ_i
```
where ``σ_i`` are the singular values of `rho`.
"""
tracenorm_nh(rho::SuperOperator{B,B,T}) where {B<:Tuple{Basis,Basis},T} =
    sum(svdvals(rho.data))

"""
    tracedistance(rho, sigma)
Trace distance between `rho` and `sigma`.
It is defined as
```math
T(ρ,σ) = \\frac{1}{2} \\operatorname{tr}\\{\\sqrt{(ρ - σ)^† (ρ - σ)}\\}.
```
It calls [`tracenorm`](@ref) which in turn either uses [`tracenorm_h`](@ref)
or [`tracenorm_nh`](@ref) depending if ``ρ-σ`` is hermitian or not.
"""
tracedistance(
    rho::SuperOperator{B,B,T},
    sigma::SuperOperator{B,B,T},
) where {B<:Tuple{Basis,Basis},T} = 0.5 * tracenorm(rho - sigma)

"""
    tracedistance_h(rho, sigma)
Trace distance between `rho` and `sigma`.
It uses the identity
```math
T(ρ,σ) = \\frac{1}{2} \\operatorname{tr}\\{\\sqrt{(ρ - σ)^† (ρ - σ)}\\} = \\frac{1}{2} \\sum_i |λ_i|
```
where ``λ_i`` are the eigenvalues of `rho` - `sigma`.
"""
tracedistance_h(
    rho::SuperOperator{B,B,T},
    sigma::SuperOperator{B,B,T},
) where {B<:Tuple{Basis,Basis},T} = 0.5 * tracenorm_h(rho - sigma)

"""
    tracedistance_nh(rho, sigma)
Trace distance between `rho` and `sigma`.
Note that in this case `rho` and `sigma` don't have to be represented by square
matrices (i.e. they can have different left-hand and right-hand bases).
It uses the identity
```math
    T(ρ,σ) = \\frac{1}{2} \\operatorname{tr}\\{\\sqrt{(ρ - σ)^† (ρ - σ)}\\}
         = \\frac{1}{2} \\sum_i σ_i
```
where ``σ_i`` are the singular values of `rho` - `sigma`.
"""
tracedistance_nh(
    rho::SuperOperator{B1,B2,T},
    sigma::SuperOperator{B1,B2,T},
) where {B1<:Tuple{Basis,Basis},B2<:Tuple{Basis,Basis},T} = 0.5 * tracenorm_nh(rho - sigma)


@doc raw"""
    InvariantMetric{G<:Metric,D<:ActionDirection} <: Metric

Extend a metric on the Lie algebra of an [`AbstractGroupManifold`](@ref) to the whole group
via translation in the specified direction.

Given a group $\mathcal{G}$ and a left- or right group translation map $τ$ on the group, a
metric $g$ is $τ$-invariant if it has the inner product

````math
g_p(X, Y) = g_{τ_q p}((\mathrm{d}τ_q)_p X, (\mathrm{d}τ_q)_p Y),
````

for all $p,q ∈ \mathcal{G}$ and $X,Y ∈ T_p \mathcal{G}$, where $(\mathrm{d}τ_q)_p$ is the
differential of translation by $q$ evaluated at $p$ (see [`translate_diff`](@ref)).

`InvariantMetric` constructs an (assumed) $τ$-invariant metric by extending the inner
product of a metric $h_e$ on the Lie algebra to the whole group:

````math
g_p(X, Y) = h_e((\mathrm{d}τ_p^{-1})_p X, (\mathrm{d}τ_p^{-1})_p Y).
````

!!! warning
    The invariance condition is not checked and must be verified for the entire group.
    To verify the condition for a set of points numerically, use
    [`has_approx_invariant_metric`](@ref).

The convenient aliases [`LeftInvariantMetric`](@ref) and [`RightInvariantMetric`](@ref) are
provided.

# Constructor

    InvariantMetric(metric::Metric, conv::ActionDirection = LeftAction())
"""
struct InvariantMetric{G<:Metric,D<:ActionDirection} <: Metric
    metric::G
end

function InvariantMetric(metric, conv = LeftAction())
    return InvariantMetric{typeof(metric),typeof(conv)}(metric)
end

const LeftInvariantMetric{G} = InvariantMetric{G,LeftAction} where {G<:Metric}

"""
    LeftInvariantMetric(metric::Metric)

Alias for a left-[`InvariantMetric`](@ref).
"""
LeftInvariantMetric(metric) = InvariantMetric{typeof(metric),LeftAction}(metric)

const RightInvariantMetric{G} = InvariantMetric{G,RightAction} where {G<:Metric}

"""
    RightInvariantMetric(metric::Metric)

Alias for a right-[`InvariantMetric`](@ref).
"""
RightInvariantMetric(metric) = InvariantMetric{typeof(metric),RightAction}(metric)

@doc raw"""
    has_approx_invariant_metric(
        G::AbstractGroupManifold,
        p,
        X,
        Y,
        qs::AbstractVector,
        conv::ActionDirection = LeftAction();
        kwargs...,
    ) -> Bool

Check whether the metric on the group $\mathcal{G}$ is (approximately) invariant using a set of predefined
points. Namely, for $p ∈ \mathcal{G}$, $X,Y ∈ T_p \mathcal{G}$, a metric $g$, and a
translation map $τ_q$ in the specified direction, check for each $q ∈ \mathcal{G}$ that the
following condition holds:

````math
g_p(X, Y) ≈ g_{τ_q p}((\mathrm{d}τ_q)_p X, (\mathrm{d}τ_q)_p Y).
````

This is necessary but not sufficient for invariance.

Optionally, `kwargs` passed to `isapprox` may be provided.
"""
has_approx_invariant_metric(
    ::AbstractGroupManifold,
    ::Any,
    ::Any,
    ::Any,
    ::Any,
    ::ActionDirection,
)
@decorator_transparent_function function has_approx_invariant_metric(
    M::AbstractGroupManifold,
    p,
    X,
    Y,
    qs,
    conv::ActionDirection = LeftAction();
    kwargs...,
)
    gpXY = inner(M, p, X, Y)
    for q in qs
        τq = translate(M, q, p, conv)
        dτqX = translate_diff(M, q, p, X, conv)
        dτqY = translate_diff(M, q, p, Y, conv)
        isapprox(gpXY, inner(M, τq, dτqX, dτqY); kwargs...) || return false
    end
    return true
end

direction(::InvariantMetric{G,D}) where {G,D} = D()

function exp!(
    M::MetricManifold{𝔽,<:AbstractGroupManifold,<:InvariantMetric},
    q,
    p,
    X,
) where {𝔽}
    if has_biinvariant_metric(M)
        conv = direction(metric(M))
        return retract!(base_group(M), q, p, X, GroupExponentialRetraction(conv))
    end
    return invoke(exp!, Tuple{MetricManifold,typeof(q),typeof(p),typeof(X)}, M, q, p, X)
end

"""
    biinvariant_metric_dispatch(G::AbstractGroupManifold) -> Val

Return `Val(true)` if the metric on the manifold is bi-invariant, that is, if the metric
is both left- and right-invariant (see [`invariant_metric_dispatch`](@ref)).
"""
function biinvariant_metric_dispatch(M::Manifold)
    return Val(
        invariant_metric_dispatch(M, LeftAction()) === Val(true) &&
        invariant_metric_dispatch(M, RightAction()) === Val(true),
    )
end

has_biinvariant_metric(M::Manifold) = _extract_val(biinvariant_metric_dispatch(M))

@doc raw"""
    invariant_metric_dispatch(G::AbstractGroupManifold, conv::ActionDirection) -> Val

Return `Val(true)` if the metric on the group $\mathcal{G}$ is invariant under translations
by the specified direction, that is, given a group $\mathcal{G}$, a left- or right group
translation map $τ$, and a metric $g_e$ on the Lie algebra, a $τ$-invariant metric at
any point $p ∈ \mathcal{G}$ is defined as a metric with the inner product

````math
g_p(X, Y) = g_{τ_q p}((\mathrm{d}τ_q)_p X, (\mathrm{d}τ_q)_p Y),
````

for $X,Y ∈ T_q \mathcal{G}$ and all $q ∈ \mathcal{G}$, where $(\mathrm{d}τ_q)_p$ is the
differential of translation by $q$ evaluated at $p$ (see [`translate_diff`](@ref)).
"""
invariant_metric_dispatch(::MetricManifold, ::ActionDirection)

@decorator_transparent_signature invariant_metric_dispatch(
    M::AbstractDecoratorManifold,
    conv::ActionDirection,
)
function invariant_metric_dispatch(M::MetricManifold, conv::ActionDirection)
    is_default_metric(M) && return invariant_metric_dispatch(M.manifold, conv)
    return Val(false)
end
function invariant_metric_dispatch(
    M::MetricManifold{𝔽,<:Manifold,<:InvariantMetric},
    conv::ActionDirection,
) where {𝔽}
    direction(metric(M)) === conv && return Val(true)
    return invoke(invariant_metric_dispatch, Tuple{MetricManifold,typeof(conv)}, M, conv)
end
invariant_metric_dispatch(M::Manifold, ::ActionDirection) = Val(false)

function has_invariant_metric(M::Manifold, conv::ActionDirection)
    return _extract_val(invariant_metric_dispatch(M, conv))
end

function inner(M::MetricManifold{𝔽,<:Manifold,<:InvariantMetric}, p, X, Y) where {𝔽}
    imetric = metric(M)
    conv = direction(imetric)
    N = MetricManifold(M.manifold, imetric.metric)
    Xₑ = inverse_translate_diff(M, p, p, X, conv)
    Yₑ = inverse_translate_diff(M, p, p, Y, conv)
    return inner(N, Identity(N, p), Xₑ, Yₑ)
end

function default_metric_dispatch(
    M::MetricManifold{𝔽,<:Manifold,<:InvariantMetric},
) where {𝔽}
    imetric = metric(M)
    N = MetricManifold(M.manifold, imetric.metric)
    default_metric_dispatch(N) === Val(true) || return Val(false)
    return invariant_metric_dispatch(N, direction(imetric))
end

function log!(
    M::MetricManifold{𝔽,<:AbstractGroupManifold,<:InvariantMetric},
    X,
    p,
    q,
) where {𝔽}
    if has_biinvariant_metric(M)
        imetric = metric(M)
        conv = direction(imetric)
        return inverse_retract!(
            base_group(M),
            X,
            p,
            q,
            GroupLogarithmicInverseRetraction(conv),
        )
    end
    return invoke(log!, Tuple{MetricManifold,typeof(X),typeof(p),typeof(q)}, M, X, p, q)
end

function LinearAlgebra.norm(
    M::MetricManifold{𝔽,<:Manifold,<:InvariantMetric},
    p,
    X,
) where {𝔽}
    imetric = metric(M)
    conv = direction(imetric)
    N = MetricManifold(M.manifold, imetric.metric)
    Xₑ = inverse_translate_diff(M, p, p, X, conv)
    return norm(N, Identity(N, p), Xₑ)
end

function Base.show(io::IO, metric::LeftInvariantMetric)
    return print(io, "LeftInvariantMetric($(metric.metric))")
end
function Base.show(io::IO, metric::RightInvariantMetric)
    return print(io, "RightInvariantMetric($(metric.metric))")
end


################################################################################
#                          TEST: DCM <=> Quaternions
################################################################################

for k = 1:samples
    # Conversion test
    # ===============

    # Create a random quaternion.
    q1 = Quaternion(randn(), randn(), randn(), rand())
    q1 = q1/norm(q1)

    # Make sure that the real part is positive.
    (q1.q0 < 0) && (q1 *= -1)

    # Create a random DCM using Euler Angles.
    dcm2 = angle2dcm(randn(), randn(), randn(), rot_seq_array[rand(1:9)])

    # Convert the quaternion in DCM and the DCM in quaternion.
    q2   = dcm2quat(dcm2)
    dcm1 = quat2dcm(q1)

    # Do the inverse conversion and compare.
    error_dcm = dcm2 - quat2dcm(q2)
    @test norm(error_dcm) < 1e-10

    error_q = q1 - dcm2quat(dcm1)
    @test norm(error_q) < 1e-10

    # Multiplication test
    # ===================

    # Sample a 3x1 vector.
    v = randn(3)

    # Compute the rotated vector using the quaternions.
    v_rot_q = vect(inv(q1*q2)*v*(q1*q2))

    # Compute the rotated vector using the DCMs.
    v_rot_dcm = dcm2*dcm1*v

    # Compare.
    @test norm(v_rot_q-v_rot_dcm) < 1e-10
end


using DifferentialEquations

"""
Calculate the Froude number.

"""
function froude(Q::Real, xs::CrossSection)
    g = 9.806
    Fr = Q / (depth(xs) * width(xs) * sqrt(g * depth(xs)))
end

"""
Ordinary differential equation describing the Gradually-Varied-Flow model.

"""
function dydx(y, p, x)
    i = trunc(Int, ceil(x / (p[3] / length(p[1]))))
    i = i > 0 ? i : 1
    # TODO: Change this dynamically depending on whether the flow is subcritical or supercritical
    pm = -1 # integrate upstream
    Q = p[1]
    xs = p[2][i]
    xs.Yₘ = y * (xs.r + 1) / xs.r
    S0 = xs.S0
    Sf = xs.n^2 * Q^2 / (width(xs)^2 * depth(xs)^(10/3))
    Fr = froude(Q, xs)
    pm * (S0 - Sf) / (1 - Fr^2)
end

"""
Calculate a water surface profile by solving the Gradually-Varied-Flow equation.

- `Q`: discharge
- `ybc`: downstream boundary condition for depth
- `z0`: bed elevation at downstream cross section
- `S0`: bed slope for reach
- `n`: Manning's roughness coefficient
- `x`: downstream distance for each cross section
- `b`: bottom width for each cross section
- `r`: channel geometry coefficient

"""
function gvf(Q, ybc, w, S0, n, x, wbf, ybf, r)
    c = [Dingman(wbf[i], ybf[i], ybc, r[i], S0[i], n) for i in 1:length(x)]
    prob = DiscreteProblem(dydx, ybc, (x[1], x[end]), (Q, c, x[end]))
    h = try
        sol = solve(prob, Tsit5(), abstol=1e-2, saveat=x)
        sol.u
    catch
        zeros(length(x)) .- 9999.
    end
    h
end


using LeastSquaresEstimation
using Plots

function main()
    # Define basis function vector
    h = [t->sin(t), t->cos(t), t->t, t->cos(t)*sin(t), t->t^2]

    # Instantiate Linear Least Squares Estimator
    llse = LinearLeastSquares(h)

    # Generate Measurements
    m    = 101
    y(t) = 0.3*sin(t) + 0.5*cos(t) + 0.1*t + sqrt(0.001)*randn()

    ts = zeros(m)
    ys = zeros(m)
    @inbounds for i in 1:m
        if i != 1
            ts[i] = ts[i-1] + 0.1
        end
        ys[i] = y(ts[i])
    end

    # Feed estimator batch of measurements
    FeedMeasurementBatch!(llse, ys, ts)

    # Compute estimate
    ComputeEstimate!(llse)

    # Compute estimated measurements and plot
    yhat = zeros(m)
    @inbounds for i in 1:m
        yhat[i] = GetMeasurementEstimate(llse, ts[i])
    end

    # Display
    println("Estimated parameters:")
    display(GetEstimate(llse))

    # Plot generated and estimated measurements
    plot(ts, ys; label = "Simulated measurements")
    plot!(ts, yhat; label = "Estimated measurements")
end

main()


using ATA
using JuMP
# add Cbc by running import Pkg; Pkg.add("Cbc")
using Cbc
# cd("where your input files are")
cd("examples")
# 1. Start ATA and add file with custom settings (Needed)
ata_model = start_ata(;
    settings_file = "settings_ata_maximin.jl",
    bank_file = "data/bank.csv",
    bank_delim = ";",
)

# 2. Add categorical constraints (Optional)
add_constraints!(ata_model; constraints_file = "constraints.csv", constraints_delim = ";")

# 3. Add overlap maxima (Optional)
add_overlap!(ata_model; overlap_file = "overlap_matrix.csv", overlap_delim = ";")

# 4. Add objective function (Optional)
add_obj_fun!(ata_model)
print_infos(ata_model)

# Assembly settings

# Set the solver, "siman" for simulated annealing, "jump" for MILP solver.
solver = "jump"

# MILP (Not suggested for large scale ATA)
# Select the solver, Cbc as open-source is a good option.
optimizer_constructor = "Cbc"
# #Optimizer attributes
optimizer_attributes = [("seconds", 100), ("logLevel", 1)]

# 5. assemble
assemble!(
    ata_model;
    solver = solver,
    optimizer_attributes = optimizer_attributes,
    optimizer_constructor = optimizer_constructor,
)

# All the settings and outputs from optimization are in ata_model object.
# See the struct in ATA.jl to understand how to retrieve all the information.
# A summary of the resulting tests is available in results_folder/results.txt
# If siman is chosen, the optimality and feasibility of the best neighbourhood
# is reported in "results/results_ata.jl"

print_results(ata_model; results_folder = "results")

#]add https://github.com/giadasp/ATAPlot.jl
using ATAPlot
plot_results(ata_model; plots_folder = "plots")


using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

compute(n::Int)::Int = Int(ceil((log(10) * (n - 1) + 0.5log(5)) / log(0.5(1 + √5))))

compute(3)

compute(1_000)

@benchmark compute(1_000)


## rand.jl.  Can we generate random GMMs and sample from them?
## (c) 2013--2014 David A. van Leeuwen

## This function initializes a random GMM, with random means and random covariances
## The variances are somewhat arbitrarily chosen.  This can certainly be improved.
function Base.rand(::Type{GMM}, ng::Int, d::Int; sep=2.0, kind=:full)
    μ = sep * randn(ng, d)
    if kind==:diag
        Σ = hcat([rand(Chisq(1.0), ng) for i=1:d]...)
    elseif kind == :full
        Σ = Array{eltype(FullCov{Float64})}(undef, ng)
        for i=1:ng
            T = randn(d,d)
            Σ[i] = cholinv(T' * T / d)
        end
    else
        error("Unknown kind")
    end
    w = ones(ng)/ng
    hist = History(@sprintf("Generated random %s covariance GMM with %d Gaussians of dimension %d",
                            kind, ng, d))
    GMM(w, μ, Σ, [hist], 0)
end

## local helper
function binsearch(x::T, a::Vector{T}) where {T}
    issorted(a) || error("Array needs to be sorted")
    mi = 1
    ma = length(a)
    if x < a[mi]
        return 0
    elseif x >= a[ma]
        return ma
    end
    while ma - mi > 1
        h = mi + div(ma-mi,2)
        if x > a[h]
            mi = h
        else
            ma = h
        end
    end
    return mi
end

forcesymmetric(c::Matrix) = (Symmetric(c))

## This function samples n data points from a GMM.  This is pretty slow, probably due to the array assignments.
function Base.rand(gmm::GMM, n::Int)
    x = Array{Float64}(undef, n, gmm.d)
    ## generate indices distriuted according to weights
    index::Vector{Int} = mapslices(findall, rand(Multinomial(1, gmm.w), n).!=0,
                                   dims=1)[:]
    gmmkind = kind(gmm)
    for i=1:gmm.n
        ind = findall(index.==i)
        nₓ = length(ind)
        if gmmkind == :diag
            x[ind,:] .= (vec(gmm.μ[i,:]) .+ vec(sqrt.(gmm.Σ[i,:])) .* randn(gmm.d, nₓ))' ## v0.5 arraymageddon
        elseif gmmkind == :full
            x[ind,:] .= rand(MvNormal(vec(gmm.μ[i,:]), forcesymmetric(covar(gmm.Σ[i]))), nₓ)'
        else
            error("Unknown kind")
        end
    end
    x
end


function sample(μs::Vector{<:Real}, σs::Vector{<:Real})::Vector{Float64}
    return max.(rand.(Normal.(μs, σs)), zero(Float64))
end

"""
    VariablePoissonInput{N, L}

Iterate over a Poisson input for `N` neurons with rate `N(λs[n], σs[n])`
for each neuron `n` `L` times, where the rate is resampled every
`resample` iterations.
"""
struct VariablePoissonInput{N, L} <: AbstractPoissonInput{N, L}
    μs::Vector{Float64}
    σs::Vector{Float64}
    λs::Vector{Float64}
    buffer::Vector{Int64}
    start::Int64
    stop::Int64
    resample::Int64

    function VariablePoissonInput{N, L}(μs, σs, λs, buffer, start, stop, resample) where {N, L}
        if !(typeof(N) <: Integer)
            error("Number of neurons N must be an integer")
        elseif !(typeof(L) <: Integer)
            error("Input length L must be an integer")
        elseif length(μs) != N
            error("μs must have length N.")
        elseif length(σs) != N
            error("σs must have length N.")
        elseif length(λs) != N
            error("λs must have length N.")
        elseif length(buffer) != N
            error("buffer must have length N.")
        end

        new(μs, σs, λs, buffer, start, stop, resample)
    end
end

VariablePoissonInput(μs::Vector{<:Real}, σs::Vector{<:Real}, L::Integer; resample=one(Int64), start=one(Int64), stop=L) = VariablePoissonInput{length(μs), L}(μs, σs, sample(μs, σs), zeros(Int64, length(μs)), start, stop, resample)
VariablePoissonInput(μ::Real, σ::Real, N::Integer, L::Integer; resample=one(Int64), start=one(Int64), stop=L) = VariablePoissonInput(fill(μ, N), fill(σ, N), N, L; resample=resample, start=start, stop=stop)

"""
    rates(iter::VariablePoissonInput{N, L}, i::Integer)::Vector{Float64} where {N, L}

Obtain the firing rate for each neuron in `iter` for iteration `i`.
"""
function rates(iter::VariablePoissonInput{N, L}, i::Integer)::Vector{Float64} where {N, L}
    if i % iter.resample == 0
        iter.λs .= sample(iter.μs, iter.σs)
    end
    return iter.λs
end


using PyPlot
using Triangulate

N = Float64[ 0  -2  2  2  1 -1 -2
             0  -2 -2  0  2  2  0 ]

T = [ 7  2  3  4  5  6
      2  3  4  5  6  7
      1  1  1  1  1  1 ]

E = [ 7  2  3  4  5  6
      2  3  4  5  6  7 ]

tri = Triangulation(N, T, E, 4)

figure(1)
draw_triangles(tri)
enumerate_triangles(tri)
enumerate_vertices(tri)
enumerate_bdry_edges(tri, 0.1)
#local_enumerate_vertices(tri, 0.3)
mark_first_local_vertex(tri, 0.3)
axis("off")
axis("equal")
savefig("good_triangulation.pdf")


for N in [Float64, Float32, Rational{Int}]
    # BallInf approximation of a 3D unit ball in the 1-norm centered at [1,2,0]
    b = Ball1(N[1, 2, 0], N(1))
    bi = ballinf_approximation(b)
    biexp = BallInf(N[1, 2, 0], N(1))
    @test bi.center ≈ biexp.center
    @test bi.radius ≈ biexp.radius

    # BallInf approximation of a 2D polygon whose vertices are
    # (0,1), (1,1), (2,-1), (-1,0)
    p = HPolygon{N}()
    addconstraint!(p, LinearConstraint(N[0, 1], N(1)))
    addconstraint!(p, LinearConstraint(N[-1, 1], N(1)))
    addconstraint!(p, LinearConstraint(N[-1, -3], N(1)))
    addconstraint!(p, LinearConstraint(N[2, 1], N(3)))
    bi = ballinf_approximation(p)
    biexp = BallInf(N[0.5, 0], N(1.5))
    @test bi.center ≈ biexp.center
    @test bi.radius ≈ biexp.radius

    # empty set
    E = EmptySet{N}(2)
    @test ballinf_approximation(E) == E

    if N == Float64
        # robustness (see issue #2532): the set has two contradicting constraints
        # => the set is empty, but requires to use high precision
        # x >= 1.0000000000000002
        # x <= 1
        s1 = HalfSpace([-1.0], -1.0000000000000002)
        s2 = HalfSpace([1.0], 1.0)
        P = s1 ∩ s2
        @test ballinf_approximation(P) == BallInf([1.0], 0.0)

        s1big = HalfSpace([-big(1.0)], -big(1.0000000000000002))
        s2big = HalfSpace([big(1.0)], big(1.0))
        Pbig = HPolytope([s1big, s2big])
        @test ballinf_approximation(Pbig) == EmptySet{BigFloat}(1)
    end
end


# Define and register smooth functions
_cos_wave(t, f, A, st, ϕ) = A*cos(2*π*f*(t - st) + ϕ)
_sin_wave(t, f, A, st, ϕ) = A*sin(2*π*f*(t - st) + ϕ)
_damped_sine_wave(t, f, A, st, ϕ, d) = exp((st-t)*d)*A*sin(2*π*f*(t-st) + ϕ)
_ramp(t, δ, st, et, h) = h/(et-st)*(_xH(t, δ, st) - _xH(t, δ, et))
_square_wave(t, δ, f, A, st) = A*2atan(sin(2π*(t-st)*f)/δ)/π
_step(t, δ, h, a) = h*(atan((t-a)/δ)/π + 0.5)
_triangular_wave(t, δ, f, A, st) = A*(1-2acos((1 - δ)sin(2π*(t-st)*f))/π)
_xH(t, δ, tₒ) = (t-tₒ)*(1+((t-tₒ)/sqrt((t-tₒ)^2+δ^2)))/2

@register_symbolic _cos_wave(t, f, A, st, ϕ)
@register_symbolic _sin_wave(t, f, A, st, ϕ)
@register_symbolic _damped_sine_wave(t, f, A, st, ϕ, damping)
@register_symbolic _ramp(t, δ, st, et, h)
@register_symbolic _square_wave(t, δ, f, A, st)
@register_symbolic _step(t, δ, h, a)
@register_symbolic _triangular_wave(t, δ, f, A, st)

"""
```julia
function ConstantVoltage(; name, V)
```

The source for an ideal constant voltage.

# States
- `v(t)`: [`V`]
  The voltage across this source, given by `p.v - n.v` and is always constant

# Connectors
- `p`
  Positive pin
- `n`
  Negative pin

# Parameters:
- `V`: [`V`]
  The constant voltage across the terminals of this source
"""
function ConstantVoltage(;name, V)
    @named oneport = OnePort()
    @unpack v, i = oneport
    pars = @parameters V=V
    eqs = [
        v ~ V
    ]

    extend(ODESystem(eqs, t, [], pars; name=name), oneport)
end

"""
```julia
function CosineVoltage(; name, offset=0.0, amplitude=1.0, frequency=1.0, starttime=0.0, phase=0.0)
```

A source in which the voltage across its terminals is a cosine function of time.

# States
- `v(t)`: [`V`]
  The voltage across this source, given by `p.v - n.v`

# Connectors
- `p`
  Positive port
- `n`
  Negative port

# Observables
- `offset`: [`V`]
  A constant offset added to the voltage output
- `amplitude`: [`V`]
  The amplitude of the cosine function
- `frequency`: [`Hz`]
  The frequency of the cosine function
- `starttime`: [`s`]
  The time at which the source starts functioning. Before this time, the voltage across
  its terminals is 0.
- `phase`: [`rad`]
  The phase offset of the cosine function
"""
function CosineVoltage(;name, offset=0.0, amplitude=1.0, frequency=1.0, start_time=0.0, phase=0.0)
    δ = 0.00001

    @named oneport = OnePort()
    @unpack v, i = oneport
    pars = @parameters begin
        offset=offset
        amplitude=amplitude
        frequency=frequency
        start_time=start_time
        phase=phase
    end
    eqs = [
        v ~ _cos_wave(t, frequency, amplitude, start_time, phase) * _step(t, δ, 1.0, start_time) + offset
    ]

    extend(ODESystem(eqs, t, [], pars; name=name), oneport)
end

"""
```julia
function ExpSineVoltage(; name, offset=0.0, amplitude=1.0, frequency=1.0, start_time=0.0, phase=0.0, damping=0.0)
```

A source in which the voltage across its terminals is a damped sine function of time.

# States
- `v(t)`: [`V`]
  The voltage across this source, given by `p.v - n.v`

# Connectors
- `p`
  Positive port
- `n`
  Negative port

# Parameters
- `offset`: [`V`]
  A constant offset added to the voltage output
- `amplitude`: [`V`]
  The amplitude of the damped sine function
- `frequency`: [`Hz`]
  The frequency of the damped sine function
- `start_time`: [`s`]
  The time at which the source starts functioning. Before this time, the voltage across
  its terminals is `offset`.
- `phase`: [`rad`]
  The phase offset of the damped sine function
- `damping_coef`: [`1/s`]
  Damping coefficient of the damped sine function
"""
function ExpSineVoltage(;name, offset=0.0, amplitude=1.0, frequency=1.0, start_time=0.0, phase=0.0, damping=0.0)
    δ = 0.00001
    @named oneport = OnePort()
    @unpack v, i = oneport
    pars = @parameters begin
        offset=offset
        amplitude=amplitude
        frequency=frequency
        start_time=start_time
        phase=phase
        damping=damping
    end
    eqs = [
        v ~ _damped_sine_wave(t, frequency, amplitude, start_time, phase, damping) * _step(t, δ, 1.0, start_time)
    ]

    extend(ODESystem(eqs, t, [], pars; name=name), oneport)
end

"""
```julia
function RampVoltage(; name, offset=0.0, start_time=0.0, duration=1.0, height=1.0)
```

A source in which the voltage across grows linearly from `offset` to `offset+height` over
the time interval `duration` starting at `start_time`

# States
- `v(t)`: [`V`]
  The voltage across this source, given by `p.v - n.v`

# Connectors
- `p`
  Positive port
- `n`
  Negative port

# Parameters
- `offset`: [`V`]
  A constant offset added to the voltage output
- `start_time`: [`s`]
  The time at which the voltage starts growing
- `duration`: [`s`]
  The duration of the ramp (`0.0` gives a step)
- `height`: [`V`]
  The amount that the voltage grows in the time interval
"""
function RampVoltage(;name, offset=0.0, start_time=0.0, duration=1.0, height=1.0)
    δ = 0.00001
    @named oneport = OnePort()
    @unpack v, i = oneport
    pars = @parameters begin
        offset=offset
        height=height
        start_time=start_time
        duration=duration
    end
    eqs = [
        v ~ _ramp(t, δ, start_time, start_time + duration, height) + offset
    ]

    extend(ODESystem(eqs, t, [], pars; name=name), oneport)
end

"""
```julia
function SineVoltage(; name, offset=0.0, amplitude=1.0, frequency=1.0, start_time=0.0, phase=0.0)
```

A source in which the voltage across its terminals is a sine function of time.

# States
- `v(t)`: [`V`]
  The voltage across this source, given by `p.v - n.v`

# Connectors
- `p`
  Positive port
- `n`
  Negative port

# Parameters
- `offset`: [`V`]
  A constant offset added to the voltage output
- `amplitude`: [`V`]
  The amplitude of the sine function
- `frequency`: [`Hz`]
  The frequency of the sine function
- `start_time`: [`s`]
  The time at which the source starts functioning. Before this time, the voltage across
  its terminals is `offset`.
- `phase`: [`rad`]
  The phase offset of the sine function
"""
function SineVoltage(;name, offset=0.0, amplitude=1.0, frequency=1.0, start_time=0.0, phase=0.0)
    δ = 0.00001

    @named oneport = OnePort()
    @unpack v, i = oneport
    pars = @parameters begin
        offset=offset
        amplitude=amplitude
        frequency=frequency
        start_time=start_time
        phase=phase
    end
    eqs = [
        v ~ _sin_wave(t, frequency, amplitude, start_time, phase) * _step(t, δ, 1.0, start_time) + offset
    ]

    extend(ODESystem(eqs, t, [], pars; name=name), oneport)
end

"""
```julia
function SquareVoltage(; name, offset=0.0, amplitude=1.0, frequency=1.0, start_time=0.0)
```

A source in which the voltage across its terminals is a square function of time.

# States
- `v(t)`: [`V`]
  The voltage across this source, given by `p.v - n.v`

# Connectors
- `p`
  Positive port
- `n`
  Negative port

# Parameters
- `offset`: [`V`]
  A constant offset added to the voltage output
- `amplitude`: [`V`]
  The amplitude of the square wave function
- `frequency`: [`Hz`]
  The frequency of the square wave function
- `start_time`: [`s`]
  The time at which the source starts functioning. Before this time, the voltage across
  its terminals is `offset`.
"""
function SquareVoltage(; name, offset=0.0, amplitude=1.0, frequency=1.0, start_time=0.0)
    δ = 0.0001

    @named oneport = OnePort()
    @unpack v, i = oneport
    pars = @parameters begin
        offset=offset
        amplitude=amplitude
        start_time=start_time
    end
    eqs = [
        v ~ _square_wave(t, δ, frequency, amplitude, start_time) * _step(t, δ, 1.0, start_time) + offset
    ]

    extend(ODESystem(eqs, t, [], pars; name=name), oneport)
end

"""
```julia
function StepVoltage(; name, offset=0.0, start_time=0.0, height=1.0)
```

A source in which the voltage across its terminals increases from `offset` to `offset+height` at
`start_time`

# States
- `v(t)`: [`V`]
  The voltage across this source, given by `p.v - n.v`

# Connectors
- `p`
  Positive port
- `n`
  Negative port

# Observables
- `offset`: [`V`]
  A constant offset added to the voltage output
- `start_time`: [`s`]
  The time at which the source starts functioning, and the voltage jumps
- `height`: [`V`]
  Magnitude of increase in voltage
"""
function StepVoltage(;name, offset=0.0, start_time=0.0, height=1.0)
    δ = 0.0001

    @named oneport = OnePort()
    @unpack v, i = oneport
    pars = @parameters begin
        offset=offset
        height=height
        start_time=start_time
    end
    eqs = [
        v ~ _step(t, δ, height, start_time) + offset
    ]

    extend(ODESystem(eqs, t, [], pars; name=name), oneport)
end

"""
```julia
function TriangularVoltage(; name, offset=0.0, amplitude=1.0, frequency=1.0, start_time=0.0)
```

A source in which the voltage across its terminals is a triangular function of time.

# States
- `v(t)`: [`V`]
  The voltage across this source, given by `p.v - n.v`

# Connectors
- `p`
  Positive port
- `n`
  Negative port

# Observables
- `offset`: [`V`]
  A constant offset added to the voltage output
- `amplitude`: [`V`]
  Amplitude of the triangular wave function
- `frequency`: [`Hz`]
  Frequency of the triangular wave function
- `start_time`: [`s`]
  The time at which the source starts functioning. Before this, the output of the source is
  `offset`
"""
function TriangularVoltage(; name, offset=0.0, amplitude=1.0, frequency=1.0, start_time=0.0)
    δ = 0.00001

    @named oneport = OnePort()
    @unpack v, i = oneport
    pars = @parameters begin
        offset=offset
        amplitude=amplitude
        frequency=frequency
        start_time=start_time
    end
    eqs = [
        v ~ _triangular_wave(t, δ, frequency, amplitude, start_time) * _step(t, δ, 1.0, start_time) + offset
    ]

    extend(ODESystem(eqs, t, [], pars; name=name), oneport)
end

# Current Sources ######################################################################################################
"""
```julia
function ConstantCurrent(; name, I = 1.0)
```

The source for an ideal constant current.

# States
- `i(t)`: [`A`]
  The current through this source, which is always constant

# Connectors
- `p`
  Positive pin
- `n`
  Negative pin

# Parameters:
- `I`: [`A`]
  The constant current through the terminals of this source
"""
function ConstantCurrent(;name, I = 1.0)
    @named oneport = OnePort()
    @unpack v, i = oneport
    pars = @parameters I=I
    eqs = [
        i ~ I
    ]

    extend(ODESystem(eqs, t, [], pars; name=name), oneport)
end

"""
```julia
function CosineCurrent(; name, offset=0.0, amplitude=1.0, frequency=1.0, starttime=0.0, phase=0.0)
```

A source in which the current through its terminals is a cosine function of time.

# States
- `i(t)`: [`A`]
  The current through this source

# Connectors
- `p`
  Positive port
- `n`
  Negative port

# Observables
- `offset`: [`A`]
  A constant offset added to the current output
- `amplitude`: [`A`]
  The amplitude of the cosine function
- `frequency`: [`Hz`]
  The frequency of the cosine function
- `starttime`: [`s`]
  The time at which the source starts functioning. Before this time, the current through
  its terminals is 0.
- `phase`: [`rad`]
  The phase offset of the cosine function
"""
function CosineCurrent(;name, offset=0.0, amplitude=1.0, frequency=1.0, start_time=0.0, phase=0.0)
    δ = 0.00001

    @named oneport = OnePort()
    @unpack v, i = oneport
    pars = @parameters begin
        offset=offset
        amplitude=amplitude
        frequency=frequency
        start_time=start_time
        phase=phase
    end
    eqs = [
        i ~ _cos_wave(t, frequency, amplitude, start_time, phase) * _step(t, δ, 1.0, start_time) + offset
    ]

    extend(ODESystem(eqs, t, [], pars; name=name), oneport)
end

"""
```julia
function ExpSineCurrent(; name, offset=0.0, amplitude=1.0, frequency=1.0, start_time=0.0, phase=0.0, damping=0.0)
```

A source in which the current through its terminals is a damped sine function of time.

# States
- `i(t)`: [`A`]
  The current through this source

# Connectors
- `p`
  Positive port
- `n`
  Negative port

# Parameters
- `offset`: [`A`]
  A constant offset added to the current output
- `amplitude`: [`A`]
  The amplitude of the damped sine function
- `frequency`: [`Hz`]
  The frequency of the damped sine function
- `start_time`: [`s`]
  The time at which the source starts functioning. Before this time, the current through
  its terminals is `offset`.
- `phase`: [`rad`]
  The phase offset of the damped sine function
- `damping_coef`: [`1/s`]
  Damping coefficient of the damped sine function
"""
function ExpSineCurrent(;name, offset=0.0, amplitude=1.0, frequency=1.0, start_time=0.0, phase=0.0, damping=0.0)
    δ = 0.00001
    @named oneport = OnePort()
    @unpack v, i = oneport
    pars = @parameters begin
        offset=offset
        amplitude=amplitude
        frequency=frequency
        start_time=start_time
        phase=phase
        damping=damping
    end
    eqs = [
        i ~ _damped_sine_wave(t, frequency, amplitude, start_time, phase, damping) * _step(t, δ, 1.0, start_time)
    ]

    extend(ODESystem(eqs, t, [], pars; name=name), oneport)
end

"""
```julia
function RampCurrent(; name, offset=0.0, start_time=0.0, duration=1.0, height=1.0)
```

A source in which the current grows linearly from `offset` to `offset+height` over
the time interval `duration` starting at `start_time`

# States
- `i(t)`: [`A`]
  The current through this source

# Connectors
- `p`
  Positive port
- `n`
  Negative port

# Parameters
- `offset`: [`A`]
  A constant offset added to the current output
- `start_time`: [`s`]
  The time at which the current starts growing
- `duration`: [`s`]
  The duration of the ramp (`0.0` gives a step)
- `height`: [`A`]
  The amount that the current grows in the time interval
"""
function RampCurrent(;name, offset=0.0, start_time=0.0, duration=1.0, height=1.0)
    δ = 0.00001
    @named oneport = OnePort()
    @unpack v, i = oneport
    pars = @parameters begin
        offset=offset
        height=height
        start_time=start_time
        duration=duration
    end
    eqs = [
        i ~ _ramp(t, δ, start_time, start_time + duration, height) + offset
    ]

    extend(ODESystem(eqs, t, [], pars; name=name), oneport)
end

"""
```julia
function SineCurrent(; name, offset=0.0, amplitude=1.0, frequency=1.0, start_time=0.0, phase=0.0)
```

A source in which the current through its terminals is a sine function of time.

# States
- `i(t)`: [`A`]
  The current through this source

# Connectors
- `p`
  Positive port
- `n`
  Negative port

# Parameters
- `offset`: [`A`]
  A constant offset added to the current output
- `amplitude`: [`V`]
  The amplitude of the sine function
- `frequency`: [`Hz`]
  The frequency of the sine function
- `start_time`: [`s`]
  The time at which the source starts functioning. Before this time, the current through
  its terminals is `offset`.
- `phase`: [`rad`]
  The phase offset of the sine function
"""
function SineCurrent(;name, offset=0.0, amplitude=1.0, frequency=1.0, start_time=0.0, phase=0.0)
    δ = 0.00001

    @named oneport = OnePort()
    @unpack v, i = oneport
    pars = @parameters begin
        offset=offset
        amplitude=amplitude
        frequency=frequency
        start_time=start_time
        phase=phase
    end
    eqs = [
        i ~ _sin_wave(t, frequency, amplitude, start_time, phase) * _step(t, δ, 1.0, start_time) + offset
    ]

    extend(ODESystem(eqs, t, [], pars; name=name), oneport)
end

"""
```julia
function SquareCurrent(; name, offset=0.0, amplitude=1.0, frequency=1.0, start_time=0.0)
```

A source in which the current through its terminals is a square function of time.

# States
- `i(t)`: [`A`]
  The current through this source

# Connectors
- `p`
  Positive port
- `n`
  Negative port

# Parameters
- `offset`: [`A`]
  A constant offset added to the current output
- `amplitude`: [`A`]
  The amplitude of the square wave function
- `frequency`: [`Hz`]
  The frequency of the square wave function
- `start_time`: [`s`]
  The time at which the source starts functioning. Before this time, the current through
  its terminals is `offset`.
"""
function SquareCurrent(; name, offset=0.0, amplitude=1.0, frequency=1.0, start_time=0.0)
    δ = 0.0001

    @named oneport = OnePort()
    @unpack v, i = oneport
    pars = @parameters begin
        offset=offset
        amplitude=amplitude
        start_time=start_time
    end
    eqs = [
        i ~ _square_wave(t, δ, frequency, amplitude, start_time) * _step(t, δ, 1.0, start_time) + offset
    ]

    extend(ODESystem(eqs, t, [], pars; name=name), oneport)
end

"""
```julia
function StepCurrent(; name, offset=0.0, start_time=0.0, height=1.0)
```

A source in which the current through its terminals increases from `offset` to `offset+height` at
`start_time`

# States
- `i(t)`: [`A`]
  The current through this source

# Connectors
- `p`
  Positive port
- `n`
  Negative port

# Observables
- `offset`: [`A`]
  A constant offset added to the current output
- `start_time`: [`s`]
  The time at which the source starts functioning, and the current jumps
- `height`: [`A`]
  Magnitude of increase in current
"""
function StepCurrent(;name, offset=0.0, start_time=0.0, height=1.0)
    δ = 0.0001

    @named oneport = OnePort()
    @unpack v, i = oneport
    pars = @parameters begin
        offset=offset
        height=height
        start_time=start_time
    end
    eqs = [
        i ~ _step(t, δ, height, start_time) + offset
    ]

    extend(ODESystem(eqs, t, [], pars; name=name), oneport)
end

"""
```julia
function TriangularCurrent(; name, offset=0.0, amplitude=1.0, frequency=1.0, start_time=0.0)
```

A source in which the current through its terminals is a triangular function of time.

# States
- `i(t)`: [`A`]
  The current through this source

# Connectors
- `p`
  Positive port
- `n`
  Negative port

# Observables
- `offset`: [`A`]
  A constant offset added to the current output
- `amplitude`: [`A`]
  Amplitude of the triangular wave function
- `frequency`: [`Hz`]
  Frequency of the triangular wave function
- `start_time`: [`s`]
  The time at which the source starts functioning. Before this, the output of the source is
  `offset`
"""
function TriangularCurrent(; name, offset=0.0, amplitude=1.0, frequency=1.0, start_time=0.0)
    δ = 0.00001

    @named oneport = OnePort()
    @unpack v, i = oneport
    pars = @parameters begin
        offset=offset
        amplitude=amplitude
        frequency=frequency
        start_time=start_time
    end
    eqs = [
        i ~ _triangular_wave(t, δ, frequency, amplitude, start_time) * _step(t, δ, 1.0, start_time) + offset
    ]

    extend(ODESystem(eqs, t, [], pars; name=name), oneport)
end


@testset "Saddle-Point Systems" begin

  @testset "Matrix tests" begin

    A1 = Float64[1 2; 2 1]
    B2 = Float64[2 3;-1 -1]
    B1 = B2'
    C = Matrix{Float64}(undef,2,2)
    C .= [5 -2; 3 -4];

    A = SaddleSystem(A1,B2,B1,C)

    # test inputs and outputs as tuples
    rhs = ([1.0,2.0],[3.0,4.0])
    sol = (zeros(2),zeros(2))

    ldiv!(sol,A,rhs)

    @test norm((A*sol)[1]-rhs[1]) < 1e-14
    @test norm((A*sol)[2]-rhs[2]) < 1e-14

    sol2 = A\rhs

    @test norm((A*sol2)[1]-rhs[1]) < 1e-14
    @test norm((A*sol2)[2]-rhs[2]) < 1e-14

    # test inputs and outputs as vectors
    rhs1, rhs2 = rhs
    rhsvec = [rhs1;rhs2]
    solvec = similar(rhsvec)

    solvec = A\rhsvec

    @test norm(A*solvec-rhsvec) < 1e-14

    # test with the other constructor
    Aother = SaddleSystem(A1,B2,B1,C,zeros(Float64,2),zeros(Float64,2))

    sol_other = Aother\rhs

    @test norm((Aother*sol_other)[1]-rhs[1]) < 1e-14
    @test norm((Aother*sol_other)[2]-rhs[2]) < 1e-14

    # Test with SaddleVector
    rhs = SaddleVector(rhs1,rhs2)
    As = SaddleSystem(A1,B2,B1,C,rhs)

    sol = As\rhs

    @test norm(state(As*sol)-state(rhs)) < 1e-14
    @test norm(constraint(As*sol)-constraint(rhs)) < 1e-14

  end

  nx = 130; ny = 130
  Lx = 2.0
  dx = Lx/(nx-2)
  w = Nodes(Dual,(nx,ny))
  q = Edges(Primal,w)

  L = plan_laplacian(size(w),with_inverse=true)

  n = 128
  θ = range(0,stop=2π,length=n+1)
  R = 0.5
  xb = 1.0 .+ R*cos.(θ[1:n])
  yb = 1.0 .+ R*sin.(θ[1:n])
  ds = (2π/n)*R
  X = VectorData(xb,yb)
  f = ScalarData(X)

  E = Regularize(X,dx;issymmetric=true)
  Hmat,Emat = RegularizationMatrix(E,f,w)

  @testset "Construction of linear maps" begin

    u = similar(w)
    wvec = vec(w)

    w .= rand(size(w)...)

    uvec = zeros(length(w))

    Lop = ConstrainedSystems.linear_map(L,w)

    u = L*w

    uvec = Lop*wvec

    @test ConstrainedSystems._wrap_vec(uvec,u) == u

    Linv = ConstrainedSystems.linear_inverse_map(L,w)

    yvec = Linv*wvec

    y = L\w

    @test ConstrainedSystems._wrap_vec(yvec,y) == y

    # point-wise operators
    fvec = vec(f)

    f[10] = 1.0
    Hop = ConstrainedSystems.linear_map(Hmat,f,w);

    y = Hmat*f

    yvec = Hop*fvec

    @test ConstrainedSystems._wrap_vec(yvec,y) == y

    Eop = ConstrainedSystems.linear_map(Emat,w,f)

    g = Emat*w

    gvec = Eop*vec(w)

    @test ConstrainedSystems._wrap_vec(gvec,g) == g

    # Test in-place function wrapping
    my_ip_fcn!(outp::Vector,inp::Vector) = outp .= 2.0*inp
    inp = rand(1000)
    outp = zero(inp)
    my_ip_fcn!(outp,inp)

    my_ip_lm! = ConstrainedSystems._create_vec_function!(my_ip_fcn!,outp,inp)
    outp2 = zero(inp)
    my_ip_lm!(outp2,inp)
    @test outp2 == outp

    curl_lm! = ConstrainedSystems._create_vec_function!(CartesianGrids.curl!,q,w)
    CartesianGrids.curl!(q,w)

    uvec .= wvec
    qvec = zeros(length(q))
    curl_lm!(qvec,uvec)

    @test ConstrainedSystems._wrap_vec(qvec,q) == q

    curl_lm! = ConstrainedSystems.linear_map(CartesianGrids.curl!,w,q)
    mul!(qvec,curl_lm!,uvec)

    @test ConstrainedSystems._wrap_vec(qvec,q) == q

  end

  ψb = ScalarData(X)
  w = Nodes(Dual,(nx,ny))
  ψb .= -(xb .- 1)
  f .= ones(Float64,n)*ds
  ψ = Nodes(Dual,w)

  nada = empty(f)

  @testset "Field operators" begin

    rhs = SaddleVector(w,ψb)

    A = SaddleSystem(L,Emat,Hmat,rhs)

    A.S*vec(f)

    A.S⁻¹*vec(f)

    sol = deepcopy(rhs)

    ldiv!(sol,A,rhs)

    sol2 = A\rhs

    @test state(sol2) == state(sol)
    @test constraint(sol2) == constraint(sol)

    rhs2 = A*sol2

    @test norm(constraint(rhs2)-ψb) < 1e-14

    fex = -2*cos.(θ[1:n])
    @test norm(constraint(sol)-fex*ds) < 0.02
    @test ψ[nx,65] ≈ -ψ[1,65]
    @test ψ[65,ny] ≈ ψ[65,1]

  end

  fv = VectorData(X)
  q = Edges(Primal,w)
  Hvmat,Evmat = RegularizationMatrix(E,fv,q)

  @testset "Vector force data" begin

    B₁ᵀ(f) = curl(Hvmat*f)
    B₂(w) = -(Evmat*(curl(L\w)))

    rhsv = SaddleVector(w,fv)
    A = SaddleSystem(I,B₂,B₁ᵀ,rhsv)

    sol = A\rhsv
    @test state(sol) == zero(w) && constraint(sol) == zero(fv)

  end

  Ẽ = Regularize(X,dx;weights=ds,filter=true)
  H̃mat = RegularizationMatrix(Ẽ,ψb,w)
  Ẽmat = InterpolationMatrix(Ẽ,w,ψb)
  Pmat = Ẽmat*H̃mat

  @testset "Filtering" begin

    rhs = SaddleVector(w,ψb)
    Afilt = SaddleSystem(L,Emat,Hmat,rhs,filter=Pmat)

    sol = Afilt\rhs

    fex = -2*cos.(θ[1:n])
    @test norm(constraint(sol)-fex*ds) < 0.01

  end

  @testset "Reduction to unconstrained system" begin

    op = linear_map(nothing,nada)

    @test size(op) == (0,0)
    @test op*nada == ()

    op = linear_map(nothing,w,nada)

    @test size(op) == (0,length(w))
    @test op*vec(w) == ()

    op = linear_map(nothing,nada,w)

    @test size(op) == (length(w),0)
    @test op*nada == vec(zero(w))

    rhsnc = SaddleVector(w,nada)

    Anc = SaddleSystem(L,rhsnc)

    sol = Anc\rhsnc
    ψ = state(sol)
    fnull = constraint(sol)

    @test ψ == L\w
    @test fnull == nada

    q.u .= 3;
    rhsqnc = SaddleVector(q,nada)
    Aqnc = SaddleSystem(I,rhsqnc)

    sol = Aqnc\rhsqnc

    q2 = state(sol)
    fnull = constraint(sol)

    @test all(q2.u .== 3.0) && all(q2.v .== 0.0)

  end

  @testset "Tuple of saddle point systems" begin

    rhs = SaddleVector(w,ψb)
    A = SaddleSystem(L,Emat,Hmat,rhs)
    sol = A\rhs
    ψ = state(sol)
    f = constraint(sol)

    rhsnc = SaddleVector(w,nada)
    Anc = SaddleSystem(L,rhsnc)

    sys = (A,Anc)
    sol1, sol2 = sys\(rhs,rhsnc)

    @test state(sol1) == ψ && constraint(sol1) == f && state(sol2) == zero(w) && constraint(sol2) == nada

    newrhs1,newrhs2 = sys*(sol1,sol2)

    @test newrhs1 == A*sol1 && newrhs2 == Anc*sol2

  end

  @testset "Recursive saddle point with vectors" begin
    A1 = Float64[1 2; 2 1]
    B21 = Float64[2 3]
    B11 = B21'
    C1 = Matrix{Float64}(undef,1,1)
    C1.= 5

    B22 = Float64[-1 -1 3]
    B12 = Float64[-1 -1 -2]'
    C2 = Matrix{Float64}(undef,1,1)
    C2.= -4

    rhs11 = [1.0,2.0]
    rhs12 = Vector{Float64}(undef,1)
    rhs12 .= 3.0
    #rhs1 = (rhs11,rhs12)
    rhs1 = [rhs11;rhs12]

    rhs2 = Vector{Float64}(undef,1)
    rhs2 .= 4.0

    rhs = (rhs1,rhs2)

    A = SaddleSystem(A1,B21,B11,C1)

    Abig = SaddleSystem(A,B22,B12,C2,rhs1,rhs2)

    sol = Abig\rhs

    out = Abig*sol

    @test norm(out[1]-rhs1) < 1e-14

  end

end