import os
import argparse
import subprocess
import re
import pandas as pd
from tqdm import tqdm

def list_directories(path="plp/programs"):
    return [d for d in os.listdir(path) if os.path.isdir(os.path.join(path, d))]

def get_programs(programs_file):
    with open(programs_file, 'r') as file:
        programs_names = [line.strip() for line in file if line.strip() and not line.strip().startswith(('#', '%', '//'))]
    return programs_names

def parse_output(output):
    compilation_time = 0.0

    for line in output.split('\n'):
        if line.startswith('c o Finished outputting dDNNF.'):
            times = re.findall(r"[-+]?\d*\.\d+|\d+", line)
            compilation_time = float(times[0])

    return {
        "compilation_time": compilation_time,
        "model_count": -1,
        "compression_rate": -1
    }

def parse_nnf_file(nnf_file):
    with open(nnf_file, 'r') as f:
        first_line = f.readline().strip()
        parts = first_line.split()
        num_nodes = int(parts[1])
        num_edges = int(parts[2])
    return num_nodes, num_edges

def write_broken_run(results, program_name, timeout):
    results["program"].append(program_name)
    results["config"].append("sharpsat_td")
    results["circuit_node_size"].append(-1)
    results["circuit_edge_size"].append(-1)
    results["model_count"].append(-1)
    results["compression_rate"].append(-1)
    results["compilation_time"].append(timeout)

def main():
    parser = argparse.ArgumentParser(description="Run SharpSAT-TD compilation experiments.")
    parser.add_argument("base_path", type=str, help="Path to the base directory containing program directories.")
    parser.add_argument("programs_file", nargs='?', help="File containing list of programs to process. If not provided, all programs in the base directory will be processed.")
    parser.add_argument("output_dir", type=str, help="Path to the directory where results will be saved.")
    parser.add_argument("--sharpsat_path", type=str, default="sharpsat-td/bin/", help="Path to the SharpSAT-TD binary directory.")
    parser.add_argument("--time_wall", type=int, default=1800, help="Maximum time (in seconds) for each execution.")
    args = parser.parse_args()

    base_path = args.base_path
    programs = get_programs(args.programs_file) if args.programs_file else list_directories(base_path)
    sharpsat_path = args.sharpsat_path
    time_wall = args.time_wall

    results = {
        "program": [],
        "config": [],
        "circuit_node_size": [],
        "circuit_edge_size": [],
        "model_count": [],
        "compression_rate": [],
        "compilation_time": []
    }

    time_wall_broken = False

    for program in tqdm(programs, desc="Programs"):
        if time_wall_broken:
            write_broken_run(results, program, time_wall)
            continue

        dir_path = os.path.join(base_path, program)
        program_path = os.path.join(dir_path, f"{program}.sharpsat_td")
        output_path = os.path.join(dir_path, f"{program}.nnf.sharpsat_td")
        command = f"./sharpSAT -dDNNF -decot 1 -decow 100 -tmpdir . -cs 16000 ../../{program_path} -dDNNF_out ../../{output_path}"

        try:
            output = subprocess.run(command, shell=True, text=True, capture_output=True, timeout=time_wall, cwd=sharpsat_path)
            if output.returncode != 0:
                raise subprocess.CalledProcessError(output.returncode, command)

            stats = parse_output(output.stdout)
            num_nodes, num_edges = parse_nnf_file(output_path)

            results["program"].append(program)
            results["config"].append("sharpsat_td")
            results["circuit_node_size"].append(num_nodes)
            results["circuit_edge_size"].append(num_edges)
            results["model_count"].append(stats["model_count"])
            results["compression_rate"].append(stats["compression_rate"])
            results["compilation_time"].append(stats["compilation_time"])
        except subprocess.TimeoutExpired:
            print(f"Command timed out for {program}: {command}")
            time_wall_broken = True
            write_broken_run(results, program, time_wall)
        except subprocess.CalledProcessError as e:
            print(f"Error processing {program}: {e}")
            write_broken_run(results, program, -1)

    df = pd.DataFrame(results)

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    subset_name = os.path.splitext(os.path.basename(args.programs_file))[0] if args.programs_file else "all"
    experiments_dir = os.path.join(args.output_dir, subset_name)
    if not os.path.exists(experiments_dir):
        os.makedirs(experiments_dir)

    df.to_csv(f"{experiments_dir}/sharpsat_td.csv", index=False)

if __name__ == "__main__":
    main()
