use super::layers::*;
use crate::nn::loss_functions::get_loss_function;
use crate::nn::loss_functions::LossFunctionType;
use crate::numeric::FloatingPoint;
use crate::tensor::math::TensorMath;
use crate::tensor::Tensor;
use crate::utils::get_current_lr;

use crate::nn::LayerData;
use crate::nn::ModelData;
use std::fs;
use std::fs::File;
use std::io;
use std::io::Write;
use std::path::Path;

use crate::nn::types::TrainingConfig;
use crate::nn::types::TrainingHook;
use colored::*;

/// Feed-forward neural network container.
///
/// Holds an ordered list of `Layer` trait objects, the configured loss
/// function, and training metadata (parameter count, name/label, and
/// learning state). Instances are used to `predict`, `fit`, and
/// `save_model`.
pub struct NeuralNet<T, D>
where
    T: Tensor<D> + TensorMath<D, MathOutput = T> + 'static,
    D: FloatingPoint,
{
    pub layers: Vec<Box<dyn Layer<T, D>>>,
    pub parameter_count: u64,
    pub label: String,
    pub name: String,
    current_epoch: usize,
    current_lr: D,
    last_train_loss: D,
    last_val_loss: D,
    loss_fn_type: LossFunctionType,
    epoch_vs_loss: Vec<(usize, D, D)>,
}

impl<T, D> NeuralNet<T, D>
where
    T: TensorMath<D, MathOutput = T> + 'static,
    D: FloatingPoint,
{
    pub fn new(layers: Vec<Box<dyn Layer<T, D>>>, model: ModelData<D>) -> Self {
        Self {
            layers,
            loss_fn_type: model.loss_fn_type,
            parameter_count: model.parameter_count,
            label: model.label,
            name: model.name,
            current_epoch: model.epoch,
            current_lr: model.saved_lr,
            last_train_loss: D::zero(),
            last_val_loss: D::zero(),
            epoch_vs_loss: model.epoch_error,
        }
    }

    /// Get the current epoch (useful for resuming training)
    pub fn get_current_epoch(&self) -> usize {
        self.current_epoch
    }

    /// Set the current epoch (useful for resuming training)
    pub fn set_current_epoch(&mut self, epoch: usize) {
        self.current_epoch = epoch;
    }

    /// Run a forward pass and return the network output for `input`.
    pub fn predict(&mut self, input: &T) -> Result<T, String> {
        let mut output = input.add(&T::zeroes(input.get_shape())).unwrap();

        for layer in &mut self.layers {
            output = layer.forward(&output, false).unwrap();
        }
        Ok(output)
    }

    pub fn fit<F>(
        &mut self,
        x_train: &T,
        y_train: &T,
        x_val: &T,
        y_val: &T,
        config: TrainingConfig<D>,
        mut hook_config: TrainingHook<F, Self, D>,
    ) -> Result<(), String>
    where
        F: FnMut(usize, D, D, D, &mut Self),
    {
        let lr_min = D::from_f64(1e-6);

        let total_timeline = config.epochs;
        let hook_interval = match config.epochs > hook_config.interval {
            true => hook_config.interval,
            false => config.epochs,
        };

        let mut last_good_model = self.get_model();
        let mut best_val_loss = D::from_f64(f64::MAX);
        let mut patience_counter = 0;
        let patience = 5; // For now, later may take it from config
        let epsilon = D::from_f64(1e-6);

        for i in config.epoch_offset..config.epochs {
            let global_i = i;

            self.current_epoch = i;

            let current_lr = get_current_lr(
                config.base_lr,
                config.lr_adjustment,
                lr_min,
                total_timeline,
                global_i,
            );

            self.current_lr = current_lr;

            print!(
                "\rProcessing epoch: {}/{}, current lr: {current_lr:.6}",
                self.current_epoch, config.epochs
            );
            io::stdout().flush().unwrap();

            let mut output = x_train.add(&T::zeroes(x_train.get_shape())).unwrap();
            for layer in &mut self.layers {
                output = layer.forward(&output, true).unwrap();
            }
            T::synchronize();

            let (loss, loss_prime) = get_loss_function::<T, D>(&self.loss_fn_type);

            let err = (loss)(y_train, &output).unwrap().sum().unwrap().get_data()[0];
            T::synchronize();

            let mut error_prime = loss_prime(y_train, &output).unwrap();

            for layer in self.layers.iter_mut().rev() {
                error_prime = layer
                    .backward(&error_prime, current_lr, config.weight_normalization)
                    .unwrap();
            }
            T::synchronize();

            let mut v_output = x_val.add(&T::zeroes(x_val.get_shape())).unwrap();
            for layer in &mut self.layers {
                v_output = layer.forward(&v_output, false).unwrap();
            }
            let err_val = (loss)(y_val, &v_output).unwrap().sum().unwrap().get_data()[0];
            T::synchronize();

            self.epoch_vs_loss.push((i, err, err_val));

            if err_val < best_val_loss - epsilon {
                // We found a new best model
                best_val_loss = err_val;
                last_good_model = self.get_model();
                patience_counter = 0;

                self.last_train_loss = err;
                self.last_val_loss = err_val;
            } else {
                // No improvement
                // patience_counter += 1;

                if patience_counter >= patience {
                    println!(
                        "\n{}",
                        "Early stopping: No improvement in validation loss."
                            .bold()
                            .red()
                    );
                    Self::write_model_to_disk(
                        last_good_model,
                        format!("model_outputs/{}/model.json", self.name).as_str(),
                    );
                    break;
                }
            }

            // Hook (Periodic Reporting)
            if i == 0 || i % hook_interval == 0 {
                T::synchronize();
                (hook_config.callback)(i, err, err_val, current_lr, self);
            }
        }

        T::synchronize();
        Ok(())
    }

    fn get_model(&self) -> ModelData<D> {
        let mut model_storage = ModelData {
            name: self.name.clone(),
            parameter_count: self.parameter_count,
            layers: Vec::new(),
            epoch: self.current_epoch,
            saved_lr: self.current_lr,
            loss_fn_type: self.loss_fn_type.clone(),
            epoch_error: self.epoch_vs_loss.clone(),
            label: self.label.clone(),
        };

        for (i, layer) in self.layers.iter().enumerate() {
            let (w, s) = match layer.get_parameters() {
                Some(wt) => (wt.get_data().to_vec(), wt.get_shape().to_vec()),
                None => (Vec::<D>::new(), Vec::<u32>::new()),
            };

            let layer_info = LayerData {
                name: layer.name().to_string(),
                index: i,
                weights: w,
                shape: s,
                layer_type: layer.layer_type().clone(),
            };

            model_storage.layers.push(layer_info);
        }
        model_storage
    }

    fn write_model_to_disk(model: ModelData<D>, filepath: &str) {
        let json_data =
            serde_json::to_string_pretty(&model).expect("Failed to serialize model weights");

        let path = Path::new(filepath);
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent).unwrap(); // Creates all directories if they don't exist
        }

        let mut file = File::create(filepath)
            .unwrap_or_else(|_| panic!("Model successfully saved to {}", filepath));
        let _ = file.write_all(json_data.as_bytes());

        let _ = file.flush();

        let status_text = format!("Model successfully saved to {}", filepath);
        println!("{}", status_text.green());
    }

    /// Serialize and write the model weights and metadata to `filepath`.
    pub fn save_model(&self, filepath: &str) {
        let model_storage = self.get_model();
        Self::write_model_to_disk(model_storage, filepath);
    }

    /// Get the epoch vs loss data for plotting
    pub fn get_epoch_error(&self) -> Vec<(usize, D, D)> {
        self.epoch_vs_loss.clone()
    }
}