srsadmm-core 0.1.2

use clap::Parser;
use futures::future;
use nalgebra as na;
use rand::prelude::*;
use rand_distr::{Normal, StandardNormal};
use rayon::prelude::*;
use srsadmm_core::ops::mm;
use srsadmm_core::utils::{LassoError, fused_lasso_factor};
use srsadmm_core::variable::{DataMatrixVariable, MatrixStorageType};
use srsadmm_core::{
    problem::{ADMMConfig, ADMMContext},
    problem::{ADMMProblem, ADMMSolver},
    resource::{LocalConfig, MemoryConfig, ResourceLocation, S3Config, StorageConfig},
    variable::MatrixVariable,
};
use std::error::Error;
use std::fs::File;
use std::io::Write;
use std::path::Path;
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::{Mutex, RwLock};
use tokio::task::JoinHandle;

use tokio::sync::Semaphore;

/// Program to solve a LASSO problem using proximal gradient descent.
/// You can either provide a a_file and b_file in the format generated by generate_problem.rs,
/// or generate a problem instance on-the-fly with n, m, and k.
#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
pub struct Args {
    /// The number of threads to use for parallelization
    #[arg(short, long, default_value_t = 8)]
    threads: usize,

    /// The prefix for the output files
    #[arg(short, long, default_value_t = ("lasso-pgd").to_string())]
    prefix: String,

    /// The number of iterations to run
    #[arg(short, long, default_value_t = 200)]
    iterations: usize,

    /// File path to preprocessed data matrix A (see generate_problem.rs)
    #[arg(short, long)]
    a_file: Option<String>,

    /// File path to preprocessed data matrix b (see generate_problem.rs)
    #[arg(short, long)]
    b_file: Option<String>,

    /// The number of features to use if generating a problem instance
    #[arg(short, long, default_value_t = 10_000)]
    n: usize,

    /// The number of samples to use if generating a problem instance
    #[arg(short, long, default_value_t = 400_000)]
    m: usize,

    /// The number of non-zero elements to use if generating a problem instance
    #[arg(short, long, default_value_t = 1000)]
    k: usize,

    /// PGD initial step size parameter
    #[arg(short, long, default_value_t = 0.25)]
    rho: f32,

    /// PGD absolute tolerance
    #[arg(short, long, default_value_t = 1e-6)]
    abs_eps: f32,

    /// L1 regularization parameter
    #[arg(short, long, default_value_t = 0.32)]
    lambda: f32,

    /// Beta parameter for backtracking line search
    #[arg(short, long, default_value_t = 0.5)]
    beta: f32,

    /// S3 bucket to use for storage
    #[arg(short, long, default_value_t = ("srsadmm").to_string())]
    s3_bucket: String,
}

// Context for LASSO problem
pub struct PGDContext {
    pub a: DataMatrixVariable,
    pub b: DataMatrixVariable,
    pub primal_residuals: Vec<f32>,
    pub dual_residuals: Vec<f32>,
    pub x_norm: f32,
    pub args: Args, // Add args to the context
    pub total_iters: usize,
}

// Subproblem state
#[derive(Clone)]
pub struct PGDSubproblem {
    // Will contain all the matrix variables for the problem
    pub ata_var: MatrixVariable,
    pub atb_var: MatrixVariable,
    pub grad_var: MatrixVariable,
    pub grad_var_local: MatrixVariable,
    pub x_var: MatrixVariable,
    pub x_old_var: MatrixVariable,
}

impl PGDContext {
    pub fn new(a: DataMatrixVariable, b: DataMatrixVariable, args: Args) -> Self {
        PGDContext {
            a,
            b,
            primal_residuals: Vec::new(),
            dual_residuals: Vec::new(),
            x_norm: 0.0,
            args,
            total_iters: 0,
        }
    }

    fn n_subproblems(&self) -> usize {
        1
    }

    fn max_iterations(&self) -> usize {
        self.args.iterations
    }

    fn absolute_tolerance(&self) -> f32 {
        self.args.abs_eps
    }

    fn rho(&self) -> f32 {
        self.args.rho
    }

    fn lambda(&self) -> f32 {
        self.args.lambda
    }

    fn n_features(&self) -> usize {
        self.args.n
    }

    fn n_samples(&self) -> usize {
        self.args.m
    }

    fn n_threads(&self) -> usize {
        self.args.threads
    }

    fn beta(&self) -> f32 {
        self.args.beta
    }
}

impl PGDSubproblem {
    pub fn new(
        ata_var: MatrixVariable,
        atb_var: MatrixVariable,
        grad_var: MatrixVariable,
        grad_var_local: MatrixVariable,
        x_var: MatrixVariable,
        x_old_var: MatrixVariable,
    ) -> Self {
        PGDSubproblem {
            ata_var,
            atb_var,
            grad_var,
            grad_var_local,
            x_var,
            x_old_var,
        }
    }
}

// LASSO problem implementation
pub struct LassoProblem {
    config: ADMMConfig,
    context: Arc<ADMMContext<PGDContext, Vec<PGDSubproblem>>>,
}

impl LassoProblem {
    pub fn new(
        config: ADMMConfig,
        context: Arc<ADMMContext<PGDContext, Vec<PGDSubproblem>>>,
    ) -> Self {
        LassoProblem { config, context }
    }
}

impl ADMMProblem<PGDContext, Vec<PGDSubproblem>> for LassoProblem {
    fn context(&self) -> &ADMMContext<PGDContext, Vec<PGDSubproblem>> {
        &self.context
    }
    async fn precompute(&self) -> Result<(), LassoError> {
        let mut context_global = self.context.global.write().await;
        let config = &self.config;
        let a_matrix = context_global
            .a
            .read_chunk(0, context_global.n_samples())
            .unwrap();
        let b_matrix = context_global
            .b
            .read_chunk(0, context_global.n_samples())
            .unwrap();
        println!("[Main] determining Lipschitz constant...");
        let ata = &a_matrix.transpose() * &a_matrix;
        let l = ata.singular_values().max();
        println!("[Main] Lipschitz constant: {}", l);
        let rho = 1.0 / l;
        println!("[Main] Rho (LR): {}", rho);
        context_global.args.rho = rho;
        let ata_var = MatrixVariable::from_matrix(
            "ata".to_string(),
            ResourceLocation::S3,
            &config.storage,
            ata,
        )
        .await;

        let atb_var = MatrixVariable::from_matrix(
            "atb".to_string(),
            ResourceLocation::Local,
            &config.storage,
            a_matrix.transpose() * b_matrix,
        )
        .await;

        let grad_var = MatrixVariable::new(
            "grad".to_string(),
            ResourceLocation::S3,
            &config.storage,
        );

        let grad_var_local = MatrixVariable::new(
            "grad-local".to_string(),
            ResourceLocation::Local,
            &config.storage,
        );

        let x_var = MatrixVariable::zeros(
            "x".to_string(),
            context_global.n_features(),
            1,
            ResourceLocation::Local,
            &config.storage,
        )
        .await;

        let x_old_var = MatrixVariable::zeros(
            "x_old".to_string(),
            context_global.n_features(),
            1,
            ResourceLocation::Local,
            &config.storage,
        )
        .await;

        let subproblem = PGDSubproblem::new(
            ata_var,
            atb_var,
            grad_var,
            grad_var_local,
            x_var,
            x_old_var,
        );
        self.context.local.lock().await.push(subproblem);

        println!("[LassoProblem] Subproblem created.");

        Ok(())
    }

    async fn update_x(&self) -> Result<(), LassoError> {
        println!("[LassoProblem] Updating x...");
        let context_global = self.context.global.write().await;
        let m = context_global.n_samples();

        let rho = context_global.rho();
        let mut local = self.context.local.lock().await;

        let mut x_var = local[0].x_var.clone();
        let mut x_old_var = local[0].x_old_var.clone();

        x_old_var.write(x_var.read().await.unwrap()).await.unwrap();

        let mut ata_var = local[0].ata_var.clone();
        let atb_var = local[0].atb_var.clone();
        let mut grad_var = local[0].grad_var.clone();


        println!("[LassoProblem] Computing A^T Ax...");
        // A^TA x
        mm(&mut ata_var, &mut x_var, &mut grad_var).await.unwrap();

        println!("[LassoProblem] Computing A^T (Ax - b)...");
        // A^T(Ax - b) = A^T Ax - A^T b
        let grad = grad_var.read().await.unwrap() - atb_var.read().await.unwrap();


        println!("[LassoProblem] Computing x_k+1...");
        let x_k_plus_1 = x_var.read().await.unwrap() - rho * &grad;
        println!("[LassoProblem] Writing x_k+1...");
        x_var.write(x_k_plus_1).await.unwrap();

        Ok(())
    }

    async fn update_z(&self) -> Result<(), LassoError> {
        let context_global = self.context.global.write().await;
        let local = self.context.local.lock().await;
        let rho = context_global.rho();
        let lambda = context_global.lambda();

        let threshold_val = lambda * rho;
        let mut x_var = local[0].x_var.clone();
        let x_k = x_var
            .read()
            .await
            .map_err(|e| LassoError::from_string(format!("Failed to read x: {}", e)))?;
        let x_k_sign = x_k.map(|x| {
            if x > 0.0 {
                1.0
            } else if x < 0.0 {
                -1.0
            } else {
                0.0
            }
        });

        let x_k_plus_1 = x_k_sign
            .map_with_location(|i, j, x_s| x_s * (x_k[(i, j)].abs() - threshold_val).max(0.0));
        x_var
            .write(x_k_plus_1)
            .await
            .map_err(|e| LassoError::from_string(format!("Failed to write x: {}", e)))?;
        Ok(())
    }

    async fn update_y(&self) -> Result<(), LassoError> {
        Ok(())
    }

    async fn update_residuals(&self) -> Result<(), LassoError> {
        let mut context_global = self.context.global.write().await;
        let local = self.context.local.lock().await;
        let x_var = local[0].x_var.clone();
        let x_old_var = local[0].x_old_var.clone();
        let x_k = x_old_var
            .read()
            .await
            .map_err(|e| LassoError::from_string(format!("Failed to read x_old: {}", e)))?;
        let x_k_plus_1 = x_var
            .read()
            .await
            .map_err(|e| LassoError::from_string(format!("Failed to read x: {}", e)))?;
        let primal_residual = (&x_k_plus_1 - &x_k).norm();
        context_global.primal_residuals.push(primal_residual);
        Ok(())
    }

    async fn check_stopping_criteria(&self) -> Result<bool, LassoError> {
        let context = self.context.global.read().await;

        let primal_residual = context.primal_residuals.last().ok_or_else(|| {
            LassoError::from_string("Primal residuals vector is empty".to_string())
        })?;

        let abs_eps = context.absolute_tolerance();


        let eps_primal = abs_eps;
        println!("[LassoProblem] Primal residual: {:.4e}", primal_residual);
        let should_stop = *primal_residual <= eps_primal;
        if should_stop {
            println!(
                "[LassoProblem] CONVERGED! Primal_res: {:.4e} <= {:.4e}",
                primal_residual, eps_primal
            );
        }

        Ok(should_stop)
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let main_start_time = Instant::now();
    println!("[Main] LASSO PGD starting up...");
    let mut rng = rand::rngs::SmallRng::seed_from_u64(0);

    let mut args = Args::parse();

    let prefix = args.prefix.clone();
    let a_file = args.a_file.clone();
    let b_file = args.b_file.clone();
    let n = args.n;
    let m = args.m;
    let k = args.k;

    // Set up storage configuration
    let storage_config = StorageConfig {
        local: LocalConfig {
            root: "./tmp".to_string(),
            prefix: prefix.clone(),
        },
        s3: S3Config {
            bucket: args.s3_bucket.clone(),
            // This actually doesn't matter, the AWS client uses env vars to determine the region
            region: "us-west-2".to_string(),
            prefix: prefix.clone(),
        },
        memory: MemoryConfig::new(),
    };

    let (a, b) = if a_file.is_some() && b_file.is_some() {
        let a = DataMatrixVariable::from_problem_file(
            'a'.to_string(),
            Path::new(&a_file.unwrap()),
            &storage_config,
        );
        let b = DataMatrixVariable::from_problem_file(
            'b'.to_string(),
            Path::new(&b_file.unwrap()),
            &storage_config,
        );
        (a, b)
    } else {
        println!("[Main] Starting data generation...");
        let data_gen_start_time = Instant::now();

        // Optimized matrix generation using parallel column generation
        println!("[Main] Generating matrix A using parallel processing...");
        let a_gen_start = Instant::now();

        // Pre-allocate the matrix
        let mut a_full_nalgebra = na::DMatrix::<f32>::zeros(m, n);

        let mut rng_seeds = Vec::with_capacity(n);
        for i in 0..n {
            rng_seeds.push(42 + i as u64);
        }

        // Use parallel iterator to generate columns
        let columns: Vec<na::DVector<f32>> = (0..n)
            .into_par_iter()
            .map(|i| {
                // Each thread gets its own RNG to avoid contention
                let mut thread_rng = rand::rngs::SmallRng::seed_from_u64(rng_seeds[i]);

                // Generate the entire column at once
                let mut column = na::DVector::<f32>::zeros(m);
                for i in 0..m {
                    column[i] = thread_rng.sample(StandardNormal);
                }

                // Normalize the column to unit length
                let norm = column.norm();
                if norm > 0.0 {
                    column /= norm;
                }
                column
            })
            .collect();

        // Copy columns into the matrix (this is fast, sequential memory access)
        for (j, column) in columns.into_iter().enumerate() {
            a_full_nalgebra.set_column(j, &column);
        }

        println!(
            "[Main] Matrix A generated in {:?} (improved from element-by-element)",
            a_gen_start.elapsed()
        );

        // Optimized x_true generation
        let x_true_start = Instant::now();
        let mut x_true_nalgebra = na::DVector::<f32>::zeros(n);
        let indices: Vec<usize> = rand::seq::index::sample(&mut rng, n, k).into_vec();

        // Vectorized assignment for selected indices
        let random_values: Vec<f32> = (0..k).map(|_| rng.sample(StandardNormal)).collect();

        for (&idx, &val) in indices.iter().zip(random_values.iter()) {
            x_true_nalgebra[idx] = val;
        }
        println!("[Main] x_true generated in {:?}", x_true_start.elapsed());

        // Optimized noise generation
        let noise_start = Instant::now();
        let noise_dist = Normal::new(0.0, 0.03162)
            .map_err(|e| LassoError::from_string(format!("Failed to create normal dist: {}", e)))?;

        // Generate all noise values at once
        let noise_values: Vec<f32> = (0..m).map(|_| rng.sample(noise_dist)).collect();
        let v_nalgebra = na::DVector::<f32>::from_vec(noise_values);

        println!(
            "[Main] Noise vector generated in {:?}",
            noise_start.elapsed()
        );

        // Matrix-vector operations (these are already optimized in nalgebra)
        let b_computation_start = Instant::now();
        let b_full_nalgebra_vec = &a_full_nalgebra * &x_true_nalgebra + v_nalgebra;
        let b_full_nalgebra = na::DMatrix::from_column_slice(m, 1, b_full_nalgebra_vec.as_slice());
        println!(
            "[Main] b = Ax + v computed in {:?}",
            b_computation_start.elapsed()
        );

        // Uncomment to print lambda_max
        // println!(
        //     "[Main] Data generation complete in {:?}. Lambda_max: {}",
        //     data_gen_start_time.elapsed(),
        //     (&a_full_nalgebra.transpose() * &b_full_nalgebra)
        //         .abs()
        //         .max()
        // );

        println!(
            "[Main] Data generation complete in {:?}",
            data_gen_start_time.elapsed()
        );

        println!("[Main] Writing A and b to DataMatrixVariables...");
        let a = DataMatrixVariable::from_matrix(
            'a'.to_string(),
            a_full_nalgebra,
            MatrixStorageType::Rows,
            &storage_config,
        );
        println!("[Main] A written to DataMatrixVariable.");
        let b = DataMatrixVariable::from_matrix(
            'b'.to_string(),
            b_full_nalgebra,
            MatrixStorageType::Rows,
            &storage_config,
        );
        println!("[Main] B written to DataMatrixVariable.");
        (a, b)
    };

    args.n = a.ncols();
    args.m = a.nrows();

    println!("[Main] n: {}, m: {}", args.n, args.m);

    args.rho = 0.25;
    let config = ADMMConfig {
        storage: storage_config.clone(),
    };

    // Create context
    let context = Arc::new(ADMMContext {
        config: config.clone(),
        global: Arc::new(RwLock::new(PGDContext::new(a, b, args))),
        local: Arc::new(Mutex::new(Vec::new())),
    });
    println!("[Main] ADMMContext created.");

    // Create problem and solver
    let problem = LassoProblem::new(config, context.clone());
    let max_iterations = context.global.read().await.max_iterations();
    let mut solver = ADMMSolver::new(problem, max_iterations);
    println!("[Main] LassoProblem and ADMMSolver created.");

    // Solve
    println!("[Main] Starting ADMM (but it's actually PGD) solver...");
    let solver_start_time = Instant::now();
    solver.solve().await?;
    println!(
        "[Main] PGD solver finished in {:?}.",
        solver_start_time.elapsed()
    );

    let context_global = context.global.read().await;
    println!("[Main] Total iterations: {}", context_global.total_iters);

    let z = context.local.lock().await[0].x_var.read().await.unwrap();
    let root = context.config.storage.local.root.clone();
    let prefix = context.config.storage.local.prefix.clone();

    let mut file = File::create(format!("{}/{}-final-z.csv", root, prefix)).unwrap();
    for i in 0..z.ncols() {
        writeln!(
            file,
            "{}",
            z.column(i)
                .to_owned()
                .iter()
                .map(|x| x.to_string())
                .collect::<Vec<String>>()
                .join(",")
        )
        .unwrap();
    }
    file.flush().unwrap();

    println!(
        "[Main] lasso_prox.rs finished. Final z is stored in {}",
        format!("{}/{}-final-z.csv", root, prefix)
    );
    let _ = solver.export_step_timings(format!("{}/{}-timings.csv", root, prefix).as_str());
    println!("[Main] Step timings exported to {}/{}-timings.csv", root, prefix);
    Ok(())
}