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

namespace GlinerPoc.Neural;

/// <summary>
/// Phase 6.20 — model-facing classification result. No UI/service concerns.
/// "scores" are softmax candidate-distribution values, not calibrated
/// correctness probabilities.
/// </summary>
public sealed class GlinerDecisionResult
{
	public string[] Labels { get; }
	public float[] RawLogits { get; }
	public float[] ScaledLogits { get; }
	public float[] Scores { get; }
	public int SelectedIndex { get; }
	public string SelectedLabel => Labels[SelectedIndex];
	public int EncodedTokenCount { get; }

	public double PreprocessMs { get; internal set; }
	public double EmbeddingMs { get; internal set; }
	public double EncoderMs { get; internal set; }
	public double ClassifierMs { get; internal set; }
	public double TotalNeuralMs { get; internal set; }

	public GlinerDecisionResult( string[] labels, float[] rawLogits, float[] scaledLogits,
		float[] scores, int selectedIndex, int encodedTokenCount )
	{
		Labels = labels;
		RawLogits = rawLogits;
		ScaledLogits = scaledLogits;
		Scores = scores;
		SelectedIndex = selectedIndex;
		EncodedTokenCount = encodedTokenCount;
	}
}

/// <summary>
/// Phase 6.21/6.22 — top-level native GLiNER classification engine.
/// Pure managed inference: no Razor/scene/UI/NPC dependencies. Orchestrates
/// preprocessing → embeddings → 12-layer encoder → [L] marker gather →
/// classifier (384→768, ReLU, 768→1) → temperature → softmax → argmax.
/// The giant word embedding table stays on-demand (Phase 4 strategy);
/// layer weights use Strategy B: decoded once, retained (~85 MB).
/// </summary>
public sealed class GlinerClassificationEngine
{
	private readonly GlinerProcessor _processor;
	private readonly GlinerModelWeights _weights;
	private readonly GlinerDebertaLayer _layer;
	private readonly GlinerDebertaLayer.LayerWeights[] _layerWeights;
	private readonly float[] _normRel;
	private readonly float _temperature;
	private readonly float[] _classifier0W, _classifier0B, _classifier3W, _classifier3B;

	public double InitializationMs { get; private set; }

	public GlinerClassificationEngine( GlinerProcessor processor, GlinerModelWeights weights,
		float temperature, float[] embeddingLnW, float[] embeddingLnB,
		float[] relLnW, float[] relLnB )
	{
		_processor = processor;
		_weights = weights;
		_temperature = temperature;
		_layer = new GlinerDebertaLayer();

		var sw = System.Diagnostics.Stopwatch.StartNew();
		_normRel = GlinerMath.LayerNorm(
			weights.GetTensor( "encoder.encoder.rel_embeddings.weight" ), 512, 384,
			relLnW, relLnB, 1e-7f );
		EmbeddingLnW = embeddingLnW;
		EmbeddingLnB = embeddingLnB;

		_layerWeights = new GlinerDebertaLayer.LayerWeights[12];
		for ( int i = 0; i < 12; i++ )
		{
			_layerWeights[i] = LoadLayer( i );
		}

		_classifier0W = weights.GetTensor( "classifier.0.weight" );
		_classifier0B = weights.GetTensor( "classifier.0.bias" );
		_classifier3W = weights.GetTensor( "classifier.3.weight" );
		_classifier3B = weights.GetTensor( "classifier.3.bias" );
		InitializationMs = sw.ElapsedMilliseconds;
	}

	private GlinerDebertaLayer.LayerWeights LoadLayer( int i )
	{
		string p = $"encoder.encoder.layer.{i}.";
		return new GlinerDebertaLayer.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" ),
		};
	}

	public float[] EmbeddingLnW { get; }
	public float[] EmbeddingLnB { get; }

	/// <summary>Shared normalized relative embeddings (computed once, reused by all layers).</summary>
	public float[] NormalizedRelEmbeddings => _normRel;

	/// <summary>Full request path: preprocessing + inference.</summary>
	public GlinerDecisionResult Classify( GlinerPoc.Preprocessing.GlinerClassificationRequest request,
		Dictionary<int, float[]> layerTrace = null )
	{
		var sw = System.Diagnostics.Stopwatch.StartNew();
		var encoded = _processor.Encode( request, collectDiagnostics: false );
		double preMs = sw.Elapsed.TotalMilliseconds;
		var result = ClassifyEncoded( encoded, layerTrace );
		result.PreprocessMs = preMs;
		return result;
	}

		/// <summary>Encoded-request path: neural inference only (Phase 6.22).</summary>
	public GlinerDecisionResult ClassifyEncoded(
		GlinerPoc.Preprocessing.GlinerEncodedRequest encoded,
		Dictionary<int, float[]> layerTrace = null )
	{
		bool ok = TryClassifyEncoded( encoded, out var result, layerTrace );
		return ok ? result : null;
	}

	/// <summary>
	/// Phase 7.10 — cancellation-capable inference. `cancelled` is polled at
	/// coarse safe boundaries only (after embeddings and between encoder
	/// layers); when it returns true the method returns false with result null
	/// and all scratch becomes garbage. No threading APIs are used.
	/// </summary>
	public bool TryClassifyEncoded(
		GlinerPoc.Preprocessing.GlinerEncodedRequest encoded,
		out GlinerDecisionResult result,
		Dictionary<int, float[]> layerTrace = null,
		Func<bool> cancelled = null )
	{
		result = null;
		var sw = System.Diagnostics.Stopwatch.StartNew();
		int S = encoded.InputIds.Length;

		// ---- embeddings (on-demand rows, never the full table) ----
		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[] hidden = GlinerMath.LayerNorm( gathered, S, 384, EmbeddingLnW, EmbeddingLnB, 1e-7f );
		if ( cancelled is not null && cancelled() )
		{
			return false;
		}
		double embMs = sw.Elapsed.TotalMilliseconds;

		// ---- relative setup (shared across all layers) ----
		int[,] relPos = GlinerRelativePositions.BuildRelativePositions( S, S, 256, -1 );

		// ---- 12-layer encoder (one generic layer implementation) ----
		for ( int i = 0; i < 12; i++ )
		{
			hidden = _layer.Forward( _layerWeights[i], hidden, S, encoded.AttentionMask, relPos, _normRel );
			layerTrace?.Add( i, hidden );
			if ( cancelled is not null && cancelled() )
			{
				return false;
			}
		}
		double encMs = sw.Elapsed.TotalMilliseconds - embMs;

		// ---- marker gather (Phase 3 indices, candidate order) ----
		int n = encoded.ClassificationMarkerIndices.Length;
		float[] markerStates = new float[n * 384];
		for ( int c = 0; c < n; c++ )
		{
			int row = encoded.ClassificationMarkerIndices[c];
			for ( int j = 0; j < 384; j++ )
			{
				markerStates[c * 384 + j] = hidden[row * 384 + j];
			}
		}

		// ---- classifier: 384→768, ReLU, 768→1 ----
		float[] clsHidden = GlinerMath.Linear( markerStates, n, _classifier0W, _classifier0B, 384, 768 );
		float[] relu = GlinerMath.Relu( clsHidden );
		float[] rawLogits = GlinerMath.Linear( relu, n, _classifier3W, _classifier3B, 768, 1 );
		double clsMs = sw.Elapsed.TotalMilliseconds - embMs - encMs;

		// ---- temperature + softmax + argmax ----
		float[] scaled = new float[n];
		for ( int i = 0; i < n; i++ )
		{
			scaled[i] = rawLogits[i] / _temperature;
		}
		float[] scores = Softmax1D( scaled );
		int best = 0;
		for ( int i = 1; i < n; i++ )
		{
			if ( scores[i] > scores[best] )
			{
				best = i; // strict > matches torch.argmax first-index tie behaviour
			}
		}

		result = new GlinerDecisionResult(
			encoded.CandidateOrder as string[] ?? CopyLabels( encoded.CandidateOrder ),
			rawLogits, scaled, scores, best, S )
		{
			EmbeddingMs = embMs,
			EncoderMs = encMs,
			ClassifierMs = clsMs,
			TotalNeuralMs = sw.Elapsed.TotalMilliseconds,
		};
		return true;
	}

	private static string[] CopyLabels( System.Collections.Generic.IReadOnlyList<string> order )
	{
		var copy = new string[order.Count];
		for ( int i = 0; i < order.Count; i++ )
		{
			copy[i] = order[i];
		}
		return copy;
	}

	/// <summary>Stable 1-D softmax over the candidate dimension.</summary>
	private static float[] Softmax1D( float[] x )
	{
		float max = x[0];
		for ( int i = 1; i < x.Length; i++ )
		{
			if ( x[i] > max )
			{
				max = x[i];
			}
		}
		var y = new float[x.Length];
		float sum = 0f;
		for ( int i = 0; i < x.Length; i++ )
		{
			float e = MathF.Exp( x[i] - max );
			y[i] = e;
			sum += e;
		}
		float inv = 1f / sum;
		for ( int i = 0; i < x.Length; i++ )
		{
			y[i] *= inv;
		}
		return y;
	}
}