# modules/server_mifa.py

import numpy as np
import torch
from copy import deepcopy

from src.modules.server_base import ServerBase
from src.modules.worker import Worker
from src.other_utils.utils import get_dataloader
from src.other_utils.synthetic import get_synthetic_dataloader


class MIFAServer(ServerBase):
    def run(self):
        self.init_net()
        self.init_data()

        worker = Worker(self.args, self.logger)
        agg_weight = 1.0 / self.n_workers
        if self.args.init_gradient:
            global_para, last_grad_table = self.init_gradient(worker)
            agg_grad = {key: torch.zeros_like(val).to(self.device) for key, val in self.net.state_dict().items()}
            for key, grad in agg_grad.items():
                for cum_grad in last_grad_table:
                    grad.add_(cum_grad[key], alpha=agg_weight)
        else:
            global_para = deepcopy(self.net.state_dict())
            last_grad_table = [None for _ in range(self.n_workers)]
            agg_grad = {key: torch.zeros_like(val).to(self.device) for key, val in self.net.state_dict().items()}

        net_weights = [deepcopy(global_para) for _ in range(self.n_workers)]

        self.init_worker_set()

        while self.comput_step <= self.args.tot_step:
            self.logger.info(f"In computation step {self.comput_step}")

            workers = self.get_current_worker()
            for id in workers:  # update models
                net_weights[id] = deepcopy(global_para)

            for worker_id in workers:
                self.net.load_state_dict(net_weights[worker_id])

                # if self.args.dataset == "synthetic":
                #     train_dl = get_synthetic_dataloader(self.args, return_idx=worker_id)
                # else:
                #     train_dl, _ = get_dataloader(self.args.dataset, self.args.datadir, self.args.batch_size, dataidxs=self.net_dataidx_map[worker_id])
                train_dl = self.train_dl_dict[worker_id]
                gradient = worker.train_net(worker_id, self.net, train_dl)
                last_grad_table[worker_id] = deepcopy(gradient)
                del gradient

            # compute aggregated gradient
            self.logger.info(f"In update iter {self.update_iter}")
            for key in agg_grad.keys():  # g^{t-1} + 1/n \nabla f_i(w^{t-1}) current gradient
                agg_grad[key] = torch.zeros_like(agg_grad[key]).to(self.device)
                for id, grad in enumerate(last_grad_table):
                    agg_grad[key].add_(grad[key], alpha=agg_weight)

            # update global model
            for key, grad in agg_grad.items():
                global_para[key].sub_(grad, alpha=self.args.lr)

            # test global model
            if (self.update_iter + 1) % self.test_freq == 0:
                self.test_net(global_para)
                if self.args.save_net:
                    self.save_net(global_para)

            self.update_iter += 1

        if self.args.save_net:
            self.save_net(global_para)

        # --- Added: save all logged data after training finishes ---
        self.save_results_to_csv()
        # --- End of added block ---

    def init_worker_set(self):
        self.n_sets = int(self.n_workers // self.buffer_size)
        self.workers_set = np.array_split(np.arange(self.n_workers), self.n_sets)

    def get_current_worker(self):
        x = self.update_iter % self.n_sets
        cur_workers = self.workers_set[x]

        max_speed = 0
        for w, speed in enumerate(self.worker_speeds):
            if w in cur_workers:
                max_speed = max(max_speed, speed)
        self.comput_step += max_speed

        return cur_workers