package h264dec

import (
	"reflect"
	"testing"
)

func TestMbTypeBinarization(t *testing.T) {
	tests := []struct {
		v, slice int
		want     []int
		err      error
	}{
		{v: 6, slice: sliceTypeI, want: []int{1, 0, 0, 1, 0, 0, 1}},
		{v: 26, slice: sliceTypeI, err: errBadMbType},
		{v: -1, slice: sliceTypeI, err: errBadMbType},
		{v: 4, slice: sliceTypeSI, want: []int{0}},
		{v: 6, slice: sliceTypeSI, want: []int{1, 1, 0, 0, 1, 0, 0, 0}},
		{v: 0, slice: sliceTypeSI, err: errBadMbType},
		{v: 27, slice: sliceTypeSI, err: errBadMbType},
		{v: 2, slice: sliceTypeP, want: []int{0, 1, 0}},
		{v: 3, slice: sliceTypeSP, want: []int{0, 0, 1}},
		{v: 7, slice: sliceTypeP, want: []int{1, 1, 0, 0, 0, 0, 1}},
		{v: 7, slice: sliceTypeSP, want: []int{1, 1, 0, 0, 0, 0, 1}},
		{v: -1, slice: sliceTypeP, err: errBadMbType},
		{v: 31, slice: sliceTypeP, err: errBadMbType},
		{v: 8, slice: sliceTypeB, want: []int{1, 1, 0, 1, 0, 1}},
		{v: 30, slice: sliceTypeB, want: []int{1, 1, 1, 1, 0, 1, 1, 0, 0, 1, 0, 1, 0}},
		{v: -1, slice: sliceTypeB, err: errBadMbType},
		{v: 49, slice: sliceTypeB, err: errBadMbType},
		{v: 6, slice: 20, err: errBadMbSliceType},
	}

	for i, test := range tests {
		got, err := mbTypeBinarization(test.v, test.slice)
		if err != test.err {
			t.Errorf("did not get expected error for test %d\nGot: %v\nWant: %v", i, err, test.err)
		}

		if !reflect.DeepEqual(got, test.want) {
			t.Errorf("did not get expected result for test %d\nGot: %v\nWant: %v", i, got, test.want)
		}
	}
}

func TestSubMbTypeBinarization(t *testing.T) {
	tests := []struct {
		v, slice int
		want     []int
		err      error
	}{
		{v: 2, slice: sliceTypeP, want: []int{0, 1, 1}},
		{v: 2, slice: sliceTypeSP, want: []int{0, 1, 1}},
		{v: -1, slice: sliceTypeSP, err: errBadMbType},
		{v: 4, slice: sliceTypeSP, err: errBadMbType},
		{v: 9, slice: sliceTypeB, want: []int{1, 1, 1, 0, 1, 0}},
		{v: 9, slice: 40, err: errBadSubMbSliceType},
	}

	for i, test := range tests {
		got, err := subMbTypeBinarization(test.v, test.slice)
		if err != test.err {
			t.Errorf("did not get expected error for test %d\nGot: %v\nWant: %v", i, err, test.err)
		}

		if !reflect.DeepEqual(got, test.want) {
			t.Errorf("did not get expected result for test %d\nGot: %v\nWant: %v", i, got, test.want)
		}
	}
}

var ctxIdxTests = []struct {
	binIdx       int
	maxBinIdxCtx int
	ctxIdxOffset int
	want         int
}{
	{0, 0, 0, 10000},
	{999, 0, 0, 10000},

	{0, 0, 3, 10000},
	{1, 0, 3, 276},
	{2, 0, 3, 3},
	{3, 0, 3, 4},
	{4, 0, 3, 10000},
	{5, 0, 3, 10000},
	{999, 0, 3, 7},

	{0, 0, 11, 10000},
	{999, 0, 11, 10000},

	{0, 0, 14, 0},
	{1, 0, 14, 1},
	{2, 0, 14, 10000},
	{999, 0, 14, 10000},

	{0, 0, 17, 0},
	{1, 0, 17, 276},
	{2, 0, 17, 1},
	{3, 0, 17, 2},
	{4, 0, 17, 10000},
	{999, 0, 17, 3},

	{0, 0, 21, 0},
	{1, 0, 21, 1},
	{2, 0, 21, 2},
	{999, 0, 21, 10000},

	{0, 0, 24, 10000},
	{999, 0, 24, 10000},

	{0, 0, 27, 10000},
	{1, 0, 27, 3},
	{2, 0, 27, 10000},
	{999, 0, 27, 5},

	{0, 0, 32, 0},
	{1, 0, 32, 276},
	{2, 0, 32, 1},
	{3, 0, 32, 2},
	{4, 0, 32, 10000},
	{999, 0, 32, 3},

	{0, 0, 36, 0},
	{1, 0, 36, 1},
	{2, 0, 36, 10000},
	{3, 0, 36, 3},
	{4, 0, 36, 3},
	{5, 0, 36, 3},

	{0, 0, 40, 10000},

	{0, 0, 47, 10000},
	{1, 0, 47, 3},
	{2, 0, 47, 4},
	{3, 0, 47, 5},
	{999, 0, 47, 6},

	{0, 0, 54, 10000},
	{1, 0, 54, 4},
	{999, 0, 54, 5},

	{0, 0, 60, 10000},
	{1, 0, 60, 2},
	{999, 0, 60, 3},

	{0, 0, 64, 10000},
	{1, 0, 64, 3},
	{2, 0, 64, 3},
	{999, 0, 64, 10000},

	{0, 0, 68, 0},
	{999, 0, 68, 10000},

	{0, 0, 69, 0},
	{1, 0, 69, 0},
	{2, 0, 69, 0},

	{0, 0, 70, 10000},
	{999, 0, 70, 10000},

	{0, 0, 73, 10000},
	{1, 0, 73, 10000},
	{2, 0, 73, 10000},
	{3, 0, 73, 10000},
	{4, 0, 73, 10000},
	{999, 0, 73, 10000},

	{0, 0, 77, 10000},
	{1, 0, 77, 10000},
	{999, 0, 77, 10000},

	{0, 0, 276, 0},
	{999, 0, 276, 10000},

	{0, 0, 399, 10000},
	{999, 0, 399, 10000},
}

// TestCtxIdx tests that the CtxIdx function returns the correct
// value given binIdx and ctxIdxOffset.
func TestCtxIdx(t *testing.T) {
	for _, tt := range ctxIdxTests {
		if got := CtxIdx(tt.binIdx, tt.maxBinIdxCtx, tt.ctxIdxOffset); got != tt.want {
			t.Errorf("CtxIdx(%d, %d, %d) = %d, want %d", tt.binIdx, tt.maxBinIdxCtx, tt.ctxIdxOffset, got, tt.want)
		}
	}
}

func TestUnaryBinarization(t *testing.T) {
	// Test data has been extracted from table 9-35 of the specifications.
	tests := []struct {
		in   int
		want []int
		err  error
	}{
		{in: 0, want: []int{0}, err: nil},
		{in: 1, want: []int{1, 0}, err: nil},
		{in: 2, want: []int{1, 1, 0}, err: nil},
		{in: 3, want: []int{1, 1, 1, 0}, err: nil},
		{in: 4, want: []int{1, 1, 1, 1, 0}, err: nil},
		{in: 5, want: []int{1, 1, 1, 1, 1, 0}, err: nil},
		{in: -3, want: nil, err: errNegativeSyntaxVal},
	}

	for i, test := range tests {
		got, err := unaryBinarization(test.in)
		if err != test.err {
			t.Errorf("did not get expected error for test %d\nGot: %v\nWant: %v", i, err, test.err)
		}

		if !reflect.DeepEqual(test.want, got) {
			t.Errorf("did not get expected result for test %d\nGot: %v\nWant: %v\n", i, got, test.want)
		}
	}
}