import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import textwrap


def manhattan(
    df: pd.DataFrame,
    *,
    p_col: str,
    category_col: str,
    label_col: str,
    # effect direction: provide ONE of these
    beta_col: str | None = None,
    odds_ratio_col: str | None = None,
    # optional:
    y_col: str | None = None,          # if you already have -log10(p) in a column
    id_col: str | None = None,         # used only if label_col missing/empty
    title: str | None = None,
    figsize=(8, 4),
    gap: int = 10,
    alpha: float = 0.85,
    s: float = 18,
    sig_p: float | None = None,        # draw horizontal line at -log10(sig_p)
    top_k_labels: int = 5,             # set >0 to annotate top hits
    ax=None,
):
    """
    General Manhattan plot with triangle direction.

    Required columns:
      - p_col: p-values
      - category_col: category/chapter (GeneBass: keep most specific string)
      - label_col: phenotype description to annotate

    Direction:
      - if beta_col provided: up if beta>0 else down if beta<0
      - elif odds_ratio_col provided: up if OR>1 else down if OR<1
    """

    if (beta_col is None) == (odds_ratio_col is None):
        raise ValueError("Provide exactly one of beta_col or odds_ratio_col.")

    d = df.copy()

    # ---- y = -log10(p) ----
    if y_col is not None and y_col in d.columns:
        d["_y"] = pd.to_numeric(d[y_col], errors="coerce")
    else:
        p = pd.to_numeric(d[p_col], errors="coerce")
        p = p.clip(lower=np.nextafter(0, 1))  # avoid log(0)
        d["_y"] = -np.log10(p)

    # ---- direction (+1 up, -1 down) ----
    if beta_col is not None:
        b = pd.to_numeric(d[beta_col], errors="coerce")
        d["_dir"] = np.where(b > 0, 1, np.where(b < 0, -1, np.nan))
    else:
        orv = pd.to_numeric(d[odds_ratio_col], errors="coerce")
        d["_dir"] = np.where(orv > 1, 1, np.where(orv < 1, -1, np.nan))

    # ---- basic cleaning ----
    d[category_col] = d[category_col].astype(str).fillna("Unknown")
    d[label_col] = d[label_col].astype(str).fillna("")
    if id_col is not None and id_col in d.columns:
        d[id_col] = d[id_col].astype(str)

    d = d.replace([np.inf, -np.inf], np.nan)
    d = d.dropna(subset=["_y", category_col, "_dir"])

    # # ---- order categories by max signal ----
    # cat_order = (
    #     d.groupby(category_col)["_y"].max()
    #     .sort_values(ascending=False)
    #     .index.tolist()
    # )
    # ---- order categories alphabetically for consistent display ----
    cat_order = sorted(d[category_col].dropna().unique().tolist())


    # ---- build x positions ----
    parts = []
    x_offset = 0
    xticks, xticklabels = [], []

    for cat in cat_order:
        sub = d[d[category_col] == cat].copy()
        sub = sub.sort_values("_y", ascending=False).reset_index(drop=True)
        sub["_x"] = np.arange(len(sub)) + x_offset
        if len(sub) > 0:
            xticks.append(sub["_x"].iloc[len(sub)//2])
            xticklabels.append(cat)
        parts.append(sub)
        x_offset += len(sub) + gap

    plot_df = pd.concat(parts, ignore_index=True) if parts else d.assign(_x=np.arange(len(d)))

    # ---- plot ----
    if ax is None:
        fig = plt.figure(figsize=figsize, constrained_layout=True)
        ax = plt.gca()
    else:
        fig = ax.figure

    # colors per category (no explicit palette; uses matplotlib defaults)
    cats = cat_order if cat_order else plot_df[category_col].unique().tolist()
    cmap = plt.get_cmap("tab20", max(len(cats), 1))
    cat_to_color = {c: cmap(i) for i, c in enumerate(cats)}
    colors = plot_df[category_col].map(cat_to_color).values

    up = plot_df["_dir"] == 1
    down = plot_df["_dir"] == -1

    ax.scatter(plot_df.loc[up, "_x"], plot_df.loc[up, "_y"],
               s=s, marker="^", c=colors[up.values], alpha=alpha, linewidths=0)
    ax.scatter(plot_df.loc[down, "_x"], plot_df.loc[down, "_y"],
               s=s, marker="v", c=colors[down.values], alpha=alpha, linewidths=0)

    if sig_p is not None:
        y = -np.log10(sig_p)
        ax.axhline(y, linestyle="--", linewidth=1)
        ax.text(plot_df["_x"].max() if len(plot_df) else 0, y, f"  p={sig_p:g}", va="bottom")

    ax.set_xticks(xticks)
    xticklabels_wrapped = [textwrap.shorten(str(x), width=40) for x in xticklabels]
    ax.set_xticklabels(xticklabels_wrapped, rotation=25, ha="right", fontsize=6)
    ax.set_ylabel(r"$-\log_{10}(p)$")
    ax.set_xlabel("")
    ax.spines[["top", "right"]].set_visible(False)
    if title:
        ax.set_title(title)

    # ---- annotate top hits ----
    if top_k_labels and top_k_labels > 0 and len(plot_df) > 0:
        top = plot_df.nlargest(top_k_labels, "_y")
        for _, r in top.iterrows():
            lab = r[label_col]
            if (not lab or lab == "nan") and id_col is not None and id_col in plot_df.columns:
                lab = r[id_col]
            if not lab or lab == "nan":
                continue
            # ax.text(r["_x"], r["_y"], str(lab), fontsize=6, rotation=20, ha="left", va="bottom")
            lab_wrapped = textwrap.shorten(str(lab), width=40)
            ax.text(r["_x"], r["_y"], lab_wrapped, fontsize=6, rotation=20, ha="left", va="bottom")

    fig.tight_layout()
    return fig, ax, plot_df


if __name__ == "__main__":
    # simple AoU test
    df = pd.read_csv("/home/dbanerjee/deepro/rwe/data/aou/INHBE/INHBE_hetz_phewas.tsv", sep="\t")
    df = df.loc[(df.ancestry == "all")&(df.converged==True)]
    fig, ax, plot_df = manhattan(
        df,
        p_col="p_value",
        category_col="phecode_category",
        label_col="phecode_string",
        odds_ratio_col="odds_ratio",
        sig_p=5e-8,
        top_k_labels=5,
        title="INHBE PheWAS AoU",
    )
    fig.savefig("/home/dbanerjee/deepro/rwe/data/aou/INHBE/phwas_manhattan.png", dpi=300)
    print(plot_df.loc[plot_df.odds_ratio < 1].sort_values("p_value").head(10))

    # simple genebass test
    df = pd.read_csv("/home/dbanerjee/deepro/rwe/data/gbs/ANGPTL3/ANGPTL3_genebass.csv", sep=",")
    df = df.loc[df["Trait type"]=="icd_first_occurrence"]
    name_col = "Description"
    p_col    = "P-Value (Burden)"
    beta_col = "Beta"
    cat_col  = "Category"
    # ---- clean ----
    df[p_col] = pd.to_numeric(df[p_col], errors="coerce")
    df = df[df[p_col].notna() & (df[p_col] > 0) & np.isfinite(df[p_col])]
    prefix = "Health-related outcomes > First occurrences > "
    df[cat_col] = (df[cat_col].astype(str)
                    .str.replace(prefix, "", regex=False)
                    .fillna("Unknown"))

    fig, ax, plot_df = manhattan(
        df,
        p_col=p_col,
        category_col=cat_col,
        label_col=name_col,
        beta_col=beta_col,
        sig_p=5e-8,
        top_k_labels=5,
        title="ANGPTL3 PheWAS GeneBass",
    )
    fig.savefig("/home/dbanerjee/deepro/rwe/data/gbs/ANGPTL3/phewas_manhattan.png", dpi=300)
    print(plot_df.loc[plot_df[beta_col] < 0].sort_values(p_col).head(10))

    # simple astrazeneca test
    df = pd.read_csv("/home/dbanerjee/deepro/rwe/data/azn/MC4R/MC4R_azphewas.csv", sep=",")
    df = df.loc[df["Collapsing model"]=="ptv"]
    name_col = "Phenotype"
    p_col    = "P value"
    odds_ratio_col = "Odds ratio"
    cat_col  = "Phenotypic category"

    fig, ax, plot_df = manhattan(
        df,
        p_col=p_col,
        category_col=cat_col,
        label_col=name_col,
        odds_ratio_col=odds_ratio_col,
        sig_p=5e-8,
        top_k_labels=5,
        title="MC4R PheWAS AstraZeneca",
    )
    fig.savefig("/home/dbanerjee/deepro/rwe/data/azn/MC4R/phewas_manhattan.png", dpi=300)
    print(plot_df.loc[plot_df[odds_ratio_col] < 1].sort_values(p_col).head(10))
