# 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:
    """
    Streaming version, aligned with Hao's intent:

      * Does NOT load PED into memory.
      * Builds a dict of PED line offsets keyed by IID (per-animal ID).
      * For each phenotype row (measnum, sex, animal_id), pulls the PED line
        for that IID, overwrites SEX and PHE (V6) with sex + zscore, and
        writes a new per-(measnum, sex) PED.
      * Writes per-(measnum, sex) PHENO: FID IID zscore value
        with FID/IID taken from the PED line, not invented.

    Assumptions:
      - pheno has columns: 'strain', 'sex', 'measnum', 'value'
      - plus: 'zscore' and 'animal_id' (animal_id matches PED IID).
    """
    os.makedirs(outdir, exist_ok=True)
    if pheno is None or pheno.empty:
        logging.info("[generate_pheno_plink_fast] empty pheno; nothing to do.")
        return pd.DataFrame()

    # Basic sanity
    required_cols = ["strain", "sex", "measnum", "value"]
    for col in required_cols:
        if col not in pheno.columns:
            raise ValueError(
                f"pheno must have at least columns: {required_cols} (+ 'zscore', 'animal_id'). "
                f"Missing: {col}"
            )

    if "zscore" not in pheno.columns:
        logging.info("[generate_pheno_plink_fast] no 'zscore' column; filling with NaN.")
        pheno = pheno.copy()
        pheno["zscore"] = np.nan

    if "animal_id" not in pheno.columns:
        raise ValueError(
            "pheno must contain 'animal_id' column matching PED IID (e.g. '129S1/SvImJ_212')."
        )

    # Clean up pheno
    ph = pheno.copy()
    ph["animal_id"] = ph["animal_id"].astype(str)
    ph["sex"] = ph["sex"].astype(str)

    # Keep only m/f
    ph = ph[ph["sex"].isin(["m", "f"])].copy()
    if ph.empty:
        logging.info("[generate_pheno_plink_fast] no m/f rows after filtering.")
        return ph

    # Deduplicate at animal level per measure/sex
    ph = ph.drop_duplicates(
        subset=["measnum", "sex", "animal_id"], keep="first"
    ).reset_index(drop=True)

    # Build PED offsets keyed by IID (second column), streaming
    ped_offsets: Dict[str, int] = {}
    with open(ped_file, "rb") as f:
        while True:
            pos = f.tell()
            line = f.readline()
            if not line:
                break
            s = line.strip()
            if not s:
                continue
            toks = s.split()
            if len(toks) < 2:
                continue
            # FID = toks[0], IID = toks[1]
            iid = toks[1].decode(errors="replace")
            iid_clean = iid.replace("?", "").replace(" ", "")
            if iid_clean and iid_clean not in ped_offsets:
                ped_offsets[iid_clean] = pos

    if not ped_offsets:
        logging.info("[generate_pheno_plink_fast] no IID offsets indexed from PED.")
        return pd.DataFrame()

    # Normalize animal_id to the same key space
    ph["iid_key"] = (
        ph["animal_id"]
        .astype(str)
        .str.replace("?", "", regex=False)
        .str.replace(" ", "", regex=False)
    )

    # Keep only phenotypes whose IID exists in PED
    ph = ph[ph["iid_key"].isin(ped_offsets.keys())].reset_index(drop=True)
    if ph.empty:
        logging.info(
            "[generate_pheno_plink_fast] no overlap between pheno animal_id and PED IID."
        )
        return ph

    # Load MAP once (this is cheap compared to PED) and patch rs == "." like Hao
    map_df = pd.read_csv(map_file, header=None, sep=r"\s+")
    map_df[1] = np.where(
        map_df[1] == ".",
        map_df[0].astype(str) + "_" + map_df[3].astype(str),
        map_df[1].astype(str),
    )

    # Open PED once; we'll seek() into it for each animal
    with open(ped_file, "rb") as f_ped:
        # Group by (measnum, sex) and write outputs
        for (measnum, sex), df in ph.groupby(["measnum", "sex"], sort=False):
            measnum = int(measnum)
            sex = str(sex)

            ped_out = os.path.join(outdir, f"{measnum}.{sex}.ped")
            map_out = os.path.join(outdir, f"{measnum}.{sex}.map")
            phe_out = os.path.join(outdir, f"{measnum}.{sex}.pheno")

            df = df.reset_index(drop=True)

            # Write MAP once per group
            map_df.to_csv(map_out, sep="\t", index=False, header=False)

            with open(ped_out, "w", encoding="utf-8") as f_out_ped, open(
                phe_out, "w", encoding="utf-8"
            ) as f_out_ph:
                for _, r in df.iterrows():
                    key = r["iid_key"]
                    if key not in ped_offsets:
                        # shouldn't happen after filtering, but be defensive
                        continue

                    pos = ped_offsets[key]
                    f_ped.seek(pos)
                    raw = f_ped.readline().decode(errors="replace").rstrip("\n")

                    if not raw:
                        continue

                    parts = raw.split()
                    if len(parts) < 7:
                        raise ValueError(
                            f"Malformed PED line for IID {key!r}: expected >=7 columns, got {len(parts)}"
                        )

                    # Clean FID/IID like Hao
                    fid = parts[0].replace("?", "").replace(" ", "")
                    iid = parts[1].replace("?", "").replace(" ", "")

                    # Parse zscore/value
                    z = r.get("zscore", np.nan)
                    val = r.get("value", np.nan)
                    try:
                        z = float(z)
                    except Exception:
                        z = math.nan
                    try:
                        val = float(val)
                    except Exception:
                        val = math.nan

                    z_str = f"{z}" if math.isfinite(z) else "-9"
                    val_str = f"{val}" if math.isfinite(val) else "-9"

                    # Meta fields: FID, IID, PID, MID, SEX, PHE
                    meta = parts[:6]
                    meta[0] = fid
                    meta[1] = iid
                    meta[2] = "0"  # PID
                    meta[3] = "0"  # MID
                    meta[4] = "2" if sex == "f" else "1"  # SEX
                    meta[5] = z_str  # PHE = zscore

                    # Reconstruct PED row
                    out = io.StringIO()
                    out.write(" ".join(meta))
                    for gp in parts[6:]:
                        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_out_ped.write(out.getvalue() + "\n")

                    # PHENO: FID IID zscore value
                    f_out_ph.write(f"{fid} {iid} {z_str} {val_str}\n")

            logging.info(
                "[generate_pheno_plink_fast] wrote %s, %s, %s",
                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"),
            )

