const std = @import("std");

/// Temperature + top-k sampling for autoregressive generation.
pub const Sampler = struct {
    temperature: f32,
    top_k: u32,
    xoshiro: std.Random.Xoshiro256,

    pub fn init(temperature: f32, top_k: u32, seed: u64) Sampler {
        return .{
            .temperature = temperature,
            .top_k = top_k,
            .xoshiro = std.Random.Xoshiro256.init(seed),
        };
    }

    /// Sample a token index from logits.
    /// Applies temperature scaling, then top-k filtering, then categorical sampling.
    pub fn sample(self: *Sampler, logits: []f32) u32 {
        const n = logits.len;
        if (n == 0) return 0;

        // Temperature scaling
        if (self.temperature != 1.0 and self.temperature > 0.0) {
            const inv_temp = 1.0 / self.temperature;
            for (logits) |*l| {
                l.* *= inv_temp;
            }
        }

        // Greedy (temperature ~0)
        if (self.temperature <= 0.0 or self.top_k == 1) {
            return argmax(logits);
        }

        // Top-k: find the k-th largest value and zero out everything below
        const k = @min(self.top_k, @as(u32, @intCast(n)));

        // Find k-th largest using partial sort via repeated max finding
        // For small k this is efficient enough
        var threshold: f32 = -std.math.inf(f32);
        if (k < n) {
            // Use a simple approach: find the k-th largest value
            // by maintaining the minimum of the top-k seen so far
            var top_k_min: f32 = std.math.inf(f32);
            var top_k_vals: [256]f32 = undefined;
            const actual_k: usize = @min(k, 256);
            var count: usize = 0;

            for (logits) |l| {
                if (count < actual_k) {
                    top_k_vals[count] = l;
                    count += 1;
                    if (count == actual_k) {
                        // Find min of top_k_vals
                        top_k_min = top_k_vals[0];
                        for (top_k_vals[1..actual_k]) |v| {
                            if (v < top_k_min) top_k_min = v;
                        }
                    }
                } else if (l > top_k_min) {
                    // Replace the min
                    for (&top_k_vals, 0..) |*v, i| {
                        _ = i;
                        if (v.* == top_k_min) {
                            v.* = l;
                            break;
                        }
                    }
                    top_k_min = top_k_vals[0];
                    for (top_k_vals[1..actual_k]) |v| {
                        if (v < top_k_min) top_k_min = v;
                    }
                }
            }
            threshold = top_k_min;
        }

        // Softmax with top-k filtering
        var max_val: f32 = -std.math.inf(f32);
        for (logits) |l| {
            if (l >= threshold and l > max_val) max_val = l;
        }

        var sum: f32 = 0.0;
        for (logits) |*l| {
            if (l.* >= threshold) {
                l.* = @exp(l.* - max_val);
                sum += l.*;
            } else {
                l.* = 0.0;
            }
        }

        if (sum <= 0.0) return argmax(logits);

        // Normalize
        const inv_sum = 1.0 / sum;
        for (logits) |*l| {
            l.* *= inv_sum;
        }

        // Categorical sampling
        var r = self.xoshiro.random().float(f32);
        for (logits, 0..) |p, i| {
            r -= p;
            if (r <= 0.0) return @intCast(i);
        }

        return @intCast(n - 1);
    }
};

fn argmax(vals: []const f32) u32 {
    var best: u32 = 0;
    var best_val: f32 = -std.math.inf(f32);
    for (vals, 0..) |v, i| {
        if (v > best_val) {
            best_val = v;
            best = @intCast(i);
        }
    }
    return best;
}

test "sampler greedy" {
    var logits = [_]f32{ 1.0, 3.0, 2.0, 0.5 };
    var s = Sampler.init(0.0, 1, 42);
    const token = s.sample(&logits);
    try std.testing.expectEqual(@as(u32, 1), token);
}

test "sampler top_k" {
    var logits = [_]f32{ 1.0, 10.0, 9.0, 0.5 };
    var s = Sampler.init(1.0, 2, 42);
    const token = s.sample(&logits);
    // Should be either 1 or 2 (top-2)
    try std.testing.expect(token == 1 or token == 2);
}