#!/usr/bin/env python

# This file is used for running AxCaliber analysis of diffusion MRI data
# This file get command line arguments and start the analysis

import argparse
import logging
import multiprocessing
from pathlib import Path

from AxSI import config
from AxSI.AxSI import axsi_main
from AxSI.toolbox import set_log_config, monitor_memory

logger = logging.getLogger(__name__)


def set_multiprocessing_method():
    # For Windows/macOS. Ensures child processes inherit a copy of the Config with updated dynamic values.
    multiprocessing.set_start_method('fork', True)
    logger.debug(f'Multiprocessing method is: {multiprocessing.get_start_method()}')  # Verify the current method.


# Run with the default values of argparse
def run_from_python(subj_folder, data_path, mask_path, bval_path, bvec_path, small_delta=15, big_delta=45, gmax=7.9,
                    gamma_val=4257):
    subj_folder = assert_path(subj_folder)
    file_names = build_files_dict(data_path, mask_path, bval_path, bvec_path)

    if subj_folder:
        logger.info(f"subj_folder: {subj_folder} \nfilenames: {file_names} \nsmall_delta: {small_delta} \nbig_delta: "
                    f"{big_delta} \ngmax: {gmax} \ngamma_val: {gamma_val}")
        axsi_main(subj_folder, file_names, small_delta, big_delta, gmax, gamma_val)


def build_files_dict(*paths):
    keys = ["data", "mask", "bval", "bvec"]
    filenames = {}
    if len(keys) != len(paths):
        logger.error("Wrong number of input paths, make sure paths to data, mask, bval and bvec exists")
    for i in range(len(keys)):
        filenames[keys[i]] = assert_path(paths[i])
    return filenames


def assert_path(s: str):
    p = Path(s)
    if not p.exists():
        logger.error(f"Path given is not valid: {s}")
        raise ValueError(f"Path given is not valid: {s}")
    return p


def validate_args(args):
    if (args.num_processes_pred > 1 and args.num_threads_pred > 1) or (
            args.num_processes_axsi > 1 and args.num_threads_axsi > 1):
        raise ValueError("Input parameters are invalid: you cannot have both more than one process and more than one "
                         "thread simultaneously in prediction step or in axsi step.")
    if args.num_threads_axsi > 1 and (
            args.linear_lsq_method == 'R-quadprog' or args.nonlinear_lsq_method == 'R-minpack'):
        raise ValueError("Invalid input parameters: currently, running R packages with multiple threads is not "
                         "supported.")
    args.subj_folder = assert_path(args.subj_folder)


def update_config_file(args):
    config.SUBJ_FOLDER = args.subj_folder
    config.NUM_PROCESSES_PRED = args.num_processes_pred
    config.NUM_THREADS_PRED = args.num_threads_pred
    config.NUM_PROCESSES_AXSI = args.num_processes_axsi
    config.NUM_THREADS_AXSI = args.num_threads_axsi
    config.DEBUG_MODE = args.debug_mode
    config.LINEAR_LSQ_METHOD = args.linear_lsq_method
    config.NON_LINEAR_LSQ_METHOD = args.nonlinear_lsq_method
    # config.R_SERVER_PORT = args.rserver_port # for Rserver - not in use, it is very slow


def run_from_terminal():
    # Create ArgumentParser object
    parser = argparse.ArgumentParser(description='AxSI for MRI data')

    # Define command-line arguments
    parser.add_argument('--subj-folder', type=Path, help='Path to the subject folder', required=True)
    parser.add_argument('--data', type=Path, help='Path to the data file', required=True)
    parser.add_argument('--bval', type=Path, help='Path to the bval file', required=True)
    parser.add_argument('--bvec', type=Path, help='Path to the bvec file', required=True)
    parser.add_argument('--mask', type=Path, help='Path to the mask file', required=True)
    parser.add_argument('--small-delta', type=float, default=15, help='Gradient duration in miliseconds')
    parser.add_argument('--big-delta', type=float, default=45, help='Time to scan (time interval) in milisecond')
    parser.add_argument('--gmax', type=float, default=7.9, help='Gradient maximum amplitude in G/cm')
    parser.add_argument('--gamma-val', type=int, default=4257, help='Gyromagnetic ratio')
    parser.add_argument('--num-processes-pred', type=int, default=1,
                        help='Number of processes to run in parallel in prediction step')
    parser.add_argument('--num-threads-pred', type=int, default=1,
                        help='Number of threads to run in parallel in prediction step')
    parser.add_argument('--num-processes-axsi', type=int, default=1,
                        help='Number of processes to run in parallel in AxSI step')
    parser.add_argument('--num-threads-axsi', type=int, default=1,
                        help='Number of threads to run in parallel in AxSI step')
    parser.add_argument('--nonlinear-lsq-method', type=str, default='R-minpack',
                        choices=['R-minpack', 'scipy', 'lsq-axsi'],
                        help="Method for linear least squares. Choose from 'R-minpack', 'scipy', or 'lsq-axsi'. "
                             "Default is 'R-minpack'.")
    parser.add_argument('--linear-lsq-method', type=str, default='R-quadprog',
                        choices=['R-quadprog', 'gurobi', 'scipy', 'cvxpy'],
                        help="Method for linear least squares. Choose from 'R-quadprog', 'gurobi', 'scipy', or 'cvxpy'. "
                             "Default is 'R-quadprog'.")
    parser.add_argument('--debug-mode', action='store_true', help='Enable debug mode (default is disabled).')
    # parser.add_argument('--rserver-port', type=int, help='Port number for R server.', default=6311) # for Rserver - not in use, it is very slow

    # Parse the command-line arguments
    args = parser.parse_args()
    validate_args(args)

    file_names = build_files_dict(args.data, args.mask, args.bval, args.bvec)
    small_delta = args.small_delta
    big_delta = args.big_delta
    gmax = args.gmax
    gamma_val = args.gamma_val

    update_config_file(args)
    set_log_config()
    set_multiprocessing_method()

    # R
    # set_r_environment()
    # for Rserver - not in use, it is very slow
    # server_socket_rserver = upload_rserver(config.R_SERVER_PORT)
    # if not server_socket_rserver:
    #     raise ValueError(f"The R server port {config.R_SERVER_PORT} is not available. Please select a different.")

    if args.debug_mode:
        # Start the monitor_memory function as a separate process
        monitor_process = multiprocessing.Process(target=monitor_memory)
        monitor_process.start()
    try:
        if args.subj_folder:
            logger.info("Starting AxSI")
            logger.info(f"The parameters are: subj_folder: {args.subj_folder} \nfilenames: {file_names} \nsmall_delta: "
                        f"{small_delta} \nbig_delta: {big_delta} \ngmax: {gmax} \ngamma_val: {gamma_val}"
                        f"\nnonlinear-lsq-method: {args.nonlinear_lsq_method} \nlinear-lsq-method: "
                        f"{args.linear_lsq_method} \nnum-processes-pred: {args.num_processes_pred} "
                        f"\nnum-threads-pred: {args.num_threads_pred}  \nnum-processes-axsi: {args.num_processes_axsi} "
                        f"\nnum-threads-axsi: {args.num_threads_axsi}")
            axsi_main(args.subj_folder, file_names, small_delta, big_delta, gmax, gamma_val)
    finally:
        # for Rserver - not in use, it is very slow
        # print(f"Port {config.R_SERVER_PORT} has been closed.")
        # server_socket_rserver.close()  # Ensure the port of R server is freed by closing the socket
        if args.debug_mode:
            # Ensure the memory monitor stops when the script ends
            monitor_process.terminate()
            monitor_process.join()
            logger.debug("[Main] Memory monitor terminated.")


if __name__ == "__main__":
    run_from_terminal()

# Example of run commenad (run from AxSI_python folder):

# Full input files:
# =================
# python AxSI_main.py --subj-folder full-output-temp --bval '/Users/user/Library/CloudStorage/GoogleDrive-refaelkohen@mail.tau.ac.il/My Drive/TLV-U-drive/BrainWork/AxSI-pipeline/input-data/sub-CLMC10/ses-202407110849/dwi/sub-CLMC10_ses-202407110849_space-dwi_desc-preproc_dwi.bval' --bvec '/Users/user/Library/CloudStorage/GoogleDrive-refaelkohen@mail.tau.ac.il/My Drive/TLV-U-drive/BrainWork/AxSI-pipeline/input-data/sub-CLMC10/ses-202407110849/dwi/sub-CLMC10_ses-202407110849_space-dwi_desc-preproc_dwi.bvec' --mask '/Users/user/Library/CloudStorage/GoogleDrive-refaelkohen@mail.tau.ac.il/My Drive/TLV-U-drive/BrainWork/AxSI-pipeline/input-data/sub-CLMC10/ses-202407110849/dwi/sub-CLMC10_ses-202407110849_space-dwi_desc-brain_mask.nii.gz' --data '/Users/user/Library/CloudStorage/GoogleDrive-refaelkohen@mail.tau.ac.il/My Drive/TLV-U-drive/BrainWork/AxSI-pipeline/input-data/sub-CLMC10/ses-202407110849/dwi/sub-CLMC10_ses-202407110849_space-dwi_desc-preproc_dwi.nii.gz' --debug-mode --num-processes 35 --num-threads 1

# 4-timeframes (0,1,5,8) input files (bvals: 0,0, 1,1):
# ======================================================
# python AxSI_main.py --subj-folder tests/short-output --bval 'tests/test_data/part_times/test_vol_0_1_5_8_bval.txt' --bvec 'tests/test_data/part_times/test_vol_0_1_5_8_bvec.txt' --data 'tests/test_data/part_times/test_vol_0_1_5_8_data.nii' --mask 'tests/test_data/part_times/test_vol_0_1_5_8_mask.nii'

# 6-timeframes (0,1,5,8,80,81) input files (bvals: 0,0, 1,1, 2,2):
# ======================================================
# python AxSI_main.py --subj-folder tests/short-output --bval 'tests/test_data/part_times/test_vol_0_1_5_8_80_81_bval.txt' --bvec 'tests/test_data/part_times/test_vol_0_1_5_8_80_81_bvec.txt' --data 'tests/test_data/part_times/test_vol_0_1_5_8_80_81_data.nii' --mask 'tests/test_data/part_times/test_vol_0_1_5_8_80_81_mask.nii'

# 15000 voxels:
# ============

# windows:
# -------
# $s = 50
# $e = 52

# mac:
# ---
# s=50
# e=52

# python AxSI_main.py --subj-folder 15-voxels-output --bval "toy_tests/test_data/part_voxels/voxels_${s}_${e}_bval.txt" --bvec "toy_tests/test_data/part_voxels/voxels_${s}_${e}_bvec.txt" --data "toy_tests/test_data/part_voxels/voxels_${s}_${e}_data.nii" --mask "toy_tests/test_data/part_voxels/voxels_${s}_${e}_mask.nii"
