import pandas as pd
import numpy as np

import os
import sys
import json
import pickle

from pathlib import Path
from tqdm import tqdm

DATA_PATH = Path("../input/Predicting_the_number_of_citations_to_a_paper/")
FEATURE_PATH = Path("../features/")
FEATURE_PATH.mkdir(parents=True, exist_ok=True)

def get_data_iter(fpath):
    with open(fpath, 'r') as f:
        for l in f:
            yield l

columns = ['id', 'submitter', 'authors', 'title', 'comments', 'journal-ref',
           'doi', 'report-no', 'categories', 'license', 'abstract', 'versions',
           'update_date', 'authors_parsed', 'doi_cites']
TARGET = 'cites'

tmp = []
train_iter = get_data_iter(DATA_PATH / "train_data.json")
for line in tqdm(train_iter, total=851_524):
    dict=json.loads(line)
    tmp.append(list(dict.values()))

train_df = pd.DataFrame(tmp, columns=columns+[TARGET])
with open(FEATURE_PATH / "train.pickle", 'wb') as f:
    pickle.dump(train_df, f)

tmp = []
test_iter = get_data_iter(DATA_PATH / "test_data.json")
for line in tqdm(test_iter, total=59_084):
    tmp.append(list(json.loads(line).values()))

test_df = pd.DataFrame(tmp, columns=columns)
with open(FEATURE_PATH / "test.pickle", 'wb') as f:
    pickle.dump(test_df, f)

display(train_df.shape)
display(train_df.head())

display(test_df.shape)
display(test_df.head())