tradingbot/src/value_estimation_team/indicators/heatmap_volume.rs

use crate::value_estimation_team::datapoints::price_data::RealtimePriceData;

#[derive(Debug, Clone, PartialEq)]
pub enum HeatMapLevel {
    ExtraHigh,
    High,
    Medium,
    Normal,
    Low,
}

#[derive(Debug, Clone)]
pub struct HeatmapVolumeData {
    pub heatmap_value: f64,
    pub heatmap_level: HeatMapLevel,
    pub close_time: i64,
}

// Implementation from TradingView Script (HeatMap Volume by xdecow)
pub async fn heatmap_volume(
    symbol: &String,
    input_rt_data: &Vec<(String, Vec<RealtimePriceData>)>,
    ma_len: usize,
    std_len: usize,
    extra_high_thold: f64,
    high_thold: f64,
    medium_thold: f64,
    normal_thold: f64,
) -> Option<Vec<HeatmapVolumeData>> {
    let symbol_search_result = input_rt_data.iter().position(|x| x.0 == *symbol);

    match symbol_search_result {
        Some(T) => {
            if input_rt_data[T].1.len() >= ma_len && input_rt_data[T].1.len() >= std_len {
                // calc mean
                #[derive(Debug, Clone)]
                struct MeanData {
                    mean_value: f64,
                    close_time: i64,
                }

                let mut mean_vec: Vec<MeanData> = Vec::new();
                let mut mean_data = MeanData {
                    mean_value: 0.0,
                    close_time: 0,
                };
                let window = input_rt_data[T].1.windows(ma_len);
                for buffer in window {
                    // calculate SMA of volume
                    mean_data.mean_value = 0.0;
                    mean_data.close_time = 0;
                    for element in buffer {
                        mean_data.mean_value += element.quote_asset_volume;
                    }
                    mean_data.mean_value /= ma_len as f64;
                    mean_data.close_time = buffer.last().unwrap().close_time;

                    mean_vec.push(mean_data.clone());
                }

                // calc pstdev
                #[derive(Debug, Clone)]
                struct PstdevData {
                    pstdev_value: f64, // population standard deviation value
                    close_time: i64,
                }
                let mut pstdev_vec: Vec<PstdevData> = Vec::new();
                let mut pstdev_data = PstdevData {
                    pstdev_value: 0.0,
                    close_time: 0,
                };
                let mut mean = 0.0;
                let mut summation = 0.0;
                let mut sample_minus_mean = 0.0;
                let window = input_rt_data[T].1.windows(std_len);
                for buffer in window {
                    pstdev_data.pstdev_value = 0.0;
                    pstdev_data.close_time = 0;
                    mean = 0.0;
                    summation = 0.0;
                    sample_minus_mean = 0.0;
                    for element in buffer {
                        mean += element.quote_asset_volume;
                    }
                    mean /= std_len as f64;

                    for element in buffer {
                        sample_minus_mean = element.quote_asset_volume - mean;
                        summation += sample_minus_mean.powi(2);
                    }

                    pstdev_data.pstdev_value = (summation / std_len as f64).sqrt();
                    pstdev_data.close_time = buffer.last().unwrap().close_time;

                    pstdev_vec.push(pstdev_data.clone());
                }

                // calc stdbar and heatmap volume
                let mut heatmap_vol_vec: Vec<HeatmapVolumeData> = Vec::new();
                let mut heatmap_vol_data = HeatmapVolumeData {
                    heatmap_value: 0.0,
                    heatmap_level: HeatMapLevel::Normal,
                    close_time: 0,
                };

                if ma_len == std_len {
                    let mut rt_data_trunc_vec =
                        input_rt_data[T].1.get(ma_len - 1..).unwrap().iter();

                    let zipped = mean_vec.iter().zip(pstdev_vec.iter());
                    for element in zipped {
                        heatmap_vol_data.heatmap_value =
                            (rt_data_trunc_vec.next().unwrap().quote_asset_volume
                                - element.0.mean_value)
                                / element.1.pstdev_value;
                        heatmap_vol_data.close_time = element.0.close_time;
                        if heatmap_vol_data.heatmap_value >= extra_high_thold {
                            heatmap_vol_data.heatmap_level = HeatMapLevel::ExtraHigh;
                        } else if heatmap_vol_data.heatmap_value > high_thold {
                            heatmap_vol_data.heatmap_level = HeatMapLevel::High;
                        } else if heatmap_vol_data.heatmap_value > medium_thold {
                            heatmap_vol_data.heatmap_level = HeatMapLevel::Medium;
                        } else if heatmap_vol_data.heatmap_value > normal_thold {
                            heatmap_vol_data.heatmap_level = HeatMapLevel::Normal;
                        } else {
                            heatmap_vol_data.heatmap_level = HeatMapLevel::Low;
                        }
                        heatmap_vol_vec.push(heatmap_vol_data.clone());
                    }
                } else if ma_len > std_len {
                    let mut rt_data_trunc_vec =
                        input_rt_data[T].1.get(std_len - 1..).unwrap().iter();
                    let mut mean_trunc_vec =
                        mean_vec.get(mean_vec.len() - std_len..).unwrap().iter();
                    let zipped = mean_trunc_vec.zip(pstdev_vec.iter());
                    for element in zipped {
                        heatmap_vol_data.heatmap_value =
                            (rt_data_trunc_vec.next().unwrap().quote_asset_volume
                                - element.0.mean_value)
                                / element.1.pstdev_value;
                        heatmap_vol_data.close_time = element.0.close_time;
                        heatmap_vol_vec.push(heatmap_vol_data.clone());
                    }
                } else {
                    let mut rt_data_trunc_vec =
                        input_rt_data[T].1.get(ma_len - 1..).unwrap().iter();
                    let mut pstdev_trunc_vec =
                        pstdev_vec.get(pstdev_vec.len() - ma_len..).unwrap().iter();
                    let zipped = mean_vec.iter().zip(pstdev_trunc_vec);
                    for element in zipped {
                        heatmap_vol_data.heatmap_value =
                            (rt_data_trunc_vec.next().unwrap().quote_asset_volume
                                - element.0.mean_value)
                                / element.1.pstdev_value;
                        heatmap_vol_data.close_time = element.0.close_time;
                        heatmap_vol_vec.push(heatmap_vol_data.clone());
                    } // level 구현
                }
                Some(heatmap_vol_vec)
            } else {
                None
            }
        }
        None => None,
    }
}