#!/usr/bin/env python3
"""
Drop-in, faster + lower-RAM replacements for PLINK text generation, with
end-to-end prepare that produces the *same files* your app expects.


Exports
-------
- generate_pheno_plink_fast(...): streaming writer for .ped/.map/.pheno
- fast_prepare_pylmm_inputs(...): end-to-end (also runs PLINK2 to make
  .bed/.bim/.fam + .kin.*), matching your current outputs.


Compatibility notes
-------------------
- Exactly matches your current semantics:
  * Strain normalization (strip spaces + '?' and uppercase)
  * Filter to strains present in PED
  * Stable sort by strain
  * Replicate PED rows per pheno row
  * Set PED col5 (SEX) = 2 for 'f', 1 for 'm'
  * Set PED col6 (PHE) = zscore
  * .pheno written as: FID IID zscore value (space-separated)
- Accepts measure IDs with suffixes (e.g. "131063_BXD"). We strip the
  suffix when looking up CSVs (so it reads 131063.csv) but keep the base
  id when naming outputs, just like your existing pipeline.
- Accepts a ped_pheno_field arg on fast_prepare_pylmm_inputs for API
  compatibility, but it is ignored (we always write zscore to PED col6
  to match your current behavior).
"""
from __future__ import annotations
import os
import io
import math
from collections import defaultdict
from typing import Dict, Iterable, List, Tuple, Optional


import numpy as np
import pandas as pd


# ----------------------------
# Helpers: normalization & IO
# ----------------------------


def _sanitize_strain(s: str) -> str:
    if not isinstance(s, str):
        s = str(s)
    return s.replace("?", "").replace(" ", "").upper()




def _base_id(mid) -> str:
    s = str(mid)
    return s.split("_", 1)[0]




def _read_map_sanitized(map_file: str) -> pd.DataFrame:
    """Read MAP and apply the same rs sanitization as your code.
    If rs==".", replace with f"chr_bp".
    """
    df = pd.read_csv(map_file, header=None, sep="\t")
    # columns: 0=chr, 1=rs, 2=cM, 3=bp38
    df[1] = np.where(df[1] == ".", df[0].astype(str) + "_" + df[3].astype(str), df[1].astype(str))
    return df




def _iter_ped_strain_offsets(ped_path: str) -> Dict[str, int]:
    """Single pass to collect byte offsets for each strain (FID) without storing lines.
    Returns { sanitized_FID: offset }.
    Assumes unique strain rows.
    """
    offsets: Dict[str, int] = {}
    with open(ped_path, "rb") as f:
        while True:
            pos = f.tell()
            line = f.readline()
            if not line:
                break
            # PED file appears TAB-delimited between fields; genotype pairs are within-field "A B".
            first_tab = line.find(b"\t")
            fid_bytes = (line.strip().split()[0] if first_tab <= 0 else line[:first_tab])
            name = _sanitize_strain(fid_bytes.decode(errors="replace"))
            if name and name not in offsets:
                offsets[name] = pos
    return offsets




def _read_ped_line_at_offset(ped_path: str, offset: int) -> str:
    with open(ped_path, "rb") as f:
        f.seek(offset)
        raw = f.readline()
    return raw.decode(errors="replace").rstrip("\n")




def _parse_tab_ped_line_to_parts(line: str) -> List[str]:
    """Split a TAB-delimited PED line into fields.
    Note: genotype pairs remain as one field each, like "2 2", "1 2".
    """
    parts = line.split("\t")
    if len(parts) <= 6:  # guard for space-delimited inputs
        parts = line.split()
    return parts




def _flatten_pairs_to_space_stream(parts: List[str], sex_flag: str, phe_value: Optional[float]) -> str:
    """Create a space-delimited PED record from TAB parts, updating SEX and PHE (col6).
    SEX: 2 for 'f', 1 for 'm'. PHE: zscore (None -> keep original).
    """
    if len(parts) < 7:
        raise ValueError("Malformed PED: expected >=7 columns (6 meta + genotypes)")


    meta = parts[:6]
    meta[4] = "2" if sex_flag == "f" else "1"
    if phe_value is not None:
        meta[5] = "-9" if (isinstance(phe_value, float) and math.isnan(phe_value)) else (str(phe_value))


    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]}")
    return out.getvalue()




# ----------------------------
# Core: fast generator
# ----------------------------


def generate_pheno_plink_fast(
    ped_file: str,
    map_file: str,
    pheno: pd.DataFrame,
    outdir: str,
    ncore: int = 1,
) -> pd.DataFrame:
    """
    Streaming, memory-safe replacement for generate_pheno_plink.


    Matches original semantics:
      - normalize strains, keep only f/m
      - filter strains by PED presence
      - stable sort by strain
      - replicate PED row per pheno row
      - set PED col5 (SEX) = 2 for 'f', 1 for 'm'
      - set PED col6 (PHE) = zscore
      - write .pheno with: FID IID zscore value


    Returns the filtered+sorted pheno (pheno_s) for parity with your API.
    """
    os.makedirs(outdir, exist_ok=True)


    # Validate / normalize pheno
    if pheno is None or pheno.empty:
        return pd.DataFrame()
    for col in ("strain", "sex", "measnum"):
        if col not in pheno.columns:
            raise ValueError("pheno must have columns: 'strain', 'sex', 'measnum' (plus 'zscore'/'value')")


    ph = pheno.copy()
    ph["strain"] = ph["strain"].astype(str).str.replace(" ", "", regex=False).str.upper()
    ph = ph[ph["sex"].isin(["f", "m"])].copy()
    if ph.empty:
        return ph


    # Index PED by strain (byte offsets)
    ped_index = _iter_ped_strain_offsets(ped_file)
    ped_strains = set(ped_index.keys())


    # Filter pheno to strains present in PED, stable sort by strain
    ph = ph[ph["strain"].isin(ped_strains)].sort_values(by="strain", kind="stable").reset_index(drop=True)
    if ph.empty:
        return ph


    # Load and sanitize MAP once, copy per group
    map_df = _read_map_sanitized(map_file)


    # Group by measnum/sex (preserve overall order)
    groups: Dict[Tuple[int, str], pd.DataFrame] = {}
    for (measnum, sex), df in ph.groupby(["measnum", "sex"], sort=False):
        groups[(int(measnum), str(sex))] = df


    for (measnum, sex), df in groups.items():
        # Write MAP
        map_out = os.path.join(outdir, f"{measnum}.{sex}.map")
        map_df.to_csv(map_out, sep="\t", index=False, header=False)


        # Build per-strain queues (zscore,value) in this group's order
        queues: Dict[str, List[Tuple[Optional[float], Optional[float]]]] = defaultdict(list)
        for row in df.itertuples(index=False):
            queues[str(row.strain)].append((getattr(row, "zscore", np.nan), getattr(row, "value", np.nan)))


        unique_strains = list(dict.fromkeys(df["strain"].tolist()))  # preserves order


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


        with open(ped_out, "w") as f_ped, open(pheno_out, "w") as f_ph:
            for strain in unique_strains:
                off = ped_index.get(strain)
                if off is None:
                    continue  # should not happen after filtering
                line = _read_ped_line_at_offset(ped_file, off)
                parts = _parse_tab_ped_line_to_parts(line)
                if len(parts) >= 2:
                    parts[0] = _sanitize_strain(parts[0])
                    parts[1] = _sanitize_strain(parts[1])


                # For each phenotyped row in this strain, replicate:
                for z, v in queues[strain]:
                    # PED: set SEX & PHE=zscore
                    space_line = _flatten_pairs_to_space_stream(parts, sex_flag=sex, phe_value=z)
                    f_ped.write(space_line + "\n")


                    # PHENO: FID IID zscore value
                    fid = iid = strain
                    z_out = "" if (z is None or (isinstance(z, float) and math.isnan(z))) else str(z)
                    v_out = "" if (v is None or (isinstance(v, float) and math.isnan(v))) else str(v)
                    f_ph.write(f"{fid} {iid} {z_out} {v_out}\n")


    return ph




# ------------------------------------------------------------------
# Fast end-to-end prepare (writes BED/BIM/FAM + KIN like your app)
# ------------------------------------------------------------------


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",  # accepted for API compatibility; ignored (we always write zscore to PED col6)
) -> None:
    """
    Drop-in replacement for prepare_pylmm_inputs that uses the streaming PED writer
    but produces the *same set of downstream PLINK outputs* as your current pipeline:
      - {base_id}.{sex}.ped/.map/.pheno
      - {base_id}.{sex}.bed/.bim/.fam
      - {base_id}.{sex}.kin.rel and .rel.id


    - base_id follows your current logic: split on first underscore (e.g., "131063_BXD" -> "131063").
    - Only runs PLINK for sexes that have corresponding .ped files.


    ped_pheno_field is accepted for compatibility but ignored (this function always writes zscore to PED col6).
    """
    os.makedirs(outdir, exist_ok=True)


    # 1) Extract phenotype using the measure_ids AS GIVEN (e.g., "131063_BXD" -> read 131063_BXD.csv)
    from plinkformatter.generate_pheno_plink import extract_pheno_measure
    pheno = extract_pheno_measure(measure_id_directory, measure_ids)


    if pheno is None or pheno.empty:
        return


    # 2) Generate .ped/.map/.pheno with streaming (same semantics; PED col6=zscore)
    generate_pheno_plink_fast(
        ped_file=ped_file,
        map_file=map_file,
        pheno=pheno,
        outdir=outdir,
        ncore=ncore,
    )


    # 3) For each requested measure (original ids), derive base_id and produce BED/BIM/FAM + KIN
    from plinkformatter.plink_utils import generate_bed_bim_fam, calculate_kinship_matrix


    for measure_id in measure_ids:
        base_id = _base_id(measure_id)
        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):
                # No data for this sex → match your current behavior
                continue


            # Create BED/BIM/FAM via PLINK2 (same as your code)
            generate_bed_bim_fam(
                plink2_path=plink2_path,
                ped_file=ped_path,
                map_file=map_path,
                output_prefix=out_prefix,
            )


            # Kinship matrix
            kin_prefix = os.path.join(outdir, f"{base_id}.{sex}.kin")
            calculate_kinship_matrix(
                plink2_path=plink2_path,
                input_prefix=out_prefix,
                output_prefix=kin_prefix,
            )




__all__ = [
    "generate_pheno_plink_fast",
    "fast_prepare_pylmm_inputs",
]






