using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Threading.Tasks;
using Sandbox;

namespace GlinerPoc.Gpu;

/// <summary>
/// Phase 8D.2/8D.3 — model-level GPU weight owner. Validates the tensor
/// schema for ALL 12 encoder layers + classifier (mirroring CPU Phase 6),
/// uploads everything ONCE to persistent GpuBuffers, and prepares the shared
/// normalized relative table plus LAYER-SPECIFIC posQ/posK projections
/// (share_att_key: each layer projects normRel with ITS OWN Q/K weights).
/// The word-embedding table is NOT GPU-resident (P8B strategy B: the CPU
/// gathers the needed rows per request — ≤384 KiB upload).
/// </summary>
public sealed class GlinerGpuModelWeights : IDisposable
{
	private readonly GlinerGpuRuntime _rt;
	public GlinerGpuLayer.LayerGpuWeights[] Layers { get; } = new GlinerGpuLayer.LayerGpuWeights[12];
	public GlinerGpuBufferDesc Cls0W, Cls0B, Cls3W, Cls3B;
	public GlinerGpuBufferDesc RelLnW, RelLnB;
	/// <summary>Shared normalized relative table [512,384] (GPU-prepared).</summary>
	public GpuBuffer<float> NormRel;
	/// <summary>Per-layer positional states [512,384] each (layer-specific projections).</summary>
	public GpuBuffer<float>[] PosQ { get; } = new GpuBuffer<float>[12];
	public GpuBuffer<float>[] PosK { get; } = new GpuBuffer<float>[12];

	public long WeightBytes { get; private set; }
	public int BufferCount { get; private set; }
	public double UploadMs { get; private set; }
	public double PrepareMs { get; private set; }

	public GlinerGpuModelWeights( GlinerGpuRuntime rt ) => _rt = rt;

	private static readonly (string Suffix, int Rank, long Elements)[] LayerSchema = new[]
	{
		("attention.self.query_proj.weight", 2, 384L * 384),
		("attention.self.query_proj.bias", 1, 384),
		("attention.self.key_proj.weight", 2, 384L * 384),
		("attention.self.key_proj.bias", 1, 384),
		("attention.self.value_proj.weight", 2, 384L * 384),
		("attention.self.value_proj.bias", 1, 384),
		("attention.output.dense.weight", 2, 384L * 384),
		("attention.output.dense.bias", 1, 384),
		("attention.output.LayerNorm.weight", 1, 384),
		("attention.output.LayerNorm.bias", 1, 384),
		("intermediate.dense.weight", 2, 1536L * 384),
		("intermediate.dense.bias", 1, 1536),
		("output.dense.weight", 2, 384L * 1536),
		("output.dense.bias", 1, 384),
		("output.LayerNorm.weight", 1, 384),
		("output.LayerNorm.bias", 1, 384),
	};

	/// <summary>8D.1 schema validation + one-time upload of every model tensor.</summary>
	public void LoadAndUpload( GlinerPoc.Neural.GlinerModelWeights cpu )
	{
		var swTotal = Stopwatch.StartNew();
		var manifest = cpu.Manifest;

		void Validate( string name, int rank, long elements )
		{
			var rec = manifest.GetRequiredTensor( name );
			if ( rec.Rank != rank || rec.ElementCount != elements )
			{
				throw new InvalidOperationException(
					$"[GLI:ERROR] tensor '{name}' shape mismatch: rank {rec.Rank}({string.Join(',', rec.Shape)}) " +
					$"elements {rec.ElementCount}, expected rank {rank} elements {elements}" );
			}
		}

		// schema validation: all 12 layers + classifier + rel tensors
		for ( int i = 0; i < 12; i++ )
		{
			string p = $"encoder.encoder.layer.{i}.";
			foreach ( var (suffix, rank, elements) in LayerSchema )
			{
				Validate( p + suffix, rank, elements );
			}
		}
		Validate( "classifier.0.weight", 2, 768L * 384 );
		Validate( "classifier.0.bias", 1, 768 );
		Validate( "classifier.3.weight", 2, 768 );
		Validate( "classifier.3.bias", 1, 1 );
		Validate( "encoder.encoder.rel_embeddings.weight", 2, 512L * 384 );
		Validate( "encoder.encoder.LayerNorm.weight", 1, 384 );
		Validate( "encoder.encoder.LayerNorm.bias", 1, 384 );

		// upload (one-time; CPU-side staging comes from the CPU decode cache)
		var swUpload = Stopwatch.StartNew();
		for ( int i = 0; i < 12; i++ )
		{
			string p = $"encoder.encoder.layer.{i}.";
			var lw = new GlinerGpuLayer.LayerGpuWeights();
			(lw.QueryW, _) = _rt.UploadPersistent( $"l{i}_qw", cpu.GetTensor( p + LayerSchema[0].Suffix ), 384, 384, $"l{i} Q" );
			(lw.QueryB, _) = _rt.UploadPersistent( $"l{i}_qb", cpu.GetTensor( p + LayerSchema[1].Suffix ), 1, 384, $"l{i} Qb" );
			(lw.KeyW, _) = _rt.UploadPersistent( $"l{i}_kw", cpu.GetTensor( p + LayerSchema[2].Suffix ), 384, 384, $"l{i} K" );
			(lw.KeyB, _) = _rt.UploadPersistent( $"l{i}_kb", cpu.GetTensor( p + LayerSchema[3].Suffix ), 1, 384, $"l{i} Kb" );
			(lw.ValueW, _) = _rt.UploadPersistent( $"l{i}_vw", cpu.GetTensor( p + LayerSchema[4].Suffix ), 384, 384, $"l{i} V" );
			(lw.ValueB, _) = _rt.UploadPersistent( $"l{i}_vb", cpu.GetTensor( p + LayerSchema[5].Suffix ), 1, 384, $"l{i} Vb" );
			(lw.AttnDenseW, _) = _rt.UploadPersistent( $"l{i}_adw", cpu.GetTensor( p + LayerSchema[6].Suffix ), 384, 384, $"l{i} attn" );
			(lw.AttnDenseB, _) = _rt.UploadPersistent( $"l{i}_adb", cpu.GetTensor( p + LayerSchema[7].Suffix ), 1, 384, $"l{i} attn b" );
			(lw.AttnLnW, _) = _rt.UploadPersistent( $"l{i}_alnw", cpu.GetTensor( p + LayerSchema[8].Suffix ), 1, 384, $"l{i} attn LN" );
			(lw.AttnLnB, _) = _rt.UploadPersistent( $"l{i}_alnb", cpu.GetTensor( p + LayerSchema[9].Suffix ), 1, 384, $"l{i} attn LN b" );
			(lw.FfnUpW, _) = _rt.UploadPersistent( $"l{i}_fuw", cpu.GetTensor( p + LayerSchema[10].Suffix ), 1536, 384, $"l{i} FFN up" );
			(lw.FfnUpB, _) = _rt.UploadPersistent( $"l{i}_fub", cpu.GetTensor( p + LayerSchema[11].Suffix ), 1, 1536, $"l{i} FFN up b" );
			(lw.FfnDownW, _) = _rt.UploadPersistent( $"l{i}_fdw", cpu.GetTensor( p + LayerSchema[12].Suffix ), 384, 1536, $"l{i} FFN dn" );
			(lw.FfnDownB, _) = _rt.UploadPersistent( $"l{i}_fdb", cpu.GetTensor( p + LayerSchema[13].Suffix ), 1, 384, $"l{i} FFN dn b" );
			(lw.OutLnW, _) = _rt.UploadPersistent( $"l{i}_olnw", cpu.GetTensor( p + LayerSchema[14].Suffix ), 1, 384, $"l{i} out LN" );
			(lw.OutLnB, _) = _rt.UploadPersistent( $"l{i}_olnb", cpu.GetTensor( p + LayerSchema[15].Suffix ), 1, 384, $"l{i} out LN b" );
			Layers[i] = lw;
		}
		(Cls0W, _) = _rt.UploadPersistent( "cls0w", cpu.GetTensor( "classifier.0.weight" ), 768, 384, "classifier.0" );
		(Cls0B, _) = _rt.UploadPersistent( "cls0b", cpu.GetTensor( "classifier.0.bias" ), 1, 768, "classifier.0 b" );
		(Cls3W, _) = _rt.UploadPersistent( "cls3w", cpu.GetTensor( "classifier.3.weight" ), 1, 768, "classifier.3" );
		(Cls3B, _) = _rt.UploadPersistent( "cls3b", cpu.GetTensor( "classifier.3.bias" ), 1, 1, "classifier.3 b" );
		(RelLnW, _) = _rt.UploadPersistent( "rlnw", cpu.GetTensor( "encoder.encoder.LayerNorm.weight" ), 1, 384, "rel LN" );
		(RelLnB, _) = _rt.UploadPersistent( "rlnb", cpu.GetTensor( "encoder.encoder.LayerNorm.bias" ), 1, 384, "rel LN b" );
		swUpload.Stop();
		UploadMs = swUpload.Elapsed.TotalMilliseconds;

		// GPU-side relative preparation: shared normRel + 12 LAYER-SPECIFIC pos
		// projections (one executor job)
		var rawRelDesc = _rt.UploadPersistent( "relraw", cpu.GetTensor( "encoder.encoder.rel_embeddings.weight" ), 512, 384, "rel raw" ).Desc;
		var swPrep = Stopwatch.StartNew();
		NormRel = new GpuBuffer<float>( 512 * 384 );
		for ( int i = 0; i < 12; i++ )
		{
			PosQ[i] = new GpuBuffer<float>( 512 * 384 );
			PosK[i] = new GpuBuffer<float>( 512 * 384 );
		}
		_rt.Executor.Submit( "model_rel_prep", () =>
		{
			GlinerGpuOps.LayerNorm( _rt.LayerNorm, rawRelDesc.Buffer, RelLnW.Buffer, RelLnB.Buffer, NormRel, 512, 384, 1e-7f );
			Graphics.UavBarrier( NormRel );
			for ( int i = 0; i < 12; i++ )
			{
				GlinerGpuOps.LinearTiled16( _rt.GemmTiled16, NormRel, Layers[i].QueryW.Buffer, Layers[i].QueryB.Buffer, PosQ[i], 512, 384, 384, true );
				Graphics.UavBarrier( PosQ[i] );
				GlinerGpuOps.LinearTiled16( _rt.GemmTiled16, NormRel, Layers[i].KeyW.Buffer, Layers[i].KeyB.Buffer, PosK[i], 512, 384, 384, true );
				Graphics.UavBarrier( PosK[i] );
			}
		} );
		swPrep.Stop();
		PrepareMs = swPrep.Elapsed.TotalMilliseconds; // submit time; completion awaited by first use barriers

		WeightBytes = _rt.PersistentBytes;
		BufferCount = 0;
		foreach ( var kv in _rt.Buffers )
		{
			BufferCount++;
		}
		swTotal.Stop();
		_ = swTotal;
	}

	public void Dispose()
	{
		NormRel?.Dispose();
		for ( int i = 0; i < 12; i++ )
		{
			PosQ[i]?.Dispose();
			PosK[i]?.Dispose();
		}
	}
}

/// <summary>
/// Phase 8D.40/.41 — the GPU classification engine (P8E integration target).
/// Receives a GlinerEncodedRequest, performs the CPU embedding gather+LN
/// (strategy B — exact Phase-4 parity by construction), submits the COMPLETE
/// model as ONE render-context job (12 layers ping-pong → marker gather →
/// classifier), reads back ONLY the candidate logits, and applies the
/// validated CPU postprocessing (temperature → softmax → argmax). No Razor/
/// service dependencies; async, non-blocking; cancellation-before-dispatch
/// via the executor; stale results discarded by the caller's request id.
/// Backend identity: "Native GPU FP32".
/// </summary>
public sealed class GlinerGpuClassificationEngine
{
	public const string BackendName = "Native GPU FP32";

	private readonly GlinerGpuRuntime _rt;
	private readonly GlinerGpuModelWeights _model;
	private readonly GlinerGpuLayer _layer;
	private readonly GlinerPoc.Neural.GlinerModelWeights _cpuWeights;
	private readonly float[] _embLnW, _embLnB;
	private readonly float _temperature;

	// engine-owned persistent buffers (S ≤ 256, candidates ≤ 8)
	private readonly GpuBuffer<float> _hiddenA, _hiddenB;
	private readonly GpuBuffer<float> _markerStates, _clsHidden, _clsRelu, _logits;
	private readonly GpuBuffer<int> _markerIdx;
	/// <summary>Diagnostic per-layer snapshots (populated ONLY in diagnostic mode).</summary>
	private readonly GpuBuffer<float>[] _diagLayers = new GpuBuffer<float>[12];

	public int LastDispatchCount { get; private set; }
	public int LastBarrierCount { get; private set; }
	/// <summary>Intermediate CPU readbacks performed by production runs — always 0.</summary>
	public int ProductionIntermediateReadbacks => 0;
	public double InitializationMs => _model.UploadMs + _model.PrepareMs;
	public GlinerGpuModelWeights Model => _model;

	public GlinerGpuClassificationEngine( GlinerGpuRuntime rt, GlinerGpuModelWeights model,
		GlinerPoc.Neural.GlinerModelWeights cpuWeights, float[] embeddingLnW, float[] embeddingLnB,
		float temperature, int maxSeq = 256 )
	{
		_rt = rt;
		_model = model;
		_cpuWeights = cpuWeights;
		_embLnW = embeddingLnW;
		_embLnB = embeddingLnB;
		_temperature = temperature;
		_layer = new GlinerGpuLayer( rt, maxSeq );

		_hiddenA = new GpuBuffer<float>( maxSeq * 384 );
		_hiddenB = new GpuBuffer<float>( maxSeq * 384 );
		_markerStates = new GpuBuffer<float>( 8 * 384 );
		_clsHidden = new GpuBuffer<float>( 8 * 768 );
		_clsRelu = new GpuBuffer<float>( 8 * 768 );
		_logits = new GpuBuffer<float>( 8 );
		_markerIdx = new GpuBuffer<int>( 8 );
		for ( int i = 0; i < 12; i++ )
		{
			_diagLayers[i] = new GpuBuffer<float>( maxSeq * 384 );
		}
	}

	public long ScratchBytes =>
		(2L * 256 * 384 + 8L * 384 + 2L * 8 * 768 + 8L) * 4 + _layer.ScratchBytes
		+ 12L * 256 * 384 * 4 /* diagnostic snapshots */;

	/// <summary>
	/// Full GPU classification of an encoded request. Embedding prep runs on
	/// the CPU (exact oracle stage); everything from hidden upload to logits
	/// is ONE render job; only the n candidate logits cross back to the CPU.
	/// diagnostic=true additionally copies each layer output into snapshot
	/// buffers (read back afterwards via ReadDiagnostics — never in-line).
	/// </summary>
	public async Task<GlinerPoc.Neural.GlinerDecisionResult> ClassifyEncodedAsync(
		GlinerPoc.Preprocessing.GlinerEncodedRequest encoded, bool diagnostic = false,
		Func<bool> cancelled = null )
	{
		int s = encoded.InputIds.Length;
		int n = encoded.ClassificationMarkerIndices.Length;
		var swAll = Stopwatch.StartNew();

		// ---- CPU embedding stage (strategy B; exact Phase-4 parity) ----
		var swEmb = Stopwatch.StartNew();
		float[] hidden;
		float[] gathered = new float[s * 384];
		for ( int i = 0; i < s; i++ )
		{
			float[] row = _cpuWeights.GetRows( "encoder.embeddings.word_embeddings.weight", encoded.InputIds[i], 1 );
			for ( int j = 0; j < 384; j++ )
			{
				gathered[i * 384 + j] = row[j];
			}
		}
		hidden = GlinerPoc.Neural.GlinerMath.LayerNorm( gathered, s, 384, _embLnW, _embLnB, 1e-7f );
		swEmb.Stop();

		// ---- uploads (request-local only; weights/pos already resident) ----
		var swUp = Stopwatch.StartNew();
		GpuBuffer<float> hiddenIn = _hiddenA;
		hiddenIn.SetData( hidden, 0 );
		var mask = new uint[s];
		for ( int i = 0; i < s; i++ )
		{
			mask[i] = encoded.AttentionMask[i] != 0 ? 1u : 0u;
		}
		_layer.SetMask( encoded.AttentionMask, s );
		var idx = new int[n];
		for ( int i = 0; i < n; i++ )
		{
			idx[i] = encoded.ClassificationMarkerIndices[i];
		}
		_markerIdx.SetData( idx );
		swUp.Stop();

		// ---- ONE job: 12 layers + marker gather + classifier ----
		bool done = false;
		float[] logitsOut = null;
		var swGpu = Stopwatch.StartNew();
		_rt.Executor.SubmitReadback( "p8d_full", () =>
		{
			int d = 0, b = 0;
			GpuBuffer<float> finalHidden = null;
			for ( int i = 0; i < 12; i++ )
			{
				GpuBuffer<float> xin = (i % 2 == 0) ? _hiddenA : _hiddenB;
				GpuBuffer<float> xout = (i % 2 == 0) ? _hiddenB : _hiddenA;
				_layer.Run( xin, s, _model.Layers[i], xout, _model.PosQ[i], _model.PosK[i] );
				d += _layer.LastDispatchCount;
				b += _layer.LastBarrierCount;
				// layer-boundary barrier: layer N's final LN write → layer N+1
				// Q/K/V reads (audited — internal layer barriers do NOT cover
				// the external consumer)
				Graphics.UavBarrier( xout );
				b++;
				if ( diagnostic )
				{
					GlinerGpuOps.Copy( _rt.Copy, xout, _diagLayers[i], s * 384 );
					d++;
				}
				finalHidden = xout;
			}
			// marker gather (candidate order preserved)
			GlinerGpuOps.GatherRows( _rt.Gather, finalHidden, _markerIdx, _markerStates, n, 384 );
			d++;
			Graphics.UavBarrier( _markerStates ); b++;
			// classifier 384→768 → ReLU → 768→1
			GlinerGpuOps.LinearTiled16( _rt.GemmTiled16, _markerStates, _model.Cls0W.Buffer, _model.Cls0B.Buffer, _clsHidden, n, 384, 768, true );
			d++;
			Graphics.UavBarrier( _clsHidden ); b++;
			GlinerGpuOps.Relu( _rt.Relu, _clsHidden, _clsRelu, n * 768 );
			d++;
			Graphics.UavBarrier( _clsRelu ); b++;
			GlinerGpuOps.LinearTiled16( _rt.GemmTiled16, _clsRelu, _model.Cls3W.Buffer, _model.Cls3B.Buffer, _logits, n, 768, 1, true );
			d++;
			LastDispatchCount = d;
			LastBarrierCount = b;
		}, _logits, n, r =>
		{
			logitsOut = r;
			done = true;
		} );
		int waits = 0;
		bool cancelSeen = false;
		while ( !done && waits < 1200 )
		{
			await Task.Delay( 8 );
			waits++;
			if ( cancelled is not null && cancelled() )
			{
				// GPU cancellation semantics (P8E.12): before the render job
				// starts it is DROPPED (CancelQueued); once dispatched it always
				// completes — either way this caller reports Cancelled and any
				// late completion is discarded upstream by request ownership.
				_rt.Executor.CancelQueued();
				cancelSeen = true;
				break;
			}
		}
		swGpu.Stop();
		if ( cancelSeen )
		{
			return null;
		}
		if ( !done )
		{
			throw new InvalidOperationException( "[GLI:ERROR] full-model job never completed." );
		}

		// ---- CPU postprocessing (the P8D.22 boundary) ----
		var swPost = Stopwatch.StartNew();
		float[] raw = new float[n];
		float[] scaled = new float[n];
		for ( int i = 0; i < n; i++ )
		{
			raw[i] = logitsOut[i];
			scaled[i] = raw[i] / _temperature;
		}
		float[] scores = Softmax1D( scaled );
		int best = 0;
		for ( int i = 1; i < n; i++ )
		{
			if ( scores[i] > scores[best] )
			{
				best = i;
			}
		}
		swPost.Stop();
		swAll.Stop();

		var labels = new string[n];
		for ( int i = 0; i < n; i++ )
		{
			labels[i] = encoded.CandidateOrder[i];
		}
		return new GlinerPoc.Neural.GlinerDecisionResult( labels, raw, scaled, scores, best, s )
		{
			EmbeddingMs = swEmb.Elapsed.TotalMilliseconds + swUp.Elapsed.TotalMilliseconds,
			EncoderMs = swGpu.Elapsed.TotalMilliseconds,
			ClassifierMs = swPost.Elapsed.TotalMilliseconds,
			TotalNeuralMs = swAll.Elapsed.TotalMilliseconds,
		};
	}

	/// <summary>Diagnostic readback of the per-layer snapshots (post-run only).</summary>
	public async Task<float[]> ReadDiagnosticLayerAsync( int layerIndex, int count )
	{
		bool done = false;
		float[] result = null;
		_rt.Executor.SubmitReadback( $"p8d_diag_{layerIndex}", () => { }, _diagLayers[layerIndex], count, r =>
		{
			result = r;
			done = true;
		} );
		int waits = 0;
		while ( !done && waits < 400 )
		{
			await Task.Delay( 8 );
			waits++;
		}
		return result;
	}

	/// <summary>Diagnostic readback of marker states / classifier hidden / relu buffers.</summary>
	public async Task<float[]> ReadBufferAsync( GpuBuffer<float> buffer, int count )
	{
		bool done = false;
		float[] result = null;
		_rt.Executor.SubmitReadback( "p8d_read", () => { }, buffer, count, r =>
		{
			result = r;
			done = true;
		} );
		int waits = 0;
		while ( !done && waits < 400 )
		{
			await Task.Delay( 8 );
			waits++;
		}
		return result;
	}

	public GpuBuffer<float> MarkerStatesBuffer => _markerStates;
	public GpuBuffer<float> ClsHiddenBuffer => _clsHidden;
	public GpuBuffer<float> ClsReluBuffer => _clsRelu;

	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;
	}
}