"""Preprocessing utilities for tkmlr."""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict, Iterable, Optional

import numpy as np


@dataclass
class Standardizer:
    """Simple standardizer independent of sklearn.

    sigma=0 columns are left unchanged by using sigma=1.0.
    """

    mean_: Optional[np.ndarray] = None
    scale_: Optional[np.ndarray] = None

    def fit(self, X: Any) -> "Standardizer":
        X_np = np.asarray(X, dtype=float)
        self.mean_ = np.mean(X_np, axis=0)
        scale = np.std(X_np, axis=0)
        scale = np.where(scale == 0.0, 1.0, scale)
        self.scale_ = scale
        return self

    def transform(self, X: Any) -> np.ndarray:
        if self.mean_ is None or self.scale_ is None:
            raise RuntimeError("Standardizer is not fitted.")
        return (np.asarray(X, dtype=float) - self.mean_) / self.scale_

    def fit_transform(self, X: Any) -> np.ndarray:
        return self.fit(X).transform(X)

    def inverse_transform(self, X: Any) -> np.ndarray:
        if self.mean_ is None or self.scale_ is None:
            raise RuntimeError("Standardizer is not fitted.")
        return np.asarray(X, dtype=float) * self.scale_ + self.mean_


@dataclass
class LogTransformer:
    """Column-wise log transform.

    ``log_columns`` can be column names or integer column indices.
    """

    log_columns: Iterable[Any]

    def transform_dataframe(self, df):
        import numpy as np

        out = df.copy()
        for col in self.log_columns:
            out[col] = np.log(out[col])
        return out