struct vs_input {
    float4 pos              : POSITION;
    float line_thickness    : THICKNESS;
    float line_spacing      : SPACING;
    float2 offset           : OFFSET;
    float4 bg0_srgb         : COLOR0;
    float4 bg1_srgb         : COLOR1;
    float4 line_srgb        : COLOR2;
    float4 x_srgb           : COLOR3;
    float4 y_srgb           : COLOR4;
};

struct ps_input {
    float4 screen_pos       : SV_POSITION;
    float line_thickness    : THICKNESS;
    float line_spacing      : SPACING;
    float2 offset           : OFFSET;
    float4 bg0_lin         : COLOR0;
    float4 bg1_lin         : COLOR1;
    float4 line_lin        : COLOR2;
    float4 x_lin           : COLOR3;
    float4 y_lin           : COLOR4;
};

cbuffer vs_constants : register(b0)
{
    float4x4 projection;
};

float4 linear_from_srgb(float4 srgb)
{
    return float4(pow(srgb.rgb, 2.2), srgb.a);
}

ps_input vs_main(vs_input input)
{
    ps_input output;

    output.screen_pos = mul(projection, float4(input.pos.xy, 0.f, 1.f));
    output.line_thickness = input.line_thickness;
    output.line_spacing = input.line_spacing;
    output.offset = input.offset;
    output.bg0_lin = linear_from_srgb(input.bg0_srgb);
    output.bg1_lin = linear_from_srgb(input.bg1_srgb);
    output.line_lin = linear_from_srgb(input.line_srgb);
    output.x_lin = linear_from_srgb(input.x_srgb);
    output.y_lin = linear_from_srgb(input.y_srgb);

    return output;
}

float4 ps_main(ps_input input) : SV_TARGET
{
    float2 grid_pos = input.screen_pos.xy + input.offset;
    float half_thickness = input.line_thickness / 2;
    float spacing = input.line_spacing;

    float4 color = input.bg0_lin;

    float2 v = abs(round(grid_pos / spacing) * spacing - grid_pos);
    float dist = min(v.x, v.y);

    if (grid_pos.y <= half_thickness && grid_pos.y >= -half_thickness) {
        color = input.x_lin;
    } else if (grid_pos.x <= half_thickness && grid_pos.x >= -half_thickness) {
        color = input.y_lin;
    } else if (dist < half_thickness) {
        color = input.line_lin;
    } else {
        bool checker = false;
        uint cell_x = (uint)(abs(grid_pos.x) / spacing) + (grid_pos.x < 0);
        uint cell_y = (uint)(abs(grid_pos.y) / spacing) + (grid_pos.y < 0);
        if (cell_x % 2 == 0) {
            checker = cell_y % 2 == 0;
        } else {
            checker = cell_y % 2 == 1;
        }
        if (checker) {
            color = input.bg1_lin;
        }
    }

    return color;
}