import numpy as np
import torch
from sklearn.model_selection import StratifiedKFold
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline


def get_hs(model, train, test):
    Xs = []
    y = []
    y2 = []

    for idx in range(len(train)):
        data = train[idx]
        pos = data["eol_positions"]
        cr = data["correct_idx"]
        embs = torch.from_numpy(data["embeds_after_l0"]).unsqueeze(0).cuda().half()
        with torch.no_grad():
            out = model.model(inputs_embeds=embs, output_hidden_states=True)
        hs = torch.cat(out.hidden_states)
        for i in range(4):
            Xs.append(hs[:, pos[i]].cpu().numpy())
            y.append(int(i == cr))
        y2.append(cr)

    Xs = np.stack(Xs).reshape(len(Xs) // 4, 4, 24, -1)
    y = np.array(y).reshape(len(y) // 4, 4)
    y2 = np.array(y2)

    Xs_test = []

    for idx in range(len(test)):
        data = test[idx]
        pos = data["eol_positions"]
        embs = torch.from_numpy(data["embeds_after_l0"]).unsqueeze(0).cuda().half()
        with torch.no_grad():
            out = model.model(inputs_embeds=embs, output_hidden_states=True)
        hs = torch.cat(out.hidden_states)
        for i in range(4):
            Xs_test.append(hs[:, pos[i]].cpu().numpy())

    Xs_test = np.stack(Xs_test).reshape(len(Xs_test) // 4, 4, 24, -1)

    return Xs, y, y2, Xs_test


def evaluate(Xs, y, y2, ly):
    X_layer = Xs[:, :, ly]

    train_accs = []
    val_accs = []

    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=0)

    for train_idx, val_idx in skf.split(X_layer, y2):
        X_train = X_layer[train_idx].reshape(4 * len(train_idx), -1)
        X_val = X_layer[val_idx].reshape(4 * len(val_idx), -1)

        y_train = y[train_idx].reshape(-1)

        clf = Pipeline([
            ("scaler", StandardScaler()),
            ("logreg", LogisticRegression(random_state=0, max_iter=1000, C=1e-2)),
        ])

        clf.fit(X_train, y_train)

        pred = clf.predict_proba(X_train)[:, 1]
        train_acc = (pred.reshape(-1, 4).argmax(axis=1) == y2[train_idx]).mean()

        pred = clf.predict_proba(X_val)[:, 1]
        val_acc = (pred.reshape(-1, 4).argmax(axis=1) == y2[val_idx]).mean()

        train_accs.append(train_acc)
        val_accs.append(val_acc)

    return np.mean(train_accs), np.mean(val_accs)


def experiment(Xs, y, y2):
    for ly in range(24):
        train_acc, val_acc = evaluate(Xs, y, y2, ly)

        print(f"Layer {ly}" f" {train_acc:.3f}" f" {val_acc:.3f}")

    for ly in [[13, 14], [13, 15], [14, 15], [13, 14, 15]]:
        train_acc, val_acc = evaluate(Xs, y, y2, ly)

        print(f"Layer {ly}" f" {train_acc:.3f}" f" {val_acc:.3f}")


def solution(model, train, test):
    Xs, y, y2, Xs_test = get_hs(model, train, test)

    ly = [13, 14]

    X = Xs[:, :, ly]
    X = X.reshape(len(X) * 4, -1)

    clf = Pipeline([
        ("scaler", StandardScaler()),
        ("logreg", LogisticRegression(random_state=0, max_iter=1000, C=1e-2)),
    ]).fit(X, y.reshape(-1))

    X_test = Xs_test[:, :, ly]
    X_test = X_test.reshape(len(X_test) * 4, -1)

    test_preds = clf.predict_proba(X_test)[:, 1].reshape(-1, 4).argmax(axis=1)

    return test_preds