import sys
import torchaudio
import torch
import torch.nn.functional as F

def load_fx_encoder_plusplus_2048(model_args, device, *args, **kwargs):
    from utils.feature_extractors.fx_encoder_plus_plus import load_model 

    assert model_args is not None, "model_args must be provided for fx_encoder type"

    ckpt_path=model_args.ckpt_path

    model=load_model(
        model_path=ckpt_path,
        device=device,
    )

    def effects_encoder_fn(x):
        assert x.ndim == 3, f"Input tensor x must be 2D, got {x.ndim}D"
        assert x.shape[1] == 2, f"Input tensor x must have 2 channels, got {x.shape[1]} channels"

        emb=model.fx_encoder(x)
        emb=emb["embedding"]  # Extract the embedding from the dictionary

        return emb
        

    return lambda x: effects_encoder_fn(x)

def add_isotropic_noise(z, sigma=0.1):
        """
        z: [..., D] normalized embeddings (e.g., from CLAP or a regressor)
        sigma: scale of noise to inject
        Returns: z with orthogonal Gaussian noise added
        """
        n=torch.randn_like(z)  # isotropic noise
        z_noisy = F.normalize(z + sigma * n, dim=-1)
        return z_noisy




def load_CLAP(model_args, device, *args, **kwargs):

    #original_path = sys.path.copy()
    from utils.laion_clap.hook import CLAP_Module
    model= CLAP_Module(enable_fusion=False, amodel= 'HTSAT-base')
    #sys.path = original_path


    print("checkpoint",model_args.ckpt_path)
    #print current sys.path
    print("sys.path", sys.path)
    model.load_ckpt(model_args.ckpt_path)
    model.to(device)

    normalize = model_args.normalize

    if model_args.use_adaptor:
        if model_args.adaptor_type == "MLP_CLAP_regressor":
            from networks.MLP_CLAP_regressor import MLP_CLAP_regressor
            adaptor=MLP_CLAP_regressor()
            ckpt=torch.load(model_args.adaptor_checkpoint, map_location=device, weights_only=False)
            adaptor.load_state_dict(ckpt["network"], strict=True)
            adaptor.to(device)

        


    def clap_fn(x, type=None):
        B, C, T = x.shape
        if C > 1:
            x= x.mean(dim=1, keepdim=True)  # Convert to mono if stereo

        with torch.no_grad():
            x=torchaudio.functional.resample(x, orig_freq=44100, new_freq=48000)
            x= x.squeeze(1)  # Remove channel dimension for CLAP
            emb=model.get_audio_embedding_from_data(x,use_tensor=True)

            if type is not None:
                if type == "wet":
                    #print("wet mode")
                    if model_args.use_adaptor:
                        emb= adaptor(emb)  # Apply the adaptor if specified
    
            if model_args.add_noise:
                emb= torch.nn.functional.normalize(emb, p=2, dim=-1)  # Normalize before adding noise
                emb = add_isotropic_noise(emb, sigma=model_args.noise_sigma)

            # Normalize the embeddings
            if normalize:
                emb = torch.nn.functional.normalize(emb, p=2, dim=-1)

            return emb

    return lambda x, type: clap_fn(x, type=type)