env/indicators_env.py

"""
indicators_env.py — IndicatorsEnv v4.1: Multi-stock Portfolio MDP

OpenEnv-compatible RL environment for NSE equity analysis.

Full OpenEnv interface:
  HTTP  : GET  /health  /tasks
          POST /reset  /step  /grader
          GET  /state  /baseline
  WebSocket: /ws

Env design (v4.1) — Multi-stock Portfolio MDP:
  At each step the agent observes 3 stocks from the same NSE sector.
  It picks ONE stock (or passes with NONE) and declares a direction.

  Observation:
    - 3 stocks from the same sector (same stocks for the whole episode)
    - Full technical indicator snapshot per stock
    - RSI + price-momentum signal history (accumulated across steps)
    - Macro context in Task 3 (NIFTY50 trend, market regime)

  Action:
    {"stock": "HDFCBANK", "direction": "Bullish"|"Bearish"|"NONE", "conviction": 0.8}
    NONE = skip this step (preserve selective participation)

  Reward (market-neutral alpha):
    alpha  = chosen_stock_period_return − sector_avg_period_return
    reward = alpha × direction_sign × conviction × REWARD_SCALE
    → random policy earns ~0 expected reward (sector avg cancels market beta)
    → skilled policy earns consistently positive reward (correct stock selection)
    → Kelly conviction semantics: fraction of virtual wealth wagered per step

  Step spacing:
    short  :  1 trading day per step → GT = 1-day return   (5 steps  = 1 week)
    medium :  5 trading days per step → GT = 5-day return  (10 steps = 10 weeks)
    long   : 20 trading days per step → GT = 20-day return (15 steps = 15 months)
    GT window = step spacing → zero overlap, reward and GT measure identical period.

  Task 3 extras:
    - Drawdown limit: virtual capital terminates at 10% drawdown
    - Macro context (NIFTY50) injected into each observation

Built for the Meta × PyTorch Hackathon.
"""

from __future__ import annotations

import json
import logging
import random
import sys
import uuid
from statistics import mean
from typing import Any, Dict, List, Optional
from datetime import datetime, timedelta

try:
    import yf_patch
    yf_patch.patch_yfinance_globally()
except ImportError:
    pass

import pandas as pd
import uvicorn
import numpy as np
from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
from pydantic import BaseModel, Field

from data_loader import (
    NSE_UNIVERSE, 
    SECTOR_GROUPS, 
    TERM_WINDOWS, 
    TERM_THRESHOLDS,
    STEP_SPACING,
    PERIOD_THRESHOLDS,
    build_multi_stock_episode,
    fetch_macro_context
)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="IndicatorsEnv",
    description=(
        "IndicatorsEnv v4.1 — Multi-stock Portfolio MDP for NSE equity analysis. "
        "At each step the agent observes 3 stocks from the same NSE sector and picks one "
        "(or passes). Reward = (chosen stock return − sector average) × direction × conviction. "
        "Market-neutral alpha with portfolio capital tracking: holding a position avoids "
        "transaction cost (0.1%); switching incurs it. Capital causally tracks cumulative alpha P&L. "
        "Drawdown limits apply to all tasks. Three tasks of increasing difficulty. "
        "Built for the Meta × PyTorch Hackathon."
    ),
    version="4.1.0",
)

# ─── Environment constants ────────────────────────────────────────────────────

TASK_MAX_STEPS: Dict[str, int] = {
    "short":  5,    # 5 × 1-day  = 1 trading week
    "medium": 10,   # 10 × 5-day = 10 trading weeks
    "long":   15,   # 15 × 20-day = 15 months
}

REWARD_SCALE     = 50.0   # ×50: 2% alpha × conviction 0.7 × 50 ≈ 0.7 reward
TRANSACTION_COST = 0.001  # 0.1% of capital on holding switch (exploration cost)

DRAWDOWN_LIMITS: Dict[str, float] = {
    "short":  0.05,   # 5%  — 5 steps, tight risk
    "medium": 0.10,   # 10% — 10 steps
    "long":   0.15,   # 15% — 15 steps, more recovery time
}

# ─── Task definitions ─────────────────────────────────────────────────────────

TASKS = [
    {
        "id": "short_term_direction",
        "name": "Short-term Relative Alpha (Easy)",
        "description": (
            "5 steps, 1 trading day apart (= 1 week of daily observations). "
            "At each step observe 3 stocks from the same NSE sector. "
            "Pick one stock and declare Bullish (long), Bearish (short), or NONE (pass). "
            "Reward = (chosen_return − sector_avg) × direction × conviction × 50. "
            "Market-neutral: sector beta cancels — the agent profits from within-sector "
            "relative momentum, not from broad market moves. "
            "GT = 1-day forward return (±0.3% threshold). "
            "Virtual capital tracked across steps; 0.1% transaction cost when switching holdings; "
            "episode ends early at 5% drawdown."
        ),
        "difficulty": "easy",
        "term": "short",
        "episode_steps": TASK_MAX_STEPS["short"],
        "step_spacing_days": STEP_SPACING["short"],
        "grader_note": "Directional accuracy on active (non-NONE) steps.",
        "action_schema": {
            "stock":     {"type": "string", "description": "NSE symbol or 'NONE' to skip"},
            "direction": {"type": "string", "enum": ["Bullish", "Bearish", "NONE"]},
            "conviction":{"type": "number", "minimum": 0.0, "maximum": 1.0},
        },
    },
    {
        "id": "medium_term_direction",
        "name": "Medium-term Relative Alpha (Medium)",
        "description": (
            "10 steps, 5 trading days apart (= 10 weeks, one observation per week). "
            "Same 3-stock multi-stock structure. "
            "Signal history accumulates across steps — RSI trends and price momentum "
            "are visible for all 3 stocks at each step, enabling sequential inference. "
            "Reward = (chosen_return − sector_avg) × direction × conviction × 50. "
            "GT = 5-day forward return (±1.5% threshold). "
            "Grader gives extra weight to Bearish/Neutral correct calls (anti-majority-bias). "
            "Virtual capital tracked; 0.1% switching cost; 10% drawdown terminates episode early."
        ),
        "difficulty": "medium",
        "term": "medium",
        "episode_steps": TASK_MAX_STEPS["medium"],
        "step_spacing_days": STEP_SPACING["medium"],
        "grader_note": (
            "Weighted accuracy on active steps: "
            "Bearish/Neutral correct = 1.5× weight. "
            "Participation bonus: score × (0.9 + 0.1 × active_rate)."
        ),
        "action_schema": {
            "stock":     {"type": "string", "description": "NSE symbol or 'NONE' to skip"},
            "direction": {"type": "string", "enum": ["Bullish", "Bearish", "NONE"]},
            "conviction":{"type": "number", "minimum": 0.0, "maximum": 1.0},
        },
    },
    {
        "id": "long_term_conviction",
        "name": "Long-term Risk-Constrained Alpha (Hard)",
        "description": (
            "15 steps, 20 trading days apart (= 15 months, one observation per month). "
            "Spans multiple market regimes — signal history essential for detecting "
            "multi-month relative momentum shifts within the sector. "
            "Macro context (NIFTY50 trend, market regime) added to each observation. "
            "Capital-tracked risk management: 0.1% switching cost + 15% drawdown limit. "
            "The agent must manage risk across months as well as generate alpha. "
            "Correct calls with conviction ≥ 0.7 score highest; "
            "overconfident wrong predictions penalized (conviction ≥ 0.8 on wrong). "
            "GT = 20-day forward return (±2.5% threshold)."
        ),
        "difficulty": "hard",
        "term": "long",
        "episode_steps": TASK_MAX_STEPS["long"],
        "step_spacing_days": STEP_SPACING["long"],
        "grader_note": (
            "Conviction-calibrated accuracy on active steps: "
            "correct+conviction≥0.7 = 1.0, correct+conviction<0.7 = 0.5, "
            "wrong = 0.0, overconfident wrong (conviction≥0.8) = −0.1. "
            "Normalized to (0, 1)."
        ),
        "action_schema": {
            "stock":     {"type": "string", "description": "NSE symbol or 'NONE' to skip"},
            "direction": {"type": "string", "enum": ["Bullish", "Bearish", "NONE"]},
            "conviction":{"type": "number", "minimum": 0.0, "maximum": 1.0},
        },
    },
]

TASK_BY_ID      = {t["id"]: t for t in TASKS}
TERM_TO_TASK_ID = {t["term"]: t["id"] for t in TASKS}

# ─── Pydantic schemas ─────────────────────────────────────────────────────────


class StockState(BaseModel):
    """Per-stock snapshot at one episode step."""
    symbol:             str
    current_price:      float
    rsi_14:             float
    rsi_trend:          str          # "up" | "down" | "flat"
    price_momentum_pct: float        # cumulative return vs episode start (%)
    indicators:         Dict[str, Any]


class MultiStockObservation(BaseModel):
    """Full observation returned by reset() and step()."""
    step:             int               # 1-indexed current step
    max_steps:        int
    term:             str               # SHORT | MEDIUM | LONG
    sector:           str               # e.g. "banking"
    available_stocks: List[str]         # 3 symbols
    stocks:           Dict[str, StockState]   # symbol → full state
    signal_history:   List[Dict[str, Any]]    # chronological pick log
    macro:            Optional[Dict[str, Any]] = None   # Task 3 only
    # v4.1 portfolio fields (all tasks)
    current_holding:  str   = "NONE"   # stock held from last active step, or "NONE"
    capital:          float = 1.0      # virtual capital (starts 1.0, tracks cumulative alpha P&L)
    drawdown:         float = 0.0      # current drawdown from peak capital


class MultiStockAction(BaseModel):
    stock:     str   = Field(..., description="NSE symbol to pick, or 'NONE' to skip")
    direction: str   = Field(..., description="Bullish | Bearish | NONE")
    conviction: float = Field(0.5, ge=0.0, le=1.0, description="Confidence 0–1")


class StepResult(BaseModel):
    observation: Optional[MultiStockObservation]
    reward:      float
    done:        bool
    info:        Dict[str, Any]


class ResetResult(BaseModel):
    observation: MultiStockObservation
    info:        Dict[str, Any]


class StateResult(BaseModel):
    session_id:          str
    current_observation: Optional[MultiStockObservation]
    episodes_completed:  int
    current_task:        Optional[str]


class GraderRequest(BaseModel):
    task_id: str
    episode_results: List[Dict[str, Any]] = Field(
        ...,
        description=(
            "List of step dicts, each with keys: "
            "'predicted' (str: Bullish|Bearish|NONE), "
            "'ground_truth' (str: Bullish|Bearish|Neutral|N/A), "
            "'conviction' (float)"
        ),
    )


class GraderResult(BaseModel):
    task_id:      str
    score:        float = Field(..., gt=0.0, lt=1.0)
    num_episodes: int
    breakdown:    Dict[str, Any]


# ─── Grader logic ─────────────────────────────────────────────────────────────


def _clamp_score(score: float) -> float:
    """Clamp to strictly open interval (0, 1) as required by the OpenEnv validator."""
    return round(max(0.001, min(0.999, score)), 4)


def grade_task(task_id: str, episode_results: List[Dict[str, Any]]) -> GraderResult:
    """
    Score episode results for a given task.

    Active steps = steps where predicted != 'NONE'.
    Grader evaluates only active steps; passing silently does not harm the score
    but forfeits the opportunity to score points.
    """
    if not episode_results:
        return GraderResult(
            task_id=task_id, score=0.001, num_episodes=0,
            breakdown={"note": "No results submitted"}
        )

    n_total  = len(episode_results)
    active   = [
        r for r in episode_results
        if str(r.get("predicted", "NONE")).capitalize() != "None"
        and str(r.get("predicted", "NONE")).upper() != "NONE"
    ]
    n_active = len(active)

    if n_active == 0:
        return GraderResult(
            task_id=task_id, score=0.001, num_episodes=n_total,
            breakdown={"note": "All steps passed (NONE)", "active_steps": 0}
        )

    participation_rate = n_active / n_total

    if task_id == "short_term_direction":
        correct = sum(
            1 for r in active
            if str(r.get("predicted","")).capitalize() == str(r.get("ground_truth","")).capitalize()
        )
        raw_score = correct / n_active
        score = _clamp_score(raw_score)
        breakdown = {
            "correct": correct,
            "active_steps": n_active,
            "total_steps": n_total,
            "participation_rate": round(participation_rate, 3),
            "metric": "directional_accuracy",
        }

    elif task_id == "medium_term_direction":
        weighted_score = 0.0
        weighted_total = 0.0
        for r in active:
            gt   = str(r.get("ground_truth", "")).capitalize()
            pred = str(r.get("predicted", "")).capitalize()
            w = 1.5 if gt in ("Bearish", "Neutral") else 1.0
            weighted_total += w
            if pred == gt:
                weighted_score += w
        base = (weighted_score / weighted_total) if weighted_total > 0 else 0.0
        # Participation bonus: agents that selectively engage vs. pass every step
        adjusted = base * (0.9 + 0.1 * participation_rate)
        score = _clamp_score(adjusted)
        breakdown = {
            "weighted_score": round(weighted_score, 3),
            "weighted_total": round(weighted_total, 3),
            "participation_rate": round(participation_rate, 3),
            "active_steps": n_active,
            "total_steps": n_total,
            "metric": "weighted_accuracy_with_participation",
        }

    elif task_id == "long_term_conviction":
        per_step: List[float] = []
        for r in active:
            gt         = str(r.get("ground_truth", "")).capitalize()
            pred       = str(r.get("predicted",    "")).capitalize()
            conviction = float(r.get("conviction", 0.5))
            correct    = (pred == gt)
            if correct and conviction >= 0.7:
                per_step.append(1.0)
            elif correct and conviction < 0.7:
                per_step.append(0.5)
            elif not correct and conviction >= 0.8:
                per_step.append(-0.1)   # overconfident wrong
            else:
                per_step.append(0.0)
        raw = sum(per_step) / n_active
        # Normalize [-0.1, 1.0] → (0, 1) exclusive
        score = _clamp_score((raw + 0.1) / 1.1)
        breakdown = {
            "per_step_scores": per_step,
            "raw_mean": round(raw, 4),
            "active_steps": n_active,
            "total_steps": n_total,
            "participation_rate": round(participation_rate, 3),
            "metric": "conviction_calibrated_accuracy",
        }

    else:
        raise HTTPException(status_code=404, detail=f"Unknown task_id: {task_id}")

    return GraderResult(
        task_id=task_id,
        score=round(score, 4),
        num_episodes=n_total,
        breakdown=breakdown,
    )


# ─── Session state ────────────────────────────────────────────────────────────


class EnvSession:
    """
    Manages a single agent session across multi-stock episodes.

    State per episode:
      episode_steps    : list of n_steps dicts from build_multi_stock_episode
      current_step_idx : 0-indexed pointer into episode_steps
      signal_history   : accumulated RSI/momentum log shown in each observation
      virtual_capital  : starts 1.0; updated by scaled alpha reward (Task 3 drawdown)

    Reward (per step):
      alpha  = chosen_stock_period_return − mean(all_3_stocks_period_return)
      reward = alpha × direction_sign × conviction × REWARD_SCALE
      NONE action → reward = 0.0
    """

    def __init__(self, session_id: str, term: str = "medium"):
        self.session_id     = session_id
        self.term           = term.lower()
        self.MAX_STEPS      = TASK_MAX_STEPS.get(self.term, 5)

        # Episode data: list of step dicts from build_multi_stock_episode
        self.episode_steps:     List[Dict[str, Any]] = []
        self.current_step_idx:  int = 0
        self.current_obs:       Optional[MultiStockObservation] = None
        self.sector:            str = ""
        self.symbols:           List[str] = []

        self.episodes_completed: int = 0
        self.episode_history:    List[Dict[str, Any]] = []
        self.signal_history:     List[Dict[str, Any]] = []

        # Portfolio: virtual capital and holding tracked across all tasks
        self.virtual_capital: float = 1.0
        self.peak_capital:    float = 1.0
        self.current_holding: str   = "NONE"

        # Lazy scenario pool
        self.scenario_pool: List[Dict[str, Any]] = []

    # ── helpers ──────────────────────────────────────────────────────────────

    def _reset_episode_state(self) -> None:
        self.episode_steps    = []
        self.current_step_idx = 0
        self.signal_history   = []
        self.virtual_capital  = 1.0
        self.peak_capital     = 1.0
        self.current_holding  = "NONE"

    def _get_scenario(self) -> Optional[Dict[str, Any]]:
        """Pop a (sector, symbols, date) scenario from the pool, refilling if needed."""
        if not self.scenario_pool:
            sectors   = list(SECTOR_GROUPS.keys())
            # Date range safely within available yfinance data
            # For long term: leave 15×20=300 days of future data → cap at 2022-12-31
            end_date = {
                "short":  "2024-06-30",
                "medium": "2023-12-31",
                "long":   "2022-06-30",
            }.get(self.term, "2023-12-31")

            dates = (
                pd.bdate_range("2020-01-01", end_date, freq="15B")
                .strftime("%Y-%m-%d")
                .tolist()
            )
            for sector in sectors:
                syms = SECTOR_GROUPS[sector]
                sampled_dates = random.sample(dates, min(60, len(dates)))
                for date in sampled_dates:
                    self.scenario_pool.append({
                        "sector":  sector,
                        "symbols": random.sample(syms, min(3, len(syms))),
                        "date":    date,
                    })
            random.shuffle(self.scenario_pool)

        return self.scenario_pool.pop() if self.scenario_pool else None

    def _build_obs(self, step_idx: int) -> MultiStockObservation:
        """Construct MultiStockObservation for the current step."""
        raw = self.episode_steps[step_idx]
        stocks_dict: Dict[str, StockState] = {}
        for s in raw["stocks"]:
            od = s["obs_dict"]
            stocks_dict[s["symbol"]] = StockState(
                symbol             = od["symbol"],
                current_price      = od["current_price"],
                rsi_14             = od["rsi_14"],
                rsi_trend          = od["rsi_trend"],
                price_momentum_pct = od["price_momentum_pct"],
                indicators         = od["indicators"],
            )

        drawdown = (
            (self.peak_capital - self.virtual_capital) / self.peak_capital
            if self.peak_capital > 0 else 0.0
        )

        return MultiStockObservation(
            step             = step_idx + 1,          # 1-indexed
            max_steps        = self.MAX_STEPS,
            term             = self.term.upper(),
            sector           = self.sector,
            available_stocks = self.symbols,
            stocks           = stocks_dict,
            signal_history   = list(self.signal_history),
            macro            = raw.get("macro"),
            current_holding  = self.current_holding,
            capital          = round(self.virtual_capital, 6),
            drawdown         = round(drawdown, 4),
        )

    # ── core interface ────────────────────────────────────────────────────────

    def reset(self) -> Optional[ResetResult]:
        """
        Start a new episode. Picks a sector + 3 stocks, fetches data in 3 yfinance
        calls (one per stock), builds the full episode in one pass.
        """
        self._reset_episode_state()

        for _ in range(10):
            sc = self._get_scenario()
            if sc is None:
                return None

            steps = build_multi_stock_episode(
                symbols      = sc["symbols"],
                start_date   = sc["date"],
                n_steps      = self.MAX_STEPS,
                term         = self.term,
                include_macro= (self.term == "long"),
            )
            if steps is not None:
                self.episode_steps = steps
                self.sector        = sc["sector"]
                self.symbols       = sc["symbols"]
                self.current_obs   = self._build_obs(0)
                return ResetResult(
                    observation = self.current_obs,
                    info = {
                        "session_id": self.session_id,
                        "term":       self.term,
                        "task_id":    TERM_TO_TASK_ID.get(self.term),
                        "sector":     self.sector,
                        "symbols":    self.symbols,
                        "step":       0,
                        "max_steps":  self.MAX_STEPS,
                    },
                )
        return None

    def step(self, action: MultiStockAction) -> StepResult:
        if not self.episode_steps or self.current_step_idx >= len(self.episode_steps):
            return StepResult(
                observation=None, reward=0.0, done=True,
                info={"error": "Call reset() first"},
            )

        raw         = self.episode_steps[self.current_step_idx]
        stock_rows  = {s["symbol"]: s for s in raw["stocks"]}

        # ── Parse action ──────────────────────────────────────────────────────
        stock     = action.stock.strip().upper()
        direction = action.direction.strip().capitalize()
        if direction not in ("Bullish", "Bearish"):
            direction = "NONE"
        if stock == "NONE" or direction == "NONE":
            stock, direction = "NONE", "NONE"
        elif stock not in stock_rows:
            # Invalid symbol → treat as NONE
            stock, direction = "NONE", "NONE"

        conviction = action.conviction

        # ── Reward: market-neutral alpha + portfolio dynamics ─────────────────
        all_returns = [s["actual_period_return"] for s in raw["stocks"]]
        sector_avg  = mean(all_returns) if all_returns else 0.0

        if direction == "NONE":
            reward        = 0.0
            alpha         = 0.0
            chosen_return = 0.0
            chosen_gt     = "N/A"
            tx_cost       = 0.0
            # NONE pass: keep existing holding, no capital change
        else:
            chosen_return  = stock_rows[stock]["actual_period_return"]
            alpha          = chosen_return - sector_avg
            direction_sign = 1 if direction == "Bullish" else -1
            raw_reward     = alpha * direction_sign * conviction * REWARD_SCALE
            chosen_gt      = stock_rows[stock]["gt"]

            # Transaction cost: 0.1% of capital when switching holdings (all tasks)
            tx_cost = 0.0
            if (self.current_holding != "NONE"
                    and stock != self.current_holding):
                tx_cost = TRANSACTION_COST * self.virtual_capital

            reward = round(max(-1.5, min(1.5, raw_reward)) - tx_cost, 4)

            # Update capital — ALL tasks (not just long)
            pnl_fraction         = alpha * direction_sign * conviction
            self.virtual_capital = max(0.0, self.virtual_capital + pnl_fraction)
            self.peak_capital    = max(self.peak_capital, self.virtual_capital)

            # Update holding
            self.current_holding = stock

        correct   = (direction != "NONE") and (direction == chosen_gt)
        drawdown_limit = DRAWDOWN_LIMITS.get(self.term, 0.10)
        drawdown  = (
            (self.peak_capital - self.virtual_capital) / self.peak_capital
            if self.peak_capital > 0 else 0.0
        )

        # ── Update signal history ─────────────────────────────────────────────
        self.signal_history.append({
            "step":              self.current_step_idx + 1,
            "picked_stock":      stock,
            "direction":         direction,
            "conviction":        conviction,
            "ground_truth":      chosen_gt,
            "correct":           correct if direction != "NONE" else None,
            "alpha_pct":         round(alpha * 100, 3),
            "reward":            reward,
            # RSI snapshot of all stocks at this step
            "rsi_snapshot": {
                s["symbol"]: round(s["obs_dict"]["rsi_14"], 1)
                for s in raw["stocks"]
            },
        })

        self.episode_history.append({
            "step":         self.current_step_idx + 1,
            "stock":        stock,
            "direction":    direction,
            "ground_truth": chosen_gt,
            "correct":      correct,
            "conviction":   conviction,
            "reward":       reward,
            "tx_cost":      round(tx_cost, 6),
            "alpha_pct":    round(alpha * 100, 3),
            "chosen_return_pct": round(chosen_return * 100, 4),
            "sector_avg_pct":    round(sector_avg * 100, 4),
        })

        # ── Advance step ──────────────────────────────────────────────────────
        self.current_step_idx += 1

        early_stop = drawdown > drawdown_limit
        done       = (self.current_step_idx >= self.MAX_STEPS) or early_stop

        info = {
            "step":              self.current_step_idx,
            "max_steps":         self.MAX_STEPS,
            "session_id":        self.session_id,
            "task_id":           TERM_TO_TASK_ID.get(self.term),
            "sector":            self.sector,
            "chosen_stock":      stock,
            "direction":         direction,
            "ground_truth":      chosen_gt,
            "correct":           correct,
            "conviction":        conviction,
            "tx_cost":           round(tx_cost, 6),
            "current_holding":   self.current_holding,
            "alpha_pct":         round(alpha * 100, 3),
            "chosen_return_pct": round(chosen_return * 100, 4),
            "sector_avg_pct":    round(sector_avg * 100, 4),
            "virtual_capital":   round(self.virtual_capital, 6),
            "drawdown":          round(drawdown, 4),
            "drawdown_limit":    drawdown_limit,
            "early_stop":        early_stop,
        }

        if not done:
            self.current_obs = self._build_obs(self.current_step_idx)
        else:
            self.episodes_completed += 1
            self.current_obs = None

        return StepResult(
            observation = self.current_obs,
            reward      = reward,
            done        = done,
            info        = info,
        )

    def state(self) -> StateResult:
        return StateResult(
            session_id          = self.session_id,
            current_observation = self.current_obs,
            episodes_completed  = self.episodes_completed,
            current_task        = TERM_TO_TASK_ID.get(self.term),
        )


# ─── Session store ────────────────────────────────────────────────────────────

_sessions: Dict[str, EnvSession] = {}


def _get_or_create(session_id: Optional[str] = None, term: str = "medium") -> EnvSession:
    if session_id is None:
        session_id = str(uuid.uuid4())
    if session_id not in _sessions:
        _sessions[session_id] = EnvSession(session_id, term=term)
    return _sessions[session_id]


# ─── HTTP endpoints (OpenEnv standard) ───────────────────────────────────────


@app.get("/health")
def health():
    return {
        "status":           "ok",
        "env":              "IndicatorsEnv",
        "version":          "4.1.0",
        "tasks":            len(TASKS),
        "episode_steps":    TASK_MAX_STEPS,
        "step_spacing":     STEP_SPACING,
        "sectors":          list(SECTOR_GROUPS.keys()),
        "drawdown_limits":  DRAWDOWN_LIMITS,
        "transaction_cost": TRANSACTION_COST,
    }


@app.get("/tasks")
def get_tasks():
    """Returns all task definitions, action schemas, and episode lengths."""
    return {
        "tasks":         TASKS,
        "total":         len(TASKS),
        "action_schema": {
            "stock":     {"type": "string", "description": "NSE symbol or 'NONE' to skip"},
            "direction": {"type": "string", "enum": ["Bullish", "Bearish", "NONE"]},
            "conviction":{"type": "number", "minimum": 0.0, "maximum": 1.0},
        },
        "episode_steps":  TASK_MAX_STEPS,
        "step_spacing":   STEP_SPACING,
    }


@app.post("/reset", response_model=ResetResult)
def reset(session_id: Optional[str] = None, term: str = "medium"):
    sess   = _get_or_create(session_id, term=term)
    result = sess.reset()
    if result is None:
        raise HTTPException(
            status_code=503,
            detail="Could not fetch market data. Retry or try a different term.",
        )
    return result


@app.post("/step", response_model=StepResult)
def step(action: MultiStockAction, session_id: Optional[str] = None):
    sess = _get_or_create(session_id)
    return sess.step(action)


@app.get("/state", response_model=StateResult)
def state(session_id: Optional[str] = None):
    sess = _get_or_create(session_id)
    return sess.state()


@app.post("/grader", response_model=GraderResult)
def grader(request: GraderRequest):
    """
    Score a completed episode set against a specific task grader.
    episode_results: list of {predicted, ground_truth, conviction}.
    Returns score in (0.0, 1.0).
    """
    if request.task_id not in TASK_BY_ID:
        raise HTTPException(
            status_code=404,
            detail=f"Unknown task_id: {request.task_id}. "
                   f"Valid: {list(TASK_BY_ID.keys())}",
        )
    return grade_task(request.task_id, request.episode_results)


@app.get("/baseline")
def baseline():
    """
    Run the built-in random baseline across all 3 tasks (10 episodes each).
    Expected: overall_mean ≈ 0.33 (random agent on 3-class directional problem).
    short  : 5×10=50  steps
    medium : 10×10=100 steps
    long   : ≤15×10=150 steps (may be fewer with early termination)
    """
    results = {}
    random.seed(42)

    for task in TASKS:
        term    = task["term"]
        task_id = task["id"]
        sess    = EnvSession(session_id=f"baseline-{task_id}", term=term)
        episode_results: List[Dict[str, Any]] = []

        for _ in range(10):
            reset_result = sess.reset()
            if reset_result is None:
                continue
            done = False
            while not done:
                # Random agent: random stock, random direction (including NONE)
                available = reset_result.observation.available_stocks
                if hasattr(sess.current_obs, "available_stocks") and sess.current_obs:
                    available = sess.current_obs.available_stocks
                chosen_stock = random.choice(available + ["NONE"])
                direction    = "NONE" if chosen_stock == "NONE" else random.choice(["Bullish", "Bearish"])
                action = MultiStockAction(
                    stock      = chosen_stock,
                    direction  = direction,
                    conviction = round(random.uniform(0.3, 0.9), 2),
                )
                step_result = sess.step(action)
                episode_results.append({
                    "predicted":    direction,
                    "ground_truth": step_result.info.get("ground_truth", "N/A"),
                    "conviction":   action.conviction,
                })
                done = step_result.done

        grader_result = grade_task(task_id, episode_results)
        results[task_id] = {
            "task_name":    task["name"],
            "difficulty":   task["difficulty"],
            "score":        grader_result.score,
            "num_episodes": grader_result.num_episodes,
            "breakdown":    grader_result.breakdown,
        }

    return {
        "agent":        "random_baseline",
        "seed":         42,
        "tasks":        results,
        "overall_mean": round(sum(r["score"] for r in results.values()) / len(results), 4),
    }


# ─── WebSocket endpoint (OpenEnv /ws) ────────────────────────────────────────


@app.websocket("/ws")
async def websocket_env(ws: WebSocket, session_id: Optional[str] = None, term: str = "medium"):
    await ws.accept()
    sess = _get_or_create(session_id, term=term)
    logger.info(f"[WS] Session {sess.session_id} connected (term={term})")
    try:
        while True:
            raw = await ws.receive_text()
            msg = json.loads(raw)
            method = msg.get("method", "")

            if method == "reset":
                result = sess.reset()
                if result is None:
                    await ws.send_json({"error": "No data available"})
                else:
                    await ws.send_text(result.model_dump_json())

            elif method == "step":
                action = MultiStockAction(**msg.get("action", {}))
                result = sess.step(action)
                await ws.send_text(result.model_dump_json())

            elif method == "state":
                result = sess.state()
                await ws.send_text(result.model_dump_json())

            else:
                await ws.send_json({"error": f"Unknown method: {method}"})

    except WebSocketDisconnect:
        logger.info(f"[WS] Session {sess.session_id} disconnected")
    except Exception as e:
        logger.error(f"[WS] Error: {e}")
        await ws.close()


# ─── Entry point ─────────────────────────────────────────────────────────────


if __name__ == "__main__":
    # workers=1 required: _sessions is an in-memory dict, not shared across processes
    uvicorn.run("indicators_env:app", host="0.0.0.0", port=7860, reload=False, workers=1)