import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
from IPython.display import display
from qutip import (about, basis, coherent, coherent_dm, displace, fock, ket2dm,
                   plot_wigner, squeeze, thermal_dm, wigner_cmap, wigner)

import scipy.ndimage
import os
from tqdm import tqdm

def rotate_and_crop(array, angle, xvec, yvec):
    rotated_array = scipy.ndimage.rotate(array, -angle, reshape=False)
    rows, cols = rotated_array.shape
    center_row, center_col = rows // 2, cols // 2
    target_rows, target_cols = len(yvec), len(xvec)
    start_row = center_row - target_rows // 2
    end_row = start_row + target_rows
    start_col = center_col - target_cols // 2
    end_col = start_col + target_cols
    return rotated_array[start_row:end_row, start_col:end_col]

def plot_wigner_marginals(W, xvec, yvec, marginal_max, resolution=200, angle=0):
    wmap = wigner_cmap(W)
    wlim = np.abs(W).max()
    cmap = plt.colormaps['RdBu']

    fig = plt.figure()
    n, m = 5, 1
    fig, axes = plt.subplot_mosaic(
        [ ["top"] * n + ["3d"] * m ] * m + [ ["mid"] * n + ["right"] * m] * n,
    figsize=(20, 20),
    layout="constrained",
    width_ratios=[1.05] * (n+m))

    ax = axes["mid"]
    norm = mpl.colors.Normalize(-wlim, wlim)
    ax.contourf(xvec, yvec, W, resolution // 3, norm=norm, cmap=cmap)
    

    ax = axes["top"]
    x_marginal = np.sum(W, axis=0)
    y_marginal = np.sum(W, axis=1)
    ax.fill_between(xvec, x_marginal, 0, color='#938fba', alpha=0.5)
    ax.plot(xvec, x_marginal, color='#4a5a90')
    ax.set_xlim(min(xvec), max(xvec))
    ax.set_ylim(0, marginal_max * 1.05)
    ax.set_xticks([])
    ax.set_yticks([])

    ax = axes["right"]
    ax.fill_betweenx(yvec, np.sum(W, axis=1), 0, color='#938fba', alpha=0.5)
    ax.plot(y_marginal, yvec, color='#4a5a90')
    ax.set_xlim(0, marginal_max * 1.05)
    ax.set_ylim(min(yvec), max(yvec))
    ax.set_xticks([])
    ax.set_yticks([])

    ax = axes["3d"]
    ax.axis('off')

    return fig

def plot_wigner_with_marginals(psi, **kwargs):
    
    radius = kwargs.get('radius', 5) 
    resolution = kwargs.get('resolution', 500)
    angles = kwargs.get('angles', np.linspace(0, 2*np.pi, 100))
    dir_path = kwargs.get('dir_path', './output')
    
    xvec_upscaled = np.linspace(-radius*1.5, radius*1.5, int(resolution*1.5))
    yvec_upscaled = np.linspace(-radius*1.5, radius*1.5, int(resolution*1.5))
    xvec = np.linspace(-radius, radius, int(resolution))
    yvec = np.linspace(-radius, radius, int(resolution))

    W_upscaled = wigner(psi, xvec_upscaled, yvec_upscaled)
    marginal_max = max(max(np.sum(W_upscaled, axis=0)), max(np.sum(W_upscaled, axis=1)))
    print(f"outputting to {dir_path}")
    for N, angle in tqdm(enumerate(angles)):
        W = rotate_and_crop(W_upscaled, angle, xvec, yvec)
        fig = plot_wigner_marginals(W, xvec, yvec, marginal_max=marginal_max, resolution=resolution, angle=angle)

        if not os.path.exists(dir_path):
            os.makedirs(dir_path)
        fig.savefig(f"{dir_path}/{N:03d}.png",bbox_inches='tight')
        plt.close(fig)
    
mpl.use('agg')

configs = {
    "N_dim" : 40,
    "radius" : 3.5,
    "resolution" : 500,
    "angles" : [i * 2 for i in range(180)],
    "dir_path" : ""
}

for n1, n2 in [(0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2)]:
    psi = (fock(configs["N_dim"], n1) + fock(configs["N_dim"], n2)).unit()
    configs["dir_path"] = f"./fock/fock_{n1}_{n2}"
    plot_wigner_with_marginals(psi, **configs)