"""Federated Learning Server for Myanmar Ghost project.

Coordinates model training across distributed clients
using Flower framework.
"""

import flwr as fl
import numpy as np
import torch
import torch.nn as nn
import yaml
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

logger = __import__("loguru").logger


class FederatedServer:
    """Flower server for federated learning."""
    
    def __init__(
        self,
        model: nn.Module,
        strategy: Optional[fl.server.strategy.Strategy] = None,
        output_dir: str = "outputs/federated",
    ):
        self.model = model
        self.strategy = strategy or self._default_strategy()
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        self.server = None
        self.history = None
    
    def _default_strategy(self) -> fl.server.strategy.Strategy:
        """Create default federated averaging strategy."""
        return fl.server.strategy.FedAvg(
            fraction_fit=0.5,
            fraction_evaluate=0.5,
            min_fit_clients=2,
            min_evaluate_clients=2,
            min_available_clients=2,
        )
    
    def get_model_parameters(self) -> List[np.ndarray]:
        """Get current model parameters."""
        return [val.cpu().numpy() for _, val in self.model.state_dict().items()]
    
    def set_model_parameters(self, parameters: List[np.ndarray]) -> None:
        """Set model parameters from numpy arrays."""
        state_dict = dict(self.model.state_dict())
        for i, (key, _) in enumerate(state_dict.items()):
            state_dict[key] = torch.from_numpy(parameters[i])
        self.model.load_state_dict(state_dict)
    
    def aggregate_results(
        self,
        results: List[Tuple[List[np.ndarray], int, Dict]],
    ) -> Tuple[List[np.ndarray], Dict]:
        """Aggregate client results using weighted averaging.
        
        Args:
            results: List of (parameters, num_samples, metrics)
        
        Returns:
            Aggregated parameters, aggregated metrics
        """
        total_samples = sum(r[1] for r in results)
        
        # Weighted average of parameters
        weighted_params = None
        
        for params, n_samples, _ in results:
            weight = n_samples / total_samples
            if weighted_params is None:
                weighted_params = [p * weight for p in params]
            else:
                weighted_params = [
                    wp + p * weight
                    for wp, p in zip(weighted_params, params)
                ]
        
        # Aggregate metrics
        aggregated_metrics = {}
        for _, _, metrics in results:
            for key, value in metrics.items():
                if key not in aggregated_metrics:
                    aggregated_metrics[key] = []
                aggregated_metrics[key].append(value)
        
        # Average metrics
        avg_metrics = {
            key: np.mean(values)
            for key, values in aggregated_metrics.items()
        }
        
        logger.info(
            f"Aggregated {len(results)} client results "
            f"(total samples: {total_samples})"
        )
        
        return weighted_params, avg_metrics
    
    def save_global_model(self, path: Optional[str] = None) -> str:
        """Save the global model."""
        if path is None:
            from datetime import datetime
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            path = self.output_dir / f"global_model_{timestamp}.pt"
        
        torch.save({
            "model_state_dict": self.model.state_dict(),
        }, path)
        
        logger.info(f"Global model saved to {path}")
        return str(path)
    
    def load_global_model(self, path: str) -> None:
        """Load a global model checkpoint."""
        checkpoint = torch.load(path)
        self.model.load_state_dict(checkpoint["model_state_dict"])
        logger.info(f"Global model loaded from {path}")
    
    def start(
        self,
        server_address: str = "[::]:8080",
        num_rounds: int = 5,
    ) -> fl.server.Server:
        """Start the federated learning server.
        
        Args:
            server_address: Address to bind the server
            num_rounds: Number of federated rounds
        
        Returns:
            Server instance
        """
        from flwr.server import ServerConfig
        
        config = ServerConfig(num_rounds=num_rounds)
        
        def on_fit_config_fn(rnd: int) -> Dict[str, Any]:
            """Generate config for each round."""
            return {
                "local_epochs": 1,
                "batch_size": 32,
                "learning_rate": 1e-4,
                "round": rnd,
            }
        
        def on_evaluate_config_fn(rnd: int) -> Dict[str, Any]:
            """Generate config for evaluation."""
            return {
                "batch_size": 32,
                "round": rnd,
            }
        
        # Create strategy with callbacks
        strategy = fl.server.strategy.FedAvg(
            fraction_fit=0.5,
            fraction_evaluate=0.5,
            min_fit_clients=2,
            min_evaluate_clients=2,
            min_available_clients=2,
            on_fit_config_fn=on_fit_config_fn,
            on_evaluate_config_fn=on_evaluate_config_fn,
        )
        
        self.server = fl.server.start_server(
            server_address=server_address,
            server=config,
            strategy=strategy,
            model=self.model,
        )
        
        return self.server
    
    def get_history(self) -> Optional[List[Dict]]:
        """Get training history from server."""
        if self.server and hasattr(self.server, "history"):
            return self.server.history
        return None


def load_server_config(config_path: str) -> Dict:
    """Load server configuration from YAML."""
    with open(config_path, "r", encoding="utf-8") as f:
        return yaml.safe_load(f)


def create_server(
    model: nn.Module,
    strategy: Optional[str] = "fedavg",
    output_dir: str = "outputs/federated",
) -> FederatedServer:
    """Factory function to create federated server."""
    return FederatedServer(
        model=model,
        output_dir=output_dir,
    )


if __name__ == "__main__":
    print("FederatedServer module loaded")
    print("Use create_server() to create a server and start() to begin")