import os, sys, argparse

def parse_args():
    parser = argparse.ArgumentParser(description="SpaDOT: Optimal transport modeling uncovers spatial domain dynamics in spatiotemporal transcriptomics.",
        prog='SpaDOT')
    subparsers = parser.add_subparsers(help='sub-command help.', dest='cmd_choice')

    # ===========================
    # Parser for performing preprocess for spatiotemporal transcriptomics data
    # -i, --data
    # -o, --output_dir
    # --prefix
    # --feature_selection
    # ============================
    preprocess_parser = subparsers.add_parser('preprocess', 
            help='(Recommended but optional) Perform data preprocessing and feature selection.')
    preprocess_parser.add_argument('-i', '--data', dest='data',
            required=True, type=str, help="We preprocess the data ")
    preprocess_parser.add_argument('-o', '--output_dir', dest='output_dir',
            type=str, help="Output directory to store preprocessed data. Default: the same as where the data locates.")
    preprocess_parser.add_argument('--prefix', dest='prefix',
            type=str, default='preprocessed_', help="Prefix for preprocessed data. Default: preprocessed_")
    preprocess_parser.add_argument('--feature_selection', dest='feature_selection',
            default=True, action="store_false",
            help="Whether performing feature selection. If not specified, feature selection option is chosen.")

    # ===========================
    # Parser for training SpaDOT model to obtain latent representations
    # -i, --data
    # -o, --output_dir
    # --prefix
    # --config
    # --save_model
    # --device
    # ============================
    train_parser = subparsers.add_parser('train', help='Train a SpaDOT model.')
    train_parser.add_argument('-i', '--data', dest='data', 
            type=str, help="A anndata object containing time point information and spatial coordinates.")
    train_parser.add_argument('-o', '--output_dir', dest='output_dir', 
            type=str, help="Output directory to store latent representations. Default: the same as where the data locates.")
    train_parser.add_argument('--prefix', dest='prefix', 
            type=str, default='', help="Prefix for output latent representations. Default: ''")
    train_parser.add_argument('--config', dest='config',
            type=str, help="Path to the config file, in a yaml format.")
    train_parser.add_argument('--device', dest='device',
            type=str, default='cuda:0', help="Device to use for training. Default: cuda:0")
    train_parser.add_argument('--save_model', dest='save_model',
            default=False, action="store_true",
            help="Whether saving the trained model. If specified, the trained model will be stored in the output_dir as 'model.pt'.")

    ## ===========================
    # Parser for analyzing latent representations
    # -i, --data
    # -o, --output_dir
    # --prefix
    # ============================
    analyze_parser = subparsers.add_parser('analyze', help='Analyze the latent representations generated by SpaDOT model.')
    analyze_parser.add_argument('-i', '--data', dest='data', 
            type=str, required=True,
            help="An anndata object containing latent representations.")
    analyze_parser.add_argument('-o', '--output_dir', dest='output_dir', 
            type=str, help="Output directory to store analyses results. Default: the same as where the data locates.")
    analyze_parser.add_argument('--prefix', dest='prefix', 
            type=str, default='', help="Prefix for output analyses results. Default: ''")
    analyze_parser.add_argument('--n_clusters', dest='n_clusters',
                                type=lambda s: [int(item) for item in s.split(',')],
                                help="A comma-separated list of integers representing the number of clusters for analysis.")
    args = parser.parse_args()
    print("User input arguments: ", args)
    return args


def main():
    args = parse_args()
    ## preprocessing
    if "preprocess" == args.cmd_choice:
        from SpaDOT import preprocess
        if not os.path.exists(args.data):
            sys.exit("The data does not exist! Please make sure the data is correctly specified.")
        preprocess(args)
        
    ## train to obtain latent representations
    if "train" == args.cmd_choice:
        from SpaDOT import train
        if not os.path.exists(args.data):
            sys.exit("The preprocessed data does not exist! Please make sure the data is correctly specified.")
        train(args)

    ## analyze latent representations
    if "analyze" == args.cmd_choice:
        from SpaDOT import analyze
        if not os.path.exists(args.data):
            sys.exit("The latent representations does not exist! Please make sure the latent representations are correctly generated from the last step.")
        analyze(args)

if __name__ == "__main__":
    main()