import numpy as np

from qiskit.circuit.library import TwoLocal

from qiskit_nature.drivers import UnitsType
from qiskit_nature.drivers.second_quantization.electronic_structure_driver import MethodType
from qiskit_nature.converters.second_quantization import QubitConverter
from qiskit_nature.mappers.second_quantization import ParityMapper
from qiskit_nature.circuit.library import HartreeFock
from qiskit_nature.drivers.second_quantization import PySCFDriver
from qiskit_nature.problems.second_quantization.electronic import ElectronicStructureProblem

# ansatz

def get_ansatz_test():
    converter = QubitConverter(mapper=ParityMapper(), two_qubit_reduction=True)
    init_st = HartreeFock(4, (1,1), converter)
    
    ansatz = TwoLocal(num_qubits=4,
                      rotation_blocks='ry', 
                      entanglement_blocks='cx',
                      entanglement=[(0,1),(1,2),(1,3)],
                      reps=2,
                      name='TwoLocal',
                      initial_state=init_st)
    return ansatz

# hamiltonian

def get_hamiltonian_test(dist):
    converter = QubitConverter(mapper=ParityMapper(), two_qubit_reduction=True)
    alt = np.sqrt(dist**2 - (dist/2)**2)
    geom = f'H .0 .0 .0; H .0 .0 {dist}; H .0 {alt} {dist/2}'
    driver = PySCFDriver(atom=geom,
                         unit=UnitsType.ANGSTROM,
                         basis='sto6g',
                         spin=0,
                         charge=1,
                         method=MethodType.RHF)
    problem = ElectronicStructureProblem(driver)
    hamiltonian = problem.second_q_ops()[0]
    circuit_hamiltonian = converter.convert(hamiltonian, 2)
    return circuit_hamiltonian

# DISTANZA 0.5

   	# NUMERO 1
    
vqe_pars_05_a = [-0.11550176,  0.12474728,  2.59668356,  3.8576121,
			0.08466277, -0.08561239, 0.33366492, -0.09046052,
			-0.01809542,  0.03835604, -0.19827222,  0.61663081]

spvqe_pars_05_a = [-0.14645914, 0.24085774, 2.4493363, 4.39635615,
			  0.0412351, -0.037009016, 0.13252825, -0.08387666,
			  -0.06751748, -0.1701828, -0.44594079, 1.11174508]

   	# NUMERO 2 

vqe_pars_05_b = [0.05880473, -0.49585711,  1.38931493,  1.22058033,
			0.04743988, -0.35168815, 1.23629905,  1.31486832,
			-0.04791204, -0.12251855, -0.42800105, -0.54385726]

spvqe_pars_05_b = [-0.0524159,  -0.35065509,  1.50978572,  1.44223811,
			  0.10572807, -0.3610393, 1.09925545,  1.16282622,
			  0.0417312,  -0.03228112, -0.52763223, -0.54086796]

   	# NUMERO 3

vqe_pars_05_c = [-0.13742697,  0.06666297,  1.00157735,  1.26190522,
			0.11471872,  0.04703104, 1.76370044,  1.61734347,
			0.12522114,  0.01627249, -0.5024544, -0.2438825 ]

spvqe_pars_05_c = [-0.00932172, -0.26494813,  1.39315548,  1.32781559,
			  0.08164483, -0.23497221, 1.17428293,  1.23188309,
			  -0.00308427, -0.00491096, -0.60690943, -0.58952378]

   	# NUMERO 4

vqe_pars_05_d = [0.16870004, -0.43813479,  1.29677422,  1.25914724,
			-0.04299792, -0.35431175,1.29298694,  1.26077492,
			-0.15053771, -0.19731699, -0.51811108, -0.6556562]

spvqe_pars_05_d = [5.02394671, 9.41745152, 1.68675252, 7.85071438,
			  3.19550059, 9.42166769, 5.09519325, 2.59998515,
			  1.88735871, 3.1531819,  5.93912936, 2.11561298]

vqes_05 = [vqe_pars_05_a, vqe_pars_05_b, vqe_pars_05_c, vqe_pars_05_d]
spvqes_05 = [spvqe_pars_05_a, spvqe_pars_05_b, spvqe_pars_05_c, spvqe_pars_05_d]

# DISTANZA 1.0

   	# NUMERO 1
    
vqe_pars_10_a = [6.16740344, 3.20411669, 5.67528188 ,3.67933879,
			2.95053784, 6.2426416, 5.32863549, 6.09733424, 
			3.16169489, 0.10689151, 5.8707268,  2.40834679]

spvqe_pars_10_a = [ 5.52205337,  6.33947666,  4.99314629,  1.56646707,
			  6.42004486,  9.34654435, 1.64389086,  1.53953422, 
			  0.73439052,  9.42814649, -0.27327807,  9.44921354]


   	# NUMERO 2

vqe_pars_10_b = [-0.3512596,  -0.0434752,   1.0358355,   1.46491864 ,
			0.20682505, -0.00289901, 1.43619352,  1.23273737,  
			0.29697954 ,-0.07741656, -0.64436651, -0.4226421 ]

spvqe_pars_10_b = [-0.46588444 , 0.03577865 , 1.12642329,  1.60349416,
			  0.16878282,  0.03519882,  1.58751012,  1.25972376, 
			  0.41562211,  0.01466055, -0.38381825, -0.28252069]

   	# NUMERO 3

vqe_pars_10_c = [ 0.62969728,  0.11563254,  1.90128434,  1.46585069,
			0.18945206,  0.26035458, 0.80622198,  0.43826824, 
			-0.67497233 , 0.1055656 , -0.53464882, -1.27380774]

spvqe_pars_10_c = [ 0.5916648 ,  1.06775659 , 1.57289366 , 1.56449836,
			  0.74571656,  1.14593536, 1.64660689,  1.73306393, 
			  -0.25437019, -0.0159303 ,  0.05439476,  0.16162612]


   	# NUMERO 4

vqe_pars_10_d = [ 0.99493485,  0.46749555,  1.66718638,  1.57712555,
			 0.76606805,  0.69293642, 1.88378101,  1.78376671,
			 -0.8351959,  -0.06788017,  0.33365866,  0.21094284]

spvqe_pars_10_d = [ 1.27976371, -0.48962493, -1.35773093,  1.44466152,
			  2.0692902,  -1.85807479, -0.80085892 , 0.50904985,
			   2.18264396,  0.25188056,  0.82224704, -1.06459749]

vqes_10 = [vqe_pars_10_a, vqe_pars_10_b, vqe_pars_10_c, vqe_pars_10_d]
spvqes_10 = [spvqe_pars_10_a, spvqe_pars_10_b, spvqe_pars_10_c, spvqe_pars_10_d]

# DISTANZA 1.5    (19_04_10_47.log)

   	# NUMERO 1

vqe_pars_15_a =  [0.09991746,  0.62611588,  1.89744194,  4.32594859,
             -0.07035062, -0.31686141, 0.54348101,  4.58529642,
             0.05761779, -0.09718047, -0.71666062,  3.03730765]

spvqe_pars_15_a = [0.21380667,  0.2427542,   2.65108444,  8.19294383,
              -0.25609082,  0.09342074, 0.35346509,  3.82999585,
              -0.06775606,  0.17041844, -0.20121347,  2.58884638]

   	# NUMERO 2    
    
vqe_pars_15_b = [-0.21548311, 0.20120314, 2.30422198, 4.48161905,
			0.13964768, -0.19722156, 0.09335863, 0.23036292,
			-0.38731748, -0.07424302, -1.00785697, 0.98497606]

spvqe_pars_15_b = [0.61714577,  2.47534277,  3.17142515,  4.70746919,
			  1.35411788, -1.58468248,  1.60239857, -0.02513338,
			  -0.97090705, -1.55142321, -0.79341562,  1.60062894]

   	# NUMERO 3

vqe_pars_15_c = [0.97968128, 1.57688493, 1.1570894,  1.60995628,
			1.34060246, 1.39458515, 1.92058373, 0.14457375,
			0.15368186, 0.51329187, 0.19821282, 0.04170353]

spvqe_pars_15_c = [ 0.48976517, -0.08104914,  0.72184541, -1.75849263,
			  0.25320102,  0.03368859, 2.59738978, -0.78635837,
			  0.28449398, -0.04249177,  0.2943293,   0.62921035]

   	# NUMERO 4

vqe_pars_15_d = [ 0.84450869,  1.25491968, -0.44514397,  2.38703763,
			0.94495797,  0.15694397, 2.98751578, -0.692738 , 
			-0.11166693,  0.72361479, -0.50367291,  0.49284549]

spvqe_pars_15_d = [ 3.90221684e-01, -4.15178472e-02, -2.99633796e-01 ,-1.58264586e+00,
 			  2.13748167e-01, -5.61087214e-04,  3.11964764e+00, -7.96008889e-01,
      		 -1.51861563e-01, -5.91138869e-02, -2.71376710e-01,  7.52653356e-01]

vqes_15 = [vqe_pars_15_a, vqe_pars_15_b, vqe_pars_15_c, vqe_pars_15_d]
spvqes_15 = [spvqe_pars_15_a, spvqe_pars_15_b, spvqe_pars_15_c, spvqe_pars_15_d]

# DISTANZA 2.0

    # NUMERO 1

vqe_pars_20_a = [0.23225322,  4.20509329,  1.82170197, 4.27582049,
		    0.05729869,  0.15604462, 0.06000301, -0.13579914,
		    0.08299985, -1.07408837, -1.21879918, 1.24899782]

spvqe_pars_20_a = [1.43809554, 7.81734322, 3.17165507, 4.75015336,
              2.33703869, -0.38781625, 0.44597784, -0.12498502,
              -0.12032009, -1.55407031, -0.10911732,  1.61769171]

    # NUMERO 2

vqe_pars_20_b = [6.65022149e-03, -9.50528847e-02,  1.60330579e+00,  7.19883316e+00,
 			-1.57724593e-02, -1.29639965e+00, -2.27676219e-02,  3.24072370e+00,
 			-7.34213970e-02, -1.27990560e+00, -1.48335812e+00,  1.75404689e+00]

spvqe_pars_20_b = [0.79587883, 0.32760775,  3.35288526, 7.85590383,
			  -0.0390306, -0.69694168, -1.4426973, 3.36124759,
			  -0.01221064, 0.74481681, -1.12900567, 1.70637381]
	# NUMERO 3

vqe_pars_20_c = [-0.08454793,  5.11004342,  1.83282241,  4.5149173,
 			0.09799665, -0.57644001, -0.15790267, -0.14204233,
 			0.08811245, -1.19742568, -1.15342555,  1.45476172]

spvqe_pars_20_c = [0.11532082,  8.09680516,  2.37427009,  4.68877452,
			  1.42678813, -0.14232414, 0.17262165, -0.04088128,
			  -1.4429427,  -1.6389454,  -1.21152173,  1.53310367]

	# NUMERO 4

vqe_pars_20_d = [-0.10993742,  1.08577653, -0.05778733,  1.43690763,
			0.8346593,  -0.64097313, -0.55031438, -0.02333615,
			0.35142381,  1.30165111,  0.38822238, 1.1910818 ]


spvqe_pars_20_d = [-0.35471067,  1.3383629,   0.01617288,  1.44855139,
			  0.44430146,  1.60495214, 1.48593211,  0.08656238,
			  -0.34956568,  1.54163135, -0.18345526, -1.55837365]

vqes_20 = [vqe_pars_20_a, vqe_pars_20_b, vqe_pars_20_c, vqe_pars_20_d]
spvqes_20 = [spvqe_pars_20_a, spvqe_pars_20_b, spvqe_pars_20_c, spvqe_pars_20_d]

# DISTANZA 2.5

	# NUMERO 1
    
vqe_pars_25_a = [-0.30538821,  4.60997242,  7.66433589,  4.17372787,
			0.10927006, -0.51266663, 0.23441709, -0.334758,
			-0.22351772, -0.53313997, -1.533864,   1.46004641]

spvqe_pars_25_a = [-1.24650601,  5.78497934,  9.18265278,  4.73156897,
			  -0.08129257, -0.41884356, -0.65460038, -0.08477229,
			  -0.97708064,  0.75940663, -1.0149008,   1.42565896]

	# NUMERO 2

vqe_pars_25_b = [-0.57303308,  4.98891235,  1.94492016 , 4.53017413,
			0.62503729, -0.61011673, -0.20483872, -0.11683771,
			0.10371051, -1.2520706,  -0.91139273, 1.4457862 ]

spvqe_pars_25_b = [-6.40316565e-01,  8.15161506e+00,  3.06074748e+00,  4.67190248e+00,
  			  6.19902164e-01, -1.27453265e+00,  1.29407199e+00, -6.22156917e-03,
              -7.40268451e-01, -1.58162040e+00, -3.93077842e-01,  1.55191480e+00]

	# NUMERO 3

vqe_pars_25_c = [ 2.71938179e-02,  5.12982224e-01,  1.17797400e+00, -9.06135220e-02,
  			 7.21795803e-04, -3.56296414e-01,  1.32625535e+00,  6.68248236e-03,
  			 8.74291139e-02, -1.45434574e-01, -6.86335893e-01, -1.34254368e-01]

spvqe_pars_25_c = [ 0.99843238,  0.44052177, -0.1961437,  -1.02581547,
			  -0.77302264, -0.7235378, -0.82344231,  0.26436617, 
			  1.26044456, -0.31363713, -1.2597413,   1.28040742]

	# NUMERO 4

vqe_pars_25_d = [ 0.18502226,  0.81530577,  0.73862034,  1.64122226, 
			 0.21070779,  0.89669377,  1.56170187, -0.22372298, 
			 0.18602616 , 0.63184611, -0.52464817,  1.38476156]

spvqe_pars_25_d = [ 1.53050786,  0.19805966,  0.09745304, -1.13801992, 
			  0.4664388,   0.42438396, 1.14798148,  0.63809267, 
			  2.067302,    0.01598754,  1.04750262,  1.75522881]

vqes_25 = [vqe_pars_25_a, vqe_pars_25_b, vqe_pars_25_c, vqe_pars_25_d]
spvqes_25 = [spvqe_pars_25_a, spvqe_pars_25_b, spvqe_pars_25_c, spvqe_pars_25_d]

vqe = [vqes_05, vqes_10, vqes_15, vqes_20, vqes_25]
spvqe = [spvqes_05, spvqes_10, spvqes_15, spvqes_20, spvqes_25]

from mitiq import zne
from functools import partial

from qiskit import Aer, IBMQ
from qiskit.utils import QuantumInstance
from qiskit.providers.aer.noise import NoiseModel
from qiskit.utils.mitigation import CompleteMeasFitter
from qiskit.opflow import CircuitStateFn, StateFn, PauliExpectation, CircuitSampler

#provider = IBMQ.load_account()

dist = 2.5

qubit_op = get_hamiltonian_test(dist)
my_ansatz = get_ansatz_test().assign_parameters(spvqe_pars_25_d).decompose()

def get_expectation(ansatz):
    tot = 0    
    for pauli in qubit_op:
        psi = CircuitStateFn( ansatz )  
        measurable_expression = StateFn( pauli, is_measurement = True ).compose( psi )
        expectation = PauliExpectation().convert( measurable_expression )  
        sampler_qasm = CircuitSampler( quantuminstance ).convert( expectation )
        expect_sampling_qasm = sampler_qasm.eval().real
        tot += expect_sampling_qasm
    return tot

shots = 10000
backend = Aer.get_backend('qasm_simulator' )
noise_model = NoiseModel.from_backend(provider.get_backend('ibmq_quito'))
measurement_error_mitigation_cls=CompleteMeasFitter
optimization_level = 0
quantuminstance = QuantumInstance(backend, 
                                  shots=shots,
                                  noise_model=noise_model,
                                  optimization_level=optimization_level,
                                  measurement_error_mitigation_cls=measurement_error_mitigation_cls)

scale_factors = [1.0, 2.0, 3.0, 3.5]

mitigated = zne.execute_with_zne(my_ansatz, get_expectation, 
                                 factory=zne.inference.RichardsonFactory(scale_factors=scale_factors),
                                 num_to_average=3,
                                 scale_noise=partial(zne.scaling.fold_gates_from_left))                        
not_mitigated = get_expectation(my_ansatz)

quantuminstance = QuantumInstance(Aer.get_backend('statevector_simulator'))
exact = get_expectation(my_ansatz)

print(mitigated, not_mitigated, exact)

