import pandas as pd

data = pd.read_csv("data/breast-cancer.csv")
diagnosis = data["diagnosis"]
labels = diagnosis.map({"M": 1, "B": -1}).values
features = data.drop(["id", "diagnosis"], axis=1)

#VGG coomet next three lines to run with the originla data above
data = pd.read_csv("data/formatted_titanic.csv")
labels = data["survived"]+2*data["sex"]
features = data.drop(["survived","sex"], axis=1)

data.shape

# Data visualization
import matplotlib.pyplot as plt
import umap
from sklearn.preprocessing import StandardScaler


scaler = StandardScaler()
scaled_features = scaler.fit_transform(features)

reducer = umap.UMAP(n_components=3)
embedding = reducer.fit_transform(scaled_features)

embedding.shape


plt.scatter(
    embedding[:, 0],embedding[:, 1], c=labels)

plt.scatter(
    embedding[:, 0],embedding[:, 2], c=labels)


plt.scatter(
    embedding[:, 1],embedding[:, 2], c=labels)

fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.scatter(embedding[:, 0],embedding[:, 1],embedding[:, 2], c=labels)
plt.show()

from sklearn.decomposition import PCA

pca = PCA(n_components=0.75)
features_pca = pca.fit_transform(features)
pca.n_components_



features_pca.shape

plt.scatter(features_pca[:, 0],features_pca[:, 1], c=labels)

plt.scatter(features_pca[:, 0],features_pca[:, 2], c=labels)

plt.scatter(features_pca[:, 1],features_pca[:, 2], c=labels)

fig3d = plt.figure()
ax3d = fig3d.add_subplot(projection='3d')
ax3d.scatter(features_pca[:, 0],features_pca[:, 1],features_pca[:, 2], c=labels)
plt.show()

from qcfl.preprocessing import scale_for_angle_encoding

features_scaled = scale_for_angle_encoding(features_pca)

from utils import split_data, write_shards


dataset = pd.DataFrame(data=features_scaled)
dataset['target'] = labels
dataset = dataset.sample(frac=1)
shards = split_data(dataset, [0.3, 0.1, 0.2, 0.2, 0.2])
write_shards(shards, 'cancer')

from pennylane import AdamOptimizer
from qcfl.penny.classifier import QuantumClassifier
from qcfl.penny.models import SimplQMLModel
from qcfl.penny.federated_qml import create_clients, FederatedServer, to_numpy_dataset

num_qubits = features_pca.shape[1]
num_layers = 3
batch_size = 5

clients = create_clients(shards[0:4], lambda: QuantumClassifier(num_qubits, num_layers, batch_size,
                                                                AdamOptimizer(),
                                                                SimplQMLModel(num_qubits)))

classifier = QuantumClassifier(num_qubits, num_layers, batch_size,
                                   AdamOptimizer(),
                                   SimplQMLModel(num_qubits))
server = FederatedServer(classifier, clients, client_epochs=10)

# Train
from pennylane import numpy as np

weights = 0.01 * np.random.randn(num_layers, num_qubits, 3, requires_grad=True)
bias = np.array(0.0, requires_grad=True)
trained_weights, trained_bias = server.train(weights, bias, iterations=15)

from sklearn.metrics import classification_report

test_features, test_labels = to_numpy_dataset(shards[-1])
test_predictions = [classifier.classify(trained_weights, trained_bias, f) for f in test_features]

target_names = ['Malignant', 'Benign']
target_names = [str(i) for i in set(labels.values).union()]#VGG comment this line for cancer

cr = classification_report(test_labels, test_predictions, target_names=target_names)
print(cr)

set(labels.values).union()

from typing import List, Tuple
from flwr.common import Metrics

def fit_round(server_round: int):
    """Send round number to client."""
    return {"server_round": server_round}


def weighted_average(metrics: List[Tuple[int, Metrics]]) -> Metrics:
    # Multiply accuracy of each client by number of examples used
    accuracies = [num_examples * m["accuracy"] for num_examples, m in metrics]
    examples = [num_examples for num_examples, _ in metrics]

    # Aggregate and return custom metric (weighted average)
    return {"accuracy": sum(accuracies) / sum(examples)}


def get_evaluate_fn(classifier, X_test, y_test):
    """Return an evaluation function for server-side evaluation."""

    # The `evaluate` function will be called after every round
    def evaluate(server_round, parameters: fl.common.NDArrays, config):
        # Update model with the latest parameters
        weights, bias = parameters[0], parameters[1]
        classifier.set_parameters(weights, bias)
        loss, accuracy = classifier.evaluate(server_round, X_test, y_test)
        return loss, {"accuracy": accuracy}

    return evaluate


from qcfl.flower.flower_qml_client import PennylaneClient

num_qubits = features_pca.shape[1]#VGG
NUM_CLIENTS = 4
client_resources = None


def client_fn(cid: str) -> PennylaneClient:
    """Create a Flower client representing a single organization."""
    print(cid)
    index = int(cid)+1
    # Load a shard corresponding to this client
    X_train, y_train = load_shard('cancer', index)
    X_train = np.array(X_train, requires_grad=False)
    y_train = np.array(y_train, requires_grad=False)

    classifier = QuantumClassifier(num_qubits, 3, 5,
                                   AdamOptimizer(),
                                   SimplQMLModel(num_qubits))
    # Define Flower client
    return PennylaneClient(f'cancer{index}',
                           classifier,
                           X_train, y_train)


from utils import load_shard

# Load test data here to avoid the overhead of doing it in `evaluate` itself
X_test, y_test = load_shard('cancer', 5)

# Create a classifier to hold the final weights
classifier = QuantumClassifier(num_qubits, 3, 10,
                                   AdamOptimizer(),
                                   SimplQMLModel(num_qubits))

import flwr as fl


# Create FedAvg strategy
strategy = fl.server.strategy.FedAvg(
    fraction_fit=1.0,
    fraction_evaluate=0.5,
    min_available_clients=NUM_CLIENTS,
    on_fit_config_fn=fit_round,
    evaluate_fn=get_evaluate_fn(classifier, X_test, y_test),
    evaluate_metrics_aggregation_fn=weighted_average
)

# Run flower in simulation mode

fl.simulation.start_simulation(
    client_fn=client_fn,
    num_clients=NUM_CLIENTS,
    config=fl.server.ServerConfig(num_rounds=15),
    strategy=strategy,
    client_resources=client_resources,
)

#TODO: Plot loss, accuracy



# Compute metrics on final model

test_features, test_labels = to_numpy_dataset(shards[-1])
test_predictions = [classifier.classify(trained_weights, trained_bias, f) for f in test_features]
target_names = ['Malignant', 'Benign']
cr = classification_report(test_labels, test_predictions, target_names=target_names)
print(cr)

