-module(viva_tensor@vision@augmentations).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/vision/augmentations.gleam").
-export([color_jitter_init/4, color_jitter_forward/2, mixup/4, cutmix/4]).
-export_type([color_jitter_config/0]).

-if(?OTP_RELEASE >= 27).
-define(MODULEDOC(Str), -moduledoc(Str)).
-define(DOC(Str), -doc(Str)).
-else.
-define(MODULEDOC(Str), -compile([])).
-define(DOC(Str), -compile([])).
-endif.

?MODULEDOC(false).

-type color_jitter_config() :: {color_jitter_config,
        float(),
        float(),
        float(),
        float()}.

-file("src/viva_tensor/vision/augmentations.gleam", 43).
?DOC(false).
-spec color_jitter_init(float(), float(), float(), float()) -> color_jitter_config().
color_jitter_init(Brightness, Contrast, Saturation, Hue) ->
    {color_jitter_config, Brightness, Contrast, Saturation, Hue}.

-file("src/viva_tensor/vision/augmentations.gleam", 525).
?DOC(false).
-spec uniform_in(float(), float()) -> float().
uniform_in(Lo, Hi) ->
    Lo + ((Hi - Lo) * viva_tensor@core@ffi:random_uniform()).

-file("src/viva_tensor/vision/augmentations.gleam", 204).
?DOC(false).
-spec apply_hue(list(float()), list(float()), list(float()), float()) -> {list(float()),
    list(float()),
    list(float())}.
apply_hue(R, G, B, Strength) ->
    case Strength =< +0.0 of
        true ->
            {R, G, B};

        false ->
            Theta = uniform_in(+0.0 - Strength, Strength),
            C = viva_tensor@core@ffi:cos(Theta),
            S = viva_tensor@core@ffi:sin(Theta),
            One_minus_c_third = (1.0 - C) / 3.0,
            Sqrt3 = viva_tensor@core@ffi:sqrt(3.0),
            S_over_sqrt3 = case Sqrt3 of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> S / Gleam@denominator
            end,
            Diag = C + One_minus_c_third,
            Off1 = One_minus_c_third - S_over_sqrt3,
            Off2 = One_minus_c_third + S_over_sqrt3,
            Triples = gleam@list:zip(R, gleam@list:zip(G, B)),
            Mixed = gleam@list:map(
                Triples,
                fun(T) ->
                    {Rv, {Gv, Bv}} = T,
                    Rn = ((Diag * Rv) + (Off1 * Gv)) + (Off2 * Bv),
                    Gn = ((Off2 * Rv) + (Diag * Gv)) + (Off1 * Bv),
                    Bn = ((Off1 * Rv) + (Off2 * Gv)) + (Diag * Bv),
                    {Rn, Gn, Bn}
                end
            ),
            R_out = gleam@list:map(
                Mixed,
                fun(T@1) -> erlang:element(1, T@1) end
            ),
            G_out = gleam@list:map(
                Mixed,
                fun(T@2) -> erlang:element(2, T@2) end
            ),
            B_out = gleam@list:map(
                Mixed,
                fun(T@3) -> erlang:element(3, T@3) end
            ),
            {R_out, G_out, B_out}
    end.

-file("src/viva_tensor/vision/augmentations.gleam", 177).
?DOC(false).
-spec apply_saturation(list(float()), list(float()), list(float()), float()) -> {list(float()),
    list(float()),
    list(float())}.
apply_saturation(R, G, B, Strength) ->
    case Strength =< +0.0 of
        true ->
            {R, G, B};

        false ->
            S = uniform_in(1.0 - Strength, 1.0 + Strength),
            Luma_for = gleam@list:map(
                gleam@list:zip(R, gleam@list:zip(G, B)),
                fun(Triple) ->
                    {Rv, {Gv, Bv}} = Triple,
                    ((0.299 * Rv) + (0.587 * Gv)) + (0.114 * Bv)
                end
            ),
            Lerp_one = fun(Channel) ->
                gleam@list:map(
                    gleam@list:zip(Channel, Luma_for),
                    fun(Pair) ->
                        {V, Y} = Pair,
                        Y + (S * (V - Y))
                    end
                )
            end,
            {Lerp_one(R), Lerp_one(G), Lerp_one(B)}
    end.

-file("src/viva_tensor/vision/augmentations.gleam", 152).
?DOC(false).
-spec apply_contrast(list(float()), list(float()), list(float()), float()) -> {list(float()),
    list(float()),
    list(float())}.
apply_contrast(R, G, B, Strength) ->
    case Strength =< +0.0 of
        true ->
            {R, G, B};

        false ->
            Factor = uniform_in(1.0 - Strength, 1.0 + Strength),
            N = erlang:float(erlang:length(R) * 3),
            Total = (gleam@list:fold(R, +0.0, fun(Acc, V) -> Acc + V end) + gleam@list:fold(
                G,
                +0.0,
                fun(Acc@1, V@1) -> Acc@1 + V@1 end
            ))
            + gleam@list:fold(B, +0.0, fun(Acc@2, V@2) -> Acc@2 + V@2 end),
            Mean = case N > +0.0 of
                true ->
                    case N of
                        +0.0 -> +0.0;
                        -0.0 -> -0.0;
                        Gleam@denominator -> Total / Gleam@denominator
                    end;

                false ->
                    +0.0
            end,
            Shift = fun(V@3) -> ((V@3 - Mean) * Factor) + Mean end,
            {gleam@list:map(R, Shift),
                gleam@list:map(G, Shift),
                gleam@list:map(B, Shift)}
    end.

-file("src/viva_tensor/vision/augmentations.gleam", 133).
?DOC(false).
-spec apply_brightness(list(float()), list(float()), list(float()), float()) -> {list(float()),
    list(float()),
    list(float())}.
apply_brightness(R, G, B, Strength) ->
    case Strength =< +0.0 of
        true ->
            {R, G, B};

        false ->
            Factor = uniform_in(1.0 - Strength, 1.0 + Strength),
            {gleam@list:map(R, fun(V) -> V * Factor end),
                gleam@list:map(G, fun(V@1) -> V@1 * Factor end),
                gleam@list:map(B, fun(V@2) -> V@2 * Factor end)}
    end.

-file("src/viva_tensor/vision/augmentations.gleam", 458).
?DOC(false).
-spec slice(list(float()), integer(), integer()) -> list(float()).
slice(Data, Start, Length) ->
    _pipe = Data,
    _pipe@1 = gleam@list:drop(_pipe, Start),
    gleam@list:take(_pipe@1, Length).

-file("src/viva_tensor/vision/augmentations.gleam", 393).
?DOC(false).
-spec parse_image_shape(binary(), list(integer())) -> {ok,
        {integer(), integer(), integer(), integer()}} |
    {error, viva_tensor@core@error:tensor_error()}.
parse_image_shape(Op, Shape) ->
    case Shape of
        [B, C, H, W] ->
            {ok, {B, C, H, W}};

        [C@1, H@1, W@1] ->
            {ok, {1, C@1, H@1, W@1}};

        _ ->
            {error,
                {operand_shape_mismatch,
                    Op,
                    <<"image"/utf8>>,
                    <<"[C, H, W] or [B, C, H, W]"/utf8>>,
                    Shape}}
    end.

-file("src/viva_tensor/vision/augmentations.gleam", 79).
?DOC(false).
-spec color_jitter_forward(color_jitter_config(), viva_tensor@tensor:tensor()) -> {ok,
        viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
color_jitter_forward(Config, Image) ->
    Shape = viva_tensor@tensor:shape(Image),
    gleam@result:'try'(
        parse_image_shape(<<"color_jitter_forward"/utf8>>, Shape),
        fun(_use0) ->
            {Batch, Channels, Height, Width} = _use0,
            case Channels =:= 3 of
                true ->
                    Data = viva_tensor@tensor:to_list(Image),
                    Plane = Height * Width,
                    Stride = Channels * Plane,
                    Total_strength = ((gleam@float:absolute_value(
                        erlang:element(2, Config)
                    )
                    + gleam@float:absolute_value(erlang:element(3, Config)))
                    + gleam@float:absolute_value(erlang:element(4, Config)))
                    + gleam@float:absolute_value(erlang:element(5, Config)),
                    case Total_strength =< +0.0 of
                        true ->
                            {ok, Image};

                        false ->
                            Processed = begin
                                _pipe = gleam@list:range(0, Batch - 1),
                                gleam@list:flat_map(
                                    _pipe,
                                    fun(B) ->
                                        Start = B * Stride,
                                        R = slice(Data, Start, Plane),
                                        G = slice(Data, Start + Plane, Plane),
                                        B_ch = slice(
                                            Data,
                                            Start + (2 * Plane),
                                            Plane
                                        ),
                                        {R1, G1, B1} = apply_brightness(
                                            R,
                                            G,
                                            B_ch,
                                            erlang:element(2, Config)
                                        ),
                                        {R2, G2, B2} = apply_contrast(
                                            R1,
                                            G1,
                                            B1,
                                            erlang:element(3, Config)
                                        ),
                                        {R3, G3, B3} = apply_saturation(
                                            R2,
                                            G2,
                                            B2,
                                            erlang:element(4, Config)
                                        ),
                                        {R4, G4, B4} = apply_hue(
                                            R3,
                                            G3,
                                            B3,
                                            erlang:element(5, Config)
                                        ),
                                        lists:append(R4, lists:append(G4, B4))
                                    end
                                )
                            end,
                            {ok, {tensor, Processed, Shape}}
                    end;

                false ->
                    {error,
                        {operand_shape_mismatch,
                            <<"color_jitter_forward"/utf8>>,
                            <<"image"/utf8>>,
                            <<"channel dim = 3 (RGB)"/utf8>>,
                            Shape}}
            end
        end
    ).

-file("src/viva_tensor/vision/augmentations.gleam", 462).
?DOC(false).
-spec mix_batched(list(float()), list(integer()), integer(), float()) -> list(float()).
mix_batched(Data, Perm, Per_sample, Lambda) ->
    _pipe = gleam@list:index_map(
        Perm,
        fun(P, I) ->
            A = slice(Data, I * Per_sample, Per_sample),
            B = slice(Data, P * Per_sample, Per_sample),
            gleam@list:map(
                gleam@list:zip(A, B),
                fun(Pair) ->
                    {Va, Vb} = Pair,
                    (Lambda * Va) + ((1.0 - Lambda) * Vb)
                end
            )
        end
    ),
    lists:append(_pipe).

-file("src/viva_tensor/vision/augmentations.gleam", 454).
?DOC(false).
-spec element_count(list(integer())) -> integer().
element_count(Shape) ->
    gleam@list:fold(Shape, 1, fun(Acc, Dim) -> Acc * Dim end).

-file("src/viva_tensor/vision/augmentations.gleam", 563).
?DOC(false).
-spec list_get(list(integer()), integer(), integer()) -> integer().
list_get(Xs, I, Default) ->
    case gleam@list:drop(Xs, I) of
        [V | _] ->
            V;

        [] ->
            Default
    end.

-file("src/viva_tensor/vision/augmentations.gleam", 546).
?DOC(false).
-spec swap_at(list(integer()), integer(), integer()) -> list(integer()).
swap_at(Xs, I, J) ->
    case I =:= J of
        true ->
            Xs;

        false ->
            Xi = list_get(Xs, I, 0),
            Xj = list_get(Xs, J, 0),
            gleam@list:index_map(Xs, fun(V, K) -> case {K =:= I, K =:= J} of
                        {true, _} ->
                            Xj;

                        {_, true} ->
                            Xi;

                        {_, _} ->
                            V
                    end end)
    end.

-file("src/viva_tensor/vision/augmentations.gleam", 535).
?DOC(false).
-spec fisher_yates(list(integer()), integer()) -> list(integer()).
fisher_yates(Xs, I) ->
    case I =< 0 of
        true ->
            Xs;

        false ->
            J = gleam@int:random(I + 1),
            Swapped = swap_at(Xs, I, J),
            fisher_yates(Swapped, I - 1)
    end.

-file("src/viva_tensor/vision/augmentations.gleam", 529).
?DOC(false).
-spec random_permutation(integer()) -> list(integer()).
random_permutation(N) ->
    Initial = gleam@list:range(0, N - 1),
    fisher_yates(Initial, N - 1).

-file("src/viva_tensor/vision/augmentations.gleam", 638).
?DOC(false).
-spec standard_normal() -> float().
standard_normal() ->
    U1 = gleam@float:max(viva_tensor@core@ffi:random_uniform(), 1.0e-12),
    U2 = viva_tensor@core@ffi:random_uniform(),
    viva_tensor@core@ffi:sqrt(-2.0 * viva_tensor@core@ffi:log(U1)) * viva_tensor@core@ffi:cos(
        (2.0 * 3.14159265358979323846) * U2
    ).

-file("src/viva_tensor/vision/augmentations.gleam", 614).
?DOC(false).
-spec marsaglia_loop(float(), float(), integer()) -> float().
marsaglia_loop(D, C, Attempts) ->
    case Attempts >= 100 of
        true ->
            D;

        false ->
            X = standard_normal(),
            V_base = 1.0 + (C * X),
            case V_base =< +0.0 of
                true ->
                    marsaglia_loop(D, C, Attempts + 1);

                false ->
                    V = (V_base * V_base) * V_base,
                    U = gleam@float:max(
                        viva_tensor@core@ffi:random_uniform(),
                        1.0e-12
                    ),
                    Lhs = U,
                    Rhs = viva_tensor@core@ffi:exp(
                        ((((0.5 * X) * X) + D) - (D * V)) + (D * viva_tensor@core@ffi:log(
                            V
                        ))
                    ),
                    case Lhs < Rhs of
                        true ->
                            D * V;

                        false ->
                            marsaglia_loop(D, C, Attempts + 1)
                    end
            end
    end.

-file("src/viva_tensor/vision/augmentations.gleam", 608).
?DOC(false).
-spec marsaglia_tsang(float()) -> float().
marsaglia_tsang(Shape) ->
    D = Shape - (1.0 / 3.0),
    C = case viva_tensor@core@ffi:sqrt(9.0 * D) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> 1.0 / Gleam@denominator
    end,
    marsaglia_loop(D, C, 0).

-file("src/viva_tensor/vision/augmentations.gleam", 595).
?DOC(false).
-spec sample_gamma(float()) -> float().
sample_gamma(Shape) ->
    case Shape < 1.0 of
        true ->
            Boosted = marsaglia_tsang(Shape + 1.0),
            U = gleam@float:max(viva_tensor@core@ffi:random_uniform(), 1.0e-12),
            Boosted * viva_tensor@core@ffi:pow(U, case Shape of
                    +0.0 -> +0.0;
                    -0.0 -> -0.0;
                    Gleam@denominator -> 1.0 / Gleam@denominator
                end);

        false ->
            marsaglia_tsang(Shape)
    end.

-file("src/viva_tensor/vision/augmentations.gleam", 580).
?DOC(false).
-spec sample_beta(float(), float()) -> float().
sample_beta(A, B) ->
    case (A =< +0.0) orelse (B =< +0.0) of
        true ->
            1.0;

        false ->
            X = sample_gamma(A),
            Y = sample_gamma(B),
            Total = X + Y,
            case Total =< +0.0 of
                true ->
                    0.5;

                false ->
                    case Total of
                        +0.0 -> +0.0;
                        -0.0 -> -0.0;
                        Gleam@denominator -> X / Gleam@denominator
                    end
            end
    end.

-file("src/viva_tensor/vision/augmentations.gleam", 441).
?DOC(false).
-spec one_hot_from_indices(list(float()), integer()) -> list(float()).
one_hot_from_indices(Indices, Num_classes) ->
    gleam@list:flat_map(
        Indices,
        fun(Idx_f) ->
            Idx = erlang:round(Idx_f),
            _pipe = gleam@list:range(0, Num_classes - 1),
            gleam@list:map(_pipe, fun(K) -> case K =:= Idx of
                        true ->
                            1.0;

                        false ->
                            +0.0
                    end end)
        end
    ).

-file("src/viva_tensor/vision/augmentations.gleam", 410).
?DOC(false).
-spec normalize_labels(
    binary(),
    viva_tensor@tensor:tensor(),
    integer(),
    integer()
) -> {ok, list(float())} | {error, viva_tensor@core@error:tensor_error()}.
normalize_labels(Op, Labels, Batch, Num_classes) ->
    case Num_classes =< 0 of
        true ->
            {error,
                {invalid_shape,
                    <<<<Op/binary, ": num_classes must be positive, got "/utf8>>/binary,
                        (erlang:integer_to_binary(Num_classes))/binary>>}};

        false ->
            Label_shape = viva_tensor@tensor:shape(Labels),
            Label_data = viva_tensor@tensor:to_list(Labels),
            case Label_shape of
                [B] when B =:= Batch ->
                    {ok, one_hot_from_indices(Label_data, Num_classes)};

                [B@1, K] when (B@1 =:= Batch) andalso (K =:= Num_classes) ->
                    {ok, Label_data};

                _ ->
                    {error,
                        {operand_shape_mismatch,
                            Op,
                            <<"labels"/utf8>>,
                            <<"[B] indices or [B, num_classes] one-hot"/utf8>>,
                            Label_shape}}
            end
    end.

-file("src/viva_tensor/vision/augmentations.gleam", 267).
?DOC(false).
-spec mixup(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    integer(),
    float()
) -> {ok, {viva_tensor@tensor:tensor(), viva_tensor@tensor:tensor()}} |
    {error, viva_tensor@core@error:tensor_error()}.
mixup(Images, Labels, Num_classes, Alpha) ->
    Image_shape = viva_tensor@tensor:shape(Images),
    gleam@result:'try'(
        parse_image_shape(<<"mixup"/utf8>>, Image_shape),
        fun(_use0) ->
            {Batch, _, _, _} = _use0,
            gleam@result:'try'(
                normalize_labels(<<"mixup"/utf8>>, Labels, Batch, Num_classes),
                fun(Label_matrix) ->
                    Lambda = sample_beta(Alpha, Alpha),
                    Perm = random_permutation(Batch),
                    Image_data = viva_tensor@tensor:to_list(Images),
                    Per_image_size = case Batch of
                        0 -> 0;
                        Gleam@denominator -> element_count(Image_shape) div Gleam@denominator
                    end,
                    Mixed_image_data = mix_batched(
                        Image_data,
                        Perm,
                        Per_image_size,
                        Lambda
                    ),
                    Mixed_label_data = mix_batched(
                        Label_matrix,
                        Perm,
                        Num_classes,
                        Lambda
                    ),
                    {ok,
                        {{tensor, Mixed_image_data, Image_shape},
                            {tensor, Mixed_label_data, [Batch, Num_classes]}}}
                end
            )
        end
    ).

-file("src/viva_tensor/vision/augmentations.gleam", 479).
?DOC(false).
-spec paste_box(
    list(float()),
    list(integer()),
    integer(),
    integer(),
    integer(),
    integer(),
    integer(),
    integer(),
    integer(),
    integer()
) -> list(float()).
paste_box(Data, Perm, Batch, Channels, Height, Width, X1, X2, Y1, Y2) ->
    Plane = Height * Width,
    Stride = Channels * Plane,
    _pipe = gleam@list:range(0, Batch - 1),
    gleam@list:flat_map(
        _pipe,
        fun(I) ->
            P = case gleam@list:drop(Perm, I) of
                [Head | _] ->
                    Head;

                [] ->
                    I
            end,
            _pipe@1 = gleam@list:range(0, Channels - 1),
            gleam@list:flat_map(
                _pipe@1,
                fun(C) ->
                    Dst_base = (I * Stride) + (C * Plane),
                    Src_base = (P * Stride) + (C * Plane),
                    _pipe@2 = gleam@list:range(0, Height - 1),
                    gleam@list:flat_map(
                        _pipe@2,
                        fun(Y) -> _pipe@3 = gleam@list:range(0, Width - 1),
                            gleam@list:map(
                                _pipe@3,
                                fun(X) ->
                                    Pick_src = (((Y >= Y1) andalso (Y < Y2))
                                    andalso (X >= X1))
                                    andalso (X < X2),
                                    Offset = (Y * Width) + X,
                                    Base = case Pick_src of
                                        true ->
                                            Src_base;

                                        false ->
                                            Dst_base
                                    end,
                                    case gleam@list:drop(Data, Base + Offset) of
                                        [V | _] ->
                                            V;

                                        [] ->
                                            +0.0
                                    end
                                end
                            ) end
                    )
                end
            )
        end
    ).

-file("src/viva_tensor/vision/augmentations.gleam", 645).
?DOC(false).
-spec int_clamp(integer(), integer(), integer()) -> integer().
int_clamp(Value, Lo, Hi) ->
    case Value < Lo of
        true ->
            Lo;

        false ->
            case Value > Hi of
                true ->
                    Hi;

                false ->
                    Value
            end
    end.

-file("src/viva_tensor/vision/augmentations.gleam", 656).
?DOC(false).
-spec float_to_round(float()) -> integer().
float_to_round(Value) ->
    erlang:round(Value).

-file("src/viva_tensor/vision/augmentations.gleam", 322).
?DOC(false).
-spec cutmix(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    integer(),
    float()
) -> {ok, {viva_tensor@tensor:tensor(), viva_tensor@tensor:tensor()}} |
    {error, viva_tensor@core@error:tensor_error()}.
cutmix(Images, Labels, Num_classes, Alpha) ->
    Image_shape = viva_tensor@tensor:shape(Images),
    gleam@result:'try'(
        parse_image_shape(<<"cutmix"/utf8>>, Image_shape),
        fun(_use0) ->
            {Batch, Channels, Height, Width} = _use0,
            gleam@result:'try'(
                normalize_labels(<<"cutmix"/utf8>>, Labels, Batch, Num_classes),
                fun(Label_matrix) ->
                    Lambda_initial = sample_beta(Alpha, Alpha),
                    Cut_ratio = viva_tensor@core@ffi:sqrt(
                        gleam@float:max(+0.0, 1.0 - Lambda_initial)
                    ),
                    Cut_w = float_to_round(erlang:float(Width) * Cut_ratio),
                    Cut_h = float_to_round(erlang:float(Height) * Cut_ratio),
                    Rx = case Width =< 0 of
                        true ->
                            0;

                        false ->
                            gleam@int:random(Width)
                    end,
                    Ry = case Height =< 0 of
                        true ->
                            0;

                        false ->
                            gleam@int:random(Height)
                    end,
                    X1 = int_clamp(Rx - (Cut_w div 2), 0, Width),
                    X2 = int_clamp(Rx + (Cut_w div 2), 0, Width),
                    Y1 = int_clamp(Ry - (Cut_h div 2), 0, Height),
                    Y2 = int_clamp(Ry + (Cut_h div 2), 0, Height),
                    Box_w = X2 - X1,
                    Box_h = Y2 - Y1,
                    Lambda = case (Width * Height) =< 0 of
                        true ->
                            1.0;

                        false ->
                            1.0 - (case erlang:float(Width * Height) of
                                +0.0 -> +0.0;
                                -0.0 -> -0.0;
                                Gleam@denominator -> erlang:float(Box_w * Box_h)
                                / Gleam@denominator
                            end)
                    end,
                    Perm = random_permutation(Batch),
                    Image_data = viva_tensor@tensor:to_list(Images),
                    Mixed_image_data = case (Box_w =:= 0) orelse (Box_h =:= 0) of
                        true ->
                            Image_data;

                        false ->
                            paste_box(
                                Image_data,
                                Perm,
                                Batch,
                                Channels,
                                Height,
                                Width,
                                X1,
                                X2,
                                Y1,
                                Y2
                            )
                    end,
                    Mixed_label_data = mix_batched(
                        Label_matrix,
                        Perm,
                        Num_classes,
                        Lambda
                    ),
                    {ok,
                        {{tensor, Mixed_image_data, Image_shape},
                            {tensor, Mixed_label_data, [Batch, Num_classes]}}}
                end
            )
        end
    ).