import os
import json
import shutil
import re
from pathlib import Path
import pandas as pd

BASE_DIR = Path("/work/ratul1/supantha/glycan-SD-VS/bert_training_v3/v3.1_cluster_training")
CONSOLIDATED_ROOT = BASE_DIR / "bert_training_v4" / "downstream_tasks" / "consolidated_results"
MODES = ["strict", "strict_3"]
TASKS = ["domain", "kingdom", "phylum", "class", "order", "family", "genus", "species", "immunogenicity", "link"]

def copy_file(src, dst):
    if src.exists():
        os.makedirs(dst.parent, exist_ok=True)
        shutil.copy2(src, dst)
        return True
    return False

def save_json(data, path):
    os.makedirs(path.parent, exist_ok=True)
    with open(path, 'w') as f:
        json.dump(data, f, indent=4)

def parse_gaa_log(log_path, task, mode):
    try:
        content = log_path.read_text()
    except:
        return None
    metrics = {"status": "completed", "task": task, "mode": mode}
    found = False
    for line in content.splitlines():
        match = re.search(r"(\w+)\s*\[([\w-]+)\]:\s*([\d\.\-]+)", line)
        if match:
            m_name, m_task, m_val = match.groups()
            if m_task == task or m_task == "average" or m_task == task.replace("_", "-"):
                try: metrics[m_name.lower()] = float(m_val); found = True
                except: pass
    if "macrof1" in metrics: metrics["test_f1_macro"] = metrics["macrof1"]
    if "auprc" in metrics: metrics["test_auprc"] = metrics["auprc"]
    return metrics if found else None

def main():
    print(f"Starting consolidation in: {CONSOLIDATED_ROOT}")
    for mode in MODES:
        print(f"\nProcessing mode: {mode}")
        
        # 1. OURS
        ours_src_root = BASE_DIR / "bert_training_v4" / "downstream_tasks" / "results_v4_BERT_topology_FINAL"
        for task in TASKS:
            src = ours_src_root / f"{task}_{mode}" / "results.json"
            dst = CONSOLIDATED_ROOT / mode / "ours" / f"{task}.json"
            if not copy_file(src, dst):
                src = ours_src_root / f"{task}" / "results.json"
                copy_file(src, dst)

        # 2. SweetTalk
        st_root = BASE_DIR / "bert_training_v4" / "sweettalk"
        st_dir = f"results_{mode}" if mode == "strict_3" else "results"
        for task in TASKS:
            src = st_root / st_dir / task / f"results_{task}.csv"
            if src.exists():
                df = pd.read_csv(src)
                metrics = {"status": "completed", "task": task, "mode": mode}
                for _, row in df.iterrows():
                    if "accuracy" in row: metrics["test_accuracy"] = row["accuracy"]
                    if "macro_f1" in row: metrics["test_f1_macro"] = row["macro_f1"]
                save_json(metrics, CONSOLIDATED_ROOT / mode / "sweettalk" / f"{task}.json")

        # 3. GlycanML (Corrected Path)
        ml_base = BASE_DIR / "baseline_comparisons" / "results" / "glycanml_paper_batch256" / mode / "single_task" / "PropertyPrediction" / "GlycanClassification"
        if ml_base.exists():
            for model_dir in ml_base.iterdir():
                if model_dir.is_dir():
                    model = model_dir.name
                    # Find any status.json in subfolders
                    for date_dir in model_dir.iterdir():
                        if date_dir.is_dir():
                            src = date_dir / "status.json"
                            if src.exists():
                                with open(src, 'r') as f:
                                    data = json.load(f)
                                    task = data.get("task", "").lower()
                                    if task in TASKS:
                                        dst = CONSOLIDATED_ROOT / mode / "glycanml" / model / f"{task}.json"
                                        save_json(data, dst)

        # 4. GlycanAA / GearNet
        gaa_logs_dir = BASE_DIR / "bert_training_v4" / "GlycanAA" / "logs"
        if gaa_logs_dir.exists():
            for log_file in gaa_logs_dir.glob("*.out"):
                name = log_file.stem
                if mode in name:
                    parts = name.split("_")
                    if len(parts) >= 5:
                        idx = parts.index("3") + 1 if "3" in parts else parts.index("strict") + 1
                        model = parts[idx]
                        task = parts[-2]
                        if task in TASKS:
                            metrics = parse_gaa_log(log_file, task, mode)
                            if metrics:
                                save_json(metrics, CONSOLIDATED_ROOT / mode / "glycanaa" / model / f"{task}.json")

        # 5. GIFFLAR
        gif_root = BASE_DIR / "baseline_comparisons" / "results" / f"gifflar_{mode}"
        if gif_root.exists():
            for task_dir in gif_root.iterdir():
                if task_dir.is_dir():
                    task = task_dir.name
                    src = task_dir / "results.json"
                    dst = CONSOLIDATED_ROOT / mode / "gifflar" / f"{task}.json"
                    copy_file(src, dst)

        # 6. GlycanGT
        gt_root = BASE_DIR / "bert_training_v4" / "GlycanGT" / "analysis" / "2_linear_probing"
        if gt_root.exists():
            for type_dir in gt_root.iterdir():
                if type_dir.is_dir():
                    model_base = f"GlycanGT_{type_dir.name}"
                    for variant_dir in type_dir.iterdir():
                        if variant_dir.is_dir():
                            for csv_file in variant_dir.glob("*_results.csv"):
                                try:
                                    df = pd.read_csv(csv_file)
                                    for task in TASKS:
                                        row = df[(df['Task'] == task) & (df['Classifier'] == 'SVM_HPO') & (df['Metric'] == 'Macro-F1')]
                                        if not row.empty:
                                            metrics = {"status": "completed", "task": task, "mode": mode, "test_f1_macro": float(row['Mean_Score'].iloc[0])}
                                            save_json(metrics, CONSOLIDATED_ROOT / mode / "glycangt" / f"{model_base}_{variant_dir.name}" / f"{task}.json")
                                except: pass

    print("\nConsolidation complete.")

if __name__ == "__main__":
    main()