train.py

# train.py
import os, json, glob, math, argparse
from pathlib import Path
from typing import Dict, Any, List, Tuple

import numpy as np
import pandas as pd

# Try LightGBM; fallback to sklearn if not available
USING_LGBM = True
try:
    from lightgbm import LGBMRegressor
except Exception:
    USING_LGBM = False
    from sklearn.experimental import enable_hist_gradient_boosting  # noqa: F401
    from sklearn.ensemble import HistGradientBoostingRegressor as LGBMRegressor  # type: ignore

from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt


def stem_to_colname(path: Path) -> str:
    # normalize filename into a safe column name
    return path.stem.strip().lower().replace("-", "_").replace(" ", "_")


def parse_series_from_obj(obj: Any, preferred_name: str) -> pd.Series | None:
    """
    Try multiple known JSON shapes and return a Series indexed by datetime.
    All timestamps assumed to be milliseconds since epoch unless clearly ISO.
    """
    # 1) Blockchain.com-like: {"status":"ok", ... "values":[{"x":1406073600, "y":589.75}, ...]}
    if isinstance(obj, dict) and "values" in obj and isinstance(obj["values"], list):
        vals = obj["values"]
        xs, ys = [], []
        for row in vals:
            if not isinstance(row, dict): 
                continue
            x, y = row.get("x"), row.get("y")
            if x is None or y is None:
                continue
            # many sources are in seconds; yours are mostly in ms — try ms first, then sec
            ts = pd.to_datetime(x, unit="ms", errors="coerce", utc=True)
            if pd.isna(ts):
                ts = pd.to_datetime(x, unit="s", errors="coerce", utc=True)
            if pd.isna(ts):
                ts = pd.to_datetime(x, errors="coerce", utc=True)
            if pd.isna(ts):
                continue
            try:
                y = float(y)
            except Exception:
                continue
            xs.append(ts)
            ys.append(y)
        if xs:
            return pd.Series(ys, index=pd.DatetimeIndex(xs, name="timestamp"), name=preferred_name).sort_index()

    # 2) Custom shape seen in your files: {"metric1":"foo","metric2":"bar","foo":[{x,y},...]}
    if isinstance(obj, dict) and "metric1" in obj:
        metric1 = obj.get("metric1")
        if isinstance(metric1, str) and isinstance(obj.get(metric1), list):
            xs, ys = [], []
            for row in obj[metric1]:
                if not isinstance(row, dict):
                    continue
                x, y = row.get("x"), row.get("y")
                if x is None or y is None:
                    continue
                ts = pd.to_datetime(x, unit="ms", errors="coerce", utc=True)
                if pd.isna(ts):
                    ts = pd.to_datetime(x, unit="s", errors="coerce", utc=True)
                if pd.isna(ts):
                    ts = pd.to_datetime(x, errors="coerce", utc=True)
                if pd.isna(ts):
                    continue
                try:
                    y = float(y)
                except Exception:
                    continue
                xs.append(ts); ys.append(y)
            if xs:
                return pd.Series(ys, index=pd.DatetimeIndex(xs, name="timestamp"), name=preferred_name).sort_index()

    # 3) Plain list of {x,y}
    if isinstance(obj, list) and obj and isinstance(obj[0], dict) and "x" in obj[0] and "y" in obj[0]:
        xs, ys = [], []
        for row in obj:
            x, y = row.get("x"), row.get("y")
            ts = pd.to_datetime(x, unit="ms", errors="coerce", utc=True)
            if pd.isna(ts):
                ts = pd.to_datetime(x, unit="s", errors="coerce", utc=True)
            if pd.isna(ts):
                ts = pd.to_datetime(x, errors="coerce", utc=True)
            if pd.isna(ts):
                continue
            try:
                y = float(y)
            except Exception:
                continue
            xs.append(ts); ys.append(y)
        if xs:
            return pd.Series(ys, index=pd.DatetimeIndex(xs, name="timestamp"), name=preferred_name).sort_index()

    # Not a recognized time series
    return None


def load_all_series(data_dir: Path) -> pd.DataFrame:
    series_list: List[pd.Series] = []
    for p in sorted(data_dir.glob("*.json")):
        colname = stem_to_colname(p)
        try:
            with p.open("r", encoding="utf-8") as f:
                obj = json.load(f)
            s = parse_series_from_obj(obj, preferred_name=colname)
            if s is None or s.empty:
                print(f"[SKIP] {p.name}: not a time-series or could not parse; skipping.")
                continue
            # keep first occurrence of a timestamp (drop dups) and sort
            s = s[~s.index.duplicated(keep="first")].sort_index()
            series_list.append(s)
            print(f"[OK]   {p.name:35s} -> column '{colname}' with {s.size} rows")
        except Exception as e:
            print(f"[WARN] {p.name}: {e}; skipping.")
    if not series_list:
        raise RuntimeError("No valid time series were loaded from the JSON files.")
    df = pd.concat(series_list, axis=1).sort_index()
    # daily frequency & forward fill to align features
    # daily index + fill both forward and backward so no leading holes remain
    df = df.asfreq("D").ffill().bfill()
    return df


def build_dataset(df: pd.DataFrame, target_hint: str, horizon: int = 1):
    target_col = target_hint
    if target_col not in df.columns:
        # try to auto-find a market price column
        candidates = [c for c in df.columns if "market" in c and "price" in c]
        if candidates:
            target_col = candidates[0]
            print(f"[INFO] target '{target_hint}' not found; using '{target_col}'")
        else:
            raise KeyError(f"Target '{target_hint}' not found in columns: {list(df.columns)[:10]} ...")

    # Basic lag features for all columns (including target)
    feat = df.copy()
    for lag in [1, 2, 7, 14, 30]:
        feat = feat.join(df.shift(lag).add_suffix(f"_lag{lag}"))

    # Rolling stats for the target
    for win in [7, 14, 30, 60, 90]:
        feat[f"{target_col}_rollmean_{win}"] = df[target_col].rolling(win).mean()
        feat[f"{target_col}_rollstd_{win}"] = df[target_col].rolling(win).std()

        # Predict target at t+horizon
    y = df[target_col].shift(-horizon).rename(f"{target_col}_t_plus_{horizon}")

    # Build features
    feat = df.copy()
    for lag in [1, 2, 7, 14, 30]:
        feat = feat.join(df.shift(lag).add_suffix(f"_lag{lag}"))
    for win in [7, 14, 30, 60, 90]:
        feat[f"{target_col}_rollmean_{win}"] = df[target_col].rolling(win).mean()
        feat[f"{target_col}_rollstd_{win}"]  = df[target_col].rolling(win).std()

    # Combine features + target and fill remaining holes
    combo = feat.join(y)

    # FIRST pass: forward fill; SECOND pass: backfill — removes isolated gaps without leaking far
    combo = combo.ffill().bfill()

    # Now split again
    y = combo[y.name]
    X = combo.drop(columns=[y.name])

    # Force numeric; drop all-NaN and constant columns
    X = X.apply(pd.to_numeric, errors="coerce")
    X = X.dropna(axis=1, how="all")
    nunique = X.nunique(dropna=False)
    constant_cols = nunique[nunique <= 1].index.tolist()
    if constant_cols:
        X = X.drop(columns=constant_cols)

    # Final clean (should be no NaNs left, but just in case)
    keep = X.join(y).dropna()
    X = keep.drop(columns=[y.name])
    y = keep[y.name]

    return X, y, target_col


def evaluate_and_plot(y_true: pd.Series, y_pred: np.ndarray, out_png: Path):
    mae = mean_absolute_error(y_true, y_pred)
    rmse = math.sqrt(mean_squared_error(y_true, y_pred))
    eps = 1e-8
    mape = float(np.mean(np.abs((y_true.values - y_pred) / np.maximum(np.abs(y_true.values), eps))) * 100.0)
    r2 = r2_score(y_true, y_pred)

    plt.figure(figsize=(12, 5), dpi=150)
    plt.plot(y_true.index, y_true.values, label="Actual")
    plt.plot(y_true.index, y_pred, label="Predicted")
    plt.title("Bitcoin Market Price — Actual vs Predicted")
    plt.xlabel("Date")
    plt.ylabel("Price (USD)")
    plt.legend()
    plt.tight_layout()
    plt.savefig(out_png)
    plt.close()

    print("\nMetrics:")
    print(f"  MAE    : {mae:,.4f}")
    print(f"  RMSE   : {rmse:,.4f}")
    print(f"  MAPE % : {mape:,.2f}")
    print(f"  R²     : {r2:,.4f}")


def chronological_split(X: pd.DataFrame, y: pd.Series, test_size: float = 0.2):
    split_idx = int((1.0 - test_size) * len(X))
    return X.iloc[:split_idx], X.iloc[split_idx:], y.iloc[:split_idx], y.iloc[split_idx:]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data_dir", type=str, default="./data")
    ap.add_argument("--target", type=str, default="market_price", help="Normalized target name, e.g., market_price")
    ap.add_argument("--horizon", type=int, default=1)
    ap.add_argument("--test_size", type=float, default=0.2)
    ap.add_argument("--out_png", type=str, default="prediction.png")
    args = ap.parse_args()

    data_dir = Path(args.data_dir)
    df = load_all_series(data_dir)

    X, y, target_col = build_dataset(df, target_hint=args.target, horizon=args.horizon)

    # Sanity checks
    print(f"[INFO] Feature matrix shape (pre-split): {X.shape}")
    print(f"[INFO] Target vector length: {len(y)}")

    if X.empty or X.shape[1] == 0:
        raise ValueError("No usable numeric features after coercion. "
                         "Check JSON schemas and ensure series are numeric.")

    # Chronological split
    X_train, X_test, y_train, y_test = chronological_split(X, y, test_size=args.test_size)

    print(f"[INFO] Train shape: {X_train.shape} | Test shape: {X_test.shape}")

    if X_train.empty or X_train.shape[1] == 0:
        raise ValueError("Training features are empty. "
                         "This usually means all columns were dropped as non-numeric or constant.")

    if len(X_train) == 0 or len(X_test) == 0:
        raise ValueError("Split produced empty train/test sets. Reduce test_size or check data length.")

    X_train, X_test, y_train, y_test = chronological_split(X, y, test_size=args.test_size)

    if USING_LGBM:
        model = LGBMRegressor(
            n_estimators=1000,
            learning_rate=0.03,
            subsample=0.8,
            colsample_bytree=0.8,
            random_state=42,
        )
    else:
        model = LGBMRegressor(
            max_iter=1000,
            learning_rate=0.05,
            random_state=42
        )

    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)

    evaluate_and_plot(y_test, y_pred, out_png=Path(args.out_png))

    print("\nInfo:")
    print(f"  Using LightGBM: {USING_LGBM}")
    print(f"  Data rows      : {len(df):,} (after daily alignment/ffill)")
    print(f"  Train/Test     : {len(X_train):,} / {len(X_test):,}")
    print(f"  Target column  : {target_col}")


if __name__ == "__main__":
    main()