import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import multivariate_normal


def importance_sampling(target_pdf, proposal_pdf, proposal_sample, num_samples):
    """
    Perform importance sampling.

    Parameters:
    target_pdf: The PDF of the target distribution.
    proposal_pdf: The PDF of the proposal distribution.
    proposal_sample: A function that generates a sample from the proposal distribution.
    num_samples: The number of samples to generate.

    Returns:
    samples: The generated samples.
    weights: The importance weights of the samples.
    """
    # Generate samples from the proposal distribution
    samples = np.array([proposal_sample() for _ in range(num_samples)])

    # Calculate the importance weights
    weights = np.array([target_pdf(sample) / proposal_pdf(sample) for sample in samples])

    # Normalize the weights
    weights /= np.sum(weights)

    return samples, weights



# Define the target and proposal distributions
target_pdf = lambda x: multivariate_normal.pdf(x, mean=[0, 0], cov=[[1, 0], [0, 1]])
proposal_pdf = lambda x: multivariate_normal.pdf(x, mean=[1, 1], cov=[[2, 0], [0, 2]])
proposal_sample = lambda: np.random.multivariate_normal([1, 1], [[2, 0], [0, 2]])

# Perform importance sampling
samples, weights = importance_sampling(target_pdf, proposal_pdf, proposal_sample, 1000)



# Define the target and proposal distributions
target_pdf = lambda x: multivariate_normal.pdf(x, mean=[0, 0], cov=[[1, 0], [0, 1]])
proposal_pdf = lambda x: multivariate_normal.pdf(x, mean=[1, 1], cov=[[2, 0], [0, 2]])

# Generate a grid of points
x = np.linspace(-3, 3, 100)
y = np.linspace(-3, 3, 100)
X, Y = np.meshgrid(x, y)
pos = np.dstack((X, Y))

# Calculate the PDFs of the target and proposal distributions at each point
target_Z = np.array([[target_pdf([x, y]) for x, y in zip(X_row, Y_row)] for X_row, Y_row in zip(X, Y)])
proposal_Z = np.array([[proposal_pdf([x, y]) for x, y in zip(X_row, Y_row)] for X_row, Y_row in zip(X, Y)])

new_samples = np.zeros(samples.shape)
for i in range(weights.shape[0]):
    new_samples[i, :] = samples[i, 0] * weights[i], samples[i, 1] * weights[i]

print(new_samples)

# Plot the target distribution
plt.figure()
plt.contourf(X, Y, target_Z, cmap='viridis')
plt.scatter(new_samples[:, 0], new_samples[:, 1], color='r', s=10)  # Add the sample points
plt.title('Target Distribution')
plt.xlabel('x')
plt.ylabel('y')
plt.colorbar(label='PDF')
plt.show()

# Plot the proposal distribution
plt.figure()
plt.contourf(X, Y, proposal_Z, cmap='viridis')
plt.scatter(samples[:, 0], samples[:, 1], color='r', s=10)  # Add the sample points
plt.title('Proposal Distribution')
plt.xlabel('x')
plt.ylabel('y')
plt.colorbar(label='PDF')
plt.show()