import logging
import json
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from stable_baselines3 import TD3
from stable_baselines3.common.noise import NormalActionNoise
from rl_env import CooptimEnv
from src.cooptim.orchestrator import Orchestrator

# Configuration du logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
logger = logging.getLogger(__name__)

# --- RL INFERENCE ENGINE ---
def run_rl_inference(model, env, config):
    """Executes a full simulation day using the trained RL model with deterministic actions."""
    obs, _ = env.reset()
    done = False
    history = []
    input_data = env.input_data 
    energy_col = config["columns"]["energy"]
    fcr_col = config["columns"]["fcr"]
    
    step = 0
    while not done and step < len(input_data):
        # Use deterministic=True to see the learned strategy without exploration noise
        action, _ = model.predict(obs, deterministic=True)
        
        current_time = input_data.index[step]
        prices = input_data.iloc[step]
        
        # Environment Step
        obs, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated
        
        # Map Logic based on the 0.6 threshold defined in rl_env.py
        market = "Ancillary" if action[0] >= 0.6 else "Arbitrage"
        decision = "Idle"
        if market == "Arbitrage":
            if action[1] > 0.1: decision = "Sell"
            elif action[1] < -0.1: decision = "Buy"

        history.append({
            "Timestamp": current_time,
            "Market": market,
            "Decision": decision,
            "Power": action[1],
            "Energy_Price": prices.get(energy_col, 0),
            "FCR_Price": prices.get(fcr_col, 0),
            "SoC": env.battery.soc,
            "Revenue": reward
        })
        step += 1

    return pd.DataFrame(history)

# --- REFINED DASHBOARD ---
def plot_rl_dashboard(df):
    """Generates the multi-panel strategy dashboard highlighting market switches and buy/sell timing."""
    if df.empty:
        logger.error("No data collected for dashboard.")
        return

    plt.style.use('seaborn-v0_8-muted')
    fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(15, 12), sharex=True)

    # 1. SoC & Market Choice (Contextual Background)
    ax1.plot(df["Timestamp"], df["SoC"], color='black', linewidth=2.5, label="Battery SoC")
    for i in range(len(df)-1):
        # Color background based on market choice
        color = 'palegreen' if df["Market"].iloc[i] == "Ancillary" else 'lightskyblue'
        ax1.axvspan(df["Timestamp"].iloc[i], df["Timestamp"].iloc[i+1], color=color, alpha=0.4)
    ax1.set_ylabel("SoC (MWh)")
    ax1.set_title("RL Strategy: Market Selection (Green=Ancillary, Blue=Arbitrage)")
    ax1.legend(loc='upper left')

    # 2. Trading Decisions vs Energy Price
    ax2.plot(df["Timestamp"], df["Energy_Price"], color='darkorange', label="Energy Price (€/MWh)", alpha=0.8)
    ax2_twin = ax2.twinx()
    # Purple bars show the power flow (negative for charging/buying, positive for discharging/selling)
    ax2_twin.bar(df["Timestamp"], df["Power"], width=0.008, color='purple', alpha=0.5, label="Action (+ Disch, - Ch)")
    ax2_twin.set_ylim(-1.1, 1.1)
    ax2.set_ylabel("Price")
    ax2_twin.set_ylabel("Power Action")
    ax2.set_title("Trading Decisions: Buy (Down) / Sell (Up) vs Price")

    # 3. Financial Performance
    df["CumRev"] = df["Revenue"].cumsum()
    ax3.plot(df["Timestamp"], df["CumRev"], color='forestgreen', linewidth=2.5, label="Cumulative Profit")
    ax3.fill_between(df["Timestamp"], df["CumRev"], color='forestgreen', alpha=0.15)
    ax3.set_ylabel("Total Revenue (€)")
    ax3.set_title("Real-Time Profit Accumulation")

    plt.tight_layout()
    plt.savefig("dashboard.png")
    plt.show(block=True)

# --- MAIN BLOCK ---
if __name__ == "__main__":
    # Load configuration
    with open("config.json", "r") as f:
        global_config = json.load(f)
    
    START = global_config["run"]["start_date"]
    END = global_config["run"]["end_date"]

    # 1. Generate Baseline Data (Classical Co-optimization)
    logger.info(f"--- Running Baseline Co-optimization: {START} to {END} ---")
    orchestrator = Orchestrator(config=global_config)
    sols_coopt = orchestrator.run()

    # 2. RL Training Phase
    logger.info("=== Initializing RL Multi-Market Agent Training ===")
    try:
        # Use first day's input as training data
        training_df = sols_coopt[0].input
        
        # Setup Custom Environment
        env = CooptimEnv(training_df, global_config)
        
        # AGGRESSIVE EXPLORATION: Sigma 0.3 forces the agent to try negative power (Buying)
        n_actions = env.action_space.shape[0]
        action_noise = NormalActionNoise(mean=np.zeros(n_actions), sigma=0.3 * np.ones(n_actions))

        # Train TD3 Model with terminal SoC penalty awareness
        model = TD3(
            "MlpPolicy", 
            env, 
            action_noise=action_noise, 
            verbose=1, 
            learning_rate=0.0005,
            batch_size=256
        )
        
        # Increased timesteps (20k) to allow agent to learn from terminal penalties
        logger.info("Starting learning process...")
        model.learn(total_timesteps=20000)
        
        # 3. Decision Dashboard & Reporting
        logger.info("=== Running Post-Training Deterministic Inference ===")
        df_rl = run_rl_inference(model, env, global_config)
        
        # Print Market Summary
        print("\n" + "="*40)
        print(f"RL MARKET ACTIVITY SUMMARY")
        if not df_rl.empty:
            print(df_rl["Market"].value_counts())
        print("="*40)
        
        # Generate the Final Dashboard
        plot_rl_dashboard(df_rl)
        
        # Save results to CSV
        df_rl.to_csv("rl_trading_results.csv", index=False)
        logger.info("Full decision report saved to rl_trading_results.csv")

    except Exception as e:
        logger.error(f"RL Pipeline failure: {e}")

    logger.info("Simulation complete.")