"""
Main caller of hapla.
"""

__author__ = "Jonas Meisner"

# Libraries
import argparse
import sys
from hapla import __version__


##### hapla caller ######
def main():
    # Argparser
    parser = argparse.ArgumentParser(prog="hapla")
    parser.add_argument("--version", action="version", version=f"v{__version__}")
    subparsers = parser.add_subparsers(title="hapla commands")

    ### Commands
    # hapla cluster
    parser_cluster = subparsers.add_parser("cluster")
    parser_cluster.add_argument(
        "--version", action="version", version=f"v{__version__}"
    )
    parser_cluster.add_argument(
        "-g",
        "--vcf",
        "--bcf",
        metavar="FILE",
        help="Input phased genotype file in VCF/BCF format",
    )
    parser_cluster.add_argument(
        "-f", "--size", type=int, metavar="INT", help="Use fixed sized windows"
    )
    parser_cluster.add_argument(
        "-l",
        "--length",
        type=int,
        metavar="INT",
        help="Use length-based windows (in BP)",
    )
    parser_cluster.add_argument(
        "-w", "--windows", metavar="FILE", help="Use provided start indices for windows"
    )
    parser_cluster.add_argument(
        "-p",
        "--lmbda",
        type=float,
        default=0.1,
        metavar="FLOAT",
        help="Set lambda hyperparameter (0.1)",
    )
    parser_cluster.add_argument(
        "-s",
        "--step",
        type=int,
        metavar="INT",
        help="Step-size for sliding window using fixed sized windows",
    )
    parser_cluster.add_argument(
        "-t",
        "--threads",
        type=int,
        default=1,
        metavar="INT",
        help="Number of threads (1)",
    )
    parser_cluster.add_argument(
        "-o", "--out", default="hapla.cluster", metavar="OUTPUT", help="Output prefix"
    )
    parser_cluster.add_argument(
        "--min-freq",
        type=float,
        default=0.005,
        metavar="FLOAT",
        help="Minimum haplotype cluster frequency (0.005)",
    )
    parser_cluster.add_argument(
        "--min-mac",
        type=int,
        metavar="INT",
        help="Minimum haplotype cluster allele count",
    )
    parser_cluster.add_argument(
        "--max-clusters",
        type=int,
        default=256,
        metavar="INT",
        help="Maximum number of haplotype clusters per window (256)",
    )
    parser_cluster.add_argument(
        "--max-iterations",
        type=int,
        default=1000,
        metavar="INT",
        help="Maximum number of iterations (1000)",
    )
    parser_cluster.add_argument(
        "--medians", action="store_true", help="Save haplotype cluster medians"
    )
    parser_cluster.add_argument(
        "--prune",
        action="store_true",
        help="Prune outlier clusters prior to reclustering",
    )
    parser_cluster.add_argument(
        "--plink", action="store_true", help="Generate binary PLINK output"
    )
    parser_cluster.add_argument(
        "--duplicate-fid",
        action="store_true",
        help="Use sample list as family ID (PLINK 1.9 compatibility)",
    )
    parser_cluster.add_argument(
        "--memory", action="store_true", help="Store haplotypes in 1-bit matrix"
    )

    # hapla predict
    parser_predict = subparsers.add_parser("predict")
    parser_predict.add_argument(
        "--version", action="version", version=f"v{__version__}"
    )
    parser_predict.add_argument(
        "-g",
        "--vcf",
        "--bcf",
        metavar="FILE",
        help="Input phased/unphased genotype file in VCF/BCF format",
    )
    parser_predict.add_argument(
        "-b",
        "--bfile",
        metavar="FILE",
        help="Input unphased genotype file in binary PLINK format",
    )
    parser_predict.add_argument(
        "-r",
        "--ref",
        metavar="FILE",
        help="Input reference prefix of pre-estimated cluster medians",
    )
    parser_predict.add_argument(
        "-t",
        "--threads",
        type=int,
        default=1,
        metavar="INT",
        help="Number of threads (1)",
    )
    parser_predict.add_argument(
        "-o", "--out", default="hapla.predict", metavar="OUTPUT", help="Output prefix"
    )
    parser_predict.add_argument(
        "--plink", action="store_true", help="Generate binary PLINK output"
    )
    parser_predict.add_argument(
        "--duplicate-fid",
        action="store_true",
        help="Use sample list as family ID (PLINK 1.9 compatibility)",
    )
    parser_predict.add_argument(
        "--memory", action="store_true", help="Store haplotypes in 2-bit matrix"
    )

    # hapla struct
    parser_struct = subparsers.add_parser("struct")
    parser_struct.add_argument("--version", action="version", version=f"v{__version__}")
    parser_struct.add_argument(
        "-f",
        "--filelist",
        metavar="FILE",
        help="Filelist with paths to haplotype cluster alleles files",
    )
    parser_struct.add_argument(
        "-z",
        "--clusters",
        metavar="FILE",
        help="Path to a single haplotype cluster alleles file",
    )
    parser_struct.add_argument(
        "-t",
        "--threads",
        type=int,
        default=1,
        metavar="INT",
        help="Number of threads (1)",
    )
    parser_struct.add_argument(
        "-o", "--out", default="hapla.pca", metavar="OUTPUT", help="Output prefix"
    )
    parser_struct.add_argument(
        "--grm",
        action="store_true",
        help="Estimate genome-wide relationship matrix (GRM)",
    )
    parser_struct.add_argument(
        "--projection", metavar="FILE", help="Project samples on to existing PC space"
    )
    parser_struct.add_argument(
        "--chunk",
        type=int,
        default=4096,
        metavar="INT",
        help="Number of cluster alleles in batches (4096)",
    )
    parser_struct.add_argument(
        "--power",
        type=int,
        default=11,
        metavar="INT",
        help="Number of power iterations to perform (11)",
    )
    parser_struct.add_argument(
        "--no-centering",
        action="store_true",
        help="Do not perform Gower and data centering of GRM",
    )
    parser_struct.add_argument(
        "--duplicate-fid",
        action="store_true",
        help="Use sample list as family ID for GCTA format",
    )
    parser_struct.add_argument(
        "--pca", type=int, metavar="INT", help="Perform PCA and extract eigenvectors"
    )
    parser_struct.add_argument(
        "--memory",
        action="store_true",
        help="Perform memory efficient SVD for large-scale data",
    )
    parser_struct.add_argument(
        "--seed", type=int, default=42, metavar="INT", help="Random seed (42)"
    )
    parser_struct.add_argument(
        "--loadings", action="store_true", help="Save frequencies and loadings of SVD"
    )
    parser_struct.add_argument(
        "--raw", action="store_true", help="Raw output without '*.fam' info"
    )

    # hapla admix
    parser_admix = subparsers.add_parser("admix")
    parser_admix.add_argument("--version", action="version", version=f"v{__version__}")
    parser_admix.add_argument(
        "-f",
        "--filelist",
        metavar="FILE",
        help="Filelist with paths to haplotype cluster alleles files",
    )
    parser_admix.add_argument(
        "-z",
        "--clusters",
        metavar="FILE",
        help="Path to a single haplotype cluster alleles file",
    )
    parser_admix.add_argument(
        "-k", "--K", type=int, metavar="INT", help="Number of ancestral components"
    )
    parser_admix.add_argument(
        "-t",
        "--threads",
        type=int,
        default=1,
        metavar="INT",
        help="Number of threads (1)",
    )
    parser_admix.add_argument(
        "-o", "--out", default="hapla.admix", metavar="OUTPUT", help="Output prefix"
    )
    parser_admix.add_argument(
        "--seed", type=int, default=42, metavar="INT", help="Random seed (42)"
    )
    parser_admix.add_argument(
        "--iter",
        type=int,
        default=1000,
        metavar="INT",
        help="Maximum number of iterations (1000)",
    )
    parser_admix.add_argument(
        "--tole",
        type=float,
        default=1e-9,
        metavar="FLOAT",
        help="Tolerance in scaled log-likelihood units (1e-9)",
    )
    parser_admix.add_argument(
        "--batches",
        type=int,
        default=16,
        metavar="INT",
        help="Number of initial mini-batches (16)",
    )
    parser_admix.add_argument(
        "--check",
        type=int,
        default=5,
        metavar="INT",
        help="Number of iterations between convergence checks (5)",
    )
    parser_admix.add_argument(
        "--supervised", metavar="FILE", help="Path to population assignment file"
    )
    parser_admix.add_argument(
        "--projection",
        metavar="FILE",
        help="Path to ancestral haplotype cluster allele frequencies file",
    )
    parser_admix.add_argument(
        "--chunk",
        type=int,
        default=4096,
        metavar="INT",
        help="Number of expected cluster alleles in batches (4096)",
    )
    parser_admix.add_argument(
        "--power",
        type=int,
        default=11,
        metavar="INT",
        help="Number of power iterations to perform (11)",
    )
    parser_admix.add_argument(
        "--als-iter",
        metavar="INT",
        type=int,
        default=1000,
        help="Maximum number of iterations in ALS (1000)",
    )
    parser_admix.add_argument(
        "--als-tole",
        metavar="FLOAT",
        type=float,
        default=1e-4,
        help="Tolerance for RMSE of P between iterations (1e-4)",
    )
    parser_admix.add_argument(
        "--subsampling",
        metavar="INT",
        type=int,
        default=4,
        help="Subsampling factor for ALS/SVD initialization (4)",
    )
    parser_admix.add_argument(
        "--no-freqs",
        action="store_true",
        help="Do not save haplotype cluster frequencies",
    )
    parser_admix.add_argument(
        "--random-init", action="store_true", help="Random initialization of parameters"
    )
    parser_admix.add_argument(
        "--prefix",
        default="chr",
        metavar="OUTPUT",
        help="Prefix for multiple haplotype cluster frequency files",
    )

    # hapla fatash
    parser_fatash = subparsers.add_parser("fatash")
    parser_fatash.add_argument("--version", action="version", version=f"v{__version__}")
    parser_fatash.add_argument(
        "-f",
        "--filelist",
        metavar="FILE",
        help="Filelist with paths to haplotype cluster alleles files",
    )
    parser_fatash.add_argument(
        "-z",
        "--clusters",
        metavar="FILE",
        help="Path to a single haplotype cluster alleles file",
    )
    parser_fatash.add_argument(
        "-e",
        "--pfilelist",
        metavar="FILE",
        help="Filelist with paths to haplotype cluster frequencies files",
    )
    parser_fatash.add_argument(
        "-p",
        "--pfile",
        metavar="FILE",
        help="Path to file with haplotype cluster frequencies",
    )
    parser_fatash.add_argument(
        "-q", "--qfile", metavar="FILE", help="Path to file with ancestry proportions"
    )
    parser_fatash.add_argument(
        "-t",
        "--threads",
        type=int,
        default=1,
        metavar="INT",
        help="Number of threads (1)",
    )
    parser_fatash.add_argument(
        "-o", "--out", default="hapla.fatash", metavar="OUTPUT", help="Output prefix"
    )
    parser_fatash.add_argument(
        "--block",
        type=int,
        metavar="INT",
        default=1,
        help="Number of genomic windows in block to increase smoothness"
    )
    parser_fatash.add_argument(
        "--alpha",
        type=float,
        metavar="FLOAT",
        help="Single alpha rate in transition matrix",
    )
    parser_fatash.add_argument(
        "--medians", action="store_true", help="Utilize haplotype cluster probabilities"
    )
    parser_fatash.add_argument(
        "--viterbi",
        action="store_true",
        help="Perform Viterbi decoding for LAI inference (HMM)",
    )
    parser_fatash.add_argument(
        "--save-posteriors",
        action="store_true",
        help="Save posterior probabilities from posterior decoding",
    )
    parser_fatash.add_argument(
        "--prefix",
        default="chr",
        metavar="OUTPUT",
        help="Prefix for multiple path files",
    )
    parser_fatash.add_argument(
        "--genome-wide",
        action="store_true",
        help="Use genome-wide ancestry proportions across all chromosomes",
    )
    parser_fatash.add_argument(
        "--alpha-min",
        type=int,
        default=4,
        help="Lowest alpha exponent (negative log10 scale) to use in majority voting (4)",
    )
    parser_fatash.add_argument(
        "--alpha-max",
        type=int,
        default=9,
        help="Highest alpha exponent (negative log10 scale) to use in majority voting (9)",
    )
    parser_fatash.add_argument(
        "--alpha-num",
        type=int,
        default=5,
        help="Number of alpha values to evaluate"
    )
    parser_fatash.add_argument(
        "--admix-seed",
        type=int,
        default=42,
        metavar="INT",
        help="Random seed in ancestry estimation (42)",
    )
    parser_fatash.add_argument(
        "--admix-iter",
        type=int,
        default=1000,
        metavar="INT",
        help="Maximum number of iterations (1000)",
    )
    parser_fatash.add_argument(
        "--admix-tole",
        type=float,
        default=1e-9,
        metavar="FLOAT",
        help="Tolerance in scaled log-likelihood units (1e-9)",
    )
    parser_fatash.add_argument(
        "--admix-batches",
        type=int,
        default=8,
        metavar="INT",
        help="Number of initial mini-batches (8)",
    )
    parser_fatash.add_argument(
        "--admix-check",
        type=int,
        default=5,
        metavar="INT",
        help="Number of iterations between convergence checks (5)",
    )
    parser_fatash.add_argument(
        "--simple",
        action="store_true",
        help="Use simplified transition probabibilities in HMM",
    )

    # Parse arguments
    args = parser.parse_args()
    if len(sys.argv) < 2:
        parser.print_help()
        sys.exit()

    ### Run specified command
    # hapla cluster
    if sys.argv[1] == "cluster":
        if len(sys.argv) < 3:
            parser_cluster.print_help()
            sys.exit()
        else:
            from hapla import cluster

            deaf = vars(parser_cluster.parse_args([]))
            cluster.main(args, deaf)

    # hapla predict
    if sys.argv[1] == "predict":
        if len(sys.argv) < 3:
            parser_predict.print_help()
            sys.exit()
        else:
            from hapla import predict

            deaf = vars(parser_predict.parse_args([]))
            predict.main(args, deaf)

    # hapla struct
    if sys.argv[1] == "struct":
        if len(sys.argv) < 3:
            parser_struct.print_help()
            sys.exit()
        else:
            from hapla import struct

            deaf = vars(parser_struct.parse_args([]))
            struct.main(args, deaf)

    # hapla admix
    if sys.argv[1] == "admix":
        if len(sys.argv) < 3:
            parser_admix.print_help()
            sys.exit()
        else:
            from hapla import admix

            deaf = vars(parser_admix.parse_args([]))
            admix.main(args, deaf)

    # hapla fatash
    if sys.argv[1] == "fatash":
        if len(sys.argv) < 3:
            parser_fatash.print_help()
            sys.exit()
        else:
            from hapla import fatash

            deaf = vars(parser_fatash.parse_args([]))
            fatash.main(args, deaf)


##### Define main #####
if __name__ == "__main__":
    main()
