# Load the FashionMNIST training dataset
train_data = torchvision.datasets.FashionMNIST(root='./data', train=True, download=True)
images = train_data.data.numpy().astype(float)  # Convert images to numpy array
targets = train_data.targets.numpy()  # Extract labels
class_dict = {i:class_name for i,class_name in enumerate(train_data.classes)}
labels = np.array([class_dict[t] for t in targets])  # Map labels to class names
n = len(images)  # Total samples
class_names = list(class_dict.values())
print(f"Loaded FashionMNIST with {n} samples. Classes: {class_dict}")
print("Classes: {}".format(class_dict))
print("Image shape: {}".format(images[0].shape))
print("Image dtype: {}".format(images[0].dtype))


# Creating a DataFrame with the images and labels
df = pd.DataFrame({"image": images.tolist(), "label": labels})
# Cast image as numpy array
df['image'] = df['image'].apply(lambda x: np.array(x).reshape(-1))

def compute_mae(y_true, y_pred):
  length=len(y_true)
  mae=np.sum(np.abs(y_true-y_pred))/length

  return mae


def compute_r2(y_true, y_pred):
  """Compute R-squared score"""
  R_Square=1-np.sum((y_true-y_pred)**2)/np.sum((y_true-np.mean(y_true))**2)

  return R_Square

def print_metrics(y_true, y_pred, dataset=""):
  """Print all regression metrics for a dataset"""
  rmse = compute_rmse(y_true, y_pred)
  mae = compute_mae(y_true, y_pred)
  r2 = compute_r2(y_true, y_pred)

  print(f"=== {dataset} Metrics ===")
  print(f"RMSE: {rmse:.2f}")
  print(f"MAE: {mae:.2f}")
  print(f"R²: {r2:.3f}")
  return rmse, mae, r2

# Calculate predictions
y_train_pred = price_model.predict(X_train_sc)
y_test_pred = price_model.predict(X_test_sc)

# Compute and print metrics
train_rmse, train_mae, train_r2 = print_metrics(prices_train['Price'], y_train_pred, "Training")
test_rmse, test_mae, test_r2 = print_metrics(prices_test['Price'], y_test_pred, "Test")

# Visualize predictions vs actual
fig = px.scatter(
    x=prices_test['Price'],
    y=y_test_pred,
    title='Predicted vs Actual Prices',
    labels={'x': 'Actual Price', 'y': 'Predicted Price'}
)
fig.add_trace(px.line(x=[prices_test['Price'].min(), prices_test['Price'].max()], 
                      y=[prices_test['Price'].min(), prices_test['Price'].max()]).data[0])
fig.show()

# TODO: Make a two-colored bar plot of the model's accuracy on the test set and the secret test set
result_df=pd.concat([test_df.groupby('label')['correct'].mean(),test_secret_df.groupby('label')['correct'].mean()],axis=1
,keys=['test','secret_test']).reset_index()

result_df_long = result_df.melt(
    id_vars='label',
    value_vars=['test', 'secret_test'],
    var_name='dataset',
    value_name='accuracy'
)

result_df_long.plot(
    x='label',
    y='accuracy',
    color='dataset',  # 按 dataset 列分组着色
    kind='bar',
    title='Class-wise Accuracy Comparison',
    barmode='group'  # 关键：设置为 'group' 实现分组显示
)

def randomly_rotate_images(images, num_rotations_per_image=5):
  """
  Create training data by rotating original images and storing the rotation angles as labels.
  Params:
    - images: numpy array of shape (n_images, 784)
    - num_rotations_per_image: int, number of rotations to perform per image
  Returns:
    - X_augmented: numpy array of shape (n_images * num_rotations_per_image, 784)
    - y_rotations: numpy array of shape (n_images * num_rotations_per_image,)
  """
  # TODO: Implement this function

  X_augmented=[]
  y_rotations=[]

  for image in images:
    angles=np.random.uniform(low=0,high=360,size=num_rotations_per_image)

    image_2d=image.reshape(28,28)

    for angle in angles:
      rotated_2d=rotate(image_2d,angle,reshape=False,cval=0,order=1)

      rotated_1d=rotated_2d.flatten()

      X_augmented.append(rotated_1d)
      y_rotations.append(angle)

  
  return np.array(X_augmented),np.array(y_rotations)
  # raise NotImplementedError("Not implemented")
  
X_train = X_train[:3000]
y_train = y_train[:3000]

num_rotations_per_image = 4
X_train_rotated, y_train_rotations = randomly_rotate_images(X_train, num_rotations_per_image)
y_train_augmented = np.array([y for label in y_train for y in [label]*num_rotations_per_image])
# y_train_augmented = np.repeat(y_train, num_rotations_per_image)
print(X_train_rotated.shape)
print(y_train_augmented.shape)

show_images(X_train_rotated, max_images=5, ncols=5, labels=y_train_augmented, reshape=True)

X_train_rotated_sc = scaler.transform(X_train_rotated)
X_test_secret_sc = scaler.transform(X_test_secret)

if load_saved_models:
  model_rotated = joblib.load('mlp_fashionmnist_rotated_model.joblib')
else:
  # TODO: initialize and train a new model on the rotated images
  model_rotated =  MLPClassifier()
  model_rotated.fit(X=np.stack(X_train_rotated_sc),y=y_train_augmented)

  if save_models:
    joblib.dump(model_rotated, "mlp_fashionmnist_rotated_model.joblib")
    print("Rotated model saved to disk.")

print(f"Training accuracy (training on rotated images): {model_rotated.score(X_train_rotated_sc, y_train_augmented):.3f}")
print(f"Test accuracy (original): {model.score(X_test_sc, y_test):.3f}")
print(f"Test accuracy (training on rotated images): {model_rotated.score(X_test_secret_sc, y_test_secret):.3f}")

loss_df = pd.DataFrame({
    'epoch': range(1, len(model_rotated.loss_curve_) + 1),
    'loss': model_rotated.loss_curve_
})
loss_df.plot(x='epoch', y='loss', title="Training Error")

if (load_saved_models or IS_GRADING_ENV) and os.path.exists('model_rotation_regression.joblib'):
  model_rotation_regression = joblib.load("model_rotation_regression.joblib")
else:
  # Train a larger regression model (MLP) to predict rotation angles
  model_rotation_regression = MLPRegressor(
      hidden_layer_sizes=(256, 128),
      activation='relu',
      solver='adam',
      max_iter=100,
      random_state=SEED,
      early_stopping=True,
      verbose=True
  )
  # TODO: train a MLP regressor to predict rotation angles
  model_rotation_regression.fit(X=X_train_rotated_sc,y=y_train_rotations)
  # Save model_rotation_regression and scaler
  if save_models:
    joblib.dump(model_rotation_regression, "model_rotation_regression.joblib")
    joblib.dump(scaler, "scaler.joblib")
    print("Rotation regression model and scaler saved to disk.")

loss_df = pd.DataFrame({
    'epoch': range(1, len(model_rotation_regression.loss_curve_) + 1),
    'loss': model_rotation_regression.loss_curve_
})
loss_df.plot(x='epoch', y='loss', title="Training Loss")

Iteration 1, loss = 9611.26410087
Validation score: 0.142763
Iteration 2, loss = 3766.69187363
Validation score: 0.385691
Iteration 3, loss = 2920.89570158
Validation score: 0.484820
Iteration 4, loss = 2530.66131818
Validation score: 0.523250
Iteration 5, loss = 2275.49896874
Validation score: 0.559317
Iteration 6, loss = 2085.74463183
Validation score: 0.584974
Iteration 7, loss = 1924.98972761
Validation score: 0.602003
Iteration 8, loss = 1804.65769176
Validation score: 0.618211
Iteration 9, loss = 1682.16687650
Validation score: 0.631943
Iteration 10, loss = 1571.02563443
Validation score: 0.640702
Iteration 11, loss = 1478.80508676
Validation score: 0.653446
Iteration 12, loss = 1392.98262424
Validation score: 0.655017
Iteration 13, loss = 1316.21452457
Validation score: 0.674004
Iteration 14, loss = 1236.47938882
Validation score: 0.679636
Iteration 15, loss = 1161.99735384
Validation score: 0.687010
Iteration 16, loss = 1105.28316375
Validation score: 0.699418
Iteration 17, loss = 1048.20584135
Validation score: 0.693873
Iteration 18, loss = 977.81792593
Validation score: 0.711631
Iteration 19, loss = 925.44572624
Validation score: 0.699191
Iteration 20, loss = 899.88230979
Validation score: 0.702843
Iteration 21, loss = 840.36804019
Validation score: 0.729254
Iteration 22, loss = 780.98734272
Validation score: 0.728786
Iteration 23, loss = 765.77998552
Validation score: 0.734939
Iteration 24, loss = 705.92240284
Validation score: 0.737799
Iteration 25, loss = 674.86031295
Validation score: 0.735162
Iteration 26, loss = 634.87954343
Validation score: 0.745473
Iteration 27, loss = 589.66117117
Validation score: 0.745895
Iteration 28, loss = 579.90028943
Validation score: 0.744533
Iteration 29, loss = 543.20496845
Validation score: 0.734867
Iteration 30, loss = 505.46502073
Validation score: 0.747184
Iteration 31, loss = 475.18828419
Validation score: 0.750737
Iteration 32, loss = 451.34752457
Validation score: 0.752262
Iteration 33, loss = 432.74789467
Validation score: 0.735379
Iteration 34, loss = 403.09219879
Validation score: 0.756212
Iteration 35, loss = 375.65089303
Validation score: 0.752666
Iteration 36, loss = 356.99911305
Validation score: 0.763152
Iteration 37, loss = 333.24136299
Validation score: 0.755819
Iteration 38, loss = 317.56468084
Validation score: 0.763254
Iteration 39, loss = 300.31424482
Validation score: 0.752810
Iteration 40, loss = 292.59236521
Validation score: 0.763904
Iteration 41, loss = 283.44042869
Validation score: 0.748904
Iteration 42, loss = 274.85827460
Validation score: 0.762908
Iteration 43, loss = 242.64105699
Validation score: 0.762618
Iteration 44, loss = 224.33508193
Validation score: 0.755256
Iteration 45, loss = 209.78245760
Validation score: 0.764000
Iteration 46, loss = 198.72396655
Validation score: 0.760627
Iteration 47, loss = 195.43882927
Validation score: 0.752333
Iteration 48, loss = 184.33454222
Validation score: 0.762145
Iteration 49, loss = 176.41046927
Validation score: 0.762408
Iteration 50, loss = 165.02544010
Validation score: 0.752812
Iteration 51, loss = 165.35008826
Validation score: 0.752220
Validation score did not improve more than tol=0.000100 for 10 consecutive epochs. Stopping.
Rotation regression model and scaler saved to disk.

# TODO: Use the model_rotation_regression to predict the rotation angles on X_test_rotated_sc
X_test_rotated, y_test_rotations = randomly_rotate_images(X_test,num_rotations_per_image=1)
X_test_rotated_sc = scaler.transform(X_test_rotated)

y_pred_angles = model_rotation_regression.predict(X_test_rotated_sc)
mse = mean_squared_error(y_test_rotations,y_pred_angles)
rmse = np.sqrt(mse)
print(f"Test MSE: {mse:.2f}")
print(f"Test RMSE: {rmse:.2f} degrees")

# Show some predictions
print("\nSample predictions:")
for i in range(5):
    print(f"True rotation: {y_test_rotations[i]:.1f}°, Predicted: {y_pred_angles[i]:.1f}°")

print("------------------------------")
print("------------------------------")
print(f"Test MSE (angle):  {mse:.2f}")
print(f"Test RMSE:        {rmse:.2f}°")
print("------------------------------")
print("------------------------------")

# Create scatter plot of predictions vs actual
scatter_fig = go.Figure()
scatter_fig.add_trace(
    go.Scatter(
        x=y_test_rotations[:200],  # First 200 samples for clarity
        y=y_pred_angles[:200],
        mode='markers',
        name='Predictions',
        marker=dict(size=6, opacity=0.7)
    )
)

# Add perfect prediction line
max_angle = max(y_test_rotations[:200])
scatter_fig.add_trace(
    go.Scatter(
        x=[0, max_angle],
        y=[0, max_angle],
        mode='lines',
        name='Perfect Prediction',
        line=dict(dash='dash', color='red')
    )
)

scatter_fig.update_layout(
    title='Predicted vs Actual Rotation Angles',
    xaxis_title='Actual Angle (degrees)',
    yaxis_title='Predicted Angle (degrees)',
    width=600,
    height=500
)

scatter_fig.show()

# Demonstrate unrotation using predicted angles
def unrotate_image(rotated_image, predicted_angle):
    """
    Unrotate an image by rotating it by the negative of the predicted angle.
    """
    return rotate_image_by_angle(rotated_image, -predicted_angle)

# Test unrotation on a sample
sample_idx = 0
original_test_image = X_test[sample_idx // 1]  # Get original unrotated image
rotated_test_image = X_test_rotated[sample_idx] # Get the rotated image
true_angle = y_test_rotations[sample_idx] # Get the true angle the image was rotated by
predicted_angle = y_pred_angles[sample_idx] # Get the model's prediction
unrotated_image = unrotate_image(rotated_test_image, predicted_angle)

# Prepare images and titles for display
images = [
    original_test_image.reshape(28, 28)[::-1],
    rotated_test_image.reshape(28, 28)[::-1],
    unrotated_image.reshape(28, 28)[::-1]
]
titles = [
    'Original',
    f'Rotated by {true_angle:.1f}°',
    f'After unrotating the image with the model\'s prediction (pred: {predicted_angle:.1f}°)'
]

# Create a subplot grid using plotly express imshow
fig = px.imshow(
    np.stack(images),
    facet_col=0,
    facet_col_wrap=3,
    color_continuous_scale='gray',
    aspect='auto'
)

# Update facet titles
for i, title in enumerate(titles):
    fig.layout.annotations[i]['text'] = title

fig.update_layout(
    height=300,
    width=1200,
    title_text="Image Rotation Prediction and Correction",
    coloraxis_showscale=False
)
fig.update_xaxes(showticklabels=False, showgrid=False, zeroline=False)
fig.update_yaxes(showticklabels=False, showgrid=False, zeroline=False)
fig.show()

# Compare all three approaches
baseline_accuracy = model.score(X_test_secret_sc, y_test_secret) # method: baseline
rotated_training_accuracy = model_rotated.score(X_test_secret_sc, y_test_secret) # method: train on rotated images
unrotated_accuracy = model.score(X_test_unrotated_sc, y_test_secret) # method: predict & unrotate, then classify the images we tried to unrotate

# Create comparison DataFrame
comparison_df = pd.DataFrame({
    'Method': ['Baseline (no handling)', 'Train on rotated images', 'Predict & unrotate'],
    'Accuracy': [baseline_accuracy, rotated_training_accuracy, unrotated_accuracy]
})

# Plot comparison
fig = px.bar(
    comparison_df,
    x='Method',
    y='Accuracy',
    title='Comparison of Rotation Handling Methods',
    color='Accuracy',
    color_continuous_scale='RdYlGn'
)
fig.update_layout(
    xaxis_title="Method",
    yaxis_title="Test Accuracy",
    yaxis=dict(range=[0, 1])
)
fig.show()

print("\n=== Summary of All Methods ===")
for _, row in comparison_df.iterrows():
    print(f"{row['Method']}: {row['Accuracy']:.3f}")

┌─────────────────────────────────────────────────────────────────┐
│                      test_time_augmentation 函数                │
├─────────────────────────────────────────────────────────────────┤
│                                                                 │
│   输入: image (原始测试图片)                                     │
│                                                                 │
│   Step 1: 预处理                                                 │
│   ┌─────────────────────────────────────────────────────────┐  │
│   │ image → 展平 → 标准化                                    │  │
│   └─────────────────────────────────────────────────────────┘  │
│                          │                                      │
│                          ▼                                      │
│   Step 2: 预测旋转角度并纠正                                     │
│   ┌─────────────────────────────────────────────────────────┐  │
│   │ model_rotation_regressor.predict() → 预测角度           │  │
│   │ rotate_image_by_angle(..., -angle) → 转回来            │  │
│   └─────────────────────────────────────────────────────────┘  │
│                          │                                      │
│                          ▼                                      │
│   Step 3: 多角度 TTA 预测                                       │
│   ┌─────────────────────────────────────────────────────────┐  │
│   │ 对 unrotated 图像应用 9 种角度旋转:                     │  │
│   │ -20°, -15°, -10°, -5°, 0°, 5°, 10°, 15°, 20°          │  │
│   │ 对每种旋转后的图像 → 标准化 → model.predict_proba()     │  │
│   │ 收集所有概率向量                                          │  │
│   └─────────────────────────────────────────────────────────┘  │
│                          │                                      │
│                          ▼                                      │
│   Step 4: 额外预测                                               │
│   ┌─────────────────────────────────────────────────────────┐  │
│   │ 原始图像的预测 + 转回来图像的预测                        │  │
│   └─────────────────────────────────────────────────────────┘  │
│                          │                                      │
│                          ▼                                      │
│   Step 5: 聚合与最终预测                                         │
│   ┌─────────────────────────────────────────────────────────┐  │
│   │ avg_probs = mean(all_probs)  # 11个概率向量取平均       │  │
│   │ prediction = argmax(avg_probs)  # 选择概率最大的类      │  │
│   └─────────────────────────────────────────────────────────┘  │
│                                                                 │
│   输出: class_dict[prediction]  # 返回类别标签                   │
└─────────────────────────────────────────────────────────────────┘

# TODO: Write a function (or functions!) that can be used to improve the accuracy of the original MLPClassifer model using test time augmentations
# Feel free to add additional functions, arguments, etc. as needed!
def test_time_augmentation(model, scaler, image):
    """
    Predict a label using test-time augmentation
    Params:
        - model: the MLPClassifier model
        - scaler: the StandardScaler used to scale the images
        - image: the image to be augmented
    Returns:
        - prediction: the predicted label
    """
    ...
    # raise NotImplementedError("Not implemented")
    image = np.array(image, dtype=np.float64).reshape(-1)
    image_scaled=scaler.transform([image])
    predicted_angle=model_rotation_regression.predict(image_scaled)[0]

    unrotated=rotate_image_by_angle(image,-predicted_angle)

    test_angles=[-20,-15,-10,-5,0,5,10,15,20]

    all_probs=[]

    for angle in test_angles:
        rotated=rotate_image_by_angle(unrotated,angle)
        rotated_scaled=scaler.transform([rotated])
        probs=model.predict_proba(rotated_scaled)[0]
        all_probs.append(probs)

    original_scaled=scaler.transform([image])
    unrotated_scaled=scaler.transform([unrotated])
    all_probs.append(model.predict_proba(original_scaled)[0])
    all_probs.append(model.predict_proba(unrotated_scaled)[0])

    avg_probs=np.mean(all_probs,axis=0)
    prediction=np.argmax(avg_probs)

    return class_dict[prediction]



# Make a copy of the test secret dataframe and apply the test time augmentation to the images to get new predictions!
part_2_df = test_secret_df.copy()
part_2_df["image"] = part_2_df["image"].apply(lambda x: np.array(x).reshape(-1))
part_2_df["prediction"] = part_2_df["image"].apply(lambda x: test_time_augmentation(model_rotated, scaler, x))

# Check the accuracy of the new predictions
correct = part_2_df["prediction"] == part_2_df["label"]
print("Accuracy:", correct.mean())