import torch
import numpy as np

def compute_dataset_embeds_svd(all_embeds, rank):

    # Perform SVD on the combined matrix
    u, s, vh = np.linalg.svd(all_embeds, full_matrices=False)

    # Select the top `rank` singular vectors to construct the projection matrix
    vh = vh[:rank]  # Top `rank` right singular vectors
    projection_matrix = vh.T @ vh  # Shape: (feature_dim, feature_dim)

    return projection_matrix

def get_embedding_composition(embed, projections_data):
    # Initialize the combined embedding with the input embed
    combined_embeds = embed.copy()

    for proj_data in projections_data:

        # Add the combined projection to the result
        combined_embeds -= embed @ proj_data["projection_matrix"]
        combined_embeds += proj_data["embed"] @ proj_data["projection_matrix"]

    return combined_embeds


def get_modified_images_embeds_composition(embed, projections_data, ip_model, prompt=None, scale=1.0, num_samples=3, seed=420, num_inference_steps=50):
    
    final_embeds = get_embedding_composition(embed, projections_data)
    clip_embeds = torch.from_numpy(final_embeds)

    images = ip_model.generate(clip_image_embeds=clip_embeds, prompt=prompt, num_samples=num_samples, num_inference_steps=num_inference_steps, seed=seed, guidance_scale=7.5, scale=scale)
    return images