import os
import gc
import datetime
import torch  # import torch before lapse
try:
    import lapse
except ImportError:
    pass
from typing import Optional, Dict
from copy import deepcopy
from torch import multiprocessing as mp
from torch import distributed as dist

from kge.misc import set_seeds
from kge.job import Job
from kge.util.io import load_checkpoint
from .parameter_client import KgeParameterClient
from .misc import get_min_rank


class WorkerProcessPool:
    """
    Creates all the train-workers for distributed training
    """
    def __init__(
        self,
        num_total_workers,
        num_workers_machine,
        already_init_workers,
        num_keys,
        num_meta_keys,
        embedding_dim,
        optimizer_dim,
        config,
        dataset,
        checkpoint_name: Optional[str] = None,
    ):
        config.log(f"creating worker process pool with {num_total_workers} worker")
        self.workers = []
        configs = {}
        parameters=None
        if config.get("job.distributed.parameter_server") == "shared":
            parameters = torch.empty((num_keys, embedding_dim + optimizer_dim), dtype=torch.float32, requires_grad=False).share_memory_()
        for rank in range(num_workers_machine):
            if rank == 0:
                self.recv_end, send_end = mp.Pipe(False)
            else:
                send_end = None
            configs[rank] = deepcopy(config)
            #configs[rank].set(config.get("model") + ".create_complete", False)
            configs[rank].init_folder()
            worker = WorkerProcess(
                rank + already_init_workers,
                num_total_workers,
                num_keys,
                num_meta_keys,
                embedding_dim,
                optimizer_dim,
                configs[rank],
                dataset,
                parameters=parameters,
                checkpoint_name=checkpoint_name,
                result_pipe=send_end
            )
            worker.start()
            self.workers.append(worker)

    def join(self):
        """Wait for all workers"""
        try:
            valid_trace = self.recv_end.recv()
        except:
            valid_trace = None
        for worker in self.workers:
            worker.join()
        return valid_trace

    def terminate(self):
        print("terminating worker process pool")
        for worker in self.workers:
            worker.terminate()

    def kill(self):
        print("killing worker process pool")
        for worker in self.workers:
            worker.kill()


class WorkerProcess(mp.get_context("spawn").Process):
    """Train worker"""
    def __init__(
        self,
        rank,
        num_total_workers,
        num_keys,
        num_meta_keys,
        embedding_dim,
        optimizer_dim,
        config,
        dataset,
        parameters=None,
        checkpoint_name: Optional[str] = None,
        result_pipe=None,
    ):
        # rank = rank + 1
        daemon = config.get("train.num_workers") <= 0 and config.get("eval.num_workers") <= 0
        super().__init__(daemon=daemon, name=f"Worker #{rank}")
        self.rank = rank
        self.num_total_workers = num_total_workers
        self.num_keys = num_keys
        self.num_meta_keys = num_meta_keys
        self.embedding_dim = embedding_dim
        self.optimizer_dim = optimizer_dim
        self.config = config
        self.dataset = dataset
        self.parameters = parameters
        self.checkpoint_name = checkpoint_name
        self.result_pipe = result_pipe

    def run(self):
        torch_device = self.config.get("job.device")
        if self.config.get("job.device") == "cuda":
            torch_device = "cuda:0"
        if torch_device != "cpu":
            torch.cuda.set_device(torch_device)
        # seeds need to be set in every process
        set_seeds(self.config, self.rank)

        os.environ["MASTER_ADDR"] = self.config.get("job.distributed.master_ip")
        os.environ["MASTER_PORT"] = str(self.config.get("job.distributed.master_port"))
        min_rank = get_min_rank(self.config)
        print("before init", self.rank + min_rank)
        dist.init_process_group(
            backend="gloo",
            init_method="env://",
            world_size=self.num_total_workers + min_rank,
            rank=self.rank + min_rank,
            timeout=datetime.timedelta(hours=6),
        )
        worker_ranks = list(range(min_rank, self.num_total_workers+min_rank))
        worker_group = dist.new_group(worker_ranks, timeout=datetime.timedelta(hours=6))
        num_eval_workers = self.config.get("job.distributed.num_eval_workers")
        eval_worker_ranks = list(range(min_rank, num_eval_workers+min_rank))
        eval_worker_group = dist.new_group(eval_worker_ranks, timeout=datetime.timedelta(hours=6))

        # create parameter server
        server = None
        if self.config.get("job.distributed.parameter_server") == "lapse":
            os.environ["DMLC_NUM_WORKER"] = "0"
            os.environ["DMLC_NUM_SERVER"] = str(self.num_total_workers)
            os.environ["DMLC_ROLE"] = "server"
            os.environ["DMLC_PS_ROOT_URI"] = self.config.get(
                "job.distributed.master_ip"
            )
            os.environ["DMLC_PS_ROOT_PORT"] = str(self.config.get(
                "job.distributed.lapse_port"
            ))

            num_workers_per_server = 1
            lapse.setup(self.num_keys, num_workers_per_server)
            server = lapse.Server(self.num_keys, self.embedding_dim + self.optimizer_dim)
        elif self.config.get("job.distributed.parameter_server") == "shared":
            server = self.parameters

        # create train-worker config, dataset and folder
        device_pool: list = self.config.get("job.device_pool")
        if len(device_pool) == 0:
            device_pool.append(self.config.get("job.device"))
        worker_id = self.rank
        config = deepcopy(self.config)
        config.set("job.device", device_pool[worker_id % len(device_pool)])
        config.folder = os.path.join(self.config.folder, f"worker-{self.rank}")
        config.init_folder()

        parameter_client = KgeParameterClient.create(
            client_type=self.config.get("job.distributed.parameter_server"),
            server_id=0,
            client_id=worker_id + min_rank,
            embedding_dim=self.embedding_dim + self.optimizer_dim,
            server=server,
            num_keys=self.num_keys,
            num_meta_keys=self.num_meta_keys,
            worker_group=worker_group,
            eval_worker_group=eval_worker_group
        )
        # don't re-initialize the model after loading checkpoint
        init_for_load_only = self.checkpoint_name is not None
        job = Job.create(
            config=config,
            dataset=self.dataset,
            parameter_client=parameter_client,
            init_for_load_only=init_for_load_only,
        )
        if self.checkpoint_name is not None:
            checkpoint = load_checkpoint(self.checkpoint_name)
            job._load(checkpoint)
            job.load_distributed(checkpoint_name=self.checkpoint_name)

        job.run()

        # all done, clean up
        print("shut down everything")
        parameter_client.barrier()
        if hasattr(job, "work_scheduler_client"):
            job.work_scheduler_client.shutdown()
        parameter_client.shutdown()
        # delete all occurrences of the parameter client to properly shutdown lapse
        # del job
        del job.parameter_client
        del job.model.get_s_embedder().parameter_client
        del job.model.get_p_embedder().parameter_client
        del job.model
        if hasattr(job, "optimizer"):
            del job.optimizer
        del parameter_client
        gc.collect()  # make sure lapse-worker destructor is called
        # shutdown server
        if server is not None and type(server) != torch.Tensor:
            server.shutdown()
        if self.result_pipe is not None:
            if hasattr(job, "valid_trace"):
                # if we valid from checkpoint there is no valid trace
                self.result_pipe.send(job.valid_trace)
            else:
                self.result_pipe.send(None)

import os
import gc
import datetime
import torch  # import torch before lapse
try:
    import lapse
except ImportError:
    pass
from typing import Optional, Dict
from copy import deepcopy
from torch import multiprocessing as mp
from torch import distributed as dist

from kge.misc import set_seeds
from kge.job import Job
from kge.util.io import load_checkpoint
from .parameter_client import KgeParameterClient
from .misc import get_min_rank


class WorkerProcessPool:
    """
    Creates all the train-workers for distributed training
    """
    def __init__(
        self,
        num_total_workers,
        num_workers_machine,
        already_init_workers,
        num_keys,
        num_meta_keys,
        embedding_dim,
        optimizer_dim,
        config,
        dataset,
        checkpoint_name: Optional[str] = None,
    ):
        config.log(f"creating worker process pool with {num_total_workers} worker")
        self.workers = []
        configs = {}
        parameters=None
        if config.get("job.distributed.parameter_server") == "shared":
            parameters = torch.empty((num_keys, embedding_dim + optimizer_dim), dtype=torch.float32, requires_grad=False).share_memory_()
        for rank in range(num_workers_machine):
            if rank == 0:
                self.recv_end, send_end = mp.Pipe(False)
            else:
                send_end = None
            configs[rank] = deepcopy(config)
            #configs[rank].set(config.get("model") + ".create_complete", False)
            configs[rank].init_folder()
            worker = WorkerProcess(
                rank + already_init_workers,
                num_total_workers,
                num_keys,
                num_meta_keys,
                embedding_dim,
                optimizer_dim,
                configs[rank],
                dataset,
                parameters=parameters,
                checkpoint_name=checkpoint_name,
                result_pipe=send_end
            )
            worker.start()
            self.workers.append(worker)

    def join(self):
        """Wait for all workers"""
        try:
            valid_trace = self.recv_end.recv()
        except:
            valid_trace = None
        for worker in self.workers:
            worker.join()
        return valid_trace

    def terminate(self):
        print("terminating worker process pool")
        for worker in self.workers:
            worker.terminate()

    def kill(self):
        print("killing worker process pool")
        for worker in self.workers:
            worker.kill()


class WorkerProcess(mp.get_context("spawn").Process):
    """Train worker"""
    def __init__(
        self,
        rank,
        num_total_workers,
        num_keys,
        num_meta_keys,
        embedding_dim,
        optimizer_dim,
        config,
        dataset,
        parameters=None,
        checkpoint_name: Optional[str] = None,
        result_pipe=None,
    ):
        # rank = rank + 1
        daemon = config.get("train.num_workers") <= 0 and config.get("eval.num_workers") <= 0
        super().__init__(daemon=daemon, name=f"Worker #{rank}")
        self.rank = rank
        self.num_total_workers = num_total_workers
        self.num_keys = num_keys
        self.num_meta_keys = num_meta_keys
        self.embedding_dim = embedding_dim
        self.optimizer_dim = optimizer_dim
        self.config = config
        self.dataset = dataset
        self.parameters = parameters
        self.checkpoint_name = checkpoint_name
        self.result_pipe = result_pipe

    def run(self):
        torch_device = self.config.get("job.device")
        if self.config.get("job.device") == "cuda":
            torch_device = "cuda:0"
        if torch_device != "cpu":
            torch.cuda.set_device(torch_device)
        # seeds need to be set in every process
        set_seeds(self.config, self.rank)

        os.environ["MASTER_ADDR"] = self.config.get("job.distributed.master_ip")
        os.environ["MASTER_PORT"] = str(self.config.get("job.distributed.master_port"))
        min_rank = get_min_rank(self.config)
        print("before init", self.rank + min_rank)
        dist.init_process_group(
            backend="gloo",
            init_method="env://",
            world_size=self.num_total_workers + min_rank,
            rank=self.rank + min_rank,
            timeout=datetime.timedelta(hours=6),
        )
        worker_ranks = list(range(min_rank, self.num_total_workers+min_rank))
        worker_group = dist.new_group(worker_ranks, timeout=datetime.timedelta(hours=6))
        num_eval_workers = self.config.get("job.distributed.num_eval_workers")
        eval_worker_ranks = list(range(min_rank, num_eval_workers+min_rank))
        eval_worker_group = dist.new_group(eval_worker_ranks, timeout=datetime.timedelta(hours=6))

        # create parameter server
        server = None
        if self.config.get("job.distributed.parameter_server") == "lapse":
            os.environ["DMLC_NUM_WORKER"] = "0"
            os.environ["DMLC_NUM_SERVER"] = str(self.num_total_workers)
            os.environ["DMLC_ROLE"] = "server"
            os.environ["DMLC_PS_ROOT_URI"] = self.config.get(
                "job.distributed.master_ip"
            )
            os.environ["DMLC_PS_ROOT_PORT"] = str(self.config.get(
                "job.distributed.lapse_port"
            ))

            num_workers_per_server = 1
            lapse.setup(self.num_keys, num_workers_per_server)
            server = lapse.Server(self.num_keys, self.embedding_dim + self.optimizer_dim)
        elif self.config.get("job.distributed.parameter_server") == "shared":
            server = self.parameters

        # create train-worker config, dataset and folder
        device_pool: list = self.config.get("job.device_pool")
        if len(device_pool) == 0:
            device_pool.append(self.config.get("job.device"))
        worker_id = self.rank
        config = deepcopy(self.config)
        config.set("job.device", device_pool[worker_id % len(device_pool)])
        config.folder = os.path.join(self.config.folder, f"worker-{self.rank}")
        config.init_folder()

        parameter_client = KgeParameterClient.create(
            client_type=self.config.get("job.distributed.parameter_server"),
            server_id=0,
            client_id=worker_id + min_rank,
            embedding_dim=self.embedding_dim + self.optimizer_dim,
            server=server,
            num_keys=self.num_keys,
            num_meta_keys=self.num_meta_keys,
            worker_group=worker_group,
            eval_worker_group=eval_worker_group
        )
        # don't re-initialize the model after loading checkpoint
        init_for_load_only = self.checkpoint_name is not None
        job = Job.create(
            config=config,
            dataset=self.dataset,
            parameter_client=parameter_client,
            init_for_load_only=init_for_load_only,
        )
        if self.checkpoint_name is not None:
            checkpoint = load_checkpoint(self.checkpoint_name)
            job._load(checkpoint)
            job.load_distributed(checkpoint_name=self.checkpoint_name)

        job.run()

        # all done, clean up
        print("shut down everything")
        parameter_client.barrier()
        if hasattr(job, "work_scheduler_client"):
            job.work_scheduler_client.shutdown()
        parameter_client.shutdown()
        # delete all occurrences of the parameter client to properly shutdown lapse
        # del job
        del job.parameter_client
        del job.model.get_s_embedder().parameter_client
        del job.model.get_p_embedder().parameter_client
        del job.model
        if hasattr(job, "optimizer"):
            del job.optimizer
        del parameter_client
        gc.collect()  # make sure lapse-worker destructor is called
        # shutdown server
        if server is not None and type(server) != torch.Tensor:
            server.shutdown()
        if self.result_pipe is not None:
            if hasattr(job, "valid_trace"):
                # if we valid from checkpoint there is no valid trace
                self.result_pipe.send(job.valid_trace)
            else:
                self.result_pipe.send(None)

# import torch before lapse Creates all the train-workers for distributed training #configs[rank].set(config.get("model") + ".create_complete", False) Wait for all workers Train worker # rank = rank + 1 #{rank}") # seeds need to be set in every process # create parameter server # create train-worker config, dataset and folder # don't re-initialize the model after loading checkpoint # all done, clean up # delete all occurrences of the parameter client to properly shutdown lapse # del job # make sure lapse-worker destructor is called # shutdown server # if we valid from checkpoint there is no valid trace