# generate_pheno_plink_fast.py
from __future__ import annotations
import os
import io
import math
import logging
from typing import Dict, List

import numpy as np
import pandas as pd

from plinkformatter.plink_utils import (
    generate_bed_bim_fam,
    calculate_kinship_matrix,
    calculate_kinship_from_pedmap,
    rewrite_pheno_ids_from_fam
)
from plinkformatter.generate_pheno_plink import extract_pheno_measure




def generate_pheno_plink_fast(
    ped_file: str,
    map_file: str,
    pheno: pd.DataFrame,
    outdir: str,
    ncore: int = 1,
) -> pd.DataFrame:
    """
    Build per-(measnum, sex) PLINK inputs (PED/MAP/PHENO) mirroring Hao's R:
      - replicate-level rows
      - FID = STRAIN
      - IID = STRAIN (duplicates allowed; PLINK will auto-suffix in FAM)
      - PID/MID = 0
      - SEX = 2 if 'f' else 1
      - PHE (col6) = zscore (or -9 if missing)
      - PHENO file: 'FID IID zscore value' (space-delimited)


    Notes
    -----
    * Exact IID alignment to FAM is handled later by rewriting PHENO IIDs from .fam,
      after BED/BIM/FAM creation (see instructions below).
    * MAP rsids "." are replaced with "chr_bp" like Hao.
    """
    os.makedirs(outdir, exist_ok=True)
    if pheno is None or pheno.empty:
        return pd.DataFrame()


    # Required columns
    need = ("strain", "sex", "measnum", "value")
    missing = [c for c in need if c not in pheno.columns]
    if missing:
        raise ValueError(f"pheno missing required columns: {missing} (need {list(need)})")


    # Hao removes spaces from strain in pheno; preserve case
    ph = pheno.copy()
    ph["strain"] = ph["strain"].astype(str).str.replace(" ", "", regex=False)


    # only f/m
    ph = ph[ph["sex"].isin(["f", "m"])].copy()
    if ph.empty:
        return ph


    # De-duplicate replicates conservatively.
    # Prefer unique animals if available, else fall back to value/zscore signature.
    dedup_keys = [c for c in ["strain", "sex", "measnum", "animal_id"] if c in ph.columns]
    if dedup_keys:
        ph = ph.drop_duplicates(subset=dedup_keys, keep="first")
    else:
        sig_cols = [c for c in ["strain", "sex", "measnum", "zscore", "value"] if c in ph.columns]
        ph = ph.drop_duplicates(subset=sig_cols, keep="first")


    if ph.empty:
        return ph


    # Build strain->byte offset index from reference PED (first occurrence per strain)
    ped_offsets: Dict[str, int] = {}
    with open(ped_file, "rb") as f:
        while True:
            pos = f.tell()
            line = f.readline()
            if not line:
                break
            # robust FID extraction (tab or space)
            first_tab = line.find(b"\t")
            fid_bytes = (line.strip().split()[0] if first_tab <= 0 else line[:first_tab])
            # Hao removes "?" artifacts; also strip spaces just in case
            name = fid_bytes.decode(errors="replace").replace("?", "").replace(" ", "")
            if name and name not in ped_offsets:
                ped_offsets[name] = pos


    # Keep only strains present in reference PED
    ped_strains = set(ped_offsets.keys())
    ph = ph[ph["strain"].isin(ped_strains)].reset_index(drop=True)
    if ph.empty:
        return ph


    # Read/sanitize MAP; replace '.' rsids with 'chr_bp'
    map_df = pd.read_csv(map_file, header=None, sep=r"\s+", engine="python")
    # columns: 0=chr, 1=rs, 2=cM, 3=bp
    map_df[1] = np.where(
        map_df[1].astype(str) == ".",
        map_df[0].astype(str) + "_" + map_df[3].astype(str),
        map_df[1].astype(str),
    )


    # Ensure zscore column exists (PHE col6 == zscore in PED)
    if "zscore" not in ph.columns:
        logging.info("[generate_pheno_plink_fast] 'zscore' missing in pheno; filling NaN (becomes -9 in PED/PHENO).")
        ph["zscore"] = np.nan


    # Write groups
    for (measnum, sex), df in ph.groupby(["measnum", "sex"], sort=False):
        measnum = int(measnum)
        sex = str(sex)


        # MAP (space-delimited to mirror Hao)
        map_out = os.path.join(outdir, f"{measnum}.{sex}.map")
        map_df.to_csv(map_out, sep=" ", index=False, header=False)


        # PED / PHENO
        ped_out = os.path.join(outdir, f"{measnum}.{sex}.ped")
        phe_out = os.path.join(outdir, f"{measnum}.{sex}.pheno")


        # Stable strain order; preserve original replicate order within strain
        df = df.sort_values(["strain"], kind="stable").reset_index(drop=True)


        with open(ped_out, "w", encoding="utf-8") as f_ped, open(phe_out, "w", encoding="utf-8") as f_ph:
            for strain, sdf in df.groupby("strain", sort=False):
                # Load reference PED line for this strain
                with open(ped_file, "rb") as fp:
                    fp.seek(ped_offsets[strain])
                    raw = fp.readline().decode(errors="replace").rstrip("\n")


                # Split line robustly
                parts = raw.split("\t")
                if len(parts) <= 6:
                    parts = raw.split()
                if len(parts) < 7:
                    raise ValueError("Malformed PED: need >=7 columns (6 meta + genotypes)")


                # Normalize FID/IID text like Hao (strip '?' artifacts and spaces)
                parts[0] = parts[0].replace("?", "").replace(" ", "")
                parts[1] = parts[1].replace("?", "").replace(" ", "")


                # Emit one row per replicate
                for _, r in sdf.iterrows():
                    # zscore/value coercion
                    z = r.get("zscore", np.nan)
                    v = r.get("value", np.nan)
                    try:
                        z = float(z)
                    except Exception:
                        z = np.nan
                    try:
                        v = float(v)
                    except Exception:
                        v = np.nan


                    # meta (FID/IID = strain, PID/MID=0, SEX=2/1, PHE=zscore or -9)
                    meta = parts[:6]
                    meta[0] = strain
                    meta[1] = strain
                    meta[2] = "0"
                    meta[3] = "0"
                    meta[4] = "2" if sex == "f" else "1"
                    meta[5] = f"{z}" if math.isfinite(z) else "-9"


                    # Reconstruct PED with space-delimited genotype pairs
                    out = io.StringIO()
                    out.write(" ".join(meta))
                    for gp in parts[6:]:
                        a_b = gp.split(" ")
                        if len(a_b) != 2:
                            a_b = gp.split()
                            if len(a_b) != 2:
                                raise ValueError(f"Genotype pair not splitable into two alleles: {gp!r}")
                        out.write(f" {a_b[0]} {a_b[1]}")
                    f_ped.write(out.getvalue() + "\n")


                    # PHENO row: FID IID zscore value  (IID will be rewritten to match .fam later)
                    f_ph.write(
                        f"{strain} {strain} "
                        f"{(z if math.isfinite(z) else -9)} "
                        f"{(v if math.isfinite(v) else -9)}\n"
                    )


        logging.info(f"[generate_pheno_plink_fast] wrote {ped_out}, {map_out}, {phe_out}")


    return ph








def fast_prepare_pylmm_inputs(
    ped_file: str,
    map_file: str,
    measure_id_directory: str,
    measure_ids: List,
    outdir: str,
    ncore: int,
    plink2_path: str,
    *,
    ped_pheno_field: str = "zscore",
) -> None:
    """
    Orchestrate extraction and PLINK file generation **like Hao**:

      1) Extract phenotype rows for requested measure_ids.
      2) Generate per-(measnum, sex):
           - {base}.{sex}.map
           - {base}.{sex}.ped   (replicate-level rows; FID=IID=STRAIN; PED V6=zscore)
           - {base}.{sex}.pheno (FID IID zscore value)
      3) Convert PED/MAP -> BED/BIM/FAM with: --geno 0.1 --mind 0.1
         (NO --maf; NO --keep).
      4) Compute kinship with: --make-rel square.
    """
    os.makedirs(outdir, exist_ok=True)

    pheno = extract_pheno_measure(measure_id_directory, measure_ids)
    if pheno is None or pheno.empty:
        logging.info("[fast_prepare_pylmm_inputs] no phenotype rows extracted; nothing to do.")
        return

    used = generate_pheno_plink_fast(
        ped_file=ped_file,
        map_file=map_file,
        pheno=pheno,
        outdir=outdir,
        ncore=ncore,
    )
    if used is None or used.empty:
        logging.info("[fast_prepare_pylmm_inputs] no usable phenotypes after PED intersection; nothing to do.")
        return

    for measure_id in measure_ids:
        base_id = str(measure_id).split("_", 1)[0]
        for sex in ("f", "m"):
            ped_path   = os.path.join(outdir, f"{base_id}.{sex}.ped")
            map_path   = os.path.join(outdir, f"{base_id}.{sex}.map")
            out_prefix = os.path.join(outdir, f"{base_id}.{sex}")

            if not (os.path.exists(ped_path) and os.path.exists(map_path)):
                continue

            logging.info(f"[fast_prepare_pylmm_inputs] make BED/BIM/FAM for {base_id}.{sex}")
            generate_bed_bim_fam(
                plink2_path=plink2_path,
                ped_file=ped_path,
                map_file=map_path,
                output_prefix=out_prefix,
                relax_mind_threshold=False,
                maf_threshold=None,     # match Hao: no --maf
                sample_keep_path=None,  # no --keep
                autosomes_only=False,   # match Hao: no --chr constraint
            )

            # >>> rewrite PHENO IIDs to EXACTLY match FAM (order + suffixes)
            fam_path   = f"{out_prefix}.fam"
            pheno_path = os.path.join(outdir, f"{base_id}.{sex}.pheno")
            rewrite_pheno_ids_from_fam(pheno_path, fam_path, pheno_path)

            logging.info(f"[fast_prepare_pylmm_inputs] compute kinship for {base_id}.{sex}")
            # calculate_kinship_matrix(
            #     plink2_path=plink2_path,
            #     input_prefix=out_prefix,
            #     output_prefix=os.path.join(outdir, f"{base_id}.{sex}.kin"),
            #     sample_keep_path=None,
            # )

            # Alternatively, compute kinship directly from PED/MAP:
            calculate_kinship_from_pedmap(
                plink2_path=plink2_path,
                pedmap_prefix=out_prefix,
                kin_prefix=os.path.join(outdir, f"{base_id}.{sex}.kin"),
            )

