onnx_graph 0.1.2

use std::{any::Any, collections::HashMap};

use crate::{
    fill_from_elem,
    nodes::{node::Node, onnx_operation_trait::FromOnnxOperation, unique_ids::UniqueId},
    tensor_map::TensorMap,
    typed_array::TypedArray,
};
use anyhow::Result;
use ndarray::{ArrayD, IxDyn};
use onnx_extractor::{AttributeValue, OnnxOperation};

#[derive(Default)]
pub struct ConstantOfShapeNode<T: Default> {
    shape_array: String,

    value: Option<TypedArray>,

    o: String,

    unique_id: UniqueId,

    next_node: Option<Vec<Box<dyn Node<T>>>>,
}

impl<T: Default> FromOnnxOperation for ConstantOfShapeNode<T> {
    fn from_onnx_operation(elem: &OnnxOperation) -> Result<Self> {
        let mut constant_of_shape = Self {
            shape_array: String::new(),
            value: None,
            o: String::new(),
            unique_id: UniqueId::ConstantOfShape,
            next_node: None,
        };

        let value = elem
            .attributes
            .get("value")
            .and_then(|val| AttributeValue::as_tensor(val))
            .map(|tensor| TypedArray::from_tensor(&tensor));
        constant_of_shape.value = value;

        if elem.input_count() != 0 {
            constant_of_shape.add_input_strings(elem.inputs[0].clone());
        } else {
            constant_of_shape.add_input_strings(elem.name.clone());
        }
        constant_of_shape.add_output_strings(elem.outputs[0].clone());
        Ok(constant_of_shape)
    }
}

impl<T: Default> ConstantOfShapeNode<T> {
    pub fn new(elem: &OnnxOperation) -> Self {
        let mut constant_of_shape = Self {
            shape_array: String::new(),
            value: Some(TypedArray::Float(ArrayD::zeros(IxDyn(&[1])))),
            o: String::new(),
            unique_id: UniqueId::ConstantOfShape,
            next_node: None,
        };
        constant_of_shape.add_input_strings(elem.inputs[0].clone());
        constant_of_shape.add_output_strings(elem.outputs[0].clone());
        constant_of_shape
    }

    pub fn add_input_strings(&mut self, x: String) {
        self.shape_array = x;
    }

    pub fn add_output_strings(&mut self, o: String) {
        self.o = o;
    }
}

impl<T: Default + 'static> Node<T> for ConstantOfShapeNode<T> {
    fn as_any_mut(&mut self) -> &mut dyn Any {
        self
    }

    fn get_unique_id(&self) -> UniqueId {
        self.unique_id
    }
    fn get_unique_id_mut(&mut self) -> UniqueId {
        self.unique_id
    }

    fn get_next(&self) -> Option<&Vec<Box<dyn Node<T>>>> {
        self.next_node.as_ref()
    }

    fn execute(&self, omap: &mut TensorMap) {
        let [shape_array, o] = omap.get_disjoint_mut([&self.shape_array, &self.o]);
        let shape_array = shape_array.map(|inner| &*inner);

        match (shape_array, &self.value, o) {
            (Some(x), Some(value), Some(result)) => {
                if let Err(e) = x.constant_of_shape(value, result) {
                    println!("!!!{e}");
                }
            }
            (None, Some(value), Some(result)) => {
                *result = value.clone();
            }
            _ => panic!("ConstantOfShapeNode: missing input {}", self.shape_array),
        }
    }

    fn output_names(&self) -> Vec<String> {
        vec![self.o.clone()]
    }

    fn input_names(&self) -> Vec<String> {
        vec![self.shape_array.clone()]
    }

    fn take_next(&mut self) -> Option<Vec<Box<dyn Node<T>>>> {
        self.next_node.take()
    }
    fn get_next_mut(&mut self) -> Option<&mut Vec<Box<dyn Node<T>>>> {
        self.next_node.as_mut()
    }

    fn set_next(&mut self, next: Option<Vec<Box<dyn Node<T>>>>) {
        self.next_node = next;
    }

    fn print(&self) {
        if let Some(list) = &self.next_node {
            print!("{}-", list.len());
        }
        println!("ConstantOfShape-{},{}", self.shape_array, self.o);
        if let Some(next) = &self.next_node {
            next.iter().for_each(|v| v.print());
        }
    }

    fn determine_output_shape(&mut self, omap: &mut TensorMap) {
        let [x, o] = omap.get_disjoint_mut([&self.shape_array, &self.o]);
        let x = x.map(|inner| &*inner);

        match (x, &self.value, o) {
            (Some(x), Some(value), Some(result)) => {
                if let Err(e) = x.constant_of_shape(value, result) {
                    println!("!!!{e}");
                }
            }
            (None, Some(value), Some(result)) => {
                *result = value.clone();
            }
            _ => panic!("ConstantOfShapeNode: missing input {}", self.shape_array),
        }

        if let Some(list) = &mut self.next_node {
            for next in list {
                next.determine_output_shape(omap);
            }
        }
    }
}

impl TypedArray {
    pub fn constant_of_shape(&self, value: &TypedArray, o: &mut TypedArray) -> anyhow::Result<()> {
        let shape: Vec<usize> = match &self {
            TypedArray::Int64(s) => s.iter().map(|&v| v as usize).collect(),
            _ => anyhow::bail!("ConstantOfShape: input shape must be I64"),
        };

        fill_from_elem!(
            value,
            &shape,
            o,
            [
                (Float, f32),
                (Double, f64),
                (Int8, i8),
                (Int16, i16),
                (Int32, i32),
                (Int64, i64),
                (Uint8, u8),
                (Uint16, u16),
                (Uint32, u32),
                (Uint64, u64),
                (Bool, bool)
            ]
        );

        Ok(())
    }
}