#!/usr/bin/python3
import os
import sys
from rivuletpy.utils.io import loadswc, saveswc, loadimg
from rivuletpy.trace import estimate_radius
import argparse

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Arguments for see anisotropic filters.')
    parser.add_argument(
        '-f', '--file', type=str, default=None, help='The file to filter')
    parser.add_argument(
        '-i',
        '--img',
        type=str,
        required=False,
        default=None,
        help='The image file to estimate radius from')
    parser.add_argument(
        '-t',
        '--threshold',
        type=float,
        required=False,
        default=0,
        help='The threshold to segment the image for radius estimation.')
    parser.add_argument(
        '-z',
        '--zoom_factor',
        type=float,
        required=False,
        default=1.,
        help='Used to rescale swc points propotionally. It is only used when --action rescale'
    )

    parser.add_argument(
        '--shiftx',
        type=float,
        required=False,
        default=0.,
    )

    parser.add_argument(
        '--shifty',
        type=float,
        required=False,
        default=0.
    )

    parser.add_argument(
        '--shiftz',
        type=float,
        required=False,
        default=0.
    )
    parser.add_argument('--radius', dest='radius', action='store_true')
    parser.set_defaults(radius=False)
    parser.add_argument('--overwrite', dest='overwrite', action='store_true')
    parser.set_defaults(overwrite=False)
    parser.add_argument(
        '--action',
        type=str,
        default=None,
        required=True,
        help='The processing of swc to perform. Possible actions are [radius, rescale]'
    )
    args = parser.parse_args()

    swc = loadswc(args.file)  # Read in the swc
    outfile = args.file if args.overwrite else os.path.splitext(args.file)[
        0] + '.' + args.action + '.swc'  # Make the output path

    if swc.shape[0] == 0:
        print('SWC in ', args.file, ' is empty...')
        sys.exit(1)

    if args.action == 'radius':
        img = loadimg(args.img)
        outswc = estimate_radius(swc[:, 2:5], img > args.threshold)
    elif args.action == 'rescale':
        swc[:, 2:5] *= args.zoom_factor
        outswc = swc
    elif args.action == 'shift':
        swc[:, 2] += args.shiftx
        swc[:, 3] += args.shifty
        swc[:, 4] += args.shiftz
        outswc = swc
    else:
        print(
            'Action not defined. Available actions are [clean, flip, radius, rescale]'
        )

    saveswc(outfile, outswc)
