from sfi import Frame
import numpy as np
import sys

# get the training and test frame names from Stata
train_frame = sys.argv[1]
test_frame = sys.argv[2]

# use the sfi Frame class to pull data from Stata training frame into Python
train = Frame.connect(train_frame)
train_X = np.array(train.get("seplen sepwid petlen petwid"))
train_y = np.array(train.get("iris"))

# train the model with the training dataset using SVC classifier
from sklearn.svm import SVC
svc_clf = SVC(gamma='auto', random_state=16)
svc_clf.fit(train_X, train_y)

# make predictions with the test dataset and store results back into the test frame
test = Frame.connect(test_frame)
test_X = np.array(test.get("seplen sepwid petlen petwid"))
irispr = svc_clf.predict(test_X)
test.addVarByte('irispr')
test.store('irispr', None, irispr)
