import csv
import json
import unittest

import mockredis
from nose.tools import eq_
from numpy import array
from numpy.testing import assert_array_equal
import time

from helix.calculators.calculator import Calculator
from helix.constants.anchor_type import AnchorType
from helix.constants.exposure_category import ExposureCategory
from helix.constants.inverter_type import InverterType
from helix.constants.module_type import ModuleType
from helix.constants.panel_type import PanelType
from helix.constants.system_type import SystemType
from helix.models.coordinate import Coordinate
from helix.models.panel import Panel
from helix.models.sql.inverters import Inverter
from helix.models.sql.power_stations import PowerStation
from helix.models.sql.sites import Site
from helix.models.subarray import Subarray
from helix.store import Store
from helix.user_values import UserValues
from test.test_helpers import assert_array_is_close


class CalculatorTest(unittest.TestCase):
    def setUp(self):
        self.store = Store(mockredis.mock_redis_client(), "foo")

        self.site = Site(
            building_height=35,
            building_width=450,
            building_length=500,
            parapet_height=5,
            wind_speed=120,
            exposure_category='B',
            ballast_block_weight=14,
            max_psf=10,
            system_type=SystemType.dualTilt.value,
            module_type=ModuleType.Cell96.value,
            spectral_response=1.5,
            anchor_type=AnchorType.OMG_PowerGrip_Plus.value,
            seismic_importance_factor=1,
        )

        self.values = UserValues(self.store, self.site)

    def test_compute_subarray_summary_dual_tilt(self):
        self.site.system_type = SystemType.dualTilt.value
        with open('test/fixtures/input_dual_tilt_csv_for_bom.csv', 'r', newline='') as csv_file:
            csv_content = csv_file.read()
        self.site.cad_file = csv_content
        self.subject = Calculator(self.values)

        expected = [
            Subarray(subarray_number=1, origin=Coordinate(0, 0), required_seismic_anchors=436, weight=170698.330000,
                     start_row=0, size=863, column_counted_geometrically=False, row_counted_geometrically=False,
                     row_count=130.5, column_count=98)
        ]
        assert_array_is_close(self.subject.subarray_summary(), expected, decimal=0)

    def test_compute_subarray_summary_single_tilt(self):
        self.site.system_type = SystemType.singleTilt.value
        with open('test/fixtures/input_single_tilt_csv_for_bom.csv', 'r', newline='') as csv_file:
            csv_content = csv_file.read()
        self.site.cad_file = csv_content
        self.subject = Calculator(self.values)

        expected = [
            Subarray(subarray_number=1, origin=Coordinate(0, 0), required_seismic_anchors=11, weight=23804, start_row=0,
                     size=192, column_counted_geometrically=False, row_counted_geometrically=False,
                     row_count=18, column_count=20),
            Subarray(subarray_number=2, origin=Coordinate(0, 0), required_seismic_anchors=0, weight=13271,
                     start_row=192, size=96, column_counted_geometrically=False, row_counted_geometrically=False,
                     row_count=4, column_count=24),
            Subarray(subarray_number=3, origin=Coordinate(0, 0), required_seismic_anchors=110, weight=40328,
                     start_row=288, size=312, column_counted_geometrically=False, row_counted_geometrically=False,
                     row_count=25, column_count=29),
            Subarray(subarray_number=4, origin=Coordinate(0, 0), required_seismic_anchors=71, weight=23324,
                     start_row=600, size=168, column_counted_geometrically=False, row_counted_geometrically=False,
                     row_count=22, column_count=13),
            Subarray(subarray_number=5, origin=Coordinate(0, 0), required_seismic_anchors=74, weight=35170,
                     start_row=768, size=234, column_counted_geometrically=False, row_counted_geometrically=False,
                     row_count=26, column_count=14),
            Subarray(subarray_number=6, origin=Coordinate(0, 0), required_seismic_anchors=1, weight=11461,
                     start_row=1002, size=96, column_counted_geometrically=False, row_counted_geometrically=False,
                     row_count=10, column_count=12),
            Subarray(subarray_number=7, origin=Coordinate(0, 0), required_seismic_anchors=28, weight=10865,
                     start_row=1098, size=96, column_counted_geometrically=False, row_counted_geometrically=False,
                     row_count=8, column_count=16),
            Subarray(subarray_number=8, origin=Coordinate(0, 0), required_seismic_anchors=12, weight=10780,
                     start_row=1194, size=96, column_counted_geometrically=False, row_counted_geometrically=False,
                     row_count=15, column_count=17.5),
            Subarray(subarray_number=9, origin=Coordinate(0, 0), required_seismic_anchors=0, weight= 10989,
                     start_row=1290, size=96, column_counted_geometrically=False, row_counted_geometrically=False,
                     row_count=8, column_count=16),
            Subarray(subarray_number=10, origin=Coordinate(0, 0), required_seismic_anchors=0, weight=4901,
                     start_row=1386, size=48, column_counted_geometrically=False, row_counted_geometrically=False,
                     row_count=8, column_count=6),
            Subarray(subarray_number=11, origin=Coordinate(0, 0), required_seismic_anchors=0, weight=10766,
                     start_row=1434, size=96, column_counted_geometrically=False, row_counted_geometrically=False,
                     row_count=13, column_count=13),
            Subarray(subarray_number=12, origin=Coordinate(0, 0), required_seismic_anchors=0, weight=6350,
                     start_row=1530, size=54, column_counted_geometrically=False, row_counted_geometrically=False,
                     row_count=8, column_count=8),
            Subarray(subarray_number=13, origin=Coordinate(0, 0), required_seismic_anchors=0, weight=4239,
                     start_row=1584, size=48, column_counted_geometrically=False, row_counted_geometrically=False,
                     row_count=6, column_count=8)
        ]
        assert_array_is_close(self.subject.subarray_summary(), expected, decimal=0)

    def test_does_not_operate_on_panels_if_told_not_to(self):
        with open('test/fixtures/input_dual_tilt_coordinates_seismic_anchors.txt', 'r', newline='') as csv_file:
            csv_content = csv_file.read()
        self.site.cad_file = csv_content
        self.subject = Calculator(self.values, calculate_panel_data=False)

        for panel in self.subject.panels:
            assert panel.seismic_anchors is None
        assert self.subject.subarrays is None

    def test_compute_summary_values(self):
        self.site.building_width = 200
        self.site.building_height = 30
        self.site.building_length = 75.5
        self.site.wind_speed = 130
        self.site.exposure_category = ExposureCategory.C.value
        self.site.ballast_block_weight = 13
        self.site.parapet_height = 0
        self.site.max_psf = 10
        self.site.anchor_type = AnchorType.OMG_PowerGrip_Plus.value
        self.site.spectral_response = 1.02
        self.site.system_type = SystemType.singleTilt.value
        with open('test/fixtures/input_single_tilt.csv', 'r', newline='') as csv_file:
            csv_content = csv_file.read()
        self.site.cad_file = csv_content
        self.subject = Calculator(self.values)

        received_values = self.subject.summary_values()
        
        expected_values = array([
            {'label': 'Total System Weight (lbs)', 'value': 2694},
            {'label': 'Max PSF', 'value': 7.9},
            {'label': 'Avg PSF', 'value': 5.41},
            {'label': 'Total Anchors', 'value': 5},
            {'label': 'Total Ballast', 'value': 135},
            {'label': 'Max Possible System Weight', 'value': 2792.0},
            {'label': 'Max System Weight Ballast Block', 'value': 17},
            {'label': 'Seismic Anchor Max. Spacing', 'value': 12},
        ])
        assert_array_equal(received_values, expected_values)

    def test_compute_documentation_summary_values(self):
        self.site.building_width = 200
        self.site.building_height = 30
        self.site.building_length = 75.5
        self.site.wind_speed = 130
        self.site.exposure_category = ExposureCategory.C.value
        self.site.ballast_block_weight = 13
        self.site.parapet_height = 0
        self.site.max_psf = 10
        self.site.anchor_type = AnchorType.OMG_PowerGrip_Plus.value
        self.site.spectral_response = 1.02
        self.site.system_type = SystemType.singleTilt.value
        with open('test/fixtures/input_single_tilt.csv', 'r', newline='') as csv_file:
            csv_content = csv_file.read()
        self.site.cad_file = csv_content
        self.subject = Calculator(self.values)

        received_values = self.subject.documentation_summary_values()

        expected_values = {
            'total_system_weight': 2694,
            'max_psf': 7.9,
            'ave_psf': 5.41,
            'total_anchors': 5,
            'total_ballast': 135,
            'max_possible_system_weight': 2792.0,
            'max_system_weight_ballast_block': 17,
            'seismic_anchor_max_spacing': 12
        }
        eq_(received_values, expected_values)

    def test_recomputes_ballast_after_seismic_calculations(self):
        self.site.building_width = 200
        self.site.building_height = 30
        self.site.building_length = 75.5
        self.site.wind_speed = 110
        self.site.exposure_category = ExposureCategory.C.value
        self.site.ballast_block_weight = 13
        self.site.parapet_height = 0
        self.site.max_psf = 10
        self.site.anchor_type = AnchorType.EcoFasten.value
        self.site.spectral_response = 1.0
        self.site.system_type = SystemType.dualTilt.value
        with open('test/fixtures/input_dual_tilt_coordinates.csv', 'r', newline='') as csv_file:
            csv_content = csv_file.read()
        self.site.cad_file = csv_content
        self.subject = Calculator(self.values)

        required_seismic_anchors = sum(subarray.required_seismic_anchors for subarray in self.subject.subarrays)

        eq_(required_seismic_anchors, 2)

        self.site.spectral_response = 1.7
        self.subject = Calculator(self.values)

        required_seismic_anchors = sum(subarray.required_seismic_anchors for subarray in self.subject.subarrays)

        eq_(required_seismic_anchors, 6)

    def test_uses_user_provided_seismic_anchor_placement_if_available(self):
        self.site.building_width = 200
        self.site.building_height = 30
        self.site.building_length = 75.5
        self.site.wind_speed = 110
        self.site.exposure_category = ExposureCategory.C.value
        self.site.ballast_block_weight = 13
        self.site.parapet_height = 0
        self.site.max_psf = 10
        self.site.anchor_type = AnchorType.EcoFasten.value
        self.site.spectral_response = 1.0
        self.site.system_type = SystemType.dualTilt.value
        with open('test/fixtures/input_dual_tilt_coordinates.csv', 'r', newline='') as csv_file:
            csv_content = csv_file.read()
        self.site.cad_file = csv_content
        self.store.set('user_override_seismic_anchors', json.dumps([
            {'panel_id': 1, 'seismic_anchors': 0},
            {'panel_id': 2, 'seismic_anchors': 1},
            {'panel_id': 3, 'seismic_anchors': 1},
            {'panel_id': 4, 'seismic_anchors': 0},
            {'panel_id': 5, 'seismic_anchors': 1},
            {'panel_id': 6, 'seismic_anchors': 0},
            {'panel_id': 7, 'seismic_anchors': 0},
            {'panel_id': 8, 'seismic_anchors': 1},
            {'panel_id': 9, 'seismic_anchors': 1},
            {'panel_id': 10, 'seismic_anchors': 1},
            {'panel_id': 11, 'seismic_anchors': 1},
            {'panel_id': 12, 'seismic_anchors': 1},
            {'panel_id': 13, 'seismic_anchors': 1},
            {'panel_id': 14, 'seismic_anchors': 1},
            {'panel_id': 15, 'seismic_anchors': 1},
            {'panel_id': 16, 'seismic_anchors': 0},
        ]))
        self.subject = Calculator(self.values)

        eq_(sum(panel.seismic_anchors for panel in self.subject.panels), 11)

        expected = [
            Panel(seismic_anchors=0, ballast=0),
            Panel(seismic_anchors=1, ballast=0),
            Panel(seismic_anchors=1, ballast=0),
            Panel(seismic_anchors=0, ballast=0),
            Panel(seismic_anchors=1, ballast=0),
            Panel(seismic_anchors=0, ballast=10),
            Panel(seismic_anchors=0, ballast=10),
            Panel(seismic_anchors=1, ballast=0),
            Panel(seismic_anchors=1, ballast=0),
            Panel(seismic_anchors=1, ballast=0),
            Panel(seismic_anchors=1, ballast=0),
            Panel(seismic_anchors=1, ballast=0),
            Panel(seismic_anchors=1, ballast=0),
            Panel(seismic_anchors=1, ballast=0),
            Panel(seismic_anchors=1, ballast=0),
            Panel(seismic_anchors=0, ballast=0)
        ]
        for idx, received_panel in enumerate(self.subject.panels):
            received_seismic = received_panel.seismic_anchors
            expected_seismic = expected[idx].seismic_anchors
            eq_(received_seismic, expected_seismic,
                "Panel at index %d had %d seismic anchors, expected %d" % (
                idx, received_seismic, expected_seismic))
            received_ballast = received_panel.ballast
            expected_ballast = expected[idx].ballast

            eq_(received_ballast, expected_ballast,
                "Panel at index %d had %d ballast, expected %d" % (
                idx, received_ballast, expected_ballast))

        eq_(sum(subarray.required_seismic_anchors for subarray in self.subject.subarrays), 2)

    def test_get_computed_csv_columns(self):
        self.site.system_type = SystemType.dualTilt.value

        self.subject = Calculator(self.values)

        panels = [
            Panel(panel_type=PanelType.Corner, link_tray=3),
            Panel(panel_type=PanelType.NorthSouth, link_tray=3),
            Panel(panel_type=PanelType.EastWest, link_tray=3),
            Panel(panel_type=PanelType.Middle, link_tray=3),
        ]
        self.subject.panels = panels

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

        eq_(self.subject.get_computed_csv_columns(), expected)

    def test_computes_bom(self):
        self.site.building_width = 200
        self.site.building_height = 30
        self.site.building_length = 75.5
        self.site.wind_speed = 125
        self.site.exposure_category = ExposureCategory.C.value
        self.site.ballast_block_weight = 13
        self.site.parapet_height = 0
        self.site.max_psf = 10
        self.site.anchor_type = AnchorType.OMG_PowerGrip_Plus.value
        self.site.spectral_response = 1

        self.site.system_type = SystemType.dualTilt.value
        with open('test/fixtures/input_dual_tilt_csv_for_bom.csv', 'r', newline='') as csv_file:
            csv_content = csv_file.read()
        self.site.cad_file = csv_content
        self.subject = Calculator(self.values)

        with open('test/fixtures/expected_dual_tilt_ebom.csv', 'r', newline='') as expected_file:
            expected_csv = expected_file.read()
        reader = csv.reader(expected_csv.splitlines(), dialect='excel-tab')
        expected = array([row for row in reader])
        print("EXPECTED")
        print(expected)

        power_station_1 = PowerStation(
            description='1',
            quantity=1,
            ac_run_length=55,
            inverters=[
                Inverter(
                    model=str(InverterType.SMA.MODEL_12KW.value),
                    strings_per_inverter=2
                ),
                Inverter(
                    model=str(InverterType.SMA.MODEL_15KW.value),
                    strings_per_inverter=5
                )
            ]
        )
        power_station_2 = PowerStation(
            description='2',
            quantity=3,
            ac_run_length=89,
            inverters=[
                Inverter(
                    model=str(InverterType.SMA.MODEL_20KW.value),
                    strings_per_inverter=6
                ),
                Inverter(
                    model=str(InverterType.SMA.MODEL_24KW.value),
                    strings_per_inverter=8
                ),
                Inverter(
                    model=str(InverterType.SMA.MODEL_24KW.value),
                    strings_per_inverter=8
                ),
                Inverter(
                    model=str(InverterType.SMA.MODEL_24KW.value),
                    strings_per_inverter=7
                )
            ]
        )

        self.site.power_stations = [power_station_1, power_station_2]
        print("COMPUTE")
        print(self.subject.compute_bom())
        assert_array_equal(self.subject.compute_bom(), expected)

    def test_documentation_bom(self):
        self.site.building_width = 200
        self.site.building_height = 30
        self.site.building_length = 75.5
        self.site.wind_speed = 125
        self.site.exposure_category = ExposureCategory.C.value
        self.site.ballast_block_weight = 13
        self.site.parapet_height = 0
        self.site.max_psf = 10
        self.site.anchor_type = AnchorType.OMG_PowerGrip_Plus.value
        self.site.spectral_response = 1

        self.site.power_stations = [
            PowerStation(
                quantity=1,
                ac_run_length=10,
                description='Test',
                id=30,
                inverters=[
                    Inverter(
                        model=str(InverterType.SMA.MODEL_12KW.value),
                        sunshade=True,
                        dc_switch=True,
                        strings_per_inverter=4
                    )
                ]
            )
        ]

        self.site.system_type = SystemType.dualTilt.value
        with open('test/fixtures/input_dual_tilt_csv_for_bom.csv', 'r', newline='') as csv_file:
            csv_content = csv_file.read()
        self.site.cad_file = csv_content
        self.subject = Calculator(self.values)

        expected = [('521794', 196), 
                    ('521795', 196), 
                    ('514056', 1000), 
                    ('modules', 1726), 
                    ('513843', 263), 
                    ('anchors', 263), 
                    ('518477', 275), 
                    ('513833', 670), 
                    ('513844', 214), 
                    ('ballast', 6777), 
                    ('515928', 261), 
                    ('517871', 139), 
                    ('514057', 1000), 
                    ('515063', 4000), 
                    ('512200', 3480), 
                    ('514265', 179), 
                    ('513303', 1), 
                    ('512660', 2), 
                    ('512661', 2), 
                    ('512662', 4), 
                    ('512663', 2), 
                    ('518331', 2), 
                    ('518058', 2), 
                    ('104813', 50), 
                    ('107551', 50), 
                    ('514865', 50), 
                    ('106925', 50), 
                    ('523921', 1), 
                    ('514438', 2), 
                    ('514437', 2), 
                    ('512910', 1), 
                    ('805615', 2), 
                    ('521031', 2), 
                    ('512575', 1), 
                    ('514698', 1), 
                    ('513007', 50), 
                    ('114961', 50), 
                    ('107549', 100), 
                    ('107586', 100), 
                    ('512021', 9), 
                    ('512199', 140), 
                    ('512511', 196),
                    ('512510', 196), 
                    ('515929', 0), 
                    ('514477', 2), 
                    ('515059', 2), 
                    ('514697', 1), 
                    ('513831', 0), 
                    ('513836', 0), 
                    ('520301', 0), 
                    ('520302', 0), 
                    ('520303', 0), 
                    ('520306', 0), 
                    ('513832', 0), 
                    ('514435', 0), 
                    ('514436', 0), 
                    ('514439', 0), 
                    ('514440', 0), 
                    ('523922', 0), 
                    ('523923', 0), 
                    ('523924', 0), 
                    ('513299', 0), 
                    ('513301', 0), 
                    ('514478', 0),
                    ('513300', 0), 
                    ('513302', 0), 
                    ('513304', 0), 
                    ('519008', 0), 
                    ('518059', 0), 
                    ('105317', 0), 
                    ('111147', 0), 
                    ('107538', 0), 
                    ('516045', 0), 
                    ('516043', 0), 
                    ('513586', 0), 
                    ('507985', 0), 
                    ('522020', 0), 
                    ('521798', 0), 
                    ('521797', 0), 
                    ('521363', 0), 
                    ('517463', 0)]

        print("===")
        print(self.subject.documentation_bom())
        assert_array_equal(sorted(self.subject.documentation_bom()), sorted(expected))

        # Performance Tests

    def test_performance_with_1000ish_module_site(self):
        store = Store(mockredis.mock_redis_client(), "foo")

        site = Site()

        values = UserValues(store, site)

        site.building_width = 200
        site.building_height = 30
        site.building_length = 75.5
        site.wind_speed = 110
        site.exposure_category = ExposureCategory.C.value
        site.ballast_block_weight = 13
        site.parapet_height = 0
        site.max_psf = 10
        site.anchor_type = AnchorType.EcoFasten.value
        site.spectral_response = 1.0
        site.system_type = SystemType.dualTilt.value
        site.module_type = ModuleType.Cell96.value
        with open('test/fixtures/input_dual_tilt_coordinates_seismic_anchors.txt', 'r', newline='') as csv_file:
            csv_content = csv_file.read()
        site.cad_file = csv_content

        run_count = 1
        runtimes = []

        for _ in range(run_count):
            start = time.clock()
            Calculator(values)  # Do all calculations, including seismic anchor placement
            end = time.clock()
            runtimes.append(end - start)
        print("%d runs took %f time on average" % (run_count, sum(runtimes) / run_count))

    def test_performance_with_giant_module_site(self):
        store = Store(mockredis.mock_redis_client(), "foo")

        site = Site()

        values = UserValues(store, site)

        site.building_width = 200
        site.building_height = 30
        site.building_length = 75.5
        site.wind_speed = 110
        site.exposure_category = ExposureCategory.C.value
        site.ballast_block_weight = 13
        site.parapet_height = 0
        site.max_psf = 10
        site.anchor_type = AnchorType.EcoFasten.value
        site.spectral_response = 1.0
        site.system_type = SystemType.singleTilt.value
        site.module_type = ModuleType.Cell96.value
        with open('test/fixtures/input_single_tilt_96_cell_giant_site.txt', 'r', newline='') as csv_file:
            csv_content = csv_file.read()
        site.cad_file = csv_content

        run_count = 1
        runtimes = []

        for _ in range(run_count):
            start = time.clock()
            Calculator(values)  # Do all calculations, including seismic anchor placement
            end = time.clock()
            runtimes.append(end - start)
        import sys
        print("%d runs took %f time on average" % (run_count, sum(runtimes) / run_count), file=sys.stderr)
