# -- Code Cell --
import pandas as pd
import numpy as np
train = pd.read_csv('train.csv')
test = pd.read_csv("test.csv")

# -- Code Cell --
text = train['text'].unique()
text

# -- Code Cell --
groups = train.groupby('id').agg(list)
groups['coords'][1][0]

# -- Code Cell --
from nltk.stem import PorterStemmer,WordNetLemmatizer
from nltk.tokenize import TweetTokenizer
import nltk
nltk.download('averaged_perceptron_tagger_eng')
from nltk.corpus import stopwords,wordnet
def wordnet_tag(word):
    tag = nltk.pos_tag([word.lower()])[0][1][0].lower()
 
    tag_dict = {
        'n': wordnet.NOUN,
        'r' : wordnet.ADV,
        'v' : wordnet.VERB,
        'j': wordnet.ADJ,
 
    }
    return tag_dict.get(tag,wordnet.NOUN)
lemmatizer = WordNetLemmatizer()
stemmer = PorterStemmer()
tokenizer = TweetTokenizer(preserve_case=False,reduce_len=True)
stop_words = stopwords.words('english')
 
def process_text(text):
 
    tokens = tokenizer.tokenize(text.lower())
    tokens = [token for token in tokens if token.isalpha() and token not in stop_words]
 
    tokens = [lemmatizer.lemmatize(word,wordnet_tag(word)) for word in tokens]
 
    return ' '.join(tokens)

# -- Code Cell --
from sklearn.decomposition import TruncatedSVD
svd= TruncatedSVD(n_components=300)
from sklearn.feature_extraction.text import TfidfVectorizer
tfidf = TfidfVectorizer(ngram_range=(1,3))
processed_text = [process_text(t) for t in text]
text_tf = tfidf.fit_transform(processed_text)
text_svd = svd.fit_transform(text_tf)

# -- Code Cell --
train_avg = train.groupby('id')['coords'].apply(lambda x: np.mean(np.vstack([[float(v) for v in coord.split('|') if v != '']for coord in x]),axis=0))

# -- Code Cell --
from sklearn.linear_model import Ridge
model = Ridge(alpha=1)
X_train = np.vstack(train_avg.values)
model.fit(X_train, text_svd)

# -- Code Cell --
groups2 = test.groupby('datapointID').agg(list)
groups2['coords'][1][0]

# -- Code Cell --
test_avg = test.groupby('datapointID')['coords'].apply(lambda x: np.mean(np.vstack([[float(v) for v in coord.split('|') if v != '']for coord in x]),axis=0))
X_test = np.vstack(test_avg.values)

# -- Code Cell --
preds = model.predict(X_test)

# -- Code Cell --
candidates = pd.read_csv('candidates.csv')

# -- Code Cell --
candidates.head()

# -- Code Cell --
texts_candidate = candidates["text"]

# -- Code Cell --
processed_candidates = [process_text(t) for t in texts_candidate]
texts_candidate_tf = tfidf.transform(processed_candidates)

# -- Code Cell --
texts_candidate_svd = svd.transform(texts_candidate_tf)
len(texts_candidate_svd)

# -- Code Cell --


# -- Code Cell --
from sklearn.metrics.pairwise import euclidean_distances

dist = euclidean_distances(preds, texts_candidate_svd)
labels = dist.argmin(axis=1)

# -- Code Cell --
from scipy.spatial.distance import cdist
import numpy as np
 
D = cdist(preds, texts_candidate_svd, metric='euclidean')
 
n_pred, n_cand = D.shape
pred = np.full(n_pred, -1, dtype=int)
used = set()
 
for i in range(n_pred):
    sorted_idx = np.argsort(D[i])
    
    for idx in sorted_idx:
        if idx not in used:
            pred[i] = idx
            used.add(idx)
            break

# -- Code Cell --
pred

# -- Code Cell --
sub = pd.DataFrame({
    "subtaskID":1,
    "datapointID":[i + 1 for i in range(50)],
    'answer':pred
}).to_csv("subs.csv",index=False)

# -- Code Cell --
