use std::fs;
use std::fs::File;
use std::io::{Read, Write};

use serde::{Deserialize, Serialize};

use tuphium::activations::Activation;
use tuphium::error::{TuphiumError, TuphiumResult};
use tuphium::frame::csv;
use tuphium::frame::dataframe::{Column, DataFrame};
use tuphium::frame::dataset::Dataset;
use tuphium::layers::{Layer, Module};
use tuphium::losses::{BinaryCrossEntropy, Loss};
use tuphium::optimizers::{Optimizer, SGD};

/*
* Model architecture
*/
#[derive(Serialize, Deserialize)]
struct MLP {
    layers: Vec<Layer>,
    input_size: usize,
}

impl MLP {
    fn new(sizes: &[usize], activations: &[Activation]) -> TuphiumResult<Self> {
        if sizes.len().saturating_sub(1) != activations.len() {
            return Err(TuphiumError::LengthMismatch {
                operation: "MLP::new",
                expected: sizes.len().saturating_sub(1),
                actual: activations.len(),
            });
        }
        let layers: Vec<Layer> = sizes
            .windows(2)
            .zip(activations)
            .map(|(s, a)| Layer::new(s[0], s[1], a.clone()))
            .collect::<TuphiumResult<_>>()?;
        let input_size = sizes[0];
        Ok(Self { layers, input_size })
    }

    fn save(&self, path: &str) -> Result<(), Box<dyn std::error::Error>> {
        let json = serde_json::to_string_pretty(self)?;
        let mut file = File::create(path)?;
        file.write_all(json.as_bytes())?;
        Ok(())
    }

    #[allow(dead_code)]
    fn load(path: &str) -> Result<Self, Box<dyn std::error::Error>> {
        let mut file = File::open(path)?;
        let mut contents = String::new();
        file.read_to_string(&mut contents)?;
        let mut model: MLP = serde_json::from_str(&contents)?;
        for layer in model.layers.iter_mut() {
            layer.rebuild_buffers();
        }
        Ok(model)
    }
}

impl Module for MLP {
    fn forward(&mut self, x: &[f32]) -> TuphiumResult<Vec<f32>> {
        if x.len() != self.input_size {
            return Err(TuphiumError::ShapeMismatch {
                operation: "MLP::forward",
                expected: vec![self.input_size],
                actual: vec![x.len()],
            });
        }
        let mut h = x.to_vec();
        for layer in self.layers.iter_mut() {
            h = layer.forward(&h)?;
        }
        Ok(h)
    }

    fn backward(&mut self, grad_output: &[f32]) -> TuphiumResult<Vec<f32>> {
        let mut grad = grad_output.to_vec();
        for layer in self.layers.iter_mut().rev() {
            grad = layer.backward(&grad)?;
        }
        Ok(grad)
    }

    fn params(&mut self) -> Vec<&mut f32> {
        self.layers.iter_mut().flat_map(|l| l.params()).collect()
    }

    fn grads(&self) -> Vec<f32> {
        self.layers.iter().flat_map(|l| l.grads()).collect()
    }

    fn zero_grad(&mut self) {
        self.layers.iter_mut().for_each(|l| l.zero_grad());
    }

    fn step(&mut self, opt: &mut dyn Optimizer) -> TuphiumResult<()> {
        for layer in self.layers.iter_mut() {
            layer.step(opt)?;
        }
        Ok(())
    }
}

// Normalize features to [0, 1] range
fn normalize_features(df: &DataFrame) -> TuphiumResult<DataFrame> {
    let mut out_columns: Vec<Column> = Vec::with_capacity(df.shape().1);
    for name in df.column_names() {
        let col = match df.as_f64(name) {
            Ok(values) => {
                let mut min = f64::MAX;
                let mut max = f64::MIN;
                for &v in &values {
                    min = min.min(v);
                    max = max.max(v);
                }
                let range = max - min;
                if range > 1e-8 {
                    Column::Float(values.iter().map(|v| (v - min) / range).collect())
                } else {
                    Column::Float(values)
                }
            }
            Err(_) => df.column(name).unwrap().clone(),
        };
        out_columns.push(col);
    }
    DataFrame::from_columns(df.column_names().to_vec(), out_columns)
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    const EPOCHS: usize = 100;
    const LR: f32 = 0.1;

    // 1. Load breast cancer dataset (569 samples, 22 features + diagnosis)
    let csv_data = csv::read_csv("data/breast_cancer.csv")?;
    let df = DataFrame::from_csv_data(csv_data)?;

    // 2. Separate features and labels
    let mut features_df = df.drop("diagnosis")?;
    features_df = normalize_features(&features_df)?;
    let labels_df = df.select(&["diagnosis"])?;

    // 3. Create and split datasets
    let mut dataset = Dataset::new(features_df, labels_df)?;
    dataset.shuffle();
    let (train, val, test) = dataset.split(0.7, 0.1, 0.2)?;

    let n_features = train.x.shape().1;

    println!("[!] Training samples: {}", train.len());
    println!("[!] Validation samples: {}", val.len());
    println!("[!] Test samples: {}", test.len());
    println!("[!] Features: {}", n_features);
    println!("[!] Backend: {:?}\n", tuphium::backend::BACKEND);

    // Convert all splits to numeric matrices
    let (train_x, _) = train.into_numeric()?;
    let train_labels = train.y.as_f64("diagnosis")?;
    let (val_x, _) = val.into_numeric()?;
    let val_labels = val.y.as_f64("diagnosis")?;
    let (test_x, _) = test.into_numeric()?;
    let test_labels = test.y.as_f64("diagnosis")?;

    // 4. Build MLP (binary classification)
    let mut model = MLP::new(
        &[n_features, 16, 8, 1],
        &[Activation::ReLU, Activation::ReLU, Activation::Sigmoid],
    )?;
    let mut optimizer = SGD { lr: LR };

    // 5. Loss function
    let bce = BinaryCrossEntropy;

    // 6. Training loop
    for epoch in 1..EPOCHS + 1 {
        let mut train_loss = 0.0;

        for i in 0..train_x.rows {
            let x = train_x.row(i)?;
            let y_true = train_labels[i] as f32;

            let y_hat = model.forward(x)?;
            let loss = bce.forward(&[y_true], &y_hat)?;
            train_loss += loss;

            let grad = bce.backward(&[y_true], &y_hat)?;
            model.backward(&grad)?;
            model.step(&mut optimizer as &mut dyn Optimizer)?;
            model.zero_grad();
        }

        train_loss /= train_x.rows as f32;

        // Validation set
        let mut val_correct = 0;
        for i in 0..val_x.rows {
            let x = val_x.row(i)?;
            let y_true = val_labels[i] as f32;
            let y_hat = model.forward(x)?;
            let pred = if y_hat[0] >= 0.5 { 1.0 } else { 0.0 };
            if pred == y_true {
                val_correct += 1;
            }
        }
        let val_acc = val_correct as f32 / val_x.rows as f32;

        println!(
            "Epoch {}/{} | train_loss = {:.4} | val_acc = {:.3}",
            epoch, EPOCHS, train_loss, val_acc,
        );
    }

    // 7. Evaluate on test set
    let mut test_correct = 0;
    for i in 0..test_x.rows {
        let x = test_x.row(i)?;
        let y_true = test_labels[i] as f32;
        let y_hat = model.forward(x)?;
        let pred = if y_hat[0] >= 0.5 { 1.0 } else { 0.0 };
        if pred == y_true {
            test_correct += 1;
        }
    }
    let test_acc = test_correct as f32 / test_x.rows as f32;
    println!(
        "\nTest accuracy: {:.3} | (correct {} of {})",
        test_acc, test_correct, test_x.rows
    );

    // 8. Save trained model
    fs::create_dir_all("exports")?;
    model.save("exports/mlp_model.json")?;
    println!("\nModel saved to exports/mlp_model.json");

    Ok(())
}
