import torch
import numpy as np
from scipy.optimize import minimize
from .Strategy import Strategy
import torch.nn as nn

class CertifiedRemoval(Strategy):
    """Implements Certified Removal for machine unlearning."""

    def __init__(self, model, data, labels, removal_bound, epsilon):
        super().__init__()
        self.model = model
        self.data = data
        self.labels = labels
        self.removal_bound = removal_bound
        self.epsilon = epsilon

    def _run(self, model: nn.Module) -> nn.Module:
        """Runs the certified removal algorithm."""
        # 1. Linear Model Creation
        # This is a simplification for demonstration purposes.  In a real implementation,
        # you'd use more sophisticated methods to learn the parameters of the
        # 'removal' model based on the example being removed.

        def linear_model(x):
            return torch.dot(x, torch.tensor([1, 1])) # Simplified Linear Model

        # 2. Optimization for Parameter Adjustment
        # Optimize the parameter values to minimize the loss while staying within bounds.
        original_params = torch.tensor([0.0, 0.0])  # Initial parameters for linear model

        def objective_function(params):
            new_model = linear_model #use same function as defined above
            return torch.sum(((new_model(self.data[0]) - self.labels)**2))
                
        result = minimize(objective_function, original_params, method='L-BFGS-B', bounds=[(-self.removal_bound, self.removal_bound)], options={'maxiter': 100})

        if not result.success:
            print("Warning: Optimization failed!")
            print(result.message)
            return model #Return original if optimization fails

        new_params = result.x
        # 3. New Model Creation

        new_model = lambda x: torch.dot(x, new_params)
        return new_model


if __name__ == '__main__':
    # Example Usage - Synthetic Data for Demonstration
    np.random.seed(42)  # For reproducibility
    n_samples = 100
    X = np.random.randn(n_samples, 2)
    y = (X[:, 0] + X[:, 1] > 0).astype(int)

    # Create a simple linear model for demonstration
    model = nn.Linear(2, 1)  # Simple linear classifier - PyTorch Version
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01) # Optimizer for training the linear model

    # Train a Linear Model
    for _ in range(100): #training loop
        optimizer.zero_grad()
        predictions = model(X)
        loss = torch.sum((predictions - y)**2)
        loss.backward()
        optimizer.step()

    # Define parameters for Certified Removal
    removal_bound = 1.0
    epsilon = 0.1

    # Create the CertifiedRemoval object with the trained model, data and labels
    certified_removal_obj = CertifiedRemoval(model, X, y, removal_bound, epsilon)

    # Run Certified Removal
    new_model = certified_removal_obj.apply(model)