use super::GroupNorm;
use crate::model::load::*;

use std::error::Error;

use burn::{
    config::Config, 
    module::{Module, Param},
    nn,
    tensor::{
        backend::Backend,
        Tensor,
    },
};

use super::*;
use crate::model::groupnorm::load::load_group_norm;

fn load_conv_self_attention_block<B: Backend>(path: &str, device: &B::Device) -> Result<ConvSelfAttentionBlock<B>, Box<dyn Error>> {
    let norm = load_group_norm(&format!("{}/{}", path, "norm"), device)?;
    let q = load_conv2d(&format!("{}/{}", path, "q"), device)?;
    let k = load_conv2d(&format!("{}/{}", path, "k"), device)?;
    let v = load_conv2d(&format!("{}/{}", path, "v"), device)?;
    let proj_out = load_conv2d(&format!("{}/{}", path, "proj_out"), device)?;

    Ok(ConvSelfAttentionBlock { norm, q, k, v, proj_out })
}

fn load_resnet_block<B: Backend>(path: &str, device: &B::Device) -> Result<ResnetBlock<B>, Box<dyn Error>> {
    let norm1 = load_group_norm(&format!("{}/{}", path, "norm1"), device)?;
    let silu1 = SILU {};
    let conv1 = load_conv2d(&format!("{}/{}", path, "conv1"), device)?;
    let norm2 = load_group_norm(&format!("{}/{}", path, "norm2"), device)?;
    let silu2 = SILU {};
    let conv2 = load_conv2d(&format!("{}/{}", path, "conv2"), device)?;
    let nin_shortcut = load_conv2d(&format!("{}/{}", path, "nin_shortcut"), device).ok();

    Ok(ResnetBlock { norm1, silu1, conv1, norm2, silu2, conv2, nin_shortcut })
}

fn load_mid<B: Backend>(path: &str, device: &B::Device) -> Result<Mid<B>, Box<dyn Error>> {
    let block_1 = load_resnet_block(&format!("{}/{}", path, "block_1"), device)?;
    let attn = load_conv_self_attention_block(&format!("{}/{}", path, "attn"), device)?;
    let block_2 = load_resnet_block(&format!("{}/{}", path, "block_2"), device)?;

    Ok(Mid { block_1, attn, block_2 })
}

fn load_padded_conv2d<B: Backend>(path: &str, device: &B::Device) -> Result<PaddedConv2d<B>, Box<dyn Error>> {
    let conv = load_conv2d(&format!("{}/{}", path, "conv"), device)?;

    let channels = load_tensor::<B, 1>("channels", path, device)?;
    let channels = tensor_to_array_2(channels); 

    let kernel_size = load_usize::<B>("kernel_size", path, device)?;
    let stride = load_usize::<B>("stride", path, device)?;

    let padding = load_tensor::<B, 1>("padding", path, device)?;
    let padding: [usize; 4] = tensor_to_array(padding);
    let padding = Padding::new(padding[0], padding[1], padding[2], padding[3]);

    let mut record = conv.into_record();

    let mut padded_conv: PaddedConv2d<B> = PaddedConv2dConfig::new(channels, kernel_size, padding).with_stride(stride).init();
    let padding_actual = PaddingConfig2d::Explicit(padded_conv.padding_actual[0], padded_conv.padding_actual[1]);

    record.padding = <PaddingConfig2d as Module<B>>::into_record(padding_actual);
    padded_conv.conv = padded_conv.conv.load_record(record);
    

    Ok(padded_conv)
}

fn load_decoder_block<B: Backend>(path: &str, device: &B::Device) -> Result<DecoderBlock<B>, Box<dyn Error>> {
    let res1 = load_resnet_block(&format!("{}/{}", path, "res1"), device)?;
    let res2 = load_resnet_block(&format!("{}/{}", path, "res2"), device)?;
    let res3 = load_resnet_block(&format!("{}/{}", path, "res3"), device)?;
    let upsampler = load_conv2d(&format!("{}/{}", path, "upsampler"), device).ok();

    Ok(DecoderBlock { res1, res2, res3, upsampler })
}

fn load_encoder_block<B: Backend>(path: &str, device: &B::Device) -> Result<EncoderBlock<B>, Box<dyn Error>> {
    let res1 = load_resnet_block(&format!("{}/{}", path, "res1"), device)?;
    let res2 = load_resnet_block(&format!("{}/{}", path, "res2"), device)?;
    let downsampler = load_padded_conv2d(&format!("{}/{}", path, "downsampler"), device).ok();

    Ok(EncoderBlock { res1, res2, downsampler })
}

fn load_decoder<B: Backend>(path: &str, device: &B::Device) -> Result<Decoder<B>, Box<dyn Error>> {
    let conv_in = load_conv2d(&format!("{}/{}", path, "conv_in"), device)?;
    let mid = load_mid(&format!("{}/{}", path, "mid"), device)?;

    let n_block = load_usize::<B>("n_block", path, device)?;
    let mut blocks = (0..n_block)
        .into_iter()
        .map(|i| {
            load_decoder_block::<B>(&format!("{}/blocks/{}", path, i), device)
        }).collect::<Result<Vec<_>, _>>()?;

    let norm_out = load_group_norm(&format!("{}/{}", path, "norm_out"), device)?;
    let silu = SILU {};
    let conv_out = load_conv2d(&format!("{}/{}", path, "conv_out"), device)?;

    Ok(Decoder { conv_in, mid, blocks, norm_out, silu, conv_out })
}

fn load_encoder<B: Backend>(path: &str, device: &B::Device) -> Result<Encoder<B>, Box<dyn Error>> {
    let conv_in = load_conv2d(&format!("{}/{}", path, "conv_in"), device)?;
    let mid = load_mid(&format!("{}/{}", path, "mid"), device)?;

    let n_block = load_usize::<B>("n_block", path, device)?;
    let mut blocks = (0..n_block)
        .into_iter()
        .map(|i| {
            load_encoder_block::<B>(&format!("{}/blocks/{}", path, i), device)
        }).collect::<Result<Vec<_>, _>>()?;

    let norm_out = load_group_norm(&format!("{}/{}", path, "norm_out"), device)?;
    let silu = SILU {};
    let conv_out = load_conv2d(&format!("{}/{}", path, "conv_out"), device)?;

    Ok(Encoder { conv_in, mid, blocks, norm_out, silu, conv_out })
}

pub fn load_autoencoder<B: Backend>(path: &str, device: &B::Device) -> Result<Autoencoder<B>, Box<dyn Error>> {
    let encoder = load_encoder(&format!("{}/{}", path, "encoder"), device)?;
    let decoder = load_decoder(&format!("{}/{}", path, "decoder"), device)?;
    let quant_conv = load_conv2d(&format!("{}/{}", path, "quant_conv"), device)?;
    let post_quant_conv = load_conv2d(&format!("{}/{}", path, "post_quant_conv"), device)?;

    Ok(Autoencoder { encoder, decoder, quant_conv, post_quant_conv })
}