"""Harmony teacher for batch correction.

Simple wrapper around the harmonypy library.
"""

import numpy as np
import pandas as pd
from typing import Optional
import logging

from .base import BaseTeacher

logger = logging.getLogger(__name__)


class HarmonyTeacher(BaseTeacher):
    """Harmony batch correction using harmonypy.

    Parameters
    ----------
    theta : float, default=2.0
        Diversity penalty for batch correction.
    max_iter : int, default=10
        Maximum Harmony iterations.

    Examples
    --------
    >>> teacher = HarmonyTeacher(theta=2.0)
    >>> Z_corrected = teacher.fit_transform(X_pca, batch_labels)
    """

    name = "HarmonyTeacher"

    def __init__(
        self,
        theta: float = 2.0,
        max_iter: int = 10,
    ):
        super().__init__()
        self.theta = theta
        self.max_iter = max_iter

    def fit_transform(
        self,
        X_pca: np.ndarray,
        batch_labels: np.ndarray,
        obs: Optional[pd.DataFrame] = None,
        verbose: bool = True,
        **kwargs,
    ) -> np.ndarray:
        """Apply Harmony batch correction.

        Parameters
        ----------
        X_pca : np.ndarray
            PCA embeddings (n_cells, n_pcs).
        batch_labels : np.ndarray
            Batch labels (n_cells,).
        obs : pd.DataFrame, optional
            Not used (for API compatibility).
        verbose : bool
            Print progress.
        **kwargs
            Additional arguments (ignored).

        Returns
        -------
        Z_corrected : np.ndarray
            Harmony-corrected embeddings (n_cells, n_pcs).
        """
        import harmonypy

        if verbose:
            logger.info(f"{self.name}: theta={self.theta}, max_iter={self.max_iter}")

        # Ensure X_pca is (N, d)
        if X_pca.shape[0] < X_pca.shape[1]:
            X_pca = X_pca.T

        # Create DataFrame for batch labels
        meta_data = pd.DataFrame({'batch': batch_labels})

        # Run harmonypy
        harmony_out = harmonypy.run_harmony(
            X_pca,
            meta_data,
            'batch',
            theta=self.theta,
            max_iter_harmony=self.max_iter,
            verbose=verbose,
        )

        # Get corrected embedding (harmonypy returns Z_corr as (d, N))
        Z_corr = harmony_out.Z_corr.T  # Convert to (N, d)

        self.Z_corrected = Z_corr
        self._is_fitted = True
        return self.Z_corrected

    def __repr__(self) -> str:
        return f"HarmonyTeacher(theta={self.theta}, max_iter={self.max_iter})"
