import os
import re
from functools import lru_cache
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 matplotlib.pyplot as plt
from importlib.resources import as_file, files

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

_AOU_PTV_COUNTS_RESOURCE = "nsamples_ptv_gene.csv.gz"

############### Variant Info and Demographics ###############
def keep_most_del_plof(consequences):
    consequences =  set([c.strip() for c in consequences.split(",")])
    rank = {
        "frameshift_variant": 0,
        "stop_gained": 1,
        "splice_acceptor_variant": 2,
        "splice_donor_variant": 3,
        "start_lost": 4,
        "stop_lost": 5
    }
    return min(consequences, key=lambda x: rank.get(x, float('inf')))

def get_individual_info(bucket, chrm, gene, zygosity="hetz"):
    # TODO: need to change bucket folder based on biobanking updates
    person_df = pd.read_csv(f"{bucket}/data/rwe_info/raw/person.csv.gz")
    # TODO: need to change bucket folder based on biobanking updates
    variant_df = pd.read_csv(f"{bucket}/data/rwe_info/genes/chr{chrm}/{gene}/lof_{zygosity}.csv.gz")
    variant_df["consequence"] = variant_df.consequence.apply(keep_most_del_plof)
    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
    doc.add_paragraph()  # spacing
    doc.add_heading("Variant level information in All of Us cohort", level=2)
    fig1, ax_bar, top_show, bottom_show = create_variant_frequency_plots(variant_df)
    fig1_path = uth._save_fig_to_tmp(fig1, basename="variant_information", dpi=300)
    plt.close(fig1)  # close the figure to free memory
    doc.add_paragraph()  # spacing
    doc.add_picture(fig1_path, width=Inches(6.5))
    uth._add_figure_caption(doc, "Variant carrier counts by consequence.")
    doc.add_paragraph()  # spacing
    # Add tables for most/least frequent variants
    top_show.columns = ['Locus', 'Alleles', 'Consequence', 'Carriers (N)']
    bottom_show.columns = ['Locus', 'Alleles', 'Consequence', 'Carriers (N)']
    uth._add_table_to_doc(doc, top_show, uth._add_table_title(doc, f"Top 10 most frequent pLoF variants for {gene} in AoU"),
                    ['Locus', 'Alleles', 'Consequence', 'Carriers (N)'], list(map(Inches, [1.35, 2, 1.85, 0.8]))
                    )
    uth._add_table_to_doc(doc, bottom_show, uth._add_table_title(doc, f"Top 10 least frequent pLoF variants for {gene} in AoU"),
                    ['Locus', 'Alleles', 'Consequence', 'Carriers (N)'], list(map(Inches, [1.35, 2, 1.85, 0.8]))
                    )
    doc.add_page_break() 
    # 2) Demographics figure
    doc.add_paragraph()  # spacing
    doc.add_heading("Demographics information in All of Us cohort", level=2)
    fig2, _axes2 = demographics_plot(person_df)
    fig2_path = uth._save_fig_to_tmp(fig2, basename="demographics", dpi=300)
    plt.close(fig2)  # close the figure to free memory
    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).")
    doc.add_page_break() 
    return doc


############### Clinical Records ###############
def clean_aou_phewas(phewas_file, tool="regenie", version="1.2"):
    if tool=="phetk":
        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)]
        else:
            raise ValueError(f"Unsupported AoU PheWAS version: {version}")
    elif tool=="regenie":
        df = pd.read_csv(phewas_file)
        # keep plof strict mask only, aaf bin remove singletons, then merge by the most significant aaf_bin
        df = df.loc[(df["mask"]=="pLoF_strict")&(df.aaf_bin!="singleton")].copy()
        df = df.loc[df.groupby("pheno")["pvalue"].idxmin()].sort_values("pvalue")
        # from assets add the category of each phecode
        labels_resource = files("rwe").joinpath("assets/phecodeX_R_labels.csv")
        with as_file(labels_resource) as labels_path:
            labels_df = pd.read_csv(labels_path, usecols=["phenotype", "group"], dtype={"phenotype": str})
        labels_df = labels_df.drop_duplicates(subset=["phenotype"])
        pheno_to_group = labels_df.set_index("phenotype")["group"]
        df["category"] = df["pheno"].astype(str).map(pheno_to_group)
    else:
        raise ValueError(f"Unsupported AoU PheWAS tool: {tool}")
    return df

def get_aou_manhattan(df, gene):
    name_col = "description"
    p_col    = "pvalue"
    beta_col = "beta"
    cat_col  = "category"
    sig_p = 0.05/len(df)
    fig, ax, plot_df = manhattan(
        df,
        p_col=p_col,
        category_col=cat_col,
        label_col=name_col,
        beta_col=beta_col,
        sig_p=sig_p,
        top_k_labels=5,
        title=f"{gene} PheWAS AoU",
    )
    return fig, plot_df

def generate_aou_clinical_report(doc, chrm, gene, zygosity, output_dir):
    from rwe.parsers.aou.config import BUCKET
    phewas_file = f"{BUCKET}/data/associations/results/burden/plof/aou_phecodeX/chr{chrm}/{gene}.csv.gz"
    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)
        # Save table to output dir
        df.to_csv(os.path.join(output_dir, f"{gene}_aou_phewas.csv"), index=False)
        fig, plot_df = get_aou_manhattan(df, gene)
        fig_path = uth._save_fig_to_tmp(fig, basename="aou_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 All of Us.")
        doc.add_paragraph()  # spacing

        # Table: Top 10 significant hits (sorted by p-value)
        top_significant = plot_df.nsmallest(10, 'pvalue')
        columns = ['description', 'category', 'beta', 'pvalue']
        # Rename columns for display
        display_df = top_significant[columns].copy()
        display_df.columns = ['Phecode', 'Category', 'Effect', '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', 'Effect', 'P-value'], list(map(Inches, [2.75, 1.75, 0.75, 0.75]))
                        )

        # Table: Top 10 negative beta hits (beta < 0, sorted by p-value)
        negative_beta = plot_df[plot_df['beta'] < 0].nsmallest(10, 'pvalue')
        if len(negative_beta) > 0:
            display_df_neg = negative_beta[columns].copy()
            display_df_neg.columns = ['Phecode', 'Category', 'Effect', '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', 'Effect', '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

def generate_aou_indication_specific_report(doc, chrm, gene, output_dir, keywords):
    doc.add_heading("Indication specific PheWAS results in All of Us cohort", level=2)
    doc.add_paragraph("The keywords searched to generate the report were: " + ", ".join(keywords))
    if not os.path.exists(os.path.join(output_dir, f"{gene}_aou_phewas.csv")):
        doc.add_paragraph("No PheWAS results available to generate indication-specific report.")
        doc.add_page_break()
        return doc
    # TODO: update phewas results to change indication specific report
    phewas_df = pd.read_csv(os.path.join(output_dir, f"{gene}_aou_phewas.csv"))
    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', 'pvalue']
    display_df = phewas_df[columns].copy().sort_values("pvalue").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 = ['Phecode', 'Category', 'Effect', 'P-value']
        uth._add_table_to_doc(doc, display_df, uth._add_table_title(doc, f"Indication-specific associations for {gene} in AoU"),
                        ['Phecode', 'Category', 'Effect', 'P-value'], list(map(Inches, [2.75, 1.75, 0.75, 0.75]))
                        )
        doc.add_page_break()
    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, 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 compare_aou_labs_measurements_raw(numerical_measurements_df, measurements=None):
    res_df = run_parallel(numerical_measurements_df, measurements=measurements, n_jobs=None, chunksize=20)
    all_lab = res_df.sort_values("p_mwu").loc[:, ["measurement", "p_mwu", "u_stat", "n_cases", "n_ctrls"]]
    all_lab.columns = ["Measurement", "P-value", "U statistic", "N cases", "N controls"]
    return res_df, all_lab

def get_aou_labs_measurements_burden(chrm, gene):
    from rwe.parsers.aou.config import BUCKET, MEASUREMENT_NAMES_MAP
    res_df = pd.read_csv(f"{BUCKET}/data/associations/results/burden/plof/measurements/chrm{chrm}/{gene}.csv.gz")
    mn_map = {v:k for k,v in MEASUREMENT_NAMES_MAP.items()}
    res_df["pheno"] = res_df["pheno"].map(mn_map).fillna(res_df["pheno"])
    res_df = res_df.loc[res_df["mask"]!="singleton"]
    res_df = res_df.loc[res_df.groupby("pheno")["pvalue"].idxmin()]
    all_lab = res_df.sort_values("pvalue").loc[:, ["pheno", "pvalue", "beta", "se", "n"]]
    all_lab.columns = ["Measurement", "P-value", "Beta", "SE", "N"]
    return res_df, all_lab

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()
    # Burden tests from REGENIE
    res_df, all_lab = get_aou_labs_measurements_burden(chrm, gene)
    doc.add_heading("Association study results of labs and measurements in All of Us cohort", level=2)
    doc.add_paragraph()  # spacing
    uth._add_table_to_doc(
        doc, all_lab, uth._add_table_title(
            doc, f"Lab measurements for {gene} pLoF carriers in AoU"), 
            all_lab.columns.tolist(), list(map(Inches, [3.5, 0.75, 0.5, 0.5, 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, res_name_col="pheno", res_p_col="pvalue",
        )
        fig_path = uth._save_fig_to_tmp(f, basename=f"aou_measurements_{k}", dpi=300)
        plt.close(f)  # close the figure to free memory
        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)
    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} pLoF carrier survey results from All of Us.")
    doc.add_paragraph()  # spacing
    doc.add_page_break()
    return doc

############### Homozygous pLoF carriers ###############
def generate_aou_homozygous_lof_carriers_report_summary(doc, chrm, gene):
    from rwe.parsers.aou.config import BUCKET
    pdf, vdf = get_individual_info(BUCKET, chrm, gene, zygosity="homo")
    samples = set(pdf.loc[pdf.cases==True, "person_id"].values)
    # Demographics table for homozygous carriers
    doc.add_heading("Demographic information", level=2)
    info_df = pdf.loc[pdf.cases==True].groupby(
        ["sex_at_birth", "ancestry_pred_other"]
        ).agg({"person_id": "count"}).reset_index().rename(
            columns={"person_id": "n_individuals"}
    )
    doc.add_paragraph()  # spacing
    uth._add_table_to_doc(
        doc, info_df, uth._add_table_title(
            doc, f"Demographic information of {gene} homozygous pLoF carriers in AoU"), 
            info_df.columns.tolist(), list(map(Inches, [1, 1, 1]))
    )
    doc.add_paragraph()  # spacing
    # TODO: add clinical records, add labs and measurements
    # Survey table for homozygous carriers
    doc.add_heading("Survey responses", level=2)
    survey_df = pd.read_parquet(f"{BUCKET}/data/rwe_info/processed/selected_surveys.parquet")
    sdf = survey_df.loc[survey_df.person_id.isin(samples)].groupby("question").agg(
        {"answer": "value_counts"}
        ).rename(columns={"answer": "n_individuals"}).reset_index()
    doc.add_paragraph()  # spacing
    uth._add_table_to_doc(
        doc, sdf, uth._add_table_title(
            doc, f"Survey responses of {gene} homozygous pLoF carriers in AoU"), 
            sdf.columns.tolist(), list(map(Inches, [3, 2, 1]))
    )
    doc.add_paragraph()  # spacing
    doc.add_page_break()
    return doc

def get_person_information(sample_id, pdf):
    assert len(pdf.loc[pdf.person_id==sample_id])==1
    description_columns = ["sex_at_birth", "age_at_cdr", "ancestry_pred_other", "race", "ethnicity"]
    info_columns = ["has_ehr_data"]
    info_dict = pdf.loc[pdf.person_id==sample_id, description_columns + info_columns].iloc[0, :].to_dict()
    return info_dict

def get_variant_information(sample_id, vdf):
    variant_info_column = ["locus", "alleles", "gene", "consequence"]
    svdf = vdf.loc[vdf.homo_samples==sample_id].drop_duplicates(subset=variant_info_column).copy()
    svdf = svdf.loc[:, variant_info_column]
    info_dict = svdf.reset_index(drop=True).to_dict()
    return len(svdf), info_dict

def get_clinical_information(sample_id, phe_df, phe_info_df, ntop=10):
    phe_df = phe_df.loc[phe_df.person_id==sample_id].sort_values("count", ascending=False)
    nphenos = len(phe_df)
    phe_df = phe_df.head(ntop)
    phe_df = phe_df.merge(phe_info_df, left_on="phecode", right_on="PHECODE")
    return nphenos, phe_df.PHENOTYPE.unique()

def clean_numerical_measurements(df):
    df = df.loc[(df.unit_concept_id!=0)].copy()
    return df

def get_numerical_measurements(sample_id, nm_df, trait_range_dict):
    snm_df = nm_df.loc[
        (nm_df.person_id==sample_id)&(nm_df.measurement.isin(trait_range_dict.keys())),
        ["measurement", "unit", "median_value"]
    ].copy()
    snm_df["normal_range"] = snm_df.measurement.map(trait_range_dict)
    snm_df.columns = [c.capitalize().replace("_", " ") for c in snm_df.columns]
    return snm_df.sort_values("Measurement").rename(columns={"Median value": "Observed value"})

def get_survey_answers(sample_id, survey_df):
    sdf = survey_df.loc[survey_df.person_id==sample_id, ["question", "answer_category"]]
    sinfo_dict = dict(zip(sdf.question, sdf.answer_category))
    return sinfo_dict

def get_case_study_info_helper(doc, sample_id, pdf, vdf, phe_df, phe_info_df, nm_df, trait_range_dict, survey_df):
    ###--- Variant information and demographics ---###
    nvars, vinfo = get_variant_information(sample_id, vdf)
    case_study = doc.add_paragraph(
        f"An individual was identified in the AoU cohort who carries biallelic variant(s) "
        f"in {vinfo['gene'][0]} at loci {','.join(vinfo['locus'][i] for i in range(nvars))} with "
        f"{', '.join(vinfo['consequence'][i] for i in range(nvars))} consequence. "
    )
    person_info = get_person_information(sample_id, pdf)
    case_study.add_run(
        f"The individual is a {person_info['sex_at_birth']} of age {person_info['age_at_cdr']}. "
        f"Their genetic ancestry is predicted as {person_info['ancestry_pred_other']}. "
        f"They identify themselves of {person_info['race']} race and {person_info['ethnicity']} ethnicity. "
    )
    ###--- Clinical records ---###
    if person_info["has_ehr_data"]:
        case_study.add_run(f"EHR data of that individual is available. ")
        nphenos, phenos = get_clinical_information(sample_id, phe_df, phe_info_df)
        case_study.add_run(f"The individual has {nphenos} phenotypes in their clinical record. ")
        if nphenos>0:
            phenos = [f"{i+1}. {str(p).replace(',', ';')}" for i,p in enumerate(phenos)]
            case_study.add_run(f"The top phenotypes based on evidence include {', '.join(phenos)}. ")
    else:
        case_study.add_run(f"EHR data of that individual is not available.")
    ###--- Labs and measurements ---###
    snm_df = get_numerical_measurements(sample_id, nm_df, trait_range_dict)
    if len(snm_df)>0:
        uth._add_table_to_doc(
            doc, snm_df, uth._add_table_title(doc, f"Labs and measurement readings of the individual in AoU"),
                        ['Measurement', 'Unit', 'Observed value', 'Normal range'], list(map(Inches, [2.5, 1.25, 1, 1.25]))
                        )
    ###--- Surveys ---###
    sinfo_dict = get_survey_answers(sample_id, survey_df)
    if sinfo_dict:
        qa = [f"{i+1}. {q}: {a}" for i, (q,a) in enumerate(sinfo_dict.items())]
        case_study.add_run(f"The answers to the health related survey questions were: {', '.join(qa)}. ")
    return doc

def generate_aou_homozygous_lof_carriers_case_study(doc, chrm, gene, phe_info_file=""):
    from rwe.parsers.aou.config import BUCKET, NORMAL_TRAIT_RANGES
    # variant information and demographics
    pdf, vdf = get_individual_info(BUCKET, chrm, gene, zygosity="homo")
    samples = sorted(set(vdf.homo_samples.dropna().values))
    # clinical records
    phecode_df = pd.read_csv(f"{BUCKET}/data/phecodes/aou_phecodev1.2_counts.csv", dtype={"phecode": str})
    if not phe_info_file:
        assets_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'assets')
        phe_info_file = os.path.join(assets_dir, 'ICD_to_Phecode_mapping.csv')
    phe_info_df = pd.read_csv(phe_info_file, dtype=str)
    # labs and measurements
    numerical_measurements_df = pd.read_parquet(f"{BUCKET}/data/rwe_info/processed/selected_numerical_measurements.parquet")
    numerical_measurements_df = clean_numerical_measurements(numerical_measurements_df)
    # surveys
    survey_df = pd.read_parquet(f"{BUCKET}/data/rwe_info/processed/selected_surveys.parquet")
    survey_df = clean_aou_surveys(survey_df, zygosity="homo")
    doc.add_paragraph()  # spacing
    doc.add_heading("Case studies of homozygous LoF carriers in All of Us cohort", level=2)
    for i,sample in enumerate(samples):
        doc.add_paragraph()
        cs = doc.add_paragraph(f"Case Study {i+1}")
        cs.runs[0].bold = True
        doc = get_case_study_info_helper(doc, sample, pdf, vdf, phecode_df, phe_info_df, numerical_measurements_df, NORMAL_TRAIT_RANGES, survey_df)
    return doc

def generate_aou_homozygous_lof_carriers_report(doc: Document, chrm: str, gene: str) -> Document:
    from rwe.parsers.aou.config import BUCKET
    variant_df = pd.read_csv(f"{BUCKET}/data/rwe_info/genes/chr{chrm}/{gene}/lof_homo.csv.gz")
    if variant_df.empty:
        doc.add_paragraph(f"No AoU individuals with homozygous loss of function variants found for {gene}.")
        doc.add_page_break()  # spacing
    else:
        if variant_df["homo_samples"].dropna().nunique()>20:
            doc = generate_aou_homozygous_lof_carriers_report_summary(doc, chrm, gene)
        else:
            doc = generate_aou_homozygous_lof_carriers_case_study(doc, chrm, gene)
    doc.add_page_break()  # spacing
    return doc

@lru_cache(maxsize=1)
def _load_aou_ptv_counts() -> pd.DataFrame:
    resource = files("rwe") / "assets" / _AOU_PTV_COUNTS_RESOURCE
    with as_file(resource) as path:
        df = pd.read_csv(
            path,
            usecols=["gene", "nhetz_samples", "nhomz_samples", "chrm"],
        )
    df["gene"] = df["gene"].astype(str).str.upper()
    return df


def get_aou_carrier_info(gene: str) -> dict:
    normalized_gene = gene.strip().upper()
    df = _load_aou_ptv_counts()
    hits = df.loc[df["gene"] == normalized_gene]

    if hits.empty:
        nhetz = None
        nhomz = None
        chrm = None
    else:
        row = hits.iloc[0]
        nhetz = None if pd.isna(row["nhetz_samples"]) else int(row["nhetz_samples"])
        nhomz = None if pd.isna(row["nhomz_samples"]) else int(row["nhomz_samples"])
        chrm = None if pd.isna(row["chrm"]) else str(row["chrm"])

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

if __name__ == "__main__":
    gene = "ANGPTL3"
    info = get_aou_carrier_info(gene)
    print(info) 
