# -*- coding: utf-8 -*-
# @Time    : 2023/2/24 20:14
# @Author  : AI悦创
# @FileName: hw4.py
# @Software: PyCharm
# @Blog    ：https://bornforthis.cn/
from qutip import *
from scipy.linalg import *
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import matplotlib.pyplot as plt

# 第一步：混态的Bloch坐标
def getCoordFromRho(Rho):
    coord = np.array([0.0000, 0.0000, 1.0000], dtype = float)
    # 计算Bloch向量的x,y,z坐标
    coord[0] = np.real(np.trace(np.dot(sigmax().full(), Rho)))
    coord[1] = np.real(np.trace(np.dot(sigmay().full(), Rho)))
    coord[2] = np.real(np.trace(np.dot(sigmaz().full(), Rho)))
    return coord

def getRhoFromCoord(coord):
    Rho = np.array([[1.0000 + 0.0000j, 0.0000 + 0.0000j], [0.0000 + 0.0000j, 0.0000 + 0.0000j]], dtype = complex)
    # 计算混态的密度矩阵
    Rho = 0.5 * (np.identity(2) + coord[0] * sigmax().full() + coord[1] * sigmay().full() + coord[2] * sigmaz().full())
    return Rho

# # 测试代码
# rho = rand_dm(2)
# print("随机混态的Bloch坐标为：", getCoordFromRho(rho))
# print("恢复后的混态为：\n", getRhoFromCoord(getCoordFromRho(rho)))

# don't modify the code in this block
n_ck = 3
rho_ck = np.array([[[0.5000 + 0.0000j, 0.0000 + 0.0000j],
                    [0.0000 + 0.0000j, 0.5000 + 0.0000j]],
                   [[0.7500 + 0.0000j, 0.2500 - 0.2500j],
                    [0.2500 + 0.2500j, 0.2500 + 0.0000j]],
                   [[0.5000 + 0.0000j, 0.3535 + 0.3535j],
                    [0.3535 - 0.3535j, 0.5000 + 0.0000j]]], dtype=complex)
coord_ck = np.array([[0.0000, 0.0000, 0.0000],
                     [0.5000, 0.5000, 0.5000],
                     [0.7071, -0.7071, 0.0000]], dtype=float)


def checkCoord2Rho():
    print('Checking converting coordinate to state vector...')
    err = [np.sum(abs(getRhoFromCoord(coord_ck[c]) - rho_ck[c])) for c in range(n_ck)]
    if np.sum(err) < 0.01:
        print('Pass!')
    else:
        print('Wrong Answer err = %.3f! Please Correct the code.' % np.sum(err))


def checkRho2Coord():
    print('Checking converting state vector to coordinate...')
    err = [np.sum(abs(getCoordFromRho(rho_ck[c]) - coord_ck[c])) for c in range(n_ck)]
    if np.sum(err) < 0.01:
        print('Pass!')
    else:
        print('Wrong Answer err = %.3f! Please Correct the code.' % np.sum(err))


checkCoord2Rho()
checkRho2Coord()

from qutip import *
from scipy.linalg import *
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import matplotlib.pyplot as plt


def getBasisState(O):
    # 初始化态矢量
    psi_0 = np.matrix([[1.0000], [0.0000]], dtype=complex)
    psi_1 = np.matrix([[0.0000], [1.0000]], dtype=complex)

    # 对角化算符
    eig_v, eig_k = eig(O)
    eig_v = np.real(eig_v)

    # 对角化后的基矢量
    m_base = [np.matrix(eig_k[:, i]).T for i in range(len(eig_v))]

    # 对角化后基矢量的顺序
    order = np.argsort(eig_v)[::-1]
    m_base = [m_base[i] for i in order]

    return m_base


def Meas(psi, O):
    m_base = getBasisState(O)

    # 计算投影概率
    p0 = (abs(m_base[0].H * psi) ** 2)[0, 0]
    p1 = (abs(m_base[1].H * psi) ** 2)[0, 0]

    return [p0, p1]