// SPDX-FileCopyrightText: 2022 Thomas Kramer <code@tkramer.ch>
//
// SPDX-License-Identifier: GPL-3.0-or-later

//! Convert the lookup-table to and from a compact binary representation.


use std::io::{Write, Read};
use std::fmt;
use smallvec::SmallVec;
use num_traits::{Num, PrimInt, Signed, Unsigned, FromPrimitive, Zero, One, ToPrimitive};
use tree::*;
use point::*;
use super::*;
use super::marker_types::Canonical;
use super::lut::LookupTable;
use lut::ValuesOrRedirect;
use std::collections::HashMap;

type UInt = u32;
type SInt = i32;

/// Beginning of the LUT file.
const MAGIC: &[u8] = "fast-steiner lookup-table\r\n".as_bytes();

/// Serialize the lookup table.
pub fn write_lut<W: Write>(writer: &mut W, lut: &LookupTable) -> Result<(), LUTWriteError> {

    // Extract all trees.
    let mut trees: Vec<_> = lut.sub_lut_by_degree.iter()
        .flat_map(|sub_lut| sub_lut.values.iter())
        .filter_map(|v| match v {
            ValuesOrRedirect::Values(v) => Some(v),
            ValuesOrRedirect::Redirect { .. } => None
        })
        .flat_map(|v| v.iter())
        .map(|v| &v.potential_optimal_steiner_tree)
        .collect();
    trees.sort_unstable();

    // Reorder trees such that minimal storage is needed when storing only difference between trees.
    println!("Compress {} trees.", trees.len());
    let trees = {
        // let tour = super::travelling_salesperson::solve_tsp(
        //     &trees,
        //     &|a, b| a.symmetric_difference(b).count(),
        // );
        let tour: Vec<_> = (0..trees.len()).collect();

        let trees: Vec<_> = tour.into_iter()
            .map(|i| trees[i])
            .collect();
        trees
    };

    // Table to later find the index of a tree by the tree itself.
    let tree_indices: HashMap<_, _> = trees.iter()
        .enumerate()
        .map(|(i, t)| (*t, i))
        .collect();

    write_trees(writer, &trees)?;

    Ok(())
}

fn write_trees<W: Write>(writer: &mut W, trees: &Vec<&Tree<Canonical>>) -> Result<(), LUTWriteError> {

    // Write number of trees.
    write_unsigned_integer(writer, trees.len())?;

    let empty_tree = Tree::empty_canonical();
    let mut prev_tree = &empty_tree;
    for current_tree in trees {

        let (edge_diff_horiz, edge_diff_vert): (Vec<_>, Vec<_>) = prev_tree.symmetric_difference(current_tree)
            .partition(|e| e.is_horizontal());

        // Number of differing edges.
        write_unsigned_integer(writer, edge_diff_horiz.len())?;
        write_unsigned_integer(writer, edge_diff_vert.len())?;

        // Write differences in horizontal edges.
        for e in edge_diff_horiz {
            write_point_u4(writer, e.start())?;
        }
        // Write differences in vertical edges.
        for e in edge_diff_vert {
            write_point_u4(writer, e.start())?;
        }

        prev_tree = current_tree;
    }

    Ok(())
}

/// Write coordinates of the point as two 4-bit integers packed into a byte.
fn write_point_u4<W: Write>(writer: &mut W, p: Point<HananCoord>) -> Result<(), LUTWriteError> {
    assert!(p.x < 16 && p.y < 16, "edge coordinate out of range for storage");

    let nibbles = ((p.x << 4) | (p.y & 0xf)) as u8;
    writer.write_all(&[nibbles])?;
    Ok(())
}

fn read_point<R: Read>(reader: &mut R) -> Result<Point<HananCoord>, LUTReadError> {
    let nibbles = read_byte(reader)?;
    let x = nibbles >> 4;
    let y = nibbles & 0xf;
    Ok(Point::new(x as HananCoord, y as HananCoord))
}


fn read_byte<R: Read>(reader: &mut R) -> Result<u8, LUTReadError> {
    let mut b = [0u8];
    reader.read_exact(&mut b)?;
    Ok(b[0])
}

/// Convert bit strings into byte strings.
/// This is used only for tests.
#[cfg(test)]
fn bits(bit_string: &str) -> Vec<u8> {
    bit_string.split_ascii_whitespace()
        .map(|s| u8::from_str_radix(s, 2).unwrap())
        .collect()
}

/// Shorthand notation for creating byte strings from the bit representation.
/// This is used only for tests.
#[cfg(test)]
macro_rules! bits {
 ($x:expr) => {&mut bits($x)[..].as_ref()}
}

/// Error during reading of the lookup table.
#[derive(Debug)]
pub enum LUTReadError {
    /// General IO error.
    IOError(std::io::Error),
    /// End of file reached.
    UnexpectedEndOfFile,
    /// Something is wrong with the file format.
    FormatError,
}

impl fmt::Display for LUTReadError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        use LUTReadError::*;
        match self {
            LUTReadError::FormatError => write!(f, "illegal table format"),
            other => other.fmt(f)
        }
    }
}

/// Error types for writing LUT.
#[derive(Debug)]
pub enum LUTWriteError {
    /// General IO error.
    IOError(std::io::Error),
}

impl fmt::Display for LUTWriteError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        use LUTWriteError::*;
        match self {
            other => other.fmt(f)
        }
    }
}

impl From<std::io::Error> for LUTReadError {
    fn from(err: std::io::Error) -> Self {
        match err.kind() {
            std::io::ErrorKind::UnexpectedEof => LUTReadError::UnexpectedEndOfFile,
            _ => LUTReadError::IOError(err)
        }
    }
}

impl From<std::io::Error> for LUTWriteError {
    fn from(err: std::io::Error) -> Self {
        match err.kind() {
            _ => LUTWriteError::IOError(err)
        }
    }
}

/// Read the magic string at the beginning of the LUT file.
/// Returns an error if the magic value is wrong.
fn read_magic<R: Read>(reader: &mut R) -> Result<(), LUTReadError> {
    // Read the magic string.
    let mut buf = [0u8; MAGIC.len()];
    reader.read_exact(&mut buf)?;

    // Compare the buffer with the magic string.
    match buf.as_ref() {
        MAGIC => Ok(()), // Matic string is correct.
        _ => {
            let s = String::from_utf8_lossy(&buf);
            log::error!("wrong file format, magic string is wrong: '{:?}', hex = '{:02x?}'", s, buf);
            Err(LUTReadError::FormatError)
        } // Magic string is wrong.
    }
}

/// Write the magic string at the beginning of the LUT file.
fn write_magic<W: Write>(writer: &mut W) -> Result<(), LUTWriteError> {
    writer.write_all(MAGIC)?;
    Ok(())
}

/// Read an unsigned integer.
/// Unsigned integers are stored as variable-length byte continuations.
/// The all but last bytes of the continuation have the MSB set to one.
/// The result is the concatenation of the 7 lower-value bits.
/// The low-order byte appears first.
fn read_unsigned_integer<T, R: Read>(reader: &mut R) -> Result<T, LUTReadError>
    where T: PrimInt + Zero + FromPrimitive {
    // Read byte continuation until the most significant bit is zero.
    let mut bytes = <SmallVec<[u8; 8]>>::new();
    for b in reader.bytes() {
        let b = b?;
        bytes.push(b);
        if b & 0x80 == 0 {
            break;
        }
    }
    if bytes.last().map(|b| b & 0x80) != Some(0) { // Last byte must have MSB set to zero.
        // Reached EOF.
        Err(LUTReadError::UnexpectedEndOfFile)
    } else {
        Ok(
            bytes.iter()
                .rev()
                .map(|b| b & 0x7f) // Set MSB to zero.
                .fold(T::zero(), |acc, b| T::from_u8(b).unwrap() + (acc << 7))
        )
    }
}


#[test]
fn test_read_unsigned_integer() {
    // Test values from LUT specification draft.
    assert_eq!(read_unsigned_integer::<u32, _>(bits!("00000000")).unwrap(), 0);
    assert_eq!(read_unsigned_integer::<i32, _>(bits!("01111111")).unwrap(), 127);
    assert_eq!(read_unsigned_integer::<i32, _>(bits!("10000000 00000001")).unwrap(), 128);
    assert_eq!(read_unsigned_integer::<i32, _>(bits!("11111111 01111111")).unwrap(), 16383);
    assert_eq!(read_unsigned_integer::<i32, _>(bits!("10000000 10000000 00000001")).unwrap(), 16384);

    // Empty bytes should raise an error.
    assert!(read_unsigned_integer::<i32, _>(bits!("")).is_err());
    // The last byte must always have the MSB set to zero.
    assert!(read_unsigned_integer::<i32, _>(bits!("10000000")).is_err());
}

/// Write an unsigned integer.
/// Unsigned integers are stored as variable-length byte continuations.
/// The all but last bytes of the continuation have the MSB set to one.
/// The result is the concatenation of the 7 lower-value bits.
/// The low-order byte appears first.
fn write_unsigned_integer<T, W: Write>(writer: &mut W, value: T) -> Result<(), LUTWriteError>
    where T: PrimInt + Unsigned + FromPrimitive + ToPrimitive {
    let mut value = value;

    let mask = T::from_u8(0x7f).unwrap();
    let _0x80 = T::from_u8(0x80).unwrap();
    while value > mask {
        let lowest = (value & mask).to_u8().unwrap();
        writer.write_all(&[0x80 | lowest])?; // Set the MSB.
        value = value.shr(7);
    }
    debug_assert!(value & _0x80 == T::zero());
    writer.write_all(&[value.to_u8().unwrap()])?;
    Ok(())
}


#[test]
fn test_write_unsigned_integer() {
    // Test values from LUT specification draft.

    fn test(num: UInt, expected: Vec<u8>) {
        let mut buf = Vec::new();
        write_unsigned_integer(&mut buf, num).unwrap();
        assert_eq!(buf, expected);
    }

    test(0, vec![0x00]);
    test(1, vec![0x01]);
    test(127, vec![0x7f]);
    test(128, vec![0x80, 0x01]);
    test(16383, vec![0xff, 0x7f]);
    test(16384, vec![0x80, 0x80, 0x01]);
}


/// Read a signed integer.
/// The format is the same as for unsigned integers except that the sign bit
/// is stored in the LSB of the lowest-order byte.
fn read_signed_integer<T, R: Read>(reader: &mut R) -> Result<T, LUTReadError>
    where T: PrimInt + Num + Zero + One + Signed + FromPrimitive {
    let u: T = read_unsigned_integer(reader)?;
    let sign = u & One::one();

    let magnitude = u.unsigned_shr(1);

    Ok(
        if sign.is_one() {
            T::zero() - magnitude
        } else {
            magnitude
        }
    )
}

#[test]
fn test_read_signed_integer() {
    // Test values from LUT specification draft.
    assert_eq!(read_signed_integer::<i32, _>(bits!("00")).unwrap(), 0);
    assert_eq!(read_signed_integer::<i32, _>(bits!("10")).unwrap(), 1);
    assert_eq!(read_signed_integer::<i32, _>(bits!("11")).unwrap(), -1);
    assert_eq!(read_signed_integer::<i32, _>(bits!("01111110")).unwrap(), 63);
    assert_eq!(read_signed_integer::<i32, _>(bits!("10000001 00000001")).unwrap(), -64);
    assert_eq!(read_signed_integer::<i32, _>(bits!("11111110 01111111")).unwrap(), 8191);
    assert_eq!(read_signed_integer::<i32, _>(bits!("10000001 10000000 00000001")).unwrap(), -8192);

    // Empty bytes should raise an error.
    assert!(read_signed_integer::<i32, _>(bits!("")).is_err());
    // The last byte must always have the MSB set to zero.
    assert!(read_signed_integer::<i32, _>(bits!("10000000")).is_err());
}

/// Write a signed integer.
/// The format is the same as for unsigned integers except that the sign bit
/// is stored in the LSB of the lowest-order byte.
fn write_signed_integer<W: Write>(writer: &mut W, value: SInt) -> Result<(), LUTWriteError> {
    // Determine the sign bit.
    let (sign, value) = if value < 0 {
        (1, -value)
    } else {
        (0, value)
    };

    // Insert the sign bit at the lowest value bit position.
    let magnitude = (value << 1) as UInt;
    let u = magnitude | sign;

    write_unsigned_integer(writer, u)
}

#[test]
fn test_write_signed_integer() {
    // Test values from LUT specification draft.

    fn test(num: SInt, expected: Vec<u8>) {
        let mut buf = Vec::new();
        write_signed_integer(&mut buf, num).unwrap();
        assert_eq!(buf, expected);
    }

    test(0, vec![0x00]);
    test(1, vec![0b10]);
    test(-1, vec![0b11]);
    test(63, vec![0b01111110]);
    test(-64, vec![0x81, 0x01]);
    test(8191, vec![0b11111110, 0b01111111]);
    test(-8192, vec![0b10000001, 0b10000000, 0b1]);
}

#[test]
fn test_write_lut() {
    let lut = super::gen_lut::gen_full_lut(4);

    let mut buffer = Vec::new();

    let write_result = write_lut(&mut buffer, &lut);
    assert!(write_result.is_ok());
    dbg!(buffer.len());
    assert!(buffer.len() < 100);
}