import pandas as pd
from tkinter import filedialog
from tkinter import Tk



def wecc_sort():

    # This prompts you to pick a file (a csv) that you want to filter
    Tk().withdraw()
    pickled_ahead = filedialog.askopenfilename()

    # This reads in the file and tells you that it is loading
    print('Loading file...')
    df = pd.read_csv(pickled_ahead)

    # This filters the dataframe by the Buyer Ownership Type. Right now it's set to "not equal to private"
    df = df[df['Buyer Ownership Type'] != 'Private']
    # df = df[df['Buyer Ownership Type'] == 'Private'] # Whereas this one for example filters for only equal to private

    # This filters by the various parameters Rate Unit Code as $/KW OR $/MW
    df = df[(df['Rate Unit Code'] == '$/KW') | (df['Rate Unit Code'] == '$/MW')]

    # This resets the index because we dropped a bunch of data and need to get rid of the empty space they left
    df.reset_index(drop=True, inplace=True)

    # This creates a new dataframe with only the columns we want from the filtered dataframe
    df_save = df[['Buyer Name', 'Seller Name', 'Rate Unit Code', 'Buyer Ownership Type', 'Price']]

    # This saves the new dataframe to a csv. Change the name to what you want. It will output to wherever this script
    # is saved.
    df_save.to_csv('WECC Sort For Zach.csv')

### Run your functions here

wecc_sort()