import argparse
from pathlib import Path

class Tree:
    def __init__(self, dirs):
        self.dirs = dirs;
        self.tree = {}
        self._build()

    # Builds a tree represented as dict, with structure of files
    # and directories.
    def _build(self):
        for di in self.dirs:
            path_list = list(Path(di).glob('**/*.cpp'))
            path_list += list(Path(di).glob('**/*.vim'))
            path_list = sorted(path_list)

            for path in path_list:
                branch = str(path).split('/')
                node = self.tree

                for i in branch[:-2]:
                    if i not in node:
                        node[i] = {}
                    node = node[i]

                if branch[-2] not in node:
                    node[branch[-2]] = []
                node[branch[-2]].append(str(branch[-1]))

    # Allows access to tree dict from instance of class.
    def __getitem__(self, arg):
        return self.tree[arg]

    # Allows iteration on tree dict from instance of class.
    def __iter__(self):
        return iter(self.tree)


class LatexGenerator:

    class SourceFile:
        def __init__(self, path):
            source = open(path, 'r')
            
            self.type = path.split('.')[-1]
            self.comment = '///' if self.type == 'cpp' else '"""'

            self.header = {}
            self.code, self.raw_header = [], []
            self.sections = ['Description', 'Time', 'Space', 'Include', 
                    'Status']

            self._parse_source(source)
            self._parse_header()

        # Splits content of the file between header and code.
        def _parse_source(self, source):
            lines = source.readlines()
            in_header = True

            for i, line in enumerate(lines):
                if line[0:3] != self.comment:
                    in_header = False

                if in_header:
                    self.raw_header.append(line[4:-1])
                else:
                    self.code = lines[(i+1):]
                    break
        
        # Returns section name of line or None if nothing is found.
        def _get_section(self, line):
            comm = line.split(':')[0]
            return comm if comm in self.sections else None

        # Builds header dict by separating the header on the 
        # found sections.
        def _parse_header(self):
            self.name = self.raw_header[0]

            curr = None
            for i in self.raw_header[1:]:
                section = self._get_section(i)

                if section != None:
                    curr = section 
                    self.header[curr] = []

                    # Add rest of line (right side of section name)
                    rest = ':'.join(i.split(':')[1:]).strip()
                    if len(rest) > 0:
                        self.header[curr].append(rest)
                else:
                    if curr != None:
                        line = i.strip()
                        if len(line) > 0:
                            self.header[curr].append(line)

            # Make description a single line
            if 'Description' in self.header:
                text = self.header['Description']
                self.header['Description'] = ' '.join(text)

        # Writes own code to LaTeX output.
        def output_code(self, write):
            write('\\begin{lstlisting}[style=custom%s]\n' % self.type)
            for i in self.code:
                write(i)
            write('\\end{lstlisting}\n')
            write('\n')

        # Writes own header to LaTeX output.
        def output_header(self, write):
            if len(self.header) > 0:
                write('\\begin{adjustwidth}{5pt}{5pt}\n')

            if 'Description' in self.header:
                write('\\textbf{\\footnotesize Description: }\n')
                write('\\footnotesize{\\par %s}' % self.header['Description'])
                write('\\\\\n')
                write('~\\\\\n')

            # Replaces text with LaTeX on Big-O notations. 
            def to_big_o(expr):
                expr = expr.replace('O', '\\mathcal{O}')
                expr = expr.replace('log', '\\log')
                expr = expr.replace('*', '\\times')
                expr = expr.replace('sqrt', '\\sqrt')
                return expr

            # Outputs time or space complexity.
            def output_complexity(comp):
                if comp in self.header:
                    write('\\textbf{\\footnotesize %s: }' % comp)
                    if len(self.header[comp]) > 1:
                        write('\n')
                        write('\\begin{compactitem}\n')
                        for i in self.header[comp]:
                            line = i[1:].split(':')
                            write('\\item{\\footnotesize{%s: $%s$}}\n' % 
                                    (line[0].strip().replace('_', '\\_'), 
                                     to_big_o(line[1].strip())))
                        write('\\end{compactitem}\n')
                    else:
                        write('\\footnotesize{$%s$}\n' % 
                                to_big_o(self.header[comp][0]))
                    return True
                return False

            if output_complexity('Time'):
                write('~\\\\\n')
            output_complexity('Space')

            if len(self.header) > 0:
                write('\\end{adjustwidth}\n')

    def __init__(self, tree, output, header):
        self.output = open(output, 'w')
        self.header = open(header, 'r')

        self.tree = tree
        self.files = {}
        self.hierarchy = [i*'sub' + 'section' for i in range(3)]

        self._gen_latex()

    # Generates full LaTeX.
    def _gen_latex(self):
        self._add_header()

        self._write('\\begin{document}\n')
        self._write('\\tableofcontents\n')
        self._write('\\newpage\n\n')

        self._gen_tree_latex(self.tree)

        self._write('\\end{document}\n')

    # Print text to output file.
    def _write(self, text):
        self.output.write(text)

    # Adds Latex header to the output.
    def _add_header(self):
        lines = self.header.readlines()
        for i in lines:
            self._write(i)

    # Prints section title in LaTeX format.
    def _gen_title_latex(self, content, depth):
        self._write('\\%s{%s}\n' % (self.hierarchy[depth], content))

    # Prints code in LaTeX format.
    def _gen_code_latex(self, source):
        source.output_header(self._write)
        source.output_code(self._write)
        self._write('\\hrule\n')

    # Generates LaTeX for entire tree recursively.
    def _gen_tree_latex(self, sub, path = '', depth = 0):
        if type(sub) == list:
            for i in sub:
                source = self.SourceFile(path + i)
                self._gen_title_latex(source.name, depth)
                self._gen_code_latex(source)
        else:
            if depth == 1:
                self._write('\\begin{multicols}{3}\n')
            for i in sub:
                self._gen_title_latex(self._parse_dir_name(i), depth)
                self._gen_tree_latex(sub[i], path + i + '/', depth + 1)
            if depth == 1:
                self._write('\\end{multicols}\n')

        self._write('\n')

    # Parses name of the section (capitalize).
    def _parse_dir_name(self, name):
        return ' '.join(list(map(lambda x : x.capitalize(), name.split('_'))))


# Parses command line arguments and returns them.
def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--header', action='store', type=str, 
            help='Header of Latex file')
    parser.add_argument('--output', action='store', type=str, 
            help='Output Latex file')
    return parser.parse_args()

def main():
    args = get_args()
    tree = Tree(['algorithms', 'misc'])
    tex = LatexGenerator(tree, args.output, args.header)

if __name__ == "__main__":
    main()