"""
Utility functions and classes for Easy LoRA Merger.
"""

import warnings
import time
import uuid
import hashlib
import gc
import json
import struct
import threading
from pathlib import Path
from contextlib import contextmanager
from typing import Optional, Dict, Any, List, Tuple, Union

import torch
from safetensors.torch import load_file, save_file
from safetensors import safe_open
import folder_paths
import comfy.utils
import comfy.sd


# Suppress warnings
warnings.filterwarnings("ignore", message="lora key not loaded")


def comfyui_yield():
    """Yield control to ComfyUI if possible, otherwise sleep briefly."""
    try:
        # Try different ComfyUI yield methods
        if hasattr(comfy.utils, 'yield_for_comfyui'):
            comfy.utils.yield_for_comfyui()
        elif hasattr(comfy.utils, 'yield_current'):
            comfy.utils.yield_current()
        else:
            time.sleep(0.001)
    except:
        time.sleep(0.001)


def read_safetensors_header_only(path: Path) -> Tuple[Dict[str, Any], Dict[str, str]]:
    """
    Read ONLY the header JSON from a safetensors file, without loading any tensors.
    
    The safetensors format stores tensor metadata (names, shapes, dtypes, offsets)
    in a JSON header at the beginning of the file. This function reads only that
    header, avoiding the tensors themselves entirely.
    
    Memory: ~50 KB for a 425-key checkpoint (vs 17+ GiB for full tensor load).
    
    Args:
        path: Path to the .safetensors file
    
    Returns:
        (header_dict, metadata_dict) where header_dict maps tensor names to their
        shape/dtype/data_offsets info, and metadata_dict is the __metadata__ field.
    """
    with open(path, 'rb') as f:
        header_len_bytes = f.read(8)
        if len(header_len_bytes) != 8:
            raise ValueError(f"File too short: {path}")
        header_len = int.from_bytes(header_len_bytes, 'little')
        header_json_bytes = f.read(header_len)
        if len(header_json_bytes) != header_len:
            raise ValueError(f"Header length mismatch in {path}: expected {header_len}, got {len(header_json_bytes)}")
        header = json.loads(header_json_bytes.decode('utf-8'))
    
    # Separate metadata from tensor headers
    metadata = header.pop('__metadata__', {})
    return header, metadata


def cleanup_memory(*objs, skip_gc=False):
    """
    Run garbage collection (optional) and clear CUDA cache.

    NOTE: Passing objects to this function does NOT delete them from the
    caller's scope.  Python's `del obj` inside a function only deletes the
    local parameter reference.  To actually free memory, callers MUST use
    explicit `del` at the call site:

        del batch_sds, batch_mappings, merged_dict
        cleanup_memory()  # gc.collect() + torch.cuda.empty_cache()

    Args:
        skip_gc: When True, skip gc.collect() and only clear CUDA cache.
                 Use this at the end of long pipelines where GC already ran
                 naturally, to avoid triggering __del__ at unpredictable times
                 (which can cause unexpected file I/O on lazy mappings).
    """
    if not skip_gc:
        gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


class NodeCache:
    """MD5-hash based change detection for ComfyUI IS_CHANGED."""

    _CACHE: Dict[tuple, str] = {}

    @classmethod
    def check_changed(cls, *cache_values) -> str:
        """
        Hash-based change detection for IS_CHANGED.

        Returns the MD5 hash of cache_values.
        If the same values produce the same hash, ComfyUI treats it as unchanged.

        Usage in IS_CHANGED:
            @classmethod
            def IS_CHANGED(cls, **kwargs):
                return NodeCache.check_changed(
                    cls.__name__,
                    kwargs.get('param1', ''),
                    kwargs.get('param2', ''),
                )
        """
        key_str = str(cache_values).encode('utf-8')
        current_hash = hashlib.md5(key_str).hexdigest()

        if cache_values in cls._CACHE and cls._CACHE[cache_values] == current_hash:
            return cls._CACHE[cache_values]

        cls._CACHE[cache_values] = current_hash
        return current_hash

    @classmethod
    def is_changed(cls, class_name: str, **kwargs) -> str:
        """
        Unified IS_CHANGED helper for ComfyUI nodes.

        Data-guard params (wired inputs that can't be cheaply hashed) should
        be checked inline before calling this method — see baker_node.py or
        checkpoint_merger_node.py for the recommended pattern.

        Usage:
            @classmethod
            def IS_CHANGED(cls, **kwargs):
                # Inline data guard FIRST
                if kwargs.get('lora_data_a') is not None:
                    return float("nan")
                # Then hash-based caching for the rest
                return NodeCache.is_changed(
                    cls.__name__,
                    lora_a=kwargs.get('lora_a', 'None'),
                    method=kwargs.get('method', 'linear'),
                )

        Args:
            class_name: cls.__name__ for namespacing the cache key.
            **kwargs: Named parameters to hash for change detection.

        Returns:
            str: MD5 hash of parameters for cache hit detection.
        """
        key_str = str((class_name, kwargs)).encode('utf-8')
        current_hash = hashlib.md5(key_str).hexdigest()

        cache_key = (class_name, tuple(sorted(kwargs.items())))
        if cache_key in cls._CACHE and cls._CACHE[cache_key] == current_hash:
            return cls._CACHE[cache_key]

        cls._CACHE[cache_key] = current_hash
        return current_hash


class ProgressTracker:
    """
    Wraps ComfyUI's ProgressBar with console fallback + ETA.

    - If comfy.utils.ProgressBar is available, uses the visual UI progress bar.
    - Otherwise, prints percentage updates to console at intervals.
    - Displays ETA for operations exceeding 5 seconds.
    - Thread-safe guard: single-threaded use within ComfyUI pipeline.

    Singleton awareness:
    - ComfyUI only supports ONE comfy.utils.ProgressBar instance at a time.
    - The class tracks _active_pbar so nested ProgressTrackers don't compete.
    - When a child tracker is created (via parent=), it delegates to the parent.
    - When a standalone tracker is created while one is already active,
      it uses console fallback without creating a competing UI bar.
    """
    # Class-level: tracks the ONE active ComfyUI ProgressBar instance.
    # Prevents nested ProgressTrackers from creating competing UI bars.
    _active_pbar = None

    def __init__(self, total: int, desc: str = "", update_interval: float = 0.5,
                 parent: 'ProgressTracker' = None):
        self.total = max(total, 1)
        self.desc = desc
        self.update_interval = update_interval
        self.current = 0
        self.start_time = time.time()
        self.last_print = 0.0
        self._pbar = None
        self._parent = parent
        self._owns_pbar = False

        if parent is not None:
            # Child tracker: delegate to parent, never create own ComfyUI bar
            self._pbar = None
        elif ProgressTracker._active_pbar is None:
            # Top-level tracker: create the ONE and only ComfyUI bar
            try:
                if hasattr(comfy.utils, 'ProgressBar'):
                    self._pbar = comfy.utils.ProgressBar(self.total)
                    ProgressTracker._active_pbar = self
                    self._owns_pbar = True
            except Exception:
                self._pbar = None
        else:
            # Another standalone tracker is already active — use console fallback
            # to avoid competing with the existing ComfyUI ProgressBar.
            self._pbar = None

    def update(self, n: int = 1):
        """Advance progress by n steps. Updates UI bar or prints to console."""
        self.current += n
        if self._parent is not None:
            # Delegate to parent tracker
            self._parent.update(n)
        elif self._pbar is not None:
            self._pbar.update(n)
        else:
            # Console fallback — print at intervals
            now = time.time()
            if now - self.last_print >= self.update_interval:
                pct = self.current / self.total * 100
                elapsed = now - self.start_time
                eta_str = ""
                if self.current > 0 and elapsed > 5:
                    rate = elapsed / self.current
                    remaining = (self.total - self.current) * rate
                    if remaining > 5:
                        eta_str = f" [ETA: {remaining:.0f}s]"
                print(f"\r{self.desc}: {self.current}/{self.total} ({pct:.0f}%){eta_str}", end="", flush=True)
                self.last_print = now

    def complete(self):
        """Mark progress as complete. Ensures final display."""
        if self._parent is not None:
            # Child: nothing to finalize — parent handles its own display
            return
        if self._pbar is None:
            elapsed = time.time() - self.start_time
            remaining = self.total - self.current
            if remaining > 0:
                self.current = self.total
            print(f"\r{self.desc}: {self.total}/{self.total} (100%) — {elapsed:.1f}s")
        else:
            remaining = self.total - self.current
            if remaining > 0:
                self._pbar.update(remaining)
                self.current = self.total
            # Release class-level ref so the next standalone tracker can create a bar
            if self._owns_pbar:
                ProgressTracker._active_pbar = None

    def __enter__(self):
        """Context manager entry — returns self for use as `with ProgressTracker(...) as p:`."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit — auto-completes progress on normal exit or exception."""
        self.complete()
        # Do not suppress exceptions
        return False

    def __iadd__(self, n: int):
        """Support `progress += 1` syntax as sugar for `progress.update(1)`."""
        self.update(n)
        return self


class NullProgressTracker:
    """Silent progress tracker — accepts all calls but does nothing.
    
    Used as a default parameter value in place of None to eliminate
    `if tracker is None` conditional branches throughout the codebase.
    Implements the same public interface as ProgressTracker.
    Stateless — all instances behave identically, safe as default param.
    """
    def update(self, n: int = 1) -> None:
        pass

    def __iadd__(self, n: int) -> 'NullProgressTracker':
        return self

    def __enter__(self) -> 'NullProgressTracker':
        return self

    def __exit__(self, *args) -> None:
        pass

def load_lora_with_metadata(path: Path):
    """Load LoRA with metadata and validation."""
    print(f"📥 Loading {path.name}...")
    
    # First try standard safetensors loading
    try:
        with safe_open(path, framework="pt") as f:
            tensors = {}
            failed_keys = []
            
            keys = list(f.keys())
            print(f"   Found {len(keys)} keys in safetensors")
            
            # Show sample of keys for debugging
            # Filter to show actual weight/bias keys instead of noise schedule params
            weight_keys = [k for k in keys if '.weight' in k or '.bias' in k]
            if weight_keys:
                print(f"   Sample keys: {weight_keys[:3]}")
            else:
                print(f"   Sample keys: {keys[:3]}")
            if len(keys) > 3:
                print(f"   ... and {len(keys) - 3} more")
            
            for k in keys:
                try:
                    tensors[k] = f.get_tensor(k)
                except Exception as e:
                    print(f"⚠️ Failed to load tensor {k}: {e}")
                    failed_keys.append(k)
            
            if failed_keys:
                print(f"⚠️ Failed to load {len(failed_keys)} keys")
            
            try:
                metadata = f.metadata() or {}
                if metadata:
                    print(f"   Metadata: {list(metadata.keys())[:5]}...")
            except Exception as e:
                print(f"⚠️ Failed to load metadata: {e}")
                metadata = {}
            
            print(f"✅ Loaded {len(tensors)} tensors from {path.name}")
            return tensors, metadata
    except Exception as e:
        print(f"⚠️ Failed to open {path} with safe_open: {e}")
        
        # Try fallback to load_file (handles more formats)
        try:
            print(f"🔧 Trying fallback loading for {path.name}...")
            tensors = load_file(str(path))
            print(f"   Loaded {len(tensors)} tensors via fallback")
            
            # Try to extract metadata if possible
            metadata = {}
            try:
                # Some LoRA files store metadata in special keys
                if '__metadata__' in tensors:
                    metadata = tensors['__metadata__']
                    del tensors['__metadata__']
                    print(f"   Found metadata with {len(metadata)} entries")
            except:
                pass
            
            return tensors, metadata
        except Exception as e2:
            # Last resort: try as PyTorch pickle file
            try:
                print(f"🔧 Trying PyTorch loading for {path.name}...")
                tensors = torch.load(str(path), map_location='cpu', weights_only=False)
                
                # Convert to dictionary if it's not already
                if not isinstance(tensors, dict):
                    print(f"❌ Unknown format in {path.name}")
                    raise ValueError(f"Unknown file format: {path.name}")
                
                print(f"✅ Loaded as PyTorch file: {len(tensors)} tensors")
                
                metadata = {}
                return tensors, metadata
            except Exception as e3:
                # All loading methods failed
                error_msg = f"Cannot load LoRA file {path.name}:\n"
                error_msg += f"1. safetensors error: {e}\n"
                error_msg += f"2. load_file error: {e2}\n"
                if 'weights_only' in str(e3):
                    error_msg += f"3. torch.load error: File may be unsafe (try enabling weights_only=False)\n"
                else:
                    error_msg += f"3. torch.load error: {e3}\n"
                raise ValueError(error_msg)




class DeviceManager:
    """Manages device and precision settings with optimal defaults."""

    # Minimum free VRAM (bytes) to prefer CUDA over CPU in "auto" mode
    VRAM_AUTO_THRESHOLD_BYTES: int = 2 * (1024**3)  # 2 GB

    @staticmethod
    def get_device(device_type: str = "auto", architecture: str = None) -> torch.device:
        """Get torch device with fallback logic.

        Args:
            device_type: ``"auto"``, ``"cuda"``, or ``"cpu"``.
            architecture: Optional model architecture name for architecture-aware VRAM
                threshold. If provided, uses the architecture-specific minimum VRAM
                instead of the flat 2 GB threshold.

        Returns:
            torch.device: Resolved device.
        """
        if device_type == "cuda" and torch.cuda.is_available():
            return torch.device("cuda")
        elif device_type == "auto":
            if torch.cuda.is_available():
                free_vram = DeviceManager.get_free_vram()
                if free_vram is not None:
                    # Use architecture-specific threshold if provided
                    threshold = (DeviceManager.get_vram_threshold_for_architecture(architecture)
                                 if architecture else DeviceManager.VRAM_AUTO_THRESHOLD_BYTES)
                    if free_vram < threshold:
                        arch_label = f" for {architecture}" if architecture else ""
                        print(f"⚠️ [DeviceManager] Only {free_vram / (1024**3):.2f} GB free VRAM "
                              f"(minimum {threshold / (1024**3):.1f} GB{arch_label}) — "
                              f"falling back to CPU")
                        return torch.device("cpu")
                return torch.device("cuda")
            return torch.device("cpu")
        else:
            if device_type == "cuda" and not torch.cuda.is_available():
                print("⚠️ [DeviceManager] CUDA explicitly selected but not available — falling back to CPU")
            return torch.device("cpu")

    @staticmethod
    def get_free_vram(device: Optional[torch.device] = None) -> Optional[int]:
        """Return free VRAM in bytes, or None if not CUDA / info unavailable."""
        if device is None:
            if not torch.cuda.is_available():
                return None
            device = torch.device("cuda")
        if device.type != "cuda":
            return None
        try:
            total = torch.cuda.get_device_properties(device).total_memory
            reserved = torch.cuda.memory_reserved(device)
            return total - reserved
        except Exception:
            return None

    @staticmethod
    def get_vram_threshold_for_architecture(architecture: str) -> int:
        """Return minimum free VRAM bytes required for given model architecture.

        Used by :meth:`get_device` when an architecture name is available.
        Thresholds account for: model weights on GPU + peak merge overhead + safety margin.

        Args:
            architecture: Model architecture name (e.g. ``"sdxl"``, ``"flux"``).

        Returns:
            int: Minimum free VRAM in bytes.
        """
        thresholds = {
            "sd15":     3 * (1024**3),   # 3 GB  — SD1.5: ~2 GB model + 1 GB overhead
            "sdxl":     8 * (1024**3),   # 8 GB  — SDXL: ~6 GB model + 2 GB overhead
            "flux":    16 * (1024**3),   # 16 GB — Flux: ~12 GB model + 4 GB overhead
            "z_image":  8 * (1024**3),   # 8 GB  — Z-Image: varies 4-8 GB
            "anima":    4 * (1024**3),   # 4 GB  — Anima: varies 2-4 GB
            "lumina2": 12 * (1024**3),   # 12 GB — Lumina2: ~8 GB model + 4 GB overhead
        }
        return thresholds.get(architecture, 6 * (1024**3))  # default 6 GB for unknown

    @staticmethod
    def suggest_batch_size(
        device: torch.device,
        architecture: Optional[str] = None,
    ) -> int:
        """Suggest an appropriate batch_size based on free VRAM and architecture.

        Args:
            device: Target torch device.
            architecture: Optional model architecture name for per-tensor size estimation.

        Returns:
            int: Suggested batch_size (clamped to 1-128).
        """
        if device.type != "cuda":
            return 64  # CPU: sequential is fine, batch matters less

        free_vram = DeviceManager.get_free_vram(device)
        if free_vram is None:
            return 32  # safe default when VRAM info unavailable

        # Estimate per-tensor memory (architecture-dependent)
        if architecture == "flux":
            per_tensor_estimate = 8 * 1024 * 1024  # 8 MB per tensor
        elif architecture in ("sdxl", "z_image"):
            per_tensor_estimate = 4 * 1024 * 1024  # 4 MB per tensor
        else:
            per_tensor_estimate = 2 * 1024 * 1024  # 2 MB per tensor (SD1.5 typical)

        # 3 sources + 1 output = 4× per batch during merge
        # Baker: 1 source + 1 delta = 2× per batch
        # Use 3× as a middle-ground estimate
        batch_overhead = per_tensor_estimate * 3
        max_batch = free_vram // batch_overhead if batch_overhead > 0 else 32

        # Clamp to reasonable range
        return max(1, min(max_batch, 128))
    
    @staticmethod
    def get_dtype(precision: str = "auto", device: torch.device = None) -> torch.dtype:
        """Get optimal dtype with FP8/NVFP4 support detection and fallback."""
        if device is None:
            device = DeviceManager.get_device()
        
        # Convert string device type to torch.device
        if isinstance(device, str):
            device = DeviceManager.get_device(device)
        
        if precision == "float32":
            return torch.float32
        elif precision == "bfloat16":
            if device.type == "cuda" and not torch.cuda.is_bf16_supported():
                print("⚠️ bfloat16 not supported on this GPU, falling back to float16")
                return torch.float16
            return torch.bfloat16
        elif precision == "float16":
            return torch.float16
        elif precision == "fp8_e4m3fn":
            result = _try_fp8_dtype('float8_e4m3fn', "fp8_e4m3fn", device)
            if result is not None:
                return result
            return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
        elif precision == "fp8_e5m2":
            result = _try_fp8_dtype('float8_e5m2', "fp8_e5m2", device)
            if result is not None:
                return result
            return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
        elif precision == "fp8":
            result = _try_fp8_dtype('float8_e4m3fn', "fp8", device)
            if result is not None:
                return result
            return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
        elif precision in ("int8", "int8_convrot"):
            # INT8 is universally supported on all CUDA GPUs (CC 6.1+)
            # and CPU. No capability checks needed.
            # int8_convrot uses the same dtype — ConvRot is a pre-processing
            # step before quantization, not a different storage format.
            return torch.int8
        elif precision == "device_optimized":
            # GPU-aware cascade: best precision THIS GPU can run natively.
            # Same logic as the original old_gpu — optimized for the LOCAL machine.
            if device.type == "cuda":
                cap = torch.cuda.get_device_capability(device)
                if cap[0] >= 9:          # Hopper+: native FP8
                    return torch.float8_e4m3fn
                elif cap[0] >= 8:        # Ampere+: native BF16
                    return torch.bfloat16
                else:                     # Pascal/Turing CC 6.1-7.5: INT8
                    return torch.int8
            else:
                return torch.float32
        elif precision == "old_gpu":
            # Target-optimized for old GPUs (Pascal/Turing CC 6.1-7.5).
            # Always produces INT8 (1 byte/param) with FP16 dequant —
            # regardless of the LOCAL machine's GPU.
            # This is the CORRECT option for creators who want to
            # share models on CivitAI for old-GPU users.
            return torch.int8
        elif precision == "svd_only":
            # SVD-only mode: apply SVD compression to float32, skip quantization.
            # Returns float32 tensors with reduced rank (smaller file size).
            # No companion scales, no quantization metadata needed.
            return torch.float32
        else:  # "auto"
            if device.type == "cuda":
                # Auto-select best precision
                # Default to bfloat16 for supported GPUs, otherwise float16
                if torch.cuda.is_bf16_supported():
                    return torch.bfloat16
                else:
                    return torch.float16
            else:
                return torch.float32
    
    @staticmethod
    def get_fallback_dtype(current_dtype: torch.dtype) -> torch.dtype:
        """Get a fallback dtype when current dtype operations fail."""
        if current_dtype == torch.float8_e4m3fn:
            # FP8 failed, try bfloat16
            if torch.cuda.is_bf16_supported():
                print("🔄 FP8 (e4m3fn) operation failed, falling back to bfloat16")
                return torch.bfloat16
            else:
                print("🔄 FP8 (e4m3fn) operation failed, falling back to float16")
                return torch.float16
        elif current_dtype == torch.float8_e5m2:
            # FP8 failed, try bfloat16
            if torch.cuda.is_bf16_supported():
                print("🔄 FP8 (e5m2) operation failed, falling back to bfloat16")
                return torch.bfloat16
            else:
                print("🔄 FP8 (e5m2) operation failed, falling back to float16")
                return torch.float16
        elif current_dtype == torch.bfloat16:
            # bfloat16 failed, try float16
            print("🔄 bfloat16 operation failed, falling back to float16")
            return torch.float16
        else:
            # float16 failed, try float32
            print("🔄 float16 operation failed, falling back to float32")
            return torch.float32


def _try_fp8_dtype(dtype_attr: str, label: str, device: torch.device) -> Optional[torch.dtype]:
    """Attempt to return an FP8 dtype if hardware supports it, otherwise None."""
    if device.type == "cuda" and hasattr(torch, dtype_attr):
        device_cap = torch.cuda.get_device_capability(device)
        # FP8 requires compute capability 8.9+ (Hopper)
        if device_cap[0] >= 9 or (device_cap[0] == 8 and device_cap[1] >= 9):
            print(f"⚠️ {label} selected but operations may be limited")
            print("   Will attempt FP8 but fall back to bfloat16 if operations fail")
            return getattr(torch, dtype_attr)
        else:
            # BF16 requires CC 8.0+ (Ampere); on older GPUs the caller will fall back to FP16
            _fallback = "bfloat16" if device_cap[0] >= 8 else "float16"
            print(f"⚠️ {label} not supported on device (compute {device_cap}), falling back to {_fallback}")
    else:
        print(f"⚠️ {label} not available in this PyTorch version, falling back to bfloat16")
    return None


@contextmanager
def memory_optimized_merge():
    """Context manager for memory-optimized merging."""
    original_state = {
        "grad_enabled": torch.is_grad_enabled(),
        "cudnn_benchmark": torch.backends.cudnn.benchmark,
    }
    
    try:
        # Disable gradients for merging
        torch.set_grad_enabled(False)
        
        # Enable cudnn benchmarking
        torch.backends.cudnn.benchmark = True
        
        yield
    finally:
        # Restore original state
        torch.set_grad_enabled(original_state["grad_enabled"])
        torch.backends.cudnn.benchmark = original_state["cudnn_benchmark"]
        
        # Clean up GPU memory
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


def memory_guard():
    """
    Check CUDA VRAM usage and empty cache if above 80% threshold.
    
    Called between batches during baking / checkpoint merging to prevent
    OOM from accumulated intermediate tensors.
    Safe to call on CPU (no-op when CUDA is unavailable).
    """
    if torch.cuda.is_available():
        try:
            allocated = torch.cuda.memory_allocated()
            total = torch.cuda.get_device_properties(0).total_memory
            if total > 0 and allocated / total > 0.8:
                torch.cuda.empty_cache()
        except Exception:
            pass


def get_available_ram() -> Optional[int]:
    """
    Returns available physical RAM in bytes, or None if detection fails.
    Tries psutil first (cross-platform), falls back to Windows ctypes.
    """
    try:
        import psutil
        return psutil.virtual_memory().available
    except ImportError:
        pass

    try:
        # Windows fallback via ctypes kernel32.GlobalMemoryStatusEx
        import ctypes
        kernel32 = ctypes.windll.kernel32
        class MEMORYSTATUSEX(ctypes.Structure):
            _fields_ = [
                ("dwLength", ctypes.c_ulong),
                ("dwMemoryLoad", ctypes.c_ulong),
                ("ullTotalPhys", ctypes.c_ulonglong),
                ("ullAvailPhys", ctypes.c_ulonglong),
                ("ullTotalPageFile", ctypes.c_ulonglong),
                ("ullAvailPageFile", ctypes.c_ulonglong),
                ("ullTotalVirtual", ctypes.c_ulonglong),
                ("ullAvailVirtual", ctypes.c_ulonglong),
                ("ullAvailExtendedVirtual", ctypes.c_ulonglong),
            ]
        memoryStatus = MEMORYSTATUSEX()
        memoryStatus.dwLength = ctypes.sizeof(MEMORYSTATUSEX)
        kernel32.GlobalMemoryStatusEx(ctypes.byref(memoryStatus))
        return memoryStatus.ullAvailPhys
    except Exception:
        pass

    return None  # Detection unavailable


def estimate_ram_mode_peak(total_data_size: int, num_tensors: int, batch_size: int, use_dict: bool = False) -> int:
    """
    Estimate peak memory for in-memory modes.
    
    Peak depends on which output mode is active:
    - Dict mode (use_dict=True): ~1.5x total_data_size
      No BytesIO/clone overhead — state dict IS the merge output.
      During weave: output dict grows to ~1x total_data_size (system RAM).
      During model loading: output dict (~1x) partially overlaps with
      model objects (~1x GPU, negligible system RAM impact via get_available_ram).
      Conservative peak: 1.5x accounts for brief overlap during ComfyUI loading
      when state dict is partially consumed and model objects partially built.
      When precision=bf16/fp16, dict uses ~half the fp32 memory.
      Phase 3 (zero-weight auto-omit) further reduces dict size proportionally.
    - Batch processing during weave: negligible vs parse/load phases
    
    We conservatively estimate 1.5x for dict mode,
    plus batch overhead. Dict mode users get ~25% lower estimate automatically
    (more with low-precision or zero-weight omission), reducing unnecessary
    fallbacks to disk mode on moderately RAM-constrained systems.
    """
    # Batch processing overhead during weave
    batch_fraction = min(batch_size / max(num_tensors, 1), 1.0)
    batch_memory = int(total_data_size * batch_fraction * 0.5)
    
    if use_dict:
        # Dict mode: state dict IS the merge output (no BytesIO copy).
        # Model objects use GPU VRAM, not system RAM (measured by get_available_ram).
        # During ComfyUI loading, ~0.5x state_dict + ~1x model objects overlap briefly
        # on CPU, but model objects are GPU-resident so system RAM impact is minimal.
        # Conservative: 1.5x accounts for partial system RAM overlap.
        operation_peak = int(total_data_size * 1.5)
    else:
        # Streaming mode (no BytesIO buffer): state_dict (~1x) + model objects (~1x)
        # can briefly coexist during parse+load on system RAM. Conservative: 2x.
        operation_peak = int(total_data_size * 2.0)
    
    return operation_peak + batch_memory


def check_ram_guard(
    total_bytes: int,
    num_tensors: int,
    batch_size: int,
    label: str = "operation",
    threshold_ratio: Optional[float] = None,
    use_dict: bool = False,
) -> bool:
    """
    Check if RAM is sufficient for an in-memory operation.

    Returns True if safe (within threshold), False if fallback recommended.

    Args:
        total_bytes: Total size of data to load (in bytes)
        num_tensors: Number of tensors (for peak estimation)
        batch_size: Batch size for peak estimation
        label: Human-readable label for diagnostic messages
        threshold_ratio: Fraction of available RAM to use as threshold
                         (default sourced from config.RAM_GUARD_THRESHOLD = 0.85)
        use_dict: If True, use dict-mode peak estimate (~1.5x) instead of
                  legacy buffer mode (~2x). Dict mode avoids the BytesIO
                  serialization spike, lowering system RAM peak.
    """
    if threshold_ratio is None:
        from .config import RAM_GUARD_THRESHOLD
        threshold_ratio = RAM_GUARD_THRESHOLD

    available_ram = get_available_ram()
    if available_ram is None:
        print(f"⚠️ RAM Guard ({label}): Cannot detect available RAM — proceeding without guard")
        return True  # Can't determine, assume safe

    peak_estimate = estimate_ram_mode_peak(total_bytes, num_tensors, batch_size, use_dict=use_dict)
    threshold_ram = int(available_ram * threshold_ratio)

    print(f"🧠 RAM Guard ({label}): Estimated peak ~{peak_estimate / (1024**3):.2f} GB, "
          f"Available RAM ~{available_ram / (1024**3):.2f} GB, "
          f"{threshold_ratio*100:.0f}% threshold ~{threshold_ram / (1024**3):.2f} GB")

    if peak_estimate > threshold_ram:
        print(f"   ⚠️ RAM Guard ({label}): Peak estimate exceeds {threshold_ratio*100:.0f}% threshold — fallback recommended")
        return False
    else:
        print(f"✅ RAM Guard ({label}): Peak estimate within safe limits")
        return True


def print_ram_delta(
    pre_ram: Optional[int],
    post_ram: Optional[int],
    label: str = "load"
) -> None:
    """Print pre/post RAM delta with consistent formatting emoji prefix.

    Baker reference: baker_node.py:450 + baker_node.py:607-614.
    """
    if pre_ram is not None and post_ram is not None:
        delta_gb = (pre_ram - post_ram) / (1024**3)
        print(f"📊 RAM: Pre-{label} {pre_ram/(1024**3):.2f} GB \u2192 "
              f"Post-{label} {post_ram/(1024**3):.2f} GB "
              f"(delta: {delta_gb:+.2f} GB)")
    elif post_ram is not None:
        print(f"📊 Post-{label} available RAM: {post_ram/(1024**3):.2f} GB")


def is_lazy_mapping(obj) -> bool:
    """Detect if object is a _LazyCheckpointMapping (mmap-backed).

    Used to skip RAM Guard for mmap-backed mappings.
    Baker reference: baker_node.py:523, baking_processor.py:278, 691.
    """
    return hasattr(obj, 'filepath')


def estimate_disk_mode_peak(total_data_size: int, num_tensors: int, batch_size: int) -> int:
    """
    Estimate peak memory for disk mode.
    - Streamed to file (disk buffers, negligible RAM for output)
    - Batch processing: batch_size/num_tensors * total_data_size
    - Lazy loading: negligible until ComfyUI access
    """
    batch_fraction = min(batch_size / max(num_tensors, 1), 1.0)
    return int(total_data_size * batch_fraction * 1.2)  # 20% overhead


def get_free_disk_space(path: Path) -> Optional[int]:
    """
    Returns free disk space in bytes on the drive containing `path`,
    or None if detection fails.
    """
    try:
        if hasattr(Path, 'statvfs'):  # Unix-like
            stat = path.statvfs()
            return stat.f_frsize * stat.f_bavail
        else:  # Windows
            import ctypes
            free_bytes = ctypes.c_ulonglong(0)
            ctypes.windll.kernel32.GetDiskFreeSpaceExW(
                ctypes.c_wchar_t(str(path.absolute())),
                None, None, ctypes.pointer(free_bytes)
            )
            return free_bytes.value
    except Exception:
        return None

class ThreadSafeCleanup:
    """Thread-safe temporary file cleanup."""
    _lock = threading.Lock()
    _registered_files = set()
    
    @classmethod
    def cleanup_temp_directory(cls):
        """Main cleanup function for all temp files with thread safety."""
        with cls._lock:
            temp_dir = Path(folder_paths.get_temp_directory())
            
            # 1. Clean easy_lora folder (keep last 10)
            easy_lora_dir = temp_dir / "easy_lora"
            if easy_lora_dir.exists():
                files = list(easy_lora_dir.glob("*.safetensors"))
                if len(files) > 10:
                    files.sort(key=lambda x: x.stat().st_mtime)
                    for old_file in files[:-10]:
                        try:
                            old_file.unlink()
                        except Exception as e:
                            print(f"⚠️ Failed to delete {old_file}: {e}")
            
            # 2. Clean temp_* files from LORA data conversion (keep last 5)
            temp_files = list(temp_dir.glob("temp_*.safetensors"))
            if len(temp_files) > 5:
                temp_files.sort(key=lambda x: x.stat().st_mtime)
                for old_file in temp_files[:-5]:
                    try:
                        old_file.unlink()
                    except Exception as e:
                        print(f"⚠️ Failed to delete {old_file}: {e}")
            
            # 3. Clean checkpoint_temp_* files from chained data (keep last 5)
            ckpt_temp_files = list(temp_dir.glob("checkpoint_temp_*.safetensors"))
            if len(ckpt_temp_files) > 5:
                ckpt_temp_files.sort(key=lambda x: x.stat().st_mtime)
                for old_file in ckpt_temp_files[:-5]:
                    try:
                        old_file.unlink()
                    except Exception as e:
                        print(f"⚠️ Failed to delete {old_file}: {e}")
            
            # 4. Clean old cache files (older than 1 hour)
            cache_files = list(temp_dir.glob("*.cache"))
            current_time = time.time()
            for cache_file in cache_files:
                try:
                    if current_time - cache_file.stat().st_mtime > 3600:
                        cache_file.unlink()
                except Exception:
                    pass

    @classmethod
    def register_temp_file(cls, path):
        """Register a temporary file for later cleanup."""
        with cls._lock:
            cls._registered_files.add(Path(path))


def get_experiment_temp_path(node_type: str = "main") -> Path:
    """Get a temp file path AND auto-register it with ThreadSafeCleanup.

    Every temp file created through this function is automatically
    registered for cleanup — callers should NOT register it again.
    Uses the ``easy_lora`` subdirectory under ComfyUI's temp directory.
    """
    temp_dir = Path(folder_paths.get_temp_directory()) / "easy_lora"
    temp_dir.mkdir(parents=True, exist_ok=True)
    
    timestamp = int(time.time())
    random_str = uuid.uuid4().hex[:6]
    
    if node_type == "main":
        filename = f"experiment_main_{timestamp}_{random_str}.safetensors"
    else:
        filename = f"experiment_only_{timestamp}_{random_str}.safetensors"
    
    path = temp_dir / filename
    ThreadSafeCleanup.register_temp_file(path)
    return path


def _get_output_path(
    save_folder: str,
    filename: str,
    default_name: str,
    folder_type: str,
) -> Path:
    """Shared output path resolution with auto-increment for LoRAs and checkpoints."""
    # Default filename
    if not filename or filename.strip() == "":
        filename = default_name

    # Ensure .safetensors extension
    if not filename.endswith(".safetensors"):
        filename += ".safetensors"

    # Clean filename
    safe_filename = ''.join(c for c in filename
                           if c.isalnum() or c in '._- ').rstrip()

    # Get default base folder
    folders = folder_paths.get_folder_paths(folder_type)
    default_base = Path(folders[0]) if folders else Path.cwd()

    # Resolve output folder
    if save_folder and isinstance(save_folder, str) and save_folder.strip():
        user_path = save_folder.strip()
        print(f"📁 User requested save folder: '{user_path}'")

        is_windows_absolute = (len(user_path) >= 2 and
                              user_path[1] == ':' and
                              user_path[2] == '\\')
        is_network_path = user_path.startswith('\\\\')

        if is_windows_absolute or is_network_path:
            try:
                output_folder = Path(user_path)
                if '..' in user_path or user_path.count(':\\') > 1:
                    print(f"⚠️ Suspicious path pattern, using default")
                    output_folder = default_base
                else:
                    output_folder.mkdir(parents=True, exist_ok=True)
                    print(f"✅ Using absolute path: {output_folder}")
            except Exception as e:
                print(f"⚠️ Could not use absolute path '{user_path}': {e}")
                output_folder = default_base
        else:
            try:
                clean_path = user_path.strip('\\/')
                output_folder = default_base / clean_path
                output_folder.mkdir(parents=True, exist_ok=True)
                print(f"✅ Using relative path: {output_folder}")
            except Exception as e:
                print(f"⚠️ Could not create subfolder '{user_path}': {e}")
                output_folder = default_base
    else:
        output_folder = default_base

    # Auto-increment logic
    base_path = output_folder / safe_filename
    counter = 1
    final_path = base_path
    while final_path.exists():
        stem = base_path.stem
        parts = stem.split('_')
        if len(parts) > 1 and parts[-1].isdigit():
            stem = '_'.join(parts[:-1])
        final_path = base_path.parent / f"{stem}_{counter}{base_path.suffix}"
        counter += 1

    print(f"💾 Final save path: {final_path}")
    return final_path


def get_user_output_path(save_folder: str, filename: str) -> Path:
    """Get output path for a merged LoRA with auto-increment."""
    return _get_output_path(save_folder, filename, "merged_lora", "loras")


def get_checkpoint_output_path(save_folder: str, filename: str) -> Path:
    """Get output path for a merged checkpoint with auto-increment."""
    return _get_output_path(save_folder, filename, "merged_checkpoint", "checkpoints")


# ── Combined model list + path resolution for UNet/checkpoint support ──
# These functions allow nodes to show BOTH checkpoints and diffusion models
# (UNets) in their dropdown lists, resolving paths from either folder type.
# Diffusion models from models/unet/ and models/diffusion_models/ are appended
# to the checkpoint list, with deduplication. When resolving paths, checkpoints
# are tried first, then diffusion_models as fallback.

def get_combined_model_list() -> List[str]:
    """Return a combined, deduplicated list of filenames from both
    ``checkpoints`` and ``diffusion_models`` ComfyUI folder types.

    Checkpoint files come first; diffusion model files that are not already
    in the checkpoint list are appended at the end.  This prevents duplicate
    entries when the same filename exists in both ``models/checkpoints/`` and
    ``models/unet/``.

    Returns:
        List of filename strings suitable for use in ComfyUI dropdowns.
    """
    checkpoints = folder_paths.get_filename_list("checkpoints")
    diffusion_models = folder_paths.get_filename_list("diffusion_models")
    seen = set(checkpoints)
    for model in diffusion_models:
        if model not in seen:
            checkpoints.append(model)
            seen.add(model)
    return checkpoints


def resolve_model_path(name: str) -> Optional[str]:
    """Resolve a model filename to its full filesystem path.

    Tries the ``checkpoints`` folder first, then falls back to
    ``diffusion_models``.  This gives priority to checkpoint paths when
    the same filename exists in both directories.

    Args:
        name: Filename (e.g. ``"my_model.safetensors"``) or ``"None"``.

    Returns:
        Full path string if found in either folder, or ``None`` if not found.
    """
    if not name or name == "None":
        return None
    path = folder_paths.get_full_path("checkpoints", name)
    if path is not None:
        return path
    return folder_paths.get_full_path("diffusion_models", name)


# ── Streaming safetensors writer (avoids BytesIO 2× memory spike) ──

_DTYPE_SIZE = {
    "F64": 8, "F32": 4, "F16": 2, "BF16": 2,
    "I64": 8, "I32": 4, "I16": 2, "I8": 1, "U8": 1,
    "F8_E4M3": 1, "F8_E5M2": 1, "BOOL": 1,
}

DTYPE_REVERSE_STREAM = {
    torch.float64: "F64", torch.float32: "F32",
    torch.float16: "F16", torch.bfloat16: "BF16",
    torch.int64: "I64",   torch.int32: "I32",
    torch.int16: "I16",   torch.int8: "I8",
    torch.uint8: "U8",
    torch.float8_e4m3fn: "F8_E4M3",
    torch.float8_e5m2: "F8_E5M2",
    torch.bool: "BOOL",
}


def _tensor_to_bytes(tensor: torch.Tensor) -> bytes:
    """Convert a tensor to raw bytes for safetensors storage.

    Automatically moves the tensor to CPU first — safetensors always
    stores CPU bytes. Supports all dtypes including bfloat16 and fp8
    (which numpy cannot represent natively) via .view(torch.uint8)
    reinterpretation. Handles scalar tensors (dim=0) for the .view() path.
    """
    # Defensive CPU transfer — safetensors stores CPU bytes only
    tensor = tensor.cpu()

    # Defensive contiguity check — prevents cryptic numpy errors
    if not tensor.is_contiguous():
        tensor = tensor.contiguous()

    # Dtypes needing .view(torch.uint8) workaround:
    # bfloat16, fp8_e4m3fn, fp8_e5m2 — numpy has no native support
    if tensor.dtype in (torch.bfloat16, torch.float8_e4m3fn, torch.float8_e5m2):
        if tensor.dim() == 0:
            tensor = tensor.reshape(-1)  # view() requires at least 1 dim
        return tensor.view(torch.uint8).numpy().tobytes()

    return tensor.numpy().tobytes()


def save_safetensors_stream(
    tensors: Dict[str, torch.Tensor],
    filename: Union[str, Path],
    metadata: Optional[Dict[str, str]] = None,
) -> None:
    """
    Write safetensors file by streaming tensors directly to disk,
    avoiding the BytesIO intermediate buffer used by save_file().

    Peak memory: state_dict only (no 2× BytesIO spike).

    Args:
        tensors: Dict of key -> contiguous CPU tensor.
        filename: Output file path.
        metadata: Optional metadata dict (must be JSON-serializable).

    Raises:
        RuntimeError: On write failure after partial file creation.
    """
    # Phase 1: Build header and compute cumulative offsets
    # NOTE: BFL original files store tensors contiguously with NO padding.
    # We match that layout exactly — no 8-byte alignment padding.
    sorted_keys = sorted(tensors.keys())
    header: Dict[str, Any] = {"__metadata__": metadata or {}}
    offset = 0
    for key in sorted_keys:
        tensor = tensors[key]
        dtype_str = DTYPE_REVERSE_STREAM.get(tensor.dtype, "F32")
        shape = list(tensor.shape)
        numel = tensor.numel()
        elem_size = _DTYPE_SIZE.get(dtype_str, 4)
        data_size = numel * elem_size
        header[key] = {
            "dtype": dtype_str,
            "shape": shape,
            "data_offsets": [offset, offset + data_size],
        }
        offset += data_size  # no padding — matches BFL's contiguous layout

    # Phase 2: Serialize to file, one tensor at a time
    header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8")
    header_size = struct.pack("<Q", len(header_bytes))

    try:
        with open(filename, "wb") as f:
            f.write(header_size)
            f.write(header_bytes)
            for key in sorted_keys:
                tensor = tensors[key]
                raw_bytes = _tensor_to_bytes(tensor)
                f.write(raw_bytes)
                # No padding — tensors stored contiguously (BFL-compatible)

        # Phase 3: Validate file with safe_open — catches corruption immediately
        try:
            with safe_open(filename, framework="pt") as sf:
                for key in sf.keys():
                    _ = sf.get_tensor(key)  # triggers offset validation
        except Exception as e:
            # Validation failed — delete corrupted file and re-raise
            try:
                Path(filename).unlink(missing_ok=True)
            except Exception:
                pass
            raise RuntimeError(
                f"Post-save validation failed for {filename}: {e}"
            ) from e
    except BaseException:
        # Clean up partial file — catch BaseException to handle
        # KeyboardInterrupt (Ctrl+C) during long writes
        try:
            Path(filename).unlink(missing_ok=True)
        except Exception:
            pass
        raise


def save_safetensors_file(
    state_dict: Dict[str, torch.Tensor],
    save_folder: str,
    filename: str,
    metadata: Optional[Dict[str, str]] = None,
    folder_type: str = "loras",
    default_name: str = "output",
) -> Tuple[Optional[Path], str]:
    """
    Unified save utility for safetensors files with error handling.

    Resolves the output path via _get_output_path, creates parent dirs,
    attempts save_file, and returns the result.

    Args:
        state_dict: The tensor dict to save.
        save_folder: User-specified folder override (empty = default).
        filename: Desired filename (without extension or with .safetensors).
        metadata: Optional metadata dict to embed in the safetensors file.
        folder_type: ComfyUI folder type ('loras', 'checkpoints', etc.).
        default_name: Fallback name if filename is empty.

    Returns:
        Tuple of (output_path or None on failure, path_string or error message).
    """
    output_path = _get_output_path(save_folder, filename, default_name, folder_type)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    from safetensors.torch import save_file
    try:
        save_file(state_dict, str(output_path), metadata=metadata)
        file_size_mb = output_path.stat().st_size / (1024 * 1024)
        print(f"   ✅ Saved: {output_path} ({file_size_mb:.1f} MB)")
        return output_path, str(output_path)
    except Exception as e:
        error_msg = f"❌ Failed to save {default_name}: {e}"
        print(error_msg)
        return None, error_msg


def safe_get_rank(tensor: torch.Tensor, key: str) -> int:
    """Safely get rank from any LoRA tensor."""
    if len(tensor.shape) < 2:
        return 1  # Alpha or bias

    key_lower = key.lower()

    # lora_A types
    if any(x in key_lower for x in ['lora_a', 'lora.down', 'lora_down']):
        return tensor.shape[0]  # [rank, in_dim]

    # lora_B types
    elif any(x in key_lower for x in ['lora_b', 'lora.up', 'lora_up']):
        if len(tensor.shape) > 1:
            return tensor.shape[1]
        else:
            return tensor.shape[0]

    # Default fallback
    else:
        return min(tensor.shape[0], tensor.shape[1])


def categorize_key(key: str) -> str:
    """
    Categorize a key into 'model', 'te', or 'other'.

    **Note:** This is a backward-compatible wrapper that delegates to
    :func:`engine.key_utils.categorize_key` but maps ``'unet'`` → ``'model'``
    to preserve the original API (callers expect ``'model'`` rather than
    ``'unet'``).
    """
    from .engine.key_utils import categorize_key as _cat
    comp = _cat(key)
    # backward compat: old API returned 'model' not 'unet'
    if comp == 'unet':
        return 'model'
    return comp


def get_tensor_energy(tensor: torch.Tensor, mode: str) -> float:
    """
    Compute the energy (RMS or percentile) of a tensor.
    Modes: "none" (returns 1.0), "rms", "top_5%", "top_10%", "top_20%", "top_30%".
    """
    if mode == "none":
        return 1.0
    abs_tensor = torch.abs(tensor)
    if mode == "rms":
        return torch.sqrt(torch.mean(abs_tensor ** 2))
    # Parse percentile from mode string, e.g., "top_5%" -> quantile = 0.95
    if mode.startswith("top_"):
        try:
            percent = float(mode.split('_')[1].rstrip('%'))
            quantile = 1.0 - (percent / 100.0)
            # Flatten tensor for quantile computation
            flat = abs_tensor.flatten().float()
            return torch.quantile(flat, quantile)
        except (ValueError, IndexError):
            raise ValueError(f"Invalid percentile mode: {mode}")
    raise ValueError(f"Unknown magnitude scaling mode: {mode}")


def silent_pad_or_truncate(tensor: torch.Tensor, target_rank: int, key: str) -> torch.Tensor:
    """Pad or truncate tensor to target rank."""
    current_rank = safe_get_rank(tensor, key)

    if current_rank == target_rank:
        return tensor

    # Determine if this is lora_A or lora_B
    is_lora_b = any(x in key.lower() for x in ['lora_b', 'lora.up', 'lora_up'])

    if len(tensor.shape) <= 1:
        # Scalar or 1D tensor (alpha?) - no change
        return tensor

    if not is_lora_b:
        # lora_A: rank is axis 0
        new_shape = (target_rank,) + tensor.shape[1:]
        new_tensor = torch.zeros(new_shape, device=tensor.device, dtype=tensor.dtype)
        min_rank = min(current_rank, target_rank)
        # Slice along axis 0
        new_tensor[:min_rank, ...] = tensor[:min_rank, ...]
    else:
        # lora_B: rank is axis 1
        new_shape = (tensor.shape[0], target_rank) + tensor.shape[2:]
        new_tensor = torch.zeros(new_shape, device=tensor.device, dtype=tensor.dtype)
        min_rank = min(current_rank, target_rank)
        # Slice along axis 1
        new_tensor[:, :min_rank, ...] = tensor[:, :min_rank, ...]

    return new_tensor


def bake_alphas(state_dict, naming_style='lora_a_b'):
    """
    Bake alpha values into weight tensors and delete .alpha keys.
    Applies sqrt(alpha/rank) scaling to both down and up weight tensors.
    
    Delegates tensor scaling to :func:`engine.scale_utils.bake_alpha_key`.
    
    naming_style: either 'lora_a_b' (lora_A.weight, lora_B.weight) or
                  'lora_down_up' (lora_down.weight, lora_up.weight).
    """
    from .engine.scale_utils import resolve_scale, bake_alpha_key
    
    # Determine suffixes for down and up weight keys
    if naming_style == 'lora_a_b':
        down_suffix = '.lora_A.weight'
        up_suffix = '.lora_B.weight'
    else:
        down_suffix = '.lora_down.weight'
        up_suffix = '.lora_up.weight'
    
    # Collect alpha keys to delete
    alpha_keys = [k for k in state_dict.keys() if k.endswith('.alpha')]
    if not alpha_keys:
        # No alpha keys, nothing to bake
        return state_dict
    
    print(f"🧪 Baking {len(alpha_keys)} alpha keys...")
    scaled_count = 0
    total_scale = 0.0
    mismatch_count = 0
    mismatch_details = []
    
    # Iterate over alpha keys
    for alpha_key in alpha_keys:
        # Derive weight keys
        base = alpha_key[:-6]  # remove '.alpha'
        down_key = base + down_suffix
        up_key = base + up_suffix
        
        # Ensure both keys exist
        if down_key not in state_dict or up_key not in state_dict:
            print(f"⚠️ Missing weight tensors for alpha key {alpha_key}, skipping")
            continue
        
        down_tensor = state_dict[down_key]
        up_tensor = state_dict[up_key]
        alpha_tensor = state_dict[alpha_key]
        # Alpha value (scalar)
        if isinstance(alpha_tensor, torch.Tensor):
            alpha_val = alpha_tensor.item()
        else:
            alpha_val = float(alpha_tensor)
        
        # Rank from weight shapes (min dimension)
        rank = safe_get_rank(down_tensor, down_key)
        if rank == 0:
            rank = 1  # safety
        
        # Scaling factor and tensor baking via shared scale_utils
        scale = resolve_scale(alpha_val, rank, mode="sqrt")
        orig_down_norm = torch.norm(down_tensor).item()
        orig_up_norm = torch.norm(up_tensor).item()
        state_dict[down_key], state_dict[up_key] = bake_alpha_key(
            down_tensor, up_tensor, alpha_val, rank
        )
        actual_down_scale = torch.norm(state_dict[down_key]).item() / orig_down_norm
        actual_up_scale = torch.norm(state_dict[up_key]).item() / orig_up_norm
        
        # Update statistics
        total_scale += scale
        scaled_count += 1
        
        # Check for significant mismatch (relative tolerance 1%)
        tolerance = 0.01
        rel_down = abs(actual_down_scale - scale) / scale
        rel_up = abs(actual_up_scale - scale) / scale
        if rel_down > tolerance or rel_up > tolerance:
            mismatch_count += 1
            mismatch_details.append((alpha_key, scale, actual_down_scale, actual_up_scale, alpha_val, rank))
    
    # Delete alpha keys
    for alpha_key in alpha_keys:
        if alpha_key in state_dict:
            del state_dict[alpha_key]
    
    # Summary
    if scaled_count > 0:
        avg_scale = total_scale / scaled_count
        if abs(avg_scale - 1.0) > 1e-6:
            print(f"   Average scale detected: {avg_scale:.6f}")
    
    if mismatch_count > 0:
        print(f"   ⚠️  {mismatch_count} keys exceeded 1% tolerance")
        # Print up to 5 mismatches for debugging
        for i, (alpha_key, scale, down_act, up_act, alpha_val, rank) in enumerate(mismatch_details[:5]):
            print(f"      {alpha_key}: expected {scale:.6f}, down actual {down_act:.6f}, up actual {up_act:.6f} (alpha={alpha_val}, rank={rank})")
        if mismatch_count > 5:
            print(f"      ... and {mismatch_count - 5} more")
    # No mismatches: remain silent (no extra line)
    
    print(f"✅ Baked {len(alpha_keys)} alpha keys ({scaled_count} scaled)")
    return state_dict


def categorize_checkpoint_key(key: str) -> str:
    """
    Categorize a checkpoint tensor key into component: 'unet', 'clip', 'vae', 'te', or 'other'.

    Delegates to :func:`engine.key_utils.categorize_key` with the same return
    values (``'unet'``, ``'clip'``, ``'vae'``, ``'te'``, ``'other'``).
    """
    from .engine.key_utils import categorize_key as _cat
    return _cat(key)


def load_state_dict_as_model_objects(
    state_dict,
    metadata=None,
    output_vae=True,
    output_clip=True,
) -> Tuple[Optional[Any], Optional[Any], Optional[Any]]:
    """
    Convert a state dict or lazy mapping to ComfyUI MODEL, CLIP, VAE objects.

    Returns (model, clip, vae).
    NOTE: ComfyUI pops keys from the input state_dict. Use .copy() if preserving the original.
    """
    model, clip, vae, _ = comfy.sd.load_state_dict_guess_config(
        state_dict,
        output_vae=output_vae,
        output_clip=output_clip,
        output_clipvision=False,
        embedding_directory=None,
        model_options={},
        te_model_options={},
        metadata=metadata,
        disable_dynamic=False,
    )
    return model, clip, vae


def load_checkpoint_with_metadata(path: Path, categorize: bool = False):
    """
    Load checkpoint tensors and metadata from a safetensors file.
    Returns (tensors, metadata) if categorize=False.
    If categorize=True, returns (tensors, metadata, component_counts) where
    component_counts is a dict mapping component to number of keys.
    """
    tensors, metadata = load_lora_with_metadata(path)
    if not categorize:
        return tensors, metadata
    # Compute component counts
    counts = {'unet': 0, 'clip': 0, 'vae': 0, 'te': 0, 'other': 0}
    for key in tensors.keys():
        comp = categorize_checkpoint_key(key)
        counts[comp] += 1
    return tensors, metadata, counts


def apply_component_scaling(tensor: torch.Tensor, key: str,
                            weight_unet: float = 1.0,
                            weight_clip: float = 1.0,
                            weight_vae: float = 1.0,
                            weight_te: float = 1.0) -> torch.Tensor:
    """
    Scale a tensor according to its component category.
    Used for Weight Block Map scaling before merging.
    """
    comp = categorize_checkpoint_key(key)
    if comp == 'unet':
        scale = weight_unet
    elif comp == 'clip':
        scale = weight_clip
    elif comp == 'vae':
        scale = weight_vae
    elif comp == 'te':
        scale = weight_te
    else:
        # 'other' components: no scaling (or default 1.0)
        scale = 1.0
    if scale != 1.0:
        tensor = tensor * scale
    return tensor




def compute_component_energy_ratios(
    norm_sds: List[Dict[str, torch.Tensor]],
    common_keys: List[str],
    original_sds: Optional[List[Optional[Dict[str, torch.Tensor]]]] = None,
    mappings: Optional[List[Optional[Dict[str, str]]]] = None,
    converted_flags: Optional[List[bool]] = None,
    key_categorizer: callable = categorize_key,
) -> Dict[str, List[float]]:
    """
    Compute per-component mean-squared-energy for N input models across common keys.

    Uses per-element mean(tensor^2) instead of sum, making the metric rank-independent
    (no zero-padding bias from rank mismatches). Only shared keys (common_keys) are
    analyzed, preventing architectural mismatches from affecting the energy ratio.

    Args:
        norm_sds: List of normalized state dicts (one per input model).
        common_keys: Keys present in ALL input state dicts (shared layers only).
        original_sds: Optional original SDs for alpha/rank scaling lookup.
        mappings: Optional key mappings (normalized -> original) for alpha lookup.
        converted_flags: Optional list of bools indicating if each model is 'converted'.
        key_categorizer: Function to categorize keys into components.

    Returns:
        Dict mapping component name -> list of accumulated mean energies
        [energy_0, energy_1, ...] for each input model.
    """
    num_inputs = len(norm_sds)
    if original_sds is None:
        original_sds = [None] * num_inputs
    if mappings is None:
        mappings = [None] * num_inputs
    if converted_flags is None:
        converted_flags = [False] * num_inputs

    energy_by_component: Dict[str, List[float]] = {}

    for key in common_keys:
        tensors = []
        for i in range(num_inputs):
            tensor = norm_sds[i][key]
            # Apply LoRA alpha/rank scaling if original data is available
            # (for unconverted LoRAs with known alpha values)
            rank = safe_get_rank(tensor, key)
            rank = max(1, rank)

            if original_sds[i] is not None:
                from .engine.scale_utils import find_alpha_value
                alpha_value = find_alpha_value(original_sds[i], key, mapping=mappings[i])
                if not converted_flags[i] and alpha_value is not None:
                    scale_factor = alpha_value / rank
                    if abs(scale_factor - 1.0) > 1e-6:
                        # Apply scaling only to up-weights (same logic as merge engine)
                        is_down = any(s in key for s in ['.lora_A.weight', '.lora_down.weight'])
                        if not is_down:
                            tensor = tensor * scale_factor

            tensors.append(tensor)

        # Categorize key
        component = key_categorizer(key)
        if component not in energy_by_component:
            energy_by_component[component] = [0.0] * num_inputs

        # Accumulate per-element mean energy (rank-independent)
        # Using mean(t^2) instead of sum(t^2) avoids bias from rank mismatches
        # Skip integer tensors (e.g. position_ids) which have no meaningful energy
        for i, t in enumerate(tensors):
            if t.dtype in (torch.long, torch.int, torch.int32, torch.int64, torch.int16, torch.int8, torch.uint8):
                continue
            energy_by_component[component][i] += torch.mean(t ** 2).item()

    return energy_by_component


def compute_primary_driver_intensity_metric(
    norm_sds: List[Dict[str, torch.Tensor]],
    common_keys: List[str],
    energy_concentration: float = 0.80,
    original_sds: Optional[List[Optional[Dict[str, torch.Tensor]]]] = None,
    mappings: Optional[List[Optional[Dict[str, str]]]] = None,
    converted_flags: Optional[List[bool]] = None,
    key_categorizer: callable = categorize_key,
) -> Dict[str, dict]:
    """
    Compute per-LoRA 'Primary Driver Intensity' using adaptive energy concentration.

    Instead of a fixed top_percent (like compute_lora_intensity_metric), this function
    dynamically determines which keys carry the majority (80%) of each LoRA's total
    energy and compares them on a JOINT key set (union) for a fair apples-to-apples
    comparison.

    Why this is better than fixed top-percentile:
      - Sparse LoRAs (Anima/Z-Image): 80% energy concentrated in ~5-10% of keys
        → focuses on truly important keys only
      - Dense LoRAs: 80% energy spread over ~30-40% of keys
        → captures adequate context
      - Cross-concept merges: joint union set ensures same-key comparison,
        preventing the "apples to oranges" bias where each LoRA's peak
        is measured on its own strongest keys where the other is weak.

    Args:
        norm_sds: List of normalized state dicts (one per input model).
        common_keys: Keys present in ALL input state dicts (shared layers only).
        energy_concentration: Fraction of total energy to capture (default 0.80 = 80%).
        original_sds: Optional original SDs for alpha/rank scaling lookup.
        mappings: Optional key mappings (normalized -> original) for alpha lookup.
        converted_flags: Optional list of bools indicating if each model is 'converted'.
        key_categorizer: Function to categorize keys into components.

    Returns:
        Dict mapping component name -> dict with keys:
            'peaks': [peak_A, peak_B]
                Mean energy on each LoRA's OWN primary driver set.
            'joint_peaks': [joint_A, joint_B]
                Mean energy on UNION of both LoRAs' primary driver sets (fair comparison).
            'primary_driver_counts': [count_A, count_B]
                Number of primary driver keys for each LoRA.
            'joint_count': int
                Number of keys in the union set.
            'energy_concentration': [conc_A, conc_B]
                Actual fraction of total energy captured by primary drivers.
            'ratio_joint': float
                sqrt(joint_A / joint_B) — THE metric to use for global scaling.
            'overlap_count': int
                Number of keys that are primary drivers for BOTH LoRAs.
    """
    import math
    num_inputs = len(norm_sds)
    if original_sds is None:
        original_sds = [None] * num_inputs
    if mappings is None:
        mappings = [None] * num_inputs
    if converted_flags is None:
        converted_flags = [False] * num_inputs

    # Phase 1: Collect per-key energies per component per model
    # {component: [[(key, energy), ...], [(key, energy), ...]]}
    per_key_energies: Dict[str, List[List[tuple]]] = {}

    for key in common_keys:
        component = key_categorizer(key)
        if component not in per_key_energies:
            per_key_energies[component] = [[] for _ in range(num_inputs)]

        for i in range(num_inputs):
            tensor = norm_sds[i][key]
            rank = safe_get_rank(tensor, key)
            rank = max(1, rank)

            # Alpha/rank scaling for unconverted LoRAs (same logic as existing functions)
            if original_sds[i] is not None:
                from .engine.scale_utils import find_alpha_value
                alpha_value = find_alpha_value(original_sds[i], key, mapping=mappings[i])
                if not converted_flags[i] and alpha_value is not None:
                    scale_factor = alpha_value / rank
                    if abs(scale_factor - 1.0) > 1e-6:
                        is_down = any(s in key for s in ['.lora_A.weight', '.lora_down.weight'])
                        if not is_down:
                            tensor = tensor * scale_factor

            # Per-element mean energy (rank-independent)
            energy = torch.mean(tensor ** 2).item()
            per_key_energies[component][i].append((key, energy))

    # Phase 2-4: Find Primary Drivers, form Joint Set, compute metrics
    result: Dict[str, dict] = {}
    for component, model_key_energies in per_key_energies.items():
        primary_sets: List[set] = []   # list of sets of primary driver keys per model
        peaks: List[float] = []        # mean energy on own primary drivers
        concentrations: List[float] = []# actual energy fraction captured
        all_energies_list: List[dict] = []  # key->energy lookup per model

        for key_energies in model_key_energies:
            # Sort by energy descending
            sorted_items = sorted(key_energies, key=lambda x: x[1], reverse=True)
            total_energy = sum(e for _, e in sorted_items)

            # Build lookup dict for later use
            energy_lookup = {k: e for k, e in key_energies}
            all_energies_list.append(energy_lookup)

            # Find smallest N where cumulative energy >= concentration_threshold * total
            cumsum = 0.0
            cutoff_idx = 0
            for idx, (k, e) in enumerate(sorted_items):
                cumsum += e
                if cumsum >= energy_concentration * total_energy:
                    cutoff_idx = idx + 1
                    break
            else:
                cutoff_idx = len(sorted_items)

            # Primary driver keys = first cutoff_idx keys
            primary_keys = set(k for k, _ in sorted_items[:cutoff_idx])
            primary_sets.append(primary_keys)

            # Peak = mean energy of primary drivers
            if primary_keys:
                peak = sum(e for k, e in sorted_items[:cutoff_idx]) / cutoff_idx
            else:
                peak = 0.0
            peaks.append(peak)
            concentrations.append(cumsum / total_energy if total_energy > 0 else 0.0)

        # Phase 3: Joint union set
        joint_set: set = set()
        for ps in primary_sets:
            joint_set |= ps

        # Phase 4: Compute joint mean energies on the union set
        joint_peaks: List[float] = []
        for i, lookup in enumerate(all_energies_list):
            if joint_set:
                joint_total = sum(lookup.get(k, 0.0) for k in joint_set)
                joint_peak = joint_total / len(joint_set)
            else:
                joint_peak = 0.0
            joint_peaks.append(joint_peak)

        # Ratio on joint set
        eps = 1e-12
        if joint_peaks[1] > eps:
            ratio_joint = math.sqrt(joint_peaks[0] / joint_peaks[1])
        else:
            ratio_joint = 1.0

        # Overlap count
        overlap_count = 0
        if len(primary_sets) >= 2:
            overlap_count = len(primary_sets[0] & primary_sets[1])

        result[component] = {
            'peaks': peaks,
            'joint_peaks': joint_peaks,
            'primary_driver_counts': [len(ps) for ps in primary_sets],
            'joint_count': len(joint_set),
            'energy_concentration': concentrations,
            'ratio_joint': ratio_joint,
            'overlap_count': overlap_count,
        }

    return result


def _lookup_alpha_value(orig_key: str, original_sd: Dict[str, torch.Tensor]) -> Optional[float]:
    """DEPRECATED: Use engine.scale_utils.find_alpha_value instead."""
    import warnings
    warnings.warn(
        "_lookup_alpha_value is deprecated, use engine.scale_utils.find_alpha_value",
        DeprecationWarning, stacklevel=2,
    )
    """
    Look up the alpha value for a given original key.
    Tries multiple candidate alpha key patterns to handle various naming conventions.
    """
    candidates = []
    # Pattern 1: replace .weight with .alpha
    if orig_key.endswith(".weight"):
        candidates.append(orig_key.replace(".weight", ".alpha"))
    # Pattern 2: replace lora_A/B.weight with alpha
    for pat in ['lora_A.weight', 'lora_B.weight', 'lora_down.weight', 'lora_up.weight']:
        if pat in orig_key:
            candidates.append(orig_key.replace(pat, "alpha"))
    # Pattern 3: strip suffix after last .lora_
    if ".lora_" in orig_key:
        base = orig_key.split(".lora_")[0]
        candidates.append(f"{base}.alpha")
    if "_lora_" in orig_key:
        base = orig_key.split("_lora_")[0]
        candidates.append(f"{base}.alpha")

    for ak in candidates:
        if ak in original_sd:
            alpha_tensor = original_sd[ak]
            if isinstance(alpha_tensor, torch.Tensor):
                return alpha_tensor.item() if alpha_tensor.numel() == 1 else alpha_tensor.mean().item()
    return None