from __future__ import print_function
import unittest
import numpy as np
import sqaod as sq
import sqaod.common as common
from tests.example_problems import *
from math import log, exp

class TestBipartiteGraphAnnealerBase:

    def __init__(self, anpkg, dtype) :
        self.anpkg = anpkg
        self.dtype = dtype
        self.epu = 1.e-6 if dtype == np.float32 else 1.e-12

    def new_annealer(self, N0, N1, m) :
        an = self.anpkg.bipartite_graph_annealer(dtype=self.dtype)
        b0, b1, W = bipartite_graph_random(N0, N1, self.dtype)
        an.set_qubo(b0, b1, W)
        an.set_preferences(n_trotters = m)
        return an

    def test_calling_sequence(self) :
        N0, N1 = 200, 100
        m = 100
        an = self.new_annealer(N0, N1, m)
        an.set_preferences(algorithm = sq.algorithm.default)
        an.seed(0)
        an.prepare()
        an.randomize_spin()
        an.anneal_one_step(1, 1)
        an.calculate_E()
        an.make_solution()
        an.get_E()
        an.get_problem_size()
        an.get_preferences()
        xpair = an.get_x()
        h0, h1, J, c = an.get_hamiltonian()
        qpair = an.get_q()

    def test_problem_size(self) :
        N0, N1 = 100, 110
        m = 10
        an = self.new_annealer(N0, N1, m)
        N0out, N1out = an.get_problem_size()
        self.assertEqual(N0, N0out)
        self.assertEqual(N1, N1out)

    def test_set_n_trotters(self) :
        an = self.new_annealer(10, 10, 10)
        an.set_preferences(n_trotters = 3)
        prefs = an.get_preferences()
        self.assertEqual(prefs['n_trotters'], 3)
        
    def test_get_hamiltonian(self) :
        N0, N1 = 10, 11
        m = 10
        an = self.anpkg.bipartite_graph_annealer(dtype=self.dtype)