import numpy as np
import pandas as pd
import pytest

import statsmodels.api as sm
from statsmodels.sandbox.predict_functional import predict_functional

# If true, the output is written to a multi-page pdf file.
pdf_output = False

try:
    import matplotlib.pyplot as plt
except ImportError:
    pass


def pctl(q):
    return lambda x: np.percentile(x, 100 * q)


class TestPredFunc:

    @classmethod
    def setup_class(cls):
        if pdf_output:
            from matplotlib.backends.backend_pdf import PdfPages

            cls.pdf = PdfPages("predict_functional.pdf")

    @classmethod
    def teardown_class(cls):
        if pdf_output:
            cls.pdf.close()

    def close_or_save(self, fig):
        if pdf_output:
            self.pdf.savefig(fig)

    @pytest.mark.matplotlib
    def test_formula(self, close_figures):

        np.random.seed(542)
        n = 500
        x1 = np.random.normal(size=n)
        x2 = np.random.normal(size=n)
        x3 = np.random.normal(size=n)
        x4 = np.random.randint(0, 5, size=n)
        x4 = np.asarray(["ABCDE"[i] for i in x4])
        x5 = np.random.normal(size=n)
        y = (
            0.3 * x2**2
            + (x4 == "B")
            + 0.1 * (x4 == "B") * x2**2
            + x5
            + np.random.normal(size=n)
        )

        df = pd.DataFrame({"y": y, "x1": x1, "x2": x2, "x3": x3, "x4": x4, "x5": x5})

        fml = "y ~ x1 + bs(x2, df=4) + x3 + x2*x3 + I(x1**2) + C(x4) + C(x4)*bs(x2, df=4) + x5"
        model = sm.OLS.from_formula(fml, data=df)
        result = model.fit()

        summaries = {"x1": np.mean, "x3": pctl(0.75), "x5": np.mean}

        values = {"x4": "B"}
        pr1, ci1, fvals1 = predict_functional(result, "x2", summaries, values)

        values = {"x4": "C"}
        pr2, ci2, fvals2 = predict_functional(result, "x2", summaries, values)

        plt.clf()
        fig = plt.figure()
        ax = plt.axes([0.1, 0.1, 0.7, 0.8])
        plt.plot(fvals1, pr1, "-", label="x4=B")
        plt.plot(fvals2, pr2, "-", label="x4=C")
        ha, lb = ax.get_legend_handles_labels()
        plt.figlegend(ha, lb, loc="center right")
        plt.xlabel("Focus variable", size=15)
        plt.ylabel("Fitted mean", size=15)
        plt.title("Linear model prediction")
        self.close_or_save(fig)

        plt.clf()
        fig = plt.figure()
        ax = plt.axes([0.1, 0.1, 0.7, 0.8])
        plt.plot(fvals1, pr1, "-", label="x4=B")
        plt.fill_between(fvals1, ci1[:, 0], ci1[:, 1], color="grey")
        plt.plot(fvals2, pr2, "-", label="x4=C")
        plt.fill_between(fvals2, ci2[:, 0], ci2[:, 1], color="grey")
        ha, lb = ax.get_legend_handles_labels()
        plt.figlegend(ha, lb, loc="center right")
        plt.xlabel("Focus variable", size=15)
        plt.ylabel("Fitted mean", size=15)
        plt.title("Linear model prediction")
        self.close_or_save(fig)

    @pytest.mark.matplotlib
    def test_lm_contrast(self, close_figures):

        np.random.seed(542)
        n = 200
        x1 = np.random.normal(size=n)
        x2 = np.random.normal(size=n)
        x3 = np.random.normal(size=n)
        y = x1 + 2 * x2 + x3 - x1 * x2 + x2 * x3 + np.random.normal(size=n)

        df = pd.DataFrame({"y": y, "x1": x1, "x2": x2, "x3": x3})

        fml = "y ~ x1 + x2 + x3 + x1*x2 + x2*x3"
        model = sm.OLS.from_formula(fml, data=df)
        result = model.fit()

        values = {"x2": 1, "x3": 1}  # y = 4
        values2 = {"x2": 0, "x3": 0}  # y = x1
        pr, cb, fvals = predict_functional(
            result, "x1", values=values, values2=values2, ci_method="scheffe"
        )

        plt.clf()
        fig = plt.figure()
        ax = plt.axes([0.1, 0.1, 0.67, 0.8])
        plt.plot(fvals, pr, "-", label="Estimate", color="orange", lw=4)
        plt.plot(fvals, 4 - fvals, "-", label="Truth", color="lime", lw=4)
        plt.fill_between(fvals, cb[:, 0], cb[:, 1], color="grey")
        ha, lb = ax.get_legend_handles_labels()
        leg = plt.figlegend(ha, lb, loc="center right")
        leg.draw_frame(False)
        plt.xlabel("Focus variable", size=15)
        plt.ylabel("Mean contrast", size=15)
        plt.title("Linear model contrast")
        self.close_or_save(fig)

    @pytest.mark.matplotlib
    def test_glm_formula_contrast(self, close_figures):

        np.random.seed(542)
        n = 50
        x1 = np.random.normal(size=n)
        x2 = np.random.normal(size=n)
        x3 = np.random.normal(size=n)
        mn = 5 + 0.1 * x1 + 0.1 * x2 + 0.1 * x3 - 0.1 * x1 * x2
        y = np.random.poisson(np.exp(mn), size=len(mn))

        df = pd.DataFrame({"y": y, "x1": x1, "x2": x2, "x3": x3})

        fml = "y ~ x1 + x2 + x3 + x1*x2"
        model = sm.GLM.from_formula(fml, data=df, family=sm.families.Poisson())
        result = model.fit()

        values = {"x2": 1, "x3": 1}  # y = 5.2
        values2 = {"x2": 0, "x3": 0}  # y = 5 + 0.1*x1
        pr, cb, fvals = predict_functional(
            result, "x1", values=values, values2=values2, ci_method="simultaneous"
        )

        plt.clf()
        fig = plt.figure()
        ax = plt.axes([0.1, 0.1, 0.67, 0.8])
        plt.plot(fvals, pr, "-", label="Estimate", color="orange", lw=4)
        plt.plot(fvals, 0.2 - 0.1 * fvals, "-", label="Truth", color="lime", lw=4)
        plt.fill_between(fvals, cb[:, 0], cb[:, 1], color="grey")
        ha, lb = ax.get_legend_handles_labels()
        leg = plt.figlegend(ha, lb, loc="center right")
        leg.draw_frame(False)
        plt.xlabel("Focus variable", size=15)
        plt.ylabel("Linear predictor contrast", size=15)
        plt.title("Poisson regression contrast")
        self.close_or_save(fig)

    @pytest.mark.matplotlib
    def test_scb(self, close_figures):

        np.random.seed(473)
        n = 100
        x = np.random.normal(size=(n, 4))
        x[:, 0] = 1

        for fam_name in "poisson", "binomial", "gaussian":

            if fam_name == "poisson":
                y = np.random.poisson(20, size=n)
                fam = sm.families.Poisson()
                true_mean = 20
                true_lp = np.log(20)
            elif fam_name == "binomial":
                y = 1 * (np.random.uniform(size=n) < 0.5)
                fam = sm.families.Binomial()
                true_mean = 0.5
                true_lp = 0
            elif fam_name == "gaussian":
                y = np.random.normal(size=n)
                fam = sm.families.Gaussian()
                true_mean = 0
                true_lp = 0

            model = sm.GLM(y, x, family=fam)
            result = model.fit()

            # CB is for linear predictor or mean response
            for linear in False, True:

                true = true_lp if linear else true_mean

                values = {"const": 1, "x2": 0}
                summaries = {"x3": np.mean}
                pred1, cb1, fvals1 = predict_functional(
                    result, "x1", values=values, summaries=summaries, linear=linear
                )
                pred2, cb2, fvals2 = predict_functional(
                    result,
                    "x1",
                    values=values,
                    summaries=summaries,
                    ci_method="simultaneous",
                    linear=linear,
                )

                plt.clf()
                fig = plt.figure()
                ax = plt.axes([0.1, 0.1, 0.58, 0.8])
                plt.plot(fvals1, pred1, "-", color="black", label="Estimate")
                plt.plot(
                    fvals1,
                    true * np.ones(len(pred1)),
                    "-",
                    color="purple",
                    label="Truth",
                )
                plt.plot(fvals1, cb1[:, 0], color="blue", label="Pointwise CB")
                plt.plot(fvals1, cb1[:, 1], color="blue")
                plt.plot(fvals2, cb2[:, 0], color="green", label="Simultaneous CB")
                plt.plot(fvals2, cb2[:, 1], color="green")
                ha, lb = ax.get_legend_handles_labels()
                leg = plt.figlegend(ha, lb, loc="center right")
                leg.draw_frame(False)
                plt.xlabel("Focus variable", size=15)
                if linear:
                    plt.ylabel("Linear predictor", size=15)
                else:
                    plt.ylabel("Fitted mean", size=15)
                plt.title("%s family prediction" % fam_name.capitalize())

                self.close_or_save(fig)

    @pytest.mark.matplotlib
    def test_glm_formula(self, close_figures):

        np.random.seed(542)
        n = 500
        x1 = np.random.normal(size=n)
        x2 = np.random.normal(size=n)
        x3 = np.random.randint(0, 3, size=n)
        x3 = np.asarray(["ABC"[i] for i in x3])
        lin_pred = -1 + 0.5 * x1**2 + (x3 == "B")
        prob = 1 / (1 + np.exp(-lin_pred))
        y = 1 * (np.random.uniform(size=n) < prob)

        df = pd.DataFrame({"y": y, "x1": x1, "x2": x2, "x3": x3})

        fml = "y ~ x1 + I(x1**2) + x2 + C(x3)"
        model = sm.GLM.from_formula(fml, family=sm.families.Binomial(), data=df)
        result = model.fit()
        summaries = {"x2": np.mean}

        for linear in False, True:

            values = {"x3": "B"}
            pr1, ci1, fvals1 = predict_functional(
                result, "x1", summaries, values, linear=linear
            )

            values = {"x3": "C"}
            pr2, ci2, fvals2 = predict_functional(
                result, "x1", summaries, values, linear=linear
            )

            exact1 = -1 + 0.5 * fvals1**2 + 1
            exact2 = -1 + 0.5 * fvals2**2

            if not linear:
                exact1 = 1 / (1 + np.exp(-exact1))
                exact2 = 1 / (1 + np.exp(-exact2))

            plt.clf()
            fig = plt.figure()
            ax = plt.axes([0.1, 0.1, 0.7, 0.8])
            plt.plot(fvals1, pr1, "-", label="x3=B")
            plt.plot(fvals2, pr2, "-", label="x3=C")
            plt.plot(fvals1, exact1, "-", label="x3=B (exact)")
            plt.plot(fvals2, exact2, "-", label="x3=C (exact)")
            ha, lb = ax.get_legend_handles_labels()
            plt.figlegend(ha, lb, loc="center right")
            plt.xlabel("Focus variable", size=15)
            if linear:
                plt.ylabel("Fitted linear predictor", size=15)
            else:
                plt.ylabel("Fitted probability", size=15)
            plt.title("Binomial GLM prediction")
            self.close_or_save(fig)

            plt.clf()
            fig = plt.figure()
            ax = plt.axes([0.1, 0.1, 0.7, 0.8])
            plt.plot(fvals1, pr1, "-", label="x3=B", color="orange")
            plt.fill_between(fvals1, ci1[:, 0], ci1[:, 1], color="grey")
            plt.plot(fvals2, pr2, "-", label="x3=C", color="lime")
            plt.fill_between(fvals2, ci2[:, 0], ci2[:, 1], color="grey")
            ha, lb = ax.get_legend_handles_labels()
            plt.figlegend(ha, lb, loc="center right")
            plt.xlabel("Focus variable", size=15)
            if linear:
                plt.ylabel("Fitted linear predictor", size=15)
            else:
                plt.ylabel("Fitted probability", size=15)
            plt.title("Binomial GLM prediction")
            self.close_or_save(fig)

    @pytest.mark.matplotlib
    def test_noformula_prediction(self, close_figures):

        np.random.seed(6434)
        n = 200
        x1 = np.random.normal(size=n)
        x2 = np.random.normal(size=n)
        x3 = np.random.normal(size=n)
        y = x1 - x2 + np.random.normal(size=n)

        exog = np.vstack((x1, x2, x3)).T

        model = sm.OLS(y, exog)
        result = model.fit()

        summaries = {"x3": pctl(0.75)}
        values = {"x2": 1}
        pr1, ci1, fvals1 = predict_functional(result, "x1", summaries, values)

        values = {"x2": -1}
        pr2, ci2, fvals2 = predict_functional(result, "x1", summaries, values)

        plt.clf()
        fig = plt.figure()
        ax = plt.axes([0.1, 0.1, 0.7, 0.8])
        plt.plot(fvals1, pr1, "-", label="x2=1", lw=4, alpha=0.6, color="orange")
        plt.plot(fvals2, pr2, "-", label="x2=-1", lw=4, alpha=0.6, color="lime")
        ha, lb = ax.get_legend_handles_labels()
        leg = plt.figlegend(ha, lb, loc="center right")
        leg.draw_frame(False)
        plt.xlabel("Focus variable", size=15)
        plt.ylabel("Fitted mean", size=15)
        plt.title("Linear model prediction")
        self.close_or_save(fig)

        plt.clf()
        fig = plt.figure()
        ax = plt.axes([0.1, 0.1, 0.7, 0.8])
        plt.plot(fvals1, pr1, "-", label="x2=1", lw=4, alpha=0.6, color="orange")
        plt.fill_between(fvals1, ci1[:, 0], ci1[:, 1], color="grey")
        plt.plot(fvals1, pr2, "-", label="x2=1", lw=4, alpha=0.6, color="lime")
        plt.fill_between(fvals2, ci2[:, 0], ci2[:, 1], color="grey")
        ha, lb = ax.get_legend_handles_labels()
        plt.figlegend(ha, lb, loc="center right")
        plt.xlabel("Focus variable", size=15)
        plt.ylabel("Fitted mean", size=15)
        plt.title("Linear model prediction")
        self.close_or_save(fig)
