dataset.go

package jsonstat

import (
	"slices"

	"github.com/jsonstat/go/internal/stride"
)

// ClassDataset, ClassCollection, ClassDimension are the JSON-stat 2.0
// response class values. The pre-2.0 bundle has no class value; documents
// without a class are treated as [ClassDataset] by tolerant decoders (see
// [Decode]).
const (
	ClassDataset    = "dataset"
	ClassCollection = "collection"
	ClassDimension  = "dimension"
	// ClassBundle is the legacy pre-2.0 container. It is not a valid 2.0
	// class but is recognised by [Decode] for backwards compatibility.
	ClassBundle = "bundle"
)

// Cell is a single value-status pair read out of a dataset cube. It is the
// return type of the cell-access methods on [Dataset].
//
// A cell may be missing: dense value arrays encode missing observations as
// JSON null, and sparse value objects omit them. Missing cells have Value
// set to the zero float64 (0) and Missing set to true; callers should test
// Missing before reading Value.
type Cell struct {
	// Value is the observation's numeric value. Always 0 when Missing is
	// true.
	Value float64

	// Status is the observation-level status code. May be empty even when
	// the value is present.
	Status string

	// Missing reports whether the cell has no value (JSON null in a dense
	// array, or absent key in a sparse object).
	Missing bool

	// Flat is the row-major flat index of the cell within the dataset's
	// value array. Always ≥ 0.
	Flat int

	// Coords is the per-dimension category index of the cell, in
	// [Dataset.ID] order.
	Coords []int
}

// HasValue reports whether the cell carries a numeric observation.
func (c Cell) HasValue() bool { return !c.Missing }

// Dataset is the in-memory representation of a JSON-stat dataset response
// (class "dataset"). It is the primary type of this package: every read,
// traverse, subset, transform, and serve operation is a function over a
// *Dataset.
//
// Dataset normalises the polymorphic parts of the wire format:
//
//   - The "value" property can be a dense JSON array (with nulls) or a sparse
//     JSON object; both are represented as a [valueStore] with a uniform
//     float64 slice and a "missing" set.
//   - The "status" property can be a string, an array, or an object; it is
//     represented as a [statusStore] that resolves the effective status per
//     cell.
//
// On encode, Dataset re-emits the same form it was decoded from (or the form
// chosen by the Builder), preserving round-trip fidelity on canonical input.
type Dataset struct {
	// Class is always [ClassDataset] for a Dataset. Present so that a
	// *Dataset satisfies the shared [Response] contract.
	Class string

	// Version is the JSON-stat version declared on the source document, or
	// [FormatVersion] for documents built with the [Builder].
	Version string

	// Label is the dataset's human-readable title.
	Label string

	// Href is the dataset's canonical URL, if declared.
	Href string

	// Source is the dataset's source attribution text.
	Source string

	// Updated is the dataset's last-updated timestamp in ISO 8601 format.
	// Stored as the raw string to avoid timezone pitfalls; callers that need
	// a time.Time should parse it themselves.
	Updated string

	// ID is the ordered list of dimension IDs that define the cube's axes.
	// The order matters: it determines row-major flat-index math.
	ID []string

	// Size is the number of categories per dimension, in the same order as
	// ID.
	Size []int

	// Role maps declared roles (time, geo, metric) to the dimension IDs that
	// carry them. May be empty.
	Role map[Role][]string

	// Dimensions holds each dimension by ID. Order is preserved from ID.
	Dimensions map[string]*Dimension

	// Note is the dataset-level notes array. Empty when absent.
	Note []string

	// Link is the dataset's related-resource links.
	Link []Link

	// Extension collects provider-specific extra fields attached to the
	// dataset itself. Anything that does not map onto a typed field ends up
	// here.
	Extension map[string]any

	// values and status normalise the polymorphic wire forms. They are
	// populated by Decode and by the Builder, and read by every traversal
	// method. The backing slice in values is always length Total(size).
	values valueStore
	status statusStore

	// strides caches stride.Strides(Size) for fast flat-index math. nil until
	// [Dataset.reindex] is called.
	strides []int

	// total caches stride.Total(Size). 0 until [Dataset.reindex] is called.
	total int

	// dimPos maps dimension ID → its index in [Dataset.ID]. Used by traversal
	// methods to resolve IDs without a linear scan.
	dimPos map[string]int
}

// N returns the total number of cells in the cube (the product of Size),
// including missing ones. This is the Go counterpart of the JS Toolkit's `n`
// property. Returns 0 if the dataset has no dimensions or any dimension has
// size 0.
func (d *Dataset) N() int {
	if d == nil {
		return 0
	}
	return d.total
}

// DimCount returns the number of dimensions.
func (d *Dataset) DimCount() int {
	if d == nil {
		return 0
	}
	return len(d.ID)
}

// Dimension returns the dimension with the given ID, or [ErrDimensionNotFound]
// when no such dimension exists.
func (d *Dataset) Dimension(id string) (*Dimension, error) {
	if d == nil {
		return nil, ErrDimensionNotFound
	}
	dim, ok := d.Dimensions[id]
	if !ok {
		return nil, &ValueError{Op: "Dataset.Dimension", Dim: id, Err: ErrDimensionNotFound}
	}
	return dim, nil
}

// DimensionByIndex returns the dimension at position i in [Dataset.ID].
// Returns [ErrIndexOutOfRange] when i is out of range.
func (d *Dataset) DimensionByIndex(i int) (*Dimension, error) {
	if d == nil || i < 0 || i >= len(d.ID) {
		return nil, &ValueError{Op: "Dataset.DimensionByIndex", Flat: i, Err: ErrIndexOutOfRange}
	}
	return d.Dimensions[d.ID[i]], nil
}

// DimensionByRole returns the first dimension assigned the given role, or
// [ErrDimensionNotFound] when no dimension carries that role. For the (rare)
// case of multiple dimensions sharing a role, use [Dataset.DimensionsByRole].
func (d *Dataset) DimensionByRole(role Role) (*Dimension, error) {
	if d == nil {
		return nil, ErrDimensionNotFound
	}
	ids := d.Role[role]
	if len(ids) == 0 {
		return nil, &ValueError{Op: "Dataset.DimensionByRole", Err: ErrDimensionNotFound}
	}
	return d.Dimension(ids[0])
}

// DimensionsByRole returns every dimension assigned the given role, in the
// order they appear in [Dataset.ID].
func (d *Dataset) DimensionsByRole(role Role) []*Dimension {
	if d == nil {
		return nil
	}
	ids := d.Role[role]
	out := make([]*Dimension, 0, len(ids))
	for _, id := range ids {
		if dim, ok := d.Dimensions[id]; ok {
			out = append(out, dim)
		}
	}
	return out
}

// DimensionPosition returns the index in [Dataset.ID] of the given dimension
// ID, or [ErrDimensionNotFound] when absent.
func (d *Dataset) DimensionPosition(id string) (int, error) {
	if d == nil {
		return 0, ErrDimensionNotFound
	}
	pos, ok := d.dimPos[id]
	if !ok {
		return 0, &ValueError{Op: "Dataset.DimensionPosition", Dim: id, Err: ErrDimensionNotFound}
	}
	return pos, nil
}

// DimensionIDs returns a defensive copy of [Dataset.ID].
func (d *Dataset) DimensionIDs() []string {
	if d == nil {
		return nil
	}
	return slices.Clone(d.ID)
}

// SizeSlice returns a defensive copy of [Dataset.Size].
func (d *Dataset) SizeSlice() []int {
	if d == nil {
		return nil
	}
	return slices.Clone(d.Size)
}

// Roles returns the set of roles declared on the dataset.
func (d *Dataset) Roles() []Role {
	if d == nil || len(d.Role) == 0 {
		return nil
	}
	out := make([]Role, 0, len(d.Role))
	for _, r := range rolesOrder {
		if _, ok := d.Role[r]; ok {
			out = append(out, r)
		}
	}
	return out
}

// HasRole reports whether any dimension is assigned the given role.
func (d *Dataset) HasRole(role Role) bool {
	if d == nil {
		return false
	}
	_, ok := d.Role[role]
	return ok
}

// Strides returns the cached strides for this dataset. The returned slice is
// a defensive copy; callers may mutate it. Returns nil before [reindex] is
// called (which Decode and Builder do automatically).
func (d *Dataset) Strides() []int {
	if d == nil {
		return nil
	}
	return slices.Clone(d.strides)
}

// reindex recomputes derived state from ID/Size/Dimensions: strides, total,
// dimPos, and each dimension's category index map. It must be called after
// any structural mutation (decode, builder assembly, dice). Returns an error
// if Size is invalid.
func (d *Dataset) reindex() error {
	if d == nil {
		return ErrInvalidSize
	}
	total, err := stride.Total(d.Size)
	if err != nil {
		return &ValueError{Op: "Dataset.reindex", Err: err}
	}
	strides, err := stride.Strides(d.Size)
	if err != nil {
		return &ValueError{Op: "Dataset.reindex", Err: err}
	}
	d.total = total
	d.strides = strides
	d.dimPos = make(map[string]int, len(d.ID))
	for i, id := range d.ID {
		d.dimPos[id] = i
		if dim, ok := d.Dimensions[id]; ok {
			dim.buildIndex()
		}
	}
	// Ensure value/status stores are sized consistently with total.
	d.values.ensureCapacity(total)
	d.status.ensureCapacity(total)
	return nil
}

// ValueFormat reports which wire form the dataset's value store would emit on
// encode. See [valueStore].
func (d *Dataset) ValueFormat() string {
	if d == nil {
		return ""
	}
	if d.values.sparse {
		return "sparse"
	}
	return "dense"
}

// StatusFormat reports which wire form the dataset's status store would emit
// on encode.
func (d *Dataset) StatusFormat() string {
	if d == nil {
		return "none"
	}
	switch d.status.form {
	case statusUniform:
		return "uniform"
	case statusArray:
		return "array"
	case statusObject:
		return "object"
	default:
		return "none"
	}
}

// valueStore is the normalised representation of the polymorphic "value"
// property. Dense arrays decode into dense=true with all cells populated
// (missing cells kept track of via [missing]); sparse objects decode into
// sparse=true with only present cells in [sparse]. Both forms share the same
// flat row-major indexing scheme.
type valueStore struct {
	sparse  bool
	dense   []float64
	missing []bool // dense only; true at flat positions holding JSON null
	sparse_ map[int]float64
}

// ensureCapacity grows the dense slice to n entries if needed. It does not
// shrink. For sparse stores it is a no-op.
func (v *valueStore) ensureCapacity(n int) {
	if v == nil || v.sparse || n <= 0 {
		return
	}
	if cap(v.dense) < n {
		grown := make([]float64, len(v.dense), n)
		copy(grown, v.dense)
		v.dense = grown
	}
	if len(v.dense) < n {
		extra := make([]float64, n-len(v.dense))
		v.dense = append(v.dense, extra...)
	}
	if len(v.missing) < n {
		extra := make([]bool, n-len(v.missing))
		v.missing = append(v.missing, extra...)
	}
}

// at returns the value at flat index i and whether it is present.
func (v *valueStore) at(i int) (float64, bool) {
	if v == nil {
		return 0, false
	}
	if v.sparse {
		val, ok := v.sparse_[i]
		return val, ok
	}
	if i < 0 || i >= len(v.dense) {
		return 0, false
	}
	if len(v.missing) > 0 && v.missing[i] {
		return 0, false
	}
	return v.dense[i], true
}

// statusForm names the wire form of a statusStore.
type statusForm byte

const (
	statusNone    statusForm = iota // no status property at all
	statusUniform                   // "status": "e"
	statusArray                     // "status": ["a","b",...]
	statusObject                    // "status": {"5":"p"}
)

// statusStore normalises the three "status" wire forms. Resolution per cell
// follows the JSON-stat rules:
//   - uniform: every cell shares [uniform].
//   - array: status is per-flat-index via [perCell].
//   - object: status is per-flat-index via [perCell] (sparse — absent keys
//     mean no status).
type statusStore struct {
	form    statusForm
	uniform string
	perCell map[int]string
}

// ensureCapacity is a no-op placeholder kept for symmetry with valueStore.
// status stores are map-backed and grow on demand.
func (s *statusStore) ensureCapacity(_ int) {}

// at returns the effective status code at flat index i. Returns "" when the
// dataset has no status for that cell.
func (s *statusStore) at(i int) string {
	if s == nil {
		return ""
	}
	switch s.form {
	case statusUniform:
		return s.uniform
	case statusArray, statusObject:
		return s.perCell[i]
	default:
		return ""
	}
}

// hasStatus reports whether the dataset declared any status at all.
func (s *statusStore) hasStatus() bool {
	return s != nil && s.form != statusNone
}

// flatToString formats a flat index for sparse-object keys. JSON-stat's