#!/usr/bin/env python3
"""usage: generate_weighted_network_linklist [-h] -r -c -o

Generate weighted network matrix files for clustering use with infomap.

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

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

  -c, --config      Name of precomputed connectome config file to use.

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

"""

import os
import argparse

from tqdm import trange
from pfctoolkit import tools, config, datasets, clustering

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Generate weighted network link-lists for clustering with infomap."
    )

    parser.add_argument(
        "-r",
        "--roi",
        metavar="\b",
        help="Path to a single binary Nifti ROI.",
        type=str,
        required=True,
    )

    parser.add_argument(
        "-c",
        "--config",
        metavar="\b",
        help="Name of precomputed connectome config file to use.",
        type=str,
        required=True,
    )

    parser.add_argument(
        "-o",
        "--output-dir",
        metavar="\b",
        help="Path to output directory.",
        type=str,
        required=True,
    )

    # Parse arguments
    args = parser.parse_args()

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

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

    # 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, pcc_config)

    roi_weight_matrix, roi_size = clustering.generate_weighted_network_matrix(
        roi, chunks, pcc_config
    )

    with open(
        os.path.join(
            output_dir, f"{os.path.basename(roi[0]).split('.nii')[0]}_linklist.txt"
        ),
        "w",
    ) as fp:
        for i in trange(roi_size - 1):
            for j in range(i + 1, roi_size):
                fp.write(f"{i+1} {j+1} {roi_weight_matrix[i,j]}\n")
