!conda install -c conda-forge -q -y squarify

from datetime import datetime
from typing import List, Optional, Union

import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
import pandas as pd
from scipy.stats import probplot

import seaborn as sns
import squarify

# Settings
%config InlineBackend.figure_format = 'retina'
sns.set_theme(font_scale=1.25)
random_seed = 87

cars_rawdf = pd.read_csv('../data/autoscout24_germany_dataset.csv')
cars_rawdf.info()

cars_rawdf.sample(frac=1, random_state=random_seed).head(n=10)

print(
    'Counts of unique offer types:',
    cars_rawdf['offerType'].value_counts(),
    sep='\n'
)

cars = cars_rawdf[cars_rawdf['offerType'] == 'Used'].copy()
cars.sample(frac=1).head(n=10)

cars.drop(columns='offerType', inplace=True)
cars.info()

cars['age'] = datetime.now().year - cars['year']
cars.drop(columns='year', inplace=True)

cars.sample(frac=1).head(n=10)

numerical_features = ['mileage', 'hp', 'age']

for feature in numerical_features + ['price']:
    cars[feature].mask(cars[feature] < 0, inplace=True)

categorical_features = ['make', 'fuel', 'gear']

for feature in categorical_features:
    print(
        f"Categorical feature: '{feature}'\n"
        f"No. of unique elemnts: {cars[feature].nunique()}\n"
        f"Unique values:\n{cars[feature].unique()}"
    )
    print('=' * 40)

replace_mapping = {'Electric': ['Electric/Gasoline', 'Electric/Diesel'],
                   'Others': ['-/- (Fuel)', 'CNG', 'LPG', 'Hydrogen', 'Ethanol']}

for value, to_replace in replace_mapping.items():
    cars['fuel'].replace(to_replace, value, inplace=True)

def countplot_missing_data(
    df: pd.DataFrame,
    ax: mpl.axes.Axes,
    **kwargs
) -> None:
    """
    Utility function to plot the counts of missing values in each
    column of the given dataframe `df` using bars.
    """
    missings = df.isna().sum(axis=0).rename('count')
    missings = missings.rename_axis(index='feature').reset_index()
    missings.sort_values('count', ascending=False, inplace=True)

    x = range(missings.shape[0])
    bars = ax.bar(x, missings['count'], **kwargs)

    ax.bar_label(bars, fmt='%d', label_type='edge', padding=3)
    ax.set(xticks=x, xticklabels=missings['feature'])

all_features = {
    'numerical': ['mileage', 'hp', 'age'],
    'categorical': ['make', 'fuel', 'gear']
}

colors = ['teal', 'skyblue']

fig, axes = plt.subplots(1, 2, figsize=(12, 6))
for (item, ax, color) in zip(all_features.items(), axes, colors):
    features_type, features_list = item
    countplot_missing_data(cars[features_list], ax, color=color, width=0.35)
    ax.set(
        title=f'{features_type.capitalize()} Features',
        ylabel='Missing Values, Count [#]'
    )

agg_data = cars.groupby(['make', 'model']).aggregate(
    {'hp': np.nanmedian, 'gear': pd.Series.mode}
)

agg_data.sample(frac=1, random_state=912).head(n=20)

def imputation_transform(
    df: pd.DataFrame,
    label: str,
    group_by: Union[str, List[str]],
    strategy: str
) -> pd.Series:
    """
    Imputation transformer for replacing missing values.

    Parameters
    ----------
    df:
        Main data set.
    label:
        Column for which the missing values are imputed. After groupby,
        the data is aggregated over this column.
    group_by:
        Column, or list of columns used to group data `df`.
    strategy : str, {'mean', 'median', 'most_frequent'}
        The imputation strategy.
        If 'mean' or 'median', then uses equivalent NumPy functions
        ignoring NaNs (`np.nanmean` and `np.nanmedian`, respectively).
        If 'most_frequent', it uses `pd.Series.mode` method.

    Returns
    -------
    Values that can be used to update `df` to replace its missing
    values for the column `label`.
    """
    strategy_maps = {
        'mean': np.nanmean,
        'median': np.nanmedian,
        'most_frequent': pd.Series.mode
    }

    func = strategy_maps[strategy]
    agg_data = df.groupby(group_by).aggregate({label: func})

    def imputer(row):
        """
        Imputation function to apply to each row of the `df` dataframe
        to replace missing values for the feature `agg_over` using the
        imputation strategy `strategy`.
        """
        try:
            x = agg_data.loc[tuple(row[group_by]), label]
            if isinstance(x, np.ndarray):
                fill_value = x[0] if x.size != 0 else np.nan
            else:
                fill_value = x
        except KeyError:
            fill_value = np.nan

        return fill_value

    nan_idx = df[label].isna()
    return df.loc[nan_idx].apply(lambda row: imputer(row), axis=1)

hp_imp = imputation_transform(
    df=cars,
    label='hp',
    strategy='median',
    group_by=['make', 'model']
)

cars['hp'].update(hp_imp)

gear_imp = imputation_transform(
    df=cars,
    label='gear',
    strategy='most_frequent',
    group_by=['make', 'model']
)

cars['gear'].update(gear_imp)

all_features = {
    'numerical': ['mileage', 'hp', 'age'],
    'categorical': ['make', 'fuel', 'gear']
}

colors = ['teal', 'skyblue']

fig, axes = plt.subplots(1, 2, figsize=(12, 6))
for (item, ax, color) in zip(all_features.items(), axes, colors):
    features_type, features_list = item
    countplot_missing_data(cars[features_list], ax, color=color, width=0.35)
    ax.set(
        title=f'{features_type.capitalize()} Features',
        ylabel='Missing Values, Count [#]'
    )

cars.dropna(axis=0, subset=['hp', 'gear'], inplace=True)
cars.info()

def plot_treemap(
    df: pd.DataFrame,
    size_key: str,
    color_key: str,
    ax: mpl.axes.Axes,
    cmap: mpl.colors.Colormap = mpl.cm.Blues,
    font_size: int = 8,
    title_prefix: Optional[str] = None,
    **kwargs
) -> None:
    """
    Utility function to plot treemaps for categorical features using squarify.

    Parameters
    ----------
    size_key:
        Column name whose numeric values specify sizes of rectangles.
    color_key:
        Column name whose numeric values specify colors of rectangles.
    ax:
        Matplotlib Axes instance.
    cmap:
        Matplotlib colormap.
    font_size:
        Set font size of the label text.
    title_prefix:
        Text to prepend to the figure title.
    **kwargs:
        Additional keyword arguments passed to matplotlib.Axes.bar by
        squarify (e.g. `edgecolor`, `linewidth` etc).
    """
    df_sorted = df.sort_values(by=size_key, ascending=False)
    sizes = df_sorted[size_key]

    norm = mpl.colors.Normalize(vmin=df[color_key].min(), vmax=df[color_key].max())

    colors = [cmap(norm(x)) for x in df_sorted[color_key]]

    labels = [
        f"{entry[0]}\nCount: {entry[1]}\nAvg. Pr.: {entry[2]:.0f}"
        for entry in df_sorted.values
    ]

    squarify.plot(sizes=sizes, color=colors, label=labels, ax=ax, **kwargs)

    ax.axis('off')
    for text_obj in ax.texts:
        text_obj.set_fontsize(font_size)

    fig = ax.get_figure()
    cbar = fig.colorbar(mpl.cm.ScalarMappable(norm, cmap), ax=ax)
    cbar.set_label(color_key)

    title = f"Size: {size_key}, Color: {color_key}"
    if title_prefix:
        title = title_prefix + '\n' + title
    ax.set_title(title)

def get_category_meanprice(
    df_in: pd.DataFrame,
    categ_key: str
) -> pd.DataFrame:
    """
    Utility function to get mean price for a given categorical feature.

    Parameters
    ----------
    df_in:
        Input dataframe.
    categ_key:
        Categorical feature name (column name).

    Returns
    -------
    df_out:
        Output dataframe with column names `['count', 'mean_price']`.
    """
    df_out = pd.concat(
        objs=[
            df_in[categ_key].value_counts().rename('count'),
            df_in.groupby(categ_key)['price'].mean().rename('mean_price')
          ],
        axis=1
    )

    df_out.reset_index(drop=False, inplace=True)
    df_out.rename(columns={'index': categ_key}, inplace=True)

    return df_out

print(
    f"Number of unique car-makes: {cars['make'].nunique()}",
    f"Unique car makes listed in the data set:",
    sep='\n'
)

cars['make'].unique()

# Count up cars by `make` feature & average `price`s for each `make` type
makes = get_category_meanprice(cars, 'make')

# Select data w/ a minimum number of items
min_nitems = 200
indexer = makes['count'] >= min_nitems

treemap_kwargs = dict(edgecolor='gray', linewidth=1, font_size=10, alpha=0.8)

fig, ax = plt.subplots(figsize=(16, 9))
plot_treemap(
    makes.loc[indexer, :],
    'count',
    'mean_price',
    ax,
    cmap=mpl.cm.PuBu,
    title_prefix=f'Car Makes (# >= {min_nitems})',
    **treemap_kwargs
)

# Count up cars by `fuel` feature & average `price`s for each `fuel` type
fuels = get_category_meanprice(cars, 'fuel')

fig, ax = plt.subplots(figsize=(12, 7))
plot_treemap(
    fuels,
    'count',
    'mean_price',
    ax,
    cmap=mpl.cm.BuPu,
    title_prefix='Fuel Types',
    **treemap_kwargs
)

# Count up cars by `gear` feature & average `price`s for each `gear` type
gears = get_category_meanprice(cars, 'gear')

fig, ax = plt.subplots(figsize=(12, 7))
plot_treemap(
    gears,
    'count',
    'mean_price',
    ax,
    cmap=mpl.cm.Reds,
    title_prefix='Gear Types',
    **treemap_kwargs
)

min_nitems = 200
s = cars['make'].value_counts() > min_nitems
indexer = cars['make'].isin(s[s].index)
order = cars[indexer].groupby('make')['price'].median().sort_values(ascending=False).index

fig = plt.figure(figsize=(12, 7), tight_layout=True)
gs = mpl.gridspec.GridSpec(2, 2)
boxplot_kwargs = dict(orient='h', showfliers=False, linewidth=1)

sns.boxplot(
    y=cars.loc[indexer, 'make'],
    x=cars.loc[indexer, 'price'] / 1e+3,
    ax=fig.add_subplot(gs[:, 0]),
    order=order,
    palette='pink',
    **boxplot_kwargs
)

for ifeature, feature in enumerate(['fuel', 'gear']):
    order = cars.groupby(feature)['price'].median().sort_values(ascending=False).index
    sns.boxplot(
        y=cars[feature],
        x=cars['price'] / 1e+3,
        ax=fig.add_subplot(gs[ifeature, 1]),
        order=order,
        **boxplot_kwargs
    )

for ax in fig.get_axes():
    ax.set_xlabel(r'price ($\times$ 1000)')

fig, axes = plt.subplots(1, 3, figsize=(14, 4.5))
fig.tight_layout(w_pad=2.5)

for (feature, ax) in zip(all_features['numerical'], axes):
    ax.scatter(cars[feature], cars['price'] / 1e+3, edgecolor='w', alpha=0.8)
    ax.set(xlabel=feature, ylabel=r'price ($\times$ 1000)')

axes[0].set(ylabel=r'price ($\times$ 1000)');

fig, axes = plt.subplots(1, 3, figsize=(14, 4.5))
fig.tight_layout(w_pad=2.5)

for (feature, ax) in zip(all_features['numerical'], axes):
    ax.hist(cars[feature], bins=25)
    ax.set(xlabel=feature, ylabel='Count [#]')

corr = cars.corr()
new_order = ['price'] + numerical_features
corr = corr.reindex(index=new_order, columns=new_order)

mask = np.zeros_like(corr, dtype=bool)
mask[np.triu_indices_from(mask)] = True

fig, ax = plt.subplots(figsize=(6, 5))
ax.set_facecolor('none')
sns.heatmap(corr, ax=ax, mask=mask, annot=True, cmap='BrBG', vmin=-1);

prices = cars['price'] / 1e+3

bin_width = 5   # in €1000
nbins = int(round((prices.max() - prices.min()) / bin_width)) + 1
bins = np.linspace(prices.min(), prices.max(), nbins)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
ax1.hist(prices, bins=bins)
ax1.set(xlabel=r'Price ($\times$ 1000)', ylabel='Count [#]')
_ = probplot(prices, plot=ax2)

cars.to_csv(
    '../data/autoscout24_germany_dataset_cleaned.csv',
    header=True,
    index=False
)