BeginPackage["Unflatten`"]

Unflatten::usage = "Unflatten[h[elems], posints, heads] creates a nested structure by wrapping
the elements in positions intervals posints in heads heads";

UnflattenNested::usage = "UnflattenNested[expr, {startposlist, endposlist}, head] wraps elements
of expr at positions starting at startposlist_i and ending at endposlist_i in head head";

Begin["`Private`"]

ClearAll[listSplit,reconstructIntervals,groupElements,groupPositions,processPosList,groupElementsNested];

listSplit[x_List, lengthlist_List,headlist_List]:=
	MapThread[
		#1@@Take[x,#2]&,
		{headlist,Transpose[{Most[#]+1,Rest[#]}&[FoldList[Plus,0,lengthlist]]]}
	];

reconstructIntervals[listlen_Integer,ints_List]:=
	Module[{missed,startint,lastint},
		startint = If[ints[[1,1]]==1,{},{1,ints[[1,1]]-1}];
		lastint = If[ints[[-1,-1]]==listlen, {},{ints[[-1,-1]]+1,listlen}];
		missed =
			Map[
				If[#[[2,1]]-#[[1,2]]>1,{#[[1,2]]+1,#[[2,1]]-1},{}]&,
				Partition[ints,2,1]
			];
		missed = Join[missed,{lastint}];
		Prepend[Flatten[Transpose[{ints,missed}],1],startint]
	];

groupElements[
	lst_List,poslist_List,headlist_List
] /; And[OrderedQ[Flatten[Sort[poslist]]],Length[headlist]==Length[poslist]]:=
	Module[{totalheadlist,allints,llist},
		totalheadlist =
			Append[
				Flatten[Transpose[{Array[Sequence&,{Length[headlist]}],headlist}],1],
				Sequence
			];
		allints = reconstructIntervals[Length[lst],poslist];
		llist = Map[If[#==={},0,1-Subtract@@#]&,allints];
		listSplit[lst, llist,totalheadlist]
	];

(* To work on general heads, we need this *)

groupElements[h_[x__],poslist_List,headlist_List]:=
	h[Sequence@@groupElements[{x},poslist,headlist]];

(* If we have a single head *)
groupElements[expr_,poslist_List,head_]:=
	groupElements[expr,poslist,Table[head,{Length[poslist]}]];


groupPositions[plist_List]:=
	Reap[Sow[Last[#],{Most[#]}]&/@plist,_,List][[2]];

processPosList::unmtch = "The starting and ending position lists `1` and `2` don't match";
processPosList[{openlist_List,closelist_List}]:=
	Module[{opengroup,closegroup,poslist},
		{opengroup,closegroup} = groupPositions/@{openlist,closelist};
		poslist = Transpose[Transpose[Sort[#]]&/@{opengroup,closegroup}];
		If[UnsameQ @@ poslist[[1]],
			Message[processPosList::unmtch,openlist,closelist];
			{},
			(* else *)
			Transpose[{poslist[[1,1]],Transpose/@ Transpose[poslist[[2]]]}]
		]
	];

groupElementsNested[nested_,{openposlist_List,closeposlist_List},head_]/;Head[head]=!=List:=
	Fold[
		Function[{x,y},MapAt[groupElements[#,y[[2]],head]&,x,{y[[1]]}]],
		nested,
		Sort[processPosList[{openposlist,closeposlist}],Length[#2[[1]]] < Length[#1[[1]]]&]
	];

Unflatten[args___]:= groupElements[args];

UnflattenNested[args___]:=groupElementsNested[args]

End[]

EndPackage[]