Adding some doc and generalyzing the code