use crate::{
    coder::Coder,
    widemuldiv::{ptngc_largeint_add, ptngc_largeint_div, ptngc_largeint_mul},
    xtc3::{self, positive_int},
};
use log::{self, debug};

// Difference in indices used for determining whether to store as large or small
const QUITE_LARGE: u32 = 3;
const IS_LARGE: u32 = 6;

// This magic value is generated by max_magic.py in the original c code
pub(crate) const MAX_MAGIC: usize = 92;

#[rustfmt::skip]
const MAGIC: [u32; MAX_MAGIC] = [
    2,          3,         4,          5,          6,          8,          10,
    12,         16,        20,         25,         32,         40,         50,
    64,         80,        101,        128,        161,        203,        256,
    322,        406,       512,        645,        812,        1024,       1290,
    1625,       2048,      2580,       3250,       4096,       5160,       6501,
    8192,       10321,     13003,      16384,      20642,      26007,      32768,
    41_285,      52_015,     65_536,      82_570,      104_031,     131_072,     165_140,
    208_063,     262_144,    330_280,     416_127,     524_288,     660_561,     832_255,
    1_048_576,   1_321_122,  1_664_510,   2_097_152,   2_642_245,   3_329_021,   4_194_304,
    5_284_491,   6_658_042,  8_388_608,   10_568_983,  13_316_085,  16_777_216,  21_137_967,
    26_632_170,  33_554_432, 42_275_935,  53_264_340,  67_108_864,  84_551_870,  106_528_681,
    134_217_728, 169_103_740, 213_057_362, 268_435_456, 338_207_481, 426_114_725, 536_870_912,
    676_414_963, 852_229_450, 1_073_741_824, 1_352_829_926, 1_704_458_900, 2_147_483_648, 2_705_659_852,
    3_408_917_801,
];

pub(crate) const MAGIC_BITS: [[u32; 8]; MAX_MAGIC] = [
    [3, 6, 9, 12, 15, 18, 21, 24],
    [5, 10, 15, 20, 24, 29, 34, 39],
    [6, 12, 18, 24, 30, 36, 42, 48],
    [7, 14, 21, 28, 35, 42, 49, 56],
    [8, 16, 24, 32, 39, 47, 55, 63],
    [9, 18, 27, 36, 45, 54, 63, 72],
    [10, 20, 30, 40, 50, 60, 70, 80],
    [11, 22, 33, 44, 54, 65, 76, 87],
    [12, 24, 36, 48, 60, 72, 84, 97],
    [13, 26, 39, 52, 65, 78, 91, 104],
    [14, 28, 42, 56, 70, 84, 98, 112],
    [15, 30, 45, 60, 75, 90, 105, 120],
    [16, 32, 48, 64, 80, 96, 112, 128],
    [17, 34, 51, 68, 85, 102, 119, 136],
    [18, 36, 54, 72, 90, 108, 127, 144],
    [19, 38, 57, 76, 95, 114, 133, 152],
    [20, 40, 60, 80, 100, 120, 140, 160],
    [21, 42, 63, 84, 105, 127, 147, 168],
    [22, 44, 66, 88, 110, 132, 154, 176],
    [23, 46, 69, 92, 115, 138, 161, 184],
    [24, 48, 72, 97, 120, 144, 168, 192],
    [25, 50, 75, 100, 125, 150, 175, 200],
    [26, 52, 78, 104, 130, 156, 182, 208],
    [27, 54, 81, 108, 135, 162, 190, 216],
    [28, 56, 84, 112, 140, 168, 196, 224],
    [29, 58, 87, 116, 145, 174, 203, 232],
    [30, 60, 90, 120, 150, 180, 210, 240],
    [31, 62, 93, 124, 155, 186, 217, 248],
    [32, 64, 96, 128, 160, 192, 224, 256],
    [33, 66, 99, 132, 165, 198, 231, 264],
    [34, 68, 102, 136, 170, 204, 238, 272],
    [35, 70, 105, 140, 175, 210, 245, 280],
    [36, 72, 108, 144, 180, 216, 252, 288],
    [37, 74, 111, 148, 185, 222, 259, 296],
    [38, 76, 114, 152, 190, 228, 266, 304],
    [39, 78, 117, 157, 195, 234, 273, 312],
    [40, 80, 120, 160, 200, 240, 280, 320],
    [41, 82, 123, 164, 205, 246, 287, 328],
    [42, 84, 127, 168, 210, 252, 294, 336],
    [43, 86, 129, 172, 215, 258, 301, 344],
    [44, 88, 132, 176, 220, 264, 308, 352],
    [45, 90, 135, 180, 225, 270, 315, 360],
    [46, 92, 138, 184, 230, 276, 322, 368],
    [47, 94, 141, 188, 235, 282, 329, 376],
    [48, 97, 144, 192, 240, 288, 336, 384],
    [49, 98, 147, 196, 245, 294, 343, 392],
    [50, 100, 150, 200, 250, 300, 350, 400],
    [52, 102, 153, 204, 255, 306, 357, 408],
    [52, 104, 156, 208, 260, 312, 364, 416],
    [53, 106, 159, 212, 265, 318, 371, 424],
    [54, 108, 162, 216, 270, 324, 378, 432],
    [55, 110, 165, 220, 275, 330, 385, 440],
    [56, 112, 168, 224, 280, 336, 392, 448],
    [57, 114, 172, 228, 285, 342, 399, 456],
    [58, 116, 174, 232, 290, 348, 406, 464],
    [59, 118, 177, 236, 295, 354, 413, 472],
    [60, 120, 180, 240, 300, 360, 420, 480],
    [61, 122, 183, 244, 305, 366, 427, 488],
    [62, 124, 186, 248, 310, 372, 434, 496],
    [63, 127, 190, 252, 315, 378, 442, 505],
    [64, 128, 192, 256, 320, 384, 448, 512],
    [65, 130, 195, 260, 325, 390, 455, 520],
    [66, 132, 198, 264, 330, 396, 462, 528],
    [67, 134, 201, 268, 335, 402, 469, 536],
    [68, 136, 204, 272, 340, 408, 476, 544],
    [69, 138, 207, 276, 345, 414, 483, 552],
    [70, 140, 210, 280, 350, 420, 490, 560],
    [71, 142, 213, 284, 355, 426, 497, 568],
    [72, 144, 216, 288, 360, 432, 505, 576],
    [73, 146, 219, 292, 365, 438, 511, 584],
    [74, 148, 222, 296, 370, 444, 518, 592],
    [75, 150, 225, 300, 375, 451, 525, 600],
    [76, 152, 228, 304, 380, 456, 532, 608],
    [77, 154, 231, 308, 385, 462, 539, 616],
    [78, 157, 234, 312, 390, 469, 546, 625],
    [79, 158, 237, 316, 395, 474, 553, 632],
    [80, 160, 240, 320, 400, 480, 560, 640],
    [81, 162, 243, 324, 406, 486, 568, 648],
    [82, 164, 246, 328, 410, 492, 574, 656],
    [83, 166, 249, 332, 415, 498, 581, 664],
    [84, 168, 252, 336, 420, 505, 588, 672],
    [85, 170, 255, 340, 425, 510, 595, 680],
    [86, 172, 258, 344, 430, 516, 602, 688],
    [87, 174, 261, 348, 435, 522, 609, 696],
    [88, 176, 264, 352, 440, 528, 616, 704],
    [89, 178, 267, 356, 445, 534, 623, 712],
    [90, 180, 270, 360, 451, 540, 631, 720],
    [91, 182, 273, 364, 455, 546, 637, 728],
    [92, 184, 276, 368, 460, 552, 644, 736],
    [94, 187, 279, 373, 466, 558, 651, 745],
    [94, 188, 282, 376, 470, 564, 658, 752],
    [95, 190, 285, 380, 475, 570, 665, 760],
];

/* Sequence instructions */
pub(crate) const INSTR_DEFAULT: i32 = 0;
pub(crate) const INSTR_BASE_RUNLENGTH: i32 = 1;
pub(crate) const INSTR_ONLY_LARGE: i32 = 2;
pub(crate) const INSTR_ONLY_SMALL: i32 = 3;
pub(crate) const INSTR_LARGE_BASE_CHANGE: i32 = 4;
pub(crate) const INSTR_FLIP: i32 = 5;
pub(crate) const INSTR_LARGE_RLE: i32 = 6;

const MAXINSTR: usize = 7;

// TODO: enable compile-time feature
pub(crate) const INSTRNAMES: [&str; MAXINSTR] = [
    "large+small",
    "base+run",
    "large",
    "small",
    "large base change",
    "flip",
    "large rle",
];

const SEQ_INSTR: [[u32; 2]; MAXINSTR] = [
    [1, 1], /* 1 : one large atom + runlength encoded small integers. Use same settings as before. */
    [0, 2], /* 00 : set base and runlength in next four bits (x). base (increase/keep/decrease)=x%3-1. runlength=1+x/3.
            The special value 1111 in the four bits means runlength=6 and base change=0 */
    [4, 4], /* 0100 : next only a large atom comes. */
    [5, 4], /* 0101 : next only runlength encoded small integers. Use same settings as before. */
    [6, 4], /* 0110 : Large change in base. Change is encoded in the
            following 2 bits. change direction (sign) is the first
            bit. The next bit +1 is the actual change. This
            allows the change of up to +/- 2 indices. */
    [14, 5], /* 01110 : flip whether integers should be modified to compress water better */
    [15, 5], /* 01111 : Large rle. The next 4 bits encode how many
             large atoms are in the following sequence: 3-18. (2 is
             more efficiently coded with two large instructions. */
];

pub(crate) fn readbits(ptr: &mut &[u8], bitptr: &mut i32, nbits: i32) -> u32 {
    let mut nbits = nbits;
    let mut val = 0u32;
    let mut extract_mask = 0x80_u32 >> *bitptr;
    let mut thisval = ptr[0];

    debug!("Read nbits={nbits}");

    while nbits != 0 {
        nbits -= 1;
        val <<= 1;
        val |= u32::from((extract_mask & u32::from(thisval)) != 0);
        *bitptr += 1;
        extract_mask >>= 1;
        if extract_mask == 0 {
            extract_mask = 0x80_u32;
            *ptr = &ptr[1..];
            *bitptr = 0;
            if nbits != 0 {
                thisval = ptr[0];
            }
        }
    }

    debug!("val={val}");

    val
}

pub(crate) fn readmanybits(ptr: &mut &[u8], bitptr: &mut i32, mut nbits: i32, buffer: &mut [u8]) {
    let mut buf_idx = 0;
    while nbits >= 8 {
        buffer[buf_idx] = readbits(ptr, bitptr, 8) as u8;
        debug!("Read value {:02x}", buffer[buf_idx]);
        buf_idx += 1;
        nbits -= 8;
    }
    if nbits != 0 {
        buffer[buf_idx] = readbits(ptr, bitptr, nbits) as u8;
        debug!("Read value {:02x}", buffer[buf_idx]);
    }
}

pub(crate) const fn ptngc_magic(i: usize) -> u32 {
    MAGIC[i]
}

pub(crate) fn ptngc_find_magic_index(maxval: u32) -> u32 {
    // Binary search: find the first index where MAGIC[i] > maxval.
    // MAGIC is sorted ascending, so partition_point gives us exactly what the
    // linear scan computed (the first i where MAGIC[i] > maxval), clamped to MAX_MAGIC-1
    // to avoid OOB (matching the original fix for MAGIC[91] == maxval edge case).
    let i = MAGIC.partition_point(|&m| m <= maxval);
    i.min(MAX_MAGIC - 1) as u32
}

#[allow(clippy::too_many_arguments)]
fn buffer_large(
    coder: &mut Coder,
    has_large: &mut usize,
    has_large_ints: &mut [i32],
    new_large_ints: &[i32],
    large_index: &[u32],
    large_nbits: u32,
    compress_buffer: &mut [u8],
    output: &mut Vec<u8>,
) {
    // If it is full we must write them all
    if *has_large == 18 {
        flush_large(
            coder,
            has_large,
            has_large_ints,
            *has_large,
            large_index,
            large_nbits,
            compress_buffer,
            output,
        ); // Flush them all
    }
    has_large_ints[*has_large * 3] = new_large_ints[0];
    has_large_ints[*has_large * 3 + 1] = new_large_ints[1];
    has_large_ints[*has_large * 3 + 2] = new_large_ints[2];
    *has_large += 1;
}

#[allow(clippy::too_many_arguments)]
fn flush_large(
    coder: &mut Coder,
    has_large: &mut usize,
    has_large_ints: &mut [i32],
    n: usize,
    large_index: &[u32],
    large_nbits: u32,
    compress_buffer: &mut [u8],
    output: &mut Vec<u8>,
) {
    if n < 3 {
        for i in 0..n {
            write_instruction(coder, INSTR_ONLY_LARGE, output);
            write_three_large(
                coder,
                &has_large_ints[i * 3..],
                large_index,
                large_nbits,
                compress_buffer,
                output,
            );
        }
    } else {
        write_instruction(coder, INSTR_LARGE_RLE, output);
        coder.ptngc_writebits(u32::try_from(n - 3).expect("u32 from usize"), 4, output); // Number of large atoms in this sequence: 3 to 18
        for i in 0..n {
            write_three_large(
                coder,
                &has_large_ints[i * 3..],
                large_index,
                large_nbits,
                compress_buffer,
                output,
            );
        }
    }
    if *has_large - n != 0 {
        for i in 0..(*has_large - n) {
            for j in 0..3 {
                has_large_ints[i * 3 + j] = has_large_ints[(i + n) * 3 + j];
            }
        }
    }
    *has_large -= n; // Number of remaining large atoms in buffer
}

fn write_three_large(
    coder: &mut Coder,
    encode_ints: &[i32],
    large_index: &[u32],
    nbits: u32,
    compress_buffer: &mut [u8],
    output: &mut Vec<u8>,
) {
    trajcoder_base_compress(encode_ints, 3, large_index, compress_buffer);
    coder.ptngc_write_many_bits(compress_buffer, nbits, output);
    debug!("{:?}", &compress_buffer[..10]);
    debug!("nbits={nbits} ({})", f64::from(nbits) / 3.);
    debug!("Large: {:?}", &encode_ints[..3]);
}

// Convert a sequence of (hopefully) small positive integers using the base pointed to
// by index (x base, y base and z base can be different)
// The largest number of integers supported is 18 (29 to handle/detect overflow)
fn trajcoder_base_compress(input: &[i32], n: usize, index: &[u32], result: &mut [u8]) {
    let mut largeint = [0; 19];
    let mut largeint_tmp = [0; 19];

    if n > 0 {
        // we're explicit about this conversion whereas the c version just... prays
        ptngc_largeint_add(input[0] as u32, &mut largeint, 19);
    }

    for i in 1..n {
        // We must do the multiplication of the largeint with the integer base
        ptngc_largeint_mul(
            MAGIC[usize::try_from(index[i % 3]).expect("usize from u32")],
            &largeint,
            &mut largeint_tmp,
            19,
        );
        largeint.copy_from_slice(&largeint_tmp);
        ptngc_largeint_add(input[i] as u32, &mut largeint, 19);
    }

    assert!(largeint[18] == 0, "BUG! Overflow in compression detected!");

    // Convert the largeint to a sequence of bytes
    for i in 0..18 {
        result[i * 4..i * 4 + 4].copy_from_slice(&largeint[i].to_le_bytes());
    }
}

/// The opposite of [`base_compress`]
pub(crate) fn trajcoder_base_decompress(input: &[u8], n: i32, index: &[u32], output: &mut [i32]) {
    // Convert the sequence of bytes to a largeint
    let mut largeint = [0; 19];
    let mut largeint_tmp = [0; 19];
    for i in 0..18 {
        largeint[i] = u32::from_le_bytes([
            input[i * 4],
            input[i * 4 + 1],
            input[i * 4 + 2],
            input[i * 4 + 3],
        ]);
    }
    largeint[18] = 0;
    debug!("Largeint");
    debug!("{largeint:?}");

    for i in (0..n).rev() {
        let remainder = ptngc_largeint_div(
            MAGIC[index[(i % 3) as usize] as usize],
            &largeint,
            &mut largeint_tmp,
            19,
        );
        debug!("Remainder: {remainder}");

        largeint = largeint_tmp;
        output[i as usize] = remainder as i32;
    }
}

fn write_instruction(coder: &mut Coder, instr: i32, output: &mut Vec<u8>) {
    debug!(
        "INSTR: {} ({} bits)",
        INSTRNAMES[instr as usize], SEQ_INSTR[instr as usize][1]
    );
    coder.ptngc_writebits(
        SEQ_INSTR[instr as usize][0],
        SEQ_INSTR[instr as usize][1],
        output,
    );
}

pub(crate) fn read_instruction(ptr: &mut &[u8], bitptr: &mut i32) -> i32 {
    let mut bits = readbits(ptr, bitptr, 1);

    let instr;
    if bits != 0 {
        instr = INSTR_DEFAULT;
    } else {
        bits = readbits(ptr, bitptr, 1);
        if bits == 0 {
            instr = INSTR_BASE_RUNLENGTH;
        } else {
            bits = readbits(ptr, bitptr, 2);
            instr = match bits {
                0 => INSTR_ONLY_LARGE,
                1 => INSTR_ONLY_SMALL,
                2 => INSTR_LARGE_BASE_CHANGE,
                3 => {
                    bits = readbits(ptr, bitptr, 1);
                    if bits == 0 {
                        INSTR_FLIP
                    } else {
                        INSTR_LARGE_RLE
                    }
                }
                _ => {
                    debug!("bits was unmatched with instruction. Got bits={bits}");
                    -1
                }
            };
        }
    }
    instr
}

pub(crate) fn ptngc_pack_array_xtc2(
    coder: &mut Coder,
    input: &mut [i32],
    length: &mut usize,
) -> (Vec<u8>, usize) {
    let ntriplets = *length / 3;
    let mut maxint = [input[0], input[1], input[2]];
    let mut minint = [input[0], input[1], input[2]];
    let mut output = Vec::with_capacity(8 * *length);
    let mut has_large = 0;
    let mut has_large_ints = [0; 54]; // Large cache. Up to 18 large atoms
    let mut runlength = 0; // Initial runlength. "Stupidly" set to zero for simplicity and explicity
    let mut swapatoms = false; // Initial guess is that we should not swap the first two atoms in each large+small transition
    let mut didswap; // Whether swapping was actually done.
    let mut encode_ints = [0; 21]; // Up to 3 large + 18 small ints can be encoded at once
    let mut nencode;
    let mut compress_buffer = [0; 18 * 4]; // Holds compressed result for 3 large ints or up to 18 small ints.
    let mut refused = 0;
    let mut input_ptr = 0;

    for i in 1..ntriplets {
        for j in 0..3 {
            if input[i * 3 + j] > maxint[j] {
                maxint[j] = input[i * 3 + j];
            }
            if input[i * 3 + j] < minint[j] {
                minint[j] = input[i * 3 + j];
            }
        }
    }

    let large_index = [
        ptngc_find_magic_index((maxint[0].wrapping_sub(minint[0]).wrapping_add(1)) as u32),
        ptngc_find_magic_index((maxint[1].wrapping_sub(minint[1]).wrapping_add(1)) as u32),
        ptngc_find_magic_index((maxint[2].wrapping_sub(minint[2]).wrapping_add(1)) as u32),
    ];
    let large_nbits = compute_magic_bits(large_index);
    let max_large_index = *large_index.iter().max().unwrap();

    for j in 0..3 {
        debug!(
            "minint[{}]={}. maxint[{}]={} large_index[{}]={} value={}",
            j, minint[j], j, maxint[j], j, large_index[j], MAGIC[large_index[j] as usize]
        );
    }
    debug!("large_nbits={large_nbits}");

    // Guess the initial small index
    let mut small_index = max_large_index / 2;

    // Find the largest value that is not large. Not large is half index of large
    let max_small = MAGIC[usize::try_from(small_index).expect("usize from u32")];
    let mut intmax = 0;
    for item in &input[..*length] {
        let s = positive_int(*item);
        if s > intmax && s < max_small {
            intmax = s;
        }
    }
    // This value is not critical, since if I guess wrong, the code will
    // just insert instructions to increase this value
    small_index = ptngc_find_magic_index(intmax);
    debug!(
        "initial_small intmax={intmax},. small_index={small_index}, value={}",
        MAGIC[small_index as usize]
    );

    // Store min integers
    coder.pack_temporary_bits = 32;
    coder.pack_temporary = positive_int(minint[0]);
    coder.out8bits(&mut output);
    coder.pack_temporary_bits = 32;
    coder.pack_temporary = positive_int(minint[1]);
    coder.out8bits(&mut output);
    coder.pack_temporary_bits = 32;
    coder.pack_temporary = positive_int(minint[2]);
    coder.out8bits(&mut output);
    // Store max indices
    coder.pack_temporary_bits = 8;
    coder.pack_temporary = large_index[0];
    coder.out8bits(&mut output);
    coder.pack_temporary_bits = 8;
    coder.pack_temporary = large_index[1];
    coder.out8bits(&mut output);
    coder.pack_temporary_bits = 8;
    coder.pack_temporary = large_index[2];
    coder.out8bits(&mut output);
    // Store initial small index
    coder.pack_temporary_bits = 8;
    coder.pack_temporary = small_index;
    coder.out8bits(&mut output);

    // Initial prevcoord is the minimum integers
    let mut prevcoord = minint;

    let mut ntriplets_left = ntriplets;
    while ntriplets_left > 0 {
        // If less than three atoms left we just write them all as large integers. Here no swapping is done!
        if ntriplets_left < 3 {
            // c code: line 1755
            // TODO: I really think the `ienc` loop can only run once
            let ienc = 0;
            for jenc in 0..3 {
                encode_ints[jenc] = input[input_ptr..][ienc * 3 + jenc].wrapping_sub(minint[jenc]);
            }
            buffer_large(
                coder,
                &mut has_large,
                &mut has_large_ints,
                &encode_ints,
                &large_index,
                large_nbits,
                &mut compress_buffer,
                &mut output,
            );
            input_ptr += 3;
            ntriplets_left -= 1;
            let n = has_large;
            flush_large(
                coder,
                &mut has_large,
                &mut has_large_ints,
                n,
                &large_index,
                large_nbits,
                &mut compress_buffer,
                &mut output,
            );
        } else {
            let mut min_runlength = 0;
            let mut largest_required_base;
            let mut largest_runlength_base;

            let mut largest_runlength_index;
            let mut new_runlength;
            let mut new_small_index;
            let mut iter_runlength;
            let mut iter_small_index;
            let mut rle_index_dep;
            didswap = false;
            debug!("Initial batch");
            // Insert the next batch of integers to be encoded into the buffer
            nencode = insert_batch(
                &input[input_ptr..],
                ntriplets_left,
                &prevcoord,
                &minint,
                &mut encode_ints,
                0,
            );
            //First we must decide if the next value is large (does not reasonably fit in current small encoding)
            // Also, if we have not written any values yet, we must begin by writing a large atom.
            if xtc3::is_quite_large(&encode_ints, small_index, max_large_index) || refused > 0 {
                // If any of the next two atoms are large we should probably write them as large and not swap them
                let mut no_swap = false;
                if xtc3::is_quite_large(&encode_ints[3..], small_index, max_large_index)
                    || xtc3::is_quite_large(&encode_ints[6..], small_index, max_large_index)
                {
                    no_swap = true;
                }
                if !no_swap {
                    // Next we must decide if we should swap the first two values
                    swapdecide(
                        coder,
                        &input[input_ptr..],
                        &mut swapatoms,
                        &minint,
                        &mut output,
                    );
                    // If we should do the integer swapping manipulation we should do it now
                    if swapatoms {
                        didswap = true;
                        for i in 0..3 {
                            let inp = [
                                input[input_ptr..][i].wrapping_sub(minint[i]),
                                input[input_ptr..][3 + i].wrapping_sub(input[input_ptr..][i]), // minint[i]-minint[i] cancels out
                                input[input_ptr..][6 + i].wrapping_sub(input[input_ptr..][3 + i]),
                            ];
                            let mut out = [0; 3];
                            xtc3::swap_ints(&inp, &mut out);
                            encode_ints[i] = out[0];
                            encode_ints[3 + i] = out[1];
                            encode_ints[6 + i] = out[2];
                        }
                        // We have swapped atoms, so the minimum run-length is 2
                        if log::log_enabled!(log::Level::Debug) {
                            debug!("Swap atoms results in:");
                            for i in 0..3 {
                                debug!(
                                    "{}: {:6} {:6} {:6}\t\t{:6} {:6} {:6}",
                                    i * 3,
                                    encode_ints[i * 3],
                                    encode_ints[i * 3 + 1],
                                    encode_ints[i * 3 + 2],
                                    positive_int(encode_ints[i * 3]),
                                    positive_int(encode_ints[i * 3 + 1]),
                                    positive_int(encode_ints[i * 3 + 2]),
                                );
                            }
                        }
                        min_runlength = 2;
                    }
                }
                if swapatoms && didswap {
                    prevcoord.copy_from_slice(&encode_ints[..3]);
                } else {
                    for ienc in 0..3 {
                        prevcoord[ienc] = input[input_ptr..][ienc].wrapping_sub(minint[ienc]);
                    }
                }
                // Cache large value for later possible combination with
                // a sequence of small integers.
                buffer_large(
                    coder,
                    &mut has_large,
                    &mut has_large_ints,
                    &prevcoord,
                    &large_index,
                    large_nbits,
                    &mut compress_buffer,
                    &mut output,
                );
                debug!("Prevcoord after packing of large: {prevcoord:?}");
                // We have written a large integer so we have one less atoms to worry about
                input_ptr += 3;
                ntriplets_left -= 1;
                refused = 0;

                // Insert the next batch of integers to be encoded into the buffer
                debug!("Update batch due to large int");
                if swapatoms && didswap {
                    // Keep swapped values
                    for i in 0..2 {
                        for ienc in 0..3 {
                            encode_ints[i * 3 + ienc] = encode_ints[(i + 1) * 3 + ienc];
                        }
                    }
                }
                nencode = insert_batch(
                    &input[input_ptr..],
                    ntriplets_left,
                    &prevcoord,
                    &minint,
                    &mut encode_ints,
                    min_runlength,
                );
            }
            // Here we should only have differences for the atom coordinates.
            // Convert the ints to positive ints
            for val in encode_ints.iter_mut().take(nencode) {
                // Match the C encoder, which stores the `positive_int` result in `int`
                // and therefore wraps when the value exceeds `i32::MAX`.
                *val = positive_int(*val) as i32;
            }

            // Now we must decide what base and runlength to do. If we have swapped atoms it will be
            // at least 2. If even the next atom is large, we will not do anything.
            largest_required_base = 0;
            // Determine required base
            for &item in encode_ints.iter().take(min_runlength * 3) {
                if item > largest_required_base {
                    largest_required_base = item;
                }
            }

            // Also compute what the largest base is for the current runlength setting!
            largest_runlength_base = 0;
            for &item in encode_ints.iter().take((runlength * 3).min(nencode)) {
                if item > largest_runlength_base {
                    largest_runlength_base = item;
                }
            }

            let largest_required_index = ptngc_find_magic_index(largest_required_base as u32);
            largest_runlength_index = ptngc_find_magic_index(largest_runlength_base as u32);

            if largest_required_index < largest_runlength_index {
                new_runlength = min_runlength;
                new_small_index = largest_required_index;
            } else {
                new_runlength = runlength;
                new_small_index = largest_runlength_index;
            }

            // Only allow increase of runlength wrt min_runlength
            if new_runlength < min_runlength {
                new_runlength = min_runlength;
            }

            // If the current runlength is longer than the number of
            // triplets left stop it from being so.
            if new_runlength > ntriplets_left {
                new_runlength = ntriplets_left;
            }

            // We must at least try to get some small integers going.
            if new_runlength == 0 {
                new_runlength = 1;
                new_small_index = small_index;
            }

            iter_runlength = new_runlength;
            iter_small_index = new_small_index;

            // Iterate to find optimal encoding and runlength
            debug!("Entering iterative loop");
            loop {
                new_runlength = iter_runlength;
                new_small_index = iter_small_index;

                debug!(
                    "Test new_small_index={new_small_index} Base={}",
                    MAGIC[new_small_index as usize]
                );

                // What is the largest runlength we can do with the currently selected encoding?
                // Also the max supported runlength is 6 triplets!
                let mut ienc = 0;
                while ienc < nencode && ienc < 18 {
                    let test_index = ptngc_find_magic_index(encode_ints[ienc] as u32);
                    if test_index > new_small_index {
                        break;
                    }
                    ienc += 1;
                }
                if ienc / 3 > new_runlength {
                    iter_runlength = ienc / 3;
                    debug!("I found a new possible runlength: {iter_runlength}");
                }

                // How large encoding do we have to use?
                largest_runlength_base = 0;
                for &item in encode_ints.iter().take(iter_runlength * 3) {
                    if item > largest_runlength_base {
                        largest_runlength_base = item;
                    }
                }

                largest_runlength_index = ptngc_find_magic_index(largest_runlength_base as u32);
                if largest_runlength_index != new_small_index {
                    iter_small_index = largest_runlength_index;
                    debug!(
                        "I found a new possible small index: {iter_small_index} Base={}",
                        MAGIC[iter_small_index as usize]
                    );
                }

                if new_runlength == iter_runlength && new_small_index == iter_small_index {
                    break;
                }
            }
            debug!("Exit iterative loop");
            // Verify that we got something good. We may have caught a
            // substantially larger atom. If so we should just bail
            // out and let the loop get on another lap. We may have a
            // minimum runlength though and then we have to fulfill
            // the request to write out these atoms!
            rle_index_dep = 0;
            if new_runlength < 3 {
                rle_index_dep = IS_LARGE;
            } else if new_runlength < 6 {
                rle_index_dep = QUITE_LARGE;
            }
            if (min_runlength > 0)
                || ((new_small_index < small_index + IS_LARGE)
                    && (new_small_index + rle_index_dep < max_large_index))
                || (new_small_index + IS_LARGE < max_large_index)
            {
                if new_runlength != runlength || new_small_index != small_index {
                    let mut change = new_small_index as i32 - small_index as i32;
                    if new_small_index == 0 {
                        change = 0;
                    }

                    if change < 0 {
                        for ixx in 0..new_runlength {
                            let mut rejected;
                            loop {
                                let mut isum = 0.0; // ints can be almost 32 bit so multiplication will overflow. So do doubles
                                for ixyz in 0..3 {
                                    // encode_ints is already positive (and multiplied by 2 versus the original, just as magic ints)
                                    let id = f64::from(encode_ints[ixx * 3 + ixyz]);
                                    isum += id * id;
                                }
                                rejected = false;

                                if isum
                                    > f64::from(
                                        MAGIC[usize::try_from(small_index as i32 + change)
                                            .expect("usize from i32")],
                                    ) * f64::from(
                                        MAGIC[usize::try_from(small_index as i32 + change)
                                            .expect("usize from i32")],
                                    )
                                {
                                    rejected = true;
                                    change += 1;
                                }
                                if !(change < 0 && rejected) {
                                    break;
                                }
                            }
                            if change == 0 {
                                break;
                            }
                        }
                    }

                    // If the only thing to do is to change the base by only one -1 it is probably not worth it
                    if !((change == -1) && (runlength == new_runlength)) {
                        /* If we have a very short runlength we do not
                        want to do large base changes.  It costs 6
                        extra bits to do -2. We gain 2/3
                        bits per value to decrease the index by -2,
                        ie 2 bits, so to any changes down we must
                        have a runlength of 3 or more to do it for
                        one molecule! If we have several molecules we
                        will gain of course, so not be so strict. */
                        if (change == -2) && (new_runlength < 3) {
                            if runlength == new_runlength {
                                change = 0;
                            } else {
                                change = -1;
                            }
                        }

                        // First adjust base using large base change instruction (if necessary)
                        while !(-1..=1).contains(&change) || ((new_runlength == 6) && (change != 0))
                        {
                            let mut code = 0;
                            let mut this_change = change;
                            this_change = this_change.clamp(-2, 2);
                            change -= this_change;
                            small_index = u32::try_from(small_index as i32 + this_change)
                                .expect("u32 from i32");
                            if this_change < 0 {
                                code |= 2;
                                this_change = -this_change;
                            }
                            code |= u32::try_from(this_change - 1).expect("u32 from i32");
                            write_instruction(coder, INSTR_LARGE_BASE_CHANGE, &mut output);
                            coder.ptngc_writebits(code, 2, &mut output);
                        }
                        // If there is still base change or runlength changes to do, we do them now
                        if (new_runlength != runlength) || (change != 0) {
                            let mut ichange = u32::try_from(change + 1).expect("u32 from i32");

                            let irun =
                                u32::try_from((new_runlength - 1) * 3).expect("u32 from i32");
                            if new_runlength == 6 {
                                ichange = 0;
                            } /* Means no change. The change has been taken care of explicitly by a
                            large base change instruction above. */
                            let code = ichange + irun;
                            debug!("Small base change: {change} Runlength change: {new_runlength}");
                            small_index =
                                u32::try_from(small_index as i32 + change).expect("u32 from i32");
                            write_instruction(coder, INSTR_BASE_RUNLENGTH, &mut output);
                            coder.ptngc_writebits(code, 4, &mut output);
                            runlength = new_runlength;
                        }
                    }
                    debug!(
                        "Current small index: {small_index} Base={}",
                        MAGIC[small_index as usize]
                    );
                }
                // If we have a large previous integer we can combine it with a sequence of small ints.
                if has_large > 0 {
                    // If swapatoms is set to 1 but we did actually not
                    // do any swapping, we must first write out the
                    // large atom and then the small. If swapatoms is 1
                    // and we did swapping we can use the efficient
                    // encoding.
                    if swapatoms && !didswap {
                        // Flush all large atoms.
                        let has_large_c = has_large;
                        flush_large(
                            coder,
                            &mut has_large,
                            &mut has_large_ints,
                            has_large_c,
                            &large_index,
                            large_nbits,
                            &mut compress_buffer,
                            &mut output,
                        );
                        write_instruction(coder, INSTR_ONLY_SMALL, &mut output);
                    } else {
                        debug!("Sequence of one large and small integers (good compression)");
                        // Flush all large atoms but one!
                        if has_large > 1 {
                            let has_large_c = has_large;
                            flush_large(
                                coder,
                                &mut has_large,
                                &mut has_large_ints,
                                has_large_c - 1,
                                &large_index,
                                large_nbits,
                                &mut compress_buffer,
                                &mut output,
                            );
                        }
                        write_instruction(coder, INSTR_DEFAULT, &mut output);
                        write_three_large(
                            coder,
                            &has_large_ints,
                            &large_index,
                            large_nbits,
                            &mut compress_buffer,
                            &mut output,
                        );
                        has_large = 0;
                    }
                } else {
                    write_instruction(coder, INSTR_ONLY_SMALL, &mut output);
                }
                // Base compress small integers using the current parameters.
                let nbits = MAGIC_BITS[usize::try_from(small_index).expect("usize from u32")]
                    [runlength - 1];
                // The same base is used for the small changes.
                let small_idx = [small_index, small_index, small_index];
                trajcoder_base_compress(
                    &encode_ints,
                    runlength * 3,
                    &small_idx,
                    &mut compress_buffer,
                );
                debug!(
                    "nbits={nbits}, ({})",
                    f64::from(nbits) / (runlength as f64 * 3.)
                );
                // write out base compressed small integers
                coder.ptngc_write_many_bits(&compress_buffer, nbits, &mut output);
                if log::log_enabled!(log::Level::Debug) {
                    for ienc in 0..runlength {
                        debug!(
                            "Small: {} {} {}",
                            encode_ints[ienc * 3],
                            encode_ints[ienc * 3 + 1],
                            encode_ints[ienc * 3 + 2],
                        );
                    }
                }

                // Update `prevcoord`
                for ienc in 0..runlength {
                    debug!("Prevcoord in packing: {prevcoord:?}");
                    prevcoord[0] =
                        prevcoord[0].wrapping_add(xtc3::unpositive_int(encode_ints[ienc * 3]));
                    prevcoord[1] =
                        prevcoord[1].wrapping_add(xtc3::unpositive_int(encode_ints[ienc * 3 + 1]));
                    prevcoord[2] =
                        prevcoord[2].wrapping_add(xtc3::unpositive_int(encode_ints[ienc * 3 + 2]));
                }

                debug!("Prevcoord in packing: {prevcoord:?}");
                input_ptr += 3 * runlength;
                ntriplets_left -= runlength;
            } else {
                debug!(
                    "Refused value: {new_small_index} old is {small_index} max is {max_large_index}"
                );
                refused = 1;
            }
        }
        debug!("Number of triplets left is {ntriplets_left}");
    }
    // If we have large previous integers we must flush them now.
    if has_large > 0 {
        let has_large_c = has_large;
        flush_large(
            coder,
            &mut has_large,
            &mut has_large_ints,
            has_large_c,
            &large_index,
            large_nbits,
            &mut compress_buffer,
            &mut output,
        );
    }
    coder.ptngc_pack_flush(&mut output);

    let length = output.len();
    (output, length)
}

// An almost identical implementation exists in xtc3.rs, but differs in how it writes
// the instructions (last couple of lines)
fn swapdecide(
    coder: &mut Coder,
    input: &[i32],
    swapatoms: &mut bool,
    minint: &[i32; 3],
    output: &mut Vec<u8>,
) {
    let mut didswap = false;

    let (normal, swapped) = xtc3::swap_is_better(input, minint);
    // We have to determine if it is worth to change the behaviour.
    // If diff is positive it means that it is worth something to
    // swap. But it costs 4 bits to do the change. If we assume that
    // we gain 0.17 bit by the swap per value, and the runlength>2
    // for four molecules in a row, we gain something. So check if we
    // gain at least 0.17 bits to even attempt the swap.

    if (swapped < normal) && ((f64::from(swapped) / f64::from(normal)).abs() < xtc3::IFLIPGAINCHECK)
        || ((normal < swapped)
            && (f64::from(normal) / f64::from(swapped)).abs() < xtc3::IFLIPGAINCHECK)
    {
        if swapped < normal {
            if !*swapatoms {
                *swapatoms = true;
                didswap = true;
            }
        } else if *swapatoms {
            *swapatoms = false;
            didswap = true;
        }
    }

    if didswap {
        write_instruction(coder, INSTR_FLIP, output);
    }
}

fn insert_batch(
    input: &[i32],
    ntriplets_left: usize,
    prevcoord: &[i32],
    minint: &[i32],
    encode_ints: &mut [i32],
    startenc: usize,
) -> usize {
    let mut nencode = startenc * 3;
    let mut tmp_prevcoord = [0; 3];
    tmp_prevcoord.copy_from_slice(prevcoord);

    if startenc > 0 {
        for i in 0..startenc {
            tmp_prevcoord[0] = tmp_prevcoord[0].wrapping_add(encode_ints[i * 3]);
            tmp_prevcoord[1] = tmp_prevcoord[1].wrapping_add(encode_ints[i * 3 + 1]);
            tmp_prevcoord[2] = tmp_prevcoord[2].wrapping_add(encode_ints[i * 3 + 2]);
        }
    }

    debug!("New batch");
    while (nencode < 21) && (nencode < ntriplets_left * 3) {
        encode_ints[nencode] = input[nencode]
            .wrapping_sub(minint[0])
            .wrapping_sub(tmp_prevcoord[0]);
        encode_ints[nencode + 1] = input[nencode + 1]
            .wrapping_sub(minint[1])
            .wrapping_sub(tmp_prevcoord[1]);
        encode_ints[nencode + 2] = input[nencode + 2]
            .wrapping_sub(minint[2])
            .wrapping_sub(tmp_prevcoord[2]);
        debug!(
            "{:6}: {:6} {:6} {:6}\t\t{:6} {:6} {:6}\t\t{:6} {:6} {:6}",
            nencode,
            input[nencode].wrapping_sub(minint[0]),
            input[nencode + 1].wrapping_sub(minint[1]),
            input[nencode + 2].wrapping_sub(minint[2]),
            encode_ints[nencode],
            encode_ints[nencode + 1],
            encode_ints[nencode + 2],
            positive_int(encode_ints[nencode]),
            positive_int(encode_ints[nencode + 1]),
            positive_int(encode_ints[nencode + 2]),
        );

        tmp_prevcoord[0] = input[nencode].wrapping_sub(minint[0]);
        tmp_prevcoord[1] = input[nencode + 1].wrapping_sub(minint[1]);
        tmp_prevcoord[2] = input[nencode + 2].wrapping_sub(minint[2]);
        nencode += 3;
    }
    nencode
}

// Compute number of bits required to store values using three different bases in the index array
pub(crate) fn compute_magic_bits(index: [u32; 3]) -> u32 {
    let mut largeint = [0; 4];
    let mut largeint_tmp = [0; 4];

    for i in 0..3 {
        if i != 0 {
            // We must do the multiplication of the largeint with the integer base
            ptngc_largeint_mul(
                MAGIC[usize::try_from(index[i]).expect("usize from u32")],
                &largeint,
                &mut largeint_tmp,
                4,
            );
            largeint[..4].copy_from_slice(&largeint_tmp[..4]);
        }
        ptngc_largeint_add(
            MAGIC[usize::try_from(index[i]).expect("usize from u32")] - 1,
            &mut largeint,
            4,
        );
    }
    // Find last bit
    let mut onebit: u32 = 0;
    for (i, &item) in largeint.iter().enumerate().take(3) {
        if item != 0 {
            onebit = (i as u32) * 32 + (32 - item.leading_zeros());
        }
    }
    onebit
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Once;

    static INIT: Once = Once::new();

    fn init_logger() {
        INIT.call_once(|| {
            env_logger::builder()
                .format_timestamp(None)
                .is_test(true)
                .try_init()
                .ok();
        });
    }

    fn run_case(mut data: Vec<i32>) -> Vec<u8> {
        let mut coder = Coder::default();
        let mut length = data.len();
        let (out, _) = ptngc_pack_array_xtc2(&mut coder, &mut data, &mut length);
        out
    }

    #[test]
    fn single_triplet_zeros() {
        let output = run_case(vec![0, 0, 0]);
        let expected = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 64];
        assert_eq!(output, expected);
    }

    #[test]
    fn two_triplets() {
        let output = run_case(vec![1, 2, 3, 4, 5, 6]);
        let expected = [0, 0, 0, 1, 0, 0, 0, 3, 0, 0, 0, 5, 3, 3, 3, 0, 64, 9, 116];
        assert_eq!(output, expected);
    }

    #[test]
    fn three_triplets() {
        let output = run_case(vec![-1, -1, -1, 0, 0, 0, 1, 1, 1]);
        let expected = [
            0, 0, 0, 2, 0, 0, 0, 2, 0, 0, 0, 2, 2, 2, 2, 1, 64, 17, 84, 168,
        ];
        assert_eq!(output, expected);
    }

    #[test]
    fn six_triplets() {
        let output = run_case(vec![
            10, 20, 30, 10, 20, 30, 10, 20, 30, 10, 20, 30, 10, 20, 30, 10, 20, 30,
        ]);
        let expected = [
            0, 0, 0, 19, 0, 0, 0, 39, 0, 0, 0, 59, 0, 0, 0, 0, 121, 0, 0, 64,
        ];
        assert_eq!(output, expected);
    }

    #[test]
    fn mixed_magnitudes() {
        init_logger();
        let output = run_case(vec![
            100000, -100000, 50000, -25000, 25000, -50000, 0, 0, 0, 1, 2, 3, 9999, -9999, 12345,
            -32768, 32767, 0,
        ]);
        let expected = [
            0, 1, 0, 0, 0, 3, 13, 64, 0, 1, 134, 160, 48, 48, 46, 4, 5, 8, 36, 170, 8, 94, 134,
            161, 54, 7, 204, 107, 216, 158, 67, 96, 240, 147, 242, 0, 0, 71, 105, 88, 192, 72, 148,
            121, 56, 72, 162, 196, 130, 110, 6, 0, 0,
        ];
        assert_eq!(output, expected);
    }

    #[test]
    fn zigzag_deltas() {
        init_logger();
        let output = run_case(vec![
            0, 0, 0, 10, -10, 20, -20, 30, -40, 50, -60, 70, -80, 90, -100,
        ]);
        let expected = [
            0, 0, 0, 160, 0, 0, 0, 120, 0, 0, 0, 200, 18, 18, 19, 9, 100, 43, 64, 170, 128, 119,
            164, 174, 51, 60, 144, 26, 176, 19, 201, 216, 0,
        ];
        assert_eq!(output, expected);
    }

    #[test]
    fn deterministic_random() {
        init_logger();
        let output = run_case(vec![
            -39240, -2939, 10066, -8293, 9777, -35408, -4547, 12916, -28241, -41983, -20421,
            -49871, 18396, 28205, 14827, 16744, 15179, -32093, -10434, 41494, -12693, 3788, -20011,
            17559, 20087, -34234, -23929, -10866, -5351, 36192,
        ]);
        let expected = [
            0, 1, 71, 254, 0, 1, 11, 116, 0, 1, 133, 158, 44, 45, 46, 0, 101, 150, 89, 101, 150,
            89, 101, 150, 89, 101, 150, 89, 101, 150, 65, 65, 195, 248, 155, 38, 225, 84, 193, 45,
            41, 72, 208, 31, 250, 238, 66, 109, 168, 188, 250, 228, 233, 149, 64, 0, 20, 167, 17,
            229, 182, 218, 240, 127, 50, 30, 108, 220, 170, 70, 120, 145, 27, 223, 177, 39, 212,
            237, 86, 85, 150, 117, 17, 165, 65, 213, 201, 67, 161, 208, 192, 49, 175, 48,
        ];
        assert_eq!(output, expected);
    }
}
