// 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 (
	"bytes"
	"compress/gzip"
	"encoding/binary"
	"io"
	"math"
	"sort"

	"gemma.intevation.de/gemma/pkg/log"
)

// 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
}

// 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) serializeIndex(w io.Writer) error {

	if err := binary.Write(w, binary.LittleEndian, int32(len(s.index))); err != nil {
		return err
	}

	var buf [binary.MaxVarintLen32]byte

	var last int32
	var written int

	for _, x := range s.index {
		delta := x - last
		n := binary.PutVarint(buf[:], int64(delta))
		for p := buf[:n]; len(p) > 0; p = p[n:] {
			var err error
			if n, err = w.Write(p); err != nil {
				return err
			}
			written += n
		}
		last = x
	}

	log.Infof("compressed index in bytes: %d %.2f (%d %.2f)\n",
		written,
		float64(written)/(1024*1024),
		4*len(s.index),
		float64(4*len(s.index))/(1024*1024),
	)

	return nil
}

func (s *STRTree) serializeBBoxes(w io.Writer) (rr error) {

	if err := binary.Write(w, binary.LittleEndian, int32(len(s.bboxes))); err != nil {
		return err
	}

	var err error

	write := func(v float64) {
		if err == nil {
			err = binary.Write(w, binary.LittleEndian, math.Float64bits(v))
		}
	}
	for _, box := range s.bboxes {
		write(box.X1)
		write(box.Y1)
		write(box.X2)
		write(box.Y2)
	}

	return err
}

// Bytes serializes this tree to a byte slice.
func (s *STRTree) Bytes() ([]byte, error) {

	var buf bytes.Buffer
	w, err := gzip.NewWriterLevel(&buf, gzip.BestSpeed)
	if err != nil {
		return nil, err
	}

	if err := s.tin.serialize(w); err != nil {
		return nil, err
	}

	if err := binary.Write(w, binary.LittleEndian, uint8(s.Entries)); err != nil {
		return nil, err
	}

	if err := s.serializeIndex(w); err != nil {
		return nil, err
	}

	if err := s.serializeBBoxes(w); err != nil {
		return nil, err
	}

	if err := w.Close(); err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

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 {
	sort.Slice(items, func(i, j int) bool {
		ti := s.tin.Triangles[items[i]]
		tj := s.tin.Triangles[items[j]]
		return 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) bool {
			ti := s.tin.Triangles[i]
			tj := s.tin.Triangles[j]
			return 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)
	}

	sort.Slice(items, func(i, j int) bool {
		return s.bbox(items[i]).X1 < s.bbox(items[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) bool { return 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(
	slices [][]int32, S int,
	less func(int32, int32) bool,
	alloc func([]int32) int32,
) []int32 {
	nodes := make([]int32, 0, S*S)

	for _, slice := range slices {
		sort.Slice(slice, func(i, j int) bool {
			return less(slice[i], slice[j])
		})

		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))
}