import os
import numpy as np
import json

import util_amira

def getEdgeLabelName(label):
    if(label == 6):
        return "axon"
    elif(label == 4):
        return "apical"
    elif(label == 5):
        return "basal"
    elif(label == 7):
        return "soma"
    else:
        return "other"


def getSomaPosition(points):
    somaPos = []
    for p in points:
        if(p["edge_label"] == "soma"):
            somaPos.append(p["position"])
    return np.mean(np.vstack(tuple(somaPos)), axis=0)


def loadAmiraExport(filename):
    with open(filename) as f:
        lines = f.readlines()
        labels = lines[0].rstrip().split(",")
        points = []
        for i in range(1, len(lines)):
            line = lines[i].rstrip().split(",")

            point = {}
            point["edge_id"] = int(line[labels.index("edge_id")])
            point["source_node_id"] = int(line[labels.index("source_node")])
            point["target_node_id"] = int(line[labels.index("target_node")])
            point["edge_label"] = getEdgeLabelName(
                int(line[labels.index("edge_label")]))
            point["edge_point_id"] = int(line[labels.index("edge_point")])
            point["position"] = np.array([float(line[labels.index("x")]), float(
                line[labels.index("y")]), float(line[labels.index("z")])])
            point["radius"] = float(line[labels.index("radius")])
            point["inside_vS1"] = int(line[labels.index("inside_vS1")])
            if(point["edge_label"] != "other"):
                points.append(point)

    return points


def separateCompartments(edgePoints):
    apical = []
    basal = []
    axon = []
    for edgePoint in edgePoints:
        if(edgePoint["edge_label"] == "apical"):
            apical.append(edgePoint)
        elif(edgePoint["edge_label"] == "basal"):
            basal.append(edgePoint)
        elif(edgePoint["edge_label"] == "axon"):
            axon.append(edgePoint)
    compartments = {}
    compartments["apical"] = apical
    compartments["basal"] = basal
    compartments["axon"] = axon
    return compartments


def loadGraphset(networkDir):
    if(os.path.exists(os.path.join(networkDir, "morphologies", "Morphologies.am"))):
        graphset = util_amira.readSpatialGraphSet(os.path.join(networkDir, "morphologies", "Morphologies.am"), legacy=False)
    else:    
        graphset = util_amira.readSpatialGraphSet(os.path.join(networkDir, "morphologies", "MorphologiesWithNeuronIDs.am"), legacy=True)  
    return graphset


def writeToCache(filename, transformation, neuronId):
    transformationFile = "/tmp/transformation_{}".format(neuronId)
    np.savetxt(transformationFile, transformation)

    meta = {
        "morphologyFile" : filename,
        "transformationFile" : transformationFile
    }
    metaFile = "/tmp/meta_{}.json".format(neuronId)
    with open(metaFile, "w") as f:
        print("meta", meta)
        json.dump(meta, f)
    

def readFromCache(neuronId):
    metaFile = "/tmp/meta_{}.json".format(neuronId)
    with open(metaFile) as f:
        meta = json.load(f)

    transformationFile = meta["transformationFile"]
    T = np.loadtxt(transformationFile)
    morphologyFile = meta["morphologyFile"]

    return morphologyFile, T


def loadAxon(graphset, neuronId, saveToCache = False, loadFromCache = False):    
    if(loadFromCache):
        filename, T = readFromCache(neuronId)
    else:
        idx = len(graphset[neuronId]) - 1                   
        filename = graphset[neuronId][idx]["file"]
        T = graphset[neuronId][idx]["transformation"] 
        if(saveToCache):
            writeToCache(filename, T, neuronId)
    
    return util_amira.readSpatialGraph(filename, T)


def loadDendrite(graphset, neuronId, saveToCache = False, loadFromCache = False):
    if(loadFromCache):
        filename, T = readFromCache(neuronId)
    else:
        filename = graphset[neuronId][0]["file"]
        T = graphset[neuronId][0]["transformation"] 
        if(saveToCache):
            writeToCache(filename, T, neuronId)
    
    return util_amira.readSpatialGraph(filename, T)
[((2638, 2684), 'numpy.savetxt', 'np.savetxt', (['transformationFile', 'transformation'], {}), '(transformationFile, transformation)\n', (2648, 2684), True, 'import numpy as np\n'), ((3140, 3170), 'numpy.loadtxt', 'np.loadtxt', (['transformationFile'], {}), '(transformationFile)\n', (3150, 3170), True, 'import numpy as np\n'), ((3664, 3704), 'util_amira.readSpatialGraph', 'util_amira.readSpatialGraph', (['filename', 'T'], {}), '(filename, T)\n', (3691, 3704), False, 'import util_amira\n'), ((4059, 4099), 'util_amira.readSpatialGraph', 'util_amira.readSpatialGraph', (['filename', 'T'], {}), '(filename, T)\n', (4086, 4099), False, 'import util_amira\n'), ((2150, 2209), 'os.path.join', 'os.path.join', (['networkDir', '"""morphologies"""', '"""Morphologies.am"""'], {}), "(networkDir, 'morphologies', 'Morphologies.am')\n", (2162, 2209), False, 'import os\n'), ((2915, 2933), 'json.dump', 'json.dump', (['meta', 'f'], {}), '(meta, f)\n', (2924, 2933), False, 'import json\n'), ((3066, 3078), 'json.load', 'json.load', (['f'], {}), '(f)\n', (3075, 3078), False, 'import json\n'), ((2263, 2322), 'os.path.join', 'os.path.join', (['networkDir', '"""morphologies"""', '"""Morphologies.am"""'], {}), "(networkDir, 'morphologies', 'Morphologies.am')\n", (2275, 2322), False, 'import os\n'), ((2402, 2474), 'os.path.join', 'os.path.join', (['networkDir', '"""morphologies"""', '"""MorphologiesWithNeuronIDs.am"""'], {}), "(networkDir, 'morphologies', 'MorphologiesWithNeuronIDs.am')\n", (2414, 2474), False, 'import os\n')]