from clearml import Task
from src.policies.simple_baseline import BaselinePolicy
from src.policies.baselines.YesterdayBaselinePolicyExecutor import (
    YesterdayBaselinePolicyEvaluator,
)
import torch
import numpy as np


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

    def evaluate_for_date(
        self,
        date,
        charge_thresholds=np.arange(-300, 300, 5),
        discharge_thresholds=np.arange(-300, 300, 5),
        penalty: int = 0,
        current_state_of_charge=0.0,
    ):

        real_imbalance_prices = self.get_imbanlance_prices_for_date(date.date())

        real_imbalance_prices_tensor = torch.tensor(
            np.array([real_imbalance_prices]), device="cpu"
        )

        best_charge_thresholds, best_discharge_thresholds = (
            self.baseline_policy.get_optimal_thresholds(
                real_imbalance_prices_tensor,
                charge_thresholds,
                discharge_thresholds,
                penalty,
                battery_state_of_charge=current_state_of_charge,
            )
        )

        best_profit, best_charge_cycles, new_state_of_charge = (
            self.baseline_policy.simulate(
                torch.tensor([[real_imbalance_prices]]),
                torch.tensor([best_charge_thresholds.mean(axis=0)]),
                torch.tensor([best_discharge_thresholds.mean(axis=0)]),
                battery_state_of_charge=torch.tensor([current_state_of_charge]),
            )
        )

        return (
            best_profit[0][0].item(),
            best_charge_cycles[0][0].item(),
            new_state_of_charge.squeeze(0).item(),
            best_charge_thresholds.mean(axis=0),
            best_discharge_thresholds.mean(axis=0),
        )