/*
    Project  : H.264 encoder test    
    Module   : Intra coding related functions

    Author   : Marko 'Fador' Viitanen

    Date     : 30/06/2009
    Modified : 01/07/2009

*/

#include <stdio.h>
#include "common.h"
#include "psnr.h"
#include "intra.h"


//Macro for result comparing
#define COMPARE(MODE) \
    tempsad=SAD(cur4x4L, pred4x4L,(uint32)4,(uint32)4); \
    if(tempsad<bestcost) \
    {\
        bestcost=tempsad;\
        bestIntra4x4PredMode=MODE;\
        for(temp_x=0;temp_x<4;temp_x++) \
        {\
            for(temp_y=0;temp_y<4;temp_y++)\
            {\
                residual[temp_x+temp_y*4]=(cur4x4L[temp_x+temp_y*4]);\
                residual[temp_x+temp_y*4]-=(pred4x4L[temp_x+temp_y*4]);\
            }\
        }\
    }


#define RETURNPRED() \
        for(temp_x=0;temp_x<4;temp_x++) \
        {\
            for(temp_y=0;temp_y<4;temp_y++)\
            {\
                residual[temp_x+temp_y*4]=pred4x4L[temp_x+temp_y*4];\
            }\
        }\
        return 1;



#define P(X,Y) p[(X)+1+((Y)+1)*9]


typedef uint8 (modefunc)(uint8 *pred4x4L, uint8 *p, uint8 *pred4x4L_avail);

//Recommendation page 126
//Mode -1 = encode all and select best
sint8 intrapred_4x4(uint8 *frame, sint32 *residual, uint32 x, uint32 y, uint32 xoff, uint32 yoff, uint32 *cost, sint8 mode)
{
    
    uint32 temp_x,temp_y,tempsad;
    uint32 bestcost=0xffffffff;
    sint8 bestIntra4x4PredMode=-1;
    uint8 Intra4x4PredMode;    

    //Extended sub-macroblock
    uint8 p[9*5];

    uint8 pred4x4L[4*4];
    uint8 cur4x4L[4*4];    

    //Different prediction mode functions here
    modefunc *funcs[9]={&intraPred4x4Mode0,&intraPred4x4Mode1,&intraPred4x4Mode2,
                        &intraPred4x4Mode3,&intraPred4x4Mode4,&intraPred4x4Mode5,
                        &intraPred4x4Mode6,&intraPred4x4Mode7,&intraPred4x4Mode8};

    //Fill the availability array
    uint8 pred4x4L_avail[9*5];
    //Fill with ones
    memset(pred4x4L_avail,1,9*5);

    //First column not available at xoff==0
    if(xoff==0)
    {
        pred4x4L_avail[0]=pred4x4L_avail[9]=pred4x4L_avail[18]=pred4x4L_avail[28]=pred4x4L_avail[36]=0;
    }
    else
    {
        temp_x=0;
        //Fill I,J,K,L
        for(temp_y=1;temp_y<5;temp_y++)
            p[temp_y*9+temp_x]=frame[xoff+(temp_x-1)+(yoff+(temp_y-1))*x];
    }

    //First row not available on yoff==0
    if(yoff==0)
    {
         pred4x4L_avail[0]=pred4x4L_avail[1]=pred4x4L_avail[2]=pred4x4L_avail[3]=
             pred4x4L_avail[4]=pred4x4L_avail[5]=pred4x4L_avail[6]=pred4x4L_avail[7]=pred4x4L_avail[8]=0;
    }
    else
    {
        temp_y=0;
        //Fill (M),A,B,C & D
        for(temp_x=(pred4x4L_avail[0]?0:1);temp_x<5;temp_x++)
            p[temp_y*9+temp_x]=frame[xoff+(temp_x-1)+(yoff+(temp_y-1))*x];
    }


    //E-H not available at right side
    if((xoff%16 == 15 && yoff%16!=0) || (x-xoff)==4)
    {        
        pred4x4L_avail[5]=pred4x4L_avail[6]=pred4x4L_avail[7]=pred4x4L_avail[8]=0;
    }
    else
    {
        temp_y=0;
        //Fill E,F,G & H
        for(temp_x=5;temp_x<9;temp_x++)
            p[temp_y*9+temp_x]=frame[xoff+(temp_x-1)+(yoff+(temp_y-1))*x];
    }
    


    //Current MB
    for(temp_x=0;temp_x<4;temp_x++)
            for(temp_y=0;temp_y<4;temp_y++)
                cur4x4L[temp_x+temp_y*4]=frame[xoff+temp_x+(yoff+temp_y)*x];

    //If we return only one prediction result
    if(mode!=-1)
    {
        (*funcs[mode])(pred4x4L,p,pred4x4L_avail);
        RETURNPRED();       
    }
    else
    {
        //Loop through all prediction modes calling the functions
        //and if prediction is executed, compare the SAD of the results to
        //current best cost
        for(Intra4x4PredMode=0;Intra4x4PredMode<8;Intra4x4PredMode++)
        {
            if((*funcs[Intra4x4PredMode])(pred4x4L,p,pred4x4L_avail))
            {
                COMPARE(Intra4x4PredMode);
            }
        }
    }

    *cost=bestcost;

    return bestIntra4x4PredMode;
}

//MODE0 - Intra_4x4_Vertical
// This mode shall be used only when the samples p[ x, -1 ] with x = 0..3 are marked as 
//"available for Intra_4x4prediction".
uint8 intraPred4x4Mode0(uint8 *pred4x4L, uint8 *p, uint8 *pred4x4L_avail)
{
    uint32 temp_x;
    uint32 temp_y;

    if(pred4x4L_avail[1] //&& pred4x4L_avail[2] &&
        //pred4x4L_avail[3] &&pred4x4L_avail[4]
        )
    {
        for(temp_x=0;temp_x<4;temp_x++)
        {
            for(temp_y=0;temp_y<4;temp_y++)
            {
                pred4x4L[ temp_x+temp_y*4 ] = P(temp_x, -1);
            }
        }      
        return 1;
    }
    return 0;
}

//MODE1 - Intra_4x4_Horizontal
//This mode shall be used only when the samples p[ -1, y ], 
//with y = 0..3 are marked as "available for Intra_4x4 prediction". 
uint8 intraPred4x4Mode1(uint8 *pred4x4L, uint8 *p, uint8 *pred4x4L_avail)
{
    uint32 temp_x;
    uint32 temp_y;
   
    if(pred4x4L_avail[9] // && pred4x4L_avail[18] &&
       //pred4x4L_avail[27] && pred4x4L_avail[36]
        )
    {
        for(temp_x=0;temp_x<4;temp_x++)
            for(temp_y=0;temp_y<4;temp_y++)
                pred4x4L[ temp_x+temp_y*4 ] = P(-1,temp_y);
        return 1;
    }
    return 0;
}

//MODE2 - Intra_4x4_DC
//If all samples p[ x, -1 ], with x = 0..3 and p[ -1, y ], 
//with y = 0..3 are marked as "available for Intra_4x4 prediction"
uint8 intraPred4x4Mode2(uint8 *pred4x4L, uint8 *p, uint8 *pred4x4L_avail)
{
    uint32 temp_x;
    uint32 temp_y;

    if(pred4x4L_avail[9] && //pred4x4L_avail[18] &&
        //pred4x4L_avail[27] &&pred4x4L_avail[36] &&
        pred4x4L_avail[1] //&& pred4x4L_avail[2] &&
        //pred4x4L_avail[3] &&pred4x4L_avail[4]
        )
    {
        for(temp_x=0;temp_x<4;temp_x++)
            for(temp_y=0;temp_y<4;temp_y++)
                pred4x4L[ temp_x+temp_y*4 ] = (P(0,-1) + P(1,-1) + P(2,-1) + 
                                               P(3,-1) + P(-1,0) + P(-1,1) +
                                               P(-1,2) + P(-1,3) + 4) >> 3;
        return 1;
    }
    //- Otherwise, if any samples p[ x, -1 ], with x = 0..3 are marked as 
    //"not available for Intra_4x4 prediction" and all
    //samples p[ -1, y ], with y = 0..3 are marked as "available for Intra_4x4 prediction"
    else if(pred4x4L_avail[9]// && pred4x4L_avail[18] &&
            //pred4x4L_avail[27] &&pred4x4L_avail[36]
            )
    {
        for(temp_y=0;temp_y<4;temp_y++)
            for(temp_x=0;temp_x<4;temp_x++)
                pred4x4L[ temp_x+temp_y*4 ] = (P(-1,0)+P(-1,1)+P(-1,2)+P(-1,3)+2)>>2;
        return 1;
    }
    //- Otherwise, if any samples p[ -1, y ], with y = 0..3 are marked as 
    //"not available for Intra_4x4 prediction" and all
    //samples p[ x, -1 ], with x = 0 .. 3 are marked as "available for Intra_4x4 prediction"
    else if(pred4x4L_avail[1] //&& pred4x4L_avail[2] &&
        //pred4x4L_avail[3] &&pred4x4L_avail[4]
            )
    {
        for(temp_y=0;temp_y<4;temp_y++)
            for(temp_x=0;temp_x<4;temp_x++)
                pred4x4L[ temp_x+temp_y*4 ] = (P(0,-1)+P(1,-1)+P(2,-1)+P(3,-1)+2)>>2; 
        return 1;
    }
    return 0;
}

//MODE3 - Intra_4x4_Diagonal_Down_Left
//This mode shall be used only when the samples p[ x, -1 ] with x = 0..7 are marked as 
//"available for Intra_4x4 prediction".
uint8 intraPred4x4Mode3(uint8 *pred4x4L, uint8 *p, uint8 *pred4x4L_avail)
{
    uint32 temp_x;
    uint32 temp_y;
    if(pred4x4L_avail[1] && //pred4x4L_avail[2] &&
       //pred4x4L_avail[3] && pred4x4L_avail[4] &&
       pred4x4L_avail[5] //&& pred4x4L_avail[6] &&
       //pred4x4L_avail[7] && pred4x4L_avail[8]
        )
    {
        for(temp_y=0;temp_y<4;temp_y++)
        {
            for(temp_x=0;temp_x<4;temp_x++)
            {
                if(temp_y==3&&temp_x==3)
                {
                    pred4x4L[temp_x+ temp_y*4] = ( P(6,-1) + 3 * P(7,-1) + 2 ) >> 2;
                }
                else
                {
                    pred4x4L[temp_x+ temp_y*4] = ( P((temp_x+temp_y),-1) + 2 * P(temp_x+temp_y+1,-1) + P(temp_x+temp_y+2,-1)+2) >> 2;
                }
            }
        }
        return 1;
    }
    return 0;
}

//MODE4 - Intra_4x4_Diagonal_Down_Right
//This mode shall be used only when the samples 
//p[ x, -1 ] with x = 0..3 and p[ -1, y ] with y = -1..3 are marked as
//"available for Intra_4x4 prediction
uint8 intraPred4x4Mode4(uint8 *pred4x4L, uint8 *p, uint8 *pred4x4L_avail)
{
    uint32 temp_x;
    uint32 temp_y;
    if(pred4x4L_avail[1] &&// pred4x4L_avail[2] &&
       //pred4x4L_avail[3] && pred4x4L_avail[4] &&
       pred4x4L_avail[0] //&& pred4x4L_avail[9] &&
       //pred4x4L_avail[18]&& pred4x4L_avail[27]&&
       //pred4x4L_avail[36]
       )
    {

        for(temp_x=0;temp_x<4;temp_x++)
        {
            for(temp_y=0;temp_y<4;temp_y++)
            {
                if(temp_x>temp_y)
                    pred4x4L[temp_x+temp_y*4] = ( P(temp_x-temp_y-2,-1) + 2 * P(temp_x-temp_y-1,-1) + P(temp_x-temp_y,-1) + 2 ) >> 2;
                else if(temp_x<temp_y)
                    pred4x4L[temp_x+temp_y*4] = ( P(-1,temp_y-temp_x-2) + 2 * P(-1,temp_y-temp_x-1) + P(-1,temp_y-temp_x) + 2 ) >> 2;
                else
                    pred4x4L[temp_x+temp_y*4] =( P(0,-1) + 2 * P(-1,-1) + P(-1,0) + 2 ) >> 2;
            }
        }
        return 1;
    }
    return 0;
}

//MODE5 - Intra_4x4_Vertical_Right
//This mode shall be used only when the samples 
//p[ x, -1 ] with x = 0..3 and p[ -1, y ] with y = -1..3 are marked as
//"available for Intra_4x4 prediction".
uint8 intraPred4x4Mode5(uint8 *pred4x4L, uint8 *p, uint8 *pred4x4L_avail)
{
    uint32 temp_x;
    uint32 temp_y;
    if(pred4x4L_avail[1] &&// pred4x4L_avail[2] &&
       //pred4x4L_avail[3] && pred4x4L_avail[4] &&
       pred4x4L_avail[0] //&& pred4x4L_avail[9] &&
       //pred4x4L_avail[18] && pred4x4L_avail[27]&&
       //pred4x4L_avail[36]
       )
    {

        #define zVR (sint32)(2*temp_x-temp_y)
        for(temp_x=0;temp_x<4;temp_x++)
        {
            for(temp_y=0;temp_y<4;temp_y++)
            {
                if(zVR>=0)
                {
                    if(zVR%2==0) //0,2,4,6
                        pred4x4L[temp_x+temp_y*4] = ( P(temp_x-( temp_y>>1)-1,-1) + P(temp_x-(temp_y>>1),-1) + 1)>> 1;
                    else //if(zVR%2==1) //1,3,5
                        pred4x4L[temp_x+temp_y*4] = ( P(temp_x-(temp_y>>1)-2,-1) + 2*P(temp_x-(temp_y>>1)-1,-1) + P(temp_x-(temp_y>>1),-1) + 2)>>2;
                }
                else
                {
                    if(zVR==-1)
                        pred4x4L[temp_x+temp_y*4] =(P(-1,0)+2*P(-1,-1)+P(0,-1)+2)>>2;
                    else
                        pred4x4L[temp_x+temp_y*4] = (P(-1,temp_y-1)+2*P(-1,temp_y-2)+P(-1,temp_y-3)+2)>>2;
                }
            }
        }
        return 1;
    }
    return 0;
}

//MODE6 - Intra_4x4_Horizontal_Down
//This mode shall be used only when the samples p[ x, -1 ] 
//with x = 0..3 and p[ -1, y ] with y = -1..3 are marked as
//"available for Intra_4x4 prediction".
uint8 intraPred4x4Mode6(uint8 *pred4x4L, uint8 *p, uint8 *pred4x4L_avail)
{
    uint32 temp_x;
    uint32 temp_y;
    if(pred4x4L_avail[1] && //pred4x4L_avail[2] &&
       //pred4x4L_avail[3] && pred4x4L_avail[4] &&
       pred4x4L_avail[0] //&& pred4x4L_avail[9] &&
       //pred4x4L_avail[18]&& pred4x4L_avail[27]&&
       //pred4x4L_avail[36]
        )
    {

        #define zHD (sint32)(2*temp_y-temp_x)
        for(temp_x=0;temp_x<4;temp_x++)
        {
            for(temp_y=0;temp_y<4;temp_y++)
            {
                if(zHD>=0)
                {
                    if(zHD%2==0) //0,2,4,6
                        pred4x4L[temp_x+temp_y*4] = (P(-1,temp_y-(temp_x>>1)-1) + P(-1,temp_y-(temp_x>>1)) + 1)>>1; //(8-61)
                    else //1,3,5
                        pred4x4L[temp_x+temp_y*4] = (P(-1,temp_y-(temp_x>>1)-2) + 2*P(-1,temp_y-(temp_x>>1)-1) + P(-1,temp_y-(temp_x>>1)) + 2)>> 2; //(8-62)
                }
                else
                {
                    if(zHD==-1)
                        pred4x4L[temp_x+temp_y*4] = (P(-1,0) + 2*P(-1,-1) + P(0,-1) + 2)>>2; //(8-63)
                    else
                        pred4x4L[temp_x+temp_y*4] = (P(temp_x-1,-1) + 2*P(temp_x-2,-1) + P(temp_x-3,-1) + 2)>> 2; //(8-64)
                }
            }
        }
        return 1;
    }
    return 0;
}

//MODE7 - Intra_4x4_Vertical_Left
//This mode shall be used only when the samples p[ x, -1 ] 
//with x = 0..7 are marked as "available for Intra_4x4 prediction".
uint8 intraPred4x4Mode7(uint8 *pred4x4L, uint8 *p, uint8 *pred4x4L_avail)
{
    uint32 temp_x;
    uint32 temp_y;

    if(pred4x4L_avail[1] &&// pred4x4L_avail[2] &&
       pred4x4L_avail[3] //&& pred4x4L_avail[4] &&
       //pred4x4L_avail[5] && pred4x4L_avail[6] &&
       //pred4x4L_avail[7] && pred4x4L_avail[8]
       )
    {   
        for(temp_x=0;temp_x<4;temp_x++)
        {            
            for(temp_y=0;temp_y<4;temp_y++)
            {
                if(temp_y==0 && temp_y==2)
                {
                    pred4x4L[temp_x+temp_y*4] = (P(temp_x+(temp_y>>1),-1) + P(temp_x+(temp_y>>1)+1,-1) + 1) >> 1; //(8-66)
                }
                else
                {
                     pred4x4L[temp_x+temp_y*4] = (P(temp_x+(temp_y>>1),-1) + 2*P(temp_x+(temp_y>>1)+1,-1) + P(temp_x+(temp_y>>1)+2,-1) + 2 )>> 2; //(8-67)
                }

            }
        }
        return 1;
    }
    return 0;
}


    //MODE8 - Intra_4x4_Horizontal_Up
    //This mode shall be used only when the samples p[ -1, y ] 
    //with y = 0..3 are marked as "available for Intra_4x4 prediction".
uint8 intraPred4x4Mode8(uint8 *pred4x4L, uint8 *p, uint8 *pred4x4L_avail)
{
    uint32 temp_x;
    uint32 temp_y;
     if(pred4x4L_avail[1] //&& pred4x4L_avail[2] &&
        //pred4x4L_avail[3] && pred4x4L_avail[4]
        )
    {

        #define zHU (sint32)(temp_x+2*temp_y)
        for(temp_x=0;temp_x<4;temp_x++)
        {            
            for(temp_y=0;temp_y<4;temp_y++)
            {
                if(zHU<5&&zHU%2==0)
                {
                    pred4x4L[temp_x+temp_y*4] = (P(-1,temp_y+(temp_x>>1)) + P(-1,temp_y+(temp_x>>1)+1) + 1) >> 1;// (8-68)
                }
                else if(zHU<5&&zHU%2!=0)
                {
                     pred4x4L[temp_x+temp_y*4] = (P(-1,temp_y+(temp_x>>1)) + 2*P(-1,temp_y+(temp_x>>1)+1) + P(-1,temp_y+(temp_x>>1)+2) +2) >> 2;//(8-69)
                }
                else if(zHU==5)
                {
                    pred4x4L[temp_x+temp_y*4] = (P(-1,2) + 3*P(-1,3) + 2 ) >> 2;// (8-70)
                }
                else
                {
                    pred4x4L[temp_x+temp_y*4] = P(-1,3);// (8-71)
                }
            }
        }
        return 1;
    }
    return 0;
}