import cvxpy as cp
import numpy as np


def solve_battery_optimization_hourly(
    hourly_prices,      # Array of prices for each hour [0, 1, ..., n-1]
    initial_B,        # Current state of charge (MWh)
    max_capacity=1.0,   # MWh
    max_rate=1.0        # MW (+ve discharge / -ve charge)
):
    """
    Solves the battery scheduling optimization problem assuming hourly steps. We want to decide at the start of each hour t=0..n-1
    how much power to buy/sell (P_net_t) and therefore the state of charge at the start of each next hour (B_t+1).

    Args:
        hourly_prices: Prices (€/MWh) for each hour t=0..n-1.
        initial_B: The state of charge at the beginning (time t=0).
        max_capacity: Maximum battery energy capacity (MWh).
        max_rate: Maximum charge/discharge power rate (MW).

    Returns:
        Tuple: (status, optimal_profit, power_schedule, B_schedule)
               Returns (status, None, None, None) if optimization fails.
    """
    n_hours = len(hourly_prices)

    # --- CVXPY Variables ---
    # Power flow for each hour t=0..n-1 (-discharge, +charge)
    P = cp.Variable(n_hours, name="Power_Flow_MW")
    # State of charge at the START of each hour t=0..n (B[t] is B at hour t)
    B = cp.Variable(n_hours + 1, name="State_of_Charge_MWh")

    # --- Objective Function ---
    # Profit = sum(price[t] * Power[t])
    prices = np.array(hourly_prices)
    profit = prices @ P # Equivalent to cp.sum(cp.multiply(prices, P)) / prices.dot(P)
    objective = cp.Maximize(profit)

    # --- Constraints ---
    constraints = []

    # 1. Initial B
    constraints.append(B[0] == initial_B)

    # 2. B Dynamics: B[t+1] = B[t] - P[t] * 1 hour
    constraints.append(B[1:] == B[:-1] + P)

    # 3. Power Rate Limits: -max_rate <= P[t] <= max_rate
    constraints.append(cp.abs(P) <= max_rate)

    # 4. B Limits: 0 <= B[t] <= max_capacity (applies to B[0]...B[n])
    constraints.append(B >= 0)
    constraints.append(B <= max_capacity)

    # --- Problem Definition and Solving ---
    problem = cp.Problem(objective, constraints)
    try:
        # Alternative solvers are ECOS, MOSEK, and SCS
        optimal_profit = problem.solve(solver=cp.CLARABEL, verbose=False)

        if problem.status in [cp.OPTIMAL, cp.OPTIMAL_INACCURATE]:
            return (
                problem.status,
                optimal_profit,
                P.value,  # NumPy array of optimal power flows per hour
                B.value   # NumPy array of optimal B at start of each hour
            )
        else:
            print(f"Optimization failed. Solver status: {problem.status}")
            return problem.status, None, None, None

    except cp.error.SolverError as e:
        print(f"Solver Error: {e}")
        return "Solver Error", None, None, None