import torch
import numpy as scinp
from QuantumGenerator import *
import matplotlib.pyplot as plt


# This function is useful when monitoring gradients during training
def print_gradients(model):
    grads = []
    for param in model.parameters():
        grads.append(param.grad.view(-1))
    grads = torch.cat(grads)
    print(grads)


# This is also useful when monitoring training
def print_parameters(model):
    print(list(model.parameters()))


# This is a simple function. It simply prints the critic's
# output on a true sample from X, and a fake sample from the
# generator. If the critic output is very high for the
# true sample, and close to zero for the generator, then
# that impedes the generator's learning.
def check_both_models_strength(X, critic, generator, device):
    true_sample = torch.Tensor(X[0]).to(device)
    output_true_sample = critic(true_sample).cpu().detach().numpy()[0]
    fake_sample = generator()
    output_fake_sample = critic(fake_sample).cpu().detach().numpy()[0]
    print(
        "Critic output for:\tTrue Sample: {:0.2f}\tFake Sample: {:0.2f}".format(output_true_sample, output_fake_sample))


# This function is useful if we simply want to see what the generated
# vectors are.
def monitor_generated_samples(generator, n_monitoring_samples=10):
    generated_samples = []
    for i in range(n_monitoring_samples):
        s = generator().cpu().detach().numpy()
        generated_samples.append(s)
    generated_samples = scinp.array(generated_samples)
    average_sample = scinp.mean(generated_samples, axis=0)
    shortest_vector_index = scinp.argmin(scinp.linalg.norm(generated_samples, axis=1, ord=2))
    shortest_vector = generated_samples[shortest_vector_index]
    longest_vector_index = scinp.argmax(scinp.linalg.norm(generated_samples, axis=1, ord=2))
    longest_vector = generated_samples[longest_vector_index]
    print("Generation Info: Average Vector:{}\tShortest Vector:{}\tLongest Vector:{}".format(average_sample,
                                                                                             shortest_vector,
                                                                                             longest_vector))


# Generate some simple data with sines and cosines as building blocks.
# We can use this as a simple sanity check for the algorithm, since it
# is 3D data and thus we can plot the samples generated by the hybrid
# quantum classical generator and see if they are starting to look
# like our the true samples.
def generate_sinusoidal_data(n_points):
    samples = []
    for i in range(n_points):
        x = scinp.random.uniform(low=-3, high=3)
        y = scinp.random.uniform(low=-3, high=3)
        z = scinp.cos(x) ** 2 + scinp.sin(y)
        samples.append(scinp.array([x, y, z]))
    return scinp.array(samples)


# Plot fake samples generated by a generator. Only works for
# 3D settings.
def plot_generator_3D(generator: QuantumGenerator, n_points=2000):
    generated_samples = [generator().cpu().detach().numpy() for i in range(n_points)]
    generated_samples = scinp.array(generated_samples)
    fig = plt.figure(figsize=(12, 12))
    ax = fig.add_subplot(111, projection='3d')
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('f(x,y)')
    ax.scatter(generated_samples[:, 0], generated_samples[:, 1], generated_samples[:, 2])
    plt.show()
