# NOTE: This version has been edited in an effort to provide more thorough debugging logs.
# plink_utils

import os
import subprocess
import logging
from typing import Optional


def run_plink2(command: str):
    """
    Run a shell command invoking PLINK2, capturing stdout/stderr for logging.

    On non-zero exit, raise RuntimeError with PLINK2's stderr so callers
    fail fast instead of hitting missing .fam/.bed later.
    """
    logging.info("[plink_utils:run_plink2] %s", command)
    result = subprocess.run(
        command,
        shell=True,
        capture_output=True,
        text=True,
        check=False,
    )

    if result.returncode != 0:
        stderr = (result.stderr or "").strip()
        stdout = (result.stdout or "").strip()
        logging.error(
            "[plink_utils:run_plink2] PLINK2 failed (code %d).\nSTDERR:\n%s\nSTDOUT:\n%s",
            result.returncode,
            stderr,
            stdout,
        )
        raise RuntimeError(
            f"PLINK2 failed (exit {result.returncode}).\n"
            f"Command: {command}\n"
            f"STDERR:\n{stderr}\n"
            f"STDOUT:\n{stdout}"
        )
    else:
        logging.debug("[plink_utils:run_plink2] Output:\n%s", (result.stdout or "").strip())


def _create_sorted_pgen_from_pedmap(
    plink2_path: str,
    ped_path: str,
    map_path: str,
    temp_prefix: str,
) -> None:
    """
    Fallback step for PLINK2 'split chromosome' errors.

    1) Take the already-written PED/MAP (per-measure, per-sex).
    2) Normalize variant ordering via --make-pgen --sort-vars.
    """
    logging.info(
        "[plink_utils:_create_sorted_pgen_from_pedmap] "
        "Making sorted PGEN from PED/MAP: ped=%s map=%s out=%s",
        ped_path,
        map_path,
        temp_prefix,
    )
    cmd = (
        f"{plink2_path} "
        f"--ped {ped_path} "
        f"--map {map_path} "
        f"--make-pgen --sort-vars "
        f"--out {temp_prefix}"
    )
    run_plink2(cmd)


def generate_bed_bim_fam(
    plink2_path: str,
    ped_file: str,
    map_file: str,
    output_prefix: str,
    relax_mind_threshold: bool = False,
    maf_threshold: Optional[float] = None,
    sample_keep_path: Optional[str] = None,
    autosomes_only: bool = False,
):
    """
    Generates BED/BIM/FAM from PED/MAP using PLINK2, matching Hao's R:

        plink2 --pedmap <prefix> --make-bed --geno 0.1 --mind 0.1 --out <prefix>

    Notes:
      - Primary path: exactly the above, via --pedmap <output_prefix>.
      - Fallback: if PLINK2 complains about 'split chromosome',
        we normalize variants via --make-pgen --sort-vars, then
        re-run as:

            plink2 --pfile <temp_prefix> --make-bed ...

      - We still expect <output_prefix>.ped/.map to exist before calling.
    """
    ped_expect = f"{output_prefix}.ped"
    map_expect = f"{output_prefix}.map"

    if not os.path.exists(ped_expect):
        raise FileNotFoundError(f"Missing PED for --pedmap: {ped_expect}")
    if not os.path.exists(map_expect):
        raise FileNotFoundError(f"Missing MAP for --pedmap: {map_expect}")

    mind = "" if relax_mind_threshold else "--mind 0.1"
    maf = f"--maf {maf_threshold}" if maf_threshold is not None else ""   # usually None
    keep = f"--keep {sample_keep_path}" if sample_keep_path else ""       # usually none
    chrflag = "--chr 1-19" if autosomes_only else ""                      # usually off

    logging.info(
        "[plink_utils:generate_bed_bim_fam] --pedmap %s -> BED with %s %s %s %s",
        output_prefix, mind or "no --mind", maf or "no --maf",
        keep or "no --keep", chrflag or "all chr",
    )

    # 1) Try Hao-style direct pedmap -> bed
    pedmap_cmd = (
        f"{plink2_path} --pedmap {output_prefix} "
        f"--make-bed --geno 0.1 {mind} {maf} {keep} {chrflag} --out {output_prefix}"
    )

    try:
        raise RuntimeError("split chromosome")  # Force fallback for testing
        run_plink2(pedmap_cmd)
    except RuntimeError as e:
        msg = str(e)
        # Only trigger fallback for the specific split-chromosome error
        if "split chromosome" not in msg:
            # it's some other PLINK2 error; re-raise
            raise

        logging.warning(
            "[plink_utils:generate_bed_bim_fam] PLINK2 reported a split chromosome. "
            "Falling back to make-pgen + sort-vars + pfile->bed pipeline."
        )

        # 2) Fallback: sorted PGEN pipeline
        temp_prefix = f"{output_prefix}_sorted"

        # Normalize variants from the per-measure PED/MAP (not the big reference files)
        _create_sorted_pgen_from_pedmap(
            plink2_path=plink2_path,
            ped_path=ped_expect,
            map_path=map_expect,
            temp_prefix=temp_prefix,
        )

        # Now convert sorted PGEN to BED/BIM/FAM with same filters
        pfile_cmd = (
            f"{plink2_path} --pfile {temp_prefix} "
            f"--make-bed --geno 0.1 {mind} {maf} {keep} {chrflag} "
            f"--out {output_prefix}"
        )
        run_plink2(pfile_cmd)

    # 3) Sanity-check that PLINK2 actually produced the expected BED/BIM/FAM
    fam_expect = f"{output_prefix}.fam"
    if not os.path.exists(fam_expect):
        raise FileNotFoundError(
            f"PLINK2 finished but FAM file not found: {fam_expect}. "
            "Check PLINK2 STDERR/STDOUT for details."
        )


def generate_bed_bim_fam_V0(
    plink2_path: str,
    ped_file: str,
    map_file: str,
    output_prefix: str,
    relax_mind_threshold: bool = False,
    maf_threshold: Optional[float] = None,
    sample_keep_path: Optional[str] = None,
    autosomes_only: bool = False,
):
    """
    Generates BED/BIM/FAM from PED/MAP using PLINK2, matching Hao's R:

        plink2 --pedmap <prefix> --make-bed --geno 0.1 --mind 0.1 --out <prefix>

    Notes:
      - We *ignore* ped_file/map_file here; PLINK reads <prefix>.ped/.map via --pedmap.
      - Ensure <output_prefix>.ped and <output_prefix>.map exist before calling.
      - Default: no --maf, no --keep, no --chr filter.
      - [Working] Matches Hao's UCLA results; however, on the combined grid, generates
        split chromosome errors due to unsorted variants.
    """
    ped_expect = f"{output_prefix}.ped"
    map_expect = f"{output_prefix}.map"

    if not os.path.exists(ped_expect):
        raise FileNotFoundError(f"Missing PED for --pedmap: {ped_expect}")
    if not os.path.exists(map_expect):
        raise FileNotFoundError(f"Missing MAP for --pedmap: {map_expect}")

    mind = "" if relax_mind_threshold else "--mind 0.1"
    maf = f"--maf {maf_threshold}" if maf_threshold is not None else ""   # usually None
    keep = f"--keep {sample_keep_path}" if sample_keep_path else ""       # usually none
    chrflag = "--chr 1-19" if autosomes_only else ""                      # usually off

    logging.info(
        "[plink_utils:generate_bed_bim_fam] --pedmap %s -> BED with %s %s %s %s",
        output_prefix, mind or "no --mind", maf or "no --maf",
        keep or "no --keep", chrflag or "all chr",
    )

    cmd = (
        f"{plink2_path} --pedmap {output_prefix} "
        f"--make-bed --geno 0.1 {mind} {maf} {keep} {chrflag} --out {output_prefix}"
    )
    run_plink2(cmd)

    # Sanity-check that PLINK2 actually produced the expected BED/BIM/FAM
    fam_expect = f"{output_prefix}.fam"
    if not os.path.exists(fam_expect):
        raise FileNotFoundError(
            f"PLINK2 finished but FAM file not found: {fam_expect}. "
            "Check PLINK2 STDERR/STDOUT for details."
        )


def calculate_kinship_matrix(
    plink2_path: str,
    input_prefix: str,
    output_prefix: str,
    sample_keep_path: Optional[str] = None,
):
    """
    Create PLINK .kin files from BED/BIM/FAM files.
    """
    keep = f"--keep {sample_keep_path}" if sample_keep_path else ""
    cmd = f"{plink2_path} --bfile {input_prefix} {keep} --make-rel square --out {output_prefix}"
    run_plink2(cmd)


def calculate_kinship_from_pedmap(
    plink2_path: str,
    pedmap_prefix: str,
    kin_prefix: str,
):
    """
    Compute kinship, preferring Hao-style --pedmap, with a safe fallback
    for split-chromosome errors.

    Normal case (UCLA1 etc.):
        plink2 --pedmap <pedmap_prefix> --make-rel square --out <kin_prefix>

    If PLINK2 complains that the temporary .pvar has a split chromosome,
    we do *not* die. Instead we:
        1) If <pedmap_prefix>.bed/.bim/.fam exist (generated earlier by
           generate_bed_bim_fam), use them:
               plink2 --bfile <pedmap_prefix> --make-rel square --out <kin_prefix>
    """
    # First: Hao-style pedmap route
    cmd = (
        f"{plink2_path} --pedmap {pedmap_prefix} "
        f"--make-rel square --out {kin_prefix}"
    )
    try:
        run_plink2(cmd)
        return
    except RuntimeError as e:
        msg = str(e)
        if "has a split chromosome" not in msg:
            # It's some other PLINK error; propagate it.
            raise

        logging.warning(
            "[plink_utils:calculate_kinship_from_pedmap] PLINK2 reported a split "
            "chromosome for %s; trying fallback via BED/BIM/FAM or sorted PGEN.",
            pedmap_prefix,
        )

    # Fallback 1: if BED/BIM/FAM already exist (from generate_bed_bim_fam),
    # use them. This keeps kinship on the same SNP set as association.
    bed = pedmap_prefix + ".bed"
    bim = pedmap_prefix + ".bim"
    fam = pedmap_prefix + ".fam"
    if os.path.exists(bed) and os.path.exists(bim) and os.path.exists(fam):
        logging.info(
            "[plink_utils:calculate_kinship_from_pedmap] Using existing BED/BIM/FAM "
            "via --bfile %s for kinship.",
            pedmap_prefix,
        )
        run_plink2(
            f"{plink2_path} --bfile {pedmap_prefix} "
            f"--make-rel square --out {kin_prefix}"
        )
    else:
        raise RuntimeError(
            "[plink_utils:calculate_kinship_from_pedmap]"
            "Fallback via BED/BIM/FAM not possible; files missing. "
        )


def calculate_kinship_from_pedmap_V0(
    plink2_path: str,
    pedmap_prefix: str,
    kin_prefix: str,
):
    """
    Compute kinship directly from --pedmap <prefix>.
    Assumes <pedmap_prefix>.ped and <pedmap_prefix>.map exist.

    Notes:
    - working for UCLA, however for combined grid generates split chromosome error.
    """
    cmd = f"{plink2_path} --pedmap {pedmap_prefix} --make-rel square --out {kin_prefix}"
    run_plink2(cmd)


def rewrite_pheno_ids_from_fam(pheno_path: str, fam_path: str, out_path: str) -> None:
    """
    Make PHENO rows match FAM in both order AND IID values (no de-duplication).
      - PHENO: FID IID zscore value   (we REPLACE IID with FAM's IID in FAM order)
      - FAM:   FID IID PID MID SEX PHE

    Counts must already match per FID if PED and PHENO were written together.
    """
    import pandas as pd

    fam = pd.read_csv(
        fam_path,
        sep=r"\s+",
        header=None,
        names=["FID", "IID", "PID", "MID", "SEX", "PHE"],
        engine="python",
    )
    phe = pd.read_csv(
        pheno_path,
        sep=r"\s+",
        header=None,
        names=["FID", "IID", "zscore", "value"],
        engine="python",
    )

    out_chunks = []
    phe_groups = {k: g for k, g in phe.groupby("FID", sort=False)}
    for fid, fam_grp in fam.groupby("FID", sort=False):
        if fid not in phe_groups:
            raise ValueError(f"FID present in FAM but missing in PHENO: {fid}")

        phe_grp = phe_groups[fid].copy()

        # Strict 1:1 check (DO NOT drop duplicates; we want exact replicate counts)
        if len(phe_grp) != len(fam_grp):
            raise ValueError(
                f"PHENO vs FAM row-count mismatch for FID={fid}: "
                f"pheno={len(phe_grp)} fam={len(fam_grp)}"
            )

        phe_grp = phe_grp.reset_index(drop=True)
        phe_grp["IID"] = fam_grp["IID"].reset_index(drop=True)  # copy FAM IIDs (with suffixes)
        out_chunks.append(phe_grp[["FID", "IID", "zscore", "value"]])

    out = pd.concat(out_chunks, axis=0)
    out.to_csv(out_path, sep=" ", header=False, index=False)


