#!/usr/bin/env python
# coding: utf-8

# # Running germ selection on multiple processors
# The code below should be put into a script and run using mpiexec.  It's primary function is to pass a MPI Comm object to `pygsti.algorithms.germselection.build_up_breadth`.

# In[1]:


from __future__ import print_function
import time

import pygsti
import pygsti.construction as pc
from pygsti.construction import std2Q_XYICNOT
from pygsti.algorithms import germselection as germsel

from mpi4py import MPI
comm = MPI.COMM_WORLD

def do_greedy_germsel(target_model, forced_germs, candidate_counts,
                      seedStart, outFilename, comm):
    #candidate_counts is a dict of keys = germ lengths, values = # of germs at that length                                                            

    tStart = time.time()

    candidate_germs = []
    for i,(germLength, count) in enumerate(candidate_counts.items()):
        if count == "all upto":
            candidate_germs.extend( pc.list_all_circuits_without_powers_and_cycles(
                    target_model.operations.keys(), maxLength=germLength) )
        else:
            candidate_germs.extend( pc.list_random_circuits_onelen(
                    target_model.operations.keys(), germLength, count, seed=seedStart+i))

    available_germs = pygsti.tools.remove_duplicates( forced_germs + candidate_germs )
    print("%d available germs" % len(available_germs))
    germs = germsel.build_up_breadth(target_model, available_germs,
                     randomizationStrength=1e-3, numCopies=3, seed=1234,
                     opPenalty=10.0, scoreFunc='all', tol=1e-6, threshold=1e5,
                     pretest=False, force=forced_germs, verbosity=5, comm=comm, memLimit=0.5*(1024**3))

    if comm is None or comm.Get_rank() == 0:
        print("Germs (%d) = \n" % len(germs), "\n".join(map(str,germs)))
        print("Total time = %mdl" % (time.time()-tStart))
        pickle.dump(germs,open(outFilename,"wb"))
    return germs
                                                                                                                                         
#2Q case                                                                                                                                              
target_model = std2Q_XYICNOT.target_model()
forced_germs = pygsti.construction.circuit_list([(gl,) for gl in target_model.operations.keys()]) #singletons                                                                                      
candidate_counts = { 3:"all upto", 4:30, 5:20, 6:20, 7:20, 8:20} # germLength:num_candidates                                                          
seedStart = 4
do_greedy_germsel(target_model, forced_germs, candidate_counts,
                  seedStart, "germs_EXAMPLE.pkl", comm)


# Above is **keyboard-interrupted on purpose**, as this output was produced with a single processor and it would have taken a very long time.
# 

# In[ ]: