#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Jun  4 14:57:19 2026

@author: leonhardtm
"""

import pandas as pd
import numpy as np
from pathlib import Path
import qutip
import matplotlib.pyplot as plt
from matplotlib import colormaps
from matplotlib.colors import ListedColormap
from matplotlib.patches import FancyArrowPatch
from mpl_toolkits.mplot3d.proj3d import proj_transform
from mpl_toolkits.mplot3d.axes3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Line3DCollection
import matplotlib.colors as colors

from typing import Optional


# create a new Arrow3D class and add corresponding method to Axes3D
class Arrow3D(FancyArrowPatch):

    def __init__(self, x, y, z, dx, dy, dz, *args, **kwargs):
        super().__init__((0, 0), (0, 0), *args, **kwargs)
        self._xyz = (x, y, z)
        self._dxdydz = (dx, dy, dz)

    def draw(self, renderer):
        x1, y1, z1 = self._xyz
        dx, dy, dz = self._dxdydz
        x2, y2, z2 = (x1 + dx, y1 + dy, z1 + dz)

        xs, ys, zs = proj_transform((x1, x2), (y1, y2), (z1, z2), self.axes.M)
        self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
        super().draw(renderer)

    def do_3d_projection(self, renderer=None):
        x1, y1, z1 = self._xyz
        dx, dy, dz = self._dxdydz
        x2, y2, z2 = (x1 + dx, y1 + dy, z1 + dz)

        xs, ys, zs = proj_transform((x1, x2), (y1, y2), (z1, z2), self.axes.M)
        self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))

        return np.min(zs)


def _arrow3D(ax, x, y, z, dx, dy, dz, *args, **kwargs):
    """Add an 3d arrow to an `Axes3D` instance."""

    arrow = Arrow3D(x, y, z, dx, dy, dz, *args, **kwargs)
    ax.add_artist(arrow)


setattr(Axes3D, "arrow3D", _arrow3D)


class CustomBloch(qutip.Bloch):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.q_grid = None
        self.line_kwargs = []
        self.point_kwargs = []
        self.arrow_vecs = []
        self.arrow_kwargs = []
        self.surface_args = None
        self.surface_kwargs = {}
        self.lcs = []
        self.lc_kwargs = []

    def add_points(self, points, meth="s", **kwargs):
        if meth == "l":
            self.line_kwargs.append(kwargs)
        elif meth == "s":
            self.point_kwargs.append(kwargs)
        super().add_points(points, meth=meth)

    def add_line_with_depth_effect(
        self,
        r_t,
        view_vector: Optional[np.ndarray] = None,
        alpha_min: float = 0.2,
        alpha_max: float = 1,
        color: str = "black",
        **kwargs,
    ):
        """Add line that with a transparency (alpha value) that depends on how far
        inside the curve is in the sphere. Points on the backside of the sphere
        have alpha=alpha_min and points on the frontside have alpha=alpha_max.

        :param r_t: The curve r(t) given as numpy array with shape (len(t), 3)
        :param view_vector: view vector of the Bloch sphere, if None use current
            view_vector, defaults to None
        :param alpha_min: Minimal alpha value, used for points on the backside of
            the sphere, defaults to 0.2
        :param alpha_max: Maximal alpha value, used for points on the backside of
            the sphere, defaults to 1
        :param color: Color of the line, defaults to "black"
        :return: the alpha values add each point
        """
        # create the colormap
        color_fade = colors.to_rgb(color) + (alpha_min,)
        color_full = colors.to_rgb(color) + (alpha_max,)

        cmap = colors.LinearSegmentedColormap.from_list("cmf", [color_fade, color_full])

        if view_vector is None:
            view = np.asarray(self.view) * np.pi / 180
            view_vector = [
                -np.sin(view[0]) * np.cos(view[1]),
                np.cos(view[0]) * np.cos(view[1]),
                np.sin(view[1]),
            ]

        # create new coordinate system, x',y',z' where z'=view_vector
        if (np.abs(view_vector) != np.asarray([0, 0, 1])).all():
            ey = np.cross(view_vector, np.asarray([0, 0, 1]))
            ey /= np.linalg.norm(ey)
            ez = view_vector / np.linalg.norm(view_vector)
            ex = np.cross(ey, ez)

            # tranform to new basis
            T = np.asarray([ex, ey, ez])
        else:
            T = np.eye(3)

        r_prime_t = T @ r_t.T

        delta = 1 - alpha_min
        # calculate alpha value for each point
        z_max_sq = (
            1 - r_prime_t[0, :] ** 2 - r_prime_t[1, :] ** 2
        )  # basically z'**2 value of the sphere at the x'(t),y'(t)
        z_max_sq[(z_max_sq) < 1e-6] = 1  # regularize zeros & negatives
        alphas = delta * (r_prime_t[2, :] / np.sqrt(z_max_sq) + 1) / 2 + alpha_min

        alphas = np.minimum(alphas, np.ones_like(alphas))

        # switch x & -y because they plot the sphere in a rotated way for some reason ...
        points = np.vstack((r_t[:, 1], -r_t[:, 0], r_t[:, 2])).T.reshape(-1, 1, 3)
        segments = np.hstack((points[:-1], points[1:]))

        lc = Line3DCollection(segments, array=alphas, cmap=cmap, **kwargs)

        self.lcs.append(lc)
        self.lc_kwargs.append({})

        return alphas

    def plot_points(self):
        # -X and Y data are switched for plotting purposes
        num_s = 0
        num_l = 0
        for k in range(len(self.points)):
            num = len(self.points[k][0])
            dist = [
                np.sqrt(
                    self.points[k][0][j] ** 2
                    + self.points[k][1][j] ** 2
                    + self.points[k][2][j] ** 2
                )
                for j in range(num)
            ]
            if any(abs(dist - dist[0]) / dist[0] > 1e-12):
                # combine arrays so that they can be sorted together
                zipped = list(zip(dist, range(num)))
                zipped.sort()  # sort rates from lowest to highest
                dist, indperm = zip(*zipped)
                indperm = np.array(indperm)
            else:
                indperm = np.arange(num)
            if self.point_style[k] == "s":
                self.axes.scatter(
                    np.real(self.points[k][1][indperm]),
                    -np.real(self.points[k][0][indperm]),
                    np.real(self.points[k][2][indperm]),
                    s=self.point_size[np.mod(k, len(self.point_size))],
                    # alpha=0.1,#1
                    edgecolor="none",
                    zdir="z",
                    **self.point_kwargs[num_s],
                )
                # color="blue",#self.point_color[mod(k, len(self.point_color))],
                # marker="o")#self.point_marker[mod(k, len(self.point_marker))])
                num_s += 1

            elif self.point_style[k] == "m":
                pnt_colors = np.array(
                    self.point_color * int(np.ceil(num / float(len(self.point_color))))
                )

                pnt_colors = pnt_colors[0:num]
                pnt_colors = list(pnt_colors[indperm])
                marker = self.point_marker[np.mod(k, len(self.point_marker))]
                s = self.point_size[np.mod(k, len(self.point_size))]
                self.axes.scatter(
                    np.real(self.points[k][1][indperm]),
                    -np.real(self.points[k][0][indperm]),
                    np.real(self.points[k][2][indperm]),
                    s=s,
                    alpha=1,
                    edgecolor="none",
                    zdir="z",
                    color=pnt_colors,
                    marker=marker,
                )

            # ONLY HERE WE CHANGED SOMETHING
            elif self.point_style[k] == "l":
                line_kwargs = self.line_kwargs[num_l]  # [mod(k, len(self.point_color))]
                self.axes.plot(
                    np.real(self.points[k][1]),
                    -np.real(self.points[k][0]),
                    np.real(self.points[k][2]),
                    zdir="z",
                    **line_kwargs,
                )
                num_l += 1

        for lc, lc_kwarg in zip(self.lcs, self.lc_kwargs):
            self.axes.add_collection(lc, **lc_kwarg)

    def render(self, **kwargs):
        super().render(**kwargs)
        self.plot_arrows()
        if self.surface_args is not None:
            self.axes.plot_surface(*self.surface_args, **self.surface_kwargs)

    def add_arrow(self, x0, xf, **arrow_kwargs):
        """Add an arrow that points from ``x0`` to ``xf`` and has the properties
        defined by arrow_kwargs

        :param x0: Starting position of the arrow
        :param xf: Position arrow points towards
        """
        self.arrow_vecs.append((x0, xf))
        self.arrow_kwargs.append(arrow_kwargs)

    def plot_arrows(self):
        """Plot the arrows."""
        # -X and Y data are switched for plotting purposes
        for vec, kwargs in zip(self.arrow_vecs, self.arrow_kwargs):
            dr = vec[1] - vec[0]
            self.axes.arrow3D(
                vec[0][1], -vec[0][0], vec[0][2], dr[1], -dr[0], dr[2], **kwargs
            )


    def plot_spin_distribution(self, H, theta_edges, phi_edges,normalize_sphere=True):
        """Plot the spin (Wigner)-distribution on the Bloch sphere. 
    
        :param H: Histogram-like distribtuion to be plotted on the Bloch sphere
        :param theta_edges: edges of the grid in theta variable, should span range 0 to pi
        :param phi_edges: edges of the grid in phi variable, should span range 0 to 2pi
        
        :param normalize_sphere: Whether or not the distribution still needs to be normalised based on the volume of the grid elements, defaults to true (i.e. the input distribution is not normalised yet)
        """
    
        if normalize_sphere==True:
            # normalize distribution according to grid bin volumes
            # Bin widths
            dphi =np.diff(phi_edges) #finding the spaces between each edge
            theta_low, theta_high = theta_edges[:-1], theta_edges[1:] #shifts the spacings by one in different directions/removes one end value
            # Bin areas on sphere: ∫∫ sinθ dθ dφ = dφ * (cos θ_low - cos θ_high)
            area_theta = np.cos(theta_low) - np.cos(theta_high)
            area = np.outer(area_theta, dphi)
            # Counts → density
            P = np.where(area > 0, H / area, 0.0) ## dividing by area but making sure it is always non-zero when dividing, where is like a vectorised if else
            # Normalize so that ∫ P(θ,φ) sinθ dθ dφ = 1,m shouldnt be necessary after already having normalised with the number of samples, but maybe due to rounding etc
            total_prob = np.sum(P * area)
            if total_prob > 0:
                P /= total_prob
            H=P # renaming variable so it has the same name even if it was already normalised
    
        print("plotting...")
        # create a colormap that fades
        # cmap_orig = cm.get_cmap("viridis")
        cmap_orig=colormaps.get_cmap("viridis")
        # Get the colormap colors
        col_vals = cmap_orig(np.arange(cmap_orig.N))
    
        # make values below 5% transparent
        alpha = np.ones(cmap_orig.N)
        alpha[: int(cmap_orig.N / 10)] = np.linspace(0, 1, int(cmap_orig.N / 10)) ** 2
        col_vals[:, -1] = alpha  # np.sqrt(np.linspace(0, 1, cmap_orig.N))
    
        # Create new colormap
        cmap_fading = ListedColormap(col_vals)
    
        
        theta_edges,phi_edges=np.meshgrid(theta_edges,phi_edges,indexing='ij')
        
        x = np.sin(theta_edges) * np.cos(phi_edges - np.pi / 2) ## WHY THE -NP.PI/2???
        y = np.sin(theta_edges) * np.sin(phi_edges - np.pi / 2)
        z = np.cos(theta_edges)
    
        norm = colors.Normalize(np.min(H),np.max(H))    
    
        # fig = plt.figure(figsize=plt.figaspect(1.))
        # ax = fig.add_subplot(111, projection='3d')
        if self.surface_args is not None:
            print("Overwriting previous surface.")
    
        self.surface_args = (x, y, z)
        self.surface_kwargs = {
            "rstride": 1,
            "cstride": 1,
            "facecolors": cmap_fading(norm(H)),
        }


## Angles
A=pd.read_csv(Path(__file__).resolve().parent.parent/"DATA"/"FIG6"/'Angles.csv')
A_data=A.to_numpy()


###### (a)
W=pd.read_csv(Path(__file__).resolve().parent.parent/"DATA"/"FIG6"/'SR_Wigner_0T.csv',header=None)
W_data=W.to_numpy()

O=pd.read_csv(Path(__file__).resolve().parent.parent/"DATA"/"FIG6"/'Superradiant_Orbit.csv')
O_data=O.to_numpy()

sphere=CustomBloch()
sphere.plot_spin_distribution(W_data,A_data[:,0],A_data[:,1],normalize_sphere=False)
sphere.add_points([2*O_data[:,0],2*O_data[:,1],2*O_data[:,2]],meth='l')
sphere.view=[120,30]
sphere.show()


###### (b)
W=pd.read_csv(Path(__file__).resolve().parent.parent/"DATA"/"FIG6"/'SR_Wigner_2.05T.csv',header=None)
W_data=W.to_numpy()

sphere=CustomBloch()
sphere.plot_spin_distribution(W_data,A_data[:,0],A_data[:,1],normalize_sphere=False)
sphere.view=[120,30]
sphere.show()


###### (c)
W=pd.read_csv(Path(__file__).resolve().parent.parent/"DATA"/"FIG6"/'SR_Wigner_8.18T.csv',header=None)
W_data=W.to_numpy()

sphere=CustomBloch()
sphere.plot_spin_distribution(W_data,A_data[:,0],A_data[:,1],normalize_sphere=False)
sphere.view=[120,30]
sphere.show()


###### (d)
W=pd.read_csv(Path(__file__).resolve().parent.parent/"DATA"/"FIG6"/'NP_Wigner_0T.csv',header=None)
W_data=W.to_numpy()

O=pd.read_csv(Path(__file__).resolve().parent.parent/"DATA"/"FIG6"/'NormalPhase_Orbit.csv')
O_data=O.to_numpy()

sphere=CustomBloch()
sphere.plot_spin_distribution(W_data,A_data[:,0],A_data[:,1],normalize_sphere=False)
sphere.add_points([2*O_data[:,0],2*O_data[:,1],2*O_data[:,2]],meth='l')
sphere.view=[120,30]
sphere.show()


###### (e)
W=pd.read_csv(Path(__file__).resolve().parent.parent/"DATA"/"FIG6"/'NP_Wigner_2.1T.csv',header=None)
W_data=W.to_numpy()

sphere=CustomBloch()
sphere.plot_spin_distribution(W_data,A_data[:,0],A_data[:,1],normalize_sphere=False)
sphere.view=[120,30]
sphere.show()


###### (f)
W=pd.read_csv(Path(__file__).resolve().parent.parent/"DATA"/"FIG6"/'NP_Wigner_5.25T.csv',header=None)
W_data=W.to_numpy()

sphere=CustomBloch()
sphere.plot_spin_distribution(W_data,A_data[:,0],A_data[:,1],normalize_sphere=False)
sphere.view=[120,30]
sphere.show()







