import torch
import numpy as np 
import pandas as pd
import gradio as gr
import matplotlib.pyplot as plt

from model import LSTMRegressor
from torch.utils.data import TensorDataset, DataLoader
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from pathlib import Path 

# Taking datas
train_df = pd.read_csv("train_FD001.txt", sep = " ", header = None)
test_df = pd.read_csv("test_FD001.txt", sep = " ", header = None)
rul_df = pd.read_csv("RUL_FD001.txt", sep = " ", header = None)

# Drop empty columns
train_df.dropna(axis = 1, inplace = True)
test_df.dropna(axis = 1, inplace = True)

# Preparing the column names
columns = ["unit_number", 
           "time_in_cycles"] + [f"operational_setting_{i}" for i in range(1, 4)] + [f"sensor_measurement_{i}" for i in range(1, 22)]

train_df.columns = columns
test_df.columns = columns

rul_df = rul_df[[0]]
rul_df.columns = ["RUL"]

# Prepare RUL column for train set
rul_per_engine = train_df.groupby("unit_number")["time_in_cycles"].max().reset_index()
rul_per_engine.columns = ["unit_number", "max_cycle"]

train_df = train_df.merge(rul_per_engine, on = "unit_number", how = "left")
train_df["RUL"] = train_df["max_cycle"] - train_df["time_in_cycles"]
train_df.drop(columns = ["max_cycle"], axis = 1, inplace = True)

def generate_sequence(data, sequence_length, feature_cols):
    sequences = []
    targets = []

    for engine_id in data["unit_number"].unique():
        engine_data = data[data["unit_number"] == engine_id]
        engine_data = engine_data.reset_index(drop = True)

        for i in range(len(engine_data) - sequence_length + 1):
            seq_x = engine_data.loc[i : i + sequence_length -1, feature_cols].values
            seq_y = engine_data.loc[i + sequence_length - 1, "RUL"]
            sequences.append(seq_x)
            targets.append(seq_y)

    return np.array(sequences), np.array(targets)

feature_columns = [f"operational_setting_{i}" for i in range(1, 4)] + [f"sensor_measurement_{i}" for i in range(1, 22)]
train_sequence_length = 160
X_train, y_train = generate_sequence(train_df, train_sequence_length, feature_columns)

X_flat = X_train.reshape(-1, X_train.shape[-1])

# Scale data between 1 - 0
feature_scaler = MinMaxScaler()
X_scaled_flat = feature_scaler.fit_transform(X_flat)
X_train_scaled = X_scaled_flat.reshape(X_train.shape)
y_train_reshaped = y_train.reshape(-1, 1)
target_scaler = MinMaxScaler()
y_train_scaled = target_scaler.fit_transform(y_train_reshaped)

# Prepare model
input_size = 24
hidden_size = 128
num_layers = 2
dropout = 0.3
model = LSTMRegressor(input_size = input_size, hidden_size = hidden_size, num_layers = num_layers, dropout = dropout)

# Load model
model.load_state_dict(torch.load(f = "LSTM_CMAPSS_with_test_and_rul_data_MAE_14.31_RMSE_18.22_r2_0.81.pth",
                                 weights_only = True,
                                 map_location=torch.device('cpu')))

# Prepare test datas
test_sequence_length = 150
X_test_list = []
y_test_list = []

for i, engine_id in enumerate(test_df["unit_number"].unique()):
    engine_data = test_df[test_df["unit_number"] == engine_id].reset_index(drop = True)

    if len(engine_data) >= test_sequence_length:
        seq = engine_data.iloc[-test_sequence_length:][feature_columns].values
    else:
        padding = np.zeros((test_sequence_length - len(engine_data), len(feature_columns)))
        seq = np.vstack((padding, engine_data[feature_columns].values))

    X_test_list.append(seq)
    y_test_list.append(rul_df.iloc[i].values[0])

X_test = np.array(X_test_list)    
y_test = np.array(y_test_list).reshape(-1, 1)

# Scale test data
X_test_flat = X_test.reshape(-1, X_test.shape[-1])
X_test_scaled_flat = feature_scaler.transform(X_test_flat)
X_test_scaled = X_test_scaled_flat.reshape(X_test.shape)
y_test_scaled = target_scaler.transform(y_test)

# Make predictions
model.eval()
X_test_tensor = torch.tensor(X_test_scaled, dtype = torch.float32)

with torch.inference_mode():
    y_pred_scaled = model(X_test_tensor).cpu().numpy()

y_pred = target_scaler.inverse_transform(y_pred_scaled)

y_pred = y_pred.flatten()
y_true = y_test.flatten()

mae = mean_absolute_error(y_true, y_pred)
mse = mean_squared_error(y_true, y_pred)
rmse = np.sqrt(mse)
r2 = r2_score(y_true, y_pred)

def plot_all_engines():
    fig, ax = plt.subplots(figsize = (12, 6), dpi = 300)
    ax.plot(y_true, label = "Real RUL", marker = 'o', color = "green")
    ax.plot(y_pred, label = "Predicted RUL", marker = 'x', color = "orange")
    ax.set_title("RUL Predictions for All Motors")
    ax.set_xlabel("Motor No")
    ax.set_ylabel("RUL")
    ax.legend()
    ax.grid(True)
    return fig

def analyze_engine(engine_id):
    idx = int(engine_id.split()[-1]) - 1
    true_rul = y_true[idx]
    pred_rul = y_pred[idx]

    return f"""
### 🔍 [INFO] {engine_id} 

- Real RUL: **{true_rul:.2f}**
- Predicted RUL: **{pred_rul:.2f}**
- Error (MAE): **{abs(true_rul - pred_rul):.2f}**
"""


#Gradio UI
with gr.Blocks() as demo:

    article = "Created by Ahmet Hakan Karadana"
    gr.Markdown("## ✈️ Jet Engine RUL Prediction (CMAPSS)")
    gr.Markdown(f"""
    **Model Metrics:**  
    ✅ MAE: {mae:.2f}  
    ✅ RMSE: {rmse:.2f}  
    ✅ R²: {r2:.2f}
    """)

    gr.Markdown("### 📈 Actual & Predicted RUL Values of All Engines")
    gr.Plot(plot_all_engines)

    gr.Markdown("### 🔍 Select Engine for Numerical Values")
    dropdown = gr.Dropdown(choices=[f"Engine {i}" for i in range(1, 101)], label="Select Motor")
    output_text = gr.Markdown()
    dropdown.change(analyze_engine, inputs=dropdown, outputs=output_text)
    gr.Markdown(f"*{article}*")
    
demo.launch()