from .backend import Backend
import torch.multiprocessing as mp
from hypergan.gan_component import ValidationException, GANComponent
import torch.utils.data as data
import hyperchamber as hc
import hypergan as hg
import copy
import torch
import time

def create_input(input_config):
    klass = GANComponent.lookup_function(None, input_config['class'])
    return klass(input_config)


def train(device, head_device, gan, inputs, loaded_event, report_weights_queue, set_weights_queue, report_weights_event, save_file):
    gan.inputs = inputs
    gan.device = device
    #torch.manual_seed(device)
    gan = gan.to(device)
    #optimizer_state_dict = gan.trainer.optimizer.state_dict()
    gan.trainer = gan.create_component("trainer")
    #gan.trainer.optimizer.load_state_dict(optimizer_state_dict)
    if gan.load(save_file):
        print("Model loaded")
    else:
        print("Initializing new model")


    loaded_event.set()
    while(True):
        gan.step()
        if report_weights_event.is_set():
            if device == head_device:
                report_weights_queue.put(list([p.clone().detach() for p in gan.parameters()]))
            else:
                report_weights_queue.put(list(gan.parameters()))
            params = set_weights_queue.get()
            for p1, p2 in zip(gan.parameters(), params):
                with torch.no_grad():
                    if(p2.device != p1.device):
                        p2 = p2.to(p1.device)
                    p1.set_(p2)
            del params
            report_weights_event.clear()


class RoundrobinBackend(Backend):
    """Trains separately then syncs each card, round robin style"""
    def __init__(self, gan, cli, devices):
        self.gan = gan
        self.cli = cli
        self.sync = 0
        self.processes = []
        self.report_weights_event = []
        self.report_weights_queue = []
        self.set_weights_queue = []
        loaded_events = []
        head_device = torch.device(list(self.gan.parameters())[0].device).index
        self.head_device=head_device
        mp.set_start_method('spawn', force=True)
        if devices == "-1":
            print("Running on all available devices: ", torch.cuda.device_count())
            devices = list(range(torch.cuda.device_count()))
        else:
            devices = [int(d) for d in devices.split(",")]
        self.devices = devices
        print("Devices:", devices)
        for device in devices:
            loaded_event = mp.Event()
            report_weights_event = mp.Event()
            set_weights_queue = mp.Queue()
            report_weights_queue = mp.Queue()
            inputs = cli.create_input(rank=device)
            p = mp.Process(target=train, args=(device, head_device, gan, inputs, loaded_event, report_weights_queue, set_weights_queue, report_weights_event, self.cli.save_file))
            p.start()
            self.processes.append(p)
            self.report_weights_event.append(report_weights_event)
            self.report_weights_queue.append(report_weights_queue)
            self.set_weights_queue.append(set_weights_queue)
            loaded_event.wait()

    def step(self):
        time.sleep(2.0)
        selected = self.sync % len(self.processes)
        print("Syncing", selected)
        self.report_weights_event[selected].set()
        params = self.report_weights_queue[selected].get()
        for p1, p2 in zip(self.gan.parameters(), params):
            with torch.no_grad():
                if(p2.device != p1.device):
                    p2 = p2.to(p1.device)
                p1.set_((p1+p2)/2.0)
        del params
        if self.devices[selected] == self.head_device:
            self.set_weights_queue[selected].put(list([p.clone().detach() for p in self.gan.parameters()]))
        else:
            self.set_weights_queue[selected].put(list(self.gan.parameters()))
        self.sync += 1
        self.cli.sample()
        if self.sync % 600 == 0:
            print("Saving from roundrobin")
            self.gan.save(self.cli.save_file)

@@ -568,236 +568,77 @@
     
-#torch.manual_seed(device)%0A    gan = gan.to(device)%0A    #optimizer_state_dict = gan.trainer.optimizer.state_dict()%0A    gan.trainer = gan.create_component(%22trainer%22)%0A    #gan.trainer.optimizer.load_state_dict(optimizer_state_dict
+gan = gan.to(device)%0A    gan.trainer = gan.create_component(%22trainer%22
 )%0A  
@@ -744,17 +744,16 @@
 odel%22)%0A%0A
-%0A
     load
@@ -1392,17 +1392,16 @@
 lear()%0A%0A
-%0A
 class Ro

remove unused code