use std::fs::File;
use std::io::prelude::*;
use std::io::BufReader;

fn parse_line(line: String) -> Vec<f64> {
    let line = line.split_whitespace();
    line.map(|x| x.parse::<f64>().unwrap())
        .collect::<Vec<f64>>()
}

#[allow(dead_code)]
#[derive(Debug)]
pub struct Atom {
    pub an: usize,
    pub charge: f32,
    pub pos: Vec<f64>,
}

#[derive(Debug)]
pub struct CubeData {
    pub n_atoms: usize,
    pub origin: Vec<f64>,
    pub shape: Vec<usize>,
    pub data: Vec<f64>,
    pub cell: Vec<Vec<f64>>,
    pub atoms: Vec<Atom>,
}

impl CubeData {
    fn new() -> Self {
        CubeData {
            n_atoms: 0,
            origin: vec![0f64; 3],
            shape: vec![0usize; 3],
            cell: vec![vec![0f64; 3]; 3],
            atoms: Vec::new(),
            data: Vec::new(),
        }
    }
}

pub fn read_cube(file_name: &str) -> std::io::Result<CubeData> {
    let f = File::open(file_name)?;
    let reader = BufReader::new(f);
    let mut cube_data = CubeData::new();

    let mut lines = reader.lines();
    lines.next();
    lines.next();
    let line3 = parse_line(lines.next().unwrap()?);
    cube_data.n_atoms = line3[0] as usize;
    cube_data.origin = Vec::from(&line3[1..]);

    let line4 = parse_line(lines.next().unwrap()?);
    let line5 = parse_line(lines.next().unwrap()?);
    let line6 = parse_line(lines.next().unwrap()?);

    cube_data.shape[0] = line4[0] as usize;
    cube_data.cell[0] = Vec::from(&line4[1..])
        .iter()
        .map(|x| x * cube_data.shape[0] as f64)
        .collect();
    cube_data.shape[1] = line5[0] as usize;
    cube_data.cell[1] = Vec::from(&line5[1..])
        .iter()
        .map(|x| x * cube_data.shape[1] as f64)
        .collect();
    cube_data.shape[2] = line6[0] as usize;
    cube_data.cell[2] = Vec::from(&line6[1..])
        .iter()
        .map(|x| x * cube_data.shape[2] as f64)
        .collect();

    let mut idx = 0;
    while idx < cube_data.n_atoms {
        let line_atom = parse_line(lines.next().unwrap()?);
        cube_data.atoms.push(Atom {
            an: line_atom[0] as usize,
            charge: line_atom[1] as f32,
            pos: Vec::from(&line_atom[2..]),
        });
        idx += 1;
    }

    for line in lines {
        cube_data.data.append(&mut parse_line(line?));
    }

    Ok(cube_data)
}

pub fn write_cube(file_name: &str, cube_data: &CubeData) -> std::io::Result<()> {
    let mut file = File::create(file_name)?;
    let delimiter = "    ";
    file.write_all(b"Cube file generated by cubeio.rs\n")?;
    file.write_all(b"OUTER LOOP: X, MIDDLE LOOP: Y, INNER LOOP: Z\n")?;

    file.write_all(cube_data.n_atoms.to_string().as_bytes())?;
    file.write_all(delimiter.as_bytes())?;
    file.write_all(
        cube_data
            .origin
            .iter()
            .map(|x| {
                let mut x = x.to_string();
                x.push_str(delimiter);
                return x;
            })
            .collect::<String>()
            .as_bytes(),
    )?;
    file.write_all(b"\n")?;

    for i in 0..3 {
        file.write_all(cube_data.shape[i].to_string().as_bytes())?;
        file.write_all(delimiter.as_bytes())?;
        file.write_all(
            cube_data.cell[i]
                .iter()
                .map(|x| {
                    let mut x = (x / cube_data.shape[i] as f64).to_string();
                    x.push_str(delimiter);
                    return x;
                })
                .collect::<String>()
                .as_bytes(),
        )?;
        file.write_all(b"\n")?;
    }

    for i in 0..(cube_data.n_atoms as usize) {
        file.write_all(cube_data.atoms[i].an.to_string().as_bytes())?;
        file.write_all(delimiter.as_bytes())?;
        file.write_all(cube_data.atoms[i].charge.to_string().as_bytes())?;
        file.write_all(delimiter.as_bytes())?;
        file.write_all(
            cube_data.atoms[i]
                .pos
                .iter()
                .map(|x| {
                    let mut x = x.to_string();
                    x.push_str(delimiter);
                    return x;
                })
                .collect::<String>()
                .as_bytes(),
        )?;
        file.write_all(b"\n")?;
    }
    for dat in cube_data.data.iter() {
        file.write_all(dat.to_string().as_bytes())?;
        file.write_all(b"\n")?;
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use float_cmp::approx_eq;
    #[test]
    fn read_file() {
        let cube_data = super::read_cube("tdc.cube").unwrap();
        assert_eq!(cube_data.n_atoms, 30);
        assert_eq!(cube_data.origin, vec![0f64; 3]);
        assert_eq!(cube_data.shape, vec![79usize, 52usize, 85usize]);

        assert!(approx_eq!(f64, cube_data.cell[0][0], 14.831539));
        assert!(approx_eq!(f64, cube_data.cell[0][1], 0.0));
        assert!(approx_eq!(f64, cube_data.cell[0][2], 0.0));

        assert!(approx_eq!(f64, cube_data.cell[1][0], 0.0));
        assert!(approx_eq!(f64, cube_data.cell[1][1], 9.67538));
        assert!(approx_eq!(f64, cube_data.cell[1][2], 0.0));

        assert!(approx_eq!(f64, cube_data.cell[2][0], -1.263185));
        assert!(approx_eq!(f64, cube_data.cell[2][1], 0.0));
        assert!(approx_eq!(f64, cube_data.cell[2][2], 15.872644999999999));

        assert_eq!(cube_data.atoms[0].an, 3usize);
        assert_eq!(cube_data.atoms[cube_data.atoms.len() - 1].an, 16usize);
        assert!(approx_eq!(f32, cube_data.atoms[0].charge, 0f32));
        assert_eq!(
            cube_data.data.len(),
            (cube_data.shape[0] * cube_data.shape[1] * cube_data.shape[2]) as usize
        );
    }
    #[test]
    fn write_file() {
        let cube_data = super::read_cube("tdc.cube").unwrap();
        super::write_cube("tdc_out.cube", &cube_data).unwrap();
        let cube_data_read = super::read_cube("tdc_out.cube").unwrap();
        assert_eq!(cube_data.n_atoms, cube_data_read.n_atoms);
        assert_eq!(cube_data.origin, cube_data_read.origin);
        assert_eq!(cube_data.shape, cube_data_read.shape);
        assert_eq!(cube_data.atoms[0].an, cube_data_read.atoms[0].an);
        assert_eq!(
            cube_data.atoms[cube_data.atoms.len() - 1].an,
            cube_data_read.atoms[cube_data_read.atoms.len() - 1].an
        );
        approx_eq!(
            f32,
            cube_data.atoms[0].charge,
            cube_data_read.atoms[0].charge
        );
        assert_eq!(
            cube_data_read.data.len(),
            (cube_data_read.shape[0] * cube_data_read.shape[1] * cube_data_read.shape[2]) as usize
        );
    }
}