From f0a3a3dd44f854946923f9664d2325aaf1eceb36 Mon Sep 17 00:00:00 2001
From: WithoutPants <53250216+WithoutPants@users.noreply.github.com>
Date: Fri, 25 Nov 2022 08:20:23 +1100
Subject: [PATCH] Handle bad funscript at values (#3182)

---
 .../generator_interactive_heatmap_speed.go    | 47 +++++++++++++------
 ...task_generate_interactive_heatmap_speed.go |  2 +-
 ui/v2.5/src/docs/en/Changelog/v0180.md        |  1 +
 3 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/internal/manager/generator_interactive_heatmap_speed.go b/internal/manager/generator_interactive_heatmap_speed.go
index 9155f7b1f..c9b295983 100644
--- a/internal/manager/generator_interactive_heatmap_speed.go
+++ b/internal/manager/generator_interactive_heatmap_speed.go
@@ -11,16 +11,18 @@ import (
 	"sort"
 
 	"github.com/lucasb-eyer/go-colorful"
+	"github.com/stashapp/stash/pkg/logger"
 )
 
 type InteractiveHeatmapSpeedGenerator struct {
-	InteractiveSpeed int
-	Funscript        Script
-	FunscriptPath    string
-	HeatmapPath      string
-	Width            int
-	Height           int
-	NumSegments      int
+	sceneDurationMilli int64
+	InteractiveSpeed   int
+	Funscript          Script
+	FunscriptPath      string
+	HeatmapPath        string
+	Width              int
+	Height             int
+	NumSegments        int
 }
 
 type Script struct {
@@ -52,13 +54,14 @@ type GradientTable []struct {
 	Pos float64
 }
 
-func NewInteractiveHeatmapSpeedGenerator(funscriptPath string, heatmapPath string) *InteractiveHeatmapSpeedGenerator {
+func NewInteractiveHeatmapSpeedGenerator(funscriptPath string, heatmapPath string, sceneDuration float64) *InteractiveHeatmapSpeedGenerator {
 	return &InteractiveHeatmapSpeedGenerator{
-		FunscriptPath: funscriptPath,
-		HeatmapPath:   heatmapPath,
-		Width:         320,
-		Height:        15,
-		NumSegments:   150,
+		sceneDurationMilli: int64(sceneDuration * 1000),
+		FunscriptPath:      funscriptPath,
+		HeatmapPath:        heatmapPath,
+		Width:              320,
+		Height:             15,
+		NumSegments:        150,
 	}
 }
 
@@ -69,6 +72,10 @@ func (g *InteractiveHeatmapSpeedGenerator) Generate() error {
 		return err
 	}
 
+	if len(funscript.Actions) == 0 {
+		return fmt.Errorf("no valid actions in funscript")
+	}
+
 	g.Funscript = funscript
 	g.Funscript.UpdateIntensityAndSpeed()
 
@@ -102,14 +109,20 @@ func (g *InteractiveHeatmapSpeedGenerator) LoadFunscriptData(path string) (Scrip
 	sort.SliceStable(funscript.Actions, func(i, j int) bool { return funscript.Actions[i].At < funscript.Actions[j].At })
 
 	// trim actions with negative timestamps to avoid index range errors when generating heatmap
-
-	isValid := func(x int64) bool { return x >= 0 }
+	// #3181 - also trim actions that occur after the scene duration
+	loggedBadTimestamp := false
+	isValid := func(x int64) bool {
+		return x >= 0 && x < g.sceneDurationMilli
+	}
 
 	i := 0
 	for _, x := range funscript.Actions {
 		if isValid(x.At) {
 			funscript.Actions[i] = x
 			i++
+		} else if !loggedBadTimestamp {
+			loggedBadTimestamp = true
+			logger.Warnf("Invalid timestamp %d in %s: subsequent invalid timestamps will not be logged", x.At, path)
 		}
 	}
 
@@ -215,6 +228,10 @@ func (funscript Script) getGradientTable(numSegments int) GradientTable {
 
 	for _, a := range funscript.Actions {
 		segment := int(float64(a.At) / float64(maxts+1) * float64(numSegments))
+		// #3181 - sanity check. Clamp segment to numSegments-1
+		if segment >= numSegments {
+			segment = numSegments - 1
+		}
 		segments[segment].count++
 		segments[segment].intensity += int(a.Intensity)
 	}
diff --git a/internal/manager/task_generate_interactive_heatmap_speed.go b/internal/manager/task_generate_interactive_heatmap_speed.go
index 8bb80354c..414584a77 100644
--- a/internal/manager/task_generate_interactive_heatmap_speed.go
+++ b/internal/manager/task_generate_interactive_heatmap_speed.go
@@ -30,7 +30,7 @@ func (t *GenerateInteractiveHeatmapSpeedTask) Start(ctx context.Context) {
 	funscriptPath := video.GetFunscriptPath(t.Scene.Path)
 	heatmapPath := instance.Paths.Scene.GetInteractiveHeatmapPath(videoChecksum)
 
-	generator := NewInteractiveHeatmapSpeedGenerator(funscriptPath, heatmapPath)
+	generator := NewInteractiveHeatmapSpeedGenerator(funscriptPath, heatmapPath, t.Scene.Files.Primary().Duration)
 
 	err := generator.Generate()
 
diff --git a/ui/v2.5/src/docs/en/Changelog/v0180.md b/ui/v2.5/src/docs/en/Changelog/v0180.md
index 38ad878d8..404e3e61a 100644
--- a/ui/v2.5/src/docs/en/Changelog/v0180.md
+++ b/ui/v2.5/src/docs/en/Changelog/v0180.md
@@ -13,6 +13,7 @@
 * Added tag description filter criterion. ([#3011](https://github.com/stashapp/stash/pull/3011))
 
 ### 🎨 Improvements
+* Generated heatmaps now only show ranges within the duration of the scene. ([#3182](https://github.com/stashapp/stash/pull/3182))
 * Added File Modification Time to File Info panels. ([#3054](https://github.com/stashapp/stash/pull/3054))
 * Added counter to File Info tabs for objects with multiple files. ([#3054](https://github.com/stashapp/stash/pull/3054))
 * Added file count in Scene Duplicate Checker for scenes with multiple files. ([#3054](https://github.com/stashapp/stash/pull/3054))