import OpenPNM
import numpy as np
import OpenPNM.Physics.models as pm


class GenericLinearTransportTest:
    def setup_class(self):
        self.net = OpenPNM.Network.Cubic(shape=[5, 5, 5])
        self.phase = OpenPNM.Phases.GenericPhase(network=self.net)
        Ps = self.net.Ps
        Ts = self.net.Ts
        self.phys = OpenPNM.Physics.GenericPhysics(network=self.net,
                                                   phase=self.phase,
                                                   pores=Ps, throats=Ts)
        self.phys['throat.cond'] = 5e-8
        self.alg = OpenPNM.Algorithms.GenericLinearTransport(network=self.net,
                                                             phase=self.phase)

    def test_set_BC_modes_pores(self):
        BC1_pores = np.arange(25, 35)
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         bcvalue=0.8,
                                         pores=BC1_pores)
        ptest = self.alg.pores('pore.Dirichlet')
        assert np.all(ptest == BC1_pores)
        BC2_pores = np.arange(43, 50)
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         bcvalue=0.8,
                                         pores=BC2_pores,
                                         mode='merge')
        ptest = self.alg.pores('pore.Dirichlet')
        assert np.all(ptest == np.concatenate((BC1_pores, BC2_pores)))
        BC3_pores = np.arange(4, 9)
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         bcvalue=0.8,
                                         pores=BC3_pores,
                                         mode='overwrite')
        ptest = self.alg.pores('pore.Dirichlet')
        assert np.all(ptest == BC3_pores)
        BC4_pores = [11, 90]
        self.alg.set_boundary_conditions(bctype='Neumann',
                                         bcvalue=0.5,
                                         pores=BC4_pores,
                                         mode='overwrite')
        ptest = self.alg.pores('pore.Neumann')
        assert np.all(ptest == BC4_pores)
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         pores=BC1_pores,
                                         bcvalue=0.3)
        ptest = self.alg.pores('pore.Dirichlet')
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         pores=self.alg.Ps,
                                         mode='remove')
        Dp = np.sum(self.alg['pore.Dirichlet'])
        assert Dp == 0
        self.alg.set_boundary_conditions(bctype='Neumann',
                                         mode='remove')
        label = 'pore.Neumann'
        assert (label not in self.alg.labels())

    def test_set_BC_modes_throats(self):
        BC1_throats = np.arange(25, 35)
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         bcvalue=0.8,
                                         throats=BC1_throats)
        t_test = self.alg.throats('throat.Dirichlet')
        assert np.all(t_test == BC1_throats)
        BC2_throats = np.arange(43, 50)
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         bcvalue=0.8,
                                         throats=BC2_throats,
                                         mode='merge')
        t_test = self.alg.throats('throat.Dirichlet')
        assert np.all(t_test == np.concatenate((BC1_throats, BC2_throats)))
        BC3_throats = np.arange(4, 9)
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         bcvalue=0.8,
                                         throats=BC3_throats,
                                         mode='overwrite')
        t_test = self.alg.throats('throat.Dirichlet')
        assert np.all(t_test == BC3_throats)
        BC4_throats = [11, 90]
        self.alg.set_boundary_conditions(bctype='Neumann',
                                         bcvalue=0.5,
                                         throats=BC4_throats,
                                         mode='overwrite')
        t_test = self.alg.throats('throat.Neumann')
        assert np.all(t_test == BC4_throats)
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         throats=BC1_throats,
                                         bcvalue=0.3)
        t_test = self.alg.throats('throat.Dirichlet')
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         throats=self.alg.Ts,
                                         mode='remove')
        Dp = np.sum(self.alg['throat.Dirichlet'])
        assert Dp == 0
        self.alg.set_boundary_conditions(bctype='Neumann',
                                         mode='remove')
        label = 'throat.Neumann'
        assert (label not in self.alg.labels())

    def test_set_BC_modes_with_boolean_masks_pores(self):
        BC1_pores = np.zeros(self.alg.Np, dtype='bool')
        BC1_pores[np.arange(25, 35)] = True
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         bcvalue=0.8,
                                         pores=BC1_pores)
        ptest = self.alg.pores('pore.Dirichlet')
        assert np.all(ptest == self.alg._parse_locations(BC1_pores))
        BC2_pores = np.zeros(self.alg.Np, dtype='bool')
        BC2_pores[np.arange(43, 50)] = True
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         bcvalue=0.8,
                                         pores=BC2_pores,
                                         mode='merge')
        ptest = self.alg.pores('pore.Dirichlet')
        B1 = self.alg._parse_locations(BC1_pores)
        B2 = self.alg._parse_locations(BC2_pores)
        assert np.all(ptest == np.concatenate((B1, B2)))
        BC3_pores = np.zeros(self.alg.Np, dtype='bool')
        BC3_pores[np.arange(4, 9)] = True
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         bcvalue=0.8,
                                         pores=BC3_pores,
                                         mode='overwrite')
        ptest = self.alg.pores('pore.Dirichlet')
        assert np.all(ptest == self.alg._parse_locations(BC3_pores))
        BC4_pores = np.zeros(self.alg.Np, dtype='bool')
        BC4_pores[[11, 90]] = True
        self.alg.set_boundary_conditions(bctype='Neumann',
                                         bcvalue=0.5,
                                         pores=BC4_pores,
                                         mode='overwrite')
        ptest = self.alg.pores('pore.Neumann')
        assert np.all(ptest == self.alg._parse_locations(BC4_pores))
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         pores=BC1_pores,
                                         bcvalue=0.3)
        ptest = self.alg.pores('pore.Dirichlet')
        removed_p = self.alg._parse_locations(self.alg.Ps)
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         pores=removed_p,
                                         mode='remove')
        Dp = np.sum(self.alg['pore.Dirichlet'])
        assert Dp == 0
        self.alg.set_boundary_conditions(bctype='Neumann',
                                         mode='remove')
        label = 'pore.Neumann'
        assert (label not in self.alg.labels())

    def test_set_BC_modes_with_boolean_masks_throats(self):
        BC1_throats = np.zeros(self.alg.Nt, dtype='bool')
        BC1_throats[np.arange(25, 35)] = True
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         bcvalue=0.8,
                                         throats=BC1_throats)
        t_test = self.alg.throats('throat.Dirichlet')
        assert np.all(t_test == self.alg._parse_locations(BC1_throats))
        BC2_throats = np.zeros(self.alg.Nt, dtype='bool')
        BC2_throats[np.arange(43, 50)] = True
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         bcvalue=0.8,
                                         throats=BC2_throats,
                                         mode='merge')
        t_test = self.alg.throats('throat.Dirichlet')
        B1 = self.alg._parse_locations(BC1_throats)
        B2 = self.alg._parse_locations(BC2_throats)
        assert np.all(t_test == np.concatenate((B1, B2)))
        BC3_throats = np.zeros(self.alg.Nt, dtype='bool')
        BC3_throats[np.arange(4, 9)] = True
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         bcvalue=0.8,
                                         throats=BC3_throats,
                                         mode='overwrite')
        t_test = self.alg.throats('throat.Dirichlet')
        assert np.all(t_test == self.alg._parse_locations(BC3_throats))
        BC4_throats = np.zeros(self.alg.Nt, dtype='bool')
        BC4_throats[[11, 90]] = True
        self.alg.set_boundary_conditions(bctype='Neumann',
                                         bcvalue=0.5,
                                         throats=BC4_throats,
                                         mode='overwrite')
        t_test = self.alg.throats('throat.Neumann')
        assert np.all(t_test == self.alg._parse_locations(BC4_throats))
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         throats=BC1_throats,
                                         bcvalue=0.3)
        t_test = self.alg.throats('throat.Dirichlet')
        removed_t = self.alg._parse_locations(self.alg.Ts)
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         throats=removed_t,
                                         mode='remove')
        Dp = np.sum(self.alg['throat.Dirichlet'])
        assert Dp == 0
        self.alg.set_boundary_conditions(bctype='Neumann',
                                         mode='remove')
        label = 'pore.Neumann'
        assert (label not in self.alg.labels())

    def test_super_pore_conductance(self):
        g_super = []
        BC1_pores = np.arange(20, 30)
        self.alg.set_boundary_conditions(bctype='Dirichlet',
                                         bcvalue=0.4,
                                         pores=BC1_pores)
        BC2_pores = np.arange(45, 66)
        self.alg.set_boundary_conditions(bctype='Neumann_group',
                                         bcvalue=1.4e-10,
                                         pores=BC2_pores)
        g_super.append(2e-12)
        BC3_pores = np.arange(87, 94)
        self.alg.set_boundary_conditions(bctype='Neumann_group',
                                         bcvalue=-0.9e-10,
                                         pores=BC3_pores)
        g_super.append(np.ones(len(BC3_pores)) * 1.5e-12)
        BC4_pores = np.arange(3, 7)
        self.alg.set_boundary_conditions(bctype='Neumann_group',
                                         bcvalue=0.1e-10,
                                         pores=BC4_pores)
        g_super.append(np.array([6.42e-13]))
        self.alg.run(conductance='throat.cond',
                     quantity='pore.mole_fraction',
                     super_pore_conductance=g_super)
        self.alg.return_results()
        r1 = self.alg.rate(BC1_pores)[0]
        r2 = self.alg.rate(BC2_pores)[0]
        r3 = self.alg.rate(BC3_pores)[0]
        r4 = self.alg.rate(BC4_pores)[0]
        assert np.absolute(r1 + r2 + r3 + r4) < 1e-20
        assert np.size(self.alg.super_pore_conductance[0]) == 1
        assert np.size(self.alg.super_pore_conductance[1]) == 7
        assert np.size(self.alg.super_pore_conductance[2]) == 1

    def test_source_term_modes(self):
        self.phys['pore.item1'] = 0.5e-12
        self.phys['pore.item2'] = 2.5
        self.phys['pore.item3'] = -1.4e-11
        self.phys.models.add(propname='pore.A',
                             model=pm.generic_source_term.power_law,
                             A1='pore.item1',
                             A2='pore.item2',
                             A3='pore.item3',
                             x='mole_fraction',
                             return_rate=False,
                             regen_mode='on_demand')
        self.phys.models.add(propname='pore.B',
                             model=pm.generic_source_term.linear,
                             A1='pore.item1',
                             A2='pore.item3',
                             x='mole_fraction',
                             return_rate=False,
                             regen_mode='on_demand')
        S1_pores = np.arange(25, 35)
        self.alg.set_source_term(source_name=['pore.A', 'pore.B'],
                                 pores=S1_pores)
        mask1 = ~np.isnan(self.alg['pore.source_nonlinear_s1_A'])
        mask2 = ~np.isnan(self.alg['pore.source_nonlinear_s2_A'])
        assert np.all(self.alg.Ps[mask1] == S1_pores)
        assert np.all(self.alg.Ps[mask2] == S1_pores)
        self.alg.set_source_term(source_name='pore.A',
                                 pores=[26], x0=np.ones(self.phys.Np),
                                 mode='update')
        assert self.alg['pore.source_nonlinear_s1_A'][26] == 1.25e-12
        S2_pores = np.array([30, 31])
        self.alg.set_source_term(source_name='pore.A',
                                 pores=S2_pores,
                                 mode='overwrite')
        mask1 = ~np.isnan(self.alg['pore.source_nonlinear_s1_A'])
        assert np.all(self.alg.Ps[mask1] == S2_pores)
        self.alg.set_source_term(source_name='pore.B',
                                 pores=S1_pores,
                                 mode='remove')
        mask1 = np.isnan(self.alg['pore.source_nonlinear_s1_B'])
        assert np.all(self.alg.Ps[mask1] == self.alg.Ps)
        self.alg.set_source_term(source_name=['pore.A', 'pore.B'],
                                 pores=self.alg.Ps,
                                 mode='remove')
        assert ('pore.source_B' in self.alg.labels())
        assert ('pore.source_A' in self.alg.labels())
        self.alg.set_source_term(source_name=['pore.A', 'pore.B'],
                                 mode='remove')
        assert ('pore.source_B' not in self.alg.labels())
        assert ('pore.source_A' not in self.alg.labels())

[(8, 'OpenPNM.Network.Cubic', 'OpenPNM.Network.Cubic', False, 'import OpenPNM\n'), (9, 'OpenPNM.Phases.GenericPhase', 'OpenPNM.Phases.GenericPhase', False, 'import OpenPNM\n'), (12, 'OpenPNM.Physics.GenericPhysics', 'OpenPNM.Physics.GenericPhysics', False, 'import OpenPNM\n'), (16, 'OpenPNM.Algorithms.GenericLinearTransport', 'OpenPNM.Algorithms.GenericLinearTransport', False, 'import OpenPNM\n'), (20, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (25, 'numpy.all', 'np.all', True, 'import numpy as np\n'), (26, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (33, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (39, 'numpy.all', 'np.all', True, 'import numpy as np\n'), (46, 'numpy.all', 'np.all', True, 'import numpy as np\n'), (54, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (62, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (67, 'numpy.all', 'np.all', True, 'import numpy as np\n'), (68, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (75, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (81, 'numpy.all', 'np.all', True, 'import numpy as np\n'), (88, 'numpy.all', 'np.all', True, 'import numpy as np\n'), (96, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (104, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (111, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (121, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (129, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (145, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (153, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (160, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (170, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (178, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (194, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (203, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (207, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (212, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (217, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (254, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (259, 'numpy.all', 'np.all', True, 'import numpy as np\n'), (260, 'numpy.all', 'np.all', True, 'import numpy as np\n'), (265, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (270, 'numpy.all', 'np.all', True, 'import numpy as np\n'), (274, 'numpy.isnan', 'np.isnan', True, 'import numpy as np\n'), (275, 'numpy.all', 'np.all', True, 'import numpy as np\n'), (221, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (230, 'numpy.absolute', 'np.absolute', True, 'import numpy as np\n'), (231, 'numpy.size', 'np.size', True, 'import numpy as np\n'), (232, 'numpy.size', 'np.size', True, 'import numpy as np\n'), (233, 'numpy.size', 'np.size', True, 'import numpy as np\n'), (257, 'numpy.isnan', 'np.isnan', True, 'import numpy as np\n'), (258, 'numpy.isnan', 'np.isnan', True, 'import numpy as np\n'), (269, 'numpy.isnan', 'np.isnan', True, 'import numpy as np\n'), (32, 'numpy.concatenate', 'np.concatenate', True, 'import numpy as np\n'), (74, 'numpy.concatenate', 'np.concatenate', True, 'import numpy as np\n'), (105, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (112, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (120, 'numpy.concatenate', 'np.concatenate', True, 'import numpy as np\n'), (122, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (154, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (161, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (169, 'numpy.concatenate', 'np.concatenate', True, 'import numpy as np\n'), (171, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (262, 'numpy.ones', 'np.ones', True, 'import numpy as np\n')]