import os
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np


def save_to_xlsx(project_name, file_name, data):
  folder_path = f'Result/{project_name}'
  os.makedirs(folder_path, exist_ok=True)
  data.to_excel(f'{folder_path}/{file_name}.xlsx', index=True)
  print('Save successed to ' + f'{folder_path}/{file_name}.xlsx')
  save_to_pic(project_name=project_name, file_name=file_name)
  return


def save_to_pic(project_name, file_name):
  os.makedirs(f'Result/{project_name}', exist_ok=True)
  if file_name == 'pca_2d':
    draw_pca_2d(f'Result/{project_name}/{file_name}.xlsx')
    print('Save successed to ' + f'Result/{project_name}/{file_name}.png')
  elif file_name == 'pca_3d':
    draw_pca_3d(f'Result/{project_name}/{file_name}.xlsx')
    print('Save successed to ' + f'Result/{project_name}/{file_name}.png')
  elif file_name == 'acc_and_loss':
    draw_epoch_data(f'Result/{project_name}/{file_name}.xlsx')
    draw_last_epoch_bar_chart(f'Result/{project_name}/{file_name}.xlsx')
    print('Save successed to line graph and bar graph')
  elif file_name == 'cm':
    draw_and_save_cm(f'Result/{project_name}/{file_name}.xlsx')
    print('Save successed cm')
  elif file_name == 'cmn':
    draw_and_save_cm(f'Result/{project_name}/{file_name}.xlsx')
    print('Save successed cmn')
  else:
    print('unknow picture type')


def draw_pca_2d(file_path):
  df = pd.read_excel(file_path)
  plt.figure(figsize=(8, 6))
  plt.scatter(df['PC1'], df['PC2'], c=df['labels'], cmap='viridis', edgecolor='k', alpha=0.6)
  plt.xlabel('PC1')
  plt.ylabel('PC2')
  plt.title('2D PCA')
  plt.colorbar(label='Labels')
  plt.savefig(file_path.replace('.xlsx', '.png'))
  plt.close()


def draw_pca_3d(file_path):
  df = pd.read_excel(file_path)
  fig = plt.figure(figsize=(8, 6))
  ax = fig.add_subplot(111, projection='3d')
  scatter = ax.scatter(df['PC1'], df['PC2'], df['PC3'], c=df['labels'], cmap='viridis', edgecolor='k', alpha=0.6)
  ax.set_xlabel('PC1')
  ax.set_ylabel('PC2')
  ax.set_zlabel('PC3')
  ax.set_title('3D PCA')
  fig.colorbar(scatter, ax=ax, label='Labels')
  plt.savefig(file_path.replace('.xlsx', '.png'))


def draw_epoch_data(file_path):
  df = pd.read_excel(file_path)
  epochs = df['epoch']
  train_loss = df['train_loss']
  train_accuracy = df['train_accuracy'] * 100
  test_accuracy = df['test_accuracy'] * 100
  f1_score = df['f1_score']
  precision = df['precision']
  recall = df['recall']

  fig, axs = plt.subplots(2, 3, figsize=(18, 12))

  axs[0, 0].plot(epochs, train_loss, 'b-', label='Train Loss')
  axs[0, 0].set_xlabel('Epoch')
  axs[0, 0].set_ylabel('Loss')
  axs[0, 0].set_title('Training Loss over Epochs')
  axs[0, 0].legend()

  axs[0, 1].plot(epochs, train_accuracy, 'g-', label='Train Accuracy')
  axs[0, 1].plot(epochs, test_accuracy, 'r-', label='Test Accuracy')
  axs[0, 1].set_xlabel('Epoch')
  axs[0, 1].set_ylabel('Accuracy (%)')
  axs[0, 1].set_title('Train and Test Accuracy over Epochs')
  axs[0, 1].legend()

  axs[0, 2].plot(epochs, f1_score, 'm-', label='F1 Score')
  axs[0, 2].set_xlabel('Epoch')
  axs[0, 2].set_ylabel('F1 Score')
  axs[0, 2].set_title('F1 Score over Epochs')
  axs[0, 2].legend()

  axs[1, 0].plot(epochs, precision, 'c-', label='Precision')
  axs[1, 0].set_xlabel('Epoch')
  axs[1, 0].set_ylabel('Precision')
  axs[1, 0].set_title('Precision over Epochs')
  axs[1, 0].legend()

  axs[1, 1].plot(epochs, recall, 'y-', label='Recall')
  axs[1, 1].set_xlabel('Epoch')
  axs[1, 1].set_ylabel('Recall')
  axs[1, 1].set_title('Recall over Epochs')
  axs[1, 1].legend()

  axs[1, 2].axis('off')

  plt.tight_layout()
  plt.savefig(file_path.replace('.xlsx', '_epoch.png'))
  plt.close()


def draw_last_epoch_bar_chart(file_path):
  df = pd.read_excel(file_path)
  last_epoch_data = df.iloc[-1]

  metrics = ['train_loss', 'train_accuracy', 'test_accuracy', 'f1_score', 'precision', 'recall']
  values = [last_epoch_data[metric] for metric in metrics]
  labels = ['Train Loss', 'Train Accuracy', 'Test Accuracy', 'F1 Score', 'Precision', 'Recall']

  values[1] *= 100
  values[2] *= 100

  plt.figure(figsize=(10, 6))
  plt.bar(labels, values, color=['blue', 'green', 'red', 'magenta', 'cyan', 'yellow'])
  plt.xlabel('Metrics')
  plt.ylabel('Values')
  plt.title('Last Epoch Metrics')
  plt.ylim(bottom=0)

  for i, value in enumerate(values):
    plt.text(i, value + 0.01, f'{value:.2f}', ha='center')

  plt.tight_layout()
  plt.savefig(file_path.replace('.xlsx', '_last_epoch_bar.png'))
  plt.close()


def draw_and_save_cm(file_path):
  df_cm = pd.read_excel(file_path)

  labels = df_cm.columns[1:].tolist()
  cm = df_cm.values[:, 1:]

  fig, axs = plt.subplots(1, 2, figsize=(12, 6))

  axs[0].imshow(cm, interpolation='nearest', cmap='Blues')
  axs[0].set_title('Confusion Matrix')
  axs[0].set_xlabel('Predicted')
  axs[0].set_ylabel('True')
  axs[0].set_xticks(np.arange(len(labels)))
  axs[0].set_yticks(np.arange(len(labels)))
  axs[0].set_xticklabels(labels)
  axs[0].set_yticklabels(labels)

  for i in range(len(labels)):
    for j in range(len(labels)):
      axs[0].text(j, i, f'{cm[i, j]}', ha='center', va='center')

  plt.tight_layout()
  plt.savefig(file_path.replace('.xlsx', '.png'))
  plt.close()