-module(gleastsq).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch]).

-export([least_squares/8]).
-export_type([fit_errors/0]).

-opaque fit_errors() :: non_converged.

-spec convert_func_params(fun((float(), list(float())) -> float())) -> fun((gleam@dynamic:dynamic_(), gleam@dynamic:dynamic_()) -> float()).
convert_func_params(Func) ->
    fun(X, Params) ->
        Func('Elixir.Nx':to_number(X), 'Elixir.Nx':to_list(Params))
    end.

-spec compute_jacobian(
    gleam@dynamic:dynamic_(),
    fun((gleam@dynamic:dynamic_(), gleam@dynamic:dynamic_()) -> float()),
    gleam@dynamic:dynamic_(),
    gleam@dynamic:dynamic_(),
    float(),
    integer(),
    integer()
) -> gleam@dynamic:dynamic_().
compute_jacobian(X, Func, Params, Jac, Epsilon, N, I) ->
    case I of
        I@1 when I@1 >= N ->
            Jac;

        _ ->
            Mask = 'Elixir.Nx':indexed_put(
                'Elixir.Nx':broadcast(+0.0, {N}),
                'Elixir.Nx':tensor([I]),
                Epsilon
            ),
            Up_params = 'Elixir.Nx':add(Params, Mask),
            Down_params = 'Elixir.Nx':subtract(Params, Mask),
            Up_f = 'Elixir.Nx':map(
                X,
                fun(_capture) -> Func(_capture, Up_params) end
            ),
            Down_f = 'Elixir.Nx':map(
                X,
                fun(_capture@1) -> Func(_capture@1, Down_params) end
            ),
            Deriv = 'Elixir.Nx':new_axis(
                'Elixir.Nx':divide(
                    'Elixir.Nx':subtract(Up_f, Down_f),
                    2.0 * Epsilon
                ),
                1
            ),
            Updated_jac = 'Elixir.Nx':put_slice(Jac, [0, I], Deriv),
            compute_jacobian(X, Func, Params, Updated_jac, Epsilon, N, I + 1)
    end.

-spec jacobian(
    gleam@dynamic:dynamic_(),
    fun((gleam@dynamic:dynamic_(), gleam@dynamic:dynamic_()) -> float()),
    gleam@dynamic:dynamic_(),
    float()
) -> gleam@dynamic:dynamic_().
jacobian(X, Func, Params, Epsilon) ->
    {N} = 'Elixir.Nx':shape(Params),
    {M} = 'Elixir.Nx':shape(X),
    Jac = 'Elixir.Nx':broadcast(+0.0, {M, N}),
    compute_jacobian(X, Func, Params, Jac, Epsilon, N, 0).

-spec do_least_squares(
    gleam@dynamic:dynamic_(),
    gleam@dynamic:dynamic_(),
    fun((gleam@dynamic:dynamic_(), gleam@dynamic:dynamic_()) -> float()),
    gleam@dynamic:dynamic_(),
    integer(),
    float(),
    float(),
    float()
) -> {ok, gleam@dynamic:dynamic_()} | {error, fit_errors()}.
do_least_squares(
    X,
    Y,
    Func,
    Params,
    Max_iterations,
    Epsilon,
    Tolerance,
    Lambda_reg
) ->
    M = erlang:element(1, 'Elixir.Nx':shape(Params)),
    case Max_iterations of
        0 ->
            {error, non_converged};

        Iterations ->
            R = begin
                _pipe = X,
                _pipe@1 = 'Elixir.Nx':map(
                    _pipe,
                    fun(_capture) -> Func(_capture, Params) end
                ),
                'Elixir.Nx':subtract(Y, _pipe@1)
            end,
            J = jacobian(X, Func, Params, Epsilon),
            Jt = 'Elixir.Nx':transpose(J),
            Lambda_eye = begin
                _pipe@2 = 'Elixir.Nx':eye(M),
                'Elixir.Nx':multiply(_pipe@2, Lambda_reg)
            end,
            H = 'Elixir.Nx':add('Elixir.Nx':dot(Jt, J), Lambda_eye),
            G = 'Elixir.Nx':dot(Jt, R),
            Delta = 'Elixir.Nx.LinAlg':solve(H, G),
            case 'Elixir.Nx':to_number('Elixir.Nx.LinAlg':norm(Delta)) of
                X@1 when X@1 < Tolerance ->
                    {ok, Params};

                _ ->
                    do_least_squares(
                        X,
                        Y,
                        Func,
                        'Elixir.Nx':add(Params, Delta),
                        Iterations - 1,
                        Epsilon,
                        Tolerance,
                        Lambda_reg
                    )
            end
    end.

-spec least_squares(
    list(float()),
    list(float()),
    fun((float(), list(float())) -> float()),
    list(float()),
    gleam@option:option(integer()),
    gleam@option:option(float()),
    gleam@option:option(float()),
    gleam@option:option(float())
) -> {ok, list(float())} | {error, fit_errors()}.
least_squares(
    X,
    Y,
    Func,
    Initial_params,
    Iterations,
    Epsilon,
    Tolerance,
    Lambda_reg
) ->
    P = 'Elixir.Nx':tensor(Initial_params),
    X@1 = 'Elixir.Nx':tensor(X),
    Y@1 = 'Elixir.Nx':tensor(Y),
    Func@1 = convert_func_params(Func),
    Iter = case Iterations of
        {some, X@2} ->
            X@2;

        none ->
            100
    end,
    Eps = case Epsilon of
        {some, X@3} ->
            X@3;

        none ->
            0.0001
    end,
    Reg = case Lambda_reg of
        {some, X@4} ->
            X@4;

        none ->
            0.0001
    end,
    Tol = case Tolerance of
        {some, X@5} ->
            X@5;

        none ->
            0.0001
    end,
    gleam@result:'try'(
        do_least_squares(X@1, Y@1, Func@1, P, Iter, Eps, Tol, Reg),
        fun(Fitted) ->
            {ok,
                begin
                    _pipe = Fitted,
                    'Elixir.Nx':to_list(_pipe)
                end}
        end
    ).