clean_data.py

"""Clean raw OHLCV files into one long-form price table."""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pandas as pd

from data_fetch import ALL_LOCAL_SYMBOLS

CANONICAL_COLUMNS = [
    "date",
    "ticker",
    "open",
    "high",
    "low",
    "close",
    "adj_close",
    "volume",
    "source",
    "vendor_symbol",
    "is_adjusted",
]


def run_clean_data(root: str | Path = ".") -> pd.DataFrame:
    root_path = Path(root)
    raw_dir = root_path / "data" / "raw"
    clean_dir = root_path / "data" / "clean"
    clean_dir.mkdir(parents=True, exist_ok=True)

    frames: list[pd.DataFrame] = []
    for symbol in ALL_LOCAL_SYMBOLS:
        frame = _load_preferred_symbol(raw_dir, symbol)
        if frame is not None and not frame.empty:
            frames.append(frame)

    if not frames:
        empty = pd.DataFrame(columns=CANONICAL_COLUMNS)
        empty.to_csv(clean_dir / "clean_prices.csv", index=False)
        _write_quality_report(empty, clean_dir)
        return empty

    prices = pd.concat(frames, ignore_index=True)
    prices = prices.dropna(subset=["date", "ticker", "close"])
    prices = prices.sort_values(["ticker", "date"], ignore_index=True)
    prices = prices.drop_duplicates(["ticker", "date"], keep="last")
    prices["date"] = pd.to_datetime(prices["date"]).dt.strftime("%Y-%m-%d")
    prices.to_csv(clean_dir / "clean_prices.csv", index=False)
    _write_quality_report(prices, clean_dir)
    return prices


def _load_preferred_symbol(raw_dir: Path, symbol: str) -> pd.DataFrame | None:
    candidates = [
        raw_dir / f"vnstock_{symbol}.csv",
        raw_dir / f"yfinance_{symbol}.csv",
    ]
    for path in candidates:
        if path.exists():
            frame = pd.read_csv(path)
            source = path.stem.split("_", 1)[0]
            normalized = _normalize_raw_frame(frame, symbol=symbol, source=source)
            if not normalized.empty:
                return normalized
    return None


def _normalize_raw_frame(frame: pd.DataFrame, *, symbol: str, source: str) -> pd.DataFrame:
    data = frame.copy()
    data.columns = [_normalize_column_name(col) for col in data.columns]
    rename = {
        "time": "date",
        "datetime": "date",
        "date": "date",
        "open": "open",
        "high": "high",
        "low": "low",
        "close": "close",
        "adj close": "adj_close",
        "adj_close": "adj_close",
        "adjusted close": "adj_close",
        "adjusted_close": "adj_close",
        "volume": "volume",
        "vol": "volume",
        "ticker": "ticker",
        "vendor_symbol": "vendor_symbol",
    }
    data = data.rename(columns={col: rename[col] for col in data.columns if col in rename})

    if "date" not in data.columns:
        return pd.DataFrame(columns=CANONICAL_COLUMNS)
    if "ticker" not in data.columns:
        data["ticker"] = symbol
    if "vendor_symbol" not in data.columns:
        data["vendor_symbol"] = symbol

    for column in ("open", "high", "low", "close", "adj_close", "volume"):
        if column not in data.columns:
            data[column] = np.nan
    if data["adj_close"].isna().all():
        data["adj_close"] = data["close"]
        is_adjusted = False
    else:
        is_adjusted = True

    data["date"] = pd.to_datetime(data["date"], errors="coerce")
    for column in ("open", "high", "low", "close", "adj_close", "volume"):
        data[column] = pd.to_numeric(data[column], errors="coerce")
    data["source"] = source
    data["is_adjusted"] = is_adjusted
    return data.loc[:, CANONICAL_COLUMNS]


def _normalize_column_name(column: object) -> str:
    return str(column).strip().replace("_", " ").lower()


def _write_quality_report(prices: pd.DataFrame, clean_dir: Path) -> None:
    rows: list[dict[str, object]] = []
    if prices.empty:
        quality = pd.DataFrame(
            columns=[
                "ticker",
                "rows",
                "start_date",
                "end_date",
                "missing_values",
                "zero_volume_days",
                "stale_close_days",
                "last_close",
                "last_volume",
            ]
        )
        quality.to_csv(clean_dir / "data_quality.csv", index=False)
        return

    work = prices.copy()
    work["date"] = pd.to_datetime(work["date"])
    numeric_cols = ["open", "high", "low", "close", "adj_close", "volume"]
    for ticker, group in work.groupby("ticker"):
        ordered = group.sort_values("date")
        rows.append(
            {
                "ticker": ticker,
                "rows": len(ordered),
                "start_date": ordered["date"].min().strftime("%Y-%m-%d"),
                "end_date": ordered["date"].max().strftime("%Y-%m-%d"),
                "missing_values": int(ordered[numeric_cols].isna().sum().sum()),
                "zero_volume_days": int((ordered["volume"].fillna(0) == 0).sum()),
                "stale_close_days": int((ordered["close"].diff().fillna(1) == 0).sum()),
                "last_close": float(ordered["close"].iloc[-1]),
                "last_volume": float(ordered["volume"].iloc[-1])
                if pd.notna(ordered["volume"].iloc[-1])
                else np.nan,
            }
        )
    pd.DataFrame(rows).to_csv(clean_dir / "data_quality.csv", index=False)


if __name__ == "__main__":
    run_clean_data()