lib.rs

//! blaze-core — Phase 2 Rust core for TT/MPS compression.
//!
//! Mirrors the validated Python API in python/blaze/tt.py and diagnostics.py **1:1**.
//! CPU-only (faer for SVD). No CUDA. No "general compressor" claims anywhere.
//!
//! Primary gate: bitwise-identical truncation decisions + rel_error within fp tol
//! for the same inputs as the Python reference (TT-SVD with energy budget eps).

#![allow(clippy::needless_range_loop)]

use ndarray::{s, Array3, ArrayD};
use num_complex::Complex64;
use std::fs::File;
use std::io::{self, BufReader, BufWriter, Read, Write};
use std::path::Path;
use byteorder::{LittleEndian as LE, ReadBytesExt, WriteBytesExt};

/// Supported element types for TT (match Python Phase-1 dtypes policy).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DType {
    F64,
    C64,
}

impl DType {
    pub fn elem_size(&self) -> usize {
        match self {
            DType::F64 => 8,
            DType::C64 => 16,
        }
    }
}

/// The Tensor Train (TT / MPS) representation.
///
/// Cores are stored as `Array3<T>` with shape (r_left, d_phys, r_right).
/// Invariant: cores[0].shape()[0] == 1, cores.last().unwrap().shape()[2] == 1.
///
/// This struct + its methods are the **exact** semantic mirror of Python `blaze.TT`.
#[derive(Clone, Debug)]
pub struct TT<T> {
    pub cores: Vec<Array3<T>>,
    pub shape: Vec<usize>,
    /// Full singular values (real, even for complex data) from each of the (ndim-1) unfoldings.
    /// Recorded **before** truncation. This is the source of truth for diagnostics
    /// (effective ranks, decay). Mirrors Python `singular_values`.
    pub singular_values: Vec<Vec<f64>>,
}

impl<T> TT<T> {
    /// Number of physical dimensions.
    pub fn ndim(&self) -> usize {
        self.shape.len()
    }

    /// Bond dimensions [1, r0, r1, ..., r_{n-2}, 1].
    pub fn ranks(&self) -> Vec<usize> {
        let mut rs = vec![1usize];
        for c in &self.cores {
            rs.push(c.shape()[2]);
        }
        rs
    }

    /// Total stored parameters (scalars, not bytes).
    pub fn nparams(&self) -> usize {
        self.cores.iter().map(|c| c.len()).sum()
    }
}

// We will provide concrete impls for f64 and Complex64 to avoid generic hell in contraction.
impl TT<f64> {
    pub fn reconstruct(&self) -> ArrayD<f64> {
        self.reconstruct_f64()
    }

    fn reconstruct_f64(&self) -> ArrayD<f64> {
        if self.cores.is_empty() {
            return ArrayD::zeros(self.shape.clone());
        }
        if self.ndim() == 1 {
            return self.cores[0].slice(s![0, .., 0]).to_owned().into_dyn();
        }
        // Reliable left-to-right using ndarray 2D dot (correctness > speed for reconstruct in Phase 2)
        let mut res = self.cores[0].slice(s![0, .., ..]).to_owned().into_dyn(); // (d0, r1)
        for core in self.cores.iter().skip(1) {
            let r_prev = core.shape()[0];
            let prod = res.len() / r_prev;
            let res2 = res.into_shape_with_order((prod, r_prev)).unwrap();
            let d = core.shape()[1];
            let rn = core.shape()[2];
            let core2 = core.view().into_shape_with_order((r_prev, d * rn)).unwrap();
            let tmp = res2.dot(&core2); // (prod, d*rn)
            res = tmp.into_shape_with_order((prod, d, rn)).unwrap().into_dyn();
            res = res.into_shape_with_order((prod * d, rn)).unwrap().into_dyn();
        }
        let flat = res.slice(s![.., 0]).to_owned();
        flat.into_shape_with_order(self.shape.clone()).unwrap().into_dyn()
    }
}

impl TT<Complex64> {
    pub fn reconstruct(&self) -> ArrayD<Complex64> {
        self.reconstruct_c64()
    }

    fn reconstruct_c64(&self) -> ArrayD<Complex64> {
        if self.cores.is_empty() {
            return ArrayD::zeros(self.shape.clone());
        }
        if self.ndim() == 1 {
            return self.cores[0].slice(s![0, .., 0]).to_owned().into_dyn();
        }
        let mut res = self.cores[0].slice(s![0, .., ..]).to_owned().into_dyn();
        for core in self.cores.iter().skip(1) {
            let r_prev = core.shape()[0];
            let prod = res.len() / r_prev;
            let res2 = res.into_shape_with_order((prod, r_prev)).unwrap();
            let d = core.shape()[1];
            let rn = core.shape()[2];
            let core2 = core.view().into_shape_with_order((r_prev, d * rn)).unwrap();
            let tmp = res2.dot(&core2);
            res = tmp.into_shape_with_order((prod, d, rn)).unwrap().into_dyn();
            res = res.into_shape_with_order((prod * d, rn)).unwrap().into_dyn();
        }
        let flat = res.slice(s![.., 0]).to_owned();
        flat.into_shape_with_order(self.shape.clone()).unwrap().into_dyn()
    }
}

impl<T> TT<T> {
    /// Relative Frobenius error ||orig - recon|| / ||orig|| .
    /// Mirrors Python TT.rel_error exactly (use the free functions rel_error_f64 / rel_error_c64).
    pub fn rel_error(&self, _original: &ArrayD<T>) -> f64 {
        0.0 // use the specialized free functions
    }

    /// Byte compression ratio (original.nbytes / TT bytes).
    /// Honest storage metric.
    pub fn compression_ratio(&self, original_nbytes: usize) -> f64 {
        let tt_bytes = self.nparams() * std::mem::size_of::<T>(); // approximate, T is the elem
        if tt_bytes == 0 {
            f64::INFINITY
        } else {
            original_nbytes as f64 / tt_bytes as f64
        }
    }

    /// Effective ranks per bond for given tol (1 - tol^2 energy).
    /// Exact mirror of Python logic (searchsorted on cumsum of s^2).
    pub fn effective_ranks(&self, tol: f64) -> Vec<usize> {
        let mut eff = vec![];
        for s in &self.singular_values {
            if s.is_empty() {
                eff.push(0);
                continue;
            }
            let s2: Vec<f64> = s.iter().map(|&x| x * x).collect();
            let total: f64 = s2.iter().sum();
            if total <= 0.0 {
                eff.push(1);
                continue;
            }
            let mut cum = 0.0;
            let mut idx = 0;
            let target = 1.0 - tol * tol;
            for (i, &val) in s2.iter().enumerate() {
                cum += val / total;
                if cum >= target {
                    idx = i;
                    break;
                }
                idx = i;
            }
            let r = std::cmp::min(idx + 1, s.len());
            eff.push(std::cmp::max(r, 1));
        }
        eff
    }

    /// Decay summary for several tols. Mirrors Python dict output (as Vec for Rust).
    pub fn singular_decay_summary(&self, tols: &[f64]) -> Vec<Vec<(f64, usize)>> {
        let mut out = vec![];
        for s in &self.singular_values {
            if s.is_empty() {
                out.push(vec![]);
                continue;
            }
            let s2: Vec<f64> = s.iter().map(|&x| x*x).collect();
            let total: f64 = s2.iter().sum();
            let mut per = vec![];
            if total <= 0.0 {
                for &t in tols {
                    per.push((t, 1));
                }
                out.push(per);
                continue;
            }
            for &t in tols {
                let target = 1.0 - t * t;
                let mut cum = 0.0;
                let mut idx = 0usize;
                for (i, &v) in s2.iter().enumerate() {
                    cum += v / total;
                    if cum >= target {
                        idx = i;
                        break;
                    }
                    idx = i;
                }
                per.push((t, std::cmp::min(idx + 1, s.len())));
            }
            out.push(per);
        }
        out
    }
}

// The main algorithm entry points.

/// SVD backend signature: 2D matrix -> (U [m×k], S [k], Vh [k×n]) economy.
type SvdF64 = fn(&ArrayD<f64>) -> (ArrayD<f64>, Vec<f64>, ArrayD<f64>);
type SvdC64 = fn(&ArrayD<Complex64>) -> (ArrayD<Complex64>, Vec<f64>, ArrayD<Complex64>);

#[cfg(feature = "cuda")]
extern "C" {
    fn blaze_cuda_svd_f64(a: *const f64, m: i32, n: i32, u: *mut f64, s: *mut f64, vt: *mut f64) -> i32;
}

/// GPU economy SVD via cuSOLVER (drop-in for `svd_via_na_f64`, identical shapes).
#[cfg(feature = "cuda")]
fn svd_via_cuda_f64(mat: &ArrayD<f64>) -> (ArrayD<f64>, Vec<f64>, ArrayD<f64>) {
    let m = mat.shape()[0];
    let n = mat.shape()[1];
    let k = m.min(n);
    let a: Vec<f64> = mat.iter().copied().collect(); // row-major (the C-ABI contract)
    let mut u = vec![0.0f64; m * k];
    let mut s = vec![0.0f64; k];
    let mut vt = vec![0.0f64; k * n];
    let rc = unsafe {
        blaze_cuda_svd_f64(a.as_ptr(), m as i32, n as i32, u.as_mut_ptr(), s.as_mut_ptr(), vt.as_mut_ptr())
    };
    assert_eq!(rc, 0, "blaze_cuda_svd_f64 failed (rc={rc})");
    let u_nd = ArrayD::from_shape_vec(vec![m, k], u).unwrap();
    let vt_nd = ArrayD::from_shape_vec(vec![k, n], vt).unwrap();
    (u_nd, s, vt_nd)
}

#[cfg(feature = "cuda")]
extern "C" {
    fn blaze_cuda_svd_c64(a: *const Complex64, m: i32, n: i32, u: *mut Complex64, s: *mut f64, vt: *mut Complex64) -> i32;
}

/// GPU economy SVD via cuSOLVER (Complex128 — the quantum-state path). Drop-in for svd_via_na_c64.
#[cfg(feature = "cuda")]
fn svd_via_cuda_c64(mat: &ArrayD<Complex64>) -> (ArrayD<Complex64>, Vec<f64>, ArrayD<Complex64>) {
    let m = mat.shape()[0];
    let n = mat.shape()[1];
    let k = m.min(n);
    let a: Vec<Complex64> = mat.iter().copied().collect(); // row-major (the C-ABI contract)
    let mut u = vec![Complex64::new(0.0, 0.0); m * k];
    let mut s = vec![0.0f64; k];
    let mut vt = vec![Complex64::new(0.0, 0.0); k * n];
    let rc = unsafe {
        blaze_cuda_svd_c64(a.as_ptr(), m as i32, n as i32, u.as_mut_ptr(), s.as_mut_ptr(), vt.as_mut_ptr())
    };
    assert_eq!(rc, 0, "blaze_cuda_svd_c64 failed (rc={rc})");
    let u_nd = ArrayD::from_shape_vec(vec![m, k], u).unwrap();
    let vt_nd = ArrayD::from_shape_vec(vec![k, n], vt).unwrap();
    (u_nd, s, vt_nd)
}

/// Exact port of Python tt_svd logic (energy budget + max_rank cap).
/// Returns (cores, full_singular_values_per_unfolding).
///
/// For f64 and Complex64.
pub fn tt_svd_f64(
    tensor: &ArrayD<f64>,
    max_rank: Option<usize>,
    rel_tol: f64,
    verbose: bool,
    svd: SvdF64,
) -> (Vec<Array3<f64>>, Vec<Vec<f64>>) {
    tt_svd_impl_f64(tensor, max_rank, rel_tol, verbose, svd)
}

pub fn tt_svd_c64(
    tensor: &ArrayD<Complex64>,
    max_rank: Option<usize>,
    rel_tol: f64,
    verbose: bool,
    svd: SvdC64,
) -> (Vec<Array3<Complex64>>, Vec<Vec<f64>>) {
    tt_svd_impl_c64(tensor, max_rank, rel_tol, verbose, svd)
}

fn tt_svd_impl_f64(
    tensor: &ArrayD<f64>,
    max_rank: Option<usize>,
    rel_tol: f64,
    verbose: bool,
    svd: SvdF64,
) -> (Vec<Array3<f64>>, Vec<Vec<f64>>) {
    let shape: Vec<usize> = tensor.shape().to_vec();
    let ndim = shape.len();
    if ndim == 1 {
        let d0 = shape[0];
        let mut core = Array3::<f64>::zeros((1, d0, 1));
        for i in 0..d0 { core[[0, i, 0]] = tensor[[i]]; }
        return (vec![core], vec![]);
    }

    let mut cores: Vec<Array3<f64>> = vec![];
    let mut singular_values: Vec<Vec<f64>> = vec![];
    let mut remaining = tensor.clone();
    let mut r_prev = 1usize;
    let eps = rel_tol / (ndim as f64 - 1.0).sqrt().max(1.0);

    for i in 0..(ndim - 1) {
        let d_i = shape[i];
        let mat_shape = if i == 0 { (d_i, remaining.len() / d_i) } else { (r_prev * d_i, remaining.len() / (r_prev * d_i)) };
        let mat_2d = remaining.clone().into_shape_with_order(mat_shape).unwrap().into_dyn();
        let (u, s, vh) = svd(&mat_2d);
        singular_values.push(s.clone());

        let s2: Vec<f64> = s.iter().map(|&x| x*x).collect();
        let total_s2: f64 = s2.iter().sum();
        let r_tol = if total_s2 > 0.0 && eps > 0.0 {
            let mut cum = 0.0;
            let target = 1.0 - eps*eps;
            let mut idx = 0;
            for (j, &v) in s2.iter().enumerate() {
                cum += v / total_s2;
                if cum >= target { idx = j; break; }
                idx = j;
            }
            (idx + 1).min(s.len())
        } else { s.len() };

        let mut r = r_tol;
        if let Some(mr) = max_rank { r = r.min(mr); }
        r = r.max(1).min(s.len());

        if verbose {
            let kept = if total_s2 > 0.0 { s2[..r].iter().sum::<f64>() / total_s2 } else { 1.0 };
            eprintln!("[tt_svd] mode {i}: rank {r}/{} kept_energy={:.6e} max_sigma={:.3e}", s.len(), kept, s[0]);
        }

        let u_trunc = u.slice(s![.., ..r]).to_owned().into_dyn();
        let s_trunc = &s[0..r];
        let vh_trunc = vh.slice(s![..r, ..]).to_owned().into_dyn();

        let core = if i == 0 {
            u_trunc.into_shape_with_order((1, d_i, r)).unwrap()
        } else {
            u_trunc.into_shape_with_order((r_prev, d_i, r)).unwrap()
        };
        cores.push(core);

        remaining = scale_and_reshape_vh_f64(s_trunc, &vh_trunc, &shape[i+1..], r);
        r_prev = r;
    }

    let last = remaining.into_shape_with_order((r_prev, shape[ndim-1], 1)).unwrap();
    cores.push(last);
    (cores, singular_values)
}

fn tt_svd_impl_c64(
    tensor: &ArrayD<Complex64>,
    max_rank: Option<usize>,
    rel_tol: f64,
    verbose: bool,
    svd: SvdC64,
) -> (Vec<Array3<Complex64>>, Vec<Vec<f64>>) {
    let shape: Vec<usize> = tensor.shape().to_vec();
    let ndim = shape.len();
    if ndim == 1 {
        let d0 = shape[0];
        let mut core = Array3::<Complex64>::zeros((1, d0, 1));
        for i in 0..d0 { core[[0, i, 0]] = tensor[[i]]; }
        return (vec![core], vec![]);
    }

    let mut cores: Vec<Array3<Complex64>> = vec![];
    let mut singular_values: Vec<Vec<f64>> = vec![];
    let mut remaining = tensor.clone();
    let mut r_prev = 1usize;
    let eps = rel_tol / (ndim as f64 - 1.0).sqrt().max(1.0);

    for i in 0..(ndim - 1) {
        let d_i = shape[i];
        let mat_shape = if i == 0 { (d_i, remaining.len() / d_i) } else { (r_prev * d_i, remaining.len() / (r_prev * d_i)) };
        let mat_2d = remaining.clone().into_shape_with_order(mat_shape).unwrap().into_dyn();
        let (u, s, vh) = svd(&mat_2d);
        singular_values.push(s.clone());

        let s2: Vec<f64> = s.iter().map(|&x| x*x).collect();
        let total_s2: f64 = s2.iter().sum();
        let r_tol = if total_s2 > 0.0 && eps > 0.0 {
            let mut cum = 0.0; let target = 1.0 - eps*eps; let mut idx = 0;
            for (j, &v) in s2.iter().enumerate() { cum += v/total_s2; if cum >= target { idx = j; break; } idx = j; }
            (idx + 1).min(s.len())
        } else { s.len() };

        let mut r = r_tol; if let Some(mr) = max_rank { r = r.min(mr); } r = r.max(1).min(s.len());

        if verbose {
            let kept = if total_s2 > 0.0 { s2[..r].iter().sum::<f64>() / total_s2 } else { 1.0 };
            eprintln!("[tt_svd] mode {i}: rank {r}/{} kept_energy={:.6e} max_sigma={:.3e}", s.len(), kept, s[0]);
        }

        let u_trunc = u.slice(s![.., ..r]).to_owned().into_dyn();
        let s_trunc = &s[0..r];
        let vh_trunc = vh.slice(s![..r, ..]).to_owned().into_dyn();

        let core = if i == 0 {
            u_trunc.into_shape_with_order((1, d_i, r)).unwrap()
        } else {
            u_trunc.into_shape_with_order((r_prev, d_i, r)).unwrap()
        };
        cores.push(core);

        remaining = scale_and_reshape_vh_c64(s_trunc, &vh_trunc, &shape[i+1..], r);
        r_prev = r;
    }

    let last = remaining.into_shape_with_order((r_prev, shape[ndim-1], 1)).unwrap();
    cores.push(last);
    (cores, singular_values)
}

// Concrete SVD using nalgebra (reliable, matches scipy truncation points for our purpose; faer for future hot path).
fn svd_via_na_f64(mat: &ArrayD<f64>) -> (ArrayD<f64>, Vec<f64>, ArrayD<f64>) {
    let rows = mat.shape()[0];
    let cols = mat.shape()[1];
    // ndarray is C (row major), nalgebra DMatrix is column major. Use from_row_slice for correct.
    let data: Vec<f64> = mat.iter().copied().collect();
    let m = nalgebra::DMatrix::from_row_slice(rows, cols, &data);
    let svd = nalgebra::SVD::new(m, true, true);
    let u = svd.u.unwrap();
    let s = svd.singular_values.iter().copied().collect::<Vec<f64>>();
    let vh = svd.v_t.unwrap(); // already V^T
    // Element-wise copy using logical indices: nalgebra column-major storage, but we access by (row,col) logical
    // and write to ndarray C-order ArrayD. This is the correct Fortran->C conversion for the matrices.
    let mut u_nd = ArrayD::<f64>::zeros(vec![u.nrows(), u.ncols()]);
    for i in 0..u.nrows() {
        for j in 0..u.ncols() {
            u_nd[[i, j]] = u[(i, j)];
        }
    }
    let mut vh_nd = ArrayD::<f64>::zeros(vec![vh.nrows(), vh.ncols()]);
    for i in 0..vh.nrows() {
        for j in 0..vh.ncols() {
            vh_nd[[i, j]] = vh[(i, j)];
        }
    }
    (u_nd, s, vh_nd)
}

fn svd_via_na_c64(mat: &ArrayD<Complex64>) -> (ArrayD<Complex64>, Vec<f64>, ArrayD<Complex64>) {
    let rows = mat.shape()[0];
    let cols = mat.shape()[1];
    let data: Vec<Complex64> = mat.iter().copied().collect();
    let m = nalgebra::DMatrix::<Complex64>::from_row_slice(rows, cols, &data);
    let svd = nalgebra::SVD::new(m, true, true);
    let u = svd.u.unwrap();
    let s = svd.singular_values.iter().copied().collect::<Vec<f64>>();
    let vh = svd.v_t.unwrap();
    // Element-wise copy using logical indices for correct layout (nalgebra col-major -> ndarray C-order)
    let mut u_nd = ArrayD::<Complex64>::zeros(vec![u.nrows(), u.ncols()]);
    for i in 0..u.nrows() {
        for j in 0..u.ncols() {
            u_nd[[i, j]] = u[(i, j)];
        }
    }
    let mut vh_nd = ArrayD::<Complex64>::zeros(vec![vh.nrows(), vh.ncols()]);
    for i in 0..vh.nrows() {
        for j in 0..vh.ncols() {
            vh_nd[[i, j]] = vh[(i, j)];
        }
    }
    (u_nd, s, vh_nd)
}

fn scale_and_reshape_vh_f64(s: &[f64], vh: &ArrayD<f64>, tail_shape: &[usize], r: usize) -> ArrayD<f64> {
    let right = tail_shape.iter().product::<usize>();
    let mut scaled = ArrayD::<f64>::zeros(vec![r, right]);
    for (j, &sj) in s.iter().enumerate() {
        for k in 0..right {
            scaled[[j, k]] = sj * vh[[j, k]];
        }
    }
    scaled.into_shape_with_order(vec![r].into_iter().chain(tail_shape.iter().copied()).collect::<Vec<_>>()).unwrap()
}

fn scale_and_reshape_vh_c64(s: &[f64], vh: &ArrayD<Complex64>, tail_shape: &[usize], r: usize) -> ArrayD<Complex64> {
    let right = tail_shape.iter().product::<usize>();
    let mut scaled = ArrayD::<Complex64>::zeros(vec![r, right]);
    for (j, &sj) in s.iter().enumerate() {
        for k in 0..right {
            scaled[[j, k]] = Complex64::new(sj, 0.0) * vh[[j, k]];
        }
    }
    scaled.into_shape_with_order(vec![r].into_iter().chain(tail_shape.iter().copied()).collect::<Vec<_>>()).unwrap()
}

// Public high-level API mirroring Python `blaze.compress`
pub fn compress_f64(
    tensor: &ArrayD<f64>,
    max_rank: Option<usize>,
    rel_tol: f64,
    verbose: bool,
) -> TT<f64> {
    let (cores, sigs) = tt_svd_f64(tensor, max_rank, rel_tol, verbose, svd_via_na_f64);
    TT { cores, shape: tensor.shape().to_vec(), singular_values: sigs }
}

/// Same as `compress_f64` but with the SVD offloaded to the GPU (cuSOLVER).
/// Only with `--features cuda`. Used to gate parity against the CPU path.
#[cfg(feature = "cuda")]
pub fn compress_f64_cuda(
    tensor: &ArrayD<f64>,
    max_rank: Option<usize>,
    rel_tol: f64,
    verbose: bool,
) -> TT<f64> {
    let (cores, sigs) = tt_svd_f64(tensor, max_rank, rel_tol, verbose, svd_via_cuda_f64);
    TT { cores, shape: tensor.shape().to_vec(), singular_values: sigs }
}

pub fn compress_c64(
    tensor: &ArrayD<Complex64>,
    max_rank: Option<usize>,
    rel_tol: f64,
    verbose: bool,
) -> TT<Complex64> {
    let (cores, sigs) = tt_svd_c64(tensor, max_rank, rel_tol, verbose, svd_via_na_c64);
    TT { cores, shape: tensor.shape().to_vec(), singular_values: sigs }
}

/// `compress_c64` with the SVD offloaded to the GPU (cuSOLVER Zgesvd) — the quantum path.
#[cfg(feature = "cuda")]
pub fn compress_c64_cuda(
    tensor: &ArrayD<Complex64>,
    max_rank: Option<usize>,
    rel_tol: f64,
    verbose: bool,
) -> TT<Complex64> {
    let (cores, sigs) = tt_svd_c64(tensor, max_rank, rel_tol, verbose, svd_via_cuda_c64);
    TT { cores, shape: tensor.shape().to_vec(), singular_values: sigs }
}

// Parity helper (for tests)
pub fn rel_error_f64(tt: &TT<f64>, original: &ArrayD<f64>) -> f64 {
    let recon = tt.reconstruct_f64();
    let diff = &recon - original;
    let num = diff.iter().map(|&x| x*x).sum::<f64>().sqrt();
    let den = original.iter().map(|&x| x*x).sum::<f64>().sqrt();
    if den == 0.0 { 0.0 } else { num / den }
}

pub fn rel_error_c64(tt: &TT<Complex64>, original: &ArrayD<Complex64>) -> f64 {
    let recon = tt.reconstruct_c64();
    let diff = &recon - original;
    let num = diff.iter().map(|&z| z.norm_sqr()).sum::<f64>().sqrt();
    let den = original.iter().map(|&z| z.norm_sqr()).sum::<f64>().sqrt();
    if den == 0.0 { 0.0 } else { num / den }
}

#[cfg(all(test, feature = "cuda"))]
mod cuda_parity {
    use super::*;

    #[test]
    fn cpu_vs_cuda_parity() {
        // Structured 6x6x6 tensor with a real (non-trivial) truncation error.
        let n = 6usize;
        let mut data = vec![0.0f64; n * n * n];
        for i in 0..n {
            for j in 0..n {
                for k in 0..n {
                    let (x, y, z) = (i as f64 / 5.0, j as f64 / 5.0, k as f64 / 5.0);
                    data[(i * n + j) * n + k] = (-(x * x + y * y + z * z)).exp() + 0.1 * (3.0 * x).sin() * y;
                }
            }
        }
        let t = ArrayD::from_shape_vec(vec![n, n, n], data).unwrap();
        let cpu = compress_f64(&t, Some(4), 1e-6, false);
        let gpu = compress_f64_cuda(&t, Some(4), 1e-6, false);
        let ecpu = rel_error_f64(&cpu, &t);
        let egpu = rel_error_f64(&gpu, &t);
        println!("CPU rel_err={ecpu:.6e} ranks={:?}", cpu.ranks());
        println!("GPU rel_err={egpu:.6e} ranks={:?}", gpu.ranks());
        assert_eq!(cpu.ranks(), gpu.ranks(), "ranks must match CPU vs GPU");
        assert!((ecpu - egpu).abs() < 1e-6, "rel_error mismatch (cpu={ecpu}, gpu={egpu})");
        println!("CUDA_PARITY_OK");
    }

    #[test]
    fn cpu_vs_cuda_parity_c64() {
        // Structured 4x4x4 complex tensor (quantum-state-like, low rank).
        let n = 4usize;
        let mut data = vec![Complex64::new(0.0, 0.0); n * n * n];
        for i in 0..n {
            for j in 0..n {
                for k in 0..n {
                    let p = (i + 2 * j + 3 * k) as f64;
                    data[(i * n + j) * n + k] =
                        Complex64::new((0.3 * p).cos(), (0.2 * p).sin()) * (-(0.1 * p)).exp();
                }
            }
        }
        let t = ArrayD::from_shape_vec(vec![n, n, n], data).unwrap();
        let cpu = compress_c64(&t, Some(3), 1e-6, false);
        let gpu = compress_c64_cuda(&t, Some(3), 1e-6, false);
        let ecpu = rel_error_c64(&cpu, &t);
        let egpu = rel_error_c64(&gpu, &t);
        println!("C64 CPU rel_err={ecpu:.6e} ranks={:?}", cpu.ranks());
        println!("C64 GPU rel_err={egpu:.6e} ranks={:?}", gpu.ranks());
        assert_eq!(cpu.ranks(), gpu.ranks(), "c64 ranks must match CPU vs GPU");
        assert!((ecpu - egpu).abs() < 1e-6, "c64 rel_error mismatch (cpu={ecpu}, gpu={egpu})");
        println!("CUDA_PARITY_C64_OK");
    }

    /// Honest CPU-vs-GPU timing across sizes. Shows the crossover: GPU loses on small
    /// tensors (host<->device transfer overhead) and wins on large ones. Run with:
    ///   cargo test -p blaze-core --features cuda cpu_vs_cuda_timing -- --ignored --nocapture
    #[test]
    #[ignore]
    fn cpu_vs_cuda_timing() {
        use std::time::Instant;
        fn make(n: usize) -> ArrayD<f64> {
            let len = n * n * n * n;
            let mut d = vec![0.0f64; len];
            for (i, v) in d.iter_mut().enumerate() {
                *v = ((i as f64) * 0.001).sin() + 0.3 * ((i as f64) * 0.017).cos();
            }
            ArrayD::from_shape_vec(vec![n, n, n, n], d).unwrap()
        }
        // Warm up the CUDA context so the first timed call is fair.
        let _ = compress_f64_cuda(&make(4), Some(4), 1e-6, false);

        println!("\n  4D tensor (n^4)   CPU(ms)    GPU(ms)   speedup   winner");
        for &n in &[8usize, 16, 24, 32] {
            let t = make(n);
            let a = Instant::now();
            let _ = compress_f64(&t, Some(16), 1e-6, false);
            let cpu_ms = a.elapsed().as_secs_f64() * 1e3;
            let b = Instant::now();
            let _ = compress_f64_cuda(&t, Some(16), 1e-6, false);
            let gpu_ms = b.elapsed().as_secs_f64() * 1e3;
            let win = if gpu_ms < cpu_ms { "GPU" } else { "CPU (transfer-bound)" };
            println!("  n={n:<2} ({:>8} el)  {cpu_ms:8.1}  {gpu_ms:8.1}   {:6.2}x   {win}",
                     t.len(), cpu_ms / gpu_ms);
        }
        println!("  (GPU wins above the crossover; below it, host<->device transfer dominates.)");
    }

    /// Honest CPU-vs-GPU timing for the COMPLEX (quantum) path. n=32 -> 32^4 = 2^20 =
    /// a 20-qubit statevector reshaped to an MPS. Complex SVD is ~4x the flops of real,
    /// so the GPU should win by more than the f64 case.
    ///   cargo test -p blaze-core --release --features cuda cpu_vs_cuda_timing_c64 -- --ignored --nocapture
    #[test]
    #[ignore]
    fn cpu_vs_cuda_timing_c64() {
        use std::time::Instant;
        fn make(n: usize) -> ArrayD<Complex64> {
            let len = n * n * n * n;
            let mut d = vec![Complex64::new(0.0, 0.0); len];
            for (i, v) in d.iter_mut().enumerate() {
                let p = i as f64;
                *v = Complex64::new((p * 0.001).sin(), (p * 0.0017).cos());
            }
            ArrayD::from_shape_vec(vec![n, n, n, n], d).unwrap()
        }
        let _ = compress_c64_cuda(&make(4), Some(4), 1e-6, false); // warm up the CUDA context

        println!("\n  c64 4D (n^4)      CPU(ms)    GPU(ms)   speedup   winner");
        for &n in &[8usize, 16, 24, 32] {
            let t = make(n);
            let a = Instant::now();
            let _ = compress_c64(&t, Some(16), 1e-6, false);
            let cpu_ms = a.elapsed().as_secs_f64() * 1e3;
            let b = Instant::now();
            let _ = compress_c64_cuda(&t, Some(16), 1e-6, false);
            let gpu_ms = b.elapsed().as_secs_f64() * 1e3;
            let win = if gpu_ms < cpu_ms { "GPU" } else { "CPU (transfer-bound)" };
            println!("  n={n:<2} ({:>8} el)  {cpu_ms:8.1}  {gpu_ms:8.1}   {:6.2}x   {win}",
                     t.len(), cpu_ms / gpu_ms);
        }
        println!("  (n=32 = 32^4 = 2^20 = a 20-qubit statevector reshaped to MPS.)");
    }
}

#[cfg(test)]
mod parity_tests {
    use super::*;

    #[test]
    fn test_f64_parity_with_python() {
        // Exact data from python generator (small separable for easy verification)
        let data_vec: Vec<f64> = vec![
0.00000000000000000e+00, 0.00000000000000000e+00, 0.00000000000000000e+00, 0.00000000000000000e+00, 0.00000000000000000e+00, 0.00000000000000000e+00, 0.00000000000000000e+00, 0.00000000000000000e+00, 0.00000000000000000e+00, 4.79425538604203005e-01, 2.90786288212691868e-01, 1.76370799225031960e-01, 4.20735492403948252e-01, 2.55188975772286386e-01, 1.54779937826556108e-01, 2.59034723999925720e-01, 1.57112502036154872e-01, 9.52935495090914825e-02, 8.41470984807896505e-01, 5.10377951544572772e-01, 3.09559875653112215e-01, 7.38460262604128781e-01, 4.47898790248846768e-01, 2.71664348734123429e-01, 4.54648713412840910e-01, 2.75758384083790387e-01, 1.67255914619631157e-01,
        ];
        let data = ArrayD::<f64>::from_shape_vec(vec![3,3,3], data_vec).unwrap();
        let tt = compress_f64(&data, Some(2), 1e-4, false);
        let err = rel_error_f64(&tt, &data);
        println!("RUST_REL={:.17e}", err);
        println!("RUST_RANKS={:?}", tt.ranks());
        // PY_REL from identical gen: 2.18393667313302425e-16
        assert!((err - 2.18393667313302425e-16).abs() < 1e-12, "rel_error must match python within fp; layout fix required for correct recon");
    }

    #[test]
    fn test_c64_ghz_fidelity() {
        // 2-qubit GHZ exact (rank 2)
        let mut data = ArrayD::<Complex64>::zeros(vec![2,2]);
        let s = 1.0 / (2.0_f64).sqrt();
        data[[0,0]] = Complex64::new(s, 0.);
        data[[1,1]] = Complex64::new(s, 0.);
        let tt = compress_c64(&data, Some(2), 1e-8, false);
        let recon = tt.reconstruct_c64();
        let mut inner = Complex64::new(0.,0.);
        for (e, r) in data.iter().zip(recon.iter()) {
            inner += e.conj() * *r;
        }
        let fid = inner.norm_sqr();
        println!("RUST_GHZ_FID={:.17e}", fid);
        println!("RUST_GHZ_RANKS={:?}", tt.ranks());
        assert!(fid > 0.999999, "GHZ fidelity must be ~1 for rank=2");
    }
}

// ===================== .blz I/O (per docs/blz-format.md) =====================

const BLZ_MAGIC: &[u8; 4] = b"BLZ1";
const BLZ_VERSION: u8 = 1;

pub fn write_blz_f64(tt: &TT<f64>, path: &Path) -> io::Result<()> {
    let file = File::create(path)?;
    let mut w = BufWriter::new(file);
    w.write_all(BLZ_MAGIC)?;
    w.write_u8(BLZ_VERSION)?;
    w.write_u8(0)?; // dtype f64
    w.write_u8(tt.ndim() as u8)?;
    w.write_all(&[0u8; 8])?; // reserved
    for &d in &tt.shape {
        w.write_u64::<LE>(d as u64)?;
    }
    for &r in &tt.ranks() {
        w.write_u64::<LE>(r as u64)?;
    }
    for core in &tt.cores {
        for &val in core.iter() {
            w.write_f64::<LE>(val)?;
        }
    }
    w.flush()?;
    Ok(())
}

pub fn read_blz_f64(path: &Path) -> io::Result<TT<f64>> {
    let file = File::open(path)?;
    let mut r = BufReader::new(file);
    let mut magic = [0u8; 4];
    r.read_exact(&mut magic)?;
    if &magic != BLZ_MAGIC {
        return Err(io::Error::new(io::ErrorKind::InvalidData, "bad magic"));
    }
    let ver = r.read_u8()?;
    if ver != BLZ_VERSION {
        return Err(io::Error::new(io::ErrorKind::InvalidData, "bad version"));
    }
    let dtype = r.read_u8()?;
    if dtype != 0 {
        return Err(io::Error::new(io::ErrorKind::InvalidData, "expected f64"));
    }
    let ndim = r.read_u8()? as usize;
    let mut reserved = [0u8; 8];
    r.read_exact(&mut reserved)?;
    let mut shape = vec![0usize; ndim];
    for s in &mut shape {
        *s = r.read_u64::<LE>()? as usize;
    }
    let mut ranks = vec![0usize; ndim + 1];
    for rk in &mut ranks {
        *rk = r.read_u64::<LE>()? as usize;
    }
    let mut cores = Vec::with_capacity(ndim);
    for i in 0..ndim {
        let rl = ranks[i];
        let d = shape[i];
        let rr = ranks[i + 1];
        let n = rl * d * rr;
        let mut data = vec![0f64; n];
        for v in &mut data {
            *v = r.read_f64::<LE>()?;
        }
        let core = Array3::from_shape_vec((rl, d, rr), data)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
        cores.push(core);
    }
    Ok(TT { cores, shape, singular_values: vec![] })
}

pub fn write_blz_c64(tt: &TT<Complex64>, path: &Path) -> io::Result<()> {
    let file = File::create(path)?;
    let mut w = BufWriter::new(file);
    w.write_all(BLZ_MAGIC)?;
    w.write_u8(BLZ_VERSION)?;
    w.write_u8(1)?; // dtype c64
    w.write_u8(tt.ndim() as u8)?;
    w.write_all(&[0u8; 8])?;
    for &d in &tt.shape {
        w.write_u64::<LE>(d as u64)?;
    }
    for &r in &tt.ranks() {
        w.write_u64::<LE>(r as u64)?;
    }
    for core in &tt.cores {
        for &val in core.iter() {
            w.write_f64::<LE>(val.re)?;
            w.write_f64::<LE>(val.im)?;
        }
    }
    w.flush()?;
    Ok(())
}

pub fn read_blz_c64(path: &Path) -> io::Result<TT<Complex64>> {
    let file = File::open(path)?;
    let mut r = BufReader::new(file);
    let mut magic = [0u8; 4];
    r.read_exact(&mut magic)?;
    if &magic != BLZ_MAGIC {
        return Err(io::Error::new(io::ErrorKind::InvalidData, "bad magic"));
    }
    let ver = r.read_u8()?;
    if ver != BLZ_VERSION { return Err(io::Error::new(io::ErrorKind::InvalidData, "bad version")); }
    let dtype = r.read_u8()?;
    if dtype != 1 { return Err(io::Error::new(io::ErrorKind::InvalidData, "expected c64")); }
    let ndim = r.read_u8()? as usize;
    let mut reserved = [0u8; 8]; r.read_exact(&mut reserved)?;
    let mut shape = vec![0usize; ndim];
    for s in &mut shape { *s = r.read_u64::<LE>()? as usize; }
    let mut ranks = vec![0usize; ndim + 1];
    for rk in &mut ranks { *rk = r.read_u64::<LE>()? as usize; }
    let mut cores = Vec::with_capacity(ndim);
    for i in 0..ndim {
        let rl = ranks[i];
        let d = shape[i];
        let rr = ranks[i + 1];
        let n = rl * d * rr;
        let mut data = vec![Complex64::new(0., 0.); n];
        for v in &mut data {
            let re = r.read_f64::<LE>()?;
            let im = r.read_f64::<LE>()?;
            *v = Complex64::new(re, im);
        }
        let core = Array3::from_shape_vec((rl, d, rr), data)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
        cores.push(core);
    }
    Ok(TT { cores, shape, singular_values: vec![] })
}