-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 |
    {wrong_parameters, binary()} |
    jacobian_task_error.

-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_col(
    gleam@dynamic:dynamic_(),
    fun((gleam@dynamic:dynamic_(), gleam@dynamic:dynamic_()) -> float()),
    gleam@dynamic:dynamic_(),
    float(),
    integer(),
    integer()
) -> gleam@dynamic:dynamic_().
compute_jacobian_col(X, Func, Params, Epsilon, N, I) ->
    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
    ),
    'Elixir.Nx':new_axis(
        'Elixir.Nx':divide('Elixir.Nx':subtract(Up_f, Down_f), 2.0 * Epsilon),
        1
    ).

-spec jacobian(
    gleam@dynamic:dynamic_(),
    fun((gleam@dynamic:dynamic_(), gleam@dynamic:dynamic_()) -> float()),
    gleam@dynamic:dynamic_(),
    float()
) -> {ok, gleam@dynamic:dynamic_()} | {error, fit_errors()}.
jacobian(X, Func, Params, Epsilon) ->
    {N} = 'Elixir.Nx':shape(Params),
    Jac_result = begin
        _pipe = gleam@list:range(0, N - 1),
        _pipe@1 = gleam@list:map(
            _pipe,
            fun(I) ->
                gleam@otp@task:async(
                    fun() ->
                        compute_jacobian_col(X, Func, Params, Epsilon, N, I)
                    end
                )
            end
        ),
        _pipe@2 = gleam@list:map(
            _pipe@1,
            fun(_capture) -> gleam@otp@task:try_await_forever(_capture) end
        ),
        gleam@result:all(_pipe@2)
    end,
    case Jac_result of
        {ok, Jac_cols} ->
            {ok, 'Elixir.NxBindings':concatenate(Jac_cols, 1)};

        {error, _} ->
            {error, jacobian_task_error}
    end.

-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,
            gleam@result:'try'(
                jacobian(X, Func, Params, Epsilon),
                fun(J) ->
                    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
            )
    end.

-spec compare_list_sizes(
    list(IWO),
    list(IWO),
    binary(),
    fun(() -> {ok, IWR} | {error, fit_errors()})
) -> {ok, IWR} | {error, fit_errors()}.
compare_list_sizes(A, B, Msg, Callback) ->
    case erlang:length(A) =:= erlang:length(B) of
        true ->
            Callback();

        false ->
            {error, {wrong_parameters, Msg}}
    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
) ->
    compare_list_sizes(
        X,
        Y,
        <<"x and y must have the same length"/utf8>>,
        fun() ->
            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,
            Tol = case Tolerance of
                {some, X@4} ->
                    X@4;

                none ->
                    0.0001
            end,
            Reg = case Lambda_reg 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
            )
        end
    ).