import os
import random
import re
from dataclasses import dataclass
from functools import lru_cache
from importlib.resources import as_file, files

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from docx.shared import Inches
from playwright.sync_api import Playwright, sync_playwright

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

_PLOF_SUMMARY_RESOURCE = "Combined_plof_summary_Gnomad_GeneBass_Guo_2026.csv"

_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",
}


@dataclass
class GeneBassClient:
    gene: str
    phewas_dir: str
    phewas_filename: str = ""

    def __post_init__(self) -> None:
        self.input_gene = self.gene
        self._approved_symbol = None
        self._ensembl_id = None

    @staticmethod
    @lru_cache(maxsize=1)
    def load_plof_summary() -> pd.DataFrame:
        resource = files("rwe") / "assets" / _PLOF_SUMMARY_RESOURCE
        with as_file(resource) as path:
            df = pd.read_csv(
                path,
                usecols=["Gene", "Homozygote_Count", "Heterozygote_Count", "source"],
            )
        df["Gene"] = df["Gene"].astype(str).str.upper()
        df["source"] = df["source"].astype(str)
        return df.loc[df["source"] == "UKB_WES"].copy()

    @property
    def hgnc_path(self) -> str:
        return os.path.join(self.phewas_dir, "gene_with_protein_product.txt")

    def resolve_gene(self) -> tuple[str, str]:
        if self._approved_symbol is None or self._ensembl_id is None:
            approved_symbol, ensembl_id, _entrez_id = get_gene_info(self.input_gene, self.hgnc_path)
            self._approved_symbol = approved_symbol
            self._ensembl_id = ensembl_id
        return self._approved_symbol, self._ensembl_id

    @property
    def approved_symbol(self) -> str:
        approved_symbol, _ensembl_id = self.resolve_gene()
        return approved_symbol

    @property
    def ensembl_id(self) -> str:
        _approved_symbol, ensembl_id = self.resolve_gene()
        return ensembl_id

    @property
    def display_gene(self) -> str:
        return self.approved_symbol

    @property
    def requested_phewas_filename(self) -> str:
        if self.phewas_filename:
            return self.phewas_filename
        return f"{self.input_gene}_genebass_phewas.csv"

    @property
    def requested_phewas_file(self) -> str:
        return os.path.join(self.phewas_dir, self.requested_phewas_filename)

    @property
    def approved_phewas_filename(self) -> str:
        return f"{self.approved_symbol}_genebass_phewas.csv"

    @property
    def approved_phewas_file(self) -> str:
        return os.path.join(self.phewas_dir, self.approved_phewas_filename)

    def existing_phewas_file(self) -> str | None:
        if os.path.exists(self.requested_phewas_file):
            return self.requested_phewas_file
        if self.requested_phewas_file != self.approved_phewas_file and os.path.exists(self.approved_phewas_file):
            return self.approved_phewas_file
        return None

    def run(self, playwright: Playwright) -> None:
        browser = playwright.chromium.launch(headless=True)
        context = browser.new_context(accept_downloads=True)
        page = context.new_page()
        query_id = f"https://app.genebass.org/gene/{self.ensembl_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))
        with page.expect_download() as download_info:
            page.click("text=Export data to CSV")
        download = download_info.value
        os.makedirs(self.phewas_dir, exist_ok=True)
        download.save_as(path=self.approved_phewas_file)
        page.wait_for_timeout(random.uniform(100, 1000))
        context.close()
        browser.close()

    def fetch_phewas_file(self) -> str:
        with sync_playwright() as playwright:
            self.run(playwright)
        return self.approved_phewas_file

    def ensure_phewas_file(self) -> str:
        phewas_file = self.existing_phewas_file()
        if phewas_file:
            return phewas_file
        return self.fetch_phewas_file()

    @staticmethod
    def clean_phewas(phewas_file: str) -> pd.DataFrame:
        df = pd.read_csv(phewas_file, sep=",")
        return df.loc[df["Trait type"] == "icd_first_occurrence"]

    def get_manhattan(self, phewas_file: str):
        df = self.clean_phewas(phewas_file)
        p_col = "P-Value (Burden)"
        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")
        df[cat_col] = df[cat_col].map(_GENEBASS_CATEGORY_SHORT).fillna(df[cat_col])
        sig_p = 0.05 / len(df)

        fig, ax, plot_df = manhattan(
            df,
            p_col=p_col,
            category_col=cat_col,
            label_col="Description",
            beta_col="Beta",
            sig_p=sig_p,
            top_k_labels=5,
            title=f"{self.display_gene} PheWAS GeneBass",
        )
        return fig, plot_df

    def generate_clinical_report(self, doc):
        phewas_file = self.ensure_phewas_file()
        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 = self.get_manhattan(phewas_file)
            fig_path = uth._save_fig_to_tmp(fig, basename="genebass_phewas", dpi=300)
            plt.close(fig)
            doc.add_paragraph()
            doc.add_picture(fig_path, width=Inches(6))
            uth._add_figure_caption(doc, f"{self.display_gene} PheWAS results from GeneBass.")
            doc.add_paragraph()

            top_significant = plot_df.nsmallest(10, "P-Value (Burden)")
            columns = ["Description", "Category", "Beta", "P-Value (Burden)"]
            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 {self.display_gene} in GeneBass"),
                ["Description", "Category", "Beta", "P-value"],
                list(map(Inches, [2.75, 2.0, 0.5, 0.75])),
            )

            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 {self.display_gene} in GeneBass"),
                    ["Description", "Category", "Beta", "P-value"],
                    list(map(Inches, [2.75, 2.0, 0.5, 0.75])),
                )
            doc.add_page_break()
        else:
            doc.add_paragraph(f"No GeneBass PheWAS results found for {self.display_gene}.")
            doc.add_page_break()
        return doc

    def generate_labs_measurements_report(self, doc):
        phewas_file = self.ensure_phewas_file()
        doc.add_heading("Association study results of labs and measurements in UK Biobank cohort collected from GeneBass", level=2)
        if os.path.exists(phewas_file):
            phewas_df = pd.read_csv(phewas_file, sep=",")
            phewas_df = phewas_df.loc[phewas_df["Trait type"] == "continuous"]
            columns = ["Description", "Category", "Beta", "P-Value (Burden)"]
            display_df = phewas_df[columns].copy()
            display_df.columns = ["Phenotype", "Category", "Beta", "P-value"]
            display_df = display_df.nsmallest(50, "P-value")
            doc.add_paragraph()
            uth._add_table_to_doc(
                doc,
                display_df,
                uth._add_table_title(doc, f"Top 50 associations for {self.display_gene} in GeneBass"),
                ["Phenotype", "Category", "Beta", "P-value"],
                list(map(Inches, [2.75, 1.75, 0.75, 0.75])),
            )
            doc.add_page_break()
        else:
            doc.add_paragraph()
            doc.add_paragraph(f"No GeneBass labs and measurements results found for {self.display_gene}.")
            doc.add_page_break()
        return doc

    def generate_indication_specific_report(self, doc, keywords):
        phewas_file = self.ensure_phewas_file()
        doc.add_heading("Indication specific PheWAS results in UK Biobank cohort collected from GeneBass", level=2)
        doc.add_paragraph("The keywords searched to generate the report were: " + ", ".join(keywords))
        if os.path.exists(phewas_file):
            phewas_df = pd.read_csv(phewas_file, sep=",")
            escaped_keywords = [re.escape(k) for k in keywords if k]
            if not escaped_keywords:
                doc.add_paragraph("No valid indication keywords were provided.")
                doc.add_page_break()
                return doc
            keyword_pattern = "|".join(escaped_keywords)
            phewas_df = phewas_df.loc[phewas_df.Description.str.contains(keyword_pattern, case=False, na=False)]
            columns = ["Description", "Category", "Beta", "P-Value (Burden)"]
            display_df = phewas_df[columns].copy().sort_values("P-Value (Burden)").reset_index(drop=True)
            if len(display_df) == 0:
                doc.add_paragraph("No associations found for the specified indications. Please broaden the indication keywords and try again.")
                doc.add_page_break()
            else:
                display_df.columns = ["Description", "Category", "Beta", "P-value"]
                uth._add_table_to_doc(
                    doc,
                    display_df,
                    uth._add_table_title(doc, f"Indication-specific associations for {self.display_gene} in GeneBass"),
                    ["Description", "Category", "Beta", "P-value"],
                    list(map(Inches, [2.75, 1.75, 0.75, 0.75])),
                )
                doc.add_page_break()
        else:
            doc.add_paragraph(f"No GeneBass PheWAS results found for {self.display_gene}.")
            doc.add_page_break()
        return doc

    def get_carrier_info(self) -> dict:
        normalized_gene = self.input_gene.strip().upper()
        df = self.load_plof_summary()
        hits = df.loc[df["Gene"] == normalized_gene]

        if hits.empty:
            nhomz = None
            nhetz = None
        else:
            row = hits.iloc[0]
            nhomz = None if pd.isna(row["Homozygote_Count"]) else int(row["Homozygote_Count"])
            nhetz = None if pd.isna(row["Heterozygote_Count"]) else int(row["Heterozygote_Count"])

        return {
            "gene": normalized_gene,
            "nhetz": nhetz,
            "nhomz": nhomz,
        }


def run(playwright: Playwright, gene: str, gene_id: str, save_path: str) -> None:
    client = GeneBassClient(gene=gene, phewas_dir=save_path)
    client._approved_symbol = gene
    client._ensembl_id = gene_id
    client.run(playwright)


def clean_genebass_phewas(phewas_file):
    return GeneBassClient.clean_phewas(phewas_file)


def get_genebass_manhattan(gene, phewas_file):
    client = GeneBassClient(gene=gene, phewas_dir=os.path.dirname(phewas_file))
    client._approved_symbol = gene
    return client.get_manhattan(phewas_file)


def generate_genebass_clinical_report(doc, gene, phewas_dir, phewas_filename=""):
    return GeneBassClient(gene=gene, phewas_dir=phewas_dir, phewas_filename=phewas_filename).generate_clinical_report(doc)


def generate_genebass_labs_measurements_report(doc, gene, phewas_dir, phewas_filename=""):
    return GeneBassClient(gene=gene, phewas_dir=phewas_dir, phewas_filename=phewas_filename).generate_labs_measurements_report(doc)


def generate_genebass_indication_specific_report(doc, gene, phewas_dir, keywords, phewas_filename=""):
    return GeneBassClient(gene=gene, phewas_dir=phewas_dir, phewas_filename=phewas_filename).generate_indication_specific_report(doc, keywords)


def get_genebass_carrier_info(gene: str) -> dict:
    return GeneBassClient(gene=gene, phewas_dir="").get_carrier_info()


if __name__ == "__main__":
    from docx import Document

    doc = Document()
    gene = "TNNC1"
    phewas_dir = "/home/dbanerjee/deepro/rwe/data/test/"
    client = GeneBassClient(gene=gene, phewas_dir=phewas_dir)
    doc = client.generate_clinical_report(doc)
    doc = client.generate_labs_measurements_report(doc)
    doc = client.generate_indication_specific_report(doc, keywords=["hypertrophic_cardiomyopathy"])
    doc.save(f"/home/dbanerjee/deepro/rwe/data/test/{gene}_genebass_report.docx")
    print(get_genebass_carrier_info(gene))
