import os
import argparse
import glob
import statistics

SECTION_TASKS = {
    1: ['put_rubbish_in_bin', 'move_tomatoes_in_plate'],
    2: ['open_drawer', "open_drawer_2"],
    3: ['auxiliary_manipulation_1', 'auxiliary_manipulation_2'],
    4: ['long_horizon_1', 'long_horizon_2'],
    5: ['put_rubbish_in_bin', 'move_tomatoes_in_plate', 'open_drawer',"open_drawer_2", 'auxiliary_manipulation_1', 'auxiliary_manipulation_2', 'long_horizon_1', 'long_horizon_2'],
    6: ["___gpt_o4_long_horizon_1", "___gpt_o4_long_horizon_2"],
    7: ["___gpt_4_1_long_horizon_1", "___gpt_4_1_long_horizon_2"],
    8: ['move_item_1', 'move_item_2'],
    9: ['open_drawer_1', 'open_drawer_2'],
    10: ['___gpt_o3_long_horizon_1', '___gpt_o3_long_horizon_2'],
    11: ['_llama_long_horizon_1', '_llama_long_horizon_2'],
    12: ["switch_move_item_1","switch_open_drawer_1"],
    13: ['long_horizon_1', 'long_horizon_2','switch_move_item_1','switch_open_drawer_1',  'open_drawer_1', 'open_drawer_2', 'move_item_1', 'move_item_2'],
    14: ['autogen_open_drawer', "autogen_open_drawer_2"],
    15: ["_llama_8b_long_horizon_1", "_llama_8b_long_horizon_2"],
    16: ["_qwen_30b_long_horizon_1", "_qwen_30b_long_horizon_2"],
}

SECTION_NAMES = {
    1: 'Object_Relocation',
    2: 'Object_Interaction',
    3: 'Auxiliary_Manipulation',
    4: 'Long_Horizon',
    5: 'All_Tasks',
    6: 'GPT_o4_mini',
    7: 'GPT_4_1',
    8: 'Realworld_Object_Relocation',
    9: 'Realworld_Object_Interaction',
    10: 'GPT_o3',
    11: 'Llama_1b',
    12: 'Auxiliary_Manipulation',
    13: 'All_Tasks',
    14: 'autogen_Object_Interaction',
    15: 'Llama_8b',
    16: 'Qwen_30b',
}

def parse_number(val_str: str) -> float:
    if '/' in val_str:
        numerator, denominator = val_str.split('/')
        return float(numerator) / float(denominator)
    else:
        return float(val_str)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--env", nargs='+', required=True)
    parser.add_argument("--type", type=int, choices=list(range(1, 18)), required=True)
    parser.add_argument("--task", type=int, choices=list(range(1, 18)), required=True)
    parser.add_argument("--eval", choices=["all", "hl", "o"], default="all")
    args = parser.parse_args()

    task_name_array = SECTION_TASKS.get(args.task, [])
    section_name = SECTION_NAMES.get(args.task, f"Section_{args.task}")

    for env_name in args.env:
        aggregated = {"high": {}, "middle": {}, "low": {}, "oracle": {}}
        base_env_dir = os.path.join("Input_Your_Path", env_name)
        task_pattern = os.path.join(base_env_dir, f"*_type{args.type}")
        task_dirs = glob.glob(task_pattern)
        task_dirs = [d for d in task_dirs if os.path.basename(d).split('_type')[0] in task_name_array]

        if not task_dirs:
            print(f"[WARNING] No directories for section '{section_name}' type{args.type} in {env_name}.")
            continue

        for task_dir in task_dirs:
            result_file_path = os.path.join(task_dir, "result.txt")
            if not os.path.isfile(result_file_path):
                print(f"[WARNING] result.txt not found in {task_dir}.")
                continue

            with open(result_file_path, 'r') as f:
                lines = f.readlines()

            for line in lines:
                line = line.strip()
                if not line:
                    continue
                try:
                    left, right = line.split('->')
                    sc_str, gc_str = right.strip().split()
                    sc = parse_number(sc_str)
                    gc = parse_number(gc_str)

                    parts = left.strip().split('_')
                    obs_type = parts[0]
                    seed_part = parts[1]
                    if not seed_part.startswith('seed'):
                        print(f"[WARNING] Unexpected seed format: {seed_part}")
                        continue
                    seed = int(seed_part.replace('seed', ''))

                    if obs_type in aggregated:
                        aggregated[obs_type].setdefault(seed, {"sc": [], "gc": []})
                        aggregated[obs_type][seed]["sc"].append(sc)
                        aggregated[obs_type][seed]["gc"].append(gc)
                except Exception as e:
                    print(f"[ERROR] Failed to parse line: '{line}'. Error: {e}")

        summary_lines = []
        if args.eval == "hl":
            seeds = sorted(set(aggregated["high"]) | set(aggregated["low"]))
            seed_sc_means, seed_gc_means = [], []
            for seed in seeds:
                sc_vals = aggregated["high"].get(seed, {}).get("sc", []) + \
                          aggregated["low"].get(seed, {}).get("sc", [])
                gc_vals = aggregated["high"].get(seed, {}).get("gc", []) + \
                          aggregated["low"].get(seed, {}).get("gc", [])
                sc_mean = statistics.mean(sc_vals) if sc_vals else 0.0
                gc_mean = statistics.mean(gc_vals) if gc_vals else 0.0
                seed_sc_means.append(sc_mean)
                seed_gc_means.append(gc_mean)

            sc_avg  = statistics.mean(seed_sc_means) if seed_sc_means else 0.0
            gc_avg  = statistics.mean(seed_gc_means) if seed_gc_means else 0.0
            sc_std  = statistics.stdev(seed_sc_means) if len(seed_sc_means) > 1 else 0.0
            gc_std  = statistics.stdev(seed_gc_means) if len(seed_gc_means) > 1 else 0.0
            summary_lines.append(
                f"high+low: sc_avg={sc_avg:.2f}, sc_std={sc_std:.2f} | gc_avg={gc_avg:.2f}, gc_std={gc_std:.2f}"
            )
        elif args.eval == "o":
            seeds = sorted(aggregated["oracle"])
            seed_sc_means, seed_gc_means = [], []
            for seed in seeds:
                sc_vals = aggregated["oracle"][seed]["sc"]
                gc_vals = aggregated["oracle"][seed]["gc"]
                sc_mean = statistics.mean(sc_vals) if sc_vals else 0.0
                gc_mean = statistics.mean(gc_vals) if gc_vals else 0.0
                seed_sc_means.append(sc_mean)
                seed_gc_means.append(gc_mean)

            sc_avg = statistics.mean(seed_sc_means) if seed_sc_means else 0.0
            gc_avg = statistics.mean(seed_gc_means) if seed_gc_means else 0.0
            sc_std = statistics.stdev(seed_sc_means) if len(seed_sc_means) > 1 else 0.0
            gc_std = statistics.stdev(seed_gc_means) if len(seed_gc_means) > 1 else 0.0
            summary_lines.append(
                f"oracle: sc_avg={sc_avg:.2f}, sc_std={sc_std:.2f} | gc_avg={gc_avg:.2f}, gc_std={gc_std:.2f}"
            )
        else:
            for obs, seeds_dict in aggregated.items():
                if not seeds_dict: 
                    continue
                all_sc, all_gc, sc_seed_means, gc_seed_means = [], [], [], []
                for seed, metrics in seeds_dict.items():
                    mean_sc = statistics.mean(metrics["sc"])
                    mean_gc = statistics.mean(metrics["gc"])
                    sc_seed_means.append(mean_sc)
                    gc_seed_means.append(mean_gc)
                    all_sc.extend(metrics["sc"])
                    all_gc.extend(metrics["gc"])

                overall_sc_avg = statistics.mean(all_sc)
                overall_gc_avg = statistics.mean(all_gc)
                overall_sc_std = statistics.stdev(sc_seed_means) if len(sc_seed_means) > 1 else 0.0
                overall_gc_std = statistics.stdev(gc_seed_means) if len(gc_seed_means) > 1 else 0.0
                summary_lines.append(
                    f"{obs}: sc_avg={overall_sc_avg:.2f}, sc_std={overall_sc_std:.2f} | gc_avg={overall_gc_avg:.2f}, gc_std={overall_gc_std:.2f}"
                )

        result_dir = os.path.join(base_env_dir, 'result')
        os.makedirs(result_dir, exist_ok=True)
        summary_file_path = os.path.join(result_dir, f"{section_name}_type{args.type}_{args.eval}.txt")
        with open(summary_file_path, 'w') as f:
            f.write("\n".join(summary_lines))

        print(f"[INFO] Summary saved to {summary_file_path} for section '{section_name}' type{args.type} in {env_name}.")

if __name__ == '__main__':
    main()
