"""
Advanced Aggregation Strategies

Additional aggregation strategies for federated learning:
- FedProx: Proximal term for handling heterogeneity
- FedAvgM: Momentum-based aggregation
- TrimmedMean: Byzantine-resilient aggregation
- Krum: Byzantine-resilient aggregation
"""

from typing import List, Optional
import numpy as np
from copy import deepcopy

from quinkgl.aggregation.base import (
    AggregationStrategy,
    ModelUpdate,
    AggregatedModel
)
from quinkgl.aggregation.fedavg import FedAvg


class FedProx(FedAvg):
    """
    FedProx: Federated Learning with Proximal Term.

    Adds a proximal term to the aggregation to handle heterogeneity
    and prevent client drift from the global model.

    Reference: https://arxiv.org/abs/1812.06127
    """

    def __init__(self, mu: float = 0.01, **kwargs):
        """
        Initialize FedProx aggregator.

        Args:
            mu: Proximal term coefficient (higher = stricter adherence to global model)
            **kwargs: Additional arguments passed to FedAvg
        """
        super().__init__(**kwargs)
        self.mu = mu
        self.global_weights = None  # Will be set during first aggregation

    async def aggregate(
        self,
        updates: List[ModelUpdate]
    ) -> AggregatedModel:
        """
        Aggregate with proximal term regularization.

        The proximal term penalizes deviations from the global model,
        helping to handle heterogeneous data distributions.
        """
        self._validate_updates(updates)

        # Store global model weights for proximal term
        if self.global_weights is None:
            # First round: use simple FedAvg
            result = await super().aggregate(updates)
            self.global_weights = deepcopy(result.weights)
            return result

        # Apply proximal term to updates before aggregation
        proximal_updates = self._apply_proximal_term(updates)

        # Aggregate using FedAvg
        result = await super().aggregate(proximal_updates)

        # Update global weights
        self.global_weights = deepcopy(result.weights)
        result.metadata["aggregation_method"] = "fedprox"
        result.metadata["mu"] = self.mu

        return result

    def _apply_proximal_term(self, updates: List[ModelUpdate]) -> List[ModelUpdate]:
        """Apply proximal term correction to updates."""
        if self.global_weights is None:
            return updates

        corrected_updates = []
        for update in updates:
            new_update = deepcopy(update)
            new_update.weights = self._proximal_correction(
                update.weights,
                self.global_weights
            )
            corrected_updates.append(new_update)

        return corrected_updates

    def _proximal_correction(self, weights, global_weights):
        """Apply proximal correction to weights."""
        if isinstance(weights, np.ndarray):
            # w - mu * (w - w_global) = (1-mu)*w + mu*w_global
            return (1 - self.mu) * weights + self.mu * global_weights
        elif isinstance(weights, dict):
            result = {}
            for key in weights:
                if key in global_weights and hasattr(weights[key], '__array__'):
                    result[key] = (
                        (1 - self.mu) * weights[key] + self.mu * global_weights[key]
                    ).astype(weights[key].dtype)
                else:
                    result[key] = weights[key]
            return result
        else:
            return weights


class FedAvgM(AggregationStrategy):
    """
    FedAvgM: Federated Averaging with Momentum.

    Uses server momentum to stabilize training and improve convergence.

    Reference: https://arxiv.org/abs/1909.03083
    """

    def __init__(self, server_momentum: float = 0.9, **kwargs):
        """
        Initialize FedAvgM aggregator.

        Args:
            server_momentum: Momentum coefficient (0-1, default 0.9)
            **kwargs: Additional configuration parameters
        """
        super().__init__(**kwargs)
        self.server_momentum = server_momentum
        self.momentum_buffer = None

    async def aggregate(
        self,
        updates: List[ModelUpdate]
    ) -> AggregatedModel:
        """
        Aggregate with server momentum.

        Momentum buffer is updated as:
            buffer = momentum * buffer + averaged_update
        """
        self._validate_updates(updates)

        # First, compute simple average of updates
        weights_list = [self.compute_weight(u) for u in updates]
        total_weight = sum(weights_list)

        if total_weight == 0:
            raise ValueError("Total weight is zero, cannot aggregate")

        first_weights = updates[0].weights

        if isinstance(first_weights, np.ndarray):
            averaged = self._average_numpy(updates, weights_list, total_weight)
        elif isinstance(first_weights, dict):
            averaged = self._average_dict(updates, weights_list, total_weight)
        else:
            # Fallback to FedAvg
            fedavg = FedAvg()
            return await fedavg.aggregate(updates)

        # Apply momentum
        if self.momentum_buffer is None:
            # First round: no momentum
            self.momentum_buffer = deepcopy(averaged)
        else:
            # Apply momentum: buffer = momentum * buffer + (1-momentum) * averaged
            self.momentum_buffer = self._apply_momentum(
                self.momentum_buffer, averaged
            )

        return AggregatedModel(
            weights=deepcopy(self.momentum_buffer),
            contributing_peers=[u.peer_id for u in updates],
            total_samples=sum(u.sample_count for u in updates),
            metadata={"aggregation_method": "fedavgm", "momentum": self.server_momentum},
            updates=updates
        )

    def _apply_momentum(self, buffer, new_weights):
        """Apply momentum to the update."""
        if isinstance(buffer, np.ndarray):
            return (
                self.server_momentum * buffer + (1 - self.server_momentum) * new_weights
            ).astype(buffer.dtype)
        elif isinstance(buffer, dict):
            result = {}
            for key in buffer:
                if key in new_weights and hasattr(buffer[key], '__array__'):
                    result[key] = (
                        self.server_momentum * buffer[key] +
                        (1 - self.server_momentum) * new_weights[key]
                    ).astype(buffer[key].dtype)
                else:
                    result[key] = buffer[key]
            return result
        else:
            return buffer

    def _average_numpy(self, updates: List[ModelUpdate], weights_list: List[float], total_weight: float):
        """Average numpy array weights."""
        result = np.zeros_like(updates[0].weights, dtype=np.float64)
        for update, weight in zip(updates, weights_list):
            normalized_weight = weight / total_weight
            result += update.weights.astype(np.float64) * normalized_weight
        return result.astype(updates[0].weights.dtype)

    def _average_dict(self, updates: List[ModelUpdate], weights_list: List[float], total_weight: float):
        """Average dict weights."""
        result = {}
        all_keys = set()
        for update in updates:
            if isinstance(update.weights, dict):
                all_keys.update(update.weights.keys())

        for key in all_keys:
            values = []
            corresponding_weights = []
            for update, weight in zip(updates, weights_list):
                if isinstance(update.weights, dict) and key in update.weights:
                    values.append(update.weights[key])
                    corresponding_weights.append(weight)

            if values and hasattr(values[0], '__array__'):
                key_total = sum(corresponding_weights)
                averaged = np.zeros_like(values[0], dtype=np.float64)
                for val, w in zip(values, corresponding_weights):
                    averaged += val.astype(np.float64) * (w / key_total)
                result[key] = averaged.astype(values[0].dtype)
            elif values:
                result[key] = values[0]

        return result


class TrimmedMean(AggregationStrategy):
    """
    Trimmed Mean: Byzantine-resilient aggregation.

    Removes the smallest and largest values before averaging,
    providing robustness against malicious/faulty peers.

    Reference: https://arxiv.org/abs/1803.09877
    """

    def __init__(self, trim_ratio: float = 0.1, **kwargs):
        """
        Initialize TrimmedMean aggregator.

        Args:
            trim_ratio: Fraction of smallest/largest values to trim (0-0.5)
            **kwargs: Additional configuration parameters
        """
        super().__init__(**kwargs)
        if not 0 <= trim_ratio < 0.5:
            raise ValueError("trim_ratio must be in [0, 0.5)")
        self.trim_ratio = trim_ratio

    async def aggregate(
        self,
        updates: List[ModelUpdate]
    ) -> AggregatedModel:
        """
        Aggregate using trimmed mean.
        """
        self._validate_updates(updates)

        if len(updates) < 3:
            # Not enough updates to trim, fall back to simple average
            fedavg = FedAvg()
            result = await fedavg.aggregate(updates)
            result.metadata["aggregation_method"] = "trimmed_mean_fallback"
            return result

        first_weights = updates[0].weights

        if isinstance(first_weights, np.ndarray):
            aggregated = self._trim_numpy(updates)
        elif isinstance(first_weights, dict):
            aggregated = self._trim_dict(updates)
        else:
            # Fallback
            fedavg = FedAvg()
            return await fedavg.aggregate(updates)

        return AggregatedModel(
            weights=aggregated,
            contributing_peers=[u.peer_id for u in updates],
            total_samples=sum(u.sample_count for u in updates),
            metadata={"aggregation_method": "trimmed_mean", "trim_ratio": self.trim_ratio},
            updates=updates
        )

    def _trim_numpy(self, updates: List[ModelUpdate]) -> np.ndarray:
        """Apply trimmed mean to numpy arrays."""
        weights_matrix = np.stack([u.weights.flatten() for u in updates])
        n = len(updates)
        k = int(n * self.trim_ratio)

        if k == 0:
            # No trimming, use mean
            return np.mean(weights_matrix, axis=0).astype(updates[0].weights.dtype)

        # For each parameter, sort and trim
        trimmed_result = np.zeros_like(weights_matrix[0])
        for i in range(weights_matrix.shape[1]):
            values = weights_matrix[:, i]
            # Sort and trim k smallest and k largest
            sorted_values = np.sort(values)
            trimmed_values = sorted_values[k:n-k]
            trimmed_result[i] = np.mean(trimmed_values)

        return trimmed_result.reshape(updates[0].weights.shape).astype(updates[0].weights.dtype)

    def _trim_dict(self, updates: List[ModelUpdate]) -> dict:
        """Apply trimmed mean to dict weights."""
        result = {}
        all_keys = set()
        for update in updates:
            if isinstance(update.weights, dict):
                all_keys.update(update.weights.keys())

        for key in all_keys:
            values = []
            for update in updates:
                if isinstance(update.weights, dict) and key in update.weights:
                    val = update.weights[key]
                    if hasattr(val, '__array__'):
                        values.append(val.flatten())

            if values:
                weights_matrix = np.stack(values)
                n = len(values)
                k = int(n * self.trim_ratio)

                if k == 0:
                    trimmed = np.mean(weights_matrix, axis=0)
                else:
                    trimmed_result = np.zeros_like(weights_matrix[0])
                    for i in range(weights_matrix.shape[1]):
                        sorted_values = np.sort(weights_matrix[:, i])
                        trimmed_values = sorted_values[k:n-k]
                        trimmed_result[i] = np.mean(trimmed_values)
                    trimmed = trimmed_result

                result[key] = trimmed.reshape(values[0].shape).astype(values[0].dtype)
            else:
                # Use first available value
                for update in updates:
                    if isinstance(update.weights, dict) and key in update.weights:
                        result[key] = update.weights[key]
                        break

        return result


class Krum(AggregationStrategy):
    """
    Krum: Byzantine-resilient aggregation.

    Selects the update closest to the majority of updates,
    providing robustness against malicious peers.

    Reference: https://arxiv.org/abs/1703.02857
    """

    def __init__(self, num_byzantines: int = 1, **kwargs):
        """
        Initialize Krum aggregator.

        Args:
            num_byzantines: Expected number of malicious peers (f)
            **kwargs: Additional configuration parameters
        """
        super().__init__(**kwargs)
        self.num_byzantines = num_byzantines

    async def aggregate(
        self,
        updates: List[ModelUpdate]
    ) -> AggregatedModel:
        """
        Aggregate using Krum - select the most central update.
        """
        self._validate_updates(updates)

        n = len(updates)
        if n <= 2 * self.num_byzantines:
            raise ValueError(
                f"Krum requires n > 2*f updates (n={n}, f={self.num_byzantines})"
            )

        # Compute distances between all pairs of updates
        distances = self._compute_distances(updates)

        # Compute scores (sum of smallest n-2*f-1 distances for each update)
        n_closest = n - 2 * self.num_byzantines - 1
        scores = []
        for i in range(n):
            # Sort distances from update i to all others
            sorted_dist = np.sort(distances[i])
            # Sum of n-2*f-1 smallest distances
            scores.append(np.sum(sorted_dist[:n_closest]))

        # Select update with minimum score
        selected_idx = int(np.argmin(scores))
        selected_update = updates[selected_idx]

        return AggregatedModel(
            weights=deepcopy(selected_update.weights),
            contributing_peers=[selected_update.peer_id],
            total_samples=selected_update.sample_count,
            metadata={
                "aggregation_method": "krum",
                "num_byzantines": self.num_byzantines,
                "selected_peer": selected_update.peer_id
            },
            updates=[selected_update]
        )

    def _compute_distances(self, updates: List[ModelUpdate]) -> np.ndarray:
        """Compute pairwise Euclidean distances between updates."""
        n = len(updates)
        distances = np.zeros((n, n))

        # Flatten all weights to vectors
        weight_vectors = []
        for update in updates:
            if isinstance(update.weights, np.ndarray):
                weight_vectors.append(update.weights.flatten())
            elif isinstance(update.weights, dict):
                # Concatenate all array values
                parts = []
                for key in sorted(update.weights.keys()):
                    val = update.weights[key]
                    if hasattr(val, '__array__'):
                        parts.append(val.flatten())
                weight_vectors.append(np.concatenate(parts))
            else:
                weight_vectors.append(np.array(update.weights).flatten())

        weight_vectors = [w.astype(np.float64) for w in weight_vectors]

        # Compute pairwise distances
        for i in range(n):
            for j in range(i + 1, n):
                dist = np.linalg.norm(weight_vectors[i] - weight_vectors[j])
                distances[i, j] = dist
                distances[j, i] = dist

        return distances


class MultiKrum(Krum):
    """
    Multi-Krum: Krum with averaging.

    Instead of selecting a single update, selects the n-2*f closest updates
    and averages them for better stability.
    """

    async def aggregate(
        self,
        updates: List[ModelUpdate]
    ) -> AggregatedModel:
        """
        Aggregate using Multi-Krum - average the n-2*f most central updates.
        """
        self._validate_updates(updates)

        n = len(updates)
        if n <= 2 * self.num_byzantines:
            raise ValueError(
                f"MultiKrum requires n > 2*f updates (n={n}, f={self.num_byzantines})"
            )

        # Compute distances
        distances = self._compute_distances(updates)

        # Compute scores
        n_closest = n - 2 * self.num_byzantines - 1
        scores = []
        for i in range(n):
            sorted_dist = np.sort(distances[i])
            scores.append(np.sum(sorted_dist[:n_closest]))

        # Select n-2*f updates with minimum scores
        num_selected = n - 2 * self.num_byzantines
        selected_indices = np.argsort(scores)[:num_selected].tolist()
        selected_updates = [updates[i] for i in selected_indices]

        # Average the selected updates
        fedavg = FedAvg()
        result = await fedavg.aggregate(selected_updates)

        result.metadata["aggregation_method"] = "multikrum"
        result.metadata["num_byzantines"] = self.num_byzantines
        result.metadata["selected_peers"] = [u.peer_id for u in selected_updates]

        return result
