using System.Text.Json.Serialization;

namespace LlmPoc.Llm;

public sealed class ForwardReferenceDocument
{
	[JsonPropertyName( "format" )]
	public string Format { get; set; }

	[JsonPropertyName( "version" )]
	public int Version { get; set; }

	[JsonPropertyName( "model" )]
	public string Model { get; set; }

	[JsonPropertyName( "prompt" )]
	public string Prompt { get; set; }

	[JsonPropertyName( "input_token_ids" )]
	public int[] InputTokenIds { get; set; }

	[JsonPropertyName( "position_ids" )]
	public int[] PositionIds { get; set; }

	[JsonPropertyName( "sequence_length" )]
	public int SequenceLength { get; set; }

	[JsonPropertyName( "hidden_size" )]
	public int HiddenSize { get; set; }

	[JsonPropertyName( "layer_norm_epsilon" )]
	public float LayerNormEpsilon { get; set; }

	[JsonPropertyName( "layer_norm_variance" )]
	public string LayerNormVariance { get; set; }

	[JsonPropertyName( "model_training" )]
	public bool ModelTraining { get; set; }

	[JsonPropertyName( "embedding_dropout" )]
	public float EmbeddingDropout { get; set; }

	[JsonPropertyName( "dtype" )]
	public string Dtype { get; set; }

	[JsonPropertyName( "stages" )]
	public List<ForwardReferenceStage> Stages { get; set; }

	public ForwardReferenceStage GetRequiredStage( string stageName )
	{
		if ( Stages is null )
		{
			throw new InvalidOperationException( "[LLM:ERROR] Intermediate-reference stages are missing." );
		}

		ForwardReferenceStage result = null;
		foreach ( ForwardReferenceStage stage in Stages )
		{
			if ( stage?.Stage != stageName )
			{
				continue;
			}
			if ( result is not null )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] Intermediate-reference stage '{stageName}' is duplicated." );
			}
			result = stage;
		}

		return result ?? throw new InvalidOperationException(
			$"[LLM:ERROR] Required intermediate-reference stage '{stageName}' is missing." );
	}
}

public sealed class ForwardReferenceStage
{
	[JsonPropertyName( "stage" )]
	public string Stage { get; set; }

	[JsonPropertyName( "file" )]
	public string File { get; set; }

	[JsonPropertyName( "dtype" )]
	public string Dtype { get; set; }

	[JsonPropertyName( "shape" )]
	public int[] Shape { get; set; }

	[JsonPropertyName( "elements" )]
	public int Elements { get; set; }

	[JsonPropertyName( "bytes" )]
	public int Bytes { get; set; }

	[JsonPropertyName( "original_projection_shape" )]
	public int[] OriginalProjectionShape { get; set; }

	[JsonPropertyName( "head_count" )]
	public int HeadCount { get; set; }

	[JsonPropertyName( "head_dimension" )]
	public int HeadDimension { get; set; }

	[JsonPropertyName( "sequence_length" )]
	public int SequenceLength { get; set; }

	[JsonPropertyName( "logical_axis_order" )]
	public string[] LogicalAxisOrder { get; set; }

	[JsonPropertyName( "flattened_storage_order" )]
	public string FlattenedStorageOrder { get; set; }

	[JsonPropertyName( "source_feature_formula" )]
	public string SourceFeatureFormula { get; set; }

	[JsonPropertyName( "projection" )]
	public string Projection { get; set; }

	[JsonPropertyName( "query_sequence_length" )]
	public int QuerySequenceLength { get; set; }

	[JsonPropertyName( "key_sequence_length" )]
	public int KeySequenceLength { get; set; }

	[JsonPropertyName( "formula" )]
	public string Formula { get; set; }

	[JsonPropertyName( "scale" )]
	public float Scale { get; set; }

	[JsonPropertyName( "scaling_operation" )]
	public string ScalingOperation { get; set; }

	[JsonPropertyName( "causal_mask_applied" )]
	public bool CausalMaskApplied { get; set; }

	[JsonPropertyName( "local_or_global_mask_applied" )]
	public bool LocalOrGlobalMaskApplied { get; set; }

	[JsonPropertyName( "softmax_applied" )]
	public bool SoftmaxApplied { get; set; }

	[JsonPropertyName( "dropout_applied" )]
	public bool DropoutApplied { get; set; }

	[JsonPropertyName( "attention_layer" )]
	public int AttentionLayer { get; set; }

	[JsonPropertyName( "attention_type" )]
	public string AttentionType { get; set; }

	[JsonPropertyName( "mask_reference_file" )]
	public string MaskReferenceFile { get; set; }

	[JsonPropertyName( "mask_semantics" )]
	public string MaskSemantics { get; set; }

	[JsonPropertyName( "masked_sentinel" )]
	public float MaskedSentinel { get; set; }

	[JsonPropertyName( "masked_sentinel_float32_bits" )]
	public string MaskedSentinelFloat32Bits { get; set; }

	[JsonPropertyName( "masked_sentinel_is_finite" )]
	public bool MaskedSentinelIsFinite { get; set; }

	[JsonPropertyName( "allowed_count_per_head" )]
	public int AllowedCountPerHead { get; set; }

	[JsonPropertyName( "masked_count_per_head" )]
	public int MaskedCountPerHead { get; set; }

	[JsonPropertyName( "allowed_count_all_heads" )]
	public int AllowedCountAllHeads { get; set; }

	[JsonPropertyName( "masked_count_all_heads" )]
	public int MaskedCountAllHeads { get; set; }

	[JsonPropertyName( "pre_softmax" )]
	public bool PreSoftmax { get; set; }

	[JsonPropertyName( "source_stage" )]
	public string SourceStage { get; set; }

	[JsonPropertyName( "softmax_function" )]
	public string SoftmaxFunction { get; set; }

	[JsonPropertyName( "softmax_dimension" )]
	public int SoftmaxDimension { get; set; }

	[JsonPropertyName( "softmax_axis" )]
	public string SoftmaxAxis { get; set; }

	[JsonPropertyName( "input_dtype" )]
	public string InputDtype { get; set; }

	[JsonPropertyName( "output_dtype" )]
	public string OutputDtype { get; set; }

	[JsonPropertyName( "softmax_dtype_argument" )]
	public string SoftmaxDtypeArgument { get; set; }

	[JsonPropertyName( "post_softmax_cast" )]
	public string PostSoftmaxCast { get; set; }

	[JsonPropertyName( "model_eval" )]
	public bool ModelEval { get; set; }

	[JsonPropertyName( "attention_dropout_probability" )]
	public float AttentionDropoutProbability { get; set; }

	[JsonPropertyName( "masked_probability" )]
	public float MaskedProbability { get; set; }

	[JsonPropertyName( "masked_probability_float32_bits" )]
	public string MaskedProbabilityFloat32Bits { get; set; }

	[JsonPropertyName( "masked_zero_count" )]
	public int MaskedZeroCount { get; set; }

	[JsonPropertyName( "nonzero_probability_count" )]
	public int NonzeroProbabilityCount { get; set; }

	[JsonPropertyName( "max_row_sum_deviation" )]
	public float MaximumRowSumDeviation { get; set; }

	[JsonPropertyName( "post_softmax" )]
	public bool PostSoftmax { get; set; }

	[JsonPropertyName( "attention_times_value_applied" )]
	public bool AttentionTimesValueApplied { get; set; }

	[JsonPropertyName( "source_probability_stage" )]
	public string SourceProbabilityStage { get; set; }

	[JsonPropertyName( "value_stage" )]
	public string ValueStage { get; set; }

	[JsonPropertyName( "value_dtype" )]
	public string ValueDtype { get; set; }

	[JsonPropertyName( "query0_identity_elements" )]
	public int Query0IdentityElements { get; set; }

	[JsonPropertyName( "query0_identity_bit_exact" )]
	public bool Query0IdentityBitExact { get; set; }

	[JsonPropertyName( "convex_range_tolerance" )]
	public float ConvexRangeTolerance { get; set; }

	[JsonPropertyName( "convex_range_checked_elements" )]
	public int ConvexRangeCheckedElements { get; set; }

	[JsonPropertyName( "convex_range_violations" )]
	public int ConvexRangeViolations { get; set; }

	[JsonPropertyName( "heads_merged" )]
	public bool HeadsMerged { get; set; }

	[JsonPropertyName( "output_projection_applied" )]
	public bool OutputProjectionApplied { get; set; }

	[JsonPropertyName( "source_shape" )]
	public int[] SourceShape { get; set; }

	[JsonPropertyName( "source_axis_order" )]
	public string[] SourceAxisOrder { get; set; }

	[JsonPropertyName( "permutation" )]
	public int[] Permutation { get; set; }

	[JsonPropertyName( "contiguous_called" )]
	public bool ContiguousCalled { get; set; }

	[JsonPropertyName( "mapping" )]
	public string Mapping { get; set; }

	[JsonPropertyName( "mapping_elements" )]
	public int MappingElements { get; set; }

	[JsonPropertyName( "mapping_bit_exact" )]
	public bool MappingBitExact { get; set; }

	[JsonPropertyName( "source_implementation" )]
	public string SourceImplementation { get; set; }

	[JsonPropertyName( "module_type" )]
	public string ModuleType { get; set; }

	[JsonPropertyName( "weight_tensor" )]
	public string WeightTensor { get; set; }

	[JsonPropertyName( "bias_tensor" )]
	public string BiasTensor { get; set; }

	[JsonPropertyName( "weight_shape" )]
	public int[] WeightShape { get; set; }

	[JsonPropertyName( "bias_exists" )]
	public bool BiasExists { get; set; }

	[JsonPropertyName( "bias_shape" )]
	public int[] BiasShape { get; set; }

	[JsonPropertyName( "weight_layout" )]
	public string WeightLayout { get; set; }

	[JsonPropertyName( "weight_dtype" )]
	public string WeightDtype { get; set; }

	[JsonPropertyName( "bias_dtype" )]
	public string BiasDtype { get; set; }

	[JsonPropertyName( "additional_cast" )]
	public bool AdditionalCast { get; set; }

	[JsonPropertyName( "residual_dropout_module" )]
	public string ResidualDropoutModule { get; set; }

	[JsonPropertyName( "residual_dropout_probability" )]
	public float ResidualDropoutProbability { get; set; }

	[JsonPropertyName( "dropout_changed_bits" )]
	public bool DropoutChangedBits { get; set; }

	[JsonPropertyName( "attention_branch_after_dropout_bit_equal" )]
	public bool AttentionBranchAfterDropoutBitEqual { get; set; }

	[JsonPropertyName( "residual_addition_applied" )]
	public bool ResidualAdditionApplied { get; set; }

	[JsonPropertyName( "attention_branch_stage" )]
	public string AttentionBranchStage { get; set; }

	[JsonPropertyName( "residual_source_stage" )]
	public string ResidualSourceStage { get; set; }

	[JsonPropertyName( "residual_source_boundary" )]
	public string ResidualSourceBoundary { get; set; }

	[JsonPropertyName( "addition_order" )]
	public string AdditionOrder { get; set; }

	[JsonPropertyName( "residual_source_bit_exact_to_block_input" )]
	public bool ResidualSourceBitExactToBlockInput { get; set; }

	[JsonPropertyName( "wrong_ln1_residual_max_abs" )]
	public float WrongLn1ResidualMaximumAbsoluteError { get; set; }

	[JsonPropertyName( "next_stage" )]
	public string NextStage { get; set; }

	[JsonPropertyName( "ln_2_applied" )]
	public bool Layer0Ln2Applied { get; set; }

	[JsonPropertyName( "epsilon" )]
	public float Epsilon { get; set; }

	[JsonPropertyName( "normalization_dimension" )]
	public string NormalizationDimension { get; set; }

	[JsonPropertyName( "variance" )]
	public string Variance { get; set; }

	[JsonPropertyName( "input_mutated" )]
	public bool InputMutated { get; set; }

	[JsonPropertyName( "intermediate_size" )]
	public int IntermediateSize { get; set; }

	[JsonPropertyName( "pre_activation" )]
	public bool PreActivation { get; set; }

	[JsonPropertyName( "activation_applied" )]
	public bool ActivationApplied { get; set; }

	[JsonPropertyName( "activation_function" )]
	public string ActivationFunction { get; set; }

	[JsonPropertyName( "activation_module" )]
	public string ActivationModule { get; set; }

	[JsonPropertyName( "activation_formula" )]
	public string ActivationFormula { get; set; }

	[JsonPropertyName( "tanh_coefficient_float32" )]
	public float TanhCoefficientFloat32 { get; set; }

	[JsonPropertyName( "tanh_coefficient_float32_bits" )]
	public string TanhCoefficientFloat32Bits { get; set; }

	[JsonPropertyName( "cubic_coefficient_float32" )]
	public float CubicCoefficientFloat32 { get; set; }

	[JsonPropertyName( "cubic_coefficient_float32_bits" )]
	public string CubicCoefficientFloat32Bits { get; set; }

	[JsonPropertyName( "pow_exponent" )]
	public float PowExponent { get; set; }

	[JsonPropertyName( "operation_dtype" )]
	public string OperationDtype { get; set; }

	[JsonPropertyName( "pre_dropout" )]
	public bool PreDropout { get; set; }

	[JsonPropertyName( "mlp_dropout_module" )]
	public string MlpDropoutModule { get; set; }

	[JsonPropertyName( "mlp_dropout_probability" )]
	public float MlpDropoutProbability { get; set; }

	[JsonPropertyName( "mlp_branch_after_dropout_bit_equal" )]
	public bool MlpBranchAfterDropoutBitEqual { get; set; }

	[JsonPropertyName( "mlp_branch_stage" )]
	public string MlpBranchStage { get; set; }

	[JsonPropertyName( "residual_source_bit_exact_to_attention_residual" )]
	public bool ResidualSourceBitExactToAttentionResidual { get; set; }

	[JsonPropertyName( "wrong_ln2_residual_max_abs" )]
	public float WrongLn2ResidualMaximumAbsoluteError { get; set; }

	[JsonPropertyName( "layer_index" )]
	public int LayerIndex { get; set; }

	[JsonPropertyName( "layer_complete" )]
	public bool LayerComplete { get; set; }

	[JsonPropertyName( "input_shape" )]
	public int[] InputShape { get; set; }

	[JsonPropertyName( "output_shape" )]
	public int[] OutputShape { get; set; }

	[JsonPropertyName( "final_position_index" )]
	public int FinalPositionIndex { get; set; }

	[JsonPropertyName( "vocabulary_size" )]
	public int VocabularySize { get; set; }

	[JsonPropertyName( "full_sequence_projection_mathematically_equivalent" )]
	public bool FullSequenceProjectionMathematicallyEquivalent { get; set; }

	[JsonPropertyName( "last_position_only" )]
	public bool LastPositionOnly { get; set; }

	[JsonPropertyName( "weight_tied_to" )]
	public string WeightTiedTo { get; set; }

	[JsonPropertyName( "tied_weight_same_parameter_object" )]
	public bool TiedWeightSameParameterObject { get; set; }

	[JsonPropertyName( "tied_weight_bit_exact" )]
	public bool TiedWeightBitExact { get; set; }

	[JsonPropertyName( "argmax_token_id" )]
	public int ArgmaxTokenId { get; set; }

	[JsonPropertyName( "argmax_token_text" )]
	public string ArgmaxTokenText { get; set; }

	[JsonPropertyName( "top_k" )]
	public ForwardReferenceTopLogit[] TopK { get; set; }
}

public sealed class ForwardReferenceTopLogit
{
	[JsonPropertyName( "rank" )]
	public int Rank { get; set; }

	[JsonPropertyName( "token_id" )]
	public int TokenId { get; set; }

	[JsonPropertyName( "logit" )]
	public float Logit { get; set; }

	[JsonPropertyName( "decoded_token" )]
	public string DecodedToken { get; set; }
}

public sealed class AttentionMaskReferenceDocument
{
	[JsonPropertyName( "format" )]
	public string Format { get; set; }

	[JsonPropertyName( "version" )]
	public int Version { get; set; }

	[JsonPropertyName( "model" )]
	public string Model { get; set; }

	[JsonPropertyName( "layer" )]
	public int Layer { get; set; }

	[JsonPropertyName( "attention_type" )]
	public string AttentionType { get; set; }

	[JsonPropertyName( "window_size" )]
	public int WindowSize { get; set; }

	[JsonPropertyName( "source_bias_shape" )]
	public int[] SourceBiasShape { get; set; }

	[JsonPropertyName( "source_bias_dtype" )]
	public string SourceBiasDtype { get; set; }

	[JsonPropertyName( "source_bias_persistent" )]
	public bool SourceBiasPersistent { get; set; }

	[JsonPropertyName( "source_slice" )]
	public string SourceSlice { get; set; }

	[JsonPropertyName( "query_length" )]
	public int QueryLength { get; set; }

	[JsonPropertyName( "key_length" )]
	public int KeyLength { get; set; }

	[JsonPropertyName( "shape" )]
	public int[] Shape { get; set; }

	[JsonPropertyName( "logical_axis_order" )]
	public string[] LogicalAxisOrder { get; set; }

	[JsonPropertyName( "flattened_storage_order" )]
	public string FlattenedStorageOrder { get; set; }

	[JsonPropertyName( "mask_semantics" )]
	public string MaskSemantics { get; set; }

	[JsonPropertyName( "allowed" )]
	public bool[][] Allowed { get; set; }

	[JsonPropertyName( "allowed_count_per_head" )]
	public int AllowedCountPerHead { get; set; }

	[JsonPropertyName( "masked_count_per_head" )]
	public int MaskedCountPerHead { get; set; }

	[JsonPropertyName( "head_count" )]
	public int HeadCount { get; set; }

	[JsonPropertyName( "allowed_count_all_heads" )]
	public int AllowedCountAllHeads { get; set; }

	[JsonPropertyName( "masked_count_all_heads" )]
	public int MaskedCountAllHeads { get; set; }

	[JsonPropertyName( "shared_across_heads" )]
	public bool SharedAcrossHeads { get; set; }

	[JsonPropertyName( "masked_sentinel" )]
	public float MaskedSentinel { get; set; }

	[JsonPropertyName( "masked_sentinel_float32_bits" )]
	public string MaskedSentinelFloat32Bits { get; set; }

	[JsonPropertyName( "masked_sentinel_is_finite" )]
	public bool MaskedSentinelIsFinite { get; set; }

	[JsonPropertyName( "mask_operation" )]
	public string MaskOperation { get; set; }

	[JsonPropertyName( "first_local_attention_layer" )]
	public int FirstLocalAttentionLayer { get; set; }

	[JsonPropertyName( "local_mask_construction" )]
	public string LocalMaskConstruction { get; set; }

	[JsonPropertyName( "local_allowed_concept" )]
	public string LocalAllowedConcept { get; set; }

	[JsonPropertyName( "status" )]
	public string Status { get; set; }
}

public sealed class LocalAttentionMaskReferenceDocument
{
	[JsonPropertyName( "format" )]
	public string Format { get; set; }

	[JsonPropertyName( "version" )]
	public int Version { get; set; }

	[JsonPropertyName( "model" )]
	public string Model { get; set; }

	[JsonPropertyName( "layer_index" )]
	public int LayerIndex { get; set; }

	[JsonPropertyName( "attention_type" )]
	public string AttentionType { get; set; }

	[JsonPropertyName( "sequence_length" )]
	public int SequenceLength { get; set; }

	[JsonPropertyName( "query_length" )]
	public int QueryLength { get; set; }

	[JsonPropertyName( "key_length" )]
	public int KeyLength { get; set; }

	[JsonPropertyName( "window_size" )]
	public int WindowSize { get; set; }

	[JsonPropertyName( "source_implementation" )]
	public string SourceImplementation { get; set; }

	[JsonPropertyName( "source_construction" )]
	public string SourceConstruction { get; set; }

	[JsonPropertyName( "allowed_rule" )]
	public string AllowedRule { get; set; }

	[JsonPropertyName( "earliest_allowed_formula" )]
	public string EarliestAllowedFormula { get; set; }

	[JsonPropertyName( "binary_file" )]
	public string BinaryFile { get; set; }

	[JsonPropertyName( "binary_encoding" )]
	public string BinaryEncoding { get; set; }

	[JsonPropertyName( "binary_bytes" )]
	public int BinaryBytes { get; set; }

	[JsonPropertyName( "binary_sha256" )]
	public string BinarySha256 { get; set; }

	[JsonPropertyName( "total_coordinates" )]
	public int TotalCoordinates { get; set; }

	[JsonPropertyName( "allowed_count" )]
	public int AllowedCount { get; set; }

	[JsonPropertyName( "masked_count" )]
	public int MaskedCount { get; set; }

	[JsonPropertyName( "first_allowed_key" )]
	public int[] FirstAllowedKey { get; set; }

	[JsonPropertyName( "last_allowed_key" )]
	public int[] LastAllowedKey { get; set; }

	[JsonPropertyName( "boundary_coordinates" )]
	public LocalAttentionMaskCoordinate[] BoundaryCoordinates { get; set; }

	[JsonPropertyName( "mismatches" )]
	public int Mismatches { get; set; }

	[JsonPropertyName( "status" )]
	public string Status { get; set; }
}

public sealed class LocalAttentionMaskCoordinate
{
	[JsonPropertyName( "query" )]
	public int Query { get; set; }

	[JsonPropertyName( "key" )]
	public int Key { get; set; }

	[JsonPropertyName( "allowed" )]
	public bool Allowed { get; set; }
}