import copy
import unittest

import networkx as nx
import numpy as np
from scipy.special import erf

from dfn import Fluid, FractureNetworkThermal


class TestFractureNetworkThermal(unittest.TestCase):

    def __init__(self, *args, **kwargs):
        super(TestFractureNetworkThermal, self).__init__(*args, **kwargs)

        # fluid properties
        cp_w = 4300.0
        rho_w = 1000.0
        mu_w = 1E-3
        self.fluid = Fluid(density=rho_w, viscosity=mu_w, heat_capacity=cp_w)

        # reservoir properties
        k_r = 2.9
        cp_r = 1050.0
        rho_r = 2700.0
        alpha_r = k_r / (rho_r * cp_r)

        # first network
        conn_1 = [(0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)]
        L_1 = [100, 500, 500, 500, 500, 100]
        H_1 = [500, 500, 500, 500, 500, 500]
        w_1 = [1E-3, 1E-3, 1E-3, 1E-3, 1E-3, 1E-3]
        self.network_1 = FractureNetworkThermal(conn_1, L_1, H_1, w_1, k_r,
                                                alpha_r)

        # second network
        conn_2 = [(0, 1), (1, 2), (2, 3), (1, 4), (2, 5), (3, 6), (4, 5),
                  (5, 6), (4, 7), (5, 8), (6, 9), (7, 8), (8, 9), (9, 10)]
        L_2 = 250 * np.ones(len(conn_2))
        L_2[0] = 100
        L_2[-1] = 100
        H_2 = 500 * np.ones(len(conn_2))
        w_2 = 1E-3 * np.ones(len(conn_2))
        self.network_2 = FractureNetworkThermal(conn_2, L_2, H_2, w_2, k_r,
                                                alpha_r)

    def copy_networks(self):
        """Return a copy of the fracture networks."""
        return copy.copy(self.network_1), copy.copy(self.network_2)

    def networks_with_flow(self):
        """Return networks with the mass flow calculated."""

        network_1, network_2 = self.copy_networks()

        P_0 = 0.0
        m_inj = 50.0

        network_1.calculate_flow(self.fluid, {0: P_0}, {5: -m_inj})
        network_2.calculate_flow(self.fluid, {0: P_0}, {10: -m_inj})

        return network_1, network_2

    def reverse_nodes(self, network, segments):
        """Reverse the node order for given segments."""

        conn = network.connectivity

        for seg in segments:
            inlet, outlet = conn[seg]
            conn[seg, :] = outlet, inlet

        network.connectivity = conn
        return network

    def test_no_mass_flow(self):
        """Test if TypeError is raised for networks without flow calculated."""

        with self.assertRaises(TypeError):
            self.network_1._check_if_calculated()

        with self.assertRaises(TypeError):
            self.network_2._check_if_calculated()

    def test_neg_mass_flow(self):
        """Test if valueError is raised for networks with negative flow."""

        network_1, network_2 = self.networks_with_flow()

        network_1 = self.reverse_nodes(network_1, [1])
        network_2 = self.reverse_nodes(network_2, [1])

        network_1.calculate_flow(self.fluid, {0: 0}, {5: -1.0})
        network_2.calculate_flow(self.fluid, {0: 0}, {10: -1.0})

        with self.assertRaises(ValueError):
            network_1.calculate_temperature(self.fluid, 0, [0], [1])

        with self.assertRaises(ValueError):
            network_2.calculate_temperature(self.fluid, 0, [0], [1])

    def test_construct_graph(self):
        """Test _construct_graph method."""

        network_1, network_2 = self.networks_with_flow()

        network_1._construct_graph()
        network_2._construct_graph()

        # construct graph for network 1
        G_1 = nx.MultiDiGraph()
        edge_data_1 = [(0, 1, {'index': 0}), (1, 2, {'index': 1}),
                       (1, 3, {'index': 2}), (2, 4, {'index': 3}),
                       (3, 4, {'index': 4}), (4, 5, {'index': 5})]
        G_1.add_edges_from(edge_data_1)

        # construct graph for network 2
        G_2 = nx.MultiDiGraph()
        edge_data_2 = [(0, 1, {'index': 0}), (1, 2, {'index': 1}),
                       (2, 3, {'index': 2}), (1, 4, {'index': 3}),
                       (2, 5, {'index': 4}), (3, 6, {'index': 5}),
                       (4, 5, {'index': 6}), (5, 6, {'index': 7}),
                       (4, 7, {'index': 8}), (5, 8, {'index': 9}),
                       (6, 9, {'index': 10}), (7, 8, {'index': 11}),
                       (8, 9, {'index': 12}), (9, 10, {'index': 13})]
        G_2.add_edges_from(edge_data_2)

        # return True if graphs are the same
        is_isomorphic_1 = nx.is_isomorphic(network_1.graph, G_1)
        is_isomorphic_2 = nx.is_isomorphic(network_2.graph, G_2)

        self.assertTrue(is_isomorphic_1)
        self.assertTrue(is_isomorphic_2)

    def test_find_injection_nodes(self):
        """Test _find_injection_nodes method."""

        network_1, network_2 = self.networks_with_flow()
        network_1._construct_graph()
        network_2._construct_graph()

        self.assertEqual(network_1._find_injection_nodes(), [0])
        self.assertEqual(network_2._find_injection_nodes(), [0])

    def test_mass_contribution(self):
        """Test _mass_contribution method."""

        network_1, network_2 = self.networks_with_flow()

        chi_1 = network_1._mass_contribution()
        chi_2 = network_2._mass_contribution()

        # first network
        for i in (0, 1, 2, 5):
            self.assertAlmostEqual(chi_1[i], 1.0, 12)

        self.assertAlmostEqual(chi_1[3] + chi_1[4], 1.0, 12)

        # second network
        for i in (0, 1, 2, 3, 8, 13):
            self.assertAlmostEqual(chi_2[i], 1.0, 12)

        for i, j in [(4, 6), (5, 7), (9, 11), (10, 12)]:
            self.assertAlmostEqual(chi_2[i] + chi_2[j], 1.0, 12)

    def test_find_paths(self):
        """Test find_paths method."""

        # .find_paths method calls .construct_graph if needed. Manually call
        # .construct_graph() on one network for testing both True and False
        # conditions
        network_1, network_2 = self.networks_with_flow()
        network_1._construct_graph()

        path_1 = {(0, 1, 3), (0, 2, 4)}
        path_2 = {(0, 1, 2, 5, 10), (0, 1, 4, 7, 10), (0, 3, 6, 7, 10),
                  (0, 3, 6, 9, 12), (0, 3, 8, 11, 12), (0, 1, 4, 9, 12)}

        self.assertEqual(path_1, set(network_1.find_paths(0, 4)))
        self.assertEqual(path_2, set(network_2.find_paths(0, 9)))

    def test_calculate_temperature_inlet_segment(self):
        """Test calculate_temperature ability to handle the inlet segment."""

        # operational parameters for temperature
        t_end = 86400 * 365.25 * 20
        time = t_end * np.linspace(1.0 / 100, 1.0, 100)

        distance = np.linspace(0.0, 100.0, 100)
        z, t = np.meshgrid(distance, time)

        network_1, network_2 = self.networks_with_flow()

        # create parameters for temperature manually
        m_1 = network_1.mass_flow[0]
        m_2 = network_2.mass_flow[0]
        beta_1 = 2 * network_1.thermal_cond * network_1.thickness[0] / \
            (m_1 * network_1.fluid.c_f)
        beta_2 = 2 * network_2.thermal_cond * network_2.thickness[0] / \
            (m_2 * network_2.fluid.c_f)

        xi_1 = beta_1 * z / (2 * np.sqrt(network_1.thermal_diff * t))
        xi_2 = beta_2 * z / (2 * np.sqrt(network_2.thermal_diff * t))

        Theta_1 = erf(xi_1)
        Theta_2 = erf(xi_2)

        # difference between manual and automatic construction
        diff_1 = Theta_1 - network_1.calculate_temperature(self.fluid, 0,
                                                           distance, time)
        diff_2 = Theta_2 - network_2.calculate_temperature(self.fluid, 0,
                                                           distance, time)

        self.assertAlmostEqual((diff_1**2).sum() / (Theta_1**2).sum(), 0, 12)
        self.assertAlmostEqual((diff_2**2).sum() / (Theta_2**2).sum(), 0, 12)

    def test_calculate_temperature(self):
        """Test calculate_temperature by constructing manual the equations."""

        # operational parameters for temperature
        t_end = 86400 * 365.25 * 20
        time = t_end * np.linspace(1.0 / 100, 1.0, 100)

        distance = np.linspace(0.0, 100.0, 100)
        z, t = np.meshgrid(distance, time)

        network_1, network_2 = self.networks_with_flow()

        # create parameters for temperature manually
        chi_1 = np.array([1.0, 1.0, 1.0, 0.5, 0.5, 1.0])
        chi_2 = np.ones(network_2.n_segments)
        chi_2[4:8] = 0.5
        chi_2[9:13] = 0.5

        m_1 = network_1.mass_flow
        m_2 = network_2.mass_flow
        beta_1 = 2 * network_1.thermal_cond * network_1.thickness / \
            (m_1 * network_1.fluid.c_f)
        beta_2 = 2 * network_2.thermal_cond * network_2.thickness / \
            (m_2 * network_2.fluid.c_f)

        xi_1 = np.einsum('i,jk->ijk', beta_1 * network_1.length,
                         1 / (2 * np.sqrt(network_1.thermal_diff * t)))
        xi_2 = np.einsum('i,jk->ijk', beta_2 * network_2.length,
                         1 / (2 * np.sqrt(network_2.thermal_diff * t)))
        a = xi_1[[0, 2, 4], :, :].sum(axis=0)
        b = xi_1[[0, 1, 3], :, :].sum(axis=0)
        xi_seg = beta_1[-1] * z / (2 * np.sqrt(network_1.thermal_diff * t))

        Theta_1 = chi_1[0] * chi_1[2] * chi_1[4] * erf(a + xi_seg) + \
            chi_1[0] * chi_1[1] * chi_1[3] * erf(b + xi_seg)

        a = xi_2[[0, 1, 2, 5, 10], :, :].sum(axis=0)
        b = xi_2[[0, 1, 4, 7, 10], :, :].sum(axis=0)
        c = xi_2[[0, 3, 6, 7, 10], :, :].sum(axis=0)
        d = xi_2[[0, 3, 6, 9, 12], :, :].sum(axis=0)
        e = xi_2[[0, 3, 8, 11, 12], :, :].sum(axis=0)
        f = xi_2[[0, 1, 4, 9, 12], :, :].sum(axis=0)

        C_1 = chi_2[0] * chi_2[1] * chi_2[2] * chi_2[5] * chi_2[10]
        C_2 = chi_2[0] * chi_2[1] * chi_2[4] * chi_2[7] * chi_2[10]
        C_3 = chi_2[0] * chi_2[3] * chi_2[6] * chi_2[7] * chi_2[10]
        C_4 = chi_2[0] * chi_2[3] * chi_2[6] * chi_2[9] * chi_2[12]
        C_5 = chi_2[0] * chi_2[3] * chi_2[8] * chi_2[11] * chi_2[12]
        C_6 = chi_2[0] * chi_2[1] * chi_2[4] * chi_2[9] * chi_2[12]
        xi_seg = beta_2[-1] * z / (2 * np.sqrt(network_2.thermal_diff * t))

        Theta_2 =  C_1 * erf(a + xi_seg) + C_2 * erf(b + xi_seg) + \
            C_3 * erf(c + xi_seg) + C_4 * erf(d + xi_seg) + \
            C_5 * erf(e + xi_seg) + C_6 * erf(f + xi_seg)

        # difference between manual and automatic construction
        diff_1 = Theta_1 - network_1.calculate_temperature(self.fluid, 5,
                                                           distance, time)
        diff_2 = Theta_2 - network_2.calculate_temperature(self.fluid, 13,
                                                           distance, time)

        self.assertAlmostEqual((diff_1**2).sum() / (Theta_1**2).sum(), 0, 12)
        self.assertAlmostEqual((diff_2**2).sum() / (Theta_2**2).sum(), 0, 12)

if __name__ == '__main__':
    unittest.main()

[((282, 4, 282, 19), 'unittest.main', 'unittest.main', ({}, {}), '()', False, 'import unittest\n'), ((20, 21, 20, 77), 'dfn.Fluid', 'Fluid', (), '', False, 'from dfn import Fluid, FractureNetworkThermal\n'), ((33, 25, 34, 56), 'dfn.FractureNetworkThermal', 'FractureNetworkThermal', ({(33, 48, 33, 54): 'conn_1', (33, 56, 33, 59): 'L_1', (33, 61, 33, 64): 'H_1', (33, 66, 33, 69): 'w_1', (33, 71, 33, 74): 'k_r', (34, 48, 34, 55): 'alpha_r'}, {}), '(conn_1, L_1, H_1, w_1, k_r, alpha_r)', False, 'from dfn import Fluid, FractureNetworkThermal\n'), ((44, 25, 45, 56), 'dfn.FractureNetworkThermal', 'FractureNetworkThermal', ({(44, 48, 44, 54): 'conn_2', (44, 56, 44, 59): 'L_2', (44, 61, 44, 64): 'H_2', (44, 66, 44, 69): 'w_2', (44, 71, 44, 74): 'k_r', (45, 48, 45, 55): 'alpha_r'}, {}), '(conn_2, L_2, H_2, w_2, k_r, alpha_r)', False, 'from dfn import Fluid, FractureNetworkThermal\n'), ((111, 14, 111, 31), 'networkx.MultiDiGraph', 'nx.MultiDiGraph', ({}, {}), '()', True, 'import networkx as nx\n'), ((118, 14, 118, 31), 'networkx.MultiDiGraph', 'nx.MultiDiGraph', ({}, {}), '()', True, 'import networkx as nx\n'), ((129, 26, 129, 64), 'networkx.is_isomorphic', 'nx.is_isomorphic', ({(129, 43, 129, 58): 'network_1.graph', (129, 60, 129, 63): 'G_1'}, {}), '(network_1.graph, G_1)', True, 'import networkx as nx\n'), ((130, 26, 130, 64), 'networkx.is_isomorphic', 'nx.is_isomorphic', ({(130, 43, 130, 58): 'network_2.graph', (130, 60, 130, 63): 'G_2'}, {}), '(network_2.graph, G_2)', True, 'import networkx as nx\n'), ((189, 19, 189, 47), 'numpy.linspace', 'np.linspace', ({(189, 31, 189, 34): '0.0', (189, 36, 189, 41): '100.0', (189, 43, 189, 46): '100'}, {}), '(0.0, 100.0, 100)', True, 'import numpy as np\n'), ((190, 15, 190, 42), 'numpy.meshgrid', 'np.meshgrid', ({(190, 27, 190, 35): 'distance', (190, 37, 190, 41): 'time'}, {}), '(distance, time)', True, 'import numpy as np\n'), ((205, 18, 205, 27), 'scipy.special.erf', 'erf', ({(205, 22, 205, 26): 'xi_1'}, {}), '(xi_1)', False, 'from scipy.special import erf\n'), ((206, 18, 206, 27), 'scipy.special.erf', 'erf', ({(206, 22, 206, 26): 'xi_2'}, {}), '(xi_2)', False, 'from scipy.special import erf\n'), ((224, 19, 224, 47), 'numpy.linspace', 'np.linspace', ({(224, 31, 224, 34): '0.0', (224, 36, 224, 41): '100.0', (224, 43, 224, 46): '100'}, {}), '(0.0, 100.0, 100)', True, 'import numpy as np\n'), ((225, 15, 225, 42), 'numpy.meshgrid', 'np.meshgrid', ({(225, 27, 225, 35): 'distance', (225, 37, 225, 41): 'time'}, {}), '(distance, time)', True, 'import numpy as np\n'), ((230, 16, 230, 56), 'numpy.array', 'np.array', ({(230, 25, 230, 55): '[1.0, 1.0, 1.0, 0.5, 0.5, 1.0]'}, {}), '([1.0, 1.0, 1.0, 0.5, 0.5, 1.0])', True, 'import numpy as np\n'), ((231, 16, 231, 45), 'numpy.ones', 'np.ones', ({(231, 24, 231, 44): 'network_2.n_segments'}, {}), '(network_2.n_segments)', True, 'import numpy as np\n'), ((49, 15, 49, 40), 'copy.copy', 'copy.copy', ({(49, 25, 49, 39): 'self.network_1'}, {}), '(self.network_1)', False, 'import copy\n'), ((49, 42, 49, 67), 'copy.copy', 'copy.copy', ({(49, 52, 49, 66): 'self.network_2'}, {}), '(self.network_2)', False, 'import copy\n'), ((187, 23, 187, 55), 'numpy.linspace', 'np.linspace', ({(187, 35, 187, 44): '(1.0 / 100)', (187, 46, 187, 49): '(1.0)', (187, 51, 187, 54): '(100)'}, {}), '(1.0 / 100, 1.0, 100)', True, 'import numpy as np\n'), ((222, 23, 222, 55), 'numpy.linspace', 'np.linspace', ({(222, 35, 222, 44): '(1.0 / 100)', (222, 46, 222, 49): '(1.0)', (222, 51, 222, 54): '(100)'}, {}), '(1.0 / 100, 1.0, 100)', True, 'import numpy as np\n'), ((202, 33, 202, 68), 'numpy.sqrt', 'np.sqrt', ({(202, 41, 202, 67): '(network_1.thermal_diff * t)'}, {}), '(network_1.thermal_diff * t)', True, 'import numpy as np\n'), ((203, 33, 203, 68), 'numpy.sqrt', 'np.sqrt', ({(203, 41, 203, 67): '(network_2.thermal_diff * t)'}, {}), '(network_2.thermal_diff * t)', True, 'import numpy as np\n'), ((248, 39, 248, 74), 'numpy.sqrt', 'np.sqrt', ({(248, 47, 248, 73): '(network_1.thermal_diff * t)'}, {}), '(network_1.thermal_diff * t)', True, 'import numpy as np\n'), ((250, 51, 250, 66), 'scipy.special.erf', 'erf', ({(250, 55, 250, 65): '(a + xi_seg)'}, {}), '(a + xi_seg)', False, 'from scipy.special import erf\n'), ((251, 45, 251, 60), 'scipy.special.erf', 'erf', ({(251, 49, 251, 59): '(b + xi_seg)'}, {}), '(b + xi_seg)', False, 'from scipy.special import erf\n'), ((266, 39, 266, 74), 'numpy.sqrt', 'np.sqrt', ({(266, 47, 266, 73): '(network_2.thermal_diff * t)'}, {}), '(network_2.thermal_diff * t)', True, 'import numpy as np\n'), ((270, 42, 270, 57), 'scipy.special.erf', 'erf', ({(270, 46, 270, 56): '(f + xi_seg)'}, {}), '(f + xi_seg)', False, 'from scipy.special import erf\n'), ((243, 34, 243, 69), 'numpy.sqrt', 'np.sqrt', ({(243, 42, 243, 68): 'network_1.thermal_diff * t'}, {}), '(network_1.thermal_diff * t)', True, 'import numpy as np\n'), ((245, 34, 245, 69), 'numpy.sqrt', 'np.sqrt', ({(245, 42, 245, 68): 'network_2.thermal_diff * t'}, {}), '(network_2.thermal_diff * t)', True, 'import numpy as np\n'), ((270, 18, 270, 33), 'scipy.special.erf', 'erf', ({(270, 22, 270, 32): '(e + xi_seg)'}, {}), '(e + xi_seg)', False, 'from scipy.special import erf\n'), ((269, 42, 269, 57), 'scipy.special.erf', 'erf', ({(269, 46, 269, 56): '(d + xi_seg)'}, {}), '(d + xi_seg)', False, 'from scipy.special import erf\n'), ((269, 18, 269, 33), 'scipy.special.erf', 'erf', ({(269, 22, 269, 32): '(c + xi_seg)'}, {}), '(c + xi_seg)', False, 'from scipy.special import erf\n'), ((268, 25, 268, 40), 'scipy.special.erf', 'erf', ({(268, 29, 268, 39): '(a + xi_seg)'}, {}), '(a + xi_seg)', False, 'from scipy.special import erf\n'), ((268, 49, 268, 64), 'scipy.special.erf', 'erf', ({(268, 53, 268, 63): '(b + xi_seg)'}, {}), '(b + xi_seg)', False, 'from scipy.special import erf\n')]