# File: AUTD3.jl
# Project: src
# Created Date: 11/02/2020
# Author: Shun Suzuki
# -----
# Last Modified: 17/12/2021
# Modified By: Shun Suzuki (suzuki@hapis.k.u-tokyo.ac.jp)
# -----
# Copyright (c) 2020 Hapis Lab. All rights reserved.
# 

module AUTD3

using StaticArrays

export Gain, Modulation, PointSequence, GainSequence, Link, AUTD, GainMode, duty_phase_full, phase_full, phase_half
export open_autd, dispose, add_device, synchronize, clear, set_silent_mode, stop, pause, resume, is_open, get_silent_mode, wavelength, set_wavelength
export num_devices
export device_direction_x, device_direction_y, device_direction_z, trans_position
export enumerate_adapters, firmware_info_list
export last_error
export focal_point_gain, focal_point_gain_with_duty, grouped_gain, bessel_beam_gain, bessel_beam_gain_with_duty, plane_wave_gain, plane_wave_gain_with_duty, custom_gain, holo_gain, transducer_test_gain, null_gain
export holo_gain_sdp, holo_gain_evd, holo_gain_gs, holo_gain_gspat, holo_gain_lm, holo_gain_naive
export static_modulation, sine_modulation, sine_pressure_modulation, sine_legacy_modulation, custom_modulation, saw_modulation, square_modulation
export sequence, gain_sequence, add_point, add_gain, set_freq, get_freq, get_sampling_freq, get_sampling_freq_div
export soem_link, twincat_link
export send
export stm, add_stm_gain, start_stm, stop_stm, finish_stm
export NUM_TRANS_IN_UNIT, NUM_TRANS_X, NUM_TRANS_Y, TRANS_SPACING_MM, DEVICE_WIDTH, DEVICE_HEIGHT

include("NativeMethods.jl")

const NUM_TRANS_IN_UNIT = 249
const NUM_TRANS_X = 18
const NUM_TRANS_Y = 14
const TRANS_SPACING_MM = 10.16
const DEVICE_WIDTH = 192.0
const DEVICE_HEIGHT = 151.4

mutable struct Gain
    _ptr::Ptr{Cvoid}
    _disposed::Bool
    function Gain(gain_ptr::Ptr{Cvoid})
        gain = new(gain_ptr, false)
        finalizer(gain -> dispose(gain), gain)
        gain
    end
end

mutable struct Modulation
    _ptr::Ptr{Cvoid}
    _disposed::Bool
    function Modulation(mod_ptr::Ptr{Cvoid})
        modulation = new(mod_ptr, false)
        finalizer(modulation -> dispose(modulation), modulation)
        modulation
    end
end

mutable struct PointSequence
    _ptr::Ptr{Cvoid}
    _disposed::Bool
    function PointSequence(seq_ptr::Ptr{Cvoid})
        seq = new(seq_ptr, false)
        finalizer(seq -> dispose(seq), seq)
        seq
    end
end

mutable struct GainSequence
    _ptr::Ptr{Cvoid}
    _disposed::Bool
    function GainSequence(seq_ptr::Ptr{Cvoid})
        seq = new(seq_ptr, false)
        finalizer(seq -> dispose(seq), seq)
        seq
    end
end

mutable struct Link
    _link_ptr::Ptr{Cvoid}
    function Link(link_ptr::Ptr{Cvoid})
        new(link_ptr)
    end
end

mutable struct AUTD
    _handle::Ptr{Cvoid}
    _disposed::Bool
    function AUTD()
        chandle = Ref(Ptr{Cvoid}(0))
        autd_create_controller(chandle)
        autd = new(chandle[], false)
        finalizer(autd -> dispose(autd), autd)
        autd
    end
end

mutable struct STMController
    _handle::Ptr{Cvoid}
    function STMController(handle::Ptr{Cvoid})
        new(handle)
    end
end

function open_autd(autd::AUTD, link::Link)
    autd_open_controller(autd._handle, link._link_ptr)
end

function dispose(self::AUTD)
    if !self._disposed
        autd_close_controller(self._handle)
        autd_free_controller(self._handle)
        self._disposed = true
    end
end

function dispose(self::Gain)
    if !self._disposed
        autd_delete_gain(self._ptr)
        self._disposed = true
    end
end

function dispose(self::Modulation)
    if !self._disposed
        autd_delete_modulation(self._ptr)
        self._disposed = true
    end
end

function dispose(self::PointSequence)
    if !self._disposed
        autd_delete_sequence(self._ptr)
        self._disposed = true
    end
end

function dispose(self::GainSequence)
    if !self._disposed
        autd_delete_sequence(self._ptr)
        self._disposed = true
    end
end

function enumerate_adapters()
    res = []
    phandle = Ref(Ptr{Cvoid}(0))
    size = autd_get_adapter_pointer(phandle)
    handle::Ptr{Cvoid} = phandle[]

    for i = 0:size-1
        sb_desc = zeros(UInt8, 128)
        sb_name = zeros(UInt8, 128)
        autd_get_adapter(handle, i, sb_desc, sb_name)
        push!(res, [String(strip(String(sb_desc), '\0')), String(strip(String(sb_name), '\0'))])
    end

    autd_free_adapter_pointer(handle)
    res
end

function firmware_info_list(autd::AUTD)
    res = []
    phandle = Ref(Ptr{Cvoid}(0))
    size = autd_get_firmware_info_list_pointer(autd._handle, phandle)
    handle::Ptr{Cvoid} = phandle[]

    for i = 0:size-1
        sb_cpu = zeros(UInt8, 128)
        sb_fpga = zeros(UInt8, 128)
        autd_get_firmware_info(handle, i, sb_cpu, sb_fpga)
        push!(res, [String(strip(String(sb_cpu), '\0')), String(strip(String(sb_fpga), '\0'))])
    end

    autd_free_firmware_info_list_pointer(handle)
    res
end

function add_device(autd::AUTD, pos::SVector{3,Float64}, rot::SVector{3,Float64})
    x, y, z = pos
    az1, ay, az2 = rot
    autd_add_device(autd._handle, x, y, z, az1, ay, az2)
end

function add_device(autd::AUTD, pos::SVector{3,Float64}, qua::SVector{4,Float64})
    x, y, z = pos
    qw, qx, qy, qz = qua
    autd_add_device_quaternion(autd._handle, x, y, z, qw, qx, qy, qz)
end

function clear(autd::AUTD)
    autd_clear(autd._handle)
end

function update_ctrl_flags(autd::AUTD)
    autd_update_ctrl_flags(autd._handle)
end

function stop(autd::AUTD)
    autd_stop(autd._handle)
end

function pause(autd::AUTD)
    autd_pause(autd._handle)
end

function resume(autd::AUTD)
    autd_resume(autd._handle)
end

function is_open(autd::AUTD)
    autd_is_open(autd._handle)
end

function last_error()
    size = autd_get_last_error(Ptr{Cvoid}(0))
    err = zeros(UInt8, size)
    autd_get_last_error(err)
    String(strip(String(err), '\0'))
end

function set_silent_mode(autd::AUTD, silent::Bool)
    autd_set_silent_mode(autd._handle, silent)
end

function get_silent_mode(autd::AUTD)
    autd_get_silent_mode(autd._handle)
end

function set_force_fan(autd::AUTD, force_fan::Bool)
    autd_set_force_fan(autd._handle, force_fan)
end

function get_force_fan(autd::AUTD)
    autd_get_force_fan(autd._handle)
end

function wavelength(autd::AUTD)
    autd_get_wavelength(autd._handle)
end

function set_wavelength(autd::AUTD, wavelength::Float64)
    autd_set_wavelength(autd._handle, wavelength)
end

function attenuation(autd::AUTD)
    autd_get_attenuation(autd._handle)
end

function set_attenuation(autd::AUTD, attenuation::Float64)
    autd_set_attenuation(autd._handle, attenuation)
end

function num_devices(autd::AUTD)
    autd_num_devices(autd._handle)
end

function adjust_amp(amp::Float64)
    d = asin(amp) / pi
    UInt8(510.0 * d)
end

function focal_point_gain_with_duty(position::SVector{3,Float64}; duty::UInt8)
    x, y, z = position
    chandle = Ref(Ptr{Cvoid}(0))
    autd_gain_focal_point(chandle, x, y, z, duty)
    Gain(chandle[])
end

function focal_point_gain(position::SVector{3,Float64}; amp::Float64 = 1.0)
    focal_point_gain_with_duty(position; duty = adjust_amp(amp))
end

function grouped_gain(autd, group_ids::Array{Tuple{Int32,Gain},1})
    chandle = Ref(Ptr{Cvoid}(0))
    autd_gain_grouped(chandle, autd._ptr)
    for (id, gain) in group_ids
        autd_gain_grouped_add(chandle[], id, gain)
    end
    Gain(chandle[])
end

function bessel_beam_gain_with_duty(position::SVector{3,Float64}, direction::SVector{3,Float64}, theta_z::Float64; duty::UInt8)
    x, y, z = position
    nx, ny, nz = direction
    chandle = Ref(Ptr{Cvoid}(0))
    autd_gain_bessel_beam(chandle, x, y, z, nx, ny, nz, theta_z, duty)
    Gain(chandle[])
end

function bessel_beam_gain(position::SVector{3,Float64}, direction::SVector{3,Float64}, theta_z::Float64; amp::Float64 = 1.0)
    bessel_beam_gain_with_duty(position, direction, theta_z; duty = adjust_amp(amp))
end

function plane_wave_gain_with_duty(direction::SVector{3,Float64}; duty::UInt8)
    nx, ny, nz = direction
    chandle = Ref(Ptr{Cvoid}(0))
    autd_gain_plane_wave(chandle, nx, ny, nz, duty)
    Gain(chandle[])
end

function plane_wave_gain(direction::SVector{3,Float64}; amp::Float64 = 1.0)
    plane_wave_gain_with_duty(direction; duty = adjust_amp(amp))
end

function custom_gain(data::Array{UInt16,1})
    len = length(data)
    chandle = Ref(Ptr{Cvoid}(0))
    autd_gain_custom(chandle, data, len)
    Gain(chandle[])
end

function _pack_foci(foci::Array{SVector{3,Float64},1})
    len = length(foci)
    foci_array = zeros(Float64, len * 3)
    for (i, focus) in enumerate(foci)
        foci_array[3*(i-1)+1] = focus[1]
        foci_array[3*(i-1)+2] = focus[2]
        foci_array[3*(i-1)+3] = focus[3]
    end
    foci_array
end

function holo_gain_sdp(foci::Array{SVector{3,Float64},1}, amps::Array{Float64,1}; alpha::Float64 = 1e-3, lambda::Float64 = 0.9, repeat::UInt64 = UInt64(100), normalize::Bool = true)
    len = length(foci)

    foci_array = _pack_foci(foci)
    backend = Ref(Ptr{Cvoid}(0))
    autd_eigen_backend(backend)

    chandle = Ref(Ptr{Cvoid}(0))
    autd_gain_holo_sdp(chandle, backend[], foci_array, amps, len, alpha, lambda, repeat, normalize)
    autd_delete_backend(backend[])
    Gain(chandle[])
end

function holo_gain_evd(foci::Array{SVector{3,Float64},1}, amps::Array{Float64,1}; gamma::Float64 = 1.0, normalize::Bool = true)
    len = length(foci)

    foci_array = _pack_foci(foci)
    backend = Ref(Ptr{Cvoid}(0))
    autd_eigen_backend(backend)

    chandle = Ref(Ptr{Cvoid}(0))
    autd_gain_holo_evd(chandle, backend[], foci_array, amps, len, gamma, normalize)
    autd_delete_backend(backend[])
    Gain(chandle[])
end

function holo_gain_naive(foci::Array{SVector{3,Float64},1}, amps::Array{Float64,1})
    len = length(foci)

    foci_array = _pack_foci(foci)
    backend = Ref(Ptr{Cvoid}(0))
    autd_eigen_backend(backend)

    chandle = Ref(Ptr{Cvoid}(0))
    autd_gain_holo_naive(chandle, backend[], foci_array, amps, len)
    autd_delete_backend(backend[])
    Gain(chandle[])
end

function holo_gain_gs(foci::Array{SVector{3,Float64},1}, amps::Array{Float64,1}; repeat::UInt64 = UInt64(100))
    len = length(foci)

    foci_array = _pack_foci(foci)
    backend = Ref(Ptr{Cvoid}(0))
    autd_eigen_backend(backend)

    chandle = Ref(Ptr{Cvoid}(0))
    autd_gain_holo_gs(chandle, backend[], foci_array, amps, len, repeat)
    autd_delete_backend(backend[])
    Gain(chandle[])
end

function holo_gain_gspat(foci::Array{SVector{3,Float64},1}, amps::Array{Float64,1}; repeat::UInt64 = UInt64(100))
    len = length(foci)

    foci_array = _pack_foci(foci)
    backend = Ref(Ptr{Cvoid}(0))
    autd_eigen_backend(backend)

    chandle = Ref(Ptr{Cvoid}(0))
    autd_gain_holo_gspat(chandle, backend[], foci_array, amps, len, repeat)
    autd_delete_backend(backend[])
    Gain(chandle[])
end

function holo_gain_lm(foci::Array{SVector{3,Float64},1}, amps::Array{Float64,1}; eps_1::Float64 = 1e-8, eps_2::Float64 = 1e-8, tau::Float64 = 1e-3, k_max::UInt64 = UInt64(5), initial::Array{Float64,1} = [])
    len = length(foci)

    foci_array = _pack_foci(foci)
    backend = Ref(Ptr{Cvoid}(0))
    autd_eigen_backend(backend)

    chandle = Ref(Ptr{Cvoid}(0))
    initial_len = length(initial)
    init = initial_len == 0 ? Ptr{Cvoid}(0) : initial
    autd_gain_holo_lm(chandle, backend[], foci_array, amps, len, eps_1, eps_2, tau, k_max, init, initial_len)
    autd_delete_backend(backend[])
    Gain(chandle[])
end

function holo_gain_greedy(foci::Array{SVector{3,Float64},1}, amps::Array{Float64,1}, phase_div::Int32)
    len = length(foci)

    foci_array = _pack_foci(foci)
    backend = Ref(Ptr{Cvoid}(0))
    autd_eigen_backend(backend)

    chandle = Ref(Ptr{Cvoid}(0))
    autd_gain_holo_greedy(chandle, backend, foci_array, amps, len, phase_div)
    autd_delete_backend(backend[])
    Gain(chandle[])
end

function transducer_test_gain(trans_idx::Int32, duty::UInt8, phase::UInt8)
    chandle = Ref(Ptr{Cvoid}(0))
    autd_gain_transducer_test(chandle, trans_idx, duty, phase)
    Gain(chandle[])
end

function null_gain()
    chandle = Ref(Ptr{Cvoid}(0))
    autd_gain_null(chandle)
    Gain(chandle[])
end

function static_modulation(amp::UInt8 = 0xFF)
    chandle = Ref(Ptr{Cvoid}(0))
    autd_modulation_static(chandle, amp)
    Modulation(chandle[])
end

function custom_modulation(data::Array{UInt8,1})
    len = length(data)
    chandle = Ref(Ptr{Cvoid}(0))
    autd_modulation_custom(chandle, data, len)
    Modulation(chandle[])
end

function sine_modulation(freq::Int32, amp::Float64 = 1.0, offset::Float64 = 0.5)
    chandle = Ref(Ptr{Cvoid}(0))
    autd_modulation_sine(chandle, freq, amp, offset)
    Modulation(chandle[])
end

function sine_squared_modulation(freq::Int32, amp::Float64 = 1.0, offset::Float64 = 0.5)
    chandle = Ref(Ptr{Cvoid}(0))
    autd_modulation_sine_squared(chandle, freq, amp, offset)
    Modulation(chandle[])
end

function sine_legacy_modulation(freq::Float64, amp::Float64 = 1.0, offset::Float64 = 0.5)
    chandle = Ref(Ptr{Cvoid}(0))
    autd_modulation_sine_legacy(chandle, freq, amp, offset)
    Modulation(chandle[])
end

function square_modulation(freq::Int32, low::UInt8 = 0x00, high::UInt8 = 0xFF, duty::Float64 = 0.5)
    chandle = Ref(Ptr{Cvoid}(0))
    autd_modulation_square(chandle, freq, low, high, duty)
    Modulation(chandle[])
end

function sequence()
    chandle = Ref(Ptr{Cvoid}(0))
    autd_sequence(chandle)
    PointSequence(chandle[])
end

function gain_sequence(autd, gain_mode::GainMode)
    chandle = Ref(Ptr{Cvoid}(0))
    autd_gain_sequence(chandle, autd._handle, UInt16(gain_mode))
    GainSequence(chandle[])
end

function add_point(seq::PointSequence, point::SVector{3,Float64}, duty::UInt8 = 0xFF)
    x, y, z = point
    autd_sequence_add_point(seq._ptr, x, y, z, duty)
end

function add_gain(seq::GainSequence, gain::Gain)
    autd_sequence_add_gain(seq._ptr, gain._ptr)
end

function set_freq(seq, freq::Float64)
    autd_sequence_set_freq(seq._ptr, freq)
end

function get_freq(seq)
    autd_sequence_freq(seq._ptr)
end

function get_sampling_freq(seq)
    autd_sequence_sampling_freq(seq._ptr)
end

function get_sampling_freq_div(seq)
    autd_sequence_sampling_freq_div(seq._ptr)
end

function soem_link(ifname::String, device_num::Int32, cycle_ticks::UInt32 = UInt32(1))
    chandle = Ref(Ptr{Cvoid}(0))
    autd_link_soem(chandle, ifname, device_num, cycle_ticks)
    Link(chandle[])
end

function twincat_link()
    chandle = Ref(Ptr{Cvoid}(0))
    autd_link_twincat(chandle)
    Link(chandle[])
end

function emulator_link(port::UInt16, autd::AUTD)
    chandle = Ref(Ptr{Cvoid}(0))
    autd_link_emulator(chandle, port, autd._handle)
    Link(chandle[])
end

function send(autd::AUTD, header, body)
    if header !== Nothing && body !== Nothing
        autd_send_header_body(autd._handle, header._ptr, body._ptr)
    elseif header !== Nothing
        autd_send_header(autd._handle, header._ptr)
    elseif body !== Nothing
        autd_send_body(autd._handle, body._ptr)
    end
end

function stm(autd::AUTD)
    chandle = Ref(Ptr{Cvoid}(0))
    autd_stm_controller(chandle, autd._handle)
    STMController(chandle[])
end

function add_stm_gain(stm::STMController, gain::Gain)
    autd_add_stm_gain(stm._handle, gain._ptr)
end

function start_stm(stm::STMController, freq::Float64)
    autd_start_stm(stm._handle, freq)
end

function stop_stm(stm::STMController)
    autd_stop_stm(stm._handle)
end

function finish_stm(stm::STMController)
    autd_finish_stm(stm._handle)
end

function trans_position(autd::AUTD, device_idx::Int32, local_trans_idx::Int32)
    x = Ref{Float64}(0)
    y = Ref{Float64}(0)
    z = Ref{Float64}(0)
    autd_trans_position(autd._handle, device_idx, local_trans_idx, x, y, z)
    SVector(x[], y[], z[])
end

function device_direction_x(autd::AUTD, device_idx::Int32)
    x = Ref{Float64}(0)
    y = Ref{Float64}(0)
    z = Ref{Float64}(0)
    autd_device_x_direction(autd._handle, device_idx, x, y, z)
    SVector(x[], y[], z[])
end

function device_direction_y(autd::AUTD, device_idx::Int32)
    x = Ref{Float64}(0)
    y = Ref{Float64}(0)
    z = Ref{Float64}(0)
    autd_device_y_direction(autd._handle, device_idx, x, y, z)
    SVector(x[], y[], z[])
end

function device_direction_z(autd::AUTD, device_idx::Int32)
    x = Ref{Float64}(0)
    y = Ref{Float64}(0)
    z = Ref{Float64}(0)
    autd_device_z_direction(autd._handle, device_idx, x, y, z)
    SVector(x[], y[], z[])
end

end
