#!/usr/bin/env python
# coding: utf-8

# # White matter challenge prediction
# 
# This notebook generates a prediction for the [ISBI2015 White Matter Challenge](http://cmic.cs.ucl.ac.uk/wmmchallenge/)

# Import all the necessary components from other libraries, including `numpy`, a module from `scikit-learn` and functions from `dipy` for managing gradient tables:

# In[1]:


import numpy as np
import sklearn.linear_model as lm
from dipy.core import gradients as grad


# The following imports are from our current implementation of the RS-SFM model

# In[2]:


import utils
from model import Model, BiExponentialIsotropicModel


# A utility function reads the data from the files. We read in the 'seen' data

# In[3]:


data = utils.read_data()
seen_bvals = data['seen']['bvals']
seen_bvecs = data['seen']['bvecs']
seen_delta = data['seen']['delta']
seen_Delta = data['seen']['Delta']
seen_te = data['seen']['TE']
seen_g = data['seen']['g']


# And create a gradient table to describe the conditions used to collect the seen data

# In[4]:


seen_gtab = grad.gradient_table(seen_bvals, 
                                seen_bvecs, 
                                big_delta=seen_Delta, 
                                small_delta=seen_delta)


# Similarly, a gradient table is created for the unseen data

# In[5]:


unseen_bvals = data['unseen']['bvals']
unseen_bvecs = data['unseen']['bvecs']
unseen_delta = data['unseen']['delta']
unseen_Delta = data['unseen']['Delta']
unseen_te = data['unseen']['TE']
unseen_g = data['unseen']['g']

unseen_gtab = grad.gradient_table(unseen_bvals, 
                                  unseen_bvecs, 
                                  big_delta=unseen_Delta, 
                                  small_delta=unseen_delta)


# The model object used to fit the data is initialized with the gradient table for the seen data

# In[6]:


model = Model(seen_gtab, 
              isotropic=BiExponentialIsotropicModel, 
              solver=lm.ElasticNet(l1_ratio=0.4, 
                                   alpha=5e-7, 
                                   positive=True, 
                                   warm_start=True, 
                                   fit_intercept=True, 
                                   normalize=True))


# We initialize a list to contain the predictions

# In[7]:


predict = []


# In each iteration of the following loop, the model is fit to the seen data from a different voxel, and a prediction is made for this voxel in the unseen conditions

# In[8]:


for vox_idx in range(6):
    signal = data['seen']['signal'][:, vox_idx]
    fit = model.fit(signal, seen_te, seen_g)
    predict.append(fit.predict(unseen_gtab, unseen_te))


# The prediction is saved in the expected format

# In[9]:


np.savetxt('unseenSignal.txt', np.array(predict).T, fmt='%.4f', header='\%-voxel1---voxel2---voxel3---voxel4---voxel5---voxel6')