import dislib as ds
import seaborn as sns
from dislib.cluster import KMeans
from dislib.data import load_txt_file
import sys
import pandas as pd
from pycompss.api.parameter import *
from pycompss.api.task import task
from pycompss.api.api import compss_wait_on_file
import seaborn as sns
from sklearn.preprocessing import LabelEncoder
import numpy as np


@task(original_dataset=FILE_IN, output_dataset=FILE_OUT)
def preprocess_dataset(original_dataset, output_dataset):
    ori_dataset = pd.read_csv(original_dataset)
    encoded_data = LabelEncoder().fit_transform(np.array(ori_dataset.ocean_proximity))
    encoded_data = encoded_data + 1
    ori_dataset = ori_dataset.drop(columns=["total_bedrooms", "ocean_proximity"])
    full_data = pd.concat([ori_dataset, pd.DataFrame({'ocean_proximity': encoded_data})], axis=1)
    full_data.to_csv(output_dataset)


@task(longitude_latitude_plot=FILE_OUT)
def first_plot(longitude_latitude_plot, x_data):
    seaborn_plot = sns.relplot(
    x="latitude", y="longitude", hue="cluster", data=x_data,
    )
    seaborn_plot.savefig(longitude_latitude_plot)
    return
    ax = seaborn_plot.figure
    ax.save_figure(longitude_latitude_plot)


@task(median_income_plot=FILE_OUT)
def plot_median_income_value(median_income_plot, x_data):
    seaborn_plot = sns.catplot(x="median_income", y="cluster", data=x_data, kind="boxen")
    seaborn_plot.savefig(median_income_plot)
    return
    ax = seaborn_plot.figure
    ax.save_figure(median_income_plot)


def main():
    original_dataset = sys.argv[1]
    output_file = sys.argv[2]
    output_lon_lat_file = sys.argv[3]
    output_med_income = sys.argv[4]
    preprocess_dataset(original_dataset, output_file)
    output_file = compss_wait_on_file(output_file)
    data = load_txt_file(output_file, block_size=(414, 5),
            col_of_index=True, discard_first_row=True)
    kmeans = KMeans(n_clusters=6, random_state=0)
    kmeans.fit(data)
    ori_dataset = pd.read_csv(original_dataset)
    ori_dataset["cluster"] = kmeans.fit_predict(data).collect()
    ori_dataset["cluster"] = ori_dataset["cluster"].astype("category")
    first_plot(output_lon_lat_file, ori_dataset)
    plot_median_income_value(output_med_income, ori_dataset)




if __name__ == '__main__':
    main()

