# This file is part of the ARG-LMM suite for scalable variance
# component analysis using ancestral recombination graphs.
# Copyright (C) 2025 ARG-LMM Developers.

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import argparse
import logging
import time

from arg_lmm.arg_lmm import arg_lmm_main
from arg_lmm.arg_rhe import arg_rhe_main

logging.basicConfig(
    format="%(asctime)s %(levelname)-8s %(message)s",
    level=logging.INFO,
    datefmt="%Y-%m-%d %H:%M:%S"
)
logger = logging.getLogger(__name__)

DEFAULT_ALPHA = -1.0
DEFAULT_MAC = 1
DEFAULT_MU = 1e-6
DEFAULT_CALIB = 1
DEFAULT_NTHREADS = 1
DEFAULT_SEED = 42

def arg_lmm(force_rhe=False):
    """
    Main CLI entrypoint

    By default arg-lmm mixes in all possbile command line args including RHE. Alternatively, force_rhe only accepts
    arguments applicable to RHE and is used for the arg-rhe alias.
    """
    parser = argparse.ArgumentParser(description="ARG-LMM: estimate total-h2, compute LOCO residuals and BLUP")

    parser.add_argument("--argfiles", nargs="+", help="list of ARG files, one per chromosome (e.g. chr1.arg chr2.arg ...)")
    parser.add_argument("--argfile", action="store", help="input ARG file (RHE only)")
    parser.add_argument("--pheno", required=True, help="phenotype file (space-separated, FID and IID as first two columns)")
    parser.add_argument("-o", "--out", required=True, help="output prefix")
    parser.add_argument("--alpha", type=float, default=DEFAULT_ALPHA, help="alpha parameter for scaling")
    parser.add_argument("--seed", type=int, default=DEFAULT_SEED, help="random seed")
    parser.add_argument("--ncalib", type=int, default=DEFAULT_CALIB, help="number of calibration variants per chromosome")
    parser.add_argument("--nthreads", type=int, default=DEFAULT_NTHREADS, help="number of threads")
    parser.add_argument("--blup", action="store_true", help="also compute genome-wide BLUPs")
    parser.add_argument("--rhe", action="store_true", help="enable Haseman-Elston estimation")
    parser.add_argument("--h2", action="store_true", help="estimate heritability and standard errors with jackknife (one block per ARG)")
    parser.add_argument("--mu", action="store", dest="mu", default=DEFAULT_MU, type=float, help="resampled mutation rate (RHE only)")
    parser.add_argument("--mac", action="store", dest="mac", default=DEFAULT_MAC, type=float, help="minimum minor allele count for resampled mutations (RHE only)")
    parser.add_argument("--assoc", action="store_true", help="enable gene-based association testing (RHE only)")

    # arg-lmm by default processes all possible CLI args, but can also be run from an alias arg-rhe with reduced args
    args = parser.parse_args()
    rhe = force_rhe or args.rhe

    # Check for either --assoc or --h2 when when running `arg-lmm --rhe`. When running --h2 the behaviour falls back to
    # arg_lmm_main, so this is made explicit via CLI args; only rhe with assoc goes into arg_rhe_main entrypoint.
    if rhe:
        if args.h2 == args.assoc:
            logger.error("--h2 and --assoc are mutually exclusive but both specified")
        if not args.assoc:
            logger.error("--assoc must be specified when not using --h2 with RHE")

    start_time = time.time()
    if rhe and args.assoc:
        if not "argfile" in args:
            logger.error("--argfile argument expected for assoc")
            exit()

        if "argfiles" in args and args.argfiles:
            logger.warning("--argfiles argument not applicable to assoc")

        if "nthreads" in args and args.nthreads != DEFAULT_NTHREADS:
            logger.warning("--nthreads argument not applicable to assoc")

        if "ncalib" in args and args.ncalib != DEFAULT_CALIB:
            logger.warning("--ncalib argument not applicable to assoc")

        if "blup" in args and args.blup:
            logger.warning("--blup flag not applicable to assoc")

        logger.info(f"Starting arg-lmm RHE with following parameters:")
        logger.info(f"  arg:        {args.argfile}")
        logger.info(f"  pheno:      {args.pheno}")
        logger.info(f"  out:        {args.out}")
        logger.info(f"  mu:         {args.mu}")
        logger.info(f"  alpha:      {args.alpha}")
        logger.info(f"  mac:        {args.mac}")
        logger.info(f"  seed:       {args.seed}")

        arg_rhe_main(args.argfile, args.pheno, args.out, args.mu, args.alpha, args.mac, args.seed)
    else:
        if not args.argfiles or len(args.argfiles) < 2:
            logger.error("arg-lmm expects at least two ARG paths")
            exit()

        if "argfile" in args and args.argfile:
            logger.warning("--argfile argument only applicable to RHE")

        if "mu" in args and args.mu != DEFAULT_MU:
            logger.warning("--mu argument only applicable to RHE")

        if "mac" in args and args.mac != DEFAULT_MAC:
            logger.warning("--mac argument only applicable to RHE")

        logger.info(f"Starting arg-lmm with following parameters:")
        logger.info(f"  arg paths:  {args.argfiles}")
        logger.info(f"  pheno:      {args.pheno}")
        logger.info(f"  out prefix: {args.out}")
        logger.info(f"  ncalib:     {args.ncalib}")
        logger.info(f"  alpha:      {args.alpha}")
        logger.info(f"  seed:       {args.seed}")
        logger.info(f"  nthreads:   {args.nthreads}")
        logger.info(f"  blup:       {args.blup}")
        logger.info(f"  h2:         {args.h2}")

        arg_lmm_main(args.argfiles, args.pheno, args.out, args.ncalib, args.alpha, args.seed, args.nthreads, args.blup, args.h2)

    logger.info(f"Done, in {time.time() - start_time} seconds")


def arg_rhe():
    arg_lmm(force_rhe=True)
