#![allow(dead_code)]

mod common; 
use rand::prelude::*; 
use wmath::*; 

fn rand_i32_tensor3(dim: u32) -> Tensor3<i32> {
    let mut rng = rand::thread_rng(); 
    let size = 3usize.pow(dim); 
    let mut vec = Vec::<i32>::with_capacity(size); 
    for _ in 0..size { vec.push(rng.gen_range(-20..=20)); }
    Tensor3::new(vec)
}

fn rand_f64_tensor3(dim: u32) -> Tensor3<f64> {
    let mut rng = rand::thread_rng(); 
    let size = 3usize.pow(dim); 
    let mut vec = Vec::<f64>::with_capacity(size); 
    for _ in 0..size { vec.push(rng.gen::<f64>() * 10.0 - 5.0); }
    Tensor3::new(vec)
}

fn rand_frac_tensor3(dim: u32) -> Tensor3<Fraction> {
    let mut rng = rand::thread_rng(); 
    let size = 3usize.pow(dim); 
    let mut vec = Vec::<Fraction>::with_capacity(size); 
    for _ in 0..size { vec.push(frac!((rng.gen::<f64>() * 10.0 - 5.0) as i128, (rng.gen::<f64>() * 8.0 + 1.0) as i128)); }
    Tensor3::new(vec)
}

fn rand_cplx_tensor3(dim: u32) -> Tensor3<Complex> {
    let mut rng = rand::thread_rng(); 
    let size = 3usize.pow(dim); 
    let mut vec = Vec::<Complex>::with_capacity(size); 
    for _ in 0..size { vec.push(cplx!(rng.gen::<f64>() * 10.0 - 5.0, rng.gen::<f64>() * 10.0 - 5.0)); }
    Tensor3::new(vec)
}

#[test]
fn test_tensor3() {
    let t0 = rand_i32_tensor3(0); 
    let t1 = rand_i32_tensor3(3); 
    let t2 = t1.get_zero(); 
    let t3 = t1.slice(&[1, 0, 2]); 
    let t4 = t1.multiply_by(3); 
    print_expr!(t0, t4, t2, t3, t1); 
    let t5 = t1.contract2(1, 2); 
    let t6 = t1.dyadic(&t5); 
    let t7 = rand_i32_tensor3(2); 
    let t8 = rand_i32_tensor3(2); 
    let t9 = t7.dot(&t8, 1); 
    let t10 = t7.dot_co(&t8, 2); 
    let t11 = rand_i32_tensor3(1); 
    let t12 = rand_i32_tensor3(1); 
    let t13 = (t11.dyadic(&t12)).contract_cross2(1, 2, 1); 
    let t14 = t11.cross(&t12); 
    let t15 = rand_i32_tensor3(5); 
    let t16 = rand_i32_tensor3(5); 
    let t17 = t15.dyadic(&t16).contract2(4, 5 + 3); 
    let t18 = t15.dot2(&t16, 4, 3); 
    print_expr!(t5, t6, t7, t8, t9, t10, t11, t12); 
    assert_eq!(t13, t14); 
    assert_eq!(t17, t18); 

    let t0 = rand_f64_tensor3(0); 
    let t1 = rand_f64_tensor3(3); 
    let t2 = t1.get_zero(); 
    let t3 = t1.slice(&[1, 0, 2]); 
    let t4 = t1.multiply_by(3.0); 
    print_expr!(t0, t4, t2, t3, t1); 
    let t5 = t1.contract2(1, 2); 
    let t6 = t1.dyadic(&t5); 
    let t7 = rand_f64_tensor3(2); 
    let t8 = rand_f64_tensor3(2); 
    let t9 = t7.dot(&t8, 1); 
    let t10 = t7.dot_co(&t8, 2); 
    let t11 = rand_f64_tensor3(1); 
    let t12 = rand_f64_tensor3(1); 
    let t13 = (t11.dyadic(&t12)).contract_cross2(1, 2, 1); 
    let t14 = t11.cross(&t12); 
    let t15 = rand_f64_tensor3(5); 
    let t16 = rand_f64_tensor3(5); 
    let t17 = t15.dyadic(&t16).contract2(4, 5 + 3); 
    let t18 = t15.dot2(&t16, 4, 3); 
    print_expr!(t5, t6, t7, t8, t9, t10, t11, t12); 
    assert_eq!(t13, t14); 
    assert_eq!(t17, t18); 

    let t0 = rand_frac_tensor3(0); 
    let t1 = rand_frac_tensor3(3); 
    let t2 = t1.get_zero(); 
    let t3 = t1.slice(&[1, 0, 2]); 
    let t4 = t1.multiply_by(frac!(3, 4)); 
    print_expr!(t0, t4, t2, t3, t1); 
    let t5 = t1.contract2(1, 2); 
    let t6 = t1.dyadic(&t5); 
    let t7 = rand_frac_tensor3(2); 
    let t8 = rand_frac_tensor3(2); 
    let t9 = t7.dot(&t8, 1); 
    let t10 = t7.dot_co(&t8, 2); 
    let t11 = rand_frac_tensor3(1); 
    let t12 = rand_frac_tensor3(1); 
    let t13 = (t11.dyadic(&t12)).contract_cross2(1, 2, 1); 
    let t14 = t11.cross(&t12); 
    let t15 = rand_frac_tensor3(5); 
    let t16 = rand_frac_tensor3(5); 
    let t17 = t15.dyadic(&t16).contract2(4, 5 + 3); 
    let t18 = t15.dot2(&t16, 4, 3); 
    print_expr!(t5, t6, t7, t8, t9, t10, t11, t12); 
    assert_eq!(t13, t14); 
    assert_eq!(t17, t18); 

    let t0 = rand_cplx_tensor3(0); 
    let t1 = rand_cplx_tensor3(3); 
    let t2 = t1.get_zero(); 
    let t3 = t1.slice(&[1, 0, 2]); 
    let t4 = t1.multiply_by(cplx!(3.0, -4.0)); 
    print_expr!(t0, t4, t2, t3, t1); 
    let t5 = t1.contract2(1, 2); 
    let t6 = t1.dyadic(&t5); 
    let t7 = rand_cplx_tensor3(2); 
    let t8 = rand_cplx_tensor3(2); 
    let t9 = t7.dot(&t8, 1); 
    let t10 = t7.dot_co(&t8, 2); 
    let t11 = rand_cplx_tensor3(1); 
    let t12 = rand_cplx_tensor3(1); 
    let t13 = (t11.dyadic(&t12)).contract_cross2(1, 2, 1); 
    let t14 = t11.cross(&t12); 
    let t15 = rand_cplx_tensor3(5); 
    let t16 = rand_cplx_tensor3(5); 
    let t17 = t15.dyadic(&t16).contract2(4, 5 + 3); 
    let t18 = t15.dot2(&t16, 4, 3); 
    print_expr!(t5, t6, t7, t8, t9, t10, t11, t12); 
    assert_eq!(t13, t14); 
    assert_eq!(t17, t18); 

    // transpose
    use std::time::Instant; 
    let dim = 8; 
    let dim1 = 3; 
    let dim2 = 6; 
    let t0 = rand_f64_tensor3(dim); 
    let start = Instant::now(); 
    let t1 = t0.transpose(dim1, dim2); 
    println!("---[{} ms]---", start.elapsed().as_millis()); 
    let idx = vec![0; dim as usize]; 
    for i in 1..=3 {
        for j in 1..=3 {
            let mut t0_idx = idx.clone(); 
            t0_idx[dim1 as usize - 1] = i; 
            t0_idx[dim2 as usize - 1] = j; 
            let mut t1_idx = idx.clone(); 
            t1_idx[dim1 as usize - 1] = j; 
            t1_idx[dim2 as usize - 1] = i; 
            assert_eq!(t0.slice(&t0_idx), t1.slice(&t1_idx)); 
        }
    }
    
    // insert
    let mut t0 = rand_f64_tensor3(10); 
    let t1 = rand_f64_tensor3(5); 
    let idx = [0, 1, 0, 2, 0, 1, 0, 0, 2, 1]; 
    t0.insert(&t1, &idx); 
    print_expr!(t0.slice(&idx) == t1); 

    // concat
    let mut t0s = vec![]; 
    for _ in 0..3usize.pow(4) {
        t0s.push(rand_f64_tensor3(3)); 
    }
    let t0 = Tensor3::concat(&t0s[..], &[1, 4, 2, 6]); 
    print_expr!(t0.slice(&[2, 3, 0, 1, 0, 3, 0]) == t0s[1 * 27 + 0 * 9 + 2 * 3 + 2]); 

    let mut t0s = vec![]; 
    for _ in 0..3usize.pow(1) {
        t0s.push(rand_i32_tensor3(1)); 
    }
    let t0 = Tensor3::concat(&t0s[..], &[2]); 
    for i in t0s.iter() { print_expr!(i); }
    print_expr!(t0); 
}

fn main() {}