use alloc::format;
use burn::tensor::ops::DeformConvOptions;

use burn_core as burn;

use crate::PaddingConfig2d;
use burn::config::Config;
use burn::module::Initializer;
use burn::module::{Content, DisplaySettings, Module, ModuleDisplay, Param};
use burn::tensor::Device;
use burn::tensor::Tensor;
use burn::tensor::module::deform_conv2d;

use crate::conv::checks;

/// Configuration to create a [deformable 2D convolution](DeformConv2d) layer, using the [init function](DeformConv2dConfig::init).
#[derive(Config, Debug)]
pub struct DeformConv2dConfig {
    /// The number of channels.
    pub channels: [usize; 2],
    /// The size of the kernel.
    pub kernel_size: [usize; 2],
    /// The stride of the convolution.
    #[config(default = "[1, 1]")]
    pub stride: [usize; 2],
    /// Spacing between kernel elements.
    #[config(default = "[1, 1]")]
    pub dilation: [usize; 2],
    /// Controls the connections between input and output channels.
    #[config(default = "1")]
    pub weight_groups: usize,
    /// Offset groups.
    #[config(default = "1")]
    pub offset_groups: usize,
    /// The padding configuration.
    ///
    /// ### Warning
    /// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel
    /// size is not supported as it will not produce the same output size.
    #[config(default = "PaddingConfig2d::Valid")]
    pub padding: PaddingConfig2d,
    /// If bias should be added to the output.
    #[config(default = true)]
    pub bias: bool,
    /// The type of function used to initialize neural network parameters
    #[config(
        default = "Initializer::KaimingUniform{gain:1.0/num_traits::Float::sqrt(3.0),fan_out_only:false}"
    )]
    pub initializer: Initializer,
}

/// Applies a deformable 2D convolution over input tensors.
///
/// Should be created with [DeformConv2dConfig].
#[derive(Module, Debug)]
#[module(custom_display)]
pub struct DeformConv2d {
    /// Tensor of shape `[channels_out, channels_in / groups, kernel_size_1, kernel_size_2]`
    pub weight: Param<Tensor<4>>,
    /// Tensor of shape `[channels_out]`
    pub bias: Option<Param<Tensor<1>>>,
    /// Stride of the convolution.
    pub stride: [usize; 2],
    /// Size of the kernel.
    pub kernel_size: [usize; 2],
    /// Spacing between kernel elements.
    pub dilation: [usize; 2],
    /// Controls the connections between input and output channels.
    pub weight_groups: usize,
    /// Offset groups.
    pub offset_groups: usize,
    /// The padding configuration.
    #[module(skip)]
    pub padding: PaddingConfig2d,
}

impl DeformConv2dConfig {
    /// Initialize a new [DeformConv2d](DeformConv2d) module.
    pub fn init(&self, device: &Device) -> DeformConv2d {
        checks::checks_channels_div_groups(self.channels[0], self.channels[1], self.weight_groups);
        if self.padding == PaddingConfig2d::Same {
            checks::check_same_padding_support(&self.kernel_size);
        }

        let shape = [
            self.channels[1],
            self.channels[0] / self.weight_groups,
            self.kernel_size[0],
            self.kernel_size[1],
        ];

        let k = self.kernel_size.iter().product::<usize>();
        let fan_in = self.channels[0] / self.weight_groups * k;
        let fan_out = self.channels[1] / self.weight_groups * k;

        let weight = self
            .initializer
            .init_with(shape, Some(fan_in), Some(fan_out), device);
        let mut bias = None;

        if self.bias {
            bias = Some(self.initializer.init_with(
                [self.channels[1]],
                Some(fan_in),
                Some(fan_out),
                device,
            ));
        }

        DeformConv2d {
            weight,
            bias,
            stride: self.stride,
            kernel_size: self.kernel_size,
            dilation: self.dilation,
            padding: self.padding.clone(),
            weight_groups: self.weight_groups,
            offset_groups: self.weight_groups,
        }
    }
}

impl ModuleDisplay for DeformConv2d {
    fn custom_settings(&self) -> Option<DisplaySettings> {
        DisplaySettings::new()
            .with_new_line_after_attribute(false)
            .optional()
    }

    fn custom_content(&self, content: Content) -> Option<Content> {
        // Format the stride, kernel_size and dilation as strings, formatted as arrays instead of indexed.
        let stride = format!("{:?}", self.stride);
        let kernel_size = format!("{:?}", self.kernel_size);
        let dilation = format!("{:?}", self.dilation);

        content
            .add("stride", &stride)
            .add("kernel_size", &kernel_size)
            .add("dilation", &dilation)
            .add("weight_groups", &self.weight_groups)
            .add("offset_groups", &self.offset_groups)
            .add_debug_attribute("padding", &self.padding)
            .optional()
    }
}

impl DeformConv2d {
    /// Applies the forward pass on the input tensor.
    ///
    /// See [deform_conv2d](burn::tensor::module::deform_conv2d) for more information.
    ///
    /// # Shapes
    ///
    /// - input: `[batch_size, channels_in, height_in, width_in]`
    /// - offset: `[batch_size, 2 * offset_groups * kernel_height * kernel_width, height_out, width_out]`
    /// - mask: `[batch_size, offset_groups * kernel_height * kernel_width, height_out, width_out]`
    /// - output: `[batch_size, channels_out, height_out, width_out]`
    pub fn forward(
        &self,
        input: Tensor<4>,
        offset: Tensor<4>,
        mask: Option<Tensor<4>>,
    ) -> Tensor<4> {
        let [_batch_size, _channels_in, height_in, width_in] = input.dims();
        let padding =
            self.padding
                .calculate_padding_2d(height_in, width_in, &self.kernel_size, &self.stride);
        deform_conv2d(
            input,
            offset,
            self.weight.val(),
            mask,
            self.bias.as_ref().map(|bias| bias.val()),
            DeformConvOptions::new(
                self.stride,
                padding,
                self.dilation,
                self.weight_groups,
                self.offset_groups,
            ),
        )
    }
}

#[cfg(test)]
mod tests {
    use burn::tensor::{ElementConversion, Tolerance};
    type FT = f32;

    use super::*;
    use burn::tensor::TensorData;

    #[test]
    fn initializer_default() {
        let device = Device::default();
        device.seed(0);

        let config = DeformConv2dConfig::new([5, 1], [5, 5]);
        let k = (config.channels[0] * config.kernel_size[0] * config.kernel_size[1]) as f64;
        let k = (config.offset_groups as f64 / k).sqrt().elem::<FT>();
        let conv = config.init(&device);

        conv.weight.to_data().assert_within_range(-k..k);
    }

    #[test]
    fn initializer_zeros() {
        let device = Device::default();
        device.seed(0);

        let config = DeformConv2dConfig::new([5, 2], [5, 5]).with_initializer(Initializer::Zeros);
        let conv = config.init(&device);

        assert_eq!(config.initializer, Initializer::Zeros);
        conv.weight.to_data().assert_approx_eq::<FT>(
            &TensorData::zeros::<f32, _>(conv.weight.shape()),
            Tolerance::default(),
        );
    }

    #[test]
    fn initializer_fan_out() {
        let device = Device::default();
        device.seed(0);

        let init = Initializer::KaimingUniform {
            gain: 1.0 / 3.0f64.sqrt(),
            fan_out_only: true, // test that fan_out is passed to `init_with()`
        };

        let config = DeformConv2dConfig::new([5, 1], [5, 5]).with_initializer(init.clone());
        let c = config.init(&device);
        let _ = c.weight.val(); // initializes param

        assert_eq!(config.initializer, init);
    }

    #[test]
    fn initializer_fan_with_groups_is_valid() {
        let device = Device::default();
        device.seed(0);

        let init = Initializer::KaimingUniform {
            gain: 1.0 / 3.0f64.sqrt(),
            fan_out_only: true,
        };

        let config = DeformConv2dConfig::new([4, 4], [1, 1])
            .with_initializer(init.clone())
            .with_weight_groups(4);
        let _ = config.init(&device);

        assert_eq!(config.initializer, init);
    }

    #[test]
    #[should_panic = "Both channels must be divisible by the number of groups."]
    fn channels_with_groups_is_invalid() {
        let device = Default::default();
        let config = DeformConv2dConfig::new([1, 4], [1, 1]).with_weight_groups(4);
        let _ = config.init(&device);
    }

    #[test]
    #[should_panic = "Same padding with an even kernel size is not supported"]
    fn same_with_even_kernel_is_invalid() {
        let device = Default::default();
        let config = DeformConv2dConfig::new([4, 4], [2, 2]).with_padding(PaddingConfig2d::Same);
        let _ = config.init(&device);
    }

    #[test]
    fn display() {
        let config = DeformConv2dConfig::new([5, 1], [5, 5]);
        let conv = config.init(&Default::default());

        assert_eq!(
            alloc::format!("{conv}"),
            "DeformConv2d {stride: [1, 1], kernel_size: [5, 5], dilation: [1, 1], weight_groups: 1, offset_groups: 1, padding: Valid, params: 126}"
        );
    }

    #[test]
    #[should_panic = "Number of channels in input tensor and input channels of convolution must be equal. got: 4, expected: 5"]
    fn input_channels_mismatch() {
        let config = DeformConv2dConfig::new([5, 3], [3, 3]);
        let conv = config.init(&Default::default());

        let input = Tensor::<4>::zeros([1, 4, 10, 10], &Default::default());
        let offset = Tensor::<4>::zeros([1, 2 * 3 * 3, 10, 10], &Default::default());
        let _ = conv.forward(input, offset, None);
    }
}
