from ai.FileIO import FileIO
from ai.NeuralNet import NeutalNet
from ai.visualization import drawScatterPlot2, drawFunctionPlot
import numpy;

def doClassification(trainFile, testFile, testHasAns=True):

    trainData = FileIO.readFile2(trainFile)
    testData = FileIO.readFile2(testFile, train=False, testHasAns=testHasAns)

    res_list = [x[-1] for x in trainData]
    unique = numpy.unique(res_list)
    if testHasAns:
        evalData = testData
    else:
        evalData = None

    print('Learning...')
    nn = NeutalNet([2, 10,10, len(unique)], 'classification', 'sigmoid')
    nn.SGD(trainData, 30, 10, 0.5,evalData)
    print('Testing...')
    nnGuesses = []
    for x_t in testData:
        y_t = nn.feedforward(x_t)
        nnGuesses.append((x_t, y_t))

    if not testHasAns: # test does not have ans so we guess output and save to file
        FileIO.saveFile(testData, nn.problemType, testFile)

    drawScatterPlot2(trainData, nnGuesses, nn.errors)


def doRegression(trainFile, testFile, testHasAns):

    trainData = FileIO.readFile2(trainFile)
    res_list = [x[-1] for x in trainData]
    maxValue = float(max(res_list))
    minValue = float(min(res_list))
    print('Learning...')
    nn = NeutalNet([1, 6, 1], 'regression', 'tanh',maxValue, minValue)
    nn.SGD(trainData, 50, 10, 0.1)

    print('Testing...')
    listInput = FileIO.readFile2(testFile, train=False, testHasAns=testHasAns)
    testData = []
    for x_t in listInput:
        y_t = nn.feedforward(x_t)
        testData.append((x_t, y_t))
    drawFunctionPlot(trainData, testData, nn.errors)
    FileIO.saveFile(testData, nn.problemType)

def main():
    trainFileCls = './ai/data.train.csv'
    testFileCls = './ai/data.test.csv'
    doClassification(trainFileCls, testFileCls, testHasAns=False)
    trainFileReg = './ai/data.xsq.train.csv'
    testFileReg = './ai/data.xsq.test.csv'
    #doRegression(trainFileReg, testFileReg, testHasAns=False)


if __name__ == '__main__':
    main()