import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import argparse

def main(args):
    df = pd.read_csv(args.table, header=None)
    full = df.values

    z1 = full[:6]
    z2 = full[6:]

    x = np.arange(full.shape[1])
    y1 = np.array([0, 0.05, 0.1, 0.2, 0.3, 0.5])
    y2 = np.array([0, 0.05, 0.1, 0.2, 0.3, 0.5])
    X1, Y1 = np.meshgrid(x, y1)
    X2, Y2 = np.meshgrid(x, y2)

    fig = plt.figure(figsize=(10, 10))

    # 첫 번째 subplot
    ax1 = fig.add_subplot(2, 1, 1, projection='3d')
    surf1 = ax1.plot_surface(X1, Y1, z1, cmap='RdBu', alpha=0.95)
    # fig.colorbar(surf1, ax=ax1, shrink=0.5, aspect=10)
    # ax1.set_title('Top Half (z1)')

    # 두 번째 subplot
    ax2 = fig.add_subplot(2, 1, 2, projection='3d')
    surf2 = ax2.plot_surface(X2, Y2, z2, cmap='RdBu', alpha=0.95)
    # fig.colorbar(surf2, ax=ax2, shrink=0.5, aspect=10)
    # ax2.set_title('Bottom Half (z2)')

    plt.tight_layout()
    plt.savefig('figure3.png', dpi=300, bbox_inches='tight')

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='3D Surface Plot with Two Subplots')
    parser.add_argument('--table', type=str, default='ucfcrime_w_img_change.csv')
    args = parser.parse_args()
    main(args)
