from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from process_data import split_train_test
import torch
import pickle


def train_predict(
    positive_train_states,
    positive_test_states,
    negative_train_states,
    negative_test_states,
    classifier_save_name="model",
):

    X_train = torch.cat((positive_train_states, negative_train_states), dim=0)

    y_train = [1 for _ in range(positive_train_states.size(0))] + [
        0 for _ in range(negative_train_states.size(0))
    ]

    X_test = torch.cat((positive_test_states, negative_test_states), dim=0)

    y_test = [1 for _ in range(positive_test_states.size(0))] + [
        0 for _ in range(negative_test_states.size(0))
    ]

    model = LogisticRegression(max_iter=2000)

    # Train LogisticRegression
    model.fit(X_train, y_train)

    # Predict
    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    print(f"Accuracy: {accuracy}")

    # Save probe
    with open(classifier_save_name + ".pkl", "wb") as f:
        pickle.dump(model, f)

    return y_pred
