using Sandbox.Diagnostics;

namespace LlmPoc.Llm;

public sealed class GptNeoLayerParityResult
{
	public int LayerIndex { get; init; }
	public NumericComparison FinalOutputComparison { get; init; }
	public double ValidationMilliseconds { get; init; }
}

public sealed class LocalMaskParityResult
{
	public int TotalCoordinates { get; init; }
	public int AllowedCoordinates { get; init; }
	public int MaskedCoordinates { get; init; }
	public int Mismatches { get; init; }
	public double ElapsedMilliseconds { get; init; }
}

public static class GptNeoLayerParity
{
	private const double StandardAbsoluteTolerance = 1.0e-5;
	private const double StandardRelativeTolerance = 1.0e-5;
	private const double SensitiveAbsoluteTolerance = 1.0e-6;

	public static GptNeoLayerParityResult Validate(
		GptNeoLayerForwardResult result,
		ForwardReferenceDocument reference,
		TinyStoriesConfig config )
	{
		if ( result is null ) throw new ArgumentNullException( nameof( result ) );
		if ( reference is null ) throw new ArgumentNullException( nameof( reference ) );
		if ( config is null ) throw new ArgumentNullException( nameof( config ) );
		if ( result.InputMutated )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Reusable layer {result.LayerIndex} mutated its input buffer." );
		}
		if ( result.AttentionType != config.AttentionLayers[result.LayerIndex] )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Layer {result.LayerIndex} execution type " +
				$"'{result.AttentionType}' disagrees with config " +
				$"'{config.AttentionLayers[result.LayerIndex]}'." );
		}

		FastTimer timer = FastTimer.StartNew();
		int layer = result.LayerIndex;
		StageValue[] stages =
		{
			new( $"layer{layer}_ln_1", result.Ln1 ),
			new( $"layer{layer}_q", result.Query ),
			new( $"layer{layer}_k", result.Key ),
			new( $"layer{layer}_v", result.Value ),
			new( $"layer{layer}_q_heads", result.QueryHeads ),
			new( $"layer{layer}_k_heads", result.KeyHeads ),
			new( $"layer{layer}_v_heads", result.ValueHeads ),
			new( $"layer{layer}_scores_scaled_unmasked", result.ScoresUnmasked ),
			new( $"layer{layer}_scores_masked_pre_softmax", result.ScoresMasked ),
			new( $"layer{layer}_attention_probs", result.AttentionProbabilities, sensitive: true ),
			new( $"layer{layer}_attention_context_heads", result.AttentionContextHeads, sensitive: true ),
			new( $"layer{layer}_attention_merged", result.AttentionMerged ),
			new( $"layer{layer}_attention_out_proj", result.AttentionOutput ),
			new( $"layer{layer}_attention_residual", result.AttentionResidual ),
			new( $"layer{layer}_ln_2", result.Ln2 ),
			new( $"layer{layer}_mlp_fc", result.MlpFc ),
			new( $"layer{layer}_mlp_gelu", result.MlpGelu, sensitive: true ),
			new( $"layer{layer}_mlp_proj", result.MlpProjection ),
			new( $"layer{layer}_output", result.Output )
		};

		ValidateHeadCopies( result );
		ValidateHeadMerge( result );
		ValidateFourTokenMask( result, reference, config );
		ValidateSoftmaxInvariants( result, config );

		NumericComparison finalComparison = null;
		foreach ( StageValue stageValue in stages )
		{
			ForwardReferenceStage stage = reference.GetRequiredStage( stageValue.StageName );
			ValidateStageMetadata( stage, stageValue.Value, result );
			float[] expected = ReferenceFloatData.LoadFromMounted(
				LlmPaths.ReferenceStage( stage.File ), stage.Elements );
			if ( stageValue.StageName == $"layer{layer}_attention_probs" && layer == 1 )
			{
				Tensor expectedTensor = new(
					$"python.{stageValue.StageName}", stage.Shape, expected );
				LlmLog.Info(
					"ATTN",
					"layer=1 targeted_softmax " +
					TinyStoriesForwardStages.DescribeAttentionSoftmaxRow(
						result.ScoresMasked,
						result.AttentionProbabilities,
						expectedTensor,
						head: 1,
						query: 3 ) );
			}
			if ( stageValue.StageName == $"layer{layer}_scores_scaled_unmasked" && layer == 1 )
			{
				int targetIndex = (1 * result.ScoresUnmasked.Shape[1] + 3) *
					result.ScoresUnmasked.Shape[2] + 1;
				LlmLog.Info(
					"ATTN",
					"layer=1 targeted_score " +
					TinyStoriesForwardStages.DescribeAttentionScore(
						result.QueryHeads,
						result.KeyHeads,
						result.ScoresUnmasked,
						head: 1,
						query: 3,
						key: 1,
						scale: 1.0f,
						expected: expected[targetIndex] ) );
			}
			double absoluteTolerance = stageValue.Sensitive
				? SensitiveAbsoluteTolerance
				: StandardAbsoluteTolerance;
			NumericComparison comparison = TensorDiagnostics.Compare(
				expected,
				stageValue.Value.Data,
				absoluteTolerance,
				StandardRelativeTolerance );
			TensorSummary summary = TensorDiagnostics.Summarize(
				stageValue.StageName,
				stageValue.Value.ShapeText,
				stageValue.Value.Data );
			string worst = FormatLogicalIndex(
				comparison.MaximumErrorIndex, stageValue.Value.Shape );
			LlmLog.Info(
				"PARITY",
				$"layer={layer} stage={stageValue.StageName} shape={stageValue.Value.ShapeText} " +
				$"finite={summary.FiniteCount}/{summary.ElementCount} " +
				$"min={summary.Minimum:G9} max={summary.Maximum:G9} mean={summary.Mean:G12} " +
				$"stddev={summary.StandardDeviation:G12} rms={summary.Rms:G12} " +
				$"maxAbs={comparison.MaximumAbsoluteError:G12} " +
				$"meanAbs={comparison.MeanAbsoluteError:G12} " +
				$"maxRel={comparison.MaximumRelativeError:G12} worst={worst} " +
				$"expected={comparison.ExpectedAtMaximumError:G9} " +
				$"actual={comparison.ActualAtMaximumError:G9} " +
				$"absTol={absoluteTolerance:G1} relTol={StandardRelativeTolerance:G1} " +
				$"{(comparison.Passed ? "PASS" : "FAIL")}" );
			if ( !comparison.Passed )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] Layer {layer} stage '{stageValue.StageName}' parity failed: " +
					$"{comparison}, worst={worst}." );
			}
			if ( stageValue.StageName == $"layer{layer}_output" )
			{
				finalComparison = comparison;
			}
		}

		return new GptNeoLayerParityResult
		{
			LayerIndex = layer,
			FinalOutputComparison = finalComparison ?? throw new InvalidOperationException(
				$"[LLM:ERROR] Layer {layer} final output comparison was not executed." ),
			ValidationMilliseconds = timer.ElapsedMilliSeconds
		};
	}

	public static GptNeoLayerParityResult ValidateOutputOnly(
		GptNeoLayerForwardResult result,
		ForwardReferenceDocument reference,
		TinyStoriesConfig config )
	{
		if ( result is null ) throw new ArgumentNullException( nameof( result ) );
		if ( reference is null ) throw new ArgumentNullException( nameof( reference ) );
		if ( config is null ) throw new ArgumentNullException( nameof( config ) );
		int layer = result.LayerIndex;
		if ( layer < 2 || layer >= config.LayerCount )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Output-only layer parity is reserved for layers 2-{config.LayerCount - 1}, " +
				$"found layer {layer}." );
		}
		if ( result.InputMutated )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Reusable layer {layer} mutated its input buffer." );
		}
		if ( result.AttentionType != config.AttentionLayers[layer] )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Layer {layer} execution type '{result.AttentionType}' " +
				$"disagrees with config '{config.AttentionLayers[layer]}'." );
		}

		FastTimer timer = FastTimer.StartNew();
		ForwardReferenceStage stage = reference.GetRequiredStage( $"layer{layer}_output" );
		ValidateStageMetadata( stage, result.Output, result );
		string expectedSource = $"layer{layer - 1}_output";
		string expectedNext = layer + 1 < config.LayerCount
			? $"layer{layer + 1}"
			: "final_layer_norm";
		if ( stage.SourceStage != expectedSource || stage.NextStage != expectedNext ||
			!stage.LayerComplete || !stage.ResidualAdditionApplied || stage.DropoutApplied ||
			stage.InputMutated )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Layer {layer} output reference expected source={expectedSource}, " +
				$"next={expectedNext}, complete=true, residual=true, dropout=false, " +
				$"input_mutated=false; found source={stage.SourceStage}, next={stage.NextStage}." );
		}

		float[] expected = ReferenceFloatData.LoadFromMounted(
			LlmPaths.ReferenceStage( stage.File ), stage.Elements );
		TensorSummary summary = TensorDiagnostics.Summarize( result.Output );
		if ( !summary.IsFinite )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Layer {layer} output contains NaN or Infinity." );
		}
		NumericComparison comparison = TensorDiagnostics.Compare(
			expected,
			result.Output.Data,
			StandardAbsoluteTolerance,
			StandardRelativeTolerance );
		string worst = FormatLogicalIndex(
			comparison.MaximumErrorIndex, result.Output.Shape );
		LlmLog.Info(
			"LAYER",
			$"layer={layer} type={result.AttentionType} shape={result.Output.ShapeText} " +
			$"finite={summary.FiniteCount}/{summary.ElementCount} min={summary.Minimum:G9} " +
			$"max={summary.Maximum:G9} mean={summary.Mean:G12} " +
			$"stddev={summary.StandardDeviation:G12} rms={summary.Rms:G12}" );
		LlmLog.Info(
			"PARITY",
			$"layer{layer}.output maxAbs={comparison.MaximumAbsoluteError:G12} " +
			$"meanAbs={comparison.MeanAbsoluteError:G12} " +
			$"maxRel={comparison.MaximumRelativeError:G12} worst={worst} " +
			$"expected={comparison.ExpectedAtMaximumError:G9} " +
			$"actual={comparison.ActualAtMaximumError:G9} " +
			$"absTol={StandardAbsoluteTolerance:G1} relTol={StandardRelativeTolerance:G1} " +
			$"{(comparison.Passed ? "PASS" : "FAIL")}" );
		if ( !comparison.Passed )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Layer {layer} complete output parity failed: " +
				$"{comparison}, worst={worst}." );
		}

		return new GptNeoLayerParityResult
		{
			LayerIndex = layer,
			FinalOutputComparison = comparison,
			ValidationMilliseconds = timer.ElapsedMilliSeconds
		};
	}

	public static LocalMaskParityResult ValidateLocalMaskReference(
		LocalAttentionMaskReferenceDocument metadata,
		TinyStoriesConfig config )
	{
		if ( metadata is null ) throw new ArgumentNullException( nameof( metadata ) );
		if ( config is null ) throw new ArgumentNullException( nameof( config ) );
		if ( metadata.Format != "SBOXLLM_LOCAL_ATTENTION_MASK_REFERENCE" ||
			metadata.Version != 1 || metadata.LayerIndex != 1 ||
			metadata.AttentionType != "local" || metadata.SequenceLength <= config.AttentionWindowSize ||
			metadata.QueryLength != metadata.SequenceLength ||
			metadata.KeyLength != metadata.SequenceLength ||
			metadata.WindowSize != config.AttentionWindowSize ||
			metadata.BinaryFile != "reference_layer1_local_mask_len260.bin" ||
			metadata.BinaryEncoding != "uint8 row-major [query,key], 1=allowed, 0=masked" ||
			metadata.AllowedRule != "key <= query and query-key < window_size" ||
			metadata.EarliestAllowedFormula != "max(0, query-window_size+1)" ||
			metadata.Status != "PASS" || metadata.Mismatches != 0 )
		{
			throw new InvalidOperationException(
				"[LLM:ERROR] Layer 1 local-mask reference metadata is structurally invalid." );
		}

		byte[] expected = FileSystem.Mounted.ReadAllBytes(
			LlmPaths.ReferenceLayer1LocalMaskLen260 ).ToArray();
		if ( expected.Length != metadata.BinaryBytes ||
			expected.Length != metadata.TotalCoordinates ||
			expected.Length != metadata.QueryLength * metadata.KeyLength )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Local-mask binary expected {metadata.TotalCoordinates} bytes, " +
				$"found {expected.Length}." );
		}

		FastTimer timer = FastTimer.StartNew();
		int allowedCount = 0;
		int maskedCount = 0;
		int mismatches = 0;
		int firstMismatchQuery = -1;
		int firstMismatchKey = -1;
		for ( int query = 0; query < metadata.QueryLength; query++ )
		{
			int firstAllowed = -1;
			int lastAllowed = -1;
			for ( int key = 0; key < metadata.KeyLength; key++ )
			{
				byte referenceValue = expected[query * metadata.KeyLength + key];
				if ( referenceValue != 0 && referenceValue != 1 )
				{
					throw new InvalidOperationException(
						$"[LLM:ERROR] Local-mask reference byte at [{query},{key}] " +
						$"must be 0 or 1, found {referenceValue}." );
				}
				bool expectedAllowed = referenceValue == 1;
				bool actualAllowed = TinyStoriesForwardStages.IsAttentionAllowed(
					"local",
					config.AttentionWindowSize,
					query,
					key,
					metadata.QueryLength,
					metadata.KeyLength );
				if ( expectedAllowed )
				{
					allowedCount++;
					if ( firstAllowed < 0 ) firstAllowed = key;
					lastAllowed = key;
				}
				else
				{
					maskedCount++;
				}
				if ( expectedAllowed != actualAllowed )
				{
					mismatches++;
					if ( firstMismatchQuery < 0 )
					{
						firstMismatchQuery = query;
						firstMismatchKey = key;
					}
				}
			}

			if ( metadata.FirstAllowedKey is null || metadata.LastAllowedKey is null ||
				metadata.FirstAllowedKey.Length != metadata.QueryLength ||
				metadata.LastAllowedKey.Length != metadata.QueryLength ||
				metadata.FirstAllowedKey[query] != firstAllowed ||
				metadata.LastAllowedKey[query] != lastAllowed )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] Local-mask boundary row {query} expected first/last " +
					$"[{metadata.FirstAllowedKey?[query]},{metadata.LastAllowedKey?[query]}], " +
					$"actual [{firstAllowed},{lastAllowed}]." );
			}
		}

		if ( allowedCount != metadata.AllowedCount || maskedCount != metadata.MaskedCount ||
			allowedCount + maskedCount != metadata.TotalCoordinates || mismatches != 0 )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Layer 1 local-mask parity failed: total={metadata.TotalCoordinates}, " +
				$"allowed={allowedCount}/{metadata.AllowedCount}, " +
				$"masked={maskedCount}/{metadata.MaskedCount}, mismatches={mismatches}, " +
				$"first_mismatch=[{firstMismatchQuery},{firstMismatchKey}]." );
		}

		foreach ( int query in new[] { 254, 255, 256, 257, 259 } )
		{
			LlmLog.Info(
				"MASK",
				$"layer=1 type=local seq=260 query={query} " +
				$"first_allowed={metadata.FirstAllowedKey[query]} " +
				$"last_allowed={metadata.LastAllowedKey[query]} window={metadata.WindowSize}" );
		}
		if ( metadata.BoundaryCoordinates is null || metadata.BoundaryCoordinates.Length == 0 )
		{
			throw new InvalidOperationException(
				"[LLM:ERROR] Local-mask reference has no boundary coordinates." );
		}
		foreach ( LocalAttentionMaskCoordinate coordinate in metadata.BoundaryCoordinates )
		{
			bool actual = TinyStoriesForwardStages.IsAttentionAllowed(
				"local",
				config.AttentionWindowSize,
				coordinate.Query,
				coordinate.Key,
				metadata.QueryLength,
				metadata.KeyLength );
			LlmLog.Trace(
				"MASK",
				$"layer=1 local_boundary query={coordinate.Query} key={coordinate.Key} " +
				$"python={coordinate.Allowed} csharp={actual} " +
				$"{(actual == coordinate.Allowed ? "PASS" : "FAIL")}" );
			if ( actual != coordinate.Allowed )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] Local-mask boundary mismatch at " +
					$"[{coordinate.Query},{coordinate.Key}]." );
			}
		}

		LlmLog.Info(
			"MASK",
			$"layer=1 type=local seq=260 coordinates={metadata.TotalCoordinates} " +
			$"allowed={allowedCount} masked={maskedCount} mismatches=0 " +
			$"elapsed_ms={timer.ElapsedMilliSeconds:N4} PASS" );
		return new LocalMaskParityResult
		{
			TotalCoordinates = metadata.TotalCoordinates,
			AllowedCoordinates = allowedCount,
			MaskedCoordinates = maskedCount,
			Mismatches = mismatches,
			ElapsedMilliseconds = timer.ElapsedMilliSeconds
		};
	}

	private static void ValidateStageMetadata(
		ForwardReferenceStage stage,
		Tensor actual,
		GptNeoLayerForwardResult result )
	{
		bool attentionTypeMatches = stage.AttentionType == result.AttentionType ||
			(result.LayerIndex == 0 && string.IsNullOrEmpty( stage.AttentionType ));
		if ( stage.File is null || stage.Dtype != "float32-le" ||
			stage.Shape is null || stage.Shape.Length != actual.Rank ||
			stage.Elements != actual.Data.Length || stage.Bytes != actual.Data.Length * sizeof( float ) ||
			stage.LayerIndex != result.LayerIndex || !attentionTypeMatches )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Layer {result.LayerIndex} reference metadata for " +
				$"'{stage.Stage}' disagrees with actual tensor {actual.ShapeText}." );
		}
		for ( int axis = 0; axis < actual.Rank; axis++ )
		{
			if ( stage.Shape[axis] != actual.Shape[axis] )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] Stage '{stage.Stage}' expected shape " +
					$"[{string.Join( ",", stage.Shape )}], actual {actual.ShapeText}." );
			}
		}
	}

	private static void ValidateHeadCopies( GptNeoLayerForwardResult result )
	{
		ValidateHeadCopy( result.Query, result.QueryHeads, result.LayerIndex, "Q" );
		ValidateHeadCopy( result.Key, result.KeyHeads, result.LayerIndex, "K" );
		ValidateHeadCopy( result.Value, result.ValueHeads, result.LayerIndex, "V" );
	}

	private static void ValidateHeadCopy(
		Tensor projection,
		Tensor heads,
		int layer,
		string label )
	{
		int headCount = heads.Shape[0];
		int sequenceLength = heads.Shape[1];
		int headDimension = heads.Shape[2];
		int mismatches = 0;
		for ( int head = 0; head < headCount; head++ )
		{
			for ( int token = 0; token < sequenceLength; token++ )
			{
				for ( int component = 0; component < headDimension; component++ )
				{
					int feature = head * headDimension + component;
					float source = projection.Data[token * projection.Shape[1] + feature];
					float destination = heads.Data[
						(head * sequenceLength + token) * headDimension + component];
					if ( BitConverter.SingleToInt32Bits( source ) !=
						BitConverter.SingleToInt32Bits( destination ) ) mismatches++;
				}
			}
		}
		LlmLog.Info(
			"ATTN",
			$"layer={layer} projection={label} head_split mappings={heads.Data.Length} " +
			$"exact={heads.Data.Length - mismatches} mismatches={mismatches} " +
			$"{(mismatches == 0 ? "PASS" : "FAIL")}" );
		if ( mismatches != 0 )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Layer {layer} {label} head split had {mismatches} copy mismatches." );
		}
	}

	private static void ValidateHeadMerge( GptNeoLayerForwardResult result )
	{
		Tensor context = result.AttentionContextHeads;
		Tensor merged = result.AttentionMerged;
		int headCount = context.Shape[0];
		int sequenceLength = context.Shape[1];
		int headDimension = context.Shape[2];
		int mismatches = 0;
		for ( int head = 0; head < headCount; head++ )
		{
			for ( int token = 0; token < sequenceLength; token++ )
			{
				for ( int component = 0; component < headDimension; component++ )
				{
					int feature = head * headDimension + component;
					float source = context.Data[
						(head * sequenceLength + token) * headDimension + component];
					float destination = merged.Data[token * merged.Shape[1] + feature];
					if ( BitConverter.SingleToInt32Bits( source ) !=
						BitConverter.SingleToInt32Bits( destination ) ) mismatches++;
				}
			}
		}
		LlmLog.Info(
			"ATTN",
			$"layer={result.LayerIndex} head_merge mappings={context.Data.Length} " +
			$"exact={context.Data.Length - mismatches} mismatches={mismatches} " +
			$"{(mismatches == 0 ? "PASS" : "FAIL")}" );
		if ( mismatches != 0 )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Layer {result.LayerIndex} head merge had " +
				$"{mismatches} copy mismatches." );
		}
	}

	private static void ValidateFourTokenMask(
		GptNeoLayerForwardResult result,
		ForwardReferenceDocument reference,
		TinyStoriesConfig config )
	{
		ForwardReferenceStage stage = reference.GetRequiredStage(
			$"layer{result.LayerIndex}_scores_masked_pre_softmax" );
		float[] expected = ReferenceFloatData.LoadFromMounted(
			LlmPaths.ReferenceStage( stage.File ), stage.Elements );
		int heads = result.ScoresMasked.Shape[0];
		int queryLength = result.ScoresMasked.Shape[1];
		int keyLength = result.ScoresMasked.Shape[2];
		int allowed = 0;
		int masked = 0;
		int mismatches = 0;
		for ( int head = 0; head < heads; head++ )
		{
			for ( int query = 0; query < queryLength; query++ )
			{
				for ( int key = 0; key < keyLength; key++ )
				{
					int index = (head * queryLength + query) * keyLength + key;
					bool isAllowed = TinyStoriesForwardStages.IsAttentionAllowed(
						result.AttentionType,
						config.AttentionWindowSize,
						query,
						key,
						queryLength,
						keyLength );
					if ( isAllowed )
					{
						allowed++;
						if ( BitConverter.SingleToInt32Bits( result.ScoresMasked.Data[index] ) !=
							BitConverter.SingleToInt32Bits( result.ScoresUnmasked.Data[index] ) )
						{
							mismatches++;
						}
					}
					else
					{
						masked++;
						if ( BitConverter.SingleToInt32Bits( result.ScoresMasked.Data[index] ) !=
							BitConverter.SingleToInt32Bits( GptNeoTransformerLayer.MaskedScoreSentinel ) ||
							BitConverter.SingleToInt32Bits( expected[index] ) !=
							BitConverter.SingleToInt32Bits( GptNeoTransformerLayer.MaskedScoreSentinel ) )
						{
							mismatches++;
						}
					}
				}
			}
		}
		if ( allowed != 160 || masked != 96 || mismatches != 0 )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Layer {result.LayerIndex} four-token " +
				$"{result.AttentionType} mask expected allowed=160 masked=96 mismatches=0, " +
				$"found allowed={allowed}, masked={masked}, mismatches={mismatches}." );
		}
		LlmLog.Info(
			"MASK",
			$"layer={result.LayerIndex} type={result.AttentionType} seq={queryLength} " +
			$"allowed={allowed} masked={masked} sentinel_exact={masked} mismatches=0 PASS " +
			$"(window_cutoff_exercised={queryLength > config.AttentionWindowSize})" );
	}

	private static void ValidateSoftmaxInvariants(
		GptNeoLayerForwardResult result,
		TinyStoriesConfig config )
	{
		Tensor probabilities = result.AttentionProbabilities;
		int heads = probabilities.Shape[0];
		int queries = probabilities.Shape[1];
		int keys = probabilities.Shape[2];
		double maximumRowError = 0;
		int maskedPositiveZero = 0;
		for ( int head = 0; head < heads; head++ )
		{
			for ( int query = 0; query < queries; query++ )
			{
				float sum = 0;
				for ( int key = 0; key < keys; key++ )
				{
					float probability = probabilities.Data[(head * queries + query) * keys + key];
					if ( !float.IsFinite( probability ) || probability < 0 )
					{
						throw new InvalidOperationException(
							$"[LLM:ERROR] Layer {result.LayerIndex} invalid probability " +
							$"at [{head},{query},{key}]: {probability}." );
					}
					bool allowed = TinyStoriesForwardStages.IsAttentionAllowed(
						result.AttentionType,
						config.AttentionWindowSize,
						query,
						key,
						queries,
						keys );
					if ( !allowed )
					{
						if ( BitConverter.SingleToInt32Bits( probability ) != 0 )
						{
							throw new InvalidOperationException(
								$"[LLM:ERROR] Layer {result.LayerIndex} masked probability " +
								$"[{head},{query},{key}] is not exact positive zero." );
						}
						maskedPositiveZero++;
					}
					sum += probability;
				}
				maximumRowError = Math.Max( maximumRowError, Math.Abs( (double)sum - 1.0 ) );
			}
		}
		if ( maximumRowError > 1.0e-6 )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Layer {result.LayerIndex} softmax max row-sum error " +
				$"{maximumRowError:G12} exceeds 1e-6." );
		}
		LlmLog.Info(
			"ATTN",
			$"layer={result.LayerIndex} softmax rows={heads * queries} " +
			$"max_row_sum_error={maximumRowError:G12} " +
			$"masked_positive_zero={maskedPositiveZero} PASS" );
	}

	private static string FormatLogicalIndex( int flatIndex, int[] shape )
	{
		int[] coordinates = new int[shape.Length];
		int remaining = flatIndex;
		for ( int axis = shape.Length - 1; axis >= 0; axis-- )
		{
			coordinates[axis] = remaining % shape[axis];
			remaining /= shape[axis];
		}
		return $"[{string.Join( ",", coordinates )}]";
	}

	private sealed class StageValue
	{
		public string StageName { get; }
		public Tensor Value { get; }
		public bool Sensitive { get; }

		public StageValue( string stageName, Tensor value, bool sensitive = false )
		{
			StageName = stageName;
			Value = value;
			Sensitive = sensitive;
		}
	}
}