//! Agilent (Oxford) Esperanto format reader and writer.
//!
//! The reader supports uncompressed (only 32 bit signed integers) and Bitfield compressed frames.
//!
//! Example usage:
//! ```rust
//! use std::path::Path;
//! use cryiorust::esperanto::Esperanto;
//! use cryiorust::frame::Frame;
//!
//! fn test_read<P: AsRef<Path>>(path: P) {
//!     let frame = Esperanto::read_file(path).unwrap();
//!     assert_eq!(frame.dim1(), 2048, "dim1 != 2048");
//!     assert_eq!(frame.dim2(), 2048, "dim2 != 2048");
//!     assert_eq!(frame.sum(), 424962850., "frame sum != 424962850");
//!     assert_eq!(frame.min(), 86., "frame min != 86");
//!     assert_eq!(frame.max(), 62570., "frame max != 62570");
//!     assert_eq!(
//!         frame.array()[frame.dim2() * 1272 + 303],
//!         58608.,
//!         "frame[1272, 303] != 58608"
//!     );
//!     let temp = tempfile::NamedTempFile::new().unwrap();
//!     frame.write_file(temp.path()).unwrap();
//! }
//! ```
use crate::frame::{Array, Frame, FrameError, FrameResult, Header};
use askama::Template;
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use cryiorust_derive::Frame;
use reinterpret::reinterpret_slice;
use std::fs::File;
use std::io::{BufWriter, Cursor, Write};
use std::mem::{size_of, MaybeUninit};
use std::num::Wrapping;
use std::path::Path;
use std::{cmp, fs, io};

/// Main exported struct.
#[derive(Frame)]
pub struct Esperanto {
    header: Header,
    array: Array,
    dim1: usize,
    dim2: usize,
    compression: Compression,
    blocks: usize,
    tail: usize,
    i: usize,
}

enum Compression {
    Uncompressed,
    AgiBitfield,
}

static MASK: &'static [u64; 9] = &[0x00, 0x01, 0x03, 0x07, 0x0f, 0x1f, 0x3f, 0x7f, 0xff];

impl Esperanto {
    const HEADER_LINES: usize = 25;
    const HEADER_LINE_SIZE: usize = 256;
    const HEADER_SIZE: usize = Esperanto::HEADER_LINES * Esperanto::HEADER_LINE_SIZE;
    const HEADER_KEY_IMAGE: &'static str = "IMAGE";
    const HEADER_START: &'static str =
        "ESPERANTO FORMAT 1 CONSISTING OF 25 LINES OF 256 BYTES EACH";
    const HEADER_IMAGE_ENTRIES: usize = 6;
    const HEADER_UNCOMPRESSED: &'static str = "\"4BYTE_LONG\"";
    const HEADER_AGI_BITFIELD: &'static str = "\"AGI_BITFIELD\"";
    const OVERFLOW_16: u8 = 0xfe;
    const OVERFLOW_32: u8 = 0xff;

    /// Reads a file as an Esperanto frame.
    pub fn read_file<P: AsRef<Path>>(path: P) -> FrameResult<Esperanto> {
        let mut file = fs::File::open(path)?;
        let mut frame = Esperanto {
            header: Default::default(),
            array: Array::new(),
            dim1: 0,
            dim2: 0,
            compression: Compression::Uncompressed,
            blocks: 0,
            tail: 0,
            i: 0,
        };
        frame.read_header(&mut file)?;
        frame.read_data(&mut file)?;
        Ok(frame)
    }

    /// Writes Esperanto into a file with ```path```.
    pub fn write_file<P: AsRef<Path>>(&self, path: P) -> FrameResult<()> {
        let mut writer = BufWriter::new(File::create(path)?);
        writer.write_all(self.render_header().as_bytes())?;
        let dim = cmp::max(self.dim1, self.dim2);
        let dim = (dim + 3) & !3;
        let (packed, addresses) = self.compress(dim)?;
        writer.write_i32::<LittleEndian>(packed.len() as i32)?;
        writer.write_all(&packed)?;
        for value in addresses {
            writer.write_i32::<LittleEndian>(value)?;
        }
        Ok(())
    }

    fn compress(&self, dim: usize) -> FrameResult<(Vec<u8>, Vec<i32>)> {
        let blocks = (dim - 1) / (2 * 8);
        let tail = dim - blocks * 2 * 8;
        let mut packed = Vec::with_capacity(dim * dim);
        let mut addresses = Vec::with_capacity(dim);
        let mut oft = Vec::with_capacity(size_of::<i32>() * 8 * 2);
        let mut block1: [i32; 8] = unsafe { MaybeUninit::uninit().assume_init() };
        let mut block2: [i32; 8] = unsafe { MaybeUninit::uninit().assume_init() };
        for i in 0..dim {
            addresses.push(packed.len() as i32);
            let value = self.value(i, 0);
            value.pack8(&mut packed)?;
            for j in (1..dim - tail).step_by(tail) {
                block1[0] = self.value(i, j) - self.value(i, j - 1);
                block2[0] = self.value(i, j + 8) - self.value(i, j + 7);
                let mut max1 = block1[0];
                let mut max2 = block2[0];
                let mut min1 = block1[0];
                let mut min2 = block2[0];
                for k in 1..8 {
                    let y1 = j + k;
                    let y0 = y1 - 1;
                    block1[k] = self.value(i, y1) - self.value(i, y0);
                    block2[k] = self.value(i, y1 + 8) - self.value(i, y0 + 8);
                    max1 = cmp::max(max1, block1[k]);
                    max2 = cmp::max(max2, block2[k]);
                    min1 = cmp::min(min1, block1[k]);
                    min2 = cmp::min(min2, block2[k]);
                }
                let size1 = cmp::max(min1.bit_size(), max1.bit_size());
                let size2 = cmp::max(min2.bit_size(), max2.bit_size());
                packed.push((size2 << 4) | size1);
                let size1 = size1 as usize;
                let value1 = block1.pack(size1, &mut oft)?;
                packed.extend_from_slice(&value1.to_le_bytes()[..size1]);
                let size2 = size2 as usize;
                let value2 = block2.pack(size2, &mut oft)?;
                packed.extend_from_slice(&value2.to_le_bytes()[..size2]);
                packed.extend_from_slice(&oft);
                oft.clear();
            }
            for j in dim - tail + 1..dim {
                let value = self.value(i, j) - self.value(i, j - 1);
                value.pack8(&mut packed)?;
            }
        }
        Ok((packed, addresses))
    }

    fn value(&self, x: usize, y: usize) -> i32 {
        if y < self.dim2 && x < self.dim1 {
            self.array[x * self.dim2 + y].round() as i32
        } else {
            -1
        }
    }

    fn render_header(&self) -> String {
        let tpl = EsperantoHeaderTemplate {
            dim1: self.dim1,
            dim2: self.dim2,
            datatype: match self.compression {
                Compression::Uncompressed => Esperanto::HEADER_UNCOMPRESSED,
                Compression::AgiBitfield => Esperanto::HEADER_AGI_BITFIELD,
            },
            exposure_time: 0.0,
            flux: 0.,
            pixel1: 0,
            pixel2: 0,
            datetime: chrono::Local::now()
                .format("%Y-%m-%dT%H:%M:%S%.f")
                .to_string(),
            omega: 0.0,
            domega: 0.0,
            theta: 0.0,
            dtheta: 0.0,
            kappa: 0.0,
            dkappa: 0.0,
            phi: 0.0,
            dphi: 0.0,
            center_x: 0.0,
            center_y: 0.0,
            alpha: 0.0,
            dist: 0.0,
            l1: 0.0,
            l2: 0.0,
            l12: 0.0,
            b: 0.0,
        };
        let mut header = String::with_capacity(Esperanto::HEADER_SIZE);
        for line in tpl.render().unwrap().lines() {
            header.push_str(line);
            for _ in line.len()..Esperanto::HEADER_LINE_SIZE - 2 {
                header.push(' ');
            }
            header.push_str("\r\n");
        }
        header
    }

    fn read_header<P: io::Read + io::Seek>(&mut self, mut reader: P) -> FrameResult<()> {
        let mut buf: [u8; Esperanto::HEADER_SIZE] = unsafe { MaybeUninit::uninit().assume_init() };
        reader.read_exact(&mut buf)?;
        let header = std::str::from_utf8(&buf)?;
        if !header.starts_with(Esperanto::HEADER_START) {
            return Err(FrameError::HeaderError("bad esperanto header".into()));
        }
        for line in header.lines() {
            if line.starts_with(Esperanto::HEADER_KEY_IMAGE) {
                self.parse_header_image_info(&line)?;
                break; // for now, we don't care about the rest
            }
        }
        Ok(())
    }

    fn parse_header_image_info(&mut self, entry: &str) -> FrameResult<()> {
        let entries: Vec<&str> = entry.split_ascii_whitespace().collect();
        if entries.len() != Esperanto::HEADER_IMAGE_ENTRIES {
            return Err(FrameError::HeaderError(
                format!("bad header entry: {}", entry).into(),
            ));
        }
        self.dim1 = entries[1].parse::<usize>()?;
        self.dim2 = entries[2].parse::<usize>()?;
        self.compression = match entries[5] {
            Esperanto::HEADER_UNCOMPRESSED => Compression::Uncompressed,
            Esperanto::HEADER_AGI_BITFIELD => Compression::AgiBitfield,
            _ => {
                return Err(FrameError::HeaderError(
                    format!("unsupported compression type: {}", entries[5]).into(),
                ));
            }
        };
        Ok(())
    }

    fn read_data<P: io::Read + io::Seek>(&mut self, mut reader: P) -> FrameResult<()> {
        let mut buf = vec![];
        reader.read_to_end(&mut buf)?;
        match self.compression {
            Compression::Uncompressed => {
                self.array = Array::from_slice_i32::<LittleEndian>(self.dim1, self.dim2, &buf)?
            }
            Compression::AgiBitfield => self.uncompress_agi_bitfield(&buf)?,
        };
        Ok(())
    }

    fn uncompress_agi_bitfield(&mut self, buf: &[u8]) -> FrameResult<()> {
        if buf.len() < size_of::<i32>() {
            return Err(FrameError::FormatError("not enough data to unpack".into()));
        }
        let packed_size: usize = unsafe { std::ptr::read(buf.as_ptr() as *const i32) } as usize;
        if buf.len() < packed_size + self.dim2 * size_of::<i32>() {
            return Err(FrameError::FormatError("not enough data to unpack".into()));
        }
        let packed = &buf[size_of::<i32>()..];
        let addresses =
            unsafe { reinterpret_slice::<u8, i32>(&buf[size_of::<i32>() + packed_size..]) };
        if addresses[0] != 0 {
            return Err(FrameError::FormatError("bad address array".into()));
        }
        self.blocks = (self.dim1 - 1) / (2 * size_of::<u64>());
        self.tail = self.dim1 - self.blocks * 2 * size_of::<u64>() - 1;
        self.array = Array::with_dims(self.dim1, self.dim2);
        unsafe { self.array.set_len() };
        for i in 0..self.dim2 {
            self.i = i * self.dim1;
            self.unpack_row(&packed[addresses[i] as usize..])?;
        }
        Ok(())
    }

    fn unpack_row(&mut self, row: &[u8]) -> FrameResult<()> {
        let mut data = Cursor::new(row);
        let mut of_table = data.clone();
        of_table.set_position(1);
        let mut prev = data.unpack_i32(&mut of_table)?;
        data.set_position(of_table.position());
        self.array.data_mut()[self.i] = prev as f64;
        self.i += 1;
        for _ in 0..self.blocks {
            let block_size = data.read_u8()?;
            let size1 = (block_size & 0xf) as usize;
            let size2 = (block_size >> 0x4) as usize;
            if size1 == 0 || size1 > size_of::<u64>() || size2 == 0 || size2 > size_of::<u64>() {
                return Err(FrameError::FormatError(
                    "garbage in compressed array".into(),
                ));
            }
            of_table.set_position(data.position() + (size1 + size2) as u64);
            self.unpack_block(&mut data, size1, &mut of_table, &mut prev)?;
            self.unpack_block(&mut data, size2, &mut of_table, &mut prev)?;
            data.set_position(of_table.position());
        }
        for _ in 0..self.tail {
            of_table.set_position(data.position() + size_of::<u64>() as u64);
            let next = data.unpack_i32(&mut of_table)? + prev;
            self.array.data_mut()[self.i] = next as f64;
            self.i += 1;
            prev = next;
        }
        Ok(())
    }

    fn unpack_block(
        &mut self,
        row: &mut Cursor<&[u8]>,
        size: usize,
        of_table: &mut Cursor<&[u8]>,
        prev: &mut i32,
    ) -> FrameResult<()> {
        if size == size_of::<u64>() {
            for _ in 0..size_of::<u64>() {
                let next = row.unpack_i32(of_table)? + *prev;
                self.array.data_mut()[self.i] = next as f64;
                self.i += 1;
                *prev = next;
            }
        } else {
            let end = row.position() + size as u64;
            let line = row.read_u64::<LittleEndian>()?;
            row.set_position(end);
            let mut j = 0;
            for _ in 0..size_of::<u64>() {
                let next =
                    (((line & (MASK[size] << j)) >> j) as i32 - MASK[size - 1] as i32) + *prev;
                self.array.data_mut()[self.i] = next as f64;
                self.i += 1;
                *prev = next;
                j += size;
            }
        }
        Ok(())
    }
}

trait UnpackI32 {
    fn unpack_i32(&mut self, of_table: &mut Cursor<&[u8]>) -> FrameResult<i32>;
}

impl UnpackI32 for Cursor<&[u8]> {
    fn unpack_i32(&mut self, of_table: &mut Cursor<&[u8]>) -> FrameResult<i32> {
        let value = self.read_u8()?;
        Ok(match value {
            Esperanto::OVERFLOW_16 => of_table.read_i16::<LittleEndian>()? as i32,
            Esperanto::OVERFLOW_32 => of_table.read_i32::<LittleEndian>()?,
            _ => ((Wrapping(value) - Wrapping(0x7f)).0 as i8) as i32,
        })
    }
}

trait PackInt {
    fn pack8(self, buf: &mut Vec<u8>) -> FrameResult<()>;
}

const I8_MIN: i32 = i8::MIN as i32;
const I8_MAX: i32 = i8::MAX as i32;
const I16_MIN: i32 = i16::MIN as i32;
const I16_MAX: i32 = i16::MAX as i32;

impl PackInt for i32 {
    fn pack8(self, buf: &mut Vec<u8>) -> FrameResult<()> {
        if self > I8_MIN && self < I8_MAX {
            buf.push((self as i8).pack(8) as u8);
        } else if self > I16_MIN && self < I16_MAX {
            buf.push(Esperanto::OVERFLOW_16);
            buf.write_i16::<LittleEndian>(self as i16)?;
        } else {
            buf.push(Esperanto::OVERFLOW_32);
            buf.write_i32::<LittleEndian>(self)?;
        }
        Ok(())
    }
}

trait PackI8 {
    fn pack(self, size: usize) -> i8;
}

impl PackI8 for i8 {
    fn pack(self, size: usize) -> i8 {
        (Wrapping(self) + Wrapping(MASK[size - 1] as i8)).0
    }
}

trait BitSize {
    fn bit_size(self) -> u8;
}

impl BitSize for i32 {
    fn bit_size(self) -> u8 {
        if self < -63 {
            return 8;
        }
        if self < -31 {
            return 7;
        }
        if self < -15 {
            return 6;
        }
        if self < -7 {
            return 5;
        }
        if self < -3 {
            return 4;
        }
        if self < -1 {
            return 3;
        }
        if self < 0 {
            return 2;
        }
        if self < 2 {
            return 1;
        }
        if self < 3 {
            return 2;
        }
        if self < 5 {
            return 3;
        }
        if self < 9 {
            return 4;
        }
        if self < 17 {
            return 5;
        }
        if self < 33 {
            return 6;
        }
        if self < 65 {
            return 7;
        }
        return 8;
    }
}

trait PackI32Block {
    fn pack(&self, size: usize, oft: &mut Vec<u8>) -> FrameResult<u64>;
}

impl PackI32Block for [i32; 8] {
    fn pack(&self, size: usize, oft: &mut Vec<u8>) -> FrameResult<u64> {
        let mut res = 0;
        for (i, value) in self.iter().enumerate() {
            let value = *value;
            let buf = if value > I8_MIN && value < I8_MAX {
                (value as i8).pack(size) as u64
            } else if value > I16_MIN && value < I16_MAX {
                oft.write_i16::<LittleEndian>(value as i16)?;
                Esperanto::OVERFLOW_16 as u64
            } else {
                oft.write_i32::<LittleEndian>(value)?;
                Esperanto::OVERFLOW_32 as u64
            };
            let topos = i * size;
            let newbits = (buf & MASK[size]) << topos;
            let oldbits = res & (!(MASK[size] << topos));
            res = newbits | oldbits;
        }
        Ok(res)
    }
}

#[derive(Template)]
#[template(path = "esp_header.jinja2", escape = "none")]
struct EsperantoHeaderTemplate<'a> {
    dim1: usize,
    dim2: usize,
    datatype: &'a str,
    exposure_time: f64,
    flux: f64,
    pixel1: usize,
    pixel2: usize,
    datetime: String,
    omega: f64,
    domega: f64,
    theta: f64,
    dtheta: f64,
    kappa: f64,
    dkappa: f64,
    phi: f64,
    dphi: f64,
    center_x: f64,
    center_y: f64,
    alpha: f64,
    dist: f64,
    l1: f64,
    l2: f64,
    l12: f64,
    b: f64,
}

#[cfg(test)]
mod tests {
    use crate::esperanto::Esperanto;
    use crate::frame::Frame;
    use std::path::Path;

    static TEST_FILE: &'static str = "testdata/sucros.esperanto";

    fn test_read<P: AsRef<Path>>(path: P) {
        let frame = Esperanto::read_file(path).unwrap();
        assert_eq!(frame.dim1(), 2048, "dim1 != 2048");
        assert_eq!(frame.dim2(), 2048, "dim2 != 2048");
        assert_eq!(frame.sum(), 424962850., "frame sum != 424962850");
        assert_eq!(frame.min(), 86., "frame min != 86");
        assert_eq!(frame.max(), 62570., "frame max != 62570");
        assert_eq!(
            frame.array[frame.dim2() * 1272 + 303],
            58608.,
            "frame[1272, 303] != 58608"
        );
    }

    #[test]
    fn test_esperanto_reader() {
        test_read(TEST_FILE)
    }

    #[test]
    fn test_esperanto_writer() {
        let frame = Esperanto::read_file(TEST_FILE).unwrap();
        let temp = tempfile::NamedTempFile::new().unwrap();
        frame.write_file(temp.path()).unwrap();
        test_read(temp.path());
    }
}