From a08aa77afc4919fd00d2c3ef3176595e7d22486b Mon Sep 17 00:00:00 2001
From: DarkPhoenix <phoenix@mail.ru>
Date: Thu, 11 Apr 2019 00:54:42 +0300
Subject: [PATCH] Restore state of command fits on undo

---
 gui/fitCommands/calc/fitAddCommand.py    | 32 ++++++++++++++++--------
 gui/fitCommands/calc/fitRemoveCommand.py | 31 ++++++++++++++++-------
 gui/fitCommands/guiAddCommand.py         |  2 +-
 3 files changed, 45 insertions(+), 20 deletions(-)

diff --git a/gui/fitCommands/calc/fitAddCommand.py b/gui/fitCommands/calc/fitAddCommand.py
index 660d7b5a0..607788b14 100644
--- a/gui/fitCommands/calc/fitAddCommand.py
+++ b/gui/fitCommands/calc/fitAddCommand.py
@@ -1,6 +1,10 @@
 import wx
-import eos.db
 from logbook import Logger
+
+import eos.db
+from service.fit import Fit
+
+
 pyfalog = Logger(__name__)
 
 
@@ -8,28 +12,36 @@ class FitAddCommandCommand(wx.Command):  # well that's an unfrtunate name
     """"
     from sFit.addCommand
     """
-    def __init__(self, fitID, commandFitID):
+    def __init__(self, fitID, commandFitID, state):
         wx.Command.__init__(self, True)
         self.fitID = fitID
         self.commandFitID = commandFitID
+        self.state = state
 
     def Do(self):
-        pyfalog.debug("Projecting command fit ({0}) onto: {1}", self.fitID, self.commandFitID)
-        fit = eos.db.getFit(self.fitID)
-        command = eos.db.getFit(self.commandFitID)
+        pyfalog.debug("Projecting command fit ({0}) onto: {1}".format(self.fitID, self.commandFitID))
+        sFit = Fit.getInstance()
+        fit = sFit.getFit(self.fitID)
+        commandFit = sFit.getFit(self.commandFitID)
 
-        if not command:
+        if not commandFit:
             # if redoing when the command fit has been deleted, simply fail this command
             return False
 
-        if command in fit.commandFits:
-            return
+        if commandFit in fit.commandFits:
+            return False
 
-        fit.commandFitDict[command.ID] = command
+        fit.commandFitDict[commandFit.ID] = commandFit
 
         # this bit is required -- see GH issue # 83
         eos.db.saveddata_session.flush()
-        eos.db.saveddata_session.refresh(command)
+        eos.db.saveddata_session.refresh(commandFit)
+
+        if self.state is not None:
+            commandInfo = commandFit.getCommandInfo(self.fitID)
+            if not commandInfo:
+                return False
+            commandInfo.active = self.state
 
         eos.db.commit()
         return True
diff --git a/gui/fitCommands/calc/fitRemoveCommand.py b/gui/fitCommands/calc/fitRemoveCommand.py
index 29e681b60..360ce9bba 100644
--- a/gui/fitCommands/calc/fitRemoveCommand.py
+++ b/gui/fitCommands/calc/fitRemoveCommand.py
@@ -1,6 +1,10 @@
 import wx
-import eos.db
 from logbook import Logger
+
+import eos.db
+from service.fit import Fit
+
+
 pyfalog = Logger(__name__)
 
 
@@ -11,27 +15,36 @@ class FitRemoveCommandCommand(wx.Command):  # well that's an unfortunate name
     def __init__(self, fitID, commandFitID):
         wx.Command.__init__(self, True)
         self.fitID = fitID
-        self.commandFitID = commandFitID
+        self.savedCommandFitID = commandFitID
+        self.savedState = None
 
     def Do(self):
-        pyfalog.debug("Removing command projection from fit ({0}) for: {1}", self.fitID, self.commandFitID)
-        fit = eos.db.getFit(self.fitID)
-        command = eos.db.getFit(self.commandFitID)
-        if not command:
+        pyfalog.debug("Removing command projection from fit ({0}) for: {1}".format(self.fitID, self.savedCommandFitID))
+        sFit = Fit.getInstance()
+        fit = sFit.getFit(self.fitID)
+        commandFit = sFit.getFit(self.savedCommandFitID)
+        if not commandFit:
             return False
-        del fit.commandFitDict[command.ID]
+
+        commandInfo = commandFit.getCommandInfo(self.fitID)
+        if not commandInfo:
+            return False
+
+        self.savedState = commandInfo.active
+
+        del fit.commandFitDict[commandFit.ID]
 
         eos.db.commit()
         return True
 
     def Undo(self):
-        command = eos.db.getFit(self.commandFitID)
+        command = eos.db.getFit(self.savedCommandFitID)
 
         if not command:
             # can't find the command fit, it must have been deleted. Just skip this undo
             return True
 
         from .fitAddCommand import FitAddCommandCommand
-        cmd = FitAddCommandCommand(self.fitID, self.commandFitID)
+        cmd = FitAddCommandCommand(self.fitID, self.savedCommandFitID, self.savedState)
         cmd.Do()
         return True
diff --git a/gui/fitCommands/guiAddCommand.py b/gui/fitCommands/guiAddCommand.py
index 06cfcbe98..958f57378 100644
--- a/gui/fitCommands/guiAddCommand.py
+++ b/gui/fitCommands/guiAddCommand.py
@@ -16,7 +16,7 @@ class GuiAddCommandCommand(wx.Command):
         self.commandFitID = commandFitID
 
     def Do(self):
-        if self.internal_history.Submit(FitAddCommandCommand(self.fitID, self.commandFitID)):
+        if self.internal_history.Submit(FitAddCommandCommand(self.fitID, self.commandFitID, None)):
             wx.PostEvent(self.mainFrame, GE.FitChanged(fitID=self.fitID))
             self.sFit.recalc(self.fitID)
             return True