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

# <a href="https://colab.research.google.com/github/NeuromatchAcademy/course-content/blob/main/projects/fMRI/load_hcp_retino.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a> &nbsp; <a href="https://kaggle.com/kernels/welcome?src=https://raw.githubusercontent.com/NeuromatchAcademy/course-content/main/projects/fMRI/load_hcp_retino.ipynb" target="_parent"><img src="https://kaggle.com/static/images/open-in-kaggle.svg" alt="Open in Kaggle"/></a>

# # HCP Retinotopy Data Loader
# 
# The HCP 7T retinotopyy dataset comprises fMRI retinotopy from 183 human subjects. The NMA-curated dataset includes the average data over all those subjects. 
# 
# In order to use this dataset, please electronically sign the HCP data use terms at [ConnectomeDB](https://db.humanconnectome.org). Instructions for this are on pp. 24-25 of the [HCP Reference Manual](https://www.humanconnectome.org/storage/app/media/documentation/s1200/HCP_S1200_Release_Reference_Manual.pdf).
# 
# The data and experiment are decribed in detail in [Benson et al.](https://jov.arvojournals.org/article.aspx?articleid=2719988#207329261)

# In[ ]:


# @title Install dependencies
get_ipython().system('pip install nilearn --quiet')


# In[ ]:


import os
import numpy as np
import nibabel as nib
import matplotlib.pyplot as plt
from nilearn import plotting


# In[ ]:


# @title Figure settings
get_ipython().run_line_magic('matplotlib', 'inline')
get_ipython().run_line_magic('config', "InlineBackend.figure_format = 'retina'")
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/course-content/main/nma.mplstyle")


# # Basic parameters

# In[ ]:


# The download cells will store the data in nested directories starting here:
DATA_DIR = "./hcp_retino"
if not os.path.isdir(DATA_DIR):
  os.mkdir(DATA_DIR)

# The data acquisition rate
TR = 1  # Time resolution, in sec

# Time series data are organized by experiment, with each experiment
# having an LR and RL (phase-encode direction) acquistion
RUN_NAMES = [
  "BAR1",   # Sweeping Bars repeat 1
  "BAR2",   # Sweeping Bars repeat 2
  "CCW",    # Counter Clockwise rotating wedge
  "CW",     # Clockwise rotating wedge
  "EXP",    # Expanding ring
  "CON"     # Contracting ring
]


# # Downloading data
# 
# The files that we provide are:
# 
#  - The whole brain time series data in [CIFTI](https://nipy.org/nibabel/reference/nibabel.cifti2.html#module-nibabel.cifti2.cifti2) format
#  - The stimulus image files as numpy arrays (time x height x width x 3)
#  - surface files for visualisation on the brain 

# In[ ]:


# @title Download the data

import os, requests, tarfile

fname = "hcp_retino.tgz"
url = "https://osf.io/d25b4/download"

if not os.path.isfile(fname):
  try:
    r = requests.get(url)
  except requests.ConnectionError:
    print("!!! Failed to download data !!!")
  else:
    if r.status_code != requests.codes.ok:
      print("!!! Failed to download data !!!")
    else:
      print(f"Downloading {fname}...")
      with open(fname, "wb") as fid:
        fid.write(r.content)
      print(f"Download {fname} completed!")


# In[ ]:


# @title Extract the data in `DATA_DIR`
fname_ex = "HCP7T_retino"
path_name = os.path.join(DATA_DIR, fname_ex)
if not os.path.exists(path_name):
  print(f"Extracting {fname}...")
  with tarfile.open(f"{fname}") as fzip:
    fzip.extractall(DATA_DIR)
else:
  print(f"File {fname}.tgz has already been extracted.")


# # Load an example run
# 
# Pick one of the runs explained above and load the timeseries data

# In[ ]:


# pick a run
run = 'EXP'
# define filename
filename = os.path.join(DATA_DIR, fname_ex, f'tfMRI_RET{run}_7T_AP.dtseries.nii')
# load the file with nibabel
cifti = nib.load(filename)

print(cifti.get_fdata().shape)


# From the shape of the data matrix, you should be able to see that there are 300 time points and 91282 brain locations (this is a mixture of cortical surface vertices and subcortical voxels, which the cifti format allows to store together).
# 
# 

# # Simple model fitting
# 
# Now let's run a simple analysis. We know that some of the runs are presenting a period stimulus (EXP, CON, CW, and CCW). We will design a linear model to fit the data with a sinusoidal signal with the same period.
# 
# Let's first define this model.

# In[ ]:


# Simple sinusoidal model
# The task timing is:
# 22 sec, 8 cycles x 32sec, 22 sec
n_before = 22
period = 32
cycles = 8
n_after = 22
t = np.arange(0, period * cycles, 1/TR)
time_var_cos = np.concatenate((np.zeros(n_before),
                               np.cos(2 * np.pi * t / period),
                               np.zeros(n_after)))
time_var_sin = np.concatenate((np.zeros(n_before),
                               np.sin(2 * np.pi * t / period),
                               np.zeros(n_after)))
# Create design matrix
design_matrix = np.asarray([time_var_cos,
                            time_var_sin,
                            np.ones_like(time_var_sin)]).T


# The design matrrix above is made of three columns. One for a cosine wave, one for a sine wave, and one constant columns. 
# 
# The first two columns together can fit a sinusoid of aritrary phase. The last column will help fit the mean of the data.
# 
# This is a linear model of the form $y = M\beta$ which we can invert using $\hat{\beta}=M^+y$ where $M^+$ is the pseudoinverse of $M$.
# 

# In[ ]:


# Invert model
beta = np.linalg.pinv(design_matrix) @ cifti.get_fdata()


# A little bit more maths. Our model is:
# 
# \begin{equation}
# y=\beta_0 cos(2\pi t/T) + \beta_1 sin(2\pi t/T) + \beta_2
# \end{equation}
# 
# To relate $\beta_0$ and $\beta_1$ to the amplitude $A$ and phase $\Phi$ of the data, we can use the trigonometric identity:
# 
# \begin{equation}
# A cos(2\pi t/T + \Phi) = \underbrace{Acos(\Phi)}_{\beta_0} cos(2\pi t/T) + \underbrace{Asin(\Phi)}_{\beta_1} sin(2\pi t/T)
# \end{equation}
# 
# Which means that:
# 
# \begin{equation}
# A = \sqrt{\beta_0^2 + \beta_1^2}
# \end{equation}
# 
# and
# 
# \begin{equation}
# \tan(\Phi) = \frac{\beta_1}{\beta_0}
# \end{equation}

# In[ ]:


angle = np.arctan2(beta[1, :],beta[0, :])
amplitude = np.sqrt(np.sum(beta[:1, :]**2, axis=0))


# Let's look at how this model fits the data. We will look at the brain location with largest amplitude from thee model fit.

# In[ ]:


idx = np.argmax(amplitude)
data = cifti.get_fdata()[:, idx]
# Model prediction is M*beta
pred = design_matrix @ beta[:, idx]

plt.figure()
plt.plot(data, label='data')
plt.plot(pred, label='model')
plt.legend()
plt.show()


# It's an okay fit. Clearly the shape of the data is not correctly captured, but at least we can get the phase information (although this phase is probably wrong due to the fact that the BOLD signal lags behind the neural signal by about 6-10 seconds, which is why we need Haemodynamic response modelling. That it something you can try to do as part of your project.
# 
# Now let's visualise the phase on a brain surface.
# 
# 
# 

# # Visualisation

# In[ ]:


# @title Helper function
# Helper function
def surf_data_from_cifti(data, cifti, surf_name):
  """Maps data from cifti file onto a surface

  Args:
      data : 2D array (nvertices x p)
      cifti : Cifti2Image object
      surf_name : str (either 'LEFT' or 'RIGHT')
  """
  surf_name_long = 'CIFTI_STRUCTURE_CORTEX_'+surf_name
  axis = cifti.header.get_axis(1)  # The brain axis
  for name, data_indices, model in axis.iter_structures():
    if name == surf_name_long:
      vtx_indices = model.vertex
      surf_data = np.zeros(axis.nvertices[surf_name_long])
      surf_data[vtx_indices] = data[data_indices]
      return surf_data
  raise ValueError(f"No structure named {surf_name_long}")


# In[ ]:


# This uses the nilearn package

# Plot the angle
data = angle
data[amplitude < 20] = 0

# If you want you can change the surface. Have a look in hcp_retino/surf for more options
filename = "Q1-Q6_R440.L.inflated.32k_fs_LR.surf.gii"
pathf = os.path.join(DATA_DIR, fname_ex, "surf", filename)
plotting.view_surf(surf_mesh=pathf,
                   surf_map=surf_data_from_cifti(data,cifti,'LEFT'),
                   cmap='hsv', threshold=.0001)


# That's it for now. Note that the stimulus files are stored as Numpy arrays in `hcp_retino/stims`. They can be used to do [Population Receptive Field modelling](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3073038/).