from datetime import timedelta
from clearml import Task
from src.policies.simple_baseline import BaselinePolicy
from src.policies.PolicyEvaluator import PolicyEvaluator
import numpy as np
import pandas as pd
from tqdm import tqdm
import torch


class YesterdayBaselinePolicyEvaluator(PolicyEvaluator):
    def __init__(self, baseline_policy: BaselinePolicy, task: Task = None):
        super(YesterdayBaselinePolicyEvaluator, self).__init__(baseline_policy, task)

    def evaluate_for_date(
        self,
        date,
        charge_thresholds=np.arange(-100, 250, 25),
        discharge_thresholds=np.arange(-100, 250, 25),
    ):

        real_imbalance_prices = self.get_imbanlance_prices_for_date(date.date())
        yesterday_imbalance_prices = self.get_imbanlance_prices_for_date(
            date.date() - timedelta(days=1)
        )
        yesterday_imbalance_prices = torch.tensor(
            np.array([yesterday_imbalance_prices]), device="cpu"
        )

        for penalty in self.penalties:
            yesterday_charge_thresholds, yesterday_discharge_thresholds = (
                self.baseline_policy.get_optimal_thresholds(
                    yesterday_imbalance_prices,
                    charge_thresholds,
                    discharge_thresholds,
                    penalty,
                )
            )

            yesterday_profit, yesterday_charge_cycles = self.baseline_policy.simulate(
                torch.tensor([[real_imbalance_prices]]),
                torch.tensor([yesterday_charge_thresholds.mean(axis=0)]),
                torch.tensor([yesterday_discharge_thresholds.mean(axis=0)]),
            )

            self.profits.append(
                [
                    date,
                    penalty,
                    yesterday_profit[0][0].item(),
                    yesterday_charge_cycles[0][0].item(),
                    yesterday_charge_thresholds.mean(axis=0).item(),
                    yesterday_discharge_thresholds.mean(axis=0).item(),
                ]
            )

    def evaluate_test_set(self, data_processor):

        if data_processor:
            filtered_dates = []
            _, test_loader = data_processor.get_dataloaders()
            for date in self.dates:
                try:
                    test_loader.dataset.get_idx_for_date(date.date())
                    filtered_dates.append(date)
                except:
                    pass
            self.dates = filtered_dates

        self.profits = []
        for date in tqdm(self.dates):
            try:
                self.evaluate_for_date(date)
            except Exception as e:
                print(e)
                pass

        self.profits = pd.DataFrame(
            self.profits,
            columns=[
                "Date",
                "Penalty",
                "Profit",
                "Charge Cycles",
                "Charge Threshold",
                "Discharge Threshold",
            ],
        )