# Federated learning framework
pip install flwr torch torchvision

# Blockchain interaction (optional)
pip install web3 eth-brownie

# Privacy-preserving ML (optional)
pip install pysyft

"""
Federated Learning Server using Flower
Aggregates model updates from distributed clients
"""
import flwr as fl
from typing import List, Tuple
from flwr.common import Metrics

def weighted_average(metrics: List[Tuple[int, Metrics]]) -> Metrics:
    """Aggregate accuracy metrics from clients weighted by dataset size"""
    accuracies = [num_examples * m["accuracy"] for num_examples, m in metrics]
    examples = [num_examples for num_examples, _ in metrics]

    return {"accuracy": sum(accuracies) / sum(examples)}

# Configure federated learning strategy
strategy = fl.server.strategy.FedAvg(
    fraction_fit=0.3,  # Sample 30% of clients for training each round
    fraction_evaluate=0.2,  # Sample 20% of clients for evaluation
    min_fit_clients=3,  # Minimum clients required for training
    min_evaluate_clients=2,  # Minimum clients for evaluation
    min_available_clients=3,  # Wait for at least 3 clients
    evaluate_metrics_aggregation_fn=weighted_average,
)

# Start Flower server
fl.server.start_server(
    server_address="0.0.0.0:8080",
    config=fl.server.ServerConfig(num_rounds=10),  # 10 training rounds
    strategy=strategy,
)

"""
Federated Learning Client using Flower
Trains model locally on private data
"""
import flwr as fl
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# Define a simple CNN model
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.max_pool2d(x, 2)
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, 2)
        x = torch.flatten(x, 1)
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

# Load local training data (simulating private client data)
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

trainset = datasets.MNIST('./data', train=True, download=True, transform=transform)
trainloader = DataLoader(trainset, batch_size=32, shuffle=True)

testset = datasets.MNIST('./data', train=False, download=True, transform=transform)
testloader = DataLoader(testset, batch_size=32)

# Initialize model
model = SimpleCNN()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

# Define Flower client
class FlowerClient(fl.client.NumPyClient):
    def get_parameters(self, config):
        """Return current model parameters"""
        return [val.cpu().numpy() for _, val in model.state_dict().items()]

    def set_parameters(self, parameters):
        """Update model with aggregated parameters from server"""
        params_dict = zip(model.state_dict().keys(), parameters)
        state_dict = {k: torch.tensor(v) for k, v in params_dict}
        model.load_state_dict(state_dict, strict=True)

    def fit(self, parameters, config):
        """Train model on local data"""
        self.set_parameters(parameters)

        optimizer = optim.Adam(model.parameters(), lr=0.001)
        criterion = nn.CrossEntropyLoss()

        model.train()
        for epoch in range(1):  # Train for 1 local epoch per round
            for images, labels in trainloader:
                images, labels = images.to(device), labels.to(device)

                optimizer.zero_grad()
                outputs = model(images)
                loss = criterion(outputs, labels)
                loss.backward()
                optimizer.step()

        return self.get_parameters(config={}), len(trainloader.dataset), {}

    def evaluate(self, parameters, config):
        """Evaluate model on local test data"""
        self.set_parameters(parameters)

        model.eval()
        loss = 0.0
        correct = 0

        criterion = nn.CrossEntropyLoss()

        with torch.no_grad():
            for images, labels in testloader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                loss += criterion(outputs, labels).item()
                correct += (outputs.argmax(1) == labels).sum().item()

        accuracy = correct / len(testloader.dataset)
        return loss, len(testloader.dataset), {"accuracy": accuracy}

# Start Flower client
fl.client.start_numpy_client(
    server_address="localhost:8080",
    client=FlowerClient()
)

# Terminal 1: Start server
python server.py

# Terminal 2-4: Start 3 clients (simulating different institutions)
python client.py  # Client 1
python client.py  # Client 2
python client.py  # Client 3

"""
Publish and monetize AI models using Ocean Protocol
"""
from ocean_lib.ocean.ocean import Ocean
from ocean_lib.web3_internal.wallet import Wallet
from ocean_lib.data_provider.data_service_provider import DataServiceProvider
import os

# Configuration
config = {
    'network': 'https://rpc.polygon.technology/',  # Polygon mainnet
    'aquarius': 'https://v4.aquarius.oceanprotocol.com',
    'provider': 'https://v4.provider.polygon.oceanprotocol.com'
}

# Initialize Ocean instance
ocean = Ocean(config)

# Load wallet (never hardcode private keys in production!)
private_key = os.getenv('PRIVATE_KEY')
wallet = Wallet(ocean.web3, private_key=private_key)

# Publish an AI model as a data asset
def publish_model(model_file_url, name, description):
    """
    Publish ML model to Ocean Protocol marketplace

    Args:
        model_file_url: URL to model file (IPFS, Arweave, etc.)
        name: Model name
        description: Model description

    Returns:
        DID (Decentralized Identifier) of published model
    """
    metadata = {
        "created": "2025-01-01T00:00:00Z",
        "updated": "2025-01-01T00:00:00Z",
        "description": description,
        "name": name,
        "type": "dataset",
        "author": "YourOrganization",
        "license": "MIT",
        "files": [{
            "url": model_file_url,
            "method": "GET",
            "contentType": "application/octet-stream"
        }]
    }

    # Create data asset
    data_nft = ocean.create_data_nft(name, symbol="MODEL", wallet=wallet)
    datatoken = data_nft.create_datatoken(name + " Token", symbol="MODELT", wallet=wallet)

    # Set pricing (e.g., 0.1 OCEAN per access)
    datatoken.create_dispenser(wallet=wallet)

    # Publish metadata
    ddo = ocean.create_asset(metadata, wallet=wallet, data_nft_address=data_nft.address)

    return ddo.did

# Access a published model
def access_model(did):
    """
    Purchase and download model from Ocean Protocol

    Args:
        did: Decentralized Identifier of the model

    Returns:
        Path to downloaded model file
    """
    asset = ocean.resolve_asset(did)

    # Purchase access (requires OCEAN tokens)
    order_tx_id = ocean.pay_for_access_service(asset, wallet=wallet)

    # Download model
    file_path = ocean.download_asset_files(order_tx_id, wallet=wallet)

    return file_path

# Example usage
model_did = publish_model(
    model_file_url="https://ipfs.io/ipfs/Qm...",
    name="Fraud Detection Model v1.0",
    description="XGBoost model trained on 1M transactions"
)

print(f"Model published with DID: {model_did}")

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

/**
 * @title AIModelRegistry
 * @notice Decentralized registry for AI model metadata and access control
 */
contract AIModelRegistry {
    struct Model {
        string name;
        string ipfsHash;  // IPFS hash of model file
        address owner;
        uint256 version;
        uint256 timestamp;
        bool isPublic;
    }

    mapping(bytes32 => Model) public models;
    mapping(bytes32 => mapping(address => bool)) public accessControl;

    event ModelRegistered(bytes32 indexed modelId, address indexed owner, string name);
    event AccessGranted(bytes32 indexed modelId, address indexed user);

    /**
     * @notice Register a new AI model
     * @param name Model name
     * @param ipfsHash IPFS content identifier
     * @param isPublic Whether model is publicly accessible
     */
    function registerModel(
        string memory name,
        string memory ipfsHash,
        bool isPublic
    ) external returns (bytes32) {
        bytes32 modelId = keccak256(abi.encodePacked(msg.sender, name, block.timestamp));

        models[modelId] = Model({
            name: name,
            ipfsHash: ipfsHash,
            owner: msg.sender,
            version: 1,
            timestamp: block.timestamp,
            isPublic: isPublic
        });

        emit ModelRegistered(modelId, msg.sender, name);
        return modelId;
    }

    /**
     * @notice Grant access to a private model
     * @param modelId Model identifier
     * @param user Address to grant access
     */
    function grantAccess(bytes32 modelId, address user) external {
        require(models[modelId].owner == msg.sender, "Not model owner");
        accessControl[modelId][user] = true;
        emit AccessGranted(modelId, user);
    }

    /**
     * @notice Check if user has access to model
     * @param modelId Model identifier
     * @param user User address
     */
    function hasAccess(bytes32 modelId, address user) external view returns (bool) {
        return models[modelId].isPublic || accessControl[modelId][user];
    }
}

"""
Distributed model inference using Ray
"""
import ray
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# Initialize Ray cluster
ray.init(address='auto')  # Connect to existing cluster, or ray.init() for local

@ray.remote(num_gpus=0.25)  # Each replica gets 1/4 GPU
class ModelReplica:
    def __init__(self, model_name):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
        self.model.eval()

    def predict(self, text):
        """Run inference on input text"""
        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=512)

        with torch.no_grad():
            outputs = self.model(**inputs)
            predictions = torch.softmax(outputs.logits, dim=-1)

        return predictions.tolist()

# Deploy 4 model replicas across cluster
model_name = "distilbert-base-uncased-finetuned-sst-2-english"
replicas = [ModelReplica.remote(model_name) for _ in range(4)]

# Load balancing: round-robin across replicas
def distributed_inference(texts):
    """Distribute inference requests across replicas"""
    futures = []

    for i, text in enumerate(texts):
        replica = replicas[i % len(replicas)]
        future = replica.predict.remote(text)
        futures.append(future)

    # Wait for all predictions
    results = ray.get(futures)
    return results

# Example usage
texts = [
    "This product is amazing!",
    "Terrible experience, would not recommend.",
    "Average quality, nothing special.",
    "Best purchase I've made this year!"
]

predictions = distributed_inference(texts)
for text, pred in zip(texts, predictions):
    print(f"Text: {text[:50]}... -> Prediction: {pred}")

"""
Add differential privacy to federated learning using Opacus
"""
from opacus import PrivacyEngine
import torch.nn as nn
import torch.optim as optim

model = SimpleCNN()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# Attach privacy engine
privacy_engine = PrivacyEngine()

model, optimizer, trainloader = privacy_engine.make_private(
    module=model,
    optimizer=optimizer,
    data_loader=trainloader,
    noise_multiplier=1.1,  # Noise level for privacy
    max_grad_norm=1.0,  # Gradient clipping
)

# Training loop now has differential privacy guarantees
# epsilon and delta can be computed via privacy_engine.get_epsilon(delta=1e-5)

"""
Privacy-preserving inference using PySyft
"""
import torch
import syft as sy

# Create virtual workers (simulating different institutions)
hook = sy.TorchHook(torch)
alice = sy.VirtualWorker(hook, id="alice")
bob = sy.VirtualWorker(hook, id="bob")

# Split data between workers (each only sees their share)
data = torch.tensor([[1, 2], [3, 4], [5, 6]])
data_alice = data[0].send(alice)
data_bob = data[1].send(bob)

# Model performs computation on encrypted data
model = nn.Linear(2, 1)
encrypted_model = model.fix_precision().share(alice, bob)

# Inference on encrypted data
result_alice = encrypted_model(data_alice.fix_precision())
result_bob = encrypted_model(data_bob.fix_precision())

# Decrypt results
result_alice = result_alice.get().float_precision()
result_bob = result_bob.get().float_precision()