from common import layers

class Transferer():
    def __init__(self):
        self.registered_types = {}

    def register(self, input_class, output_class):
        self.registered_types[input_class] = output_class

    def transfer(self, input_model, batch_size, quantization_interval):
        input_class = input_model.__class__
        if not input_class in self.registered_types:
            raise Exception(f"No transfer function is registered for class {input_class}")
        transfered_model = self.transfer_model(input_model, self.registered_types[input_class], batch_size, quantization_interval)
        return transfered_model

    def transfer_model(self, model, output_model_class, batch_size, quantization_interval):
        qmodel = output_model_class(config = model.config)
        model.config.quantization_interval = quantization_interval
        for attr, value in model.named_children():
            if isinstance(value, layers.UpscaleBlock):
                getattr(qmodel, attr).stage = getattr(model, attr).stage.get_lut_model(quantization_interval=quantization_interval, batch_size=batch_size)

        return qmodel

TRANSFERER = Transferer()