#!/usr/bin/env python3

import copy
from typing import Dict, Optional, Tuple

import flwr as fl
import torch

from torch.utils.data import ConcatDataset, DataLoader  # ✅ NEW

from .client import FlowerClient
from .federated_dataset import FederatedDataset
from .model import Model
from .strategy import Strategy


class FlowerServer:
    """
    A class representing the central server for Federated Learning using Flower.
    """

    def __init__(
        self,
        global_model: Model,
        strategy: Strategy,
        num_rounds: int,
        num_clients: int,
        fed_dataset: FederatedDataset,
        diff_privacy: bool = False,
        client_resources: Optional[Dict[str, float]] = {'num_cpus': 1, 'num_gpus': 0.0}
    ) -> None:

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.global_model = global_model
        self.params = global_model.get_parameters()
        self.global_model.model = global_model.model.to(self.device)
        self.num_rounds = num_rounds
        self.num_clients = num_clients
        self.fed_dataset = fed_dataset
        self.strategy = strategy
        self.client_resources = client_resources

        setattr(self.strategy.strategy_object, "min_available_clients", self.num_clients)
        setattr(self.strategy.strategy_object, "initial_parameters", fl.common.ndarrays_to_parameters(self.params))
        setattr(self.strategy.strategy_object, "evaluate_fn", self.evaluate)

        self.diff_priv = diff_privacy

        # ✅ tracking
        self.accuracies = []
        self.losses = []
        self.auc = []
        self.rmse = []
        self.mae = []
        self.r2 = []
        self.adj_r2 = []

        self.train_losses = []
        self.train_accuracies = []
        self.train_auc = []
        self.train_rmse = []
        self.train_mae = []
        self.train_r2 = []
        self.train_adj_r2 = []

        self.flower_clients = []
        self.device = torch.device("cpu")
        self.global_model = global_model

        # move model to CPU
        self.global_model.model = self.global_model.model.to(self.device)

        # also move optimizer state to CPU (important if optimizer had any tensors)
        for state in self.global_model.optimizer.state.values():
            for k, v in state.items():
                if torch.is_tensor(v):
                    state[k] = v.detach().cpu()

        # if criterion has tensor buffers (e.g., pos_weight), ensure CPU too
        if hasattr(self.global_model.criterion, "pos_weight") and self.global_model.criterion.pos_weight is not None:
            self.global_model.criterion.pos_weight = self.global_model.criterion.pos_weight.detach().cpu()
        self.validate()

    def validate(self) -> None:
        if not isinstance(self.global_model, Model):
            raise TypeError("global_model argument must be a Model instance")
        if not isinstance(self.num_clients, int):
            raise TypeError("num_clients argument must be an int")
        if not isinstance(self.num_rounds, int):
            raise TypeError("num_rounds argument must be an int")
        if not isinstance(self.diff_priv, bool):
            raise TypeError("diff_priv argument must be a bool")

    def client_fn(self, cid) -> FlowerClient:
        # NOTE: you might want to replace %4 with device_count() later
        device = torch.device(f"cuda:{int(cid) % 4}" if torch.cuda.is_available() else "cpu")

        client_model = copy.deepcopy(self.global_model)

        trainloader = self.fed_dataset.trainloaders[int(cid)]
        valloader = self.fed_dataset.valloaders[int(cid)]

        client_model = copy.deepcopy(self.global_model)
        client = FlowerClient(cid, client_model, trainloader, valloader, self.diff_priv)
        self.flower_clients.append(client)
        return client

    # ✅ NEW helper: build one global loader from all client loaders (simulation)
    def _concat_loader(self, loaders, shuffle: bool = False) -> DataLoader:
        # Use each loader's dataset and build one union dataset
        ds = ConcatDataset([ldr.dataset for ldr in loaders])

        # Try to reuse an existing batch_size if possible; fallback to 64
        batch_size = getattr(loaders[0], "batch_size", None) or 64
        num_workers = getattr(loaders[0], "num_workers", 0)
        pin_memory = getattr(loaders[0], "pin_memory", False)

        return DataLoader(
            ds,
            batch_size=batch_size,
            shuffle=shuffle,
            num_workers=num_workers,
            pin_memory=pin_memory,
        )

    def evaluate(
        self,
        server_round: int,
        parameters: fl.common.NDArrays,
        config: Dict[str, fl.common.Scalar],
    ) -> Optional[Tuple[float, Dict[str, fl.common.Scalar]]]:

        # Update model with latest weights
        self.global_model.set_parameters(parameters)
        self.last_parameters = parameters


        # ✅ GLOBAL loaders across ALL clients (simulation)
        global_valloader = self._concat_loader(self.fed_dataset.valloaders, shuffle=False)
        global_trainloader = self._concat_loader(self.fed_dataset.trainloaders, shuffle=False)

        # ---- Validation (global) ----
        val_loss, val_metrics = self.global_model.evaluate(global_valloader, self.device)
        self.losses.append(val_loss)

        if isinstance(val_metrics, dict):
            if "accuracy" in val_metrics: self.accuracies.append(val_metrics["accuracy"])
            if "auc" in val_metrics: self.auc.append(val_metrics["auc"])
            if "rmse" in val_metrics: self.rmse.append(val_metrics["rmse"])
            if "mae" in val_metrics: self.mae.append(val_metrics["mae"])
            if "r2" in val_metrics: self.r2.append(val_metrics["r2"])
            if "adj_r2" in val_metrics: self.adj_r2.append(val_metrics["adj_r2"])

        # ---- Train (global) ----
        train_loss, train_metrics = self.global_model.evaluate(global_trainloader, self.device)
        self.train_losses.append(train_loss)

        if isinstance(train_metrics, dict):
            if "accuracy" in train_metrics: self.train_accuracies.append(train_metrics["accuracy"])
            if "auc" in train_metrics: self.train_auc.append(train_metrics["auc"])
            if "rmse" in train_metrics: self.train_rmse.append(train_metrics["rmse"])
            if "mae" in train_metrics: self.train_mae.append(train_metrics["mae"])
            if "r2" in train_metrics: self.train_r2.append(train_metrics["r2"])
            if "adj_r2" in train_metrics: self.train_adj_r2.append(train_metrics["adj_r2"])

        # Return validation loss + a metrics dict that also includes train_*
        out_metrics = dict(val_metrics or {})
        out_metrics["train_loss"] = float(train_loss) if train_loss is not None else float("nan")

        if isinstance(train_metrics, dict):
            for k, v in train_metrics.items():
                out_metrics[f"train_{k}"] = float(v) if v is not None else float("nan")

        return val_loss, out_metrics

    def score_each_client(self, parameters=None, split: str = "val"):
        """
        Evaluate the (final) global model on each client's data.
        split: "val" or "train"
        """
        if parameters is not None:
            self.global_model.set_parameters(parameters)

        loaders = self.fed_dataset.valloaders if split == "val" else self.fed_dataset.trainloaders

        results = {}
        for cid, loader in enumerate(loaders):
            loss, metrics = self.global_model.evaluate(loader, self.device)
            results[str(cid)] = {
                "num_examples": len(loader.dataset),
                "loss": float(loss),
                **{k: float(v) for k, v in (metrics or {}).items()},
            }
        return results



    def run(self) -> None:
        ray_init_args = {"include_dashboard": False, "object_store_memory": 78643200 , 
                        }
        self.fed_dataset.eng = None

        history = fl.simulation.start_simulation(
            client_fn=self.client_fn,
            num_clients=self.num_clients,
            config=fl.server.ServerConfig(self.num_rounds),
            strategy=self.strategy.strategy_object,
            ray_init_args=ray_init_args,
            client_resources=self.client_resources
        )
        return history
