import sys
import warnings
from itertools import product

import numpy as np
from graspy.simulations import sample_edges
from rpy2.robjects.packages import importr
from tqdm import tqdm

# Import tests from R
barnard = importr("Barnard")
exact = importr("exact2x2")

# Try to mute any warnings
if not sys.warnoptions:
    warnings.simplefilter("ignore")


class IndependentEdge:
    def __init__(self, sample_size, n_vertices, epsilon, delta):

        self.sample_size = int(sample_size)
        self.n_vertices = int(n_vertices)
        self.epsilon = epsilon
        self.delta = delta

        self.p1, self.p2 = self._generate_p_matrices()

    def _generate_p_matrices(self):

        # Generate p1
        b = np.linspace(self.epsilon, 1 - self.epsilon, num=self.n_vertices ** 2)
        p1 = b.reshape(self.n_vertices, self.n_vertices)

        # Generate p2
        p2 = np.copy(p1)
        with np.nditer(p2, op_flags=["readwrite"]) as it:
            for x in it:
                if x < 0.5:
                    x[...] = x + self.delta
                else:
                    x[...] = x - self.delta

        return p1, p2

    def _sample(self):
        """
        Sample graphs from IE model

        Returns
        -------
        x, y : np.ndarray, shapes (m, n, n)
            Samples from two models
        """

        x = [
            sample_edges(self.p1, directed=True, loops=True)
            for _ in range(self.sample_size)
        ]
        y = [
            sample_edges(self.p2, directed=True, loops=True)
            for _ in range(self.sample_size)
        ]

        x = np.stack(x)
        y = np.stack(y)

        return x, y

    def calculate_pvals(self, scipy_methods=[], r_methods=[], n_iter=1000):
        """
        Calculate the power of a given test

        Parameters
        ----------
        scipy_methods : list of functions
            Statistical tests from scipy.stats
            Assumes function returns are of the form (statistic, p-value)
        r_methods : list of strings
            Strings corresponding to methods in R
        n_iter : int (default = 100)
            Number of Monte Carlo runs.

        Returns
        -------
        pvals : np.ndarray, shape (n_tests, n_vertices, n_vertices, n_iter)
            All p-values generated by test
        """

        # Construct a tensor to store all of the p-values
        tests = scipy_methods + r_methods
        n_tests = len(tests)
        pvals = np.zeros(shape=(n_tests, self.n_vertices, self.n_vertices, n_iter))

        for iteration in tqdm(range(n_iter)):

            # Sample the independent edge model
            x, y = self._sample()

            for idx, test in enumerate(tests):
                for i, j in product(range(self.n_vertices), range(self.n_vertices)):

                    # Get the number nonzero edges
                    xi = x[:, i, j]
                    yi = y[:, i, j]
                    xi_n_zero = np.count_nonzero(xi)
                    yi_n_zero = np.count_nonzero(yi)

                    if np.array_equal(xi, yi):
                        pval = 1
                    elif test == "boschloo":
                        pval = exact.boschloo(
                            xi_n_zero,
                            self.sample_size,
                            yi_n_zero,
                            self.sample_size,
                            alternative="two.sided",
                        ).rx2("p.value")[0]
                    elif test == "barnard":
                        pval = barnard.barnard_test(
                            xi_n_zero,
                            self.sample_size - xi_n_zero,
                            yi_n_zero,
                            self.sample_size - yi_n_zero,
                        ).rx2("p.value")[1]
                    elif test.__name__ == "fisher_exact":
                        data = [
                            [xi_n_zero, self.sample_size - xi_n_zero],
                            [yi_n_zero, self.sample_size - yi_n_zero],
                        ]
                        _, pval = test(data)
                    else:
                        try:
                            _, pval = test(xi, yi)
                        except:
                            print("Test not recognized: {}".format(test))

                    # Add pvalue to corresponding element in matrix
                    pvals[idx, i, j, iteration] = pval

        return pvals

    def calculate_proportion_positive(self, pvals, alpha=0.05):

        """
        If the null is true (ie p1 = p2), then proportion of positive results
        is Type 1 Error. If the null is false (ie p1 != p2), then the
        proportion of positive results is power.
        """

        n_tests, n_vertices, n_vertices, n_iter = pvals.shape
        assert n_vertices == self.n_vertices

        positives = pvals < alpha
        positives = np.sum(positives, axis=3)
        positives = positives / n_iter

        return positives
