#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((9.9,0),(10,1),(45,1),(50,0))
predkosc_srednia = Trapezoid((40,0),(50,1),(80,1),(95,0))
predkosc_szybka = Trapezoid((85,0),(90,1),(140,1),(145,0))
predkosc_bardzo_szybka = Trapezoid((135,0),(140,1),(200,1),(200.1,0))

A_predkosc_mala = Adjective('predkosc_mala', predkosc_mala)
A_predkosc_srednia = Adjective('predkosc_srednia', predkosc_srednia)
A_predkosc_szybka = Adjective('predkosc_szybka', predkosc_szybka)
A_predkosc_bardzo_szybka = Adjective('predkosc_bardzo_szybka', 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('widocznosc_bardzo_slaba', widocznosc_bardzo_slaba)
A_widocznosc_srednia = Adjective('widocznosc_srednia', widocznosc_srednia)
A_widocznosc_dobra = Adjective('widocznosc_dobra', 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('wypadek_bardzo_male', wypadek_bardzo_male)
A_wypadek_male = Adjective('wypadek_male', wypadek_male)
A_wypadek_srednie = Adjective('wypadek_srednie', wypadek_srednie)
A_wypadek_duze = Adjective('wypadek_duze', 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='predkosc_mala')
plot_fuzzyset(ax1, predkosc_srednia, x, 'b', label='predkosc_srednia')
plot_fuzzyset(ax1, predkosc_szybka, x, 'y', label='predkosc_szybka')
plot_fuzzyset(ax1, predkosc_bardzo_szybka, x, 'r', label='predkosc_bardzo_szybka')

#WYKRES WIDOCZNOŚCI
plot_fuzzyset(ax2, widocznosc_bardzo_slaba, x1, 'r', label='widocznosc_bardzo_slaba')
plot_fuzzyset(ax2, widocznosc_srednia, x1, 'b', label='widocznosc_srednia')
plot_fuzzyset(ax2, widocznosc_dobra, x1, 'g', label='widocznosc_dobra')

#WYKRES MOŻLIWOŚCI SPOWODOWANIA WYPADKU
plot_fuzzyset(ax3, wypadek_bardzo_male, x2, 'g', label='wypadek_bardzo_male')
plot_fuzzyset(ax3, wypadek_male, x2, 'b', label='wypadek_male')
plot_fuzzyset(ax3, wypadek_srednie, x2, 'y', label='wypadek_srednie')
plot_fuzzyset(ax3, wypadek_duze, x2, 'r', label='wypadek_duze')

#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':190, '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)