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 BallastCalculatorWhenSingleTiltAndPSeriesTest(unittest.TestCase):
    def setUp(self):
        self.values = MagicMock()
        """
        Generated using site parameters:
        height: 100
        length: 1000
        width: 1500
        parapet: 0
        wind_speed: 150
        exposure_category: D
        """

        self.c_p_matrix = array([[1.475479, 1.345176, 1.120980, 1.088447],
                                 [1.161129, 1.055610, 0.879675, 0.853329],
                                 [0.664800, 0.615030, 0.439307, 0.428656],
                                 [0.669430, 0.601241, 0.429458, 0.414865],
                                 [0.450104, 0.272016, 0.408024, 0.263611],
                                 [0.428892, 0.387706, 0.298236, 0.288744],
                                 [0.202426, 0.185933, 0.185933, 0.180991],
                                 [0.547213, 0.491058, 0.491058, 0.474234],
                                 [1.503719, 1.317213, 1.317213, 1.261333],
                                 [0.897331, 0.798273, 0.399137, 0.384297],
                                 [0.172217, 0.151210, 0.108007, 0.103511]])
        self.system_type = SystemType.singleTilt
        self.module_type = ModuleType.PSeries
        self.values.max_system_pressure.return_value = 1000
        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 = 70.155851
        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': [7, 5, 3, 3, 2, 1, 0, 2, 7, 4, 0],
                'ballast blocks': [5, 11, 3, 4, 2, 15, 13, 9, 7, 6, 11],
                'pressure': ['5.07', '8.69', '3.77', '4.37', '3.17', '11.10', '9.80', '7.39', '6.37', '5.67', '8.59']
            },
            PanelType.NorthSouth: {
                'anchors': [6, 5, 2, 2, 0, 1, 0, 2, 6, 3, 0],
                'ballast blocks': [10, 3, 15, 14, 18, 12, 12, 5, 8, 14, 9],
                'pressure': ['8.09', '3.78', '11.11', '10.41', '12.82', '9.20', '9.20', '4.89', '6.89', '10.50', '7.30']
            },
            PanelType.EastWest: {
                'anchors': [5, 4, 2, 1, 1, 0, 0, 2, 6, 1, 0],
                'ballast blocks': [8, 5, 1, 15, 14, 20, 12, 5, 8, 13, 6],
                'pressure': ['6.87', '4.97', '2.50', '11.09', '10.49', '14.11', '9.19', '4.97', '6.87', '9.89', '5.51']
            },
            PanelType.Middle: {
                'anchors': [5, 4, 1, 1, 0, 0, 0, 2, 6, 1, 0],
                'ballast blocks': [6, 3, 15, 14, 18, 20, 11, 4, 4, 12, 6],
                'pressure': ['5.59', '3.68', '11.01', '10.41', '12.82', '14.12', '8.60', '4.29', '4.38', '9.20', '5.44']
            },
        }

        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):
        panels = array([
            Panel(wind_zone=5, panel_type=PanelType.Corner),
            Panel(wind_zone=6, panel_type=PanelType.NorthSouth),
            Panel(wind_zone=7, panel_type=PanelType.EastWest),
            Panel(wind_zone=0, panel_type=PanelType.Corner),
            Panel(wind_zone=8, panel_type=PanelType.Middle),
            Panel(wind_zone=9, panel_type=PanelType.NorthSouth),
            Panel(wind_zone=6, panel_type=PanelType.Corner),
            Panel(wind_zone=6, panel_type=PanelType.EastWest)
        ])

        self.values.ballast_block_weight.return_value = 20

        expected_value = [
            Panel(wind_zone=5, panel_type=PanelType.Corner, ballast=15, link_tray=0, cross_tray=1, wind_anchors=1,
                  pressure=11.101436),
            Panel(wind_zone=6, panel_type=PanelType.NorthSouth, ballast=12, link_tray=0, cross_tray=1, wind_anchors=0,
                  pressure=9.202700),
            Panel(wind_zone=7, panel_type=PanelType.EastWest, ballast=5, link_tray=2, cross_tray=0, wind_anchors=2,
                  pressure=4.967106, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=0, panel_type=PanelType.Corner, ballast=5, link_tray=0, cross_tray=1, wind_anchors=7,
                  pressure=5.070833, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=8, panel_type=PanelType.Middle, ballast=4, link_tray=1, cross_tray=1, wind_anchors=6,
                  pressure=4.380027, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=9, panel_type=PanelType.NorthSouth, ballast=14, link_tray=0, cross_tray=2, wind_anchors=3,
                  pressure=10.503803, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=6, panel_type=PanelType.Corner, ballast=13, link_tray=0, cross_tray=0, wind_anchors=0,
                  pressure=9.800333),
            Panel(wind_zone=6, panel_type=PanelType.EastWest, ballast=12, link_tray=2, cross_tray=0, wind_anchors=0,
                  pressure=9.188528)
            ]

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

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

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

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

    def test_pressure_exceeding_max_system_pressure(self):
        self.values.ballast_block_weight.return_value = 20
        self.values.max_system_pressure.return_value = 25

        panel_data = [
            Panel(wind_zone=7, panel_type=PanelType.Corner),
            Panel(wind_zone=7, panel_type=PanelType.NorthSouth),
            Panel(wind_zone=9, panel_type=PanelType.EastWest),
            Panel(wind_zone=5, panel_type=PanelType.Corner),
            Panel(wind_zone=6, panel_type=PanelType.Middle),
            Panel(wind_zone=8, panel_type=PanelType.NorthSouth),
            Panel(wind_zone=6, panel_type=PanelType.Corner)
        ]

        expected_value = [
            Panel(wind_zone=7, panel_type=PanelType.Corner, ballast=9, link_tray=0, cross_tray=0,
                  wind_anchors=2, pressure=7.388092, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=7, panel_type=PanelType.NorthSouth, ballast=5, link_tray=0, cross_tray=0,
                  wind_anchors=2, pressure=4.886296, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=9, panel_type=PanelType.EastWest, ballast=28, link_tray=2, cross_tray=2,
                  wind_anchors=0, pressure=19.027457),
            Panel(wind_zone=5, panel_type=PanelType.Corner, ballast=15, link_tray=0, cross_tray=1,
                  wind_anchors=1, pressure=11.101436),
            Panel(wind_zone=6, panel_type=PanelType.Middle, ballast=11, link_tray=1, cross_tray=1,
                  wind_anchors=0, pressure= 8.601449),
            Panel(wind_zone=8, panel_type=PanelType.NorthSouth, ballast=8, link_tray=0, cross_tray=2,
                  wind_anchors=6, pressure=6.885441, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=6, panel_type=PanelType.Corner, ballast=13, link_tray=0, cross_tray=0,
                  wind_anchors=0, pressure=9.800333)
        ]

        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):
        # Need to regenerate a c_p matrix to get to this case

        """
        Generated using site parameters:
        height: 30
        length: 100
        width: 100
        parapet: 0
        wind_speed: 110
        exposure_category: B
        """
        self.c_p_matrix = array([[0.831094, 0.670280, 0.558567, 0.384751],
                                 [0.639306, 0.509079, 0.424233, 0.283477],
                                 [0.418671, 0.357246, 0.255176, 0.198270],
                                 [0.332216, 0.248060, 0.177186, 0.122119],
                                 [0.242006, 0.158462, 0.190073, 0.089168],
                                 [0.225219, 0.174390, 0.134146, 0.089168],
                                 [0.184388, 0.158462, 0.113187, 0.089168],
                                 [0.269511, 0.200207, 0.200207, 0.127913],
                                 [0.581392, 0.351215, 0.351215, 0.150139],
                                 [0.407464, 0.285212, 0.142606, 0.089168],
                                 [0.184388, 0.158462, 0.113187, 0.089168]])
        self.q_z = 18.446284

        expected = Panel(wind_zone=10, panel_type=PanelType.Middle, ballast=0, link_tray=0, cross_tray=0,
                         wind_anchors=0, pressure=1.81)

        assert self.subject.ballast_tray_and_anchor_count(10, PanelType.Middle, 14, 1000, self.c_p_matrix,
                                                          self.q_z).almost_equal(expected, decimal=2)

    def test_ballast_and_tray_count_iterates_on_link_and_cross_trays_changing_weight_and_ballast_requirements(self):
        panel = self.subject.ballast_tray_and_anchor_count(3, PanelType.NorthSouth, 20, 17, self.c_p_matrix, self.q_z)
        assert_almost_equal(panel.pressure, 10.408820882088209, 2)

    def test_individual_ballast_block_uplift_greater_than_base_weight(self):
        expected = Panel(wind_zone=6, panel_type=PanelType.Middle, ballast=16, link_tray=1, cross_tray=1,
                         wind_anchors=0, pressure=8.722061)
        assert self.subject.ballast_tray_and_anchor_count(6, PanelType.Middle, 14, 1000, self.c_p_matrix,
                                                          self.q_z).almost_equal(expected, decimal=2)

        expected = Panel(wind_zone=6, panel_type=PanelType.Middle, ballast=11, link_tray=1, cross_tray=1,
                         wind_anchors=0, pressure=8.601449)
        assert self.subject.ballast_tray_and_anchor_count(6, PanelType.Middle, 20, 1000, 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.Corner, ballast=5, link_tray=0, cross_tray=1,
                         wind_anchors=7, pressure=5.070833, warnings=[PanelWarnings.MaxPsf])
        assert self.subject.ballast_tray_and_anchor_count(0, PanelType.Corner, 20, 1000, self.c_p_matrix,
                                                          self.q_z).almost_equal(expected, decimal=2)

        expected = Panel(wind_zone=8, panel_type=PanelType.Middle, ballast=4, link_tray=1, cross_tray=1,
                         wind_anchors=6, pressure=4.39, warnings=[PanelWarnings.MaxPsf])
        assert self.subject.ballast_tray_and_anchor_count(8, PanelType.Middle, 20, 1000, self.c_p_matrix,
                                                          self.q_z).almost_equal(expected, decimal=2)

    def test_ballast_blocks_when_panel_exceeds_max_system_pressure(self):
        expected = Panel(wind_zone=5, panel_type=PanelType.Middle, ballast=6, link_tray=1, cross_tray=0,
                         wind_anchors=1, pressure=4.41)

        assert self.subject.ballast_tray_and_anchor_count(5, PanelType.Middle, 14, 10, self.c_p_matrix,
                                                          self.q_z).almost_equal(expected, decimal=2)

        expected = Panel(wind_zone=0, panel_type=PanelType.Corner, ballast=7, link_tray=0, cross_tray=1,
                         wind_anchors=7, pressure=5.010527, warnings=[PanelWarnings.MaxPsf])
        assert self.subject.ballast_tray_and_anchor_count(0, PanelType.Corner, 14, 10, 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=5, panel_type=PanelType.Middle, ballast=0, link_tray=0, cross_tray=0,
                         wind_anchors=2, pressure=1.81)
        assert self.subject.ballast_tray_and_anchor_count(5, 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(1.506920, self.q_z), 2349.08, 2)

    def test_calculate_single_unit_pressure_on_roof(self): # Do manually, don't need excel
        assert_almost_equal(self.subject.calculate_pressure_on_roof(90, 20, 59.1), 56.06, 2)
        assert_almost_equal(self.subject.calculate_pressure_on_roof(15, 14, 54.5), 7.98, 2)
