"""
Stain Separation Module for Digital Pathology

Provides color deconvolution functionality for separating individual stains
in histopathology images, particularly H&E (Hematoxylin & Eosin) stained slides.
"""

from typing import Dict, Optional, Tuple

import numpy as np
from skimage.color import hed2rgb, rgb2hed

# Import stain matrix extraction from normalization module
from .stain_normalization import od_to_rgb, rgb_to_od


class StainSeparator:
    """
    Stain separation using color deconvolution.

    Separates individual stain contributions in histopathology images
    using various color deconvolution methods.
    """

    def __init__(self, method: str = "hed", custom_stain_matrix: Optional[np.ndarray] = None):
        """
        Initialize stain separator.

        Args:
            method: Separation method ("hed", "custom", "macenko")
            custom_stain_matrix: Custom stain matrix for separation (3x3)
        """
        self.method = method.lower()
        self.custom_stain_matrix = custom_stain_matrix

        # Default H&E stain matrix from Ruifrok & Johnston
        self.default_he_matrix = np.array(
            [
                [0.65, 0.70, 0.29],  # Hematoxylin
                [0.07, 0.99, 0.11],  # Eosin
                [0.29, 0.11, 0.90],  # Background/DAB
            ]
        )

    def separate(self, image: np.ndarray) -> Dict[str, np.ndarray]:
        """
        Separate stains from RGB image.

        Args:
            image: RGB image (HxWx3)

        Returns:
            Dictionary containing:
                - 'hematoxylin': Hematoxylin channel
                - 'eosin': Eosin channel
                - 'background': Background/residual channel
                - 'rgb_h': RGB visualization of hematoxylin
                - 'rgb_e': RGB visualization of eosin
                - 'concentrations': Raw stain concentrations
        """
        if self.method == "hed":
            return self._separate_hed(image)
        elif self.method == "macenko":
            return self._separate_macenko(image)
        elif self.method == "custom":
            if self.custom_stain_matrix is None:
                raise ValueError("Custom stain matrix required for custom method")
            return self._separate_custom(image, self.custom_stain_matrix)
        else:
            raise ValueError(f"Unknown separation method: {self.method}")

    def _separate_hed(self, image: np.ndarray) -> Dict[str, np.ndarray]:
        """Separate using built-in HED color space conversion."""
        # Convert to HED space
        hed = rgb2hed(image)

        # Extract channels
        h_channel = hed[:, :, 0]
        e_channel = hed[:, :, 1]
        d_channel = hed[:, :, 2]

        # Create RGB visualizations
        h_rgb = self._create_single_stain_rgb(h_channel, 0, hed.shape)
        e_rgb = self._create_single_stain_rgb(e_channel, 1, hed.shape)
        d_rgb = self._create_single_stain_rgb(d_channel, 2, hed.shape)

        return {
            "hematoxylin": h_channel,
            "eosin": e_channel,
            "background": d_channel,
            "rgb_h": h_rgb,
            "rgb_e": e_rgb,
            "rgb_d": d_rgb,
            "concentrations": hed,
            "method": "hed",
        }

    def _separate_macenko(self, image: np.ndarray) -> Dict[str, np.ndarray]:
        """Separate using Macenko stain matrix estimation."""
        from .stain_normalization import get_stain_matrix_macenko

        stain_matrix = get_stain_matrix_macenko(image)

        # Add third vector for completeness (orthogonal to first two)
        v1, v2 = stain_matrix[:, 0], stain_matrix[:, 1]
        v3 = np.cross(v1, v2)
        v3 = v3 / np.linalg.norm(v3)
        stain_matrix_full = np.column_stack([v1, v2, v3])

        return self._separate_custom(image, stain_matrix_full)

    def _separate_custom(
        self, image: np.ndarray, stain_matrix: np.ndarray
    ) -> Dict[str, np.ndarray]:
        """Separate using custom stain matrix."""
        # Convert to optical density
        od = rgb_to_od(image)
        od_flat = od.reshape(-1, 3)

        # Calculate concentrations using pseudo-inverse
        concentrations = od_flat @ np.linalg.pinv(stain_matrix).T
        concentrations = np.maximum(concentrations, 0)  # Non-negative
        concentrations = concentrations.reshape(image.shape[:2] + (stain_matrix.shape[1],))

        # Extract individual channels
        h_channel = concentrations[:, :, 0]
        e_channel = concentrations[:, :, 1]
        d_channel = (
            concentrations[:, :, 2] if concentrations.shape[2] > 2 else np.zeros_like(h_channel)
        )

        # Create RGB visualizations
        h_rgb = self._reconstruct_single_stain(h_channel, stain_matrix[:, 0])
        e_rgb = self._reconstruct_single_stain(e_channel, stain_matrix[:, 1])
        d_rgb = (
            self._reconstruct_single_stain(d_channel, stain_matrix[:, 2])
            if stain_matrix.shape[1] > 2
            else np.zeros_like(image)
        )

        return {
            "hematoxylin": h_channel,
            "eosin": e_channel,
            "background": d_channel,
            "rgb_h": h_rgb,
            "rgb_e": e_rgb,
            "rgb_d": d_rgb,
            "concentrations": concentrations,
            "stain_matrix": stain_matrix,
            "method": "custom",
        }

    def _create_single_stain_rgb(
        self, channel: np.ndarray, channel_idx: int, shape: Tuple
    ) -> np.ndarray:
        """Create RGB visualization of a single stain using HED space."""
        hed_single = np.zeros(shape)
        hed_single[:, :, channel_idx] = channel
        return hed2rgb(hed_single)

    def _reconstruct_single_stain(
        self, concentration: np.ndarray, stain_vector: np.ndarray
    ) -> np.ndarray:
        """Reconstruct RGB image from single stain concentration."""
        # Reshape concentration to flat array
        conc_flat = concentration.flatten()

        # Reconstruct OD for single stain
        od_single = np.outer(conc_flat, stain_vector)
        od_single = od_single.reshape(concentration.shape + (3,))

        # Convert to RGB
        rgb_single = od_to_rgb(od_single)
        return rgb_single

    def compute_stain_statistics(self, image: np.ndarray) -> Dict[str, Dict[str, float]]:
        """
        Compute statistics for each stain channel.

        Args:
            image: RGB image

        Returns:
            Dictionary with statistics for each stain
        """
        # Separate stains
        separation = self.separate(image)

        stats = {}
        for stain in ["hematoxylin", "eosin", "background"]:
            channel = separation[stain]
            stats[stain] = {
                "mean": float(np.mean(channel)),
                "std": float(np.std(channel)),
                "min": float(np.min(channel)),
                "max": float(np.max(channel)),
                "median": float(np.median(channel)),
                "percentile_25": float(np.percentile(channel, 25)),
                "percentile_75": float(np.percentile(channel, 75)),
                "percentile_99": float(np.percentile(channel, 99)),
            }

        return stats

    def enhance_stain(
        self, image: np.ndarray, h_factor: float = 1.0, e_factor: float = 1.0, d_factor: float = 1.0
    ) -> np.ndarray:
        """
        Enhance or suppress specific stains.

        Args:
            image: RGB image
            h_factor: Multiplication factor for hematoxylin
            e_factor: Multiplication factor for eosin
            d_factor: Multiplication factor for background/DAB

        Returns:
            Modified RGB image
        """
        # Separate stains
        separation = self.separate(image)

        if self.method == "hed":
            # Modify concentrations
            hed_modified = separation["concentrations"].copy()
            hed_modified[:, :, 0] *= h_factor
            hed_modified[:, :, 1] *= e_factor
            hed_modified[:, :, 2] *= d_factor

            # Clip to valid range
            hed_modified = np.clip(hed_modified, 0, None)

            # Reconstruct
            return hed2rgb(hed_modified)
        else:
            # For custom/macenko methods
            concentrations = separation["concentrations"]
            stain_matrix = separation.get("stain_matrix", self.default_he_matrix)

            # Modify concentrations
            modified_conc = concentrations.copy()
            modified_conc[:, :, 0] *= h_factor
            modified_conc[:, :, 1] *= e_factor
            if modified_conc.shape[2] > 2:
                modified_conc[:, :, 2] *= d_factor

            # Reconstruct
            od_reconstructed = modified_conc.reshape(-1, modified_conc.shape[2]) @ stain_matrix.T
            od_reconstructed = od_reconstructed.reshape(image.shape)

            return od_to_rgb(od_reconstructed)


def separate_stains(image: np.ndarray, method: str = "hed") -> Dict[str, np.ndarray]:
    """
    Convenience function for quick stain separation.

    Args:
        image: RGB image
        method: Separation method ("hed", "macenko")

    Returns:
        Dictionary with separated stain channels
    """
    separator = StainSeparator(method=method)
    return separator.separate(image)


def get_stain_concentrations(image: np.ndarray, method: str = "hed") -> np.ndarray:
    """
    Get raw stain concentrations.

    Args:
        image: RGB image
        method: Separation method

    Returns:
        Stain concentrations array
    """
    separator = StainSeparator(method=method)
    result = separator.separate(image)
    return result["concentrations"]


def visualize_stains(image: np.ndarray, method: str = "hed") -> Dict[str, np.ndarray]:
    """
    Get RGB visualizations of individual stains.

    Args:
        image: RGB image
        method: Separation method

    Returns:
        Dictionary with RGB visualizations
    """
    separator = StainSeparator(method=method)
    result = separator.separate(image)
    return {
        "original": image,
        "hematoxylin": result["rgb_h"],
        "eosin": result["rgb_e"],
        "background": result.get("rgb_d", np.zeros_like(image)),
    }
