import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

import numpy as np
import matplotlib.pyplot as plt

from rl.environment import EnergyEnv
from rl.agent import RLAgent




# -------------------------
# Config
# -------------------------

PRICE_DATA_PATH = "data/price_data.csv"
NUM_EPISODES = 300
SAVE_MODEL_EVERY = 50
MODEL_PATH = "pvsai/PowerGrid"

os.makedirs("models", exist_ok=True)
os.makedirs("plots", exist_ok=True)


# -------------------------
# Training Loop
# -------------------------

def train():
    env = EnergyEnv(PRICE_DATA_PATH)
    state_dim = env.reset().shape[0]

    agent = RLAgent(state_dim=state_dim)

    episode_rewards = []
    episode_profits = []

    print("\n🚀 Starting Training...\n")

    for episode in range(1, NUM_EPISODES + 1):

        state = env.reset()
        total_reward = 0.0
        done = False

        while not done:
            action = agent.act(state)
            next_state, reward, done, info = env.step(action)

            agent.store(state, action, reward, next_state, done)
            agent.learn()

            state = next_state
            total_reward += reward

        episode_rewards.append(total_reward)
        episode_profits.append(info["total_profit"])

        if episode % 10 == 0:
            print(
                f"Episode {episode:4d} | "
                f"Reward: {total_reward:8.2f} | "
                f"Profit: ₹{info['total_profit']:8.2f} | "
                f"Epsilon: {agent.epsilon:.3f}"
            )

        if episode % SAVE_MODEL_EVERY == 0:
            agent.save(MODEL_PATH)
            print(f"💾 Model saved at episode {episode}")

    agent.save(MODEL_PATH)

    plot_training_curves(episode_rewards, episode_profits)

    print("\n✅ Training Complete!\n")


# -------------------------
# Evaluation Loop
# -------------------------

def evaluate(num_episodes=5, render=True):
    env = EnergyEnv(PRICE_DATA_PATH)
    state_dim = env.reset().shape[0]
    agent = RLAgent(state_dim=state_dim)
    agent.load(MODEL_PATH)
    agent.epsilon = 0.0  # purely greedy policy
    print("\n🔍 Evaluating trained policy...\n")
    for ep in range(num_episodes):
        state = env.reset()
        done = False

        soc_trace = []
        price_trace = []
        stress_trace = []
        profit_trace = []

        while not done:
            action = agent.act(state)
            next_state, reward, done, info = env.step(action)

            soc_trace.append(info["soc"])
            price_trace.append(info["price"])
            stress_trace.append(info["grid_stress"])
            profit_trace.append(info["total_profit"])

            state = next_state

        print(f"Episode {ep+1}: Final Profit = ₹{profit_trace[-1]:.2f}")

        if render:
            plot_episode(
                soc_trace,
                price_trace,
                stress_trace,
                profit_trace,
                episode=ep+1
            )

# -------------------------
# Plotting Utilities
# -------------------------

def plot_training_curves(rewards, profits):
    plt.figure(figsize=(14, 5))

    plt.subplot(1, 2, 1)
    plt.plot(rewards)
    plt.title("Training Reward Curve")
    plt.xlabel("Episode")
    plt.ylabel("Total Reward")

    plt.subplot(1, 2, 2)
    plt.plot(profits)
    plt.title("Training Profit Curve")
    plt.xlabel("Episode")
    plt.ylabel("Total Profit (₹)")

    plt.tight_layout()
    plt.savefig("plots/training_curves.png")
    plt.show()


def plot_episode(soc, price, stress, profit, episode):
    t = np.arange(len(soc))

    plt.figure(figsize=(14, 10))

    plt.subplot(4, 1, 1)
    plt.plot(price)
    plt.title("Market Price (₹/MWh)")
    plt.ylabel("Price")

    plt.subplot(4, 1, 2)
    plt.plot(soc)
    plt.title("Battery SoC (MWh)")
    plt.ylabel("SoC")

    plt.subplot(4, 1, 3)
    plt.plot(stress)
    plt.title("Grid Stress Index")
    plt.ylabel("Stress")

    plt.subplot(4, 1, 4)
    plt.plot(profit)
    plt.title("Cumulative Profit (₹)")
    plt.ylabel("Profit")
    plt.xlabel("15-min Time Steps")

    plt.tight_layout()
    plt.savefig(f"plots/episode_{episode}.png")
    plt.show()


# -------------------------
# Main
# -------------------------

if __name__ == "__main__":
    train()
    evaluate(num_episodes=3)