{{{
#!/bin/env python
''' Run models on ROI data '''

import os

import numpy as N
from scipy import io as sio
import scipy.sandbox.models as SSM
import scipy.interpolate as SI

from neuroimaging.modalities.fmri import protocol
from neuroimaging.modalities.fmri import hrf
from neuroimaging.modalities.fmri import functions

data_file = os.path.join('..', 'data', 'roi_data.mat')
roi_data = sio.loadmat(data_file)['roi_data']

# Specify type of linear model to estimate
model_type = SSM.regression.ar_model
ar_rho = 3
ar_niter = 6

# Specify drift term for all models
drift_window = (0,256)
drift_df = 7
drift = functions.SplineConfound(df=drift_df,
                                 window=drift_window)
drift_term = protocol.ExperimentalQuantitative(
    'spline_drift',
    drift)

analyses = {}
for sb in range(roi_data.shape[0]):
    analyses[sb] = {}
    for ss in range(roi_data.shape[1]):
        data = roi_data[sb, ss]
        onsets = data.ons
        tr = data.TR
        y = data.y
        movements = data.movements
        # Adust onsets to seconds
        onsets *= tr
        # Create time vector
        T = N.arange(len(y)) * tr
        # Make event factor
        flashes = zip(['flash']*len(onsets), onsets)
        flash_factor = protocol.ExperimentalFactor('flash',
                                                   flashes,
                                                   delta=True)
        flash_factor.convolve(hrf.canonical)
        # Add movements by providing interpolator
        moves_interpolator = functions.InterpolatedConfound(T, movements.T)
        movement_term = protocol.ExperimentalQuantitative(
            'movement',
            moves_interpolator)
        # Add factor, moves and drift term to model
        formula = flash_factor + movement_term + drift_term
        # Calculate design from model and time vector
        design = formula(time=T).T
        # Estimate design on time series data
        model = model_type(design, rho=ar_rho)
        result = model.iterative_fit(y, niter=ar_niter)
        analyses[sb][ss] = {'model': model,
                            'result': result}
}}}