"""md-agent candidate models (rebuilt post-outage). Registered by run.py."""
import numpy as np
import pandas as pd
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler, RobustScaler
from sklearn.pipeline import make_pipeline
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import BaggingClassifier, HistGradientBoostingClassifier

import batteryswap_example.models_cv as mc

# risk-direction of p42 in each feature (0 = unconstrained)
MONO = {
    'v_last': -1, 'v_drop': +1, 'slope_30': -1, 'slope_90': -1,
    'v_min_30': -1, 'staleness': +1, 'age_days': +1, 'slope_lastk': -1,
    'drop_recent': +1, 'below_frac': +1, 'accel': -1, 'n_30': 0,
    'tau_projected': -1,
    'vrel': -1, 'vmin7_rank': -1, 'r_vlast': -1, 'r_tau': -1,
}


class _DaysMixin:
    def _fit_days(self, train):
        dying = train[train['dies42'] == 1]
        self.days = mc.DayQuantiles().fit(dying[mc.FEATURES], dying['days_to_eol'])

    def _out(self, rows, p42):
        out = self.days.predict(rows[mc.FEATURES])
        out['p42'] = p42
        return out


class HGBMono(_DaysMixin):
    """HistGradientBoosting with monotonic constraints (risk-direction by name)."""

    def __init__(self, **kw):
        self.kw = dict(max_iter=200, max_depth=3, learning_rate=0.1,
                       min_samples_leaf=20, random_state=0)
        self.kw.update(kw)

    def fit(self, train):
        feats = list(mc.FEATURES)
        cst = [MONO.get(f, 0) for f in feats]
        self.feats = feats
        self.clf = HistGradientBoostingClassifier(monotonic_cst=cst, **self.kw)
        self.clf.fit(train[feats], train['dies42'])
        self._fit_days(train)
        return self

    def predict(self, rows):
        return self._out(rows, self.clf.predict_proba(rows[self.feats])[:, 1])


class MeanBlend(_DaysMixin):
    """Mean of logbag p42 and a partner model's p42."""

    def __init__(self, partner, w=0.5):
        self.partner_factory = partner
        self.w = w

    def fit(self, train):
        self.a = mc.SklearnRUL('logbag').fit(train)
        self.b = self.partner_factory().fit(train)
        return self

    def predict(self, rows):
        out = self.a.predict(rows)
        out['p42'] = ((1 - self.w) * out['p42']
                      + self.w * self.b.predict(rows)['p42'])
        return out


class MeanBlendCM(MeanBlend):
    """Mean blend, count-matched: p42 rescaled so sum(p42) equals logbag's."""

    def predict(self, rows):
        pa = self.a.predict(rows)
        pb = self.b.predict(rows)['p42']
        blend = (1 - self.w) * pa['p42'] + self.w * pb
        s = blend.sum()
        if s > 0:
            blend = (blend * (pa['p42'].sum() / s)).clip(upper=1.0)
        out = pa
        out['p42'] = blend
        return out


class OrderBlend3(_DaysMixin):
    """Value-preserving 3-way ORDER blend: logbag p42 values kept (count
    machinery untouched), ordering = weighted mean of pct-ranks of
    logbag / partner2 / partner3."""

    def __init__(self, p2, w2, p3=None, w3=0.0):
        self.p2f, self.w2 = p2, w2
        self.p3f, self.w3 = p3, w3

    def fit(self, train):
        self.a = mc.SklearnRUL('logbag').fit(train)
        self.b = self.p2f().fit(train)
        self.c = self.p3f().fit(train) if self.p3f is not None else None
        return self

    def predict(self, rows):
        out = self.a.predict(rows)
        w1 = 1.0 - self.w2 - self.w3
        r = (w1 * out['p42'].rank(pct=True)
             + self.w2 * self.b.predict(rows)['p42'].rank(pct=True))
        if self.c is not None:
            r = r + self.w3 * self.c.predict(rows)['p42'].rank(pct=True)
        vals = np.sort(out['p42'].values)[::-1]
        out['p42'] = pd.Series(
            vals, index=r.sort_values(ascending=False).index).reindex(out.index)
        return out


def register():
    M = mc.MODELS
    M['md_hgb_d2'] = lambda: HGBMono(max_depth=2, l2_regularization=1.0)
    M['md_hgb_d2l2'] = lambda: HGBMono(max_depth=2, l2_regularization=2.0)
    M['md_hgb_d2l10'] = lambda: HGBMono(max_depth=2, l2_regularization=10.0)
    M['md_hgb_d1'] = lambda: HGBMono(max_depth=1, l2_regularization=1.0,
                                     learning_rate=0.05, max_iter=400)
    M['md_hgb_d1l25'] = lambda: HGBMono(max_depth=1, l2_regularization=5.0,
                                        learning_rate=0.05, max_iter=400)
    M['md_bl_hgbd2'] = lambda: MeanBlend(M['md_hgb_d2'], 0.5)
    M['md_blcm_hgbd2'] = lambda: MeanBlendCM(M['md_hgb_d2'], 0.5)
    M['md_blcm_hgbd1'] = lambda: MeanBlendCM(M['md_hgb_d1'], 0.5)
    # 3-way order blend on the frontier recipe: logbag .55 / hgb .25 / wiener .20
    M['md_ord3_d2'] = lambda: OrderBlend3(M['md_hgb_d2'], 0.25, mc.WienerFHT, 0.20)
    M['md_ord3_d1l25'] = lambda: OrderBlend3(M['md_hgb_d1l25'], 0.25,
                                             mc.WienerFHT, 0.20)
    M['md_ord3_d2w15'] = lambda: OrderBlend3(M['md_hgb_d2'], 0.15, mc.WienerFHT, 0.20)