import yfinance as yf
import pandas as pd
import numpy as np
from sklearn.metrics import accuracy_score


import numpy as np
import pandas as pd
import yfinance as yf

from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, mean_absolute_error

from catboost import CatBoostClassifier


TEST_SIZE = 0.2

SEQ_LENGTH = 180
SEQ_INTERVAL = 10

import datetime


def get_forecast_data(
    ticker, q, feature_cols, best_seq_len, selected_model, start="2017-01-01"
):
    df = yf.download(ticker, start=start, auto_adjust=True, progress=False)

    d_ = yf.download(
        ticker,
        start=datetime.datetime.now().strftime("%Y-%m-%d"),
        auto_adjust=True,
        progress=False,
    )

    df.update(d_)

    if isinstance(df.columns, pd.MultiIndex):
        df.columns = df.columns.get_level_values(0)

    df["ret"] = np.log(df["Close"] / df["Close"].shift(1))
    df["volatility"] = df["ret"].rolling(20).std()
    df["hl_ratio"] = np.log(df["High"] / df["Low"])
    df["oc_ratio"] = np.log(df["Close"] / df["Open"])
    df["volume_change"] = np.log(df["Volume"] / df["Volume"].shift(1))

    sign = np.sign(df["ret"])
    streak = []
    s = 0
    for x in sign:
        if pd.isna(x):
            streak.append(np.nan)
            continue
        if x > 0:
            s = s + 1 if s > 0 else 1
        elif x < 0:
            s = s - 1 if s < 0 else -1
        else:
            s = 0
        streak.append(s)
    df["streak"] = streak

    df["target_ret"] = df["ret"].shift(-1)

    # Re-apply quantile binning using previously fitted bin_edges

    quantile_features = [
        "ret",
        "volatility",
        "hl_ratio",
        "oc_ratio",
        "volume_change",
    ]

    bin_edges = {}

    for col in quantile_features:

        _, bins = pd.qcut(df[col], q=q, labels=False, retbins=True, duplicates="drop")

        bins[0] = -np.inf
        bins[-1] = np.inf

        bin_edges[col] = bins

        df[col + "_q"] = pd.cut(df[col], bins=bins, labels=False)  # .astype(int)

    for col in quantile_features:
        df[col + "_q"] = pd.cut(
            df[col], bins=bin_edges[col], labels=False, duplicates="drop"
        )

    _, target_bins = pd.qcut(
        df["target_ret"], q=q, labels=False, retbins=True, duplicates="drop"
    )

    target_bins[0] = -np.inf
    target_bins[-1] = np.inf

    df["target"] = pd.cut(df["target_ret"], bins=target_bins, labels=False)

    # Drop rows with NaN values introduced by feature engineering
    df_processed = df.dropna(subset=feature_cols + ["target"])

    # df_for_tomorrow_prediction = df.dropna()#.tail(best_seq_len_multiclass)

    # Create the sequence. The make_sequences function expects a target column, but for actual prediction, it won't be used.
    # We will manually extract X_tomorrow.
    X_tomorrow_raw, _ = make_sequences(
        df_processed,
        feature_cols,
        "target",  # Placeholder, not actually used for a single prediction point
        best_seq_len,
    )

    # The make_sequences returns sequences ending at the last day of the input df.
    # We need the *last* sequence from X_tomorrow_raw for the actual prediction.
    X_tomorrow = X_tomorrow_raw[-1].reshape(1, -1)
    pred_tomorrow = selected_model["model"].predict(X_tomorrow)[0]

    # First, generate predictions for all historical data using the best model
    X_full, y_full = make_sequences(df_processed, feature_cols, "target", best_seq_len)
    X_full_flat = X_full.reshape(X_full.shape[0], -1)

    pred_full = selected_model["model"].predict(X_full_flat)
    pred_full = pred_full.astype(int).ravel()

    # Create tmp_current dataframe using predicted bins and actual historical returns
    tmp_current = pd.DataFrame(
        {
            "pred": pred_full,
            "ret": df_processed["target_ret"].iloc[best_seq_len:].values,
        }
    )

    # Calculate rolling metrics on tmp_current
    TREND_WINDOW = 15
    for val in range(q):
        conditional_ret_current = tmp_current["ret"].where(tmp_current["pred"] == val)
        tmp_current[
            f"rolling_ret_{TREND_WINDOW}_mean_pred_{val}"
        ] = conditional_ret_current.rolling(window=TREND_WINDOW, min_periods=1).mean()

    # Calculate the trend signal for each day in tmp_current
    trend_signal_current = []
    for idx, row in tmp_current.fillna(0).iterrows():
        res = 0
        for n in range(q):
            res += row[f"rolling_ret_{TREND_WINDOW}_mean_pred_{n}"]
        trend_signal_current.append(np.sign(res))

    # The trend signal for tomorrow is the last calculated signal
    trend_signal_for_tomorrow = trend_signal_current[-1]

    TOP_BINS = 2

    signal_tomorrow = (
        (pred_tomorrow >= (q - TOP_BINS)) and (trend_signal_for_tomorrow >= 0)
    ).astype(int)

    return (
        pred_tomorrow,
        trend_signal_for_tomorrow,
        signal_tomorrow,
        df_processed.index[-1],
    )


def get_dataframe(ticker, q=3, start="2017-01-01", end="2026-02-01"):
    df = yf.download(ticker, start=start, end=None, auto_adjust=True)
    if isinstance(df.columns, pd.MultiIndex):
        df.columns = df.columns.get_level_values(0)

    # ============================================================
    # FEATURES
    # ============================================================

    df["ret"] = np.log(df["Close"] / df["Close"].shift(1))

    df["volatility"] = df["ret"].rolling(20).std()

    df["hl_ratio"] = np.log(df["High"] / df["Low"])

    df["oc_ratio"] = np.log(df["Close"] / df["Open"])

    df["volume_change"] = np.log(df["Volume"] / df["Volume"].shift(1))

    # ============================================================
    # STREAK
    # ============================================================

    sign = np.sign(df["ret"])

    streak = []
    s = 0

    for x in sign:

        if pd.isna(x):
            streak.append(np.nan)
            continue

        if x > 0:
            s = s + 1 if s > 0 else 1

        elif x < 0:
            s = s - 1 if s < 0 else -1

        else:
            s = 0

        streak.append(s)

    df["streak"] = streak

    # ============================================================
    # TARGET
    # ============================================================

    df["target_ret"] = df["ret"].shift(-1)

    # ============================================================
    # TRAIN/TEST SPLIT FIRST
    # (important to avoid leakage)
    # ============================================================

    split_idx = int(len(df) * (1 - TEST_SIZE))

    train_df = df.iloc[:split_idx].copy()
    test_df = df.iloc[split_idx:].copy()

    # ============================================================
    # QCUT FEATURES USING TRAIN ONLY
    # ============================================================

    quantile_features = [
        "ret",
        "volatility",
        "hl_ratio",
        "oc_ratio",
        "volume_change",
    ]

    bin_edges = {}

    for col in quantile_features:

        _, bins = pd.qcut(
            train_df[col], q=q, labels=False, retbins=True, duplicates="drop"
        )

        bins[0] = -np.inf
        bins[-1] = np.inf

        bin_edges[col] = bins

        train_df[col + "_q"] = pd.cut(
            train_df[col], bins=bins, labels=False
        )  # .astype(int)

        test_df[col + "_q"] = pd.cut(
            test_df[col], bins=bins, labels=False
        )  # .astype(int)

    # ============================================================
    # TARGET BINS
    # ============================================================

    _, target_bins = pd.qcut(
        train_df["target_ret"], q=q, labels=False, retbins=True, duplicates="drop"
    )  # .astype(int)

    target_bins[0] = -np.inf
    target_bins[-1] = np.inf

    train_df["target"] = pd.cut(
        train_df["target_ret"], bins=target_bins, labels=False
    )  # .astype(int)

    test_df["target"] = pd.cut(
        test_df["target_ret"], bins=target_bins, labels=False
    )  # .astype(int)

    # ============================================================
    # KEEP ONLY VALID ROWS
    # ============================================================

    feature_cols = [
        "ret_q",
        "volatility_q",
        "hl_ratio_q",
        "oc_ratio_q",
        "volume_change_q",
        "streak",
    ]

    train_df = train_df.dropna(subset=feature_cols + ["target"])

    test_df = test_df.dropna(subset=feature_cols + ["target"])

    return train_df, test_df, feature_cols


def make_sequences(df, feature_cols, target_col, seq_len):

    X = []
    y = []

    values = df[feature_cols].values
    target = df[target_col].values

    for i in range(seq_len, len(df)):

        X.append(values[i - seq_len : i])

        y.append(target[i])

    return np.array(X), np.array(y)


def prepare_data(train_df, test_df, feature_cols):
    models = {}
    for seq in range(SEQ_INTERVAL, SEQ_LENGTH + SEQ_INTERVAL, SEQ_INTERVAL):

        X_train, y_train = make_sequences(train_df, feature_cols, "target", seq)

        X_test, y_test = make_sequences(test_df, feature_cols, "target", seq)

        print(X_train.shape)
        print(X_test.shape)

        # ============================================================
        # FLATTEN FOR CATBOOST
        # ============================================================

        X_train_flat = X_train.reshape(X_train.shape[0], -1)

        X_test_flat = X_test.reshape(X_test.shape[0], -1)

        # ============================================================
        # CATBOOST
        # ============================================================

        model = CatBoostClassifier(
            loss_function="MultiClass",
            iterations=5000,
            learning_rate=0.01,
            depth=7,
            l2_leaf_reg=50,
            od_type="Iter",
            od_wait=100,
            use_best_model=True,
            verbose=False,
        )
        #     iterations=5000,
        #     depth=6,
        #     learning_rate=0.03,
        #     random_seed=42,
        #     verbose=100
        # )

        model.fit(X_train_flat, y_train, eval_set=(X_test_flat, y_test))

        pred = model.predict(X_test_flat)
        pred = pred.astype(int).ravel()

        # ============================================================
        # METRICS
        # ============================================================

        acc = accuracy_score(y_test, pred)

        mae_bins = mean_absolute_error(y_test, pred)

        # print()
        # print("Accuracy :", acc)
        # print("MAE bins :", mae_bins)

        # exact ±1 bin accuracy
        adj_acc = np.mean(np.abs(pred - y_test) <= 1)

        models[seq] = {
            "acc": acc,
            "mae_bins": mae_bins,
            "adj_acc": adj_acc,
            "model": model,
            "pred": pred,
            "y_test": y_test,
            "val_loss": model.best_score_["validation"]["MultiClass"],
        }

        # for x in range(1,q-1):
        #     adj_acc = np.mean(
        #         np.abs(pred - y_test) <= x
        #     )
        #     models[seq].update({
        #         f"adj_acc_{x}": adj_acc
        #     })

        # print(f"Within {x} bin :", adj_acc)

    return models