model.py

# rl_ascii_console.py
"""
Entrée: image 8x8 (fichier ou générée pour test) -> prédiction caractère ASCII -> feedback humain
Mode: boucle console interactive
Actions:
 - y : bien (reward positif)
 - n : mal (reward négatif ou 0 selon config)
 - c : fournir le caractère correct (supervised update)
 - t : tester avec une image rendue d'un caractère (debug)
 - s : sauvegarder modèle
 - q : quitter
Dépendances: pip install torch pillow numpy
"""

import os
import random
import numpy as np
from PIL import Image, ImageDraw, ImageFont

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical

# -------- Configuration --------
ASCII_START = 32            # premier ASCII considéré
ASCII_END = 127             # exclusif (on prend 32..126)
N_ACTIONS = ASCII_END - ASCII_START
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
CHECKPOINT = "policy_console.pth"

# Choix de reward pour "mal"
MALUS = -1.0    # mettre 0.0 si tu préfères ne rien appliquer
REWARD_OK = 1.0

# -------- Modèle --------
class PolicyNet(nn.Module):
    """Tiny CNN (B,1,8,8) -> logits (B, N_ACTIONS)"""
    def __init__(self, n_actions=N_ACTIONS):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(16)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(32)
        self.pool = nn.MaxPool2d(2)  # 8x8 -> 4x4
        self.fc1 = nn.Linear(32 * 4 * 4, 128)
        self.fc2 = nn.Linear(128, n_actions)

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = self.pool(F.relu(self.bn2(self.conv2(x))))
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        return self.fc2(x)

# -------- Utilitaires --------
def state_to_tensor(state_np):
    """numpy (8,8) -> torch tensor (1,1,8,8) sur DEVICE, normalisé 0..1"""
    arr = np.array(state_np, dtype=np.float32)
    if arr.max() > 2.0:
        arr = arr / 255.0
    arr = np.clip(arr, 0.0, 1.0)
    t = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0)  # (1,1,8,8)
    return t.to(DEVICE)

def idx_to_char(idx):
    return chr(ASCII_START + int(idx))

def char_to_idx(ch):
    code = ord(ch)
    if code < ASCII_START or code >= ASCII_END:
        raise ValueError("Caractère hors plage supportée")
    return code - ASCII_START

def tensor_logits_to_probs(logits):
    return F.softmax(logits, dim=-1).squeeze(0).cpu().detach().numpy()

# console print 8x8
_PALETTE = "@%#*+=-:. "  # sombre -> clair
def print_8x8_console(arr):
    a = np.array(arr, dtype=np.float32)
    if a.max() > 2.0:
        a = a / 255.0
    a = np.clip(a, 0.0, 1.0)
    pal = _PALETTE
    n = len(pal)
    for row in a:
        line = "".join(pal[int((1.0 - v) * (n - 1))] for v in row)
        print(line)

# load image file and force to 8x8 grayscale numpy
def load_image(path):
    im = Image.open(path).convert("L")
    return np.array(im, dtype=np.uint8)

# quick renderer for debug: draw char, resize to 8x8
def render_char_8x8(code, jitter=True):
    img = Image.new("L", (16,16), color=255)
    draw = ImageDraw.Draw(img)
    try:
        font = ImageFont.truetype("DejaVuSans.ttf", 20)
    except Exception:
        font = ImageFont.load_default()
    ch = chr(code)
    # taille texte compatible PIL
    try:
        bbox = draw.textbbox((0,0), ch, font=font)
        w = bbox[2] - bbox[0]; h = bbox[3] - bbox[1]
    except Exception:
        try:
            w,h = font.getsize(ch)
        except Exception:
            w,h = 10,10
    draw.text(((16-w)/2,(16-h)/2), ch, font=font, fill=0)
    if jitter:
        angle = random.uniform(-12,12)
        img = img.rotate(angle, fillcolor=255)
    img = img.resize((8,8), Image.BILINEAR)
    arr = np.array(img).astype(np.uint8)
    if jitter and random.random() < 0.25:
        arr = np.clip(arr + np.random.randint(-15,15,(8,8)), 0, 255)
    return arr

# -------- Prédiction et mises à jour --------
def predict_sample(model, state_np):
    """échantillonne selon la policy, renvoie idx, char, log_prob, probs_numpy"""
    model.eval()
    with torch.no_grad():
        logits = model(state_to_tensor(state_np))
    probs = F.softmax(logits, dim=-1).squeeze(0)
    dist = Categorical(probs)
    action = dist.sample()
    log_prob = dist.log_prob(action)
    idx = int(action.item())
    return idx, idx_to_char(idx), log_prob, probs.cpu().detach().numpy()

def predict_greedy(model, state_np):
    model.eval()
    with torch.no_grad():
        logits = model(state_to_tensor(state_np))
    probs = tensor_logits_to_probs(logits)
    idx = int(np.argmax(probs))
    return idx, idx_to_char(idx), probs

def reinforce_update_from_logprob(model, optimizer, log_prob, reward, baseline=0.0):
    """REINFORCE single-step: loss = -log_prob * (reward - baseline)"""
    model.train()
    advantage = torch.tensor(float(reward - baseline), device=DEVICE)
    loss = - log_prob * advantage
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()

def supervised_update(model, optimizer, state_np, true_idx):
    model.train()
    x = state_to_tensor(state_np)
    logits = model(x)
    target = torch.tensor([true_idx], dtype=torch.long, device=DEVICE)
    loss = F.cross_entropy(logits, target)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()

def image_to_blocks(
    arr: np.ndarray,
    block: int = 8,
    resize_mode: str = "bilinear"
):
    """
    Découpe arr (HxW ou HxWxC) en blocs non-chevauchants de taille `block x block`.
    Retour:
      - blocks_grid: numpy array shape (n_rows, n_cols, block, block) pour niveau de gris
                     ou (n_rows, n_cols, block, block, C) pour couleur
      - info: dict { 'method': 'none'|'resized'|'cropped',
                     'orig_size': (H,W),
                     'new_size': (H2,W2),
                     'n_rows': n_rows, 'n_cols': n_cols, 'block': block }
    """
    # validation
    if arr.ndim not in (2, 3):
        raise ValueError("arr doit être 2D (HxW) ou 3D (HxWxC)")

    # interpolation PIL
    interp_map = {"nearest": Image.NEAREST, "bilinear": Image.BILINEAR, "bicubic": Image.BICUBIC}
    interp = interp_map.get(resize_mode, Image.BILINEAR)

    orig_h, orig_w = arr.shape[:2]

    # nearest multiple (arrondi au plus proche), au moins block
    def nearest_multiple(x, block):
        return max(block, int(round(x / block)) * block)

    target_h = nearest_multiple(orig_h, block)
    target_w = nearest_multiple(orig_w, block)

    # transformation si nécessaire
    if orig_h % block == 0 and orig_w % block == 0:
        arr_t = arr.copy()
        method = "none"
    else:
        try:
            need_restore_float = False
            if np.issubdtype(arr.dtype, np.floating):
                pil_in = Image.fromarray((np.clip(arr, 0.0, 1.0) * 255).astype(np.uint8))
                need_restore_float = True
            else:
                pil_in = Image.fromarray(arr.astype(np.uint8))
            pil_resized = pil_in.resize((target_w, target_h), resample=interp)
            arr_t = np.array(pil_resized)
            if need_restore_float:
                arr_t = arr_t.astype(np.float32) / 255.0
            method = "resized"
            if arr_t.shape[0] < block or arr_t.shape[1] < block:
                raise RuntimeError("Taille trop petite après resize")
        except Exception:
            crop_h = (orig_h // block) * block or block
            crop_w = (orig_w // block) * block or block
            top = max(0, (orig_h - crop_h) // 2)
            left = max(0, (orig_w - crop_w) // 2)
            arr_t = arr[top:top+crop_h, left:left+crop_w].copy()
            method = "cropped"

    new_h, new_w = arr_t.shape[:2]
    n_rows = new_h // block
    n_cols = new_w // block

    # construire grille 2D
    if arr_t.ndim == 2:
        blocks_grid = np.zeros((n_rows, n_cols, block, block), dtype=arr_t.dtype)
        for i in range(n_rows):
            for j in range(n_cols):
                y = i * block
                x = j * block
                blocks_grid[i, j] = arr_t[y:y+block, x:x+block]
    else:
        C = arr_t.shape[2]
        blocks_grid = np.zeros((n_rows, n_cols, block, block, C), dtype=arr_t.dtype)
        for i in range(n_rows):
            for j in range(n_cols):
                y = i * block
                x = j * block
                blocks_grid[i, j] = arr_t[y:y+block, x:x+block, :]

    info = {
        "method": method,
        "orig_size": (orig_h, orig_w),
        "new_size": (new_h, new_w),
        "n_rows": n_rows,
        "n_cols": n_cols,
        "block": block
    }
    return blocks_grid, info

def _blocks_grid_to_batch_tensor(blocks_grid: np.ndarray, device: torch.device):
    """
    Convertit blocks_grid (n_rows, n_cols, block, block) ou (n_rows,n_cols,block,block,C)
    en tenseur torch shape (N, C, block, block) prêt pour le modèle et indique si c'est grayscale.
    Normalise en float32 0..1.
    """
    # détecter canaux
    if blocks_grid.ndim == 4:  # (n_rows, n_cols, block, block) -> grayscale
        n_rows, n_cols, b1, b2 = blocks_grid.shape
        N = n_rows * n_cols
        arr = blocks_grid.reshape(N, b1, b2)               # (N,block,block)
        arr = arr.astype(np.float32) / 255.0
        # ajouter channel dim -> (N,1,block,block)
        t = torch.from_numpy(arr).unsqueeze(1)            # (N,1,block,block)
        is_gray = True
    else:  # color (n_rows,n_cols,block,block,C)
        n_rows, n_cols, b1, b2, C = blocks_grid.shape
        N = n_rows * n_cols
        arr = blocks_grid.reshape(N, b1, b2, C)           # (N,block,block,C)
        # convertir en float32 0..1
        arr = arr.astype(np.float32) / 255.0
        # transpose -> (N,C,block,block)
        t = torch.from_numpy(arr).permute(0, 3, 1, 2)
        is_gray = False
    return t.to(device), is_gray

def create_char_table_from_image(
    model: torch.nn.Module,
    image_arr: np.ndarray,
    block: int = 8,
    device: torch.device = None,
    mode: str = "greedy"
):
    """
    Crée une table 2D de caractères prédits pour chaque bloc de l'image.
    Args:
      - model: réseau (doit renvoyer logits shape (B, n_actions))
      - image_arr: numpy array HxW (grayscale) ou HxWxC (color), dtype uint8 ou float
      - block: taille du bloc (ex. 8)
      - device: torch.device (si None, utilise model.device ou cpu)
      - mode: "greedy" ou "sample"
    Returns:
      - grid_chars: liste de listes (n_rows x n_cols) de caractères (str)
      - info: dict retourné par image_to_blocks (method, sizes, n_rows, n_cols...)
    """
    # choisir device
    if device is None:
        try:
            device = next(model.parameters()).device
        except StopIteration:
            device = torch.device("cpu")

    # découper en grille 2D (utilise ta fonction image_to_blocks définie ailleurs)
    blocks_grid, info = image_to_blocks(image_arr, block=block)  # blocks_grid shape (n_rows,n_cols,8,8) ou (...,C)
    n_rows = info["n_rows"]
    n_cols = info["n_cols"]

    # convertir tous les blocs en un batch tensor (N, C, block, block)
    batch_t, is_gray = _blocks_grid_to_batch_tensor(blocks_grid, device)

    model.eval()
    with torch.no_grad():
        logits = model(batch_t)                      # (N, n_actions)
        probs = F.softmax(logits, dim=-1)            # (N, n_actions)

        if mode == "greedy":
            actions = torch.argmax(probs, dim=-1)   # (N,)
            # actions est tensor CPU pour transformation
            actions = actions.cpu().numpy()
        elif mode == "sample":
            # échantillonne indépendamment pour chaque exemple
            actions = []
            for p in probs:
                dist = Categorical(p)
                a = int(dist.sample().cpu().item())
                actions.append(a)
            actions = np.array(actions, dtype=np.int32)
        else:
            raise ValueError("mode doit être 'greedy' ou 'sample'")

    # convertir indices en caractères
    # idx_to_char doit être définie dans ton code (ex: ascii_start = 32)
    chars = [idx_to_char(int(a)) for a in actions]   # liste longueur N

    # reconstruire grille 2D (n_rows x n_cols)
    grid_chars = []
    idx = 0
    for i in range(n_rows):
        row = []
        for j in range(n_cols):
            row.append(chars[idx])
            idx += 1
        grid_chars.append(row)

    return grid_chars, info



# -------- Checkpoint --------
def save_checkpoint(model, path=CHECKPOINT):
    torch.save(model.state_dict(), path)

def load_checkpoint(model, path=CHECKPOINT):
    if os.path.exists(path):
        model.load_state_dict(torch.load(path, map_location=DEVICE))
        return True
    return False

# -------- Boucle interactive --------

def print_char_grid(grid_chars, sep: str = "") -> None:
    """
    Affiche une grille 2D de caractères ligne par ligne.
    - grid_chars: list[list[str]] ou numpy array 2D de caractères
    - sep: chaîne insérée entre les caractères d'une même ligne (ex: " " pour espace)
    Exemple:
        print_char_grid(grid_chars)        # affichage collé
        print_char_grid(grid_chars, sep=" ")  # affichage avec espace entre les caractères
    """
    # Convertir en liste de listes si ndarray
    if isinstance(grid_chars, np.ndarray):
        grid = grid_chars.tolist()
    else:
        grid = grid_chars

    for row in grid:
        # remplacer None ou valeurs vides par espace
        out_chars = [(ch if (ch is not None and ch != "") else " ") for ch in row]
        print(sep.join(out_chars))

def interactive(model, optimizer):
    """
    Interactive loop where user assigns rewards based on the final 2D char table.
    Options to assign rewards:
      a) same feedback for all (y/n)
      b) list good coordinates, others get malus (format: r,c; r,c; ...)
      c) provide a grid of y/n per row (format: 'y n; n y')
    """
    print("Mode interactif: donne une image (ou 't' test, 'q' pour quitter).")
    print("Après prédiction: tu choisis une carte de récompenses pour la grille entière.")
    while True:
        cmd = input("\nEntrer chemin image (ou 't' test, 'q' quit, 's' save): ").strip()
        if cmd == "q":
            print("Quit.")
            break
        if cmd == "s":
            save_checkpoint(model)
            print("Sauvegardé.")
            continue

        if cmd == "t":
            code = random.randint(65, 90)
            img = render_char_8x8(code, jitter=True)
            print(f"(image test générée du caractère cible {chr(code)})")
        else:
            if not os.path.isfile(cmd):
                print("Fichier introuvable. Réessayez ou tapez 't' pour test.")
                continue
            img = load_image(cmd)  # ou load_image(cmd) selon ton nom

        # découper en grille 2D
        blocks_grid, info = image_to_blocks(img, block=8)
        n_rows = info["n_rows"]; n_cols = info["n_cols"]
        print(f"Image traitée: méthode={info['method']}  orig={info['orig_size']} -> new={info['new_size']}")
        print(f"Grille: {n_rows} x {n_cols}  (total blocs {n_rows*n_cols})")

        # convertir en batch tensor et prédire par batch (sample pour RL)
        batch_t, _ = _blocks_grid_to_batch_tensor(blocks_grid, device=next(model.parameters()).device)
        model.eval()
        with torch.no_grad():
            logits = model(batch_t)                   # (N, A)
            probs = F.softmax(logits, dim=-1)         # (N, A)
            dist = Categorical(probs)
            actions = dist.sample()                   # tensor (N,) on device
            log_probs = dist.log_prob(actions)        # tensor (N,) on device
            actions_np = actions.cpu().numpy().astype(int)

        # construire la grille de caractères
        chars = [idx_to_char(a) for a in actions_np]
        grid_chars = []
        idx = 0
        for i in range(n_rows):
            row = []
            for j in range(n_cols):
                row.append(chars[idx])
                idx += 1
            grid_chars.append(row)

        # afficher la grille
        print_char_grid(grid_chars)

        # optionnel : afficher aperçu des blocs
        if input("Afficher aperçu blocs ? (y/N): ").strip().lower() == "y":
            k = 0
            for i in range(n_rows):
                for j in range(n_cols):
                    print(f"\nBloc [{i},{j}] (char='{grid_chars[i][j]}'):")
                    blk = blocks_grid[i, j]
                    print_8x8_console(blk if blk.ndim==2 else blk.mean(axis=2))
                    top5 = np.argsort(probs[k].cpu().numpy())[-5:][::-1]
                    print("Top-5:", [(idx_to_char(t), float(probs[k, t].cpu().numpy())) for t in top5])
                    k += 1

        # choix du mode d'attribution des récompenses
        print("\nChoisir mode d'attribution des récompenses:")
        print(" a) même feedback pour tous (y=ok / n=mal)")
        print(" b) lister coordonnées bonnes, ex: 0,1;1,2 (les autres reçoivent malus)")
        print(" c) saisir grille y/n ligne par ligne, ex: 'y n n; y y n'")

        mode = input("Mode (a/b/c) ou 'q' pour annuler: ").strip().lower()
        if mode == "q":
            continue

        N = n_rows * n_cols
        rewards = np.zeros(N, dtype=float)
        supervised_corrections = {}  # optional corrections: idx -> char

        if mode == "a":
            fb = input("Entrer feedback pour tous (y/n) ou 'c' pour supervised: ").strip().lower()
            if fb == "y":
                rewards[:] = REWARD_OK
            elif fb == "n":
                rewards[:] = MALUS
            elif fb == "c":
                corr = input("Entrer caractère correct pour TOUS: ").strip()
                if corr:
                    try:
                        true_idx = char_to_idx(corr[0])
                        # supervised update for all blocks
                        for k in range(N):
                            r = k // n_cols; c = k % n_cols
                            blk_np = blocks_grid[r, c]
                            loss = supervised_update(model, optimizer, blk_np, true_idx)
                        print("Supervised update appliquée à tous les blocs.")
                    except Exception as e:
                        print("Erreur:", e)
                else:
                    print("Aucun caractère fourni.")
            else:
                print("Entrée inconnue, aucune récompense.")

        elif mode == "b":
            s = input("Coordonnées bonnes (format r,c;... ), ex: 0,1;1,2  (vide = none): ").strip()
            good = set()
            if s:
                for p in [p.strip() for p in s.split(";") if p.strip()]:
                    try:
                        r, c = [int(x.strip()) for x in p.split(",")]
                        if 0 <= r < n_rows and 0 <= c < n_cols:
                            good.add((r, c))
                    except Exception:
                        print("Ignoré coord invalide:", p)
            k = 0
            for i in range(n_rows):
                for j in range(n_cols):
                    rewards[k] = REWARD_OK if (i, j) in good else MALUS
                    k += 1

        elif mode == "c":
            s = input("Saisir grille (lignes séparées par ';', ex 'y n; n y'): ").strip()
            if not s:
                print("Aucune grille fournie, aucune récompense.")
            else:
                rows = [row.strip() for row in s.split(";") if row.strip()]
                if len(rows) != n_rows:
                    print(f"Nombre de lignes donné != {n_rows}, aucune récompense appliquée.")
                else:
                    k = 0
                    for i, row in enumerate(rows):
                        items = row.split()
                        for j in range(n_cols):
                            it = items[j] if j < len(items) else "n"
                            if it.lower() == "y":
                                rewards[k] = REWARD_OK
                            else:
                                rewards[k] = MALUS
                            k += 1

        else:
            print("Mode inconnu, aucune récompense appliquée.")

        # appliquer mises à jour REINFORCE pour chaque bloc (log_probs tensor précédemment calculé)
        # log_probs is tensor shape (N,) on device
        if log_probs is None:
            print("Pas de log_probs disponibles (mode prédiction différent). Aucune mise à jour RL appliquée.")
        else:
            for k in range(N):
                r = float(rewards[k])
                if r != 0.0:
                    lp = log_probs[k]
                    loss = reinforce_update_from_logprob(model, optimizer, lp, r, baseline=0.0)
                    i = k // n_cols; j = k % n_cols
                    print(f"Applied RL update [{i},{j}] reward={r} loss={loss:.6f}")

        # afficher la grille finale (après corrections éventuelles)
        print("\nGrille finale (après corrections si fournies):")
        for row in grid_chars:
            print("".join(row))

        # sauvegarde optionnelle
        if input("\nSauvegarder modèle ? (y/N): ").strip().lower() == "y":
            save_checkpoint(model)
            print("Sauvegardé.")

        # continuer ou quitter
        if input("Nouvelle image ? (y/N): ").strip().lower() != "y":
            print("Quit.")
            break


# -------- Main --------
def main():
    model = PolicyNet(n_actions=N_ACTIONS).to(DEVICE)
    opt = optim.Adam(model.parameters(), lr=5e-4, weight_decay=1e-6)
    if load_checkpoint(model):
        print("Checkpoint chargé:", CHECKPOINT)
    else:
        print("Pas de checkpoint trouvé. Nouveau modèle.")
    try:
        interactive(model, opt)
    finally:
        save_checkpoint(model)
        print("Checkpoint sauvegardé à la sortie.")

if __name__ == "__main__":
    main()