import matplotlib.pyplot as plt
from matplotlib import gridspec
import numpy as np
import pandas as pd
import seaborn as sns
from qp.metrics.pit import PIT
from IPython.display import Markdown
import h5py
import os
from qp import Ensemble
from qp import interp


def read_pz_output(
    pdfs_file,
    ztrue_file,
    pdfs_key="photoz_pdf",
    zgrid_key="zgrid",
    photoz_mode_key="photoz_mode",
    ztrue_key="redshift",
):
    _, ext = os.path.splitext(pdfs_file)

    if ext == ".hdf5":
        with h5py.File(ztrue_file, "r") as zf:
            try:
                ztrue = np.array(zf["photometry"][ztrue_key])
            except:
                try:
                    ztrue = np.array(zf[ztrue_key])
                except:
                    raise ValueError(
                        "Invalid key for true redshift column in ztrue file."
                    )
        with h5py.File(pdfs_file, "r") as pf:
            pdfs = np.array(pf[pdfs_key])
            zgrid = np.array(pf[zgrid_key]).flatten()
            photoz_mode = np.array(pf[photoz_mode_key])
    elif ext == ".out":
        print("Validation file from DC1 paper (ascii format).")
        ztrue = np.loadtxt(ztrue_file, unpack=True, usecols=[2])
        pdfs = np.loadtxt(pdfs_file)
        path = "/".join(pdfs_file.split("/")[:-1])
        zgrid = np.loadtxt(path + "/zarrayfile.out")
        photoz_mode = np.array([zgrid[np.argmax(pdf)] for pdf in pdfs])  # qp mode?
    else:
        raise ValueError(f"PDFs input file format {ext} is not supported.")

    return pdfs, zgrid, ztrue, photoz_mode


def plot_pdfs(sample, gals, show_ztrue=True, show_photoz_mode=False):
    """Plot a list of individual PDFs using qp plotting function for illustration.
    Ancillary function to be used by class Sample.

    Parameters
    ----------
    gals: `list`
        list of galaxies' indexes
    show_ztrue: `bool`, optional
        if True (default=True), show ztrue as dashed vertical line
    show_photoz_mode: `bool`, optional
        if True (default=False), show photoz_mode as dotted vertical line

    Returns
    -------
    colors: `list`
        list of HTML codes for colors used in the plot lines
    """
    colors = []
    peaks = []
    for i, gal in enumerate(gals):
        peaks.append(sample.pdf(sample.photoz_mode[gal])[gal])
        if i == 0:
            axes = sample.plot(key=gal, xlim=(0.0, 2.2), label=f"Galaxy {gal}")
        else:
            _ = sample.plot(key=gal, axes=axes, label=f"Galaxy {gal}")
        colors.append(axes.get_lines()[-1].get_color())
        if show_ztrue:
            axes.vlines(sample.ztrue[gal], ymin=0, ymax=100, colors=colors[-1], ls="--")
        if show_photoz_mode:
            axes.vlines(
                sample.photoz_mode[gal], ymin=0, ymax=100, colors=colors[-1], ls=":"
            )
    plt.ylim(0, np.max(peaks) * 1.05)
    axes.figure.legend()
    return colors


def plot_old_valid(photoz, ztrue, gals=None, colors=None, code=""):
    """Plot traditional Zphot X Zspec and N(z) plots for illustration
    Ancillary function to be used by class Sample.

    Parameters
    ----------
    gals: `list`, (optional)
        list of galaxies' indexes
    colors: `list`, (optional)
        list of HTML codes for colors used in the plot highlighted points
    """
    df = pd.DataFrame({"z$_{true}$": ztrue, "z$_{phot}$": photoz})
    fig = plt.figure(figsize=(10, 4), dpi=100)
    # fig.suptitle(name, fontsize=16)
    ax = plt.subplot(121)
    # sns.jointplot(data=df, x='z$_{true}$', y='z$_{phot}$', kind="kde")
    plt.plot(ztrue, photoz, "k,", label=code)
    leg = ax.legend(fancybox=True, handlelength=0, handletextpad=0, loc="upper left")
    if gals:
        if not colors:
            colors = ["r"] * len(gals)
        for i, gal in enumerate(gals):
            plt.plot(
                ztrue[gal], photoz[gal], "o", color=colors[i], label=f"Galaxy {gal}"
            )
    zmax = np.max(ztrue) * 1.01
    plt.xlim(0, zmax)
    plt.ylim(0, zmax)
    plt.xlabel("z$_{true}$")
    plt.ylabel("z$_{phot}$")
    plt.subplot(122)
    sns.kdeplot(ztrue, shade=True, label="z$_{true}$")
    sns.kdeplot(photoz, shade=True, label="z$_{phot}$")
    plt.xlim(0, zmax)
    plt.xlabel("z")
    plt.legend()
    plt.tight_layout()


def old_metrics(photoz, ztrue):
    point = EvaluatePointStats(photoz, ztrue)
    sigma_iqr = point.CalculateSigmaIQR()
    bias = point.CalculateBias()
    frac = point.CalculateOutlierRate()
    sigma_mad = point.CalculateSigmaMAD()
    return sigma_iqr, bias, frac, sigma_mad


def old_metrics_table(photoz, ztrue, name="", show_dc1=True):
    rows = ["|Metric |", "|:---|", "|scatter |", "|bias |", "|outlier rate |"]
    sigma_iqr, bias, frac, sigma_mad = old_metrics(photoz, ztrue)
    rows[0] += f"{name} |"
    rows[1] += "---:|"
    rows[2] += f"{sigma_iqr:11.4f} |"
    rows[3] += f"{bias:11.5f} |"
    rows[4] += f"{frac:11.3f} |"
    if show_dc1:
        rows[0] += "DC1 paper"
        rows[1] += "---:"
        rows[2] += f"  0.0154"
        rows[3] += f" -0.00027"
        rows[4] += f"  0.020"
    table = ("\n").join(rows)
    return Markdown(table)


def plot_pit_qq(
    pdfs,
    zgrid,
    ztrue,
    bins=None,
    title=None,
    code=None,
    show_pit=True,
    show_qq=True,
    pit_out_rate=None,
    savefig=False,
) -> str:
    """Quantile-quantile plot
    Ancillary function to be used by class Metrics.

    Parameters
    ----------
    pit: `PIT` object
        class from metrics.py
    bins: `int`, optional
        number of PIT bins
        if None, use the same number of quantiles (sample.n_quant)
    title: `str`, optional
        if None, use formatted sample's name (sample.name)
    label: `str`, optional
        if None, use formatted code's name (sample.code)
    show_pit: `bool`, optional
        include PIT histogram (default=True)
    show_qq: `bool`, optional
        include QQ plot (default=True)
    pit_out_rate: `ndarray`, optional
        print metric value on the plot panel (default=None)
    savefig: `bool`, optional
        save plot in .png file (default=False)
    """

    if bins is None:
        bins = 100
    if title is None:
        title = ""

    if code is None:
        code = ""
        label = ""
    else:
        label = code + "\n"

    if pit_out_rate is not None:
        try:
            label += "PIT$_{out}$: "
            label += f"{float(pit_out_rate):.4f}"
        except:
            print("Unsupported format for pit_out_rate.")

    plt.figure(figsize=[4, 5])
    gs = gridspec.GridSpec(2, 1, height_ratios=[3, 1])
    ax0 = plt.subplot(gs[0])
    sample = Sample(pdfs, zgrid, ztrue)

    if show_qq:
        ax0.plot(
            sample.qq[0], sample.qq[1], c="r", linestyle="-", linewidth=3, label=label
        )
        ax0.plot([0, 1], [0, 1], color="k", linestyle="--", linewidth=2)
        ax0.set_ylabel("Q$_{data}$", fontsize=18)
        plt.ylim(-0.001, 1.001)
    plt.xlim(-0.001, 1.001)
    plt.title(title)
    if show_pit:
        fzdata = Ensemble(interp, data=dict(xvals=zgrid, yvals=pdfs))
        pitobj = PIT(fzdata, ztrue)
        pit_vals = np.array(pitobj.pit_samps)
        pit_out_rate = pitobj.evaluate_PIT_outlier_rate()

        try:
            y_uni = float(len(pit_vals)) / float(bins)
        except:
            y_uni = float(len(pit_vals)) / float(len(bins))
        if not show_qq:
            ax0.hist(pit_vals, bins=bins, alpha=0.7, label=label)
            ax0.set_ylabel("Number")
            ax0.hlines(y_uni, xmin=0, xmax=1, color="k")
            plt.ylim(
                0,
            )  # -0.001, 1.001)
        else:
            ax1 = ax0.twinx()
            ax1.hist(pit_vals, bins=bins, alpha=0.7)
            ax1.set_ylabel("Number")
            ax1.hlines(y_uni, xmin=0, xmax=1, color="k")
    leg = ax0.legend(handlelength=0, handletextpad=0, fancybox=True)
    if show_qq:
        ax2 = plt.subplot(gs[1])
        ax2.plot(
            sample.qq[0],
            (sample.qq[1] - sample.qq[0]),
            c="r",
            linestyle="-",
            linewidth=3,
        )
        plt.ylabel("$\Delta$Q", fontsize=18)
        ax2.plot([0, 1], [0, 0], color="k", linestyle="--", linewidth=2)
        plt.xlim(-0.001, 1.001)
        plt.ylim(
            np.min([-0.12, np.min(sample.qq[1] - sample.qq[0]) * 1.05]),
            np.max([0.12, np.max(sample.qq[1] - sample.qq[0]) * 1.05]),
        )
    if show_pit:
        if show_qq:
            plt.xlabel("Q$_{theory}$ / PIT Value", fontsize=18)
        else:
            plt.xlabel("PIT Value", fontsize=18)
    else:
        if show_qq:
            plt.xlabel("Q$_{theory}$", fontsize=18)
    if savefig:
        fig_filename = str("plot_pit_qq_" + f"{(code).replace(' ', '_')}.png")
        plt.savefig(fig_filename)
    else:
        fig_filename = None

    return fig_filename


def ks_plot(pitobj, n_quant=100):
    """KS test illustration.
    Ancillary function to be used by class KS."""
    pits = np.array(pitobj.pit_samps)
    stat_and_pval = pitobj.evaluate_PIT_KS()
    xvals = np.linspace(0.0, 1.0, n_quant)
    yvals = np.array([np.histogram(pits, bins=len(xvals))[0]])
    pit_cdf = Ensemble(interp, data=dict(xvals=xvals, yvals=yvals)).cdf(xvals)
    uniform_yvals = np.array([np.full(n_quant, 1.0 / float(n_quant))])
    uniform_cdf = Ensemble(interp, data=dict(xvals=xvals, yvals=uniform_yvals)).cdf(
        xvals
    )

    plt.figure(figsize=[4, 4])
    plt.plot(xvals, uniform_cdf, "r-", label="uniform")
    plt.plot(xvals, pit_cdf, "b-", label="sample PIT")
    bin_stat = np.argmax(np.abs(pit_cdf - uniform_cdf))

    plt.vlines(
        x=xvals[bin_stat],
        ymin=np.min([pit_cdf[bin_stat], uniform_cdf[bin_stat]]),
        ymax=np.max([pit_cdf[bin_stat], uniform_cdf[bin_stat]]),
        colors="k",
    )
    plt.plot(xvals[bin_stat], pit_cdf[bin_stat], "k.")
    plt.plot(xvals[bin_stat], uniform_cdf[bin_stat], "k.")
    ymean = (pit_cdf[bin_stat] + uniform_cdf[bin_stat]) / 2.0
    plt.text(xvals[bin_stat] + 0.05, ymean, "max", fontsize=16)
    plt.xlabel("PIT value")
    plt.ylabel("CDF(PIT)")
    xtext = 0.63
    ytext = 0.03
    plt.text(xtext, ytext, f"KS={stat_and_pval.statistic:.4f}", fontsize=16)
    plt.xlim(0, 1)
    plt.ylim(0, 1)
    plt.legend()
    plt.tight_layout()


class EvaluatePointStats(object):
    """Copied from PZDC1paper repo. Adapted to remove the cut based on magnitude."""

    def __init__(self, pzvec, szvec):
        """An object that takes in the vectors of the point photo-z
        the spec-z, and the i-band magnitudes for calculating the
        point statistics
        Parameters:
        pzvec: Numpy 1d array of the point photo-z values
        szvec: Numpy 1d array of the spec-z values
        magvec: Numpy 1d array of the i-band magnitudes
        imagcut: float: i-band magnitude cut for the sample
        Calculates:
        ez_all: (pz-sz)/(1+sz), the quantity will be useful for calculating statistics
        ez_magcut: ez sample trimmed with imagcut
        """
        self.pzs = pzvec
        self.szs = szvec
        ez = (pzvec - szvec) / (1.0 + szvec)
        self.ez_all = ez

    def CalculateSigmaIQR(self):
        """Calculate the width of the e_z distribution
        using the Interquartile range
        Parameters:
        imagcut: float: i-band magnitude cut for the sample
        Returns:
        sigma_IQR_all float: width of ez distribution for full sample
        sigma_IQR_magcut float: width of ez distribution for magcut sample
        """
        x75, x25 = np.percentile(self.ez_all, [75.0, 25.0])
        iqr_all = x75 - x25
        sigma_iqr_all = iqr_all / 1.349
        self.sigma_iqr_all = sigma_iqr_all

        return sigma_iqr_all

    def CalculateBias(self):
        """calculates the bias of the ez and ez_magcut samples.  In
        keeping with the Science Book, this is just the median of the
        ez values
        Returns:
        bias_all: median of the full ez sample
        bias_magcut: median of the magcut ez sample
        """
        bias_all = np.median(self.ez_all)
        return bias_all

    def CalculateOutlierRate(self):
        """Calculates the catastrophic outlier rate, defined in the
        Science Book as the number of galaxies with ez larger than
        max(0.06,3sigma).  This keeps the fraction reasonable when
        sigma is very small.
        Returns:
        frac_all: fraction of catastrophic outliers for full sample
        frac_magcut: fraction of catastrophic outliers for magcut
        sample
        """
        num_all = len(self.ez_all)
        threesig_all = 3.0 * self.sigma_iqr_all
        cutcriterion_all = np.maximum(0.06, threesig_all)
        mask_all = self.ez_all > np.fabs(cutcriterion_all)
        outlier_all = np.sum(mask_all)
        frac_all = float(outlier_all) / float(num_all)
        return frac_all

    def CalculateSigmaMAD(self):
        """Function to calculate median absolute deviation and sigma
        based on MAD (just scaled up by 1.4826) for the full and
        magnitude trimmed samples of ez values
        Returns:
        sigma_mad_all: sigma_MAD for full sample
        sigma_mad_cut: sigma_MAD for the magnitude cut sample
        """
        tmpmed_all = np.median(self.ez_all)
        tmpx_all = np.fabs(self.ez_all - tmpmed_all)
        mad_all = np.median(tmpx_all)
        sigma_mad_all = mad_all * 1.4826
        return sigma_mad_all


class DC1:

    def __init__(self):
        # Reference values:
        self.codes = (
            "ANNz2",
            "BPZ",
            "CMNN",
            "Delight",
            "EAZY",
            "FlexZBoost",
            "GPz",
            "LePhare",
            "METAPhoR",
            "SkyNet",
            "TPZ",
        )
        self.metrics = ("PIT out rate", "CDE loss", "KS", "CvM", "AD")
        self.pit_out_rate = [
            0.0265,
            0.0192,
            0.0034,
            0.0006,
            0.0154,
            0.0202,
            0.0058,
            0.0486,
            0.0229,
            0.0001,
            0.0130,
        ]
        self.cde_loss = [
            -6.88,
            -7.82,
            -10.43,
            -8.33,
            -7.07,
            -10.60,
            -9.93,
            -1.66,
            -6.28,
            -7.89,
            -9.55,
        ]
        self.ks = [
            0.0200,
            0.0388,
            0.0795,
            0.08763,
            0.0723,
            0.0240,
            0.0241,
            0.0663,
            0.0438,
            0.0747,
            0.1138,
            0.0047,
        ]
        self.cvm = [
            52.25,
            280.79,
            1011.11,
            1075.17,
            1105.58,
            68.83,
            66.80,
            473.05,
            298.56,
            763.00,
            1.16,
        ]
        self.ad = [
            759.2,
            1557.5,
            6307.5,
            6167.5,
            4418.6,
            478.8,
            670.9,
            383.8,
            715.5,
            4216.4,
            10565.7,
            7.7,
        ]

    @property
    def results(self):
        results = {
            "PIT out rate": dict(
                [(code, value) for code, value in zip(self.codes, self.pit_out_rate)]
            ),
            "CDE loss": dict(
                [(code, value) for code, value in zip(self.codes, self.cde_loss)]
            ),
            "KS": dict([(code, value) for code, value in zip(self.codes, self.ks)]),
            "CvM": dict([(code, value) for code, value in zip(self.codes, self.cvm)]),
            "AD": dict([(code, value) for code, value in zip(self.codes, self.ad)]),
        }
        return results

    @property
    def table(self):
        table = "Code | PIT out rate | KS | CvM | AD | CDE loss  \n ---|---:|---:|---:|---:|---: \n "
        for code in self.codes:
            table += f"{code} | {self.results['PIT out rate'][code]:11.4f} "
            table += f" | {self.results['KS'][code]:11.4f}"
            table += f" | {self.results['CvM'][code]:11.4f}"
            table += f" | {self.results['AD'][code]:11.4f}"
            table += f" | {self.results['CDE loss'][code]:11.4f}\n"
        return Markdown(table)


class Sample(Ensemble):
    """Expand qp.Ensemble to append true redshifts
    array, metadata, and specific plots."""

    def __init__(
        self, pdfs, zgrid, ztrue, photoz_mode=None, code="", name="", n_quant=100
    ):
        """Class constructor

        Parameters
        ----------
        pdfs: `ndarray`
            photo-z PDFs array, shape=(Ngals, Nbins)
        zgrid: `ndarray`
            PDF bins centers, shape=(Nbins,)
        ztrue: `ndarray`
            true redshifts, shape=(Ngals,)
        photoz_mode: `ndarray`
            photo-z (PDF mode), shape=(Ngals,)
        code: `str`, (optional)
            algorithm name (for plot legends)
        name: `str`, (optional)
            sample name (for plot legends)
        """

        super().__init__(interp, data=dict(xvals=zgrid, yvals=pdfs))
        self._pdfs = pdfs
        self._zgrid = zgrid
        self._ztrue = ztrue
        self._photoz_mode = photoz_mode
        self._code = code
        self._name = name
        self._n_quant = n_quant
        self._pit = None
        self._qq = None

    @property
    def code(self):
        """Photo-z code/algorithm name"""
        return self._code

    @property
    def name(self):
        """Sample name"""
        return self._name

    @property
    def ztrue(self):
        """True redshifts array"""
        return self._ztrue

    @property
    def zgrid(self):
        """Redshift grid (binning)"""
        return self._zgrid

    @property
    def photoz_mode(self):
        """Photo-z (mode) array"""
        return self._photoz_mode

    @property
    def n_quant(self):
        return self._n_quant

    @property
    def pit(self):
        if self._pit is None:
            pit_array = np.array([self[i].cdf(self.ztrue[i]) for i in range(len(self))])
            self._pit = pit_array
        return self._pit

    @property
    def qq(self, n_quant=100):
        q_theory = np.linspace(0.0, 1.0, n_quant)
        q_data = np.quantile(self.pit, q_theory)
        self._qq = (q_theory, q_data)
        return self._qq

    def __len__(self):
        if len(self._ztrue) != len(self._pdfs):
            raise ValueError("Number of pdfs and true redshifts do not match!!!")
        return len(self._ztrue)

    def __str__(self):
        code_str = f"Algorithm: {self._code}"
        name_str = f"Sample: {self._name}"
        line_str = "-" * (max(len(code_str), len(name_str)))
        text = str(
            line_str
            + "\n"
            + name_str
            + "\n"
            + code_str
            + "\n"
            + line_str
            + "\n"
            + f"{len(self)} PDFs with {len(self.zgrid)} probabilities each \n"
            + f"qp representation: {self.gen_class.name} \n"
            + f"z grid: {len(self.zgrid)} z values from {np.min(self.zgrid)} to {np.max(self.zgrid)} inclusive"
        )
        return text

    def plot_pdfs(self, gals, show_ztrue=True, show_photoz_mode=False):
        colors = plot_pdfs(
            self, gals, show_ztrue=show_ztrue, show_photoz_mode=show_photoz_mode
        )
        return colors

    def plot_old_valid(self, gals=None, colors=None):
        old_metrics_table = plot_old_valid(self, gals=gals, colors=colors)
        return old_metrics_table

    def plot_pit_qq(
        self,
        bins=None,
        label=None,
        title=None,
        show_pit=True,
        show_qq=True,
        show_pit_out_rate=True,
        savefig=False,
    ):
        """Make plot PIT-QQ as Figure 2 from Schmidt et al. 2020."""
        fig_filename = plot_pit_qq(
            self,
            bins=bins,
            label=label,
            title=title,
            show_pit=show_pit,
            show_qq=show_qq,
            show_pit_out_rate=show_pit_out_rate,
            savefig=savefig,
        )
        return fig_filename
