from statsmodels.compat.python import iterkeys, itervalues, zip, range
from statsmodels.stats.correlation_tools import cov_nearest
import numpy as np
import pandas as pd
from scipy import linalg as spl
from collections import defaultdict
from statsmodels.tools.sm_exceptions import (ConvergenceWarning,
                                             IterationLimitWarning)
import warnings

"""
Some details for the covariance calculations can be found in the Stata
docs:

http://www.stata.com/manuals13/xtxtgee.pdf
"""

class CovStruct(object):
    """
    A base class for correlation and covariance structures of grouped
    data.

    Each implementation of this class takes the residuals from a
    regression model that has been fitted to grouped data, and uses
    them to estimate the within-group dependence structure of the
    random errors in the model.

    The state of the covariance structure is represented through the
    value of the class variable `dep_params`.  The default state of a
    newly-created instance should correspond to the identity
    correlation matrix.
    """

    def __init__(self, cov_nearest_method="clipped"):

        # Parameters describing the dependency structure
        self.dep_params = None

        # Keep track of the number of times that the covariance was
        # adjusted.
        self.cov_adjust = []

        # Method for projecting the covariance matrix if it not SPD.
        self.cov_nearest_method = cov_nearest_method


    def initialize(self, model):
        """
        Called by GEE, used by implementations that need additional
        setup prior to running `fit`.

        Parameters
        ----------
        model : GEE class
            A reference to the parent GEE class instance.
        """
        self.model = model

    def update(self, params):
        """
        Updates the association parameter values based on the current
        regression coefficients.

        Parameters
        ----------
        params : array-like
            Working values for the regression parameters.
        """
        raise NotImplementedError

    def covariance_matrix(self, endog_expval, index):
        """
        Returns the working covariance or correlation matrix for a
        given cluster of data.

        Parameters
        ----------
        endog_expval: array-like
           The expected values of endog for the cluster for which the
           covariance or correlation matrix will be returned
        index: integer
           The index of the cluster for which the covariane or
           correlation matrix will be returned

        Returns
        -------
        M: matrix
            The covariance or correlation matrix of endog
        is_cor: bool
            True if M is a correlation matrix, False if M is a
            covariance matrix
        """
        raise NotImplementedError

    def covariance_matrix_solve(self, expval, index, stdev, rhs):
        """
        Solves matrix equations of the form `covmat * soln = rhs` and
        returns the values of `soln`, where `covmat` is the covariance
        matrix represented by this class.

        Parameters
        ----------
        expval: array-like
           The expected value of endog for each observed value in the
           group.
        index: integer
           The group index.
        stdev : array-like
            The standard deviation of endog for each observation in
            the group.
        rhs : list/tuple of array-like
            A set of right-hand sides; each defines a matrix equation
            to be solved.

        Returns
        -------
        soln : list/tuple of array-like
            The solutions to the matrix equations.

        Notes
        -----
        Returns None if the solver fails.

        Some dependence structures do not use `expval` and/or `index`
        to determine the correlation matrix.  Some families
        (e.g. binomial) do not use the `stdev` parameter when forming
        the covariance matrix.

        If the covariance matrix is singular or not SPD, it is
        projected to the nearest such matrix.  These projection events
        are recorded in the fit_history member of the GEE model.

        Systems of linear equations with the covariance matrix as the
        left hand side (LHS) are solved for different right hand sides
        (RHS); the LHS is only factorized once to save time.

        This is a default implementation, it can be reimplemented in
        subclasses to optimize the linear algebra according to the
        struture of the covariance matrix.
        """

        vmat, is_cor = self.covariance_matrix(expval, index)
        if is_cor:
            vmat *= np.outer(stdev, stdev)

        # Factor the covariance matrix.  If the factorization fails,
        # attempt to condition it into a factorizable matrix.
        threshold = 1e-2
        success = False
        cov_adjust = 0
        for itr in range(20):
            try:
                vco = spl.cho_factor(vmat)
                success = True
                break
            except np.linalg.LinAlgError:
                vmat = cov_nearest(vmat, method=self.cov_nearest_method,
                                   threshold=threshold)
                threshold *= 2
                cov_adjust += 1

        self.cov_adjust.append(cov_adjust)

        # Last resort if we still can't factor the covariance matrix.
        if success == False:
            warnings.warn("Unable to condition covariance matrix to an SPD matrix using cov_nearest",
                          ConvergenceWarning)
            vmat = np.diag(np.diag(vmat))
            vco = spl.cho_factor(vmat)

        soln = [spl.cho_solve(vco, x) for x in rhs]
        return soln

    def summary(self):
        """
        Returns a text summary of the current estimate of the
        dependence structure.
        """
        raise NotImplementedError


class Independence(CovStruct):
    """
    An independence working dependence structure.
    """

    # Nothing to update
    def update(self, params):
        return

    def covariance_matrix(self, expval, index):
        dim = len(expval)
        return np.eye(dim, dtype=np.float64), True

    def covariance_matrix_solve(self, expval, index, stdev, rhs):