// 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) 2024 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 (
	"bufio"
	"cmp"
	"encoding/binary"
	"io"
	"math"
	"slices"

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

const quantScale = 1_000

func quant(x float64) int64   { return int64(math.Round(x * quantScale)) }
func unquant(x int64) float64 { return float64(x) / quantScale }

func (s *STRTree) serializeV2(w io.Writer) error {
	return serializer(w, s,
		serializeHeader(),
		serializeVersion(2),
		(*STRTree).serializeTinV2,
		(*STRTree).serializeEntries,
		(*STRTree).serializeIndexV2,
		(*STRTree).serializeBBoxesV2,
	)
}

func (s *STRTree) serializeTinV2(w io.Writer) error {
	return s.tin.serializeV2(w)
}

func (t *Tin) serializeV2(w io.Writer) error {
	return serializer(w, t,
		(*Tin).serializeEPSG,
		(*Tin).serializeExtent,
		(*Tin).serializeVerticesV2,
		(*Tin).serializeTrianglesV2,
	)
}

func zigZag(buf []byte, w io.Writer, x int64) error {
	_, err := w.Write(buf[:binary.PutVarint(buf, x)])
	return err
}

func (t *Tin) serializeVerticesV2(w io.Writer) error {
	vertices := t.Vertices
	if err := binary.Write(w, binary.LittleEndian, uint32(len(vertices))); err != nil {
		return err
	}
	m := t.Min
	buf := make([]byte, binary.MaxVarintLen64)
	for dim := 0; dim <= 2; dim++ {
		delta := common.Delta()
		for i := range vertices {
			var v float64
			switch dim {
			case 0:
				v = vertices[i].X - m.X
			case 1:
				v = vertices[i].Y - m.Y
			case 2:
				v = vertices[i].Z - m.Z
			}
			q := quant(v)
			d := delta(q)
			if err := zigZag(buf, w, d); err != nil {
				return err
			}
		}
	}
	return nil
}

func (t *Tin) serializeTrianglesV2(w io.Writer) error {
	triangles := t.Triangles
	if err := binary.Write(w, binary.LittleEndian, uint32(len(triangles))); err != nil {
		return err
	}
	buf := make([]byte, binary.MaxVarintLen64)
	delta := common.Delta()
	for _, tri := range triangles {
		for _, idx := range tri {
			d := delta(int64(idx))
			if err := zigZag(buf, w, d); err != nil {
				return err
			}
		}
	}
	return nil
}

func (s *STRTree) serializeIndexV2(w io.Writer) error {
	index := s.index
	if err := binary.Write(w, binary.LittleEndian, uint32(len(index))); err != nil {
		return err
	}
	buf := make([]byte, binary.MaxVarintLen64)
	delta := common.Delta()
	for _, idx := range index {
		d := delta(int64(idx))
		if err := zigZag(buf, w, d); err != nil {
			return err
		}
	}
	return nil
}

func (s *STRTree) serializeBBoxesV2(w io.Writer) error {
	boxes := s.bboxes
	if err := binary.Write(w, binary.LittleEndian, uint32(len(boxes))); err != nil {
		return err
	}
	var (
		m      = s.Min()
		x1     = func(b Box2D) int64 { return quant(b.X1 - m.X) }
		y1     = func(b Box2D) int64 { return quant(b.Y1 - m.Y) }
		width  = func(b Box2D) int64 { return quant(b.X2 - b.X1) }
		height = func(b Box2D) int64 { return quant(b.Y2 - b.Y1) }
		buf    = make([]byte, binary.MaxVarintLen64)
	)
	for _, proj := range []func(Box2D) int64{x1, y1, width, height} {
		delta := common.Delta()
		for _, b := range boxes {
			v := delta(proj(b))
			if err := zigZag(buf, w, v); err != nil {
				return err
			}
		}
	}
	return nil
}

func (s *STRTree) deserializeV2(r *bufio.Reader) error {
	// The header is already skipped here.
	return serializer(r, s,
		(*STRTree).deserializeTinV2,
		(*STRTree).deserializeEntries,
		(*STRTree).deserializeIndexV2,
		(*STRTree).deserializeBBoxesV2,
	)
}

func (s *STRTree) deserializeTinV2(r *bufio.Reader) error {
	s.tin = new(Tin)
	return s.tin.deserializeV2(r)
}

func (s *STRTree) deserializeBBoxesV2(r *bufio.Reader) error {
	var numBoxes uint32
	if err := binary.Read(r, binary.LittleEndian, &numBoxes); err != nil {
		return err
	}
	boxes := make([]Box2D, numBoxes)
	s.bboxes = boxes
	var (
		m      = s.Min()
		x1     = func(b *Box2D, v float64) { b.X1 = v + m.X }
		y1     = func(b *Box2D, v float64) { b.Y1 = v + m.Y }
		width  = func(b *Box2D, v float64) { b.X2 = b.X1 + v }
		height = func(b *Box2D, v float64) { b.Y2 = b.Y1 + v }
	)
	for _, unproj := range []func(*Box2D, float64){x1, y1, width, height} {
		invDelta := common.InvDelta()
		for i := range boxes {
			v, err := binary.ReadVarint(r)
			if err != nil {
				return err
			}
			unproj(&boxes[i], unquant(invDelta(v)))
		}
	}
	return nil
}

func (s *STRTree) deserializeIndexV2(r *bufio.Reader) error {
	var numIndices uint32
	if err := binary.Read(r, binary.LittleEndian, &numIndices); err != nil {
		return err
	}
	index := make([]int32, numIndices)
	s.index = index
	invDelta := common.InvDelta()
	for i := range index {
		d, err := binary.ReadVarint(r)
		if err != nil {
			return err
		}
		index[i] = int32(invDelta(d))
	}
	return nil
}

func (t *Tin) deserializeV2(r *bufio.Reader) error {
	return serializer(r, t,
		(*Tin).deserializeEPSG,
		(*Tin).deserializeExtent,
		(*Tin).deserializeVerticesV2,
		(*Tin).deserializeTrianglesV2,
	)
}

func (t *Tin) deserializeVerticesV2(r *bufio.Reader) error {
	var numVertices uint32
	if err := binary.Read(r, binary.LittleEndian, &numVertices); err != nil {
		return err
	}
	vertices := make([]Vertex, numVertices)
	t.Vertices = vertices
	m := t.Min
	for dim := 0; dim <= 2; dim++ {
		invDelta := common.InvDelta()
		for i := range vertices {
			d, err := binary.ReadVarint(r)
			if err != nil {
				return err
			}
			v := unquant(invDelta(d))
			switch dim {
			case 0:
				vertices[i].X = v + m.X
			case 1:
				vertices[i].Y = v + m.Y
			case 2:
				vertices[i].Z = v + m.Z
			}
		}
	}
	return nil
}

func (t *Tin) deserializeTrianglesV2(r *bufio.Reader) error {
	var numTriangles uint32
	if err := binary.Read(r, binary.LittleEndian, &numTriangles); err != nil {
		return err
	}
	indices := make([]int32, 3*numTriangles)
	triangles := make([][]int32, numTriangles)
	t.Triangles = triangles
	invDelta := common.InvDelta()
	for i := range triangles {
		tri := indices[:3]
		indices = indices[3:]
		triangles[i] = tri
		for j := range tri {
			v, err := binary.ReadVarint(r)
			if err != nil {
				return err
			}
			tri[j] = int32(invDelta(v))
		}
	}
	return nil
}

// optimizeForSerializationV2 prepares the mesh to be stored in V2.
func (s *STRTree) optimizeForSerializationV2() {
	s.removeUnusedTriangles()
	s.sortVertices()
	s.sortIndices()
}

// removeUnusedTriangles removes all triangles from the
// TIN that are not registered in the spatial index.
func (s *STRTree) removeUnusedTriangles() {
	used := make([]bool, len(s.tin.Triangles))
	unused := len(used)
	s.allTriangleIndices(func(indices []int32) {
		for _, idx := range indices {
			used[idx] = true
			unused-- // They only occur once in the spatial index.
		}
	})
	if unused <= 0 {
		return
	}
	newTris := make([][]int32, 0, len(s.tin.Triangles)-unused)
	remap := map[int32]int32{}
	for idx, tri := range s.tin.Triangles {
		if used[idx] {
			remap[int32(idx)] = int32(len(newTris))
			newTris = append(newTris, tri)
		}
	}
	s.tin.Triangles = newTris
	// Update the spatial index as the gaps are now closed.
	s.allTriangleIndices(func(indices []int32) {
		for i, idx := range indices {
			indices[i] = remap[idx]
		}
	})
}

// sortVertices so that the deltas between neighbors are minimized.
func (s *STRTree) sortVertices() {
	vertices := s.tin.Vertices
	indices := make([]int32, len(vertices))
	for i := range indices {
		indices[i] = int32(i)
	}

	// Take aspect ratio of bbox into account.
	width := s.tin.Max.X - s.tin.Min.X
	height := s.tin.Max.Y - s.tin.Min.Y
	if width == 0 || height == 0 {
		return
	}

	// v = len(vertices)
	// r = max(w, h)/min(w, h)
	// n * t = v <=> n = v/t
	// n / t = r <=> n = r*t
	// v/t = r*t
	// v = r*t^2
	// v/r = t^2
	// t = sqrt(v/r)
	r := float64(max(width, height)) / float64(min(width, height))
	t := int(math.Ceil(math.Sqrt(float64(len(vertices)) / r)))

	var d1, d2 func(int32, int32) int

	if width > height {
		// Sort in X first and slices alternating up and down
		slices.SortFunc(indices, func(a, b int32) int {
			return cmp.Compare(vertices[a].X, vertices[b].X)
		})
		d1 = func(a, b int32) int { return cmp.Compare(vertices[a].Y, vertices[b].Y) }
		d2 = func(a, b int32) int { return cmp.Compare(vertices[b].Y, vertices[a].Y) }
	} else {
		// Sort in Y first and slices alternating left and right
		slices.SortFunc(indices, func(a, b int32) int {
			return cmp.Compare(vertices[a].Y, vertices[b].Y)
		})
		d1 = func(a, b int32) int { return cmp.Compare(vertices[a].X, vertices[b].X) }
		d2 = func(a, b int32) int { return cmp.Compare(vertices[b].X, vertices[a].X) }
	}

	for p := indices; len(p) > 0; d1, d2 = d2, d1 {
		l := min(len(p), t)
		slices.SortStableFunc(p[:l], d1)
		p = p[l:]
	}

	v2 := make([]Vertex, len(vertices))
	reindices := make([]int32, len(indices))
	for i, idx := range indices {
		v2[i] = vertices[idx]
		reindices[idx] = int32(i)
	}
	copy(vertices, v2)

	// Update the indices in the triangles.
	for _, tri := range s.tin.Triangles {
		for i, idx := range tri {
			tri[i] = reindices[idx]
		}
	}
}