/*

 * ALAC audio encoder

 * Copyright (c) 2008  Jaikrishnan Menon <realityman@gmx.net>

 *

 * This file is part of FFmpeg.

 *

 * FFmpeg is free software; you can redistribute it and/or

 * modify it under the terms of the GNU Lesser General Public

 * License as published by the Free Software Foundation; either

 * version 2.1 of the License, or (at your option) any later version.

 *

 * FFmpeg 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

 * Lesser General Public License for more details.

 *

 * You should have received a copy of the GNU Lesser General Public

 * License along with FFmpeg; if not, write to the Free Software

 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA

 */

#include "libavutil/mem.h"

#include "libavutil/opt.h"

#include "avcodec.h"

#include "codec_internal.h"

#include "encode.h"

#include "put_bits.h"

#include "lpc.h"

#include "mathops.h"

#include "alac_data.h"

#define DEFAULT_FRAME_SIZE        4096

#define ALAC_EXTRADATA_SIZE       36

#define ALAC_FRAME_HEADER_SIZE    55

#define ALAC_FRAME_FOOTER_SIZE    3

#define ALAC_ESCAPE_CODE          0x1FF

#define ALAC_MAX_LPC_ORDER        30

#define DEFAULT_MAX_PRED_ORDER    6

#define DEFAULT_MIN_PRED_ORDER    4

#define ALAC_MAX_LPC_PRECISION    9

#define ALAC_MIN_LPC_SHIFT        0

#define ALAC_MAX_LPC_SHIFT        9

#define ALAC_CHMODE_LEFT_RIGHT    0

#define ALAC_CHMODE_LEFT_SIDE     1

#define ALAC_CHMODE_RIGHT_SIDE    2

#define ALAC_CHMODE_MID_SIDE      3

typedef struct RiceContext {

    int history_mult;

    int initial_history;

    int k_modifier;

    int rice_modifier;

} RiceContext;

typedef struct AlacLPCContext {

    int lpc_order;

    int lpc_coeff[ALAC_MAX_LPC_ORDER+1];

    int lpc_quant;

} AlacLPCContext;

typedef struct AlacEncodeContext {

    const AVClass *class;

    AVCodecContext *avctx;

    int frame_size;                     /**< current frame size               */

    int verbatim;                       /**< current frame verbatim mode flag */

    int compression_level;

    int min_prediction_order;

    int max_prediction_order;

    int max_coded_frame_size;

    int write_sample_size;

    int extra_bits;

    int32_t sample_buf[2][DEFAULT_FRAME_SIZE];

    int32_t predictor_buf[2][DEFAULT_FRAME_SIZE];

    int interlacing_shift;

    int interlacing_leftweight;

    PutBitContext pbctx;

    RiceContext rc;

    AlacLPCContext lpc[2];

    LPCContext lpc_ctx;

} AlacEncodeContext;

static void init_sample_buffers(AlacEncodeContext *s, int channels,

                                const uint8_t *samples[2])

{

    int ch, i;

    int shift = av_get_bytes_per_sample(s->avctx->sample_fmt) * 8 -

                s->avctx->bits_per_raw_sample;

#define COPY_SAMPLES(type) do {                             \

        for (ch = 0; ch < channels; ch++) {                 \

            int32_t       *bptr = s->sample_buf[ch];        \

            const type *sptr = (const type *)samples[ch];   \

            for (i = 0; i < s->frame_size; i++)             \

                bptr[i] = sptr[i] >> shift;                 \

        }                                                   \

    } while (0)

    if (s->avctx->sample_fmt == AV_SAMPLE_FMT_S32P)

        COPY_SAMPLES(int32_t);

    else

        COPY_SAMPLES(int16_t);

}

static void encode_scalar(AlacEncodeContext *s, int x,

                          int k, int write_sample_size)

{

    int divisor, q, r;

    k = FFMIN(k, s->rc.k_modifier);

    divisor = (1<<k) - 1;

    q = x / divisor;

    r = x % divisor;

    if (q > 8) {

        // write escape code and sample value directly

        put_bits(&s->pbctx, 9, ALAC_ESCAPE_CODE);

        put_bits(&s->pbctx, write_sample_size, x);

    } else {

        if (q)

            put_bits(&s->pbctx, q, (1<<q) - 1);

        put_bits(&s->pbctx, 1, 0);

        if (k != 1) {

            if (r > 0)

                put_bits(&s->pbctx, k, r+1);

            else

                put_bits(&s->pbctx, k-1, 0);

        }

    }

}

static void write_element_header(AlacEncodeContext *s,

                                 enum AlacRawDataBlockType element,

                                 int instance)

{

    int encode_fs = 0;

    if (s->frame_size < DEFAULT_FRAME_SIZE)

        encode_fs = 1;

    put_bits(&s->pbctx, 3,  element);               // element type

    put_bits(&s->pbctx, 4,  instance);              // element instance

    put_bits(&s->pbctx, 12, 0);                     // unused header bits

    put_bits(&s->pbctx, 1,  encode_fs);             // Sample count is in the header

    put_bits(&s->pbctx, 2,  s->extra_bits >> 3);    // Extra bytes (for 24-bit)

    put_bits(&s->pbctx, 1,  s->verbatim);           // Audio block is verbatim

    if (encode_fs)

        put_bits32(&s->pbctx, s->frame_size);       // No. of samples in the frame

}

static void calc_predictor_params(AlacEncodeContext *s, int ch)

{

    int32_t coefs[MAX_LPC_ORDER][MAX_LPC_ORDER];

    int shift[MAX_LPC_ORDER];

    int opt_order;

    if (s->compression_level == 1) {

        s->lpc[ch].lpc_order = 6;

        s->lpc[ch].lpc_quant = 6;

        s->lpc[ch].lpc_coeff[0] =  160;

        s->lpc[ch].lpc_coeff[1] = -190;

        s->lpc[ch].lpc_coeff[2] =  170;

        s->lpc[ch].lpc_coeff[3] = -130;

        s->lpc[ch].lpc_coeff[4] =   80;

        s->lpc[ch].lpc_coeff[5] =  -25;

    } else {

        opt_order = ff_lpc_calc_coefs(&s->lpc_ctx, s->sample_buf[ch],

                                      s->frame_size,

                                      s->min_prediction_order,

                                      s->max_prediction_order,

                                      ALAC_MAX_LPC_PRECISION, coefs, shift,

                                      FF_LPC_TYPE_LEVINSON, 0,

                                      ORDER_METHOD_EST, ALAC_MIN_LPC_SHIFT,

                                      ALAC_MAX_LPC_SHIFT, 1);

        s->lpc[ch].lpc_order = opt_order;

        s->lpc[ch].lpc_quant = shift[opt_order-1];

        memcpy(s->lpc[ch].lpc_coeff, coefs[opt_order-1], opt_order*sizeof(int));

    }

}

static int estimate_stereo_mode(int32_t *left_ch, int32_t *right_ch, int n)

{

    int i, best;

    int32_t lt, rt;

    uint64_t sum[4];

    uint64_t score[4];

    /* calculate sum of 2nd order residual for each channel */

    sum[0] = sum[1] = sum[2] = sum[3] = 0;

    for (i = 2; i < n; i++) {

        lt =  left_ch[i] - 2 *  left_ch[i - 1] +  left_ch[i - 2];

        rt = right_ch[i] - 2 * right_ch[i - 1] + right_ch[i - 2];

        sum[2] += FFABS((lt + rt) >> 1);

        sum[3] += FFABS(lt - rt);

        sum[0] += FFABS(lt);

        sum[1] += FFABS(rt);

    }

    /* calculate score for each mode */

    score[0] = sum[0] + sum[1];

    score[1] = sum[0] + sum[3];

    score[2] = sum[1] + sum[3];

    score[3] = sum[2] + sum[3];

    /* return mode with lowest score */

    best = 0;

    for (i = 1; i < 4; i++) {

        if (score[i] < score[best])

            best = i;

    }

    return best;

}

static void alac_stereo_decorrelation(AlacEncodeContext *s)

{

    int32_t *left = s->sample_buf[0], *right = s->sample_buf[1];

    int i, mode, n = s->frame_size;

    int32_t tmp;

    mode = estimate_stereo_mode(left, right, n);

    switch (mode) {

    case ALAC_CHMODE_LEFT_RIGHT:

        s->interlacing_leftweight = 0;

        s->interlacing_shift      = 0;

        break;

    case ALAC_CHMODE_LEFT_SIDE:

        for (i = 0; i < n; i++)

            right[i] = left[i] - right[i];

        s->interlacing_leftweight = 1;

        s->interlacing_shift      = 0;

        break;

    case ALAC_CHMODE_RIGHT_SIDE:

        for (i = 0; i < n; i++) {

            tmp = right[i];

            right[i] = left[i] - right[i];

            left[i]  = tmp + (right[i] >> 31);

        }

        s->interlacing_leftweight = 1;

        s->interlacing_shift      = 31;

        break;

    default:

        for (i = 0; i < n; i++) {

            tmp = left[i];

            left[i]  = (tmp + right[i]) >> 1;

            right[i] =  tmp - right[i];

        }

        s->interlacing_leftweight = 1;

        s->interlacing_shift      = 1;

        break;

    }

}

static void alac_linear_predictor(AlacEncodeContext *s, int ch)

{

    int i;

    AlacLPCContext lpc = s->lpc[ch];

    int32_t *residual = s->predictor_buf[ch];

    if (lpc.lpc_order == 31) {

        residual[0] = s->sample_buf[ch][0];

        for (i = 1; i < s->frame_size; i++) {

            residual[i] = s->sample_buf[ch][i    ] -

                          s->sample_buf[ch][i - 1];

        }

        return;

    }

    // generalised linear predictor

    if (lpc.lpc_order > 0) {

        int32_t *samples  = s->sample_buf[ch];

        // generate warm-up samples

        residual[0] = samples[0];

        for (i = 1; i <= lpc.lpc_order; i++)

            residual[i] = sign_extend(samples[i] - samples[i-1], s->write_sample_size);

        // perform lpc on remaining samples

        for (i = lpc.lpc_order + 1; i < s->frame_size; i++) {

            int sum = 1 << (lpc.lpc_quant - 1), res_val, j;

            for (j = 0; j < lpc.lpc_order; j++) {

                sum += (samples[lpc.lpc_order-j] - samples[0]) *

                       lpc.lpc_coeff[j];

            }

            sum >>= lpc.lpc_quant;

            sum += samples[0];

            residual[i] = sign_extend(samples[lpc.lpc_order+1] - sum,

                                      s->write_sample_size);

            res_val = residual[i];

            if (res_val) {

                int index = lpc.lpc_order - 1;

                int neg = (res_val < 0);

                while (index >= 0 && (neg ? (res_val < 0) : (res_val > 0))) {

                    int val  = samples[0] - samples[lpc.lpc_order - index];

                    int sign = (val ? FFSIGN(val) : 0);

                    if (neg)

                        sign *= -1;

                    lpc.lpc_coeff[index] -= sign;

                    val *= sign;

                    res_val -= (val >> lpc.lpc_quant) * (lpc.lpc_order - index);

                    index--;

                }

            }

            samples++;

        }

    }

}

static void alac_entropy_coder(AlacEncodeContext *s, int ch)

{

    unsigned int history = s->rc.initial_history;

    int sign_modifier = 0, i, k;

    int32_t *samples = s->predictor_buf[ch];

    for (i = 0; i < s->frame_size;) {

        int x;

        k = av_log2((history >> 9) + 3);

        x  = -2 * (*samples) -1;

        x ^= x >> 31;

        samples++;

        i++;

        encode_scalar(s, x - sign_modifier, k, s->write_sample_size);

        history += x * s->rc.history_mult -

                   ((history * s->rc.history_mult) >> 9);

        sign_modifier = 0;

        if (x > 0xFFFF)

            history = 0xFFFF;

        if (history < 128 && i < s->frame_size) {

            unsigned int block_size = 0;

            k = 7 - av_log2(history) + ((history + 16) >> 6);

            while (*samples == 0 && i < s->frame_size) {

                samples++;

                i++;

                block_size++;

            }

            encode_scalar(s, block_size, k, 16);

            sign_modifier = (block_size <= 0xFFFF);

            history = 0;

        }

    }

}

static void write_element(AlacEncodeContext *s,

                          enum AlacRawDataBlockType element, int instance,

                          const uint8_t *samples0, const uint8_t *samples1)

{

    const uint8_t *samples[2] = { samples0, samples1 };

    int i, j, channels;

    int prediction_type = 0;

    PutBitContext *pb = &s->pbctx;

    channels = element == TYPE_CPE ? 2 : 1;

    if (s->verbatim) {

        write_element_header(s, element, instance);

        /* samples are channel-interleaved in verbatim mode */

        if (s->avctx->sample_fmt == AV_SAMPLE_FMT_S32P) {

            int shift = 32 - s->avctx->bits_per_raw_sample;

            const int32_t *samples_s32[2] = { (const int32_t *)samples0,

                                              (const int32_t *)samples1 };

            for (i = 0; i < s->frame_size; i++)

                for (j = 0; j < channels; j++)

                    put_sbits(pb, s->avctx->bits_per_raw_sample,

                              samples_s32[j][i] >> shift);

        } else {

            const int16_t *samples_s16[2] = { (const int16_t *)samples0,

                                              (const int16_t *)samples1 };

            for (i = 0; i < s->frame_size; i++)

                for (j = 0; j < channels; j++)

                    put_sbits(pb, s->avctx->bits_per_raw_sample,

                              samples_s16[j][i]);

        }

    } else {

        s->write_sample_size = s->avctx->bits_per_raw_sample - s->extra_bits +

                               channels - 1;

        init_sample_buffers(s, channels, samples);

        write_element_header(s, element, instance);

        // extract extra bits if needed

        if (s->extra_bits) {

            uint32_t mask = (1 << s->extra_bits) - 1;

            for (j = 0; j < channels; j++) {

                int32_t *extra = s->predictor_buf[j];

                int32_t *smp   = s->sample_buf[j];

                for (i = 0; i < s->frame_size; i++) {

                    extra[i] = smp[i] & mask;

                    smp[i] >>= s->extra_bits;

                }

            }

        }

        if (channels == 2)

            alac_stereo_decorrelation(s);

        else

            s->interlacing_shift = s->interlacing_leftweight = 0;

        put_bits(pb, 8, s->interlacing_shift);

        put_bits(pb, 8, s->interlacing_leftweight);

        for (i = 0; i < channels; i++) {

            calc_predictor_params(s, i);

            put_bits(pb, 4, prediction_type);

            put_bits(pb, 4, s->lpc[i].lpc_quant);

            put_bits(pb, 3, s->rc.rice_modifier);

            put_bits(pb, 5, s->lpc[i].lpc_order);

            // predictor coeff. table

            for (j = 0; j < s->lpc[i].lpc_order; j++)

                put_sbits(pb, 16, s->lpc[i].lpc_coeff[j]);

        }

        // write extra bits if needed

        if (s->extra_bits) {

            for (i = 0; i < s->frame_size; i++) {

                for (j = 0; j < channels; j++) {

                    put_bits(pb, s->extra_bits, s->predictor_buf[j][i]);

                }

            }

        }

        // apply lpc and entropy coding to audio samples

        for (i = 0; i < channels; i++) {

            alac_linear_predictor(s, i);

            // TODO: determine when this will actually help. for now it's not used.

            if (prediction_type == 15) {

                // 2nd pass 1st order filter

                int32_t *residual = s->predictor_buf[i];

                for (j = s->frame_size - 1; j > 0; j--)

                    residual[j] -= residual[j - 1];

            }

            alac_entropy_coder(s, i);

        }

    }

}

static int write_frame(AlacEncodeContext *s, AVPacket *avpkt,

                       uint8_t * const *samples)

{

    PutBitContext *pb = &s->pbctx;

    int channels = s->avctx->ch_layout.nb_channels;

    const enum AlacRawDataBlockType *ch_elements = ff_alac_channel_elements[channels - 1];

    const uint8_t *ch_map = ff_alac_channel_layout_offsets[channels - 1];

    int ch, element, sce, cpe;

    init_put_bits(pb, avpkt->data, avpkt->size);

    ch = element = sce = cpe = 0;

    while (ch < channels) {

        if (ch_elements[element] == TYPE_CPE) {

            write_element(s, TYPE_CPE, cpe, samples[ch_map[ch]],

                          samples[ch_map[ch + 1]]);

            cpe++;

            ch += 2;

        } else {

            write_element(s, TYPE_SCE, sce, samples[ch_map[ch]], NULL);

            sce++;

            ch++;

        }

        element++;

    }

    put_bits(pb, 3, TYPE_END);

    flush_put_bits(pb);

    return put_bytes_output(pb);

}

static av_always_inline int get_max_frame_size(int frame_size, int ch, int bps)

{

    int header_bits = 23 + 32 * (frame_size < DEFAULT_FRAME_SIZE);

    return FFALIGN(header_bits + bps * ch * frame_size + 3, 8) / 8;

}

static av_cold int alac_encode_close(AVCodecContext *avctx)

{

    AlacEncodeContext *s = avctx->priv_data;

    ff_lpc_end(&s->lpc_ctx);

    return 0;

}

static av_cold int alac_encode_init(AVCodecContext *avctx)

{

    AlacEncodeContext *s = avctx->priv_data;

    int ret;

    uint8_t *alac_extradata;

    avctx->frame_size = s->frame_size = DEFAULT_FRAME_SIZE;

    if (avctx->sample_fmt == AV_SAMPLE_FMT_S32P) {

        if (avctx->bits_per_raw_sample != 24)

            av_log(avctx, AV_LOG_WARNING, "encoding as 24 bits-per-sample\n");

        avctx->bits_per_raw_sample = 24;

    } else {

        avctx->bits_per_raw_sample = 16;

        s->extra_bits              = 0;

    }

    // Set default compression level

    if (avctx->compression_level == FF_COMPRESSION_DEFAULT)

        s->compression_level = 2;

    else

        s->compression_level = av_clip(avctx->compression_level, 0, 2);

    // Initialize default Rice parameters

    s->rc.history_mult    = 40;

    s->rc.initial_history = 10;

    s->rc.k_modifier      = 14;

    s->rc.rice_modifier   = 4;

    s->max_coded_frame_size = get_max_frame_size(avctx->frame_size,

                                                 avctx->ch_layout.nb_channels,

                                                 avctx->bits_per_raw_sample);

    avctx->extradata = av_mallocz(ALAC_EXTRADATA_SIZE + AV_INPUT_BUFFER_PADDING_SIZE);

    if (!avctx->extradata)

        return AVERROR(ENOMEM);

    avctx->extradata_size = ALAC_EXTRADATA_SIZE;

    alac_extradata = avctx->extradata;

    AV_WB32(alac_extradata,    ALAC_EXTRADATA_SIZE);

    AV_WB32(alac_extradata+4,  MKBETAG('a','l','a','c'));

    AV_WB32(alac_extradata+12, avctx->frame_size);

    AV_WB8 (alac_extradata+17, avctx->bits_per_raw_sample);

    AV_WB8 (alac_extradata+21, avctx->ch_layout.nb_channels);

    AV_WB32(alac_extradata+24, s->max_coded_frame_size);

    AV_WB32(alac_extradata+28,

            avctx->sample_rate * avctx->ch_layout.nb_channels * avctx->bits_per_raw_sample); // average bitrate

    AV_WB32(alac_extradata+32, avctx->sample_rate);

    // Set relevant extradata fields

    if (s->compression_level > 0) {

        AV_WB8(alac_extradata+18, s->rc.history_mult);

        AV_WB8(alac_extradata+19, s->rc.initial_history);

        AV_WB8(alac_extradata+20, s->rc.k_modifier);

    }

    if (s->max_prediction_order < s->min_prediction_order) {

        av_log(avctx, AV_LOG_ERROR,

               "invalid prediction orders: min=%d max=%d\n",

               s->min_prediction_order, s->max_prediction_order);

        return AVERROR(EINVAL);

    }

    s->avctx = avctx;

    if ((ret = ff_lpc_init(&s->lpc_ctx, avctx->frame_size,

                           s->max_prediction_order,

                           FF_LPC_TYPE_LEVINSON)) < 0) {

        return ret;

    }

    return 0;

}

static int alac_encode_frame(AVCodecContext *avctx, AVPacket *avpkt,

                             const AVFrame *frame, int *got_packet_ptr)

{

    AlacEncodeContext *s = avctx->priv_data;

    int out_bytes, max_frame_size, ret;

    s->frame_size = frame->nb_samples;

    if (frame->nb_samples < DEFAULT_FRAME_SIZE)

        max_frame_size = get_max_frame_size(s->frame_size, avctx->ch_layout.nb_channels,

                                            avctx->bits_per_raw_sample);

    else

        max_frame_size = s->max_coded_frame_size;

    if ((ret = ff_alloc_packet(avctx, avpkt, 4 * max_frame_size)) < 0)

        return ret;

    /* use verbatim mode for compression_level 0 */

    if (s->compression_level) {

        s->verbatim   = 0;

        s->extra_bits = avctx->bits_per_raw_sample - 16;

    } else {

        s->verbatim   = 1;

        s->extra_bits = 0;

    }

    out_bytes = write_frame(s, avpkt, frame->extended_data);

    if (out_bytes > max_frame_size) {

        /* frame too large. use verbatim mode */

        s->verbatim = 1;

        s->extra_bits = 0;

        out_bytes = write_frame(s, avpkt, frame->extended_data);

    }

    avpkt->size = out_bytes;

    *got_packet_ptr = 1;

    return 0;

}

#define OFFSET(x) offsetof(AlacEncodeContext, x)

#define AE AV_OPT_FLAG_AUDIO_PARAM | AV_OPT_FLAG_ENCODING_PARAM

static const AVOption options[] = {

    { "min_prediction_order", NULL, OFFSET(min_prediction_order), AV_OPT_TYPE_INT, { .i64 = DEFAULT_MIN_PRED_ORDER }, MIN_LPC_ORDER, ALAC_MAX_LPC_ORDER, AE },

    { "max_prediction_order", NULL, OFFSET(max_prediction_order), AV_OPT_TYPE_INT, { .i64 = DEFAULT_MAX_PRED_ORDER }, MIN_LPC_ORDER, ALAC_MAX_LPC_ORDER, AE },

    { NULL },

};

static const AVClass alacenc_class = {

    .class_name = "alacenc",

    .item_name  = av_default_item_name,

    .option     = options,

    .version    = LIBAVUTIL_VERSION_INT,

};

const FFCodec ff_alac_encoder = {

    .p.name         = "alac",

    CODEC_LONG_NAME("ALAC (Apple Lossless Audio Codec)"),

    .p.type         = AVMEDIA_TYPE_AUDIO,

    .p.id           = AV_CODEC_ID_ALAC,

    .p.capabilities = AV_CODEC_CAP_DR1 | AV_CODEC_CAP_SMALL_LAST_FRAME |

                      AV_CODEC_CAP_ENCODER_REORDERED_OPAQUE,

    .priv_data_size = sizeof(AlacEncodeContext),

    .p.priv_class   = &alacenc_class,

    .init           = alac_encode_init,

    FF_CODEC_ENCODE_CB(alac_encode_frame),

    .close          = alac_encode_close,

    CODEC_CH_LAYOUTS_ARRAY(ff_alac_ch_layouts),

    CODEC_SAMPLEFMTS(AV_SAMPLE_FMT_S32P, AV_SAMPLE_FMT_S16P),

};