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)
		}
	}
}

func TestTruncUnaryBinarization(t *testing.T) {
	tests := []struct {
		v    int
		cMax int
		want []int
		err  error
	}{
		{v: 0, cMax: 10, want: []int{0}, err: nil},
		{v: 1, cMax: 10, want: []int{1, 0}, err: nil},
		{v: 2, cMax: 10, want: []int{1, 1, 0}, err: nil},
		{v: 0, cMax: 0, want: []int{}, err: nil},
		{v: 4, cMax: 4, want: []int{1, 1, 1, 1}, err: nil},
		{v: 1, cMax: 10, want: []int{1, 0}, err: nil},
		{v: 2, cMax: 10, want: []int{1, 1, 0}, err: nil},
		{v: -3, cMax: 10, want: nil, err: errNegativeSyntaxVal},
		{v: 5, cMax: 4, want: nil, err: errInvalidSyntaxVal},
	}

	for i, test := range tests {
		got, err := truncUnaryBinarization(test.v, test.cMax)
		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)
		}
	}
}

func TestUEGkSuffix(t *testing.T) {
	// Data from https://patents.google.com/patent/US20070092150
	tests := []struct {
		v, uCoff, k   int
		signedValFlag bool
		want          []int
	}{
		0:  {v: 14, uCoff: 14, want: []int{0}},
		1:  {v: 15, uCoff: 14, want: []int{1, 0, 0}},
		2:  {v: 16, uCoff: 14, want: []int{1, 0, 1}},
		3:  {v: 17, uCoff: 14, want: []int{1, 1, 0, 0, 0}},
		4:  {v: 18, uCoff: 14, want: []int{1, 1, 0, 0, 1}},
		5:  {v: 19, uCoff: 14, want: []int{1, 1, 0, 1, 0}},
		6:  {v: 20, uCoff: 14, want: []int{1, 1, 0, 1, 1}},
		7:  {v: 21, uCoff: 14, want: []int{1, 1, 1, 0, 0, 0, 0}},
		8:  {v: 22, uCoff: 14, want: []int{1, 1, 1, 0, 0, 0, 1}},
		9:  {v: 23, uCoff: 14, want: []int{1, 1, 1, 0, 0, 1, 0}},
		10: {v: 24, uCoff: 14, want: []int{1, 1, 1, 0, 0, 1, 1}},
		11: {v: 25, uCoff: 14, want: []int{1, 1, 1, 0, 1, 0, 0}},
	}

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

func TestUnaryExpGolombBinarization(t *testing.T) {
	tests := []struct {
		v, uCoff, k   int
		signedValFlag bool
		want          []int
	}{
		0: {v: 7, uCoff: 14, want: []int{1, 1, 1, 1, 1, 1, 1, 0}},
		1: {v: 17, uCoff: 14, want: []int{1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0}},
		2: {v: 15, uCoff: 14, signedValFlag: true, want: []int{1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0}},
		3: {v: -15, uCoff: 14, signedValFlag: true, want: []int{1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1}},
	}

	for i, test := range tests {
		got, err := unaryExpGolombBinarization(test.v, test.uCoff, test.k, test.signedValFlag)
		if err != nil {
			t.Errorf("did not expect error %v for test %d", err, i)
		}

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