import numpy as np
import openpnm as op
from numpy.testing import assert_allclose
from openpnm.algorithms import AdvectionDiffusion, StokesFlow
import pytest


class AdvectionDiffusionTest:

    def setup_class(self):
        np.random.seed(0)
        self.net = op.network.Cubic(shape=[4, 3, 1], spacing=1.)
        self.geo = op.geometry.GenericGeometry(network=self.net,
                                               pores=self.net.Ps,
                                               throats=self.net.Ts)
        self.geo['throat.conduit_lengths.pore1'] = 0.1
        self.geo['throat.conduit_lengths.throat'] = 0.6
        self.geo['throat.conduit_lengths.pore2'] = 0.1

        self.phase = op.phase.GenericPhase(network=self.net)
        self.phys = op.physics.GenericPhysics(network=self.net,
                                              phase=self.phase,
                                              geometry=self.geo)
        self.phys['throat.diffusive_conductance'] = 1e-15
        self.phys['throat.hydraulic_conductance'] = 1e-15

        self.sf = StokesFlow(network=self.net, phase=self.phase)
        self.sf.set_value_BC(pores=self.net.pores('right'), values=1)
        self.sf.set_value_BC(pores=self.net.pores('left'), values=0)
        self.sf.run()

        self.phase.update(self.sf.results())

        self.ad = AdvectionDiffusion(network=self.net, phase=self.phase)
        self.ad.settings._update({"cache": False})
        self.ad.set_value_BC(pores=self.net.pores('right'), values=2)
        self.ad.set_value_BC(pores=self.net.pores('left'), values=0)

    def test_settings(self):
        self.ad.settings._update({'quantity': "pore.blah",
                                  'conductance': "throat.foo",
                                  'diffusive_conductance': "throat.foo2",
                                  'hydraulic_conductance': "throat.foo3",
                                  'pressure': "pore.bar",
                                  })
        assert self.ad.settings["quantity"] == "pore.blah"
        assert self.ad.settings["conductance"] == "throat.foo"
        assert self.ad.settings["diffusive_conductance"] == "throat.foo2"
        assert self.ad.settings["hydraulic_conductance"] == "throat.foo3"
        assert self.ad.settings["pressure"] == "pore.bar"
        # Reset back to previous values for other tests to run
        self.ad.settings._update({'quantity': "pore.concentration",
                                  'conductance': "throat.ad_dif_conductance",
                                  "diffusive_conductance": "throat.diffusive_conductance",
                                  "hydraulic_conductance": "throat.hydraulic_conductance",
                                  "pressure": "pore.pressure"
                                  })

    def test_conductance_gets_updated_when_pressure_changes(self):
        mod = op.models.physics.ad_dif_conductance.ad_dif
        self.phys.add_model(propname='throat.ad_dif_conductance',
                            model=mod, s_scheme='powerlaw')
        g_old = self.phys["throat.ad_dif_conductance"]
        # Run StokesFlow with a different BC to change pressure field
        self.sf.set_value_BC(pores=self.net.pores('right'), values=1.5)
        # Running the next line should update "throat.ad_dif_conductance"
        self.sf.run()
        self.phase.update(self.sf.results())
        self.ad.settings["conductance"] = "throat.ad_dif_conductance"
        self.ad.run()
        g_updated = self.phys["throat.ad_dif_conductance"]
        # Ensure conductance values are updated
        assert g_old.mean() != g_updated.mean()
        assert_allclose(g_updated.mean(), 1.01258990e-15)
        # Reset BCs for other tests to run properly
        self.sf.set_value_BC(pores=self.net.pores('right'), values=1)
        self.sf.run()

    def test_powerlaw_advection_diffusion(self):
        mod = op.models.physics.ad_dif_conductance.ad_dif
        self.phys.add_model(propname='throat.ad_dif_conductance_powerlaw',
                            model=mod, s_scheme='powerlaw')
        self.phys.regenerate_models()
        self.ad.settings['conductance'] = 'throat.ad_dif_conductance_powerlaw'
        self.ad.run()
        x = [
            0., 0., 0.,
            0.89653, 0.89653, 0.89653,
            1.53924, 1.53924, 1.53924,
            2., 2., 2.
        ]
        y = self.ad['pore.concentration']
        assert_allclose(actual=y, desired=x, rtol=1e-5)

    def test_upwind_advection_diffusion(self):
        mod = op.models.physics.ad_dif_conductance.ad_dif
        self.phys.add_model(propname='throat.ad_dif_conductance_upwind',
                            model=mod, s_scheme='upwind')
        self.phys.regenerate_models()
        self.ad.settings['conductance'] = 'throat.ad_dif_conductance_upwind'
        self.ad.run()
        x = [
            0., 0., 0.,
            0.86486, 0.86486, 0.86486,
            1.51351, 1.51351, 1.51351,
            2., 2., 2.
        ]
        y = self.ad['pore.concentration']
        assert_allclose(actual=y, desired=x, rtol=1e-5)

    def test_hybrid_advection_diffusion(self):
        mod = op.models.physics.ad_dif_conductance.ad_dif
        self.phys.add_model(propname='throat.ad_dif_conductance_hybrid',
                            model=mod, s_scheme='hybrid')
        self.phys.regenerate_models()

        self.ad.settings['conductance'] = 'throat.ad_dif_conductance_hybrid'
        self.ad.run()
        x = [
            0., 0., 0.,
            0.89908, 0.89908, 0.89908,
            1.54128, 1.54128, 1.54128,
            2., 2., 2.
        ]
        y = self.ad['pore.concentration']
        assert_allclose(actual=y, desired=x, rtol=1e-5)

    def test_exponential_advection_diffusion(self):
        mod = op.models.physics.ad_dif_conductance.ad_dif
        self.phys.add_model(propname='throat.ad_dif_conductance_exponential',
                            model=mod, s_scheme='exponential')
        self.phys.regenerate_models()

        self.ad.settings['conductance'] = 'throat.ad_dif_conductance_exponential'
        self.ad.run()
        x = [
            0., 0., 0.,
            0.89688173, 0.89688173, 0.89688173,
            1.53952557, 1.53952557, 1.53952557,
            2., 2., 2.
        ]
        y = self.ad['pore.concentration']