""" Code to grab the data from NeuroVault, and compute a map of
frequency of activation in the brain.
"""
# Authors: Chris Filo Gorgolewski, Gael Varoquaux
# License: BSD

import json
import urllib, os, errno
from urllib2 import Request, urlopen, HTTPError

import pandas as pd
from pandas.io.json import json_normalize
import numpy as np
import pylab as plt

from nipype.utils.filemanip import split_filename
import nibabel as nb

from joblib import Memory

from nilearn.image import resample_img
from nilearn.masking import compute_background_mask, _extrapolate_out_mask
from nilearn.plotting.img_plotting import plot_anat

# Use a joblib memory, to avoid depending on an Internet connection
mem = Memory(cachedir=os.path.join(os.getenv('HOME'), 'neurovault_analysis/cache'))


def mkdir_p(path):
    try:
        os.makedirs(path)
    except OSError as exc: # Python >2.5
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else: raise


def get_collections_df():
    """Downloads metadata about collections/papers stored in NeuroVault and
    return it as a pandas DataFrame"""

    request = Request('http://neurovault.org/api/collections/?format=json')
    response = urlopen(request)
    elevations = response.read()
    data = json.loads(elevations)
    collections_df = json_normalize(data)
    collections_df.rename(columns={'id':'collection_id'}, inplace=True)
    collections_df.set_index("collection_id")

    return collections_df


def get_images_df():
    """Downloads metadata about images/statistical maps stored in NeuroVault and 
    return it as a pandas DataFrame"""

    request=Request('http://neurovault.org/api/images/?format=json')
    response = urlopen(request)
    elevations = response.read()
    data = json.loads(elevations)
    images_df = json_normalize(data)
    images_df['collection'] = images_df['collection'].apply(lambda x: int(x.split("/")[-2]))
    images_df['image_id'] = images_df['url'].apply(lambda x: int(x.split("/")[-2]))
    images_df.rename(columns={'collection':'collection_id'}, inplace=True)
    return images_df


def get_images_with_collections_df():
    """Downloads metadata about images/statistical maps stored in NeuroVault and
    and enriches it with metadata of the corresponding collections. The result 
    is returned as a pandas DataFrame"""
    
    collections_df = get_collections_df()
    images_df = get_images_df()

    combined_df = pd.merge(images_df, collections_df, how='left', on='collection_id',
                      suffixes=('_image', '_collection'))
    return combined_df


def download_and_resample(images_df, dest_dir, target):
    """Downloads all stat maps and resamples them to a common space.
    """

    target_nii = nb.load(target)
    orig_path = os.path.join(dest_dir, "original")
    mkdir_p(orig_path)
    resampled_path = os.path.join(dest_dir, "resampled")
    mkdir_p(resampled_path)
    out_df = combined_df.copy()

    for row in combined_df.iterrows():
        # Downloading the file to the "original" subfolder
        _, _, ext = split_filename(row[1]['file'])
        orig_file = os.path.join(orig_path, "%04d%s" % (row[1]['image_id'], ext))
        if not os.path.exists(orig_file):
            print "Downloading %s" % orig_file
            urllib.urlretrieve(row[1]['file'], orig_file)

        try:
            # Compute the background and extrapolate outside of the mask
            print "Extrapolating %s" % orig_file
            niimg = nb.load(orig_file)
            data = niimg.get_data().squeeze()
            niimg = nb.Nifti1Image(data, niimg.affine,
                                header=niimg.get_header())
            bg_mask = compute_background_mask(niimg).get_data()
            # Test if the image has been masked:
            out_of_mask = data[np.logical_not(bg_mask)]
            if np.all(np.isnan(out_of_mask)) or len(np.unique(out_of_mask)) == 1:
                # Need to extrapolate
                data = _extrapolate_out_mask(data.astype(np.float), bg_mask,
                                            iterations=3)[0]
            niimg = nb.Nifti1Image(data, niimg.affine,
                                header=niimg.get_header())
            del out_of_mask, bg_mask
            # Resampling the file to target and saving the output in the "resampled"
            # folder
            resampled_file = os.path.join(resampled_path,
                                        "%06d%s" % (row[1]['image_id'], ext))

            print "Resampling %s" % orig_file
            resampled_nii = resample_img(niimg, target_nii.get_affine(),
                target_nii.shape)
            resampled_nii = nb.Nifti1Image(resampled_nii.get_data().squeeze(),
                                        resampled_nii.get_affine(),
                                        header=niimg.get_header())
            if len(resampled_nii.shape) == 3:
                resampled_nii.to_filename(resampled_file)
            else:
                # We have a 4D file
                assert len(resampled_nii.shape) == 4
                resampled_data = resampled_nii.get_data()
                affine = resampled_nii.affine
                for index in range(resampled_nii.shape[-1]):
                    # First save the files separately
                    this_nii = nb.Nifti1Image(resampled_data[..., index],
                                                affine)
                    this_id = int("%i%i" % (-row[1]['image_id'], index))
                    this_file = os.path.join(resampled_path,
                        "%06d%s" % (this_id, ext))
                    this_nii.to_filename(this_file)

                    # Second, fix the dataframe
                    out_df = out_df[out_df.image_id != row[1]['image_id']]
                    this_row = row[1].copy()
                    this_row.image_id = this_id
                    out_df = out_df.append(this_row)
        except IOError as e:
            # Fix the dataframe
            out_df = out_df[out_df.image_id != row[1]['image_id']]
            print "Could not load %s " % orig_file
            print e

    return out_df


def get_frequency_map(images_df, dest_dir, target):
    """
    """
    mask_img = 'gm_mask.nii.gz'
    mask = nb.load(mask_img).get_data().astype(np.bool)


    target_nii = nb.load(target)
    resampled_path = os.path.join(dest_dir, "resampled")
    freq_map_data = np.zeros(target_nii.shape)

    n_images = 0
    for row in combined_df.iterrows():
        _, _, ext = split_filename(row[1]['file'])
        orig_file = os.path.join(resampled_path,
                                 "%06d%s" % (row[1]['image_id'], ext))
        nb.load(orig_file)
        if not os.path.exists(orig_file):
            urllib.urlretrieve(row[1]['file'], orig_file)

        resampled_nii = resample_img(orig_file, target_nii.get_affine(),
                                     target_nii.shape,
                                     interpolation="nearest")
        data = resampled_nii.get_data().squeeze()
        data[np.isnan(data)] = 0
        data[np.logical_not(mask)] = 0
        data = np.abs(data)

        # Keep only things that are very significant
        data = data > 3
        if len(data.shape) == 4:
            for d in np.rollaxis(data, -1):
                freq_map_data += (d != 0)
                n_images +=1
        else:
            freq_map_data += data
            n_images += 1

    freq_map_data *= 100. / n_images

    return nb.Nifti1Image(freq_map_data, target_nii.get_affine())


def url_get(url):
    request = Request(url)
    response = urlopen(request)
    return response.read()


def get_neurosynth_terms(combined_df):
    """ Grab terms for each image, decoded with neurosynth"""
    terms = list()
    from sklearn.feature_extraction import DictVectorizer
    vectorizer = DictVectorizer()
    image_ids = list()
    for row in combined_df.iterrows():
        image_id = row[1]['image_id']
        image_ids.append(int(image_id))
        print "Fetching terms for image %i" % image_id
        image_url = row[1]['url_image'].split('/')[-2]

        try: