plots.py

import matplotlib.pyplot as plt
from typing import Dict, Any, List, Optional, Tuple
import pandas as pd
import numpy as np
from matplotlib.gridspec import GridSpec
import matplotlib.pyplot as plt
import glob

def _is_nonempty_str(x) -> bool:
    if x is None or (isinstance(x, float) and np.isnan(x)):
        return False
    s = str(x).strip()
    return s != "" and s.lower() != "nan"


def _to_bool_series(s: pd.Series) -> pd.Series:
    """
    Robust boolean conversion for columns that might be:
    True/False, 'True'/'False', 1/0, '1'/'0', NaN, ''.
    Missing -> False (you can change this if you prefer).
    """
    if s is None:
        return pd.Series([], dtype=bool)
    if s.dtype == bool:
        return s.fillna(False)

    def conv(v):
        if v is None or (isinstance(v, float) and np.isnan(v)):
            return False
        if isinstance(v, bool):
            return v
        sv = str(v).strip().lower()
        if sv in {"true", "t", "1", "yes", "y"}:
            return True
        if sv in {"false", "f", "0", "no", "n"}:
            return False
        # fallback: treat unknown as False
        return False

    return s.map(conv).astype(bool)


def _add_flags(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df["has_medical_gene"] = df["medical_relevant_gene"].map(_is_nonempty_str)
    #df["OMIM_bool"] = _to_bool_series(df.get("OMIM", pd.Series([False]*len(df))))
    df["OMIM_bool"] = [True if d == "True" else False for d in df.get("OMIM")]
    
    
    df["mena_bool"] = _to_bool_series(df.get("mena", pd.Series([False]*len(df))))

    #df["in_mrg_bool"] = _to_bool_series(df.get("in_mrg", pd.Series([False]*len(df))))
    #df["in_mrg_bool"] = _to_bool_series(df.get("in_mrg", pd.Series([False]*len(df))))
    df["in_mrg_bool"] = [True if d == "TRUE" else False for d in df.get("in_mrg")]
    return df
    
    

def plot_summary_panels(df: pd.DataFrame):

    df = _add_flags(df)
    figs = []

    def pie_with_counts(ax, values, labels, title=None):
        total = int(sum(values))

        def _autopct(pct):
            count = int(round(pct * total / 100.0))
            return f"{count}" if count > 0 else ""

        ax.pie(values, labels=labels, autopct=_autopct, startangle=90)
        ax.axis("equal")
        if title:
            ax.set_title(title)

    counts1 = df["has_medical_gene"].value_counts().reindex([True, False], fill_value=0)
    labels1 = ["overlapping medical relevant genes", "not overlapping medical relevant genes"]
    values1 = [int(counts1[True]), int(counts1[False])]

    fig1, (ax1b, ax1p) = plt.subplots(1, 2, figsize=(12, 4))
    ax1b.bar(labels1, values1)
    ax1b.set_ylabel("Number of SV ids")
    ax1b.set_title("SVs overlapping with medical relevant genes")
    ax1b.tick_params(axis="x", rotation=15)
    pie_with_counts(ax1p, values1, labels1)
    figs.append(fig1)


    sub = df[df["has_medical_gene"] == True]
    counts2 = sub["in_mrg_bool"].value_counts().reindex([True, False], fill_value=0)
    labels2 = ["medically difficult", "not medically difficult"]
    values2 = [int(counts2[True]), int(counts2[False])]

    fig2, (ax2b, ax2p) = plt.subplots(1, 2, figsize=(12, 4))
    ax2b.bar(labels2, values2)
    ax2b.set_ylabel("Number of SV ids")
    ax2b.set_title("Subset: has_medical_gene (medically difficult)")
    ax2b.tick_params(axis="x", rotation=15)
    pie_with_counts(ax2p, values2, labels2)
    figs.append(fig2)
    

    counts1 = df["OMIM_bool"].value_counts().reindex([True, False], fill_value=0)
    labels1 = ["overlapping OMIMs", "not overlapping OMIMs"]
    values1 = [int(counts1[True]), int(counts1[False])]

    fig3, (ax3b, ax3p) = plt.subplots(1, 2, figsize=(12, 4))
    ax3b.bar(labels1, values1)
    ax3b.set_ylabel("Number of SV ids")
    ax3b.set_title("SVs overlapping with OMIM exons")
    ax3b.tick_params(axis="x", rotation=15)
    pie_with_counts(ax3p, values1, labels1)
    figs.append(fig3)
    """
    # ---------- 3) OMIM True/False cross-referenced with medical relevance ----------
    ct = pd.crosstab(df["OMIM_bool"], df["has_medical_gene"])
    ct = ct.reindex(index=[True, False], columns=[True, False], fill_value=0)

    # Bar plot: grouped bars (OMIM True/False on x, medical yes/no as groups)
    x = np.arange(len(ct.index))
    width = 0.35
    omim_labels = ["OMIM=True", "OMIM=False"]
    med_present = ct[True].values.astype(int)
    med_absent = ct[False].values.astype(int)

    fig3, (ax3b, ax3p) = plt.subplots(1, 2, figsize=(12, 4))
    ax3b.bar(x - width / 2, med_present, width, label="medical_relevant_gene present")
    ax3b.bar(x + width / 2, med_absent, width, label="medical_relevant_gene absent")
    ax3b.set_xticks(x)
    ax3b.set_xticklabels(omim_labels)
    ax3b.set_ylabel("Number of SV ids")
    ax3b.set_title("OMIM vs medical relevance (2×2)")
    ax3b.legend()

    # Pie plot: total OMIM True vs False (overall OMIM distribution)
    omim_counts = df["OMIM_bool"].value_counts().reindex([True, False], fill_value=0)
    labels3 = ["OMIM exon", "no OMIM exon"]
    values3 = [int(omim_counts[True]), int(omim_counts[False])]
    pie_with_counts(ax3p, values3, labels3)
    figs.append(fig3)

    # ---------- 4) Effect of excluding mena=True ----------
    before = df["has_medical_gene"].value_counts().reindex([True, False], fill_value=0)
    after_df = df[df["mena_bool"] == False]
    after = after_df["has_medical_gene"].value_counts().reindex([True, False], fill_value=0)

    labels4 = ["medical gene present", "medical gene absent"]
    before_vals = [int(before[True]), int(before[False])]
    after_vals = [int(after[True]), int(after[False])]

    x = np.arange(len(labels4))
    width = 0.35

    fig4, (ax4b, ax4p) = plt.subplots(1, 2, figsize=(12, 4))
    ax4b.bar(x - width / 2, before_vals, width, label="All SVs (incl. MENA)")
    ax4b.bar(x + width / 2, after_vals, width, label="Excluding mena=True")
    ax4b.set_xticks(x)
    ax4b.set_xticklabels(labels4)
    ax4b.set_ylabel("Number of SV ids")
    ax4b.set_title("Impact of excluding MENA on medical relevance counts")
    ax4b.legend()

    # Pie plot: show AFTER distribution (excluding MENA) — same categories as the bar's x-axis
    pie_with_counts(ax4p, after_vals, labels4, title="Excluding MENA (counts)")
    figs.append(fig4)
    """
    return figs
    
    
def plot_mena_delta_summary(df: pd.DataFrame, *, count_unique_ids: bool = True):

    df = _add_flags(df)

    def _count(mask: pd.Series, frame: pd.DataFrame) -> int:
        if count_unique_ids and "id" in frame.columns:
            return int(frame.loc[mask, "id"].nunique())
        return int(mask.sum())

    all_df = df
    no_mena_df = df[df["mena_bool"] == False]


    def masks(frame: pd.DataFrame):
        med = frame["has_medical_gene"] == True
        omim = frame["OMIM_bool"] == True
        both = med & omim
        med_only = med & (~omim)
        omim_only = (~med) & omim
        neither = (~med) & (~omim)
        return {
            "Medical": med,
            "OMIM": omim,
            "Medical ∩ OMIM": both,
            "Medical only": med_only,
            "OMIM only": omim_only,
            "Neither": neither,
        }

    m_all = masks(all_df)
    m_nom = masks(no_mena_df)

    cats = list(m_all.keys())
    before = [ _count(m_all[c], all_df) for c in cats ]
    after  = [ _count(m_nom[c], no_mena_df) for c in cats ]
    delta  = [ a - b for a, b in zip(after, before) ]  # negative if reduced

    x = np.arange(len(cats))
    width = 0.38

    fig, ax = plt.subplots(figsize=(12, 4.5))
    ax.bar(x - width/2, before, width, label="All (incl. MENA)")
    ax.bar(x + width/2, after,  width, label="Excluding mena=True")

    # Annotate deltas above the "after" bars
    for i, (b, a, d) in enumerate(zip(before, after, delta)):
        # Place label above the taller of the two bars
        y = max(b, a)
        ax.text(i, y + max(1, 0.01*y), f"Δ {d:+d}", ha="center", va="bottom")

    ax.set_xticks(x)
    ax.set_xticklabels(cats, rotation=20, ha="right")
    ax.set_ylabel("Count" + (" (unique ids)" if count_unique_ids else " (rows)"))
    ax.set_title("Effect of excluding MENA on key categories")
    ax.legend()

    fig.tight_layout()
    return fig
    

def plot_omim_medical_heatmaps(df: pd.DataFrame, *, count_unique_ids: bool = True):

    df = _add_flags(df)

    def _cell_counts(frame: pd.DataFrame) -> np.ndarray:

        mat = np.zeros((2, 2), dtype=int)

        for i, omim_val in enumerate([True, False]):
            for j, med_val in enumerate([True, False]):
                mask = (frame["OMIM_bool"] == omim_val) & (frame["has_medical_gene"] == med_val)
                if count_unique_ids and "id" in frame.columns:
                    mat[i, j] = int(frame.loc[mask, "id"].nunique())
                else:
                    mat[i, j] = int(mask.sum())
        return mat

    all_df = df
    no_mena_df = df[df["mena_bool"] == False]

    mat_all = _cell_counts(all_df)
    mat_nom = _cell_counts(no_mena_df)

    fig, axes = plt.subplots(1, 2, figsize=(10.5, 4.5))

    def _plot_heat(ax, mat: np.ndarray, title: str):
        im = ax.imshow(mat)
        ax.set_title(title)

        ax.set_xticks([0, 1])
        ax.set_xticklabels(["Medical=True", "Medical=False"], rotation=20, ha="right")
        ax.set_yticks([0, 1])
        ax.set_yticklabels(["OMIM exon", "no OMIM exon"])


        for (i, j), v in np.ndenumerate(mat):
            ax.text(j, i, str(int(v)), ha="center", va="center")

        plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)

    _plot_heat(axes[0], mat_all, "OMIM exons × Medical (All)")
    _plot_heat(axes[1], mat_nom, "OMIM exons × Medical (Excluding mena=True)")

    fig.suptitle("2×2 category counts" + (" (unique ids)" if count_unique_ids else " (rows)"), y=1.02)
    fig.tight_layout()
    return fig
    
def plot_pie_grid_3x5(
    data,
    labels,
    titles,
    *,
    figsize=(18, 10),
    startangle: int = 90,
):
    """
    Plot a 3x5 grid of pie charts showing raw counts.

    Parameters
    ----------
    data : list of sequences of int
        Length must be <= 15.
        Each element is a list/tuple of counts for one pie.
        Example: [[10, 5], [3, 7, 2], ...]
    labels : list of sequences of str
        Same structure as data. Labels for each pie slice.
    titles : list of str, optional
        One title per pie (length <= 15).
    figsize : tuple
        Figure size.
    startangle : int
        Starting angle for pies.

    Returns
    -------
    fig : matplotlib.figure.Figure
    """

    if len(data) > 15:
        raise ValueError("Maximum of 15 pie charts (3x5 grid).")

    if len(data) != len(labels):
        raise ValueError("data and labels must have the same length.")

    if titles is not None and len(titles) != len(data):
        raise ValueError("titles must have same length as data.")

    fig, axes = plt.subplots(3, 5, figsize=figsize)
    axes = axes.flatten()

    def autopct_factory(values):
        total = sum(values)

        def _autopct(pct):
            count = int(round(pct * total / 100.0))
            return f"{count}" if count > 0 else ""

        return _autopct

    for i, ax in enumerate(axes):
        if i >= len(data):
            ax.axis("off")
            continue

        values = list(data[i])
        lbls = list(labels[i])

        if sum(values) == 0:
            ax.text(0.5, 0.5, "No data", ha="center", va="center")
            ax.axis("off")
            continue

        ax.pie(
            values,
            labels=lbls,
            autopct=autopct_factory(values),
            startangle=startangle,
        )
        ax.axis("equal")

        if titles is not None:
            ax.set_title(titles[i], fontsize=10)
    plt.suptitle("reduction of SVs after overlapping with different databases")
    fig.tight_layout()

    return fig
    
def plot_bar_grid_3x5(
    data,
    labels,
    titles,
    *,
    figsize=(18, 10),
    bar_width: float = 0.6,
    rotate_xticks: int = 20,
):
    """
    Plot a 3x5 grid of bar charts showing raw counts.

    Parameters
    ----------
    data : list of sequences of int
        Length must be <= 15.
        Each element is a list/tuple of counts for one bar plot.
    labels : list of sequences of str
        Same structure as data. X-axis labels per bar plot.
    titles : list of str, optional
        One title per panel.
    figsize : tuple
        Figure size.
    bar_width : float
        Width of bars.
    rotate_xticks : int
        Rotation angle for x tick labels.

    Returns
    -------
    fig : matplotlib.figure.Figure
    """

    if len(data) > 15:
        raise ValueError("Maximum of 15 bar plots (3x5 grid).")

    if len(data) != len(labels):
        raise ValueError("data and labels must have the same length.")

    if titles is not None and len(titles) != len(data):
        raise ValueError("titles must have same length as data.")

    fig, axes = plt.subplots(3, 5, figsize=figsize)
    axes = axes.flatten()

    for i, ax in enumerate(axes):
        if i >= len(data):
            ax.axis("off")
            continue

        values = list(data[i])
        lbls = list(labels[i])

        if sum(values) == 0:
            ax.text(0.5, 0.5, "No data", ha="center", va="center")
            ax.axis("off")
            continue

        x = range(len(values))
        bars = ax.bar(x, values, width=bar_width)

        ax.set_xticks(x)
        ax.set_xticklabels(lbls, rotation=rotate_xticks, ha="right")


        for bar, val in zip(bars, values):
            ax.text(
                bar.get_x() + bar.get_width() / 2,
                bar.get_height(),
                str(val),
                ha="center",
                va="bottom",
                fontsize=9,
            )

        if titles is not None:
            ax.set_title(titles[i], fontsize=10)

        ax.set_ylabel("Count")

    fig.tight_layout()
    return fig
    

def upset_plot_matplotlib(
    df: pd.DataFrame,
    set_columns: List[str],
    *,
    id_col: str = "id",
    count_unique_ids: bool = True,
    min_subset_size: int = 1,
    max_intersections: int = 20,
    title: str = "UpSet plot",
):
    """
    Draw an UpSet plot without upsetplot (matplotlib only).

    Parameters
    ----------
    df : DataFrame
    set_columns : list[str]
        Columns defining sets; values may be bool or boolean-ish strings.
    id_col : str
        Used when count_unique_ids=True.
    count_unique_ids : bool
        If True, count unique ids per intersection; else count rows.
    min_subset_size : int
        Hide intersections smaller than this.
    max_intersections : int
        Show at most this many intersections (largest first).
    title : str

    Returns
    -------
    fig : matplotlib.figure.Figure
    """
    data = df.copy()

    def to_bool_series(s: pd.Series) -> pd.Series:
        if s.dtype == bool:
            return s.fillna(False)
        def conv(v):
            if v is None or (isinstance(v, float) and np.isnan(v)):
                return False
            if isinstance(v, bool):
                return v
            return str(v).strip().lower() in {"true", "t", "1", "yes", "y"}
        return s.map(conv).astype(bool)


    for c in set_columns:
        if c not in data.columns:
            raise KeyError(f"Missing set column: {c}")
        data[c] = to_bool_series(data[c])


    if count_unique_ids:
        if id_col not in data.columns:
            raise KeyError(f"id_col='{id_col}' not found in df")
        agg = {c: "max" for c in set_columns}
        data = data.groupby(id_col, as_index=False).agg(agg)


    patterns = data[set_columns].apply(lambda r: tuple(bool(x) for x in r), axis=1)
    counts = patterns.value_counts()


    counts = counts[counts >= min_subset_size]
    counts = counts.sort_values(ascending=False).head(max_intersections)

    if counts.empty:
        fig, ax = plt.subplots(figsize=(8, 3))
        ax.text(0.5, 0.5, "No intersections meet min_subset_size", ha="center", va="center")
        ax.axis("off")
        return fig


    set_sizes = data[set_columns].sum().astype(int)  # number of True per set


    patterns_kept = list(counts.index)  # list of boolean tuples
    mat = np.array(patterns_kept, dtype=int).T  # shape: (n_sets, n_intersections)

    n_sets = len(set_columns)
    n_int = len(patterns_kept)


    fig = plt.figure(figsize=(max(10, 0.5 * n_int + 4), max(6, 0.35 * n_sets + 4)))
    gs = GridSpec(
        nrows=2, ncols=2,
        width_ratios=[1.2, 4.0],
        height_ratios=[2.0, 3.0],
        wspace=0.05, hspace=0.05
    )

    ax_top = fig.add_subplot(gs[0, 1])
    ax_left = fig.add_subplot(gs[1, 0])
    ax_mat = fig.add_subplot(gs[1, 1])


    x = np.arange(n_int)
    top_vals = counts.values.astype(int)
    ax_top.bar(x, top_vals)
    ax_top.set_ylabel("Intersection size")
    ax_top.set_xticks([])

    for i, v in enumerate(top_vals):
        ax_top.text(i, v, str(int(v)), ha="center", va="bottom", fontsize=9)


    
    
    y = np.arange(n_sets)
    left_vals = set_sizes.values.astype(int)
    bars = ax_left.barh(y, left_vals)
    ax_left.set_yticks(y)
    ax_left.set_yticklabels(set_columns)
    ax_left.invert_yaxis()
    ax_left.set_xlabel("Set size")
    

    for bar, val in zip(bars, left_vals):
        ax_left.text(
        val,
        bar.get_y() + bar.get_height() / 2,
        str(val),
        va="center",
        ha="left",
        fontsize=9,
        )


    ax_mat.set_xlim(-0.5, n_int - 0.5)
    ax_mat.set_ylim(-0.5, n_sets - 0.5)
    ax_mat.invert_yaxis()

    ax_mat.set_xticks(x)
    ax_mat.set_xticklabels([""] * n_int)
    ax_mat.set_yticks(y)
    ax_mat.set_yticklabels([""] * n_sets)


    for i in range(n_sets):
        ax_mat.hlines(i, -0.5, n_int - 0.5, linewidth=0.5, alpha=0.3)


    for j in range(n_int):
        present = np.where(mat[:, j] == 1)[0]

        ax_mat.scatter([j] * len(present), present, s=40)

        if len(present) >= 2:
            ax_mat.plot([j, j], [present.min(), present.max()], linewidth=2)

    ax_mat.set_xlabel("Intersections (sorted by size)")

    fig.suptitle(title, y=0.98)
    fig.tight_layout()
    return fig