use crate::errors::*;
use crate::facedet::*;
use error_stack::ResultExt;
use mnn_bridge::ndarray::*;
use ndarray_resize::NdFir;
use std::path::Path;

#[derive(Debug)]
pub struct FaceDetection {
    handle: mnn_sync::SessionHandle,
}

pub struct FaceDetectionBuilder {
    schedule_config: Option<mnn::ScheduleConfig>,
    backend_config: Option<mnn::BackendConfig>,
    model: mnn::Interpreter,
}

impl FaceDetectionBuilder {
    pub fn new(model: impl AsRef<[u8]>) -> Result<Self> {
        Ok(Self {
            schedule_config: None,
            backend_config: None,
            model: mnn::Interpreter::from_bytes(model.as_ref())
                .map_err(|e| e.into_inner())
                .change_context(Error)
                .attach_printable("Failed to load model from bytes")?,
        })
    }

    pub fn with_forward_type(mut self, forward_type: mnn::ForwardType) -> Self {
        self.schedule_config
            .get_or_insert_default()
            .set_type(forward_type);
        self
    }

    pub fn with_schedule_config(mut self, config: mnn::ScheduleConfig) -> Self {
        self.schedule_config = Some(config);
        self
    }

    pub fn with_backend_config(mut self, config: mnn::BackendConfig) -> Self {
        self.backend_config = Some(config);
        self
    }

    pub fn build(self) -> Result<FaceDetection> {
        let model = self.model;
        let sc = self.schedule_config.unwrap_or_default();
        let handle = mnn_sync::SessionHandle::new(model, sc)
            .change_context(Error)
            .attach_printable("Failed to create session handle")?;
        Ok(FaceDetection { handle })
    }
}

impl FaceDetection {
    pub fn builder<T: AsRef<[u8]>>(
        model: T,
    ) -> std::result::Result<FaceDetectionBuilder, error_stack::Report<crate::errors::Error>> {
        FaceDetectionBuilder::new(model)
    }
}

impl FaceDetector for FaceDetection {
    fn run_model(&mut self, image: ndarray::ArrayView3<u8>) -> Result<FaceDetectionModelOutput> {
        #[rustfmt::skip]
        let mut resized = image
            .fast_resize(1024, 1024, None)
            .change_context(Error)?
            .mapv(|f| f as f32);

        // Apply mean subtraction: [104, 117, 123]
        resized
            .axis_iter_mut(ndarray::Axis(2))
            .zip([104, 117, 123])
            .for_each(|(mut array, pixel)| {
                let pixel = pixel as f32;
                array.map_inplace(|v| *v -= pixel);
            });

        let mut resized = resized
            .permuted_axes((2, 0, 1))
            .insert_axis(ndarray::Axis(0))
            .as_standard_layout()
            .into_owned();

        use ::tap::*;
        let output = self
            .handle
            .run(move |sr| {
                let tensor = resized
                    .as_mnn_tensor_mut()
                    .attach_printable("Failed to convert ndarray to mnn tensor")
                    .change_context(mnn::error::ErrorKind::TensorError)?;
                tracing::trace!("Image Tensor shape: {:?}", tensor.shape());
                let (intptr, session) = sr.both_mut();
                tracing::trace!("Copying input tensor to host");
                unsafe {
                    let mut input = intptr.input_unresized::<f32>(session, "input")?;
                    tracing::trace!("Input shape: {:?}", input.shape());
                    intptr.resize_tensor_by_nchw::<mnn::View<&mut f32>, _>(
                        input.view_mut(),
                        1,
                        3,
                        1024,
                        1024,
                    );
                }
                intptr.resize_session(session);
                let mut input = intptr.input::<f32>(session, "input")?;
                tracing::trace!("Input shape: {:?}", input.shape());
                input.copy_from_host_tensor(tensor.view())?;

                tracing::info!("Running face detection session");
                intptr.run_session(&session)?;
                let output_tensor = intptr
                    .output::<f32>(&session, "bbox")?
                    .create_host_tensor_from_device(true)
                    .as_ndarray()
                    .to_owned();
                tracing::trace!("Output Bbox: \t\t{:?}", output_tensor.shape());
                let output_confidence = intptr
                    .output::<f32>(&session, "confidence")?
                    .create_host_tensor_from_device(true)
                    .as_ndarray::<ndarray::Ix3>()
                    .to_owned();
                tracing::trace!("Output Confidence: \t{:?}", output_confidence.shape());
                let output_landmark = intptr
                    .output::<f32>(&session, "landmark")?
                    .create_host_tensor_from_device(true)
                    .as_ndarray::<ndarray::Ix3>()
                    .to_owned();
                tracing::trace!("Output Landmark: \t{:?}", output_landmark.shape());
                Ok(FaceDetectionModelOutput {
                    bbox: output_tensor,
                    confidence: output_confidence,
                    landmark: output_landmark,
                })
            })
            .map_err(|e| e.into_inner())
            .change_context(Error)?;
        Ok(output)
    }
}