"""LoRA trainer for Fibo 2.0 (edit and text-to-image) on the diffusers `BriaFibo2Pipeline`.

Self-contained: needs only torch, diffusers (with `BriaFibo2Pipeline`: until PR
huggingface/diffusers#14991 is merged, branch ali/fibo2 of muhammadali-afridi/diffusers, tested
at e158bdea), torchvision (Qwen3-VL's processor needs it), transformers, peft, safetensors,
pillow, numpy and click.

What it does
------------
1. Pre-encodes every training sample once with the pipeline's own methods (`encode_prompt`,
   `prepare_image_latents`, `image_processor`, the pipeline's caption cleaning and its text
   images, drawn in the base repo's fonts), writes the result to
   `<cache-dir>/<idx>.safetensors`, then frees Qwen3-VL and the VAE.
2. Loads the transformer in bf16 (frozen), adds a LoRA to every `nn.Linear`, keeps the LoRA
   weights in fp32, and trains them with flow matching (batch 1, gradient accumulation).
3. Saves `<output-dir>/checkpoint-<step>/` every `--save-every` steps and `<output-dir>/final/`
   (`pytorch_lora_weights.safetensors`, see "Load the result"). Also writes `loss.csv`.

Data
----
edit mode (`<data>/edits.jsonl` exists, or `--mode edit`): one JSON object per line,
    {"sources": ["a/src.jpg", ...1 to 4 paths], "target": "a/tgt.jpg", "instruction": "..."}
paths are relative to `<data>`. The instruction may refer to the sources as <image_1>, <image_2>.
"instruction" is a plain string, or a structured JSON caption of the target (an object, or that
object serialised) with an "edit_instruction" key; the pipeline splits it into the two.
t2i mode: images (png/jpg/jpeg/webp) anywhere under `<data>`. Each image's caption is a file next
to it with the same name: `<name>.json` (a structured caption object, re-serialised compact) or
`<name>.txt` (a plain sentence, used as is). The .json wins if both exist; uncaptioned images are
skipped. The README has an example.

Text images: a JSON caption with `text_render` entries gets an image of each entry's text, drawn
and shown to Qwen3-VL as the pipeline does at inference. Text-to-image samples also cache the
caption without them, used above noise level 0.9, where inference leaves them out.

Sizing: the result size is chosen on the 16-pixel grid within 10% of `--resolution`^2 pixels,
cropping the least from the first source (edit) or the target (t2i), exactly as the pipeline does
at inference (`image_processor.get_image_size` with area `--resolution`^2). Targets and extra
sources are cover-cropped to their sizes by `image_processor.preprocess`.

Run (1024 / rank 128 / grad accum 8 peaked at 31 GiB, about 58 s per step on an L40S)
--------------------------------------------------------------------------------------
    python train_lora.py --base <hf-repo-or-local-dir> --data my_edits/ --output-dir out/ \\
        --resolution 1024 --rank 64 --lr 1e-4 --steps 1000 --grad-accum 8

Load the result (`weight_name` is required: without it diffusers looks for a `.bin`):
    pipe = BriaFibo2Pipeline.from_pretrained(base, dtype=torch.bfloat16)
    pipe.transformer.load_lora_adapter("out/final", weight_name="pytorch_lora_weights.safetensors")

Recipe: sigma ~ sigmoid(N(shift, 1)) with prob 0.9 else U(0, 1), shift = 0.5 + 0.65 * (n - 256) /
3840 for n target tokens; x_t = sigma * noise + (1 - sigma) * x0; target = noise - x0; fp32 MSE.
Text dropout swaps in the null caption embedding; context dropout (edit) removes the ref latents
but keeps what Qwen3-VL saw. AdamW (wd 0.01), constant LR, grad clip 1.0.
"""

import csv
import gc
import json
import math
import random
import time
from pathlib import Path
from typing import Any

import click
import numpy as np
import torch
from PIL import Image, ImageOps
from safetensors import SafetensorError, safe_open
from safetensors.torch import load_file, save_file
from torch import nn

DEFAULT_BASE = "briaai/Fibo-2.0-dev"
EXPECTED_ADAPTED_MODULES = 415  # nn.Linear layers of the Fibo 2.0 transformer
IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp"}
FONT_FILES = ("Aileron-Regular.otf", "DejaVuSans.ttf")
SHIFT_BASE, SHIFT_MAX, SEQ_BASE, SEQ_MAX = 0.5, 1.15, 256, 4096
# text-to-image leaves out the text images above this noise level, as the pipeline does
TEXT_IMAGES_MAX_SIGMA = 0.9


# ------------------------------------------------------------------------------------ data
def compact(caption: dict) -> str:
    return json.dumps(caption, separators=(",", ":"), ensure_ascii=False)


def as_prompt(value: str | dict) -> str:
    """A caption object serialised compactly, as Fibo 2.0 was trained, or a plain string as is."""
    if isinstance(value, dict):
        return compact(value)
    return value


def edit_prompt(instruction: str | dict) -> str:
    """An edit row's prompt. A JSON caption must hold the "edit_instruction" key."""
    prompt = as_prompt(instruction)
    try:
        caption = json.loads(prompt)
    except json.JSONDecodeError:
        return prompt
    if not isinstance(caption, dict):
        return prompt
    if "edit_instruction" not in caption:
        raise click.ClickException(
            f'JSON instruction has no "edit_instruction" key: {prompt[:200]}'
        )
    if "structured_caption" in caption and (
        not isinstance(caption["structured_caption"], (dict, str)) or len(caption) > 2
    ):
        raise click.ClickException(
            'a "structured_caption" instruction holds only "structured_caption" (an object or a'
            f' string) and "edit_instruction": {prompt[:200]}'
        )
    return prompt


def load_samples(data: Path, mode: str) -> list[dict[str, Any]]:
    """Return [{'sources': [Path], 'target': Path, 'prompt': str}] (sources empty for t2i)."""
    samples: list[dict[str, Any]] = []
    if mode == "edit":
        for line in (data / "edits.jsonl").read_text().splitlines():
            if not line.strip():
                continue
            row = json.loads(line)
            if not 1 <= len(row["sources"]) <= 4:
                raise click.ClickException(f"need 1-4 sources, got {len(row['sources'])}")
            samples.append(
                {
                    "sources": [data / s for s in row["sources"]],
                    "target": data / row["target"],
                    "prompt": edit_prompt(row["instruction"]),
                }
            )
    else:
        for img in sorted(p for p in data.rglob("*") if p.suffix.lower() in IMAGE_SUFFIXES):
            if img.with_suffix(".json").exists():
                prompt = as_prompt(json.loads(img.with_suffix(".json").read_text()))
            elif img.with_suffix(".txt").exists():
                prompt = img.with_suffix(".txt").read_text().strip()
            else:
                continue
            samples.append({"sources": [], "target": img, "prompt": prompt})
    if not samples:
        raise click.ClickException(f"no samples found in {data} (mode={mode})")
    return samples


def load_image(path: Path) -> Image.Image:
    return ImageOps.exif_transpose(Image.open(path)).convert("RGB")


# ----------------------------------------------------------------------- pre-encode phase
def cache_ok(path: Path, meta: dict[str, str]) -> bool:
    if not path.exists():
        return False
    try:
        with safe_open(path, "pt") as f:
            got = f.metadata() or {}
    except (OSError, SafetensorError):  # partial or corrupt file: re-encode
        return False
    return all(got.get(k) == v for k, v in meta.items())


def load_fonts(base: str) -> list[Any]:
    """The fonts the base model's training drew text images with, from its `fonts` folder."""
    from diffusers.pipelines.bria_fibo2.text_images import load_font
    from huggingface_hub import hf_hub_download

    paths = [
        str(Path(base) / "fonts" / name)
        if Path(base).is_dir()
        else hf_hub_download(base, name, subfolder="fonts")
        for name in FONT_FILES
    ]
    return [load_font(path) for path in paths]


def text_images(caption: dict | str | None, fonts: list[Any]) -> list[Image.Image] | None:
    """Images of the caption's text_render entries that have text, drawn as the pipeline does."""
    from diffusers.pipelines.bria_fibo2.text_images import draw_text_image

    entries = (caption.get("text_render") or []) if isinstance(caption, dict) else []
    images = [
        draw_text_image(e["text"], fonts)
        for e in entries
        if isinstance(e, dict) and isinstance(e.get("text"), str) and e["text"]
    ]
    return images or None


def pre_encode(
    base: str,
    samples: list[dict[str, Any]],
    mode: str,
    resolution: int,
    cache_dir: Path,
    dev: torch.device,
) -> None:
    from diffusers import BriaFibo2Pipeline
    from diffusers.pipelines.bria_fibo2.pipeline_bria_fibo2 import (
        DEFAULT_NEGATIVE_PROMPT,
    )

    cache_dir.mkdir(parents=True, exist_ok=True)
    # "encoding" changes whenever the cached embeddings would: older caches are re-encoded
    meta = {"resolution": str(resolution), "mode": mode, "encoding": "3"}
    todo = [i for i in range(len(samples)) if not cache_ok(cache_dir / f"{i}.safetensors", meta)]
    need_null = not cache_ok(cache_dir / "null.safetensors", meta)
    print(f"[encode] {len(samples)} samples, {len(todo)} to encode (cache: {cache_dir})")
    if not todo and not need_null:
        return

    pipe = BriaFibo2Pipeline.from_pretrained(base, transformer=None, dtype=torch.bfloat16)
    pipe.to(dev)
    pipe.text_encoder.eval()
    proc = pipe.image_processor
    area = resolution * resolution
    fonts = load_fonts(base)
    tdtype = torch.bfloat16

    def pack(emb: torch.Tensor, mask: torch.Tensor | None) -> dict[str, torch.Tensor]:
        out = {"emb": emb.to(tdtype).cpu().contiguous()}
        if mask is not None:
            out["mask"] = mask.to(torch.uint8).cpu().contiguous()
        return out

    with torch.no_grad():
        if need_null:
            emb, mask = pipe.encode_prompt(DEFAULT_NEGATIVE_PROMPT, None, dev)
            save_file(pack(emb, mask), cache_dir / "null.safetensors", metadata=meta)
        for n, i in enumerate(todo):
            s = samples[i]
            target = load_image(s["target"])
            caption, _ = pipe._parse_prompt(s["prompt"])
            texts = text_images(caption, fonts)
            extra = {}
            if mode == "edit":
                srcs = [load_image(p) for p in s["sources"]]
                width, height = proc.get_image_size(srcs[0].width / srcs[0].height, area=area)
                # one source: cropped to the result size; several: each to its own aspect's size
                vae_srcs, thumbs = proc.preprocess(
                    srcs,
                    height=height,
                    width=width,
                    area=area,
                    text_encoder_image_size=pipe.text_encoder_image_size,
                )
                emb, mask = pipe.encode_prompt(s["prompt"], thumbs, dev, text_images=texts)
                refs = pipe.prepare_image_latents(vae_srcs, 1, dev, torch.float32)
            else:
                width, height = proc.get_image_size(target.width / target.height, area=area)
                # a JSON caption in its trained form; a plain sentence as is
                prompt = s["prompt"] if isinstance(caption, str) else compact(caption or {})
                emb, mask = pipe.encode_prompt(prompt, None, dev, text_images=texts)
                refs = []
                if texts:
                    # the caption without its text images, which inference uses at high noise
                    base_emb, base_mask = pipe.encode_prompt(
                        compact(pipe._mark_text_images(json.loads(prompt), "")), None, dev
                    )
                    extra = {f"base_{k}": v for k, v in pack(base_emb, base_mask).items()}
            tgt, _ = proc.preprocess(target, height=height, width=width)
            x0 = pipe.prepare_image_latents(tgt, 1, dev, torch.float32)[0]
            tensors = {"x0": x0.cpu().contiguous(), **pack(emb, mask), **extra}
            for k, r in enumerate(refs):
                tensors[f"ref{k}"] = r.cpu().contiguous()
            save_file(tensors, cache_dir / f"{i}.safetensors", metadata=meta)
            if (n + 1) % 20 == 0 or n + 1 == len(todo):
                print(f"[encode] {n + 1}/{len(todo)}")
    del pipe
    gc.collect()
    torch.cuda.empty_cache()


def load_cached(path: Path, dev: torch.device) -> dict[str, Any]:
    t = load_file(path)
    refs = [t[f"ref{k}"].to(dev) for k in range(4) if f"ref{k}" in t]

    def mask(key: str) -> torch.Tensor | None:
        return t[key].to(dev).bool() if key in t else None

    return {
        "x0": t["x0"].to(dev),
        "emb": t["emb"].to(dev),
        "mask": mask("mask"),
        "base_emb": t["base_emb"].to(dev) if "base_emb" in t else None,
        "base_mask": mask("base_mask"),
        "refs": refs or None,
    }


# ---------------------------------------------------------------------------------- model
def build_model(
    base: str, rank: int, alpha: int, grad_ckpt: bool, expect: int, dev: torch.device
) -> tuple[nn.Module, list[nn.Parameter]]:
    from diffusers import BriaFibo2Transformer2DModel
    from peft import LoraConfig
    from peft.tuners.lora import LoraLayer

    model = BriaFibo2Transformer2DModel.from_pretrained(
        base, subfolder="transformer", dtype=torch.bfloat16
    )
    model.to(dev).requires_grad_(False)
    linears = [n for n, m in model.named_modules() if isinstance(m, nn.Linear)]
    model.add_adapter(
        LoraConfig(r=rank, lora_alpha=alpha, init_lora_weights=True, target_modules=linears)
    )
    adapted = sum(isinstance(m, LoraLayer) for m in model.modules())
    print(f"[model] nn.Linear layers: {len(linears)}, LoRA-adapted modules: {adapted}")
    assert adapted == len(linears), "some Linear layers were not adapted"
    if expect:
        assert adapted == expect, f"expected {expect} adapted modules, got {adapted}"
    params = []
    for name, p in model.named_parameters():
        if "lora_" in name:
            # fp32 LoRA weights: bf16 would round small AdamW updates to zero
            p.data = p.data.float()
            p.requires_grad_(True)
            params.append(p)
        else:
            assert not p.requires_grad, name
    assert all(p.dtype == torch.float32 for p in params)
    print(f"[model] LoRA params: {sum(p.numel() for p in params) / 1e6:.1f}M, all fp32")
    if grad_ckpt:
        model.enable_gradient_checkpointing()
    model.train()
    return model, params


def save_adapter(model: nn.Module, directory: Path) -> None:
    """Write the LoRA in the diffusers pipeline format: keys and metadata prefixed `transformer.`.

    (`transformer.save_lora_adapter` writes unprefixed keys, which `load_lora_adapter` with its
    default prefix silently ignores, leaving the model un-adapted.)
    """
    from peft.utils import get_peft_model_state_dict

    state = get_peft_model_state_dict(model, adapter_name="default")
    state = {f"transformer.{k}": v.detach().contiguous() for k, v in state.items()}
    cfg = {
        k: sorted(v) if isinstance(v, set) else v
        for k, v in model.peft_config["default"].to_dict().items()
    }
    meta = {f"transformer.{k}": v for k, v in cfg.items()}
    directory.mkdir(parents=True, exist_ok=True)
    save_file(
        state,
        directory / "pytorch_lora_weights.safetensors",
        metadata={
            "format": "pt",
            "lora_adapter_metadata": json.dumps(meta, indent=2, sort_keys=True),
        },
    )


def sample_sigma(gen: torch.Generator, n_tokens: int) -> float:
    if torch.rand(1, generator=gen).item() < 0.1:
        return torch.rand(1, generator=gen).item()
    shift = SHIFT_BASE + (SHIFT_MAX - SHIFT_BASE) * (n_tokens - SEQ_BASE) / (SEQ_MAX - SEQ_BASE)
    z = torch.randn(1, generator=gen).item() + shift
    return 1.0 / (1.0 + math.exp(-z))


# ------------------------------------------------------------------------------------ CLI
@click.command()
@click.option("--base", default=DEFAULT_BASE, show_default=True, help="HF repo id or local dir.")
@click.option(
    "--data",
    type=click.Path(exists=True, file_okay=False, path_type=Path),
    required=True,
)
@click.option("--mode", type=click.Choice(["edit", "t2i"]), default=None, help="Default: auto.")
@click.option("--output-dir", type=click.Path(path_type=Path), required=True)
@click.option("--resolution", default=1024, show_default=True, help="Target side; area = side^2.")
@click.option("--rank", default=64, show_default=True)
@click.option("--alpha", type=int, default=None, help="LoRA alpha. Default: rank.")
@click.option("--lr", default=1e-4, show_default=True)
@click.option("--steps", default=1000, show_default=True, help="Optimizer steps.")
@click.option("--grad-accum", default=8, show_default=True)
@click.option("--save-every", default=250, show_default=True)
@click.option("--text-drop", default=0.1, show_default=True)
@click.option("--context-drop", default=0.2, show_default=True, help="Edit mode only.")
@click.option("--seed", default=0, show_default=True)
@click.option("--cache-dir", type=click.Path(path_type=Path), default=None)
@click.option("--gradient-checkpointing/--no-gradient-checkpointing", default=True)
@click.option(
    "--expect-modules",
    default=EXPECTED_ADAPTED_MODULES,
    show_default=True,
    help="Assert this many adapted Linear layers (0 disables; for other bases).",
)
def main(
    base: str,
    data: Path,
    mode: str | None,
    output_dir: Path,
    resolution: int,
    rank: int,
    alpha: int | None,
    lr: float,
    steps: int,
    grad_accum: int,
    save_every: int,
    text_drop: float,
    context_drop: float,
    seed: int,
    cache_dir: Path | None,
    gradient_checkpointing: bool,
    expect_modules: int,
) -> None:
    """Train a LoRA on Fibo 2.0 from a folder of edit pairs or captioned images."""
    if not torch.cuda.is_available():
        raise click.ClickException("a CUDA GPU is required")
    dev = torch.device("cuda")
    mode = mode or ("edit" if (data / "edits.jsonl").exists() else "t2i")
    alpha = alpha if alpha is not None else rank
    cache_dir = cache_dir or output_dir / "cache"
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / "train_args.json").write_text(
        json.dumps({k: str(v) for k, v in locals().items() if k != "dev"}, indent=2)
    )
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

    samples = load_samples(data, mode)
    print(f"[data] mode={mode}, {len(samples)} samples, resolution={resolution}")
    pre_encode(base, samples, mode, resolution, cache_dir, dev)
    nt = load_file(cache_dir / "null.safetensors")
    null_emb = nt["emb"].to(dev)
    null_mask = nt["mask"].to(dev).bool() if "mask" in nt else None

    model, params = build_model(base, rank, alpha, gradient_checkpointing, expect_modules, dev)
    opt = torch.optim.AdamW(params, lr=lr, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01)
    torch.cuda.reset_peak_memory_stats()

    gen = torch.Generator().manual_seed(seed)
    order: list[int] = []
    log_path = output_dir / "loss.csv"
    with open(log_path, "w", newline="") as f:
        csv.writer(f).writerow(["step", "loss", "mean_sigma", "seconds_per_step"])
    patch = model.config.patch_size

    for step in range(1, steps + 1):
        t0 = time.time()
        losses: list[float] = []
        sigmas: list[float] = []
        for _ in range(grad_accum):
            if not order:
                order = torch.randperm(len(samples), generator=gen).tolist()
            s = load_cached(cache_dir / f"{order.pop()}.safetensors", dev)
            x0 = s["x0"]
            n_tokens = (x0.shape[2] // patch) * (x0.shape[3] // patch)
            sigma = sample_sigma(gen, n_tokens)
            emb, mask, refs = s["emb"], s["mask"], s["refs"]
            if s["base_emb"] is not None and sigma > TEXT_IMAGES_MAX_SIGMA:
                emb, mask = s["base_emb"], s["base_mask"]
            if torch.rand(1, generator=gen).item() < text_drop:
                emb, mask = null_emb, null_mask
            if mode == "edit" and torch.rand(1, generator=gen).item() < context_drop:
                refs = None
            noise = torch.randn(x0.shape, generator=gen).to(dev)
            x_t = sigma * noise + (1.0 - sigma) * x0
            pred = model(
                hidden_states=x_t.to(torch.bfloat16),
                timestep=torch.full((1,), sigma, dtype=torch.float32, device=dev),
                encoder_hidden_states=emb,
                encoder_attention_mask=mask,
                context_latents=None if refs is None else [r.to(torch.bfloat16) for r in refs],
                return_dict=False,
            )[0]
            loss = torch.nn.functional.mse_loss(pred.float(), noise - x0)
            (loss / grad_accum).backward()
            losses.append(loss.item())
            sigmas.append(sigma)
        torch.nn.utils.clip_grad_norm_(params, 1.0)
        opt.step()
        opt.zero_grad(set_to_none=True)

        dt = time.time() - t0
        mean_loss = float(np.mean(losses))
        if not math.isfinite(mean_loss):
            raise RuntimeError(f"non-finite loss at step {step}")
        with open(log_path, "a", newline="") as f:
            csv.writer(f).writerow(
                [step, f"{mean_loss:.6f}", f"{np.mean(sigmas):.4f}", f"{dt:.2f}"]
            )
        if step == 1 or step % 10 == 0 or step == steps:
            peak = torch.cuda.max_memory_allocated() / 2**30
            print(f"step {step}/{steps} loss {mean_loss:.4f} {dt:.1f}s/step peak {peak:.1f} GiB")
        if step % save_every == 0 and step != steps:
            save_adapter(model, output_dir / f"checkpoint-{step}")
    save_adapter(model, output_dir / "final")
    if steps % save_every == 0:
        save_adapter(model, output_dir / f"checkpoint-{steps}")
    print(f"peak memory allocated: {torch.cuda.max_memory_allocated() / 2**30:.2f} GiB")
    print(f"saved to {output_dir}")


if __name__ == "__main__":
    main()