import numpy as np
import pandas as pd
import tensorflow as tf

# ============== CONFIGURATION ==============
CSV_PATH_1 = "nonlinear_full_truth_data.csv"
CSV_PATH_2 = "model4_queue_aware_dcf_tgax.csv"
OUTPUT_CSV = "mape_errors.csv"
OUTPUT_COLUMN_1 = "y_p3"
OUTPUT_COLUMN_2 = "targetDelay"
ROW_INDICES_1 = list(range(9101, 9150))
ROW_INDICES_2 = list(range(1801, 1850))

# Scale factor for models 1-3 (e.g., m -> mm)
# This shifts values away from zero before MAPE calculation
# Set to 1.0 for no scaling
SCALE_FACTOR_1 = 1.0   # meters -> millimeters for y_p3

# Minimum absolute value threshold to avoid division by near-zero
# Samples with |y_true_scaled| < this are excluded from MAPE
ZERO_THRESHOLD = 1e-12

MODELS = {
    "model1": {
        "path": "model1.tflite",
        "input_cols": ["p1", "p2", "p3", "v1", "v2", "v3", "i_a", "V"],
        "csv": 1,
        "target_output_index": -3,
    },
    "model2": {
        "path": "model2.tflite",
        "input_cols": ["p1", "p2", "p3", "v1", "v2", "v3", "temp1", "temp2", "temp3", "i_a", "V"],
        "csv": 1,
        "target_output_index": -3,
    },
    "model3": {
        "path": "model3.tflite",
        "input_cols": ["p1", "p2", "p3", "v1", "v2", "v3", "temp1", "temp2", "temp3", "wear1", "wear2", "wear3", "i_a", "V"],
        "csv": 1,
        "target_output_index": -3,
    },
    "model4": {"path": "model4.tflite", "input_cols": ["delayLag1","delayLag2","delayLag3","delayLag4","delayLag5","delayLag6","arrivalGapMs","packetBytes","snrDb","mcs","txTimeMs"], "csv": 2},
    "model5": {"path": "model5.tflite", "input_cols": ["delayLag1","delayLag2","delayLag3","delayLag4","delayLag5","delayLag6","arrivalGapMs","packetBytes","snrDb","mcs","packetDetected","retryCount","successFlag","snrMarginDb","txTimeMs"], "csv": 2},
    "model6": {"path": "model6.tflite", "input_cols": ["delayLag1","delayLag2","delayLag3","delayLag4","delayLag5","delayLag6","arrivalGapMs","packetBytes","snrDb","mcs","packetDetected","retryCount","successFlag","snrMarginDb","rxPowerdB","txTimeMs"], "csv": 2},
    "model7": {"path": "model7.tflite", "input_cols": ["delayLag1","delayLag2","delayLag3","delayLag4","delayLag5","delayLag6","arrivalGapMs","arrivalTimeMs","packetBytes","nStations","busyRatio","snrDb","mcs","collisionProb","packetDetected","backoff0Slots","retryCount","queueLen","waitMs","serviceMs","txTimeMs","successFlag","rxPowerdB","snrMarginDb"], "csv": 2},
}
# ===========================================

df1 = pd.read_csv(CSV_PATH_1)
df2 = pd.read_csv(CSV_PATH_2)
if ROW_INDICES_1 is not None:
    df1 = df1.iloc[ROW_INDICES_1].reset_index(drop=True)
if ROW_INDICES_2 is not None:
    df2 = df2.iloc[ROW_INDICES_2].reset_index(drop=True)

csv_map = {1: df1, 2: df2}
output_col_map = {1: OUTPUT_COLUMN_1, 2: OUTPUT_COLUMN_2}
scale_map = {1: SCALE_FACTOR_1, 2: 1.0}

results = []

for name, config in MODELS.items():
    print(f"\nEvaluating {name}...")
    df = csv_map[config["csv"]]
    output_col = output_col_map[config["csv"]]
    scale = scale_map[config["csv"]]

    X_test = df[config["input_cols"]].values.astype(np.float32)
    y_true = df[output_col].values.astype(np.float64)

    interpreter = tf.lite.Interpreter(model_path=config["path"])
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    is_multi_output = "target_output_index" in config

    y_pred = []
    for sample in X_test:
        sample = sample.reshape(1, -1).astype(np.float32)
        interpreter.resize_tensor_input(input_details[0]['index'], sample.shape)
        interpreter.allocate_tensors()
        interpreter.set_tensor(input_details[0]['index'], sample)
        interpreter.invoke()
        output = interpreter.get_tensor(output_details[0]['index']).flatten()

        if is_multi_output:
            y_pred.append(output[config["target_output_index"]])
        else:
            y_pred.append(output[0])

    y_pred = np.array(y_pred, dtype=np.float64)

    # Apply scaling (e.g., m -> mm for models 1-3)
    y_true_scaled = y_true * scale
    y_pred_scaled = y_pred * scale

    # Exclude near-zero actual values to avoid MAPE explosion
    mask = np.abs(y_true_scaled) > ZERO_THRESHOLD
    n_excluded = np.sum(~mask)
    n_total = len(y_true)

    if np.any(mask):
        mape = np.mean(np.abs((y_true_scaled[mask] - y_pred_scaled[mask]) / y_true_scaled[mask])) * 100
    else:
        mape = float('inf')

    print(f"  Scale factor: {scale}")
    print(f"  Samples used: {np.sum(mask)}/{n_total} (excluded {n_excluded} near-zero)")
    print(f"  y_true range (scaled): [{y_true_scaled.min():.4f}, {y_true_scaled.max():.4f}]")
    print(f"  y_pred range (scaled): [{y_pred_scaled.min():.4f}, {y_pred_scaled.max():.4f}]")
    print(f"  MAPE: {mape:.4f}%")

    results.append({
        "Model": name,
        "MAPE_%": round(mape, 4),
        "Samples_Used": int(np.sum(mask)),
        "Samples_Total": n_total,
    })

results_df = pd.DataFrame(results)
results_df.to_csv(OUTPUT_CSV, index=False)
print(f"\n✅ Results saved to {OUTPUT_CSV}")
