# modules/server_buffer.py

import random
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 BufferedServer(ServerBase):
    def run(self):
        self.init_net()
        self.init_data()

        worker = Worker(self.args, self.logger)

        if self.args.init_gradient:
            global_para, _ = self.init_gradient(worker)
        else:
            global_para = deepcopy(self.net.state_dict())

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

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

            workers = self.get_current_worker()

            grad_list = []
            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)
                grad_list.append(deepcopy(gradient))
                del gradient

            # update global model
            self.logger.info(f"In update iter {self.update_iter}")
            for key, val in global_para.items():
                agg_grad = torch.zeros_like(val).to(self.device)
                for grad in grad_list:
                    agg_grad.add_(grad[key], alpha=1.0 / len(workers))
                val.sub_(agg_grad, alpha=self.args.lr)
            del grad_list

            # print("After update", {k: v.norm() for k, v in global_para.items()})

            # update local models
            for worker_id in workers:
                net_weights[worker_id] = deepcopy(global_para)

            # 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 get_current_worker(self):
        buffer_size = min(self.args.buffer_size, self.n_workers)
        if len(self.avail_workers) >= buffer_size:
            workers = random.sample(self.avail_workers, buffer_size)
            for w in workers:
                self.avail_workers.remove(w)
            return workers

        workers = self.avail_workers
        while len(workers) < buffer_size:
            for w, speed in enumerate(self.worker_speeds):
                if (self.comput_step + 1) % speed == 0 and w not in workers:
                    # worker w is available
                    if len(workers) < buffer_size:
                        workers.append(w)
                    else:  # # buffer is full
                        self.avail_workers.append(w)  # Add to the next available list
            self.comput_step += 1

        return workers