1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
/*
 * File: point.rs
 * Project: stm
 * Created Date: 05/05/2022
 * Author: Shun Suzuki
 * -----
 * Last Modified: 15/01/2023
 * Modified By: Shun Suzuki (suzuki@hapis.k.u-tokyo.ac.jp)
 * -----
 * Copyright (c) 2022 Shun Suzuki. All rights reserved.
 *
 */

use crate::{
    datagram::{DatagramBody, Empty, Filled, Sendable},
    geometry::{Geometry, Matrix3, Transducer, Vector3},
};

use anyhow::{Ok, Result};
use autd3_driver::*;

use super::STM;

pub struct FocusSTM {
    control_points: Vec<(Vector3, u8)>,
    props: autd3_driver::FocusSTMProps,
}

impl FocusSTM {
    pub fn new() -> Self {
        Self {
            control_points: vec![],
            props: Default::default(),
        }
    }

    pub fn add(&mut self, point: Vector3) -> Result<()> {
        self.control_points.push((point, 0));
        Ok(())
    }

    pub fn add_with_shift(&mut self, point: Vector3, shift: u8) -> Result<()> {
        self.control_points.push((point, shift));
        Ok(())
    }

    pub fn control_points(&self) -> &[(Vector3, u8)] {
        &self.control_points
    }
}

impl<T: Transducer> DatagramBody<T> for FocusSTM {
    type O = autd3_driver::FocusSTM;

    fn operation(&mut self, geometry: &Geometry<T>) -> Result<Self::O> {
        let points = geometry
            .device_map()
            .iter()
            .scan(0, |state, tr_num| {
                let r = Some(*state);
                *state += tr_num;
                r
            })
            .map(|origin_idx| {
                let tr = &geometry[origin_idx];
                let origin = tr.position();
                let trans_inv =
                    Matrix3::from_columns(&[tr.x_direction(), tr.y_direction(), tr.z_direction()])
                        .transpose();
                self.control_points
                    .iter()
                    .map(|(p, shift)| {
                        let lp = trans_inv * (p - origin);
                        STMFocus::new(lp.x, lp.y, lp.z, *shift)
                    })
                    .collect()
            })
            .collect();
        let tr_num_min = geometry.device_map().iter().min().unwrap();
        self.props.sound_speed = geometry.sound_speed;
        Ok(Self::O::new(points, *tr_num_min, self.props))
    }
}

impl<T: Transducer> Sendable<T> for FocusSTM {
    type H = Empty;
    type B = Filled;
    type O = <Self as DatagramBody<T>>::O;

    fn operation(&mut self, geometry: &Geometry<T>) -> Result<Self::O> {
        <Self as DatagramBody<T>>::operation(self, geometry)
    }
}

impl STM for FocusSTM {
    fn size(&self) -> usize {
        self.control_points.len()
    }

    fn set_sampling_freq_div(&mut self, freq_div: u32) {
        self.props.freq_div = freq_div;
    }

    fn sampling_freq_div(&self) -> u32 {
        self.props.freq_div
    }

    fn set_start_idx(&mut self, idx: Option<u16>) {
        self.props.start_idx = idx;
    }

    fn start_idx(&self) -> Option<u16> {
        self.props.start_idx
    }

    fn set_finish_idx(&mut self, idx: Option<u16>) {
        self.props.finish_idx = idx;
    }

    fn finish_idx(&self) -> Option<u16> {
        self.props.finish_idx
    }
}

impl Default for FocusSTM {
    fn default() -> Self {
        Self::new()
    }
}