use tempfile::tempdir;

use crate::{pack_compact_int_matrix, PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder};
use crate::traits::CountPartials;

fn make_matrix(cols: &[&[u32]]) -> (tempfile::TempDir, PersistentCompactIntMatrix) {
    let n = cols.first().map_or(0, |c| c.len());
    let dir = tempdir().unwrap();
    let mut b = PersistentCompactIntMatrixBuilder::new(n, &dir.path().join("counts")).unwrap();
    for &col in cols {
        let mut cb = b.add_col().unwrap();
        for (slot, &v) in col.iter().enumerate() {
            cb.set(slot, v);
        }
        cb.close().unwrap();
    }
    b.close().unwrap();
    let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
    (dir, m)
}

#[test]
fn single_col_roundtrip() {
    let dir = tempdir().unwrap();
    let mut b = PersistentCompactIntMatrixBuilder::new(4, &dir.path().join("counts")).unwrap();
    let mut col = b.add_col().unwrap();
    col.set(0, 10);
    col.set(1, 200);
    col.set(2, 300);
    col.set(3, 1000);
    col.close().unwrap();
    b.close().unwrap();

    let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
    assert_eq!(m.n_cols(), 1);
    assert_eq!(m.n(), 4);
    assert_eq!(&*m.row(0), &[10u32]);
    assert_eq!(&*m.row(1), &[200u32]);
    assert_eq!(&*m.row(2), &[300u32]);
    assert_eq!(&*m.row(3), &[1000u32]);
}

#[test]
fn two_cols_roundtrip() {
    let dir = tempdir().unwrap();
    let mut b = PersistentCompactIntMatrixBuilder::new(3, &dir.path().join("counts")).unwrap();
    let mut col0 = b.add_col().unwrap();
    col0.set(0, 1); col0.set(1, 2); col0.set(2, 3);
    col0.close().unwrap();
    let mut col1 = b.add_col().unwrap();
    col1.set(0, 10); col1.set(1, 20); col1.set(2, 30);
    col1.close().unwrap();
    b.close().unwrap();

    let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
    assert_eq!(m.n_cols(), 2);
    assert_eq!(&*m.row(0), &[1u32, 10]);
    assert_eq!(&*m.row(1), &[2u32, 20]);
    assert_eq!(&*m.row(2), &[3u32, 30]);
}

#[test]
fn col_accessor() {
    let dir = tempdir().unwrap();
    let mut b = PersistentCompactIntMatrixBuilder::new(2, &dir.path().join("counts")).unwrap();
    let mut col0 = b.add_col().unwrap();
    col0.set(0, 5); col0.set(1, 7);
    col0.close().unwrap();
    b.close().unwrap();

    let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
    assert_eq!(m.col(0).get(0), 5);
    assert_eq!(m.col(0).get(1), 7);
}

#[test]
fn zero_cols_roundtrip() {
    let dir = tempdir().unwrap();
    let b = PersistentCompactIntMatrixBuilder::new(10, &dir.path().join("counts")).unwrap();
    b.close().unwrap();

    let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
    assert_eq!(m.n_cols(), 0);
    assert_eq!(m.n(), 10);
}

// ── Distance matrix tests ─────────────────────────────────────────────────────

#[test]
fn bray_dist_matrix_symmetry_and_diagonal() {
    // col0=[1,0,1], col1=[1,1,0], col2=[0,1,1]
    let (_d, m) = make_matrix(&[&[1, 0, 1], &[1, 1, 0], &[0, 1, 1]]);
    let dm = m.bray_dist_matrix();
    let n = m.n_cols();
    for i in 0..n { assert_eq!(dm[[i, i]], 0.0, "diagonal"); }
    for i in 0..n {
        for j in 0..n {
            assert!((dm[[i, j]] - dm[[j, i]]).abs() < 1e-12, "symmetry");
        }
    }
}

#[test]
fn bray_dist_matrix_values_match_pairwise() {
    let (_d, m) = make_matrix(&[&[1, 0, 1], &[1, 1, 0], &[0, 1, 1]]);
    let dm = m.bray_dist_matrix();
    for i in 0..m.n_cols() {
        for j in 0..m.n_cols() {
            let expected = m.col(i).bray_dist(m.col(j));
            assert!((dm[[i, j]] - expected).abs() < 1e-12, "[{i},{j}]");
        }
    }
}

#[test]
fn jaccard_dist_matrix_values_match_pairwise() {
    let (_d, m) = make_matrix(&[&[1, 0, 2], &[0, 1, 1], &[1, 1, 0]]);
    let dm = m.jaccard_dist_matrix();
    for i in 0..m.n_cols() {
        for j in 0..m.n_cols() {
            let expected = m.col(i).jaccard_dist(m.col(j));
            assert!((dm[[i, j]] - expected).abs() < 1e-12, "[{i},{j}]");
        }
    }
}

#[test]
fn partial_bray_dist_matrix_consistent() {
    let (_d, m) = make_matrix(&[&[1, 0, 1], &[1, 1, 0], &[0, 1, 1]]);
    let sum_min  = m.partial_bray_dist_matrix();
    let col_sums = m.sum();
    let n = m.n_cols();

    // symmetry
    for i in 0..n {
        for j in 0..n {
            assert_eq!(sum_min[[i, j]], sum_min[[j, i]]);
        }
    }

    // reconstruct distance from partials and compare to direct method
    for i in 0..n {
        for j in i + 1..n {
            let denom = col_sums[i] + col_sums[j];
            let dist = if denom == 0 { 0.0 } else { 1.0 - 2.0 * sum_min[[i, j]] as f64 / denom as f64 };
            let expected = m.col(i).bray_dist(m.col(j));
            assert!((dist - expected).abs() < 1e-12, "[{i},{j}]");
        }
    }
}

#[test]
fn partial_euclidean_dist_matrix_consistent() {
    let (_d, m) = make_matrix(&[&[3, 0], &[0, 4], &[1, 1]]);
    let sq = m.partial_euclidean_dist_matrix();
    let n = m.n_cols();
    for i in 0..n {
        for j in i + 1..n {
            let expected = m.col(i).partial_euclidean_dist(m.col(j));
            assert!((sq[[i, j]] - expected).abs() < 1e-12, "[{i},{j}]");
            assert!((sq[[j, i]] - expected).abs() < 1e-12, "symmetry [{j},{i}]");
        }
    }
}

#[test]
fn partial_threshold_jaccard_consistent() {
    let (_d, m) = make_matrix(&[&[3, 0, 2], &[1, 2, 0], &[0, 3, 1]]);
    let threshold = 2u32;
    let (inter, union) = m.partial_threshold_jaccard_dist_matrix(threshold);
    let n = m.n_cols();
    for i in 0..n {
        for j in i + 1..n {
            let (ei, eu) = m.col(i).partial_threshold_jaccard_dist(m.col(j), threshold);
            assert_eq!(inter[[i, j]], ei);
            assert_eq!(union[[i, j]], eu);
            assert_eq!(inter[[j, i]], ei, "symmetry inter");
            assert_eq!(union[[j, i]], eu, "symmetry union");
        }
    }
}

#[test]
fn single_col_distance_matrix_is_zero() {
    let (_d, m) = make_matrix(&[&[1, 2, 3]]);
    let dm = m.bray_dist_matrix();
    assert_eq!(dm.shape(), &[1, 1]);
    assert_eq!(dm[[0, 0]], 0.0);
}

#[test]
fn sum_returns_column_sums() {
    let (_d, m) = make_matrix(&[&[1, 2, 3], &[10, 20, 30], &[0, 0, 5]]);
    let s = m.sum();
    assert_eq!(s[0], 6);
    assert_eq!(s[1], 60);
    assert_eq!(s[2], 5);
}

#[test]
fn partial_relfreq_bray_matches_full() {
    let (_d, m) = make_matrix(&[&[1, 0, 2], &[0, 1, 1], &[1, 1, 0]]);
    let col_sums = m.sum();
    let partial = m.partial_relfreq_bray_dist_matrix(&col_sums);
    let full    = m.relfreq_bray_dist_matrix();
    let n = m.n_cols();
    // partial[i,j] = sum_min_relfreq; full[i,j] = 1 - sum_min_relfreq (off-diagonal only)
    for i in 0..n {
        for j in 0..n {
            if i == j { continue; }
            assert!((partial[[i, j]] - (1.0 - full[[i, j]])).abs() < 1e-12, "[{i},{j}]");
        }
    }
}

#[test]
fn partial_relfreq_euclidean_matches_full() {
    let (_d, m) = make_matrix(&[&[3, 0], &[0, 4], &[1, 1]]);
    let col_sums = m.sum();
    let partial = m.partial_relfreq_euclidean_dist_matrix(&col_sums);
    let full    = m.relfreq_euclidean_dist_matrix();
    let n = m.n_cols();
    for i in 0..n {
        for j in 0..n {
            // partial = squared euclidean; full = sqrt(partial)
            assert!((partial[[i, j]].sqrt() - full[[i, j]]).abs() < 1e-12, "[{i},{j}]");
        }
    }
}

#[test]
fn partial_hellinger_matches_full() {
    let (_d, m) = make_matrix(&[&[3, 0], &[0, 4], &[1, 1]]);
    let col_sums = m.sum();
    let partial = m.partial_hellinger_euclidean_dist_matrix(&col_sums);
    let full    = m.hellinger_dist_matrix();
    let n = m.n_cols();
    for i in 0..n {
        for j in 0..n {
            // partial / sqrt(2) gives the Hellinger distance
            assert!((partial[[i, j]].sqrt() / std::f64::consts::SQRT_2 - full[[i, j]]).abs() < 1e-12, "[{i},{j}]");
        }
    }
}

#[test]
fn col_view_packed_values() {
    // Build Columnar with overflow values (≥ 255), pack, reopen as Packed, exercise col_view().
    let (dir, _col) = make_matrix(&[&[10, 300, 500], &[200, 50, 1000]]);
    pack_compact_int_matrix(&dir.path().join("counts")).unwrap();
    let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();

    // col 0: [10, 300, 500] — two overflow slots
    let v0 = m.col_view(0);
    assert_eq!(v0.get(0), 10);
    assert_eq!(v0.get(1), 300);
    assert_eq!(v0.get(2), 500);
    assert_eq!(v0.sum(), 810);
    assert_eq!(v0.count_nonzero(), 3);
    let mut ov0: Vec<(usize, u32)> = v0.overflow_entries().collect();
    ov0.sort_unstable_by_key(|&(s, _)| s);
    assert_eq!(ov0, vec![(1, 300), (2, 500)]);

    // col 1: [200, 50, 1000] — one overflow slot
    let v1 = m.col_view(1);
    assert_eq!(v1.get(0), 200);
    assert_eq!(v1.get(1), 50);
    assert_eq!(v1.get(2), 1000);
    let mut ov1: Vec<(usize, u32)> = v1.overflow_entries().collect();
    ov1.sort_unstable_by_key(|&(s, _)| s);
    assert_eq!(ov1, vec![(2, 1000)]);
}

#[test]
fn col_view_packed_matches_columnar() {
    // Same data, compare col_view() on Packed against col() on Columnar slot-by-slot.
    let data: &[&[u32]] = &[&[0, 255, 1, 300, 128], &[500, 3, 0, 700, 42]];
    let (dir_col, m_col) = make_matrix(data);
    // Re-build in a separate dir so we can pack without touching m_col's files.
    let (dir_pack, _) = make_matrix(data);
    pack_compact_int_matrix(&dir_pack.path().join("counts")).unwrap();
    let m_pack = PersistentCompactIntMatrix::open(dir_pack.path()).unwrap();

    for c in 0..data.len() {
        let col_ref  = m_col.col(c);
        let col_view = m_pack.col_view(c);
        assert_eq!(col_view.len(), col_ref.len());
        for s in 0..col_ref.len() {
            assert_eq!(col_view.get(s), col_ref.get(s), "col={c} slot={s}");
        }
        assert_eq!(col_view.sum(), col_ref.sum(), "col={c} sum");
        let mut ov_view: Vec<(usize, u32)> = col_view.overflow_entries().collect();
        let mut ov_ref:  Vec<(usize, u32)> = col_ref.view().overflow_entries().collect();
        ov_view.sort_unstable_by_key(|&(s, _)| s);
        ov_ref.sort_unstable_by_key(|&(s, _)| s);
        assert_eq!(ov_view, ov_ref, "col={c} overflow_entries");
    }
    drop(dir_col);
}

#[test]
fn partial_relfreq_bray_additive_across_split() {
    // Split rows [1,2,3,4,5] between two matrices; partial sums should add up.
    // col0=[1,2,3,4,5], col1=[5,4,3,2,1]
    // global sums: col0=15, col1=15
    // split: first=[1,2], second=[3,4,5]
    let (_d1, m1) = make_matrix(&[&[1, 2], &[5, 4]]);
    let (_d2, m2) = make_matrix(&[&[3, 4, 5], &[3, 2, 1]]);
    let global_sums = ndarray::array![15u64, 15u64];
    let p1 = m1.partial_relfreq_bray_dist_matrix(&global_sums);
    let p2 = m2.partial_relfreq_bray_dist_matrix(&global_sums);
    let combined = &p1 + &p2;

    let (_d, m_full) = make_matrix(&[&[1, 2, 3, 4, 5], &[5, 4, 3, 2, 1]]);
    let full_partial = m_full.partial_relfreq_bray_dist_matrix(&global_sums);
    let n = m_full.n_cols();
    for i in 0..n {
        for j in 0..n {
            assert!((combined[[i, j]] - full_partial[[i, j]]).abs() < 1e-12, "[{i},{j}]");
        }
    }
}
