const std = @import("std");
const Thread = std.Thread;
const c = @cImport({
    @cInclude("ggml.h");
});

fn is_close(a: f32, b: f32, epsilon: f32) bool {
    return @abs(a - b) < epsilon;
}

pub fn main() !void {
    const params = .{
        .mem_size = 128 * 1024 * 1024,
        .mem_buffer = null,
        .no_alloc = false,
    };

    var opt_params = c.ggml_opt_default_params(c.GGML_OPT_TYPE_LBFGS);

    const nthreads = try Thread.getCpuCount();
    opt_params.n_threads = @intCast(nthreads);
    std.debug.print("test2: n_threads:{}\n", .{opt_params.n_threads});

    const xi = [_]f32{ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0 };
    const yi = [_]f32{ 15.0, 25.0, 35.0, 45.0, 55.0, 65.0, 75.0, 85.0, 95.0, 105.0 };

    const n = xi.len;

    const ctx0 = c.ggml_init(params);
    defer c.ggml_free(ctx0);

    const x = c.ggml_new_tensor_1d(ctx0, c.GGML_TYPE_F32, n);
    const y = c.ggml_new_tensor_1d(ctx0, c.GGML_TYPE_F32, n);

    for (0..n) |i| {
        const x_data_pointer: [*]f32 = @ptrCast(@alignCast(x.*.data));
        x_data_pointer[i] = xi[i];
        const y_data_pointer: [*]f32 = @ptrCast(@alignCast(y.*.data));
        y_data_pointer[i] = yi[i];
    }

    {
        const t0 = c.ggml_new_f32(ctx0, 0.0);
        const t1 = c.ggml_new_f32(ctx0, 0.0);

        // initialize auto-diff parameters:
        _ = c.ggml_set_param(ctx0, t0);
        _ = c.ggml_set_param(ctx0, t1);

        // f = sum_i[(t0 + t1*x_i - y_i)^2]/(2n)
        const f =
            c.ggml_div(ctx0, c.ggml_sum(ctx0, c.ggml_sqr(ctx0, c.ggml_sub(ctx0, c.ggml_add(ctx0, c.ggml_mul(ctx0, x, c.ggml_repeat(ctx0, t1, x)), c.ggml_repeat(ctx0, t0, x)), y))), c.ggml_new_f32(ctx0, @as(f32, 2.0) * n));

        const res = c.ggml_opt(null, opt_params, f);

        std.debug.print("t0 = {d:.6}\n", .{c.ggml_get_f32_1d(t0, 0)});
        std.debug.print("t1 = {d:.6}\n", .{c.ggml_get_f32_1d(t1, 0)});

        try std.testing.expect(res == c.GGML_OPT_RESULT_OK);
        try std.testing.expect(is_close(c.ggml_get_f32_1d(t0, 0), 5.0, 1e-3));
        try std.testing.expect(is_close(c.ggml_get_f32_1d(t1, 0), 10.0, 1e-3));
    }

    {
        const t0 = c.ggml_new_f32(ctx0, -1.0);
        const t1 = c.ggml_new_f32(ctx0, 9.0);

        _ = c.ggml_set_param(ctx0, t0);
        _ = c.ggml_set_param(ctx0, t1);

        // f = 0.5*sum_i[abs(t0 + t1*x_i - y_i)]/n
        const f =
            c.ggml_mul(ctx0, c.ggml_new_f32(ctx0, @as(f32, 1.0) / (2 * n)), c.ggml_sum(ctx0, c.ggml_abs(ctx0, c.ggml_sub(ctx0, c.ggml_add(ctx0, c.ggml_mul(ctx0, x, c.ggml_repeat(ctx0, t1, x)), c.ggml_repeat(ctx0, t0, x)), y))));

        const res = c.ggml_opt(null, opt_params, f);

        try std.testing.expect(res == c.GGML_OPT_RESULT_OK);
        try std.testing.expect(is_close(c.ggml_get_f32_1d(t0, 0), 5.0, 1e-2));
        try std.testing.expect(is_close(c.ggml_get_f32_1d(t1, 0), 10.0, 1e-2));
    }

    {
        const t0 = c.ggml_new_f32(ctx0, 5.0);
        const t1 = c.ggml_new_f32(ctx0, -4.0);

        _ = c.ggml_set_param(ctx0, t0);
        _ = c.ggml_set_param(ctx0, t1);

        // f = t0^2 + t1^2
        const f =
            c.ggml_add(ctx0, c.ggml_sqr(ctx0, t0), c.ggml_sqr(ctx0, t1));

        const res = c.ggml_opt(null, opt_params, f);

        try std.testing.expect(res == c.GGML_OPT_RESULT_OK);
        try std.testing.expect(is_close(c.ggml_get_f32_1d(f, 0), 0.0, 1e-3));
        try std.testing.expect(is_close(c.ggml_get_f32_1d(t0, 0), 0.0, 1e-3));
        try std.testing.expect(is_close(c.ggml_get_f32_1d(t1, 0), 0.0, 1e-3));
    }

    /////////////////////////////////////////

    {
        const t0 = c.ggml_new_f32(ctx0, -7.0);
        const t1 = c.ggml_new_f32(ctx0, 8.0);

        _ = c.ggml_set_param(ctx0, t0);
        _ = c.ggml_set_param(ctx0, t1);

        // f = (t0 + 2*t1 - 7)^2 + (2*t0 + t1 - 5)^2
        const f =
            c.ggml_add(ctx0, c.ggml_sqr(ctx0, c.ggml_sub(ctx0, c.ggml_add(ctx0, t0, c.ggml_mul(ctx0, t1, c.ggml_new_f32(ctx0, 2.0))), c.ggml_new_f32(ctx0, 7.0))), c.ggml_sqr(ctx0, c.ggml_sub(ctx0, c.ggml_add(ctx0, c.ggml_mul(ctx0, t0, c.ggml_new_f32(ctx0, 2.0)), t1), c.ggml_new_f32(ctx0, 5.0))));

        const res = c.ggml_opt(null, opt_params, f);

        try std.testing.expect(res == c.GGML_OPT_RESULT_OK);
        try std.testing.expect(is_close(c.ggml_get_f32_1d(f, 0), 0.0, 1e-3));
        try std.testing.expect(is_close(c.ggml_get_f32_1d(t0, 0), 1.0, 1e-3));
        try std.testing.expect(is_close(c.ggml_get_f32_1d(t1, 0), 3.0, 1e-3));
    }

    _ = try std.io.getStdIn().reader().readByte();
}
