Multiresolution Analysis

def mra(data, wavelet, level: int = None, axis: int = -1, 
        transform: str = 'swt', mode: str = 'periodization'):
    """
    1D multiresolution analysis.
    
    Parameters:
    - data: Input 1D signal
    - wavelet: Wavelet specification
    - level: Number of decomposition levels (default: maximum possible)
    - axis: Axis along which to perform MRA (default: -1)
    - transform: Transform type ('swt' for stationary WT, 'dwt' for discrete WT)
    - mode: Signal extension mode (default: 'periodization')
    
    Returns:
    List of approximation components [A1, A2, ..., An] where:
    - A1: Finest scale approximation
    - An: Coarsest scale approximation
    """

def imra(mra_coeffs):
    """
    Inverse 1D MRA via summation.
    
    Parameters:
    - mra_coeffs: List of approximation components from mra()
    
    Returns:
    Reconstructed signal (sum of all approximation levels)
    """

import pywt
import numpy as np
import matplotlib.pyplot as plt

# Create test signal with multiple scales
np.random.seed(42)
t = np.linspace(0, 4, 2048)
signal = (2 * np.sin(2 * np.pi * 1 * t) +          # Slow trend
          1.5 * np.sin(2 * np.pi * 8 * t) +        # Medium frequency
          1 * np.sin(2 * np.pi * 32 * t) +         # Fast oscillation
          0.5 * np.sin(2 * np.pi * 100 * t) +      # High frequency
          0.3 * np.random.randn(len(t)))           # Noise

print(f"Signal length: {len(signal)}")

# Perform multiresolution analysis
level = 8
mra_coeffs = pywt.mra(signal, 'db8', level=level, transform='swt')
print(f"MRA decomposition levels: {len(mra_coeffs)}")
print(f"Each approximation shape: {mra_coeffs[0].shape}")

# Verify perfect reconstruction
reconstructed = pywt.imra(mra_coeffs)  
reconstruction_error = np.max(np.abs(signal - reconstructed))
print(f"MRA reconstruction error: {reconstruction_error:.2e}")

# Analyze frequency content at each scale
def analyze_frequency_content(approx_component, sampling_rate=1.0, title=""):
    """Analyze dominant frequency in approximation component."""
    # Simple frequency analysis using FFT
    fft_vals = np.fft.fft(approx_component)
    freqs = np.fft.fftfreq(len(approx_component), 1/sampling_rate)
    
    # Find dominant frequency (excluding DC component)
    magnitude = np.abs(fft_vals[1:len(fft_vals)//2])
    freq_bins = freqs[1:len(freqs)//2]
    dominant_idx = np.argmax(magnitude)
    dominant_freq = freq_bins[dominant_idx]
    
    return dominant_freq, magnitude, freq_bins

sampling_rate = len(t) / (t[-1] - t[0])  # Samples per second
print(f"\nFrequency analysis (sampling rate: {sampling_rate:.1f} Hz):")

dominant_frequencies = []
for i, approx in enumerate(mra_coeffs):
    dom_freq, _, _ = analyze_frequency_content(approx, sampling_rate)
    dominant_frequencies.append(dom_freq)
    scale_factor = 2**(i+1)
    expected_max_freq = sampling_rate / (2 * scale_factor)
    print(f"Level {i+1}: dominant freq = {dom_freq:.2f} Hz, max expected = {expected_max_freq:.2f} Hz")

# Visualization of multiresolution components
fig, axes = plt.subplots(min(6, len(mra_coeffs) + 1), 1, figsize=(15, 12))

# Original signal
axes[0].plot(t, signal, 'k-', linewidth=1)
axes[0].set_title('Original Signal')
axes[0].set_ylabel('Amplitude')
axes[0].grid(True, alpha=0.3)

# Plot first 5 approximation levels
for i in range(min(5, len(mra_coeffs))):
    axes[i+1].plot(t, mra_coeffs[i], 'b-', linewidth=1.5)
    axes[i+1].set_title(f'Approximation Level {i+1} (Scale 2^{i+1})')
    axes[i+1].set_ylabel('Amplitude')
    axes[i+1].grid(True, alpha=0.3)

axes[-1].set_xlabel('Time')
plt.tight_layout()
plt.show()

# Compare different transform backends
transforms = ['swt', 'dwt']
mra_comparison = {}

for transform_type in transforms:
    try:
        if transform_type == 'dwt':
            # DWT-based MRA may have different requirements
            mra_dwt = pywt.mra(signal, 'db8', level=6, transform='dwt')
            mra_comparison[transform_type] = mra_dwt
        else:
            mra_comparison[transform_type] = mra_coeffs
        
        print(f"\n{transform_type.upper()}-based MRA:")
        mra_result = mra_comparison[transform_type]
        print(f"  Levels: {len(mra_result)}")
        print(f"  Component shapes: {[comp.shape for comp in mra_result[:3]]}...")
        
        # Reconstruction test
        reconstructed_comp = pywt.imra(mra_result)
        if len(reconstructed_comp) == len(signal):
            error = np.max(np.abs(signal - reconstructed_comp))
            print(f"  Reconstruction error: {error:.2e}")
        else:
            print(f"  Reconstruction shape: {reconstructed_comp.shape} (vs {signal.shape})")
            
    except Exception as e:
        print(f"Error with {transform_type}: {e}")

# Adaptive signal analysis using MRA
def adaptive_denoising_mra(signal, wavelet, noise_level_estimate=None):
    """Adaptive denoising using MRA with scale-dependent thresholding."""
    
    # Perform MRA
    mra_components = pywt.mra(signal, wavelet, transform='swt')
    
    if noise_level_estimate is None:
        # Estimate noise from finest scale approximation
        finest_approx = mra_components[0]
        noise_level_estimate = np.std(finest_approx - np.median(finest_approx))
    
    print(f"Estimated noise level: {noise_level_estimate:.4f}")
    
    # Apply scale-dependent denoising
    denoised_components = []
    for i, component in enumerate(mra_components):
        scale = 2**(i+1)
        # Threshold decreases with scale (coarser scales need less denoising)
        threshold = noise_level_estimate / np.sqrt(scale)
        
        # Apply soft thresholding
        denoised = pywt.threshold(component, threshold, mode='soft')
        denoised_components.append(denoised)
        
        print(f"Level {i+1}: threshold = {threshold:.4f}")
    
    # Reconstruct denoised signal
    return pywt.imra(denoised_components)

# Test adaptive denoising
noisy_signal = signal + 0.5 * np.random.randn(len(signal))
denoised_mra = adaptive_denoising_mra(noisy_signal, 'db8', noise_level_estimate=0.5)

# Compare with standard wavelet denoising
coeffs_std = pywt.wavedec(noisy_signal, 'db8', level=8)
coeffs_thresh = [coeffs_std[0]]  # Keep approximation
for detail in coeffs_std[1:]:
    threshold = 0.1 * np.std(detail)
    coeffs_thresh.append(pywt.threshold(detail, threshold, mode='soft'))
denoised_std = pywt.waverec(coeffs_thresh, 'db8')

# Performance comparison
def calculate_snr(clean, noisy):
    return 10 * np.log10(np.var(clean) / np.var(clean - noisy))

original_snr = calculate_snr(signal, noisy_signal)
mra_snr = calculate_snr(signal, denoised_mra)
std_snr = calculate_snr(signal, denoised_std)

print(f"\nDenoising Performance:")
print(f"Original SNR: {original_snr:.2f} dB")
print(f"MRA denoising SNR: {mra_snr:.2f} dB")
print(f"Standard denoising SNR: {std_snr:.2f} dB")

# Plot denoising comparison
plt.figure(figsize=(15, 10))

plt.subplot(2, 2, 1)
plt.plot(t, signal, 'b-', label='Clean signal', linewidth=2)
plt.plot(t, noisy_signal, 'r-', alpha=0.7, label='Noisy signal')
plt.title('Original vs Noisy Signal')
plt.legend()
plt.grid(True)

plt.subplot(2, 2, 2)
plt.plot(t, signal, 'b-', alpha=0.7, label='Clean')
plt.plot(t, denoised_mra, 'g-', linewidth=2, label='MRA denoised')
plt.title(f'MRA Denoising (SNR: {mra_snr:.2f} dB)')
plt.legend()
plt.grid(True)

plt.subplot(2, 2, 3)
plt.plot(t, signal, 'b-', alpha=0.7, label='Clean')
plt.plot(t, denoised_std, 'm-', linewidth=2, label='Standard denoised')
plt.title(f'Standard Denoising (SNR: {std_snr:.2f} dB)')
plt.legend()
plt.grid(True)

plt.subplot(2, 2, 4)
plt.plot(t, signal - denoised_mra, 'g-', label='MRA error')
plt.plot(t, signal - denoised_std, 'm-', alpha=0.7, label='Standard error')
plt.title('Denoising Errors')
plt.legend()
plt.grid(True)

plt.tight_layout()
plt.show()

def mra2(data, wavelet, level: int = None, axes=(-2, -1), 
         transform: str = 'swt2', mode: str = 'periodization'):
    """
    2D multiresolution analysis.
    
    Parameters:
    - data: Input 2D array (image)
    - wavelet: Wavelet specification
    - level: Number of decomposition levels (default: maximum possible)
    - axes: Pair of axes for 2D transform (default: last two axes)
    - transform: Transform type ('swt2' for 2D stationary WT)
    - mode: Signal extension mode (default: 'periodization')
    
    Returns:
    List of 2D approximation components at different scales
    """

def imra2(mra_coeffs):
    """
    Inverse 2D MRA via summation.
    
    Parameters:
    - mra_coeffs: List of 2D approximation components from mra2()
    
    Returns:
    Reconstructed 2D array (sum of all approximation levels)
    """

import pywt
import numpy as np
import matplotlib.pyplot as plt

# Create test image with multiple scales
size = 256
x, y = np.mgrid[0:size, 0:size]

# Multi-scale image with different frequency content
image = (0.5 +                                                    # DC component
         0.4 * np.sin(2 * np.pi * x / 64) * np.cos(2 * np.pi * y / 64) +   # Low frequency
         0.3 * np.sin(2 * np.pi * x / 16) * np.cos(2 * np.pi * y / 16) +   # Medium frequency
         0.2 * np.sin(2 * np.pi * x / 4) * np.cos(2 * np.pi * y / 4))      # High frequency

# Add texture and noise
texture = 0.1 * np.random.randn(size, size)
image = image + texture

print(f"Image shape: {image.shape}")
print(f"Image value range: [{np.min(image):.3f}, {np.max(image):.3f}]")

# 2D Multiresolution analysis
level = 5
mra2_coeffs = pywt.mra2(image, 'db4', level=level, transform='swt2')
print(f"2D MRA levels: {len(mra2_coeffs)}")
print(f"Approximation shapes: {[approx.shape for approx in mra2_coeffs]}")

# Perfect reconstruction
reconstructed_2d = pywt.imra2(mra2_coeffs)
reconstruction_error_2d = np.max(np.abs(image - reconstructed_2d))
print(f"2D MRA reconstruction error: {reconstruction_error_2d:.2e}")

# Visualize 2D multiresolution components  
fig, axes = plt.subplots(2, 3, figsize=(18, 12))
axes = axes.flatten()

# Original image
axes[0].imshow(image, cmap='gray')
axes[0].set_title('Original Image')
axes[0].axis('off')

# Show first 5 approximation levels
for i in range(min(5, len(mra2_coeffs))):
    axes[i+1].imshow(mra2_coeffs[i], cmap='gray')
    scale = 2**(i+1)
    axes[i+1].set_title(f'Approximation Level {i+1} (Scale 2^{i+1})')
    axes[i+1].axis('off')

plt.tight_layout()
plt.show()

# Scale-space analysis
def analyze_image_scales(mra_components):
    """Analyze image content at different scales."""
    analysis = {}
    
    for i, component in enumerate(mra_components):
        scale = 2**(i+1)
        
        # Calculate statistics
        mean_val = np.mean(component)
        std_val = np.std(component)
        energy = np.sum(component**2)
        
        # Edge content (simple gradient measure)
        grad_x = np.diff(component, axis=1)
        grad_y = np.diff(component, axis=0)
        edge_strength = np.mean(np.sqrt(grad_x[:-1,:]**2 + grad_y[:,:-1]**2))
        
        analysis[i+1] = {
            'scale': scale,
            'mean': mean_val,
            'std': std_val,
            'energy': energy,
            'edge_strength': edge_strength,
            'shape': component.shape
        }
    
    return analysis

scale_analysis = analyze_image_scales(mra2_coeffs)

print("\n2D Scale-Space Analysis:")
print(f"{'Level':<6} {'Scale':<6} {'Mean':<8} {'Std':<8} {'Energy':<10} {'Edges':<8}")
print("-" * 55)

for level, stats in scale_analysis.items():
    print(f"{level:<6} {stats['scale']:<6} {stats['mean']:<8.4f} {stats['std']:<8.4f} "
          f"{stats['energy']:<10.2f} {stats['edge_strength']:<8.4f}")

# Multi-scale image enhancement
def multiscale_enhancement(image, wavelet, enhancement_factors=None):
    """Enhance image using multi-scale processing."""
    
    mra_components = pywt.mra2(image, wavelet, transform='swt2')
    
    if enhancement_factors is None:
        # Default: enhance fine details more than coarse
        enhancement_factors = [1.5, 1.3, 1.1, 1.0, 0.9]
    
    # Ensure we have enough factors
    enhancement_factors = enhancement_factors[:len(mra_components)]
    enhancement_factors += [1.0] * (len(mra_components) - len(enhancement_factors))
    
    # Apply enhancement
    enhanced_components = []
    for i, (component, factor) in enumerate(zip(mra_components, enhancement_factors)):
        if factor != 1.0:
            # Enhance by scaling around the mean
            mean_val = np.mean(component)
            enhanced = mean_val + factor * (component - mean_val)
        else:
            enhanced = component
        
        enhanced_components.append(enhanced)
        print(f"Level {i+1}: enhancement factor = {factor}")
    
    return pywt.imra2(enhanced_components)

# Apply enhancement
enhanced_image = multiscale_enhancement(image, 'db4', [2.0, 1.5, 1.2, 1.0, 0.8])

# Compare original and enhanced
plt.figure(figsize=(15, 5))

plt.subplot(1, 3, 1)
plt.imshow(image, cmap='gray')
plt.title('Original Image')
plt.axis('off')

plt.subplot(1, 3, 2)
plt.imshow(enhanced_image, cmap='gray')
plt.title('Enhanced Image')
plt.axis('off')

plt.subplot(1, 3, 3)
plt.imshow(enhanced_image - image, cmap='RdBu_r')
plt.title('Enhancement Difference')
plt.colorbar()
plt.axis('off')

plt.tight_layout()
plt.show()

# Calculate enhancement metrics
original_edge_strength = np.mean(np.abs(np.gradient(image)[0]) + np.abs(np.gradient(image)[1]))
enhanced_edge_strength = np.mean(np.abs(np.gradient(enhanced_image)[0]) + np.abs(np.gradient(enhanced_image)[1]))

print(f"\nEnhancement Results:")
print(f"Original edge strength: {original_edge_strength:.4f}")
print(f"Enhanced edge strength: {enhanced_edge_strength:.4f}")
print(f"Enhancement ratio: {enhanced_edge_strength / original_edge_strength:.2f}")

def mran(data, wavelet, level: int = None, axes=None, 
         transform: str = 'swtn', mode: str = 'periodization'):
    """
    nD multiresolution analysis.
    
    Parameters:
    - data: Input nD array
    - wavelet: Wavelet specification
    - level: Number of decomposition levels (default: maximum possible)
    - axes: Axes along which to perform transform (default: all axes)
    - transform: Transform type ('swtn' for nD stationary WT)
    - mode: Signal extension mode (default: 'periodization')
    
    Returns:
    List of nD approximation components at different scales
    """

def imran(mra_coeffs):
    """
    Inverse nD MRA via summation.
    
    Parameters:
    - mra_coeffs: List of nD approximation components from mran()
    
    Returns:
    Reconstructed nD array (sum of all approximation levels)
    """

import pywt
import numpy as np

# 3D volume multiresolution analysis
volume_shape = (64, 64, 64)
x, y, z = np.mgrid[0:volume_shape[0], 0:volume_shape[1], 0:volume_shape[2]]

# Create 3D test volume with multiple scales
volume = (0.5 * np.sin(2 * np.pi * x / 32) * np.cos(2 * np.pi * y / 32) * np.sin(2 * np.pi * z / 32) +
          0.3 * np.sin(2 * np.pi * x / 8) * np.cos(2 * np.pi * y / 8) * np.sin(2 * np.pi * z / 8) +
          0.1 * np.random.randn(*volume_shape))

print(f"3D volume shape: {volume.shape}")

# 3D MRA
mra3d_coeffs = pywt.mran(volume, 'haar', level=4, transform='swtn')
print(f"3D MRA levels: {len(mra3d_coeffs)}")
print(f"3D approximation shapes: {[approx.shape for approx in mra3d_coeffs]}")

# Perfect reconstruction
reconstructed_3d = pywt.imran(mra3d_coeffs)
reconstruction_error_3d = np.max(np.abs(volume - reconstructed_3d))
print(f"3D MRA reconstruction error: {reconstruction_error_3d:.2e}")

# Analyze 3D scales
print("\n3D Scale Analysis:")
for i, component in enumerate(mra3d_coeffs):
    scale = 2**(i+1)
    energy = np.sum(component**2)
    mean_val = np.mean(component)
    std_val = np.std(component)
    print(f"Level {i+1} (scale {scale}): energy = {energy:.2f}, mean = {mean_val:.4f}, std = {std_val:.4f}")

# Time series MRA (1D signal in higher dimensional context)
time_series_length = 10000
time_series = np.cumsum(np.random.randn(time_series_length)) * 0.1  # Random walk
time_series += np.sin(2 * np.pi * np.arange(time_series_length) / 100)  # Add trend

# Reshape as pseudo-2D for demonstration
ts_reshaped = time_series.reshape(100, 100)
mra_ts = pywt.mra2(ts_reshaped, 'db8', level=6)

print(f"\nTime series MRA:")
print(f"Original time series length: {time_series_length}")
print(f"Reshaped as: {ts_reshaped.shape}")
print(f"MRA levels: {len(mra_ts)}")

# Reconstruct and verify
reconstructed_ts = pywt.imra2(mra_ts)
ts_error = np.max(np.abs(ts_reshaped - reconstructed_ts))
print(f"Time series MRA error: {ts_error:.2e}")

# Demonstrate memory efficiency of MRA vs full wavelet decomposition
def compare_memory_usage(data_shape, wavelet, level):
    """Compare memory usage between MRA and full wavelet decomposition."""
    
    # Create dummy data
    data = np.random.randn(*data_shape)
    
    # MRA memory (approximations only)
    mra_result = pywt.mra(data.ravel(), wavelet, level=level) if len(data_shape) == 1 else \
                 pywt.mra2(data, wavelet, level=level) if len(data_shape) == 2 else \
                 pywt.mran(data, wavelet, level=level)
    
    mra_memory = sum(component.nbytes for component in mra_result)
    
    # Full decomposition memory (approximation + all details)
    if len(data_shape) == 1:
        full_coeffs = pywt.wavedec(data.ravel(), wavelet, level=level)
        full_memory = sum(coeff.nbytes for coeff in full_coeffs)
    elif len(data_shape) == 2:
        full_coeffs = pywt.wavedec2(data, wavelet, level=level)
        full_memory = full_coeffs[0].nbytes + sum(sum(detail.nbytes for detail in level_details) 
                                                 for level_details in full_coeffs[1:])
    else:
        full_coeffs = pywt.wavedecn(data, wavelet, level=level)
        full_memory = full_coeffs[0].nbytes + sum(sum(detail.nbytes for detail in level_details.values()) 
                                                 for level_details in full_coeffs[1:])
    
    original_memory = data.nbytes
    
    return {
        'original': original_memory,
        'mra': mra_memory,
        'full': full_memory,
        'mra_ratio': mra_memory / original_memory,
        'full_ratio': full_memory / original_memory
    }

# Memory comparison for different data sizes
test_shapes = [(2048,), (512, 512), (64, 64, 64)]

print("\nMemory Usage Comparison:")
print(f"{'Shape':<15} {'Original (MB)':<12} {'MRA (MB)':<10} {'Full (MB)':<10} {'MRA Ratio':<10} {'Full Ratio':<10}")
print("-" * 80)

for shape in test_shapes:
    memory_info = compare_memory_usage(shape, 'db4', 5)
    print(f"{str(shape):<15} {memory_info['original']/1024**2:<12.2f} "
          f"{memory_info['mra']/1024**2:<10.2f} {memory_info['full']/1024**2:<10.2f} "
          f"{memory_info['mra_ratio']:<10.2f} {memory_info['full_ratio']:<10.2f}")

# MRA transform backends
MRATransform1D = Literal['swt', 'dwt']
MRATransform2D = Literal['swt2', 'dwt2']  
MRATransformND = Literal['swtn', 'dwtn']

# MRA coefficient format (list of approximation components)
MRACoeffs1D = List[np.ndarray]  # [A1, A2, ..., An] - finest to coarsest
MRACoeffs2D = List[np.ndarray]  # [A1, A2, ..., An] - 2D approximations
MRACoeffsND = List[np.ndarray]  # [A1, A2, ..., An] - nD approximations