"""Shared Model A pipeline helpers."""

from __future__ import annotations

import warnings
from contextlib import contextmanager

import numpy as np
from sklearn.linear_model import Ridge
from sklearn.preprocessing import PolynomialFeatures, StandardScaler

from prepare import (
    CASE_FEATURE_NAMES,
    DEFAULT_CASE_FEATURE_ROW,
    FEATURE_NAMES,
    build_case_feature_index,
    case_feature_dict,
    compute_case_features,
    extract_features,
)

ALPHA = 5.2
RIDGE_SOLVER = "lsqr"
TOP_FEAT_IDX = [18, 19, 17, 6, 24, 1, 10, 2]
SUBSCORE_NAMES = ("team", "traction", "market", "risk")


@contextmanager
def suppress_matmul_warnings():
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", message=".*encountered in matmul", category=RuntimeWarning)
        yield


def fit_ridge(x, y, sample_weight=None):
    with suppress_matmul_warnings():
        model = Ridge(alpha=ALPHA, solver=RIDGE_SOLVER)
        model.fit(x, y, sample_weight=sample_weight)
    return model


def predict_ridge(model, x):
    with suppress_matmul_warnings():
        return model.predict(x)


def build_base_feature_matrix(texts, inputs):
    return np.array([extract_features(t, inp) for t, inp in zip(texts, inputs)], dtype=np.float64)


def build_case_feature_block(texts, examples, query_texts=None, query_matrix=None, self_indices=None):
    query_texts = texts if query_texts is None else query_texts
    case_index = build_case_feature_index(texts, examples)
    if case_index is None:
        return case_index, np.tile(DEFAULT_CASE_FEATURE_ROW, (len(query_texts), 1))

    return case_index, compute_case_features(
        query_texts,
        case_index,
        query_matrix=query_matrix,
        self_indices=self_indices,
    )


def append_case_features(base_features, case_features):
    return np.hstack([base_features, case_features])


def fit_poly_block(base_features, core_features):
    poly = PolynomialFeatures(degree=2, include_bias=False)
    poly_features = poly.fit_transform(base_features[:, TOP_FEAT_IDX])
    return poly, np.hstack([core_features, poly_features])


def transform_poly_block(base_features, core_features, poly):
    poly_features = poly.transform(base_features[:, TOP_FEAT_IDX])
    return np.hstack([core_features, poly_features])


def fit_sub_models(feat_train_aug, sub_train_np, sample_weight):
    scaler_base = StandardScaler()
    feat_train_base_n = scaler_base.fit_transform(feat_train_aug)

    sub_models = {}
    for i, name in enumerate(SUBSCORE_NAMES):
        y_sub = sub_train_np[:, i]
        mask = y_sub >= 0
        if mask.sum() > 0:
            sub_models[name] = fit_ridge(feat_train_base_n[mask], y_sub[mask], sample_weight=sample_weight[mask])
    return scaler_base, sub_models, feat_train_base_n


def predict_sub_scores(sub_models, feat_base_n):
    sub_preds = np.zeros((feat_base_n.shape[0], len(SUBSCORE_NAMES)), dtype=np.float64)
    sub_scores = {}
    for i, name in enumerate(SUBSCORE_NAMES):
        if name in sub_models:
            values = np.clip(predict_ridge(sub_models[name], feat_base_n), 1, 10)
        else:
            values = np.full(feat_base_n.shape[0], 5.0, dtype=np.float64)
        sub_preds[:, i] = values
        if feat_base_n.shape[0] == 1:
            sub_scores[name] = float(values[0])
    return sub_preds, sub_scores


def fit_model_a_bundle(texts_train, inputs_train, train_examples, sub_train_np, s_train, sample_weight):
    feat_train = build_base_feature_matrix(texts_train, inputs_train)
    case_index, case_train = build_case_feature_block(
        texts_train,
        train_examples,
        query_matrix=None,
        self_indices=np.arange(len(texts_train)),
    )
    core_train = append_case_features(feat_train, case_train)
    poly, feat_train_aug = fit_poly_block(feat_train, core_train)

    scaler_base, sub_models, feat_train_base_n = fit_sub_models(feat_train_aug, sub_train_np, sample_weight)
    sub_preds_train, _ = predict_sub_scores(sub_models, feat_train_base_n)

    feat_train_stacked = np.hstack([feat_train_aug, sub_preds_train])
    scaler = StandardScaler()
    feat_train_n = scaler.fit_transform(feat_train_stacked)
    model = fit_ridge(feat_train_n, s_train, sample_weight=sample_weight)

    bundle = {
        "model": model,
        "sub_models": sub_models,
        "scaler_base": scaler_base,
        "scaler": scaler,
        "poly": poly,
        "poly_top_feat_idx": TOP_FEAT_IDX,
        "use_embeddings": False,
        "use_stacking": True,
        "alpha": ALPHA,
        "ridge_solver": RIDGE_SOLVER,
        "feature_names": FEATURE_NAMES,
        "case_feature_names": CASE_FEATURE_NAMES,
        "base_feature_count": len(FEATURE_NAMES),
        "case_feature_index": case_index,
    }
    return bundle


def transform_model_a_batch(bundle, texts, inputs, query_matrix=None, self_indices=None):
    feat = build_base_feature_matrix(texts, inputs)
    case_index = bundle.get("case_feature_index")
    if case_index is not None:
        case_rows = compute_case_features(
            texts,
            case_index,
            query_matrix=query_matrix,
            self_indices=self_indices,
        )
    else:
        case_rows = np.tile(DEFAULT_CASE_FEATURE_ROW, (len(texts), 1))

    feat_core = append_case_features(feat, case_rows)
    feat_aug = transform_poly_block(feat, feat_core, bundle["poly"])
    feat_base_n = bundle["scaler_base"].transform(feat_aug)
    sub_preds, _ = predict_sub_scores(bundle["sub_models"], feat_base_n)
    feat_stacked = np.hstack([feat_aug, sub_preds])
    feat_n = bundle["scaler"].transform(feat_stacked)
    case_contexts = [case_feature_dict(row) for row in case_rows]
    return feat, case_rows, case_contexts, feat_n, sub_preds


def transform_model_a_inputs(bundle, text, inp):
    feat, case_rows, case_contexts, feat_n, sub_preds = transform_model_a_batch(bundle, [text], [inp])
    sub_scores = {
        name: float(sub_preds[0, i])
        for i, name in enumerate(SUBSCORE_NAMES)
    }
    return feat, case_rows[0], case_contexts[0], feat_n, sub_scores


def predict_model_a_bundle(bundle, texts, inputs, query_matrix=None, self_indices=None):
    feat, case_rows, case_contexts, feat_n, sub_preds = transform_model_a_batch(
        bundle,
        texts,
        inputs,
        query_matrix=query_matrix,
        self_indices=self_indices,
    )
    preds = np.clip(predict_ridge(bundle["model"], feat_n), 1, 10)
    return preds, case_rows, case_contexts, sub_preds, feat, feat_n
