#!python

__author__ = 'SungHo Lee (shlee@unc.edu)'
__version_info__ = ('2017', '07', '25')

import os, argparse
from pynit.tools import methods, messages
from pynit.tools.visualizers import BrainPlot
from pynit.handler import images


def path_parser(args):
    print(args.base)
    print(args.input)
    b_img = load(args.base)
    i_img = load(args.input)
    o_path = check_path_avail(args.output)
    flip = check_flip(args.flip)
    return b_img, i_img, o_path, flip


def check_flip(args):
    flip_handler = dict()
    if isinstance(args, list):
        if 'x' in args:
            flip_handler['invertx']=True
        if 'y' in args:
            flip_handler['inverty']=True
        if 'z' in args:
            flip_handler['invertz']=True
        if flip_handler.keys():
            return flip_handler
        else:
            return None
    else:
        return None


def load(path):
    if os.path.exists(path):
        if os.path.isfile(path):
            try:
                img = images.load(path)
                return img
            except:
                methods.raiseerror(messages.Errors.InputTypeError, '{} is not available format, \
                please use Nifti-1 format instead.'.format(path))
        else:
            methods.raiseerror(messages.Errors.InputTypeError, '{} is not file, \
            please input correct path.'.format(path))
    else:
        methods.raiseerror(messages.InputPathError, 'File is not exist, please double check the input path.')


def check_path_avail(path):
    if os.path.exists(path):
        methods.raiseerror(messages.InputFileError, "File is exist, please use different filename.")
    else:
        output_dir = os.path.dirname(path)
        filename = os.path.basename(path)
        if os.path.exists(output_dir):
            pass
        else:
            try:
                methods.mkdir(output_dir)
            except:
                methods.raiseerror(messages.InputPathError, '"{}" is not exist, \
                cannot make output folder!'.format(os.path.dirname(output_dir)))
        if '.png' in filename:
            pass
        else:
            filename = filename+'.png'
        return os.path.join(output_dir, filename)


def main():
    parser = argparse.ArgumentParser(prog='check_reg', description="Check registration results")
    parser.add_argument("base", help="Path of under-layer image", type=str)
    parser.add_argument("input", help="Path of over-layer image", type=str)
    parser.add_argument("output", help="Output path", type=str)
    parser.add_argument("-f", "--flip", nargs='*', help='Filp axis', default=None)
    args = parser.parse_args()
    b_img, i_img, o_path, flip = path_parser(args)
    print(type(b_img), type(i_img), o_path)
    if flip:
        fig, _ = BrainPlot.check_reg(b_img, i_img, **flip)
    else:
        fig, _ = BrainPlot.check_reg(b_img, i_img)
    fig.savefig(o_path, facecolor='black', edgecolor='black')

if __name__ == '__main__':
    main()