#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Apr  7 08:41:07 2020

@author: joshs
"""

import os

import numpy as np
import xarray as xr
import pandas as pd
import matplotlib.pyplot as plt

from allensdk.brain_observatory.ecephys.ecephys_project_cache import EcephysProjectCache

# %%

cache_dir = '/mnt/nvme0/ecephys_cache_dir' 

manifest_path = os.path.join(cache_dir, "manifest.json")
cache = EcephysProjectCache.from_warehouse(manifest=manifest_path)

import warnings
warnings.filterwarnings("ignore")

sessions = cache.get_session_table()

# %%

session = cache.get_session_data(sessions.index.values[0])

# %%

import nrrd

streamlines, header = nrrd.read('/home/joshs/Downloads/laplacian_10.nrrd')

annotations = np.load('/mnt/md0/data/opt/annotation_volume_10um_by_index.npy')

structure_tree = pd.read_csv('/mnt/md0/data/opt/ccf_structure_tree.csv', index_col=0)

# %%

import re

def get_layer_name(acronym):
    
    try:
        layer = int(re.findall(r'\d+', acronym)[0])
        if layer == 3:
            layer = 0
        return layer
    except IndexError:
        return 0
    
def get_structure_ids(df, annotations):
    
    x = (df.anterior_posterior_ccf_coordinate.values / 10).astype('int')
    y = (df.dorsal_ventral_ccf_coordinate.values / 10).astype('int')
    z = (df.left_right_ccf_coordinate.values / 10).astype('int')
    
    x[x < 0] = 0
    y[y < 0] = 0
    z[z < 0] = 0
    
    structure_ids = annotations[x, y, z] - 1 # annotation volume is Matlab-indexed
    
    return structure_ids

# %%

df = session.channels[session.channels.anterior_posterior_ccf_coordinate > 0]

x = (df.anterior_posterior_ccf_coordinate.values / 10).astype('int')
y = (df.dorsal_ventral_ccf_coordinate.values / 10).astype('int')
z = (df.left_right_ccf_coordinate.values / 10).astype('int')

cortical_depth = streamlines[x, y, z]

df['cortical_depth'] = 0

df.loc[df.anterior_posterior_ccf_coordinate > 0, 'cortical_depth'] = cortical_depth


# %%

structure_ids = get_structure_ids(df, annotations)
structure_acronyms = structure_tree.loc[structure_ids].acronym

layers = [get_layer_name(acronym) for acronym in structure_acronyms]
        
df['cortical_layer'] = layers