From b9c1ff13dc7e60363573a23754bc06b250d7c020 Mon Sep 17 00:00:00 2001
From: Viktor Sip <viktor.sip@univ-amu.fr>
Date: Mon, 20 Mar 2017 14:40:10 +0100
Subject: [PATCH] Save surfaces to TVB format

---
 bin/run_surfaces_to_struct_dataset.py         |  5 ++--
 .../flow/surfaces_to_structural_datasets.py   | 30 ++++++++++++++++++-
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/bin/run_surfaces_to_struct_dataset.py b/bin/run_surfaces_to_struct_dataset.py
index f73f973..b2bc15e 100644
--- a/bin/run_surfaces_to_struct_dataset.py
+++ b/bin/run_surfaces_to_struct_dataset.py
@@ -8,7 +8,7 @@
 from tvb.recon.cli.runner import SimpleRunner
 
 def main():
-    subjects_dir, subjid, source_lut, target_lut, weights_file, tract_lengths_file, out_file = sys.argv[1:]
+    subjects_dir, subjid, source_lut, target_lut, weights_file, tract_lengths_file, out_file, out_surf = sys.argv[1:]
 
     logging.basicConfig(level=logging.INFO)
     runner = SimpleRunner()
@@ -21,7 +21,8 @@ def main():
         target_lut,
         weights_file,
         tract_lengths_file,
-        out_file
+        out_file,
+        out_surf
     )
 
     flow.run(runner)
diff --git a/tvb/recon/flow/surfaces_to_structural_datasets.py b/tvb/recon/flow/surfaces_to_structural_datasets.py
index bc80165..40f3e30 100644
--- a/tvb/recon/flow/surfaces_to_structural_datasets.py
+++ b/tvb/recon/flow/surfaces_to_structural_datasets.py
@@ -4,8 +4,11 @@
 import os.path
 import logging
 import re
+import tempfile
 import time
 from typing import List, Optional
+from zipfile import ZipFile
+
 from tvb.recon.cli.runner import Runner, File
 from tvb.recon.cli import fs
 from tvb.recon.flow.core import Flow
@@ -37,6 +40,22 @@ def __init__(self, vertices: np.array, triangles: np.array, region_mapping: np.a
         self.triangles = triangles
         self.region_mapping = region_mapping
 
+    def save_surf_zip(self, filename):
+        tmpdir = tempfile.TemporaryDirectory()
+
+        file_vertices = os.path.join(tmpdir.name, 'vertices.txt')
+        file_triangles = os.path.join(tmpdir.name, 'triangles.txt')
+
+        np.savetxt(file_vertices, self.vertices, fmt='%.2f %.2f %.2f')
+        np.savetxt(file_triangles, self.triangles, fmt='%d %d %d')
+
+        with ZipFile(filename, 'w') as zip_file:
+            zip_file.write(file_vertices, os.path.basename(file_vertices))
+            zip_file.write(file_triangles, os.path.basename(file_triangles))
+
+    def save_region_mapping_txt(self, filename):
+        np.savetxt(filename, self.region_mapping, fmt="%d")
+
 
 class ColorLut:
     def __init__(self, filename: os.PathLike):
@@ -356,7 +375,8 @@ def __init__(self,
                  target_lut: os.PathLike,
                  weights_file: os.PathLike,
                  tract_lengths_file: os.PathLike,
-                 struct_zip_file: os.PathLike):
+                 struct_zip_file: os.PathLike,
+                 out_surfaces_dir: os.PathLike=None):
         """
 
         Parameters
@@ -375,6 +395,7 @@ def __init__(self,
         weights_file: text file with weights matrix (which should be upper triangular)
         tract_lengths_file: text file with tract length matrix (which should be upper triangular)
         struct_zip_file: zip file containing TVB structural dataset to be created
+        out_surfaces_dir: directory where to put the surfaces and region mappings in TVB format
         """
         self.cort_surf_direc = cort_surf_direc
         self.subcort_surf_direc = subcort_surf_direc
@@ -384,6 +405,7 @@ def __init__(self,
         self.weights_file = weights_file
         self.tract_lenghts_file = tract_lengths_file
         self.struct_zip_file = struct_zip_file
+        self.out_surfaces_dir = out_surfaces_dir
 
     def run(self, runner: Runner, include_unknown=False):
 
@@ -395,6 +417,12 @@ def run(self, runner: Runner, include_unknown=False):
         surf_subcort = get_subcortical_surfaces(self.subcort_surf_direc, region_index_mapping)
         surf_cort = get_cortical_surfaces(runner, self.cort_surf_direc, self.label_direc, region_index_mapping)
 
+        if self.out_surfaces_dir:
+            surf_subcort.save_region_mapping_txt(os.path.join(self.out_surfaces_dir, "region_mapping_subcort.txt"))
+            surf_subcort.save_surf_zip(os.path.join(self.out_surfaces_dir, "surface_subcort.zip"))
+            surf_cort.save_region_mapping_txt(os.path.join(self.out_surfaces_dir, "region_mapping_cort.txt"))
+            surf_cort.save_surf_zip(os.path.join(self.out_surfaces_dir, "surface_cort.zip"))
+
         region_params_subcort = compute_region_params(surf_subcort, True)
         region_params_cort = compute_region_params(surf_cort, False)