namespace LlmPoc.Llm;

public sealed class LogitRank
{
	public int Rank { get; init; }
	public int TokenId { get; init; }
	public float Logit { get; init; }
}

/// <summary>
/// The model-specific final vocabulary projection. It intentionally computes only
/// the final sequence row required for next-token prediction.
/// </summary>
public static class TinyStoriesModelHead
{
	/// <summary>
	/// Returns the lowest vocabulary index containing the maximum finite logit,
	/// matching torch.argmax's documented tie behavior.
	/// </summary>
	public static int ArgmaxFinite( ReadOnlySpan<float> logits, string logicalName )
	{
		if ( logits.Length == 0 )
		{
			throw new ArgumentException(
				$"[LLM:ERROR] {logicalName} cannot be argmaxed because it is empty." );
		}

		int bestTokenId = 0;
		float bestValue = logits[0];
		if ( !float.IsFinite( bestValue ) )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] {logicalName} contains non-finite logit {bestValue} at token 0." );
		}

		for ( int tokenId = 1; tokenId < logits.Length; tokenId++ )
		{
			float value = logits[tokenId];
			if ( !float.IsFinite( value ) )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] {logicalName} contains non-finite logit {value} " +
					$"at token {tokenId}." );
			}
			if ( value > bestValue )
			{
				bestValue = value;
				bestTokenId = tokenId;
			}
		}
		return bestTokenId;
	}

	public static Tensor ProjectLastPositionNoBias(
		Tensor normalizedHiddenStates,
		Tensor weight,
		string outputName )
	{
		if ( normalizedHiddenStates is null )
		{
			throw new ArgumentNullException( nameof( normalizedHiddenStates ) );
		}
		if ( weight is null )
		{
			throw new ArgumentNullException( nameof( weight ) );
		}
		if ( string.IsNullOrWhiteSpace( outputName ) )
		{
			throw new ArgumentException(
				"[LLM:ERROR] LM-head output name cannot be empty.",
				nameof( outputName ) );
		}
		if ( normalizedHiddenStates.Rank != 2 || normalizedHiddenStates.Shape[0] <= 0 ||
			normalizedHiddenStates.Shape[1] <= 0 )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] {outputName} expected normalized hidden states " +
				$"[sequence,hidden], found {normalizedHiddenStates.ShapeText}." );
		}
		if ( weight.Rank != 2 || weight.Shape[0] <= 0 ||
			weight.Shape[1] != normalizedHiddenStates.Shape[1] )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] {outputName} LM-head weight '{weight.Name}' expected " +
				$"[vocabulary,{normalizedHiddenStates.Shape[1]}], found {weight.ShapeText}." );
		}

		int sequenceLength = normalizedHiddenStates.Shape[0];
		int hiddenSize = normalizedHiddenStates.Shape[1];
		int vocabularySize = weight.Shape[0];
		int inputRow = (sequenceLength - 1) * hiddenSize;
		float[] logits = new float[vocabularySize];
		for ( int vocabulary = 0; vocabulary < vocabularySize; vocabulary++ )
		{
			int weightRow = vocabulary * hiddenSize;
			float sum = 0;
			for ( int hidden = 0; hidden < hiddenSize; hidden++ )
			{
				sum = MathF.FusedMultiplyAdd(
					normalizedHiddenStates.Data[inputRow + hidden],
					weight.Data[weightRow + hidden],
					sum );
			}
			if ( !float.IsFinite( sum ) )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] {outputName} produced non-finite logit {sum} " +
					$"at vocabulary index {vocabulary}." );
			}
			logits[vocabulary] = sum;
		}

		return new Tensor( outputName, new[] { vocabularySize }, logits );
	}

	public static LogitRank[] TopKFinite(
		ReadOnlySpan<float> logits,
		int count,
		string logicalName )
	{
		if ( count <= 0 || count > logits.Length )
		{
			throw new ArgumentOutOfRangeException(
				nameof( count ), count,
				$"[LLM:ERROR] {logicalName} top-k count must be within [1,{logits.Length}]." );
		}

		int[] tokenIds = new int[count];
		float[] values = new float[count];
		for ( int rank = 0; rank < count; rank++ )
		{
			tokenIds[rank] = -1;
			values[rank] = float.NegativeInfinity;
		}

		for ( int tokenId = 0; tokenId < logits.Length; tokenId++ )
		{
			float value = logits[tokenId];
			if ( !float.IsFinite( value ) )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] {logicalName} contains non-finite logit {value} " +
					$"at token {tokenId}." );
			}
			for ( int rank = 0; rank < count; rank++ )
			{
				if ( value <= values[rank] ) continue;
				for ( int shift = count - 1; shift > rank; shift-- )
				{
					values[shift] = values[shift - 1];
					tokenIds[shift] = tokenIds[shift - 1];
				}
				values[rank] = value;
				tokenIds[rank] = tokenId;
				break;
			}
		}

		LogitRank[] result = new LogitRank[count];
		for ( int rank = 0; rank < count; rank++ )
		{
			result[rank] = new LogitRank
			{
				Rank = rank + 1,
				TokenId = tokenIds[rank],
				Logit = values[rank]
			};
		}
		return result;
	}
}