xmipp

#!/usr/bin/env python3
# ***************************************************************************
# * Authors:		Alberto García (alberto.garcia@cnb.csic.es)
# *					Martín Salinas (martin.salinas@cnb.csic.es)
# *
# *
# * 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 2 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, write to the Free Software
# * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
# * 02111-1307 USA
# *
# * All comments concerning this program package may be sent to the
# * e-mail address 'scipion@cnb.csic.es'
# ***************************************************************************/

# General imports
import argparse, sys, os, multiprocessing
from typing import Dict

# Installer imports
from installer.constants import (MODE_ALL, MODE_COMPILE_AND_INSTALL, MODE_CONFIG, MODE_CLEAN_BIN,
	MODE_CLEAN_ALL, MODE_VERSION, MODE_GET_MODELS, MODE_TEST, MODE_GIT, MODE_ADD_MODEL, MODE_CONFIG_BUILD,
	COMMON_USAGE_HELP_MESSAGE, DEFAULT_MODELS_DIR, CONFIG_FILE, CMAKE_INSTALL_PREFIX,
	CMAKE_CONFIGURE_ERROR, CMAKE_COMPILE_ERROR, CMAKE_INSTALL_ERROR, PARAM_LOGIN, PARAM_SHORT,
	PARAM_JOBS, PARAM_BRANCH, PARAM_GIT_COMMAND, PARAM_TEST_PRO, PARAM_TEST_PRO, PARAM_TEST_FUNC, PARAM_MODEL_PATH,
	PARAM_MODELS_DIRECTORY, PARAM_KEEP_OUTPUT, PARAM_SHOW_TESTS, PARAM_TEST_NAME, PARAM_UPDATE,
	PARAM_OVERWRITE, BUILD_PATH, INSTALL_PATH, BUILD_TYPE, SOURCES_PATH, XMIPP_SOURCES, XMIPP,
	LOG_FILE, CMAKE_ERROR, MODE_GET_SOURCES, VERSION_FILE, PARAMS, LONG_VERSION, SHORT_VERSION, SEND_INSTALLATION_STATISTICS)
from installer.utils import runStreamingJob, runJob
from installer.parser import ModeHelpFormatter, GeneralHelpFormatter, ErrorHandlerArgumentParser, getParamNames
from installer.config import readConfig, writeConfig
from installer.cmake import getCMake, getCMakeVarsStr
from installer.main import getSources, exitXmipp, handleRetCode, getSectionMessage, getSuccessMessage, getVersionMessage
from installer.logger import logger, yellow, red, blue, green
from installer.api import sendApiPOST
from installer.test import runTests
from installer import modelsDLTK

####################### EXECUTION MODES #######################
def __getProjectRootDir() -> str:
	"""
	### Returns the root directory of Xmipp.

	#### Returns:
	- (str): Absolute path to Xmipp's root directory.
	"""
	return os.path.dirname(os.path.abspath(__file__))

####################### EXECUTION MODES #######################
def modeAddModel(args: argparse.Namespace):
	"""
	### Checks the params for execution mode "addModel" and then runs it.

	#### Params:
	- args (Namespace): Command line arguments parsed by argparse library.
	"""
	if not os.path.isdir(args.modelPath):
		logger("<modelsPath> is not a directory. Please, check the path. \n"
		      "The name of the model will be the name of that folder.\n")
	else:
		modelsDLTK.addModels(login=args.login, modelPath=args.modelPath, update=args.update)

def modeCleanBin():
	"""
	### Removes all compiled binaries.
	
	#### Params:
	- args (Namespace): Command line arguments parsed by argparse library.
	"""
	confirmationText = 'y'
	warningStr = "WARNING: This will DELETE from {} all *.so, *.os and *.o files. Also the *.pyc and *.dblite files".format(SOURCES_PATH)
	warningStr += f"\nIf you are sure you want to do this, type '{confirmationText}' (case sensitive):"
	logger(yellow(warningStr), forceConsoleOutput=True)
	try:
		userInput = input()
	except KeyboardInterrupt:
		userInput = ''
		logger("", forceConsoleOutput=True)
	if userInput == confirmationText:
		runJob(f"find {SOURCES_PATH}/* -name \"*.so\" -exec rm -rf {{}} \;", showCommand=True)
		runJob(f"find {SOURCES_PATH}/* -name \"*.os\" -exec rm -rf {{}} \;", showCommand=True)
		runJob(f"find {SOURCES_PATH}/* -name \"*.o\" -exec rm -rf {{}} \;", showCommand=True)
		runJob("find . -iname \"*.pyc\" -delete", showCommand=True)
		runJob("find . -iname \"*.dblite\" -delete", showCommand=True)
		runJob(f"find {os.path.join(SOURCES_PATH, XMIPP, 'applications', 'programs')} --type d -empty", showCommand=True)
		runJob(f"rm -rf {BUILD_PATH}", showCommand=True)
	else:
		logger(red("Operation cancelled."), forceConsoleOutput=True)

def modeCleanAll():
	"""
	### Removes all compiled binaries and cloned sources.
	
	#### Params:
	- args (Namespace): Command line arguments parsed by argparse library.
	"""
	# Print warning text and await input
	confirmationText = 'YeS'
	warningStr = "WARNING: This will DELETE ALL content from src and build, also de xmipp.conf file.\n"
	warningStr += "\tNotice that if you have unpushed changes, they will be deleted.\n"
	warningStr += f"\nIf you are sure you want to do this, type '{confirmationText}' (case sensitive):"
	logger(yellow(warningStr), forceConsoleOutput=True)

	try:
		userInput = input()
	except KeyboardInterrupt:
		userInput = ''
		logger("", forceConsoleOutput=True)
	if userInput == confirmationText:
		# Get xmipp sources
		xmippSources = [os.path.join(SOURCES_PATH, source) for source in XMIPP_SOURCES]
		# Get installation path
		configDict = readConfig(CONFIG_FILE) if os.path.exists(CONFIG_FILE) else {}
		installDir = configDict.get(CMAKE_INSTALL_PREFIX, INSTALL_PATH)
		runJob(f"rm -rf {BUILD_PATH} {installDir} {' '.join(xmippSources)} {CONFIG_FILE}", showCommand=True)
		runJob(f"rm -rf {CONFIG_FILE} {BUILD_PATH}", showCommand=True)
	else:
		logger(red("Operation cancelled."), forceConsoleOutput=True)

def modeCompileAndInstall(args: argparse.Namespace, configDict: Dict={}):
	"""
	### Checks the params for execution mode "compileAndInstall" and then runs it.
	
	#### Params:
	- args (Namespace): Command line arguments parsed by argparse library.
	- configDict (dict): Optional. If provided, it will use it's variables. If not, it will read from config file.
	"""
	# If config variable dictionary is not received this mode is the first being executed
	if not configDict:
		logger.startLogFile(LOG_FILE)
		configDict = readConfig(CONFIG_FILE)
		if not getCMake(configDict):
			apiPost = True if configDict.get(SEND_INSTALLATION_STATISTICS) == 'ON' else False
			handleRetCode(CMAKE_ERROR, predefinedErrorCode=CMAKE_ERROR, sendAPI=apiPost)

	# Compile with CMake
	cmakeExecutable = getCMake(configDict)
	logger(getSectionMessage("Compiling with CMake"), forceConsoleOutput=True)
	retCode = runStreamingJob(f"{cmakeExecutable} --build {BUILD_PATH} --config {BUILD_TYPE} -j {args.jobs}",
													 showOutput=True, substitute=True)
	apiPost = True if configDict.get(
		SEND_INSTALLATION_STATISTICS) == 'ON' else False
	handleRetCode(retCode, predefinedErrorCode=CMAKE_COMPILE_ERROR, sendAPI=apiPost)
	
	# Install with CMake
	logger(getSectionMessage("Installing with CMake"), forceConsoleOutput=True)
	retCode = runStreamingJob(f"{cmakeExecutable} --install {BUILD_PATH} --config {BUILD_TYPE}",
													 showOutput=True, substitute=True)
	handleRetCode(retCode, predefinedErrorCode=CMAKE_INSTALL_ERROR, sendAPI=apiPost)
	
def modeConfigBuild(configDict: Dict={}, sendAPIPost: bool=False):
	"""
	### Configures the project using CMake.
	
	#### Params:
	- args (Namespace): Command line arguments parsed by argparse library.
	- configDict (dict): Optional. If provided, it will use it's variables. If not, it will read from config file.
	"""
	# If config variable dictionary is not received this mode is the first being executed
	if not configDict:
		logger.startLogFile(LOG_FILE)
		configDict = readConfig(CONFIG_FILE) if os.path.exists(CONFIG_FILE) else modeConfig()
	
	# Check if CMake exists
	cmakeExecutable = getCMake(configDict)
	if not getCMake(configDict):
		handleRetCode(CMAKE_ERROR, predefinedErrorCode=CMAKE_ERROR, sendAPI=sendAPIPost)
	
	logger(getSectionMessage("Configuring with CMake"), forceConsoleOutput=True)
	configureCmd = f"{cmakeExecutable} -S . -B {BUILD_PATH} -D CMAKE_BUILD_TYPE={BUILD_TYPE}"
	configureCmd += f" {getCMakeVarsStr(configDict)}"
	retCode = runStreamingJob(configureCmd, showOutput=True, substitute=True)
	handleRetCode(retCode, predefinedErrorCode=CMAKE_CONFIGURE_ERROR, sendAPI=sendAPIPost)

def modeConfig(overwrite: bool=False) -> Dict:
	"""
	### Generates a template config file.
	
	#### Params:
	- overwrite (bool): If True, file is created from scratch with default values.

	#### Returns:
	- (dict): Dictionary containig all config variables.
	"""
	configDict = {}
	logger(getSectionMessage("Managing config file"), forceConsoleOutput=True)
	if not overwrite and os.path.exists(CONFIG_FILE):
		logger('Reading config file...', forceConsoleOutput=True)
		configDict = readConfig(CONFIG_FILE)
	else:
		logger('Generating config file from scratch with default values...', forceConsoleOutput=True)
	writeConfig(CONFIG_FILE, configDict=configDict)
	logger(green("Done"), forceConsoleOutput=True)
	return readConfig(CONFIG_FILE)

def modeGetModels(args: argparse.Namespace):
	"""
	### Checks the params for execution mode "getModels" and then runs it.
	
	#### Params:
	- args (Namespace): Command line arguments parsed by argparse library.
	"""
	distPath = os.path.join(__getProjectRootDir(), DEFAULT_MODELS_DIR)
	modelsDLTK.downloadDeepLearningModels(args.directory, distPath)

def modeGetSources(args: argparse.Namespace):
	"""
	### Downloads all Xmipp's sources.

	#### Params:
	- args (Namespace): Command line arguments parsed by argparse library.
	"""
	# Clone sources from specified branch
	getSources(branch=args.branch)

def modeGit(args: argparse.Namespace):
	"""
	### Executes the given git command into all xmipp source repositories.
	
	#### Params:
	- args (Namespace): Command line arguments parsed by argparse library.
	"""
	cmd = f"git {' '.join(args.command)}"
	logger(f"Running command '{cmd}' for all xmipp sources...", forceConsoleOutput=True)
	XMIPP_SOURCES.insert(0, XMIPP)
	for source in XMIPP_SOURCES:
		logger("", forceConsoleOutput=True)
		# Check if source exists to perform command, else skip
		sourcePath = os.path.abspath(os.path.join(SOURCES_PATH, source))
		if not os.path.exists(sourcePath):
			logger(yellow(f"WARNING: Source {source} does not exist in path {sourcePath}. Skipping."), forceConsoleOutput=True)
			continue
		logger(blue(f"Running command for {source} in path {sourcePath}..."), forceConsoleOutput=True)
		runJob(cmd, cwd=sourcePath, showOutput=True, showError=True)

def modeTest(args: argparse.Namespace):
	"""
	### Checks the params for execution mode "test" and then runs it.
	
	#### Params:
	- parser (ErrorHandlerArgumentParser): Parser object used to parse the arguments.
	- args (Namespace): Command line arguments parsed by argparse library.
	"""
	if args.show == True:
		logger("Showing test--------------------------------------", forceConsoleOutput=True)
		runTests(PARAMS[PARAM_SHOW_TESTS][LONG_VERSION])
	elif args.allPrograms:
		logger("Running all tests--------------------------------------", forceConsoleOutput=True)
		runTests(PARAMS[PARAM_TEST_PRO][LONG_VERSION])
	elif args.allFuncs:
		runTests(PARAMS[PARAM_TEST_FUNC][LONG_VERSION])

	elif args.testName:
		logger("Running test {}-------------------------------".format(args.testName), forceConsoleOutput=True)
		runTests(args.testName)

def modeVersion(args: argparse.Namespace):
	"""
	### Checks the params for execution mode "version" and then runs it.

	#### Params:
	- args (Namespace): Command line arguments parsed by argparse library.
	"""
	logger(getVersionMessage(short=args.short), forceConsoleOutput=True)

def modeAll(args: argparse.Namespace):
	"""
	### Runs get sources & modes config, configBuild, and compileAndInstall
	
	#### Params:
	- args (Namespace): Command line arguments parsed by argparse library.
	"""
	# Initiate logger
	logger.startLogFile(LOG_FILE)

	# Read and/or write variables from config file
	configDict = modeConfig()

	# Clone sources from specified branch
	modeGetSources(args)

	# Configure with CMake
	modeConfigBuild(configDict=configDict, sendAPIPost=True)

	# Compile and install with CMake
	modeCompileAndInstall(args, configDict=configDict)

def runSelectedMode(parser: ErrorHandlerArgumentParser, args: argparse.Namespace):
	"""
	### This function runs the selected execution mode.

	#### Params:
	- parser (ErrorHandlerArgumentParser): Parser object used to parse the arguments.
	- args (Namespace): Command line arguments parsed by argparse library.
	"""
	sendAPI = False
	if args.mode == MODE_ADD_MODEL:
		modeAddModel(args)
	elif args.mode == MODE_ALL:
		sendAPI = True
		modeAll(args)
	elif args.mode == MODE_CLEAN_ALL:
		modeCleanAll()
	elif args.mode == MODE_CLEAN_BIN:
		modeCleanBin()
	elif args.mode == MODE_COMPILE_AND_INSTALL:
		modeCompileAndInstall(args)
		sendAPI = True
	elif args.mode == MODE_CONFIG_BUILD:
		sendAPI = True
		modeConfigBuild()
	elif args.mode == MODE_CONFIG:
		modeConfig(overwrite=args.overwrite)
	elif args.mode == MODE_GET_MODELS:
		modeGetModels(args)
	elif args.mode == MODE_GET_SOURCES:
		modeGetSources(args)
	elif args.mode == MODE_GIT:
		modeGit(args)
	elif args.mode == MODE_TEST:
		modeTest(args)
	elif args.mode == MODE_VERSION:
		modeVersion(args)
	else:
		# If method was none of the above, exit with error
		logger(red(f"Mode \"{args.mode}\" not recognized. {COMMON_USAGE_HELP_MESSAGE}"), forceConsoleOutput=True)
		exitXmipp(retCode=1)
	
	# Send API message
	try:
		apiPost = True if readConfig(CONFIG_FILE).get(
			SEND_INSTALLATION_STATISTICS) == 'ON' else False
		if sendAPI and apiPost and os.path.exists(VERSION_FILE):
			sendApiPOST()
	except FileNotFoundError:
		pass

	# Print success message for specific modes
	if args.mode == MODE_ALL or args.mode == MODE_COMPILE_AND_INSTALL:
		logger(getSuccessMessage(), forceConsoleOutput=True)
	exitXmipp()
	
####################### MAIN EXECUTION THREAD #######################
if __name__ == "__main__":
	""" Calls main function when executed. """
	# Defining default jobs: 120% current thread count (not all jobs take 100% of CPU time continuously)
	JOBS = multiprocessing.cpu_count() + int(multiprocessing.cpu_count() * 0.2)

	# Creating parser to parse the command-line arguments
	parser = ErrorHandlerArgumentParser(formatter_class=GeneralHelpFormatter, prog="xmipp")

	# Adding subparsers to have other variables deppending on the value of the mode
	subparsers = parser.add_subparsers(dest="mode")

	# Arguments for mode 'addModel'
	addModelSubparser = subparsers.add_parser(MODE_ADD_MODEL, formatter_class=ModeHelpFormatter)
	addModelSubparser.add_argument(*getParamNames(PARAM_LOGIN))
	addModelSubparser.add_argument(*getParamNames(PARAM_MODEL_PATH))
	addModelSubparser.add_argument(*getParamNames(PARAM_UPDATE), action='store_true')

	# Arguments for mode 'all'
	allSubparser = subparsers.add_parser(MODE_ALL, formatter_class=ModeHelpFormatter)
	allSubparser.add_argument(*getParamNames(PARAM_JOBS), type=int, default=JOBS)
	allSubparser.add_argument(*getParamNames(PARAM_BRANCH))
	allSubparser.add_argument(*getParamNames(PARAM_KEEP_OUTPUT), action='store_true')

	# Arguments for mode 'cleanAll'
	cleanAllSubparser = subparsers.add_parser(MODE_CLEAN_ALL, formatter_class=ModeHelpFormatter)

	# Arguments for mode 'cleanBin'
	cleanBinSubparser = subparsers.add_parser(MODE_CLEAN_BIN, formatter_class=ModeHelpFormatter)

	# Arguments for mode 'compileAndInstall'
	compileAndInstallSubparser = subparsers.add_parser(MODE_COMPILE_AND_INSTALL, formatter_class=ModeHelpFormatter)
	compileAndInstallSubparser.add_argument(*getParamNames(PARAM_JOBS), type=int, default=JOBS)
	compileAndInstallSubparser.add_argument(*getParamNames(PARAM_BRANCH))
	compileAndInstallSubparser.add_argument(*getParamNames(PARAM_KEEP_OUTPUT), action='store_true')

	# Arguments for mode 'configBuild'
	buildConfigSubparser = subparsers.add_parser(MODE_CONFIG_BUILD, formatter_class=ModeHelpFormatter)
	buildConfigSubparser.add_argument(*getParamNames(PARAM_KEEP_OUTPUT), action='store_true')

	# Arguments for mode 'config'
	configSubparser = subparsers.add_parser(MODE_CONFIG, formatter_class=ModeHelpFormatter)
	configSubparser.add_argument(*getParamNames(PARAM_OVERWRITE), action='store_true')

	# Arguments for mode 'getModels'
	getModelsSubparser = subparsers.add_parser(MODE_GET_MODELS, formatter_class=ModeHelpFormatter)
	getModelsSubparser.add_argument(*getParamNames(PARAM_MODELS_DIRECTORY), default=os.path.join(__getProjectRootDir(), DEFAULT_MODELS_DIR))

	# Arguments for mode 'getSources'
	getSourcesSubparser = subparsers.add_parser(MODE_GET_SOURCES, formatter_class=ModeHelpFormatter)
	getSourcesSubparser.add_argument(*getParamNames(PARAM_BRANCH))
	getSourcesSubparser.add_argument(*getParamNames(PARAM_KEEP_OUTPUT), action='store_true')

	# Arguments for mode 'git'
	gitSubparser = subparsers.add_parser(MODE_GIT, formatter_class=ModeHelpFormatter)
	gitSubparser.add_argument(*getParamNames(PARAM_GIT_COMMAND), nargs='+')

	# Arguments for mode 'test'
	testSubparser = subparsers.add_parser(MODE_TEST, formatter_class=ModeHelpFormatter)
	testSubparser.add_argument(*getParamNames(PARAM_TEST_NAME), nargs='?', default=None)
	testSubparser.add_argument(*getParamNames(PARAM_TEST_PRO), action='store_true')
	testSubparser.add_argument(*getParamNames(PARAM_TEST_FUNC), action='store_true')
	testSubparser.add_argument(*getParamNames(PARAM_SHOW_TESTS), action='store_true')

	# Arguments for mode 'version'
	versionSubparser = subparsers.add_parser(MODE_VERSION, formatter_class=ModeHelpFormatter)
	versionSubparser.add_argument(*getParamNames(PARAM_SHORT), action='store_true')

	# Applying default mode value if needed
	if len(sys.argv) == 1 or (
		len(sys.argv) > 1 and 
		sys.argv[1].startswith('-') and 
		'-h' not in sys.argv and
		'--help' not in sys.argv):
		sys.argv.insert(1, MODE_ALL)

	# Parse arguments
	args = parser.parse_args()

	# Error control for number of jobs
	if hasattr(args, 'jobs') and args.jobs < 1:
		parser.error(f"Wrong job number \"{args.jobs}\". Number of jobs has to be 1 or greater.")
	
	# Error control for branch
	if hasattr(args, "branch") and args.branch is not None and len(args.branch.split(' ')) > 1:
		parser.error(f"Incorrect branch name \"{args.branch}\". Branch names can only be one word long.")
	
	if hasattr(args, "keep_output") and args.keep_output:
		logger.setAllowSubstitution(False)

	# Running always under this own directory.
	os.chdir(__getProjectRootDir())

	# Running installer in selected mode
	runSelectedMode(parser, args)