package stats

import (
	"math"
	"sort"

	"salty/terrain/internal/field"
)

// Gathering a world's statistics one piece at a time.
//
// The geology is solved one landmass at a time (D-53) and a planet's regions never exist together, so a
// planet-wide statistic has to be assembled rather than computed. Terrain.md's rule for that is "statistics
// pool across regions rather than being computed per region and averaged", and until now it was a rule with
// no implementation: the whole package took a grid and sorted it, so a planet bake printed its elevation
// range and nothing else - no slope distribution, no per-uplift-class breakdown, no drainage density. The
// block the documentation calls the one that matters most was the one that could not be afforded.
//
// An Accumulator is what makes the rule true. Every quantity in it is either a counter, an exact running
// extreme, or a Histogram, and all three are **additive**: merging two regions and reading the result gives
// exactly what one pass over both would have. See histogram.go for why that is the whole design and not an
// implementation detail.
//
// Add takes a grid. It does not care whether that grid is one region of a planet or the whole square canvas,
// which is the other half of the point: `generate` and `bake` now compute their statistics with the same
// code, so a number measured on one is comparable with the same number measured on the other.

// Options are the constants a world is judged against. They have to be the same for every region of a planet,
// which is why they live on the accumulator rather than being passed to each Add.
type Options struct {
	// ElevMin and ElevMax bound the elevation histogram, and they are the manifest's encoding range on
	// purpose rather than the data's own extremes. A histogram's bounds have to be known before the first
	// value arrives, or two regions would bin against different scales and could not be merged - and the
	// encoding range is the one bound that is a property of the world rather than of whatever happens to be
	// in front of it. Anything outside is counted as out of range, which is also what the clip fraction is
	// about.
	ElevMin, ElevMax float64

	TalusDeg      float64 // the angle of repose, for the "pinned against the clamp" share
	ReliefWindowM float64 // the side of the square local relief is taken over
	ChannelM2     float64 // drainage area at which a cell counts as a channel
	K, M, N       float64 // the stream-power constants, for the slope-area normalisation
}

// slopeBins and elevBins are the resolutions. A twentieth of a degree and a metre or two of elevation are far
// finer than any verdict in Summary turns on, and the whole structure is a few tens of kilobytes either way.
const (
	slopeBins = 2048
	elevBins  = 4096
	logSABins = 1024
)

// bucketAcc is one uplift class's share of the accumulator.
type bucketAcc struct {
	slope, relief, elev *Histogram
	near, total         int64
}

// saBin is one decade-fraction of drainage area in the slope-area plot.
type saBin struct{ norm, raw *Histogram }

// Accumulator gathers one world's statistics, a grid at a time.
type Accumulator struct {
	opt Options

	Cells, Land, Clip int64
	MinM, MaxM        float64 // the whole field, sea floor included: what the 16-bit encoding has to hold

	elev  *Histogram // land only
	slope *Histogram // land only, degrees

	buckets []bucketAcc
	sa      map[int]*saBin

	saChannels   int64
	channelCells int64
	leafCells    int64
}

// New returns an empty accumulator.
func New(opt Options) *Accumulator {
	if opt.ElevMax <= opt.ElevMin {
		opt.ElevMin, opt.ElevMax = -1024, 2048
	}
	a := &Accumulator{
		opt:  opt,
		MinM: math.Inf(1), MaxM: math.Inf(-1),
		elev:  NewHistogram(opt.ElevMin, opt.ElevMax, elevBins),
		slope: NewHistogram(0, 90, slopeBins),
		sa:    map[int]*saBin{},
	}
	span := opt.ElevMax - opt.ElevMin
	a.buckets = make([]bucketAcc, len(bucketDefs))
	for i := range a.buckets {
		a.buckets[i] = bucketAcc{
			slope:  NewHistogram(0, 90, slopeBins),
			relief: NewHistogram(0, span, elevBins),
			elev:   NewHistogram(opt.ElevMin, opt.ElevMax, elevBins),
		}
	}
	return a
}

// Input is one grid and everything known about it. Everything but H and Land is optional; a caller with no
// flow topology gets the statistics that do not need one.
type Input struct {
	H    *field.Field
	Land []bool // nil means every cell is land

	// WrapX says whether this grid's left and right edges are the same meridian. A region of a planet is a
	// rectangle cut out of the cylinder with water all round it, so it does *not* wrap; the whole square
	// canvas does not either. It is here because the local relief window is the one thing that reads
	// neighbours, and being wrong about it would put a seam in one column of the relief map.
	WrapX bool

	UpliftMYr []float32 // per cell; without it there is no per-class breakdown
	KLocal    []float32 // the lithology multiplier, for the slope-area normalisation

	// The flow topology, for slope-area and drainage density. All three or none.
	Area     []float32
	Receiver []int32
	Length   []float32
}

// AddExtent records what a *finished* grid covers: how many cells, how many of them are land, how many fall
// outside the encoding range, and the extremes over everything including the sea floor.
//
// It is separate from Add because on a planet the two are measured in different places, and measuring them in
// the wrong one is silently wrong rather than obviously so. A region is a rectangle cut out of the cylinder
// with an ocean margin round it, and neighbouring regions' margins overlap - so pooling "cells" across regions
// counts the same water more than once and reports a land fraction that means nothing. The extent is a
// property of the composited planet and is measured once, on it. Land statistics are the opposite: they are
// per landmass, disjoint by construction, and never see the finished cylinder at all.
func (a *Accumulator) AddExtent(data []float32, land []bool, clipCells int64) {
	a.Clip += clipCells
	for i, v := range data {
		a.Cells++
		f := float64(v)
		if f < a.MinM {
			a.MinM = f
		}
		if f > a.MaxM {
			a.MaxM = f
		}
		if land == nil || land[i] {
			a.Land++
		}
	}
}

// Add folds one grid's land statistics in. It reads only the cells the mask calls land, and it deliberately
// records nothing about the grid's extent - see AddExtent.
func (a *Accumulator) Add(in Input) {
	h := in.H
	if h == nil || len(h.Data) == 0 {
		return
	}

	// Local relief first, because it is the one quantity that needs a neighbourhood and therefore a whole
	// field of its own. Two sliding passes, O(1) a cell whatever the window: the loop this replaces was
	// 1.1e11 comparisons on a planet, which is why no planet bake has ever printed this block.
	var relief *field.Field
	if a.opt.ReliefWindowM > 0 && in.UpliftMYr != nil {
		r := int(math.Round(a.opt.ReliefWindowM / h.CellM / 2))
		if r < 1 {
			r = 1
		}
		relief = field.LocalRelief(h, r, in.WrapX)
	}

	inv := 1.0 / (2.0 * h.CellM)
	cellArea := h.CellM * h.CellM
	for y := 0; y < h.H; y++ {
		for x := 0; x < h.W; x++ {
			i := y*h.W + x
			if in.Land != nil && !in.Land[i] {
				continue
			}
			v := float64(h.Data[i])
			a.elev.Add(v)

			// The slope inline rather than through Field.Slope: that allocates a whole field, which at
			// planet scale is 300 MB per call and there would be two of them.
			gx := float64(h.AtClamped(x+1, y)-h.AtClamped(x-1, y)) * inv
			gy := float64(h.AtClamped(x, y+1)-h.AtClamped(x, y-1)) * inv
			deg := math.Atan(math.Hypot(gx, gy)) * 180 / math.Pi
			a.slope.Add(deg)

			if in.UpliftMYr != nil {
				if b := bucketOf(float64(in.UpliftMYr[i]) * 1000); b >= 0 {
					acc := &a.buckets[b]
					acc.total++
					acc.slope.Add(deg)
					acc.elev.Add(v)
					if relief != nil {
						acc.relief.Add(float64(relief.Data[i]))
					}
					if deg >= a.opt.TalusDeg-2 { // pinned against the clamp rather than shaped by erosion
						acc.near++
					}
				}
			}

			if in.Area == nil {
				continue
			}
			if float64(in.Area[i]) >= a.opt.ChannelM2 {
				a.channelCells++
			}
			// A leaf is a cell that drains nothing but itself, and what it measures is the router when the
			// ground is smooth - not the landscape when it is finished. On a planar ramp with no erosion at
			// all, D8 leaves 29.5 % of the grid draining nothing, because a cell either sits on one of its
			// parallel flow lines or it does not; multiple-flow leaves 0.4 %, which is the strict local
			// maxima. After three hundred steps of solving the same ramp both come back near 8 %: the
			// terrain has dissected itself by then and its own divides dominate the count. So read this on
			// young ground, on a stage dump, or against another run of the same age, and do not read it as a
			// verdict on a mature one. It is a count, so it pools across regions exactly.
			if float64(in.Area[i]) <= cellArea*1.001 {
				a.leafCells++
			}
			a.addSlopeArea(in, i, h)
		}
	}
}

// addSlopeArea records one channel cell in the slope-area plot.
//
// S is the gradient *along the flow path*, not the magnitude of the topographic gradient: on a valley floor
// the central difference is dominated by the walls across the channel, which reads far steeper than the water
// actually runs and bends the fitted exponent well past -m/n. And the slope is normalised by (U/K)^(1/n) with
// the *local* K, because erodibility correlates with drainage area by construction - soft rock is cut down,
// sits low and collects flow - so one global K mis-corrects the large-A end systematically.
func (a *Accumulator) addSlopeArea(in Input, i int, h *field.Field) {
	if in.Receiver == nil || in.Length == nil || in.UpliftMYr == nil {
		return
	}
	r := in.Receiver[i]
	if int(r) == i { // a root drains to itself and has no gradient to measure
		return
	}
	area := float64(in.Area[i])
	s := float64(h.Data[i]-h.Data[r]) / float64(in.Length[i])
	if area < a.opt.ChannelM2 || s <= 1e-6 {
		return
	}
	u := float64(in.UpliftMYr[i])
	kk := a.opt.K
	if in.KLocal != nil {
		kk *= float64(in.KLocal[i])
	}
	if u <= 0 || kk <= 0 || a.opt.N <= 0 {
		return // no steady state to normalise against
	}
	a.saChannels++
	key := int(math.Floor(math.Log10(area) * binsPerDecade))
	b := a.sa[key]
	if b == nil {
		b = &saBin{norm: NewHistogram(-8, 4, logSABins), raw: NewHistogram(-8, 4, logSABins)}
		a.sa[key] = b
	}
	b.norm.Add(math.Log10(s / math.Pow(u/kk, 1/a.opt.N)))
	b.raw.Add(math.Log10(s))
}

const binsPerDecade = 4

// Merge folds another accumulator in. Every field is additive by construction; see histogram.go.
func (a *Accumulator) Merge(o *Accumulator) {
	if o == nil {
		return
	}
	a.Cells += o.Cells
	a.Land += o.Land
	a.Clip += o.Clip
	a.saChannels += o.saChannels
	a.channelCells += o.channelCells
	a.leafCells += o.leafCells
	a.MinM = math.Min(a.MinM, o.MinM)
	a.MaxM = math.Max(a.MaxM, o.MaxM)
	a.elev.Merge(o.elev)
	a.slope.Merge(o.slope)
	for i := range a.buckets {
		if i >= len(o.buckets) {
			break
		}
		a.buckets[i].slope.Merge(o.buckets[i].slope)
		a.buckets[i].relief.Merge(o.buckets[i].relief)
		a.buckets[i].elev.Merge(o.buckets[i].elev)
		a.buckets[i].near += o.buckets[i].near
		a.buckets[i].total += o.buckets[i].total
	}
	// Sorted, because Go randomises map iteration and cross-cutting rule 12 says the answer must not depend
	// on it. Here it would only change the order two float sums happen in, which is exactly the sort of "it
	// does not matter this time" the rule exists to refuse.
	keys := make([]int, 0, len(o.sa))
	for k := range o.sa {
		keys = append(keys, k)
	}
	sort.Ints(keys)
	for _, k := range keys {
		b := a.sa[k]
		if b == nil {
			b = &saBin{norm: NewHistogram(-8, 4, logSABins), raw: NewHistogram(-8, 4, logSABins)}
			a.sa[k] = b
		}
		b.norm.Merge(o.sa[k].norm)
		b.raw.Merge(o.sa[k].raw)
	}
}

// bucketDefs are the uplift classes the breakdown splits on.
//
// Absolute rather than percentiles of this map's own field: the point is to compare one run against the next,
// and a percentile split would redefine "plain" every time the uplift field was retuned. They are reporting
// buckets and not a description of terrain - 0.1 mm/yr is a fourteen-degree hillslope at an 8 m cell, which
// is hill country wherever it is painted, and reading this axis as guidance is how a legend once ended up ten
// times too hot (D-55).
var bucketDefs = []struct {
	name   string
	lo, hi float64
}{
	{"plain", 0, 0.1},
	{"rolling", 0.1, 0.5},
	// The top bound is finite rather than +Inf only because the report is marshalled to meta.json and
	// encoding/json refuses an infinity. 100 mm/yr is an order of magnitude above anything on Earth.
	{"mountain", 0.5, 100},
}

func bucketOf(mmYr float64) int {
	for i, d := range bucketDefs {
		if mmYr >= d.lo && mmYr < d.hi {
			return i
		}
	}
	return -1
}

// Report turns everything gathered into the numbers a run is judged by.
func (a *Accumulator) Report(cellM float64) Report {
	r := Report{
		MinM: a.MinM, MaxM: a.MaxM, ReliefM: a.MaxM - a.MinM,
	}
	if a.Cells > 0 {
		r.LandFraction = float64(a.Land) / float64(a.Cells)
		r.ClipFraction = float64(a.Clip) / float64(a.Cells)
	}
	r.LandCells = a.Land
	r.MeasuredLandCells = a.elev.Count
	if a.elev.Count == 0 {
		return r
	}
	r.LandMinM, r.LandMaxM = a.elev.MinV, a.elev.MaxV
	r.LandReliefM = r.LandMaxM - r.LandMinM

	r.Slopes = Slopes{
		Under15Deg: a.slope.FracBelow(15),
		Under30Deg: a.slope.FracBelow(30),
		Over50Deg:  1 - a.slope.FracBelow(50),
		MedianDeg:  a.slope.Quantile(0.5),
	}

	// The hypsometric integral is a *mean* of the normalised elevation, so it comes off the exact running sum
	// rather than out of the bins: (sum - n*lo) / (n*span). The curve is the binned part, which is what it
	// should be - it is eleven fractions and nobody reads the third decimal of one.
	if span := r.LandMaxM - r.LandMinM; span > 1e-6 {
		r.Hypsometry.Integral = (a.elev.Sum - float64(a.elev.Count)*r.LandMinM) /
			(float64(a.elev.Count) * span)
		curve := make([]float64, 11)
		for i := 0; i <= 10; i++ {
			curve[i] = 1 - a.elev.FracBelow(r.LandMinM+span*float64(i)/10)
		}
		r.Hypsometry.Curve = curve
	}

	// Channel length over the area the channels were *counted* in, which is the land Add walked and not the
	// land the planet has. On a full bake the two are the same number. On a partial one - `bake --only` - the
	// extent is still the whole cylinder while the land statistics cover three islands, and dividing one by
	// the other would report a drainage density an order of magnitude low with nothing to say it had.
	if measured := a.elev.Count; measured > 0 && (a.channelCells > 0 || a.saChannels > 0) {
		lengthKm := float64(a.channelCells) * cellM / 1000
		areaKm2 := float64(measured) * cellM * cellM / 1e6
		if areaKm2 > 0 {
			r.DrainageDensity = lengthKm / areaKm2
		}
	}
	if measured := a.elev.Count; measured > 0 {
		r.LeafFraction = float64(a.leafCells) / float64(measured)
	}
	r.SlopeArea = a.slopeArea()
	r.Buckets = a.bucketReport(cellM)
	return r
}

func (a *Accumulator) slopeArea() SlopeArea {
	out := SlopeArea{Expected: expectedGradient, Channels: int(a.saChannels),
		ThreshKm2: a.opt.ChannelM2 / 1e6}
	keys := make([]int, 0, len(a.sa))
	for k := range a.sa {
		keys = append(keys, k)
	}
	sort.Ints(keys)

	var xs, normYs, rawYs []float64
	for _, key := range keys {
		b := a.sa[key]
		if b.norm.Count < 8 { // a bin with a handful of cells is noise, not a data point
			continue
		}
		logA := (float64(key) + 0.5) / binsPerDecade
		med := b.norm.Quantile(0.5)
		out.Bins = append(out.Bins, Bin{LogA: logA, LogS: med, N: int(b.norm.Count)})
		xs = append(xs, logA)
		normYs = append(normYs, med)
		rawYs = append(rawYs, b.raw.Quantile(0.5))
	}
	out.Exponent, out.R2 = fitLine(xs, normYs)
	out.RawExponent, out.RawR2 = fitLine(xs, rawYs)
	return out
}

func (a *Accumulator) bucketReport(cellM float64) []UpliftBucket {
	total := int64(0)
	for i := range a.buckets {
		total += a.buckets[i].total
	}
	if total == 0 {
		return nil
	}
	out := make([]UpliftBucket, 0, len(bucketDefs))
	for i, d := range bucketDefs {
		b := &a.buckets[i]
		if b.total == 0 {
			continue
		}
		out = append(out, UpliftBucket{
			Name: d.name, LoMmYr: d.lo, HiMmYr: d.hi,
			LandFrac:    float64(b.total) / float64(total),
			MedianDeg:   b.slope.Quantile(0.5),
			P90Deg:      b.slope.Quantile(0.9),
			MedianRelM:  b.relief.Quantile(0.5),
			WindowM:     a.opt.ReliefWindowM,
			NearTalus:   float64(b.near) / float64(b.total),
			MedianElevM: b.elev.Quantile(0.5),
			Cells:       int(b.total),
		})
	}
	return out
}