"""
Test of basic project saving and loading
"""
import mantidplottests
from mantidplottests import *
import shutil
import numpy as np
import re
from PyQt4 import QtGui, QtCore


class MantidPlotProjectSerialiseTest(unittest.TestCase):

    def setUp(self):
        self._project_name = "MantidPlotTestProject"
        self._project_folder = os.path.join(os.path.expanduser("~"),
                                            self._project_name)
        file_name = "%s.mantid" % self._project_name
        self._project_file = os.path.join(self._project_folder, file_name)

    def tearDown(self):
        # Clean up project files
        if os.path.isdir(self._project_folder):
            remove_folder(self._project_folder)

        clear_mantid()

    def test_project_file_with_no_data(self):
        workspace_name = "fake_workspace"
        create_dummy_workspace(workspace_name)

        saveProjectAs(self._project_folder)

        self.assertTrue(os.path.isdir(self._project_folder))
        self.assertTrue(os.path.isfile(self._project_file))

        file_text = "MantidPlot 0.9.5 project file\n" \
                    "<scripting-lang>\tPython\n" \
                    "<windows>\t0\n" \
                    "<mantidworkspaces>\n" \
                    "WorkspaceNames\tfake_workspace\n" \
                    "</mantidworkspaces>"

        exp_contents = parse_project_file(file_text)
        contents = read_project_file(self._project_folder)
        self.assertEqual(contents, exp_contents)

    def test_project_file_with_plotted_spectrum(self):
        workspace_name = "fake_workspace"
        create_dummy_workspace(workspace_name)
        plotSpectrum(workspace_name, 1)

        saveProjectAs(self._project_folder)

        self.assert_project_files_saved(workspace_name)
        contents = read_project_file(self._project_folder)

        # Check corrent number of windows
        self.assertEqual(int(contents['<windows>']), 1)

        # Check workspace list was written
        workspace_list = contents['mantidworkspaces']['WorkspaceNames']
        self.assertEqual(workspace_list, workspace_name)

        # Check plot was written
        plot_titles = contents['multiLayer']['graph']['PlotTitle']
        self.assertEqual(len(plot_titles), 3)
        self.assertEqual(plot_titles[0], workspace_name)

    def test_project_file_1D_plot_with_labels_modified(self):
        workspace_name = "fake_workspace"
        create_dummy_workspace(workspace_name
)
        plotSpectrum(workspace_name, [0, 1])

        # modify axes labels
        graph = windows()[0]
        layer = gra
ph.layer(1)
        # call using threadsafe_call to ensure things are executed on the GUI
        # thread, otherwise we get segfaults.
        threadsafe_call(layer.setTitle, "Hello World")
        threadsafe_call(layer.setAxisTitle, 0, "Y Axis Modified")
        threadsafe_call(layer.setAxisTitle, 2, "X Axis Modified")

        saveProjectAs(self._project_folder)

        self.assert_project_files_saved(workspace_name)
        contents = read_project_file(self._project_folder)

        # Check corrent number of windows
        self.assertEqual(int(contents['<windows>']), 1)

        # Check plot title is correct

        plot_title = contents['multiLayer']['graph']['PlotTitle']
        self.assertEqual(len(plot_title), 3)
        self.assertEqual(plot_title[0], "Hello World")

        # Check axes titles are correct
        axes_titles = contents['multiLayer']['graph']['AxesTitles']
        self.assertEqual(len(axes_titles), 2)
        self.assertEqual(axes_titles[0], 'X Axis Modified')
        self.assertEqual(axes_titles[1], 'Y Axis Modified')

    def test_project_file_1D_plot_with_error_bars(self):
        workspace_name = "fake_workspace"
        create_dummy_workspace(workspace_name)
        plotSpectrum(workspace_name, 0, error_bars=True)

        saveProjectAs(self._project_folder)

        self.assert_project_files_saved(workspace_name)
        contents = read_project_file(self._project_folder)
        error_bars = contents['multiLayer']['graph']['MantidYErrors']['1']
        self.assertEqual(len(error_bars), 5)

    def test_project_file_1D_plot_with_axes_scaling(self):
        workspace_name = "fake_workspace"
        create_dummy_workspace(workspace_name)
        plotSpectrum(workspace_name, 0)

        # modify axes scales
        graph = windows()[0]
        layer = graph.layer(1)
        # call using threadsafe_call to ensure things are executed on the GUI
        # thread. Otherwise we get segfaults.
        threadsafe_call(layer.setAxisScale, 0, 10, 10)
        threadsafe_call(layer.logYlinX)

        saveProjectAs(self._project_folder)

        self.assert_project_files_saved(workspace_name)
        contents = read_project_file(self._project_folder)

        # Check axis scales are as expected
        scales = contents['multiLayer']['graph']['scale']
        scale1, scale2, scale3, scale4 = scales[0], scales[1], scales[2], scales[3]

        self.assertAlmostEqual(float(scale1[1]), 110.6670313)
        self.assertEqual(int(scale1[2]), 1000)

        self.assertAlmostEqual(float(scale2[1]), 110.6670313)
        self.assertEqual(int(scale2[2]), 1000)

        self.assertEqual(int(scale3[1]), 0)
        self.assertEqual(int(scale3[2]), 12)

        self.assertEqual(int(scale4[1]), 0)
        self.assertEqual(int(scale4[2]), 12)

    def test_serialise_with_no_data(self):
        workspace_name = "fake_workspace"
        create_dummy_workspace(workspace_name)
        self.save_and_reopen_project()

        # Check that objects were reloaded
        self.assertEqual(rootFolder().name(), self._project_name)
        self.assertEqual(len(windows()), 0)
        self.assertEqual(len(mtd.getObjectNames()), 1)

    def test_serialise_1D_plot_with_plotted_spectrum(self):
        workspace_name = "fake_workspace"
        create_dummy_workspace(workspace_name)

        plotSpectrum(workspace_name, 1)

        self.save_and_reopen_project()

        # Check that objects were reloaded
        self.assertEqual(rootFolder().name(), self._project_name)
        self.assertEqual(len(windows()), 1)
        self.assertEqual(len(mtd.getObjectNames()), 1)

    def test_serialise_1D_plot_with_two_plot_windows(self):
        create_dummy_workspace("ws1")
        create_dummy_workspace("ws2")

        plotSpectrum("ws1", 1)
        plotSpectrum("ws2", 1)

        self.save_and_reopen_project()

        # Check that objects were reloaded
        self.assertEqual(rootFolder().name(), self._project_name)
        self.assertEqual(len(windows()), 2)
        self.assertEqual(len(mtd.getObjectNames()), 2)

        # Check both windows are graph objects
        for window in windows():
            # slight hack as 'type' only returns
            # an MDIWindow instance
            self.assertTrue('Graph' in str(window))

    def test_serialise_1D_plot_with_one_plot_and_multiple_spectra(self):
        workspace_name = "fake_workspace"
        create_dummy_workspace(workspace_name)
        plotSpectrum(workspace_name, [0, 1])

        self.save_and_reopen_project()

        self.assertEqual(rootFolder().name(), self._project_name)
        self.assertEqual(len(mtd.getObjectNames()), 1)
        self.assertEqual(len(windows()), 1)

        graph = windows()[0]
        layer = graph.layer(1)

        # Check graph and layer exist
        self.assertTrue('Graph' in str(graph))
        self.assertTrue(layer is not None)

        # Check plot curves exist
        curve1 = layer.curve(0)
        curve2 = layer.curve(1)
        self.assertTrue('QwtPlotCurve', str(type(curve1)))
        self.assertTrue('QwtPlotCurve', str(type(curve2)))

    def test_serialise_waterfall_plot(self):
        workspace_name = "fake_workspace"
        create_dummy_workspace(workspace_name)
        plotSpectrum(workspace_name, [0, 1], waterfall=True)

        self.save_and_reopen_project()

        # Check that objects were reloaded
        self.assertEqual(rootFolder().name(), self._project_name)
        self.assertEqual(len(windows()), 1)
        self.assertEqual(len(mtd.getObjectNames()), 1)

        # Check window exis