import GPy
import numpy as np
import polars as pl
import matplotlib.pyplot as plt
import matplotlib.patches as patches

from matplotlib import rcParams

rcParams["font.family"] = "sans-serif"
rcParams["font.sans-serif"] = ["Yu Gothic"]

MAX_PX = 3840
AR = 1
FIG_WIDTH = 8
FIG_HEIGHT = FIG_WIDTH / AR
DPI = MAX_PX / FIG_WIDTH
WAFER_RADIUS_MM = 100


def load_data():
    FILE_PATH = "data/data.csv"
    df = pl.read_csv(FILE_PATH)
    X = df.select(["wafer_x_mm", "wafer_y_mm"]).to_numpy()
    Y = df.select(["thickness_nm"]).to_numpy()
    return X, Y


def build_model(X, Y):
    kern = GPy.kern.RBF(input_dim=2, lengthscale=100)
    model = GPy.models.GPRegression(X, Y, kernel=kern, normalizer=True)
    model.optimize()
    print(model)
    return model


def plot_2d(model):
    fig = plt.figure(figsize=[FIG_WIDTH, FIG_HEIGHT])
    ax = fig.add_subplot(1, 1, 1)

    model.plot(projection="2d", ax=ax, levels=16)
    ax.grid()
    ax.set_aspect("equal")

    ax.set_xlim(-WAFER_RADIUS_MM, WAFER_RADIUS_MM)
    ax.set_ylim(-WAFER_RADIUS_MM, WAFER_RADIUS_MM)
    ax.set_xlabel("x(mm)")
    ax.set_ylabel("y(mm)")

    circle = patches.Circle(
        xy=(0, 0), radius=WAFER_RADIUS_MM, fill=False, linewidth=2, edgecolor="black"
    )
    ax.add_patch(circle)

    fig.savefig("map_lpteos_thickness_2d.png", dpi=DPI, facecolor="white")

    plt.show()


def plot_3d(model):
    fig = plt.figure(figsize=[FIG_WIDTH, FIG_HEIGHT])
    ax = fig.add_subplot(1, 1, 1, projection="3d")

    model.plot(projection="3d", ax=ax)
    ax.grid()

    ax.set_xlim(-WAFER_RADIUS_MM, WAFER_RADIUS_MM)
    ax.set_ylim(-WAFER_RADIUS_MM, WAFER_RADIUS_MM)
    ax.set_xlabel("x(mm)")
    ax.set_ylabel("y(mm)")
    ax.set_zlabel("thickness(nm)")

    fig.savefig("map_lpteos_thickness_3d.png", dpi=DPI, facecolor="white")

    plt.show()


def predict_point(model, x, y):
    X = np.array([[x, y]])
    Y_mu, Y_sigma = model.predict(X)
    print(f"wafer_x = {x} (mm)")
    print(f"wafer_y = {y} (mm)")
    print(f"thickness μ = {Y_mu[0, 0]:.1f} (nm)")
    print(f"thickness σ = {Y_sigma[0, 0]:.1f} (nm)")
    return Y_mu, Y_sigma


def main():
    X, Y = load_data()
    model = build_model(X, Y)
    plot_2d(model)
    plot_3d(model)
    _ = predict_point(model, 50, 70)


if __name__ == "__main__":
    main()