using System;
using System.Collections;
using System.Collections.Generic;
using UnityEngine;



interface IBitStreamWriter
{
    void SeekZero();
    uint ReadBits(int numbits);
    void WriteBits(uint value, int numbits);
    void FlushBits();
}


class BitStreamWriter8 : IBitStreamWriter
{
    byte[] m_Buffer;
    UInt64 m_BitStage;
    int m_CurrentBitIdx;
    int m_CurrentByteIdx;

    public BitStreamWriter8(int capacity)
    {
        m_Buffer = new byte[capacity];
        m_CurrentBitIdx = 0;
        m_CurrentByteIdx = 0;
        m_BitStage = 0;
    }

    public void SeekZero()
    {
        m_CurrentBitIdx = 0;
        m_CurrentByteIdx = 0;
        m_BitStage = 0;
    }

    // Write the lower numbits from value into stream.
    public void WriteBits(UInt32 value, int numbits)
    {
        //Debug.Assert(numbits > 0 && numbits <= 32);
        //Debug.Assert((UInt64.MaxValue << numbits & value) == 0);

        m_BitStage |= ((UInt64)value << m_CurrentBitIdx);
        m_CurrentBitIdx += numbits;

        while (m_CurrentBitIdx >= 8)
        {
            m_Buffer[m_CurrentByteIdx++] = (byte)(m_BitStage & 0xff);
            m_CurrentBitIdx -= 8;
            m_BitStage >>= 8;
        }
    }

    // Read numbits from stream, return in lower bits
    public UInt32 ReadBits(int numbits)
    {
        //Debug.Assert(numbits > 0 && numbits <= 32);

        while (m_CurrentBitIdx < 32)
        {
            m_BitStage |= (UInt64)m_Buffer[m_CurrentByteIdx++] << m_CurrentBitIdx;
            m_CurrentBitIdx += 8;
        }

        var res = m_BitStage & (((UInt64)1 << numbits) - 1);
        m_BitStage >>= numbits;
        m_CurrentBitIdx -= numbits;
        return (UInt32)res;
    }

    public void FlushBits()
    {
        if (m_CurrentBitIdx > 0)
        {
            WriteBits(0, 8 - m_CurrentBitIdx);
        }
    }
}

class BitStreamWriter32 : IBitStreamWriter
{
    UInt32[] m_Buffer;
    UInt64 m_BitStage;
    int m_CurrentBitIdx;
    int m_CurrentWordIdx;

    // capacity must be multiple of 4
    public BitStreamWriter32(int capacity)
    {
        Debug.Assert((capacity % 4) == 0);

        m_Buffer = new UInt32[capacity / 4];
        m_CurrentBitIdx = 0;
        m_CurrentWordIdx = 0;
        m_BitStage = 0;
    }

    public void SeekZero()
    {
        m_CurrentBitIdx = 0;
        m_CurrentWordIdx = 0;
        m_BitStage = 0;
    }

    // Write the lower numbits from value into stream.
    public void WriteBits(UInt32 value, int numbits)
    {
        //Debug.Assert(numbits > 0 && numbits <= 32);
        //Debug.Assert((UInt64.MaxValue << numbits & value) == 0);

        m_BitStage |= ((UInt64)value << m_CurrentBitIdx);
        m_CurrentBitIdx += numbits;

        if (m_CurrentBitIdx >= 32)
        {
            int outgoing = (int)(m_BitStage & 0xffffffff);
            m_Buffer[m_CurrentWordIdx++] = (UInt32)System.Net.IPAddress.HostToNetworkOrder(outgoing);
            m_CurrentBitIdx -= 32;
            m_BitStage >>= 32;
        }
    }

    // Read numbits from stream, return in lower bits
    public UInt32 ReadBits(int numbits)
    {
        //Debug.Assert(numbits > 0 && numbits <= 32);

        if (m_CurrentBitIdx < 32)
        {
            m_BitStage |= (UInt64)((UInt32)System.Net.IPAddress.NetworkToHostOrder((int)m_Buffer[m_CurrentWordIdx++])) << m_CurrentBitIdx;
            m_CurrentBitIdx += 32;
        }

        var res = m_BitStage & (((UInt64)1 << numbits) - 1);
        m_BitStage >>= numbits;
        m_CurrentBitIdx -= numbits;
        return (UInt32)res;
    }

    public void FlushBits()
    {
        if (m_CurrentBitIdx > 0)
        {
            WriteBits(0, 32 - m_CurrentBitIdx);
        }
    }
}

class Program
{
    public static void Main(string[] args)
    {
        var num_tests = 1000000;

        var stream = new BitStreamWriter32(num_tests * 4);
        ProfileIt("Stream32", stream, num_tests);

        var stream8 = new BitStreamWriter8(num_tests * 4);
        ProfileIt("Stream8", stream8, num_tests);
    }

    struct TestCase
    {
        public int numbits;
        public UInt32 value;
    }

    static void ProfileIt(string name, IBitStreamWriter stream, int num_tests)
    {
        Debug.Log("\n==================");
        Debug.Log("Profiling " + name);
        Debug.Log("==================");
        Debug.Log("Creating " + num_tests + " testcases");

        var rand = new System.Random();

        var bits = 0;
        var cases = new TestCase[num_tests];
        for (var i = 0; i < num_tests; i++)
        {
            cases[i].numbits = rand.Next(1, 33);
            bits += cases[i].numbits;
            cases[i].value = (UInt32)rand.Next(0, 65536) | ((UInt32)rand.Next(0, 65536) << 16);
            cases[i].value &= (UInt32)(((UInt64)1 << cases[i].numbits) - 1);
        }
        Debug.Log("Done");

        var t = new System.Diagnostics.Stopwatch();
        t.Start();
        for (var i = 0; i < num_tests; i++)
        {
            var tc = cases[i];
            stream.WriteBits(tc.value, tc.numbits);
        }
        stream.FlushBits();

        t.Stop();
        Debug.Log("Wrote " + num_tests + " ints with 1-32 bits. Total " + bits + " bits");
        Debug.Log("Time: " + t.ElapsedMilliseconds + " ms ("+(bits/8/1024/1024*1000/t.ElapsedMilliseconds)+") mb/s");

        stream.SeekZero();

        t.Reset();
        t.Start();

        foreach (var tc in cases)
        {
            var r = stream.ReadBits(tc.numbits);
            //Debug.Assert(r == tc.value);
        }
        t.Stop();
        Debug.Log("Read " + cases.Length + " chunks");
        Debug.Log("Time: " + t.ElapsedMilliseconds + " ms ("+(bits/8/1024/1024*1000/t.ElapsedMilliseconds)+" mb/s)");
    }
}


public class BitStream : MonoBehaviour {

	IEnumerator Start ()
    {
        // Wait a bit to let stuff settle
        Debug.developerConsoleVisible = true;
        yield return new WaitForSeconds(2.0f);

        var args = new string[0];
        Program.Main(args);
    }
}