from __future__ import annotations

from typing import Any, Callable, Optional, Union

import numpy as np

# import pandas as pd
import statsmodels.api as sm

from samplics.estimation.expansion import TaylorEstimator
from samplics.utils.formats import dict_to_dataframe, fpc_as_dict, numpy_array, remove_nans
from samplics.utils.types import Array, Number, Series, StringNumber


class SurveyGLM:
    """General linear models under complex survey sampling"""

    def __init__(self):
        self.beta: np.ndarray

    @staticmethod
    def _residuals(e: np.ndarray, psu: np.ndarray, nb_vars: Number) -> tuple(np.ndarray, Number):

        psus = np.unique(psu)
        if psus.shape[0] == 1 and e.shape[0] == 1:
            raise AssertionError("Only one observation in the stratum")
        if psus.shape[0] == 1:
            psu = np.arange(e.shape[0])
            psus = np.unique(psu)
        e_values = np.zeros((psus.shape[0], nb_vars))

        for i, p in enumerate(np.unique(psus)):
            e_values[i, :] += np.sum(e[psu == p, :], axis=0)
        e_means = np.sum(e_values, axis=0) / psus.shape[0]

        return np.transpose(e_values - e_means) @ (e_values - e_means), psus.shape[0]

    def _calculate_g(
        self,
        samp_weight: np.ndarray,
        resid: np.ndarray,
        x: np.ndarray,
        stratum: Optional[np.ndarray],
        psu: Optional[np.ndarray],
        fpc: Union[dict[StringNumber, Number], Number],
        glm_scale=Number,
    ) -> np.ndarray:

        e = (samp_weight * resid)[:, None] * x / glm_scale
        if psu is None:
            psu = np.arange(e.shape[0])
        if stratum is None:
            e_h, n_h = self._residuals(e=e, psu=psu, nb_vars=x.shape[1])
            return fpc * (n_h / (n_h - 1)) * e_h
        else:
            g_h = np.zeros((x.shape[1], x.shape[1]))
            for s in np.unique(stratum):
                e_s = e[stratum == s, :]
                psu_s = psu[stratum == s]
                e_h, n_h = self._residuals(e=e_s, psu=psu_s, nb_vars=x.shape[1])
                g_h += fpc[s] * (n_h / (n_h - 1)) * e_h
            return g_h

    def estimate(
        self,
        y: Array,
        x: Optional[Array] = None,
        samp_weight: Optional[Array] = None,
        stratum: Optional[Series] = None,
        psu: Optional[Series] = None,
        fpc: Union[dict[StringNumber, Number], Series, Number] = 1.0,
        remove_nan: bool = False,
    ) -> None:

        y = numpy_array(y)
        y_temp = y.copy()

        x = numpy_array(x) if x is not None else None
        psu = numpy_array(psu) if psu is not None else None

        if samp_weight is None:
            weight_temp = np.ones(y.shape[0])
        elif isinstance(samp_weight, (float, int)):
            weight_temp = samp_weight * np.ones(y_temp.shape[0])
        elif isinstance(samp_weight, np.ndarray):
            weight_temp = samp_weight.copy()
        else:
            weight_temp = np.asarray(samp_weight)

        if not isinstance(fpc, dict):
            self.fpc = fpc_as_dict(stratum, fpc)
        else:
            if list(np.unique(stratum)) != list(fpc.keys()):
                raise AssertionError("fpc dictionary keys must be the same as the strata!")
            else:
                self.fpc = fpc

        glm_model = sm.GLM(endog=y_temp, exog=x, var_weights=weight_temp)
        glm_results = glm_model.fit()

        g = self._calculate_g(
            samp_weight=samp_weight,
            resid=glm_results.resid_response,
            x=x,
            stratum=stratum,
            psu=psu,
            fpc=self.fpc,
            glm_scale=glm_results.scale,
        )

        d = glm_results.cov_params()

        self.beta = glm_results.params
        self.cov_beta = (d @ g) @ d

[((24, 15, 24, 29), 'numpy.unique', 'np.unique', ({(24, 25, 24, 28): 'psu'}, {}), '(psu)', True, 'import numpy as np\n'), ((30, 19, 30, 53), 'numpy.zeros', 'np.zeros', ({(30, 28, 30, 52): '(psus.shape[0], nb_vars)'}, {}), '((psus.shape[0], nb_vars))', True, 'import numpy as np\n'), ((75, 12, 75, 26), 'samplics.utils.formats.numpy_array', 'numpy_array', ({(75, 24, 75, 25): 'y'}, {}), '(y)', False, 'from samplics.utils.formats import dict_to_dataframe, fpc_as_dict, numpy_array, remove_nans\n'), ((98, 20, 98, 73), 'statsmodels.api.GLM', 'sm.GLM', (), '', True, 'import statsmodels.api as sm\n'), ((28, 18, 28, 39), 'numpy.arange', 'np.arange', ({(28, 28, 28, 38): 'e.shape[0]'}, {}), '(e.shape[0])', True, 'import numpy as np\n'), ((29, 19, 29, 33), 'numpy.unique', 'np.unique', ({(29, 29, 29, 32): 'psu'}, {}), '(psu)', True, 'import numpy as np\n'), ((32, 30, 32, 45), 'numpy.unique', 'np.unique', ({(32, 40, 32, 44): 'psus'}, {}), '(psus)', True, 'import numpy as np\n'), ((33, 30, 33, 60), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((34, 18, 34, 42), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((51, 18, 51, 39), 'numpy.arange', 'np.arange', ({(51, 28, 51, 38): 'e.shape[0]'}, {}), '(e.shape[0])', True, 'import numpy as np\n'), ((56, 18, 56, 52), 'numpy.zeros', 'np.zeros', ({(56, 27, 56, 51): '(x.shape[1], x.shape[1])'}, {}), '((x.shape[1], x.shape[1]))', True, 'import numpy as np\n'), ((57, 21, 57, 39), 'numpy.unique', 'np.unique', ({(57, 31, 57, 38): 'stratum'}, {}), '(stratum)', True, 'import numpy as np\n'), ((78, 12, 78, 26), 'samplics.utils.formats.numpy_array', 'numpy_array', ({(78, 24, 78, 25): 'x'}, {}), '(x)', False, 'from samplics.utils.formats import dict_to_dataframe, fpc_as_dict, numpy_array, remove_nans\n'), ((79, 14, 79, 30), 'samplics.utils.formats.numpy_array', 'numpy_array', ({(79, 26, 79, 29): 'psu'}, {}), '(psu)', False, 'from samplics.utils.formats import dict_to_dataframe, fpc_as_dict, numpy_array, remove_nans\n'), ((82, 26, 82, 45), 'numpy.ones', 'np.ones', ({(82, 34, 82, 44): 'y.shape[0]'}, {}), '(y.shape[0])', True, 'import numpy as np\n'), ((91, 23, 91, 48), 'samplics.utils.formats.fpc_as_dict', 'fpc_as_dict', ({(91, 35, 91, 42): 'stratum', (91, 44, 91, 47): 'fpc'}, {}), '(stratum, fpc)', False, 'from samplics.utils.formats import dict_to_dataframe, fpc_as_dict, numpy_array, remove_nans\n'), ((36, 15, 36, 47), 'numpy.transpose', 'np.transpose', ({(36, 28, 36, 46): '(e_values - e_means)'}, {}), '(e_values - e_means)', True, 'import numpy as np\n'), ((84, 40, 84, 64), 'numpy.ones', 'np.ones', ({(84, 48, 84, 63): 'y_temp.shape[0]'}, {}), '(y_temp.shape[0])', True, 'import numpy as np\n'), ((88, 26, 88, 49), 'numpy.asarray', 'np.asarray', ({(88, 37, 88, 48): 'samp_weight'}, {}), '(samp_weight)', True, 'import numpy as np\n'), ((93, 20, 93, 38), 'numpy.unique', 'np.unique', ({(93, 30, 93, 37): 'stratum'}, {}), '(stratum)', True, 'import numpy as np\n')]