#!/usr/bin/env python3

"""Simulates performing an approximate modular exponentiation."""

from __future__ import annotations

import argparse
import pathlib
import sys


src_path = pathlib.Path(__file__).parent.parent / 'src'
assert src_path.exists()
sys.path.append(str(src_path))

from facto.algorithm.prep import table_str, ExecutionConfig
from scatter_script import QPU, rvalue_multi_int
from facto.algorithm import approx_modexp


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--exec_config_dir", required=True, type=str)
    parser.add_argument("--samples", default=10, type=int)
    args = parser.parse_args()

    conf = ExecutionConfig.from_data_directory(args.exec_config_dir)
    qpu = QPU(num_branches=args.samples)
    Q_exponent = qpu.alloc_quint(
        scatter=True,
        length=conf.num_input_qubits,
    )
    g = rvalue_multi_int.from_value(conf.generator, expected_count=qpu.num_branches)
    initial_value = Q_exponent.UNPHYSICAL_copy()
    result = approx_modexp(
        qpu=qpu,
        conf=conf,
        Q_exponent=Q_exponent,
    )
    assert Q_exponent == initial_value

    exact = pow(g, Q_exponent, conf.modulus)

    table = {}
    table[0-1j] = 'deviation'
    table[1-1j] = 'approx_result'
    table[2-1j] = 'exponent'
    for row, (e, a, m) in enumerate(zip(exact.UNPHYSICAL_branch_vals, result.UNPHYSICAL_branch_vals, Q_exponent.UNPHYSICAL_branch_vals)):
        err = (e - (a << conf.dropped_bits)) % conf.modulus
        err = min(err, -err % conf.modulus)
        deviation = err / conf.modulus
        table[0 + row*1j] = f'{deviation:0.3g}'
        ak = a
        dk = conf.dropped_bits
        while dk % 4 != 0:
            dk -= 1
            ak <<= 1
        table[1 + row*1j] = f'{hex(ak)} << {dk}'
        table[2 + row*1j] = hex(m)
    print(f"generator: {hex(conf.generator)}")
    print(f"modulus: {hex(conf.modulus)}")
    print(table_str(table))

    Q_exponent.UNPHYSICAL_force_del(dealloc=True)
    result.UNPHYSICAL_force_del(dealloc=True)
    qpu.verify_clean_finish()


if __name__ == '__main__':
    main()
