# modules/central.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 CentralServer(ServerBase):
    def run(self):
        if self.args.n_workers != 1:
            self.args.n_workers = 1

        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())

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

            self.net.load_state_dict(global_para)

            if self.args.dataset == "synthetic":
                train_dl = get_synthetic_dataloader(self.args, return_idx=0)
            else:
                train_dl, _ = get_dataloader(self.args.dataset, self.args.datadir, self.args.batch_size,
                                             dataidxs=self.net_dataidx_map[0])

            gradient = worker.train_net(0, self.net, train_dl)

            # update global model
            for key, grad in gradient.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 ---