# 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
)
from plinkformatter.generate_pheno_plink import extract_pheno_measure


from __future__ import annotations
import os
import logging
from typing import List


import numpy as np
import pandas as pd


from plinkformatter.plink_utils import (
    generate_bed_bim_fam,
    calculate_kinship_from_pedmap,
)
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:
    """
    Python port of Hao's generate_pheno_plink (non-DO branch):

      * Reads the existing PED/MAP from GCS (same ones Hao used).
      * Intersects phenotypes with PED by (FID=strain-cleaned, IID=animal_id).
      * For each (measnum, sex), writes:
          - {measnum}.{sex}.ped
          - {measnum}.{sex}.map
          - {measnum}.{sex}.pheno  (FID IID zscore value)

    Assumptions:
      - pheno has columns: 'strain', 'sex', 'measnum', 'value', 'zscore', 'animal_id'.
      - ped_file has FID in col 0, IID in col 1 (matching strain-cleaned and animal_id).
    """
    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()


    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 in pheno; 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 strain names to match PED FID
    ph = pheno.copy()
    ph["FID"] = ph["strain"].astype(str).str.replace(" ", "", regex=False)
    ph["IID"] = ph["animal_id"].astype(str)


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


    # Load PED (same as Hao's, from GCS)
    ped = pd.read_csv(ped_file, sep=r"\s+", header=None, dtype=str)


    # V1/V2 cleanup like Hao (remove '?')
    ped[0] = ped[0].str.replace("?", "", regex=False)
    ped[1] = ped[1].str.replace("?", "", regex=False)


    # Attach explicit FID/IID columns for merging
    ped = ped.rename(columns={0: "PED_FID", 1: "PED_IID"})
    ped_cols = ped.columns  # remember all columns, including PED_FID/PED_IID


    # Inner join: only animals present in both pheno and ped
    merged = ph.merge(
        ped,
        left_on=["FID", "IID"],
        right_on=["PED_FID", "PED_IID"],
        how="inner",
        sort=False,
    )
    if merged.empty:
        logging.info(
            "[generate_pheno_plink_fast] no overlap between pheno (FID/IID) and PED."
        )
        return merged


    # Load MAP once, patch rs == "." -> chr_bp 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),
    )


    # Now for each (measnum, sex), write ped/map/pheno
    for (measnum, sex), df in merged.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)


        # Extract PED columns (in original order)
        ped_part = df[ped_cols].copy()


        # In this ped_part:
        #   PED_FID = column "PED_FID" (FID)
        #   PED_IID = column "PED_IID" (IID)
        # The original PED columns 2..5 are PID, MID, SEX, PHE, etc.
        # We want to overwrite SEX (V5) and PHE (V6) like Hao:
        #   SEX (V5) = 2 for 'f', 1 for 'm'
        #   PHE (V6) = zscore
        # But since we renamed cols 0 and 1, the original 2.. are shifted.
        # PED structure was: 0:FID,1:IID,2:PID,3:MID,4:SEX,5:PHE,...
        # After renaming, the positional order is the same; we can use iloc.


        # Ensure columns are in original order for writing
        # (PED_FID, PED_IID, then the rest)
        ped_part = ped_part[
            ["PED_FID", "PED_IID"]
            + [c for c in ped_part.columns if c not in ("PED_FID", "PED_IID")]
        ]


        # Overwrite SEX (col index 4) and PHE (col index 5)
        # col 0: PED_FID, col 1: PED_IID, col 2: PID, col 3: MID, col 4: SEX, col 5: PHE
        sex_code = 2 if sex == "f" else 1
        ped_part.iloc[:, 4] = str(sex_code)
        ped_part.iloc[:, 5] = df["zscore"].astype(float).astype(str)


        # Write PED with TAB separator like Hao
        ped_part.to_csv(ped_out, sep="\t", header=False, index=False)


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


        # Write PHENO: FID IID zscore value, using PED_FID/PED_IID
        phe_df = pd.DataFrame(
            {
                "FID": ped_part["PED_FID"],
                "IID": ped_part["PED_IID"],
                "zscore": df["zscore"].astype(float),
                "value": df["value"].astype(float),
            }
        )
        phe_df.to_csv(phe_out, sep=" ", header=False, index=False)


        logging.info(
            "[generate_pheno_plink_fast] wrote %s, %s, %s",
            ped_out,
            map_out,
            phe_out,
        )


    return merged


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
            )

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

