# modeling an energy supplier for the purposes of peak shaving

import numpy as np
import pandas as pd
import datetime
import unittest


class Supplier:
    # price [HUF/kWh]
    # peak_demand kW
    # surcharge_per_kw [HUF/kW for each 15 minute timeframe]
    def __init__(self, price):
        self.hourly_prices = np.ones(168) * price
        self.peak_demand = np.inf # no demand_charge by default
        self.surcharge_per_kw = 0

    # start and end are indices of hours starting from Monday 00:00.
    def set_price_for_interval(self, start, end, price):
        self.hourly_prices[start:end] = price

    # start and end are indices of hours of the day. for each day, this interval is set to price
    def set_price_for_daily_interval(self, start, end, price):
        for day in range(7):
            h = day * 24
            self.set_price_for_interval(h + start, h + end, price)

    def set_price_for_daily_interval_on_workdays(self, start, end, price):
        for day in range(5):
            h = day * 24
            self.set_price_for_interval(h + start, h + end, price)

    def set_demand_charge(self, peak_demand, surcharge_per_kw):
        self.peak_demand = peak_demand # [kW]
        # the HUF charged per kW of demand exceeding peak_demand during a 15 minutes timeframe.
        self.surcharge_per_kw = surcharge_per_kw # [HUF/kW]

    @staticmethod
    def hour_of_date(date):
        hours_since_midnight = (date - datetime.datetime(date.year, date.month, date.day, 0, 0, 0)).total_seconds() / 3600
        # weekday() calculates from sunday morning:
        hungarian_weekday = (date.weekday() + 0) % 7
        hours_elapsed_in_previous_days = hungarian_weekday * 24
        return int(hours_since_midnight) + hours_elapsed_in_previous_days

    def price(self, date):
        return self.hourly_prices[self.hour_of_date(date)]

    # demand is the maximum demand in kW during a 15 minute interval
    def demand_charge(self, demand):
        if demand <= self.peak_demand:
            return 0.0
        else:
            return (demand - self.peak_demand) * self.surcharge_per_kw

    # demand_series is pandas series indexed by time.
    # during each time step demand [kW] is assumed to be constant.
    def fee(self, demand_series):
        prices = [self.price(date) for date in demand_series.index]
        prices_series = pd.Series(data=prices, index=demand_series.index)
        # prices are HUF/kWh, demand is kW. note the missing h.

        step_in_hour = demand_series.index.freq.n / 60 # [hour], the length of a time step.
        # for each step the product tells the fee IF the step was 1 hour long. it's actually step_in_hour long:
        consumption_charge = demand_series.dot(prices_series) * step_in_hour

        # 15 minutes (the demand charge calculation interval) should be a multiple of the series time step.
        assert 15 % demand_series.index.freq.n == 0
        time_steps_per_demand_charge_evaluation = 15 // demand_series.index.freq.n
        # fifteen_minute_peaks [kW] tells the maximum demand in a 15 minutes timeframe:
        fifteen_minute_peaks = demand_series.resample('15T').max()
        demand_charges = [self.demand_charge(demand) for demand in fifteen_minute_peaks]
        total_demand_charge = sum(demand_charges)
        return consumption_charge + total_demand_charge


class TestSupplier(unittest.TestCase):

    def setUp(self):
        self.constant_price = 10
        self.supplier = Supplier(self.constant_price)

    def test_hourly_prices(self):
        expected_hourly_prices = np.ones(168) * self.constant_price
        self.assertTrue(np.array_equal(self.supplier.hourly_prices, expected_hourly_prices))

    def test_set_price_for_interval(self):
        self.supplier.set_price_for_interval(0, 24, 20)
        expected_hourly_prices = np.ones(168) * self.constant_price
        expected_hourly_prices[0:24] = 20
        self.assertTrue(np.array_equal(self.supplier.hourly_prices, expected_hourly_prices))

    def test_price(self):
        increased_price = 20
        self.supplier.set_price_for_interval(0, 24, increased_price)

        date = datetime.datetime(2023, 4, 30, 12, 0, 0)  # Sunday noon
        expected_price = self.constant_price
        self.assertEqual(self.supplier.price(date), expected_price)

        date = datetime.datetime(2023, 5, 1, 12, 0, 0)  # Monday noon
        expected_price = increased_price
        self.assertEqual(self.supplier.price(date), expected_price)

        date = datetime.datetime(2023, 5, 2, 12, 0, 0)  # Tuesday noon
        expected_price = self.constant_price
        self.assertEqual(self.supplier.price(date), expected_price)

    def test_fee(self):
        start = pd.Timestamp('2021-04-28')
        end = start + pd.Timedelta(days=1)
        freq = '5T' # 5 minutes
        time_index = pd.date_range(start=start, end=end, freq=freq, inclusive='left')
        constant_demand = 100
        demand_in_kw = [constant_demand] * len(time_index)

        demand_series = pd.Series(data=demand_in_kw, index=time_index)
        # 24 because it's a 24 hour period with constant demand:
        self.assertEqual(self.supplier.fee(demand_series), constant_demand * 24 * self.constant_price)

        extreme_demand = 1000
        demand_series[12:24] = extreme_demand # in second hour we set extreme demand.

        expected_fee = (constant_demand * 23 + extreme_demand) * self.constant_price
        self.assertEqual(self.supplier.fee(demand_series), expected_fee)

        # now the (1000-500) kW above 500 kW is surcharged for (1000-500 kW) * 10 HUF/kW/15mins, for 1 hour,
        # that is 500*10*4=20000 demand_charge.
        self.supplier.set_demand_charge(peak_demand=500, surcharge_per_kw=10)
        expected_fee += 20000
        self.assertEqual(self.supplier.fee(demand_series), expected_fee)



if __name__ == '__main__':
    unittest.main()