import unittest
from unittest.mock import MagicMock

from nose.tools import assert_almost_equal, eq_
from numpy import array
from numpy.testing import assert_array_equal

from helix.calculators.ballast_calculator import BallastCalculator
from helix.constants.anchor_type import AnchorType
from helix.constants.module_type import ModuleType
from helix.constants.panel_type import PanelType
from helix.constants.system_type import SystemType
from helix.models.panel import Panel, PanelWarnings
from test.test_helpers import assert_array_is_close


class BallastCalculatorWhenDualTiltAndPSeriesTest(unittest.TestCase):
    def setUp(self):
        self.values = MagicMock()
        """
        Generated using site parameters:
        height: 100
        length: 1000
        width: 1500
        parapet: 0
        wind_speed: 110
        exposure_category: D
        """
        self.c_p_matrix = array([[0.964647, 0.843638, 0.850871, 0.733149],
                                 [0.730023, 0.631703, 0.637580, 0.541930],
                                 [0.452508, 0.391163, 0.394830, 0.335151],
                                 [0.253272, 0.220498, 0.222457, 0.190574],
                                 [0.060000, 0.050000, 0.052000, 0.039000]])

        self.system_type = SystemType.dualTilt
        self.module_type = ModuleType.PSeries
        self.values.max_system_pressure.return_value = 100
        self.values.system_type.return_value = self.system_type
        self.values.anchor_type.return_value = AnchorType.OMG_PowerGrip
        self.values.module_system_constants.return_value = self.system_type.module_constants(self.module_type)
        self.q_z = 37.72825742
        self.subject = BallastCalculator(self.values)

    def test_summary_table(self):
        self.values.ballast_block_weight.return_value = 20
        self.values.max_system_pressure.return_value = 17

        expected_value = {
            PanelType.Corner: {
                'anchors': [4, 3, 0, 0, 0],
                'ballast blocks': [15, 10, 33, 16, 0],
                'pressure': ['8.43', '6.38', '15.93', '8.83', '2.12']
            },
            PanelType.NorthSouth: {
                'anchors': [4, 3, 0, 0, 0],
                'ballast blocks': [5, 2, 28, 14, 0],
                'pressure': ['4.22', '2.97', '13.77', '7.90', '2.10']
            },
            PanelType.EastWest: {
                'anchors': [4, 3, 0, 0, 0],
                'ballast blocks': [6, 3, 28, 14, 0],
                'pressure': ['4.58', '3.29', '13.72', '7.86', '2.01']
            },
            PanelType.Middle: {
                'anchors': [3, 2, 0, 0, 0],
                'ballast blocks': [11, 10, 23, 11, 0],
                'pressure': ['6.61', '6.20', '11.59', '6.61', '1.99']
            },
        }

        received_table = self.subject.summary_table(self.c_p_matrix, self.q_z)
        eq_(received_table.keys(), expected_value.keys())
        for key in expected_value.keys():
            received_table[key].pop('warnings')
            self.assertDictEqual(received_table[key], expected_value[key])
        # Because the above is way more explicit than simply 'eq_(received_table, expected_value)'

    def test_summary_table_when_special_small_building(self):
        """
        Generated using site parameters:
        height: 15
        length: 100
        width: 100
        parapet: 0
        wind_speed: 110
        exposure_category: B
        """
        self.c_p_matrix = array([[0.277131, 0.182052, 0.187735, 0.129774],
                                 [0.182934, 0.115506, 0.118940, 0.084465],
                                 [0.111800, 0.075909, 0.077290, 0.066873],
                                 [0.068781, 0.050000, 0.052000, 0.039068],
                                 [0.060000, 0.050000, 0.052000, 0.039000]])

        self.system_type = SystemType.dualTilt
        self.module_type = ModuleType.PSeries
        self.values.max_system_pressure.return_value = 8
        self.values.ballast_block_weight.return_value = 16
        self.values.system_type.return_value = self.system_type
        self.values.anchor_type.return_value = AnchorType.OMG_PowerGrip
        self.values.module_system_constants.return_value = self.system_type.module_constants(self.module_type)
        self.q_z = 18.43072
        self.subject = BallastCalculator(self.values)

        expected_value = {
            PanelType.Corner: {
                'anchors': [0, 0, 0, 0, 0],
                'ballast blocks': [8, 3, 0, 0, 0],
                'pressure': ['4.79', '3.11', '2.12', '2.12', '2.12']
            },
            PanelType.NorthSouth: {
                'anchors': [0, 0, 0, 0, 0],
                'ballast blocks': [3, 0, 0, 0, 0],
                'pressure': ['3.08', '2.10', '2.10', '2.10', '2.10']
            },
            PanelType.EastWest: {
                'anchors': [0, 0, 0, 0, 0],
                'ballast blocks': [4, 0, 0, 0, 0],
                'pressure': ['3.32', '2.01', '2.01', '2.01', '2.01']
            },
            PanelType.Middle: {
                'anchors': [0, 0, 0, 0, 0],
                'ballast blocks': [1, 0, 0, 0, 0],
                'pressure': ['2.32', '1.99', '1.99', '1.99', '1.99']
            },
        }

        received_table = self.subject.summary_table(self.c_p_matrix, self.q_z)
        eq_(received_table.keys(), expected_value.keys())
        for key in expected_value.keys():
            received_table[key].pop('warnings')
            self.assertDictEqual(received_table[key], expected_value[key])

    def test_ballast_and_trays_matrix(self):
        raw_data = [
            Panel(wind_zone=2, panel_type=PanelType.Corner),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth),
            Panel(wind_zone=4, panel_type=PanelType.EastWest),
            Panel(wind_zone=0, panel_type=PanelType.Corner),
            Panel(wind_zone=1, panel_type=PanelType.Middle),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth),
            Panel(wind_zone=1, panel_type=PanelType.Corner)
        ]
        self.values.ballast_block_weight.return_value = 20

        expected_value = [
            Panel(wind_zone=2, panel_type=PanelType.Corner, ballast=33, link_tray=2, cross_tray=3, wind_anchors=0,
                  pressure=15.93),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth, ballast=28, link_tray=2, cross_tray=2, wind_anchors=0,
                  pressure=13.77),
            Panel(wind_zone=4, panel_type=PanelType.EastWest, ballast=0, link_tray=0, cross_tray=0, wind_anchors=0,
                  pressure=2.01),
            Panel(wind_zone=0, panel_type=PanelType.Corner, ballast=15, link_tray=2, cross_tray=1, wind_anchors=4,
                  pressure=8.43, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=1, panel_type=PanelType.Middle, ballast=40, link_tray=2, cross_tray=4, wind_anchors=0,
                  pressure=18.68),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth, ballast=14, link_tray=2, cross_tray=0, wind_anchors=0,
                  pressure=7.90),
            Panel(wind_zone=1, panel_type=PanelType.Corner, ballast=10, link_tray=2, cross_tray=1, wind_anchors=3,
                  pressure=6.38),
            ]

        assert_array_is_close(self.subject.ballast_and_trays_matrix(self.c_p_matrix, self.q_z, raw_data),
                              expected_value, decimal=2)

    def test_update_ballast(self):
        self.values.ballast_block_weight.return_value = 20
        self.values.max_system_pressure.return_value = 17

        panels = [
            Panel(wind_zone=2, panel_type=PanelType.Corner, wind_anchors=0, seismic_anchors=1),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth, wind_anchors=0, seismic_anchors=1),
            Panel(wind_zone=4, panel_type=PanelType.EastWest, wind_anchors=0, seismic_anchors=1),
            Panel(wind_zone=0, panel_type=PanelType.Corner, wind_anchors=4, seismic_anchors=0),
            Panel(wind_zone=1, panel_type=PanelType.Middle, wind_anchors=0, seismic_anchors=1),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth, wind_anchors=0, seismic_anchors=1),
            Panel(wind_zone=1, panel_type=PanelType.Corner, wind_anchors=3, seismic_anchors=0)
        ]

        expected = [
            Panel(wind_zone=2, panel_type=PanelType.Corner, wind_anchors=0, seismic_anchors=1,
                  ballast=18, link_tray=2, cross_tray=1, pressure=9.65),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth, wind_anchors=0, seismic_anchors=1,
                  ballast=13, link_tray=2, cross_tray=0, pressure=7.50),
            Panel(wind_zone=4, panel_type=PanelType.EastWest, wind_anchors=0, seismic_anchors=1,
                  ballast=0, link_tray=0, cross_tray=0, pressure=2.01),
            Panel(wind_zone=0, panel_type=PanelType.Corner, wind_anchors=4, seismic_anchors=0,
                  ballast=15, link_tray=2, cross_tray=1, pressure=8.43),
            Panel(wind_zone=1, panel_type=PanelType.Middle, wind_anchors=0, seismic_anchors=1,
                  ballast=25, link_tray=2, cross_tray=3, pressure=12.47),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth, wind_anchors=0, seismic_anchors=1,
                  ballast=0, link_tray=0, cross_tray=0, pressure=2.10),
            Panel(wind_zone=1, panel_type=PanelType.Corner, wind_anchors=3, seismic_anchors=0,
                  ballast=10, link_tray=2, cross_tray=1, pressure=6.38),
        ]

        result = self.subject.update_ballast(self.c_p_matrix, self.q_z, panels)
        assert_array_is_close(result, expected, decimal=2)

    def test_presented_link_trays(self):
        panels = [
            Panel(panel_type=PanelType.Corner, link_tray=2),
            Panel(panel_type=PanelType.NorthSouth, link_tray=2),
            Panel(panel_type=PanelType.EastWest, link_tray=0),
            Panel(panel_type=PanelType.Corner, link_tray=2),
            Panel(panel_type=PanelType.Middle, link_tray=2),
            Panel(panel_type=PanelType.NorthSouth, link_tray=1),
            Panel(panel_type=PanelType.Corner, link_tray=2),
        ]

        expected = [
            Panel(presented_link_tray=2, panel_type=PanelType.Corner, link_tray=2),
            Panel(presented_link_tray=2, panel_type=PanelType.NorthSouth, link_tray=2),
            Panel(presented_link_tray=1, panel_type=PanelType.EastWest, link_tray=0),
            Panel(presented_link_tray=2, panel_type=PanelType.Corner, link_tray=2),
            Panel(presented_link_tray=1, panel_type=PanelType.Middle, link_tray=2),
            Panel(presented_link_tray=2, panel_type=PanelType.NorthSouth, link_tray=1),
            Panel(presented_link_tray=2, panel_type=PanelType.Corner, link_tray=2)
        ]

        assert_array_equal(self.subject.show_presented_link_trays(panels), expected)

    def test_pressure_exceeeding_max_system_pressure(self):
        self.values.ballast_block_weight.return_value = 20
        self.values.max_system_pressure.return_value = 17

        panel_data = [
            Panel(wind_zone=2, panel_type=PanelType.Corner),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth),
            Panel(wind_zone=4, panel_type=PanelType.EastWest),
            Panel(wind_zone=0, panel_type=PanelType.Corner),
            Panel(wind_zone=1, panel_type=PanelType.Middle),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth),
            Panel(wind_zone=1, panel_type=PanelType.Corner)
        ]

        expected_value = [
            Panel(wind_zone=2, panel_type=PanelType.Corner, ballast=33, link_tray=2, cross_tray=3, wind_anchors=0, pressure=15.93),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth, ballast=28, link_tray=2, cross_tray=2, wind_anchors=0, pressure=13.77),
            Panel(wind_zone=4, panel_type=PanelType.EastWest, ballast=0, link_tray=0, cross_tray=0, wind_anchors=0, pressure=2.01),
            Panel(wind_zone=0, panel_type=PanelType.Corner, ballast=15, link_tray=2, cross_tray=1, wind_anchors=4, pressure=8.43, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=1, panel_type=PanelType.Middle, ballast=10, link_tray=2, cross_tray=1, wind_anchors=2, pressure=6.20),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth, ballast=14, link_tray=2, cross_tray=0, wind_anchors=0, pressure=7.90),
            Panel(wind_zone=1, panel_type=PanelType.Corner, ballast=10, link_tray=2, cross_tray=1, wind_anchors=3, pressure=6.38)
        ]

        assert_array_is_close(self.subject.ballast_and_trays_matrix(self.c_p_matrix, self.q_z, panel_data), expected_value, decimal=2)

    def test_ballast_count_when_special_small_building(self):
        """
        Generated using site parameters:
        height: 15
        length: 100
        width: 100
        parapet: 0
        wind_speed: 110
        exposure_category: B

        """
        self.c_p_matrix = array([[0.277131, 0.182052, 0.187735, 0.129774],
                                 [0.182934, 0.115506, 0.118940, 0.084465],
                                 [0.111800, 0.075909, 0.077290, 0.066873],
                                 [0.068781, 0.050000, 0.052000, 0.039068],
                                 [0.060000, 0.050000, 0.052000, 0.039000]])

        self.system_type = SystemType.dualTilt
        self.module_type = ModuleType.PSeries
        self.values.max_system_pressure.return_value = 8
        self.values.ballast_block_weight.return_value = 16
        self.values.system_type.return_value = self.system_type
        self.values.anchor_type.return_value = AnchorType.OMG_PowerGrip
        self.values.module_system_constants.return_value = self.system_type.module_constants(self.module_type)
        self.q_z = 18.43072
        self.subject = BallastCalculator(self.values)

        expected = Panel(wind_zone=0, panel_type=PanelType.NorthSouth, ballast=3, link_tray=0, cross_tray=0, wind_anchors=0, pressure=3.08)
        received = self.subject.ballast_tray_and_anchor_count(0, PanelType.NorthSouth, 16, 8, self.c_p_matrix, self.q_z)
        assert received.almost_equal(expected, decimal=2)

    def test_ballast_count_when_base_weight_greater_than_uplift(self):
        expected = Panel(wind_zone=4, panel_type=PanelType.Middle, ballast=0, link_tray=0, cross_tray=0, wind_anchors=0, pressure=1.99)
        assert self.subject.ballast_tray_and_anchor_count(4, PanelType.Middle, 14, 100, self.c_p_matrix, self.q_z).almost_equal(expected, decimal=2)

    def test_individual_ballast_block_uplift_greater_than_base_weight(self):
        expected = Panel(wind_zone=0, panel_type=PanelType.Middle, ballast=15, link_tray=2, cross_tray=2, wind_anchors=3, pressure=6.47)
        assert self.subject.ballast_tray_and_anchor_count(0, PanelType.Middle, 14, 100, self.c_p_matrix, self.q_z).almost_equal(expected, decimal=2)

        expected = Panel(wind_zone=0, panel_type=PanelType.Middle, ballast=11, link_tray=2, cross_tray=1, wind_anchors=3, pressure=6.61)
        assert self.subject.ballast_tray_and_anchor_count(0, PanelType.Middle, 20, 100, self.c_p_matrix, self.q_z).almost_equal(expected, decimal=2)

    def test_ballast_blocks_exceeding_tray_capacity(self):
        expected = Panel(wind_zone=0, panel_type=PanelType.Middle, ballast=18, link_tray=2, cross_tray=2, wind_anchors=3, pressure=6.59)
        assert self.subject.ballast_tray_and_anchor_count(0, PanelType.Middle, 12, 100, self.c_p_matrix, self.q_z).almost_equal(expected, decimal=2)

    def test_when_max_system_pressure_is_lower_than_base_weight_pressure(self):
        max_system_pressure = 0
        expected = Panel(wind_zone=0, panel_type=PanelType.Middle, ballast=0, link_tray=0, cross_tray=0, wind_anchors=4, pressure=1.99)

        assert self.subject.ballast_tray_and_anchor_count(0, PanelType.Middle, 12, max_system_pressure, self.c_p_matrix, self.q_z).almost_equal(
                                  expected,
                                  decimal=2)

    def test_uplift(self):
        assert_almost_equal(self.subject.uplift(0.62717, self.q_z), 1051.54, 1)

    def test_calculate_single_unit_pressure_on_roof(self):
        assert_almost_equal(self.subject.calculate_pressure_on_roof(78, 12, 105.23), 21.3213163825, 6)
