using System.Text;

namespace LlmPoc.Llm;

public sealed class TensorSummary
{
	public string Name { get; init; }
	public string Shape { get; init; }
	public long ElementCount { get; init; }
	public float Minimum { get; init; }
	public float Maximum { get; init; }
	public double Mean { get; init; }
	public double StandardDeviation { get; init; }
	public double Rms { get; init; }
	public int NaNCount { get; init; }
	public int PositiveInfinityCount { get; init; }
	public int NegativeInfinityCount { get; init; }
	public string FirstValues { get; init; }
	public string LastValues { get; init; }
	public ulong Checksum { get; init; }

	public bool IsFinite => NaNCount == 0 && PositiveInfinityCount == 0 && NegativeInfinityCount == 0;
	public long FiniteCount => ElementCount - NaNCount - PositiveInfinityCount - NegativeInfinityCount;

	public override string ToString()
	{
		return $"name={Name} shape={Shape} count={ElementCount:N0} " +
			$"min={Minimum:G9} max={Maximum:G9} mean={Mean:G12} " +
			$"std={StandardDeviation:G12} rms={Rms:G12} finite={FiniteCount:N0}/{ElementCount:N0} " +
			$"nan={NaNCount} +inf={PositiveInfinityCount} -inf={NegativeInfinityCount} " +
			$"first={FirstValues} last={LastValues} fnv1a64={Checksum:X16}";
	}
}

public sealed class NumericComparison
{
	public int ElementCount { get; init; }
	public double MaximumAbsoluteError { get; init; }
	public double MeanAbsoluteError { get; init; }
	public double MaximumRelativeError { get; init; }
	public int MaximumErrorIndex { get; init; }
	public float ExpectedAtMaximumError { get; init; }
	public float ActualAtMaximumError { get; init; }
	public int FirstFailingIndex { get; init; }
	public float ExpectedAtFirstFailure { get; init; }
	public float ActualAtFirstFailure { get; init; }
	public double AbsoluteTolerance { get; init; }
	public double RelativeTolerance { get; init; }
	public bool Passed { get; init; }

	public override string ToString()
	{
		string failure = FirstFailingIndex < 0
			? ""
			: $" first_fail_index={FirstFailingIndex} " +
				$"first_expected={ExpectedAtFirstFailure:G9} first_actual={ActualAtFirstFailure:G9}";
		return $"count={ElementCount:N0} max_abs={MaximumAbsoluteError:G12} " +
			$"mean_abs={MeanAbsoluteError:G12} max_rel={MaximumRelativeError:G12} " +
			$"max_index={MaximumErrorIndex} expected={ExpectedAtMaximumError:G9} " +
			$"actual={ActualAtMaximumError:G9} abs_tol={AbsoluteTolerance:G6} " +
			$"rel_tol={RelativeTolerance:G6}{failure} {(Passed ? "PASS" : "FAIL")}";
	}
}

public static class TensorDiagnostics
{
	private const ulong FnvOffsetBasis = 14695981039346656037UL;
	private const ulong FnvPrime = 1099511628211UL;

	public static TensorSummary Summarize( Tensor tensor, int edgeCount = 4 )
	{
		if ( tensor is null )
		{
			throw new ArgumentNullException( nameof( tensor ) );
		}
		return Summarize( tensor.Name, tensor.ShapeText, tensor.Data, edgeCount );
	}

	public static TensorSummary Summarize(
		string name,
		string shape,
		ReadOnlySpan<float> values,
		int edgeCount = 4 )
	{
		if ( edgeCount < 0 || edgeCount > 32 )
		{
			throw new ArgumentOutOfRangeException(
				nameof( edgeCount ), edgeCount, "[LLM:ERROR] Edge count must be between 0 and 32." );
		}
		if ( values.Length == 0 )
		{
			throw new ArgumentException( $"[LLM:ERROR] {name} cannot be summarized because it is empty." );
		}

		float minimum = float.PositiveInfinity;
		float maximum = float.NegativeInfinity;
		double sum = 0;
		double sumSquares = 0;
		int finiteCount = 0;
		int nanCount = 0;
		int positiveInfinityCount = 0;
		int negativeInfinityCount = 0;
		ulong checksum = FnvOffsetBasis;

		for ( int index = 0; index < values.Length; index++ )
		{
			float value = values[index];
			uint bits = unchecked( (uint)BitConverter.SingleToInt32Bits( value ) );
			checksum = HashByte( checksum, (byte)bits );
			checksum = HashByte( checksum, (byte)(bits >> 8) );
			checksum = HashByte( checksum, (byte)(bits >> 16) );
			checksum = HashByte( checksum, (byte)(bits >> 24) );

			if ( float.IsNaN( value ) )
			{
				nanCount++;
				continue;
			}
			if ( float.IsPositiveInfinity( value ) )
			{
				positiveInfinityCount++;
				continue;
			}
			if ( float.IsNegativeInfinity( value ) )
			{
				negativeInfinityCount++;
				continue;
			}

			minimum = Math.Min( minimum, value );
			maximum = Math.Max( maximum, value );
			sum += value;
			sumSquares += (double)value * value;
			finiteCount++;
		}

		double mean = finiteCount == 0 ? double.NaN : sum / finiteCount;
		double rms = finiteCount == 0 ? double.NaN : Math.Sqrt( sumSquares / finiteCount );
		double variance = finiteCount == 0 ? double.NaN : Math.Max( 0, sumSquares / finiteCount - mean * mean );

		return new TensorSummary
		{
			Name = name,
			Shape = shape,
			ElementCount = values.Length,
			Minimum = finiteCount == 0 ? float.NaN : minimum,
			Maximum = finiteCount == 0 ? float.NaN : maximum,
			Mean = mean,
			StandardDeviation = Math.Sqrt( variance ),
			Rms = rms,
			NaNCount = nanCount,
			PositiveInfinityCount = positiveInfinityCount,
			NegativeInfinityCount = negativeInfinityCount,
			FirstValues = FormatEdge( values, 0, Math.Min( edgeCount, values.Length ) ),
			LastValues = FormatEdge( values, Math.Max( 0, values.Length - edgeCount ), Math.Min( edgeCount, values.Length ) ),
			Checksum = checksum
		};
	}

	public static NumericComparison Compare(
		ReadOnlySpan<float> expected,
		ReadOnlySpan<float> actual,
		double absoluteTolerance = 1e-4,
		double relativeTolerance = 1e-3 )
	{
		if ( expected.Length != actual.Length )
		{
			throw new ArgumentException(
				$"[LLM:ERROR] Parity comparison shape mismatch: expected {expected.Length:N0} " +
				$"elements, actual {actual.Length:N0}." );
		}
		if ( expected.Length == 0 )
		{
			throw new ArgumentException( "[LLM:ERROR] Parity comparison cannot use empty arrays." );
		}
		if ( absoluteTolerance < 0 || relativeTolerance < 0 )
		{
			throw new ArgumentOutOfRangeException(
				nameof( absoluteTolerance ), "[LLM:ERROR] Parity tolerances cannot be negative." );
		}

		double maxAbsolute = -1;
		double maxRelative = 0;
		double absoluteSum = 0;
		int maxIndex = 0;
		int firstFailingIndex = -1;
		bool passed = true;

		for ( int index = 0; index < expected.Length; index++ )
		{
			float expectedValue = expected[index];
			float actualValue = actual[index];
			bool finite = float.IsFinite( expectedValue ) && float.IsFinite( actualValue );
			if ( !finite )
			{
				passed = false;
				if ( firstFailingIndex < 0 )
				{
					firstFailingIndex = index;
				}
			}

			double absolute = Math.Abs( (double)actualValue - expectedValue );
			double scale = Math.Max( Math.Abs( expectedValue ), 1e-12 );
			double relative = absolute / scale;
			absoluteSum += absolute;
			maxRelative = Math.Max( maxRelative, relative );
			if ( absolute > maxAbsolute )
			{
				maxAbsolute = absolute;
				maxIndex = index;
			}

			if ( absolute > absoluteTolerance + relativeTolerance * Math.Abs( expectedValue ) )
			{
				passed = false;
				if ( firstFailingIndex < 0 )
				{
					firstFailingIndex = index;
				}
			}
		}

		return new NumericComparison
		{
			ElementCount = expected.Length,
			MaximumAbsoluteError = maxAbsolute,
			MeanAbsoluteError = absoluteSum / expected.Length,
			MaximumRelativeError = maxRelative,
			MaximumErrorIndex = maxIndex,
			ExpectedAtMaximumError = expected[maxIndex],
			ActualAtMaximumError = actual[maxIndex],
			FirstFailingIndex = firstFailingIndex,
			ExpectedAtFirstFailure = firstFailingIndex < 0 ? 0 : expected[firstFailingIndex],
			ActualAtFirstFailure = firstFailingIndex < 0 ? 0 : actual[firstFailingIndex],
			AbsoluteTolerance = absoluteTolerance,
			RelativeTolerance = relativeTolerance,
			Passed = passed
		};
	}

	private static ulong HashByte( ulong hash, byte value )
	{
		return (hash ^ value) * FnvPrime;
	}

	private static string FormatEdge( ReadOnlySpan<float> values, int start, int count )
	{
		StringBuilder builder = new();
		builder.Append( '[' );
		for ( int index = 0; index < count; index++ )
		{
			if ( index > 0 )
			{
				builder.Append( ',' );
			}
			builder.Append( values[start + index].ToString( "G9" ) );
		}
		builder.Append( ']' );
		return builder.ToString();
	}
}