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 BallastCalculatorWhenSingleTiltAnd128CellTest(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.464330, 1.334027, 1.111689, 1.079155],
                                 [1.152100, 1.046580, 0.872150, 0.845805],
                                 [0.660541, 0.610771, 0.436265, 0.425614],
                                 [0.663595, 0.595406, 0.425290, 0.410697],
                                 [0.446503, 0.269616, 0.404423, 0.261210],
                                 [0.425368, 0.384182, 0.295525, 0.286033],
                                 [0.201015, 0.184521, 0.184521, 0.179580],
                                 [0.542408, 0.486253, 0.486253, 0.469428],
                                 [1.487760, 1.301254, 1.301254, 1.245374],
                                 [0.888854, 0.789797, 0.394899, 0.380059],
                                 [0.172217, 0.151210, 0.108007, 0.103511]])
        self.system_type = SystemType.singleTilt
        self.module_type = ModuleType.Cell128
        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, 2, 0, 2, 7, 4, 0],
                'ballast blocks': [10, 15, 5, 6, 3, 1, 13, 11, 12, 9, 11],
                'pressure': ['7.89', '10.77', '4.83', '5.41', '3.68', '2.53', '9.53', '8.38', '9.04', '7.23', '8.29']
            },
            PanelType.NorthSouth: {
                'anchors': [6, 5, 3, 2, 0, 1, 0, 2, 6, 4, 0],
                'ballast blocks': [14, 6, 1, 15, 19, 13, 12, 6, 12, 1, 10],
                'pressure': ['10.11', '5.41', '2.45', '10.68', '12.98', '9.44', '8.87', '5.32', '8.96', '2.45', '7.62']
            },
            PanelType.EastWest: {
                'anchors': [5, 4, 2, 2, 1, 0, 0, 2, 6, 1, 0],
                'ballast blocks': [12, 7, 2, 1, 15, 21, 12, 6, 12, 14, 6],
                'pressure': ['8.94', '6.07', '3.05', '2.47', '10.67', '14.12', '8.85', '5.40', '9.03', '10.09', '5.35']
            },
            PanelType.Middle: {
                'anchors': [5, 4, 2, 1, 0, 0, 0, 2, 6, 1, 0],
                'ballast blocks': [9, 5, 2, 15, 18, 20, 12, 5, 7, 13, 6],
                'pressure': ['7.14', '4.84', '2.97', '10.59', '12.32', '13.56', '8.87', '4.75', '5.99', '9.44', '5.27']
            },
        }

        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):
        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=1, link_tray=0, cross_tray=0, wind_anchors=2,
                  pressure=2.532134),
            Panel(wind_zone=6, panel_type=PanelType.NorthSouth, ballast=12, link_tray=0, cross_tray=1, wind_anchors=0,
                  pressure=8.866209),
            Panel(wind_zone=7, panel_type=PanelType.EastWest, ballast=6, link_tray=2, cross_tray=0, wind_anchors=2,
                  pressure=5.400562, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=0, panel_type=PanelType.Corner, ballast=10, link_tray=0, cross_tray=2, wind_anchors=7,
                  pressure=7.891559, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=8, panel_type=PanelType.Middle, ballast=7, link_tray=1, cross_tray=1, wind_anchors=6,
                  pressure=5.991164, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=9, panel_type=PanelType.NorthSouth, ballast=1, link_tray=0, cross_tray=0, wind_anchors=4,
                  pressure=2.446694, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=6, panel_type=PanelType.Corner, ballast=13, link_tray=0, cross_tray=1, wind_anchors=0,
                  pressure=9.527003),
            Panel(wind_zone=6, panel_type=PanelType.EastWest, ballast=12, link_tray=2, cross_tray=0, wind_anchors=0,
                  pressure=8.852688)
            ]

        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=11, link_tray=0, cross_tray=1,
                  wind_anchors=2, pressure=8.376295, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=7, panel_type=PanelType.NorthSouth, ballast=6, link_tray=0, cross_tray=0,
                  wind_anchors=2, pressure=5.323465, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=9, panel_type=PanelType.EastWest, ballast=29, link_tray=2, cross_tray=2,
                  wind_anchors=0, pressure=18.814946),
            Panel(wind_zone=5, panel_type=PanelType.Corner, ballast=1, link_tray=0, cross_tray=0,
                  wind_anchors=2, pressure=2.532134),
            Panel(wind_zone=6, panel_type=PanelType.Middle, ballast=12, link_tray=1, cross_tray=1,
                  wind_anchors=0, pressure=8.867935),
            Panel(wind_zone=8, panel_type=PanelType.NorthSouth, ballast=12, link_tray=0, cross_tray=2,
                  wind_anchors=6, pressure=8.956827, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=6, panel_type=PanelType.Corner, ballast=13, link_tray=0, cross_tray=1,
                  wind_anchors=0, pressure=9.527003)
        ]

        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.788942, 0.628471, 0.523726, 0.350229],
                                 [0.605171, 0.475222, 0.396018, 0.255521],
                                 [0.402570, 0.341277, 0.243769, 0.186968],
                                 [0.310157, 0.226181, 0.161558, 0.114584],
                                 [0.228394, 0.153519, 0.176571, 0.085682],
                                 [0.211896, 0.161175, 0.123981, 0.085682],
                                 [0.179390, 0.153519, 0.109657, 0.085682],
                                 [0.251345, 0.182189, 0.182189, 0.116534],
                                 [0.521059, 0.291372, 0.291372, 0.127381],
                                 [0.375420, 0.253428, 0.126714, 0.085682],
                                 [0.179390, 0.153519, 0.109657, 0.085682]])
        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.819845)

        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_individual_ballast_block_uplift_greater_than_base_weight(self):
        expected = Panel(wind_zone=6, panel_type=PanelType.Middle, ballast=17, link_tray=1, cross_tray=1,
                         wind_anchors=0, pressure=8.810399)
        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=12, link_tray=1, cross_tray=1,
                         wind_anchors=0, pressure=8.867935)
        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=10, link_tray=0, cross_tray=2,
                         wind_anchors=7, pressure=7.891559, 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=7, link_tray=1, cross_tray=1,
                         wind_anchors=6, pressure=5.991164, 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=7, link_tray=1, cross_tray=0,
                         wind_anchors=1, pressure=4.692301)
        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=14, link_tray=0, cross_tray=2,
                         wind_anchors=7, pressure=7.776488, 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.819845)
        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), 2462.20, 2)

    def test_calculate_single_unit_pressure_on_roof(self):
        assert_almost_equal(self.subject.calculate_pressure_on_roof(90, 20, 59.1), 53.48, 2)
        assert_almost_equal(self.subject.calculate_pressure_on_roof(15, 14, 54.5), 7.61, 2)
