"""
(c) 2014 Brendan Bulik-Sullivan and Hilary Finucane

Estimators of heritability and genetic correlation.

Shape convention is (n_snp, n_annot) for all classes.
Last column = intercept.

"""

from collections import namedtuple

import numpy as np
import pandas as pd
from scipy.stats import chi2, norm
from scipy.stats import t as tdist

import ldsc.jackknife as jk
from ldsc.irwls import IRWLS

np.seterr(divide="raise", invalid="raise")

s = lambda x: remove_brackets(str(np.array([[x]])))


def update_separators(s, ii):
    """s are separators with ii masked. Returns unmasked separators."""
    maplist = np.arange(len(ii))[np.squeeze(ii)]
    mask_to_unmask = lambda i: maplist[i]
    t = np.apply_along_axis(mask_to_unmask, 0, s[1:-1])
    t = np.hstack(((0), t, (len(ii))))
    return t


def p_z_norm(est, se):
    """Convert estimate and se to Z-score and P-value."""
    try:
        Z = est / se
    except (FloatingPointError, ZeroDivisionError):
        Z = float("inf")

    P = chi2.sf(Z**2, 1, loc=0, scale=1)  # 0 if Z=inf
    return P, Z


def remove_brackets(x):
    """Get rid of brackets and trailing whitespace in numpy arrays."""
    return x.replace("[", "").replace("]", "").strip()


def append_intercept(x):
    """
    Appends an intercept term to the design matrix for a linear regression.

    Parameters
    ----------
    x : np.array with shape (n_row, n_col)
        Design matrix. Columns are predictors; rows are observations.

    Returns
    -------
    x_new : np.array with shape (n_row, n_col+1)
        Design matrix with intercept term appended.

    """
    n_row = x.shape[0]
    intercept = np.ones((n_row, 1))
    x_new = np.concatenate((x, intercept), axis=1)
    return x_new


def remove_intercept(x):
    """Removes the last column."""
    n_col = x.shape[1]
    return x[:, 0 : n_col - 1]


def gencov_obs_to_liab(gencov_obs, P1, P2, K1, K2):
    """
    Converts genetic covariance on the observed scale in an ascertained sample to genetic
    covariance on the liability scale in the population

    Parameters
    ----------
    gencov_obs : float
        Genetic covariance on the observed scale in an ascertained sample.
    P1, P2 : float in (0,1)
        Prevalences of phenotypes 1,2 in the sample.
    K1, K2 : float in (0,1)
        Prevalences of phenotypes 1,2 in the population.

    Returns
    -------
    gencov_liab : float
        Genetic covariance between liabilities in the population.

    Note: if a trait is a QT, set P = K = None.

    """
    c1 = 1
    c2 = 1
    if P1 is not None and K1 is not None:
        c1 = np.sqrt(h2_obs_to_liab(1, P1, K1))
    if P2 is not None and K2 is not None:
        c2 = np.sqrt(h2_obs_to_liab(1, P2, K2))

    return gencov_obs * c1 * c2


def h2_obs_to_liab(h2_obs, P, K):
    """
    Converts heritability on the observed scale in an ascertained sample to heritability
    on the liability scale in the population.

    Parameters
    ----------
    h2_obs : float
        Heritability on the observed scale in an ascertained sample.
    P : float in (0,1)
        Prevalence of the phenotype in the sample.
    K : float in (0,1)
        Prevalence of the phenotype in the population.

    Returns
    -------
    h2_liab : float
        Heritability of liability in the population.

    """
    if np.isnan(P) and np.isnan(K):
        return h2_obs
    if K <= 0 or K >= 1:
        raise ValueError("K must be in the range (0,1)")
    if P <= 0 or P >= 1:
        raise ValueError("P must be in the range (0,1)")

    thresh = norm.isf(K)
    conversion_factor = K**2 * (1 - K) ** 2 / (P * (1 - P) * norm.pdf(thresh) ** 2)
    return h2_obs * conversion_factor


class LD_Score_Regression:
    def __init__(
        self,
        y,
        x,
        w,
        N,
        M,
        n_blocks,
        intercept=None,
        slow=False,
        step1_ii=None,
        old_weights=False,
    ):
        """
        Parameters
        ----------
        y : np.array with shape (n_snp, 1)
            Chi-squared statistics (or z1z2?).
        x : np.array with shape (n_snp, 1)
            "ref_ld", - LD scores to use as independent variables.
        w : np.array with shape (n_snp, 1)
            "w_ld_cname" - LD scores to use as weights.
        N : np.array with shape (n_snp, 1)
            Sample size
        M : np.array with shape (1, n_annot)
            Number of SNPs in each annotation category.
        n_blocks : int
            Number of jackknife blocks.
        """
        for i in [y, x, w, M, N]:
            try:
                if len(i.shape) != 2:
                    raise TypeError("Arguments must be 2D arrays.")
            except AttributeError as e:
                raise TypeError("Arguments must be arrays.") from e

        n_snp, self.n_annot = x.shape
        if any(i.shape != (n_snp, 1) for i in [y, w, N]):
            raise ValueError(
                "N, weights and response (z1z2 or chisq) must have shape (n_snp, 1)."
            )
        if M.shape != (1, self.n_annot):
            raise ValueError("M must have shape (1, n_annot).")

        M_tot = float(np.sum(M))
        x_tot = np.sum(x, axis=1).reshape((n_snp, 1))
        self.constrain_intercept = intercept is not None
        self.intercept = intercept
        self.n_blocks = n_blocks
        tot_agg = self.aggregate(y, x_tot, N, M_tot, intercept)
        initial_w = self._update_weights(x_tot, w, N, M_tot, tot_agg, intercept)
        Nbar = np.mean(N)  # keep condition number low
        x = np.multiply(N, x) / Nbar
        if not self.constrain_intercept:
            x, x_tot = append_intercept(x), append_intercept(x_tot)
            yp = y
        else:
            yp = y - intercept
            self.intercept_se = "NA"
        del y
        self.twostep_filtered = None
        if step1_ii is not None and self.constrain_intercept:
            raise ValueError("twostep is not compatible with constrain_intercept.")
        elif step1_ii is not None and self.n_annot > 1:
            raise ValueError("twostep not compatible with partitioned LD Score yet.")
        elif step1_ii is not None:
            n1 = np.sum(step1_ii)
            self.twostep_filtered = n_snp - n1
            x1 = x[np.squeeze(step1_ii), :]
            yp1, w1, N1, initial_w1 = map(
                lambda a: a[step1_ii].reshape((n1, 1)), (yp, w, N, initial_w)
            )
            update_func1 = lambda a: self._update_func(
                a, x1, w1, N1, M_tot, Nbar, ii=step1_ii
            )
            step1_jknife = IRWLS(
                x1, yp1, update_func1, n_blocks, slow=slow, w=initial_w1
            )
            step1_int, _ = self._intercept(step1_jknife)
            yp = yp - step1_int
            x = remove_intercept(x)
            x_tot = remove_intercept(x_tot)
            update_func2 = lambda a: self._update_func(
                a, x_tot, w, N, M_tot, Nbar, step1_int
            )
            s = update_separators(step1_jknife.separators, step1_ii)
            step2_jknife = IRWLS(
                x, yp, update_func2, n_blocks, slow=slow, w=initial_w, separators=s
            )
            c = np.sum(np.multiply(initial_w, x)) / np.sum(
                np.multiply(initial_w, np.square(x))
            )
            jknife = self._combine_twostep_jknives(
                step1_jknife, step2_jknife, M_tot, c, Nbar
            )
        elif old_weights:
            initial_w = np.sqrt(initial_w)
            x = IRWLS._weight(x, initial_w)
            y = IRWLS._weight(yp, initial_w)
            jknife = jk.LstsqJackknifeFast(x, y, n_blocks)
        else:
            update_func = lambda a: self._update_func(
                a, x_tot, w, N, M_tot, Nbar, intercept
            )
            jknife = IRWLS(x, yp, update_func, n_blocks, slow=slow, w=initial_w)

        self.coef, self.coef_cov, self.coef_se = self._coef(jknife, Nbar)
        self.cat, self.cat_cov, self.cat_se = self._cat(
            jknife, M, Nbar, self.coef, self.coef_cov
        )

        self.tot, self.tot_cov, self.tot_se = self._tot(self.cat, self.cat_cov)
        self.prop, self.prop_cov, self.prop_se = self._prop(
            jknife, M, Nbar, self.cat, self.tot
        )

        self.enrichment, self.M_prop = self._enrichment(M, M_tot, self.cat, self.tot)
        if not self.constrain_intercept:
            self.intercept, self.intercept_se = self._intercept(jknife)

        self.jknife = jknife
        self.tot_delete_values = self._delete_vals_tot(jknife, Nbar, M)
        self.part_delete_values = self._delete_vals_part(jknife, Nbar, M)
        if not self.constrain_intercept:
            self.intercept_delete_values = jknife.delete_values[:, self.n_annot]

        self.M = M

    @classmethod
    def aggregate(cls, y, x, N, M, intercept=None):
        if intercept is None:
            intercept = cls.__null_intercept__

        num = M * (np.mean(y) - intercept)
        denom = np.mean(np.multiply(x, N))
        return num / denom

    def _update_func(self, x, ref_ld_tot, w_ld, N, M, Nbar, intercept=None, ii=None):
        raise NotImplementedError

    def _delete_vals_tot(self, jknife, Nbar, M):
        """Get delete values for total h2 or gencov."""
        n_annot = self.n_annot
        tot_delete_vals = jknife.delete_values[
            :, 0:n_annot
        ]  # shape (n_blocks, n_annot)
        # shape (n_blocks, 1)
        tot_delete_vals = np.dot(tot_delete_vals, M.T) / Nbar
        return tot_delete_vals

    def _delete_vals_part(self, jknife, Nbar, M):
        """Get delete values for partitioned h2 or gencov."""
        n_annot = self.n_annot
        return jknife.delete_values[:, 0:n_annot] / Nbar

    def _coef(self, jknife, Nbar):
        """Get coefficient estimates + cov from the jackknife."""
        n_annot = self.n_annot
        coef = jknife.est[0, 0:n_annot] / Nbar
        coef_cov = jknife.jknife_cov[0:n_annot, 0:n_annot] / Nbar**2
        coef_se = np.sqrt(np.diag(coef_cov))
        return coef, coef_cov, coef_se

    def _cat(self, jknife, M, Nbar, coef, coef_cov):
        """Convert coefficients to per-category h2 or gencov."""
        cat = np.multiply(M, coef)
        cat_cov = np.multiply(np.dot(M.T, M), coef_cov)
        cat_se = np.sqrt(np.diag(cat_cov))
        return cat, cat_cov, cat_se

    def _tot(self, cat, cat_cov):
        """Convert per-category h2 to total h2 or gencov."""
        tot = np.sum(cat)
        tot_cov = np.sum(cat_cov)
        tot_se = np.sqrt(tot_cov)
        return tot, tot_cov, tot_se

    def _prop(self, jknife, M, Nbar, cat, tot):
        """Convert total h2 and per-category h2 to per-category proportion h2 or gencov."""
        n_annot = self.n_annot
        n_blocks = jknife.delete_values.shape[0]
        numer_delete_vals = (
            np.multiply(M, jknife.delete_values[:, 0:n_annot]) / Nbar
        )  # (n_blocks, n_annot)
        denom_delete_vals = np.sum(numer_delete_vals, axis=1).reshape((n_blocks, 1))
        denom_delete_vals = np.dot(denom_delete_vals, np.ones((1, n_annot)))
        prop = jk.RatioJackknife(cat / tot, numer_delete_vals, denom_delete_vals)
        return prop.est, prop.jknife_cov, prop.jknife_se

    def _enrichment(self, M, M_tot, cat, tot):
        """Compute proportion of SNPs per-category enrichment for h2 or gencov."""
        M_prop = M / M_tot
        enrichment = np.divide(cat, M) / (tot / M_tot)
        return enrichment, M_prop

    def _intercept(self, jknife):
        """Extract intercept and intercept SE from block jackknife."""
        n_annot = self.n_annot
        intercept = jknife.est[0, n_annot]
        intercept_se = jknife.jknife_se[0, n_annot]
        return intercept, intercept_se

    def _combine_twostep_jknives(self, step1_jknife, step2_jknife, M_tot, c, Nbar=1):
        """Combine free intercept and constrained intercept jackknives for --two-step."""
        n_blocks, n_annot = step1_jknife.delete_values.shape
        n_annot -= 1
        if n_annot > 2:
            raise ValueError("twostep not yet implemented for partitioned LD Score.")

        step1_int, _ = self._intercept(step1_jknife)
        est = np.hstack((step2_jknife.est, np.array(step1_int).reshape((1, 1))))
        delete_values = np.zeros((n_blocks, n_annot + 1))
        delete_values[:, n_annot] = step1_jknife.delete_values[:, n_annot]
        delete_values[:, 0:n_annot] = step2_jknife.delete_values - c * (
            step1_jknife.delete_values[:, n_annot] - step1_int
        ).reshape((n_blocks, n_annot))  # check this
        pseudovalues = jk.Jackknife.delete_values_to_pseudovalues(delete_values, est)
        jknife_est, jknife_var, jknife_se, jknife_cov = jk.Jackknife.jknife(
            pseudovalues
        )
        jknife = namedtuple(
            "jknife",
            [
                "est",
                "jknife_se",
                "jknife_est",
                "jknife_var",
                "jknife_cov",
                "delete_values",
            ],
        )
        return jknife(est, jknife_se, jknife_est, jknife_var, jknife_cov, delete_values)


class Hsq(LD_Score_Regression):
    __null_intercept__ = 1

    def __init__(
        self,
        y,
        x,
        w,
        N,
        M,
        n_blocks=200,
        intercept=None,
        slow=False,
        twostep=None,
        old_weights=False,
    ):
        step1_ii = None
        if twostep is not None:
            step1_ii = y < twostep

        LD_Score_Regression.__init__(
            self,
            y,
            x,
            w,
            N,
            M,
            n_blocks,
            intercept=intercept,
            slow=slow,
            step1_ii=step1_ii,
            old_weights=old_weights,
        )
        self.mean_chisq, self.lambda_gc = self._summarize_chisq(y)
        if not self.constrain_intercept:
            self.ratio, self.ratio_se = self._ratio(
                self.intercept, self.intercept_se, self.mean_chisq
            )

    def _update_func(self, x, ref_ld_tot, w_ld, N, M, Nbar, intercept=None, ii=None):
        """
        Update function for IRWLS

        x is the output of np.linalg.lstsq.
        x[0] is the regression coefficients
        x[0].shape is (# of dimensions, 1)
        the last element of x[0] is the intercept.

        intercept is None --> free intercept
        intercept is not None --> constrained intercept
        """
        hsq = M * x[0][0] / Nbar
        if intercept is None:
            intercept = max(x[0][1])  # divide by zero error if intercept < 0
        else:
            if ref_ld_tot.shape[1] > 1:
                raise ValueError(
                    "Design matrix has intercept column for constrained intercept regression!"
                )

        ld = ref_ld_tot[:, 0].reshape(w_ld.shape)  # remove intercept
        w = self.weights(ld, w_ld, N, M, hsq, intercept, ii)
        return w

    def _summarize_chisq(self, chisq):
        """Compute mean chi^2 and lambda_GC."""
        mean_chisq = np.mean(chisq)
        # median and matrix don't play nice
        lambda_gc = np.median(np.asarray(chisq)) / 0.4549
        return mean_chisq, lambda_gc

    def _ratio(self, intercept, intercept_se, mean_chisq):
        """Compute ratio (intercept - 1) / (mean chi^2 -1 )."""
        if mean_chisq > 1:
            ratio_se = intercept_se / (mean_chisq - 1)
            ratio = (intercept - 1) / (mean_chisq - 1)
        else:
            ratio = "NA"
            ratio_se = "NA"

        return ratio, ratio_se

    def _overlap_output(
        self, category_names, overlap_matrix, M_annot, M_tot, print_coefficients
    ):
        """LD Score regression summary for overlapping categories."""
        overlap_matrix_prop = np.zeros([self.n_annot, self.n_annot])
        for i in range(self.n_annot):
            overlap_matrix_prop[i, :] = overlap_matrix[i, :] / M_annot

        prop_hsq_overlap = np.dot(overlap_matrix_prop, self.prop.T).reshape(
            (1, self.n_annot)
        )
        prop_hsq_overlap_var = np.diag(
            np.dot(np.dot(overlap_matrix_prop, self.prop_cov), overlap_matrix_prop.T)
        )
        prop_hsq_overlap_se = np.sqrt(np.maximum(0, prop_hsq_overlap_var)).reshape(
            (1, self.n_annot)
        )
        one_d_convert = lambda x: np.array(x).reshape(np.prod(x.shape))
        prop_M_overlap = M_annot / M_tot
        enrichment = prop_hsq_overlap / prop_M_overlap
        enrichment_se = prop_hsq_overlap_se / prop_M_overlap
        overlap_matrix_diff = np.zeros([self.n_annot, self.n_annot])
        for i in range(self.n_annot):
            if not M_tot == M_annot[0, i]:
                overlap_matrix_diff[i, :] = overlap_matrix[i, :] / M_annot[0, i] - (
                    M_annot - overlap_matrix[i, :]
                ) / (M_tot - M_annot[0, i])

        diff_est = np.dot(overlap_matrix_diff, self.coef)
        diff_cov = np.dot(
            np.dot(overlap_matrix_diff, self.coef_cov), overlap_matrix_diff.T
        )
        diff_se = np.sqrt(np.diag(diff_cov))
        diff_p = [
            "NA"
            if diff_se[i] == 0
            else 2 * tdist.sf(abs(diff_est[i] / diff_se[i]), self.n_blocks)
            for i in range(self.n_annot)
        ]

        df = pd.DataFrame(
            {
                "Category": category_names,
                "Prop._SNPs": one_d_convert(prop_M_overlap),
                "Prop._h2": one_d_convert(prop_hsq_overlap),
                "Prop._h2_std_error": one_d_convert(prop_hsq_overlap_se),
                "Enrichment": one_d_convert(enrichment),
                "Enrichment_std_error": one_d_convert(enrichment_se),
                "Enrichment_p": diff_p,
                "Coefficient": one_d_convert(self.coef),
                "Coefficient_std_error": self.coef_se,
                "Coefficient_z-score": one_d_convert(self.coef)
                / one_d_convert(self.coef_se),
            }
        )
        if print_coefficients:
            df = df[
                [
                    "Category",
                    "Prop._SNPs",
                    "Prop._h2",
                    "Prop._h2_std_error",
                    "Enrichment",
                    "Enrichment_std_error",
                    "Enrichment_p",
                    "Coefficient",
                    "Coefficient_std_error",
                    "Coefficient_z-score",
                ]
            ]
        else:
            df = df[
                [
                    "Category",
                    "Prop._SNPs",
                    "Prop._h2",
                    "Prop._h2_std_error",
                    "Enrichment",
                    "Enrichment_std_error",
                    "Enrichment_p",
                ]
            ]
        return df

    def summary(self, ref_ld_colnames=None, P=None, K=None, overlap=False):
        """Print summary of the LD Score Regression."""
        if P is not None and K is not None:
            T = "Liability"
            c = h2_obs_to_liab(1, P, K)
        else:
            T = "Observed"
            c = 1

        out = [
            "Total "
            + T
            + " scale h2: "
            + s(c * self.tot)
            + " ("
            + s(c * self.tot_se)
            + ")"
        ]
        if self.n_annot > 1:
            if ref_ld_colnames is None:
                ref_ld_colnames = ["CAT_" + str(i) for i in range(self.n_annot)]

            out.append("Categories: " + " ".join(ref_ld_colnames))

            if not overlap:
                out.append(T + " scale h2: " + s(c * self.cat))
                out.append(T + " scale h2 SE: " + s(c * self.cat_se))
                out.append("Proportion of SNPs: " + s(self.M_prop))
                out.append("Proportion of h2g: " + s(self.prop))
                out.append("Enrichment: " + s(self.enrichment))
                out.append("Coefficients: " + s(self.coef))
                out.append("Coefficient SE: " + s(self.coef_se))

        out.append("Lambda GC: " + s(self.lambda_gc))
        out.append("Mean Chi^2: " + s(self.mean_chisq))
        if self.constrain_intercept:
            out.append("Intercept: constrained to {C}".format(C=s(self.intercept)))
        else:
            out.append(
                "Intercept: " + s(self.intercept) + " (" + s(self.intercept_se) + ")"
            )
            if self.mean_chisq > 1:
                if self.ratio < 0:
                    out.append("Ratio < 0 (usually indicates GC correction).")
                else:
                    out.append(
                        "Ratio: " + s(self.ratio) + " (" + s(self.ratio_se) + ")"
                    )
            else:
                out.append("Ratio: NA (mean chi^2 < 1)")

        return remove_brackets("\n".join(out))

    def _update_weights(self, ld, w_ld, N, M, hsq, intercept, ii=None):
        if intercept is None:
            intercept = self.__null_intercept__

        return self.weights(ld, w_ld, N, M, hsq, intercept, ii)

    @classmethod
    def weights(cls, ld, w_ld, N, M, hsq, intercept=None, ii=None):
        """
        Regression weights.

        Parameters
        ----------
        ld : np.array with shape (n_snp, 1)
            LD Scores (non-partitioned).
        w_ld : np.array with shape (n_snp, 1)
            LD Scores (non-partitioned) computed with sum r^2 taken over only those SNPs included
            in the regression.
        N :  np.array of ints > 0 with shape (n_snp, 1)
            Number of individuals sampled for each SNP.
        M : float > 0
            Number of SNPs used for estimating LD Score (need not equal number of SNPs included in
            the regression).
        hsq : float in [0,1]
            Heritability estimate.

        Returns
        -------
        w : np.array with shape (n_snp, 1)
            Regression weights. Approx equal to reciprocal of conditional variance function.

        """
        if isinstance(M, np.ndarray):
            M = M.item()
        M = float(M)
        if intercept is None:
            intercept = 1

        hsq = max(hsq, 0.0)
        hsq = min(hsq, 1.0)
        ld = np.fmax(ld, 1.0)
        w_ld = np.fmax(w_ld, 1.0)
        c = hsq * N / M
        het_w = 1.0 / (2 * np.square(intercept + np.multiply(c, ld)))
        oc_w = 1.0 / w_ld
        w = np.multiply(het_w, oc_w)
        return w


class Gencov(LD_Score_Regression):
    __null_intercept__ = 0

    def __init__(
        self,
        z1,
        z2,
        x,
        w,
        N1,
        N2,
        M,
        hsq1,
        hsq2,
        intercept_hsq1,
        intercept_hsq2,
        n_blocks=200,
        intercept_gencov=None,
        slow=False,
        twostep=None,
    ):
        self.intercept_hsq1 = intercept_hsq1
        self.intercept_hsq2 = intercept_hsq2
        self.hsq1 = hsq1
        self.hsq2 = hsq2
        self.N1 = N1
        self.N2 = N2
        y = z1 * z2
        step1_ii = None
        if twostep is not None:
            step1_ii = np.logical_and(z1**2 < twostep, z2**2 < twostep)

        LD_Score_Regression.__init__(
            self,
            y,
            x,
            w,
            np.sqrt(N1 * N2),
            M,
            n_blocks,
            intercept=intercept_gencov,
            slow=slow,
            step1_ii=step1_ii,
        )
        self.p, self.z = p_z_norm(self.tot, self.tot_se)
        self.mean_z1z2 = np.mean(np.multiply(z1, z2))

    def summary(self, ref_ld_colnames, P=None, K=None):
        """Print summary of the LD Score regression."""
        out = []
        if (
            P is not None
            and K is not None
            and all((i is not None for i in P))
            and all((i is not None for i in K))
        ):
            T = "Liability"
            c = gencov_obs_to_liab(1, P[0], P[1], K[0], K[1])
        else:
            T = "Observed"
            c = 1

        out.append(
            "Total "
            + T
            + " scale gencov: "
            + s(c * self.tot)
            + " ("
            + s(c * self.tot_se)
            + ")"
        )
        if self.n_annot > 1:
            out.append("Categories: " + str(" ".join(ref_ld_colnames)))
            out.append(T + " scale gencov: " + s(c * self.cat))
            out.append(T + " scale gencov SE: " + s(c * self.cat_se))
            out.append("Proportion of SNPs: " + s(self.M_prop))
            out.append("Proportion of gencov: " + s(self.prop))
            out.append("Enrichment: " + s(self.enrichment))

        out.append("Mean z1*z2: " + s(self.mean_z1z2))
        if self.constrain_intercept:
            out.append("Intercept: constrained to {C}".format(C=s(self.intercept)))
        else:
            out.append(
                "Intercept: " + s(self.intercept) + " (" + s(self.intercept_se) + ")"
            )

        return remove_brackets("\n".join(out))

    def _update_func(self, x, ref_ld_tot, w_ld, N, M, Nbar, intercept=None, ii=None):
        """
        Update function for IRWLS
        x is the output of np.linalg.lstsq.
        x[0] is the regression coefficients
        x[0].shape is (# of dimensions, 1)
        the last element of x[0] is the intercept.

        """
        rho_g = M * x[0][0] / Nbar
        if intercept is None:  # if the regression includes an intercept
            intercept = x[0][1]

        # remove intercept if we have one
        ld = ref_ld_tot[:, 0].reshape(w_ld.shape)
        if ii is not None:
            N1 = self.N1[ii].reshape((w_ld.shape))
            N2 = self.N2[ii].reshape((w_ld.shape))
        else:
            N1 = self.N1
            N2 = self.N2

        return self.weights(
            ld,
            w_ld,
            N1,
            N2,
            np.sum(M),
            self.hsq1,
            self.hsq2,
            rho_g,
            intercept,
            self.intercept_hsq1,
            self.intercept_hsq2,
            ii,
        )

    def _update_weights(self, ld, w_ld, sqrt_n1n2, M, rho_g, intercept, ii=None):
        """Weight function with the same signature for Hsq and Gencov."""
        w = self.weights(
            ld,
            w_ld,
            self.N1,
            self.N2,
            M,
            self.hsq1,
            self.hsq2,
            rho_g,
            intercept,
            self.intercept_hsq1,
            self.intercept_hsq2,
        )
        return w

    @classmethod
    def weights(
        cls,
        ld,
        w_ld,
        N1,
        N2,
        M,
        h1,
        h2,
        rho_g,
        intercept_gencov=None,
        intercept_hsq1=None,
        intercept_hsq2=None,
        ii=None,
    ):
        """
        Regression weights.

        Parameters
        ----------
        ld : np.array with shape (n_snp, 1)
            LD Scores (non-partitioned)
        w_ld : np.array with shape (n_snp, 1)
            LD Scores (non-partitioned) computed with sum r^2 taken over only those SNPs included
            in the regression.
        M : float > 0
            Number of SNPs used for estimating LD Score (need not equal number of SNPs included in
            the regression).
        N1, N2 :  np.array of ints > 0 with shape (n_snp, 1)
            Number of individuals sampled for each SNP for each study.
        h1, h2 : float in [0,1]
            Heritability estimates for each study.
        rhog : float in [0,1]
            Genetic covariance estimate.
        intercept : float
            Genetic covariance intercept, on the z1*z2 scale (so should be Ns*rho/sqrt(N1*N2)).

        Returns
        -------
        w : np.array with shape (n_snp, 1)
            Regression weights. Approx equal to reciprocal of conditional variance function.

        """
        M = float(M)
        if intercept_gencov is None:
            intercept_gencov = 0
        if intercept_hsq1 is None:
            intercept_hsq1 = 1
        if intercept_hsq2 is None:
            intercept_hsq2 = 1

        h1, h2 = max(h1, 0.0), max(h2, 0.0)
        h1, h2 = min(h1, 1.0), min(h2, 1.0)
        rho_g = min(rho_g, 1.0)
        rho_g = max(rho_g, -1.0)
        ld = np.fmax(ld, 1.0)
        w_ld = np.fmax(w_ld, 1.0)
        a = np.multiply(N1, h1 * ld) / M + intercept_hsq1
        b = np.multiply(N2, h2 * ld) / M + intercept_hsq2
        sqrt_n1n2 = np.sqrt(np.multiply(N1, N2))
        c = np.multiply(sqrt_n1n2, rho_g * ld) / M + intercept_gencov
        try:
            het_w = 1.0 / (np.multiply(a, b) + np.square(c))
        except FloatingPointError:  # bizarre error; should never happen
            raise FloatingPointError("Why did you set hsq intercept <= 0?")

        oc_w = 1.0 / w_ld
        w = np.multiply(het_w, oc_w)
        return w


class RG(object):
    def __init__(
        self,
        z1,
        z2,
        x,
        w,
        N1,
        N2,
        M,
        intercept_hsq1=None,
        intercept_hsq2=None,
        intercept_gencov=None,
        n_blocks=200,
        slow=False,
        twostep=None,
    ):
        """
        Parameters
        ----------
        z1 : np.array
            Z scores for phenotype 1.
        z2 : np.array
            Z scores for phenotype 2.
        x : np.array
            LD scores.
        w : np.array
        N1 : np.array
            Sample size for phenotype 1.
        N2 : np.array
            Sample size for phenotype 2.
        M : np.array
        intercept_hsq1 : float
        intercept_hsq2 : float
        intercept_gencov : float
        n_blocks : int
        slow : bool
        twostep : float
        """
        self.intercept_gencov = intercept_gencov
        self._negative_hsq = None
        n_snp, n_annot = x.shape
        hsq1 = Hsq(
            np.square(z1),
            x,
            w,
            N1,
            M,
            n_blocks=n_blocks,
            intercept=intercept_hsq1,
            slow=slow,
            twostep=twostep,
        )
        hsq2 = Hsq(
            np.square(z2),
            x,
            w,
            N2,
            M,
            n_blocks=n_blocks,
            intercept=intercept_hsq2,
            slow=slow,
            twostep=twostep,
        )
        gencov = Gencov(
            z1,
            z2,
            x,
            w,
            N1,
            N2,
            M,
            hsq1.tot,
            hsq2.tot,
            hsq1.intercept,
            hsq2.intercept,
            n_blocks,
            intercept_gencov=intercept_gencov,
            slow=slow,
            twostep=twostep,
        )
        gencov.N1 = None  # save memory
        gencov.N2 = None
        self.hsq1, self.hsq2, self.gencov = hsq1, hsq2, gencov
        if hsq1.tot <= 0 or hsq2.tot <= 0:
            self._negative_hsq = True
            self.rg_ratio = self.rg = self.rg_se = "NA"
            self.p = self.z = "NA"
        else:
            rg_ratio = np.array(gencov.tot / np.sqrt(hsq1.tot * hsq2.tot)).reshape(
                (1, 1)
            )
            denom_delete_values = np.sqrt(
                np.multiply(hsq1.tot_delete_values, hsq2.tot_delete_values)
            )
            rg = jk.RatioJackknife(
                rg_ratio, gencov.tot_delete_values, denom_delete_values
            )
            if isinstance(rg.jknife_est, np.ndarray):
                rg.jknife_est = rg.jknife_est.item()
            self.rg_jknife = float(rg.jknife_est)
            if isinstance(rg.jknife_se, np.ndarray):
                rg.jknife_se = rg.jknife_se.item()
            self.rg_se = float(rg.jknife_se)
            if isinstance(rg_ratio, np.ndarray):
                rg_ratio = rg_ratio.item()
            self.rg_ratio = float(rg_ratio)
            self.p, self.z = p_z_norm(self.rg_ratio, self.rg_se)

    def summary(self, silly=False):
        """Print output of Gencor object."""
        out = []
        if self._negative_hsq:
            out.append("Genetic Correlation: nan (nan) (h2  out of bounds) ")
            out.append("Z-score: nan (nan) (h2  out of bounds)")
            out.append("P: nan (nan) (h2  out of bounds)")
            out.append("WARNING: One of the h2's was out of bounds.")
            out.append(
                "This usually indicates a data-munging error "
                + "or that h2 or N is low."
            )
        elif (self.rg_ratio > 1.2 or self.rg_ratio < -1.2) and not silly:
            out.append("Genetic Correlation: nan (nan) (rg out of bounds) ")
            out.append("Z-score: nan (nan) (rg out of bounds)")
            out.append("P: nan (nan) (rg out of bounds)")
            out.append("WARNING: rg was out of bounds.")
            if self.intercept_gencov is None:
                out.append(
                    "This often means that h2 is not significantly "
                    + "different from zero."
                )
            else:
                out.append(
                    "This often means that you have constrained"
                    + " the intercepts to the wrong values."
                )
        else:
            out.append(
                "Genetic Correlation: " + s(self.rg_ratio) + " (" + s(self.rg_se) + ")"
            )
            out.append("Z-score: " + s(self.z))
            out.append("P: " + s(self.p))
        return remove_brackets("\n".join(out))
