mapper.py

import numpy as np
from matplotlib.colors import hsv_to_rgb
from typing import Protocol, runtime_checkable

from normalizers import (
    MagnitudeNormalizer,
    RationalNormalizer,
    LinearNormalizer,
    LogNormalizer,
)
from strategies import ColorCombineStrategy, ValueStrategy


@runtime_checkable
class SupportsToRgb(Protocol):
    def to_rgb(self, mag_norm: np.ndarray) -> np.ndarray:
        ...


class ComplexColorMapper:
    def __init__(
        self,
        normalizer: MagnitudeNormalizer | None = None,
        combine_strategy: ColorCombineStrategy | None = None,
    ):
        self.normalizer = normalizer or RationalNormalizer()
        self.combine_strategy = combine_strategy or ValueStrategy()

    def to_rgb(self, arr: np.ndarray, reference_max: float | None = None) -> np.ndarray:
        z = np.asarray(arr)

        if z.ndim != 2:
            raise ValueError(f"Expected 2D array, got shape {z.shape}")

        phase = np.angle(z)
        mag = np.abs(z)

        if reference_max is None:
            mag_norm = np.clip(self.normalizer(mag), 0.0, 1.0)
        elif isinstance(self.normalizer, LinearNormalizer):
            mag_norm = mag / reference_max if reference_max > 0 else mag
        elif isinstance(self.normalizer, LogNormalizer):
            mag_norm = (
                np.log1p(mag) / np.log1p(reference_max)
                if reference_max > 0
                else mag
            )
        else:
            mag_norm = np.clip(self.normalizer(mag), 0.0, 1.0)

        if isinstance(self.combine_strategy, SupportsToRgb):
            return self.combine_strategy.to_rgb(mag_norm)

        hue = (phase / (2 * np.pi)) % 1.0
        saturation, value = self.combine_strategy.compute_sv(mag_norm)

        hsv = np.stack([hue, saturation, value], axis=-1)
        return hsv_to_rgb(hsv)