shell:
    'c3d {input.ct} -scale 0 -landmarks-to-spheres {input.txt} 1 -o {output.mask}'

shell:
    'export SINGULARITYENV_FS_LICENSE=$HOME/.freesurfer.txt&&\
    singularity run --cleanenv \
    --bind {params.bids_dir}:/tmp/input \
    --bind {params.out_dir}:/tmp/output \
    --bind {params.license}:/tmp/{params.license_name} \
    {params.fmriprep_img} /tmp/input  /tmp/output participant --skip_bids_validation \
    --participant_label {params.sub} --anat-only \
    --fs-license-file /tmp/{params.license_name} \
    --bids-filter-file {params.bids_filter}'

shell:
    'greedy -d 3 -threads 4 -a -ia-image-centers -m MI -i {input.ref} {input.flo} -o {output.affine_xfm_ras} -n {params.n_iterations_affine}&&'
    'greedy -d 3 -threads 4 -m MI -i {input.ref} {input.flo} -it {output.affine_xfm_ras}  -o {output.xfm_deform} -oinv {output.xfm_deform_inv} -n {params.n_iterations_deform} -s {params.grad_sigma} {params.warp_sigma} &&'
    'greedy -d 3 -threads 4 -rf {input.ref} -rm {input.flo} {output.warped_subj_affine} -r {output.affine_xfm_ras}&&'
    'greedy -d 3 -threads 4 -rf {input.ref} -rm {input.flo} {output.warped_subj_greedy} -r {output.xfm_deform} {output.affine_xfm_ras}'

import numpy as np
import pandas as pd
from scipy.io import loadmat


# filen=r'/home/greydon/Documents/data/SEEG_peds/derivatives/atlasreg/sub-P010/sub-P010_acq-noncontrast_desc-rigid_from-noncontrast_to-contrast_type-ras_xfm.mat'
# filen_out=r'/home/greydon/Documents/data/SEEG_peds/derivatives/atlasreg/sub-P010/sub-P010_acq-noncontrast_desc-rigid_from-noncontrast_to-contrast_type-ras_xfm_1.txt'
# filen=r'/home/greydon/Documents/data/SEEG/derivatives/atlasreg/sub-P097/sub-P097_desc-rigid_from-ct_to-T1w_type-ras_xfm.txt'


transformMatrix = np.loadtxt(snakemake.input.xfm)
lps2ras=np.diag([-1, -1, 1, 1])
ras2lps=np.diag([-1, -1, 1, 1])
transform_lps=np.dot(ras2lps, np.dot(transformMatrix,lps2ras))

Parameters = " ".join([str(x) for x in np.concatenate((transform_lps[0:3,0:3].reshape(9), transform_lps[0:3,3]))])
#output_matrix_txt = filen.split('.txt')[0] + '.tfm'

with open(snakemake.output.tfm, 'w') as fid:
	fid.write("#Insight Transform File V1.0\n")
	fid.write("#Transform 0\n")
	fid.write("Transform: AffineTransform_double_3_3\n")
	fid.write("Parameters: " + Parameters + "\n")
	fid.write("FixedParameters: 0 0 0\n")

import pandas as pd
import numpy as np
import nibabel as nib
import os

#read fcsv electrodes file
df_elec = pd.read_table(snakemake.input.fcsv,sep=',',header=2)
df_elec
df_atlas = pd.read_table(snakemake.input.dseg_tsv)
df_atlas

#load up tissue probability, warped from template
tissue_prob_vol = dict()
tissue_prob_elec = dict()

for label,nii in zip(snakemake.config['tissue_labels'], snakemake.input.tissue_seg):
	print(label)
	print(nii)
	tissue_prob_vol[label] = nib.load(nii).get_fdata()
	tissue_prob_elec[label] = list()

#load dseg nii (as integer)
dseg_nii = nib.load(snakemake.input.dseg_nii)
dseg_vol = dseg_nii.get_fdata().astype('int')

#get affine from image, so we can go from RAS coords to array indices
dseg_affine = dseg_nii.affine
dseg_affine

#get coords from fcsv
coords = df_elec[['x','y','z']].to_numpy()


labelnames = []

for i in range(len(coords)):

	vec = np.hstack([coords[i,:],1])

	#dseg_affine is used to xfm indices to RAS coords, 
	# so we use the inverse to go the other way
	tvec = np.linalg.inv(dseg_affine) @ vec.T   
	inds = np.round(tvec[:3]).astype('int')
	labelnum = dseg_vol[inds[0],inds[1],inds[2]]


	if labelnum >0:
		labelnames.append(df_atlas.loc[df_atlas['label']==labelnum,'name'].to_list()[0])
	else:
		labelnames.append('None')

	for label in snakemake.config['tissue_labels']:
		tissue_prob_elec[label].append(tissue_prob_vol[label][inds[0],inds[1],inds[2]])

#add new columns to existing dataframe
df_elec['atlas_label'] = labelnames
for label in snakemake.config['tissue_labels']:
	df_elec[label] = tissue_prob_elec[label]

#create new dataframe with selected variables and save it
out_df = df_elec[['label','atlas_label'] + snakemake.config['tissue_labels'] + ['x','y','z']]
out_df.to_csv(snakemake.output.tsv,sep='\t',float_format='%.3f',index=False)
out_df.to_excel(os.path.splitext(snakemake.output.tsv)[0]+'.xlsx',float_format='%.3f',index=False)

out_df

import numpy as np
import nibabel as nib
import json

#load up tissue probability, warped from template
tissue_prob_vol = dict()

for label,nii in zip(snakemake.config['tissue_labels'], snakemake.input.tissue_priors):
	tissue_prob_vol[label] = nib.load(nii).get_fdata()


#load up k-class tissue segmentation
tissue_k_seg = nib.load(snakemake.input.seg_channels_4d)
tissue_k_seg.shape

sim_prior_k = np.zeros([len(snakemake.config['tissue_labels']),tissue_k_seg.shape[3]])

#for each prior, need to find the channel that best fits
for i,label in enumerate(snakemake.config['tissue_labels']):
	for k in range(tissue_k_seg.shape[3]):

		print(f'Computing overlap of {label} prior and channel {k}... ')
		#compute intersection over union
		s1 = tissue_prob_vol[label] >0.5
		s2 = tissue_k_seg.slicer[:,:,:,k].get_fdata() >0.5
		sim_prior_k[i,k] = np.sum(np.logical_and(s1,s2).flat) / np.sum(np.logical_or(s1,s2).flat) 

label_to_k_dict = dict()

for i,label in enumerate(snakemake.config['tissue_labels']):
	label_to_k_dict[label] = int(np.argmax(sim_prior_k[i,:]))
	#write nii to file
	print('writing image at channel {} to output file: {}'.format(label_to_k_dict[label], \
													snakemake.output.tissue_segs[i]))
	nib.save(tissue_k_seg.slicer[:,:,:,label_to_k_dict[label]],\
					snakemake.output.tissue_segs[i])


with open(snakemake.output.mapping_json, 'w') as outfile:
	json.dump(label_to_k_dict, outfile,indent=4)

import matplotlib
matplotlib.use('Agg')
import ants
from nilearn import plotting,image
import nibabel as nib
import numpy as np

# snakemake.input.
# snakemake.output.
# html_view.open_in_browser()

template = ants.image_read(ants.get_ants_data('mni'))

ct_img=nib.load(snakemake.input.ct)
if (np.isnan(ct_img.get_fdata())).any():
	ct_img=nib.Nifti1Image(np.nan_to_num(ct_img.get_fdata()), header=ct_img.header, affine=ct_img.affine)
	nib.save(ct_img,snakemake.input.ct)

ct_ants = ants.image_read(snakemake.input.ct)
mask_ants = ants.image_read(snakemake.input.mask)

ct_ants_reg = ants.registration(template, ct_ants, type_of_transform='QuickRigid')
ct_ants_reg_applied=ants.apply_transforms(template, ct_ants, transformlist=ct_ants_reg['fwdtransforms'])
ct_resample = ants.to_nibabel(ct_ants_reg_applied)

mask_ants_reg_applied = ants.apply_transforms(ct_ants, mask_ants, transformlist=ct_ants_reg['fwdtransforms'])
mask_resample = ants.to_nibabel(mask_ants_reg_applied)

mask_params = {
			'symmetric_cmap': True,
			'cut_coords':[0,0,0],
			'dim': 1,
			'cmap':'viridis',
			'opacity':0.7
			}

html_view = plotting.view_img(stat_map_img=mask_resample,bg_img=ct_resample,**mask_params)
html_view.save_as_html(snakemake.output.html)

import os
import nibabel as nb
import numpy as np
import pandas as pd
import regex as re
import matplotlib.pyplot as plt
from nilearn.plotting.displays import PlotlySurfaceFigure
import plotly.graph_objs as go
from mne.transforms import apply_trans


AXIS_CONFIG = {
    "showgrid": False,
    "showline": False,
    "ticks": "",
    "title": "",
    "showticklabels": False,
    "zeroline": False,
    "showspikes": False,
    "spikesides": False,
    "showbackground": False,
}

LAYOUT = {
	"scene": {f"{dim}axis": AXIS_CONFIG for dim in ("x", "y", "z")},
	"paper_bgcolor": "#fff",
	"hovermode": False,
	"showlegend":True,
	"legend":{
		"itemsizing": "constant",
		"groupclick":"togglegroup",
		"yanchor":"top",
		"y":0.8,
		"xanchor":"left",
		"x":0.05,
		"title_font_family":"Times New Roman",
		"font":{
			"size":20
		},
		"bordercolor":"Black",
		"borderwidth":1
	},
	"margin": {"l": 0, "r": 0, "b": 0, "t": 0, "pad": 0},
}

CAMERAS = {
    "left": {
        "eye": {"x": -1.5, "y": 0, "z": 0},
        "up": {"x": 0, "y": 0, "z": 1},
        "center": {"x": 0, "y": 0, "z": 0},
    },
    "right": {
        "eye": {"x": 1.5, "y": 0, "z": 0},
        "up": {"x": 0, "y": 0, "z": 1},
        "center": {"x": 0, "y": 0, "z": 0},
    },
    "dorsal": {
        "eye": {"x": 0, "y": 0, "z": 1.5},
        "up": {"x": 0, "y": 1, "z": 0},
        "center": {"x": 0, "y": 0, "z": 0},
    },
    "ventral": {
        "eye": {"x": 0, "y": 0, "z": -1.5},
        "up": {"x": 0, "y": 1, "z": 0},
        "center": {"x": 0, "y": 0, "z": 0},
    },
    "anterior": {
        "eye": {"x": 0, "y": 1.5, "z": 0},
        "up": {"x": 0, "y": 0, "z": 1},
        "center": {"x": 0, "y": 0, "z": 0},
    },
    "posterior": {
        "eye": {"x": 0, "y": -1.5, "z": 0},
        "up": {"x": 0, "y": 0, "z": 1},
        "center": {"x": 0, "y": 0, "z": 0},
    },
}

lighting_effects = dict(ambient=0.4, diffuse=0.5, roughness = 0.9, specular=0.6, fresnel=0.2)

def determine_groups(iterable, numbered_labels=False):
	values = []
	for item in iterable:
		temp=None
		if re.findall(r"([a-zA-Z]+)([0-9]+)([a-zA-Z]+)", item):
			temp = "".join(list(re.findall(r"([a-zA-Z]+)([0-9]+)([a-zA-Z]+)", item)[0]))
		elif '-' in item:
			temp=item.split('-')[0]
		else:
			if numbered_labels:
				temp=''.join([x for x in item if not x.isdigit()])
				for sub in ("T1","T2"):
					if sub in item:
						temp=item.split(sub)[0] + sub
			else:
				temp=item
		if temp is None:
			temp=item

		values.append(temp)

	vals,indexes,count = np.unique(values, return_index=True, return_counts=True)
	vals=vals[indexes.argsort()]
	count=count[indexes.argsort()]
	return vals,count


hemi = ["lh", "rh"]
surf_suffix = ["pial", "white", "inflated"]

def readRegMatrix(trsfPath):
	with open(trsfPath) as (f):
		return np.loadtxt(f.readlines())

#%%


debug = False
if debug:
	class dotdict(dict):
		"""dot.notation access to dictionary attributes"""
		__getattr__ = dict.get
		__setattr__ = dict.__setitem__
		__delattr__ = dict.__delitem__

	class Namespace:
		def __init__(self, **kwargs):
			self.__dict__.update(kwargs)

	isub="098"
	datap=r'/home/greydon/Documents/data/SEEG/derivatives'

	input=dotdict({
		't1_fname':datap+f'/fastsurfer/sub-P{isub}/mri/orig.mgz',
		'fcsv':datap+ f'/seega_coordinates/sub-P{isub}/sub-P{isub}_space-native_SEEGA.tsv',
		'xfm_noncontrast':datap+f'/atlasreg/sub-P{isub}/sub-P{isub}_desc-rigid_from-noncontrast_to-contrast_type-ras_xfm.txt',
	})

	output=dotdict({
		'html':datap+f'/atlasreg/sub-P{isub}/sub-P{isub}_space-native_electrodes.html',
	})

	params=dotdict({
		'lh_pial':datap+f'/fastsurfer/sub-P{isub}/surf/lh.pial',
		'rh_pial':datap+f'/fastsurfer/sub-P{isub}/surf/rh.pial',
		'lh_sulc':datap+f'/fastsurfer/sub-P{isub}/surf/lh.sulc',
		'rh_sulc':datap+f'/fastsurfer/sub-P{isub}/surf/rh.sulc',
	})

	snakemake = Namespace(output=output, input=input,params=params)

t1_obj = nb.load(snakemake.input.t1_fname)
Torig = t1_obj.header.get_vox2ras_tkr()
#fs_transform=(t1_obj.affine-Torig)+np.eye(4)
fs_transform=np.dot(t1_obj.affine, np.linalg.inv(Torig))

verl,facel=nb.freesurfer.read_geometry(snakemake.params.lh_pial)
verr,facer=nb.freesurfer.read_geometry(snakemake.params.rh_pial)

all_ver = np.concatenate([verl, verr], axis=0)
all_face = np.concatenate([facel, facer+verl.shape[0]], axis=0)
surf_mesh = [all_ver, all_face]

all_ver_shift=(apply_trans(fs_transform, all_ver))

if len(snakemake.input.xfm_noncontrast)>0:
	if os.path.exists(snakemake.input.xfm_noncontrast):
		t1_transform=readRegMatrix(snakemake.input.xfm_noncontrast)
		all_ver_shift=(apply_trans(np.linalg.inv(t1_transform), all_ver_shift))
		#all_ver_shift=(apply_trans(t1_transform, all_ver_shift))


lh_sulc_data = nb.freesurfer.read_morph_data(snakemake.params.lh_sulc)
rh_sulc_data = nb.freesurfer.read_morph_data(snakemake.params.rh_sulc)
bg_map = np.concatenate((lh_sulc_data, rh_sulc_data))


mesh_3d = go.Mesh3d(x=all_ver_shift[:,0], y=all_ver_shift[:,1], z=all_ver_shift[:,2], i=all_face[:,0], j=all_face[:,1], k=all_face[:,2],opacity=.1,color='grey',alphahull=-10)

value=[np.round(x,2) for x in np.arange(.1,.6-.05,.05)]

df = pd.read_table(os.path.splitext(snakemake.input.fcsv)[0]+".tsv",sep='\t',header=0)
groups,n_members=determine_groups(df['label'].tolist(), True)
df['group']=np.repeat(groups,n_members)

cmap = plt.get_cmap('rainbow')
color_maps=cmap(np.linspace(0, 1, len(groups))).tolist()
res = dict(zip(groups, color_maps))

colors=[]
for igroup in df['group']:
	colors.append(res[igroup])

colors=np.vstack(colors)

data=[mesh_3d]
for igroup in groups:
	idx = [i for i,x in enumerate(df['label'].tolist()) if igroup in x]
	data.append(go.Scatter3d(
		x = df['x'][idx].values,
		y = df['y'][idx].values,
		z = df['z'][idx].values,
		name=igroup,
		mode = "markers+text",
		text=df['label'][idx].tolist(),
		textfont=dict(
			family="sans serif",
			size=16,
			color="black"
		),
		textposition = "middle left",
		marker=dict(
			size=5,
			line=dict(
				width=1,
			),
			color=['rgb({},{},{})'.format(int(r*256),int(g*256),int(b*256)) for r,g,b,h in colors[idx]],
			opacity=1
			)))

fig = go.Figure(data=data)
fig.update_layout(scene_camera=CAMERAS['left'],
				  legend_title_text="Electrodes",
				  **LAYOUT)

steps = []
for i in range(len(value)):
	step = dict(
		label = str(f"{value[i]:.2f}"),
		method="restyle",
		args=[{'opacity': [value[i]]+(len(data)-1)*[1],
			 'alphahull': [-10]+(len(data)-1)*[1]
		 }]
	)
	steps.append(step)

sliders = [dict(
	currentvalue={"visible": True,"prefix": "Opacity: ","font":{"size":16}},
	active=0,
	steps=steps,
	x=.35,y=.1,len=.3,
	pad={"t": 8},
)]

fig.update_layout(sliders=sliders)
fig.write_html(snakemake.output.html)

from nilearn import plotting
import matplotlib.pyplot as plt
import matplotlib
from scipy import ndimage
from nilearn import plotting, image
import matplotlib.pyplot as plt
import matplotlib
import nibabel as nib
from nibabel.affines import apply_affine
matplotlib.use('Qt5Agg')
import numpy as np
import base64
import os
from io import BytesIO
import base64
from svgutils.transform import SVGFigure, GroupElement,fromstring
from svgutils.compose import Unit
from tempfile import TemporaryDirectory
from pathlib import Path
from uuid import uuid4
import re
import numpy as np
from matplotlib import gridspec

np.set_printoptions(precision=6,suppress=True)

from nilearn.datasets import load_mni152_template


def svg2str(display_object, dpi):
	"""Serialize a nilearn display object to string."""
	from io import StringIO

	image_buf = StringIO()
	display_object.frame_axes.figure.savefig(
		image_buf, dpi=dpi, format="svg", facecolor="k", edgecolor="k"
	)
	return image_buf.getvalue()

def extract_svg(display_object, dpi=300):
	"""Remove the preamble of the svg files generated with nilearn."""
	image_svg = svg2str(display_object, dpi)

	image_svg = re.sub(' height="[0-9]+[a-z]*"', "", image_svg, count=1)
	image_svg = re.sub(' width="[0-9]+[a-z]*"', "", image_svg, count=1)
	image_svg = re.sub(
		" viewBox", ' preseveAspectRation="xMidYMid meet" viewBox', image_svg, count=1
	)
	start_tag = "<svg "
	start_idx = image_svg.find(start_tag)
	end_tag = "</svg>"
	end_idx = image_svg.rfind(end_tag)

	# rfind gives the start index of the substr. We want this substr
	# included in our return value so we add its length to the index.
	end_idx += len(end_tag)
	return image_svg[start_idx:end_idx]

def clean_svg(fg_svgs, bg_svgs, ref=0):
	# Find and replace the figure_1 id.
	svgs = bg_svgs+fg_svgs
	roots = [f.getroot() for f in svgs]

	sizes = []
	for f in svgs:
		viewbox = [float(v) for v in f.root.get("viewBox").split(" ")]
		width = int(viewbox[2])
		height = int(viewbox[3])
		sizes.append((width, height))
	nsvgs = len([bg_svgs])

	sizes = np.array(sizes)

	# Calculate the scale to fit all widths
	width = sizes[ref, 0]
	scales = width / sizes[:, 0]
	heights = sizes[:, 1] * scales

	# Compose the views panel: total size is the width of
	# any element (used the first here) and the sum of heights
	fig = SVGFigure(Unit(f"{width}px"), Unit(f"{heights[:nsvgs].sum()}px"))

	yoffset = 0
	for i, r in enumerate(roots):
		r.moveto(0, yoffset, scale_x=scales[i])
		if i == (nsvgs - 1):
			yoffset = 0
		else:
			yoffset += heights[i]

	# Group background and foreground panels in two groups
	if fg_svgs:
		newroots = [
			GroupElement(roots[:nsvgs], {"class": "background-svg"}),
			GroupElement(roots[nsvgs:], {"class": "foreground-svg"}),
		]
	else:
		newroots = roots

	fig.append(newroots)
	fig.root.attrib.pop("width", None)
	fig.root.attrib.pop("height", None)
	fig.root.set("preserveAspectRatio", "xMidYMid meet")

	with TemporaryDirectory() as tmpdirname:
		out_file = Path(tmpdirname) / "tmp.svg"
		fig.save(str(out_file))
		# Post processing
		svg = out_file.read_text().splitlines()

	# Remove <?xml... line
	if svg[0].startswith("<?xml"):
		svg = svg[1:]

	# Add styles for the flicker animation
	if fg_svgs:
		svg.insert(
			2,
			"""\
<style type="text/css">
@keyframes flickerAnimation%s { 0%% {opacity: 1;} 100%% { opacity:0; }}
.foreground-svg { animation: 1s ease-in-out 0s alternate none infinite running flickerAnimation%s;}
.foreground-svg:hover { animation-play-state: running;}
</style>"""
			% tuple([uuid4()] * 2),
		)

	return svg

debug = False
if debug:
	class dotdict(dict):
		"""dot.notation access to dictionary attributes"""
		__getattr__ = dict.get
		__setattr__ = dict.__setitem__
		__delattr__ = dict.__delitem__

	class Namespace:
		def __init__(self, **kwargs):
			self.__dict__.update(kwargs)

	isub="070"
	datap=r'/media/veracrypt6/projects/SEEG/derivatives/atlasreg/'

	input=dotdict({
		'img':datap+f'sub-P{isub}/sub-P{isub}_desc-masked_from-atropos3seg_T1w.nii.gz',
		'seg':datap+f'sub-P{isub}/sub-P{isub}_atlas-CerebrA_from-MNI152NLin2009cSym_reg-SyN_dseg.nii.gz',
	})

	output=dotdict({
		#'html':'/home/greydon/Downloads/' + f'sub-P{isub}_from-contrast_to-noncontrast_regqc.html',
		'html':'/home/greydon/Downloads/' + f'sub-P{isub}_desc-affine_from-subject_to-MNI152NLin2009cSym_regqc.html',
		#'png':'/home/greydon/Downloads/' + f'sub-P{isub}_from-contrast_to-noncontrast_regqc.png'
		'png':'/home/greydon/Downloads/' + f'sub-P{isub}_desc-affine_from-subject_to-MNI152NLin2009cSym_regqc.png'
	})

	snakemake = Namespace(output=output, input=input)

	title = 'sub-P001'


#%%

ref_img=nib.load(snakemake.input.img)
ref_resamp = nib.nifti1.Nifti1Image(ref_img.get_fdata(), affine=ref_img.affine,header=ref_img.header)
#ref_resamp = image.resample_img(ref_img, target_affine=np.eye(3), interpolation='continuous')

flo_img=nib.load(snakemake.input.seg)
flo_resamp = nib.nifti1.Nifti1Image(flo_img.get_fdata(), affine=flo_img.affine,header=flo_img.header)
#flo_resamp = image.resample_img(flo_img, target_affine=np.eye(3), interpolation='continuous')


display = plotting.plot_anat(ref_resamp, display_mode='ortho', draw_cross=False,dim=-1,bg_img=None, cut_coords=[0,0,30])
fg_svgs = [fromstring(extract_svg(display,450))]
display.close()

display = plotting.plot_anat(flo_resamp, display_mode='ortho', draw_cross=False,bg_img=None,black_bg=True,cmap=plt.cm.cubehelix, cut_coords=[0,0,30])
bg_svgs = [fromstring(extract_svg(display,450))]
display.close()

final_svg="\n".join(clean_svg(fg_svgs, bg_svgs))

# make figure of thalamic contours
display = plotting.plot_roi(roi_img=flo_resamp, bg_img=ref_resamp, display_mode='ortho', draw_cross=False, cut_coords=[0,0,30])
display.savefig(snakemake.output.png,dpi=300)
display.close()

tmpfile_ref = BytesIO()
display.savefig(tmpfile_ref,dpi=300)
display.close()
tmpfile_ref.seek(0)
data_uri = base64.b64encode(tmpfile_ref.getvalue()).decode('utf-8')
img_tag = '<center><img src="data:image/png;base64,{0}"/></center>'.format(data_uri)

htmlbase='<!DOCTYPE html> <html lang="en"> <head> <title>Slice viewer</title>  <meta charset="UTF-8" /> </head> <body>'
htmlend='</body> </html>'

htmlfull=htmlbase + final_svg + img_tag + htmlend

# Write HTML String to file.html
with open(snakemake.output.html, "w") as file:
	file.write(htmlfull)