using Sandbox.Diagnostics;

namespace LlmPoc.Llm;

public sealed class GptNeoLayerTensorNames
{
	public int LayerIndex { get; }
	public string Prefix { get; }
	public string Ln1Weight { get; }
	public string Ln1Bias { get; }
	public string QueryWeight { get; }
	public string KeyWeight { get; }
	public string ValueWeight { get; }
	public string AttentionOutputWeight { get; }
	public string AttentionOutputBias { get; }
	public string Ln2Weight { get; }
	public string Ln2Bias { get; }
	public string MlpFcWeight { get; }
	public string MlpFcBias { get; }
	public string MlpProjectionWeight { get; }
	public string MlpProjectionBias { get; }

	public GptNeoLayerTensorNames( int layerIndex )
	{
		if ( layerIndex < 0 )
		{
			throw new ArgumentOutOfRangeException( nameof( layerIndex ) );
		}

		LayerIndex = layerIndex;
		Prefix = $"transformer.h.{layerIndex}";
		Ln1Weight = $"{Prefix}.ln_1.weight";
		Ln1Bias = $"{Prefix}.ln_1.bias";
		QueryWeight = $"{Prefix}.attn.attention.q_proj.weight";
		KeyWeight = $"{Prefix}.attn.attention.k_proj.weight";
		ValueWeight = $"{Prefix}.attn.attention.v_proj.weight";
		AttentionOutputWeight = $"{Prefix}.attn.attention.out_proj.weight";
		AttentionOutputBias = $"{Prefix}.attn.attention.out_proj.bias";
		Ln2Weight = $"{Prefix}.ln_2.weight";
		Ln2Bias = $"{Prefix}.ln_2.bias";
		MlpFcWeight = $"{Prefix}.mlp.c_fc.weight";
		MlpFcBias = $"{Prefix}.mlp.c_fc.bias";
		MlpProjectionWeight = $"{Prefix}.mlp.c_proj.weight";
		MlpProjectionBias = $"{Prefix}.mlp.c_proj.bias";
	}
}

public sealed class GptNeoLayerForwardResult
{
	public int LayerIndex { get; init; }
	public string AttentionType { get; init; }
	public GptNeoLayerTensorNames TensorNames { get; init; }
	public Tensor Input { get; init; }
	public Tensor Ln1 { get; init; }
	public Tensor Query { get; init; }
	public Tensor Key { get; init; }
	public Tensor Value { get; init; }
	public Tensor QueryHeads { get; init; }
	public Tensor KeyHeads { get; init; }
	public Tensor ValueHeads { get; init; }
	public Tensor ScoresUnmasked { get; init; }
	public Tensor ScoresMasked { get; init; }
	public Tensor AttentionProbabilities { get; init; }
	public Tensor AttentionContextHeads { get; init; }
	public Tensor AttentionMerged { get; init; }
	public Tensor AttentionOutput { get; init; }
	public Tensor AttentionResidual { get; init; }
	public Tensor Ln2 { get; init; }
	public Tensor MlpFc { get; init; }
	public Tensor MlpGelu { get; init; }
	public Tensor MlpProjection { get; init; }
	public Tensor Output { get; init; }
	public bool InputMutated { get; init; }
	public double ElapsedMilliseconds { get; init; }
}

/// <summary>
/// The scalar, model-specific GPT-Neo transformer block proven first for layer 0.
/// Layer identity changes tensor resolution and mask type; the mathematical kernels remain shared.
/// </summary>
public static class GptNeoTransformerLayer
{
	public static readonly float MaskedScoreSentinel =
		BitConverter.Int32BitsToSingle( unchecked((int)0xFF7FFFFF) );

	public static GptNeoLayerForwardResult Forward(
		SboxLlmModel model,
		TinyStoriesConfig config,
		Tensor hiddenStates,
		int layerIndex,
		bool logDiagnostics = true,
		bool validateInputMutation = true )
	{
		if ( model is null ) throw new ArgumentNullException( nameof( model ) );
		if ( config is null ) throw new ArgumentNullException( nameof( config ) );
		if ( hiddenStates is null ) throw new ArgumentNullException( nameof( hiddenStates ) );
		if ( layerIndex < 0 || layerIndex >= config.LayerCount )
		{
			throw new IndexOutOfRangeException(
				$"[LLM:ERROR] Transformer layer index {layerIndex} is outside " +
				$"[0,{config.LayerCount})." );
		}
		hiddenStates.RequireShape( hiddenStates.Shape[0], config.HiddenSize );
		if ( hiddenStates.Shape[0] <= 0 )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Layer {layerIndex} requires a non-empty sequence." );
		}

		string attentionType = config.AttentionLayers[layerIndex];
		if ( attentionType != "global" && attentionType != "local" )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] attention_layers[{layerIndex}] must be global or local, " +
				$"found '{attentionType}'." );
		}
		if ( attentionType == "local" && config.AttentionWindowSize <= 0 )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] Layer {layerIndex} local attention requires a positive " +
				$"window, found {config.AttentionWindowSize}." );
		}

		float[] inputSnapshot = validateInputMutation
			? new float[hiddenStates.Data.Length]
			: null;
		if ( inputSnapshot is not null )
		{
			Array.Copy( hiddenStates.Data, inputSnapshot, hiddenStates.Data.Length );
		}
		GptNeoLayerTensorNames names = new( layerIndex );
		ValidateLearnedTensorShapes( model, config, names );
		FastTimer timer = FastTimer.StartNew();
		string outputPrefix = $"forward.layer{layerIndex}";
		if ( logDiagnostics )
		{
			LlmLog.Info(
				"LAYER",
				$"layer={layerIndex} input={hiddenStates.ShapeText} attention_type={attentionType} " +
				$"window={config.AttentionWindowSize} implementation=reusable_scalar_gpt_neo" );
		}

		Tensor ln1 = TinyStoriesForwardStages.ApplyLayerNorm(
			model, config, hiddenStates, layerIndex, 1, $"{outputPrefix}.ln_1", logDiagnostics );
		Tensor query = TinyStoriesForwardStages.LinearNoBias(
			ln1, model.GetRequiredTensor( names.QueryWeight ), $"{outputPrefix}.q" );
		Tensor key = TinyStoriesForwardStages.LinearNoBias(
			ln1, model.GetRequiredTensor( names.KeyWeight ), $"{outputPrefix}.k" );
		Tensor value = TinyStoriesForwardStages.LinearNoBias(
			ln1, model.GetRequiredTensor( names.ValueWeight ), $"{outputPrefix}.v" );
		Tensor queryHeads = TinyStoriesForwardStages.SplitHeads(
			query, config.AttentionHeadCount, $"{outputPrefix}.q_heads" );
		Tensor keyHeads = TinyStoriesForwardStages.SplitHeads(
			key, config.AttentionHeadCount, $"{outputPrefix}.k_heads" );
		Tensor valueHeads = TinyStoriesForwardStages.SplitHeads(
			value, config.AttentionHeadCount, $"{outputPrefix}.v_heads" );
		Tensor scoresUnmasked = TinyStoriesForwardStages.ComputeScaledUnmaskedAttentionScores(
			queryHeads, keyHeads, scale: 1.0f, $"{outputPrefix}.scores_unmasked" );
		Tensor scoresMasked = TinyStoriesForwardStages.ApplyAttentionMask(
			scoresUnmasked,
			attentionType,
			config.AttentionWindowSize,
			MaskedScoreSentinel,
			$"{outputPrefix}.scores_masked_pre_softmax" );
		Tensor probabilities = TinyStoriesForwardStages.ComputeAttentionProbabilities(
			scoresMasked, $"{outputPrefix}.attention_probs" );
		Tensor contextHeads = TinyStoriesForwardStages.ComputeAttentionContextHeads(
			probabilities, valueHeads, $"{outputPrefix}.attention_context_heads" );
		Tensor merged = TinyStoriesForwardStages.MergeHeads(
			contextHeads, $"{outputPrefix}.attention_merged" );
		Tensor attentionOutput = TinyStoriesForwardStages.LinearWithBias(
			merged,
			model.GetRequiredTensor( names.AttentionOutputWeight ),
			model.GetRequiredTensor( names.AttentionOutputBias ),
			$"{outputPrefix}.attention_out_proj" );
		// Both attention residual dropout and MLP dropout are p=0 and no-ops in eval mode.
		Tensor attentionResidual = TinyStoriesForwardStages.AddResidual(
			attentionOutput, hiddenStates, $"{outputPrefix}.attention_residual" );
		Tensor ln2 = TinyStoriesForwardStages.ApplyLayerNorm(
			model, config, attentionResidual, layerIndex, 2, $"{outputPrefix}.ln_2", logDiagnostics );
		Tensor mlpFc = TinyStoriesForwardStages.LinearWithBias(
			ln2,
			model.GetRequiredTensor( names.MlpFcWeight ),
			model.GetRequiredTensor( names.MlpFcBias ),
			$"{outputPrefix}.mlp_fc" );
		Tensor mlpGelu = TinyStoriesForwardStages.ApplyGeluNew(
			mlpFc, $"{outputPrefix}.mlp_gelu" );
		Tensor mlpProjection = TinyStoriesForwardStages.LinearWithBias(
			mlpGelu,
			model.GetRequiredTensor( names.MlpProjectionWeight ),
			model.GetRequiredTensor( names.MlpProjectionBias ),
			$"{outputPrefix}.mlp_proj" );
		Tensor output = TinyStoriesForwardStages.AddMlpResidual(
			attentionResidual, mlpProjection, $"{outputPrefix}.output" );

		bool inputMutated = false;
		for ( int index = 0; inputSnapshot is not null && index < inputSnapshot.Length; index++ )
		{
			if ( BitConverter.SingleToInt32Bits( inputSnapshot[index] ) !=
				BitConverter.SingleToInt32Bits( hiddenStates.Data[index] ) )
			{
				inputMutated = true;
				break;
			}
		}
		if ( logDiagnostics )
		{
			LlmLog.Info(
				"PERF",
				$"stage=layer{layerIndex}.reusable_forward attention_type={attentionType} " +
				$"shape={output.ShapeText} elapsed_ms={timer.ElapsedMilliSeconds:N4}" );
		}

		return new GptNeoLayerForwardResult
		{
			LayerIndex = layerIndex,
			AttentionType = attentionType,
			TensorNames = names,
			Input = hiddenStates,
			Ln1 = ln1,
			Query = query,
			Key = key,
			Value = value,
			QueryHeads = queryHeads,
			KeyHeads = keyHeads,
			ValueHeads = valueHeads,
			ScoresUnmasked = scoresUnmasked,
			ScoresMasked = scoresMasked,
			AttentionProbabilities = probabilities,
			AttentionContextHeads = contextHeads,
			AttentionMerged = merged,
			AttentionOutput = attentionOutput,
			AttentionResidual = attentionResidual,
			Ln2 = ln2,
			MlpFc = mlpFc,
			MlpGelu = mlpGelu,
			MlpProjection = mlpProjection,
			Output = output,
			InputMutated = inputMutated,
			ElapsedMilliseconds = timer.ElapsedMilliSeconds
		};
	}

	private static void ValidateLearnedTensorShapes(
		SboxLlmModel model,
		TinyStoriesConfig config,
		GptNeoLayerTensorNames names )
	{
		int hidden = config.HiddenSize;
		int intermediate = hidden * 4;
		model.GetRequiredTensor( names.Ln1Weight ).RequireShape( hidden );
		model.GetRequiredTensor( names.Ln1Bias ).RequireShape( hidden );
		model.GetRequiredTensor( names.QueryWeight ).RequireShape( hidden, hidden );
		model.GetRequiredTensor( names.KeyWeight ).RequireShape( hidden, hidden );
		model.GetRequiredTensor( names.ValueWeight ).RequireShape( hidden, hidden );
		model.GetRequiredTensor( names.AttentionOutputWeight ).RequireShape( hidden, hidden );
		model.GetRequiredTensor( names.AttentionOutputBias ).RequireShape( hidden );
		model.GetRequiredTensor( names.Ln2Weight ).RequireShape( hidden );
		model.GetRequiredTensor( names.Ln2Bias ).RequireShape( hidden );
		model.GetRequiredTensor( names.MlpFcWeight ).RequireShape( intermediate, hidden );
		model.GetRequiredTensor( names.MlpFcBias ).RequireShape( intermediate );
		model.GetRequiredTensor( names.MlpProjectionWeight ).RequireShape( hidden, intermediate );
		model.GetRequiredTensor( names.MlpProjectionBias ).RequireShape( hidden );
	}
}