import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import statsmodels.api as sm
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.path import Path



def piecewise_linear_regression_3d(data, candidate_knots_x, candidate_knots_y):
    best_knot_x = None
    best_knot_y = None
    best_rss = np.inf 
    best_model = None

    for knot_x in candidate_knots_x:
        for knot_y in candidate_knots_y:
            data['D_x'] = (data['X'] > knot_x).astype(int)
            data['D_y'] = (data['Y'] > knot_y).astype(int)
            data['kink_x'] = (data['X'] - knot_x) * data['D_x']
            data['kink_y'] = (data['Y'] - knot_y) * data['D_y']
            
            X_vars = sm.add_constant(data[['X', 'Y', 'D_x', 'D_y', 'kink_x', 'kink_y']])
            
            model = sm.OLS(data['Z'], X_vars).fit()
            
            rss = np.sum(model.resid ** 2)
            
            if rss < best_rss:
                best_rss = rss
                best_knot_x = knot_x
                best_knot_y = knot_y
                best_model = model
                
    return best_knot_x, best_knot_y, best_model



def piecewise_linear_regression(data, candidate_knots):
    best_knot = None
    best_rss = np.inf
    best_model = None

    for knot in candidate_knots:

        data['D'] = (data['X'] > knot).astype(int)
        data['kink'] = (data['X'] - knot) * data['D']
        
        X_vars = sm.add_constant(data[['X', 'D', 'kink']])
        
        model = sm.OLS(data['Y'], X_vars).fit()
        
        rss = np.sum(model.resid ** 2)
        
        if rss < best_rss:
            best_rss = rss
            best_knot = knot
            best_model = model
            
    return best_knot, best_model



def kink_regression_2d(file_path,xl,xh,iplot=True):
    data = pd.read_csv(file_path)
    candidate_knots_x = np.linspace(xl, xh, 20)
    best_knot_x, best_model = piecewise_linear_regression(data, candidate_knots_x)
    if iplot:
        x_vals = np.linspace(xl, xh, 100)
        y_vals = best_model.predict(sm.add_constant(pd.DataFrame({
            'X': x_vals,
            'D': (x_vals > best_knot).astype(int),
            'kink': (x_vals - best_knot) * (x_vals > best_knot)
        })))
        plt.scatter(data['X'], data['Y'], alpha=0.5, label='Data')
        plt.plot(x_vals, y_vals, color='#F09148', label='Piecewise Linear Fit')
        plt.axvline(best_knot, color='grey', linestyle='--', label='Estimated Knot')
        plt.xlabel('X')
        plt.ylabel('Y')
        plt.title('Piecewise Linear Fit with Estimated Knot')
        plt.legend()
        plt.show()
    return best_knot_x, best_model



def kink_regression_3d(file_path,xl,xh,yl,yh,iplot=True):
    data = pd.read_csv(file_path)
    candidate_knots_x = np.linspace(xl, xh, 20)  
    candidate_knots_y = np.linspace(yl, yh, 20)
    best_knot_x, best_knot_y, best_model = piecewise_linear_regression_3d(data, candidate_knots_x, candidate_knots_y)
    x_vals = np.linspace(xl, xh, 100)
    y_vals = np.linspace(yl, yh, 100)
    X_grid, Y_grid = np.meshgrid(x_vals, y_vals)
    data_grid = pd.DataFrame({
        'X': X_grid.ravel(),
        'Y': Y_grid.ravel(),
        'D_x': (X_grid.ravel() > best_knot_x).astype(int),
        'D_y': (Y_grid.ravel() > best_knot_y).astype(int),
        'kink_x': (X_grid.ravel() - best_knot_x) * (X_grid.ravel() > best_knot_x),
        'kink_y': (Y_grid.ravel() - best_knot_y) * (Y_grid.ravel() > best_knot_y)
    })
    min_idx = data['Z'].idxmin()
    min_point = data.loc[min_idx]
    min_x, min_y, min_z = min_point['X'], min_point['Y'], min_point['Z']
    Z_pred = best_model.predict(sm.add_constant(data_grid))
    Z_grid = Z_pred.values.reshape(X_grid.shape)
    intercept = best_model.params['const']
    coef_X = best_model.params['X']
    coef_Y = best_model.params['Y']
    coef_D_x = best_model.params['D_x']
    coef_D_y = best_model.params['D_y']
    coef_kink_x = best_model.params['kink_x']
    coef_kink_y = best_model.params['kink_y']

    D_x = int(min_x > best_knot_x)
    D_y = int(min_y > best_knot_y)
    kink_x = (min_x - best_knot_x) * D_x
    kink_y = (min_y - best_knot_y) * D_y

    min_point_pred = (intercept +
                    coef_X * min_x +
                    coef_Y * min_y +
                    coef_D_x * D_x +
                    coef_D_y * D_y +
                    coef_kink_x * kink_x +
                    coef_kink_y * kink_y)
    if iplot:
        fig = plt.figure(figsize=(14, 8))

        ax = fig.add_subplot(121, projection='3d')
        ax.scatter(data['X'], data['Y'], data['Z'], alpha=0.5, label='Data', color='blue')

        ax.plot_surface(X_grid, Y_grid, Z_grid, color='gray', alpha=0.5, label='Piecewise Linear Fit')

        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        ax.set_title('Regression Kink Design')
        ax.legend()
        plt.show()

        fig2 = plt.figure(figsize=(14, 6))
        ax2 = fig2.add_subplot(122)
        ax2.scatter(data['X'], data['Y'], alpha=0.5, label='Total Sample', color='#427AB2')

        contour_plot = ax2.contour(X_grid, Y_grid, Z_grid, levels=[max(min_point_pred * 1.1, min_z)], colors='black', linewidths=2, label='Contour at Z=minz')

        contour_path = contour_plot.get_paths()[0]
        inside_mask = contour_path.contains_points(data[['X', 'Y']].values)

        ax2.scatter(data.loc[inside_mask, 'X'], data.loc[inside_mask, 'Y'], color='#F8D5E4', label='Likely Optm Points', edgecolor='black')
        ax2.scatter(min_x, min_y, color='#CC88B0', s=100, label='Min Z Point', edgecolor='black')
        ax2.scatter(best_knot_x, best_knot_y, color='#96D1C6', s=100, label='Kink Point', edgecolor='black')

        ax2.plot([], [], color='black', linewidth=2, label='O(minz) in Kink Eq')
        ax2.legend()
        ax2.set_xlabel('X')
        ax2.set_ylabel('Y')
        ax2.set_title('Response Surface Projection')
        ax2.legend()
        plt.show()
    return best_knot_x, best_knot_y, best_model, pd.concat([data.loc[inside_mask, 'X'], data.loc[inside_mask, 'Y']],axis=1)
