package dataset

import (
	"testing"
)

func TestDatasetShape(t *testing.T) {
	if NumHours != 65 {
		t.Errorf("NumHours = %d, want 65", NumHours)
	}
	if NumLevels != 47 {
		t.Errorf("NumLevels = %d, want 47", NumLevels)
	}
	if NumVariables != 3 {
		t.Errorf("NumVariables = %d, want 3", NumVariables)
	}
	if NumLatitudes != 361 {
		t.Errorf("NumLatitudes = %d, want 361", NumLatitudes)
	}
	if NumLongitudes != 720 {
		t.Errorf("NumLongitudes = %d, want 720", NumLongitudes)
	}
}

func TestDatasetSize(t *testing.T) {
	// 65 * 47 * 3 * 361 * 720 * 4 = 9,528,667,200
	want := int64(9_528_667_200)
	if DatasetSize != want {
		t.Errorf("DatasetSize = %d, want %d", DatasetSize, want)
	}
}

func TestPressureLevels(t *testing.T) {
	if len(Pressures) != 47 {
		t.Fatalf("len(Pressures) = %d, want 47", len(Pressures))
	}

	// First should be 1000 (highest pressure, near surface)
	if Pressures[0] != 1000 {
		t.Errorf("Pressures[0] = %d, want 1000", Pressures[0])
	}
	// Last should be 1 (lowest pressure, high atmosphere)
	if Pressures[46] != 1 {
		t.Errorf("Pressures[46] = %d, want 1", Pressures[46])
	}

	// Should be sorted descending
	for i := 1; i < len(Pressures); i++ {
		if Pressures[i] >= Pressures[i-1] {
			t.Errorf("Pressures not descending at [%d]: %d >= %d", i, Pressures[i], Pressures[i-1])
		}
	}

	// Total levels: 26 from pgrb2 + 21 from pgrb2b = 47
	if len(PressuresPgrb2) != 26 {
		t.Errorf("len(PressuresPgrb2) = %d, want 26", len(PressuresPgrb2))
	}
	if len(PressuresPgrb2b) != 21 {
		t.Errorf("len(PressuresPgrb2b) = %d, want 21", len(PressuresPgrb2b))
	}
}

func TestPressureIndex(t *testing.T) {
	if PressureIndex(1000) != 0 {
		t.Errorf("PressureIndex(1000) = %d, want 0", PressureIndex(1000))
	}
	if PressureIndex(1) != 46 {
		t.Errorf("PressureIndex(1) = %d, want 46", PressureIndex(1))
	}
	if PressureIndex(500) != 20 {
		t.Errorf("PressureIndex(500) = %d, want 20", PressureIndex(500))
	}
	if PressureIndex(9999) != -1 {
		t.Errorf("PressureIndex(9999) = %d, want -1", PressureIndex(9999))
	}
}

func TestPressureLevelSet(t *testing.T) {
	// 1000 mb should be in pgrb2 (A)
	ls, ok := PressureLevelSet(1000)
	if !ok || ls != LevelSetA {
		t.Errorf("PressureLevelSet(1000) = %v, %v; want A, true", ls, ok)
	}

	// 125 mb should be in pgrb2b (B)
	ls, ok = PressureLevelSet(125)
	if !ok || ls != LevelSetB {
		t.Errorf("PressureLevelSet(125) = %v, %v; want B, true", ls, ok)
	}

	// 1, 2, 3, 5, 7 should be in pgrb2b (B)
	for _, p := range []int{1, 2, 3, 5, 7} {
		ls, ok := PressureLevelSet(p)
		if !ok || ls != LevelSetB {
			t.Errorf("PressureLevelSet(%d) = %v, %v; want B, true", p, ls, ok)
		}
	}

	// Every pressure level should have a level set assignment
	for _, p := range Pressures {
		_, ok := PressureLevelSet(p)
		if !ok {
			t.Errorf("PressureLevelSet(%d) not found", p)
		}
	}
}

func TestHourIndex(t *testing.T) {
	if HourIndex(0) != 0 {
		t.Errorf("HourIndex(0) = %d, want 0", HourIndex(0))
	}
	if HourIndex(3) != 1 {
		t.Errorf("HourIndex(3) = %d, want 1", HourIndex(3))
	}
	if HourIndex(192) != 64 {
		t.Errorf("HourIndex(192) = %d, want 64", HourIndex(192))
	}
	if HourIndex(1) != -1 {
		t.Errorf("HourIndex(1) = %d, want -1 (not multiple of 3)", HourIndex(1))
	}
	if HourIndex(195) != -1 {
		t.Errorf("HourIndex(195) = %d, want -1 (out of range)", HourIndex(195))
	}
}

func TestHours(t *testing.T) {
	hours := Hours()
	if len(hours) != NumHours {
		t.Fatalf("len(Hours()) = %d, want %d", len(hours), NumHours)
	}
	if hours[0] != 0 {
		t.Errorf("Hours()[0] = %d, want 0", hours[0])
	}
	if hours[len(hours)-1] != MaxHour {
		t.Errorf("Hours()[last] = %d, want %d", hours[len(hours)-1], MaxHour)
	}
}

func TestVariableIndex(t *testing.T) {
	if VariableIndex(3, 5) != VarHeight {
		t.Errorf("HGT: got %d, want %d", VariableIndex(3, 5), VarHeight)
	}
	if VariableIndex(2, 2) != VarWindU {
		t.Errorf("UGRD: got %d, want %d", VariableIndex(2, 2), VarWindU)
	}
	if VariableIndex(2, 3) != VarWindV {
		t.Errorf("VGRD: got %d, want %d", VariableIndex(2, 3), VarWindV)
	}
	if VariableIndex(0, 0) != -1 {
		t.Errorf("unknown: got %d, want -1", VariableIndex(0, 0))
	}
}