backend.py

import modal
import os
import uuid
from datetime import datetime
import yfinance as yf
import pandas as pd
import requests
import io
import os
import math
import json
import warnings
import random
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from datetime import datetime
from typing import Optional, Tuple, Dict, List, Union
from dataclasses import dataclass, asdict
from pathlib import Path
from torch.utils.data import Dataset, DataLoader
from torch.amp import autocast, GradScaler
from transformers import AutoModelForSeq2SeqLM, AutoConfig
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import matplotlib.pyplot as plt

# --- Configuration ---
DOWNLOAD_DATA = False
RAW_DATA_DIR = "raw_data"
os.makedirs(RAW_DATA_DIR, exist_ok=True)

# Global feature order for training consistency
# LCS Upgrade (Phase 2): All covariates are now scale-invariant ratios
COV_COLS = [
    # Price ratios: (price / SMA) - 1  → mean-reverting, scale-invariant
    "ratio_mva5",  "ratio_mva10", "ratio_mva15", "ratio_mva20", "ratio_mva30",
    # Volume normalized by SMA-20 (added inside add_technical_indicators)
    "Volume_ratio",
    # RSI is already [0,100] and scale-invariant
    "rsi14",
    # Bollinger %B: (Close - Lower) / (Upper - Lower) → [0,1] range, scale-invariant
    "BB_PercentB",
    # MACD normalised by Close price → scale-invariant
    "MACD_norm", "Signal_norm", "Hist_norm",
]

START_DATE = "2013-01-01"
HEADERS = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
}


def run_download_pipeline():
    """Encapsulated downloader logic."""
    if not DOWNLOAD_DATA:
        print(f"Skipping download (DOWNLOAD_DATA={DOWNLOAD_DATA}).")
        return

    # 1. Get Symbols
    symbols = get_combined_market_tickers()

    if len(symbols) > 0:
        # 2. Download Data
        historical_data = get_stock_data(symbols, start_date=START_DATE)

        # 3. Save Files
        print(f"\nSaving {len(historical_data)} files to {RAW_DATA_DIR}...")
        for ticker, df in historical_data.items():
            file_path = os.path.join(RAW_DATA_DIR, f"{ticker}.csv")
            df.to_csv(file_path)
        print("Done.")
    else:
        print("Failed to retrieve symbol list.")


def get_wiki_tickers(url, table_keywords):
    """
    Helper to fetch ticker symbols from a Wikipedia URL table.
    """
    tickers = []
    try:
        response = requests.get(url, headers=HEADERS)
        response.raise_for_status()

        tables = pd.read_html(io.StringIO(response.text))

        target_table = None
        for table in tables:
            cols = [c.lower() for c in table.columns]
            if any(k in cols for k in ['symbol', 'ticker']):
                target_table = table
                break

        if target_table is not None:
            # Find the specific column name (e.g. "Symbol" or "Ticker")
            col_name = next(c for c in target_table.columns if c.lower() in ['symbol', 'ticker'])
            tickers = target_table[col_name].tolist()
    except Exception as e:
        print(f" -> Error scraping {url}: {e}")

    return tickers

def get_combined_market_tickers():
    """
    Aggregates S&P 500, NASDAQ 100, and S&P 400 (MidCap).
    """
    print("--- Generaring Symbol List ---")
    all_tickers = set()

    # 1. S&P 500
    sp500 = get_wiki_tickers("https://en.wikipedia.org/wiki/List_of_S%26P_500_companies", ['symbol'])
    all_tickers.update(sp500)
    print(f"S&P 500 count: {len(sp500)}")

    # 2. NASDAQ 100
    ndx100 = get_wiki_tickers("https://en.wikipedia.org/wiki/Nasdaq-100", ['ticker', 'symbol'])
    all_tickers.update(ndx100)
    print(f"NASDAQ 100 count: {len(ndx100)}")

    # 3. S&P 400 MidCap (Ensures >250 quality stocks)
    sp400 = get_wiki_tickers("https://en.wikipedia.org/wiki/List_of_S%26P_400_companies", ['ticker', 'symbol'])
    all_tickers.update(sp400)
    print(f"S&P 400 count: {len(sp400)}")

    # Clean: Replace '.' with '-' (e.g. BRK.B -> BRK-B) and remove non-strings
    clean_list = [str(t).replace('.', '-') for t in all_tickers if isinstance(t, str)]
    unique_list = sorted(list(set(clean_list)))

    print(f"\nTotal unique symbols found: {len(unique_list)}")
    return unique_list

def get_stock_data(tickers, start_date=None, interval="1d", end_date=None):
    """
    Downloads data. If start_date is provided, it uses that.
    Otherwise, it uses period="max".
    """
    if not tickers:
        print("No tickers provided.")
        return {}

    if end_date is None:
        end_date = datetime.now().strftime('%Y-%m-%d')

    print(f"Starting download for {len(tickers)} symbols...")
    if start_date:
        print(f"Data Range: {start_date} to {end_date}")
    else:
        print(f"Data Range: MAX available history to {end_date}")

    all_data = {}
    BATCH_SIZE = 50

    for i in range(0, len(tickers), BATCH_SIZE):
        batch = tickers[i:i + BATCH_SIZE]
        print(f"Processing batch {i} - {i+len(batch)}...")

        try:
            # --- NEW LOGIC FOR START DATE ---
            # We construct the arguments dynamically.
            download_args = {
                "tickers": ' '.join(batch),
                "interval": interval,
                "end": end_date,
                "group_by": 'ticker',
                "auto_adjust": True,
                "progress": False,
                "threads": True
            }

            # If a start date is given, use 'start', otherwise use 'period="max"'
            if start_date:
                download_args["start"] = start_date
            else:
                download_args["period"] = "max"

            data = yf.download(**download_args)

            if not data.empty:
                if len(batch) == 1:
                    all_data[batch[0]] = data.dropna(how='all')
                else:
                    for ticker in batch:
                        try:
                            t_data = data.get(ticker)
                            if t_data is not None and not t_data.empty:
                                # Filter out rows that are all NaN
                                t_data = t_data.dropna(how='all')
                                if len(t_data) > 0:
                                    all_data[ticker] = t_data
                        except Exception:
                            continue
        except Exception as e:
            print(f"Batch error: {e}")

    return all_data


TRAIN_END = pd.Timestamp("2021-12-31")
VAL_END   = pd.Timestamp("2023-12-31")


# =============================================================================
# SECTION A: TECHNICAL INDICATORS
# =============================================================================

def simple_moving_average(
    df: pd.DataFrame,
    period: int,
    price_col: str = "Close",
    return_list: bool = False
):
    """Simple Moving Average (SMA) over a right-aligned rolling window."""
    if price_col not in df.columns or not isinstance(period, int) or period <= 0:
        raise ValueError("Invalid inputs.")
    sma = df[price_col].rolling(window=period, min_periods=period).mean()
    return sma.tolist() if return_list else sma


def ema(
    df: pd.DataFrame,
    period: int,
    price_col: str = "Close",
    seed: str = "price",  # kept for API compatibility
    min_periods: Optional[int] = None,
    return_list: bool = False
):
    """Exponential Moving Average (EMA)."""
    if price_col not in df.columns or not isinstance(period, int) or period <= 0:
        raise ValueError("Invalid inputs.")
    x = df[price_col].astype(float)

    # Standard EMA (recursive, no adjust)
    ser = x.ewm(span=period, adjust=False, min_periods=min_periods).mean()
    return ser.tolist() if return_list else ser


def rsi_wilder(
    df: pd.DataFrame,
    period: int = 14,
    price_col: str = "Close",
    return_list: bool = False
):
    """Wilder's Relative Strength Index (RSI)."""
    if price_col not in df.columns or not isinstance(period, int) or period <= 1:
        raise ValueError("Invalid inputs.")

    price = df[price_col].astype(float)
    delta = price.diff()

    gains = delta.clip(lower=0)
    losses = (-delta).clip(lower=0)

    avg_gain = gains.ewm(alpha=1/period, adjust=False, min_periods=period).mean()
    avg_loss = losses.ewm(alpha=1/period, adjust=False, min_periods=period).mean()

    rs = avg_gain / avg_loss
    rsi = 100.0 - (100.0 / (1.0 + rs))

    rsi = rsi.where(~((avg_loss == 0) & (avg_gain > 0)), 100.0)
    rsi = rsi.where(~((avg_gain == 0) & (avg_loss > 0)), 0.0)
    rsi = rsi.where(~((avg_gain == 0) & (avg_loss == 0)), 50.0)

    return rsi.tolist() if return_list else rsi


def bollinger_bands(
    df: pd.DataFrame,
    period: int = 20,
    num_std_dev: float = 2.0,
    price_col: str = "Close",
    ddof: int = 0,
    return_extras: bool = False
) -> pd.DataFrame:
    """Bollinger Bands over a right-aligned rolling window."""
    if price_col not in df.columns or not isinstance(period, int) or period <= 1:
        raise ValueError("Invalid inputs.")

    s = df[price_col].astype(float)
    middle = s.rolling(window=period, min_periods=period).mean()
    sigma = s.rolling(window=period, min_periods=period).std(ddof=ddof)

    upper = middle + num_std_dev * sigma
    lower = middle - num_std_dev * sigma

    out = pd.DataFrame({
        "BB_Middle": middle,
        "BB_Upper": upper,
        "BB_Lower": lower
    }, index=df.index)

    if return_extras:
        width = upper - lower
        percent_b = (s - lower) / width
        percent_b = percent_b.where(width > 0)
        bandwidth = width / middle
        out["BB_PercentB"] = percent_b
        out["BB_Bandwidth"] = bandwidth

    return out


def macd(
    df: pd.DataFrame,
    price_col: str = "Close",
    fast_period: int = 12,
    slow_period: int = 26,
    signal_period: int = 9,
    min_periods: Optional[int] = None
) -> pd.DataFrame:
    """MACD (Moving Average Convergence Divergence)."""
    if price_col not in df.columns or not (fast_period < slow_period and slow_period > 0):
        raise ValueError("Invalid periods/columns.")

    x = pd.to_numeric(df[price_col], errors="coerce").astype(float)
    ema_fast = x.ewm(span=fast_period, adjust=False).mean()
    ema_slow = x.ewm(span=slow_period, adjust=False).mean()

    macd_line = ema_fast - ema_slow
    signal = macd_line.ewm(span=signal_period, adjust=False).mean()
    histogram = macd_line - signal

    return pd.DataFrame(
        {"MACD": macd_line, "Signal_Line": signal, "MACD_Histogram": histogram},
        index=df.index
    )


def pivot_points_fibonacci_FIXED(
    df: pd.DataFrame,
    high_col: str = "High",
    low_col: str = "Low",
    close_col: str = "Close",
    session: str = "D"
) -> pd.DataFrame:
    """
    Fibonacci Pivot Points, shifted one period forward to avoid leakage.
    """
    df = df.sort_index()
    ohlc = df[[high_col, low_col, close_col]].copy()

    sess = ohlc.resample(session, label="left", closed="left").agg({
        high_col: "max",
        low_col: "min",
        close_col: "last"
    })

    H, L, C = sess[high_col], sess[low_col], sess[close_col]
    P = (H + L + C) / 3.0
    R = (H - L)

    R1 = P + 0.382 * R
    R2 = P + 0.618 * R
    R3 = P + 1.000 * R
    S1 = P - 0.382 * R
    S2 = P - 0.618 * R
    S3 = P - 1.000 * R

    fib_sess = pd.DataFrame({
        "PP": P,
        "R1": R1,
        "R2": R2,
        "R3": R3,
        "S1": S1,
        "S2": S2,
        "S3": S3,
    }, index=sess.index)

    fib_sess_shifted = fib_sess.shift(1)
    return fib_sess_shifted.reindex(df.index).ffill()


def add_technical_indicators(df: pd.DataFrame) -> pd.DataFrame:
    out = df.copy()

    # Raw SMAs (needed for ratio computation, not stored as raw price cols)
    sma5  = simple_moving_average(out, 5)
    sma10 = simple_moving_average(out, 10)
    sma15 = simple_moving_average(out, 15)
    sma20 = simple_moving_average(out, 20)
    sma30 = simple_moving_average(out, 30)

    # LCS Upgrade (Phase 2): Price/SMA ratios → scale-invariant
    out["ratio_mva5"]  = (out["Close"] / sma5)  - 1
    out["ratio_mva10"] = (out["Close"] / sma10) - 1
    out["ratio_mva15"] = (out["Close"] / sma15) - 1
    out["ratio_mva20"] = (out["Close"] / sma20) - 1
    out["ratio_mva30"] = (out["Close"] / sma30) - 1

    # Volume normalized by 20-day SMA of Volume (scale-invariant)
    if "Volume" in out.columns:
        vol_sma20 = out["Volume"].rolling(window=20, min_periods=20).mean()
        out["Volume_ratio"] = out["Volume"] / (vol_sma20 + 1e-8)
    else:
        out["Volume_ratio"] = 0.0  # fallback if Volume absent

    # RSI: already [0, 100], scale-invariant
    out["rsi14"] = rsi_wilder(out, 14)

    # Bollinger Bands: use %B instead of raw upper/lower
    bb = bollinger_bands(out, period=20, num_std_dev=2.0, price_col="Close", return_extras=True)
    out = out.join(bb)
    # BB_PercentB already computed by return_extras=True; handle NaN (flat price sections)
    out["BB_PercentB"] = out["BB_PercentB"].fillna(0.5)

    # MACD normalised by Close price → scale-invariant
    macd_df = macd(out, price_col="Close")
    out["MACD_norm"]   = macd_df["MACD"]         / (out["Close"] + 1e-8)
    out["Signal_norm"] = macd_df["Signal_Line"]   / (out["Close"] + 1e-8)
    out["Hist_norm"]   = macd_df["MACD_Histogram"]/ (out["Close"] + 1e-8)

    return out


# =============================================================================
# SECTION B: WINDOWING AND MULTI-TICKER LOADER
# =============================================================================

def build_windows_from_df(df: pd.DataFrame, T: int = 64, H: int = 16, columns: list = None):
    """
    LCS Upgrade (Phase 3): Produces sliding windows with Local Context Scaling.

    For each window:
      1. Extract raw price history (T steps) + future (H steps).
      2. Compute local scale = mean(abs(history prices)).
      3. Scale the entire window by this local scale → scaled prices ≈ N(1, σ).
      4. Tokenize scaled_history as input_ids, scaled_future as labels.
      5. Return raw prices, covariates, raw labels, AND local scales.

    Windows where scale <= 0.001 (flat/zero data) are dropped.
    """
    prices = df["price"].values.astype("float32")

    # Use passed columns (from checkpoint) or global fallback
    target_cols = columns if columns is not None else COV_COLS

    missing = [c for c in target_cols if c not in df.columns]
    if missing:
        raise ValueError(f"Missing columns in dataframe: {missing}")

    cov = df[target_cols].values.astype("float32")

    N_total = len(df)
    windows = N_total - T - H + 1
    if windows <= 0:
        raise ValueError("Dataset too small for given T and H.")

    price_windows = []
    cov_windows = []
    label_windows = []
    scale_windows = []

    for i in range(windows):
        p_window = prices[i:i+T]
        f_window = prices[i+T:i+T+H]

        # Filter non-positive windows early
        if np.any(p_window <= 0) or np.any(f_window <= 0):
            continue

        # LCS: compute local scale from history
        local_scale = float(np.mean(np.abs(p_window)))
        if local_scale <= 0.001:
            continue  # Drop flat/near-zero windows

        price_windows.append(p_window)
        cov_windows.append(cov[i:i+T])
        label_windows.append(f_window)
        scale_windows.append(local_scale)

    if not price_windows:
        return np.array([]), np.array([]), np.array([]), np.array([])

    return (
        np.stack(price_windows),
        np.stack(cov_windows),
        np.stack(label_windows),
        np.array(scale_windows, dtype="float32"),
    )

def load_multi_csvs_stream(
    data_dir: str,
    T: int = 64,
    H: int = 16,
    min_len: int = 120,
    max_files: int = 200      # <-- limit processed symbols per session
):
    """
    Stream CSVs in small batches to avoid RAM crashes.

    Returns concatenated arrays but processes at most `max_files` tickers.
    """
    all_prices = []
    all_cov = []
    all_labels = []

    files = [f for f in os.listdir(data_dir) if f.endswith(".csv")]
    files = files[:max_files]   # HARD LIMIT

    print(f"Processing {len(files)} tickers instead of full dataset...")

    for fname in files:
        path = os.path.join(data_dir, fname)
        try:
            df = pd.read_csv(path, parse_dates=["Date"])
        except Exception:
            continue

        if "Date" not in df.columns:
            continue

        df = df.sort_values("Date")
        df = df.drop_duplicates()
        df = df.set_index("Date")

        # Required cols
        if not {"Open","High","Low","Close"}.issubset(df.columns):
            continue

        # Business day alignment
        #df = df.resample("B").asfreq().ffill().bfill()
        df = df.resample("B").asfreq().ffill()
        df = df.dropna()

        df["price"] = df["Close"]

        if len(df) < min_len:
            continue

        # Compute indicators (heavy! but OK per file)
        df = add_technical_indicators(df)
        df = df.dropna()

        if len(df) < T + H:
            continue

        # Build windows — now returns (prices, cov, labels, scales)
        result = build_windows_from_df(df, T=T, H=H)
        p_i, c_i, y_i = result[0], result[1], result[2]

        # Only append if windows were generated
        if len(p_i) > 0:
            all_prices.append(p_i)
            all_cov.append(c_i)
            all_labels.append(y_i)

    if not all_prices:
        raise ValueError("No data after streaming.")

    prices = np.concatenate(all_prices, axis=0)
    covs   = np.concatenate(all_cov, axis=0)
    labels = np.concatenate(all_labels, axis=0)

    return prices, covs, labels



def compute_mase_scale_multi(data_dir: str) -> float:
    """
    Compute naive random-walk MAE across all tickers in data_dir
    for MASE scaling.
    """
    required_cols = ["Open", "High", "Low", "Close"]
    diffs = []

    for fname in os.listdir(data_dir):
        if not fname.endswith(".csv"):
            continue
        path = os.path.join(data_dir, fname)
        try:
            df = pd.read_csv(path, parse_dates=["Date"])
        except Exception:
            continue

        if "Date" not in df.columns:
            continue

        df = df.sort_values("Date")
        df = df.drop_duplicates()
        df = df.set_index("Date")

        if not set(required_cols).issubset(df.columns):
            continue

        df = df[required_cols].copy()
        #df = df.resample("B").asfreq().ffill().bfill()
        df = df.resample("B").asfreq().ffill()
        df = df.dropna()
        df["price"] = df["Close"]
        p = df["price"].values.astype("float32")

        if len(p) < 2:
            continue

        diffs.append(np.abs(p[1:] - p[:-1]))

    if not diffs:
        return 1.0

    diffs_all = np.concatenate(diffs, axis=0)
    mase_scale = float(np.mean(diffs_all))
    return max(mase_scale, 1e-8)


# =============================================================================
# SECTION C: TOKENIZER & NORMALIZER
# =============================================================================

class SimpleChronosTokenizer:
    """
    LCS Upgrade: Local Context Scaling tokenizer.
    Operates on scaled prices (price / local_scale), covering [0.0, 5.0].
    Replaces the old log-return tokenizer to eliminate compounding errors.
    """
    def __init__(self, num_bins: int = 1024, pad_token_id: int = 0,
                 min_val: float = 0.0, max_val: float = 5.0):
        self.num_bins = num_bins
        self.pad_token_id = pad_token_id
        self.min_val = min_val
        self.max_val = max_val
        self.range = self.max_val - self.min_val + 1e-8

        self.bos_token_id = 1
        self.eos_token_id = 2
        self.vocab_offset = 3
        self.vocab_size = num_bins + self.vocab_offset
        self.fitted = True

    def fit(self, prices: np.ndarray):
        # LCS: No-op. Fixed range [0.0, 5.0] covers scaled prices well.
        print(f"✓ Tokenizer using fixed LCS range [{self.min_val}, {self.max_val}]")

    def encode(self, arr: np.ndarray) -> np.ndarray:
        arr = np.array(arr, dtype=float)
        # Clamp; warn if many values are out-of-range (sign of bad scale)
        clamp_count = np.sum(arr > self.max_val)
        if clamp_count > len(arr) * 0.1:
            warnings.warn(f"⚠️ {clamp_count}/{len(arr)} values clamped above {self.max_val}. "
                          "Consider re-checking local scale computation.")
        arr = np.clip(arr, self.min_val, self.max_val)
        scaled = (arr - self.min_val) / self.range
        scaled = np.clip(scaled, 0, 1)
        bins = (scaled * (self.num_bins - 1)).astype(np.int64)
        tokens = bins + self.vocab_offset
        tokens[~np.isfinite(arr)] = self.pad_token_id
        return tokens

    def decode(self, tokens: np.ndarray) -> np.ndarray:
        tokens = np.array(tokens, dtype=np.int64)
        mask_special = (
            (tokens == self.pad_token_id) |
            (tokens == self.bos_token_id) |
            (tokens == self.eos_token_id)
        )
        bins = np.clip(tokens - self.vocab_offset, 0, self.num_bins - 1)
        scaled = bins.astype(float) / (self.num_bins - 1)
        values = scaled * self.range + self.min_val
        values[mask_special] = np.nan
        return values

    def save_config(self, path: str):
        config = {
            "num_bins": self.num_bins,
            "pad_token_id": self.pad_token_id,
            "min_val": self.min_val,
            "max_val": self.max_val,
            "mode": "local_scale"   # LCS Upgrade: renamed from "log_return"
        }
        with open(path, "w") as f:
            json.dump(config, f, indent=2)
        print(f"✓ Tokenizer config saved to {path}")

    @classmethod
    def load_config(cls, path: str):
        with open(path, "r") as f:
            config = json.load(f)
        return cls(
            num_bins=config["num_bins"],
            pad_token_id=config["pad_token_id"],
            min_val=config["min_val"],
            max_val=config["max_val"]
        )

class ZScoreNormalizer:
    """
    Anti-leakage Z-score normalizer with persistence.
    """

    def __init__(self, epsilon: float = 1e-8):
        self.mean: Optional[torch.Tensor] = None
        self.std: Optional[torch.Tensor] = None
        self.num_features: Optional[int] = None
        self.epsilon = epsilon
        self.fitted = False

    def fit(self, train_data: torch.Tensor):
        if train_data.dim() != 3:
            raise ValueError(f"Expected (N, T, F), got {train_data.shape}")
        self.mean = train_data.mean(dim=(0, 1), keepdim=True)
        self.std = train_data.std(dim=(0, 1), keepdim=True) + self.epsilon
        self.num_features = train_data.shape[2]
        self.fitted = True

        low_var = (self.std.squeeze() < (self.epsilon * 10)).cpu()
        if low_var.any():
            idx = torch.where(low_var)[0].tolist()
            warnings.warn(f"⚠️ Low variance features at indices: {idx}")

        print(f"✓ Normalizer fitted on training covariates (F={self.num_features})")

    def transform(self, data: torch.Tensor) -> torch.Tensor:
        if not self.fitted:
            raise ValueError("Normalizer not fitted. Call fit() first.")
        if data.shape[-1] != self.num_features:
            raise ValueError(f"Feature mismatch: expected {self.num_features}, got {data.shape[-1]}")
        return (data - self.mean.to(data.device)) / self.std.to(data.device)

    def inverse_transform(self, data: torch.Tensor) -> torch.Tensor:
        if not self.fitted:
            raise ValueError("Normalizer not fitted.")
        return data * self.std.to(data.device) + self.mean.to(data.device)

    # --- NEW PERSISTENCE METHODS ---
    def save(self, dir_path: str):
        if not self.fitted:
            raise ValueError("Cannot save unfitted normalizer")

        # Save tensors
        torch.save(self.mean, os.path.join(dir_path, "norm_mean.pt"))
        torch.save(self.std, os.path.join(dir_path, "norm_std.pt"))

        # Save metadata
        meta = {"num_features": self.num_features, "epsilon": self.epsilon}
        with open(os.path.join(dir_path, "norm_meta.json"), "w") as f:
            json.dump(meta, f, indent=2)
        print(f"✓ Normalizer stats saved to {dir_path}")

    @classmethod
    def load(cls, dir_path: str):
        meta_path = os.path.join(dir_path, "norm_meta.json")
        if not os.path.exists(meta_path):
            raise FileNotFoundError(f"Missing normalizer metadata: {meta_path}")

        with open(meta_path, "r") as f:
            meta = json.load(f)

        obj = cls(epsilon=meta["epsilon"])
        obj.num_features = meta["num_features"]

        # Load tensors (map to CPU initially)
        obj.mean = torch.load(os.path.join(dir_path, "norm_mean.pt"), map_location="cpu")
        obj.std = torch.load(os.path.join(dir_path, "norm_std.pt"), map_location="cpu")
        obj.fitted = True

        print(f"✓ Normalizer loaded (F={obj.num_features})")
        return obj

# =============================================================================
# SECTION D: DUAL-PATH CROSS-ATTENTION ADAPTER (PHASE 2)
# =============================================================================

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class xLSTMCovariateEncoder(nn.Module):
    """
    Strategy 1: Deep xLSTM Encoder — 3 stacked Pre-Norm residual mLSTM blocks.

    Layer 1 (Perception):  raw indicators -> baseline market states
    Layer 2 (Context):     combine states across time
    Layer 3 (Translation): map to T5 cross-attention space

    Each layer: residual highway + branch(LayerNorm -> mLSTMBlock)
    """
    def __init__(self, input_dim=11, d_model=256, num_heads=8,
                 num_layers=3, expansion_factor=2, dropout=0.1):
        super().__init__()
        self.input_dim  = input_dim
        self.d_model    = d_model
        self.num_heads  = num_heads
        self.num_layers = num_layers

        # Input projection: input_dim -> d_model
        self.projection = nn.Linear(input_dim, d_model)

        # Sinusoidal positional encoding
        self.pos_encoding = SinusoidalPositionalEncoding(d_model)

        # Stack of Pre-Norm mLSTM blocks
        self.norms  = nn.ModuleList([nn.LayerNorm(d_model) for _ in range(num_layers)])
        self.blocks = nn.ModuleList([
            MultiHeadmLSTMBlock(d_model=d_model, num_heads=num_heads,
                                expansion_factor=expansion_factor)
            for _ in range(num_layers)
        ])

        self.dropout = nn.Dropout(dropout)

    def forward(self, covariates):
        """
        Args:    covariates: (B, T, input_dim)
        Returns: output:     (B, T, d_model)
        """
        # 1. Project + positional encoding
        x = self.projection(covariates)    # (B, T, d_model)
        x = self.pos_encoding(x)           # (B, T, d_model)

        # 2. Hierarchical Pre-Norm residual stack
        for norm, block in zip(self.norms, self.blocks):
            # Branch: LayerNorm -> mLSTM -> Dropout
            branch = self.dropout(block(norm(x)))
            # Residual highway (untouched)
            x = x + branch

        return x                           # (B, T, d_model)

class MultiHeadmLSTMBlock(nn.Module):
    """
    Multi-head mLSTM block with Up-Projection + SiLU gate (Strategy 1).
    Operates in expanded space (d_model * expansion_factor) internally,
    then projects back down to d_model.
    """
    def __init__(self, d_model=256, num_heads=8, expansion_factor=2):
        super().__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_expanded = d_model * expansion_factor   # 512 default
        assert self.d_expanded % num_heads == 0
        self.d_head = self.d_expanded // num_heads     # 64 per head

        # Up-projection: d_model -> d_expanded (Cover's Theorem)
        self.up_proj = nn.Linear(d_model, self.d_expanded)

        # Per-head projections operating in EXPANDED space
        self.W_q     = nn.Parameter(torch.randn(num_heads, self.d_head, self.d_head))
        self.W_k     = nn.Parameter(torch.randn(num_heads, self.d_head, self.d_head))
        self.W_v     = nn.Parameter(torch.randn(num_heads, self.d_head, self.d_head))
        self.W_alpha = nn.Parameter(torch.randn(num_heads, self.d_head, 1))
        self.b_alpha = nn.Parameter(torch.zeros(num_heads, 1))
        self.W_out   = nn.Parameter(torch.randn(num_heads, self.d_head, self.d_head))

        # Down-projection: d_expanded -> d_model
        self.down_proj = nn.Linear(self.d_expanded, d_model)

        # Post-memory norm (operates on d_head in expanded space)
        self.post_norm = nn.LayerNorm(self.d_head)

        self._init_weights()

    def _init_weights(self):
        nn.init.xavier_uniform_(self.W_q)
        nn.init.xavier_uniform_(self.W_k)
        nn.init.xavier_uniform_(self.W_v)
        nn.init.xavier_uniform_(self.W_alpha)
        nn.init.xavier_uniform_(self.W_out)
        nn.init.uniform_(self.b_alpha, 0.0, 0.5)
        nn.init.xavier_uniform_(self.up_proj.weight)
        nn.init.xavier_uniform_(self.down_proj.weight)

    def forward(self, x):
        """
        Args:  x: (B, T, d_model)
        Returns: (B, T, d_model)
        """
        B, T, _ = x.shape

        # 1. Up-project + SiLU filter (information gate before memory write)
        x_up = F.silu(self.up_proj(x))              # (B, T, d_expanded)

        # 2. Split into heads in expanded space
        x_heads = x_up.view(B, T, self.num_heads, self.d_head).transpose(1, 2)  # (B, H, T, d_head)

        # 3. Q, K, V projections
        Q = torch.einsum('bhtd,hde->bhte', x_heads, self.W_q)
        K = torch.einsum('bhtd,hde->bhte', x_heads, self.W_k)
        V = torch.einsum('bhtd,hde->bhte', x_heads, self.W_v)

        # 4. Gating
        alpha_logits = torch.einsum('bhtd,hdf->bhtf', x_heads, self.W_alpha)
        alpha_logits = alpha_logits + self.b_alpha.view(1, self.num_heads, 1, 1)
        alpha = torch.clamp(torch.sigmoid(alpha_logits), min=1e-6, max=1.0 - 1e-6).unsqueeze(-1)

        # 5. Rank-1 memory updates
        updates = torch.matmul(K.unsqueeze(-1), V.unsqueeze(-2))  # (B, H, T, d, d)

        # 6. Sequential memory update
        M = self._sequential_memory_update(alpha.transpose(1, 2), updates.transpose(1, 2))
        M = M.transpose(1, 2)  # (B, H, T, d, d)

        # 7. Memory read + post-norm
        h = torch.matmul(M, Q.unsqueeze(-1)).squeeze(-1)  # (B, H, T, d_head)
        h = self.post_norm(h)

        # 8. Output projection with SiLU
        y = F.silu(torch.einsum('bhtd,hde->bhte', h, self.W_out))  # (B, H, T, d_head)

        # 9. Reassemble heads -> d_expanded
        y = y.transpose(1, 2).contiguous().view(B, T, self.d_expanded)

        # 10. Down-project back to d_model
        return self.down_proj(y)                      # (B, T, d_model)

    def _sequential_memory_update(self, alpha, updates):
        B, T, H, d, _ = updates.shape
        memory_states = []
        m_curr = alpha[:, 0] * updates[:, 0]
        memory_states.append(m_curr)
        for t in range(1, T):
            m_prev = memory_states[-1]
            m_curr = (1 - alpha[:, t]) * m_prev + alpha[:, t] * updates[:, t]
            memory_states.append(m_curr)
        return torch.stack(memory_states, dim=1)


class SinusoidalPositionalEncoding(nn.Module):
    """Standard sinusoidal positional encoding."""
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)  # (1, max_len, d_model)
        self.register_buffer('pe', pe)

    def forward(self, x):
        """
        Args:
            x: (B, T, d_model)
        Returns:
            x + positional encoding: (B, T, d_model)
        """
        return x + self.pe[:, :x.size(1), :]



class SinusoidalTimeEncoding(nn.Module):
    """
    Robust Sinusoidal PE that handles odd/even dimensions correctly.
    """
    def __init__(self, d_model: int, max_len: int = 500):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))

        pe[:, 0::2] = torch.sin(position * div_term)
        # Handle odd d_model by slicing the cosine term
        n_cos = pe[:, 1::2].shape[1]
        pe[:, 1::2] = torch.cos(position * div_term)[:, :n_cos]

        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        T = x.size(1)
        return x + self.pe[:, :T, :]


class CovariateEncoder(nn.Module):
    def __init__(self, input_dim, d_model, nhead=4, num_layers=1, dropout=0.1):
        """
        Encodes raw covariates into a dense, time-aware representation.
        Input: (B, T, F) -> Output: (B, T, d_model)
        """
        super().__init__()
        # 1. Project features to model dimension
        self.input_proj = nn.Linear(input_dim, d_model)

        # 2. Add Time Awareness (Crucial for Transformer)
        self.pos_encoder = SinusoidalTimeEncoding(d_model)

        # 3. Transformer Encoder to mix temporal features
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=d_model * 2,
            dropout=dropout,
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

    def forward(self, covariates):
        # covariates: [B, T, F]
        x = self.input_proj(covariates)  # [B, T, D]
        x = self.pos_encoder(x)          # Add PE
        return self.transformer(x)       # [B, T, D]


class CrossAttentionAdapter(nn.Module):
    def __init__(self, d_model, nhead=4, dropout=0.1):
        """
        Injects covariate context into encoder hidden states via Cross-Attention.
        Query = Encoder Hidden States
        Key/Value = Covariate Context
        """
        super().__init__()
        self.cross_attn = nn.MultiheadAttention(d_model, nhead, batch_first=True, dropout=dropout)
        self.layernorm = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

        # 🟢 Zero-Init Output Projection (Optional but recommended for stability)
        # This helps the model start close to Identity behavior
        nn.init.zeros_(self.cross_attn.out_proj.weight)
        nn.init.zeros_(self.cross_attn.out_proj.bias)

    def forward(self, encoder_hidden, cov_context):
        # encoder_hidden: [B, T, D] (Query)
        # cov_context:    [B, T, D] (Key/Value)

        residual = encoder_hidden

        # Cross Attention: Q=Hidden, K=Cov, V=Cov
        attn_output, _ = self.cross_attn(
            query=encoder_hidden,
            key=cov_context,
            value=cov_context,
            need_weights=False
        )

        # Residual Connection + Norm
        out = self.layernorm(residual + self.dropout(attn_output))
        return out


class DualPathAdapter(nn.Module):
    def __init__(self, cov_dim, d_model, num_layers=3, expansion_factor=2):
        super().__init__()

        self.cov_encoder = xLSTMCovariateEncoder(
            input_dim=cov_dim,
            d_model=d_model,
            num_heads=8,
            num_layers=num_layers,             # <--- Dynamic variable (Only once)
            expansion_factor=expansion_factor, # <--- Dynamic variable (Only once)
            dropout=0.1
        )
        # No cov_to_t5_projection needed: encoder output is already d_model
        self.cross_adapter = CrossAttentionAdapter(d_model=d_model)

    def forward(self, covariates, encoder_hidden):
        cov_encoded = self.cov_encoder(covariates)   # (B, T, d_model) directly
        return self.cross_adapter(encoder_hidden, cov_encoded)


class ChronosDualPath(nn.Module):
    def __init__(
        self,
        t5_model: AutoModelForSeq2SeqLM,
        num_features: int,
        freeze_encoder: bool = True,
        num_layers: int = 3,
        expansion_factor: int = 2
    ):
        super().__init__()
        self.t5_model = t5_model
        self.d_model = t5_model.config.d_model

        # Initialize the Dual-Path Adapter
        self.adapter = DualPathAdapter(
            cov_dim=num_features,
            d_model=self.d_model,
            num_layers=num_layers,
            expansion_factor=expansion_factor
        )

        # Apply freezing configuration immediately
        self.configure_t5_freezing(freeze_encoder)

    def configure_t5_freezing(self, freeze_encoder: bool):
        """
        Option A (freeze_encoder=True): Encoder Frozen, Decoder Unfrozen.
        Option B (freeze_encoder=False): ALL Unfrozen (Full Fine-Tuning).
        """
        # 1. Handle Encoder
        for p in self.t5_model.encoder.parameters():
            p.requires_grad = not freeze_encoder

        # 2. Handle Decoder (Always unfrozen to learn new target distribution)
        for p in self.t5_model.decoder.parameters():
            p.requires_grad = True

        if freeze_encoder:
            print("✓ T5 Freezing: Option A (Encoder FROZEN, Decoder UNFROZEN)")
        else:
            print("✓ T5 Freezing: Option B (Full Fine-Tuning, ALL UNFROZEN)")

    #change
    def forward(
        self,
        input_ids: torch.Tensor,
        price: torch.Tensor,
        cov: torch.Tensor,
        labels: Optional[torch.Tensor] = None,
        decoder_input_ids: Optional[torch.Tensor] = None, # Required for Step E1
        attention_mask: Optional[torch.Tensor] = None,
        decoder_attention_mask: Optional[torch.Tensor] = None, # Future-proofing
    ):
        encoder_outputs = self.t5_model.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            return_dict=True
        )

        encoder_outputs.last_hidden_state = self.adapter(
            covariates=cov,
            encoder_hidden=encoder_outputs.last_hidden_state
        )

        # Pass explicit decoder inputs and masks
        outputs = self.t5_model(
            encoder_outputs=encoder_outputs,
            labels=labels,
            decoder_input_ids=decoder_input_ids,
            attention_mask=attention_mask,
            decoder_attention_mask=decoder_attention_mask,
            return_dict=True
        )

        return outputs
    # --- Generation Utilities ---

    def _restrict_logits_to_numeric(self, logits):
        start = self.t5_model.config.vocab_offset
        end = start + self.t5_model.config.num_bins
        mask = torch.full_like(logits, float("-inf"))
        mask[..., start:end] = logits[..., start:end]
        return mask

    def _sample_numeric(self, logits, temperature=1.0, top_p=0.9):
        logits = logits / max(temperature, 1e-8)
        probs = torch.softmax(logits, dim=-1)
        sorted_probs, sorted_idx = torch.sort(probs, descending=True)
        cum = sorted_probs.cumsum(dim=-1)
        cutoff = cum > top_p
        cutoff[..., 1:] = cutoff[..., :-1].clone()
        cutoff[..., 0] = False  # NEVER mask the top token
        sorted_probs = sorted_probs.masked_fill(cutoff, 0.0)
        probs = torch.zeros_like(probs).scatter_(1, sorted_idx, sorted_probs)
        probs = probs / probs.sum(dim=-1, keepdim=True).clamp(min=1e-8)  # Safe division
        return torch.multinomial(probs, num_samples=1)

    def generate(
            self,
            input_ids: torch.Tensor,
            price: torch.Tensor,
            cov: torch.Tensor,
            attention_mask: Optional[torch.Tensor] = None,
            max_length: int = 64,
            **generate_kwargs
        ):

        # CRITICAL FIX: Ensure pipeline contract matches model capability
        num_seq = generate_kwargs.get("num_return_sequences", 1)
        assert num_seq == 1, "Current custom generate() only supports num_return_sequences=1"

        # 1. Run Encoder
        enc_outputs = self.t5_model.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            return_dict=True
        )
        # 2. 🟢 Apply Adapter to Encoder Outputs
        enc_outputs.last_hidden_state = self.adapter(
            covariates=cov,
            encoder_hidden=enc_outputs.last_hidden_state
        )

        # 3. Standard Decoder Loop
        device = input_ids.device
        B = input_ids.shape[0]
        dec_input = torch.full(
            (B, 1),
            self.t5_model.config.bos_token_id,
            dtype=torch.long,
            device=device
        )

        outputs = []
        for _ in range(max_length):
            out = self.t5_model(
                encoder_outputs=enc_outputs,
                decoder_input_ids=dec_input,
                use_cache=True,
                return_dict=True
            )
            logits = out.logits[:, -1, :]
            logits = self._restrict_logits_to_numeric(logits)
            next_tok = self._sample_numeric(
                logits=logits,
                temperature=generate_kwargs.get("temperature", 1.0),
                top_p=generate_kwargs.get("top_p", 0.9),
            )
            # Remove early break to ensure exactly max_length outputs
            # if (next_tok == self.t5_model.config.eos_token_id).all():
            #     break
            dec_input = torch.cat([dec_input, next_tok], dim=1)
            outputs.append(next_tok)

        return torch.cat(outputs, dim=1)

# =============================================================================
# SECTION E: DATASET
# =============================================================================

class TimeSeriesDataset(Dataset):
    """
    LCS Upgrade: Stores raw prices, scaled input_ids/labels, and local scale.
    p_last is removed; reconstruction is now: decoded_token * scale.
    """
    def __init__(
        self,
        prices: np.ndarray,
        covariates: np.ndarray,
        input_ids: np.ndarray,      # LCS: scaled-price tokens
        labels: np.ndarray,         # LCS: scaled-price tokens
        future_prices: np.ndarray,
        scale: np.ndarray,          # LCS: local scale per window (replaces p_last)
        attention_mask: Optional[np.ndarray] = None,
        pad_token_id: int = 0,
    ):
        self.prices = torch.FloatTensor(prices)
        self.covariates = torch.FloatTensor(covariates)
        self.input_ids = torch.LongTensor(input_ids)
        self.labels = torch.LongTensor(labels)
        self.future_prices = torch.FloatTensor(future_prices)
        self.scale = torch.FloatTensor(scale)          # LCS

        if attention_mask is None:
            self.attention_mask = (self.input_ids != pad_token_id).long()
        else:
            self.attention_mask = torch.LongTensor(attention_mask)

    def __len__(self):
        return self.prices.shape[0]

    def __getitem__(self, idx):
        return {
            "price": self.prices[idx],
            "cov": self.covariates[idx],
            "input_ids": self.input_ids[idx],
            "labels": self.labels[idx],
            "future_prices": self.future_prices[idx],
            "scale": self.scale[idx],               # LCS (replaces p_last)
            "attention_mask": self.attention_mask[idx],
        }

# =============================================================================
# SECTION F: TRAINING CONFIG & SCHEDULER
# =============================================================================

@dataclass
class TrainingConfig:
    lr: float = 1e-4
    lr_end: float = 3e-5
    warmup_ratio: float = 0.05
    weight_decay: float = 0.01
    epochs: int = 12
    batch_size: int = 32
    gradient_clip: float = 1.0
    gradient_accumulation_steps: int = 1
    mixed_precision: bool = False
    patience: int = 6  # Stop if no improvement after 6 epochs
    device: str = "cuda" if torch.cuda.is_available() else "cpu"

    lambda_price: float = 10000.0  # weight for price MSE term

    # --- NEW TOGGLE ---
    # True = Option A (Only Decoder + Adapter learn)
    # False = Option B (Encoder + Decoder + Adapter learn)
    freeze_encoder: bool = True

    save_dir: str = "./checkpoints"
    log_every_n_steps: int = 50

    def save(self, path: str):
        with open(path, "w") as f:
            json.dump(asdict(self), f, indent=2)

    @classmethod
    def load(cls, path: str):
        with open(path, "r") as f:
            return cls(**json.load(f))


def get_custom_cosine_schedule(
    optimizer,
    num_warmup_steps: int,
    num_training_steps: int,
    lr_end_factor: float = 0.3,
):
    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        progress = float(current_step - num_warmup_steps) / float(
            max(1, num_training_steps - num_warmup_steps)
        )
        cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
        return max(lr_end_factor, cosine)

    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)


# =============================================================================
# SECTION G: TRAINER (HYBRID LOSS + METRICS)
# =============================================================================

class CGCXTrainer:
    def __init__(
        self,
        model,
        config: TrainingConfig,
        normalizer: ZScoreNormalizer,
        train_loader: DataLoader,
        val_loader: Optional[DataLoader],
        tokenizer: SimpleChronosTokenizer,
        mase_scale: float = 1.0,
        cov_cols: Optional[List[str]] = None,
    ):
        self.model = model.to(config.device)
        self.config = config
        self.normalizer = normalizer
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.device = config.device
        self.tokenizer = tokenizer

        assert cov_cols is not None, "cov_cols must be provided"
        self.cov_cols = cov_cols

        # Collect ALL parameters that have requires_grad=True
        # This includes Adapter + Decoder + (optionally) Encoder
        trainable_params = []
        for name, p in model.named_parameters():
            if p.requires_grad:
                trainable_params.append(p)

        self.optimizer = torch.optim.AdamW(
            trainable_params,
            lr=config.lr, weight_decay=config.weight_decay,
        )

        steps_per_epoch = math.ceil(len(train_loader) / config.gradient_accumulation_steps)
        num_training_steps = steps_per_epoch * config.epochs

        self.scheduler = get_custom_cosine_schedule(
            self.optimizer,
            int(num_training_steps * config.warmup_ratio),
            num_training_steps,
            lr_end_factor=config.lr_end / config.lr # Fix Bug 5 here too
        )
        self.scaler = GradScaler('cuda') if config.mixed_precision else None
        self.mase_scale = float(max(mase_scale, 1e-8))

        # Precompute bin VALUES (scaled prices in [0.0, 5.0])
        num_bins = tokenizer.num_bins
        vocab_size = model.t5_model.config.vocab_size
        offset = model.t5_model.config.vocab_offset
        bin_idx = torch.arange(num_bins, dtype=torch.float32)
        scaled = bin_idx / (num_bins - 1)
        bin_values = scaled * tokenizer.range + tokenizer.min_val  # [0.0 … 5.0]

        full_bin_values = torch.zeros(vocab_size, dtype=torch.float32)
        full_bin_values[offset:offset + num_bins] = bin_values
        self.bin_values = full_bin_values.view(1, 1, -1).to(self.device)
        self.bin_values.requires_grad_(False)

        print("✓ Trainer initialized (LCS Mode: scaled prices in [0.0, 5.0])")

    def _expected_scaled_price_from_logits(self, logits: torch.Tensor) -> torch.Tensor:
        """Returns expected scaled-price value from logits (LCS mode)."""
        masked_logits = self.model._restrict_logits_to_numeric(logits)
        probs = torch.softmax(masked_logits, dim=-1)
        exp_scaled_price = (probs * self.bin_values).sum(dim=-1)
        return exp_scaled_price

    def _scaled_prices_from_labels(self, labels: torch.Tensor) -> torch.Tensor:
        """
        Maps token IDs (labels) to their numeric scaled-price values.
        """
        # self.bin_values is shape (1, 1, vocab_size) -> flatten to (V,)
        vocab_vals = self.bin_values.view(-1)

        # Use gather without clamping. If labels are invalid, this will crash (Good).
        return vocab_vals.gather(0, labels.view(-1)).view_as(labels)

    #CHANGE
    def train_epoch(self, epoch: int) -> Dict[str, float]:
        self.model.train()
        total_loss = 0.0
        total_ret_mse = 0.0
        total_count = 0

        pad_id = self.tokenizer.pad_token_id
        bos_id = self.model.t5_model.config.bos_token_id
        self.optimizer.zero_grad()

        for step, batch in enumerate(self.train_loader):
            input_ids = batch["input_ids"].to(self.device)
            price = batch["price"].to(self.device)
            cov = batch["cov"].to(self.device)
            labels = batch["labels"].to(self.device)
            mask = batch["attention_mask"].to(self.device)

            cov_norm = self.normalizer.transform(cov)

            # --- Prepare Decoder Inputs (Teacher Forcing) ---
            # Shift labels right: [BOS, L1, L2, ..., Ln-1]
            B, H = labels.shape
            decoder_input_ids = torch.cat(
                [torch.full((B, 1), bos_id, device=self.device, dtype=torch.long), labels[:, :-1]],
                dim=1
            )

            # Build decoder mask (crucial if padding ever exists)
            decoder_attention_mask = (decoder_input_ids != pad_id).long()

            with autocast('cuda', enabled=self.config.mixed_precision):
                # Call model with explicit decoder inputs and NO labels
                outputs = self.model(
                    input_ids=input_ids, price=price, cov=cov_norm,
                    labels=None,
                    decoder_input_ids=decoder_input_ids,
                    attention_mask=mask,
                    decoder_attention_mask=decoder_attention_mask
                )
                logits = outputs.logits # (B, H, V)

                # --- 1. Numeric-Only Cross Entropy ---
                masked_logits = self.model._restrict_logits_to_numeric(logits)

                ce_loss = F.cross_entropy(
                    masked_logits.reshape(-1, masked_logits.size(-1)),
                    labels.reshape(-1),
                    ignore_index=pad_id
                )

                # --- 2. LCS Scaled-Price MSE Regression ---
                # Predict expected scaled price, compare with true scaled-price tokens
                pred_scaled = self._expected_scaled_price_from_logits(logits)
                true_scaled = self._scaled_prices_from_labels(labels)

                # Masked MSE (valid tokens only)
                label_mask = (labels != pad_id).float()
                mse_elem = (pred_scaled - true_scaled) ** 2
                ret_mse_loss = (mse_elem * label_mask).sum() / (label_mask.sum() + 1e-8)

                # Hybrid Loss
                hybrid_loss = ce_loss + self.config.lambda_price * ret_mse_loss
                loss = hybrid_loss / self.config.gradient_accumulation_steps

            if self.scaler:
                self.scaler.scale(loss).backward()
            else:
                loss.backward()

            if (step + 1) % self.config.gradient_accumulation_steps == 0 or (step + 1) == len(self.train_loader):
                if self.scaler:
                    self.scaler.unscale_(self.optimizer)
                    torch.nn.utils.clip_grad_norm_([p for p in self.model.parameters() if p.requires_grad], self.config.gradient_clip)
                    self.scaler.step(self.optimizer)
                    self.scaler.update()
                else:
                    #
                    torch.nn.utils.clip_grad_norm_([p for p in self.model.parameters() if p.requires_grad], self.config.gradient_clip)
                    # NaN guard: skip the step entirely if any gradient is NaN
                    has_nan_grad = any(
                            p.grad is not None and not torch.isfinite(p.grad).all()
                            for p in self.model.parameters() if p.requires_grad
                    )
                    if has_nan_grad:
                            print(f"  ⚠️ NaN gradient detected at step {step}. Skipping optimizer step.")
                            self.optimizer.zero_grad()
                    else:
                            self.optimizer.step()

                self.scheduler.step()
                self.optimizer.zero_grad()

            # Logging (unscaled)
            total_loss += hybrid_loss.item() * self.config.gradient_accumulation_steps
            total_ret_mse += ret_mse_loss.item() * self.config.gradient_accumulation_steps
            total_count += 1

            if step % self.config.log_every_n_steps == 0:
                print(f"  [Epoch {epoch+1} Step {step}] Loss: {hybrid_loss.item() * self.config.gradient_accumulation_steps:.4f} (CE:{ce_loss.item():.3f} ScaledMSE:{ret_mse_loss.item():.5f})")

        return {
            "train_loss": total_loss / max(1, total_count),
            "train_scaled_rmse": math.sqrt(total_ret_mse / max(1, total_count))
        }

    @torch.no_grad()
    def validate(self):
        self.model.eval()
        total_price_sse = 0.0
        total_tokens = 0
        pad_id = self.tokenizer.pad_token_id
        bos_id = self.model.t5_model.config.bos_token_id

        for batch in self.val_loader:
            input_ids = batch["input_ids"].to(self.device)
            cov = batch["cov"].to(self.device)
            future_prices = batch["future_prices"].to(self.device)
            scale = batch["scale"].to(self.device)   # LCS: local scale (replaces p_last)
            mask = batch["attention_mask"].to(self.device)
            labels = batch["labels"].to(self.device)

            # Prepare Decoder Inputs
            B, H = labels.shape
            decoder_input_ids = torch.cat(
                [torch.full((B, 1), bos_id, device=self.device, dtype=torch.long), labels[:, :-1]],
                dim=1
            )
            decoder_attention_mask = (decoder_input_ids != pad_id).long()

            outputs = self.model(
                input_ids=input_ids, price=batch["price"].to(self.device),
                cov=self.normalizer.transform(cov),
                labels=None,
                decoder_input_ids=decoder_input_ids,
                attention_mask=mask,
                decoder_attention_mask=decoder_attention_mask
            )

            # LCS: expected scaled price → multiply by local scale to get absolute price
            pred_scaled = self._expected_scaled_price_from_logits(outputs.logits)
            pred_prices = pred_scaled * scale.unsqueeze(1)   # NO cumsum, NO exp

            valid_mask = (labels != pad_id)
            error_sq = (pred_prices - future_prices) ** 2

            total_price_sse += (error_sq * valid_mask).sum().item()
            total_tokens += valid_mask.sum().item()

        avg_mse = total_price_sse / total_tokens if total_tokens > 0 else 0.0
        return {"val_rmse": math.sqrt(avg_mse)}

    def train(self):
        print(f"🚀 Training Started (Max Epochs: {self.config.epochs}, Patience: {self.config.patience})")

        best_rmse = float("inf")
        patience_counter = 0  # <--- NEW: Track epochs without improvement

        for epoch in range(self.config.epochs):
            metrics = self.train_epoch(epoch)
            val_metrics = self.validate()
            #fix the bug for training missmatch
            #print(f"Epoch {epoch+1}: Train RMSE {metrics['train_rmse']:.4f} | Val RMSE {val_metrics['val_rmse']:.4f}")
            print(f"Epoch {epoch+1}: Train ScaledRMSE {metrics['train_scaled_rmse']:.4f} | Val PriceRMSE {val_metrics['val_rmse']:.4f}")

            # --- EARLY STOPPING LOGIC ---
            current_val_rmse = val_metrics["val_rmse"]

            if val_metrics["val_rmse"] < best_rmse:
                best_rmse = val_metrics["val_rmse"]
                patience_counter = 0  # Reset counter
                print(f"   🔥 New Best Model! Saving... (RMSE: {best_rmse:.4f})")

                Path(self.config.save_dir).mkdir(parents=True, exist_ok=True)
                torch.save(self.model.adapter.state_dict(), os.path.join(self.config.save_dir, "best_adapter.pt"))
                torch.save(self.model.t5_model.state_dict(), os.path.join(self.config.save_dir, "best_t5.pt"))
                self.tokenizer.save_config(os.path.join(self.config.save_dir, "tokenizer_config.json"))
                self.normalizer.save(self.config.save_dir)

                with open(os.path.join(self.config.save_dir, "model_config.json"), "w") as f:
                    #
                    json.dump({
                        "t5_base_model": "amazon/chronos-t5-base",
                        "vocab_size": self.model.t5_model.config.vocab_size,
                        "pad_token_id": self.model.t5_model.config.pad_token_id,
                        "eos_token_id": self.model.t5_model.config.eos_token_id,
                        "bos_token_id": self.model.t5_model.config.bos_token_id,
                        "vocab_offset": self.model.t5_model.config.vocab_offset,
                        "num_bins": self.model.t5_model.config.num_bins,
                        "num_features": self.model.adapter.cov_encoder.projection.in_features,
                        "d_model": self.model.d_model,
                        "embeddings_resized": True,
                        "freeze_encoder": self.config.freeze_encoder,
                        "cov_cols": self.cov_cols,
                        "target_representation": "local_scale",
                        # Strategy 1 architecture params (needed for correct reload)
                        "xlstm_num_layers": self.model.adapter.cov_encoder.num_layers,
                        "xlstm_expansion_factor": self.model.adapter.cov_encoder.blocks[0].up_proj.out_features // self.model.d_model,
                    }, f, indent=2)
            else:
                # Case B: No improvement
                patience_counter += 1
                print(f"   ⚠️ No improvement. Patience: {patience_counter}/{self.config.patience}")

                if patience_counter >= self.config.patience:
                    print(f"\n⏹️ Early Stopping Triggered at Epoch {epoch+1}!")
                    print(f"   Best Validation RMSE was: {best_rmse:.4f}")
                    break  # <--- Stop the loop immediately

        return {"best_val_rmse": best_rmse}

# =============================================================================
# SECTION H: INFERENCE WRAPPER
# =============================================================================

class CGCXInference:
    def __init__(
        self,
        model,
        normalizer,
        tokenizer,
        device: Optional[str] = None,
    ):
        if device is None:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device

        self.model = model.to(self.device)
        self.model.eval()
        self.normalizer = normalizer
        self.tokenizer = tokenizer


    #
    #

    @torch.no_grad()
    def predict(
        self,
        input_ids: torch.Tensor,
        price: torch.Tensor,
        covariates: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        num_samples: int = 20,
        max_length: int = 1,
        temperature: float = 1.0,
        top_p: float = 0.9,
    ) -> np.ndarray:

        cov_norm = self.normalizer.transform(covariates.to(self.device))
        price_np = price.cpu().numpy()
        local_scale = np.mean(np.abs(price_np), axis=1, keepdims=True)
        local_scale = np.clip(local_scale, 1e-8, None)

        B = input_ids.shape[0]
        all_pred_scaled = []

        # --- THE FIX: MAXIMUM INFERENCE CHUNKING ---
        # Process exactly 1 window at a time to prevent OOM
        for i in range(B):

            # Slice the current window (Batch size = 1)
            ids_single = input_ids[i:i+1].to(self.device)
            cov_single = cov_norm[i:i+1].to(self.device)
            mask_single = attention_mask[i:i+1].to(self.device) if attention_mask is not None else None

            window_samples = []

            # Generate the 20 samples one-by-one to save memory!
            for _ in range(num_samples):
                gen_tokens = self.model.generate(
                    input_ids=ids_single,
                    price=None,
                    cov=cov_single,
                    attention_mask=mask_single,
                    max_length=max_length,
                    num_return_sequences=1,
                    temperature=temperature,
                    top_p=top_p
                )

                # Decode the single token array
                gen_tokens_np = gen_tokens.cpu().numpy()
                pred_scaled_exp = self.tokenizer.decode(gen_tokens_np)

                if np.isnan(pred_scaled_exp).any():
                    pred_scaled_exp = np.nan_to_num(pred_scaled_exp, nan=1.0)

                window_samples.append(pred_scaled_exp)

            # We now have 20 samples for this 1 window. Stack them.
            # window_samples shape: (num_samples, 1, max_length)
            stacked_samples = np.stack(window_samples, axis=0)

            # Take the median across the samples dimension
            median_path = np.median(stacked_samples, axis=0) # Shape: (1, max_length)
            all_pred_scaled.append(median_path)

            # --- FORCE MEMORY CLEARING ---
            # Delete intermediate tensors and empty the CUDA cache
            # so the next chunk starts with a fresh 15GB.
            del ids_single, cov_single, mask_single, gen_tokens
            torch.cuda.empty_cache()

        # Reassemble all windows
        final_pred_scaled = np.concatenate(all_pred_scaled, axis=0)

        # Scale back to absolute prices
        pred_prices = final_pred_scaled * local_scale
        return pred_prices


###############################################
# NEW GLOBAL TIME-BASED MULTI-ASSET LOADER
###############################################

def load_multi_asset_time_splits(data_dir: str, T: int, H: int,
                                 max_files: int = 200,
                                 min_len: int = 120):

    train_prices_list = []
    train_cov_list = []
    train_labels_list = []

    val_prices_list = []
    val_cov_list = []
    val_labels_list = []

    train_scales_list = []
    val_scales_list   = []
    test_scales_list  = []

    test_prices_list = []
    test_cov_list = []
    test_labels_list = []

    files = [f for f in os.listdir(data_dir) if f.endswith(".csv")]
    files = files[:max_files]       # LIMIT TO PROTECT RAM

    print(f"Processing {len(files)} tickers...")

    for fname in files:
        path = os.path.join(data_dir, fname)

        # ---------------------- SAFE CSV LOAD ----------------------
        try:
            df = pd.read_csv(path)
        except Exception as e:
            print(f"⚠️ Skipping {fname}: cannot load ({e})")
            continue

        if "Date" not in df.columns:
            print(f"⚠️ Skipping {fname}: no Date column")
            continue
        if not {"Open","High","Low","Close"}.issubset(df.columns):
            print(f"⚠️ Skipping {fname}: missing OHLC")
            continue

        df["Date"] = pd.to_datetime(df["Date"], errors="coerce")
        df = df.dropna(subset=["Date"])
        df = df.sort_values("Date").drop_duplicates(subset=["Date"])
        df = df.set_index("Date")

        # Business-day alignment
        #df = df.resample("B").asfreq().ffill().bfill()
        df = df.resample("B").asfreq().ffill()
        df = df.dropna()
        df["price"] = df["Close"]

        if len(df) < min_len:
            continue

        # ------------- INDICATORS (heavy, per ticker only) ----------
        df = add_technical_indicators(df)
        df = df.dropna()

        if len(df) < T + H:
            continue

        # ------------- GLOBAL TIME SPLITS (NO DATA LEAKAGE) ---------
        df_train = df.loc[:TRAIN_END]
        df_val   = df.loc[TRAIN_END + pd.Timedelta(days=1): VAL_END]
        df_test  = df.loc[VAL_END + pd.Timedelta(days=1):]

        # Helper: turn df → sliding windows (LCS: returns 4-tuple including scales)
        def process(split_df, P_list, C_list, L_list, S_list):
            if len(split_df) >= T + H:
                result = build_windows_from_df(split_df, T=T, H=H)
                p, c, y, s = result[0], result[1], result[2], result[3]
                # Only append if windows were actually generated
                if len(p) > 0:
                    P_list.append(p); C_list.append(c)
                    L_list.append(y); S_list.append(s)

        process(df_train, train_prices_list, train_cov_list, train_labels_list, train_scales_list)
        process(df_val,   val_prices_list,   val_cov_list,   val_labels_list,   val_scales_list)
        process(df_test,  test_prices_list,  test_cov_list,  test_labels_list,  test_scales_list)

    # ----------------- CONCAT WITH SAFETY CHECKS ----------------------
    def safe_concat(arr_list, name):
        if not arr_list:
            raise ValueError(f"No data generated for {name} split. Check Date range/Data quality.")
        return np.concatenate(arr_list, axis=0)

    # We enforce training data exists, but val/test might technically be empty in some edge cases
    # (though usually we want them to exist).
    train_prices = safe_concat(train_prices_list, "Train")
    train_cov    = safe_concat(train_cov_list, "Train")
    train_labels = safe_concat(train_labels_list, "Train")
    train_scales = safe_concat(train_scales_list, "Train")

    val_prices   = safe_concat(val_prices_list, "Validation")
    val_cov      = safe_concat(val_cov_list, "Validation")
    val_labels   = safe_concat(val_labels_list, "Validation")
    val_scales   = safe_concat(val_scales_list, "Validation")

    test_prices  = safe_concat(test_prices_list, "Test")
    test_cov     = safe_concat(test_cov_list, "Test")
    test_labels  = safe_concat(test_labels_list, "Test")
    test_scales  = safe_concat(test_scales_list, "Test")

    return (
        train_prices, train_cov, train_labels, train_scales,
        val_prices,   val_cov,   val_labels,   val_scales,
        test_prices,  test_cov,  test_labels,  test_scales,
    )

# =============================================================================
# SECTION J: POST-TRAINING TEST PIPELINE (METRICS + PLOTTING)
# =============================================================================

def calculate_metrics(y_true, y_pred):
    """Calculates comprehensive regression metrics."""
    # Handle NaNs if any
    mask = np.isfinite(y_true) & np.isfinite(y_pred)
    y_true = y_true[mask]
    y_pred = y_pred[mask]

    if len(y_true) == 0:
        return {}

    mse = mean_squared_error(y_true, y_pred)
    mae = mean_absolute_error(y_true, y_pred)
    rmse = np.sqrt(mse)

    # Mean Absolute Percentage Error (avoid div by zero)
    mape = np.mean(np.abs((y_true - y_pred) / (np.abs(y_true) + 1e-8))) * 100

    # R-squared
    r2 = r2_score(y_true, y_pred)

    return {
        "MSE": mse,
        "RMSE": rmse,
        "MAE": mae,
        "MAPE": f"{mape:.2f}%",
        "R2": r2
    }

######here
def load_checkpoint_with_config(
    checkpoint_dir: str,
    device: str = "cuda",
    strict_schema: bool = True,
    override_freeze_encoder: Optional[bool] = None  # <--- Allows switching fine-tuning modes
):
    print(f"\n🔧 Loading checkpoint from {checkpoint_dir}...")

    tok_path = os.path.join(checkpoint_dir, "tokenizer_config.json")
    mod_path = os.path.join(checkpoint_dir, "model_config.json")
    adp_path = os.path.join(checkpoint_dir, "best_adapter.pt")

    if not all(os.path.exists(p) for p in [tok_path, mod_path, adp_path]):
        raise FileNotFoundError("Missing checkpoint files.")

    with open(tok_path, "r") as f: tconfig = json.load(f)
    with open(mod_path, "r") as f: mconfig = json.load(f)

    saved_cov_cols = mconfig.get("cov_cols", None)
    assert saved_cov_cols is not None, "Checkpoint missing 'cov_cols'"

    if strict_schema:
        saved_mode = mconfig.get("target_representation", "price")
        tok_mode = tconfig.get("mode", "price")
        assert saved_mode == "local_scale", f"Model mode '{saved_mode}' != 'local_scale'"
        assert tok_mode == "local_scale", f"Tokenizer mode '{tok_mode}' != 'local_scale'"
        assert saved_cov_cols == COV_COLS, "Global COV_COLS mismatch"

    tokenizer = SimpleChronosTokenizer.load_config(tok_path)
    normalizer = ZScoreNormalizer.load(checkpoint_dir)

    t5_model = AutoModelForSeq2SeqLM.from_pretrained(mconfig['t5_base_model'])
    if mconfig.get('embeddings_resized', False):
        t5_model.resize_token_embeddings(mconfig['vocab_size'])

    for k in ["vocab_size", "pad_token_id", "eos_token_id", "bos_token_id", "vocab_offset", "num_bins"]:
        setattr(t5_model.config, k, mconfig[k])

    # --- THE DYNAMIC MODE SWITCHER ---
    if override_freeze_encoder is not None:
        # User is forcing a mode switch (e.g., from Standard to Full Fine-Tuning)
        freeze_setting = override_freeze_encoder
        print(f"⚠️ Overriding saved freeze state. Forcing freeze_encoder={freeze_setting}")
    else:
        # Fallback to whatever was saved in JSON (handles the legacy "freeze_backbone" key)
        freeze_setting = mconfig.get("freeze_encoder", mconfig.get("freeze_backbone", True))

    num_layers = mconfig.get("xlstm_num_layers", 3)
    expansion_factor = mconfig.get("xlstm_expansion_factor", 2)

    model = ChronosDualPath(
        t5_model=t5_model,
        num_features=mconfig['num_features'],
        freeze_encoder=freeze_setting,
        num_layers=num_layers,                 # <--- LOADED FROM JSON
        expansion_factor=expansion_factor      # <--- LOADED FROM JSON
    )

    # Load the raw state dict
    raw_state_dict = torch.load(adp_path, map_location=torch.device(device))

    # --- torch.compile Fix ---
    # Strip the "_orig_mod." prefix from the compiled weights so they match the standard model
    clean_state_dict = {}
    for key, value in raw_state_dict.items():
        clean_key = key.replace("_orig_mod.", "")
        clean_state_dict[clean_key] = value

    model.adapter.load_state_dict(clean_state_dict)

    t5_path = os.path.join(checkpoint_dir, "best_t5.pt")
    if os.path.exists(t5_path):
        model.t5_model.load_state_dict(torch.load(t5_path, map_location=torch.device(device)))
        print("✓ T5 weights restored from checkpoint")
    else:
        print("⚠️ No best_t5.pt found — using pretrained T5 weights (old checkpoint)")

    model.to(device)
    model.eval()

    print(f"✓ Model loaded successfully (Mode: local_scale)")
    return model, tokenizer, normalizer, saved_cov_cols


def run_test_pipeline(
    data_dir: str,
    model,
    normalizer,
    tokenizer,
    cov_cols: list,
    T: int = 64,
    H: int = 16,
    num_symbols: int = 5
):
    print("\n" + "="*60)
    print("🚀 STARTING AUTOMATED TEST PIPELINE (LCS: Local Context Scaling)")
    print("="*60)

    all_files = [f for f in os.listdir(data_dir) if f.endswith(".csv")]
    if not all_files: return

    selected_files = random.sample(all_files, min(num_symbols, len(all_files)))
    infer = CGCXInference(model, normalizer, tokenizer)

    for fname in selected_files:
        symbol = fname.replace('.csv', '')
        print(f"\n🔍 Analyzing Symbol: {symbol}")
        path = os.path.join(data_dir, fname)

        try:
            df = pd.read_csv(path, parse_dates=["Date"])
            df = df.sort_values("Date").drop_duplicates(subset=["Date"]).set_index("Date")
            #df = df.resample("B").asfreq().ffill().bfill()
            df = df.resample("B").asfreq().ffill()
            df = df.dropna()
            df["price"] = df["Close"]
            df = add_technical_indicators(df).dropna()

            df_test = df.loc[VAL_END + pd.Timedelta(days=1):]
            if len(df_test) < T + H: continue

            prices_test, cov_test, labels_test, scales_test = build_windows_from_df(
                df_test, T=T, H=H, columns=cov_cols
            )
            if len(prices_test) == 0: continue

            if np.any(prices_test <= 0):
                print(f"Skipping {fname}: non-positive prices")
                continue

            # LCS: tokenize scaled prices (price / local_scale)
            local_scales = np.mean(np.abs(prices_test), axis=1, keepdims=True)  # (N, 1)
            local_scales = np.clip(local_scales, 1e-8, None)
            scaled_input = prices_test / local_scales

            input_ids_test = np.stack([tokenizer.encode(row) for row in scaled_input])
            input_ids_t = torch.LongTensor(input_ids_test)
            price_t = torch.FloatTensor(prices_test)
            cov_t = torch.FloatTensor(cov_test)

            # Pass Mask (Fix Bug 3)
            input_mask = (input_ids_t != tokenizer.pad_token_id).long()

            # LCS Diagnostics: token diversity (scaled prices)
            pad_id = tokenizer.pad_token_id
            unique_tokens = [len(np.unique(row[row != pad_id])) for row in input_ids_test]
            print(f"   LCS Token Diversity: {np.mean(unique_tokens):.1f}/{T}")

            print(f"   🤖 Predicting {len(prices_test)} sequences...")
            preds = infer.predict(
                input_ids=input_ids_t,
                price=price_t,
                covariates=cov_t,
                attention_mask=input_mask,
                num_samples=20,
                temperature=1.0,
                max_length=H
            )

            metrics = calculate_metrics(labels_test.flatten(), preds.flatten())
            rmse = metrics.get("RMSE")
            print(f"   📊 All-Window RMSE: {rmse:.4f}" if rmse else "   ⚠️ RMSE: N/A")

        except Exception as e:
            print(f"Error: {e}")
            import traceback
            traceback.print_exc()


####################complete visualizer############################

def run_evaluations(cursor, current_date, actuals_df, model_version="v1.0"):
    print("\n--- Running Daily Evaluations ---")
    import pandas as pd   
    import numpy as np    
    
    # 1. Get previous predictions that haven't been evaluated yet
    cursor.execute("""
        SELECT date, entity_id, predicted_value 
        FROM predictions 
        WHERE model_version = %s AND date < %s
        ORDER BY date DESC LIMIT 400
    """, (model_version, current_date))
    
    #
    past_preds = cursor.fetchall()
    if not past_preds:
        print("No previous predictions to evaluate yet.")
        return

    # Sanitize actuals index and compute exact grading dates
    actuals_df.index = pd.to_datetime(actuals_df.index).tz_localize(None).date
    valid_dates = sorted([d for d in actuals_df.index if d < current_date])
    if len(valid_dates) < 2:
        print("Not enough market data to evaluate.")
        return
    target_date = valid_dates[-1]   # yesterday's close
    prev_date   = valid_dates[-2]   # day before that


    evals_to_insert = []
    
    # Variables for the daily summary
    abs_errors_pct = []
    abs_errors_dollar = []
    direction_hits = 0
    pred_returns = []
    act_returns = []
    
    for row in past_preds:
        pred_date, symbol, predicted_price = row
        
        # Check if we have the actual data for this symbol today
        if symbol in actuals_df.columns:
            # Get actual price and previous close from yfinance data
            try:
                # Assuming actuals_df has the latest closes
                actual_price = float(actuals_df.loc[target_date, symbol])
                prev_close = float(actuals_df.loc[prev_date, symbol])
                
                # Math
                error_dollar = predicted_price - actual_price
                error_pct = error_dollar / actual_price
                
                pred_ret = (predicted_price - prev_close) / prev_close
                act_ret = (actual_price - prev_close) / prev_close
                
                # Did we guess the direction correctly?
                direction_hit = (pred_ret > 0 and act_ret > 0) or (pred_ret < 0 and act_ret < 0)
                
                evals_to_insert.append((
                    target_date, symbol, predicted_price, actual_price, 
                    error_dollar, error_pct, pred_ret, act_ret, direction_hit, model_version
                ))
                
                # Store for summary math
                abs_errors_dollar.append(abs(error_dollar))
                abs_errors_pct.append(abs(error_pct))
                pred_returns.append(pred_ret)
                act_returns.append(act_ret)
                if direction_hit:
                    direction_hits += 1
                    
            except Exception as e:
                continue
    
    #
    # --- NEW: Save raw actuals to the DB (Stage B) ---
    actuals_to_insert = []
    for symbol in actuals_df.columns:
        try:
            today_actual_price = float(actuals_df.loc[target_date, symbol])
            actuals_to_insert.append((target_date, symbol, today_actual_price))
        except:
            continue
            
    if actuals_to_insert:
        cursor.executemany("""
            INSERT INTO actuals (date, symbol, actual_value)
            VALUES (%s, %s, %s)
            ON CONFLICT DO NOTHING;
        """, actuals_to_insert)

    # 2. Save Evaluations to DB
    if evals_to_insert:
        cursor.executemany("""
            INSERT INTO evaluations (date, symbol, predicted_value, actual_value, error_dollar, error_pct, predicted_return, actual_return, direction_hit, model_version)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            ON CONFLICT DO NOTHING;
        """, evals_to_insert)
        
        # 3. Calculate and Save Daily Summary
        
        summary_df = pd.DataFrame({"pred": pred_returns, "act": act_returns})
        
        # --- THE FIX: Strip Numpy types and handle NaNs safely ---
        try:
            spearman = float(summary_df.corr(method='spearman').iloc[0, 1])
            if np.isnan(spearman): spearman = 0.0
        except:
            spearman = 0.0
            
        try:
            mae_dollar = float(np.nanmean(abs_errors_dollar))
            if np.isnan(mae_dollar): mae_dollar = 0.0
        except:
            mae_dollar = 0.0
            
        try:
            mae_pct = float(np.nanmean(abs_errors_pct))
            if np.isnan(mae_pct): mae_pct = 0.0
        except:
            mae_pct = 0.0
            
        dir_acc = float(direction_hits / len(evals_to_insert)) if evals_to_insert else 0.0
        
        cursor.execute("""
            INSERT INTO daily_summary (date, model_version, symbols_count, mae_dollar, mae_pct, directional_accuracy, spearman_corr)
            VALUES (%s, %s, %s, %s, %s, %s, %s)
            ON CONFLICT DO NOTHING;
        """, (target_date, model_version, len(evals_to_insert), mae_dollar, mae_pct, dir_acc, spearman))
        
        print(f"✅ Evaluated {len(evals_to_insert)} symbols. MAE: {mae_pct:.2%}, Accuracy: {dir_acc:.2%}")


# =============================================================================
# SECTION I: MAIN PIPELINE (MULTI-ASSET)
# =============================================================================
def main():

    torch.set_float32_matmul_precision('high')   # Enables TF32, safe for compile
    torch.backends.cuda.matmul.allow_tf32 = True  # Belt-and-suspenders

    print("Start LCS (Local Context Scaling)")
    DATA_DIR = "raw_data"
    T = 64
    H = 1
    epochs= 12
    batch_size=16
    lambda_price=10000

    # NOTE: Set this to False if you want Option B (Full Fine-Tuning)
    # If False, LOWER the LR to 1e-5 to avoid destroying weights.
    freeze_encoder_setting = True
    lr = 1e-4 if freeze_encoder_setting else 1e-5
    #########


    print("Load multi-asset windows")
    (train_prices, train_cov, train_labels, train_scales,
     val_prices,   val_cov,   val_labels,   val_scales,
     test_prices,  test_cov,  test_labels,  test_scales) = load_multi_asset_time_splits(DATA_DIR, T=T, H=H)

    print("Tokenizer (LCS: scaled prices in [0.0, 5.0])")
    tokenizer = SimpleChronosTokenizer(num_bins=1024, min_val=0.0, max_val=5.0)

    # --- HELPER: Prepare LCS Dataset ---
    def prepare_dataset_tensors(prices, future_prices, scales, tokenizer_obj):
        """
        LCS Upgrade: Encodes windows as scaled prices, not log-returns.
        scale = mean(abs(history)) per window.
        input_ids = tokenize(history / scale)
        label_ids = tokenize(future / scale)
        """
        # 1. Filter invalid windows
        valid_mask = (np.min(prices, axis=1) > 0) & (np.min(future_prices, axis=1) > 0)

        if not np.all(valid_mask):
            drop_count = np.sum(~valid_mask)
            print(f"⚠️ Dropping {drop_count}/{len(prices)} windows with non-positive prices.")
            prices = prices[valid_mask]
            future_prices = future_prices[valid_mask]
            scales = scales[valid_mask]

            if len(prices) == 0:
                raise ValueError("All data dropped due to non-positive prices.")

        # 2. LCS: scale each window by its local scale
        scaled_input  = prices       / scales[:, None]        # (N, T)
        scaled_future = future_prices / scales[:, None]        # (N, H)

        # 3. Tokenize scaled prices
        input_ids = np.stack([tokenizer_obj.encode(row) for row in scaled_input])
        label_ids = np.stack([tokenizer_obj.encode(row) for row in scaled_future])

        # 4. Mask
        attention_mask = (input_ids != tokenizer_obj.pad_token_id).astype(np.int64)

        return input_ids, label_ids, attention_mask, scales, prices, future_prices, valid_mask

    tr_in, tr_lb, tr_mk, tr_sc, tr_p_safe, tr_l_safe, tr_valid = prepare_dataset_tensors(train_prices, train_labels, train_scales, tokenizer)
    va_in, va_lb, va_mk, va_sc, va_p_safe, va_l_safe, va_valid = prepare_dataset_tensors(val_prices, val_labels, val_scales, tokenizer)

    # Apply mask to covariates to maintain alignment
    train_cov = train_cov[tr_valid]
    val_cov   = val_cov[va_valid]

    # FIX: Assert labels are valid numeric bins
    vocab_start = tokenizer.vocab_offset
    vocab_end = vocab_start + tokenizer.num_bins

    assert tr_lb.min() >= vocab_start and tr_lb.max() < vocab_end, "Train labels contain invalid tokens"
    assert va_lb.min() >= vocab_start and va_lb.max() < vocab_end, "Val labels contain invalid tokens"

    for special in [tokenizer.pad_token_id, tokenizer.bos_token_id, tokenizer.eos_token_id]:
        assert not np.any(tr_lb == special), f"Train labels contain special token {special}"
        assert not np.any(va_lb == special), f"Val labels contain special token {special}"

    print("Normalizer")
    normalizer = ZScoreNormalizer()
    normalizer.fit(torch.FloatTensor(train_cov))

    config = TrainingConfig(
        lr=lr, epochs=epochs, batch_size=batch_size,
        save_dir="./checkpoints_multi", lambda_price=lambda_price,
        freeze_encoder=freeze_encoder_setting,
        mixed_precision= False
    )

    # LCS Dataset: uses scale instead of p_last
    train_dataset = TimeSeriesDataset(
        prices=tr_p_safe, covariates=train_cov, input_ids=tr_in,
        labels=tr_lb, future_prices=tr_l_safe, scale=tr_sc,
        attention_mask=tr_mk, pad_token_id=tokenizer.pad_token_id,
    )
    val_dataset = TimeSeriesDataset(
        prices=va_p_safe, covariates=val_cov, input_ids=va_in,
        labels=va_lb, future_prices=va_l_safe, scale=va_sc,
        attention_mask=va_mk, pad_token_id=tokenizer.pad_token_id,
    )

    train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True)
    val_loader   = DataLoader(val_dataset, batch_size=config.batch_size)

    print("Load chronos-T5")
    model_name = "amazon/chronos-t5-base"
    t5_model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
    t5_model.resize_token_embeddings(tokenizer.vocab_size)

    # Apply config
    t5_model.config.vocab_size = tokenizer.vocab_size
    t5_model.config.pad_token_id = tokenizer.pad_token_id
    t5_model.config.eos_token_id = tokenizer.eos_token_id
    t5_model.config.bos_token_id = tokenizer.bos_token_id
    t5_model.config.vocab_offset = tokenizer.vocab_offset
    t5_model.config.num_bins = tokenizer.num_bins

    model = ChronosDualPath(
        t5_model=t5_model,
        num_features=train_cov.shape[2],
        freeze_encoder=config.freeze_encoder
    )

    model.adapter = torch.compile(model.adapter,  dynamic=False)

    trainer = CGCXTrainer(
        model=model, config=config, normalizer=normalizer,
        train_loader=train_loader, val_loader=val_loader,
        tokenizer=tokenizer, mase_scale=1.0,
        cov_cols=COV_COLS
    )

    trainer.train()

    print("\n🔄 RELOADING Best Model...")
    reloaded_model, reloaded_tokenizer, reloaded_normalizer, reloaded_cov_cols = load_checkpoint_with_config(
        config.save_dir, device=config.device
    )

    run_test_pipeline(
        data_dir=DATA_DIR,
        model=reloaded_model,
        normalizer=reloaded_normalizer,
        tokenizer=reloaded_tokenizer,
        cov_cols=reloaded_cov_cols,
        T=T, H=H, num_symbols=5
    )


####tester code####
def evaluate_and_plot_symbol(symbol: str, checkpoint_dir: str, data_dir: str, T=64, H=1):
    """
    Loads the BEST model checkpoint and visualizes predictions for a specific symbol.
    """
    print(f"\n🎨 VISUALIZATION: Analyzing {symbol}...")

    # 1. Load the BEST model (Crucial: reloads the saved state, not the last epoch)
    # This uses your existing load_checkpoint_with_config which correctly handles xLSTM
    try:
        model, tokenizer, normalizer, cov_cols = load_checkpoint_with_config(
            checkpoint_dir, device="cuda" if torch.cuda.is_available() else "cpu"
        )
    except FileNotFoundError:
        print("❌ Checkpoint not found. Skipping visualization.")
        return

    infer = CGCXInference(model, normalizer, tokenizer)

    # 2. Load Data for the Symbol
    file_path = os.path.join(data_dir, f"{symbol}.csv")
    if not os.path.exists(file_path):
        print(f"❌ Data file for {symbol} not found.")
        return

    # Standard Preprocessing (Must match training)
    df = pd.read_csv(file_path, parse_dates=["Date"]).sort_values("Date").set_index("Date")
    df = df.resample("B").asfreq().ffill().dropna()
    df["price"] = df["Close"]
    df = add_technical_indicators(df).dropna()

    # Use only Test portion (Data the model has likely not seen)
    VAL_END_DATE = pd.Timestamp("2023-12-31")
    df_test = df.loc[VAL_END_DATE + pd.Timedelta(days=1):]

    if len(df_test) < T + H:
        print("❌ Not enough data for testing this symbol.")
        return

    # 3. Build Windows
    result = build_windows_from_df(df_test, T, H, columns=cov_cols)
    prices, cov, labels, scales = result[0], result[1], result[2], result[3]
    if len(prices) == 0: return

    # 4. LCS: Tokenize scaled input prices
    local_scales = np.mean(np.abs(prices), axis=1, keepdims=True)  # (N, 1)
    local_scales = np.clip(local_scales, 1e-8, None)
    scaled_input = prices / local_scales

    input_ids = np.stack([tokenizer.encode(row) for row in scaled_input])

    # Convert to Tensor
    input_ids_t = torch.LongTensor(input_ids)
    price_t = torch.FloatTensor(prices)
    cov_t = torch.FloatTensor(cov)
    mask_t = (input_ids_t != tokenizer.pad_token_id).long()

    # 5. Predict
    print(f"   Generating predictions for {len(prices)} windows (Median of 20 samples)...")

    # --- ADD THIS WRAPPER TO PREVENT MEMORY LEAKS ---
    with torch.no_grad():
        preds = infer.predict(
            input_ids=input_ids_t,
            price=price_t,
            covariates=cov_t,
            attention_mask=mask_t,
            num_samples=20,
            temperature=1.0,
            max_length=H
        )

    # 6. Calculate Metrics
    metrics = calculate_metrics(labels.flatten(), preds.flatten())
    print(f"   📊 RMSE: {metrics.get('RMSE', 'N/A'):.4f}")
    print(f"   📊 R2:   {metrics.get('R2', 'N/A'):.4f}")

    # 7. Plotting (The new part you wanted)
    # Plot a random window or the last window
    true_curve = labels.flatten()
    pred_curve = preds.flatten()

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

    x_axis = np.arange(len(true_curve))

    # Plot the full continuous year of actual prices
    plt.plot(x_axis, true_curve, label='True Price (Entire Test Set)', color='black', linewidth=2)

    # Overlay the model's rolling 1-day predictions for the whole year
    plt.plot(x_axis, pred_curve, label='Model 1-Day Predictions', color='red', linestyle='--', linewidth=1.5, alpha=0.8)

    plt.title(f"Full Test Set Evaluation: {symbol} ({len(true_curve)} Trading Days)")
    plt.xlabel("Days in Test Set")
    plt.ylabel("Absolute Price ($)")
    plt.legend()
    plt.grid(True, alpha=0.3)

    # Save plot
    save_path = f"prediction_{symbol}_FULL_YEAR.png"
    plt.savefig(save_path)
    print(f"   📸 Full timeline plot saved to {save_path}")
    plt.close()
####tester end#####


# ==========================================
# 2. PREDICTION LOGIC
# ==========================================
def fetch_latest_data(symbol: str, save_dir: str = "daily_test"):
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    ticker = yf.Ticker(symbol)
    df = ticker.history(period="150d")
    if df.empty:
        raise ValueError(f"Could not find data for symbol: {symbol}")
    df = df.reset_index()
    file_path = os.path.join(save_dir, f"{symbol}_latest.csv")
    df.to_csv(file_path, index=False)
    return file_path

def predict_tomorrow(csv_path: str, infer, tokenizer, cov_cols, T: int = 64):
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Notice we NO LONGER load the model here. It's passed in via 'infer', 'tokenizer', and 'cov_cols'

    df = pd.read_csv(csv_path)
    date_col = "Date" if "Date" in df.columns else df.columns[0]
    df[date_col] = pd.to_datetime(df[date_col], utc=True).dt.tz_localize(None).dt.normalize()
    df = df.sort_values(date_col).set_index(date_col)
    df = df.resample("B").asfreq().ffill().dropna()

    if "Close" in df.columns: df["price"] = df["Close"]
    elif "close" in df.columns: df["price"] = df["close"]
    else: return None

    df = add_technical_indicators(df).dropna()
    df_window = df.iloc[-T:]

    prices_np = df_window["price"].values.astype("float32")
    cov_np = df_window[cov_cols].values.astype("float32")
    local_scale = np.mean(np.abs(prices_np))
    scaled_prices = prices_np / local_scale

    input_ids = tokenizer.encode(scaled_prices)
    input_ids_t = torch.LongTensor(input_ids).unsqueeze(0).to(device)
    price_t = torch.FloatTensor(prices_np).unsqueeze(0).to(device)
    cov_t = torch.FloatTensor(cov_np).unsqueeze(0).to(device)
    mask_t = (input_ids_t != tokenizer.pad_token_id).long().to(device)

    with torch.no_grad():
        pred_prices = infer.predict(
            input_ids=input_ids_t, price=price_t, covariates=cov_t,
            attention_mask=mask_t, num_samples=20, max_length=1, temperature=1.0
        )

    return float(pred_prices[0, 0])


# 1. Setup Modal Image with your specific libraries
# Added 'lxml' because pandas.read_html requires it!
image = modal.Image.debian_slim().pip_install(
    "torch", 
    "pandas", 
    "numpy",
    "yfinance",
    "requests",
    "lxml", 
    "psycopg2-binary",
    "transformers",
    "scikit-learn",
    "matplotlib"
)

app = modal.App("daily-ml-job")
volume = modal.Volume.from_name("model-volume")

@app.function(
    image=image,
    volumes={"/data": volume}, 
    schedule=modal.Cron("0 11 * * 1-5"), # 8:00 AM UTC
    gpu="T4", 
    secrets=[modal.Secret.from_name("neon-db-secret")],
    timeout=1800 # Gives your script up to 30 mins to run since yfinance can be slow
)
def run_daily_inference():
    import yfinance as yf
    import pandas as pd
    import numpy as np
    import torch
    import requests
    import io
    import psycopg2
    from datetime import datetime, timezone
    print(f"⏰ SERVER WAKE-UP TIME (UTC): {datetime.now(timezone.utc)}")
    
    # ==========================================
    # PASTE YOUR CUSTOM CLASSES & FUNCTIONS HERE
    # (e.g., load_checkpoint_with_config, CGCXInference, add_technical_indicators)
    # ==========================================

    # --- HELPER FUNCTIONS ---
    HEADERS = {"User-Agent": "Mozilla/5.0"}

    def get_wiki_tickers(url, table_keywords):
        tickers = []
        try:
            response = requests.get(url, headers=HEADERS)
            response.raise_for_status()
            tables = pd.read_html(io.StringIO(response.text))
            target_table = next((t for t in tables if any(k in [c.lower() for c in t.columns] for k in table_keywords)), None)
            if target_table is not None:
                col_name = next(c for c in target_table.columns if c.lower() in table_keywords)
                tickers = target_table[col_name].tolist()
        except Exception as e:
            print(f" -> Error scraping {url}: {e}")
        return [str(t).replace('.', '-') for t in tickers if isinstance(t, str)]

    def get_all_nasdaq_tickers():
        try:
            url = "ftp://ftp.nasdaqtrader.com/symboldirectory/nasdaqtraded.txt"
            df = pd.read_csv(url, sep='|')
            df = df[(df['Test Issue'] == 'N') & (df['NASDAQ Symbol'].notnull())]
            return [str(t).replace('.', '-') for t in df['Symbol'].dropna().tolist()]
        except Exception:
            return get_wiki_tickers("https://en.wikipedia.org/wiki/Nasdaq-100", ['ticker', 'symbol'])

    def fetch_market_caps(tickers):
        caps = {}
        for i, ticker in enumerate(tickers):
            try:
                mc = yf.Ticker(ticker).fast_info.get('market_cap')
                if mc: caps[ticker] = mc
            except: continue
        return caps

    def fetch_latest_data(symbol: str, save_dir: str = "/tmp/daily_test"):
        # We use /tmp because it is the standard temporary folder in Linux/Modal
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        ticker = yf.Ticker(symbol)

        df = ticker.history(period="150d")
        if df.empty:
            raise ValueError(f"Could not find data for symbol: {symbol}")
        df = df.reset_index()

        # TEMPORAL FIREWALL: strip timezone from date index and cut today's partial row
        df["Date"] = pd.to_datetime(df["Date"]).dt.tz_localize(None).dt.date
        server_today = datetime.utcnow().date()
        df = df[df["Date"] < server_today]
        if df.empty:
            raise ValueError(f"No historical data for {symbol} after firewall cut.")

        file_path = os.path.join(save_dir, f"{symbol}_latest.csv")
        df.to_csv(file_path, index=False)
        return file_path

    # ==========================================
    # CLOUD-NATIVE LOGIC UPDATES
    # ==========================================
    
    # Connect to Neon to establish state
    print("Connecting to Neon Database...")
    conn = psycopg2.connect(os.environ["DATABASE_URL"])
    cursor = conn.cursor()

    ###
    def get_target_symbols_from_db():
        print("--- 1. Reading Golden Symbol List from Volume ---")
        # Read the 400 symbols from your uploaded file
        target_symbols_path = "/data/my_model_dir/symbols.txt"
        
        with open(target_symbols_path, "r") as f:
            pooled_400 = set(line.strip().upper() for line in f if line.strip())

        print(f"Loaded {len(pooled_400)} target symbols from file.")

        print("--- 2. Checking Neon Database for existing predictions ---")
        # Ask Neon which ones we already processed today/historically
        cursor.execute("SELECT DISTINCT entity_id FROM predictions")
        downloaded_symbols = {row[0] for row in cursor.fetchall()}

        # Split them into existing and new
        existing_list = list(downloaded_symbols.intersection(pooled_400))
        remaining_pool = pooled_400 - downloaded_symbols
        new_list = list(remaining_pool)

        return existing_list, new_list

    # ==========================================
    # MAIN EXECUTION BLOCK
    # ==========================================
    run_id = str(uuid.uuid4())
    today = datetime.now().date()
    model_version = "v1.0"
    
    existing_symbols, new_symbols = get_target_symbols_from_db()
    print(f"\nTargeting {len(existing_symbols)} existing symbols and {len(new_symbols)} new symbols.")

    # --- NEW: Close the DB connection so it doesn't time out while doing math ---
    cursor.close()
    conn.close()
    print("Closed initial DB connection. Starting ML heavy lifting...")

    print("\n⚙️ Loading neural network model into memory...")

    print("\n⚙️ Loading neural network model into memory...")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    # --- CLOUD UPDATE: Point to the Modal Volume Path ---
    checkpoint_dir = "/data/my_model_dir" 
    
    model, tokenizer, normalizer, cov_cols = load_checkpoint_with_config(checkpoint_dir, device=device)
    infer = CGCXInference(model, normalizer, tokenizer, device=device)
    print("✅ Model loaded successfully!")

    all_predictions = []

    # Combine existing and new into one loop for cleaner database insertion
    for sym in existing_symbols + new_symbols:
        try:
            path = fetch_latest_data(sym, save_dir="/tmp/daily_test")
            price_pred = predict_tomorrow(path, infer, tokenizer, cov_cols)
            if price_pred:
                # Format exactly to match the Neon 'predictions' SQL table
                all_predictions.append((run_id, today, sym, price_pred, model_version))
                print(f"{sym}: ${price_pred:.2f}")
        except Exception as e:
            print(f"Skipping {sym} - Error: {e}")

    ##
    # --- CLOUD UPDATE: SAVE RESULTS TO DATABASE ---
    print(f"\n💾 Saving {len(all_predictions)} predictions to Neon Database...")
    if all_predictions:
        print("Waking up Neon Database for saving...")
        conn = psycopg2.connect(os.environ["DATABASE_URL"])
        cursor = conn.cursor()

        insert_query = """
            INSERT INTO predictions (run_id, date, entity_id, predicted_value, model_version) 
            VALUES (%s, %s, %s, %s, %s)
            ON CONFLICT (date, entity_id, model_version) DO NOTHING;
        """
        cursor.executemany(insert_query, all_predictions)

        # --- THE FIX: Fetch a clean DataFrame of Actuals ---
        print("Fetching actual closing prices for evaluation...")
        all_syms = existing_symbols + new_symbols
        try:
            # Download the last 5 days of closing prices for all symbols
            raw_actuals = yf.download(all_syms, period="10d")["Close"]
            actuals_df = pd.DataFrame(raw_actuals)
            actuals_df.index = pd.to_datetime(actuals_df.index).tz_localize(None).date
            server_today = datetime.utcnow().date()
            actuals_df = actuals_df[actuals_df.index < server_today]
            actuals_df = actuals_df.dropna(how='all')
            
            # Run evaluations using the correct 'today' and 'actuals_df'
            run_evaluations(cursor, current_date=today, actuals_df=actuals_df)
        except Exception as e:
            print(f"⚠️ Warning: Evaluations failed, but predictions are safe. Error: {e}")
        
        # Finally, commit everything
        conn.commit()
        print("✅ Database commit successful. Pipeline complete.")
        
        cursor.close()
        conn.close()
        print("✅ Database save complete!")
    else:
        print("⚠️ No predictions generated today.")

@app.local_entrypoint()
def main():
    run_daily_inference.remote()