import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)
import matplotlib.pyplot as plt # plotting graphs
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
from sklearn.tree import DecisionTreeClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.svm import SVC
from xgboost import XGBClassifier
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.optimizers import Adam

import os
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.join(BASE_DIR, "data")

'''
        os.path.join(BASE_DIR),                    # Same folder
        os.path.join(BASE_DIR, "data"),            # data subfolder
        os.path.join(BASE_DIR, "..", "data"),      # One level up
        os.path.join(BASE_DIR, "..", "..", "data") # Two levels up
'''

file_path = os.path.join(DATA_DIR, "data.csv")
df = pd.read_csv(file_path)

#Time-Domain Behavior of Phase-A Current and Voltage for Different Fault Types
fs = 3200  # Sampling frequency (Hz)
ts = 1 / fs
N_window = 10

df_sample = df.copy()
df_sample['time_stamp'] = df_sample.groupby('Class_lable').cumcount() * ts
# Convert to long format for faceting
df_long = df_sample.melt(
    id_vars=['time_stamp', 'Class_lable'],
    value_vars=['Ia_Mag', 'Va_Mag'],
    var_name='Signal',
    value_name='Magnitude'
)

# Rename for clean titles
df_long['Signal'] = df_long['Signal'].map({
    'Ia_Mag': 'Phase-A Current (Ia)',
    'Va_Mag': 'Phase-A Voltage (Va)'
})

# Create Facet Grid
g = sns.FacetGrid(
    df_long,
    row='Signal',
    col='Class_lable',
    height=4,
    aspect=1.5,
    margin_titles=True,
    sharex=True,
    sharey=False
)

# Map scatter plots
g.map_dataframe(
    sns.scatterplot,
    x='time_stamp',
    y='Magnitude',
    s=10
)

# Axis labels
g.set_axis_labels("Time (s)", "Magnitude")

# Titles
g.set_titles(row_template='{row_name}', col_template='Fault Type: {col_name}')

# Overall figure title
g.fig.suptitle(
    'Time-Domain Behavior of Phase-A Current and Voltage for Different Fault Types',
    fontsize=14,
    y=1.03,
    ha='center'
)

plt.tight_layout()
# plt.show()
plt.savefig("plot1_current_voltage_behavior_fault_types.png", dpi=600, bbox_inches='tight')

#CORRELATION MATRIX
correlation_matrix = df.corr()
# Plot the heatmap
fig2 = plt.figure(figsize=(12, 10))
sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', fmt='.2f')
plt.title('Correlation Matrix of Numerical Features')
# plt.show()
fig2.savefig('plot2_correlation_matrix.png', dpi=600, bbox_inches='tight')


#Voltage–Current Feature Space
normal = df[df['Class_lable'] == 0]
fault = df[df['Class_lable'] == 1]

fig3 = plt.figure(figsize=(7, 5))

plt.scatter(
    normal['Va_Mag'], normal['Ia_Mag'],
    alpha=0.4, label='Normal', s=10
)

plt.scatter(
    fault['Va_Mag'], fault['Ia_Mag'],
    alpha=0.6, label='Fault', s=10
)

plt.xlabel('Voltage Magnitude (Va)')
plt.ylabel('Current Magnitude (Ia)')
plt.title('Voltage–Current Feature Space')
plt.legend()
plt.grid(True)

plt.tight_layout()
# plt.show()
fig3.savefig('plot3_feature_space.png', dpi=600, bbox_inches='tight')


# Sliding Window-Based Sub-Cycle Observation
# Example time vector (one AC cycle ~20 ms)
t = np.arange(0, 0.02, ts)

# Example signal (replace with real Ia or Va)
signal = np.sin(2 * np.pi * 50 * t)

# Simulated fault deformation
fault_start = int(0.008 / ts)
signal[fault_start:fault_start+15] *= 1.8  # abrupt distortion

# Sliding window location (first few samples after fault)
window_start = fault_start + 2
window_end = window_start + N_window

fig4 = plt.figure(figsize=(8, 4))

plt.plot(t * 1000, signal, label='PMU Signal', linewidth=1.5)

# Highlight sliding window
plt.axvspan(
    t[window_start] * 1000,
    t[window_end] * 1000,
    color='red',
    alpha=0.3,
    label='Sliding Window (10 samples)'
)

# Fault marker
plt.axvline(
    t[fault_start] * 1000,
    color='black',
    linestyle='--',
    label='Fault Inception'
)

plt.xlabel('Time (ms)')
plt.ylabel('Signal Magnitude')
plt.title('Sliding Window-Based Sub-Cycle Observation')
plt.legend()
plt.grid(True)
plt.tight_layout()
# plt.show()
fig4.savefig('plot4_sliding_window_architecture.png', dpi=600, bbox_inches='tight')
 
#Function to generate Stream for Sliding Window Architecture
def generate_realistic_stream(df, n_events=50):
    """
    Creates a continuous timeline by stitching random Normal segments 
    with random Fault segments.
    """
    # 1. Separate Normal (NF) and Fault data
    nf_data = df[df['Class_lable'] == 0]
    fault_data = df[df['Class_lable'] != 0]
    
    stream_segments = []
    labels = []
    
    # 2. Loop to create 'n' events (e.g., 50 different faults happening over time)
    for _ in range(n_events):
        # A. Normal Operation Phase
        # Pick a random duration (e.g., between 200 and 500 steps)
        n_duration = np.random.randint(200, 500)
        # Pick a random start point in the NF data
        start_idx = np.random.randint(0, len(nf_data) - n_duration)
        # Cut that slice out
        segment = nf_data.iloc[start_idx : start_idx + n_duration]
        stream_segments.append(segment)
        
        # B. Fault Phase
        # Pick a random fault type (LG, LL, etc.)
        random_fault_type = np.random.choice(fault_data['Class_lable'].unique())
        f_segment_full = fault_data[fault_data['Class_lable'] == random_fault_type]
        
        # We simulate the fault JUST starting (first 100 rows)
        # This teaches the model to recognize the "onset" of the fault
        start_f = np.random.randint(0, len(f_segment_full) - 100)
        segment_fault = f_segment_full.iloc[start_f : start_f + 100] 
        stream_segments.append(segment_fault)
        
    # 3. Glue it all together into one long dataframe
    full_stream = pd.concat(stream_segments).reset_index(drop=True)
    return full_stream

# Generate two separate streams: one for teaching the model, one for testing it
train_stream = generate_realistic_stream(df, n_events=100) # Long stream to learn
test_stream = generate_realistic_stream(df, n_events=20)   # Short stream to test
print(f"Training Stream Length: {len(train_stream)} rows")


#Function to create Windows
def create_windows(data, window_size=10, stride=1):
    X = []
    y = []

    features = data[['Va_Mag','Ia_Mag','Vb_Mag','Ib_Mag','Vc_Mag','Ic_Mag','source_id']].values
    # features = data[['Va_Mag','Ia_Mag','Vb_Mag','Ib_Mag','Vc_Mag','Ic_Mag']].values
    targets = data["Class_lable"].values

    for i in range(0, len(data) - window_size + 1, stride):
        window = features[i : i + window_size]
        X.append(window.flatten())

        # Onset-aware labeling (IMPORTANT)
        window_labels = targets[i : i + window_size]
        label = 1 if np.any(window_labels == 1) else 0
        y.append(label)

    return np.array(X), np.array(y)

# Sliding window parameters
WINDOW_SIZE = 15   
STRIDE = 1

X_train, y_train = create_windows(train_stream, WINDOW_SIZE)
X_test, y_test = create_windows(test_stream, WINDOW_SIZE)

scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

rf = RandomForestClassifier(
    n_estimators=300,
    max_depth=20,
    min_samples_split=5,
    class_weight="balanced",
    random_state=42,
    n_jobs=-1
)

rf.fit(X_train, y_train)

y_pred = rf.predict(X_test)

print(confusion_matrix(y_test, y_pred))
print(classification_report(y_test, y_pred, digits=4))

acc = accuracy_score(y_test, y_pred)
print("Accuracy:", acc)

# Get probability (0.0 to 1.0) instead of just 0 or 1
probs = rf.predict_proba(X_test)[:, 1]

fig5 = plt.figure(figsize=(15, 6))
# Zoom in on a specific section (e.g., samples 0 to 1000)
plt.plot(y_test[0:1000], label='Actual Fault Status', color='black', linewidth=2)
plt.plot(probs[0:1000], label='Model Prediction Probability', color='red', alpha=0.7)
plt.axhline(0.5, linestyle='--', color='gray') # Threshold
plt.title("Early Warning System Response")
plt.ylabel("Fault Probability")
plt.xlabel("Time Steps")
plt.legend()
# plt.show()
fig5.savefig('plot5_early_warning.png', dpi=600, bbox_inches='tight')

# Calculate Probability
probs = rf.predict_proba(X_test)[:, 1]

# --- THE FIX: SMOOTHING ---
# We take the average of the last 5 predictions.
# If the model spikes for just 1 step, the average won't move much.
# If the model spikes and STAYS high, the average will go up.
probs_smooth = pd.Series(probs).rolling(window=5).mean()

fig6 = plt.figure(figsize=(15, 6))
plt.plot(y_test[:1000], label='Actual Fault', color='black', linewidth=2)

# Plot the noisy raw predictions (faint)
plt.plot(probs[:1000], label='Raw Model Output (Nervous)', color='red', alpha=0.3, linewidth=1)

# Plot the smooth predictions (solid)
plt.plot(probs_smooth[:1000], label='Smoothed Output (Reliable)', color='blue', linewidth=2)

plt.axhline(0.5, linestyle='--', color='gray')
plt.title("Raw Spikes vs. Smoothed Early Warning")
plt.ylabel("Fault Probability")
plt.xlabel("Time Steps")
plt.legend()
# plt.show()
fig6.savefig('plot6_probability_smoothing.png', dpi=600, bbox_inches='tight')


dt = DecisionTreeClassifier(
    max_depth=20,
    min_samples_split=5,
    class_weight="balanced",
    random_state=42
)

dt.fit(X_train, y_train)
y_pred_dt = dt.predict(X_test)

print(confusion_matrix(y_test, y_pred_dt))
print(classification_report(y_test, y_pred_dt, digits=4))

lr = LogisticRegression(
    max_iter=1000,
    class_weight="balanced",
    n_jobs=-1
)

lr.fit(X_train, y_train)
y_pred_lr = lr.predict(X_test)

print(confusion_matrix(y_test, y_pred_lr))
print(classification_report(y_test, y_pred_lr, digits=4))

svm = SVC(kernel='rbf', random_state=42)
svm.fit(X_train, y_train)
y_pred_svm = svm.predict(X_test)

print(confusion_matrix(y_test, y_pred_svm))
print(classification_report(y_test, y_pred_svm, digits=4))

xgb = XGBClassifier(
    n_estimators=300,
    max_depth=8,
    learning_rate=0.05,
    subsample=0.8,
    colsample_bytree=0.8,
    eval_metric="logloss",
    random_state=42
)

xgb.fit(X_train, y_train)
y_pred_xgb = xgb.predict(X_test)

print(confusion_matrix(y_test, y_pred_xgb))
print(classification_report(y_test, y_pred_xgb, digits=4))

model = Sequential([
    Dense(32, activation='relu', input_shape=(X_train.shape[1],)),
    Dense(16, activation='relu'),
    Dense(1, activation='sigmoid')   
])

model.compile(
    optimizer=Adam(learning_rate=0.001),
    loss='binary_crossentropy',     
    metrics=['accuracy']
)

model.summary()

history = model.fit(
    X_train, y_train,
    epochs=20,
    batch_size=32,
    validation_split=0.1,
    verbose=1
)

loss, accuracy = model.evaluate(X_test, y_test, verbose=0)
print(f"Test Accuracy: {accuracy:.4f}")

# Correct predictions
y_pred_probs = model.predict(X_test)
y_pred = (y_pred_probs > 0.5).astype(int).ravel()

print(classification_report(y_test, y_pred, digits=4))

# Confusion matrices (TN, FP / FN, TP)
models = [
    "Random Forest",
    "Decision Tree",
    "SVM (Linear)",
    "XGBoost Classifier",
    "Artificial Neural Network",
    "Logistic Regression"
]

confusion_matrices = [
    np.array([[6143, 0],
              [  91, 2080]]),

    np.array([[6084, 59],
              [  54, 2117]]),

    np.array([[6143, 0],
              [ 105, 2066]]),

    np.array([[6143, 0],
              [   0, 2171]]),

    np.array([[6220, 0],
              [ 316, 1950]]),

    np.array([[4747, 1396],
              [ 596, 1575]])
]

fig7, axes = plt.subplots(2, 3, figsize=(12, 7))

for ax, model, cm in zip(axes.flatten(), models, confusion_matrices):
    im = ax.imshow(cm, cmap="Blues")
    ax.set_title(model, fontsize=11)
    ax.set_xticks([0, 1])
    ax.set_yticks([0, 1])
    ax.set_xticklabels(["Normal", "Fault"])
    ax.set_yticklabels(["Normal", "Fault"])
    ax.set_xlabel("Predicted")
    ax.set_ylabel("True")

    threshold = cm.max() / 2.0

    for i in range(2):
        for j in range(2):
            color = "white" if cm[i, j] > threshold else "black"
            ax.text(j, i, cm[i, j],
                    ha="center", va="center",
                    color=color, fontsize=10)

plt.tight_layout()
fig7.savefig('plot7_confusion_matrix.png', dpi=600, bbox_inches='tight')

plt.show()