epicthefamily/tree.py

import pygraphviz as pgv
import json
import os
import hashlib
from PIL import Image, ImageDraw, ImageFont, ImageOps

def get_file_hash(filename: str) -> str:
    """Calcule le hash SHA-256 d'un fichier"""
    hash_sha256 = hashlib.sha256()
    with open(filename, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_sha256.update(chunk)
    return hash_sha256.hexdigest()

def create_avatar_with_name(user_id, name, avatar_path, output_path):
    """Crée une image combinant avatar rond et pseudo en dessous"""
    try:
        avatar = Image.open(avatar_path).convert("RGBA")
        avatar = avatar.resize((150, 150))
        
        mask = Image.new('L', (150, 150), 0)
        draw = ImageDraw.Draw(mask)
        draw.ellipse((0, 0, 150, 150), fill=255)
        rounded_avatar = ImageOps.fit(avatar, mask.size, centering=(0.5, 0.5))
        rounded_avatar.putalpha(mask)

        combined = Image.new('RGBA', (200, 200), (0, 0, 0, 0))
        combined.paste(rounded_avatar, (25, 0), rounded_avatar)

        draw = ImageDraw.Draw(combined)
        try:
            font = ImageFont.truetype("arial.ttf", 14)
        except:
            font = ImageFont.load_default()

        text_width = draw.textlength(name, font=font)
        draw.text(((200 - text_width) / 2, 160), name, font=font, fill="white")

        combined.save(output_path)
        return True
    except Exception as e:
        print(f"Erreur création avatar+texte {user_id}: {e}")
        return False

def generate_tree(input_file, output_file, root_id=None):
    with open(input_file) as f:
        data = json.load(f)

    G = pgv.AGraph(
        directed=True,
        rankdir="TB",
        nodesep="0.8",
        ranksep="1.5",
        bgcolor="#000000",
        splines="true"
    )

    node_style = {
        "shape": "none",
        "imagescale": "false",
        "labelloc": "b",
        "width": "2",
        "height": "2",
        "fixedsize": "true"
    }

    edge_style = {
        "color": "#FFFFFF",
        "penwidth": "3",
        "arrowsize": "1.2"
    }

    couple_style = {
        "style": "dashed",
        "color": "#FF69B4",
        "penwidth": "2.5",
        "dir": "none"
    }

    os.makedirs("temp_avatars", exist_ok=True)

    # 1. Création de tous les nœuds
    for user_id, info in data["members"].items():
        original_avatar = f"avatars/{user_id}.png"
        combined_avatar = f"temp_avatars/combined_{user_id}.png"

        if os.path.exists(original_avatar):
            create_avatar_with_name(user_id, info["name"], original_avatar, combined_avatar)
            G.add_node(user_id, image=combined_avatar, **node_style)
        else:
            G.add_node(user_id, label=info["name"], **{**node_style, "fontcolor": "white"})

    # 2. Calcul des générations avec la nouvelle logique
    generations = {}
    
    # a. Identifier les racines
    if root_id:
        root_members = [root_id]
    else:
        root_members = data.get("roots", [k for k,v in data["members"].items() if not v.get("parents")])
    
    # b. Première passe: membres sans parents = gen 0
    for user_id in data["members"]:
        if not data["members"][user_id].get("parents"):
            generations[user_id] = 0
    
    # c. Deuxième passe: aligner les couples
    for couple in data.get("couples", []):
        m1, m2 = couple
        if m1 in generations and m2 in generations:
            gen = min(generations[m1], generations[m2])
            generations[m1] = gen
            generations[m2] = gen
    
    # d. Troisième passe: calcul descendant
    changed = True
    while changed:
        changed = False
        for user_id in data["members"]:
            if user_id in generations:
                continue
                
            parents = data["members"][user_id].get("parents", [])
            if parents:
                parent_gens = [generations[p["id"]] for p in parents if p["id"] in generations]
                if parent_gens:
                    new_gen = min(parent_gens) + 1
                    generations[user_id] = new_gen
                    changed = True
                    
                    # Aligner avec partenaire
                    partner_id = next((c[0] if c[1]==user_id else c[1] for c in data.get("couples",[]) if user_id in c), None)
                    if partner_id and partner_id in data["members"]:
                        generations[partner_id] = new_gen

    # 3. Organisation des niveaux
    max_gen = max(generations.values()) if generations else 0
    
    # Niveau 0: Uniquement les racines
    roots = [m for m in root_members if m in data["members"]]
    if roots:
        G.add_subgraph(roots, name="rank0", rank="same")

    # Niveaux suivants
    for gen in range(1, max_gen + 1):
        members = [m for m,g in generations.items() if g == gen]
        
        # Séparer les couples
        couples = []
        for couple in data.get("couples", []):
            if all(m in members for m in couple):
                couples.append(couple)
                members = [m for m in members if m not in couple]
        
        # Créer les sous-graphes
        if members:
            G.add_subgraph(members, name=f"rank{gen}", rank="same")
        
        for couple in couples:
            G.add_subgraph(couple, name=f"couple_{couple[0]}_{couple[1]}", rank="same")
            G.add_edge(couple[0], couple[1], **couple_style)

    # 4. Liens parent-enfant
    for user_id, info in data["members"].items():
        for parent in info.get("parents", []):
            if parent["id"] in data["members"]:
                G.add_edge(parent["id"], user_id, **edge_style)

    G.layout(prog="dot")
    G.draw(output_file)

    # Nettoyage
    for f in os.listdir("temp_avatars"):
        os.remove(f"temp_avatars/{f}")