OpenNest/OpenNest.Engine/Fill/AngleCandidateBuilder.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using OpenNest.Engine.ML;
using OpenNest.Geometry;
using OpenNest.Math;

namespace OpenNest.Engine.Fill
{
    public class AngleCandidateBuilder
    {
        private readonly HashSet<double> knownGoodAngles = new();
        private readonly Func<bool> isPredictionAvailable;
        private readonly Func<Drawing, bool, PartFeatures> extractFeatures;
        private readonly Func<PartFeatures, double, double, List<double>> predictAngles;

        public AngleCandidateBuilder()
            : this(
                () => AnglePredictor.IsAvailable,
                FeatureExtractor.Extract,
                (features, width, height) => AnglePredictor.PredictAngles(features, width, height)
            )
        { }

        // Per-instance dependencies keep tests independent of the process-wide ONNX session.
        internal AngleCandidateBuilder(
            Func<bool> isPredictionAvailable,
            Func<Drawing, bool, PartFeatures> extractFeatures,
            Func<PartFeatures, double, double, List<double>> predictAngles
        )
        {
            this.isPredictionAvailable = isPredictionAvailable;
            this.extractFeatures = extractFeatures;
            this.predictAngles = predictAngles;
        }

        public bool ForceFullSweep { get; set; }

        public List<double> Build(NestItem item, ClassificationResult classification, Box workArea)
        {
            // User constraints always take precedence over classification.
            if (HasExplicitConstraints(item))
                return BuildFromConstraints(item);

            switch (classification.Type)
            {
                case PartType.Circle:
                    return new List<double> { 0 };

                case PartType.Rectangle:
                    return new List<double>
                    {
                        classification.PrimaryAngle,
                        classification.PrimaryAngle + Angle.HalfPI,
                    };

                default:
                    return BuildIrregularAngles(item, classification.PrimaryAngle, workArea);
            }
        }

        private static bool HasExplicitConstraints(NestItem item)
        {
            // Default NestConstraints: Start=0, End=0. Both zero = no constraints.
            return !(item.RotationStart.IsEqualTo(0) && item.RotationEnd.IsEqualTo(0));
        }

        private static List<double> BuildFromConstraints(NestItem item)
        {
            var angles = new List<double>();
            var step = item.StepAngle > Tolerance.Epsilon ? item.StepAngle : Angle.ToRadians(5);

            for (var a = item.RotationStart; a <= item.RotationEnd + Tolerance.Epsilon; a += step)
            {
                if (!ContainsAngle(angles, a))
                    angles.Add(a);
            }

            if (angles.Count == 0)
                angles.Add(item.RotationStart);

            return angles;
        }

        private List<double> BuildIrregularAngles(NestItem item, double primaryAngle, Box workArea)
        {
            var baseAngles = new[] { primaryAngle, primaryAngle + Angle.HalfPI };

            if (knownGoodAngles.Count > 0 && !ForceFullSweep)
                return BuildPrunedList(baseAngles);

            var angles = new List<double>(baseAngles);

            // Full 5-degree sweep for irregular parts.
            AddSweepAngles(angles);

            // ML prediction complements the sweep when available.
            angles = ApplyMlPrediction(item, workArea, baseAngles, angles);

            return angles;
        }

        private static void AddSweepAngles(List<double> angles)
        {
            var step = Angle.ToRadians(5);
            for (var a = 0.0; a < System.Math.PI; a += step)
            {
                if (!ContainsAngle(angles, a))
                    angles.Add(a);
            }
        }

        private List<double> ApplyMlPrediction(
            NestItem item,
            Box workArea,
            double[] baseAngles,
            List<double> fallback
        )
        {
            if (!isPredictionAvailable())
                return fallback;

            // Inference needs only scalar features, never the training bitmap.
            var features = extractFeatures(item.Drawing, false);
            if (features == null)
                return fallback;

            var predicted = predictAngles(features, workArea.Width, workArea.Length);
            if (predicted == null)
                return fallback;

            var mlAngles = new List<double>(predicted);
            foreach (var b in baseAngles)
            {
                if (!ContainsAngle(mlAngles, b))
                    mlAngles.Add(b);
            }

            // Merge ML angles into the existing sweep so both contribute.
            foreach (var a in fallback)
            {
                if (!ContainsAngle(mlAngles, a))
                    mlAngles.Add(a);
            }

            Debug.WriteLine(
                $"[AngleCandidateBuilder] ML: {fallback.Count} sweep + {predicted.Count} predicted = {mlAngles.Count} total"
            );
            return mlAngles;
        }

        private List<double> BuildPrunedList(double[] baseAngles)
        {
            var pruned = new List<double>(baseAngles);
            foreach (var a in knownGoodAngles)
            {
                if (!ContainsAngle(pruned, a))
                    pruned.Add(a);
            }

            Debug.WriteLine(
                $"[AngleCandidateBuilder] Pruned to {pruned.Count} angles (known-good)"
            );
            return pruned;
        }

        private static bool ContainsAngle(List<double> angles, double angle)
        {
            return angles.Any(existing => existing.IsEqualTo(angle));
        }

        public void RecordProductive(List<AngleResult> angleResults)
        {
            foreach (var ar in angleResults)
            {
                if (ar.PartCount > 0)
                    knownGoodAngles.Add(Angle.ToRadians(ar.AngleDeg));
            }
        }
    }
}