import pandas as pd
from sklearn.model_selection import train_test_split

crime = pd.read_csv('crime_conv.csv')

# Większość danych o strzelaninach jest pusta więc zakładam, że są to inceydenty bez strzelanin
crime["SHOOTING"].fillna("N", inplace=True)

# Kolumna location powtarza wartości z Lat i Long
crime.drop(columns=["Location"], inplace=True)

# Usuwam błędne/brakujące wartości współrzędnych
crime = crime[(crime["Lat"] > 35) & (crime["Long"] < -65)]

# Lowercase na polach tekstowych
crime["OFFENSE_CODE_GROUP"] = crime["OFFENSE_CODE_GROUP"].str.lower()
crime["OFFENSE_DESCRIPTION"] = crime["OFFENSE_DESCRIPTION"].str.lower()
crime["DAY_OF_WEEK"] = crime["DAY_OF_WEEK"].str.lower()
crime["UCR_PART"] = crime["UCR_PART"].str.lower()
crime["STREET"] = crime["STREET"].str.lower()

# Usuwam pozostałe wiersze zawierające nulle
crime.dropna()

# Zbiór jest całkiem duży - 300k wierszy po oczyszczeniu, więc wybieram podział 8:1:1, czyli w przybliżeniu 30k na dev i test
crime_train, crime_test = train_test_split(crime, test_size=60000, random_state=1)
crime_test, crime_dev = train_test_split(crime_test, test_size=30000, random_state=1)

crime_test.to_csv("crime_test.csv", encoding="utf-8", index=False)
crime_dev.to_csv("crime_dev.csv", encoding="utf-8", index=False)
crime_train.to_csv("crime_train.csv", encoding="utf-8", index=False)