package gfs

import "fmt"

// Variant describes one configuration of a NOAA dataset family (GFS or GEFS).
//
// The dataset cube is a 5-D float32 array with shape
// (NumHours, NumLevels, NumVariables, NumLatitudes, NumLongitudes) where
// NumVariables and ElementSize are fixed across all GFS variants but the
// other dimensions depend on the resolution and forecast cadence.
type Variant struct {
	// ID is a stable identifier ("gfs-0p50-3h", "gefs-0p50-3h", ...).
	ID string
	// Family identifies the dataset family the variant belongs to.
	Family Family

	// Resolution token used in NOAA URLs ("0p50", "0p25").
	ResToken string
	// Grid step in degrees (0.5, 0.25). 180 / Resolution + 1 latitudes and
	// 360 / Resolution longitudes.
	Resolution float64

	HourStep int // hours between forecast steps
	MaxHour  int // largest forecast hour (inclusive)

	// Pressures lists every pressure level in dataset index order, descending.
	Pressures []int
	// PressuresPgrb2 / PressuresPgrb2b split the pressures between the two
	// downloaded GRIB files. Their union must equal Pressures.
	PressuresPgrb2  []int
	PressuresPgrb2b []int

	pressureIndex    map[int]int
	pressureLevelSet map[int]LevelSet
}

// NumHours returns MaxHour/HourStep + 1.
func (v *Variant) NumHours() int { return v.MaxHour/v.HourStep + 1 }

// NumLevels returns len(Pressures).
func (v *Variant) NumLevels() int { return len(v.Pressures) }

// NumLatitudes returns 180/Resolution + 1.
func (v *Variant) NumLatitudes() int { return int(180.0/v.Resolution) + 1 }

// NumLongitudes returns 360/Resolution.
func (v *Variant) NumLongitudes() int { return int(360.0 / v.Resolution) }

// DatasetSize returns the canonical file size in bytes.
func (v *Variant) DatasetSize() int64 {
	return int64(v.NumHours()) * int64(v.NumLevels()) * int64(NumVariables) *
		int64(v.NumLatitudes()) * int64(v.NumLongitudes()) * int64(ElementSize)
}

// Hours returns the full list of forecast hours [0, HourStep, ..., MaxHour].
func (v *Variant) Hours() []int {
	out := make([]int, 0, v.NumHours())
	for h := 0; h <= v.MaxHour; h += v.HourStep {
		out = append(out, h)
	}
	return out
}

// HourIndex returns the dataset time index for an hour, or -1 if invalid.
func (v *Variant) HourIndex(hour int) int {
	if hour < 0 || hour > v.MaxHour || hour%v.HourStep != 0 {
		return -1
	}
	return hour / v.HourStep
}

// PressureIndex returns the dataset index for a pressure level in hPa,
// or -1 when the level is unknown to this variant.
func (v *Variant) PressureIndex(hPa int) int {
	v.indexLazyInit()
	if i, ok := v.pressureIndex[hPa]; ok {
		return i
	}
	return -1
}

// PressureLevelSet returns the GRIB file set carrying a pressure level.
func (v *Variant) PressureLevelSet(hPa int) (LevelSet, bool) {
	v.indexLazyInit()
	ls, ok := v.pressureLevelSet[hPa]
	return ls, ok
}

// VariableIndex maps a GRIB (category, number) pair to a dataset variable index.
func (v *Variant) VariableIndex(parameterCategory, parameterNumber int) int {
	switch {
	case parameterCategory == 3 && parameterNumber == 5:
		return VarHeight
	case parameterCategory == 2 && parameterNumber == 2:
		return VarWindU
	case parameterCategory == 2 && parameterNumber == 3:
		return VarWindV
	default:
		return -1
	}
}

// GribURL returns the S3 URL for the primary (pgrb2) GRIB file.
func (v *Variant) GribURL(date string, runHour, forecastStep int) string {
	return fmt.Sprintf("%s/gfs.%s/%02d/atmos/gfs.t%02dz.pgrb2.%s.f%03d",
		S3BaseURL, date, runHour, runHour, v.ResToken, forecastStep)
}

// GribURLB returns the S3 URL for the secondary (pgrb2b) GRIB file.
func (v *Variant) GribURLB(date string, runHour, forecastStep int) string {
	return fmt.Sprintf("%s/gfs.%s/%02d/atmos/gfs.t%02dz.pgrb2b.%s.f%03d",
		S3BaseURL, date, runHour, runHour, v.ResToken, forecastStep)
}

func (v *Variant) indexLazyInit() {
	if v.pressureIndex != nil {
		return
	}
	v.pressureIndex = make(map[int]int, len(v.Pressures))
	for i, p := range v.Pressures {
		v.pressureIndex[p] = i
	}
	v.pressureLevelSet = make(map[int]LevelSet, len(v.Pressures))
	for _, p := range v.PressuresPgrb2 {
		v.pressureLevelSet[p] = LevelSetA
	}
	for _, p := range v.PressuresPgrb2b {
		v.pressureLevelSet[p] = LevelSetB
	}
}

// Standard variants -- these mirror what NOAA publishes today.
//
// GFS0p50_3h is the historical Tawhiri default: 0.5° resolution, 3-hour
// forecast cadence, 0..192h horizon, 47 pressure levels split across the
// primary and secondary GRIB files.
//
// GFS0p25_3h mirrors the same 3-hour cadence at 0.25° resolution (the
// horizon is larger in practice but we keep 192h for parity with 0p50).
//
// GFS0p25_1h targets the 1-hourly portion NOAA publishes out to 120h.
var (
	GFS0p50_3h = &Variant{
		ID:              "gfs-0p50-3h",
		ResToken:        "0p50",
		Resolution:      0.5,
		HourStep:        3,
		MaxHour:         192,
		Pressures:       []int{1000, 975, 950, 925, 900, 875, 850, 825, 800, 775, 750, 725, 700, 675, 650, 625, 600, 575, 550, 525, 500, 475, 450, 425, 400, 375, 350, 325, 300, 275, 250, 225, 200, 175, 150, 125, 100, 70, 50, 30, 20, 10, 7, 5, 3, 2, 1},
		PressuresPgrb2:  []int{10, 20, 30, 50, 70, 100, 150, 200, 250, 300, 350, 400, 450, 500, 550, 600, 650, 700, 750, 800, 850, 900, 925, 950, 975, 1000},
		PressuresPgrb2b: []int{1, 2, 3, 5, 7, 125, 175, 225, 275, 325, 375, 425, 475, 525, 575, 625, 675, 725, 775, 825, 875},
	}

	GFS0p25_3h = &Variant{
		ID:              "gfs-0p25-3h",
		ResToken:        "0p25",
		Resolution:      0.25,
		HourStep:        3,
		MaxHour:         192,
		Pressures:       GFS0p50_3h.Pressures,
		PressuresPgrb2:  GFS0p50_3h.PressuresPgrb2,
		PressuresPgrb2b: GFS0p50_3h.PressuresPgrb2b,
	}

	GFS0p25_1h = &Variant{
		ID:              "gfs-0p25-1h",
		ResToken:        "0p25",
		Resolution:      0.25,
		HourStep:        1,
		MaxHour:         120,
		Pressures:       GFS0p50_3h.Pressures,
		PressuresPgrb2:  GFS0p50_3h.PressuresPgrb2,
		PressuresPgrb2b: GFS0p50_3h.PressuresPgrb2b,
	}
)

// VariantByID returns one of the predefined variants by its ID.
func VariantByID(id string) (*Variant, error) {
	switch id {
	case GFS0p50_3h.ID:
		return GFS0p50_3h, nil
	case GFS0p25_3h.ID:
		return GFS0p25_3h, nil
	case GFS0p25_1h.ID:
		return GFS0p25_1h, nil
	case GEFS0p50_3h.ID:
		return GEFS0p50_3h, nil
	default:
		return nil, fmt.Errorf("unknown variant %q", id)
	}
}