#Example 5 Gradio Version
# Example 4 Decision Treen Analysis
# Example 2 rework on classification of flow regime type
import pandas as pd
import gradio as gr
df = pd.read_csv("olga.tab" ,sep="\t")
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
plt.figure(figsize=(16,8))
import numpy as np
from sklearn.tree import DecisionTreeClassifier
plt.rcParams['font.size'] = 16
def rs(array):
    return np.array(array).reshape(-1, 1)
#drop empty values
df.dropna(inplace=True)
#Redcue to stable/non stable regime 1 = Unstable, 2
df = df.dropna()
df = df[ (df["QOST [STB/d]"]>= 0)  & (df["QGST [MMscf/d]"] >=0) ] 
#create training values
X_train = df[["QOST [STB/d]", "QGST [MMscf/d]"]]
y_train = df["ID"]
#initialize mode
DT = DecisionTreeClassifier(max_depth=5)
#train and predict
DT.fit(X_train, y_train)
y_pred = list(DT.predict(X_train))

color_map = {1:"red", 2:"blue", 3:"green", 4:"magenta"}
y_color_original = pd.Series(y_train).map(color_map)
y_pred_color = pd.Series(y_pred).map(color_map)
#create dual pltis plot
#tcf = plt.tricontourf(df["QOST [STB/d]"], df["QGST [MMscf/d]"],y_pred , cmap="rainbow")
#plt[0].scatter(df["QOST [STB/d]"], df["QGST [MMscf/d]"], color=y_color_original, s=100, alpha=0.5)

#Create custom legend
legend_items = []
regime = ["Stratified", "Bubble", "Slug", "Annular"]
for i in range(1,5):
    #Create Custom Scatter Points
    scatt = Line2D([0], [0], color=color_map[i],marker="o", markersize=15, lw=0,label=regime[i-1])
    legend_items.append(scatt)
#Setup Plot

def predict_model(qost, qgst):
    x_pred = pd.DataFrame({"QOST [STB/d]":[qost], "QGST [MMscf/d]":[qgst]})
    _pred = DT.predict(x_pred)
    plt.scatter(df["QOST [STB/d]"], df["QGST [MMscf/d]"], color=y_pred_color, s=100, alpha=0.5)
    plt.scatter(x_pred["QOST [STB/d]"], x_pred["QGST [MMscf/d]"], color="black",marker="X", s=200, label="Your Point")
    plt.xlabel("QOST") ;plt.ylabel("QGST")
    plt.legend(True,handles=legend_items)
    plt.title("Original Dataset"); plt.title("Predicted Dataset")
    plt.ylim(0,200)
    return plt, regime[ int(_pred[0])-1  ]
interface = gr.Interface(predict_model, inputs=[gr.Number(label="QOil[bbl/day]", value=20000), gr.Number(label="QGas[mmscf/day]", value=100)], 
                         outputs=[gr.Plot(label="Flow Map"),gr.Textbox(label="Identified Regime")])
interface.launch()