/*
 * 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 <mosek.h>
#include "common.h"

static void MSKAPI printstr( void *handle, char *str )
{
    printf( "%s", str );
}

#define ALLOC(type, name, size)\
    type *name = malloc( (size) * sizeof(type) );\
    assert(name);

#define MB(f,x,y) h->fdecs[(f)+1]->macroblocks[((x)>>4) + ((y)>>4)*h->mb_width]

/* order of variables (and thus also order of rows in the matrices):
 * dct coefs for the first frame, ordered first by block (bitstream order) then by raster scan within a block
 * dct coefs for the next frame, ...
 * reconstructed pixels for the first frame, in raster scan order
 * reconstructed pixels for the next frame, ...
 * bits estimate, in the same order as dct coefs */

void xlsc_lookahead_quant( xlsc_t *h )
{
    int frames = 1;
    // FIXME should be <= lookahead, but this way I can compare to the non-lookahead optimization problem
    while( frames < h->param.lookahead && h->fdec[frames+1].frame_num >= 0 )
        frames++;

    const int framepel = h->mb_count * 256; // number of pixels in one frame
    const int numpel = framepel * frames; // number of pixels in all frames
    const int numvar = numpel * 3; // number of variables = dct coefs + pixels + bits
    const int numcon = numpel * 3; // number of constraints = dct + 2x bits
    const int numanz = numpel * (64 + 4 + 1 + 4); // max nonzeros in constraints = dct + mc + pixel + bits
    const int numqnz = numpel; // number of nonzeros in quadratic objective

    ALLOC( int, bkc, numcon );
    ALLOC( int, bkx, numvar );
    ALLOC( int, aptrb, numvar );
    ALLOC( int, aptre, numvar );
    ALLOC( int, asub, numanz );
    ALLOC( int, qsubi, numqnz );
    ALLOC( int, qsubj, numqnz );
    ALLOC( double, blc, numcon );
    ALLOC( double, buc, numcon );
    ALLOC( double, blx, numvar );
    ALLOC( double, bux, numvar );
    ALLOC( double, aval, numanz );
    ALLOC( double, qval, numqnz );
    ALLOC( double, cx, numvar );
    MSKenv_t env;
    MSKrescodee r;
    MSKtask_t task;

    int dct2pel[framepel>>6];
    const int width = h->mb_width * 16;
    const int height = h->mb_height * 16;
    const int stride = h->stride;
    int f, i, j, k, t, x, y, b, mb_x, mb_y;

    /* logically we should multiply the idct_tab entries by the quantizer, but Mosek's
     * presolve simplifier runs faster if we divide the pixel values instead.
     * if different frames could have different qp, then we'd have to use the slow way. */
    const float quantf = 1.0 / (h->fdec->qp * 4);

    /* tables: could be precomputed */
    for( j=0, mb_y=0; mb_y<h->mb_height; mb_y++ )
        for( mb_x=0; mb_x<h->mb_width; mb_x++ )
            for( b=0; b<4; b++, j++ )
                dct2pel[j] = mb_x*16 + mb_y*16*width + (b&1)*8 + (b>>1)*8*width;
    assert( j == framepel>>6 );

    /* Mosek setup */
    r = MSK_makeenv( &env, NULL, NULL, NULL, NULL );
    assert( r == MSK_RES_OK );
    MSK_linkfunctoenvstream( env, MSK_STREAM_LOG, NULL, printstr );
    r = MSK_initenv( env );
    assert( r == MSK_RES_OK );
    r = MSK_maketask( env, numcon, numvar, &task );
    assert( r == MSK_RES_OK );
    r = MSK_linkfunctotaskstream( task, MSK_STREAM_LOG, NULL, printstr );
    assert( r == MSK_RES_OK );
    MSK_putintparam( task, MSK_IPAR_INTPNT_NUM_THREADS, h->param.threads );

    // disabling presolve makes it 1.2x faster.
    // there aren't any linearly dependent constraints, so that check is useless.
    // the eliminatior converts the abs(dct) construct from 2 constraints + 1 var
    // to 1 constraint + 2 vars, dunno if that matters.
    // most of the speedup comes from needing fewer iterations. but I don't know
    // if it actaully converges quicker, or if the problem is just scaled differently
    // so the same termination threshold means a different real precision.
    MSK_putintparam( task, MSK_IPAR_PRESOLVE_USE, 0 );


    /* unconstrained variables */
    // will constrain bits to be positive, but for now that's implied by the other inequalities
    for( i=0; i<numvar; i++ )
    {
        bkx[i] = MSK_BK_FR;
        blx[i] = bux[i] = 0.0;
    }

    /* linear equality constraints: includes the component of mc that comes from
     * previous frames outside the scope of optimization */
    for( i=0, f=0; f<frames; f++ )
        for( y=0; y<height; y++ )
            for( x=0; x<width; x++, i++ )
            {
                int pred = 0;
                mb_t mb = MB(f,x,y);
                if( mb.intra )
                    pred = 128;
                else if( f == 0 )
                {
                    // only the first frame gets constant mc values,
                    // the rest will get entries in the constraint matrix and 0 here
                    assert(!((mb.mv.x|mb.mv.y)&1)); // interpolation not implemented yet
                    pred = h->fref->plane[(x+(mb.mv.x>>1))+(y+(mb.mv.y>>1))*stride];
                }
                bkc[i] = MSK_BK_FX;
                blc[i] = buc[i] = - pred * quantf;
            }
    assert( i == numpel );

    /* linear inequality constraints: emulate bits=abs(dct) */
    for( i=numpel; i<numpel*3; i++ )
    {
        bkc[i] = MSK_BK_LO;
        blc[i] = 0.0;
        buc[i] = MSK_INFINITY;
    }

    /* idct */
    // one row of A contains one row of the idct matrix.
    // one column of A contains an entry for each pixel affected by this dct coef.
    // Mosek stores A in column-major order. If pixels were arranged into dct blocks
    // then the column would be contiguous, but I have pixels in raster order.
    t=0;
    for( j=0, f=0; f<frames; f++ )
        for( b=0; b<h->mb_count*4; b++ )
        {
            int pel = f*framepel + dct2pel[b];
            for( k=0; k<64; k++, j++ )
            {
                aptrb[j] = t;
                /* one column of idct */
                for( i=0; i<64; i++, t++ )
                {
                    asub[t] = pel + (i&7) + (i>>3)*width;
                    aval[t] = xlsc_idct_tab[i][k];
                }
                /* half of bits=abs(dct) */
                for( i=0; i<2; i++, t++ )
                {
                    asub[t] = numpel + 2*j + i;
                    aval[t] = i ? 1.0 : -1.0;
                }
                aptre[j] = t;
            }
        }

    /* motion compensation */
    for( j=0, f=0; f<frames; f++ )
        for( y=0; y<height; y++ )
            for( x=0; x<width; x++, j++ )
            {
                aptrb[numpel+j] = t;
                if( f > 0 )
                {
                    // mc
                    mb_t mb = MB(f,x,y);
                    if( !mb.intra )
                    {
                        assert(!((mb.mv.x|mb.mv.y)&1)); // interpolation not implemented yet
                        asub[t] = j - framepel + (mb.mv.x>>1) + (mb.mv.y>>1)*stride;
                        aval[t] = 1.0;
                        t++;
                    }
                }
                asub[t] = j;
                aval[t] = -1.0; // set mc+idct equal to the reconstructed pixel
                t++;
                aptre[numpel+j] = t;
            }

    /* other half of bits=abs(dct) */
    for( j=0; j<numpel; j++ )
    {
        aptrb[numpel*2+j] = t;
        for( i=0; i<2; i++, t++ )
        {
            asub[t] = numpel + 2*j + i;
            aval[t] = 1.0;
        }
        aptre[numpel*2+j] = t;
    }

    /* objective function - quadratic */
    for( i=0; i<numqnz; i++ )
    {
        qsubi[i] = numpel+i;
        qsubj[i] = numpel+i;
        qval[i] = 1.0; // 1/lambda
    }

    /* objective function - linear */
    for( i=0; i<numpel; i++ )
        cx[i] = 0.0; // DCT coefs aren't included in the objective
    for( y=0; y<height; y++ )
        for( x=0; x<width; x++, i++ )
            cx[i] = - h->fenc->plane[x+y*stride] * quantf; // * 1/lambda
    for( ; i<numpel*3; i++ )
        cx[i] = h->param.lookahead_lambda; // bits

    r = MSK_inputdata( task, numcon, numvar, numcon, numvar,
                       cx, 0.0, aptrb, aptre, asub, aval, bkc, blc, buc, bkx, blx, bux );
    assert( r == MSK_RES_OK );

    r = MSK_putqobj( task, numqnz, qsubi, qsubj, qval );
    assert( r == MSK_RES_OK );

    free( bkc );
    free( bkx );
    free( aptrb );
    free( aptre );
    free( asub );
    free( qsubi );
    free( qsubj );
    free( blc );
    free( buc );
    free( blx );
    free( bux );
    free( aval );
    free( qval );
    free( cx );

    r = MSK_optimize( task );
    assert( r == MSK_RES_OK );

    ALLOC( double, dct_coef, framepel );
    MSK_getsolutionslice( task, MSK_SOL_ITR, MSK_SOL_ITEM_XX, 0, framepel, dct_coef );

    MSK_deletetask( &task );
    MSK_deleteenv( &env );

    for( i=0; i<framepel; i++ )
        h->fdec->dct[0][0][i] = rint(dct_coef[i]);

    free( dct_coef );
}
