from typing import Optional
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, classification_report
from sklearn.model_selection import TimeSeriesSplit
import joblib
from config import RANDOM_SEED, TEST_SIZE


MODEL_FEATURES = [
    'return_1', 'return_5', 'return_20',
    'close_to_sma_5', 'close_to_sma_10', 'close_to_sma_20', 'close_to_sma_50',
    'rsi', 'bb_width', 'bb_position',
    'macd', 'macd_signal', 'macd_hist',
    'atr_pct', 'volume_ratio',
]


def prepare_data(df: pd.DataFrame, feature_cols: Optional[list[str]] = None) -> tuple[pd.DataFrame, pd.Series]:
    if feature_cols is None:
        feature_cols = MODEL_FEATURES
    df = df.dropna().copy()
    # Валидация
    missing = [c for c in feature_cols if c not in df.columns]
    if missing:
        raise KeyError(f"prepare_data: отсутствуют колонки {missing} в датафрейме")
    if 'target_direction' not in df.columns:
        raise KeyError("prepare_data: отсутствует колонка 'target_direction'. Вызовите add_target() перед prepare_data()")
    X = df[feature_cols]
    y = df['target_direction']
    return X, y


def train_test_split_time(X: pd.DataFrame, y: pd.Series, test_size: float = TEST_SIZE):
    if len(X) < 5:
        raise ValueError(f"train_test_split_time: слишком мало данных ({len(X)} строк), нужно минимум 5")
    split_idx = int(len(X) * (1 - test_size))
    if split_idx >= len(X):
        split_idx = len(X) - 1
    if split_idx == 0:
        split_idx = 1
    X_train = X.iloc[:split_idx]
    X_test = X.iloc[split_idx:]
    y_train = y.iloc[:split_idx]
    y_test = y.iloc[split_idx:]
    return X_train, X_test, y_train, y_test


def train_model(
    X_train: pd.DataFrame,
    y_train: pd.Series,
    model_type: str = 'rf',
    **kwargs,
):
    params = {
        'n_estimators': kwargs.get('n_estimators', 200),
        'max_depth': kwargs.get('max_depth', 10),
        'min_samples_split': kwargs.get('min_samples_split', 5),
        'min_samples_leaf': kwargs.get('min_samples_leaf', 2),
        'random_state': RANDOM_SEED,
        'n_jobs': -1,
    }
    model = RandomForestClassifier(**params)
    model.fit(X_train, y_train)
    return model


def evaluate_model(model, X_test: pd.DataFrame, y_test: pd.Series) -> dict:
    y_pred = model.predict(X_test)
    return {
        'accuracy': accuracy_score(y_test, y_pred),
        'precision': precision_score(y_test, y_pred, zero_division=0),
        'recall': recall_score(y_test, y_pred, zero_division=0),
        'f1': f1_score(y_test, y_pred, zero_division=0),
        'classification_report': classification_report(y_test, y_pred, zero_division=0),
    }


def cross_validate(X: pd.DataFrame, y: pd.Series, n_splits: int = 5) -> list[dict]:
    if len(X) <= n_splits:
        raise ValueError(f"cross_validate: слишком мало данных ({len(X)} строк) для {n_splits} фолдов")
    tscv = TimeSeriesSplit(n_splits=min(n_splits, len(X) - 1))
    results = []
    for fold, (train_idx, val_idx) in enumerate(tscv.split(X)):
        X_train = X.iloc[train_idx]
        X_val = X.iloc[val_idx]
        y_train = y.iloc[train_idx]
        y_val = y.iloc[val_idx]
        model = train_model(X_train, y_train)
        metrics = evaluate_model(model, X_val, y_val)
        metrics['fold'] = fold
        results.append(metrics)
    return results


def save_model(model, scaler: Optional[StandardScaler], path: str):
    joblib.dump({'model': model, 'scaler': scaler, 'features': MODEL_FEATURES}, path)


def load_model(path: str) -> dict:
    return joblib.load(path)