#!python

"""
Expand a SNP list with LD neighbors from pre-computed magenpy LD matrices.
"""

import argparse
import logging
import os
import sys

import numpy as np
import pandas as pd

import magenpy as mgp
from magenpy.parsers.misc_parsers import read_snp_filter_file
from magenpy.utils.system_utils import setup_logger

logger = logging.getLogger("magenpy.cli.mgp_expand_ld")


def read_snp_file(path, column=0):
    """
    Read a plink-style SNP filter file.
    """
    return (
        pd.Series(read_snp_filter_file(path, snp_id_col=column))
        .dropna()
        .astype(str)
        .tolist()
    )


def unique_preserving_order(values):
    """
    Remove duplicates while preserving input order.
    """
    return list(dict.fromkeys(values))


def make_data_loader(args):
    """
    Load LD matrices through GWADataLoader.
    """
    gdl = mgp.GWADataLoader(ld_store_files=args.ld_path)
    if gdl.ld is None or len(gdl.ld) < 1:
        raise FileNotFoundError(f"No LD matrices were found at: {args.ld_path}")
    return gdl


def expand_chromosome(ld_mat, focal_snps, corr_threshold):
    """
    Return SNPs tagged by focal SNPs in one LD matrix.
    """
    ld_snps = ld_mat.snps.astype(str)
    snp_to_idx = {snp: i for i, snp in enumerate(ld_snps)}
    focal_indices = [snp_to_idx[snp] for snp in focal_snps if snp in snp_to_idx]
    if not focal_indices:
        return []

    expanded = ld_mat.find_tagging_variants(
        np.asarray(focal_indices, dtype=np.int32),
        threshold=corr_threshold,
        return_value="snps",
    ).astype(str)
    return [snp for snp in ld_snps if snp in set(expanded)]


def write_variants(snps, output_file):
    """
    Write expanded SNP IDs to a tab-delimited output table.
    """
    os.makedirs(os.path.dirname(os.path.abspath(output_file)), exist_ok=True)
    pd.DataFrame({"SNP": snps}).to_csv(output_file, sep="\t", index=False)


def parse_args():
    parser = argparse.ArgumentParser(
        description="Expand a SNP list with LD neighbors from pre-computed magenpy LD matrices."
    )
    parser.add_argument(
        "-l", "--ld", dest="ld_path", required=True, help="LD matrix path or glob."
    )
    parser.add_argument(
        "--snp-file", required=True, help="Plink-style SNP filter file."
    )
    parser.add_argument(
        "--snp-column", type=int, default=0, help="Zero-based SNP column index."
    )
    parser.add_argument("--r2-threshold", type=float, default=0.1, help="Default: 0.1.")
    parser.add_argument("--output-file", required=True, help="Output table path.")
    parser.add_argument(
        "--log-level",
        default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
        help="Logging level.",
    )
    return parser.parse_args()


def main():
    """
    Parse CLI inputs, expand focal SNPs to LD neighbors, and write SNP IDs.
    """
    logo = mgp.make_ascii_logo(
        desc="< Expand variants using pre-computed LD matrices >",
        left_padding=10,
    )
    print(f"\n{logo}\n", flush=True)

    args = parse_args()
    setup_logger(modules=["magenpy"], log_level=args.log_level)

    if not 0.0 < args.r2_threshold <= 1.0:
        raise ValueError("--r2-threshold must be in (0, 1].")

    focal_snps = unique_preserving_order(read_snp_file(args.snp_file, args.snp_column))
    gdl = make_data_loader(args)
    corr_threshold = float(np.sqrt(args.r2_threshold))

    expanded = list(focal_snps)
    for _, ld_mat in sorted(gdl.ld.items(), key=lambda x: str(x[0])):
        expanded.extend(expand_chromosome(ld_mat, focal_snps, corr_threshold))
        ld_mat.release()

    output_snps = unique_preserving_order(expanded)
    write_variants(output_snps, args.output_file)
    logger.info("Wrote %s expanded variants to %s", len(output_snps), args.output_file)


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        logger.error("Error: %s", e)
        sys.exit(1)
