package radarfetch

import (
	"math"
)

// SegmentClusters finds 8-connected regions where dBZ >= thr (e.g., 40),
// computes stats and recommended sampling points per cluster.
// Input grid: 256x256, invalid as NaN; bounds/res used to compute lon/lat.
func SegmentClusters(grid [][]*float64, bounds Bounds, resDeg float64, thr float64) []Cluster {
	const W, H = 256, 256
	if len(grid) != H || (len(grid) > 0 && len(grid[0]) != W) {
		return nil
	}
	// Mask of eligible pixels
	mask := make([][]bool, H)
	for r := 0; r < H; r++ {
		mask[r] = make([]bool, W)
		for c := 0; c < W; c++ {
			if grid[r][c] == nil {
				continue
			}
			v := *grid[r][c]
			if v >= thr {
				mask[r][c] = true
			}
		}
	}

	// Visited flags
	vis := make([][]bool, H)
	for r := 0; r < H; r++ {
		vis[r] = make([]bool, W)
	}

	// 8-neighborhood
	nbr := [8][2]int{{-1, -1}, {-1, 0}, {-1, 1}, {0, -1}, {0, 1}, {1, -1}, {1, 0}, {1, 1}}

	var clusters []Cluster
	clusterID := 0
	for r := 0; r < H; r++ {
		for c := 0; c < W; c++ {
			if !mask[r][c] || vis[r][c] {
				continue
			}
			// BFS/DFS stack
			stack := [][2]int{{r, c}}
			vis[r][c] = true
			// stats
			area := 0
			sumW := 0.0
			sumWR := 0.0
			sumWC := 0.0
			maxDBZ := -math.MaxFloat64
			sumDBZ := 0.0
			minR, minC := r, c
			maxR, maxC := r, c
			pixels := make([][2]int, 0, 512)
			for len(stack) > 0 {
				cur := stack[len(stack)-1]
				stack = stack[:len(stack)-1]
				rr, cc := cur[0], cur[1]
				area++
				dbz := *grid[rr][cc]
				w := dbz // dBZ-weighted centroid
				sumW += w
				sumWR += float64(rr) * w
				sumWC += float64(cc) * w
				if dbz > maxDBZ {
					maxDBZ = dbz
				}
				sumDBZ += dbz
				if rr < minR {
					minR = rr
				}
				if cc < minC {
					minC = cc
				}
				if rr > maxR {
					maxR = rr
				}
				if cc > maxC {
					maxC = cc
				}
				pixels = append(pixels, [2]int{rr, cc})
				for _, d := range nbr {
					nr, nc := rr+d[0], cc+d[1]
					if nr < 0 || nr >= H || nc < 0 || nc >= W {
						continue
					}
					if vis[nr][nc] || !mask[nr][nc] {
						continue
					}
					vis[nr][nc] = true
					stack = append(stack, [2]int{nr, nc})
				}
			}
			if area == 0 {
				continue
			}
			// centroid (row/col)
			cr, cc := 0.0, 0.0
			if sumW > 0 {
				cr = sumWR / sumW
				cc = sumWC / sumW
			} else {
				// fallback to geometric center
				cr = float64(minR+maxR) / 2.0
				cc = float64(minC+maxC) / 2.0
			}
			// Convert centroid to lon/lat (pixel center)
			clon := bounds.West + (cc+0.5)*resDeg
			clat := bounds.South + (cr+0.5)*resDeg

			// Sample points: center + four rays (N,S,E,W) until boundary
			samples := make([]Sample, 0, 5)
			samples = append(samples, Sample{Row: int(math.Round(cr)), Col: int(math.Round(cc)), Lon: clon, Lat: clat, Role: "center"})
			// helper to step ray and clamp to last in-mask pixel
			stepRay := func(dr, dc int, role string) {
				rr := int(math.Round(cr))
				cc2 := int(math.Round(cc))
				lastR, lastC := rr, cc2
				for {
					rr += dr
					cc2 += dc
					if rr < 0 || rr >= H || cc2 < 0 || cc2 >= W {
						break
					}
					if !mask[rr][cc2] {
						break
					}
					lastR, lastC = rr, cc2
				}
				lon := bounds.West + (float64(lastC)+0.5)*resDeg
				lat := bounds.South + (float64(lastR)+0.5)*resDeg
				if lastR != samples[0].Row || lastC != samples[0].Col {
					samples = append(samples, Sample{Row: lastR, Col: lastC, Lon: lon, Lat: lat, Role: role})
				}
			}
			stepRay(-1, 0, "ray_n")
			stepRay(1, 0, "ray_s")
			stepRay(0, 1, "ray_e")
			stepRay(0, -1, "ray_w")

			avgDBZ := sumDBZ / float64(area)
			cluster := Cluster{
				ID:     clusterID,
				AreaPx: area,
				MaxDBZ: maxDBZ,
				AvgDBZ: avgDBZ,
				Row:    int(math.Round(cr)),
				Col:    int(math.Round(cc)),
				Lon:    clon,
				Lat:    clat,
				MinRow: minR, MinCol: minC, MaxRow: maxR, MaxCol: maxC,
				Samples: samples,
			}
			clusters = append(clusters, cluster)
			clusterID++
		}
	}
	return clusters
}