# modules/server_base.py

import torch
# import wandb  # Remove or comment out the wandb import if unused
import os
from copy import deepcopy
import pandas as pd  # Added for exporting metrics to CSV files

from src.models.model import *
from src.models.vggmodel import *
from src.other_utils.utils import compute_accuracy, partition_data, get_dataloader
from src.other_utils.delay import set_worker_speed
from src.other_utils.synthetic import get_synthetic_dataloader


class ServerBase:
    def __init__(self, args, logger):
        self.logger = logger
        self.device = torch.device(args.device)
        self.args = args
        self.tot_step = args.tot_step
        self.test_freq = args.test_freq
        self.buffer_size = args.buffer_size

        self.n_workers = args.n_workers
        self.worker_speeds = set_worker_speed(args, logger)

        # init
        self.update_iter = 0  # current
        self.comput_step = 0  # current
        self.avail_workers = []

        self.net_optimal = None

        # --- Added: containers for tracking metrics during training ---
        self.train_loss_history = []
        self.test_acc_history = []
        self.comput_step_history = []  # Records the corresponding computation steps
        self.update_iter_history = []  # Records the corresponding update iterations

        # Define the CSV file paths
        # Ensure that args.logdir exists and that the filename distinguishes different experiments
        self.log_file_path_acc = os.path.join(self.args.logdir, f"{self.args.log_file_name}-test-acc.csv")
        self.log_file_path_loss = os.path.join(self.args.logdir, f"{self.args.log_file_name}-train-loss.csv")

        # Initialize the CSV files and write the header rows
        # The 'w' mode overwrites any existing file for a fresh run.
        # To append to existing files instead, adjust the naming logic in run.py.
        pd.DataFrame(columns=['Step/Comput', 'Update/Iter', 'Test/Acc']).to_csv(self.log_file_path_acc, index=False)
        pd.DataFrame(columns=['Step/Comput', 'Update/Iter', 'Train/Loss']).to_csv(self.log_file_path_loss, index=False)
        # --- End of added block ---

    def init_net(self):
        self.logger.info("Initializing net")
        model = self.args.model
        dataset = self.args.dataset

        if model == "vgg11":  # for cifar
            net = vgg11()
        elif model == "cnn":
            if dataset in ("fmnist", "mnist"):
                net = SimpleCNNMNIST(input_dim=(16 * 4 * 4), hidden_dims=[120, 84], output_dim=10)
            elif dataset == "cifar10":
                net = SimpleCNN(input_dim=(16 * 5 * 5), hidden_dims=[120, 84], output_dim=10)
        elif model == "cnn4":
            net = CNN4Conv()
        elif model == "linear":
            if dataset == "synthetic":
                net = LinearModel(input_dim=self.args.in_dim)
            elif dataset in ("fmnist", "mnist"):
                net = LinearModel(input_dim=28 * 28, output_dim=10)
        else:
            print(f"Not supported model {model} yet")
            exit(1)

        if self.args.resume:
            net = self.resume(net)

        net.to(self.device)
        self.net = net
        return

    def init_data(self):
        self.logger.info("Initializing data (%s)" % self.args.partition)
        if self.args.dataset == "synthetic":
            # Important: make sure self.train_dl_dict is initialized here
            self.train_dl_dict = {}
            for i in range(self.n_workers):
                # Build the synthetic-data loader for each worker
                self.train_dl_dict[i] = get_synthetic_dataloader(self.args, return_idx=i)
            # Build the global test-data loader
            _, self.test_dl = get_synthetic_dataloader(self.args, return_idx=-1)

            self.train_dl_global, self.test_dl_global = get_synthetic_dataloader(self.args, return_idx=-1)
            return

        net_dataidx_map = partition_data(self.args.dataset, self.args.datadir, self.args.logdir, self.args.partition,
                                         self.args.n_workers, self.args.n_sample, alpha=self.args.alpha)

        train_dl_global, test_dl_global, train_dl_dict = get_dataloader(self.args.dataset, self.args.datadir,
                                                                        self.args.batch_size, dataidxs=-1,
                                                                        net_dataidx_map=net_dataidx_map)

        self.net_dataidx_map = net_dataidx_map
        self.train_dl_global = train_dl_global
        self.test_dl_global = test_dl_global
        self.train_dl_dict = train_dl_dict
        return

    def init_gradient(self, worker):
        self.logger.info("Initialize the gradient table")
        global_para = deepcopy(self.net.state_dict())
        last_grad_table = []
        for worker_id in range(self.n_workers):
            # 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.append(deepcopy(gradient))

            for key, grad in gradient.items():
                global_para[key].sub_(grad, alpha=self.args.lr * 1.0 / self.n_workers)
            del gradient

        return global_para, last_grad_table

    def get_current_worker(self):
        pass

    def run(self):
        # Let workers train, the server collects the results according to different rules
        # key function, please refer to implementations
        pass

    def test_net(self, net_weight=None):
        if net_weight:
            self.net.load_state_dict(net_weight)

        train_acc, train_loss = compute_accuracy(self.net, self.train_dl_global, device=self.device,
                                                 dataset=self.args.dataset)
        test_acc, test_loss = compute_accuracy(self.net, self.test_dl_global, device=self.device,
                                               dataset=self.args.dataset)

        dist = 0.0
        if self.args.model == "linear":
            if self.net_optimal is not None:
                for k in self.net_optimal.keys():
                    if self.net_optimal[k].device != net_weight[k].device:
                        self.net_optimal[k] = self.net_optimal[k].to(net_weight[k].device)
                    dist += torch.norm(self.net_optimal[k] - net_weight[k])
                self.logger.info('>> Distance to optimal model: %f' % dist)
            else:
                try:
                    self.net_optimal = torch.load(
                        os.path.join(self.args.modeldir, f"300-mnist-{self.args.lamb}-model.pth"))
                    print("Loading optimal model")
                    dist = 0.0
                    for k, v in self.net_optimal.items():
                        v = v.to(net_weight[k].device)
                        dist += torch.norm(v - net_weight[k])
                    self.logger.info('>> Distance to optimal model: %f' % dist)
                except:
                    pass

        self.logger.info('>> Global Model Train accuracy: %f' % train_acc)
        self.logger.info('>> Global Model Train loss: %f' % train_loss)
        self.logger.info('>> Global Model Test accuracy: %f' % test_acc)
        self.logger.info('>> Global Model Test loss: %f' % test_loss)

        # --- Added: append the metrics to the tracking lists ---
        self.train_loss_history.append(train_loss)
        self.test_acc_history.append(test_acc)
        self.comput_step_history.append(self.comput_step)
        self.update_iter_history.append(self.update_iter)
        # --- End of added block ---

        # # wandb.log({"Iter/Update": self.update_iter, "Step/Comput": self.comput_step, \
        # #             "Train/Acc": train_acc, "Train/Loss": train_loss, \
        # #             "Test/Acc": test_acc, "Test/Loss": test_loss, \
        # #             "Dist/Opt": dist})  # Remove or comment out the wandb.log call if unused

    def save_net(self, net_weight):
        # 1. Build the directory and full path in the standard way
        #    Pass each path component to os.path.join as a separate argument
        model_dir = os.path.join(self.args.modeldir, self.args.dataset)
        model_path = os.path.join(model_dir, f"{self.args.run_name}-model.pth")

        # 2. Ensure that the target directory exists before saving
        #    os.makedirs(..., exist_ok=True) creates any missing intermediate directories safely
        os.makedirs(model_dir, exist_ok=True)

        # 3. The model can now be saved safely
        torch.save(net_weight, model_path)
        self.logger.info(f"Model saved to {model_path}")
        return

    def resume(self, net):
        model_path = os.path.join(self.args.modeldir + self.args.dataset, f"{self.args.run_name}-model.pth")
        if os.path.exists(model_path):
            self.logger.info(f"Resume from {model_path}")
            net.load_state_dict(torch.load(model_path))
        return net

    # --- Added: helper for saving results to CSV files ---
    def save_results_to_csv(self):
        self.logger.info("Saving results to CSV files...")

        # Save the test accuracy history
        df_acc = pd.DataFrame({
            'Step/Comput': self.comput_step_history,
            'Update/Iter': self.update_iter_history,
            'Test/Acc': self.test_acc_history
        })
        # Append rows using mode='a'; header=False avoids writing the header repeatedly.
        df_acc.to_csv(self.log_file_path_acc, mode='a', header=False, index=False)
        self.logger.info(f"Test accuracy saved to {self.log_file_path_acc}")

        # Save the training loss history
        df_loss = pd.DataFrame({
            'Step/Comput': self.comput_step_history,
            'Update/Iter': self.update_iter_history,
            'Train/Loss': self.train_loss_history
        })
        df_loss.to_csv(self.log_file_path_loss, mode='a', header=False, index=False)
        self.logger.info(f"Training loss saved to {self.log_file_path_loss}")
    # --- End of added block ---