use super::{Node, NodeCodegen};
use crate::burn::{OtherType, Scope, TensorType, ToTokens, Type};
use burn::record::PrecisionSettings;
use proc_macro2::TokenStream;
use quote::quote;

#[derive(Debug, Clone)]
pub struct ResizeNode {
    pub field: OtherType,
    pub input: TensorType,
    pub output: TensorType,
    mode: String,
    scales: Vec<f32>,
    sizes: Vec<usize>,
}

impl ResizeNode {
    pub fn new<S: AsRef<str>>(
        name: S,
        input: TensorType,
        output: TensorType,
        mode: String,
        scales: Vec<f32>,
        sizes: Vec<usize>,
    ) -> Self {
        let ty = if input.dim == 3 {
            quote! {
                Interpolate1d
            }
        } else if input.dim == 4 {
            quote! {
                Interpolate2d
            }
        } else {
            panic!("Unsupported input dimension for resize node");
        };

        Self {
            field: OtherType::new(name, ty),
            input,
            output,
            mode,
            scales,
            sizes,
        }
    }
}

impl<PS: PrecisionSettings> NodeCodegen<PS> for ResizeNode {
    fn output_types(&self) -> Vec<Type> {
        vec![Type::Tensor(self.output.clone())]
    }

    fn input_types(&self) -> Vec<Type> {
        vec![Type::Tensor(self.input.clone())]
    }

    fn field_type(&self) -> Option<Type> {
        Some(Type::Other(self.field.clone()))
    }

    fn field_init(&self) -> Option<TokenStream> {
        let name = &self.field.name;

        let mode = match self.mode.as_str() {
            "nearest" => quote! { InterpolateMode::Nearest },
            "linear" => quote! { InterpolateMode::Linear },
            "cubic" => quote! { InterpolateMode::Cubic },
            _ => panic!("Unsupported mode for resize node"),
        };

        let tokens = if self.input.dim == 3 {
            let size = if let Some(size) = self.sizes.first() {
                let size = size.to_tokens();
                quote! { Some(#size) }
            } else {
                quote! { None }
            };

            let scale_factor = if let Some(scale) = self.scales.first() {
                let scale = scale.to_tokens();
                quote! { Some(#scale) }
            } else {
                quote! { None }
            };

            quote! {
                let #name = Interpolate1dConfig::new()
                    .with_output_size(#size)
                    .with_scale_factor(#scale_factor)
                    .with_mode(#mode)
                    .init();
            }
        } else if self.input.dim == 4 {
            let size = if self.sizes.len() == 2 {
                let h = self.sizes[0].to_tokens();
                let w = self.sizes[1].to_tokens();
                quote! { Some([#h, #w]) }
            } else {
                quote! { None }
            };

            let scale_factor = if self.scales.len() == 2 {
                let h = self.scales[0].to_tokens();
                let w = self.scales[1].to_tokens();
                quote! { Some([#h, #w]) }
            } else {
                quote! { None }
            };

            quote! {
                let #name = Interpolate2dConfig::new()
                    .with_output_size(#size)
                    .with_scale_factor(#scale_factor)
                    .with_mode(#mode)
                    .init();
            }
        } else {
            panic!("Unsupported input dimension for resize node");
        };

        Some(tokens)
    }

    fn register_imports(&self, imports: &mut crate::burn::BurnImports) {
        imports.register("burn::nn::interpolate::InterpolateMode");
        if self.input.dim == 3 {
            imports.register("burn::nn::interpolate::Interpolate1dConfig");
            imports.register("burn::nn::interpolate::Interpolate1d");
        } else if self.input.dim == 4 {
            imports.register("burn::nn::interpolate::Interpolate2dConfig");
            imports.register("burn::nn::interpolate::Interpolate2d");
        } else {
            panic!("Unsupported input dimension for resize node");
        }
    }

    fn field_serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        S::serialize_none(serializer)
    }

    fn forward(&self, scope: &mut Scope, node_position: usize) -> TokenStream {
        let input = scope.tensor_use_owned(&self.input, node_position);
        let output = &self.output.name;
        let field = &self.field.name;

        quote! {
            let #output = self.#field.forward(#input);
        }
    }

    fn into_node(self) -> Node<PS> {
        Node::Resize(self)
    }
}

#[cfg(test)]
mod tests {
    use burn::record::FullPrecisionSettings;

    use super::*;
    use crate::burn::{
        graph::BurnGraph,
        node::{resize::ResizeNode, test::assert_tokens},
        TensorType,
    };

    #[test]
    fn test_codegen_nodes_2d() {
        let mut graph = BurnGraph::<FullPrecisionSettings>::default();

        graph.register(ResizeNode::new(
            "resize",
            TensorType::new_float("tensor1", 4),
            TensorType::new_float("tensor2", 4),
            "nearest".to_string(),
            vec![0.5, 0.5],
            vec![],
        ));

        graph.register_input_output(vec!["tensor1".to_string()], vec!["tensor2".to_string()]);

        let expected = quote! {
            use burn::nn::interpolate::Interpolate2d;
            use burn::nn::interpolate::Interpolate2dConfig;
            use burn::nn::interpolate::InterpolateMode;
            use burn::{
                module::Module,
                tensor::{backend::Backend, Tensor},
            };
            #[derive(Module, Debug)]
            pub struct Model<B: Backend> {
                resize: Interpolate2d,
                phantom: core::marker::PhantomData<B>,
                device: burn::module::Ignored<B::Device>,
            }
            impl<B: Backend> Model<B> {
                #[allow(unused_variables)]
                pub fn new(device: &B::Device) -> Self {
                    let resize = Interpolate2dConfig::new()
                        .with_output_size(None)
                        .with_scale_factor(Some([0.5, 0.5]))
                        .with_mode(InterpolateMode::Nearest)
                        .init();
                    Self {
                        resize,
                        phantom: core::marker::PhantomData,
                        device: burn::module::Ignored(device.clone()),
                    }
                }
                #[allow(clippy::let_and_return, clippy::approx_constant)]
                pub fn forward(&self, tensor1: Tensor<B, 4>) -> Tensor<B, 4> {
                    let tensor2 = self.resize.forward(tensor1);
                    tensor2
                }
            }
        };

        assert_tokens(graph.codegen(), expected);
    }

    #[test]
    fn test_codegen_nodes_1d() {
        let mut graph = BurnGraph::<FullPrecisionSettings>::default();

        graph.register(ResizeNode::new(
            "resize",
            TensorType::new_float("tensor1", 3),
            TensorType::new_float("tensor2", 3),
            "cubic".to_string(),
            vec![],
            vec![20],
        ));

        graph.register_input_output(vec!["tensor1".to_string()], vec!["tensor2".to_string()]);

        let expected = quote! {
            use burn::nn::interpolate::Interpolate1d;
            use burn::nn::interpolate::Interpolate1dConfig;
            use burn::nn::interpolate::InterpolateMode;
            use burn::{
                module::Module,
                tensor::{backend::Backend, Tensor},
            };
            #[derive(Module, Debug)]
            pub struct Model<B: Backend> {
                resize: Interpolate1d,
                phantom: core::marker::PhantomData<B>,
                device: burn::module::Ignored<B::Device>,
            }
            impl<B: Backend> Model<B> {
                #[allow(unused_variables)]
                pub fn new(device: &B::Device) -> Self {
                    let resize = Interpolate1dConfig::new()
                        .with_output_size(Some(20))
                        .with_scale_factor(None)
                        .with_mode(InterpolateMode::Cubic)
                        .init();
                    Self {
                        resize,
                        phantom: core::marker::PhantomData,
                        device: burn::module::Ignored(device.clone()),
                    }
                }
                #[allow(clippy::let_and_return, clippy::approx_constant)]
                pub fn forward(&self, tensor1: Tensor<B, 3>) -> Tensor<B, 3> {
                    let tensor2 = self.resize.forward(tensor1);
                    tensor2
                }
            }
        };

        assert_tokens(graph.codegen(), expected);
    }
}
