#!/usr/bin/env python3
import flwr as fl
from opacus import PrivacyEngine
from torch.utils.data import DataLoader

from .model import Model
from .utils import params
import torch
import numpy as np

class FlowerClient(fl.client.NumPyClient):
    """
    FlowerClient class for creating MEDfl clients.

    Attributes:
        cid (str): Client ID.
        local_model (Model): Local model of the federated learning network.
        trainloader (DataLoader): DataLoader for training data.
        valloader (DataLoader): DataLoader for validation data.
        diff_priv (bool): Flag indicating whether to use differential privacy.
    """
    def __init__(self, cid: str, local_model: Model, trainloader: DataLoader, valloader: DataLoader, diff_priv: bool = params["diff_privacy"]):
        """
        Initializes the FlowerClient instance.

        Args:
            cid (str): Client ID.
            local_model (Model): Local model of the federated learning network.
            trainloader (DataLoader): DataLoader for training data.
            valloader (DataLoader): DataLoader for validation data.
            diff_priv (bool): Flag indicating whether to use differential privacy.
        """
        self.cid = cid
        self.local_model = local_model
        self.trainloader = trainloader
        self.valloader = valloader 
        """      if torch.cuda.is_available():
            num_cuda_devices = torch.cuda.device_count()
            if num_cuda_devices > 0:
                device_idx = int(self.cid) % num_cuda_devices
                self.device = torch.device(f"cuda:{device_idx}")
                self.local_model.model.to(self.device)
            else:
                # Handle case where CUDA is available but no CUDA devices are found
                raise RuntimeError("CUDA is available, but no CUDA devices are found.")
        else:
            # Handle case where CUDA is not available
            self.device = torch.device("cpu")
            self.local_model.model.to(self.device) """

        self.device = torch.device("cpu")
        self.local_model.model.to(self.device)
        
        self.privacy_engine = PrivacyEngine(secure_mode=False)
        self.diff_priv = diff_priv
        self.epsilons = []
        self.accuracies = []
        self.losses = []
        self.train_round_losses = []
        self.train_round_metrics = []

        if self.diff_priv:
            model, optimizer, self.trainloader = self.privacy_engine.make_private_with_epsilon(
                module=self.local_model.model.train(),
                optimizer=self.local_model.optimizer,
                data_loader=self.trainloader,
                epochs=params["train_epochs"],
                target_epsilon=float(params["EPSILON"]),
                target_delta= float(params["DELTA"]),
                max_grad_norm=params["MAX_GRAD_NORM"],
            )
            setattr(self.local_model, "model", model)
            setattr(self.local_model, "optimizer", optimizer)
        self.validate()

    def validate(self):
        """Validates cid, local_model, trainloader, valloader."""
        if not isinstance(self.cid, str):
            raise TypeError("cid argument must be a string")

        if not isinstance(self.local_model, Model):
            raise TypeError("local_model argument must be a MEDfl.LearningManager.model.Model")

        if not isinstance(self.trainloader, DataLoader):
            raise TypeError("trainloader argument must be a torch.utils.data.dataloader")

        if not isinstance(self.valloader, DataLoader):
            raise TypeError("valloader argument must be a torch.utils.data.dataloader")

        if not isinstance(self.diff_priv, bool):
            raise TypeError("diff_priv argument must be a bool")

    def get_parameters(self, config):
        """
        Returns the current parameters of the local model.

        Args:
            config: Configuration information.

        Returns:
            Numpy array: Parameters of the local model.
        """
        print(f"[Client {self.cid}] get_parameters")
        return self.local_model.get_parameters()

    def fit(self, parameters, config):
        """
        Fits the local model to the received parameters using federated learning.

        Args:
            parameters: Parameters received from the server.
            config: Configuration information.

        Returns:
            Tuple: Parameters of the local model, number of training examples, and privacy information.
        """
        print('\n -------------------------------- \n  this is the config of the client')
        print(f"[Client {self.cid}] fit, config: {config}")
        # print(config['epochs'])
        print('\n -------------------------------- \n  ')
        self.local_model.set_parameters(parameters)

        epoch_losses, epoch_metrics = [], []

        for _ in range(params["train_epochs"]):
            epsilon, tr_loss, tr_metric = self.local_model.train(
                self.trainloader,
                epoch=_,
                device=self.device,
                privacy_engine=self.privacy_engine,
                diff_priv=self.diff_priv,
            )
            self.epsilons.append(epsilon)
            epoch_losses.append(tr_loss)
            epoch_metrics.append(tr_metric)

        print(f"epsilon of client {self.cid} : eps = {epsilon}")
        round_train_loss = float(np.nanmean(epoch_losses)) if len(epoch_losses) else float("nan") 
         
        round_train_metric = float(np.nanmean(epoch_metrics)) if len(epoch_metrics) else float("nan")

        self.train_round_losses.append(round_train_loss)
        self.train_round_metrics.append(round_train_metric)

        metric_key = "train_accuracy" if self.local_model.task_type == "binary" else "train_rmse"

        return (
            self.local_model.get_parameters(),
            len(self.trainloader.dataset),   # better than len(dataloader)
            {
                "epsilon": float(epsilon),
                "train_loss": float(round_train_loss),
                metric_key: float(round_train_metric),
            },
        )


    def evaluate(self, parameters, config):
            print(f"[Client {self.cid}] evaluate, config: {config}")
            self.local_model.set_parameters(parameters)

            # NEW: evaluate returns (loss, metrics_dict)
            loss, metrics = self.local_model.evaluate(self.valloader, device=self.device)

            self.losses.append(loss)

            # Keep old behavior if accuracy exists
            if isinstance(metrics, dict) and "accuracy" in metrics:
                self.accuracies.append(metrics["accuracy"])

            # Return metrics as-is (works for binary + regression)
            # Flower expects: (loss, num_examples, metrics_dict)
            return float(loss), len(self.valloader.dataset), {k: float(v) for k, v in (metrics or {}).items()}
