#-----------------------------------------------------------------------------
# Copyright 2006-2008, Andre Gouws and York NeuroImaging Centre
#
# Name :          ynicDV3D - ynic Data Viewer 3D
# Description :   multi-modal imaging visualisation toolkit
# Author :        Andre Gouws, Mark Hymers, Will Woods
# Created :       2006-October-06
# Last Update:    2008-Apr-06
# Notes :         requires the following additional files:
#                     dv3dRenderWindowInteractor
#                     dv3dNifti2VtkImageData
#                     dv3dImginfo
#                     dv3dNiftiReader
#                     buttonPanel_planes
#                     buttonPanel_functional
#                     buttonPanel_surfaces
#                     buttonPanel_markers
#                     buttonPanel_colours
#                     buttonPanel_export
#                     buttonPanel_MEG
#                     globalvar, utils
#
# History:        windows/linux alpha 0.1 (Apr 2008)
# Dependencies:   Python 2.4.1 or later, VTK 5.0 or later, currently wxPython 2.6 (not later!)
# 
#
# Major Changes:  uses a main wx app unlike OS X version (pipes between processes)
#
# TODO :
#      build 'loaded object tree ctrl'
#      build main button notebook
#            - add functionality
#      translate major load routines from OS X version
#            - functional data
#            - surfaces
#            - MEG
#            - other (?)
#      Safe exit routine
#      scrollbars for tree ctrl?
#      mouse-tracking for mm update on planewidgets
#      TEST transform code for reference comparison (MNI/TAL)
#      treectrl resize errors- sizers?
#-----------------------------------------------------------------------------


# import usual libraries
import math, os
import vtk
import wx,sys
import time
import random
from scipy import *


# import custom modules
from dv3dImginfo import *
from dv3dNiftiReader import *
from dv3dRenderWindowInteractor import wxVTKRenderWindowInteractor
from dv3dNifti2VtkImageData import nifti2vtkImageData
from dv3dPlanesTextPos import *
from dv3dTrackMouse import *
from dv3dKeyCapture import ProcessKey
from dv3dCoordLookups import *
from dv3dPlaneWidgets import *
from dv3dDataVolumes import *
from dv3dTreeCtrl import *
from dv3dObjectsToRender import *
from dv3d_setup_initial_view import *

#utilities / modules for button panels
from utils import *
import globalvar as globalvar
from buttonPanel_planes import *
from buttonPanel_refvol import *
from buttonPanel_functional import *
from buttonPanel_surfaces import *
from buttonPanel_markers import *
from buttonPanel_colours import *
from buttonPanel_export import *
from buttonPanel_MEG import *


class ProcessKeyRoutine(ProcessKey):
    def __init__(self, parent, frame_parent):
        ProcessKey.__init__(self, parent, frame_parent)

class PlanesPage(PlanesTab):
    def __init__(self, parent, frame_parent):
        PlanesTab.__init__(self, parent, frame_parent)

class FunctionalPage(FunctionalTab):
    def __init__(self, parent, frame_parent):
        FunctionalTab.__init__(self, parent, frame_parent)

class SurfacesPage(SurfacesTab):
    def __init__(self, parent, frame_parent):
        SurfacesTab.__init__(self, parent, frame_parent)

class MarkersPage(MarkersTab):
    def __init__(self, parent, frame_parent):
        MarkersTab.__init__(self, parent, frame_parent)

class ColoursPage(ColoursTab):
    def __init__(self, parent, frame_parent):
        ColoursTab.__init__(self, parent, frame_parent)

class ExportPage(ExportTab):
    def __init__(self, parent, frame_parent):
        ExportTab.__init__(self, parent, frame_parent)

class MEGPage(MEGTab):
    def __init__(self, parent, frame_parent):
        MEGTab.__init__(self, parent, frame_parent)

class REFVolPage(RefTab):
    def __init__(self, parent, frame_parent):
        RefTab.__init__(self, parent, frame_parent)


class DataObjectTree(CreateObjectTree):
    def __init__(self, parent, frame_parent):
        CreateObjectTree.__init__(self, parent, frame_parent)



###################################################################################
# test paramaters for development testing

#file_to_load = 'C:/Documents and Settings/andre/Desktop/DV3D_examples/MNI152_T1_1mm_brain.nii.gz'
file_to_load = 'C:/MNI_no_cer'
#file_to_load = 'C:/avg152T1_brain'
#file_to_load = 'C:/1.nii.gz'
#file_to_load = 'C:/Documents and Settings/andre/Desktop/DV3D_examples/Structural_MRI.nii.gz'
#file_to_load = 'C:/Documents and Settings/andre/Desktop/DV3D_examples/MNI_structural.nii.gz'
#file_to_load = 'C:/Documents and Settings/andre/Desktop/DV3D_examples/dti_FA.nii.gz'


#path to the talairach volume
tal_file_to_load = 'C:/Documents and Settings/andre/Desktop/DV3D_examples/talairach.nii.gz'
tal_lookup_list = 'C:/Documents and Settings/andre/Desktop/DV3D_examples/talairach_labels.txt'

#test stat vols
stat_file_to_load = 'c:/stat.nii.gz' #'C:/Documents and Settings/andre/Desktop/DV3D_examples/MNI_1mm_No_cerebellum'#
###################################################################################




#MAIN APPLICATION CLASS
class Layout(wx.Frame):
    def __init__(self, parent, id, title):

        #?#we need to set the planewidgets as global to initialise them outside the
        #?# main loop ... a quirk of using the vtkRenderwindowinteractor with wx
        # global planeWidgetX, planeWidgetY, planeWidgetz
        
        #empty lists to hold all objects we create
        self.ListOfObjects = []
                
        self.my_loaded_volumes = [[0],[0],[0]] #LVD   # LoadedVolumeData - a list of vtkImageData objects returned from Nifti2VtkImageData
                                    # 1st (0) is always the base structural volume
                                    # 2nd is always the talairach volume
                                    # 3rd is always a place-holder in case a ref volume is loaded
                                    # 4.... are the overlay volumes
                                        
        #a few variable we wilkl reference throughout the program
        self.first_load = 0 # is this load the first or 'base' image
        self.base_data_min, self.base_data_max = 0, 0
        self.ref_loaded = 0 # when we try to load a reference image .. check if we already have one
        self.ShowRefTransform = 0
        self.ShowTalTransform = 0
        self.overlay_count = 1 #allows us to track and reference extra volumes loaded later


        #the input or 'base' data set
        if self.first_load == 0:
            base_vol = create_new_volume(file_to_load,'base data')
            self.first_load =1
            self.my_loaded_volumes[0] = [base_vol]
            #method to load the Talairach reference list automaticall on startup

            f = open(tal_lookup_list, 'rb')
            self.tal_ref_list = f.readlines()
            f.close()
            
            #load the talairach data into and array
            talairach_vol = create_new_volume(tal_file_to_load,'talairach data')
            self.first_load =1
            self.my_loaded_volumes[1] = [talairach_vol]
        else:
            vol, ign1, ign2, ign3, ign4, ign5, ign6,ign7, ign8, ign9, ign10  = nifti2vtkImageData(file_to_load)



        #------------- WX WINDOW ------------------------------------

        # -- Set up the main window with its sub-regions --
        # initialise a wx frame to act as the main app window
        wx.Frame.__init__(self, parent, id, title, size=(800,600))

        self.StatusBar = self.CreateStatusBar()
        self.SetStatusText('hello')

        #the main window has a bakground panel and a sizer
        sizer = wx.BoxSizer(wx.HORIZONTAL)
        panel = wx.Panel(self,-1)
        #split the main window into two side by side boxes
        splitter = wx.SplitterWindow(panel)
        splitter.Bind(wx.EVT_SPLITTER_SASH_POS_CHANGED, self.midChange,id=splitter.GetId())

        #LEFT
        #add a sizer for the left side, a panel for its background
        sizer_left = wx.BoxSizer(wx.VERTICAL)
        panel_left = wx.Panel(splitter,-1)

        #split it vertically and add panels for the top and bottom
        splitter_left = wx.SplitterWindow(panel_left, style=1)
        splitter_left.Bind(wx.EVT_SPLITTER_SASH_POS_CHANGED,self.leftChange,id=splitter_left.GetId())

        panel_left_upper = wx.Panel(splitter_left,style= wx.NO_BORDER)
        panel_left_upper.SetBackgroundColour("WHITE")

        self.panel_left_lower = wx.Panel(splitter_left,style= wx.NO_BORDER)
        splitter_left.SplitHorizontally(panel_left_upper,self.panel_left_lower)
        sizer_left.Add(splitter_left,1,wx.EXPAND)

        #RIGHT
        #add a sizer for the rigth side, a panel for its background
        sizer_right = wx.BoxSizer(wx.VERTICAL)
        panel_right = wx.Panel(splitter,-1)

        #split it vertically and add panels for the top and bottom
        splitter_right =wx.SplitterWindow(panel_right)
        splitter_right.Bind(wx.EVT_SPLITTER_SASH_POS_CHANGED, self.rightChange,id=splitter_right.GetId())
        panel_right_upper = wx.Panel(splitter_right,style= wx.NO_BORDER)



        # -------- Set up the objects with populate each of the sub-regions --
        # Top left has logo(?)
        # Top right has the main button panel - a notebook with multiple tabs
        # Bottom Left has a Tree ctrl which holds all loaded objects
        # Bottom Right hold the VTK window


        # a text box to display messages and coords in
        self.textCtrl1 = wx.TextCtrl(id=-1, name='msgBox1',
               parent=panel_left_upper, pos=wx.Point(0, 0), size = wx.Size(600,600),
               style=wx.TE_MULTILINE|wx.NO_BORDER, value='')
        self.textCtrl1.SetHelpText('')
        self.textCtrl1.SetFont(wx.Font(7, wx.SWISS, wx.NORMAL, wx.NORMAL, False,
              'Arial'))


        # -------- Main Button Panel - Top right --------
        # the notebook object
        self.MainButtonNotebook = wx.Notebook(id=-1, name='self.MainButtonNotebook',
              parent=panel_right_upper, pos=wx.Point(0, 0), style=0)

        #1.) manipulating the main image planes
        self.planes_page = PlanesPage(self.MainButtonNotebook, self)
        self.MainButtonNotebook.AddPage(self.planes_page, "Planes")

        #2.) load a volume to reference against int real time
        self.ref_page = REFVolPage(self.MainButtonNotebook, self)
        self.MainButtonNotebook.AddPage(self.ref_page, "Reference")

        #3.) loading functional data
        self.functional_page = FunctionalPage(self.MainButtonNotebook, self)
        self.MainButtonNotebook.AddPage(self.functional_page, "Functional")

        #4.) loading surface files
        self.surfaces_page = SurfacesPage(self.MainButtonNotebook, self)
        self.MainButtonNotebook.AddPage(self.surfaces_page, "Surfaces")

        #5.) placing markers
        self.markers_page = MarkersPage(self.MainButtonNotebook, self)
        self.MainButtonNotebook.AddPage(self.markers_page, "Markers")

        #6.) manipulating colours
        self.colours_page = ColoursPage(self.MainButtonNotebook, self)
        self.MainButtonNotebook.AddPage(self.colours_page, "Colours")

        #7.) export routines
        self.export_page = ExportPage(self.MainButtonNotebook, self)
        self.MainButtonNotebook.AddPage(self.export_page, "Export")

        #8.) YNiC MEG
        self.MEG_page = MEGPage(self.MainButtonNotebook, self)
        self.MainButtonNotebook.AddPage(self.MEG_page, "MEG")


        # Set the sizer for the main button panel and give it a background
        sizer_panel_ru = wx.GridSizer(rows = 1, cols = 1)
        sizer_panel_ru.Add(self.MainButtonNotebook, 0, wx.EXPAND)
        panel_right_upper.SetSizer(sizer_panel_ru)


        # END Main Button Panel --------
        
        #Add a Tree Control object to hold all the loaded data objects for later
        # interaction
        DataObjectTree(self.panel_left_lower, self)


        panel_right_lower = wx.Panel(splitter_right,style= wx.BORDER_SUNKEN)
        panel_right_lower.SetBackgroundColour("RED")
        splitter_right.SplitHorizontally(panel_right_upper,panel_right_lower)
        sizer_right.Add(splitter_right,1,wx.EXPAND)

        sizer_panel_rl = wx.GridSizer(rows = 1, cols = 1)
        panel_right_lower.SetSizer(sizer_panel_rl)        
        
        panel_right_upper.SetSizer(sizer_panel_ru)        
        splitter.SplitVertically(panel_left,panel_right)
        sizer.Add(splitter,1,wx.EXPAND)
        panel.SetSizer(sizer)
        panel_left.SetSizer(sizer_left)
        panel_right.SetSizer(sizer_right)

        self.splitter = splitter        
        self.splitter_left = splitter_left
        self.splitter_right = splitter_right
        self.splitter.SetMinimumPaneSize(150)
        self.splitter_left.SetMinimumPaneSize(135)
        self.splitter_right.SetMinimumPaneSize(135)
        
        self.splitter_left.SetSashPosition(135)
        self.splitter_right.SetSashPosition(135)
        self.splitter.SetSashPosition(150)
        # END OF WX STUFF



        #------------- VTK WINDOW ------------------------------------
        #the interactor
        self.widget = wxVTKRenderWindowInteractor(panel_right_lower ,-1)

        #allow key bindings to be sent to the interactor
        self.widget.AddObserver("KeyPressEvent", self.myKeyEvents)

        # The shared picker enables us to use 3 planes at one time
        # and gets the picking order right
        self.picker = vtk.vtkCellPicker()
        self.picker.SetTolerance(0.005)
        

        #add a safe exit routine
        self.widget.AddObserver("ExitEvent", lambda o,e,f=Layout: sys.exit())

        #define the vtk renderer to do the drawing
        self.ren = vtk.vtkRenderer()

        #add the renderwindowinteractor to the wx application
        sizer_panel_rl.Add(self.widget, 0, wx.EXPAND)

        #tell the renderer to draw the widget's output
        self.widget.GetRenderWindow().AddRenderer(self.ren)

        #now enable and initialise the vtk components
        self.widget.Enable(1)



        #the base image planes x, y and z
        base_vol = self.my_loaded_volumes[0][0].volume_data.GetOutput()
        #???nb WE SET THE VOL FOR ALL WIDGETS TO REFERENCE BEFORE (EASE)???pissed?
        # X
        
        initial_planes_setup(self)

        #planeWidgetX = create_new_planewidget(self.widget,                              \
        #                                    base_vol,                                   \
        #                                    self.my_loaded_volumes[0][0].x_slice_pos,   \
        #                                    0,self.picker,0)
        ## hard coded for now - TODO - load dynamically
        #planeWidgetX.SetWindowLevel(4000,8000)
        ##planeWidgetX.On()
        #AddVTKObjectWithAttributes(planeWidgetX,\
        #                    self,\
        #                    'base_imageplane',\
        #                    'Sagittal plane',\
        #                    'Base volume',\
        #                    1,\
        #                    -1,\
        #                    -1,\
        #                    'structural',\
        #                    base_vol)
        #
        #planeWidgetY = create_new_planewidget(self.widget,                              \
        #                                    base_vol,                                   \
        #                                    self.my_loaded_volumes[0][0].y_slice_pos,   \
        #                                    1,self.picker,0)
        ##planeWidgetY.On()
        #AddVTKObjectWithAttributes(planeWidgetY,\
        #                    self,\
        #                    'base_imageplane',\
        #                    'Coronal plane',\
        #                    'Base volume',\
        #                    1,\
        #                    -1,\
        #                    -1,\
        #                    'structural',\
        #                    base_vol)
        #
        #planeWidgetZ = create_new_planewidget(self.widget,                              \
        #                                    base_vol,                                   \
        #                                    self.my_loaded_volumes[0][0].z_slice_pos,   \
        #                                    2,self.picker,0)
        ##planeWidgetZ.On()
        #AddVTKObjectWithAttributes(planeWidgetZ,\
        #                    self,\
        #                    'base_imageplane',\
        #                    'Axial plane',\
        #                    'Base volume',\
        #                    1,\
        #                    -1,\
        #                    -1,\
        #                    'structural',\
        #                    base_vol)
        #
        #self.planeWidget1 = planewidget_group(planeWidgetX,planeWidgetY,planeWidgetZ)
        #AddVTKObjectWithAttributes(planeWidgetY,\
        #                    self,\
        #                    'base_imageplanegroup',\
        #                    'Base volume',\
        #                    None,\
        #                    1,\
        #                    -1,\
        #                    -1,\
        #                    'structural',\
        #                    base_vol)
        
        
        
        #self.PWG[0] = [self.planeWidget1,0]

        
        ### initial base image setup
        #TODO -? - only everonce - for the main base-data planes
        #self.PWG[0][0].planewidgets[0].AddObserver("InteractionEvent", self.MouseTracker)
        #self.PWG[0][0].planewidgets[1].AddObserver("InteractionEvent", self.MouseTracker)
        #self.PWG[0][0].planewidgets[2].AddObserver("InteractionEvent", self.MouseTracker)

        # we may want to reference this later to reset
        #self.original_LU_table = self.PWG[0][0].planewidgets[0].GetLookupTable()


        
        
        #for i in range(3):
        #    self.PWG[0][0].planewidgets[i].On()
        #

        #self.PWG[0][0].planewidgets[0].SetSliceIndex(self.my_loaded_volumes[0][0].x_slice_pos,)
        #self.PWG[0][0].planewidgets[1].SetSliceIndex(self.my_loaded_volumes[0][0].y_slice_pos,)
        #self.PWG[0][0].planewidgets[2].SetSliceIndex(self.my_loaded_volumes[0][0].z_slice_pos,)        
        

    def myKeyEvents(self,  *args, **kwargs):
        self.key=self.widget.GetKeyCode() #why GetKeyCode on Windows? = GetKeySym on mac
        ProcessKeyRoutine(self.MainButtonNotebook, self)
    
        
    #------------
    #Functions for interaction in the main window

    # track the mouse interaction with gthe planes and ouput the slice number, mm and ref details if loaded
    def MouseTracker(self, *args, **kwargs):
        #scan the list and find groups of planes
        for i in range(len(self.ListOfObjects)):
            if self.ListOfObjects[i].my_type == 'overlay_imageplanegroup':
                self.ListOfObjects[i].setlocation(self.ListOfObjects[3])
            #for j in range(3):
                #TODO - dont do this here
                #y = vtk.vtkLookupTable()
                #y.SetNumberOfTableValues(256)
                #y.SetTableRange(0,20)
                #y.Build()
                #if i == 1:
                #    self.PWG[i][0].planewidgets[j].GetColorMap().SetLookupTable(y)
                #    self.PWG[i][0].planewidgets[j].GetColorMap().GetLookupTable().SetTableValue(0, (0.0,0.0,0.0,0.0))
                #    self.PWG[i][0].planewidgets[j].GetColorMap().SetOutputFormatToRGBA()
                #elif i == 2:
                #    self.PWG[i][0].planewidgets[j].GetColorMap().SetLookupTable(y)
                #    self.PWG[i][0].planewidgets[j].GetColorMap().GetLookupTable().SetTableValue(0, (0.0,0.0,0.0,0.0))
                #    self.PWG[i][0].planewidgets[j].GetColorMap().SetOutputFormatToRGBA()                   
                #else:
                #    self.PWG[i][0].planewidgets[j].GetColorMap().SetLookupTable(y)
                #    self.PWG[i][0].planewidgets[j].GetColorMap().SetOutputFormatToRGBA()
                #    self.PWG[i][0].planewidgets[j].GetColorMap().GetLookupTable().SetTableValue(0, (0.0,0.0,0.0,0.0))
                        
        #if self.PWG[0][0].planewidgets[0].GetCursorDataStatus() == 1:
        #    self.cursor_curr_slice_num = self.PWG[0][0].planewidgets[0].GetCurrentCursorPosition()
        #    display_coords(self, self)
        #
        #elif self.PWG[0][0].planewidgets[1].GetCursorDataStatus() == 1:
        #    self.cursor_curr_slice_num = self.PWG[0][0].planewidgets[1].GetCurrentCursorPosition()
        #    display_coords(self, self)
        #
        #elif self.PWG[0][0].planewidgets[2].GetCursorDataStatus() == 1:
        #    self.cursor_curr_slice_num = self.PWG[0][0].planewidgets[2].GetCurrentCursorPosition()
        #    display_coords(self, self)
        #
        #else:
        #    return

    #control routines for the sub-window sizers - thes make sure
    # that if we resize the left the right resizes to the same
    ## weird? - have to loop twice and tell each ibject to load
    ## it own parameters otherwise panels dont size properly(?)
    def leftChange(self, event):
        for i in range(2):
            pos = self.splitter_left.GetSashPosition()
            self.splitter_right.SetSashPosition(pos)
            pos = self.splitter_right.GetSashPosition()
            self.splitter_left.SetSashPosition(pos)
            event.Skip()
            self.tree.SetSize(self.treepanel.GetSize())
            self.treepanel.Refresh()

    def midChange(self, event):
        for i in range(2):
            pos = self.splitter.GetSashPosition()
            self.splitter.SetSashPosition(pos)
            pos = self.splitter_left.GetSashPosition()
            self.splitter_right.SetSashPosition(pos)
            pos = self.splitter_right.GetSashPosition()
            self.splitter_left.SetSashPosition(pos)
            event.Skip()
            self.tree.SetSize(self.treepanel.GetSize())
            self.treepanel.Refresh()

    def rightChange(self, event):
        for i in range(2):
            pos = self.splitter_right.GetSashPosition()
            self.splitter_left.SetSashPosition(pos)
            pos = self.splitter_left.GetSashPosition()
            self.splitter_right.SetSashPosition(pos)
            event.Skip()
            self.tree.SetSize(self.treepanel.GetSize())
            self.treepanel.Refresh()


#The main application loop
app = wx.App(0)
k = Layout(None, -1, 'v0.1')
k.Show(True)
app.MainLoop()
