import os
import re
import shutil
import zipfile
from pathlib import Path
from playwright.sync_api import Playwright, sync_playwright
from rwe.plots.clinical import manhattan
import matplotlib.pyplot as plt
import pandas as pd

import rwe.utils.helpers as uth
from docx.shared import Inches

AZN_CHAPTER_SHORT = {
    "Chapter I Certain infectious and parasitic diseases": "Infectious",
    "Chapter II Neoplasms": "Neoplasms",
    "Chapter III Diseases of the blood and blood-forming organs and certain disorders involving the immune mechanism": "Blood / Immune",
    "Chapter IV Endocrine nutritional and metabolic diseases": "Endocrine / Metabolic",
    "Chapter V Mental and behavioural disorders": "Mental / Behavioral",
    "Chapter VI Diseases of the nervous system": "Nervous System",
    "Chapter VII Diseases of the eye and adnexa": "Eye",
    "Chapter VIII Diseases of the ear and mastoid process": "Ear / Mastoid",
    "Chapter IX Diseases of the circulatory system": "Circulatory",
    "Chapter X Diseases of the respiratory system": "Respiratory",
    "Chapter XI Diseases of the digestive system": "Digestive",
    "Chapter XII Diseases of the skin and subcutaneous tissue": "Skin",
    "Chapter XIII Diseases of the musculoskeletal system and connective tissue": "Musculoskeletal",
    "Chapter XIV Diseases of the genitourinary system": "Genitourinary",
    "Chapter XV Pregnancy childbirth and the puerperium": "Pregnancy / Childbirth",
    "Chapter XVII Congenital malformations deformations and chromosomal abnormalities": "Congenital",
    "Chapter XVIII Symptoms signs and abnormal clinical and laboratory findings not elsewhere classified": "Symptoms / Findings",
    "Chapter XXI Factors influencing health status and contact with health services": "Health / Services",
}

def extract_only_csv_rename(zip_path: str, gene: str, trait_type: str) -> str:
    zip_path = Path(zip_path)
    gene_dir = zip_path.parent
    target_csv = gene_dir / f"{gene}_azn_{trait_type}.csv"

    # Extract into a dedicated folder
    extract_dir = gene_dir / f"__extract_{zip_path.stem}"
    if extract_dir.exists():
        shutil.rmtree(extract_dir)
    extract_dir.mkdir(parents=True, exist_ok=True)

    with zipfile.ZipFile(zip_path, "r") as z:
        z.extractall(extract_dir)

    csvs = list(extract_dir.rglob("*.csv"))
    if not csvs:
        raise FileNotFoundError(f"No CSV found after extracting {zip_path}")

    csvs.sort(key=lambda p: p.stat().st_size, reverse=True)
    chosen = csvs[0]

    if target_csv.exists():
        target_csv.unlink()
    chosen.replace(target_csv)

    # Cleanup only what we created
    shutil.rmtree(extract_dir)
    if zip_path.exists():
        zip_path.unlink()

    return str(target_csv)


def run(playwright: Playwright, gene: str, save_path: str, trait_type="binary") -> None:
    browser = playwright.chromium.launch(headless=True)
    context = browser.new_context(accept_downloads=True)
    page = context.new_page()

    url = f"https://azphewas.com/geneView/6319c068-fd59-46d8-85ee-82d82482eb14/{gene}/glr/{trait_type}"
    page.goto(url, wait_until="networkidle")

    export_control = page.locator('[data-testid="export-settings-control"] button:has-text("Export")')
    export_control.click()

    inner_export_btn = page.get_by_test_id("export_button")
    inner_export_btn.wait_for(state="visible")

    plot_checkbox = page.locator('[data-testid="plot-export-type-toggler"] input[type="checkbox"]')
    plot_checkbox.set_checked(False, force=True)
    page.wait_for_timeout(300)

    os.makedirs(save_path, exist_ok=True)

    with page.expect_download(timeout=120_000) as dl_info:
        inner_export_btn.click(force=True)

    download = dl_info.value
    fname = download.suggested_filename
    out_dir = save_path
    os.makedirs(out_dir, exist_ok=True)
    out_path = os.path.join(out_dir, fname)
    download.save_as(out_path)
    print("Saved zip:", out_path)

    final_csv = extract_only_csv_rename(out_path, gene=gene, trait_type=trait_type)
    print("Final CSV:", final_csv)

    context.close()
    browser.close()
    return

def clean_azn_phewas(phewas_file):
    df = pd.read_csv(phewas_file, sep=",")
    df = df.loc[df["Collapsing model"]=="ptv"]
    return df

def get_azn_manhattan(gene, phewas_file):
    df = clean_azn_phewas(phewas_file)
    name_col = "Phenotype"
    p_col    = "P value"
    odds_ratio_col = "Odds ratio"
    cat_col  = "Phenotypic category"

    df[name_col] = df[name_col].astype(str).str.rsplit("#", n=1).str[-1]
    df[cat_col] = df[cat_col].map(AZN_CHAPTER_SHORT).fillna("Other")

    sig_p = 0.05/len(df)
    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=sig_p,
        top_k_labels=5,
        title=f"{gene} PheWAS AstraZeneca",
    )
    return fig, plot_df

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

    if not os.path.exists(phewas_file):
        with sync_playwright() as playwright:
            run(playwright, gene=gene, save_path=phewas_dir, trait_type="binary")

    doc.add_heading("Phenome-wide Association Study (PheWAS) in UK Biobank cohort collected from AstraZeneca", level=2)
    if os.path.exists(phewas_file):
        fig, plot_df = get_azn_manhattan(gene, phewas_file)
        fig_path = uth._save_fig_to_tmp(fig, basename="azn_phewas", dpi=300)
        plt.close(fig)  # close the figure to free memory
        doc.add_paragraph()  # spacing
        doc.add_picture(fig_path, width=Inches(6.5))
        uth._add_figure_caption(doc, f"{gene} PheWAS results from AstraZeneca.")
        doc.add_paragraph()  # spacing

        # Table: Top 10 significant hits (sorted by p-value)
        top_significant = plot_df.nsmallest(10, 'P value')
        columns = ['Phenotype', 'Phenotypic category', 'Odds ratio', 'P value', 'Ancestry']
        # Rename columns for display
        display_df = top_significant[columns].copy()
        display_df.columns = ['Phenotype', 'Category', 'OR', 'P-value', 'Ancestry']
        uth._add_table_to_doc(doc, display_df, uth._add_table_title(doc, f"Top 10 significant associations for {gene} in AstraZeneca"),
                        ['Phenotype', 'Category', 'OR', 'P-value', 'Ancestry'], list(map(Inches, [2.25, 1.4, 0.7, 0.85, 0.8]))
                        )
        # Table: Top 10 negative beta hits (odds_ratio < 0, sorted by p-value)
        negative_beta = plot_df[plot_df['Odds ratio'] < 1].nsmallest(10, 'P value')
        if len(negative_beta) > 0:
            display_df_neg = negative_beta[columns].copy()
            display_df_neg.columns = ['Phenotype', 'Category', 'OR', 'P-value', 'Ancestry']
            uth._add_table_to_doc(doc, display_df_neg, uth._add_table_title(doc, f"Top 10 protective associations for {gene} in AstraZeneca"),
                            ['Phenotype', 'Category', 'OR', 'P-value', 'Ancestry'], list(map(Inches, [2.25, 1.4, 0.7, 0.85, 0.8]))
                            )
        doc.add_page_break()  # spacing

    else:
        doc.add_paragraph(f"No AstraZeneca PheWAS results found for {gene}.")
        doc.add_page_break()  # spacing
    return doc


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

    if not os.path.exists(phewas_file):
        with sync_playwright() as playwright:
            run(playwright, gene=gene, save_path=phewas_dir, trait_type="continuous")

    doc.add_heading("Association study results of labs and measurements in UK Biobank cohort collected from AstraZeneca", level=2)
    if os.path.exists(phewas_file):
        phewas_df = clean_azn_phewas(phewas_file)
        columns = ['Phenotype', 'Phenotypic category', 'Effect size', 'P value', 'Ancestry']
        display_df = phewas_df[columns].copy()
        display_df.columns = ['Phenotype', 'Category', 'Beta', 'P-value', 'Ancestry']
        display_df = display_df.loc[display_df["Category"].isin(AZN_CHAPTER_SHORT.keys())].sort_values("P-value")
        display_df["Category"] = display_df["Category"].map(AZN_CHAPTER_SHORT).fillna(display_df["Category"])
        display_df = display_df.head(50)
        doc.add_paragraph()  # spacing
        uth._add_table_to_doc(doc, display_df, uth._add_table_title(doc, f"Top 50 associations for {gene} in AstraZeneca"),
                        ['Phenotype', 'Category', 'Beta', 'P-value', 'Ancestry'], list(map(Inches, [2.25, 1.4, 0.7, 0.85, 0.8]))
                        )
        doc.add_page_break()  # spacing
    else:
        doc.add_paragraph()  # spacing
        doc.add_paragraph(f"No AstraZeneca labs and measurements results found for {gene}.")
        doc.add_page_break()  # spacing
    return doc


def generate_azn_indication_specific_report(doc, gene, phewas_dir, keywords, trait_types=["binary", "continuous"], phewas_filenames=["", ""]):
    trait_to_effect_dict = {"binary": "Odds ratio", "continuous": "Effect size"}
    display_dfs = []
    doc.add_heading("Indication specific PheWAS results in UK Biobank cohort collected from AstraZeneca", level=2)
    doc.add_paragraph("The keywords searched to generate the report were: " + ", ".join(keywords))
    for trait, fn in zip(trait_types, phewas_filenames):
        if not fn:
            phewas_filename = f"{gene}_azn_{trait}.csv"
        else:
            phewas_filename = fn
        phewas_file = os.path.join(phewas_dir, phewas_filename)
        
        if os.path.exists(phewas_file):
            phewas_df = clean_azn_phewas(phewas_file)
            escaped_keywords = [re.escape(k) for k in keywords if k]
            if not escaped_keywords:
                continue
            keyword_pattern = "|".join(escaped_keywords)
            phewas_df = phewas_df.loc[(phewas_df.Phenotype.str.contains(keyword_pattern, case=False, na=False))]
            columns = ['Phenotype', 'Phenotypic category', trait_to_effect_dict[trait], 'P value', 'Ancestry']
            display_df = phewas_df[columns].copy()
            display_df.columns = ['Phenotype', 'Category', 'Effect', 'P-value', 'Ancestry']
            display_df["Trait type"] = trait
            display_dfs.append(display_df)
    
    if not display_dfs:
        display_df = pd.DataFrame()
    else:
        display_df = pd.concat(display_dfs, ignore_index=True).sort_values("P-value").reset_index(drop=True)
        display_df["Category"] = display_df["Category"].map(AZN_CHAPTER_SHORT).fillna(display_df["Category"])

    if len(display_df) == 0:
        doc.add_paragraph()  # spacing
        doc.add_paragraph("No associations found for the specified indications. Please broaden the indication keywords and try again.")
        doc.add_page_break()
    else:
        doc.add_paragraph()  # spacing
        uth._add_table_to_doc(doc, display_df, uth._add_table_title(doc, f"Indication-specific associations for {gene} in AstraZeneca"),
                        ['Phenotype', 'Category', 'Effect', 'P-value', 'Trait type', 'Ancestry'], list(map(Inches, [1.8, 1, 0.75, 0.75, 0.9, 0.8]))
                        )
        doc.add_page_break()
    return doc

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

    if not os.path.exists(phewas_file):
        with sync_playwright() as playwright:
            run(playwright, gene=gene, save_path=phewas_dir, trait_type="continuous")

    doc.add_heading("Association study results of plasma proteomics data in UK Biobank cohort collected from AstraZeneca", level=2)
    if os.path.exists(phewas_file):
        phewas_df = clean_azn_phewas(phewas_file)
        columns = ['Phenotype', 'Phenotypic category', 'Effect size', 'P value']
        display_df = phewas_df[columns].copy()
        display_df.columns = ['Phenotype', 'Category', 'Beta', 'P-value']
        display_df = display_df.loc[display_df["Category"]=="Proteomics"].sort_values("P-value")
        display_df = display_df.head(50)
        display_df["Phenotype"] = display_df["Phenotype"].str.split("#").str[2] + " (" + display_df["Phenotype"].str.split("#").str[3] + ")"
        doc.add_paragraph()  # spacing
        uth._add_table_to_doc(doc, display_df, uth._add_table_title(doc, f"Top 50 plasma proteomics biomarkers for {gene} in AstraZeneca"),
                        ['Phenotype', 'Category', 'Beta', 'P-value'], list(map(Inches, [2.75, 1.75, 0.75, 0.75]))
                        )
        doc.add_page_break()  # spacing
    else:
        doc.add_paragraph()  # spacing
        doc.add_paragraph(f"No AstraZeneca proteomics data found for {gene}.")
        doc.add_page_break()  # spacing
    return doc

if __name__ == "__main__":
    from docx import Document
    doc = Document()
    gene = "MBL2"
    phewas_dir = "/home/dbanerjee/deepro/rwe/data/test/"
    doc = generate_azn_clinical_report(doc, gene, phewas_dir)
    doc = generate_azn_labs_measurements_report(doc, gene, phewas_dir)
    doc = generate_azn_plasma_protein_biomarkers_report(doc, gene, phewas_dir)
    doc = generate_azn_indication_specific_report(doc, gene, phewas_dir, keywords=["diab", "obes", "kid", "ren", "fat", "weig"])
    doc.save("/home/dbanerjee/deepro/rwe/data/test/MBL2_azn_report.docx")
