-module(viva_tensor@core@shape).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/core/shape.gleam").
-export([reshape/2, flatten/1, squeeze/1, unsqueeze/2, expand_dims/2, take_first/2, take_last/2, concat/1, concat_axis/2, stack/2, squeeze_axis/2, slice/3]).

-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(
    " Shape manipulation - reshape, slice, concat, stack.\n"
    "\n"
    " The \"plumbing\" of tensor operations. Not glamorous, but essential.\n"
    "\n"
    " Philosophy: these ops should be zero-copy when possible (via strides),\n"
    " but we're not there yet. Current implementation copies data.\n"
    " TODO: implement as views where safe (reshape of contiguous tensor, slice, etc.)\n"
    "\n"
    " Reshape is particularly tricky because it changes how we interpret memory.\n"
    " A [2,3] tensor reshaped to [3,2] has the same data but different semantics.\n"
    " This only works if the tensor is contiguous (strides match row-major order).\n"
).

-file("src/viva_tensor/core/shape.gleam", 21).
?DOC(" Reshape to new dimensions. Total size must match (obviously).\n").
-spec reshape(viva_tensor@core@tensor:tensor(), list(integer())) -> {ok,
        viva_tensor@core@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
reshape(T, New_shape) ->
    Old_size = viva_tensor@core@tensor:size(T),
    New_size = gleam@list:fold(New_shape, 1, fun(Acc, Dim) -> Acc * Dim end),
    case Old_size =:= New_size of
        true ->
            Data = viva_tensor@core@tensor:to_list(T),
            viva_tensor@core@tensor:new(Data, New_shape);

        false ->
            {error,
                {invalid_shape,
                    <<<<<<<<"Cannot reshape: size mismatch ("/utf8,
                                    (erlang:integer_to_binary(Old_size))/binary>>/binary,
                                " vs "/utf8>>/binary,
                            (erlang:integer_to_binary(New_size))/binary>>/binary,
                        ")"/utf8>>}}
    end.

-file("src/viva_tensor/core/shape.gleam", 42).
?DOC(" Flatten to 1D tensor\n").
-spec flatten(viva_tensor@core@tensor:tensor()) -> viva_tensor@core@tensor:tensor().
flatten(T) ->
    Data = viva_tensor@core@tensor:to_list(T),
    viva_tensor@core@tensor:from_list(Data).

-file("src/viva_tensor/core/shape.gleam", 48).
?DOC(" Remove all dimensions of size 1\n").
-spec squeeze(viva_tensor@core@tensor:tensor()) -> viva_tensor@core@tensor:tensor().
squeeze(T) ->
    Data = viva_tensor@core@tensor:to_list(T),
    New_shape = gleam@list:filter(
        viva_tensor@core@tensor:shape(T),
        fun(D) -> D /= 1 end
    ),
    Final_shape = case New_shape of
        [] ->
            [1];

        _ ->
            New_shape
    end,
    case viva_tensor@core@tensor:new(Data, Final_shape) of
        {ok, Result} ->
            Result;

        {error, _} ->
            T
    end.

-file("src/viva_tensor/core/shape.gleam", 84).
?DOC(" Add dimension of size 1 at specified axis\n").
-spec unsqueeze(viva_tensor@core@tensor:tensor(), integer()) -> viva_tensor@core@tensor:tensor().
unsqueeze(T, Axis) ->
    Data = viva_tensor@core@tensor:to_list(T),
    Shp = viva_tensor@core@tensor:shape(T),
    Rnk = erlang:length(Shp),
    Insert_at = case Axis < 0 of
        true ->
            (Rnk + Axis) + 1;

        false ->
            Axis
    end,
    {Before, After} = gleam@list:split(Shp, Insert_at),
    New_shape = lists:append([Before, [1], After]),
    case viva_tensor@core@tensor:new(Data, New_shape) of
        {ok, Result} ->
            Result;

        {error, _} ->
            T
    end.

-file("src/viva_tensor/core/shape.gleam", 102).
?DOC(" Alias for unsqueeze - expand dimensions\n").
-spec expand_dims(viva_tensor@core@tensor:tensor(), integer()) -> viva_tensor@core@tensor:tensor().
expand_dims(T, Axis) ->
    unsqueeze(T, Axis).

-file("src/viva_tensor/core/shape.gleam", 109).
?DOC(" Take first n elements (along axis 0).\n").
-spec take_first(viva_tensor@core@tensor:tensor(), integer()) -> viva_tensor@core@tensor:tensor().
take_first(T, N) ->
    Data = viva_tensor@core@tensor:to_list(T),
    case viva_tensor@core@tensor:shape(T) of
        [] ->
            T;

        [First_dim | Rest_dims] ->
            Take_n = gleam@int:min(N, First_dim),
            Stride = gleam@list:fold(Rest_dims, 1, fun(Acc, D) -> Acc * D end),
            New_data = gleam@list:take(Data, Take_n * Stride),
            New_shape = [Take_n | Rest_dims],
            case viva_tensor@core@tensor:new(New_data, New_shape) of
                {ok, Result} ->
                    Result;

                {error, _} ->
                    T
            end
    end.

-file("src/viva_tensor/core/shape.gleam", 127).
?DOC(" Take last N elements along first axis\n").
-spec take_last(viva_tensor@core@tensor:tensor(), integer()) -> viva_tensor@core@tensor:tensor().
take_last(T, N) ->
    Data = viva_tensor@core@tensor:to_list(T),
    case viva_tensor@core@tensor:shape(T) of
        [] ->
            T;

        [First_dim | Rest_dims] ->
            Take_n = gleam@int:min(N, First_dim),
            Stride = gleam@list:fold(Rest_dims, 1, fun(Acc, D) -> Acc * D end),
            Skip = (First_dim - Take_n) * Stride,
            New_data = gleam@list:drop(Data, Skip),
            New_shape = [Take_n | Rest_dims],
            case viva_tensor@core@tensor:new(New_data, New_shape) of
                {ok, Result} ->
                    Result;

                {error, _} ->
                    T
            end
    end.

-file("src/viva_tensor/core/shape.gleam", 201).
?DOC(" Concat 1D tensors. Just appends the data, nothing fancy.\n").
-spec concat(list(viva_tensor@core@tensor:tensor())) -> viva_tensor@core@tensor:tensor().
concat(Tensors) ->
    Data = gleam@list:flat_map(
        Tensors,
        fun(T) -> viva_tensor@core@tensor:to_list(T) end
    ),
    viva_tensor@core@tensor:from_list(Data).

-file("src/viva_tensor/core/shape.gleam", 397).
-spec flat_to_multi(integer(), list(integer())) -> list(integer()).
flat_to_multi(Flat, Shape) ->
    Reversed = lists:reverse(Shape),
    {Indices, _} = gleam@list:fold(
        Reversed,
        {[], Flat},
        fun(Acc, Dim) ->
            {Idxs, Remaining} = Acc,
            Idx = case Dim of
                0 -> 0;
                Gleam@denominator -> Remaining rem Gleam@denominator
            end,
            Next = case Dim of
                0 -> 0;
                Gleam@denominator@1 -> Remaining div Gleam@denominator@1
            end,
            {[Idx | Idxs], Next}
        end
    ),
    Indices.

-file("src/viva_tensor/core/shape.gleam", 418).
-spec compute_strides(list(integer())) -> list(integer()).
compute_strides(Shape) ->
    Reversed = lists:reverse(Shape),
    {Strides, _} = gleam@list:fold(
        Reversed,
        {[], 1},
        fun(Acc, Dim) ->
            {S, Running} = Acc,
            {[Running | S], Running * Dim}
        end
    ),
    Strides.

-file("src/viva_tensor/core/shape.gleam", 212).
?DOC(
    " Concat along arbitrary axis.\n"
    "\n"
    " This function is gnarly. The general case requires computing which source\n"
    " tensor each output index maps to, then translating coordinates. O(n) where\n"
    " n is total output size, but the constant factor is high due to all the\n"
    " index arithmetic. For axis=0, we fast-path to simple concatenation.\n"
).
-spec concat_axis(list(viva_tensor@core@tensor:tensor()), integer()) -> {ok,
        viva_tensor@core@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
concat_axis(Tensors, Axis) ->
    case Tensors of
        [] ->
            {error, {invalid_shape, <<"Cannot concatenate empty list"/utf8>>}};

        [Single] ->
            {ok, Single};

        [First | Rest] ->
            Base_shape = viva_tensor@core@tensor:shape(First),
            R = erlang:length(Base_shape),
            case (Axis >= 0) andalso (Axis < R) of
                false ->
                    {error,
                        {dimension_error,
                            <<"Invalid axis for concatenation"/utf8>>}};

                true ->
                    Shapes_ok = gleam@list:all(
                        Rest,
                        fun(T) ->
                            T_shape = viva_tensor@core@tensor:shape(T),
                            case erlang:length(T_shape) =:= R of
                                false ->
                                    false;

                                true ->
                                    _pipe = gleam@list:zip(Base_shape, T_shape),
                                    _pipe@1 = gleam@list:index_map(
                                        _pipe,
                                        fun(Pair, I) -> {Pair, I} end
                                    ),
                                    gleam@list:all(
                                        _pipe@1,
                                        fun(X) ->
                                            {{Dim_a, Dim_b}, I@1} = X,
                                            (I@1 =:= Axis) orelse (Dim_a =:= Dim_b)
                                        end
                                    )
                            end
                        end
                    ),
                    case Shapes_ok of
                        false ->
                            {error,
                                {invalid_shape,
                                    <<"Shapes must match except on concat axis"/utf8>>}};

                        true ->
                            Concat_dim = gleam@list:fold(
                                Tensors,
                                0,
                                fun(Acc, T@1) ->
                                    case begin
                                        _pipe@2 = gleam@list:drop(
                                            viva_tensor@core@tensor:shape(T@1),
                                            Axis
                                        ),
                                        gleam@list:first(_pipe@2)
                                    end of
                                        {ok, D} ->
                                            Acc + D;

                                        {error, _} ->
                                            Acc
                                    end
                                end
                            ),
                            New_shape = begin
                                _pipe@3 = Base_shape,
                                gleam@list:index_map(
                                    _pipe@3,
                                    fun(D@1, I@2) -> case I@2 =:= Axis of
                                            true ->
                                                Concat_dim;

                                            false ->
                                                D@1
                                        end end
                                )
                            end,
                            case Axis =:= 0 of
                                true ->
                                    Data = gleam@list:flat_map(
                                        Tensors,
                                        fun(T@2) ->
                                            viva_tensor@core@tensor:to_list(T@2)
                                        end
                                    ),
                                    viva_tensor@core@tensor:new(Data, New_shape);

                                false ->
                                    Total_size = gleam@list:fold(
                                        New_shape,
                                        1,
                                        fun(Acc@1, D@2) -> Acc@1 * D@2 end
                                    ),
                                    Result = begin
                                        _pipe@4 = gleam@list:range(
                                            0,
                                            Total_size - 1
                                        ),
                                        gleam@list:map(
                                            _pipe@4,
                                            fun(Flat_idx) ->
                                                Indices = flat_to_multi(
                                                    Flat_idx,
                                                    New_shape
                                                ),
                                                Axis_idx = case begin
                                                    _pipe@5 = gleam@list:drop(
                                                        Indices,
                                                        Axis
                                                    ),
                                                    gleam@list:first(_pipe@5)
                                                end of
                                                    {ok, I@3} ->
                                                        I@3;

                                                    {error, _} ->
                                                        0
                                                end,
                                                {Tensor_idx, Local_axis_idx, _} = gleam@list:fold(
                                                    Tensors,
                                                    {-1, Axis_idx, 0},
                                                    fun(Acc@2, T@3) ->
                                                        {Found_t,
                                                            Remaining,
                                                            T_idx} = Acc@2,
                                                        case Found_t >= 0 of
                                                            true ->
                                                                Acc@2;

                                                            false ->
                                                                T_axis_size = case begin
                                                                    _pipe@6 = gleam@list:drop(
                                                                        viva_tensor@core@tensor:shape(
                                                                            T@3
                                                                        ),
                                                                        Axis
                                                                    ),
                                                                    gleam@list:first(
                                                                        _pipe@6
                                                                    )
                                                                end of
                                                                    {ok, D@3} ->
                                                                        D@3;

                                                                    {error, _} ->
                                                                        0
                                                                end,
                                                                case Remaining < T_axis_size of
                                                                    true ->
                                                                        {T_idx,
                                                                            Remaining,
                                                                            T_idx};

                                                                    false ->
                                                                        {-1,
                                                                            Remaining
                                                                            - T_axis_size,
                                                                            T_idx
                                                                            + 1}
                                                                end
                                                        end
                                                    end
                                                ),
                                                Local_indices = begin
                                                    _pipe@7 = Indices,
                                                    gleam@list:index_map(
                                                        _pipe@7,
                                                        fun(Idx, I@4) ->
                                                            case I@4 =:= Axis of
                                                                true ->
                                                                    Local_axis_idx;

                                                                false ->
                                                                    Idx
                                                            end
                                                        end
                                                    )
                                                end,
                                                case begin
                                                    _pipe@8 = gleam@list:drop(
                                                        Tensors,
                                                        Tensor_idx
                                                    ),
                                                    gleam@list:first(_pipe@8)
                                                end of
                                                    {ok, T@4} ->
                                                        T_strides = compute_strides(
                                                            viva_tensor@core@tensor:shape(
                                                                T@4
                                                            )
                                                        ),
                                                        Local_flat = begin
                                                            _pipe@9 = gleam@list:zip(
                                                                Local_indices,
                                                                T_strides
                                                            ),
                                                            gleam@list:fold(
                                                                _pipe@9,
                                                                0,
                                                                fun(A, P) ->
                                                                    A + (erlang:element(
                                                                        1,
                                                                        P
                                                                    )
                                                                    * erlang:element(
                                                                        2,
                                                                        P
                                                                    ))
                                                                end
                                                            )
                                                        end,
                                                        T_data = viva_tensor@core@tensor:to_list(
                                                            T@4
                                                        ),
                                                        case begin
                                                            _pipe@10 = gleam@list:drop(
                                                                T_data,
                                                                Local_flat
                                                            ),
                                                            gleam@list:first(
                                                                _pipe@10
                                                            )
                                                        end of
                                                            {ok, V} ->
                                                                V;

                                                            {error, _} ->
                                                                +0.0
                                                        end;

                                                    {error, _} ->
                                                        +0.0
                                                end
                                            end
                                        )
                                    end,
                                    viva_tensor@core@tensor:new(
                                        Result,
                                        New_shape
                                    )
                            end
                    end
            end
    end.

-file("src/viva_tensor/core/shape.gleam", 353).
?DOC(" Stack tensors along a new axis\n").
-spec stack(list(viva_tensor@core@tensor:tensor()), integer()) -> {ok,
        viva_tensor@core@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
stack(Tensors, Axis) ->
    case Tensors of
        [] ->
            {error, {invalid_shape, <<"Cannot stack empty list"/utf8>>}};

        [First | Rest] ->
            Base_shape = viva_tensor@core@tensor:shape(First),
            Shapes_ok = gleam@list:all(
                Rest,
                fun(T) -> viva_tensor@core@tensor:shape(T) =:= Base_shape end
            ),
            case Shapes_ok of
                false ->
                    {error, {shape_mismatch, Base_shape, []}};

                true ->
                    _ = erlang:length(Tensors),
                    R = erlang:length(Base_shape),
                    Insert_axis = case Axis < 0 of
                        true ->
                            (R + Axis) + 1;

                        false ->
                            Axis
                    end,
                    case (Insert_axis >= 0) andalso (Insert_axis =< R) of
                        false ->
                            {error,
                                {dimension_error,
                                    <<"Invalid axis for stacking"/utf8>>}};

                        true ->
                            Unsqueezed = gleam@list:map(
                                Tensors,
                                fun(T@1) -> unsqueeze(T@1, Insert_axis) end
                            ),
                            concat_axis(Unsqueezed, Insert_axis)
                    end
            end
    end.

-file("src/viva_tensor/core/shape.gleam", 409).
-spec multi_to_flat(list(integer()), list(integer())) -> integer().
multi_to_flat(Indices, Shape) ->
    Strides = compute_strides(Shape),
    _pipe = gleam@list:zip(Indices, Strides),
    gleam@list:fold(
        _pipe,
        0,
        fun(Acc, Pair) ->
            {Idx, Stride} = Pair,
            Acc + (Idx * Stride)
        end
    ).

-file("src/viva_tensor/core/shape.gleam", 428).
-spec list_at(list(KQY), integer()) -> {ok, KQY} | {error, nil}.
list_at(Lst, Index) ->
    case Index < 0 of
        true ->
            {error, nil};

        false ->
            _pipe = Lst,
            _pipe@1 = gleam@list:drop(_pipe, Index),
            gleam@list:first(_pipe@1)
    end.

-file("src/viva_tensor/core/shape.gleam", 62).
?DOC(" Remove dimension at specific axis if it's 1\n").
-spec squeeze_axis(viva_tensor@core@tensor:tensor(), integer()) -> {ok,
        viva_tensor@core@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
squeeze_axis(T, Axis) ->
    Shp = viva_tensor@core@tensor:shape(T),
    case list_at(Shp, Axis) of
        {error, _} ->
            {error, {dimension_error, <<"Axis out of bounds"/utf8>>}};

        {ok, D} ->
            case D =:= 1 of
                false ->
                    {error,
                        {invalid_shape, <<"Dimension at axis is not 1"/utf8>>}};

                true ->
                    Data = viva_tensor@core@tensor:to_list(T),
                    New_shape = begin
                        _pipe = Shp,
                        _pipe@1 = gleam@list:index_map(
                            _pipe,
                            fun(Dim, I) -> {Dim, I} end
                        ),
                        _pipe@2 = gleam@list:filter(
                            _pipe@1,
                            fun(Pair) -> erlang:element(2, Pair) /= Axis end
                        ),
                        gleam@list:map(
                            _pipe@2,
                            fun(Pair@1) -> erlang:element(1, Pair@1) end
                        )
                    end,
                    viva_tensor@core@tensor:new(Data, New_shape)
            end
    end.

-file("src/viva_tensor/core/shape.gleam", 438).
-spec list_at_float(list(float()), integer()) -> {ok, float()} | {error, nil}.
list_at_float(Lst, Index) ->
    list_at(Lst, Index).

-file("src/viva_tensor/core/shape.gleam", 147).
?DOC(
    " General slice - specify start indices and lengths for each dimension.\n"
    " This one's tricky for n-dimensional tensors.\n"
).
-spec slice(viva_tensor@core@tensor:tensor(), list(integer()), list(integer())) -> {ok,
        viva_tensor@core@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
slice(T, Start, Lengths) ->
    Data = viva_tensor@core@tensor:to_list(T),
    Shp = viva_tensor@core@tensor:shape(T),
    R = viva_tensor@core@tensor:rank(T),
    case (erlang:length(Start) =:= R) andalso (erlang:length(Lengths) =:= R) of
        false ->
            {error,
                {dimension_error,
                    <<"Slice dimensions must match tensor rank"/utf8>>}};

        true ->
            case R of
                1 ->
                    S = case gleam@list:first(Start) of
                        {ok, V} ->
                            V;

                        {error, _} ->
                            0
                    end,
                    Len = case gleam@list:first(Lengths) of
                        {ok, V@1} ->
                            V@1;

                        {error, _} ->
                            0
                    end,
                    Sliced = begin
                        _pipe = Data,
                        _pipe@1 = gleam@list:drop(_pipe, S),
                        gleam@list:take(_pipe@1, Len)
                    end,
                    viva_tensor@core@tensor:new(Sliced, [Len]);

                _ ->
                    New_size = gleam@list:fold(
                        Lengths,
                        1,
                        fun(Acc, D) -> Acc * D end
                    ),
                    Result = begin
                        _pipe@2 = gleam@list:range(0, New_size - 1),
                        gleam@list:map(
                            _pipe@2,
                            fun(Flat_idx) ->
                                Local_indices = flat_to_multi(Flat_idx, Lengths),
                                Global_indices = gleam@list:map2(
                                    Local_indices,
                                    Start,
                                    fun(L, S@1) -> L + S@1 end
                                ),
                                Global_flat = multi_to_flat(Global_indices, Shp),
                                case list_at_float(Data, Global_flat) of
                                    {ok, V@2} ->
                                        V@2;

                                    {error, _} ->
                                        +0.0
                                end
                            end
                        )
                    end,
                    viva_tensor@core@tensor:new(Result, Lengths)
            end
    end.