import numpy as np
import nibabel as nib
from intensity_normalization import normalize_image

# Load MRI image and brain mask
img = nib.load("brain_t1.nii.gz")
mask = nib.load("brain_mask.nii.gz")

# Normalize using FCM (default, recommended for T1)
normalized = normalize_image(img, method="fcm", mask=mask)

# Different methods for different modalities
t2_normalized = normalize_image(img, method="kde", modality="t2")
zscore_normalized = normalize_image(img, method="zscore")

from intensity_normalization import FCMNormalizer, ZScoreNormalizer
import numpy as np

# Create synthetic brain data
brain_data = np.random.normal(1000, 200, (64, 64, 32))  # WM ~1000
brain_data[20:40, 20:40, 10:20] = np.random.normal(600, 100, (20, 20, 10))  # GM ~600

# FCM normalization (tissue-based)
fcm = FCMNormalizer(tissue_type="wm")  # white matter reference
normalized = fcm.fit_transform(brain_data)

# Z-score normalization
zscore = ZScoreNormalizer()
standardized = zscore.fit_transform(brain_data)

print(f"Original mean: {brain_data[brain_data > 0].mean():.1f}")
print(f"FCM normalized WM mean: {normalized[40:60, 40:60, 10:20].mean():.2f}")
print(f"Z-score mean: {standardized[brain_data > 0].mean():.2f}")

from intensity_normalization import NyulNormalizer, LSQNormalizer
from intensity_normalization.adapters import create_image

# Load multiple subjects
image_paths = ["subject1_t1.nii.gz", "subject2_t1.nii.gz", "subject3_t1.nii.gz"]
images = [create_image(path) for path in image_paths]

# Nyúl histogram matching
nyul = NyulNormalizer(output_min_value=0, output_max_value=100)
nyul.fit_population(images)

# Normalize all images to the same scale
normalized_images = [nyul.transform(img) for img in images]

# LSQ tissue mean harmonization
lsq = LSQNormalizer()
lsq.fit_population(images)
harmonized_images = [lsq.transform(img) for img in images]

# T1-weighted images (structural)
normalize_image(t1_image, method="fcm", tissue_type="wm")

# T2-weighted or FLAIR
normalize_image(t2_image, method="kde", modality="t2")

# Quick standardization
normalize_image(image, method="zscore")

# Multi-site harmonization (requires multiple subjects)
from intensity_normalization.services.normalization import NormalizationService
config = NormalizationConfig(method="nyul")
harmonized = NormalizationService.normalize_images(all_images, config)

from intensity_normalization.domain.protocols import BaseNormalizer
from intensity_normalization.domain.protocols import ImageProtocol

class CustomNormalizer(BaseNormalizer):
    def fit(self, image: ImageProtocol, mask=None) -> 'CustomNormalizer':
        # Implement fitting logic
        self.is_fitted = True
        return self

    def transform(self, image: ImageProtocol, mask=None) -> ImageProtocol:
        # Implement normalization logic
        data = image.get_data()
        normalized_data = your_normalization_function(data)
        return image.with_data(normalized_data)

from intensity_normalization.domain.models import NormalizationConfig, Modality, TissueType
from intensity_normalization.services.normalization import NormalizationService

# Create configuration
config = NormalizationConfig(
    method="fcm",
    modality=Modality.T1,
    tissue_type=TissueType.WM
)

# Validate configuration
from intensity_normalization.services import ValidationService
ValidationService.validate_normalization_config(config)

# Apply normalization
result = NormalizationService.normalize_image(image, config, mask)

from pathlib import Path
from intensity_normalization.adapters.images import create_image, save_image
from intensity_normalization.services.normalization import NormalizationService

def process_directory(input_dir: Path, output_dir: Path, method: str = "fcm"):
    """Process all NIfTI files in a directory."""
    output_dir.mkdir(exist_ok=True)

    for img_file in input_dir.glob("*.nii.gz"):
        # Load image
        image = create_image(img_file)

        # Create configuration
        config = NormalizationConfig(method=method)

        # Normalize
        normalized = NormalizationService.normalize_image(image, config)

        # Save result
        output_file = output_dir / f"{img_file.stem}_normalized.nii.gz"
        save_image(normalized, output_file)
        print(f"Processed: {img_file.name}")

# Usage
process_directory(Path("raw_images/"), Path("normalized/"))