zfin/src/charts/compare_chart.zig

//! Projection comparison chart: overlays two percentile-band
//! envelopes - a "then" projection (computed as of a past snapshot)
//! and a "now" projection - so the viewer can see how the forecast
//! envelope shifted between the two dates.
//!
//! Both projections are aligned at year 0 (each one's own start), so
//! the x-axis is "years from the projection's start" and the overlay
//! answers "how did my projected envelope move between then and now?".
//!
//! Each side draws a light p10-p90 fill plus a solid median line in
//! its own hue (then = `theme.info` / cyan, now = `theme.accent` /
//! purple), keyed by a small top-left "then"/"now" color legend.
//!
//! Visual layers (bottom to top):
//!   - Background
//!   - Horizontal grid lines
//!   - "then" envelope fill, then "now" envelope fill
//!   - "then" median, then "now" median (both on top of both fills)
//!   - Zero line (if visible)
//!   - Panel border
//!   - then/now color legend (top-left)
//!   - Axis labels (export only)

const std = @import("std");
const z2d = @import("z2d");
const theme = @import("../tui/theme.zig");
const projections = @import("../analytics/projections.zig");
const draw = @import("draw.zig");
const axis = @import("axis.zig");
const text = @import("text.zig");

const Surface = z2d.Surface;
const Context = z2d.Context;

const margin_left: f64 = 4;
const margin_right: f64 = 4;
const margin_top: f64 = 4;
const margin_bottom: f64 = 4;

/// Comparison chart render result (RGB for kitty graphics).
pub const CompareChartResult = struct {
    rgb_data: []const u8,
    width: u16,
    height: u16,
    value_min: f64,
    value_max: f64,
};

/// Owned by the caller - call `result.deinit(alloc)` when done. The
/// surface seam shared between RGB extraction (kitty) and PNG export.
/// Mirrors `projection_chart.RenderedProjection`.
pub const RenderedCompare = struct {
    surface: Surface,
    width: u16,
    height: u16,
    value_min: f64,
    value_max: f64,

    pub fn deinit(self: *RenderedCompare, alloc: std.mem.Allocator) void {
        self.surface.deinit(alloc);
        self.* = undefined;
    }

    pub fn extractRgb(self: *const RenderedCompare, alloc: std.mem.Allocator) ![]u8 {
        return draw.extractRgb(alloc, &self.surface);
    }
};

/// Render the "then" vs "now" comparison overlay into a `Surface`.
/// Both band slices are aligned at year 0; the x-axis spans the
/// longer of the two horizons. With `axis_labels`, reserves margins
/// and stamps y-axis dollar ticks + x-axis year endpoints (export
/// path); the kitty wrapper passes `false`.
pub fn renderToSurface(
    io: std.Io,
    alloc: std.mem.Allocator,
    then_bands: []const projections.YearPercentiles,
    now_bands: []const projections.YearPercentiles,
    width_px: u32,
    height_px: u32,
    th: theme.Theme,
    axis_labels: bool,
) !RenderedCompare {
    if (then_bands.len < 2 or now_bands.len < 2) return error.InsufficientData;

    const w: i32 = @intCast(width_px);
    const h: i32 = @intCast(height_px);
    var sfc = try Surface.init(.image_surface_rgb, alloc, w, h);
    errdefer sfc.deinit(alloc);

    var ctx = Context.init(io, alloc, &sfc);
    defer ctx.deinit();

    ctx.setAntiAliasingMode(.none);
    ctx.setOperator(.src);

    const bg = th.bg;
    const fwidth: f64 = @floatFromInt(width_px);
    const fheight: f64 = @floatFromInt(height_px);

    try draw.fillBackground(&ctx, fwidth, fheight, bg);

    // Chart area. With axis labels, reserve a right margin for the
    // y-axis dollar ticks and a bottom margin for the year endpoints.
    const label_scale: i32 = axis.labelScale(h);
    const label_char_h: f64 = axis.charHeight(label_scale);
    const m_left: f64 = if (axis_labels) label_char_h else margin_left;
    const m_right: f64 = if (axis_labels) axis.yAxisMargin(label_scale) else margin_right;
    const m_top: f64 = if (axis_labels) (label_char_h / 2 + 4) else margin_top;
    const m_bottom: f64 = if (axis_labels) axis.bottomMargin(label_scale) else margin_bottom;
    const chart_left = m_left;
    const chart_right = fwidth - m_right;
    const chart_w = chart_right - chart_left;
    const chart_top = m_top;
    const chart_bottom = fheight - m_bottom;

    // Value range across BOTH envelopes (p10 floor, p90 ceiling).
    var value_min: f64 = then_bands[0].p10;
    var value_max: f64 = then_bands[0].p90;
    for (then_bands) |bp| {
        if (bp.p10 < value_min) value_min = bp.p10;
        if (bp.p90 > value_max) value_max = bp.p90;
    }
    for (now_bands) |bp| {
        if (bp.p10 < value_min) value_min = bp.p10;
        if (bp.p90 > value_max) value_max = bp.p90;
    }
    const pad = (value_max - value_min) * 0.05;
    value_min -= pad;
    value_max += pad;
    if (value_min < 0) value_min = 0;

    // X step: align both at year 0; the longer horizon spans the full
    // width. `bands[i].year == i`, so index doubles as the year offset.
    const n = @max(then_bands.len, now_bands.len);
    const x_step = chart_w / @as(f64, @floatFromInt(n - 1));

    // Grid lines.
    try draw.drawHorizontalGridLines(&ctx, chart_left, chart_right, chart_top, chart_bottom, 5, draw.blendColor(th.text_muted, 40, bg));

    // Envelopes: draw BOTH light fills first, then BOTH medians on
    // top. Rendering uses the `.src` operator (replace, not blend), so
    // drawing a fill after a median would occlude that median - hence
    // the two-pass order. "now" fill goes on top of "then" fill.
    try drawEnvelopeFill(&ctx, then_bands, chart_left, x_step, value_min, value_max, chart_top, chart_bottom, th.info, bg);
    try drawEnvelopeFill(&ctx, now_bands, chart_left, x_step, value_min, value_max, chart_top, chart_bottom, th.accent, bg);
    try drawEnvelopeMedian(&ctx, then_bands, chart_left, x_step, value_min, value_max, chart_top, chart_bottom, th.info);
    try drawEnvelopeMedian(&ctx, now_bands, chart_left, x_step, value_min, value_max, chart_top, chart_bottom, th.accent);

    // Zero line (if visible).
    if (value_min <= 0 and value_max > 0) {
        const zero_y = draw.mapY(0, value_min, value_max, chart_top, chart_bottom);
        try draw.drawHLine(&ctx, chart_left, chart_right, zero_y, draw.blendColor(th.negative, 120, bg), 1.0);
    }

    // Panel border.
    try draw.drawRect(&ctx, chart_left, chart_top, chart_right, chart_bottom, draw.blendColor(th.border, 80, bg), 1.0);

    // Color legend (top-left): keys the two envelope hues. Two
    // unlabeled colored medians would otherwise be ambiguous.
    {
        const lgh: i32 = @intFromFloat(axis.charHeight(label_scale));
        const lx: i32 = @as(i32, @intFromFloat(chart_left)) + 4 * label_scale;
        const ly: i32 = @as(i32, @intFromFloat(chart_top)) + 2 * label_scale;
        text.drawText(&sfc, lx, ly, label_scale, th.info, "then");
        text.drawText(&sfc, lx, ly + lgh + 2 * label_scale, label_scale, th.accent, "now");
    }

    // Axis labels (export only): y dollar ticks + x year endpoints.
    if (axis_labels) {
        axis.drawYTicks(&sfc, label_scale, th.text_muted, chart_right, chart_top, chart_bottom, value_min, value_max, 5, .dollars);
        var fbuf: [8]u8 = undefined;
        var lbuf: [8]u8 = undefined;
        const first_s = std.fmt.bufPrint(&fbuf, "{d}", .{0}) catch "0";
        const last_s = std.fmt.bufPrint(&lbuf, "{d}", .{n - 1}) catch "";
        const yr_y = chart_bottom + axis.labelGap(label_scale);
        axis.drawXEndpoints(&sfc, label_scale, th.text_muted, chart_left, chart_right, yr_y, first_s, last_s);
    }

    return .{
        .surface = sfc,
        .width = @intCast(width_px),
        .height = @intCast(height_px),
        .value_min = value_min,
        .value_max = value_max,
    };
}

/// Thin RGB wrapper over `renderToSurface` for the inline kitty path:
/// renders without axis labels, extracts RGB, frees the surface.
pub fn renderCompareChart(
    io: std.Io,
    alloc: std.mem.Allocator,
    then_bands: []const projections.YearPercentiles,
    now_bands: []const projections.YearPercentiles,
    width_px: u32,
    height_px: u32,
    th: theme.Theme,
) !CompareChartResult {
    var rendered = try renderToSurface(io, alloc, then_bands, now_bands, width_px, height_px, th, false);
    defer rendered.deinit(alloc);
    return .{
        .rgb_data = try rendered.extractRgb(alloc),
        .width = rendered.width,
        .height = rendered.height,
        .value_min = rendered.value_min,
        .value_max = rendered.value_max,
    };
}

/// Draw one envelope's light p10-p90 fill in `hue`. Indices map to x
/// via `chart_left + i * x_step`.
fn drawEnvelopeFill(
    ctx: *Context,
    bands: []const projections.YearPercentiles,
    chart_left: f64,
    x_step: f64,
    value_min: f64,
    value_max: f64,
    chart_top: f64,
    chart_bottom: f64,
    hue: [3]u8,
    bg: [3]u8,
) !void {
    ctx.setSourceToPixel(draw.blendColor(hue, 22, bg));
    ctx.resetPath();
    for (bands, 0..) |bp, i| {
        const x = chart_left + @as(f64, @floatFromInt(i)) * x_step;
        const y = draw.mapY(bp.p90, value_min, value_max, chart_top, chart_bottom);
        if (i == 0) try ctx.moveTo(x, y) else try ctx.lineTo(x, y);
    }
    var j: usize = bands.len;
    while (j > 0) {
        j -= 1;
        const x = chart_left + @as(f64, @floatFromInt(j)) * x_step;
        const y = draw.mapY(bands[j].p10, value_min, value_max, chart_top, chart_bottom);
        try ctx.lineTo(x, y);
    }
    try ctx.closePath();
    try ctx.fill();
}

/// Draw one envelope's solid p50 median line in `hue`.
fn drawEnvelopeMedian(
    ctx: *Context,
    bands: []const projections.YearPercentiles,
    chart_left: f64,
    x_step: f64,
    value_min: f64,
    value_max: f64,
    chart_top: f64,
    chart_bottom: f64,
    hue: [3]u8,
) !void {
    ctx.setSourceToPixel(draw.opaqueColor(hue));
    ctx.setLineWidth(2.0);
    ctx.resetPath();
    for (bands, 0..) |bp, i| {
        const x = chart_left + @as(f64, @floatFromInt(i)) * x_step;
        const y = draw.mapY(bp.p50, value_min, value_max, chart_top, chart_bottom);
        if (i == 0) try ctx.moveTo(x, y) else try ctx.lineTo(x, y);
    }
    try ctx.stroke();
}

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

fn syntheticBands(buf: []projections.YearPercentiles, base: f64, growth: f64) []projections.YearPercentiles {
    for (buf, 0..) |*b, i| {
        const v = base * (1.0 + growth * @as(f64, @floatFromInt(i)));
        b.* = .{
            .year = @intCast(i),
            .p10 = v * 0.6,
            .p25 = v * 0.8,
            .p50 = v,
            .p75 = v * 1.2,
            .p90 = v * 1.5,
        };
    }
    return buf;
}

test "renderCompareChart produces valid RGB output for two envelopes" {
    const alloc = std.testing.allocator;
    var then_buf: [11]projections.YearPercentiles = undefined;
    var now_buf: [11]projections.YearPercentiles = undefined;
    const then_bands = syntheticBands(&then_buf, 1_000_000, 0.05);
    const now_bands = syntheticBands(&now_buf, 1_200_000, 0.06);

    const th = theme.default_theme;
    const result = try renderCompareChart(std.testing.io, alloc, then_bands, now_bands, 240, 120, th);
    defer alloc.free(result.rgb_data);

    try std.testing.expectEqual(@as(u16, 240), result.width);
    try std.testing.expectEqual(@as(u16, 120), result.height);
    try std.testing.expectEqual(@as(usize, 240 * 120 * 3), result.rgb_data.len);
    try std.testing.expect(result.value_max > result.value_min);
}

test "renderToSurface draws both envelope hues" {
    const alloc = std.testing.allocator;
    var then_buf: [11]projections.YearPercentiles = undefined;
    var now_buf: [11]projections.YearPercentiles = undefined;
    const then_bands = syntheticBands(&then_buf, 1_000_000, 0.03);
    const now_bands = syntheticBands(&now_buf, 1_500_000, 0.07);

    const th = theme.default_theme;
    var rendered = try renderToSurface(std.testing.io, alloc, then_bands, now_bands, 200, 100, th, false);
    defer rendered.deinit(alloc);

    // Both median hues should have painted opaque pixels.
    try std.testing.expect(draw.countColor(&rendered.surface, th.info) > 0);
    try std.testing.expect(draw.countColor(&rendered.surface, th.accent) > 0);
}

test "renderToSurface insufficient data on a single-year band" {
    const alloc = std.testing.allocator;
    var then_buf: [1]projections.YearPercentiles = undefined;
    var now_buf: [11]projections.YearPercentiles = undefined;
    const then_bands = syntheticBands(&then_buf, 1_000_000, 0.05);
    const now_bands = syntheticBands(&now_buf, 1_000_000, 0.05);

    const th = theme.default_theme;
    try std.testing.expectError(
        error.InsufficientData,
        renderToSurface(std.testing.io, alloc, then_bands, now_bands, 200, 100, th, false),
    );
}