package Lugh::Inference;

use strict;
use warnings;

=head1 NAME

Lugh::Inference - Transformer Forward Pass and Token Generation

=head1 VERSION

Version 0.01

=cut

our $VERSION = '0.03';

=head1 SYNOPSIS

    use Lugh;
    
    # Load model and create inference engine
    my $model = Lugh::Model->new(model => '/path/to/model.gguf');
    my $tokenizer = Lugh::Tokenizer->new(model => $model);
    my $inference = Lugh::Inference->new(model => $model);
    
    # Encode a prompt
    my @tokens = $tokenizer->encode("The capital of France is");
    
    # Run forward pass to get logits
    my @logits = $inference->forward(\@tokens);
    
    # Find the most likely next token (greedy)
    my $max_idx = 0;
    my $max_val = $logits[0];
    for my $i (1..$#logits) {
        if ($logits[$i] > $max_val) {
            $max_val = $logits[$i];
            $max_idx = $i;
        }
    }
    
    # Decode the predicted token
    my $next_token = $tokenizer->decode([$max_idx]);
    print "Next token: $next_token\n";  # " Paris"
    
    # Or use top-p sampling
    my $sampled = $inference->sample_top_p(
        \@logits,
        temperature => 0.8,
        top_p => 0.95
    );

=head1 DESCRIPTION

Lugh::Inference implements the transformer forward pass for autoregressive
language model inference. Given a sequence of input tokens, it computes
the probability distribution (logits) over the vocabulary for the next token.

=head2 Transformer Architecture

The forward pass implements the standard transformer decoder architecture
used by LLaMA, Mistral, and similar models:

    Input Tokens
         │
         ▼
    ┌─────────────┐
    │  Token      │
    │  Embeddings │
    └─────────────┘
         │
         ▼
    ┌─────────────────────────────────┐
    │  Transformer Layer (× N)        │
    │  ┌─────────────────────────────┐│
    │  │ RMSNorm                     ││
    │  │      ↓                      ││
    │  │ Multi-Head Attention (GQA)  ││
    │  │      ↓                      ││
    │  │ Residual Add                ││
    │  │      ↓                      ││
    │  │ RMSNorm                     ││
    │  │      ↓                      ││
    │  │ FFN (SwiGLU)                ││
    │  │      ↓                      ││
    │  │ Residual Add                ││
    │  └─────────────────────────────┘│
    └─────────────────────────────────┘
         │
         ▼
    ┌─────────────┐
    │  Final      │
    │  RMSNorm    │
    └─────────────┘
         │
         ▼
    ┌─────────────┐
    │  Output     │
    │  Projection │ → Logits [vocab_size]
    └─────────────┘

=head2 Key Components

=over 4

=item * B<RMSNorm> - Root Mean Square Layer Normalization

=item * B<RoPE> - Rotary Position Embeddings for position encoding

=item * B<GQA> - Grouped Query Attention (multiple Q heads per KV head)

=item * B<SwiGLU> - Gated activation in feed-forward network

=back

=head1 CONSTRUCTOR

=head2 new

    my $inference = Lugh::Inference->new(
        model      => $model,
        flash_attn => 0,  # optional, use flash attention
    );

Creates a new Inference engine from a loaded model.

B<Parameters:>

=over 4

=item * C<model> (required) - A Lugh::Model object

=item * C<flash_attn> (optional) - Use flash attention if set to 1 (default: 0)

=back

B<Returns:> A Lugh::Inference object.

B<Throws:> Dies if no model is provided.

B<Example:>

    my $model = Lugh::Model->new(model => 'model.gguf');
    my $inference = Lugh::Inference->new(model => $model);
    
    # Or with flash attention enabled
    my $fast_inference = Lugh::Inference->new(
        model      => $model,
        flash_attn => 1,
    );

=head1 METHODS

=head2 forward

    my @logits = $inference->forward(\@tokens);

Runs the transformer forward pass on input tokens.

B<Parameters:>

=over 4

=item * C<\@tokens> - Array reference of token IDs (integers)

=back

B<Returns:> A list of logits (one per vocabulary token).

B<Details:>

The forward pass:

=over 4

=item 1. Looks up token embeddings

=item 2. Applies N transformer layers with attention and FFN

=item 3. Applies final normalization

=item 4. Projects to vocabulary size

=item 5. Returns logits for the I<last> token position

=back

B<Performance Notes:>

=over 4

=item * Each call creates a new computation graph

=item * Memory is allocated and freed for each call

=item * For multi-token generation, consider batching

=back

B<Example:>

    my @tokens = (1, 450, 7483, 310, 3444, 338);  # "The capital of France is"
    my @logits = $inference->forward(\@tokens);
    
    # logits has 32000 elements (vocab size)
    print "Vocab size: ", scalar(@logits), "\n";

=head2 sample_top_p

    my $token_id = $inference->sample_top_p(
        \@logits,
        temperature => 0.8,
        top_p => 0.95
    );

Samples a token from logits using nucleus (top-p) sampling.

B<Parameters:>

=over 4

=item * C<\@logits> - Array reference of logits from forward()

=item * C<temperature> - Sampling temperature (default: 0.8)

=item * C<top_p> - Cumulative probability threshold (default: 0.95)

=back

B<Returns:> A single token ID.

B<Algorithm:>

=over 4

=item 1. Apply temperature scaling: logit / temperature

=item 2. Convert to probabilities via softmax

=item 3. Sort tokens by probability (descending)

=item 4. Keep tokens until cumulative probability >= top_p

=item 5. Randomly sample from this "nucleus"

=back

B<Temperature Effects:>

=over 4

=item * B<< temperature < 1 >> - More deterministic (sharper distribution)

=item * B<temperature = 1> - Use raw probabilities

=item * B<< temperature > 1 >> - More random (flatter distribution)

=item * B<temperature → 0> - Approaches greedy (argmax)

=back

B<Top-p Effects:>

=over 4

=item * B<top_p = 0.1> - Very focused, only most likely tokens

=item * B<top_p = 0.9> - Typical setting, good balance

=item * B<top_p = 1.0> - Consider all tokens (like top-k with k=all)

=back

B<Example:>

    my @logits = $inference->forward(\@tokens);
    
    # Creative generation
    my $token = $inference->sample_top_p(\@logits,
        temperature => 1.0,
        top_p => 0.95
    );
    
    # More focused generation
    my $token = $inference->sample_top_p(\@logits,
        temperature => 0.3,
        top_p => 0.5
    );

=head2 sample_top_k

    my $token_id = $inference->sample_top_k(
        \@logits,
        temperature => 0.8,
        top_k => 40
    );

Samples a token from logits using top-k sampling.

B<Parameters:>

=over 4

=item * C<\@logits> - Array reference of logits from forward()

=item * C<temperature> - Sampling temperature (default: 0.8)

=item * C<top_k> - Number of top tokens to consider (default: 40)

=back

B<Returns:> A single token ID.

B<Algorithm:>

=over 4

=item 1. Apply temperature scaling: logit / temperature

=item 2. Convert to probabilities via softmax

=item 3. Select top-k tokens by probability

=item 4. Renormalize probabilities

=item 5. Randomly sample from top-k tokens

=back

B<Example:>

    my @logits = $inference->forward(\@tokens);
    
    # Sample from top 50 tokens
    my $token = $inference->sample_top_k(\@logits,
        temperature => 0.9,
        top_k => 50
    );

=head2 generate

    my @tokens = $inference->generate(
        \@prompt_tokens,
        max_tokens  => 100,
        temperature => 0.8,
        top_p       => 0.95,
        top_k       => 40,
        greedy      => 0,
        eos_token   => 2,
        callback    => sub { ... },
    );

Generates multiple tokens autoregressively from a prompt.

B<Parameters:>

=over 4

=item * C<\@prompt_tokens> (required) - Array reference of prompt token IDs

=item * C<max_tokens> - Maximum tokens to generate (default: 128)

=item * C<temperature> - Sampling temperature (default: 0.8)

=item * C<top_p> - Top-p (nucleus) sampling threshold (default: 0.95)

=item * C<top_k> - Top-k sampling limit (default: 40). If < 1000, uses top_k; otherwise uses top_p

=item * C<greedy> - If true, use greedy decoding (argmax) (default: 0)

=item * C<eos_token> - Token ID to stop generation (default: from model, typically 2)

=item * C<callback> - Optional subroutine called for each generated token

=back

B<Returns:> A list of generated token IDs (not including the prompt).

B<Callback:>

The callback receives (token_id, count) and should return true to stop generation:

    callback => sub {
        my ($token, $count) = @_;
        print $tokenizer->decode([$token]);
        return 0;  # Continue (return 1 to stop)
    }

B<Stopping Conditions:>

Generation stops when:

=over 4

=item * max_tokens is reached

=item * EOS token is generated

=item * Callback returns true

=back

B<Example:>

    use Lugh;
    
    my $model = Lugh::Model->new(model => 'model.gguf');
    my $tokenizer = Lugh::Tokenizer->new(model => $model);
    my $inference = Lugh::Inference->new(model => $model);
    
    my @prompt = $tokenizer->encode("Once upon a time");
    
    # Greedy generation
    my @tokens = $inference->generate(\@prompt,
        max_tokens => 50,
        greedy     => 1,
    );
    print $tokenizer->decode(\@tokens);
    
    # Creative generation with streaming
    @tokens = $inference->generate(\@prompt,
        max_tokens  => 100,
        temperature => 1.0,
        top_p       => 0.95,
        callback    => sub {
            my ($tok, $n) = @_;
            print $tokenizer->decode([$tok]);
            STDOUT->flush();
            return 0;
        },
    );

=head1 ATTENTION MECHANISM

=head2 Scaled Dot-Product Attention

    Attention(Q, K, V) = softmax(QK^T / √d_k) × V

Where:

=over 4

=item * Q - Query vectors [head_dim, n_tokens, n_heads]

=item * K - Key vectors [head_dim, n_tokens, n_kv_heads]

=item * V - Value vectors [head_dim, n_tokens, n_kv_heads]

=item * d_k - Head dimension (typically 64-128)

=back

=head2 Grouped Query Attention (GQA)

GQA uses fewer KV heads than query heads to reduce memory:

    Model       n_head  n_kv_head  Ratio
    LLaMA 7B    32      32         1:1 (MHA)
    LLaMA 2 70B 64      8          8:1 (GQA)
    TinyLlama   32      4          8:1 (GQA)
    Mistral 7B  32      8          4:1 (GQA)

The implementation broadcasts KV heads to match query heads using
ggml's native broadcasting.

=head2 Causal Masking

The attention uses causal (autoregressive) masking so each position
can only attend to itself and previous positions:

    Position:  0  1  2  3
    0          ✓  ✗  ✗  ✗
    1          ✓  ✓  ✗  ✗
    2          ✓  ✓  ✓  ✗
    3          ✓  ✓  ✓  ✓

This is implemented using C<ggml_diag_mask_inf> which sets the upper
triangle to -infinity before softmax.

=head2 RoPE (Rotary Position Embeddings)

Position information is encoded by rotating Q and K vectors:

    RoPE(x, pos) = x × cos(pos × θ) + rotate(x) × sin(pos × θ)

Where θ depends on the dimension and base frequency (typically 10000).

Parameters are read from model metadata:

=over 4

=item * C<llama.rope.dimension_count> - Dimensions to rotate

=item * C<llama.rope.freq_base> - Base frequency

=item * C<llama.context_length> - Original context length

=back

=head1 FEED-FORWARD NETWORK

The FFN uses SwiGLU activation:

    FFN(x) = down(gate(x) × SiLU(up(x)))

Where:

=over 4

=item * gate, up - Linear projections to intermediate dimension

=item * SiLU - Sigmoid Linear Unit: x × sigmoid(x)

=item * down - Linear projection back to model dimension

=back

Typical dimensions:

    Model       n_embd  FFN_dim   Ratio
    TinyLlama   2048    5632      2.75×
    LLaMA 7B    4096    11008     2.69×
    LLaMA 13B   5120    13824     2.70×

=head1 GENERATION LOOP

The C<generate()> method handles the complete generation loop internally.
For simple use cases:

    use Lugh;
    
    my $model = Lugh::Model->new(model => 'model.gguf');
    my $tokenizer = Lugh::Tokenizer->new(model => $model);
    my $inference = Lugh::Inference->new(model => $model);
    
    my @prompt = $tokenizer->encode("Once upon a time");
    my @generated = $inference->generate(\@prompt,
        max_tokens  => 100,
        temperature => 0.8,
        top_p       => 0.95,
    );
    print $tokenizer->decode(\@generated);

For streaming output:

    my @generated = $inference->generate(\@prompt,
        max_tokens  => 100,
        temperature => 0.8,
        callback    => sub {
            my ($token, $count) = @_;
            print $tokenizer->decode([$token]);
            STDOUT->flush();
            return 0;  # Continue
        },
    );

For manual control (building your own loop):

    my @tokens = $tokenizer->encode($prompt);
    my @generated;
    
    for (1..$max_tokens) {
        my @logits = $inference->forward(\@tokens);
        my $next = $inference->sample_top_p(\@logits,
            temperature => 0.8,
            top_p => 0.9
        );
        
        last if $next == $tokenizer->eos_id;
        
        push @tokens, $next;
        push @generated, $next;
        
        print $tokenizer->decode([$next]);
        STDOUT->flush();
    }

=head1 PERFORMANCE

=head2 Computation

A single forward pass performs approximately:

    FLOPs ≈ 2 × n_params × n_tokens

For TinyLlama (1.1B params) with 6 tokens:

    2 × 1.1e9 × 6 ≈ 13 GFLOPs

=head2 Memory

During inference, memory is needed for:

=over 4

=item * Model weights (quantized) - Depends on model size and quantization

=item * Activations - O(n_tokens × n_embd × n_layers)

=item * Attention scores - O(n_tokens² × n_heads × n_layers)

=back

=head2 Optimizations

Current implementation:

=over 4

=item * Uses ggml's Metal GPU backend on macOS

=item * Uses Accelerate BLAS for matrix operations

=item * Quantized weights stay quantized during computation

=back

Potential improvements:

=over 4

=item * KV cache for incremental decoding

=item * Batch processing multiple sequences

=item * Flash attention for long contexts

=back

=head1 THREAD SAFETY

Lugh::Inference objects are NOT thread-safe. Each Perl thread must
create its own Inference object.

=head1 SEE ALSO

L<Lugh>, L<Lugh::Model>, L<Lugh::Tokenizer>

L<https://arxiv.org/abs/1706.03762> - "Attention Is All You Need"

L<https://arxiv.org/abs/2104.09864> - RoPE paper

L<https://arxiv.org/abs/2002.05202> - SwiGLU activation

L<https://arxiv.org/abs/2305.13245> - GQA paper

=head1 AUTHOR

lnation E<lt>email@lnation.orgE<gt>

=head1 LICENSE

This library is free software; you can redistribute it and/or modify
it under the same terms as Perl itself.

=cut

1;
