using System;

namespace GlinerPoc.Neural;

/// <summary>
/// Phase 4.17/4.18 — DeBERTa relative-position construction, exact port of
/// transformers 4.57.6 `build_relative_position` + `make_log_bucket_position`.
///
/// Pinned checkpoint: position_buckets = 256, max_relative_positions = -1.
/// Because bucketing only activates when bucket_size &gt; 0 AND max_position &gt; 0,
/// THIS CHECKPOINT USES RAW DIFFERENCES: matrix[q][k] = q - k, and the
/// attention path later offsets by `position_buckets` (256) to index the
/// 512-row rel embedding table: tableIndex = (q - k) + 256 ∈ [1, 511] for
/// sequences up to 256 tokens (guaranteed by the V1 length limit).
/// The log-bucket branch is implemented and unit-tested for completeness but
/// is not exercised by this checkpoint.
/// </summary>
public static class GlinerRelativePositions
{
	/// <summary>Matrix [querySize, keySize] with matrix[q][k] = q - k (exact ints).</summary>
	public static int[,] BuildRelativePositions( int querySize, int keySize, int bucketSize, int maxPosition )
	{
		if ( querySize <= 0 || keySize <= 0 )
		{
			throw new InvalidOperationException( "[GLI:ERROR] Relative positions require positive sizes." );
		}
		var result = new int[querySize, keySize];
		if ( bucketSize > 0 && maxPosition > 0 )
		{
			for ( int q = 0; q < querySize; q++ )
			{
				for ( int k = 0; k < keySize; k++ )
				{
					result[q, k] = LogBucket( q - k, bucketSize, maxPosition );
				}
			}
		}
		else
		{
			for ( int q = 0; q < querySize; q++ )
			{
				for ( int k = 0; k < keySize; k++ )
				{
					result[q, k] = q - k;
				}
			}
		}
		return result;
	}

	/// <summary>
	/// Exact port of make_log_bucket_position for a scalar. Torch computes:
	/// sign = sign(x); mid = bucket//2;
	/// abs_pos = mid-1 if (-mid &lt; x &lt; mid) else |x|;
	/// log_pos = ceil( log(abs_pos/mid) / log((max_position-1)/mid) * (mid-1) ) + mid;
	/// bucket = abs_pos &lt;= mid ? x : log_pos * sign.
	/// </summary>
	public static int LogBucket( int relativePos, int bucketSize, int maxPosition )
	{
		int sign = Math.Sign( relativePos );
		int mid = bucketSize / 2;
		int absPos = (relativePos < mid && relativePos > -mid) ? mid - 1 : Math.Abs( relativePos );
		double logPos = Math.Ceiling(
			Math.Log( absPos / (double)mid ) / Math.Log( (maxPosition - 1) / (double)mid ) * (mid - 1) ) + mid;
		bool useRaw = absPos <= mid;
		return useRaw ? relativePos : (int)(logPos * sign);
	}

	/// <summary>Table index for the rel embedding lookup used by attention: (q-k) + span.</summary>
	public static int ToTableIndex( int relativePos, int positionBucketSpan ) =>
		relativePos + positionBucketSpan;
}