from charged_shells import expansion, parameters
import numpy as np
import matplotlib.pyplot as plt
import plotly.graph_objects as go
from config import *
import quadrupole_model_mappings

Expansion = expansion.Expansion


def plot_theta_profile(ex: Expansion, phi: float = 0, num: int = 100, theta_start: float = 0, theta_end: float = np.pi):
    theta_vals = np.linspace(theta_start, theta_end, num)
    charge = ex.charge_value(theta_vals, phi)
    plt.plot(theta_vals, charge.T)
    plt.show()


def plot_theta_profile_multiple(ex_list: list[Expansion], label_list, phi: float = 0, num: int = 100,
                                theta_start: float = 0, theta_end: float = np.pi):
    theta_vals = np.linspace(theta_start, theta_end, num)

    fig, ax = plt.subplots()
    for ex, label in zip(ex_list, label_list):
        ax.plot(theta_vals, ex.charge_value(theta_vals, phi).T, label=label)
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
    ax.set_xlabel(r'$\theta$', fontsize=13)
    ax.set_ylabel(r'$\sigma$', fontsize=13)
    plt.legend(fontsize=12)
    plt.tight_layout()
    plt.savefig(FIGURES_PATH.joinpath("charge_shape_comparison.png"), dpi=600)
    plt.show()


def plot_charge_3d(ex: Expansion, num_theta=100, num_phi=100, save_as: Path = None):
    theta = np.linspace(0, np.pi, num_theta)
    phi = np.linspace(0, 2 * np.pi, num_phi)
    theta, phi = np.meshgrid(theta, phi)
    r = ex.charge_value(theta.flatten(), phi.flatten()).reshape(theta.shape)

    # Convert spherical coordinates to Cartesian coordinates
    x = np.sin(theta) * np.cos(phi)
    y = np.sin(theta) * np.sin(phi)
    z = np.cos(theta)

    # Create a heatmap on the sphere
    fig = go.Figure(data=go.Surface(x=x, y=y, z=z, surfacecolor=r,
                                    colorscale='RdBu', reversescale=True))
    fig.update_layout(scene=dict(aspectmode='data'))
    fig.update_layout(scene=dict(xaxis_title='', yaxis_title='', zaxis_title=''))

    # Remove axes planes, background, ticks, and labels
    fig.update_layout(scene=dict(xaxis=dict(showbackground=False, gridcolor='white', showticklabels=False, ticks=''),
                                 yaxis=dict(showbackground=False, gridcolor='white', showticklabels=False, ticks=''),
                                 zaxis=dict(showbackground=False, gridcolor='white', showticklabels=False, ticks='')))

    # Adjust the width and height for higher resolution
    fig.update_layout(width=1200, height=1200)

    # Save as PNG with higher resolution
    if save_as is not None:
        fig.write_image(save_as, scale=3)  # Adjust the scale as needed

    fig.show()


def main():

    params = parameters.ModelParams(kappaR=3, R=150)
    # ex = expansion.MappedExpansionQuad(0.328, params.kappaR, 0.001, 30)
    # ex = expansion.Expansion(np.arange(3), np.array([1, -1, 0, 1, 2, 0, 3, 0, 2]))
    # ex = expansion.GaussianCharges(omega_k=np.array([[0, 0], [np.pi, 0]]), lambda_k=2.676, sigma1=0.00044, l_max=30)
    # ex = expansion.SphericalCap(np.array([[0, 0], [np.pi, 0]]), 0.894, 0.00132, 50)
    # ex = quadrupole_model_mappings.ic_to_gauss(0.001, 0.328, params, l_max=30)
    ex = quadrupole_model_mappings.ic_to_cap(0.001, 0.328, params, l_max=50)
    # print(np.real(ex.coefs))
    # plot_theta_profile(ex, num=1000, theta_end=2 * np.pi, phi=0)
    plot_charge_3d(ex, save_as=FIGURES_PATH.joinpath("model_3D_cap.png"))

    # new_coeffs = expanison.expansion_rotation(Quaternion(np.arange(20).reshape(5, 4)).normalized, ex.coeffs, ex.l_array)
    # print(new_coeffs.shape)
    #
    # newnew_coeffs = expansion.expansion_rotation(Quaternion(np.arange(16).reshape(4, 4)).normalized, new_coeffs, ex.l_array)
    # print(newnew_coeffs.shape)


if __name__ == '__main__':

    main()