from itertools import product
import numpy as np
import pandas as pd
import csv
import id3test
import pprint
import sys
import pandas
from sklearn import tree
import pydotplus
from sklearn.tree import DecisionTreeClassifier
import matplotlib.pyplot as plt
import matplotlib.image as pltimg

field_states = ['toPlow', 'toWater', 'toSeed', 'toFertilize', 'toCut']
header = ['Equipped', 'Not Equipped']
hitch = ['Tillage Unit', 'Crop Trailer']
tillage_unit = ['Seeds', 'Water', 'Fertilizer', 'Nothing']
engine_working = ['Yes', 'No']
in_base = ['Yes', 'No']
fuel = ['Enough', 'Not_Enough']

# output = list(product(field_states, header, hitch, tillage_unit, engine_working, in_base, fuel))
output = list(product(field_states, header, hitch, tillage_unit, in_base))

dict = []
# for x in range(len(output)):
#     # dict.append({'field': output[x][0], 'header' : output[x][1], 'hitch' : output[x][2],
#     #              'tillage_unit' : output[x][3], 'engine_working' : output[x][4],
#     #              'in_base' : output[x][5], 'fuel': output[x][6]})
#     print(x+1, output[x])

decisions = []
curr_decision = "Make Action"

for x in range(len(output)):
    aField = output[x][0]
    aHeader = output[x][1]
    aHitch = output[x][2]
    aTillage_unit = output[x][3]
    aIn_base = output[x][4]

    while True:
        curr_decision = "Make Action"

        if aField == 'toCut':
            if not aHeader == 'Equipped':
                if aIn_base == 'Yes':
                    curr_decision = 'Change Header'
                    break
                else:
                    curr_decision = 'Go To Base'
                    break

            if not aHitch == 'Crop Trailer':
                if aIn_base == 'Yes':
                    curr_decision = 'Change Hitch'
                    break
                else:
                    curr_decision = 'Go To Base'
                    break

        if aField == 'toPlow':
            if aHeader == 'Equipped':
                if aIn_base == "Yes":
                    curr_decision = 'Change Header'
                    break
                else:
                    curr_decision = 'Go To Base'
                    break

            if not aHitch == "Tillage Unit":
                if aIn_base == "Yes":
                    curr_decision = "Change Hitch"
                    break
                else:
                    curr_decision = "Go To Base"
                    break
            if not aTillage_unit == "Nothing":
                if aIn_base == "Yes":
                    curr_decision = "Change Load"
                    break
                else:
                    curr_decision = "Go To Base"
                    break

        if aField == 'toSeed':
            if aHeader == 'Equipped':
                if aIn_base == "Yes":
                    curr_decision = 'Change Header'
                    break
                else:
                    curr_decision = 'Go To Base'
                    break

            if not aHitch == "Tillage Unit":
                if aIn_base == "Yes":
                    curr_decision = "Change Hitch"
                    break
                else:
                    curr_decision = "Go To Base"
                    break
            if not aTillage_unit == "Seeds":
                if aIn_base == "Yes":
                    curr_decision = "Change Load"
                    break
                else:
                    curr_decision = "Go To Base"
                    break

        if aField == 'toWater':
            if aHeader == 'Equipped':
                if aIn_base == "Yes":
                    curr_decision = 'Change Header'
                    break
                else:
                    curr_decision = 'Go To Base'
                    break

            if not aHitch == "Tillage Unit":
                if aIn_base == "Yes":
                    curr_decision = "Change Hitch"
                    break
                else:
                    curr_decision = "Go To Base"
                    break
            if not aTillage_unit == "Water":
                if aIn_base == "Yes":
                    curr_decision = "Change Load"
                    break
                else:
                    curr_decision = "Go To Base"
                    break

        if aField == 'toFertilize':
            if aHeader == 'Equipped':
                if aIn_base == "Yes":
                    curr_decision = 'Change Header'
                    break
                else:
                    curr_decision = 'Go To Base'
                    break

            if not aHitch == "Tillage Unit":
                if aIn_base == "Yes":
                    curr_decision = "Change Hitch"
                    break
                else:
                    curr_decision = "Go To Base"
                    break
            if not aTillage_unit == "Fertilizer":
                if aIn_base == "Yes":
                    curr_decision = "Change Load"
                    break
                else:
                    curr_decision = "Go To Base"
                    break
        if aIn_base == 'Yes':
            curr_decision = 'Go To Field'
        break
    dict.append({'Field': aField, 'Header': aHeader, 'Hitch': aHitch, 'Tillage_Unit': aTillage_unit,
                 'In_Base': aIn_base, 'Decision': curr_decision})

print(dict)

fields = ['Field', 'Header', 'Hitch', 'Tillage_Unit', 'In_Base', 'Decision']

filename = "treedata\\data2.csv"

with open(filename, 'w') as csvfile:
    writer = csv.DictWriter(csvfile, fieldnames=fields)
    writer.writeheader()
    writer.writerows(dict)

df = pandas.read_csv("treedata\\data2.csv")
print(df)

# Map text values to number values
d = {'toPlow': 0, 'toWater': 1, 'toSeed': 2, 'toFertilize': 3, 'toCut': 4}
df['Field'] = df['Field'].map(d)

d = {'Not Equipped': 0, 'Equipped': 1}
df['Header'] = df['Header'].map(d)

d = {'Tillage Unit': 0, 'Crop Trailer': 1}
df['Hitch'] = df['Hitch'].map(d)

d = {'Nothing': 0, 'Seeds': 1, 'Water': 2, 'Fertilizer': 3}
df['Tillage_Unit'] = df['Tillage_Unit'].map(d)

d = {'No': 0, 'Yes': 1}
df['In_Base'] = df['In_Base'].map(d)

d = {'Make Action': 0, 'Change Header': 1, 'Go To Base': 2, 'Change Hitch': 3, 'Change Load': 4, 'Go To Field': 5}
df['Decision'] = df['Decision'].map(d)

# Separate the feature columns from targert columns

features = ['Field', 'Header', 'Hitch', 'Tillage_Unit', 'In_Base']

X = df[features]
y = df['Decision']

# FIELD 'toPlow' : 0, 'toWater' : 1, 'toSeed' : 2, 'toFertilize' : 3, 'toCut' : 4
# HEADER 'Not Equipped' : 0, 'Equipped' : 1
# HITCH 'Tillage Unit' : 0, 'Crop Trailer' : 1
# TILLAGE 'Nothing' : 0, 'Seeds' : 1, 'Water': 2, 'Fertilizer': 3
# IN BASE 'No' : 0, 'Yes' : 1


dtree = DecisionTreeClassifier()
dtree = dtree.fit(X, y)
data = tree.export_graphviz(dtree, out_file=None, feature_names=features)
graph = pydotplus.graph_from_dot_data(data)
graph.write_png('mydecisiontree.png')

img = pltimg.imread('mydecisiontree.png')
imgplot = plt.imshow(img)
plt.show()

print(dtree.predict([[0, 1, 1, 3, 0]]))