#!/usr/bin/env python2
# -*- coding: utf-8 -*-

import sys
import logging
import os.path as op

from optparse import OptionParser
from pprint import pformat

import numpy as np

import pyhrf
from pyhrf.tools import add_suffix, has_ext
from pyhrf.tools._io import read_volume, write_volume, read_texture, write_texture

logger = logging.getLogger(__name__)

usage = 'usage: %%prog [options] PARCELLATION_FILE PARCEL_ID1 [PARCEL_ID2 ...]'

description = 'Extract a sub parcellation from input parcellation'

parser = OptionParser(usage=usage, description=description)

parser.add_option('-v', '--verbose', dest='verbose', metavar='VERBOSELEVEL',
                  type='int', default=0,
                  help=pformat(pyhrf.verbose_levels))

parser.add_option('-o', '--output', metavar='FILE',
                  default=None, help='Output parcellation file. '
                  'Default is <input_file>_<PARCEL_IDS>.(nii|gii)')

parser.add_option('-c', '--contiguous', action='store_true',
                  default=False, help='Make output parcel ids be contiguous')


(options, args) = parser.parse_args()
# pyhrf.verbose.set_verbosity(options.verbose)
pyhrf.logger.setLevel(options.verbose)

# Treat result of option parsing:
if len(args) < 2:
    parser.print_help()
    sys.exit(1)


pfile = args[0]
p_ids = map(int, args[1:])

if options.output is None:
    options.output = add_suffix(pfile, '_p%s' % '_'.join(args[1:]))

if has_ext(pfile, 'nii') or has_ext(pfile, 'img'):
    data_type = 'volume'
    p, meta = read_volume(pfile)
elif has_ext(pfile, 'gii'):
    data_type = 'surface'
    p, meta = read_texture(pfile)
else:
    raise Exception('Unrecognised extension %s' % op.splitext(pfile)[1])

p = p.squeeze()

new_p = np.zeros_like(p)
all_ids = np.unique(p)
logger.info("Parcel IDs: %s", str(all_ids))
for i, p_id in enumerate(p_ids):
    if p_id not in all_ids:
        print 'Parcel %d not found in input parcellation' % p_id
        sys.exit(1)
    if not options.contiguous or p_id == 0:
        new_p[np.where(p == p_id)] = p_id
    else:
        new_p[np.where(p == p_id)] = i + 1  # 0 is background

if data_type == 'volume':
    write_volume(new_p, options.output, meta)
else:
    write_texture(new_p, options.output)
