use super::neural_types::*;
use super::neural_activations::*;

pub struct DenseLayer {
    pub input_dim: usize,
    pub output_dim: usize,
    pub weights: Array2<f64>,
    pub bias: Array1<f64>,
    pub activation: Option<CNNActivation>,
    pub use_bias: bool,
}

impl DenseLayer {
    pub fn new(input_dim: usize, output_dim: usize, activation: Option<CNNActivation>) -> Self {
        let mut rng = Random::seed_from_u64(42);
        let scale = (2.0 / (input_dim + output_dim) as f64).sqrt();

        let mut weights = Array2::zeros((input_dim, output_dim));
        for i in 0..input_dim {
            for j in 0..output_dim {
                weights[(i, j)] = rng.gen_range(-scale..scale);
            }
        }

        Self {
            input_dim,
            output_dim,
            weights,
            bias: Array1::zeros(output_dim),
            activation,
            use_bias: true,
        }
    }

    pub fn with_bias(mut self, use_bias: bool) -> Self {
        self.use_bias = use_bias;
        if !use_bias {
            self.bias = Array1::zeros(self.output_dim);
        }
        self
    }

    pub fn forward(&self, input: &Array2<f64>) -> Array2<f64> {
        let linear_output = input.dot(&self.weights);

        let output_with_bias = if self.use_bias {
            &linear_output + &self.bias
        } else {
            linear_output
        };

        if let Some(activation) = &self.activation {
            self.apply_activation(&output_with_bias, activation)
        } else {
            output_with_bias
        }
    }

    fn apply_activation(&self, input: &Array2<f64>, activation: &CNNActivation) -> Array2<f64> {
        let mut result = Array2::zeros(input.raw_dim());

        for i in 0..input.nrows() {
            for j in 0..input.ncols() {
                result[(i, j)] = activation.apply(input[(i, j)]);
            }
        }

        result
    }

    pub fn update_weights(&mut self, weight_updates: &Array2<f64>, bias_updates: &Array1<f64>, learning_rate: f64) {
        self.weights = &self.weights - learning_rate * weight_updates;
        if self.use_bias {
            self.bias = &self.bias - learning_rate * bias_updates;
        }
    }

    pub fn get_parameters(&self) -> (Array2<f64>, Array1<f64>) {
        (self.weights.clone(), self.bias.clone())
    }

    pub fn set_parameters(&mut self, weights: Array2<f64>, bias: Array1<f64>) {
        self.weights = weights;
        self.bias = bias;
    }
}

pub struct Conv2DLayer {
    pub in_channels: usize,
    pub out_channels: usize,
    pub kernel_size: usize,
    pub stride: usize,
    pub padding: usize,
    pub weights: Array4<f64>,
    pub bias: Array1<f64>,
    pub activation: Option<CNNActivation>,
    pub use_bias: bool,
}

impl Conv2DLayer {
    pub fn new(
        in_channels: usize,
        out_channels: usize,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        activation: Option<CNNActivation>,
    ) -> Self {
        let mut rng = Random::seed_from_u64(42);
        let fan_in = in_channels * kernel_size * kernel_size;
        let fan_out = out_channels * kernel_size * kernel_size;
        let scale = (2.0 / (fan_in + fan_out) as f64).sqrt();

        let mut weights = Array4::zeros((out_channels, in_channels, kernel_size, kernel_size));
        for i in 0..out_channels {
            for j in 0..in_channels {
                for k in 0..kernel_size {
                    for l in 0..kernel_size {
                        weights[(i, j, k, l)] = rng.gen_range(-scale..scale);
                    }
                }
            }
        }

        Self {
            in_channels,
            out_channels,
            kernel_size,
            stride,
            padding,
            weights,
            bias: Array1::zeros(out_channels),
            activation,
            use_bias: true,
        }
    }

    pub fn forward(&self, input: &Array3<f64>) -> Array3<f64> {
        let padded_input = self.apply_padding(input);
        let convolved = self.convolve(&padded_input);

        if let Some(activation) = &self.activation {
            self.apply_activation(&convolved, activation)
        } else {
            convolved
        }
    }

    fn apply_padding(&self, input: &Array3<f64>) -> Array3<f64> {
        if self.padding == 0 {
            return input.clone();
        }

        let (height, width, channels) = input.dim();
        let padded_height = height + 2 * self.padding;
        let padded_width = width + 2 * self.padding;
        let mut padded = Array3::zeros((padded_height, padded_width, channels));

        for c in 0..channels {
            for h in 0..height {
                for w in 0..width {
                    padded[(h + self.padding, w + self.padding, c)] = input[(h, w, c)];
                }
            }
        }

        padded
    }

    fn convolve(&self, input: &Array3<f64>) -> Array3<f64> {
        let (input_height, input_width, _) = input.dim();
        let output_height = (input_height - self.kernel_size) / self.stride + 1;
        let output_width = (input_width - self.kernel_size) / self.stride + 1;

        let mut output = Array3::zeros((output_height, output_width, self.out_channels));

        for f in 0..self.out_channels {
            for oh in 0..output_height {
                for ow in 0..output_width {
                    let mut sum = if self.use_bias { self.bias[f] } else { 0.0 };

                    for kh in 0..self.kernel_size {
                        for kw in 0..self.kernel_size {
                            for c in 0..self.in_channels {
                                let ih = oh * self.stride + kh;
                                let iw = ow * self.stride + kw;
                                sum += input[(ih, iw, c)] * self.weights[(f, c, kh, kw)];
                            }
                        }
                    }

                    output[(oh, ow, f)] = sum;
                }
            }
        }

        output
    }

    fn apply_activation(&self, input: &Array3<f64>, activation: &CNNActivation) -> Array3<f64> {
        let mut result = Array3::zeros(input.raw_dim());

        for i in 0..input.dim().0 {
            for j in 0..input.dim().1 {
                for k in 0..input.dim().2 {
                    result[(i, j, k)] = activation.apply(input[(i, j, k)]);
                }
            }
        }

        result
    }
}

pub struct MaxPool2DLayer {
    pub pool_size: usize,
    pub stride: usize,
}

impl MaxPool2DLayer {
    pub fn new(pool_size: usize, stride: usize) -> Self {
        Self { pool_size, stride }
    }

    pub fn forward(&self, input: &Array3<f64>) -> Array3<f64> {
        let (input_height, input_width, channels) = input.dim();
        let output_height = (input_height - self.pool_size) / self.stride + 1;
        let output_width = (input_width - self.pool_size) / self.stride + 1;

        let mut output = Array3::zeros((output_height, output_width, channels));

        for c in 0..channels {
            for oh in 0..output_height {
                for ow in 0..output_width {
                    let mut max_val = f64::NEG_INFINITY;

                    for ph in 0..self.pool_size {
                        for pw in 0..self.pool_size {
                            let ih = oh * self.stride + ph;
                            let iw = ow * self.stride + pw;

                            if ih < input_height && iw < input_width {
                                max_val = max_val.max(input[(ih, iw, c)]);
                            }
                        }
                    }

                    output[(oh, ow, c)] = max_val;
                }
            }
        }

        output
    }
}

pub struct AvgPool2DLayer {
    pub pool_size: usize,
    pub stride: usize,
}

impl AvgPool2DLayer {
    pub fn new(pool_size: usize, stride: usize) -> Self {
        Self { pool_size, stride }
    }

    pub fn forward(&self, input: &Array3<f64>) -> Array3<f64> {
        let (input_height, input_width, channels) = input.dim();
        let output_height = (input_height - self.pool_size) / self.stride + 1;
        let output_width = (input_width - self.pool_size) / self.stride + 1;

        let mut output = Array3::zeros((output_height, output_width, channels));

        for c in 0..channels {
            for oh in 0..output_height {
                for ow in 0..output_width {
                    let mut sum = 0.0;
                    let mut count = 0;

                    for ph in 0..self.pool_size {
                        for pw in 0..self.pool_size {
                            let ih = oh * self.stride + ph;
                            let iw = ow * self.stride + pw;

                            if ih < input_height && iw < input_width {
                                sum += input[(ih, iw, c)];
                                count += 1;
                            }
                        }
                    }

                    output[(oh, ow, c)] = if count > 0 { sum / count as f64 } else { 0.0 };
                }
            }
        }

        output
    }
}

pub struct DropoutLayer {
    pub dropout_rate: f64,
    pub training: bool,
}

impl DropoutLayer {
    pub fn new(dropout_rate: f64) -> Self {
        Self {
            dropout_rate,
            training: true,
        }
    }

    pub fn set_training(&mut self, training: bool) {
        self.training = training;
    }

    pub fn forward_2d(&self, input: &Array2<f64>) -> Array2<f64> {
        if !self.training || self.dropout_rate == 0.0 {
            return input.clone();
        }

        let mut rng = Random::seed_from_u64(42);
        let mut result = input.clone();

        for elem in result.iter_mut() {
            if rng.random_range(0.0, 1.0) < self.dropout_rate {
                *elem = 0.0;
            } else {
                *elem /= 1.0 - self.dropout_rate;
            }
        }

        result
    }

    pub fn forward_3d(&self, input: &Array3<f64>) -> Array3<f64> {
        if !self.training || self.dropout_rate == 0.0 {
            return input.clone();
        }

        let mut rng = Random::seed_from_u64(42);
        let mut result = input.clone();

        for elem in result.iter_mut() {
            if rng.random_range(0.0, 1.0) < self.dropout_rate {
                *elem = 0.0;
            } else {
                *elem /= 1.0 - self.dropout_rate;
            }
        }

        result
    }
}

pub struct BatchNormLayer {
    pub num_features: usize,
    pub eps: f64,
    pub momentum: f64,
    pub affine: bool,
    pub running_mean: Array1<f64>,
    pub running_var: Array1<f64>,
    pub weight: Array1<f64>,
    pub bias: Array1<f64>,
    pub training: bool,
}

impl BatchNormLayer {
    pub fn new(num_features: usize, eps: f64, momentum: f64, affine: bool) -> Self {
        Self {
            num_features,
            eps,
            momentum,
            affine,
            running_mean: Array1::zeros(num_features),
            running_var: Array1::ones(num_features),
            weight: Array1::ones(num_features),
            bias: Array1::zeros(num_features),
            training: true,
        }
    }

    pub fn set_training(&mut self, training: bool) {
        self.training = training;
    }

    pub fn forward_2d(&mut self, input: &Array2<f64>) -> Array2<f64> {
        let (batch_size, features) = input.dim();

        if features != self.num_features {
            panic!("Input features {} do not match layer features {}", features, self.num_features);
        }

        let (mean, var) = if self.training {
            let mean = self.compute_mean_2d(input);
            let var = self.compute_variance_2d(input, &mean);

            self.update_running_stats(&mean, &var);
            (mean, var)
        } else {
            (self.running_mean.clone(), self.running_var.clone())
        };

        let mut normalized = Array2::zeros((batch_size, features));

        for i in 0..batch_size {
            for j in 0..features {
                let normalized_val = (input[(i, j)] - mean[j]) / (var[j] + self.eps).sqrt();
                normalized[(i, j)] = if self.affine {
                    self.weight[j] * normalized_val + self.bias[j]
                } else {
                    normalized_val
                };
            }
        }

        normalized
    }

    pub fn forward_3d(&mut self, input: &Array3<f64>) -> Array3<f64> {
        let (height, width, channels) = input.dim();

        if channels != self.num_features {
            panic!("Input channels {} do not match layer features {}", channels, self.num_features);
        }

        let (mean, var) = if self.training {
            let mean = self.compute_mean_3d(input);
            let var = self.compute_variance_3d(input, &mean);

            self.update_running_stats(&mean, &var);
            (mean, var)
        } else {
            (self.running_mean.clone(), self.running_var.clone())
        };

        let mut normalized = Array3::zeros((height, width, channels));

        for h in 0..height {
            for w in 0..width {
                for c in 0..channels {
                    let normalized_val = (input[(h, w, c)] - mean[c]) / (var[c] + self.eps).sqrt();
                    normalized[(h, w, c)] = if self.affine {
                        self.weight[c] * normalized_val + self.bias[c]
                    } else {
                        normalized_val
                    };
                }
            }
        }

        normalized
    }

    fn compute_mean_2d(&self, input: &Array2<f64>) -> Array1<f64> {
        let (batch_size, features) = input.dim();
        let mut mean = Array1::zeros(features);

        for j in 0..features {
            let mut sum = 0.0;
            for i in 0..batch_size {
                sum += input[(i, j)];
            }
            mean[j] = sum / batch_size as f64;
        }

        mean
    }

    fn compute_variance_2d(&self, input: &Array2<f64>, mean: &Array1<f64>) -> Array1<f64> {
        let (batch_size, features) = input.dim();
        let mut var = Array1::zeros(features);

        for j in 0..features {
            let mut sum_sq_diff = 0.0;
            for i in 0..batch_size {
                let diff = input[(i, j)] - mean[j];
                sum_sq_diff += diff * diff;
            }
            var[j] = sum_sq_diff / batch_size as f64;
        }

        var
    }

    fn compute_mean_3d(&self, input: &Array3<f64>) -> Array1<f64> {
        let (height, width, channels) = input.dim();
        let mut mean = Array1::zeros(channels);

        for c in 0..channels {
            let mut sum = 0.0;
            for h in 0..height {
                for w in 0..width {
                    sum += input[(h, w, c)];
                }
            }
            mean[c] = sum / (height * width) as f64;
        }

        mean
    }

    fn compute_variance_3d(&self, input: &Array3<f64>, mean: &Array1<f64>) -> Array1<f64> {
        let (height, width, channels) = input.dim();
        let mut var = Array1::zeros(channels);

        for c in 0..channels {
            let mut sum_sq_diff = 0.0;
            for h in 0..height {
                for w in 0..width {
                    let diff = input[(h, w, c)] - mean[c];
                    sum_sq_diff += diff * diff;
                }
            }
            var[c] = sum_sq_diff / (height * width) as f64;
        }

        var
    }

    fn update_running_stats(&mut self, batch_mean: &Array1<f64>, batch_var: &Array1<f64>) {
        for i in 0..self.num_features {
            self.running_mean[i] = (1.0 - self.momentum) * self.running_mean[i] + self.momentum * batch_mean[i];
            self.running_var[i] = (1.0 - self.momentum) * self.running_var[i] + self.momentum * batch_var[i];
        }
    }
}