"""
## Creating training df for "Coleridge - Show US the data" competition
"""
import numpy as np 
import pandas as pd 
import json
import seaborn as sns
import re
no_of_examples = None #reducing the number for experimentation, using None for final code (applied on papers dict generation only)

train_path = '../input/coleridgeinitiative-show-us-the-data/train.csv'
train = pd.read_csv(train_path)

train_folder = '../input/coleridgeinitiative-show-us-the-data/train'
papers = {}
for paper_id in train[:no_of_examples]['Id'].unique():
    with open(f'{train_folder}/{paper_id}.json', 'r') as f:
        paper = json.load(f)
        papers[paper_id] = paper
train.describe()
#creating a dataframe with papers info and text divided into sections

df_list = []

for paper_id in train[:no_of_examples]['Id'].unique():
    
    for count, s in enumerate(papers[paper_id]):
        extended_train = {}
        
        extended_train['Id'] = paper_id
        extended_train['pub_title'] = train.loc[train.Id == paper_id, 'pub_title']
        extended_train['dataset_title'] = train.loc[train.Id == paper_id, 'dataset_title']
        extended_train['dataset_label'] = train.loc[train.Id == paper_id, 'dataset_label']
        extended_train['cleaned_label'] = train.loc[train.Id == paper_id, 'cleaned_label']
        
        extended_train['section_title'] = s['section_title']
        extended_train['section_number'] = count
        extended_train['text'] = s['text']
        
        df_list.append(pd.DataFrame(extended_train))
        
df_extended = pd.concat(df_list, ignore_index = True)
df_extended = df_extended.sort_values(['Id', 'section_number'])
df_extended.shape
df_extended
#adding cleaned_text column

def clean_text(txt):
    return re.sub('[^A-Za-z0-9]+', ' ', str(txt).lower())

df_extended['cleaned_text'] = df_extended['text'].apply(clean_text)
#adding a column indicating if a cleaned label is in the cleaned text with bool
df_extended['label_match'] = df_extended.apply(lambda x:x.cleaned_label in x.cleaned_text, axis =1)
df_extended.groupby('dataset_title').dataset_label.unique().apply(len)
df_extended.groupby('dataset_title').cleaned_label.unique().apply(len)
df_extended.head()
df_extended.to_pickle("coleridge_train_extended.pkl")