mpeg2enc: Fix the picture header for B frame
[profile/ivi/libva.git] / test / encode / mpeg2enc.c
1 /*
2  * Copyright (c) 2012 Intel Corporation. All Rights Reserved.
3  *
4  * Permission is hereby granted, free of charge, to any person obtaining a
5  * copy of this software and associated documentation files (the
6  * "Software"), to deal in the Software without restriction, including
7  * without limitation the rights to use, copy, modify, merge, publish,
8  * distribute, sub license, and/or sell copies of the Software, and to
9  * permit persons to whom the Software is furnished to do so, subject to
10  * the following conditions:
11  *
12  * The above copyright notice and this permission notice (including the
13  * next paragraph) shall be included in all copies or substantial portions
14  * of the Software.
15  *
16  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
17  * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
18  * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT.
19  * IN NO EVENT SHALL PRECISION INSIGHT AND/OR ITS SUPPLIERS BE LIABLE FOR
20  * ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
21  * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
22  * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
23  */
24 /*
25  * Simple MPEG-2 encoder based on libVA.
26  *
27  */  
28
29 #include "sysdeps.h"
30
31 #include <getopt.h>
32 #include <unistd.h>
33
34 #include <sys/time.h>
35 #include <sys/types.h>
36 #include <fcntl.h>
37 #include <time.h>
38 #include <pthread.h>
39
40 #include <va/va.h>
41 #include <va/va_enc_mpeg2.h>
42
43 #include "va_display.h"
44
45 #define START_CODE_PICUTRE      0x00000100
46 #define START_CODE_SLICE        0x00000101
47 #define START_CODE_USER         0x000001B2
48 #define START_CODE_SEQ          0x000001B3
49 #define START_CODE_EXT          0x000001B5
50 #define START_CODE_GOP          0x000001B8
51
52 #define CHROMA_FORMAT_RESERVED  0
53 #define CHROMA_FORMAT_420       1
54 #define CHROMA_FORMAT_422       2
55 #define CHROMA_FORMAT_444       3
56
57 #define MAX_SLICES              128
58
59 enum {
60     MPEG2_MODE_I = 0,
61     MPEG2_MODE_IP,
62     MPEG2_MODE_IPB,
63 };
64
65 enum {
66     MPEG2_LEVEL_LOW = 0,
67     MPEG2_LEVEL_MAIN,
68     MPEG2_LEVEL_HIGH,
69 };
70
71 #define CHECK_VASTATUS(va_status, func)                                 \
72     if (va_status != VA_STATUS_SUCCESS) {                               \
73         fprintf(stderr, "%s:%s (%d) failed, exit\n", __func__, func, __LINE__); \
74         exit(1);                                                        \
75     }
76
77 static VAProfile mpeg2_va_profiles[] = {
78     VAProfileMPEG2Simple,
79     VAProfileMPEG2Main
80 };
81
82 struct mpeg2enc_context {
83     /* args */
84     int rate_control_mode;
85     int fps;
86     int mode; /* 0:I, 1:I/P, 2:I/P/B */
87     VAProfile profile;
88     int level;
89     int width;
90     int height;
91     int frame_size;
92     int num_pictures;
93     int qp;
94     FILE *ifp;
95     FILE *ofp;
96     unsigned char *frame_data_buffer;
97     int intra_period;
98     int ip_period;
99     int bit_rate; /* in kbps */
100     VAEncPictureType next_type;
101     int next_display_order;
102     int next_bframes;
103     int new_sequence;
104     int new_gop_header;
105     int gop_header_in_display_order;
106
107     /* VA resource */
108     VADisplay va_dpy;
109     VAEncSequenceParameterBufferMPEG2 seq_param;
110     VAEncPictureParameterBufferMPEG2 pic_param;
111     VAEncSliceParameterBufferMPEG2 slice_param[MAX_SLICES];
112     VAContextID context_id;
113     VAConfigID config_id;
114     VABufferID seq_param_buf_id;                /* Sequence level parameter */
115     VABufferID pic_param_buf_id;                /* Picture level parameter */
116     VABufferID slice_param_buf_id[MAX_SLICES];  /* Slice level parameter, multil slices */
117     VABufferID codedbuf_buf_id;                 /* Output buffer, compressed data */
118     VABufferID packed_seq_header_param_buf_id;
119     VABufferID packed_seq_buf_id;
120     VABufferID packed_pic_header_param_buf_id;
121     VABufferID packed_pic_buf_id;
122     int num_slice_groups;
123     int codedbuf_i_size;
124     int codedbuf_pb_size;
125
126     /* thread */
127     pthread_t upload_thread_id;
128     int upload_thread_value;
129     int current_input_surface;
130     int current_upload_surface;
131 };
132
133 /*
134  * mpeg2enc helpers
135  */
136 #define BITSTREAM_ALLOCATE_STEPPING     4096
137
138 struct __bitstream {
139     unsigned int *buffer;
140     int bit_offset;
141     int max_size_in_dword;
142 };
143
144 typedef struct __bitstream bitstream;
145
146 static unsigned int 
147 swap32(unsigned int val)
148 {
149     unsigned char *pval = (unsigned char *)&val;
150
151     return ((pval[0] << 24)     |
152             (pval[1] << 16)     |
153             (pval[2] << 8)      |
154             (pval[3] << 0));
155 }
156
157 static void
158 bitstream_start(bitstream *bs)
159 {
160     bs->max_size_in_dword = BITSTREAM_ALLOCATE_STEPPING;
161     bs->buffer = calloc(bs->max_size_in_dword * sizeof(int), 1);
162     bs->bit_offset = 0;
163 }
164
165 static void
166 bitstream_end(bitstream *bs)
167 {
168     int pos = (bs->bit_offset >> 5);
169     int bit_offset = (bs->bit_offset & 0x1f);
170     int bit_left = 32 - bit_offset;
171
172     if (bit_offset) {
173         bs->buffer[pos] = swap32((bs->buffer[pos] << bit_left));
174     }
175 }
176  
177 static void
178 bitstream_put_ui(bitstream *bs, unsigned int val, int size_in_bits)
179 {
180     int pos = (bs->bit_offset >> 5);
181     int bit_offset = (bs->bit_offset & 0x1f);
182     int bit_left = 32 - bit_offset;
183
184     if (!size_in_bits)
185         return;
186
187     if (size_in_bits < 32)
188         val &= ((1 << size_in_bits) - 1);
189
190     bs->bit_offset += size_in_bits;
191
192     if (bit_left > size_in_bits) {
193         bs->buffer[pos] = (bs->buffer[pos] << size_in_bits | val);
194     } else {
195         size_in_bits -= bit_left;
196         bs->buffer[pos] = (bs->buffer[pos] << bit_left) | (val >> size_in_bits);
197         bs->buffer[pos] = swap32(bs->buffer[pos]);
198
199         if (pos + 1 == bs->max_size_in_dword) {
200             bs->max_size_in_dword += BITSTREAM_ALLOCATE_STEPPING;
201             bs->buffer = realloc(bs->buffer, bs->max_size_in_dword * sizeof(unsigned int));
202         }
203
204         bs->buffer[pos + 1] = val;
205     }
206 }
207
208 static void
209 bitstream_byte_aligning(bitstream *bs, int bit)
210 {
211     int bit_offset = (bs->bit_offset & 0x7);
212     int bit_left = 8 - bit_offset;
213     int new_val;
214
215     if (!bit_offset)
216         return;
217
218     assert(bit == 0 || bit == 1);
219
220     if (bit)
221         new_val = (1 << bit_left) - 1;
222     else
223         new_val = 0;
224
225     bitstream_put_ui(bs, new_val, bit_left);
226 }
227
228 static struct mpeg2_frame_rate {
229     int code;
230     float value;
231 } frame_rate_tab[] = {
232     {1, 23.976},
233     {2, 24.0},
234     {3, 25.0},
235     {4, 29.97},
236     {5, 30},
237     {6, 50},
238     {7, 59.94},
239     {8, 60}
240 };
241
242 static int
243 find_frame_rate_code(const VAEncSequenceParameterBufferMPEG2 *seq_param)
244 {
245     unsigned int delta = -1;
246     int code = 1, i;
247     float frame_rate_value = seq_param->frame_rate * 
248         (seq_param->sequence_extension.bits.frame_rate_extension_d + 1) / 
249         (seq_param->sequence_extension.bits.frame_rate_extension_n + 1);
250
251     for (i = 0; i < sizeof(frame_rate_tab) / sizeof(frame_rate_tab[0]); i++) {
252
253         if (abs(1000 * frame_rate_tab[i].value - 1000 * frame_rate_value) < delta) {
254             code = frame_rate_tab[i].code;
255             delta = abs(1000 * frame_rate_tab[i].value - 1000 * frame_rate_value);
256         }
257     }
258
259     return code;
260 }
261
262 static void 
263 sps_rbsp(struct mpeg2enc_context *ctx,
264          const VAEncSequenceParameterBufferMPEG2 *seq_param,
265          bitstream *bs)
266 {
267     int frame_rate_code = find_frame_rate_code(seq_param);
268
269     if (ctx->new_sequence) {
270         bitstream_put_ui(bs, START_CODE_SEQ, 32);
271         bitstream_put_ui(bs, seq_param->picture_width, 12);
272         bitstream_put_ui(bs, seq_param->picture_height, 12);
273         bitstream_put_ui(bs, seq_param->aspect_ratio_information, 4);
274         bitstream_put_ui(bs, frame_rate_code, 4); /* frame_rate_code */
275         bitstream_put_ui(bs, (seq_param->bits_per_second + 399) / 400, 18); /* the low 18 bits of bit_rate */
276         bitstream_put_ui(bs, 1, 1); /* marker_bit */
277         bitstream_put_ui(bs, seq_param->vbv_buffer_size, 10);
278         bitstream_put_ui(bs, 0, 1); /* constraint_parameter_flag, always 0 for MPEG-2 */
279         bitstream_put_ui(bs, 0, 1); /* load_intra_quantiser_matrix */
280         bitstream_put_ui(bs, 0, 1); /* load_non_intra_quantiser_matrix */
281
282         bitstream_byte_aligning(bs, 0);
283
284         bitstream_put_ui(bs, START_CODE_EXT, 32);
285         bitstream_put_ui(bs, 1, 4); /* sequence_extension id */
286         bitstream_put_ui(bs, seq_param->sequence_extension.bits.profile_and_level_indication, 8);
287         bitstream_put_ui(bs, seq_param->sequence_extension.bits.progressive_sequence, 1);
288         bitstream_put_ui(bs, seq_param->sequence_extension.bits.chroma_format, 2);
289         bitstream_put_ui(bs, seq_param->picture_width >> 12, 2);
290         bitstream_put_ui(bs, seq_param->picture_height >> 12, 2);
291         bitstream_put_ui(bs, ((seq_param->bits_per_second + 399) / 400) >> 18, 12); /* bit_rate_extension */
292         bitstream_put_ui(bs, 1, 1); /* marker_bit */
293         bitstream_put_ui(bs, seq_param->vbv_buffer_size >> 10, 8);
294         bitstream_put_ui(bs, seq_param->sequence_extension.bits.low_delay, 1);
295         bitstream_put_ui(bs, seq_param->sequence_extension.bits.frame_rate_extension_n, 2);
296         bitstream_put_ui(bs, seq_param->sequence_extension.bits.frame_rate_extension_d, 5);
297
298         bitstream_byte_aligning(bs, 0);
299     }
300
301     if (ctx->new_gop_header) {
302         bitstream_put_ui(bs, START_CODE_GOP, 32);
303         bitstream_put_ui(bs, seq_param->gop_header.bits.time_code, 25);
304         bitstream_put_ui(bs, seq_param->gop_header.bits.closed_gop, 1);
305         bitstream_put_ui(bs, seq_param->gop_header.bits.broken_link, 1);
306
307         bitstream_byte_aligning(bs, 0);
308     }
309 }
310
311 static void 
312 pps_rbsp(const VAEncSequenceParameterBufferMPEG2 *seq_param,
313          const VAEncPictureParameterBufferMPEG2 *pic_param,
314          bitstream *bs)
315 {
316     int i;
317     int chroma_420_type;
318
319     if (seq_param->sequence_extension.bits.chroma_format == CHROMA_FORMAT_420)
320         chroma_420_type = pic_param->picture_coding_extension.bits.progressive_frame;
321     else
322         chroma_420_type = 0;
323
324     bitstream_put_ui(bs, START_CODE_PICUTRE, 32);
325     bitstream_put_ui(bs, pic_param->temporal_reference, 10);
326     bitstream_put_ui(bs,
327                      pic_param->picture_type == VAEncPictureTypeIntra ? 1 :
328                      pic_param->picture_type == VAEncPictureTypePredictive ? 2 : 3,
329                      3);
330     bitstream_put_ui(bs, 0xFFFF, 16); /* vbv_delay, always 0xFFFF */
331     
332     if (pic_param->picture_type == VAEncPictureTypePredictive ||
333         pic_param->picture_type == VAEncPictureTypeBidirectional) {
334         bitstream_put_ui(bs, 0, 1); /* full_pel_forward_vector, always 0 for MPEG-2 */
335         bitstream_put_ui(bs, 7, 3); /* forward_f_code, always 7 for MPEG-2 */
336     }
337
338     if (pic_param->picture_type == VAEncPictureTypeBidirectional) {
339         bitstream_put_ui(bs, 0, 1); /* full_pel_backward_vector, always 0 for MPEG-2 */
340         bitstream_put_ui(bs, 7, 3); /* backward_f_code, always 7 for MPEG-2 */
341     }
342      
343     bitstream_put_ui(bs, 0, 1); /* extra_bit_picture, 0 */
344
345     bitstream_byte_aligning(bs, 0);
346
347     bitstream_put_ui(bs, START_CODE_EXT, 32);
348     bitstream_put_ui(bs, 8, 4); /* Picture Coding Extension ID: 8 */
349     bitstream_put_ui(bs, pic_param->f_code[0][0], 4);
350     bitstream_put_ui(bs, pic_param->f_code[0][1], 4);
351     bitstream_put_ui(bs, pic_param->f_code[1][0], 4);
352     bitstream_put_ui(bs, pic_param->f_code[1][1], 4);
353
354     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.intra_dc_precision, 2);
355     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.picture_structure, 2);
356     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.top_field_first, 1);
357     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.frame_pred_frame_dct, 1);
358     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.concealment_motion_vectors, 1);
359     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.q_scale_type, 1);
360     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.intra_vlc_format, 1);
361     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.alternate_scan, 1);
362     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.repeat_first_field, 1);
363     bitstream_put_ui(bs, chroma_420_type, 1);
364     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.progressive_frame, 1);
365     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.composite_display_flag, 1);
366
367     bitstream_byte_aligning(bs, 0);
368
369     if (pic_param->user_data_length) {
370         bitstream_put_ui(bs, START_CODE_USER, 32);
371
372         for (i = 0; i < pic_param->user_data_length; i++) {
373             bitstream_put_ui(bs, pic_param->user_data[i], 8);
374         }
375
376         bitstream_byte_aligning(bs, 0);
377     }
378 }
379
380 static int
381 build_packed_pic_buffer(const VAEncSequenceParameterBufferMPEG2 *seq_param,
382                         const VAEncPictureParameterBufferMPEG2 *pic_param,
383                         unsigned char **header_buffer)
384 {
385     bitstream bs;
386
387     bitstream_start(&bs);
388     pps_rbsp(seq_param, pic_param, &bs);
389     bitstream_end(&bs);
390
391     *header_buffer = (unsigned char *)bs.buffer;
392     return bs.bit_offset;
393 }
394
395 static int
396 build_packed_seq_buffer(struct mpeg2enc_context *ctx,
397                         const VAEncSequenceParameterBufferMPEG2 *seq_param,
398                         unsigned char **header_buffer)
399 {
400     bitstream bs;
401
402     bitstream_start(&bs);
403     sps_rbsp(ctx, seq_param, &bs);
404     bitstream_end(&bs);
405
406     *header_buffer = (unsigned char *)bs.buffer;
407     return bs.bit_offset;
408 }
409
410 /*
411  * mpeg2enc
412  */
413 #define SID_INPUT_PICTURE_0                     0
414 #define SID_INPUT_PICTURE_1                     1
415 #define SID_REFERENCE_PICTURE_L0                2
416 #define SID_REFERENCE_PICTURE_L1                3
417 #define SID_RECON_PICTURE                       4
418 #define SID_NUMBER                              SID_RECON_PICTURE + 1
419
420 static VASurfaceID surface_ids[SID_NUMBER];
421
422 /*
423  * upload thread function
424  */
425 static void *
426 upload_yuv_to_surface(void *data)
427 {
428     struct mpeg2enc_context *ctx = data;
429     VAImage surface_image;
430     VAStatus va_status;
431     void *surface_p = NULL;
432     unsigned char *y_src, *u_src, *v_src;
433     unsigned char *y_dst, *u_dst, *v_dst;
434     int y_size = ctx->width * ctx->height;
435     int u_size = (ctx->width >> 1) * (ctx->height >> 1);
436     int row, col;
437     size_t n_items;
438
439     do {
440         n_items = fread(ctx->frame_data_buffer, ctx->frame_size, 1, ctx->ifp);
441     } while (n_items != 1);
442
443     va_status = vaDeriveImage(ctx->va_dpy, surface_ids[ctx->current_upload_surface], &surface_image);
444     CHECK_VASTATUS(va_status,"vaDeriveImage");
445
446     vaMapBuffer(ctx->va_dpy, surface_image.buf, &surface_p);
447     assert(VA_STATUS_SUCCESS == va_status);
448         
449     y_src = ctx->frame_data_buffer;
450     u_src = ctx->frame_data_buffer + y_size; /* UV offset for NV12 */
451     v_src = ctx->frame_data_buffer + y_size + u_size;
452
453     y_dst = surface_p + surface_image.offsets[0];
454     u_dst = surface_p + surface_image.offsets[1]; /* UV offset for NV12 */
455     v_dst = surface_p + surface_image.offsets[2];
456
457     /* Y plane */
458     for (row = 0; row < surface_image.height; row++) {
459         memcpy(y_dst, y_src, surface_image.width);
460         y_dst += surface_image.pitches[0];
461         y_src += ctx->width;
462     }
463
464     if (surface_image.format.fourcc == VA_FOURCC_NV12) { /* UV plane */
465         for (row = 0; row < surface_image.height / 2; row++) {
466             for (col = 0; col < surface_image.width / 2; col++) {
467                 u_dst[col * 2] = u_src[col];
468                 u_dst[col * 2 + 1] = v_src[col];
469             }
470
471             u_dst += surface_image.pitches[1];
472             u_src += (ctx->width / 2);
473             v_src += (ctx->width / 2);
474         }
475     } else {
476         for (row = 0; row < surface_image.height / 2; row++) {
477             for (col = 0; col < surface_image.width / 2; col++) {
478                 u_dst[col] = u_src[col];
479                 v_dst[col] = v_src[col];
480             }
481
482             u_dst += surface_image.pitches[1];
483             v_dst += surface_image.pitches[2];
484             u_src += (ctx->width / 2);
485             v_src += (ctx->width / 2);
486         }
487     }
488
489     vaUnmapBuffer(ctx->va_dpy, surface_image.buf);
490     vaDestroyImage(ctx->va_dpy, surface_image.image_id);
491
492     return NULL;
493 }
494
495 static void 
496 mpeg2enc_exit(struct mpeg2enc_context *ctx, int exit_code)
497 {
498     if (ctx->frame_data_buffer) {
499         free(ctx->frame_data_buffer);
500         ctx->frame_data_buffer = NULL;
501     }
502
503     if (ctx->ifp) {
504         fclose(ctx->ifp);
505         ctx->ifp = NULL;
506     }
507
508     if (ctx->ofp) {
509         fclose(ctx->ofp);
510         ctx->ofp = NULL;
511     }
512
513     exit(exit_code);
514 }
515
516 static void 
517 usage(char *program)
518 {   
519     fprintf(stderr, "Usage: %s --help\n", program);
520     fprintf(stderr, "\t--help   print this message\n");
521     fprintf(stderr, "Usage: %s <width> <height> <ifile> <ofile> [options]\n", program);
522     fprintf(stderr, "\t<width>  specifies the frame width\n");
523     fprintf(stderr, "\t<height> specifies the frame height\n");
524     fprintf(stderr, "\t<ifile>  specifies the I420/IYUV YUV file\n");
525     fprintf(stderr, "\t<ofile>  specifies the encoded MPEG-2 file\n");
526     fprintf(stderr, "where options include:\n");
527     fprintf(stderr, "\t--cqp <QP>       const qp mode with specified <QP>\n");
528     fprintf(stderr, "\t--fps <FPS>      specify the frame rate\n");
529     fprintf(stderr, "\t--mode <MODE>    specify the mode 0 (I), 1 (I/P) and 2 (I/P/B)\n");
530     fprintf(stderr, "\t--profile <PROFILE>      specify the profile 0(Simple), or 1(Main, default)\n");
531     fprintf(stderr, "\t--level <LEVEL>  specify the level 0(Low), 1(Main, default) or 2(High)\n");    
532 }
533
534 static void 
535 parse_args(struct mpeg2enc_context *ctx, int argc, char **argv)
536 {
537     int c, tmp;
538     int option_index = 0;
539     long file_size;
540     static struct option long_options[] = {
541         {"help",        no_argument,            0,      'h'},
542         {"cqp",         required_argument,      0,      'c'},
543         {"fps",         required_argument,      0,      'f'},
544         {"mode",        required_argument,      0,      'm'},
545         {"profile",     required_argument,      0,      'p'},
546         {"level",       required_argument,      0,      'l'},
547         { NULL,         0,                      NULL,   0 }
548     };
549
550     if ((argc == 2 && strcmp(argv[1], "--help") == 0) ||
551         (argc < 5))
552         goto print_usage;
553
554     ctx->width = atoi(argv[1]);
555     ctx->height = atoi(argv[2]);
556
557     if (ctx->width <= 0 || ctx->height <= 0) {
558         fprintf(stderr, "<width> and <height> must be greater than 0\n");
559         goto err_exit;
560     }
561
562     ctx->ifp = fopen(argv[3], "rb");
563
564     if (ctx->ifp == NULL) {
565         fprintf(stderr, "Can't open the input file\n");
566         goto err_exit;
567     }
568
569     fseek(ctx->ifp, 0l, SEEK_END);
570     file_size = ftell(ctx->ifp);
571     ctx->frame_size = ctx->width * ctx->height * 3 / 2;
572
573     if ((file_size < ctx->frame_size) ||
574         (file_size % ctx->frame_size)) {
575         fprintf(stderr, "The input file size %ld isn't a multiple of the frame size %d\n", file_size, ctx->frame_size);
576         goto err_exit;
577     }
578
579     ctx->num_pictures = file_size / ctx->frame_size;
580     fseek(ctx->ifp, 0l, SEEK_SET);
581     
582     ctx->ofp = fopen(argv[4], "wb");
583     
584     if (ctx->ofp == NULL) {
585         fprintf(stderr, "Can't create the output file\n");
586         goto err_exit;
587     }
588
589     opterr = 0;
590     ctx->fps = 30;
591     ctx->qp = 8;
592     ctx->rate_control_mode = VA_RC_CQP;
593     ctx->mode = MPEG2_MODE_IP;
594     ctx->profile = VAProfileMPEG2Main;
595     ctx->level = MPEG2_LEVEL_MAIN;
596
597     optind = 5;
598
599     while((c = getopt_long(argc, argv,
600                            "",
601                            long_options, 
602                            &option_index)) != -1) {
603         switch(c) {
604         case 'c':
605             tmp = atoi(optarg);
606
607             /* only support q_scale_type = 0 */
608             if (tmp > 62 || tmp < 2) {
609                 fprintf(stderr, "Warning: QP must be in [2, 62]\n");
610
611                 if (tmp > 62)
612                     tmp = 62;
613
614                 if (tmp < 2)
615                     tmp = 2;
616             }
617
618             ctx->qp = tmp & 0xFE;
619             ctx->rate_control_mode = VA_RC_CQP;
620
621             break;
622
623         case 'f':
624             tmp = atoi(optarg);
625
626             if (tmp <= 0)
627                 fprintf(stderr, "Warning: FPS must be greater than 0\n");
628             else
629                 ctx->fps = tmp;
630
631             ctx->rate_control_mode = VA_RC_CBR;
632
633             break;
634
635         case 'm':
636             tmp = atoi(optarg);
637             
638             if (tmp < MPEG2_MODE_I || tmp > MPEG2_MODE_IPB)
639                 fprintf(stderr, "Waning: MODE must be 0, 1, or 2\n");
640             else
641                 ctx->mode = tmp;
642
643             break;
644
645         case 'p':
646             tmp = atoi(optarg);
647             
648             if (tmp < 0 || tmp > 1)
649                 fprintf(stderr, "Waning: PROFILE must be 0 or 1\n");
650             else
651                 ctx->profile = mpeg2_va_profiles[tmp];
652
653             break;
654
655         case 'l':
656             tmp = atoi(optarg);
657             
658             if (tmp < MPEG2_LEVEL_LOW || tmp > MPEG2_LEVEL_HIGH)
659                 fprintf(stderr, "Waning: LEVEL must be 0, 1, or 2\n");
660             else
661                 ctx->level = tmp;
662
663             break;
664
665         case '?':
666             fprintf(stderr, "Error: unkown command options\n");
667
668         case 'h':
669             goto print_usage;
670         }
671     }
672
673     return;
674
675 print_usage:    
676     usage(argv[0]);
677 err_exit:
678     mpeg2enc_exit(ctx, 1);
679 }
680
681 /*
682  * init
683  */
684 static void
685 mpeg2enc_init_sequence_parameter(struct mpeg2enc_context *ctx,
686                                 VAEncSequenceParameterBufferMPEG2 *seq_param)
687 {
688     int profile = 4, level = 8;
689
690     switch (ctx->profile) {
691     case VAProfileMPEG2Simple:
692         profile = 5;
693         break;
694
695     case VAProfileMPEG2Main:
696         profile = 4;
697         break;
698
699     default:
700         assert(0);
701         break;
702     }
703
704     switch (ctx->level) {
705     case MPEG2_LEVEL_LOW:
706         level = 10;
707         break;
708
709     case MPEG2_LEVEL_MAIN:
710         level = 8;
711         break;
712
713     case MPEG2_LEVEL_HIGH:
714         level = 4;
715         break;
716
717     default:
718         assert(0);
719         break;
720     }
721         
722     seq_param->intra_period = ctx->intra_period;
723     seq_param->ip_period = ctx->ip_period;   /* FIXME: ??? */
724     seq_param->picture_width = ctx->width;
725     seq_param->picture_height = ctx->height;
726
727     if (ctx->bit_rate > 0)
728         seq_param->bits_per_second = 1024 * ctx->bit_rate; /* use kbps as input */
729     else
730         seq_param->bits_per_second = 0x3FFFF * 400;
731
732     seq_param->frame_rate = ctx->fps;
733     seq_param->aspect_ratio_information = 1;
734     seq_param->vbv_buffer_size = 3; /* B = 16 * 1024 * vbv_buffer_size */
735
736     seq_param->sequence_extension.bits.profile_and_level_indication = profile << 4 | level;
737     seq_param->sequence_extension.bits.progressive_sequence = 1; /* progressive frame-pictures */
738     seq_param->sequence_extension.bits.chroma_format = CHROMA_FORMAT_420; /* 4:2:0 */
739     seq_param->sequence_extension.bits.low_delay = 0; /* FIXME */
740     seq_param->sequence_extension.bits.frame_rate_extension_n = 0;
741     seq_param->sequence_extension.bits.frame_rate_extension_d = 0;
742
743     seq_param->gop_header.bits.time_code = (1 << 12); /* bit12: marker_bit */
744     seq_param->gop_header.bits.closed_gop = 0;
745     seq_param->gop_header.bits.broken_link = 0;    
746 }
747
748 static void
749 mpeg2enc_init_picture_parameter(struct mpeg2enc_context *ctx,
750                                VAEncPictureParameterBufferMPEG2 *pic_param)
751 {
752     pic_param->forward_reference_picture = VA_INVALID_ID;
753     pic_param->backward_reference_picture = VA_INVALID_ID;
754     pic_param->reconstructed_picture = VA_INVALID_ID;
755     pic_param->coded_buf = VA_INVALID_ID;
756     pic_param->picture_type = VAEncPictureTypeIntra;
757
758     pic_param->temporal_reference = 0;
759     pic_param->f_code[0][0] = 0xf;
760     pic_param->f_code[0][1] = 0xf;
761     pic_param->f_code[1][0] = 0xf;
762     pic_param->f_code[1][1] = 0xf;
763
764     pic_param->picture_coding_extension.bits.intra_dc_precision = 0; /* 8bits */
765     pic_param->picture_coding_extension.bits.picture_structure = 3; /* frame picture */
766     pic_param->picture_coding_extension.bits.top_field_first = 0; 
767     pic_param->picture_coding_extension.bits.frame_pred_frame_dct = 1; /* FIXME */
768     pic_param->picture_coding_extension.bits.concealment_motion_vectors = 0;
769     pic_param->picture_coding_extension.bits.q_scale_type = 0;
770     pic_param->picture_coding_extension.bits.intra_vlc_format = 0;
771     pic_param->picture_coding_extension.bits.alternate_scan = 0;
772     pic_param->picture_coding_extension.bits.repeat_first_field = 0;
773     pic_param->picture_coding_extension.bits.progressive_frame = 1;
774     pic_param->picture_coding_extension.bits.composite_display_flag = 0;
775
776     pic_param->user_data_length = 0;
777 }
778
779 static void 
780 mpeg2enc_alloc_va_resources(struct mpeg2enc_context *ctx)
781 {
782     VAEntrypoint *entrypoint_list;
783     VAConfigAttrib attrib_list[2];
784     VAStatus va_status;
785     int max_entrypoints, num_entrypoints, entrypoint;
786     int major_ver, minor_ver;
787
788     ctx->va_dpy = va_open_display();
789     va_status = vaInitialize(ctx->va_dpy,
790                              &major_ver,
791                              &minor_ver);
792     CHECK_VASTATUS(va_status, "vaInitialize");
793
794     max_entrypoints = vaMaxNumEntrypoints(ctx->va_dpy);
795     entrypoint_list = malloc(max_entrypoints * sizeof(VAEntrypoint));
796     vaQueryConfigEntrypoints(ctx->va_dpy,
797                              ctx->profile,
798                              entrypoint_list,
799                              &num_entrypoints);
800
801     for (entrypoint = 0; entrypoint < num_entrypoints; entrypoint++) {
802         if (entrypoint_list[entrypoint] == VAEntrypointEncSlice)
803             break;
804     }
805
806     free(entrypoint_list);
807
808     if (entrypoint == num_entrypoints) {
809         /* not find Slice entry point */
810         assert(0);
811     }
812
813     /* find out the format for the render target, and rate control mode */
814     attrib_list[0].type = VAConfigAttribRTFormat;
815     attrib_list[1].type = VAConfigAttribRateControl;
816     vaGetConfigAttributes(ctx->va_dpy,
817                           ctx->profile,
818                           VAEntrypointEncSlice,
819                           &attrib_list[0],
820                           2);
821
822     if ((attrib_list[0].value & VA_RT_FORMAT_YUV420) == 0) {
823         /* not find desired YUV420 RT format */
824         assert(0);
825     }
826
827     if ((attrib_list[1].value & ctx->rate_control_mode) == 0) {
828         /* Can't find matched RC mode */
829         fprintf(stderr, "RC mode %d isn't found, exit\n", ctx->rate_control_mode);
830         assert(0);
831     }
832
833     attrib_list[0].value = VA_RT_FORMAT_YUV420; /* set to desired RT format */
834     attrib_list[1].value = ctx->rate_control_mode; /* set to desired RC mode */
835
836     va_status = vaCreateConfig(ctx->va_dpy,
837                                ctx->profile,
838                                VAEntrypointEncSlice,
839                                attrib_list,
840                                2,
841                                &ctx->config_id);
842     CHECK_VASTATUS(va_status, "vaCreateConfig");
843
844     /* Create a context for this decode pipe */
845     va_status = vaCreateContext(ctx->va_dpy,
846                                 ctx->config_id,
847                                 ctx->width,
848                                 ctx->height,
849                                 VA_PROGRESSIVE, 
850                                 0,
851                                 0,
852                                 &ctx->context_id);
853     CHECK_VASTATUS(va_status, "vaCreateContext");
854
855     va_status = vaCreateSurfaces(ctx->va_dpy,
856                                  VA_RT_FORMAT_YUV420,
857                                  ctx->width,
858                                  ctx->height,
859                                  surface_ids,
860                                  SID_NUMBER,
861                                  NULL,
862                                  0);
863     CHECK_VASTATUS(va_status, "vaCreateSurfaces");
864 }
865
866 static void 
867 mpeg2enc_init(struct mpeg2enc_context *ctx)
868 {
869     int i;
870
871     ctx->frame_data_buffer = (unsigned char *)malloc(ctx->frame_size);
872     ctx->seq_param_buf_id = VA_INVALID_ID;
873     ctx->pic_param_buf_id = VA_INVALID_ID;
874     ctx->packed_seq_header_param_buf_id = VA_INVALID_ID;
875     ctx->packed_seq_buf_id = VA_INVALID_ID;
876     ctx->packed_pic_header_param_buf_id = VA_INVALID_ID;
877     ctx->packed_pic_buf_id = VA_INVALID_ID;
878     ctx->codedbuf_buf_id = VA_INVALID_ID;
879     ctx->codedbuf_i_size = ctx->frame_size;
880     ctx->codedbuf_pb_size = 0;
881     ctx->next_display_order = 0;
882     ctx->next_type = VAEncPictureTypeIntra;
883
884     if (ctx->mode == MPEG2_MODE_I) {
885         ctx->intra_period = 1;
886         ctx->ip_period = 0;
887     } else if (ctx->mode == MPEG2_MODE_IP) {
888         ctx->intra_period = 16;
889         ctx->ip_period = 0;
890     } else {
891         ctx->intra_period = 16;
892         ctx->ip_period = 2;
893     }
894
895     ctx->next_bframes = ctx->ip_period;
896
897     ctx->new_sequence = 1;
898     ctx->new_gop_header = 1;
899     ctx->gop_header_in_display_order = 0;
900
901     ctx->bit_rate = -1;
902
903     for (i = 0; i < MAX_SLICES; i++) {
904         ctx->slice_param_buf_id[i] = VA_INVALID_ID;
905     }
906
907     mpeg2enc_init_sequence_parameter(ctx, &ctx->seq_param);
908     mpeg2enc_init_picture_parameter(ctx, &ctx->pic_param);
909     mpeg2enc_alloc_va_resources(ctx);
910
911     /* thread */
912     ctx->current_input_surface = SID_INPUT_PICTURE_0;
913     ctx->current_upload_surface = SID_INPUT_PICTURE_1;
914     ctx->upload_thread_value = pthread_create(&ctx->upload_thread_id,
915                                               NULL,
916                                               upload_yuv_to_surface,
917                                               ctx);
918 }
919
920 static int 
921 mpeg2enc_time_code(VAEncSequenceParameterBufferMPEG2 *seq_param,
922                    int num_frames)
923 {
924     int fps = (int)(seq_param->frame_rate + 0.5);
925     int time_code = 0;
926     int time_code_pictures, time_code_seconds, time_code_minutes, time_code_hours;
927     int drop_frame_flag = 0;
928
929     assert(fps <= 60);
930
931     time_code_seconds = num_frames / fps;
932     time_code_pictures = num_frames % fps;
933     time_code |= time_code_pictures;
934
935     time_code_minutes = time_code_seconds / 60;
936     time_code_seconds = time_code_seconds % 60;
937     time_code |= (time_code_seconds << 6);
938
939     time_code_hours = time_code_minutes / 60;
940     time_code_minutes = time_code_minutes % 60;
941
942     time_code |= (1 << 12);     /* marker_bit */
943     time_code |= (time_code_minutes << 13);
944
945     time_code_hours = time_code_hours % 24;
946     time_code |= (time_code_hours << 19);
947
948     time_code |= (drop_frame_flag << 24);
949
950     return time_code;
951 }
952
953 /*
954  * run
955  */
956 static void
957 mpeg2enc_update_sequence_parameter(struct mpeg2enc_context *ctx,
958                                    VAEncPictureType picture_type,
959                                    int coded_order,
960                                    int display_order)
961 {
962     VAEncSequenceParameterBufferMPEG2 *seq_param = &ctx->seq_param;
963
964     /* update the time_code info for the new GOP */
965     if (ctx->new_gop_header) {
966         seq_param->gop_header.bits.time_code = mpeg2enc_time_code(seq_param, display_order);
967     }
968 }
969
970 static void
971 mpeg2enc_update_picture_parameter(struct mpeg2enc_context *ctx,
972                                   VAEncPictureType picture_type,
973                                   int coded_order,
974                                   int display_order)
975 {
976     VAEncPictureParameterBufferMPEG2 *pic_param = &ctx->pic_param;
977
978     pic_param->picture_type = picture_type;
979     pic_param->temporal_reference = (display_order - ctx->gop_header_in_display_order) & 0x3FF;
980     pic_param->reconstructed_picture = surface_ids[SID_RECON_PICTURE];
981     pic_param->forward_reference_picture = surface_ids[SID_REFERENCE_PICTURE_L0];
982     pic_param->backward_reference_picture = surface_ids[SID_REFERENCE_PICTURE_L1];
983
984     if (pic_param->picture_type == VAEncPictureTypeIntra) {
985         pic_param->f_code[0][0] = 0xf;
986         pic_param->f_code[0][1] = 0xf;
987         pic_param->f_code[1][0] = 0xf;
988         pic_param->f_code[1][1] = 0xf;
989         pic_param->forward_reference_picture = VA_INVALID_SURFACE;
990         pic_param->backward_reference_picture = VA_INVALID_SURFACE;
991
992     } else if (pic_param->picture_type == VAEncPictureTypePredictive) {
993         pic_param->f_code[0][0] = 0x1;
994         pic_param->f_code[0][1] = 0x1;
995         pic_param->f_code[1][0] = 0xf;
996         pic_param->f_code[1][1] = 0xf;
997         pic_param->forward_reference_picture = surface_ids[SID_REFERENCE_PICTURE_L0];
998         pic_param->backward_reference_picture = VA_INVALID_SURFACE;
999     } else if (pic_param->picture_type == VAEncPictureTypeBidirectional) {
1000         pic_param->f_code[0][0] = 0x1;
1001         pic_param->f_code[0][1] = 0x1;
1002         pic_param->f_code[1][0] = 0x1;
1003         pic_param->f_code[1][1] = 0x1;
1004         pic_param->forward_reference_picture = surface_ids[SID_REFERENCE_PICTURE_L0];
1005         pic_param->backward_reference_picture = surface_ids[SID_REFERENCE_PICTURE_L1];
1006     } else {
1007         assert(0);
1008     }
1009 }
1010
1011 static void
1012 mpeg2enc_update_picture_parameter_buffer(struct mpeg2enc_context *ctx,
1013                                          VAEncPictureType picture_type,
1014                                          int coded_order,
1015                                          int display_order)
1016 {
1017     VAEncPictureParameterBufferMPEG2 *pic_param = &ctx->pic_param;
1018     VAStatus va_status;
1019
1020     /* update the coded buffer id */
1021     pic_param->coded_buf = ctx->codedbuf_buf_id;
1022     va_status = vaCreateBuffer(ctx->va_dpy,
1023                                ctx->context_id,
1024                                VAEncPictureParameterBufferType,
1025                                sizeof(*pic_param),
1026                                1,
1027                                pic_param,
1028                                &ctx->pic_param_buf_id);
1029     CHECK_VASTATUS(va_status, "vaCreateBuffer");
1030 }
1031
1032 static void 
1033 mpeg2enc_update_slice_parameter(struct mpeg2enc_context *ctx, VAEncPictureType picture_type)
1034 {
1035     VAEncSequenceParameterBufferMPEG2 *seq_param;
1036     VAEncPictureParameterBufferMPEG2 *pic_param;
1037     VAEncSliceParameterBufferMPEG2 *slice_param;
1038     VAStatus va_status;
1039     int i, width_in_mbs, height_in_mbs;
1040
1041     pic_param = &ctx->pic_param;
1042     assert(pic_param->picture_coding_extension.bits.q_scale_type == 0);
1043
1044     seq_param = &ctx->seq_param;
1045     width_in_mbs = (seq_param->picture_width + 15) / 16;
1046     height_in_mbs = (seq_param->picture_height + 15) / 16;
1047     ctx->num_slice_groups = 1;
1048
1049     for (i = 0; i < height_in_mbs; i++) {
1050         slice_param = &ctx->slice_param[i];
1051         slice_param->macroblock_address = i * width_in_mbs;
1052         slice_param->num_macroblocks = width_in_mbs;
1053         slice_param->is_intra_slice = (picture_type == VAEncPictureTypeIntra);
1054         slice_param->quantiser_scale_code = ctx->qp / 2;
1055     }
1056
1057     va_status = vaCreateBuffer(ctx->va_dpy,
1058                                ctx->context_id,
1059                                VAEncSliceParameterBufferType,
1060                                sizeof(*slice_param),
1061                                height_in_mbs,
1062                                ctx->slice_param,
1063                                ctx->slice_param_buf_id);
1064     CHECK_VASTATUS(va_status, "vaCreateBuffer");;
1065 }
1066
1067 static int 
1068 begin_picture(struct mpeg2enc_context *ctx,
1069               int coded_order,
1070               int display_order,
1071               VAEncPictureType picture_type)
1072 {
1073     VAStatus va_status;
1074     int tmp;
1075     VAEncPackedHeaderParameterBuffer packed_header_param_buffer;
1076     unsigned int length_in_bits;
1077     unsigned char *packed_seq_buffer = NULL, *packed_pic_buffer = NULL;
1078
1079     if (ctx->upload_thread_value != 0) {
1080         fprintf(stderr, "FATAL error!!!\n");
1081         exit(1);
1082     }
1083     
1084     pthread_join(ctx->upload_thread_id, NULL);
1085
1086     ctx->upload_thread_value = -1;
1087     tmp = ctx->current_input_surface;
1088     ctx->current_input_surface = ctx->current_upload_surface;
1089     ctx->current_upload_surface = tmp;
1090
1091     mpeg2enc_update_sequence_parameter(ctx, picture_type, coded_order, display_order);
1092     mpeg2enc_update_picture_parameter(ctx, picture_type, coded_order, display_order);
1093
1094     if (ctx->new_sequence || ctx->new_gop_header) {
1095         assert(picture_type == VAEncPictureTypeIntra);
1096         length_in_bits = build_packed_seq_buffer(ctx, &ctx->seq_param, &packed_seq_buffer);
1097         packed_header_param_buffer.type = VAEncPackedHeaderMPEG2_SPS;
1098         packed_header_param_buffer.has_emulation_bytes = 0;
1099         packed_header_param_buffer.bit_length = length_in_bits;
1100         va_status = vaCreateBuffer(ctx->va_dpy,
1101                                    ctx->context_id,
1102                                    VAEncPackedHeaderParameterBufferType,
1103                                    sizeof(packed_header_param_buffer), 1, &packed_header_param_buffer,
1104                                    &ctx->packed_seq_header_param_buf_id);
1105         CHECK_VASTATUS(va_status,"vaCreateBuffer");
1106
1107         va_status = vaCreateBuffer(ctx->va_dpy,
1108                                    ctx->context_id,
1109                                    VAEncPackedHeaderDataBufferType,
1110                                    (length_in_bits + 7) / 8, 1, packed_seq_buffer,
1111                                    &ctx->packed_seq_buf_id);
1112         CHECK_VASTATUS(va_status,"vaCreateBuffer");
1113
1114         free(packed_seq_buffer);
1115     }
1116
1117     length_in_bits = build_packed_pic_buffer(&ctx->seq_param, &ctx->pic_param, &packed_pic_buffer);
1118     packed_header_param_buffer.type = VAEncPackedHeaderMPEG2_PPS;
1119     packed_header_param_buffer.has_emulation_bytes = 0;
1120     packed_header_param_buffer.bit_length = length_in_bits;
1121
1122     va_status = vaCreateBuffer(ctx->va_dpy,
1123                                ctx->context_id,
1124                                VAEncPackedHeaderParameterBufferType,
1125                                sizeof(packed_header_param_buffer), 1, &packed_header_param_buffer,
1126                                &ctx->packed_pic_header_param_buf_id);
1127     CHECK_VASTATUS(va_status,"vaCreateBuffer");
1128
1129     va_status = vaCreateBuffer(ctx->va_dpy,
1130                                ctx->context_id,
1131                                VAEncPackedHeaderDataBufferType,
1132                                (length_in_bits + 7) / 8, 1, packed_pic_buffer,
1133                                &ctx->packed_pic_buf_id);
1134     CHECK_VASTATUS(va_status,"vaCreateBuffer");
1135
1136     free(packed_pic_buffer);
1137
1138     /* sequence parameter set */
1139     VAEncSequenceParameterBufferMPEG2 *seq_param = &ctx->seq_param;
1140     va_status = vaCreateBuffer(ctx->va_dpy,
1141                                ctx->context_id,
1142                                VAEncSequenceParameterBufferType,
1143                                sizeof(*seq_param),
1144                                1,
1145                                seq_param,
1146                                &ctx->seq_param_buf_id);
1147     CHECK_VASTATUS(va_status,"vaCreateBuffer");;
1148
1149     /* slice parameter */
1150     mpeg2enc_update_slice_parameter(ctx, picture_type);
1151
1152     return 0;
1153 }
1154
1155 static int 
1156 mpeg2enc_render_picture(struct mpeg2enc_context *ctx)
1157 {
1158     VAStatus va_status;
1159     VABufferID va_buffers[16];
1160     unsigned int num_va_buffers = 0;
1161
1162     va_buffers[num_va_buffers++] = ctx->seq_param_buf_id;
1163     va_buffers[num_va_buffers++] = ctx->pic_param_buf_id;
1164
1165     if (ctx->packed_seq_header_param_buf_id != VA_INVALID_ID)
1166         va_buffers[num_va_buffers++] = ctx->packed_seq_header_param_buf_id;
1167
1168     if (ctx->packed_seq_buf_id != VA_INVALID_ID)
1169         va_buffers[num_va_buffers++] = ctx->packed_seq_buf_id;
1170
1171     if (ctx->packed_pic_header_param_buf_id != VA_INVALID_ID)
1172         va_buffers[num_va_buffers++] = ctx->packed_pic_header_param_buf_id;
1173
1174     if (ctx->packed_pic_buf_id != VA_INVALID_ID)
1175         va_buffers[num_va_buffers++] = ctx->packed_pic_buf_id;
1176
1177     va_status = vaBeginPicture(ctx->va_dpy,
1178                                ctx->context_id,
1179                                surface_ids[ctx->current_input_surface]);
1180     CHECK_VASTATUS(va_status,"vaBeginPicture");
1181
1182     va_status = vaRenderPicture(ctx->va_dpy,
1183                                 ctx->context_id,
1184                                 va_buffers,
1185                                 num_va_buffers);
1186     CHECK_VASTATUS(va_status,"vaRenderPicture");
1187
1188     va_status = vaRenderPicture(ctx->va_dpy,
1189                                 ctx->context_id,
1190                                 &ctx->slice_param_buf_id[0],
1191                                 ctx->num_slice_groups);
1192     CHECK_VASTATUS(va_status,"vaRenderPicture");
1193
1194     va_status = vaEndPicture(ctx->va_dpy, ctx->context_id);
1195     CHECK_VASTATUS(va_status,"vaEndPicture");
1196
1197     return 0;
1198 }
1199
1200 static int 
1201 mpeg2enc_destroy_buffers(struct mpeg2enc_context *ctx, VABufferID *va_buffers, unsigned int num_va_buffers)
1202 {
1203     VAStatus va_status;
1204     unsigned int i;
1205
1206     for (i = 0; i < num_va_buffers; i++) {
1207         if (va_buffers[i] != VA_INVALID_ID) {
1208             va_status = vaDestroyBuffer(ctx->va_dpy, va_buffers[i]);
1209             CHECK_VASTATUS(va_status,"vaDestroyBuffer");
1210             va_buffers[i] = VA_INVALID_ID;
1211         }
1212     }
1213
1214     return 0;
1215 }
1216
1217 static void
1218 end_picture(struct mpeg2enc_context *ctx, VAEncPictureType picture_type, int next_is_bpic)
1219 {
1220     VABufferID tempID;
1221
1222     /* Prepare for next picture */
1223     tempID = surface_ids[SID_RECON_PICTURE];  
1224
1225     if (picture_type != VAEncPictureTypeBidirectional) {
1226         if (next_is_bpic) {
1227             surface_ids[SID_RECON_PICTURE] = surface_ids[SID_REFERENCE_PICTURE_L1]; 
1228             surface_ids[SID_REFERENCE_PICTURE_L1] = tempID;     
1229         } else {
1230             surface_ids[SID_RECON_PICTURE] = surface_ids[SID_REFERENCE_PICTURE_L0]; 
1231             surface_ids[SID_REFERENCE_PICTURE_L0] = tempID;
1232         }
1233     } else {
1234         if (!next_is_bpic) {
1235             surface_ids[SID_RECON_PICTURE] = surface_ids[SID_REFERENCE_PICTURE_L0]; 
1236             surface_ids[SID_REFERENCE_PICTURE_L0] = surface_ids[SID_REFERENCE_PICTURE_L1];
1237             surface_ids[SID_REFERENCE_PICTURE_L1] = tempID;
1238         }
1239     }
1240
1241     mpeg2enc_destroy_buffers(ctx, &ctx->seq_param_buf_id, 1);
1242     mpeg2enc_destroy_buffers(ctx, &ctx->pic_param_buf_id, 1);
1243     mpeg2enc_destroy_buffers(ctx, &ctx->packed_seq_header_param_buf_id, 1);
1244     mpeg2enc_destroy_buffers(ctx, &ctx->packed_seq_buf_id, 1);
1245     mpeg2enc_destroy_buffers(ctx, &ctx->packed_pic_header_param_buf_id, 1);
1246     mpeg2enc_destroy_buffers(ctx, &ctx->packed_pic_buf_id, 1);
1247     mpeg2enc_destroy_buffers(ctx, &ctx->slice_param_buf_id[0], ctx->num_slice_groups);
1248     mpeg2enc_destroy_buffers(ctx, &ctx->codedbuf_buf_id, 1);
1249     memset(ctx->slice_param, 0, sizeof(ctx->slice_param));
1250     ctx->num_slice_groups = 0;
1251 }
1252
1253 static int
1254 store_coded_buffer(struct mpeg2enc_context *ctx, VAEncPictureType picture_type)
1255 {
1256     VACodedBufferSegment *coded_buffer_segment;
1257     unsigned char *coded_mem;
1258     int slice_data_length;
1259     VAStatus va_status;
1260     VASurfaceStatus surface_status;
1261     size_t w_items;
1262
1263     va_status = vaSyncSurface(ctx->va_dpy, surface_ids[ctx->current_input_surface]);
1264     CHECK_VASTATUS(va_status,"vaSyncSurface");
1265
1266     surface_status = 0;
1267     va_status = vaQuerySurfaceStatus(ctx->va_dpy, surface_ids[ctx->current_input_surface], &surface_status);
1268     CHECK_VASTATUS(va_status,"vaQuerySurfaceStatus");
1269
1270     va_status = vaMapBuffer(ctx->va_dpy, ctx->codedbuf_buf_id, (void **)(&coded_buffer_segment));
1271     CHECK_VASTATUS(va_status,"vaMapBuffer");
1272     coded_mem = coded_buffer_segment->buf;
1273
1274     if (coded_buffer_segment->status & VA_CODED_BUF_STATUS_SLICE_OVERFLOW_MASK) {
1275         if (picture_type == VAEncPictureTypeIntra)
1276             ctx->codedbuf_i_size *= 2;
1277         else
1278             ctx->codedbuf_pb_size *= 2;
1279
1280         vaUnmapBuffer(ctx->va_dpy, ctx->codedbuf_buf_id);
1281         return -1;
1282     }
1283
1284     slice_data_length = coded_buffer_segment->size;
1285
1286     do {
1287         w_items = fwrite(coded_mem, slice_data_length, 1, ctx->ofp);
1288     } while (w_items != 1);
1289
1290     if (picture_type == VAEncPictureTypeIntra) {
1291         if (ctx->codedbuf_i_size > slice_data_length * 3 / 2) {
1292             ctx->codedbuf_i_size = slice_data_length * 3 / 2;
1293         }
1294         
1295         if (ctx->codedbuf_pb_size < slice_data_length) {
1296             ctx->codedbuf_pb_size = slice_data_length;
1297         }
1298     } else {
1299         if (ctx->codedbuf_pb_size > slice_data_length * 3 / 2) {
1300             ctx->codedbuf_pb_size = slice_data_length * 3 / 2;
1301         }
1302     }
1303
1304     vaUnmapBuffer(ctx->va_dpy, ctx->codedbuf_buf_id);
1305
1306     return 0;
1307 }
1308
1309 static void
1310 encode_picture(struct mpeg2enc_context *ctx,
1311                int coded_order,
1312                int display_order,
1313                VAEncPictureType picture_type,
1314                int next_is_bpic,
1315                int next_display_order)
1316 {
1317     VAStatus va_status;
1318     int ret = 0, codedbuf_size;
1319     
1320     begin_picture(ctx, coded_order, display_order, picture_type);
1321
1322     if (1) {
1323         /* upload YUV data to VA surface for next frame */
1324         if (next_display_order >= ctx->num_pictures)
1325             next_display_order = ctx->num_pictures - 1;
1326
1327         fseek(ctx->ifp, ctx->frame_size * next_display_order, SEEK_SET);
1328         ctx->upload_thread_value = pthread_create(&ctx->upload_thread_id,
1329                                                   NULL,
1330                                                   upload_yuv_to_surface,
1331                                                   ctx);
1332     }
1333
1334     do {
1335         mpeg2enc_destroy_buffers(ctx, &ctx->codedbuf_buf_id, 1);
1336         mpeg2enc_destroy_buffers(ctx, &ctx->pic_param_buf_id, 1);
1337
1338
1339         if (VAEncPictureTypeIntra == picture_type) {
1340             codedbuf_size = ctx->codedbuf_i_size;
1341         } else {
1342             codedbuf_size = ctx->codedbuf_pb_size;
1343         }
1344
1345         /* coded buffer */
1346         va_status = vaCreateBuffer(ctx->va_dpy,
1347                                    ctx->context_id,
1348                                    VAEncCodedBufferType,
1349                                    codedbuf_size, 1, NULL,
1350                                    &ctx->codedbuf_buf_id);
1351         CHECK_VASTATUS(va_status,"vaCreateBuffer");
1352
1353         /* picture parameter set */
1354         mpeg2enc_update_picture_parameter_buffer(ctx, picture_type, coded_order, display_order);
1355
1356         mpeg2enc_render_picture(ctx);
1357
1358         ret = store_coded_buffer(ctx, picture_type);
1359     } while (ret);
1360
1361     end_picture(ctx, picture_type, next_is_bpic);
1362 }
1363
1364 static void
1365 update_next_frame_info(struct mpeg2enc_context *ctx,
1366                        VAEncPictureType curr_type,
1367                        int curr_coded_order,
1368                        int curr_display_order)
1369 {
1370     if (((curr_coded_order + 1) % ctx->intra_period) == 0) {
1371         ctx->next_type = VAEncPictureTypeIntra;
1372         ctx->next_display_order = curr_coded_order + 1;
1373         
1374         return;
1375     }
1376
1377     if (curr_type == VAEncPictureTypeIntra) {
1378         assert(curr_display_order == curr_coded_order);
1379         ctx->next_type = VAEncPictureTypePredictive;
1380         ctx->next_bframes = ctx->ip_period;
1381         ctx->next_display_order = curr_display_order + ctx->next_bframes + 1;
1382     } else if (curr_type == VAEncPictureTypePredictive) {
1383         if (ctx->ip_period == 0) {
1384             assert(curr_display_order == curr_coded_order);
1385             ctx->next_type = VAEncPictureTypePredictive;
1386             ctx->next_display_order = curr_display_order + 1;
1387         } else {
1388             ctx->next_type = VAEncPictureTypeBidirectional;
1389             ctx->next_display_order = curr_display_order - ctx->next_bframes;
1390             ctx->next_bframes--;
1391         }
1392     } else if (curr_type == VAEncPictureTypeBidirectional) {
1393         if (ctx->next_bframes == 0) {
1394             ctx->next_type = VAEncPictureTypePredictive;
1395             ctx->next_bframes = ctx->ip_period;
1396             ctx->next_display_order = curr_display_order + ctx->next_bframes + 2;
1397         } else {
1398             ctx->next_type = VAEncPictureTypeBidirectional;
1399             ctx->next_display_order = curr_display_order + 1;
1400             ctx->next_bframes--;
1401         }
1402     }
1403
1404     if (ctx->next_display_order >= ctx->num_pictures) {
1405         int rtmp = ctx->next_display_order - (ctx->num_pictures - 1);
1406         ctx->next_display_order = ctx->num_pictures - 1;
1407         ctx->next_bframes -= rtmp;
1408     }
1409 }
1410
1411 static void
1412 mpeg2enc_run(struct mpeg2enc_context *ctx)
1413 {
1414     int display_order = 0, coded_order = 0;
1415     VAEncPictureType type;
1416
1417     ctx->new_sequence = 1;
1418     ctx->new_gop_header = 1;
1419     ctx->gop_header_in_display_order = display_order;
1420
1421     while (coded_order < ctx->num_pictures) {
1422         type = ctx->next_type;
1423         display_order = ctx->next_display_order;
1424         /* follow the IPBxxBPBxxB mode */
1425         update_next_frame_info(ctx, type, coded_order, display_order);
1426         encode_picture(ctx,
1427                        coded_order,
1428                        display_order,
1429                        type,
1430                        ctx->next_type == VAEncPictureTypeBidirectional,
1431                        ctx->next_display_order);
1432
1433         /* update gop_header */
1434         ctx->new_sequence = 0;
1435         ctx->new_gop_header = ctx->next_type == VAEncPictureTypeIntra;
1436
1437         if (ctx->new_gop_header)
1438             ctx->gop_header_in_display_order += ctx->intra_period;
1439
1440         coded_order++;
1441
1442         fprintf(stderr, "\r %d/%d ...", coded_order, ctx->num_pictures);
1443         fflush(stdout);
1444     }
1445 }
1446
1447 /*
1448  * end
1449  */
1450 static void
1451 mpeg2enc_release_va_resources(struct mpeg2enc_context *ctx)
1452 {
1453     vaDestroySurfaces(ctx->va_dpy, surface_ids, SID_NUMBER);    
1454     vaDestroyContext(ctx->va_dpy, ctx->context_id);
1455     vaDestroyConfig(ctx->va_dpy, ctx->config_id);
1456     vaTerminate(ctx->va_dpy);
1457     va_close_display(ctx->va_dpy);
1458 }
1459
1460 static void
1461 mpeg2enc_end(struct mpeg2enc_context *ctx)
1462 {
1463     pthread_join(ctx->upload_thread_id, NULL);
1464     mpeg2enc_release_va_resources(ctx);
1465 }
1466
1467 int 
1468 main(int argc, char *argv[])
1469 {
1470     struct mpeg2enc_context ctx;
1471     struct timeval tpstart, tpend; 
1472     float timeuse;
1473
1474     gettimeofday(&tpstart, NULL);
1475
1476     memset(&ctx, 0, sizeof(ctx));
1477     parse_args(&ctx, argc, argv);
1478     mpeg2enc_init(&ctx);
1479     mpeg2enc_run(&ctx);
1480     mpeg2enc_end(&ctx);
1481
1482     gettimeofday(&tpend, NULL);
1483     timeuse = 1000000 * (tpend.tv_sec - tpstart.tv_sec) + tpend.tv_usec - tpstart.tv_usec;
1484     timeuse /= 1000000;
1485     fprintf(stderr, "\ndone!\n");
1486     fprintf(stderr, "encode %d frames in %f secondes, FPS is %.1f\n", ctx.num_pictures, timeuse, ctx.num_pictures / timeuse);
1487
1488     mpeg2enc_exit(&ctx, 0);
1489
1490     return 0;
1491 }