from fractions import Fraction

def bernoulli_tree( hauteur,\
                    proba,\
                    succes = 'S',\
                    espv = 0.5,\
                    espvinter = 1,\
                    espx = 2,\
                    fraction = False,\
                    root = '$\\Omega$',\
                    stylebranche = '',\
                    stylenoeud = '',\
                    styleproba = 'midway,fill=white,scale=.5',
                    filename = 'arbre'):
    
    # Stéphane Pasquet
    # mathweb.fr
    # 2022/07/02
    # https://www.mathweb.fr/euclide/2022/07/02/latex-et-python-creer-un-arbre-de-bernoulli-automatiquement/
    
    if fraction:
        proba_num = Fraction(str(proba)).numerator
        proba_denom = Fraction(str(proba)).denominator
        if isinstance(proba,float):
            proba_echec_num = Fraction(str(1-proba)).numerator
            proba_echec_denom = Fraction(str(1-proba)).denominator
        else:
            proba_echec_num = Fraction(str(1-Fraction(proba))).numerator
            proba_echec_denom = Fraction(str(1-Fraction(proba))).denominator
        proba = '$\\frac{'+str(proba_num)+'}{'+str(proba_denom)+'}$'
        probaechec = '$\\frac{'+str(proba_echec_num)+'}{'+str(proba_echec_denom)+'}$'
    else:
        probaechec = str(1-proba).replace(".",",")
        proba = str(proba).replace(".",",")
        
    tex = '\\documentclass{standalone}\n'
    tex += '\\usepackage{tikz}\n'
    tex += '\\usetikzlibrary{calc}\n'
    tex += '\\setlength{\\parindent}{0pt}\n'
    tex += '\\begin{document}\n'

    tex += '\\begin{tikzpicture}\n'
    tex += '\\tikzstyle{branche} = [ '+stylebranche+' ]\n'
    tex += '\\tikzstyle{noeud} = [ '+stylenoeud+' ]\n'
    tex += '\\tikzstyle{proba} = [ '+styleproba+' ]\n'
    
    level = hauteur
    echec = '\\overline{' + succes + '}'
    
    # placement des noeuds
    
    while level != 0:
        if level == hauteur:
            y = 0
            for i in range(2**level):
                if i%2 == 0:
                    tex += '\\node[noeud] (N' + str(level) + '-' + str(i) + ') at ({'+str(espx)+'*'+str(level)+'},'+str(y)+') {$'+succes+'$};\n'
                    y -= espv
                else:
                    tex += '\\node[noeud] (N' + str(level) + '-' + str(i) + ') at ({'+str(espx)+'*'+str(level)+'},'+str(y)+') {$'+echec+'$};\n'
                    y -= espvinter
        else:
            for i in range(2**level):
                if i%2 == 0:
                    tex += '\\node[noeud] (N' + str(level) + '-' + str(i) + ') at ($0.5*(N'+str(level+1)+'-'+str(2*i)+')+0.5*(N'+str(level+1)+'-'+str(2*i+1)+')+(-'+str(espx)+',0)$) \
{$'+succes+'$};\n'
                else:
                     tex += '\\node[noeud] (N' + str(level) + '-' + str(i) + ') at ($0.5*(N'+str(level+1)+'-'+str(2*i)+')+0.5*(N'+str(level+1)+'-'+str(2*i+1)+')+(-'+str(espx)+',0)$) \
{$'+echec+'$};\n'
        
        level -= 1
    
    # tracé des branches
    
    tex += '\\node[left,noeud] at ($0.5*(N1-0)+0.5*(N1-1)+(-'+str(espx)+',0)$) {'+root+'};\n'
    tex += '\\draw[branche] ($0.5*(N1-0)+0.5*(N1-1)+(-'+str(espx)+',0)$) -- (N1-0) node[proba] {'+proba+'};\n'
    tex += '\\draw[branche] ($0.5*(N1-0)+0.5*(N1-1)+(-'+str(espx)+',0)$) -- (N1-1) node[proba] {'+probaechec+'};\n'
    
    for lev in range(1,hauteur):
        for i in range(2**lev):
            tex += '\\draw[branche] (N'+str(lev)+'-'+str(i)+'.east) -- (N'+str(lev+1)+'-'+str(2*i)+') node[proba] {'+proba+'};\n'
            tex += '\\draw[branche] (N'+str(lev)+'-'+str(i)+'.east) -- (N'+str(lev+1)+'-'+str(2*i+1)+') node[proba] {'+probaechec+'};\n'

    tex += '\\end{tikzpicture}\n'
    tex += '\\end{document}'
    
    # sauvegarde du fichier LaTeX
    
    from os.path import isfile
    from os import remove, system
    
    if isfile(filename+".tex"):
        remove(filename+".tex")
    
    fichier = open(filename+".tex","x")
    fichier.write(tex)
    fichier.close()

    # compilation PdfLaTeX dans le répertoire courant

    cmd = "pdflatex  --shell-escape -synctex=1 -interaction=nonstopmode "+filename+".tex"
    system(cmd)
    cmd = "START "+filename+".pdf"
    system(cmd)