#YOUR CODE HERE
import numpy as np
import matplotlib.pyplot as plt
from fuzzython.fsets.trapezoid import Trapezoid
from fuzzython.fsets.triangular import Triangular
from fuzzython.variable import Variable
from fuzzython.adjective import Adjective
from fuzzython.ruleblock import RuleBlock
from fuzzython.systems.mamdani import MamdaniSystem
from fuzzython.fsets.fuzzy_set import FuzzySet
from mpl_toolkits.mplot3d import Axes3D

def plot_fuzzyset(ax, fuzzy_set, x, *args, **kwargs):
    y = np.array([fuzzy_set(e) for e in x])
    ax.plot(x, y,  *args, **kwargs)
    ax.set_ylim(-0.1, 1.1)
    ax.legend()

#PRĘDKOŚĆ
predkosc_mala = Trapezoid((8,0),(10,1),(50,1),(60,0))
predkosc_srednia = Trapezoid((50,0),(60,1),(75,1),(90,0))
predkosc_szybka = Trapezoid((75,0),(85,1),(155,1),(160,0))
predkosc_bardzo_szybka = Trapezoid((160,0),(175,1),(200,1),(200.1,0))

A_predkosc_mala = Adjective('Mała prędkość', predkosc_mala)
A_predkosc_srednia = Adjective('Średnia prędkość', predkosc_srednia)
A_predkosc_szybka = Adjective('Szybka prędkość', predkosc_szybka )
A_predkosc_bardzo_szybka = Adjective('Bardzo szybka prękość', predkosc_bardzo_szybka)

predkosc = Variable('predkosc', 'km/h', predkosc_mala, predkosc_srednia, predkosc_szybka, predkosc_bardzo_szybka)


#WIDOCZNOŚĆ
widocznosc_bardzo_slaba = Triangular((0.04,0),(0.05,1),(2,0))
widocznosc_srednia = Triangular((0.05,0),(2,1),(4,0))
widocznosc_dobra = Triangular((2,0),(4,1),(4.1,0))

A_widocznosc_bardzo_slaba = Adjective('Bardzo słaba widoczność', widocznosc_bardzo_slaba)
A_widocznosc_srednia = Adjective('Średnia widoczność', widocznosc_srednia)
A_widocznosc_dobra = Adjective('Dobra widoczność', widocznosc_dobra)

widocznosc = Variable('widocznosc', 'km', widocznosc_bardzo_slaba, widocznosc_srednia, widocznosc_dobra)

#MOŻLIWOŚĆ SPOWODOWANIA WYPADKU
wypadek_bardzo_male = Trapezoid((-0.01,0),(0,1),(0.25,1),(0.3,0))
wypadek_male = Trapezoid((0.2,0),(0.25,1),(0.4,1),(0.45,0))
wypadek_srednie = Trapezoid((0.4,0),(0.45,1),(0.6,1),(0.7,0))
wypadek_duze = Trapezoid((0.65,0),(0.7,1),(1,1),(1.01,0))

A_wypadek_bardzo_male = Adjective('Bardzo mała możliwość spowodowania wypadku', wypadek_bardzo_male)
A_wypadek_male = Adjective('Mała możliwość spowodowania wypadku', wypadek_male)
A_wypadek_srednie = Adjective('Średnia możliwość spowodowania wypadku', wypadek_srednie)
A_wypadek_duze = Adjective('Duża możiwośc spowodowania wypadku', wypadek_duze)

wypadek = Variable('wypadek', '0-1', wypadek_bardzo_male, wypadek_male, wypadek_srednie, wypadek_duze, defuzzification='COG', default=0)

#GRAF
x = np.linspace(10,200,10000)
x1 = np.linspace(0.05,4,1000)
x2 = np.linspace(0,1,1000)
fig, axs = plt.subplots(nrows=3, ncols=1, figsize=(12,8))
((ax1), (ax2), (ax3)) = axs

#WYKRES PRĘDKOŚCI
plot_fuzzyset(ax1, predkosc_mala, x, 'g', label='Mała prędkość')
plot_fuzzyset(ax1, predkosc_srednia, x, 'b', label='Średnia prędkość')
plot_fuzzyset(ax1, predkosc_szybka, x, 'y', label='Szybka prędkość')
plot_fuzzyset(ax1, predkosc_bardzo_szybka, x, 'r', label='Bardzo szybka prędkość')

#WYKRES WIDOCZNOŚCI
plot_fuzzyset(ax2, widocznosc_bardzo_slaba, x1, 'r', label='Bardzo słaba widoczność')
plot_fuzzyset(ax2, widocznosc_srednia, x1, 'b', label='Średnia widoczność')
plot_fuzzyset(ax2, widocznosc_dobra, x1, 'g', label='Dobra widoczność')

#WYKRES MOŻLIWOŚCI SPOWODOWANIA WYPADKU
plot_fuzzyset(ax3, wypadek_bardzo_male, x2, 'g', label='Bardzo mała możliwość wypadku')
plot_fuzzyset(ax3, wypadek_male, x2, 'b', label='Mała możliwość wypadku')
plot_fuzzyset(ax3, wypadek_srednie, x2, 'y', label='Średnia możliwość wypadku')
plot_fuzzyset(ax3, wypadek_duze, x2, 'r', label='Duża możiwośc wypadku')

#RYSOWANIE WYKRESY
plt.show()

#ZASADY
scope = locals()

rule1 = 'if predkosc is A_predkosc_mala and widocznosc is A_widocznosc_bardzo_slaba then wypadek is A_wypadek_male'
rule2 = 'if predkosc is A_predkosc_mala and widocznosc is A_widocznosc_srednia  then wypadek is A_wypadek_male'
rule3 = 'if predkosc is A_predkosc_mala and widocznosc is A_widocznosc_dobra  then wypadek is A_wypadek_bardzo_male'
rule4 = 'if predkosc is A_predkosc_srednia  and widocznosc is A_widocznosc_bardzo_slaba then wypadek is A_wypadek_srednie'
rule5 = 'if predkosc is A_predkosc_srednia  and widocznosc is A_widocznosc_srednia  then wypadek is A_wypadek_srednie'
rule6 = 'if predkosc is A_predkosc_srednia  and widocznosc is A_widocznosc_dobra  then wypadek is A_wypadek_male'
rule7 = 'if predkosc is A_predkosc_szybka  and widocznosc is A_widocznosc_bardzo_slaba then wypadek is A_wypadek_duze'
rule8 = 'if predkosc is A_predkosc_szybka  and widocznosc is A_widocznosc_srednia  then wypadek is A_wypadek_srednie'
rule9 = 'if predkosc is A_predkosc_szybka  and widocznosc is A_widocznosc_dobra  then wypadek is A_wypadek_srednie'
rule10 = 'if predkosc is A_predkosc_bardzo_szybka and widocznosc is A_widocznosc_bardzo_slaba then wypadek is A_wypadek_duze'
rule11 = 'if predkosc is A_predkosc_bardzo_szybka and widocznosc is A_widocznosc_srednia  then wypadek is A_wypadek_duze'
rule12 = 'if predkosc is A_predkosc_bardzo_szybka and widocznosc is A_widocznosc_dobra  then wypadek is A_wypadek_srednie'



block = RuleBlock('rb_mamdani', operators=('MIN', 'MAX', 'ZADEH'), activation='MIN', accumulation='MAX')
block.add_rules(rule1,rule2,rule3,rule4,rule5,rule6,rule7,rule8,rule9,rule10,rule11,rule12,scope=scope)

#TWORZENIE MODELU MAMDANI
mamdani = MamdaniSystem('mamdani model', block)

#DANE WEJŚCIOWE
inputs = {'predkosc':170, 'widocznosc': 0.1}
res = mamdani.compute(inputs)
print(res)

#PRZYGOTOWANIE SIATKI
sampled_one = np.linspace(10, 200, 20)
sampled_two = np.linspace(0.05,4, 20)
x, y = np.meshgrid(sampled_one, sampled_two)
z = np.zeros((len(sampled_one), len(sampled_two)))


for i in range(len(sampled_one)):
    for j in range(len(sampled_two)):
        inputs = {'predkosc': x[i, j], 'widocznosc': y[i, j]}
        res = mamdani.compute(inputs)
        z[i, j] = res['rb_mamdani']['wypadek']

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
surf = ax.plot_surface(x, y, z, rstride=1, cstride=1, cmap='viridis', linewidth=0.4, antialiased=True)
cset = ax.contourf(x, y, z, zdir='z', offset= -1, cmap='viridis', alpha=0.5)
cset = ax.contourf(x, y, z, zdir='x', offset= 11, cmap='viridis', alpha=0.5)
cset = ax.contourf(x, y, z, zdir='y', offset= 11, cmap='viridis', alpha=0.5)
ax.set_xlabel('predkosc')
ax.set_ylabel('widocznosc')
ax.set_zlabel('wypadek')
ax.view_init(30, 200)