"""
DISTRIBUTION STATEMENT A. Approved for public release. Distribution is unlimited.

This material is based upon work supported by the Under Secretary of Defense for
Research and Engineering under Air Force Contract No. FA8702-15-D-0001. Any opinions,
findings, conclusions or recommendations expressed in this material are those of the
author(s) and do not necessarily reflect the views of the Under Secretary of Defense
for Research and Engineering.

© 2023 Massachusetts Institute of Technology.

The software/firmware is provided to you on an As-Is basis

Delivered to the U.S. Government with Unlimited Rights, as defined in DFARS Part
252.227-7013 or 7014 (Feb 2014). Notwithstanding any copyright notice, U.S. Government
rights in this work are defined by DFARS 252.227-7013 or DFARS 252.227-7014 as detailed
above. Use of this work other than as specifically authorized by the U.S. Government
may violate any copyrights that exist in this work.
"""
import  numpy  as  np
import  cirq   as  cirq

from    pyLIQTR.clam.operator_strings              import  op_strings
from    pyLIQTR.ProblemInstances.ProblemInstance   import  ProblemInstance
from    openfermion.chem                           import  MolecularData
from    openfermion                                import  jordan_wigner
from    openfermion                                import  InteractionOperator

from    functools   import   cache

class ChemicalHamiltonian(ProblemInstance):
    
    """
    This ProblemInstance generates information for the PauliLCU encoding
    of a molecular Hamiltonian.
    
    Args:
        - mol_ham: molecular Hamiltonian in the form of an openfermion
            InteractionOperator. This object can be generated by 
            molecule.get_molecular_hamiltonian().
        - mol_name (Optional): str, name or molecular formula.
        
    """

    def __init__(self, mol_ham: InteractionOperator, mol_name=None, **kwargs):
        self._mol_ham   =  mol_ham
        self._mol_name  =  mol_name

        super(ProblemInstance, self).__init__(**kwargs)

        self._terms_jw  =  jordan_wigner(self._mol_ham).terms
        self._ops       =  op_strings( N_qb=self.n_qubits() )
                

        ct = 0
        for term in self._terms_jw:

            bits    =  []
            ops     =  []
            coeff   =  self._terms_jw[term]

            for p in term:
                op = p[1]
                qubit = p[0]
                bits.append(qubit)
                ops.append(op)

            bits = tuple(bits)

            op_str = ''
            for op in ops:
                op_str += op

            if (ct >= 1):
                self._ops.append_tuple( (bits,op_str,coeff) )
            ct += 1



    def __str__(self):
        if self._mol_name == None:
            return "Molecule"
        else:
            return f"{self._mol_name}"
    


    def n_qubits(self):
        return self._mol_ham.n_qubits

    
#  #   @property
#     def n_terms(self,**kwargs):
#         return len(list(self._terms_jw))



    @property
    @cache
    def lam(self):
        lam = 0
        for term in self._terms_jw:
            coeff  =  self._terms_jw[term]
            lam   +=  abs(coeff)
        return lam
    
    def get_alpha(self):
        return(self._ops.get_alpha())


    # def isPowerOfTwo(self,n):
    #     return (np.ceil(np.log2(n)) == math.floor(np.log2(n)));



    def yield_PauliLCU_Info(self,return_as='arrays',do_pad=0,pad_value=1.0):
        if (return_as == 'arrays'):
            terms = self._ops.terms(do_pad=do_pad,pad_value=pad_value)
        elif (return_as == 'strings'):
            terms = self._ops.strings(do_pad=do_pad,pad_value=pad_value)

        for term in terms:
            yield term

