#!/usr/bin/env python3

import pandas as pd
import numpy as np
from matplotlib import pyplot as plt

from sklearn import svm
from sklearn.model_selection import cross_val_score


def readData(filename, types):
    """(str, (str, str))->(pd.DataFrame, pd.DataFrame, pd.DataFrame)
    read .xlsx file, and return data"""
    df = pd.read_excel(filename)
    # set type1 to 0, set type2 to 1
    df.loc[:, "TYPE"] = df.loc[:, "TYPE"].apply(
        lambda x: 0 if x == types[0] else 1)
    dataTrain, trainTarget = df.iloc[:, 2:4], df.iloc[:, 1]

    return dataTrain, trainTarget, df


def plotData(df, modelDf=None, output=None, types=None):
    """(pd.DataFrame, pd.DataFrame, None or pd.DataFrame)->None
    plot type1, type2 and prediction of trained model"""
    # get title
    blueDf, redDf = df.loc[df.loc[:, "TYPE"]
                        == 0], df.loc[df.loc[:, "TYPE"] == 1]
    x, y = blueDf.columns[2:4]

    # plot predictions
    if modelDf is not None:
        # type 1 prediction
        plt.scatter(modelDf[modelDf.loc[:, "c"] <= .5].loc[:, x],
                    modelDf.loc[modelDf.loc[:, "c"] <= .5].loc[:, y], c="purple")
        # type 2 prediction background
        plt.scatter(modelDf[modelDf.loc[:, "c"] > .5].loc[:, x],
                    modelDf.loc[modelDf.loc[:, "c"] > .5].loc[:, y], c="yellow")
    # type1 and type2
    plt.scatter(x=blueDf.loc[:, x], y=blueDf.loc[:, y], c="b")
    plt.scatter(x=redDf.loc[:, x], y=redDf.loc[:, y], c="r")
    plt.xlabel(x)
    plt.ylabel(y)
    plt.title(" and ".join(types))
    if output is None:
        plt.show()
    else:
        plt.savefig(output)
        plt.clf()


def _svmMethod(dataTrain, trainTarget, gamma, kernel):
    """(pd.DataFrame, pd.DataFrame, float, str)->(float, sklearn.svm.SVC)
    use ten fold cross validation to evaluate model with different parameters
    """
    kernelSvm = svm.SVC(gamma=gamma, kernel=kernel)
    # ten fold cross validation
    scores = cross_val_score(kernelSvm, dataTrain,
                             trainTarget, cv=10, scoring="accuracy")
    return scores.mean(), kernelSvm


def svmMethod(dataTrain, trainTarget):
    """(pd.DataFrame, pd.DataFrame)->sklearn.svm.SVC
    first use some methods to choose a model with best parameters,
    then return the trained model
    """
    bestModel, bestScore = None, 0
    # choose parameters in (kernel, gamma)
    for kernel in ("rbf", "linear", "poly", "sigmoid"):
        for gamma in ("auto", "scale", *np.arange(0.1, 1.0, 0.1)):
            score, model = _svmMethod(
                dataTrain, trainTarget, gamma, kernel)
            if score > bestScore:
                bestModel, bestScore= model, score
    # fit data
    trainedModel = bestModel.fit(dataTrain, trainTarget)
    return trainedModel


def generateBackground(df):
    """(pd.DataFrame)->pd.DataFrame
    generate data used as prediction background
    """
    x, y = df.columns[2:4]
    minX, maxX = df.iloc[:, 2].min(), df.iloc[:, 2].max()
    minY, maxY = df.iloc[:, 3].min(), df.iloc[:, 3].max()
    # here are about 90000 points, it will slow down the program
    xRan, yRan = np.arange(minX, maxX*1.001, maxX /
                           300), np.arange(minY, maxY*1.001, maxY/300)
    lenX, lenY = len(xRan), len(yRan)
    xRan = xRan.repeat(lenY)
    yRan = np.tile(yRan, lenX)
    
    return pd.DataFrame([xRan, yRan], index=[x, y]).transpose()


def classifier(filename, types, output):
    """(str, (str, str), str)->(sklearn.svm.SVC, float)
    with the given data (filename), classify the data into types,
    and save the result to a picture named output
    """
    dataTrain, trainTarget, df = readData(filename, types)
    # svm method
    bestModel = svmMethod(dataTrain, trainTarget)
    # background prediction
    background = generateBackground(df)
    backgroundResult = bestModel.predict(background)
    modelDf = background.join(pd.DataFrame(
        backgroundResult.reshape(1, -1), index=["c"]).transpose())

    plotData(df, modelDf, output=output, types=types)

    return bestModel, bestModel.score(dataTrain, trainTarget)