"""
LoRA Merging Methods for EasyLoRAMerger - Version 1.3.0
Includes Active Region Blending for better sparse LoRA compatibility
"""

import hashlib
import torch
from typing import List

from ._merge_kernels import (
    compute_per_channel_alignment,
    apply_svd_reduction,
    apply_slerp,
    make_dare_mask,
)

# Import get_tensor_energy for apply_magnitude_scaling
try:
    from ..utils import get_tensor_energy
except ImportError:
    from utils import get_tensor_energy

# ==================== SHAPE ADAPTATION UTILITIES ====================

def ensure_shape_match(a, b, device=None):
    """
    Adjusts tensors a and b to match the larger shape.
    Pads the smaller tensor with zeros to prevent math errors.
    
    Handles ndim mismatches (e.g., 1D bias vs 2D weight) by reshaping
    the lower-dimensional tensor with trailing size-1 dimensions before
    padding, preventing silent data corruption from zip truncation.

    NOTE: When padding occurs, a log message is printed showing the
    original and target shapes for transparency.

    Args:
        a, b: Input tensors.
        device: Optional target device. When provided, tensors are moved to
                this device regardless of their current location.
    """
    a, b = _unify_dtype_and_device(a, b, device=device)
    if a.shape == b.shape:
        return a, b
    
    # Log shape mismatch before handling — transparency for users
    orig_a_shape = a.shape
    orig_b_shape = b.shape
    
    # Handle ndim mismatch: reshape smaller tensor to match ndim of larger
    # by adding trailing dimensions of size 1 (e.g., [4096] -> [4096, 1] for 2D target)
    if a.ndim != b.ndim:
        if a.ndim < b.ndim:
            a = a.reshape(a.shape + (1,) * (b.ndim - a.ndim))
        else:
            b = b.reshape(b.shape + (1,) * (a.ndim - b.ndim))
    
    # Decide which shape is the 'target' based on number of elements
    if a.numel() >= b.numel():
        target_shape = a.shape
        # Create a zero-filled version of B in the shape of A
        new_b = torch.zeros(target_shape, device=b.device, dtype=b.dtype)
        # Calculate the intersection (min size for each dimension) to copy data
        slices = tuple(slice(0, min(i, j)) for i, j in zip(b.shape, target_shape))
        new_b[slices] = b[slices]
        print(f"   [pad] Shape mismatch: {orig_a_shape} vs {orig_b_shape} — padded B {b.shape} -> {target_shape}")
        return a, new_b
    else:
        target_shape = b.shape
        # Create a zero-filled version of A in the shape of B
        new_a = torch.zeros(target_shape, device=a.device, dtype=a.dtype)
        slices = tuple(slice(0, min(i, j)) for i, j in zip(a.shape, target_shape))
        new_a[slices] = a[slices]
        print(f"   [pad] Shape mismatch: {orig_a_shape} vs {orig_b_shape} — padded A {a.shape} -> {target_shape}")
        return new_a, b

def _unify_dtype_and_device(a, b, device=None):
    """
    Ensure two tensors have the same dtype and device.
    Returns (a_unified, b_unified) where dtype is promoted to the higher precision,
    and device is the same (respects user's device choice when provided).

    Args:
        a, b: Input tensors.
        device: Optional target device. When provided, tensors are moved to
                this device regardless of their current location. When None,
                infers device (legacy behavior: CUDA if any tensor is on CUDA,
                otherwise CPU).
    """
    # Determine target dtype: promote to higher precision
    target_dtype = torch.promote_types(a.dtype, b.dtype)
    # Determine target device: respect explicit device parameter when provided
    if device is not None:
        target_device = torch.device(device)
    elif (a.device.type == 'cuda' or b.device.type == 'cuda') and torch.cuda.is_available():
        target_device = torch.device('cuda')
    else:
        target_device = torch.device('cpu')
    
    a_unified = a.to(device=target_device, dtype=target_dtype)
    b_unified = b.to(device=target_device, dtype=target_dtype)
    return a_unified, b_unified

def universal_merge_executor(method_fn, a, b, wa, wb, pbar=None, device=None, **kwargs):
    """
    Safety wrapper that ensures shapes match before executing any merge method.
    Prevents "RuntimeError: size mismatch" between different model architectures.

    Args:
        method_fn: Merge function to apply.
        a, b: Input tensors.
        wa, wb: Weights for a and b.
        pbar: Optional progress tracker; advanced by 1 after each merge.
        device: Optional target device. When provided, tensors are moved to
                this device regardless of their current location. Respects
                user's "auto"/"cuda"/"cpu" choice.
        **kwargs: Additional arguments passed to method_fn.
    """
    a_unified, b_unified = _unify_dtype_and_device(a, b, device=device)
    a_safe, b_safe = ensure_shape_match(a_unified, b_unified, device=device)
    merged = method_fn(a_safe, b_safe, wa, wb, **kwargs)
    if pbar is not None:
        pbar.update(1)
    return merged

# ==================== ACTIVE REGION BLENDING HELPER ====================

def apply_active_blending(method_fn, a, b, wa, wb, blend_mode="active", device=None, **kwargs):
    """
    Wrapper that applies merge methods only where tensors are active (non-zero)
    This preserves sparsity patterns and improves cross-trainer blending

    Args:
        device: Optional target device. Passed through for pipeline consistency;
                upstream callers ensure tensors arrive on the correct device.
    """
    dtype = a.dtype  # Store original dtype
    
    if blend_mode != "active":
        # Standard mode - apply to all values
        return method_fn(a, b, wa, wb, **kwargs)
    
    # Find active regions (above threshold)
    threshold = kwargs.get('active_threshold', 1e-8)
    a_active = (torch.abs(a) > threshold)
    b_active = (torch.abs(b) > threshold)
    
    both_active = a_active & b_active
    only_a = a_active & ~b_active
    only_b = b_active & ~a_active
    neither = ~a_active & ~b_active
    
    # Initialize result
    result = torch.zeros_like(a)
    
    # Where both are active: apply the merge method with float32 upcasting
    if both_active.any():
        a_slice = a[both_active]
        b_slice = b[both_active]
        # Upcast to float32 for numerical stability
        a_fp32 = a_slice.float()
        b_fp32 = b_slice.float()
        merged_fp32 = method_fn(a_fp32, b_fp32, wa, wb, **kwargs)
        # NaN/Inf safety
        if torch.isnan(merged_fp32).any() or torch.isinf(merged_fp32).any():
            print(f"⚠️ NaN/Inf detected in active blending, replacing with zeros")
            merged_fp32 = torch.nan_to_num(merged_fp32, nan=0.0, posinf=1.0, neginf=-1.0)
        # Downcast to original dtype
        merged = merged_fp32.to(dtype)
        result[both_active] = merged
    
    # Where only A is active: just take A (weighted) with upcasting
    if only_a.any():
        a_slice = a[only_a]
        a_fp32 = a_slice.float()
        weighted = a_fp32 * wa
        # NaN/Inf safety
        if torch.isnan(weighted).any() or torch.isinf(weighted).any():
            weighted = torch.nan_to_num(weighted, nan=0.0, posinf=1.0, neginf=-1.0)
        result[only_a] = weighted.to(dtype)
    
    # Where only B is active: just take B (weighted) with upcasting
    if only_b.any():
        b_slice = b[only_b]
        b_fp32 = b_slice.float()
        weighted = b_fp32 * wb
        if torch.isnan(weighted).any() or torch.isinf(weighted).any():
            weighted = torch.nan_to_num(weighted, nan=0.0, posinf=1.0, neginf=-1.0)
        result[only_b] = weighted.to(dtype)
    
    # Where neither is active: already zero

    return result

# ==================== ACTIVE BLEND SCAFFOLDING HELPERS ====================

def _active_blend_setup(a, b, kwargs):
    """
    Extract active-region masks shared by all merge methods.
    
    Returns a dict with:
      'both':   regions where both a and b are active (above threshold)
      'only_a': regions where only a is active
      'only_b': regions where only b is active
      'result': zero-filled tensor of the same shape as a
    """
    threshold = kwargs.get('active_threshold', 1e-8)
    a_active = torch.abs(a) > threshold
    b_active = torch.abs(b) > threshold
    return {
        'both': a_active & b_active,
        'only_a': a_active & ~b_active,
        'only_b': b_active & ~a_active,
        'result': torch.zeros_like(a),
    }


def _active_blend_finalize(result, a, b, wa, wb, regions):
    """
    Apply only_a / only_b weighting — common tail for all active-blend methods.
    """
    if regions['only_a'].any():
        result[regions['only_a']] = a[regions['only_a']] * wa
    if regions['only_b'].any():
        result[regions['only_b']] = b[regions['only_b']] * wb
    return result


# ==================== PROMOTED CRAZY MODE VARIANTS ====================

def merge_cross(a, b, wa, wb, blend_mode="active", **kwargs):
    """
    Cross-magnitude merge — linear blend plus pairwise interaction term.
    
    Computes: result = a*wa + b*wb + 0.3 * sqrt(|a*wa| * |b*wb|) * sign(result)
    
    The cross term captures pairwise interactions between LoRAs, creating
    a non-linear mixing effect. Extends naturally to N-way via sum of
    all pairwise cross terms in the triple variant.
    
    Checkpoint compatibility:
        ✅ Works as expected on absolute checkpoint weights.
        ✅ 4D conv weights handled generically — no reshape needed.
        ✅ 1D bias tensors work without special handling.
    
    Originated from merge_linear's deterministic crazy_mode variant.
    """
    def _cross_impl(a, b, wa, wb, **kwargs):
        a_w = a * wa
        b_w = b * wb
        base = a_w + b_w
        cross_mag = (a_w.abs() * b_w.abs()).sqrt() * 0.3
        sign = torch.sign(base)
        # If base is zero, fall back to sign of a_w
        sign = torch.where(sign == 0, torch.sign(a_w), sign)
        cross = cross_mag * sign
        return base + cross
    
    return apply_active_blending(_cross_impl, a, b, wa, wb, blend_mode=blend_mode, **kwargs)


def merge_ties_contrast(a, b, wa, wb, agreement_threshold=0.3, blend_mode="active", **kwargs):
    """
    TIES contrast merge — amplifies disagreements, mutes agreements.
    
    Where signs disagree: amplify by 2x.
    Where signs agree: mute by 0.5x.
    This emphasizes divergent features between LoRAs, useful for
    combining strongly different styles.
    
    NOTE: agreement_threshold is accepted in the signature for API compatibility
    with the 3-way version, but is unused in the binary version which always
    uses a hard sign-equality test (threshold=0). This is because sign-based
    comparison is naturally binary — two signs either match or don't.
    
    Checkpoint compatibility:
        ⚠️ Sign-based comparison is near-no-op on all-positive checkpoint
           weights (common in ReLU-activated layers). The amplify/mute
           effect that makes TIES contrast distinctive is largely absent
           on absolute weights. Best results on delta weights (LoRA).
        ✅ 4D conv weights handled generically — no reshape needed.
        ✅ 1D bias tensors work without special handling.
    
    Originated from merge_ties_gentle's deterministic crazy_mode variant.
    """
    dtype = a.dtype
    
    def _contrast_impl(a, b, wa, wb, **kwargs):
        a_w = a * wa
        b_w = b * wb
        a_sign = torch.sign(a_w).float()
        b_sign = torch.sign(b_w).float()
        
        # Where they disagree, boost the values
        disagreement = (a_sign != b_sign).to(dtype)
        agreement = (a_sign == b_sign).to(dtype)
        
        # Disagreement areas get amplified, agreement areas get muted
        result = (a_w + b_w) * (1.0 + disagreement) * 0.5
        return result
    
    return apply_active_blending(_contrast_impl, a, b, wa, wb, blend_mode=blend_mode, **kwargs)


def merge_block_swap(a, b, wa, wb, block_size=8, seed=42, blend_mode="active", **kwargs):
    """
    Deterministic block-swapping merge with seeded randomness.
    
    Divides 2D tensors into block_size x block_size blocks. Each block
    is randomly chosen from A or B using a seeded RNG (not global random).
    For 1D tensors, falls back to element-wise random selection.
    For 4D conv weights, flattens to 2D [C_out, C_in*H*W], performs block
    swapping, then reshapes back to the original 4D shape.
    For other dimensions, falls back to weighted blend with a warning.

    Checkpoint compatibility:
        ⚠️ 4D conv weights now properly block-swapped via 2D flatten.
        ⚠️ 1D bias tensors use element-wise random selection.
        ✅ 2D linear weights work optimally.
    
    Related to DBIT/block-wise model merging techniques.
    
    Originated from merge_svd_preserve's crazy_mode variant, now
    made reproducible via seed-controlled generator.
    """
    def _block_impl(a, b, wa, wb, **kwargs):
        a_w = a * wa
        b_w = b * wb
        
        # Use seed-based generator for reproducibility
        # NOTE: Python's built-in hash() is randomized per process on 3.6+ (PYTHONHASHSEED),
        # so we use hashlib.md5 for a deterministic, cross-run-stable seed offset.
        shape_bytes = str(a.shape).encode()
        digest = int(hashlib.md5(shape_bytes).hexdigest()[:8], 16)
        rng = torch.Generator(device=a.device).manual_seed(seed + digest % 10000)
        
        if a_w.dim() == 2:
            h, w = a_w.shape
            block_h = max(1, h // block_size)
            block_w = max(1, w // block_size)
            
            # Create block pattern
            result = a_w.clone()
            for i in range(0, h, block_h):
                for j in range(0, w, block_w):
                    rand_val = torch.rand(1, generator=rng, device=a.device).item()
                    if rand_val > 0.5:
                        # Swap this block from B
                        i_end = min(i + block_h, h)
                        j_end = min(j + block_w, w)
                        result[i:i_end, j:j_end] = b_w[i:i_end, j:j_end]
            return result
        elif a_w.dim() == 1:
            # 1D: element-wise random selection with seeded RNG
            mask = torch.rand(a_w.shape, generator=rng, device=a.device) > 0.5
            return torch.where(mask, a_w, b_w)
        elif a_w.dim() == 4:
            # 4D conv weights: flatten to 2D [C_out, C_in*H*W], block swap, reshape back
            orig_shape = a_w.shape
            a_2d = a_w.reshape(orig_shape[0], -1)
            b_2d = b_w.reshape(orig_shape[0], -1)
            h, w = a_2d.shape
            block_h = max(1, h // block_size)
            block_w = max(1, w // block_size)
            result_2d = a_2d.clone()
            for i in range(0, h, block_h):
                for j in range(0, w, block_w):
                    rand_val = torch.rand(1, generator=rng, device=a.device).item()
                    if rand_val > 0.5:
                        i_end = min(i + block_h, h)
                        j_end = min(j + block_w, w)
                        result_2d[i:i_end, j:j_end] = b_2d[i:i_end, j:j_end]
            return result_2d.reshape(orig_shape)
        else:
            # Fallback for other dims — a_w and b_w already have wa/wb applied
            print(f"   ⚠️ block_swap: unsupported tensor dim {a_w.dim()} ({a_w.shape}), "
                  f"falling back to weighted sum")
            return a_w + b_w
    
    return apply_active_blending(_block_impl, a, b, wa, wb, blend_mode=blend_mode, **kwargs)


# ==================== CORE METHODS (UPDATED WITH ACTIVE BLENDING) ====================

def merge_linear(a, b, wa, wb, blend_mode="active", **kwargs):
    """Simple weighted average with active region preservation.

    Checkpoint compatibility:
        ✅ Works as expected on absolute checkpoint weights.
        ✅ 4D conv weights handled generically — no reshape needed.
        ✅ 1D bias tensors work without special handling.
    """
    if blend_mode == "active":
        regions = _active_blend_setup(a, b, kwargs)
        result = regions['result']
        
        if regions['both'].any():
            a_w = a[regions['both']] * wa
            b_w = b[regions['both']] * wb
            result[regions['both']] = a_w + b_w
        
        return _active_blend_finalize(result, a, b, wa, wb, regions)
    
    else:
        # Simple weighted sum (dense mode)
        return a * wa + b * wb

# ==================== TIES METHODS ====================

def merge_ties_strict(a, b, wa, wb, blend_mode="active", **kwargs):
    """TIES merging: only keep weights where signs agree.

    Checkpoint compatibility:
        ⚠️ Sign-based comparison is near-no-op on all-positive checkpoint
           weights (common in ReLU-activated layers). The disagreement
           zeroing that makes TIES distinctive is largely absent on
           absolute weights. Best results on delta weights (LoRA).
        ✅ 4D conv weights handled generically — no reshape needed.
        ✅ 1D bias tensors work without special handling.
    """
    dtype = a.dtype
    
    if blend_mode == "active":
        regions = _active_blend_setup(a, b, kwargs)
        result = regions['result']
        
        if regions['both'].any():
            a_w = a[regions['both']] * wa
            b_w = b[regions['both']] * wb
            sign_agreement = (torch.sign(a_w) == torch.sign(b_w)).to(dtype)
            result[regions['both']] = (a_w + b_w) * sign_agreement
        
        return _active_blend_finalize(result, a, b, wa, wb, regions)
    
    else:
        a_weighted = a * wa
        b_weighted = b * wb
        sign_agreement = (torch.sign(a_weighted) == torch.sign(b_weighted)).to(dtype)
        return (a_weighted + b_weighted) * sign_agreement


def merge_ties_gentle(a, b, wa, wb, agreement_threshold=0.3, blend_mode="active", **kwargs):
    """Gentle TIES: Suppress elements where weighted tensors disagree in sign.

    Applies per-element sign agreement unconditionally — if a_i * wa and b_i * wb
    have opposite signs, that element is zeroed out.  This avoids the scalar
    cos_sim gate which could spuriously disable per-element masking (C4 fix).

    Checkpoint compatibility:
        ⚠️ Sign-based comparison is near-no-op on all-positive checkpoint
           weights (common in ReLU-activated layers). The disagreement
           zeroing that makes TIES distinctive is largely absent on
           absolute weights. Best results on delta weights (LoRA).
        ✅ 4D conv weights handled generically — no reshape needed.
        ✅ 1D bias tensors work without special handling.
    """
    dtype = a.dtype  # Store original dtype
    
    if blend_mode == "active":
        regions = _active_blend_setup(a, b, kwargs)
        result = regions['result']
        
        if regions['both'].any():
            a_w = a[regions['both']] * wa
            b_w = b[regions['both']] * wb
            
            # Per-element sign agreement — always apply (removed scalar cos_sim gate)
            sign_agreement = (torch.sign(a_w) == torch.sign(b_w)).to(dtype)
            result[regions['both']] = (a_w + b_w) * sign_agreement
        
        return _active_blend_finalize(result, a, b, wa, wb, regions)
    
    else:
        a_weighted = a * wa
        b_weighted = b * wb
        
        # Per-element sign agreement — always apply (removed scalar cos_sim gate)
        sign_agreement = (torch.sign(a_weighted) == torch.sign(b_weighted)).to(dtype)
        return (a_weighted + b_weighted) * sign_agreement


# ==================== DARE METHODS ====================

def merge_dare_lite(a, b, wa, wb, drop_rate=0.1, seed=None, blend_mode="active", **kwargs):
    """DARE: Random dropout per tensor before summing.

    Applies an independent dropout mask to each weighted tensor **before**
    summing, so that dropped elements from one tensor are not contaminated
    by the other tensor's signal.  This matches the DARE-lite behaviour in
    the triple-merge core (H6 fix).

    If seed is provided, uses a seeded torch.Generator for reproducible results.

    Checkpoint compatibility:
        ⚠️ Random sparsification is designed for delta weights (LoRA).
           On absolute checkpoint weights, dropout destroys meaningful
           signal. Only use DARE on LoRAs, or on checkpoints that have
           been pre-converted to delta weights.
        ✅ 4D conv weights handled generically — no reshape needed.
        ✅ 1D bias tensors work without special handling.
    """
    dtype = a.dtype
    rng = torch.Generator(device=a.device).manual_seed(seed) if seed is not None else None
    
    if blend_mode == "active":
        regions = _active_blend_setup(a, b, kwargs)
        result = regions['result']
        
        if regions['both'].any():
            a_slice = a[regions['both']]
            b_slice = b[regions['both']]
            # Apply dropout to each weighted tensor individually, then sum
            mask_a = make_dare_mask(a_slice, drop_rate, dtype, rng)
            mask_b = make_dare_mask(b_slice, drop_rate, dtype, rng)
            result[regions['both']] = a_slice * wa * mask_a + b_slice * wb * mask_b
        
        return _active_blend_finalize(result, a, b, wa, wb, regions)
    
    else:
        # Apply dropout to each weighted tensor individually, then sum
        mask_a = make_dare_mask(a, drop_rate, dtype, rng)
        mask_b = make_dare_mask(b, drop_rate, dtype, rng)
        return a * wa * mask_a + b * wb * mask_b


def merge_dare_rescale(a, b, wa, wb, drop_rate=0.1, seed=None, blend_mode="active", **kwargs):
    """DARE with proper rescaling.

    If seed is provided, uses a seeded torch.Generator for reproducible results.

    Checkpoint compatibility:
        ⚠️ Random sparsification is designed for delta weights (LoRA).
           On absolute checkpoint weights, dropout destroys meaningful
           signal. Only use DARE on LoRAs, or on checkpoints that have
           been pre-converted to delta weights.
        ✅ 4D conv weights handled generically — no reshape needed.
        ✅ 1D bias tensors work without special handling.
    """
    dtype = a.dtype
    drop_rate = max(0.0, min(0.9, drop_rate))
    rng = torch.Generator(device=a.device).manual_seed(seed) if seed is not None else None
    
    if blend_mode == "active":
        regions = _active_blend_setup(a, b, kwargs)
        result = regions['result']
        rescale = 1.0 / (1.0 - drop_rate) if drop_rate < 1.0 else 1.0
        
        if regions['both'].any():
            mask_a = make_dare_mask(a[regions['both']], drop_rate, dtype, rng)
            mask_b = make_dare_mask(b[regions['both']], drop_rate, dtype, rng)
            result[regions['both']] = (a[regions['both']] * mask_a * rescale) * wa + (b[regions['both']] * mask_b * rescale) * wb
        
        if regions['only_a'].any():
            mask_a = make_dare_mask(a[regions['only_a']], drop_rate, dtype, rng)
            result[regions['only_a']] = (a[regions['only_a']] * mask_a * rescale) * wa
        
        if regions['only_b'].any():
            mask_b = make_dare_mask(b[regions['only_b']], drop_rate, dtype, rng)
            result[regions['only_b']] = (b[regions['only_b']] * mask_b * rescale) * wb
        
        return result
    
    else:
        mask_a = make_dare_mask(a, drop_rate, dtype, rng)
        mask_b = make_dare_mask(b, drop_rate, dtype, rng)
        rescale = 1.0 / (1.0 - drop_rate) if drop_rate < 1.0 else 1.0
        return (a * mask_a * rescale) * wa + (b * mask_b * rescale) * wb


# ==================== INTUITIVE METHODS ====================

def merge_subtract(a, b, wa, wb, threshold=0.2, blend_mode="active", **kwargs):
    """Subtract B from A where B is significant.

    Checkpoint compatibility:
        ⚠️ Subtractive merging can produce negative weights on absolute
           checkpoints (especially after rescaling). This is mathematically
           valid but unusual for standard checkpoints. Consider using with
           lower weights or on delta/LoRA weights for best results.
        ✅ 4D conv weights handled generically — no reshape needed.
        ✅ 1D bias tensors work without special handling.
    """
    dtype = a.dtype
    
    if blend_mode == "active":
        regions = _active_blend_setup(a, b, kwargs)
        result = regions['result']
        
        if regions['both'].any():
            a_w = a[regions['both']] * wa
            b_w = b[regions['both']] * wb
            mask = (b_w.abs() > (threshold * a_w.abs())).to(dtype)
            result[regions['both']] = (a_w - b_w) * mask + (a_w * (1 - mask))
        
        if regions['only_a'].any():
            result[regions['only_a']] = a[regions['only_a']] * wa
        
        if regions['only_b'].any():
            result[regions['only_b']] = -b[regions['only_b']] * wb
        
        return result
    
    else:
        a_w = a * wa
        b_w = b * wb
        mask = (b_w.abs() > (threshold * a_w.abs())).to(dtype)
        return (a_w - b_w) * mask + (a_w * (1 - mask))


def merge_magnitude(a, b, wa, wb, blend=0.5, blend_mode="active", **kwargs):
    """Keep larger magnitude from either LoRA.

    Checkpoint compatibility:
        ✅ Works as expected on absolute checkpoint weights. Picks the
           dominant feature map per-element regardless of weight type.
        ✅ 4D conv weights handled generically — no reshape needed.
        ✅ 1D bias tensors work without special handling.
    """
    if blend_mode == "active":
        regions = _active_blend_setup(a, b, kwargs)
        result = regions['result']
        
        if regions['both'].any():
            a_w = a[regions['both']] * wa
            b_w = b[regions['both']] * wb
            a_larger = a_w.abs() > b_w.abs()
            selected = torch.where(a_larger, a_w, b_w)
            averaged = (a_w + b_w) / 2
            result[regions['both']] = selected * blend + averaged * (1 - blend)
        
        return _active_blend_finalize(result, a, b, wa, wb, regions)
    
    else:
        a_w = a * wa
        b_w = b * wb
        a_larger = a_w.abs() > b_w.abs()
        selected = torch.where(a_larger, a_w, b_w)
        averaged = (a_w + b_w) / 2
        return selected * blend + averaged * (1 - blend)


def merge_feature_mix(a, b, wa, wb, uniqueness=0.7, blend_mode="active", **kwargs):
    """Preserve unique features from each LoRA.

    Checkpoint compatibility:
        ✅ Works as expected on absolute checkpoint weights. Magnitude-ratio
           based selection is agnostic to absolute vs delta weight semantics.
        ✅ 4D conv weights handled generically — no reshape needed.
        ✅ 1D bias tensors work without special handling.
    """
    dtype = a.dtype
    
    if blend_mode == "active":
        regions = _active_blend_setup(a, b, kwargs)
        result = regions['result']
        
        if regions['both'].any():
            a_w = a[regions['both']] * wa
            b_w = b[regions['both']] * wb
            total = a_w.abs() + b_w.abs() + 1e-8
            a_share = a_w.abs() / total
            
            mask_a = (a_share > uniqueness).to(dtype)
            mask_b = (a_share < (1.0 - uniqueness)).to(dtype)
            mask_shared = torch.clamp(1.0 - mask_a - mask_b, 0.0, 1.0).to(dtype)
            
            result[regions['both']] = (a_w * mask_a) + (b_w * mask_b) + ((a_w + b_w) / 2) * mask_shared
        
        return _active_blend_finalize(result, a, b, wa, wb, regions)
    
    else:
        a_w = a * wa
        b_w = b * wb
        total = a_w.abs() + b_w.abs() + 1e-8
        a_share = a_w.abs() / total
        
        mask_a = (a_share > uniqueness).to(dtype)
        mask_b = (a_share < (1.0 - uniqueness)).to(dtype)
        mask_shared = torch.clamp(1.0 - mask_a - mask_b, 0.0, 1.0).to(dtype)
        
        result = (a_w * mask_a) + (b_w * mask_b) + ((a_w + b_w) / 2) * mask_shared
        return result


# ==================== ADVANCED METHODS ====================

def merge_svd_preserve(a, b, wa, wb, preserve_ratio=0.8, blend_mode="active", **kwargs):
    """SVD-based merge with rank reduction.

    Reshapes 4D conv weights to 2D [C_out, C_in*H*W] before SVD for a
    meaningful per-output-channel low-rank decomposition (instead of
    PyTorch's default batched SVD over spatial dims).  Reshapes back
    to the original 4D shape after decomposition.

    1D tensors (biases) and 1x1 matrices skip SVD and fall back to
    linear blend, since their structure has no meaningful low-rank
    approximation.

    Checkpoint compatibility:
        ⚠️ 4D conv weights are now properly flattened for meaningful SVD.
        ⚠️ 1D bias tensors use linear blend (logged).
        ✅ 2D linear weights work optimally.
    """
    dtype = a.dtype
    combined = a.float() * wa + b.float() * wb
    result = apply_svd_reduction(combined, preserve_ratio)
    return result.to(dtype)


def merge_noise_aware(a, b, wa, wb, noise_threshold=0.01, blend_mode="active", **kwargs):
    """Reduce noise before merging.

    NOTE: In active blend mode, the noise threshold is computed per-region
    (from the 'both' region's max), not from the full tensor's max. This is
    intentional — active blending already isolates regions — but differs from
    dense mode which uses the full tensor's max for threshold computation.
    """
    
    if blend_mode == "active":
        regions = _active_blend_setup(a, b, kwargs)
        result = regions['result']
        
        if regions['both'].any():
            threshold_a = noise_threshold * a[regions['both']].abs().max()
            threshold_b = noise_threshold * b[regions['both']].abs().max()
            a_clean = torch.where(a[regions['both']].abs() < threshold_a, a[regions['both']] * 0.1, a[regions['both']])
            b_clean = torch.where(b[regions['both']].abs() < threshold_b, b[regions['both']] * 0.1, b[regions['both']])
            result[regions['both']] = a_clean * wa + b_clean * wb
        
        if regions['only_a'].any():
            threshold_a = noise_threshold * a[regions['only_a']].abs().max()
            a_clean = torch.where(a[regions['only_a']].abs() < threshold_a, a[regions['only_a']] * 0.1, a[regions['only_a']])
            result[regions['only_a']] = a_clean * wa
        
        if regions['only_b'].any():
            threshold_b = noise_threshold * b[regions['only_b']].abs().max()
            b_clean = torch.where(b[regions['only_b']].abs() < threshold_b, b[regions['only_b']] * 0.1, b[regions['only_b']])
            result[regions['only_b']] = b_clean * wb
        
        return result
    
    else:
        threshold_a = noise_threshold * a.abs().max()
        threshold_b = noise_threshold * b.abs().max()
        a_clean = torch.where(a.abs() < threshold_a, a * 0.1, a)
        b_clean = torch.where(b.abs() < threshold_b, b * 0.1, b)
        return a_clean * wa + b_clean * wb


def merge_gradient_alignment(a, b, wa, wb, blend_mode="active", **kwargs):
    """Merge based on per-channel directional alignment.

    Reshapes to [channel_dim, -1] for per-channel cosine similarity instead of
    flattening everything, preserving per-filter directional information for
    diverse layers like conv weights (H5 fix).

    Checkpoint compatibility:
        ✅ Works as expected on absolute checkpoint weights. Directional
           alignment is based on vector geometry, not weight semantics.
        ✅ 4D conv weights handled via per-channel reshape [C_out, -1].
        ✅ 1D bias tensors fall back to neutral alignment (0.5).
    """
    dtype = a.dtype

    if blend_mode == "active":
        # Note: uses != 0 rather than threshold-based active detection
        # because gradient directions are meaningful even at small magnitudes
        a_active = (a != 0)
        b_active = (b != 0)
        both_active = a_active & b_active
        only_a = a_active & ~b_active
        only_b = b_active & ~a_active

        result = torch.zeros_like(a)

        if both_active.any():
            alignment = compute_per_channel_alignment(a[both_active], b[both_active])
            # Merge in float32 then downcast
            a_fp32 = a[both_active].float()
            b_fp32 = b[both_active].float()
            merged_fp32 = a_fp32 * wa * alignment + b_fp32 * wb * alignment
            result[both_active] = merged_fp32.to(dtype)

        if only_a.any():
            result[only_a] = a[only_a] * wa

        if only_b.any():
            result[only_b] = b[only_b] * wb

        return result

    else:
        # Per-channel alignment (no full flatten)
        a_weighted = a * wa
        b_weighted = b * wb
        alignment = compute_per_channel_alignment(a_weighted, b_weighted)
        # Merge in float32 then downcast
        a_fp32 = a_weighted.float()
        b_fp32 = b_weighted.float()
        merged_fp32 = a_fp32 * alignment + b_fp32 * alignment
        return merged_fp32.to(dtype)

def merge_slerp(a, b, wa, wb, blend_mode="active", **kwargs):
    """
    Spherical Linear Interpolation (SLERP) merge — per-channel implementation.

    Reshapes to [channel_dim, -1] for per-channel SLERP instead of full flat 1D.
    This avoids OOM on large tensors (4D conv with millions of elements) and
    preserves per-filter directional information (C2 fix).

    Formula:
      t = |wa| / (|wa| + |wb|)
      result = sin((1-t)*θ)/sin(θ) * a + sin(t*θ)/sin(θ) * b

    NOTE: SLERP uses |weight| for the interpolation factor t, discarding
    the weight sign.  If you need subtractive effects, use the standard
    merge methods instead.

    Checkpoint compatibility:
        ✅ Works as expected on absolute checkpoint weights. Per-channel
           SLERP interpolates directions which is meaningful regardless
           of absolute vs delta weight representation.
        ✅ 4D conv weights handled via per-channel reshape [C_out, -1].
        ✅ 1D bias tensors fall back to linear interpolation.

    Edge cases (per-channel):
    - Zero-norm channels → linear interpolation fallback.
    - Near-parallel channels (cos_θ > 0.9995) → linear interpolation.
    - Antipodal channels (cos_θ < -0.9995) → linear interpolation (C3 fix).
    - NaN from 0/0 division → torch.nan_to_num safety.
    """
    dtype = a.dtype
    orig_shape = a.shape

    # Normalise weights to interpolation factor t ∈ [0, 1]
    total_abs = abs(wa) + abs(wb)
    if total_abs == 0:
        return torch.zeros_like(a)
    t = abs(wa) / total_abs

    # 1D tensors (e.g., bias): early return with linear interpolation
    # Per-channel SLERP is meaningless for single-element channels.
    if a.ndim <= 1:
        return a * (1 - t) + b * t

    result = apply_slerp(a, b, t)
    return result.to(dtype)


def apply_density(tensor, density):
    """Keep top 'density' percentage of weights, rescale to preserve L1 norm."""
    dtype = tensor.dtype
    
    if density >= 1:
        return tensor
    
    # Use float32 for the threshold calculation to avoid dtype issues
    flat = tensor.abs().float().flatten()
    k = max(1, int(flat.numel() * density))
    threshold = torch.topk(flat, k).values.min()
    
    # Create mask in original dtype
    # Use > instead of >= to avoid including more than k elements due to ties
    mask = (tensor.abs() > threshold).to(dtype)
    pruned = tensor * mask
    
    # Rescale to preserve L1 norm (sum of absolute values)
    original_l1 = tensor.abs().sum()
    new_l1 = pruned.abs().sum()
    if new_l1 > 0:
        scale = original_l1 / new_l1
        pruned = pruned * scale
    
    return pruned


# ==================== SHARED UTILITIES (consolidated from checkpoint_methods / triple_methods) ====================

def resolve_blend_mode_triple(blend_mode: str, metas: list) -> str:
    """
    Translate UI blend_mode to internal value for triple merge.

    Mapping:
    - "dense" -> "dense" (fall‑back weighted sum)
    - "auto" -> "active" if any trainer mismatch, "dense" if all three match.
    """
    if blend_mode == "dense":
        return "dense"

    if blend_mode == "auto":
        # Determine if all three trainers match
        if len(metas) < 3:
            # fallback
            return "active"
        trainers_match = True
        # Compare ss_network_module and ss_base_model across all metas
        nm_vals = [meta.get("ss_network_module") for meta in metas]
        bm_vals = [meta.get("ss_base_model") for meta in metas]
        # If any metadata missing, assume mismatch
        if any(v is None for v in nm_vals) and any(v is None for v in bm_vals):
            trainers_match = False
        else:
            # Check network_module equality (if present)
            if all(v is not None for v in nm_vals):
                if len(set(nm_vals)) != 1:
                    trainers_match = False
            # Check base_model equality (if present)
            if all(v is not None for v in bm_vals):
                if len(set(bm_vals)) != 1:
                    trainers_match = False
        if trainers_match:
            return "dense"
        else:
            return "active"

    # Pass through any unrecognised value
    return blend_mode


def apply_magnitude_scaling(tensors: List[torch.Tensor], valid_weights: List[float],
                            magnitude_scaling: str, max_scaling_factor: float = 10.0,
                            ref_idx: int = 0) -> List[torch.Tensor]:
    """
    Apply signal magnitude equalization (RMS/percentile scaling) to tensors.
    Scales each tensor to match the energy of the reference tensor (ref_idx).
    Returns scaled tensors (original tensors may be modified).
    """
    if magnitude_scaling == "none" or len(tensors) < 2:
        return tensors

    epsilon = 1e-8
    silence_threshold = 1e-6
    warning_threshold = 5.0

    energy_ref = get_tensor_energy(tensors[ref_idx], magnitude_scaling)
    scaling_factors = []
    skipped_count = 0
    non_ref_count = 0

    for j in range(len(tensors)):
        if j == ref_idx:
            continue
        non_ref_count += 1
        energy_j = get_tensor_energy(tensors[j], magnitude_scaling)
        if energy_j < silence_threshold:
            print(f"   🔇 [magnitude scaling] Energy_{j} ({energy_j:.2e}) below silence threshold; skipping scaling")
            skipped_count += 1
            continue
        raw_scale = energy_ref / (energy_j + epsilon)
        scale = raw_scale
        # Clamping to max scaling factor
        if raw_scale > max_scaling_factor:
            scale = max_scaling_factor
            print(f"   ⚠️ [magnitude scaling] Gain clamped to {max_scaling_factor:.1f}x (Original requirement: {raw_scale:.1f}x)")
        # Clipping warning (based on raw scale)
        if raw_scale > warning_threshold:
            print(f"   ⚠️ [magnitude scaling] Tensor {j} is significantly quieter than reference (energy={energy_j:.2e}, energy_ref={energy_ref:.2e}, scale={raw_scale:.1f}x)")
        tensors[j] = tensors[j] * scale
        if abs(scale - 1.0) > 0.01:
            scaling_factors.append(scale.item() if torch.is_tensor(scale) else scale)

    if skipped_count == non_ref_count and non_ref_count > 0:
        print(f"   ⚠️ [magnitude scaling] ALL non-reference tensors were below silence threshold ({silence_threshold:.0e}); no scaling applied")

    return tensors