import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial.distance import cdist
import ipywidgets as widgets
from matplotlib.patches import Circle

np.random.seed(1)

# Positions fixes X des points
points_x = {'A': 30, 'B': 45, 'E': 50, 'C': 55, 'D': 70}
labels = ['A', 'B', 'E', 'C', 'D']

def exponential_cov(h, range_, sill, nugget=0):
    return nugget * (h == 0) + (sill - nugget) * np.exp(-h / range_)

def compute_simple_weights(observed_points, range_, sill):
    X = np.array(list(observed_points.values()))
    target = np.array([points_x['E'], 50]).reshape(1,2)  # Point E fixe
    
    if len(X) == 0:
        return {}
    
    d_matrix = cdist(X, X)
    d_vector = cdist(X, target).flatten()
    
    C = exponential_cov(d_matrix, range_, sill)
    c0 = exponential_cov(d_vector, range_, sill)
    
    try:
        weights = np.linalg.solve(C, c0)
    except np.linalg.LinAlgError:
        weights = np.full(len(X), np.nan)
    
    return {label: w for label, w in zip(observed_points.keys(), weights)}

def compute_ordinary_weights(observed_points, range_, sill):
    X = np.array(list(observed_points.values()))
    target = np.array([points_x['E'], 50]).reshape(1,2)  # Point E fixe
    
    if len(X) == 0:
        return {}
    
    d_matrix = cdist(X, X)
    d_vector = cdist(X, target).flatten()
    
    C = exponential_cov(d_matrix, range_, sill)
    c0 = exponential_cov(d_vector, range_, sill)
    
    n = len(X)
    C_ext = np.ones((n+1, n+1))
    C_ext[:n, :n] = C
    C_ext[-1, -1] = 0
    c0_ext = np.ones(n+1)
    c0_ext[:n] = c0
    
    try:
        sol = np.linalg.solve(C_ext, c0_ext)
        weights = sol[:-1]
    except np.linalg.LinAlgError:
        weights = np.full(n, np.nan)
    
    return {label: w for label, w in zip(observed_points.keys(), weights)}

def plot_screening_effect(yA, yB, yC, yD, show_A, show_B, show_C, show_D, range_, sill):
    observed_points = {}
    if show_A:
        observed_points['A'] = np.array([points_x['A'], yA])
    if show_B:
        observed_points['B'] = np.array([points_x['B'], yB])
    if show_C:
        observed_points['C'] = np.array([points_x['C'], yC])
    if show_D:
        observed_points['D'] = np.array([points_x['D'], yD])
    
    w_simple = compute_simple_weights(observed_points, range_, sill)
    w_ordinary = compute_ordinary_weights(observed_points, range_, sill)
    
    fig, axs = plt.subplots(1, 2, figsize=(14,6), sharey=True)
    
    def plot_weights(ax, weights, title):
        ax.set_title(title, fontsize=14)
        ax.set_xlim(15, 75)
        ax.set_ylim(30, 75)
        ax.set_aspect('equal')
        ax.grid(True)
        ax.scatter(points_x['E'], 50, c='black', marker='X', s=150)  # Point fixe E
        
        # Cercle portée autour de E
        ax.add_patch(Circle((points_x['E'], 50), radius=range_, edgecolor='gray', linestyle='--', fill=False))
        
        # Affichage des points observés
        for i, label in enumerate(['A', 'B', 'C', 'D']):
            if label in observed_points:
                w = weights.get(label, 0)
                # normalisation couleur
                norm_w = 0.5
                if weights:
                    w_vals = list(weights.values())
                    if max(w_vals) != min(w_vals):
                        norm_w = (w - min(w_vals)) / (max(w_vals) - min(w_vals))
                color = plt.cm.coolwarm(norm_w)
                ax.scatter(observed_points[label][0], observed_points[label][1], c=[color], s=300, edgecolor='k')
                # Décalage variable du texte pour éviter chevauchement
                x_text = observed_points[label][0] + 0.5
                y_text = observed_points[label][1] + 1.5 + 1.2
                ax.text(
                    x_text,
                    y_text,
                    f"{label}\n{w:.2f}",
                    fontsize=12,
                    bbox=dict(facecolor='white', alpha=0.7, edgecolor='none', pad=1)
                )
            else:
                # points cachés en gris clair petits
                x = points_x[label]
                y = {'A': yA, 'B': yB, 'C': yC, 'D': yD}[label]
                ax.scatter(x, y, c='lightgray', s=60, alpha=0.3)
                ax.text(x+0.3, y+0.3, label, fontsize=12, color='gray')
    
    plot_weights(axs[0], w_simple, "Krigeage simple")
    plot_weights(axs[1], w_ordinary, "Krigeage ordinaire")
    
    plt.tight_layout()
    plt.show()


# Sliders
yA_slider = widgets.FloatSlider(value=50, min=35, max=70, step=1, description='A')
yB_slider = widgets.FloatSlider(value=50, min=35, max=70, step=1, description='B')
yC_slider = widgets.FloatSlider(value=50, min=35, max=70, step=1, description='C')
yD_slider = widgets.FloatSlider(value=50, min=35, max=70, step=1, description='D')

show_A = widgets.Checkbox(value=True, description='Afficher A')
show_B = widgets.Checkbox(value=True, description='Afficher B')
show_C = widgets.Checkbox(value=True, description='Afficher C')
show_D = widgets.Checkbox(value=True, description='Afficher D')

range_slider = widgets.FloatSlider(value=15, min=5, max=40, step=1, description='Portée')
sill_slider = widgets.FloatSlider(value=1.0, min=0.5, max=2.0, step=0.1, description='Sill')

widgets.interact(
    plot_screening_effect,
    yA=yA_slider,
    yB=yB_slider,
    yC=yC_slider,
    yD=yD_slider,
    show_A=show_A,
    show_B=show_B,
    show_C=show_C,
    show_D=show_D,
    range_=range_slider,
    sill=sill_slider,
)