plot3d 0.1.14

use std::fs::File;
use std::io::{self, BufReader, Read};

use crate::block::Block;
use crate::utils::{self, read_fortran_record, Endian};
use crate::Float;

#[derive(Copy, Clone, Debug)]
pub enum BinaryFormat {
    Fortran,
    Raw,
}

#[derive(Copy, Clone, Debug)]
pub enum FloatPrecision {
    F32,
    F64,
}

// If you want to re-export Endian from this module, do it from the original path:
pub use crate::utils::Endian as EndianOrder;

pub fn read_plot3d_ascii(path: &str) -> io::Result<Vec<Block>> {
    // Bulk-read the entire file into memory for faster tokenization
    let mut contents = String::new();
    File::open(path).and_then(|f| {
        let mut rdr = BufReader::with_capacity(8 * 1024 * 1024, f);
        rdr.read_to_string(&mut contents)
    })?;

    let mut tokens = contents.split_whitespace();

    let nblocks: usize = next_parse(&mut tokens, "bad nblocks")?;

    let mut dims = Vec::with_capacity(nblocks);
    for _ in 0..nblocks {
        let imax: usize = next_parse(&mut tokens, "bad dims")?;
        let jmax: usize = next_parse(&mut tokens, "bad dims")?;
        let kmax: usize = next_parse(&mut tokens, "bad dims")?;
        dims.push((imax, jmax, kmax));
    }

    let mut blocks = Vec::with_capacity(nblocks);
    for (imax, jmax, kmax) in dims {
        let n = imax * jmax * kmax;
        let x = parse_n_floats(&mut tokens, n)?;
        let y = parse_n_floats(&mut tokens, n)?;
        let z = parse_n_floats(&mut tokens, n)?;
        blocks.push(Block::new(imax, jmax, kmax, x, y, z));
    }
    Ok(blocks)
}

/// Parse the next whitespace-delimited token as type `T`.
fn next_parse<T: std::str::FromStr>(
    tokens: &mut std::str::SplitWhitespace<'_>,
    msg: &str,
) -> io::Result<T> {
    tokens
        .next()
        .ok_or_else(|| ioerr(msg))?
        .parse::<T>()
        .map_err(|_| ioerr(msg))
}

/// Parse exactly `n` float tokens from the iterator.
fn parse_n_floats(tokens: &mut std::str::SplitWhitespace<'_>, n: usize) -> io::Result<Vec<Float>> {
    let mut out = Vec::with_capacity(n);
    for _ in 0..n {
        let t = tokens
            .next()
            .ok_or_else(|| ioerr("unexpected EOF in payload"))?;
        out.push(t.parse::<Float>().map_err(|_| ioerr("bad float"))?);
    }
    Ok(out)
}

pub fn read_plot3d_binary(
    path: &str,
    format: BinaryFormat,
    precision: FloatPrecision,
    endian: Endian,
) -> io::Result<Vec<Block>> {
    let f = File::open(path)?;
    let mut r = BufReader::with_capacity(8 * 1024 * 1024, f);
    match format {
        BinaryFormat::Raw => read_binary_raw(&mut r, precision, endian),
        BinaryFormat::Fortran => read_binary_fortran(&mut r, precision, endian),
    }
}

fn read_binary_raw(
    r: &mut impl Read,
    precision: FloatPrecision,
    endian: Endian,
) -> io::Result<Vec<Block>> {
    use byteorder::{BigEndian, LittleEndian, ReadBytesExt};

    // header
    let nblocks = match endian {
        Endian::Little => r.read_u32::<LittleEndian>()?,
        Endian::Big => r.read_u32::<BigEndian>()?,
    } as usize;

    let mut dims = Vec::with_capacity(nblocks);
    for _ in 0..nblocks {
        let imax = match endian {
            Endian::Little => r.read_u32::<LittleEndian>()?,
            Endian::Big => r.read_u32::<BigEndian>()?,
        } as usize;
        let jmax = match endian {
            Endian::Little => r.read_u32::<LittleEndian>()?,
            Endian::Big => r.read_u32::<BigEndian>()?,
        } as usize;
        let kmax = match endian {
            Endian::Little => r.read_u32::<LittleEndian>()?,
            Endian::Big => r.read_u32::<BigEndian>()?,
        } as usize;
        dims.push((imax, jmax, kmax));
    }

    // payload
    let mut blocks = Vec::with_capacity(nblocks);
    for (imax, jmax, kmax) in dims {
        let n = imax * jmax * kmax;
        let x = read_vec_num(r, n, precision, endian)?;
        let y = read_vec_num(r, n, precision, endian)?;
        let z = read_vec_num(r, n, precision, endian)?;
        blocks.push(Block::new(imax, jmax, kmax, x, y, z));
    }
    Ok(blocks)
}

fn read_binary_fortran(
    r: &mut impl Read,
    precision: FloatPrecision,
    endian: Endian,
) -> io::Result<Vec<Block>> {
    // nblocks record
    let nb_rec = read_fortran_record(r, endian)?;
    if nb_rec.len() < 4 {
        return Err(ioerr("short nblocks record"));
    }
    let nblocks = utils::Endian::read_u32(&nb_rec[..4], endian) as usize;

    // dims per block (one record per block)
    let mut dims = Vec::with_capacity(nblocks);
    for _ in 0..nblocks {
        let rec = read_fortran_record(r, endian)?;
        if rec.len() < 12 {
            return Err(ioerr("short dims record"));
        }
        let imax = utils::Endian::read_u32(&rec[0..4], endian) as usize;
        let jmax = utils::Endian::read_u32(&rec[4..8], endian) as usize;
        let kmax = utils::Endian::read_u32(&rec[8..12], endian) as usize;
        dims.push((imax, jmax, kmax));
    }

    // Decode one Fortran record's raw bytes into `Float`s, honoring the
    // file's precision + endianness.
    let decode = |bytes: &[u8]| -> Vec<Float> {
        match precision {
            FloatPrecision::F32 => utils::Endian::read_f32_slice(bytes, endian)
                .into_iter()
                .map(|v| v as Float)
                .collect(),
            FloatPrecision::F64 => utils::Endian::read_f64_slice(bytes, endian)
                .into_iter()
                .map(|v| v as Float)
                .collect(),
        }
    };

    // Payload. The standard PLOT3D convention writes X, Y, Z as a single
    // concatenated record per block; older plot3d-rs / Plot3D_utilities
    // releases wrote them as three separate records. Detect which layout
    // is present from the size of the first payload record per block, so
    // existing files written either way still read correctly.
    let mut blocks = Vec::with_capacity(nblocks);
    for (imax, jmax, kmax) in dims {
        let n = imax * jmax * kmax;

        let first = decode(&read_fortran_record(r, endian)?);
        let (x, y, z) = if first.len() == 3 * n {
            // Single concatenated record: [X | Y | Z].
            let x = first[..n].to_vec();
            let y = first[n..2 * n].to_vec();
            let z = first[2 * n..].to_vec();
            (x, y, z)
        } else if first.len() == n {
            // Legacy three-record layout: this record is X; Y and Z each
            // follow in their own record.
            let y = decode(&read_fortran_record(r, endian)?);
            if y.len() != n {
                return Err(ioerr("Y size mismatch"));
            }
            let z = decode(&read_fortran_record(r, endian)?);
            if z.len() != n {
                return Err(ioerr("Z size mismatch"));
            }
            (first, y, z)
        } else {
            return Err(ioerr(&format!(
                "unexpected Fortran record size for block {}: got {} reals, \
                 expected {} (legacy 3-record) or {} (concatenated)",
                blocks.len(),
                first.len(),
                n,
                3 * n
            )));
        };

        blocks.push(Block::new(imax, jmax, kmax, x, y, z));
    }

    Ok(blocks)
}

fn read_vec_num(
    r: &mut impl Read,
    n: usize,
    precision: FloatPrecision,
    endian: Endian,
) -> io::Result<Vec<Float>> {
    let bytes_per = match precision {
        FloatPrecision::F32 => 4,
        FloatPrecision::F64 => 8,
    };
    let mut buf = vec![0u8; n * bytes_per];
    r.read_exact(&mut buf)?;

    let mut out = Vec::with_capacity(n);
    match (precision, endian) {
        (FloatPrecision::F32, Endian::Little) => {
            for chunk in buf.chunks_exact(4) {
                out.push(f32::from_le_bytes(chunk.try_into().unwrap()) as Float);
            }
        }
        (FloatPrecision::F32, Endian::Big) => {
            for chunk in buf.chunks_exact(4) {
                out.push(f32::from_be_bytes(chunk.try_into().unwrap()) as Float);
            }
        }
        (FloatPrecision::F64, Endian::Little) => {
            for chunk in buf.chunks_exact(8) {
                out.push(f64::from_le_bytes(chunk.try_into().unwrap()) as Float);
            }
        }
        (FloatPrecision::F64, Endian::Big) => {
            for chunk in buf.chunks_exact(8) {
                out.push(f64::from_be_bytes(chunk.try_into().unwrap()) as Float);
            }
        }
    }
    Ok(out)
}

fn ioerr(msg: &str) -> io::Error {
    io::Error::new(io::ErrorKind::InvalidData, msg)
}

/// Read an AP NASA Fortran unformatted file.
///
/// The AP NASA format stores a single block with header integers
/// `[il, jl, kl, ile, ite, jtip, nbld]` followed by `jl` records
/// of interleaved `(x, r, theta)` data (f32).
///
/// Returns `(block, nbld)` where `nbld` is the number of blades.
/// Coordinates are converted from cylindrical `(x, r, theta)` to
/// Cartesian `(x, y, z)`.
pub fn read_ap_nasa(path: &str, endian: Endian) -> io::Result<(Block, i32)> {
    let mut f = File::open(path)?;

    // Record 1: 7 integers [il, jl, kl, ile, ite, jtip, nbld]
    let int_rec = read_fortran_record(&mut f, endian)?;
    if int_rec.len() < 28 {
        return Err(ioerr("AP NASA header too short (expected 7 i32)"));
    }
    let il = utils::Endian::read_u32(&int_rec[0..4], endian) as usize;
    let jl = utils::Endian::read_u32(&int_rec[4..8], endian) as usize;
    let kl = utils::Endian::read_u32(&int_rec[8..12], endian) as usize;
    // ile, ite, jtip are at offsets 12, 16, 20 (not needed for block)
    let nbld = utils::Endian::read_u32(&int_rec[24..28], endian) as i32;

    // Read jl records, each containing 3 * (il * kl) floats (f32)
    let stride = il * kl;
    let total = il * jl * kl;
    let mut meshx = Vec::with_capacity(total);
    let mut meshr = Vec::with_capacity(total);
    let mut mesht = Vec::with_capacity(total);

    for _j in 0..jl {
        let rec = read_fortran_record(&mut f, endian)?;
        let floats = utils::Endian::read_f32_slice(&rec, endian);
        if floats.len() < 3 * stride {
            return Err(ioerr("AP NASA data record too short"));
        }
        // Layout per record: [x(il*kl), r(il*kl), theta(il*kl)]
        for idx in 0..stride {
            meshx.push(floats[idx] as Float);
            meshr.push(floats[stride + idx] as Float);
            mesht.push(floats[2 * stride + idx] as Float);
        }
    }

    // Data arrives in (j, k, i) order from reading jl records each of kl*il entries.
    // Reorder to (i, j, k) with i-fastest for Rust Block convention.
    // Convert from (x, r, theta) to (x, y, z):
    //   y = r * cos(theta), z = r * sin(theta)
    let mut x = Vec::with_capacity(total);
    let mut y = Vec::with_capacity(total);
    let mut z = Vec::with_capacity(total);

    for k in 0..kl {
        for j in 0..jl {
            for i in 0..il {
                let src = j * (kl * il) + k * il + i;
                let r = meshr[src];
                let theta = mesht[src];
                x.push(meshx[src]);
                y.push(r * theta.cos());
                z.push(r * theta.sin());
            }
        }
    }

    Ok((Block::new(il, jl, kl, x, y, z), nbld))
}