launch

#!/usr/bin/env python3

import os
import sys
import subprocess
import shutil
from pathlib import Path


SCRIPT_DIR = Path(__file__).parent

def load_ascii_art():
    return """
    ███████╗███╗   ███╗██████╗ ██╗██████╗ ██████╗ ███████╗██████╗                
    ██╔════╝████╗ ████║██╔══██╗██║██╔══██╗██╔══██╗██╔════╝██╔══██╗               
    █████╗  ██╔████╔██║██████╔╝██║██████╔╝██████╔╝█████╗  ██████╔╝               
    ██╔══╝  ██║╚██╔╝██║██╔══██╗██║██╔═══╝ ██╔══██╗██╔══╝  ██╔═══╝                
    ██║     ██║ ╚═╝ ██║██║  ██║██║██║     ██║  ██║███████╗██║                    
    ╚═╝     ╚═╝     ╚═╝╚═╝  ╚═╝╚═╝╚═╝     ╚═╝  ╚═╝╚══════╝╚═╝                    
                                                                                 
    ██╗    ██╗ ██████╗ ██████╗ ██╗  ██╗██████╗ ███████╗███╗   ██╗ ██████╗██╗  ██╗
    ██║    ██║██╔═══██╗██╔══██╗██║ ██╔╝██╔══██╗██╔════╝████╗  ██║██╔════╝██║  ██║
    ██║ █╗ ██║██║   ██║██████╔╝█████╔╝ ██████╔╝█████╗  ██╔██╗ ██║██║     ███████║
    ██║███╗██║██║   ██║██╔══██╗██╔═██╗ ██╔══██╗██╔══╝  ██║╚██╗██║██║     ██╔══██║
    ╚███╔███╔╝╚██████╔╝██║  ██║██║  ██╗██████╔╝███████╗██║ ╚████║╚██████╗██║  ██║
     ╚══╝╚══╝  ╚═════╝ ╚═╝  ╚═╝╚═╝  ╚═╝╚═════╝ ╚══════╝╚═╝  ╚═══╝ ╚═════╝╚═╝  ╚═╝
    """

class Colors:
    HEADER = '\033[95m'
    BLUE = '\033[94m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

def clear_screen():
    """Clear the screen using ANSI escape sequences for better compatibility with headless systems."""
    # Flush stdout to ensure any pending output is written
    sys.stdout.flush()
    # Use ANSI escape sequences for more reliable clearing on headless systems
    # \033[2J clears the entire screen
    # \033[H moves cursor to home position (0,0)
    sys.stdout.write('\033[2J\033[H')
    sys.stdout.flush()

def get_terminal_size():
    try:
        columns, rows = shutil.get_terminal_size()
        return columns, rows
    except:
        return 80, 24  # default fallback

def print_centered(text, color=None):
    columns, _ = get_terminal_size()
    text = text.strip()

    if color:
        formatted_text = f"{color}{text}{Colors.ENDC}"
    else:
        formatted_text = text

    print(formatted_text.center(columns))

def print_header():
    columns, _ = get_terminal_size()
    print('\n')
    print('=' * columns)
    print_centered("FMRIPREP WORKBENCH: Main Menu", Colors.BOLD + Colors.BLUE)
    print('=' * columns)
    print('\n')

def display_ascii_art():
    ascii_art = load_ascii_art()
    print(ascii_art)
    input(f"\n{Colors.GREEN}Press Enter to continue...{Colors.ENDC}")
    clear_screen()

def get_user_input(prompt, options=None, default=None, yes_no=False):
    if yes_no:
        while True:
            if default is not None:
                default_text = "Y/n" if default else "y/N"
                user_input = input(f"{prompt} [{default_text}]: ").strip().lower()
                if user_input == "":
                    return default
            else:
                user_input = input(f"{prompt} [y/n]: ").strip().lower()

            if user_input in ['y', 'yes']:
                return True
            elif user_input in ['n', 'no']:
                return False
            else:
                print(f"{Colors.YELLOW}Please enter 'y' or 'n'.{Colors.ENDC}")
    else:
        while True:
            if default is not None:
                user_input = input(f"{prompt} [{default}]: ").strip()
                if user_input == "":
                    return default
            else:
                user_input = input(f"{prompt}: ").strip()

            if options is None:
                if user_input:  # ensure non-empty input
                    return user_input
                else:
                    print(f"{Colors.YELLOW}Input cannot be empty.{Colors.ENDC}")
            elif user_input in [str(o) for o in options]:
                return user_input
            else:
                print(f"{Colors.YELLOW}Invalid input. Please enter one of: {', '.join([str(o) for o in options])}{Colors.ENDC}")

def display_menu():
    print_header()

    menu_items = [
        ("1", "Step 1: Download from FlyWheel -> Server"),
        ("2", "Step 2: Run dcm2niix (DICOM -> BIDS format conversion)"),
        ("3", "Step 3: Prep for fMRIPrep (remove dummy scans, update fmap JSON metadata, config fmap SDC)"),
        ("4", "Step 4: QC - Verify dcm -> nii -> bids metadata"),
        ("5", "Step 5: QC - Verify number of volumes per scan file"),
        ("6", "Step 6: Run fMRIPrep anatomical workflows only (if doing manual edits, otherwise skip to step 9)"),
        ("7", "Step 7: Download Freesurfer outputs for manual editing"),
        ("8", "Step 8: Upload edited Freesurfer outputs back to server"),
        ("9", "Step 9: Run remaining fMRIPrep steps (full anatomical + functional workflows)"),
        ("10", "FSL GLM: Setup new statistical model"),
        ("11", "FSL GLM: Run Level 1 analysis (individual runs)"),
        ("12", "FSL GLM: Run Level 2 analysis (subject-level)"),
        ("13", "FSL GLM: Run Level 3 analysis (group-level)"),
        ("14", "Utilities: Tarball/Untar utility for sourcedata/ BIDS directories")
    ]

    print(f"{Colors.BOLD}Select the prep step you want to run:{Colors.ENDC}\n")
    for key, description in menu_items:
        print(f"{Colors.GREEN}{key}.{Colors.ENDC} {description}")
    print()

    options = [item[0] for item in menu_items]
    step_choice = get_user_input("Enter your choice", options)

    return step_choice

def get_parameters(step_choice):
    step_params = {
        "1": {  # Step 1 (Flywheel Downloader)
            "keys": ["fw_subid", "fw_session_id"],
            "prompts": ["Subject ID on Flywheel?", "Session ID on Flywheel?"]
        },
        "2": {  # Step 2 (dcm2niix BIDS)
            "keys": ["fw_session_id", "new_bids_id_number", "skip_tar"],
            "prompts": ["Session ID on Flywheel?", "What subject ID number would you like to assign for BIDS?", "Skip tar extraction (for manually configured scan directories)?"]
        },
        "3": {  # Step 3 (Prep for fMRIPrep)
            "keys": [],
            "prompts": []
        },
        "4": {  # Step 4 (QC metadata)
            "keys": [],
            "prompts": []
        },
        "5": {  # Step 5 (QC volumes)
            "keys": [],
            "prompts": []
        },
        "6": {  # Step 6 (fmriprep --anat-only)
            "keys": [],
            "prompts": []
        },
        "7": {  # Freesurfer download
            "keys": [],
            "prompts": []
        },
        "8": {  # Freesurfer upload
            "keys": [],
            "prompts": []
        },
        "9": {  # Step 9 (fmriprep full)
            "keys": [],
            "prompts": []
        },
        "10": {  # FSL GLM: Setup model
            "keys": [],
            "prompts": []
        },
        "11": {  # FSL GLM: Level 1
            "keys": ["model_name", "no_feat"],
            "prompts": ["Model name?", "Only create FSF files without running FEAT?"]
        },
        "12": {  # FSL GLM: Level 2
            "keys": ["model_name", "no_feat"],
            "prompts": ["Model name?", "Only create FSF files without running FEAT?"]
        },
        "13": {  # FSL GLM: Level 3
            "keys": ["model_name", "no_feat"],
            "prompts": ["Model name?", "Only create FSF files without running FEAT?"]
        },
        "14": {  # Tarball/Untar utility
            "keys": ["operation", "sourcedata_dir", "subjects_spec", "keep_original"],
            "prompts": [
                "Operation (1=tar-all, 2=tar-subjects, 3=untar-all, 4=untar-subjects)",
                "Path to sourcedata directory",
                "Subject list (comma-separated IDs or file path, press Enter to skip for tar-all/untar-all)",
                "Keep original directories after tarballing? (only for tar operations)"
            ]
        }
    }

    param_keys = step_params[step_choice]["keys"]
    param_prompts = step_params[step_choice]["prompts"]

    user_inputs = []

    # Special handling for tarball utility
    if step_choice == "14":
        # Get operation type
        operation = get_user_input(param_prompts[0], options=["1", "2", "3", "4"])
        operation_map = {
            "1": "--tar-all",
            "2": "--tar-subjects",
            "3": "--untar-all",
            "4": "--untar-subjects"
        }
        
        # Build command based on operation
        if operation in ["2", "4"]:
            # For tar-subjects/untar-subjects, need subject list
            subjects_spec = get_user_input(param_prompts[2])
            user_inputs.extend([operation_map[operation], subjects_spec])
        else:
            # For tar-all/untar-all
            user_inputs.append(operation_map[operation])
        
        # Get sourcedata directory
        sourcedata_dir = get_user_input(param_prompts[1])
        user_inputs.extend(["--sourcedata-dir", sourcedata_dir])
        
        # Ask about keeping original only for tar operations
        if operation in ["1", "2"]:
            keep_original = get_user_input(param_prompts[3], yes_no=True, default=False)
            if keep_original:
                user_inputs.append("--keep-original")
        
        return user_inputs, param_keys
    
    # Regular parameter handling for other steps
    for i in range(len(param_keys)):
        key = param_keys[i]
        prompt = param_prompts[i]

        # special case for anat-only yes/no
        if key == "anat_only":
            response = get_user_input(prompt, yes_no=True)
            if response:
                user_inputs.append("--anat-only")
        # special case for skip-tar yes/no
        elif key == "skip_tar":
            response = get_user_input(prompt, yes_no=True, default=False)
            if response:
                user_inputs.append("--skip-tar")
        # special case for no_feat yes/no (FSL GLM)
        elif key == "no_feat":
            response = get_user_input(prompt, yes_no=True, default=False)
            if response:
                user_inputs.append("--no-feat")
        else:
            input_value = get_user_input(prompt)
            user_inputs.append(input_value)

    return user_inputs, param_keys

def confirm_parameters(user_inputs, param_keys):
    if not user_inputs:
        return True

    print(f"\n{Colors.BOLD}You have entered the following parameters:{Colors.ENDC}")

    for i, value in enumerate(user_inputs):
        if i < len(param_keys):
            key = param_keys[i]
            print(f"  {key}: {value}")
        else:
            print(f"  Parameter {i}: {value}")

    print()
    return get_user_input("Proceed with execution?", yes_no=True, default=True)

def run_script(step_choice, user_inputs):
    # map step choice to script name
    scripts = {
        "1": "01-run.sbatch",
        "2": "02-run.sbatch",
        "3": "03-run.sbatch",
        "4": "04-run.sbatch",
        "5": "05-run.sbatch",
        "6": "06-run.sbatch",
        "7": "toolbox/download_freesurfer.sh",
        "8": "toolbox/upload_freesurfer.sh",
        "9": "09-run.sbatch",
        "10": "10-fsl-glm/setup_glm.sh",
        "11": "11-run.sbatch",
        "12": "12-run.sbatch",
        "13": "13-run.sbatch",
        "14": "toolbox/tarball_sourcedata.sh"
    }

    selected_script = scripts[step_choice]

    cmd = ["bash", selected_script] + user_inputs

    print(f"\n{Colors.BOLD}Executing:{Colors.ENDC} {' '.join(cmd)}")

    try:
        subprocess.run(cmd)
    except Exception as e:
        print(f"{Colors.RED}Error running script: {e}{Colors.ENDC}")
        return False

    return True

def main():
    clear_screen()
    display_ascii_art()

    step_choice = display_menu()
    clear_screen()

    user_inputs, param_keys = get_parameters(step_choice)

    if confirm_parameters(user_inputs, param_keys):
        success = run_script(step_choice, user_inputs)

        if success:
            print(f"\n{Colors.GREEN}Script executed successfully.{Colors.ENDC}")
        else:
            print(f"\n{Colors.RED}Script execution failed.{Colors.ENDC}")
    else:
        print(f"\n{Colors.YELLOW}Operation cancelled.{Colors.ENDC}")

if __name__ == "__main__":
    main()