Blame for maze_model.py

"""
maze_model.py — semua yang dibutuhkan untuk melatih & menjalankan model penyelesai maze.

Format maze:  1 = dinding, 0 = jalan, 2 = start, 3 = exit (list of lists / numpy array).

Ide model (DeepThinking-style recurrent CNN):
  - Model membaca seluruh peta, lalu "berpikir" berulang kali (iterasi) memakai
    blok konvolusi yang SAMA. Setiap iterasi, informasi "arah ke exit" menyebar
    beberapa kotak lebih jauh dari exit.
  - Output: untuk setiap kotak jalan, arah terbaik (atas/bawah/kiri/kanan) menuju exit.
  - Solver lalu mulai dari start dan mengikuti arah-arah itu sampai exit.
  - Karena modelnya fully-convolutional dan jumlah iterasi bisa ditambah saat dipakai,
    model bisa dipakai untuk peta ukuran berapa pun (makin besar → makin banyak iterasi).
"""
from __future__ import annotations

import json
import random
from collections import deque

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

WALL, OPEN, START, EXIT = 1, 0, 2, 3
DIRS = [(-1, 0), (1, 0), (0, -1), (0, 1)]  # atas, bawah, kiri, kanan
DIR_NAMES = ["U", "D", "L", "R"]


# --------------------------------------------------------------------------- #
#  1. Generator maze acak (data latih tak terbatas)
# --------------------------------------------------------------------------- #
def _carve_perfect_maze(h: int, w: int, rng: random.Random) -> np.ndarray:
    """Recursive backtracker pada grid ganjil h x w (h, w ganjil)."""
    g = np.ones((h, w), dtype=np.int64)
    sr, sc = rng.randrange(1, h, 2), rng.randrange(1, w, 2)
    g[sr, sc] = OPEN
    stack = [(sr, sc)]
    while stack:
        r, c = stack[-1]
        nbrs = [(r + 2 * dr, c + 2 * dc, dr, dc) for dr, dc in DIRS
                if 0 < r + 2 * dr < h - 1 and 0 < c + 2 * dc < w - 1
                and g[r + 2 * dr, c + 2 * dc] == WALL]
        if not nbrs:
            stack.pop()
            continue
        nr, nc, dr, dc = rng.choice(nbrs)
        g[r + dr, c + dc] = OPEN
        g[nr, nc] = OPEN
        stack.append((nr, nc))
    return g


def bfs_dist(maze: np.ndarray, src: tuple[int, int]) -> np.ndarray:
    """Jarak BFS dari src ke semua kotak non-dinding (-1 = tidak terjangkau)."""
    h, w = maze.shape
    dist = np.full((h, w), -1, dtype=np.int64)
    dist[src] = 0
    q = deque([src])
    while q:
        r, c = q.popleft()
        for dr, dc in DIRS:
            nr, nc = r + dr, c + dc
            if 0 <= nr < h and 0 <= nc < w and maze[nr, nc] != WALL and dist[nr, nc] < 0:
                dist[nr, nc] = dist[r, c] + 1
                q.append((nr, nc))
    return dist


def random_maze(h: int, w: int, rng: random.Random | None = None,
                loop_prob: float | None = None, noise_prob: float | None = None,
                min_path: int = 4) -> np.ndarray:
    """Buat maze acak berukuran tepat h x w (boleh genap/ganjil), dengan 2=start & 3=exit.

    Variasi supaya model tahan banting terhadap berbagai gaya peta:
      - loop_prob : sebagian dinding antar-lorong dibuka → ada jalan memutar/loop
      - noise_prob: sebagian dinding sembarang dibuka (seperti contoh peta user)
      - start/exit diletakkan di kotak jalan acak; kadang exit ada di bingkai luar
    """
    rng = rng or random.Random()
    loop_prob = rng.uniform(0.0, 0.12) if loop_prob is None else loop_prob
    noise_prob = rng.uniform(0.0, 0.04) if noise_prob is None else noise_prob
    for _ in range(50):
        gh, gw = h + (1 - h % 2), w + (1 - w % 2)  # ukuran ganjil ≥ target
        g = _carve_perfect_maze(gh, gw, rng)[:h, :w].copy()
        g[0, :] = g[-1, :] = g[:, 0] = g[:, -1] = WALL
        # buka dinding untuk membuat loop / noise (hanya bagian dalam)
        for r in range(1, h - 1):
            for c in range(1, w - 1):
                if g[r, c] != WALL:
                    continue
                horiz = g[r, c - 1] != WALL and g[r, c + 1] != WALL
                vert = g[r - 1, c] != WALL and g[r + 1, c] != WALL
                if (horiz or vert) and rng.random() < loop_prob:
                    g[r, c] = OPEN
                elif rng.random() < noise_prob:
                    g[r, c] = OPEN
        open_cells = list(zip(*np.nonzero(g == OPEN)))
        if len(open_cells) < 2:
            continue
        start = rng.choice(open_cells)
        # kadang exit di bingkai luar (pintu keluar di tepi)
        exit_ = None
        if rng.random() < 0.25:
            border = [(r, c) for r in range(h) for c in range(w)
                      if (r in (0, h - 1) or c in (0, w - 1))
                      and not ((r in (0, h - 1)) and (c in (0, w - 1)))
                      and any(0 <= r + dr < h and 0 <= c + dc < w and g[r + dr, c + dc] == OPEN
                              for dr, dc in DIRS)]
            if border:
                exit_ = rng.choice(border)
        dist = bfs_dist(g, start)
        if exit_ is None:
            reach = [p for p in open_cells if dist[p] >= min_path]
            if not reach:
                continue
            # condong memilih exit yang jauh (lebih menantang)
            reach.sort(key=lambda p: dist[p])
            exit_ = reach[int(len(reach) * rng.uniform(0.3, 1.0)) - 1]
        g[exit_] = EXIT
        g[start] = START
        if bfs_dist(g, start)[exit_] < 1:
            continue
        return g
    raise RuntimeError("gagal membuat maze")


# --------------------------------------------------------------------------- #
#  2. Encoding input & label
# --------------------------------------------------------------------------- #
def encode(maze) -> torch.Tensor:
    """maze (H,W) → tensor (3,H,W): [dinding, jalan, exit]."""
    m = np.asarray(maze)
    wall = (m == WALL)
    exit_ = (m == EXIT)
    return torch.from_numpy(np.stack([wall, ~wall, exit_]).astype(np.float32))


def direction_targets(maze: np.ndarray) -> np.ndarray:
    """Label (4,H,W): 1 jika bergerak ke arah itu memperpendek jarak ke exit.
    Kotak dinding / exit / tak terjangkau → semua 0 (diabaikan saat loss)."""
    h, w = maze.shape
    ex = tuple(np.argwhere(maze == EXIT)[0])
    dist = bfs_dist(maze, ex)
    tgt = np.zeros((4, h, w), dtype=np.float32)
    for k, (dr, dc) in enumerate(DIRS):
        nd = np.full((h, w), -1, dtype=np.int64)
        rs, re_ = max(0, -dr), h - max(0, dr)
        cs, ce = max(0, -dc), w - max(0, dc)
        nd[rs:re_, cs:ce] = dist[rs + dr:re_ + dr, cs + dc:ce + dc]
        tgt[k] = ((dist > 0) & (nd >= 0) & (nd == dist - 1)).astype(np.float32)
    return tgt


def make_batch(batch_size: int, h: int, w: int, rng: random.Random, device="cpu"):
    mazes = [random_maze(h, w, rng) for _ in range(batch_size)]
    x = torch.stack([encode(m) for m in mazes]).to(device)
    y = torch.from_numpy(np.stack([direction_targets(m) for m in mazes])).to(device)
    return x, y, mazes


# --------------------------------------------------------------------------- #
#  3. Model
# --------------------------------------------------------------------------- #
class ResBlock(nn.Module):
    def __init__(self, ch: int):
        super().__init__()
        self.c1 = nn.Conv2d(ch, ch, 3, padding=1, bias=False)
        self.c2 = nn.Conv2d(ch, ch, 3, padding=1, bias=False)

    def forward(self, x):
        return F.relu(x + self.c2(F.relu(self.c1(x))))


class MazeNet(nn.Module):
    """Recurrent CNN: blok yang sama dipakai berulang (weight-tied).
    'Input recall': peta asli selalu disuntikkan lagi tiap iterasi agar model tidak lupa."""

    def __init__(self, width: int = 64, n_blocks: int = 2, in_ch: int = 3):
        super().__init__()
        self.config = {"width": width, "n_blocks": n_blocks, "in_ch": in_ch}
        self.proj = nn.Conv2d(in_ch, width, 3, padding=1, bias=False)
        self.recall = nn.Conv2d(width + in_ch, width, 3, padding=1, bias=False)
        self.blocks = nn.Sequential(*[ResBlock(width) for _ in range(n_blocks)])
        self.head = nn.Sequential(
            nn.Conv2d(width, 32, 3, padding=1, bias=False), nn.ReLU(),
            nn.Conv2d(32, 4, 3, padding=1, bias=False),
        )

    def init_state(self, x):
        return F.relu(self.proj(x))

    def step(self, x, h):
        return self.blocks(F.relu(self.recall(torch.cat([h, x], 1))))

    def forward(self, x, iters: int, h=None):
        if h is None:
            h = self.init_state(x)
        for _ in range(iters):
            h = self.step(x, h)
        return self.head(h), h


def masked_direction_loss(logits, target):
    """-log(peluang memilih salah satu arah yang benar), hanya di kotak berlabel.
    Arah ke dinding di-mask agar model fokus memilih di antara langkah yang mungkin."""
    valid = target.sum(1) > 0                                # (B,H,W)
    logp = F.log_softmax(logits, dim=1)
    good = torch.logsumexp(logp.masked_fill(target == 0, -1e4), dim=1)
    return -(good[valid]).mean()


def direction_accuracy(logits, target):
    valid = target.sum(1) > 0
    pred = logits.argmax(1, keepdim=True)
    hit = target.gather(1, pred).squeeze(1) > 0
    return hit[valid].float().mean().item()


# --------------------------------------------------------------------------- #
#  4. Solver: model "berpikir" lalu jalan dari start mengikuti arah
# --------------------------------------------------------------------------- #
def follow_directions(maze: np.ndarray, logits: np.ndarray):
    """Ikuti arah terbaik (yang bukan dinding) dari start. Return path atau None."""
    h, w = maze.shape
    start = tuple(int(v) for v in np.argwhere(maze == START)[0])
    pos, path, seen = start, [start], {start}
    for _ in range(h * w):
        if maze[pos] == EXIT:
            return path
        r, c = pos
        best, best_k = None, -np.inf
        for k, (dr, dc) in enumerate(DIRS):
            nr, nc = r + dr, c + dc
            if 0 <= nr < h and 0 <= nc < w and maze[nr, nc] != WALL and logits[k, r, c] > best_k:
                best, best_k = (nr, nc), logits[k, r, c]
        if best is None or best in seen:
            return None  # buntu / berputar-putar
        pos = best
        path.append(pos)
        seen.add(pos)
    return None


def validate_maze(maze) -> np.ndarray:
    m = np.asarray(maze, dtype=np.int64)
    if m.ndim != 2 or m.shape[0] < 3 or m.shape[1] < 3:
        raise ValueError("maze harus list 2D minimal 3x3")
    if not np.isin(m, [0, 1, 2, 3]).all():
        raise ValueError("isi maze hanya boleh 0, 1, 2, 3")
    if (m == START).sum() != 1 or (m == EXIT).sum() != 1:
        raise ValueError("harus ada tepat satu start (2) dan satu exit (3)")
    return m


def follow_with_corrections(maze: np.ndarray, logits: np.ndarray, max_corrections: int):
    """Seperti follow_directions, tapi boleh 'koreksi' kecil: bila panah terbaik menunjuk ke
    kotak yang sudah dilewati, ambil arah terbaik berikutnya menurut model; bila buntu,
    mundur satu kotak. Setiap penyimpangan dari panah utama dihitung sebagai 1 koreksi.
    Return (path, jumlah_koreksi) atau (None, jumlah_koreksi)."""
    h, w = maze.shape
    start = tuple(int(v) for v in np.argwhere(maze == START)[0])
    stack, seen, corr = [start], {start}, 0
    while stack:
        pos = stack[-1]
        if maze[pos] == EXIT:
            return list(stack), corr
        r, c = pos
        opts = sorted(((logits[k, r, c], (r + dr, c + dc)) for k, (dr, dc) in enumerate(DIRS)
                       if 0 <= r + dr < h and 0 <= c + dc < w and maze[r + dr, c + dc] != WALL),
                      key=lambda t: -t[0])
        nxt = next(((rank, p) for rank, (_, p) in enumerate(opts) if p not in seen), None)
        if nxt is None:            # buntu → mundur
            stack.pop()
            corr += 1
        else:
            corr += nxt[0] > 0     # tidak mengikuti panah terbaik
            stack.append(nxt[1])
            seen.add(nxt[1])
        if corr > max_corrections:
            return None, corr
    return None, corr


@torch.no_grad()
def solve(model: MazeNet, maze, max_iters: int | None = None, check_every: int = 5,
          allow_corrections: bool = True, max_corrections: int | None = None,
          fallback_bfs: bool = False, device="cpu"):
    """Selesaikan maze. Model berpikir bertahap; setiap `check_every` iterasi dicek
    apakah panah dari start sudah sampai exit. Berhenti begitu berhasil, atau begitu
    'pikiran' model tidak berubah lagi (sudah konvergen).

    Urutan usaha:
      1. "model"                 : murni mengikuti panah model
      2. "model_with_corrections": panah model + sedikit koreksi (bila allow_corrections)
      3. "bfs_fallback"          : algoritma BFS klasik (bila fallback_bfs)
    """
    m = validate_maze(maze)
    h, w = m.shape
    max_iters = max_iters or min(3000, max(60, h * w))
    if max_corrections is None:
        max_corrections = max(10, int((m != WALL).sum()) // 10)
    model.eval()
    x = encode(m).unsqueeze(0).to(device)
    hstate = model.init_state(x)
    open_mask = m != WALL
    settle_limit = max(2, int(0.02 * open_mask.sum()))   # "hampir tidak berubah"
    path, it, corrections, solved_by = None, 0, 0, None
    prev_arrows, calm = None, 0
    while it < max_iters:
        n = min(check_every, max_iters - it)
        for _ in range(n):
            hstate = model.step(x, hstate)
        it += n
        logits = model.head(hstate)[0].float().cpu().numpy()
        path = follow_directions(m, logits)
        if path is not None:
            solved_by = "model"
            break
        # Sudah "selesai berpikir"? Kalau hanya segelintir panah yang masih berubah
        # selama 4 pengecekan berturut-turut, berpikir lebih lama tidak akan membantu.
        arrows = logits.argmax(0)
        if prev_arrows is not None:
            changed = int(((arrows != prev_arrows) & open_mask).sum())
            calm = calm + 1 if changed <= settle_limit else 0
            if calm >= 4:
                break
        prev_arrows = arrows
    if path is None and allow_corrections:
        path, corrections = follow_with_corrections(m, logits, max_corrections)
        solved_by = "model_with_corrections" if path is not None else None
    if path is None and fallback_bfs:
        path = bfs_path(m)
        solved_by = "bfs_fallback" if path is not None else None
    return {
        "solved": path is not None,
        "solved_by": solved_by,
        "iterations": it,
        "corrections": corrections if solved_by == "model_with_corrections" else 0,
        "path": [list(map(int, p)) for p in path] if path else None,
        "moves": "".join(_moves(path)) if path else None,
        "steps": len(path) - 1 if path else None,
    }


def bfs_path(m: np.ndarray):
    """Jalur terpendek klasik (untuk evaluasi & cadangan)."""
    start = tuple(np.argwhere(m == START)[0])
    ex = tuple(np.argwhere(m == EXIT)[0])
    dist = bfs_dist(m, ex)
    if dist[start] < 0:
        return None
    path, pos = [start], start
    while pos != ex:
        r, c = pos
        for dr, dc in DIRS:
            nr, nc = r + dr, c + dc
            if 0 <= nr < m.shape[0] and 0 <= nc < m.shape[1] and dist[nr, nc] == dist[r, c] - 1:
                pos = (nr, nc)
                break
        path.append(pos)
    return path


def _moves(path):
    out = []
    for (r0, c0), (r1, c1) in zip(path, path[1:]):
        out.append(DIR_NAMES[DIRS.index((r1 - r0, c1 - c0))])
    return out


def render(maze, path=None) -> str:
    """Gambar maze di terminal. '#'=dinding, 'S'=start, 'E'=exit, '•'=jalur."""
    m = np.asarray(maze)
    on = set(map(tuple, path or []))
    chars = {WALL: "█", OPEN: " ", START: "S", EXIT: "E"}
    lines = []
    for r in range(m.shape[0]):
        row = ""
        for c in range(m.shape[1]):
            v = int(m[r, c])
            row += "•" if (r, c) in on and v == OPEN else chars[v]
        lines.append(row)
    return "\n".join(lines)


# --------------------------------------------------------------------------- #
#  5. Simpan / muat (format safetensors + config.json, cocok untuk HuggingFace)
# --------------------------------------------------------------------------- #
def save_model(model: MazeNet, folder: str):
    import os
    from safetensors.torch import save_file
    os.makedirs(folder, exist_ok=True)
    save_file({k: v.contiguous().cpu() for k, v in model.state_dict().items()},
              f"{folder}/model.safetensors")
    with open(f"{folder}/config.json", "w") as f:
        json.dump(model.config, f, indent=2)


def load_model(folder: str, device="cpu") -> MazeNet:
    from safetensors.torch import load_file
    with open(f"{folder}/config.json") as f:
        cfg = json.load(f)
    model = MazeNet(**cfg)
    model.load_state_dict(load_file(f"{folder}/model.safetensors", device=str(device)))
    return model.to(device).eval()