import time
from sys import argv

from matplotlib.pyplot import *
from pandas import DataFrame
from seaborn import clustermap, set_style, set_theme, pairplot
from sklearn.metrics import silhouette_score, calinski_harabasz_score
from sklearn.cluster import KMeans, Birch, SpectralClustering, MeanShift, DBSCAN

from preprocessing import parse_data, filter_dataframe


def choose_model(model, cluster_number):
    if model == "kmeans":
        return KMeans(n_clusters=cluster_number, random_state=42)
    elif model == "birch":
        return Birch(n_clusters=cluster_number)
    elif model == "spectral":
        return SpectralClustering()
    elif model == "meanshift":
        return MeanShift()
    elif model == "dbscan":
        return DBSCAN()


def predict_data(data, model, cluster_number, results):
    model_name = model
    model = choose_model(model=model, cluster_number=cluster_number)
    start_time = time.time()
    prediction = model.fit_predict(data)
    execution_time = time.time() - start_time
    cluster_number = len(set(prediction))
    calinski = calinski_harabasz_score(X=data, labels=prediction)
    silhouette = silhouette_score(
        X=data,
        labels=prediction,
        metric="euclidean",
        random_state=42,
    )
    populated_results = populate_results(
        df=results,
        model=model_name,
        prediction=prediction,
        clusters=cluster_number,
        calinski=calinski,
        silhouette=silhouette,
        time=execution_time,
    )
    return populated_results


def plot_heatmap(results):
    fig = figure(figsize=(20, 10))
    results.reset_index()
    matrix = results["prediction"]
    print(matrix)
    clustermap(
        data=matrix,
        cmap="mako",
        metric="euclidean",
        annot=True,
    )
    fig_title = "Heatmap"
    title(fig_title)
    show()
    fig.savefig(f"docs/assets/{fig_title.replace(' ', '_').lower()}.png")


def plot_scatter_plot(results):
    fig = figure(figsize=(20, 10))
    matrix = results.filter(items=["input", "prediction"])
    pairplot(
        data=results,
        vars=matrix,
        hue="prediction",
        palette="Paired",
        diag_kind="hist",
    )
    fig_title = "Scatter plot"
    title(fig_title)
    show()
    fig.savefig(f"docs/assets/{fig_title.replace(' ', '_').lower()}.png")


def show_results(results):
    set_theme()
    set_style("white")
    plot_heatmap(results=results)
    plot_scatter_plot(results=results)
    print(results)


def create_result_dataframes():
    results = DataFrame(
        columns=[
            "model",
            "clusters",
            "prediction",
            "silhouette",
            "calinski-harabasz",
            "time",
        ]
    )
    return results, results


def populate_results(df, model, clusters, prediction, calinski, silhouette, time):
    columns = [
        "model",
        "clusters",
        "prediction",
        "silhouette",
        "calinski-harabasz",
        "time",
    ]
    values = [model, clusters, prediction, silhouette, calinski, time]
    dictionary = dict(zip(columns, values))
    populated_df = df.append(dictionary, ignore_index=True)
    return populated_df


def construct_case(df, choice):
    cases = {
        "case1": df.loc[(df["LUMINOSIDAD"].str.contains("NOCHE"))],
        "case2": df.loc[
            (df["ISLA"].str.contains("NO_ES_ISLA") == False)
            & (df["FACTORES_ATMOSFERICOS"].str.contains("LLUVIA|LLOVIZNA"))
        ],
        "case3": df.loc[(df["HORA"] > 19) & (df["TIPO_VIA"] == "AUTOPISTA")],
        "case4": df.loc[
            (df["COMUNIDAD_AUTONOMA"] == "Andalucía")
            & (df["LUMINOSIDAD"].str.contains("SIN ILUMINACIÓN"))
        ],
        "case5": df.loc[
            (df["DIASEMANA"] == 7)
            & (df["COMUNIDAD_AUTONOMA"] == "Madrid, Comunidad de")
        ],
    }
    return cases[choice]


def usage():
    print("Usage: " + argv[0] + "<preprocessing action> <case> <number of clusters>")
    print("preprocessing actions:")
    print("fill: fills the na values with the mean")
    print("drop: drops the na values")
    print("cases: choice of case study")
    print(
        "number of clusters: number of clusters for the algorithms that use a fixed number"
    )
    exit()


def main():
    models = ["kmeans", "birch", "spectral", "meanshift", "dbscan"]
    if len(argv) != 4:
        usage()
    case, cluster_number = argv[2], int(argv[3])
    data = parse_data(source="data/accidentes_2013.csv", action=str(argv[1]))
    individual_result, complete_results = create_result_dataframes()
    case_data = construct_case(df=data, choice=case)
    filtered_data = filter_dataframe(df=case_data)
    for model in models:
        model_results = predict_data(
            data=filtered_data,
            model=model,
            results=individual_result,
            cluster_number=cluster_number,
        )
        complete_results = complete_results.append(
            individual_result.append(model_results)
        )
    complete_results.set_index("model")
    show_results(results=complete_results)


if __name__ == "__main__":
    main()
