import warnings
warnings.filterwarnings('ignore')

from collections import Counter
from pprint import pprint

import pandas as pd
import numpy as np

import gzip
import zipfile

import seaborn as sns

import nltk 
from nltk.corpus import stopwords

import matplotlib
import matplotlib.pyplot as plt 

import os
import re
import tqdm as tq
from tqdm import tqdm
tqdm.pandas()

import networkx as nx

import sklearn
from sklearn.linear_model import LinearRegression

import watermark

%load_ext watermark
%matplotlib inline

plt.style.use('./d4sci.mplstyle')
colors = plt.rcParams['axes.prop_cycle'].by_key()['color']

text = 'Mary had a little lamb little lamb little lamb'

text_regex = re.compile(r"[a-zA-Z]+")
words = text_regex.findall(text)

filename = "data/2-gram_full.dat.gz"

# If we don't have the processed file
if not os.path.exists(filename):
    edges = {}
    
    text = re.compile(r"^[a-zA-Z]+$")
    
    # Go through the list of 2-gram files from google books
    for url in tqdm(open("data/2-grams.txt", "rt"), total=100):
        filename = os.path.basename(url.strip())
        localfile = os.path.join('data', '2-grams', filename)
        
        # If we don't already have it, download it
        # Depending on your internet connection this might take a while...
        # and you should make sure that the "2-grams" directory exists
        if not os.path.exists(localfile):
            wget.download(url, out='data/2-grams')
        
        # Process it, line by line
        with zipfile.ZipFile(localfile) as zf:
            # About 3m per file in my MBP, about 6h total 
            with zf.open(filename[:-4], 'r') as fp:
                for line in fp:
                    fields = line.decode('utf8').strip().split()
                    year = int(fields[2])

                    word1 = fields[0].lower().strip()
                    word2 = fields[1].lower().strip()
                    count = int(fields[3])

                    # Make sure we have two words with just letters
                    if  text.match(word1) is not None and text.match(word2) is not None:
                        edge = (word1, word2)

                        edges[edge] = edges.get(edge, 0) + count
    
    new_edges = {}
    
    # Write down the edges
    with gzip.open(filename, "wt") as fp:
        print('node_i', 'node_j', 'weight', file=fp)
        
        for edge in tqdm(edges, total=len(edges)):
            print(edge[0], edge[1], edges[edge], file=fp)
            new_edges[edge] = edges[edge]
    
    del new_edges
    del edges

edges = pd.read_csv(filename, sep=' ')

degree_out = edges[['node_i', 'node_j']].groupby('node_i').count()
degree_in = edges[['node_i', 'node_j']].groupby('node_j').count()

def prob_dist(values, resolution=1, normed=True):
    counts = Counter((1+np.round(values.flatten()/resolution))*resolution)
    
    counts = pd.DataFrame(np.array(list(counts.items())))
    counts = counts.sort_values(0)    
    
    if normed:
        counts[1] /= counts[1].sum()
    
    return counts

K_out = prob_dist(degree_out.values, 1)
K_in = prob_dist(degree_in.values, 1)

node_info = degree_in.join(degree_out, how='outer').fillna(0).sort_values('node_i', ascending=False).copy()
node_info.columns=['kin', 'kout']

corr = node_info.corr()
corr.rename(columns={'kin': r'$k_{in}$', 
                    'kout': r'$k_{out}$'}, inplace=True)
corr.rename(index={'kin': r'$k_{in}$', 
                    'kout': r'$k_{out}$'}, inplace=True)

stop_en = stopwords.words('english')
overlap = set(stop_en) & set(node_info.index)

Pw = prob_dist(edges['weight'].values)

edges_small = edges[edges['weight'] >= 1000].dropna().copy()

if not os.path.exists('data/2-gram.dat.gz'):
    edges_small.to_csv('data/2-gram.dat.gz', sep=' ', index=False)

Pw_small = prob_dist(edges_small['weight'].values)

strength_out = edges[['node_i', 'weight']].groupby(['node_i']).sum()
strength_in = edges[['node_j', 'weight']].groupby(['node_j']).sum()

S_out = prob_dist(strength_out.values, 1)
S_in = prob_dist(strength_in.values, 1)

strength_in.columns=['sin']
strength_out.columns=['sout']

node_info = node_info.join(strength_in, how='outer').fillna(1)
node_info = node_info.join(strength_out, how='outer').fillna(1)

corr = node_info.corr()
corr.rename(columns={'kin': r'$k_{in}$', 
                    'sin': r'$s_{in}$',
                    'kout': r'$k_{out}$',
                    'sout': r'$s_{out}$'}, inplace=True)
corr.rename(index={'kin': r'$k_{in}$', 
                    'sin': r'$s_{in}$',
                    'kout': r'$k_{out}$',
                    'sout': r'$s_{out}$'}, inplace=True)

def normalize(x):
    x['weight'] /= x['weight'].sum()
    return x