import pandas as pd
import numpy as np
import torch
from sklearn.preprocessing import StandardScaler

try:
    import yfinance as yf
except ImportError:
    raise ImportError("yfinance must be installed to fetch financial data.")

def load_data(data_src="yahoo", ticker="AAPL", file_upload=None, start="2020-01-01", end="2023-01-01", horizon=1):
    main_df = None
    future_df = None
    
    if data_src == "yahoo":
        try:
            info = yf.Ticker(ticker).info
            if not info:
                raise ValueError(f"Ticker '{ticker}' not found.")
            # Fetch data up to end_date + horizon for future actuals
            extended_end = (pd.to_datetime(end) + pd.Timedelta(days=horizon)).strftime('%Y-%m-%d')
            df = yf.download(ticker, start=start, end=extended_end, progress=False)
            if df.empty:
                raise ValueError(f"No data found for ticker '{ticker}' in the specified date range. Please check the symbol and dates.")
            df = df[['Close']].dropna().rename(columns={'Close': 'value'})
            df.reset_index(inplace=True)
            # Split into main_df (up to end_date) and future_df (beyond end_date)
            main_df = df[df['Date'] <= pd.to_datetime(end)].copy()
            future_df = df[(df['Date'] > pd.to_datetime(end)) & (df['Date'] <= pd.to_datetime(extended_end))].copy()
        except Exception as e:
            raise ValueError(f"Error fetching data for ticker '{ticker}': {e}")
    elif data_src == "csv":
        if file_upload is None:
            raise ValueError("CSV file upload required but not provided.")
        try:
            df = pd.read_csv(file_upload)
        except Exception as e:
            raise ValueError(f"Failed to read uploaded CSV file: {e}")
        if 'value' not in df.columns:
            if 'Close' in df.columns:
                df = df[['Close']].rename(columns={'Close': 'value'})
            else:
                raise ValueError("CSV must contain a 'value' or 'Close' column.")
        df['Date'] = pd.to_datetime(df.get('Date', df.index))
        df = df[['Date', 'value']].dropna().reset_index(drop=True)
        # Split into main_df (up to end_date) and future_df (beyond end_date)
        main_df = df[df['Date'] <= pd.to_datetime(end)].copy()
        future_df = df[(df['Date'] > pd.to_datetime(end)) & (df['Date'] <= pd.to_datetime(end) + pd.Timedelta(days=horizon))].copy()
    else:
        raise ValueError("Invalid data source. 'csv' or 'yahoo' expected.")
    
    return main_df, future_df


def preprocess_data(df, column, window_size=30):
    scaler = StandardScaler()
    data = df[[column]].values.astype(float)
    scaled = scaler.fit_transform(data)

    X, y = [], []
    for i in range(len(scaled) - window_size):
        X.append(scaled[i:i + window_size])
        y.append(scaled[i + window_size])
    X = np.array(X)
    y = np.array(y)

    return X, y, scaler