from . import adelie_core as core
from .adelie_core.glm import (
    GlmBase64,
    GlmBase32,
    GlmMultiBase64,
    GlmMultiBase32,
)
from typing import Union
import numpy as np


def _coerce_dtype(y, dtype):
    dtype_map = {
        np.dtype("float32"): np.float32,
        np.dtype("float64"): np.float64,
    }
    valid_dtypes = list(dtype_map.keys())
    y = np.array(y, copy=False, order="C")  # important!
    if dtype is None:
        if not (y.dtype in valid_dtypes):
            raise RuntimeError(
                "y must have an underlying type of np.float32 or np.float64, "
                "or dtype must be explicitly specified."
            )
        dtype = dtype_map[y.dtype]
    else:
        if not (dtype in valid_dtypes):
            raise RuntimeError("dtype must be either np.float32 or np.float64.")
    y = y.astype(dtype, copy=False)
    return y, dtype


class glm_base:
    """Base wrapper GLM class.

    All Python wrapper classes for core GLM classes must inherit from this class.
    The purpose of this class is to expose extra member interface
    that are more easily written in Python (and not speed-critical).
    """
    def __init__(self, y, weights, core_base, dtype):
        self.core_base = core_base
        self.y = np.array(y, copy=True, dtype=dtype)
        self.dtype = dtype
        if len(y.shape) != 1:
            raise RuntimeError("y must be 1-dimensional.")
        n = y.shape[0]
        if not (weights is None):
            if weights.shape != (n,):
                raise RuntimeError("y and weights must have same length.")
            weights_sum = np.sum(weights)
            if not np.allclose(weights_sum, 1):
                weights = weights / weights_sum
        else:
            weights = np.full(n, 1/n, dtype=dtype)
        self.weights = np.array(weights, copy=True, dtype=dtype)


class multiglm_base:
    """Base wrapper multi-response GLM class.

    All Python wrapper classes for core multi-response GLM classes must inherit from this class.
    The purpose of this class is to expose extra member interface
    that are more easily written in Python (and not speed-critical).
    """
    def __init__(self, y, weights, core_base, dtype):
        self.core_base = core_base
        self.y = np.array(y, copy=True, dtype=dtype)
        self.dtype = dtype
        if len(y.shape) != 2:
            raise RuntimeError("y must be 2-dimensional.")
        n = y.shape[0]
        if not (weights is None):
            if weights.shape != (n,):
                raise RuntimeError("y rows and weights must have same length.")
            weights_sum = np.sum(weights)
            if not np.allclose(weights_sum, 1):
                weights = weights / np.sum(weights)
        else:
            weights = np.full(n, 1/n, dtype=dtype)
        self.weights = np.array(weights, copy=True, dtype=dtype)


def binomial(
    y: np.ndarray,
    *,
    weights: np.ndarray =None,
    link: str ="logit",
    dtype: Union[np.float32, np.float64] =None,
):
    """Creates a Binomial GLM family object.

    The Binomial GLM family with the logit link function 
    specifies the loss function as:

    .. math::
        \\begin{align*}
            \\ell(\\eta)
            =
            \\sum\\limits_{i=1}^n w_i \\left(
                -y_i \\eta_i + \\log(1 + e^{\\eta_i})
            \\right)
        \\end{align*}

    The link function is given by

    .. math::
        \\begin{align*}
            g(\\mu)_i
            &=
            \\log\\left(\\frac{\\mu_i}{1 - \\mu_i}\\right)
        \\end{align*}

    The Binomial GLM family with the probit link function 
    specifies the loss function as:

    .. math::
        \\begin{align*}
            \\ell(\\eta)
            =
            -\\sum\\limits_{i=1}^n w_i \\left(
                y_i \\log(\\Phi(\\eta_i)) + (1-y_i) \\log(1-\\Phi(\\eta_i))
            \\right)
        \\end{align*}

    where :math:`\\Phi` is the standard normal CDF.
    The link function is given by

    .. math::
        \\begin{align*}
            g(\\mu)_i
            &=
            \\Phi^{-1}(\\mu_i)
        \\end{align*}

    We assume that :math:`y_i \\in [0,1]`.

    Parameters
    ----------
    y : (n,) ndarray 
        Response vector :math:`y`.
    weights : (n,) ndarray, optional
        Observation weights :math:`W`.
        Weights are normalized such that they sum to ``1``.
        Default is ``None``, in which case, it is set to ``np.full(n, 1/n)``.
    link : str, optional
        The link function type.
        It must be one of the following:

            - ``"logit"``: the logit link function.
            - ``"probit"``: the probit link function.

        Default is ``"logit"``.
    dtype : Union[float32, float64], optional
        The underlying data type.
        If ``None``, it is inferred from ``y``,
        in which case ``y`` must have an underlying data type of
        :class:`numpy.float32` or :class:`numpy.float64`.
        Default is ``None``.

    Returns
    -------
    glm
        Binomial GLM object.

    See Also
    --------
    adelie.adelie_core.glm.GlmBinomialLogit32
    adelie.adelie_core.glm.GlmBinomialLogit64
    adelie.adelie_core.glm.GlmBinomialProbit32
    adelie.adelie_core.glm.GlmBinomialProbit64
    """
    dispatcher = {
        "logit": {
            np.float64: core.glm.GlmBinomialLogit64,
            np.float32: core.glm.GlmBinomialLogit32,
        },
        "probit": {
            np.float64: core.glm.GlmBinomialProbit64,
            np.float32: core.glm.GlmBinomialProbit32,
        },
    }

    y, dtype = _coerce_dtype(y, dtype)

    core_base = dispatcher[link][dtype]

    class _binomial(glm_base, core_base):
        def __init__(self):
            glm_base.__init__(self, y, weights, core_base, dtype)
            core_base.__init__(self, self.y, self.weights)

        def reweight(self, weights=None):
            weights = self.weights if weights is None else weights
            return binomial(y=y, weights=weights, dtype=dtype)

    return _binomial()


def cox(
    start: np.ndarray,
    stop: np.ndarray,
    status: np.ndarray,
    *,
    strata: np.ndarray =None,
    weights: np.ndarray =None,
    tie_method: str ="efron",
    dtype: Union[np.float32, np.float64] =None,
):
    """Creates a Cox GLM family object.

    The Cox GLM family specifies the loss function as:

    .. math::
        \\begin{align*}
            \\ell(\\eta)
            &=
            -\\sum\\limits_{i=1}^n w_i \\delta_i \\eta_i
            +\\sum\\limits_{i=1}^n \\overline{w}_i \\delta_i A_i(\\eta)
            \\\\
            A_i(\\eta)
            &=
            \\log\\left(
                \\sum\\limits_{k \\in R(t_i)} w_k e^{\\eta_k}
                -
                \\sigma_i \\sum\\limits_{k \\in H(t_i)} w_k e^{\\eta_k}
            \\right)
        \\end{align*}

    where
    
    .. math::
        \\begin{align*}
            R(u)
            &=
            \\{i : u \\in (s_i, t_i]\\}
            \\\\
            H(u)
            &=
            \\{i : t_i = u, \\delta_i = 1\\}
            \\\\
            \\overline{w}_i 
            &= 
            \\frac{\\sum_{k \\in H(t_i)} w_k}{\\sum_{k \\in H(t_i)} 1_{w_k > 0}} 
            1_{\\delta_i = 1, w_i > 0}
        \\end{align*}

    Here,
    :math:`\\delta` is the status (``1`` for event, ``0`` for censored) vector,
    :math:`s` is the vector of start times,
    :math:`t` is the vector of stop times,
    :math:`R(u)` is the at-risk set at time :math:`u`,
    :math:`H(u)` is the set of ties at event time :math:`u`,
    :math:`\\overline{w}` is the vector of average weights within ties with positive weights,
    :math:`\\sigma` is the correction scale for tie-breaks, 
    which is determined by the type of correction method (Breslow or Efron).
    Note that :math:`\\overline{w}_i` and :math:`A_i(\\eta)` are only well-defined 
    whenever :math:`\\delta_i=1`, which is not an issue in the computation of :math:`\\ell(\\eta)`.

    The link function is given by

    .. math::
        \\begin{align*}
            g(\\mu)_i = \\log(\\mu_i)
        \\end{align*}

    If ``strata`` is specified, then the loss function 
    is simply the sum of the losses for each strata.
    Namely, given the strata vector :math:`S`,

    .. math::
        \\begin{align*}
            \\ell(\\eta)
            &=
            \\sum_{m=0}^{M-1}
            \\ell_m(\\eta_{\\mathcal{I}_m})
        \\end{align*}

    where
    :math:`M` is the number of strata,
    :math:`\\mathcal{I}_m = \\{i : S_i = m\\}` is the set of individuals in stratum :math:`m`,
    :math:`\\eta_I` is the subset of :math:`\\eta` given by the indices in :math:`I`,
    and :math:`\\ell_m` is the usual Cox loss function as above
    using only the input data in :math:`\\mathcal{I}_m`.

    .. note::
        The strata indicator :math:`S_i` *must* take on values in the set :math:`\\{0, \\ldots, M-1\\}`.

    Parameters
    ----------
    start : (n,) ndarray
        Start time vector :math:`s`.
    stop : (n,) ndarray
        Stop time vector :math:`t`.
    status : (n,) ndarray 
        Status vector :math:`\\delta`.
    strata : (n,) ndarray, optional
        Strata vector :math:`S`.
        If ``None``, there is only one stratum.
        Default is ``None``.
    weights : (n,) ndarray, optional
        Observation weights :math:`W`.
        Weights are normalized such that they sum to ``1``.
        Default is ``None``, in which case, it is set to ``np.full(n, 1/n)``.
    tie_method : str, optional
        The tie-breaking method that determines the scales :math:`\\sigma`.
        It must be one of the following:

            - ``"efron"``
            - ``"breslow"``

        Default is ``"efron"``.
    dtype : Union[float32, float64], optional
        The underlying data type.
        If ``None``, it is inferred from ``status``,
        in which case ``status`` must have an underlying data type of
        :class:`numpy.float32` or :class:`numpy.float64`.
        Default is ``None``.

    Returns
    -------
    glm
        Cox GLM object.

    See Also
    --------
    adelie.adelie_core.glm.GlmCox32
    adelie.adelie_core.glm.GlmCox64
    """
    dispatcher = {
        np.float64: core.glm.GlmCox64,
        np.float32: core.glm.GlmCox32,
    }

    status, dtype = _coerce_dtype(status, dtype)

    core_base = dispatcher[dtype]

    if strata is None:
        strata = np.zeros(status.size, dtype=int)

    class _cox(glm_base, core_base):
        def __init__(self):
            self.start = np.array(start, copy=True, dtype=dtype)
            self.stop = np.array(stop, copy=True, dtype=dtype)
            glm_base.__init__(self, status, weights, core_base, dtype)
            self.status = self.y
            self.strata = np.array(strata, copy=True, dtype=int)
            self.tie_method = tie_method
            core_base.__init__(
                self, 
                self.start, 
                self.stop, 
                self.status, 
                self.strata,
                self.weights, 
                self.tie_method,
            )

        def reweight(self, weights=None):
            weights = self.weights if weights is None else weights
            return cox(
                start=start, 
                stop=stop,
                status=status,
                strata=strata,
                weights=weights, 
                tie_method=tie_method,
                dtype=dtype,
            )

    return _cox()


def gaussian(
    y: np.ndarray,
    *,
    weights: np.ndarray =None,
    dtype: Union[np.float32, np.float64] =None,
    opt: bool =True,
):
    """Creates a Gaussian GLM family object.

    The Gaussian GLM family specifies the loss function as:

    .. math::
        \\begin{align*}
            \\ell(\\eta)
            =
            \\sum\\limits_{i=1}^n w_i \\left(
                -y_i \\eta_i + \\frac{\\eta_i^2}{2}
            \\right) 
        \\end{align*}

    The link function is given by

    .. math::
        \\begin{align*}
            g(\\mu)_i
            &=
            \\mu_i
        \\end{align*}

    Parameters
    ----------
    y : (n,) ndarray 
        Response vector :math:`y`.
    weights : (n,) ndarray, optional
        Observation weights :math:`W`.
        Weights are normalized such that they sum to ``1``.
        Default is ``None``, in which case, it is set to ``np.full(n, 1/n)``.
    dtype : Union[float32, float64], optional
        The underlying data type.
        If ``None``, it is inferred from ``y``,
        in which case ``y`` must have an underlying data type of
        :class:`numpy.float32` or :class:`numpy.float64`.
        Default is ``None``.
    opt : bool, optional
        If ``True``, an optimized routine is used when passed into :func:`adelie.solver.grpnet`.
        Otherwise, a general routine with IRLS is used.
        This flag is mainly for developers for testing purposes.
        We advise users to use the default value.
        Default is ``True``.

    Returns
    -------
    glm
        Gaussian GLM object.

    See Also
    --------
    adelie.adelie_core.glm.GlmGaussian32
    adelie.adelie_core.glm.GlmGaussian64
    """
    dispatcher = {
        np.float64: core.glm.GlmGaussian64,
        np.float32: core.glm.GlmGaussian32,
    }

    y, dtype = _coerce_dtype(y, dtype)

    core_base = dispatcher[dtype]

    class _gaussian(glm_base, core_base):
        def __init__(self):
            self.opt = opt
            glm_base.__init__(self, y, weights, core_base, dtype)
            core_base.__init__(self, self.y, self.weights)

        def reweight(self, weights=None):
            weights = self.weights if weights is None else weights
            return gaussian(y=y, weights=weights, dtype=dtype, opt=opt)

    return _gaussian()


def multigaussian(
    y: np.ndarray,
    *,
    weights: np.ndarray =None,
    dtype: Union[np.float32, np.float64] =None,
    opt: bool =True,
):
    """Creates a MultiGaussian GLM family object.

    The MultiGaussian GLM family specifies the loss function as:

    .. math::
        \\begin{align*}
            \\ell(\\eta)
            =
            \\frac{1}{K}
            \\sum\\limits_{i=1}^n 
            w_{i} \\left(
                -\\sum\\limits_{k=1}^K y_{ik} \\eta_{ik} 
                +\\frac{\\|\\eta_{i\\cdot}\\|^2}{2}
            \\right)
        \\end{align*}

    The link function is given by

    .. math::
        \\begin{align*}
            g(\\mu)_{ik}
            &=
            \\mu_{ik}
        \\end{align*}

    Parameters
    ----------
    y : (n, K) ndarray 
        Response matrix :math:`y`.
    weights : (n,) ndarray, optional
        Observation weights :math:`W`.
        Weights are normalized such that they sum to ``1``.
        Default is ``None``, in which case, it is set to ``np.full(n, 1/n)``.
    dtype : Union[float32, float64], optional
        The underlying data type.
        If ``None``, it is inferred from ``y``,
        in which case ``y`` must have an underlying data type of
        :class:`numpy.float32` or :class:`numpy.float64`.
        Default is ``None``.
    opt : bool, optional
        If ``True``, an optimized routine is used when passed into :func:`adelie.solver.grpnet`.
        Otherwise, a general routine with IRLS is used.
        This flag is mainly for developers for testing purposes.
        We advise users to use the default value.
        Default is ``True``.

    Returns
    -------
    glm
        MultiGaussian GLM object.

    See Also
    --------
    adelie.adelie_core.glm.GlmMultiGaussian32
    adelie.adelie_core.glm.GlmMultiGaussian64
    """
    dispatcher = {
        np.float64: core.glm.GlmMultiGaussian64,
        np.float32: core.glm.GlmMultiGaussian32,
    }

    y, dtype = _coerce_dtype(y, dtype)

    core_base = dispatcher[dtype]

    class _multigaussian(multiglm_base, core_base):
        def __init__(self):
            self.opt = opt
            multiglm_base.__init__(self, y, weights, core_base, dtype)
            core_base.__init__(self, self.y, self.weights)

        def reweight(self, weights=None):
            weights = self.weights if weights is None else weights
            return multigaussian(y=y, weights=weights, dtype=dtype, opt=opt)

    return _multigaussian()


def multinomial(
    y: np.ndarray,
    *,
    weights: np.ndarray =None,
    dtype: Union[np.float32, np.float64] =None,
):
    """Creates a Multinomial GLM family object.

    The Multinomial GLM family specifies the loss function as:

    .. math::
        \\begin{align*}
            \\ell(\\eta)
            =
            \\frac{1}{K}
            \\sum\\limits_{i=1}^n 
            w_i
            \\left(
            -\\sum\\limits_{k=1}^{K} y_{ik} \\eta_{ik} 
            + \\log\\left(
                \\sum\\limits_{\\ell=1}^{K} e^{\\eta_{i\\ell}}
            \\right)
            \\right)
        \\end{align*}

    The link function is given by

    .. math::
        \\begin{align*}
            g(\\mu)_{ik}
            &=
            \\log(\\mu_{ik}) + C_i
        \\end{align*}

    for any arbitrary constants :math:`C_i`.

    We assume that every :math:`y_{ik} \\in [0,1]` and
    for each fixed :math:`i`, :math:`\\sum_{k=1}^K y_{ik} = 1`.

    .. note::
        The :func:`~adelie.adelie_core.glm.GlmMultiBase64.hessian` method computes 
        :math:`2 \\mathrm{diag}(\\nabla^2 \\ell(\\eta))`
        as the diagonal majorization.

    Parameters
    ----------
    y : (n, K) ndarray 
        Response matrix :math:`y`.
    weights : (n,) ndarray, optional
        Observation weights :math:`W`.
        Weights are normalized such that they sum to ``1``.
        Default is ``None``, in which case, it is set to ``np.full(n, 1/n)``.
    dtype : Union[float32, float64], optional
        The underlying data type.
        If ``None``, it is inferred from ``y``,
        in which case ``y`` must have an underlying data type of
        :class:`numpy.float32` or :class:`numpy.float64`.
        Default is ``None``.

    Returns
    -------
    glm
        Multinomial GLM object.

    See Also
    --------
    adelie.adelie_core.glm.GlmMultinomial32
    adelie.adelie_core.glm.GlmMultinomial64
    """
    dispatcher = {
        np.float64: core.glm.GlmMultinomial64,
        np.float32: core.glm.GlmMultinomial32,
    }

    y, dtype = _coerce_dtype(y, dtype)

    core_base = dispatcher[dtype]

    class _multinomial(multiglm_base, core_base):
        def __init__(self):
            multiglm_base.__init__(self, y, weights, core_base, dtype)
            core_base.__init__(self, self.y, self.weights)

        def reweight(self, weights=None):
            weights = self.weights if weights is None else weights
            return multinomial(y=y, weights=weights, dtype=dtype)

    return _multinomial()


def poisson(
    y: np.ndarray,
    *,
    weights: np.ndarray =None,
    dtype: Union[np.float32, np.float64] =None,
):
    """Creates a Poisson GLM family object.

    The Poisson GLM family specifies the loss function as:

    .. math::
        \\begin{align*}
            \\ell(\\eta)
            =
            \\sum\\limits_{i=1}^n w_i \\left(
                -y_i \\eta_i + e^{\\eta_i}
            \\right) 
        \\end{align*}

    The link function is given by

    .. math::
        \\begin{align*}
            g(\\mu)_i
            &=
            \\log(\\mu_i)
        \\end{align*}

    We assume that :math:`y_i \\in \\mathbb{N}_0`.

    Parameters
    ----------
    y : (n,) ndarray 
        Response vector :math:`y`.
    weights : (n,) ndarray, optional
        Observation weights :math:`W`.
        Weights are normalized such that they sum to ``1``.
        Default is ``None``, in which case, it is set to ``np.full(n, 1/n)``.
    dtype : Union[float32, float64], optional
        The underlying data type.
        If ``None``, it is inferred from ``y``,
        in which case ``y`` must have an underlying data type of
        :class:`numpy.float32` or :class:`numpy.float64`.
        Default is ``None``.

    Returns
    -------
    glm
        Poisson GLM object.

    See Also
    --------
    adelie.adelie_core.glm.GlmPoisson32
    adelie.adelie_core.glm.GlmPoisson64
    """
    dispatcher = {
        np.float64: core.glm.GlmPoisson64,
        np.float32: core.glm.GlmPoisson32,
    }

    y, dtype = _coerce_dtype(y, dtype)

    core_base = dispatcher[dtype]

    class _poisson(glm_base, core_base):
        def __init__(self):
            glm_base.__init__(self, y, weights, core_base, dtype)
            core_base.__init__(self, self.y, self.weights)

        def reweight(self, weights=None):
            weights = self.weights if weights is None else weights
            return poisson(y=y, weights=weights, dtype=dtype)

    return _poisson()
