import pandas as pd

from constants import DATA_SIZE


def load_data():
    data = read_data()
    Y = data['Model']
    X = data[['Year', 'Price', 'Mileage']]

    return X, Y


def fit_model(model, x):
    transformed = model.fit_transform(x)

    return transformed[:, 0], transformed[:, 1]


def create_data():
    data = read_data()
    return data[['Year', 'Price']]


def read_data():
    data = pd.read_csv('true_car_listings.csv')[:DATA_SIZE]

    unique_numbers = list(set(data['Model']))
    data['Model'] = data['Model'].apply(unique_numbers.index)
    return data