data_utils.py

import os
os.environ.setdefault("MPLBACKEND", "Agg")
from datetime import datetime, timedelta
from typing import Optional

import pandas as pd
import yfinance as yf
from fmp_python.fmp import FMP
from tqdm import tqdm


CACHE_DIR = "data_cache"
os.makedirs(CACHE_DIR, exist_ok=True)


def _get_cache_filename_yf(tckr_symbl, interval, start_date, end_date, adjust_prices):
    start_str = start_date.replace("-", "_")
    end_str = end_date.replace("-", "_")
    adjust_str = "adj" if adjust_prices else "raw"
    return os.path.join(CACHE_DIR, f"{tckr_symbl}_{interval}_{start_str}_to_{end_str}_{adjust_str}.csv")


def _get_cache_filename_fmp(tckr_symbl: str, interval: str) -> str:
    return os.path.join(CACHE_DIR, f"{tckr_symbl}_{interval}_fmp.csv")


def _is_valid_cache(df: pd.DataFrame) -> bool:
    if df.empty or len(df) < 2:
        return False
    if not isinstance(df.index, pd.DatetimeIndex):
        return False
    required_upper = ["Open", "High", "Low", "Close", "Volume"]
    required_lower = ["open", "high", "low", "close", "volume"]
    has_upper = all(col in df.columns for col in required_upper)
    has_lower = all(col in df.columns for col in required_lower)
    if not (has_upper or has_lower):
        return False
    return True


def _load_cached_data(cache_file: str) -> Optional[pd.DataFrame]:
    if os.path.exists(cache_file):
        try:
            df = pd.read_csv(cache_file, index_col=0, parse_dates=True, header=0)
            if isinstance(df.columns, pd.MultiIndex):
                df.columns = df.columns.get_level_values(0)
            if not isinstance(df.index, pd.DatetimeIndex):
                try:
                    df.index = pd.to_datetime(df.index)
                except Exception as e:
                    print(f"Could not parse dates in cache: {e}, will re-download")
                    return None
            expected_upper = ["Open", "High", "Low", "Close", "Volume"]
            expected_lower = ["open", "high", "low", "close", "volume"]
            actual_cols = list(df.columns)
            has_upper = any(col in actual_cols for col in expected_upper)
            has_lower = any(col in actual_cols for col in expected_lower)
            if not (has_upper or has_lower):
                print(f"Unexpected columns in cache: {actual_cols}, will re-download")
                return None
            print(f"Loaded cached data: {len(df)} rows, columns: {list(df.columns)} from {cache_file}")
            return df
        except Exception as e:
            print(f"Error loading cache: {e}, will re-download")
            import traceback
            traceback.print_exc()
            return None
    return None


def _save_cached_data(df: pd.DataFrame, cache_file: str):
    try:
        df.to_csv(cache_file)
        print(f"Cached data saved: {cache_file}")
    except Exception as e:
        print(f"Error saving cache: {e}")


# FMP interval -> max days per API chunk
FMP_INTERVAL_DAYS = {"1min": 2, "5min": 7, "15min": 38, "1hour": 70, "4hour": 160}


def _fetch_fmp_range(fmp, tckr_symbl, interval, start_dt, end_dt, progress=None):
    """Download FMP data for a date range. Returns DataFrame with 'date' index."""
    chunk_span = timedelta(days=FMP_INTERVAL_DAYS[interval])
    frames = []

    # Build chunk list
    chunks = []
    temp = start_dt
    while temp <= end_dt:
        chunks.append(temp)
        temp = min(temp + chunk_span, end_dt) + timedelta(days=1)

    # Download with progress
    if progress:
        chunk_iter = progress.tqdm(chunks, desc=f"FMP {interval}")
    else:
        chunk_iter = tqdm(chunks, desc=f"FMP {interval}")

    for chunk_start in chunk_iter:
        chunk_end = min(chunk_start + chunk_span, end_dt)
        try:
            chunk = fmp.get_historical_chart(
                interval, tckr_symbl,
                _from=chunk_start.strftime("%Y-%m-%d"),
                _to=chunk_end.strftime("%Y-%m-%d")
            )
        except Exception as e:
            raise ValueError(f"FMP download failed ({chunk_start.date()} to {chunk_end.date()}): {e}")

        if chunk is not None and not chunk.empty:
            chunk["date"] = pd.to_datetime(chunk["date"], errors="coerce")
            chunk = chunk.dropna(subset=["date"])
            if not chunk.empty:
                frames.append(chunk)

    if not frames:
        return pd.DataFrame()

    df = pd.concat(frames, ignore_index=True)
    df["date"] = pd.to_datetime(df["date"], errors="coerce")
    df = df.dropna(subset=["date"])

    # Timezone handling
    if not df.empty and df["date"].dt.tz is None:
        df["date"] = df["date"].dt.tz_localize("America/New_York", nonexistent="shift_forward", ambiguous="NaT")
    df["date"] = df["date"].dt.tz_convert("UTC").dt.tz_localize(None)

    return df.sort_values("date").set_index("date")


def _download_data_fmp(tckr_symbl, interval, date, progress=None, replay=False, compression=None):
    if interval not in FMP_INTERVAL_DAYS:
        raise ValueError(f"Unsupported FMP interval '{interval}'")

    end_dt = datetime.strptime(date["end"], "%Y-%m-%d")
    start_dt = datetime.strptime(date["start"], "%Y-%m-%d")

    # Clamp to 15-year limit
    max_lookback = end_dt - timedelta(days=15 * 365)
    if start_dt < max_lookback:
        print(f"Warning: start date clamped to {max_lookback.date()} (15y limit).")
        start_dt = max_lookback

    fmp = FMP(output_format="pandas", write_to_file=False)
    cache_file = _get_cache_filename_fmp(tckr_symbl, interval)
    cached_df = _load_cached_data(cache_file)

    if cached_df is not None and _is_valid_cache(cached_df):
        cache_start = cached_df.index.min().date()
        cache_end = cached_df.index.max().date()
        frames = []

        # Download past data if needed
        if start_dt.date() < cache_start:
            past_end = datetime.combine(cache_start - timedelta(days=1), datetime.min.time())
            print(f"Fetching past data: {start_dt.date()} to {past_end.date()}")
            chunk_past = _fetch_fmp_range(fmp, tckr_symbl, interval, start_dt, past_end, progress)
            if not chunk_past.empty:
                frames.append(chunk_past)

        frames.append(cached_df)

        # Download current data if needed
        if end_dt.date() > cache_end:
            current_start = datetime.combine(cache_end + timedelta(days=1), datetime.min.time())
            print(f"Fetching current data: {current_start.date()} to {end_dt.date()}")
            chunk_current = _fetch_fmp_range(fmp, tckr_symbl, interval, current_start, end_dt, progress)
            if not chunk_current.empty:
                frames.append(chunk_current)

        # Merge, dedupe, sort, save
        if len(frames) > 1:
            df = pd.concat(frames).sort_index()
            df = df[~df.index.duplicated(keep='last')]
            _save_cached_data(df, cache_file)
        else:
            df = cached_df
            print(f"Using cached data: {len(df)} rows (no download needed)")
    else:
        # No cache - download full range
        print(f"No cache, downloading: {start_dt.date()} to {end_dt.date()}")
        df = _fetch_fmp_range(fmp, tckr_symbl, interval, start_dt, end_dt, progress)
        if not df.empty:
            _save_cached_data(df, cache_file)

    if df.empty:
        return df

    # Filter to user's requested range
    df = df[(df.index >= pd.Timestamp(start_dt)) & (df.index <= pd.Timestamp(end_dt))]
    return df


def _download_data_yf(tckr_symbl, interval, date, adjust_prices, auto_period=True, period="60d"):
    try:
        print("Interval: ", interval)
        start_dt = datetime.strptime(date["start"], "%Y-%m-%d")
        end_dt = datetime.strptime(date["end"], "%Y-%m-%d")

        cache_file = _get_cache_filename_yf(tckr_symbl, interval, date["start"], date["end"], adjust_prices)
        cached_df = _load_cached_data(cache_file)
        if cached_df is not None and _is_valid_cache(cached_df):
            df = cached_df
            print(f"Using cached data: {len(df)} rows (no download needed)")
        else:
            print("No valid cache, downloading...")
            if interval in ["1m", "2m", "5m", "15m", "30m", "60m", "1h"] and auto_period:
                if interval in ["1m"]:
                    max_days = 7
                elif interval in ["2m", "5m", "15m", "30m"]:
                    max_days = 60
                else:
                    max_days = 730
                desired_days = max(1, (end_dt - start_dt).days or 1)
                clamped_days = min(desired_days, max_days)
                period = f"{clamped_days}d"
                df = yf.download(tckr_symbl, period=period, interval=interval, auto_adjust=adjust_prices)
                print(f"Downloaded {interval} data for {period}")
            else:
                df = yf.download(tckr_symbl, start=date["start"], end=date["end"], interval=interval, auto_adjust=adjust_prices)
                print(f"Downloaded data from {date['start']} to {date['end']} with {interval} interval")

            if isinstance(df.columns, pd.MultiIndex):
                df.columns = df.columns.get_level_values(0)
            _save_cached_data(df, cache_file)
            if isinstance(df.columns, pd.MultiIndex):
                df.columns = df.columns.get_level_values(0)
            if df.empty:
                raise ValueError("No data available for the specified parameters!")
            if df.index.tz is not None:
                df.index = df.index.tz_localize(None)
            print(f"Data points: {len(df)}")
        return df
    except Exception as e:
        raise ValueError(f"Error downloading data: {e}")


def get_data(
    data_source: str,
    tckr_symbl: str,
    interval: str,
    date: dict,
    adjust_prices: bool = True,
    auto_period: bool = True,
    period: str = "60d",
    upload_data: bool = False,
    upload_data_path: str = None,
    progress=None,
):
    if upload_data:
        if not upload_data_path:
            raise ValueError("upload_data_path is required when upload_data is True.")
        df = pd.read_csv(upload_data_path)
        if df.shape[1] >= 2:
            dt = pd.to_datetime(df.iloc[:, 0].astype(str) + " " + df.iloc[:, 1].astype(str), errors="coerce")
            df = df.drop(columns=df.columns[:2])
        else:
            dt = pd.to_datetime(df.iloc[:, 0], errors="coerce")
            df = df.drop(columns=df.columns[:1])
        df.insert(0, "datetime", dt)
        df = df.dropna(subset=["datetime"]).set_index("datetime")
        expected_cols = ["open", "high", "low", "close", "volume"]
        if len(df.columns) >= 5:
            df.columns = list(expected_cols) + list(df.columns[len(expected_cols) :])
        df.columns = [c.capitalize() for c in df.columns]
        df.index = df.index.tz_localize(None)
        return df

    source = (data_source or "").lower()
    loaders = {
        "yahoofinance": lambda: _download_data_yf(tckr_symbl, interval, date, adjust_prices, auto_period, period),
        "yf": lambda: _download_data_yf(tckr_symbl, interval, date, adjust_prices, auto_period, period),
        "yahoo": lambda: _download_data_yf(tckr_symbl, interval, date, adjust_prices, auto_period, period),
        "fmp": lambda: _download_data_fmp(tckr_symbl, interval, date, progress),
    }
    if source not in loaders:
        raise ValueError(f"Invalid data source: {data_source}")
    return loaders[source]()