import os
import sys

import cv2
import numpy as np
import torch
from PIL import Image, ImageOps
from tqdm import tqdm

# Settings --------------------------------------------------------------------

DATA_DIR   = "/home/nobus/Raid0/DataSet/Images1"
OUT_DIR    = "/home/nobus/Raid0/DataSet/hed_maps_768"   # saved as {stem}.jpg alongside originals
IMG_SIZE   = 768
DEVICE     = "cuda:0"
BATCH_SIZE = 8    # images processed in parallel through HED

HED_CKPT   = os.path.join(os.path.dirname(os.path.abspath(__file__)), "checkpoints/ControlNetHED.pth")

# -----------------------------------------------------------------------------

_D    = os.path.dirname(os.path.abspath(__file__))
_ROOT = os.path.abspath(os.path.join(_D, "../.."))
sys.path.insert(0, os.path.join(_ROOT, "Sana"))

from einops import rearrange
from hed import ControlNetHED_Apache2


def load_hed(device):
    net = ControlNetHED_Apache2().float().to(device).eval()
    ckpt = torch.load(HED_CKPT, map_location="cpu", weights_only=False)
    if "state_dict" in ckpt:
        ckpt = ckpt["state_dict"]
    net.load_state_dict(ckpt)
    return net


@torch.no_grad()
def run_hed(net, arr_hwc, device):
    H, W = arr_hwc.shape[:2]
    t = torch.from_numpy(arr_hwc.copy()).float().to(device)
    t = rearrange(t, "h w c -> 1 c h w")
    edges = net(t)
    edges = [e.detach().cpu().numpy().astype(np.float32)[0, 0] for e in edges]
    edges = [cv2.resize(e, (W, H), interpolation=cv2.INTER_LINEAR) for e in edges]
    edges = np.stack(edges, axis=2)
    edge = 1.0 / (1.0 + np.exp(-np.mean(edges, axis=2).astype(np.float64)))
    return (edge * 255).clip(0, 255).astype(np.uint8)


def main():
    os.makedirs(OUT_DIR, exist_ok=True)

    exts = {".jpg", ".jpeg", ".png", ".webp"}
    paths = sorted(
        os.path.join(r, f)
        for r, _, files in os.walk(DATA_DIR)
        for f in files
        if os.path.splitext(f)[1].lower() in exts
    )
    print(f"Found {len(paths)} images")

    # Skip already processed
    pending = []
    for p in paths:
        stem = os.path.splitext(os.path.basename(p))[0]
        out  = os.path.join(OUT_DIR, f"{stem}.jpg")
        if not os.path.exists(out):
            pending.append(p)
    print(f"Pending: {len(pending)}  (already done: {len(paths) - len(pending)})")

    if not pending:
        print("All done!")
        return

    print(f"Loading HED from {HED_CKPT}...")
    net = load_hed(DEVICE)

    for path in tqdm(pending, unit="img"):
        stem = os.path.splitext(os.path.basename(path))[0]
        out  = os.path.join(OUT_DIR, f"{stem}.jpg")

        img = ImageOps.fit(
            Image.open(path).convert("RGB"),
            (IMG_SIZE, IMG_SIZE),
            method=Image.LANCZOS,
        )
        arr = np.asarray(img, dtype=np.uint8)

        edge = run_hed(net, arr, DEVICE)

        Image.fromarray(edge).save(out, quality=90)

    print(f"Done -> {OUT_DIR}/")


if __name__ == "__main__":
    main()