#!/usr/bin/env python3
"""Converts a 16-bit PCM WAV (mono or stereo) to an IMA ADPCM WAV, about four times smaller.

Used by build-parallel-overhead.sh for the two music tracks of Parallel Overhead (CC0, 20 MB each as PCM).
SDL2's SDL_LoadWAV decodes IMA ADPCM back to 16-bit samples, so the game needs no change. Pure Python,
deterministic (no floating point), same licence as the rest of this build (MIT).

  python3 -I wav2adpcm.py in.wav out.wav
"""
import struct
import sys

STEPS = [7, 8, 9, 10, 11, 12, 13, 14, 16, 17, 19, 21, 23, 25, 28, 31, 34, 37, 41, 45, 50, 55, 60, 66, 73, 80, 88, 97,
         107, 118, 130, 143, 157, 173, 190, 209, 230, 253, 279, 307, 337, 371, 408, 449, 494, 544, 598, 658, 724,
         796, 876, 963, 1060, 1166, 1282, 1411, 1552, 1707, 1878, 2066, 2272, 2499, 2749, 3024, 3327, 3660, 4026,
         4428, 4871, 5358, 5894, 6484, 7132, 7845, 8630, 9493, 10442, 11487, 12635, 13899, 15289, 16818, 18500,
         20350, 22385, 24623, 27086, 29794, 32767]
INDEX = [-1, -1, -1, -1, 2, 4, 6, 8]
BLOCK = 512  # bytes per channel and block -> 1017 samples per block


def read_wav(path):
    data = open(path, 'rb').read()
    assert data[:4] == b'RIFF' and data[8:12] == b'WAVE'
    pos, fmt, pcm = 12, None, None
    while pos + 8 <= len(data):
        cid, size = data[pos:pos + 4], struct.unpack('<I', data[pos + 4:pos + 8])[0]
        body = pos + 8
        if cid == b'fmt ':
            fmt = struct.unpack('<HHIIHH', data[body:body + 16])
        elif cid == b'data':
            pcm = data[body:min(body + size, len(data))]
            break
        pos = body + size + (size & 1)
    assert fmt and fmt[0] == 1 and fmt[5] == 16, 'only 16-bit PCM'
    return fmt[1], fmt[2], pcm


def encode_channel(samples):
    """samples: list of int16. Returns a list of blocks (bytes), each BLOCK bytes."""
    per_block = (BLOCK - 4) * 2 + 1
    blocks = []
    pred, index = 0, 0
    n = len(samples)
    for start in range(0, n, per_block):
        chunk = samples[start:start + per_block]
        if len(chunk) < per_block:
            chunk = chunk + [chunk[-1]] * (per_block - len(chunk))
        pred = chunk[0]
        out = bytearray(struct.pack('<hBB', pred, index, 0))
        nibbles = []
        for s in chunk[1:]:
            step = STEPS[index]
            diff = s - pred
            code = 0
            if diff < 0:
                code = 8
                diff = -diff
            delta = step >> 3
            if diff >= step:
                code |= 4; diff -= step; delta += step
            step >>= 1
            if diff >= step:
                code |= 2; diff -= step; delta += step
            step >>= 1
            if diff >= step:
                code |= 1; delta += step
            pred += -delta if code & 8 else delta
            pred = -32768 if pred < -32768 else 32767 if pred > 32767 else pred
            index += INDEX[code & 7]
            index = 0 if index < 0 else 88 if index > 88 else index
            nibbles.append(code)
        for i in range(0, len(nibbles), 2):
            out.append(nibbles[i] | (nibbles[i + 1] << 4))
        blocks.append(bytes(out))
    return blocks


def main(src, dst):
    ch, rate, pcm = read_wav(src)
    count = len(pcm) // (2 * ch)
    allsamples = struct.unpack('<%dh' % (count * ch), pcm[:count * ch * 2])
    chans = [list(allsamples[c::ch]) for c in range(ch)]
    encoded = [encode_channel(c) for c in chans]
    per_block = (BLOCK - 4) * 2 + 1
    body = bytearray()
    for b in range(len(encoded[0])):
        parts = [encoded[c][b] for c in range(ch)]
        for c in range(ch):
            body += parts[c][:4]  # block headers
        # nibble data interleaved in groups of 4 bytes per channel
        for off in range(4, BLOCK, 4):
            for c in range(ch):
                body += parts[c][off:off + 4]
    block_align = BLOCK * ch
    fmt = struct.pack('<HHIIHHHH', 0x11, ch, rate, rate // per_block * block_align + (rate % per_block) * block_align // per_block,
                      block_align, 4, 2, per_block)
    fact = struct.pack('<I', count)
    riff = (b'WAVE' + b'fmt ' + struct.pack('<I', len(fmt)) + fmt + b'fact' + struct.pack('<I', 4) + fact +
            b'data' + struct.pack('<I', len(body)) + bytes(body))
    open(dst, 'wb').write(b'RIFF' + struct.pack('<I', len(riff)) + riff)


if __name__ == '__main__':
    main(sys.argv[1], sys.argv[2])