using QuantumComputer
using Random
using StatsBase
using Test

@testset "Register" begin
  register = QuantumComputer.Register(1)
  @test register.qubits == [1 0]

  register = QuantumComputer.Register(1, 1)
  @test register.qubits == [0 1]

  register = QuantumComputer.Register(2, 2)
  @test register.qubits == [0 1; 1 0]

  register = QuantumComputer.Register(10)
  @test size(register.qubits) == (10, 2)
end

@testset "Superposition" begin
  register = QuantumComputer.Register(3, 7)
  superposition = QuantumComputer.Superposition(register.qubits)
  @test length(superposition.state) == 2^3
  @test superposition.state[8] == 1
  @test sum(superposition.state) == 1
end

@testset "Measurement" begin
  @test_throws DomainError QuantumComputer.Measurement([1, 1], [1])
  @test_throws DomainError QuantumComputer.Measurement([1], [1, 1])
  @test_throws DomainError QuantumComputer.Measurement([1, 1], [1, 2])
  @test_throws DomainError QuantumComputer.Measurement([1, 2], [1, 1])
end

@testset "qubits_operated_on_by_unitary_matrix" begin
  identity_matrix = [1 0 0 0; 0 1 0 0; 0 0 1 0; 0 0 0 1]
  @test QuantumComputer.qubits_operated_on_by_unitary_matrix(identity_matrix) == 2
  @test_throws DomainError QuantumComputer.qubits_operated_on_by_unitary_matrix(identity_matrix[1:3, :])
  @test_throws DomainError QuantumComputer.qubits_operated_on_by_unitary_matrix(identity_matrix[:, 1:3])
end

@testset "Gate" begin
  identity_matrix = [1 0 0 0; 0 1 0 0; 0 0 1 0; 0 0 0 1]
  gate::QuantumComputer.Gate = QuantumComputer.Gate(identity_matrix)
  @test gate.matrix == identity_matrix
  @test gate.superposed_qubits_required == 2
end

@testset "gate_x" begin
  state::Array{Complex{Float64}, 1} = [1, 0]
  @test QuantumComputer.gate_x.matrix * state == [0, 1]

  state = [0, 1]
  @test QuantumComputer.gate_x.matrix * state == [1, 0]
end

@testset "gate_cx" begin
  gate::QuantumComputer.Gate = QuantumComputer.gate_cx(1, 3, 3)

  qubits::Matrix{Complex{Float64}} = [1 0; 1 0; 1 0]
  state = QuantumComputer.qubit_tensor_product(qubits)
  @test gate.matrix * state == state

  qubits = [1 0; 1 0; 0 1]
  state = QuantumComputer.qubit_tensor_product(qubits)
  @test gate.matrix * state == state

  qubits = [0 1; 1 0; 1 0]
  state = QuantumComputer.qubit_tensor_product(qubits)
  @test gate.matrix * state == [0, 0, 0, 0, 0, 1, 0, 0]

  qubits = [0 1; 1 0; 0 1]
  state = QuantumComputer.qubit_tensor_product(qubits)
  @test gate.matrix * state == [0, 0, 0, 0, 1, 0, 0, 0]

  gate = QuantumComputer.gate_cx(1, 2, 3)

  qubits = [0 1; 1 0; 0 1]
  state = QuantumComputer.qubit_tensor_product(qubits)
  @test gate.matrix * state == [0, 0, 0, 0, 0, 0, 0, 1]
end

@testset "gate_cnx" begin
  gate::QuantumComputer.Gate = QuantumComputer.gate_cnx([1, 2], 3, 3)

  qubits::Matrix{Complex{Float64}} = [0 1; 1 0; 1 0]
  state = QuantumComputer.qubit_tensor_product(qubits)
  @test gate.matrix * state == state

  qubits = [1 0; 0 1; 1 0]
  state = QuantumComputer.qubit_tensor_product(qubits)
  @test gate.matrix * state == state

  qubits = [1 0; 1 0; 1 0]
  state = QuantumComputer.qubit_tensor_product(qubits)
  @test gate.matrix * state == state

  qubits = [0 1; 0 1; 1 0]
  state = QuantumComputer.qubit_tensor_product(qubits)
  @test gate.matrix * state == [0, 0, 0, 0, 0, 0, 0, 1]
end

@testset "gate_swap" begin
  gate = QuantumComputer.gate_swap(1, 3, 3)

  qubits::Matrix{Complex{Float64}} = [1 0; 1 0; 1 0]
  state::Array{Complex{Float64}} = QuantumComputer.qubit_tensor_product(qubits)
  @test gate.matrix * state == state

  qubits = [0 1; 1 0; 0 1]
  state = QuantumComputer.qubit_tensor_product(qubits)
  @test gate.matrix * state == state

  qubits = [0 1; 1 0; 1 0]
  state = QuantumComputer.qubit_tensor_product(qubits)
  @test gate.matrix * state == [0, 1, 0, 0, 0, 0, 0, 0]

  qubits = [1 0; 1 0; 0 1]
  state = QuantumComputer.qubit_tensor_product(qubits)
  @test gate.matrix * state == [0, 0, 0, 0, 1, 0, 0, 0]
end

@testset "gate_multi_control" begin
  inner_gate = QuantumComputer.gate_swap(1, 3, 3)
  gate = QuantumComputer.gate_multi_control(inner_gate, 2, 5)

  qubits::Matrix{Complex{Float64}} = [1 0; 0 1; 0 1; 1 0; 1 0]
  state::Array{Complex{Float64}} = QuantumComputer.qubit_tensor_product(qubits)
  @test gate.matrix * state == state

  qubits = [0 1; 1 0; 1 0; 1 0; 0 1]
  state = QuantumComputer.qubit_tensor_product(qubits)
  @test gate.matrix * state == state

  qubits = [0 1; 0 1; 0 1; 1 0; 1 0]
  state = QuantumComputer.qubit_tensor_product(qubits)
  expected_qubits::Matrix{Complex{Float64}} = [0 1; 0 1; 1 0; 1 0; 0 1]
  expected_state::Array{Complex{Float64}} = QuantumComputer.qubit_tensor_product(expected_qubits)
  @test gate.matrix * state == expected_state

  qubits = [0 1; 0 1; 1 0; 0 1; 0 1]
  state = QuantumComputer.qubit_tensor_product(qubits)
  expected_qubits = [0 1; 0 1; 0 1; 0 1; 1 0]
  expected_state = QuantumComputer.qubit_tensor_product(expected_qubits)
  @test gate.matrix * state == expected_state

  extended_gate = QuantumComputer.gate_extension(gate, 2, 6)

  qubits = [1 0; 0 1; 0 1; 1 0; 0 1; 0 1]
  state = QuantumComputer.qubit_tensor_product(qubits)
  expected_qubits = [1 0; 0 1; 0 1; 0 1; 0 1; 1 0]
  expected_state = QuantumComputer.qubit_tensor_product(expected_qubits)
  @test extended_gate.matrix * state == expected_state

  qubits = [1 0; 0 1; 1 0; 1 0; 1 0; 0 1]
  state = QuantumComputer.qubit_tensor_product(qubits)
  @test extended_gate.matrix * state == state
end

@testset "gate_fourier_transform" begin
  gate = QuantumComputer.gate_fourier_transform(2)

  qubits::Matrix{Complex{Float64}} = [1 0; 1 0]
  state::Array{Complex{Float64}} = QuantumComputer.qubit_tensor_product(qubits)
  @test gate.matrix * state == [0.5, 0.5, 0.5, 0.5]

  qubits = [0 1; 0 1]
  state = QuantumComputer.qubit_tensor_product(qubits)
  @test gate.matrix * state ≈ [0.5, -0.5im, -0.5, 0.5im] atol=0.000000000000001
end

@testset "circuit_convert_to_gate" begin
  initial_value = 13
  constant = 12
  qubit_count = 5

  register = QuantumComputer.Register(qubit_count, initial_value)
  classical_register = QuantumComputer.ClassicalRegister(qubit_count)
  superposition = QuantumComputer.Superposition(register.qubits)
  adder = QuantumComputer.Circuits.constant_adder(constant, qubit_count)
  gate = QuantumComputer.circuit_convert_to_gate(adder)
  measurement = QuantumComputer.Measurement(Array(1:qubit_count), Array(1:qubit_count))

  circuit = QuantumComputer.Circuit()
  QuantumComputer.add_gate_to_circuit!(circuit, gate)
  QuantumComputer.add_measurement_to_circuit!(circuit, measurement)

  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)

  @test classical_register.value == (initial_value + constant) % 2^qubit_count
end

function test_constant_adder(initial_value, constant, qubit_count)
  register = QuantumComputer.Register(qubit_count, initial_value)
  classical_register = QuantumComputer.ClassicalRegister(qubit_count)
  superposition = QuantumComputer.Superposition(register.qubits)
  adder = QuantumComputer.Circuits.constant_adder(constant, qubit_count)
  measurement = QuantumComputer.Measurement(Array(1:qubit_count), Array(1:qubit_count))

  circuit = QuantumComputer.Circuit()
  QuantumComputer.add_subcircuit_to_circuit!(circuit, adder)
  QuantumComputer.add_measurement_to_circuit!(circuit, measurement)

  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)

  classical_register.value == (initial_value + constant) % 2^qubit_count
end

@testset "constant_adder" begin
  @test test_constant_adder(0, 3, 2)
  @test test_constant_adder(1, 1, 2)
  @test test_constant_adder(0, 63, 6)
  @test test_constant_adder(31, 31, 6)
end

function test_constant_adder_core(initial_value, constant, qubit_count)
  register = QuantumComputer.Register(qubit_count, initial_value)
  classical_register = QuantumComputer.ClassicalRegister(qubit_count)
  superposition = QuantumComputer.Superposition(register.qubits)
  qft = QuantumComputer.gate_fourier_transform(qubit_count)
  adder_core = QuantumComputer.Circuits.constant_adder_core(constant, qubit_count)
  inverse_qft = QuantumComputer.gate_invert(qft)
  measurement = QuantumComputer.Measurement(Array(1:qubit_count), Array(1:qubit_count))

  circuit = QuantumComputer.Circuit()
  QuantumComputer.add_gate_to_circuit!(circuit, qft)
  QuantumComputer.add_subcircuit_to_circuit!(circuit, adder_core)
  QuantumComputer.add_gate_to_circuit!(circuit, inverse_qft)
  QuantumComputer.add_measurement_to_circuit!(circuit, measurement)

  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)

  classical_register.value == (initial_value + constant) % 2^qubit_count
end

@testset "constant_adder_core" begin
  @test test_constant_adder_core(0, 3, 2)
  @test test_constant_adder_core(1, 1, 2)
  @test test_constant_adder_core(0, 63, 6)
  @test test_constant_adder_core(31, 31, 6)
end

@testset "shor2n3_controlled_controlled_modular_adder" begin
  a = 5
  n = 9

  modular_adder = QuantumComputer.Circuits.shor2n3_controlled_controlled_modular_adder(n, a)
  gate_qft = QuantumComputer.gate_fourier_transform(5)
  gate_inverse_qft = QuantumComputer.gate_invert(gate_qft)
  qft = QuantumComputer.gate_extension(gate_qft, 7, 11)
  inverse_qft = QuantumComputer.gate_extension(gate_inverse_qft, 7, 11)
  measurement = QuantumComputer.Measurement(Array(8:11), Array(8:11))

  circuit = QuantumComputer.Circuit()
  QuantumComputer.add_gate_to_circuit!(circuit, qft)
  QuantumComputer.add_subcircuit_to_circuit!(circuit, modular_adder)
  QuantumComputer.add_gate_to_circuit!(circuit, inverse_qft)
  QuantumComputer.add_measurement_to_circuit!(circuit, measurement)

  classical_register = QuantumComputer.ClassicalRegister(11)

  qubits::Matrix{Complex{Float64}} = [0 1; 0 1; 1 0; 1 0; 1 0; 1 0; 1 0; 1 0; 0 1; 1 0; 0 1]
  superposition = QuantumComputer.Superposition(qubits)
  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)
  @test classical_register.value == 1

  qubits = [0 1; 1 0; 1 0; 1 0; 1 0; 1 0; 1 0; 1 0; 0 1; 1 0; 0 1]
  superposition = QuantumComputer.Superposition(qubits)
  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)
  @test classical_register.value == 5

  qubits = [1 0; 0 1; 1 0; 1 0; 1 0; 1 0; 1 0; 1 0; 0 1; 1 0; 0 1]
  superposition = QuantumComputer.Superposition(qubits)
  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)
  @test classical_register.value == 5
end

@testset "shor2n3_controlled_modular_multiplier" begin
  a = 5
  n = 9

  modular_multiplier = QuantumComputer.Circuits.shor2n3_controlled_modular_multiplier(n, a)
  measurement = QuantumComputer.Measurement(Array(8:11), Array(8:11), 1)
  circuit = QuantumComputer.Circuit()

  QuantumComputer.add_subcircuit_to_circuit!(circuit, modular_multiplier)
  QuantumComputer.add_measurement_to_circuit!(circuit, measurement)

  classical_register = QuantumComputer.ClassicalRegister(11)

  qubits::Matrix{Complex{Float64}} = [1 0; 1 0; 1 0; 1 0; 1 0; 1 0; 1 0; 1 0; 0 1; 1 0; 0 1]
  superposition = QuantumComputer.Superposition(qubits)
  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)
  @test classical_register.value == 5

  qubits = [1 0; 1 0; 1 0; 1 0; 0 1; 1 0; 1 0; 1 0; 0 1; 1 0; 0 1]
  superposition = QuantumComputer.Superposition(qubits)
  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)
  @test classical_register.value == 5

  qubits = [0 1; 1 0; 1 0; 1 0; 0 1; 1 0; 1 0; 1 0; 0 1; 1 0; 0 1]
  superposition = QuantumComputer.Superposition(qubits)
  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)
  @test classical_register.value == 1

  qubits = [0 1; 1 0; 1 0; 0 1; 1 0; 1 0; 1 0; 1 0; 0 1; 1 0; 0 1]
  superposition = QuantumComputer.Superposition(qubits)
  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)
  @test classical_register.value == 6
end

@testset "shor2n3_controlled_ua" begin
  a = 5
  n = 9

  classical_register = QuantumComputer.ClassicalRegister(11)
  modular_multiplier = QuantumComputer.Circuits.shor2n3_controlled_ua(n, a)

  # first measure that we take no action when the control is not active
  circuit = QuantumComputer.Circuit()
  QuantumComputer.add_subcircuit_to_circuit!(circuit, modular_multiplier)
  measurement = QuantumComputer.Measurement(Array(8:11), Array(2:5))
  QuantumComputer.add_measurement_to_circuit!(circuit, measurement)

  qubits::Matrix{Complex{Float64}} = [1 0; 1 0; 1 0; 1 0; 0 1; 1 0; 1 0; 1 0; 1 0; 1 0; 1 0]
  superposition = QuantumComputer.Superposition(qubits)
  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)
  @test classical_register.value == 1

  circuit = QuantumComputer.Circuit()
  QuantumComputer.add_subcircuit_to_circuit!(circuit, modular_multiplier)
  measurement = QuantumComputer.Measurement(Array(8:11), Array(8:11))
  QuantumComputer.add_measurement_to_circuit!(circuit, measurement)

  qubits = [1 0; 1 0; 1 0; 1 0; 0 1; 1 0; 1 0; 1 0; 1 0; 1 0; 1 0]
  superposition = QuantumComputer.Superposition(qubits)
  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)
  @test classical_register.value == 0

  # next we test that b is reset to 0 after the operation
  circuit = QuantumComputer.Circuit()
  QuantumComputer.add_subcircuit_to_circuit!(circuit, modular_multiplier)
  measurement = QuantumComputer.Measurement(Array(8:11), Array(8:11))
  QuantumComputer.add_measurement_to_circuit!(circuit, measurement)

  qubits = [0 1; 1 0; 1 0; 1 0; 0 1; 1 0; 1 0; 1 0; 1 0; 1 0; 1 0]
  superposition = QuantumComputer.Superposition(qubits)
  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)
  @test classical_register.value == 0

  # now we test that that x is multiplied by a
  circuit = QuantumComputer.Circuit()
  QuantumComputer.add_subcircuit_to_circuit!(circuit, modular_multiplier)
  measurement = QuantumComputer.Measurement(Array(8:11), Array(2:5))
  QuantumComputer.add_measurement_to_circuit!(circuit, measurement)

  qubits = [0 1; 1 0; 1 0; 1 0; 0 1; 1 0; 1 0; 1 0; 1 0; 1 0; 1 0]
  superposition = QuantumComputer.Superposition(qubits)
  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)
  @test classical_register.value == 5

  # finally we test that the modulus is applied
  qubits = [0 1; 1 0; 1 0; 0 1; 1 0; 1 0; 1 0; 1 0; 1 0; 1 0; 1 0]
  superposition = QuantumComputer.Superposition(qubits)
  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)
  @test classical_register.value == 1

  # one more for good measure
  qubits = [0 1; 1 0; 1 0; 0 1; 0 1; 1 0; 1 0; 1 0; 1 0; 1 0; 1 0]
  superposition = QuantumComputer.Superposition(qubits)
  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)
  @test classical_register.value == 6
end

@testset "11^x mod 13 period finder" begin
  qubit_count = 5
  bit_count = 3
  sample_size = 256

  register = QuantumComputer.Register(qubit_count)
  classical_register = QuantumComputer.ClassicalRegister(bit_count)
  superposition = QuantumComputer.Superposition(register.qubits)
  period_finder = QuantumComputer.Circuits.period_finding_for_11x_mod_15()
  measurement = QuantumComputer.Measurement([1, 2, 3], [3, 4, 5], sample_size)

  circuit = QuantumComputer.Circuit()
  QuantumComputer.add_subcircuit_to_circuit!(circuit, period_finder)
  QuantumComputer.add_measurement_to_circuit!(circuit, measurement)

  Random.seed!(0)
  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)

  # this circuit doesn't find the period, but instead a related value (4).
  # here we test that over 256 samples, at least 128 resulted in binary '100'.
  # success of measurement depends on the random seed above
  @test countmap(measurement.samples)[4] > sample_size / 2
  @test classical_register.value == 4
  @test sum(map(x -> abs(x^2), superposition.state)) ≈ 1 atol=0.00000000001
end

@testset "grover" begin
  qubit_count = 10
  # should work for any valid input
  target_input = rand(0:(2^qubit_count - 1))

  circuit_oracle = QuantumComputer.Circuits.identification_oracle(qubit_count + 1, target_input)
  oracle = QuantumComputer.circuit_convert_to_gate(circuit_oracle)
  grover = QuantumComputer.Circuits.grover(oracle)

  register = QuantumComputer.Register(qubit_count + 1, 1)
  classical_register = QuantumComputer.ClassicalRegister(qubit_count, 0)
  superposition = QuantumComputer.Superposition(register.qubits)
  measurement = QuantumComputer.Measurement(Array(1:qubit_count), Array(1:qubit_count))

  circuit = QuantumComputer.Circuit()
  QuantumComputer.add_subcircuit_to_circuit!(circuit, grover)
  QuantumComputer.add_measurement_to_circuit!(circuit, measurement)

  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)
  @test classical_register.value == target_input
end

@testset "deutsch-jozsa, constant oracle f(x) = 1" begin
  qubit_count = 10

  oracle_circuit = QuantumComputer.Circuits.constant_oracle(qubit_count, true)
  oracle = QuantumComputer.circuit_convert_to_gate(oracle_circuit)
  deutsch_jozsa = QuantumComputer.Circuits.deutsch_jozsa(oracle)
  measurement = QuantumComputer.Measurement(Array(1:(qubit_count-1)), Array(1:(qubit_count-1)))

  register = QuantumComputer.Register(qubit_count, 1)
  classical_register = QuantumComputer.ClassicalRegister(qubit_count - 1, 0)
  superposition = QuantumComputer.Superposition(register.qubits)

  circuit = QuantumComputer.Circuit()
  QuantumComputer.add_subcircuit_to_circuit!(circuit, deutsch_jozsa)
  QuantumComputer.add_measurement_to_circuit!(circuit, measurement)
  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)

  @test classical_register.value == 0
end

@testset "deutsch-jozsa, constant oracle f(x) = 0" begin
  qubit_count = 10

  oracle_circuit = QuantumComputer.Circuits.constant_oracle(qubit_count, false)
  oracle = QuantumComputer.circuit_convert_to_gate(oracle_circuit)
  deutsch_jozsa = QuantumComputer.Circuits.deutsch_jozsa(oracle)
  measurement = QuantumComputer.Measurement(Array(1:(qubit_count-1)), Array(1:(qubit_count-1)))

  register = QuantumComputer.Register(qubit_count, 1)
  classical_register = QuantumComputer.ClassicalRegister(qubit_count - 1, 0)
  superposition = QuantumComputer.Superposition(register.qubits)

  circuit = QuantumComputer.Circuit()
  QuantumComputer.add_subcircuit_to_circuit!(circuit, deutsch_jozsa)
  QuantumComputer.add_measurement_to_circuit!(circuit, measurement)
  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)

  @test classical_register.value == 0
end

@testset "deutsch-jozsa, balanced oracle" begin
  qubit_count = 10

  oracle_circuit = QuantumComputer.Circuits.cx_balanced_oracle(qubit_count)
  oracle = QuantumComputer.circuit_convert_to_gate(oracle_circuit)
  deutsch_jozsa = QuantumComputer.Circuits.deutsch_jozsa(oracle)
  measurement = QuantumComputer.Measurement(Array(1:(qubit_count-1)), Array(1:(qubit_count-1)))

  register = QuantumComputer.Register(qubit_count, 1)
  classical_register = QuantumComputer.ClassicalRegister(qubit_count - 1, 0)
  superposition = QuantumComputer.Superposition(register.qubits)

  circuit = QuantumComputer.Circuit()
  QuantumComputer.add_subcircuit_to_circuit!(circuit, deutsch_jozsa)
  QuantumComputer.add_measurement_to_circuit!(circuit, measurement)
  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)

  @test classical_register.value != 0
end

@testset "bernstein-vazirani" begin
  qubit_count = 10
  secret::Int64 = rand(0:(2^(qubit_count-1)-1))

  oracle_circuit = QuantumComputer.Circuits.bitwise_product_oracle(secret, qubit_count)
  oracle = QuantumComputer.circuit_convert_to_gate(oracle_circuit)
  bernstein_vazirani = QuantumComputer.Circuits.bernstein_vazirani(oracle)
  measurement = QuantumComputer.Measurement(Array(1:(qubit_count-1)), Array(1:(qubit_count-1)))

  register = QuantumComputer.Register(qubit_count, 0)
  classical_register = QuantumComputer.ClassicalRegister(qubit_count - 1, 0)
  superposition = QuantumComputer.Superposition(register.qubits)

  circuit = QuantumComputer.Circuit()
  QuantumComputer.add_subcircuit_to_circuit!(circuit, bernstein_vazirani)
  QuantumComputer.add_measurement_to_circuit!(circuit, measurement)
  QuantumComputer.apply_circuit_to_superposition!(superposition, circuit, classical_register)

  @test classical_register.value == secret
end


using Interact, Distributed
import Plots, Printf, SearchBall
const pl = Plots
const pr = Printf
const sb = SearchBall

pl.gr()
addprocs(8)
@everywhere using SharedArrays

"""
robot lies in (0,0,0)
0 Rad -> orientation in positive x direction
"""
const v = 200.0
const ω = π/16.0
@everywhere function duration(x::Number, y::Number, start::Tuple{<:Number, <:Number, <:Number}=(0.0, 0.0, 0.0); velocity::Number=v, angular_velocity::Number=ω)
    (xi, yi, ai) = start
    a = mod(atan(y-yi, x-xi) - ai + π, 2π) - π
    sqrt((x-xi)^2 + (y-yi)^2)/velocity + abs(a)/angular_velocity
end

@everywhere const xs = -200:200
@everywhere const ys = -200:200
const maxxs = maximum(xs)
const maxys = maximum(ys)
const hmap1 = SharedArray(zeros(Float64, length(xs), length(ys)))
const hmap2 = SharedArray(zeros(Float64, length(xs), length(ys)))
const hmap3 = SharedArray(zeros(Float64, length(xs), length(ys)))

const init1 = (-100, 0, -π/2.0)
const init2 = (100, 0, π/2.0)

function calc_area!(hmap1, hmap2, hmap3, init1, init2, v=v, ω=ω)
    @inbounds @sync @distributed for ix in 1:length(xs)
        for iy in 1:length(ys)
            x = xs[ix]
            y = ys[iy]
            hmap1[iy, ix] = duration(x, y, init1, velocity=v, angular_velocity=ω)
            hmap2[iy, ix] = duration(x, y, init2, velocity=v, angular_velocity=ω)
            hmap3[iy, ix] = hmap1[iy, ix] < hmap2[iy, ix] ? 0.0 : 1.0
        end
    end
    # hmap1[3.5.<hmap1.<3.6] .= maximum(hmap1)
    # hmap2[3.5.<hmap2.<3.6] .= maximum(hmap2)
    return
end

pl.heatmap(hmap1, color=:ice_r, aspect_ratio=:equal, title=pr.@sprintf("v=%.1fm/s ω=%.0f°/s", v/100.0, rad2deg(ω)))
pl.heatmap(hmap2, color=:ice_r, aspect_ratio=:equal, title=pr.@sprintf("v=%.1fm/s ω=%.0f°/s", v/100.0, rad2deg(ω)))
pl.heatmap(hmap3, color=:ice, aspect_ratio=:equal, title=pr.@sprintf("v=%.1fm/s ω=%.0f°/s", v/100.0, rad2deg(ω)))

const mx = [maxxs, maxxs]
const my = [maxys, maxys]
const i11 = [init1[1], init1[1]]
const i12 = [init1[2], init1[2]]
const i21 = [init2[1], init2[1]]
const i22 = [init2[2], init2[2]]
ui = @manipulate for angle_red in -π:π/32:π, angle_blue in -π:π/32:π, vel in 1:400, ang in π/32:π/32:2π, show_field in [:voronoi, :red_player, :blue_player]
    calc_area!(hmap1, hmap2, hmap3, (init1[1], init1[2], angle_red), (init2[1], init2[2], angle_blue), vel, ang)
    if show_field == :voronoi
        myplot = pl.heatmap(hmap3, color=:ice, aspect_ratio=:equal, title=pr.@sprintf("v=%.1fm/s ω=%.0f°/s", vel/100.0, rad2deg(ang)), colorbar=false)
        # myplot = pl.plot(colorview(Gray, hmap3), aspect_ratio=:equal, title=pr.@sprintf("v=%.1fm/s ω=%.0f°/s", vel/100.0, rad2deg(ang)))
        dir1 = sb.rotate2d([30.0, 0.0], angle_red)
        pl.plot!(myplot, [0.0, dir1[1]] .+ mx .+ i11, [0.0, dir1[2]] .+ my .+ i12, lw=3, legend=false, color=:red)
        # pl.scatter!(myplot, init1[1], init1[2], lw=8, col=:white)
        dir2 = sb.rotate2d([30.0, 0.0], angle_blue)
        pl.plot!(myplot, [0.0, dir2[1]] .+ mx .+ i21, [0.0, dir2[2]] .+ my .+ i22, lw=3, legend=false, color=:blue)
        # pl.scatter!(myplot, init2[1], init2[2], lw=8, col=:black)
    elseif show_field == :red_player
        pl.heatmap(hmap1, color=:ice, aspect_ratio=:equal)
    elseif show_field == :blue_player
        pl.heatmap(hmap2, color=:ice, aspect_ratio=:equal)
    end
end


module GravitationalAcceleration

# This implements the gravitational acceleration

using Parameters, LaTeXStrings, Unitful
using ..Units
using GeoParams: AbstractMaterialParam
import Base.show, GeoParams.param_info
using ..MaterialParameters: MaterialParamsInfo

abstract type AbstractGravity{_T} <: AbstractMaterialParam end

export  compute_gravity,        # calculation routines
        ConstantGravity,        # constant
        param_info

# Constant Gravity -------------------------------------------------------
"""
    GravityConstant(g=9.81m/s^2)

Set a constant value for the gravitational acceleration:
```math
    g  = 9.81 m s^{-2}
```
"""
@with_kw_noshow struct ConstantGravity{_T,U}   <: AbstractGravity{_T}
    g::GeoUnit{_T,U}              =   9.81m/s^2               # gravitational acceleration
end
ConstantGravity(args...) = ConstantGravity(convert.(GeoUnit,args)...)

function param_info(s::ConstantGravity) # info about the struct
    return MaterialParamsInfo(Equation = L"g = 9.81 m s^{-2}" )
end

# Calculation routine
function compute_gravity(s::ConstantGravity{_T}) where _T
    @unpack_val g   = s

    return g
end

# Print info
function show(io::IO, d::ConstantGravity{_T})  where _T
    print(io, "Gravitational acceleration: g=$(UnitValue(d.g))")
end
#-------------------------------------------------------------------------

# Help info for the calculation routines
"""
compute_gravity(s:<AbstractGravity)

Returns the gravitational acceleration

"""
compute_gravity

end


import Observables
import Makie

q1 = Biquaternion(Quaternion(rand(4)), ℝ³(rand(3)))
scene = Makie.Scene()
len = 1 + rand()
width = 5
color = [:red, :green, :blue]
transparency = false
triad = Triad(q1,
              scene,
              length = len,
              width = width,
              color = color)
q2 = q1 * Biquaternion(Quaternion(rand(4)), ℝ³(rand(3)))
value1 = Observables.to_value(triad.observable)
update(triad, q2)
value2 = Observables.to_value(triad.observable)

@test isapprox(triad.q, q2)
@test isapprox(value1, value2) == false


#   Copyright (c) 2020-2022 Velexi Corporation
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

"""
AbstractScalar.jl defines the AbstractScalar type, interface, and core methods
"""

# --- Exports types and interface methods

# Types
export AbstractScalar

# Properties
export value

# --- Type definitions

"""
    AbstractScalar{<:AbstractFloat}

Supertype for all scalar types.

Interface
=========

Note: the return value of all methods should preserve the precision of its
AbstractScalar arguments (when possible).

Properties
----------

    value(B::AbstractScalar{T})::T where {T<:AbstractFloat}
"""
abstract type AbstractScalar{T<:AbstractFloat} <: AbstractBlade{T} end

# --- Method definitions
#
# Note: the following method definitions are no-op place holders and intended
#       to be extended.

"""
    value(B::AbstractScalar)::AbstractFloat

Return the value of `B` (with the same precision as `B`).
"""
function value end

# --- Method definitions for AbstractBlade interface functions

grade(B::AbstractScalar) = 0

basis(B::AbstractScalar) = one(value(B))

volume(B::AbstractScalar) = value(B)

-(B::AbstractScalar) = Scalar(-value(B))

inv(B::AbstractScalar) = 1 / B

# --- Method definitions for AbstractMultivector interface functions

"""
    dim(B::AbstractScalar)::Int

Return 0.

Notes
=====

* The convention that scalars have zero dimension is adopted because
  (1) scalars exist independently of all geometric algebras and (2) scalars
  are 0-dimensional entities.
"""
dim(B::AbstractScalar) = 0

Base.reverse(B::AbstractScalar) = B

"""
    dual(B::AbstractScalar, dim::Integer)::Pseudoscalar

Compute the dual of `B` relative to a real vector space having dimension `dim`.

Notes
=====

* An error is raised if the dimension of the embedding space is not explicitly
  specified.
"""
dual(B::AbstractScalar, dim::Integer) =
    mod(dim, 4) < 2 ?
        Pseudoscalar(dim, value(B)) :
        Pseudoscalar(dim, -value(B))

dual(B::AbstractScalar) =
    error("The dual of a scalar is not well-defined if `dim` is not " *
          "specified")

# --- Comparison methods

# ------ ==(B, C)

# B::AbstractScalar, C::AbstractScalar
==(B::AbstractScalar, C::AbstractScalar) = (value(B) == value(C))

# B::AbstractScalar, x::Real
# x::Real, B::AbstractScalar
==(B::AbstractScalar, x::Real) = (x == value(B))
==(x::Real, B::AbstractScalar) = (value(B) == x)

# ------ isapprox(B, C)

# B::AbstractScalar, C::AbstractScalar
isapprox(B::AbstractScalar{T1}, C::AbstractScalar{T2};
  atol::Real=0,
  rtol::Real=atol>0 ? 0 : max(√eps(T1), √eps(T2))) where {T1<:AbstractFloat,
                                                          T2<:AbstractFloat} =
    isapprox(value(B), value(C), atol=atol, rtol=rtol)

# B::AbstractScalar, x::Real
# x::Real, B::AbstractScalar
isapprox(B::AbstractScalar{T}, x::Real;
  atol::Real=0, rtol::Real=atol>0 ? 0 : sqrt(eps(T))) where {T<:AbstractFloat} =
    isapprox(x, value(B), atol=atol, rtol=rtol)

isapprox(x::Real, B::AbstractScalar{T};
  atol::Real=0, rtol::Real=atol>0 ? 0 : sqrt(eps(T))) where {T<:AbstractFloat} =
    isapprox(B, x, atol=atol, rtol=rtol)

# --- Utility methods

convert(::Type{T}, B::AbstractScalar) where {T<:AbstractFloat} =
    T == typeof(value(B)) ? B : Scalar{T}(value(B))


## gradient methods

"""
    gradient(inFLD::MeshArray,Γ::NamedTuple)

Compute spatial derivatives. Other methods:
```
gradient(inFLD::MeshArray,Γ::NamedTuple,doDIV::Bool)
gradient(inFLD::MeshArray,iDXC::MeshArray,iDYC::MeshArray)
```
"""
function gradient(inFLD::MeshArray,Γ::NamedTuple)
(dFLDdx, dFLDdy)=gradient(inFLD,Γ,true)
return dFLDdx, dFLDdy
end

function gradient(inFLD::MeshArray,Γ::NamedTuple,doDIV::Bool)

exFLD=exchange(inFLD,1)
dFLDdx=similar(inFLD)
dFLDdy=similar(inFLD)

for a=1:inFLD.grid.nFaces
  (s1,s2)=size(exFLD.f[a])
  tmpA=view(exFLD.f[a],2:s1-1,2:s2-1)
  tmpB=tmpA-view(exFLD.f[a],1:s1-2,2:s2-1)
  tmpC=tmpA-view(exFLD.f[a],2:s1-1,1:s2-2)
  if doDIV
    dFLDdx.f[a]=tmpB./Γ.DXC.f[a]
    dFLDdy.f[a]=tmpC./Γ.DYC.f[a]
  else
    dFLDdx.f[a]=tmpB
    dFLDdy.f[a]=tmpC
  end
end

return dFLDdx, dFLDdy
end

function gradient(inFLD::MeshArray,iDXC::MeshArray,iDYC::MeshArray)

exFLD=exchange(inFLD,1)
dFLDdx=similar(inFLD)
dFLDdy=similar(inFLD)

for a=1:inFLD.grid.nFaces
  (s1,s2)=size(exFLD.f[a])
  tmpA=view(exFLD.f[a],2:s1-1,2:s2-1)
  tmpB=tmpA-view(exFLD.f[a],1:s1-2,2:s2-1)
  tmpC=tmpA-view(exFLD.f[a],2:s1-1,1:s2-2)
  dFLDdx.f[a]=tmpB.*iDXC.f[a]
  dFLDdy.f[a]=tmpC.*iDYC.f[a]
end

return dFLDdx, dFLDdy
end

##

"""
    curl(u::MeshArray,v::MeshArray,Γ::NamedTuple)

Compute curl of a velocity field.
"""
function curl(u::MeshArray,v::MeshArray,Γ::NamedTuple)

	uvcurl=similar(Γ.XC)
	fac=exchange(1.0 ./Γ.RAZ,1)
	(U,V)=exchange(u,v,1)
    (DXC,DYC)=exchange(Γ.DXC,Γ.DYC,1)
	[DXC[i].=abs.(DXC[i]) for i in eachindex(U)]
	[DYC[i].=abs.(DYC[i]) for i in eachindex(V)]

	for i in eachindex(U)
    ucur=U[i][2:end,:]
    vcur=V[i][:,2:end]
    tmpcurl=ucur[:,1:end-1]-ucur[:,2:end]
    tmpcurl=tmpcurl-(vcur[1:end-1,:]-vcur[2:end,:])
    tmpcurl=tmpcurl.*fac[i][1:end-1,1:end-1]

		##still needed:
		##- deal with corners
		##- if putCurlOnTpoints

		tmpcurl=1/4*(tmpcurl[1:end-1,2:end]+tmpcurl[1:end-1,1:end-1]+
					tmpcurl[2:end,2:end]+tmpcurl[2:end,1:end-1])

		uvcurl[i]=tmpcurl
	end

	return uvcurl
end

## mask methods

function mask(fld::MeshArray)
fldmsk=mask(fld,NaN)
return fldmsk
end

"""
    mask(fld::MeshArray, val::Number)

Replace non finite values with val. Other methods:
```
mask(fld::MeshArray)
mask(fld::MeshArray, val::Number, noval::Number)
```
"""
function mask(fld::MeshArray, val::Number)
  fldmsk=similar(fld)
  for a=1:fld.grid.nFaces
    tmp1=copy(fld.f[a])
    replace!(x -> !isfinite(x) ? val : x, tmp1 )
    fldmsk.f[a]=tmp1
  end
  return fldmsk
end

function mask(fld::MeshArray, val::Number, noval::Number)
  fldmsk=similar(fld)
  for a=1:fld.grid.nFaces
    tmp1=copy(fld.f[a])
    replace!(x -> x==noval ? val : x, tmp1  )
    fldmsk.f[a]=tmp1
  end
  return fldmsk
end

## convergence methods

"""
    convergence(uFLD::MeshArray,vFLD::MeshArray)

Compute convergence of a vector field
"""
function convergence(uFLD::MeshArray,vFLD::MeshArray)

#important note:
#  Normally uFLD, vFLD should not contain any NaN;
#  if otherwise then something this may be needed:
#  uFLD=mask(uFLD,0.0); vFLD=mask(vFLD,0.0);

CONV=similar(uFLD)

(tmpU,tmpV)=exch_UV(uFLD,vFLD)
for a=1:tmpU.grid.nFaces
  (s1,s2)=size(uFLD.f[a])
  tmpU1=view(tmpU.f[a],1:s1,1:s2)
  tmpU2=view(tmpU.f[a],2:s1+1,1:s2)
  tmpV1=view(tmpV.f[a],1:s1,1:s2)
  tmpV2=view(tmpV.f[a],1:s1,2:s2+1)
  CONV.f[a]=tmpU1-tmpU2+tmpV1-tmpV2
end

return CONV
end

## smooth function

"""
    smooth(FLD::MeshArray,DXCsm::MeshArray,DYCsm::MeshArray,Γ::NamedTuple)

Smooth out scales below DXCsm / DYCsm via diffusion
"""
function smooth(FLD::MeshArray,DXCsm::MeshArray,DYCsm::MeshArray,Γ::NamedTuple)

#important note:
#input FLD should be land masked (NaN/1) by caller if needed

#get land masks (NaN/1):
mskC=fill(1.0,FLD) + 0.0 * mask(FLD)
(mskW,mskS)=gradient(FLD,Γ,false)
mskW=fill(1.0,FLD) + 0.0 * mask(mskW)
mskS=fill(1.0,FLD) + 0.0 * mask(mskS)

#replace NaN with 0. in FLD and land masks:
FLD=mask(FLD,0.0)
mskC=mask(mskC,0.0)
mskW=mask(mskW,0.0)
mskS=mask(mskS,0.0)

#get inverse grid spacing:
iDXC=similar(FLD)
iDYC=similar(FLD)
for a=1:FLD.grid.nFaces
  iDXC.f[a]=1.0./Γ.DXC.f[a]
  iDYC.f[a]=1.0./Γ.DYC.f[a]
end

#Before scaling the diffusive operator ...
tmp0=DXCsm*iDXC*mskW;
tmp00=maximum(tmp0);
tmp0=DYCsm*iDYC*mskS;
tmp00=max(tmp00,maximum(tmp0));

#... determine a suitable time period:
nbt=ceil(1.1*2*tmp00^2);
dt=1.;
T=nbt*dt;
#println("nbt="*"$nbt")

#diffusion operator times DYG / DXG
KuxFac=mskW*DXCsm*DXCsm/T/2.0*Γ.DYG
KvyFac=mskS*DYCsm*DYCsm/T/2.0*Γ.DXG

#time steping factor:
dtFac=dt*mskC/Γ.RAC

#loop:
for it=1:nbt
  (dTdxAtU,dTdyAtV)=gradient(FLD,iDXC,iDYC);
  tmpU=similar(FLD)
  tmpV=similar(FLD)
  for a=1:FLD.grid.nFaces
      tmpU.f[a]=dTdxAtU.f[a].*KuxFac.f[a];
      tmpV.f[a]=dTdyAtV.f[a].*KvyFac.f[a];
  end
  tmpC=convergence(tmpU,tmpV);
  for a=1:FLD.grid.nFaces
      FLD.f[a]=FLD.f[a]-dtFac.f[a].*tmpC.f[a];
  end
end

#Apply land mask (NaN/1) to end result:
mskC=mask(mskC,NaN,0.0)
FLD=mskC*FLD

return FLD

end

## ThroughFlow function

"""
    ThroughFlow(VectorField,IntegralPath,Γ::NamedTuple)

Compute transport through an integration path
"""
function ThroughFlow(VectorField,IntegralPath,Γ::NamedTuple)

    #Note: vertical intergration is not always wanted; left for user to do outside

    U=VectorField["U"]
    V=VectorField["V"]

    nd=ndims(U)
    #println("nd=$nd and d=$d")

    n=fill(1,4)
    tmp=size(U)
    n[1:nd].=tmp[1:nd]

    haskey(VectorField,"factors") ? f=VectorField["factors"] : f=Array{String,1}(undef,0)
    haskey(VectorField,"dimensions") ? d=VectorField["dimensions"] : d=Array{String,1}(undef,nd)

    #a bit of a hack to distinguish gcmfaces v gcmarray indexing:
    isdefined(U,:fIndex) ? ndoffset=1 : ndoffset=0
    length(d)!=nd+ndoffset ? error("inconsistent specification of dims") : nothing

    trsp=Array{Float64}(undef,1,n[3],n[4])
    do_dz=sum(f.=="dz")
    do_dxory=sum(f.=="dxory")

    for i3=1:n[3]
        #method 1: quite slow
        #mskW=IntegralPath.mskW
        #do_dxory==1 ? mskW=mskW*Γ.DYG : nothing
        #do_dz==1 ? mskW=Γ.DRF[i3]*mskW : nothing
        #mskS=IntegralPath.mskS
        #do_dxory==1 ? mskS=mskS*Γ.DXG : nothing
        #do_dz==1 ? mskS=Γ.DRF[i3]*mskS : nothing
        #
        #method 2: less slow
        tabW=IntegralPath.tabW
        tabS=IntegralPath.tabS
        for i4=1:n[4]
            #method 1: quite slow
            #trsp[1,i3,i4]=sum(mskW*U[:,:,i3,i4])+sum(mskS*V[:,:,i3,i4])
            #
            #method 2: less slow
            trsp[1,i3,i4]=0.0
            for k=1:size(tabW,1)
                (a,i1,i2,w)=tabW[k,:]
                do_dxory==1 ? w=w*Γ.DYG.f[a][i1,i2] : nothing
                do_dz==1 ? w=w*Γ.DRF[i3] : nothing
                isdefined(U,:fIndex) ? u=U.f[a,i3,i4][i1,i2] : u=U.f[a][i1,i2,i3,i4]
                trsp[1,i3,i4]=trsp[1,i3,i4]+w*u
            end
            for k=1:size(tabS,1)
                (a,i1,i2,w)=tabS[k,:]
                do_dxory==1 ? w=w*Γ.DXG.f[a][i1,i2] : nothing
                do_dz==1 ? w=w*Γ.DRF[i3] : nothing
                isdefined(V,:fIndex) ? v=V.f[a,i3,i4][i1,i2] : v=V.f[a][i1,i2,i3,i4]
                trsp[1,i3,i4]=trsp[1,i3,i4]+w*v
            end
        end
    end

    nd+ndoffset<4 ? trsp=dropdims(trsp,dims=3) : nothing
    nd+ndoffset<3 ? trsp=dropdims(trsp,dims=2) : nothing
    nd+ndoffset==2 ? trsp=trsp[1] : nothing

    return trsp
end

## LatitudeCircles function

"""
    LatitudeCircles(LatValues,Γ::NamedTuple)

Compute integration paths that follow latitude circles
"""
function LatitudeCircles(LatValues,Γ::NamedTuple)
    LatitudeCircles=Array{NamedTuple}(undef,length(LatValues))
    for j=1:length(LatValues)
        mskCint=1*(Γ.YC .>= LatValues[j])
        mskC,mskW,mskS=edge_mask(mskCint)
        LatitudeCircles[j]=(lat=LatValues[j],
        tabC=MskToTab(mskC),tabW=MskToTab(mskW),tabS=MskToTab(mskS))
    end
    return LatitudeCircles
end

function MskToTab(msk::MeshArray)
  n=Int(sum(msk .!= 0)); k=0
  tab=Array{Int,2}(undef,n,4)
  for i in eachindex(msk)
    a=msk[i]
    b=findall( a .!= 0)
    for ii in eachindex(b)
      k += 1
      tab[k,:]=[i,b[ii][1],b[ii][2],a[b[ii]]]
    end
  end
  return tab
end

"""
    edge_mask(mskCint::MeshArray)

Compute edge mask (mskC,mskW,mskS) from domain interior mask (mskCint).
This is used in `LatitudeCircles` and `Transect`.
"""
function edge_mask(mskCint::MeshArray)
  mskCint=1.0*mskCint

  #treat the case of blank tiles:
  #mskCint[findall(RAC.==0)].=NaN

  mskC=similar(mskCint)
  mskW=similar(mskCint)
  mskS=similar(mskCint)

  mskCint=exchange(mskCint,1)

  for i in eachindex(mskCint)
      tmp1=mskCint[i]
      # tracer mask:
      tmp2=tmp1[2:end-1,1:end-2]+tmp1[2:end-1,3:end]+
      tmp1[1:end-2,2:end-1]+tmp1[3:end,2:end-1]
      mskC[i]=1((tmp2.>0).&(tmp1[2:end-1,2:end-1].==0))
      # velocity masks:
      mskW[i]=tmp1[2:end-1,2:end-1] - tmp1[1:end-2,2:end-1]
      mskS[i]=tmp1[2:end-1,2:end-1] - tmp1[2:end-1,1:end-2]
  end

  #treat the case of blank tiles:
  #mskC[findall(isnan.(mskC))].=0.0
  #mskW[findall(isnan.(mskW))].=0.0
  #mskS[findall(isnan.(mskS))].=0.0

  return mskC,mskW,mskS
end

##

"""
    Transect(name,lons,lats,Γ)

Compute integration paths that follow a great circle between two geolocations give by `lons`, `lats`.
"""
function Transect(name,lons,lats,Γ)
  x0,y0,z0,R=rotate_points(lons,lats)
  x,y,z=rotate_XCYC(Γ,R)
  mskCint=1.0*(z.>0)
  mskCedge,mskWedge,mskSedge=edge_mask(mskCint)
  mskCedge,mskWedge,mskSedge=shorter_paths!((x,y,z),(x0,y0,z0),(mskCedge,mskWedge,mskSedge))
  tabC=MskToTab(mskCedge)
  tabW=MskToTab(mskWedge)
  tabS=MskToTab(mskSedge)
  return (name=name,tabC=tabC,tabW=tabW,tabS=tabS)
end

##

nanmean(x) = mean(filter(!isnan,x))
nanmean(x,y) = mapslices(nanmean,x,dims=y)

"""
    UVtoUEVN(u,v,G::NamedTuple)

1. Interpolate to grid cell centers (uC,vC)
2. Convert to `Eastward/Northward` components (uE,vN)
"""
function UVtoUEVN(u::MeshArray,v::MeshArray,G::NamedTuple)
    u[findall(G.hFacW[:,1].==0)].=NaN
    v[findall(G.hFacS[:,1].==0)].=NaN

    (u,v)=exch_UV(u,v); uC=similar(u); vC=similar(v)
    for iF=1:u.grid.nFaces
        tmp1=u[iF][1:end-1,:]; tmp2=u[iF][2:end,:]
        uC[iF]=reshape(nanmean([tmp1[:] tmp2[:]],2),size(tmp1))
        tmp1=v[iF][:,1:end-1]; tmp2=v[iF][:,2:end]
        vC[iF]=reshape(nanmean([tmp1[:] tmp2[:]],2),size(tmp1))
    end

    return uC.*G.AngleCS-vC.*G.AngleSN, uC.*G.AngleSN+vC.*G.AngleCS
end

"""
    UVtoTransport(U,V,G::NamedTuple)

Convert e.g. velocity (m/s) to transport (m^3/s) by multiplying by `DRF` and `DXG`,`DYG`.
"""
function UVtoTransport(U::MeshArray,V::MeshArray,G::NamedTuple)
  uTr=deepcopy(U)
  vTr=deepcopy(V)
  for i in eachindex(U)
      tmp1=U[i]; tmp1[(!isfinite).(tmp1)] .= 0.0
      tmp1=V[i]; tmp1[(!isfinite).(tmp1)] .= 0.0
      uTr[i]=G.DRF[i[2]]*U[i].*G.DYG[i[1]]
      vTr[i]=G.DRF[i[2]]*V[i].*G.DXG[i[1]]
  end
  return uTr,vTr
end


module TestUnits

using Modia
@usingModiaPlot

UnitTest = Model(

    T = 0.2u"s",
    u0 = 8.5u"kg",
    k = 1u"kg^-1",
    m = 10u"kg",
    F = 20u"N",

    x = Var(init = 5.5u"kg"),    # 5500.0u"g"),
    v = Var(init = 1u"m/s"),
    equations = :[
        T*der(x) + x = u0
        y = x
        z = x*u"kg^-1"

        F = m * a
        der(v) = a ]
)

model = @instantiateModel(UnitTest)

simulate!(model, Tsit5(), stopTime = 1.0, requiredFinalStates = [8.479786015016273, 2.999999999999999])

#plot(model, ["T", "x", "der(x)", "y", "a", "der(v)"])
plot(model, ["x", "der(x)", "y", "a", "der(v)"])

end


# nonparametric regression, using pre-generated weights
# can be local constant, local linear, or local quadratic
#
# y is n X p
# x is n X dimx
# xeval is neval X dimx
# weights is n X neval (different weights for each eval. point)
#
# weights should sum to one by columns, they are typical
# nonparametric weights from a kernel.

function npreg()
    println("npreg(), with no arguments, runs Pkg.dir/Econometrics/examples/npreg_example.jl")
    println("examine that file for information on how to call npreg.jl")
    include(Pkg.dir()*"/Econometrics/examples/npreg_example.jl")
end

function npreg(y::Array{Float64,2}, x::Array{Float64,2}, xeval::Array{Float64,2}, weights::Array{Float64,2}, order::Int64=1, do_median::Bool=false, do_ci::Bool=false)
    weights = sqrt.(weights)
    neval, dimx = size(xeval)
    n, dimy = size(y)
    # local constant
    if order==0
        X = ones(n,1)
        Xeval = ones(neval,1)
    elseif order==1
    # local linear
        X = [ones(n,1) x]
        Xeval = [ones(neval,1) xeval]
    else
    # local quadratic
        stacked = [x; xeval]
        # cross products
        CP = Array{Float64}(n+neval, Int((dimx-1)*dimx/2))
        cpind = 0
        for i = 1:dimx-1
            for j = (i+1):dimx
                cpind += 1
                CP[:,cpind] = stacked[:,[i]].*stacked[:,[j]]
            end
        end
        ZZ = [ones(n+neval,1) stacked CP]
        X = view(ZZ,1:n,:)
        Xeval = view(ZZ,(n+1):n+neval,:)
    end
    # do the fit
    yhat = Array{Float64}(neval, dimy)
    for i = 1:neval
        WX = weights[:,i] .* X
        Wy = weights[:,i] .* y
        b = WX\Wy
        yhat[i,:] = Xeval[[i],:]*b
    end
    return yhat
end


module EMD

using Dierckx

export emd

function emd(x::Vector{Float64})
    c = x[:]
    N = length(c)
    imf = []
    while true
        h = c
        sd = 1

        maxmin = []

        while sd > 0.3
            d = diff(h)
            maxmin = []
            for i in 1:N-2
                if d[i] == 0
                    push!(maxmin, i)
                elseif sign(d[i]) != sign(d[i+1])
                    push!(maxmin, i + 1)
                end
            end

            if length(maxmin) < 2
                break
            end

            if maxmin[1] > maxmin[2]
                maxes = maxmin[1:2:end]
                mins  = maxmin[2:2:end]
            else
                maxes = maxmin[2:2:end]
                mins  = maxmin[1:2:end]
            end

            maxes = [1, maxes..., N];
            mins  = [1, mins..., N];

            k = min(3, length(maxes)-1)
            maxenv = Spline1D(maxes, h[maxes]; k=k)(1:N)
            minenv = Spline1D(mins, h[mins]; k=k)(1:N)

            m = (maxenv + minenv) / 2
            prevh = h[:]
            h = h - m

            eps = 1e-7
            sd = sum(((prevh - h).^2) ./ (prevh.^2 .+ eps))
        end

        push!(imf, h)

        if length(maxmin) < 2
            break
        end

        c -= h
    end

    imf, c
end

end # module


nt = 1000
seed = 99
m = ACQR_simula(nt,seed)
y = m["y"]
A = m["A"]
C = m["C"]
Q = m["Q"]
R = m["R"]
x10 = m["m1"]
P10 = m["P1"]

###########################
print("Testing ACQR_ksmoother: ")
xtt,Ptt,xtt1,Ptt1,et,St,Kt,loglik = ACQR_kfilter(y,A,C,Q,R,x10,P10)
xtN,PtN,Pt1tN = ACQR_ksmoother(A,xtt,Ptt,xtt1,Ptt1)
if sum(xtN.^2) > 1e-6
    println("OK")
end


### Global surface temperature map

for i = 1:length(casenames)

    fig, ax = plt[:subplots](1, 1, figsize=(12,8))

    casename = casenames[i]
    cmap = plt[:get_cmap]("bwr")
    cmap[:set_over]("r")
    cmap[:set_under]("b")
    cntrs = range(900, step=4, stop=1100)
    Dataset("$nc_file_dir/$casename.h0.nc", "r") do ds

        PSL = mean(mreplace(ds["PSL"][:, :, t_rng]), dims=(3,))[:, :, 1] / 100.0
        #plt[:contourf](lon, lat, mask'[:,:], [.45,.55], cmap=cmap, antialiased=false, extend="both")
        #plt[:contourf](lon, lat, mask'[:,:], [.45,.55], cmap=cmap, antialiased=false, extend="both")
#        cmapping = plt[:contourf](lon, lat, PSL'[:,:], cntrs, cmap=cmap, antialised=true, extend="both")
        plt[:contour](lon, lat, mask'[:,:], [.5], colors="#cccccc")
        cs = plt[:contour](lon, lat, PSL'[:,:], cntrs, antialised=true, colors="k")
        #cs = plt[:contour](lon, lat, TREFHT'[:,:], cntrs, colors="k")
        plt[:clabel](cs, fmt="%d")

        ax[:set_title](format("Mean sea surface pressure ({:d} years): $casename", years))
        ax[:set_xlabel]("Lat [deg]")
        ax[:set_ylabel]("Lon [deg]")
    end

    fig[:savefig]("$(casename)_map_sea_surface_pressure.png", dpi=200)

end


### DIC ###

# Batch DIC
function deviance(param::T, loglike::Function) where T
  return -2 * loglike(param)
end

function DIC(params::Vector{T}, loglike::Function, paramMean::Function) where T
  D(param) = deviance(param, loglike)

  Dbar = mean(D.(params))
  paramBar = paramMean(params)
  pD = Dbar - D(paramBar)

  return Dbar + pD
end

# Streaming DIC
"""
Methods to implement:
- updateParams(d::DICstream{T}, param::P) where {T, P}, to add param to d.paramSum
- paramMeanCompute(d::DICstream{T}) where T, to compute the mean of the parameters
- loglike(param::P) computes the loglikelihood based on the parameters (param)
- convert(state::State) convets the current state to param::Param
"""
mutable struct DICstream{T}
  paramSum::T
  Dsum::Float64
  counter::Int

  DICstream{T}(tmp) where T = new(tmp, 0, 0)
end

"""
convert: converts state to param
"""
function updateDIC(x::DICstream{T}, state::S, updateParams::Function,
                   loglike::Function, convert::Function) where {T, S}
  @assert x.counter >= 0

  param = convert(state)

  # update Dsum
  x.Dsum += deviance(param, loglike)

  # update paramSum
  if x.counter == 0
    x.paramSum = deepcopy(param)
  else
    updateParams(x, param)
  end

  # update counter
  x.counter += 1
end

"""
DIC is a measure of fit. Higher is better.

DIC = Dmean + pD, where
- Dmean is a measure of fit. Higher is better.
- pD is a measure of (non) complexity. Higher is better.
"""
function computeDIC(x::DICstream{T}, loglike::Function,
                    paramMeanCompute::Function;
                    return_Dmean_pD::Bool=false) where {T}
  Dmean = x.Dsum / x.counter
  paramMean = paramMeanCompute(x)
  pD = Dmean - deviance(paramMean, loglike)

  if return_Dmean_pD
    return Dmean, pD
  else
    return Dmean + pD
  end
end


struct ComteButucea <: LinearEstimator
    Q::BinnedCalibrator
    h::Float64
    m::Int64
    target::InferenceTarget
end

function (c::ComteButucea)(x)
    (c.Q)(x)
end

pretty_label(cb::ComteButucea)  = "Butucea-Comte"

inference_target(cb::ComteButucea) = cb.target

default_bandwidth(::Type{ComteButucea}, target::MarginalDensityTarget, n) = sqrt(log(n))
default_bandwidth(::Type{ComteButucea}, target::LFSRNumerator, n) = sqrt(log(n)*2/3)

default_bandwidth(::Type{ComteButucea}, n) = sqrt(log(n))
default_bandwidth(::Type{ComteButucea}, target, n) = default_bandwidth(ComteButucea, n)

function comte_butucea(z, h, target::InferenceTarget)
   f(t) = real(exp(im*t*z)*cf(target,-t)/cf(Normal(),t))
   res = hquadrature(f, -h, +h)[1]
   res/2/π
end

function ComteButucea(target::InferenceTarget, m::Int64, marginal_grid;
    h= default_bandwidth(ComteButucea, target, m))
    Qc = BinnedCalibrator(marginal_grid, comte_butucea.(marginal_grid, h, target))
    ComteButucea(Qc, h, m, target)
end

function estimate(Xs,cb::ComteButucea)
    mean(cb.Q.(Xs))
end

function estimate(Xs, ::Type{ComteButucea}, target::InferenceTarget, marginal_grid)
    cb = ComteButucea(target, length(Xs), marginal_grid)
    estimate(Xs, cb)
end
# probably should learn how to do the above properly via FFT?

function estimate(Xs, target::InferenceTarget, ::Type{ComteButucea}; K_grid=10_000,
    marginal_grid = collect(range(minimum(Xs), stop=maximum(Xs), length=K_grid)), args...)
    n = length(Xs)
    Qc = BinnedCalibrator(marginal_grid, comte_butucea.(marginal_grid, target, n; args...))
    mean(Qc.(Xs))
end


module BeamColumn

using DiffEqOperators: CenteredDifference
using LinearAlgebra
using NLsolve

using ..Mesh

export define, solve, Model, governing_equations

Base.@kwdef mutable struct Model

   A::Union{Array{Float64}, Nothing} = nothing
   Ix::Union{Array{Float64}, Nothing} = nothing
   Iy::Union{Array{Float64}, Nothing} = nothing
   Io::Union{Array{Float64}, Nothing} = nothing
   J::Union{Array{Float64}, Nothing} = nothing
   Cw::Union{Array{Float64}, Nothing} = nothing
   E::Union{Array{Float64}, Nothing} = nothing
   ν::Union{Array{Float64}, Nothing} = nothing
   G::Union{Array{Float64}, Nothing} = nothing
   xc::Union{Array{Float64}, Nothing} = nothing
   yc::Union{Array{Float64}, Nothing} = nothing
   xs::Union{Array{Float64}, Nothing} = nothing
   ys::Union{Array{Float64}, Nothing} = nothing
   xo::Union{Array{Float64}, Nothing} = nothing  #distance from centroid to shear center
   yo::Union{Array{Float64}, Nothing} = nothing  #distance from centroid to shear center
   ax::Union{Array{Float64}, Nothing} = nothing #distance from shear center to qy applied load location
   ay::Union{Array{Float64}, Nothing} = nothing #distance from shear center to qx applied load location
   kx::Union{Array{Float64}, Nothing} = nothing
   ky::Union{Array{Float64}, Nothing} = nothing
   kϕ::Union{Array{Float64}, Nothing} = nothing
   hx::Union{Array{Float64}, Nothing} = nothing #distance from centroid to ky spring location
   hy::Union{Array{Float64}, Nothing} = nothing #distance from centroid to kx spring location

   qx::Union{Array{Float64}, Nothing} = nothing
   qy::Union{Array{Float64}, Nothing} = nothing
   P::Union{Array{Float64}, Nothing} = nothing

   end_boundary_conditions::Union{Array{Int64}, Nothing} = nothing
   supports::Union{Array{Float64}, Nothing} = nothing

   z::Union{Array{Float64}, Nothing} = nothing
   m::Union{Array{Int64}, Nothing} = nothing

   K::Union{Matrix{Float64}, Nothing} = nothing
   F::Union{Array{Float64}, Nothing} = nothing

   free_dof::Union{Array{Int64}, Nothing} = nothing

   u::Union{Array{Float64}, Nothing} = nothing
   v::Union{Array{Float64}, Nothing} = nothing
   ϕ::Union{Array{Float64}, Nothing} = nothing

end

function calculate_derivative_operators(dz)

   #Define the number of nodes.
   num_nodes = length(dz) + 1

   #Add extra dz on each end for padding nodes used in the CenteredDifference function.
   dz = [dz[1]; dz; dz[end]]

   #Calculate the 4th derivative operator.
   nth_derivative = 4
   derivative_order = 2
   Azzzz = CenteredDifference(nth_derivative, derivative_order, dz, num_nodes)

   #Convert the operator to a matrix.
   Azzzz = Array(Azzzz)

   #Trim off the ghost nodes.
   Azzzz = Azzzz[:,2:end-1]

   #Calculate the 2nd derivative operator.
   nth_derivative = 2
   derivative_order = 2
   Azz = CenteredDifference(nth_derivative, derivative_order, dz, num_nodes)
   Azz = Array(Azz)
   Azz = Azz[:,2:end-1]

   return Azzzz, Azz

end

function calculate_boundary_stencils(bc_flag, h, nth_derivative)

   #Calculate boundary conditions stencils without ghost nodes using
   #Jorge M. Souza, "Boundary Conditions in the Finite Difference Method"
   #https://cimec.org.ar/ojs/index.php/mc/article/download/2662/2607

   taylor_coeffs =  [1; h; h^2/2; h^3/6; h^4/24]

   RHS = zeros(Float64, 5)
   #row1*[A;B;C;D;E]*u2
   #row2*[A;B;C;D;E]*u2'
   #row3*[A;B;C;D;E]*u2''
   #row4*[A;B;C;D;E]*u2'''
   #row5*[A;B;C;D;E]*u2''''

   if bc_flag == 1 #simply supported end

      LHS = [1 1 1 1 0
         -1 0 1 2 0
         1 0 1 4 2
         -1 0 1 8 -6
         1 0 1 16 12]

      RHS[nth_derivative + 1] = (1/taylor_coeffs[nth_derivative + 1])
      boundary_stencil = LHS \ RHS
      boundary_stencil = ((boundary_stencil[1:4]),(zeros(4)))  #since u''=0

   elseif bc_flag == 2 #fixed end

      LHS = [1 1 1 1 0
         -1 0 1 2 1
         1 0 1 4 -2
         -1 0 1 8 3
         1 0 1 16 -4]

      RHS[nth_derivative+1] = (1/taylor_coeffs[nth_derivative + 1])
      boundary_stencil = LHS\RHS
      boundary_stencil = ((boundary_stencil[1:4]),(zeros(4)))  #since u'=0

   elseif bc_flag == 3 #free end
                #u'' u'''
      LHS = [1 1 1  0   0
           0 1 2  0   0
           0 1 4  0.5 0
           0 1 8  0   6
           0 1 16 0  0]

      RHS[nth_derivative+1] = (1/taylor_coeffs[nth_derivative+1])
      boundary_stencil1 = LHS\RHS  #at free end
      boundary_stencil1 = boundary_stencil1[1:3]   #u''=u'''=0

      # use simply supported BC to find stencil at one node in from free end
      LHS = [1 1 1 1 0
         -1 0 1 2 0
         1 0 1 4 2
         -1 0 1 8 -6
         1 0 1 16 12]

      RHS[nth_derivative+1] = (1/taylor_coeffs[nth_derivative+1])
      boundary_stencil2 = LHS\RHS #at one node over from free end
      boundary_stencil2 = boundary_stencil2[1:4]
      boundary_stencil = ((boundary_stencil1), (boundary_stencil2))  #two stencils are calculated

   end

   return boundary_stencil

end

function apply_end_boundary_conditions(A, end_boundary_conditions, nth_derivative, dz)

   #Consider the left end boundary condition.  Swap out interior stencil for boundary stencil.
   h = dz[1]
   bc_flag = end_boundary_conditions[1]
   boundary_stencil = calculate_boundary_stencils(bc_flag, h, nth_derivative)

   A[1,:] .= 0.0
   A[2,:] .= 0.0

   if (bc_flag == 1) | (bc_flag == 2)   #make this cleaner, combine
      A[2,1:length(boundary_stencil[1])] = boundary_stencil[1]
   else
      A[1,1:length(boundary_stencil[1])] = boundary_stencil[1]
      A[2,1:length(boundary_stencil[2])] = boundary_stencil[2]
   end

   #Consider the right end boundary condition.
   h = dz[end]
   bc_flag = end_boundary_conditions[2]
   boundary_stencil = calculate_boundary_stencils(bc_flag,h,nth_derivative)

   A[end,:] .= 0.0
   A[end-1,:] .= 0.0

   if (bc_flag == 1) | (bc_flag == 2)
      A[end-1,(end-(length(boundary_stencil[1])-1)):end] = reverse(boundary_stencil[1])
   else
      A[end,end-(length(boundary_stencil[1])-1):end] = reverse(boundary_stencil[1])
      A[end-1,end-(length(boundary_stencil[2])-1):end] = reverse(boundary_stencil[2])
   end

   return A

end

function discretize(member_definitions)

   #Define discretization along beam (z) and assign a member type to each node (m).
   dz, z, m = Mesh.define_line_element(member_definitions)

   return z, m

end

function define(z, m, member_definitions, section_properties, material_properties, kx, ky, kϕ, hx, hy, qx, qy, P, ax, ay, end_boundary_conditions, supports)

   #This is the user interface mapping.
   #L(1) dL(2) section_properties(3) material_properties(4)

   dz = diff(z)

   #define property vectors
   A =  Mesh.create_line_element_property_array(member_definitions, m, dz, section_properties, 3, 1)
   Ix = Mesh.create_line_element_property_array(member_definitions, m, dz, section_properties, 3, 2)
   Iy = Mesh.create_line_element_property_array(member_definitions, m, dz, section_properties, 3, 3)
   J = Mesh.create_line_element_property_array(member_definitions, m, dz, section_properties, 3, 4)
   Cw = Mesh.create_line_element_property_array(member_definitions, m, dz, section_properties, 3, 5)
   xc = Mesh.create_line_element_property_array(member_definitions, m, dz, section_properties, 3, 6)
   yc = Mesh.create_line_element_property_array(member_definitions, m, dz, section_properties, 3, 7)
   xs = Mesh.create_line_element_property_array(member_definitions, m, dz, section_properties, 3, 8)
   ys = Mesh.create_line_element_property_array(member_definitions, m, dz, section_properties, 3, 9)

   xo = -(xc .- xs)
   yo = yc .- ys

   Io = Ix .+ Iy .+ A .* (xo.^2 + yo.^2)

   E = Mesh.create_line_element_property_array(member_definitions, m, dz, material_properties, 4, 1)
   ν = Mesh.create_line_element_property_array(member_definitions, m, dz, material_properties, 4, 2)
   G = E./(2 .*(1 .+ ν))

   model = Model(A, Ix, Iy, Io, J, Cw, E, ν, G, xc, yc, xs, ys, xo, yo, ax, ay, kx, ky, kϕ, hx, hy, qx, qy, P, end_boundary_conditions, supports, z, m, nothing, nothing, nothing, nothing, nothing, nothing)

   return model

end

function governing_equations(model)

   #Define the number of nodes.
   num_nodes = length(model.z)

   #Calculate the derivative operators.
   dz = diff(model.z)
   Azzzz,Azz = calculate_derivative_operators(dz)

   #Apply left and right end boundary condition stencils to derivative operators.
   nth_derivative = 4
   Azzzz = apply_end_boundary_conditions(Azzzz, model.end_boundary_conditions, nth_derivative, dz)

   nth_derivative = 2
   Azz = apply_end_boundary_conditions(Azz, model.end_boundary_conditions, nth_derivative, dz)

   #build identity matrix for ODE operations
   AI = Matrix(1.0I,num_nodes,num_nodes)

   #build operator matrix that doesn't update with load

   #start building operator matrix, LHS first
   A11 = zeros(num_nodes,num_nodes)
   A12 = zeros(num_nodes,num_nodes)
   A13 = zeros(num_nodes,num_nodes)
   A21 = zeros(num_nodes,num_nodes)
   A22 = zeros(num_nodes,num_nodes)
   A23 = zeros(num_nodes,num_nodes)
   A31 = zeros(num_nodes,num_nodes)
   A32 = zeros(num_nodes,num_nodes)
   A33 = zeros(num_nodes,num_nodes)

   #calculate operator quantities on LHS  AU=B
   for i = 1:num_nodes
      A11[i,:] = model.E .* model.Iy .* Azzzz[i,:] .+ model.P .* Azz[i,:] .+ model.kx .* AI[i,:]
      A13[i,:] = -model.kx .* model.hy .* AI[i,:]
      A22[i,:] = model.E .* model.Ix .* Azzzz[i,:] .+ model.P .* Azz[i,:] .+ model.ky .* AI[i,:]
      A23[i,:] = model.ky .* model.hx .* AI[i,:]
      A31[i,:] = -model.kx .* model.hy .* AI[i,:]
      A32[i,:] = model.ky .* model.hx .* AI[i,:]
      A33[i,:] = model.E .* model.Cw .* Azzzz[i,:] .-(model.G .* model.J .- (model.P .* model.Io ./ model.A)).*Azz[i,:] .+model.kx .* model.hy.^2 .*AI[i,:] .+ model.ky .* model.hx.^2 .*AI[i,:] .+ model.kϕ .* AI[i,:] .+  model.qx .* model.ax .* AI[i,:] - .+ model.qy .* model.ay .* AI[i,:]
   end

   #calculate RHS of AU=B
   B1 = model.qx
   B2 = model.qy
   B3 = model.qx .* model.ay .+ model.qy .* model.ax

   #reduce problem to free dof
   #Find fixed dof.
   fixed_dof = []
   for i = 1:length(model.supports)

      if i == 1

            fixed_dof = findall(x->abs(x-model.supports[i])<=10e-6, model.z)

      elseif i != 1

            new_fixed_dof = findall(x->abs(x-model.supports[i])<=10e-6, model.z)
            fixed_dof = [fixed_dof; new_fixed_dof]

      end

   end

   model.free_dof = setdiff(1:num_nodes,fixed_dof)

   #Define stiffness matrix.
   model.K = [A11[model.free_dof,model.free_dof] A12[model.free_dof,model.free_dof] A13[model.free_dof,model.free_dof];
      A21[model.free_dof,model.free_dof] A22[model.free_dof,model.free_dof] A23[model.free_dof,model.free_dof];
      A31[model.free_dof,model.free_dof] A32[model.free_dof,model.free_dof] A33[model.free_dof,model.free_dof]]

   #Define external force vector.
   model.F = [B1[model.free_dof]; B2[model.free_dof]; B3[model.free_dof]]

   return model

end

function residual!(R, U, K, F)

   for i=1:length(F)
    R[i] = transpose(K[i,:]) * (U) - F[i]
   end

   return R

end

function solve(model)

   #Set up solution matrices from governing equations.
   model = governing_equations(model)

   #Define the number of nodes.
   num_nodes=length(model.z)

   #Define the deformation vectors.
   u=zeros(Float64, num_nodes)
   v=zeros(Float64, num_nodes)
   ϕ=zeros(Float64, num_nodes)

   #Define the deformation initial guess for the nonlinear solver.
   deformation_guess = model.K \ model.F    #consider revising this for large systems, it might be slow...

   #Solve for the deformations.
   solution = nlsolve((R, U) ->residual!(R, U, model.K, model.F), deformation_guess)

   #Pull the displacements and twist from the solution results.
   u[model.free_dof] = solution.zero[1:length(model.free_dof)]
   v[model.free_dof] = solution.zero[length(model.free_dof)+1:2*length(model.free_dof)]
   ϕ[model.free_dof] = solution.zero[2*length(model.free_dof)+1:3*length(model.free_dof)]

   #Add deformations to data structure.
   model.u = u
   model.v = v
   model.ϕ = ϕ

   return model

end

end #module


using LinearAlgebra
using Images
using Printf

# load data
xs = load("./examples/ChessBoard/inpainting_sparse.png")
xr = load("./examples/ChessBoard/inpainting_low_rank.png")

# size
m, n = size(xs)

# preprocessing on xs
xs = Gray.(xs); xs = convert(Array{Float64}, xs)
γ = xs[2,2]; xs = xs .- γ # normalize
idx = findall(x->abs(x)<=5e-2, xs); xs[idx] .= 0
idx = findall(x->x!=0, xs); xs[idx] .= maximum(xs)
ks = count(x->x!=0, xs)
xs = 3*vec(xs)

# preprocessing on xr
xr = Gray.(xr); xr = convert(Array{Float64}, xr)
F = svd(xr); kr = 2
xr = F.U[:, 1:kr] * Diagonal(F.S[1:kr]) * F.Vt[1:kr, :]
xr = 0.5*vec(xr)

# observation
b = xs + xr

# write to file
f = open("/home/zhenan/Github/AtomicOpt.jl/examples/ChessBoard/ChessBoardData.txt", "w")
l = length(b)
@printf(f, "xs = ")
for i in 1:l
    @printf(f, "%.3f ", xs[i])
end
println(f, "\n")
@printf(f, "xr = ")
for i in 1:l
    @printf(f, "%.3f ", xr[i])
end
println(f, "\n")
@printf(f, "b = ")
for i in 1:l
    @printf(f, "%.3f ", b[i])
end
println(f, "\n")
@printf(f, "ks = %d\n", ks)
@printf(f, "kr = %d\n", kr)
@printf(f, "m = %d\n", m)
@printf(f, "n = %d\n", n)
close(f)


include("../util.jl")

using JuMP, Gurobi

function TEAVAR(env,
                edges,
                capacity,
                flows,
                demand,
                beta,
                k,
                T,
                Tf,
                scenarios,
                scenario_probs,
				outputfilename;
                explain=false,
                verbose=false,
                utilization=false,
                average=false)

    nedges = length(edges)
    nflows = length(flows)
    ntunnels = length(T)
    nscenarios = length(scenarios)
    p = scenario_probs

	println(Dates.format(now(), "HH:MM:SS"), ": #edges ", nedges, " #flows ", nflows, " #tunnels ", ntunnels, " #scenarios ", nscenarios)

    #CREATE TUNNEL SCENARIO MATRIX
    X  = ones(nscenarios,ntunnels)
    for s in 1:nscenarios
        for t in 1:ntunnels
            if size(T[t],1) == 0
                X[s,t] = 0
            else
                for e in 1:nedges
                    if scenarios[s][e] == 0
                        back_edge = findfirst(x -> x == (edges[e][2],edges[e][1]), edges)
                        if in(e, T[t]) || in(back_edge, T[t])
                        # if in(e, T[t])
                            X[s,t] = 0
                        end
                    end
                end
            end
        end
    end

	println(Dates.format(now(), "HH:MM:SS"), ": created tunnel scenario matrix")

    #CREATE TUNNEL EDGE MATRIX
    L = zeros(ntunnels, nedges)
    for t in 1:ntunnels
        for e in 1:nedges
            if in(e, T[t])
                L[t,e] = 1
            end
        end
    end

    println(Dates.format(now(), "HH:MM:SS"), ": created tunnel edge matrix")

    model = Model(solver=GurobiSolver(env, OutputFlag=1))
    # flow per commodity per path variables
    @variable(model, a[1:nflows, 1:k] >= 0, basename="a", category=:SemiCont)
    # alpha variable
    @variable(model, alpha >= 0, basename="alpha", category=:SemiCont)
    # maximum flow lost in that scenario
    @variable(model, umax[1:nscenarios] >= 0, basename="umax")
    # flow lost per commod per scenario
    @variable(model, u[1:nscenarios, 1:nflows] >= 0, basename="u")

    # capacity constraints for final flow assigned to "a" variables
    for e in 1:nedges
        @constraint(model, sum(a[f,t] * L[Tf[f][t],e] for f in 1:nflows, t in 1:size(Tf[f],1)) <= capacity[e])
    end

      # FLOW LEVEL LOSS
      @expression(model, satisfied[s=1:nscenarios, f=1:nflows], sum(a[f,t] * X[s,Tf[f][t]] for t in 1:size(Tf[f],1)) / demand[f])

      for s in 1:nscenarios
          for f in 1:nflows
              # @constraint(model, (demand[f] - sum(a[f,t] * X[s,Tf[f][t]] for t in 1:size(Tf[f],1))) / demand[f] <= u[s,f])
              @constraint(model, u[s,f] >= 1 - satisfied[s,f])
          end
      end

      for s in 1:nscenarios
          if average
              @constraint(model, umax[s] + alpha >= (sum(u[s,f] for f in 1:nflows)) / nflows)
              # @constraint(model, umax[s] + alpha >= avg_loss[s])
          else
              for f in 1:nflows
                  @constraint(model, umax[s] + alpha >= u[s,f])
              end
          end
      end
      @objective(model, Min, alpha + (1 / (1 - beta)) * sum((p[s] * umax[s] for s in 1:nscenarios)))

    println(Dates.format(now(), "HH:MM:SS"), ": ready to solve")

    solve(model)

    println(Dates.format(now(), "HH:MM:SS"), ": solver finished; explaining")

    if (explain)
        println("Runtime: ", getsolvetime(model))
	println("beta: ", beta)
	println("#flows: ", nflows)
	println("#edges: ", nedges)
	println("#tunnels: ", ntunnels)
	println("#demands: ", nflows, " total demand=", sum(demand[f] for f in 1:nflows))
	println("#scenarios: ", nscenarios, " total_prob= ", sum(p[s] for s in 1:nscenarios))
	println("total_scenario_prob= ", sum(p[s] for s in 1:nscenarios))

	result_u = getvalue(u)
	result_umax = getvalue(umax)
	result_satisfied = getvalue(satisfied)
	result_alpha = getvalue(alpha)
	result_a = getvalue(a)

	# compute allocations and weights
	flow_allocation_1 = Array{Float64}(undef, nflows)
	numtunnel_perflow = 0
	for f in 1:nflows
		numtunnel_perflow = max(numtunnel_perflow, size(Tf[f], 1))
	end
	flow_tunnel_weights = Matrix{Float64}(undef, nflows, numtunnel_perflow)
	all_flow_alloc = 0

	for f in 1:nflows
		flow_totalalloc = 0
		for t in 1:size(Tf[f], 1)
			flow_totalalloc += result_a[f,t]
		end

		all_flow_alloc += flow_totalalloc
		flow_allocation_1[f] = (1 - result_alpha) * min(flow_totalalloc, demand[f])
		for t in 1:size(Tf[f], 1)
			weight = 0
			if flow_totalalloc > 0
				weight = result_a[f, t]/ flow_totalalloc
			end
			flow_tunnel_weights[f, t] = weight
		end
	end
	flow_allocation_1[flow_allocation_1 .< 0] .= 0

	println("totalalloc= ", all_flow_alloc, " netalloc1= ", sum(flow_allocation_1[f] for f in 1:nflows))

	output = open(outputfilename, "w")
	write(output, "---flow allocation with alpha---\n")
	writedlm(output, flow_allocation_1, ',')
	write(output, "---weights per flow tunnel---\n")
	writedlm(output, flow_tunnel_weights, ',')
	flush(output)
	close(output)

        printResults(getobjectivevalue(model), result_alpha, result_a, result_u, result_umax, edges, scenarios, T, Tf, L, capacity, p, demand, verbose=verbose, utilization=utilization)
    end

    return (result_alpha, getobjectivevalue(model), result_a, result_umax, getsolvetime(model))
end


include(joinpath(pathof(GroupsCore), "..", "..", "test", "conformance_test.jl"))

@testset "GroupsCore interface conformance" begin
    G = PermGroup(perm"(1,2,3,4,5)", perm"(1,2)")
    test_Group_interface(G)
    test_GroupElement_interface(rand(G, 2)...)

    @test gens(SymmetricGroup(1)) == Perm{Int}[]
    @test gens(SymmetricGroup(2)) == [perm"(1,2)"]
    @test gens(SymmetricGroup(3)) == [perm"(1,2,3)", perm"(1,2)(3)"]
end


using Lasso
using Random, SparseArrays
DATADIR = joinpath(dirname(@__FILE__), "data")

function diffmatslow(order, n)
    D = sparse(SparseArrays.spdiagm_internal(0 => fill(-1., n-1), 1=>fill(1., n-1))...)
    for i = 1:order
        D = sparse(SparseArrays.spdiagm_internal(0 => fill(-1., n-i-1), 1=>fill(1., n-i-1))...)*D
    end
    D
end

# Test that DifferenceMatrix behaves as expected
@testset "DifferenceMatrix" begin
    @testset "order = $(order)" for order in (1, 2, 3)
        D1 = diffmatslow(order, 100)
        D2 = Lasso.TrendFiltering.DifferenceMatrix{Float64}(order, 100)
        x = randn(testrng, 100)
        @test D2'D2 == D1'D1
        @test D2*x ≈ D1*x
        @test D2'*x[1:size(D1, 1)] ≈ D1'*x[1:size(D1, 1)]
    end
end

# Test against results from glmgen
lakehuron = readcsvmat(joinpath(DATADIR, "LakeHuron.csv"); header=true)[:, 3]
@testset "TrendFiltering" begin
    @testset "order = $(order)" for order in (1, 2, 3)
        @testset "lambda = $(lambda)" for lambda in (1., 10., 100.)
            @test coef(fit(TrendFilter, lakehuron, order, lambda; tol=1e-9)) ≈
                  vec(readcsvmat(joinpath(DATADIR, "LakeHuron_order_$(order)_lambda_$(lambda).csv")))
        end
    end
end


"""
xxx
# Examples
```@example

```
"""
struct LTISource
  position::Vector{Float64}
  transmission ::Function
end

# Method
function (𝚽::LTISource)(𝛏₀::Vector{Float64}, t₀::Float64)
   𝐩ₛ, p = 𝚽.position, 𝚽.transmission
   delay = distBetween(𝐩ₛ,𝛏₀)/lightSpeed
   return A(delay) * p(t₀-delay)
end

# DISPLAY
Base.show(io::IO, x::LTISource) = print(io, "LTI omnidirectional source")


product_manifold_split(::Type{PM{M1,M2}}) where {M1,M2} = M1, M2

dim(::Union{PM{M1,M2},Type{PM{M1,M2}}}) where {M1,M2} = dim(M1) + dim(M2)
embdim(::Union{PM{M1,M2},Type{PM{M1,M2}}}) where {M1,M2} = embdim(M1) + embdim(M2)

dims(::Union{PM{M1,M2},Type{PM{M1,M2}}}) where {M1,M2} = dim(M1), dim(M2)
embdims(::Union{PM{M1,M2},Type{PM{M1,M2}}}) where {M1,M2} = embdim(M1), embdim(M2)

product_chart_split(C::Chart{Tuple{I,J},PM{M1,M2}}) where {M1,M2,I,J} =
    Chart{I,M1}(), Chart{J,M2}()
product_chart_merge(C1::Chart{I,M1}, C2::Chart{J,M2}) where {M1,M2,I,J} =
    Chart{Tuple{I,J},PM{M1,M2}}()

isglobal(::Chart{Tuple{I,J},PM{M1,M2}}) where {M1,M2,I,J} =
    isglobal(Chart{I,M1}()) && isglobal(Chart{J,M2}())

function product_manifold_emb_splitview(pe, ::Type{M}) where M <: ProductManifold
    m = embdim(M)
    m1, m2 = embdims(M)
    x1e, x2e = pe[SVector{m1}(1:m1)], pe[SVector{m2}(m1+1:m)]
end

function product_manifold_chart_splitview(p, ::Type{M}) where M <: ProductManifold
    m = dim(M)
    m1, m2 = dims(M)
    x1, x2 = p[SVector{m1}(1:m1)], p[SVector{m2}(m1+1:m)]
end

function emb_to_chart(xe, C::Chart{I,M}) where {M<:PM,I}
    C1, C2 = product_chart_split(C)
    x1e, x2e = product_manifold_emb_splitview(xe, M)
    x1, x2 = emb_to_chart(x1e, C1), emb_to_chart(x2e, C2)
    x = [x1; x2]
end

function chart_to_emb(x, C::Chart{I,M}) where {M<:PM,I}
    C1, C2 = product_chart_split(C)
    x1, x2 = product_manifold_chart_splitview(x, M)
    x1e, x2e = chart_to_emb(x1, C1), chart_to_emb(x2, C2)
    xe = [x1e; x2e]
end

function chart_transition(xa, Ca::Chart{Tuple{Ia,Ja},M},
        Cb::Chart{Tuple{Ib,Jb},M}) where {M<:PM,Ia,Ja,Ib,Jb}
    C1a, C2a = product_chart_split(Ca)
    C1b, C2b = product_chart_split(Cb)
    x1a, x2a = product_manifold_chart_splitview(xa, M)
    x1b, x2b = chart_transition(x1a, C1a, C1b), chart_transition(x2a, C2a, C2b)
    xb = [x1b; x2b]
end

function choose_chart_emb(xe, C::Chart{Tuple{I,J},M}) where {M<:PM,I,J}
    C1, C2 = product_chart_split(C)
    x1e, x2e = product_manifold_emb_splitview(xe, M)
    C1c, C2c = choose_chart_emb(x1e, C1), choose_chart_emb(x2e, C2)
    product_chart_merge(C1c, C2c)
end

function choose_chart_chart(x, C::Chart{Tuple{I,J},M}) where {M<:PM,I,J}
    C1, C2 = product_chart_split(C)
    x1, x2 = product_manifold_chart_splitview(x, M)
    C1c, C2c = choose_chart_chart(x1, C1), choose_chart_chart(x2, C2)
    product_chart_merge(C1c, C2c)
end


## Exercise 7-3
## The built-in function Meta.parse takes a string and transforms it into an expression. This expression can be evaluated in Julia with the function Core.eval. For example:

## julia> expr = Meta.parse("1+2*3")
## :(1 + 2 * 3)
## julia> eval(expr)
## 7
## julia> expr = Meta.parse("sqrt(π)")
## :(sqrt(π))
## julia> eval(expr)
## 1.7724538509055159

## Write a function called evalloop that iteratively prompts the user, takes the resulting input and evaluates it using eval, and prints the result. It should continue until the user enters done, and then return the value of the last expression it evaluated.
println("Ans: ")

function evalloop()
    while true
        print("eval>> ")
        expr = readline()
        if lowercase(expr) == "done"
            break
        end
        expr = Meta.parse(expr)
        println(eval(expr))
    end
end

evalloop()

println("End.")


# Original implementations by @nantonel
# https://github.com/maxmouchet/HMMBase.jl/pull/6

# *log! methods use the samples log-likelihood instead of the likelihood.

# In-place forward pass, where α and c are allocated beforehand.
function forward!(α::AbstractMatrix, c::AbstractVector, a::AbstractVector, A::AbstractMatrix, L::AbstractMatrix)
    @argcheck size(α, 1) == size(L, 1) == size(c, 1)
    @argcheck size(α, 2) == size(L, 2) == size(a, 1) == size(A, 1) == size(A, 2)

    T, K = size(L)
    (T == 0) && return

    fill!(α, 0.0)
    fill!(c, 0.0)

    for j in OneTo(K)
        α[1,j] = a[j] * L[1,j]
        c[1] += α[1,j]
    end

    for j in OneTo(K)
        α[1,j] /= c[1]
    end

    @inbounds for t in 2:T
        for j in OneTo(K)
            for i in OneTo(K)
                α[t,j] += α[t-1,i] * A[i,j]
            end
            α[t,j] *= L[t,j]
            c[t] += α[t,j]
        end

        for j in OneTo(K)
            α[t,j] /= c[t]
        end
    end
end

# In-place backward pass, where β and c are allocated beforehand.
function backward!(β::AbstractMatrix, c::AbstractVector, a::AbstractVector, A::AbstractMatrix, L::AbstractMatrix)
    @argcheck size(β, 1) == size(L, 1) == size(c, 1)
    @argcheck size(β, 2) == size(L, 2) == size(a, 1) == size(A, 1) == size(A, 2)

    T, K = size(L)
    (T == 0) && return

    fill!(β, 0.0)
    fill!(c, 0.0)

    for j in OneTo(K)
        β[end,j] = 1.0
    end

    @inbounds for t in T-1:-1:1
        for j in OneTo(K)
            for i in OneTo(K)
                β[t,j] += β[t+1,i] * A[j,i] * L[t+1,i]
            end
            c[t+1] += β[t,j]
        end

        for j in OneTo(K)
            β[t,j] /= c[t+1]
        end
    end

    for j in OneTo(K)
        c[1] += a[j] * L[1,j] * β[1,j]
    end
end

# In-place forward pass, where α and c are allocated beforehand.
function forwardlog!(α::AbstractMatrix, c::AbstractVector, a::AbstractVector, A::AbstractMatrix, LL::AbstractMatrix)
    @argcheck size(α, 1) == size(LL, 1) == size(c, 1)
    @argcheck size(α, 2) == size(LL, 2) == size(a, 1) == size(A, 1) == size(A, 2)

    T, K = size(LL)
    (T == 0) && return

    fill!(α, 0.0)
    fill!(c, 0.0)

    m = vec_maximum(view(LL, 1, :))

    for j in OneTo(K)
        α[1,j] = a[j] * exp(LL[1,j] - m)
        c[1] += α[1,j]
    end

    for j in OneTo(K)
        α[1,j] /= c[1]
    end

    c[1] *= exp(m) + eps()

    @inbounds for t in 2:T
        m = vec_maximum(view(LL, t, :))

        for j in OneTo(K)
            for i in OneTo(K)
                α[t,j] += α[t-1,i] * A[i,j]
            end
            α[t,j] *= exp(LL[t,j] - m)
            c[t] += α[t,j]
        end

        for j in OneTo(K)
            α[t,j] /= c[t]
        end

        c[t] *= exp(m) + eps()
    end
end

# In-place backward pass, where β and c are allocated beforehand.
function backwardlog!(β::AbstractMatrix, c::AbstractVector, a::AbstractVector, A::AbstractMatrix, LL::AbstractMatrix)
    @argcheck size(β, 1) == size(LL, 1) == size(c, 1)
    @argcheck size(β, 2) == size(LL, 2) == size(a, 1) == size(A, 1) == size(A, 2)

    T, K = size(LL)
    L = zeros(K)
    (T == 0) && return

    fill!(β, 0.0)
    fill!(c, 0.0)

    for j in OneTo(K)
        β[end,j] = 1.0
    end

    @inbounds for t in T-1:-1:1
        m = vec_maximum(view(LL, t+1, :))

        for i in OneTo(K)
            L[i] = exp(LL[t+1,i] - m)
        end

        for j in OneTo(K)
            for i in OneTo(K)
                β[t,j] += β[t+1,i] * A[j,i] * L[i]
            end
            c[t+1] += β[t,j]
        end

        for j in OneTo(K)
            β[t,j] /= c[t+1]
        end

        c[t+1] *= exp(m) + eps()
    end

    m = vec_maximum(view(LL, 1,:))

    for j in OneTo(K)
        c[1] += a[j] * exp(LL[1,j] - m) * β[1,j]
    end

    c[1] *= exp(m) + eps();
end

# In-place posterior computation, where γ is allocated beforehand.
function posteriors!(γ::AbstractMatrix, α::AbstractMatrix, β::AbstractMatrix)
    @argcheck size(γ) == size(α) == size(β)
    T, K = size(α)
    for t in OneTo(T)
        c = 0.0
        for i = OneTo(K)
            γ[t,i] = α[t,i] * β[t,i]
            c += γ[t,i]
        end

        for i in OneTo(K)
            γ[t,i] /= c
        end
    end
end


function oem_fit(x::AbstractMatrix{<:Real},
                 y::AbstractVector{<:Real};
                 intercept::Bool = true,
                 standardize::Bool = true,
                 nlambda::Int = 100,
                 maxit::Int = 1000,
                 tol::BLAS.BlasFloat = 1e-5,
                 lambda_min_ratio::BLAS.BlasFloat = 1e-3,
                 penalty_factor::AbstractVector{<:Real} = ones(Float64, size(x)[2]))

    n, p = size(x)

    if intercept
        p += 1
        penalty_factor = vcat(0.0, penalty_factor)
    end

    xtx, xty, colsquare_inv, intscale = compute_xtx_y(x, y, intercept, standardize)

    # get eigenvalues, stepsize, A
    steps, A, d = make_steps_A(xtx)

    # initialize beta as vector of zeros
    beta  = zeros(p)

    # compute maximum lambda value
    lambda_max = compute_maxlam(xty, penalty_factor)

    # compute sequence of values for lambda
    lambda = exp.(range(log(lambda_max),
                        stop   = log(lambda_min_ratio * lambda_max),
                        length = nlambda))

    beta_mat  = zeros(p, nlambda)
    iters_vec = repeat([maxit], outer = nlambda)

    # loop over each lambda value
    for l in 1:nlambda
        # set up penalty amounts
        penalty_val   = repeat([lambda[l]], outer = p)
        # adjust penalty by user-specified factors
        penalty_val .*= penalty_factor

        iters = maxit

        # start oem iterations
        for i in 1:maxit
            beta_prev = beta

            u = xty + A * beta_prev

            # M-step
            beta = softthresh(u, penalty_val) ./ steps

            #for j in 1:p
            #    if beta[j] != beta_prev[j]
            #        resid .+= x[:,j] * (beta[j] - beta_prev[j])
            #    end
            #end

            if all(abs.(beta - beta_prev) .<= tol)
                iters = i
                break
            end
        end # end loop over iterations

        ## un-standardize coefficients
        if standardize
            if intercept
                beta[2:p] .*= colsquare_inv[1,:]
            else
                beta .*= colsquare_inv[1,:]
            end
        end

        beta_mat[:,l] = beta
        iters_vec[l] = iters
    end # end loop over lambda

    if intercept
        beta_mat[1,:] ./= intscale
    end

    beta_mat, iters_vec, lambda, d
end


"""
    cleanSC(E[;maxiter=50,ϕ=0.5,stopn=10,peak_removal=false,trust_region=nothing])

CLEAN-SC algorithm for source identification and quantification optionally setting maximum iterations `maxiter` (default 50), and loop gain `ϕ` (default 0.5). Additionally, a stopping criterion `max_peak[i] > max_peak[i-stopn]` can be set with parameter `stopn`. Subtraction of peak sources in the dirty map is supported by `peak_removal=true/false` to remove peak sources outside `trust_region` with limits given as `[xmin,xmax,ymin,ymax]` e.g. `trust_region=[-0.75;1.5;-1.0;1.0]`.

References:
-	Sijtsma, P. (2007). CLEAN based on spatial source coherence. International Journal of Aeroacoustics, 6(4), 357–374.

With inspiration from https://github.com/acoular/acoular/blob/66cba3cffb3bc72602c869f99347be76798f4ac1/acoular/fbeamform.py#L1496
"""
function cleanSC(E;maxiter=50,ϕ=0.5,stopn=10,peak_removal=false,trust_region=nothing)
    @unpack CSM_s,steeringvec,M,N,Nf,fn = E
    x = zeros(N,Nf)

    if peak_removal
        xi = findall((E.rx.>=trust_region[1]) .& (E.rx.<=trust_region[2]))
        yi = findall((E.ry.>=trust_region[3]) .& (E.ry.<=trust_region[4]))
        trust_indices = LinearIndices((E.Nx,E.Ny))[xi,yi]
    else
        trust_indices = LinearIndices((E.Nx,E.Ny))
    end

    @views @inbounds for j = 1:Nf
        _cleanSC!(x[:,j],steeringvec.arr[:,:,j],CSM_s.arr[:,:,j],maxiter,ϕ,stopn,peak_removal,trust_indices)
    end
    return FreqArray(x,E.fn)
end

function _cleanSC!(x,st,csm,maxiter,ϕ,stopn,peak_removal,trust_indices)
    M,N = size(st)
    g = similar(st[:,1])
    h = similar(g)
    h_old = similar(h)
    H = similar(csm)
    max_val = zeros(maxiter)
    dirtyMap = zeros(N)
    for i in 1:maxiter
        AeroAcoustics.bf_col!(dirtyMap,st,csm)
        # Numerical issue with Amiet correction can cause NaNs
        if isempty(filter(!isnan,dirtyMap))
            @goto stop
        else
            dirtyMap[ isnan.(dirtyMap)] .= -Inf
            max_val[i],max_idx = findmax(dirtyMap)
        end
        g .= st[:,max_idx]
        h .= g
        for _ in 1:20
            h_old .= h
            H .= h*h'
            H .= diagm(diag(H))
            h .= (csm*g/max_val[i] + H*g) ./sqrt.(1 .+ mydot(g',H,g))
        end
        Pmax = zeros(N)
        Pmax[max_idx] = 1

        if peak_removal
            if max_idx ∈ trust_indices
                x .+= ϕ*max_val[i]*Pmax
                csm -= ϕ*max_val[i]*(h*h')
            else
                csm -= max_val[i]*(h*h')
            end
        else
            x .+= ϕ*max_val[i]*Pmax
            csm -= ϕ*max_val[i]*(h*h')
        end

        csm[Diagonal(ones(Bool,M))] .= 0
        if (i > stopn) && (max_val[i] > max_val[i-stopn])
            @goto stop
        end
    end
    @label stop
    return nothing
end


struct JohnsonSU{T1<:Real,T2<:Real,T3<:Real,T4<:Real} <: Distributions.Distribution{Univariate,Continuous}
    gamma::T1 #positive: left tail, negativ: right tail, 0: no tail
    delta::T2 #height and width
    xi::T3 #shift along the x-axis
    lambda::T4 #height and width
end

function JohnsonSU(; gamma::Real = 10, delta::Real = 3.5, xi::Real = 10, lambda::Real = 1)
    gamma, delta, xi, lambda = promote(gamma, delta, xi, lambda)
    JohnsonSU(gamma, delta, xi, lambda)
end

Base.minimum(d::JohnsonSU) = -Inf
Base.maximum(d::JohnsonSU) = +Inf

StatsBase.params(d::JohnsonSU) = (d.gamma, d.delta, d.xi, d.lambda)
@inline Distributions.partype(d::JohnsonSU{T}) where {T<:Real} = T

Distributions.location(d::JohnsonSU) = mean(d)
Distributions.scale(d::JohnsonSU) = var(d)

Statistics.mean(d::JohnsonSU) = d.xi - d.lambda * exp(d.delta^(-2)/2)*sinh(d.gamma/d.delta)
Statistics.median(d::JohnsonSU) = d.xi + d.lambda * sinh(-d.gamma/d.delta)
Statistics.var(d::JohnsonSU) = (d.lambda^2)/2 * (exp(d.delta^(-2)) - 1) * (exp(d.delta^(-2))*cosh(2*d.gamma/d.delta) + 1)

function johnsontrafo(x::T1, gamma::T2, delta::T3, xi::T4, lambda::T5) where {T1<:Real,T2<:Real,T3<:Real,T4<:Real, T5<:Real}
    R = float(promote_type(T1, T2, T3, T4, T5))
    convert(R, (gamma + delta*asinh((x - xi)/lambda)))
end

function johnsontrafo_inv(x::T1, gamma::T2, delta::T3, xi::T4, lambda::T5) where {T1<:Real,T2<:Real,T3<:Real,T4<:Real, T5<:Real}
    R = float(promote_type(T1, T2, T3, T4, T5))
    convert(R, lambda*sinh((x - gamma)/delta) + xi)
end

function deriv_johnsontrafo(x::T1, gamma::T2, delta::T3, xi::T4, lambda::T5) where {T1<:Real,T2<:Real,T3<:Real,T4<:Real, T5<:Real}
    R = float(promote_type(T1, T2, T3, T4, T5))
    convert(R, (delta/lambda) * (1/(sqrt(1 + ((x - xi)/lambda)^2))))
end

function deriv_johnsontrafo_inv(x::T1, gamma::T2, delta::T3, xi::T4, lambda::T5) where {T1<:Real,T2<:Real,T3<:Real,T4<:Real, T5<:Real}
    R = float(promote_type(T1, T2, T3, T4, T5))
    convert(R, lambda*cosh((x - gamma)/delta)/delta)
end

function johnsontrafo_ladj(x::T1, gamma::T2, delta::T3, xi::T4, lambda::T5) where {T1<:Real,T2<:Real,T3<:Real,T4<:Real, T5<:Real}
    R = float(promote_type(T1, T2, T3, T4, T5))
    convert(R, log(abs(deriv_johnsontrafo(x,gamma,delta,xi,lambda))))
end

function johnsontrafo_inv_ladj(x::T1, gamma::T2, delta::T3, xi::T4, lambda::T5) where {T1<:Real,T2<:Real,T3<:Real,T4<:Real, T5<:Real}
    R = float(promote_type(T1, T2, T3, T4, T5))
    convert(R, log(abs(deriv_johnsontrafo_inv(x,gamma,delta,xi,lambda))))
end

@with_kw struct JohnsonTrafo{T<:Union{Real,AbstractVector{<:Real}}} <: Function
    gamma::T = 10.0
    delta::T = 3.5
    xi::T = 10.0
    lambda::T = 1.0
end

Base.:(==)(a::JohnsonTrafo, b::JohnsonTrafo) = a.gamma == b.gamma && a.delta == b.delta && a.xi == b.xi && a.lambda == b.lambda
Base.isequal(a::JohnsonTrafo, b::JohnsonTrafo) = isequal(a.gamma, b.gamma) && isequal(a.delta, b.delta) && isequal(a.xi, b.xi) && isequal(a.lambda, b.lambda)
Base.hash(x::JohnsonTrafo, h::UInt) = hash(x.lambda, hash(x.xi, hash(x.delta, hash(x.gamma, hash(:JohnsonTrafo, hash(:EuclidianNormalizingFlows, h))))))

@functor JohnsonTrafo

(f::JohnsonTrafo)(x::Union{Real,AbstractVecOrMat{<:Real}}) = johnsontrafo.(x, f.gamma, f.delta, f.xi, f.lambda)

function with_logabsdet_jacobian(f::JohnsonTrafo, x::Union{Real,AbstractVecOrMat{<:Real}})
    y = f(x)
    ladjs = johnsontrafo_ladj.(x, f.gamma, f.delta, f.xi, f.lambda)
    (y, sum_ladjs(ladjs))
end

inverse(f::JohnsonTrafo) = JohnsonTrafoInv(f.gamma, f.delta, f.xi, f.lambda)

@with_kw struct JohnsonTrafoInv{T<:Union{Real,AbstractVector{<:Real}}} <: Function
    gamma::T = 10.0
    delta::T = 3.5
    xi::T = 10.0
    lambda::T = 1.0
end

@functor JohnsonTrafoInv

Base.:(==)(a::JohnsonTrafoInv, b::JohnsonTrafoInv) = a.gamma == b.gamma && a.delta == b.delta && a.xi == b.xi && a.lambda == b.lambda
Base.isequal(a::JohnsonTrafoInv, b::JohnsonTrafoInv) = isequal(a.gamma, b.gamma) && isequal(a.delta, b.delta) && isequal(a.xi, b.xi) && isequal(a.lambda, b.lambda)
Base.hash(x::JohnsonTrafoInv, h::UInt) = hash(x.lambda, hash(x.xi, hash(x.delta, hash(x.gamma, hash(:JohnsonTrafoInv, hash(:EuclidianNormalizingFlows, h))))))

(f::JohnsonTrafoInv)(x::Union{Real,AbstractVecOrMat{<:Real}}) = johnsontrafo_inv.(x, f.gamma, f.delta, f.xi, f.lambda)

function with_logabsdet_jacobian(f::JohnsonTrafoInv, x::Union{Real,AbstractVecOrMat{<:Real}})
    y = f(x)
    neg_ladjs = johnsontrafo_ladj.(y, f.gamma, f.delta, f.xi, f.lambda)
    (y, - sum_ladjs(neg_ladjs))
end

inverse(f::JohnsonTrafoInv) = JohnsonTrafo(f.gamma, f.delta, f.xi, f.lambda)

#johnsontrafo(d::JohnsonSU, x::Real) = (d.gamma + d.delta*asinh((x - d.xi)/d.lambda))
#johnsontrafo_inv(d::JohnsonSU, x::Real) = d.lambda*sinh((x - d.gamma)/d.delta) + d.xi
#deriv_johnsontrafo(d::JohnsonSU, x::Real) = (d.delta/d.lambda) * (1/(sqrt(1 + ((x - d.xi)/d.lambda)^2)))
#deriv_johnsontrafo_inv(d::JohnsonSU, x::Real) = d.lambda*cosh((x - d.gamma)/d.delta)/d.delta
#johnsontrafo_ladj(d::JohnsonSU, x::Real) = log(abs(deriv_johnsontrafo(d::JohnsonSU, x::Real)))
#johnsontrafo_inv_ladj(d::JohnsonSU, x::Real) = log(abs(deriv_johnsontrafo_inv(d::JohnsonSU, x::Real)))

#Distributions.pdf(d::JohnsonSU, x::Real) = deriv_johnsontrafo(d,x) * 1/sqrt(2*pi) * exp(-0.5 * johnsontrafo(d,x)^2)
Distributions.pdf(d::JohnsonSU, x::Real) = deriv_johnsontrafo(x, d.gamma, d.delta, d.xi, d.lambda) * pdf(Normal(), johnsontrafo(x, d.gamma, d.delta, d.xi, d.lambda))
Distributions.cdf(d::JohnsonSU, x::Real) = cdf(Normal(), johnsontrafo(x, d.gamma, d.delta, d.xi, d.lambda))
#Distributions.logpdf(d::JohnsonSU, x::Real) = log(deriv_johnsontrafo(d,x) * 1/sqrt(2*pi) * exp(-0.5 * johnsontrafo(d,x)^2))
Distributions.logpdf(d::JohnsonSU, x::Real) = log(deriv_johnsontrafo(x, d.gamma, d.delta, d.xi, d.lambda) * pdf(Normal(), johnsontrafo(x, d.gamma, d.delta, d.xi, d.lambda)))
Distributions.logcdf(d::JohnsonSU, x::Real) = logcdf(Normal(), johnsontrafo(x, d.gamma, d.delta, d.xi, d.lambda))
Distributions.ccdf(d::JohnsonSU, x::Real) = 1 - Distributions.cdf(d, x)
Distributions.logccdf(d::JohnsonSU, x::Real) = log(1 - Distributions.cdf(d, x))

Statistics.quantile(d::JohnsonSU, x::Real) = johnsontrafo_inv(quantile(Normal(),x), d.gamma, d.delta, d.xi, d.lambda)