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 BallastCalculatorWhenSingleTiltAnd96CellTest(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.506920, 1.347220, 1.122684, 0.948795],
                                 [1.186590, 1.057265, 0.881054, 0.740239],
                                 [0.676809, 0.615810, 0.439865, 0.382935],
                                 [0.685883, 0.602310, 0.430222, 0.352224],
                                 [0.460258, 0.272456, 0.408684, 0.227532],
                                 [0.438829, 0.388352, 0.298733, 0.247999],
                                 [0.206406, 0.186191, 0.186191, 0.159779],
                                 [0.560763, 0.491939, 0.491939, 0.402013],
                                 [1.548721, 1.320138, 1.320138, 1.021469],
                                 [0.921232, 0.799827, 0.399914, 0.320599],
                                 [0.168333, 0.142587, 0.101848, 0.078692]])
        self.system_type = SystemType.singleTilt
        self.module_type = ModuleType.Cell96
        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': [5, 4, 2, 2, 1, 1, 0, 2, 6, 3, 0],
                'ballast blocks': [14, 10, 9, 10, 11, 10, 11, 2, 2, 9, 8],
                'pressure': ['12.81', '9.67', '8.90', '9.67', '10.43', '9.58', '10.34', '3.47', '3.57', '8.90', '8.05']
            },
            PanelType.NorthSouth: {
                'anchors': [5, 4, 2, 2, 0, 1, 0, 1, 5, 3, 0],
                'ballast blocks': [5, 2, 5, 5, 15, 7, 9, 13, 3, 2, 7],
                'pressure': ['5.74', '3.45', '5.74', '5.65', '13.37', '7.17', '8.79', '11.84', '4.22', '3.36', '7.17']
            },
            PanelType.EastWest: {
                'anchors': [4, 3, 1, 1, 1, 0, 0, 1, 5, 1, 0],
                'ballast blocks': [6, 6, 10, 9, 8, 16, 9, 13, 3, 7, 4],
                'pressure': ['6.62', '6.52', '9.57', '8.81', '8.05', '14.24', '8.81', '11.96', '4.24', '7.29', '4.93']
            },
            PanelType.Middle: {
                'anchors': [3, 2, 1, 0, 0, 0, 0, 1, 3, 0, 0],
                'ballast blocks': [11, 13, 6, 19, 12, 13, 8, 8, 15, 18, 3],
                'pressure': ['10.40', '11.92', '6.49', '16.59', '11.16', '11.92', '8.02', '8.02', '13.54', '15.73', '4.13']
            },
        }

        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 = [
            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=25, link_tray=0, cross_tray=2, wind_anchors=0, pressure=21.200142),
            Panel(wind_zone=6, panel_type=PanelType.NorthSouth, ballast=9, link_tray=0, cross_tray=1, wind_anchors=0, pressure=8.791605),
            Panel(wind_zone=7, panel_type=PanelType.EastWest, ballast=28, link_tray=2, cross_tray=2, wind_anchors=0, pressure=23.487752),
            Panel(wind_zone=0, panel_type=PanelType.Corner, ballast=14, link_tray=0, cross_tray=2, wind_anchors=5, pressure=12.810842, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=8, panel_type=PanelType.Middle, ballast=15, link_tray=1, cross_tray=2, wind_anchors=3, pressure=13.538805, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=9, panel_type=PanelType.NorthSouth, ballast=2, link_tray=0, cross_tray=0, wind_anchors=3, pressure=3.360677, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=6, panel_type=PanelType.Corner, ballast=11, link_tray=0, cross_tray=0, wind_anchors=0, pressure=10.337906),
            Panel(wind_zone=6, panel_type=PanelType.EastWest, ballast=9, link_tray=2, cross_tray=0, wind_anchors=0, pressure=8.812197)
        ]

        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),
        ]
        self.values.ballast_block_weight.return_value = 20

        expected_value = [
            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_value)

    def test_pressure_exceeding_max_system_pressure(self):
        self.values.ballast_block_weight.return_value = 20
        self.values.max_system_pressure.return_value = 25
        panels = [
            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 = array([
            Panel(wind_zone=7, panel_type=PanelType.Corner, ballast=2, link_tray=0, cross_tray=0, wind_anchors=2, pressure=3.473933),
            Panel(wind_zone=7, panel_type=PanelType.NorthSouth, ballast=13, link_tray=0, cross_tray=1, wind_anchors=1, pressure=11.842260),
            Panel(wind_zone=9, panel_type=PanelType.EastWest, ballast=22, link_tray=2, cross_tray=1, wind_anchors=0, pressure=18.819488),
            Panel(wind_zone=5, panel_type=PanelType.Corner, ballast=25, link_tray=0, cross_tray=2, wind_anchors=0, pressure=21.200142),
            Panel(wind_zone=6, panel_type=PanelType.Middle, ballast=8, link_tray=1, cross_tray=0, wind_anchors=0, pressure=8.015213),
            Panel(wind_zone=8, panel_type=PanelType.NorthSouth, ballast=3, link_tray=0, cross_tray=1, wind_anchors=5, pressure=4.215623, warnings=[PanelWarnings.MaxPsf]),
            Panel(wind_zone=6, panel_type=PanelType.Corner, ballast=11, link_tray=0, cross_tray=0, wind_anchors=0, pressure=10.337906),
        ])

        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(10, PanelType.Middle, 14, 1000, self.c_p_matrix,
                                                          30.155851).almost_equal(
            Panel(wind_zone=10, panel_type=PanelType.Middle, ballast=0, link_tray=0, cross_tray=0,
                  wind_anchors=0, pressure=1.842977),
            decimal=2)

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

    def test_ballast_blocks_exceeding_tray_capacity(self):
        assert self.subject.ballast_tray_and_anchor_count(0, PanelType.Corner, 20, 1000, self.c_p_matrix,
                                                          self.q_z).almost_equal(
            Panel(wind_zone=0, panel_type=PanelType.Corner, ballast=14, link_tray=0, cross_tray=2,
                  wind_anchors=5, pressure=12.810842, warnings=[PanelWarnings.MaxPsf]),
            decimal=2)
        assert self.subject.ballast_tray_and_anchor_count(8, PanelType.Middle, 20, 1000, self.c_p_matrix,
                                                          self.q_z).almost_equal(
            Panel(wind_zone=8, panel_type=PanelType.Middle, ballast=15, link_tray=1, cross_tray=2,
                  wind_anchors=3, pressure=13.538805, warnings=[PanelWarnings.MaxPsf]),
            decimal=2)

    def test_ballast_blocks_when_panel_exceeds_max_system_pressure(self):
        assert self.subject.ballast_tray_and_anchor_count(5, PanelType.Middle, 14, 10, self.c_p_matrix,
                                                          self.q_z).almost_equal(
            Panel(wind_zone=5, panel_type=PanelType.Middle, ballast=0, link_tray=0, cross_tray=0,
                  wind_anchors=1, pressure=1.842977),
            decimal=2)
        assert self.subject.ballast_tray_and_anchor_count(0, PanelType.Corner, 14, 10, self.c_p_matrix,
                                                          self.q_z).almost_equal(
            Panel(wind_zone=0, panel_type=PanelType.Corner, ballast=0, link_tray=0, cross_tray=0,
                  wind_anchors=6, pressure=1.948606, warnings=[PanelWarnings.MaxPsf]),
            decimal=2)

    def test_when_max_system_pressure_is_lower_than_base_weight_pressure(self):
        max_system_pressure = 0

        assert self.subject.ballast_tray_and_anchor_count(5, PanelType.Middle, 12, max_system_pressure, self.c_p_matrix,
                                                          self.q_z).almost_equal(
            Panel(wind_zone=5, panel_type=PanelType.Middle, ballast=0, link_tray=0, cross_tray=0,
                  wind_anchors=1, pressure=1.842977),
            decimal=2)

    def test_uplift(self):
        assert_almost_equal(self.subject.uplift(1.506920, self.q_z), 1857.49, 1)

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