from typing import Tuple, Optional
import numpy as np
import pandas as pd
from pathlib import Path


def load_timeseries_data(dataset_name: str, data_path: str = "dataset/") -> Tuple[np.ndarray, Optional[pd.DataFrame]]:
    """Load universal time series data (electricity, traffic, etc.)"""

    # First check if dataset_name is already a valid file path
    if Path(dataset_name).exists():
        data_file = Path(dataset_name)
    else:
        # Try different path patterns
        possible_paths = [
            Path(data_path) / f"{dataset_name}.csv",  # dataset/electricity.csv
            Path(data_path) / dataset_name / f"{dataset_name}.csv",  # dataset/electricity/electricity.csv
        ]

        data_file = None
        for path in possible_paths:
            if path.exists():
                data_file = path
                break

        if data_file is None:
            # Try alternative extensions
            for ext in [".txt", ".h5", ".hdf5"]:
                for base_path in [Path(data_path) / f"{dataset_name}{ext}",
                                Path(data_path) / dataset_name / f"{dataset_name}{ext}"]:
                    if base_path.exists():
                        data_file = base_path
                        break
                if data_file:
                    break
            else:
                raise FileNotFoundError(f"Dataset file not found: {dataset_name}. Tried paths: {possible_paths}")

    # Load file based on extension
    ext = str(data_file).lower().rsplit(".", 1)[-1]
    if ext == "txt":
        raw_data = np.loadtxt(data_file, delimiter=",")
        dates_df = None
    elif ext == "csv":
        df = pd.read_csv(data_file, parse_dates=[0])
        dates_df = pd.DataFrame(df.iloc[:, [0]].rename(columns={df.columns[0]: "date"}))
        raw_data = df.iloc[:, 1:].values
    elif ext in ("h5", "hdf5"):
        d = pd.read_hdf(data_file).reset_index()
        dates_df = pd.DataFrame(d["index"].rename("date"))
        raw_data = d.iloc[:, 1:].values
    else:
        raise ValueError(f"Unsupported file extension: {ext}")

    return raw_data.astype(np.float32), dates_df