#!python

"""
Compute Linkage-Disequilibrium (LD) matrices and store in Zarr array format
----------------------------

This is a commandline script that facilitates the computation of LD matrices
from genotype data stored in plink BED format. The script supports various
estimators for computing the LD matrix, including windowed, shrinkage, block,
and sample-based estimators. The script outputs the computed LD matrices in
Zarr array format, which is a compressed, parallelized, and scalable format
for storing large numerical arrays.

Usage:

    python -m mgp_compute_ld --bfile <bed_file> --estimator <estimator> --output-dir <output_dir>

For larger genotype matrices, we recommend using `plink1.9` as a backend to compute the LD matrices.
You can do that by specifying the `--backend` parameter:

    python -m mgp_compute_ld --bfile <bed_file> --estimator <estimator> --output-dir <output_dir> --backend plink

"""

import argparse
import logging
import os.path as osp
import time
from datetime import timedelta

import magenpy as mgp
from magenpy.utils.bin_utils import (
    format_cli_block,
    format_rows,
    format_sections,
    print_cli_logo,
    setup_cli_logger,
)
from magenpy.utils.system_utils import valid_url

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

LOGO_DESCRIPTION = "< Compute LD matrix and store in Zarr format >"


def validate_args(args):
    """
    Validate estimator-specific command-line arguments.

    :param args: Parsed command-line arguments.
    :raises ValueError: If required estimator inputs are missing.
    :raises FileNotFoundError: If an LD block file path does not exist.
    """
    if args.estimator == "windowed":
        if (
            args.ld_window is None
            and args.ld_window_kb is None
            and args.ld_window_cm is None
        ):
            raise ValueError(
                "For the windowed estimator, provide --ld-window, "
                "--ld-window-kb, or --ld-window-cm."
            )
    elif args.estimator == "block":
        if args.ld_blocks is None:
            raise ValueError(
                "The block LD estimator requires an LD blocks file via --ld-blocks."
            )
        if not osp.isfile(args.ld_blocks) and not valid_url(args.ld_blocks):
            raise FileNotFoundError("The LD blocks file does not exist.")
    elif args.estimator == "shrinkage":
        if args.genmap_ne is None:
            raise ValueError("The shrinkage estimator requires --genmap-Ne.")
        if args.genmap_ss is None:
            raise ValueError("The shrinkage estimator requires --genmap-sample-size.")


def get_ld_kwargs(args):
    """
    Convert estimator-specific command-line arguments to LD estimator keywords.

    :param args: Parsed command-line arguments.
    :return: A dictionary of keyword arguments for ``GWADataLoader.compute_ld``.
    """
    ld_kwargs = {}

    if args.estimator == "windowed":
        if args.ld_window is not None:
            ld_kwargs["window_size"] = args.ld_window
        if args.ld_window_kb is not None:
            ld_kwargs["kb_window_size"] = args.ld_window_kb
        if args.ld_window_cm is not None:
            ld_kwargs["cm_window_size"] = args.ld_window_cm
    elif args.estimator == "block":
        ld_kwargs["ld_blocks_file"] = args.ld_blocks
    elif args.estimator == "shrinkage":
        ld_kwargs["genetic_map_ne"] = args.genmap_ne
        ld_kwargs["genetic_map_sample_size"] = args.genmap_ss
        if args.shrink_cutoff is not None:
            ld_kwargs["threshold"] = args.shrink_cutoff

    return ld_kwargs


def parse_metadata(metadata):
    """
    Parse comma-separated ``key=value`` metadata supplied by the user.

    :param metadata: A comma-separated string of metadata key-value pairs.
    :return: A dictionary of parsed metadata.
    :raises ValueError: If any metadata entry is not formatted as ``key=value``.
    """
    if metadata is None:
        return {}

    parsed_metadata = {}
    for entry in metadata.split(","):
        entry = entry.strip()
        if len(entry) < 1:
            continue
        if "=" not in entry:
            raise ValueError(f"Metadata entry '{entry}' is not formatted as key=value.")
        key, value = entry.split("=", 1)
        parsed_metadata[key.strip()] = value.strip()

    return parsed_metadata


def format_run_summary(args, ld_kwargs, metadata):
    """
    Build a compact summary of the requested LD computation.

    :param args: Parsed command-line arguments.
    :param ld_kwargs: Estimator keyword arguments.
    :param metadata: Parsed metadata dictionary.
    :return: A formatted multi-line string.
    """
    filter_rows = [
        ("Keep samples", args.keep_file),
        ("Extract variants", args.extract_file),
        ("Minimum MAF", args.min_maf),
        ("Minimum MAC", args.min_mac),
    ]

    sections = [
        (
            "Estimator",
            [
                ("Method", args.estimator),
                ("Parameters", ld_kwargs if ld_kwargs else "default"),
            ],
        ),
        (
            "Genotype source",
            [
                ("BED path", args.bed_file),
                ("Backend", args.backend),
                ("Genome build", args.genome_build),
            ],
        ),
        ("Filters", filter_rows),
        (
            "Storage",
            [
                ("Output directory", args.output_dir),
                ("Temporary directory", args.temp_dir),
                ("Storage dtype", args.storage_dtype),
                ("Compressor", args.compressor),
                ("Compression level", args.compression_level),
                ("Spectral properties", args.compute_spectral),
            ],
        ),
        ("Metadata", [("Attributes", metadata if metadata else "Date only")]),
    ]

    return format_sections(sections)


def set_ld_metadata(ld_matrices, metadata):
    """
    Store metadata attributes on each computed LD matrix.

    :param ld_matrices: Dictionary of chromosome to ``LDMatrix`` objects.
    :param metadata: Metadata dictionary to store.
    """
    if "Date" not in metadata:
        metadata["Date"] = time.strftime("%Y-%m-%d")

    for ld_mat in ld_matrices.values():
        for key, value in metadata.items():
            ld_mat.set_store_attr(key, value)


print_cli_logo(mgp, LOGO_DESCRIPTION)

# --------- Options ---------

parser = argparse.ArgumentParser(
    description="""
    Commandline arguments for LD matrix computation and storage
"""
)

# General options:
parser.add_argument(
    "--estimator",
    dest="estimator",
    type=str,
    default="windowed",
    choices={"windowed", "shrinkage", "block", "sample"},
    help="The LD estimator (windowed, shrinkage, block, sample)",
)
parser.add_argument(
    "--bfile",
    dest="bed_file",
    type=str,
    required=True,
    help="The path to a plink BED file.",
)
parser.add_argument(
    "--keep",
    dest="keep_file",
    type=str,
    help="A plink-style keep file to select a subset of individuals to compute the LD matrices.",
)
parser.add_argument(
    "--extract",
    dest="extract_file",
    type=str,
    help="A plink-style extract file to select a subset of SNPs to compute the LD matrix for.",
)
parser.add_argument(
    "--backend",
    dest="backend",
    type=str,
    default="magenpy",
    choices={"magenpy", "bed-reader", "xarray", "plink"},
    help="The backend software used to compute the Linkage-Disequilibrium between variants.",
)
parser.add_argument(
    "--temp-dir",
    dest="temp_dir",
    type=str,
    default="temp",
    help="The temporary directory where we store intermediate files.",
)
parser.add_argument(
    "--output-dir",
    dest="output_dir",
    type=str,
    required=True,
    help="The output directory where the Zarr formatted LD matrices will be stored.",
)
parser.add_argument(
    "--min-maf",
    dest="min_maf",
    type=float,
    help="The minimum minor allele frequency for variants included in the LD matrix.",
)
parser.add_argument(
    "--min-mac",
    dest="min_mac",
    type=float,
    help="The minimum minor allele count for variants included in the LD matrix.",
)

# Metadata / reproducibility options:
parser.add_argument(
    "--genome-build",
    dest="genome_build",
    type=str,
    help="The genome build for the genotype data (recommend storing as metadata).",
)
parser.add_argument(
    "--metadata",
    dest="metadata",
    type=str,
    help="A comma-separated string with metadata keys and values. This is used to store "
    "information about the genotype data from which the LD matrix was computed, such as "
    "the biobank/samples, cohort characteristics (e.g. ancestry), etc. Keys and values "
    "should be separated by =, such that inputs are in the form of:"
    "--metadata Biobank=UKB,Ancestry=EUR,Date=April2024",
)

# Argument for the float precision:
parser.add_argument(
    "--storage-dtype",
    dest="storage_dtype",
    type=str,
    default="int16",
    help="The data type for the entries of the LD matrix.",
    choices={"float32", "float64", "int16", "int8"},
)

# Other options:
parser.add_argument(
    "--compute-spectral-properties",
    dest="compute_spectral",
    default=False,
    action="store_true",
    help="Compute and store the spectral properties of the "
    "LD matrix (e.g. eigenvalues, eigenvectors).",
)

# Add arguments for the compressor:
parser.add_argument(
    "--compressor",
    dest="compressor",
    type=str,
    default="zstd",
    help="The compressor name or compression algorithm to use for the LD matrix.",
    choices={"lz4", "zstd", "gzip", "zlib"},
)

parser.add_argument(
    "--compression-level",
    dest="compression_level",
    type=int,
    default=7,
    help="The compression level to use for the entries of the LD matrix (1-9).",
)

# Options for the various LD estimators:

# For the windowed estimator:
parser.add_argument(
    "--ld-window",
    dest="ld_window",
    type=int,
    help="Maximum number of neighboring SNPs to consider when computing LD.",
)
parser.add_argument(
    "--ld-window-kb",
    dest="ld_window_kb",
    type=float,
    help="Maximum distance (in kilobases) between pairs of variants when computing LD.",
)
parser.add_argument(
    "--ld-window-cm",
    dest="ld_window_cm",
    type=float,
    help="Maximum distance (in centi Morgan) between pairs of variants when computing LD.",
)

# For the block estimator:
parser.add_argument(
    "--ld-blocks",
    dest="ld_blocks",
    type=str,
    help="Path to the file with the LD block boundaries, "
    "in LDetect format (e.g. chr start stop, tab-separated)",
)

# For the shrinkage estimator:
parser.add_argument(
    "--genmap-Ne",
    dest="genmap_ne",
    type=int,
    help="The effective population size for the population from which the genetic map was derived.",
)
parser.add_argument(
    "--genmap-sample-size",
    dest="genmap_ss",
    type=int,
    help="The sample size for the dataset used to infer the genetic map.",
)
parser.add_argument(
    "--shrinkage-cutoff",
    dest="shrink_cutoff",
    type=float,
    help="The cutoff value below which we assume that the correlation between variants is zero.",
)
parser.add_argument(
    "--log-level",
    dest="log_level",
    type=str,
    default="INFO",
    choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
    help="Logging level.",
)

args = parser.parse_args()
setup_cli_logger(logger, args.log_level)
validate_args(args)
ld_kwargs = get_ld_kwargs(args)
metadata = parse_metadata(args.metadata)

logger.info(
    "%s",
    format_cli_block(
        "LD computation request", format_run_summary(args, ld_kwargs, metadata)
    ),
)

start_time = time.time()

logger.info("Reading genotype data with GWADataLoader...")
gdl = mgp.GWADataLoader(
    bed_files=args.bed_file,
    keep_file=args.keep_file,
    extract_file=args.extract_file,
    min_maf=args.min_maf,
    min_mac=args.min_mac,
    backend=args.backend,
    temp_dir=args.temp_dir,
    genome_build=args.genome_build,
)

logger.info("Computing LD matrices...")
gdl.compute_ld(
    args.estimator,
    args.output_dir,
    dtype=args.storage_dtype,
    compressor_name=args.compressor,
    compression_level=args.compression_level,
    compute_spectral_properties=args.compute_spectral,
    **ld_kwargs,
)
ld_matrices = gdl.get_ld_matrices()

logger.info("Storing LD metadata...")
set_ld_metadata(ld_matrices, metadata)

gdl.cleanup()

end_time = time.time()
chromosomes = ", ".join(map(str, sorted(ld_matrices)))
logger.info(
    "LD computation completed\n%s",
    format_rows(
        [
            ("Output directory", args.output_dir),
            ("Chromosomes", chromosomes),
            ("Total runtime", timedelta(seconds=end_time - start_time)),
        ]
    ),
)
