import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
from sklearn.datasets import make_moons, make_circles, make_classification
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import sys
import os

# Ensure ellice is in path
try:
    import ellice
except ImportError:
    sys.path.append(os.path.abspath("."))
    import ellice

# 1. Generate Synthetic Data (2 overlapping Gaussians)
def generate_data(n_samples=500):
    # Use make_moons or classification for interesting boundaries
    # Overlapping Gaussians
    np.random.seed(42)
    X1 = np.random.randn(n_samples // 2, 2) + np.array([1.0, 1.0])
    X2 = np.random.randn(n_samples // 2, 2) + np.array([-1.0, -1.0])
    X = np.vstack([X1, X2])
    y = np.hstack([np.zeros(n_samples // 2), np.ones(n_samples // 2)])
    
    # Standardize
    scaler = StandardScaler()
    X = scaler.fit_transform(X)
    
    return X, y

# 2. Define Simple NN
class SimpleNN(nn.Sequential):
    def __init__(self):
        super().__init__(
            nn.Linear(2, 16),
            nn.ReLU(),
            nn.Linear(16, 16),
            nn.ReLU(),
            nn.Linear(16, 1)
        )

def train_model(model, X, y, epochs=200):
    X_t = torch.FloatTensor(X)
    y_t = torch.FloatTensor(y).unsqueeze(1)
    
    criterion = nn.BCEWithLogitsLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.01)
    
    model.train()
    for _ in range(epochs):
        optimizer.zero_grad()
        out = model(X_t)
        loss = criterion(out, y_t)
        loss.backward()
        optimizer.step()
    
    model.eval()
    return model

# 3. Plotting Helper
def plot_decision_boundary(model, X, y, title="Decision Boundary"):
    # Define grid
    x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1
    y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1
    xx, yy = np.meshgrid(np.arange(x_min, x_max, 0.05),
                         np.arange(y_min, y_max, 0.05))
    
    grid_tensor = torch.FloatTensor(np.c_[xx.ravel(), yy.ravel()])
    
    with torch.no_grad():
        Z = torch.sigmoid(model(grid_tensor)).reshape(xx.shape).numpy()
        
    plt.contourf(xx, yy, Z, alpha=0.8, cmap=plt.cm.RdBu, levels=np.linspace(0, 1, 11))
    plt.scatter(X[:, 0], X[:, 1], c=y, cmap=plt.cm.RdBu_r, edgecolors='k', alpha=0.6)
    plt.contour(xx, yy, Z, levels=[0.5], colors='black', linewidths=2)
    plt.title(title)

def get_worst_model_boundary(exp, query_x, base_model, predicted_class=None):
    """
    Reconstructs the 'worst-case model' corresponding to a specific query point.
    
    If predicted_class is 1, worst-case means minimize logit (push to 0).
    If predicted_class is 0, worst-case means maximize logit (push to 1).
    """
    
    # 1. Get Penultimate Features for Query to calculate direction
    q_tensor = torch.FloatTensor(query_x).unsqueeze(0).to(exp.generator.device)
    with torch.no_grad():
        h_flat = exp.generator._get_penult_features(q_tensor)
        bias = torch.ones(h_flat.size(0), 1, device=exp.generator.device, dtype=exp.generator.dtype)
        h_aug = torch.cat([h_flat, bias], dim=1)
        
        # 2. Compute Worst Theta (minimizes logit by default: theta_star - direction)
        worst_theta = exp.generator._compute_worst_model(h_aug) # [dim_h + 1]
        
        if predicted_class == 0:
            # If predicted class is 0, we want to MAXIMIZE the logit to flip it to 1
            # The standard worst_theta = center - direction (minimizes)
            # We want theta = center + direction
            # direction = center - worst_theta
            # target = center + (center - worst_theta) = 2*center - worst_theta
            worst_theta = 2 * exp.generator.omega_c - worst_theta
    
    # 3. Construct New Model with Worst Theta
    # Clone the base model
    import copy
    worst_model = copy.deepcopy(base_model)
    
    # The last layer of SimpleNN is at index 4 because it's a Sequential
    last_layer = worst_model[4]
    
    # worst_theta contains [weights..., bias]
    w_new = worst_theta[:-1]
    b_new = worst_theta[-1]
    
    with torch.no_grad():
        last_layer.weight.data = w_new.unsqueeze(0)
        last_layer.bias.data = b_new.unsqueeze(0)
        
    return worst_model

def main():
    # Setup
    X, y = generate_data(n_samples=2000)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    model = SimpleNN()
    model = train_model(model, X_train, y_train)
    
    # Initialize ElliCE
    df_full = pd.DataFrame(X_train, columns=['f1', 'f2'])
    df_full['target'] = y_train
    data = ellice.Data(df_full, target_column='target')
    
    # Use backend='pytorch' explicitly and pass initialized generator to access internals easily later
    # Or just use standard API and access internal generator via private attribute if needed (not recommended usually but fine for demo)
    # Actually, let's instantiate the Explainer which creates the Generator
    exp = ellice.Explainer(model, data, backend='pytorch')
    
    # Hack: Initialize generator once by calling generate (or manually) to setup Q matrix
    # This is needed because generator is created inside generate_counterfactuals usually.
    # We'll just instantiate a ContinuousGenerator manually to access its internals.
    from ellice.generators.continuous import ContinuousGenerator
    gen = ContinuousGenerator(exp.model, data, eps=0.02, reg_coef=0.001)
    
    # Attach this generator to our explainer or just use it
    exp.generator = gen # Mock attachment for helper function
    
    # Plotting
    plt.figure(figsize=(15, 8)) # Adjusted size for better layout
    
    # 1. Plot Original Decision Boundary
    plt.subplot(2, 3, 1)
    plot_decision_boundary(model, X_train, y_train, title="Original Model Boundary")
    
    # 2. Select 5 Random Points from Test Set
    indices = np.random.choice(len(X_test), 5, replace=False)
    random_points = X_test[indices]
    
    for i, query_pt in enumerate(random_points):
        # Predict class to determine worst direction
        with torch.no_grad():
            logit = model(torch.FloatTensor(query_pt).unsqueeze(0)).item()
            prob = 1 / (1 + np.exp(-logit))
            pred_class = 1 if prob > 0.5 else 0
            
        # Get worst model for this specific point
        worst_model = get_worst_model_boundary(exp, query_pt, model, predicted_class=pred_class)
        
        target_class = 1 - pred_class
        plt.subplot(2, 3, i + 2)
        plot_decision_boundary(worst_model, X_train, y_train, title=f"Worst Model (Pred {pred_class}->{target_class})")
        
        # Highlight the query point
        plt.scatter(query_pt[0], query_pt[1], c='gold', s=200, marker='*', edgecolors='black', label='Query', zorder=5)
        # Add arrow pointing to point
        # plt.annotate('Query', xy=query_pt, xytext=(query_pt[0]+0.5, query_pt[1]+0.5),
        #             arrowprops=dict(facecolor='black', shrink=0.05))
        # plt.legend() # Legend might clutter small plots

    plt.tight_layout()
    plt.savefig('boundaries_demo.png')
    print("Saved boundaries_demo.png")
    plt.show()

if __name__ == "__main__":
    main()

