import os
import subprocess
import json
import pandas as pd
from biobanking.association.helpers import (
    list_gcs_files,
    save_files_in_regenie_fmt, rint_normalization
)


# REGENIE helper files preparation
def prepare_covariate_file():
    from biobanking.utils.config import AllofUs
    aou = AllofUs()
    cov_df = pd.read_csv(f"{aou.bucket}/data/phenotypes/individuals.csv.gz")
    cov_df = cov_df.loc[
        (cov_df.duplicates==False)&
        (cov_df.qc_failed==False)&
        (cov_df.sex_at_birth.notna())
    ].copy()
    # assign sex
    cov_df["genetic_sex"] = cov_df.sex_at_birth.astype(int)
    # calculate age related covariates
    cov_df["age_2"] = cov_df.age**2
    cov_df["age_sex"] = cov_df.age*cov_df.genetic_sex
    # save in regenie fomrat
    save_file = f"{aou.bucket}/data/associations/phenotypes/covariates.tsv.gz"
    save_files_in_regenie_fmt(
        cov_df, save_file, phenos=[], 
        covariates=[
            'genetic_sex', 'age', 'age_2', 'age_sex', 'ancestry_pred_other', 
            'genetic_pca1', 'genetic_pca2', 'genetic_pca3', 'genetic_pca4', 
            'genetic_pca5', 'genetic_pca6', 'genetic_pca7', 'genetic_pca8', 
            'genetic_pca9', 'genetic_pca10'
            ]
        )
    return

def prepare_continous_traits(filename, traits=[], id_col="person_id"):
    from biobanking.utils.config import AllofUs
    aou = AllofUs()
    cov_df = pd.read_csv(f"{aou.bucket}/data/phenotypes/individuals.csv.gz")
    pheno_df = pd.read_csv(f"{aou.bucket}/data/phenotypes/measurements.csv.gz")
    if not traits:
        traits = [c for c in pheno_df.columns if c!=id_col]
    pheno_df = pheno_df.merge(cov_df, on=id_col)
    for trait in traits:
        pheno_df[f"{trait}"] = pheno_df.groupby(
            ["ancestry_pred_other", "sex_at_birth"]
            )[trait].transform(rint_normalization)
    save_file = f"{aou.bucket}/data/associations/phenotypes/{filename}.tsv.gz"
    save_files_in_regenie_fmt(pheno_df, save_file, phenos=traits, covariates=[])
    return

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 create_regenie_burden_files_helper(burden_df):
    # need to add chr prefix to variants  as in all of us chr stays
    burden_df["variants"] = burden_df.locus + ":" + burden_df.alleles.str.replace("_", ":")
    # create annotation df
    annot_df = burden_df.loc[:, ["variants", "gene", "consequence"]]
    # this annotates the same locus for the same gene with the most severe consequence
    annot_df =  annot_df.groupby(["variants", "gene"]).agg({"consequence": lambda x: ",".join(x)}).reset_index()
    annot_df["annotation"] = annot_df.consequence.apply(keep_most_del_plof)
    # Only keep required columns
    annot_df = annot_df.loc[:, ["variants", "gene", "annotation"]]
    # this gets rid of duplicates due to transcripts of same gene with same consequence
    annot_df = annot_df.dropna().drop_duplicates()
    # create set list df
    set_df = annot_df.groupby("gene").agg({"variants": lambda x: ",".join(x)})
    set_df[["chrm", "location"]] = set_df.variants.apply(lambda x: pd.Series(dict(zip(["chrm", "location"], x.split(",")[0].split(":")[:2]))))
    set_df = set_df.reset_index().loc[:, ["gene", "chrm", "location", "variants"]]
    # create aaf df
    aaf_df = burden_df.loc[:, ["variants", "maf"]]
    aaf_df = aaf_df.dropna().drop_duplicates()
    # create masks
    mask_names = ["pLoF_strict", "pLoF_lenient"]
    categories =  [
        "frameshift_variant,stop_gained,splice_acceptor_variant,splice_donor_variant",
        "frameshift_variant,stop_gained,splice_acceptor_variant,splice_donor_variant,start_lost,stop_lost"]
    mask_df = pd.DataFrame({"mask_name": mask_names, "categories": categories})
    return annot_df, set_df, aaf_df, mask_df

def create_regenie_burden_files(burden_type="plof"):
    from biobanking.utils.config import AllofUs
    aou = AllofUs()
    for chrm in list(range(1, 23))+["X"]:
        exome_df = pd.read_csv(f"{aou.bucket}/data/exome/annot/chr{chrm}_{burden_type}.tsv.gz", sep="\t")
        annot_df, set_df, aaf_df, mask_df = create_regenie_burden_files_helper(exome_df)
        annot_df_name = f"chr{chrm}_annotations_{burden_type}.tsv.gz"
        set_df_name = f"chr{chrm}_sets_{burden_type}.tsv.gz"
        aaf_df_name = f"chr{chrm}_aafs_{burden_type}.tsv.gz"
        mask_df_name = f"chr{chrm}_masks_{burden_type}.tsv.gz"
        proj_dir = f"{aou.bucket}/data/associations/burden/"

        for df, name in zip(
            [annot_df, set_df, aaf_df, mask_df],
            [annot_df_name, set_df_name, aaf_df_name, mask_df_name]
        ):
            df.to_csv(os.path.join(proj_dir, name), sep='\t', index=False, header=False)
    return


# REGENIE job submission
def create_input_dict(phase_prefix, run_step1=True, isbinary=False, burden_type="plof", step1_workflow_id="", step1_glob_id=""):
    from biobanking.utils.config import AllofUs
    aou = AllofUs()
    input_dict = dict()
    input_dict["regenie_workflow.run_step1"] = run_step1
    input_dict["regenie_workflow.isbinary"] = isbinary
    input_dict["regenie_workflow.phenotype_file"] = f"{aou.bucket}/data/associations/phenotypes/{phase_prefix}.tsv.gz"
    input_dict["regenie_workflow.covariate_file"] = f"{aou.bucket}/data/associations/phenotypes/covariates.tsv.gz"
    # Check phenotypes directly from the phenotype file columns and load them
    pheno_df =  pd.read_csv(f"{aou.bucket}/data/associations/phenotypes/{phase_prefix}.tsv.gz", sep="\t", nrows=0)
    phenos = [c for c in pheno_df.columns if not c in ["FID", "IID"]]

    input_dict["regenie_workflow.phenoColList"] = phenos
    input_dict["regenie_workflow.phenoColString"] = ",".join(phenos)
    input_dict["regenie_workflow.bed_file"] = f"{aou.bucket}/data/array/arrays.bed"
    input_dict["regenie_workflow.bim_file"] = f"{aou.bucket}/data/array/arrays.bim"
    input_dict["regenie_workflow.fam_file"] = f"{aou.bucket}/data/array/arrays.fam"
    input_dict["regenie_workflow.keep_variant_file"] = f"{aou.bucket}/data/array/qc/final_array_snps_GRCh38_qc_pass_pruned.prune.in"
    input_dict["regenie_workflow.keep_sample_file"] = f"{aou.bucket}/data/array/qc/final_array_snps_GRCh38_qc_pass.id"
    input_dict["regenie_workflow.prefix"] = phase_prefix

    input_dict["regenie_workflow.chrs"] = [f"{i}" for i in range(1, 23)]
    input_dict["regenie_workflow.bgen_files"] = [f"{aou.bucket}/data/bgen/chr{i}.bgen" for i in range(1, 23)]
    input_dict["regenie_workflow.sample_files"] = [f"{aou.bucket}/data/bgen/chr{i}.sample" for i in range(1, 23)]
    input_dict["regenie_workflow.index_files"] = [f"{aou.bucket}/data/bgen/chr{i}.bgen.bgi" for i in range(1, 23)]
    input_dict["regenie_workflow.set_files"] = [f"{aou.bucket}/data/associations/burden/chr{i}_sets_{burden_type}.tsv.gz" for i in range(1, 23)]
    input_dict["regenie_workflow.anno_files"] = [f"{aou.bucket}/data/associations/burden/chr{i}_annotations_{burden_type}.tsv.gz" for i in range(1, 23)]
    input_dict["regenie_workflow.aaf_files"] = [f"{aou.bucket}/data/associations/burden/chr{i}_aafs_{burden_type}.tsv.gz" for i in range(1, 23)]
    input_dict["regenie_workflow.mask_files"] = [f"{aou.bucket}/data/associations/burden/chr{i}_masks_{burden_type}.tsv.gz" for i in range(1, 23)]

    if not run_step1:
        assert step1_workflow_id
        assert step1_glob_id

        pred_file = f"{aou.bucket}/cromwell-execution/regenie_workflow/{step1_workflow_id}/call-regenie_step1/{phase_prefix}_pheno_pred.list"
        input_dict["regenie_workflow.pred_file"] = pred_file

        loco_file_prefix = f"cromwell-execution/regenie_workflow/{step1_workflow_id}/call-regenie_step1/{step1_glob_id}"
        input_dict["regenie_workflow.loco_files"] = [f for f in list_gcs_files(aou.bucket, loco_file_prefix) if f.endswith("loco")]

    with open(f"{phase_prefix}.json", "w") as f:
        json.dump(input_dict, f, indent=4)

    return


def read_cromshell_output(result):
    lines = result.stdout.strip().splitlines()
    for i, line in enumerate(lines):
        if line == "{":
            start = i
        elif line == "}":
            end = i

    submission_info = json.loads("".join(lines[start : end + 1]))
    workflow_id = submission_info.get("id")
    return workflow_id


def track_workflow_id(phase_prefix, workflow_json, result):
    if result.returncode == 0:
        with open(workflow_json, "r") as f:
            workflow_dict = json.load(f)
        workflow_id = read_cromshell_output(result)

        if workflow_id:
            workflow_dict[phase_prefix] = workflow_id
            with open(workflow_json, "w") as f:
                json.dump(workflow_dict, f, indent=4)
        else:
            print(f"Warning: No workflow ID found in JSON output for {phase_prefix}. File not updated.")
            print(result.stdout)

    else:
        print(f"Submission failed for {phase_prefix}:")
        print(result.stderr)
    return

def submit_and_track_workflow(phase_prefix, workflow_json="workflow_ids.json"):
    result = subprocess.run(
        ["cromshell", "submit", "regenie.wdl", f"{phase_prefix}.json"],
        capture_output=True,
        text=True
    )
    if not os.path.exists(workflow_json):
        with open(workflow_json, "w") as f:
            json.dump(dict(), f)
    track_workflow_id(phase_prefix, workflow_json, result)
    return
