# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import csv
import matplotlib.pyplot as plt
import pandas as pd
import re

# +
# '../data/cleaned/clean_middlesex.csv'
# '../data/cleaned/clean_northwestern.csv'
# '../data/cleaned/clean_suffolk.csv'

middlesex_df = pd.read_csv('../data/cleaned/clean_middlesex.csv', low_memory=False)
northwestern_df = pd.read_csv('../data/cleaned/clean_northwestern.csv', low_memory=False)
suffolk_df = pd.read_csv('../data/cleaned/clean_suffolk.csv', low_memory=False)

# Csv containing Chapter-Section titles
info_df = pd.read_csv('Part-Title-Chapter-Section.csv', low_memory=False)


# +
# Inputs:
#   - df: dataframe to be cleaned;
#         mostly removes Paragraphs from Section column
#         some Chapter entries end in commas, so commas are removed
# Outputs:
#   - df: cleaned dataframe
def clean_data(df):
    df = df.fillna('NaN')
    df['Chapter'] = [re.split('[,]', row, maxsplit=1)[0] for row in df['Chapter']]
    df['Section'] = [re.split('[(Ã \-]', row, maxsplit=1)[0] for row in df['Section']]
    return df

# Inputs:
#   - df: dataframe to be analyzed and displayed
#   - DA_name: district name; added to title of plots
#   - num_bars: number of bars to be displayed; first 10 bars are always red
# Outputs:
#   - data_df: dataframe containing Chapter-Section pairs, # of Cases, Percent Frequency, Cumulative Frequency
def plot_frequency(df, DA_name='DA', num_bars=25):
    # Generate column for Chapter-Section pairs and find # of entries for each pair
    # Calculate Percent Frequency and Cumulative Frequency
    df['ChapterSection'] = df['Chapter'] + '-' + df['Section']
    df_size = df['ChapterSection'].size
    percent_frequency = df['ChapterSection'].value_counts().sort_values(ascending=False) / df_size * 100
    cumulative_frequency = percent_frequency.cumsum()
    
    # Returns dataframe w/ columns for Chapter-Section, Absolute Number of Cases, Percent Frequency, & Cumulative Frequency
    data_df = pd.concat([df['ChapterSection'].value_counts().rename('Absolute # of Cases'),
                         percent_frequency.rename('Percent Frequency (%)'),
                         cumulative_frequency.rename('Cumulative Frequency (%)')], axis=1)
    data_df.reset_index(level=data_df.index.names, inplace=True)
    data_df.rename(columns={'index': 'Chapter-Section'}, inplace=True)
    
    # Add Chapter and Section data to use as keys to pair w/ Section_Heading
    data_df['Chapter'] = [cs.split('-')[0] for cs in data_df['Chapter-Section']]
    data_df['Section'] = [cs.split('-')[1] for cs in data_df['Chapter-Section']]
    data_df = data_df.merge(info_df[['Chapter', 'Section', 'Section_Heading']], how='left', on=['Chapter', 'Section'])
    data_df.fillna('', inplace=True)
    
    # Plot!!!
    fig, axes = plt.subplots(2, 1, figsize=(15, 15))
    colors = ['#d62728']*10 + ['#1f77b4']*(num_bars-10)*(num_bars>10)
    yticklabels = [f'{cs}: {sh[:27]}..' if len(sh) > 30 else f'{cs}: {sh}' for cs, sh in zip(data_df['Chapter-Section'].iloc[:num_bars], data_df['Section_Heading'].iloc[:num_bars])]

    percent_frequency.iloc[:num_bars].plot(kind='barh', ax=axes[0], color=colors)
    axes[0].set_title(DA_name + ' Chapter-Section Percent Frequency', fontsize=15)
    axes[0].set_yticklabels(yticklabels)
    axes[0].set_xlabel('Percent Frequency (%)')
    axes[0].set_ylabel('Chapter-Section')
    axes[0].invert_yaxis()
    axes[0].legend(['Top 10'], frameon=False)
    axes[0].spines['top'].set_visible(False)
    axes[0].spines['right'].set_visible(False)
    # axes[0].spines['bottom'].set_visible(False)
    axes[0].spines['left'].set_visible(False)
    axes[0].yaxis.set_tick_params(bottom=False)

    cumulative_frequency.iloc[:num_bars].plot(kind='barh', ax=axes[1], color=colors)
    axes[1].set_title(DA_name + ' Chapter-Section Cummulative Frequency', fontsize=15)
    axes[1].set_yticklabels(yticklabels)
    axes[1].set_xlabel('Cummulative Frequency (%)')
    axes[1].set_ylabel('Chapter-Section')
    axes[1].invert_yaxis()
    axes[1].legend(['Top 10'], frameon=False)
    axes[1].spines['top'].set_visible(False)
    axes[1].spines['right'].set_visible(False)
    # axes[1].spines['bottom'].set_visible(False)
    axes[1].spines['left'].set_visible(False)
    axes[1].yaxis.set_tick_params(bottom=False)
    
    plt.tight_layout()
    
    return data_df



# +
# df = northwestern_df.copy()

# # Generate column for Chapter-Section pairs and find # of entries for each pair
# # Calculate Percent Frequency and Cumulative Frequency
# df['ChapterSection'] = df['Chapter'] + '-' + df['Section']
# df_size = df['ChapterSection'].size
# percent_frequency = df['ChapterSection'].value_counts().sort_values(ascending=False) / df_size * 100
# cumulative_frequency = percent_frequency.cumsum()

# # Returns dataframe w/ columns for Chapter-Section, Absolute Number of Cases, Percent Frequency, & Cumulative Frequency
# data_df = pd.concat([df['ChapterSection'].value_counts().rename('Absolute # of Cases'),
#                      percent_frequency.rename('Percent Frequency (%)'),
#                      cumulative_frequency.rename('Cumulative Frequency (%)')], axis=1)
# data_df.reset_index(level=data_df.index.names, inplace=True)
# data_df.rename(columns={'index': 'Chapter-Section'}, inplace=True)

# data_df['Chapter'] = [cs.split('-')[0] for cs in data_df['Chapter-Section']]
# data_df['Section'] = [cs.split('-')[1] for cs in data_df['Chapter-Section']]
# data_df = data_df.merge(info_df[['Chapter', 'Section', 'Section_Heading']], how='left', on=['Chapter', 'Section'])
# data_df.fillna('NaN', inplace=True)

# num_bars = 25

# fig, axes = plt.subplots(2, 1, figsize=(15, 15))
# colors = ['#d62728'] * 10 + ['#1f77b4'] * (num_bars - 10) * (num_bars > 10)
# yticklabels = [f'{cs}: {sh[:27]}..' if len(sh) > 30 else f'{cs}: {sh}' for cs, sh in zip(data_df['Chapter-Section'].iloc[:num_bars], data_df['Section_Heading'].iloc[:num_bars])]

# percent_frequency.iloc[:num_bars].plot(kind='barh', ax=axes[0], color=colors)
# axes[0].set_title('Northwestern Chapter-Section Percent Frequency', fontsize=15)
# axes[0].set_yticklabels(yticklabels)
# axes[0].set_xlabel('Percent Frequency (%)')
# axes[0].set_ylabel('Chapter-Section')
# axes[0].invert_yaxis()
# axes[0].legend(['Top 10'], frameon=False)
# axes[0].spines['top'].set_visible(False)
# axes[0].spines['right'].set_visible(False)
# # axes[0].spines['bottom'].set_visible(False)
# axes[0].spines['left'].set_visible(False)
# axes[0].yaxis.set_tick_params(bottom=False)

# cumulative_frequency.iloc[:num_bars].plot(kind='barh', ax=axes[1], color=colors)
# axes[1].set_title('Northwestern Chapter-Section Cummulative Frequency', fontsize=15)
# axes[1].set_yticklabels(yticklabels)
# axes[1].set_xlabel('Cummulative Frequency (%)')
# axes[1].set_ylabel('Chapter-Section')
# axes[1].invert_yaxis()
# axes[1].legend(['Top 10'], frameon=False)
# axes[1].spines['top'].set_visible(False)
# axes[1].spines['right'].set_visible(False)
# # axes[1].spines['bottom'].set_visible(False)
# axes[1].spines['left'].set_visible(False)
# axes[1].yaxis.set_tick_params(bottom=False)

# plt.tight_layout()

# display(data_df.iloc[:10])
# -

# # Middlesex

# +
middlesex_cleaned_df = clean_data(middlesex_df)

middlesex_freq_df = plot_frequency(middlesex_cleaned_df, DA_name='Middlesex', num_bars=25)
display(middlesex_freq_df.iloc[:10])
# -

# # Northwestern

# +
northwestern_cleaned_df = clean_data(northwestern_df)

northwestern_freq_df = plot_frequency(northwestern_cleaned_df, DA_name='Northwestern', num_bars=25)
display(northwestern_freq_df.iloc[:10])
# -

# # Suffolk

# +
suffolk_cleaned_df = clean_data(suffolk_df)

suffolk_freq_df = plot_frequency(suffolk_cleaned_df, DA_name='Suffolk', num_bars=25)
display(suffolk_freq_df.iloc[:10])
# -



StarcoderdataJupyterScriptsDedupFiltered