%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import glob
import os
import scipy as sp
from scipy import stats

from tools.plt import color2d #from the 'srcole/tools' repo
from matplotlib import cm
pd.options.display.max_rows = 1000
pd.options.display.max_columns = 100

"""
Explanation: https://docs.google.com/document/d/15GMxvJYAUO-b96c18QmfcF278IVVBBtCFk8nayDp9oY/edit
End of explanation
"""


# Load cities info
df_cities = pd.read_csv('/gh/data2/yelp/city_pop.csv', index_col=0)
df_cities.head()

# Load restaurants
df_restaurants = pd.read_csv('/gh/data2/yelp/food_by_city/df_restaurants.csv', index_col=0)
df_restaurants.head()

# Load categories by restaurant
df_categories = pd.read_pickle('/gh/data2/yelp/food_by_city/df_categories_sparse.pkl')
df_categories.head()

# These are used for the 'category' input to the search function
df_categories_info = pd.read_json('/gh/data2/yelp/categories.json')
df_categories_info.head()

"""
Explanation: Load dataframes
End of explanation
"""


# New dataframe: For each cuisine, compute the average rating, average price, and # restaurants
all_cuisines = df_categories.keys()
cuisine_dict = {'cuisine': [],
                'avg_rating': [],
                'avg_cost': [],
                'N': []}
for k in all_cuisines:
    df_temp = df_restaurants[df_categories[k]==1]
    cuisine_dict['cuisine'].append(k)
    cuisine_dict['avg_rating'].append(df_temp['rating'].mean())
    cuisine_dict['avg_cost'].append(df_temp['cost'].mean())
    cuisine_dict['N'].append(len(df_temp))
df_cuisine = pd.DataFrame.from_dict(cuisine_dict)

# Determine cuisines of interest
# Only look at cuisines with at least 2000 restaurants
min_N = 1000
category_counts = df_categories.sum()
categories_keep = category_counts[category_counts > min_N]
cuisines_rmv = ['bars', 'beer_and_wine', 'beerbar', 'breweries', 'butcher', 'cafes', 'catering',
                'chickenshop', 'cocktailbars', 'convenience', 'cosmetics', 'customcakes',
                'deptstores', 'divebars', 'drugstores', 'eventplanning', 'farmersmarket', 'fooddeliveryservices',
                'foodstands', 'gastropubs', 'gourmet', 'grocery', 'healthmarkets', 'importedfood', 'intlgrocery',
                'karaoke', 'lounges', 'markets', 'meats', 'musicvenues', 'personalchefs', 'pubs',
                'restaurants', 'salvadoran', 'seafoodmarkets', 'servicestations', 'sportsbars', 'streetvendors',
                'tapasmallplates', 'venues', 'wine_bars', 'wineries']
categories_keep.drop(cuisines_rmv, inplace=True)
categories_keep = categories_keep.keys()

df_categories.loc[:,categories_keep]

"""
Explanation: Cuisines by city
End of explanation
"""


# Set up dataframe for restaurants with categories of interest
restaurant_have_category = df_categories.loc[:,categories_keep].sum(axis=1).to_dict()
df_restaurants_keep_idx = [k for k in restaurant_have_category.keys() if restaurant_have_category[k]]
df_restaurants_temp = df_restaurants.loc[df_restaurants_keep_idx].reset_index(drop=True)
df_categories_temp = df_categories.loc[df_restaurants_keep_idx,categories_keep].reset_index(drop=True)
df_restaurants_temp = df_restaurants_temp.merge(df_categories_temp, left_index=True, right_index=True)

# Compute fraction of each cuisine by city
df_city_cuisines = df_restaurants_temp.groupby('city').mean()
df_state_cuisines = df_restaurants_temp.groupby('state').mean()

"""
Explanation: IN THIS NOTEBOOK AND OTHERS, NEED TO FIX HOW USE DF_CAT IN BELOW DF NOW THAT ITS SPARSE
End of explanation
"""


df_city_cuisines.head(5)

"""
Explanation: Explore features by city

rating, review_count, cost, has_delivery, has_pickup
each cuisine
End of explanation
"""


df_city_cuisines.sort_values('rating', ascending=False, inplace=True)

N=60
plt.figure(figsize=(30,5))
plt.bar(np.arange(N), df_city_cuisines['rating'].values[:N], color='k', ecolor='.5')
plt.xticks(np.arange(N), df_city_cuisines.index[:N])
plt.ylabel('Average rating', size=20)
plt.xlabel('City', size=20)
plt.xticks(size=15, rotation='vertical')
# plt.yticks([10**3, 10**4, 10**5], size=15)
plt.ylim((3.5, 4.5))
plt.xlim((-1, N))

"""
Explanation: Highest average rating

Highest average rating are the most popular cities because yelp will return the top ones in each city
End of explanation
"""


c = 'bubbletea'
df_city_cuisines.sort_values(c, ascending=False, inplace=True)

N=60
plt.figure(figsize=(30,5))
plt.bar(np.arange(N), df_city_cuisines[c].values[:N], color='k', ecolor='.5')
plt.xticks(np.arange(N), df_city_cuisines.index[:N])
plt.ylabel('Fraction of restaurants are\n'+c, size=20)
plt.xlabel('City', size=20)
plt.xticks(size=15, rotation='vertical')
plt.xlim((-1, N))

c = 'mexican'
df_state_cuisines.sort_values(c, ascending=False, inplace=True)

plt.figure(figsize=(30,5))
plt.bar(np.arange(len(df_state_cuisines)), df_state_cuisines[c].values, color='k', ecolor='.5')
plt.xticks(np.arange(len(df_state_cuisines)), df_state_cuisines.index)
plt.ylabel('Fraction of restaurants are\n'+c, size=20)
plt.xlabel('State', size=20)
plt.xticks(size=15, rotation='vertical')
plt.xlim((-1, len(df_state_cuisines)))

c = 'italian'
df_state_cuisines.sort_values(c, ascending=False, inplace=True)

plt.figure(figsize=(30,5))
plt.bar(np.arange(len(df_state_cuisines)), df_state_cuisines[c].values, color='k', ecolor='.5')
plt.xticks(np.arange(len(df_state_cuisines)), df_state_cuisines.index)
plt.ylabel('Fraction of restaurants are\n'+c, size=20)
plt.xlabel('State', size=20)
plt.xticks(size=15, rotation='vertical')
plt.xlim((-1, len(df_state_cuisines)))

"""
Explanation: boba
End of explanation
"""


