import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

from cartoframes.auth import set_default_credentials
from cartoframes.data.observatory import *
from cartoframes.viz import *
from shapely import wkt

pd.set_option('display.max_columns', None)
sns.set_style('whitegrid')

%matplotlib inline

"""
Explanation: Combining geosocial and financial data to understand retail performance
Geosocial data is location-based social media data that can be interpreted and analyzed as part of any location-oriented business decision. 
In this notebook, Spatial.ai geosocial data is combined with Mastercard financial data to quantify how much of retail performance can be explained by geosocial behavior. We compare the cities of Chicago and Los Angeles to show how retail performance is driven by different social segments depending on the city.
Note this use case leverages premium datasets from CARTO Data Observatory.
The notebook is organized in the following sections:
1. Download data from the Data Observatory
2. Analyzing geosocial drivers
  - Identify where top performers are
  - Characterization of top performers
  - Correlation analysis
0. Setup
0.1. Import packages
End of explanation
"""


from cartoframes.auth import set_default_credentials

set_default_credentials('creds.json')

"""
Explanation: 0.2. Set CARTO default credentials
In order to be able to use the Data Observatory via CARTOframes, you need to set your CARTO account credentials first.
Please, visit the Authentication guide for further detail.
End of explanation
"""


LA_CBG_PATH = 'https://libs.cartocdn.com/cartoframes/samples/la_cbg.csv'
LA_BB = '-118.673619,33.553967,-117.997960,34.360425'
CHICACO_CBG_PATH = 'https://libs.cartocdn.com/cartoframes/samples/chicago_cbg.csv'
CHICAGO_BB = '-88.638285,41.434892,-87.487468,42.502873'

def read_cbg(city):
    cbg_list = pd.DataFrame()
    bbox = ''
    if city == 'la':
        cbg_list = pd.read_csv(LA_CBG_PATH, dtype={'geoid':str})['geoid'].tolist()
        bbox = LA_BB
    elif city == 'chicago':
        cbg_list = pd.read_csv(CHICACO_CBG_PATH, dtype={'geoid':str})['geoid'].tolist()
        bbox = CHICAGO_BB
    return cbg_list, bbox

"""
Explanation: Note about credentials
For security reasons, we recommend storing your credentials in an external file to prevent publishing them by accident when sharing your notebooks. You can get more information in the section Setting your credentials of the Authentication guide.
<a id='section1'></a>
1. Download data from the Data Observatory
In this section, we'll download the two datasets we're interested in and combine them into a single dataframe.
For more information on how to access Data Observatory datasets using CARTOframes visit the Guides or take a look at the <a href='https://carto.com/developers/cartoframes/examples/#example-access-premium-data-from-the-data-observatory' target='_blank'>Access Premium Data</a> template.
Data loading functions
The following function loads the census block group ids for the city of interest. This is used afterwards to filter the data downloaded from the Data Observatory.
End of explanation
"""


def download_social(bbox, cbg_list):
    dataset = Dataset.get('spa_geosocial_s_d5dc42ae')
    sql_query = f"SELECT * FROM $dataset$ WHERE CAST(do_date AS date) >= (SELECT MAX(CAST(do_date AS date)) FROM $dataset$) AND ST_IntersectsBox(geom, {bbox})"
    social = dataset.to_dataframe(sql_query=sql_query)
    social = social[social['geoid'].isin(cbg_list)]
    social.drop(columns=['do_label', 'do_area', 'do_perimeter', 'do_num_vertices'], inplace=True)
    return social

"""
Explanation: The following function downloads the latest Spatial.ai Geosocial Segments data for the specified bounding box.
End of explanation
"""


def download_mastercard(bbox, cbg_list):
    dataset = Dataset.get('mc_geographic__7980c5c3')
    sql_query = f"SELECT * FROM $dataset$ WHERE geo_type = 'm' AND CAST(do_date AS date) >= (SELECT MAX(CAST(do_date AS date)) FROM $dataset$) AND ST_IntersectsBox(geom, {bbox})"
    mrli = dataset.to_dataframe(sql_query=sql_query)
    mrli = mrli[mrli['geoid'].isin(cbg_list)]
    mrli.drop(columns=['do_label', 'do_area', 'do_perimeter', 'do_num_vertices'], inplace=True)

    mrli = mrli.sort_values(['geoid', 'industry', 'segment', 'geo_type', 'do_date']).reset_index(drop=True)
    return mrli

"""
Explanation: The following function downloads the latest Mastercard Geographic Insights data for the specified bounding box.
Note the data is also filtered by geo_type. There are three types of geo_types depending if the indices shown represent a comparison with regards to country (c), province (p) or metropolitan area (m).
End of explanation
"""


def download_data(city):
    cbg_list, bbox = read_cbg(city)
    social = download_social(bbox, cbg_list)
    social.drop(columns='geom', inplace=True)
    mrli = download_mastercard(bbox, cbg_list)
    mrli = mrli[(mrli['segment'] == 'o')]
    mrli = mrli.merge(social, on='geoid')
    mrli = gpd.GeoDataFrame(mrli, crs='epsg:4326')
    return mrli

"""
Explanation: The following function loads the three datasets and merges them into a single one.
Note we can merge the dataframes because both datasets are defined at the census block group level, so the geoid's are common to both of them.
End of explanation
"""


chicago = download_data('chicago')
chicago.head(3)

la = download_data('la')
la.head(3)

"""
Explanation: Download data into dataframes
Here we'll load the data for Chicago and LA.
End of explanation
"""


target_var = 'txn_amt'  # Total transaction amount
industry = 'eap'  # Eating places

ch_ret = chicago[(chicago['industry'] == industry) & (chicago['comparison_level'] == 'DMA')]
la_ret = la[(la['industry'] == industry) & (la['comparison_level'] == 'DMA')]

"""
Explanation: <a id='section2'></a>
2. Analysis
Once we have the data ready for analysis, we will start by selecting a target variable out of the Mastercard financial variables. This is the variable we are interested in explaining. We'll select the variable txn_amt which is the total transaction amount by census block group. In addition, we can select an industry. Mastercard provides data for 5 different industries: eating places, groceries, apparel, automotive fuel, and accomodation. They also provide the industry retail which is the sum of all the latter. For this use case, we'll select eating places, eap.
End of explanation
"""


def classify_finance(value, thres_95):
    if value >= thres_95:
        return 'Top 5%'
    else:
        return 'Rest'

ch_ret.loc[:, f'{target_var}_class'] = list(map(classify_finance, ch_ret[target_var], [ch_ret[target_var].quantile(0.95)]*ch_ret.shape[0]))
la_ret.loc[:, f'{target_var}_class'] = list(map(classify_finance, la_ret[target_var], [la_ret[target_var].quantile(0.95)]*la_ret.shape[0]))

Layout([Map(Layer(ch_ret[ch_ret[f'{target_var}_class'] != 'Rest'], 
                  geom_col='geom',
                  style=color_category_style(f'{target_var}_class', 
                                             cat=['Top 5%'], 
                                             palette=['#009B9E']),
                  legends=color_category_legend('Performers'))),
        Map(Layer(la_ret[la_ret[f'{target_var}_class'] != 'Rest'], 
                  geom_col='geom',
                  style=color_category_style(f'{target_var}_class', 
                                             cat=['Top 5%'],
                                             palette=['#009B9E']),
                  legends=color_category_legend('Performers')))], 
       map_height=420
      )

"""
Explanation: <a id='section21'></a>
2.1 Top performers. Location visualization.
Here we identify the top 5% performers and visualize them. 
An interesting insight here is how in Chicago top performers tend to concentrate in the downtown area, while in LA they are more spread throughout the city.
End of explanation
"""


social_segment_columns = ch_ret.columns[17:-22]

top_ch_ret = ch_ret[ch_ret[f'{target_var}_class'] == 'Top 5%']
top_ch_ret = top_ch_ret[social_segment_columns].describe().transpose()
top_ch_ret.sort_values(['50%'], ascending=False, inplace=True)

top_la_ret = la_ret[la_ret[f'{target_var}_class'] == 'Top 5%']
top_la_ret = top_la_ret[social_segment_columns].describe().transpose()
top_la_ret.sort_values(['50%'], ascending=False, inplace=True)

"""
Explanation: <a id='section22'></a>
2.2 Characterization of top performers
In this section, we'll identify what social segments characterize top performers in Chicago and LA and compare the two cities.
Characterization based on social segments
<b>Main insights</b>:
 - Both cities share 4 out of 7 of their most important social segments (although not with the same order). These segments are related to food and drinks, and lgbtq culture.
 - Regarding the social segments which are different from one city to the other, while in Chicago Asian food and culture, whiskey business, and sweet treats are on the top 7 segments, in LA it is film lovers, fitness fashion, and heartfelt sharing. This shows two very different cities in terms of social behavior patterns.
End of explanation
"""


def palette(segments, color_dict):
    return [color_dict[segment] for segment in segments]

unique_segments = np.unique(top_ch_ret.head(7).index.tolist() + top_la_ret.head(7).index.tolist())
palette_c=['#7F3C8D','#11A579','#3969AC','#F2B701','#E73F74','#80BA5A','#E68310','#008695','#CF1C90','#f97b72', '#4b4b8f', '#A5AA99']
color_dict = dict(zip(unique_segments, palette_c))

def plot_top_segments(df, ax, color_dict):
    sns.barplot(x='50%', y='social_segment', data=df.head(7).reset_index().rename(columns={'index':'social_segment'}), 
                alpha=0.96, ax=ax, palette=palette(df.head(7).index.tolist(), color_dict))
    ax.set_title('Chicago', fontsize=18, fontweight='bold')
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.tick_params(labelsize=15)

fig, axs = plt.subplots(1, 2, figsize=(18, 6))
plot_top_segments(top_ch_ret, axs[0], color_dict)
plot_top_segments(top_la_ret, axs[1], color_dict)
fig.tight_layout()

"""
Explanation: We'll use the following function to make sure social segments have same color in both cities for an easier comparison.
Note we'll use CARTO's bold palette colors. Here you can explore all our palettes.
End of explanation
"""


social_indices_columns = ch_ret.columns[-22:-3]

top_ch_ret_ix = ch_ret[ch_ret[f'{target_var}_class'] == 'Top 5%']
top_ch_ret_ix = top_ch_ret_ix[social_indices_columns].describe().transpose()
top_ch_ret_ix.sort_values(['50%'], ascending=False, inplace=True)

top_la_ret_ix = la_ret[la_ret[f'{target_var}_class'] == 'Top 5%']
top_la_ret_ix = top_la_ret_ix[social_indices_columns].describe().transpose()
top_la_ret_ix.sort_values(['50%'], ascending=False, inplace=True)

unique_segments = np.unique(top_ch_ret_ix.head(7).index.tolist() + top_la_ret_ix.head(7).index.tolist())
palette_c=['#7F3C8D','#11A579','#3969AC','#F2B701','#E73F74','#80BA5A','#E68310','#008695','#CF1C90','#f97b72']
color_dict = dict(zip(unique_segments, palette_c))

def plot_top_indices(df, ax, color_dict):
    sns.barplot(x='50%', y='index', data=df.head(7).reset_index(), 
                alpha=0.96, ax=ax, palette=palette(df.head(7).index.tolist(), color_dict))
    ax.set_title('Chicago', fontsize=18, fontweight='bold')
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.tick_params(labelsize=11)

fig, axs = plt.subplots(1, 2, figsize=(18, 6))
plot_top_indices(top_ch_ret_ix, axs[0], color_dict)
plot_top_indices(top_la_ret_ix, axs[1], color_dict)
fig.tight_layout()

"""
Explanation: Characterization based on indices
<b>Main insights</b>:
- Both cities share 4 out of the 7 most impactful social indices: Breakfast+brunch, coffe, foodie, and high end affinity.
- In Chicago top eating places performers are located where late night, fashion, and entertainment affinity are high, whereas in LA they are located where politically liberal, personal care, and organic+local affinity are high.
End of explanation
"""


def calculate_corrcoef(target_var, df, columns):
    corr_coefs = []
    for ssegment in columns:
        corr_aux = df[~df[ssegment].isnull()]
        corr_coefs.append(np.corrcoef(corr_aux[target_var], corr_aux[ssegment])[0][1])

    corr_df_ch = pd.DataFrame(data={'social_segment':columns, 'corr_coef':corr_coefs})
    corr_df_ch['corr_coef_abs'] = np.abs(corr_df_ch['corr_coef'])
    corr_df_ch.sort_values('corr_coef_abs', ascending=False, inplace=True)
    return corr_df_ch

corr_df_ch = calculate_corrcoef(target_var, ch_ret, social_segment_columns)
corr_df_la = calculate_corrcoef(target_var, la_ret, social_segment_columns)

corr_df_ch['city'] = 'Chicago'
corr_df_la['city'] = 'LA'
corr_df = pd.concat([corr_df_ch, corr_df_la], ignore_index=True)

plt.figure(figsize=(18, 6))
sns.barplot(x='social_segment', y='corr_coef', hue='city', data=corr_df, alpha=0.99)
plt.xticks(rotation=90)
plt.xlabel('')
plt.title(f"Correlation strength of social segments with target var '{target_var}' - industry '{industry}'", 
          fontsize=15, fontweight='light', pad=15)
plt.tight_layout()

"""
Explanation: <a id='section23'></a>
2.3 Correlation analysis
In this section we'll carry out a deeper analysis on how geosocial data can help explain retail performance. We'll compare the most important features for Chicago and LA.
Social segments
Here we calculate the correlation coefficient between every social segment and the total transaction amount for restaurants. This allows us to identify which segments have a strongest impact and the differences between both cities.
This analysis provides very interesting insights. We can see how the social segment "sites to see" impacts positively in Chicago, while it barely impacts in LA.
End of explanation
"""


corr_df_ch_ix = calculate_corrcoef(target_var, ch_ret, social_indices_columns)
corr_df_la_ix = calculate_corrcoef(target_var, la_ret, social_indices_columns)

corr_df_ch_ix['city'] = 'Chicago'
corr_df_la_ix['city'] = 'LA'
corr_df_ix = pd.concat([corr_df_ch_ix, corr_df_la_ix], ignore_index=True)

plt.figure(figsize=(10, 6))
sns.barplot(y='social_segment', x='corr_coef', hue='city', data=corr_df_ix, alpha=0.95)
plt.xticks(rotation=90)
plt.ylabel('')
plt.yticks(fontsize=12)
plt.title(f"Correlation strength of indices with target var '{target_var}' - industry '{industry}'", 
          fontsize=14, fontweight='light', pad=15)
plt.tight_layout()

"""
Explanation: Indices
Now, we'll calculate the correlation coefficient between every social index and the metric total transaction amount for restaurants. This allows us to identify which indices have a strongest impact and the differences between both cities.
It is interesting to see the different correlation coefficients for older affinity and discount affinity in Chicago and LA.
End of explanation
"""


