import numpy as np
import matplotlib.pyplot as plt

from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

from qiskit import Aer
from qiskit.utils import QuantumInstance
from qiskit.circuit.library import ZZFeatureMap, RealAmplitudes
from qiskit_machine_learning.neural_networks import EstimatorQNN
from qiskit_machine_learning.algorithms import NeuralNetworkRegressor

from qiskit.primitives import Estimator
from qiskit.circuit import QuantumCircuit

# =========================
# 1. 构造数据
# =========================
np.random.seed(42)

X = np.linspace(0, 2*np.pi, 50)
y = np.sin(X)

X = X.reshape(-1, 1)
y = y.reshape(-1, 1)

# 归一化
scaler = MinMaxScaler((-1, 1))
X = scaler.fit_transform(X)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

# =========================
# 2. 构建量子电路
# =========================

num_qubits = 1

# 特征映射
feature_map = ZZFeatureMap(feature_dimension=num_qubits, reps=2)

# 可训练参数层
ansatz = RealAmplitudes(num_qubits, reps=2)

qc = QuantumCircuit(num_qubits)
qc.compose(feature_map, inplace=True)
qc.compose(ansatz, inplace=True)

# =========================
# 3. 构建量子神经网络
# =========================

estimator = Estimator()

qnn = EstimatorQNN(
    circuit=qc,
    input_params=feature_map.parameters,
    weight_params=ansatz.parameters,
    estimator=estimator
)

regressor = NeuralNetworkRegressor(
    neural_network=qnn,
    optimizer=None,  # 默认使用 L-BFGS
)

# =========================
# 4. 训练模型
# =========================
regressor.fit(X_train, y_train)

# =========================
# 5. 预测
# =========================
y_pred = regressor.predict(X_test)

mse = mean_squared_error(y_test, y_pred)
print("测试集 MSE:", mse)

# =========================
# 6. 可视化
# =========================
plt.scatter(X_test, y_test, label="真实值")
plt.scatter(X_test, y_pred, label="预测值")
plt.legend()
plt.title("量子机器学习回归预测")
plt.show()