#!/usr/bin/env python3
"""Mind2Web: Kev tarzı pointer kafası + LoRA, Gemma 4 E2B, METİN + EKRAN GÖRÜNTÜSÜ.

m2w_train_lora.py'nin görüntülü eşi. Fark yalnızca girdi:
    [BOS][tam sayfa ekran görüntüsü, 1280x1280'lik şeritlere bölünmüş][görev + geçmiş][aday satırları][görev + geçmiş + "Next element:"]
Sorgu, aday anahtarları, kafa, kayıp ve metrikler metin sürümüyle birebir aynıdır; böylece iki sürüm
yalnızca "görüntü var / yok" bakımından ayrılır.

Görüntü kuralları:
  * Hedefin yerine bakılarak KIRPMA YAPILMAZ (sızıntı olur). Sayfa baştan sona, sırayla şeritlere bölünür.
  * Şerit başına token bütçesi --img_tokens (70/140/280/560/1120). 1280x1280 şerit 280'de yaklaşık %63 ölçeğe iner.
  * En fazla --max_strips şerit (varsayılan 5 = 6400 piksel); daha uzun sayfaların altı kesilir (sayısı raporlanır).
  * --blank_image: aynı boyutta düz gri şeritler verir. "Kazanç görüntü bilgisinden mi, yoksa sadece fazladan token'dan mı" kontrolü.

!! Sandbox'ta (torch yok) çalıştırılamadı; Gemma 4 kaynak koduna göre yazıldı. Önce küçük deneme: --max_train 40 --max_val 40.
"""
import argparse
import json
import math
import os
import shutil
import time

import numpy as np
import torch
from torch import nn
from PIL import Image

import m2w_train_lora as T


def strip_bounds(H, strip, max_strips, counter=None):
    """Sayfayı yukarıdan aşağı şeritlere böler. Çok kısa (<200 px) son parça öncekine katılır.
    max_strips'i aşan sayfanın altı kesilir."""
    n = max(1, math.ceil(H / strip))
    if n > 1 and H - (n - 1) * strip < 200:
        n -= 1
    cut = n > max_strips
    if cut:
        n = max_strips
    if counter is not None:
        counter.n_pages += 1
        counter.n_cut += int(cut)
    out = []
    for i in range(n):
        top = i * strip
        bot = H if (i == n - 1 and not cut) else (i + 1) * strip
        out.append((top, min(bot, H)))
    return out


class Sayfa:
    """Ekran görüntüsü -> şeritler -> işlemci -> (input_ids öneki, mm_token_type_ids, pixel_values, image_position_ids)."""

    def __init__(self, proc, tok, shots, strip=1280, max_strips=5, img_tokens=280, blank=False):
        self.proc, self.tok, self.shots = proc, tok, shots
        self.strip, self.max_strips, self.blank = strip, max_strips, blank
        proc.image_processor.max_soft_tokens = img_tokens
        self.n_missing = 0
        self.n_cut = 0
        self.n_pages = 0

    def load_strips(self, uid):
        p = os.path.join(self.shots, uid + ".png")
        if not os.path.exists(p):
            return []
        im = Image.open(p).convert("RGB")
        W, H = im.size
        strips = [im.crop((0, t, W, bt)) for t, bt in strip_bounds(H, self.strip, self.max_strips, self)]
        if self.blank:
            strips = [Image.new("RGB", s.size, (127, 127, 127)) for s in strips]
        return strips

    def prefix(self, uid):
        strips = self.load_strips(uid)
        if not strips:
            self.n_missing += 1
            bos = self.tok.bos_token_id
            ids = [bos] if bos is not None else []
            return ids, [0] * len(ids), None, None
        out = self.proc(images=[strips], text=[" ".join([self.proc.image_token] * len(strips))], return_tensors="pt")
        return (out["input_ids"][0].tolist(), out["mm_token_type_ids"][0].tolist(),
                out["pixel_values"], out["image_position_ids"])


MAX_SEQ_V = 0   # --max_seq: görüntü token'ları DAHİL toplam uzunluk sınırı (0 = sınırsız)


def build_ids_v(tok, ex, order, prefix_ids, max_seq=None):
    max_seq = MAX_SEQ_V if max_seq is None else max_seq
    hist = " ; ".join(h[:80] for h in ex["history"][-T.MAX_HIST:]) or "(none)"
    head = f"\nTask: {ex['task']}\nPrevious actions: {hist}\nCandidates:\n"
    tail = f"Task: {ex['task']}\nPrevious actions: {hist}\nNext element to interact with:"
    ids = list(prefix_ids) + list(tok(head, add_special_tokens=False)["input_ids"])
    tail_ids = list(tok(tail, add_special_tokens=False)["input_ids"])
    toks = [tok("- " + ex["cands"][j]["text"] + "\n", add_special_tokens=False, truncation=True,
                max_length=T.CAND_TOK)["input_ids"] for j in order]
    if max_seq:
        budget = max_seq - len(ids) - len(tail_ids)
        if sum(len(t) for t in toks) > budget:
            lo, hi = 4, T.CAND_TOK
            while lo < hi:
                mid = (lo + hi + 1) // 2
                if sum(min(len(t), mid) for t in toks) <= budget:
                    lo = mid
                else:
                    hi = mid - 1
            toks = [t[:lo] for t in toks]
    spans = []
    for t in toks:
        spans.append((len(ids), len(ids) + len(t)))
        ids += list(t)
    ids += tail_ids
    return ids, spans, len(ids) - 1


def collate_v(tok, sayfa, ex, order, device):
    """Mikro-batch her zaman 1 (şerit sayısı örnekten örneğe değişiyor)."""
    pids, pmm, pv, pos = sayfa.prefix(ex["action_uid"])
    ids, spans, q = build_ids_v(tok, ex, order, pids)
    mm = pmm + [0] * (len(ids) - len(pmm))
    b = {"input_ids": torch.tensor([ids], device=device),
         "attention_mask": torch.ones((1, len(ids)), dtype=torch.long, device=device),
         "mm_token_type_ids": torch.tensor([mm], device=device)}
    if pv is not None:
        b["pixel_values"] = pv.to(device)
        b["image_position_ids"] = pos.to(device)
    return b, [spans], [q], int(sum(mm))


def run_model_v(model, b, layer):
    kw = dict(b, output_hidden_states=True)
    try:
        out = model(logits_to_keep=1, **kw)
    except TypeError:
        out = model(**kw)
    return out.hidden_states[layer]


@torch.no_grad()
def evaluate_v(model, head, tok, sayfa, rows, layer, device, mode="ranker"):
    model.eval(); head.eval()
    res = []
    for ex in rows:
        order = list(range(len(ex["cands"])))
        if mode == "reverse":
            order = order[::-1]
        b, spans, qidx, _ = collate_v(tok, sayfa, ex, order, device)
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device == "cuda"):
            h = run_model_v(model, b, layer)
        logits, opl = head(h, spans, qidx)
        p = torch.softmax(logits[0], 0).cpu().numpy()
        n = len(order)
        j = int(p.argmax())
        pos = set(ex["pos_idx"])
        res.append({"uid": ex["action_uid"], "conf": float(p[j]), "correct": bool(j < n and order[j] in pos),
                    "pred": ex["cands"][order[j]]["bid"] if j < n else "NULL", "null": j == n,
                    "has_pos": len(pos) > 0, "op": ex["op"], "op_ok": int(opl[0].argmax()) == T.OPS.get(ex["op"], 0),
                    "p_null": float(p[n]), "pos_bids": [ex["cands"][k]["bid"] for k in ex["pos_idx"]]})
    return res


def sync(src, dst_root):
    if dst_root:
        try:
            shutil.copytree(src, os.path.join(dst_root, os.path.basename(src.rstrip("/"))), dirs_exist_ok=True)
        except Exception as e:
            print("Drive'a kopyalanamadı:", e)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="m2w_v2")
    ap.add_argument("--shots", default="ekran_goruntuleri")
    ap.add_argument("--out", default="egitim_gorselli")
    ap.add_argument("--layer", type=int, default=30)
    ap.add_argument("--epochs", type=int, default=3)
    ap.add_argument("--accum", type=int, default=8)
    ap.add_argument("--lr_lora", type=float, default=2e-4)
    ap.add_argument("--lr_head", type=float, default=1e-3)
    ap.add_argument("--rank", type=int, default=16)
    ap.add_argument("--max_train", type=int, default=0)
    ap.add_argument("--max_val", type=int, default=0)
    ap.add_argument("--strip", type=int, default=1280)
    ap.add_argument("--max_strips", type=int, default=5)
    ap.add_argument("--img_tokens", type=int, default=280)
    ap.add_argument("--blank_image", action="store_true")
    ap.add_argument("--grad_ckpt", action="store_true")
    ap.add_argument("--sync_dir", default="", help="her epoch sonunda sonuçları buraya (Drive) kopyala")
    ap.add_argument("--eval_splits", nargs="+", default=["val", "test_task", "test_website", "test_domain"])
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--max_cand_tok", type=int, default=T.MAX_CAND_TOK, help="aday satırı başına en çok token")
    ap.add_argument("--max_seq", type=int, default=0, help="adım başına toplam token sınırı (görüntü dahil); aşılırsa aday satırları kısaltılır")
    ap.add_argument("--lora_mlp", action="store_true", help="LoRA'yı MLP projeksiyonlarına da koy")
    ap.add_argument("--always_pos", action="store_true", help="eğitimde süzgecin attığı doğru adayı (dropped_pos) listeye geri ekle")
    ap.add_argument("--resume", action="store_true", help="out klasöründe ckpt.pt varsa son bitmiş epoch'tan devam et")
    ap.add_argument("--no_reverse", action="store_true")
    args = ap.parse_args()
    global MAX_SEQ_V
    T.CAND_TOK = args.max_cand_tok; MAX_SEQ_V = args.max_seq

    from peft import LoraConfig, get_peft_model
    from transformers import AutoModelForMultimodalLM, AutoProcessor, AutoTokenizer

    device = "cuda" if torch.cuda.is_available() else "cpu"
    torch.manual_seed(args.seed); rng = np.random.default_rng(args.seed)
    os.makedirs(args.out, exist_ok=True)

    data = {s: T.read(os.path.join(args.data, f"{s}.jsonl.gz")) for s in ["train"] + args.eval_splits
            if os.path.exists(os.path.join(args.data, f"{s}.jsonl.gz"))}
    train = data["train"]
    if "val" not in data:
        n = max(2, int(0.2 * len(train)))
        data["val"], train = train[-n:], train[:-n]
    if args.max_train:
        train = train[: args.max_train]
    if args.always_pos:
        fixed, added = [], 0
        for e in train:
            dp = e.get("dropped_pos") or []
            if dp:
                e = dict(e); n0 = len(e["cands"])
                e["cands"] = list(e["cands"]) + list(dp)
                e["pos_idx"] = list(e["pos_idx"]) + list(range(n0, n0 + len(dp))); added += 1
            fixed.append(e)
        train = fixed
        print(f"doğruyu ekle: {added}/{len(train)} adıma geri eklendi; hâlâ doğrusu olmayan: {sum(1 for e in train if not e['pos_idx'])}")
    for s in list(data):
        if s != "train" and args.max_val:
            data[s] = data[s][: args.max_val]
    print({s: len(v) for s, v in data.items()}, "| eğitim kullanılan:", len(train))

    tok = AutoTokenizer.from_pretrained(T.MODEL_ID, revision=T.REV)
    try:
        proc = AutoProcessor.from_pretrained(T.MODEL_ID, revision=T.REV)
    except Exception as e:
        print("temel modelde işlemci yok, -it deposundan alınıyor:", type(e).__name__)
        proc = AutoProcessor.from_pretrained("google/gemma-4-E2B-it")
    sayfa = Sayfa(proc, tok, args.shots, args.strip, args.max_strips, args.img_tokens, args.blank_image)

    model = AutoModelForMultimodalLM.from_pretrained(T.MODEL_ID, revision=T.REV, dtype=torch.bfloat16).to(device)
    names = T.lora_names(model, args.layer, args.lora_mlp)
    if not names:
        raise SystemExit("LoRA hedef modül bulunamadı")
    print("LoRA hedef modül sayısı:", len(names), "| örnek:", names[0], names[-1])
    model = get_peft_model(model, LoraConfig(r=args.rank, lora_alpha=2 * args.rank, lora_dropout=0.05,
                                             target_modules=names, bias="none"))
    if args.grad_ckpt:
        model.gradient_checkpointing_enable(); model.enable_input_require_grads()
    model.print_trainable_parameters()
    try:
        d = model.config.get_text_config().hidden_size
    except Exception:
        d = 1536
    head = T.Head(d).to(device)

    lora_params = [p for p in model.parameters() if p.requires_grad]
    opt = torch.optim.AdamW([{"params": lora_params, "lr": args.lr_lora, "weight_decay": 0.0},
                             {"params": head.parameters(), "lr": args.lr_head, "weight_decay": 0.01}])
    total = math.ceil(len(train) / args.accum) * args.epochs
    warm = max(1, int(0.03 * total))
    sched = torch.optim.lr_scheduler.LambdaLR(
        opt, lambda s: (s + 1) / warm if s < warm else 0.5 * (1 + math.cos(math.pi * (s - warm) / max(1, total - warm))))

    best, log = -1.0, []
    t0 = time.time(); step = 0; n_oom = 0; start_ep = 0
    ck = os.path.join(args.out, "ckpt.pt")
    if args.resume and os.path.exists(ck):
        from peft import set_peft_model_state_dict
        from safetensors.torch import load_file
        c = torch.load(ck, map_location="cpu", weights_only=False)
        set_peft_model_state_dict(model, load_file(os.path.join(args.out, "lora_son", "adapter_model.safetensors")))
        head.load_state_dict(torch.load(os.path.join(args.out, "head_son.pt"), map_location=device))
        opt.load_state_dict(c["opt"]); sched.load_state_dict(c["sched"])
        rng.bit_generator.state = c["rng"]; best, log, step, n_oom = c["best"], c["log"], c["step"], c["n_oom"]
        start_ep = c["ep"] + 1
        print(f"DEVAM: epoch {start_ep}'dan devam ediliyor (en iyi val acc şu ana kadar {best:.3f})", flush=True)
    for ep in range(start_ep, args.epochs):
        model.train(); head.train()
        perm = rng.permutation(len(train))
        run_loss, run_n = 0.0, 0
        for i, j in enumerate(perm):
            ex = train[j]
            order = list(rng.permutation(len(ex["cands"])))
            b, spans, qidx, n_img = collate_v(tok, sayfa, ex, order, device)
            if ep == start_ep and i == 0:
                print(f"ilk batch: girdi {tuple(b['input_ids'].shape)} | görüntü token'ı {n_img} | "
                      f"şerit/pixel_values {tuple(b['pixel_values'].shape) if 'pixel_values' in b else None} | aday {len(order)}", flush=True)
            try:
                with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device == "cuda"):
                    h = run_model_v(model, b, args.layer)
                logits, opl = head(h, spans, qidx)
                l_sel = T.loss_fn(logits[0], ex, order, opl[0])
                tgt_op = torch.tensor([T.OPS.get(ex["op"], 0)], device=device)
                loss = l_sel + 0.3 * nn.functional.cross_entropy(opl, tgt_op)
                (loss / args.accum).backward()
                run_loss += float(l_sel.detach()); run_n += 1
            except torch.cuda.OutOfMemoryError:
                n_oom += 1
                h = logits = opl = l_sel = loss = None
                opt.zero_grad(set_to_none=True); torch.cuda.empty_cache()
                print(f"!!! bellek taştı (girdi {tuple(b['input_ids'].shape)}), adım atlandı (toplam {n_oom})", flush=True)
            if (i + 1) % args.accum == 0 or i + 1 == len(perm):
                nn.utils.clip_grad_norm_(lora_params + list(head.parameters()), 1.0)
                opt.step(); sched.step(); opt.zero_grad(set_to_none=True); step += 1
                if step % 25 == 0:
                    print(f"ep {ep} adım {step}/{total} | seçim kaybı {run_loss / run_n:.3f} | {time.time() - t0:.0f}s", flush=True)
                    run_loss, run_n = 0.0, 0
        mem = torch.cuda.max_memory_allocated() / 2 ** 30 if device == "cuda" else 0
        print(f"epoch {ep} eğitim süresi {time.time() - t0:.0f}s ({(time.time() - t0) / (len(perm) * (ep + 1)):.2f} sn/örnek) "
              f"| tepe bellek {mem:.1f} GB | ekran görüntüsü olmayan {sayfa.n_missing} | {sayfa.n_cut}/{max(1, sayfa.n_pages)} sayfa kesildi", flush=True)
        r = T.report(evaluate_v(model, head, tok, sayfa, data["val"], args.layer, device))
        print(f"== epoch {ep} val: acc {r['acc']:.3f} (ece {r['ece']:.3f}, auroc {r['auroc']:.3f}, "
              f"op {r['op_acc']:.3f}, liste içi {r['acc_pos_in_list']:.3f})", flush=True)
        log.append({"epoch": ep, **{k: v for k, v in r.items() if k != "acc_ci95"}})
        if r["acc"] > best:
            best = r["acc"]
            model.save_pretrained(os.path.join(args.out, "lora"))
            torch.save(head.state_dict(), os.path.join(args.out, "head.pt"))
        with open(os.path.join(args.out, "val_gunlugu.json"), "w") as f:
            json.dump(log, f, indent=1)
        model.save_pretrained(os.path.join(args.out, "lora_son")); torch.save(head.state_dict(), os.path.join(args.out, "head_son.pt"))
        torch.save({"opt": opt.state_dict(), "sched": sched.state_dict(), "ep": ep, "best": best, "log": log, "step": step,
                    "rng": rng.bit_generator.state, "n_oom": n_oom}, os.path.join(args.out, "ckpt.pt"))
        sync(args.out, args.sync_dir)

    if os.path.exists(os.path.join(args.out, "head.pt")):
        head.load_state_dict(torch.load(os.path.join(args.out, "head.pt"), map_location=device))
        from peft import set_peft_model_state_dict
        try:
            from safetensors.torch import load_file
            set_peft_model_state_dict(model, load_file(os.path.join(args.out, "lora", "adapter_model.safetensors")))
        except Exception as e:
            print("LoRA en iyi epoch yüklenemedi, son epoch ile devam:", e)
    final = {}
    for s in args.eval_splits:
        if s not in data:
            continue
        res = evaluate_v(model, head, tok, sayfa, data[s], args.layer, device)
        final[s] = T.report(res)
        if s == "val" and not args.no_reverse:
            rev = evaluate_v(model, head, tok, sayfa, data[s], args.layer, device, mode="reverse")
            final[s]["sira_degisim_orani_ters_siralama"] = float(np.mean([a["pred"] != c["pred"] for a, c in zip(res, rev)]))
        with open(os.path.join(args.out, f"tahmin_{s}.json"), "w") as f:
            json.dump(res, f)
        print(f"\n[{s}] " + json.dumps({k: (round(v, 3) if isinstance(v, float) else v) for k, v in final[s].items()
                                       if k != 'acc_ci95'}, ensure_ascii=False))
        print("   acc %95 aralığı:", [round(x, 3) for x in final[s]["acc_ci95"]])
    with open(os.path.join(args.out, "sonuc.json"), "w", encoding="utf-8") as f:
        json.dump({"args": vars(args), "val_log": log, "final": final,
                   "ekran_goruntusu_olmayan": sayfa.n_missing, "oom_atlanan_adim": n_oom, "n_lora_modul": len(names)}, f, indent=2, ensure_ascii=False)
    sync(args.out, args.sync_dir)
    print("\nyazıldı:", args.out)


if __name__ == "__main__":
    main()