#!/usr/bin/env python
import sys
import os
import argparse
from pkg_resources import require
from pathlib import Path

from cxas import CXAS


def main():
    parser = argparse.ArgumentParser(description="Segment 159 anatomical structures in X-Ray images.",
                                     epilog="Written by Constantin Seibold. If you use this tool please cite accordingly.")

    parser.add_argument("-i", "--input",metavar="filepath", dest="input",
                        help="Either path to file or to directory to be processed.", type=str, required=True)
    
    parser.add_argument("-o", "--output",metavar="directory", dest="output",
                        help="Output directory for segmentation masks", type=str, required=True)
    
    parser.add_argument("-f", "--feature", choices=["SCD", 
                                                    "CTR", 
                                                    "Spine-Center Distance",
                                                    "Cardio-Thoracic Ratio",
                                                    'perimeter',
                                                    'compactness',
                                                    'area',
                                                    'centroid',
                                                    'box',
                                                    'convexity',
                                                    ],
                    help="Select which features are supposed to be extracted.", required=True)
    
    parser.add_argument("-s", "--store_seg", action='store_true', dest="store_seg",
                        help="Wether to also store segmentation masks")
    
    parser.add_argument("-ot", "--output_type", choices=["json", "npy", "npz", "jpg", "png", "dicom-seg", 
                                                   ],
                        help="Designates the storage type of segmentations if they are stored ", default='png')
    
    parser.add_argument("-g", "--gpus",
                    help="Select specific GPU/CPU to process the input.", default="0")
    
    parser.add_argument("-m", "--model", choices=["UNet_ResNet50_default"],
                    help="Select Model used for inference.", default="UNet_ResNet50_default")

    args = parser.parse_args()

    if args.gpus != 'cpu':
        import torch
        if not torch.cuda.is_available():
            print('No GPU is available. Switching to CPU.')
            args.gpus = 'cpu'
            
    model = CXAS(
            model_name = args.model,
            gpus       = args.gpus
        )
    
    if os.path.isdir(args.input):
        model.extract_features_for_folder(
                input_directory_name = args.input,  
                output_directory = args.output,
                feat_to_extract = args.feature,
                create = True, 
                store_pred = args.store_seg,
                storage_type = args.output_type,
            )
    elif os.path.isfile(args.input):
        model.extract_features_for_file(
                filename = args.input,  
                output_directory = args.output,
                feat_to_extract = args.feature,
                create = True, 
                do_store = True,
                storage_type = args.output_type,
            )
    else:
        print('{} is neither file nor directory...'.format(args.input))

if __name__ == "__main__":
    main()