import os
import numpy as np
from playwright.sync_api import Playwright, sync_playwright
import random
import pandas as pd
from docx.shared import Inches, Pt

from rwe.plots.clinical import manhattan
import rwe.utils.helpers as uth
from rwe.clients.hgnc import get_gene_info

def run(playwright: Playwright, gene:str, gene_id:str, save_path:str) -> None:
    browser = playwright.chromium.launch(headless=True)
    context = browser.new_context(accept_downloads=True)
    # Open new page
    page = context.new_page()
    # Go to https://app.genebass.org/
   # page.goto("https://app.genebass.org/")
    query_id = "https://app.genebass.org/gene/" + gene_id +"?burdenSet=pLoF&phewasOpts=1&resultLayout=full"
    page.goto(query_id)
    page.click("label:has-text(\"Burden\")")
    page.wait_for_timeout(random.uniform(100, 1000))
    # Click text=Export data to CSV
    with page.expect_download() as download_info:
        page.click("text=Export data to CSV")
    download = download_info.value
    download_path = os.path.join(save_path, f"{gene}_genebass.csv")
    os.makedirs(os.path.dirname(download_path), exist_ok=True)
    download.save_as(path=download_path)
    page.wait_for_timeout(random.uniform(100, 1000))
    # ---------------------
    context.close()
    browser.close()
    return

def clean_genebass_phewas(phewas_file):
    df = pd.read_csv(phewas_file, sep=",")
    df = df.loc[df["Trait type"]=="icd_first_occurrence"]
    return df

def get_genebass_manhattan(gene, phewas_file):
    df = clean_genebass_phewas(phewas_file)
    name_col = "Description"
    p_col    = "P-Value (Burden)"
    odds_ratio_col = None
    beta_col = "Beta"
    cat_col  = "Category"

    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"))
    GENEBASS_CATEGORY_SHORT = {
        "Blood, blood-forming organs and certain immune disorders": "Blood/Immune",
        "Certain conditions originating in the perinatal period": "Perinatal",
        "Certain infectious and parasitic diseases": "Infectious",
        "Circulatory system disorders": "Circulatory",
        "Congenital disruptions and chromosomal abnormalities": "Congenital",
        "Digestive system disorders": "Digestive",
        "Ear and mastoid process disorders": "Ear/Mastoid",
        "Endocrine, nutritional and metabolic diseases": "Endocrine/Metabolic",
        "Eye and adnexa disorders": "Eye",
        "Genitourinary system disorders": "Genitourinary",
        "Mental and behavioural disorders": "Mental/Behavioral",
        "Musculoskeletal system and connective tissue disorders": "Musculoskeletal",
        "Nervous system disorders": "Nervous System",
        "Pregnancy, childbirth and the puerperium": "Pregnancy/Childbirth",
        "Respiratory system disorders": "Respiratory",
        "Skin and subcutaneous tissue disorders": "Skin",
    }
    df[cat_col] = df[cat_col].map(GENEBASS_CATEGORY_SHORT).fillna(df[cat_col])
    fig, ax, plot_df = manhattan(
        df,
        p_col=p_col,
        category_col=cat_col,
        label_col=name_col,
        beta_col=beta_col,
        sig_p=2.8e-4,
        top_k_labels=5,
        title=f"{gene} PheWAS GeneBass",
    )
    return fig, plot_df

def generate_genebass_clinical_report(doc, gene, phewas_dir, phewas_filename=""):
    if not phewas_filename:
        phewas_filename = f"{gene}_genebass.csv"
    phewas_file = os.path.join(phewas_dir, phewas_filename)

    if not os.path.exists(phewas_file):
        hgnc_path = os.path.join(phewas_dir, "gene_with_protein_product.txt")
        approved_symbol, ensembl_id, entrez_id = get_gene_info(gene, hgnc_path)
        with sync_playwright() as playwright:
            run(playwright, approved_symbol, ensembl_id, save_path=phewas_dir)
    doc.add_heading("Phenome-wide Association Study (PheWAS) in UK Biobank cohort collected from GeneBass", level=2)
    if os.path.exists(phewas_file):
        fig, plot_df = get_genebass_manhattan(gene, phewas_file)
        fig_path = uth._save_fig_to_tmp(fig, basename="genebass_phewas", dpi=300)
        doc.add_paragraph()  # spacing
        doc.add_picture(fig_path, width=Inches(6))
        uth._add_figure_caption(doc, f"{gene} PheWAS results from GeneBass.")
        doc.add_paragraph()  # spacing

        # Table: Top 10 significant hits (sorted by p-value)
        top_significant = plot_df.nsmallest(10, 'P-Value (Burden)')
        columns = ['Description', 'Category', 'Beta', 'P-Value (Burden)']
        # Rename columns for display
        display_df = top_significant[columns].copy()
        display_df.columns = ['Description', 'Category', 'Beta', 'P-value']
        uth._add_table_to_doc(doc, display_df, uth._add_table_title(doc, f"Top 10 significant associations for {gene} in GeneBass"),
                        ['Description', 'Category', 'Beta', 'P-value'], list(map(Inches, [2.75, 2.0, 0.5, 0.75]))
                        )

        # Table: Top 10 negative beta hits (beta < 0, sorted by p-value)
        negative_beta = plot_df[plot_df['Beta'] < 0].nsmallest(10, 'P-Value (Burden)')
        if len(negative_beta) > 0:
            display_df_neg = negative_beta[columns].copy()
            display_df_neg.columns = ['Description', 'Category', 'Beta', 'P-value']
            uth._add_table_to_doc(doc, display_df_neg, uth._add_table_title(doc, f"Top 10 protective associations for {gene} in GeneBass"),
                            ['Description', 'Category', 'Beta', 'P-value'], list(map(Inches, [2.75, 2.0, 0.5, 0.75]))
                            )
            
        doc.add_page_break()  # spacing
    else:
        doc.add_paragraph(f"No GeneBass PheWAS results found for {gene}.")
        doc.add_page_break()  # spacing
    return doc

if __name__ == "__main__":
    from docx import Document
    doc = Document()
    gene = "INHBE"
    phewas_dir = "/home/dbanerjee/deepro/rwe/data/test/"
    doc = generate_genebass_clinical_report(doc, gene, phewas_dir)
    doc.save("/home/dbanerjee/deepro/rwe/data/test/INHBE_genebass_report.docx")

