# modules/server_async_ccc.py

import random
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 ConServer(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)]

        self.init_active_workers()

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

            worker_id = self.get_current_worker()

            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)

            # update global model
            self.logger.info(f"In update iter {self.update_iter}")
            for key, grad in gradient.items():
                global_para[key].sub_(grad, alpha=self.args.lr)
            del gradient

            # update active workers
            self.update_active_workers(worker_id)

            # update local models
            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):
        self.logger.debug(f"Current active workers: {self.active_workers}")

        cur_worker = None
        while cur_worker is None:
            for w in self.active_workers:
                speed = self.worker_speeds[w]

                if (self.comput_step + 1) % speed == 0:
                    cur_worker = w
                    break

            self.comput_step += 1

        # print(f"At step {self.comput_step}: Worker {cur_worker} is selected with speed {self.worker_speeds[cur_worker]}")
        return cur_worker

    def init_active_workers(self):
        buffer_size = min(self.args.buffer_size, self.n_workers)
        self.active_workers = random.sample(range(self.n_workers), buffer_size)
        self.logger.debug(f"Initial active workers: {self.active_workers}")

    def update_active_workers(self, worker_id):
        # already updated worker
        self.active_workers.remove(worker_id)

        # sample a new worker
        sampled_worker = random.choice(range(self.n_workers))

        if sampled_worker not in self.active_workers:
            self.active_workers.append(sampled_worker)