Mondain

package com.altair.codec;

import org.apache.mina.common.ByteBuffer;

import com.jcraft.jzlib.JZlib;
import com.jcraft.jzlib.ZStream;

/**
 * Flash Screen Video encoder <br /> Converted to Java from FFmpeg source
 * authored by Alex Beregszaszi and Benjamin Larsson
 * <br />
 * Bitstream description The picture is divided into blocks that are zlib-compressed.
 *
 * The decoder is fed complete frames, the frameheader contains: 4bits of block
 * width 12bits of frame width 4bits of block height 12bits of frame height
 *
 * Directly after the header are the compressed blocks. The blocks have their
 * compressed size represented with 16bits in the begining. If the size = 0 then
 * the block is unchanged from the previous frame. All blocks are decompressed
 * until the buffer is consumed.
 *
 * Encoding ideas, a basic encoder would just use a fixed block size. Block
 * sizes can be multiples of 16, from 16 to 256. The blocks don't have to be
 * quadratic. A brute force search with a set of different block sizes should
 * give a better result than to just use a fixed size.
 *
 * TODO: Don't reencode the frame in brute force mode if the frame is a dupe.
 * Speed up. Make the difference check faster.
 *
 * @author Paul Gregoire (mondain@gmail.com)
 */
public class Encoder {

    // "Flash Screen Video"
    // flashsv PIX_FMT_BGR24

    private int image_width;
    private int image_height;
    private int block_width = 64;
    private int block_height = 64;

    private Frame previousFrame;

    private byte[] tmpblock;
    private byte[] encbuffer;

    private int block_size;

    private ZStream zstream;

    private int frameNumber = 0;

    private int copy_region_enc(byte[] sptr, byte[] dptr, int dx, int dy, int h, int w, int stride) {

        int copyLength = w * 3;

        byte[] nsptr = new byte[copyLength];
        byte[] npfptr = new byte[copyLength];
        int diff = 0;

        byte[] tmp = previousFrame.getData().asReadOnlyBuffer().array();

        int destIndex = 0;

        for (int i = dx + h; i > dx; i--) {
            int index = (i * stride) + dy * 3;
            System.arraycopy(sptr, index, nsptr, 0, copyLength);
            System.arraycopy(tmp, index, npfptr, 0, copyLength);
            for (int j = 0; j < w * 3; j++) {
                diff |= npfptr[j] ^ nsptr[j];
                dptr[j + destIndex] = nsptr[j];
            }
            destIndex += copyLength;
        }
        if (diff != 0) {
            return 1;
        }
        return 0;
    }

    public int init(int width, int height) {

        if (width > 4095 || height > 4095) {
            System.out
                    .printf("Input dimensions too large, input must be max 4096x4096 !\n");
            return -1;
        }

        zstream = new ZStream();

        int err = zstream.deflateInit(JZlib.Z_BEST_COMPRESSION);
        if (err == JZlib.Z_OK) {
            System.out.printf("Deflate initialized");
        } else {
            System.out.printf("Deflate init failed error code: %s", err);
            return -1;
        }

        image_width = width;
        image_height = height;

        tmpblock = new byte[3 * 256 * 256];
        encbuffer = new byte[image_width * image_height * 3];

        return 0;
    }

    public Frame encode(byte[] buf, int bufSize) {

        ByteBuffer outBuffer = ByteBuffer.allocate(bufSize * 8);
        outBuffer.setAutoExpand(true);

        // 4 bits of block width
        int width = (((block_width / 16) - 1) & 0x0f) << 4;
        // 12 bits of frame width
        width |= ((image_width & 0x0f00) >> 8) & 0x0f;

        outBuffer.putShort((short) width);

        // 4 bits of block height
        int height = (((block_height / 16) - 1) & 0x0f) << 4;
        // 12 bits of frame height
        height |= ((image_height & 0x0f00) >> 8) & 0x0f;

        outBuffer.putShort((short) height);

        int predBlocks = 0;

        int h_blocks = image_width / block_width;
        int h_part = image_width % block_width;
        int v_blocks = image_height / block_height;
        int v_part = image_height % block_height;

        // loop over all block columns
        for (int j = 0; j < v_blocks + (v_part != 0 ? 1 : 0); j++) {

            int hp = j * block_height; // horiz position in frame
            int hs = (j < v_blocks) ? block_height : v_part; // size of block

            // loop over all block rows
            for (int i = 0; i < h_blocks + (h_part != 0 ? 1 : 0); i++) {
                int wp = i * block_width; // vert position in frame
                int ws = (i < h_blocks) ? block_width : h_part; // size of block
                int ret = JZlib.Z_OK;

                // copy the block to the temp buffer before compression (if it
                // differs from the previous frame's block)
                int res = copy_region_enc(buf, tmpblock, image_height - (hp + hs + 1), wp, hs, ws, stride);

                if (res != 0) {
                    int zsize = 3 * block_width * block_height;

                    zstream.next_in = tmpblock;
                    zstream.next_in_index = 0;

                    zstream.next_out = encbuffer;
                    zstream.next_out_index = 0;

                    int sourceLength = 3 * ws * hs;

                    while (zstream.total_in != sourceLength && zstream.total_out < zsize) {
                        zstream.avail_in = zstream.avail_out = 1; // force small buffers
                        ret = zstream.deflate(JZlib.Z_NO_FLUSH);
                        if (ret != JZlib.Z_OK) {
                            System.out.printf("Deflate loop1 failed error code: %d\n", ret);
                        }
                    }

                    while (true) {
                        zstream.avail_out = 1;
                        ret = zstream.deflate(JZlib.Z_FINISH);
                        if (ret == JZlib.Z_STREAM_END) {
                            break;
                        }
                        if (ret != JZlib.Z_OK) {
                            System.out.printf("Deflate loop2 failed error code: %d\n", ret);
                        }
                    }

                    ret = zstream.deflateEnd();
                    if (ret != JZlib.Z_OK) {
                        System.out.printf(
                                "Error while compressing block %dx%d\n", i, j);
                    }

                    outBuffer.putShort((short) zsize);

                    //put the encoded / compressed bytes into the result
                    outBuffer.put(encbuffer);

                    System.out.printf("buf_pos = %d\n", outBuffer.position());
                } else {
                    predBlocks++;
                    outBuffer.putShort((short) 0);
                }
            }
        }

        Frame outFrame = new Frame();
        outFrame.setData(outBuffer);
        outFrame.setFrameNumber(frameNumber);

        if (predBlocks != 0) {
            outFrame.setFrameType(Frame.FRAME_TYPE_P);
        } else {
            outFrame.setFrameType(Frame.FRAME_TYPE_I);
        }

        outFrame.setEncoded(true);

        return outFrame;
    }

    public Frame encodeFrame(byte[] buf) {

        // First frame needs to be a keyframe
        if (frameNumber == 0) {
            previousFrame = new Frame();
        }

        int bufSize = buf.length;
        int bufferSizeEstimate = image_width * image_height * 3;

        if (bufSize < bufferSizeEstimate) {
            // Conservative upper bound check for compressed data
            System.out.printf("buf_size %d <  %d\n", bufSize, bufferSizeEstimate);
            return null;
        }

        Frame result = encode(buf, bufSize);

        // save the current frame
        previousFrame.setData(result.getData());
        previousFrame.setFrameNumber(frameNumber);
        previousFrame.setFrameType(result.getFrameType());

        // mark the frame type so the muxer can mux it correctly
        if (result.getFrameType() == Frame.FRAME_TYPE_I) {
            System.out.printf("Inserting key frame at frame %d\n", frameNumber);
        }

        frameNumber++;

        return result;
    }

}