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.
    """
    result = subprocess.run(
        command,
        shell=True,
        capture_output=True,
        text=True,
        check=False
    )
    if result.returncode != 0:
        logging.debug("Error: %s", result.stderr)
    else:
        logging.debug("Output: %s", result.stdout)


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, *exactly like Hao's R scripts*:

      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 already exist before calling.
      - By default: no --maf, no --keep, no --chr filter (Hao's behavior).
      - Optional --maf, --keep, and autosome filter.
    
    Args:
        plink2_path: Path to PLINK2.
        ped_file: PED path.
        map_file: MAP path.
        output_prefix: Output prefix (no extension).
        relax_mind_threshold: If True, skip per-sample missingness filter.
        maf_threshold: Minor allele frequency filter for SNPs; if None, skip --maf.
        sample_keep_path: Path to file with samples to keep; if None, skip --keep.
        autosomes_only: If True, restrict to autosomal chromosomes only.
    """
    # sanity: the files PLINK expects with --pedmap <prefix>
    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 to match Hao
    keep = f"--keep {sample_keep_path}" if sample_keep_path else ""        # usually none to match Hao
    chrflag = "--chr 1-19" if autosomes_only else ""                       # usually off to match Hao

    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"
    )

    # Exact PLINK command matching Hao’s workflow
    run_plink2(
        f"{plink2_path} --pedmap {output_prefix} "
        f"--make-bed --geno 0.1 {mind} {maf} {keep} {chrflag} --out {output_prefix}"
    )


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 directly from --pedmap <prefix>.
    Assumes <pedmap_prefix>.ped and <pedmap_prefix>.map exist.
    """
    run_plink2(
        f"{plink2_path} --pedmap {pedmap_prefix} --make-rel square --out {kin_prefix}"
    )


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.
    - PHENO: FID IID zscore value   (we REPLACE IID with FAM's IID)
    - FAM:   FID IID PID MID SEX PHE


    For each FID, counts must match (after drop_duplicates to avoid inflation).
    """
    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"
    )

    # De-dup identical PHENO rows to avoid replicated lines blowing up counts
    phe = phe.drop_duplicates(subset=["FID","zscore","value"], keep="first")

    out_chunks = []
    for fid, fam_grp in fam.groupby("FID", sort=False):
        phe_grp = phe[phe["FID"] == fid].copy()

        if len(phe_grp) != len(fam_grp):
            # Try stricter de-dup within FID, then fail if still mismatched
            phe_grp = phe_grp.drop_duplicates(subset=["FID","zscore","value"], keep="first")


        if len(phe_grp) != len(fam_grp):
            raise ValueError(
                f"PHENO vs FAM row-count mismatch for FID={fid}: 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)  # <-- critical: 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)