using System;

namespace GlinerPoc.Neural;

/// <summary>
/// Phase 5 — one complete DeBERTa-v2 encoder layer, exact port of
/// transformers 4.57.6 DebertaV2Layer (DisentangledSelfAttention +
/// DebertaV2SelfOutput + DebertaV2Intermediate + DebertaV2Output) for THIS
/// checkpoint's active configuration: pos_att_type=[p2c, c2p],
/// share_att_key=true, relative_attention=true, position_buckets=256,
/// max_relative_positions=-1 (raw q−k relative positions, no bucketing).
///
/// Pinned behaviours (see ARCHITECTURE §17):
/// - transpose_for_scores: [S,384] → [heads=6, S, 64], layer[h,s,d] = x[s,h·64+d]
/// - scale_factor = 1 + c2p + p2c = 3; scale = √(64·3); the scale divides K
///   (and the c2p/p2c CONTRIBUTIONS, after the gather) — not Q, not the scores
/// - pos projections reuse the content query_proj/key_proj on the normalized
///   rel_embeddings rows [0:512]
/// - c2p gather index: clamp(rel[q][k]+256, 0, 511)
/// - p2c gather index: clamp(−rel[k][q]+256, 0, 511) = clamp(q−k+256, 0, 511)
///   (build_rpos returns relative_pos unchanged for equal seq dims)
/// - mask: [1,1,S,S] boolean outer product mask[q]·mask[k]; masked positions
///   filled with float32 minimum before softmax
/// - residuals are POST-norm: LN(dense(x) + residual)
/// All kernels out-of-place; inputs never mutated.
/// </summary>
public sealed class GlinerDebertaLayer
{
	public sealed class LayerWeights
	{
		public float[] QueryW, QueryB, KeyW, KeyB, ValueW, ValueB;
		public float[] AttnDenseW, AttnDenseB, AttnLnW, AttnLnB;
		public float[] FfnUpW, FfnUpB, FfnDownW, FfnDownB;
		public float[] OutputLnW, OutputLnB;
	}

	private readonly int _heads = 6;
	private readonly int _headDim = 64;
	private readonly int _hidden = 384;
	private readonly float _scale;
	private readonly int _attSpan = 256;

	public GlinerDebertaLayer()
	{
		_scale = MathF.Sqrt(_headDim * 3f); // scale_factor 1 + c2p + p2c
	}

	public LayerWeights LoadLayer0(GlinerModelWeights weights)
	{
		const string p = "encoder.encoder.layer.0.";
		return new LayerWeights
		{
			QueryW = weights.GetTensor(p + "attention.self.query_proj.weight"),
			QueryB = weights.GetTensor(p + "attention.self.query_proj.bias"),
			KeyW = weights.GetTensor(p + "attention.self.key_proj.weight"),
			KeyB = weights.GetTensor(p + "attention.self.key_proj.bias"),
			ValueW = weights.GetTensor(p + "attention.self.value_proj.weight"),
			ValueB = weights.GetTensor(p + "attention.self.value_proj.bias"),
			AttnDenseW = weights.GetTensor(p + "attention.output.dense.weight"),
			AttnDenseB = weights.GetTensor(p + "attention.output.dense.bias"),
			AttnLnW = weights.GetTensor(p + "attention.output.LayerNorm.weight"),
			AttnLnB = weights.GetTensor(p + "attention.output.LayerNorm.bias"),
			FfnUpW = weights.GetTensor(p + "intermediate.dense.weight"),
			FfnUpB = weights.GetTensor(p + "intermediate.dense.bias"),
			FfnDownW = weights.GetTensor(p + "output.dense.weight"),
			FfnDownB = weights.GetTensor(p + "output.dense.bias"),
			OutputLnW = weights.GetTensor(p + "output.LayerNorm.weight"),
			OutputLnB = weights.GetTensor(p + "output.LayerNorm.bias"),
		};
	}

	/// <summary>[S,384] → [heads,S,64]: layer[h,s,d] = x[s,h·64+d].</summary>
	public void TransposeForScores(float[] x, int seq, float[] outHeads)
	{
		for ( int s = 0; s < seq; s++ )
		{
			for ( int h = 0; h < _heads; h++ )
			{
				int src = s * _hidden + h * _headDim;
				int dst = h * seq * _headDim + s * _headDim;
				for ( int d = 0; d < _headDim; d++ )
				{
					outHeads[dst + d] = x[src + d];
				}
			}
		}
	}

	/// <summary>
	/// Full layer-0 forward. hiddenStates [S,384]; attentionMask [S] (0/1);
	/// relPos [S,S] (raw q−k); normRel [512,384] (normalized rel embeddings).
	/// When trace is non-null, named intermediate stages are stored for parity
	/// gates. Returns the new hidden states [S,384].
	/// </summary>
	public float[] Forward(
		LayerWeights w,
		float[] hiddenStates,
		int seq,
		byte[] attentionMask,
		int[,] relPos,
		float[] normRel,
		Dictionary<string, float[]> trace = null,
		Dictionary<string, int[,]> intTrace = null )
	{
		int hCount = _heads;
		int hd = _headDim;

		// ---- Q/K/V projections + head split --------------------------------
		float[] qProj = GlinerMath.Linear( hiddenStates, seq, w.QueryW, w.QueryB, _hidden, _hidden );
		float[] kProj = GlinerMath.Linear( hiddenStates, seq, w.KeyW, w.KeyB, _hidden, _hidden );
		float[] vProj = GlinerMath.Linear( hiddenStates, seq, w.ValueW, w.ValueB, _hidden, _hidden );
		trace?.Add( "q_proj", qProj );
		trace?.Add( "k_proj", kProj );
		trace?.Add( "v_proj", vProj );
		float[] qh = new float[hCount * seq * hd];
		float[] kh = new float[hCount * seq * hd];
		float[] vh = new float[hCount * seq * hd];
		TransposeForScores( qProj, seq, qh );
		TransposeForScores( kProj, seq, kh );
		TransposeForScores( vProj, seq, vh );
		trace?.Add( "q_layer", qh );
		trace?.Add( "k_layer", kh );
		trace?.Add( "v_layer", vh );

		// ---- pos projections on normalized rel embeddings [0:512] -----------
		int relRows = _attSpan * 2;
		float[] posQProj = GlinerMath.Linear( normRel, relRows, w.QueryW, w.QueryB, _hidden, _hidden );
		float[] posKProj = GlinerMath.Linear( normRel, relRows, w.KeyW, w.KeyB, _hidden, _hidden );
		float[] posQ = new float[hCount * relRows * hd];
		float[] posK = new float[hCount * relRows * hd];
		TransposeForScores( posQProj, relRows, posQ );
		TransposeForScores( posKProj, relRows, posK );
		if ( trace is not null )
		{
			trace["pos_q_head_rows"] = new float[hCount * 8 * hd];
			trace["pos_k_head_rows"] = new float[hCount * 8 * hd];
			for ( int h = 0; h < hCount; h++ )
			{
				for ( int r = 0; r < 8; r++ )
				{
					for ( int d = 0; d < hd; d++ )
					{
						trace["pos_q_head_rows"][h * 8 * hd + r * hd + d] = posQ[h * relRows * hd + r * hd + d];
						trace["pos_k_head_rows"][h * 8 * hd + r * hd + d] = posK[h * relRows * hd + r * hd + d];
					}
				}
			}
		}

		// ---- c2c: Q · (Kᵀ / scale) — scale divides K BEFORE the product -----
		float[] c2c = new float[hCount * seq * seq];
		for ( int h = 0; h < hCount; h++ )
		{
			for ( int qq = 0; qq < seq; qq++ )
			{
				for ( int kk = 0; kk < seq; kk++ )
				{
					float sum = 0f;
					int qBase = h * seq * hd + qq * hd;
					int kBase = h * seq * hd + kk * hd;
					for ( int d = 0; d < hd; d++ )
					{
						sum += qh[qBase + d] * (kh[kBase + d] / _scale);
					}
					c2c[h * seq * seq + qq * seq + kk] = sum;
				}
			}
		}

		// ---- c2p / p2c -------------------------------------------------------
		float[] c2pContrib = new float[hCount * seq * seq];
		float[] p2cContrib = new float[hCount * seq * seq];
		for ( int h = 0; h < hCount; h++ )
		{
			for ( int q = 0; q < seq; q++ )
			{
				for ( int k = 0; k < seq; k++ )
				{
					int idx = Math.Clamp( relPos[q, k] + _attSpan, 0, relRows - 1 );
					float c2p = 0f;
					int qBase = h * seq * hd + q * hd;
					int pkBase = h * relRows * hd + idx * hd;
					for ( int d = 0; d < hd; d++ )
					{
						c2p += qh[qBase + d] * posK[pkBase + d];
					}
					c2pContrib[h * seq * seq + q * seq + k] = c2p / _scale;

					// p2c: p2c_att[h, k, r] = k_layer·posQ; gather at
					// clamp(−rel[k][q]+256) = clamp(q−k+256); result transposed.
					int pidx = Math.Clamp( relPos[k, q] * -1 + _attSpan, 0, relRows - 1 );
					float p2c = 0f;
					int kBase = h * seq * hd + k * hd;
					int pqBase = h * relRows * hd + pidx * hd;
					for ( int d = 0; d < hd; d++ )
					{
						p2c += kh[kBase + d] * posQ[pqBase + d];
					}
					p2cContrib[h * seq * seq + q * seq + k] = p2c / _scale;
				}
			}
		}

		// ---- combine + mask + softmax ----------------------------------------
		// mask4[q,k] = mask[q]·mask[k] (oracle get_attention_mask outer product)
		var probs = new float[hCount * seq * seq];
		float[] c2cTrace = trace is not null ? new float[hCount * seq * seq] : null;
		float[] c2pTrace = trace is not null ? new float[hCount * seq * seq] : null;
		float[] p2cTrace = trace is not null ? new float[hCount * seq * seq] : null;
		for ( int h = 0; h < hCount; h++ )
		{
			for ( int q = 0; q < seq; q++ )
			{
				float max = float.NegativeInfinity;
				int base2 = h * seq * seq;
				for ( int k = 0; k < seq; k++ )
				{
					bool allowed = attentionMask[q] != 0 && attentionMask[k] != 0;
					float s;
					if ( allowed )
					{
						float c2cV = c2c[base2 + q * seq + k];
						float c2pV = c2pContrib[base2 + q * seq + k];
						float p2cV = p2cContrib[base2 + q * seq + k];
						s = c2cV + c2pV + p2cV;
						if ( trace is not null )
						{
							c2cTrace[base2 + q * seq + k] = c2cV;
							c2pTrace[base2 + q * seq + k] = c2pV;
							p2cTrace[base2 + q * seq + k] = p2cV;
						}
					}
					else
					{
						s = float.MinValue;
					}
					probs[base2 + q * seq + k] = s;
					if ( s > max )
					{
						max = s;
					}
				}
				float sum = 0f;
				for ( int k = 0; k < seq; k++ )
				{
					float e = MathF.Exp( probs[base2 + q * seq + k] - max );
					probs[base2 + q * seq + k] = e;
					sum += e;
				}
				float inv = 1f / sum;
				for ( int k = 0; k < seq; k++ )
				{
					probs[base2 + q * seq + k] *= inv;
				}
			}
		}
		if ( trace is not null )
		{
			trace["c2c"] = c2cTrace;
			trace["c2p_contrib"] = c2pTrace;
			trace["p2c_contrib"] = p2cTrace;
		}

		trace?.Add( "probs", probs );
		// ---- context + head merge ---------------------------------------------
		float[] context = new float[hCount * seq * hd];
		for ( int h = 0; h < hCount; h++ )
		{
			for ( int q = 0; q < seq; q++ )
			{
				for ( int d = 0; d < hd; d++ )
				{
					float sum = 0f;
					for ( int k = 0; k < seq; k++ )
					{
						sum += probs[h * seq * seq + q * seq + k] * vh[h * seq * hd + k * hd + d];
					}
					context[h * seq * hd + q * hd + d] = sum;
				}
			}
		}
		float[] merged = new float[seq * _hidden];
		for ( int s = 0; s < seq; s++ )
		{
			for ( int h = 0; h < hCount; h++ )
			{
				for ( int d = 0; d < hd; d++ )
				{
					merged[s * _hidden + h * hd + d] = context[h * seq * hd + s * hd + d];
				}
			}
		}

		// ---- context + head merge trace ---------------------------------------
		// "context" is the attention module OUTPUT (merged [S,384]), matching
		// the oracle capture; per-head layout is kept separately.
		trace?.Add( "context", merged );
		trace?.Add( "context_heads", context );

		// ---- attention output dense + residual + LN (POST-norm) ---------------
		float[] attnDense = GlinerMath.Linear( merged, seq, w.AttnDenseW, w.AttnDenseB, _hidden, _hidden );
		float[] attnSum = GlinerMath.Add( attnDense, hiddenStates );
		float[] postAttn = GlinerMath.LayerNorm( attnSum, seq, _hidden, w.AttnLnW, w.AttnLnB, 1e-7f );
		trace?.Add( "attn_dense", attnDense );
		trace?.Add( "attn_ln", postAttn );

		// ---- FFN ---------------------------------------------------------------
		float[] up = GlinerMath.Linear( postAttn, seq, w.FfnUpW, w.FfnUpB, _hidden, 1536 );
		float[] act = GlinerMath.GeluErf( up );
		float[] down = GlinerMath.Linear( act, seq, w.FfnDownW, w.FfnDownB, 1536, _hidden );
		float[] outSum = GlinerMath.Add( down, postAttn );
		trace?.Add( "ffn_pre", up );
		trace?.Add( "gelu", act );
		trace?.Add( "ffn_out", down );
		if ( trace is not null )
		{
			// pre-mask combined scores (c2c + rel terms) — reconstruct from
			// stored traces so the fixture gate can compare them directly
			var pre = new float[hCount * seq * seq];
			for ( int i = 0; i < pre.Length; i++ )
			{
				pre[i] = c2c[i] + c2pContrib[i] + p2cContrib[i];
			}
			trace["scores_premask"] = pre;
		}
		return GlinerMath.LayerNorm( outSum, seq, _hidden, w.OutputLnW, w.OutputLnB, 1e-7f );
	}
}