"""
StackedMultimodal — the final multi-modal model (used for training AND deployment).

Architecture (per survival horizon):
  base_tabular : HistGradientBoosting (class-balanced) on 7 tabular features
                 (country, cyclic month, log viral loads CBPV/DWV/KBV)
  base_audio   : StandardScaler -> PCA(16) -> LogisticRegression (balanced) on the
                 191 librosa acoustic features
  meta         : LogisticRegression on [p_tabular, p_audio]  (learned late fusion)

fit() trains the meta on inner GroupKFold out-of-fold base predictions (no leakage),
then refits both bases on all supplied data and tunes the decision threshold on the
inner-OOF meta probabilities. Because fit() is self-contained, the same object can
be honestly evaluated with an outer GroupKFold and then fit on all data for release.
"""
import numpy as np
from sklearn.model_selection import GroupKFold
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from sklearn.impute import SimpleImputer
from sklearn.decomposition import PCA
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.metrics import f1_score

RNG = 42


def make_tabular():
    return Pipeline([
        ('imp', SimpleImputer(strategy='median')),
        ('clf', HistGradientBoostingClassifier(
            max_iter=400, learning_rate=0.05, max_leaf_nodes=31,
            l2_regularization=1.0, class_weight='balanced', random_state=RNG))])


def make_audio(k=16):
    return Pipeline([
        ('imp', SimpleImputer(strategy='median')),
        ('sc', StandardScaler()),
        ('pca', PCA(n_components=k, random_state=RNG)),
        ('clf', LogisticRegression(max_iter=2000, class_weight='balanced', C=0.5))])


def tune_threshold(y, p):
    ts = np.clip(np.unique(np.round(p, 4)), 1e-3, 1 - 1e-3)
    best_t, best_f1 = 0.5, -1.0
    for t in ts:
        f = f1_score(y, (p >= t).astype(int), zero_division=0)
        if f > best_f1:
            best_f1, best_t = f, t
    return float(best_t)


class StackedMultimodal:
    """Fixed-weight late fusion: p = w_tab * p_tabular + (1 - w_tab) * p_audio.

    A weight scan under GroupKFold showed w_tab ~= 0.8 dominates tabular-only on
    BOTH ROC-AUC and average precision for both horizons (audio helps rank
    borderline cases without diluting the concentrated metadata signal). A learned
    meta-learner over-weighted audio and hurt precision, so we keep the robust,
    transparent fixed weight.
    """
    def __init__(self, pca_k=16, inner_splits=5, w_tab=0.8):
        self.pca_k = pca_k
        self.inner_splits = inner_splits
        self.w_tab = w_tab

    def fit(self, Xtab, Xaud, y, groups):
        y = np.asarray(y)
        n = len(y)
        oof_t = np.zeros(n)
        oof_a = np.zeros(n)
        gkf = GroupKFold(n_splits=self.inner_splits)
        for tr, te in gkf.split(Xtab, y, groups):
            mt = make_tabular().fit(Xtab[tr], y[tr])
            ma = make_audio(self.pca_k).fit(Xaud[tr], y[tr])
            oof_t[te] = mt.predict_proba(Xtab[te])[:, 1]
            oof_a[te] = ma.predict_proba(Xaud[te])[:, 1]
        # refit bases on ALL data
        self.tab_ = make_tabular().fit(Xtab, y)
        self.aud_ = make_audio(self.pca_k).fit(Xaud, y)
        # threshold on inner-OOF fused probs (honest, no test leakage)
        self.oof_meta_ = self.w_tab * oof_t + (1 - self.w_tab) * oof_a
        self.threshold_ = tune_threshold(y, self.oof_meta_)
        self.base_oof_ = (oof_t, oof_a)
        return self

    def predict_proba(self, Xtab, Xaud):
        pt = self.tab_.predict_proba(Xtab)[:, 1]
        pa = self.aud_.predict_proba(Xaud)[:, 1]
        p = self.w_tab * pt + (1 - self.w_tab) * pa
        return p, pt, pa

    def predict(self, Xtab, Xaud):
        p, _, _ = self.predict_proba(Xtab, Xaud)
        return (p >= self.threshold_).astype(int)