/*
NAME
  adpcm.go

DESCRIPTION
  adpcm.go contains functions for encoding/compressing pcm into adpcm and decoding/decompressing back to pcm.

AUTHOR
  Trek Hopton <trek@ausocean.org>

LICENSE
  adpcm.go is Copyright (C) 2018 the Australian Ocean Lab (AusOcean)

  It is free software: you can redistribute it and/or modify them
  under the terms of the GNU General Public License as published by the
  Free Software Foundation, either version 3 of the License, or (at your
  option) any later version.

  It is distributed in the hope that it will be useful, but WITHOUT
  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
  for more details.

  You should have received a copy of the GNU General Public License in gpl.txt.
  If not, see [GNU licenses](http://www.gnu.org/licenses).
*/

/*
	Original IMA/DVI ADPCM specification: (http://www.cs.columbia.edu/~hgs/audio/dvi/IMA_ADPCM.pdf).
	Reference algorithms for ADPCM compression and decompression are in part 6.
*/

package adpcm

import (
	"bytes"
	"encoding/binary"
	"fmt"
)

// encoder is used to encode to ADPCM from PCM data.
// pred and index hold state that persists between calls to encodeSample and calcHead.
// dest is the output buffer that implements io.writer and io.bytewriter, ie. where the encoded ADPCM data is written to.
type encoder struct {
	dest  *bytes.Buffer
	pred  int16
	index int16
}

// decoder is used to decode from ADPCM to PCM data.
// pred, index, and step hold state that persists between calls to decodeSample.
// dest is the output buffer that implements io.writer and io.bytewriter, ie. where the decoded PCM data is written to.
type decoder struct {
	dest  *bytes.Buffer
	pred  int16
	index int16
	step  int16
}

// BytesOutput will return the number of adpcm bytes that will be generated for the given pcm data
func BytesOutput(pcm int) int {
	// for X pcm bytes, 2 bytes are left uncompressed, the rest is compressed by a factor of 4
	// and a start index and padding byte are added.
	return (pcm-2)/4 + 2 + 1 + 1
}

// PcmBS is the size of the blocks that an encoder uses.
// 'encodeBlock' will encode PcmBS bytes at a time and the output will be AdpcmBS bytes long.
const PcmBS = 1010

// AdpcmBS is the size of the blocks that a decoder uses.
// 'decodeBlock' will decode AdpcmBS bytes at a time and the output will be PcmBS bytes long.
const AdpcmBS = 256

// Table of index changes (see spec).
var indexTable = []int16{
	-1, -1, -1, -1, 2, 4, 6, 8,
	-1, -1, -1, -1, 2, 4, 6, 8,
}

// Quantize step size table (see spec).
var stepTable = []int16{
	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,
}

// NewEncoder retuns a new ADPCM encoder.
func NewEncoder(dst *bytes.Buffer) *encoder {
	e := encoder{
		dest: dst,
	}
	return &e
}

// NewDecoder retuns a new ADPCM decoder.
func NewDecoder(dst *bytes.Buffer) *decoder {
	d := decoder{
		step: stepTable[0],
		dest: dst,
	}
	return &d
}

// encodeSample takes a single 16 bit PCM sample and
// returns a byte of which the last 4 bits are an encoded ADPCM nibble.
func (e *encoder) encodeSample(sample int16) byte {
	// Find difference of actual sample from encoder's prediction.
	delta := sample - e.pred

	// Create and set sign bit for nibble and find absolute value of difference.
	var nib byte
	if delta < 0 {
		nib = 8
		delta = -delta
	}

	step := stepTable[e.index]
	diff := step >> 3
	var mask byte = 4

	for i := 0; i < 3; i++ {
		if delta > step {
			nib |= mask
			delta -= step
			diff += step
		}
		mask >>= 1
		step >>= 1
	}

	// Adjust predicted sample based on calculated difference.
	if nib&8 != 0 {
		e.pred -= diff
	} else {
		e.pred += diff
	}

	e.index += indexTable[nib&7]

	// Check for underflow and overflow.
	if e.index < 0 {
		e.index = 0
	} else if e.index > int16(len(stepTable)-1) {
		e.index = int16(len(stepTable) - 1)
	}

	return nib
}

// decodeSample takes a byte, the last 4 bits of which contain a single
// 4 bit ADPCM nibble, and returns a 16 bit decoded PCM sample.
func (d *decoder) decodeSample(nibble byte) int16 {
	// Calculate difference.
	var diff int16
	if nibble&4 != 0 {
		diff += d.step
	}
	if nibble&2 != 0 {
		diff += d.step >> 1
	}
	if nibble&1 != 0 {
		diff += d.step >> 2
	}
	diff += d.step >> 3

	// Account for sign bit.
	if nibble&8 != 0 {
		diff = -diff
	}

	// Adjust predicted sample based on calculated difference.
	d.pred += diff

	// Adjust index into step size lookup table using nibble.
	d.index += indexTable[nibble]

	// Check for overflow and underflow.
	if d.index < 0 {
		d.index = 0
	} else if d.index > int16(len(stepTable)-1) {
		d.index = int16(len(stepTable) - 1)
	}

	// Find new quantizer step size.
	d.step = stepTable[d.index]

	return d.pred
}

// calcHead sets the state for the encoder by running the first sample through
// the encoder, and writing the first sample to the encoder's io.Writer (dest).
// It returns the number of bytes written to the encoder's io.Writer (dest) along with any errors.
func (e *encoder) calcHead(sample []byte) (int, error) {
	// Check that we are given 1 16-bit sample (2 bytes).
	const sampSize = 2
	if len(sample) != sampSize {
		return 0, fmt.Errorf("length of given byte array is: %v, expected: %v", len(sample), sampSize)
	}

	intSample := int16(binary.LittleEndian.Uint16(sample))
	e.encodeSample(intSample)

	n, err := e.dest.Write(sample)
	if err != nil {
		return n, err
	}

	err = e.dest.WriteByte(byte(uint16(e.index)))
	if err != nil {
		return n, err
	}
	n++

	err = e.dest.WriteByte(byte(0x00))
	if err != nil {
		return n, err
	}
	n++
	return n, nil
}

// encodeBlock takes a slice of 1010 bytes (505 16-bit PCM samples).
// It writes encoded (compressed) bytes (each byte containing two ADPCM nibbles) to the encoder's io.Writer (dest).
// The number of bytes written is returned along with any errors.
// Note: nibbles are output in little endian order, eg. n1n0 n3n2 n5n4...
// Note: first 4 bytes are for initializing the decoder before decoding a block.
// 		- First two bytes contain the first 16-bit sample uncompressed.
// 		- Third byte is the decoder's starting index for the block, the fourth is padding and ignored.
func (e *encoder) encodeBlock(block []byte) (int, error) {
	if len(block) != PcmBS {
		return 0, fmt.Errorf("unsupported block size. Given: %v, expected: %v, ie. 505 16-bit PCM samples", len(block), PcmBS)
	}

	n, err := e.calcHead(block[0:2])
	if err != nil {
		return n, err
	}

	for i := 3; i < PcmBS; i += 4 {
		nib1 := e.encodeSample(int16(binary.LittleEndian.Uint16(block[i-1 : i+1])))
		nib2 := e.encodeSample(int16(binary.LittleEndian.Uint16(block[i+1 : i+3])))
		err = e.dest.WriteByte(byte((nib2 << 4) | nib1))
		if err != nil {
			return n, err
		}
		n++
	}

	return n, nil
}

// decodeBlock takes a slice of 256 bytes, each byte after the first 4 should contain two ADPCM encoded nibbles.
// It writes the resulting decoded (decompressed) 16-bit PCM samples to the decoder's io.Writer (dest).
// The number of bytes written is returned along with any errors.
func (d *decoder) decodeBlock(block []byte) (int, error) {
	if len(block) != AdpcmBS {
		return 0, fmt.Errorf("unsupported block size. Given: %v, expected: %v", len(block), AdpcmBS)
	}

	// Initialize decoder with first 4 bytes of the block.
	d.pred = int16(binary.LittleEndian.Uint16(block[0:2]))
	d.index = int16(block[2])
	d.step = stepTable[d.index]
	n, err := d.dest.Write(block[0:2])
	if err != nil {
		return n, err
	}

	// For each byte, seperate it into two nibbles (each nibble is a compressed sample),
	// then decode each nibble and output the resulting 16-bit samples.
	for i := 4; i < AdpcmBS; i++ {
		twoNibs := block[i]
		nib2 := byte(twoNibs >> 4)
		nib1 := byte((nib2 << 4) ^ twoNibs)

		firstBytes := make([]byte, 2)
		binary.LittleEndian.PutUint16(firstBytes, uint16(d.decodeSample(nib1)))
		_n, err := d.dest.Write(firstBytes)
		n += _n
		if err != nil {
			return n, err
		}

		secondBytes := make([]byte, 2)
		binary.LittleEndian.PutUint16(secondBytes, uint16(d.decodeSample(nib2)))
		_n, err = d.dest.Write(secondBytes)
		n += _n
		if err != nil {
			return n, err
		}
	}

	return n, nil
}

// Write takes a slice of bytes of arbitrary length representing pcm and encodes in into adpcm.
// It writes its output to the encoder's dest.
// The number of bytes written out is returned along with any error that occured.
func (e *encoder) Write(inPcm []byte) (int, error) {
	numBlocks := len(inPcm) / PcmBS
	n := 0
	for i := 0; i < numBlocks; i++ {
		block := inPcm[PcmBS*i : PcmBS*(i+1)]
		_n, err := e.encodeBlock(block)
		n += _n
		if err != nil {
			return n, err
		}
	}

	return n, nil
}

// Write takes a slice of bytes of arbitrary length representing adpcm and decodes in into pcm.
// It writes its output to the decoder's dest.
// The number of bytes written out is returned along with any error that occured.
func (d *decoder) Write(inAdpcm []byte) (int, error) {
	numBlocks := len(inAdpcm) / AdpcmBS
	n := 0
	for i := 0; i < numBlocks; i++ {
		block := inAdpcm[AdpcmBS*i : AdpcmBS*(i+1)]
		_n, err := d.decodeBlock(block)
		n += _n
		if err != nil {
			return n, err
		}
	}

	return n, nil
}