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,
    /// Start date of max drawdown period
    drawdown_start: ?Date = null,
    /// Trough date of max drawdown
    drawdown_trough: ?Date = null,
    /// 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.
fn avgRiskFreeRate(start: Date, end: Date) f64 {
    const start_year = @as(u16, @intCast(start.year()));
    const end_year = @as(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, avgRiskFreeRate(start_1y, end_date)),
        .three_year = computeRisk(candles, start_3y, end_date, avgRiskFreeRate(start_3y, end_date)),
        .five_year = computeRisk(candles, start_5y, end_date, avgRiskFreeRate(start_5y, end_date)),
        .ten_year = computeRisk(candles, start_10y, end_date, avgRiskFreeRate(start_10y, end_date)),
    };
}

/// 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 {
    // Find the slice of 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; // no candles in range

    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;

    // If data starts more than 45 days after the requested start,
    // the period isn't sufficiently covered — don't report misleading metrics.
    if (slice[0].date.days - start.days > 45) return null;

    // Resample to month-end closes: for each calendar month, take the last available close.
    // We store up to 130 month-end prices (10+ years).
    const max_months = 130;
    var month_closes: [max_months]f64 = undefined;
    var month_dates: [max_months]Date = undefined;
    var n_months: usize = 0;

    var prev_ym: u32 = yearMonth(slice[0].date);
    var last_close: f64 = slice[0].adj_close;
    var last_date: Date = slice[0].date;

    for (slice[1..]) |c| {
        const ym = yearMonth(c.date);
        if (ym != prev_ym) {
            // Month boundary crossed — record the previous month's last close
            if (n_months < max_months) {
                month_closes[n_months] = last_close;
                month_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;
        month_dates[n_months] = last_date;
        n_months += 1;
    }

    if (n_months < 2) return null;

    // Compute monthly returns
    const n_returns = n_months - 1;
    if (n_returns < 12) return null; // need at least 12 monthly returns

    var sum: f64 = 0;
    var sum_sq: f64 = 0;

    // Max drawdown from monthly closes
    var peak: f64 = month_closes[0];
    var max_dd: f64 = 0;
    var dd_start: ?Date = null;
    var dd_trough: ?Date = null;
    var current_dd_start: Date = month_dates[0];

    for (1..n_months) |i| {
        const prev = month_closes[i - 1];
        const curr = month_closes[i];
        if (prev <= 0) continue;

        const ret = (curr / prev) - 1.0;
        sum += ret;
        sum_sq += ret * ret;

        // Drawdown tracking on monthly closes
        if (curr > peak) {
            peak = curr;
            current_dd_start = month_dates[i];
        }
        const dd = (peak - curr) / peak;
        if (dd > max_dd) {
            max_dd = dd;
            dd_start = current_dd_start;
            dd_trough = month_dates[i];
        }
    }

    const nf = @as(f64, @floatFromInt(n_returns));
    const mean = sum / nf;
    // Use sample variance (n-1) for unbiased estimate
    const variance = if (n_returns > 1)
        (sum_sq - nf * mean * mean) / (nf - 1.0)
    else
        0;
    const monthly_vol = @sqrt(@max(variance, 0));

    // Annualize using standard monthly-to-annual conversion
    const annual_vol = monthly_vol * @sqrt(months_per_year);
    // Geometric annualized return
    const annual_return = std.math.pow(f64, 1.0 + mean, months_per_year) - 1.0;
    const sharpe = if (annual_vol > 0) (annual_return - risk_free_rate) / annual_vol else 0;

    return .{
        .volatility = annual_vol,
        .sharpe = sharpe,
        .max_drawdown = max_dd,
        .drawdown_start = dd_start,
        .drawdown_trough = dd_trough,
        .sample_size = n_returns,
    };
}

/// Encode year+month as a single comparable integer (e.g. 2026*100+3 = 202603).
fn yearMonth(date: Date) u32 {
    const y: u32 = @intCast(date.year());
    const m: u32 = date.month();
    return y * 100 + m;
}

// ── 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 "yearMonth encoding" {
    try std.testing.expectEqual(@as(u32, 202603), yearMonth(Date.fromYmd(2026, 3, 16)));
    try std.testing.expectEqual(@as(u32, 202412), yearMonth(Date.fromYmd(2024, 12, 1)));
}

test "avgRiskFreeRate uses historical T-bill data" {
    // 2023-2025: average of 5.07%, 4.97%, 4.07% = 4.70%
    const rate_3y = avgRiskFreeRate(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 = avgRiskFreeRate(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
}
