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': [14, 10, 32, 16, 0],
                'pressure': ['8.26', '6.63', '15.76', '9.08', '2.39']
            },
            PanelType.NorthSouth: {
                'anchors': [4, 3, 0, 0, 0],
                'ballast blocks': [4, 2, 27, 13, 0],
                'pressure': ['4.08', '3.22', '13.63', '7.77', '2.37']
            },
            PanelType.EastWest: {
                'anchors': [4, 3, 0, 0, 0],
                'ballast blocks': [5, 2, 27, 13, 0],
                'pressure': ['4.49', '3.19', '13.63', '7.76', '2.30']
            },
            PanelType.Middle: {
                'anchors': [3, 2, 0, 0, 0],
                'ballast blocks': [10, 10, 23, 11, 0],
                'pressure': ['6.51', '6.51', '11.90', '6.92', '2.28']
            },
        }

        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': [7, 3, 0, 0, 0],
                'pressure': ['4.68', '3.37', '2.39', '2.39', '2.39']
            },
            PanelType.NorthSouth: {
                'anchors': [0, 0, 0, 0, 0],
                'ballast blocks': [3, 0, 0, 0, 0],
                'pressure': ['3.35', '2.37', '2.37', '2.37', '2.37']
            },
            PanelType.EastWest: {
                'anchors': [0, 0, 0, 0, 0],
                'ballast blocks': [3, 0, 0, 0, 0],
                'pressure': ['3.28', '2.30', '2.30', '2.30', '2.30']
            },
            PanelType.Middle: {
                'anchors': [0, 0, 0, 0, 0],
                'ballast blocks': [0, 0, 0, 0, 0],
                'pressure': ['2.28', '2.28', '2.28', '2.28', '2.28']
            },
        }

        received_table = self.subject.summary_table(self.c_p_matrix, self.q_z)
        print("====\r\n", received_table)
        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=32, link_tray=2, cross_tray=3, wind_anchors=0,
                  pressure=15.764460),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth, ballast=27, link_tray=2, cross_tray=2, wind_anchors=0,
                  pressure=13.630342),
            Panel(wind_zone=4, panel_type=PanelType.EastWest, ballast=0, link_tray=0, cross_tray=0, wind_anchors=0,
                  pressure=2.299163),
            Panel(wind_zone=0, panel_type=PanelType.Corner, ballast=14, link_tray=2, cross_tray=1, wind_anchors=4,
                  pressure=8.263513, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=1, panel_type=PanelType.Middle, ballast=40, link_tray=2, cross_tray=4, wind_anchors=0,
                  pressure=18.993076),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth, ballast=13, link_tray=2, cross_tray=0, wind_anchors=0,
                  pressure=7.767559),
            Panel(wind_zone=1, panel_type=PanelType.Corner, ballast=10, link_tray=2, cross_tray=1, wind_anchors=3,
                  pressure=6.625349),
            ]

        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=17, link_tray=2, cross_tray=1, pressure=9.492135463546356),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth, wind_anchors=0, seismic_anchors=1,
                  ballast=12, link_tray=2, cross_tray=0, pressure=7.358017551755176),
            Panel(wind_zone=4, panel_type=PanelType.EastWest, wind_anchors=0, seismic_anchors=1,
                  ballast=0, link_tray=0, cross_tray=0, pressure=2.299162916291629),
            Panel(wind_zone=0, panel_type=PanelType.Corner, wind_anchors=4, seismic_anchors=0,
                  ballast=14, link_tray=2, cross_tray=1, pressure=8.263512601260127),
            Panel(wind_zone=1, panel_type=PanelType.Middle, wind_anchors=0, seismic_anchors=1,
                  ballast=24, link_tray=2, cross_tray=3, pressure=12.375918091809181),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth, wind_anchors=0, seismic_anchors=1,
                  ballast=0, link_tray=0, cross_tray=0, pressure=2.3677610261026105),
            Panel(wind_zone=1, panel_type=PanelType.Corner, wind_anchors=3, seismic_anchors=0,
                  ballast=10, link_tray=2, cross_tray=1, pressure=6.625348784878488),
        ]

        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=32, link_tray=2, cross_tray=3, wind_anchors=0, pressure=15.764460),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth, ballast=27, link_tray=2, cross_tray=2, wind_anchors=0, pressure=13.630342),
            Panel(wind_zone=4, panel_type=PanelType.EastWest, ballast=0, link_tray=0, cross_tray=0, wind_anchors=0, pressure=2.299163),
            Panel(wind_zone=0, panel_type=PanelType.Corner, ballast=14, link_tray=2, cross_tray=1, wind_anchors=4, pressure= 8.263513, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=1, panel_type=PanelType.Middle, ballast=10, link_tray=2, cross_tray=1, wind_anchors=2, pressure=6.513135),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth, ballast=13, link_tray=2, cross_tray=0, wind_anchors=0, pressure=7.767559),
            Panel(wind_zone=1, panel_type=PanelType.Corner, ballast=10, link_tray=2, cross_tray=1, wind_anchors=3, pressure=6.625349)
        ]

        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.350659)
        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=2.277252)
        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=14, link_tray=2, cross_tray=2, wind_anchors=3, pressure=6.495729)
        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=10, link_tray=2, cross_tray=1, wind_anchors=3, pressure=6.513135)
        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=17, link_tray=2, cross_tray=2, wind_anchors=3, pressure=6.659545)
        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=2.277252)

        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)
