import os
import numpy as np
import pandas as pd
from tqdm import tqdm
import multiprocessing as mp
from scipy import stats
from docx import Document
from docx.shared import Inches

import rwe.utils.helpers as uth
from rwe.plots.variant_info import create_variant_frequency_plots
from rwe.plots.demographics import demographics_plot
from rwe.plots.clinical import manhattan
from rwe.plots.measurements import plot_measurements
from rwe.plots.surveys import plot_survey_questions

############### Variant Info and Demographics ###############
def get_individual_info(bucket, chrm, gene, zygosity="hetz"):
    person_df = pd.read_csv(f"{bucket}/data/rwe_info/raw/person.csv.gz")
    variant_df = pd.read_csv(f"{bucket}/data/rwe_info/genes/chr{chrm}/{gene}/lof_{zygosity}.csv.gz")
    person_df["cases"] = person_df.person_id.isin(variant_df[f"{zygosity}_samples"])
    person_df["group"] = np.where(person_df["cases"], "Cases", "Controls") 
    return person_df, variant_df

def generate_aou_variant_info_demographics_report(doc: Document, chrm: str, gene: str, zygosity: str ="hetz") -> Document:
    from rwe.parsers.aou.config import BUCKET
    person_df, variant_df = get_individual_info(BUCKET, chrm, gene, zygosity)
    if person_df.empty or variant_df.empty:
        doc.add_paragraph("No variant or individual level data provided to generate figures.")
        doc.add_page_break()
        return doc
    # 1) Variant information figure
    fig1, _axes1 = create_variant_frequency_plots(variant_df)
    fig1_path = uth._save_fig_to_tmp(fig1, basename="variant_information", dpi=300)
    doc.add_paragraph()  # spacing
    doc.add_picture(fig1_path, width=Inches(6.5))
    uth._add_figure_caption(doc, "Variant information summary (carrier counts by consequence; most/least frequent variants).")
    # 2) Demographics figure
    fig2, _axes2 = demographics_plot(person_df)
    fig2_path = uth._save_fig_to_tmp(fig2, basename="demographics", dpi=300)
    doc.add_paragraph()  # spacing
    doc.add_picture(fig2_path, width=Inches(4))
    uth._add_figure_caption(doc, "Demographics of cases vs controls (age, ancestry, sex at birth, ethnicity).")
    return doc


############### Clinical Records ###############
def clean_aou_phewas(phewas_file, version="1.2"):
    if version == "1.2":
        df = pd.read_csv(phewas_file)
        df = df.loc[df.converged==True]
    elif version =="X":
        df = pd.read_csv(phewas_file, sep="\t")
        df = df.loc[(df.ancestry == "all")&(df.converged==True)]
    return df

def get_aou_manhattan(df, gene):
    name_col = "phecode_string"
    p_col    = "p_value"
    odds_ratio_col = "odds_ratio"
    cat_col  = "phecode_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=2.8e-4,
        top_k_labels=5,
        title=f"{gene} PheWAS AoU",
    )
    return fig, plot_df

def generate_aou_clinical_report(doc, chrm, gene, zygosity):
    from rwe.parsers.aou.config import BUCKET
    phewas_file = f"{BUCKET}/data/phewas/results/chr{chrm}/{gene}_phewas.csv"
    doc.add_heading("Phenome-wide Association Study (PheWAS) in All of Us cohort", level=2)
    if uth._gcs_size(phewas_file, BUCKET) > 0: 
        df = clean_aou_phewas(phewas_file)
        fig, plot_df = get_aou_manhattan(df, gene)
        fig_path = uth._save_fig_to_tmp(fig, basename="aou_phewas", dpi=300)
        doc.add_paragraph()  # spacing
        doc.add_picture(fig_path, width=Inches(6.5))
        uth._add_figure_caption(doc, f"{gene} PheWAS results from All of Us.")
        doc.add_paragraph()  # spacing

        # Table: Top 10 significant hits (sorted by p-value)
        top_significant = plot_df.nsmallest(10, 'p_value')
        columns = ['phecode_string', 'phecode_category', 'odds_ratio', 'p_value']
        # Rename columns for display
        display_df = top_significant[columns].copy()
        display_df.columns = ['Phecode', 'Category', 'OR', 'P-value']
        uth._add_table_to_doc(doc, display_df, uth._add_table_title(doc, f"Top 10 significant associations for {gene} in AoU"),
                        ['Phecode', 'Category', 'OR', 'P-value'], list(map(Inches, [2.75, 1.75, 0.75, 0.75]))
                        )

        # 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 = ['Phecode', 'Category', 'OR', 'P-value']
            uth._add_table_to_doc(doc, display_df_neg, uth._add_table_title(doc, f"Top 10 protective associations for {gene} in AoU"),
                            ['Phecode', 'Category', 'OR', 'P-value'], list(map(Inches, [2.75, 1.75, 0.75, 0.75]))
                            )

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


############### Labs and Measurements ###############
def compare(df, measurement):
    mdf = df.loc[df.measurement==measurement].copy()
    ctrls = mdf.loc[mdf.cases==False].drop_duplicates(["person_id", "measurement_concept_id"])
    cases = mdf.loc[mdf.cases==True].drop_duplicates(["person_id", "measurement_concept_id"])
    x = pd.to_numeric(cases["median_value"], errors="coerce").dropna().to_numpy()
    y = pd.to_numeric(ctrls["median_value"], errors="coerce").dropna().to_numpy()
    caq1, caq2, caq3 = pd.Series(x).quantile([0.25, 0.5, 0.75])
    ctq1, ctq2, ctq3 = pd.Series(y).quantile([0.25, 0.5, 0.75])
    # guard against empty arrays
    if x.size<20 or y.size<20:
        return (measurement, np.nan, np.nan, len(cases), len(ctrls), np.nan, np.nan, np.nan)
    u_stat, p_mwu = stats.mannwhitneyu(x, y, alternative="two-sided")
    return (measurement, float(u_stat), float(p_mwu), int(len(cases)), int(len(ctrls)), caq1, caq2, caq3, ctq1, ctq2, ctq3)

_G_DF = None

def _init_worker(df):
    global _G_DF
    _G_DF = df

def _worker_compare(measurement):
    # uses global df set once per process
    return compare(_G_DF, measurement)

def run_parallel(df, measurements=None, n_jobs=None, chunksize=10):
    if measurements is None:
        measurements = sorted(df["measurement"].dropna().unique().tolist())
    if n_jobs is None:
        n_jobs = max(1, mp.cpu_count() - 1)
    with mp.Pool(processes=n_jobs, initializer=_init_worker, initargs=(df,)) as pool:
        rows = list(
            tqdm(
                pool.imap_unordered(_worker_compare, measurements, chunksize=chunksize),
                total=len(measurements),
                desc="Processing measurements"
            )
        )
    results_df = pd.DataFrame(
        rows, columns=["measurement", "u_stat", "p_mwu", "n_cases", "n_ctrls", "q1_case", "median_case", "q3_case", "q1_ctrl", "median_ctrl", "q3_ctrl"]
    )
    return results_df

def generate_aou_labs_measurements_report(doc: Document, chrm: str, gene: str, zygosity: str) -> Document:
    from rwe.parsers.aou.config import MEASUREMENT_GROUPS, BUCKET
    numerical_measurements_df = pd.read_parquet(f"{BUCKET}/data/rwe_info/processed/selected_numerical_measurements.parquet")
    variant_df = pd.read_csv(f"{BUCKET}/data/rwe_info/genes/chr{chrm}/{gene}/lof_{zygosity}.csv.gz")
    numerical_measurements_df["cases"] = numerical_measurements_df.person_id.isin(variant_df[f"{zygosity}_samples"])
    measurements = numerical_measurements_df["measurement"].unique().tolist()
    res_df = run_parallel(numerical_measurements_df, measurements=measurements, n_jobs=None, chunksize=20)
    top_ten_lab = res_df.nsmallest(10, "p_mwu").loc[:, ["measurement", "p_mwu", "u_stat", "n_cases", "n_ctrls"]]
    top_ten_lab.columns = ["Measurement", "P-value", "U statistic", "N cases", "N controls"]
    uth._add_table_to_doc(
        doc, top_ten_lab, uth._add_table_title(
            doc, f"Top 10 lab measurements for {gene} pLoF carriers in AoU"), 
            top_ten_lab.columns.tolist(), list(map(Inches, [3, 0.75, 0.75, 0.75, 0.75]))
    )
    doc.add_page_break()  # spacing
    for k,v in MEASUREMENT_GROUPS.items():
        f, a = plot_measurements(numerical_measurements_df, v, col="median_value", res_df=res_df)
        fig_path = uth._save_fig_to_tmp(f, basename=f"aou_measurements_{k}", dpi=300)
        doc.add_paragraph()  # spacing
        doc.add_heading(f"{k.capitalize} profile of {gene} pLoF carriers in All of Us cohort", level=2)
        doc.add_picture(fig_path, width=Inches(6.5))
        uth._add_figure_caption(doc, f"{gene} pLoF carrier {k} results from All of Us.")
        doc.add_page_break()  # spacing
    return doc

############### Surveys ###############
def clean_aou_surveys(df, survey_col="survey", question_col="question", answer_col="answer_category", zygosity="hetz"):
    df = df.copy()
    df[question_col] = df[question_col].str.replace(r"^.*?:\s*", "", regex=True, n=1)
    df[answer_col] = df[answer_col].str.replace(r"^.*?:\s*", "", regex=True, n=1)
    return df

def generate_aou_survey_report(doc: Document, chrm: str, gene: str, zygosity: str) -> Document:
    from rwe.parsers.aou.config import BUCKET
    survey_df = pd.read_parquet(f"{BUCKET}/data/rwe_info/processed/selected_surveys.parquet")
    variant_df = pd.read_csv(f"{BUCKET}/data/rwe_info/genes/chr{chrm}/{gene}/lof_{zygosity}.csv.gz")
    survey_df["cases"] = survey_df.person_id.isin(variant_df[f"{zygosity}_samples"])
    survey_df = clean_aou_surveys(survey_df, zygosity=zygosity)
    fig, ax = plot_survey_questions(survey_df)
    fig_path = uth._save_fig_to_tmp(fig, basename="aou_surveys", dpi=300)
    doc.add_paragraph()  # spacing
    doc.add_picture(fig_path, width=Inches(6.5))
    uth._add_figure_caption(doc, f"{gene} pLoF carrier survey results from All of Us.")
    doc.add_paragraph()  # spacing
    return doc

if __name__ == "__main__":
    from docx import Document
    doc = Document()
