#!/usr/bin/env -S uv run
import argparse
import subprocess
from pathlib import Path
import time

import numpy as np

from kyupy import verilog, bench, log, logic, batchrange, atalanta, stil
from kyupy.techlib import techlib_by_name, KYUPY

from fsim.static import LineRoles, FaultSet
from fsim.simple import SAFSimSimple
from fsim.incremental import SAFSimIncremental


def main():

    parser = argparse.ArgumentParser(description='A basic stuck-at fault simulator.')
    parser.add_argument('-t', '--tlib', default='SKY130', help=f'Techlib for verilog circuit. Default: SKY130, available: {sorted(techlib_by_name.keys())}.')
    parser.add_argument('-p', '--patterns', default='1024', help='Pattern file or number of random patterns to simulate. Default: 1024.')
    parser.add_argument('-o', '--output', default=None, help='')
    parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility. Default: 42.')
    parser.add_argument('circuit', help='Gate-level verilog, bench, or nix package to import. See available packages: "nix flake show github:s-holst/benchmark-circuits".')
    args = parser.parse_args()

    if not (circuit_path := Path(args.circuit)).exists(): # fallback to published nix package.
        nix_cmd = f"nix build github:s-holst/benchmark-circuits#{args.circuit} --print-out-paths --no-link"
        benchmark_path = Path(subprocess.check_output(nix_cmd.split(), text=True).strip())
        circuit_path = next(benchmark_path.glob("*/nl/*.nl.v"))

    log.info(f'loading {circuit_path} ...')
    if circuit_path.name.endswith('.bench'):
        args.tlib = techlib_by_name['KYUPY']
        c = bench.load(circuit_path)
    else:
        args.tlib = techlib_by_name[args.tlib]
        c = verilog.load(circuit_path, tlib=args.tlib)

    stats = {k.replace('__',''): v for k, v in c.stats(args.tlib).items() if k.startswith('__') or k.endswith('put')}
    log.info(f'circuit {stats=}')

    lr = LineRoles(c, args.tlib)
    log.info(f'line role stats={lr.stats}')

    c_resolved = c.copy()
    c_resolved.resolve_tlib_cells(args.tlib)

    fs = FaultSet(c, args.tlib, c_resolved)
    log.info(f'fault sites: {len(fs.fault_sites)}')
    log.info(f'uncollapsed stuck-at fault count: {len(fs.saf_set)}')
    log.info(f'collapsed stuck-at fault count: {len(fs.saf_equiv_classes)}')

    ffr_stems = []
    for stem, _ in c_resolved.fanout_free_regions(KYUPY):
        if len(stem.outs) > 0 and stem.outs[0] is not None:
            ffr_stems.append(stem.outs[0])
    ffr_stems = np.array(ffr_stems, dtype=np.uint32)

    log.info(f'FFR count: {len(ffr_stems)}')

    rng = np.random.default_rng(args.seed)

    if not (tests_path := Path(args.patterns)).exists(): # fallback to random patterns.
        patterns = rng.choice(
                    [logic.ZERO, logic.ONE],
                    size=(len(c_resolved.s_nodes(KYUPY)), int(args.patterns)),
                ).astype(np.uint8)
    else:
        log.info(f'loading {tests_path} ...')
        if tests_path.name.endswith('.stil'):
            patterns = stil.load(tests_path).tests(c_resolved)
        else:
            patterns = atalanta.load(tests_path).tests(c_resolved)

    saf_collapsed = np.array(list(fs.saf_equiv_classes.keys()), dtype=np.uint32)
    rng.shuffle(saf_collapsed)

    safsim = SAFSimSimple(c_resolved, min(patterns.shape[1], 10240))
    log.info(f'{safsim.sim=}')

    fclasses = safsim.classify_faults(saf_collapsed, patterns)

    log.info(f'{safsim.timers=}')
    sim_performance = stats['comb'] * len(saf_collapsed) * patterns.shape[1] / safsim.timers['sim'].s
    log.info(f'fsim performance: {sim_performance:.2e} gfp/s')
    log.info(f'detected by simulation (collapsed): {len(fclasses["DS"])}/{len(saf_collapsed)}  -  {len(fclasses["DS"])/len(saf_collapsed)*100:.2f}%')

    if args.output is not None:
        out_path = Path(args.output)
        with open(out_path, 'w') as f:
            log.info(f'Writing {out_path.absolute()} ...')
            for flt_rep in fclasses["NO"]:
                for flt in fs.saf_equiv_classes[flt_rep]:
                    f.write(f'{fs.fault_type_str(flt)}\tNO\t{fs.fault_site_str(c, args.tlib, flt)}\n')
            for flt_rep in fclasses["DS"]:
                for flt in fs.saf_equiv_classes[flt_rep]:
                    f.write(f'{fs.fault_type_str(flt)}\tDS\t{fs.fault_site_str(c, args.tlib, flt)}\n')

if __name__ == "__main__":
    main()