import os
import pandas as pd
import json

experiments = {
    "all": ["(all 35 bands)"],
    "17": ["TPI", "ROUGHNESS", "LST_winter", "NDWI", "LST_summer", "REL_RELIEF", "ASPECT", "LST_anomaly", "ELEV_COP", "B2", "HILLSHADE", "B8A", "NDMI", "B12", "B3", "B4", "B8"],
    "14": ["TPI", "ROUGHNESS", "LST_winter", "NDWI", "LST_summer", "REL_RELIEF", "ASPECT", "LST_anomaly", "ELEV_COP", "B2", "HILLSHADE", "B8A", "NDMI", "B12"],
    "opt_idx": ["B2", "B3", "B4", "B8", "B8A", "B11", "B12", "NDVI", "NDWI", "NDMI", "NDSI", "BSI", "ALBEDO", "ELEV_COP"],
    "terr_therm": ["TPI", "ROUGHNESS", "SLOPE", "ASPECT", "REL_RELIEF", "HILLSHADE", "CURVATURE", "ELEV_COP", "LST_winter", "LST_summer", "LST_annual", "LST_anomaly"],
    "mixed_16": ["TPI", "ROUGHNESS", "REL_RELIEF", "ELEV_COP", "SLOPE", "LST_winter", "LST_summer", "NDWI", "NDVI", "B2", "B3", "B4", "B8", "B12", "NDMI", "BSI"],
    "top_20": ["TPI", "ROUGHNESS", "LST_winter", "NDWI", "LST_summer", "REL_RELIEF", "ASPECT", "LST_anomaly", "ELEV_COP", "B2", "HILLSHADE", "B8A", "NDMI", "B12", "B3", "B4", "B8", "NDVI", "SLOPE", "CURVATURE"]
}

results = []

for exp, bands in experiments.items():
    if exp == "all":
        ckpt_dir = "all_data/checkpoint"
        res_dir = "all_data/result"
        n_bands = 35
    else:
        ckpt_dir = f"all_data/checkpoint_{exp}band"
        res_dir = f"all_data/result_{exp}band"
        n_bands = len(bands)
    
    # Parse val metrics
    try:
        df = pd.read_csv(f"{ckpt_dir}/training_history.csv")
        best_row = df.loc[df['val_dice_coef'].idxmax()]
        val_dice = best_row['val_dice_coef']
        val_iou = best_row['val_iou_m']
    except Exception as e:
        val_dice = 0
        val_iou = 0
        
    # Parse test metrics
    test_dice, test_iou, precision, recall = 0, 0, 0, 0
    try:
        with open(f"{res_dir}/evaluation_summary.txt", "r") as f:
            for line in f:
                if "dice      :" in line or "dice :" in line:
                    test_dice = float(line.split(":")[1].strip())
                elif "iou       :" in line or "iou :" in line:
                    test_iou = float(line.split(":")[1].strip())
                elif "precision :" in line:
                    precision = float(line.split(":")[1].strip())
                elif "recall    :" in line or "recall :" in line:
                    recall = float(line.split(":")[1].strip())
    except Exception as e:
        pass
        
    results.append({
        "run_name": exp,
        "n_bands": n_bands,
        "bands": ",".join(bands),
        "val_dice": val_dice,
        "val_iou": val_iou,
        "test_dice": test_dice,
        "test_iou": test_iou,
        "precision": precision,
        "recall": recall
    })

# Sort by test_dice descending
results.sort(key=lambda x: x['test_dice'], reverse=True)

# Print as markdown table
print("| run_name | n_bands | val_dice | val_iou | test_dice | test_iou | precision | recall | bands |")
print("|----------|---------|----------|---------|-----------|----------|-----------|--------|-------|")
for r in results:
    print(f"| {r['run_name']} | {r['n_bands']} | {r['val_dice']:.4f} | {r['val_iou']:.4f} | {r['test_dice']:.4f} | {r['test_iou']:.4f} | {r['precision']:.4f} | {r['recall']:.4f} | {r['bands']} |")