using Sandbox.Diagnostics;

namespace LlmPoc.Llm;

public sealed class GreedyGenerationParityResult
{
	public GreedyGenerationResult Generation { get; init; }
	public NumericComparison LaterStepLogitComparison { get; init; }
	public int LaterStepLogitIndex { get; init; }
	public float MinimumMargin { get; init; }
	public int MinimumMarginStep { get; init; }
	public double ValidationHarnessMilliseconds { get; init; }
}

public static class GreedyGenerationParity
{
	private const double LogitAbsoluteTolerance = 5.0e-5;
	private const double LogitRelativeTolerance = 1.0e-5;
	private const double TopLogitAbsoluteTolerance = 5.0e-5;

	public static GreedyGenerationParityResult Validate(
		SboxLlmModel model,
		TinyStoriesConfig config,
		Gpt2ByteBpeTokenizer tokenizer,
		LlmReferenceData historicalReference,
		GreedyGenerationReferenceDocument reference )
	{
		if ( model is null ) throw new ArgumentNullException( nameof( model ) );
		if ( config is null ) throw new ArgumentNullException( nameof( config ) );
		if ( tokenizer is null ) throw new ArgumentNullException( nameof( tokenizer ) );
		if ( historicalReference is null )
			throw new ArgumentNullException( nameof( historicalReference ) );
		if ( reference is null ) throw new ArgumentNullException( nameof( reference ) );

		ValidateReference( config, tokenizer, historicalReference, reference );
		float[] laterExpected = ReferenceFloatData.LoadFromMounted(
			LlmPaths.ReferenceStage( reference.LaterFullLogitReference.File ),
			reference.LaterFullLogitReference.Elements );
		NumericComparison laterComparison = null;
		FastTimer validationTimer = FastTimer.StartNew();

		GreedyGenerationResult generation = TinyStoriesGreedyGenerator.Generate(
			model,
			config,
			tokenizer,
			reference.InputTokenIds,
			reference.MaxNewTokens,
			observation =>
			{
				GreedyGenerationStepResult actual = observation.Step;
				if ( actual.Step < 0 || actual.Step >= reference.Steps.Length )
				{
					throw new InvalidOperationException(
						$"[LLM:ERROR] Generation produced unexpected step {actual.Step}." );
				}
				GreedyGenerationReferenceStep expected = reference.Steps[actual.Step];
				ValidateStepContext( expected, observation.InputTokenIds );
				bool tokenPassed = actual.TokenId == expected.ExpectedNextTokenId;
				if ( !tokenPassed )
				{
					LogMismatch( tokenizer, expected, observation );
					throw new InvalidOperationException(
						$"[LLM:ERROR] Greedy generation first mismatch at step {actual.Step}: " +
						$"expected token {expected.ExpectedNextTokenId}, actual {actual.TokenId}. " +
						"No mismatching token was appended." );
				}

				if ( actual.DecodedToken != expected.DecodedToken )
				{
					throw new InvalidOperationException(
						$"[LLM:ERROR] Generation step {actual.Step} token {actual.TokenId} decoded " +
						$"as '{EscapeVisible( actual.DecodedToken )}', Python expected " +
						$"'{EscapeVisible( expected.DecodedToken )}'." );
				}
				ValidateTopFive( tokenizer, expected, actual );
				if ( actual.EosReached != expected.EosReached )
				{
					throw new InvalidOperationException(
						$"[LLM:ERROR] Generation step {actual.Step} EOS state expected " +
						$"{expected.EosReached}, actual {actual.EosReached}." );
				}

				if ( actual.Step == reference.LaterFullLogitReference.Step )
				{
					laterComparison = TensorDiagnostics.Compare(
						laterExpected,
						observation.Forward.Logits.Data,
						LogitAbsoluteTolerance,
						LogitRelativeTolerance );
					TensorSummary summary = TensorDiagnostics.Summarize(
						observation.Forward.Logits );
					LlmLog.Info(
						"PARITY",
						$"generation.step{actual.Step}.logits count={summary.ElementCount:N0} " +
						$"finite={summary.FiniteCount:N0}/{summary.ElementCount:N0} " +
						$"maxAbs={laterComparison.MaximumAbsoluteError:G12} " +
						$"meanAbs={laterComparison.MeanAbsoluteError:G12} " +
						$"maxRel={laterComparison.MaximumRelativeError:G12} " +
						$"worst_vocab={laterComparison.MaximumErrorIndex} " +
						$"expected={laterComparison.ExpectedAtMaximumError:G9} " +
						$"actual={laterComparison.ActualAtMaximumError:G9} " +
						$"absTol={LogitAbsoluteTolerance:G1} relTol={LogitRelativeTolerance:G1} " +
						$"{(laterComparison.Passed ? "PASS" : "FAIL")}" );
					if ( !laterComparison.Passed )
					{
						throw new InvalidOperationException(
							$"[LLM:ERROR] Later generation step {actual.Step} full-logit parity " +
							$"failed: {laterComparison}." );
					}
				}

				LlmLog.Info(
					"GEN",
					$"step={actual.Step} context={actual.InputSequenceLength} " +
					$"position={actual.NewTokenPosition} expected={expected.ExpectedNextTokenId} " +
					$"actual={actual.TokenId} token='{EscapeVisible( actual.DecodedToken )}' " +
					$"margin={actual.Top1Top2Margin:G9} forward_ms={actual.ForwardMilliseconds:N4} " +
					$"lm_head_ms={actual.LmHeadMilliseconds:N4} PASS" );
			} );

		RequireExactArray(
			"generated token IDs", reference.GeneratedTokenIds, generation.GeneratedTokenIds );
		RequireExactArray(
			"full generated sequence", reference.FullSequenceTokenIds,
			generation.FullSequenceTokenIds );
		if ( generation.GeneratedText != reference.GeneratedText )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Generated text expected '{EscapeVisible( reference.GeneratedText )}', " +
				$"actual '{EscapeVisible( generation.GeneratedText )}'." );
		}
		if ( generation.StopReason != reference.StopReason ||
			generation.EosReached != reference.EosReached )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Generation stop expected reason={reference.StopReason} " +
				$"eos={reference.EosReached}, actual reason={generation.StopReason} " +
				$"eos={generation.EosReached}." );
		}
		if ( laterComparison is null )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Later-step full-logit checkpoint " +
				$"{reference.LaterFullLogitReference.Step} did not execute." );
		}

		GreedyGenerationStepResult minimum = generation.Steps[0];
		for ( int index = 1; index < generation.Steps.Length; index++ )
		{
			if ( generation.Steps[index].Top1Top2Margin < minimum.Top1Top2Margin )
			{
				minimum = generation.Steps[index];
			}
		}
		if ( minimum.Step != reference.MinimumMarginStep ||
			Math.Abs( minimum.Top1Top2Margin - reference.MinimumTop1Top2Margin ) >
				TopLogitAbsoluteTolerance )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Minimum generation margin expected step=" +
				$"{reference.MinimumMarginStep} value={reference.MinimumTop1Top2Margin:G9}, " +
				$"actual step={minimum.Step} value={minimum.Top1Top2Margin:G9}." );
		}

		LlmLog.Info(
			"PARITY",
			$"greedy_sequence generated={generation.GeneratedTokenIds.Length} " +
			$"matched={generation.GeneratedTokenIds.Length} ids=" +
			$"[{string.Join( ",", generation.GeneratedTokenIds )}] " +
			$"text='{EscapeVisible( generation.GeneratedText )}' stop={generation.StopReason} PASS" );
		LlmLog.Info(
			"GEN",
			$"minimum_margin={minimum.Top1Top2Margin:G9} step={minimum.Step} " +
			$"first_token_ms={generation.Steps[0].StepMilliseconds:N4} " +
			$"last_token_ms={generation.Steps[^1].StepMilliseconds:N4} " +
			$"total_forward_ms={generation.TotalForwardMilliseconds:N4} " +
			$"total_generation_ms={generation.TotalGenerationMilliseconds:N4} " +
			$"average_ms_per_token=" +
			$"{generation.TotalGenerationMilliseconds / generation.Steps.Length:N4} " +
			$"tokens_per_second=" +
			$"{generation.Steps.Length * 1000.0 / generation.TotalGenerationMilliseconds:N4} " +
			"kv_cache=false PASS" );

		return new GreedyGenerationParityResult
		{
			Generation = generation,
			LaterStepLogitComparison = laterComparison,
			LaterStepLogitIndex = reference.LaterFullLogitReference.Step,
			MinimumMargin = minimum.Top1Top2Margin,
			MinimumMarginStep = minimum.Step,
			ValidationHarnessMilliseconds = validationTimer.ElapsedMilliSeconds
		};
	}

	private static void ValidateReference(
		TinyStoriesConfig config,
		Gpt2ByteBpeTokenizer tokenizer,
		LlmReferenceData historical,
		GreedyGenerationReferenceDocument reference )
	{
		RequireExactArray( "reference prompt IDs", historical.InputTokenIds, reference.InputTokenIds );
		RequireExactArray(
			"historical generated IDs", historical.GeneratedTokenIds,
			reference.GeneratedTokenIds );
		if ( historical.Prompt != reference.Prompt ||
			historical.GeneratedText != reference.GeneratedText )
		{
			throw new InvalidOperationException(
				"[LLM:ERROR] Compact generation reference differs from historical reference.json." );
		}
		if ( reference.EosTokenId != config.EosTokenId ||
			reference.BosTokenId != config.BosTokenId )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Generation special IDs expected BOS/EOS " +
				$"{config.BosTokenId}/{config.EosTokenId}, found " +
				$"{reference.BosTokenId}/{reference.EosTokenId}." );
		}
		if ( reference.DoSample || reference.BosAutomaticallyAdded ||
			reference.EosAutomaticallyAdded || !reference.RawLogitsUsedForArgmax ||
			!reference.CachedAndUncachedSequencesMatch || reference.CsharpBaselineUseCache ||
			reference.LogitsProcessors is null || reference.LogitsProcessors.Length != 0 ||
			reference.ArgmaxTieBreak != "first (lowest) vocabulary index" )
		{
			throw new InvalidOperationException(
				"[LLM:ERROR] Generation reference does not describe raw, uncached, " +
				"deterministic first-index argmax semantics." );
		}
		if ( reference.MaxNewTokens <= 0 || reference.Steps is null ||
			reference.Steps.Length != reference.GeneratedTokenIds.Length ||
			reference.Steps.Length > reference.MaxNewTokens )
		{
			throw new InvalidOperationException(
				"[LLM:ERROR] Generation reference step/token counts are inconsistent." );
		}
		if ( tokenizer.Decode( reference.GeneratedTokenIds ) != reference.GeneratedText )
		{
			throw new InvalidOperationException(
				"[LLM:ERROR] C# tokenizer cannot reproduce Python generated text from " +
				"the authoritative token sequence." );
		}
		if ( reference.LaterFullLogitReference is null ||
			reference.LaterFullLogitReference.Step < 1 ||
			reference.LaterFullLogitReference.Step >= reference.Steps.Length ||
			reference.LaterFullLogitReference.Elements != config.VocabularySize ||
			reference.LaterFullLogitReference.Bytes != config.VocabularySize * sizeof( float ) )
		{
			throw new InvalidOperationException(
				"[LLM:ERROR] Later generation full-logit reference metadata is invalid." );
		}
	}

	private static void ValidateStepContext(
		GreedyGenerationReferenceStep expected,
		int[] actualInput )
	{
		if ( expected.Step < 0 || expected.InputTokenIds is null ||
			expected.InputSequenceLength != expected.InputTokenIds.Length ||
			expected.NewTokenPosition != expected.InputSequenceLength )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Generation reference step {expected.Step} context metadata is invalid." );
		}
		RequireExactArray( $"generation step {expected.Step} input", expected.InputTokenIds, actualInput );
	}

	private static void ValidateTopFive(
		Gpt2ByteBpeTokenizer tokenizer,
		GreedyGenerationReferenceStep expected,
		GreedyGenerationStepResult actual )
	{
		if ( expected.TopFive is null || expected.TopFive.Length != 5 ||
			actual.TopFive is null || actual.TopFive.Length != 5 )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Generation step {actual.Step} requires five top-logit entries." );
		}
		for ( int rank = 0; rank < 5; rank++ )
		{
			ForwardReferenceTopLogit expectedRank = expected.TopFive[rank];
			LogitRank actualRank = actual.TopFive[rank];
			string decoded = tokenizer.Decode( new[] { actualRank.TokenId } );
			if ( expectedRank.Rank != actualRank.Rank ||
				expectedRank.TokenId != actualRank.TokenId ||
				expectedRank.DecodedToken != decoded ||
				Math.Abs( expectedRank.Logit - actualRank.Logit ) > TopLogitAbsoluteTolerance )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] Generation step {actual.Step} top-5 rank {rank + 1} " +
					$"expected token={expectedRank.TokenId} logit={expectedRank.Logit:G9} " +
					$"decoded='{EscapeVisible( expectedRank.DecodedToken )}', actual " +
					$"token={actualRank.TokenId} logit={actualRank.Logit:G9} " +
					$"decoded='{EscapeVisible( decoded )}'." );
			}
		}
	}

	private static void LogMismatch(
		Gpt2ByteBpeTokenizer tokenizer,
		GreedyGenerationReferenceStep expected,
		GreedyGenerationStepObservation observation )
	{
		GreedyGenerationStepResult actual = observation.Step;
		LlmLog.Error(
			$"Greedy mismatch step={actual.Step} context={actual.InputSequenceLength} " +
			$"input=[{string.Join( ",", observation.InputTokenIds )}] " +
			$"expected={expected.ExpectedNextTokenId} " +
			$"expected_piece='{EscapeVisible( expected.DecodedToken )}' actual={actual.TokenId} " +
			$"actual_piece='{EscapeVisible( tokenizer.Decode( new[] { actual.TokenId } ) )}' " +
			$"expected_margin={expected.Top1Top2Margin:G9} " +
			$"actual_margin={actual.Top1Top2Margin:G9}." );
		LlmLog.Error(
			$"Python top5=[{string.Join( ",", expected.TopFive.Select( item => $"{item.TokenId}:{item.Logit:G9}" ) )}] " +
			$"C# top5=[{string.Join( ",", actual.TopFive.Select( item => $"{item.TokenId}:{item.Logit:G9}" ) )}]. " +
			"Generate a full Python logit reference for this step only before changing tolerances." );
	}

	private static void RequireExactArray( string label, int[] expected, int[] actual )
	{
		if ( expected is null || actual is null || expected.Length != actual.Length )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] {label} expected length {expected?.Length ?? -1}, " +
				$"actual {actual?.Length ?? -1}." );
		}
		for ( int index = 0; index < expected.Length; index++ )
		{
			if ( expected[index] != actual[index] )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] {label} mismatch at index {index}: " +
					$"expected={expected[index]}, actual={actual[index]}." );
			}
		}
	}

	private static string EscapeVisible( string value )
	{
		return (value ?? "<null>")
			.Replace( "\\", "\\\\" )
			.Replace( "\r", "\\r" )
			.Replace( "\n", "\\n" )
			.Replace( "\t", "\\t" )
			.Replace( "'", "\\'" );
	}
}