#!python
"""usage: connectome_precomputed [-h] -r -c -o

Generate AvgR, AvgR_Fz, and T functional connectivity chunks for a single chunk.

Arguments:
  -h, --help        show this help message and exit

  -r, --roi-dir     Path to directory containing binary Nifti ROIs, or path to a single 
                    binary Nifti ROI, or path to CSV containing binary Nifti ROI paths.

  -c, --config      Path or Name of precomputed connectome config file to use. If Path, 
                    assumed to be in current working directory or absolute path. If Name 
                    of config file, config file '[config].json' must be located in 
                    ~/.config/pfctoolkit/.

  -o, --output-dir  Path to output directory.

  -s, --stat        Statistical map choice (e.g., '(avgr, t, fz)').

"""
import os
import argparse
from tqdm import tqdm
from pfctoolkit import tools
from pfctoolkit import config
from pfctoolkit import mapping
from pfctoolkit import datasets

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Generate functional lesion network maps using the precomputed connectome."
    )

    parser.add_argument(
        "-r",
        "--roi-dir",
        metavar="\b",
        help="Path to directory containing binary Nifti ROIs, or path to a single binary Nifti ROI, or path to CSV containing binary Nifti ROI paths.",
        type=str,
        required=True,
    )

    parser.add_argument(
        "-c",
        "--config",
        metavar="\b",
        help="Path or Name of precomputed connectome config file to use. If Path, assumed to be in current working directory or absolute path. If Name of config file, config file '<config>.json' must be located in '~/.config/pfctoolkit/'.",
        type=str,
        required=True,
    )

    parser.add_argument(
        "-o",
        "--output-dir",
        metavar="\b",
        help="Path to output directory.",
        type=str,
        required=True,
    )
    
    parser.add_argument(
        "-s",
        "--stat",
        metavar="\b",
        help="Statistical map choice (e.g., '(avgr, t, fz)').",
        type=str,
        required=False,
        default='(avgr, t, fz)'
    )

    # Parse arguments
    args = parser.parse_args()
    
    # Convert the --stat argument to a tuple of strings
    stat = tuple(map(str.strip, args.stat.strip('()').split(',')))

    # Load ROI list
    roi_paths = tools.load_roi(os.path.abspath(args.roi_dir))

    # Load and check PCC configuration
    pcc_config = config.Config(args.config, stat)

    # Set output directory
    output_dir = os.path.abspath(args.output_dir)

    # Load brain mask
    brain_mask = datasets.get_img(pcc_config.get("mask"))

    # Get chunks
    chunks = tools.get_chunks(roi_paths, pcc_config)

    # Process Chunks
    atlas = {}
    for chunk in tqdm(chunks, desc='ROI Progress'):
        contribution = mapping.process_chunk(chunk, chunks[chunk], pcc_config, stat)
        atlas = mapping.update_atlas(contribution, atlas, stat)

    # Consolidate outputs
    mapping.publish_atlas(atlas, output_dir, pcc_config, stat)
