"""Pipeline for contrastive pre-training of an unitary encoder"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../src/pipeline/unitary_clip_pipeline.ipynb.

# %% auto #0
__all__ = ['UnitaryCLIPPipeline']

# %% ../../src/pipeline/unitary_clip_pipeline.ipynb #06272f6f-b4e3-4504-a90a-feebbf6ad821
from ..imports import *
from .pipeline import Pipeline
from ..utils.config_loader import *
from ..models.config_model import ConfigModel

# %% ../../src/pipeline/unitary_clip_pipeline.ipynb #9ea10ee8-e231-421c-8fca-3c5dc3e57743
class UnitaryCLIPPipeline(Pipeline):
    def __init__(self, 
                 model: nn.Module,
                 device: torch.device) -> None:
        super().__init__(model, device)   
        
        self.loss   = nn.CrossEntropyLoss()
        self.device = device

    #------------------------------------
    
    add_config = {}
    
    def params_config(self, save_path: str) -> dict:         
        params_config = {}
                
        params_config["model"]                = self.model.get_config(save_path=save_path+"model")
        params_config["unitary_text_encoder"] = self.model.unitary_text_encoder.get_config(save_path=None)
        params_config["circuit_encoder"]      = self.model.circuit_encoder.get_config(save_path=None)
        
        params_config["device"]                = str(self.device)
        params_config["add_config"]            = self.add_config
        
        return params_config

    def store_pipeline(self, config_path: str, save_path: str):
        super().store_pipeline(config_path, save_path)
        config = self.get_config(save_path)
        save_dict_yaml(config, config_path+"config.yaml")
               
        self.model.store_model(config_path=None, save_path=save_path+"model")

    @staticmethod
    def from_config_file(config_path, device: torch.device, save_path: str=None):    
        config = load_config(config_path+"config.yaml")   
        config = config_to_dict(config)

        def _get_save_path(config_save_path, appendix):
            _save_path = default(save_path, config_path) + appendix
            if "save_path" in config_save_path:
                _save_path = config_save_path["save_path"]
            return _save_path  

        if exists(device):
            config["params"]["device"] = device
            config["params"]["model"]["params"]["text_encoder_config"]["device"] = device

        unitary_text_encoder = ConfigModel.from_config(config["params"].pop("unitary_text_encoder", None), device, None) 
        circuit_encoder      = ConfigModel.from_config(config["params"].pop("circuit_encoder", None), device, None) 
        
        config["params"]["model"]["params"]["unitary_text_encoder"] = unitary_text_encoder
        config["params"]["model"]["params"]["circuit_encoder"]      = circuit_encoder
        config["params"]["model"] = ConfigModel.from_config(config["params"]["model"], device, _get_save_path(config["params"]["model"], "model"))

        add_config = config["params"].pop("add_config", None)

        pipeline = instantiate_from_config(config)
        
        if exists(pipeline.add_config):
            pipeline.add_config = add_config
        
        return pipeline
    
    #------------------------------------
    # Inference functions
        
    @torch.no_grad()    
    def __call__(self, tokens: torch.Tensor, params: torch.Tensor, y: torch.Tensor, U: torch.Tensor, softmax=True) -> torch.Tensor:
        #compute the score of img-label pairs for classification!!
        self.model.eval()
        
        scores = self.model(tokens=tokens, params=params, y=y, U=U) #[b, b]  
        
        if softmax: 
            scores = F.softmax(scores, dim-1)
        
        return scores
        
    #------------------------------------
    # Training functions
                
    def get_loss(self, tokens: torch.Tensor, params: torch.Tensor, y: torch.Tensor, U: torch.Tensor) -> torch.Tensor:
        
        scores = self.model(tokens=tokens, params=params, y=y, U=U) #[b, b]
        
        #scores is: I=unitary_text   T=circuit
        #--------------------------------
        #| I1*T1   I1*T2   I1*T3   ...
        #| I2*T1
        #| I3*T1
        # ...
        #--------------------------------

        target = torch.arange(scores.shape[0], device=scores.device)
        
        loss_unitary_text = self.loss(scores  , target)
        loss_circuit      = self.loss(scores.T, target)
        
        #symmetric loss
        loss = (loss_unitary_text + loss_circuit) / 2.0
                
        return loss
        
    def train_step(self, data, **kwargs): 
        tokens, y, params, U = data 
        
        tokens = tokens.to(self.device)  
        params = params.to(self.device)   
        y = y.to(self.device)    
        U = U.to(torch.float32).to(self.device)
        
        loss = self.get_loss(tokens=tokens, params=params, y=y, U=U)

        return loss
