#¡/usr/bin/env python

import os
import subprocess
import re
import random 
import argparse
from tqdm import tqdm

parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str, default="scripts/pytorch_script_rex_mask_spectral.py")
parser.add_argument("--config", type=str, default="rex_spectral.toml")
parser.add_argument("--database", type=str, default="ECG-Spectral/Results/rex_mask_seed_42_threshold_0.9.db")
parser.add_argument("--dir", type=str, default="ECG-Spectral/Spectral/exp_test_two_class")
parser.add_argument("--output_dir", type=str, default="ECG-Spectral/Results/ResNet_masking_seed_42_threshold_0.9")

args = parser.parse_args()

out = os.path.join(args.output_dir, args.output_dir.rsplit("/", 1)[-1] + ".csv")

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

random.seed(42)
all_files = []
for root, dirs, files in os.walk(args.dir):
    for file in files:
        if file.endswith('.npy'):
            all_files.append(os.path.join(root, file))

with open(out, "a") as e:
    e.write(f"Filename, actual classification, predicted classification, area, responsibility entropy, max entropy, insertion curve, deletion curve\n")

for name in tqdm(all_files):
        fp = name
        rr = None
        process = subprocess.Popen(['ReX', fp, "--script", args.model, "--config", args.config, "--analyse", "--database", args.database, '--output', f'{args.output_dir}/{fp.split("/")[-1].split(".npy")[0]}.png'], stdout=subprocess.PIPE)
        for line in process.stdout: #type: ignore
            line = line.decode().strip()
            # print(line)
            actual_class = fp.split("/")[-1].rsplit("_class_",1)[1].split(".npy")[0]
            rr = re.findall(r"[+-]? *(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?", line)
            rr = [actual_class] + rr
            rr = ",".join([r.strip() for r in rr])
        with open(out, "a") as e:
            e.write(f"{fp.split('/')[-1]},{rr}\n")
        # print("finished:", fp, rr, flush=True)

