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
from test.test_helpers import assert_array_is_close


class BallastCalculatorWhenDualTiltAnd96CellTest(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.7800, 0.70, 0.68, 0.546],
                                 [0.5949, 0.52, 0.51, 0.400],
                                 [0.3700, 0.32, 0.32, 0.250],
                                 [0.2100, 0.18, 0.18, 0.140],
                                 [0.0600, 0.05, 0.05, 0.040]])
        self.system_type = SystemType.dualTilt
        self.module_type = ModuleType.Cell96
        self.values.max_system_pressure.return_value = 100
        self.values.system_type.return_value = self.system_type
        self.values.module_type.return_value = self.module_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.729
        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': [2, 5, 20, 10, 0],
                'pressure': ['3.48', '5.03', '12.91', '7.62', '2.39']
            },
            PanelType.NorthSouth: {
                'anchors': [2, 1, 0, 0, 0],
                'ballast blocks': [12, 15, 17, 8, 0],
                'pressure': ['8.75', '10.30', '11.34', '6.56', '2.37']
            },
            PanelType.EastWest: {
                'anchors': [2, 1, 0, 0, 0],
                'ballast blocks': [10, 14, 17, 8, 0],
                'pressure': ['7.62', '9.69', '11.24', '6.52', '2.28']
            },
            PanelType.Middle: {
                'anchors': [2, 0, 0, 0, 0],
                'ballast blocks': [2, 22, 12, 5, 0],
                'pressure': ['3.34', '13.87', '8.63', '4.85', '2.26']
            },
        }

        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):
        panels = [
            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

        # Ballast, link tray, cross tray, anchor count, psf
        expected_value = [
            Panel(wind_zone=2, panel_type=PanelType.Corner, ballast=20, link_tray=2, cross_tray=1, wind_anchors=0, pressure=12.910398406374503),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth, ballast=17, link_tray=2, cross_tray=1, wind_anchors=0, pressure=11.336414342629482),
            Panel(wind_zone=4, panel_type=PanelType.EastWest, ballast=0, link_tray=0, cross_tray=0, wind_anchors=0, pressure=2.2788844621513946),
            Panel(wind_zone=0, panel_type=PanelType.Corner, ballast=47, link_tray=2, cross_tray=4, wind_anchors=0, pressure=27.082988047808765),
            Panel(wind_zone=1, panel_type=PanelType.Middle, ballast=22, link_tray=2, cross_tray=2, wind_anchors=0, pressure=13.874521912350598),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth, ballast=8, link_tray=1, cross_tray=0, wind_anchors=0, pressure=6.564223107569722),
            Panel(wind_zone=1, panel_type=PanelType.Corner, ballast=35, link_tray=2, cross_tray=3, wind_anchors=0, pressure=20.8049203187251)
        ]
        assert_array_is_close(self.subject.ballast_and_trays_matrix(self.c_p_matrix, self.q_z, panels), 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_value = [
            Panel(wind_zone=2, panel_type=PanelType.Corner, wind_anchors=0, seismic_anchors=1, ballast=5, link_tray=0, cross_tray=0, pressure=4.982729083665339),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth, wind_anchors=0, seismic_anchors=1, ballast=2, link_tray=0, cross_tray=0, pressure=3.408745019920319),
            Panel(wind_zone=4, panel_type=PanelType.EastWest, wind_anchors=0, seismic_anchors=1, ballast=0, link_tray=0, cross_tray=0, pressure=2.2788844621513946),
            Panel(wind_zone=0, panel_type=PanelType.Corner, wind_anchors=3, seismic_anchors=0, ballast=2, link_tray=1, cross_tray=0, pressure=3.4768525896414344),
            Panel(wind_zone=1, panel_type=PanelType.Middle, wind_anchors=0, seismic_anchors=1, ballast=7, link_tray=2, cross_tray=0, pressure=5.98),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth, wind_anchors=0, seismic_anchors=1, ballast=0, link_tray=0, cross_tray=0, pressure=2.3728884462151396),
            Panel(wind_zone=1, panel_type=PanelType.Corner, wind_anchors=2, seismic_anchors=0, ballast=5, link_tray=1, cross_tray=0, pressure=5.030637450199203)
        ]

        result = self.subject.update_ballast(self.c_p_matrix, self.q_z, panels)
        print("===")
        print(result)
        print("===")
        assert_array_is_close(result, expected_value, decimal=2)

    def test_update_ballast_with_fuzzy_wind_zone(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, fuzzy_wind_zone=True),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth, wind_anchors=0, seismic_anchors=1, fuzzy_wind_zone=True),
            Panel(wind_zone=4, panel_type=PanelType.EastWest, wind_anchors=0, seismic_anchors=1, fuzzy_wind_zone=True),
            Panel(wind_zone=0, panel_type=PanelType.Corner, wind_anchors=3, seismic_anchors=0, fuzzy_wind_zone=True),
            Panel(wind_zone=1, panel_type=PanelType.Middle, wind_anchors=0, seismic_anchors=1, fuzzy_wind_zone=True),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth, wind_anchors=0, seismic_anchors=1, fuzzy_wind_zone=True),
            Panel(wind_zone=1, panel_type=PanelType.Corner, wind_anchors=2, seismic_anchors=0, fuzzy_wind_zone=True)
        ]

        expected_value = [
            Panel(wind_zone=2, panel_type=PanelType.Corner, wind_anchors=0, seismic_anchors=1, ballast=9, link_tray=2, cross_tray=0, pressure=7.150517928286853, fuzzy_wind_zone=True),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth, wind_anchors=0, seismic_anchors=1, ballast=5, link_tray=0, cross_tray=0, pressure=4.962529880478088, fuzzy_wind_zone=True),
            Panel(wind_zone=4, panel_type=PanelType.EastWest, wind_anchors=0, seismic_anchors=1, ballast=0, link_tray=0, cross_tray=0, pressure=2.2788844621513946, fuzzy_wind_zone=True),
            Panel(wind_zone=0, panel_type=PanelType.Corner, wind_anchors=3, seismic_anchors=0, ballast=9, link_tray=2, cross_tray=0, pressure=7.150517928286853, fuzzy_wind_zone=True),
            Panel(wind_zone=1, panel_type=PanelType.Middle, wind_anchors=0, seismic_anchors=1, ballast=11, link_tray=2, cross_tray=1, pressure=8.114382470119523, fuzzy_wind_zone=True),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth, wind_anchors=0, seismic_anchors=1, ballast=0, link_tray=0, cross_tray=0, pressure=2.3728884462151396, fuzzy_wind_zone=True),
            Panel(wind_zone=1, panel_type=PanelType.Corner, wind_anchors=2, seismic_anchors=0, ballast=11, link_tray=2, cross_tray=0, pressure=8.186374501992033, fuzzy_wind_zone=True)
        ]

        result = self.subject.update_ballast(self.c_p_matrix, self.q_z, panels)
        assert_array_is_close(result, expected_value, 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_value = [
            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_value)

    def test_pressure_exceeeding_max_system_pressure(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),
            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=20, link_tray=2, cross_tray=1, wind_anchors=0, pressure=12.910398406374503),
            Panel(wind_zone=2, panel_type=PanelType.NorthSouth, ballast=17, link_tray=2, cross_tray=1, wind_anchors=0, pressure=11.336414342629482),
            Panel(wind_zone=4, panel_type=PanelType.EastWest, ballast=0, link_tray=0, cross_tray=0, wind_anchors=0, pressure=2.2788844621513946),
            Panel(wind_zone=0, panel_type=PanelType.Corner, ballast=2, link_tray=1, cross_tray=0, wind_anchors=3, pressure=3.4768525896414344),
            Panel(wind_zone=1, panel_type=PanelType.Middle, ballast=22, link_tray=2, cross_tray=2, wind_anchors=0, pressure=13.874521912350598),
            Panel(wind_zone=3, panel_type=PanelType.NorthSouth, ballast=8, link_tray=1, cross_tray=0, wind_anchors=0, pressure=6.564223107569722),
            Panel(wind_zone=1, panel_type=PanelType.Corner, ballast=5, link_tray=1, cross_tray=0, wind_anchors=2, pressure=5.030637450199203)
        ]
        assert_array_is_close(self.subject.ballast_and_trays_matrix(self.c_p_matrix, self.q_z, panels), expected_value, decimal=2)

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

    def test_individual_ballast_block_uplift_greater_than_base_weight(self):
        assert self.subject.ballast_tray_and_anchor_count(0, PanelType.Middle, 14, 100, self.c_p_matrix,
                                                          self.q_z).almost_equal(
            Panel(wind_zone=0, panel_type=PanelType.Middle, ballast=45, link_tray=2, cross_tray=4, wind_anchors=0, pressure=18.920438),
            decimal=2)
        assert self.subject.ballast_tray_and_anchor_count(0, PanelType.Middle, 20, 100, self.c_p_matrix,
                                                          self.q_z).almost_equal(
            Panel(wind_zone=0, panel_type=PanelType.Middle, ballast=32, link_tray=2, cross_tray=3, wind_anchors=0, pressure=19.116474),
            decimal=2)

    def test_ballast_blocks_exceeding_tray_capacity(self):
        assert self.subject.ballast_tray_and_anchor_count(0, PanelType.Middle, 12, 100, self.c_p_matrix,
                                                          self.q_z).almost_equal(
            Panel(wind_zone=0, panel_type=PanelType.Middle, ballast=3, link_tray=1, cross_tray=0, wind_anchors=2, pressure=3.23),
            decimal=2)
        assert self.subject.ballast_tray_and_anchor_count(0, PanelType.Corner, 12, 100, self.c_p_matrix,
                                                          self.q_z).almost_equal(
            Panel(wind_zone=0, panel_type=PanelType.Corner, ballast=3, link_tray=1, cross_tray=0, wind_anchors=3, pressure=3.37),
            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.26)

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

    def test_calculate_single_unit_pressure_on_roof(self):
        assert_almost_equal(self.subject.calculate_pressure_on_roof(78, 12, 105.23), 26.96, 2)
