// This is Free Software under GNU Affero General Public License v >= 3.0
// without warranty, see README.md and license for details.
//
// SPDX-License-Identifier: AGPL-3.0-or-later
// License-Filename: LICENSES/AGPL-3.0.txt
//
// Copyright (C) 2018 by via donau
//   – Österreichische Wasserstraßen-Gesellschaft mbH
// Software engineering by Intevation GmbH
//
// Author(s):
//  * Sascha L. Teichmann <sascha.teichmann@intevation.de>

package mesh

import (
	"cmp"
	"math"
	"runtime"
	"slices"
	"sync"
)

// STRTreeDefaultEntries is the default number of children per node and leaf.
const STRTreeDefaultEntries = 8

// STRTree is a Sort Tile Recurse tree.
type STRTree struct {
	// Entries is the total number of entries.
	Entries int
	// tin is a reference to the TIN this tree is build from.
	tin *Tin
	// index is the internal tree representation.
	index []int32
	// bboxes are the bounding boxes used in nodes and leaves.
	bboxes []Box2D
}

// Vertical does a vertical cross cut from (x1, y1) to (x2, y2).
func (s *STRTree) Vertical(x1, y1, x2, y2 float64, fn func(*Triangle)) {
	box := Box2D{
		X1: math.Min(x1, x2),
		Y1: math.Min(y1, y2),
		X2: math.Max(x1, x2),
		Y2: math.Max(y1, y2),
	}

	// out of bounding box
	if box.X2 < s.tin.Min.X || s.tin.Max.X < box.X1 ||
		box.Y2 < s.tin.Min.Y || s.tin.Max.Y < box.Y1 {
		return
	}

	line := NewPlane2D(x1, y1, x2, y2)

	vertices := s.tin.Vertices

	stack := []int32{s.index[0]}

	for len(stack) > 0 {
		top := stack[len(stack)-1]
		stack = stack[:len(stack)-1]

		if top > 0 { // node
			if b := s.bbox(top); b.Intersects(box) && b.IntersectsPlane(line) {
				n := s.index[top+1]
				stack = append(stack, s.index[top+2:top+2+n]...)
			}
		} else { // leaf
			top = -top - 1
			if b := s.bbox(top); !b.Intersects(box) || !b.IntersectsPlane(line) {
				continue
			}
			n := s.index[top+1]
			for _, idx := range s.index[top+2 : top+2+n] {
				ti := s.tin.Triangles[idx]
				t := Triangle{
					vertices[ti[0]],
					vertices[ti[1]],
					vertices[ti[2]],
				}
				v0 := line.Eval(t[0].X, t[0].Y)
				v1 := line.Eval(t[1].X, t[1].Y)
				v2 := line.Eval(t[2].X, t[2].Y)

				if onPlane(v0) || onPlane(v1) || onPlane(v2) ||
					sides(sides(sides(0, v0), v1), v2) == 3 {
					fn(&t)
				}
			}
		}
	}
}

// Min return the minimum vertex of the TIN this tree is build from.
func (s *STRTree) Min() Vertex { return s.tin.Min }

// Max return the maximum vertex of the TIN this tree is build from.
func (s *STRTree) Max() Vertex { return s.tin.Max }

// EPSG return the EPSF code of the TIN this tree is build from.
func (s *STRTree) EPSG() uint32 { return s.tin.EPSG }

// Value returns the Z Value for a given X/Y point.
// Second return value is false if the point is out of bound.
func (s *STRTree) Value(x, y float64) (float64, bool) {

	if len(s.index) == 0 {
		return 0, false
	}

	stack := []int32{s.index[0]}

	vertices := s.tin.Vertices

	for len(stack) > 0 {
		top := stack[len(stack)-1]
		stack = stack[:len(stack)-1]

		if top > 0 { // node
			if s.bbox(top).Contains(x, y) {
				n := s.index[top+1]
				stack = append(stack, s.index[top+2:top+2+n]...)
			}
		} else { // leaf
			top = -top - 1
			if !s.bbox(top).Contains(x, y) {
				continue
			}
			n := s.index[top+1]
			for _, idx := range s.index[top+2 : top+2+n] {
				ti := s.tin.Triangles[idx]
				t := Triangle{
					vertices[ti[0]],
					vertices[ti[1]],
					vertices[ti[2]],
				}
				if t.Contains(x, y) {
					return t.Plane3D().Z(x, y), true
				}
			}
		}
	}
	return 0, false
}

// Build builds a tree for a given TIN.
func (s *STRTree) Build(t *Tin) {

	if s.Entries == 0 {
		s.Entries = STRTreeDefaultEntries
	}

	s.tin = t

	all := make([]int32, len(t.Triangles))

	for i := range all {
		all[i] = int32(i)
	}

	s.index = append(s.index, 0)

	root := s.build(all)

	s.index[0] = root
}

// countNodes returns the number of nodes for a given index.
func (s *STRTree) countNodes() int {
	if len(s.index) == 0 {
		return 0
	}
	count := 0
	stack := []int32{s.index[0]}
	for len(stack) > 0 {
		top := stack[len(stack)-1]
		stack = stack[:len(stack)-1]
		count++
		if top > 0 { // node
			n := s.index[top+1]
			stack = append(stack, s.index[top+2:top+2+n]...)
		}
	}
	return count
}

// rebuildBBoxes rebuilds the bounding boxes for a given index and triangles.
func (s *STRTree) rebuildBBoxes() {

	num := s.countNodes()
	if num == 0 {
		return
	}
	bboxes := make([]Box2D, num)
	s.bboxes = bboxes

	var recurse func(int32, int32) Box2D
	recurse = func(top, depth int32) Box2D {
		var box Box2D
		if top > 0 { // node
			n := s.index[top+1]
			for i, child := range s.index[top+2 : top+2+n] {
				if i == 0 {
					box = recurse(child, depth+1)
				} else {
					box = box.Union(recurse(child, depth+1))
				}
			}
		} else { // leaf
			top = -top - 1
			n := s.index[top+1]
			for i, idx := range s.index[top+2 : top+2+n] {
				ti := s.tin.Triangles[idx]
				t := Triangle{
					s.tin.Vertices[ti[2]],
					s.tin.Vertices[ti[0]],
					s.tin.Vertices[ti[1]],
				}
				if i == 0 {
					box = t.BBox()
				} else {
					box = box.Union(t.BBox())
				}
			}
		}
		bboxes[s.index[top]] = box
		return box
	}

	// Check if we can boost re-building with some extra CPU cores.
	cpus := max(1, runtime.NumCPU()/2)
	if cpus == 1 {
		recurse(s.index[0], 0)
		return
	}

	type job struct {
		top   int32
		depth int32
		store func(Box2D)
	}

	jobs := make(chan job)
	var wg sync.WaitGroup

	singleThreaded := recurse

	for i := 0; i < cpus; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for j := range jobs {
				j.store(singleThreaded(j.top, j.depth))
			}
		}()
	}

	recurse = func(top, depth int32) Box2D {
		// Only handle nodes at level 1.
		if depth != 1 || top <= 0 {
			return singleThreaded(top, depth)
		}
		var (
			first = true
			box   Box2D
			done  = make(chan struct{})
			n     = s.index[top+1]
			count = int32(0)
			mu    sync.Mutex
		)
		store := func(b Box2D) {
			mu.Lock()
			defer mu.Unlock()
			if first {
				first = false
				box = b
			} else {
				box = box.Union(b)
			}
			if count++; count >= n {
				close(done)
			}
		}
		for _, child := range s.index[top+2 : top+2+n] {
			jobs <- job{top: child, depth: depth + 1, store: store}
		}
		<-done
		bboxes[s.index[top]] = box
		return box
	}
	recurse(s.index[0], 0)
	close(jobs)
	wg.Wait()
}

// BuildWithout builds a tree for a given TIN ignoring the triangles
// with the indices given in the remove map.
func (s *STRTree) BuildWithout(t *Tin, remove map[int32]struct{}) {

	if s.Entries == 0 {
		s.Entries = STRTreeDefaultEntries
	}

	s.tin = t

	all := make([]int32, 0, len(t.Triangles)-len(remove))

	for i := 0; i < len(t.Triangles); i++ {
		idx := int32(i)
		if _, found := remove[idx]; !found {
			all = append(all, idx)
		}
	}

	s.index = append(s.index, 0)

	root := s.build(all)

	s.index[0] = root
}

// Clip returns the indices of the triangles of the TIN
// which are not fully covered by the given polygon.
func (s *STRTree) Clip(p *Polygon) map[int32]struct{} {

	removed := make(map[int32]struct{})

	stack := []int32{s.index[0]}

	vertices := s.tin.Vertices

	for len(stack) > 0 {
		top := stack[len(stack)-1]
		stack = stack[:len(stack)-1]

		if top > 0 { // node
			switch p.IntersectionBox2D(s.bbox(top)) {
			case IntersectionInside:
				// all triangles are inside polygon
			case IntersectionOutSide:
				// all triangles are outside polygon
				s.allTriangles(top, removed)
			default:
				// mixed bag
				n := s.index[top+1]
				stack = append(stack, s.index[top+2:top+2+n]...)
			}
		} else { // leaf
			top = -top - 1
			switch p.IntersectionBox2D(s.bbox(top)) {
			case IntersectionInside:
				// all triangles are inside polygon
			case IntersectionOutSide:
				// all triangles are outside polygon
				n := s.index[top+1]
				for _, idx := range s.index[top+2 : top+2+n] {
					removed[idx] = struct{}{}
				}
			default:
				n := s.index[top+1]
				for _, idx := range s.index[top+2 : top+2+n] {
					ti := s.tin.Triangles[idx]
					t := Triangle{
						vertices[ti[0]],
						vertices[ti[1]],
						vertices[ti[2]],
					}
					if p.IntersectionWithTriangle(&t) != IntersectionInside {
						removed[idx] = struct{}{}
					}
				}
			}
		}
	}

	return removed
}

func (s *STRTree) allTriangles(pos int32, tris map[int32]struct{}) {
	stack := []int32{pos}
	for len(stack) > 0 {
		top := stack[len(stack)-1]
		stack = stack[:len(stack)-1]
		if top > 0 { // node
			n := s.index[top+1]
			stack = append(stack, s.index[top+2:top+2+n]...)
		} else { // leaf
			top = -top - 1
			n := s.index[top+1]
			for _, idx := range s.index[top+2 : top+2+n] {
				tris[idx] = struct{}{}
			}
		}
	}
}

func (s *STRTree) build(items []int32) int32 {
	slices.SortFunc(items, func(i, j int32) int {
		ti := s.tin.Triangles[i]
		tj := s.tin.Triangles[j]
		return cmp.Compare(s.tin.Vertices[ti[0]].X, s.tin.Vertices[tj[0]].X)
	})

	P := int(math.Ceil(float64(len(items)) / float64(s.Entries)))
	S := int(math.Ceil(math.Sqrt(float64(P))))

	slices := s.strSplit(items, S)

	leaves := s.strJoin(
		slices, S,
		func(i, j int32) int {
			ti := s.tin.Triangles[i]
			tj := s.tin.Triangles[j]
			return cmp.Compare(s.tin.Vertices[ti[0]].Y, s.tin.Vertices[tj[0]].Y)
		},
		s.allocLeaf,
	)

	return s.buildNodes(leaves)
}

func (s *STRTree) buildNodes(items []int32) int32 {

	if len(items) <= s.Entries {
		return s.allocNode(items)
	}

	slices.SortFunc(items, func(i, j int32) int {
		return cmp.Compare(s.bbox(i).X1, s.bbox(j).X1)
	})

	P := int(math.Ceil(float64(len(items)) / float64(s.Entries)))
	S := int(math.Ceil(math.Sqrt(float64(P))))

	slices := s.strSplit(items, S)

	nodes := s.strJoin(
		slices, S,
		func(i, j int32) int { return cmp.Compare(s.bbox(i).Y1, s.bbox(j).Y1) },
		s.allocNode,
	)

	return s.buildNodes(nodes)
}

func (s *STRTree) bbox(idx int32) Box2D {
	if idx < 0 { // Don't care if leaf or node.
		idx = -idx - 1
	}
	return s.bboxes[s.index[idx]]
}

func (s *STRTree) strSplit(items []int32, S int) [][]int32 {
	sm := S * s.Entries
	slices := make([][]int32, S)
	for i := range slices {
		var n int
		if l := len(items); l < sm {
			n = l
		} else {
			n = sm
		}
		slices[i] = items[:n]
		items = items[n:]
	}
	return slices
}

func (s *STRTree) strJoin(
	slics [][]int32, S int,
	cmp func(int32, int32) int,
	alloc func([]int32) int32,
) []int32 {
	nodes := make([]int32, 0, S*S)

	for _, slice := range slics {
		slices.SortFunc(slice, cmp)

		for len(slice) > 0 {
			var n int
			if l := len(slice); l >= s.Entries {
				n = s.Entries
			} else {
				n = l
			}
			nodes = append(nodes, alloc(slice[:n]))
			slice = slice[n:]
		}
	}
	return nodes
}

func (s *STRTree) allocNode(items []int32) int32 {
	pos := len(s.index)
	s.index = append(s.index, 0, int32(len(items)))
	s.index = append(s.index, items...)
	if len(items) > 0 {
		box := s.bbox(items[0])
		for i := 1; i < len(items); i++ {
			box = box.Union(s.bbox(items[i]))
		}
		s.index[pos] = int32(s.allocBBox(box))
	}
	return int32(pos)
}

func (s *STRTree) allocBBox(box Box2D) int {
	pos := len(s.bboxes)
	s.bboxes = append(s.bboxes, box)
	return pos
}

func (s *STRTree) allocLeaf(items []int32) int32 {
	pos := len(s.index)
	s.index = append(s.index, 0, int32(len(items)))
	s.index = append(s.index, items...)
	if len(items) > 0 {
		vertices := s.tin.Vertices
		ti := s.tin.Triangles[items[0]]
		t := Triangle{
			vertices[ti[0]],
			vertices[ti[1]],
			vertices[ti[2]],
		}
		box := t.BBox()
		for i := 1; i < len(items); i++ {
			it := items[i]
			ti := s.tin.Triangles[it]
			t := Triangle{
				vertices[ti[0]],
				vertices[ti[1]],
				vertices[ti[2]],
			}
			box = box.Union(t.BBox())
		}
		s.index[pos] = int32(s.allocBBox(box))
	}
	return int32(-(pos + 1))
}

// allTriangleIndices passes all triangle indices to callback fn.
func (s *STRTree) allTriangleIndices(fn func([]int32)) {
	stack := []int32{s.index[0]}
	for len(stack) > 0 {
		top := stack[len(stack)-1]
		stack = stack[:len(stack)-1]
		if top > 0 { // node
			n := s.index[top+1]
			stack = append(stack, s.index[top+2:top+2+n]...)
		} else { // leaf
			top = -top - 1
			n := s.index[top+1]
			fn(s.index[top+2 : top+2+n])
		}
	}
}

// sortIndices sorts the indices in the inner nodes and leaves of the tree.
// This helps if the index array is delta encoded in serialization.
// It also may help reduding CPU cache line misses during usage because
// redirects are ordered closer together.
func (s *STRTree) sortIndices() {
	stack := []int32{s.index[0]}
	for len(stack) > 0 {
		top := stack[len(stack)-1]
		stack = stack[:len(stack)-1]
		if top > 0 { // node
			n := s.index[top+1]
			children := s.index[top+2 : top+2+n]
			stack = append(stack, children...)
			slices.Sort(children)
		} else { // leaf
			top = -top - 1
			n := s.index[top+1]
			triangles := s.index[top+2 : top+2+n]
			slices.Sort(triangles)
		}
	}
}