radarmp/loaders/etws_loader/src/lib.rs

mod error;
mod raw;
use raw::parse_raw_data;
mod parser;
use abi_stable::{
    export_root_module,
    prefix_type::PrefixTypeTrait,
    rvec, sabi_extern_fn,
    sabi_trait::prelude::TD_Opaque,
    std_types::{
        RNone, ROk,
        ROption::RSome,
        RResult::{self, RErr},
        RStr, RString, RVec,
    },
};
use parser::{Record, ValueResult};

use radarg_plugin_interface::{
    DataLoaderPlugin, DataLoaderPlugin_TO, Error, GridDataInfo, LoadedData, PluginId, PluginMod,
    PluginMod_Ref, PluginType, ProbeDataType, RadarGridData, VecResult,
};

macro_rules! data_rvec {
    ($data:ident, $({$parsed_vec:tt => $abi_stable_rvec:tt}),+) => {
        match $data {
            $(
                ValueResult::$parsed_vec(data) => VecResult::$abi_stable_rvec(RVec::from(data)),
            )+
        }
    };
}

macro_rules! data_type {
    ($datetype:ident, $({ $($name:literal)|+ =>  $type:tt }),+) => {

        match $datetype {
            $(
                $($name)|+ => ProbeDataType::$type,
            )+
            _ => ProbeDataType::Unknown
        }

    };
}

#[export_root_module]
fn instantiate_root_module() -> PluginMod_Ref {
    PluginMod { new }.leak_into_prefix()
}

struct ETWSLoader {
    id: PluginId,
}

impl DataLoaderPlugin for ETWSLoader {
    fn plugin_id(&self) -> &PluginId {
        &self.id
    }

    fn load(&self, path: RStr<'_>) -> RResult<RVec<LoadedData>, Error> where {
        let path = path.as_str();
        let c = Record::parse_from_path(path.to_string())
            .map(|record| {
                let result = record
                    .blocks
                    .into_iter()
                    .map(|b| {
                        let data = b.data;
                        let converted_data = data_rvec!(
                            data,
                            {I16 => I16},
                            {F32 => F32},
                            {F64 => F64},
                            {I32 => I32},
                            {U32 => U32},
                            {I64 => I64},
                            {U64 => U64},
                            {I8 => I8},
                            {U8 => U8}
                        );

                        let dimension_values = convert_nested_vec_to_rvec(b.info.dimension_values);
                        let dimension_size = convert_vec_to_rvec(
                            b.info
                                .dimension_size
                                .into_iter()
                                .map(|p| p as usize)
                                .collect::<Vec<_>>(),
                        );
                        let dimension_names = convert_vec_to_rvec(
                            b.info
                                .dimension_des
                                .into_iter()
                                .map(RString::from)
                                .collect::<Vec<_>>(),
                        );

                        let fill_value = b.info.fill_value;
                        let value_name = RString::from(b.info.value_name);
                        let value_description = RString::from(b.info.value_des);
                        let datetime = record.filetime.timestamp();

                        let value_key = value_name.as_str();

                        let data_type = data_type!(
                            value_key,
                            { "ET"  => ET },
                            { "DBZ" | "CR" | "FR" | "R" | "CR0" | "CR1" | "CR2"  => DBZ},
                            { "VIL" => VIL},
                            { "EB" => EB},
                            { "V" => V},
                            { "ZDR" => ZDR},
                            { "PHIDP" => PHIDP},
                            { "KDP" => KDP},
                            { "CC" =>CC},
                            { "HCA" => HCA},
                            { "QPE" => QPE},
                            { "QPF" => QPF}
                        );

                        let grid_info = GridDataInfo {
                            shape: dimension_size,
                            dimensions: dimension_values,
                            dimension_names,
                            fill_value,
                            datetime: RSome(datetime),
                            value_description: RSome(value_description),
                            maybe_probe_data_type: RSome(data_type),
                            value_name,
                            dimension_description: "Wrong".into(),
                        };

                        LoadedData::RadarGridData(RadarGridData {
                            data: converted_data,
                            info: grid_info,
                        })
                    })
                    .collect::<RVec<_>>();
                result
            })
            .or_else(|_| {
                let result = parse_raw_data(path)
                    .map(|result| {
                        let data = result.datas;

                        let shape = rvec![
                            result.els.len() as usize,
                            result.azs.len() as usize,
                            result.ranges.len() as usize
                        ];

                        let dimensions = rvec![
                            RVec::from(result.els),
                            RVec::from(result.azs),
                            RVec::from(result.ranges)
                        ];

                        let dimension_names = rvec!["el".into(), "az".into(), "range".into()];
                        let datetime = result.datetime;

                        let loaded = data
                            .into_iter()
                            .map(move |(k, v)| {
                                let data_type = data_type!(
                                    k,
                                    { "ET"  => ET },
                                    { "DBZ" | "CR" | "FR" | "R" | "CR0" | "CR1" | "CR2" => DBZ},
                                    { "VIL" => VIL},
                                    { "EB" => EB},
                                    { "V" | "VEL" => V},
                                    { "ZDR" => ZDR},
                                    { "PHIDP" => PHIDP},
                                    { "KDP" => KDP},
                                    { "CC" =>CC},
                                    { "HCA" => HCA},
                                    { "QPE" => QPE},
                                    { "QPF" => QPF}
                                );
                                let radar_data = RadarGridData {
                                    data: VecResult::F64(RVec::from(v)),
                                    info: GridDataInfo {
                                        shape: shape.clone(),
                                        dimensions: dimensions.clone(),
                                        dimension_names: dimension_names.clone(),
                                        fill_value: 0.0,
                                        datetime: RSome(datetime),
                                        value_description: RNone,
                                        maybe_probe_data_type: RSome(data_type),
                                        value_name: RString::from(k),
                                        dimension_description: "Wrong".into(),
                                    },
                                };

                                LoadedData::RadarGridData(radar_data)
                            })
                            .collect::<RVec<_>>();

                        loaded
                    })
                    .map_err(|_| Error::UnsupportedFormat)?;
                Ok(result)
            });

        RResult::from(c)
    }

    fn supported_extensions(&self) -> RVec<RStr<'static>> where {
        convert_iter_to_rvec(vec!["dat".into(), "gz".into(), "zip".into()].into_iter())
    }

    fn plugin_info(&self) -> radarg_plugin_interface::PluginInfo where {
        radarg_plugin_interface::PluginInfo {
            plugin_type: "DataLoader".into(),
            name: "ETWS_Loader".into(),
            version: "0.1.0".into(),
            author: "Tsuki".into(),
            description: RSome("ETWS Loader".into()),
            url: RSome("https://keitsuki.club".into()),
        }
    }
}

fn convert_iter_to_rvec<T>(raw: impl Iterator<Item = T>) -> RVec<T> {
    RVec::from(raw.collect::<Vec<_>>())
}

fn convert_vec_to_rvec<T>(raw: Vec<T>) -> RVec<T> {
    RVec::from(raw)
}

fn convert_nested_vec_to_rvec<T>(raw: Vec<Vec<T>>) -> RVec<RVec<T>> {
    RVec::from(raw.into_iter().map(RVec::from).collect::<Vec<_>>())
}

#[sabi_extern_fn]
pub fn new(plugin_id: PluginId) -> RResult<PluginType, Error> {
    let this = ETWSLoader { id: plugin_id };
    ROk(DataLoaderPlugin_TO::from_value(this, TD_Opaque))
}

// 测试模块
#[cfg(test)]
mod tests {
    // 导入外部作用域中的所有项
    use super::*;

    // 一个简单的测试
    #[test]

    fn test() {
        let result =
            Record::parse_from_path("/Volumes/data2/RadarArray/ShaoXing/radarData/OutputProducts/RadarProducts/BasicProductsX/20230627/20230627163400/ZJSXAA_20230627163400_VIL.dat.gz")
                .unwrap();
    }
}