"""
Model construction and training.

build_pipeline_map()  -> dict of {ticker: (feature_type, fresh_model_clone)}
train_models()        -> dict of {ticker: (feature_type, fitted_model)}
"""

import pandas as pd
from sklearn.ensemble import (
    RandomForestClassifier,
    HistGradientBoostingClassifier,
    VotingClassifier,
)
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from sklearn.feature_selection import SelectKBest, f_classif
from sklearn.base import clone

from core.config import TICKERS, DATA_DIR, PIPELINE_MAP
from core.features import extract_semantic_features, extract_sequential_features


# ── Base model templates ─────────────────────────────────────────────────────

_RF  = RandomForestClassifier(
    random_state=42, n_estimators=300, max_depth=8,
    min_samples_leaf=5, n_jobs=-1,
)
_GBM = HistGradientBoostingClassifier(
    random_state=42, max_iter=300, l2_regularization=1.0, max_depth=8,
)
_ENSEMBLE = VotingClassifier(
    estimators=[("rf", _RF), ("gbm", _GBM)], voting="soft",
)
_LR_PIPELINE = Pipeline([
    ("scaler", StandardScaler()),
    ("lr", LogisticRegression(C=0.1, max_iter=1000)),
])
_KBEST15_LR = Pipeline([
    ("scaler", StandardScaler()),
    ("kbest", SelectKBest(f_classif, k=15)),
    ("lr", LogisticRegression(C=0.1, max_iter=1000)),
])

_MODEL_TEMPLATES = {
    "ensemble":    _ENSEMBLE,
    "lr_pipeline": _LR_PIPELINE,
    "kbest15_lr":  _KBEST15_LR,
}


# ── Public API ───────────────────────────────────────────────────────────────

def build_pipeline_map():
    """
    Return a dict of {ticker: (feature_type, fresh_model_clone)}.
    Uses PIPELINE_MAP from config to look up the architecture per ticker.
    """
    result = {}
    for ticker in TICKERS:
        feat_type, model_key = PIPELINE_MAP[ticker]
        result[ticker] = (feat_type, clone(_MODEL_TEMPLATES[model_key]))
    return result


def train_models(log_fn=None):
    """
    Load parquet data, extract features, and train all ticker models.

    Parameters
    ----------
    log_fn : callable(str), optional
        Logging function (e.g. logger.info).  Falls back to print.

    Returns
    -------
    dict  { ticker: (feature_type, fitted_model) }
    """
    if log_fn is None:
        log_fn = print

    pipeline_map = build_pipeline_map()
    models = {}

    for ticker in TICKERS:
        fpath = DATA_DIR / f"{ticker}_minute.parquet"
        if not fpath.exists():
            log_fn(f"[{ticker}] Parquet file not found: {fpath}")
            continue

        df = pd.read_parquet(fpath)
        df["date"] = pd.to_datetime(df["date"])
        df.set_index("date", inplace=True)
        df.sort_index(inplace=True)

        # Keep at most 300 trading days
        unique_days = df.index.normalize().unique()
        if len(unique_days) > 300:
            df = df[df.index.normalize().isin(unique_days[-300:])]

        feat_type, clf = pipeline_map[ticker]

        if feat_type == "semantic":
            X, y, _ = extract_semantic_features(df)
        else:
            X, y, _ = extract_sequential_features(df)

        if X is None or X.empty:
            log_fn(f"[{ticker}] Feature extraction returned empty!")
            continue

        clf.fit(X, y)
        models[ticker] = (feat_type, clf)
        log_fn(f"[{ticker}] Model trained  ({feat_type})  |  {len(X)} samples")

    return models