# ─────────────────────────────────────────────────────────────────────────────
# Cell 2 — JSON-conditioned sd15-flow-lune finetune  (225-token chunked CLIP)
#
# Adapts the Flux-teacher lune trainer to condition on structured JSON instead
# of natural-language prompts. Warm-starts from the checkpoint-00018765 UNet,
# trains on AbstractPhil/synthetic-object-relations-json, pushes to a NEW repo.
#
# Run TWICE — flip RUN — to produce the A/B couple:
#   RUN = "prompt"  → conditions on json_prompt      → sd15-flow-lune-json-prompt
#   RUN = "vit"     → conditions on vit_json_prompt  → sd15-flow-lune-json-vit
#
# CLIP CONDITIONING — 225-token kohya-style chunked encoding:
#   Cell 1b found 94% of the ViT-derived JSON exceeds the 75-token single-CLIP-
#   chunk budget (rich captions → 90-140 token JSON). So both runs encode the
#   JSON as 3 chunks of 75 content tokens: tokenize to 227, split into 3×77
#   (each re-wrapped with BOS/EOS), encode each through frozen CLIP, stitch the
#   hidden states → a 227-length conditioning sequence. The UNet cross-attention
#   takes any length, so no UNet change. Short prompt-JSON simply fills one
#   chunk and pads the rest — both runs use identical machinery, so the only
#   variable in the A/B is the conditioning text.
#
# Prereqs: Cell 1b must have added vit_json_prompt (done). This cell reads the
# dataset's parquet shards directly, so it does NOT depend on the repo README
# metadata. Full finetune, flow-matching timesteps untouched (full 0-1000),
# base_lr 1e-5, save_optimizer=False (ships only the UNet).
# ─────────────────────────────────────────────────────────────────────────────


# ═════════════════════════════════════════════════════════════════════════════
#  RUN SELECTOR — flip this for each of the two finetunes
# ═════════════════════════════════════════════════════════════════════════════

RUN = "vit"        # "prompt" → json_prompt      "vit" → vit_json_prompt


# ═════════════════════════════════════════════════════════════════════════════
# 1. INSTALL
# ═════════════════════════════════════════════════════════════════════════════

import subprocess, sys, os

def _pip(*pkgs):
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *pkgs])

print("Installing dependencies…")
_pip("-U", "diffusers>=0.30", "transformers>=4.50", "accelerate>=1.0",
     "datasets>=4.0", "huggingface_hub>=0.25", "tensorboard")
print("  done.")


# ═════════════════════════════════════════════════════════════════════════════
# 2. AUTH
# ═════════════════════════════════════════════════════════════════════════════

def _load_hf_token():
    if os.environ.get("HF_TOKEN"):
        return "env"
    try:
        from google.colab import userdata
        tok = userdata.get("HF_TOKEN")
        if tok:
            os.environ["HF_TOKEN"] = tok
            os.environ["HUGGING_FACE_HUB_TOKEN"] = tok
            return "secrets"
    except Exception:
        pass
    return None

print(f"HF token: {_load_hf_token() or 'not set'}")


# ═════════════════════════════════════════════════════════════════════════════
# 3. IMPORTS
# ═════════════════════════════════════════════════════════════════════════════

import json
import datetime
from dataclasses import dataclass, asdict
from tqdm.auto import tqdm
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader
import datasets
from diffusers import UNet2DConditionModel, AutoencoderKL
from transformers import CLIPTextModel, CLIPTokenizer
from huggingface_hub import HfApi, create_repo

assert torch.cuda.is_available(), "No GPU. Switch Colab runtime."
print(f"GPU: {torch.cuda.get_device_name(0)}  "
      f"({torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB)")


# ═════════════════════════════════════════════════════════════════════════════
# 4. CONFIG
# ═════════════════════════════════════════════════════════════════════════════

#                   condition column   length column        target repo
_RUN_MAP = {
    "prompt": ("json_prompt",     "json_token_len",     "AbstractPhil/sd15-flow-lune-json-prompt"),
    "vit":    ("vit_json_prompt", "vit_json_token_len", "AbstractPhil/sd15-flow-lune-json-vit"),
}
assert RUN in _RUN_MAP, f"RUN must be one of {list(_RUN_MAP)}, got {RUN!r}"
_COND_COL, _LEN_COL, _HF_REPO = _RUN_MAP[RUN]


@dataclass
class TrainConfig:
    # --- set by the RUN selector ---
    condition_column: str = _COND_COL
    len_column: str       = _LEN_COL
    hf_repo_id: str       = _HF_REPO

    # --- sources ---
    dataset_name: str   = "AbstractPhil/synthetic-object-relations-json"
    sd_base: str        = "stable-diffusion-v1-5/stable-diffusion-v1-5"   # VAE + CLIP
    unet_repo: str      = "AbstractPhil/sd15-flow-lune-flux"
    unet_subfolder: str = "flux_t2_6_pose_t4_6_port_t1_4/checkpoint-00018765/unet"

    output_dir: str      = "./outputs"
    max_clip_tokens: int = 225           # 3 chunks × 75 — rows above this are filtered out

    seed: int        = 42
    batch_size: int  = 8
    base_lr: float   = 1e-5              # modality shift — above the 2e-6 continued-run rate
    shift: float     = 2.0
    dropout: float   = 0.1               # CFG conditioning dropout

    num_train_epochs: int    = 4         # prototype: short, just see if it behaves
    warmup_epochs: int       = 1
    checkpointing_steps: int = 1000
    num_workers: int         = 0         # collate does VAE/CLIP on GPU — must be 0

    vae_scale: float     = 0.18215
    save_optimizer: bool = False         # prototype: ship only the UNet, skip the ~7 GB .pt
    upload_to_hub: bool  = True


# ═════════════════════════════════════════════════════════════════════════════
# 5. 225-TOKEN CHUNKED CLIP ENCODING  (kohya-style)
# ═════════════════════════════════════════════════════════════════════════════

def encode_clip_225(prompts, tokenizer, text_encoder, device):
    """Encode text as 3 × 75-token chunks → one 227-length hidden-state sequence.

    A CLIP window is BOS + 75 content + EOS. We tokenize to 227 (= 3×75 + 2),
    split into three 75-token bodies, re-wrap each with the original BOS/EOS
    into a valid 77-token chunk, encode every chunk through the frozen CLIP,
    then concatenate the hidden states — keeping only the first BOS and last
    EOS. The UNet cross-attention consumes the 227-length result directly.
    """
    chunk_len = tokenizer.model_max_length          # 77
    body_len  = chunk_len - 2                        # 75
    n_chunks  = 3
    max_tok   = body_len * n_chunks                  # 225

    ids = tokenizer(prompts, padding="max_length", max_length=max_tok + 2,
                    truncation=True, return_tensors="pt").input_ids        # [B, 227]
    bos, eos = ids[:, :1], ids[:, -1:]
    chunks = []
    for k in range(n_chunks):
        s = 1 + k * body_len
        chunks.append(torch.cat([bos, ids[:, s:s + body_len], eos], dim=1))  # [B, 77]
    ids = torch.stack(chunks, dim=1).reshape(-1, chunk_len).to(device)        # [B*3, 77]

    hs = text_encoder(ids)[0]                                                 # [B*3, 77, 768]
    hs = hs.reshape(len(prompts), n_chunks * chunk_len, -1)                    # [B, 231, 768]

    out = [hs[:, :1]]                                                         # first BOS
    for k in range(n_chunks):
        s = k * chunk_len + 1
        out.append(hs[:, s:s + body_len])                                     # 75 content tokens
    out.append(hs[:, -1:])                                                    # last EOS
    return torch.cat(out, dim=1)                                              # [B, 227, 768]


# ═════════════════════════════════════════════════════════════════════════════
# 6. LOAD UNET — straight from the diffusers checkpoint folder
# ═════════════════════════════════════════════════════════════════════════════

def load_unet(config, device="cuda"):
    print(f"\nLoading UNet from {config.unet_repo}/{config.unet_subfolder}…")
    unet = UNet2DConditionModel.from_pretrained(
        config.unet_repo, subfolder=config.unet_subfolder, torch_dtype=torch.float32)
    print(f"  ✓ UNet loaded ({sum(p.numel() for p in unet.parameters()) / 1e6:.0f}M params)")
    return unet.to(device)


# ═════════════════════════════════════════════════════════════════════════════
# 7. TRAIN
# ═════════════════════════════════════════════════════════════════════════════

def train(config):
    device = "cuda"
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.manual_seed(config.seed)
    torch.cuda.manual_seed(config.seed)

    date_time = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    real_output_dir = os.path.join(config.output_dir, f"{RUN}_{date_time}")
    os.makedirs(real_output_dir, exist_ok=True)
    t_writer = SummaryWriter(log_dir=real_output_dir, flush_secs=60)

    hf_api = None
    if config.upload_to_hub:
        try:
            hf_api = HfApi()
            create_repo(repo_id=config.hf_repo_id, repo_type="model",
                        exist_ok=True, private=False)
            print(f"✓ HF repo ready: {config.hf_repo_id}")
        except Exception as e:
            print(f"⚠ hub upload disabled: {e}")
            config.upload_to_hub = False

    config_path = os.path.join(real_output_dir, "config.json")
    with open(config_path, "w") as f:
        json.dump({"run": RUN, **asdict(config)}, f, indent=2)
    if config.upload_to_hub:
        hf_api.upload_file(path_or_fileobj=config_path, path_in_repo="config.json",
                           repo_id=config.hf_repo_id, repo_type="model")

    # ── VAE + CLIP (frozen) ──────────────────────────────────────────────────
    print(f"\nLoading VAE + CLIP from {config.sd_base}…")
    vae = AutoencoderKL.from_pretrained(config.sd_base, subfolder="vae",
                                        torch_dtype=torch.float32).to(device)
    vae.requires_grad_(False); vae.eval()
    tokenizer = CLIPTokenizer.from_pretrained(config.sd_base, subfolder="tokenizer")
    text_encoder = CLIPTextModel.from_pretrained(
        config.sd_base, subfolder="text_encoder", torch_dtype=torch.float32).to(device)
    text_encoder.requires_grad_(False); text_encoder.eval()
    print("✓ VAE + CLIP loaded (frozen)")
    print(f"  conditioning: 3×{tokenizer.model_max_length - 2}-token chunks "
          f"→ 227-length sequence (kohya-style)")

    # ── dataset: load the parquet shards directly ────────────────────────────
    # Cell 1b added vit_json_* columns to the parquet but its push left the repo
    # README's dataset_info stale (11 cols vs 13 in the parquet), so
    # load_dataset(repo) CastErrors casting parquet → README schema. The parquet
    # builder reads the parquet's own (correct) schema and never sees the README.
    print(f"\nLoading dataset {config.dataset_name} (parquet-direct)…")
    _pq = sorted(f for f in HfApi().list_repo_files(config.dataset_name, repo_type="dataset")
                 if f.endswith(".parquet"))
    if not _pq:
        raise RuntimeError(f"no parquet files found in {config.dataset_name}")
    ds = datasets.load_dataset(
        "parquet",
        data_files={"train": [f"hf://datasets/{config.dataset_name}/{f}" for f in _pq]},
        split="train")
    ds = ds.cast_column("image", datasets.Image())
    col, lcol = config.condition_column, config.len_column
    if col not in ds.column_names or lcol not in ds.column_names:
        raise RuntimeError(
            f"dataset is missing '{col}'/'{lcol}'. Run Cell 1b then Cell 1c first.")
    n_before = len(ds)
    ds = ds.filter(lambda ex: isinstance(ex[col], str) and ex[col].strip() != ""
                              and 0 < ex[lcol] <= config.max_clip_tokens)
    ds = ds.select_columns(["id", "image", col])
    print(f"✓ {len(ds)}/{n_before} rows usable for '{col}' (JSON ≤ {config.max_clip_tokens} tok)")
    if len(ds) == 0:
        raise RuntimeError(f"no usable rows for '{col}'.")

    steps_per_epoch = len(ds) // config.batch_size
    total_steps     = steps_per_epoch * config.num_train_epochs
    warmup_steps    = max(steps_per_epoch * config.warmup_epochs, 1)
    print(f"\nSchedule: {steps_per_epoch} steps/epoch × {config.num_train_epochs} "
          f"epochs = {total_steps} steps  (warmup {warmup_steps})")

    @torch.no_grad()
    def collate_fn(examples):
        """Encode images (VAE) and JSON conditioning text (chunked CLIP)."""
        images, prompts, ids = [], [], []
        for ex in examples:
            img = ex["image"].convert("RGB")
            img = torch.tensor(np.array(img)).permute(2, 0, 1).float() / 255.0
            images.append(img * 2.0 - 1.0)              # [-1, 1]
            prompts.append(ex[col])
            ids.append(ex["id"])
        images  = torch.stack(images).to(device)
        latents = vae.encode(images).latent_dist.sample() * config.vae_scale
        ehs = encode_clip_225(prompts, tokenizer, text_encoder, device)   # [B, 227, 768]
        return latents.cpu(), ehs.cpu(), ids, prompts

    train_loader = DataLoader(ds, batch_size=config.batch_size, shuffle=True,
                              collate_fn=collate_fn, num_workers=config.num_workers,
                              pin_memory=True)

    # ── UNet + fresh optimizer ───────────────────────────────────────────────
    unet = load_unet(config, device)
    unet.requires_grad_(True); unet.train()
    optimizer = torch.optim.AdamW(unet.parameters(), lr=config.base_lr,
                                  betas=(0.9, 0.999), weight_decay=0.01, eps=1e-8)

    def lr_scale(step):
        return step / warmup_steps if step < warmup_steps else 1.0
    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_scale)
    print(f"✓ fresh AdamW, lr {config.base_lr:.1e}, linear warmup")

    global_step = 0
    train_logs  = {"train_step": [], "train_loss": [], "train_timestep": [],
                   "trained_images": []}

    def get_prediction(batch, log_to=None):
        latents, ehs, ids, prompts = batch
        latents = latents.to(dtype=torch.float32, device=device)
        ehs     = ehs.to(dtype=torch.float32, device=device)
        bsz = latents.shape[0]

        # CFG conditioning dropout — zero some embeddings
        drop = torch.rand(bsz, device=device) < config.dropout
        ehs = ehs.clone(); ehs[drop] = 0

        # shifted flow-matching timesteps, full 0–1000 range (no masking)
        sigmas = torch.rand(bsz, device=device)
        sigmas = (config.shift * sigmas) / (1 + (config.shift - 1) * sigmas)
        timesteps = sigmas * 1000
        sigmas = sigmas[:, None, None, None]

        noise  = torch.randn_like(latents)
        noisy  = noise * sigmas + latents * (1 - sigmas)
        target = noise - latents                              # velocity

        pred = unet(noisy, timesteps, ehs, return_dict=False)[0]
        loss = F.mse_loss(pred, target, reduction="none")
        loss = loss.mean(dim=list(range(1, len(loss.shape))))

        if log_to is not None:
            for i in range(bsz):
                log_to["train_step"].append(global_step)
                log_to["train_loss"].append(loss[i].item())
                log_to["train_timestep"].append(timesteps[i].item())
                log_to["trained_images"].append(
                    {"step": global_step, "id": ids[i], "prompt": prompts[i]})
        return loss.mean()

    def plot_logs(d):
        plt.figure(figsize=(10, 6))
        plt.scatter(d["train_timestep"], d["train_loss"], s=3, c=d["train_step"],
                    marker=".", cmap="cool")
        plt.xlabel("timestep"); plt.ylabel("loss"); plt.yscale("log")
        plt.colorbar(label="step")

    def save_checkpoint(step, epoch):
        ckpt = os.path.join(real_output_dir, f"checkpoint-{step:08}")
        os.makedirs(ckpt, exist_ok=True)
        unet.save_pretrained(os.path.join(ckpt, "unet"), safe_serialization=True)
        meta = {"step": step, "epoch": epoch, "run": RUN, "condition_column": col,
                "trained_images": train_logs["trained_images"]}
        meta_path = os.path.join(ckpt, "trained_images.json")
        with open(meta_path, "w") as f:
            json.dump(meta, f, indent=2)
        if config.save_optimizer:
            torch.save({"cfg": asdict(config), "student": unet.state_dict(),
                        "opt": optimizer.state_dict(), "scheduler": scheduler.state_dict(),
                        "gstep": step, "epoch": epoch},
                       os.path.join(ckpt, f"sd15_flow_lune_json_e{epoch}_s{step}.pt"))
        print(f"✓ checkpoint {step} (epoch {epoch})")
        if config.upload_to_hub and hf_api is not None:
            try:
                hf_api.upload_folder(folder_path=os.path.join(ckpt, "unet"),
                                     path_in_repo=f"checkpoint-{step:08}/unet",
                                     repo_id=config.hf_repo_id, repo_type="model")
                hf_api.upload_file(path_or_fileobj=meta_path,
                                   path_in_repo=f"checkpoint-{step:08}/trained_images.json",
                                   repo_id=config.hf_repo_id, repo_type="model")
                print(f"  ✓ uploaded to {config.hf_repo_id}")
            except Exception as e:
                print(f"  ⚠ upload failed: {e}")

    # ── training loop ────────────────────────────────────────────────────────
    print(f"\nStarting training — RUN='{RUN}', conditioning on '{col}'\n")
    progress = tqdm(total=total_steps)
    for epoch in range(config.num_train_epochs):
        for batch in train_loader:
            if global_step >= total_steps:
                break
            loss = get_prediction(batch, log_to=train_logs)
            t_writer.add_scalar("train/loss", loss.item(), global_step)
            t_writer.add_scalar("train/lr", scheduler.get_last_lr()[0], global_step)

            loss.backward()
            gn = torch.nn.utils.clip_grad_norm_(unet.parameters(), 1.0)
            t_writer.add_scalar("train/grad_norm", gn.item(), global_step)
            optimizer.step(); scheduler.step(); optimizer.zero_grad()

            progress.update(1)
            progress.set_postfix({"epoch": epoch, "loss": f"{loss.item():.4f}",
                                  "lr": f"{scheduler.get_last_lr()[0]:.2e}"})
            global_step += 1

            if global_step % 100 == 0:
                plot_logs(train_logs)
                t_writer.add_figure("train_loss", plt.gcf(), global_step)
                plt.close()
            if global_step % config.checkpointing_steps == 0:
                save_checkpoint(global_step, epoch)

        save_checkpoint(global_step, epoch)        # end-of-epoch

    progress.close()
    print(f"\n✅ Training complete — RUN='{RUN}'")
    print(f"   → https://huggingface.co/{config.hf_repo_id}")
    other = "vit" if RUN == "prompt" else "prompt"
    print(f"   Now set RUN = \"{other}\" and re-run this cell for the other half "
          f"of the couple.")


# ═════════════════════════════════════════════════════════════════════════════
# 8. RUN
# ═════════════════════════════════════════════════════════════════════════════

train(TrainConfig())