# %%
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import datetime as dt
import os

plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.size'] = 18


def lee_RMM(file):
    # Saltamos las primeras líneas de metadatos y leemos solo la tabla
    lectura = pd.read_csv(
        file, 
        skiprows=2,  # Línea 0: título; línea 1: encabezado descriptivo
        sep=r"\s+",  # Reemplaza 'delim_whitespace' obsoleto
        header=None,  # No usar encabezado de archivo
        usecols=range(7)  # Solo las 7 columnas numéricas
    )
    lectura.columns = ['year', 'month', 'day', 'RMM1', 'RMM2', 'phase', 'amplitude']
    # Vector de fechas
    fechas = [dt.datetime(y, m, d) for y, m, d in zip(lectura.year, lectura.month, lectura.day)]
    # Limpiar valores extremos
    RMM1 = lectura.RMM1.copy()
    RMM2 = lectura.RMM2.copy()
    RMM1[RMM1 > 100] = np.nan
    RMM2[RMM2 > 100] = np.nan
    # Crear DataFrame final
    data = np.array([RMM1, RMM2, lectura.phase])
    df = pd.DataFrame(data.T, columns=['RMM1', 'RMM2', 'phase'], index=fechas)

    return df


# descargar manualmente el index en 
# https://www-bom-gov-au.translate.goog/clim_data/IDCKGEM000/rmm.74toRealtime.txt?_x_tr_sl=en&_x_tr_tl=es&_x_tr_hl=es&_x_tr_pto=tc&_x_tr_sch=http
    

path_ima = '/home/ccruz/0_Proyecto/Scripts/4_Analisis_MJO/Imagenes/Conteo_extremos/'
output_file = os.path.join('/home/ccruz/0_Proyecto/Scripts/4_Analisis_MJO/0_Data/', 'rmm.74toRealtime2.txt')
df_rmm = lee_RMM(output_file)


# Ajustar rango fechas RMM al rango de df_extremos
df_rmm = df_rmm['2025-06':'2025-07-31']

rmm1 = df_rmm['RMM1']
rmm2 = df_rmm['RMM2']
fechas = df_rmm.index
dias = [f.day for f in fechas]
meses = [f.month for f in fechas]

# Colores fijos a usar (rojo, azul y verde)
colores = ['red', 'blue', 'green']

# Meses únicos ordenados
meses_unicos = sorted(set(meses))

# Asignar color cíclicamente a cada mes único
colores_por_mes = {mes: colores[i % len(colores)] for i, mes in enumerate(meses_unicos)}

fig, ax = plt.subplots(figsize=(8, 8))

# Círculo de amplitud 1 en gris
circle = plt.Circle((0, 0), 1, color='gray', fill=False, linestyle='--', linewidth=1.2)
ax.add_artist(circle)

# Líneas punteadas entre fases
for angle in np.arange(0, 360, 45):
    x = [0, 4 * np.cos(np.deg2rad(angle))]
    y = [0, 4 * np.sin(np.deg2rad(angle))]
    ax.plot(x, y, linestyle='dotted', color='gray', linewidth=0.8)

# Dibujar trayectoria por mes
handles = []
for mes in meses_unicos:
    idx = [i for i, m in enumerate(meses) if m == mes]
    color = colores_por_mes[mes]
    linea, = ax.plot(rmm1.iloc[idx], rmm2.iloc[idx], '-', color=color, linewidth=1.1)
    handles.append(linea)
    
    for i in idx:
        ax.text(rmm1.iloc[i], rmm2.iloc[i], str(dias[i]), fontsize=9,  # más pequeño
                ha='center', va='center', color='black')  # color negro

# START
ax.text(rmm1.iloc[0], rmm2.iloc[0], 'START', fontsize=9, color='g', ha='left', va='bottom', weight='bold')

# Ejes
ax.axhline(0, color='black', linewidth=0.6)
ax.axvline(0, color='black', linewidth=0.6)
ax.set_xlim(-4, 4)
ax.set_ylim(-4, 4)
ax.set_xlabel('RMM1')
ax.set_ylabel('RMM2')

fecha_inicio = df_rmm.index.min().strftime('%Y-%m-%d')
fecha_fin = df_rmm.index.max().strftime('%Y-%m-%d')

ax.set_title(f'Diagrama de fases MJO\nPeriodo: {fecha_inicio} a {fecha_fin}', fontsize=16)

# Fases con etiquetas más visibles - CORREGIDO
etiquetas_fase = {
    1: 'West. Hem.\nand Africa',
    2: 'Indian\nOcean',
    3: 'Indian\nOcean',
    4: 'Maritime\nContinent',
    5: 'Maritime\nContinent',
    6: 'Western\nPacific',
    7: 'Western\nPacific',
    8: 'West. Hem.\nand Africa',
}

# CORRECCIÓN: Posiciones exactas según BOM
# Desde la izquierda en sentido antihorario: 7,8,1,2
# Fase 8 debe estar arriba-izquierda, fase 1 abajo-izquierda, etc.
angulos_fases = {
    7: 135-15,  # superior izquierda (West. Hem. and Africa)
    1: 225-15,  # inferior izquierda (West. Hem. and Africa) 
    2: 270-15,  # inferior (Indian Ocean)
    3: 315-15,  # inferior derecha (Indian Ocean)
    4: 0-15,    # derecha (Maritime Continent)  
    5: 45-15,   # superior derecha (Maritime Continent)
    6: 90-15,   # superior (Western Pacific)
    8: 180-15   # izquierda (Western Pacific)
}

for fase in range(1, 9):
    angle = np.deg2rad(angulos_fases[fase])

    # Número de fase
    x = 3.0 * np.cos(angle)
    y = 3.0 * np.sin(angle)
    ax.text(x, y, str(fase), ha='center', va='center', fontsize=14, fontweight='bold')


angulos_fases = {
    1: 180,  # inferior izquierda (West. Hem. and Africa) 
    2: 270,  # inferior (Indian Ocean)
    5: 0,   # superior derecha (Maritime Continent)
    6: 90,   # superior (Western Pacific)
}

for fase in angulos_fases.keys():
    angle = np.deg2rad(angulos_fases[fase])

    # Nombre de la fase
    x_txt = 3.6 * np.cos(angle)
    y_txt = 3.6 * np.sin(angle)
    if fase == 1:
        ax.text(x_txt, y_txt, etiquetas_fase[fase], ha='center', va='center', fontsize=14, rotation = 90)
    
    if fase == 5:
        ax.text(x_txt, y_txt, etiquetas_fase[fase], ha='center', va='center', fontsize=14, rotation = 270)
    
    if fase != 1 and fase != 5:
        ax.text(x_txt, y_txt, etiquetas_fase[fase], ha='center', va='center', fontsize=14)

# Ajustar márgenes para que todo quepa
plt.subplots_adjust(left=0.1, right=0.9, top=0.9, bottom=0.1)

# Leyenda por mes
labels_leyenda = [dt.date(1900, m, 1).strftime('%b').upper() for m in meses_unicos]
ax.legend(handles, labels_leyenda, loc='upper right', fontsize=12, title="Mes", title_fontsize=12)

plt.tight_layout()

plt.savefig(path_ima+'RMM.png', bbox_inches='tight', dpi=200)
plt.close()