/*------------------------------------------------------------------------
 *
 * 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.
 *
 * 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 BIBPBIBPB... and each GOP contains 8 frames.
 *
 * Note: the input files should be named 000.ppm, 001.ppm ...
 *------------------------------------------------------------------------
 */

#include <sys/stat.h>
#include <dvmbasic.h>
#include <dvmmpeg.h>
#include <dvmpnm.h>
#include <dvmcolor.h>

#define PREFIX_SIZE         100     /* prefix to the input files including directory */
#define NUM_OF_FRAMES       20
#define GOP_SIZE            8
#define BUFFER_SIZE         100000  /* large enough to hold 4 frames and some headers */
#define FRAMES_PER_SECOND   30
#define FORWARD_F_CODE      2
#define BACKWARD_F_CODE     2
#define WIDTH               176     /* width and height of each frame in pixels */
#define HEIGHT              120

int fsize (char *name)
{
#ifdef __WIN32__
     struct _stat s;
#else
     struct stat s;
#endif
    int result;

    result = stat(name, &s);
    return s.st_size;
}

int pictures = 0;
int seconds = 0;
int minutes = 0;
int hours = 0;

/*
 *------------------------------------------------------------------------
 * Use this to keep track of how many picture frames, seconds, minutes,
 * hours.  This information is encoded in the GOP headers.
 *------------------------------------------------------------------------
 */
void
IncrementTime()
{
    if ( pictures == FRAMES_PER_SECOND ) {
        pictures = 0;
        seconds++;
        if ( seconds == 60 ) {
            seconds = 0;
            minutes++;
            if ( minutes == 60 ) {
                minutes = 0;
                hours++;
            }
        }
    }
}


/*
 *------------------------------------------------------------------------
 * Swap two pointers
 *------------------------------------------------------------------------
 */

void Swap(a, b)
    void **a, **b;
{
    void *temp;
    temp = *a;
    *a = *b;
    *b = temp;
}

/*
 *------------------------------------------------------------------------
 * Read in a PPM file into r, g, b ByteImages, then convert it to y, u, v
 * with 4:2:0 sampling
 *------------------------------------------------------------------------
 */

#define READ_TO_YUV(bp, bs, r, g, b, y, u, v) {\
    sprintf(fileName, "%s%03d.ppm", namePrefix, temporalRef + gopStart);\
    printf("Processing %s\n", fileName);\
    inFile = fopen(fileName, "rb");\
    if (inFile == NULL) {\
	fprintf(stderr, "unable to open %s for reading.\n", fileName);\
	exit(1);\
    }\
    BitParserWrap(bp, bs);\
    BitStreamFileRead(bs, inFile, 0);\
    PnmHdrParse(bp, pnmHdr);\
    PpmParse(bp, r, g, b);\
    fclose(inFile);\
    RgbToYuv420(r, g, b, y, u, v);\
}


/*
 *------------------------------------------------------------------------
 * I Frame
 *------------------------------------------------------------------------
 */
#define I_FRAME_ENCODE {\
    READ_TO_YUV(bp, bs, r, g, b, y, u, v);\
    Swap(&nextY, &prevY);\
    Swap(&nextU, &prevU);\
    Swap(&nextV, &prevV);\
    Swap(&interNext, &interPrev);\
    \
    ByteToScI(y, qScale, MPEG_INTRA, scY);\
    ByteToScI(u, qScale, MPEG_INTRA, scU);\
    ByteToScI(v, qScale, MPEG_INTRA, scV);\
    \
    MpegPicHdrSetTemporalRef(iPicHdr, temporalRef);\
    MpegPicHdrEncode(iPicHdr, obp);\
    MpegPicIEncode (iPicHdr, scY, scU, scV, qScale, sliceInfo, sliceInfoLen, obp);\
    \
    /* decode the image just encoded to use for encoding next P, B frame */\
    ScDequantize(scY, qScale, MPEG_INTRA, scY);\
    ScDequantize(scU, qScale, MPEG_INTRA, scU);\
    ScDequantize(scV, qScale, MPEG_INTRA, scV);\
    ScIToByte(scY, nextY);\
    ScIToByte(scU, nextU);\
    ScIToByte(scV, nextV);\
    ByteComputeIntermediates(nextY, interNext);\
    IncrementTime();\
}\


/*
 *------------------------------------------------------------------------
 * P Frame
 *------------------------------------------------------------------------
 */
#define P_FRAME_ENCODE {\
    READ_TO_YUV(bp, bs, r, g, b, y, u, v);\
    Swap(&nextY, &prevY);\
    Swap(&nextU, &prevU);\
    Swap(&nextV, &prevV);\
    Swap(&interNext, &interPrev);\
    \
    BytePMotionVecSearch(pPicHdr, y, prevY, interPrev, fmv);\
    ByteYToScP (y, prevY, fmv, qScale, MPEG_INTRA, MPEG_NON_INTRA, scY);\
    ByteUVToScP(u, prevU, fmv, qScale, MPEG_INTRA, MPEG_NON_INTRA, scU);\
    ByteUVToScP(v, prevV, fmv, qScale, MPEG_INTRA, MPEG_NON_INTRA, scV);\
    \
    MpegPicHdrSetTemporalRef(pPicHdr, temporalRef);\
    MpegPicHdrEncode (pPicHdr, obp);\
    MpegPicPEncode (pPicHdr, scY, scU, scV, fmv, qScale, sliceInfo, sliceInfoLen, obp);\
    \
    /* decode the image just encoded to use for encoding next B frame */\
    ScNonIDequantize(scY, qScale, MPEG_INTRA, MPEG_NON_INTRA, scY);\
    ScNonIDequantize(scU, qScale, MPEG_INTRA, MPEG_NON_INTRA, scU);\
    ScNonIDequantize(scV, qScale, MPEG_INTRA, MPEG_NON_INTRA, scV);\
    ScPToY (scY, fmv, prevY, nextY);\
    ScPToUV(scU, fmv, prevU, nextU);\
    ScPToUV(scV, fmv, prevV, nextV);\
    ByteComputeIntermediates(nextY, interNext);\
    IncrementTime();\
}


/*
 *------------------------------------------------------------------------
 * B Frame
 *------------------------------------------------------------------------
 */
#define B_FRAME_ENCODE {\
    READ_TO_YUV(bp, bs, r, g, b, y, u, v);\
    ByteBMotionVecSearch(bPicHdr, y, prevY, nextY, interPrev, interNext, sliceInfo, sliceInfoLen, fmv, bmv);\
    ByteYToScB (y, prevY, nextY, fmv, bmv, qScale, MPEG_INTRA, MPEG_NON_INTRA, scY);\
    ByteUVToScB(u, prevU, nextU, fmv, bmv, qScale, MPEG_INTRA, MPEG_NON_INTRA, scU);\
    ByteUVToScB(v, prevV, nextV, fmv, bmv, qScale, MPEG_INTRA, MPEG_NON_INTRA, scV);\
    MpegPicHdrSetTemporalRef(bPicHdr, temporalRef);\
    MpegPicHdrEncode (bPicHdr, obp);\
    MpegPicBEncode (bPicHdr, scY, scU, scV, fmv, bmv, qScale, sliceInfo, sliceInfoLen, obp);\
    IncrementTime();\
}

int main(int argc, char *argv[])
{
    FILE *inFile, *outFile;
    char fileName[PREFIX_SIZE+3];
    char namePrefix[PREFIX_SIZE];
    BitParser *bp  = BitParserNew();        /* input bitstream */
    BitStream *bs;
    BitParser *obp = BitParserNew();        /* output bitstream */
    BitStream *obs = BitStreamNew(BUFFER_SIZE);

    ByteImage *r, *g, *b, *y, *u, *v;
    ByteImage *prevY, *prevU, *prevV;
    ByteImage *nextY, *nextU, *nextV;
    ByteImage *interPrev[3], *interNext[3];
    ByteImage *qScale;
    ScImage *scY, *scU, *scV;
    VectorImage *fmv, *bmv;

    PnmHdr *pnmHdr = PnmHdrNew();
    MpegSeqHdr *seqHdr = MpegSeqHdrNew();
    MpegGopHdr *gopHdr = MpegGopHdrNew();
    MpegPicHdr *iPicHdr = MpegPicHdrNew();
    MpegPicHdr *pPicHdr = MpegPicHdrNew();
    MpegPicHdr *bPicHdr = MpegPicHdrNew();

    int mbw = (WIDTH + 15) / 16;
    int mbh = (HEIGHT + 15) / 16;
    int w = mbw*16;
    int h = mbh*16;
    int halfw = w/2;
    int halfh = h/2;

    int sliceInfo[] = {1000};       /* this should be enough for one frame */
    int sliceInfoLen = 1;
    int gop, pic;
    int temporalRef, gopStart = 0, currentGopSize;

    r       = ByteNew(w, h);
    g       = ByteNew(w, h);
    b	    = ByteNew(w, h);
    y       = ByteNew(w, h);
    prevY   = ByteNew(w, h);
    nextY   = ByteNew(w, h);
    u       = ByteNew(halfw, halfh);
    v       = ByteNew(halfw, halfh);
    prevU   = ByteNew(halfw, halfh);
    prevV   = ByteNew(halfw, halfh);
    nextU   = ByteNew(halfw, halfh);
    nextV   = ByteNew(halfw, halfh);
    qScale  = ByteNew(mbw, mbh);
    scY     = ScNew(mbw*2, mbh*2);
    scU     = ScNew(mbw, mbh);
    scV     = ScNew(mbw, mbh);
    fmv     = VectorNew(mbw, mbh);
    bmv     = VectorNew(mbw, mbh);
    
    interPrev[0] = ByteNew(w-1, h);
    interNext[0] = ByteNew(w-1, h);
    interPrev[1] = ByteNew(w, h-1);
    interNext[1] = ByteNew(w, h-1);
    interPrev[2] = ByteNew(w-1, h-1);
    interNext[2] = ByteNew(w-1, h-1);

    if (argc < 2) {
	fprintf(stderr, "Not enough parameters : %s <input directroy> <output file name>\n", argv[0]);
	exit(1);
    }

    /* Get input file size */
    strcpy(namePrefix, argv[1]);
    sprintf(fileName, "%s000.ppm", namePrefix);
    inFile = fopen(fileName, "rb");
    if (inFile == NULL) {
	fprintf(stderr, "unable to open %s for reading.\n", fileName);
	exit(1);
    }
    bs  = BitStreamNew(fsize(fileName));
    fclose(inFile);

    /* Create outfile file */
    outFile = fopen(argv[2], "wb");
    if (outFile == NULL) {
	fprintf(stderr, "unable to open %s for writing.\n", argv[2]);
	exit(1);
    }
    BitParserWrap(obp, obs);

    ByteSet(qScale, 4);

    MpegSeqHdrSetWidth(seqHdr, WIDTH);
    MpegSeqHdrSetHeight(seqHdr, HEIGHT);
    MpegSeqHdrSetAspectRatio(seqHdr, 1.000);
    MpegSeqHdrSetPicRate(seqHdr, FRAMES_PER_SECOND);
    MpegSeqHdrSetBitRate(seqHdr, -1);
    MpegSeqHdrSetBufferSize(seqHdr, 16);
    MpegSeqHdrSetConstrained(seqHdr, 0);
    MpegSeqHdrSetDefaultIQT(seqHdr);
    MpegSeqHdrSetDefaultNIQT(seqHdr);

    MpegSeqHdrEncode(seqHdr, obp);

    /* GOP header */
    MpegGopHdrSetDropFrameFlag(gopHdr, 0);
    MpegGopHdrSetClosedGop(gopHdr, 1);
    MpegGopHdrSetBrokenLink(gopHdr, 0);
    MpegGopHdrSetHours(gopHdr, hours);
    MpegGopHdrSetMinutes(gopHdr, minutes);
    MpegGopHdrSetSeconds(gopHdr, seconds);
    MpegGopHdrSetPictures(gopHdr, pictures);

    /* I Frame picHdr */
    MpegPicHdrSetVBVDelay(iPicHdr, 0);
    MpegPicHdrSetType(iPicHdr, I_FRAME);
    MpegPicHdrSetFullPelForward(iPicHdr, 0);
    MpegPicHdrSetForwardFCode(iPicHdr, 0);
    MpegPicHdrSetFullPelBackward(iPicHdr, 0);
    MpegPicHdrSetBackwardFCode(iPicHdr, 0);

    /* P Frame picHdr */
    MpegPicHdrSetVBVDelay(pPicHdr, 0);
    MpegPicHdrSetType(pPicHdr, P_FRAME);
    MpegPicHdrSetFullPelForward(pPicHdr, 0);
    MpegPicHdrSetForwardFCode(pPicHdr, FORWARD_F_CODE);
    MpegPicHdrSetFullPelBackward(pPicHdr, 0);
    MpegPicHdrSetBackwardFCode(pPicHdr, 0);

    /* B Frame picHdr */
    MpegPicHdrSetVBVDelay(bPicHdr, 0);
    MpegPicHdrSetType(bPicHdr, B_FRAME);
    MpegPicHdrSetFullPelForward(bPicHdr, 0);
    MpegPicHdrSetForwardFCode(bPicHdr, 0);                 /* special case for 1st frame */
    MpegPicHdrSetFullPelBackward(bPicHdr, 0);
    MpegPicHdrSetBackwardFCode(bPicHdr, BACKWARD_F_CODE);

    /* Encode the first GOP, which is different because of the first B frame */
    MpegGopHdrEncode(gopHdr, obp);
    temporalRef = 1;
    I_FRAME_ENCODE;
    temporalRef += -1;
    B_FRAME_ENCODE;
    temporalRef += 3;
    P_FRAME_ENCODE;
    temporalRef += -1;
    MpegPicHdrSetForwardFCode(bPicHdr, FORWARD_F_CODE);   /* for the rest do, bi-directional */
    B_FRAME_ENCODE;
    
    /* write what we have so far to the bitstream */
    printf("Writing to file...\n");
    BitStreamFileWriteSegment(obs, outFile, 0, BitParserTell(obp));
    BitParserWrap(obp, obs);

    /* Encode the remaining GOP's */
    gopStart = 4;
    for (gop = 4; gop < NUM_OF_FRAMES; gop += GOP_SIZE) {
        MpegGopHdrSetHours(gopHdr, hours);
        MpegGopHdrSetMinutes(gopHdr, minutes);
        MpegGopHdrSetSeconds(gopHdr, seconds);
        MpegGopHdrSetPictures(gopHdr, pictures);

        MpegGopHdrEncode(gopHdr, obp);
        currentGopSize = 0;
        temporalRef = 1;

        for (pic = 0; pic < 2; pic++) {
            I_FRAME_ENCODE;
            temporalRef += -1;
            B_FRAME_ENCODE;
            temporalRef += 3;
            P_FRAME_ENCODE;
            temporalRef += -1;
            B_FRAME_ENCODE;
            temporalRef += 3;
            currentGopSize += 4;

            /* write two frames to the bitstream */
            printf("Writing to file...\n");
            BitStreamFileWriteSegment(obs, outFile, 0, BitParserTell(obp));
            BitParserWrap(obp, obs);
        }
        gopStart += currentGopSize;
    }

    /* don't forget this */
    MpegSeqEndCodeEncode(obp);
    BitStreamFileWriteSegment(obs, outFile, 0, BitParserTell(obp));
    BitParserWrap(obp, obs);

    fclose(outFile);


    /* lots of things to free up */
    BitStreamFree(bs);
    BitParserFree(bp);
    BitStreamFree(obs);
    BitParserFree(obp);

    PnmHdrFree(pnmHdr);
    MpegSeqHdrFree(seqHdr);
    MpegGopHdrFree(gopHdr);
    MpegPicHdrFree(iPicHdr);
    MpegPicHdrFree(pPicHdr);
    MpegPicHdrFree(bPicHdr);

    ByteFree(r);
    ByteFree(g);
    ByteFree(b);
    ByteFree(y);
    ByteFree(u);
    ByteFree(v);
    ByteFree(prevY);
    ByteFree(prevU);
    ByteFree(prevV);
    ByteFree(qScale);
    ScFree(scY);
    ScFree(scU);
    ScFree(scV);
    VectorFree(fmv);
    VectorFree(bmv);
    for (pic = 0; pic < 3; pic++) {
        ByteFree(interPrev[pic]);
        ByteFree(interNext[pic]);
    }

    return 0;
}