#!/usr/bin/env python3

"""
Simulate Polygenic Traits using Complex Genetic Architectures
----------------------------

This is a commandline script that facilitates the simulation of complex traits
using a linear additive model with heterogeneous genetic architectures. The script
supports simulating phenotypes with different heritabilities, levels of polygenicity,
and genetic architectures. The script outputs the simulated phenotypes in a tabular
format that can be used for downstream analyses.

The script can simulate both quantitative and case-control traits. For case-control
traits, the script requires the specification of the prevalence of cases in the population.

The script requires access to genotype data in PLINK BED format.

Usage:

    python -m magenpy_simulate --bfile <bed_files> --h2 <h2> --prop-causal <p> --output-file <output_file>

"""

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

import numpy as np

import magenpy as mgp
from magenpy.simulation.PhenotypeSimulator import PhenotypeSimulator
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 get_filenames, makedir

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

LOGO_DESCRIPTION = "< Simulate complex quantitative or case-control traits >"


print_cli_logo(mgp, LOGO_DESCRIPTION)


def format_run_summary(args, bed_files, pi, d):
    """
    Build a compact summary of the requested phenotype simulation.

    :param args: Parsed command-line arguments.
    :param bed_files: BED files identified from the user's path or glob.
    :param pi: Mixture proportions for the simulation architecture.
    :param d: Variance multipliers for the simulation architecture.
    :return: A formatted multi-line string.
    """
    sections = [
        (
            "Simulation",
            [
                ("Likelihood", args.likelihood),
                ("Heritability", args.h2),
                ("Mixing proportions", pi),
                ("Variance multipliers", d),
                ("Prevalence", args.prevalence if args.likelihood == "binomial" else None),
                ("Seed", args.seed),
            ],
        ),
        (
            "Genotype source",
            [
                ("BED path", args.bed_file),
                ("Resolved BED files", len(bed_files)),
                ("Backend", args.backend),
            ],
        ),
        (
            "Filters",
            [
                ("Keep samples", args.keep_file),
                ("Extract variants", args.extract_file),
                ("Minimum MAF", args.min_maf),
                ("Minimum MAC", args.min_mac),
            ],
        ),
        (
            "Output",
            [
                ("Output prefix", args.output_file),
                ("Temporary directory", args.temp_dir),
                ("Output simulated beta", args.output_sim_beta),
            ],
        ),
    ]

    return format_sections(sections)

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

parser = argparse.ArgumentParser(
    description="""
    Commandline arguments for the complex trait simulator
"""
)

parser.add_argument(
    "--bfile",
    dest="bed_file",
    type=str,
    required=True,
    help="The BED files containing the genotype data. "
    'You may use a wildcard here (e.g. "data/chr_*.bed")',
)
parser.add_argument(
    "--keep",
    dest="keep_file",
    type=str,
    help="A plink-style keep file to select a subset of individuals for simulation.",
)
parser.add_argument(
    "--extract",
    dest="extract_file",
    type=str,
    help="A plink-style extract file to select a subset of SNPs for simulation.",
)
parser.add_argument(
    "--backend",
    dest="backend",
    type=str,
    default="magenpy",
    choices={"magenpy", "bed-reader", "xarray", "plink"},
    help="The backend software used for the computation.",
)
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-file",
    dest="output_file",
    type=str,
    required=True,
    help="The path where the simulated phenotype will be stored (no extension needed).",
)
parser.add_argument(
    "--output-simulated-beta",
    dest="output_sim_beta",
    action="store_true",
    default=False,
    help="Output a table with the true simulated effect size for each variant.",
)
parser.add_argument(
    "--min-maf",
    dest="min_maf",
    type=float,
    help="The minimum minor allele frequency for variants included in the simulation.",
)
parser.add_argument(
    "--min-mac",
    dest="min_mac",
    type=int,
    help="The minimum minor allele count for variants included in the simulation.",
)

# Simulation parameters:
parser.add_argument(
    "--h2",
    dest="h2",
    type=float,
    required=True,
    help="Trait heritability. Ranges between 0. and 1., inclusive.",
)
parser.add_argument(
    "--mix-prop",
    dest="mix_prop",
    type=str,
    help="Mixing proportions for the mixture density (comma separated). For example, "
    "for the spike-and-slab mixture density, with the proportion of causal variants "
    'set to 0.1, you can specify: "--mix-prop 0.9,0.1 --var-mult 0,1".',
)
parser.add_argument(
    "--prop-causal",
    "-p",
    dest="prop_causal",
    type=float,
    help="The proportion of causal variants in the simulation. See --mix-prop for "
    "more complex architectures specification.",
)
parser.add_argument(
    "--var-mult",
    "-d",
    dest="var_mult",
    type=str,
    help="Multipliers on the variance for each mixture component.",
)
parser.add_argument(
    "--phenotype-likelihood",
    dest="likelihood",
    type=str,
    default="gaussian",
    choices={"gaussian", "binomial"},
    help="The likelihood for the simulated trait: "
    "gaussian (e.g. quantitative) or binomial (e.g. case-control).",
)
parser.add_argument(
    "--prevalence",
    dest="prevalence",
    type=float,
    help="The prevalence of cases (or proportion of positives) for binary traits. "
    "Ranges between 0. and 1.",
)

parser.add_argument(
    "--seed",
    dest="seed",
    type=int,
    help="The random seed to use for the random number generator.",
)
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)

# ------------------------------------------------------
# Sanity checks on the inputs:

bed_file = get_filenames(args.bed_file, extension=".bed")
if len(bed_file) < 1:
    raise FileNotFoundError(
        f"No BED files were identified at the specified location: {args.bed_file}"
    )


if args.prop_causal is not None:
    pi = [1.0 - args.prop_causal, args.prop_causal]
    d = [0.0, 1.0]
elif args.mix_prop is not None:
    pi = list(map(float, args.mix_prop.split(",")))
    if args.var_mult:
        d = list(map(float, args.var_mult.split(",")))
    else:
        raise ValueError(
            "Specifying mixing proportions without variance multipliers is not permitted."
        )
else:
    logger.warning(
        "Mixing proportions not specified. Assuming an infinitesimal architecture "
        "where all variants are causal!"
    )
    pi = [0.0, 1.0]
    d = [0.0, 1.0]

if len(pi) != len(d):
    raise ValueError(
        "The multipliers and mixing proportions must be of the same length!"
    )

logger.info(
    "%s",
    format_cli_block(
        "Phenotype simulation request", format_run_summary(args, bed_file, pi, d)
    ),
)

# ------------------------------------------------------

# Record start time:
start_time = time.time()

# Set the random seed:
if args.seed is not None:
    np.random.seed(args.seed)

# Construct the PhenotypeSimulator object:
gs = PhenotypeSimulator(
    bed_file,
    keep_file=args.keep_file,
    extract_file=args.extract_file,
    phenotype_likelihood=args.likelihood,
    h2=args.h2,
    pi=pi,
    d=d,
    min_maf=args.min_maf,
    min_mac=args.min_mac,
    backend=args.backend,
    temp_dir=args.temp_dir,
)

logger.info("Simulating phenotype...")
gs.simulate(reset_beta=True, perform_gwas=False)
pheno_table = gs.to_phenotype_table()

# Write the simulated phenotypes to file:
makedir(osp.dirname(args.output_file))
pheno_table.to_csv(args.output_file + ".SimPheno", sep="\t", index=False, header=False)

if args.output_sim_beta:
    # Output the simulated effect sizes:
    sim_effects = gs.to_true_beta_table()
    sim_effects.to_csv(args.output_file + ".SimEffect", sep="\t", index=False)

gs.cleanup()

output_files = [args.output_file + ".SimPheno"]
if args.output_sim_beta:
    output_files.append(args.output_file + ".SimEffect")
# Record the end time:
end_time = time.time()
logger.info(
    "Phenotype simulation finished\n%s",
    format_rows(
        [
            ("Output files", ", ".join(output_files)),
            ("Total runtime", timedelta(seconds=end_time - start_time)),
        ]
    ),
)
