#------------------------------------------------------------------------
#
# Copyright (c) 1997-1998 by Cornell University.
# 
# See the file "license.txt" for information on usage and redistribution
# of this file, and for a DISCLAIMER OF ALL WARRANTIES.
#
#------------------------------------------------------------------------
package require DvmBasic
package require DvmPnm
package require DvmColor
package require DvmMpeg

source ../lib/pnmlib.tcl

#------------------------------------------------------------------------
# This is a more complicated example (than ppmtompg1.c).  It demonstrates
# how to encode a series of PPM files into I, P, B frames using decoded
# reference frames and half-pel precision motion vector search.  The
# frame pattern is BIBP.BIBP...  The first GOP contains 4 frames.  The 
# other two contain 8 frames each.
#
# Note: the input files should be named <prefix>000.ppm, 
#       <prefix>001.ppm ...
#------------------------------------------------------------------------

# check arguments
if {$argc != 2} {
    puts "enter the prefix of input file names: "
    set namePrefix  [gets stdin]
    puts "enter the output file name: "
    set outFileName [gets stdin]
} else {
    set namePrefix  [lindex $argv 0]
    set outFileName [lindex $argv 1]
}

set num_of_frames       20
set gop_size            8
# buffer large enough to hold 4 frames and some headers
set buffer_size         100000
set frames_per_second   30
set forward_f_code      2
set backward_f_code     2
# width and height of each frame in pixels
set width               176
set height              120

set pictures 0
set seconds 0
set minutes 0
set hours 0
set gop_start 0

#---------------------------------------------------------------
# This procedure encode 3 byte image into a bitstream bs,
# using bitparser bp, and output it to a tcl channel called
# name.  Assumes that the header is already encoded in the 
# bitstream.  (This is an improvement over the routines in pnmlib.tcl
# since it reuse the same header and bitstream)
#---------------------------------------------------------------
proc write_ppm {pnmHdr r g b name} {
    set abp  [bitparser_new]
    set abs  [bitstream_new 300000]
    bitparser_wrap $abp $abs
    set chan [open $name w]
    fconfigure $chan -translation binary -buffersize 65536
    set curr [bitparser_tell $abp]
    pnm_hdr_encode $pnmHdr $abp
    ppm_encode $r $g $b $abp
    bitstream_channel_write $abs $chan 0
    close $chan
    bitparser_free $abp
    bitstream_free $abs
}

#------------------------------------------------------------------------
# Use this to keep track of how many picture frames, seconds, minutes,
# hours.  This information is encoded in the GOP headers.
#------------------------------------------------------------------------
proc increment_time {} {
uplevel #0 {
    incr pictures 1
    if { $pictures == $num_of_frames } {
        set pictures 0
        incr seconds 1
        if { $seconds == 60 } {
            set seconds 0
            incr minutes 1
            if { $minutes == 60 } {
                set minutes 0
                incr hours 1
            }
        }
    }
}}


#------------------------------------------------------------------------
# swap the value of two pointers.
#------------------------------------------------------------------------
proc swap {a b} {
    upvar $a aa
    upvar $b bb
    set temp $aa
    set aa $bb
    set bb $temp
}

#------------------------------------------------------------------------
# Read in a PPM file into r, g, b ByteImages, then convert it to y, u, v
# with 4:2:0 sampling.
#------------------------------------------------------------------------
proc read_to_yuv {} {
uplevel #0 {
    set fileName [format "%s%03d.ppm" $namePrefix [expr $temporalRef+$gop_start]]
    set inFile [open $fileName r]
    puts [format "Processing %s" $fileName]
    fconfigure $inFile -translation binary -buffersize $buffer_size
    bitparser_wrap $bp $bs
    bitstream_channel_read $bs $inFile 0
    pnm_hdr_parse $bp $pnmHdr
    ppm_parse $bp $r $g $b
    close $inFile
    rgb_to_yuv_420 $r $g $b $y $u $v
}}


#------------------------------------------------------------------------
# I Frame
#------------------------------------------------------------------------
proc i_frame_encode {} {
uplevel #0 {
    read_to_yuv
    swap nextY prevY
    swap nextU prevU
    swap nextV prevV
    swap interNext interPrev

    byte_to_sc_i $y $qScale mpeg-intra $scY
    byte_to_sc_i $u $qScale mpeg-intra $scU
    byte_to_sc_i $v $qScale mpeg-intra $scV

    mpeg_pic_hdr_set_temporal_ref       $iPicHdr $temporalRef
    mpeg_pic_hdr_encode $iPicHdr $obp
    mpeg_pic_i_encode $iPicHdr $scY $scU $scV $qScale $sliceInfo $obp

    # decode the image just encoded to use for encoding next P, B Frame
    sc_dequantize $scY $qScale mpeg-intra $scY
    sc_dequantize $scU $qScale mpeg-intra $scU
    sc_dequantize $scV $qScale mpeg-intra $scV
    sc_i_to_byte $scY $nextY
    sc_i_to_byte $scU $nextU
    sc_i_to_byte $scV $nextV
    #swap y nextY
    #swap u nextU
    #swap v nextV
    byte_compute_intermediates $nextY $interNext
    increment_time
}}

#------------------------------------------------------------------------
# P Frame
#------------------------------------------------------------------------
proc p_frame_encode {} {
uplevel #0 {
    read_to_yuv
    swap nextY prevY
    swap nextU prevU
    swap nextV prevV
    swap interNext interPrev

    byte_p_motion_vec_search $pPicHdr $y $prevY $interPrev $fmv
    byte_y_to_sc_p  $y $prevY $fmv $qScale mpeg-intra mpeg-non-intra $scY
    byte_uv_to_sc_p $u $prevU $fmv $qScale mpeg-intra mpeg-non-intra $scU 
    byte_uv_to_sc_p $v $prevV $fmv $qScale mpeg-intra mpeg-non-intra $scV

    mpeg_pic_hdr_set_temporal_ref       $pPicHdr $temporalRef
    mpeg_pic_hdr_encode $pPicHdr $obp
    mpeg_pic_p_encode $pPicHdr $scY $scU $scV $fmv $qScale $sliceInfo $obp

    # decode the image just encoded to use for encoding next B Frame
    sc_non_i_dequantize $scY $qScale mpeg-intra mpeg-non-intra $scY
    sc_non_i_dequantize $scU $qScale mpeg-intra mpeg-non-intra $scU
    sc_non_i_dequantize $scV $qScale mpeg-intra mpeg-non-intra $scV
    sc_p_to_y  $scY $fmv $prevY $nextY
    sc_p_to_uv $scU $fmv $prevU $nextU
    sc_p_to_uv $scV $fmv $prevV $nextV
    byte_compute_intermediates $nextY $interNext
    increment_time
}}

#------------------------------------------------------------------------
# B Frame
#------------------------------------------------------------------------
proc b_frame_encode {} {
uplevel #0 {
    read_to_yuv

    byte_b_motion_vec_search $bPicHdr $y $prevY $nextY $interPrev $interNext $sliceInfo $fmv $bmv
    byte_y_to_sc_b  $y $prevY $nextY $fmv $bmv $qScale mpeg-intra mpeg-non-intra $scY
    byte_uv_to_sc_b $u $prevU $nextU $fmv $bmv $qScale mpeg-intra mpeg-non-intra $scU 
    byte_uv_to_sc_b $v $prevV $nextV $fmv $bmv $qScale mpeg-intra mpeg-non-intra $scV

    mpeg_pic_hdr_set_temporal_ref       $bPicHdr $temporalRef
    mpeg_pic_hdr_encode $bPicHdr $obp
    mpeg_pic_b_encode $bPicHdr $scY $scU $scV $fmv $bmv $qScale $sliceInfo $obp
    increment_time
}}


#------------------------------------------------------------------------
# initialization
#------------------------------------------------------------------------
set pnmHdr [pnm_hdr_new]
set seqHdr  [mpeg_seq_hdr_new]
set gopHdr  [mpeg_gop_hdr_new]
set iPicHdr [mpeg_pic_hdr_new]
set pPicHdr [mpeg_pic_hdr_new]
set bPicHdr [mpeg_pic_hdr_new]

set mbw [expr int(($width+15)/16)]
set mbh [expr int(($height+15)/16)]
set w [expr $mbw*16]
set h [expr $mbh*16]
set halfw [expr $w/2]
set halfh [expr $h/2]

# this should be enough for one frame
# (no reason to have more than one slice/frame in this example)
set sliceInfo {1000}

set r       [byte_new $w $h]
set g       [byte_new $w $h]
set b       [byte_new $w $h]
set y       [byte_new $w $h]
set prevY   [byte_new $w $h]
set nextY   [byte_new $w $h]
set u       [byte_new $halfw $halfh]
set v       [byte_new $halfw $halfh]
set prevU   [byte_new $halfw $halfh]
set prevV   [byte_new $halfw $halfh]
set nextU   [byte_new $halfw $halfh]
set nextV   [byte_new $halfw $halfh]
set qScale   [byte_new $mbw $mbh]
set scY     [sc_new [expr $mbw*2] [expr $mbh*2]]
set scU     [sc_new $mbw $mbh]
set scV     [sc_new $mbw $mbh]
set fmv     [vector_new $mbw $mbh]
set bmv     [vector_new $mbw $mbh]
set interPX  [byte_new [expr $w-1] $h]
set interPY  [byte_new $w [expr $h-1]]
set interPXY [byte_new [expr $w-1] [expr $h-1]]
set interNX  [byte_new [expr $w-1] $h]
set interNY  [byte_new $w [expr $h-1]]
set interNXY [byte_new [expr $w-1] [expr $h-1]]
set interPrev "$interPX $interPY $interPXY"
set interNext "$interNX $interNY $interNXY"

set fileSize [file size ${namePrefix}000.ppm]
set bp  [bitparser_new]
set bs  [bitstream_new $fileSize]
set obp [bitparser_new]
set obs [bitstream_new $buffer_size]
bitparser_wrap $obp $obs
set outFile [open $outFileName w]
fconfigure $outFile -translation binary

byte_set $qScale 4

#------------------------------------------------------------------------
# encode
#------------------------------------------------------------------------
mpeg_seq_hdr_set_width          $seqHdr $width 
mpeg_seq_hdr_set_height         $seqHdr $height
mpeg_seq_hdr_set_aspect_ratio   $seqHdr 1.000
mpeg_seq_hdr_set_pic_rate       $seqHdr $frames_per_second
mpeg_seq_hdr_set_bit_rate       $seqHdr -1
mpeg_seq_hdr_set_buffer_size    $seqHdr 16
mpeg_seq_hdr_set_constrained    $seqHdr 0
mpeg_seq_hdr_set_default_iqt    $seqHdr
mpeg_seq_hdr_set_default_niqt   $seqHdr

mpeg_seq_hdr_encode $seqHdr $obp

# we are not going to change these
mpeg_gop_hdr_set_drop_frame_flag    $gopHdr 0 
mpeg_gop_hdr_set_closed_gop         $gopHdr 1
mpeg_gop_hdr_set_broken_link        $gopHdr 0
mpeg_gop_hdr_set_hours              $gopHdr $hours
mpeg_gop_hdr_set_minutes            $gopHdr $minutes
mpeg_gop_hdr_set_seconds            $gopHdr $seconds
mpeg_gop_hdr_set_pictures           $gopHdr $pictures

# I Frame picHdr
mpeg_pic_hdr_set_vbv_delay          $iPicHdr 0
mpeg_pic_hdr_set_type               $iPicHdr i-frame
mpeg_pic_hdr_set_full_pel_forward   $iPicHdr 0
mpeg_pic_hdr_set_forward_f_code     $iPicHdr 0
mpeg_pic_hdr_set_full_pel_backward  $iPicHdr 0
mpeg_pic_hdr_set_backward_f_code    $iPicHdr 0

# P Frame picHdr
mpeg_pic_hdr_set_vbv_delay          $pPicHdr 0
mpeg_pic_hdr_set_type               $pPicHdr p-frame
mpeg_pic_hdr_set_full_pel_forward   $pPicHdr 0
mpeg_pic_hdr_set_forward_f_code     $pPicHdr $forward_f_code
mpeg_pic_hdr_set_full_pel_backward  $pPicHdr 0
mpeg_pic_hdr_set_backward_f_code    $pPicHdr 0

# B Frame picHdr -- set forward_f_code to 0 for the first B frame
# (no previous image to get fmv from)
mpeg_pic_hdr_set_vbv_delay          $bPicHdr 0
mpeg_pic_hdr_set_type               $bPicHdr b-frame
mpeg_pic_hdr_set_full_pel_forward   $bPicHdr 0
mpeg_pic_hdr_set_forward_f_code     $bPicHdr 0
mpeg_pic_hdr_set_full_pel_backward  $bPicHdr 0
mpeg_pic_hdr_set_backward_f_code    $bPicHdr $backward_f_code

# Encode the first GOP of 4 frames
mpeg_gop_hdr_encode $gopHdr $obp 
set temporalRef 1
i_frame_encode
incr temporalRef -1
b_frame_encode
incr temporalRef 3
p_frame_encode
incr temporalRef -1
mpeg_pic_hdr_set_forward_f_code     $bPicHdr $forward_f_code
b_frame_encode

puts "Writing to file..."
set size [bitparser_tell $obp]
bitstream_channel_write_segment $obs $outFile 0 $size
bitparser_wrap $obp $obs

#Encode the remaining GOP's
set gop_start 4
mpeg_gop_hdr_set_closed_gop         $gopHdr 0
for {set gop 4} {$gop < $num_of_frames} {incr gop $gop_size} {
    mpeg_gop_hdr_set_hours      $gopHdr $hours
    mpeg_gop_hdr_set_minutes    $gopHdr $minutes
    mpeg_gop_hdr_set_seconds    $gopHdr $seconds
    mpeg_gop_hdr_set_pictures   $gopHdr $pictures

    mpeg_gop_hdr_encode $gopHdr $obp 
    set current_gop_size 0
    set temporalRef 1
    for {set pic 0} {$pic < 2} {incr pic 1} {
        
        i_frame_encode
        incr temporalRef -1
        b_frame_encode
        incr temporalRef 3
        p_frame_encode
        incr temporalRef -1
        b_frame_encode
        incr temporalRef 3
        incr current_gop_size 4

        puts "Writing to file..."
        set size [bitparser_tell $obp]
        bitstream_channel_write_segment $obs $outFile 0 $size
        bitparser_wrap $obp $obs
    }
    incr gop_start $current_gop_size
}
# don't forget this
mpeg_seq_end_code_encode $obp
set size [bitparser_tell $obp]
bitstream_channel_write_segment $obs $outFile 0 $size

close $outFile

#----------------------------------------------------------------
# clean up
#----------------------------------------------------------------
bitstream_free $bs
bitparser_free $bp
bitstream_free $obs
bitparser_free $obp

pnm_hdr_free $pnmHdr
mpeg_seq_hdr_free $seqHdr
mpeg_gop_hdr_free $gopHdr
mpeg_pic_hdr_free $iPicHdr
mpeg_pic_hdr_free $pPicHdr
mpeg_pic_hdr_free $bPicHdr

byte_free $r
byte_free $g
byte_free $b
byte_free $y
byte_free $u
byte_free $v
byte_free $prevY
byte_free $prevU
byte_free $prevV
byte_free $qScale
sc_free $scY
sc_free $scU
sc_free $scV
vector_free $fmv
vector_free $bmv
byte_free $interPX  
byte_free $interPY  
byte_free $interPXY 
byte_free $interNX  
byte_free $interNY
byte_free $interNXY