import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import accuracy_score, classification_report
import os
import sys

# Ensure ellice is in path (for local execution, in colab you'd install it)
# Assuming this script is run from the root of the project or ellice is installed
try:
    import ellice
except ImportError:
    # Fallback for local dev structure
    sys.path.append(os.path.abspath("."))
    import ellice

# --------------------------------------------------------------------------- #
# 1. Setup & Data Loading
# --------------------------------------------------------------------------- #
def load_mnist_binary(batch_size=64):
    """Load MNIST and filter for classes 0 and 1."""
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    # Download Data
    train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
    test_dataset = datasets.MNIST('./data', train=False, download=True, transform=transform)

    # Filter for 0 and 1
    train_indices = (train_dataset.targets == 0) | (train_dataset.targets == 1)
    test_indices = (test_dataset.targets == 0) | (test_dataset.targets == 1)

    train_subset = Subset(train_dataset, np.where(train_indices)[0])
    test_subset = Subset(test_dataset, np.where(test_indices)[0])

    train_loader = DataLoader(train_subset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_subset, batch_size=batch_size, shuffle=False)
    
    # Create ElliCE compatible Data object
    # We need a dataframe for ElliCE's Data class. 
    # For Images, ElliCE expects flattened features if passed as dataframe, 
    # OR we can pass the torch tensors directly if we adapt the workflow.
    # Standard ElliCE workflow takes a DataFrame.
    # Let's extract a subset of training data to a DataFrame for ElliCE initialization (Rashomon set calc)
    print("Extracting data for ElliCE initialization...")
    
    X_list = []
    y_list = []
    
    # Take a subset for faster Ellipsoid computation (e.g., 2000 samples)
    count = 0
    for img, label in train_loader:
        # Flatten images: 1x28x28 -> 784
        flat = img.view(img.size(0), -1).numpy()
        X_list.append(flat)
        y_list.append(label.numpy())
        count += len(label)
        if count > 2000: 
            break
            
    X_np = np.concatenate(X_list)
    y_np = np.concatenate(y_list)
    
    feature_names = [f"pixel_{i}" for i in range(X_np.shape[1])]
    df = pd.DataFrame(X_np, columns=feature_names)
    df['target'] = y_np
    
    ellice_data = ellice.Data(df, target_column='target')
    
    return train_loader, test_loader, ellice_data, feature_names

# --------------------------------------------------------------------------- #
# 2. Model Definition
# --------------------------------------------------------------------------- #
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 1) # Binary classification

    def forward(self, x):
        # x: [Batch, 1, 28, 28]
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2(x), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        # Last layer is linear, ElliCE will handle sigmoid
        x = self.fc2(x) 
        return x

def train_model(model, train_loader, epochs=3):
    criterion = nn.BCEWithLogitsLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    model.train()
    print(f"\nTraining CNN for {epochs} epochs...")
    for epoch in range(epochs):
        total_loss = 0
        for batch_idx, (data, target) in enumerate(train_loader):
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target.float().unsqueeze(1))
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
            
        print(f"Epoch {epoch+1}: Loss {total_loss / len(train_loader):.4f}")

def evaluate_model(model, test_loader):
    model.eval()
    all_preds = []
    all_targets = []
    
    with torch.no_grad():
        for data, target in test_loader:
            output = model(data)
            preds = (torch.sigmoid(output) > 0.5).long().squeeze()
            all_preds.extend(preds.numpy())
            all_targets.extend(target.numpy())
            
    acc = accuracy_score(all_targets, all_preds)
    print(f"\nTest Accuracy: {acc:.4f}")
    print(classification_report(all_targets, all_preds))
    return all_targets, all_preds

# --------------------------------------------------------------------------- #
# 3. ElliCE Wrapper for CNN
# --------------------------------------------------------------------------- #
# We need to wrap the CNN so ElliCE can interpret it as Feature Extractor + Linear Layer.
# Our SimpleCNN has `fc2` as the last linear layer.
# ElliCE's `PyTorchModel` wrapper automatically handles this if `split_model` logic works for arbitrary modules.
# The default `split_model` assumes `nn.Sequential`. We should override or ensure it works.
# Let's rely on `ellice.models.wrappers.PyTorchModel`'s `split_model` which looks for the last linear layer.
# However, `split_model` in current code tries to reconstruct the penult model by slicing children. 
# This FAILS for custom classes like SimpleCNN defined in `forward`.
# Solution: We define a custom wrapper or modify SimpleCNN to be Sequential-like or modify ElliCE wrapper.
# Easiest for this script: Define a custom Wrapper inheriting from ElliCE's ModelWrapper.

class CNNWrapper(ellice.models.wrappers.ModelWrapper):
    def __init__(self, model, device):
        super().__init__(model, backend='pytorch')
        self.device = device
        self.model.to(self.device)
        
    def get_torch_model(self):
        return self.model
        
    def split_model(self):
        # Separate the last layer (fc2) from the rest
        last_layer = self.model.fc2
        
        # Create a feature extractor module
        class FeatureExtractor(nn.Module):
            def __init__(self, original_model):
                super().__init__()
                self.conv1 = original_model.conv1
                self.conv2 = original_model.conv2
                self.fc1 = original_model.fc1
                
            def forward(self, x):
                # Reshape input if it comes flattened from ElliCE
                if x.dim() == 2 and x.shape[1] == 784:
                    x = x.view(-1, 1, 28, 28)
                    
                x = F.relu(F.max_pool2d(self.conv1(x), 2))
                x = F.relu(F.max_pool2d(self.conv2(x), 2))
                x = x.view(-1, 320)
                x = F.relu(self.fc1(x))
                return x
                
        penult = FeatureExtractor(self.model).to(self.device)
        
        # Flatten parameters of the last layer
        weight = last_layer.weight.detach().view(-1)
        bias = last_layer.bias.detach()
        theta = torch.cat([weight, bias])
        
        return penult, theta
        
    def predict_proba(self, X):
        # Handle numpy input
        if isinstance(X, np.ndarray):
            X = torch.from_numpy(X).float().to(self.device)
            
        # Reshape if flat
        if X.dim() == 2 and X.shape[1] == 784:
            X = X.view(-1, 1, 28, 28)
            
        self.model.eval()
        with torch.no_grad():
            logits = self.model(X)
            probs_1 = torch.sigmoid(logits)
            probs_0 = 1 - probs_1
            return torch.cat([probs_0, probs_1], dim=1).cpu().numpy()

# --------------------------------------------------------------------------- #
# 4. Main Execution
# --------------------------------------------------------------------------- #
def main():
    # Settings
    device = torch.device("cpu")
    print(f"Using device: {device}")
    
    # 1. Data
    train_loader, test_loader, ellice_data, feature_names = load_mnist_binary()
    
    # 2. Train CNN
    cnn = SimpleCNN()
    train_model(cnn, train_loader)
    evaluate_model(cnn, test_loader)
    
    # 3. Initialize ElliCE
    # We use our custom wrapper
    model_wrapper = CNNWrapper(cnn, device)
    
    # We must patch the Explainer to accept an already wrapped model or instantiate it manually
    # The standard Explainer constructor calls `load_model`.
    # We can manually construct the generator or modify Explainer.
    # Easiest: Just instantiate Explainer and then overwrite self.model
    exp = ellice.Explainer(cnn, ellice_data, backend='pytorch') # Dummy init
    exp.model = model_wrapper # Overwrite with correct wrapper
    
    # 4. Generate Counterfactuals
    # Generate CFs for 6 examples
    num_examples = 6
    
    # Get samples from test set
    test_dataset = test_loader.dataset
    indices = range(num_examples) # First 6 examples
    
    # Prepare plot
    rows = num_examples
    cols = 3
    fig, axes = plt.subplots(rows, cols, figsize=(15, 5 * rows))
    
    for i, idx in enumerate(indices):
        img, label = test_dataset[idx] # img is [1, 28, 28]
        flat_query = img.view(-1).numpy() # [784]
        
        # Create Query DataFrame
        query_df = pd.DataFrame([flat_query], columns=feature_names)
        query_series = query_df.iloc[0]
        
        orig_pred = cnn(img.unsqueeze(0)).item()
        orig_prob = torch.sigmoid(torch.tensor(orig_pred)).item()
        orig_class = int(orig_prob > 0.5)
        target_class = 1 - orig_class
        
        print(f"\n--- Example {i+1}/{num_examples} ---")
        print(f"Query Image Index: {idx}")
        print(f"True Label: {label}")
        print(f"Model Prediction: {orig_class} (Prob Class 1: {orig_prob:.4f})")
        print(f"Target Class: {target_class}")
        
        # Generate CF
        print("Generating Robust Counterfactual...")
        cf_df = exp.generate_counterfactuals(
            query_series,
            method='continuous',
            target_class=target_class,
            robustness_epsilon=0.01, 
            regularization_coefficient=0.001,
            return_probs=True,
            progress_bar=True,
            optimization_params={
                'max_iterations': 500, 
                'learning_rate': 0.05,
                'robustness_weight': 10.0,
                'proximity_weight': 0.5
            }
        )
        
        ax_row = axes[i]
        
        # Original
        ax_row[0].imshow(img.squeeze(), cmap='gray')
        ax_row[0].set_title(f"Ex {i}: Original (Class {orig_class})\nProb: {orig_prob:.3f}")
        ax_row[0].axis('off')
        
        if cf_df.empty:
            print("Failed to find CF.")
            ax_row[1].text(0.5, 0.5, "Failed to find CF", ha='center', va='center')
            ax_row[1].axis('off')
            ax_row[2].axis('off')
            continue

        cf_row = cf_df.iloc[0]
        cf_features = cf_row[feature_names].values
        
        # Counterfactual
        cf_img = cf_features.reshape(28, 28)
        
        # Get metrics
        model_prob_cf = cf_row['model_prob_class_1']
        if target_class == 0:
            prob_display = 1 - model_prob_cf
        else:
            prob_display = model_prob_cf
            
        robust_prob = cf_row['worst_case_prob_target']
        
        ax_row[1].imshow(cf_img, cmap='gray')
        ax_row[1].set_title(f"Counterfactual (Target {target_class})\nModel Prob: {prob_display:.3f}\nRobust Prob: {robust_prob:.3f}")
        ax_row[1].axis('off')
        
        # Difference
        diff = cf_img - img.squeeze().numpy()
        im = ax_row[2].imshow(diff, cmap='seismic', vmin=-1, vmax=1)
        ax_row[2].set_title("Difference (CF - Original)")
        ax_row[2].axis('off')
        plt.colorbar(im, ax=ax_row[2])
    
    plt.tight_layout()
    plt.savefig('mnist_cf_result_multiple.png')
    print("\nResult saved to 'mnist_cf_result_multiple.png'")
    # plt.show()

if __name__ == "__main__":
    main()

