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

namespace GlinerPoc.Neural;

/// <summary>
/// Phase 6.34 — full-model parity harness (rebuilt on engine 26.09.22):
/// native preprocessing → embeddings
/// → 12 layers (each compared against the oracle) → marker gather → classifier
/// stages → logits/scores/selected label; plus native-vs-Python decision
/// parity on a cross-domain Phase 1 benchmark subset.
/// </summary>
[Title( "GLiNER Full Model Parity" )]
[Category( "GLiNER" )]
public sealed class GlinerFullModelParity : 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:L06] full-model 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 );
			} );
			var processor = new GlinerProcessor( tokenizer );

			GlinerClassificationEngine engine = null;
			await Task.RunInThreadAsync( () =>
			{
				engine = new GlinerClassificationEngine(
					processor, weights, 1.0f,
					weights.GetTensor( "encoder.embeddings.LayerNorm.weight" ),
					weights.GetTensor( "encoder.embeddings.LayerNorm.bias" ),
					weights.GetTensor( "encoder.encoder.LayerNorm.weight" ),
					weights.GetTensor( "encoder.encoder.LayerNorm.bias" ) );
			} );
			Log.Info( $"[GLI:L06] engine ready init_ms={engine.InitializationMs:N1} " +
				"(decode-all strategy: 12 layers + classifier)" );

			// ---- detailed fixtures ----
			foreach ( var f in new[] { NeuralP6FixturesV2.Short, NeuralP6FixturesV2.S40 } )
			{
				RunFullCase( engine, f );
			}

			// ---- decision parity subset ----
			foreach ( var d in NeuralP6FixturesV2.Decisions )
			{
				RunDecision( engine, d );
			}

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

	private void RunFullCase( GlinerClassificationEngine engine, NeuralP6FixturesV2.FullCase f )
	{
		string name = f.Name;
		try
		{
			var request = BuildRequest( name );
			var trace = new Dictionary<int, float[]>();
			var result = engine.Classify( request, trace );

			// preprocessing IDs
			Gate( name, "preprocessing", result.EncodedTokenCount == f.SeqLen, "length" );

			// every layer boundary
			for ( int i = 0; i < 12; i++ )
			{
				CompareF( name, $"layer_{i}", trace[i], Decode( f.LayerB64[i] ) );
			}

			// marker states
			CompareF( name, "marker_states", GatherMarkers( trace[11], f.MarkerIdx ), Decode( f.MarkersB64 ) );

			// classifier stages
			CompareF( name, "raw_logits", result.RawLogits, f.Logits );
			CompareF( name, "scaled_logits", result.ScaledLogits, f.Logits ); // temperature 1.0
			CompareF( name, "softmax_scores", result.Scores, f.Probs );

			if ( result.SelectedIndex != f.Selected )
			{
				throw new InvalidOperationException(
					$"selected index {result.SelectedIndex} != oracle {f.Selected}" );
			}
			Gate( name, "selected_label", true,
				$"{result.SelectedLabel} (oracle {f.Labels[f.Selected]})" );
		}
		catch ( Exception error )
		{
			_failed++;
			Log.Error( $"[GLI:L06] FAIL {name}: {error.Message}" );
		}
	}

	private void RunDecision( GlinerClassificationEngine engine, NeuralP6FixturesV2.Decision d )
	{
		try
		{
			var candidates = new List<GlinerCandidate>( d.Labels.Length );
			foreach ( var l in d.Labels )
			{
				candidates.Add( new GlinerCandidate( l ) );
			}
			var request = new GlinerClassificationRequest( d.Context, d.Task, candidates );
			var result = engine.Classify( request );

			bool labelOk = result.SelectedLabel == d.OfficialLabel;
			var m = GlinerMath.Compare( result.Scores, d.Probs );
			bool scoreOk = m.Pass;
			Gate( d.Id, "decision_parity", labelOk && scoreOk,
				$"native={result.SelectedLabel} python={d.OfficialLabel} " +
				$"maxAbs={m.MaxAbs:0.###e-00} ms={result.TotalNeuralMs:N0}" );
		}
		catch ( Exception error )
		{
			_failed++;
			Log.Error( $"[GLI:L06] FAIL {d.Id}: {error.Message}" );
		}
	}

	private static float[] Decode( string b64 )
	{
		byte[] bytes = Convert.FromBase64String( b64 );
		return GlinerMath.DecodeF32( bytes, 0, bytes.Length / 4 );
	}

	private static float[] GatherMarkers( float[] hidden, int[] markerIdx )
	{
		var outBuf = new float[markerIdx.Length * 384];
		for ( int c = 0; c < markerIdx.Length; c++ )
		{
			for ( int j = 0; j < 384; j++ )
			{
				outBuf[c * 384 + j] = hidden[markerIdx[c] * 384 + j];
			}
		}
		return outBuf;
	}

	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, string b64 )
	{
		byte[] bytes = Convert.FromBase64String( b64 );
		var expected = GlinerMath.DecodeF32( bytes, 0, bytes.Length / 4 );
		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}" );
	}

	private void CompareF( string caseName, string stage, float[] actual, float[] expected )
	{
		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}" );
	}

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