from dv3dObjectsToRender import *
from dv3dPlaneWidgets import *
from dv3dDataVolumes import *
from dv3dPropertyFrames import *
from common_functions import *



def planes_setup(the_parent_frame, data_path, volume_type, source_dir = None):
        
        #if we are dealing with the 'base' plane images do this ....
        if volume_type == 'base':
            data_vol = the_parent_frame.my_loaded_volumes[0][0].volume_data.GetOutput()
            
            window, level = 2000, 4000
            DrawPlanesAndGroup(the_parent_frame, data_vol, data_path, -1)
            
        # if we are dealing with new overlay data to add over our base planes ... (e.g. fMRi stats etc.)
        else:
            
            for k in range(len(the_parent_frame.overlay_files_to_load)):  
            #now generate an object - or feature array for a stat object - to store memory for later access
            #how many stats do we already have?
                        
                #we set this counter to zero when we initialise the main window
                #now we increment the count for the data we are about to load
                the_parent_frame.overlay_count += 1
                
                #set up some constants for ease later
                number_to_assign_to_label = the_parent_frame.overlay_count
                index_in_volume_list = the_parent_frame.overlay_count+2
                index_for_planewidget_list = the_parent_frame.overlay_count
                
                overlay_data_label = 'Overlay_%s' %str(number_to_assign_to_label)
                overlay_source_file = the_parent_frame.overlay_files_to_load[k]

                #first load the data to memory
                
                if source_dir == None:
                    curr_vol = create_new_volume(overlay_source_file, overlay_data_label)
                else:
                    curr_vol = create_new_volume(overlay_source_file, overlay_data_label, source_dir)
                                                        
                the_parent_frame.my_loaded_volumes.append([0]) #create a place-holder for new volumes
                the_parent_frame.my_loaded_volumes[index_in_volume_list] = [curr_vol]
                
                voldata = the_parent_frame.my_loaded_volumes[index_in_volume_list][0].volume_data.GetOutput()
                
                ### draw the planes that go in the main window
                
                #use the original routine to draw the planes and create a group
                DrawPlanesAndGroup(the_parent_frame, voldata, overlay_source_file, index_in_volume_list)
                
                #now update the automatic labels and settings with some more appropriate ones
                curr_object_num = len(the_parent_frame.ListOfObjects)-1
                the_parent_frame.ListOfObjects[curr_object_num].my_loaded_volumesindex = index_in_volume_list
                the_parent_frame.ListOfObjects[curr_object_num].my_type = 'overlay_imageplanegroup'
                the_parent_frame.ListOfObjects[curr_object_num].my_label = overlay_data_label
                the_parent_frame.ListOfObjects[curr_object_num].my_struct_or_func = 'functional'
                the_parent_frame.ListOfObjects[curr_object_num].my_property_frame_instance.SetLabel('Properties for %s' %overlay_data_label)
                
                for t in range(1,4):
                    the_parent_frame.ListOfObjects[curr_object_num-t].my_parents_label = overlay_data_label
                    the_parent_frame.ListOfObjects[curr_object_num-t].my_struct_or_func = 'functional'
                    the_parent_frame.ListOfObjects[curr_object_num-t].SetUserControlledLookupTable(1)
                    the_parent_frame.ListOfObjects[curr_object_num-t].GetColorMap().SetLookupTable(curr_vol.lookuptable)
                    the_parent_frame.ListOfObjects[curr_object_num-t].GetColorMap().GetLookupTable().SetTableValue(0, (0.0,0.0,0.0,0.0))
                    #the_parent_frame.ListOfObjects[curr_object_num-t].GetColorMap().SetOutputFormatToRGBA()                       
                
                new_group = the_parent_frame.tree.AppendItem(the_parent_frame.tree_overlays, the_parent_frame.ListOfObjects[-1].my_label)
                
                #add the newly created treeitem as an atribute of the plane_group
                the_parent_frame.ListOfObjects[-1].my_treeitem = new_group
                
                ###
                # set a colour bar for reference of scalar value mapping
                sb = vtk.vtkScalarBarActor()
                sb.SetTitle('')
                # If the orientation is vertical there is a problem.
                #sb.SetOrientationToHorizontal()
                # Vertical is OK.
                sb.SetOrientationToVertical()
                #sb.SetWidth(0.8)
                #sb.SetHeight(0.17)
                sb.SetPosition(0.5, 0.5)
                sb.SetPosition2(0.2, 0.2)
                sb.SetLookupTable(curr_vol.lookuptable)
                
                sbw = vtk.vtkScalarBarWidget()
                sbw.SetScalarBarActor(sb)
                sbw.SetInteractor(the_parent_frame.widget)
                
                the_parent_frame.ListOfObjects[-1].myScalarBar = sbw
                
                ### draw some overlay planes to go in the ortho view window
                DrawPlanes_OverlaysForOrthoWindow(the_parent_frame, voldata, overlay_source_file, index_in_volume_list)
                
                #now update the automatic labels and settings with some more appropriate ones
                curr_object_num = len(the_parent_frame.ListOfObjects)-1
                the_parent_frame.ListOfObjects[curr_object_num].my_loaded_volumesindex = index_in_volume_list
                #the_parent_frame.ListOfObjects[curr_object_num].my_type = 'overlay_imageplanegroup'
                the_parent_frame.ListOfObjects[curr_object_num].my_label = overlay_data_label
                #the_parent_frame.ListOfObjects[curr_object_num].my_struct_or_func = 'functional'
                #the_parent_frame.ListOfObjects[curr_object_num].my_property_frame_instance.SetLabel('Properties for %s' %overlay_data_label)
                
                for t in range(1,4):
                    the_parent_frame.ListOfObjects[curr_object_num-t].my_parents_label = overlay_data_label
                    the_parent_frame.ListOfObjects[curr_object_num-t].my_struct_or_func = 'functional'
                    the_parent_frame.ListOfObjects[curr_object_num-t].SetUserControlledLookupTable(1)
                    the_parent_frame.ListOfObjects[curr_object_num-t].GetColorMap().SetLookupTable(curr_vol.lookuptable)
                    the_parent_frame.ListOfObjects[curr_object_num-t].GetColorMap().GetLookupTable().SetTableValue(0, (0.0,0.0,0.0,0.0))
                    #the_parent_frame.ListOfObjects[curr_object_num-t].GetColorMap().SetOutputFormatToRGBA()                       
                
                #add the previously created treeitem as an atribute of the plane_group
                the_parent_frame.ListOfObjects[-1].my_treeitem = the_parent_frame.ListOfObjects[-5].my_treeitem
                
                
                
                ### ------- DO NOT DELETE - may want to re-implement later
                for t in range(-4,-2):
                    pass
                    #curr_treeItemId = the_parent_frame.tree.AppendItem(new_group, the_parent_frame.ListOfObjects[t].my_label)
                    #add the newly created treeitem as an atribute of the individual plane
                    #the_parent_frame.ListOfObjects[t].my_treeitem = the_parent_frame.tree_root
                ### ------- DO NOT DELETE - may want to re-implement later
                
                
                #refresh t he tree once the items are loaded
                the_parent_frame.tree.Refresh()
                
                #.. and align the overlays to the current base planes
                align_all_planes(the_parent_frame)
                
                    




def DrawPlanes_OverlaysForOrthoWindow(the_parent_frame, data_vol, data_path, index_in_volume_list):
    #TODO -not working- for initial slice positions
    #TODO - completely duplicated code from DrawPlanesAndGroup below
    x,y,z = the_parent_frame.my_loaded_volumes[0][0].x_slice_pos, \
            the_parent_frame.my_loaded_volumes[0][0].y_slice_pos, \
            the_parent_frame.my_loaded_volumes[0][0].z_slice_pos
    
    slice_pos = [x,y,z]
    labels = ['Sagittal plane','Coronal plane','Axial plane']
    widgets = [the_parent_frame.ortho_window.widget1, the_parent_frame.ortho_window.widget3, the_parent_frame.ortho_window.widget2]

    #create the individual planes
    for i in range(3):
        #call the routine to create a plane widget
        planeWidgetX = create_new_planewidget(widgets[i],                              \
                                            data_vol,                                   \
                                            slice_pos[i],   \
                                            i,the_parent_frame.ortho_window.picker,0,the_parent_frame)
        
        # hard coded for now - TODO - load dynamically
        #planeWidgetX.SetWindowLevel(window,level)
        
        #add the individual object to our VTK object list
        AddVTKObjectWithAttributes(planeWidgetX,\
                            the_parent_frame,\
                            'overlay_imageplane_ortho',\
                            labels[i],\
                            'Planes',\
                            None,\
                            6,\
                            7,\
                            1,\
                            'functional',\
                            data_vol,\
                            data_path,\
                            0,\
                            None,\
                            None)


    #now create another vtk object - just the 3 above planes grouped
    # together so that the whole group can later be toggled on / off /
    # updated together
    planeWidgetGroup = planewidget_group(the_parent_frame.ListOfObjects[-3],\
                                         the_parent_frame.ListOfObjects[-2],\
                                         the_parent_frame.ListOfObjects[-1])
    

    # ... and add this object to our list
    AddVTKObjectWithAttributes(planeWidgetGroup,\
                        the_parent_frame,\
                        'ortho_overlay_imageplanegroup',\
                        'Planes',\
                        None,\
                        None,\
                        1,\
                        10,\
                        11,\
                        'functional',\
                        data_vol,\
                        data_path,\
                        0,\
                        None,\
                        None)
    
    
    # ... and apply it to the relevant list object
    the_parent_frame.ListOfObjects[-1].my_property_frame_instance = the_parent_frame.ListOfObjects[-5].my_property_frame_instance





def DrawPlanesAndGroup(the_parent_frame, data_vol, data_path, index_in_volume_list):
    #TODO -not working- for initial slice positions
    x,y,z = the_parent_frame.my_loaded_volumes[0][0].x_slice_pos, \
            the_parent_frame.my_loaded_volumes[0][0].y_slice_pos, \
            the_parent_frame.my_loaded_volumes[0][0].z_slice_pos
    
    slice_pos = [x,y,z]
    labels = ['Sagittal plane','Coronal plane','Axial plane']


    #create the individual planes
    for i in range(3):
        #call the routine to create a plane widget
        planeWidgetX = create_new_planewidget(the_parent_frame.widget,                              \
                                            data_vol,                                   \
                                            slice_pos[i],   \
                                            i,the_parent_frame.picker,0,the_parent_frame)
        
        # hard coded for now - TODO - load dynamically
        #planeWidgetX.SetWindowLevel(window,level)
        
        #add the individual object to our VTK object list
        AddVTKObjectWithAttributes(planeWidgetX,\
                            the_parent_frame,\
                            'base_imageplane',\
                            labels[i],\
                            'Planes',\
                            None,\
                            6,\
                            7,\
                            1,\
                            'structural',\
                            data_vol,\
                            data_path,\
                            0,\
                            None,\
                            None)

    
    #now create another vtk object - just the 3 above planes grouped
    # together so that the whole group can later be toggled on / off /
    # updated together
    planeWidgetGroup = planewidget_group(the_parent_frame.ListOfObjects[-3],\
                                         the_parent_frame.ListOfObjects[-2],\
                                         the_parent_frame.ListOfObjects[-1])
    

    # ... and add this object to our list
    AddVTKObjectWithAttributes(planeWidgetGroup,\
                        the_parent_frame,\
                        'base_imageplanegroup',\
                        'Planes',\
                        None,\
                        None,\
                        1,\
                        10,\
                        11,\
                        'structural',\
                        data_vol,\
                        data_path,\
                        1,\
                        'PlaneWidgetProperties',\
                        -1)
    

    # .. create a properties window for the plane set ...
    # we pass the parent frame and the newly created object as args
    frame = MyPropertyFrame(the_parent_frame, the_parent_frame.ListOfObjects[-1], index_in_volume_list)
    frame.Show(True)
    
    # ... and apply it to the relevant list object
    the_parent_frame.ListOfObjects[-1].my_property_frame_instance = frame
    


