//! 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.
///
/// Uses O(n) sliding window algorithm instead of O(n * period):
/// - Maintains running sum for SMA
/// - Maintains running sum of squares for variance: Var(X) = E[X²] - E[X]²
pub fn bollingerBands(
    alloc: std.mem.Allocator,
    closes: []const f64,
    period: usize,
    k: f64,
) ![]?BollingerBand {
    const result = try alloc.alloc(?BollingerBand, closes.len);

    if (closes.len < period or period == 0) {
        @memset(result, null);
        return result;
    }

    const p_f: f64 = @floatFromInt(period);

    // Initialize running sums for the first window [0..period)
    var sum: f64 = 0;
    var sum_sq: f64 = 0;
    for (0..period) |i| {
        sum += closes[i];
        sum_sq += closes[i] * closes[i];
    }

    // First period-1 values are null (not enough data points)
    for (0..period - 1) |i| {
        result[i] = null;
    }

    // Compute for index period-1 (first valid point)
    {
        const mean = sum / p_f;
        // Variance via E[X²] - E[X]² formula
        const variance = (sum_sq / p_f) - (mean * mean);
        // Use @max to guard against tiny negative values from floating point error
        const stddev = @sqrt(@max(variance, 0.0));
        result[period - 1] = .{
            .upper = mean + k * stddev,
            .middle = mean,
            .lower = mean - k * stddev,
        };
    }

    // Slide the window for remaining points: O(n - period) iterations, O(1) each
    for (period..closes.len) |i| {
        const old_val = closes[i - period];
        const new_val = closes[i];

        // Update running sums in O(1)
        sum = sum - old_val + new_val;
        sum_sq = sum_sq - (old_val * old_val) + (new_val * new_val);

        const mean = sum / p_f;
        const variance = (sum_sq / p_f) - (mean * mean);
        const stddev = @sqrt(@max(variance, 0.0));

        result[i] = .{
            .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);
}

test "bollingerBands sliding window correctness" {
    const alloc = std.testing.allocator;
    // Test with realistic price data
    const closes = [_]f64{
        100.0, 101.5, 99.8,  102.3, 103.1, 101.9, 104.2, 105.0,
        103.8, 106.2, 107.1, 105.5, 108.0, 109.2, 107.8, 110.5,
        111.3, 109.8, 112.4, 113.0,
    };
    const bands = try bollingerBands(alloc, &closes, 5, 2.0);
    defer alloc.free(bands);

    // First 4 should be null
    for (0..4) |i| {
        try std.testing.expect(bands[i] == null);
    }

    // Verify a few points manually
    // Index 4: window is [100.0, 101.5, 99.8, 102.3, 103.1]
    // Mean = 101.34, manually computed
    const b4 = bands[4].?;
    try std.testing.expectApproxEqAbs(@as(f64, 101.34), b4.middle, 0.01);
    try std.testing.expect(b4.upper > b4.middle);
    try std.testing.expect(b4.lower < b4.middle);

    // Index 19: window is [109.8, 112.4, 113.0, 111.3, 110.5] — wait, let me recalculate
    // Window at i=19 is closes[15..20] = [110.5, 111.3, 109.8, 112.4, 113.0]
    // Mean = (110.5 + 111.3 + 109.8 + 112.4 + 113.0) / 5 = 557.0 / 5 = 111.4
    const b19 = bands[19].?;
    try std.testing.expectApproxEqAbs(@as(f64, 111.4), b19.middle, 0.01);

    // Verify bands are properly ordered
    for (bands) |b_opt| {
        if (b_opt) |b| {
            try std.testing.expect(b.upper >= b.middle);
            try std.testing.expect(b.middle >= b.lower);
        }
    }
}

test "bollingerBands edge cases" {
    const alloc = std.testing.allocator;

    // Empty data
    {
        const empty: []const f64 = &.{};
        const bands = try bollingerBands(alloc, empty, 5, 2.0);
        defer alloc.free(bands);
        try std.testing.expectEqual(@as(usize, 0), bands.len);
    }

    // Data shorter than period
    {
        const short = [_]f64{ 1, 2, 3 };
        const bands = try bollingerBands(alloc, &short, 5, 2.0);
        defer alloc.free(bands);
        for (bands) |b| try std.testing.expect(b == null);
    }

    // Period = 1 (each point is its own window, stddev = 0)
    {
        const data = [_]f64{ 10, 20, 30 };
        const bands = try bollingerBands(alloc, &data, 1, 2.0);
        defer alloc.free(bands);
        // With period=1, stddev=0, so upper=middle=lower
        for (bands, 0..) |b_opt, i| {
            const b = b_opt.?;
            try std.testing.expectApproxEqAbs(data[i], b.middle, 0.001);
            try std.testing.expectApproxEqAbs(data[i], b.upper, 0.001);
            try std.testing.expectApproxEqAbs(data[i], b.lower, 0.001);
        }
    }

    // Constant data (stddev = 0)
    {
        const constant = [_]f64{ 50, 50, 50, 50, 50 };
        const bands = try bollingerBands(alloc, &constant, 3, 2.0);
        defer alloc.free(bands);
        for (bands[2..]) |b_opt| {
            const b = b_opt.?;
            try std.testing.expectApproxEqAbs(@as(f64, 50), b.middle, 0.001);
            try std.testing.expectApproxEqAbs(@as(f64, 50), b.upper, 0.001);
            try std.testing.expectApproxEqAbs(@as(f64, 50), b.lower, 0.001);
        }
    }
}