import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.optimize import linprog

# ---------------------------------------------------------
# 1. Create 3 days of hourly data
# ---------------------------------------------------------

np.random.seed(7)

n_hours = 72
hours = np.arange(n_hours)
hour_of_day = hours % 24
day = hours // 24

# PV plant parameters
pv_capacity = 80  # MW
grid_limit = 50   # MW

# Synthetic PV generation curve for 3 days
solar_shape = np.maximum(0, np.sin(np.pi * (hour_of_day - 6) / 12))

# Different solar conditions for each day
cloud_factor = np.array([1.00, 0.85, 1.10])[day]

pv_generation = pv_capacity * solar_shape * cloud_factor
pv_generation += np.random.normal(0, 2, size=n_hours)
pv_generation = np.clip(pv_generation, 0, None)

# Synthetic electricity price curve in $/MWh
# Low prices around midday, high prices in the evening
base_price = 45

evening_peak = 55 * np.exp(-0.5 * ((hour_of_day - 19) / 2.2) ** 2)
morning_peak = 15 * np.exp(-0.5 * ((hour_of_day - 8) / 2.5) ** 2)
midday_discount = 25 * np.exp(-0.5 * ((hour_of_day - 13) / 3.0) ** 2)

price = base_price + evening_peak + morning_peak - midday_discount
price += np.random.normal(0, 3, size=n_hours)
price = np.clip(price, 5, None)

# ---------------------------------------------------------
# 2. Battery and economic assumptions
# ---------------------------------------------------------

charging_efficiency = 0.95
discharging_efficiency = 0.95

# Economic assumptions
battery_power_cost = 150_000   # $/MW
battery_energy_cost = 200_000  # $/MWh

discount_rate = 0.08
battery_lifetime = 15  # years

# Capital Recovery Factor
crf = (
    discount_rate * (1 + discount_rate) ** battery_lifetime
    / ((1 + discount_rate) ** battery_lifetime - 1)
)

# ---------------------------------------------------------
# 3. Optimization model for a given BESS size
# ---------------------------------------------------------

def optimize_bess_dispatch(bess_power_mw, bess_energy_mwh):
    """
    Optimizes the hourly operation of the BESS for a fixed size.

    Decision variables:
    - battery charge
    - battery discharge
    - state of charge
    - curtailment
    - energy exported to the grid
    """

    n = n_hours

    # Variable indexes
    idx_charge = np.arange(0, n)
    idx_discharge = np.arange(n, 2 * n)
    idx_soc = np.arange(2 * n, 3 * n)
    idx_curtailment = np.arange(3 * n, 4 * n)
    idx_export = np.arange(4 * n, 5 * n)

    n_variables = 5 * n

    # Objective function:
    # maximize revenue = price * exported energy
    # linprog minimizes, so we use negative revenue
    c = np.zeros(n_variables)
    c[idx_export] = -price

    A_eq = []
    b_eq = []

    # Energy balance:
    # PV + battery discharge = export + battery charge + curtailment
    for t in range(n):
        row = np.zeros(n_variables)

        row[idx_export[t]] = 1
        row[idx_charge[t]] = 1
        row[idx_curtailment[t]] = 1
        row[idx_discharge[t]] = -1

        A_eq.append(row)
        b_eq.append(pv_generation[t])

    # Battery SOC balance:
    # SOC[t] = SOC[t-1] + charge * eta_c - discharge / eta_d
    for t in range(n):
        row = np.zeros(n_variables)

        row[idx_soc[t]] = 1
        row[idx_charge[t]] = -charging_efficiency
        row[idx_discharge[t]] = 1 / discharging_efficiency

        if t > 0:
            row[idx_soc[t - 1]] = -1

        A_eq.append(row)
        b_eq.append(0)

    # Variable bounds
    bounds = []

    # Battery charge limit
    bounds += [(0, bess_power_mw)] * n

    # Battery discharge limit
    bounds += [(0, bess_power_mw)] * n

    # Battery energy capacity limit
    bounds += [(0, bess_energy_mwh)] * n

    # Curtailment
    bounds += [(0, None)] * n

    # Grid export limit
    bounds += [(0, grid_limit)] * n

    result = linprog(
        c,
        A_eq=np.array(A_eq),
        b_eq=np.array(b_eq),
        bounds=bounds,
        method="highs"
    )

    if not result.success:
        raise RuntimeError("Optimization failed: " + result.message)

    x = result.x

    dispatch = pd.DataFrame({
        "Hour": hours,
        "PV Generation (MW)": pv_generation,
        "Price ($/MWh)": price,
        "Battery Charge (MW)": x[idx_charge],
        "Battery Discharge (MW)": x[idx_discharge],
        "Battery SOC (MWh)": x[idx_soc],
        "Curtailment (MWh)": x[idx_curtailment],
        "Energy Exported (MWh)": x[idx_export]
    })

    revenue_3_days = np.sum(
        dispatch["Energy Exported (MWh)"] * dispatch["Price ($/MWh)"]
    )

    total_curtailment = dispatch["Curtailment (MWh)"].sum()

    return dispatch, revenue_3_days, total_curtailment

# ---------------------------------------------------------
# 4. Search for the optimum BESS size
# ---------------------------------------------------------

candidate_powers = [0, 5, 10, 15, 20, 25, 30, 40]  # MW
storage_durations = [1, 2, 3, 4]  # hours

sizing_results = []

for power in candidate_powers:

    if power == 0:
        candidate_energies = [0]
    else:
        candidate_energies = [power * duration for duration in storage_durations]

    for energy in candidate_energies:

        dispatch, revenue_3_days, curtailment = optimize_bess_dispatch(
            power, energy
        )

        # Convert 3-day revenue to annual revenue
        annual_revenue = revenue_3_days * 365 / 3

        # Annualized BESS investment cost
        investment_cost = (
            power * battery_power_cost
            + energy * battery_energy_cost
        )

        annualized_investment_cost = investment_cost * crf

        # Net annual benefit
        net_annual_value = annual_revenue - annualized_investment_cost

        sizing_results.append({
            "BESS Power (MW)": power,
            "BESS Energy (MWh)": energy,
            "Storage Duration (h)": energy / power if power > 0 else 0,
            "3-Day Revenue ($)": revenue_3_days,
            "Annual Revenue ($)": annual_revenue,
            "Annualized BESS Cost ($)": annualized_investment_cost,
            "Net Annual Value ($)": net_annual_value,
            "Curtailment (MWh)": curtailment
        })

sizing_results = pd.DataFrame(sizing_results)

# Find best size
best_solution = sizing_results.loc[
    sizing_results["Net Annual Value ($)"].idxmax()
]

best_power = best_solution["BESS Power (MW)"]
best_energy = best_solution["BESS Energy (MWh)"]

print("\nSizing Results:")
print(sizing_results.sort_values("Net Annual Value ($)", ascending=False))

print("\nBest BESS Size:")
print(best_solution)

# ---------------------------------------------------------
# 5. Dispatch of the best BESS size
# ---------------------------------------------------------

best_dispatch, best_revenue, best_curtailment = optimize_bess_dispatch(
    best_power,
    best_energy
)

print("\nOptimal Dispatch for Best BESS Size:")
print(best_dispatch)

# ---------------------------------------------------------
# 6. Plot results
# ---------------------------------------------------------

plt.figure(figsize=(14, 6))
plt.plot(hours, pv_generation, label="PV Generation")
plt.plot(hours, best_dispatch["Energy Exported (MWh)"], label="Energy Exported")
plt.plot(hours, best_dispatch["Battery SOC (MWh)"], label="Battery SOC")
plt.axhline(grid_limit, linestyle="--", label="Grid Limit")
plt.xlabel("Hour")
plt.ylabel("MW / MWh")
plt.title("PV Plant with Optimized BESS Operation")
plt.legend()
plt.grid(True)
plt.show()

plt.figure(figsize=(14, 5))
plt.plot(hours, price, label="Electricity Price")
plt.xlabel("Hour")
plt.ylabel("Price ($/MWh)")
plt.title("Hourly Electricity Price Curve")
plt.legend()
plt.grid(True)
plt.show()

plt.figure(figsize=(14, 5))
plt.bar(hours, best_dispatch["Curtailment (MWh)"])
plt.xlabel("Hour")
plt.ylabel("Curtailment (MWh)")
plt.title("Remaining Curtailment After Optimized BESS Operation")
plt.grid(True)
plt.show()