import OpenPNM
import scipy as sp
import pytest
mgr = OpenPNM.Base.Workspace()
mgr.loglevel = 60


class CoreTest:
    def setup_class(self):
        self.net = OpenPNM.Network.Cubic(shape=[3, 3, 3])
        self.geo = OpenPNM.Geometry.GenericGeometry(network=self.net,
                                                    pores=self.net.Ps,
                                                    throats=self.net.Ts)
        self.geo['pore.diameter'] = sp.rand(self.net.Np)
        self.geo.models.add(propname='pore.volume',
                            model=OpenPNM.Geometry.models.pore_volume.sphere)
        self.geo['throat.diameter'] = sp.rand(self.net.Nt)
        self.geo.models.add(propname='throat.area',
                            model=OpenPNM.Geometry.models.throat_area.cylinder)
        self.geo['throat.label1'] = False
        self.geo['throat.label2'] = False
        self.geo['throat.label1'][0:6] = True
        self.geo['throat.label2'][3:9] = True
        self.net1 = OpenPNM.Network.Cubic(shape=[3, 3, 3])
        self.geo1 = OpenPNM.Geometry.GenericGeometry(network=self.net1,
                                                     pores=self.net1.Ps,
                                                     throats=self.net1.Ts)
        self.phase1 = OpenPNM.Phases.GenericPhase(network=self.net1)
        self.phase2 = OpenPNM.Phases.GenericPhase(network=self.net1)
        self.phys1 = OpenPNM.Physics.GenericPhysics(network=self.net1,
                                                    geometry=self.geo1,
                                                    phase=self.phase1)
        self.phys2 = OpenPNM.Physics.GenericPhysics(network=self.net1,
                                                    geometry=self.geo1,
                                                    phase=self.phase2)
        self.net2 = OpenPNM.Network.Cubic(shape=[3, 3, 3])
        Ps = sp.arange(0, 18)
        Ts = self.net2.find_neighbor_pores(Ps, mode='union')
        self.geo21 = OpenPNM.Geometry.GenericGeometry(network=self.net2,
                                                      pores=Ps,
                                                      throats=Ts)
        Ps = sp.arange(18, 27)
        Ts = self.net2.find_neighbor_pores(Ps, mode='intersection')
        self.geo22 = OpenPNM.Geometry.GenericGeometry(network=self.net2,
                                                      pores=Ps,
                                                      throats=Ts)

    def test_clear_complete(self):
        net = OpenPNM.Network.Cubic(shape=[3, 3, 3])
        b = sorted(list(net.keys()))
        dict_ = net.copy()
        net.clear()
        assert net.Np == 0
        assert net.Nt == 0
        a = sorted(list(net.keys()))
        assert a == ['pore.all', 'throat.all']
        net.update(dict_)
        assert net.Np == 27
        assert net.Nt == 54
        a = sorted(list(net.keys()))
        assert a == b

    def test_clear_props(self):
        net = OpenPNM.Network.Cubic(shape=[3, 3, 3])
        net.clear(mode='props')
        assert len(net.props()) == 0

    def test_clear_labels(self):
        net = OpenPNM.Network.Cubic(shape=[3, 3, 3])
        net.clear(mode='labels')
        assert len(net.labels()) == 2
        assert net.Np == 27
        assert net.Nt == 54

    def test_clear_labels_and_props(self):
        net = OpenPNM.Network.Cubic(shape=[3, 3, 3])
        net.clear(mode=['labels', 'props'])
        assert len(net.labels()) == 2
        assert net.Np == 27
        assert net.Nt == 54
        assert len(net.props()) == 0

    def test_clear_models(self):
        net = OpenPNM.Network.Cubic(shape=[3, 3, 3])
        geo = OpenPNM.Geometry.Stick_and_Ball(network=net,
                                              pores=net.Ps,
                                              throats=net.Ts)
        geo.clear(mode='props')
        assert len(geo.props()) == 0
        geo['pore.seed'] = sp.rand(geo.Np)
        geo.models.regenerate()
        assert len(geo.props()) > 0
        geo.clear(mode=['props', 'models'])
        geo.models.regenerate()
        assert len(geo.props()) == 0

    def test_props_all(self):
        a = self.geo.props()
        assert sorted(a) == ['pore.diameter', 'pore.volume',
                             'throat.area', 'throat.diameter']

    def test_props_models(self):
        a = self.geo.props(mode='models')
        b = ['pore.volume', 'throat.area']
        assert sorted(a) == sorted(b)

    def test_props_constants(self):
        a = self.geo.props(mode='constants')
        b = ['pore.diameter', 'throat.diameter']
        assert sorted(a) == sorted(b)

    def test_props_pores_all(self):
        a = self.geo.props(element='pores')
        b = ['pore.diameter', 'pore.volume']
        assert sorted(a) == sorted(b)

    def test_props_pores_models(self):
        a = self.geo.props(element='pores', mode='models')
        b = ['pore.volume']
        assert sorted(a) == sorted(b)

    def test_props_pores_constants(self):
        a = self.geo.props(element='pores', mode='constants')
        b = ['pore.diameter']
        assert sorted(a) == sorted(b)

    def test_props_hidden_keys(self):
        self.net['pore._blah'] = 1.0
        assert 'pore._blah' not in self.net.__str__()
        assert 'pore._blah' in self.net.keys()

    def test_labels(self):
        a = self.net.labels()