from util import fcidump_to_integrals
from openfermion.resource_estimates import df
from openfermion.resource_estimates.molecule import cas_to_pyscf

import pandas
import numpy as np


def load_hamiltonian(n_orbitals):
    h0, h1e, h2e = fcidump_to_integrals(f"../fcidump/CAS{n_orbitals}_min_FCIDUMP")
    return h0, h1e, h2e


def double_factorization(hamil, n_orbitals):
    _, h1e, h2e = hamil[0], hamil[1], hamil[2]
    thresh = 0.00125
    eri_rr, LR, L, Lxi = df.factorize(h2e, thresh=thresh)
    _, pyscf_mf = cas_to_pyscf(h1e, h2e, 0, 1, 1)

    total_lambda = df.compute_lambda(pyscf_mf, LR)

    print("lambda df: ", total_lambda)

    eps = 0.0016
    chi = 10
    N = 2 * n_orbitals
    beta = 16

    # Here we're using an initial calculation with a very rough estimate of the
    # number of steps to give a more accurate number of steps.
    # Then we input that into the function again.
    output = df.compute_cost(N, total_lambda, eps, L, Lxi, chi, beta, stps=20000)
    stps1 = output[0]
    output = df.compute_cost(N, total_lambda, eps, L, Lxi, chi, beta, stps1)

    toff_step = output[0]
    toff_tot = output[1]
    ancillas = output[2]

    return toff_step, toff_tot, ancillas, total_lambda


def compute_df_cost(n):

    num_orb = []
    toffolis_step = []
    toffolis_tot = []
    ancillas = []
    lambda_df = []

    try:
        data = pandas.read_csv("df_cost.csv")
        num_orb += list(data["num_orb"])
        toffolis_step += list(data["toffolis_step"])
        toffolis_tot += list(data["toffolis_tot"])
        ancillas += list(data["ancillas"])
        lambda_df += list(data["lambda_df"])
        if n in data["num_orb"]:
            return
    except FileNotFoundError:
        pass

    # if not os.path.exists(f"CAS_{n}"):
    #     create_fcidump(n)

    hamil = load_hamiltonian(n)
    cost = double_factorization(hamil, n)

    num_orb.append(n)
    toffolis_step.append(cost[0])
    toffolis_tot.append(cost[1])
    ancillas.append(cost[2])
    lambda_df.append(cost[3])
    data = {
        "num_orb": num_orb,
        "toffolis_step": toffolis_step,
        "toffolis_tot": toffolis_tot,
        "ancillas": ancillas,
        "lambda_df": lambda_df,
    }
    df = pandas.DataFrame(data=data)
    df.to_csv("df_cost.csv", index=False)


def main():
    for n in np.arange(15, 80, 5):
        compute_df_cost(n)


if __name__ == "__main__":
    main()