import time
from typing import Union
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from scipy.stats import beta
from janusx.gtools.cleaner import chrom_sort_key


def ppoints(n: int) -> np.ndarray:
    """
    Approximate expected quantiles of the Uniform(0,1) distribution.

    Equivalent to R's ppoints(n) with a simple n/(n+1) scheme:
        q_i = i / (n + 1), i = 1..n
    """
    return np.arange(1, n + 1) / (n + 1)


class GWASPLOT:
    """
    Lightweight helper for GWAS visualization (Manhattan & QQ plots).

    Parameters
    ----------
    df : pd.DataFrame
        GWAS result table. Must contain chromosome, position and p-value columns.
    chr : str, default '#CHROM'
        Column name for chromosome.
    pos : str, default 'POS'
        Column name for base-pair position.
    pvalue : str, default 'p'
        Column name for p-values.
    interval_rate : float, default 0.1
        Spacing ratio between chromosomes on the x-axis
        (absolute spacing = interval_rate * max_position).
    compression : bool, default True
        If True, down-sample SNPs to roughly ~100,000 points for plotting:
          - keep all highly significant points (p <= 10000 / n),
          - for the rest, partition into chunks and keep only the most significant
            SNP per chunk.
    chr_order : list or None, default None
        Optional chromosome order. If None, chromosome labels are ordered
        automatically by natural chromosome order (1..N, X/Y/M/MT, then others).

    Notes
    -----
    - Internally, chromosomes are mapped to integer codes 1..K for plotting.
    - The main DataFrame is stored as self.df with a MultiIndex (chr, pos).
    - A global index list self.minidx encodes which SNPs are retained when
      compression is enabled (used by both Manhattan and QQ plots).
    """

    def __init__(
        self,
        df: pd.DataFrame,
        chr: str = "chrom",
        pos: str = "pos",
        pvalue: str = "pwald",
        interval_rate: float = 0.1,
        compression: bool = True,
        chr_order: Union[list[object], None] = None,
    ) -> None:
        self.t_start = time.time()
        # Ensure positional indices are contiguous for downstream iloc usage.
        df = df.reset_index(drop=True)

        # ---- (1) Optional down-sampling for plotting ----
        if compression:
            n_snp = df.shape[0]
            # target: ~100,000 SNPs for display
            chunk_size = n_snp // 100_000 if n_snp >= 100_000 else 1

            # p-value threshold: small p => highly significant
            # we keep all SNPs with p <= p_thresh without compression
            p_thresh = 10_000 / n_snp

            # Among the "less significant" SNPs (p > p_thresh), we down-sample:
            #   - sort by p descending (largest p first, least significant),
            #   - take an integer multiple of chunk_size rows,
            #   - then in each chunk keep the smallest p (most significant in that chunk).
            mask_large_p = df[pvalue] > p_thresh
            n_large = int(mask_large_p.sum())
            n_chunks = n_large // chunk_size

            if n_chunks > 0:
                # indices of "large p" SNPs, sorted by p descending
                idx_large_sorted = (
                    df.loc[mask_large_p]
                    .sort_values(pvalue, ascending=False)
                    .iloc[: n_chunks * chunk_size, :]
                    .index
                )
                df_large = df.loc[idx_large_sorted].sort_index()

                # convert to array to find minima in each chunk of size chunk_size
                pvals_large = df_large[pvalue].values.reshape(n_chunks, chunk_size)
                local_min_pos = np.argmin(pvals_large, axis=1)  # position within each chunk

                # map chunk index + local min position back to global index
                keep_idx_large = df_large.iloc[
                    [i * chunk_size + j for i, j in enumerate(local_min_pos)]
                ].index.tolist()
            else:
                keep_idx_large = []

            # always keep all highly significant SNPs
            keep_idx_small = df.loc[~mask_large_p].index.tolist()

            # final kept SNP indices for plotting
            self.minidx = keep_idx_large + keep_idx_small
        else:
            self.minidx = df.index.to_list()

        # ---- (2) Core columns copy to avoid side effects ----
        df = df[[chr, pos, pvalue]].copy()

        # ---- (3) Map chromosome labels to integers 1..K ----
        observed_labels = df[chr].drop_duplicates().tolist()

        if chr_order is None:
            self.chr_labels = sorted(observed_labels, key=chrom_sort_key)
        else:
            observed_set = set(observed_labels)
            ordered_present = [label for label in chr_order if label in observed_set]
            ordered_set = set(ordered_present)
            missing = [label for label in observed_labels if label not in ordered_set]
            self.chr_labels = ordered_present + sorted(missing, key=chrom_sort_key)

        chr_map = dict(zip(self.chr_labels, range(1, 1 + len(self.chr_labels))))
        df[chr] = df[chr].map(chr_map).astype(int)

        # Sort by chromosome and position
        df = df.sort_values(by=[chr, pos])
        self.chr_ids = df[chr].unique()

        # ---- (4) Build cumulative genomic x-coordinate ----
        # base spacing between chromosomes
        self.interval = int(interval_rate * df[pos].max())
        df["x"] = df[pos]

        if len(self.chr_ids) > 1:
            for cid in self.chr_ids:
                if cid > 1:
                    prev_max = df.loc[df[chr] == cid - 1, "x"].max()
                    df.loc[df[chr] == cid, "x"] += prev_max

        # add extra spacing between chromosomes
        df["x"] = (df[chr] - 1) * self.interval + df["x"]

        # rename standardized y/z columns
        df["y"] = df[pvalue]        # raw p-values
        df["z"] = df[chr]           # integer chromosome ID

        # chromosome tick positions on x-axis
        self.ticks_loc = df.groupby("z")["x"].mean()

        # store as MultiIndex (chr, pos) for easier masking later
        self.df = df.set_index([chr, pos])

    # ------------------------------------------------------------------
    # Manhattan plot
    # ------------------------------------------------------------------
    def manhattan(
        self,
        threshold: Union[float,None] = None,
        color_set: Union[list[str],None] = None,
        ax: Union[plt.Axes,None] = None,
        ignore: Union[list[str],None] = None,
        min_logp: Union[float, None] = 0.5,
        max_logp: Union[float, None] = None,
        y_min: Union[float, None] = None,
        **kwargs,
    ) -> plt.Axes:
        """
        Draw a Manhattan plot.

        Parameters
        ----------
        threshold : float or None
            Genome-wide significance threshold on -log10(p).
            If provided, significant loci are highlighted in red.
        color_set : list[str] or None
            Colors for alternating chromosomes, e.g. ['black','grey'].
        ax : matplotlib.axes.Axes or None
            Existing Axes to draw on. If None, a new figure is created.
        ignore : list
            List of index keys (MultiIndex entries) to ignore in highlighting.
        **kwargs :
            Additional keyword arguments passed to ax.scatter.

        Returns
        -------
        ax : matplotlib.axes.Axes
        """
        if color_set is None or len(color_set) == 0:
            color_set = ["black", "grey"]
        if ignore is None:
            ignore = []

        # subset to plotting SNPs
        df = self.df.iloc[self.minidx, -3:].copy()
        df["y"] = -np.log10(df["y"])
        if min_logp is not None:
            df = df[df["y"] >= float(min_logp)]
        if max_logp is not None:
            df = df[df["y"] <= float(max_logp)]

        if ax is None:
            fig = plt.figure(figsize=[12, 6], dpi=300)
            gs = GridSpec(12, 1, figure=fig)
            ax = fig.add_subplot(gs[0:12, 0])

        if df.shape[0] == 0:
            ax.text(0.5, 0.5, "No SNPs", ha="center", va="center", transform=ax.transAxes)
            ax.set_xticks(self.ticks_loc, self.chr_labels)
            ax.set_xlabel("Chromosome")
            ax.set_ylabel("-log10(p-value)")
            return ax

        # color per chromosome (alternating colors)
        chr_color_map = dict(
            zip(
                self.chr_ids,
                [color_set[i % len(color_set)] for i in range(len(self.chr_ids))],
            )
        )
        plot_kwargs = dict(kwargs)
        plot_kwargs.setdefault("alpha", 0.78)
        plot_kwargs.setdefault("edgecolors", "none")
        plot_kwargs.setdefault("linewidths", 0.0)

        mask_ignore = ~df.index.isin(ignore)
        ax.scatter(
            df.loc[mask_ignore, "x"],
            df.loc[mask_ignore, "y"],
            color=df.loc[mask_ignore, "z"].map(chr_color_map),
            **plot_kwargs,
        )

        # Highlight SNPs above genome-wide threshold
        if threshold is not None and df["y"].max() >= threshold:
            df_sig = df[df["y"] >= threshold]
            sig_mask = ~df_sig.index.isin(ignore)
            sig_kwargs = dict(kwargs)
            sig_kwargs.setdefault("alpha", 0.9)
            sig_kwargs.setdefault("edgecolors", "none")
            sig_kwargs.setdefault("linewidths", 0.0)
            ax.scatter(
                df_sig.loc[sig_mask, "x"],
                df_sig.loc[sig_mask, "y"],
                color="red",
                **sig_kwargs,
            )
            ax.axhline(
                y=threshold,
                color="grey",
                linewidth=1,
                linestyle="--",
            )

        # axis cosmetics
        chr_bounds = (
            self.df.groupby("z")["x"]
            .agg(["min", "max"])
            .sort_index()
        )
        if chr_bounds.shape[0] > 1:
            mins = chr_bounds["min"].to_numpy(dtype=float)
            maxs = chr_bounds["max"].to_numpy(dtype=float)
            separators = (maxs[:-1] + mins[1:]) / 2.0
            for xsep in separators:
                if np.isfinite(xsep):
                    ax.axvline(
                        float(xsep),
                        ymin=0.0,
                        ymax=1.0 / 3.0,
                        linestyle="--",
                        color="lightgrey",
                        linewidth=0.6,
                        alpha=0.8,
                        zorder=0,
                    )

        ax.set_xticks(self.ticks_loc, self.chr_labels)
        xmin = float(df["x"].min())
        xmax = float(df["x"].max())
        if xmax > xmin:
            ax.set_xlim(xmin, xmax)
        else:
            eps = max(1e-9, abs(xmin) * 1e-9)
            ax.set_xlim(xmin - eps, xmax + eps)
        ax.margins(x=0.0)
        ymax = df["y"].max()
        base_ymin = float(min_logp) if min_logp is not None else 0.0
        if y_min is not None:
            base_ymin = float(y_min)
        top = ymax + 0.1 * ymax if ymax > 0 else (base_ymin + 1.0)
        if top <= base_ymin:
            top = base_ymin + 1.0
        ax.set_ylim([base_ymin, top])
        ax.set_xlabel("Chromosome")
        ax.set_ylabel("-log10(p-value)")
        return ax

    # ------------------------------------------------------------------
    # QQ plot
    # ------------------------------------------------------------------
    def qq(
        self,
        ax: plt.Axes = None,
        ci: int = 95,
        color_set: list[str] = None,
        scatter_size: float = 8.0,
        line_color: str = "black",
    ) -> plt.Axes:
        """
        Draw a QQ plot of observed vs expected -log10(p) with a beta-based
        confidence band.

        This implementation matches the original logic you used:
        - Use all SNPs (no extra down-sampling beyond self.minidx indices)
        - Map back to the original index before combining observed/expected
        - Use self.minidx for consistency with the Manhattan down-sampling.

        Parameters
        ----------
        ax : matplotlib.axes.Axes, optional
            Axes to draw on. If None, a new figure is created.
        ci : int, default 95
            Confidence interval (in percent) for the null band.
        color_set : list[str], optional
            QQ colors. If >=2 colors are given:
            - color_set[0]: scatter points
            - color_set[1]: confidence band
            If only one color is provided, it is used for both.
        scatter_size : float, default 8.0
            Marker size for QQ scatter points.
        line_color : str, default "black"
            Color of the ideal diagonal line y=x.

        Returns
        -------
        ax : matplotlib.axes.Axes
        """
        if color_set is None or len(color_set) == 0:
            color_set = ["black", "grey"]
        if len(color_set) >= 2:
            point_color = color_set[0]
            band_color = color_set[1]
        else:
            point_color = color_set[0]
            band_color = color_set[0]

        # Use a fresh copy so we don't accidentally modify self.df
        df = self.df.copy()

        if ax is None:
            fig = plt.figure(figsize=[12, 6], dpi=600)
            gs = GridSpec(12, 1, figure=fig)
            ax = fig.add_subplot(gs[0:12, 0])

        # Raw p-values (full set, not compressed)
        p = df["y"].fillna(1)
        n = len(p)
        if n == 0:
            raise ValueError("No p-values found for QQ plot.")

        # Step 1: sort observed p-values and assign expected uniform quantiles
        #   o_e: index = SNP index (from df),
        #        columns 'o' (observed p), 'e' (expected p under Uniform)
        o_e = p.sort_values().to_frame()
        o_e["e"] = ppoints(n)
        o_e.columns = ["o", "e"]

        # Step 2: reorder back by the original index order,
        #         then transform both observed and expected to -log10 scale
        #         so that indices align with df (this is key for self.minidx).
        o_e = -np.log10(o_e.sort_index())

        # Theoretical expected -log10(p) values sorted ascending
        # (used for drawing confidence band and diagonal reference line)
        e_theoretical = o_e["e"].sort_values().values

        # Step 3: beta-based confidence band under the null
        # For each theoretical point with -log10(p) = e, the underlying p is 10^-e,
        # so expected rank is ~p * n, and we use Beta(r, n-r+1) to get a band.
        xi = np.ceil(10 ** (-e_theoretical) * n)  # ranks
        # Make sure indices are valid for Beta; your original code did not clip,
        # but xi is always in [1, n] here, so we keep behavior the same.
        lower = -np.log10(beta.ppf(1 - ci / 100, xi, n - xi + 1))
        upper = -np.log10(beta.ppf(ci / 100, xi, n - xi + 1))

        # Draw confidence interval
        ax.fill_between(
            e_theoretical,
            lower,
            upper,
            color=band_color,
            alpha=0.3,
            rasterized=True,
        )

        # Diagonal reference line y = x
        max_lim = np.min(o_e.max(axis=0))
        ax.plot([0, max_lim], [0, max_lim], lw=1, color=line_color)

        # Scatter observed vs expected for the same subset used in Manhattan plot
        #   o_e.iloc[self.minidx, 1] -> expected -log10(p)
        #   o_e.iloc[self.minidx, 0] -> observed -log10(p)
        ax.scatter(
            o_e.iloc[self.minidx, 1],
            o_e.iloc[self.minidx, 0],
            s=scatter_size,
            alpha=0.75,
            rasterized=True,
            color=point_color,
            edgecolors="none",
            linewidths=0.0,
        )

        ax.set_xlabel("Expected -log10(p-value)")
        ax.set_ylabel("Observed -log10(p-value)")
        return ax
