// Port of jsruntime/uikit/USHeatmap.tsx — a choropleth of the 50 states + DC, plus
// optional lat/lng markers, drawn by the WebAssembly. The state boundaries (usstates.go)
// are pre-projected with d3's albersUsa; the SAME projection is reimplemented below so a
// lat/lng point lands on top of the states — a faithful port validated to 0px against
// d3-geo (Alaska and Hawaii insets included).
//
// The map scales by viewBox (there is no axis text to keep crisp), so no measurement.
// Hover reads e.Target()'s data-attribute — the state path or point marker under the
// pointer — so no coordinate maths is needed to know what is being pointed at; the tooltip
// anchors to the hovered state's centroid (or the point), drawn into the same SVG.
package webui

import (
	"math"
	"sort"
	"strconv"
	"strings"

	"kjol/wasmruntime"
	"kjol/wasmruntime/vdom"
)

// ── the albersUsa projection (ported from d3-geo; scale 1280, translate [480,300]) ──────

const (
	usRad = math.Pi / 180
	usTau = 2 * math.Pi
	usK   = 1280.0
	usTX  = 480.0
	usTY  = 300.0
	usEps = 1e-6
)

func conicEqualAreaRaw(y0, y1 float64) func(lambda, phi float64) (float64, float64) {
	sy0 := math.Sin(y0)
	n := (sy0 + math.Sin(y1)) / 2
	c := 1 + sy0*(2*n-sy0)
	r0 := math.Sqrt(c) / n
	return func(lambda, phi float64) (float64, float64) {
		r := math.Sqrt(c-2*n*math.Sin(phi)) / n
		return r * math.Sin(lambda*n), r0 - r*math.Cos(lambda*n)
	}
}

// albersLobe is one conic-equal-area lobe. center is in the rotated frame (near 0° lon),
// so it is not re-rotated; the input point is rotated by rotateLon before projecting.
func albersLobe(rotateLon, centerLon, centerLat, p0, p1, scale, tx, ty float64) func(lon, lat float64) (float64, float64) {
	raw := conicEqualAreaRaw(p0*usRad, p1*usRad)
	rot := func(lon float64) float64 {
		l := (lon + rotateLon) * usRad
		return math.Mod(math.Mod(l+math.Pi, usTau)+usTau, usTau) - math.Pi
	}
	cx, cy := raw(centerLon*usRad, centerLat*usRad)
	return func(lon, lat float64) (float64, float64) {
		x, y := raw(rot(lon), lat*usRad)
		return tx + scale*(x-cx), ty - scale*(y-cy)
	}
}

var (
	usLower48 = albersLobe(96, -0.6, 38.7, 29.5, 45.5, usK, usTX, usTY)
	usAlaska  = albersLobe(154, -2, 58.5, 55, 65, usK*0.35, usTX-0.307*usK, usTY+0.201*usK)
	usHawaii  = albersLobe(157, -3, 19.9, 8, 18, usK, usTX-0.205*usK, usTY+0.212*usK)
)

func usInBox(x, y, x0, y0, x1, y1 float64) bool {
	return x >= x0 && x <= x1 && y >= y0 && y <= y1
}

// ProjectUS maps [lng, lat] to the 960×600 map, picking the lower-48 / Alaska / Hawaii
// lobe by which one's clip box the point falls in, the way albersUsa does. ok is false
// when the point is off-map.
func ProjectUS(lng, lat float64) (x, y float64, ok bool) {
	x, y = usLower48(lng, lat)
	if usInBox(x, y, usTX-0.455*usK, usTY-0.238*usK, usTX+0.455*usK, usTY+0.238*usK) {
		return x, y, true
	}
	x, y = usAlaska(lng, lat)
	if usInBox(x, y, usTX-0.425*usK+usEps, usTY+0.120*usK+usEps, usTX-0.214*usK-usEps, usTY+0.234*usK-usEps) {
		return x, y, true
	}
	x, y = usHawaii(lng, lat)
	if usInBox(x, y, usTX-0.214*usK+usEps, usTY+0.166*usK+usEps, usTX-0.115*usK-usEps, usTY+0.234*usK-usEps) {
		return x, y, true
	}
	return 0, 0, false
}

// ── state centroids (bounding-box centre of each path, for the tooltip anchor) ──────

var usStateCodes = func() []string {
	codes := make([]string, 0, len(usStates))
	for c := range usStates {
		codes = append(codes, c)
	}
	sort.Strings(codes)
	return codes
}()

var usCentroids = map[string][2]float64{}

func stateCentroid(code string) (float64, float64) {
	if c, ok := usCentroids[code]; ok {
		return c[0], c[1]
	}
	nums := pathNums(usStates[code].D)
	minX, minY, maxX, maxY := math.Inf(1), math.Inf(1), math.Inf(-1), math.Inf(-1)
	for i := 0; i+1 < len(nums); i += 2 {
		minX, maxX = math.Min(minX, nums[i]), math.Max(maxX, nums[i])
		minY, maxY = math.Min(minY, nums[i+1]), math.Max(maxY, nums[i+1])
	}
	cx, cy := (minX+maxX)/2, (minY+maxY)/2
	usCentroids[code] = [2]float64{cx, cy}
	return cx, cy
}

// pathNums extracts every number from a path d string. The state coords are all positive
// (the albersUsa box is 0..960 × 0..600), comma/letter separated, so a simple digit-run
// scan suffices — no sign or exponent handling needed.
func pathNums(s string) []float64 {
	var out []float64
	for i := 0; i < len(s); {
		if c := s[i]; (c >= '0' && c <= '9') || c == '.' {
			j := i
			for j < len(s) && ((s[j] >= '0' && s[j] <= '9') || s[j] == '.') {
				j++
			}
			if f, err := strconv.ParseFloat(s[i:j], 64); err == nil {
				out = append(out, f)
			}
			i = j
		} else {
			i++
		}
	}
	return out
}

// ── the component ────────────────────────────────────────────────────────────────

type USHeatmapPoint struct {
	Lat, Lng float64
	Value    float64
	Label    string
}

type USHeatmapProps struct {
	// State value map: USPS code ("CA", "TX", "DC") → number. Present states are shaded on
	// the sequential ramp; absent states get the no-data neutral.
	Data map[string]float64
	// lat/lng markers, projected onto the map. Points outside the US are dropped.
	Points []USHeatmapPoint

	Height       float64 // px; 0 = size from container width (960:600 aspect).
	Class        string
	Steps        int  // choropleth buckets 1–6 (default 6, the token count).
	NoTooltip    bool //
	Proportional bool // scale each dot's AREA by its value (radius ∝ √value).

	ValueFormat  func(float64) string
	PointColor   string  // default var(--color-chart-1)
	PointRadius  float64 // fixed radius (default 5); MAX radius when proportional.
	PointOpacity float64 // default 0.85
	StateName    func(string) string
}

const choroplethSteps = 6 // must match --color-choropleth-1..N

type usHover struct {
	kind string // "" | "state" | "point"
	code string
	idx  int
}

// USHeatmap is a CONTROLLER (it holds the hover signal across renders). Build it once and
// call Render(props) each render.
type USHeatmap struct {
	hover  *vdom.Signal[usHover]
	svgRef *vdom.Ref
	tipRef *vdom.Ref // the HTML tooltip, positioned imperatively so it follows the cursor
	props  USHeatmapProps
}

func NewUSHeatmap() *USHeatmap {
	return &USHeatmap{hover: vdom.NewSignal(usHover{}), svgRef: vdom.NewRef(), tipRef: vdom.NewRef()}
}

func (c *USHeatmap) setHover(h usHover) {
	if h != c.hover.Get() {
		c.hover.Set(h)
	}
}
func (c *USHeatmap) onLeave() { c.setHover(usHover{}) }

func (c *USHeatmap) onMove(e vdom.Event) {
	if c.props.NoTooltip {
		return
	}
	var hv usHover
	if s, ok := wasmruntime.ClosestAttr(e.Target(), "[data-pt]", "data-pt"); ok {
		idx, _ := strconv.Atoi(s)
		hv = usHover{kind: "point", idx: idx}
	} else if code, ok := wasmruntime.ClosestAttr(e.Target(), "[data-state]", "data-state"); ok {
		hv = usHover{kind: "state", code: code}
	}
	c.setHover(hv)
	if hv.kind == "" {
		return
	}
	// position the tooltip at the pointer (follows the cursor, like the Solid heatmap). The
	// svg is at the wrapper's top-left, so pointer-minus-svg-rect is the tooltip's left/top.
	r := wasmruntime.Measure(c.svgRef)
	if r.Width == 0 {
		return
	}
	wasmruntime.SetStyle(c.tipRef, "left", num(clampf(float64(e.ClientX())-r.X, 8, r.Width-8))+"px")
	wasmruntime.SetStyle(c.tipRef, "top", num(clampf(float64(e.ClientY())-r.Y, 8, r.Height-8))+"px")
	wasmruntime.SetStyle(c.tipRef, "transform", "translate(-50%, calc(-100% - 12px))")
}

// tooltipNode is the HTML tooltip (absolute, pointer-events-none), rendered every pass so
// its ref stays valid; onMove positions it imperatively. Hidden when nothing is hovered.
func (c *USHeatmap) tooltipNode(p USHeatmapProps, hv usHover, mn, mx float64, steps int) *vdom.VNode {
	base := "pointer-events-none absolute z-10 min-w-28 max-w-64 rounded-default border border-line bg-surface px-3 py-2 text-xs shadow-lg"
	if p.NoTooltip || hv.kind == "" {
		return vdom.Div(vdom.WithRef(c.tipRef), vdom.Attr("class", cx(base, "hidden")))
	}
	var title, value, swatch string
	if hv.kind == "state" {
		title = c.stateName(hv.code)
		if v, has := p.Data[hv.code]; has && !math.IsNaN(v) {
			value = c.valFmt(v)
			swatch = "var(--color-choropleth-" + strconv.Itoa(heatBucket(v, mn, mx, steps)) + ")"
		} else {
			value = "no data"
			swatch = "var(--color-surface-strong)"
		}
	} else if hv.idx >= 0 && hv.idx < len(p.Points) {
		pt := p.Points[hv.idx]
		title = pt.Label
		if title == "" {
			title = strconv.FormatFloat(pt.Lat, 'f', 2, 64) + ", " + strconv.FormatFloat(pt.Lng, 'f', 2, 64)
		}
		value = c.valFmt(pt.Value)
		swatch = p.PointColor
		if swatch == "" {
			swatch = "var(--color-chart-1)"
		}
	}
	content := []*vdom.VNode{
		vdom.Div(vdom.Attr("class", "flex items-center gap-2"),
			swatchSpan(swatch),
			vdom.Span(vdom.Attr("class", "font-medium text-ink"), vdom.Text(title))),
	}
	if value != "" {
		content = append(content, vdom.Div(vdom.Attr("class", "mt-1 font-semibold text-ink"),
			vdom.Attr("style", "font-variant-numeric:tabular-nums"), vdom.Text(value)))
	}
	return vdom.Div(kids([]vdom.Mod{vdom.WithRef(c.tipRef), vdom.Attr("class", base)}, content)...)
}

func (c *USHeatmap) valFmt(v float64) string {
	if c.props.ValueFormat != nil {
		return c.props.ValueFormat(v)
	}
	return groupNum(v)
}

func (c *USHeatmap) stateName(code string) string {
	if c.props.StateName != nil {
		return c.props.StateName(code)
	}
	if s, ok := usStates[code]; ok {
		return s.Name
	}
	return code
}

func heatBucket(v, mn, mx float64, steps int) int {
	t := 1.0
	if mx > mn {
		t = (v - mn) / (mx - mn)
	}
	return clampi(int(math.Floor(t*float64(steps))), 0, steps-1) + 1
}

func (c *USHeatmap) Render(p USHeatmapProps) *vdom.VNode {
	c.props = p
	hv := c.hover.Get()
	steps := clampi(p.Steps, 1, choroplethSteps)
	if p.Steps == 0 {
		steps = choroplethSteps
	}

	mn, mx, hasData := math.Inf(1), math.Inf(-1), false
	for _, v := range p.Data {
		if !math.IsNaN(v) && !math.IsInf(v, 0) {
			mn, mx, hasData = math.Min(mn, v), math.Max(mx, v), true
		}
	}

	var b strings.Builder
	// states (sorted for a stable diff order)
	for _, code := range usStateCodes {
		v, has := p.Data[code]
		fill := "var(--color-surface-strong)"
		if has && !math.IsNaN(v) {
			fill = "var(--color-choropleth-" + strconv.Itoa(heatBucket(v, mn, mx, steps)) + ")"
		}
		op := "1"
		if hv.kind == "state" && hv.code == code {
			op = "0.82"
		}
		b.WriteString(`<path d="` + usStates[code].D + `" data-state="` + code + `" fill="` + fill + `" fill-opacity="` + op + `" stroke="var(--color-surface)" stroke-width="0.8"/>`)
	}

	// points
	pc := p.PointColor
	if pc == "" {
		pc = "var(--color-chart-1)"
	}
	pop := p.PointOpacity
	if pop == 0 {
		pop = 0.85
	}
	maxR := p.PointRadius
	if maxR == 0 {
		if p.Proportional {
			maxR = 16
		} else {
			maxR = 5
		}
	}
	minR := math.Min(3, maxR*0.35)
	maxV := 1.0
	if p.Proportional {
		for _, pt := range p.Points {
			if pt.Value > maxV {
				maxV = pt.Value
			}
		}
	}
	for i, pt := range p.Points {
		x, y, ok := ProjectUS(pt.Lng, pt.Lat)
		if !ok {
			continue
		}
		r := maxR
		if p.Proportional {
			r = minR + (maxR-minR)*math.Sqrt(clampf(pt.Value/maxV, 0, 1))
		}
		if hv.kind == "point" && hv.idx == i {
			r += 2
		}
		b.WriteString(`<circle data-pt="` + strconv.Itoa(i) + `" cx="` + num(x) + `" cy="` + num(y) + `" r="` + num(r) + `" fill="` + svgEsc(pc) + `" fill-opacity="` + num(pop) + `"/>`)
	}

	svgMods := []vdom.Mod{
		vdom.WithRef(c.svgRef),
		vdom.Attr("viewBox", usViewBox),
		vdom.Attr("role", "img"),
		vdom.Raw(b.String()),
	}
	if p.Height > 0 {
		svgMods = append(svgMods, vdom.Attr("class", "block"), vdom.Attr("style", "height:"+num(p.Height)+"px;width:auto;margin:0 auto"))
	} else {
		svgMods = append(svgMods, vdom.Attr("class", "block w-full"), vdom.Attr("style", "aspect-ratio:960/600"))
	}
	if !p.NoTooltip {
		svgMods = append(svgMods,
			vdom.OnEvent(vdom.EVENT_POINTERMOVE, c.onMove),
			vdom.On(vdom.EVENT_POINTERLEAVE, c.onLeave),
		)
	}

	children := []*vdom.VNode{vdom.Svg(svgMods...), c.tooltipNode(p, hv, mn, mx, steps)}
	if hasData {
		children = append(children, choroplethLegend(mn, mx, steps, c.valFmt))
	}
	return vdom.Div(kids([]vdom.Mod{vdom.Attr("class", cx("relative w-full", p.Class))}, children)...)
}

func choroplethLegend(mn, mx float64, steps int, fmt func(float64) string) *vdom.VNode {
	swatches := []*vdom.VNode{}
	for k := 1; k <= steps; k++ {
		swatches = append(swatches, vdom.Span(vdom.Attr("class", "h-3 w-6"),
			vdom.Attr("style", "background-color:var(--color-choropleth-"+strconv.Itoa(k)+")")))
	}
	ramp := vdom.Div(kids([]vdom.Mod{vdom.Attr("class", "flex overflow-hidden rounded-xs")}, swatches)...)
	return vdom.Div(vdom.Attr("class", "mt-3 flex items-center gap-2 text-xs text-ink-muted"),
		vdom.Span(vdom.Attr("style", "font-variant-numeric:tabular-nums"), vdom.Text(fmt(mn))),
		ramp,
		vdom.Span(vdom.Attr("style", "font-variant-numeric:tabular-nums"), vdom.Text(fmt(mx))),
	)
}