import os
import pandas as pd
import numpy as np
from scipy.signal import medfilt
from scipy.integrate import simpson
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error
import joblib
import matplotlib.pyplot as plt
from micromlgen import port


def clean_and_extract_features(df, drop_window=50, recovery_threshold=0.9):
    df["Filtered Voltage"] = medfilt(df["Voltage (V)"], kernel_size=5)
    baseline = df["Filtered Voltage"].iloc[0]
    drop_idx = df[df["Filtered Voltage"] < (baseline - 0.01)].index
    if drop_idx.empty or (drop_idx[0] + drop_window > len(df)):
        return None
    drop_start = drop_idx[0]
    window_df = df.iloc[drop_start : drop_start + drop_window].copy()
    voltages = window_df["Filtered Voltage"].values
    min_voltage = voltages.min()
    drop_height = baseline - min_voltage
    area_under = simpson(baseline - voltages)
    slope = (voltages[-1] - voltages[0]) / (drop_window * 0.1)
    recovery_voltage = min_voltage + (baseline - min_voltage) * recovery_threshold
    recovery_idx = next(
        (i for i, v in enumerate(voltages) if v >= recovery_voltage), drop_window - 1
    )
    time_to_recovery = recovery_idx * 0.1
    peak_to_peak = voltages.max() - voltages.min()
    return {
        "baseline": baseline,
        "drop_height": drop_height,
        "min_voltage": min_voltage,
        "area_under_drop": area_under,
        "slope_post_drop": slope,
        "time_to_recovery": time_to_recovery,
        "peak_to_peak_range": peak_to_peak,
    }


def train_model(log_folder):
    features_list = []
    for filename in os.listdir(log_folder):
        if filename.endswith(".csv") and filename.startswith("log_"):
            filepath = os.path.join(log_folder, filename)
            df = pd.read_csv(filepath)
            features = clean_and_extract_features(df)
            if features:
                glucose = float(filename.split("_")[1].replace("mmolL.csv", ""))
                features["glucose"] = glucose
                features_list.append(features)
    df_all = pd.DataFrame(features_list)
    X = df_all.drop(columns=["glucose"])
    y = df_all["glucose"]
    model = RandomForestRegressor(n_estimators=15, max_depth=6)
    model.fit(X, y)
    y_pred = model.predict(X)
    mae = mean_absolute_error(y, y_pred)
    print(f"Model trained. Mean Absolute Error (MAE): {mae:.2f} mmol/L")
    plt.figure(figsize=(8, 6))
    plt.scatter(y, y_pred, color="blue", label="Predicted vs Actual")
    plt.plot([y.min(), y.max()], [y.min(), y.max()], "r--", lw=2, label="Ideal Fit")
    plt.xlabel("Actual Glucose Level (mmol/L)")
    plt.ylabel("Predicted Glucose Level (mmol/L)")
    plt.title("Regression Fit: Actual vs Predicted")
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.show()
    joblib.dump(model, "glucose_rf_model.pkl")
    print("Model trained and saved as 'glucose_rf_model.pkl'")
    return model, X.columns.tolist()


def predict_glucose(csv_path, model_path="glucose_rf_model.pkl"):
    model = joblib.load(model_path)
    df = pd.read_csv(csv_path)
    features = clean_and_extract_features(df)
    if features:
        X_test = pd.DataFrame([features])
        prediction = model.predict(X_test)[0]
        print(f"Predicted Glucose Level: {prediction:.2f} mmol/L")
        return prediction
    else:
        print("Invalid or unusable log format for prediction.")
        return None


# Example usage:
model, feature_names = train_model("./train_logs")
joblib.dump(model, "glucose_rf_model.pkl")
predict_glucose("ref/reference_log.csv")
print(port(model, class_name="GlucoseModel"))