'''
- Codigos para filtrar datos simulados en un intervalo de -z a z (valor)
- Graficas para matrices con ceros y cambio de coordenadas en estas matrices
- Interpolacion de puntos 
'''

import numpy as np
import orthopoly
import numpy as np
import time as timeglobal
from scipy.spatial import KDTree
from tqdm import tqdm
from colorama import Fore, Style
from scipy.interpolate import splprep, splev
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
# from Codigo1 import mcart_msph

def public_c2s(posicionescart):
    '''
    De cartesianas a esfericas (c2s) de forma publica
    '''
    x,y,z       = posicionescart.T
    r,theta,phi = orthopoly.spherical_harmonic.cart2sph(-x,-y,z) #Bug in library
    return np.asarray([r,theta,phi]).T

def public_s2c(posicionesesf):
    '''
    De esfericas a cartesianas (s2c) de forma publica
    '''
    r,theta,phi = posicionesesf.T
    x,y,z       = orthopoly.spherical_harmonic.sph2cart(r,theta,phi)
    return np.asarray([x,y,z]).T

def mcart_msph(matrizcart):
    tiempo = matrizcart.shape[0]
    matrizsph = (public_c2s(matrizcart.reshape(-1, 3))).reshape(tiempo, -1, 3)
    return matrizsph

def msph_mcart_filtrado(matrizsph):
    """
    Transforma las filas de una matriz de coordenadas cartesianas a esféricas,
    ignorando las filas con valores [0, 0, 0].

    Parameters:
    matrizcart (numpy.ndarray): Matriz de coordenadas cartesianas [tiempo, n, 3].

    Returns:
    numpy.ndarray: Matriz transformada a coordenadas esféricas, 
                   manteniendo las filas [0, 0, 0] sin cambiar.
    """
    tiempo = matrizsph.shape[0]
    matrizcart = np.zeros_like(matrizsph)  

    for t in range(tiempo):
        frame = matrizsph[t]
        
        mask = ~np.all(frame == 0, axis=1) # solo seleccionamos filas que sean distintas de 0 (sino sale error en arcos)

        matrizcart[t, mask] = public_s2c(frame[mask])

    return matrizcart

def mcart_msph_filtrado(matrizcart):
    """
    Transforma las filas de una matriz de coordenadas cartesianas a esféricas,
    ignorando las filas con valores [0, 0, 0].

    Parameters:
    matrizcart (numpy.ndarray): Matriz de coordenadas cartesianas [tiempo, n, 3].

    Returns:
    numpy.ndarray: Matriz transformada a coordenadas esféricas, 
                   manteniendo las filas [0, 0, 0] sin cambiar.
    """
    tiempo = matrizcart.shape[0]
    matrizsph = np.zeros_like(matrizcart)  

    for t in range(tiempo):
        frame = matrizcart[t]
        
        mask = ~np.all(frame == 0, axis=1) # solo seleccionamos filas que sean distintas de 0 (sino sale error en arcos)

        matrizsph[t, mask] = public_c2s(frame[mask])

    return matrizsph

def contarpuntos(valor,matrix):

    
    tiempo, particulas, _ = matrix.shape

    contador = np.zeros(tiempo, dtype=int)

    for t in range(tiempo):

        z_coords = matrix[t, :, 2]
        contador[t] = np.sum((-valor <= z_coords) & (z_coords <= valor))
      
    return contador

def datosfiltrados(valor,matrix):
    
    contador = contarpuntos(valor,matrix)
    
    fecha_actual = timeglobal.strftime("%d-%b-%Y %H:%M:%S")
    print(f"{fecha_actual} - Análisis3 - Contando datos con Z ∈ [-{valor}, {valor}] ...")
    print("El numero de puntos esta entre ",np.min(contador)," y ",np.max(contador))
    print(f"{fecha_actual} - Análisis3 - Filtrando datos de las simulaciones...")

    tiempo = matrix.shape[0]
    max_contador = np.max(contador)
    matriz_homogenea = np.zeros((tiempo, max_contador, 3), dtype=float)

    for t in range(tiempo):
        puntos = matrix[t, :, :]
        puntos_en_rango = puntos[(-valor <= puntos[:, 2]) & (puntos[:, 2] <= valor)]
        matriz_homogenea[t, :len(puntos_en_rango), :] = puntos_en_rango

    return matriz_homogenea

def distribuirtraj_sobregrid(traj_sph, grid_sph):
    
    '''
    Funcion analoga a la de distribucion de Codigo1
    '''
    fecha_actual = timeglobal.strftime("%d-%b-%Y %H:%M:%S")
    n_traj = traj_sph.shape[1]
    n_grid_points = grid_sph.shape[0] 
    
    print("------------------------------")
    print(f"{fecha_actual} - Análisis3 - Distribuyendo {Style.BRIGHT + Fore.GREEN}{n_traj}{Style.RESET_ALL} puntos de trayectoria en {Style.BRIGHT + Fore.BLUE}{n_grid_points}{Style.RESET_ALL} puntos del grid.")

    print(f"Se está usando {Style.BRIGHT + Fore.GREEN}{traj_sph.shape[0]}{Style.RESET_ALL} valores de tiempo")

    
    theta_grid = grid_sph[:, 1]
    phi_grid = grid_sph[:, 2]

    time = traj_sph.shape[0]
    # n_grid_points = grid_sph.shape[0]

    grid_kdtree = KDTree(grid_sph)

    traj_grid_sph = np.empty((time, n_grid_points, 3))

    for t, frame_pos in enumerate(tqdm(traj_sph, total=time, desc="Procesando")): # Para cada instante de tiempo
        
        frame_pos = frame_pos[~np.all(frame_pos == 0, axis=1)] # Filtrar valores distintos de cero

        if frame_pos.size == 0:
            raise Exception(f"Ningun punto encontrado para el tiempo {t}. Revisar los datos de trayectoria.")

        r, theta, phi = frame_pos.T

        r = np.ones_like(r) # radio 1
        traj_over_sphere = np.vstack([r, theta, phi]).T

        frame_grid_positions = grid_kdtree.query(traj_over_sphere)[1]

        for i in range(n_grid_points):
            mask = (frame_grid_positions == i)

            if not np.any(mask):
                raise Exception(f"Punto del grid {i} sin datos. Pruebe a reducir el numero de puntos del grid.")

            r_vals, theta_vals, phi_vals = frame_pos[mask].T

            traj_grid_sph[t, i] = np.array([np.mean(r_vals), theta_grid[i], phi_grid[i]])

    return traj_grid_sph # devuelve en esfericas

# %% GRAFICAS

def graficapuntual(matriz_homogenea, t, rangoz,title):
    """
    Grafica los puntos almacenados en la matriz homogénea en 3D para un instante de tiempo específico,
    ignorando los valores cero.

    Parameters:
    matriz_homogenea (numpy.ndarray): Matriz de dimensiones [tiempo, max_contador, 3].
    t (int): Instante de tiempo para el cual se desea graficar los puntos.
    """
    if t < 0 or t >= matriz_homogenea.shape[0]:
        raise ValueError("El tiempo especificado está fuera del rango válido.")
    
    puntos = matriz_homogenea[t]
    
    puntos_filtrados = puntos[~np.all(puntos == 0, axis=1)]
    
    if len(puntos_filtrados) > 0:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        
        x = puntos_filtrados[:, 0]
        y = puntos_filtrados[:, 1]
        z = puntos_filtrados[:, 2]
        
        ax.scatter(x, y, z, color = "blue", linestyle='None')
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        ax.set_zlim(-rangoz, rangoz)
        ax.set_title(f"{title} en t = {t}", fontweight="bold")
        
        plt.show()
    else:
        print(f"No hay puntos válidos para el tiempo {t}.")
        
def gengraficav2(matrix, t, title, valor):
    '''
    Funcion identica a gengrafica solo que ahora tiene un input para los limites en z
    --> Util para comprobar que la funcion de filtrado actue adecuadamente
    '''

    a3 = matrix[t][:, 0]
    b3 = matrix[t][:, 1]
    c3 = matrix[t][:, 2]

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(a3, b3, c3, marker='o', color="blue")

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.set_title(f"{title} en t = {t}", fontweight="bold")

    ax.set_zlim(-valor, valor)

    plt.show()
    
# %% CONTINUO

# from Codigo1 import coefs,gencircunferencia

# def expandir(MSph,npoints,Lmax):
#     '''
#     # Es como interpolar mediante armonicos esfericos (no funciona!)
#     '''
#     _,coe = coefs(MSph,Lmax)
#     Matrizcart,Matrizsph = gencircunferencia(coe,npoints) # selecciono las esfericas en concreto
#     return Matrizcart,Matrizsph # cartesianas

# def interpolarpuntos(matrizcart, n):
#     '''
#     Interpola los puntos en coordenadas cartesianas para cada instante de tiempo.
#     ver: https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.splprep.html
#     '''
    
#     print("------------------")
#     fecha_actual = timeglobal.strftime("%d-%b-%Y %H:%M:%S")
#     print(f"{fecha_actual} - Análisis3 - Interpolando {matrizcart.shape[1]} .")
#     print(f"Se han obtenido {n} puntos tras la interpolación.")

    
#     tiempo = matrizcart.shape[0]
#     M1 = np.empty((tiempo, n+1, 3)) # cerrada con termino repetido
#     minterpolada_cart = np.empty((tiempo,n,3)) # lo que buscamos
#     u_new = np.linspace(0, 1, n+1) # siempre es cte, creamos un punto mas de lo normal (para forzar cierre) 
#     znew = np.zeros((n+1)) # vector nulo (ya que se supone que el input es este mismo)
    
#     for t in range(tiempo):

#         frame = matrizcart[t]

#         mask = ~np.all(frame == 0, axis=1)
#         frame = frame[mask]

#         if frame.shape[0] < 2:
#             raise ValueError(f"No hay suficientes puntos para interpolar en el tiempo {t}.")
        
#         # Datos sin cerrar
#         x = frame[:,0]
#         y = frame[:,1]
        
#         # Forzamos cierre
#         x_closed = np.append(x, x[0])
#         y_closed = np.append(y, y[0])
        
#         tck,_ = splprep([x_closed, y_closed], s=0) # tck generado a partir de los valores cerrados

#         newpoints = splev(u_new, tck) # generamos puntos de dimension n+1, pero el ultimo termino es el primero

#         M1[t] = np.vstack([newpoints[0], newpoints[1], znew]).T
        
#         minterpolada_cart[t][:,:] = M1[t][0:-1,:] # no contamos el ultimo termino (repetido)
    
#         if np.mean(abs(minterpolada_cart[t, :, 2])) <= 0.00001: # para que la grafica salga exactamente en 0
#             minterpolada_cart[t,:,2] = 0
        
#     minterpolada_sph = mcart_msph(minterpolada_cart)
#     return minterpolada_cart,minterpolada_sph