TKMK00

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Cereal64.Common.Utils;

//Written by mib_f8sm9c on May 2016, blah blah blah

namespace MarioKartTestingTool
{
    public static class TKMK00Encoder
    {
        /// <summary>
        /// Huffman trees have a left/right node and a value. In TKMK00, if the left or right node exists, then it's just a traversal
        ///  node, and will not be an output. Traversal nodes will have a value >= 0x20, while output nodes will have a value between
        ///  0x00 - 0x1F (all values for a 5-bit number).
        ///
        /// While it isn't necessary in decoding to store the values of traversal nodes, we DO need to have it work for the encoding
        ///  process, since the assembly code uses two arrays to simulate a binary tree, and incorrect traversal values will make it
        ///  non-functional.
        /// </summary>
        public class TKMK00HuffmanTreeNode
        {
            public TKMK00HuffmanTreeNode Left;
            public TKMK00HuffmanTreeNode Right;
            public int Value;

            public bool IsEndNode
            { get { return Left == null && Right == null; } }
        }

        /// <summary>
        /// The 0x2C size header for the TKMK00 data. Starts with "TKMK00".
        /// </summary>
        public class TKMK00Header
        {
            public byte RepeatModeMask;
            public byte Unknown;
            public ushort Width;
            public ushort Height;
            public int[] ChannelPointers = new int[8];

            public TKMK00Header()
            {
                //Empty, let the user define the values themself
            }

            public TKMK00Header(byte[] headerData)
            {
                RepeatModeMask = ByteHelper.ReadByte(headerData, 0x6);
                Unknown = ByteHelper.ReadByte(headerData, 0x7);
                Width = ByteHelper.ReadUShort(headerData, 0x8);
                Height = ByteHelper.ReadUShort(headerData, 0xA);
                for (int i = 0; i < 8; i++)
                    ChannelPointers[i] = ByteHelper.ReadInt(headerData, 0xC + 4 * i);
            }

            public byte[] GetAsBytes() //TEST THIS
            {
                return ByteHelper.CombineIntoBytes('T', 'K', 'M', 'K', '0', '0',
                    RepeatModeMask, Unknown, Width, Height, ChannelPointers);
            }

            /// <summary>
            /// Reads the RepeatModeMask for the specific channel's bit, and returns if it's flipped on or off.
            /// </summary>
            public bool RepeatEnabledFor(int channelNum)
            {
                int maskCompare = 0x1 << channelNum;
                return ((RepeatModeMask & maskCompare) != 0);
            }

            public static int DataSize { get { return 0x2C; } }
        }

        /// <summary>
        /// This class is combined from two different bit readers specified in the decoding assembly. The first form of it
        ///  (no repeat mode) reads in 4 bytes at a time, and reads 1 bit at a time. It can also return multi-bit sized
        ///  outputs (the decoding reads up to 6 at a time).
        ///
        /// The second form is repeat mode. In this form it loads 1 byte at a time and always returns 1 bit, but there's
        ///  an extra ignoreRepeat flag. If the flag is off (repeat mode is on), then when it hits the end of the byte it
        ///  re-uses the current byte. When RemainingBits reaches zero, it reads in a byte that tells it if the IgnoreRepeat
        ///  flag is on or off, and then the number of RemainingBits (will almost always be bigger than the size of a byte).
        /// </summary>
        public class TKMK00CommandReader
        {
            private byte[] _data; //Rom data, TKMK00 data, whatever you're reading from
            private int _dataPointer; //Reading offset in the data
            private int _buffer; //Last byte data loaded from _data
            private int _remainingBits; //Remaining bits before loading the next byte[s]
            private bool _repeatEnabled; //Repeat mode
            private bool _ignoreRepeatFlag; //Flag used in repeat mode

            public TKMK00CommandReader(byte[] data, int dataPointer, bool repeatEnabled)
            {
                _data = data;
                _buffer = 0;
                _dataPointer = dataPointer;
                _remainingBits = 0;
                _repeatEnabled = repeatEnabled;
                _ignoreRepeatFlag = false;

                //Here, we set back the pointer 1 load's worth so when it reads the first value it moves up and reads the correct location
                if (!repeatEnabled)
                    _dataPointer -= 4;
                else
                    _dataPointer--;
            }

            public int ReadBits()
            {
                return ReadBits(1);
            }

            public int ReadBits(int bitCount) //Bitcount only counts for non-repeating readers
            {
                int commandValue = 0;//clear the command value

                if (!_repeatEnabled) //Not repeat mode
                {
                    while (bitCount > 0) //Read out the # of bits in the parameter
                    {
                        if (_remainingBits == 0) //If the end of the current buffer has been reached
                        {
                            //Load the next 4 bytes and progress the pointer/reset the remaining bits
                            _dataPointer += 4;
                            _buffer = ByteHelper.ReadInt(_data, _dataPointer);
                            _remainingBits = 0x20;
                        }

                        //Append the next bit in the buffer to the end of the commandValue
                        commandValue = ((_buffer >> (_remainingBits - 1)) & 0x1) | (commandValue << 1);
                        _remainingBits--;

                        bitCount--;
                    }
                }
                else //Repeat mode
                {
                    if (_remainingBits == 0) //out of data
                    {
                        _dataPointer++;
                        byte repeatInfo = ByteHelper.ReadByte(_data, _dataPointer);

                        _ignoreRepeatFlag = ((repeatInfo & 0x80) != 0); //The new ignoreRepeatFlag is stored in the top bit

                        if (_ignoreRepeatFlag)
                        {
                            _remainingBits = (ushort)(((repeatInfo & 0x7F) + 1) * 8); //Read new byte (repeatInfo + 1) times
                        }
                        else
                        {
                            _remainingBits = (ushort)((repeatInfo + 3) * 8); //Repeat byte (repeatInfo + 3) times
                        }

                        _dataPointer++;
                        _buffer = ByteHelper.ReadByte(_data, _dataPointer); //Read in next byte
                    }

                    int byteWrappedRemainingBits = (_remainingBits & 0x7);
                    if (byteWrappedRemainingBits == 0)
                        byteWrappedRemainingBits = 8; //Remaining bits on CURRENT byte

                    //Set the commandValue to the bit in the buffer
                    commandValue = ((_buffer >> (byteWrappedRemainingBits - 1)) & 0x1);
                    _remainingBits--;

                    byteWrappedRemainingBits--;

                    //Check if we've reached the end of the bit but not the remaining bits
                    if (byteWrappedRemainingBits == 0 && _remainingBits > 0)
                    {
                        if (_ignoreRepeatFlag) //Load in next byte if ignoreRepeatMode
                        {
                            _dataPointer++;
                            _buffer = ByteHelper.ReadByte(_data, _dataPointer);
                        }
                    }
                }
                return commandValue;
            }
        }

        public static void Encode()
        {
            //Finish me!
        }

        #region Decode

        /// <summary>
        /// Decodes the TKMK00 format into RGBA5551 formatted data. Uses a Huffman tree to store
        ///  color values that are added/subtracted/sometimes overwrite a predicted color for the
        ///  current pixel. Look up DPCM for a conceptual idea of what it's doing (the Huffman tree
        ///  stands in place for entropy coding.
        /// </summary>
        public static byte[] Decode(byte[] data, int tkmk00Offset, ushort alphaColor)
        {
            //Initialize the header & readers
            byte[] headerBytes = new byte[TKMK00Header.DataSize];
            Array.Copy(data, tkmk00Offset, headerBytes, 0, TKMK00Header.DataSize);
            TKMK00Header header = new TKMK00Header(headerBytes);
            TKMK00CommandReader masterReader = new TKMK00CommandReader(data, tkmk00Offset + TKMK00Header.DataSize, false);
            TKMK00CommandReader[] channelReaders = new TKMK00CommandReader[8];
            for(int i = 0; i < 8; i++)
                channelReaders[i] = new TKMK00CommandReader(data, tkmk00Offset + header.ChannelPointers[i], header.RepeatEnabledFor(i));

            //Set up the image data/buffers
            ushort[] rgbaBuffer = new ushort[0x40];
            for (int i = 0; i < 0x40; i++)
                rgbaBuffer[i] = 0xFF;
            byte[] colorChangeMap = new byte[header.Width * header.Height];
            byte[] imageData = new byte[header.Width * header.Height * 2];
            int pixelIndex = 0;

            //Set up the Huffman binary tree
            TKMK00HuffmanTreeNode headTreeNode = SetUpHuffmanTree(0x20, data, channelReaders[0]);

            ushort lastPixelColor = 0;

            //Iterate through each pixel in order left to right, top to bottom
            for (int row = 0; row < header.Height; row++)
            {
                for (int col = 0; col < header.Width; col++)
                {
                    //Look at the current pixel's color. If it's not empty, then it's already been
                    // set to its correct value, and we can skip to the next pixel
                    ushort currentPixelColor = ByteHelper.ReadUShort(imageData, pixelIndex * 2);

                    if (currentPixelColor != 0) //Color already exists
                    {
                        lastPixelColor = currentPixelColor;

                        //Test to make sure that the curent color is not the alpha value with the incorrect alpha channel value
                        ushort currentPixelWithoutAlpha = (ushort)(currentPixelColor & 0xFFFE);
                        if (currentPixelWithoutAlpha == alphaColor)
                        {
                            ByteHelper.WriteUShort(alphaColor, imageData, pixelIndex * 2);
                            lastPixelColor = alphaColor;
                        }

                        //Done, go to end of the loop
                    }
                    else
                    {
                        //Load up the channel reader that is associated with the given color change value in the
                        // colorChangeMap (low values = not much change around that pixel, high values = lots of change)
                        byte channelIndex = (byte)(colorChangeMap[pixelIndex] + 1);

                        int command = channelReaders[channelIndex].ReadBits(1); // 0 - Use the previous color, 1 - Use new color

                        if (command == 0)
                        {
                            ByteHelper.WriteUShort(lastPixelColor, imageData, pixelIndex * 2);
                            //End of this line
                        }
                        else
                        {
                            command = masterReader.ReadBits(1); // 0 - Create new RGBA, 1 - Use existing RGBA
                            if (command != 0)
                            {
                                //Load in the huffman values for the new green, red and blue. These are combined
                                // with a predicted pixel value for them later on.
                                int newGreen = RetrieveHuffmanTreeValue(data, headTreeNode, channelReaders[0]);
                                int newRed = RetrieveHuffmanTreeValue(data, headTreeNode, channelReaders[0]);
                                int newBlue = RetrieveHuffmanTreeValue(data, headTreeNode, channelReaders[0]);

                                //Retreive the pixel colors from the pixel above and the pixel to the left
                                ushort rgbaTop, rgbaLeft;

                                if (row != 0)
                                {
                                    rgbaTop = ByteHelper.ReadUShort(imageData, (pixelIndex - header.Width) * 2);
                                    rgbaLeft = ByteHelper.ReadUShort(imageData, (pixelIndex - 1) * 2);
                                }
                                else
                                {
                                    rgbaTop = 0;
                                    if (col != 0)
                                        rgbaLeft = ByteHelper.ReadUShort(imageData, (pixelIndex - 1) * 2);
                                    else
                                        rgbaLeft = 0;
                                }

                                //Combine green values of the pixels to make our predicted pixel color
                                ushort greenTop = (byte)((rgbaTop & 0x7C0) >> 6);
                                ushort greenLeft = (byte)((rgbaLeft & 0x7C0) >> 6);
                                int greenPrediction = (greenTop + greenLeft) / 2;

                                //Combine the prediction & huffman value to make the output color
                                ColorCombine(greenPrediction, ref newGreen);

                                //Use the change between the old & new green values to project expected
                                // values for the red & blue colors
                                int greenChange = newGreen - greenPrediction;

                                //Combine red values of the pixels to make our predicted pixel color
                                ushort redTop = (byte)((rgbaTop & 0xF800) >> 11);
                                ushort redLeft = (byte)((rgbaLeft & 0xF800) >> 11);
                                int redPrediction = greenChange + (redTop + redLeft) / 2;
                                redPrediction = Math.Max(0, Math.Min(0x1F, redPrediction)); //Keep between 0 and 0x1F

                                //Combine the prediction & huffman value to make the output color
                                ColorCombine(redPrediction, ref newRed);

                                //Combine blue values of the pixels to make our predicted pixel color
                                ushort blueTop = (byte)((rgbaTop & 0x3E) >> 1);
                                ushort blueLeft = (byte)((rgbaLeft & 0x3E) >> 1);
                                int bluePrediction = greenChange + (blueTop + blueLeft) / 2;
                                bluePrediction = Math.Max(0, Math.Min(0x1F, bluePrediction)); //Keep between 0 and 0x1F

                                //Combine the prediction & huffman value to make the output color
                                ColorCombine(bluePrediction, ref newBlue);

                                //Make the newpixel color
                                currentPixelColor = (ushort)((newRed << 11) | (newGreen << 6) | (newBlue << 1));
                                if (currentPixelColor != alphaColor) //Only transparent if it matches the transparency pixel
                                    currentPixelColor |= 0x1;

                                //Add to the front of the color buffer
                                for (int i = rgbaBuffer.Length - 1; i > 0; i--)
                                    rgbaBuffer[i] = rgbaBuffer[i - 1];
                                rgbaBuffer[0] = currentPixelColor;
                            }
                            else //Use existing RGBA
                            {
                                command = masterReader.ReadBits(6); // Returns index of color in color buffer to use
                                currentPixelColor = rgbaBuffer[command];
                                if (command != 0)
                                {
                                    //Bump the selected color to the front of the buffer
                                    for (int i = command; i > 0; i--)
                                        rgbaBuffer[i] = rgbaBuffer[i - 1];
                                    rgbaBuffer[0] = currentPixelColor;
                                }
                            }

                            //Write the RGBA to the imageData
                            ByteHelper.WriteUShort(currentPixelColor, imageData, pixelIndex * 2);
                            lastPixelColor = currentPixelColor;

                            //Add nearby pixels to the colorChangeMap
                            bool hasLeftCol = (col != 0);
                            bool hasRightCol = (col < (header.Width - 1));
                            bool has2RightCols = (col < (header.Width - 2));
                            bool hasDownRow = (row < (header.Height - 1));
                            bool has2DownRows = (row < (header.Height - 2));

                            //Right 1
                            if (hasRightCol)
                                colorChangeMap[pixelIndex + 1]++;
                            //Right 2
                            if (has2RightCols)
                                colorChangeMap[pixelIndex + 2]++;
                            //Down 1 Left 1
                            if (hasDownRow && hasLeftCol)
                                colorChangeMap[pixelIndex + header.Width - 1]++;
                            //Down 1
                            if (hasDownRow)
                                colorChangeMap[pixelIndex + header.Width]++;
                            //Down 1 Right 1
                            if (hasDownRow && hasRightCol)
                                colorChangeMap[pixelIndex + header.Width + 1]++;
                            //Down 2
                            if (has2DownRows)
                                colorChangeMap[pixelIndex + header.Width * 2]++;

                            //Now test to see if we need to continue writing this color down the column
                            command = masterReader.ReadBits(1);//1 - repeat color, 0 - continue

                            if (command == 1) //Repeat color
                            {
                                //Basically move down one row each repeat, possibly moving to the side, and write the color again
                                int pixelOffset = 0;
                                ushort currentPixelColorOpaque = (ushort)(currentPixelColor | 0x1); //Not sure why this is the case, is it to catch it in the first if statement?

                                while(true) //I hate while(true)
                                {
                                    command = masterReader.ReadBits(2);//0 - advanced move, 1 - back one, 2 - no lateral move, 3 - forward one
                                    if (command == 0)
                                    {
                                        //Advanced move
                                        command = masterReader.ReadBits(1);//0 - stop, 1 - advanced move

                                        if (command == 0)
                                            break;

                                        command = masterReader.ReadBits(1); //0 - move back 2, 1 - move forward 2
                                        if (command == 0)
                                            pixelOffset -= 2;
                                        else
                                            pixelOffset += 2;
                                    }
                                    else if (command == 1)
                                        pixelOffset--;
                                    else if (command == 3)
                                        pixelOffset++;

                                    pixelOffset += header.Width; //move down a row
                                    ByteHelper.WriteUShort(currentPixelColorOpaque, imageData, (pixelIndex + pixelOffset) * 2);
                                }
                            }
                        }
                    }

                    //Next pixel
                    pixelIndex++;
                }
            }

            return imageData;
        }

        /// <summary>
        /// Creates a Huffman-style binary tree that holds the 5-bit color data referenced in the
        /// TKMK00 decoding process.
        /// </summary>
        private static TKMK00HuffmanTreeNode SetUpHuffmanTree(uint val, byte[] data, TKMK00CommandReader commandReader)
        {
            TKMK00HuffmanTreeNode newTree = new TKMK00HuffmanTreeNode();

            int command = commandReader.ReadBits(1);//0 - Make new branches, 1 - End current branch

            if (command != 0)
            {
                newTree.Value = (int)val; //Not used, but can't hurt to number it
                val++;
                newTree.Left = SetUpHuffmanTree(val, data,  commandReader);
                newTree.Right = SetUpHuffmanTree(val, data,  commandReader);
                return newTree;
            }

            //Else, return a node with a value
            int value = 0;
            int bitCount = 5;
            do
            {
                command = commandReader.ReadBits(1);
                value = value * 2 + command; //basically bitshifts s0 to the left and adds v0, aka it's loading 5 bytes straight from the comandReader
                bitCount -= 1;
            } while (bitCount > 0);

            newTree.Value = value;
            return newTree;
        }

        //Only called when creating new rgba value. Traverses the Huffman tree to return the correct value
        private static int RetrieveHuffmanTreeValue(byte[] data, TKMK00HuffmanTreeNode currentNode, TKMK00CommandReader commandReader)
        {
            while (!currentNode.IsEndNode)/*currentNode.Value >= 0x20*/
            {
                int command = commandReader.ReadBits(1); // 0 - left, 1 - right
                if (command == 0)
                    currentNode = currentNode.Left;
                else
                    currentNode = currentNode.Right;
            }

            return currentNode.Value;
        }

        /// <summary>
        /// Combines a predicted color with a new color. Sometimes the new color is just something added to/subtracted
        ///  from the predicted color, others it's just used as the new color. Complex behavior, but it seems to work.
        /// </summary>
        private static void ColorCombine(int predictedColor, ref int newColor)
        {
            //See if our numbers could possibly go over 0x1F or below 0x00 if we do the differential method. If they could,
            // then just use the new color.

            if (predictedColor >= 0x10) //Check if it'll go above 0x1F
            {
                int remainingSpace = (0x1F - predictedColor);
                int incrementDecrementValue = newColor >> 1;
                bool isAdding = (newColor & 0x1) != 0;
                if (remainingSpace < incrementDecrementValue || (remainingSpace == incrementDecrementValue && isAdding)) //Since an extra 1 is used when adding
                {
                    newColor = 0x1F - newColor;
                    return;
                }
            }
            else //Check if it'll go below 0
            {
                int incrementDecrementValue = newColor >> 1;
                bool isAdding = (newColor & 0x1) != 0;
                if (predictedColor < incrementDecrementValue || (predictedColor == incrementDecrementValue && isAdding))
                {
                    return;
                }
            }

            bool addToOld = (newColor & 0x1) != 0; //Last bit is used to determine if adding or subtracting the value
            newColor = newColor >> 1;
            if (addToOld)
                newColor += predictedColor + 1;
            else
                newColor = predictedColor - newColor;

            return;
        }

        #endregion
    }
}