import os
import os.path as op

import numpy as np
import matplotlib.pyplot as plt


def set_up_paths(onedrive_dir=None):
    """Prepare OneDrive BIDS anatomy paths.

    Parameters
    ----------
    onedrive_dir : str | None
        Path to OneDrive. If ``None``, will be detected automatically using
        ``sarna`` package.

    Returns
    -------
    paths : dict
        Dictionary of paths.
    """
    if onedrive_dir is None:
        try:
            import sarna
            onedrive_dir = sarna.proj.find_onedrive()
        except ImportError:
            raise RuntimeError('You either need to provide onedrive_dir, or'
                               'have sarna package to auto-locate OneDrive.')

    anat_dir = op.join(onedrive_dir, 'RESEARCH', 'anat')
    subjects_dir = op.join(anat_dir, 'derivatives', 'freesurfer')

    paths = {'subjects_dir': subjects_dir, 'anat_dir': anat_dir,
             'onedrive_dir': onedrive_dir}
    return paths


def plot_overlay(image, compare, title='', thresh=None):
    """Overlay plot for comparing MRI and CT scans."""
    import nibabel as nib

    image = nib.orientations.apply_orientation(
        np.asarray(image.dataobj), nib.orientations.axcodes2ornt(
            nib.orientations.aff2axcodes(image.affine))).astype(np.float32)
    compare = nib.orientations.apply_orientation(
        np.asarray(compare.dataobj), nib.orientations.axcodes2ornt(
            nib.orientations.aff2axcodes(compare.affine))).astype(np.float32)

    if thresh is not None:
        compare[compare < np.quantile(compare, thresh)] = np.nan

    fig, axes = plt.subplots(1, 3, figsize=(12, 4))

    if isinstance(title, str) and len(title) > 0:
        fig.suptitle(title)

    for i, ax in enumerate(axes):
        ax.imshow(np.take(image, [image.shape[i] // 2], axis=i).squeeze().T,
                  cmap='gray')
        ax.imshow(np.take(compare, [compare.shape[i] // 2],
                          axis=i).squeeze().T, cmap='gist_heat', alpha=0.5)
        ax.invert_yaxis()
        ax.axis('off')

    fig.tight_layout()
    return fig


def find_scans(subject, paths):
    """Find MRI and CT scans for given subjects in Labianca OneDrive BIDS
    structure.

    Parameters
    ----------
    subjects : str
        Subject name.
    paths : dict
        Dictionary of paths obtained via ``pylabianca.localize.set_up_paths``.

    Returns
    -------
    ct_dir, ct_file, mri_dir, mri_file
    """

    # we also need the path to the subject's CT scan:
    ct_dir = op.join(paths['anat_dir'], subject)
    ct_files = os.listdir(ct_dir)

    # filter files
    ct_files = [f for f in ct_files if f.endswith('.nii')
                or f.endswith('.nii.gz')]

    print('Found the following anatomy files:')
    print(ct_files)

    ct_files = [f for f in ct_files if '_ct' in f and 'preop' not in f]
    if len(ct_files) > 1:
        fname_lens = [len(f) for f in ct_files]
        longest = np.argmax(fname_lens)
        ct_fname = ct_files[longest]
    else:
        ct_fname = ct_files[0]
    print(f'Choosing {ct_fname} as the correct CT')
    ct_file = op.join(ct_dir, ct_fname)

    mri_dir = op.join(paths['subjects_dir'], subject, 'mri')
    mri_file = [f for f in os.listdir(mri_dir) if 'T1' in f][0]

    return ct_dir, ct_file, mri_dir, mri_file


def read_compute_ct_alignment(subject, paths, CT_orig, T1, plot=False):
    """Read CT -> MRI alignment matrix or compute it if it is not available.

    If the CT -> MRI alignment affine matrix is not available, but a
    freeview pre-alignment (generated by hand) is present in the relevant
    coreg subject subdirectory, then the prealignment is applied first
    and then the CT -> MRI affine is further optimized automatically.

    Parameters
    ----------
    subject : str
        Subject identifier.
    paths : dict
        Dictionary of paths obtained via ``pylabianca.localize.set_up_paths``.
    CT_orig : nibabel object
        Read CT scan in the original space (not resampled to MRI) to
        align to the MRI.
    T1 : nibabel object
        MRI T1 scan to align the CT to.
    plot : bool
        If True, plot CT-MRI alignment before and after applying the affine
        matrix.

    Return
    ------
    reg_affine : numpy.array
        The CT - MRI alignment affine matrix.
    """
    import mne
    import h5io
    import nibabel as nib

    # if pre-alignment was needed
    coreg_dir = op.join(paths['anat_dir'], 'derivatives', 'coreg', subject)
    ct_coreg_fname = op.join(coreg_dir, f'{subject}_ct_trans.h5')

    this_fname = f'{subject}_ct_aligned_manual.lta'
    pre_alg_fname = op.join(coreg_dir, this_fname)

    if op.exists(pre_alg_fname) and not op.exists(ct_coreg_fname):
        # read freeview pre-alignment transform
        manual_reg_affine_vox = mne.read_lta(pre_alg_fname)

        # convert from vox->vox to ras->ras
        manual_reg_affine = \
            CT_orig.affine @ np.linalg.inv(manual_reg_affine_vox) \
            @ np.linalg.inv(CT_orig.affine)

        with np.printoptions(suppress=True, precision=3):
            print('manual tranformation matrix:\n', manual_reg_affine)

        CT_aligned = mne.transforms.apply_volume_registration(
            CT_orig, T1, manual_reg_affine, cval='1%')


        # optimize further starting from pre-alignment
        reg_affine, _ = mne.transforms.compute_volume_registration(
            CT_orig, T1, pipeline=['rigid'],
            starting_affine=manual_reg_affine)

        # show transformation matrix after further optimization
        with np.printoptions(suppress=True, precision=3):
            print('tranformation matrix after further optimization:\n', reg_affine)

        # transform CT
        CT_aligned = mne.transforms.apply_volume_registration(
            CT_orig, T1, reg_affine, cval='1%')

        if plot:
            plot_overlay(T1, CT_aligned, 'Aligned CT - T1 preoptim',
                         thresh=0.95)

        # save
        h5io.write_hdf5(ct_coreg_fname, reg_affine, overwrite=True)

    # else/then - read the alignment or compute
    if op.exists(ct_coreg_fname):
        print('reading existing CT-MRI transform')
        reg_affine = h5io.read_hdf5(ct_coreg_fname)

    else:
        print('computing CT-MRI transform')
        reg_affine, _ = mne.transforms.compute_volume_registration(
            CT_orig, T1, pipeline='rigids')
        h5io.write_hdf5(ct_coreg_fname, reg_affine, overwrite=True)

    with np.printoptions(suppress=True, precision=3):
        print(reg_affine)

    return reg_affine


def read_channel_table(subject, paths):
    """Read table containing channel ranges and initial anatomy labels.

    Parameters
    ----------
    subject : str
        Subject identifier, for example ``'sub-U10'``.
    paths : dict
        Dictionary of paths obtained with ``pylabianca.localize.set_up_paths``.

    Returns
    -------
    chan_info : pandas.DataFrame
        DataFrame containing channel ranges and labels.
    """
    import pandas as pd

    ch_info_dir = op.join(paths['onedrive_dir'], 'RESEARCH', 'additional info',
                          'channels description')
    chan_info = pd.read_excel(
        op.join(ch_info_dir, 'All_channels_unified.xlsx'),
        sheet_name=subject.replace('sub-', ''))
    return chan_info


def read_anatomical_labels(subject, onedrive_dir):
    """Read anatomical labels assigned to channels based on DKT atlas.

    Parameters
    ----------
    subject : str
        Subject identifier.
    onedrive_dir : str
        Path to OneDrive.

    Return
    ------
    anat: pd.DataFrame
        Table with anatomical labels.
    """
    import pandas as pd

    anat_dir = op.join(onedrive_dir, 'RESEARCH', 'anat', 'derivatives',
                       'labels')
    fname = f'{subject}_labels-DKT.tsv'
    anat_table = pd.read_csv(op.join(anat_dir, fname), sep='\t')
    return anat_table


def is_known(val_or_str):
    if isinstance(val_or_str, str):
        return not val_or_str == '?'
    else:
        return True


def construct_info_from_channel_table(chan_info):
    """Construct mne-python Info object from channel table."""
    import mne
    ch_names = list()

    for row_ix in chan_info.index:
        current_row = chan_info.loc[row_ix]

        if (isinstance(current_row.electrode, str)
            and 'macro' in current_row.electrode):

            prefix = current_row.electrode.split('-')[0]
            known_range = (
                is_known(current_row['channel start'])
                and is_known(current_row['channel end'])
            )

            if not known_range:
                n_channels = 8 if prefix == 'BF' else 10
            else:
                n_channels = int(current_row['channel end']
                                 - current_row['channel start']) + 1

            area = current_row.area
            ch_name = f'{prefix}_{area}_'

            if prefix == 'BF':
                ch_names.append(ch_name + 'micro')

            for ch_idx in range(n_channels):
                ch_names.append(ch_name + f'{ch_idx + 1:01d}')

    info = mne.create_info(ch_names, sfreq=32_000, ch_types='seeg')
    return info


# TODO: separate function for reading channel labels from
#       positions unified
def read_create_channel_positions(subject, paths, ending='_micro'):
    import mne

    fname_base = f'{subject}_channel_positions_ieeg' + ending
    coreg_dir = op.join(paths['anat_dir'], 'derivatives', 'coreg', subject)


    info_fname = op.join(coreg_dir, fname_base + '.fif')
    info_fname2 = op.join(coreg_dir, fname_base + '_Karolina.fif')

    if op.exists(info_fname):
        print('Reading channel positions created by Mikołaj.')
        raw = mne.io.read_raw(info_fname)
        info = raw.info
    elif op.exists(info_fname2):
        print('Reading channel positions created by Karolina.')
        raw = mne.io.read_raw(info_fname2)
        info = raw.info
    else:
        # no saved positions found, creating new
        chan_info = read_channel_table(subject, paths)
        info = construct_info_from_channel_table(chan_info)
        print('No channel positions for this subject, you have to do '
              'some clicking!')

    return info


# select only specific regions micro channels from info
# regions -> str or list
def find_channels(info, regions=None, micro=True, side='both'):
    """Select relevant channel names and indices from mne Info object."""
    assert side.lower() in ['both', 'l', 'r']
    one_side = not side == 'both'
    regions = [regions] if isinstance(regions, str) else regions

    names = list()
    indices = list()
    for ch_idx, ch_name in enumerate(info.ch_names):
        ch_type, region_name, contact_type = ch_name.split('_')

        if regions is not None:
            is_region = any([region in region_name for region in regions])
        else:
            is_region = True

        if is_region and one_side:
            is_region = side.upper() in region_name
        if is_region and micro:
            is_region = 'micro' in contact_type

        if is_region:
            names.append(ch_name)
            indices.append(ch_idx)

    indices = np.array(indices)
    names = np.array(names)

    return names, indices


# TODO: make sure if mne.pick_info operates inplace or not
def pick_info(info, regions=None, micro=True, side='both'):
    """Subselect channels from mne Info object."""
    import mne

    _, idx = find_channels(info, regions=regions, micro=micro, side=side)
    info_sel = mne.pick_info(info.copy(), sel=idx)
    return info_sel


def project_channel_positions_to_voxels(T1, info, stat_map=True,
                                        stat_map_sd=3., channels='micro'):
    """Project channel positions from freesurfer to voxel coordinates."""
    import mne
    import borsar
    from scipy.stats import norm

    T_orig = T1.header.get_vox2ras_tkr()

    ch_pos = borsar.channels.get_ch_pos(info)
    ch_names = np.asarray(info.ch_names)

    # take only non-nan positions
    good_pos = ~ (np.isnan(ch_pos).any(axis=1))
    ch_pos = ch_pos[good_pos]
    ch_names = ch_names[good_pos]

    if channels == 'micro':
        # take only micro channels
        micro_ch_mask = np.array(['micro' in ch for ch in ch_names])
        ch_names = ch_names[micro_ch_mask]
        ch_pos = ch_pos[micro_ch_mask]

    # find channel positions in voxels
    ch_pos_mm = ch_pos * 1000  # meters → millimeters
    pos_vox = mne.transforms.apply_trans(np.linalg.inv(T_orig), ch_pos_mm)

    if stat_map:
        coords = np.indices(T1.shape)
        dist = np.sqrt(np.sum(
            (coords[None, :] - pos_vox[:, :, None, None, None]) ** 2,
            axis=1)
        )

        dist_gauss = norm.pdf(dist, loc=0, scale=stat_map_sd)
        dist_gauss /= dist_gauss.max(axis=(1, 2, 3), keepdims=True)
        all_distgauss = dist_gauss.sum(axis=0)

        return all_distgauss
    else:
        return ch_names, pos_vox


def autolabel_channels(montage, subject, paths):
    '''Auto-label channel positions with DKT atlas labels.
    The passed montage has to be in MRI space.'''

    # in: montage, subject, subjects_dir=None, aseg="aparc+aseg", dist=2
    from collections import OrderedDict
    from tqdm import tqdm

    from mne.channels import DigMontage
    from mne._freesurfer import read_freesurfer_lut, _get_aseg
    from mne.utils import _validate_type
    from mne.transforms import apply_trans
    from mne.surface import _voxel_neighbors, _VOXELS_MAX

    subjects_dir = paths['subjects_dir']
    _validate_type(montage, DigMontage, "montage")
    distances = np.arange(0.25, 5.5, step=0.25)

    aseg = 'aparc.DKTatlas+aseg'
    aseg, aseg_data = _get_aseg(aseg, subject, subjects_dir)

    # read freesurfer lookup table
    lut, fs_colors = read_freesurfer_lut()
    label_lut = {v: k for k, v in lut.items()}

    # assert that all the values in the aseg are in the labels
    assert all([idx in label_lut for idx in np.unique(aseg_data)])

    # get transform to surface RAS for distance units instead of voxels
    vox2ras_tkr = aseg.header.get_vox2ras_tkr()

    ch_dict = montage.get_positions()
    if ch_dict["coord_frame"] != "mri":
        raise RuntimeError(
            "Coordinate frame not supported, expected "
            '"mri", got ' + str(ch_dict["coord_frame"])
        )
    ch_coords = np.array(list(ch_dict["ch_pos"].values()))

    # convert to freesurfer voxel space
    ch_coords = apply_trans(
        np.linalg.inv(aseg.header.get_vox2ras_tkr()), ch_coords * 1000
    )

    # try various distances
    labels_dist = OrderedDict()
    found_labels = OrderedDict()

    for dist in tqdm(distances):
        for ch_name, ch_coord in zip(montage.ch_names, ch_coords):
            if ch_name not in found_labels:
                found_labels[ch_name] = list()
            if ch_name not in labels_dist:
                labels_dist[ch_name] = list()

            else:
                voxels = _voxel_neighbors(
                    ch_coord,
                    aseg_data,
                    dist=dist,
                    vox2ras_tkr=vox2ras_tkr,
                    voxels_max=_VOXELS_MAX,
                )
                label_idxs = set([aseg_data[tuple(voxel)].astype(int)
                                  for voxel in voxels])
                current_labels = [label_lut[idx] for idx in label_idxs]

                for lab in current_labels:
                    if lab not in found_labels[ch_name]:
                        found_labels[ch_name].append(lab)
                        labels_dist[ch_name].append((lab, dist))

    # now rename from DKT atlas names, to human-readable
    new_names = dict()
    for ch_name in labels_dist.keys():
        new_names[ch_name] = list()
        for (region, distance) in labels_dist[ch_name]:
            region = rename_region(region)
            if region is not None:
                new_names[ch_name].append((region, distance))

    return new_names


def parse_part(name):
    '''Parse parts from DKT atlas names.'''
    parts = ['superior', 'middle', 'inferior',
             'anterior', 'posterior', 'caudal', 'rostral',
             'lateral', 'medial', 'isthmus', 'transverse']

    for part in parts:
        if name.startswith(part):
            part_len = len(part)
            return part.capitalize(), name[part_len:]
    return name, None


def iterative_parsing(name):
    rest = name
    new_name = list()
    while rest is not None:
        part, rest = parse_part(rest)
        new_name.append(part.capitalize())
    return '-'.join(new_name)


def rename_region(region, ignore=None):
    '''Rename DKT region name to something more human-readable.'''
    ignore = ['White-Matter', 'Unknown', 'Vent'] if ignore is None else ignore
    translate_hemi = {'lh': 'Left', 'rh': 'Right'}

    if isinstance(region, np.str_):
        region = str(region)
    for ign in ignore:
        if ign in region:
            return None
    if region.startswith('ctx'):
        region_parts = region.split('-')
        hemi = translate_hemi[region_parts[1]]
        rest = region_parts[2]
        name = hemi + '-' + iterative_parsing(rest)
        return name
    else:
        return region


def construct_table_from_anatomical_labels(new_names):
    import pandas as pd

    df = pd.DataFrame(columns=['label', 'closest_anat', 'closest_distance',
                            'second_closest_anat', 'second_closest_distance'])

    # construct a table
    for ch_name, labels in new_names.items():
        _, ch_label, _ = ch_name.split('_')

        idx = df.shape[0]
        df.loc[idx, 'label'] = ch_label

        for label_idx, (anat_label, dist) in enumerate(labels):
            if label_idx > 1:
                break
            prefix = 'closest_' if label_idx == 0 else 'second_closest_'
            df.loc[idx, prefix + 'anat'] = anat_label
            df.loc[idx, prefix + 'distance'] = dist

    return df


def simplify_DKT_name(DKT_name, retain_side=False, translate=None):
    """Translate human-readable DKT atlas name to short labels.

    Translate the human-readable KDT atlas labels obtained with
    ``pylabianca.localize.rename_region`` to shorter anatomical labels
    we used before, like HIP, AMG or SMA."""
    assert retain_side == False  # retaining side currently not supported

    transl = {'Hippocampus': 'HIP', 'Caudal-Anterior-Cingulate': 'cACC',
              'Rostral-Anterior-Cingulate': 'rACC', 'Amygdala': 'AMG',
              'Posterior-Cingulate': 'PCC', 'Medial-Orbitofrontal': 'OFC',
              'Lateral-Orbitofrontal': 'OFC', 'Isthmus-Cingulate': 'PCC',
              'Superior-Frontal': 'SMA', 'Precuneus': 'RSC'}
    if translate is not None:
        transl = transl.update(translate)

    for tr_from, tr_to in transl.items():
        if tr_from in DKT_name:
            return tr_to

    # no simplifications found, retain original name
    return DKT_name


def create_info_from_pos(pos, ch_names=None, sfreq=None):
    import mne

    assert pos.ndim == 2
    assert pos.shape[1] == 3
    n_pos = pos.shape[0]
    if ch_names is None:
        ch_names = ['SEEG_{:03d}'.format(idx) for idx in range(n_pos)]

    sfreq = 32_000 if sfreq is None else sfreq
    dev_head_trans = mne.Transform(fro=1, to=4, trans=np.identity(4))
    info = mne.create_info(ch_names, sfreq=sfreq, ch_types='seeg')
    info['dev_head_t'] = dev_head_trans

    # apply positions
    # CHECK, CONSIDER: maybe constructing and applying montage is better?
    for ch_idx in range(n_pos):
        info['chs'][ch_idx]['loc'][:3] = pos[ch_idx, :]

    return info
