import numpy as np
import pandas as pd
import copy
from enum import IntEnum

# it's really just a network pricing model
from supplier import Supplier
from data_processing import read_datasets, add_production_field, interpolate_and_join, Parameters


STEPS_PER_HOUR = 12


# mock model
class BatteryModel:
    def __init__(self, capacity, efficiency=0.95):
        self.soc = 0
        self.capacity = capacity # kWh
        self.efficiency = efficiency

    def discharge(self, target):
        assert 0 <= self.soc <= 1
        if self.soc >= target:
            self.soc -= target
            amount = target
        else:
            amount = self.soc
            self.soc = 0
        assert 0 <= self.soc <= 1
        return amount

    # not very refined, whatevs.
    def charge(self, target):
        assert 0 <= self.soc <= 1
        target_to_add = target * self.efficiency
        if self.soc <= 1 - target_to_add:
            self.soc += target_to_add
            could_take = target
        else:
            could_take = (1 - self.soc) / self.efficiency
            self.soc = 1
        assert 0 <= self.soc <= 1
        return could_take


class Decision(IntEnum):
    # use solar to satisfy consumption,
    # and if it is not enough, use network.
    # BESS is not discharged in this mode,
    # but might be charged if solar has surplus.
    PASSIVE = 0

    # use the battery if possible and necessary.
    # the possible part means that there's charge in it,
    # and the necessary part means that the consumption
    # is not already covered by solar.
    DISCHARGE = 1

    # use the network to charge the battery
    # this is similar to PASSIVE, but forces the
    # BESS to be charged, even if solar does not cover
    # the whole of consumption plus BESS.
    NETWORK_CHARGE = 2


# mock class as usual
# output_window_size is not yet used, always decides one timestep.
class Decider:
    def __init__(self):
        self.parameters = None
        self.input_window_size = STEPS_PER_HOUR * 24 # day long window.
        self.output_window_size = STEPS_PER_HOUR # only output decisions for the next hour
        self.random_seed = 0

    # prod_cons_pred is a dataframe starting at now, containing
    # fields Production and Consumption.
    # this function does not mutate its inputs.
    # battery_model is just queried for capacity and current soc.
    # the method returns a pd.Series of Decisions as integers.
    def decide(self, prod_pred, cons_pred, battery_model):
        # assert len(prod_pred) == len(cons_pred) == self.input_window_size
        self.random_seed += 1
        self.random_seed %= 3 # dummy rotates between 3 Decisions
        return self.random_seed
        # dummy decider always says DISCHARGE:
        # return pd.Series([Decision.DISCHARGE] * self.output_window_size, dtype=int)


# even mock-er class than usual.
# knows the future in advance, so it predicts it very well.
# it's also unrealistic in that it takes row index instead of date.
class DummyPredictor:
    def __init__(self, series):
        self.series = series

    def predict(self, indx, window_size):
        prediction = self.series.iloc[indx: indx + window_size]
        return prediction


# this function does not mutate its inputs.
# it makes a clone of battery_model and modifies that.
# TODO even in a first mockup version, parameters should come from a single
# place, not some from the parameters dataclass and some from the battery_model.
def simulator(battery_model, supplier, prod_cons, prod_predictor, cons_predictor, decider, parameters):
    battery_model = copy.copy(battery_model)

    demand_np = prod_cons['Consumption'].to_numpy()
    production_np = prod_cons['Production'].to_numpy()
    assert len(demand_np) == len(production_np)
    step_in_minutes = prod_cons.index.freq.n
    print(step_in_minutes)
    assert step_in_minutes == 5

    print("Simulating for", len(demand_np), "time steps. Each step is", step_in_minutes, "minutes.")
    soc_series = []
    # by convention, we only call end user demand, demand,
    # and we only call end user consumption, consumption.
    # in our simple model, demand is always satisfied, hence demand=consumption.
    # BESS demand is called charge.
    consumption_from_solar_series = [] # demand satisfied by solar production
    consumption_from_network_series = [] # full demand satisfied by network, also includes consumption_from_network_to_bess.
    consumption_from_bess_series = [] # demand satisfied by BESS
    consumption_from_network_to_bess_series = [] # network used to charge BESS, also included in consumption_from_network.
    # the previous three must sum to demand_series.

    # power taken from solar by BESS.
    # note: in inital mock version power is never taken from network by BESS.
    charge_of_bess_series = []
    discarded_production_series = [] # solar power thrown away

    # 1 is not nominal but targeted (healthy) maximum charge.
    # we start with an empty battery, but not emptier than what's healthy for the batteries.

    # For the sake of simplicity 0 <= soc <=1
    # soc=0 means battery is emptied till it's 20% and soc=1 means battery is charged till 80% of its capacity
    # soc = 1 - maximal_depth_of_discharge
    # and will use only maximal_depth_of_discharge percent of the real battery capacity

    max_cap_of_battery = parameters.bess_capacity * parameters.maximal_depth_of_discharge

    time_interval = step_in_minutes / 60 # amount of time step in hours
    for i, (demand, production) in enumerate(zip(demand_np, production_np)):
        cap_of_battery = battery_model.soc * max_cap_of_battery
        # these five are modified on the appropriate codepaths:
        consumption_from_solar = 0
        consumption_from_bess = 0
        consumption_from_network = 0
        discarded_production = 0
        network_used_to_charge = 0

        unsatisfied_demand = demand
        remaining_production = production
        
        assert parameters.bess_present
        is_battery_charged_enough = battery_model.soc > 0
        is_battery_chargeable = battery_model.soc < 1.0

        prod_prediction = prod_predictor.predict(i, decider.input_window_size)
        cons_prediction = cons_predictor.predict(i, decider.input_window_size)
        decision = decider.decide(prod_prediction, cons_prediction, battery_model)

        production_used_to_charge = 0
        if unsatisfied_demand >= remaining_production:
            # all goes to demand
            consumption_from_solar = remaining_production
            unsatisfied_demand -= consumption_from_solar
            remaining_production = 0
            # we try to cover the rest from BESS
            if unsatisfied_demand > 0:
                if is_battery_charged_enough and decision == Decision.DISCHARGE:
                    # battery capacity is limited!
                    if cap_of_battery >= unsatisfied_demand * time_interval:
                        consumption_from_bess = unsatisfied_demand
                        unsatisfied_demand = 0
                        cap_of_battery -= consumption_from_bess * time_interval
                        soc = cap_of_battery / max_cap_of_battery
                    else:
                        discharge_of_bess = cap_of_battery / time_interval
                        discharge = min(parameters.bess_discharge, discharge_of_bess)
                        consumption_from_bess = discharge
                        unsatisfied_demand -= consumption_from_bess
                        cap_of_battery -= consumption_from_bess * time_interval
                        soc = cap_of_battery / max_cap_of_battery
                        consumption_from_network = unsatisfied_demand
                        unsatisfied_demand = 0
                else:
                    # we cover the rest from network
                    consumption_from_network = unsatisfied_demand
                    unsatisfied_demand = 0
        else:
            # demand fully satisfied by production
            consumption_from_solar = unsatisfied_demand
            remaining_production -= unsatisfied_demand
            unsatisfied_demand = 0
            if remaining_production > 0:
                # exploitable production still remains:
                if is_battery_chargeable:
                    # we try to specify the BESS modell
                    if parameters.bess_charge <= remaining_production :
                        energy = parameters.bess_charge * time_interval
                        remaining_production = remaining_production - parameters.bess_charge
                        production_used_to_charge = parameters.bess_charge
                    else :
                        production_used_to_charge = remaining_production
                        energy = remaining_production * time_interval
                        remaining_production = 0
                    cap_of_battery += energy
                    soc = cap_of_battery / max_cap_of_battery

        discarded_production = remaining_production

        if decision == Decision.NETWORK_CHARGE:
            # there are two things that can limit charging at this point,
            # one is a distance-like, the other is a velocity-like quantity.
            # 1. the battery is fully charged
            # 2. the battery is being charged from solar, no bandwidth left.
            is_battery_chargeable = battery_model.soc < 1.0
            remaining_network_charge = parameters.bess_charge - production_used_to_charge
            if is_battery_chargeable:
                # we try to specify the BESS modell
                if parameters.bess_charge <= remaining_network_charge :
                    energy = parameters.bess_charge * time_interval
                    remaining_network_charge = remaining_network_charge - parameters.bess_charge
                    network_used_to_charge = parameters.bess_charge
                else :
                    network_used_to_charge = remaining_production
                    energy = remaining_production * time_interval
                cap_of_battery += energy
                soc = cap_of_battery / max_cap_of_battery
            consumption_from_network += network_used_to_charge

        soc_series.append(battery_model.soc)
        consumption_from_solar_series.append(consumption_from_solar)
        consumption_from_network_series.append(consumption_from_network)
        consumption_from_network_to_bess_series.append(network_used_to_charge)
        consumption_from_bess_series.append(consumption_from_bess)
        discarded_production_series.append(discarded_production)

    soc_series = np.array(soc_series)
    consumption_from_solar_series = np.array(consumption_from_solar_series)
    consumption_from_network_series = np.array(consumption_from_network_series)
    consumption_from_bess_series = np.array(consumption_from_bess_series)
    discarded_production_series = np.array(discarded_production_series)

    results = pd.DataFrame({'soc_series': soc_series, 'consumption_from_solar': consumption_from_solar_series,
                            'consumption_from_network': consumption_from_network_series,
                            'consumption_from_bess': consumption_from_bess_series,
                            'consumption_from_network_to_bess_series': consumption_from_network_to_bess_series,
                            'discarded_production': discarded_production_series,
                            'Consumption': prod_cons['Consumption'],
                            'Production': prod_cons['Production']
                            })
    results = results.set_index(prod_cons.index)
    return results


def main():
    battery_model = BatteryModel(capacity=200, efficiency=0.95)

    supplier = Supplier(price=100) # Ft/kWh
    supplier.set_price_for_interval(9, 17, 150) # nine-to-five increased price.

    parameters = Parameters()

    met_2021_data, cons_2021_data = read_datasets()
    add_production_field(met_2021_data, parameters)
    all_2021_data = interpolate_and_join(met_2021_data, cons_2021_data)

    prod_predictor = DummyPredictor(pd.Series(all_2021_data['Production']))
    cons_predictor = DummyPredictor(pd.Series(all_2021_data['Consumption']))

    decider = Decider()

    results = simulator(battery_model, supplier, all_2021_data, prod_predictor, cons_predictor, decider, parameters)


if __name__ == '__main__':
    main()