from nabu.pipeline.config_validators import *

from .validators import *

# This should be moved to nabu.pipeline.xrdct in the future
integrator_config = {
    "dataset": {
        "location": {
            "default": "",
            "help": "Dataset(s) to integrate. This can be a comma-separated list, with wildcards. The HDF5-Nexus files have to contain entries like '1.1'.",
            "validator": no_validator,
            "type": "required",
        },
        "hdf5_entry": {
            "default": "",
            "help": "Entry in the HDF5 file, eg. '1.1'. Default (empty) means that ALL the entries will be processed.",
            "validator": optional_string_validator,
            "type": "advanced",
        },
        "layout": {
            "default": "ID15A",
            "help": "Which folders layout to use for parsing datasets and creating output files. Available are: ID11, ID15A, ID31.",
            "validator": xrdct_layout_validator,
            "type": "advanced",
        },
    },
    "azimuthal integration": {
        "n_pts": {
            "default": "2500",
            "help": "Number of azimuthal bins.",
            "validator": positive_integer_validator,
            "type": "required",
        },
        "unit": {
            "default": "r_mm",
            "help": "Radial unit. Can be r_mm, q_A^-1, 2th_deg.",
            "validator": xrdct_integration_unit_validator,
            "type": "required",
        },
        "detector": {
            "default": "",
            "help": "Detector name used in pyFAI (ex. eiger2_cdte_4m) or full path to the a detector distortion file (ex. path/id15_pilatus.h5).",
            "validator": nonempty_string_validator,
            "type": "required",
        },
        "detector_name": {
            "default": "",
            "help": "Lima name used for detector in HDF5 file(s).",
            "validator": nonempty_string_validator,
            "type": "required",
        },
        "detector_suffix": {
            "default": "",
            "help": "Suffix to use after detector name that can be found in the measurements section of the HDF5 file (Ex: raw_frame).",
            "validator": optional_string_validator,
            "type": "optional",
        },
        "mask_file": {
            "default": "",
            "help": "Path to the mask file.",
            "validator": file_location_validator,
            "type": "required",
        },
        "flatfield_file": {
            "default": "",
            "help": "Path to the flatfield file. If provided, flat-field normalization will be applied.",
            "validator": optional_file_location_validator,
            "type": "required",
        },
        "dark_file": {
            "default": "",
            "help": "Path to the dark file. If provided, dark current will be subtracted from the raw data.",
            "validator": optional_file_location_validator,
            "type": "required",
        },
        "poni_file": {
            "default": "",
            "help": "Path to the pyFAI calibration file.",
            "validator": file_location_validator,
            "type": "required",
        },
        "error_model": {
            "default": "hybrid",
            "help": "Error model for azimuthal integration. Can be hybrid, poisson, or None.",
            "validator": xrdct_error_model_validator,
            "type": "optional",
        },
        "azimuthal_range": {
            "default": "(-180., 180., 1)",
            "help": "azimuthal ranges in the form (min, max, n_slices).",
            "validator": xrdct_azim_range_validator,
            "type": "advanced",
        },
        "radial_range": {
            "default": "",
            "help": "Lower and upper range of the radial unit. If not provided, range is simply (data.min(), data.max()). Values outside the range are ignored.",
            "validator": optional_tuple_of_two_floats_validator,
            "type": "advanced",
        },
        "ai_method": {
            "default": "opencl",
            "help": "Which azimuthal integration method to use.",
            "validator": xrdct_ai_method_validator,
            "type": "advanced",
        },
        "polarization_factor": {
            "default": 0.99,
            "help": "Polarization factor for PyFAI. Default is 1.0.",
            "validator": float_validator,
            "type": "advanced",
        },
        "correct_solid_angle": {
            "default": "1",
            "help": "Whether to correct solid angle. Put this parameter if you wish to correct solid angle.",
            "validator": boolean_validator,
            "type": "advanced",
        },
        "correct_detector_attenuation": {
            "default": "1",
            "help": "Whether to correct for attenuation caused by detector non-zero thickness.",
            "validator": boolean_validator,
            "type": "advanced",
        },
        "average_xy": {
            "default": "0",
            "help": "Whether to additionally compute the mean/median of the integrated stacks. Can be 0 (disabled), 'mean' or 'median'.",
            "validator": average_xy_validator,
            "type": "advanced",
        },
        "pixel_splitting": {
            "default": "0",
            "help": "Whether to perform pixel splitting. Possible values are: no (or 0), BBox, pseudo, full.",
            "validator": pixel_splitting_validator,
            "type": "advanced",
        },
        "trim_method": {
            "default": "",
            "help": "Which outliers removal method to use. Can be none (default), median, or sigma clip. NB: when using one of the latter two, neither azimuthal caking nor uncertainty estimation are possible, only one azimuthal slice is used.",
            "validator": xrdct_trim_method_validator,
            "type": "optional",
        },
        "trim_n_pts": {
            "default": "",
            "help": "Number of azimuthal bins for trimmed mean.",
            "validator": optional_integer_validator,
            "type": "optional",
        },
        "trim_bounds": {
            "default": "",
            "help": "Bounds for trim methods:\n - For median: percentiles in the form (cut_low, cut_high). Integration uses medfilt1d with only one azimuthal slice.\n - sigma clip: keep only pixels with intensity |I - mean(I)| < thres * std(I).",
            "validator": xrdct_trim_bounds_validator,
            "type": "optional",
        },
    },
    "computations distribution": {
        "partition": {
            "default": "gpu",
            "help": "Name of the SLURM partition (queue).",
            "validator": nonempty_string_validator,
            "type": "optional",
        },
        "n_workers": {
            "default": "4",
            "help": "Number of workers to use. If partition != local, it corresponds to the number of SLURM jobs submitted.",
            "validator": integer_validator,
            "type": "required",
        },
        "cores_per_worker": {
            "default": "4",
            "help": "Number of CPU cores (threads) per workers to use, mostly for LZ4 decompression of data.",
            "validator": integer_validator,
            "type": "advanced",
        },
        "ai_engines_per_worker": {
            "default": "8",
            "help": "Number of AI engines per worker. Each AI engine is spawned in a process with 'cores_per_worker' threads.",
            "validator": integer_validator,
            "type": "advanced",
        },
        "time": {
            "default": "01:00:00",
            "help": "Time limit for SLURM job duration.",
            "validator": nonempty_string_validator,
            "type": "optional",
        },
        "memory_per_worker": {
            "default": "100GB",
            "help": "Amount of memory per worker. Default is 100 GB.",
            "validator": nonempty_string_validator,
            "type": "advanced",
        },
        "python_executable": {
            "default": "/cvmfs/tomo.esrf.fr/software/packages/ubuntu24.04/x86_64/tomotools/stable/bin/python",
            "help": "Defines a python executable to use for each SLURM partition. Can determine using `whereis python` when environment is activated.",
            "validator": no_validator,
            "type": "advanced",
        },
        "workspace_path": {
            "default": "",
            "help": "path to the 'worker space', a directory where integrator stores files useful for distribution/communication with workers. Empty means in the same directory as the configuration file.",
            "validator": optional_directory_validator,
            "type": "advanced",
        },
    },
    "output": {
        "location": {
            "default": "",
            "help": "Path to the output file. If not provided, it will be in the same directory as the input dataset. NOTA: a directory with the same name (less the extension) will be created at the same level, for storing the actual integrated data.",
            "validator": optional_output_file_or_directory_validator,
            "type": "optional",
        },
        "existing_output": {
            "default": "reprocess_if_conffile_more_recent",
            "help": "What to do if output already exists. Possible values are:\n - skip: go to next image of the output already exists (and has the same processing configuration)\n - reprocess_if_conffile_more_recent: re-do the processing if the configuration file was edited after the integration file \n - overwrite: re-do the processing, overwrite the file\n - raise: raise an error and exit.",
            "validator": xrdct_existing_output_validator,
            "type": "optional",
        },
        "repack_output": {
            "default": "1",
            "help": "Whether to repack output data, meaning transform the virtual datasets to a contiguous dataset. If activated, the partial result files are deleted.",
            "validator": boolean_validator,
            "type": "optional",
        },
        "partial_files_subfolder": {
            "default": "",
            "help": "Subfolder where the partial integration files (one per acquisition file) should be saved. If 'repack_output' is set to 1, these partial files should disappear at the end of the processing.",
            "validator": optional_string_validator,
            "type": "advanced",
        },
        "file_mode": {
            "default": "755",
            "help": "Which file mode to use when creating new files and directories. The value must be a octal number like provided to the 'chmod' command, eg. 775, 777, 755, ...",
            "validator": positive_integer_validator,
            "type": "optional",
        },
        "try_metadata": {
            "default": "",
            "help": "Which metadata (eg. motors positions, diodes readings) should be propagated into the output file. This should be a list of comma-separated values.",
            "validator": no_validator,
            "type": "optional",
        },
    },
    "pipeline": {
        "verbosity": {
            "default": "2",
            "help": "Level of verbosity of the processing. 0 = terse, 3 = more information (good for debugging).",
            "validator": logging_validator,
            "type": "optional",
        },
    },
}

renamed_keys = {
    "detector": {
        "section": "dataset",
        "new_name": "detector",
        "new_section": "azimuthal integration",
        "since": "2021.2.0",
        "message": "Option 'detector' was moved from section [dataset] to section [azimuthal integration]",
    },
    "mask_file": {
        "section": "dataset",
        "new_name": "mask_file",
        "new_section": "azimuthal integration",
        "since": "2021.1.0",
        "message": "Option 'mask_file' was moved from section [dataset] to section [azimuthal integration]",
    },
}
