%%% @doc Generic Zipper Implementation.
%%%      Zippers let you traverse immutable data structures with ease and
%%%      flexibility.
-module(zipper).

-compile({no_auto_import, [node/1]}).

%% Creation
-export([new/4]).
%% Traverse
-export([up/1, down/1, left/1, leftmost/1, right/1, rightmost/1, next/1, is_end/1, prev/1,
         root/1, traverse/2]).
%% Editing
-export([insert_left/2, insert_right/2, replace/2, edit/3, insert_child/2, append_child/2,
         remove/1]).
%% Iteration
-export([map/2, fmap/3, filter/2, fold/3, size/1]).
%% Info
-export([node/1, children/1, is_branch/1]).

-type is_branch_fun(T) :: fun((T) -> boolean()).
-type make_node_fun(T) :: fun((T, [T]) -> T).
-type children_fun(T) :: fun((T) -> [T]).
-type info(T) ::
    #{lefts => [T],
      rights => [T],
      parent_node => undefined | T,
      parent_info => undefined | info(T),
      is_modified => boolean()}.

-opaque zipper(T) ::
    #{spec =>
          #{is_branch => is_branch_fun(T),
            make_node => make_node_fun(T),
            children => children_fun(T)},
      node => T,
      info => 'end' | info(T)}.

-type operation() ::
    next | prev | up | down | left | right | root | node | rightmost | leftmost.

-export_type([zipper/1]).
-export_type([info/1, is_branch_fun/1, make_node_fun/1, children_fun/1]).
-export_type([operation/0]).

-ifdef(TEST).

-export([as_list_of_maps/1]).

-endif.

-elvis([{elvis_style, dont_repeat_yourself, disable}]).
-elvis([{elvis_style, god_modules, disable}]).
-elvis([{elvis_style, no_throw, disable}]).

-ifdef(TEST).

-spec as_list_of_maps([zipper(any())]) -> [map()].
as_list_of_maps(Children) ->
    Children.

-endif.

%% @doc Builds a new zipper with nodes of type T.
-spec new(is_branch_fun(T), children_fun(T), make_node_fun(T), T) -> zipper(T).
new(IsBranch, Children, MakeNode, Root) ->
    #{spec =>
          #{is_branch => IsBranch,
            children => Children,
            make_node => MakeNode},
      node => Root,
      info =>
          #{lefts => [],
            rights => [],
            parent_node => undefined,
            parent_info => undefined}}.

%% @doc Returns the zipper in the parent node, if possible.
-spec up(zipper(T)) -> zipper(T) | undefined.
up(#{info := #{parent_node := undefined}}) ->
    undefined;
up(#{spec := #{make_node := MakeNode},
     node := Node,
     info :=
         #{lefts := Lefts,
           rights := Rights,
           parent_node := ParentNode,
           parent_info := ParentInfo,
           is_modified := true}} =
       Zipper) ->
    Children = lists:reverse(Lefts) ++ [Node | Rights],
    NewParentNode = MakeNode(ParentNode, Children),
    Zipper#{node => NewParentNode, info => ParentInfo};
up(#{info := #{parent_node := Parent, parent_info := ParentInfo}} = Zipper) ->
    Zipper#{node => Parent, info => ParentInfo}.

%% @doc Returns the zipper in the first child, if any.
-spec down(zipper(T)) -> zipper(T) | undefined.
down(#{node := Node,
       info := Info,
       spec := #{children := Children}} =
         Zipper) ->
    case is_branch(Zipper) of
        true ->
            [NewNode | Rights] = Children(Node),
            Zipper#{node => NewNode,
                    info =>
                        #{lefts => [],
                          rights => Rights,
                          parent_node => Node,
                          parent_info => Info}};
        false ->
            undefined
    end.

%% @doc Returns the zipper on the left, if any.
-spec left(zipper(T)) -> zipper(T) | undefined.
left(#{info := #{lefts := []}}) ->
    undefined;
left(#{info := Info = #{lefts := [NewNode | Lefts], rights := Rights}, node := Node} =
         Zipper) ->
    Zipper#{info => Info#{lefts => Lefts, rights => [Node | Rights]}, node => NewNode}.

%% @doc Returns the leftmost zipper in the current zipper.
-spec leftmost(zipper(T)) -> zipper(T).
leftmost(#{info := #{lefts := []}} = Zipper) ->
    Zipper;
leftmost(#{info := Info = #{lefts := Lefts, rights := Rights}, node := Node} = Zipper) ->
    Fun = fun(Item, Acc) -> [Item | Acc] end,
    [NewNode | NewRights] = lists:foldl(Fun, [Node | Rights], Lefts),
    Zipper#{info => Info#{lefts => [], rights => NewRights}, node => NewNode}.

%% @doc Returns the zipper on the right, if any.
-spec right(zipper(T)) -> zipper(T) | undefined.
right(#{info := #{rights := []}}) ->
    undefined;
right(#{info := Info = #{rights := [NewNode | Rights], lefts := Lefts}, node := Node} =
          Zipper) ->
    Zipper#{info => Info#{rights := Rights, lefts := [Node | Lefts]}, node := NewNode}.

%% @doc Returns the rightmost zipper in the current zipper.
-spec rightmost(zipper(T)) -> zipper(T).
rightmost(#{info := Info = #{rights := Rights, lefts := Lefts}, node := Node} = Zipper) ->
    Fun = fun(Item, Acc) -> [Item | Acc] end,
    [NewNode | NewLefts] = lists:foldl(Fun, [Node | Lefts], Rights),
    Zipper#{info => Info#{lefts => NewLefts, rights => []}, node => NewNode}.

%% @doc Returns the next zipper.
-spec next(zipper(T)) -> zipper(T).
next(#{info := 'end'} = Zipper) ->
    Zipper;
next(Zipper) ->
    case {is_branch(Zipper), right(Zipper)} of
        {true, _} ->
            down(Zipper);
        {false, undefined} ->
            next_recur(Zipper);
        {false, Right} ->
            Right
    end.

-spec next_recur(zipper(T)) -> zipper(T).
next_recur(Zipper) ->
    case up(Zipper) of
        undefined ->
            Zipper#{info => 'end'};
        UpZipper ->
            case right(UpZipper) of
                undefined ->
                    next_recur(UpZipper);
                Next ->
                    Next
            end
    end.

%% @doc Is it the end of the zipper traversal.
-spec is_end(zipper(_)) -> boolean().
is_end(#{info := 'end'}) ->
    true;
is_end(_Zipper) ->
    false.

%% @doc Returns the previous zipper.
-spec prev(zipper(T)) -> zipper(T) | undefined.
prev(#{info := #{lefts := []}} = Zipper) ->
    up(Zipper);
prev(Zipper) ->
    prev_recur(left(Zipper)).

-spec prev_recur(zipper(T)) -> zipper(T).
prev_recur(Zipper) ->
    case down(Zipper) of
        undefined ->
            Zipper;
        DownZipper ->
            RightMost = rightmost(DownZipper),
            prev_recur(RightMost)
    end.

%% @doc Returns the node on the root of the zipper.
-spec root(zipper(T)) -> T.
root(Zipper) ->
    case up(Zipper) of
        undefined ->
            node(Zipper);
        Parent ->
            root(Parent)
    end.

%% @doc Traverses the zipper following the given list of operations.
%%      If, at some point, an operation is invalid, it will crash.
-spec traverse([operation()], zipper(T)) -> zipper(T) | T | undefined.
traverse([], Zipper) ->
    Zipper;
traverse([Op | Rest], Zipper) ->
    traverse(Rest, zipper:Op(Zipper)).

%% @doc Inserts a node to the left of the current one.
-spec insert_left(T, zipper(T)) -> zipper(T).
insert_left(_, #{info := #{parent_node := undefined}}) ->
    throw(insert_at_top);
insert_left(Node, #{info := Info = #{lefts := Lefts}} = Zipper) ->
    NewInfo = Info#{lefts => [Node | Lefts], is_modified => true},
    Zipper#{info => NewInfo}.

%% @doc Inserts a node to the right of the current one.
-spec insert_right(T, zipper(T)) -> zipper(T).
insert_right(_, #{info := #{parent_node := undefined}}) ->
    throw(insert_at_top);
insert_right(Node, #{info := Info = #{rights := Rights}} = Zipper) ->
    NewInfo = Info#{rights => [Node | Rights], is_modified => true},
    Zipper#{info => NewInfo}.

%% @doc Replaces the current node.
-spec replace(T, zipper(T)) -> zipper(T).
replace(Node, #{info := Info} = Zipper) ->
    Zipper#{node => Node, info => Info#{is_modified => true}}.

%% @doc Edits the current node by applying the given function.
%%      The parameters of said function will be [Node | Args].
-spec edit(fun((...) -> T), [term()], zipper(T)) -> zipper(T).
edit(Fun, Args, #{node := Node, info := Info} = Zipper) ->
    NewNode = erlang:apply(Fun, [Node | Args]),
    Zipper#{node => NewNode, info => Info#{is_modified => true}}.

%% @doc Adds a node as the leftmost child of the current one.
-spec insert_child(T, zipper(T)) -> zipper(T).
insert_child(Child, #{spec := #{make_node := MakeNode}} = Zipper) ->
    Node = node(Zipper),
    Children = children(Zipper),
    NewNode = MakeNode(Node, [Child | Children]),
    replace(NewNode, Zipper).

%% @doc Adds a node as the rightmost child of the current one.
-spec append_child(T, zipper(T)) -> zipper(T).
append_child(Child, #{spec := #{make_node := MakeNode}} = Zipper) ->
    Node = node(Zipper),
    Children = children(Zipper),
    NewNode = MakeNode(Node, Children ++ [Child]),
    replace(NewNode, Zipper).

%% @doc Removes current node from zipper.
%%      Moves down, if possible. If not, it moves to the rightmost node.
-spec remove(zipper(T)) -> zipper(T).
remove(#{info := #{parent_node := undefined}}) ->
    throw(remove_at_top);
remove(#{spec := Spec = #{make_node := MakeNode},
         info :=
             #{lefts := [],
               rights := Rights,
               parent_node := ParentNode,
               parent_info := ParentInfo}}) ->
    NewParentNode = MakeNode(ParentNode, Rights),
    #{spec => Spec,
      node => NewParentNode,
      info => ParentInfo#{is_modified => true}};
remove(#{spec := Spec, info := Info = #{lefts := [Node | Lefts]}}) ->
    NewZipper =
        #{spec => Spec,
          node => Node,
          info => Info#{lefts => Lefts, is_modified => true}},
    prev_removed(NewZipper).

prev_removed(Zipper) ->
    case down(Zipper) of
        undefined ->
            Zipper;
        Child ->
            prev_removed(rightmost(Child))
    end.

%% Iteration

%% @doc Applies a function to all nodes of the zipper.
%%      Returns a list with the results.
-spec map(fun((T) -> U), zipper(T)) -> [U].
map(Fun, Zipper) ->
    ApplyAddFun = fun(X, Acc) -> [Fun(X) | Acc] end,
    Result = fold(ApplyAddFun, [], Zipper),
    lists:reverse(Result).

%% @doc Returns the root of the tree, where the value of each node
%%      (after the current location of Zipper) is replaced with the
%%      result from applying Fun to the node as the first argument
%%      and Args as additional arguments.
-spec fmap(fun((...) -> T), [term()], zipper(T)) -> T.
fmap(Fun, Args, Zipper) ->
    NewZipper = edit(Fun, Args, Zipper),
    NextZipper = next(NewZipper),
    case is_end(NextZipper) of
        true ->
            root(NewZipper);
        false ->
            fmap(Fun, Args, NextZipper)
    end.

%% @doc Applies Fun recursively on the zipper.
%%      The arguments of Fun will be (Node, Acc) where Acc is the result of the
%%      previous call or the initial value provided.
-spec fold(fun((T, A) -> A), A, zipper(T)) -> A.
fold(Fun, Acc, Zipper) ->
    case is_end(Zipper) of
        true ->
            Acc;
        false ->
            Node = node(Zipper),
            NewAcc = Fun(Node, Acc),
            fold(Fun, NewAcc, next(Zipper))
    end.

%% @doc Returns a list of all the nodes in the zipper that match a predicate.
-spec filter(fun((T) -> boolean()), zipper(T)) -> [T].
filter(Pred, Zipper) ->
    FilterFun =
        fun(X, Acc) ->
           case Pred(X) of
               true ->
                   [X | Acc];
               false ->
                   Acc
           end
        end,
    fold(FilterFun, [], Zipper).

%% @doc Returns the size of the zipper.
-spec size(zipper(_)) -> non_neg_integer().
size(Zipper) ->
    IncFun = fun(_, Acc) -> Acc + 1 end,
    fold(IncFun, 0, Zipper).

%% Info

%% @doc Returns the value of the current node in the zipper.
-spec node(zipper(T)) -> T.
node(#{node := Node}) ->
    Node.

%% @doc Returns the list of children zippers.
-spec children(zipper(T)) -> [zipper(T)].
children(#{spec := #{children := Children}, node := Node} = Zipper) ->
    case is_branch(Zipper) of
        true ->
            Children(Node);
        false ->
            throw(children_on_leaf)
    end.

%% @doc Is this node a branch?
-spec is_branch(zipper(_)) -> boolean().
is_branch(#{spec := #{is_branch := IsBranch}, node := Node}) ->
    IsBranch(Node).