/*
 * The simplest mpeg encoder
 * Copyright (c) 2000 Gerard Lantau.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program 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
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */
#include <stdlib.h>
#include <stdio.h>
#include <netinet/in.h>
#include <math.h>
#include "avcodec.h"
#include "dsputil.h"
#include "mpegvideo.h"

/* enable all paranoid tests for rounding, overflows, etc... */
//#define PARANOID

//#define DEBUG

/* for jpeg fast DCT */
#define CONST_BITS 14

static const unsigned short aanscales[64] = {
    /* precomputed values scaled up by 14 bits */
    16384, 22725, 21407, 19266, 16384, 12873,  8867,  4520,
    22725, 31521, 29692, 26722, 22725, 17855, 12299,  6270,
    21407, 29692, 27969, 25172, 21407, 16819, 11585,  5906,
    19266, 26722, 25172, 22654, 19266, 15137, 10426,  5315,
    16384, 22725, 21407, 19266, 16384, 12873,  8867,  4520,
    12873, 17855, 16819, 15137, 12873, 10114,  6967,  3552,
    8867, 12299, 11585, 10426,  8867,  6967,  4799,  2446,
    4520,  6270,  5906,  5315,  4520,  3552,  2446,  1247
};

static void encode_picture(MpegEncContext *s, int picture_number);
static void rate_control_init(MpegEncContext *s);
static int rate_estimate_qscale(MpegEncContext *s);
static void mpeg1_skip_picture(MpegEncContext *s, int pict_num);

#include "mpegencodevlc.h"

static void put_header(MpegEncContext *s, int header)
{
    align_put_bits(&s->pb);
    put_bits(&s->pb, 32, header);
}

static void convert_matrix(int *qmat, const UINT8 *quant_matrix, int qscale)
{
    int i;

    for(i=0;i<64;i++) {
        qmat[i] = (int)((1 << 22) * 16384.0 / (aanscales[i] * qscale * quant_matrix[i]));
    }
}


int MPV_encode_init(AVEncodeContext *avctx)
{
    MpegEncContext *s = avctx->priv_data;
    int pict_size, c_size;
    UINT8 *pict;

    dsputil_init();

    s->bit_rate = avctx->bit_rate;
    s->frame_rate = avctx->rate;
    s->width = avctx->width;
    s->height = avctx->height;
    s->gop_size = avctx->gop_size;
    if (s->gop_size <= 1) {
        s->intra_only = 1;
        s->gop_size = 12;
    } else {
        s->intra_only = 0;
    }
    s->full_search = 0;
    if (avctx->flags & CODEC_FLAG_HQ)
        s->full_search = 1;

    switch(avctx->codec->id) {
    case CODEC_ID_MPEG1VIDEO:
        s->out_format = FMT_MPEG1;
        break;
    case CODEC_ID_MJPEG:
        s->out_format = FMT_MJPEG;
        s->intra_only = 1; /* force intra only for jpeg */
        if (mjpeg_init(s) < 0)
            return -1;
        break;
    case CODEC_ID_H263:
        s->out_format = FMT_H263;
        break;
    case CODEC_ID_RV10:
        s->out_format = FMT_H263;
        s->h263_rv10 = 1;
        break;
    case CODEC_ID_DIVX:
        s->out_format = FMT_H263;
        s->h263_pred = 1;
        break;
    default:
        return -1;
    }

    switch(s->frame_rate) {
    case 24:
        s->frame_rate_index = 2;
        break;
    case 25:
        s->frame_rate_index = 3;
        break;
    case 30:
        s->frame_rate_index = 5;
        break;
    case 50:
        s->frame_rate_index = 6;
        break;
    case 60:
        s->frame_rate_index = 8;
        break;
    default:
        /* we accept lower frame rates than 24 for low bit rate mpeg */
        if (s->frame_rate >= 1 && s->frame_rate < 24) {
            s->frame_rate_index = 2;
        } else {
            return -1;
        }
        break;
    }

    /* init */
    s->mb_width = s->width / 16;
    s->mb_height = s->height / 16;
    
    c_size = s->width * s->height;
    pict_size = (c_size * 3) / 2;
    pict = malloc(pict_size);
    if (pict == NULL)
        goto fail;
    s->last_picture[0] = pict;
    s->last_picture[1] = pict + c_size;
    s->last_picture[2] = pict + c_size + (c_size / 4);
    
    pict = malloc(pict_size);
    if (pict == NULL) 
        goto fail;
    s->current_picture[0] = pict;
    s->current_picture[1] = pict + c_size;
    s->current_picture[2] = pict + c_size + (c_size / 4);

    if (s->out_format == FMT_H263) {
        int size;
        /* MV prediction */
        size = (2 * s->mb_width + 2) * (2 * s->mb_height + 2);
        s->motion_val = malloc(size * 2 * sizeof(INT16));
        if (s->motion_val == NULL)
            goto fail;
        memset(s->motion_val, 0, size * 2 * sizeof(INT16));
    }

    if (s->h263_pred) {
        int y_size, c_size, i, size;

        y_size = (2 * s->mb_width + 2) * (2 * s->mb_height + 2);
        c_size = (s->mb_width + 2) * (s->mb_height + 2);
        size = y_size + 2 * c_size;
        s->dc_val[0] = malloc(size * sizeof(INT16));
        if (s->dc_val[0] == NULL)
            goto fail;
        s->dc_val[1] = s->dc_val[0] + y_size;
        s->dc_val[2] = s->dc_val[1] + c_size;
        for(i=0;i<size;i++)
            s->dc_val[0][i] = 1024;
    }

    /* rate control init */
    rate_control_init(s);

    s->picture_number = 0;
    s->fake_picture_number = 0;
    /* motion detector init */
    s->f_code = 1;

    return 0;
 fail:
    if (s->motion_val)
        free(s->motion_val);
    if (s->dc_val[0])
        free(s->dc_val[0]);
    if (s->last_picture[0])
        free(s->last_picture[0]);
    if (s->current_picture[0])
        free(s->current_picture[0]);
    return -1;
}

int MPV_encode_end(AVEncodeContext *avctx)
{
    MpegEncContext *s = avctx->priv_data;
#if 0
    /* end of sequence */
    if (s->out_format == FMT_MPEG1) {
        put_header(s, SEQ_END_CODE);
    }

    if (!s->flush_frames)
        flush_put_bits(&s->pb);
#endif    
    if (s->motion_val)
        free(s->motion_val);
    if (s->h263_pred)
        free(s->dc_val[0]);
    free(s->last_picture[0]);
    free(s->current_picture[0]);
    if (s->out_format == FMT_MJPEG)
        mjpeg_close(s);
    return 0;
}

int MPV_encode_picture(AVEncodeContext *avctx,
                       unsigned char *buf, int buf_size, void *data)
{
    MpegEncContext *s = avctx->priv_data;
    int i;

    memcpy(s->new_picture, data, 3 * sizeof(UINT8 *));

    init_put_bits(&s->pb, buf, buf_size, NULL, NULL);

    /* group of picture */
    if (s->out_format == FMT_MPEG1) {
        unsigned int vbv_buffer_size;
        unsigned int time_code, fps, n;

        if ((s->picture_number % s->gop_size) == 0) {
            /* mpeg1 header repeated every gop */
            put_header(s, SEQ_START_CODE);
            
            put_bits(&s->pb, 12, s->width);
            put_bits(&s->pb, 12, s->height);
            put_bits(&s->pb, 4, 1); /* 1/1 aspect ratio */
            put_bits(&s->pb, 4, s->frame_rate_index);
            put_bits(&s->pb, 18, 0x3ffff);
            put_bits(&s->pb, 1, 1); /* marker */
            /* vbv buffer size: slightly greater than an I frame. We add
               some margin just in case */
            vbv_buffer_size = (3 * s->I_frame_bits) / (2 * 8);
            put_bits(&s->pb, 10, (vbv_buffer_size + 16383) / 16384); 
            put_bits(&s->pb, 1, 1); /* constrained parameter flag */
            put_bits(&s->pb, 1, 0); /* no custom intra matrix */
            put_bits(&s->pb, 1, 0); /* no custom non intra matrix */

            put_header(s, GOP_START_CODE);
            put_bits(&s->pb, 1, 0); /* do drop frame */
            /* time code : we must convert from the real frame rate to a
               fake mpeg frame rate in case of low frame rate */
            fps = frame_rate_tab[s->frame_rate_index];
            time_code = s->fake_picture_number;
            s->gop_picture_number = time_code;
            put_bits(&s->pb, 5, (time_code / (fps * 3600)) % 24);
            put_bits(&s->pb, 6, (time_code / (fps * 60)) % 60);
            put_bits(&s->pb, 1, 1);
            put_bits(&s->pb, 6, (time_code / fps) % 60);
            put_bits(&s->pb, 6, (time_code % fps));
            put_bits(&s->pb, 1, 1); /* closed gop */
            put_bits(&s->pb, 1, 0); /* broken link */
        }

        if (s->frame_rate < 24 && s->picture_number > 0) {
            /* insert empty P pictures to slow down to the desired
               frame rate. Each fake pictures takes about 20 bytes */
            fps = frame_rate_tab[s->frame_rate_index];
            n = ((s->picture_number * fps) / s->frame_rate) - 1;
            while (s->fake_picture_number < n) {
                mpeg1_skip_picture(s, s->fake_picture_number - 
                                   s->gop_picture_number); 
                s->fake_picture_number++;
            }

        }
        s->fake_picture_number++;
    }
    
    
    if (!s->intra_only) {
        /* first picture of GOP is intra */
        if ((s->picture_number % s->gop_size) == 0)
            s->pict_type = I_TYPE;
        else
            s->pict_type = P_TYPE;
    } else {
        s->pict_type = I_TYPE;
    }
    avctx->key_frame = (s->pict_type == I_TYPE);
    
    encode_picture(s, s->picture_number);
    
    /* swap current and last picture */
    for(i=0;i<3;i++) {
        UINT8 *tmp;
        
        tmp = s->last_picture[i];
        s->last_picture[i] = s->current_picture[i];
        s->current_picture[i] = tmp;
    }
    s->picture_number++;

    if (s->out_format == FMT_MJPEG)
        mjpeg_picture_trailer(s);

    flush_put_bits(&s->pb);
    s->total_bits += (s->pb.buf_ptr - s->pb.buf) * 8;
    avctx->quality = s->qscale;
    return s->pb.buf_ptr - s->pb.buf;
}

/* insert a fake P picture */
static void mpeg1_skip_picture(MpegEncContext *s, int pict_num)
{
    unsigned int mb_incr;

    /* mpeg1 picture header */
    put_header(s, PICTURE_START_CODE);
    /* temporal reference */
    put_bits(&s->pb, 10, pict_num & 0x3ff); 
    
    put_bits(&s->pb, 3, P_TYPE);
    put_bits(&s->pb, 16, 0xffff); /* non constant bit rate */
    
    put_bits(&s->pb, 1, 1); /* integer coordinates */
    put_bits(&s->pb, 3, 1); /* forward_f_code */
    
    put_bits(&s->pb, 1, 0); /* extra bit picture */
    
    /* only one slice */
    put_header(s, SLICE_MIN_START_CODE);
    put_bits(&s->pb, 5, 1); /* quantizer scale */
    put_bits(&s->pb, 1, 0); /* slice extra information */
    
    mb_incr = 1;
    put_bits(&s->pb, mbAddrIncrTable[mb_incr][1], 
             mbAddrIncrTable[mb_incr][0]);
    
    /* empty macroblock */
    put_bits(&s->pb, 3, 1); /* motion only */
    
    /* zero motion x & y */
    put_bits(&s->pb, 1, 1); 
    put_bits(&s->pb, 1, 1); 

    /* output a number of empty slice */
    mb_incr = s->mb_width * s->mb_height - 1;
    while (mb_incr > 33) {
        put_bits(&s->pb, 11, 0x008);