namespace LlmPoc.Llm;

public static class ReferenceFloatData
{
	public static float[] LoadFromMounted( string path, int expectedCount )
	{
		if ( !FileSystem.Mounted.FileExists( path ) )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Reference FP32 file '{path}' is missing from FileSystem.Mounted." );
		}

		byte[] bytes = FileSystem.Mounted.ReadAllBytes( path ).ToArray();
		if ( bytes.Length % sizeof( float ) != 0 )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Reference FP32 file '{path}' has {bytes.Length:N0} bytes, " +
				"which is not divisible by four." );
		}

		int count = bytes.Length / sizeof( float );
		if ( count != expectedCount )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Reference FP32 file '{path}' expected {expectedCount:N0} values " +
				$"({expectedCount * sizeof( float ):N0} bytes), found {count:N0} values " +
				$"({bytes.Length:N0} bytes)." );
		}

		float[] values = new float[count];
		for ( int index = 0; index < count; index++ )
		{
			int offset = index * sizeof( float );
			uint bits = (uint)(
				bytes[offset]
				| (bytes[offset + 1] << 8)
				| (bytes[offset + 2] << 16)
				| (bytes[offset + 3] << 24) );
			values[index] = BitConverter.Int32BitsToSingle( unchecked( (int)bits ) );
		}
		return values;
	}

	public static int ArgmaxFinite( ReadOnlySpan<float> values, string logicalName )
	{
		if ( values.Length == 0 )
		{
			throw new ArgumentException(
				$"[LLM:ERROR] {logicalName} cannot be argmaxed because it is empty." );
		}

		int bestIndex = 0;
		float bestValue = values[0];
		if ( !float.IsFinite( bestValue ) )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] {logicalName} contains non-finite value {bestValue} at index 0." );
		}

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