/*
 * Copyright (C) 2007 Loren Merritt <lorenm@u.washington.edu>
 *
 * 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., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
 */

#include "common.h"

static float ssd_to_psnr( uint64_t ssd, uint64_t pixels )
{
    float msd = (float)ssd / (65025.0 * (float)pixels);
    if( msd <= 0.0000000001 ) /* Max 100dB */
        return 100;
    return -10.0 * log(msd) / log(10.0);
}

static void update_frame_stats( xlsc_t *h )
{
    int y;
    if( h->fdec->frame_num < 0 )
        return;
    h->stat.ssd += xlsc_ssd_plane( h->fenc->plane, h->fdec->plane, h->stride, h->param.width, h->param.height );
    h->stat.frames++;
    if( h->out.dump )
        for( y=0; y<h->param.height; y++ )
            fwrite( h->fdec->plane+y*h->stride, h->param.width, 1, h->out.dump );
}

static void print_frame_stats( xlsc_t *h )
{
#ifdef DEBUG_RAC_COUNT
    printf( "RAC mb:%10"PRIu64"\n", h->stat.rac_mb );
    printf( "RAC mv:%10"PRIu64"\n", h->stat.rac_mv );
    printf( "RAC cbp:%9"PRIu64"\n", h->stat.rac_cbp );
    printf( "RAC nnz:%9"PRIu64"\n", h->stat.rac_nnz );
    printf( "RAC eob:%9"PRIu64"\n", h->stat.rac_eob );
    printf( "RAC coef:%8"PRIu64"\n", h->stat.rac_coef );
#endif

    uint64_t mbs = h->stat.mb_type[0] + h->stat.mb_type[1];
    printf( "MB I:%4.1f%% P:%4.1f%%\n",
            h->stat.mb_type[1]*100./mbs, h->stat.mb_type[0]*100./mbs );

    uint64_t pixels = (uint64_t)h->stat.frames * h->param.width * h->param.height;
    printf( "PSNR: %5.2f\n", ssd_to_psnr( h->stat.ssd, pixels ) );
}

static void write_stream_header( xlsc_t *h, uint8_t **buf, int *size )
{
    *buf = h->out.buffer;
    memcpy( *buf, "XLSC", 4 );
    (*buf)[4] = h->param.width>>8;
    (*buf)[5] = h->param.width;
    (*buf)[6] = h->param.height>>8;
    (*buf)[7] = h->param.height;
    *size = 8;
}

static int validate_parameters( xlsc_t *h )
{
    h->param.qp = clip( h->param.qp, 1, 127 );
    h->param.deadzone = clip( h->param.deadzone, 0, 1000 );
    h->param.lookahead = clip( h->param.lookahead, 0, 16 );
#ifdef HAVE_LINALG
    if( h->param.lookahead && h->param.umv )
    {
        fprintf( stderr, "xlsc [error]: no lookahead + umv support yet\n" );
        return -1;
    }
#else
    if( h->param.lookahead )
        fprintf( stderr, "xlsc [warning]: not compiled with libmosek, lookahead will be a nop.\n" );
#endif
    return 0;
}

xlsc_t *xlsc_encoder_open( xlsc_param_t *param, uint8_t **header, int *header_size )
{
    int i;
    xlsc_t *h = malloc( sizeof(xlsc_t) );
    memset( h, 0, sizeof(xlsc_t) );
    memcpy( &h->param, param, sizeof(xlsc_param_t) );

    if( validate_parameters( h ) < 0 )
    {
        free( h );
        return NULL;
    }

    if( h->param.dump_file )
    {
        h->out.dump = fopen( h->param.dump_file, "wb" );
        if( !h->out.dump )
        {
            free( h );
            return NULL;
        }
    }

    xlsc_dsp_init( h );

    if( h->param.threads == 0 )
        h->param.threads = xlsc_cpu_num_processors() > 1 ? 2 : 1;
    else
        h->param.threads = clip( h->param.threads, 1, 16 );

    h->out.buffer_size = 10000000;
    h->out.buffer = malloc( h->out.buffer_size );

    h->mb_width  = (param->width  + 15) / 16;
    h->mb_height = (param->height + 15) / 16;
    h->mb_count = h->mb_width * h->mb_height;
    h->stride = 16*h->mb_width + 2*FRAME_PAD;

    // queue + current frame + reference + sentry
    h->fencs = malloc( (h->param.lookahead + 3) * sizeof(xlsc_frame_t*) );
    h->fdecs = malloc( (h->param.lookahead + 3) * sizeof(xlsc_frame_t*) );

    for( i=0; i < h->param.lookahead + 2; i++ )
    {
        h->fencs[i] = xlsc_frame_new( h );
        h->fdecs[i] = xlsc_frame_new( h );
    }
    h->fencs[h->param.lookahead + 2] = NULL;
    h->fdecs[h->param.lookahead + 2] = NULL;

    h->fenc = h->fencs[1];
    h->fdec = h->fdecs[1];
    h->fref = h->fdecs[0];

    xlsc_macroblock_init( h );
    xlsc_build_rac_states( &h->out.c, 0.05*(1LL<<32), 256-8 );

    write_stream_header( h, header, header_size );

    return h;
}

int xlsc_encoder_close( xlsc_t *h )
{
    print_frame_stats( h );
    while( h->fencs[0] )
        xlsc_frame_free( xlsc_frame_pop( h->fencs ) );
    while( h->fdecs[0] )
        xlsc_frame_free( xlsc_frame_pop( h->fdecs ) );
    free( h->fencs );
    free( h->fdecs );
    free( h->out.buffer );
    free( h );
    return 0;
}

static void load_rac_model( xlsc_t *h )
{
    if( h->fdec->intra )
        memset( &h->model, 128, sizeof(h->model) );
    else
        h->model = h->fref->model;
}

static void write_frame_header( xlsc_t *h )
{
    assert(h->fdec->qp >= 1 && h->fdec->qp <= 127);
    h->out.buffer[0] = (h->fdec->intra << 7) | h->fdec->qp;
    if( h->fdec->intra )
        memset( &h->model, 128, sizeof(h->model) );
    xlsc_init_range_encoder( &h->out.c, h->out.buffer+1, h->out.buffer_size-1 );
}

static void write_frame_footer( xlsc_t *h )
{
    h->out.frame_size = xlsc_rac_terminate( &h->out.c ) + 1;
}

static void update_reference_list( xlsc_t *h )
{
    xlsc_frame_push( h->fdecs, xlsc_frame_shift( h->fdecs ) );
    xlsc_frame_push( h->fencs, xlsc_frame_shift( h->fencs ) );
    h->gop_size++;
}

static void encode_frame( xlsc_t *h )
{
    write_frame_header( h );
    xlsc_macroblock_frame_start( h );

    int x, y;
    for( y = 0; y < h->mb_height; y++ )
        for( x = 0; x < h->mb_width; x++ )
        {
            xlsc_macroblock_load( h, x, y );
            if( !h->fdec->intra )
                xlsc_macroblock_analyse( h );
            if( !h->param.mb_decision || h->mb.intra )
                xlsc_macroblock_encode( h );
            xlsc_macroblock_write( h );
            xlsc_macroblock_save( h );
        }

    write_frame_footer( h );
    xlsc_frame_expand_border( h, h->fdec );
}

static void encode_frame_lookahead( xlsc_t *h )
{
    int x, y, b;

    if( h->fdec->frame_num >= 0 )
    {
        load_rac_model( h );
        xlsc_macroblock_frame_start( h );
        for( y = 0; y < h->mb_height; y++ )
            for( x = 0; x < h->mb_width; x++ )
            {
                xlsc_macroblock_load( h, x, y );
                if( !h->fdec->intra )
                    xlsc_macroblock_analyse( h );
                if( !h->param.mb_decision || h->mb.intra )
                    xlsc_macroblock_encode( h );
                xlsc_macroblock_rac_update( h );
                xlsc_macroblock_save( h );
                for( b=0; b<4; b++ )
                    xlsc_unzigzag( h->fdec->dct[h->mb.xy][b], h->mb.dct_zigzag[b] );
            }
        xlsc_frame_expand_border( h, h->fdec );
        h->fdec->model = h->model;
    }

    h->fenc = h->fencs[1];
    h->fdec = h->fdecs[1];
    h->fref = h->fdecs[0];

    if( h->fdec->frame_num >= 0 )
    {
#ifdef HAVE_LINALG
        xlsc_lookahead_quant( h );
#endif

        load_rac_model( h );
        write_frame_header( h );
        xlsc_macroblock_frame_start( h );
        h->decoding = 1;
        for( y = 0; y < h->mb_height; y++ )
            for( x = 0; x < h->mb_width; x++ )
            {
                xlsc_macroblock_load( h, x, y );
                h->mb.intra = h->fdec->macroblocks[h->mb.xy].intra;
                h->mb.mv = h->fdec->macroblocks[h->mb.xy].mv;
                memcpy( h->mb.dct, h->fdec->dct[h->mb.xy], sizeof(h->mb.dct) );
                for( b=0; b<4; b++ )
                    xlsc_zigzag( h->mb.dct_zigzag[b], h->mb.dct[b] );
                xlsc_macroblock_write( h );
                xlsc_macroblock_decode( h );
            }
        h->decoding = 0;
        write_frame_footer( h );
        xlsc_frame_expand_border( h, h->fdec );
        h->fdec->model = h->model;
    }
}

int xlsc_encoder_encode( xlsc_t *h, uint8_t **buffer, int *buffer_size, xlsc_picture_t *pic )
{
    xlsc_frame_t *fenc = h->fenc = h->fencs[1 + h->param.lookahead];
    xlsc_frame_t *fdec = h->fdec = h->fdecs[1 + h->param.lookahead];
                         h->fref = h->fdecs[0 + h->param.lookahead];

    if( pic )
    {
        xlsc_frame_from_pic( h, fenc, pic );
        fdec->frame_num =
        fenc->frame_num = h->stat.frames;
        if( h->gop_size >= h->param.keyint )
            h->gop_size = 0;
        fdec->intra =
        fenc->intra = h->gop_size == 0;
        fdec->qp =
        fenc->qp = h->param.qp;
    }
    else
    {
        fdec->frame_num =
        fenc->frame_num = -1;
    }

    if( h->param.lookahead )
        encode_frame_lookahead( h );
    else
        encode_frame( h );

    update_frame_stats( h );
    update_reference_list( h );

    if( h->fdec->frame_num >= 0 )
    {
        *buffer = h->out.buffer;
        *buffer_size = h->out.frame_size;
    }
    else
    {
        *buffer = NULL;
        *buffer_size = 0;
    }

    return 0;
}
