#¡/usr/bin/env python

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

parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str, default="scripts/pytorch_script_rex_AD_PVOC_resnet.py")
parser.add_argument("--config", type=str, default="toml_files_resnet/PVOC/rex_AD_PVOC_resnet.toml")
parser.add_argument("--database", type=str, default="PASCAL-VOC/Results_resnet/rex_AD_seed_42_threshold_0.9.db")
parser.add_argument("--dir", type=str, default="PASCAL-VOC/exp_test")
parser.add_argument("--output_dir", type=str, default="PASCAL-VOC/Results_resnet/ResNet_AD_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)

if not args.pickle:
    random.seed(42)
    all_files = []
    for root, dirs, files in os.walk(args.dir):
        for file in files:
            if file.endswith('.jpeg') or file.endswith('.jpg'):
                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(".jpeg")[0]}.png'], stdout=subprocess.PIPE)
        for line in process.stdout: #type: ignore
            line = line.decode().strip()
            # print(line)
            rr = re.findall(r"[+-]? *(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?", line)

            found_class = False
            new_rr = []
            for r in rr:
                if found_class:
                     new_rr.append(r)
                elif r.endswith("."):
                     found_class = True
                     new_rr.append(r.strip("."))

            rr = ",".join([r.strip() for r in new_rr])
        with open(out, "a") as e:
            e.write(f"{fp},{rr}\n")
        # print("finished:", fp, rr, flush=True)

