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 BallastCalculatorWhenDualTiltAnd128CellTest(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.751977, 0.656898, 0.662582, 0.570086],
                                 [0.568747, 0.491495, 0.496113, 0.420960],
                                 [0.352521, 0.304321, 0.307202, 0.260312],
                                 [0.197385, 0.171634, 0.173174, 0.148123],
                                 [0.060000, 0.050000, 0.052000, 0.039000]
                                 ])
        self.system_type = SystemType.dualTilt
        self.module_type = ModuleType.Cell128
        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': [3, 2, 0, 0, 0],
                'ballast blocks': [14, 13, 25, 12, 0],
                'pressure': ['8.00', '7.61', '12.36', '7.16', '2.40']
            },
            PanelType.NorthSouth: {
                'anchors': [3, 2, 0, 0, 0],
                'ballast blocks': [6, 7, 21, 9, 0],
                'pressure': ['4.79', '5.18', '10.72', '5.93', '2.38']
            },
            PanelType.EastWest: {
                'anchors': [3, 2, 0, 0, 0],
                'ballast blocks': [7, 7, 21, 10, 0],
                'pressure': ['5.18', '5.18', '10.71', '6.29', '2.31']
            },
            PanelType.Middle: {
                'anchors': [2, 0, 0, 0, 0],
                'ballast blocks': [14, 31, 17, 7, 0],
                'pressure': ['7.89', '14.66', '9.07', '5.10', '2.29']
            },
        }

        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_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=25, link_tray=2, cross_tray=2, wind_anchors=0,
                  pressure=12.360800772863891),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth, ballast=21, link_tray=2, cross_tray=1, wind_anchors=0,
                  pressure=10.715259768140834),
            Panel(wind_zone=4, panel_type=PanelType.EastWest, ballast=0, link_tray=0, cross_tray=0, wind_anchors=0,
                  pressure=2.3107513954486905),
            Panel(wind_zone=0, panel_type=PanelType.Corner, ballast=14, link_tray=2, cross_tray=1, wind_anchors=3,
                  pressure=8.001084156290254),
            Panel(wind_zone=1, panel_type=PanelType.Middle, ballast=31, link_tray=2, cross_tray=3, wind_anchors=0,
                  pressure=14.659635036496352),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth, ballast=9, link_tray=1, cross_tray=0, wind_anchors=0,
                  pressure=5.928870759982826),
            Panel(wind_zone=1, panel_type=PanelType.Corner, ballast=44, link_tray=2, cross_tray=4, wind_anchors=0,
                  pressure=19.907666380420782)
            ]

        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_ballast_and_trays_matrix_ups_ballast_for_fuzzy_wind_zone(self):
        raw_data = [
            Panel(wind_zone=2, panel_type=PanelType.Corner, fuzzy_wind_zone=True),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth, fuzzy_wind_zone=True),
            Panel(wind_zone=4, panel_type=PanelType.EastWest, fuzzy_wind_zone=True),
            Panel(wind_zone=0, panel_type=PanelType.Corner, fuzzy_wind_zone=True),
            Panel(wind_zone=1, panel_type=PanelType.Middle, fuzzy_wind_zone=True),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth, fuzzy_wind_zone=True),
            Panel(wind_zone=1, panel_type=PanelType.Corner, fuzzy_wind_zone=True)
        ]
        self.values.ballast_block_weight.return_value = 20

        expected_value = [
            Panel(wind_zone=2, panel_type=PanelType.Corner, ballast=29, link_tray=2, cross_tray=2, wind_anchors=0,
                  pressure=13.923703, fuzzy_wind_zone=True),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth, ballast=25, link_tray=2, cross_tray=2, wind_anchors=0,
                  pressure=12.339897, fuzzy_wind_zone=True),
            Panel(wind_zone=4, panel_type=PanelType.EastWest, ballast=0, link_tray=0, cross_tray=0, wind_anchors=0,
                  pressure=2.310751, fuzzy_wind_zone=True),
            Panel(wind_zone=0, panel_type=PanelType.Corner, ballast=9, link_tray=2, cross_tray=1, wind_anchors=4,
                  pressure=6.047456, fuzzy_wind_zone=True, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=1, panel_type=PanelType.Middle, ballast=36, link_tray=2, cross_tray=4, wind_anchors=0,
                  pressure=16.674802, fuzzy_wind_zone=True),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth, ballast=12, link_tray=2, cross_tray=0, wind_anchors=0,
                  pressure=7.137190, fuzzy_wind_zone=True),
            Panel(wind_zone=1, panel_type=PanelType.Corner, ballast=6, link_tray=2, cross_tray=0, wind_anchors=3,
                  pressure=4.813740, fuzzy_wind_zone=True)
        ]

        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=3, 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=2, seismic_anchors=0)
        ]

        expected = [
            Panel(wind_zone=2, panel_type=PanelType.Corner, wind_anchors=0, seismic_anchors=1,
                  ballast=10, link_tray=2, cross_tray=0, pressure=6.376642335766423),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth, wind_anchors=0, seismic_anchors=1,
                  ballast=6, link_tray=1, cross_tray=0, pressure=4.756693860025763),
            Panel(wind_zone=4, panel_type=PanelType.EastWest, wind_anchors=0, seismic_anchors=1,
                  ballast=0, link_tray=0, cross_tray=0, pressure=2.3107513954486905),
            Panel(wind_zone=0, panel_type=PanelType.Corner, wind_anchors=3, seismic_anchors=0,
                  ballast=14, link_tray=2, cross_tray=1, pressure=8.001084156290254),
            Panel(wind_zone=1, panel_type=PanelType.Middle, wind_anchors=0, seismic_anchors=1,
                  ballast=16, link_tray=2, cross_tray=1, pressure=8.675476599398884),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth, wind_anchors=0, seismic_anchors=1,
                  ballast=0, link_tray=0, cross_tray=0, pressure=2.3761979390296264),
            Panel(wind_zone=1, panel_type=PanelType.Corner, wind_anchors=2, seismic_anchors=0,
                  ballast=13, link_tray=2, cross_tray=1, pressure=7.610358522971233),
        ]

        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=25, link_tray=2, cross_tray=2, wind_anchors=0, pressure=12.360800772863891),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth, ballast=21, link_tray=2, cross_tray=1, wind_anchors=0, pressure=10.715259768140834),
            Panel(wind_zone=4, panel_type=PanelType.EastWest, ballast=0, link_tray=0, cross_tray=0, wind_anchors=0, pressure=2.3107513954486905),
            Panel(wind_zone=0, panel_type=PanelType.Corner, ballast=14, link_tray=2, cross_tray=1, wind_anchors=3, pressure=8.001084156290254),
            Panel(wind_zone=1, panel_type=PanelType.Middle, ballast=31, link_tray=2, cross_tray=3, wind_anchors=0, pressure=14.659635036496352),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth, ballast=9, link_tray=1, cross_tray=0, wind_anchors=0, pressure=5.928870759982826),
            Panel(wind_zone=1, panel_type=PanelType.Corner, ballast=13, link_tray=2, cross_tray=1, wind_anchors=2, pressure=7.610358522971233)
        ]

        print("===\r\n")
        print(self.subject.ballast_and_trays_matrix(self.c_p_matrix, self.q_z, panel_data))
        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_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.289848)
        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=19, link_tray=2, cross_tray=2, wind_anchors=2, pressure=7.682057)
        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=14, link_tray=2, cross_tray=1, wind_anchors=2, pressure=7.894025)
        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=22, link_tray=2, cross_tray=3, wind_anchors=2, pressure=7.704719)
        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=3, pressure=2.289848)

        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), 1102.2, 1)

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