"""Built-in linear functionals (estimands).

The central object in *genriesz* is a (typically) linear functional

    theta = E[ m(X, gamma0) ],

where ``gamma0(x) = E[Y | X=x]`` is the outcome regression and ``m`` is a
user-specified linear operator acting on functions.

For models ``v(x) = phi(x)^T beta`` (where ``phi`` is a basis), linearity means

    m(X_i, v) = M_i^T beta

for some row vector ``M_i`` that depends on ``X_i`` and the basis.

This module provides built-in functionals used in the notebooks:

- ATE, ATT, and DID (as ATT on delta outcomes)
- AME (average marginal effect / average derivative)
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

import numpy as np
from numpy.typing import ArrayLike, NDArray

from .basis import Basis


def _as_2d(X: ArrayLike) -> NDArray[np.float64]:
    X_ = np.asarray(X, dtype=float)
    if X_.ndim != 2:
        raise ValueError(f"X must be 2D. Got shape {X_.shape}.")
    return X_


def _toggle_treatment(
    X: NDArray[np.float64], *, treatment_index: int, value: float
) -> NDArray[np.float64]:
    X_cf = X.copy()
    X_cf[:, treatment_index] = float(value)
    return X_cf


PredictFn = Callable[[NDArray[np.float64]], NDArray[np.float64]]
DerivFn = Callable[[NDArray[np.float64], int], NDArray[np.float64]]


# A scalar-valued function on a single regressor row.
GammaFn = Callable[[NDArray[np.float64]], float]


def _as_1d_row(x: ArrayLike) -> NDArray[np.float64]:
    return np.asarray(x, dtype=float).reshape(-1)


@dataclass(frozen=True)
class LinearFunctional:
    """Base class for linear functionals used by GRR."""

    name: str

    def m_basis_matrix(self, X: ArrayLike, basis: Basis) -> NDArray[np.float64]:  # pragma: no cover
        raise NotImplementedError

    def m_from_predictor(self, X: ArrayLike, predict: PredictFn) -> NDArray[np.float64]:  # pragma: no cover
        raise NotImplementedError

    def m_from_function(
        self,
        X: ArrayLike,
        *,
        predict: PredictFn,
        derivative: DerivFn | None = None,
    ) -> NDArray[np.float64]:
        """Apply the functional to a generic function.

        This is mostly used for TMLE updates where we need ``m(X, alpha_hat)``.
        """

        return self.m_from_predictor(X, predict)


class CallableFunctional(LinearFunctional):
    """Wrap a Python callable ``m(x, gamma)`` as a :class:`LinearFunctional`.

    This adapter is meant for quick experiments and README-style usage.
    The callable must implement a **linear** map in ``gamma``.

    Parameters
    ----------
    m:
        A callable of the form ``m(x_row, gamma) -> float`` where:

        - ``x_row`` is a 1D regressor vector,
        - ``gamma`` is a callable that maps a 1D regressor vector to a scalar.

        The callable is assumed to be linear in the function argument ``gamma``.
        If it is not, the GRR objective will generally be invalid.
    name:
        Name used in results.

    Notes
    -----
    To solve GRR in a finite-dimensional model ``alpha(x)=phi(x)^T beta``, the solver
    needs the matrix ``M`` such that ``m(X_i, phi(\\cdot)^T beta)=M_i^T beta``.
    For a callable functional we construct ``M`` by repeatedly applying ``m`` to each
    basis coordinate function.

    This is flexible but can be slower than implementing a dedicated
    :class:`~genriesz.functionals.LinearFunctional`.
    """

    m: Callable[[NDArray[np.float64], GammaFn], float]

    def __init__(self, m: Callable[[NDArray[np.float64], GammaFn], float], *, name: str = "custom"):
        super().__init__(name=str(name))
        object.__setattr__(self, "m", m)

    def m_basis_matrix(self, X: ArrayLike, basis: Basis) -> NDArray[np.float64]:
        X_ = _as_2d(X)
        n = X_.shape[0]

        # Determine the basis dimension without touching unfitted properties.
        Phi0 = np.asarray(basis(X_[:1]), dtype=float)
        p = int(Phi0.shape[0] if Phi0.ndim == 1 else Phi0.shape[1])

        M = np.empty((n, p), dtype=float)

        # Build M_{i,j} = m(X_i, phi_j) by applying `m` to each basis coordinate.
        for j in range(p):
            def phi_j(x_row: NDArray[np.float64], *, _j: int = j) -> float:
                x1 = _as_1d_row(x_row)
                Phi = np.asarray(basis(x1.reshape(1, -1)), dtype=float)
                return float(Phi[_j] if Phi.ndim == 1 else Phi[0, _j])

            for i in range(n):
                M[i, j] = float(self.m(_as_1d_row(X_[i]), phi_j))

        return M

    def m_from_predictor(self, X: ArrayLike, predict: PredictFn) -> NDArray[np.float64]:
        X_ = _as_2d(X)
        out = np.empty(X_.shape[0], dtype=float)

        def gamma(x_row: NDArray[np.float64]) -> float:
            x1 = _as_1d_row(x_row)
            return float(np.asarray(predict(x1.reshape(1, -1)), dtype=float).reshape(-1)[0])

        for i in range(X_.shape[0]):
            out[i] = float(self.m(_as_1d_row(X_[i]), gamma))
        return out


@dataclass(frozen=True)
class ATEFunctional(LinearFunctional):
    """Average treatment effect: E[ gamma(1,Z) - gamma(0,Z) ]."""

    treatment_index: int = 0

    def __init__(self, treatment_index: int = 0):
        super().__init__(name="ATE")
        object.__setattr__(self, "treatment_index", int(treatment_index))

    def m_basis_matrix(self, X: ArrayLike, basis: Basis) -> NDArray[np.float64]:
        X_ = _as_2d(X)
        X1 = _toggle_treatment(X_, treatment_index=self.treatment_index, value=1.0)
        X0 = _toggle_treatment(X_, treatment_index=self.treatment_index, value=0.0)
        return np.asarray(basis(X1) - basis(X0), dtype=float)

    def m_from_predictor(self, X: ArrayLike, predict: PredictFn) -> NDArray[np.float64]:
        X_ = _as_2d(X)
        X1 = _toggle_treatment(X_, treatment_index=self.treatment_index, value=1.0)
        X0 = _toggle_treatment(X_, treatment_index=self.treatment_index, value=0.0)
        return np.asarray(predict(X1) - predict(X0), dtype=float).reshape(-1)


@dataclass(frozen=True)
class ATTFunctional(LinearFunctional):
    """Average treatment effect on the treated.

    theta = E[ Y(1) - Y(0) | D=1 ]
          = E[ D * (gamma(1,Z) - gamma(0,Z)) ] / E[D].

    We treat this as a *plug-in linear functional* given a fixed value of
    ``pi = E[D]`` (estimated from the sample in the wrapper).
    """

    treatment_index: int = 0
    pi: float = 0.5

    def __init__(self, *, treatment_index: int = 0, pi: float):
        if not np.isfinite(pi) or pi <= 0.0:
            raise ValueError("pi must be positive")
        super().__init__(name="ATT")
        object.__setattr__(self, "treatment_index", int(treatment_index))
        object.__setattr__(self, "pi", float(pi))

    def m_basis_matrix(self, X: ArrayLike, basis: Basis) -> NDArray[np.float64]:
        X_ = _as_2d(X)
        D = X_[:, self.treatment_index].reshape(-1, 1)
        X1 = _toggle_treatment(X_, treatment_index=self.treatment_index, value=1.0)
        X0 = _toggle_treatment(X_, treatment_index=self.treatment_index, value=0.0)
        return (D / self.pi) * (basis(X1) - basis(X0))

    def m_from_predictor(self, X: ArrayLike, predict: PredictFn) -> NDArray[np.float64]:
        X_ = _as_2d(X)
        D = X_[:, self.treatment_index].reshape(-1)
        X1 = _toggle_treatment(X_, treatment_index=self.treatment_index, value=1.0)
        X0 = _toggle_treatment(X_, treatment_index=self.treatment_index, value=0.0)
        return (D / self.pi) * (predict(X1) - predict(X0))


@dataclass(frozen=True)
class AMEFunctional(LinearFunctional):
    """Average marginal effect (average derivative) of gamma wrt x_k."""

    coordinate: int = 0

    def __init__(self, coordinate: int = 0):
        super().__init__(name=f"AME(coord={int(coordinate)})")
        object.__setattr__(self, "coordinate", int(coordinate))

    def m_basis_matrix(self, X: ArrayLike, basis: Basis) -> NDArray[np.float64]:
        return np.asarray(basis.derivative(X, self.coordinate), dtype=float)

    def m_from_predictor(self, X: ArrayLike, predict: PredictFn) -> NDArray[np.float64]:
        raise NotImplementedError(
            "AME requires a derivative-capable predictor; use m_from_function(..., derivative=...)"
        )

    def m_from_function(
        self,
        X: ArrayLike,
        *,
        predict: PredictFn,
        derivative: DerivFn | None = None,
    ) -> NDArray[np.float64]:
        if derivative is None:
            raise ValueError("AME requires derivative()")
        return np.asarray(derivative(_as_2d(X), self.coordinate), dtype=float).reshape(-1)


@dataclass(frozen=True)
class DIDFunctional(ATTFunctional):
    """Difference-in-differences as ATT on delta outcomes.

    In the notebooks we treat DID as an ATT estimand on the panel difference

        ΔY = Y_post - Y_pre.

    The functional form is identical to ATT (with the same ``pi``), but we keep
    a separate name for clarity.
    """

    def __init__(self, *, treatment_index: int = 0, pi: float):
        super().__init__(treatment_index=treatment_index, pi=pi)
        object.__setattr__(self, "name", "DID")
