#!/usr/bin/env python3

__author__ = "Anarya Ray, Caitlin Rose, Vinaya Valsan"


import os
import scipy

import numpy as np
import matplotlib.pyplot as plt

from gstlal import far


class pastro:
    def __init__(
        self,
        evidence_dict,
        rankstatpdf_file,
        likelihood,
        rate_dict,
        far_threshold,
    ):
        self.evidence_dict = evidence_dict
        self.rate_dict = rate_dict
        self.far_threshold = far_threshold
        self.rankstatpdf_file = rankstatpdf_file
        self.likelihood = likelihood

    def normalize_f_over_b(self):
        rankingstatpdf = far.marginalize_pdf_urls(
            [self.rankstatpdf_file], "RankingStatPDF"
        )

        zl = rankingstatpdf.zero_lag_lr_lnpdf.copy()
        zl.array[:40] = 0.0

        if not zl.array.any():
            raise ValueError("zero-lag counts are all zero")
        (ln_likelihood_ratio_threshold,) = zl.argmax()

        rankingstatpdf = rankingstatpdf.new_with_extinction()
        fapfar = far.FAPFAR(rankingstatpdf)

        xmin = max(
            fapfar.rank_from_far(self.far_threshold),
            ln_likelihood_ratio_threshold,
        )

        rankingstatpdf.noise_lr_lnpdf.array[
            : rankingstatpdf.noise_lr_lnpdf.bins[0][xmin]
        ] = 0.0
        rankingstatpdf.noise_lr_lnpdf.normalize()
        rankingstatpdf.signal_lr_lnpdf.array[
            : rankingstatpdf.signal_lr_lnpdf.bins[0][xmin]
        ] = 0.0
        rankingstatpdf.signal_lr_lnpdf.normalize()
        rankingstatpdf.zero_lag_lr_lnpdf.array[
            : rankingstatpdf.zero_lag_lr_lnpdf.bins[0][xmin]
        ] = 0.0
        rankingstatpdf.zero_lag_lr_lnpdf.normalize()

        noise_lr_lnpdf = rankingstatpdf.noise_lr_lnpdf.mkinterp()
        lnb = np.vectorize(lambda x: noise_lr_lnpdf(x))

        B, dB = scipy.integrate.quad(
            lambda x: np.exp(lnb(x)), xmin, np.inf, limit=1000
        )

        A, dA = scipy.integrate.quad(
            lambda x: np.exp(x + lnb(x) - np.log(B)), xmin, np.inf, limit=1000
        )

        return np.reciprocal(A)

    def pipeline_bayesfactor(self):
        return self.normalize_f_over_b() * np.exp(self.likelihood)

    def compute_pastro(self):
        R_tot = sum(self.rate_dict.values())
        print(f"pipeline_bayesfactor = {self.pipeline_bayesfactor()}")

        evidence_sum = (
            1.0
            if self.pipeline_bayesfactor() == 1.0
            else (
                self.evidence_dict["evidence_BNS"]
                + self.evidence_dict["evidence_NSBH"]
                + self.evidence_dict["evidence_BBH"]
            )
            / 3.0
        )
        pH = {
            "Terrestrial": self.far_threshold / R_tot,
            "BBH": self.rate_dict["R_bbh"] / R_tot,
            "NSBH": self.rate_dict["R_nsbh"] / R_tot,
            "BNS": self.rate_dict["R_bns"] / R_tot,
        }

        evidence_factor = self.pipeline_bayesfactor() / evidence_sum

        evidence_ratio = {
            "Terrestrial": 1.0,
            "BBH": self.evidence_dict["evidence_BBH"] * evidence_factor,
            "NSBH": self.evidence_dict["evidence_NSBH"] * evidence_factor,
            "BNS": self.evidence_dict["evidence_BNS"] * evidence_factor,
        }

        pastros_numerator = {k: z * pH[k] for k, z in evidence_ratio.items()}

        pastros_denominator = sum(pastros_numerator.values())
        pastros = {
            k: numerator / pastros_denominator
            for k, numerator in pastros_numerator.items()
        }
        return pastros


def plot_pastro(pastro_dict, save_dir, filename="p_astro.png"):
    """
    adapted from gwcelery.tasks
    """

    def _format_prob(prob):
        if prob >= 1:
            return "100%"
        elif prob <= 0:
            return "0%"
        elif prob > 0.99:
            return ">99%"
        elif prob < 0.01:
            return "<1%"
        else:
            return "{}%".format(int(np.round(100 * prob)))

    fname = os.path.join(save_dir, filename)

    probs = list(pastro_dict.values())
    names = list(pastro_dict.keys())

    with plt.style.context("seaborn-white"):
        fig, ax = plt.subplots(figsize=(2.5, 2))
        ax.barh(names, probs)
        for i, prob in enumerate(probs):
            ax.annotate(
                _format_prob(prob),
                (0, i),
                (4, 0),
                textcoords="offset points",
                ha="left",
                va="center",
            )
        ax.set_xlim(0, 1)
        ax.set_xticks([])
        ax.tick_params(left=False)
        for side in ["top", "bottom", "right"]:
            ax.spines[side].set_visible(False)
        fig.tight_layout()
        fig.savefig(fname)
    return
