//! Technical indicators for financial charting.
//! Bollinger Bands, RSI, SMA — all computed from candle close prices.

const std = @import("std");
const Candle = @import("../models/candle.zig").Candle;

/// Simple Moving Average for a window of `period` values ending at index `end` (inclusive).
/// Returns null if there aren't enough data points.
pub fn sma(closes: []const f64, end: usize, period: usize) ?f64 {
    if (end + 1 < period) return null;
    var sum: f64 = 0;
    const start = end + 1 - period;
    for (closes[start .. end + 1]) |v| sum += v;
    return sum / @as(f64, @floatFromInt(period));
}

/// Bollinger Bands output for a single data point.
pub const BollingerBand = struct {
    upper: f64,
    middle: f64, // SMA
    lower: f64,
};

/// Compute Bollinger Bands (SMA ± k * stddev) for the full series.
/// Returns a slice of optional BollingerBand — null where period hasn't been reached.
pub fn bollingerBands(
    alloc: std.mem.Allocator,
    closes: []const f64,
    period: usize,
    k: f64,
) ![]?BollingerBand {
    const result = try alloc.alloc(?BollingerBand, closes.len);
    for (result, 0..) |*r, i| {
        const mean = sma(closes, i, period) orelse {
            r.* = null;
            continue;
        };
        // Standard deviation
        const start = i + 1 - period;
        var sq_sum: f64 = 0;
        for (closes[start .. i + 1]) |v| {
            const diff = v - mean;
            sq_sum += diff * diff;
        }
        const stddev = @sqrt(sq_sum / @as(f64, @floatFromInt(period)));
        r.* = .{
            .upper = mean + k * stddev,
            .middle = mean,
            .lower = mean - k * stddev,
        };
    }
    return result;
}

/// RSI (Relative Strength Index) for the full series using Wilder's smoothing.
/// Returns a slice of optional f64 — null for the first `period` data points.
pub fn rsi(
    alloc: std.mem.Allocator,
    closes: []const f64,
    period: usize,
) ![]?f64 {
    const result = try alloc.alloc(?f64, closes.len);
    if (closes.len < period + 1) {
        @memset(result, null);
        return result;
    }

    // Seed: average gain/loss over first `period` changes
    var avg_gain: f64 = 0;
    var avg_loss: f64 = 0;
    for (1..period + 1) |i| {
        const change = closes[i] - closes[i - 1];
        if (change > 0) avg_gain += change else avg_loss += -change;
    }
    const p_f: f64 = @floatFromInt(period);
    avg_gain /= p_f;
    avg_loss /= p_f;

    // First `period` values are null
    for (0..period) |i| result[i] = null;
    // Value at index `period`
    if (avg_loss == 0) {
        result[period] = 100.0;
    } else {
        const rs = avg_gain / avg_loss;
        result[period] = 100.0 - (100.0 / (1.0 + rs));
    }

    // Wilder's smoothing for the rest
    for (period + 1..closes.len) |i| {
        const change = closes[i] - closes[i - 1];
        const gain = if (change > 0) change else 0;
        const loss = if (change < 0) -change else 0;
        avg_gain = (avg_gain * (p_f - 1.0) + gain) / p_f;
        avg_loss = (avg_loss * (p_f - 1.0) + loss) / p_f;
        if (avg_loss == 0) {
            result[i] = 100.0;
        } else {
            const rs = avg_gain / avg_loss;
            result[i] = 100.0 - (100.0 / (1.0 + rs));
        }
    }
    return result;
}

/// Extract close prices from candles into a contiguous f64 slice.
pub fn closePrices(alloc: std.mem.Allocator, candles: []const Candle) ![]f64 {
    const result = try alloc.alloc(f64, candles.len);
    for (candles, 0..) |c, i| result[i] = c.close;
    return result;
}

/// Extract volumes from candles.
pub fn volumes(alloc: std.mem.Allocator, candles: []const Candle) ![]f64 {
    const result = try alloc.alloc(f64, candles.len);
    for (candles, 0..) |c, i| result[i] = @floatFromInt(c.volume);
    return result;
}

test "sma basic" {
    const closes = [_]f64{ 1, 2, 3, 4, 5 };
    try std.testing.expectEqual(@as(?f64, null), sma(&closes, 1, 3));
    try std.testing.expectApproxEqAbs(@as(f64, 2.0), sma(&closes, 2, 3).?, 0.001);
    try std.testing.expectApproxEqAbs(@as(f64, 3.0), sma(&closes, 3, 3).?, 0.001);
}

test "rsi basic" {
    const alloc = std.testing.allocator;
    // 15 prices with a clear uptrend
    const closes = [_]f64{ 44, 44.34, 44.09, 43.61, 44.33, 44.83, 45.10, 45.42, 45.84, 46.08, 45.89, 46.03, 45.61, 46.28, 46.28 };
    const result = try rsi(alloc, &closes, 14);
    defer alloc.free(result);
    // First 14 should be null, last should have a value
    try std.testing.expect(result[13] == null);
    try std.testing.expect(result[14] != null);
}

test "bollingerBands basic" {
    const alloc = std.testing.allocator;
    const closes = [_]f64{ 10, 11, 12, 13, 14, 15, 16, 17, 18, 19 };
    const bands = try bollingerBands(alloc, &closes, 5, 2.0);
    defer alloc.free(bands);
    // First 4 (indices 0-3) should be null (period=5, need indices 0..4)
    try std.testing.expect(bands[0] == null);
    try std.testing.expect(bands[3] == null);
    // Index 4 onward should have values
    try std.testing.expect(bands[4] != null);
    const b4 = bands[4].?;
    // SMA of [10,11,12,13,14] = 12.0
    try std.testing.expectApproxEqAbs(@as(f64, 12.0), b4.middle, 0.001);
    // upper > middle > lower
    try std.testing.expect(b4.upper > b4.middle);
    try std.testing.expect(b4.middle > b4.lower);
}

test "closePrices" {
    const alloc = std.testing.allocator;
    const candles = [_]Candle{
        .{ .date = @import("../models/date.zig").Date.fromYmd(2024, 1, 2), .open = 100, .high = 105, .low = 99, .close = 102, .adj_close = 102, .volume = 1000 },
        .{ .date = @import("../models/date.zig").Date.fromYmd(2024, 1, 3), .open = 102, .high = 107, .low = 101, .close = 105, .adj_close = 105, .volume = 2000 },
    };
    const prices = try closePrices(alloc, &candles);
    defer alloc.free(prices);
    try std.testing.expectEqual(@as(usize, 2), prices.len);
    try std.testing.expectApproxEqAbs(@as(f64, 102), prices[0], 0.001);
    try std.testing.expectApproxEqAbs(@as(f64, 105), prices[1], 0.001);
}

test "volumes" {
    const alloc = std.testing.allocator;
    const candles = [_]Candle{
        .{ .date = @import("../models/date.zig").Date.fromYmd(2024, 1, 2), .open = 100, .high = 105, .low = 99, .close = 102, .adj_close = 102, .volume = 1500 },
        .{ .date = @import("../models/date.zig").Date.fromYmd(2024, 1, 3), .open = 102, .high = 107, .low = 101, .close = 105, .adj_close = 105, .volume = 3000 },
    };
    const vols = try volumes(alloc, &candles);
    defer alloc.free(vols);
    try std.testing.expectEqual(@as(usize, 2), vols.len);
    try std.testing.expectApproxEqAbs(@as(f64, 1500), vols[0], 0.001);
    try std.testing.expectApproxEqAbs(@as(f64, 3000), vols[1], 0.001);
}

test "sma edge cases" {
    // period=1: should equal the value itself
    const closes = [_]f64{ 5, 10, 15 };
    try std.testing.expectApproxEqAbs(@as(f64, 5.0), sma(&closes, 0, 1).?, 0.001);
    try std.testing.expectApproxEqAbs(@as(f64, 10.0), sma(&closes, 1, 1).?, 0.001);
    // period > data length: always null
    try std.testing.expect(sma(&closes, 2, 10) == null);
}

test "rsi all up" {
    const alloc = std.testing.allocator;
    // Prices going up by 1 each day for 20 days
    var closes: [20]f64 = undefined;
    for (0..20) |i| closes[i] = 100.0 + @as(f64, @floatFromInt(i));
    const result = try rsi(alloc, &closes, 14);
    defer alloc.free(result);
    // RSI should be 100 (all gains, no losses)
    try std.testing.expect(result[14] != null);
    try std.testing.expectApproxEqAbs(@as(f64, 100.0), result[14].?, 0.001);
}

test "rsi insufficient data" {
    const alloc = std.testing.allocator;
    const closes = [_]f64{ 1, 2, 3 };
    const result = try rsi(alloc, &closes, 14);
    defer alloc.free(result);
    // All should be null since len < period + 1
    for (result) |r| try std.testing.expect(r == null);
}