#!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.bin_utils import (
    format_cli_block,
    format_rows,
    format_sections,
    print_cli_logo,
    setup_cli_logger,
)

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

LOGO_DESCRIPTION = "< Find tagging variants using pre-computed LD matrices >"


def format_run_summary(args, n_focal_snps=None):
    """
    Build a compact summary of the requested LD expansion operation.

    :param args: Parsed command-line arguments.
    :param n_focal_snps: Number of unique focal SNPs read from the input file.
    :return: A formatted multi-line string.
    """
    sections = [
        ("LD source", [("LD path", args.ld_path)]),
        (
            "Variant input",
            [
                ("SNP file", args.snp_file),
                ("SNP column", args.snp_column),
                ("Unique focal SNPs", n_focal_snps),
            ],
        ),
        (
            "Expansion",
            [
                ("R2 threshold", args.r2_threshold),
                ("Correlation threshold", round(float(np.sqrt(args.r2_threshold)), 2)),
            ],
        ),
        ("Output", [("Output file", args.output_file)]),
    ]

    return format_sections(sections)


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 focal_snps_by_chromosome(gdl, focal_snps):
    """
    Split focal SNPs by the LD matrices where they are present.

    :param gdl: A ``GWADataLoader`` object with LD matrices.
    :param focal_snps: Ordered focal SNP IDs supplied by the user.
    :return: Dictionary mapping chromosomes to matched focal SNP IDs.
    """
    matched = {}
    for chrom, ld_mat in gdl.ld.items():
        ld_snps = set(ld_mat.snps.astype(str))
        present = [snp for snp in focal_snps if snp in ld_snps]
        if present:
            matched[chrom] = present
    return matched


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.
    """
    print_cli_logo(mgp, LOGO_DESCRIPTION)
    args = parse_args()
    setup_cli_logger(logger, 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))
    logger.info(
        "%s",
        format_cli_block(
            "LD expansion request",
            format_run_summary(args, n_focal_snps=len(focal_snps)),
        ),
    )
    gdl = make_data_loader(args)
    focal_by_chrom = focal_snps_by_chromosome(gdl, focal_snps)
    n_matched_focal = sum(len(snps) for snps in focal_by_chrom.values())
    logger.info(
        "Matched %s of %s focal variants to %s LD matrices before expansion.",
        n_matched_focal,
        len(focal_snps),
        len(focal_by_chrom),
    )
    corr_threshold = float(np.sqrt(args.r2_threshold))

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

    output_snps = unique_preserving_order(expanded)
    write_variants(output_snps, args.output_file)
    logger.info(
        "LD expansion completed.\n%s",
        format_rows(
            [
                ("Input variants", len(focal_snps)),
                ("Output variants", len(output_snps)),
                ("Output file", args.output_file),
            ]
        ),
    )


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