import pandas as pd
import os
from tqdm import tqdm
import numpy as np
import ast
# plotting imports 
import seaborn as sns
import matplotlib
import matplotlib.pyplot as plt
matplotlib.rcParams['pdf.fonttype'] = 42
matplotlib.rcParams['ps.fonttype'] = 42
# matplotlib.rcParams['font.sans-serif'] = "Arial" # missing fonts:: https://alexanderlabwhoi.github.io/post/2021-03-missingfont/
# Then, "ALWAYS use sans-serif fonts"
matplotlib.rcParams['font.family'] = "sans-serif"
matplotlib.rcParams.update({'font.size': 6, 'axes.linewidth': 1, 'xtick.major.width': 1, 'xtick.major.size': 5, 'ytick.major.width': 1, 'ytick.major.size': 5})


def get_most_impactful_consequence(consequence_str):
    """
    Given a comma-separated string of consequences,
    return the most impactful one using a fixed priority order.
    """
    if not consequence_str or pd.isna(consequence_str):
        return None
    priority = [
        "frameshift_variant", "stop_gained",
        "splice_acceptor_variant", "splice_donor_variant",
    ]
    consequences = [c.strip() for c in consequence_str.split(",")]
    for p in priority:
        if p in consequences:
            return p
    # otherwise, just return the first available consequence
    return consequences[0]



def create_variant_frequency_plots(
    variant_df,
    n_table=10,
    # columns in variant_df:
    locus_col="locus",
    alleles_col="alleles",
    consequence_col="consequence",
    person_id_col="hetz_samples",              # optional; if provided, counts unique carriers per consequence
    # appearance:
    figsize=(11, 5),
    bar_color="#D1495B",
    bar_fontsize=6,
    table_fontsize=6,
    bar_height_ratio=0.25,            # relative height of bar panel vs table panel
):
    df = variant_df.copy()

    # --- derive fields (as you started) ---
    df[consequence_col] = df[consequence_col].apply(get_most_impactful_consequence)
    df[alleles_col] = df[alleles_col].apply(lambda x: "_".join(ast.literal_eval(x)))

    # -----------------------------
    # Build vdf_plot for bar panel
    # -----------------------------
    vdf_plot = df.groupby(consequence_col).agg({person_id_col: "nunique"}).reset_index().sort_values(person_id_col)

    # -----------------------------
    # Tables: most / least frequent
    # -----------------------------

    needed = [locus_col, alleles_col, consequence_col, person_id_col]
    for c in needed:
        if c not in df.columns:
            raise ValueError(f"Missing required column: {c}")

    vdf = df.groupby(needed[:3]).agg({person_id_col: "nunique"}).sort_values(person_id_col).reset_index()
    # top/bottom frames for display
    cols = (locus_col, alleles_col, consequence_col, person_id_col)
    top_df = vdf.sort_values(person_id_col, ascending=False).head(n_table).reset_index(drop=True)
    bottom_df = vdf.sort_values(person_id_col, ascending=True).head(n_table).reset_index(drop=True)

    # -----------------------------
    # Create ONE combined figure
    # -----------------------------
    fig = plt.figure(figsize=figsize)
    gs = fig.add_gridspec(
        nrows=2, ncols=4,
        height_ratios=[bar_height_ratio, 1.0],
        width_ratios=[1, 1, 1, 1],
        hspace=0.35, wspace=0.025
    )

    ax_bar = fig.add_subplot(gs[0, 1:3])   # spans both columns on top
    ax_t1  = fig.add_subplot(gs[1, 0:2])
    ax_t2  = fig.add_subplot(gs[1, 2:4])

    # --- barplot (top) ---
    sns.barplot(
        data=vdf_plot,
        x=person_id_col,
        y=consequence_col,
        color=bar_color,
        ax=ax_bar,
        errorbar=None
    )
    ax_bar.set_xlabel("Number of heterozygous carriers")
    ax_bar.set_ylabel("")
    ax_bar.spines[['top', 'right']].set_visible(False)

    # annotate counts
    xmax = max(vdf_plot[person_id_col]) if len(vdf_plot) else 0
    for p in ax_bar.patches:
        width = p.get_width()
        ax_bar.text(
            width + (xmax * 0.02 if xmax else 0.1),
            p.get_y() + p.get_height() / 2,
            f"{int(width)}",
            va="center",
            ha="left",
            fontsize=bar_fontsize
        )

    # --- tables (bottom) ---
    ax_t1.axis("off")
    ax_t2.axis("off")

    def _prep_table_df(df_):
        out = df_.loc[:, list(cols)].copy()
        out[person_id_col] = out[person_id_col].astype(int)
        return out

    top_show = _prep_table_df(top_df)
    bottom_show = _prep_table_df(bottom_df)

    t1 = ax_t1.table(
        cellText=top_show.values,
        colLabels=list(cols),
        cellLoc="center",
        loc="center"
    )
    t2 = ax_t2.table(
        cellText=bottom_show.values,
        colLabels=list(cols),
        cellLoc="center",
        loc="center"
    )

    def _style_table(t, header_color="#F2F2F2"):
        t.auto_set_font_size(False)
        t.set_fontsize(table_fontsize)
        t.scale(1, 1.4)
        for (r, c), cell in t.get_celld().items():
            if r == 0:
                cell.set_text_props(weight="bold")
                cell.set_facecolor(header_color)

    _style_table(t1)
    _style_table(t2)

    ax_t1.set_title(f"{n_table} most frequent variants", fontsize=table_fontsize + 1, pad=1)
    ax_t2.set_title(f"{n_table} least frequent variants", fontsize=table_fontsize + 1, pad=1)

    fig.tight_layout()

    axes = {"bar": ax_bar, "top_table": ax_t1, "bottom_table": ax_t2}
    return fig, axes
