"""
SNR-Quant - ComfyUI nodes for adaptive FP8/ BF16 quantization.
"""
import os
import torch
import numpy as np
from safetensors.torch import load_file, save_file
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass
from enum import Enum

import folder_paths
from comfy.utils import ProgressBar


class ModelFragility(Enum):
    """Model fragility levels."""
    FRAGILE = "fragile"
    MODERATE = "moderate"
    ROBUST = "robust"


@dataclass
class LayerInfo:
    """Information about a single layer."""
    name: str
    numel: int
    snr: float
    outlier: float

    @property
    def bf16_size(self) -> int:
        return self.numel * 2

    @property
    def fp8_size(self) -> int:
        return self.numel


def compute_snr(weight: torch.Tensor) -> float:
    """Compute Signal-to-Noise Ratio for FP8 E4M3 quantization."""
    f32_w = weight.float()

    if torch.isnan(f32_w).any() or torch.isinf(f32_w).any():
        return float('inf')

    max_abs = torch.max(torch.abs(f32_w)).item()
    if max_abs == 0:
        return float('inf')

    scale = max_abs / 448.0
    q = (f32_w / scale).clamp(-448, 448)
    q = q.to(torch.float8_e4m3fn).to(f32_w.dtype)
    dq = q * scale

    signal_pwr = torch.mean(f32_w ** 2).item()
    noise_pwr = torch.mean((f32_w - dq) ** 2).item()

    if noise_pwr == 0:
        return float('inf')

    return 10 * np.log10(signal_pwr / noise_pwr)


def compute_outlier_index(weight: torch.Tensor) -> float:
    """Compute outlier index = max_abs / std."""
    f32_w = weight.float()
    max_abs = torch.max(torch.abs(f32_w)).item()
    std = torch.std(f32_w).item()
    if std == 0:
        return 0.0
    return max_abs / std


def analyze_model(sd: Dict[str, torch.Tensor]) -> Tuple[List[LayerInfo], int, int]:
    """Analyze all 2D weight layers.

    Returns:
        layers: List of LayerInfo for quantizable (non-FP8) layers
        total_params: Total params in 2D layers
        existing_fp8_count: Number of already quantized FP8 layers
    """
    layers = []
    total_params = 0
    existing_fp8_count = 0

    for name, tensor in sd.items():
        if tensor.dim() == 2:
            numel = tensor.numel()
            total_params += numel

            # Skip already quantized FP8 layers
            if tensor.dtype == torch.float8_e4m3fn:
                existing_fp8_count += 1
                continue

            snr = compute_snr(tensor)
            outlier = compute_outlier_index(tensor)
            layers.append(LayerInfo(name=name, numel=numel, snr=snr, outlier=outlier))

    return layers, total_params, existing_fp8_count


def detect_fragility(layers: List[LayerInfo]) -> Tuple[ModelFragility, float, float]:
    """Detect model fragility based on SNR distribution."""
    snrs = [l.snr for l in layers]
    snr_range = max(snrs) - min(snrs)
    snr_median = np.median(snrs)

    if snr_range < 2 and snr_median < 32:
        return ModelFragility.FRAGILE, snr_range, snr_median
    elif snr_range > 5 and snr_median > 34:
        return ModelFragility.ROBUST, snr_range, snr_median
    else:
        return ModelFragility.MODERATE, snr_range, snr_median


class SNRQuantNode:
    """Quantize model to target size using FP8/ BF16 mixed precision."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model_name": (folder_paths.get_filename_list("diffusion_models"), {
                    "tooltip": "Model from ComfyUI/models/diffusion_models"
                }),
                "target_size_GB": ("FLOAT", {
                    "default": 22.0,
                    "min": 1.0,
                    "step": 0.1,
                    "tooltip": "Target model size in GB (e.g., 22.0)"
                }),
            }
        }

    RETURN_TYPES = ()
    OUTPUT_NODE = True
    FUNCTION = "quantize"
    CATEGORY = "SNR-Quant"

    def quantize(self, model_name: str, target_size_GB: float):
        # Get paths
        model_path = folder_paths.get_full_path("diffusion_models", model_name)
        output_dir = folder_paths.get_output_directory()
        # Create diffusion_models folder in output if it doesn't exist
        diffusion_models_dir = os.path.join(output_dir, "diffusion_models")
        if not os.path.exists(diffusion_models_dir):
            os.makedirs(diffusion_models_dir)
        # Replace path separators with underscores for flat output filename
        base_name = os.path.splitext(model_name)[0].replace('\\', '_').replace('/', '_')
        # Include target size in filename to distinguish multiple quantizations
        output_path = os.path.join(diffusion_models_dir, f"quant_{base_name}_{target_size_GB}GB.safetensors")

        print(f"\n{'='*60}")
        print(f"[SNR-Quant] Auto Quant - FP8/BF16 Mixed Precision")
        print(f"{'='*60}")
        print(f"[SNR-Quant] Model: {model_name}")
        print(f"[SNR-Quant] Target size: {target_size_GB:.2f} GB")
        print(f"{'='*60}")

        # Load model to CPU
        print(f"\n[Step 1/6] Loading model to CPU...")
        sd = load_file(model_path, device="cpu")
        print(f"  ✓ Loaded {len(sd)} tensors")

        # Detect input dtype
        sample = next(iter(sd.values()))
        input_dtype = sample.dtype
        print(f"  ✓ Input dtype: {input_dtype}")

        # Analyze layers
        print(f"\n[Step 2/6] Analyzing layers (computing SNR and outliers)...")
        layers, total_params, existing_fp8_count = analyze_model(sd)
        if existing_fp8_count > 0:
            print(f"  ✓ Found {len(layers)} quantizable 2D layers ({existing_fp8_count} FP8 layers skipped)")
        else:
            print(f"  ✓ Found {len(layers)} 2D weight layers")

        # Detect fragility
        fragility, snr_range, snr_median = detect_fragility(layers)
        outliers = [l.outlier for l in layers]
        outlier_median = np.median(outliers)

        print(f"\n[Step 3/6] Detecting model fragility...")
        print(f"  SNR range:   {snr_range:.2f} dB")
        print(f"  SNR median:  {snr_median:.2f} dB")
        print(f"  Outlier med: {outlier_median:.1f}")
        print(f"  → Fragility: {fragility.value.upper()}")

        # Set adaptive percentile
        if fragility.value == "fragile":
            outlier_percentile = 90
            print(f"  → Using 90th percentile (conservative)")
        elif fragility.value == "robust":
            outlier_percentile = 50
            print(f"  → Using 50th percentile (aggressive)")
        else:
            outlier_percentile = 75
            print(f"  → Using 75th percentile (balanced)")

        # Calculate sizes
        # layers only contains quantizable 2D weights (excludes existing FP8)
        existing_fp8_numel = sum(t.numel() for t in sd.values() if t.dim() == 2 and t.dtype == torch.float8_e4m3fn)
        existing_fp8_size = existing_fp8_numel  # FP8 = 1 byte/param

        size_2d_bf16 = sum(l.numel * 2 for l in layers)  # quantizable 2D layers as BF16
        size_2d_fp8 = sum(l.numel for l in layers)       # quantizable 2D layers as FP8
        size_1d_bf16 = sum(t.numel() * 2 for t in sd.values() if t.dim() != 2)  # 1D layers always BF16

        # Current total size (existing FP8 + quantizable as BF16 + 1D)
        current_total_size = existing_fp8_size + size_2d_bf16 + size_1d_bf16

        # Minimum achievable (all quantizable as FP8)
        min_size = existing_fp8_size + size_2d_fp8 + size_1d_bf16

        target_bytes = target_size_GB * 1024**3
        need_save = current_total_size - target_bytes

        print(f"\n[Step 4/6] Calculating size targets...")
        print(f"  Current model size:    {current_total_size / 1024**3:.2f} GB")
        print(f"    └─ Existing FP8:     {existing_fp8_size / 1024**3:.2f} GB")
        print(f"    └─ Quantizable BF16: {size_2d_bf16 / 1024**3:.2f} GB")
        if existing_fp8_count > 0:
            print(f"  Note: {existing_fp8_count} FP8 layers already quantized (skipped)")
        print(f"  Minimum achievable:    {min_size / 1024**3:.2f} GB")
        print(f"  Target size:           {target_bytes / 1024**3:.2f} GB")
        print(f"  Need to save:          {need_save / 1024**3:.2f} GB")

        # Check if target is achievable
        if target_bytes < min_size:
            print(f"\n  ⚠ WARNING: Target {target_size_GB:.2f} GB is below minimum!")
            print(f"  Adjusting to {min_size / 1024**3:.2f} GB (all quantizable layers as FP8)")
            target_bytes = min_size
            need_save = current_total_size - min_size

        if need_save <= 0:
            print(f"\n[Step 5/6] Target >= Current size, no quantization needed")
            print(f"  Saving all layers as-is...")

            # Build layer info lookup
            layer_info_map = {l.name: l for l in layers}

            quantized = {}
            for name, tensor in sd.items():
                if tensor.dim() == 2 and tensor.dtype == torch.float8_e4m3fn:
                    # Keep existing FP8 layers as is
                    quantized[name] = tensor
                elif tensor.dim() == 2:
                    quantized[name] = tensor.to(torch.bfloat16)
                else:
                    quantized[name] = tensor.to(torch.bfloat16)
            save_file(quantized, output_path)

            # Generate report - all layers BF16
            report_path = os.path.join(diffusion_models_dir, f"quant_{base_name}_{target_size_GB}GB_report.csv")
            print(f"\n  Generating quantization report: {report_path}")

            report_lines = ["layer_name,precision,params_numel,SNR_dB,outlier,precision_change"]

            # Add existing FP8 layers (already quantized, skipped)
            for name, tensor in sd.items():
                if tensor.dim() == 2 and tensor.dtype == torch.float8_e4m3fn:
                    report_lines.append(f"{name},FP8 (existing),{tensor.numel()},N/A,N/A,already quantized")

            # Add other layers
            for name, tensor in sd.items():
                if tensor.dim() == 2 and name in layer_info_map:
                    l = layer_info_map[name]
                    report_lines.append(f"{name},BF16,{tensor.numel()},{l.snr:.2f},{l.outlier:.2f},unchanged")
                elif tensor.dim() != 2:
                    report_lines.append(f"{name},BF16,{tensor.numel()},N/A,N/A,unchanged")

            with open(report_path, 'w', encoding='utf-8') as f:
                f.write('\n'.join(report_lines))

            print(f"  ✓ Report saved: {report_path}")

            del sd, quantized
            print(f"\n[Step 6/6] Saving to: {output_path}")
            print(f"  ✓ Done! Output: quant_{base_name}.safetensors")
            print(f"{'='*60}\n")
            return ()

        # Iterative selection
        print(f"\n[Step 5/6] Selecting layers for FP8 quantization...")
        final_fp8 = []
        iteration = 0

        while outlier_percentile >= 50:
            iteration += 1
            threshold = np.percentile(outliers, outlier_percentile)
            quantizable = [l for l in layers if l.outlier < threshold]
            quantizable_sorted = sorted(quantizable, key=lambda x: x.snr, reverse=True)
            max_save = sum(l.numel for l in quantizable)

            if max_save >= need_save:
                saved = 0
                selected = 0
                for l in quantizable_sorted:
                    if saved >= need_save:
                        break
                    final_fp8.append(l.name)
                    saved += l.numel
                    selected += 1
                print(f"  Iteration {iteration}: {selected} layers selected at {outlier_percentile}th percentile")
                print(f"  Saved: {saved / 1024**3:.2f} GB / {need_save / 1024**3:.2f} GB needed")
                break
            else:
                print(f"  Iteration {iteration}: Not reachable at {outlier_percentile}th ({max_save / 1024**3:.2f} GB < {need_save / 1024**3:.2f} GB)")
                outlier_percentile -= 5

        if not final_fp8:
            quantizable_sorted = sorted(layers, key=lambda x: x.snr, reverse=True)
            final_fp8 = [l.name for l in quantizable_sorted]
            print(f"  Using ALL {len(final_fp8)} layers for maximum quantization")

        fp8_set = set(final_fp8)
        bf16_count_layers = len(layers) - len(fp8_set)

        # Calculate final sizes
        final_2d_fp8 = sum(l.numel for l in layers if l.name in fp8_set)
        final_2d_bf16 = sum(l.numel * 2 for l in layers if l.name not in fp8_set)
        final_size_gb = (existing_fp8_size + final_2d_fp8 + final_2d_bf16 + size_1d_bf16) / 1024**3

        print(f"\n  Final configuration:")
        print(f"    FP8 (existing):   {existing_fp8_count} ({existing_fp8_size / 1024**3:.2f} GB)")
        print(f"    FP8 (new):        {len(fp8_set)} ({final_2d_fp8 / 1024**3:.2f} GB)")
        print(f"    BF16 (2D):        {bf16_count_layers} ({final_2d_bf16 / 1024**3:.2f} GB)")
        print(f"    └─ Est. size:     {final_size_gb:.2f} GB")

        # Generate quantization report before quantizing
        report_path = os.path.join(diffusion_models_dir, f"quant_{base_name}_{target_size_GB}GB_report.csv")
        print(f"\n  Generating quantization report: {report_path}")

        # Build layer info lookup
        layer_info_map = {l.name: l for l in layers}

        report_lines = ["layer_name,precision,params_numel,SNR_dB,outlier,precision_change"]

        # Add existing FP8 layers (already quantized, skipped)
        for name, tensor in sd.items():
            if tensor.dim() == 2 and tensor.dtype == torch.float8_e4m3fn:
                report_lines.append(f"{name},FP8 (existing),{tensor.numel()},N/A,N/A,already quantized")

        # Add quantizable layers
        for l in layers:
            if l.name in fp8_set:
                report_lines.append(f"{l.name},FP8,{l.numel},{l.snr:.2f},{l.outlier:.2f},BF16→FP8")
            else:
                report_lines.append(f"{l.name},BF16,{l.numel},{l.snr:.2f},{l.outlier:.2f},unchanged")

        # Add 1D layers
        for name, tensor in sd.items():
            if tensor.dim() != 2:
                report_lines.append(f"{name},BF16,{tensor.numel()},N/A,N/A,unchanged")

        with open(report_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(report_lines))

        print(f"  ✓ Report saved: {report_path}")

        # Quantize with progress bar
        print(f"\n[Step 6/6] Quantizing model...")
        quantized = {}
        pbar = ProgressBar(len(sd))
        fp8_count = 0
        bf16_count = 0

        for i, (name, tensor) in enumerate(sd.items()):
            if tensor.dim() == 2 and tensor.dtype == torch.float8_e4m3fn:
                # Already FP8 - keep as is
                quantized[name] = tensor
                fp8_count += 1
            elif tensor.dim() == 2 and name in fp8_set:
                # 2D layer selected for FP8 quantization
                quantized[name] = tensor.to(torch.float8_e4m3fn)
                fp8_count += 1
            else:
                # BF16: all 1D layers + 2D layers not selected for FP8
                quantized[name] = tensor.to(torch.bfloat16)
                bf16_count += 1
            pbar.update_absolute(i + 1)
            if (i + 1) % 100 == 0 or i == len(sd) - 1:
                print(f"  Progress: {i + 1}/{len(sd)} (FP8: {fp8_count}, BF16: {bf16_count})")

        # Save
        print(f"\n  Saving to: {output_path}")
        save_file(quantized, output_path)

        # Cleanup
        del quantized, sd
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        print(f"  ✓ Quantization complete!")
        print(f"  Output: quant_{base_name}.safetensors")
        print(f"{'='*60}\n")
        return ()


class SNRAnalyzeNode:
    """Analyze model and print SNR/outlier statistics."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model_name": (folder_paths.get_filename_list("diffusion_models"), {
                    "tooltip": "Model from ComfyUI/models/diffusion_models"
                }),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("report",)
    FUNCTION = "analyze"
    CATEGORY = "SNR-Quant"

    def analyze(self, model_name: str):
        model_path = folder_paths.get_full_path("diffusion_models", model_name)

        print(f"\n{'='*60}")
        print(f"[SNR-Quant] Model Analyzer")
        print(f"{'='*60}")
        print(f"[SNR-Quant] Model: {model_name}")

        print(f"\n[Step 1/3] Loading model...")
        sd = load_file(model_path, device="cpu")
        print(f"  ✓ Loaded {len(sd)} tensors")

        print(f"\n[Step 2/3] Analyzing layers...")
        layers, total_params, existing_fp8_count = analyze_model(sd)
        if existing_fp8_count > 0:
            print(f"  ✓ Found {len(layers)} quantizable 2D layers ({existing_fp8_count} FP8 layers skipped)")
        else:
            print(f"  ✓ Found {len(layers)} 2D weight layers")

        print(f"\n[Step 3/3] Computing statistics...")
        fragility, snr_range, snr_median = detect_fragility(layers)
        outliers = [l.outlier for l in layers]
        outlier_median = np.median(outliers)
        snr_min = min([l.snr for l in layers])
        snr_max = max([l.snr for l in layers])

        print(f"  SNR:   min={snr_min:.2f}, max={snr_max:.2f}, median={snr_median:.2f} dB")
        print(f"  Outliers: min={min(outliers):.2f}, max={max(outliers):.2f}, median={outlier_median:.1f}")
        print(f"  Fragility: {fragility.value}")

        # Build report - calculate actual model sizes
        size_2d_bf16 = sum(l.numel * 2 for l in layers)  # quantizable 2D as BF16
        size_2d_fp8 = sum(l.numel for l in layers)       # quantizable 2D as FP8
        size_1d_bf16 = sum(t.numel() * 2 for t in sd.values() if t.dim() != 2)  # 1D layers BF16

        # Calculate already quantized FP8 layers size
        existing_fp8_numel = sum(t.numel() for t in sd.values() if t.dim() == 2 and t.dtype == torch.float8_e4m3fn)
        existing_fp8_size = existing_fp8_numel  # FP8 = 1 byte/param

        # Current total size (with existing FP8 + all quantizable as BF16)
        current_total_size = existing_fp8_size + size_2d_bf16 + size_1d_bf16

        # Minimum achievable (all quantizable as FP8)
        min_size = existing_fp8_size + size_2d_fp8 + size_1d_bf16

        # Analyze SNR and outlier distribution for recommendations
        snrs = [l.snr for l in layers]
        outliers = [l.outlier for l in layers]

        # Count layers by SNR quality
        snr_excellent = sum(1 for s in snrs if s >= 35)
        snr_good = sum(1 for s in snrs if 32 <= s < 35)
        snr_moderate = sum(1 for s in snrs if 30 <= s < 32)
        snr_poor = sum(1 for s in snrs if s < 30)

        # Count layers by outlier risk
        outlier_low = sum(1 for o in outliers if o < 15)
        outlier_medium = sum(1 for o in outliers if 15 <= o < 30)
        outlier_high = sum(1 for o in outliers if o >= 30)

        # Calculate safety analysis
        safe_to_quantize = sum(1 for l in layers if l.snr >= 32 and l.outlier < 20)
        moderate_risk = sum(1 for l in layers if (30 <= l.snr < 32 or 20 <= l.outlier < 30))
        high_risk = sum(1 for l in layers if l.snr < 30 or l.outlier >= 30)

        safe_ratio = safe_to_quantize / len(layers) if len(layers) > 0 else 0
        moderate_ratio = moderate_risk / len(layers) if len(layers) > 0 else 0
        high_risk_ratio = high_risk / len(layers) if len(layers) > 0 else 0

        # Weight by layer size for accurate size estimation
        safe_numel = sum(l.numel for l in layers if l.snr >= 32 and l.outlier < 20)
        moderate_numel = sum(l.numel for l in layers if (30 <= l.snr < 32 or 20 <= l.outlier < 30))
        high_risk_numel = sum(l.numel for l in layers if l.snr < 30 or l.outlier >= 30)

        # Size calculations (savings from quantizing)
        conservative_fp8_size = safe_numel  # Only safe layers as FP8
        balanced_fp8_size = safe_numel + moderate_numel * 0.6  # Safe + 60% of moderate
        aggressive_fp8_size = safe_numel + moderate_numel * 0.9 + high_risk_numel * 0.3  # Most except very risky

        # Final target sizes = current size - savings
        conservative_target = current_total_size - conservative_fp8_size
        balanced_target = current_total_size - balanced_fp8_size
        aggressive_target = current_total_size - aggressive_fp8_size

        # Ensure targets are within achievable range
        conservative_target = max(conservative_target, min_size)
        balanced_target = max(balanced_target, min_size)
        aggressive_target = max(aggressive_target, min_size)

        # Build strategy description
        if len(layers) == 0:
            strategy_note = "No quantizable layers remaining - model is already fully quantized"
        elif high_risk_ratio > 0.3:
            strategy_note = "High-risk layers detected - recommend conservative quantization"
        elif safe_ratio > 0.5:
            strategy_note = "Many safe-to-quantize layers - can use aggressive strategy"
        else:
            strategy_note = "Mixed layer characteristics - balanced approach recommended"

        # Handle case where all layers are already quantized
        if len(layers) == 0:
            lines = [
                f"Model: {model_name}",
                f"Fragility: N/A (no quantizable layers)",
                f"2D Layers: 0 + {existing_fp8_count} FP8 (already quantized)",
                f"1D Layers: {len(sd) - existing_fp8_count}",
                f"",
                f"This model is already fully quantized to FP8.",
                f"Current size: {current_total_size / 1024**3:.2f} GB",
                f"No further quantization recommended.",
            ]
            report = "\n".join(lines)
            print(f"\n{report}")
            print(f"\n{'='*60}\n")
            del sd
            return (report,)

        lines = [
            f"Model: {model_name}",
            f"Fragility: {fragility.value}",
            f"2D Layers: {len(layers)} + {existing_fp8_count} FP8 (skipped)",
            f"1D Layers: {len(sd) - len(layers) - existing_fp8_count}",
            f"",
            f"Note: {existing_fp8_count} FP8 layers already quantized (skipped in analysis)" if existing_fp8_count > 0 else "",
            f"",
            f"SNR Distribution:",
            f"  Min:     {snr_min:.2f} dB",
            f"  Max:     {snr_max:.2f} dB",
            f"  Median:  {snr_median:.2f} dB",
            f"  Range:   {snr_range:.2f} dB",
            f"",
            f"SNR Quality Breakdown:",
            f"  Excellent (>=35 dB):  {snr_excellent} layers ({snr_excellent/max(len(layers),1)*100:.0f}%)",
            f"  Good (32-35 dB):      {snr_good} layers ({snr_good/max(len(layers),1)*100:.0f}%)",
            f"  Moderate (30-32 dB):  {snr_moderate} layers ({snr_moderate/max(len(layers),1)*100:.0f}%)",
            f"  Poor (<30 dB):        {snr_poor} layers ({snr_poor/max(len(layers),1)*100:.0f}%)",
            f"",
            f"Outlier Distribution:",
            f"  Min:     {min(outliers):.2f}",
            f"  Max:     {max(outliers):.2f}",
            f"  Median:  {outlier_median:.1f}",
            f"",
            f"Outlier Risk Breakdown:",
            f"  Low (<15):     {outlier_low} layers ({outlier_low/max(len(layers),1)*100:.0f}%)",
            f"  Medium (15-30): {outlier_medium} layers ({outlier_medium/max(len(layers),1)*100:.0f}%)",
            f"  High (>=30):    {outlier_high} layers ({outlier_high/max(len(layers),1)*100:.0f}%)",
            f"",
            f"Quantization Safety Analysis:",
            f"  Safe to quantize (SNR>=32 AND outlier<20):  {safe_to_quantize} layers ({safe_ratio*100:.0f}%)",
            f"  Moderate risk:                               {moderate_risk} layers ({moderate_ratio*100:.0f}%)",
            f"  High risk (protect recommended):             {high_risk} layers ({high_risk_ratio*100:.0f}%)",
            f"",
            f"Size estimates:",
            f"  Current model size:    {current_total_size / 1024**3:.2f} GB",
            f"    └─ Existing FP8:     {existing_fp8_size / 1024**3:.2f} GB",
            f"    └─ Quantizable BF16: {size_2d_bf16 / 1024**3:.2f} GB",
            f"  Minimum achievable:    {min_size / 1024**3:.2f} GB (all quantizable as FP8)",
            f"",
            f"Quantization Recommendations:",
            f"  Strategy: {strategy_note}",
            f"",
            f"  Recommended target sizes (final model size):",
            f"    Conservative: {conservative_target / 1024**3:.2f} GB (save {(current_total_size - conservative_target) / 1024**3:.2f} GB)",
            f"      └─ FP8 layers: only safe layers (SNR>=32, outlier<20)",
            f"      └─ Expected quality: Near-lossless",
            f"      └─ Best for: Production, quality-critical work",
            f"",
            f"    Balanced: {balanced_target / 1024**3:.2f} GB (save {(current_total_size - balanced_target) / 1024**3:.2f} GB)",
            f"      └─ FP8 layers: safe + partial moderate-risk layers",
            f"      └─ Expected quality: Minimal degradation",
            f"      └─ Best for: General use, good quality/size tradeoff",
            f"",
            f"    Aggressive: {aggressive_target / 1024**3:.2f} GB (save {(current_total_size - aggressive_target) / 1024**3:.2f} GB)",
            f"      └─ FP8 layers: all except high-risk",
            f"      └─ Expected quality: Acceptable for most cases",
            f"      └─ Best for: VRAM-constrained environments",
        ]
        report = "\n".join(lines)

        print(f"\n{report}")
        print(f"\n{'='*60}\n")
        del sd

        return (report,)


NODE_CLASS_MAPPINGS = {
    "SNRQuant": SNRQuantNode,
    "SNRAnalyze": SNRAnalyzeNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SNRQuant": "SNR Quant (FP8/ BF16)",
    "SNRAnalyze": "SNR Model Analyzer",
}

__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]