# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
from __future__ import absolute_import, division, print_function, unicode_literals
import numpy as np
import scipy.sparse as sp
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.utils.validation import check_is_fitted
from sklearn.feature_extraction.text import _document_frequency

class BM25Transformer(BaseEstimator, TransformerMixin):
    """
    Parameters
    ----------
    use_idf : boolean, optional (default=True)
    k1 : float, optional (default=2.0)
    b : float, optional (default=0.75)
    References
    ----------
    Okapi BM25: a non-binary model - Introduction to Information Retrieval
    http://nlp.stanford.edu/IR-book/html/htmledition/okapi-bm25-a-non-binary-model-1.html
    """
    def __init__(self, use_idf=True, k1=2.0, b=0.75):
        self.use_idf = use_idf
        self.k1 = k1
        self.b = b

    def fit(self, X):
        """
        Parameters
        ----------
        X : sparse matrix, [n_samples, n_features]
            document-term matrix
        """
        if not sp.issparse(X):
            X = sp.csc_matrix(X)
        if self.use_idf:
            n_samples, n_features = X.shape
            df = _document_frequency(X)
            idf = np.log((n_samples - df + 0.5) / (df + 0.5))
            self._idf_diag = sp.spdiags(idf, diags=0, m=n_features, n=n_features)
        return self

    def transform(self, X, copy=True):
        """
        Parameters
        ----------
        X : sparse matrix, [n_samples, n_features]
            document-term matrix
        copy : boolean, optional (default=True)
        """
        if hasattr(X, 'dtype') and np.issubdtype(X.dtype, np.floating):
            # preserve float family dtype
            X = sp.csr_matrix(X, copy=copy)
        else:
            # convert counts or binary occurrences to floats
            X = sp.csr_matrix(X, dtype=np.float64, copy=copy)

        n_samples, n_features = X.shape

        # Document length (number of terms) in each row
        # Shape is (n_samples, 1)
        dl = X.sum(axis=1)
        # Number of non-zero elements in each row
        # Shape is (n_samples, )
        sz = X.indptr[1:] - X.indptr[0:-1]
        # In each row, repeat `dl` for `sz` times
        # Shape is (sum(sz), )
        # Example
        # -------
        # dl = [4, 5, 6]
        # sz = [1, 2, 3]
        # rep = [4, 5, 5, 6, 6, 6]
        rep = np.repeat(np.asarray(dl), sz)
        # Average document length
        # Scalar value
        avgdl = np.average(dl)
        # Compute BM25 score only for non-zero elements
        data = X.data * (self.k1 + 1) / (X.data + self.k1 * (1 - self.b + self.b * rep / avgdl))
        X = sp.csr_matrix((data, X.indices, X.indptr), shape=X.shape)

        if self.use_idf:
            check_is_fitted(self, '_idf_diag', 'idf vector is not fitted')

            expected_n_features = self._idf_diag.shape[0]
            if n_features != expected_n_features:
                raise ValueError("Input has n_features=%d while the model"
                                 " has been trained with n_features=%d" % (
                                     n_features, expected_n_features))
            # *= doesn't work
            X = X * self._idf_diag

        return X


# -

from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.pipeline import Pipeline
import numpy as np
from scipy import sparse
import pandas as pd



# %matplotlib inline

counts_tr = sparse.load_npz('../data/processed/counts_tr.npz')
counts_tst = sparse.load_npz('../data/processed/counts_tst.npz')
counts_tr

# +
df_tr = pd.read_csv('../data/processed/meta_tr.csv', index_col=0)
df_tst = pd.read_csv('../data/processed/meta_tst.csv', index_col=0)

malwares_tr = (df_tr.label == 'class1').values
malwares_tst = (df_tst.label == 'class1').values
# -

from sklearn.linear_model import LogisticRegression

from sklearn.metrics import confusion_matrix

# ## using tfidf

tfidf_tf = TfidfTransformer()

# +
tfidf_tr = tfidf_tf.fit_transform(counts_tr)
tfidf_tst = tfidf_tf.transform(counts_tst)

tfidf_tr
# -

lr_tfidf = LogisticRegression()
lr_tfidf.fit(tfidf_tr, malwares_tr)

lr_tfidf.score(tfidf_tr, malwares_tr)

lr_tfidf.score(tfidf_tst, malwares_tst)

confusion_matrix(malwares_tst, lr_tfidf.predict(tfidf_tst))

from sklearn.feature_selection import SelectFromModel

# +
sfm = SelectFromModel(lr_tfidf, prefit=True, max_features=100)

# sfm.fit(tfidf_tr, malwares)
n_features = sfm.transform(tfidf_tr).shape[1]

n_features
# -

lr_new = LogisticRegression()
lr_new.fit(sfm.transform(tfidf_tr), malwares_tr)
lr_new.score(sfm.transform(tfidf_tr), malwares_tr), lr_new.score(sfm.transform(tfidf_tst), malwares_tst)

# ## using counts



lr_counts = LogisticRegression()

lr_counts.fit(counts_tr, malwares_tr)

lr_counts.score(counts_tr, malwares_tr)

lr_counts.score(counts_tst, malwares_tst)

confusion_matrix(malwares_tst, lr_counts.predict(counts_tst))

# +
sfm = SelectFromModel(lr_counts, prefit=True, max_features=100)

# sfm.fit(counts_tr, malwares)
n_features = sfm.transform(counts_tr).shape[1]

n_features
# -

lr_new = LogisticRegression()
lr_new.fit(sfm.transform(counts_tr), malwares_tr)
lr_new.score(sfm.transform(counts_tr), malwares_tr), lr_new.score(sfm.transform(counts_tst), malwares_tst)



# ## using BM25

# +
bm = BM25Transformer()  # k1=1.2

bm_tr = bm.fit_transform(counts_tr)
bm_tst = bm.transform(counts_tst)

bm_tr
# -

lr_bm = LogisticRegression(solver='sag')
lr_bm.fit(bm_tr, malwares_tr)

lr_bm.score(bm_tr, malwares_tr)

lr_bm.score(bm_tst, malwares_tst)

confusion_matrix(malwares_tst, lr_bm.predict(bm_tst))

lr_bm.coef_

from sklearn.feature_selection import SelectFromModel

# +
sfm = SelectFromModel(lr_bm, prefit=True, max_features=1000)
# sfm = SelectFromModel(lr_bm, prefit=True, threshold=0.012)

# sfm.fit(counts_tr, malwares)
n_features = sfm.transform(counts_tr).shape[1]

n_features, lr_bm.coef_[0, sfm.get_support()].max()
# -

lr_new = LogisticRegression(C=1)
lr_new.fit(sfm.transform(bm_tr), malwares_tr)
lr_new.score(sfm.transform(bm_tr), malwares_tr), lr_new.score(sfm.transform(bm_tst), malwares_tst)

confusion_matrix(malwares_tst, lr_new.predict(sfm.transform(bm_tst)))

sfm.get_support()



# +
A_tr = sparse.load_npz('../data/processed/A_tr.npz')
B_tr = sparse.load_npz('../data/processed/B_tr.npz')
P_tr = sparse.load_npz('../data/processed/P_tr.npz')
A_tst = sparse.load_npz('../data/processed/A_tst.npz')

A_tr = sparse.csr_matrix(A_tr, dtype='uint32')
A_tst = sparse.csr_matrix(A_tst, dtype='uint32')
# -

from sklearn.svm import SVC

apis = np.random.choice(700000, replace=False, size=500000)
apis = sfm.get_support()

apis.sum()

A_tr = A_tr[:, apis]
B_tr = B_tr[apis, :][:, apis]
P_tr = P_tr[apis, :][:, apis]
A_tst = A_tst[:, apis]

sparse.save_npz('../data/processed/A_reduced_tr.npz', A_tr)
sparse.save_npz('../data/processed/B_reduced_tr.npz', B_tr)
sparse.save_npz('../data/processed/P_reduced_tr.npz', P_tr)
sparse.save_npz('../data/processed/A_reduced_tst.npz', A_tst)



# +
svm = SVC(kernel='precomputed')

gram_train = (A_tr * A_tr.T).todense()

svm.fit(gram_train, malwares_tr)
svm.score(gram_train, malwares_tr)

# +
gram_test = (A_tst * A_tr.T).todense()

svm.score(gram_test, malwares_tst)

# +
svm = SVC(kernel='precomputed')

gram_train = (A_tr * B_tr * P_tr * B_tr * A_tr.T).todense()

svm.fit(gram_train, malwares_tr)
svm.score(gram_train, malwares_tr)

# +
gram_test = (A_tst * B_tr * P_tr * B_tr * A_tr.T).todense()

svm.score(gram_test, malwares_tst)
# -



# +
svm = SVC(kernel='precomputed')

gram_train = (A_tr * P_tr * A_tr.T).todense()

svm.fit(gram_train, malwares_tr)
svm.score(gram_train, malwares_tr)

# +
gram_test = (A_tst * P_tr * A_tr.T).todense()

svm.score(gram_test, malwares_tst)
# -



# +
svm = SVC(kernel='precomputed')

gram_train = (A_tr * B_tr * A_tr.T).todense()

svm.fit(gram_train, malwares_tr)
svm.score(gram_train, malwares_tr)

# +
gram_test = (A_tst * B_tr * A_tr.T).todense()

svm.score(gram_test, malwares_tst)
# -

np.random.choice([1, 2], random_state=0)

StarcoderdataJupyterScriptsDedupFiltered