# -*- coding: utf-8 -*-

from copy import deepcopy
import csv
import gzip
import json
import logging
import os
import tempfile

import nibabel as nib
from nibabel.streamlines.array_sequence import ArraySequence
import numpy as np
try:
    import dipy
    dipy_available = True
except ImportError:
    dipy_available = False

from trx.io import get_trx_tmp_dir, load, load_sft_with_reference, save
from trx.streamlines_ops import perform_streamlines_operation, intersection
import trx.trx_file_memmap as tmm
from trx.viz import display
from trx.utils import (flip_sft, is_header_compatible,
                       get_axis_shift_vector,
                       get_reference_info_wrapper,
                       get_reverse_enum,
                       load_matrix_in_any_format,
                       split_name_with_gz)


def convert_dsi_studio(in_dsi_tractogram, in_dsi_fa, out_tractogram,
                       remove_invalid=True, keep_invalid=False):
    if not dipy_available:
        logging.error('Dipy library is missing, scripts are not available.')
        return None
    from dipy.io.stateful_tractogram import StatefulTractogram, Space
    from dipy.io.streamline import save_tractogram, load_tractogram

    in_ext = split_name_with_gz(in_dsi_tractogram)[1]
    out_ext = split_name_with_gz(out_tractogram)[1]

    if in_ext == '.trk.gz':
        with gzip.open(in_dsi_tractogram, 'rb') as f_in:
            with open('tmp.trk', 'wb') as f_out:
                f_out.writelines(f_in)
                sft = load_tractogram('tmp.trk', 'same',
                                      bbox_valid_check=False)
            os.remove('tmp.trk')
    elif in_ext == '.trk':
        sft = load_tractogram(in_dsi_tractogram, 'same',
                              bbox_valid_check=False)
    else:
        raise IOError('{} is not currently supported.'.format(in_ext))

    sft.to_vox()
    sft_fix = StatefulTractogram(sft.streamlines, in_dsi_fa, Space.VOXMM,
                                 data_per_point=sft.data_per_point,
                                 data_per_streamline=sft.data_per_streamline)
    sft_fix.to_vox()
    flip_axis = ['x', 'y']
    sft_fix.streamlines._data -= get_axis_shift_vector(flip_axis)
    sft_flip = flip_sft(sft_fix, flip_axis)

    sft_flip.to_rasmm()
    sft_flip.streamlines._data -= [0.5, 0.5, -0.5]

    if remove_invalid:
        sft_flip.remove_invalid_streamlines()

    if out_ext != '.trx':
        save_tractogram(sft_flip, out_tractogram,
                        bbox_valid_check=not keep_invalid)
    else:
        trx = tmm.TrxFile.from_sft(sft_flip)
        tmm.save(trx, out_tractogram)


def convert_tractogram(in_tractogram, out_tractogram, reference,
                       pos_dtype='float32', offsets_dtype='uint32'):
    if not dipy_available:
        logging.error('Dipy library is missing, scripts are not available.')
        return None
    from dipy.io.streamline import save_tractogram

    in_ext = split_name_with_gz(in_tractogram)[1]
    out_ext = split_name_with_gz(out_tractogram)[1]

    if in_ext == out_ext:
        raise IOError('Input and output cannot be of the same file format.')

    if in_ext != '.trx':
        sft = load_sft_with_reference(in_tractogram, reference,
                                      bbox_check=False)
    else:
        trx = tmm.load(in_tractogram)
        sft = trx.to_sft()
        trx.close()

    if out_ext != '.trx':
        if out_ext == '.vtk':
            if sft.streamlines._data.dtype.name != pos_dtype:
                sft.streamlines._data = sft.streamlines._data.astype(pos_dtype)
            if offsets_dtype == 'uint64' or offsets_dtype == 'uint32':
                offsets_dtype = offsets_dtype[1:]
            if sft.streamlines._offsets.dtype.name != offsets_dtype:
                sft.streamlines._offsets = sft.streamlines._offsets.astype(
                    offsets_dtype)
        save_tractogram(sft, out_tractogram, bbox_valid_check=False)
    else:
        trx = tmm.TrxFile.from_sft(sft)
        if trx.streamlines._data.dtype.name != pos_dtype:
            trx.streamlines._data = trx.streamlines._data.astype(pos_dtype)
        if trx.streamlines._offsets.dtype.name != offsets_dtype:
            trx.streamlines._offsets = trx.streamlines._offsets.astype(
                offsets_dtype)
        tmm.save(trx, out_tractogram)
        trx.close()


def tractogram_simple_compare(in_tractograms, reference):
    if not dipy_available:
        logging.error('Dipy library is missing, scripts are not available.')
        return
    from dipy.io.stateful_tractogram import StatefulTractogram

    tractogram_obj = load(in_tractograms[0], reference)
    if not isinstance(tractogram_obj, StatefulTractogram):
        sft_1 = tractogram_obj.to_sft()
        tractogram_obj.close()
    else:
        sft_1 = tractogram_obj

    tractogram_obj = load(in_tractograms[1], reference)
    if not isinstance(tractogram_obj, StatefulTractogram):
        sft_2 = tractogram_obj.to_sft()
        tractogram_obj.close()
    else:
        sft_2 = tractogram_obj

    if np.allclose(sft_1.streamlines._data, sft_2.streamlines._data,
                   atol=0.001):
        print('Matching tractograms in rasmm!')
    else:
        print('Average difference in rasmm of {}'.format(np.average(
            sft_1.streamlines._data - sft_2.streamlines._data, axis=0)))

    sft_1.to_voxmm()
    sft_2.to_voxmm()
    if np.allclose(sft_1.streamlines._data, sft_2.streamlines._data,
                   atol=0.001):
        print('Matching tractograms in voxmm!')
    else:
        print('Average difference in voxmm of {}'.format(np.average(
            sft_1.streamlines._data - sft_2.streamlines._data, axis=0)))

    sft_1.to_vox()
    sft_2.to_vox()
    if np.allclose(sft_1.streamlines._data, sft_2.streamlines._data,
                   atol=0.001):
        print('Matching tractograms in vox!')
    else:
        print('Average difference in vox of {}'.format(np.average(
            sft_1.streamlines._data - sft_2.streamlines._data, axis=0)))


def verify_header_compatibility(in_files):
    if not dipy_available:
        logging.error('Dipy library is missing, scripts are not available.')
        return

    all_valid = True
    for filepath in in_files:
        if not os.path.isfile(filepath):
            print('{} does not exist'.format(filepath))
        _, in_extension = split_name_with_gz(filepath)
        if in_extension not in ['.trk', '.nii', '.nii.gz', '.trx']:
            raise IOError('{} does not have a supported extension'.format(
                filepath))
        if not is_header_compatible(in_files[0], filepath):
            print('{} and {} do not have compatible header.'.format(
                in_files[0], filepath))
            all_valid = False
    if all_valid:
        print('All input files have compatible headers.')


def tractogram_visualize_overlap(in_tractogram, reference, remove_invalid=True):
    if not dipy_available:
        logging.error('Dipy library is missing, scripts are not available.')
        return None
    from dipy.io.stateful_tractogram import StatefulTractogram
    from dipy.tracking.streamline import set_number_of_points
    from dipy.tracking.utils import density_map

    tractogram_obj = load(in_tractogram, reference)
    if not isinstance(tractogram_obj, StatefulTractogram):
        sft = tractogram_obj.to_sft()
        tractogram_obj.close()
    else:
        sft = tractogram_obj
    sft.streamlines._data = sft.streamlines._data.astype(float)

    sft.data_per_point = None
    sft.streamlines = set_number_of_points(sft.streamlines, 200)

    if remove_invalid:
        sft.remove_invalid_streamlines()

    # Approach (1)
    density_1 = density_map(sft.streamlines, sft.affine, sft.dimensions)
    img = nib.load(reference)
    display(img.get_fdata(), volume_affine=img.affine,
            streamlines=sft.streamlines,  title='RASMM')

    # Approach (2)
    sft.to_vox()
    density_2 = density_map(sft.streamlines, np.eye(4), sft.dimensions)

    # Small difference due to casting of the affine as float32 or float64
    diff = density_1 - density_2
    print('Total difference of {} voxels with total value of {}'.format(
        np.count_nonzero(diff), np.sum(np.abs(diff))))

    display(img.get_fdata(), streamlines=sft.streamlines, title='VOX')

    # Try VOXMM
    sft.to_voxmm()
    affine = np.eye(4)
    affine[0:3, 0:3] *= sft.voxel_sizes

    display(img.get_fdata(), volume_affine=affine,
            streamlines=sft.streamlines,  title='VOXMM')


def validate_tractogram(in_tractogram, reference, out_tractogram,
                        remove_identical_streamlines=True, precision=1):

    if not dipy_available:
        logging.error('Dipy library is missing, scripts are not available.')
        return None
    from dipy.io.stateful_tractogram import StatefulTractogram

    tractogram_obj = load(in_tractogram, reference)

    if not isinstance(tractogram_obj, StatefulTractogram):
        sft = tractogram_obj.to_sft()
        # tractogram_obj.close()
    else:
        sft = tractogram_obj

    ori_dtype = sft.dtype_dict
    ori_len = len(sft)
    tot_remove = 0

    invalid_coord_ind, _ = sft.remove_invalid_streamlines()
    tot_remove += len(invalid_coord_ind)
    logging.warning('Removed {} streamlines with invalid coordinates.'.format(
        len(invalid_coord_ind)))

    indices = [i for i in range(len(sft)) if len(sft.streamlines[i]) <= 1]
    tot_remove = + len(indices)
    logging.warning('Removed {} invalid streamlines (1 or 0 points).'.format(
        len(indices)))

    for i in np.setdiff1d(range(len(sft)), indices):
        norm = np.linalg.norm(np.diff(sft.streamlines[i],
                                      axis=0), axis=1)

        if (norm < 0.001).any():
            indices.append(i)

    indices_val = np.setdiff1d(range(len(sft)), indices).astype(np.uint32)
    logging.warning('Removed {} invalid streamlines (overlapping points).'.format(
        ori_len - len(indices_val)))
    tot_remove += ori_len - len(indices_val)

    if remove_identical_streamlines:
        _, indices_uniq = perform_streamlines_operation(intersection,
                                                        [sft.streamlines],
                                                        precision=precision)
        indices_final = np.intersect1d(
            indices_val, indices_uniq).astype(np.uint32)
        logging.warning('Removed {} overlapping streamlines.'.format(
            ori_len - len(indices_final) - tot_remove))

        indices_final = np.intersect1d(indices_val, indices_uniq)
    else:
        indices_final = indices_val

    if out_tractogram:
        streamlines = sft.streamlines[indices_final].copy()
        dpp = {}
        for key in sft.data_per_point.keys():
            dpp[key] = sft.data_per_point[key][indices_final].copy()

        dps = {}
        for key in sft.data_per_streamline.keys():
            dps[key] = sft.data_per_streamline[key][indices_final]
        new_sft = StatefulTractogram.from_sft(streamlines, sft,
                                              data_per_point=dpp,
                                              data_per_streamline=dps)
        new_sft.dtype_dict = ori_dtype
        save(new_sft, out_tractogram)


def generate_trx_from_scratch(reference, out_tractogram, positions_csv=False,
                              positions=False, offsets=False,
                              positions_dtype='float32', offsets_dtype='uint64',
                              space_str='rasmm', origin_str='nifti',
                              verify_invalid=True, dpv=[], dps=[],
                              groups=[], dpg=[]):
    with get_trx_tmp_dir() as tmp_dir_name:
        if positions_csv:
            with open(positions_csv, newline='') as f:
                reader = csv.reader(f)
                data = list(reader)
                data = [np.reshape(i, (len(i) // 3, 3)).astype(float)
                        for i in data]
                streamlines = ArraySequence(data)
        else:
            positions = load_matrix_in_any_format(positions)
            offsets = load_matrix_in_any_format(offsets)
            lengths = tmm._compute_lengths(offsets)
            streamlines = ArraySequence()
            streamlines._data = positions
            streamlines._offsets = deepcopy(offsets)
            streamlines._lengths = lengths

        if space_str.lower() != 'rasmm' or origin_str.lower() != 'nifti' or \
                verify_invalid:
            if not dipy_available:
                logging.error('Dipy library is missing, advanced options '
                              'related to spatial transforms and invalid '
                              'streamlines are not available.')
                return
            from dipy.io.stateful_tractogram import StatefulTractogram

            space, origin = get_reverse_enum(space_str, origin_str)
            sft = StatefulTractogram(streamlines, reference, space, origin)
            if verify_invalid:
                rem, _ = sft.remove_invalid_streamlines()
                print('{} streamlines were removed becaused they were '
                      'invalid.'.format(len(rem)))
            sft.to_rasmm()
            sft.to_center()
            streamlines = sft.streamlines
            streamlines._offsets = offsets

        affine, dimensions, _, _ = get_reference_info_wrapper(reference)
        header = {
            "DIMENSIONS": dimensions.tolist(),
            "VOXEL_TO_RASMM": affine.tolist(),
            "NB_VERTICES": len(streamlines._data),
            "NB_STREAMLINES": len(streamlines)-1,
        }

        if header['NB_STREAMLINES'] <= 1:
            raise IOError('To use this script, you need at least 2'
                          'streamlines.')

        with open(os.path.join(tmp_dir_name, "header.json"), "w") as out_json:
            json.dump(header, out_json)

        curr_filename = os.path.join(tmp_dir_name, 'positions.3.{}'.format(
            positions_dtype))
        streamlines._data.astype(positions_dtype).tofile(
            curr_filename)
        curr_filename = os.path.join(tmp_dir_name, 'offsets.{}'.format(
            offsets_dtype))
        streamlines._offsets.astype(offsets_dtype).tofile(
            curr_filename)

        if dpv:
            os.mkdir(os.path.join(tmp_dir_name, 'dpv'))
            for arg in dpv:
                curr_arr = np.squeeze(load_matrix_in_any_format(arg[0]).astype(
                    arg[1]))
                if arg[1] == 'bool':
                    arg[1] = 'bit'
                if curr_arr.ndim > 2:
                    raise IOError('Maximum of 2 dimensions for dpv/dps/dpg.')
                dim = '' if curr_arr.ndim == 1 else '{}.'.format(
                    curr_arr.shape[-1])
                curr_filename = os.path.join(tmp_dir_name, 'dpv', '{}.{}{}'.format(
                    os.path.basename(os.path.splitext(arg[0])[0]), dim, arg[1]))
                curr_arr.tofile(curr_filename)

        if dps:
            os.mkdir(os.path.join(tmp_dir_name, 'dps'))
            for arg in dps:
                curr_arr = np.squeeze(load_matrix_in_any_format(arg[0]).astype(
                    arg[1]))
                if arg[1] == 'bool':
                    arg[1] = 'bit'
                if curr_arr.ndim > 2:
                    raise IOError('Maximum of 2 dimensions for dpv/dps/dpg.')
                dim = '' if curr_arr.ndim == 1 else '{}.'.format(
                    curr_arr.shape[-1])
                curr_filename = os.path.join(tmp_dir_name, 'dps', '{}.{}{}'.format(
                    os.path.basename(os.path.splitext(arg[0])[0]), dim, arg[1]))
                curr_arr.tofile(curr_filename)
        if groups:
            os.mkdir(os.path.join(tmp_dir_name, 'groups'))
            for arg in groups:
                curr_arr = load_matrix_in_any_format(arg[0]).astype(arg[1])
                if arg[1] == 'bool':
                    arg[1] = 'bit'
                if curr_arr.ndim > 2:
                    raise IOError('Maximum of 2 dimensions for dpv/dps/dpg.')
                dim = '' if curr_arr.ndim == 1 else '{}.'.format(
                    curr_arr.shape[-1])
                curr_filename = os.path.join(tmp_dir_name, 'groups', '{}.{}{}'.format(
                    os.path.basename(os.path.splitext(arg[0])[0]), dim, arg[1]))
                curr_arr.tofile(curr_filename)

        if dpg:
            os.mkdir(os.path.join(tmp_dir_name, 'dpg'))
            for arg in dpg:
                if not os.path.isdir(os.path.join(tmp_dir_name, 'dpg', arg[0])):
                    os.mkdir(os.path.join(tmp_dir_name, 'dpg', arg[0]))
                curr_arr = load_matrix_in_any_format(arg[1]).astype(arg[2])
                if arg[1] == 'bool':
                    arg[1] = 'bit'
                if curr_arr.ndim > 2:
                    raise IOError('Maximum of 2 dimensions for dpv/dps/dpg.')
                if curr_arr.shape == (1, 1):
                    curr_arr = curr_arr.reshape((1,))
                dim = '' if curr_arr.ndim == 1 else '{}.'.format(
                    curr_arr.shape[-1])
                curr_filename = os.path.join(tmp_dir_name, 'dpg', arg[0], '{}.{}{}'.format(
                    os.path.basename(os.path.splitext(arg[1])[0]), dim, arg[2]))
                curr_arr.tofile(curr_filename)

        trx = tmm.load(tmp_dir_name)
        tmm.save(trx, out_tractogram)
        trx.close()


def manipulate_trx_datatype(in_filename, out_filename, dict_dtype):
    trx = tmm.load(in_filename)

    # For each key in dict_dtype, we create a new memmap with the new dtype
    # and we copy the data from the old memmap to the new one.
    for key in dict_dtype:
        if key == 'positions':
            tmp_mm = np.memmap(tempfile.NamedTemporaryFile(),
                               dtype=dict_dtype[key],
                               mode='w+',
                               shape=trx.streamlines._data.shape)
            tmp_mm[:] = trx.streamlines._data[:]
            trx.streamlines._data = tmp_mm
        elif key == 'offsets':
            tmp_mm = np.memmap(tempfile.NamedTemporaryFile(),
                               dtype=dict_dtype[key],
                               mode='w+',
                               shape=trx.streamlines._offsets.shape)
            tmp_mm[:] = trx.streamlines._offsets[:]
            trx.streamlines._offsets = tmp_mm
        elif key == 'dpv':
            for key_dpv in dict_dtype[key]:
                tmp_mm = np.memmap(tempfile.NamedTemporaryFile(),
                                   dtype=dict_dtype[key][key_dpv],
                                   mode='w+',
                                   shape=trx.data_per_vertex[key_dpv]._data.shape)
                tmp_mm[:] = trx.data_per_vertex[key_dpv]._data[:]
                trx.data_per_vertex[key_dpv]._data = tmp_mm
        elif key == 'dps':
            for key_dps in dict_dtype[key]:
                tmp_mm = np.memmap(tempfile.NamedTemporaryFile(),
                                   dtype=dict_dtype[key][key_dps],
                                   mode='w+',
                                   shape=trx.data_per_streamline[key_dps].shape)
                tmp_mm[:] = trx.data_per_streamline[key_dps][:]
                trx.data_per_streamline[key_dps] = tmp_mm
        elif key == 'dpg':
            for key_group in dict_dtype[key]:
                for key_dpg in dict_dtype[key][key_group]:
                    tmp_mm = np.memmap(tempfile.NamedTemporaryFile(),
                                       dtype=dict_dtype[key][key_group][key_dpg],
                                       mode='w+',
                                       shape=trx.data_per_group[key_group][key_dpg].shape)
                    tmp_mm[:] = trx.data_per_group[key_group][key_dpg][:]
                    trx.data_per_group[key_group][key_dpg] = tmp_mm
        elif key == 'groups':
            for key_group in dict_dtype[key]:
                tmp_mm = np.memmap(tempfile.NamedTemporaryFile(),
                                   dtype=dict_dtype[key][key_group],
                                   mode='w+',
                                   shape=trx.groups[key_group].shape)
                tmp_mm[:] = trx.groups[key_group][:]
                trx.groups[key_group] = tmp_mm

    tmm.save(trx, out_filename)
    trx.close()