#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Sep  4 15:25:17 2020

@author: p277634
"""
import numpy as np
from scipy.io import loadmat
from itertools import combinations

class BNVfiles():
    def __init__(self):
        self.fileNodes  = './data/Desikan-Killiany68_myTest.txt'
        self.fileMapper = './data/mapFS2DTI.mat'
        self.pathsave   = './dataSave/' 
        self.nroi       = 68
        self.pairs      = np.array(list(combinations(range(self.nroi),2)), dtype=np.int32)
        self._getMapper()
        
    def _getMapper(self):
        self.map    = loadmat(self.fileMapper)['map']['raw2cross'][0][0][0]-1
               
    def doNodes(self, colors=None, sizes=None, name='nodes'):
        """ Desikan parcelation by default
        @ colors: [nRoi], scale for coloring nodes
        @ sizes: [nROi], scale for the size of nodes """

        A = np.loadtxt(self.fileNodes)
        if colors is not None:
            A[:,3] = colors[self.map]
            
        if sizes is not None:
            A[:,4] = sizes[self.map]
        
        np.savetxt(self.pathsave+name + '.node', A)
            
    def doEdges(self, edges, weights, name='edges'):
        """ @ edges: directed FC with sign
            @ weights: strength of edge """
        ## Edges are transposed with respect to brain connectivity toolbox
        fc = np.zeros((self.nroi, self.nroi))
        for i in range(self.pairs.shape[0]):
            if edges[i] > 0:
                fc[self.map[self.pairs[i,0]], self.map[self.pairs[i,1]]] = weights[i]
            elif edges[i] < 0:
                fc[self.map[self.pairs[i,1]], self.map[self.pairs[i,0]]] = weights[i]
        
        np.savetxt(self.pathsave+name + '.edge', fc.T)