from fractions import Fraction

def formater_prob(prob_str,mode):
    """Formate une probabilité saisie sous forme LaTeX."""
    if mode=="fraction":
        if "/" in prob_str:
            num, den = prob_str.split("/")
        else:
            num, den = Fraction(prob_str).numerator, Fraction(prob_str).denominator
        return f"$\\dfrac{{{num.strip()}}}{{{den.strip()}}}$"
    else:
        return str(float(Fraction(prob_str))).replace(".", ",")

def generer_nodes(noeud, compteur, prob_composee=Fraction(1), mode="decimal", decimales=3):
    code = ""
    style = "noeud"

    if not noeud["enfants"]:
        style = "noeud"
        nom = f"f{compteur[0]}"
        if mode == "fraction":
            p = prob_composee
            prob_composee_str = f"\\dfrac{{{p.numerator}}}{{{p.denominator}}}"
        else:
            prob_float = float(prob_composee)
            prob_composee_str = f"{prob_float:.{decimales}f}".replace(".", ",")
        code += f'[{style}] ({nom}) {{{noeud["label"]}}}\n'
        compteur[1].append((nom, prob_composee_str))
        compteur[0] += 1
    else:
        code += f'[{style}] {{{noeud["label"]}}}\n'
        for enfant in noeud["enfants"]:
            try:
                prob = Fraction(enfant["prob"])
            except ValueError:
                prob = Fraction(enfant["prob"].replace(",", "."))
            code += "child {\n"
            code += "node " + generer_nodes(
                enfant,
                compteur,
                prob_composee=prob_composee * prob,
                mode=mode,
                decimales=decimales
            )
            prob_label = formater_prob(enfant["prob"],mode)
            code += f'edge from parent node[prob] {{{prob_label}}}\n'
            code += "}\n"

    return code


def generer_fleches(feuilles, longueur, flechecolor):
    code = ""
    for nom, prob in feuilles:
        code += f'\\draw[->, thick, {flechecolor}] ({nom}.east) -- ++({longueur},0)\n'
        code += f'  node[right] {{\\small ${prob}$}};\n'
    return code


def arbre_tikz(arbre):
    compteur = [0, []]

    # Récupération des paramètres avec valeurs par défaut
    mode      = arbre.get("mode", "decimal")
    decimales = arbre.get("decimales", 3)
    ld1       = arbre.get("level distance 1", 3)
    ld2       = arbre.get("level distance 2", 3)
    sd1       = arbre.get("sibling distance 1", 4)
    sd2       = arbre.get("sibling distance 2", 2)
    flechelength    = arbre.get("flechelength", 1)
    flechecolor = arbre.get("flechecolor","black")
    edgestyle = arbre.get("edgestyle","black")
    nodestyle = arbre.get("nodestyle", "fill=white")
    probastyle = arbre.get("probastyle","fill=white, inner sep=1mm")
    probaend  = arbre.get("probaend",False)
    preambule = rf"""
\begin{{tikzpicture}}[
  >=latex,
  grow=right,
  edge from parent/.style={{draw, {edgestyle}}},
  level 1/.style={{sibling distance={sd1}cm, level distance={ld1}cm}},
  level 2/.style={{sibling distance={sd2}cm, level distance={ld2}cm}}
]
\tikzset{{
  noeud/.style   = {{{nodestyle}}},
  prob/.style    = {{{probastyle}}}
}}
"""

    corps  = "  \\node " + generer_nodes(
        arbre,
        compteur,
        prob_composee=Fraction(1),
        mode=mode,
        decimales=decimales
    ) + ";"
    if probaend:
        fleches = generer_fleches(compteur[1], flechelength, flechecolor)
    else:
        fleches = ""
    fin     = "\n\\end{tikzpicture}"

    return preambule + corps + "\n" + fleches + fin