# **************************************************************************
# *
# * Authors:     J.M. de la Rosa Trevin (jmdelarosa@cnb.csic.es)
# *
# * Unidad de  Bioinformatica of Centro Nacional de Biotecnologia , CSIC
# *
# * 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 3 of the License, or
# * (at your option) any later version.
# *
# * This program is distributed in the hope that it will be useful,
# * but WITHOUT ANY WARRANTY; without even the implied warranty of
# * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# * GNU General Public License for more details.
# *
# * You should have received a copy of the GNU General Public License
# * along with this program; if not, write to the Free Software
# * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
# * 02111-1307  USA
# *
# *  All comments concerning this program package may be sent to the
# *  e-mail address 'scipion@cnb.csic.es'
# *
# **************************************************************************

"""
Test the execution of the initial volume service with different datasets.
Execute:
1) Import averages
2) Run initial volume algorithms (ransac, eman, significant, prime)
3) Run align volumes

Datasets:
1) Groel
2) Ribosome
3) BPV
"""

import os
from os.path import exists

import pyworkflow.tests as pwtests

from pwem import Domain, Config
from pwem.protocols import ProtImportAverages


class TestGroel(pwtests.BaseTest):

    @classmethod
    def setUpClass(cls):
        # Create a new project
        pwtests.setupTestProject(cls)
        cls.ds = pwtests.DataSet.getDataSet('initial_volume')

    def test_groel(self):
        """ Run an Import particles protocol. """
        cpus = os.environ.get('SCIPION_TEST_CPU', 4)
        # 1. Run import of averages
        groelAvg = self.ds.getFile('groel')
        sym = 'd7'
        protImport = self.newProtocol(ProtImportAverages,
                                      objLabel='import averages (groel)',
                                      filesPath=groelAvg,
                                      samplingRate=1)
        self.launchProtocol(protImport)

        # 2. Run initial models
        # 2a. Ransac
        xmipp3 = Domain.importFromPlugin('xmipp3.protocols', doRaise=True)
        protRansac = self.newProtocol(xmipp3.XmippProtRansac,
                                      objLabel='xmipp - ransac',
                                      symmetryGroup=sym,
                                      numberOfMpi=1,
                                      numberOfThreads=cpus
                                      )
        protRansac.inputSet.set(protImport.outputAverages)
        self.launchProtocol(protRansac)

        # 2b. Eman
        eman2 = Domain.importFromPlugin('eman2.protocols', doRaise=True)
        protEmanInitVol = self.newProtocol(eman2.EmanProtInitModel,
                                           objLabel='eman - initial vol',
                                           symmetry=sym,
                                           numberOfThreads=cpus)
        protEmanInitVol.inputSet.set(protImport.outputAverages)
        self.launchProtocol(protEmanInitVol)

        # 3. Significant
        protSignificant = self.newProtocol(xmipp3.XmippProtReconstructSignificant,
                                           objLabel='xmipp - significant',
                                           symmetryGroup=sym,
                                           numberOfMpi=cpus,
                                           numberOfThreads=1,
                                           iter=15,
                                           alpha0=95
                                           )
        protSignificant.inputSet.set(protImport.outputAverages)
        self.launchProtocol(protSignificant)

        # 4. Align all volumes
        protAlign = self.newProtocol(xmipp3.XmippProtAlignVolumeForWeb,
                                     objLabel='xmipp - align volumes',
                                     numberOfMpi=1,
                                     numberOfThreads=cpus
                                     )
        protAlign.inputReference.set(protSignificant.outputVolume)
        protAlign.inputVolumes.append(protRansac.outputVolumes)
        protAlign.inputVolumes.append(protEmanInitVol.outputVolumes)
        protAlign.inputVolumes.append(protSignificant.outputVolume)
        self.launchProtocol(protAlign)


class TestBPV(pwtests.BaseTest):

    @classmethod
    def setUpClass(cls):
        # Create a new project
        pwtests.setupTestProject(cls)
        cls.ds = pwtests.DataSet.getDataSet('initial_volume')

    def test_bpv(self):
        """ Run an Import particles protocol. """
        cpus = os.environ.get('SCIPION_TEST_CPU', 4)
        # 1. Run import of averages
        groelAvg = self.ds.getFile('bpv')
        sym = 'i1'
        protImport = self.newProtocol(ProtImportAverages,
                                      objLabel='import averages (bpv)',
                                      filesPath=groelAvg,
                                      samplingRate=1)
        self.launchProtocol(protImport)

        # 2. Run initial models
        # 2a. Ransac
        xmipp3 = Domain.importFromPlugin('xmipp3.protocols', doRaise=True)
        protRansac = self.newProtocol(xmipp3.XmippProtRansac,
                                      objLabel='xmipp - ransac',
                                      objComment='Since there are only 8 projections, a dimensionality reduction '
                                                 'cannot be safely done. In this case, it is better to take only 3 '
                                                 'images in every RANSAC iterations and lower the inlier threshold to '
                                                 '0.65 so that more images can have the chances of being considered '
                                                 'during the reconstruction process.',
                                      symmetryGroup=sym,
                                      dimRed=False,
                                      numSamples=3,  # less than 8 classes provided
                                      corrThresh=0.65,
                                      numberOfMpi=1,
                                      numberOfThreads=cpus
                                      )
        protRansac.inputSet.set(protImport.outputAverages)
        self.launchProtocol(protRansac)

        # 2b. Eman

        eman2 = Domain.importFromPlugin('eman2.protocols', doRaise=True)

        protEmanInitVol = self.newProtocol(eman2.EmanProtInitModel,
                                           objLabel='eman - initial vol',
                                           symmetry='icos',
                                           numberOfThreads=cpus)
        protEmanInitVol.inputSet.set(protImport.outputAverages)
        self.launchProtocol(protEmanInitVol)

        # 3. Significant
        protSignificant = self.newProtocol(xmipp3.XmippProtReconstructSignificant,
                                           objLabel='xmipp - significant',
                                           symmetryGroup=sym,
                                           numberOfMpi=cpus,
                                           numberOfThreads=1,
                                           iter=15,
                                           alpha0=99.0
                                           )
        protSignificant.inputSet.set(protImport.outputAverages)
        self.launchProtocol(protSignificant)

        # 4. Align all volumes
        protAlign = self.newProtocol(xmipp3.XmippProtAlignVolumeForWeb,
                                     objLabel='xmipp - align volumes',
                                     numberOfMpi=1,
                                     numberOfThreads=cpus
                                     )
        protAlign.inputReference.set(protSignificant.outputVolume)
        protAlign.inputVolumes.append(protRansac.outputVolumes)
        protAlign.inputVolumes.append(protEmanInitVol.outputVolumes)
        protAlign.inputVolumes.append(protSignificant.outputVolume)
        self.launchProtocol(protAlign)


class TestRibosome(pwtests.BaseTest):

    @classmethod
    def setUpClass(cls):
        # Create a new project
        pwtests.setupTestProject(cls)
        cls.ds = pwtests.DataSet.getDataSet('initial_volume')

    def test_ribosome(self):
        """ Run an Import particles protocol. """
        cpus = os.environ.get('SCIPION_TEST_CPU', 4)
        # 1. Run import of averages
        groelAvg = self.ds.getFile('ribosome')
        sym = 'c1'
        protImport = self.newProtocol(ProtImportAverages,
                                      objLabel='import averages (ribosome)',
                                      filesPath=groelAvg,
                                      samplingRate=1)
        self.launchProtocol(protImport)

        # 2. Run initial models
        # 2a. Ransac
        xmipp3 = Domain.importFromPlugin('xmipp3.protocols', doRaise=True)
        protRansac = self.newProtocol(xmipp3.XmippProtRansac,
                                      objLabel='xmipp - ransac',
                                      symmetryGroup=sym,
                                      numberOfMpi=1,
                                      numberOfThreads=cpus
                                      )
        protRansac.inputSet.set(protImport.outputAverages)
        self.launchProtocol(protRansac)

        # 2b. Eman
        eman2 = Domain.importFromPlugin('eman2.protocols', doRaise=True)

        protEmanInitVol = self.newProtocol(eman2.EmanProtInitModel,
                                           objLabel='eman - initial vol',
                                           symmetry=sym,
                                           numberOfThreads=cpus)
        protEmanInitVol.inputSet.set(protImport.outputAverages)
        self.launchProtocol(protEmanInitVol)

        # 3. Significant
        protSignificant = self.newProtocol(xmipp3.XmippProtReconstructSignificant,
                                           objLabel='xmipp - significant',
                                           symmetryGroup=sym,
                                           numberOfMpi=cpus,
                                           numberOfThreads=1,
                                           iter=15
                                           )
        protSignificant.inputSet.set(protImport.outputAverages)
        self.launchProtocol(protSignificant)

        # 4. Align all volumes
        protAlign = self.newProtocol(xmipp3.XmippProtAlignVolumeForWeb,
                                     objLabel='xmipp - align volumes',
                                     numberOfMpi=1,
                                     numberOfThreads=cpus
                                     )
        protAlign.inputReference.set(protSignificant.outputVolume)
        protAlign.inputVolumes.append(protRansac.outputVolumes)
        protAlign.inputVolumes.append(protEmanInitVol.outputVolumes)
        protAlign.inputVolumes.append(protSignificant.outputVolume)
        self.launchProtocol(protAlign)


class TestSignificant(pwtests.BaseTest):
    """ Test only significant execution with BPV virus. """

    @classmethod
    def setUpClass(cls):
        # Create a new project
        pwtests.setupTestProject(cls)
        cls.ds = pwtests.DataSet.getDataSet('initial_volume')

    def _runSignificant(self, inputSet, args):
        myargs = dict(args)
        xmipp3 = Domain.importFromPlugin('xmipp3.protocols', doRaise=True)
        prot1 = self.newProtocol(xmipp3.XmippProtReconstructSignificant,
                                 objLabel='significant d7',
                                 **myargs
                                 )
        # prot1.inputClasses.set(inputSet)
        self.launchProtocol(prot1)

        output = prot1.outputVolume

        myargs['thereisRefVolume'] = True

        prot2 = self.newProtocol(xmipp3.XmippProtReconstructSignificant,
                                 objLabel='significant d7 (with ref)',
                                 **myargs
                                 )
        # prot2.inputClasses.set(inputSet)
        prot2.refVolume.set(output)
        self.launchProtocol(prot2)

        # Run significant with one volume and new target resolution
        myargs['useMaxRes'] = True
        prot3 = self.newProtocol(xmipp3.XmippProtReconstructSignificant,
                                 objLabel='significant d7 (ref and maxRes)',
                                 **myargs
                                 )
        prot3.refVolume.set(output)
        self.launchProtocol(prot3)

    def test_significant(self):
        """ Run an Import particles protocol. """
        cpus = os.environ.get('SCIPION_TEST_CPU', 4)
        gpuFlag = exists(Config.CUDA_BIN)
        # 1. Run import of averages
        avg = self.ds.getFile('groel')

        protImport = self.newProtocol(ProtImportAverages,
                                      filesPath=avg,
                                      samplingRate=1)
        self.launchProtocol(protImport)

        args = {'symmetryGroup': 'd7',
                'iter': 3,
                'numberOfMpi': cpus,
                'inputSet': protImport.outputAverages,
                'alpha0': 98.0
                }
        # Run significant with one volume
        self._runSignificant(protImport.outputAverages, args)

        # Run more test with GPU if it is possible
        if gpuFlag:
            args = {'symmetryGroup': 'd7',
                    'iter': 3,
                    'useGPU': True,
                    'inputSet': protImport.outputAverages,
                    'alpha0': 98.0
                    }
            # Run significant with GPU and new target resolution
            self._runSignificant(protImport.outputAverages, args)
