#!/usr/bin/env python3
"""Neural-network models with the small sklearn-style surface the pipeline needs.

Only `fit` and `predict_proba` are required, so these wrappers drop into
train_ai.fit_predict beside logistic regression and LightGBM without the rest of
the code knowing the difference.

torch on this cluster needs the conda libstdc++ ahead of the system one:
    export LD_LIBRARY_PATH=$CONDA_PREFIX/lib
The system /lib64/libstdc++.so.6 has no CXXABI_1.3.9 and importing torch without
that pointing at the conda copy fails inside numpy.
"""

from __future__ import annotations

import numpy as np


class TorchMLP:
    """Fully connected net on the tabular features, trained with class weighting.

    The event is rare even after the negative thinning, so the loss carries the
    same positive weight the boosted trees get; without it the net collapses onto
    the majority class and predicts no events at all.
    """

    def __init__(self, hidden=(256, 128, 64), epochs=12, batch=8192, lr=1e-3,
                 pos_weight=1.0, dropout=0.1, seed=0, verbose=False):
        self.hidden, self.epochs, self.batch = hidden, epochs, batch
        self.lr, self.pos_weight, self.dropout = lr, pos_weight, dropout
        self.seed, self.verbose = seed, verbose

    def _build(self, nin):
        import torch.nn as nn
        layers, prev = [], nin
        for h in self.hidden:
            layers += [nn.Linear(prev, h), nn.BatchNorm1d(h), nn.ReLU(),
                       nn.Dropout(self.dropout)]
            prev = h
        layers.append(nn.Linear(prev, 1))
        return nn.Sequential(*layers)

    def fit(self, X, y):
        import torch
        torch.manual_seed(self.seed)
        X = np.asarray(X, dtype="f4")
        y = np.asarray(y, dtype="f4").ravel()

        # standardise here rather than in a pipeline so predict_proba stays cheap
        self.mu_ = X.mean(axis=0)
        self.sd_ = X.std(axis=0)
        self.sd_[self.sd_ < 1e-6] = 1.0

        net = self._build(X.shape[1])
        opt = torch.optim.Adam(net.parameters(), lr=self.lr)
        lossf = torch.nn.BCEWithLogitsLoss(
            pos_weight=torch.tensor([float(self.pos_weight)]))
        n = X.shape[0]
        rng = np.random.default_rng(self.seed)
        net.train()
        for ep in range(self.epochs):
            order = rng.permutation(n)
            tot = 0.0
            for a in range(0, n, self.batch):
                idx = order[a:a + self.batch]
                if idx.size < 2:            # BatchNorm needs more than one row
                    continue
                xb = torch.from_numpy((X[idx] - self.mu_) / self.sd_)
                yb = torch.from_numpy(y[idx]).unsqueeze(1)
                opt.zero_grad()
                loss = lossf(net(xb), yb)
                loss.backward()
                opt.step()
                tot += float(loss) * idx.size
            if self.verbose:
                print("    [mlp] epoch %d/%d loss %.5f"
                      % (ep + 1, self.epochs, tot / n), flush=True)
        net.eval()
        self.net_ = net
        return self

    def predict_proba(self, X):
        import torch
        X = np.asarray(X, dtype="f4")
        out = np.empty(X.shape[0], dtype="f8")
        with torch.no_grad():
            for a in range(0, X.shape[0], 65536):
                xb = torch.from_numpy((X[a:a + 65536] - self.mu_) / self.sd_)
                out[a:a + 65536] = torch.sigmoid(self.net_(xb)).numpy().ravel()
        return np.column_stack([1.0 - out, out])

    def predict(self, X):
        return self.predict_proba(X)[:, 1]