#!/usr/bin/env python
 
# popmodels.py -- models for atomic and fragment orbitals
#
# Copyright (c) Adam Tenderholt, Stanford University, 2006
#                               a-tenderholt@stanford.edu
#
# This program is free software; you can redistribute it and/or modify  
# it under the terms of the GNU General Public License as published by 
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version with the exceptions noted in the 
# Help dialog.

from PyQt4 import QtCore,QtGui
import sys

class AOModel(QtGui.QStandardItemModel):
    def __init__(self,parent=None):
        QtGui.QStandardItemModel.__init__(self,parent)
        self.setHorizontalHeaderLabels(["Element/Orbital"])

        QtCore.QObject.connect(self,QtCore.SIGNAL("rowsRemoved(const QModelIndex&, int, int)"),self.slotRemoveRows)

    def flags(self, index):
        if index.isValid():
            return ( QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable
                    | QtCore.Qt.ItemIsDragEnabled | QtCore.Qt.ItemIsDropEnabled ) 

        return QtCore.Qt.ItemIsEnabled

    def mimeTypes(self):
        types = QtCore.QStringList()
        types << "text/x-atom-orbital"
        return types

    def mimeData(self,indexes):
        mimeData = QtCore.QMimeData()
        encodedData = QtCore.QByteArray()

        stream = QtCore.QDataStream(encodedData, QtCore.QIODevice.WriteOnly)
       
        atoms=[]
        self.atomIndexes=[]
        orbitals=[]
        for index in indexes:
            if index.isValid():
                item = self.itemFromIndex(index)
                #is this an orbital
                if index.parent() != QtCore.QModelIndex():
                    orbitals.append(item)
                    self.atomIndexes.append(index.parent())
                else:
                    atoms.append(item)

        keys=[]
        for j in range(len(atoms)):
            
            count=atoms[j].rowCount()
            for i in range(count):
                keys.append(atoms[j].child(i))

        for orbital in orbitals:
            try:
                index=keys.index(orbital)
            except ValueError:
                keys.append(orbital)

        for key in keys:
            parentName = key.parent().data(QtCore.Qt.DisplayRole).toString()
            name = key.data(QtCore.Qt.DisplayRole).toString()
            string=QtCore.QString(parentName+","+name)
            stream << string
           
        mimeData.setData("text/x-atom-orbital", encodedData)
        return mimeData

    def addOrbitals(self, names):

        atoms = []
        aos = []
        atomname,aoname=names[0].split('_')

#add very first atom
        name = "(%i) %s"%(len(atoms)+1,atomname)
        atoms.append(atomname)
        atomItem = QtGui.QStandardItem(name)
        self.appendRow(atomItem)
        
#add very first orbital
        aoname = "(%i) %s"%(len(aos)+1,aoname)
        aos.append(aoname)
        aoItem = QtGui.QStandardItem(aoname)
        atomItem.appendRow(aoItem)

        for i in range(1,len(names)):
            atomname,aoname=names[i].split('_')
            
            if atomname != atoms[-1]:
                name = "(%i) %s"%(len(atoms)+1,atomname)
                atoms.append(atomname)
                atomItem=QtGui.QStandardItem(name)
                self.appendRow(atomItem)

            aoname = "(%i) %s"%(len(aos)+1,aoname)
            aos.append(aoname)
            aoItem = QtGui.QStandardItem(aoname) 
            atomItem.appendRow(aoItem) #atomItem was last created atom

    def slotRemoveRows(self,parent,start,end):
        if not parent.isValid():
            return

        parentItem = self.itemFromIndex(parent)
        if parentItem.rowCount() > 0:
            return

        row = parentItem.row()
        self.beginRemoveRows(QtCore.QModelIndex(), row, row)
        self.removeRow(row,QtCore.QModelIndex())
        self.endRemoveRows()

class FragModel(QtGui.QStandardItemModel):
    def __init__(self,parent=None):
        QtGui.QStandardItemModel.__init__(self,parent)
        
        self.setHorizontalHeaderLabels(["Fragment/Orbital"])
        
        #QtCore.QObject.connect(self,QtCore.SIGNAL("dataChanged()"),self.slotDeleteMoved)

    def flags(self, index):
        if index.isValid():
            if index.parent() == QtCore.QModelIndex():
                return ( QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable \
                    | QtCore.Qt.ItemIsDropEnabled | QtCore.Qt.ItemIsEditable)
            else:
                return ( QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsDragEnabled)
        return QtCore.Qt.ItemIsEnabled

    def mimeTypes(self):
        types = QtCore.QStringList()
        types << "text/x-frag-orbital"
        types << "text/x-atom-orbital"
        return types

    def suppordedDropActions(self):
        return ( QtCore.Qt.MoveAction | QtCore.Qt.TargetMoveAction )

    def mimeData(self,indexes):
        self.movedItems=[]
        mimeData = QtCore.QMimeData()
        encodedData = QtCore.QByteArray()

        stream = QtCore.QDataStream(encodedData, QtCore.QIODevice.WriteOnly)
       
        for index in indexes:

            if index.isValid():
                item = self.itemFromIndex(index)
                #is this an orbital
                if index.parent() != QtCore.QModelIndex():
                    self.movedItems.append(index)
                    string = item.data(QtCore.Qt.DisplayRole).toString()
                    stream << string
          
        mimeData.setData("text/x-frag-orbital", encodedData)
        return mimeData

    def dropMimeData(self, data, action, row, column, parent):

        if action==QtCore.Qt.IgnoreAction:
            return True
        
        if column > 0:
            return False

        if not parent.isValid():
            beginRow = row;
        else:
            beginRow = parent.row()

        parentItem = self.itemFromIndex(parent)

        #check to make sure that we are dropping on a fragment name
        if parent.parent() != QtCore.QModelIndex():
            return False

        if data.hasFormat("text/x-frag-orbital"):
#re-arranging from inside fragment view

            encodedData = data.data("text/x-frag-orbital")
            stream = QtCore.QDataStream(encodedData, QtCore.QIODevice.ReadOnly)
            newItems = QtCore.QStringList()
            rows = 0

            while not stream.atEnd():
                text = QtCore.QString()
                stream >> text

                #insert into list
                self.beginInsertRows(parent,beginRow,beginRow)
                newItem = QtGui.QStandardItem(text)
                parentItem.appendRow(newItem)
                self.endInsertRows()

            self.emit(QtCore.SIGNAL("dataChanged()"))
            return True

        if data.hasFormat("text/x-atom-orbital"):
#dragging from aoview to fragview

            encodedData = data.data("text/x-atom-orbital")
            stream = QtCore.QDataStream(encodedData, QtCore.QIODevice.ReadOnly)
            newItems = QtCore.QStringList()
            rows = 0

            while not stream.atEnd():
                text = QtCore.QString()
                stream >> text

                #get info from text
                atom,orbital=text.split(',')
                atomName=atom.split(' ')[1]
                index=orbital.indexOf(' ')
                orbNum=int(orbital[1:index-1])
                orbName=orbital[index+1:]

                #insert into list
                self.beginInsertRows(parent,beginRow,beginRow)
                newItem = QtGui.QStandardItem("(%i) %s:%s"%(orbNum, atomName,orbName))
                rowIndex = parentItem.rowCount()
                parentItem.insertRow(rowIndex,newItem)
                self.endInsertRows()

            self.emit(QtCore.SIGNAL("dataChanged()"))
            return True

    def addFragment(self):

        newItem=QtGui.QStandardItem("New Fragment")
        self.appendRow(newItem)

    def setSelectedIndexs(self,indexList):
        
        self.selectedIndexs = indexList

    def removeFragments(self,indexes):

        for index in indexes:

            item = self.itemFromIndex(index)
            if index.parent() != QtCore.QModelIndex():
                continue

            if item.rowCount() != 0:
                text = item.data(QtCore.Qt.DisplayRole).toString()
                QtGui.QMessageBox.information(None,"","The fragment "+text+" cannot be removed until all of its orbitals have been reassigned to a different fragment.","OK")
                continue

            row=item.row()

            self.beginRemoveRows(index.parent(),row,row)
            self.removeRow(row)
            self.endRemoveRows()


    def getFragments(self):

        names=[]
        lists=[]
        
        for i in range(self.rowCount()):
            
            fragment=self.item(i)
            list=[]
            
            for j in range(fragment.rowCount()):
                orbital=fragment.child(j)
                text = str(orbital.data(QtCore.Qt.DisplayRole).toString())
                rindex = text.find(")")
                list.append(int(text[1:rindex])-1)

            if len(list) > 0:
                lists.append(list)
                names.append(fragment.data(QtCore.Qt.DisplayRole).toString())

        return names,lists

class FragView(QtGui.QTreeView):

    def __init__(self,parent=None):
        QtGui.QTreeView.__init__(self,parent)

    def slotRemoveFragments(self):

        selection = self.selectionModel()
        indexes = selection.selectedIndexes()

        self.model().removeFragments(indexes)


if __name__ == "__main__":
    app = QtGui.QApplication(sys.argv)

    view = QtGui.QTreeView()
    model = AOModel(view)
    model.addOrbitals(["C1_1s","C1_2s","N2_1s","N2_2s","C3_1s","C3_2s","C3_2py","C3_2px"])

    frags = FragModel()
    frags.addFragment()
    frags.addFragment()
    fragview = QtGui.QTreeView()
    fragview.setModel(frags)
    fragview.setSelectionMode(QtGui.QAbstractItemView.ExtendedSelection)
    fragview.setDragEnabled(True)
    fragview.setWindowTitle("Fragments")
    fragview.setAcceptDrops(True)
    fragview.setDropIndicatorShown(True)
    fragview.show()

    view.setModel(model)
    view.setSelectionMode(QtGui.QAbstractItemView.ExtendedSelection)
    view.setDragEnabled(True)
    view.setWindowTitle("Atomic Orbitals")
    view.show()

    sys.exit(app.exec_())