from charged_shells import expansion
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import plotly.graph_objects as go

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(Path("/home/andraz/ChargedShells/Figures/charge_shape_comparison.png"), dpi=600)
    plt.show()


def plot_charge_3d(ex: Expansion, num_theta=100, num_phi=100):
    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='Jet'))
    fig.update_layout(scene=dict(aspectmode='data'))
    fig.update_layout(scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'))

    fig.show()


def main():
    # ex = expansion.MappedExpansionQuad(0.44, 3, 1, 10)
    # 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=10, sigma1=0.001, l_max=10)
    ex = expansion.SphericalCap(np.array([[0, 0], [np.pi, 0]]), 0.5, 0.1, 30)
    # print(np.real(ex.coefs))
    # plot_theta_profile(ex, num=1000, theta_end=2 * np.pi, phi=0)
    plot_charge_3d(ex)

    # 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()