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

const months_per_year: f64 = 12.0;

/// Risk metrics for a specific trailing period, computed from monthly returns.
pub const RiskMetrics = struct {
    /// Annualized standard deviation of monthly returns
    volatility: f64,
    /// Sharpe ratio: (annualized return - risk-free rate) / annualized volatility
    sharpe: f64,
    /// Maximum drawdown as a positive decimal (e.g., 0.30 = 30% drawdown)
    max_drawdown: f64,
    /// Number of monthly returns used
    sample_size: usize,
};

/// Risk metrics for all standard trailing periods.
pub const TrailingRisk = struct {
    one_year: ?RiskMetrics = null,
    three_year: ?RiskMetrics = null,
    five_year: ?RiskMetrics = null,
    ten_year: ?RiskMetrics = null,
};

/// Average annual 3-month T-bill rate by year (source: FRED series DTB3).
/// Used to compute period-appropriate risk-free rates for Sharpe ratio.
/// Update annually - bump `tbill_rates_last_updated` below when you
/// refresh the table. `src/data/staleness.zig` nags on stderr every
/// invocation once it's past the annual due date (Jan 31).
///
/// To update mid-year (e.g. refresh the current year's YTD average):
///   curl -s "https://fred.stlouisfed.org/graph/fredgraph.csv?id=DTB3&cosd=2026-01-01&coed=2026-12-31&fq=Daily" \
///     | tail -n +2 | awk -F, '{s+=$2;n++} END {printf "%.4f\n", s/n/100}'
///
/// To get the prior year's final average (e.g. 2026 final, run in Jan 2027):
///   curl -s "https://fred.stlouisfed.org/graph/fredgraph.csv?id=DTB3&cosd=2026-01-01&coed=2026-12-31&fq=Annual&fam=avg" \
///     | tail -1 | awk -F, '{printf "%.4f\n", $2/100}'
///   Then add a new entry for 2027 using the mid-year command above.
///   Finally, bump `tbill_rates_last_updated` to today's date.
///
/// Registered with the staleness checker in `src/data/staleness.zig`.
pub const tbill_rates_last_updated: Date = Date.fromYmd(2026, 3, 15);

const tbill_rates = [_]struct { year: u16, rate: f64 }{
    .{ .year = 2015, .rate = 0.0005 },
    .{ .year = 2016, .rate = 0.0032 },
    .{ .year = 2017, .rate = 0.0093 },
    .{ .year = 2018, .rate = 0.0194 },
    .{ .year = 2019, .rate = 0.0206 },
    .{ .year = 2020, .rate = 0.0035 },
    .{ .year = 2021, .rate = 0.0005 },
    .{ .year = 2022, .rate = 0.0202 },
    .{ .year = 2023, .rate = 0.0507 },
    .{ .year = 2024, .rate = 0.0497 },
    .{ .year = 2025, .rate = 0.0407 },
    .{ .year = 2026, .rate = 0.0345 },
};

/// Look up the average risk-free rate for a date range from the T-bill
/// table. Returns the simple average of annual rates for all years
/// that overlap the range. Used both for per-symbol Sharpe (`computeRisk`
/// in this file) and for synthetic-portfolio Sharpe
/// (`analytics/portfolio_risk.zig`) so both paths use the same rate.
pub fn avgRiskFreeRateForRange(start: Date, end: Date) f64 {
    const start_year: u16 = @intCast(start.year());
    const end_year: u16 = @intCast(end.year());

    var sum: f64 = 0;
    var count: f64 = 0;
    for (tbill_rates) |entry| {
        if (entry.year >= start_year and entry.year <= end_year) {
            sum += entry.rate;
            count += 1;
        }
    }
    if (count > 0) return sum / count;
    // Fallback: return the latest available rate
    return tbill_rates[tbill_rates.len - 1].rate;
}

/// Compute trailing risk metrics (1Y, 3Y, 5Y, 10Y) from daily candles.
/// Uses monthly total returns (from adj_close, which includes dividends)
/// to match Morningstar's methodology. Risk-free rate is the average
/// 3-month T-bill rate over each period (from FRED historical data).
/// Candles must be sorted by date ascending.
pub fn trailingRisk(candles: []const Candle) TrailingRisk {
    if (candles.len == 0) return .{};
    const end_date = candles[candles.len - 1].date;

    const start_1y = end_date.subtractYears(1);
    const start_3y = end_date.subtractYears(3);
    const start_5y = end_date.subtractYears(5);
    const start_10y = end_date.subtractYears(10);

    return .{
        .one_year = computeRisk(candles, start_1y, end_date, avgRiskFreeRateForRange(start_1y, end_date)),
        .three_year = computeRisk(candles, start_3y, end_date, avgRiskFreeRateForRange(start_3y, end_date)),
        .five_year = computeRisk(candles, start_5y, end_date, avgRiskFreeRateForRange(start_5y, end_date)),
        .ten_year = computeRisk(candles, start_10y, end_date, avgRiskFreeRateForRange(start_10y, end_date)),
    };
}

/// Maximum monthly observations we'll resample from a 10+ year window.
/// Public so callers (notably `analytics/portfolio_risk.zig`) can size
/// their stack-allocated buffers consistently.
pub const max_months: usize = 130;

/// Stats derived from a month-end return series. Pure-data result of
/// `statsFromMonthlyReturns`; callers wrap it (or extend it with date
/// attribution) for their richer outward types like `RiskMetrics`.
pub const MonthlyStats = struct {
    /// Annualized standard deviation of monthly returns.
    volatility: f64,
    /// Sharpe ratio: (annualized return - risk-free rate) / annualized vol.
    sharpe: f64,
    /// Maximum drawdown over the synthetic compound series, as a
    /// positive decimal (e.g. 0.30 = 30% drawdown).
    max_drawdown: f64,
    /// Number of monthly returns the stats were derived from.
    sample_size: usize,
};

/// Result of `monthEndReturns`: a slice into the caller's buffer
/// containing one return per month transition, plus the parallel
/// month-end dates for date-attribution callers (e.g. drawdown
/// start/trough). The `dates` slice is one element LONGER than
/// `returns` because returns are between consecutive month-ends.
pub const MonthEndSeries = struct {
    returns: []const f64,
    dates: []const Date,
};

/// Resample daily candles to a month-end return series, scoped to
/// `[start, end]`. Stack-only - caller provides two scratch buffers
/// of size at least `max_months`. Returns null when the period
/// isn't sufficiently covered:
///
///   - no candles inside the window, OR
///   - data starts more than 45 days after `start`, OR
///   - fewer than 2 month-end observations (need at least 1 return).
///
/// Uses `Candle.adj_close` so synthetic returns include dividends
/// already baked into the adjusted price (Tiingo / Polygon
/// adj_close are split-and-dividend adjusted).
pub fn monthEndReturns(
    candles: []const Candle,
    start: Date,
    end: Date,
    out_returns: []f64,
    out_dates: []Date,
) ?MonthEndSeries {
    std.debug.assert(out_returns.len >= max_months);
    std.debug.assert(out_dates.len >= max_months);

    // Find the candles within [start, end].
    var first: usize = 0;
    for (candles, 0..) |c, i| {
        if (c.date.days >= start.days) {
            first = i;
            break;
        }
    } else return null;

    var last: usize = first;
    for (candles[first..], first..) |c, i| {
        if (c.date.days > end.days) break;
        last = i;
    }

    const slice = candles[first .. last + 1];
    if (slice.len < 2) return null;

    // Period not sufficiently covered.
    if (slice[0].date.days - start.days > 45) return null;

    // Resample: for each calendar month, take the last available
    // close. We track the closes in a local buffer to avoid stomping
    // the caller's `out_returns` (which we'll fill with returns after
    // the resample completes).
    var month_closes: [max_months]f64 = undefined;
    var n_months: usize = 0;
    var prev_ym: u32 = slice[0].date.yearMonth();
    var last_close: f64 = slice[0].adj_close;
    var last_date: Date = slice[0].date;

    for (slice[1..]) |c| {
        const ym = c.date.yearMonth();
        if (ym != prev_ym) {
            if (n_months < max_months) {
                month_closes[n_months] = last_close;
                out_dates[n_months] = last_date;
                n_months += 1;
            }
            prev_ym = ym;
        }
        last_close = c.adj_close;
        last_date = c.date;
    }
    // Record the final (possibly partial) month.
    if (n_months < max_months) {
        month_closes[n_months] = last_close;
        out_dates[n_months] = last_date;
        n_months += 1;
    }

    if (n_months < 2) return null;

    // Compute month-over-month returns into `out_returns`.
    const n_returns = n_months - 1;
    for (0..n_returns) |i| {
        const prev = month_closes[i];
        const curr = month_closes[i + 1];
        out_returns[i] = if (prev > 0) (curr / prev) - 1.0 else 0;
    }

    return .{
        .returns = out_returns[0..n_returns],
        .dates = out_dates[0..n_months],
    };
}

/// Compute volatility / Sharpe / max-drawdown from a month-end
/// return series. Pure: no allocation, no I/O. Returns
/// zero-vol / zero-sharpe when the input is empty or trivially
/// non-volatile; callers that want a "no data" sentinel should
/// gate on `returns.len` themselves.
///
/// MaxDD walks a synthetic compound series anchored at 1.0. This
/// is mathematically equivalent to walking month-end prices -
/// `(p_t / p_0)` is exactly the compounded return - but lets the
/// helper operate on returns alone, which is what the synthetic-
/// portfolio path produces.
pub fn statsFromMonthlyReturns(returns: []const f64, risk_free_rate: f64) MonthlyStats {
    var sum: f64 = 0;
    var sum_sq: f64 = 0;

    var price: f64 = 1.0;
    var peak: f64 = 1.0;
    var max_dd: f64 = 0;

    for (returns) |r| {
        sum += r;
        sum_sq += r * r;
        price *= (1.0 + r);
        if (price > peak) peak = price;
        if (peak > 0) {
            const dd = (peak - price) / peak;
            if (dd > max_dd) max_dd = dd;
        }
    }

    const n = returns.len;
    const nf: f64 = @floatFromInt(n);
    const mean: f64 = if (n > 0) sum / nf else 0;
    // Sample variance (Bessel's correction) for unbiased estimate.
    const variance: f64 = if (n > 1)
        (sum_sq - nf * mean * mean) / (nf - 1.0)
    else
        0;
    const monthly_vol = @sqrt(@max(variance, 0));
    const annual_vol = monthly_vol * @sqrt(months_per_year);
    const annual_return = std.math.pow(f64, 1.0 + mean, months_per_year) - 1.0;
    const sharpe: f64 = if (annual_vol > 0) (annual_return - risk_free_rate) / annual_vol else 0;

    return .{
        .volatility = annual_vol,
        .sharpe = sharpe,
        .max_drawdown = max_dd,
        .sample_size = n,
    };
}

/// Compute risk metrics for a specific date range using monthly returns.
/// Returns null if fewer than 12 monthly returns are available.
fn computeRisk(candles: []const Candle, start: Date, end: Date, risk_free_rate: f64) ?RiskMetrics {
    var ret_buf: [max_months]f64 = undefined;
    var date_buf: [max_months]Date = undefined;
    const series = monthEndReturns(candles, start, end, &ret_buf, &date_buf) orelse return null;
    if (series.returns.len < 12) return null; // need at least 12 monthly returns

    const stats = statsFromMonthlyReturns(series.returns, risk_free_rate);
    return .{
        .volatility = stats.volatility,
        .sharpe = stats.sharpe,
        .max_drawdown = stats.max_drawdown,
        .sample_size = stats.sample_size,
    };
}

// ── Tests ────────────────────────────────────────────────────

fn makeCandle(date: Date, price: f64) Candle {
    return .{ .date = date, .open = price, .high = price, .low = price, .close = price, .adj_close = price, .volume = 1000 };
}

test "trailingRisk 3-year period" {
    // Build 3+ years of monthly data (40 months).
    // Price goes from 100 to 140 linearly over 40 months.
    var candles: [40 * 21]Candle = undefined; // ~21 trading days per month
    var idx: usize = 0;
    var date = Date.fromYmd(2022, 1, 3);

    for (0..40) |month| {
        const base_price: f64 = 100.0 + @as(f64, @floatFromInt(month));
        for (0..21) |day| {
            if (idx >= candles.len) break;
            // Small daily variation within the month
            const daily_offset: f64 = @as(f64, @floatFromInt(day)) * 0.01;
            candles[idx] = makeCandle(date, base_price + daily_offset);
            date = date.addDays(1);
            // Skip weekends
            while (date.dayOfWeek() >= 5) date = date.addDays(1);
            idx += 1;
        }
    }

    const tr = trailingRisk(candles[0..idx]);

    // 3-year should exist (40 months > 36)
    try std.testing.expect(tr.three_year != null);
    const m3 = tr.three_year.?;
    // Monthly returns are ~1% (price goes up ~1/month from ~100)
    // Annualized return should be roughly 12-13%
    // Sharpe should be positive
    try std.testing.expect(m3.sharpe > 0);
    try std.testing.expect(m3.volatility > 0);
    try std.testing.expect(m3.sample_size >= 35); // ~36 monthly returns

    // 5-year should be null (only ~40 months of data, period not covered)
    try std.testing.expect(tr.five_year == null);
}

test "trailingRisk max drawdown" {
    // Build 2 years of data with a 20% drawdown in the middle
    var candles: [500]Candle = undefined;
    var idx: usize = 0;
    var date = Date.fromYmd(2023, 1, 3);

    // 6 months up: 100 -> 120
    for (0..126) |i| {
        const price = 100.0 + @as(f64, @floatFromInt(i)) * (20.0 / 126.0);
        candles[idx] = makeCandle(date, price);
        date = date.addDays(1);
        while (date.dayOfWeek() >= 5) date = date.addDays(1);
        idx += 1;
    }
    // 3 months down: 120 -> 96 (20% drawdown from 120)
    for (0..63) |i| {
        const price = 120.0 - @as(f64, @floatFromInt(i)) * (24.0 / 63.0);
        candles[idx] = makeCandle(date, price);
        date = date.addDays(1);
        while (date.dayOfWeek() >= 5) date = date.addDays(1);
        idx += 1;
    }
    // 6 months recovery: 96 -> 130
    for (0..126) |i| {
        const price = 96.0 + @as(f64, @floatFromInt(i)) * (34.0 / 126.0);
        candles[idx] = makeCandle(date, price);
        date = date.addDays(1);
        while (date.dayOfWeek() >= 5) date = date.addDays(1);
        idx += 1;
    }

    const tr = trailingRisk(candles[0..idx]);
    try std.testing.expect(tr.one_year != null);
    const m = tr.one_year.?;
    // Drawdown should be approximately 20% (from monthly closes, may not be exact)
    try std.testing.expect(m.max_drawdown > 0.15);
    try std.testing.expect(m.max_drawdown < 0.25);
}

test "trailingRisk insufficient data" {
    // Only 6 months of data -> 1-year should be null
    var candles: [126]Candle = undefined;
    var date = Date.fromYmd(2025, 7, 1);
    for (0..126) |i| {
        candles[i] = makeCandle(date, 100.0 + @as(f64, @floatFromInt(i)) * 0.1);
        date = date.addDays(1);
        while (date.dayOfWeek() >= 5) date = date.addDays(1);
    }
    const tr = trailingRisk(&candles);
    // Less than 12 monthly returns for any period
    try std.testing.expect(tr.one_year == null);
    try std.testing.expect(tr.three_year == null);
}

test "avgRiskFreeRateForRange uses historical T-bill data" {
    // 2023-2025: average of 5.07%, 4.97%, 4.07% = 4.70%
    const rate_3y = avgRiskFreeRateForRange(Date.fromYmd(2023, 3, 16), Date.fromYmd(2025, 12, 31));
    try std.testing.expectApproxEqAbs(@as(f64, 0.047), rate_3y, 0.002);

    // 2020-2025: includes the near-zero years
    const rate_5y = avgRiskFreeRateForRange(Date.fromYmd(2020, 1, 1), Date.fromYmd(2025, 12, 31));
    try std.testing.expect(rate_5y < rate_3y); // should be lower due to 2020-2021 near-zero rates
}

test "monthEndReturns: empty range returns null" {
    var ret_buf: [max_months]f64 = undefined;
    var date_buf: [max_months]Date = undefined;
    const candles: []const Candle = &.{};
    const series = monthEndReturns(candles, Date.fromYmd(2024, 1, 1), Date.fromYmd(2025, 1, 1), &ret_buf, &date_buf);
    try std.testing.expect(series == null);
}

test "monthEndReturns: data starting >45 days late returns null" {
    var ret_buf: [max_months]f64 = undefined;
    var date_buf: [max_months]Date = undefined;
    // Window starts 2022-01-01, but candles start 2022-04-01 (90 days late).
    var candles: [60]Candle = undefined;
    var d = Date.fromYmd(2022, 4, 1);
    for (0..candles.len) |i| {
        candles[i] = makeCandle(d, 100.0 + @as(f64, @floatFromInt(i)));
        d = d.addDays(7);
    }
    const series = monthEndReturns(&candles, Date.fromYmd(2022, 1, 1), Date.fromYmd(2024, 1, 1), &ret_buf, &date_buf);
    try std.testing.expect(series == null);
}

test "monthEndReturns: happy path returns one fewer return than month-end dates" {
    var ret_buf: [max_months]f64 = undefined;
    var date_buf: [max_months]Date = undefined;
    // 6 months of data, daily candles.
    var candles: [180]Candle = undefined;
    var d = Date.fromYmd(2024, 1, 3);
    for (0..candles.len) |i| {
        candles[i] = makeCandle(d, 100.0 + @as(f64, @floatFromInt(i)) * 0.1);
        d = d.addDays(1);
        while (d.dayOfWeek() >= 5) d = d.addDays(1);
    }
    const series = monthEndReturns(&candles, Date.fromYmd(2024, 1, 1), Date.fromYmd(2025, 1, 1), &ret_buf, &date_buf);
    try std.testing.expect(series != null);
    const s = series.?;
    // returns is one shorter than dates: returns[i] is the return
    // from dates[i] to dates[i+1].
    try std.testing.expectEqual(s.dates.len - 1, s.returns.len);
    try std.testing.expect(s.returns.len >= 5); // at least Feb..Jun returns
}

test "statsFromMonthlyReturns: zero-variance series has zero vol" {
    const constant = [_]f64{ 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01 };
    const stats = statsFromMonthlyReturns(&constant, 0.04);
    try std.testing.expectApproxEqAbs(@as(f64, 0), stats.volatility, 0.0001);
    try std.testing.expectEqual(@as(usize, 12), stats.sample_size);
    // Drawdown is zero - every month is +1%, no peak retreats.
    try std.testing.expectApproxEqAbs(@as(f64, 0), stats.max_drawdown, 0.0001);
}

test "statsFromMonthlyReturns: max drawdown matches hand calc" {
    // Up 10%, up 10%, down 20%, up 5%. Peak at month 2 (1.21);
    // trough at month 3 (1.21 × 0.80 = 0.968). Drawdown =
    // (1.21 - 0.968) / 1.21 ≈ 0.20.
    const returns = [_]f64{ 0.10, 0.10, -0.20, 0.05 };
    const stats = statsFromMonthlyReturns(&returns, 0.04);
    try std.testing.expectApproxEqAbs(@as(f64, 0.20), stats.max_drawdown, 0.001);
}

test "statsFromMonthlyReturns: empty input returns zeros" {
    const stats = statsFromMonthlyReturns(&.{}, 0.04);
    try std.testing.expectEqual(@as(f64, 0), stats.volatility);
    try std.testing.expectEqual(@as(f64, 0), stats.sharpe);
    try std.testing.expectEqual(@as(f64, 0), stats.max_drawdown);
    try std.testing.expectEqual(@as(usize, 0), stats.sample_size);
}