import numpy as _np
from scipy.linalg import expm as _expm

import pygsti
from pygsti.extras import interpygate as interp
from pygsti.extras.interpygate.process_tomography import run_process_tomography, unvec_square
from pygsti.tools import change_basis
from ..testutils import BaseTestCase

try:
    from mpi4py import MPI
    _comm = MPI.COMM_WORLD
    _rank = _comm.Get_rank()
    _size = _comm.Get_size()
except ImportError: # Reverted RuntimeError to cause Windows runners to break
    _comm = None
    _rank = 0
    _size = 1

mpi_workers_per_process = 1


class ExampleProcess(interp.PhysicalProcess):
    def __init__(self):
        self.Hx = _np.array([[0, 0, 0, 0],
                            [0, 0, 0, 0],
                            [0, 0, 0, -1],
                            [0, 0, 1, 0]], dtype='float')
        self.Hy = _np.array([[0, 0, 0, 0],
                            [0, 0, 0, 1],
                            [0, 0, 0, 0],
                            [0, -1, 0, 0]], dtype='float')
        self.Hz = _np.array([[0, 0, 0, 0],
                            [0, 0, -1, 0],
                            [0, 1, 0, 0],
                            [0, 0, 0, 0]], dtype='float')

        self.dephasing_generator = _np.diag([0, -1, -1, 0])
        self.decoherence_generator = _np.diag([0, -1, -1, -1])
        num_params = 6  # omega (0), phase (1), detuning (2), dephasing (3), decoherence (4), time (5)
        process_shape = (4, 4)
        super().__init__(num_params, process_shape,
                         aux_shape=(),  # a single float
                         num_params_evaluated_as_group=0)

    def advance(self, state, v, t):
        state = _np.array(state, dtype='complex')
        omega, phase, detuning, dephasing, decoherence = v

        H = (omega * _np.cos(phase) * self.Hx + omega * _np.sin(phase) * self.Hy + detuning * self.Hz)
        L = dephasing * self.dephasing_generator + decoherence * self.decoherence_generator

        process = change_basis(_expm((H + L) * t), 'pp', 'col')
        state = unvec_square(_np.dot(process, _np.outer(state, state.conj()).ravel(order='F')), 'F')
        return state

    def create_process_matrix(self, v, comm=None):

        t = v[5]
        vv = v[:5]

        def state_to_process_mxs(state):
            return self.advance(state, vv, t)
        #print(f'Calling process tomography as {comm.Get_rank()} of {comm.Get_size()} on {comm.Get_name()}.')
        process = run_process_tomography(state_to_process_mxs,
                                         n_qubits=1, basis='pp', time_dependent=False, comm=comm, verbose=False)
        return _np.array(process) if (process is not None) else None  # must return an *array* of appropriate shape

    def create_aux_info(self, v, comm=None):
        omega, phase, detuning, dephasing, decoherence, t = v
        return t * omega


class ExampleProcess_timedep(interp.PhysicalProcess):
    def __init__(self):
        self.Hx = _np.array([[0, 0, 0, 0],
                            [0, 0, 0, 0],
                            [0, 0, 0, -1],
                            [0, 0, 1, 0]], dtype='float')
        self.Hy = _np.array([[0, 0, 0, 0],
                            [0, 0, 0, 1],
                            [0, 0, 0, 0],
                            [0, -1, 0, 0]], dtype='float')
        self.Hz = _np.array([[0, 0, 0, 0],
                            [0, 0, -1, 0],
                            [0, 1, 0, 0],
                            [0, 0, 0, 0]], dtype='float')

        self.dephasing_generator = _np.diag([0, -1, -1, 0])
        self.decoherence_generator = _np.diag([0, -1, -1, -1])
        num_params = 6  # omega (0), phase (1), detuning (2), dephasing (3), decoherence (4), time (5)
        process_shape = (4, 4)
        super().__init__(num_params, process_shape,
                         aux_shape=(),  # a single float
                         num_params_evaluated_as_group=1)  # time values can be evaluated all at once

    def advance(self, state, v, times):
        state = _np.array(state, dtype='complex')
        omega, phase, detuning, dephasing, decoherence = v

        H = (omega * _np.cos(phase) * self.Hx + omega * _np.sin(phase) * self.Hy + detuning * self.Hz)
        L = dephasing * self.dephasing_generator + decoherence * self.decoherence_generator

        processes = [change_basis(_expm((H + L) * t), 'pp', 'col') for t in times]
        states = [unvec_square(_np.dot(process, _np.outer(state, state.conj()).ravel(order='F')),'F') for process in processes]

        return states

    def create_process_matrices(self, v, grouped_v, comm=None):
        times = grouped_v[0]
        def state_to_process_mxs(state):
            return self.advance(state, v, times)
        #print(f'Calling process tomography as {comm.Get_rank()} of {comm.Get_size()} on {comm.Get_name()}.')
        #print(f'DEBUG {comm.Get_rank()}: ', times, v)
        processes = run_process_tomography(state_to_process_mxs,
                                           n_qubits=1, basis='pp', time_dependent=True, comm=comm, verbose=False)
        return _np.array(processes) if (processes is not None) else None  # must return an *array* of appropriate shape

    def create_aux_infos(self, v, grouped_v, comm=None):
        omega, phase, detuning, dephasing, decoherence = v
        times = grouped_v[0]
        return _np.array([t * omega for t in times], 'd')


class InterpygateTestCase(BaseTestCase):

    def test_timedep_op(self):
        example_process = ExampleProcess_timedep()
        target_mxs = example_process.create_process_matrices(_np.array([1.0, 0.0, 0.0, 0.0, 0.0]), [[_np.pi / 2]], comm=_comm)
        if _comm is None or _comm.rank == 0:
            target_mx = target_mxs[0]
            target_op = pygsti.modelmembers.operations.StaticArbitraryOp(target_mx)
            print(target_op)
            if _comm is not None: _comm.bcast(target_op, root=0)
        else:
            target_op = _comm.bcast(None, root=0)

        param_ranges = ([(0.9, 1.1, 2),  # omega
                         (-0.1, 0, 2),   # phase
                         (-0.2, -0.1, 2),   # detuning
                         (0, 0.1, 2),    # dephasing
                         (0.1, 0.2, 2),    # decoherence
                         _np.linspace(_np.pi / 2, _np.pi / 2 + 0.5, 10)  # time
                        ])
        interp_op = interp.InterpolatedDenseOp.create_by_interpolating_physical_process(
            target_op, example_process, param_ranges, comm=_comm,
            mpi_workers_per_process=mpi_workers_per_process)

        self.assertEqual(interp_op.num_params, 6)
        interp_op.from_vector([1.1, -0.01, -0.11, 0.055, 0.155, 1.59])
        self.assertArraysAlmostEqual(_np.array([1.1, -0.01, -0.11, 0.055, 0.155, 1.59]), interp_op.to_vector())

        #print(interp_op.to_dense())
        expected = _np.array([[ 1.00000000e+00, -3.92320766e-17,  1.43373814e-18, -5.00802093e-18],
                              [ 6.11491323e-19,  7.07674657e-01,  6.02660161e-02, -9.37426062e-02],
                              [-1.58409463e-18, -7.75367332e-02, -1.56121712e-01, -7.31059113e-01],
                              [ 7.43524968e-18, -7.91093355e-02,  7.32786974e-01, -1.10942333e-01]])
        #KEVIN - check this: below is the OLD expected result, and now we're getting the above.  Is this ok?
        #expected = _np.array([[ 1.00000000e+00, -5.14632352e-17,  1.58551100e-17, -8.59219991e-18],
        #                      [ 6.10412172e-19,  7.07796561e-01,  6.01596594e-02, -9.41693123e-02],
        #                      [-1.58441909e-18, -7.76254825e-02, -1.56153689e-01, -7.30978833e-01],
        #                      [ 7.43467815e-18, -7.91773419e-02,  7.32730647e-01, -1.10922086e-01]])
        #print(interp_op.to_dense())
        self.assertArraysAlmostEqual(expected, interp_op.to_dense())

    def test_timedep_factory(self):
        class TargetOpFactory(pygsti.modelmembers.operations.OpFactory):
            def __init__(self):
                self.process = ExampleProcess_timedep()
                ss = pygsti.baseobjs.QubitSpace(1)
                pygsti.modelmembers.operations.OpFactory.__init__(self, ss, evotype="densitymx")

            def create_object(self, args=None, sslbls=None):
                assert(sslbls is None)  # don't worry about sslbls for now -- these are for factories that can create gates placed at arbitrary circuit locations
                assert(len(args) == 2)  # t (time), omega
                t, omega = args
                mx = self.process.create_process_matrices(_np.array([omega, 0.0, 0.0, 0.0, 0.0]), [[t]], comm=None)[0]
                return pygsti.modelmembers.operations.StaticArbitraryOp(mx)

        arg_ranges = [_np.linspace(_np.pi / 2, _np.pi / 2 + 0.5, 10),  # time
                      (0.9, 1.1, 2)  # omega
                      ]

        param_ranges = [(-0.1, 0.1, 2),  # phase
                        (-0.1, 0.1, 2),  # detuning
                        (0, 0.1, 2),   # dephasing
                        (0, 0.1, 2)    # decoherence
                        ]
        arg_indices = [5, 0]  # indices for time and omega within ExampleProcess's parameters (see ExampleProcess.__init__)

        example_process = ExampleProcess_timedep()
        opfactory = interp.InterpolatedOpFactory.create_by_interpolating_physical_process(
            TargetOpFactory(), example_process, arg_ranges, param_ranges, arg_indices, comm=_comm,
            mpi_workers_per_process=mpi_workers_per_process)

        self.assertEqual(opfactory.num_params, 4)
        v = _np.array([0.01, 0.01, 0.055, 0.055])
        opfactory.from_vector(v)
        self.assertArraysAlmostEqual(v, opfactory.to_vector())
        op = opfactory.create_op((1.59, 1.1))
        self.assertArraysAlmostEqual(v, op.to_vector())
        #print("Actual:")
        #print(op.to_dense())
        #print(op.aux_info)
        expected = _np.array([[ 1.00000000e+00, -1.79603810e-17, -3.02087736e-17, -7.29450879e-19],
                              [ 2.73166900e-19,  8.39404484e-01,  2.73637219e-03,  1.78744631e-02],
                              [-1.70329720e-17,  1.75441114e-02, -1.74656910e-01, -8.63668429e-01],
                              [ 7.78491911e-17,  5.98213820e-04,  8.63855305e-01, -1.31408966e-01]])
        #KEVIN - check this: below is the OLD expected result, and now we're getting the above.  Is this ok?
        #expected = _np.array([[ 1.00000000e+00, -3.67730279e-17, -4.28676069e-17,  6.20501240e-18],
        #                      [ 5.44719386e-19,  8.41211070e-01,  5.73783231e-03,  1.81516805e-02],
        #                      [-1.70671506e-17,  1.36836825e-02, -1.75019744e-01, -8.64632441e-01],
        #                      [ 7.80124509e-17, -7.41873094e-04,  8.65665135e-01, -1.31573407e-01]])
        #print("Expected:")
        #print(expected)
        self.assertArraysAlmostEqual(expected, op.to_dense())
        self.assertAlmostEqual(op.aux_info, 1.749)

    def test_timeindep_op(self):
        example_process = ExampleProcess()
        target_mx = example_process.create_process_matrix(_np.array([1.0, 0.0, 0.0, 0.0, 0.0, _np.pi / 2]), comm=_comm)
        if _comm is None or _comm.rank == 0:
            target_op = pygsti.modelmembers.operations.StaticArbitraryOp(target_mx)
            print(target_op)
            if _comm: _comm.bcast(target_op, root=0)
        else:
            target_op = _comm.bcast(None, root=0)

        param_ranges = ([(0.9, 1.1, 2),  # omega
                         (-0.1, 0, 2),   # phase
                         (-0.2, -0.1, 2),   # detuning
                         (0, 0.1, 2),    # dephasing
                         (0.1, 0.2, 2),    # decoherence
                         _np.linspace(_np.pi / 2, _np.pi / 2 + 0.5, 10)  # time
                         ])
        interp_op = interp.InterpolatedDenseOp.create_by_interpolating_physical_process(
            target_op, example_process, param_ranges, comm=_comm,
            mpi_workers_per_process=mpi_workers_per_process)
        
        self.assertEqual(interp_op.num_params, 6)
        interp_op.from_vector([1.1, -0.01, -0.11, 0.055, 0.155, 1.59])
        self.assertArraysAlmostEqual(_np.array([1.1, -0.01, -0.11, 0.055, 0.155, 1.59]), interp_op.to_vector())
        #print(interp_op.to_dense())

        expected = _np.array([[ 1.00000000e+00, -3.92320766e-17,  1.43373814e-18, -5.00802093e-18],
                              [ 6.11491323e-19,  7.07674657e-01,  6.02660161e-02, -9.37426062e-02],
                              [-1.58409463e-18, -7.75367332e-02, -1.56121712e-01, -7.31059113e-01],
                              [ 7.43524968e-18, -7.91093355e-02,  7.32786974e-01, -1.10942333e-01]])
        #KEVIN - check this: below is the OLD expected result, and now we're getting the above.  Is this ok?
        #expected = _np.array([[ 1.00000000e+00, -5.14632352e-17,  1.58551100e-17, -8.59219991e-18],
        #                      [ 6.10412172e-19,  7.07796561e-01,  6.01596594e-02, -9.41693123e-02],
        #                      [-1.58441909e-18, -7.76254825e-02, -1.56153689e-01, -7.30978833e-01],
        #                      [ 7.43467815e-18, -7.91773419e-02,  7.32730647e-01, -1.10922086e-01]])
        #print(interp_op.to_dense())
        self.assertArraysAlmostEqual(expected, interp_op.to_dense())

    def test_timeindep_factory(self):
        class TargetOpFactory(pygsti.modelmembers.operations.OpFactory):
            def __init__(self):
                self.process = ExampleProcess()
                ss = pygsti.baseobjs.QubitSpace(1)
                pygsti.modelmembers.operations.OpFactory.__init__(self, ss, evotype="densitymx")

            def create_object(self, args=None, sslbls=None):
                assert(sslbls is None)  # don't worry about sslbls for now -- these are for factories that can create gates placed at arbitrary circuit locations
                assert(len(args) == 2)  # t (time), omega
                t, omega = args
                mx = self.process.create_process_matrix(_np.array([omega, 0.0, 0.0, 0.0, 0.0, t]), comm=None)
                return pygsti.modelmembers.operations.StaticArbitraryOp(mx)

        arg_ranges = [_np.linspace(_np.pi / 2, _np.pi / 2 + 0.5, 10),  # time
                      (0.9, 1.1, 2)  # omega
                      ]

        param_ranges = [(-0.1, 0.1, 2),  # phase
                        (-0.1, 0.1, 2),  # detuning
                        (0, 0.1, 2),   # dephasing
                        (0, 0.1, 2)    # decoherence
                        ]
        arg_indices = [5, 0]  # indices for time and omega within ExampleProcess's parameters (see ExampleProcess.__init__)

        example_process = ExampleProcess()
        opfactory = interp.InterpolatedOpFactory.create_by_interpolating_physical_process(
            TargetOpFactory(), example_process, arg_ranges, param_ranges, arg_indices, comm=_comm,
            mpi_workers_per_process=mpi_workers_per_process)

        self.assertEqual(opfactory.num_params, 4)
        v = _np.array([0.01, 0.01, 0.055, 0.055])
        opfactory.from_vector(v)
        self.assertArraysAlmostEqual(v, opfactory.to_vector())
        op = opfactory.create_op((1.59, 1.1))
        self.assertArraysAlmostEqual(v, op.to_vector())
        #print(op.to_dense())
        #print(op.aux_info)
        expected = _np.array([[ 1.00000000e+00, -1.79603810e-17, -3.02087736e-17, -7.29450879e-19],
                              [ 2.73166900e-19,  8.39404484e-01,  2.73637219e-03,  1.78744631e-02],
                              [-1.70329720e-17,  1.75441114e-02, -1.74656910e-01, -8.63668429e-01],
                              [ 7.78491911e-17,  5.98213820e-04,  8.63855305e-01, -1.31408966e-01]])
        #KEVIN - check this: below is the OLD expected result, and now we're getting the above.  Is this ok?
        #expected = _np.array([[ 1.00000000e+00, -3.67730279e-17, -4.28676069e-17,  6.20501240e-18],
        #                      [ 5.44719386e-19,  8.41211070e-01,  5.73783231e-03,  1.81516805e-02],
        #                      [-1.70671506e-17,  1.36836825e-02, -1.75019744e-01, -8.64632441e-01],
        #                      [ 7.80124509e-17, -7.41873094e-04,  8.65665135e-01, -1.31573407e-01]])
        self.assertArraysAlmostEqual(expected, op.to_dense())
        self.assertAlmostEqual(op.aux_info, 1.749)

    def test_process_tomography(self):
        """ Demonstrate the process tomography function with (potentially) time-dependent outputs. """
        sigI = _np.array([[1, 0], [0, 1]], dtype='complex')
        sigX = _np.array([[0, 1], [1, 0]], dtype='complex')
        sigY = _np.array([[0, -1.j], [1.j, 0]], dtype='complex')
        sigZ = _np.array([[1, 0], [0, -1]], dtype='complex')
        theta = 0.32723
        u = _np.cos(theta) * sigI + 1.j * _np.sin(theta) * sigX
        v = _np.sin(theta) * sigI - 1.j * _np.cos(theta) * sigX

        U = _np.kron(u, v)
        test_process = _np.kron(U.conj().T, U)

        def single_time_test_function(pure_state, test_process=test_process):
            rho = _np.outer(pure_state, pure_state.conj()).ravel(order='F')
            return unvec_square(_np.dot(test_process, rho),'F')

        def multi_time_test_function(pure_state, test_process=test_process):
            rho = _np.outer(pure_state, pure_state.conj()).ravel(order='F')
            temp = _np.dot(_np.linalg.matrix_power(test_process, 2), rho)
            return [unvec_square(_np.dot(test_process, rho), 'F'), unvec_square(temp, 'F')]

        process_matrix = run_process_tomography(single_time_test_function, n_qubits=2, verbose=False)
        if _rank == 0:
            test_process_pp = change_basis(test_process, 'col', 'pp')
            print("\nSingle-time test result should be True:")
            print(_np.isclose(process_matrix, test_process_pp).all())

        process_matrices = run_process_tomography(multi_time_test_function, n_qubits=2, verbose=False, time_dependent=True)
        if _rank == 0:
            test_process = change_basis(test_process, 'col', 'pp')
            print("\nMulti-time test result should be [True, False]:")
            print([_np.isclose(x, test_process).all() for x in process_matrices])
