baseline.py

"""
EcoGrid-OpenEnv — Baseline Inference Script

Runs a full episode of the environment using either a smart heuristic agent
(for reliable baselines) or an LLM (OpenAI) to demonstrate reasoning capabilities.
"""

import argparse
import json
import os
import time
from pathlib import Path
from typing import Literal

HAS_LITELLM = None

from env.environment import EcoGridEnv
from env.tasks import BasicGridBalanceGrader, RenewableVariabilityGrader, CarbonConstrainedGrader
from env.action_utils import safe_grid_action
from models.schemas import GridAction, GridState

_trained_model = None
_trained_tokenizer = None
_trained_load_attempted = False
TASK_EPISODE_LENGTH = {"easy": 48, "medium": 96, "hard": 96}
FOSSIL_EMISSION_FACTOR = 0.5
LORA_DIR = Path(os.environ.get("LORA_ADAPTER_DIR", str(Path(__file__).resolve().parent / "lora_adapter"))).resolve()


def _get_litellm():
    """Lazily import litellm to avoid startup-time network side effects."""
    global HAS_LITELLM
    if HAS_LITELLM is False:
        return None
    try:
        import litellm

        HAS_LITELLM = True
        return litellm
    except ImportError:
        HAS_LITELLM = False
        return None

def is_lora_valid() -> bool:
    """Check if LoRA files exist and are not just Git LFS pointers (> 1KB)."""
    required_files = ["adapter_model.safetensors", "tokenizer.json"]
    for name in required_files:
        path = LORA_DIR / name
        if not path.exists() or path.stat().st_size < 1024: # LFS pointers are ~133 bytes
            return False
    return True

def load_trained_model():
    """Lazily load the LoRA model if available."""
    global _trained_model, _trained_tokenizer, _trained_load_attempted
    if _trained_model is not None:
        return _trained_model, _trained_tokenizer
    if _trained_load_attempted:
        return None, None
        
    adapter_config_path = LORA_DIR / "adapter_config.json"
    if not adapter_config_path.exists():
        return None, "adapter_config.json missing"
    if not is_lora_valid():
        return None, "LORA weights are incomplete (check LFS pull)"
        
    print(f"Loading LoRA adapter from {LORA_DIR} ...")
    _trained_load_attempted = True
    try:
        from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
        from peft import PeftModel
        import torch
        
        with adapter_config_path.open("r", encoding="utf-8") as f:
            peft_config = json.load(f)
        base_model_name = peft_config.get("base_model_name_or_path", "unsloth/Qwen2.5-1.5B-Instruct")
        
        _trained_tokenizer = AutoTokenizer.from_pretrained(str(LORA_DIR))
        
        device = "cuda" if torch.cuda.is_available() else "cpu"
        
        if device == "cuda":
            quant_config = BitsAndBytesConfig(load_in_4bit=True)
            base_model = AutoModelForCausalLM.from_pretrained(
                base_model_name,
                quantization_config=quant_config,
                device_map="auto"
            )
        else:
            base_model = AutoModelForCausalLM.from_pretrained(
                base_model_name,
                device_map="cpu",
                torch_dtype=torch.float32
            )
            
        _trained_model = PeftModel.from_pretrained(base_model, str(LORA_DIR))
        print("LoRA successfully loaded!")
        return _trained_model, _trained_tokenizer
    except Exception as e:
        error_msg = f"Failed to load LoRA: {str(e)}"
        print(error_msg)
        return None, error_msg

def local_llm_agent(state: GridState, task_name: str) -> GridAction:
    """Agent that runs inference using the locally trained LoRA."""
    model, tokenizer = load_trained_model()
    if model is None:
        print("No LoRA found, falling back to heuristic.")
        return heuristic_agent(state, task_name)
        
    state_json = state.model_dump_json(indent=2)
    prompt = f"You are an expert energy grid operator.\nYour goal is to balance renewable energy, fossil fuels, and battery storage to meet demand while minimising cost and carbon emissions.\n\nCURRENT STATE:\n{state_json}\n\nTASK: {task_name}\nCONSTRAINTS: \n- renewable_ratio + fossil_ratio <= 1.0\n- battery_action must be between -1.0 (discharge) and 1.0 (charge)\n\nOutput ONLY a valid JSON object:\n{{\n  \"renewable_ratio\": float,\n  \"fossil_ratio\": float,\n  \"battery_action\": float\n}}"
    
    try:
        # Standard chat template formatting
        messages = [{"role": "user", "content": prompt}]
        inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)
        
        outputs = model.generate(inputs, max_new_tokens=100, temperature=0.1, do_sample=True)
        content = tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True).strip()
        
        if content.startswith("```json"):
            content = content[7:-3]
        elif content.startswith("```"):
            content = content[3:-3]
            
        data = json.loads(content)
        return safe_grid_action(
            renewable_ratio=data.get("renewable_ratio", 0.5),
            fossil_ratio=data.get("fossil_ratio", 0.5),
            battery_action=data.get("battery_action", 0.0),
        )
    except Exception as e:
        print(f"Local LLM Error: {e}. Falling back to heuristic.")
        return heuristic_agent(state, task_name)



def _constraint_aware_hard_controller(state: GridState) -> GridAction:
    """Hard-mode controller that enforces carbon budget pacing."""
    remaining_steps = max(1, TASK_EPISODE_LENGTH["hard"] - state.time_step)
    avg_renewable_cap = (state.solar_capacity + state.wind_capacity) / 2.0
    avg_renewable_cap = min(1.0, max(0.0, avg_renewable_cap))

    # Budget-aware fossil cap:
    # carbon_per_step = fossil_ratio * demand * emission_factor
    # => fossil_ratio <= carbon_budget_remaining / (remaining_steps * demand * emission_factor)
    if state.demand > 0:
        budget_fossil_cap = state.carbon_budget_remaining / (
            remaining_steps * state.demand * FOSSIL_EMISSION_FACTOR
        )
    else:
        budget_fossil_cap = 0.0

    # Keep a safety margin to avoid late-episode budget collapse.
    budget_fossil_cap = max(0.0, min(0.14, budget_fossil_cap * 0.92))
    future_floor = remaining_steps * max(state.demand, 1.0) * FOSSIL_EMISSION_FACTOR * 0.08
    if state.grid_stability < 0.75 and state.carbon_budget_remaining > future_floor:
        budget_fossil_cap = min(0.18, budget_fossil_cap + 0.03)

    renewable_ratio = min(0.9, max(0.62, avg_renewable_cap + 0.12))
    fossil_ratio = min(max(0.02, 1.0 - renewable_ratio), budget_fossil_cap)

    # Battery dispatch policy:
    # - discharge on high demand or low stability
    # - charge when demand is light and stability is healthy
    if (state.demand > 100 or state.grid_stability < 0.8) and state.battery_level > 0.12:
        battery_action = -0.9
    elif state.demand < 78 and state.battery_level < 0.7 and avg_renewable_cap > 0.4:
        battery_action = 0.6
    else:
        battery_action = 0.0

    return safe_grid_action(
        renewable_ratio=renewable_ratio,
        fossil_ratio=fossil_ratio,
        battery_action=battery_action,
    )


def heuristic_agent(state: GridState, task_name: str) -> GridAction:
    """Constraint-aware baseline agent with strict action validity guarantees."""
    if task_name == "hard":
        return _constraint_aware_hard_controller(state)

    avg_renewable_cap = (state.solar_capacity + state.wind_capacity) / 2.0
    avg_renewable_cap = min(1.0, max(0.0, avg_renewable_cap))
    renewable_ratio = min(0.95, max(0.05, avg_renewable_cap))
    fossil_ratio = max(0.0, 1.0 - avg_renewable_cap)

    if task_name == "medium" and (state.grid_stability < 0.8 or state.demand > 105):
        fossil_ratio = min(1.0, fossil_ratio + 0.05)

    if state.demand > 100 and state.battery_level > 0.2:
        battery_action = -0.9
        if task_name == "medium":
            fossil_ratio = max(0.0, fossil_ratio - 0.05)
    elif state.demand < 70 and state.battery_level < 0.8 and avg_renewable_cap > 0.5:
        battery_action = 0.7
        if task_name == "medium":
            fossil_ratio = min(1.0, fossil_ratio + 0.03)
    else:
        battery_action = 0.0

    return safe_grid_action(
        renewable_ratio=renewable_ratio,
        fossil_ratio=fossil_ratio,
        battery_action=battery_action,
    )


def llm_agent(state: GridState, task_name: str) -> GridAction:
    """An agent that uses an LLM to make decisions via Chain-of-Thought."""
    litellm = _get_litellm()
    if litellm is None:
        return heuristic_agent(state, task_name)
    
    prompt = f"""
You are an expert energy grid operator managing a power grid.
Your goal is to balance renewable energy, fossil fuels, and battery storage to meet demand while minimising cost and carbon emissions.

CURRENT STATE:
{state.model_dump_json(indent=2)}

TASK: {task_name}
CONSTRAINTS: 
- renewable_ratio + fossil_ratio <= 1.0
- battery_action must be between -1.0 (discharge) and 1.0 (charge)
- Grid stability target: >= 0.7
- Carbon budget remaining: {state.carbon_budget_remaining} kg CO2

Reason step-by-step internally about the best strategy, considering the current demand, available renewable capacity, and carbon budget.
Then, output ONLY a valid JSON object matching this schema, with no markdown fences:
{{
  "renewable_ratio": float,
  "fossil_ratio": float,
  "battery_action": float
}}
"""
    
    try:
        response = litellm.completion(
            model="gpt-4o",  # Using best model as requested
            messages=[{"role": "user", "content": prompt}],
            temperature=0.2,
        )
        
        content = response.choices[0].message.content.strip()
        # Clean up markdown if model ignored instructions
        if content.startswith("```json"):
            content = content[7:-3]
        elif content.startswith("```"):
            content = content[3:-3]
            
        data = json.loads(content)
        return safe_grid_action(
            renewable_ratio=data.get("renewable_ratio", 0.5),
            fossil_ratio=data.get("fossil_ratio", 0.5),
            battery_action=data.get("battery_action", 0.0),
        )
        
    except Exception as e:
        print(f"LLM Error: {e}. Falling back to heuristic.")
        return heuristic_agent(state, task_name)


def main():
    from rich.console import Console
    from rich.table import Table
    from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeElapsedColumn

    console = Console()
    
    parser = argparse.ArgumentParser(description="EcoGrid-OpenEnv Baseline Inference")
    parser.add_argument("--task", type=str, choices=["easy", "medium", "hard"], default="easy")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--agent", type=str, choices=["heuristic", "llm"], default="heuristic")
    args = parser.parse_args()
    
    if args.agent == "llm" and _get_litellm() is None:
        console.print("[bold red]Error:[/bold red] litellm package not installed. Run: pip install litellm")
        return
        
    if args.agent == "llm" and not os.environ.get("OPENAI_API_KEY"):
        console.print("[bold yellow]Warning:[/bold yellow] OPENAI_API_KEY environment variable not set. Falling back to heuristic.")
        args.agent = "heuristic"
        
    # Initialize environment
    console.print(f"[bold blue]Initializing EcoGridEnv for task:[/bold blue] {args.task} (seed={args.seed})")
    env = EcoGridEnv()
    state = env.reset(task=args.task, seed=args.seed)
    
    start_time = time.time()
    total_reward = 0.0
    
    console.print("\n[bold green]Starting episode...[/bold green]")
    
    table = Table(title="Live Grid Simulation", show_header=True, header_style="bold magenta")
    table.add_column("Step", style="dim", width=6)
    table.add_column("Demand", justify="right")
    table.add_column("Renw Ratio", justify="right")
    table.add_column("Foss Ratio", justify="right")
    table.add_column("Blackout", justify="right")
    table.add_column("Reward", justify="right", style="green")
    
    episode_length = env.get_task_config(args.task)["episode_length"]

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
        TimeElapsedColumn(),
        console=console,
        transient=True
    ) as progress:
        sim_task = progress.add_task("[cyan]Simulating grid...", total=episode_length)
        
        while not env.is_done:
            if args.agent == "llm":
                action = llm_agent(state, args.task)
            else:
                action = heuristic_agent(state, args.task)
                
            try:
                result = env.step(action)
            except ValueError as e:
                console.print(f"[bold yellow]Action constraint violation:[/bold yellow] {e}. Falling back to safe action.")
                safe_action = GridAction(renewable_ratio=0.5, fossil_ratio=0.5, battery_action=0.0)
                result = env.step(safe_action)
                
            state = result.observation
            total_reward += result.reward
            
            # Print progress every 10 steps or at the end
            if env.current_step % 10 == 0 or env.is_done:
                table.add_row(
                    str(env.current_step),
                    f"{state.demand:.1f}",
                    f"{action.renewable_ratio:.2f}",
                    f"{action.fossil_ratio:.2f}",
                    f"{result.info.get('blackout_risk', 0.0):.2f}",
                    f"{result.reward:.2f}"
                )
                
            progress.update(sim_task, advance=1)
            time.sleep(0.01) # slight delay to render progress smoothly for small baselines

    console.print(table)
    elapsed = time.time() - start_time
    
    # Grade the episode
    log = env.get_episode_log()
    if args.task == "easy":
        score = BasicGridBalanceGrader.grade(log)
    elif args.task == "medium":
        score = RenewableVariabilityGrader.grade(log)
    else:
        score = CarbonConstrainedGrader.grade(log)
        
    console.print("\n[bold]==================================================[/bold]")
    console.print("[bold cyan]EPISODE COMPLETE[/bold cyan]")
    console.print("[bold]==================================================[/bold]")
    console.print(f"Task: [bold]{args.task}[/bold]")
    console.print(f"Agent: [bold]{args.agent}[/bold]")
    console.print(f"Steps: {env.current_step}")
    console.print(f"Time: {elapsed:.2f}s ({env.current_step/elapsed:.0f} steps/sec)")
    if log:
        console.print(f"Termination: [bold]{log[-1].info.get('termination_reason', 'unknown')}[/bold]")
        
    console.print("\n[bold green]FINAL SCORE:[/bold green]")
    console.print(f"[bold text]{score.score * 100:.1f} / 100.0[/bold text]")
    console.print("\n[bold]Score Breakdown:[/bold]")
    
    breakdown_table = Table(show_header=False, box=None)
    breakdown_table.add_column("Metric", style="cyan")
    breakdown_table.add_column("Value", justify="right")
    
    for k, v in score.breakdown.items():
        breakdown_table.add_row(k.replace('_', ' ').title(), f"{v:.4f}")
    
    console.print(breakdown_table)

if __name__ == "__main__":
    main()