/*
 * JPEG decoding engine for DCT-baseline
 *
 *      copyrights 2003 by nikq | nikq::club.
 */


/*
  利用法:
  printfはつぶしたほうがいいかもしれません.
  プログレは駄目です. 対応はほぼ無理です.(構造上)

  JPEG jpeg;
  jpeg->data = データ(メモリマップして直結でいけます)
  jpeg->data_size = ファイルサイズ (in byte)
  jpeg->data_index = 0;

  //デコードの用意
  jpeg_init(&jpeg);

  //ヘッダの処理
  jpeg_header(&jpeg);
  

  //画像のサイズがここで確定します.
  width  = jpeg.width;
  height = jpeg.height;

  //デコード
  // rgb : 24bit(R,G,B) * width * height

  ret = jpeg_decode(&jpeg,rgb);

  // ret: 1 -> エラー
  // ret: 0 -> 成功
  
  */

#include <stdio.h>

//ジグザグテーブル
int zigzag_table[]={
     0, 1, 8, 16,9, 2, 3,10,
    17,24,32,25,18,11, 4, 5,
    12,19,26,33,40,48,41,34,
    27,20,13, 6, 7,14,21,28,
    35,42,49,56,57,50,43,36,
    29,22,15,23,30,37,44,51,
    58,59,52,45,38,31,39,46,
    53,60,61,54,47,55,62,63,
    0
};

// ハフマンテーブル
typedef struct{
    int elem; //要素数
    unsigned short code[256];
    unsigned char  size[256];
    unsigned char  value[256];
}HUFF;

typedef struct{
    // SOF
    int width;
    int height;
    // MCU
    int mcu_width;
    int mcu_height;

    int max_h,max_v;
    int compo_count;
    int compo_id[3];
    int compo_sample[3];
    int compo_h[3];
    int compo_v[3];
    int compo_qt[3];

    // SOS
    int scan_count;
    int scan_id[3];
    int scan_ac[3];
    int scan_dc[3];
    int scan_h[3];  // サンプリング要素数
    int scan_v[3];  // サンプリング要素数
    int scan_qt[3]; // 量子化テーブルインデクス
    
    // DRI
    int interval;

    int mcu_buf[32*32*4]; //バッファ
    int *mcu_yuv[4];
    int mcu_preDC[3];
    
    // DQT
    int dqt[3][64];
    int n_dqt;
    
    // DHT
    HUFF huff[2][3];
    
    
    // i/o
    unsigned char *data;
    int data_index;
    int data_size;
    
    unsigned long bit_buff;
    int bit_remain;
    
}JPEG;


// -------------------------- I/O ----------------------------

// ビッグエンディアンっぽい
unsigned char get_byte(JPEG *jpeg)
{
    unsigned char c;
    c = jpeg->data[ jpeg->data_index ];
    jpeg->data_index++;
    return c;
}
int get_word(JPEG *jpeg)
{
    unsigned char h,l;
    h = get_byte(jpeg);
    l = get_byte(jpeg);
    return (h<<8)|l;
}

unsigned short get_bits(JPEG *jpeg,int bit)
{
    unsigned char  c;
    unsigned short ret;
    unsigned long  buff;
    int remain;
    
    buff   = jpeg->bit_buff;
    remain = jpeg->bit_remain;

    while(remain <= 16){
        c = get_byte(jpeg);
        if(c == 0xFF) { // マーカエラーを防ぐため、FF -> FF 00 にエスケープされてる
            get_byte(jpeg);
        }
        buff = (buff << 8) | c;
        remain += 8;
    }
    ret = (buff>>(remain - bit))&((1<<bit)-1);
    remain -= bit;

    jpeg->bit_remain = remain;
    jpeg->bit_buff   = buff;
    
    return ret;
}

// ------------------------ JPEG セグメント実装 -----------------

// 未対応なのはすっ飛ばす
void jpeg_skip(JPEG *jpeg)
{
    unsigned w;
    w = get_word(jpeg) - 2;
    //printf("skip %d bytes\n",w);
    fseek(jpeg->fp,w,SEEK_CUR);
}

// start of frame
int jpeg_sof(JPEG *jpeg)
{
    unsigned char c,n;
    int i,h,v;
    
    c = get_word(jpeg);

    c = get_byte(jpeg); // bpp
    jpeg->height = get_word(jpeg);
    jpeg->width  = get_word(jpeg);

    n = get_byte(jpeg);    // Num of compo
    jpeg->compo_count = n; // nf
    
    for(i=0;i<n;i++){
        jpeg->compo_id[i] = get_byte(jpeg);

        c = get_byte(jpeg);
        jpeg->compo_sample[i] = c;
        h = (c>>4) & 0x0F;
        v =  c     & 0x0F;

        if(jpeg->max_h < h)
            jpeg->max_h = h;
        if(jpeg->max_v < v)
            jpeg->max_v = v;
        
        jpeg->compo_h[i] = (c>>4)&0x0F;
        jpeg->compo_v[i] = c & 0x0F;

        jpeg->compo_qt[i] = get_byte(jpeg);
    }
    return 0;
}

// data restart interval
void jpeg_dri(JPEG *jpeg)
{
    get_word(jpeg);
    jpeg->interval = get_word(jpeg);
}

// define quantize table
int jpeg_dqt(JPEG *jpeg)
{
    unsigned char c;
    int i,j,v,size;
    
    size = get_word(jpeg) - 2;
    
    //printf("-------DQT----\n");
    
    while(size>0) {
        
        //printf("size: %d\n",size);
        
        c = get_byte(jpeg);
        size--;
        j = c & 7;
        if(j > jpeg->n_dqt)
            jpeg->n_dqt = j;

        //printf("idx : %02x\n",j);

        if((c>>3)){
            // 16 bit DQT
            //printf("DQT 16\n");
            for(i=0;i<64;i++){
                v = get_word(jpeg);
                size-=2;
                jpeg->dqt[j][ i ] = (v>>8);
            }
        }
        else{
            //  8 bit DQT
            for(i=0;i<64;i++){
                v = get_byte(jpeg);
                size--;
                jpeg->dqt[j][ i ] = v;
            }
        }
    }
    return 0;
}

// define huffman table
int jpeg_dht(JPEG *jpeg)
{
    unsigned tc,th;
    unsigned code = 0;
    unsigned char val;
    int i,j,k,num,Li[17];
    int len,max_val;
    HUFF *table;
    
    len = get_word(jpeg) - 2;
    //printf("-------DHT----\n");

    while(len > 0)
    {
        val = get_byte(jpeg);
        
        tc = (val>>4) & 0x0F; // テーブルクラス(DC/AC成分セレクタ)
        th =  val     & 0x0F; // テーブルヘッダ(何枚目のプレーンか)
        //printf(" val:%02x\n",val);
        //printf(" tc :%d(%cC)\n",tc,(tc==0) ? 'D':'A');
        //printf(" th :%d\n",th);

        table = (HUFF*)&(jpeg->huff[tc][th]);

        num = 0;
        for (i = 1; i <= 16; i++) {
            Li[i] = get_byte(jpeg);
            num += Li[i];
        }
        table->elem = num;
        //printf(" num:%d\n",num);

        // 符号生成
        k=0;
        for(i=1;i<=16;i++) {
            for(j=0;j<Li[i];j++) {
                table->size[k++] = i;
            }
        }

        k=0;
        code=0;
        i = table->size[0];
        while(k<num) {
            while(table->size[k] == i){
                table->code[k++] = code++;
            }
            if(k>=num)
                break;
            do{
                code = code << 1;
                i++;
            }while(table->size[k] != i);
        }
        
        for(k=0;k<num;k++)
            table->value[k] = get_byte(jpeg);

        len = len - 18 - num;
    }
    return 0;
}

// start of scan
void jpeg_sos(JPEG *jpeg)
{
    int i;
    unsigned char c;
    get_word(jpeg);

    //printf("-------SOS----\n");
    jpeg->scan_count = get_byte(jpeg);

    //printf("scan:%d\n",jpeg->scan_count);

    for(i=0;i<jpeg->scan_count;i++) {
        c = get_byte(jpeg);
        //printf(" id :%d\n",c);
        jpeg->scan_id[i] = c;

        c = get_byte(jpeg);
        //printf("  dc :%d\n",c>>4);
        //printf("  ac :%d\n",c&0x0F);
        jpeg->scan_dc[i] = c >> 4;   // DC Huffman Table
        jpeg->scan_ac[i] = c & 0x0F; // AC Huffman Table
    }
    //3 bytes skip
    get_byte(jpeg);
    get_byte(jpeg);
    get_byte(jpeg);
}

void jpeg_idct_init(void);
int jpeg_init(JPEG *jpeg)
{
    int i,j;
    //printf("-------init----\n");
    jpeg_idct_init();
    for(i=0;i<3;i++)
        jpeg->mcu_preDC[i]=0;
    jpeg->n_dqt = 0;

    jpeg->max_h = 0;
    jpeg->max_v = 0;
    jpeg->bit_remain = 0;
    jpeg->bit_buff   = 0;
    // DRIリセット無し
    jpeg->interval = 0;
    return 0;
}

int jpeg_header(JPEG *jpeg)
{
    unsigned char c;

    //printf("-------header----\n");

    while(1)
    {
        c = get_byte(jpeg);
        if(c != 0xFF || feof(jpeg->fp))
            return 1;
        c = get_byte(jpeg);
        //printf("%02x ",c);
        switch(c)
        {
        case 0xD8:// printf("SOI\n");
            break;
        case 0xD9:// printf("EOI\n");
            return 1;
            break;
        case 0xC0:
            jpeg_sof(jpeg);
            break;
        case 0xC4:
            jpeg_dht(jpeg);
            break;
        case 0xDB:
            jpeg_dqt(jpeg);
            break;
        case 0xDD:
            jpeg_dri(jpeg);
            break;
        case 0xDA:
            jpeg_sos(jpeg);
            return 0;
        default:
            //printf("%02x:",c);
            jpeg_skip(jpeg);
            break;
        }
    }
}




// MCU decode

// デコード
int jpeg_decode_init(JPEG *jpeg)
{
    int i,j;
    //printf("decode init\n");

    for(i=0;i<jpeg->scan_count;i++) {
        // i:scan
        for(j=0;j<jpeg->compo_count;j++) {
            // j:frame
            if(jpeg->scan_id[i] ==jpeg->compo_id[j]){
                //printf("scan %d is frame %d\n",i,j);
                jpeg->scan_h[i] = jpeg->compo_h[j];
                jpeg->scan_v[i] = jpeg->compo_v[j];
                jpeg->scan_qt[i]= jpeg->compo_qt[j];
                break;
            }
        }
        if(j>=jpeg->compo_count){
            // printf("ヘッダ異常(スキャンとの対応が変)\n");
            // プログレッシブJPEGだとこれが出ます
            return 1;
        }
    }
    // printf("MCU\n");
    // printf("max h:%d\n",jpeg->max_h);
    // printf("max v:%d\n",jpeg->max_v);
    jpeg->mcu_width  = jpeg->max_h * 8;
    jpeg->mcu_height = jpeg->max_v * 8;
    // printf("mcu w:%d\n",jpeg->mcu_width);
    // printf("mcu h:%d\n",jpeg->mcu_height);
    
    for(i=0;i<32*32*4;i++){
        jpeg->mcu_buf[i] = 0x80;
    }
    for(i=0;i<jpeg->scan_count;i++){
        jpeg->mcu_yuv[i] = jpeg->mcu_buf + i*32*32;
    }
    return 0;
}



// ハフマン 1シンボル復号
int jpeg_huff_decode(JPEG *jpeg,int tc,int th)
{
    HUFF *h = &(jpeg->huff[tc][th]);
    int code,size,k,v;
    code = 0;
    size = 0;
    k = 0;
    //printf("tc:%d,th:%d\n",tc,th);
    while( size < 16 ){
        size++;
        v = get_bits(jpeg,1);
        if(v < 0){
            return v;
        }
        code = (code << 1) | v;

        //dump_bits(code);
        //printf(" size:%d\n",size);
        
        while(h->size[k] == size){
            if(h->code[k] == code){
                //printf("value:%02x\n",h->value[k]);
                return h->value[k];
            }
            k++;
        }
    }

/*
    printf("code error:%04lx\n",code);
    dump_bits(jpeg->bit_buff);
    printf("remain:%d\n",jpeg->bit_remain);
 */
    return -1;
}

// 逆DCT（パク
// 浮動少数つかっちゃってるんで、どうにかしたい.

#ifdef USE_CHENDCT

#include "chendct.c"
void jpeg_idct_init(void)
{
    
}

void jpeg_idct(int *block,int *dest)
{
    ChenIDct(block,dest);
}
#else

// ８×８個 の『基底画像』生成
double base_img[64][64]; // 基底画像 ( [横周波数uπ][縦周波数vπ][横位相(M/8)][縦位相(N/8)]
void jpeg_idct_init(void)
{
    int u,v,m,n;
    double tmpm[8],tmpn[8],c;
    double pi = 3.1415926535897932;
    for(u=0; u<8; u++)
    {
        for(v=0; v<8; v++)
        {
            // 基底画像演算 (横×縦を乗算する)
            for(m=0; m<8; m++)
                tmpm[m] = cos( ((2.0*m+1.0)*u*pi)/16.0 ) ; // 横のCos波形
            for(n=0; n<8; n++)
                tmpn[n] = cos( ((2.0*n+1.0)*v*pi)/16.0 ) ; // 縦のCos波形

            // 掛け算して基底画像に
            for(m=0; m<8; m++){
                for(n=0; n<8; n++){
                    c = 1.0/4.0;
                    if(u==0) c *= 1.0 / sqrt(2.0);
                    if(v==0) c *= 1.0 / sqrt(2.0);
                    base_img[u*8 + v][m*8 + n] = tmpm[m] * tmpn[n] * c;
                }
            }
        }
    }
}

void jpeg_idct(int *block,int *dest)
{
    int i,j;
    double temp[64];
    for(i=0;i<64;i++)
        temp[i] = 0.0;
    
    for(i=0;i<64;i++){
        for(j=0;j<64;j++) {
            temp[j] += block[i] * base_img[i][j];
        }
    }
    for(i=0;i<64;i++)
        dest[i] = (int) temp[i];
}

#endif

// 符号化された数値を元に戻す

int jpeg_get_value(JPEG *jpeg,int size) {
    int val = 0;
    if(size == 0)
        val = 0;
    else
    {
        val = get_bits(jpeg,size);
        if (!(val & (1<<(size-1))))
            val = val - (1 << size) + 1;
    }
    return val;
}

// ---- ブロックのデコード ---
// ハフマンデコード＋逆量子化＋逆ジグザグ
int jpeg_decode_huff(JPEG *jpeg,int scan,int *block)
{
    int size, len, val, run, index;
    int *pQt = (int *)(jpeg->dqt[jpeg->scan_qt[scan]]);
    
    // DC
    size = jpeg_huff_decode(jpeg,0,jpeg->scan_dc[scan]);
    if(size < 0)
        return 0;
    val = jpeg_get_value(jpeg,size);
    jpeg->mcu_preDC[scan] += val;
    block[0] = jpeg->mcu_preDC[scan] * pQt[0];

    //AC復号
    index = 1;
    while(index<64)
    {
        size = jpeg_huff_decode(jpeg,1,jpeg->scan_ac[scan]);
        if(size < 0)
            break;
        // EOB
        if(size == 0)
            break;
        
        // RLE
        run  = (size>>4)&0xF;
        size = size & 0x0F;
        
        val = jpeg_get_value(jpeg,size);
        if(val>=0x10000) {
            //マーカ発見
            return val;
        }

        // ZRL
        while (run-- > 0)
            block[ zigzag_table[index++] ] = 0;
        
        block[ zigzag_table[index] ] = val * pQt[index];
        index++;
    }
    while(index<64)
        block[zigzag_table[index++]]=0;
    return 0;
}

// ブロック (補間かけるには、ここで)
// リサンプリング
void jpeg_mcu_bitblt(int *src,int *dest,int width,
                     int x0,int y0,int x1,int y1)
{
    int w,h;
    int x,y,x2,y2;
    int dx,dy;
    w = x1 - x0;
    h = y1 - y0;
    
    for(y=y0;y<y1;y++) {
        y2 = (y-y0) * 8 / h;
        
        for(x=x0;x<x1;x++) {
            x2 = (x-x0) * 8 / w;
            
            dest[(y*width) + x] = src[(y2*8) + x2];
        }
    }
}

// MCU一個変換
int jpeg_decode_mcu(JPEG *jpeg)
{
    int scan,val;
    int unit,i,h,v;
    int *p,hh,vv;
    int block[64],dest[64];

    // mcu_width x mcu_heightサイズのブロックを展開
    for(scan=0;scan<jpeg->scan_count;scan++)
    {
        hh = jpeg->scan_h[scan];
        vv = jpeg->scan_v[scan];
        //printf("scan %d (%dx%d)\n",scan,hh,vv);
        for(v=0;v<vv;v++){
            for(h=0;h<hh;h++){

                // ブロック(8x8)のデコード
                val = jpeg_decode_huff(jpeg , scan , block);
                
                // 逆DCT
                jpeg_idct(block,dest);
                // リサンプリング
                
                // 書き込みバッファ
                p = jpeg->mcu_buf + (scan * 32 * 32);
                
                // 拡大転送
                jpeg_mcu_bitblt(dest , p ,
                                jpeg->mcu_width,
                                jpeg->mcu_width * h / hh,
                                jpeg->mcu_height* v / vv,
                                jpeg->mcu_width * (h+1) / hh,
                                jpeg->mcu_height* (v+1) / vv);
            }
        }
    }
}

// YCrCb=>RGB

int jpeg_decode_yuv(JPEG *jpeg,int h,int v,unsigned char *rgb)
{
    int x0,y0,x,y,x1,y1;
    int *py,*pu,*pv;
    int Y,U,V,k;
    int R,G,B;
    int mw,mh,w;
    
    mw = jpeg->mcu_width;
    mh = jpeg->mcu_height;
    w = jpeg->width;
    
    x0 = h * jpeg->max_h * 8;
    y0 = v * jpeg->max_v * 8;

    x1 = jpeg->width - x0;
    if(x1 > mw)
        x1 = mw;
    y1 = jpeg->height - y0;
    if(y1 > mh)
        y1 = mh;
    
    py = jpeg->mcu_buf;
    pu = jpeg->mcu_buf + 1024;
    pv = jpeg->mcu_buf + 2048;
    

    for(y=0;y<y1;y++){
        for(x=0;x<x1;x++){
            k = y*mw+x;
            
            Y = py[ k ];
            U = pu[ k ];
            V = pv[ k ];
            
            R = 128 + ((Y*0x1000 + V*0x166E) / 4096);
            R = (R & 0xffffff00) ? (R >> 24) ^ 0xff : R;

            G = 128 + ((Y*0x1000 - V*0x0B6C) / 4096);
            G = (G & 0xffffff00) ? (G >> 24) ^ 0xff : G;
            
            B = 128 + ((Y*0x1000 - V*4 + U*0x1C59) / 4096);
            B = (B & 0xffffff00) ? (B >> 24) ^ 0xff : B;
            
            rgb[((y0+y)*w + (x0+x))*3  ] = R;
            rgb[((y0+y)*w + (x0+x))*3+1] = G;
            rgb[((y0+y)*w + (x0+x))*3+2] = B;
        }
    }
}

int jpeg_decode(JPEG *jpeg,unsigned char *rgb)
{
    int h_unit,v_unit;
    int mcu_count,h,v;
    int val;
    unsigned char m;
    
    // MCUサイズ計算
    if(jpeg_decode_init(jpeg))
        return 1; //エラー
    
    h_unit = jpeg->width / jpeg->mcu_width;
    v_unit = jpeg->height/ jpeg->mcu_height;
    if((jpeg->width  % jpeg->mcu_width) > 0){
        h_unit++;
    }
    if((jpeg->height % jpeg->mcu_height) > 0){
        v_unit++;
    }
    
    // printf("(%d x %d)\n",h_unit,v_unit);

    // 1ブロック展開するかもしれない
    mcu_count = 0;
    for(v=0;v<v_unit;v++){
        for(h=0;h<h_unit;h++){
            mcu_count++;
            jpeg_decode_mcu(jpeg);
            jpeg_decode_yuv(jpeg,h,v,rgb);
            //printf("mcu_count:%d\n",mcu_count);
            if(jpeg->interval > 0 && mcu_count >= jpeg->interval){
                
                //printf("interval\n");
                //printf("remain:%d\n",jpeg->bit_remain);

                // RSTmマーカをすっ飛ばす(FF hoge)
                // hogeは読み飛ばしてるので、FFも飛ばす
                jpeg->bit_remain -= (jpeg->bit_remain & 7);
                jpeg->bit_remain -= 8;
                
                //printf("remain:%d\n",jpeg->bit_remain);
                jpeg->mcu_preDC[0] = 0;
                jpeg->mcu_preDC[1] = 0;
                jpeg->mcu_preDC[2] = 0;
                mcu_count = 0;
                
            }
        }
    }

    return 0;
}

/*
int jpeg2beta(unsigned char *data,int size,
              unsigned char *picture,int w,int h)
{
    jpeg_load(&jpeg,fp);
    jpeg_init(&jpeg);
    jpeg_header(&jpeg);
    
}

int main(int argc,char *argv[])
{
    JPEG jpeg;
    FILE *fp;
    unsigned char *rgb;
    
    fp = fopen(argv[1],"rb");
    
    jpeg_load(&jpeg,fp);
    jpeg_init(&jpeg);
    jpeg_header(&jpeg);

    printf("size:%d\n",sizeof(JPEG));
    
    // サイズに応じたメモリを確保
    rgb = (unsigned char *)malloc(jpeg.width * jpeg.height * 3 + 1);
    if(rgb == NULL){
        printf("GET's\n");
        return 1;
    }
    
    jpeg_decode(&jpeg,rgb);
    fclose(jpeg.fp);

    printf("writing back...\n");
    
    bmp_write("test.bmp",jpeg.width,jpeg.height,rgb);
    free(rgb);
    return 0;
}
*/