using System;
using System.Collections.Generic;
using GlinerPoc.Preprocessing;
using System.Diagnostics;
using Sandbox;

namespace GlinerPoc.Neural;

/// <summary>
/// Phase 5.37 — attention/layer parity harness. For each oracle case (short +
/// S=40) runs native preprocessing → embeddings → layer 0, comparing every
/// exported stage boundary against the pinned Python reference. Integer index
/// fixtures are compared exactly; float stages use the 1e-5/1e-4 policy.
/// </summary>
[Title( "GLiNER Layer Parity" )]
[Category( "GLiNER" )]
public sealed class GlinerAttentionParity : Component
{
	[Property]
	public bool RunOnStart { get; set; } = true;

	[Property]
	public GlinerPoc.Packaging.GlinerModelResource ModelResource { get; set; }

	private int _passed;
	private int _failed;

	protected override void OnStart()
	{
		if ( RunOnStart )
		{
			_ = RunAsync();
		}
	}

	private async Task RunAsync()
	{
		var sw = Stopwatch.StartNew();
		Log.Info( "[GLI:L05] layer parity harness start" );
		try
		{
			if ( ModelResource is null )
			{
				throw new InvalidOperationException( "[GLI:ERROR] No GlinerModelResource assigned." );
			}

			GlinerTokenizer tokenizer = null;
			GlinerModelWeights weights = null;
			double tokMs = 0.0;
			await Task.RunInThreadAsync( () =>
			{
				tokenizer = GlinerTokenizer.Load( ModelResource.TokenizerData.Bytes, out tokMs );
				weights = GlinerModelWeights.FromResource( ModelResource );
			} );
			Log.Info( $"[GLI:L05] ready tokenizer_ms={tokMs:N1}" );

			var processor = new GlinerProcessor( tokenizer );
			var layer = new GlinerDebertaLayer();
			var layerWeights = layer.LoadLayer0( weights );
			float[] normRel = GlinerMath.LayerNorm(
				weights.GetTensor( "encoder.encoder.rel_embeddings.weight" ),
				512, 384,
				weights.GetTensor( "encoder.encoder.LayerNorm.weight" ),
				weights.GetTensor( "encoder.encoder.LayerNorm.bias" ), 1e-7f );
			float[] embLnW = weights.GetTensor( "encoder.embeddings.LayerNorm.weight" );
			float[] embLnB = weights.GetTensor( "encoder.embeddings.LayerNorm.bias" );

			foreach ( var fixture in NeuralP5FixturesV2.All )
			{
				RunCase( processor, layer, layerWeights, weights, embLnW, embLnB, normRel, fixture );
			}

			Log.Info( $"[GLI:L05] parity complete passed={_passed} failed={_failed} " +
				$"total_ms={sw.ElapsedMilliseconds}" );
			Log.Info( _failed == 0 ? "[GLI:L05] ALL PASS" : $"[GLI:L05] FAILURES PRESENT ({_failed})" );
		}
		catch ( Exception error )
		{
			Log.Error( $"[GLI:L05] harness failure: {error.Message}" );
			Log.Error( $"[GLI:L05] parity complete passed={_passed} failed={_failed}" );
		}
	}

	private void RunCase(
		GlinerProcessor processor,
		GlinerDebertaLayer layer,
		GlinerDebertaLayer.LayerWeights w,
		GlinerModelWeights weights,
		float[] embLnW, float[] embLnB, float[] normRel,
		NeuralP5FixturesV2.Case f )
	{
		string name = f.Name;
		try
		{
			int S = f.SeqLen;

			// ---- native preprocessing must reproduce oracle IDs ----
			var encoded = processor.Encode( BuildRequest( name ), collectDiagnostics: false );
			bool idsOk = encoded.InputIds.Length == S;
			if ( idsOk )
			{
				for ( int i = 0; i < S; i++ )
				{
					idsOk &= encoded.InputIds[i] == f.Ids[i];
				}
			}
			Gate( name, "preprocessing_ids", idsOk, "exact" );

			// ---- embeddings: gather + LN ----
			float[] gathered = new float[S * 384];
			for ( int i = 0; i < S; i++ )
			{
				float[] row = weights.GetRows( "encoder.embeddings.word_embeddings.weight", encoded.InputIds[i], 1 );
				for ( int j = 0; j < 384; j++ )
				{
					gathered[i * 384 + j] = row[j];
				}
			}
			float[] layerInput = GlinerMath.LayerNorm( gathered, S, 384, embLnW, embLnB, 1e-7f );
			CompareF( name, "layer_input", layerInput, f.T( "layer_input" ) );

			// ---- relative positions (exact ints) ----
			int[,] rel = GlinerRelativePositions.BuildRelativePositions( S, S, 256, -1 );
			bool relOk = true;
			for ( int q = 0; q < S && relOk; q++ )
			{
				for ( int k = 0; k < S; k++ )
				{
					if ( rel[q, k] != f.RelIdx[q * S + k] )
					{
						relOk = false;
						Fail( name, $"rel_idx[{q},{k}] expected {f.RelIdx[q * S + k]} found {rel[q, k]}" );
					}
					if ( Math.Clamp( rel[q, k] + 256, 0, 511 ) != f.C2pIdx[q * S + k] )
					{
						relOk = false;
						Fail( name, $"c2p_idx[{q},{k}] mismatch" );
					}
					// Fixture documents the SOURCE's index matrix before the
					// transpose: p2c_pos[0,q,k] = clamp(-(q-k)+256) = k-q+256.
					if ( Math.Clamp( rel[q, k] * -1 + 256, 0, 511 ) != f.P2cIdx[q * S + k] )
					{
						relOk = false;
						Fail( name, $"p2c_idx[{q},{k}] mismatch" );
					}
				}
			}
			Gate( name, "rel_indices", relOk, "exact (rel + c2p + p2c)" );

			// ---- full layer forward with stage trace ----
			var trace = new Dictionary<string, float[]>();
			var sw = Stopwatch.StartNew();
			float[] layerOut = layer.Forward( w, layerInput, S, encoded.AttentionMask, rel, normRel, trace );
			double layerMs = sw.Elapsed.TotalMilliseconds;

			// ---- stage gates ----
			CompareF( name, "q_proj", trace["q_proj"], f.T( "q_proj" ) );
			CompareF( name, "k_proj", trace["k_proj"], f.T( "k_proj" ) );
			CompareF( name, "v_proj", trace["v_proj"], f.T( "v_proj" ) );
			CompareF( name, "q_layer", trace["q_layer"], f.T( "q_layer" ) );
			CompareF( name, "k_layer", trace["k_layer"], f.T( "k_layer" ) );
			CompareF( name, "v_layer", trace["v_layer"], f.T( "v_layer" ) );
			CompareF( name, "pos_q_rows", trace["pos_q_head_rows"], f.T( "pos_q_head_rows" ) );
			CompareF( name, "pos_k_rows", trace["pos_k_head_rows"], f.T( "pos_k_head_rows" ) );
			CompareF( name, "c2c", trace["c2c"], f.T( "c2c" ) );
			CompareF( name, "c2p_contrib", trace["c2p_contrib"], f.T( "c2p_contrib" ) );
			CompareF( name, "p2c_contrib", trace["p2c_contrib"], f.T( "p2c_contrib" ) );
			CompareF( name, "scores_premask", trace["scores_premask"], f.T( "scores_premask" ) );
			CompareF( name, "probs", trace["probs"], f.T( "probs" ) );
			CompareF( name, "context", trace["context"], f.T( "context" ) );
			CompareF( name, "attn_dense", trace["attn_dense"], f.T( "attn_dense" ) );
			CompareF( name, "attn_ln", trace["attn_ln"], f.T( "attn_ln" ) );
			CompareF( name, "ffn_pre", trace["ffn_pre"], f.T( "ffn_pre" ) );
			CompareF( name, "gelu", trace["gelu"], f.T( "gelu" ) );
			CompareF( name, "ffn_out", trace["ffn_out"], f.T( "ffn_out" ) );
			CompareF( name, "layer_output", layerOut, f.T( "layer_out" ) );
			Log.Info( $"[GLI:L05] PERF {name} layer_ms={layerMs:N2} seq={S}" );
		}
		catch ( Exception error )
		{
			_failed++;
			Log.Error( $"[GLI:L05] FAIL {name}: {error.Message}" );
		}
	}

	private static GlinerClassificationRequest BuildRequest( string name )
	{
		if ( name == "short" )
		{
			return new GlinerClassificationRequest( "Boss close. Ammo low. Cover near.",
				"Choose the best action.",
				new[] { new GlinerCandidate( "attack" ), new GlinerCandidate( "retreat" ) } );
		}
		return new GlinerClassificationRequest(
			"Health: 25%. Ammo: 2/10. Player visible: yes. Heal available: yes.",
			"Choose the best action.",
			new[] { new GlinerCandidate( "attack" ), new GlinerCandidate( "heal" ),
				new GlinerCandidate( "reload" ), new GlinerCandidate( "retreat" ) } );
	}

	private void CompareF( string caseName, string stage, float[] actual, float[] expected )
	{
		if ( actual is null )
		{
			throw new InvalidOperationException( $"{stage}: no trace recorded" );
		}
		if ( actual.Length != expected.Length )
		{
			throw new InvalidOperationException(
				$"{stage}: length {actual.Length} vs expected {expected.Length}" );
		}
		var m = GlinerMath.Compare( actual, expected );
		if ( !m.Pass )
		{
			throw new InvalidOperationException(
				$"{stage}: maxAbs={m.MaxAbs:0.###e-00} meanAbs={m.MeanAbs:0.###e-00} " +
				$"worst={m.WorstIndex} expected={expected[m.WorstIndex]} actual={actual[m.WorstIndex]}" );
		}
		Gate( caseName, stage, true,
			$"maxAbs={m.MaxAbs:0.###e-00} meanAbs={m.MeanAbs:0.###e-00} maxRel={m.MaxRel:0.###e-00}" );
	}

	private void Gate( string caseName, string gate, bool pass, string detail )
	{
		if ( pass )
		{
			_passed++;
			Log.Info( $"[GLI:L05] PASS {caseName}/{gate} {detail}" );
		}
		else
		{
			_failed++;
			Log.Error( $"[GLI:L05] FAIL {caseName}/{gate} {detail}" );
		}
	}

	private void Fail( string caseName, string detail )
	{
		Log.Error( $"[GLI:L05] detail FAIL {caseName}: {detail}" );
	}
}